From e3006ed2d77eeebde370ca7a007bd6a5746409e7 Mon Sep 17 00:00:00 2001
From: 2877992943 <2877992943@qq.com>
Date: Fri, 18 May 2018 15:05:14 +0800
Subject: [PATCH 0001/2720] beam decode slow

---
 tensor2tensor/utils/t2t_model.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 3e6eb07e8..b71f1ed16 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -644,7 +644,9 @@ def symbols_to_logits_fn(ids):
     # Setting decode length to input length + decode_length
     decode_length = tf.constant(decode_length)
     if "partial_targets" not in features:
-      decode_length += common_layers.shape_list(features["inputs"])[1]
+      inputs = features["inputs"]
+      decode_length = (common_layers.shape_list(inputs)[1] +
+                       features.get("decode_length", decode_length))
     ids, scores = beam_search.beam_search(
         symbols_to_logits_fn,
         initial_ids,

From 79a64782b12d002b709562d4af20a2a63b20f675 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 21 May 2018 16:34:13 -0700
Subject: [PATCH 0002/2720] Correct name-clash bug in RealModality.

PiperOrigin-RevId: 197481067
---
 tensor2tensor/layers/modalities.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index a6789a7ba..871d35a12 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -693,11 +693,11 @@ class RealModality(modality.Modality):
 
   def bottom(self, x):
     with tf.variable_scope("real"):
-      return tf.layers.dense(x, self._body_input_depth)
+      return tf.layers.dense(x, self._body_input_depth, name="bottom")
 
   def top(self, body_output, _):
     with tf.variable_scope("real"):
-      return tf.layers.dense(body_output, self._vocab_size)
+      return tf.layers.dense(body_output, self._vocab_size, name="top")
 
   def loss(self, top_out, targets):
     raise NotImplementedError()

From e4c738f031f6e53fa2aefa2495cf6652aab66793 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 22 May 2018 13:40:35 -0700
Subject: [PATCH 0003/2720] Set deterministic function name.

PiperOrigin-RevId: 197615718
---
 tensor2tensor/layers/common_layers.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 18401cc28..e3451b4b3 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -23,7 +23,6 @@
 import functools
 from functools import partial
 import math
-import random
 
 # Dependency imports
 
@@ -2528,7 +2527,7 @@ def custom_grad_fn(op, *dys):
 
   @function.Defun(
       *(in_types + var_types + out_types),
-      func_name="identity_custom_grad%d" % random.randint(1, 10**9),
+      func_name="identity_custom_grad%d" % ops.uid(),
       python_grad_func=custom_grad_fn,
       shape_func=lambda _: [t.get_shape() for t in outputs])
   def identity(*args):

From 7d47c8199eea03845b989674e2006350f73b9f1e Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Tue, 22 May 2018 14:07:15 -0700
Subject: [PATCH 0004/2720] Refactor vector quantized bottlenecks to follow new
 API

PiperOrigin-RevId: 197620257
---
 tensor2tensor/layers/discretization.py      | 85 +++++++++++++++++++++
 tensor2tensor/layers/discretization_test.py | 42 ++++++++++
 2 files changed, 127 insertions(+)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index f72bc1d5a..98c19eecd 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -672,6 +672,86 @@ def discrete_bottleneck(x,
 # * The [method]_unbottleneck function moves from discretized state to dense.
 
 
+def get_vq_bottleneck(bottleneck_size, hidden_size):
+  """Get lookup table for VQ bottleneck."""
+  with tf.variable_scope("vq", reuse=tf.AUTO_REUSE):
+    means = tf.get_variable(
+        name="means",
+        shape=[bottleneck_size, hidden_size],
+        initializer=tf.uniform_unit_scaling_initializer())
+
+    ema_count = tf.get_variable(
+        name="ema_count",
+        shape=[bottleneck_size],
+        initializer=tf.constant_initializer(0),
+        trainable=False)
+
+    with tf.colocate_with(means):
+      ema_means = tf.get_variable(
+          name="ema_means",
+          initializer=means.initialized_value(),
+          trainable=False)
+  return means, ema_means, ema_count
+
+
+def vq_nearest_neighbor(x, means):
+  """Find the nearest element in means to elements in x."""
+  bottleneck_size = common_layers.shape_list(means)[0]
+  x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keepdims=True)
+  means_norm_sq = tf.reduce_sum(tf.square(means), axis=-1, keepdims=True)
+  scalar_prod = tf.matmul(x, means, transpose_b=True)
+  dist = x_norm_sq + tf.transpose(means_norm_sq) - 2 * scalar_prod
+  x_means_idx = tf.argmax(-dist, axis=-1)
+  x_means_hot = tf.one_hot(x_means_idx, bottleneck_size)
+  x_means_hot_flat = tf.reshape(x_means_hot, [-1, bottleneck_size])
+  x_means = tf.matmul(x_means_hot_flat, means)
+  e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
+  return x_means_hot, e_loss
+
+
+def vq_discrete_bottleneck(x,
+                           bottleneck_size,
+                           beta=0.25,
+                           decay=0.999,
+                           epsilon=1e-5):
+  """Simple vector quantized discrete bottleneck."""
+  hidden_size = common_layers.shape_list(x)[-1]
+  means, ema_means, ema_count = get_vq_bottleneck(bottleneck_size, hidden_size)
+  x_means_hot, e_loss = vq_nearest_neighbor(x, means)
+
+  # Update the ema variables
+  updated_ema_count = moving_averages.assign_moving_average(
+      ema_count,
+      tf.reduce_sum(
+          tf.reshape(x_means_hot, shape=[-1, bottleneck_size]), axis=0),
+      decay,
+      zero_debias=False)
+
+  dw = tf.matmul(x_means_hot, x, transpose_a=True)
+
+  updated_ema_means = moving_averages.assign_moving_average(
+      ema_means, dw, decay, zero_debias=False)
+  n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
+  updated_ema_count = (
+      (updated_ema_count + epsilon) / (n + bottleneck_size * epsilon) * n)
+  updated_ema_means /= tf.expand_dims(updated_ema_count, axis=-1)
+  with tf.control_dependencies([e_loss]):
+    update_means = tf.assign(means, updated_ema_means)
+    with tf.control_dependencies([update_means]):
+      loss = beta * e_loss
+
+  d = x_means_hot
+  return d, loss
+
+
+def vq_discrete_unbottleneck(x, hidden_size):
+  """Simple undiscretization from vector quantized representation."""
+  x = tf.to_float(x)
+  bottleneck_size = common_layers.shape_list(x)[1]
+  means, _, _ = get_vq_bottleneck(bottleneck_size, hidden_size)
+  return tf.matmul(x, means)
+
+
 def tanh_discrete_bottleneck(x, bottleneck_size, bottleneck_noise,
                              discretize_warmup_steps, mode):
   """Simple discretization through tanh, flip bottleneck_noise many bits."""
@@ -738,6 +818,9 @@ def parametrized_bottleneck(x, hparams):
         x, hparams.bottleneck_size, hparams.bottleneck_noise * 0.5,
         hparams.discretize_warmup_steps, hparams.mode,
         hparams.isemhash_noise_dev, hparams.isemhash_mix_prob)
+  if hparams.bottleneck_kind == "vq":
+    return vq_discrete_bottleneck(x, hparams.bottleneck_size, hparams.beta,
+                                  hparams.decay, hparams.epsilon)
   raise ValueError("Unsupported hparams.bottleneck_kind %s"
                    % hparams.bottleneck_kind)
 
@@ -749,5 +832,7 @@ def parametrized_unbottleneck(x, hidden_size, hparams):
   if hparams.bottleneck_kind == "isemhash":
     return isemhash_unbottleneck(
         x, hidden_size, hparams.isemhash_filter_size_multiplier)
+  if hparams.bottleneck_kind == "vq":
+    return vq_discrete_unbottleneck(x, hidden_size)
   raise ValueError("Unsupported hparams.bottleneck_kind %s"
                    % hparams.bottleneck_kind)
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index 25576e09d..8873e6b61 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -119,6 +119,48 @@ def testNearestNeighbors(self):
       self.assertEqual(np.shape(x_means_hot_eval), (1, 2, 4))
       self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
 
+  def testGetVQBottleneck(self):
+    bottleneck_size = 4
+    hidden_size = 3
+    means, _, ema_count = discretization.get_vq_bottleneck(bottleneck_size,
+                                                           hidden_size)
+    assign_op = means.assign(tf.zeros(shape=[bottleneck_size, hidden_size]))
+    means_new, _, _ = discretization.get_vq_bottleneck(bottleneck_size,
+                                                       hidden_size)
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      sess.run(assign_op)
+      self.assertTrue(np.all(sess.run(means_new) == 0))
+      self.assertTrue(np.all(sess.run(ema_count) == 0))
+
+  def testVQNearestNeighbors(self):
+    x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
+    means = tf.constant(
+        [[1, 0, 0], [0, 1, 0], [0, 0, 1], [9, 9, 9]], dtype=tf.float32)
+    x_means_hot, _ = discretization.vq_nearest_neighbor(x, means)
+    x_means_hot_test = np.array([[0, 1, 0, 0], [1, 0, 0, 0]])
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      x_means_hot_eval = sess.run(x_means_hot)
+      self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
+      self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
+
+  def testVQDiscreteBottleneck(self):
+    x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
+    x_means_hot, _ = discretization.vq_discrete_bottleneck(x, bottleneck_size=4)
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      x_means_hot_eval = sess.run(x_means_hot)
+      self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
+
+  def testVQDiscreteUnbottlenck(self):
+    x = tf.constant([[1, 0, 0, 0], [0, 0, 1, 0]], dtype=tf.int32)
+    x_means = discretization.vq_discrete_unbottleneck(x, hidden_size=3)
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      x_means_eval = sess.run(x_means)
+      self.assertEqual(np.shape(x_means_eval), (2, 3))
+
 
 if __name__ == '__main__':
   tf.test.main()

From 602b63d27ee6ab433f2468bd9334905c97449d77 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 22 May 2018 16:27:52 -0700
Subject: [PATCH 0005/2720] Make sure the types are floats in RealModality
 bottom.

PiperOrigin-RevId: 197643088
---
 tensor2tensor/layers/modalities.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 871d35a12..c136c293d 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -693,7 +693,8 @@ class RealModality(modality.Modality):
 
   def bottom(self, x):
     with tf.variable_scope("real"):
-      return tf.layers.dense(x, self._body_input_depth, name="bottom")
+      return tf.layers.dense(tf.to_float(x), self._body_input_depth,
+                             name="bottom")
 
   def top(self, body_output, _):
     with tf.variable_scope("real"):

From 6abe738d80c785c1592969051e9fbdc64fa8ec84 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 22 May 2018 22:52:59 -0700
Subject: [PATCH 0006/2720] Multi timeseries forecasting using transformer
 model.

PiperOrigin-RevId: 197676914
---
 tensor2tensor/data_generators/text_encoder.py |  30 ++++
 tensor2tensor/data_generators/timeseries.py   | 130 ++++++++++++++++++
 .../data_generators/timeseries_test.py        |  64 +++++++++
 3 files changed, 224 insertions(+)
 create mode 100644 tensor2tensor/data_generators/timeseries.py
 create mode 100644 tensor2tensor/data_generators/timeseries_test.py

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 699ad001a..c256f94b8 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -1005,3 +1005,33 @@ def decode_list(self, ids):
   @property
   def vocab_size(self):
     return 256
+
+
+class RealEncoder(object):
+  """Encoder class for saving and loading float values."""
+
+  def encode(self, s):
+    """Transform a string (space separated float values) into a float array.
+
+    Args:
+      s: space separated float values.
+
+    Returns:
+      Array of float values.
+    """
+    return [float(w) for w in s.split()]
+
+  def decode(self, ids):
+    """Transform sequence of float values into string (float values).
+
+    Args:
+      ids: array of floats to be converted.
+
+    Returns:
+      String having space separated float values.
+
+    Raises:
+      ValueError: if the ids are not of the appropriate size.
+    """
+    return " ".join(ids)
+  
\ No newline at end of file
diff --git a/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
new file mode 100644
index 000000000..07e32d39a
--- /dev/null
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -0,0 +1,130 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Multi time series forecasting problem."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import random
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.utils import metrics
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+@registry.register_problem
+class TimeSeriesToyProblem(problem.Problem):
+  """Base Problem for multi timeseries for datasets."""
+
+  def __init__(self,
+               was_reversed=False,
+               was_copy=False,
+               num_train_shards=9,
+               num_eval_shards=1,
+               num_samples=100):
+    super(TimeSeriesToyProblem, self).__init__(was_reversed, was_copy)
+    self._num_train_shards = num_train_shards
+    self._num_eval_shards = num_eval_shards
+    self._num_samples = num_samples
+
+  def feature_encoders(self, data_dir):
+    del data_dir
+    return {
+        'inputs': text_encoder.RealEncoder(),
+        'targets': text_encoder.RealEncoder()
+    }
+
+  @property
+  def is_generate_per_split(self):
+    # generate_data will shard the data into TRAIN and EVAL for us.
+    return False
+
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each."""
+    # 10% evaluation data
+    return [{
+        'split': problem.DatasetSplit.TRAIN,
+        'shards': self._num_train_shards,
+    }, {
+        'split': problem.DatasetSplit.EVAL,
+        'shards': self._num_eval_shards,
+    }]
+
+  def eval_metrics(self):
+    eval_metrics = [metrics.Metrics.RMSE]
+
+    return eval_metrics
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    del data_dir
+    del tmp_dir
+    del dataset_split
+
+    series_1 = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
+
+    # This generates _num_samples instances of each possible split of series_1;
+    # inputs & targets are of variable size.
+    for x in range(self._num_samples):
+      split_index = random.randint(1, 9)
+      inputs, targets = series_1[:split_index], series_1[split_index:]
+      example_keys = ['inputs', 'targets']
+      ex_dict = dict(zip(example_keys, [inputs, targets]))
+      print('Inputs & Targets #', x, ':', ex_dict)
+      yield ex_dict
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    p.input_modality = {'inputs': (registry.Modalities.REAL, 1)}
+    p.target_modality = (registry.Modalities.REAL, 1)
+    p.input_space_id = problem.SpaceID.REAL
+    p.target_space_id = problem.SpaceID.REAL
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    filepath_fns = {
+        problem.DatasetSplit.TRAIN: self.training_filepaths,
+        problem.DatasetSplit.EVAL: self.dev_filepaths,
+        problem.DatasetSplit.TEST: self.test_filepaths,
+    }
+
+    split_paths = [(split['split'], filepath_fns[split['split']](
+        data_dir, split['shards'], shuffled=False))
+                   for split in self.dataset_splits]
+
+    all_paths = []
+    for _, paths in split_paths:
+      all_paths.extend(paths)
+
+    if self.is_generate_per_split:
+      for split, paths in split_paths:
+        generator_utils.generate_files(
+            self.generate_samples(data_dir, tmp_dir, split), paths)
+    else:
+      generator_utils.generate_files(
+          self.generate_samples(data_dir, tmp_dir, problem.DatasetSplit.TRAIN),
+          all_paths)
+
+    generator_utils.shuffle_dataset(all_paths)
+
+  def example_reading_spec(self):
+    data_fields = {
+        'inputs': tf.VarLenFeature(tf.float32),
+        'targets': tf.VarLenFeature(tf.float32),
+    }
+    data_items_to_decoders = None
+    return (data_fields, data_items_to_decoders)
diff --git a/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
new file mode 100644
index 000000000..e0f1f3558
--- /dev/null
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -0,0 +1,64 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Timeseries generators tests."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import shutil
+# Dependency imports
+
+from tensor2tensor.data_generators import timeseries
+
+import tensorflow as tf
+
+
+class TimeseriesTest(tf.test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    cls.tmp_dir = tf.test.get_temp_dir()
+    shutil.rmtree(cls.tmp_dir)
+    os.mkdir(cls.tmp_dir)
+
+  def testTimeSeriesToyProblem(self):
+    problem = timeseries.TimeSeriesToyProblem(
+        num_train_shards=1, num_eval_shards=1, num_samples=4)
+    problem.generate_data(self.tmp_dir, self.tmp_dir)
+
+    dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, self.tmp_dir)
+    features = dataset.make_one_shot_iterator().get_next()
+
+    examples = []
+    exhausted = False
+    with self.test_session() as sess:
+      examples.append(sess.run(features))
+      examples.append(sess.run(features))
+      try:
+        sess.run(features)
+      except tf.errors.OutOfRangeError:
+        exhausted = True
+
+    self.assertTrue(exhausted)
+    self.assertEqual(2, len(examples))
+
+    self.assertNotEqual(
+        list(examples[0]["inputs"]), list(examples[1]["inputs"]))
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 41dd31aa2f72ae0c39dc42c65084876158b428d1 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 23 May 2018 11:15:53 -0700
Subject: [PATCH 0007/2720] Boilerplace for simple generated video problem with
 moving shapes.

PiperOrigin-RevId: 197755812
---
 tensor2tensor/data_generators/all_problems.py |  1 +
 .../data_generators/video_generated.py        | 87 +++++++++++++++++++
 2 files changed, 88 insertions(+)
 create mode 100644 tensor2tensor/data_generators/video_generated.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 2244697c5..93deecf61 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -57,6 +57,7 @@
     "tensor2tensor.data_generators.translate_envi",
     "tensor2tensor.data_generators.translate_enzh",
     "tensor2tensor.data_generators.twentybn",
+    "tensor2tensor.data_generators.video_generated",
     "tensor2tensor.data_generators.wiki",
     "tensor2tensor.data_generators.wikisum.wikisum",
     "tensor2tensor.data_generators.wikitext103",
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
new file mode 100644
index 000000000..2db7b45fb
--- /dev/null
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -0,0 +1,87 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generators for video problems with artificially generated frames."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# Dependency imports
+
+import numpy as np
+
+from tensor2tensor.data_generators import video_utils
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+@registry.register_problem
+class VideoStochasticShapes10k(video_utils.VideoProblem):
+  """Shapes moving in a stochastic way."""
+
+  @property
+  def num_input_frames(self):
+    """Number of frames to batch on one input."""
+    return 4
+
+  @property
+  def num_target_frames(self):
+    """Number of frames to predict in one step."""
+    return 1
+
+  @property
+  def is_generate_per_split(self):
+    """Whether we have a train/test split or just hold out data."""
+    return False  # Just hold out some generated data for evals.
+
+  @property
+  def frame_height(self):
+    return 64
+
+  @property
+  def frame_width(self):
+    return 64
+
+  @property
+  def total_number_of_frames(self):
+    return 10000
+
+  @property
+  def extra_reading_spec(self):
+    """Additional data fields to store on disk and their decoders."""
+    data_fields = {
+        "frame_number": tf.FixedLenFeature([1], tf.int64),
+    }
+    decoders = {
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
+            tensor_key="frame_number"),
+    }
+    return data_fields, decoders
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    p.input_modality = {
+        "inputs": ("video", 256),
+        "input_frame_number": ("symbol:identity", 1)
+    }
+
+  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
+    frame_number = 0
+    for _ in range(self.total_number_of_frames):
+      frame = np.zeros([self.frame_height, self.frame_width, self.num_channels],
+                       dtype=np.uint8)
+      yield {"frame": frame, "frame_number": [frame_number]}
+      frame_number += 1

From 8ef895f95fe2992e0b89822bda3a274cf7ec7878 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 23 May 2018 13:21:34 -0700
Subject: [PATCH 0008/2720] Rename basic_conv_gen to next_frame and add
 prefixed VQ hparams to autoencoders.

PiperOrigin-RevId: 197778012
---
 tensor2tensor/bin/t2t_trainer.py              |  2 +
 tensor2tensor/data_generators/gym_problems.py |  3 +-
 .../data_generators/video_generated.py        |  1 +
 tensor2tensor/layers/discretization.py        |  4 +-
 tensor2tensor/models/__init__.py              |  2 +-
 tensor2tensor/models/research/autoencoders.py |  3 +
 .../{basic_conv_gen.py => next_frame.py}      | 55 ++++++++++---------
 tensor2tensor/rl/model_rl_experiment.py       | 28 +++++-----
 tensor2tensor/rl/rl_trainer_lib.py            |  3 +-
 9 files changed, 56 insertions(+), 45 deletions(-)
 rename tensor2tensor/models/research/{basic_conv_gen.py => next_frame.py} (88%)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 1ad6c67d1..21eee7662 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -25,6 +25,7 @@
 
 from tensor2tensor import models  # pylint: disable=unused-import
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
+from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
 from tensor2tensor.utils import cloud_mlengine
 from tensor2tensor.utils import cloud_tpu
 from tensor2tensor.utils import decoding
@@ -332,6 +333,7 @@ def main(argv):
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
   log_registry()
 
+
   if FLAGS.cloud_mlengine:
     return cloud_mlengine.launch()
 
diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 3b1df390d..ffdce69dd 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -613,7 +613,8 @@ def real_env(self):
   def restore_networks(self, sess):
     super(GymSimulatedDiscreteProblemWithAgent, self).restore_networks(sess)
     # TODO(blazej): adjust regexp for different models.
-    env_model_loader = tf.train.Saver(tf.global_variables("basic_conv_gen.*"))
+    env_model_loader = tf.train.Saver(tf.global_variables(
+        "next_frame_basic.*"))
     sess = tf.get_default_session()
 
     ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 2db7b45fb..43e049f11 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -77,6 +77,7 @@ def hparams(self, defaults, unused_model_hparams):
         "inputs": ("video", 256),
         "input_frame_number": ("symbol:identity", 1)
     }
+    p.target_modality = ("video", 256)
 
   def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
     frame_number = 0
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 98c19eecd..8df2222a3 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -819,8 +819,8 @@ def parametrized_bottleneck(x, hparams):
         hparams.discretize_warmup_steps, hparams.mode,
         hparams.isemhash_noise_dev, hparams.isemhash_mix_prob)
   if hparams.bottleneck_kind == "vq":
-    return vq_discrete_bottleneck(x, hparams.bottleneck_size, hparams.beta,
-                                  hparams.decay, hparams.epsilon)
+    return vq_discrete_bottleneck(x, hparams.bottleneck_size, hparams.vq_beta,
+                                  hparams.vq_decay, hparams.vq_epsilon)
   raise ValueError("Unsupported hparams.bottleneck_kind %s"
                    % hparams.bottleneck_kind)
 
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 883193316..5b1357618 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -42,11 +42,11 @@
 from tensor2tensor.models.research import attention_lm
 from tensor2tensor.models.research import attention_lm_moe
 from tensor2tensor.models.research import autoencoders
-from tensor2tensor.models.research import basic_conv_gen
 from tensor2tensor.models.research import cycle_gan
 from tensor2tensor.models.research import gene_expression
 from tensor2tensor.models.research import lm_experiments
 from tensor2tensor.models.research import multimodel
+from tensor2tensor.models.research import next_frame
 from tensor2tensor.models.research import r_transformer
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research import super_lm
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 20b455485..d1ff2c327 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -480,6 +480,9 @@ def autoencoder_residual_discrete():
   hparams.add_hparam("isemhash_noise_dev", 0.5)
   hparams.add_hparam("isemhash_mix_prob", 0.5)
   hparams.add_hparam("isemhash_filter_size_multiplier", 2.0)
+  hparams.add_hparam("vq_beta", 0.25)
+  hparams.add_hparam("vq_decay", 0.999)
+  hparams.add_hparam("vq_epsilon", 1e-5)
   return hparams
 
 
diff --git a/tensor2tensor/models/research/basic_conv_gen.py b/tensor2tensor/models/research/next_frame.py
similarity index 88%
rename from tensor2tensor/models/research/basic_conv_gen.py
rename to tensor2tensor/models/research/next_frame.py
index 6150ca86d..dcca82856 100644
--- a/tensor2tensor/models/research/basic_conv_gen.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -32,8 +32,8 @@
 
 
 @registry.register_model
-class BasicConvGen(t2t_model.T2TModel):
-  """Basic convolutional next-frame model."""
+class NextFrameBasic(t2t_model.T2TModel):
+  """Basic next-frame model, may take actions and predict rewards too."""
 
   def make_even_size(self, x):
     """Pad x to be even-sized on axis 1 and 2, but only if necessary."""
@@ -75,13 +75,14 @@ def body(self, features):
                              strides=(2, 2), padding="SAME")
         x = common_layers.layer_norm(x)
 
-    # Add embedded action.
-    action = tf.reshape(features["input_action"][:, -1, :],
-                        [-1, 1, 1, hparams.hidden_size])
-    action_mask = tf.layers.dense(action, filters, name="action_mask")
-    zeros_mask = tf.zeros(common_layers.shape_list(x)[:-1] + [filters],
-                          dtype=tf.float32)
-    x *= action_mask + zeros_mask
+    # Add embedded action if present.
+    if "input_action" in features:
+      action = tf.reshape(features["input_action"][:, -1, :],
+                          [-1, 1, 1, hparams.hidden_size])
+      action_mask = tf.layers.dense(action, filters, name="action_mask")
+      zeros_mask = tf.zeros(common_layers.shape_list(x)[:-1] + [filters],
+                            dtype=tf.float32)
+      x *= action_mask + zeros_mask
 
     # Run a stack of convolutions.
     for i in range(hparams.num_hidden_layers):
@@ -112,7 +113,9 @@ def body(self, features):
     # Cut down to original size.
     x = x[:, :inputs_shape[1], :inputs_shape[2], :]
 
-    # Reward prediction.
+    # Reward prediction if needed.
+    if "target_reward" not in features:
+      return x
     reward_pred = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
     return {"targets": x, "target_reward": reward_pred}
 
@@ -163,7 +166,7 @@ def logits_to_samples(logits):
 
 
 @registry.register_hparams
-def basic_conv():
+def next_frame():
   """Basic 2-frame conv model."""
   hparams = common_hparams.basic_params1()
   hparams.hidden_size = 64
@@ -186,15 +189,15 @@ def basic_conv():
 
 
 @registry.register_hparams
-def basic_conv_tpu():
-  hparams = basic_conv()
+def next_frame_tpu():
+  hparams = next_frame()
   hparams.batch_size = 1
 
 
 @registry.register_hparams
-def basic_conv_ae():
+def next_frame_ae():
   """Conv autoencoder."""
-  hparams = basic_conv()
+  hparams = next_frame()
   hparams.input_modalities = "inputs:video:bitwise"
   hparams.hidden_size = 256
   hparams.batch_size = 16
@@ -205,33 +208,33 @@ def basic_conv_ae():
 
 
 @registry.register_hparams
-def basic_conv_small():
+def next_frame_small():
   """Small conv model."""
-  hparams = basic_conv()
+  hparams = next_frame()
   hparams.hidden_size = 32
   return hparams
 
 
 @registry.register_hparams
-def basic_conv_l1():
+def next_frame_l1():
   """Basic conv model with L1 modality."""
-  hparams = basic_conv()
+  hparams = next_frame()
   hparams.target_modality = "video:l1"
   hparams.video_modality_loss_cutoff = 2.4
   return hparams
 
 
 @registry.register_hparams
-def basic_conv_l2():
+def next_frame_l2():
   """Basic conv model with L2 modality."""
-  hparams = basic_conv()
+  hparams = next_frame()
   hparams.target_modality = "video:l2"
   hparams.video_modality_loss_cutoff = 2.4
   return hparams
 
 
 @registry.register_ranged_hparams
-def basic_conv_base_range(rhp):
+def next_frame_base_range(rhp):
   """Basic tuning grid."""
   rhp.set_float("dropout", 0.2, 0.6)
   rhp.set_discrete("hidden_size", [64, 128, 256])
@@ -245,27 +248,27 @@ def basic_conv_base_range(rhp):
 
 
 @registry.register_ranged_hparams
-def basic_conv_doubling_range(rhp):
+def next_frame_doubling_range(rhp):
   """Filter doubling and dropout tuning grid."""
   rhp.set_float("dropout", 0.2, 0.6)
   rhp.set_int("filter_double_steps", 2, 5)
 
 
 @registry.register_ranged_hparams
-def basic_conv_clipgrad_range(rhp):
+def next_frame_clipgrad_range(rhp):
   """Filter doubling and dropout tuning grid."""
   rhp.set_float("dropout", 0.3, 0.4)
   rhp.set_float("clip_grad_norm", 0.5, 10.0)
 
 
 @registry.register_ranged_hparams
-def basic_conv_xent_cutoff_range(rhp):
+def next_frame_xent_cutoff_range(rhp):
   """Cross-entropy tuning grid."""
   rhp.set_float("video_modality_loss_cutoff", 0.005, 0.05)
 
 
 @registry.register_ranged_hparams
-def basic_conv_ae_range(rhp):
+def next_frame_ae_range(rhp):
   """Autoencoder world model tuning grid."""
   rhp.set_float("dropout", 0.3, 0.5)
   rhp.set_int("num_compress_steps", 1, 3)
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 4babec17c..92e860412 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -463,8 +463,8 @@ def rl_modelrl_base():
       # 1/11 steps are used for evaluation data
       # 100k frames for training = 36666
       true_env_generator_num_steps=36666,
-      generative_model="basic_conv_gen",
-      generative_model_params="basic_conv",
+      generative_model="next_frame_basic",
+      generative_model_params="next_frame",
       ppo_params="ppo_pong_base",
       autoencoder_train_steps=0,
       model_train_steps=50000,
@@ -537,7 +537,7 @@ def rl_modelrl_tiny():
 def rl_modelrl_l1_base():
   """Parameter set with L1 loss."""
   hparams = rl_modelrl_base()
-  hparams.generative_model_params = "basic_conv_l1"
+  hparams.generative_model_params = "next_frame_l1"
   return hparams
 
 
@@ -545,7 +545,7 @@ def rl_modelrl_l1_base():
 def rl_modelrl_l1_medium():
   """Medium parameter set with L1 loss."""
   hparams = rl_modelrl_medium()
-  hparams.generative_model_params = "basic_conv_l1"
+  hparams.generative_model_params = "next_frame_l1"
   return hparams
 
 
@@ -553,7 +553,7 @@ def rl_modelrl_l1_medium():
 def rl_modelrl_l1_short():
   """Short parameter set with L1 loss."""
   hparams = rl_modelrl_short()
-  hparams.generative_model_params = "basic_conv_l1"
+  hparams.generative_model_params = "next_frame_l1"
   return hparams
 
 
@@ -561,7 +561,7 @@ def rl_modelrl_l1_short():
 def rl_modelrl_l1_tiny():
   """Tiny parameter set with L1 loss."""
   hparams = rl_modelrl_tiny()
-  hparams.generative_model_params = "basic_conv_l1"
+  hparams.generative_model_params = "next_frame_l1"
   return hparams
 
 
@@ -569,7 +569,7 @@ def rl_modelrl_l1_tiny():
 def rl_modelrl_l2_base():
   """Parameter set with L2 loss."""
   hparams = rl_modelrl_base()
-  hparams.generative_model_params = "basic_conv_l2"
+  hparams.generative_model_params = "next_frame_l2"
   return hparams
 
 
@@ -577,7 +577,7 @@ def rl_modelrl_l2_base():
 def rl_modelrl_l2_medium():
   """Medium parameter set with L2 loss."""
   hparams = rl_modelrl_medium()
-  hparams.generative_model_params = "basic_conv_l2"
+  hparams.generative_model_params = "next_frame_l2"
   return hparams
 
 
@@ -585,7 +585,7 @@ def rl_modelrl_l2_medium():
 def rl_modelrl_l2_short():
   """Short parameter set with L2 loss."""
   hparams = rl_modelrl_short()
-  hparams.generative_model_params = "basic_conv_l2"
+  hparams.generative_model_params = "next_frame_l2"
   return hparams
 
 
@@ -593,7 +593,7 @@ def rl_modelrl_l2_short():
 def rl_modelrl_l2_tiny():
   """Tiny parameter set with L2 loss."""
   hparams = rl_modelrl_tiny()
-  hparams.generative_model_params = "basic_conv_l2"
+  hparams.generative_model_params = "next_frame_l2"
   return hparams
 
 
@@ -602,7 +602,7 @@ def rl_modelrl_ae_base():
   """Parameter set for autoencoders."""
   hparams = rl_modelrl_base()
   hparams.ppo_params = "ppo_pong_ae_base"
-  hparams.generative_model_params = "basic_conv_ae"
+  hparams.generative_model_params = "next_frame_ae"
   hparams.autoencoder_train_steps = 30000
   return hparams
 
@@ -618,7 +618,7 @@ def rl_modelrl_ae_25k():
 def rl_modelrl_ae_l1_base():
   """Parameter set for autoencoders and L1 loss."""
   hparams = rl_modelrl_ae_base()
-  hparams.generative_model_params = "basic_conv_l1"
+  hparams.generative_model_params = "next_frame_l1"
   return hparams
 
 
@@ -626,7 +626,7 @@ def rl_modelrl_ae_l1_base():
 def rl_modelrl_ae_l2_base():
   """Parameter set for autoencoders and L2 loss."""
   hparams = rl_modelrl_ae_base()
-  hparams.generative_model_params = "basic_conv_l2"
+  hparams.generative_model_params = "next_frame_l2"
   return hparams
 
 
@@ -654,7 +654,7 @@ def rl_modelrl_ae_tiny():
   """Tiny set for testing autoencoders."""
   hparams = rl_modelrl_tiny()
   hparams.ppo_params = "ppo_pong_ae_base"
-  hparams.generative_model_params = "basic_conv_ae"
+  hparams.generative_model_params = "next_frame_ae"
   hparams.autoencoder_train_steps = 2
   return hparams
 
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index f3537946c..f79cc74c2 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -99,7 +99,8 @@ def train(hparams, environment_spec, event_dir=None, model_dir=None,
       model_saver = None
 
     if hparams.simulated_environment:
-      env_model_loader = tf.train.Saver(tf.global_variables("basic_conv_gen.*"))
+      env_model_loader = tf.train.Saver(
+          tf.global_variables("next_frame_basic.*"))
     else:
       env_model_loader = None
 

From 6a7ef7f79f56fdcb1b16ae76d7e61cb09033dc4f Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 23 May 2018 16:55:52 -0700
Subject: [PATCH 0009/2720] Add pylintrc and lint fixes

PiperOrigin-RevId: 197813640
---
 .travis.yml                                   |  20 +-
 pylintrc                                      | 221 ++++++++++++++++++
 setup.py                                      |   2 +-
 tensor2tensor/bin/make_tf_configs.py          |   3 -
 tensor2tensor/bin/t2t_avg_all.py              |   3 -
 tensor2tensor/bin/t2t_bleu.py                 |   3 -
 tensor2tensor/bin/t2t_datagen.py              |   3 -
 tensor2tensor/bin/t2t_decoder.py              |   4 +-
 tensor2tensor/bin/t2t_distill.py              |   6 +-
 tensor2tensor/bin/t2t_trainer.py              |   6 +-
 tensor2tensor/bin/t2t_trainer_test.py         |   3 -
 tensor2tensor/bin/t2t_translate_all.py        |   3 -
 tensor2tensor/data_generators/algorithmic.py  |   7 +-
 .../data_generators/algorithmic_math.py       |   3 -
 .../data_generators/algorithmic_math_test.py  |   3 -
 .../data_generators/algorithmic_test.py       |   3 -
 .../data_generators/all_problems_test.py      |   1 -
 tensor2tensor/data_generators/audio.py        |  21 +-
 tensor2tensor/data_generators/audio_test.py   |   3 -
 tensor2tensor/data_generators/babi_qa.py      |   1 -
 tensor2tensor/data_generators/celeba.py       |   3 -
 tensor2tensor/data_generators/cifar.py        |   3 -
 tensor2tensor/data_generators/cipher.py       |   3 -
 .../data_generators/cnn_dailymail.py          |   3 -
 tensor2tensor/data_generators/desc2code.py    |  11 +-
 .../data_generators/desc2code_test.py         |   1 -
 tensor2tensor/data_generators/dna_encoder.py  |   8 +-
 .../data_generators/dna_encoder_test.py       |   3 -
 tensor2tensor/data_generators/fsns.py         |   3 -
 .../data_generators/gene_expression.py        |   4 +-
 .../data_generators/gene_expression_test.py   |   7 +-
 .../data_generators/generator_utils.py        |   6 +-
 .../data_generators/generator_utils_test.py   |   3 -
 tensor2tensor/data_generators/gym_problems.py |   4 +-
 tensor2tensor/data_generators/gym_utils.py    |  21 +-
 tensor2tensor/data_generators/ice_parsing.py  |   3 -
 tensor2tensor/data_generators/image_utils.py  |   6 +-
 .../data_generators/image_utils_test.py       |   3 -
 tensor2tensor/data_generators/imagenet.py     |   1 -
 tensor2tensor/data_generators/imdb.py         |   3 -
 .../data_generators/inspect_tfrecord.py       |   3 -
 tensor2tensor/data_generators/lambada.py      |   3 -
 tensor2tensor/data_generators/librispeech.py  |  10 +-
 tensor2tensor/data_generators/lm1b.py         |   3 -
 tensor2tensor/data_generators/mnist.py        |   3 -
 tensor2tensor/data_generators/mscoco.py       |   5 +-
 tensor2tensor/data_generators/multinli.py     |   4 +-
 tensor2tensor/data_generators/ocr.py          |   3 -
 tensor2tensor/data_generators/problem.py      |   4 +-
 .../data_generators/problem_hparams.py        |   3 -
 tensor2tensor/data_generators/ptb.py          |   3 -
 tensor2tensor/data_generators/snli.py         |   4 +-
 .../data_generators/speech_recognition.py     |   8 +-
 tensor2tensor/data_generators/squad.py        |   3 -
 .../data_generators/subject_verb_agreement.py |   3 -
 tensor2tensor/data_generators/text_encoder.py |  33 +--
 .../text_encoder_build_subword.py             |   3 -
 .../data_generators/text_encoder_test.py      |   1 -
 .../data_generators/text_problems_test.py     |   3 -
 .../data_generators/timeseries_test.py        |   1 -
 tensor2tensor/data_generators/tokenizer.py    |   3 -
 .../data_generators/tokenizer_test.py         |   3 -
 tensor2tensor/data_generators/translate.py    |   3 -
 .../data_generators/translate_encs.py         |   3 -
 .../data_generators/translate_ende.py         |   3 -
 .../data_generators/translate_enet.py         |   3 -
 .../data_generators/translate_enfr.py         |   3 -
 .../data_generators/translate_enmk.py         |   3 -
 .../data_generators/translate_envi.py         |   3 -
 .../data_generators/translate_enzh.py         |   3 -
 .../data_generators/translate_test.py         |   3 -
 tensor2tensor/data_generators/twentybn.py     |   3 -
 .../data_generators/video_generated.py        |   2 -
 tensor2tensor/data_generators/video_utils.py  |   7 +-
 tensor2tensor/data_generators/wiki.py         |   3 -
 .../data_generators/wikisum/utils_test.py     |   3 -
 tensor2tensor/data_generators/wikitext103.py  |   3 -
 tensor2tensor/data_generators/wsj_parsing.py  |  31 +--
 tensor2tensor/insights/query_processor.py     |   3 -
 tensor2tensor/layers/common_attention.py      |  41 ++--
 tensor2tensor/layers/common_attention_test.py |   3 -
 tensor2tensor/layers/common_hparams.py        |   3 -
 .../layers/common_image_attention.py          |   1 -
 tensor2tensor/layers/common_layers.py         |  13 +-
 tensor2tensor/layers/common_layers_test.py    |   3 -
 tensor2tensor/layers/discretization.py        |   6 +-
 tensor2tensor/layers/discretization_test.py   |   3 +-
 tensor2tensor/layers/modalities.py            |  63 ++---
 tensor2tensor/layers/modalities_test.py       |   3 -
 tensor2tensor/layers/rev_block.py             |   9 +-
 tensor2tensor/layers/rev_block_test.py        |  37 +--
 tensor2tensor/models/__init__.py              |   3 -
 tensor2tensor/models/basic.py                 |  10 +-
 tensor2tensor/models/basic_test.py            |   3 -
 tensor2tensor/models/bytenet.py               |   3 -
 tensor2tensor/models/bytenet_test.py          |   3 -
 tensor2tensor/models/distillation.py          |   3 -
 tensor2tensor/models/image_transformer.py     |   7 +-
 tensor2tensor/models/image_transformer_2d.py  |   4 +-
 .../models/image_transformer_2d_test.py       |  11 +-
 .../models/image_transformer_test.py          |   7 +-
 tensor2tensor/models/lstm.py                  |   3 -
 tensor2tensor/models/lstm_test.py             |   3 -
 tensor2tensor/models/neural_gpu.py            |   3 -
 tensor2tensor/models/neural_gpu_test.py       |   3 -
 tensor2tensor/models/research/aligned.py      |   3 -
 tensor2tensor/models/research/attention_lm.py |   3 -
 .../models/research/attention_lm_moe.py       |   3 -
 tensor2tensor/models/research/autoencoders.py |  13 +-
 .../models/research/autoencoders_test.py      |  33 ++-
 tensor2tensor/models/research/cycle_gan.py    |   4 +-
 .../models/research/gene_expression.py        |   4 +-
 .../models/research/gene_expression_test.py   |   7 +-
 tensor2tensor/models/research/multimodel.py   |   4 +-
 .../models/research/multimodel_test.py        |   3 -
 tensor2tensor/models/research/next_frame.py   |   6 +-
 .../models/research/r_transformer.py          |  19 +-
 .../models/research/r_transformer_test.py     |   8 +-
 .../models/research/r_transformer_util.py     |   3 -
 tensor2tensor/models/research/rl.py           |   3 -
 tensor2tensor/models/research/super_lm.py     |   3 -
 .../models/research/transformer_moe.py        |   7 +-
 .../models/research/transformer_revnet.py     |   3 -
 .../research/transformer_revnet_test.py       |   3 -
 .../models/research/transformer_sketch.py     |  10 +-
 .../models/research/transformer_symshard.py   |   3 -
 .../models/research/transformer_vae.py        |  15 +-
 tensor2tensor/models/resnet.py                |   3 -
 tensor2tensor/models/resnet_test.py           |   7 +-
 tensor2tensor/models/revnet.py                |   3 -
 tensor2tensor/models/shake_shake.py           |   3 -
 tensor2tensor/models/slicenet.py              |   3 -
 tensor2tensor/models/slicenet_test.py         |   3 -
 tensor2tensor/models/transformer.py           |   8 +-
 tensor2tensor/models/transformer_test.py      |   3 -
 tensor2tensor/models/vanilla_gan.py           |  13 +-
 tensor2tensor/models/xception.py              |   6 +-
 tensor2tensor/models/xception_test.py         |   9 +-
 tensor2tensor/problems.py                     |   8 +-
 tensor2tensor/problems_test.py                |  59 -----
 tensor2tensor/rl/envs/py_func_batch_env.py    |   3 -
 tensor2tensor/rl/envs/simulated_batch_env.py  |   3 -
 tensor2tensor/rl/envs/tf_atari_wrappers.py    |   3 -
 tensor2tensor/rl/envs/utils.py                |  23 +-
 tensor2tensor/rl/model_rl_experiment.py       |   3 -
 tensor2tensor/rl/rl_trainer_lib.py            |   3 -
 tensor2tensor/rl/rl_trainer_lib_test.py       |   3 -
 tensor2tensor/rl/t2t_rl_trainer.py            |   3 -
 tensor2tensor/serving/export.py               |   3 -
 tensor2tensor/serving/serving_utils.py        |   2 -
 tensor2tensor/utils/adafactor.py              |   4 +-
 tensor2tensor/utils/avg_checkpoints.py        |   3 -
 tensor2tensor/utils/beam_search.py            |   3 -
 tensor2tensor/utils/beam_search_test.py       |   3 -
 tensor2tensor/utils/bleu_hook.py              |   5 +-
 tensor2tensor/utils/bleu_hook_test.py         |   3 -
 .../utils/checkpoint_compatibility_test.py    |   3 -
 tensor2tensor/utils/cloud_mlengine.py         |  19 +-
 tensor2tensor/utils/data_reader.py            |   3 -
 tensor2tensor/utils/data_reader_test.py       |   4 +-
 tensor2tensor/utils/decoding.py               |   3 -
 tensor2tensor/utils/devices.py                |   3 -
 tensor2tensor/utils/diet.py                   |   4 +-
 tensor2tensor/utils/diet_test.py              |   3 -
 tensor2tensor/utils/expert_utils.py           |   9 +-
 tensor2tensor/utils/expert_utils_test.py      |   2 -
 tensor2tensor/utils/flags.py                  |   1 -
 tensor2tensor/utils/get_rouge.py              |   7 +-
 tensor2tensor/utils/learning_rate.py          |   3 -
 tensor2tensor/utils/metrics.py                |   3 -
 tensor2tensor/utils/metrics_hook.py           |   9 +-
 tensor2tensor/utils/metrics_hook_test.py      |   3 -
 tensor2tensor/utils/metrics_test.py           |   3 -
 tensor2tensor/utils/modality.py               |   3 -
 tensor2tensor/utils/optimize.py               |   7 +-
 tensor2tensor/utils/quantization.py           |   3 -
 tensor2tensor/utils/registry.py               |   3 -
 tensor2tensor/utils/registry_test.py          |   3 -
 tensor2tensor/utils/rouge_test.py             |   3 -
 tensor2tensor/utils/t2t_model.py              |   9 +-
 tensor2tensor/utils/trainer_lib.py            |   3 -
 tensor2tensor/utils/trainer_lib_test.py       |   6 +-
 tensor2tensor/utils/usr_dir.py                |   3 -
 tensor2tensor/utils/yellowfin.py              |   3 -
 tensor2tensor/utils/yellowfin_test.py         |  16 +-
 tensor2tensor/visualization/visualization.py  |   3 -
 186 files changed, 545 insertions(+), 820 deletions(-)
 create mode 100644 pylintrc
 delete mode 100644 tensor2tensor/problems_test.py

diff --git a/.travis.yml b/.travis.yml
index e83c02d0c..08901475c 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -7,13 +7,17 @@ env:
     - T2T_PROBLEM=algorithmic_reverse_binary40_test
     - T2T_DATA_DIR=/tmp/t2t-data
     - T2T_TRAIN_DIR=/tmp/t2t-train
+    - TF_LATEST="1.8.*"
   matrix:
+    # We test against the last 4 versions of TensorFlow
+    # If updating, also update TF_LATEST above
     - TF_VERSION="1.5.*"
     - TF_VERSION="1.6.*"
     - TF_VERSION="1.7.*"
     - TF_VERSION="1.8.*"
 matrix:
   exclude:
+    # We test against all versions in Python 2 but only the latest in Python 3
     - python: "3.6"
       env: TF_VERSION="1.5.*"
     - python: "3.6"
@@ -38,14 +42,18 @@ script:
   # Check import
   - python -c "from tensor2tensor.models import transformer; print(transformer.Transformer.__name__)"
 
+  # Lint
+  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST*"  ]]; then
+        pylint tensor2tensor;
+    fi
+
   # Run tests
   - pytest
-    --ignore=tensor2tensor/utils/registry_test.py
-    --ignore=tensor2tensor/utils/trainer_lib_test.py
-    --ignore=tensor2tensor/visualization/visualization_test.py
-    --ignore=tensor2tensor/problems_test.py
-    --ignore=tensor2tensor/bin/t2t_trainer_test.py
-    --ignore=tensor2tensor/data_generators/algorithmic_math_test.py
+    --ignore=tensor2tensor/utils/registry_test.py  # tested separately
+    --ignore=tensor2tensor/utils/trainer_lib_test.py  # tested separately
+    --ignore=tensor2tensor/visualization/visualization_test.py  # tested separately
+    --ignore=tensor2tensor/bin/t2t_trainer_test.py  # tested separately
+    --ignore=tensor2tensor/data_generators/algorithmic_math_test.py  # flaky
     --ignore=tensor2tensor/models/research/r_transformer_test.py  # Requires new feature in tf.foldl (rm with TF 1.9)
   - pytest tensor2tensor/utils/registry_test.py
   - pytest tensor2tensor/utils/trainer_lib_test.py
diff --git a/pylintrc b/pylintrc
new file mode 100644
index 000000000..0d4e111ef
--- /dev/null
+++ b/pylintrc
@@ -0,0 +1,221 @@
+
+
+[MASTER]
+
+# Pickle collected data for later comparisons.
+persistent=no
+
+# Set the cache size for astng objects.
+cache-size=500
+
+# Ignore Py3 files
+ignore=get_references_web.py,get_references_web_single_group.py
+
+
+[REPORTS]
+
+# Set the output format.
+# output-format=sorted-text
+
+# Put messages in a separate file for each module / package specified on the
+# command line instead of printing them on stdout. Reports (if any) will be
+# written in a file name "pylint_global.[txt|html]".
+files-output=no
+
+# Tells whether to display a full report or only the messages.
+reports=no
+
+# Disable the report(s) with the given id(s).
+disable-report=R0001,R0002,R0003,R0004,R0101,R0102,R0201,R0202,R0220,R0401,R0402,R0701,R0801,R0901,R0902,R0903,R0904,R0911,R0912,R0913,R0914,R0915,R0921,R0922,R0923
+
+# Error message template (continued on second line)
+msg-template={msg_id}:{line:3} {obj}: {msg} [{symbol}]
+
+
+[MESSAGES CONTROL]
+# List of checkers and warnings to enable.
+enable=indexing-exception,old-raise-syntax
+
+# List of checkers and warnings to disable.
+disable=design,similarities,no-self-use,attribute-defined-outside-init,locally-disabled,star-args,pointless-except,bad-option-value,global-statement,fixme,suppressed-message,useless-suppression,locally-enabled,file-ignored,multiple-imports,c-extension-no-member,trailing-newlines,unsubscriptable-object,misplaced-comparison-constant,no-member,abstract-method,no-else-return,missing-docstring,wrong-import-order,protected-access,inconsistent-return-statements,invalid-unary-operand-type,import-error,no-name-in-module
+
+[BASIC]
+
+# Required attributes for module, separated by a comma
+required-attributes=
+
+# Regular expression which should only match the name
+# of functions or classes which do not require a docstring.
+no-docstring-rgx=(__.*__|main)
+
+# Min length in lines of a function that requires a docstring.
+docstring-min-length=10
+
+# Regular expression which should only match correct module names. The
+# leading underscore is sanctioned for private modules by Google's style
+# guide.
+#
+# There are exceptions to the basic rule (_?[a-z][a-z0-9_]*) to cover
+# requirements of Python's module system.
+module-rgx=^(_?[a-z][a-z0-9_]*)|__init__$
+
+# Regular expression which should only match correct module level names
+const-rgx=^(_?[A-Z][A-Z0-9_]*|__[a-z0-9_]+__|_?[a-z][a-z0-9_]*)$
+
+# Regular expression which should only match correct class attribute
+class-attribute-rgx=^(_?[A-Z][A-Z0-9_]*|__[a-z0-9_]+__|_?[a-z][a-z0-9_]*)$
+
+# Regular expression which should only match correct class names
+class-rgx=^_?[A-Z][a-zA-Z0-9]*$
+
+# Regular expression which should only match correct function names.
+# 'camel_case' and 'snake_case' group names are used for consistency of naming
+# styles across functions and methods.
+function-rgx=^(?:(?P<exempt>setUp|tearDown|setUpModule|tearDownModule)|(?P<camel_case>_?[A-Z][a-zA-Z0-9]*)|(?P<snake_case>_?[a-z][a-z0-9_]*))$
+
+
+# Regular expression which should only match correct method names.
+# 'camel_case' and 'snake_case' group names are used for consistency of naming
+# styles across functions and methods. 'exempt' indicates a name which is
+# consistent with all naming styles.
+method-rgx=(?x)
+  ^(?:(?P<exempt>_[a-z0-9_]+__|runTest|setUp|tearDown|setUpTestCase
+         |tearDownTestCase|setupSelf|tearDownClass|setUpClass
+         |(test|assert)_*[A-Z0-9][a-zA-Z0-9_]*|next)
+     |(?P<camel_case>_{0,2}[A-Z][a-zA-Z0-9_]*)
+     |(?P<snake_case>_{0,2}[a-z][a-z0-9_]*))$
+
+
+# Regular expression which should only match correct instance attribute names
+attr-rgx=^_{0,2}[a-z][a-z0-9_]*$
+
+# Regular expression which should only match correct argument names
+argument-rgx=^[a-z][a-z0-9_]*$
+
+# Regular expression which should only match correct variable names
+variable-rgx=^[a-z][a-z0-9_]*$
+
+# Regular expression which should only match correct list comprehension /
+# generator expression variable names
+inlinevar-rgx=^[a-z][a-z0-9_]*$
+
+# Good variable names which should always be accepted, separated by a comma
+good-names=main,_
+
+# Bad variable names which should always be refused, separated by a comma
+bad-names=
+
+# List of builtins function names that should not be used, separated by a comma
+bad-functions=input,apply,reduce
+
+# List of decorators that define properties, such as abc.abstractproperty.
+property-classes=abc.abstractproperty
+
+
+[TYPECHECK]
+
+# Tells whether missing members accessed in mixin class should be ignored. A
+# mixin class is detected if its name ends with "mixin" (case insensitive).
+ignore-mixin-members=yes
+
+# List of decorators that create context managers from functions, such as
+# contextlib.contextmanager.
+contextmanager-decorators=contextlib.contextmanager,contextlib2.contextmanager
+
+
+[VARIABLES]
+
+# Tells whether we should check for unused import in __init__ files.
+init-import=no
+
+# A regular expression matching names used for dummy variables (i.e. not used).
+dummy-variables-rgx=^\*{0,2}(_$|unused_|dummy_)
+
+# List of additional names supposed to be defined in builtins. Remember that
+# you should avoid to define new builtins when possible.
+additional-builtins=
+
+
+[CLASSES]
+
+# List of method names used to declare (i.e. assign) instance attributes.
+defining-attr-methods=__init__,__new__,setUp
+
+# "class_" is also a valid for the first argument to a class method.
+valid-classmethod-first-arg=cls,class_
+
+
+[EXCEPTIONS]
+
+overgeneral-exceptions=StandardError,Exception,BaseException
+
+
+[IMPORTS]
+
+# Deprecated modules which should not be used, separated by a comma
+deprecated-modules=regsub,TERMIOS,Bastion,rexec,sets
+
+
+[FORMAT]
+
+# Maximum number of characters on a single line.
+max-line-length=80
+
+# Regexp for a line that is allowed to be longer than the limit.
+# This "ignore" regex is today composed of several independent parts:
+# (1) Long import lines
+# (2) URLs in comments or pydocs. Detecting URLs by regex is a hard problem and
+#     no amount of tweaking will make a perfect regex AFAICT. This one is a good
+#     compromise.
+# (3) Constant string literals at the start of files don't need to be broken
+#     across lines. Allowing long paths and urls to be on a single
+#     line. Also requires that the string not be a triplequoted string.
+ignore-long-lines=(?x)
+  (^\s*(import|from)\s
+   |^\s*(\#\ )?<?(https?|ftp):\/\/[^\s\/$.?#].[^\s]*>?$
+   |^[a-zA-Z_][a-zA-Z0-9_]*\s*=\s*("[^"]\S+"|'[^']\S+')
+   )
+
+# Maximum number of lines in a module
+max-module-lines=99999
+
+# String used as indentation unit. We differ from PEP8's normal 4 spaces.
+indent-string='  '
+
+# Do not warn about multiple statements on a single line for constructs like
+#   if test: stmt
+single-line-if-stmt=y
+
+# Make sure : in dicts and trailing commas are checked for whitespace.
+no-space-check=
+
+
+[LOGGING]
+
+# Add logging modules.
+logging-modules=logging,absl.logging
+
+
+[MISCELLANEOUS]
+
+# List of note tags to take in consideration, separated by a comma.
+notes=
+
+
+# Maximum line length for lambdas
+short-func-length=1
+
+# List of module members that should be marked as deprecated.
+# All of the string functions are listed in 4.1.4 Deprecated string functions
+# in the Python 2.4 docs.
+deprecated-members=string.atof,string.atoi,string.atol,string.capitalize,string.expandtabs,string.find,string.rfind,string.index,string.rindex,string.count,string.lower,string.split,string.rsplit,string.splitfields,string.join,string.joinfields,string.lstrip,string.rstrip,string.strip,string.swapcase,string.translate,string.upper,string.ljust,string.rjust,string.center,string.zfill,string.replace,sys.exitfunc,sys.maxint
+
+
+# List of exceptions that do not need to be mentioned in the Raises section of
+# a docstring.
+ignore-exceptions=AssertionError,NotImplementedError,StopIteration,TypeError
+
+
+# Number of spaces of indent required when the last token on the preceding line
+# is an open (, [, or {.
+indent-after-paren=4
diff --git a/setup.py b/setup.py
index 018444555..95b91a2af 100644
--- a/setup.py
+++ b/setup.py
@@ -49,7 +49,7 @@
     extras_require={
         'tensorflow': ['tensorflow>=1.5.0'],
         'tensorflow_gpu': ['tensorflow-gpu>=1.5.0'],
-        'tests': ['pytest', 'mock'],
+        'tests': ['pytest', 'mock', 'pylint'],
     },
     classifiers=[
         'Development Status :: 4 - Beta',
diff --git a/tensor2tensor/bin/make_tf_configs.py b/tensor2tensor/bin/make_tf_configs.py
index cf3d10257..d896e6b91 100644
--- a/tensor2tensor/bin/make_tf_configs.py
+++ b/tensor2tensor/bin/make_tf_configs.py
@@ -29,9 +29,6 @@
 from __future__ import print_function
 
 import json
-
-# Dependency imports
-
 import tensorflow as tf
 
 flags = tf.flags
diff --git a/tensor2tensor/bin/t2t_avg_all.py b/tensor2tensor/bin/t2t_avg_all.py
index 7d34f4a33..a3f96d3a4 100644
--- a/tensor2tensor/bin/t2t_avg_all.py
+++ b/tensor2tensor/bin/t2t_avg_all.py
@@ -20,9 +20,6 @@
 from collections import deque
 import os
 import shutil
-
-# Dependency imports
-
 import numpy as np
 import six
 from six.moves import zip  # pylint: disable=redefined-builtin
diff --git a/tensor2tensor/bin/t2t_bleu.py b/tensor2tensor/bin/t2t_bleu.py
index 9cf789e7c..338b2adaa 100644
--- a/tensor2tensor/bin/t2t_bleu.py
+++ b/tensor2tensor/bin/t2t_bleu.py
@@ -57,9 +57,6 @@
 
 import os
 import time
-
-# Dependency imports
-
 from tensor2tensor.utils import bleu_hook
 import tensorflow as tf
 
diff --git a/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index d6c7b320d..c02897834 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -33,9 +33,6 @@
 import random
 import tempfile
 
-# Dependency imports
-
-
 import numpy as np
 
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index cb3b19032..576941fe4 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -32,9 +32,6 @@
 from __future__ import print_function
 
 import os
-
-# Dependency imports
-
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.utils import decoding
@@ -80,6 +77,7 @@ def create_decode_hparams():
 
 
 def decode(estimator, hparams, decode_hp):
+  """Decode from estimator. Interactive, from file, or from dataset."""
   if FLAGS.decode_interactive:
     if estimator.config.use_tpu:
       raise ValueError("TPU can only decode from dataset.")
diff --git a/tensor2tensor/bin/t2t_distill.py b/tensor2tensor/bin/t2t_distill.py
index 2ed3d0cb6..ef9109373 100644
--- a/tensor2tensor/bin/t2t_distill.py
+++ b/tensor2tensor/bin/t2t_distill.py
@@ -22,9 +22,6 @@
 from __future__ import print_function
 
 import os
-
-# Dependency imports
-
 from tensor2tensor import models  # pylint: disable=unused-import
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
 from tensor2tensor.bin import t2t_trainer
@@ -46,7 +43,8 @@ def main(argv):
   t2t_trainer.log_registry()
 
   if FLAGS.cloud_mlengine:
-    return cloud_mlengine.launch()
+    cloud_mlengine.launch()
+    return
 
   if FLAGS.generate_data:
     t2t_trainer.generate_data()
diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 21eee7662..fc12cce15 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -20,9 +20,6 @@
 import contextlib
 import os
 import sys
-
-# Dependency imports
-
 from tensor2tensor import models  # pylint: disable=unused-import
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
 from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
@@ -335,7 +332,8 @@ def main(argv):
 
 
   if FLAGS.cloud_mlengine:
-    return cloud_mlengine.launch()
+    cloud_mlengine.launch()
+    return
 
   if FLAGS.generate_data:
     generate_data()
diff --git a/tensor2tensor/bin/t2t_trainer_test.py b/tensor2tensor/bin/t2t_trainer_test.py
index 31b1e884b..9acb56850 100644
--- a/tensor2tensor/bin/t2t_trainer_test.py
+++ b/tensor2tensor/bin/t2t_trainer_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.utils import trainer_lib_test
 
diff --git a/tensor2tensor/bin/t2t_translate_all.py b/tensor2tensor/bin/t2t_translate_all.py
index c6f8354e5..9194fa43d 100644
--- a/tensor2tensor/bin/t2t_translate_all.py
+++ b/tensor2tensor/bin/t2t_translate_all.py
@@ -29,9 +29,6 @@
 
 import os
 import shutil
-
-# Dependency imports
-
 from tensor2tensor.utils import bleu_hook
 
 import tensorflow as tf
diff --git a/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index d98a133c5..3efd2381f 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from six.moves import range  # pylint: disable=redefined-builtin
@@ -332,7 +329,7 @@ class AlgorithmicAdditionBinary40(AlgorithmicProblem):
   def num_symbols(self):
     return 2
 
-  def generator(self, base, max_length, nbr_cases):
+  def generator(self, base, max_length, nbr_cases):  # pylint: disable=arguments-differ
     """Generator for the addition task.
 
     The length of each number is drawn uniformly at random in [1, max_length/2]
@@ -382,7 +379,7 @@ class AlgorithmicMultiplicationBinary40(AlgorithmicProblem):
   def num_symbols(self):
     return 2
 
-  def generator(self, base, max_length, nbr_cases):
+  def generator(self, base, max_length, nbr_cases):  # pylint: disable=arguments-differ
     """Generator for the multiplication task.
 
     The length of each number is drawn uniformly at random in [1, max_length/2]
diff --git a/tensor2tensor/data_generators/algorithmic_math.py b/tensor2tensor/data_generators/algorithmic_math.py
index 3edc0db19..604a51b53 100644
--- a/tensor2tensor/data_generators/algorithmic_math.py
+++ b/tensor2tensor/data_generators/algorithmic_math.py
@@ -23,9 +23,6 @@
 
 from collections import namedtuple
 import random
-
-# Dependency imports
-
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
 import sympy
diff --git a/tensor2tensor/data_generators/algorithmic_math_test.py b/tensor2tensor/data_generators/algorithmic_math_test.py
index 953415947..20eed95d2 100644
--- a/tensor2tensor/data_generators/algorithmic_math_test.py
+++ b/tensor2tensor/data_generators/algorithmic_math_test.py
@@ -18,9 +18,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import six
 import sympy
 from tensor2tensor.data_generators import algorithmic_math
diff --git a/tensor2tensor/data_generators/algorithmic_test.py b/tensor2tensor/data_generators/algorithmic_test.py
index ffa2f4b38..234c5a7b7 100644
--- a/tensor2tensor/data_generators/algorithmic_test.py
+++ b/tensor2tensor/data_generators/algorithmic_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.data_generators import algorithmic
diff --git a/tensor2tensor/data_generators/all_problems_test.py b/tensor2tensor/data_generators/all_problems_test.py
index cc899fa6f..9d9fb3316 100644
--- a/tensor2tensor/data_generators/all_problems_test.py
+++ b/tensor2tensor/data_generators/all_problems_test.py
@@ -18,7 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-# Dependency imports
 from tensor2tensor.data_generators import all_problems
 
 import tensorflow as tf
diff --git a/tensor2tensor/data_generators/audio.py b/tensor2tensor/data_generators/audio.py
index 3335cd70f..96ed93d83 100644
--- a/tensor2tensor/data_generators/audio.py
+++ b/tensor2tensor/data_generators/audio.py
@@ -21,10 +21,7 @@
 from subprocess import call
 import tarfile
 import wave
-
-# Dependency imports
-
-from tensor2tensor.data_generators import generator_utils
+# from tensor2tensor.data_generators import generator_utils
 
 import tensorflow as tf
 
@@ -127,16 +124,22 @@ def timit_generator(data_dir,
     * audio/sample_width: an integer
     * targets: an integer sequence representing the encoded sentence
   """
+  del data_dir
   eos_list = [1] if eos_list is None else eos_list
   if vocab_filename is not None:
-    vocab_symbolizer = generator_utils.get_or_generate_vocab(
-        data_dir, tmp_dir, vocab_filename, vocab_size)
+    # TODO(lukaszkaiser): Correct this call to generate a vocabulary. No data
+    # sources are being passed.
+    # vocab_symbolizer = generator_utils.get_or_generate_vocab(
+    #     data_dir, tmp_dir, vocab_filename, vocab_size)
+    del vocab_size
+    vocab_symbolizer = None
+    assert False
   _get_timit(tmp_dir)
   datasets = (_TIMIT_TRAIN_DATASETS if training else _TIMIT_TEST_DATASETS)
   i = 0
-  for data_dir, (audio_ext, transcription_ext) in datasets:
-    data_dir = os.path.join(tmp_dir, data_dir)
-    data_files = _collect_data(data_dir, audio_ext, transcription_ext)
+  for timit_data_dir, (audio_ext, transcription_ext) in datasets:
+    timit_data_dir = os.path.join(tmp_dir, timit_data_dir)
+    data_files = _collect_data(timit_data_dir, audio_ext, transcription_ext)
     data_pairs = data_files.values()
     for input_file, target_file in sorted(data_pairs)[start_from:]:
       if i == how_many:
diff --git a/tensor2tensor/data_generators/audio_test.py b/tensor2tensor/data_generators/audio_test.py
index 37e188e3f..57a377cc0 100644
--- a/tensor2tensor/data_generators/audio_test.py
+++ b/tensor2tensor/data_generators/audio_test.py
@@ -20,9 +20,6 @@
 
 import io
 import os
-
-# Dependency imports
-
 from tensor2tensor.data_generators import audio
 
 import tensorflow as tf
diff --git a/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index 0c13313d5..395952a34 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -34,7 +34,6 @@
 import shutil
 import tarfile
 
-# Dependency imports
 import six
 
 from tensor2tensor.data_generators import generator_utils
diff --git a/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index 1ca5fc04c..ccb6cd7fc 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -20,9 +20,6 @@
 
 import os
 import zipfile
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.utils import registry
diff --git a/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index e2e27b787..74532af17 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -20,9 +20,6 @@
 
 import os
 import tarfile
-
-# Dependency imports
-
 import numpy as np
 
 from six.moves import cPickle
diff --git a/tensor2tensor/data_generators/cipher.py b/tensor2tensor/data_generators/cipher.py
index d6a244d59..b4e4878b2 100644
--- a/tensor2tensor/data_generators/cipher.py
+++ b/tensor2tensor/data_generators/cipher.py
@@ -18,9 +18,6 @@
 from __future__ import print_function
 
 from collections import deque
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import algorithmic
diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index 6cc0a48ec..ce147ca75 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -22,9 +22,6 @@
 import io
 import os
 import tarfile
-
-# Dependency imports
-
 import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
diff --git a/tensor2tensor/data_generators/desc2code.py b/tensor2tensor/data_generators/desc2code.py
index 99c15882e..64e0dc219 100644
--- a/tensor2tensor/data_generators/desc2code.py
+++ b/tensor2tensor/data_generators/desc2code.py
@@ -23,9 +23,6 @@
 import random
 import re
 import zipfile
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -159,6 +156,7 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
     ))
 
     def generator_samples_content(get_source, get_target):
+      """Generate samples."""
       source, target = None, None
       # Iterate over the coding samples
       for sample in samples:
@@ -182,8 +180,11 @@ def generator_target():
 
     # Generate vocab for both source and target
 
-    source_vocab = generator_utils.get_or_generate_vocab(
-        data_dir, tmp_dir, self.vocab_input_filename, self.input_vocab_size)
+    # TODO(lukaszkaiser): Fix vocab generation call. No sources given.
+    assert not self.vocab_input_filename
+    source_vocab = None
+    # source_vocab = generator_utils.get_or_generate_vocab(
+    #     data_dir, tmp_dir, self.vocab_input_filename, self.input_vocab_size)
 
     target_vocab = generator_utils.get_or_generate_vocab_inner(
         data_dir=data_dir,
diff --git a/tensor2tensor/data_generators/desc2code_test.py b/tensor2tensor/data_generators/desc2code_test.py
index cccfef801..1109daa3d 100644
--- a/tensor2tensor/data_generators/desc2code_test.py
+++ b/tensor2tensor/data_generators/desc2code_test.py
@@ -18,7 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-# Dependency imports
 from tensor2tensor.data_generators import desc2code
 
 import tensorflow as tf
diff --git a/tensor2tensor/data_generators/dna_encoder.py b/tensor2tensor/data_generators/dna_encoder.py
index ce1d09955..62c9b24e8 100644
--- a/tensor2tensor/data_generators/dna_encoder.py
+++ b/tensor2tensor/data_generators/dna_encoder.py
@@ -23,7 +23,6 @@
 from __future__ import print_function
 
 import itertools
-# Dependency imports
 
 from six.moves import range  # pylint: disable=redefined-builtin
 from tensor2tensor.data_generators import text_encoder
@@ -116,9 +115,10 @@ def delimiter(self):
   def _tokens(self):
     return super(DelimitedDNAEncoder, self)._tokens() + [self._delimiter_key]
 
-  def encode(self, delimited_string):
+  def encode(self, s):
+    delimited_string = s
     ids = []
-    for s in delimited_string.split(self.delimiter):
-      ids.extend(super(DelimitedDNAEncoder, self).encode(s))
+    for part in delimited_string.split(self.delimiter):
+      ids.extend(super(DelimitedDNAEncoder, self).encode(part))
       ids.append(self._tokens_to_ids[self._delimiter_key])
     return ids[:-1]
diff --git a/tensor2tensor/data_generators/dna_encoder_test.py b/tensor2tensor/data_generators/dna_encoder_test.py
index 453faf6a6..afe4b568b 100644
--- a/tensor2tensor/data_generators/dna_encoder_test.py
+++ b/tensor2tensor/data_generators/dna_encoder_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.data_generators import dna_encoder
 import tensorflow as tf
 
diff --git a/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
index 46fb8c021..6dc2567da 100644
--- a/tensor2tensor/data_generators/fsns.py
+++ b/tensor2tensor/data_generators/fsns.py
@@ -19,9 +19,6 @@
 from __future__ import print_function
 
 import os
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import problem
diff --git a/tensor2tensor/data_generators/gene_expression.py b/tensor2tensor/data_generators/gene_expression.py
index 2b640dc11..3f8acb3bc 100644
--- a/tensor2tensor/data_generators/gene_expression.py
+++ b/tensor2tensor/data_generators/gene_expression.py
@@ -37,9 +37,6 @@
 import math
 import multiprocessing as mp
 import os
-
-# Dependency imports
-
 import h5py
 import numpy as np
 
@@ -235,6 +232,7 @@ def dataset_generator(filepath,
                       chunk_size=1,
                       start_idx=None,
                       end_idx=None):
+  """Generate example dicts."""
   encoder = dna_encoder.DNAEncoder(chunk_size=chunk_size)
   with h5py.File(filepath, "r") as h5_file:
     # Get input keys from h5_file
diff --git a/tensor2tensor/data_generators/gene_expression_test.py b/tensor2tensor/data_generators/gene_expression_test.py
index b70b0885d..7df355fa5 100644
--- a/tensor2tensor/data_generators/gene_expression_test.py
+++ b/tensor2tensor/data_generators/gene_expression_test.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import dna_encoder
@@ -29,7 +26,7 @@
 
 class GeneticsTest(tf.test.TestCase):
 
-  def _oneHotBases(self, bases):
+  def _one_hot_bases(self, bases):
     ref = ["A", "C", "T", "G"]
     one_hots = []
     for base in bases:
@@ -44,7 +41,7 @@ def testRecordToExample(self):
     raw_inputs = ["A", "C", "G", "N", "C", "T"]
 
     # Put in numpy arrays in the same format as in the h5 file
-    inputs = self._oneHotBases(raw_inputs)
+    inputs = self._one_hot_bases(raw_inputs)
     mask = np.array([True, False, True])
     outputs = np.array([[1.0, 2.0, 3.0], [5.0, 1.0, 0.2], [5.1, 2.3, 2.3]])
     # Convert to example dict
diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 32d5b92da..2667f927d 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -25,13 +25,11 @@
 import stat
 import tarfile
 import tempfile
-
-# Dependency imports
-
 import requests
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
-import six.moves.urllib_request as urllib  # Imports urllib on Python2, urllib.request on Python3
+# Imports urllib on Python2, urllib.request on Python3
+import six.moves.urllib_request as urllib
 
 from tensor2tensor.data_generators import text_encoder
 
diff --git a/tensor2tensor/data_generators/generator_utils_test.py b/tensor2tensor/data_generators/generator_utils_test.py
index 6276e0d3a..b053e508c 100644
--- a/tensor2tensor/data_generators/generator_utils_test.py
+++ b/tensor2tensor/data_generators/generator_utils_test.py
@@ -22,9 +22,6 @@
 import io
 import os
 import tempfile
-
-# Dependency imports
-
 from builtins import bytes  # pylint: disable=redefined-builtin
 
 from tensor2tensor.data_generators import generator_utils
diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index ffdce69dd..a13020d9a 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -20,9 +20,6 @@
 
 import math
 import os
-
-# Dependency imports
-
 import gym
 import numpy as np
 
@@ -141,6 +138,7 @@ def num_testing_steps(self):
     return None
 
   def get_action(self, observation=None):
+    del observation
     return self.env.action_space.sample()
 
   def hparams(self, defaults, unused_model_hparams):
diff --git a/tensor2tensor/data_generators/gym_utils.py b/tensor2tensor/data_generators/gym_utils.py
index 0edaefe47..46abe9bfa 100644
--- a/tensor2tensor/data_generators/gym_utils.py
+++ b/tensor2tensor/data_generators/gym_utils.py
@@ -15,9 +15,6 @@
 """Utilities for openai gym."""
 
 from collections import deque
-
-# Dependency imports
-
 import gym
 
 import numpy as np
@@ -27,6 +24,7 @@
 from tensor2tensor.data_generators import image_utils
 
 
+# pylint: disable=method-hidden
 class WarmupWrapper(gym.Wrapper):
   """Warmup wrapper."""
 
@@ -48,11 +46,11 @@ def get_starting_data(self, num_frames):
 
     return starting_observations, starting_actions, starting_rewards
 
-  def step(self, ac):
-    action = ac
+  def step(self, action):
     return self.env.step(action)
 
   def reset(self, **kwargs):
+    del kwargs
     self.env.reset()
     observation = None
     for _ in range(self.warm_up_examples):
@@ -78,10 +76,9 @@ def __init__(self, env, warm_up_examples=0,
     self.reward_skip_steps = reward_skip_steps
     self.big_ball = big_ball
 
-  def step(self, ac):
-    action = ac
+  def step(self, action):
     if self.action_space_reduction:
-      action = 2 if int(ac) == 0 else 5
+      action = 2 if int(action) == 0 else 5
     ob, rew, done, info = self.env.step(action)
     ob = self.process_observation(ob)
     if rew != 0 and self.reward_skip_steps != 0:
@@ -166,8 +163,8 @@ def __init__(self, env, warm_up_examples=0,
            "include_direction_info to work correctly")
     assert not self.include_direction_info or ball_down_skip >= 9, msg
 
-  def step(self, ac):
-    ob, rew, done, info = self.env.step(ac)
+  def step(self, action):
+    ob, rew, done, info = self.env.step(action)
 
     if BreakoutWrapper.find_ball(ob) is None and self.ball_down_skip != 0:
       for _ in range(self.ball_down_skip):
@@ -266,8 +263,8 @@ def __init__(self, env,
   def chicken_height(self, image):
     raise NotImplementedError()
 
-  def step(self, ac):
-    ob, rew, done, info = self.env.step(ac)
+  def step(self, action):
+    ob, rew, done, info = self.env.step(action)
 
     if self.easy_freeway:
       if rew > 0:
diff --git a/tensor2tensor/data_generators/ice_parsing.py b/tensor2tensor/data_generators/ice_parsing.py
index e5c39934f..2561f28bc 100644
--- a/tensor2tensor/data_generators/ice_parsing.py
+++ b/tensor2tensor/data_generators/ice_parsing.py
@@ -25,9 +25,6 @@
 from __future__ import print_function
 
 import os
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 59e7a1277..a824913cd 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -19,9 +19,6 @@
 from __future__ import print_function
 
 import os
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -62,7 +59,7 @@ def num_channels(self):
     """Number of color channels."""
     return 3
 
-  def example_reading_spec(self, label_repr=None):
+  def example_reading_spec(self):
     data_fields = {
         "image/encoded": tf.FixedLenFeature((), tf.string),
         "image/format": tf.FixedLenFeature((), tf.string),
@@ -156,6 +153,7 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
 
 def encode_images_as_png(images):
+  """Yield images encoded as pngs."""
   if tf.contrib.eager.in_eager_mode():
     for image in images:
       yield tf.image.encode_png(image).numpy()
diff --git a/tensor2tensor/data_generators/image_utils_test.py b/tensor2tensor/data_generators/image_utils_test.py
index 6c0ce9367..c47992f12 100644
--- a/tensor2tensor/data_generators/image_utils_test.py
+++ b/tensor2tensor/data_generators/image_utils_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 from tensor2tensor.data_generators import image_utils
 
diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 06206ce6f..6b5037735 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -19,7 +19,6 @@
 from __future__ import print_function
 
 import os
-# Dependency imports
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import image_utils
diff --git a/tensor2tensor/data_generators/imdb.py b/tensor2tensor/data_generators/imdb.py
index 865def20c..c914df22e 100644
--- a/tensor2tensor/data_generators/imdb.py
+++ b/tensor2tensor/data_generators/imdb.py
@@ -20,9 +20,6 @@
 
 import os
 import tarfile
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_problems
diff --git a/tensor2tensor/data_generators/inspect_tfrecord.py b/tensor2tensor/data_generators/inspect_tfrecord.py
index 0113757e6..7b29015c9 100644
--- a/tensor2tensor/data_generators/inspect_tfrecord.py
+++ b/tensor2tensor/data_generators/inspect_tfrecord.py
@@ -24,9 +24,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import six
 
 from tensor2tensor.data_generators import text_encoder
diff --git a/tensor2tensor/data_generators/lambada.py b/tensor2tensor/data_generators/lambada.py
index c33d7c599..688db95ca 100644
--- a/tensor2tensor/data_generators/lambada.py
+++ b/tensor2tensor/data_generators/lambada.py
@@ -35,9 +35,6 @@
 import csv
 import os
 import tarfile
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
diff --git a/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index 28cb7b756..37adc4aa6 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -16,9 +16,6 @@
 
 import os
 import tarfile
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import speech_recognition
@@ -119,6 +116,7 @@ def use_train_shards_for_dev(self):
 
   def generator(self, data_dir, tmp_dir, datasets,
                 eos_list=None, start_from=0, how_many=0):
+    del eos_list
     i = 0
     for url, subdir in datasets:
       filename = os.path.basename(url)
@@ -185,13 +183,13 @@ class LibrispeechTrainFullTestClean(Librispeech):
   """Problem to train on full 960h, but evaluate on clean data only."""
 
   def training_filepaths(self, data_dir, num_shards, shuffled):
-    return Librispeech.training_filepaths(data_dir, num_shards, shuffled)
+    return Librispeech.training_filepaths(self, data_dir, num_shards, shuffled)
 
   def dev_filepaths(self, data_dir, num_shards, shuffled):
-    return LibrispeechClean.dev_filepaths(data_dir, num_shards, shuffled)
+    return LibrispeechClean.dev_filepaths(self, data_dir, num_shards, shuffled)
 
   def test_filepaths(self, data_dir, num_shards, shuffled):
-    return LibrispeechClean.test_filepaths(data_dir, num_shards, shuffled)
+    return LibrispeechClean.test_filepaths(self, data_dir, num_shards, shuffled)
 
   def generate_data(self, data_dir, tmp_dir, task_id=-1):
     raise Exception("Generate librispeech and librispeech_clean data.")
diff --git a/tensor2tensor/data_generators/lm1b.py b/tensor2tensor/data_generators/lm1b.py
index a81ff02bd..84a03fef7 100644
--- a/tensor2tensor/data_generators/lm1b.py
+++ b/tensor2tensor/data_generators/lm1b.py
@@ -20,9 +20,6 @@
 
 import os
 import tarfile
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.data_generators import generator_utils
diff --git a/tensor2tensor/data_generators/mnist.py b/tensor2tensor/data_generators/mnist.py
index c3d122499..5db518461 100644
--- a/tensor2tensor/data_generators/mnist.py
+++ b/tensor2tensor/data_generators/mnist.py
@@ -21,9 +21,6 @@
 import gzip
 import os
 import random
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import generator_utils
diff --git a/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index 76b745ade..acb8bc697 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -23,9 +23,6 @@
 import os
 import random
 import zipfile
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import imagenet
@@ -96,7 +93,7 @@ def get_vocab():
         vocab_symbolizer = text_encoder.SubwordTextEncoder(vocab_filepath)
         return vocab_symbolizer
       else:
-        raise ValueError("Vocab file does not exist: %s", vocab_filepath)
+        raise ValueError("Vocab file does not exist: %s" % vocab_filepath)
     return None
 
   vocab_symbolizer = get_vocab()
diff --git a/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index 4a6649a4f..6c50cf9c7 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -22,9 +22,6 @@
 import json
 import os
 import zipfile
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -77,6 +74,7 @@ def _dev_file(self):
       return 'multinli_1.0/multinli_1.0_dev_mismatched.jsonl'
 
   def _examples(self, data_dir, tmp_dir, train):
+    del data_dir
     file_path = generator_utils.maybe_download(tmp_dir, self._ZIP, self._URL)
     zip_ref = zipfile.ZipFile(file_path, 'r')
     zip_ref.extractall(tmp_dir)
diff --git a/tensor2tensor/data_generators/ocr.py b/tensor2tensor/data_generators/ocr.py
index fcc5e07e5..cff224bd2 100644
--- a/tensor2tensor/data_generators/ocr.py
+++ b/tensor2tensor/data_generators/ocr.py
@@ -19,9 +19,6 @@
 
 import os
 import struct
-
-# Dependency imports
-
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.utils import registry
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index bce818684..83542c8c9 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -19,7 +19,7 @@
 import collections
 import os
 import random
-# Dependency imports
+
 import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import text_encoder
@@ -787,7 +787,7 @@ def define_shapes(example):
         dataset = dataset.apply(
             tf.contrib.data.batch_and_drop_remainder(batch_size))
       else:
-        num_shards = (config and config.data_parallelism.n) or 1
+        num_shards = config.data_parallelism.n if config else 1
         batch_size = hparams.batch_size * num_shards
         dataset = dataset.batch(batch_size)
     else:
diff --git a/tensor2tensor/data_generators/problem_hparams.py b/tensor2tensor/data_generators/problem_hparams.py
index 7e84257c7..a6ffd7b0c 100644
--- a/tensor2tensor/data_generators/problem_hparams.py
+++ b/tensor2tensor/data_generators/problem_hparams.py
@@ -20,9 +20,6 @@
 from __future__ import print_function
 
 import os
-
-# Dependency imports
-
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.layers import modalities  # pylint: disable=unused-import
diff --git a/tensor2tensor/data_generators/ptb.py b/tensor2tensor/data_generators/ptb.py
index de3db0e28..c40c135cd 100644
--- a/tensor2tensor/data_generators/ptb.py
+++ b/tensor2tensor/data_generators/ptb.py
@@ -22,9 +22,6 @@
 import os
 import sys
 import tarfile
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
diff --git a/tensor2tensor/data_generators/snli.py b/tensor2tensor/data_generators/snli.py
index 7e2dd067c..3cc8c2b5d 100644
--- a/tensor2tensor/data_generators/snli.py
+++ b/tensor2tensor/data_generators/snli.py
@@ -20,9 +20,6 @@
 
 import os
 import zipfile
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import tokenizer
@@ -149,6 +146,7 @@ def _get_or_generate_vocab(tmp_dir, vocab_filename, vocab_size):
 
 
 def snli_token_generator(tmp_dir, train, vocab_size):
+  """Generate example dicts."""
   _download_and_parse_dataset(tmp_dir, train)
 
   symbolizer_vocab = _get_or_generate_vocab(
diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 12f31667c..9eed9fe50 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -22,9 +22,6 @@
 import os
 from subprocess import call
 import tempfile
-
-# Dependency imports
-
 import numpy as np
 from scipy.io import wavfile
 import scipy.signal
@@ -338,15 +335,16 @@ def eval_metrics(self):
 class SpeechRecognitionModality(modality.Modality):
   """Common ASR filterbank processing."""
 
-  def bottom(self, inputs):
+  def bottom(self, x):
     """Use batchnorm instead of CMVN and shorten the stft with strided convs.
 
     Args:
-      inputs: float32 tensor with shape [batch_size, len, 1, freqs * channels]
+      x: float32 tensor with shape [batch_size, len, 1, freqs * channels]
 
     Returns:
       float32 tensor with shape [batch_size, shorter_len, 1, hidden_size]
     """
+    inputs = x
     p = self._model_hparams
 
     num_mel_bins = p.audio_num_mel_bins
diff --git a/tensor2tensor/data_generators/squad.py b/tensor2tensor/data_generators/squad.py
index e19307242..7d82df547 100644
--- a/tensor2tensor/data_generators/squad.py
+++ b/tensor2tensor/data_generators/squad.py
@@ -21,9 +21,6 @@
 
 import json
 import os
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_problems
diff --git a/tensor2tensor/data_generators/subject_verb_agreement.py b/tensor2tensor/data_generators/subject_verb_agreement.py
index 5e126e763..3b6604215 100644
--- a/tensor2tensor/data_generators/subject_verb_agreement.py
+++ b/tensor2tensor/data_generators/subject_verb_agreement.py
@@ -31,9 +31,6 @@
 import gzip
 import os
 import random
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index c256f94b8..53fb47ea8 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -28,9 +28,6 @@
 import math
 import re
 import tempfile
-
-# Dependency imports
-
 import numpy as np
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
@@ -215,10 +212,12 @@ def __init__(self, class_labels=None, class_labels_fname=None):
 
     self._class_labels = class_labels
 
-  def encode(self, label_str):
+  def encode(self, s):
+    label_str = s
     return self._class_labels.index(label_str)
 
-  def decode(self, label_id):
+  def decode(self, ids):
+    label_id = ids
     if isinstance(label_id, list):
       assert len(label_id) == 1
       label_id, = label_id
@@ -248,14 +247,15 @@ def __init__(self, class_labels=None, class_labels_fname=None):
 
     self._class_labels = class_labels
 
-  def encode(self, label_str, on_value=1, off_value=0):
+  def encode(self, label_str, on_value=1, off_value=0):  # pylint: disable=arguments-differ
     e = np.zeros(self.vocab_size, dtype=np.int32)
     if off_value != 0:
       e.fill(off_value)
     e[self._class_labels.index(label_str)] = on_value
     return e.tolist()
 
-  def decode(self, label_id):
+  def decode(self, ids):
+    label_id = ids
     if isinstance(label_id, np.ndarray):
       label_id = np.squeeze(label_id).astype(np.int8).tolist()
     assert isinstance(label_id, list)
@@ -303,8 +303,9 @@ def __init__(self,
       assert vocab_list is not None
       self._init_vocab_from_list(vocab_list)
 
-  def encode(self, sentence):
+  def encode(self, s):
     """Converts a space-separated string of tokens to a list of ids."""
+    sentence = s
     tokens = sentence.strip().split()
     if self._replace_oov is not None:
       tokens = [t if t in self._token_to_id else self._replace_oov
@@ -482,16 +483,16 @@ def __init__(self, filename=None):
       self._load_from_file(filename)
     super(SubwordTextEncoder, self).__init__(num_reserved_ids=None)
 
-  def encode(self, raw_text):
+  def encode(self, s):
     """Converts a native string to a list of subtoken ids.
 
     Args:
-      raw_text: a native string.
+      s: a native string.
     Returns:
       a list of integers in the range [0, vocab_size)
     """
     return self._tokens_to_subtoken_ids(
-        tokenizer.encode(native_to_unicode(raw_text)))
+        tokenizer.encode(native_to_unicode(s)))
 
   def encode_without_tokenizing(self, token_text):
     """Converts string to list of subtoken ids without calling tokenizer.
@@ -510,19 +511,19 @@ def encode_without_tokenizing(self, token_text):
     """
     return self._tokens_to_subtoken_ids([native_to_unicode(token_text)])
 
-  def decode(self, subtokens):
+  def decode(self, ids):
     """Converts a sequence of subtoken ids to a native string.
 
     Args:
-      subtokens: a list of integers in the range [0, vocab_size)
+      ids: a list of integers in the range [0, vocab_size)
     Returns:
       a native string
     """
     return unicode_to_native(
-        tokenizer.decode(self._subtoken_ids_to_tokens(subtokens)))
+        tokenizer.decode(self._subtoken_ids_to_tokens(ids)))
 
-  def decode_list(self, subtokens):
-    return [self._subtoken_id_to_subtoken_string(s) for s in subtokens]
+  def decode_list(self, ids):
+    return [self._subtoken_id_to_subtoken_string(s) for s in ids]
 
   @property
   def vocab_size(self):
diff --git a/tensor2tensor/data_generators/text_encoder_build_subword.py b/tensor2tensor/data_generators/text_encoder_build_subword.py
index 4c5f01f6c..ddcd524b2 100644
--- a/tensor2tensor/data_generators/text_encoder_build_subword.py
+++ b/tensor2tensor/data_generators/text_encoder_build_subword.py
@@ -30,9 +30,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import tokenizer
 
diff --git a/tensor2tensor/data_generators/text_encoder_test.py b/tensor2tensor/data_generators/text_encoder_test.py
index 1606e790c..74c401cf3 100644
--- a/tensor2tensor/data_generators/text_encoder_test.py
+++ b/tensor2tensor/data_generators/text_encoder_test.py
@@ -26,7 +26,6 @@
 import shutil
 import string
 
-# Dependency imports
 import mock
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
diff --git a/tensor2tensor/data_generators/text_problems_test.py b/tensor2tensor/data_generators/text_problems_test.py
index bddc58a6e..ab50cdefc 100644
--- a/tensor2tensor/data_generators/text_problems_test.py
+++ b/tensor2tensor/data_generators/text_problems_test.py
@@ -20,9 +20,6 @@
 
 import os
 import shutil
-
-# Dependency imports
-
 from tensor2tensor.data_generators import problem as problem_lib
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
diff --git a/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index e0f1f3558..716721c48 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -20,7 +20,6 @@
 
 import os
 import shutil
-# Dependency imports
 
 from tensor2tensor.data_generators import timeseries
 
diff --git a/tensor2tensor/data_generators/tokenizer.py b/tensor2tensor/data_generators/tokenizer.py
index 92a42382c..c84642733 100644
--- a/tensor2tensor/data_generators/tokenizer.py
+++ b/tensor2tensor/data_generators/tokenizer.py
@@ -46,9 +46,6 @@
 import collections
 import sys
 import unicodedata
-
-# Dependency imports
-
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
 import tensorflow as tf
diff --git a/tensor2tensor/data_generators/tokenizer_test.py b/tensor2tensor/data_generators/tokenizer_test.py
index 7c5ababd1..01c192327 100644
--- a/tensor2tensor/data_generators/tokenizer_test.py
+++ b/tensor2tensor/data_generators/tokenizer_test.py
@@ -21,9 +21,6 @@
 
 import os
 import random
-
-# Dependency imports
-
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
 from tensor2tensor.data_generators import tokenizer
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 200cc71c0..1edae6564 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -20,9 +20,6 @@
 
 import os
 import tarfile
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
diff --git a/tensor2tensor/data_generators/translate_encs.py b/tensor2tensor/data_generators/translate_encs.py
index fed471ea9..bc4e5df3b 100644
--- a/tensor2tensor/data_generators/translate_encs.py
+++ b/tensor2tensor/data_generators/translate_encs.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
diff --git a/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index 9991ad3f5..70bd53528 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -20,9 +20,6 @@
 
 import os
 import tarfile
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
diff --git a/tensor2tensor/data_generators/translate_enet.py b/tensor2tensor/data_generators/translate_enet.py
index a81b117ca..ec98db06e 100644
--- a/tensor2tensor/data_generators/translate_enet.py
+++ b/tensor2tensor/data_generators/translate_enet.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
diff --git a/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index 56f1f23fc..93d2e3f42 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
diff --git a/tensor2tensor/data_generators/translate_enmk.py b/tensor2tensor/data_generators/translate_enmk.py
index 8a0568f05..67dfefd99 100644
--- a/tensor2tensor/data_generators/translate_enmk.py
+++ b/tensor2tensor/data_generators/translate_enmk.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
diff --git a/tensor2tensor/data_generators/translate_envi.py b/tensor2tensor/data_generators/translate_envi.py
index f102cdff8..3d1709c79 100644
--- a/tensor2tensor/data_generators/translate_envi.py
+++ b/tensor2tensor/data_generators/translate_envi.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import translate
diff --git a/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index a1d9ecd53..1944b0fff 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -19,9 +19,6 @@
 from __future__ import print_function
 
 import os
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
diff --git a/tensor2tensor/data_generators/translate_test.py b/tensor2tensor/data_generators/translate_test.py
index 1cb3c9f36..b34cff1b7 100644
--- a/tensor2tensor/data_generators/translate_test.py
+++ b/tensor2tensor/data_generators/translate_test.py
@@ -21,9 +21,6 @@
 import os
 import shutil
 import tarfile
-
-# Dependency imports
-
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.data_generators import translate
 
diff --git a/tensor2tensor/data_generators/twentybn.py b/tensor2tensor/data_generators/twentybn.py
index 3505b17a2..64c69d169 100644
--- a/tensor2tensor/data_generators/twentybn.py
+++ b/tensor2tensor/data_generators/twentybn.py
@@ -19,9 +19,6 @@
 from __future__ import print_function
 
 import os
-
-# Dependency imports
-
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.utils import registry
 
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 43e049f11..1b8d2370e 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -18,8 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import video_utils
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 1094577c6..77ad56805 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -19,9 +19,6 @@
 from __future__ import print_function
 
 import os
-
-# Dependency imports
-
 import six
 
 from tensor2tensor.data_generators import generator_utils
@@ -120,7 +117,7 @@ def is_generate_per_split(self):
     """
     raise NotImplementedError()
 
-  def example_reading_spec(self, label_repr=None):
+  def example_reading_spec(self):
     extra_data_fields, extra_data_items_to_decoders = self.extra_reading_spec
 
     data_fields = {
@@ -323,7 +320,7 @@ def num_channels(self):
     """Number of color channels."""
     return 3
 
-  def example_reading_spec(self, label_repr=None):
+  def example_reading_spec(self):
     data_fields = {
         "image/encoded": tf.FixedLenFeature((), tf.string),
         "image/format": tf.FixedLenFeature((), tf.string),
diff --git a/tensor2tensor/data_generators/wiki.py b/tensor2tensor/data_generators/wiki.py
index 772898745..45e3387e8 100644
--- a/tensor2tensor/data_generators/wiki.py
+++ b/tensor2tensor/data_generators/wiki.py
@@ -20,9 +20,6 @@
 
 import os
 import subprocess
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import generator_utils
diff --git a/tensor2tensor/data_generators/wikisum/utils_test.py b/tensor2tensor/data_generators/wikisum/utils_test.py
index d57c187ae..2e8a5bcdb 100644
--- a/tensor2tensor/data_generators/wikisum/utils_test.py
+++ b/tensor2tensor/data_generators/wikisum/utils_test.py
@@ -19,9 +19,6 @@
 from __future__ import print_function
 
 import os
-
-# Dependency imports
-
 from tensor2tensor.data_generators.wikisum import utils
 
 import tensorflow as tf
diff --git a/tensor2tensor/data_generators/wikitext103.py b/tensor2tensor/data_generators/wikitext103.py
index 5e1d4f310..7712a717b 100644
--- a/tensor2tensor/data_generators/wikitext103.py
+++ b/tensor2tensor/data_generators/wikitext103.py
@@ -24,9 +24,6 @@
 import collections
 import os
 import zipfile
-
-# Dependency imports
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
diff --git a/tensor2tensor/data_generators/wsj_parsing.py b/tensor2tensor/data_generators/wsj_parsing.py
index 15281b4d0..30d91572b 100644
--- a/tensor2tensor/data_generators/wsj_parsing.py
+++ b/tensor2tensor/data_generators/wsj_parsing.py
@@ -14,11 +14,8 @@
 # limitations under the License.
 """Data generators for parsing data-sets."""
 
-import os
-
-# Dependency imports
-
-from tensor2tensor.data_generators import generator_utils
+# import os
+# from tensor2tensor.data_generators import generator_utils
 
 import tensorflow as tf
 
@@ -103,13 +100,17 @@ def parsing_token_generator(data_dir, tmp_dir, train, source_vocab_size,
   Returns:
     A generator to a dictionary of inputs and outputs.
   """
-  source_symbolizer_vocab = generator_utils.get_or_generate_vocab(
-      data_dir, tmp_dir, "wsj_source.vocab.%d" % source_vocab_size,
-      source_vocab_size)
-  target_symbolizer_vocab = generator_utils.get_or_generate_vocab(
-      data_dir, tmp_dir, "wsj_target.vocab.%d" % target_vocab_size,
-      target_vocab_size)
-  filename = "%s_%s.trees" % (FLAGS.parsing_path, "train" if train else "dev")
-  tree_filepath = os.path.join(tmp_dir, filename)
-  return token_generator(tree_filepath, source_symbolizer_vocab,
-                         target_symbolizer_vocab, 1)
+  # TODO(lukaszkaiser): Correct these calls to generate vocabularies. No data
+  # sources are being passed.
+  del (data_dir, tmp_dir, train, source_vocab_size, target_vocab_size)
+  assert False, "Vocabulary generation not implemented"
+  # source_symbolizer_vocab = generator_utils.get_or_generate_vocab(
+  #     data_dir, tmp_dir, "wsj_source.vocab.%d" % source_vocab_size,
+  #     source_vocab_size)
+  # target_symbolizer_vocab = generator_utils.get_or_generate_vocab(
+  #     data_dir, tmp_dir, "wsj_target.vocab.%d" % target_vocab_size,
+  #     target_vocab_size)
+  # filename = "%s_%s.trees" % (FLAGS.parsing_path, "train" if train else "dev")
+  # tree_filepath = os.path.join(tmp_dir, filename)
+  # return token_generator(tree_filepath, source_symbolizer_vocab,
+  #                        target_symbolizer_vocab, 1)
diff --git a/tensor2tensor/insights/query_processor.py b/tensor2tensor/insights/query_processor.py
index 00a2ca297..c615f4868 100644
--- a/tensor2tensor/insights/query_processor.py
+++ b/tensor2tensor/insights/query_processor.py
@@ -25,9 +25,6 @@ class QueryProcessor(object):
   protos are in better shape.
   """
 
-  def __init__(self):
-    pass
-
   def process(self, query):
     """Returns the generated visualizations for query.
 
diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 45151cd14..4b1f114d8 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -22,7 +22,6 @@
 import math
 import operator
 
-# Dependency imports
 import numpy as np
 
 from six.moves import range  # pylint: disable=redefined-builtin
@@ -2634,8 +2633,7 @@ def compute_qkv(query_antecedent,
   def _compute(inp, depth, filter_width, padding, name):
     if filter_width == 1:
       return common_layers.dense(inp, depth, use_bias=False, name=name)
-    else:
-      return common_layers.conv1d(inp, depth, filter_width, padding, name=name)
+    return common_layers.conv1d(inp, depth, filter_width, padding, name=name)
   q = _compute(
       query_antecedent, total_key_depth, q_filter_width, q_padding, "q")
   k = _compute(
@@ -3101,10 +3099,9 @@ def add_or_set_if(prev_bias, new_bias, condition):
       """Add the bias together while considering the None case."""
       if not condition:
         return prev_bias
-      elif prev_bias is None:
+      if prev_bias is None:
         return new_bias
-      else:
-        return prev_bias + new_bias
+      return prev_bias + new_bias
 
     def mask_and_call_attention(x):
       """Function applied once for each sequence of the batch."""
@@ -3294,9 +3291,9 @@ def eventually_dispatch(dispatcher, value):
   # Iterate over every dispatched group
   list_v_out = []
   for (
-      q,
-      k,
-      v,
+      q_i,
+      k_i,
+      v_i,
       qbc,
       qbo,
       kbc,
@@ -3314,9 +3311,9 @@ def eventually_dispatch(dispatcher, value):
   ):
     list_v_out.append(
         expert_dot_product(
-            q,
-            k,
-            v,
+            q_i,
+            k_i,
+            v_i,
             info_q=BatchInfo(coordinates=qbc, order=qbo),
             info_k=BatchInfo(coordinates=kbc, order=kbo)))
 
@@ -3346,11 +3343,10 @@ def map_fn_switch(fn, elems, use_map_fn=True, **kwargs):
   """
   if use_map_fn:
     return tf.map_fn(fn, elems, **kwargs)
-  else:
-    elems_unpacked = (tf.unstack(e) for e in elems)
-    out_unpacked = [fn(e) for e in zip(*elems_unpacked)]
-    out = tf.stack(out_unpacked)
-    return out
+  elems_unpacked = (tf.unstack(e) for e in elems)
+  out_unpacked = [fn(e) for e in zip(*elems_unpacked)]
+  out = tf.stack(out_unpacked)
+  return out
 
 
 @expert_utils.add_name_scope()
@@ -3388,14 +3384,15 @@ def sparse_dot_product_attention(q, k, v, bi, use_map_fn, experts_params):
 
   @expert_utils.add_name_scope()
   def flatten_first_dims(x):
+    """Reshape such that x is [num_heads, -1, depth]."""
     # Case 1: Either constant batch size of size 1 or batch already flattened
     if x.get_shape().as_list()[0] == 1:
       return tf.squeeze(x, axis=0)
+
     # Case 2: Flatten batch dimension
-    else:
-      x = tf.transpose(x, perm=[1, 0, 2, 3])
-      x = tf.reshape(x, [nb_heads, -1, depth])
-      return x
+    x = tf.transpose(x, perm=[1, 0, 2, 3])
+    x = tf.reshape(x, [nb_heads, -1, depth])
+    return x
 
   def flatten_batch(x):
     if x is None:
@@ -3471,6 +3468,7 @@ def dot_product_batched_head(q, k, v, gates_q, gates_k, mask_right=False):
 
   @expert_utils.add_name_scope()
   def get_dispatcher(gates):
+    """Construct dispatcher for gates."""
     length = common_layers.shape_list(gates)[1]
     # Count the number of ones per batch (and keep the max value)
     nb_elems_to_dispatch = tf.reduce_sum(gates, axis=[1, 2])
@@ -3958,6 +3956,7 @@ def forward_internal(x, wqkv, wo, attention_bias, norm_scale, norm_bias):
 
     @function.Defun(compiled=True)
     def grad_fn(x, wqkv, wo, attention_bias, norm_scale, norm_bias, dy):
+      """Custom gradient function."""
       with tf.control_dependencies([dy]):
         n = common_layers.layer_norm_compute_python(x, epsilon, norm_scale,
                                                     norm_bias)
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index e3c24f5b4..825cde343 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index a617a029b..d3ae28d7d 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from six.moves import zip  # pylint: disable=redefined-builtin
 from tensor2tensor.utils import registry
 
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index dcfdb6894..11c12d5d6 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Utils for attention mechanism for images."""
-# Dependency imports
 
 from six.moves import range  # pylint: disable=redefined-builtin
 
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index e3451b4b3..6a78e1136 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -24,9 +24,6 @@
 from functools import partial
 import math
 
-# Dependency imports
-
-
 import numpy as np
 from six.moves import range  # pylint: disable=redefined-builtin
 
@@ -1433,8 +1430,8 @@ def maybe_zero_out_padding(inputs, kernel_size, nonpadding_mask):
     while nonpadding_mask.get_shape().ndims < inputs.get_shape().ndims:
       nonpadding_mask = tf.expand_dims(nonpadding_mask, -1)
     return inputs * nonpadding_mask
-  else:
-    return inputs
+
+  return inputs
 
 
 def dense_relu_dense(inputs,
@@ -2225,7 +2222,7 @@ def body():
       assert outs
 
       deps = outs
-      if isinstance(outs[0], list) or isinstance(outs[0], tuple):
+      if isinstance(outs[0], (list, tuple)):
         assert len(outs) == 1
         deps = outs[0]
       fn_device_dependency_dict()[key] = deps
@@ -2499,7 +2496,7 @@ def _fn_with_custom_grad(fn, inputs, grad_fn, use_global_vars=False):
   if grad_fn is None:
     return outputs
 
-  if not (isinstance(outputs, tuple) or isinstance(outputs, list)):
+  if not isinstance(outputs, (tuple, list)):
     outputs = [outputs]
   outputs = list(outputs)
 
@@ -2772,7 +2769,7 @@ def grad_fn(inputs, variables, outputs, output_grads):
         with tf.variable_scope(cached_vs[0], reuse=True):
           outputs = fn(*inputs)
 
-    if not (isinstance(outputs, list) or isinstance(outputs, tuple)):
+    if not isinstance(outputs, (list, tuple)):
       outputs = [outputs]
     outputs = list(outputs)
     grads = tf.gradients(outputs, inputs + variables, output_grads)
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index 2d0266651..24f9448da 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 from tensor2tensor.layers import common_layers
 
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 8df2222a3..584d3b7ae 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -13,14 +13,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Discretization bottlenecks used to train discrete latent variables."""
-
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
 from functools import partial
-# Dependency imports
+
 from tensor2tensor.layers import common_layers
+
 import tensorflow as tf
+
 from tensorflow.python.training import moving_averages
 
 
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index 8873e6b61..157ae62bb 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -17,13 +17,14 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-# Dependency imports
+
 import numpy as np
 from tensor2tensor.layers import discretization
 import tensorflow as tf
 
 
 class DiscretizationTest(tf.test.TestCase):
+  """Tests for discretization layers."""
 
   def setUp(self):
     tf.set_random_seed(1234)
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index c136c293d..0c37af2d8 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_layers
@@ -119,8 +116,7 @@ def bottom(self, x):
     self._bottom_was_called = True
     if self._model_hparams.shared_embedding_and_softmax_weights:
       return self.bottom_simple(x, "shared", reuse=None)
-    else:
-      return self.bottom_simple(x, "input_emb", reuse=None)
+    return self.bottom_simple(x, "input_emb", reuse=None)
 
   def targets_bottom(self, x):
     if self._model_hparams.shared_embedding_and_softmax_weights:
@@ -184,8 +180,9 @@ def targets_weights_fn(self):
 class CTCSymbolModality(SymbolModality):
   """SymbolModality that uses CTC loss."""
 
-  def loss(self, logits, targets):
+  def loss(self, top_out, targets):
     """Compute the CTC loss."""
+    logits = top_out
     with tf.name_scope("ctc_loss", values=[logits, targets]):
       # For CTC we assume targets are 1d, [batch, length, 1, 1] here.
       targets_shape = targets.get_shape().as_list()
@@ -214,14 +211,15 @@ class ImageModality(modality.Modality):
   """Modality for images."""
   PIXEL_EMBEDDING_SIZE = 64
 
-  def bottom(self, inputs):
+  def bottom(self, x):
     with tf.variable_scope(self.name):
-      inputs = tf.to_float(inputs)
+      x = tf.to_float(x)
       if not tf.contrib.eager.in_eager_mode():
-        tf.summary.image("inputs", inputs, max_outputs=2)
-      return inputs
+        tf.summary.image("inputs", x, max_outputs=2)
+      return x
 
-  def targets_bottom(self, inputs):
+  def targets_bottom(self, x):
+    inputs = x
     with tf.variable_scope(self.name):
       if not tf.contrib.eager.in_eager_mode():
         tf.summary.image("targets_bottom",
@@ -258,8 +256,9 @@ def top(self, body_output, _):
         tf.summary.image("result", res_argmax, max_outputs=1)
       return res
 
-  def loss(self, logits, targets):
+  def loss(self, top_out, targets):
     """Compute loss numerator and denominator for one shard of output."""
+    logits = top_out
     return common_layers.padded_cross_entropy(
         logits,
         targets,
@@ -306,11 +305,11 @@ def bottom_compress(self, inputs, name="bottom"):
       x.set_shape([None, None, None, self._body_input_depth])
       return x
 
-  def bottom(self, inputs):
-    return self.bottom_compress(inputs, "input_bottom")
+  def bottom(self, x):
+    return self.bottom_compress(x, "input_bottom")
 
-  def targets_bottom(self, inputs):
-    return self.bottom_compress(inputs, "output_bottom")
+  def targets_bottom(self, x):
+    return self.bottom_compress(x, "output_bottom")
 
   def top(self, body_output, _):
     with tf.variable_scope(self.name):
@@ -357,7 +356,8 @@ def get_channel_embeddings(self, io_depth, targets, hidden_size,
 
     return tf.concat(channel_target_embs, axis=-1)
 
-  def targets_bottom(self, inputs):
+  def targets_bottom(self, x):
+    inputs = x
     io_depth = self._model_hparams.num_channels
     tshape = common_layers.shape_list(inputs)
     hidden_size = self._model_hparams.hidden_size
@@ -382,14 +382,15 @@ def top(self, body_output, _):
 class AudioModality(modality.Modality):
   """Performs strided conv compressions for audio data."""
 
-  def bottom(self, inputs):
+  def bottom(self, x):
     """Transform input from data space to model space.
 
     Args:
-      inputs: A Tensor with shape [batch, ...]
+      x: A Tensor with shape [batch, ...]
     Returns:
       body_input: A Tensor with shape [batch, ?, ?, body_input_depth].
     """
+    inputs = x
     with tf.variable_scope(self.name):
       # TODO(aidangomez): Will need to sort out a better audio pipeline
       def xnet_resblock(x, filters, res_relu, name):
@@ -426,14 +427,15 @@ def xnet_resblock(x, filters, res_relu, name):
 class AudioSpectralModality(modality.Modality):
   """Performs strided conv compressions for audio spectral data."""
 
-  def bottom(self, inputs):
+  def bottom(self, x):
     """Transform input from data space to model space.
 
     Args:
-      inputs: A Tensor with shape [batch, ...]
+      x: A Tensor with shape [batch, ...]
     Returns:
       body_input: A Tensor with shape [batch, ?, ?, body_input_depth].
     """
+    inputs = x
     with tf.variable_scope(self.name):
       # TODO(aidangomez): Will need to sort out a better audio pipeline
       def xnet_resblock(x, filters, res_relu, name):
@@ -472,7 +474,8 @@ class VideoModality(modality.Modality):
   """Modality for videos, i.e., time-sequences of frames."""
   PIXEL_EMBEDDING_SIZE = 64
 
-  def bottom(self, inputs):
+  def bottom(self, x):
+    inputs = x
     with tf.variable_scope(self.name):
       common_layers.summarize_video(inputs, "inputs")
       inputs_shape = common_layers.shape_list(inputs)
@@ -487,7 +490,8 @@ def bottom(self, inputs):
           [inputs_shape[0], inputs_shape[2], inputs_shape[3],
            inputs_shape[1] * inputs_shape[4]])
 
-  def targets_bottom(self, inputs, summary_prefix="targets_bottom"):
+  def targets_bottom(self, x, summary_prefix="targets_bottom"):  # pylint: disable=arguments-differ
+    inputs = x
     with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
       common_layers.summarize_video(inputs, summary_prefix)
       inputs_shape = common_layers.shape_list(inputs)
@@ -522,8 +526,9 @@ def top(self, body_output, _):
         tf.summary.image("result", res_argmax, max_outputs=1)
       return res
 
-  def loss(self, logits, targets):
+  def loss(self, top_out, targets):
     """Compute loss numerator and denominator for one shard of output."""
+    logits = top_out
     logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
     targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
     cutoff = getattr(self._model_hparams, "video_modality_loss_cutoff", 0.01)
@@ -539,16 +544,17 @@ def loss(self, logits, targets):
 class VideoModalityEmbed(VideoModality):
   """Video Modality where bottom embeds pixels."""
 
-  def bottom(self, inputs):
+  def bottom(self, x):
     return super(VideoModalityEmbed, self).targets_bottom(
-        inputs, summary_prefix="bottom")
+        x, summary_prefix="bottom")
 
 
 @registry.register_video_modality("bitwise")
 class VideoModalityBitwise(VideoModality):
   """Video Modality where bottom embeds pixels bitwise."""
 
-  def bottom(self, inputs):
+  def bottom(self, x):
+    inputs = x
     with tf.variable_scope(self.name):
       common_layers.summarize_video(inputs, "targets_bottom")
       # Embed bitwise.
@@ -585,8 +591,9 @@ def cutoff(self):
   def internal_loss(self, logits, targets):
     return tf.nn.relu(tf.abs(logits - targets) - self.cutoff)
 
-  def loss(self, logits, targets):
+  def loss(self, top_out, targets):
     """Compute loss numerator and denominator for one shard of output."""
+    logits = top_out
     logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:-1])
     targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
     weights = self.targets_weights_fn(targets)
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 949e0b817..04fba3bff 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.layers import common_hparams
diff --git a/tensor2tensor/layers/rev_block.py b/tensor2tensor/layers/rev_block.py
index 964b294d0..3d8ad5b2e 100644
--- a/tensor2tensor/layers/rev_block.py
+++ b/tensor2tensor/layers/rev_block.py
@@ -23,9 +23,6 @@
 from __future__ import print_function
 
 import re
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_layers
@@ -52,10 +49,10 @@ def _rev_layer_forward(xs, f, g, f_side_input, g_side_input,
   x1, x2 = xs
   y1 = x1 + (f(x2, f_side_input) if f_side_input else f(x2))
   y2 = x2 + (g(y1, g_side_input) if g_side_input else g(y1))
+  out = (y1, y2)
   if gate_outputs:
-    return tf.tuple([y1, y2])
-  else:
-    return (y1, y2)
+    out = tf.tuple(out)
+  return out
 
 
 def _rev_layer_backward(ys, grad_ys, f, g, f_vars, f_side_input, g_vars,
diff --git a/tensor2tensor/layers/rev_block_test.py b/tensor2tensor/layers/rev_block_test.py
index 6c3a10be7..93944373f 100644
--- a/tensor2tensor/layers/rev_block_test.py
+++ b/tensor2tensor/layers/rev_block_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import rev_block
 
 import tensorflow as tf
@@ -74,12 +71,12 @@ def g(x):
       self.assertAllClose(y1, y1_inv)
       self.assertAllClose(y2, y2_inv)
 
-  def _testRevBlock(self,
-                    x=None,
-                    f=None,
-                    g=None,
-                    f_side_input=None,
-                    g_side_input=None):
+  def _test_rev_block(self,
+                      x=None,
+                      f=None,
+                      g=None,
+                      f_side_input=None,
+                      g_side_input=None):
     tf.set_random_seed(1234)
 
     if f is None:
@@ -144,7 +141,7 @@ def g(x):  # pylint: disable=function-redefined
         self.assertAllClose(g1, g2)
 
   def testRevBlock(self):
-    self._testRevBlock()
+    self._test_rev_block()
 
   def testSideInput(self):
     f_side_input = tf.random_uniform([self.BATCH_SIZE, self.CHANNELS // 2])
@@ -153,7 +150,7 @@ def f(x, side_input):
       return tf.layers.dense(
           x, self.CHANNELS // 2, use_bias=True) + side_input[0]
 
-    self._testRevBlock(f=f, f_side_input=[f_side_input])
+    self._test_rev_block(f=f, f_side_input=[f_side_input])
 
   def testMultipleFns(self):
 
@@ -163,23 +160,7 @@ def f1(x):
     def f2(x):
       return tf.layers.dense(x, self.CHANNELS // 2, activation=tf.nn.relu)
 
-    self._testRevBlock(f=[f1, f2, f1, f2])
-
-  # TODO(rsepassi): Recent change to conv seems to have broken this test. Find
-  # out why.
-  def _testConvAndBatchNorm(self):
-
-    x = tf.random_uniform(
-        [self.BATCH_SIZE, 10, self.CHANNELS], dtype=tf.float32)
-
-    def f(x):
-      x = tf.layers.conv1d(x, self.CHANNELS // 2, 3, padding="same")
-      x = tf.layers.batch_normalization(x, training=True)
-      x = tf.layers.conv1d(x, self.CHANNELS // 2, 3, padding="same")
-      x = tf.layers.batch_normalization(x, training=True)
-      return x
-
-    self._testRevBlock(x=x, f=f)
+    self._test_rev_block(f=[f1, f2, f1, f2])
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 5b1357618..b6f7e1d73 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 # pylint: disable=unused-import
 
 from tensor2tensor.layers import modalities
diff --git a/tensor2tensor/models/basic.py b/tensor2tensor/models/basic.py
index c584523a2..3519b7f15 100644
--- a/tensor2tensor/models/basic.py
+++ b/tensor2tensor/models/basic.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import registry
@@ -68,6 +65,7 @@ def unbottleneck(self, x, res_size):
       return x
 
   def bottleneck_loss(self, b):
+    del b
     return 0.0
 
   def make_even_size(self, x):
@@ -150,7 +148,8 @@ def body(self, features):
                             hparams.bottleneck_warmup_steps // 2, is_training)
     return res, {"bottleneck_loss": b_loss}
 
-  def sample(self):
+  def sample(self, features=None):
+    del features
     hp = self.hparams
     div_x = 2**hp.num_hidden_layers
     div_y = 1 if self.is1d else 2**hp.num_hidden_layers
@@ -167,8 +166,9 @@ def encode(self, x):
     self._cur_bottleneck_tensor = None
     return res
 
-  def infer(self, features, *args, **kwargs):
+  def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
     """Produce predictions from the model by sampling."""
+    del args, kwargs
     # Inputs and features preparation needed to handle edge cases.
     if not features:
       features = {}
diff --git a/tensor2tensor/models/basic_test.py b/tensor2tensor/models/basic_test.py
index 1b581b718..7e9b8fd42 100644
--- a/tensor2tensor/models/basic_test.py
+++ b/tensor2tensor/models/basic_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import mnist  # pylint: disable=unused-import
diff --git a/tensor2tensor/models/bytenet.py b/tensor2tensor/models/bytenet.py
index 6f65adb03..a81bbd106 100644
--- a/tensor2tensor/models/bytenet.py
+++ b/tensor2tensor/models/bytenet.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_hparams
diff --git a/tensor2tensor/models/bytenet_test.py b/tensor2tensor/models/bytenet_test.py
index fc920487c..67fff4a50 100644
--- a/tensor2tensor/models/bytenet_test.py
+++ b/tensor2tensor/models/bytenet_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
diff --git a/tensor2tensor/models/distillation.py b/tensor2tensor/models/distillation.py
index 3630cc334..9ef341a77 100644
--- a/tensor2tensor/models/distillation.py
+++ b/tensor2tensor/models/distillation.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index b2a7d7ba9..7c1ab1e6f 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -24,9 +24,6 @@
 from __future__ import print_function
 
 import copy
-
-# Dependency imports
-
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
@@ -99,6 +96,8 @@ def body_sharded(self, sharded_features):
                                    inputs, targets, hparams)
 
     # Run decoder.
+    # TODO(nikip): Use q_padding and kv_padding
+    del q_padding, kv_padding
     decoder_output, extra_loss = cia.transformer_layers_sharded(
         dp,
         self._ps_devices,
@@ -108,8 +107,6 @@ def body_sharded(self, sharded_features):
         self_attention_bias=None,
         enc_output=None,
         attention_type=hparams.dec_attention_type,
-        q_padding=q_padding,
-        kv_padding=kv_padding,
         name="decoder")
 
     output = dp(cia.create_output, decoder_output, rows, cols, targets, hparams)
diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 07aa1231c..6bc9f7ecc 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -24,9 +24,6 @@
 from __future__ import print_function
 
 import copy
-
-# Dependency imports
-
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
@@ -501,6 +498,7 @@ def img2img_transformer2d_n31():
 
 @registry.register_hparams
 def img2img_transformer2d_n24():
+  """Set of hyperparameters."""
   hparams = img2img_transformer2d_base()
   hparams.batch_size = 1
   hparams.hidden_size = 1024
diff --git a/tensor2tensor/models/image_transformer_2d_test.py b/tensor2tensor/models/image_transformer_2d_test.py
index 42cb0ce58..b5dafdd99 100644
--- a/tensor2tensor/models/image_transformer_2d_test.py
+++ b/tensor2tensor/models/image_transformer_2d_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import celeba  # pylint: disable=unused-import
@@ -32,7 +29,7 @@
 
 class Img2imgTransformerTest(tf.test.TestCase):
 
-  def _testImg2imgTransformer(self, net):
+  def _test_img2img_transformer(self, net):
     batch_size = 3
     hparams = image_transformer_2d.img2img_transformer2d_tiny()
     hparams.data_dir = ""
@@ -52,12 +49,12 @@ def _testImg2imgTransformer(self, net):
     self.assertEqual(res.shape, (batch_size, 8, 8, 3, 256))
 
   def testImg2imgTransformer(self):
-    self._testImg2imgTransformer(image_transformer_2d.Img2imgTransformer)
+    self._test_img2img_transformer(image_transformer_2d.Img2imgTransformer)
 
 
 class Imagetransformer2dTest(tf.test.TestCase):
 
-  def _testImagetransformer2d(self, net):
+  def _test_imagetransformer_2d(self, net):
     batch_size = 3
     size = 7
     vocab_size = 256
@@ -80,7 +77,7 @@ def _testImagetransformer2d(self, net):
     self.assertEqual(res.shape, (batch_size, size, size, 3, vocab_size))
 
   def testImagetransformer2d(self):
-    self._testImagetransformer2d(image_transformer_2d.Imagetransformer2d)
+    self._test_imagetransformer_2d(image_transformer_2d.Imagetransformer2d)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/models/image_transformer_test.py b/tensor2tensor/models/image_transformer_test.py
index 9c9110d8e..ad0f45f58 100644
--- a/tensor2tensor/models/image_transformer_test.py
+++ b/tensor2tensor/models/image_transformer_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
@@ -30,7 +27,7 @@
 
 class ImagetransformerTest(tf.test.TestCase):
 
-  def _testImagetransformer(self, net):
+  def _test_imagetransformer(self, net):
     batch_size = 3
     size = 7
     vocab_size = 256
@@ -53,7 +50,7 @@ def _testImagetransformer(self, net):
     self.assertEqual(res.shape, (batch_size, size, size, 3, vocab_size))
 
   def testImagetransformer(self):
-    self._testImagetransformer(image_transformer.Imagetransformer)
+    self._test_imagetransformer(image_transformer.Imagetransformer)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index e2b23c067..05395e042 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -19,9 +19,6 @@
 from __future__ import print_function
 
 import copy
-
-# Dependency imports
-
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import registry
diff --git a/tensor2tensor/models/lstm_test.py b/tensor2tensor/models/lstm_test.py
index e22760311..4b2fcd640 100644
--- a/tensor2tensor/models/lstm_test.py
+++ b/tensor2tensor/models/lstm_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
diff --git a/tensor2tensor/models/neural_gpu.py b/tensor2tensor/models/neural_gpu.py
index e278d4606..44b4473af 100644
--- a/tensor2tensor/models/neural_gpu.py
+++ b/tensor2tensor/models/neural_gpu.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_hparams
diff --git a/tensor2tensor/models/neural_gpu_test.py b/tensor2tensor/models/neural_gpu_test.py
index 08fb2b18c..9a2dab9f1 100644
--- a/tensor2tensor/models/neural_gpu_test.py
+++ b/tensor2tensor/models/neural_gpu_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
diff --git a/tensor2tensor/models/research/aligned.py b/tensor2tensor/models/research/aligned.py
index 7aa731bd9..10a4f1c04 100644
--- a/tensor2tensor/models/research/aligned.py
+++ b/tensor2tensor/models/research/aligned.py
@@ -25,9 +25,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
diff --git a/tensor2tensor/models/research/attention_lm.py b/tensor2tensor/models/research/attention_lm.py
index 1d2d2acb0..ee3d9b9f2 100644
--- a/tensor2tensor/models/research/attention_lm.py
+++ b/tensor2tensor/models/research/attention_lm.py
@@ -23,9 +23,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_attention
diff --git a/tensor2tensor/models/research/attention_lm_moe.py b/tensor2tensor/models/research/attention_lm_moe.py
index dd0163bfb..29ff7a66f 100644
--- a/tensor2tensor/models/research/attention_lm_moe.py
+++ b/tensor2tensor/models/research/attention_lm_moe.py
@@ -25,9 +25,6 @@
 from __future__ import print_function
 
 import functools
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_attention
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index d1ff2c327..1485e04e9 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
@@ -248,7 +245,8 @@ def bottleneck(self, x):
                           hparams.mode == tf.estimator.ModeKeys.TRAIN)
     return x
 
-  def sample(self):
+  def sample(self, features=None):
+    del features
     hp = self.hparams
     div_x = 2**hp.num_hidden_layers
     div_y = 1 if self.is1d else 2**hp.num_hidden_layers
@@ -262,7 +260,7 @@ def sample(self):
 class AutoencoderResidualDiscrete(AutoencoderResidual):
   """Discrete residual autoencoder."""
 
-  def bottleneck(self, x, bottleneck_size=None):
+  def bottleneck(self, x, bottleneck_size=None):  # pylint: disable=arguments-differ
     if bottleneck_size is not None:
       old_bottleneck_size = self.hparams.bottleneck_size
       self.hparams.bottleneck_size = bottleneck_size
@@ -281,7 +279,8 @@ def bottleneck_loss(self, b):
     part_avg = tf.abs(tf.reduce_sum(b * selection)) / (selection_size + 1)
     return part_avg
 
-  def sample(self):
+  def sample(self, features=None):
+    del features
     hp = self.hparams
     div_x = 2**hp.num_hidden_layers
     div_y = 1 if self.is1d else 2**hp.num_hidden_layers
@@ -303,7 +302,7 @@ class AutoencoderOrderedDiscrete(AutoencoderResidualDiscrete):
   def bottleneck_loss(self, unused_b):
     return 0.0
 
-  def bottleneck(self, x):
+  def bottleneck(self, x):  # pylint: disable=arguments-differ
     hparams = self.hparams
     if hparams.unordered:
       return super(AutoencoderOrderedDiscrete, self).bottleneck(x)
diff --git a/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index 23c8108e9..596e8abf0 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import mnist  # pylint: disable=unused-import
@@ -32,8 +29,8 @@
 
 class AutoencoderTest(tf.test.TestCase):
 
-  def getMnistRandomOutput(self, model_name, hparams_set=None,
-                           mode=tf.estimator.ModeKeys.TRAIN):
+  def get_mnist_random_output(self, model_name, hparams_set=None,
+                              mode=tf.estimator.ModeKeys.TRAIN):
     hparams_set = hparams_set or model_name
     x = np.random.random_integers(0, high=255, size=(1, 28, 28, 1))
     y = np.random.random_integers(0, high=9, size=(1, 1))
@@ -52,32 +49,32 @@ def getMnistRandomOutput(self, model_name, hparams_set=None,
     return res
 
   @property
-  def mnistOutputShape(self):
+  def mnist_output_shape(self):
     return (1, 28, 28, 1, 256)
 
   def testAutoencoderAutoregressive(self):
-    res = self.getMnistRandomOutput("autoencoder_autoregressive")
-    self.assertEqual(res.shape, self.mnistOutputShape)
+    res = self.get_mnist_random_output("autoencoder_autoregressive")
+    self.assertEqual(res.shape, self.mnist_output_shape)
 
   def testAutoencoderResidual(self):
-    res = self.getMnistRandomOutput("autoencoder_residual")
-    self.assertEqual(res.shape, self.mnistOutputShape)
+    res = self.get_mnist_random_output("autoencoder_residual")
+    self.assertEqual(res.shape, self.mnist_output_shape)
 
   def testAutoencoderBasicDiscrete(self):
-    res = self.getMnistRandomOutput("autoencoder_basic_discrete")
-    self.assertEqual(res.shape, self.mnistOutputShape)
+    res = self.get_mnist_random_output("autoencoder_basic_discrete")
+    self.assertEqual(res.shape, self.mnist_output_shape)
 
   def testAutoencoderResidualDiscrete(self):
-    res = self.getMnistRandomOutput("autoencoder_residual_discrete")
-    self.assertEqual(res.shape, self.mnistOutputShape)
+    res = self.get_mnist_random_output("autoencoder_residual_discrete")
+    self.assertEqual(res.shape, self.mnist_output_shape)
 
   def testAutoencoderOrderedDiscrete(self):
-    res = self.getMnistRandomOutput("autoencoder_ordered_discrete")
-    self.assertEqual(res.shape, self.mnistOutputShape)
+    res = self.get_mnist_random_output("autoencoder_ordered_discrete")
+    self.assertEqual(res.shape, self.mnist_output_shape)
 
   def testAutoencoderStacked(self):
-    res = self.getMnistRandomOutput("autoencoder_stacked")
-    self.assertEqual(res.shape, self.mnistOutputShape)
+    res = self.get_mnist_random_output("autoencoder_stacked")
+    self.assertEqual(res.shape, self.mnist_output_shape)
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index bba12768c..3974ff2df 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models.research import transformer_vae
 from tensor2tensor.utils import registry
@@ -45,6 +42,7 @@ def generator(x, hparams, name, reuse=False):
 
 
 def lossfn(real_input, fake_input, compress, hparams, lsgan, name):
+  """Loss function."""
   eps = 1e-12
   with tf.variable_scope(name):
     d1 = discriminator(real_input, compress, hparams, "discriminator")
diff --git a/tensor2tensor/models/research/gene_expression.py b/tensor2tensor/models/research/gene_expression.py
index 97134184c..69999b861 100644
--- a/tensor2tensor/models/research/gene_expression.py
+++ b/tensor2tensor/models/research/gene_expression.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_hparams
@@ -99,6 +96,7 @@ def conv_layer(x,
                dropout_rate,
                dilation_rate,
                name="conv"):
+  """Single conv layer with relu, optional pooling, and dropout."""
   with tf.variable_scope(name):
     out = x
     out = common_layers.conv1d_block(
diff --git a/tensor2tensor/models/research/gene_expression_test.py b/tensor2tensor/models/research/gene_expression_test.py
index 06779c978..8e65334de 100644
--- a/tensor2tensor/models/research/gene_expression_test.py
+++ b/tensor2tensor/models/research/gene_expression_test.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import gene_expression as gene_data
@@ -37,7 +34,7 @@ def gene_expression_conv_test():
 
 class GeneExpressionModelsTest(tf.test.TestCase):
 
-  def _testModel(self, hparams, model_cls):
+  def _test_model(self, hparams, model_cls):
     batch_size = 3
     target_length = 6
     target_out = 10  # GeneExpressionProblem.num_output_predictions
@@ -70,7 +67,7 @@ def testGeneExpressionModels(self):
       hparams.add_hparam("data_dir", None)
       p_hparams = gene_data.GenomicsExpressionCage10().get_hparams(hparams)
       hparams.problem_hparams = p_hparams
-      self._testModel(hparams, model_cls)
+      self._test_model(hparams, model_cls)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/models/research/multimodel.py b/tensor2tensor/models/research/multimodel.py
index ccb62bae2..4152e20c3 100644
--- a/tensor2tensor/models/research/multimodel.py
+++ b/tensor2tensor/models/research/multimodel.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_attention
@@ -106,6 +103,7 @@ def prepare_decoder(targets, target_space_emb):
 
 @registry.register_model
 class MultiModel(t2t_model.T2TModel):
+  """Model to train on multiple tasks simultaneously."""
 
   @property
   def use_body_sharded(self):
diff --git a/tensor2tensor/models/research/multimodel_test.py b/tensor2tensor/models/research/multimodel_test.py
index 64a510dab..41febb460 100644
--- a/tensor2tensor/models/research/multimodel_test.py
+++ b/tensor2tensor/models/research/multimodel_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import cifar  # pylint: disable=unused-import
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index dcca82856..63b3edcb8 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import six
 
 from tensor2tensor.layers import common_attention
@@ -119,8 +116,9 @@ def body(self, features):
     reward_pred = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
     return {"targets": x, "target_reward": reward_pred}
 
-  def infer(self, features, *args, **kwargs):
+  def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
     """Produce predictions from the model by running it."""
+    del args, kwargs
     # Inputs and features preparation needed to handle edge cases.
     if not features:
       features = {}
diff --git a/tensor2tensor/models/research/r_transformer.py b/tensor2tensor/models/research/r_transformer.py
index 5312852ea..8f36f9861 100644
--- a/tensor2tensor/models/research/r_transformer.py
+++ b/tensor2tensor/models/research/r_transformer.py
@@ -30,9 +30,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
@@ -46,7 +43,7 @@
 class RTransformer(transformer.Transformer):
   """R-Transformer: Depth-wise recurrent transformer model."""
 
-  def encode(self, inputs, target_space, hparams, features=None):
+  def encode(self, inputs, target_space, hparams, features=None, losses=None):
     """Encode r-transformer inputs.
 
     It is similar to "transformer.encode", but it uses
@@ -60,6 +57,7 @@ def encode(self, inputs, target_space, hparams, features=None):
       hparams: hyperparmeters for model.
       features: optionally pass the entire features dictionary as well.
         This is needed now for "packed" datasets.
+      losses: optional list onto which to append extra training losses
 
     Returns:
       Tuple of:
@@ -70,6 +68,7 @@ def encode(self, inputs, target_space, hparams, features=None):
           encoder_extra_output: which is extra encoder output used in some
             variants of the model (e.g. in ACT, to pass the ponder-time to body)
     """
+    del losses
 
     inputs = common_layers.flatten4d3d(inputs)
 
@@ -96,7 +95,9 @@ def decode(self,
              encoder_decoder_attention_bias,
              decoder_self_attention_bias,
              hparams,
-             nonpadding=None):
+             cache=None,
+             nonpadding=None,
+             losses=None):
     """Decode R-Transformer outputs from encoder representation.
 
     It is similar to "transformer.decode", but it uses
@@ -113,7 +114,10 @@ def decode(self,
       decoder_self_attention_bias: Bias and mask weights for decoder
         self-attention. [batch_size, decoder_length]
       hparams: hyperparmeters for model.
+      cache: dict, containing tensors which are the results of previous
+          attentions, used for fast decoding.
       nonpadding: optional Tensor with shape [batch_size, decoder_length]
+      losses: optional list onto which to append extra training losses
 
     Returns:
        Tuple of:
@@ -123,6 +127,7 @@ def decode(self,
             variants of the model (e.g. in ACT, to pass the ponder-time to body)
 
     """
+    del cache, losses
 
     decoder_input = tf.nn.dropout(decoder_input,
                                   1.0 - hparams.layer_prepostprocess_dropout)
@@ -264,7 +269,7 @@ def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
 class RTransformerEncoder(transformer.Transformer):
   """R-Transformer Encoder: Depth-wise recurrent transformer encoder-only."""
 
-  def encode(self, inputs, target_space, hparams, features=None):
+  def encode(self, inputs, target_space, hparams, features=None, losses=None):
     """Encode transformer inputs.
 
     Args:
@@ -274,6 +279,7 @@ def encode(self, inputs, target_space, hparams, features=None):
       hparams: hyperparmeters for model.
       features: optionally pass the entire features dictionary as well.
         This is needed now for "packed" datasets.
+      losses: optional list onto which to append extra training losses
 
     Returns:
       Tuple of:
@@ -282,6 +288,7 @@ def encode(self, inputs, target_space, hparams, features=None):
           encoder_extra_output: which is extra encoder output used in some
             variants of the model (e.g. in ACT, to pass the ponder-time to body)
     """
+    del losses
     inputs = common_layers.flatten4d3d(inputs)
 
     (encoder_input, self_attention_bias, _) = (
diff --git a/tensor2tensor/models/research/r_transformer_test.py b/tensor2tensor/models/research/r_transformer_test.py
index c9fd6521d..9ec513591 100644
--- a/tensor2tensor/models/research/r_transformer_test.py
+++ b/tensor2tensor/models/research/r_transformer_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
@@ -35,7 +32,8 @@
 
 class RTransformerTest(tf.test.TestCase):
 
-  def getModel(self, hparams, mode=tf.estimator.ModeKeys.TRAIN, has_input=True):
+  def get_model(self, hparams, mode=tf.estimator.ModeKeys.TRAIN,
+                has_input=True):
     hparams.hidden_size = 8
     hparams.filter_size = 32
     hparams.num_heads = 1
@@ -59,7 +57,7 @@ def getModel(self, hparams, mode=tf.estimator.ModeKeys.TRAIN, has_input=True):
     return r_transformer.RTransformer(hparams, mode, p_hparams), features
 
   def testTransformer(self):
-    model, features = self.getModel(r_transformer.r_transformer_base())
+    model, features = self.get_model(r_transformer.r_transformer_base())
     logits, _ = model(features)
     with self.test_session() as session:
       session.run(tf.global_variables_initializer())
diff --git a/tensor2tensor/models/research/r_transformer_util.py b/tensor2tensor/models/research/r_transformer_util.py
index 93dee1160..521adbf8a 100644
--- a/tensor2tensor/models/research/r_transformer_util.py
+++ b/tensor2tensor/models/research/r_transformer_util.py
@@ -39,9 +39,6 @@
 
 import copy
 import functools
-
-# Dependency imports
-
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 5e44e2261..c8d57792c 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -17,9 +17,6 @@
 import collections
 import functools
 import operator
-
-# Dependency imports
-
 import gym
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
diff --git a/tensor2tensor/models/research/super_lm.py b/tensor2tensor/models/research/super_lm.py
index adcf5f3ad..f90e40d57 100644
--- a/tensor2tensor/models/research/super_lm.py
+++ b/tensor2tensor/models/research/super_lm.py
@@ -26,9 +26,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_attention
diff --git a/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index c0d3ed2e9..222fd3365 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -19,9 +19,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
@@ -243,8 +240,8 @@ def _extract_layer_types(self):
 
     # Extend the blocks and fill them with the default values if not specified
     final_layers = ([], [])
-    for i, blocks_str in enumerate(layer_types):
-      for blocks_str in blocks_str.split(SEP_LAYER):
+    for i, blocks_str_joined in enumerate(layer_types):
+      for blocks_str in blocks_str_joined.split(SEP_LAYER):
         if not blocks_str:
           continue
         blocks_list = blocks_str.split(SEP_FF)
diff --git a/tensor2tensor/models/research/transformer_revnet.py b/tensor2tensor/models/research/transformer_revnet.py
index 0bafb546b..5935e1c0d 100644
--- a/tensor2tensor/models/research/transformer_revnet.py
+++ b/tensor2tensor/models/research/transformer_revnet.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import rev_block
diff --git a/tensor2tensor/models/research/transformer_revnet_test.py b/tensor2tensor/models/research/transformer_revnet_test.py
index 4393943ce..38571aefa 100644
--- a/tensor2tensor/models/research/transformer_revnet_test.py
+++ b/tensor2tensor/models/research/transformer_revnet_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
diff --git a/tensor2tensor/models/research/transformer_sketch.py b/tensor2tensor/models/research/transformer_sketch.py
index d26f8e96b..bbe27d077 100644
--- a/tensor2tensor/models/research/transformer_sketch.py
+++ b/tensor2tensor/models/research/transformer_sketch.py
@@ -18,9 +18,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
 from tensor2tensor.models.research import transformer_vae
@@ -33,7 +30,7 @@
 class TransformerSketch(transformer.Transformer):
   """Transformer with strided convolutions."""
 
-  def encode(self, inputs, target_space, hparams):
+  def encode(self, inputs, target_space, hparams, features=None, losses=None):
     """Add two layers strided convolutions ontop of encode."""
     inputs = common_layers.conv_block(
         inputs,
@@ -44,12 +41,13 @@ def encode(self, inputs, target_space, hparams):
         name="small_image_conv")
 
     hparams.num_compress_steps = 2
-    compressed_inputs = transformer_vae.compress(inputs, is_2d=True,
+    compressed_inputs = transformer_vae.compress(inputs, None, is_2d=True,
                                                  hparams=hparams,
                                                  name="convolutions")
 
     return super(TransformerSketch, self).encode(
-        compressed_inputs, target_space, hparams)
+        compressed_inputs, target_space, hparams, features=features,
+        losses=losses)
 
 
 @registry.register_hparams
diff --git a/tensor2tensor/models/research/transformer_symshard.py b/tensor2tensor/models/research/transformer_symshard.py
index fe0741550..0e182fc36 100644
--- a/tensor2tensor/models/research/transformer_symshard.py
+++ b/tensor2tensor/models/research/transformer_symshard.py
@@ -42,9 +42,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_attention
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 288a54f6f..d7f6e193f 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -20,9 +20,6 @@
 
 import functools
 import math
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_attention
@@ -59,6 +56,7 @@ def residual_conv(x, repeat, k, hparams, name, reuse=None):
 
 
 def attend(x, source, hparams, name):
+  """Self-attention layer with source as memory antecedent."""
   with tf.variable_scope(name):
     x = tf.squeeze(x, axis=2)
     if len(source.get_shape()) > 3:
@@ -181,7 +179,9 @@ def decode_transformer(encoder_output,
                                      hparams.num_channels*hparams.hidden_size])
 
       # Prepare decoder inputs and bias.
-      decoder_input, _, _, bias = cia.prepare_decoder(targets, hparams)
+      # TODO(nikip): Make prepare_decoder return bias
+      decoder_input, _, _ = cia.prepare_decoder(targets, hparams)
+      bias = None
 
       # Add class label to decoder input.
       if not hparams.drop_inputs:
@@ -606,11 +606,12 @@ def prepare_features_for_infer(self, features):
     features["cache_raw"] = cache
 
   def infer(self, features=None, decode_length=50, beam_size=1, top_beams=1,
-            alpha=0.0):
+            alpha=0.0, use_tpu=False):
     """Produce predictions from the model."""
     if not self._hparams.do_mask:
-      return super(TransformerAE, self).infer(
-          features, decode_length, beam_size, top_beams, alpha)["outputs"]
+      infer_out = super(TransformerAE, self).infer(
+          features, decode_length, beam_size, top_beams, alpha, use_tpu=use_tpu)
+      return infer_out["outputs"]
     if not features:
       features = {}
     inputs_old = None
diff --git a/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index ed753225b..c2db3e064 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -18,9 +18,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import registry
diff --git a/tensor2tensor/models/resnet_test.py b/tensor2tensor/models/resnet_test.py
index 24c11a325..18baff6e2 100644
--- a/tensor2tensor/models/resnet_test.py
+++ b/tensor2tensor/models/resnet_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
@@ -38,7 +35,7 @@ def resnet_tiny_cpu():
 
 class ResnetTest(tf.test.TestCase):
 
-  def _testResnet(self, img_size, output_size):
+  def _test_resnet(self, img_size, output_size):
     vocab_size = 9
     batch_size = 2
     x = np.random.random_integers(
@@ -61,7 +58,7 @@ def _testResnet(self, img_size, output_size):
     self.assertEqual(res.shape, (batch_size,) + output_size + (1, vocab_size))
 
   def testResnetLarge(self):
-    self._testResnet(img_size=224, output_size=(1, 1))
+    self._test_resnet(img_size=224, output_size=(1, 1))
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/models/revnet.py b/tensor2tensor/models/revnet.py
index 7ddab9a2b..6db6f8940 100644
--- a/tensor2tensor/models/revnet.py
+++ b/tensor2tensor/models/revnet.py
@@ -35,9 +35,6 @@
 """
 
 import functools
-
-# Dependency imports
-
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import rev_block
 from tensor2tensor.utils import registry
diff --git a/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index f05659a89..964290d11 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import registry
diff --git a/tensor2tensor/models/slicenet.py b/tensor2tensor/models/slicenet.py
index e2c2271aa..fc9ecbf50 100644
--- a/tensor2tensor/models/slicenet.py
+++ b/tensor2tensor/models/slicenet.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 from six.moves import zip  # pylint: disable=redefined-builtin
 
diff --git a/tensor2tensor/models/slicenet_test.py b/tensor2tensor/models/slicenet_test.py
index 85a6d161b..89919aa10 100644
--- a/tensor2tensor/models/slicenet_test.py
+++ b/tensor2tensor/models/slicenet_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import cifar  # pylint: disable=unused-import
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 502a9969c..3699e50a1 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -25,9 +25,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.data_generators import librispeech
@@ -574,9 +571,10 @@ def infer(self,
             decode_length=50,
             beam_size=1,
             top_beams=1,
-            alpha=0.0):
+            alpha=0.0,
+            use_tpu=False):
     """Returns the targets and their log probabilities."""
-    del decode_length, beam_size, top_beams, alpha
+    del decode_length, beam_size, top_beams, alpha, use_tpu
     assert features is not None
 
     # Run the model
diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index fdcf86731..5da51fecf 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
diff --git a/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index 0aa32b136..dcf32586f 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -20,9 +20,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
@@ -36,6 +33,7 @@ def lrelu(input_, leak=0.2, name="lrelu"):
 
 def deconv2d(
     input_, output_shape, k_h, k_w, d_h, d_w, stddev=0.02, name="deconv2d"):
+  """Deconvolution layer."""
   with tf.variable_scope(name):
     w = tf.get_variable(
         "w", [k_h, k_w, output_shape[-1], input_.get_shape()[-1]],
@@ -188,13 +186,8 @@ def top(self, body_output, features):
 class VanillaGan(AbstractGAN):
   """Simple GAN for demonstration."""
 
-  def infer(self,
-            features=None,
-            decode_length=50,
-            beam_size=1,
-            top_beams=1,
-            last_position_only=False,
-            alpha=0.0):
+  def infer(self, *args, **kwargs):  # pylint: disable=arguments-differ
+    del args, kwargs
     with tf.variable_scope("body/vanilla_gan", reuse=tf.AUTO_REUSE):
       z = tf.random_uniform(
           shape=[1, self._hparams.random_sample_size],
diff --git a/tensor2tensor/models/xception.py b/tensor2tensor/models/xception.py
index 0730966de..2452a7d4f 100644
--- a/tensor2tensor/models/xception.py
+++ b/tensor2tensor/models/xception.py
@@ -19,9 +19,6 @@
 from __future__ import print_function
 
 import math
-
-# Dependency imports
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_hparams
@@ -73,9 +70,11 @@ def xception_internal(inputs, hparams):
 
 
 def xception_entry(inputs, hidden_dim):
+  """Xception entry flow."""
   with tf.variable_scope("xception_entry"):
 
     def xnet_resblock(x, filters, res_relu, name):
+      """Resblock."""
       with tf.variable_scope(name):
         y = common_layers.separable_conv_block(
             x,
@@ -111,6 +110,7 @@ def xnet_resblock(x, filters, res_relu, name):
 
 
 def xception_exit(inputs):
+  """Xception exit flow."""
   with tf.variable_scope("xception_exit"):
     x = inputs
     x_shape = x.get_shape().as_list()
diff --git a/tensor2tensor/models/xception_test.py b/tensor2tensor/models/xception_test.py
index 553a218b1..b57a757b9 100644
--- a/tensor2tensor/models/xception_test.py
+++ b/tensor2tensor/models/xception_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
@@ -31,7 +28,7 @@
 
 class XceptionTest(tf.test.TestCase):
 
-  def _testXception(self, img_size):
+  def _test_xception(self, img_size):
     vocab_size = 9
     batch_size = 3
     x = np.random.random_integers(
@@ -54,10 +51,10 @@ def _testXception(self, img_size):
     self.assertEqual(res.shape, (batch_size, 1, 1, 1, vocab_size))
 
   def testXceptionSmallImage(self):
-    self._testXception(img_size=9)
+    self._test_xception(img_size=9)
 
   def testXceptionLargeImage(self):
-    self._testXception(img_size=256)
+    self._test_xception(img_size=256)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/problems.py b/tensor2tensor/problems.py
index ae0a42da1..c986c8eb5 100644
--- a/tensor2tensor/problems.py
+++ b/tensor2tensor/problems.py
@@ -12,17 +12,11 @@
 # 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.
-"""Access T2T Problems.
-
-See problems_test.py for basic usage.
-"""
-
+"""Access T2T Problems."""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-# Dependency imports
-
 from tensor2tensor.data_generators import all_problems  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
 
diff --git a/tensor2tensor/problems_test.py b/tensor2tensor/problems_test.py
deleted file mode 100644
index 7c67f3de9..000000000
--- a/tensor2tensor/problems_test.py
+++ /dev/null
@@ -1,59 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-"""tensor2tensor.problems test."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-# Dependency imports
-
-from tensor2tensor import problems
-
-import tensorflow as tf
-
-MODES = tf.estimator.ModeKeys
-
-
-class ProblemsTest(tf.test.TestCase):
-
-  def testBuildDataset(self):
-    # See all the available problems
-    self.assertTrue(len(problems.available()) > 10)
-
-    # Retrieve a problem by name
-    problem = problems.problem("translate_ende_wmt8k")
-
-    # Access train and dev datasets through Problem
-    train_dataset = problem.dataset(MODES.TRAIN)
-    dev_dataset = problem.dataset(MODES.EVAL)
-
-    # Access vocab size and other info (e.g. the data encoders used to
-    # encode/decode data for the feature, used below) through feature_info.
-    feature_info = problem.feature_info
-    self.assertTrue(feature_info["inputs"].vocab_size > 0)
-    self.assertTrue(feature_info["targets"].vocab_size > 0)
-
-    train_example = train_dataset.make_one_shot_iterator().get_next()
-    dev_example = dev_dataset.make_one_shot_iterator().get_next()
-
-    with tf.Session() as sess:
-      train_ex_val, _ = sess.run([train_example, dev_example])
-      _ = feature_info["inputs"].encoder.decode(train_ex_val["inputs"])
-      _ = feature_info["targets"].encoder.decode(train_ex_val["targets"])
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index 011da95bf..92dca08ab 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -20,9 +20,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import gym
 
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 7df0f3211..eac9834d6 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -20,9 +20,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import common_layers
 from tensor2tensor.rl.envs import in_graph_batch_env
 from tensor2tensor.utils import registry
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 250e2c7de..55e4bb807 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
 
 import tensorflow as tf
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index 0a5f3c478..de2af016a 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -29,9 +29,6 @@
 import subprocess
 import sys
 import traceback
-
-# Dependency imports
-
 import gym
 
 from tensor2tensor.rl.envs import batch_env
@@ -76,13 +73,13 @@ def _reset(self, **kwargs):
     if self._reset_counter % 2 == 1:
       self._active = True
       return self.env.reset(**kwargs)
-    else:
-      self._active = False
-      self._last_returned = (self._last_returned[0],
-                             self._last_returned[1],
-                             False,  # done = False
-                             self._last_returned[3])
-      return self._last_returned[0]
+
+    self._active = False
+    self._last_returned = (self._last_returned[0],
+                           self._last_returned[1],
+                           False,  # done = False
+                           self._last_returned[3])
+    return self._last_returned[0]
 
 
 class ExternalProcessEnv(object):
@@ -207,8 +204,7 @@ def step(self, action, blocking=True):
     promise = self.call("step", action)
     if blocking:
       return promise()
-    else:
-      return promise
+    return promise
 
   def reset(self, blocking=True):
     """Reset the environment.
@@ -223,8 +219,7 @@ def reset(self, blocking=True):
     promise = self.call("reset")
     if blocking:
       return promise()
-    else:
-      return promise
+    return promise
 
   def _receive(self):
     """Wait for a message from the worker process and return its payload.
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 92e860412..09ee739ed 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -30,9 +30,6 @@
 import math
 import os
 import time
-
-# Dependency imports
-
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.layers import discretization
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index f79cc74c2..f039b9eb0 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -19,9 +19,6 @@
 
 import functools
 import os
-
-# Dependency imports
-
 import gym
 
 from tensor2tensor import models  # pylint: disable=unused-import
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index bb433fa80..9e14c4e9a 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.utils import trainer_lib
 
diff --git a/tensor2tensor/rl/t2t_rl_trainer.py b/tensor2tensor/rl/t2t_rl_trainer.py
index 2cab99d57..715dc10ae 100644
--- a/tensor2tensor/rl/t2t_rl_trainer.py
+++ b/tensor2tensor/rl/t2t_rl_trainer.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 1f65b67cb..f3dc0ac76 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -18,9 +18,6 @@
 from __future__ import print_function
 
 import os
-
-# Dependency imports
-
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import trainer_lib
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index e755d9551..c20fe8b48 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -19,8 +19,6 @@
 from __future__ import print_function
 
 import base64
-
-# Dependency imports
 from googleapiclient import discovery
 from grpc.beta import implementations
 
diff --git a/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index 6a668a2a2..1f65f5905 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -17,7 +17,6 @@
 from __future__ import division
 from __future__ import print_function
 
-# Dependency imports
 from tensor2tensor.utils import quantization
 
 import tensorflow as tf
@@ -206,7 +205,8 @@ def _parameter_scale(self, var):
     """
     return tf.maximum(reduce_rms(var), self._epsilon2)
 
-  def _resource_apply_dense(self, grad, var):
+  def _resource_apply_dense(self, grad, handle):
+    var = handle
     grad = tf.to_float(grad)
     grad_squared = tf.square(grad) + self._epsilon1
     grad_squared_mean = tf.reduce_mean(grad_squared)
diff --git a/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index 15175405d..e498bc18a 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -18,9 +18,6 @@
 from __future__ import print_function
 
 import os
-
-# Dependency imports
-
 import numpy as np
 import six
 from six.moves import zip  # pylint: disable=redefined-builtin
diff --git a/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index c01df0276..fff8bd226 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.layers import common_layers
 
 import tensorflow as tf
diff --git a/tensor2tensor/utils/beam_search_test.py b/tensor2tensor/utils/beam_search_test.py
index 2c1262724..bd6a9cc97 100644
--- a/tensor2tensor/utils/beam_search_test.py
+++ b/tensor2tensor/utils/beam_search_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 from tensor2tensor.utils import beam_search
 
diff --git a/tensor2tensor/utils/bleu_hook.py b/tensor2tensor/utils/bleu_hook.py
index a0fa7fe3e..56bff62b0 100644
--- a/tensor2tensor/utils/bleu_hook.py
+++ b/tensor2tensor/utils/bleu_hook.py
@@ -24,9 +24,6 @@
 import sys
 import time
 import unicodedata
-
-# Dependency imports
-
 import numpy as np
 import six
 # pylint: disable=redefined-builtin
@@ -242,7 +239,7 @@ def _read_stepfiles_list(path_prefix, path_suffix=".index", min_steps=0):
   """Return list of StepFiles sorted by step from files at path_prefix."""
   stepfiles = []
   for filename in _try_twice_tf_glob(path_prefix + "*-[0-9]*" + path_suffix):
-    basename = filename[:-len(path_suffix)] if len(path_suffix) else filename
+    basename = filename[:-len(path_suffix)] if path_suffix else filename
     try:
       steps = int(basename.rsplit("-")[-1])
     except ValueError:  # The -[0-9]* part is not an integer.
diff --git a/tensor2tensor/utils/bleu_hook_test.py b/tensor2tensor/utils/bleu_hook_test.py
index 522573b08..017236912 100644
--- a/tensor2tensor/utils/bleu_hook_test.py
+++ b/tensor2tensor/utils/bleu_hook_test.py
@@ -18,9 +18,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.utils import bleu_hook
 
 import tensorflow as tf
diff --git a/tensor2tensor/utils/checkpoint_compatibility_test.py b/tensor2tensor/utils/checkpoint_compatibility_test.py
index 73a2958e3..a1a121ab3 100644
--- a/tensor2tensor/utils/checkpoint_compatibility_test.py
+++ b/tensor2tensor/utils/checkpoint_compatibility_test.py
@@ -29,9 +29,6 @@
 from __future__ import print_function
 
 import os
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor import models  # pylint: disable=unused-import
diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 37158af0a..f9a4080a8 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -95,15 +95,16 @@ def get_default_master_type(num_gpus=1, use_tpu=False):
   """Returns master_type for trainingInput."""
   if use_tpu:
     return "cloud_tpu"
-  elif num_gpus <= 0:
-    return "standard"
-  elif num_gpus == 1:
-    return "standard_p100"
-  elif num_gpus == 4:
-    return "complex_model_m_p100"
-  elif num_gpus == 8:
-    return "complex_model_l_gpu"
-  assert False
+  gpus_to_master_map = {
+      0: "standard",
+      1: "standard_p100",
+      4: "complex_model_m_p100",
+      8: "complex_model_l_gpu",
+  }
+  if num_gpus not in gpus_to_master_map:
+    raise ValueError("Num gpus must be in %s" %
+                     str(sorted(list(gpus_to_master_map.keys()))))
+  return gpus_to_master_map[num_gpus]
 
 
 def configure_job():
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index ed8dad700..fa9facb64 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 import six
diff --git a/tensor2tensor/utils/data_reader_test.py b/tensor2tensor/utils/data_reader_test.py
index 176785321..a677eec98 100644
--- a/tensor2tensor/utils/data_reader_test.py
+++ b/tensor2tensor/utils/data_reader_test.py
@@ -20,9 +20,6 @@
 
 import os
 import tempfile
-
-# Dependency imports
-
 import numpy as np
 from six.moves import range  # pylint: disable=redefined-builtin
 
@@ -38,6 +35,7 @@
 class TestProblem(problem_mod.Problem):
 
   def generator(self, data_dir, tmp_dir, is_training):
+    del data_dir, tmp_dir, is_training
     for i in range(30):
       yield {"inputs": [i] * (i + 1), "targets": [i], "floats": [i + 0.5]}
 
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index ce917b0cd..15996932c 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -19,9 +19,6 @@
 
 import operator
 import os
-
-# Dependency imports
-
 import numpy as np
 import six
 
diff --git a/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index 9371a4885..12e6a3d9a 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -18,9 +18,6 @@
 from __future__ import print_function
 
 import inspect
-
-# Dependency imports
-
 from tensor2tensor.utils import expert_utils as eu
 import tensorflow as tf
 
diff --git a/tensor2tensor/utils/diet.py b/tensor2tensor/utils/diet.py
index e3db8c10a..ec8212fa8 100644
--- a/tensor2tensor/utils/diet.py
+++ b/tensor2tensor/utils/diet.py
@@ -24,7 +24,7 @@
 from collections import defaultdict
 import copy
 import math
-# Dependency imports
+
 from tensor2tensor.layers import common_layers
 import tensorflow as tf
 
@@ -259,6 +259,7 @@ def make_diet_var_getter(params):
   """Create a custom variable getter for diet variables according to params."""
 
   def diet_var_initializer(shape, dtype, partition_info=None):
+    """Initializer for a diet variable."""
     del dtype
     del partition_info
 
@@ -296,6 +297,7 @@ def _fn_with_diet_vars(fn, args, params):
   vs_ctr = []
 
   def grad_fn(inputs, variables, outputs, output_grads):
+    """Custom gradient function."""
     del outputs  # recomputing below
     with common_layers.fn_device_dependency("diet_grad",
                                             output_grads[0].device) as out_dep:
diff --git a/tensor2tensor/utils/diet_test.py b/tensor2tensor/utils/diet_test.py
index 93fcaefa5..809a5769f 100644
--- a/tensor2tensor/utils/diet_test.py
+++ b/tensor2tensor/utils/diet_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.utils import diet
 
 import tensorflow as tf
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 46461a815..323b2578b 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -24,9 +24,6 @@
 
 import functools
 import math
-
-# Dependency imports
-
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
 from six.moves import zip  # pylint: disable=redefined-builtin
@@ -1017,9 +1014,9 @@ def local_moe(x,
       expert_kwargs["x"] = dispatcher.dispatch(x_flat)
     if pass_gates:
       expert_kwargs["gates"] = dispatcher.expert_to_gates()
-    for k, v in six.iteritems(additional_dispatch_params or {}):
-      v = flatten_all_but_last(v)
-      expert_kwargs[k] = dispatcher.dispatch(v)
+    for key, val in six.iteritems(additional_dispatch_params or {}):
+      val = flatten_all_but_last(val)
+      expert_kwargs[key] = dispatcher.dispatch(val)
 
     ep = Parallelism([DEFAULT_DEV_STRING] * num_experts, reuse=None)
     expert_outputs = ep(expert_fn, **expert_kwargs)
diff --git a/tensor2tensor/utils/expert_utils_test.py b/tensor2tensor/utils/expert_utils_test.py
index b90edec8c..33209fc47 100644
--- a/tensor2tensor/utils/expert_utils_test.py
+++ b/tensor2tensor/utils/expert_utils_test.py
@@ -13,12 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Tests for tensor2tensor.utils.expert_utils."""
-
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-# Dependency imports
 from tensor2tensor.layers import common_attention
 from tensor2tensor.utils import expert_utils
 import tensorflow as tf
diff --git a/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index 7cd9b3f3f..d9b753997 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -18,7 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-# Dependency imports
 import tensorflow as tf
 
 flags = tf.flags
diff --git a/tensor2tensor/utils/get_rouge.py b/tensor2tensor/utils/get_rouge.py
index a56319154..474bb2393 100644
--- a/tensor2tensor/utils/get_rouge.py
+++ b/tensor2tensor/utils/get_rouge.py
@@ -22,10 +22,7 @@
 import os
 import shutil
 from tempfile import mkdtemp
-
-# Dependency imports
-
-from pyrouge import Rouge155
+import pyrouge
 import tensorflow as tf
 
 FLAGS = tf.flags.FLAGS
@@ -53,7 +50,7 @@ def prep_data(decode_dir, target_dir):
 
 
 def main(_):
-  rouge = Rouge155()
+  rouge = pyrouge.Rouge155()
   rouge.log.setLevel(logging.ERROR)
   rouge.system_filename_pattern = "rouge.(\\d+).txt"
   rouge.model_filename_pattern = "rouge.[A-Z].#ID#.txt"
diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index 3768bbe78..843e484ab 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 import tensorflow as tf
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 35a328dea..1d9ac96b7 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -18,9 +18,6 @@
 from __future__ import print_function
 
 import inspect
-
-# Dependency imports
-
 import numpy as np
 import six
 
diff --git a/tensor2tensor/utils/metrics_hook.py b/tensor2tensor/utils/metrics_hook.py
index b87564e0c..2482b487b 100644
--- a/tensor2tensor/utils/metrics_hook.py
+++ b/tensor2tensor/utils/metrics_hook.py
@@ -18,9 +18,6 @@
 from __future__ import print_function
 
 import os
-
-# Dependency imports
-
 import tensorflow as tf
 
 from tensorboard.backend.event_processing import event_accumulator
@@ -85,6 +82,7 @@ def after_run(self, run_context, run_values):
     self._after_run(run_context, run_values, global_step, metrics)
 
   def _after_run(self, run_context, run_values, global_step, metrics):
+    del run_values
     if self._process_metrics(global_step, metrics):
       run_context.request_stop()
 
@@ -117,6 +115,7 @@ def _process_metrics(self, global_step, metrics):
     Returns:
       should_stop: bool. If True, will request that the session stops.
     """
+    del global_step, metrics
     return False
 
 
@@ -156,10 +155,10 @@ def __init__(self,
 
   def _process_metrics(self, global_step, metrics):
     if not metrics:
-      return
+      return None
 
     if not list(metrics.values())[0]:
-      return
+      return None
 
     # Metrics should have just a single subdir and a single tag
     steps, vals = list(metrics.values())[0][self._tags[0]]
diff --git a/tensor2tensor/utils/metrics_hook_test.py b/tensor2tensor/utils/metrics_hook_test.py
index 3eac23208..35695d1ae 100644
--- a/tensor2tensor/utils/metrics_hook_test.py
+++ b/tensor2tensor/utils/metrics_hook_test.py
@@ -21,9 +21,6 @@
 import contextlib
 import os
 import shutil
-
-# Dependency imports
-
 from tensor2tensor.utils import metrics_hook
 
 import tensorflow as tf
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index a47379a65..3c19a7bf5 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 from tensor2tensor.utils import metrics
 
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index c26154783..d38ee737b 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -18,9 +18,6 @@
 from __future__ import print_function
 
 import re
-
-# Dependency imports
-
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import registry
 
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 40870a4f4..1c29ca519 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.utils import adafactor
@@ -71,7 +68,7 @@ def optimize(loss, learning_rate, hparams, use_tpu=False):
 class ConditionalOptimizer(tf.train.Optimizer):
   """Conditional optimizer."""
 
-  def __init__(self, optimizer_name, lr, hparams, use_tpu=False):
+  def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disable=super-init-not-called
     if optimizer_name == "Adam" and use_tpu:
       # LazyAdamOptimizer does not work on TPU
       optimizer_name = "TrueAdam"
@@ -105,7 +102,7 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):
     else:
       self._opt = tf.contrib.layers.OPTIMIZER_CLS_NAMES[optimizer_name](lr)
 
-  def compute_gradients(self, loss, var_list=None, **kwargs):
+  def compute_gradients(self, loss, var_list=None, **kwargs):  # pylint: disable=arguments-differ
     gradients = self._opt.compute_gradients(loss, var_list, **kwargs)
     def cast_grad(g, v):
       if v is None or g is None:
diff --git a/tensor2tensor/utils/quantization.py b/tensor2tensor/utils/quantization.py
index 19e8a8b24..6883423a0 100644
--- a/tensor2tensor/utils/quantization.py
+++ b/tensor2tensor/utils/quantization.py
@@ -16,9 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import tensorflow as tf
 
 from tensorflow.python.framework import function
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index ee1efb868..f73a92e5e 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -45,9 +45,6 @@ class MyModel(T2TModel):
 
 import inspect
 import re
-
-# Dependency imports
-
 import six
 import tensorflow as tf
 
diff --git a/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index b0c85027f..bc593edf5 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 from tensor2tensor.utils import modality
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
diff --git a/tensor2tensor/utils/rouge_test.py b/tensor2tensor/utils/rouge_test.py
index 2760c4b42..f90be3ba7 100644
--- a/tensor2tensor/utils/rouge_test.py
+++ b/tensor2tensor/utils/rouge_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 from tensor2tensor.utils import rouge
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 01e6f9590..b989e9470 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -23,9 +23,6 @@
 import functools
 import math
 import time
-
-# Dependency imports
-
 import six
 
 from tensor2tensor.data_generators import text_encoder
@@ -144,7 +141,9 @@ def _custom_getter(self):
     else:
       return None
 
-  def call(self, features):
+  def call(self, inputs, **kwargs):
+    del kwargs
+    features = inputs
     set_custom_getter_compose(self._custom_getter)
     tf.get_variable_scope().set_initializer(
         optimize.get_variable_initializer(self.hparams))
@@ -537,6 +536,7 @@ def infer(self,
           "losses": a dictionary: {loss-name (string): floating point `Scalar`
       }
     """
+    del use_tpu
     set_custom_getter_compose(self._custom_getter)
     with self._eager_var_store.as_default():
       # TODO(rsepassi): Make decoding work with real-valued model outputs
@@ -1032,6 +1032,7 @@ def estimator_spec_train(self, loss, num_async_replicas=1):
 
   def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
     """Construct EstimatorSpec for EVAL mode."""
+    del losses_dict
     hparams = self.hparams
 
     if not hasattr(hparams, "problem"):
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 4367cc6f4..f1387de8e 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -20,9 +20,6 @@
 
 import os
 import random
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.utils import devices
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index 271394fe1..e735c05b0 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -20,9 +20,6 @@
 
 import os
 import shutil
-
-# Dependency imports
-
 from tensor2tensor import models  # pylint: disable=unused-import
 from tensor2tensor.data_generators import algorithmic
 from tensor2tensor.data_generators import generator_utils
@@ -36,7 +33,8 @@
 @registry.register_problem
 class TinyAlgo(algorithmic.AlgorithmicIdentityBinary40):
 
-  def generate_data(self, data_dir, _):
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    del tmp_dir, task_id
     identity_problem = algorithmic.AlgorithmicIdentityBinary40()
     generator_utils.generate_files(
         identity_problem.generator(self.num_symbols, 40, 100000),
diff --git a/tensor2tensor/utils/usr_dir.py b/tensor2tensor/utils/usr_dir.py
index 367d98c1d..f66b21d11 100644
--- a/tensor2tensor/utils/usr_dir.py
+++ b/tensor2tensor/utils/usr_dir.py
@@ -20,9 +20,6 @@
 import importlib
 import os
 import sys
-
-# Dependency imports
-
 import tensorflow as tf
 
 
diff --git a/tensor2tensor/utils/yellowfin.py b/tensor2tensor/utils/yellowfin.py
index 6e0252882..0090c6fbf 100644
--- a/tensor2tensor/utils/yellowfin.py
+++ b/tensor2tensor/utils/yellowfin.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import tensorflow as tf
 
 
diff --git a/tensor2tensor/utils/yellowfin_test.py b/tensor2tensor/utils/yellowfin_test.py
index 914068e41..dd487689c 100644
--- a/tensor2tensor/utils/yellowfin_test.py
+++ b/tensor2tensor/utils/yellowfin_test.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.utils.yellowfin import YellowFinOptimizer
@@ -33,7 +30,8 @@
 
 class YellowFinTest(tf.test.TestCase):
 
-  def tuneEverything(self, x0squared, c, t, gmin, gmax):
+  def tune_everything(self, x0squared, c, t, gmin, gmax):
+    del t
     # First tune based on dynamic range
     if c == 0:
       dr = gmax / gmin
@@ -191,11 +189,11 @@ def testLrMu(self):
                        0.001 * g_norm_avg / g_norm_squared_avg)
 
         if i > 0:
-          lr, mu = self.tuneEverything(target_dist**2,
-                                       target_var,
-                                       1,
-                                       target_h_min,
-                                       target_h_max)
+          lr, mu = self.tune_everything(target_dist**2,
+                                        target_var,
+                                        1,
+                                        target_h_min,
+                                        target_h_max)
           target_lr = 0.999 * target_lr + 0.001 * lr
           target_mu = 0.999 * target_mu + 0.001 * mu
 
diff --git a/tensor2tensor/visualization/visualization.py b/tensor2tensor/visualization/visualization.py
index c32b44295..59d07dc33 100644
--- a/tensor2tensor/visualization/visualization.py
+++ b/tensor2tensor/visualization/visualization.py
@@ -17,9 +17,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-# Dependency imports
-
 import numpy as np
 
 # To register the hparams set

From b1b947ef129e028f3c8b43ede23902ce08a5da55 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Wed, 23 May 2018 17:13:16 -0700
Subject: [PATCH 0010/2720] Try logit normalization for latent predictions

PiperOrigin-RevId: 197816221
---
 tensor2tensor/models/research/transformer_vae.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index d7f6e193f..12e772d1d 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -220,6 +220,10 @@ def ae_latent_softmax(latents_pred, latents_discrete, hparams):
   if hparams.num_decode_blocks < 2:
     latents_logits = tf.layers.dense(latents_pred, vocab_size,
                                      name="extra_logits")
+    if hparams.logit_normalization:
+      latents_logits *= tf.rsqrt(1e-8 +
+                                 tf.reduce_mean(tf.square(latents_logits)))
+
     loss = None
     if latents_discrete is not None:
       if hparams.soft_em:
@@ -667,6 +671,7 @@ def transformer_ae_small():
   hparams.add_hparam("z_size", 14)
   hparams.add_hparam("noise_dev", 0.5)
   hparams.add_hparam("d_mix", 0.5)
+  hparams.add_hparam("logit_normalization", True)
   # Bottleneck kinds supported: dense, vae, semhash, gumbel-softmax, dvq.
   hparams.add_hparam("bottleneck_kind", "semhash")
   hparams.add_hparam("num_blocks", 1)

From 0c1d9850caba8acb8bda5dadc9c84ba9c6993a4e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 23 May 2018 18:13:27 -0700
Subject: [PATCH 0011/2720] Adding stochastic shape video data generator.

PiperOrigin-RevId: 197823190
---
 .../data_generators/video_generated.py        | 110 +++++++++++++++++-
 1 file changed, 105 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 1b8d2370e..03d62357f 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -18,6 +18,10 @@
 from __future__ import division
 from __future__ import print_function
 
+import math
+import google3
+
+import matplotlib.pyplot as plt
 import numpy as np
 
 from tensor2tensor.data_generators import video_utils
@@ -57,6 +61,10 @@ def frame_width(self):
   def total_number_of_frames(self):
     return 10000
 
+  @property
+  def video_length(self):
+    return 5
+
   @property
   def extra_reading_spec(self):
     """Additional data fields to store on disk and their decoders."""
@@ -77,10 +85,102 @@ def hparams(self, defaults, unused_model_hparams):
     }
     p.target_modality = ("video", 256)
 
+  @staticmethod
+  def get_circle(x, y, z, c, s):
+    """Draws a circle with center(x, y), color c, size s and z-order of z."""
+    cir = plt.Circle((x, y), s, fc=c, zorder=z)
+    return cir
+
+  @staticmethod
+  def get_rectangle(x, y, z, c, s):
+    """Draws a rectangle with center(x, y), color c, size s and z-order of z."""
+    rec = plt.Rectangle((x-s, y-s), s*2.0, s*2.0, fc=c, zorder=z)
+    return rec
+
+  @staticmethod
+  def get_triangle(x, y, z, c, s):
+    """Draws a triangle with center (x, y), color c, size s and z-order of z."""
+    points = np.array([[0, 0], [s, s*math.sqrt(3.0)], [s*2.0, 0]])
+    tri = plt.Polygon(points + [x-s, y-s], fc=c, zorder=z)
+    return tri
+
+  def generate_stochastic_shape_instance(self):
+    """Yields one video of a shape moving to a random direction.
+
+       The size and color of the shapes are random but
+       consistent in a single video. The speed is fixed.
+
+    Raises:
+       ValueError: The frame size is not square.
+    """
+    if self.frame_height != self.frame_width or self.frame_height % 2 != 0:
+      raise ValueError("Generator only supports square frames with even size.")
+
+    lim = 10.0
+    direction = np.array([[+1.0, +1.0],
+                          [+1.0, +0.0],
+                          [+1.0, -1.0],
+                          [+0.0, +1.0],
+                          [+0.0, -1.0],
+                          [-1.0, +1.0],
+                          [-1.0, +0.0],
+                          [-1.0, -1.0]
+                         ])
+
+    rnd = np.random.randint(len(direction))
+    sp = np.array([lim/2.0, lim/2.0])
+    di = direction[rnd]
+
+    colors = ["b", "g", "r", "c", "m", "y", "k"]
+    color = np.random.choice(colors)
+
+    shape = np.random.choice([
+        VideoStochasticShapes10k.get_circle,
+        VideoStochasticShapes10k.get_rectangle,
+        VideoStochasticShapes10k.get_triangle])
+    speed = 1.0
+
+    size = np.random.uniform(0.5, 1.5)
+
+    back_color = str(0.0)
+    plt.ioff()
+
+    xy = np.array(sp)
+    di = direction[0]
+
+    for _ in range(self.video_length):
+      fig = plt.figure()
+      fig.set_dpi(self.frame_height//2)
+      fig.set_size_inches(2, 2)
+      ax = plt.axes(xlim=(0, lim), ylim=(0, lim))
+
+      # Background
+      ax.add_patch(VideoStochasticShapes10k.get_rectangle(
+          0.0, 0.0, -1.0, back_color, 25.0))
+      # Foreground
+      ax.add_patch(shape(xy[0], xy[1], 0.0, color, size))
+
+      plt.axis("off")
+      plt.tight_layout(pad=-2.0)
+      fig.canvas.draw()
+      image = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
+      image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+      image = np.copy(np.uint8(image))
+
+      plt.close()
+      xy += speed * di
+
+      yield image
+
   def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
     frame_number = 0
-    for _ in range(self.total_number_of_frames):
-      frame = np.zeros([self.frame_height, self.frame_width, self.num_channels],
-                       dtype=np.uint8)
-      yield {"frame": frame, "frame_number": [frame_number]}
-      frame_number += 1
+    done = False
+    while not done:
+      for frame_number, frame in enumerate(
+          self.generate_stochastic_shape_instance()):
+        if frame_number >= self.total_number_of_frames:
+          done = True
+          break
+
+        yield {"frame": frame, "frame_number": [frame_number]}
+        frame_number += 1

From ff639c809c81684006bf381a0889f16a64db8be6 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 23 May 2018 18:44:56 -0700
Subject: [PATCH 0012/2720] internal

PiperOrigin-RevId: 197826090
---
 tensor2tensor/data_generators/video_generated.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 03d62357f..863f211ac 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -19,7 +19,6 @@
 from __future__ import print_function
 
 import math
-import google3
 
 import matplotlib.pyplot as plt
 import numpy as np

From e30ef62ececf6e68d95759cdbbb366b20defaafc Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 23 May 2018 19:49:25 -0700
Subject: [PATCH 0013/2720] Fix Celeba Py3, avoid unnecessary cast in optimize

PiperOrigin-RevId: 197830648
---
 tensor2tensor/data_generators/celeba.py       |  2 +-
 .../data_generators/generator_utils.py        | 41 +++++++++----------
 tensor2tensor/utils/optimize.py               |  2 +
 3 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index ccb6cd7fc..0cc1a6319 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -120,7 +120,7 @@ def process_attrs(raw_data):
       landmarks = img_landmarks[img_name]
       attrs = img_attrs[img_name]
 
-      with tf.gfile.Open(filename, "r") as f:
+      with tf.gfile.Open(filename, "rb") as f:
         encoded_image_data = f.read()
         yield {
             "image/encoded": [encoded_image_data],
diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 2667f927d..af665d16b 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -203,29 +203,28 @@ def maybe_download(directory, filename, uri):
   Returns:
     The path to the downloaded file.
   """
-  if not tf.gfile.Exists(directory):
-    tf.logging.info("Creating directory %s" % directory)
-    tf.gfile.MakeDirs(directory)
+  tf.gfile.MakeDirs(directory)
   filepath = os.path.join(directory, filename)
-  if not tf.gfile.Exists(filepath):
-    tf.logging.info("Downloading %s to %s" % (uri, filepath))
-    try:
-      tf.gfile.Copy(uri, filepath)
-    except tf.errors.UnimplementedError:
-      if uri.startswith("http"):
-        inprogress_filepath = filepath + ".incomplete"
-        inprogress_filepath, _ = urllib.urlretrieve(
-            uri, inprogress_filepath, reporthook=download_report_hook)
-        # Print newline to clear the carriage return from the download progress
-        print()
-        tf.gfile.Rename(inprogress_filepath, filepath)
-      else:
-        raise ValueError("Unrecognized URI: " + filepath)
-    statinfo = os.stat(filepath)
-    tf.logging.info("Successfully downloaded %s, %s bytes." %
-                    (filename, statinfo.st_size))
-  else:
+  if tf.gfile.Exists(filepath):
     tf.logging.info("Not downloading, file already found: %s" % filepath)
+    return filepath
+
+  tf.logging.info("Downloading %s to %s" % (uri, filepath))
+  try:
+    tf.gfile.Copy(uri, filepath)
+  except tf.errors.UnimplementedError:
+    if uri.startswith("http"):
+      inprogress_filepath = filepath + ".incomplete"
+      inprogress_filepath, _ = urllib.urlretrieve(
+          uri, inprogress_filepath, reporthook=download_report_hook)
+      # Print newline to clear the carriage return from the download progress
+      print()
+      tf.gfile.Rename(inprogress_filepath, filepath)
+    else:
+      raise ValueError("Unrecognized URI: " + filepath)
+  statinfo = os.stat(filepath)
+  tf.logging.info("Successfully downloaded %s, %s bytes." %
+                  (filename, statinfo.st_size))
   return filepath
 
 
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 1c29ca519..dcb7b91b1 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -107,6 +107,8 @@ def compute_gradients(self, loss, var_list=None, **kwargs):  # pylint: disable=a
     def cast_grad(g, v):
       if v is None or g is None:
         return (g, v)
+      if g.dtype == v.dtype:
+        return (g, v)
       return (tf.cast(g, v.dtype), v)
     gradients = [cast_grad(g, v) for g, v in gradients]
     return gradients

From 3878da3a5d1bc72415e1571d2c2b63211c2dcca1 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 24 May 2018 12:46:17 -0700
Subject: [PATCH 0014/2720] Make autoencoders work with VQ loss.

PiperOrigin-RevId: 197937379
---
 tensor2tensor/layers/discretization.py        |  35 +++---
 tensor2tensor/layers/modalities.py            |  15 ++-
 tensor2tensor/models/basic.py                 |  17 +--
 tensor2tensor/models/research/autoencoders.py | 109 +++++++++++-------
 .../models/research/autoencoders_test.py      |   5 +
 tensor2tensor/models/research/next_frame.py   |   4 +-
 tensor2tensor/rl/model_rl_experiment.py       |   8 +-
 tensor2tensor/utils/expert_utils.py           |   2 +
 8 files changed, 120 insertions(+), 75 deletions(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 584d3b7ae..02977ac19 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -693,6 +693,7 @@ def get_vq_bottleneck(bottleneck_size, hidden_size):
           name="ema_means",
           initializer=means.initialized_value(),
           trainable=False)
+
   return means, ema_means, ema_count
 
 
@@ -717,8 +718,10 @@ def vq_discrete_bottleneck(x,
                            decay=0.999,
                            epsilon=1e-5):
   """Simple vector quantized discrete bottleneck."""
-  hidden_size = common_layers.shape_list(x)[-1]
+  x_shape = common_layers.shape_list(x)
+  hidden_size = x_shape[-1]
   means, ema_means, ema_count = get_vq_bottleneck(bottleneck_size, hidden_size)
+  x = tf.reshape(x, [-1, hidden_size])
   x_means_hot, e_loss = vq_nearest_neighbor(x, means)
 
   # Update the ema variables
@@ -738,26 +741,29 @@ def vq_discrete_bottleneck(x,
       (updated_ema_count + epsilon) / (n + bottleneck_size * epsilon) * n)
   updated_ema_means /= tf.expand_dims(updated_ema_count, axis=-1)
   with tf.control_dependencies([e_loss]):
-    update_means = tf.assign(means, updated_ema_means)
+    update_means = means.assign(updated_ema_means)
     with tf.control_dependencies([update_means]):
       loss = beta * e_loss
 
-  d = x_means_hot
+  d = tf.reshape(x_means_hot, x_shape[:-1] + [bottleneck_size])
   return d, loss
 
 
 def vq_discrete_unbottleneck(x, hidden_size):
   """Simple undiscretization from vector quantized representation."""
+  x_shape = common_layers.shape_list(x)
   x = tf.to_float(x)
-  bottleneck_size = common_layers.shape_list(x)[1]
+  bottleneck_size = common_layers.shape_list(x)[-1]
   means, _, _ = get_vq_bottleneck(bottleneck_size, hidden_size)
-  return tf.matmul(x, means)
+  result = tf.matmul(tf.reshape(x, [-1, x_shape[-1]]), means)
+  return tf.reshape(result,
+                    x_shape[:-1] + [common_layers.shape_list(means)[-1]])
 
 
-def tanh_discrete_bottleneck(x, bottleneck_size, bottleneck_noise,
+def tanh_discrete_bottleneck(x, bottleneck_bits, bottleneck_noise,
                              discretize_warmup_steps, mode):
   """Simple discretization through tanh, flip bottleneck_noise many bits."""
-  x = tf.tanh(tf.layers.dense(x, bottleneck_size,
+  x = tf.tanh(tf.layers.dense(x, bottleneck_bits,
                               name="tanh_discrete_bottleneck"))
   d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
   if mode == tf.estimator.ModeKeys.TRAIN:
@@ -766,7 +772,7 @@ def tanh_discrete_bottleneck(x, bottleneck_size, bottleneck_noise,
     d *= noise
   d = common_layers.mix(d, x, discretize_warmup_steps,
                         mode == tf.estimator.ModeKeys.TRAIN)
-  return d
+  return d, 0.0
 
 
 def tanh_discrete_unbottleneck(x, hidden_size):
@@ -775,12 +781,12 @@ def tanh_discrete_unbottleneck(x, hidden_size):
   return x
 
 
-def isemhash_bottleneck(x, bottleneck_size, bottleneck_noise,
+def isemhash_bottleneck(x, bottleneck_bits, bottleneck_noise,
                         discretize_warmup_steps, mode,
                         isemhash_noise_dev=0.5, isemhash_mix_prob=0.5):
   """Improved semantic hashing bottleneck."""
   with tf.variable_scope("isemhash_bottleneck"):
-    x = tf.layers.dense(x, bottleneck_size, name="dense")
+    x = tf.layers.dense(x, bottleneck_bits, name="dense")
     y = common_layers.saturating_sigmoid(x)
     if isemhash_noise_dev > 0 and mode == tf.estimator.ModeKeys.TRAIN:
       noise = tf.truncated_normal(
@@ -795,7 +801,7 @@ def isemhash_bottleneck(x, bottleneck_size, bottleneck_noise,
       d = common_layers.mix(d, 2.0 * y - 1.0, discretize_warmup_steps,
                             mode == tf.estimator.ModeKeys.TRAIN,
                             max_prob=isemhash_mix_prob)
-    return d
+    return d, 0.0
 
 
 def isemhash_unbottleneck(x, hidden_size, isemhash_filter_size_multiplier=1.0):
@@ -813,15 +819,16 @@ def parametrized_bottleneck(x, hparams):
   """Meta-function calling all the above bottlenecks with hparams."""
   if hparams.bottleneck_kind == "tanh_discrete":
     return tanh_discrete_bottleneck(
-        x, hparams.bottleneck_size, hparams.bottleneck_noise * 0.5,
+        x, hparams.bottleneck_bits, hparams.bottleneck_noise * 0.5,
         hparams.discretize_warmup_steps, hparams.mode)
   if hparams.bottleneck_kind == "isemhash":
     return isemhash_bottleneck(
-        x, hparams.bottleneck_size, hparams.bottleneck_noise * 0.5,
+        x, hparams.bottleneck_bits, hparams.bottleneck_noise * 0.5,
         hparams.discretize_warmup_steps, hparams.mode,
         hparams.isemhash_noise_dev, hparams.isemhash_mix_prob)
   if hparams.bottleneck_kind == "vq":
-    return vq_discrete_bottleneck(x, hparams.bottleneck_size, hparams.vq_beta,
+    bottleneck_size = 2**hparams.bottleneck_bits
+    return vq_discrete_bottleneck(x, bottleneck_size, hparams.vq_beta,
                                   hparams.vq_decay, hparams.vq_epsilon)
   raise ValueError("Unsupported hparams.bottleneck_kind %s"
                    % hparams.bottleneck_kind)
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 0c37af2d8..b94ff92bc 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -555,7 +555,20 @@ class VideoModalityBitwise(VideoModality):
 
   def bottom(self, x):
     inputs = x
-    with tf.variable_scope(self.name):
+    with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
+      common_layers.summarize_video(inputs, "bottom")
+      # Embed bitwise.
+      assert self.top_dimensionality == 256
+      embedded = discretization.int_to_bit_embed(
+          inputs, 8, self.PIXEL_EMBEDDING_SIZE)
+      # Transpose and project.
+      transposed = common_layers.time_to_channels(embedded)
+      return tf.layers.dense(transposed, self._body_input_depth,
+                             name="merge_pixel_embedded_frames")
+
+  def targets_bottom(self, x):
+    inputs = x
+    with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
       common_layers.summarize_video(inputs, "targets_bottom")
       # Embed bitwise.
       assert self.top_dimensionality == 256
diff --git a/tensor2tensor/models/basic.py b/tensor2tensor/models/basic.py
index 3519b7f15..d0aded7bc 100644
--- a/tensor2tensor/models/basic.py
+++ b/tensor2tensor/models/basic.py
@@ -53,21 +53,17 @@ def __init__(self, *args, **kwargs):
   def bottleneck(self, x):
     with tf.variable_scope("bottleneck"):
       hparams = self.hparams
-      x = tf.layers.dense(x, hparams.bottleneck_size, name="bottleneck")
+      x = tf.layers.dense(x, hparams.bottleneck_bits, name="bottleneck")
       if hparams.mode == tf.estimator.ModeKeys.TRAIN:
         noise = 2.0 * tf.random_uniform(common_layers.shape_list(x)) - 1.0
-        return tf.tanh(x) + noise * hparams.bottleneck_noise
-      return tf.tanh(x)
+        return tf.tanh(x) + noise * hparams.bottleneck_noise, 0.0
+      return tf.tanh(x), 0.0
 
   def unbottleneck(self, x, res_size):
     with tf.variable_scope("unbottleneck"):
       x = tf.layers.dense(x, res_size, name="dense")
       return x
 
-  def bottleneck_loss(self, b):
-    del b
-    return 0.0
-
   def make_even_size(self, x):
     shape = [dim if dim is not None else -1 for dim in x.get_shape().as_list()]
     if shape[1] % 2 == 0 and shape[2] % 2 == 0:
@@ -119,9 +115,8 @@ def body(self, features):
       # Run encoder.
       x = self.encoder(x)
       # Bottleneck (mix during early training, not too important but stable).
-      b = self.bottleneck(x)
+      b, b_loss = self.bottleneck(x)
       self._cur_bottleneck_tensor = b
-      b_loss = self.bottleneck_loss(b)
       b = self.unbottleneck(b, common_layers.shape_list(x)[-1])
       b = common_layers.mix(b, x, hparams.bottleneck_warmup_steps, is_training)
       # With probability bottleneck_max_prob use the bottleneck, otherwise x.
@@ -154,7 +149,7 @@ def sample(self, features=None):
     div_x = 2**hp.num_hidden_layers
     div_y = 1 if self.is1d else 2**hp.num_hidden_layers
     size = [hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
-            hp.bottleneck_size]
+            hp.bottleneck_bits]
     # Sample in [-1, 1] as the bottleneck is under tanh.
     return 2.0 * tf.random_uniform(size) - 1.0
 
@@ -258,7 +253,7 @@ def basic_autoencoder():
   hparams.kernel_width = 4
   hparams.dropout = 0.1
   hparams.add_hparam("max_hidden_size", 1024)
-  hparams.add_hparam("bottleneck_size", 128)
+  hparams.add_hparam("bottleneck_bits", 128)
   hparams.add_hparam("bottleneck_noise", 0.1)
   hparams.add_hparam("bottleneck_warmup_steps", 3000)
   hparams.add_hparam("bottleneck_max_prob", 1.0)
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 1485e04e9..891787b96 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -235,7 +235,7 @@ class AutoencoderBasicDiscrete(AutoencoderAutoregressive):
 
   def bottleneck(self, x):
     hparams = self.hparams
-    x = tf.tanh(tf.layers.dense(x, hparams.bottleneck_size, name="bottleneck"))
+    x = tf.tanh(tf.layers.dense(x, hparams.bottleneck_bits, name="bottleneck"))
     d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
     if hparams.mode == tf.estimator.ModeKeys.TRAIN:
       noise = tf.random_uniform(common_layers.shape_list(x))
@@ -243,7 +243,7 @@ def bottleneck(self, x):
       d *= noise
     x = common_layers.mix(d, x, hparams.discretize_warmup_steps,
                           hparams.mode == tf.estimator.ModeKeys.TRAIN)
-    return x
+    return x, 0.0
 
   def sample(self, features=None):
     del features
@@ -251,7 +251,7 @@ def sample(self, features=None):
     div_x = 2**hp.num_hidden_layers
     div_y = 1 if self.is1d else 2**hp.num_hidden_layers
     size = [hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
-            hp.bottleneck_size]
+            hp.bottleneck_bits]
     rand = tf.random_uniform(size)
     return 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
 
@@ -260,32 +260,32 @@ def sample(self, features=None):
 class AutoencoderResidualDiscrete(AutoencoderResidual):
   """Discrete residual autoencoder."""
 
-  def bottleneck(self, x, bottleneck_size=None):  # pylint: disable=arguments-differ
-    if bottleneck_size is not None:
-      old_bottleneck_size = self.hparams.bottleneck_size
-      self.hparams.bottleneck_size = bottleneck_size
-    res = discretization.parametrized_bottleneck(x, self.hparams)
-    if bottleneck_size is not None:
-      self.hparams.bottleneck_size = old_bottleneck_size
-    return res
-
-  def unbottleneck(self, x, res_size):
-    return discretization.parametrized_unbottleneck(x, res_size, self.hparams)
-
-  def bottleneck_loss(self, b):
+  def variance_loss(self, b):
     part = tf.random_uniform(common_layers.shape_list(b))
     selection = tf.to_float(tf.less(part, tf.random_uniform([])))
     selection_size = tf.reduce_sum(selection)
     part_avg = tf.abs(tf.reduce_sum(b * selection)) / (selection_size + 1)
     return part_avg
 
+  def bottleneck(self, x, bottleneck_bits=None):  # pylint: disable=arguments-differ
+    if bottleneck_bits is not None:
+      old_bottleneck_bits = self.hparams.bottleneck_bits
+      self.hparams.bottleneck_bits = bottleneck_bits
+    res, loss = discretization.parametrized_bottleneck(x, self.hparams)
+    if bottleneck_bits is not None:
+      self.hparams.bottleneck_bits = old_bottleneck_bits
+    return res, loss
+
+  def unbottleneck(self, x, res_size):
+    return discretization.parametrized_unbottleneck(x, res_size, self.hparams)
+
   def sample(self, features=None):
     del features
     hp = self.hparams
     div_x = 2**hp.num_hidden_layers
     div_y = 1 if self.is1d else 2**hp.num_hidden_layers
     size = [hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
-            hp.bottleneck_size]
+            hp.bottleneck_bits]
     rand = tf.random_uniform(size)
     res = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
     # If you want to set some first bits to a fixed value, do this:
@@ -299,22 +299,19 @@ def sample(self, features=None):
 class AutoencoderOrderedDiscrete(AutoencoderResidualDiscrete):
   """Ordered discrete autoencoder."""
 
-  def bottleneck_loss(self, unused_b):
-    return 0.0
-
   def bottleneck(self, x):  # pylint: disable=arguments-differ
     hparams = self.hparams
     if hparams.unordered:
       return super(AutoencoderOrderedDiscrete, self).bottleneck(x)
     noise = hparams.bottleneck_noise
     hparams.bottleneck_noise = 0.0  # We'll add noise below.
-    x = discretization.parametrized_bottleneck(x, hparams)
+    x, loss = discretization.parametrized_bottleneck(x, hparams)
     hparams.bottleneck_noise = noise
     if hparams.mode == tf.estimator.ModeKeys.TRAIN:
-      # We want a number p such that p^bottleneck_size = 1 - noise.
-      # So log(p) * bottleneck_size = log(noise)
-      log_p = tf.log(1 - float(noise) / 2) / float(hparams.bottleneck_size)
-      # Probabilities of flipping are p, p^2, p^3, ..., p^bottleneck_size.
+      # We want a number p such that p^bottleneck_bits = 1 - noise.
+      # So log(p) * bottleneck_bits = log(noise)
+      log_p = tf.log(1 - float(noise) / 2) / float(hparams.bottleneck_bits)
+      # Probabilities of flipping are p, p^2, p^3, ..., p^bottleneck_bits.
       noise_mask = 1.0 - tf.exp(tf.cumsum(tf.zeros_like(x) + log_p, axis=-1))
       # Having the no-noise mask, we can make noise just uniformly at random.
       ordered_noise = tf.random_uniform(tf.shape(x))
@@ -322,24 +319,25 @@ def bottleneck(self, x):  # pylint: disable=arguments-differ
       ordered_noise = tf.to_float(tf.less(noise_mask, ordered_noise))
       # Now we flip the bits of x on the noisy positions (ordered and normal).
       x *= 2.0 * ordered_noise - 1
-    return x
+    return x, loss
 
 
 @registry.register_model
 class AutoencoderStacked(AutoencoderResidualDiscrete):
   """A stacked autoencoder."""
 
-  def stack(self, b, size, bottleneck_size, name):
+  def stack(self, b, size, bottleneck_bits, name):
     with tf.variable_scope(name + "_stack"):
       unb = self.unbottleneck(b, size)
       enc = self.encoder(unb)
-      return self.bottleneck(enc, bottleneck_size=bottleneck_size)
+      b, _ = self.bottleneck(enc, bottleneck_bits=bottleneck_bits)
+      return b
 
-  def unstack(self, b, size, bottleneck_size, name):
+  def unstack(self, b, size, bottleneck_bits, name):
     with tf.variable_scope(name + "_unstack"):
       unb = self.unbottleneck(b, size)
       dec = self.decoder(unb)
-      pred = tf.layers.dense(dec, bottleneck_size, name="pred")
+      pred = tf.layers.dense(dec, bottleneck_bits, name="pred")
       pred_shape = common_layers.shape_list(pred)
       pred1 = tf.reshape(pred, pred_shape[:-1] + [-1, 2])
       x, y = tf.split(pred1, 2, axis=-1)
@@ -357,13 +355,12 @@ def stack_loss(self, b, b_pred, name):
           labels=labels_discrete, logits=b_pred)
       return tf.reduce_mean(loss)
 
-  def full_stack(self, b, x_size, bottleneck_size, losses, is_training, i):
-    stack1_b = self.stack(b, x_size, bottleneck_size, "step%d" % i)
+  def full_stack(self, b, x_size, bottleneck_bits, losses, is_training, i):
+    stack1_b = self.stack(b, x_size, bottleneck_bits, "step%d" % i)
     if i > 1:
-      stack1_b = self.full_stack(stack1_b, 2 * x_size, 2 * bottleneck_size,
+      stack1_b = self.full_stack(stack1_b, 2 * x_size, 2 * bottleneck_bits,
                                  losses, is_training, i - 1)
-    b1, b_pred = self.unstack(stack1_b, x_size, bottleneck_size, "step%d" % i)
-    losses["bottleneck%d_loss" % i] = self.bottleneck_loss(stack1_b)
+    b1, b_pred = self.unstack(stack1_b, x_size, bottleneck_bits, "step%d" % i)
     losses["stack%d_loss" % i] = self.stack_loss(b, b_pred, "step%d" % i)
     b_shape = common_layers.shape_list(b)
     if is_training:
@@ -390,10 +387,9 @@ def body(self, features):
       x = self.encoder(x)
       x_size = common_layers.shape_list(x)[-1]
       # Bottleneck (mix during early training, not too important but stable).
-      b = self.bottleneck(x)
-      b_loss = self.bottleneck_loss(b)
+      b, b_loss = self.bottleneck(x)
       losses = {"bottleneck0_loss": b_loss}
-      b = self.full_stack(b, 2 * x_size, 2 * hparams.bottleneck_size,
+      b = self.full_stack(b, 2 * x_size, 2 * hparams.bottleneck_bits,
                           losses, is_training, num_stacks - 1)
       b = self.unbottleneck(b, x_size)
       b = common_layers.mix(b, x, hparams.bottleneck_warmup_steps, is_training)
@@ -460,7 +456,7 @@ def autoencoder_basic_discrete():
   hparams = autoencoder_autoregressive()
   hparams.num_hidden_layers = 5
   hparams.hidden_size = 64
-  hparams.bottleneck_size = 4096
+  hparams.bottleneck_bits = 4096
   hparams.bottleneck_noise = 0.1
   hparams.bottleneck_warmup_steps = 3000
   hparams.add_hparam("discretize_warmup_steps", 5000)
@@ -471,7 +467,7 @@ def autoencoder_basic_discrete():
 def autoencoder_residual_discrete():
   """Residual discrete autoencoder model."""
   hparams = autoencoder_residual()
-  hparams.bottleneck_size = 4096
+  hparams.bottleneck_bits = 4096
   hparams.bottleneck_noise = 0.1
   hparams.bottleneck_warmup_steps = 3000
   hparams.add_hparam("discretize_warmup_steps", 5000)
@@ -499,19 +495,28 @@ def autoencoder_residual_discrete_big():
 
 @registry.register_hparams
 def autoencoder_ordered_discrete():
-  """Basic autoencoder model."""
+  """Ordered discrete autoencoder model."""
   hparams = autoencoder_residual_discrete()
   hparams.bottleneck_noise = 1.0
   hparams.add_hparam("unordered", False)
   return hparams
 
 
+@registry.register_hparams
+def autoencoder_ordered_discrete_vq():
+  """Ordered discrete autoencoder model with VQ bottleneck."""
+  hparams = autoencoder_ordered_discrete()
+  hparams.bottleneck_kind = "vq"
+  hparams.bottleneck_bits = 16
+  return hparams
+
+
 @registry.register_hparams
 def autoencoder_discrete_pong():
   """Discrete autoencoder model for compressing pong frames."""
   hparams = autoencoder_ordered_discrete()
-  hparams.num_hidden_layers = 4
-  hparams.bottleneck_size = 24
+  hparams.num_hidden_layers = 2
+  hparams.bottleneck_bits = 24
   hparams.dropout = 0.1
   hparams.batch_size = 2
   hparams.bottleneck_noise = 0.2
@@ -520,6 +525,22 @@ def autoencoder_discrete_pong():
   return hparams
 
 
+@registry.register_hparams
+def autoencoder_discrete_cifar():
+  """Discrete autoencoder model for compressing cifar."""
+  hparams = autoencoder_ordered_discrete()
+  hparams.bottleneck_noise = 0.0
+  hparams.bottleneck_size = 90
+  hparams.unordered = True
+  hparams.num_hidden_layers = 2
+  hparams.hidden_size = 256
+  hparams.num_residual_layers = 4
+  hparams.batch_size = 32
+  hparams.learning_rate_constant = 2.0
+  hparams.dropout = 0.1
+  return hparams
+
+
 @registry.register_ranged_hparams
 def autoencoder_discrete_pong_range(rhp):
   """Narrow tuning grid."""
@@ -531,5 +552,5 @@ def autoencoder_discrete_pong_range(rhp):
 def autoencoder_stacked():
   """Stacked autoencoder model."""
   hparams = autoencoder_residual_discrete()
-  hparams.bottleneck_size = 128
+  hparams.bottleneck_bits = 128
   return hparams
diff --git a/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index 596e8abf0..fc45aeede 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -72,6 +72,11 @@ def testAutoencoderOrderedDiscrete(self):
     res = self.get_mnist_random_output("autoencoder_ordered_discrete")
     self.assertEqual(res.shape, self.mnist_output_shape)
 
+  def testAutoencoderOrderedDiscreteVQ(self):
+    res = self.get_mnist_random_output(
+        "autoencoder_ordered_discrete", "autoencoder_ordered_discrete_vq")
+    self.assertEqual(res.shape, self.mnist_output_shape)
+
   def testAutoencoderStacked(self):
     res = self.get_mnist_random_output("autoencoder_stacked")
     self.assertEqual(res.shape, self.mnist_output_shape)
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 63b3edcb8..0a9ac2ae0 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -198,9 +198,9 @@ def next_frame_ae():
   hparams = next_frame()
   hparams.input_modalities = "inputs:video:bitwise"
   hparams.hidden_size = 256
-  hparams.batch_size = 16
+  hparams.batch_size = 8
   hparams.num_hidden_layers = 4
-  hparams.num_compress_steps = 3
+  hparams.num_compress_steps = 4
   hparams.dropout = 0.4
   return hparams
 
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 09ee739ed..b7a21439e 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -263,7 +263,9 @@ def generator():
         except tf.errors.OutOfRangeError:
           break
 
-    generator_utils.generate_files(generator(), out_files)
+    generator_utils.generate_files(
+        generator(), out_files,
+        cycle_every_n=problem.total_number_of_frames // 10)
 
 
 def encode_env_frames(problem_name, ae_problem_name, autoencoder_path,
@@ -522,7 +524,7 @@ def rl_modelrl_tiny():
       true_env_generator_num_steps=100,
       model_train_steps=2,
       simulated_env_generator_num_steps=100,
-      ppo_epochs_num=2,
+      ppo_epochs_num=6,
       ppo_time_limit=20,
       ppo_epoch_length=20,
       ppo_num_agents=2,
@@ -600,7 +602,7 @@ def rl_modelrl_ae_base():
   hparams = rl_modelrl_base()
   hparams.ppo_params = "ppo_pong_ae_base"
   hparams.generative_model_params = "next_frame_ae"
-  hparams.autoencoder_train_steps = 30000
+  hparams.autoencoder_train_steps = 50000
   return hparams
 
 
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 323b2578b..2b446da18 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -81,6 +81,8 @@ def _add_variable_proxy_methods(var, proxy_tensor):
   """
   proxy_tensor.read_value = lambda: tf.identity(proxy_tensor)
   proxy_tensor.assign_sub = var.assign_sub
+  proxy_tensor.assign = var.assign
+  proxy_tensor.initialized_value = var.initialized_value
 
 
 class Parallelism(object):

From 051e579faf0f935e5fcdc96a516c002a6aa4f962 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 24 May 2018 13:50:58 -0700
Subject: [PATCH 0015/2720] internal

PiperOrigin-RevId: 197947305
---
 tensor2tensor/bin/t2t_datagen.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index c02897834..30fc5e50b 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -46,6 +46,7 @@
 
 import tensorflow as tf
 
+
 flags = tf.flags
 FLAGS = flags.FLAGS
 

From 979fc54b91903cb6ec6c825c8a1f8601f15ea25e Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 24 May 2018 15:49:42 -0700
Subject: [PATCH 0016/2720] Strip EOS and PAD in serving

PiperOrigin-RevId: 197967759
---
 tensor2tensor/data_generators/dna_encoder.py  |  4 +-
 tensor2tensor/data_generators/text_encoder.py | 42 +++++++++++++++----
 tensor2tensor/serving/serving_utils.py        |  2 +-
 3 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/data_generators/dna_encoder.py b/tensor2tensor/data_generators/dna_encoder.py
index 62c9b24e8..297a8975b 100644
--- a/tensor2tensor/data_generators/dna_encoder.py
+++ b/tensor2tensor/data_generators/dna_encoder.py
@@ -84,7 +84,7 @@ def encode(self, s):
       ids.append(self._tokens_to_ids[chunk])
     return ids
 
-  def decode(self, ids):
+  def decode(self, ids, strip_extraneous=False):
     bases = []
     for idx in ids:
       if idx >= self._num_reserved_ids:
@@ -92,6 +92,8 @@ def decode(self, ids):
         if self.PAD in chunk:
           chunk = chunk[:chunk.index(self.PAD)]
       else:
+        if strip_extraneous:
+          continue
         chunk = [text_encoder.RESERVED_TOKENS[idx]]
       bases.extend(chunk)
     return "".join(bases)
diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 53fb47ea8..2197bb867 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -115,17 +115,21 @@ def encode(self, s):
     """
     return [int(w) + self._num_reserved_ids for w in s.split()]
 
-  def decode(self, ids):
+  def decode(self, ids, strip_extraneous=False):
     """Transform a sequence of int ids into a human-readable string.
 
     EOS is not expected in ids.
 
     Args:
       ids: list of integers to be converted.
+      strip_extraneous: bool, whether to strip off extraneous tokens
+        (EOS and PAD).
 
     Returns:
       s: human-readable string.
     """
+    if strip_extraneous:
+      ids = strip_ids(ids, list(range(self._num_reserved_ids)))
     return " ".join(self.decode_list(ids))
 
   def decode_list(self, ids):
@@ -166,7 +170,9 @@ def encode(self, s):
     # Python3: explicitly convert to UTF-8
     return [c + numres for c in s.encode("utf-8")]
 
-  def decode(self, ids):
+  def decode(self, ids, strip_extraneous=False):
+    if strip_extraneous:
+      ids = strip_ids(ids, list(range(self._num_reserved_ids)))
     numres = self._num_reserved_ids
     decoded_ids = []
     int2byte = six.int2byte
@@ -216,7 +222,8 @@ def encode(self, s):
     label_str = s
     return self._class_labels.index(label_str)
 
-  def decode(self, ids):
+  def decode(self, ids, strip_extraneous=False):
+    del strip_extraneous
     label_id = ids
     if isinstance(label_id, list):
       assert len(label_id) == 1
@@ -254,7 +261,8 @@ def encode(self, label_str, on_value=1, off_value=0):  # pylint: disable=argumen
     e[self._class_labels.index(label_str)] = on_value
     return e.tolist()
 
-  def decode(self, ids):
+  def decode(self, ids, strip_extraneous=False):
+    del strip_extraneous
     label_id = ids
     if isinstance(label_id, np.ndarray):
       label_id = np.squeeze(label_id).astype(np.int8).tolist()
@@ -313,7 +321,7 @@ def encode(self, s):
     ret = [self._token_to_id[tok] for tok in tokens]
     return ret[::-1] if self._reverse else ret
 
-  def decode(self, ids):
+  def decode(self, ids, strip_extraneous=False):
     return " ".join(self.decode_list(ids))
 
   def decode_list(self, ids):
@@ -511,14 +519,19 @@ def encode_without_tokenizing(self, token_text):
     """
     return self._tokens_to_subtoken_ids([native_to_unicode(token_text)])
 
-  def decode(self, ids):
+  def decode(self, ids, strip_extraneous=False):
     """Converts a sequence of subtoken ids to a native string.
 
     Args:
       ids: a list of integers in the range [0, vocab_size)
+      strip_extraneous: bool, whether to strip off extraneous tokens
+        (EOS and PAD).
+
     Returns:
       a native string
     """
+    if strip_extraneous:
+      ids = strip_ids(ids, list(range(self._num_reserved_ids)))
     return unicode_to_native(
         tokenizer.decode(self._subtoken_ids_to_tokens(ids)))
 
@@ -956,11 +969,12 @@ def encode(self, s):
       raise NotImplementedError("Image reading not implemented.")
     return im.imread(s)
 
-  def decode(self, ids):
+  def decode(self, ids, strip_extraneous=False):
     """Transform a sequence of int ids into an image file.
 
     Args:
       ids: list of integers to be converted.
+      strip_extraneous: unused
 
     Returns:
       Path to the temporary file where the image was saved.
@@ -968,6 +982,7 @@ def decode(self, ids):
     Raises:
       ValueError: if the ids are not of the appropriate size.
     """
+    del strip_extraneous
     _, tmp_file_path = tempfile.mkstemp("_decode.png")
     if self._height is None or self._width is None:
       size = int(math.sqrt(len(ids) / self._channels))
@@ -1022,11 +1037,12 @@ def encode(self, s):
     """
     return [float(w) for w in s.split()]
 
-  def decode(self, ids):
+  def decode(self, ids, strip_extraneous=False):
     """Transform sequence of float values into string (float values).
 
     Args:
       ids: array of floats to be converted.
+      strip_extraneous: unused
 
     Returns:
       String having space separated float values.
@@ -1034,5 +1050,13 @@ def decode(self, ids):
     Raises:
       ValueError: if the ids are not of the appropriate size.
     """
+    del strip_extraneous
     return " ".join(ids)
-  
\ No newline at end of file
+
+
+def strip_ids(ids, ids_to_strip):
+  """Strip ids_to_strip from the end ids."""
+  ids = list(ids)
+  while ids[-1] in ids_to_strip:
+    ids.pop()
+  return ids
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index c20fe8b48..c7d44bf47 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -54,7 +54,7 @@ def _encode(inputs, encoder, add_eos=True):
 
 
 def _decode(output_ids, output_decoder):
-  return output_decoder.decode(output_ids)
+  return output_decoder.decode(output_ids, strip_extraneous=True)
 
 
 

From 51ba4798fe5959e28f33085002283296d70fd6ed Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 24 May 2018 16:26:40 -0700
Subject: [PATCH 0017/2720] Add cast_like(x, y) and update calls to tf.cast(x,
 y.dtype)

PiperOrigin-RevId: 197973639
---
 tensor2tensor/layers/common_attention.py      |  6 ++---
 .../layers/common_image_attention.py          |  2 +-
 tensor2tensor/layers/common_layers.py         | 23 +++++++++++++++----
 tensor2tensor/layers/modalities.py            |  2 +-
 tensor2tensor/utils/adafactor.py              |  3 ++-
 tensor2tensor/utils/optimize.py               |  9 ++++----
 6 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 4b1f114d8..42f2eb745 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -582,7 +582,7 @@ def add_timing_signal_1d_given_position(x,
           tf.expand_dims(inv_timescales, 0), 0))
   signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=2)
   signal = tf.pad(signal, [[0, 0], [0, 0], [0, tf.mod(channels, 2)]])
-  signal = tf.cast(signal, x.dtype)
+  signal = common_layers.cast_like(signal, x)
   return x + signal
 
 
@@ -1416,7 +1416,7 @@ def dot_product_attention(q,
     # [batch, num_heads, query_length, memory_length]
     logits = tf.matmul(q, k, transpose_b=True)
     if bias is not None:
-      bias = tf.cast(bias, logits.dtype)
+      bias = common_layers.cast_like(bias, logits)
       logits += bias
     weights = tf.nn.softmax(logits, name="attention_weights")
     if save_weights_to is not None:
@@ -1834,7 +1834,7 @@ def local(x, depth):
     good_part = common_layers.ones_matrix_band_part(block_length, local_length,
                                                     -1, block_length)
     mask = (1.0 - good_part) * -1e9
-    mask = tf.cast(mask, attention.dtype)
+    mask = common_layers.cast_like(mask, attention)
     attention += tf.reshape(mask, [1, 1, 1, block_length, local_length])
     attention = tf.nn.softmax(attention)
     # TODO(noam): figure out how to show a summary for the remaining blocks.
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 11c12d5d6..a34841a88 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -608,7 +608,7 @@ def prepare_decoder(targets, hparams):
     x = tf.reshape(x, [targets_shape[0],
                        x_shape[1], x_shape[2], hparams.hidden_size])
     x = add_pos_signals(x, hparams, "dec_pos")
-  x = tf.cast(x, targets.dtype)
+  x = common_layers.cast_like(x, targets)
   return x, x_shape[1], x_shape[2]
 
 
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 6a78e1136..54fd9623c 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -293,8 +293,8 @@ def dropout_no_scaling(x, keep_prob):
   """
   if keep_prob == 1.0:
     return x
-  return x * tf.cast(
-      tf.less(tf.random_uniform(tf.shape(x)), keep_prob), x.dtype)
+  mask = tf.less(tf.random_uniform(tf.shape(x)), keep_prob)
+  return x * cast_like(mask, x)
 
 
 def embedding(x,
@@ -599,7 +599,7 @@ def layer_norm_vars(filters):
 
 def layer_norm_compute_python(x, epsilon, scale, bias):
   """Layer norm raw computation."""
-  epsilon, scale, bias = [tf.cast(t, x.dtype) for t in [epsilon, scale, bias]]
+  epsilon, scale, bias = [cast_like(t, x) for t in [epsilon, scale, bias]]
   mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
   variance = tf.reduce_mean(tf.square(x - mean), axis=[-1], keepdims=True)
   norm_x = (x - mean) * tf.rsqrt(variance + epsilon)
@@ -651,7 +651,7 @@ def group_norm(x, filters=None, num_groups=8, epsilon=1e-5):
       "group_norm_scale", [filters], initializer=tf.ones_initializer())
   bias = tf.get_variable(
       "group_norm_bias", [filters], initializer=tf.zeros_initializer())
-  epsilon, scale, bias = [tf.cast(t, x.dtype) for t in [epsilon, scale, bias]]
+  epsilon, scale, bias = [cast_like(t, x) for t in [epsilon, scale, bias]]
   # Reshape and compute group norm.
   x = tf.reshape(x, x_shape[:-1] + [num_groups, filters // num_groups])
   # Calculate mean and variance on heights, width, channels (not groups).
@@ -3081,3 +3081,18 @@ def time_to_channels(embedded_video):
   return tf.reshape(transposed,
                     [video_shape[0], video_shape[2], video_shape[3],
                      video_shape[1] * video_shape[4]])
+
+
+def cast_like(x, y):
+  """Cast x to y's dtype, if necessary."""
+  x = tf.convert_to_tensor(x)
+  y = tf.convert_to_tensor(y)
+
+  if x.dtype.base_dtype == y.dtype.base_dtype:
+    return x
+
+  cast_x = tf.cast(x, y.dtype)
+  if cast_x.device != x.device:
+    tf.logging.warning("Cast for %s may induce copy from '%s' to '%s'",
+                       x.name, x.device, cast_x.device)
+  return cast_x
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index b94ff92bc..0c1bba2a9 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -566,7 +566,7 @@ def bottom(self, x):
       return tf.layers.dense(transposed, self._body_input_depth,
                              name="merge_pixel_embedded_frames")
 
-  def targets_bottom(self, x):
+  def targets_bottom(self, x):  # pylint: disable=arguments-differ
     inputs = x
     with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
       common_layers.summarize_video(inputs, "targets_bottom")
diff --git a/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index 1f65f5905..4f627675f 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -17,6 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import quantization
 
 import tensorflow as tf
@@ -255,7 +256,7 @@ def _resource_apply_dense(self, grad, handle):
       m = self.get_slot(var, "m")
       new_m = self._beta1 * tf.to_float(m) + (1.0 - self._beta1) * subtrahend
       subtrahend = new_m
-      new_m = tf.cast(new_m, var.dtype)
+      new_m = common_layers.cast_like(new_m, var)
       updates.append(tf.assign(m, new_m, use_locking=self._use_locking))
     new_val = tf.to_float(old_val) - subtrahend
     if var.dtype.base_dtype == tf.bfloat16:
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index dcb7b91b1..b973f9ed3 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -18,6 +18,7 @@
 from __future__ import print_function
 import numpy as np
 
+from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import adafactor
 from tensor2tensor.utils import yellowfin
 
@@ -105,11 +106,9 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
   def compute_gradients(self, loss, var_list=None, **kwargs):  # pylint: disable=arguments-differ
     gradients = self._opt.compute_gradients(loss, var_list, **kwargs)
     def cast_grad(g, v):
-      if v is None or g is None:
-        return (g, v)
-      if g.dtype == v.dtype:
-        return (g, v)
-      return (tf.cast(g, v.dtype), v)
+      if v is not None and g is not None:
+        g = common_layers.cast_like(g, v)
+      return (g, v)
     gradients = [cast_grad(g, v) for g, v in gradients]
     return gradients
 

From c4c7cf87b08c216b2d6a573ba8cc104ae531e822 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 24 May 2018 16:43:48 -0700
Subject: [PATCH 0018/2720] Use resource variables and set caching_devices to
 None.

PiperOrigin-RevId: 197976097
---
 tensor2tensor/layers/discretization.py |  4 +--
 tensor2tensor/utils/devices.py         |  3 +-
 tensor2tensor/utils/expert_utils.py    |  6 ++--
 tensor2tensor/utils/t2t_model.py       | 45 +++++++++++++-------------
 4 files changed, 31 insertions(+), 27 deletions(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 02977ac19..a09d0be95 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -734,8 +734,8 @@ def vq_discrete_bottleneck(x,
 
   dw = tf.matmul(x_means_hot, x, transpose_a=True)
 
-  updated_ema_means = moving_averages.assign_moving_average(
-      ema_means, dw, decay, zero_debias=False)
+  updated_ema_means = tf.identity(moving_averages.assign_moving_average(
+      ema_means, dw, decay, zero_debias=False))
   n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
   updated_ema_count = (
       (updated_ema_count + epsilon) / (n + bottleneck_size * epsilon) * n)
diff --git a/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index 12e6a3d9a..061be8f81 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -162,7 +162,8 @@ def _replica_device_setter(worker_device):
           _replica_device_setter(worker_job + "/GPU:%d" % d)
           for d in _gpu_order(worker_gpu)
       ]
-      caching_devices = [worker_job + "/GPU:0"] * worker_gpu
+      # caching_devices = [worker_job + "/GPU:0"] * worker_gpu
+      caching_devices = None
     else:
       datashard_devices = [_replica_device_setter(worker_job)]
       caching_devices = None
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 2b446da18..86747480b 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -181,7 +181,8 @@ def daisy_chain_getter(getter, name, *args, **kwargs):
           v = tf.identity(last_device_v)
         else:
           var = getter(name, *args, **kwargs)
-          v = tf.identity(var._ref())  # pylint: disable=protected-access
+          # v = tf.identity(var._ref())  # pylint: disable=protected-access
+          v = var.read_value()
 
         # keep track of the original variable
         tensor_to_var[v] = var
@@ -202,7 +203,8 @@ def caching_getter(getter, name, *args, **kwargs):
 
         v = getter(name, *args, **kwargs)
         with tf.device(self._caching_devices[i]):
-          ret = tf.identity(v._ref())  # pylint: disable=protected-access
+          # ret = tf.identity(v._ref())  # pylint: disable=protected-access
+          ret = v.read_value()
         _add_variable_proxy_methods(v, ret)
         cache[key] = ret
         return ret
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index b989e9470..36fd72e8d 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -229,29 +229,30 @@ def model_fn_sharded(self, sharded_features):
     return sharded_logits, losses
 
   def model_fn(self, features):
-    transformed_features = self.bottom(features)
-
-    if self.hparams.activation_dtype == "bfloat16":
-      for k, v in sorted(six.iteritems(transformed_features)):
-        if v.dtype == tf.float32:
-          transformed_features[k] = tf.cast(v, tf.bfloat16)
-
-    with tf.variable_scope("body"):
-      log_info("Building model body")
-      body_out = self.body(transformed_features)
-    output, losses = self._normalize_body_output(body_out)
-
-    if "training" in losses:
-      log_info("Skipping T2TModel top and loss because training loss "
-               "returned from body")
-      logits = output
-    else:
-      logits = self.top(output, features)
-      losses["training"] = 0.0
-      if self._hparams.mode != tf.estimator.ModeKeys.PREDICT:
-        losses["training"] = self.loss(logits, features)
+    with tf.variable_scope(tf.get_variable_scope(), use_resource=True):
+      transformed_features = self.bottom(features)
+
+      if self.hparams.activation_dtype == "bfloat16":
+        for k, v in sorted(six.iteritems(transformed_features)):
+          if v.dtype == tf.float32:
+            transformed_features[k] = tf.cast(v, tf.bfloat16)
+
+      with tf.variable_scope("body"):
+        log_info("Building model body")
+        body_out = self.body(transformed_features)
+      output, losses = self._normalize_body_output(body_out)
+
+      if "training" in losses:
+        log_info("Skipping T2TModel top and loss because training loss "
+                 "returned from body")
+        logits = output
+      else:
+        logits = self.top(output, features)
+        losses["training"] = 0.0
+        if self._hparams.mode != tf.estimator.ModeKeys.PREDICT:
+          losses["training"] = self.loss(logits, features)
 
-    return logits, losses
+      return logits, losses
 
   def bottom(self, features):
     """Transform features to feed into body."""

From e317065088f6555d66d7fb6bdc85a3a4a2063c7d Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 24 May 2018 18:43:57 -0700
Subject: [PATCH 0019/2720] Fix Travis config

PiperOrigin-RevId: 197988958
---
 .travis.yml                                   | 23 ++++++++++++-------
 .../data_generators/video_generated.py        |  6 ++++-
 tensor2tensor/models/transformer_test.py      |  7 ++++--
 3 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 08901475c..04da5d004 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -43,18 +43,25 @@ script:
   - python -c "from tensor2tensor.models import transformer; print(transformer.Transformer.__name__)"
 
   # Lint
-  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST*"  ]]; then
-        pylint tensor2tensor;
+  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
+        pylint -j 2 tensor2tensor;
     fi
 
   # Run tests
+  # Ignores:
+  # Tested separately:
+  #   * registry_test
+  #   * trainer_lib_test
+  #   * visualization_test
+  # algorithmic_math_test: flaky
+  # r_transformer_test: requires new feature in tf.foldl (rm with TF 1.9)
   - pytest
-    --ignore=tensor2tensor/utils/registry_test.py  # tested separately
-    --ignore=tensor2tensor/utils/trainer_lib_test.py  # tested separately
-    --ignore=tensor2tensor/visualization/visualization_test.py  # tested separately
-    --ignore=tensor2tensor/bin/t2t_trainer_test.py  # tested separately
-    --ignore=tensor2tensor/data_generators/algorithmic_math_test.py  # flaky
-    --ignore=tensor2tensor/models/research/r_transformer_test.py  # Requires new feature in tf.foldl (rm with TF 1.9)
+    --ignore=tensor2tensor/utils/registry_test.py
+    --ignore=tensor2tensor/utils/trainer_lib_test.py
+    --ignore=tensor2tensor/visualization/visualization_test.py
+    --ignore=tensor2tensor/bin/t2t_trainer_test.py
+    --ignore=tensor2tensor/data_generators/algorithmic_math_test.py
+    --ignore=tensor2tensor/models/research/r_transformer_test.py
   - pytest tensor2tensor/utils/registry_test.py
   - pytest tensor2tensor/utils/trainer_lib_test.py
   - pytest tensor2tensor/visualization/visualization_test.py
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 863f211ac..14f85962f 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -20,7 +20,6 @@
 
 import math
 
-import matplotlib.pyplot as plt
 import numpy as np
 
 from tensor2tensor.data_generators import video_utils
@@ -28,6 +27,11 @@
 
 import tensorflow as tf
 
+try:
+  import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top
+except ImportError:
+  pass
+
 
 @registry.register_problem
 class VideoStochasticShapes10k(video_utils.VideoProblem):
diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 5da51fecf..d3d1ef686 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -157,9 +157,12 @@ def testBeamDecodeWithRelativeAttention(self):
 
     with self.test_session():
       tf.global_variables_initializer().run()
-      beam_res = beam_result.eval()
+      beam_result.eval()
 
-    self.assertEqual(beam_res.shape, (BATCH_SIZE, INPUT_LENGTH + decode_length))
+    # TODO(petershaw): This test is flaky because the decode may hit EOS before
+    # getting to the expected length.
+    # self.assertEqual(beam_res.shape,
+    #                  (BATCH_SIZE, INPUT_LENGTH + decode_length))
 
   def testBeamVsFast(self):
     model, features = get_model(transformer.transformer_small())

From a8b6ea9619ffbdf8b0c6879371868c6a8d731c5c Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 24 May 2018 20:38:36 -0700
Subject: [PATCH 0020/2720] Corrections to interleaving to make autoencoder RL
 work better.

PiperOrigin-RevId: 197996798
---
 tensor2tensor/data_generators/problem.py         | 16 ++++++++++++----
 tensor2tensor/data_generators/video_generated.py |  6 +++---
 tensor2tensor/models/research/autoencoders.py    |  5 +++++
 tensor2tensor/models/research/next_frame.py      |  6 +++++-
 4 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 83542c8c9..b8731273a 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -564,11 +564,19 @@ def _load_records_and_preprocess(filename):
         "partition: %d num_data_files: %d" % (partition_id, len(data_files)))
     if shuffle_files:
       random.shuffle(data_files)
-    dataset = tf.data.Dataset.from_tensor_slices(tf.constant(data_files))
 
-    dataset = dataset.apply(
-        tf.contrib.data.parallel_interleave(
-            _load_records_and_preprocess, sloppy=is_training, cycle_length=8))
+    # Create data-set from files by parsing, pre-processing and interleaving.
+    if shuffle_files:
+      dataset = tf.data.Dataset.from_tensor_slices(tf.constant(data_files))
+      dataset = dataset.apply(
+          tf.contrib.data.parallel_interleave(
+              _load_records_and_preprocess, sloppy=True, cycle_length=8))
+    else:
+      dataset = None
+      for f in data_files:
+        f_data = _load_records_and_preprocess(f)
+        dataset = f_data if dataset is None else dataset.concatenate(f_data)
+
     dataset = dataset.map(
         self.maybe_reverse_and_copy, num_parallel_calls=num_threads)
     dataset = dataset.take(max_records)
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 14f85962f..014d8adc4 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -176,14 +176,14 @@ def generate_stochastic_shape_instance(self):
       yield image
 
   def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
-    frame_number = 0
+    counter = 0
     done = False
     while not done:
       for frame_number, frame in enumerate(
           self.generate_stochastic_shape_instance()):
-        if frame_number >= self.total_number_of_frames:
+        if counter >= self.total_number_of_frames:
           done = True
           break
 
         yield {"frame": frame, "frame_number": [frame_number]}
-        frame_number += 1
+        counter += 1
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 891787b96..5930e09e1 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -167,6 +167,11 @@ def encoder(self, x):
       residual_conv = tf.layers.conv2d
       if hparams.residual_use_separable_conv:
         residual_conv = tf.layers.separable_conv2d
+      # Input embedding with a non-zero bias for uniform inputs.
+      x = tf.layers.dense(
+          x, hparams.hidden_size, name="embed", activation=common_layers.belu,
+          bias_initializer=tf.random_normal_initializer(stddev=0.01))
+      x = common_attention.add_timing_signal_nd(x)
       # Down-convolutions.
       for i in range(hparams.num_hidden_layers):
         with tf.variable_scope("layer_%d" % i):
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 0a9ac2ae0..ebecefc3b 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -58,7 +58,11 @@ def body(self, features):
 
     # Embed the inputs.
     inputs_shape = common_layers.shape_list(features["inputs"])
-    x = tf.layers.dense(features["inputs"], filters, name="inputs_embed")
+    # Using non-zero bias initializer below for edge cases of uniform inputs.
+    x = tf.layers.dense(
+        features["inputs"], filters, name="inputs_embed",
+        bias_initializer=tf.random_normal_initializer(stddev=0.01))
+    x = common_attention.add_timing_signal_nd(x)
 
     # Down-stride.
     layer_inputs = [x]

From df7a2eaa9ba0642ff91aca34dcdeb479e5676f1c Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 24 May 2018 23:49:02 -0700
Subject: [PATCH 0021/2720] Allow multiple time-series in TimeseriesProblem.

PiperOrigin-RevId: 198009162
---
 tensor2tensor/data_generators/timeseries.py   | 56 +++++++++++--------
 .../data_generators/timeseries_test.py        |  4 +-
 2 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index 07e32d39a..6c670b45c 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -28,25 +28,27 @@
 
 
 @registry.register_problem
-class TimeSeriesToyProblem(problem.Problem):
+class TimeseriesToyProblem(problem.Problem):
   """Base Problem for multi timeseries for datasets."""
 
   def __init__(self,
                was_reversed=False,
                was_copy=False,
+               num_series=2,
                num_train_shards=9,
                num_eval_shards=1,
                num_samples=100):
-    super(TimeSeriesToyProblem, self).__init__(was_reversed, was_copy)
+    super(TimeseriesToyProblem, self).__init__(was_reversed, was_copy)
     self._num_train_shards = num_train_shards
     self._num_eval_shards = num_eval_shards
     self._num_samples = num_samples
+    self._num_series = num_series
 
   def feature_encoders(self, data_dir):
     del data_dir
     return {
-        'inputs': text_encoder.RealEncoder(),
-        'targets': text_encoder.RealEncoder()
+        "inputs": text_encoder.RealEncoder(),
+        "targets": text_encoder.RealEncoder()
     }
 
   @property
@@ -59,39 +61,49 @@ def dataset_splits(self):
     """Splits of data to produce and number of output shards for each."""
     # 10% evaluation data
     return [{
-        'split': problem.DatasetSplit.TRAIN,
-        'shards': self._num_train_shards,
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": self._num_train_shards,
     }, {
-        'split': problem.DatasetSplit.EVAL,
-        'shards': self._num_eval_shards,
+        "split": problem.DatasetSplit.EVAL,
+        "shards": self._num_eval_shards,
     }]
 
   def eval_metrics(self):
     eval_metrics = [metrics.Metrics.RMSE]
-
     return eval_metrics
 
+  def preprocess_example(self, example, unused_mode, unused_hparams):
+    # Time series are flat on disk, we un-flatten them back here.
+    flat_inputs = example["inputs"]
+    flat_targets = example["targets"]
+    example["inputs"] = tf.reshape(flat_inputs, [-1, self._num_series])
+    example["targets"] = tf.reshape(flat_targets, [-1, self._num_series])
+    return example
+
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
     del data_dir
     del tmp_dir
     del dataset_split
 
-    series_1 = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
+    series = [[float(i + n) for n in range(self._num_series)]
+              for i in range(10)]
 
-    # This generates _num_samples instances of each possible split of series_1;
+    # This generates _num_samples instances of each possible split of series;
     # inputs & targets are of variable size.
-    for x in range(self._num_samples):
+    for _ in range(self._num_samples):
       split_index = random.randint(1, 9)
-      inputs, targets = series_1[:split_index], series_1[split_index:]
-      example_keys = ['inputs', 'targets']
-      ex_dict = dict(zip(example_keys, [inputs, targets]))
-      print('Inputs & Targets #', x, ':', ex_dict)
+      inputs, targets = series[:split_index], series[split_index:]
+      # We need to flatten the lists on disk for tf,Example to work.
+      flat_inputs = [item for sublist in inputs for item in sublist]
+      flat_targets = [item for sublist in targets for item in sublist]
+      example_keys = ["inputs", "targets"]
+      ex_dict = dict(zip(example_keys, [flat_inputs, flat_targets]))
       yield ex_dict
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {'inputs': (registry.Modalities.REAL, 1)}
-    p.target_modality = (registry.Modalities.REAL, 1)
+    p.input_modality = {"inputs": (registry.Modalities.REAL, self._num_series)}
+    p.target_modality = (registry.Modalities.REAL, self._num_series)
     p.input_space_id = problem.SpaceID.REAL
     p.target_space_id = problem.SpaceID.REAL
 
@@ -102,8 +114,8 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
         problem.DatasetSplit.TEST: self.test_filepaths,
     }
 
-    split_paths = [(split['split'], filepath_fns[split['split']](
-        data_dir, split['shards'], shuffled=False))
+    split_paths = [(split["split"], filepath_fns[split["split"]](
+        data_dir, split["shards"], shuffled=False))
                    for split in self.dataset_splits]
 
     all_paths = []
@@ -123,8 +135,8 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
   def example_reading_spec(self):
     data_fields = {
-        'inputs': tf.VarLenFeature(tf.float32),
-        'targets': tf.VarLenFeature(tf.float32),
+        "inputs": tf.VarLenFeature(tf.float32),
+        "targets": tf.VarLenFeature(tf.float32),
     }
     data_items_to_decoders = None
     return (data_fields, data_items_to_decoders)
diff --git a/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index 716721c48..a45c80895 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -34,8 +34,8 @@ def setUpClass(cls):
     shutil.rmtree(cls.tmp_dir)
     os.mkdir(cls.tmp_dir)
 
-  def testTimeSeriesToyProblem(self):
-    problem = timeseries.TimeSeriesToyProblem(
+  def testTimeseriesToyProblem(self):
+    problem = timeseries.TimeseriesToyProblem(
         num_train_shards=1, num_eval_shards=1, num_samples=4)
     problem.generate_data(self.tmp_dir, self.tmp_dir)
 

From ef6373e73c58e83b3a6b3b86393e1ca799909dd3 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 25 May 2018 10:10:27 -0700
Subject: [PATCH 0022/2720] Place for stochastic next-frame model.

PiperOrigin-RevId: 198064797
---
 tensor2tensor/models/research/next_frame.py | 22 +++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index ebecefc3b..452a66ee0 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -167,6 +167,28 @@ def logits_to_samples(logits):
     return results
 
 
+@registry.register_model
+class NextFrameStochastic(NextFrameBasic):
+  """Stochastic next-frame model."""
+
+  def body(self, features):
+    hparams = self.hparams
+    filters = hparams.hidden_size
+
+    # Split inputs time-wise into a list of frames. Inputs are by default
+    # concatenated time-wise on channels in VideoModality, so we split on
+    # the last axis. Can do the same for target frames with num_target_frames.
+    # TODO(lukaszkaiser): should we change VideoModality to not concatenate?
+    num_frames = hparams.problem.num_input_frames
+    input_frames = tf.split(features["inputs"], num_frames, axis=-1)
+
+    # For now predict using just a linear transformation of the last frame.
+    # Here input_frames[-1] is contrast-normalized last frame.
+    prediction = tf.layers.dense(input_frames[-1], filters,
+                                 name="final_dense")
+    return prediction
+
+
 @registry.register_hparams
 def next_frame():
   """Basic 2-frame conv model."""

From aed7a4c46e24e33ec8a8eb0e1bc58edba8234205 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 25 May 2018 11:44:32 -0700
Subject: [PATCH 0023/2720] Add a sorting task.

PiperOrigin-RevId: 198080173
---
 tensor2tensor/data_generators/algorithmic.py  | 38 +++++++++++++++++++
 .../data_generators/algorithmic_test.py       |  8 ++++
 2 files changed, 46 insertions(+)

diff --git a/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index 3efd2381f..93f3b27a9 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -445,3 +445,41 @@ def dev_size(self):
   @property
   def num_shards(self):
     return 1
+
+
+@registry.register_problem
+class AlgorithmicSortProblem(AlgorithmicProblem):
+  """Problem spec for sorting numbers."""
+
+  @property
+  def num_symbols(self):
+    return 10
+
+  @property
+  def train_length(self):
+    return 10
+
+  @property
+  def dev_length(self):
+    return 10
+
+  def generator(self, nbr_symbols, max_length, nbr_cases):
+    """Generating for sorting task on sequence of symbols.
+
+    The length of the sequence is drawn uniformly at random from [1, max_length]
+    and then symbols are drawn uniformly at random from [0, nbr_symbols) until
+    nbr_cases sequences have been produced.
+
+    Args:
+      nbr_symbols: number of symbols to use in each sequence.
+      max_length: integer, maximum length of sequences to generate.
+      nbr_cases: the number of cases to generate.
+
+    Yields:
+      A dictionary {"inputs": input-list, "targets": target-list} where
+      target-list is input-list sorted.
+    """
+    for _ in range(nbr_cases):
+      l = np.random.randint(max_length) + 1
+      inputs = list(np.random.randint(nbr_symbols, size=l))
+      yield {"inputs": inputs, "targets": list(sorted(inputs))}
diff --git a/tensor2tensor/data_generators/algorithmic_test.py b/tensor2tensor/data_generators/algorithmic_test.py
index 234c5a7b7..548cc3f21 100644
--- a/tensor2tensor/data_generators/algorithmic_test.py
+++ b/tensor2tensor/data_generators/algorithmic_test.py
@@ -97,6 +97,14 @@ def testMultiplicationGenerator(self):
       self.assertEqual(d["targets"].count(5), 0)
     self.assertEqual(counter, 10)
 
+  def testSortGenerator(self):
+    sort_problem = algorithmic.AlgorithmicSortProblem()
+    counter = 0
+    for d in sort_problem.generator(10, 10, 10):
+      counter += 1
+      self.assertEqual(list(sorted(d["inputs"])), d["targets"])
+    self.assertEqual(counter, 10)
+
 
 if __name__ == "__main__":
   tf.test.main()

From 4f5fe245b6227894bad701b4cd6ab9b891d14f39 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 25 May 2018 23:02:17 -0700
Subject: [PATCH 0024/2720] Add model_rl_experiment_test

PiperOrigin-RevId: 198144637
---
 .travis.yml                                  | 16 ++++++---
 tensor2tensor/models/research/next_frame.py  | 11 ++++++
 tensor2tensor/rl/model_rl_experiment.py      | 22 ++++++------
 tensor2tensor/rl/model_rl_experiment_test.py | 38 ++++++++++++++++++++
 tensor2tensor/rl/rl_trainer_lib.py           |  2 +-
 5 files changed, 72 insertions(+), 17 deletions(-)
 create mode 100644 tensor2tensor/rl/model_rl_experiment_test.py

diff --git a/.travis.yml b/.travis.yml
index 04da5d004..e3b056bb4 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -42,17 +42,13 @@ script:
   # Check import
   - python -c "from tensor2tensor.models import transformer; print(transformer.Transformer.__name__)"
 
-  # Lint
-  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
-        pylint -j 2 tensor2tensor;
-    fi
-
   # Run tests
   # Ignores:
   # Tested separately:
   #   * registry_test
   #   * trainer_lib_test
   #   * visualization_test
+  #   * model_rl_experiment_test
   # algorithmic_math_test: flaky
   # r_transformer_test: requires new feature in tf.foldl (rm with TF 1.9)
   - pytest
@@ -62,6 +58,7 @@ script:
     --ignore=tensor2tensor/bin/t2t_trainer_test.py
     --ignore=tensor2tensor/data_generators/algorithmic_math_test.py
     --ignore=tensor2tensor/models/research/r_transformer_test.py
+    --ignore=tensor2tensor/rl/model_rl_experiment_test.py
   - pytest tensor2tensor/utils/registry_test.py
   - pytest tensor2tensor/utils/trainer_lib_test.py
   - pytest tensor2tensor/visualization/visualization_test.py
@@ -78,6 +75,15 @@ script:
   - t2t-trainer --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR --model=transformer --hparams_set=transformer_tiny --train_steps=5 --eval_steps=5 --output_dir=$T2T_TRAIN_DIR
   - t2t-decoder --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR --model=transformer --hparams_set=transformer_tiny --output_dir=$T2T_TRAIN_DIR --decode_hparams='num_samples=10'
 
+  # Do some things only on Python 2 and the latest TF version
+  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
+        pylint -j 2 tensor2tensor;
+    fi
+  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
+        pytest tensor2tensor/rl/model_rl_experiment_test.py;
+    fi
+
+
   # Export and query (on Python 2 only)
   # Bug: https://github.com/tensorflow/serving/issues/819
   #- if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "1.6.*"  ]]; then
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 452a66ee0..39e3804b4 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -239,6 +239,17 @@ def next_frame_small():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_tiny():
+  """Tiny for testing."""
+  hparams = next_frame()
+  hparams.hidden_size = 32
+  hparams.num_hidden_layers = 1
+  hparams.num_compress_steps = 2
+  hparams.filter_double_steps = 1
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_l1():
   """Basic conv model with L1 modality."""
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index b7a21439e..89ed60674 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -519,17 +519,17 @@ def rl_modelrl_short():
 @registry.register_hparams
 def rl_modelrl_tiny():
   """Tiny set for testing."""
-  tiny_hp = tf.contrib.training.HParams(
-      epochs=2,
-      true_env_generator_num_steps=100,
-      model_train_steps=2,
-      simulated_env_generator_num_steps=100,
-      ppo_epochs_num=6,
-      ppo_time_limit=20,
-      ppo_epoch_length=20,
-      ppo_num_agents=2,
-  )
-  return rl_modelrl_base().override_from_dict(tiny_hp.values())
+  return rl_modelrl_base().override_from_dict(
+      tf.contrib.training.HParams(
+          epochs=2,
+          true_env_generator_num_steps=100,
+          simulated_env_generator_num_steps=100,
+          model_train_steps=2,
+          ppo_epochs_num=2,
+          ppo_time_limit=5,
+          ppo_epoch_length=5,
+          ppo_num_agents=2,
+      ).values())
 
 
 @registry.register_hparams
diff --git a/tensor2tensor/rl/model_rl_experiment_test.py b/tensor2tensor/rl/model_rl_experiment_test.py
new file mode 100644
index 000000000..2fd327260
--- /dev/null
+++ b/tensor2tensor/rl/model_rl_experiment_test.py
@@ -0,0 +1,38 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Tiny run of model_rl_experiment. Smoke test."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.rl import model_rl_experiment
+
+import tensorflow as tf
+
+FLAGS = tf.flags.FLAGS
+
+
+class ModelRLExperimentTest(tf.test.TestCase):
+
+  def test_run(self):
+    FLAGS.output_dir = tf.test.get_temp_dir()
+    FLAGS.loop_hparams_set = "rl_modelrl_tiny"
+    FLAGS.loop_hparams = "generative_model_params=next_frame_tiny"
+    FLAGS.schedule = "train"  # skip evaluation for world model training
+    model_rl_experiment.main(None)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index f039b9eb0..ffb64dc7b 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -112,7 +112,7 @@ def train(hparams, environment_spec, event_dir=None, model_dir=None,
             model_dir, model_saver, sess)
 
       # Fail-friendly, don't train if already trained for this epoch
-      if start_step >= ((hparams.epochs_num * (epoch+1)) - 5):
+      if start_step >= ((hparams.epochs_num * (epoch + 1))):
         tf.logging.info("Skipping PPO training for epoch %d as train steps "
                         "(%d) already reached", epoch, start_step)
         return

From f385c89d5617428e522c217882f60d1264f18566 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 25 May 2018 23:31:02 -0700
Subject: [PATCH 0025/2720] Add Travis test for hello_t2t.ipynb

PiperOrigin-RevId: 198145973
---
 .travis.yml                             | 4 +++-
 setup.py                                | 2 +-
 tensor2tensor/notebooks/hello_t2t.ipynb | 8 +++++---
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index e3b056bb4..83e1f087d 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -82,7 +82,9 @@ script:
   - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
         pytest tensor2tensor/rl/model_rl_experiment_test.py;
     fi
-
+  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
+        jupyter nbconvert --to notebook --execute tensor2tensor/notebooks/hello_t2t.ipynb;
+    fi
 
   # Export and query (on Python 2 only)
   # Bug: https://github.com/tensorflow/serving/issues/819
diff --git a/setup.py b/setup.py
index 95b91a2af..59b663f9e 100644
--- a/setup.py
+++ b/setup.py
@@ -49,7 +49,7 @@
     extras_require={
         'tensorflow': ['tensorflow>=1.5.0'],
         'tensorflow_gpu': ['tensorflow-gpu>=1.5.0'],
-        'tests': ['pytest', 'mock', 'pylint'],
+        'tests': ['pytest', 'mock', 'pylint', 'jupyter'],
     },
     classifiers=[
         'Development Status :: 4 - Beta',
diff --git a/tensor2tensor/notebooks/hello_t2t.ipynb b/tensor2tensor/notebooks/hello_t2t.ipynb
index b7661fcfc..e28344d20 100644
--- a/tensor2tensor/notebooks/hello_t2t.ipynb
+++ b/tensor2tensor/notebooks/hello_t2t.ipynb
@@ -75,7 +75,8 @@
       "cell_type": "code",
       "source": [
         "# Install deps\n",
-        "!pip install -q -U tensor2tensor tensorflow"
+        "!pip install -q -U tensor2tensor\n",
+        "!pip install -q tensorflow matplotlib\n"
       ],
       "execution_count": 0,
       "outputs": []
@@ -109,7 +110,7 @@
         "from tensor2tensor.utils import metrics\n",
         "\n",
         "# Enable TF Eager execution\n",
-        "from tensorflow.contrib.eager.python import tfe\n",
+        "tfe = tf.contrib.eager\n",
         "tfe.enable_eager_execution()\n",
         "\n",
         "# Other setup\n",
@@ -422,7 +423,8 @@
         "\n",
         "# Copy the vocab file locally so we can encode inputs and decode model outputs\n",
         "# All vocabs are stored on GCS\n",
-        "vocab_file = os.path.join(gs_data_dir, \"vocab.ende.32768\")\n",
+        "vocab_name = \"vocab.ende.32768\"\n",
+        "vocab_file = os.path.join(gs_data_dir, vocab_name)\n",
         "!gsutil cp {vocab_file} {data_dir}\n",
         "\n",
         "# Get the encoders from the problem\n",

From 4e218e4d18e1d74c26ba4cab52251bb570682480 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 25 May 2018 23:35:17 -0700
Subject: [PATCH 0026/2720] Add gym[atari] to tests deps

PiperOrigin-RevId: 198146149
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 59b663f9e..72be5a0c9 100644
--- a/setup.py
+++ b/setup.py
@@ -49,7 +49,7 @@
     extras_require={
         'tensorflow': ['tensorflow>=1.5.0'],
         'tensorflow_gpu': ['tensorflow-gpu>=1.5.0'],
-        'tests': ['pytest', 'mock', 'pylint', 'jupyter'],
+        'tests': ['pytest', 'mock', 'pylint', 'jupyter', 'gym[atari]'],
     },
     classifiers=[
         'Development Status :: 4 - Beta',

From 0fed09cf762525f164eeecd5b4a7759391f77f79 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 25 May 2018 23:49:26 -0700
Subject: [PATCH 0027/2720] Update ISSUE_TEMPLATE

---
 ISSUE_TEMPLATE.md | 31 ++++++++++++++++++-------------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/ISSUE_TEMPLATE.md b/ISSUE_TEMPLATE.md
index 477fa82fc..c5ae7c6a6 100644
--- a/ISSUE_TEMPLATE.md
+++ b/ISSUE_TEMPLATE.md
@@ -1,22 +1,27 @@
-<!-- **Note** Please tag your issue with *bug*, *feature request* or *help* -->
-
 ### Description
 
-> …
-
-### *TensorFlow* and *tensor2tensor* versions
+...
 
-<!-- **Note** Run `pip freeze | grep tensor` to get versions -->
+### Environment information
 
-> …
+```
+OS: <your answer here>
 
-### In case of bug report: Steps to reproduce the problem
+$ pip freeze | grep tensor
+# your output here
 
-> …
+$ python -V
+# your output here
+```
 
-### In case of bug report: Error log
+### For bugs: reproduction and error logs
 
-<!-- Please use code markdown (```) to format output messages. -->
-<!-- See https://help.github.com/articles/creating-and-highlighting-code-blocks/ -->
+```
+# Steps to reproduce:
+...
+```
 
-> …
+```
+# Error logs:
+...
+```

From 2bd2ddddc398a13d127011de223493b4271a9614 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Sat, 26 May 2018 01:17:43 -0700
Subject: [PATCH 0028/2720] Add gsutil to [tests]

PiperOrigin-RevId: 198151553
---
 .travis.yml |  6 +++++-
 setup.py    | 12 +++++++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 83e1f087d..a57de12ab 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -8,6 +8,8 @@ env:
     - T2T_DATA_DIR=/tmp/t2t-data
     - T2T_TRAIN_DIR=/tmp/t2t-train
     - TF_LATEST="1.8.*"
+    # This is necessary to have gsutil work with Python 2.7
+    - BOTO_CONFIG=/dev/null
   matrix:
     # We test against the last 4 versions of TensorFlow
     # If updating, also update TF_LATEST above
@@ -35,6 +37,8 @@ before_install:
 install:
   - pip install -q "tensorflow==$TF_VERSION"
   - pip install -q .[tests]
+  # Make sure to install the atari extras for gym
+  - pip install "gym[atari]"
   # Make sure we have the latest version of numpy - avoid problems we were
   # seeing with Python 3
   - pip install -q -U numpy
@@ -83,7 +87,7 @@ script:
         pytest tensor2tensor/rl/model_rl_experiment_test.py;
     fi
   - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
-        jupyter nbconvert --to notebook --execute tensor2tensor/notebooks/hello_t2t.ipynb;
+        jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/hello_t2t.ipynb;
     fi
 
   # Export and query (on Python 2 only)
diff --git a/setup.py b/setup.py
index 72be5a0c9..2cda2c4c2 100644
--- a/setup.py
+++ b/setup.py
@@ -49,7 +49,17 @@
     extras_require={
         'tensorflow': ['tensorflow>=1.5.0'],
         'tensorflow_gpu': ['tensorflow-gpu>=1.5.0'],
-        'tests': ['pytest', 'mock', 'pylint', 'jupyter', 'gym[atari]'],
+        'tests': [
+            'pytest',
+            'mock',
+            'pylint',
+            'jupyter',
+            'gsutil'
+            # Need atari extras for Travis tests, but because gym is already in
+            # install_requires, pip skips the atari extras, so we instread do an
+            # explicit pip install gym[atari] for the tests.
+            # 'gym[atari]',
+        ],
     },
     classifiers=[
         'Development Status :: 4 - Beta',

From 920ca459a90decfcbdc80d81c6ce33e7489296a9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sun, 27 May 2018 16:01:24 -0700
Subject: [PATCH 0029/2720] fix typo s/instread/instead

PiperOrigin-RevId: 198240333
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 2cda2c4c2..fed6e8c20 100644
--- a/setup.py
+++ b/setup.py
@@ -56,7 +56,7 @@
             'jupyter',
             'gsutil'
             # Need atari extras for Travis tests, but because gym is already in
-            # install_requires, pip skips the atari extras, so we instread do an
+            # install_requires, pip skips the atari extras, so we instead do an
             # explicit pip install gym[atari] for the tests.
             # 'gym[atari]',
         ],

From 75e25acc6c3e81878938d791217c5bb2d44b1dc4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 29 May 2018 07:53:23 -0700
Subject: [PATCH 0030/2720] Use the dynamic shape of the input to determine the
 batch size in `serving_input_fn`.

PiperOrigin-RevId: 198395185
---
 tensor2tensor/data_generators/problem.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index b8731273a..5f7ca2808 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -872,7 +872,9 @@ def serving_input_fn(self, hparams):
     dataset = dataset.map(lambda ex: self.preprocess_example(ex, mode, hparams))
     dataset = dataset.map(self.maybe_reverse_and_copy)
     dataset = dataset.map(data_reader.cast_ints_to_int32)
-    dataset = dataset.padded_batch(1000, dataset.output_shapes)
+    dataset = dataset.padded_batch(
+        tf.shape(serialized_example, out_type=tf.int64)[0],
+        dataset.output_shapes)
     dataset = dataset.map(standardize_shapes)
     features = tf.contrib.data.get_single_element(dataset)
 

From ceeacda87912876dd6607e12ad04ac3148a2dda3 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Tue, 29 May 2018 10:16:26 -0700
Subject: [PATCH 0031/2720] Reduce decode speed from 95ms to 81ms without
 sacrificing accuracy

PiperOrigin-RevId: 198414682
---
 tensor2tensor/layers/discretization.py           | 6 ++----
 tensor2tensor/models/research/transformer_vae.py | 4 ----
 2 files changed, 2 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index a09d0be95..7b5818f96 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -291,8 +291,7 @@ def embed(x,
     else:
       raise ValueError("Unknown bottleneck kind.")
 
-    h2 = tf.layers.dense(tf.nn.relu(h1), filter_size, name="vch2")
-    return tf.layers.dense(tf.nn.relu(h2), hidden_size, name="vcfin")
+    return h1
 
 
 def vae(x, name, z_size):
@@ -650,8 +649,7 @@ def discrete_bottleneck(x,
     else:
       raise ValueError("Unknown discretization method.")
 
-    h2 = tf.layers.dense(tf.nn.relu(h1), filter_size, name="vch2")
-    res = tf.layers.dense(tf.nn.relu(h2), hidden_size, name="vcfin")
+    res = h1
 
     embed_fn = partial(
         embed,
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 12e772d1d..94cc0b67c 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -455,14 +455,10 @@ def bn_inputs():
       # reshape back to 4d here
       if hparams.task == "image":
         targets = tf.reshape(targets, original_targets_shape)
-    if hparams.task == "translate":
-      targets = tf.concat([tf.reverse(latents_dense, [1]), targets], axis=1)
 
   res = decode_transformer(inputs, ed, targets, hparams, "decoder",
                            causal=hparams.causal)
   if hparams.do_ae:
-    if hparams.task == "translate":
-      res = res[:, common_layers.shape_list(latents_dense)[1]:, :, :]
     if hparams.do_mask and hparams.do_refine:
       def refine_res():
         # return residual_conv(res, 1, (5, 1), hparams, "refine")

From 8bb874893395794fb2f373699dd50a30ad539744 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 29 May 2018 15:27:18 -0700
Subject: [PATCH 0032/2720] Fix strip_extraneous

PiperOrigin-RevId: 198467702
---
 tensor2tensor/data_generators/text_encoder.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 2197bb867..4a37e69c4 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -129,7 +129,7 @@ def decode(self, ids, strip_extraneous=False):
       s: human-readable string.
     """
     if strip_extraneous:
-      ids = strip_ids(ids, list(range(self._num_reserved_ids)))
+      ids = strip_ids(ids, list(range(self._num_reserved_ids or 0)))
     return " ".join(self.decode_list(ids))
 
   def decode_list(self, ids):
@@ -172,7 +172,7 @@ def encode(self, s):
 
   def decode(self, ids, strip_extraneous=False):
     if strip_extraneous:
-      ids = strip_ids(ids, list(range(self._num_reserved_ids)))
+      ids = strip_ids(ids, list(range(self._num_reserved_ids or 0)))
     numres = self._num_reserved_ids
     decoded_ids = []
     int2byte = six.int2byte
@@ -531,7 +531,7 @@ def decode(self, ids, strip_extraneous=False):
       a native string
     """
     if strip_extraneous:
-      ids = strip_ids(ids, list(range(self._num_reserved_ids)))
+      ids = strip_ids(ids, list(range(self._num_reserved_ids or 0)))
     return unicode_to_native(
         tokenizer.decode(self._subtoken_ids_to_tokens(ids)))
 

From a42e00148f7393f5ba3017a89997b170db488c96 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 30 May 2018 17:12:50 -0700
Subject: [PATCH 0033/2720] Fixed the issue with varying-length sequences in a
 batch for lstm.

PiperOrigin-RevId: 198645336
---
 tensor2tensor/models/lstm.py | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 05395e042..1c6c21138 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -27,7 +27,7 @@
 import tensorflow as tf
 
 
-def lstm(inputs, hparams, train, name, initial_state=None):
+def lstm(inputs, hparams, train, name, initial_state=None, lengths=None):
   """Run LSTM cell on inputs, assuming they are [batch x time x size]."""
 
   def dropout_lstm_cell():
@@ -41,12 +41,13 @@ def dropout_lstm_cell():
         tf.contrib.rnn.MultiRNNCell(layers),
         inputs,
         initial_state=initial_state,
+        sequence_length=lengths,
         dtype=tf.float32,
         time_major=False)
 
 
 def lstm_attention_decoder(inputs, hparams, train, name, initial_state,
-                           encoder_outputs):
+                           encoder_outputs, lengths=None):
   """Run LSTM cell with attention on inputs of shape [batch x time x size]."""
 
   def dropout_lstm_cell():
@@ -63,7 +64,7 @@ def dropout_lstm_cell():
     raise ValueError("Unknown hparams.attention_mechanism = %s, must be "
                      "luong or bahdanau." % hparams.attention_mechanism)
   attention_mechanism = attention_mechanism_class(
-      hparams.hidden_size, encoder_outputs)
+      hparams.hidden_size, encoder_outputs, memory_sequence_length=lengths)
 
   cell = tf.contrib.seq2seq.AttentionWrapper(
       tf.nn.rnn_cell.MultiRNNCell(layers),
@@ -116,16 +117,21 @@ def lstm_seq2seq_internal(inputs, targets, hparams, train):
 def lstm_seq2seq_internal_attention(inputs, targets, hparams, train):
   """LSTM seq2seq model with attention, main step used for training."""
   with tf.variable_scope("lstm_seq2seq_attention"):
+    # This is a temporary fix for varying-length sequences within in a batch.
+    # A more complete fix should pass a length tensor from outside so that
+    # all the lstm variants can use it.
+    lengths = tf.reduce_sum(
+        common_layers.mask_from_embedding(inputs), [1, 2, 3])
     # Flatten inputs.
     inputs = common_layers.flatten4d3d(inputs)
     # LSTM encoder.
     encoder_outputs, final_encoder_state = lstm(
-        tf.reverse(inputs, axis=[1]), hparams, train, "encoder")
+        inputs, hparams, train, "encoder", lengths=lengths)
     # LSTM decoder with attention
     shifted_targets = common_layers.shift_right(targets)
     decoder_outputs, _ = lstm_attention_decoder(
         common_layers.flatten4d3d(shifted_targets), hparams, train, "decoder",
-        final_encoder_state, encoder_outputs)
+        final_encoder_state, encoder_outputs, lengths=lengths)
     return tf.expand_dims(decoder_outputs, axis=2)
 
 

From 44defb5f8b66fffa3b0e41bf7a508af5277af2dc Mon Sep 17 00:00:00 2001
From: Mostafa Dehghani <dehghani.mostafa@gmail.com>
Date: Thu, 31 May 2018 16:19:31 +0200
Subject: [PATCH 0034/2720] Changing the names and cleaning hparams_sets of the
 Universal Transformer based on the NIPS submission.  (#837)

* Changing the names and cleaning hparams_sets based on the NIPS submission

* Adding some ranges for haparams_set (for both the Universal Transformer and Adaptive Universal Transformer)

* Updating name of r_transofmer in .travis.yml --to ignore universal_transformer_test.py until tf.foldl is updated in tensorflow 1.9

* Resolving the .travis.yml conflict.

* Update .travis.yml

* fixing lint errors

* Revert "Merge branch 'master' of https://github.com/MostafaDehghani/tensor2tensor"

This reverts commit d016ab368f2c515e855d0f3d96612c3944823fea, reversing
changes made to 41434d4a74fb8b64518c03a1a19752de2788251c.

* Revert "Merge branch 'master' into master"

This reverts commit d88575706553d2f6e3d806b62f90b3adbedc74bd, reversing
changes made to 42981f3b5c3c62e937dd4c329f73a552e1810237.

* Changing the names and cleaning hparams_sets based on the NIPS submission

* Adding some ranges for haparams_set (for both the Universal Transformer and Adaptive Universal Transformer)

* Updating name of r_transofmer in .travis.yml --to ignore universal_transformer_test.py until tf.foldl is updated in tensorflow 1.9

* fixing lint errors

* Revert "Merge branch 'master' of https://github.com/MostafaDehghani/tensor2tensor"

This reverts commit d016ab368f2c515e855d0f3d96612c3944823fea, reversing
changes made to 41434d4a74fb8b64518c03a1a19752de2788251c.

* Revert "Merge branch 'master' into master"

This reverts commit d88575706553d2f6e3d806b62f90b3adbedc74bd, reversing
changes made to 42981f3b5c3c62e937dd4c329f73a552e1810237.

* Revert "merge text_encoder"

This reverts commit 008ff4c3083f88d8824977419dcadcf66184f06b, reversing
changes made to a78722864a0097286518dc092f3fa2380fb7a0d5.

* Revert "Revert "Merge branch 'master' of https://github.com/MostafaDehghani/tensor2tensor""

This reverts commit 06eb36d43a40bdf450aceb9b10c10db4297449a3.

* Revert "Revert "Merge branch 'master' into master""

This reverts commit a78722864a0097286518dc092f3fa2380fb7a0d5.

* fixing arguments of encode and decode to match the overridden methods
---
 .travis.yml                                   |   4 +-
 tensor2tensor/models/__init__.py              |   2 +-
 ...ransformer.py => universal_transformer.py} | 511 +++++++-----------
 ..._test.py => universal_transformer_test.py} |  17 +-
 ..._util.py => universal_transformer_util.py} | 305 ++++++-----
 5 files changed, 392 insertions(+), 447 deletions(-)
 rename tensor2tensor/models/research/{r_transformer.py => universal_transformer.py} (64%)
 rename tensor2tensor/models/research/{r_transformer_test.py => universal_transformer_test.py} (82%)
 rename tensor2tensor/models/research/{r_transformer_util.py => universal_transformer_util.py} (86%)

diff --git a/.travis.yml b/.travis.yml
index a57de12ab..59065472f 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -54,14 +54,14 @@ script:
   #   * visualization_test
   #   * model_rl_experiment_test
   # algorithmic_math_test: flaky
-  # r_transformer_test: requires new feature in tf.foldl (rm with TF 1.9)
+  # universal_transformer_test: requires new feature in tf.foldl (rm with TF 1.9)
   - pytest
     --ignore=tensor2tensor/utils/registry_test.py
     --ignore=tensor2tensor/utils/trainer_lib_test.py
     --ignore=tensor2tensor/visualization/visualization_test.py
     --ignore=tensor2tensor/bin/t2t_trainer_test.py
     --ignore=tensor2tensor/data_generators/algorithmic_math_test.py
-    --ignore=tensor2tensor/models/research/r_transformer_test.py
+    --ignore=tensor2tensor/models/research/universal_transformer_test.py
     --ignore=tensor2tensor/rl/model_rl_experiment_test.py
   - pytest tensor2tensor/utils/registry_test.py
   - pytest tensor2tensor/utils/trainer_lib_test.py
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index b6f7e1d73..9b49dff70 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -44,7 +44,7 @@
 from tensor2tensor.models.research import lm_experiments
 from tensor2tensor.models.research import multimodel
 from tensor2tensor.models.research import next_frame
-from tensor2tensor.models.research import r_transformer
+from tensor2tensor.models.research import universal_transformer
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research import super_lm
 from tensor2tensor.models.research import transformer_moe
diff --git a/tensor2tensor/models/research/r_transformer.py b/tensor2tensor/models/research/universal_transformer.py
similarity index 64%
rename from tensor2tensor/models/research/r_transformer.py
rename to tensor2tensor/models/research/universal_transformer.py
index 8f36f9861..51050d8cf 100644
--- a/tensor2tensor/models/research/r_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -12,42 +12,38 @@
 # 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.
-"""Transformers with depthwise recurrency (go/r-transformer).
+"""Universal Transformers.
 
 
-A high-level explanation on the idea and the architecture:
-
-The vanilla Transformer model has no recurrence and struggles with some tasks
-that a fully recurrent model can easily solve. Instead of incorporating
-recurrence in time (which has a dependency on sequence length T),
-we apply recurrence in depth (which we can set to some fixed length D << T),
-and apply self-attention instead of sequential processing to enable the model
-to incorporate long-range dependencies.
-
-Structure of the code is explained in r_transformer_util.py
+Universal Transformer is recurrent in depth while employing self-attention
+to combine information from different parts of sequences.
+In contrast to the Transformer, given enough memory its recurrence in depth
+makes the Universal Transformer computationally universal.
 """
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
-from tensor2tensor.models.research import r_transformer_util
+from tensor2tensor.models.research import universal_transformer_util
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
 
 @registry.register_model
-class RTransformer(transformer.Transformer):
-  """R-Transformer: Depth-wise recurrent transformer model."""
+class UniversalTransformer(transformer.Transformer):
+  """Universal Transformer: Depth-wise recurrent transformer model."""
 
   def encode(self, inputs, target_space, hparams, features=None, losses=None):
-    """Encode r-transformer inputs.
+    """Encode Universal Transformer inputs.
 
     It is similar to "transformer.encode", but it uses
-    "r_transformer_util.r_transformer_encoder" instead of
+    "universal_transformer_util.universal_transformer_encoder" instead of
     "transformer.transformer_encoder".
 
     Args:
@@ -57,7 +53,6 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
       hparams: hyperparmeters for model.
       features: optionally pass the entire features dictionary as well.
         This is needed now for "packed" datasets.
-      losses: optional list onto which to append extra training losses
 
     Returns:
       Tuple of:
@@ -79,16 +74,17 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
     encoder_input = tf.nn.dropout(encoder_input,
                                   1.0 - hparams.layer_prepostprocess_dropout)
 
-    (encoder_output,
-     encoder_extra_output) = r_transformer_util.r_transformer_encoder(
-         encoder_input,
-         self_attention_bias,
-         hparams,
-         nonpadding=transformer.features_to_nonpadding(features, "inputs"),
-         save_weights_to=self.attention_weights)
+    (encoder_output, encoder_extra_output) = (
+        universal_transformer_util.universal_transformer_encoder(
+            encoder_input,
+            self_attention_bias,
+            hparams,
+            nonpadding=transformer.features_to_nonpadding(features, "inputs"),
+            save_weights_to=self.attention_weights))
 
     return encoder_output, encoder_decoder_attention_bias, encoder_extra_output
 
+
   def decode(self,
              decoder_input,
              encoder_output,
@@ -98,10 +94,10 @@ def decode(self,
              cache=None,
              nonpadding=None,
              losses=None):
-    """Decode R-Transformer outputs from encoder representation.
+    """Decode Universal Transformer outputs from encoder representation.
 
     It is similar to "transformer.decode", but it uses
-    "r_transformer_util.r_transformer_decoder" instead of
+    "universal_transformer_util.universal_transformer_decoder" instead of
     "transformer.transformer_decoder".
 
     Args:
@@ -114,10 +110,7 @@ def decode(self,
       decoder_self_attention_bias: Bias and mask weights for decoder
         self-attention. [batch_size, decoder_length]
       hparams: hyperparmeters for model.
-      cache: dict, containing tensors which are the results of previous
-          attentions, used for fast decoding.
       nonpadding: optional Tensor with shape [batch_size, decoder_length]
-      losses: optional list onto which to append extra training losses
 
     Returns:
        Tuple of:
@@ -127,26 +120,29 @@ def decode(self,
             variants of the model (e.g. in ACT, to pass the ponder-time to body)
 
     """
-    del cache, losses
+    del losses
+    # TODO(dehghani): enable caching.
+    del cache
 
     decoder_input = tf.nn.dropout(decoder_input,
                                   1.0 - hparams.layer_prepostprocess_dropout)
 
-    # No caching in r-transformers!
-    decoder_output, dec_extra_output = r_transformer_util.r_transformer_decoder(
-        decoder_input,
-        encoder_output,
-        decoder_self_attention_bias,
-        encoder_decoder_attention_bias,
-        hparams,
-        nonpadding=nonpadding,
-        save_weights_to=self.attention_weights)
+    # No caching in Universal Transformers!
+    (decoder_output, dec_extra_output) = (
+        universal_transformer_util.universal_transformer_decoder(
+            decoder_input,
+            encoder_output,
+            decoder_self_attention_bias,
+            encoder_decoder_attention_bias,
+            hparams,
+            nonpadding=nonpadding,
+            save_weights_to=self.attention_weights))
 
     # Expand since t2t expects 4d tensors.
     return tf.expand_dims(decoder_output, axis=2), dec_extra_output
 
   def body(self, features):
-    """R-Transformer main model_fn.
+    """Universal Transformer main model_fn.
 
 
     Args:
@@ -236,7 +232,7 @@ def _greedy_infer(self, features, decode_length):
     Raises:
       NotImplementedError: If there are multiple data shards.
     """
-    # TODO(dehghani): Support fast decoding for r-transformer (needs caching)
+    # TODO(dehghani): Support fast decoding for Universal Transformer
     return self._slow_greedy_infer(features, decode_length)
 
   def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
@@ -259,15 +255,15 @@ def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
               None if using greedy decoding (beam_size=1)
       }
     """
-    # Caching is not ebabled in r-transformer
-    # TODO(dehghani): Support fast decoding for r-transformer(needs caching)
+    # Caching is not ebabled in Universal Transformer
+    # TODO(dehghani): Support fast decoding for Universal Transformer
     return self._beam_decode_slow(features, decode_length, beam_size,
                                   top_beams, alpha)
 
 
 @registry.register_model
-class RTransformerEncoder(transformer.Transformer):
-  """R-Transformer Encoder: Depth-wise recurrent transformer encoder-only."""
+class UniversalTransformerEncoder(transformer.Transformer):
+  """Universal Transformer Encoder: Has no decoder (e.g.for classification)."""
 
   def encode(self, inputs, target_space, hparams, features=None, losses=None):
     """Encode transformer inputs.
@@ -279,7 +275,6 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
       hparams: hyperparmeters for model.
       features: optionally pass the entire features dictionary as well.
         This is needed now for "packed" datasets.
-      losses: optional list onto which to append extra training losses
 
     Returns:
       Tuple of:
@@ -297,18 +292,18 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
     encoder_input = tf.nn.dropout(encoder_input,
                                   1.0 - hparams.layer_prepostprocess_dropout)
 
-    (encoder_output,
-     encoder_extra_output) = r_transformer_util.r_transformer_encoder(
-         encoder_input,
-         self_attention_bias,
-         hparams,
-         nonpadding=transformer.features_to_nonpadding(features, "inputs"),
-         save_weights_to=self.attention_weights)
+    (encoder_output, encoder_extra_output) = (
+        universal_transformer_util.universal_transformer_encoder(
+            encoder_input,
+            self_attention_bias,
+            hparams,
+            nonpadding=transformer.features_to_nonpadding(features, "inputs"),
+            save_weights_to=self.attention_weights))
 
     return encoder_output, encoder_extra_output
 
   def body(self, features):
-    """R-Transformer main model_fn.
+    """Universal Transformer main model_fn.
 
     Args:
       features: Map of features to the model. Should contain the following:
@@ -322,8 +317,8 @@ def body(self, features):
     """
     hparams = self._hparams
 
-    assert self.has_input, ("r_transformer_encoder is applicable on problems"
-                            "with inputs")
+    assert self.has_input, ("universal_transformer_encoder is applicable on "
+                            "problems with inputs")
 
     inputs = features["inputs"]
     target_space = features["target_space_id"]
@@ -342,22 +337,22 @@ def body(self, features):
     return encoder_output
 
 
-def update_hparams_for_r_transformer(hparams):
-  """Adds deault hparams for all of the variants of the R-transformer.
+def update_hparams_for_universal_transformer(hparams):
+  """Adds deault hparams for all of the variants of the Universal Transformer.
 
   Args:
     hparams: default hparams (usually one of the standard hparams from
       transformer model (like "transformer_base")
 
   Returns:
-    hparams with default values for R-Transformers hyper-parameters
+    hparams with default values for Universal Transformers hyper-parameters
 
   """
-  # If not None, mixes vanilla transformer with r-transformer.
-  # Options: None, "before_rt", and "after_rt".
+  # If not None, mixes vanilla transformer with Universal Transformer.
+  # Options: None, "before_ut", and "after_ut".
   hparams.add_hparam("mix_with_transformer", None)
 
-  # Number of vanilla transformer layers used to be mixed with r-transofmer.
+  # Number of vanilla transformer layers used to be mixed with u-transofmer.
   hparams.add_hparam("num_mixedin_layers", 2)
 
   # Type of recurrency:
@@ -368,13 +363,15 @@ def update_hparams_for_r_transformer(hparams):
   hparams.add_hparam("num_rec_steps", hparams.num_hidden_layers)
 
   # Add the positional mebedding at each step(horisontal timing)
-  hparams.add_hparam("add_position_timing_signal", False)
+  hparams.add_hparam("add_position_timing_signal", True)
+  if hparams.add_position_timing_signal:
+    hparams.pos = None
   # Logic of position shifting when using timing signal:
   # None, "random", "step"
   hparams.add_hparam("position_start_index", None)
 
   # Add an step embedding at each step (vertical timing)
-  hparams.add_hparam("add_step_timing_signal", False)
+  hparams.add_hparam("add_step_timing_signal", True)
   # Either "learned" or "sinusoid"
   hparams.add_hparam("step_timing_signal_type", "learned")
 
@@ -382,7 +379,7 @@ def update_hparams_for_r_transformer(hparams):
   # Options: "add" and "concat".
   hparams.add_hparam("add_or_concat_timing_signal", "add")
 
-  # Add SRU at the beginning of each r-transformer step.
+  # Add SRU at the beginning of each Universal Transformer step.
   # This can be considered as a position timing signal
   hparams.add_hparam("add_sru", False)
 
@@ -442,23 +439,32 @@ def update_hparams_for_r_transformer(hparams):
 
 
 @registry.register_hparams
-def r_transformer_big():
+def universal_transformer_base():
   hparams = transformer.transformer_big()
-  hparams = update_hparams_for_r_transformer(hparams)
+  hparams = update_hparams_for_universal_transformer(hparams)
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_base():
+def universal_transformer_big():
+  hparams = transformer.transformer_big()
+  hparams = update_hparams_for_universal_transformer(hparams)
+  hparams.hidden_size = 2048
+  hparams.filter_size = 8192
+  return hparams
+
+
+@registry.register_hparams
+def universal_transformer_small():
   hparams = transformer.transformer_base()
-  hparams = update_hparams_for_r_transformer(hparams)
+  hparams = update_hparams_for_universal_transformer(hparams)
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_tiny():
+def universal_transformer_tiny():
   hparams = transformer.transformer_tiny()
-  hparams = update_hparams_for_r_transformer(hparams)
+  hparams = update_hparams_for_universal_transformer(hparams)
   hparams.num_rec_steps = 8
   return hparams
 
@@ -474,101 +480,101 @@ def transformer_teeny():
 
 
 @registry.register_hparams
-def r_transformer_teeny():
+def universal_transformer_teeny():
   hparams = transformer_teeny()
-  hparams = update_hparams_for_r_transformer(hparams)
+  hparams = update_hparams_for_universal_transformer(hparams)
   hparams.num_rec_steps = 10
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_base_dropconnect():
-  hparams = r_transformer_base()
+def universal_transformer_small_dropconnect():
+  hparams = universal_transformer_small()
   hparams.gate_ffn_layer = "dense_dropconnect"
   hparams.add_hparam("dropconnect_dropout", 0.5)
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_base():
-  hparams = r_transformer_base()
+def adaptive_universal_transformer_small():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "act"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_tiny():
-  hparams = r_transformer_tiny()
+def adaptive_universal_transformer_tiny():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "act"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_big():
-  hparams = r_transformer_big()
+def adaptive_universal_transformer_base():
+  hparams = universal_transformer_base()
   hparams.recurrence_type = "act"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_random_base():
-  hparams = r_transformer_base()
+def adaptive_universal_transformer_random_small():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "act"
   hparams.act_type = "random"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_accumulated_base():
-  hparams = r_transformer_base()
+def adaptive_universal_transformer_accumulated_small():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "act"
   hparams.act_type = "accumulated"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_global_base():
-  hparams = r_transformer_base()
+def adaptive_universal_transformer_global_small():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "act"
   hparams.act_type = "global"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_accumulated_tiny():
-  hparams = r_transformer_tiny()
+def adaptive_universal_transformer_accumulated_tiny():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "act"
   hparams.act_type = "accumulated"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_global_tiny():
-  hparams = r_transformer_tiny()
+def adaptive_universal_transformer_global_tiny():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "act"
   hparams.act_type = "global"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_random_tiny():
-  hparams = r_transformer_tiny()
+def adaptive_universal_transformer_random_tiny():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "act"
   hparams.act_type = "random"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_base_sb():
-  hparams = r_transformer_base()
+def adaptive_universal_transformer_small_sb():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "act"
   hparams.batch_size = 2048
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_large():
-  hparams = r_transformer_base()
+def adaptive_universal_transformer_large():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "act"
   hparams.hidden_size = 1024
   hparams.batch_size = 2048
@@ -577,8 +583,8 @@ def r_transformer_act_large():
 
 
 @registry.register_hparams
-def r_transformer_act_tall():
-  hparams = r_transformer_base()
+def adaptive_universal_transformer_tall():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "act"
   hparams.num_hidden_layers = 16
   hparams.batch_size = 1024
@@ -587,8 +593,8 @@ def r_transformer_act_tall():
 
 
 @registry.register_hparams
-def r_transformer_act_tall_actlossw0():
-  hparams = r_transformer_base()
+def adaptive_universal_transformer_tall_actlossw0():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "act"
   hparams.num_hidden_layers = 16
   hparams.batch_size = 1024
@@ -597,8 +603,8 @@ def r_transformer_act_tall_actlossw0():
 
 
 @registry.register_hparams
-def r_transformer_act_tall_actlossw001():
-  hparams = r_transformer_base()
+def adaptive_universal_transformer_tall_actlossw001():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "act"
   hparams.num_hidden_layers = 16
   hparams.batch_size = 1024
@@ -607,8 +613,8 @@ def r_transformer_act_tall_actlossw001():
 
 
 @registry.register_hparams
-def r_transformer_act_base_d03():
-  hparams = r_transformer_base()
+def adaptive_universal_transformer_small_d03():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "act"
   hparams.layer_prepostprocess_dropout = 0.3
   hparams.attention_dropout = 0.3
@@ -617,8 +623,8 @@ def r_transformer_act_base_d03():
 
 
 @registry.register_hparams
-def r_transformer_act_big_d03():
-  hparams = r_transformer_big()
+def adaptive_universal_transformer_base_d03():
+  hparams = universal_transformer_base()
   hparams.recurrence_type = "act"
   hparams.layer_prepostprocess_dropout = 0.3
   hparams.attention_dropout = 0.3
@@ -627,8 +633,8 @@ def r_transformer_act_big_d03():
 
 
 @registry.register_hparams
-def r_transformer_act_tiny_d02():
-  hparams = r_transformer_tiny()
+def adaptive_universal_transformer_tiny_d02():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "act"
   hparams.layer_prepostprocess_dropout = 0.2
   hparams.attention_dropout = 0.2
@@ -637,8 +643,8 @@ def r_transformer_act_tiny_d02():
 
 
 @registry.register_hparams
-def r_transformer_act_tiny_d02_sb():
-  hparams = r_transformer_tiny()
+def adaptive_universal_transformer_tiny_d02_sb():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "act"
   hparams.layer_prepostprocess_dropout = 0.2
   hparams.attention_dropout = 0.2
@@ -648,16 +654,16 @@ def r_transformer_act_tiny_d02_sb():
 
 
 @registry.register_hparams
-def r_transformer_act_tiny_sb():
-  hparams = r_transformer_tiny()
+def adaptive_universal_transformer_tiny_sb():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "act"
   hparams.batch_size = 2048
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_tiny_d05():
-  hparams = r_transformer_tiny()
+def adaptive_universal_transformer_tiny_d05():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "act"
   hparams.layer_prepostprocess_dropout = 0.5
   hparams.attention_dropout = 0.5
@@ -666,327 +672,226 @@ def r_transformer_act_tiny_d05():
 
 
 @registry.register_hparams
-def r_transformer_base_sb():
-  hparams = r_transformer_base()
+def universal_transformer_small_sb():
+  hparams = universal_transformer_small()
   hparams.batch_size = 2048
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_skip_base():
-  hparams = r_transformer_base()
+def universal_transformer_skip_small():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "skip"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_skip_tiny():
-  hparams = r_transformer_tiny()
+def universal_transformer_skip_tiny():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "skip"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_highway_base():
-  hparams = r_transformer_base()
+def universal_transformer_highway_small():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "highway"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_highway_tiny():
-  hparams = r_transformer_tiny()
+def universal_transformer_highway_tiny():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "highway"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_dwa_base():
-  hparams = r_transformer_base()
+def universal_transformer_dwa_small():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "dwa"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_dwa_tiny():
-  hparams = r_transformer_tiny()
+def universal_transformer_dwa_tiny():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "dwa"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_dwa_tiny_test():
-  hparams = r_transformer_tiny()
+def universal_transformer_dwa_tiny_test():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "dwa"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_rnn_base():
-  hparams = r_transformer_base()
+def universal_transformer_rnn_small():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "rnn"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_gru_base():
-  hparams = r_transformer_base()
+def universal_transformer_gru_small():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "gru"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_lstm_base():
-  hparams = r_transformer_base()
+def universal_transformer_lstm_small():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "lstm"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_position_timing_tiny():
-  hparams = r_transformer_tiny()
-  hparams.pos = None
-  hparams.add_position_timing_signal = True
-  return hparams
-
-
-@registry.register_hparams
-def r_transformer_position_random_timing_base():
-  hparams = r_transformer_base()
-  hparams.pos = None
-  hparams.add_position_timing_signal = True
+def universal_transformer_position_random_timing_small():
+  hparams = universal_transformer_small()
   hparams.position_start_index = "random"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_position_random_timing_tiny():
-  hparams = r_transformer_tiny()
-  hparams.pos = None
-  hparams.add_position_timing_signal = True
+def universal_transformer_position_random_timing_tiny():
+  hparams = universal_transformer_tiny()
   hparams.position_start_index = "random"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_position_step_timing_tiny():
-  hparams = r_transformer_tiny()
-  hparams.pos = None
-  hparams.add_position_timing_signal = True
+def universal_transformer_position_step_timing_tiny():
+  hparams = universal_transformer_tiny()
   hparams.position_start_index = "step"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_step_timing_tiny():
-  hparams = r_transformer_tiny()
-  hparams.add_step_timing_signal = True
-  return hparams
-
-
-@registry.register_hparams
-def r_transformer_step_sinusoid_timing_tiny():
-  hparams = r_transformer_tiny()
-  hparams.add_step_timing_signal = True
+def universal_transformer_step_sinusoid_timing_tiny():
+  hparams = universal_transformer_tiny()
   hparams.step_timing_signal_type = "sinusoid"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_step_position_random_timing_tiny():
-  hparams = r_transformer_tiny()
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  hparams.position_start_index = "random"
-  hparams.add_step_timing_signal = True
-  return hparams
-
-
-@registry.register_hparams
-def r_transformer_act_position_timing_tiny():
-  hparams = r_transformer_tiny()
+def adaptive_universal_transformer_position_random_timing_tiny():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "act"
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  return hparams
-
-
-@registry.register_hparams
-def r_transformer_act_position_random_timing_tiny():
-  hparams = r_transformer_tiny()
-  hparams.recurrence_type = "act"
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
   hparams.position_start_index = "random"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_position_step_timing_tiny():
-  hparams = r_transformer_tiny()
+def adaptive_universal_transformer_position_step_timing_tiny():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "act"
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
   hparams.position_start_index = "step"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_step_timing_tiny():
-  hparams = r_transformer_tiny()
+def adaptive_universal_transformer_step_sinusoid_timing_tiny():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "act"
-  hparams.add_step_timing_signal = True
+  hparams.step_timing_signal_type = "sinusoid"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_step_position_random_timing_tiny():
-  hparams = r_transformer_tiny()
-  hparams.recurrence_type = "act"
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  hparams.position_start_index = "random"
-  hparams.add_step_timing_signal = True
+def universal_transformer_mix_after_ut_small():
+  hparams = universal_transformer_small()
+  hparams.mix_with_transformer = "before_ut"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_step_sinusoid_position_random_timing_tiny():
-  hparams = r_transformer_tiny()
+def adaptive_universal_transformer_mix_before_ut_small():
+  hparams = universal_transformer_small()
+  hparams.mix_with_transformer = "before_ut"
   hparams.recurrence_type = "act"
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  hparams.position_start_index = "random"
-  hparams.add_step_timing_signal = True
-  hparams.step_timing_signal_type = "sinusoid"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_step_sinusoid_position_timing_tiny():
-  hparams = r_transformer_tiny()
+def adaptive_universal_transformer_mix_after_ut_small():
+  hparams = universal_transformer_small()
+  hparams.mix_with_transformer = "after_ut"
   hparams.recurrence_type = "act"
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  hparams.add_step_timing_signal = True
-  hparams.step_timing_signal_type = "sinusoid"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_step_position_timing_tiny():
-  hparams = r_transformer_tiny()
+def adaptive_universal_transformer_concat_tiny():
+  hparams = universal_transformer_tiny()
   hparams.recurrence_type = "act"
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  hparams.add_step_timing_signal = True
-  return hparams
-
-
-@registry.register_hparams
-def r_transformer_step_position_timing_tiny():
-  hparams = r_transformer_tiny()
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  hparams.add_step_timing_signal = True
+  hparams.add_or_concat_timing_signal = "concat"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_step_position_random_timing_base():
-  hparams = r_transformer_base()
+def adaptive_universal_transformer_concat_small():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "act"
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  hparams.position_start_index = "random"
-  hparams.add_step_timing_signal = True
+  hparams.add_or_concat_timing_signal = "concat"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_step_position_timing_base():
-  hparams = r_transformer_base()
+def adaptive_universal_transformer_with_sru_small():
+  hparams = universal_transformer_small()
   hparams.recurrence_type = "act"
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  hparams.add_step_timing_signal = True
-  return hparams
-
-
-@registry.register_hparams
-def r_transformer_step_position_timing_base():
-  hparams = r_transformer_base()
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  hparams.add_step_timing_signal = True
-  return hparams
-
-
-@registry.register_hparams
-def r_transformer_mix_after_rt_base():
-  hparams = r_transformer_base()
-  hparams.mix_with_transformer = "before_rt"
+  hparams.add_sru = True
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_step_position_timing_mix_before_rt_base():
-  hparams = r_transformer_base()
-  hparams.mix_with_transformer = "before_rt"
-  hparams.recurrence_type = "act"
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  hparams.add_step_timing_signal = True
+def universal_transformer_fc_small():
+  hparams = universal_transformer_small()
+  hparams.transformer_ffn_type = "fc"
   return hparams
 
 
 @registry.register_hparams
-def r_mix_transformer_act_step_position_timing_mix_after_rt_base():
-  hparams = r_transformer_base()
-  hparams.mix_with_transformer = "after_rt"
-  hparams.recurrence_type = "act"
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  hparams.add_step_timing_signal = True
+def universal_transformer_fc_base():
+  hparams = universal_transformer_base()
+  hparams.transformer_ffn_type = "fc"
   return hparams
 
 
 @registry.register_hparams
-def r_transformer_act_step_position_timing_big():
-  hparams = r_transformer_big()
-  hparams.batch_size //= 2
-  hparams.recurrence_type = "act"
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  hparams.add_step_timing_signal = True
+def universal_transformer_fc_big():
+  hparams = universal_transformer_big()
+  hparams.transformer_ffn_type = "fc"
   return hparams
 
 
-@registry.register_hparams
-def r_transformer_act_step_position_timing_concat_tiny():
-  hparams = r_transformer_tiny()
-  hparams.recurrence_type = "act"
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  hparams.add_step_timing_signal = True
-  hparams.add_or_concat_timing_signal = "concat"
-  return hparams
+@registry.register_ranged_hparams
+def universal_transformer_base_range(rhp):
+  """Small range of hyperparameters."""
+  # After starting from base, set intervals for some parameters.
+  rhp.set_discrete("num_rec_steps", [6, 8, 10])
+  rhp.set_discrete("hidden_size", [1024, 2048, 4096])
+  rhp.set_discrete("filter_size", [2048, 4096, 8192])
+  rhp.set_discrete("num_heads", [8, 16, 32])
+  rhp.set_discrete("transformer_ffn_type", ["sepconv", "fc"])
+  rhp.set_float("learning_rate", 0.3, 3.0, scale=rhp.LOG_SCALE)
+  rhp.set_float("weight_decay", 0.0, 2.0)
 
 
-@registry.register_hparams
-def r_transformer_act_step_position_timing_concat_base():
-  hparams = r_transformer_base()
-  hparams.recurrence_type = "act"
-  hparams.add_position_timing_signal = True
-  hparams.pos = None
-  hparams.add_step_timing_signal = True
-  hparams.add_or_concat_timing_signal = "concat"
-  return hparams
 
 
-@registry.register_hparams
-def r_transformer_act_with_sru_base():
-  hparams = r_transformer_base()
-  hparams.recurrence_type = "act"
-  hparams.add_sru = True
-  return hparams
+@registry.register_ranged_hparams
+def adaptive_universal_transformer_base_range(rhp):
+  """Small range of hyperparameters."""
+  # After starting from base, set intervals for some parameters.
+  rhp.set_discrete("act_max_steps", [8, 16, 32])
+  rhp.set_float("act_loss_weight", 0.0, 0.5)
+  rhp.set_discrete("hidden_size", [1024, 2048, 4096])
+  rhp.set_discrete("filter_size", [2048, 4096, 8192])
+  rhp.set_discrete("num_heads", [8, 16, 32])
+  rhp.set_discrete("transformer_ffn_type", ["sepconv", "fc"])
+  rhp.set_float("learning_rate", 0.3, 3.0, scale=rhp.LOG_SCALE)
+  rhp.set_float("weight_decay", 0.0, 2.0)
diff --git a/tensor2tensor/models/research/r_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
similarity index 82%
rename from tensor2tensor/models/research/r_transformer_test.py
rename to tensor2tensor/models/research/universal_transformer_test.py
index 9ec513591..6c0922da2 100644
--- a/tensor2tensor/models/research/r_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -17,10 +17,13 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+# Dependency imports
+
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
-from tensor2tensor.models.research import r_transformer
+from tensor2tensor.models.research import universal_transformer
 
 import tensorflow as tf
 
@@ -30,10 +33,10 @@
 VOCAB_SIZE = 10
 
 
-class RTransformerTest(tf.test.TestCase):
+class UniversalTransformerTest(tf.test.TestCase):
 
-  def get_model(self, hparams, mode=tf.estimator.ModeKeys.TRAIN,
-                has_input=True):
+  def get_model(self,
+                hparams, mode=tf.estimator.ModeKeys.TRAIN, has_input=True):
     hparams.hidden_size = 8
     hparams.filter_size = 32
     hparams.num_heads = 1
@@ -54,10 +57,12 @@ def get_model(self, hparams, mode=tf.estimator.ModeKeys.TRAIN,
         "target_space_id": tf.constant(1, dtype=tf.int32)
     }
 
-    return r_transformer.RTransformer(hparams, mode, p_hparams), features
+    return universal_transformer.UniversalTransformer(
+        hparams, mode, p_hparams), features
 
   def testTransformer(self):
-    model, features = self.get_model(r_transformer.r_transformer_base())
+    model, features = self.get_model(
+        universal_transformer.universal_transformer_base())
     logits, _ = model(features)
     with self.test_session() as session:
       session.run(tf.global_variables_initializer())
diff --git a/tensor2tensor/models/research/r_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
similarity index 86%
rename from tensor2tensor/models/research/r_transformer_util.py
rename to tensor2tensor/models/research/universal_transformer_util.py
index 521adbf8a..b5ffbc07b 100644
--- a/tensor2tensor/models/research/r_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -12,24 +12,32 @@
 # 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.
-"""Utilities for R-Transformer.
-
-
-R-Transformer learns a function (for instance the transformer multi-head
-attention plus a feed-forward unit) and uses this function over n-steps to
-process the input.
-In other words, we can describe this as having a vanilla transformer, in which
-the weights in the layers are shared and we have a module(the recurrency module)
-next to this transformer that controls how steps communicate with each other in
-depth.
-
-For instance, the recurrency module, can be a simple identity function
-which passes the output of a step as the input to next step (applying one layer
-of transformer n times on the input in a row --> lead to a better
-generalization!). Or as another example, the recurrent module can be an LSTM,
-(filliped vertically) next to the transformer which controls how state of the
-model changes in depth, Or even a grit transformer (a transformer which learns
-the attention over steps of an R-Transformer)
+"""Utilities for Universal Transformer.
+
+The Universal Transformer is based on the popular encoder-decoder architecture.
+However, as opposed to a fixed stack of distinct layers (as is usually the case
+for most popular neural sequence models), the Universal Transformer is
+recurrent "in depth", and repeatedly applies the same series of functions with
+the same parameters to all elements of the sequence in parallel, revising their
+representations with every step. The encoder and decoder have the same
+recurrent structure, but the decoder additionally consumes the final encoder
+representations for each position. Like the Transformer, the Universal
+Transformer is autoregressive. Trained using teacher-forcing, at generation
+time it produces its output one position at a time, with the decoder consuming
+the previously produced output positions.
+
+Given an input sequence of length m, we start with a matrix whose rows are the
+d-dimensional embeddings of the symbols at each position of the sequence.
+The Universal Transformer then iteratively computes representation of the input
+at each step by applying the multiheaded dot-product self-attention mechanism,
+followed by a recurrent transition function. We also add residual connections
+around each of these function blocks and apply dropout and layer normalization.
+
+The recurrent transition function in fact controls how steps communicate with
+each other in depth. For instance, the recurrent transition, can be a simple
+identity function which passes the output of a step as the input to next step.
+Or it can be an LSTM (filliped vertically) next to the transformer which
+controls how state  of the model changes in depth, Or even another transformer.
 
 """
 
@@ -39,6 +47,9 @@
 
 import copy
 import functools
+
+# Dependency imports
+
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
@@ -47,17 +58,17 @@
 import tensorflow as tf
 
 
-def r_transformer_encoder(encoder_input,
-                          encoder_self_attention_bias,
-                          hparams,
-                          name="encoder",
-                          nonpadding=None,
-                          save_weights_to=None,
-                          make_image_summary=True):
-  """R_transformer_encoder function.
+def universal_transformer_encoder(encoder_input,
+                                  encoder_self_attention_bias,
+                                  hparams,
+                                  name="encoder",
+                                  nonpadding=None,
+                                  save_weights_to=None,
+                                  make_image_summary=True):
+  """Universal Transformer encoder function.
 
   Prepares all the arguments and the inputs and passes it to a
-  r_transformer_layer to encode the encoder_input.
+  universal_transformer_layer to encode the encoder_input.
 
   Args:
     encoder_input: a Tensor
@@ -110,7 +121,7 @@ def r_transformer_encoder(encoder_input,
         save_weights_to=save_weights_to,
         make_image_summary=make_image_summary)
 
-    x, extra_output = r_transformer_layer(
+    x, extra_output = universal_transformer_layer(
         x, hparams, ffn_unit, attention_unit, pad_remover=pad_remover)
 
     if hparams.get("use_memory_as_last_state", False):
@@ -118,19 +129,19 @@ def r_transformer_encoder(encoder_input,
     return common_layers.layer_preprocess(x, hparams), extra_output
 
 
-def r_transformer_decoder(decoder_input,
-                          encoder_output,
-                          decoder_self_attention_bias,
-                          encoder_decoder_attention_bias,
-                          hparams,
-                          name="decoder",
-                          nonpadding=None,
-                          save_weights_to=None,
-                          make_image_summary=True):
-  """R_transformer decoder function.
+def universal_transformer_decoder(decoder_input,
+                                  encoder_output,
+                                  decoder_self_attention_bias,
+                                  encoder_decoder_attention_bias,
+                                  hparams,
+                                  name="decoder",
+                                  nonpadding=None,
+                                  save_weights_to=None,
+                                  make_image_summary=True):
+  """Universal Transformer decoder function.
 
   Prepares all the arguments and the inputs and passes it to a
-  core_r_transformer_layer to decoder.
+  core_universal_transformer_layer to decoder.
 
   Args:
     decoder_input: a Tensor
@@ -175,13 +186,18 @@ def r_transformer_decoder(decoder_input,
         save_weights_to=save_weights_to,
         make_image_summary=make_image_summary)
 
-    x, extra_output = r_transformer_layer(x, hparams, ffn_unit, attention_unit)
+    x, extra_output = universal_transformer_layer(
+        x, hparams, ffn_unit, attention_unit)
 
     return common_layers.layer_preprocess(x, hparams), extra_output
 
 
-def r_transformer_layer(x, hparams, ffn_unit, attention_unit, pad_remover=None):
-  """Core function applying the r-transforemr layer.
+def universal_transformer_layer(x,
+                                hparams,
+                                ffn_unit,
+                                attention_unit,
+                                pad_remover=None):
+  """Core function applying the universal transforemr layer.
 
   Args:
     x: input
@@ -218,33 +234,38 @@ def add_vanilla_transformer_layer(x, num_layers):
         x = ffn_unit(attention_unit(x))
     return x
 
-  with tf.variable_scope("r_transformer_%s" % hparams.recurrence_type):
+  with tf.variable_scope("universal_transformer_%s" % hparams.recurrence_type):
 
-    if hparams.mix_with_transformer == "before_rt":
+    if hparams.mix_with_transformer == "before_ut":
       x = add_vanilla_transformer_layer(x, hparams.num_mixedin_layers)
 
     if hparams.recurrence_type == "act":
-      return r_transformer_act(x, hparams, ffn_unit, attention_unit)
+      return universal_transformer_act(x, hparams, ffn_unit, attention_unit)
 
     else:  # for all the other recurrency types with fixed number of steps
-      rt_function, initializer = get_rt_layer(x, hparams, ffn_unit,
+
+      ut_function, initializer = get_ut_layer(x, hparams, ffn_unit,
                                               attention_unit, pad_remover)
 
       output, _, extra_output = tf.foldl(
-          rt_function, tf.range(hparams.num_rec_steps), initializer=initializer)
+          ut_function, tf.range(hparams.num_rec_steps), initializer=initializer)
 
-      # This can be the if we use r_transformer_lstm layer.
+      # This can be the if we use universal_transformer_lstm layer.
       if hparams.get("use_memory_as_final_state", False):
         output = extra_output
 
-    if hparams.mix_with_transformer == "after_rt":
+    if hparams.mix_with_transformer == "after_ut":
       output = add_vanilla_transformer_layer(output, hparams.num_mixedin_layers)
 
     return output, extra_output
 
 
-def get_rt_layer(x, hparams, ffn_unit, attention_unit, pad_remover=None):
-  """provides the function that is used in r-transforemr steps.
+def get_ut_layer(x,
+                 hparams,
+                 ffn_unit,
+                 attention_unit,
+                 pad_remover=None):
+  """Provides the function that is used in universal transforemr steps.
 
   Args:
     x: input
@@ -254,33 +275,33 @@ def get_rt_layer(x, hparams, ffn_unit, attention_unit, pad_remover=None):
     pad_remover: to mask out padding in convolutional layers (efficiency).
 
   Returns:
-    rt_function and the rt_initializer
+    ut_function and the ut_initializer
 
   Raises:
     ValueError: Unknown recurrence type
   """
 
   if hparams.recurrence_type == "basic":
-    rt_initializer = (x, x, x)  # (state, input, memory)
-    rt_function = functools.partial(
-        r_transformer_basic,
+    ut_initializer = (x, x, x)  # (state, input, memory)
+    ut_function = functools.partial(
+        universal_transformer_basic,
         hparams=hparams,
         ffn_unit=ffn_unit,
         attention_unit=attention_unit)
 
   elif hparams.recurrence_type == "highway":
-    rt_initializer = (x, x, x)  # (state, input, memory)
-    rt_function = functools.partial(
-        r_transformer_highway,
+    ut_initializer = (x, x, x)  # (state, input, memory)
+    ut_function = functools.partial(
+        universal_transformer_highway,
         hparams=hparams,
         ffn_unit=ffn_unit,
         attention_unit=attention_unit,
         pad_remover=pad_remover)
 
   elif hparams.recurrence_type == "skip":
-    rt_initializer = (x, x, x)  # (state, input, memory)
-    rt_function = functools.partial(
-        r_transformer_skip,
+    ut_initializer = (x, x, x)  # (state, input, memory)
+    ut_function = functools.partial(
+        universal_transformer_skip,
         hparams=hparams,
         ffn_unit=ffn_unit,
         attention_unit=attention_unit,
@@ -296,35 +317,35 @@ def get_rt_layer(x, hparams, ffn_unit, attention_unit, pad_remover=None):
     # filling the first slot with the original input
     memory = fill_memory_slot(memory_empty, x, 0)
 
-    rt_initializer = (x, x, memory)  # (state, input, memory)
-    rt_function = functools.partial(
-        r_transformer_depthwise_attention,
+    ut_initializer = (x, x, memory)  # (state, input, memory)
+    ut_function = functools.partial(
+        universal_transformer_depthwise_attention,
         hparams=hparams,
         ffn_unit=ffn_unit,
         attention_unit=attention_unit)
 
   elif hparams.recurrence_type == "rnn":
-    rt_initializer = (x, x, x)  # (state, input, memory)
-    rt_function = functools.partial(
-        r_transformer_rnn,
+    ut_initializer = (x, x, x)  # (state, input, memory)
+    ut_function = functools.partial(
+        universal_transformer_rnn,
         hparams=hparams,
         ffn_unit=ffn_unit,
         attention_unit=attention_unit,
         pad_remover=pad_remover)
 
   elif hparams.recurrence_type == "gru":
-    rt_initializer = (x, x, x)  # (state, input, memory)
-    rt_function = functools.partial(
-        r_transformer_gru,
+    ut_initializer = (x, x, x)  # (state, input, memory)
+    ut_function = functools.partial(
+        universal_transformer_gru,
         hparams=hparams,
         attention_unit=attention_unit,
         pad_remover=pad_remover)
 
   elif hparams.recurrence_type == "lstm":
     memory = tf.zeros(common_layers.shape_list(x))
-    rt_initializer = (x, x, memory)  # (state, input, memory)
-    rt_function = functools.partial(
-        r_transformer_lstm,
+    ut_initializer = (x, x, memory)  # (state, input, memory)
+    ut_function = functools.partial(
+        universal_transformer_lstm,
         hparams=hparams,
         attention_unit=attention_unit,
         pad_remover=pad_remover)
@@ -332,7 +353,7 @@ def get_rt_layer(x, hparams, ffn_unit, attention_unit, pad_remover=None):
   else:
     raise ValueError("Unknown recurrence type: %s" % hparams.recurrence_type)
 
-  return rt_function, rt_initializer
+  return ut_function, ut_initializer
 
 
 def transformer_encoder_ffn_unit(x,
@@ -372,6 +393,8 @@ def transformer_encoder_ffn_unit(x,
           common_layers.layer_preprocess(x, hparams),
           filter_size=hparams.filter_size,
           output_size=hparams.hidden_size,
+          first_kernel_size=(3, 1),
+          second_kernel_size=(5, 1),
           padding="SAME",
           nonpadding_mask=nonpadding_mask,
           dropout=hparams.relu_dropout)
@@ -424,7 +447,9 @@ def transformer_encoder_attention_unit(x,
   return x
 
 
-def transformer_decoder_ffn_unit(x, hparams, nonpadding_mask=None):
+def transformer_decoder_ffn_unit(x,
+                                 hparams,
+                                 nonpadding_mask=None):
   """Applies a feed-forward function which is parametrised for decoding.
 
   Args:
@@ -454,6 +479,8 @@ def transformer_decoder_ffn_unit(x, hparams, nonpadding_mask=None):
           common_layers.layer_preprocess(x, hparams),
           filter_size=hparams.filter_size,
           output_size=hparams.hidden_size,
+          first_kernel_size=(3, 1),
+          second_kernel_size=(5, 1),
           padding="LEFT",
           nonpadding_mask=nonpadding_mask,
           dropout=hparams.relu_dropout)
@@ -528,8 +555,11 @@ def transformer_decoder_attention_unit(x,
   return x
 
 
-def r_transformer_basic(layer_inputs, step, hparams, ffn_unit, attention_unit):
-  """Basic r_transformer.
+def universal_transformer_basic(layer_inputs,
+                                step, hparams,
+                                ffn_unit,
+                                attention_unit):
+  """Basic universal_transformer.
 
   This is in fact vanilla transformer in which weights are shared between
   layers. For some tasks, this simple idea brings a generalization that is not
@@ -556,13 +586,12 @@ def r_transformer_basic(layer_inputs, step, hparams, ffn_unit, attention_unit):
   return new_state, inputs, memory
 
 
-def r_transformer_highway(layer_inputs,
-                          step,
-                          hparams,
-                          ffn_unit,
-                          attention_unit,
-                          pad_remover=None):
-  """R_transformer with highway connection.
+def universal_transformer_highway(layer_inputs,
+                                  step, hparams,
+                                  ffn_unit,
+                                  attention_unit,
+                                  pad_remover=None):
+  """universal_transformer with highway connection.
 
 
   It transforms the state using attention and ffn and wrap this transformation
@@ -645,13 +674,13 @@ def r_transformer_highway(layer_inputs,
   return new_state, inputs, memory
 
 
-def r_transformer_skip(layer_inputs,
-                       step,
-                       hparams,
-                       ffn_unit,
-                       attention_unit,
-                       pad_remover=None):
-  """R_transformer with highway connection.
+def universal_transformer_skip(layer_inputs,
+                               step,
+                               hparams,
+                               ffn_unit,
+                               attention_unit,
+                               pad_remover=None):
+  """universal_transformer with highway connection.
 
 
   It transforms the state using attention and ffn and wrap this transformation
@@ -733,9 +762,11 @@ def r_transformer_skip(layer_inputs,
   return new_state, inputs, memory
 
 
-def r_transformer_depthwise_attention(layer_inputs, step, hparams, ffn_unit,
-                                      attention_unit):
-  """R_transformer with depth-wise attention.
+def universal_transformer_depthwise_attention(layer_inputs,
+                                              step, hparams,
+                                              ffn_unit,
+                                              attention_unit):
+  """universal_transformer with depth-wise attention.
 
   It uses an attention mechanism-flipped vertically-
   over all the states from previous steps to generate the new_state.
@@ -788,15 +819,15 @@ def r_transformer_depthwise_attention(layer_inputs, step, hparams, ffn_unit,
   return new_state, inputs, memory
 
 
-def r_transformer_rnn(layer_inputs,
-                      step,
-                      hparams,
-                      ffn_unit,
-                      attention_unit,
-                      pad_remover=None):
-  """The RT layer which models recurencey similar to basic RNN cell.
+def universal_transformer_rnn(layer_inputs,
+                              step,
+                              hparams,
+                              ffn_unit,
+                              attention_unit,
+                              pad_remover=None):
+  """The UT layer which models recurencey similar to basic RNN cell.
 
-    It's an R-transformer with an RNN applied over the stats on depth.
+    It's an U-Transformer with an RNN applied over the stats on depth.
 
   Args:
     layer_inputs:
@@ -860,11 +891,11 @@ def r_transformer_rnn(layer_inputs,
   return new_state, inputs, memory
 
 
-def r_transformer_gru(layer_inputs,
-                      step,
-                      hparams,
-                      attention_unit,
-                      pad_remover=None):
+def universal_transformer_gru(layer_inputs,
+                              step,
+                              hparams,
+                              attention_unit,
+                              pad_remover=None):
   """The RT layer which models recurencey similar to GRU cell.
 
     It's an R-transformer with a gru applied over the stats on depth.
@@ -935,12 +966,12 @@ def r_transformer_gru(layer_inputs,
   return new_state, inputs, memory
 
 
-def r_transformer_lstm(layer_inputs,
-                       step,
-                       hparams,
-                       attention_unit,
-                       pad_remover=None):
-  """The RT layer which models recurencey similar to GRU cell.
+def universal_transformer_lstm(layer_inputs,
+                               step,
+                               hparams,
+                               attention_unit,
+                               pad_remover=None):
+  """The UT layer which models recurencey similar to GRU cell.
 
   It's an R-transformer with a gru applied over the stats on depth.
   based on LSTM paper: https://arxiv.org/pdf/1409.2329.pdf
@@ -1018,7 +1049,7 @@ def r_transformer_lstm(layer_inputs,
   return new_state, inputs, memory
 
 
-def r_transformer_act(x, hparams, ffn_unit, attention_unit):
+def universal_transformer_act(x, hparams, ffn_unit, attention_unit):
   """ACT based models.
 
   Implementations of all act models are based on craffel@'s cl/160711592.
@@ -1036,25 +1067,29 @@ def r_transformer_act(x, hparams, ffn_unit, attention_unit):
     ValueError: Unknown act type
 
   """
-
+  # TODO(dehghani): Use pad_remover for the act computations.
   if hparams.act_type == "basic":
-    return r_transformer_act_basic(x, hparams, ffn_unit, attention_unit)
+    return universal_transformer_act_basic(
+        x, hparams, ffn_unit, attention_unit)
 
   elif hparams.act_type == "accumulated":
-    return r_transformer_act_accumulated(x, hparams, ffn_unit, attention_unit)
+    return universal_transformer_act_accumulated(
+        x, hparams, ffn_unit, attention_unit)
 
   elif hparams.act_type == "global":
-    return r_transformer_act_global(x, hparams, ffn_unit, attention_unit)
+    return universal_transformer_act_global(
+        x, hparams, ffn_unit, attention_unit)
 
   elif hparams.act_type == "random":
-    return r_transformer_act_random(x, hparams, ffn_unit, attention_unit)
+    return universal_transformer_act_random(
+        x, hparams, ffn_unit, attention_unit)
 
   else:
     raise ValueError("Unknown act type: %s" % hparams.act_type)
 
 
-def r_transformer_act_basic(x, hparams, ffn_unit, attention_unit):
-  """Basic r_transformer with ACT based on remainder-distribution ACT.
+def universal_transformer_act_basic(x, hparams, ffn_unit, attention_unit):
+  """Basic universal_transformer with ACT based on remainder-distribution ACT.
 
   Args:
     x: input
@@ -1097,7 +1132,7 @@ def r_transformer_act_basic(x, hparams, ffn_unit, attention_unit):
   previous_state = tf.zeros_like(state, name="previous_state")
   step = tf.constant(0, dtype=tf.int32)
 
-  def rt_function(state, step, halting_probability, remainders, n_updates,
+  def ut_function(state, step, halting_probability, remainders, n_updates,
                   previous_state):
     """implements act (position-wise halting).
 
@@ -1193,7 +1228,7 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
 
   # Do while loop iterations until predicate above is false.
   (_, _, _, remainder, n_updates, new_state) = tf.while_loop(
-      should_continue, rt_function,
+      should_continue, ut_function,
       (state, step, halting_probability, remainders, n_updates, previous_state))
 
   ponder_times = n_updates
@@ -1204,8 +1239,8 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
   return new_state, (ponder_times, remainders)
 
 
-def r_transformer_act_accumulated(x, hparams, ffn_unit, attention_unit):
-  """The RTAct layer where the final state is accumulation of all states.
+def universal_transformer_act_accumulated(x, hparams, ffn_unit, attention_unit):
+  """The UTAct layer where the final state is the accumulation of all states.
 
     (similar to the main ACT paper: --> check the issue of differentiability)
 
@@ -1249,7 +1284,7 @@ def r_transformer_act_accumulated(x, hparams, ffn_unit, attention_unit):
   accumulated_state = tf.zeros_like(state, name="previous_state")
   step = tf.constant(0, dtype=tf.int32)
 
-  def rt_function(state, step, halting_probability, remainders, n_updates,
+  def ut_function(state, step, halting_probability, remainders, n_updates,
                   accumulated_state):
     """Position-wise act.
 
@@ -1344,7 +1379,7 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
 
   # Do while loop iterations until predicate above is false.
   (_, _, _, remainder, n_updates, accumulated_state) = tf.while_loop(
-      should_continue, rt_function, (state, step, halting_probability,
+      should_continue, ut_function, (state, step, halting_probability,
                                      remainders, n_updates, accumulated_state))
 
   ponder_times = n_updates
@@ -1355,8 +1390,8 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
   return accumulated_state, (ponder_times, remainders)
 
 
-def r_transformer_act_global(x, hparams, ffn_unit, attention_unit):
-  """The RTAct  with global halting probability (not position-wise).
+def universal_transformer_act_global(x, hparams, ffn_unit, attention_unit):
+  """The UTAct  with global halting probability (not position-wise).
 
   Args:
     x: input
@@ -1385,7 +1420,7 @@ def r_transformer_act_global(x, hparams, ffn_unit, attention_unit):
   previous_state = tf.zeros_like(state, name="previous_state")
   step = tf.constant(0, dtype=tf.int32)
 
-  def rt_function(state, step, halting_probability, remainders, n_updates,
+  def ut_function(state, step, halting_probability, remainders, n_updates,
                   previous_state):
     """implements act (global halting).
 
@@ -1486,7 +1521,7 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
 
   # Do while loop iterations until predicate above is false.
   (_, _, _, remainder, n_updates, new_state) = tf.while_loop(
-      should_continue, rt_function,
+      should_continue, ut_function,
       (state, step, halting_probability, remainders, n_updates, previous_state))
 
   ponder_times = n_updates
@@ -1497,8 +1532,8 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
   return new_state, (ponder_times, remainders)
 
 
-def r_transformer_act_random(x, hparams, ffn_unit, attention_unit):
-  """r_transformer with ACT with random halting probability.
+def universal_transformer_act_random(x, hparams, ffn_unit, attention_unit):
+  """universal_transformer with ACT with random halting probability.
 
   Args:
     x: input
@@ -1541,7 +1576,7 @@ def r_transformer_act_random(x, hparams, ffn_unit, attention_unit):
   previous_state = tf.zeros_like(state, name="previous_state")
   step = tf.constant(0, dtype=tf.int32)
 
-  def rt_function(state, step, halting_probability, remainders, n_updates,
+  def ut_function(state, step, halting_probability, remainders, n_updates,
                   previous_state):
     """Implements act (position-wise halting).
 
@@ -1633,7 +1668,7 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
 
   # Do while loop iterations until predicate above is false.
   (_, _, _, remainder, n_updates, new_state) = tf.while_loop(
-      should_continue, rt_function,
+      should_continue, ut_function,
       (state, step, halting_probability, remainders, n_updates, previous_state))
 
   ponder_times = n_updates

From 95485e49e2188b5c2ee65759bb0e3281114c2fda Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 31 May 2018 07:36:30 -0700
Subject: [PATCH 0035/2720] internal merge of PR #837

PiperOrigin-RevId: 198719318
---
 tensor2tensor/models/__init__.py                           | 2 +-
 tensor2tensor/models/research/universal_transformer.py     | 7 ++++---
 .../models/research/universal_transformer_test.py          | 2 --
 .../models/research/universal_transformer_util.py          | 2 --
 4 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 9b49dff70..d0d774604 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -44,7 +44,6 @@
 from tensor2tensor.models.research import lm_experiments
 from tensor2tensor.models.research import multimodel
 from tensor2tensor.models.research import next_frame
-from tensor2tensor.models.research import universal_transformer
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research import super_lm
 from tensor2tensor.models.research import transformer_moe
@@ -52,4 +51,5 @@
 from tensor2tensor.models.research import transformer_sketch
 from tensor2tensor.models.research import transformer_symshard
 from tensor2tensor.models.research import transformer_vae
+from tensor2tensor.models.research import universal_transformer
 # pylint: enable=unused-import
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 51050d8cf..c9ee89a69 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -53,6 +53,7 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
       hparams: hyperparmeters for model.
       features: optionally pass the entire features dictionary as well.
         This is needed now for "packed" datasets.
+      losses: Unused.
 
     Returns:
       Tuple of:
@@ -84,7 +85,6 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
 
     return encoder_output, encoder_decoder_attention_bias, encoder_extra_output
 
-
   def decode(self,
              decoder_input,
              encoder_output,
@@ -110,7 +110,9 @@ def decode(self,
       decoder_self_attention_bias: Bias and mask weights for decoder
         self-attention. [batch_size, decoder_length]
       hparams: hyperparmeters for model.
+      cache: Unimplemented.
       nonpadding: optional Tensor with shape [batch_size, decoder_length]
+      losses: Unused.
 
     Returns:
        Tuple of:
@@ -275,6 +277,7 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
       hparams: hyperparmeters for model.
       features: optionally pass the entire features dictionary as well.
         This is needed now for "packed" datasets.
+      losses: Unused.
 
     Returns:
       Tuple of:
@@ -881,8 +884,6 @@ def universal_transformer_base_range(rhp):
   rhp.set_float("weight_decay", 0.0, 2.0)
 
 
-
-
 @registry.register_ranged_hparams
 def adaptive_universal_transformer_base_range(rhp):
   """Small range of hyperparameters."""
diff --git a/tensor2tensor/models/research/universal_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
index 6c0922da2..da228c2d0 100644
--- a/tensor2tensor/models/research/universal_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -18,8 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-# Dependency imports
-
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index b5ffbc07b..231cf9a60 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -48,8 +48,6 @@
 import copy
 import functools
 
-# Dependency imports
-
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer

From 8b312b55e47927d54860f88b26fb5f1c0727049b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 31 May 2018 09:53:27 -0700
Subject: [PATCH 0036/2720] Fixed the padding issue for "targets" as well.

PiperOrigin-RevId: 198735543
---
 tensor2tensor/layers/common_layers.py | 11 +++++++++++
 tensor2tensor/models/lstm.py          | 23 +++++++++++++++--------
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 54fd9623c..8fc7a94ae 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1288,6 +1288,17 @@ def mask_from_embedding(emb):
   return weights_nonzero(tf.reduce_sum(tf.abs(emb), axis=3, keepdims=True))
 
 
+def length_from_embedding(emb):
+  """Compute the length of each sequence in the batch.
+
+  Args:
+    emb: a sequence embedding Tensor with shape [batch, max_time, 1, depth].
+  Returns:
+    a Tensor with shape [batch].
+  """
+  return tf.reduce_sum(mask_from_embedding(emb), [1, 2, 3])
+
+
 def mask_leq(target_length, source_length):
   """A mask with 1.0 wherever source_pos <= target_pos and 0.0 elsewhere.
 
diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 1c6c21138..0b933b6ce 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -27,7 +27,8 @@
 import tensorflow as tf
 
 
-def lstm(inputs, hparams, train, name, initial_state=None, lengths=None):
+def lstm(inputs, hparams, train, name, initial_state=None,
+         sequence_length=None):
   """Run LSTM cell on inputs, assuming they are [batch x time x size]."""
 
   def dropout_lstm_cell():
@@ -41,13 +42,14 @@ def dropout_lstm_cell():
         tf.contrib.rnn.MultiRNNCell(layers),
         inputs,
         initial_state=initial_state,
-        sequence_length=lengths,
+        sequence_length=sequence_length,
         dtype=tf.float32,
         time_major=False)
 
 
 def lstm_attention_decoder(inputs, hparams, train, name, initial_state,
-                           encoder_outputs, lengths=None):
+                           encoder_outputs, encoder_output_length=None,
+                           decoder_input_length=None):
   """Run LSTM cell with attention on inputs of shape [batch x time x size]."""
 
   def dropout_lstm_cell():
@@ -64,7 +66,8 @@ def dropout_lstm_cell():
     raise ValueError("Unknown hparams.attention_mechanism = %s, must be "
                      "luong or bahdanau." % hparams.attention_mechanism)
   attention_mechanism = attention_mechanism_class(
-      hparams.hidden_size, encoder_outputs, memory_sequence_length=lengths)
+      hparams.hidden_size, encoder_outputs,
+      memory_sequence_length=encoder_output_length)
 
   cell = tf.contrib.seq2seq.AttentionWrapper(
       tf.nn.rnn_cell.MultiRNNCell(layers),
@@ -82,6 +85,7 @@ def dropout_lstm_cell():
         cell,
         inputs,
         initial_state=initial_state,
+        sequence_length=decoder_input_length,
         dtype=tf.float32,
         time_major=False)
 
@@ -120,18 +124,21 @@ def lstm_seq2seq_internal_attention(inputs, targets, hparams, train):
     # This is a temporary fix for varying-length sequences within in a batch.
     # A more complete fix should pass a length tensor from outside so that
     # all the lstm variants can use it.
-    lengths = tf.reduce_sum(
-        common_layers.mask_from_embedding(inputs), [1, 2, 3])
+    inputs_length = common_layers.length_from_embedding(inputs)
     # Flatten inputs.
     inputs = common_layers.flatten4d3d(inputs)
     # LSTM encoder.
     encoder_outputs, final_encoder_state = lstm(
-        inputs, hparams, train, "encoder", lengths=lengths)
+        inputs, hparams, train, "encoder", sequence_length=inputs_length)
     # LSTM decoder with attention
     shifted_targets = common_layers.shift_right(targets)
+    # Add 1 to account for the padding added to the left from shift_right
+    targets_length = common_layers.length_from_embedding(shifted_targets) + 1
     decoder_outputs, _ = lstm_attention_decoder(
         common_layers.flatten4d3d(shifted_targets), hparams, train, "decoder",
-        final_encoder_state, encoder_outputs, lengths=lengths)
+        final_encoder_state, encoder_outputs,
+        encoder_output_length=inputs_length,
+        decoder_input_length=targets_length)
     return tf.expand_dims(decoder_outputs, axis=2)
 
 

From 058b5fc0be35c260d70acf11b5b86d7e6b29daa5 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 31 May 2018 14:55:38 -0700
Subject: [PATCH 0037/2720] s/log_registry/maybe_log_registry_and_exit to make
 its intent clearer

PiperOrigin-RevId: 198784505
---
 tensor2tensor/bin/t2t_distill.py | 6 +++---
 tensor2tensor/bin/t2t_trainer.py | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/bin/t2t_distill.py b/tensor2tensor/bin/t2t_distill.py
index ef9109373..91f9caece 100644
--- a/tensor2tensor/bin/t2t_distill.py
+++ b/tensor2tensor/bin/t2t_distill.py
@@ -40,7 +40,7 @@ def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
   trainer_lib.set_random_seed(FLAGS.random_seed)
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
-  t2t_trainer.log_registry()
+  t2t_trainer.maybe_log_registry_and_exit()
 
   if FLAGS.cloud_mlengine:
     cloud_mlengine.launch()
@@ -94,7 +94,7 @@ def create_teacher_experiment(run_config, hparams, argv):
   tf.logging.set_verbosity(tf.logging.INFO)
   trainer_lib.set_random_seed(FLAGS.random_seed)
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
-  t2t_trainer.log_registry()
+  t2t_trainer.maybe_log_registry_and_exit()
 
   if FLAGS.cloud_mlengine:
     return cloud_mlengine.launch()
@@ -121,7 +121,7 @@ def create_student_experiment(run_config, hparams, argv):
   tf.logging.set_verbosity(tf.logging.INFO)
   trainer_lib.set_random_seed(FLAGS.random_seed)
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
-  t2t_trainer.log_registry()
+  t2t_trainer.maybe_log_registry_and_exit()
 
   if FLAGS.cloud_mlengine:
     return cloud_mlengine.launch()
diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index fc12cce15..f82f170e0 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -247,7 +247,7 @@ def profile_context():
     yield
 
 
-def log_registry():
+def maybe_log_registry_and_exit():
   if FLAGS.registry_help:
     tf.logging.info(registry.help_string())
     sys.exit(0)
@@ -328,7 +328,7 @@ def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
   trainer_lib.set_random_seed(FLAGS.random_seed)
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
-  log_registry()
+  maybe_log_registry_and_exit()
 
 
   if FLAGS.cloud_mlengine:

From 6c632b1ae6f90e098b20fe8b8e471a583b43e380 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Thu, 31 May 2018 17:24:05 -0700
Subject: [PATCH 0038/2720] Move vae_layers to latent_layers

PiperOrigin-RevId: 198805733
---
 tensor2tensor/layers/latent_layers.py | 527 ++++++++++++++++++++++++++
 1 file changed, 527 insertions(+)
 create mode 100644 tensor2tensor/layers/latent_layers.py

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
new file mode 100644
index 000000000..0cd5bd712
--- /dev/null
+++ b/tensor2tensor/layers/latent_layers.py
@@ -0,0 +1,527 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Utils for latent variable models."""
+
+from six.moves import range  # pylint: disable=redefined-builtin
+
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_image_attention as cia
+from tensor2tensor.layers import common_layers
+from tensor2tensor.models import transformer
+
+import tensorflow as tf
+
+DO_SUMMARIES = True
+
+
+class Latent(object):
+  DISCRETE = "discrete"
+  DENSE = "dense"
+
+  @staticmethod
+  def get_choices():
+    return [
+        Latent.DISCRETE,
+        Latent.DENSE,
+    ]
+
+
+def add_learned_positional_embeddings(x, hparams):
+  pos = tf.get_variable("pos",
+                        [1, hparams.img_len*hparams.img_len,
+                         1, hparams.hidden_size])
+  pos = pos[:, :common_layers.shape_list(x)[1], :, :]
+  x = tf.expand_dims(x, axis=2)
+  x += pos
+  return x
+
+
+def attend(x, source, hparams, name):
+  """Attend function."""
+  with tf.variable_scope(name):
+    # x = tf.squeeze(x, axis=2)
+    x, xshape, _ = cia.maybe_reshape_4d_to_3d(x)
+    if len(source.get_shape()) > 3:
+      source = tf.squeeze(source, axis=2)
+    source = common_attention.add_timing_signal_1d(source)
+    y = common_attention.multihead_attention(
+        common_layers.layer_preprocess(x, hparams),
+        source,
+        None,
+        hparams.attention_key_channels or hparams.hidden_size,
+        hparams.attention_value_channels or hparams.hidden_size,
+        hparams.hidden_size, hparams.num_heads,
+        hparams.attention_dropout)
+    res = common_layers.layer_postprocess(x, y, hparams)
+    return tf.reshape(res, xshape)
+
+
+def multinomial_sample(x, vocab_size, temperature):
+  """Multinomial sampling from a n-dimensional tensor."""
+  if temperature > 0:
+    samples = tf.multinomial(tf.reshape(x, [-1, vocab_size]) / temperature, 1)
+  else:
+    samples = tf.argmax(x, axis=-1)
+  reshaped_samples = tf.reshape(samples, common_layers.shape_list(x)[:-1])
+  return reshaped_samples
+
+
+def ae_latent_softmax(latents_pred, latents_discrete, hparams):
+  """Latent prediction and loss."""
+  vocab_size = 2 ** hparams.z_size
+  if hparams.num_decode_blocks < 2:
+    with tf.variable_scope("extra_logits"):
+      latents_logits = tf.layers.dense(latents_pred, vocab_size,
+                                       name="extra_logits")
+      if hparams.logit_normalization:
+        latents_logits *= tf.rsqrt(1e-8 +
+                                   tf.reduce_mean(tf.square(latents_logits)))
+
+      loss = None
+      if latents_discrete is not None:
+        if hparams.soft_em:
+          # latents_discrete is actually one-hot of multinomial samples
+          assert hparams.num_decode_blocks == 1
+          loss = tf.nn.softmax_cross_entropy_with_logits(
+              labels=latents_discrete, logits=latents_logits)
+        else:
+          loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
+              labels=latents_discrete, logits=latents_logits)
+      sample = multinomial_sample(latents_logits, vocab_size,
+                                  hparams.sampling_temp)
+      return sample, loss
+
+  # Multi-block case.
+  block_vocab_size = 2**(hparams.z_size // hparams.num_decode_blocks)
+  latents_logits = [
+      tf.layers.dense(
+          latents_pred, block_vocab_size, name="extra_logits_%d" % i)
+      for i in range(hparams.num_decode_blocks)
+  ]
+  loss = None
+  if latents_discrete is not None:
+    losses = []
+    for i in range(hparams.num_decode_blocks):
+      d = tf.floormod(tf.floordiv(latents_discrete,
+                                  block_vocab_size**i), block_vocab_size)
+      losses.append(tf.nn.sparse_softmax_cross_entropy_with_logits(
+          labels=d, logits=latents_logits[i]))
+    loss = sum(losses)
+  samples = [multinomial_sample(l, block_vocab_size, hparams.sampling_temp)
+             for l in latents_logits]
+  sample = sum([s * block_vocab_size**i for i, s in enumerate(samples)])
+  return sample, loss
+
+
+def residual_block_layer(inputs, hparams):
+  """Residual block over inputs.
+
+  Runs a residual block consisting of
+    conv: kernel_size x kernel_size
+    conv: 1x1
+    dropout, add and normalize according to hparams.layer_postprocess_sequence.
+
+  Args:
+    inputs: Tensor of shape [batch_size, height, width, hidden_dim].
+    hparams: Dict, hyperparameters.
+
+  Returns:
+    x: Tensor of shape [batch_size, height, width, hidden_dim]
+  """
+  kernel = (hparams.res_kernel_size, hparams.res_kernel_size)
+  x = inputs
+  for i in range(hparams.num_res_layers):
+    with tf.variable_scope("res_conv_%d" % i):
+      # kernel_size x kernel_size conv block
+      y = common_layers.conv_block(
+          common_layers.layer_norm(x, hparams.hidden_size, name="lnorm"),
+          hparams.hidden_size, [((1, 1), kernel)],
+          strides=(1, 1),
+          padding="SAME",
+          name="residual_conv")
+      # 1x1 conv block
+      y = common_layers.conv_block(
+          y,
+          hparams.hidden_size, [((1, 1), (1, 1))],
+          strides=(1, 1),
+          padding="SAME",
+          name="residual_dense")
+      x = common_layers.layer_postprocess(x, y, hparams)
+  return x
+
+
+def compress_encoder(inputs, hparams,
+                     strides=(2, 2),
+                     kernel=(3, 3),
+                     name="compress"):
+  """Encoder that compresses inputs to length/2**num_compress_steps.
+
+  Args:
+    inputs: Tensor of shape [batch, height, width, hidden_dim].
+    hparams: Dict, hyperparameters.
+    strides: Tuple, strides for conv block.
+    kernel: Tuple, kernel window size for conv block.
+    name: string, variable scope.
+
+  Returns:
+    x: Tensor of shape [batch, height*width/2**(compress_steps), hidden_dim].
+  """
+  with tf.variable_scope(name):
+    x = inputs
+    # Compress conv layers with strides and kernels as passed to the function.
+    for i in range(hparams.num_compress_steps // 2):
+      with tf.variable_scope("compress_conv_%d" % i):
+        y = common_layers.conv_block(
+            common_layers.layer_norm(
+                x, hparams.hidden_size, name="lnorm"),
+            hparams.hidden_size, [((1, 1), kernel)],
+            strides=strides,
+            padding="SAME",
+            name="compress_conv_%d" % i)
+        y = tf.nn.dropout(y, 1.0 - hparams.dropout)
+        x = y
+
+    # Residual blocks.
+    x = residual_block_layer(x, hparams)
+
+    # If using multiple copies of latents, blow up the hidden size and then
+    # reshape to increase by num_latents.
+    shape_x = common_layers.shape_list(x)
+    x = tf.layers.dense(x, hparams.num_latents*hparams.hidden_size,
+                        name=name + "_dense")
+    new_shape = [shape_x[0], shape_x[1] * shape_x[2]*hparams.num_latents,
+                 hparams.hidden_size]
+    return tf.reshape(x, new_shape)
+
+
+def compress_encoder_2d(x, hparams, name):
+  """Encoder that compresses inputs to height*width/2**num_compress_steps.
+
+  Args:
+    x: Tensor of shape [batch, height, width, hidden_dim].
+    hparams: Dict, hyperparameters.
+    name: string, variable scope.
+
+  Returns:
+    x: Tensor of shape [batch, height*width/2**(compress_steps), hidden_dim].
+  """
+  return compress_encoder(x, hparams,
+                          strides=(2, 2),
+                          kernel=(hparams.kernel_size, hparams.kernel_size),
+                          name=name)
+
+
+def compress_encoder_1d(x, hparams, name):
+  """Encoder that compresses inputs to length/2**num_compress_steps.
+
+  Args:
+    x: Tensor of shape [batch, length, hidden_dim].
+    hparams: Dict, hyperparameters.
+    name: string, variable scope.
+
+  Returns:
+    x: Tensor of shape [batch, length/2**(compress_steps), hidden_dim].
+  """
+  x = tf.expand_dims(x, axis=2)
+  return compress_encoder(x, hparams,
+                          strides=(2, 1),
+                          kernel=(hparams.kernel_size, 1),
+                          name=name)
+
+
+def decompress_decoder(inputs, hparams,
+                       strides=(2, 2),
+                       kernel=(3, 3),
+                       name="decompress"):
+  """Encoder that compresses inputs to length/2**num_compress_steps.
+
+  Args:
+    inputs: Tensor of shape [batch, compress_height, compress_width, hidden_dim]
+    hparams: Dict, hyperparameters.
+    strides: Tuple, strides for conv block.
+    kernel: Tuple, kernel window size for conv block.
+    name: string, variable scope.
+
+  Returns:
+    x: Tensor of shape [batch, height, width, hidden_dim].
+  """
+  with tf.variable_scope(name):
+    x = inputs
+    # Reshape?
+    x = tf.layers.dense(x, hparams.hidden_size, name=name + "_dense")
+    # Residual blocks.
+    x = residual_block_layer(x, hparams)
+
+    # Decompress conv layers with strides and kernels as passed to the function.
+    for i in range(hparams.num_compress_steps // 2):
+      j = hparams.num_compress_steps // 2 - i - 1
+      with tf.variable_scope(name + "_%d" % j):
+        y = tf.layers.conv2d_transpose(
+            x,
+            hparams.hidden_size,
+            kernel,
+            strides=strides,
+            padding="SAME",
+            activation=tf.nn.relu if i > 0 else None,
+            name="decompress_conv")
+        x = y
+    return x
+
+
+def decompress_decoder_2d(x, hparams, name):
+  """Dencoder that decompresses x to length height*width.
+
+  Args:
+    x: Tensor of shape [batch, compress_height, compress_width, hidden_dim].
+    hparams: Dict, hyperparameters.
+    name: string, variable scope.
+
+  Returns:
+    x: Tensor of shape [batch, height, width, hidden_dim].
+  """
+  return decompress_decoder(x, hparams,
+                            strides=(2, 2),
+                            kernel=(hparams.kernel_size, hparams.kernel_size),
+                            name=name)
+
+
+def decompress_decoder_1d(x, hparams, name):
+  """Dencoder that decompresses x to original target length.
+
+  Args:
+    x: Tensor of shape [batch, compress_length, hidden_dim].
+    hparams: Dict, hyperparameters.
+    name: string, variable scope.
+
+  Returns:
+    x: Tensor of shape [batch, length, hidden_dim].
+  """
+  x = tf.expand_dims(x, axis=2)
+  output = decompress_decoder(x, hparams,
+                              strides=(2, 1),
+                              kernel=(hparams.kernel_size, 1),
+                              name=name)
+  return tf.squeeze(output, axis=2)
+
+
+def transformer_text_encoder(inputs, space_id,
+                             hparams, name="transformer_text_enc"):
+  """Transformer text encoder."""
+  with tf.variable_scope(name):
+    x = common_layers.flatten4d3d(inputs)
+    (encoder_input, encoder_self_attention_bias,
+     ed) = transformer.transformer_prepare_encoder(x, space_id, hparams)
+    encoder_input = tf.nn.dropout(encoder_input, 1.0 - hparams.dropout)
+    return transformer.transformer_encoder(
+        encoder_input, encoder_self_attention_bias, hparams), ed
+
+
+def transformer_image_decoder(encoder_output,
+                              ed_attention_bias,
+                              targets,
+                              hparams,
+                              name="transformer_dec"):
+  """Original Transformer decoder."""
+  with tf.variable_scope(name):
+    batch_size = common_layers.shape_list(targets)[0]
+    # Reshape targets as b, 32, 32, 3*hidden size].
+    targets = tf.reshape(targets, [
+        batch_size, hparams.img_len, hparams.img_len,
+        hparams.num_channels*hparams.hidden_size])
+
+    # Prepare decoder inputs and bias. This also shifts targets and adds 2D
+    # position embeddings to target.
+    decoder_input, _, _ = cia.prepare_decoder(targets, hparams)
+    decoder_output = cia.transformer_decoder_layers(
+        decoder_input,
+        encoder_output,
+        hparams.num_decoder_layers or hparams.num_hidden_layers,
+        hparams,
+        attention_type=hparams.dec_attention_type,
+        encoder_decoder_attention_bias=ed_attention_bias,
+        name="decoder")
+    decoder_output_shape = common_layers.shape_list(decoder_output)
+    decoder_output = tf.reshape(decoder_output, [
+        decoder_output_shape[0],
+        hparams.img_len, hparams.img_len*hparams.num_channels,
+        hparams.hidden_size])
+    return decoder_output
+
+
+def transformer_latent_decoder(encoder_output,
+                               ed_attention_bias,
+                               targets,
+                               hparams,
+                               name="transformer_latent_dec"):
+  """Original Transformer decoder."""
+  with tf.variable_scope(name):
+    batch_size = common_layers.shape_list(targets)[0]
+    compress_ratio = 2**(hparams.num_compress_steps // 2)
+    # Reshape targets as b, 32, 32, 3*hidden size].
+    targets = tf.reshape(targets, [
+        batch_size, hparams.img_len / compress_ratio,
+        (hparams.img_len*hparams.num_latents) / compress_ratio,
+        hparams.hidden_size
+    ])
+
+    # Prepare decoder inputs and bias.
+    decoder_input, _, _ = cia.prepare_decoder(targets, hparams)
+    # hparams.num_channels = 3
+    decoder_output = cia.transformer_decoder_layers(
+        decoder_input,
+        encoder_output,
+        hparams.num_latent_layers or hparams.num_hidden_layers,
+        hparams,
+        attention_type=hparams.latent_attention_type,
+        encoder_decoder_attention_bias=ed_attention_bias,
+        name="decoder")
+    decoder_output_shape = common_layers.shape_list(decoder_output)
+    decoder_output = tf.reshape(decoder_output, [
+        decoder_output_shape[0],
+        (hparams.img_len * hparams.img_len *
+         hparams.num_latents) / (2**hparams.num_compress_steps),
+        hparams.hidden_size
+    ])
+    return decoder_output
+
+
+def bottleneck_layer(targets_c,
+                     hparams,
+                     name="bottlneck_d"):
+  """Compute latents from compressed targets."""
+  # TODO(nikip): Condense hparams by removing options we don't use.
+  latents_dense, latents_discrete, extra_loss, embed_func = (
+      hparams.bottleneck(
+          x=targets_c,
+          filter_size=hparams.compress_filter_size,
+          name=name,
+          mode=hparams.mode))
+  if DO_SUMMARIES:
+    tf.summary.histogram("b0", tf.reshape(latents_discrete, [-1]))
+  return latents_dense, latents_discrete, extra_loss, embed_func
+
+
+def latent_prediction_model(
+    inputs, ed_attention_bias,
+    latents_discrete, embed,
+    hparams, name="latent_pred"):
+  """Transformer based latent prediction model."""
+  with tf.variable_scope(name):
+    if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+      latents_pred = transformer_latent_decoder(
+          inputs, ed_attention_bias,
+          tf.stop_gradient(embed(latents_discrete)), hparams, name + "_extra")
+      _, latent_pred_loss = ae_latent_softmax(
+          latents_pred, tf.stop_gradient(latents_discrete), hparams)
+  return latents_pred, latent_pred_loss
+
+
+def transformer_autoencoder(inputs,
+                            targets,
+                            target_space,
+                            hparams,
+                            cache=None,
+                            predict_mask=1.0):
+  """AE Transformer, main step used for training."""
+  # Define losses
+  losses = {"extra": tf.constant(0.0), "latent_pred": tf.constant(0.0)}
+
+  # Reshape image targets as 4d tensor.
+  original_targets_shape = common_layers.shape_list(targets)
+  if len(original_targets_shape) == 4:
+    compress_fn = compress_encoder_2d
+    decompress_fn = decompress_decoder_2d
+  else:
+    compress_fn = compress_encoder_1d
+    decompress_fn = decompress_decoder_1d
+
+  # Encoder decoder attention bias.
+  ed_attention_bias = None
+
+  # Input Encoder if present.
+  if inputs is not None:
+    inputs = common_layers.flatten4d3d(inputs)
+    inputs, ed_attention_bias = transformer_text_encoder(
+        inputs, target_space, hparams, "input_enc")
+
+  # Encode targets to compute targets compressed.
+  targets_c = compress_fn(targets, hparams, "compress")
+  targets, _, _ = cia.maybe_reshape_4d_to_3d(targets)
+
+  # Following code creates an exponentially decaying variable based on which
+  # we rescale the los values.
+  batch_size = common_layers.shape_list(targets_c)[0]
+  pc = common_layers.inverse_exp_decay(hparams.startup_steps)
+  pc = pc if hparams.mode == tf.estimator.ModeKeys.TRAIN else 1.0
+  cond = tf.less(tf.random_uniform([batch_size]), pc)
+
+  # TODO(lukaszkaiser): return extra losses batchwise, multiply before mean.
+  # Call bottleneck layer to get the latents.
+  # Returns embedded latents, discrete latents, loss and the embedding function.
+  latents_dense, latents_discrete, extra_loss, embed = (
+      bottleneck_layer(targets_c, hparams))
+  extra_loss = tf.reduce_mean(extra_loss) * tf.to_float(cond)
+
+  # Call the autoregressive latent prediction model.
+  _, latents_pred_loss = latent_prediction_model(
+      targets_c, ed_attention_bias, latents_discrete,
+      embed, hparams, name="latent_pred")
+  latents_pred_loss = tf.reduce_mean(latents_pred_loss) * tf.to_float(cond)
+
+  # Assign latent loss
+  losses["latent_pred"] = latents_pred_loss
+  losses["extra_loss"] = extra_loss
+
+  latents_decoder = latents_dense
+  if len(original_targets_shape) == 4:
+    cmp_img_len = hparams.img_len / (2**(hparams.num_compress_steps // 2))
+    latents_decoder = tf.reshape(
+        latents_decoder,
+        [batch_size, cmp_img_len, cmp_img_len,
+         hparams.num_latents*hparams.hidden_size])
+
+  # Decompress either using 1D or 2D upconvs.
+  latents_decoder = decompress_fn(latents_decoder, hparams, name="decompress")
+  # if we're operating in 2d space on images, then we're assuming that the
+  # last dimension will not be a multiple of channels
+  latents_decoder = tf.reshape(
+      latents_decoder,
+      shape=[-1, hparams.img_len, hparams.img_len, hparams.hidden_size])
+
+  if hparams.use_gold_targets:
+    latents_decoder, _, _ = cia.maybe_reshape_4d_to_3d(latents_decoder)
+    masking = common_layers.inverse_exp_decay(hparams.mask_startup_steps)
+    if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+      masking = predict_mask
+    mask = tf.less(masking, tf.random_uniform(
+        common_layers.shape_list(targets)[:-1]))
+    mask = tf.expand_dims(tf.to_float(mask), 2)
+    targets = mask * targets + (1.0 - mask) * latents_decoder
+  else:
+    targets = latents_decoder
+  # reshape back to 4d here
+  targets = tf.reshape(targets, original_targets_shape)
+  if hparams.decode_autoregressive:
+    # Transformer decoder, that goes from inputs->targets
+    res = transformer_image_decoder(inputs, ed_attention_bias,
+                                    targets, hparams, "decoder")
+  else:
+    res = targets
+
+  # We'll start training the extra model of latents after mask_startup_steps.
+  latent_time = tf.less(hparams.mask_startup_steps,
+                        tf.to_int32(tf.train.get_global_step()))
+  losses["latent_pred"] *= tf.to_float(latent_time)
+  return res, losses, cache

From 9bd6b6bb960601fbfb52241d5cac6c3e685fb76f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 4 Jun 2018 12:42:01 -0700
Subject: [PATCH 0039/2720] add a comment for target_space_id in
 transformer.py.

PiperOrigin-RevId: 199176946
---
 tensor2tensor/models/transformer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 3699e50a1..f21f2936f 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -145,7 +145,7 @@ def body(self, features):
           "inputs": Transformer inputs [batch_size, input_length, hidden_dim]
           "targets": Target decoder outputs.
               [batch_size, decoder_length, hidden_dim]
-          "target_space_id"
+          "target_space_id": A scalar int from data_generators.problem.SpaceID.
 
     Returns:
       Final decoder representation. [batch_size, decoder_length, hidden_dim]

From 9878f2fd6a903b5bbcd8060cbef3a2d4fa04f64c Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Mon, 4 Jun 2018 13:08:23 -0700
Subject: [PATCH 0040/2720] Self contained NAT Transformer from
 https://arxiv.org/abs/1805.11063

PiperOrigin-RevId: 199180774
---
 tensor2tensor/data_generators/translate.py    |  14 +
 tensor2tensor/layers/discretization.py        |  41 +-
 tensor2tensor/layers/discretization_test.py   |   5 +-
 tensor2tensor/models/__init__.py              |   1 +
 .../models/research/transformer_nat.py        | 399 ++++++++++++++++++
 5 files changed, 446 insertions(+), 14 deletions(-)
 create mode 100644 tensor2tensor/models/research/transformer_nat.py

diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 1edae6564..28984b6fa 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -164,6 +164,20 @@ class TranslateDistillProblem(TranslateProblem):
   def is_generate_per_split(self):
     return True
 
+  def example_reading_spec(self):
+    data_fields = {"dist_targets": tf.VarLenFeature(tf.int64)}
+
+    if self.has_inputs:
+      data_fields["inputs"] = tf.VarLenFeature(tf.int64)
+
+    # hack: ignoring true targets and putting dist_targets in targets
+    data_items_to_decoders = {
+        "inputs": tf.contrib.slim.tfexample_decoder.Tensor("inputs"),
+        "targets": tf.contrib.slim.tfexample_decoder.Tensor("dist_targets"),
+    }
+
+    return (data_fields, data_items_to_decoders)
+
   def get_or_create_vocab(self, data_dir, tmp_dir, force_get=False):
     """Get vocab for distill problems."""
     # We assume that vocab file is present in data_dir directory where the
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 7b5818f96..bcb534d6c 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -695,15 +695,21 @@ def get_vq_bottleneck(bottleneck_size, hidden_size):
   return means, ema_means, ema_count
 
 
-def vq_nearest_neighbor(x, means):
+def vq_nearest_neighbor(x, means, soft_em=False, num_samples=10):
   """Find the nearest element in means to elements in x."""
   bottleneck_size = common_layers.shape_list(means)[0]
   x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keepdims=True)
   means_norm_sq = tf.reduce_sum(tf.square(means), axis=-1, keepdims=True)
   scalar_prod = tf.matmul(x, means, transpose_b=True)
   dist = x_norm_sq + tf.transpose(means_norm_sq) - 2 * scalar_prod
-  x_means_idx = tf.argmax(-dist, axis=-1)
-  x_means_hot = tf.one_hot(x_means_idx, bottleneck_size)
+  if soft_em:
+    x_means_idx = tf.multinomial(-dist, num_samples=num_samples)
+    x_means_hot = tf.one_hot(
+        x_means_idx, depth=common_layers.shape_list(means)[0])
+    x_means_hot = tf.reduce_sum(x_means_hot, axis=1)
+  else:
+    x_means_idx = tf.argmax(-dist, axis=-1)
+    x_means_hot = tf.one_hot(x_means_idx, bottleneck_size)
   x_means_hot_flat = tf.reshape(x_means_hot, [-1, bottleneck_size])
   x_means = tf.matmul(x_means_hot_flat, means)
   e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
@@ -711,16 +717,20 @@ def vq_nearest_neighbor(x, means):
 
 
 def vq_discrete_bottleneck(x,
-                           bottleneck_size,
+                           bottleneck_bits,
                            beta=0.25,
                            decay=0.999,
-                           epsilon=1e-5):
+                           epsilon=1e-5,
+                           soft_em=False,
+                           num_samples=10):
   """Simple vector quantized discrete bottleneck."""
+  bottleneck_size = 2**bottleneck_bits
   x_shape = common_layers.shape_list(x)
   hidden_size = x_shape[-1]
   means, ema_means, ema_count = get_vq_bottleneck(bottleneck_size, hidden_size)
   x = tf.reshape(x, [-1, hidden_size])
-  x_means_hot, e_loss = vq_nearest_neighbor(x, means)
+  x_means_hot, e_loss = vq_nearest_neighbor(
+      x, means, soft_em=soft_em, num_samples=num_samples)
 
   # Update the ema variables
   updated_ema_count = moving_averages.assign_moving_average(
@@ -731,7 +741,6 @@ def vq_discrete_bottleneck(x,
       zero_debias=False)
 
   dw = tf.matmul(x_means_hot, x, transpose_a=True)
-
   updated_ema_means = tf.identity(moving_averages.assign_moving_average(
       ema_means, dw, decay, zero_debias=False))
   n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
@@ -754,8 +763,7 @@ def vq_discrete_unbottleneck(x, hidden_size):
   bottleneck_size = common_layers.shape_list(x)[-1]
   means, _, _ = get_vq_bottleneck(bottleneck_size, hidden_size)
   result = tf.matmul(tf.reshape(x, [-1, x_shape[-1]]), means)
-  return tf.reshape(result,
-                    x_shape[:-1] + [common_layers.shape_list(means)[-1]])
+  return tf.reshape(result, x_shape[:-1] + [hidden_size])
 
 
 def tanh_discrete_bottleneck(x, bottleneck_bits, bottleneck_noise,
@@ -825,9 +833,18 @@ def parametrized_bottleneck(x, hparams):
         hparams.discretize_warmup_steps, hparams.mode,
         hparams.isemhash_noise_dev, hparams.isemhash_mix_prob)
   if hparams.bottleneck_kind == "vq":
-    bottleneck_size = 2**hparams.bottleneck_bits
-    return vq_discrete_bottleneck(x, bottleneck_size, hparams.vq_beta,
+    return vq_discrete_bottleneck(x, hparams.bottleneck_bits, hparams.vq_beta,
                                   hparams.vq_decay, hparams.vq_epsilon)
+  if hparams.bottleneck_kind == "em":
+    return vq_discrete_bottleneck(
+        x,
+        hparams.bottleneck_bits,
+        hparams.vq_beta,
+        hparams.vq_decay,
+        hparams.vq_epsilon,
+        soft_em=True,
+        num_samples=hparams.vq_num_samples)
+
   raise ValueError("Unsupported hparams.bottleneck_kind %s"
                    % hparams.bottleneck_kind)
 
@@ -839,7 +856,7 @@ def parametrized_unbottleneck(x, hidden_size, hparams):
   if hparams.bottleneck_kind == "isemhash":
     return isemhash_unbottleneck(
         x, hidden_size, hparams.isemhash_filter_size_multiplier)
-  if hparams.bottleneck_kind == "vq":
+  if hparams.bottleneck_kind in ["vq", "em"]:
     return vq_discrete_unbottleneck(x, hidden_size)
   raise ValueError("Unsupported hparams.bottleneck_kind %s"
                    % hparams.bottleneck_kind)
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index 157ae62bb..8ad25a362 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -121,7 +121,8 @@ def testNearestNeighbors(self):
       self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
 
   def testGetVQBottleneck(self):
-    bottleneck_size = 4
+    bottleneck_bits = 2
+    bottleneck_size = 2**bottleneck_bits
     hidden_size = 3
     means, _, ema_count = discretization.get_vq_bottleneck(bottleneck_size,
                                                            hidden_size)
@@ -148,7 +149,7 @@ def testVQNearestNeighbors(self):
 
   def testVQDiscreteBottleneck(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
-    x_means_hot, _ = discretization.vq_discrete_bottleneck(x, bottleneck_size=4)
+    x_means_hot, _ = discretization.vq_discrete_bottleneck(x, bottleneck_bits=2)
     with self.test_session() as sess:
       tf.global_variables_initializer().run()
       x_means_hot_eval = sess.run(x_means_hot)
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index d0d774604..2aec38a24 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -47,6 +47,7 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research import super_lm
 from tensor2tensor.models.research import transformer_moe
+from tensor2tensor.models.research import transformer_nat
 from tensor2tensor.models.research import transformer_revnet
 from tensor2tensor.models.research import transformer_sketch
 from tensor2tensor.models.research import transformer_symshard
diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
new file mode 100644
index 000000000..e9393fdc5
--- /dev/null
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -0,0 +1,399 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""NAT Transformer from https://arxiv.org/abs/1805.11063."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import functools
+from six.moves import range
+from tensor2tensor.layers import common_layers
+from tensor2tensor.models import transformer
+from tensor2tensor.utils import beam_search
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
+import tensorflow as tf
+from tensorflow.python.training import moving_averages
+
+
+def get_vq_bottleneck(bottleneck_size, hidden_size):
+  """Get lookup table for VQ bottleneck."""
+  with tf.variable_scope("vq", reuse=tf.AUTO_REUSE):
+    means = tf.get_variable(
+        name="means",
+        shape=[bottleneck_size, hidden_size],
+        initializer=tf.uniform_unit_scaling_initializer())
+
+    ema_count = tf.get_variable(
+        name="ema_count",
+        shape=[bottleneck_size],
+        initializer=tf.constant_initializer(0),
+        trainable=False)
+
+    with tf.colocate_with(means):
+      ema_means = tf.get_variable(
+          name="ema_means",
+          initializer=means.initialized_value(),
+          trainable=False)
+
+  return means, ema_means, ema_count
+
+
+def vq_nearest_neighbor(x, means, hparams):
+  """Find the nearest element in means to elements in x."""
+  bottleneck_size = common_layers.shape_list(means)[0]
+  x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keepdims=True)
+  means_norm_sq = tf.reduce_sum(tf.square(means), axis=-1, keepdims=True)
+  scalar_prod = tf.matmul(x, means, transpose_b=True)
+  dist = x_norm_sq + tf.transpose(means_norm_sq) - 2 * scalar_prod
+  if hparams.bottleneck_kind == "em":
+    x_means_idx = tf.multinomial(-dist, num_samples=hparams.num_samples)
+    x_means_hot = tf.one_hot(
+        x_means_idx, depth=common_layers.shape_list(means)[0])
+    x_means_hot = tf.reduce_sum(x_means_hot, axis=1)
+  else:
+    x_means_idx = tf.argmax(-dist, axis=-1)
+    x_means_hot = tf.one_hot(x_means_idx, bottleneck_size)
+  x_means_hot_flat = tf.reshape(x_means_hot, [-1, bottleneck_size])
+  x_means = tf.matmul(x_means_hot_flat, means)
+  e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
+  return x_means_hot, e_loss
+
+
+def vq_discrete_bottleneck(x, hparams):
+  """Simple vector quantized discrete bottleneck."""
+  bottleneck_size = 2**hparams.bottleneck_bits
+  x_shape = common_layers.shape_list(x)
+  means, ema_means, ema_count = get_vq_bottleneck(bottleneck_size,
+                                                  hparams.hidden_size)
+  x = tf.reshape(x, [-1, hparams.hidden_size])
+  x_means_hot, e_loss = vq_nearest_neighbor(
+      x, means, hparams)
+
+  # Update the ema variables
+  updated_ema_count = moving_averages.assign_moving_average(
+      ema_count,
+      tf.reduce_sum(
+          tf.reshape(x_means_hot, shape=[-1, bottleneck_size]), axis=0),
+      hparams.decay,
+      zero_debias=False)
+
+  dw = tf.matmul(x_means_hot, x, transpose_a=True)
+  updated_ema_means = tf.identity(moving_averages.assign_moving_average(
+      ema_means, dw, hparams.decay, zero_debias=False))
+  n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
+  updated_ema_count = (
+      (updated_ema_count + hparams.epsilon) /
+      (n + bottleneck_size * hparams.epsilon) * n)
+  updated_ema_means /= tf.expand_dims(updated_ema_count, axis=-1)
+  with tf.control_dependencies([e_loss]):
+    update_means = means.assign(updated_ema_means)
+    with tf.control_dependencies([update_means]):
+      loss = hparams.beta * e_loss
+
+  d = tf.reshape(x_means_hot, x_shape[:-1] + [bottleneck_size])
+  return d, loss
+
+
+def vq_discrete_unbottleneck(x, hparams):
+  """Simple undiscretization from vector quantized representation."""
+  x_shape = common_layers.shape_list(x)
+  x = tf.to_float(x)
+  bottleneck_size = 2**hparams.bottleneck_bits
+  means, _, _ = get_vq_bottleneck(bottleneck_size, hparams.hidden_size)
+  result = tf.matmul(tf.reshape(x, [-1, x_shape[-1]]), means)
+  return tf.reshape(result, x_shape[:-1] + [hparams.hidden_size])
+
+
+def residual_conv(x, repeat, k, hparams, name, reuse=None):
+  """A stack of convolution blocks with residual connections."""
+  with tf.variable_scope(name, reuse=reuse):
+    dilations_and_kernels = [((1, 1), k) for _ in range(3)]
+    for i in range(repeat):
+      with tf.variable_scope("repeat_%d" % i):
+        y = common_layers.conv_block(
+            common_layers.layer_norm(x, hparams.hidden_size, name="lnorm"),
+            hparams.hidden_size,
+            dilations_and_kernels,
+            padding="SAME",
+            name="residual_conv")
+        y = tf.nn.dropout(y, 1.0 - hparams.dropout)
+        x += y
+    return x
+
+
+def decompress_step(source, hparams, first_relu, name):
+  """Decompression function."""
+  with tf.variable_scope(name):
+    shape = common_layers.shape_list(source)
+    multiplier = 2
+    kernel = (1, 1)
+    thicker = common_layers.conv_block(
+        source,
+        hparams.hidden_size * multiplier, [((1, 1), kernel)],
+        first_relu=first_relu,
+        name="decompress_conv")
+    return tf.reshape(thicker, [shape[0], shape[1] * 2, 1, hparams.hidden_size])
+
+
+def compress(x, hparams, name):
+  """Compress."""
+  with tf.variable_scope(name):
+    # Run compression by strided convs.
+    cur = x
+    k1 = (3, 1)
+    k2 = (2, 1)
+    cur = residual_conv(cur, hparams.num_compress_steps, k1, hparams, "rc")
+    for i in range(hparams.num_compress_steps):
+      cur = common_layers.conv_block(
+          cur,
+          hparams.hidden_size, [((1, 1), k2)],
+          strides=k2,
+          name="compress_%d" % i)
+    return cur
+
+
+def encode(x, x_space, hparams, name):
+  """Transformer preparations and encoder."""
+  with tf.variable_scope(name):
+    (encoder_input, encoder_self_attention_bias,
+     ed) = transformer.transformer_prepare_encoder(x, x_space, hparams)
+    encoder_input = tf.nn.dropout(encoder_input, 1.0 - hparams.dropout)
+    return transformer.transformer_encoder(
+        encoder_input, encoder_self_attention_bias, hparams), ed
+
+
+def decode_transformer(encoder_output, encoder_decoder_attention_bias, targets,
+                       hparams, name):
+  """Original Transformer decoder."""
+  orig_hparams = hparams
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    targets = common_layers.flatten4d3d(targets)
+
+    decoder_input, decoder_self_bias = (
+        transformer.transformer_prepare_decoder(targets, hparams))
+
+    decoder_input = tf.nn.dropout(decoder_input,
+                                  1.0 - hparams.layer_prepostprocess_dropout)
+
+    decoder_output = transformer.transformer_decoder(
+        decoder_input, encoder_output, decoder_self_bias,
+        encoder_decoder_attention_bias, hparams)
+    decoder_output = tf.expand_dims(decoder_output, axis=2)
+    decoder_output_shape = common_layers.shape_list(decoder_output)
+    decoder_output = tf.reshape(
+        decoder_output, [decoder_output_shape[0], -1, 1, hparams.hidden_size])
+    # Expand since t2t expects 4d tensors.
+    hparams = orig_hparams
+    return decoder_output
+
+
+def get_latent_pred_loss(latents_pred, latents_discrete, hparams):
+  """Latent prediction and loss."""
+  latents_logits = tf.layers.dense(
+      latents_pred, 2**hparams.bottleneck_bits, name="extra_logits")
+  loss = tf.nn.softmax_cross_entropy_with_logits_v2(
+      labels=latents_discrete, logits=latents_logits)
+  return loss
+
+
+def ae_latent_sample_beam(latents_dense_in, inputs, ed, embed, hparams):
+  """Sample from the latent space in the autoencoder."""
+
+  def symbols_to_logits_fn(ids):
+    """Go from ids to logits."""
+    ids = tf.expand_dims(ids, axis=2)  # Ids start with added all-zeros.
+    latents_discrete = tf.pad(ids[:, 1:], [[0, 0], [0, 1], [0, 0]])
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=False):
+      latents_dense = embed(
+          tf.one_hot(latents_discrete, depth=2**hparams.bottleneck_bits))
+      latents_pred = decode_transformer(inputs, ed, latents_dense, hparams,
+                                        "extra")
+      logits = tf.layers.dense(
+          latents_pred, 2**hparams.bottleneck_bits, name="extra_logits")
+      current_output_position = common_layers.shape_list(ids)[1] - 1
+      logits = logits[:, current_output_position, :, :]
+    return tf.squeeze(logits, axis=[1])
+
+  initial_ids = tf.zeros([tf.shape(latents_dense_in)[0]], dtype=tf.int32)
+  length = tf.shape(latents_dense_in)[1]
+  ids, _ = beam_search.beam_search(
+      symbols_to_logits_fn,
+      initial_ids,
+      1,
+      length,
+      2**hparams.bottleneck_bits,
+      alpha=0.0,
+      eos_id=-1,
+      stop_early=False)
+
+  res = tf.expand_dims(ids[:, 0, :], axis=2)  # Pick first beam.
+  return res[:, 1:]  # Remove the added all-zeros from ids.
+
+
+def ae_transformer_internal(inputs, targets, target_space, hparams, cache=None):
+  """Main step used for training."""
+  # Prepare.
+  if inputs is not None:
+    batch_size = common_layers.shape_list(inputs)[0]
+  else:
+    batch_size = common_layers.shape_list(targets)[0]
+  targets = tf.reshape(targets, [batch_size, -1, 1, hparams.hidden_size])
+
+  # Encoder.
+  if inputs is not None:
+    inputs = common_layers.flatten4d3d(inputs)
+    inputs, ed = encode(inputs, target_space, hparams, "input_enc")
+    inputs_ex, ed_ex = inputs, ed
+  else:
+    ed, inputs_ex, ed_ex = None, None, None
+
+  # Autoencoding.
+  losses = {"extra": tf.constant(0.0), "latent_pred": tf.constant(0.0)}
+
+  max_targets_len_from_inputs = tf.concat([inputs, inputs], axis=1)
+
+  targets, _ = common_layers.pad_to_same_length(
+      targets,
+      max_targets_len_from_inputs,
+      final_length_divisible_by=2**hparams.num_compress_steps)
+  targets_c = compress(targets, hparams, "compress")
+
+  if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+    # Compress and bottleneck.
+    latents_discrete_hot, extra_loss = vq_discrete_bottleneck(
+        x=targets_c, hparams=hparams)
+    latents_dense = vq_discrete_unbottleneck(latents_discrete_hot, hparams)
+    latents_discrete = tf.argmax(latents_discrete_hot, axis=-1)
+    tf.summary.histogram("codes", tf.reshape(latents_discrete[:, 0, :], [-1]))
+    pc = common_layers.inverse_exp_decay(hparams.startup_steps)
+    pc = pc if hparams.mode == tf.estimator.ModeKeys.TRAIN else 1.0
+    cond = tf.less(tf.random_uniform([batch_size]), pc)
+    latents_dense = tf.where(cond, latents_dense, targets_c)
+    losses["extra"] = extra_loss * tf.reduce_mean(tf.to_float(cond))
+
+    # Extra loss predicting latent code from input. Discrete only.
+    latents_pred = decode_transformer(inputs_ex, ed_ex, latents_dense, hparams,
+                                      "extra")
+    latent_pred_loss = get_latent_pred_loss(latents_pred, latents_discrete_hot,
+                                            hparams)
+    losses["latent_pred"] = tf.reduce_mean(latent_pred_loss * tf.to_float(cond))
+  else:
+    latent_len = common_layers.shape_list(targets_c)[1]
+    embed = functools.partial(vq_discrete_unbottleneck, hparams=hparams)
+    latents_dense = tf.zeros_like(targets_c[:, :latent_len, :, :])
+    if cache is None:
+      cache = ae_latent_sample_beam(latents_dense, inputs_ex, ed_ex, embed,
+                                    hparams)
+    latents_dense = embed(tf.one_hot(cache, depth=2**hparams.bottleneck_bits))
+
+  # Postprocess.
+  d = latents_dense
+  pos = tf.get_variable("pos", [1, 1000, 1, hparams.hidden_size])
+  pos = pos[:, :common_layers.shape_list(latents_dense)[1] + 1, :, :]
+  latents_dense = tf.pad(latents_dense, [[0, 0], [1, 0], [0, 0], [0, 0]]) + pos
+
+  # Decompressing the dense latents
+  for i in range(hparams.num_compress_steps):
+    j = hparams.num_compress_steps - i - 1
+    d = residual_conv(d, 1, (3, 1), hparams, "decompress_rc_%d" % j)
+    d = decompress_step(d, hparams, i > 0, "decompress_%d" % j)
+
+  res = decode_transformer(inputs, ed, targets, hparams, "decoder")
+  # We'll start training the extra model of latents after mask_startup_steps.
+  nonlatent_steps = hparams.mask_startup_steps
+  latent_time = tf.less(nonlatent_steps,
+                        tf.to_int32(tf.train.get_global_step()))
+  losses["latent_pred"] *= tf.to_float(latent_time)
+  return res, losses, cache
+
+
+@registry.register_model
+class TransformerNAT(t2t_model.T2TModel):
+  """Nonautoregressive Transformer from https://arxiv.org/abs/1805.11063."""
+
+  @property
+  def has_input(self):
+    return self._problem_hparams.input_modality
+
+  def body(self, features):
+    inputs = features["inputs"] if "inputs" in features else None
+    reuse = "cache_raw" in features
+    with tf.variable_scope(tf.get_variable_scope(), reuse=reuse):
+      res, loss, _ = ae_transformer_internal(
+          inputs, features["targets"], features["target_space_id"],
+          self._hparams, features.get("cache_raw", None))
+      return res, loss
+
+  def prepare_features_for_infer(self, features):
+    beam_batch_size = self._decode_hparams.beam_size
+    beam_batch_size *= self._decode_hparams.batch_size
+    inputs = tf.zeros([beam_batch_size, 1, 1, self._hparams.hidden_size])
+    inputs = inputs if "inputs" in features else None
+    targets = tf.zeros([beam_batch_size, 1, 1, self._hparams.hidden_size])
+    with tf.variable_scope("transformer_nat/body"):
+      _, _, cache = ae_transformer_internal(
+          inputs, targets, features["target_space_id"], self._hparams)
+    features["cache_raw"] = cache
+
+  def infer(self,
+            features=None,
+            decode_length=50,
+            beam_size=1,
+            top_beams=1,
+            alpha=0.0,
+            use_tpu=False):
+    """Produce predictions from the model."""
+    infer_out = super(TransformerNAT, self).infer(
+        features, decode_length, beam_size, top_beams, alpha, use_tpu=use_tpu)
+    return infer_out["outputs"]
+
+
+@registry.register_hparams
+def transformer_nat_small():
+  """Set of hyperparameters."""
+  hparams = transformer.transformer_small()
+  hparams.batch_size = 2048
+  hparams.learning_rate = 0.2
+  hparams.learning_rate_warmup_steps = 4000
+  hparams.num_hidden_layers = 3
+  hparams.hidden_size = 384
+  hparams.filter_size = 2048
+  hparams.label_smoothing = 0.0
+  hparams.optimizer = "Adam"  # Can be unstable, maybe try Adam.
+  hparams.optimizer_adam_epsilon = 1e-9
+  hparams.optimizer_adam_beta1 = 0.9
+  hparams.optimizer_adam_beta2 = 0.997  # Needs tuning, try 0.98 to 0.999.
+  hparams.add_hparam("bottleneck_kind", "em")
+  hparams.add_hparam("bottleneck_bits", 12)
+  hparams.add_hparam("num_compress_steps", 3)
+  hparams.add_hparam("startup_steps", 10000)
+  hparams.add_hparam("mask_startup_steps", 50000)
+  hparams.add_hparam("beta", 0.25)
+  hparams.add_hparam("epsilon", 1e-5)
+  hparams.add_hparam("decay", 0.999)
+  hparams.add_hparam("num_samples", 10)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_nat_base():
+  """Set of hyperparameters."""
+  hparams = transformer_nat_small()
+  hparams.batch_size = 2048
+  hparams.hidden_size = 512
+  hparams.filter_size = 4096
+  hparams.num_hidden_layers = 6
+  return hparams

From 64e1df15e649c6db97856f7ea6aea6508d52b0cf Mon Sep 17 00:00:00 2001
From: fstahlberg <fstahlberg@gmail.com>
Date: Tue, 5 Jun 2018 12:56:39 -0700
Subject: [PATCH 0041/2720] Add MultistepAdamOptimizer: Large training batches
 on limited GPU hardware (#754)

Simulates n times more GPUs at cost of n times more training iterations
---
 tensor2tensor/layers/common_hparams.py        |   2 +
 tensor2tensor/models/transformer.py           |   9 ++
 tensor2tensor/utils/learning_rate.py          |  17 ++-
 tensor2tensor/utils/multistep_optimizer.py    | 139 ++++++++++++++++++
 .../utils/multistep_optimizer_test.py         | 106 +++++++++++++
 tensor2tensor/utils/optimize.py               |   8 +
 6 files changed, 278 insertions(+), 3 deletions(-)
 create mode 100644 tensor2tensor/utils/multistep_optimizer.py
 create mode 100644 tensor2tensor/utils/multistep_optimizer_test.py

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index d3ae28d7d..986a220e3 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -64,6 +64,8 @@ def basic_params1():
       optimizer_adafactor_memory_exponent=0.8,
       optimizer_adafactor_clipping_threshold=1.0,
       optimizer_adafactor_multiply_by_parameter_scale=True,
+      # Number of accumulating steps for multi step optimizers.
+      optimizer_multistep_accumulate_steps=None,
       weight_decay=1e-6,
       weight_noise=0.0,
       # Defines the learning rate as a product of named functions.
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index f21f2936f..86ed71eb7 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1136,6 +1136,15 @@ def transformer_base_single_gpu():
   return hparams
 
 
+@registry.register_hparams
+def transformer_base_multistep8():
+  """HParams for simulating 8 GPUs with MultistepAdam optimizer."""
+  hparams = transformer_base()
+  hparams.optimizer = "MultistepAdam"
+  hparams.optimizer_multistep_accumulate_steps = 8
+  return hparams
+
+
 @registry.register_hparams
 def transformer_parsing_base():
   """HParams for parsing on WSJ only."""
diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index 843e484ab..bc4894e20 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -40,7 +40,7 @@ def learning_rate_factor(name, step_num, hparams):
 
 def learning_rate_schedule(hparams):
   """Learning rate schedule based on hparams."""
-  step_num = tf.to_float(tf.train.get_or_create_global_step())
+  step_num = _global_step(hparams)
   schedule_string = hparams.learning_rate_schedule
   names = schedule_string.split("*")
   names = [name.strip() for name in names if name.strip()]
@@ -52,7 +52,7 @@ def learning_rate_schedule(hparams):
 
 def legacy_learning_rate_schedule(hparams):
   """Backwards-compatible learning-rate schedule."""
-  step_num = tf.to_float(tf.train.get_or_create_global_step())
+  step_num = _global_step(hparams)
   warmup_steps = tf.to_float(hparams.learning_rate_warmup_steps)
   if hparams.learning_rate_decay_scheme == "noam":
     ret = 5000.0 * hparams.hidden_size**-0.5 * tf.minimum(
@@ -67,6 +67,17 @@ def legacy_learning_rate_schedule(hparams):
   return ret * optimizer_correction * hparams.learning_rate
 
 
+def _global_step(hparams):
+  """Adjust global step if a multi-step optimizer is used."""
+  step = tf.to_float(tf.train.get_or_create_global_step())
+  multiplier = hparams.optimizer_multistep_accumulate_steps
+  if multiplier:
+    step = step / tf.to_float(multiplier)
+    tf.logging.info("Divided global step by %d for multi-step optimizer."
+                    % multiplier)
+  return step
+
+
 def _legacy_sqrt_decay(step):
   """Decay like 1 / sqrt(step), multiplied by 500 to normalize."""
   return 500.0 / tf.sqrt(tf.maximum(step, 1.0))
@@ -95,7 +106,7 @@ def _learning_rate_decay(hparams, warmup_steps=0):
   """Learning rate decay multiplier."""
   scheme = hparams.learning_rate_decay_scheme
   warmup_steps = tf.to_float(warmup_steps)
-  global_step = tf.to_float(tf.train.get_or_create_global_step())
+  global_step = _global_step(hparams)
 
   if not scheme or scheme == "none":
     return tf.constant(1.)
diff --git a/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
new file mode 100644
index 000000000..9f25349eb
--- /dev/null
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -0,0 +1,139 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Optimizer variants which make it possible to use very large batch sizes with
+limited GPU memory. Optimizers in this module accumulate the gradients for n
+batches, and call the optimizer's update rule every n batches with the
+accumulated gradients.
+
+See [Saunders et al., 2018](https://arxiv.org/abs/1805.00456) for details.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# Dependency imports
+
+import tensorflow as tf
+
+
+class MultistepAdamOptimizer(tf.train.AdamOptimizer):
+  """Adam with SGD updates every n steps with accumulated gradients."""
+
+  def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8,
+               use_locking=False, name="Adam", n=1):
+    super(MultistepAdamOptimizer, self).__init__(
+        learning_rate=learning_rate, beta1=beta1, beta2=beta2, epsilon=epsilon,
+        use_locking=use_locking, name=name)
+    self._n = n  # Call Adam optimizer every n batches with accumulated grads
+    self._n_t = None  # n as tensor
+
+  def _create_slots(self, var_list):
+    """Create slot variables for Adam with accumulated gradients.
+
+    Like super class method, but additionally creates slots for the gradient
+    accumulator `acc_grad` and the counter variable.
+    """
+    super(MultistepAdamOptimizer, self)._create_slots(var_list)
+    first_var = min(var_list, key=lambda x: x.name)
+    self._create_non_slot_variable(initial_value=0 if self._n == 1 else 1,
+                                   name="iter",
+                                   colocate_with=first_var)
+    for v in var_list:
+      self._zeros_slot(v, "grad_acc", self._name)
+
+  def _get_iter_variable(self):
+    if tf.contrib.eager.in_eager_mode():
+      graph = None
+    else:
+      graph = tf.get_default_graph()
+    return self._get_non_slot_variable("iter", graph=graph)
+
+  def _prepare(self):
+    super(MultistepAdamOptimizer, self)._prepare()
+    self._n_t = tf.convert_to_tensor(self._n, name="n")
+
+  def _apply_cond(self, apply_fn, grad, var, *args, **kwargs):
+    """Conditionally apply or accumulate gradient.
+
+    Call `apply_fn only if the current counter value (iter) is zero. This
+    method couples common functionality for all _apply_*() implementations
+    in Adam.
+    """
+    grad_acc = self.get_slot(var, "grad_acc")
+
+    def apply_adam(grad_acc, apply_fn, grad, var, *args, **kwargs):
+      total_grad = (grad_acc + grad) / tf.cast(self._n_t, grad.dtype)
+      adam_op = apply_fn(total_grad, var, *args, **kwargs)
+      with tf.control_dependencies([adam_op]):
+        grad_acc_to_zero_op = grad_acc.assign(tf.zeros_like(grad_acc),
+                                              use_locking=self._use_locking)
+      return tf.group(adam_op, grad_acc_to_zero_op)
+
+    def accumulate_gradient(grad_acc, grad):
+      assign_op = tf.assign_add(grad_acc, grad, use_locking=self._use_locking)
+      return tf.group(assign_op)  # Strip return value
+
+    return tf.cond(tf.equal(self._get_iter_variable(), 0),
+                   lambda: apply_adam(
+                       grad_acc, apply_fn, grad, var, *args, **kwargs),
+                   lambda: accumulate_gradient(grad_acc, grad))
+
+  def _apply_dense(self, grad, var):
+    return self._apply_cond(
+        super(MultistepAdamOptimizer, self)._apply_dense, grad, var)
+
+  def _resource_apply_dense(self, grad, var):
+    return self._apply_cond(
+        super(MultistepAdamOptimizer, self)._resource_apply_dense, grad, var)
+
+  def _apply_sparse_shared(self, grad, var, indices, scatter_add):
+    return self._apply_cond(
+        super(MultistepAdamOptimizer, self)._apply_sparse_shared, grad, var,
+        indices, scatter_add)
+
+  def _apply_sparse(self, grad, var):
+    # TODO: Implement a sparse version
+    dense_grad = tf.convert_to_tensor(grad)
+    return self._apply_cond(
+        super(MultistepAdamOptimizer, self)._apply_dense, dense_grad, var)
+
+  def _finish(self, update_ops, name_scope):
+    """Like super class method, but updates beta_power variables only every
+    n batches. The iter variable is updated with
+
+       iter <- iter + 1 mod n
+    """
+    iter_ = self._get_iter_variable()
+    beta1_power, beta2_power = self._get_beta_accumulators()
+    with tf.control_dependencies(update_ops):
+      with tf.colocate_with(iter_):
+
+        def update_beta_op():
+          update_beta1 = beta1_power.assign(
+              beta1_power * self._beta1_t,
+              use_locking=self._use_locking)
+          update_beta2 = beta2_power.assign(
+              beta2_power * self._beta2_t,
+              use_locking=self._use_locking)
+          return tf.group(update_beta1, update_beta2)
+        maybe_update_beta = tf.cond(
+            tf.equal(iter_, 0), update_beta_op, tf.no_op)
+        with tf.control_dependencies([maybe_update_beta]):
+          update_iter = iter_.assign(tf.mod(iter_ + 1, self._n_t),
+                                     use_locking=self._use_locking)
+    return tf.group(
+        *update_ops + [update_iter, maybe_update_beta], name=name_scope)
+
diff --git a/tensor2tensor/utils/multistep_optimizer_test.py b/tensor2tensor/utils/multistep_optimizer_test.py
new file mode 100644
index 000000000..0cfc60482
--- /dev/null
+++ b/tensor2tensor/utils/multistep_optimizer_test.py
@@ -0,0 +1,106 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Multi-step Optimizer Test Module for TensorFlow."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# Dependency imports
+
+import numpy as np
+import tensorflow as tf
+from tensor2tensor.utils.multistep_optimizer import MultistepAdamOptimizer
+
+
+class MultistepAdamOptimizerTest(tf.test.TestCase):
+
+  def testMultistep(self):
+    ver = tf.__version__.split('.')
+    # TODO: Remove version check once 1.5 is not tested anymore
+    if int(ver[0]) <= 1 and int(ver[1]) < 6:
+      # MultistepAdamOptimizer requires TF >= 1.6
+      return
+    dtype = tf.float32
+    beta1 = 0.2
+    beta2 = 0.99
+    alpha = 10.0
+    grads0_np_lst = [
+        np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype),
+        np.array([0.2, -0.1], dtype=dtype.as_numpy_dtype),
+        np.array([0.3, 0.1], dtype=dtype.as_numpy_dtype),
+        np.array([0.4, -0.1], dtype=dtype.as_numpy_dtype)
+    ]
+    grads1_np_lst = [
+        np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype),
+        np.array([0.02, 0.02], dtype=dtype.as_numpy_dtype),
+        np.array([-0.04, 0.04], dtype=dtype.as_numpy_dtype),
+        np.array([-0.04, 0.06], dtype=dtype.as_numpy_dtype)
+    ]
+    var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+    var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+    # Test accumulating gradients for n=1..4 steps
+    for n in range(1, 5):
+      with self.test_session():
+        with self.test_session(graph=tf.Graph()):
+          singlestep_var0 = tf.Variable(var0_np)
+          singlestep_var1 = tf.Variable(var1_np)
+
+          multistep_var0 = tf.Variable(var0_np)
+          multistep_var1 = tf.Variable(var1_np)
+
+          singlestep_opt = tf.train.AdamOptimizer(
+              beta1=beta1, beta2=beta2, learning_rate=alpha)
+          multistep_opt = MultistepAdamOptimizer(
+              n=n, beta1=beta1, beta2=beta2, learning_rate=alpha)
+
+          singlestep_update = singlestep_opt.apply_gradients([
+              (tf.constant(sum(grads0_np_lst[:n]) / n), singlestep_var0),
+              (tf.constant(sum(grads1_np_lst[:n]) / n), singlestep_var1)])
+          multistep_updates = [
+              multistep_opt.apply_gradients([(tf.constant(g0), multistep_var0),
+                                             (tf.constant(g1), multistep_var1)])
+              for g0, g1 in zip(grads0_np_lst, grads1_np_lst)][:n]
+
+          self.evaluate(tf.global_variables_initializer())
+          (singlestep_beta1_power,
+           singlestep_beta2_power) = singlestep_opt._get_beta_accumulators()
+          (multistep_beta1_power,
+           multistep_beta2_power) = multistep_opt._get_beta_accumulators()
+
+          # Run 3 steps of Adam
+          for _ in range(1, 4):
+            self.evaluate(singlestep_update)
+            for multistep_update in multistep_updates:
+              self.evaluate(multistep_update)
+
+            self.assertAllCloseAccordingToType(
+                self.evaluate(singlestep_beta1_power),
+                self.evaluate(multistep_beta1_power))
+            self.assertAllCloseAccordingToType(
+                self.evaluate(singlestep_beta2_power),
+                self.evaluate(multistep_beta2_power))
+            # Validate updated params
+            self.assertAllCloseAccordingToType(
+                self.evaluate(singlestep_var0),
+                self.evaluate(multistep_var0))
+            self.assertAllCloseAccordingToType(
+                self.evaluate(singlestep_var1),
+                self.evaluate(multistep_var1))
+
+
+if __name__ == "__main__":
+  tf.test.main()
+
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index b973f9ed3..a64a2869a 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -20,6 +20,7 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import adafactor
+from tensor2tensor.utils import multistep_optimizer
 from tensor2tensor.utils import yellowfin
 
 import tensorflow as tf
@@ -84,6 +85,13 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
           beta1=hparams.optimizer_adam_beta1,
           beta2=hparams.optimizer_adam_beta2,
           epsilon=hparams.optimizer_adam_epsilon)
+    elif optimizer_name == "MultistepAdam":
+      self._opt = multistep_optimizer.MultistepAdamOptimizer(
+          lr,
+          beta1=hparams.optimizer_adam_beta1,
+          beta2=hparams.optimizer_adam_beta2,
+          epsilon=hparams.optimizer_adam_epsilon,
+          n=hparams.optimizer_multistep_accumulate_steps)
     elif optimizer_name == "Momentum":
       self._opt = tf.train.MomentumOptimizer(
           lr,

From 5626d0609444c907b9d77ec5dca5a9cd02546f8b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 5 Jun 2018 13:13:04 -0700
Subject: [PATCH 0042/2720] Internal change.

PiperOrigin-RevId: 199346216
---
 tensor2tensor/layers/common_hparams.py        |   2 -
 tensor2tensor/models/transformer.py           |   9 --
 tensor2tensor/utils/learning_rate.py          |  17 +--
 tensor2tensor/utils/multistep_optimizer.py    | 139 ------------------
 .../utils/multistep_optimizer_test.py         | 106 -------------
 tensor2tensor/utils/optimize.py               |   8 -
 6 files changed, 3 insertions(+), 278 deletions(-)
 delete mode 100644 tensor2tensor/utils/multistep_optimizer.py
 delete mode 100644 tensor2tensor/utils/multistep_optimizer_test.py

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 986a220e3..d3ae28d7d 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -64,8 +64,6 @@ def basic_params1():
       optimizer_adafactor_memory_exponent=0.8,
       optimizer_adafactor_clipping_threshold=1.0,
       optimizer_adafactor_multiply_by_parameter_scale=True,
-      # Number of accumulating steps for multi step optimizers.
-      optimizer_multistep_accumulate_steps=None,
       weight_decay=1e-6,
       weight_noise=0.0,
       # Defines the learning rate as a product of named functions.
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 86ed71eb7..f21f2936f 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1136,15 +1136,6 @@ def transformer_base_single_gpu():
   return hparams
 
 
-@registry.register_hparams
-def transformer_base_multistep8():
-  """HParams for simulating 8 GPUs with MultistepAdam optimizer."""
-  hparams = transformer_base()
-  hparams.optimizer = "MultistepAdam"
-  hparams.optimizer_multistep_accumulate_steps = 8
-  return hparams
-
-
 @registry.register_hparams
 def transformer_parsing_base():
   """HParams for parsing on WSJ only."""
diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index bc4894e20..843e484ab 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -40,7 +40,7 @@ def learning_rate_factor(name, step_num, hparams):
 
 def learning_rate_schedule(hparams):
   """Learning rate schedule based on hparams."""
-  step_num = _global_step(hparams)
+  step_num = tf.to_float(tf.train.get_or_create_global_step())
   schedule_string = hparams.learning_rate_schedule
   names = schedule_string.split("*")
   names = [name.strip() for name in names if name.strip()]
@@ -52,7 +52,7 @@ def learning_rate_schedule(hparams):
 
 def legacy_learning_rate_schedule(hparams):
   """Backwards-compatible learning-rate schedule."""
-  step_num = _global_step(hparams)
+  step_num = tf.to_float(tf.train.get_or_create_global_step())
   warmup_steps = tf.to_float(hparams.learning_rate_warmup_steps)
   if hparams.learning_rate_decay_scheme == "noam":
     ret = 5000.0 * hparams.hidden_size**-0.5 * tf.minimum(
@@ -67,17 +67,6 @@ def legacy_learning_rate_schedule(hparams):
   return ret * optimizer_correction * hparams.learning_rate
 
 
-def _global_step(hparams):
-  """Adjust global step if a multi-step optimizer is used."""
-  step = tf.to_float(tf.train.get_or_create_global_step())
-  multiplier = hparams.optimizer_multistep_accumulate_steps
-  if multiplier:
-    step = step / tf.to_float(multiplier)
-    tf.logging.info("Divided global step by %d for multi-step optimizer."
-                    % multiplier)
-  return step
-
-
 def _legacy_sqrt_decay(step):
   """Decay like 1 / sqrt(step), multiplied by 500 to normalize."""
   return 500.0 / tf.sqrt(tf.maximum(step, 1.0))
@@ -106,7 +95,7 @@ def _learning_rate_decay(hparams, warmup_steps=0):
   """Learning rate decay multiplier."""
   scheme = hparams.learning_rate_decay_scheme
   warmup_steps = tf.to_float(warmup_steps)
-  global_step = _global_step(hparams)
+  global_step = tf.to_float(tf.train.get_or_create_global_step())
 
   if not scheme or scheme == "none":
     return tf.constant(1.)
diff --git a/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
deleted file mode 100644
index 9f25349eb..000000000
--- a/tensor2tensor/utils/multistep_optimizer.py
+++ /dev/null
@@ -1,139 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-"""Optimizer variants which make it possible to use very large batch sizes with
-limited GPU memory. Optimizers in this module accumulate the gradients for n
-batches, and call the optimizer's update rule every n batches with the
-accumulated gradients.
-
-See [Saunders et al., 2018](https://arxiv.org/abs/1805.00456) for details.
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-# Dependency imports
-
-import tensorflow as tf
-
-
-class MultistepAdamOptimizer(tf.train.AdamOptimizer):
-  """Adam with SGD updates every n steps with accumulated gradients."""
-
-  def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8,
-               use_locking=False, name="Adam", n=1):
-    super(MultistepAdamOptimizer, self).__init__(
-        learning_rate=learning_rate, beta1=beta1, beta2=beta2, epsilon=epsilon,
-        use_locking=use_locking, name=name)
-    self._n = n  # Call Adam optimizer every n batches with accumulated grads
-    self._n_t = None  # n as tensor
-
-  def _create_slots(self, var_list):
-    """Create slot variables for Adam with accumulated gradients.
-
-    Like super class method, but additionally creates slots for the gradient
-    accumulator `acc_grad` and the counter variable.
-    """
-    super(MultistepAdamOptimizer, self)._create_slots(var_list)
-    first_var = min(var_list, key=lambda x: x.name)
-    self._create_non_slot_variable(initial_value=0 if self._n == 1 else 1,
-                                   name="iter",
-                                   colocate_with=first_var)
-    for v in var_list:
-      self._zeros_slot(v, "grad_acc", self._name)
-
-  def _get_iter_variable(self):
-    if tf.contrib.eager.in_eager_mode():
-      graph = None
-    else:
-      graph = tf.get_default_graph()
-    return self._get_non_slot_variable("iter", graph=graph)
-
-  def _prepare(self):
-    super(MultistepAdamOptimizer, self)._prepare()
-    self._n_t = tf.convert_to_tensor(self._n, name="n")
-
-  def _apply_cond(self, apply_fn, grad, var, *args, **kwargs):
-    """Conditionally apply or accumulate gradient.
-
-    Call `apply_fn only if the current counter value (iter) is zero. This
-    method couples common functionality for all _apply_*() implementations
-    in Adam.
-    """
-    grad_acc = self.get_slot(var, "grad_acc")
-
-    def apply_adam(grad_acc, apply_fn, grad, var, *args, **kwargs):
-      total_grad = (grad_acc + grad) / tf.cast(self._n_t, grad.dtype)
-      adam_op = apply_fn(total_grad, var, *args, **kwargs)
-      with tf.control_dependencies([adam_op]):
-        grad_acc_to_zero_op = grad_acc.assign(tf.zeros_like(grad_acc),
-                                              use_locking=self._use_locking)
-      return tf.group(adam_op, grad_acc_to_zero_op)
-
-    def accumulate_gradient(grad_acc, grad):
-      assign_op = tf.assign_add(grad_acc, grad, use_locking=self._use_locking)
-      return tf.group(assign_op)  # Strip return value
-
-    return tf.cond(tf.equal(self._get_iter_variable(), 0),
-                   lambda: apply_adam(
-                       grad_acc, apply_fn, grad, var, *args, **kwargs),
-                   lambda: accumulate_gradient(grad_acc, grad))
-
-  def _apply_dense(self, grad, var):
-    return self._apply_cond(
-        super(MultistepAdamOptimizer, self)._apply_dense, grad, var)
-
-  def _resource_apply_dense(self, grad, var):
-    return self._apply_cond(
-        super(MultistepAdamOptimizer, self)._resource_apply_dense, grad, var)
-
-  def _apply_sparse_shared(self, grad, var, indices, scatter_add):
-    return self._apply_cond(
-        super(MultistepAdamOptimizer, self)._apply_sparse_shared, grad, var,
-        indices, scatter_add)
-
-  def _apply_sparse(self, grad, var):
-    # TODO: Implement a sparse version
-    dense_grad = tf.convert_to_tensor(grad)
-    return self._apply_cond(
-        super(MultistepAdamOptimizer, self)._apply_dense, dense_grad, var)
-
-  def _finish(self, update_ops, name_scope):
-    """Like super class method, but updates beta_power variables only every
-    n batches. The iter variable is updated with
-
-       iter <- iter + 1 mod n
-    """
-    iter_ = self._get_iter_variable()
-    beta1_power, beta2_power = self._get_beta_accumulators()
-    with tf.control_dependencies(update_ops):
-      with tf.colocate_with(iter_):
-
-        def update_beta_op():
-          update_beta1 = beta1_power.assign(
-              beta1_power * self._beta1_t,
-              use_locking=self._use_locking)
-          update_beta2 = beta2_power.assign(
-              beta2_power * self._beta2_t,
-              use_locking=self._use_locking)
-          return tf.group(update_beta1, update_beta2)
-        maybe_update_beta = tf.cond(
-            tf.equal(iter_, 0), update_beta_op, tf.no_op)
-        with tf.control_dependencies([maybe_update_beta]):
-          update_iter = iter_.assign(tf.mod(iter_ + 1, self._n_t),
-                                     use_locking=self._use_locking)
-    return tf.group(
-        *update_ops + [update_iter, maybe_update_beta], name=name_scope)
-
diff --git a/tensor2tensor/utils/multistep_optimizer_test.py b/tensor2tensor/utils/multistep_optimizer_test.py
deleted file mode 100644
index 0cfc60482..000000000
--- a/tensor2tensor/utils/multistep_optimizer_test.py
+++ /dev/null
@@ -1,106 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-"""Multi-step Optimizer Test Module for TensorFlow."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-# Dependency imports
-
-import numpy as np
-import tensorflow as tf
-from tensor2tensor.utils.multistep_optimizer import MultistepAdamOptimizer
-
-
-class MultistepAdamOptimizerTest(tf.test.TestCase):
-
-  def testMultistep(self):
-    ver = tf.__version__.split('.')
-    # TODO: Remove version check once 1.5 is not tested anymore
-    if int(ver[0]) <= 1 and int(ver[1]) < 6:
-      # MultistepAdamOptimizer requires TF >= 1.6
-      return
-    dtype = tf.float32
-    beta1 = 0.2
-    beta2 = 0.99
-    alpha = 10.0
-    grads0_np_lst = [
-        np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype),
-        np.array([0.2, -0.1], dtype=dtype.as_numpy_dtype),
-        np.array([0.3, 0.1], dtype=dtype.as_numpy_dtype),
-        np.array([0.4, -0.1], dtype=dtype.as_numpy_dtype)
-    ]
-    grads1_np_lst = [
-        np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype),
-        np.array([0.02, 0.02], dtype=dtype.as_numpy_dtype),
-        np.array([-0.04, 0.04], dtype=dtype.as_numpy_dtype),
-        np.array([-0.04, 0.06], dtype=dtype.as_numpy_dtype)
-    ]
-    var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
-    var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
-    # Test accumulating gradients for n=1..4 steps
-    for n in range(1, 5):
-      with self.test_session():
-        with self.test_session(graph=tf.Graph()):
-          singlestep_var0 = tf.Variable(var0_np)
-          singlestep_var1 = tf.Variable(var1_np)
-
-          multistep_var0 = tf.Variable(var0_np)
-          multistep_var1 = tf.Variable(var1_np)
-
-          singlestep_opt = tf.train.AdamOptimizer(
-              beta1=beta1, beta2=beta2, learning_rate=alpha)
-          multistep_opt = MultistepAdamOptimizer(
-              n=n, beta1=beta1, beta2=beta2, learning_rate=alpha)
-
-          singlestep_update = singlestep_opt.apply_gradients([
-              (tf.constant(sum(grads0_np_lst[:n]) / n), singlestep_var0),
-              (tf.constant(sum(grads1_np_lst[:n]) / n), singlestep_var1)])
-          multistep_updates = [
-              multistep_opt.apply_gradients([(tf.constant(g0), multistep_var0),
-                                             (tf.constant(g1), multistep_var1)])
-              for g0, g1 in zip(grads0_np_lst, grads1_np_lst)][:n]
-
-          self.evaluate(tf.global_variables_initializer())
-          (singlestep_beta1_power,
-           singlestep_beta2_power) = singlestep_opt._get_beta_accumulators()
-          (multistep_beta1_power,
-           multistep_beta2_power) = multistep_opt._get_beta_accumulators()
-
-          # Run 3 steps of Adam
-          for _ in range(1, 4):
-            self.evaluate(singlestep_update)
-            for multistep_update in multistep_updates:
-              self.evaluate(multistep_update)
-
-            self.assertAllCloseAccordingToType(
-                self.evaluate(singlestep_beta1_power),
-                self.evaluate(multistep_beta1_power))
-            self.assertAllCloseAccordingToType(
-                self.evaluate(singlestep_beta2_power),
-                self.evaluate(multistep_beta2_power))
-            # Validate updated params
-            self.assertAllCloseAccordingToType(
-                self.evaluate(singlestep_var0),
-                self.evaluate(multistep_var0))
-            self.assertAllCloseAccordingToType(
-                self.evaluate(singlestep_var1),
-                self.evaluate(multistep_var1))
-
-
-if __name__ == "__main__":
-  tf.test.main()
-
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index a64a2869a..b973f9ed3 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -20,7 +20,6 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import adafactor
-from tensor2tensor.utils import multistep_optimizer
 from tensor2tensor.utils import yellowfin
 
 import tensorflow as tf
@@ -85,13 +84,6 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
           beta1=hparams.optimizer_adam_beta1,
           beta2=hparams.optimizer_adam_beta2,
           epsilon=hparams.optimizer_adam_epsilon)
-    elif optimizer_name == "MultistepAdam":
-      self._opt = multistep_optimizer.MultistepAdamOptimizer(
-          lr,
-          beta1=hparams.optimizer_adam_beta1,
-          beta2=hparams.optimizer_adam_beta2,
-          epsilon=hparams.optimizer_adam_epsilon,
-          n=hparams.optimizer_multistep_accumulate_steps)
     elif optimizer_name == "Momentum":
       self._opt = tf.train.MomentumOptimizer(
           lr,

From 1d6cb552c4d637a6cb241f6ba7e19a10a0a632bc Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 5 Jun 2018 14:05:02 -0700
Subject: [PATCH 0043/2720] internal merge of PR #754

PiperOrigin-RevId: 199354554
---
 tensor2tensor/layers/common_hparams.py        |   2 +
 tensor2tensor/models/transformer.py           |   9 ++
 tensor2tensor/utils/learning_rate.py          |  24 +++-
 tensor2tensor/utils/multistep_optimizer.py    | 123 ++++++++++++++++++
 .../utils/multistep_optimizer_test.py         | 102 +++++++++++++++
 tensor2tensor/utils/optimize.py               |   8 ++
 6 files changed, 262 insertions(+), 6 deletions(-)
 create mode 100644 tensor2tensor/utils/multistep_optimizer.py
 create mode 100644 tensor2tensor/utils/multistep_optimizer_test.py

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index d3ae28d7d..986a220e3 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -64,6 +64,8 @@ def basic_params1():
       optimizer_adafactor_memory_exponent=0.8,
       optimizer_adafactor_clipping_threshold=1.0,
       optimizer_adafactor_multiply_by_parameter_scale=True,
+      # Number of accumulating steps for multi step optimizers.
+      optimizer_multistep_accumulate_steps=None,
       weight_decay=1e-6,
       weight_noise=0.0,
       # Defines the learning rate as a product of named functions.
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index f21f2936f..86ed71eb7 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1136,6 +1136,15 @@ def transformer_base_single_gpu():
   return hparams
 
 
+@registry.register_hparams
+def transformer_base_multistep8():
+  """HParams for simulating 8 GPUs with MultistepAdam optimizer."""
+  hparams = transformer_base()
+  hparams.optimizer = "MultistepAdam"
+  hparams.optimizer_multistep_accumulate_steps = 8
+  return hparams
+
+
 @registry.register_hparams
 def transformer_parsing_base():
   """HParams for parsing on WSJ only."""
diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index 843e484ab..ba2b7920c 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -40,7 +40,7 @@ def learning_rate_factor(name, step_num, hparams):
 
 def learning_rate_schedule(hparams):
   """Learning rate schedule based on hparams."""
-  step_num = tf.to_float(tf.train.get_or_create_global_step())
+  step_num = _global_step(hparams)
   schedule_string = hparams.learning_rate_schedule
   names = schedule_string.split("*")
   names = [name.strip() for name in names if name.strip()]
@@ -52,14 +52,14 @@ def learning_rate_schedule(hparams):
 
 def legacy_learning_rate_schedule(hparams):
   """Backwards-compatible learning-rate schedule."""
-  step_num = tf.to_float(tf.train.get_or_create_global_step())
+  step_num = _global_step(hparams)
   warmup_steps = tf.to_float(hparams.learning_rate_warmup_steps)
   if hparams.learning_rate_decay_scheme == "noam":
     ret = 5000.0 * hparams.hidden_size**-0.5 * tf.minimum(
         (step_num + 1) * warmup_steps**-1.5, (step_num + 1)**-0.5)
   else:
     warmup_steps = hparams.learning_rate_warmup_steps
-    warmup = _learning_rate_warmup(warmup_steps)
+    warmup = _learning_rate_warmup(warmup_steps, hparams=hparams)
     decay = _learning_rate_decay(hparams, warmup_steps)
     ret = tf.where(step_num < warmup_steps, warmup, decay)
   optimizer_correction = 0.002 if "Adam" in hparams.optimizer else 1.0
@@ -67,6 +67,18 @@ def legacy_learning_rate_schedule(hparams):
   return ret * optimizer_correction * hparams.learning_rate
 
 
+def _global_step(hparams):
+  """Adjust global step if a multi-step optimizer is used."""
+  step = tf.to_float(tf.train.get_or_create_global_step())
+  multiplier = hparams.optimizer_multistep_accumulate_steps
+  if not multiplier:
+    return step
+
+  tf.logging.info("Dividing global step by %d for multi-step optimizer."
+                  % multiplier)
+  return step / tf.to_float(multiplier)
+
+
 def _legacy_sqrt_decay(step):
   """Decay like 1 / sqrt(step), multiplied by 500 to normalize."""
   return 500.0 / tf.sqrt(tf.maximum(step, 1.0))
@@ -95,7 +107,7 @@ def _learning_rate_decay(hparams, warmup_steps=0):
   """Learning rate decay multiplier."""
   scheme = hparams.learning_rate_decay_scheme
   warmup_steps = tf.to_float(warmup_steps)
-  global_step = tf.to_float(tf.train.get_or_create_global_step())
+  global_step = _global_step(hparams)
 
   if not scheme or scheme == "none":
     return tf.constant(1.)
@@ -136,7 +148,7 @@ def _learning_rate_decay(hparams, warmup_steps=0):
                    hparams.learning_rate_decay_scheme)
 
 
-def _learning_rate_warmup(warmup_steps, warmup_schedule="exp"):
+def _learning_rate_warmup(warmup_steps, warmup_schedule="exp", hparams=None):
   """Learning rate warmup multiplier."""
   if not warmup_steps:
     return tf.constant(1.)
@@ -145,7 +157,7 @@ def _learning_rate_warmup(warmup_steps, warmup_schedule="exp"):
                   warmup_schedule, warmup_steps)
 
   warmup_steps = tf.to_float(warmup_steps)
-  global_step = tf.to_float(tf.train.get_or_create_global_step())
+  global_step = _global_step(hparams)
 
   if warmup_schedule == "exp":
     return tf.exp(tf.log(0.01) / warmup_steps)**(warmup_steps - global_step)
diff --git a/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
new file mode 100644
index 000000000..d4c5c6b86
--- /dev/null
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -0,0 +1,123 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Multi-step optimizers simulating large batches.
+
+Optimizer variants which make it possible to use very large batch sizes with
+limited GPU memory. Optimizers in this module accumulate the gradients for n
+batches, and call the optimizer's update rule every n batches with the
+accumulated gradients.
+
+See [Saunders et al., 2018](https://arxiv.org/abs/1805.00456) for details.
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+
+class MultistepAdamOptimizer(tf.train.AdamOptimizer):
+  """Adam with SGD updates every n steps with accumulated gradients."""
+
+  def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8,
+               use_locking=False, name="Adam", n=1):
+    super(MultistepAdamOptimizer, self).__init__(
+        learning_rate=learning_rate, beta1=beta1, beta2=beta2, epsilon=epsilon,
+        use_locking=use_locking, name=name)
+    self._n = n  # Call Adam optimizer every n batches with accumulated grads
+    self._n_t = None  # n as tensor
+
+  def _create_slots(self, var_list):
+    """Create slot variables for Adam with accumulated gradients."""
+    super(MultistepAdamOptimizer, self)._create_slots(var_list)
+    first_var = min(var_list, key=lambda x: x.name)
+    self._create_non_slot_variable(initial_value=0 if self._n == 1 else 1,
+                                   name="iter",
+                                   colocate_with=first_var)
+    for v in var_list:
+      self._zeros_slot(v, "grad_acc", self._name)
+
+  def _get_iter_variable(self):
+    graph = (
+        None if tf.contrib.eager.in_eager_mode() else tf.get_default_graph())
+    return self._get_non_slot_variable("iter", graph=graph)
+
+  def _prepare(self):
+    super(MultistepAdamOptimizer, self)._prepare()
+    self._n_t = tf.convert_to_tensor(self._n, name="n")
+
+  def _apply_cond(self, apply_fn, grad, var, *args, **kwargs):
+    """Apply conditionally if counter is zero."""
+    grad_acc = self.get_slot(var, "grad_acc")
+
+    def apply_adam(grad_acc, apply_fn, grad, var, *args, **kwargs):
+      total_grad = (grad_acc + grad) / tf.cast(self._n_t, grad.dtype)
+      adam_op = apply_fn(total_grad, var, *args, **kwargs)
+      with tf.control_dependencies([adam_op]):
+        grad_acc_to_zero_op = grad_acc.assign(tf.zeros_like(grad_acc),
+                                              use_locking=self._use_locking)
+      return tf.group(adam_op, grad_acc_to_zero_op)
+
+    def accumulate_gradient(grad_acc, grad):
+      assign_op = tf.assign_add(grad_acc, grad, use_locking=self._use_locking)
+      return tf.group(assign_op)  # Strip return value
+
+    return tf.cond(
+        tf.equal(self._get_iter_variable(), 0),
+        lambda: apply_adam(grad_acc, apply_fn, grad, var, *args, **kwargs),
+        lambda: accumulate_gradient(grad_acc, grad))
+
+  def _apply_dense(self, grad, var):
+    return self._apply_cond(
+        super(MultistepAdamOptimizer, self)._apply_dense, grad, var)
+
+  def _resource_apply_dense(self, grad, var):
+    return self._apply_cond(
+        super(MultistepAdamOptimizer, self)._resource_apply_dense, grad, var)
+
+  def _apply_sparse_shared(self, grad, var, indices, scatter_add):
+    return self._apply_cond(
+        super(MultistepAdamOptimizer, self)._apply_sparse_shared, grad, var,
+        indices, scatter_add)
+
+  def _apply_sparse(self, grad, var):
+    # TODO(fstahlberg): Implement a sparse version
+    tf.logging.warning("MultistepAdamOptimizer does not support sparse updates")
+    dense_grad = tf.convert_to_tensor(grad)
+    return self._apply_cond(
+        super(MultistepAdamOptimizer, self)._apply_dense, dense_grad, var)
+
+  def _finish(self, update_ops, name_scope):
+    """Updates beta_power variables every n batches and incrs counter."""
+    iter_ = self._get_iter_variable()
+    beta1_power, beta2_power = self._get_beta_accumulators()
+    with tf.control_dependencies(update_ops):
+      with tf.colocate_with(iter_):
+
+        def update_beta_op():
+          update_beta1 = beta1_power.assign(
+              beta1_power * self._beta1_t,
+              use_locking=self._use_locking)
+          update_beta2 = beta2_power.assign(
+              beta2_power * self._beta2_t,
+              use_locking=self._use_locking)
+          return tf.group(update_beta1, update_beta2)
+        maybe_update_beta = tf.cond(
+            tf.equal(iter_, 0), update_beta_op, tf.no_op)
+        with tf.control_dependencies([maybe_update_beta]):
+          update_iter = iter_.assign(tf.mod(iter_ + 1, self._n_t),
+                                     use_locking=self._use_locking)
+    return tf.group(
+        *update_ops + [update_iter, maybe_update_beta], name=name_scope)
diff --git a/tensor2tensor/utils/multistep_optimizer_test.py b/tensor2tensor/utils/multistep_optimizer_test.py
new file mode 100644
index 000000000..50affc12b
--- /dev/null
+++ b/tensor2tensor/utils/multistep_optimizer_test.py
@@ -0,0 +1,102 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Multi-step Optimizer Test Module for TensorFlow."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.utils import multistep_optimizer
+import tensorflow as tf
+
+
+class MultistepAdamOptimizerTest(tf.test.TestCase):
+
+  def testMultistep(self):
+    ver = tf.__version__.split('.')
+    # TODO(rsepassi): Remove version check once 1.5 is not tested anymore
+    if int(ver[0]) <= 1 and int(ver[1]) < 6:
+      # MultistepAdamOptimizer requires TF >= 1.6
+      return
+    dtype = tf.float32
+    beta1 = 0.2
+    beta2 = 0.99
+    alpha = 10.0
+    grads0_np_lst = [
+        np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype),
+        np.array([0.2, -0.1], dtype=dtype.as_numpy_dtype),
+        np.array([0.3, 0.1], dtype=dtype.as_numpy_dtype),
+        np.array([0.4, -0.1], dtype=dtype.as_numpy_dtype)
+    ]
+    grads1_np_lst = [
+        np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype),
+        np.array([0.02, 0.02], dtype=dtype.as_numpy_dtype),
+        np.array([-0.04, 0.04], dtype=dtype.as_numpy_dtype),
+        np.array([-0.04, 0.06], dtype=dtype.as_numpy_dtype)
+    ]
+    var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+    var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+    # Test accumulating gradients for n=1..4 steps
+    for n in range(1, 5):
+      with tf.Graph().as_default():
+        with tf.Session():
+          singlestep_var0 = tf.Variable(var0_np)
+          singlestep_var1 = tf.Variable(var1_np)
+
+          multistep_var0 = tf.Variable(var0_np)
+          multistep_var1 = tf.Variable(var1_np)
+
+          singlestep_opt = tf.train.AdamOptimizer(
+              beta1=beta1, beta2=beta2, learning_rate=alpha)
+          multistep_opt = multistep_optimizer.MultistepAdamOptimizer(
+              n=n, beta1=beta1, beta2=beta2, learning_rate=alpha)
+
+          singlestep_update = singlestep_opt.apply_gradients([
+              (tf.constant(sum(grads0_np_lst[:n]) / n), singlestep_var0),
+              (tf.constant(sum(grads1_np_lst[:n]) / n), singlestep_var1)])
+          multistep_updates = [
+              multistep_opt.apply_gradients([(tf.constant(g0), multistep_var0),
+                                             (tf.constant(g1), multistep_var1)])
+              for g0, g1 in zip(grads0_np_lst, grads1_np_lst)][:n]
+
+          self.evaluate(tf.global_variables_initializer())
+          (singlestep_beta1_power,
+           singlestep_beta2_power) = singlestep_opt._get_beta_accumulators()
+          (multistep_beta1_power,
+           multistep_beta2_power) = multistep_opt._get_beta_accumulators()
+
+          # Run 3 steps of Adam
+          for _ in range(1, 4):
+            self.evaluate(singlestep_update)
+            for multistep_update in multistep_updates:
+              self.evaluate(multistep_update)
+
+            self.assertAllCloseAccordingToType(
+                self.evaluate(singlestep_beta1_power),
+                self.evaluate(multistep_beta1_power))
+            self.assertAllCloseAccordingToType(
+                self.evaluate(singlestep_beta2_power),
+                self.evaluate(multistep_beta2_power))
+            # Validate updated params
+            self.assertAllCloseAccordingToType(
+                self.evaluate(singlestep_var0),
+                self.evaluate(multistep_var0))
+            self.assertAllCloseAccordingToType(
+                self.evaluate(singlestep_var1),
+                self.evaluate(multistep_var1))
+
+
+if __name__ == '__main__':
+  tf.test.main()
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index b973f9ed3..a64a2869a 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -20,6 +20,7 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import adafactor
+from tensor2tensor.utils import multistep_optimizer
 from tensor2tensor.utils import yellowfin
 
 import tensorflow as tf
@@ -84,6 +85,13 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
           beta1=hparams.optimizer_adam_beta1,
           beta2=hparams.optimizer_adam_beta2,
           epsilon=hparams.optimizer_adam_epsilon)
+    elif optimizer_name == "MultistepAdam":
+      self._opt = multistep_optimizer.MultistepAdamOptimizer(
+          lr,
+          beta1=hparams.optimizer_adam_beta1,
+          beta2=hparams.optimizer_adam_beta2,
+          epsilon=hparams.optimizer_adam_epsilon,
+          n=hparams.optimizer_multistep_accumulate_steps)
     elif optimizer_name == "Momentum":
       self._opt = tf.train.MomentumOptimizer(
           lr,

From c16ac87132fe6c3c1618d49ce63832f108b0deaa Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 5 Jun 2018 16:18:34 -0700
Subject: [PATCH 0044/2720] Fixing the Shapes dataset bugs:

PiperOrigin-RevId: 199376240
---
 tensor2tensor/data_generators/video_generated.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 014d8adc4..bac927168 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -40,12 +40,12 @@ class VideoStochasticShapes10k(video_utils.VideoProblem):
   @property
   def num_input_frames(self):
     """Number of frames to batch on one input."""
-    return 4
+    return 1
 
   @property
   def num_target_frames(self):
     """Number of frames to predict in one step."""
-    return 1
+    return 3
 
   @property
   def is_generate_per_split(self):
@@ -66,7 +66,11 @@ def total_number_of_frames(self):
 
   @property
   def video_length(self):
-    return 5
+    return self.num_input_frames + self.num_target_frames
+
+  @property
+  def random_skip(self):
+    return False
 
   @property
   def extra_reading_spec(self):
@@ -130,11 +134,11 @@ def generate_stochastic_shape_instance(self):
                           [-1.0, -1.0]
                          ])
 
-    rnd = np.random.randint(len(direction))
     sp = np.array([lim/2.0, lim/2.0])
+    rnd = np.random.randint(len(direction))
     di = direction[rnd]
 
-    colors = ["b", "g", "r", "c", "m", "y", "k"]
+    colors = ["b", "g", "r", "c", "m", "y"]
     color = np.random.choice(colors)
 
     shape = np.random.choice([
@@ -149,7 +153,6 @@ def generate_stochastic_shape_instance(self):
     plt.ioff()
 
     xy = np.array(sp)
-    di = direction[0]
 
     for _ in range(self.video_length):
       fig = plt.figure()

From 66a1db26367930af7c8c0f44edb274ee669292f5 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Tue, 5 Jun 2018 17:24:03 -0700
Subject: [PATCH 0045/2720] Add predict mode to latent layers and use the new
 api

PiperOrigin-RevId: 199385746
---
 tensor2tensor/layers/latent_layers.py | 141 ++++++++++++++++----------
 1 file changed, 89 insertions(+), 52 deletions(-)

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 0cd5bd712..75ea4fe6a 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -14,12 +14,16 @@
 # limitations under the License.
 """Utils for latent variable models."""
 
+import functools
+
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import discretization
 from tensor2tensor.models import transformer
+from tensor2tensor.utils import beam_search
 
 import tensorflow as tf
 
@@ -80,7 +84,7 @@ def multinomial_sample(x, vocab_size, temperature):
 
 def ae_latent_softmax(latents_pred, latents_discrete, hparams):
   """Latent prediction and loss."""
-  vocab_size = 2 ** hparams.z_size
+  vocab_size = 2**hparams.bottleneck_bits
   if hparams.num_decode_blocks < 2:
     with tf.variable_scope("extra_logits"):
       latents_logits = tf.layers.dense(latents_pred, vocab_size,
@@ -103,26 +107,41 @@ def ae_latent_softmax(latents_pred, latents_discrete, hparams):
                                   hparams.sampling_temp)
       return sample, loss
 
-  # Multi-block case.
-  block_vocab_size = 2**(hparams.z_size // hparams.num_decode_blocks)
-  latents_logits = [
-      tf.layers.dense(
-          latents_pred, block_vocab_size, name="extra_logits_%d" % i)
-      for i in range(hparams.num_decode_blocks)
-  ]
-  loss = None
-  if latents_discrete is not None:
-    losses = []
-    for i in range(hparams.num_decode_blocks):
-      d = tf.floormod(tf.floordiv(latents_discrete,
-                                  block_vocab_size**i), block_vocab_size)
-      losses.append(tf.nn.sparse_softmax_cross_entropy_with_logits(
-          labels=d, logits=latents_logits[i]))
-    loss = sum(losses)
-  samples = [multinomial_sample(l, block_vocab_size, hparams.sampling_temp)
-             for l in latents_logits]
-  sample = sum([s * block_vocab_size**i for i, s in enumerate(samples)])
-  return sample, loss
+
+def ae_latent_sample_beam(latents_dense_in, inputs, ed, embed, hparams):
+  """Sample from the latent space in the autoencoder."""
+
+  def symbols_to_logits_fn(ids):
+    """Go from ids to logits."""
+    ids = tf.expand_dims(ids, axis=2)  # Ids start with added all-zeros.
+    latents_discrete = tf.pad(ids[:, 1:], [[0, 0], [0, 1], [0, 0]])
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=False):
+      latents_dense = embed(
+          tf.one_hot(latents_discrete, depth=2**hparams.bottleneck_bits),
+          hparams.hidden_size)
+      latents_pred = transformer_latent_decoder(inputs, ed, latents_dense,
+                                                hparams, "extra")
+      logits = tf.layers.dense(
+          latents_pred, 2**hparams.bottleneck_bits, name="extra_logits")
+      current_output_position = common_layers.shape_list(ids)[1] - 1
+      logits = logits[:, current_output_position, :]
+    return logits
+
+  initial_ids = tf.zeros([tf.shape(latents_dense_in)[0]], dtype=tf.int32)
+  length = tf.shape(latents_dense_in)[1]
+  ids, _ = beam_search.beam_search(
+      symbols_to_logits_fn,
+      initial_ids,
+      1,
+      length,
+      2**hparams.bottleneck_bits,
+      alpha=0.0,
+      eos_id=-1,
+      stop_early=False)
+
+  res = tf.expand_dims(ids[:, 0, :], axis=2)  # Pick first beam.
+  return res[:, 1:]  # Remove the added all-zeros from ids.
 
 
 def residual_block_layer(inputs, hparams):
@@ -397,32 +416,33 @@ def transformer_latent_decoder(encoder_output,
     return decoder_output
 
 
-def bottleneck_layer(targets_c,
-                     hparams,
-                     name="bottlneck_d"):
+def bottleneck_layer(targets_c, hparams):
   """Compute latents from compressed targets."""
-  # TODO(nikip): Condense hparams by removing options we don't use.
-  latents_dense, latents_discrete, extra_loss, embed_func = (
-      hparams.bottleneck(
-          x=targets_c,
-          filter_size=hparams.compress_filter_size,
-          name=name,
-          mode=hparams.mode))
+  latents_discrete_hot, extra_loss = discretization.parametrized_bottleneck(
+      targets_c, hparams)
+  latents_dense = discretization.parametrized_unbottleneck(
+      latents_discrete_hot, hparams.hidden_size, hparams)
+  latents_discrete = tf.argmax(latents_discrete_hot, axis=-1)
+
   if DO_SUMMARIES:
     tf.summary.histogram("b0", tf.reshape(latents_discrete, [-1]))
-  return latents_dense, latents_discrete, extra_loss, embed_func
+  return latents_dense, latents_discrete, extra_loss
 
 
-def latent_prediction_model(
-    inputs, ed_attention_bias,
-    latents_discrete, embed,
-    hparams, name="latent_pred"):
+def latent_prediction_model(inputs,
+                            ed_attention_bias,
+                            latents_discrete,
+                            latents_dense,
+                            hparams,
+                            name="latent_pred"):
   """Transformer based latent prediction model."""
   with tf.variable_scope(name):
+    # latents_dense = discretization.parametrized_unbottleneck(
+    #     latents_discrete, hparams.hidden_size, hparams)
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
-      latents_pred = transformer_latent_decoder(
-          inputs, ed_attention_bias,
-          tf.stop_gradient(embed(latents_discrete)), hparams, name + "_extra")
+      latents_pred = transformer_latent_decoder(inputs, ed_attention_bias,
+                                                tf.stop_gradient(latents_dense),
+                                                hparams, name + "_extra")
       _, latent_pred_loss = ae_latent_softmax(
           latents_pred, tf.stop_gradient(latents_discrete), hparams)
   return latents_pred, latent_pred_loss
@@ -470,19 +490,36 @@ def transformer_autoencoder(inputs,
   # TODO(lukaszkaiser): return extra losses batchwise, multiply before mean.
   # Call bottleneck layer to get the latents.
   # Returns embedded latents, discrete latents, loss and the embedding function.
-  latents_dense, latents_discrete, extra_loss, embed = (
-      bottleneck_layer(targets_c, hparams))
-  extra_loss = tf.reduce_mean(extra_loss) * tf.to_float(cond)
-
-  # Call the autoregressive latent prediction model.
-  _, latents_pred_loss = latent_prediction_model(
-      targets_c, ed_attention_bias, latents_discrete,
-      embed, hparams, name="latent_pred")
-  latents_pred_loss = tf.reduce_mean(latents_pred_loss) * tf.to_float(cond)
-
-  # Assign latent loss
-  losses["latent_pred"] = latents_pred_loss
-  losses["extra_loss"] = extra_loss
+  if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+    latents_dense, latents_discrete, extra_loss = (
+        bottleneck_layer(targets_c, hparams))
+    extra_loss = tf.reduce_mean(extra_loss) * tf.to_float(cond)
+
+    # Call the autoregressive latent prediction model.
+    _, latents_pred_loss = latent_prediction_model(
+        targets_c,
+        ed_attention_bias,
+        latents_discrete,
+        latents_dense,
+        hparams,
+        name="latent_pred")
+    latents_pred_loss = tf.reduce_mean(latents_pred_loss) * tf.to_float(cond)
+    # Assign latent loss
+    losses["latent_pred"] = latents_pred_loss
+    losses["extra_loss"] = extra_loss
+  else:
+    latent_len = (
+        hparams.img_len * hparams.img_len * hparams.num_latents) / 2**(
+            hparams.num_compress_steps)
+    embed = functools.partial(
+        discretization.parametrized_unbottleneck, hparams=hparams)
+    latents_dense = tf.zeros([batch_size, latent_len, 1, hparams.hidden_size])
+    if cache is None:
+      cache = ae_latent_sample_beam(latents_dense, inputs, ed_attention_bias,
+                                    embed, hparams)
+    latents_dense = embed(
+        tf.one_hot(cache, depth=2**hparams.bottleneck_bits),
+        hparams.hidden_size)
 
   latents_decoder = latents_dense
   if len(original_targets_shape) == 4:

From 9304249f8efa20a60565b1f3a4fde43732e89556 Mon Sep 17 00:00:00 2001
From: LucienWang <lucienwang@qq.com>
Date: Wed, 6 Jun 2018 08:58:51 +0800
Subject: [PATCH 0046/2720] add log_step_count_steps (#821)

---
 tensor2tensor/bin/t2t_trainer.py   | 4 ++++
 tensor2tensor/utils/trainer_lib.py | 2 ++
 2 files changed, 6 insertions(+)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index f82f170e0..998bd2584 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -114,6 +114,9 @@
 flags.DEFINE_string("job-dir", None,
                     "DO NOT USE. Exists only for Cloud ML Engine to pass in "
                     "during hyperparameter tuning. Overrides --output_dir.")
+flags.DEFINE_integer("log_step_count_steps", 100,
+                     "Number of local steps after which progress is printed "
+                     "out")
 
 
 def set_hparams_from_args(args):
@@ -220,6 +223,7 @@ def create_run_config(hp):
       random_seed=FLAGS.random_seed,
       tpu_infeed_sleep_secs=FLAGS.tpu_infeed_sleep_secs,
       inter_op_parallelism_threads=FLAGS.inter_op_parallelism_threads,
+      log_step_count_steps=FLAGS.log_step_count_steps,
       intra_op_parallelism_threads=FLAGS.intra_op_parallelism_threads)
 
 
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index f1387de8e..ad17b1da9 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -110,6 +110,7 @@ def create_run_config(master="",
                       tpu_infeed_sleep_secs=None,
                       use_tpu=False,
                       inter_op_parallelism_threads=0,
+                      log_step_count_steps=100,
                       intra_op_parallelism_threads=0):
   """Create RunConfig, TPUConfig, and Parallelism object."""
   session_config = create_session_config(
@@ -129,6 +130,7 @@ def create_run_config(master="",
       "keep_checkpoint_max": keep_checkpoint_max,
       "keep_checkpoint_every_n_hours": keep_checkpoint_every_n_hours,
       "tf_random_seed": random_seed,
+      "log_step_count_steps": log_step_count_steps
   }
   if save_checkpoints_secs:
     del run_config_args["save_checkpoints_steps"]

From f76d5ea8f2b90d5ec7fd1c4a5a6fe1bd6d9c95cf Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Tue, 5 Jun 2018 21:17:04 -0700
Subject: [PATCH 0047/2720] Fix but with latent pred model

PiperOrigin-RevId: 199404887
---
 tensor2tensor/layers/latent_layers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 75ea4fe6a..bd8824414 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -497,7 +497,7 @@ def transformer_autoencoder(inputs,
 
     # Call the autoregressive latent prediction model.
     _, latents_pred_loss = latent_prediction_model(
-        targets_c,
+        inputs,
         ed_attention_bias,
         latents_discrete,
         latents_dense,

From 4c9bd0d74915dc0d9b568463fcea36f8b74b6318 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 6 Jun 2018 11:15:54 -0700
Subject: [PATCH 0048/2720] Adding the Stochastic Variational Video Prediction
 (SV2P) model.

PiperOrigin-RevId: 199495273
---
 .../data_generators/video_generated.py        |  18 +-
 tensor2tensor/data_generators/video_utils.py  |  10 +-
 tensor2tensor/layers/common_layers.py         |  12 +-
 tensor2tensor/layers/modalities.py            |  35 ++
 tensor2tensor/models/research/next_frame.py   | 592 +++++++++++++++++-
 5 files changed, 645 insertions(+), 22 deletions(-)

diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index bac927168..6aee3d787 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -22,6 +22,7 @@
 
 import numpy as np
 
+from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.utils import registry
 
@@ -45,7 +46,7 @@ def num_input_frames(self):
   @property
   def num_target_frames(self):
     """Number of frames to predict in one step."""
-    return 3
+    return 4
 
   @property
   def is_generate_per_split(self):
@@ -72,6 +73,17 @@ def video_length(self):
   def random_skip(self):
     return False
 
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each."""
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 1,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
   @property
   def extra_reading_spec(self):
     """Additional data fields to store on disk and their decoders."""
@@ -87,10 +99,10 @@ def extra_reading_spec(self):
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.input_modality = {
-        "inputs": ("video", 256),
+        "inputs": ("video:raw", 256),
         "input_frame_number": ("symbol:identity", 1)
     }
-    p.target_modality = ("video", 256)
+    p.target_modality = ("video:raw", 256)
 
   @staticmethod
   def get_circle(x, y, z, c, s):
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 77ad56805..c3e928eef 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -79,6 +79,11 @@ def num_target_frames(self):
     """Number of frames to batch on one target."""
     return 1
 
+  @property
+  def random_skip(self):
+    """Whether to skip random inputs at the beginning or not."""
+    return True
+
   @property
   def extra_reading_spec(self):
     """Additional data fields to store on disk and their decoders."""
@@ -191,8 +196,9 @@ def _preprocess(example):
 
     num_frames = self.num_input_frames + self.num_target_frames
     # We jump by a random position at the beginning to add variety.
-    random_skip = tf.random_uniform([], maxval=num_frames, dtype=tf.int64)
-    preprocessed_dataset = preprocessed_dataset.skip(random_skip)
+    if self.random_skip:
+      random_skip = tf.random_uniform([], maxval=num_frames, dtype=tf.int64)
+      preprocessed_dataset = preprocessed_dataset.skip(random_skip)
     batch_dataset = preprocessed_dataset.apply(
         tf.contrib.data.batch_and_drop_remainder(num_frames))
     dataset = batch_dataset.map(features_from_batch).shuffle(8)
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 8fc7a94ae..9f847640b 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -202,10 +202,14 @@ def convert_rgb_to_real(x):
   """Conversion of pixel values to real numbers."""
   with tf.name_scope("rgb_to_real", values=[x]):
     x = tf.to_float(x)
-    # Use the formula (value/127.5) - 1 to convert each channel value into a
-    # real number in the range -1 to 1. We use 127.5 instead of 128 because
-    # the intensities are in the range 0 to 255
-    x = (x / 127.5) - 1
+    x /= 255.0
+    return x
+
+
+def convert_real_to_rgb(x):
+  """Conversion of real numbers to pixel values."""
+  with tf.name_scope("real_to_rgb", values=[x]):
+    x *= 255.0
     return x
 
 
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 0c1bba2a9..74772cf11 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -540,6 +540,41 @@ def loss(self, top_out, targets):
         weights_fn=self.targets_weights_fn)
 
 
+@registry.register_video_modality("raw")
+class VideoModalityRaw(modality.Modality):
+  """Modality for raw videos, i.e., time-sequences of frames."""
+
+  def bottom(self, x):
+    common_layers.summarize_video(x, "inputs")
+    return common_layers.convert_rgb_to_real(x)
+
+  def targets_bottom(self, x):
+    common_layers.summarize_video(x, "targets_bottom")
+    return common_layers.convert_rgb_to_real(x)
+
+  def top(self, body_output, _):
+    frames = tf.stack(body_output, axis=1)
+    rgb_frames = common_layers.convert_real_to_rgb(frames)
+    common_layers.summarize_video(rgb_frames, "body_output")
+    return frames
+
+  def loss(self, top_out, targets):
+    assert(top_out.shape.as_list() == targets.shape.as_list()), \
+           "The dimensions doesn't match."
+
+    common_layers.summarize_video(targets, "targets_top")
+    targets = common_layers.convert_rgb_to_real(targets)
+
+    num_frames = top_out.shape[1].value
+    loss = 0.0
+    for frame_id in range(num_frames):
+      frame = tf.to_float(top_out[:, frame_id])
+      target = tf.to_float(targets[:, frame_id])
+      loss += tf.reduce_mean(tf.square(frame - target))
+    loss /= num_frames
+    return loss, tf.zeros_like(loss)
+
+
 @registry.register_video_modality("embed")
 class VideoModalityEmbed(VideoModality):
   """Video Modality where bottom embeds pixels."""
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 39e3804b4..3b4b9b8e5 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -17,6 +17,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+import numpy as np
 import six
 
 from tensor2tensor.layers import common_attention
@@ -26,6 +27,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow as tf
+slim = tf.contrib.slim
 
 
 @registry.register_model
@@ -171,22 +173,566 @@ def logits_to_samples(logits):
 class NextFrameStochastic(NextFrameBasic):
   """Stochastic next-frame model."""
 
+  def construct_latent_tower(self, images):
+    """Builds convolutional latent tower for stochastic model.
+
+    At training time this tower generates a latent distribution (mean and std)
+    conditioned on the entire video. This latent variable will be fed to the
+    main tower as an extra variable to be used for future frames prediction.
+    At inference time, the tower is disabled and only returns latents sampled
+    from N(0,1).
+    If the multi_latent flag is on, a different latent for every timestep would
+    be generated.
+
+    Args:
+      images: tensor of ground truth image sequences
+    Returns:
+      latent_mean: predicted latent mean
+      latent_std: predicted latent standard deviation
+      latent_loss: loss of the latent twoer
+      samples: random samples sampled from standard guassian
+    """
+    sequence_length = len(images)
+
+    with slim.arg_scope([slim.conv2d], reuse=False):
+      stacked_images = tf.concat(images, 3)
+
+      latent_enc1 = slim.conv2d(
+          stacked_images,
+          32, [3, 3],
+          stride=2,
+          scope="latent_conv1")
+      latent_enc1 = slim.batch_norm(latent_enc1, scope="latent_bn1")
+
+      latent_enc2 = slim.conv2d(
+          latent_enc1,
+          64, [3, 3],
+          stride=2,
+          scope="latent_conv2")
+      latent_enc2 = slim.batch_norm(latent_enc2, scope="latent_bn2")
+
+      latent_enc3 = slim.conv2d(
+          latent_enc2,
+          64, [3, 3],
+          stride=1,
+          scope="latent_conv3")
+      latent_enc3 = slim.batch_norm(latent_enc3, scope="latent_bn3")
+
+      latent_mean = slim.conv2d(
+          latent_enc3,
+          self.hparams.latent_channels, [3, 3],
+          stride=2,
+          activation_fn=None,
+          scope="latent_mean")
+
+      latent_std = slim.conv2d(
+          latent_enc3,
+          self.hparams.latent_channels, [3, 3],
+          stride=2,
+          scope="latent_std")
+
+      latent_std += self.hparams.latent_std_min
+
+    if self.hparams.multi_latent:
+      # timestep x batch_size x latent_size
+      samples = tf.random_normal(
+          [sequence_length-1] + latent_mean.shape, 0, 1,
+          dtype=tf.float32)
+    else:
+      # batch_size x latent_size
+      samples = tf.random_normal(tf.shape(latent_mean), 0, 1, dtype=tf.float32)
+
+    if self.hparams.mode == "train":
+      return latent_mean, latent_std, samples
+    else:
+      # No latent tower at inference time, just standard gaussian.
+      return None, None, samples
+
+  def construct_model(self,
+                      images,
+                      actions,
+                      states,
+                      k=-1,
+                      use_state=False,
+                      num_masks=10,
+                      cdna=True,
+                      dna=False,
+                      context_frames=2):
+    """Build convolutional lstm video predictor using CDNA, or DNA.
+
+    Args:
+      images: tensor of ground truth image sequences
+      actions: tensor of action sequences
+      states: tensor of ground truth state sequences
+      k: constant used for scheduled sampling. -1 to feed in own prediction.
+      use_state: True to include state and action in prediction
+      num_masks: the number of different pixel motion predictions (and
+                 the number of masks for each of those predictions)
+      cdna: True to use Convoluational Dynamic Neural Advection (CDNA)
+      dna: True to use Dynamic Neural Advection (DNA)
+      context_frames: number of ground truth frames to pass in before
+                      feeding in own predictions
+    Returns:
+      gen_images: predicted future image frames
+      gen_states: predicted future states
+
+    Raises:
+      ValueError: if more than one network option specified or more than 1 mask
+      specified for DNA model.
+    """
+    # Each image is being used twice, in latent tower and main tower.
+    # This is to make sure we are using the *same* image for both, ...
+    # ... given how TF queues work.
+    images = [tf.identity(image) for image in images]
+
+    if cdna + dna != 1:
+      raise ValueError("More than one, or no network option specified.")
+
+    batch_size, img_height, img_width, color_channels = \
+    images[0].get_shape()[0:4]
+    batch_size = tf.shape(images[0])[0]
+    lstm_func = self.basic_conv_lstm_cell
+
+    # Generated robot states and images.
+    gen_states, gen_images = [], []
+    current_state = states[0]
+
+    if k == -1:
+      feedself = True
+    else:
+      # Scheduled sampling:
+      # Calculate number of ground-truth frames to pass in.
+      iter_num = tf.train.get_or_create_global_step()
+      num_ground_truth = tf.to_int32(
+          tf.round(
+              tf.to_float(batch_size) *
+              (k / (k + tf.exp(tf.to_float(iter_num) / tf.to_float(k))))))
+      feedself = False
+
+    # LSTM state sizes and states.
+    lstm_size = np.int32(np.array([32, 32, 64, 64, 128, 64, 32]))
+    lstm_state1, lstm_state2, lstm_state3, lstm_state4 = None, None, None, None
+    lstm_state5, lstm_state6, lstm_state7 = None, None, None
+
+    # Latent tower
+    if self.hparams.stochastic_model:
+      latent_tower_outputs = self.construct_latent_tower(images)
+      latent_mean, latent_std, samples = latent_tower_outputs
+
+    # Main tower
+    timestep = 0
+    layer_norm = tf.contrib.layers.layer_norm
+
+    for image, action in zip(images[:-1], actions[:-1]):
+      # Reuse variables after the first timestep.
+      reuse = bool(gen_images)
+
+      done_warm_start = len(gen_images) > context_frames - 1
+      with slim.arg_scope(
+          [
+              lstm_func, slim.layers.conv2d, slim.layers.fully_connected,
+              layer_norm, slim.layers.conv2d_transpose
+          ],
+          reuse=reuse):
+
+        if feedself and done_warm_start:
+          # Feed in generated image.
+          prev_image = gen_images[-1]
+        elif done_warm_start:
+          # Scheduled sampling
+          prev_image = self.scheduled_sample(
+              image, gen_images[-1], batch_size, num_ground_truth)
+        else:
+          # Always feed in ground_truth
+          prev_image = image
+
+        # Predicted state is always fed back in
+        state_action = tf.concat(axis=1, values=[action, current_state])
+
+        enc0 = slim.layers.conv2d(
+            prev_image,
+            32, [5, 5],
+            stride=2,
+            scope="scale1_conv1",
+            normalizer_fn=layer_norm,
+            normalizer_params={"scope": "layer_norm1"})
+
+        hidden1, lstm_state1 = lstm_func(
+            enc0, lstm_state1, lstm_size[0], scope="state1")
+        hidden1 = layer_norm(hidden1, scope="layer_norm2")
+        hidden2, lstm_state2 = lstm_func(
+            hidden1, lstm_state2, lstm_size[1], scope="state2")
+        hidden2 = layer_norm(hidden2, scope="layer_norm3")
+        enc1 = slim.layers.conv2d(
+            hidden2, hidden2.get_shape()[3], [3, 3], stride=2, scope="conv2")
+
+        hidden3, lstm_state3 = lstm_func(
+            enc1, lstm_state3, lstm_size[2], scope="state3")
+        hidden3 = layer_norm(hidden3, scope="layer_norm4")
+        hidden4, lstm_state4 = lstm_func(
+            hidden3, lstm_state4, lstm_size[3], scope="state4")
+        hidden4 = layer_norm(hidden4, scope="layer_norm5")
+        enc2 = slim.layers.conv2d(
+            hidden4, hidden4.get_shape()[3], [3, 3], stride=2, scope="conv3")
+
+        # Pass in state and action.
+        smear = tf.reshape(
+            state_action,
+            [-1, 1, 1, int(state_action.get_shape()[1])])
+        smear = tf.tile(
+            smear, [1, int(enc2.get_shape()[1]),
+                    int(enc2.get_shape()[2]), 1])
+        if use_state:
+          enc2 = tf.concat(axis=3, values=[enc2, smear])
+
+        # Setup latent
+        if self.hparams.stochastic_model:
+          latent = samples
+          if self.hparams.multi_latent:
+            latent = samples[timestep]
+          if self.hparams.mode == "train":
+            # TODO(mbz): put 1st stage of training back in if necessary
+            latent = latent_mean + tf.exp(latent_std / 2.0) * latent
+          with tf.control_dependencies([latent]):
+            enc2 = tf.concat([enc2, latent], 3)
+
+        enc3 = slim.layers.conv2d(
+            enc2, hidden4.get_shape()[3], [1, 1], stride=1, scope="conv4")
+
+        hidden5, lstm_state5 = lstm_func(
+            enc3, lstm_state5, lstm_size[4], scope="state5")  # last 8x8
+        hidden5 = layer_norm(hidden5, scope="layer_norm6")
+        enc4 = slim.layers.conv2d_transpose(
+            hidden5, hidden5.get_shape()[3], 3, stride=2, scope="convt1")
+
+        hidden6, lstm_state6 = lstm_func(
+            enc4, lstm_state6, lstm_size[5], scope="state6")  # 16x16
+        hidden6 = layer_norm(hidden6, scope="layer_norm7")
+        # Skip connection.
+        hidden6 = tf.concat(axis=3, values=[hidden6, enc1])  # both 16x16
+
+        enc5 = slim.layers.conv2d_transpose(
+            hidden6, hidden6.get_shape()[3], 3, stride=2, scope="convt2")
+        hidden7, lstm_state7 = lstm_func(
+            enc5, lstm_state7, lstm_size[6], scope="state7")  # 32x32
+        hidden7 = layer_norm(hidden7, scope="layer_norm8")
+
+        # Skip connection.
+        hidden7 = tf.concat(axis=3, values=[hidden7, enc0])  # both 32x32
+
+        enc6 = slim.layers.conv2d_transpose(
+            hidden7,
+            hidden7.get_shape()[3],
+            3,
+            stride=2,
+            scope="convt3",
+            activation_fn=None,
+            normalizer_fn=layer_norm,
+            normalizer_params={"scope": "layer_norm9"})
+
+        if dna:
+          # Using largest hidden state for predicting untied conv kernels.
+          enc7 = slim.layers.conv2d_transpose(
+              enc6,
+              self.hparams.DNA_KERN_SIZE**2,
+              1,
+              stride=1,
+              scope="convt4",
+              activation_fn=None)
+        else:
+          # Using largest hidden state for predicting a new image layer.
+          enc7 = slim.layers.conv2d_transpose(
+              enc6,
+              color_channels,
+              1,
+              stride=1,
+              scope="convt4",
+              activation_fn=None)
+          # This allows the network to also generate one image from scratch,
+          # which is useful when regions of the image become unoccluded.
+          transformed = [tf.nn.sigmoid(enc7)]
+
+        if cdna:
+          # cdna_input = tf.reshape(hidden5, [int(batch_size), -1])
+          cdna_input = tf.contrib.layers.flatten(hidden5)
+          transformed += self.cdna_transformation(
+              prev_image, cdna_input, num_masks, int(color_channels))
+        elif dna:
+          # Only one mask is supported (more should be unnecessary).
+          if num_masks != 1:
+            raise ValueError("Only one mask is supported for DNA model.")
+          transformed = [self.dna_transformation(prev_image, enc7)]
+
+        masks = slim.layers.conv2d_transpose(
+            enc6, num_masks + 1, 1,
+            stride=1, scope="convt7", activation_fn=None)
+        masks = tf.reshape(
+            tf.nn.softmax(tf.reshape(masks, [-1, num_masks + 1])),
+            [batch_size,
+             int(img_height),
+             int(img_width), num_masks + 1])
+        mask_list = tf.split(
+            axis=3, num_or_size_splits=num_masks + 1, value=masks)
+        output = mask_list[0] * prev_image
+        for layer, mask in zip(transformed, mask_list[1:]):
+          output += layer * mask
+        gen_images.append(output)
+
+        current_state = slim.layers.fully_connected(
+            state_action,
+            int(current_state.get_shape()[1]),
+            scope="state_pred",
+            activation_fn=None)
+        gen_states.append(current_state)
+        timestep += 1
+
+    return gen_images, gen_states, latent_mean, latent_std
+
+  def cdna_transformation(self,
+                          prev_image,
+                          cdna_input,
+                          num_masks,
+                          color_channels):
+    """Apply convolutional dynamic neural advection to previous image.
+
+    Args:
+      prev_image: previous image to be transformed.
+      cdna_input: hidden lyaer to be used for computing CDNA kernels.
+      num_masks: number of masks and hence the number of CDNA transformations.
+      color_channels: the number of color channels in the images.
+    Returns:
+      List of images transformed by the predicted CDNA kernels.
+    """
+    batch_size = tf.shape(cdna_input)[0]
+    height = int(prev_image.get_shape()[1])
+    width = int(prev_image.get_shape()[2])
+
+    # Predict kernels using linear function of last hidden layer.
+    cdna_kerns = slim.layers.fully_connected(
+        cdna_input,
+        self.hparams.DNA_KERN_SIZE *
+        self.hparams.DNA_KERN_SIZE * num_masks,
+        scope="cdna_params",
+        activation_fn=None)
+
+    # Reshape and normalize.
+    cdna_kerns = tf.reshape(
+        cdna_kerns, [batch_size, self.hparams.DNA_KERN_SIZE,
+                     self.hparams.DNA_KERN_SIZE, 1, num_masks])
+    cdna_kerns = tf.nn.relu(cdna_kerns - self.hparams.RELU_SHIFT) \
+               + self.hparams.RELU_SHIFT
+    norm_factor = tf.reduce_sum(cdna_kerns, [1, 2, 3], keep_dims=True)
+    cdna_kerns /= norm_factor
+
+    # Treat the color channel dimension as the batch dimension since the same
+    # transformation is applied to each color channel.
+    # Treat the batch dimension as the channel dimension so that
+    # depthwise_conv2d can apply a different transformation to each sample.
+    cdna_kerns = tf.transpose(cdna_kerns, [1, 2, 0, 4, 3])
+    cdna_kerns = tf.reshape(cdna_kerns,
+                            [self.hparams.DNA_KERN_SIZE,
+                             self.hparams.DNA_KERN_SIZE,
+                             batch_size,
+                             num_masks])
+    # Swap the batch and channel dimensions.
+    prev_image = tf.transpose(prev_image, [3, 1, 2, 0])
+
+    # Transform image.
+    transformed = tf.nn.depthwise_conv2d(prev_image, cdna_kerns, [1, 1, 1, 1],
+                                         "SAME")
+
+    # Transpose the dimensions to where they belong.
+    transformed = tf.reshape(
+        transformed, [color_channels, height, width, batch_size, num_masks])
+    transformed = tf.transpose(transformed, [3, 1, 2, 0, 4])
+    transformed = tf.unstack(transformed, axis=-1)
+    return transformed
+
+  def dna_transformation(self,
+                         prev_image,
+                         dna_input):
+    """Apply dynamic neural advection to previous image.
+
+    Args:
+      prev_image: previous image to be transformed.
+      dna_input: hidden lyaer to be used for computing DNA transformation.
+    Returns:
+      List of images transformed by the predicted CDNA kernels.
+    """
+    # Construct translated images.
+    prev_image_pad = tf.pad(prev_image, [[0, 0], [2, 2], [2, 2], [0, 0]])
+    image_height = int(prev_image.get_shape()[1])
+    image_width = int(prev_image.get_shape()[2])
+
+    inputs = []
+    for xkern in range(self.hparams.DNA_KERN_SIZE):
+      for ykern in range(self.hparams.DNA_KERN_SIZE):
+        inputs.append(
+            tf.expand_dims(
+                tf.slice(prev_image_pad, [0, xkern, ykern, 0],
+                         [-1, image_height, image_width, -1]), [3]))
+    inputs = tf.concat(axis=3, values=inputs)
+
+    # Normalize channels to 1.
+    kernel = tf.nn.relu(dna_input -self.hparams.RELU_SHIFT) \
+           + self.hparams.RELU_SHIFT
+    kernel = tf.expand_dims(kernel / tf.reduce_sum(kernel, [3], keep_dims=True),
+                            [4])
+    return tf.reduce_sum(kernel * inputs, [3], keep_dims=False)
+
+  def scheduled_sample(self,
+                       ground_truth_x,
+                       generated_x,
+                       batch_size,
+                       num_ground_truth):
+    """Sample batch with specified mix of groundtruth and generated data points.
+
+    Args:
+      ground_truth_x: tensor of ground-truth data points.
+      generated_x: tensor of generated data points.
+      batch_size: batch size
+      num_ground_truth: number of ground-truth examples to include in batch.
+    Returns:
+      New batch with num_ground_truth sampled from ground_truth_x and the rest
+      from generated_x.
+    """
+    idx = tf.random_shuffle(tf.range(int(batch_size)))
+    ground_truth_idx = tf.gather(idx, tf.range(num_ground_truth))
+    generated_idx = tf.gather(idx, tf.range(num_ground_truth, int(batch_size)))
+
+    ground_truth_examps = tf.gather(ground_truth_x, ground_truth_idx)
+    generated_examps = tf.gather(generated_x, generated_idx)
+    return tf.dynamic_stitch([ground_truth_idx, generated_idx],
+                             [ground_truth_examps, generated_examps])
+
+  def init_state(self,
+                 inputs,
+                 state_shape,
+                 state_initializer=tf.zeros_initializer(),
+                 dtype=tf.float32):
+    """Helper function to create an initial state given inputs.
+
+    Args:
+      inputs: input Tensor, at least 2D, the first dimension being batch_size
+      state_shape: the shape of the state.
+      state_initializer: Initializer(shape, dtype) for state Tensor.
+      dtype: Optional dtype, needed when inputs is None.
+    Returns:
+       A tensors representing the initial state.
+    """
+    # recoded by @mbz
+    initial_state = tf.zeros([tf.shape(inputs)[0]] + state_shape)
+    return initial_state
+
+  # TODO(mbz): use tf.distributions.kl_divergence instead.
+  def kl_divergence(self, mu, log_sigma):
+    """KL divergence of diagonal gaussian N(mu,exp(log_sigma)) and N(0,1).
+
+    Args:
+      mu: mu parameter of the distribution.
+      log_sigma: log(sigma) parameter of the distribution.
+    Returns:
+      the KL loss.
+    """
+
+    return -.5 * tf.reduce_sum(
+        1. + log_sigma - tf.square(mu) - tf.exp(log_sigma),
+        axis=1)
+
+  @slim.add_arg_scope
+  def basic_conv_lstm_cell(self,
+                           inputs,
+                           state,
+                           num_channels,
+                           filter_size=5,
+                           forget_bias=1.0,
+                           scope=None,
+                           reuse=None):
+    """Basic LSTM recurrent network cell, with 2D convolution connctions.
+
+    We add forget_bias (default: 1) to the biases of the forget gate in order to
+    reduce the scale of forgetting in the beginning of the training.
+    It does not allow cell clipping, a projection layer, and does not
+    use peep-hole connections: it is the basic baseline.
+    Args:
+      inputs: input Tensor, 4D, batch x height x width x channels.
+      state: state Tensor, 4D, batch x height x width x channels.
+      num_channels: the number of output channels in the layer.
+      filter_size: the shape of the each convolution filter.
+      forget_bias: the initial value of the forget biases.
+      scope: Optional scope for variable_scope.
+      reuse: whether or not the layer and the variables should be reused.
+    Returns:
+       a tuple of tensors representing output and the new state.
+    """
+    spatial_size = [v.value for v in inputs.get_shape()[1:3]]
+
+    if state is None:
+      state = self.init_state(inputs, spatial_size + [2 * num_channels])
+    with tf.variable_scope(scope,
+                           "BasicConvLstmCell",
+                           [inputs, state],
+                           reuse=reuse):
+      inputs.get_shape().assert_has_rank(4)
+      state.get_shape().assert_has_rank(4)
+      c, h = tf.split(axis=3, num_or_size_splits=2, value=state)
+      inputs_h = tf.concat(axis=3, values=[inputs, h])
+      # Parameters of gates are concatenated into one conv for efficiency.
+      i_j_f_o = slim.layers.conv2d(inputs_h,
+                                   4 * num_channels, [filter_size, filter_size],
+                                   stride=1,
+                                   activation_fn=None,
+                                   scope="Gates")
+
+      # i = input_gate, j = new_input, f = forget_gate, o = output_gate
+      i, j, f, o = tf.split(axis=3, num_or_size_splits=4, value=i_j_f_o)
+
+      new_c = c * tf.sigmoid(f + forget_bias) + tf.sigmoid(i) * tf.tanh(j)
+      new_h = tf.tanh(new_c) * tf.sigmoid(o)
+
+      return new_h, tf.concat(axis=3, values=[new_c, new_h])
+
   def body(self, features):
     hparams = self.hparams
-    filters = hparams.hidden_size
-
-    # Split inputs time-wise into a list of frames. Inputs are by default
-    # concatenated time-wise on channels in VideoModality, so we split on
-    # the last axis. Can do the same for target frames with num_target_frames.
-    # TODO(lukaszkaiser): should we change VideoModality to not concatenate?
-    num_frames = hparams.problem.num_input_frames
-    input_frames = tf.split(features["inputs"], num_frames, axis=-1)
 
-    # For now predict using just a linear transformation of the last frame.
-    # Here input_frames[-1] is contrast-normalized last frame.
-    prediction = tf.layers.dense(input_frames[-1], filters,
-                                 name="final_dense")
-    return prediction
+    # Split inputs and targets time-wise into a list of frames.
+    input_frames = tf.unstack(features["inputs"], axis=1)
+    target_frames = tf.unstack(features["targets"], axis=1)
+
+    num_frames = hparams.problem.num_input_frames + \
+                 hparams.problem.num_target_frames
+    batch_size = tf.shape(input_frames)[0]
+    fake_zeros = [tf.zeros((batch_size, 1), dtype=tf.float32)
+                  for _ in range(num_frames)]
+
+    gen_images, _, latent_mean, latent_std = self.construct_model(
+        images=input_frames + target_frames,
+        actions=fake_zeros,
+        states=fake_zeros,
+        k=900.0 if self.hparams.mode == "training" else -1.0,
+        use_state=False,
+        num_masks=10,
+        cdna=True,
+        dna=False,
+        context_frames=hparams.problem.num_input_frames)
+
+    kl_loss = 0.0
+    step_num = tf.train.get_or_create_global_step()
+    beta = tf.cond(step_num > self.hparams.num_iterations_2nd_stage,
+                   lambda: self.hparams.latent_loss_multiplier,
+                   lambda: 0.0)
+
+    tf.summary.scalar("beta", beta)
+    tf.summary.histogram("posterior_mean", latent_mean)
+    tf.summary.histogram("posterior_std", latent_std)
+
+    if self.hparams.mode == "train":
+      kl_loss = self.kl_divergence(latent_mean, latent_std)
+      tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
+    kl_loss *= beta
+
+    predictions = gen_images[hparams.problem.num_input_frames-1:]
+    return predictions, kl_loss
 
 
 @registry.register_hparams
@@ -212,6 +758,25 @@ def next_frame():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_stochastic():
+  """SV2P model."""
+  hparams = common_hparams.basic_params1()
+  hparams.batch_size = 8
+  hparams.learning_rate_constant = 1e-3
+  hparams.learning_rate_schedule = "constant"
+  hparams.weight_decay = 0.0
+  hparams.add_hparam("stochastic_model", True)
+  hparams.add_hparam("latent_channels", 1)
+  hparams.add_hparam("latent_std_min", -5.0)
+  hparams.add_hparam("num_iterations_2nd_stage", 10000)
+  hparams.add_hparam("latent_loss_multiplier", 1e-4)
+  hparams.add_hparam("multi_latent", False)
+  hparams.add_hparam("RELU_SHIFT", 1e-12)
+  hparams.add_hparam("DNA_KERN_SIZE", 5)
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_tpu():
   hparams = next_frame()
@@ -311,3 +876,4 @@ def next_frame_ae_range(rhp):
   rhp.set_float("learning_rate_constant", 1., 2.)
   rhp.set_float("initializer_gain", 0.8, 1.5)
   rhp.set_int("filter_double_steps", 2, 3)
+

From 015d1a0fa70e7d4680ecad6fc2e6303c60d647b9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 6 Jun 2018 16:46:11 -0700
Subject: [PATCH 0049/2720] Cleanup
 s/softmax_cross_entropy_with_logits/softmax_cross_entropy_with_logits_v2

PiperOrigin-RevId: 199554042
---
 tensor2tensor/layers/common_layers.py            | 2 +-
 tensor2tensor/layers/latent_layers.py            | 2 +-
 tensor2tensor/models/research/transformer_vae.py | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 9f847640b..103943b6a 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1890,7 +1890,7 @@ def smoothing_cross_entropy(logits,
           depth=vocab_size,
           on_value=confidence,
           off_value=low_confidence)
-    xentropy = tf.nn.softmax_cross_entropy_with_logits(
+    xentropy = tf.nn.softmax_cross_entropy_with_logits_v2(
         logits=logits, labels=soft_targets)
     return xentropy - normalizing
 
diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index bd8824414..0ea5b30f4 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -98,7 +98,7 @@ def ae_latent_softmax(latents_pred, latents_discrete, hparams):
         if hparams.soft_em:
           # latents_discrete is actually one-hot of multinomial samples
           assert hparams.num_decode_blocks == 1
-          loss = tf.nn.softmax_cross_entropy_with_logits(
+          loss = tf.nn.softmax_cross_entropy_with_logits_v2(
               labels=latents_discrete, logits=latents_logits)
         else:
           loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 94cc0b67c..dcf19b5d6 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -229,7 +229,7 @@ def ae_latent_softmax(latents_pred, latents_discrete, hparams):
       if hparams.soft_em:
         # latents_discrete is actually one-hot of multinomial samples
         assert hparams.num_decode_blocks == 1
-        loss = tf.nn.softmax_cross_entropy_with_logits(
+        loss = tf.nn.softmax_cross_entropy_with_logits_v2(
             labels=latents_discrete, logits=latents_logits)
       else:
         loss = tf.nn.sparse_softmax_cross_entropy_with_logits(

From f8cf3582d2b2f7a30e69f02a31da2acd44e7f7c6 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 6 Jun 2018 17:04:53 -0700
Subject: [PATCH 0050/2720] Make stochastic video model compatible with current
 eval metrics.

PiperOrigin-RevId: 199557602
---
 tensor2tensor/data_generators/gym_problems.py |  5 +-
 .../data_generators/video_generated.py        |  6 ++
 tensor2tensor/data_generators/video_utils.py  | 12 +++-
 tensor2tensor/layers/modalities.py            | 18 ++---
 tensor2tensor/models/research/next_frame.py   | 66 +++++++++----------
 5 files changed, 56 insertions(+), 51 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index a13020d9a..5c8669cd6 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -395,8 +395,7 @@ def autoencode_tensor(self, x, batch_size=1):
       autoencoded = self.autoencoder_model.encode(
           tf.reshape(x, [batch_size, 1] + shape))
     autoencoded = tf.reshape(
-        autoencoded, [batch_size, self.frame_height, self.frame_width,
-                      self.num_channels, 8])  # 8-bit groups.
+        autoencoded, [batch_size] + self.frame_shape + [8])  # 8-bit groups.
     if batch_size == 1:
       autoencoded = tf.squeeze(autoencoded, axis=0)
     return discretization.bit_to_int(autoencoded, 8)
@@ -447,7 +446,7 @@ def _setup(self):
         self.autoencoder_feed = tf.placeholder(tf.int32, shape=shape)
         self.autoencoder_result = self.autoencode_tensor(self.autoencoder_feed)
         # Now for autodecoding.
-        shape = [self.frame_height, self.frame_width, self.num_channels]
+        shape = self.frame_shape
         self.autodecoder_feed = tf.placeholder(tf.int32, shape=shape)
         bottleneck = tf.reshape(
             discretization.int_to_bit(self.autodecoder_feed, 8),
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 6aee3d787..fd60bd5c6 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -24,6 +24,7 @@
 
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
+from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -73,6 +74,11 @@ def video_length(self):
   def random_skip(self):
     return False
 
+  def eval_metrics(self):
+    eval_metrics = [metrics.Metrics.ACC, metrics.Metrics.ACC_PER_SEQ,
+                    metrics.Metrics.IMAGE_RMSE]
+    return eval_metrics
+
   @property
   def dataset_splits(self):
     """Splits of data to produce and number of output shards for each."""
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index c3e928eef..8a2e2c113 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -64,6 +64,11 @@ def frame_width(self):
     """Width of each frame."""
     raise NotImplementedError
 
+  @property
+  def frame_shape(self):
+    """Shape of a frame: a list [height , width , channels]."""
+    return [self.frame_height, self.frame_width, self.num_channels]
+
   @property
   def total_number_of_frames(self):
     """The total number of frames, needed for sharding."""
@@ -79,6 +84,11 @@ def num_target_frames(self):
     """Number of frames to batch on one target."""
     return 1
 
+  @property
+  def num_input_and_target_frames(self):
+    """Number of frames on input and target added."""
+    return self.num_input_frames + self.num_target_frames
+
   @property
   def random_skip(self):
     """Whether to skip random inputs at the beginning or not."""
@@ -194,7 +204,7 @@ def _preprocess(example):
       return self.preprocess_example(example, mode, hparams)
     preprocessed_dataset = dataset.map(_preprocess)
 
-    num_frames = self.num_input_frames + self.num_target_frames
+    num_frames = self.num_input_and_target_frames
     # We jump by a random position at the beginning to add variety.
     if self.random_skip:
       random_skip = tf.random_uniform([], maxval=num_frames, dtype=tf.int64)
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 74772cf11..506b78d4d 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -556,23 +556,15 @@ def top(self, body_output, _):
     frames = tf.stack(body_output, axis=1)
     rgb_frames = common_layers.convert_real_to_rgb(frames)
     common_layers.summarize_video(rgb_frames, "body_output")
-    return frames
+    # TODO(lukaszkaiser): remove the need for the last dimension of 1 in eval.
+    return tf.expand_dims(rgb_frames, axis=-1)
 
   def loss(self, top_out, targets):
+    top_out = tf.squeeze(top_out, axis=[-1])
     assert(top_out.shape.as_list() == targets.shape.as_list()), \
            "The dimensions doesn't match."
-
-    common_layers.summarize_video(targets, "targets_top")
-    targets = common_layers.convert_rgb_to_real(targets)
-
-    num_frames = top_out.shape[1].value
-    loss = 0.0
-    for frame_id in range(num_frames):
-      frame = tf.to_float(top_out[:, frame_id])
-      target = tf.to_float(targets[:, frame_id])
-      loss += tf.reduce_mean(tf.square(frame - target))
-    loss /= num_frames
-    return loss, tf.zeros_like(loss)
+    loss = tf.square(top_out - tf.to_float(targets))
+    return tf.reduce_sum(loss), tf.reduce_sum(tf.ones_like(loss))
 
 
 @registry.register_video_modality("embed")
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 3b4b9b8e5..569b747d8 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -145,13 +145,13 @@ def logits_to_samples(logits):
 
     # Get predictions.
     try:
-      num_channels = self._hparams.problem.num_channels
+      num_channels = self.hparams.problem.num_channels
     except AttributeError:
       num_channels = 1
     features["targets"] = tf.zeros(
-        [self._hparams.batch_size, 1, 1, 1, num_channels], dtype=tf.int32)
+        [self.hparams.batch_size, 1, 1, 1, num_channels], dtype=tf.int32)
     features["target_reward"] = tf.zeros(
-        [self._hparams.batch_size, 1, 1], dtype=tf.int32)
+        [self.hparams.batch_size, 1, 1], dtype=tf.int32)
     logits, _ = self(features)  # pylint: disable=not-callable
     if isinstance(logits, dict):
       results = {}
@@ -242,7 +242,7 @@ def construct_latent_tower(self, images):
       # batch_size x latent_size
       samples = tf.random_normal(tf.shape(latent_mean), 0, 1, dtype=tf.float32)
 
-    if self.hparams.mode == "train":
+    if self.hparams.mode == tf.estimator.ModeKeys.TRAIN:
       return latent_mean, latent_std, samples
     else:
       # No latent tower at inference time, just standard gaussian.
@@ -288,9 +288,8 @@ def construct_model(self,
     if cdna + dna != 1:
       raise ValueError("More than one, or no network option specified.")
 
-    batch_size, img_height, img_width, color_channels = \
-    images[0].get_shape()[0:4]
-    batch_size = tf.shape(images[0])[0]
+    img_height, img_width, color_channels = self.hparams.problem.frame_shape
+    batch_size = common_layers.shape_list(images[0])[0]
     lstm_func = self.basic_conv_lstm_cell
 
     # Generated robot states and images.
@@ -341,7 +340,7 @@ def construct_model(self,
         elif done_warm_start:
           # Scheduled sampling
           prev_image = self.scheduled_sample(
-              image, gen_images[-1], batch_size, num_ground_truth)
+              image, gen_images[-1], self.hparams.batch_size, num_ground_truth)
         else:
           # Always feed in ground_truth
           prev_image = image
@@ -390,7 +389,7 @@ def construct_model(self,
           latent = samples
           if self.hparams.multi_latent:
             latent = samples[timestep]
-          if self.hparams.mode == "train":
+          if self.hparams.mode == tf.estimator.ModeKeys.TRAIN:
             # TODO(mbz): put 1st stage of training back in if necessary
             latent = latent_mean + tf.exp(latent_std / 2.0) * latent
           with tf.control_dependencies([latent]):
@@ -434,7 +433,7 @@ def construct_model(self,
           # Using largest hidden state for predicting untied conv kernels.
           enc7 = slim.layers.conv2d_transpose(
               enc6,
-              self.hparams.DNA_KERN_SIZE**2,
+              self.hparams.dna_kernel_size**2,
               1,
               stride=1,
               scope="convt4",
@@ -510,17 +509,17 @@ def cdna_transformation(self,
     # Predict kernels using linear function of last hidden layer.
     cdna_kerns = slim.layers.fully_connected(
         cdna_input,
-        self.hparams.DNA_KERN_SIZE *
-        self.hparams.DNA_KERN_SIZE * num_masks,
+        self.hparams.dna_kernel_size *
+        self.hparams.dna_kernel_size * num_masks,
         scope="cdna_params",
         activation_fn=None)
 
     # Reshape and normalize.
     cdna_kerns = tf.reshape(
-        cdna_kerns, [batch_size, self.hparams.DNA_KERN_SIZE,
-                     self.hparams.DNA_KERN_SIZE, 1, num_masks])
-    cdna_kerns = tf.nn.relu(cdna_kerns - self.hparams.RELU_SHIFT) \
-               + self.hparams.RELU_SHIFT
+        cdna_kerns, [batch_size, self.hparams.dna_kernel_size,
+                     self.hparams.dna_kernel_size, 1, num_masks])
+    cdna_kerns = (tf.nn.relu(cdna_kerns - self.hparams.relu_shift)
+                  + self.hparams.relu_shift)
     norm_factor = tf.reduce_sum(cdna_kerns, [1, 2, 3], keep_dims=True)
     cdna_kerns /= norm_factor
 
@@ -530,8 +529,8 @@ def cdna_transformation(self,
     # depthwise_conv2d can apply a different transformation to each sample.
     cdna_kerns = tf.transpose(cdna_kerns, [1, 2, 0, 4, 3])
     cdna_kerns = tf.reshape(cdna_kerns,
-                            [self.hparams.DNA_KERN_SIZE,
-                             self.hparams.DNA_KERN_SIZE,
+                            [self.hparams.dna_kernel_size,
+                             self.hparams.dna_kernel_size,
                              batch_size,
                              num_masks])
     # Swap the batch and channel dimensions.
@@ -565,8 +564,8 @@ def dna_transformation(self,
     image_width = int(prev_image.get_shape()[2])
 
     inputs = []
-    for xkern in range(self.hparams.DNA_KERN_SIZE):
-      for ykern in range(self.hparams.DNA_KERN_SIZE):
+    for xkern in range(self.hparams.dna_kernel_size):
+      for ykern in range(self.hparams.dna_kernel_size):
         inputs.append(
             tf.expand_dims(
                 tf.slice(prev_image_pad, [0, xkern, ykern, 0],
@@ -574,8 +573,8 @@ def dna_transformation(self,
     inputs = tf.concat(axis=3, values=inputs)
 
     # Normalize channels to 1.
-    kernel = tf.nn.relu(dna_input -self.hparams.RELU_SHIFT) \
-           + self.hparams.RELU_SHIFT
+    kernel = (tf.nn.relu(dna_input -self.hparams.relu_shift)
+              + self.hparams.relu_shift)
     kernel = tf.expand_dims(kernel / tf.reduce_sum(kernel, [3], keep_dims=True),
                             [4])
     return tf.reduce_sum(kernel * inputs, [3], keep_dims=False)
@@ -699,17 +698,16 @@ def body(self, features):
     input_frames = tf.unstack(features["inputs"], axis=1)
     target_frames = tf.unstack(features["targets"], axis=1)
 
-    num_frames = hparams.problem.num_input_frames + \
-                 hparams.problem.num_target_frames
-    batch_size = tf.shape(input_frames)[0]
+    num_frames = hparams.problem.num_input_and_target_frames
+    batch_size = common_layers.shape_list(input_frames)[0]
     fake_zeros = [tf.zeros((batch_size, 1), dtype=tf.float32)
                   for _ in range(num_frames)]
-
+    is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
     gen_images, _, latent_mean, latent_std = self.construct_model(
         images=input_frames + target_frames,
         actions=fake_zeros,
         states=fake_zeros,
-        k=900.0 if self.hparams.mode == "training" else -1.0,
+        k=900.0 if is_training else -1.0,
         use_state=False,
         num_masks=10,
         cdna=True,
@@ -722,11 +720,12 @@ def body(self, features):
                    lambda: self.hparams.latent_loss_multiplier,
                    lambda: 0.0)
 
-    tf.summary.scalar("beta", beta)
-    tf.summary.histogram("posterior_mean", latent_mean)
-    tf.summary.histogram("posterior_std", latent_std)
+    if is_training:
+      tf.summary.scalar("beta", beta)
+      tf.summary.histogram("posterior_mean", latent_mean)
+      tf.summary.histogram("posterior_std", latent_std)
 
-    if self.hparams.mode == "train":
+    if is_training:
       kl_loss = self.kl_divergence(latent_mean, latent_std)
       tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
     kl_loss *= beta
@@ -772,8 +771,8 @@ def next_frame_stochastic():
   hparams.add_hparam("num_iterations_2nd_stage", 10000)
   hparams.add_hparam("latent_loss_multiplier", 1e-4)
   hparams.add_hparam("multi_latent", False)
-  hparams.add_hparam("RELU_SHIFT", 1e-12)
-  hparams.add_hparam("DNA_KERN_SIZE", 5)
+  hparams.add_hparam("relu_shift", 1e-12)
+  hparams.add_hparam("dna_kernel_size", 5)
   return hparams
 
 
@@ -876,4 +875,3 @@ def next_frame_ae_range(rhp):
   rhp.set_float("learning_rate_constant", 1., 2.)
   rhp.set_float("initializer_gain", 0.8, 1.5)
   rhp.set_int("filter_double_steps", 2, 3)
-

From 74ff7072c069b01cfa85bad7da7bf6770fd2802f Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 6 Jun 2018 20:12:05 -0700
Subject: [PATCH 0051/2720] Change weight_decay range in hparams
 transformer_base_range because the default is quite large.

PiperOrigin-RevId: 199575307
---
 tensor2tensor/models/transformer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 86ed71eb7..1c7ca4553 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1422,7 +1422,7 @@ def transformer_base_range(rhp):
   rhp.set_float("initializer_gain", 0.5, 2.0)
   rhp.set_float("optimizer_adam_beta1", 0.85, 0.95)
   rhp.set_float("optimizer_adam_beta2", 0.97, 0.99)
-  rhp.set_float("weight_decay", 0.0, 2.0)
+  rhp.set_float("weight_decay", 0.0, 1e-4)
 
 
 @registry.register_hparams

From 9cd8f61e2e02ed31fd838a934e7894c5265910b4 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 7 Jun 2018 10:30:30 -0700
Subject: [PATCH 0052/2720] Internal change

PiperOrigin-RevId: 199656880
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 .../data_generators/program_search.py         | 124 ++++++++++++++++++
 .../data_generators/program_search_test.py    | 111 ++++++++++++++++
 3 files changed, 236 insertions(+)
 create mode 100644 tensor2tensor/data_generators/program_search.py
 create mode 100644 tensor2tensor/data_generators/program_search_test.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 93deecf61..98dbc2740 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -43,6 +43,7 @@
     "tensor2tensor.data_generators.mnist",
     "tensor2tensor.data_generators.mscoco",
     "tensor2tensor.data_generators.multinli",
+    "tensor2tensor.data_generators.program_search",
     "tensor2tensor.data_generators.ocr",
     "tensor2tensor.data_generators.problem_hparams",
     "tensor2tensor.data_generators.ptb",
diff --git a/tensor2tensor/data_generators/program_search.py b/tensor2tensor/data_generators/program_search.py
new file mode 100644
index 000000000..8ba661b86
--- /dev/null
+++ b/tensor2tensor/data_generators/program_search.py
@@ -0,0 +1,124 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Program Search Problems."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gzip
+import json
+import os
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+@registry.register_problem
+class ProgramSearchAlgolisp(text_problems.Text2TextProblem):
+  """Problem class for Program Search Algolisp task.
+
+  Synthesizing programs from description and examples.
+
+  Please see: https://arxiv.org/pdf/1802.04335.pdf for the full description.
+  """
+
+  # The locations of the train, dev, and test set.
+  DROPBOX = "https://www.dropbox.com"
+  DATA_URLS = {
+      problem.DatasetSplit.TRAIN: (
+          DROPBOX + "/s/wep81pcrar5fttl/metaset3.train.jsonl.gz?dl=1"),
+      problem.DatasetSplit.EVAL: (
+          DROPBOX + "/s/h3mn0abeiqy6foz/metaset3.dev.jsonl.gz?dl=1"),
+      problem.DatasetSplit.TEST: (
+          DROPBOX + "/s/f1x9ybkjpf371cp/metaset3.test.jsonl.gz?dl=1"),
+  }
+
+  @staticmethod
+  def _extract_filename_from_url(url):
+    # Ex: TRAIN_URL --> metaset3.train.jsonl.gz
+
+    # Get everything from the last / onwards.
+    filename = os.path.basename(url)
+
+    # Get rid of everything after the first ?
+    return filename.split("?")[0]
+
+  @staticmethod
+  def _flatten_target_programs(iterable):
+    # The target programs are read as nested lists, we should flatten them.
+    yield unicode("[")
+    it = iter(iterable)
+    for e in it:
+      if isinstance(e, (list, tuple)):
+        for f in ProgramSearchAlgolisp._flatten_target_programs(e):
+          yield f
+      else:
+        yield e
+    yield unicode("]")
+
+  @staticmethod
+  def _parse_json_to_dict(json_line):
+    # First parse it through json.
+    line_json_dict = json.loads(json_line)
+
+    # The features of interest "text" and "short_tree" are stored as lists in
+    # this dictionary -- "short_tree" is a nested list. We flatten and join the
+    # lists on space, to return a string in both these cases.
+
+    # Make another dictionary, to return only the features we want.
+    return {
+        "inputs":
+            " ".join(line_json_dict["text"]),
+        "targets":
+            " ".join([
+                i for i in ProgramSearchAlgolisp._flatten_target_programs(
+                    line_json_dict["short_tree"])
+            ])
+    }
+
+  @property
+  def is_generate_per_split(self):
+    # Return True since we already have the train and the dev set separated out.
+    return True
+
+  def maybe_download_dataset(self, tmp_dir, dataset_split):
+    """Downloads the appropriate dataset file and returns its path."""
+    # Get the dataset url for the split requested.
+    url = self.DATA_URLS.get(dataset_split, None)
+
+    # Sanity check.
+    if url is None:
+      tf.logging.fatal("Unknown dataset_split passed: {}".format(dataset_split))
+
+    # Download the data, if it doesn't already exist.
+    return generator_utils.maybe_download(tmp_dir,
+                                          self._extract_filename_from_url(url),
+                                          url)
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    del data_dir
+
+    # Download the data, if it doesn't already exist.
+    downloaded_filepath = self.maybe_download_dataset(tmp_dir, dataset_split)
+
+    # Decompress the file and iterate through it.
+    with gzip.open(downloaded_filepath, "rb") as data_fp:
+      for line in data_fp:
+        yield self._parse_json_to_dict(line.strip())
diff --git a/tensor2tensor/data_generators/program_search_test.py b/tensor2tensor/data_generators/program_search_test.py
new file mode 100644
index 000000000..84f4a1e03
--- /dev/null
+++ b/tensor2tensor/data_generators/program_search_test.py
@@ -0,0 +1,111 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.data_generators.program_search."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gzip
+import os
+import shutil
+import tempfile
+
+from builtins import bytes  # pylint: disable=redefined-builtin
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import program_search
+
+import tensorflow as tf
+
+
+class ProgramSearchAlgolispStub(program_search.ProgramSearchAlgolisp):
+  """Stub of ProgramSearchAlgolisp that stubs out maybe_download_dataset.
+
+  The maybe_download_dataset writes one predetermined example in a zip file
+  self.n number of times and returns the file path.
+  """
+
+  EXAMPLE = ('{"funcs": [], "tests": [{"output": 0, "input": {"a": 5}}, '
+             '{"output": 1, "input": {"a": 20}}, {"output": 2, "input": '
+             '{"a": 28}}, {"output": 1, "input": {"a": 13}}, {"output": 1, '
+             '"input": {"a": 27}}, {"output": 1, "input": {"a": 13}}, '
+             '{"output": 1, "input": {"a": 20}}, {"output": 0, '
+             '"input": {"a": 8}}, {"output": 0, "input": {"a": 8}}, '
+             '{"output": 0, "input": {"a": 4}}], "short_tree": ["invoke1", '
+             '["lambda1", ["if", ["==", ["len", ["digits", "arg1"]], "1"], "0",'
+             ' ["+", "1", ["self", ["reduce", ["digits", "arg1"], "0", '
+             '"+"]]]]], "a"], "tags": [], "text": ["given", "a", "number", "a",'
+             ' ",", "find", "how", "many", "times", "you", "can", "replace", '
+             '"a", "with", "sum", "of", "its", "digits", "before", "it", '
+             '"becomes", "a", "single", "digit", "number"], "return_type": '
+             '"int", "args": {"a": "int"}, "nodes": ["l1_recursive_digits"]}')
+
+  EXAMPLE_INPUT = ('given a number a , find how many times you can replace a '
+                   'with sum of its digits before it becomes a single digit '
+                   'number')
+
+  EXAMPLE_TARGET = ('[ invoke1 [ lambda1 [ if [ == [ len [ digits arg1 ] ] 1 ]'
+                    ' 0 [ + 1 [ self [ reduce [ digits arg1 ] 0 + ] ] ] ] ] a '
+                    ']')
+
+  N = 10
+
+  def maybe_download_dataset(self, tmp_dir, dataset_split):
+    (_, data_file) = tempfile.mkstemp(
+        suffix='.gz', prefix=str(dataset_split) + '-', dir=tmp_dir)
+
+    with gzip.open(data_file, 'wb') as gz_file:
+      content = '\n'.join([self.EXAMPLE] * self.N)
+      gz_file.write(bytes(content, 'utf-8'))
+    return data_file
+
+
+class ProgramSearchAlgolispTest(tf.test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    # Setup the temp directory tree.
+    cls.tmp_dir = tf.test.get_temp_dir()
+    shutil.rmtree(cls.tmp_dir)
+    os.mkdir(cls.tmp_dir)
+
+  @classmethod
+  def tearDownClass(cls):
+    # Cleanup the temp directory tree.
+    shutil.rmtree(cls.tmp_dir)
+
+  def testEndToEnd(self):
+    # End-to-end test, the stub problem class creates a .gz file with nps_stub.N
+    # example and we check if we're able to process it correctly.
+    nps_stub = ProgramSearchAlgolispStub()
+    num = 0
+    for example in nps_stub.generate_samples(None, self.tmp_dir,
+                                             problem.DatasetSplit.TRAIN):
+
+      # Only one example in 'file', so this is OK.
+      self.assertEqual(example['inputs'],
+                       ProgramSearchAlgolispStub.EXAMPLE_INPUT)
+
+      self.assertEqual(example['targets'],
+                       ProgramSearchAlgolispStub.EXAMPLE_TARGET)
+
+      num += 1
+
+    # assert that we have as many examples as there are in the file.
+    self.assertEqual(num, nps_stub.N)
+
+
+if __name__ == '__main__':
+  tf.test.main()

From 6969fab42200a7da11bc40c9537b76b0a204b46a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 7 Jun 2018 11:39:59 -0700
Subject: [PATCH 0053/2720] Move input and target frame batching for video to
 hparams to make it tunable easier.

PiperOrigin-RevId: 199669699
---
 tensor2tensor/data_generators/gym_problems.py | 17 ++---------
 .../data_generators/video_generated.py        | 12 +-------
 tensor2tensor/data_generators/video_utils.py  | 29 +++++--------------
 tensor2tensor/layers/common_hparams.py        |  3 ++
 tensor2tensor/layers/modalities.py            |  4 +--
 tensor2tensor/models/research/next_frame.py   | 11 +++++--
 tensor2tensor/rl/envs/simulated_batch_env.py  |  4 +--
 7 files changed, 26 insertions(+), 54 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 5c8669cd6..61214bb14 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -63,14 +63,11 @@ def __init__(self, *args, **kwargs):
 
   @property
   def num_input_frames(self):
-    """Number of frames to batch on one input."""
+    """Number of frames on input for real environment."""
+    # TODO(lukaszkaiser): This must be equal to hparams.video_num_input_frames,
+    # we should automate this to avoid bug in the future.
     return 4
 
-  @property
-  def num_target_frames(self):
-    """Number of frames to batch on one target."""
-    return 1
-
   def eval_metrics(self):
     eval_metrics = [metrics.Metrics.ACC, metrics.Metrics.ACC_PER_SEQ,
                     metrics.Metrics.IMAGE_RMSE]
@@ -732,14 +729,6 @@ def frame_width(self):
       return 160
     return int(math.ceil(160 / self.autoencoder_factor))
 
-  @property
-  def raw_frame_height(self):
-    return self.frame_height
-
-  @property
-  def raw_frame_width(self):
-    return self.frame_width
-
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnFreewayAe(  # with autoencoder
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index fd60bd5c6..2ce9c318b 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -39,16 +39,6 @@
 class VideoStochasticShapes10k(video_utils.VideoProblem):
   """Shapes moving in a stochastic way."""
 
-  @property
-  def num_input_frames(self):
-    """Number of frames to batch on one input."""
-    return 1
-
-  @property
-  def num_target_frames(self):
-    """Number of frames to predict in one step."""
-    return 4
-
   @property
   def is_generate_per_split(self):
     """Whether we have a train/test split or just hold out data."""
@@ -68,7 +58,7 @@ def total_number_of_frames(self):
 
   @property
   def video_length(self):
-    return self.num_input_frames + self.num_target_frames
+    return 5
 
   @property
   def random_skip(self):
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 8a2e2c113..d5a06435a 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -74,21 +74,6 @@ def total_number_of_frames(self):
     """The total number of frames, needed for sharding."""
     raise NotImplementedError
 
-  @property
-  def num_input_frames(self):
-    """Number of frames to batch on one input."""
-    return 1
-
-  @property
-  def num_target_frames(self):
-    """Number of frames to batch on one target."""
-    return 1
-
-  @property
-  def num_input_and_target_frames(self):
-    """Number of frames on input and target added."""
-    return self.num_input_frames + self.num_target_frames
-
   @property
   def random_skip(self):
     """Whether to skip random inputs at the beginning or not."""
@@ -157,7 +142,7 @@ def preprocess(self, dataset, mode, hparams):
     def split_on_batch(x):
       """Split x on batch dimension into x[:size, ...] and x[size:, ...]."""
       length = len(x.get_shape())
-      size = self.num_input_frames
+      size = hparams.video_num_input_frames
       if length < 1:
         raise ValueError("Batched tensor of length < 1.")
       if length == 1:
@@ -187,10 +172,10 @@ def features_from_batch(batched_prefeatures):
         if k == "frame":  # We rename past frames to inputs and targets.
           s1, s2 = split_on_batch(v)
           # Reshape just to make sure shapes are right and set.
-          s1 = tf.reshape(s1, [self.num_input_frames, self.frame_height,
-                               self.frame_width, self.num_channels])
-          s2 = tf.reshape(s2, [self.num_target_frames, self.frame_height,
-                               self.frame_width, self.num_channels])
+          s1 = tf.reshape(
+              s1, [hparams.video_num_input_frames] + self.frame_shape)
+          s2 = tf.reshape(
+              s2, [hparams.video_num_target_frames] + self.frame_shape)
           features["inputs"] = s1
           features["targets"] = s2
         else:
@@ -203,8 +188,8 @@ def features_from_batch(batched_prefeatures):
     def _preprocess(example):
       return self.preprocess_example(example, mode, hparams)
     preprocessed_dataset = dataset.map(_preprocess)
-
-    num_frames = self.num_input_and_target_frames
+    num_frames = (hparams.video_num_input_frames +
+                  hparams.video_num_target_frames)
     # We jump by a random position at the beginning to add variety.
     if self.random_skip:
       random_skip = tf.random_uniform([], maxval=num_frames, dtype=tf.int64)
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 986a220e3..bc12b2629 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -185,6 +185,9 @@ def basic_params1():
       # examples.  e.g.  The examples may be written with length 65536, but we
       # want to split each example into 64 examples of length 1024.
       split_to_length=0,
+      # Video settings: how many frames to batch on input and targets.
+      video_num_input_frames=1,
+      video_num_target_frames=1,
       # This flag allows us to optionally treat a seq-to-seq problem
       # as a language model.  Legal values are:
       #
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 506b78d4d..c1c307651 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -511,7 +511,7 @@ def targets_bottom(self, x, summary_prefix="targets_bottom"):  # pylint: disable
 
   def top(self, body_output, _):
     num_channels = self._model_hparams.problem.num_channels
-    num_frames = self._model_hparams.problem.num_target_frames
+    num_frames = self._model_hparams.video_num_target_frames
     with tf.variable_scope("rgb_softmax"):
       body_output_shape = common_layers.shape_list(body_output)
       reshape_shape = body_output_shape[:3]
@@ -613,7 +613,7 @@ class VideoModalityL1(VideoModality):
 
   def top(self, body_output, _):
     num_channels = self._model_hparams.problem.num_channels
-    num_frames = self._model_hparams.problem.num_target_frames
+    num_frames = self._model_hparams.video_num_target_frames
     with tf.variable_scope("rgb"):
       body_output_shape = common_layers.shape_list(body_output)
       res = tf.layers.dense(body_output, num_channels * num_frames, name="cast")
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 569b747d8..1c7b54ae3 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -698,7 +698,8 @@ def body(self, features):
     input_frames = tf.unstack(features["inputs"], axis=1)
     target_frames = tf.unstack(features["targets"], axis=1)
 
-    num_frames = hparams.problem.num_input_and_target_frames
+    num_frames = (hparams.video_num_input_frames +
+                  hparams.video_num_target_frames)
     batch_size = common_layers.shape_list(input_frames)[0]
     fake_zeros = [tf.zeros((batch_size, 1), dtype=tf.float32)
                   for _ in range(num_frames)]
@@ -712,7 +713,7 @@ def body(self, features):
         num_masks=10,
         cdna=True,
         dna=False,
-        context_frames=hparams.problem.num_input_frames)
+        context_frames=hparams.video_num_input_frames)
 
     kl_loss = 0.0
     step_num = tf.train.get_or_create_global_step()
@@ -730,7 +731,7 @@ def body(self, features):
       tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
     kl_loss *= beta
 
-    predictions = gen_images[hparams.problem.num_input_frames-1:]
+    predictions = gen_images[hparams.video_num_input_frames-1:]
     return predictions, kl_loss
 
 
@@ -738,6 +739,8 @@ def body(self, features):
 def next_frame():
   """Basic 2-frame conv model."""
   hparams = common_hparams.basic_params1()
+  hparams.video_num_input_frames = 4
+  hparams.video_num_target_frames = 1
   hparams.hidden_size = 64
   hparams.batch_size = 4
   hparams.num_hidden_layers = 2
@@ -761,6 +764,8 @@ def next_frame():
 def next_frame_stochastic():
   """SV2P model."""
   hparams = common_hparams.basic_params1()
+  hparams.video_num_input_frames = 1
+  hparams.video_num_target_frames = 4
   hparams.batch_size = 8
   hparams.learning_rate_constant = 1e-3
   hparams.learning_rate_schedule = "constant"
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index eac9834d6..2ed6f2b65 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -122,11 +122,11 @@ def __init__(self, environment_lambda, length, problem,
 
     if simulation_random_starts:
       dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
-                                shuffle_files=True)
+                                shuffle_files=True, hparams=hparams)
       dataset = dataset.shuffle(buffer_size=100)
     else:
       dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
-                                shuffle_files=False).take(1)
+                                shuffle_files=False, hparams=hparams).take(1)
 
     dataset = dataset.map(lambda x: x["inputs"]).repeat()
     self.history_buffer = HistoryBuffer(dataset, self.length)

From 480de38319f97f90193464a1c728f4e827aecf4c Mon Sep 17 00:00:00 2001
From: Prasasto Adi <prasastoadi@users.noreply.github.com>
Date: Fri, 8 Jun 2018 04:26:47 +0700
Subject: [PATCH 0054/2720] Add English-Indonesian (#851)

---
 tensor2tensor/data_generators/all_problems.py |  1 +
 .../data_generators/translate_enid.py         | 85 +++++++++++++++++++
 2 files changed, 86 insertions(+)
 create mode 100644 tensor2tensor/data_generators/translate_enid.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 98dbc2740..caaf61069 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -54,6 +54,7 @@
     "tensor2tensor.data_generators.translate_ende",
     "tensor2tensor.data_generators.translate_enet",
     "tensor2tensor.data_generators.translate_enfr",
+    "tensor2tensor.data_generators.translate_enid",
     "tensor2tensor.data_generators.translate_enmk",
     "tensor2tensor.data_generators.translate_envi",
     "tensor2tensor.data_generators.translate_enzh",
diff --git a/tensor2tensor/data_generators/translate_enid.py b/tensor2tensor/data_generators/translate_enid.py
new file mode 100644
index 000000000..a1a94fb49
--- /dev/null
+++ b/tensor2tensor/data_generators/translate_enid.py
@@ -0,0 +1,85 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generators for En-Id translation."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# Dependency imports
+
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import translate
+from tensor2tensor.utils import registry
+
+# End-of-sentence marker.
+EOS = text_encoder.EOS_ID
+
+# IWSLT17 :
+# 	109335 sentences
+#	https://wit3.fbk.eu/mt.php?release=2017-01-more
+# PANL-BPPT :
+#	24024 sentences
+# 	http://www.panl10n.net/english/outputs/Indonesia/BPPT/0902/BPPTIndToEngCorpusHalfM.zip
+_ENID_TRAIN_DATASETS = [
+    [
+        "https://github.com/prasastoadi/parallel-corpora-en-id/raw/master/IWSLT17.train.en-id.tok.tgz",
+        ("IWSLT17.train.en-id.tok.en", "IWSLT17.train.en-id.tok.id")
+    ],
+    [
+        "https://github.com/prasastoadi/parallel-corpora-en-id/raw/master/PANL-BPPT-ECO-EN-ID-150Kw.tok.tgz",
+        ("PANL-BPPT-ECO-EN-150Kw.tok.txt", "PANL-BPPT-ECO-ID-150Kw.tok.txt")
+    ],
+    [
+        "https://github.com/prasastoadi/parallel-corpora-en-id/raw/master/PANL-BPPT-INT-EN-ID-150Kw.tok.tgz",
+        ("PANL-BPPT-INT-EN-150Kw.tok.txt", "PANL-BPPT-INT-ID-150Kw.tok.txt")
+    ],
+    [
+        "https://github.com/prasastoadi/parallel-corpora-en-id/raw/master/PANL-BPPT-SCI-EN-ID-100Kw.tok.tgz",
+        ("PANL-BPPT-SCI-EN-100Kw.tok.txt", "PANL-BPPT-SCI-ID-100Kw.tok.txt")
+    ],
+    [
+        "https://github.com/prasastoadi/parallel-corpora-en-id/raw/master/PANL-BPPT-SPO-EN-ID-100Kw.tok.tgz",
+        ("PANL-BPPT-SPO-EN-100Kw.tok.txt", "PANL-BPPT-SPO-ID-100Kw.tok.txt")
+    ],
+]
+
+# IWSLT17 :
+# 1478 sentences
+# https://wit3.fbk.eu/mt.php?release=2017-01-more
+_ENID_TEST_DATASETS = [
+    [
+        "https://github.com/prasastoadi/parallel-corpora-en-id/raw/master/IWSLT17.TED.tst2017plus.en-id.tok.tgz",
+        ("IWSLT17.TED.tst2017plus.en-id.tok.en", "IWSLT17.TED.tst2017plus.en-id.tok.id")
+    ]
+]
+
+
+@registry.register_problem
+class TranslateEnidIwslt32k(translate.TranslateProblem):
+  """Problem spec for IWSLT'15 En-Vi translation."""
+
+  @property
+  def approx_vocab_size(self):
+    return 2**15  # 32768
+
+  @property
+  def vocab_filename(self):
+    return "vocab.enid.%d" % self.approx_vocab_size
+
+  def source_data_files(self, dataset_split):
+    train = dataset_split == problem.DatasetSplit.TRAIN
+    return _ENID_TRAIN_DATASETS if train else _ENID_TEST_DATASETS

From 1d643e4c789a8a0dd419e24c8130c78e4f247637 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Thu, 7 Jun 2018 14:29:36 -0700
Subject: [PATCH 0055/2720] Wrap constant learning rate as Tensor (#842)

* Wrap learning rate as tf Variable.

* Wrap constant learning rate as Tensor.
---
 tensor2tensor/utils/learning_rate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index ba2b7920c..e5e05c90e 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -44,7 +44,7 @@ def learning_rate_schedule(hparams):
   schedule_string = hparams.learning_rate_schedule
   names = schedule_string.split("*")
   names = [name.strip() for name in names if name.strip()]
-  ret = 1.0
+  ret = tf.constant(1.0)
   for name in names:
     ret *= learning_rate_factor(name, step_num, hparams)
   return ret

From 9234f492a60dbbcd5076fed162769c968f5061bb Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 7 Jun 2018 14:27:04 -0700
Subject: [PATCH 0056/2720] internal merge of PR #851

PiperOrigin-RevId: 199695936
---
 .../data_generators/translate_enid.py         | 25 +++++++++++--------
 tensor2tensor/utils/learning_rate.py          |  2 +-
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/data_generators/translate_enid.py b/tensor2tensor/data_generators/translate_enid.py
index a1a94fb49..df49d31ec 100644
--- a/tensor2tensor/data_generators/translate_enid.py
+++ b/tensor2tensor/data_generators/translate_enid.py
@@ -28,31 +28,33 @@
 # End-of-sentence marker.
 EOS = text_encoder.EOS_ID
 
+_REPO = "https://github.com/prasastoadi/parallel-corpora-en-id/raw/master/"
+
 # IWSLT17 :
-# 	109335 sentences
-#	https://wit3.fbk.eu/mt.php?release=2017-01-more
+# 109335 sentences
+# https://wit3.fbk.eu/mt.php?release=2017-01-more
 # PANL-BPPT :
-#	24024 sentences
-# 	http://www.panl10n.net/english/outputs/Indonesia/BPPT/0902/BPPTIndToEngCorpusHalfM.zip
+# 24024 sentences
+# http://www.panl10n.net/english/outputs/Indonesia/BPPT/0902/BPPTIndToEngCorpusHalfM.zip # pylint: disable=line-too-long
 _ENID_TRAIN_DATASETS = [
     [
-        "https://github.com/prasastoadi/parallel-corpora-en-id/raw/master/IWSLT17.train.en-id.tok.tgz",
+        _REPO + "IWSLT17.train.en-id.tok.tgz",
         ("IWSLT17.train.en-id.tok.en", "IWSLT17.train.en-id.tok.id")
     ],
     [
-        "https://github.com/prasastoadi/parallel-corpora-en-id/raw/master/PANL-BPPT-ECO-EN-ID-150Kw.tok.tgz",
+        _REPO + "PANL-BPPT-ECO-EN-ID-150Kw.tok.tgz",
         ("PANL-BPPT-ECO-EN-150Kw.tok.txt", "PANL-BPPT-ECO-ID-150Kw.tok.txt")
     ],
     [
-        "https://github.com/prasastoadi/parallel-corpora-en-id/raw/master/PANL-BPPT-INT-EN-ID-150Kw.tok.tgz",
+        _REPO + "PANL-BPPT-INT-EN-ID-150Kw.tok.tgz",
         ("PANL-BPPT-INT-EN-150Kw.tok.txt", "PANL-BPPT-INT-ID-150Kw.tok.txt")
     ],
     [
-        "https://github.com/prasastoadi/parallel-corpora-en-id/raw/master/PANL-BPPT-SCI-EN-ID-100Kw.tok.tgz",
+        _REPO + "PANL-BPPT-SCI-EN-ID-100Kw.tok.tgz",
         ("PANL-BPPT-SCI-EN-100Kw.tok.txt", "PANL-BPPT-SCI-ID-100Kw.tok.txt")
     ],
     [
-        "https://github.com/prasastoadi/parallel-corpora-en-id/raw/master/PANL-BPPT-SPO-EN-ID-100Kw.tok.tgz",
+        _REPO + "PANL-BPPT-SPO-EN-ID-100Kw.tok.tgz",
         ("PANL-BPPT-SPO-EN-100Kw.tok.txt", "PANL-BPPT-SPO-ID-100Kw.tok.txt")
     ],
 ]
@@ -62,8 +64,9 @@
 # https://wit3.fbk.eu/mt.php?release=2017-01-more
 _ENID_TEST_DATASETS = [
     [
-        "https://github.com/prasastoadi/parallel-corpora-en-id/raw/master/IWSLT17.TED.tst2017plus.en-id.tok.tgz",
-        ("IWSLT17.TED.tst2017plus.en-id.tok.en", "IWSLT17.TED.tst2017plus.en-id.tok.id")
+        _REPO + "IWSLT17.TED.tst2017plus.en-id.tok.tgz",
+        ("IWSLT17.TED.tst2017plus.en-id.tok.en",
+         "IWSLT17.TED.tst2017plus.en-id.tok.id")
     ]
 ]
 
diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index e5e05c90e..ba2b7920c 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -44,7 +44,7 @@ def learning_rate_schedule(hparams):
   schedule_string = hparams.learning_rate_schedule
   names = schedule_string.split("*")
   names = [name.strip() for name in names if name.strip()]
-  ret = tf.constant(1.0)
+  ret = 1.0
   for name in names:
     ret *= learning_rate_factor(name, step_num, hparams)
   return ret

From e4ddc53a6c4c3f4f60b089d0f62f0891385aab54 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 7 Jun 2018 14:29:52 -0700
Subject: [PATCH 0057/2720] internal merge of PR #842

PiperOrigin-RevId: 199696404
---
 tensor2tensor/utils/learning_rate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index ba2b7920c..e5e05c90e 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -44,7 +44,7 @@ def learning_rate_schedule(hparams):
   schedule_string = hparams.learning_rate_schedule
   names = schedule_string.split("*")
   names = [name.strip() for name in names if name.strip()]
-  ret = 1.0
+  ret = tf.constant(1.0)
   for name in names:
     ret *= learning_rate_factor(name, step_num, hparams)
   return ret

From f2393154af25fd9d5372a40ae67c0557787b9448 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 7 Jun 2018 15:56:57 -0700
Subject: [PATCH 0058/2720] Fix some pylint errors.

PiperOrigin-RevId: 199711177
---
 tensor2tensor/models/research/next_frame.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 1c7b54ae3..1e5524d4f 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -619,6 +619,10 @@ def init_state(self,
     Returns:
        A tensors representing the initial state.
     """
+
+    del state_initializer
+    del dtype
+
     # recoded by @mbz
     initial_state = tf.zeros([tf.shape(inputs)[0]] + state_shape)
     return initial_state

From 42f8b2677dc4ff1bd60accda4a0a0946cbae23af Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 7 Jun 2018 17:58:51 -0700
Subject: [PATCH 0059/2720] s/os.mkdir/tf.gfile.MakeDirs -- this seems to make
 Travis happy.

PiperOrigin-RevId: 199727255
---
 tensor2tensor/data_generators/text_encoder_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/text_encoder_test.py b/tensor2tensor/data_generators/text_encoder_test.py
index 74c401cf3..f7c7ed4e2 100644
--- a/tensor2tensor/data_generators/text_encoder_test.py
+++ b/tensor2tensor/data_generators/text_encoder_test.py
@@ -69,7 +69,7 @@ def setUpClass(cls):
     """Make sure the test dir exists and is empty."""
     cls.test_temp_dir = os.path.join(tf.test.get_temp_dir(), "encoder_test")
     shutil.rmtree(cls.test_temp_dir, ignore_errors=True)
-    os.mkdir(cls.test_temp_dir)
+    tf.gfile.MakeDirs(cls.test_temp_dir)
 
   def test_save_and_reload(self):
     """Test that saving and reloading doesn't change the vocab.
@@ -113,7 +113,7 @@ def setUpClass(cls):
     """Make sure the test dir exists and is empty."""
     cls.test_temp_dir = os.path.join(tf.test.get_temp_dir(), "encoder_test")
     shutil.rmtree(cls.test_temp_dir, ignore_errors=True)
-    os.mkdir(cls.test_temp_dir)
+    tf.gfile.MakeDirs(cls.test_temp_dir)
 
   def test_encode_decode(self):
     corpus = (

From bed06aade213b01a565fefe50c17cb545ffc1686 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 7 Jun 2018 19:50:56 -0700
Subject: [PATCH 0060/2720] unicode in program_search wasn't needed, removed
 it. This fixes the py3 Travis issue.

PiperOrigin-RevId: 199736437
---
 tensor2tensor/data_generators/program_search.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/program_search.py b/tensor2tensor/data_generators/program_search.py
index 8ba661b86..d1846dadc 100644
--- a/tensor2tensor/data_generators/program_search.py
+++ b/tensor2tensor/data_generators/program_search.py
@@ -63,7 +63,7 @@ def _extract_filename_from_url(url):
   @staticmethod
   def _flatten_target_programs(iterable):
     # The target programs are read as nested lists, we should flatten them.
-    yield unicode("[")
+    yield "["
     it = iter(iterable)
     for e in it:
       if isinstance(e, (list, tuple)):
@@ -71,7 +71,7 @@ def _flatten_target_programs(iterable):
           yield f
       else:
         yield e
-    yield unicode("]")
+    yield "]"
 
   @staticmethod
   def _parse_json_to_dict(json_line):

From a7382de8d17dc0e7cef9fdaf4e5811d71c608488 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 7 Jun 2018 20:02:08 -0700
Subject: [PATCH 0061/2720] Make stochastic video model run on pong.

PiperOrigin-RevId: 199737031
---
 tensor2tensor/data_generators/gym_problems.py | 14 ++++
 tensor2tensor/layers/common_layers.py         | 26 +++++++
 tensor2tensor/models/basic.py                 | 12 +--
 tensor2tensor/models/research/next_frame.py   | 74 +++++--------------
 4 files changed, 64 insertions(+), 62 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 61214bb14..722acaae5 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -170,6 +170,20 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
 class GymPongRandom5k(GymDiscreteProblem):
   """Pong game, random actions."""
 
+  # Hard-coding num_actions, frame_height, frame_width to avoid loading
+  # libale.so file.
+  @property
+  def num_actions(self):
+    return 6
+
+  @property
+  def frame_height(self):
+    return 210
+
+  @property
+  def frame_width(self):
+    return 160
+
   @property
   def env_name(self):
     return "PongDeterministic-v4"
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 103943b6a..afbb0f3d9 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3111,3 +3111,29 @@ def cast_like(x, y):
     tf.logging.warning("Cast for %s may induce copy from '%s' to '%s'",
                        x.name, x.device, cast_x.device)
   return cast_x
+
+
+def make_even_size(x):
+  """Pad x to be even-sized on axis 1 and 2, but only if necessary."""
+  x_shape = x.get_shape().as_list()
+  assert len(x_shape) > 2, "Only 3+-dimensional tensors supported."
+  shape = [dim if dim is not None else -1 for dim in x_shape]
+  new_shape = x_shape  # To make sure constant shapes remain constant.
+  if x_shape[1] is not None:
+    new_shape[1] = 2 * int(math.ceil(x_shape[1] * 0.5))
+  if x_shape[2] is not None:
+    new_shape[2] = 2 * int(math.ceil(x_shape[2] * 0.5))
+  if shape[1] % 2 == 0 and shape[2] % 2 == 0:
+    return x
+  if shape[1] % 2 == 0:
+    x, _ = pad_to_same_length(x, x, final_length_divisible_by=2, axis=2)
+    x.set_shape(new_shape)
+    return x
+  if shape[2] % 2 == 0:
+    x, _ = pad_to_same_length(x, x, final_length_divisible_by=2, axis=1)
+    x.set_shape(new_shape)
+    return x
+  x, _ = pad_to_same_length(x, x, final_length_divisible_by=2, axis=1)
+  x, _ = pad_to_same_length(x, x, final_length_divisible_by=2, axis=2)
+  x.set_shape(new_shape)
+  return x
diff --git a/tensor2tensor/models/basic.py b/tensor2tensor/models/basic.py
index d0aded7bc..4d9c7b763 100644
--- a/tensor2tensor/models/basic.py
+++ b/tensor2tensor/models/basic.py
@@ -65,17 +65,13 @@ def unbottleneck(self, x, res_size):
       return x
 
   def make_even_size(self, x):
-    shape = [dim if dim is not None else -1 for dim in x.get_shape().as_list()]
-    if shape[1] % 2 == 0 and shape[2] % 2 == 0:
-      return x
-    if shape[1] % 2 == 0 and self.is1d:
+    if not self.is1d:
+      return common_layers.make_even_size(x)
+    shape1 = x.get_shape().as_list()[1]
+    if shape1 is not None and shape1 % 2 == 0:
       return x
     x, _ = common_layers.pad_to_same_length(
         x, x, final_length_divisible_by=2, axis=1)
-    if self.is1d:
-      return x
-    x, _ = common_layers.pad_to_same_length(
-        x, x, final_length_divisible_by=2, axis=2)
     return x
 
   def encoder(self, x):
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 1e5524d4f..fbec036c5 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -17,6 +17,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
 import numpy as np
 import six
 
@@ -34,25 +35,6 @@
 class NextFrameBasic(t2t_model.T2TModel):
   """Basic next-frame model, may take actions and predict rewards too."""
 
-  def make_even_size(self, x):
-    """Pad x to be even-sized on axis 1 and 2, but only if necessary."""
-    shape = [dim if dim is not None else -1 for dim in x.get_shape().as_list()]
-    if shape[1] % 2 == 0 and shape[2] % 2 == 0:
-      return x
-    if shape[1] % 2 == 0:
-      x, _ = common_layers.pad_to_same_length(
-          x, x, final_length_divisible_by=2, axis=2)
-      return x
-    if shape[2] % 2 == 0:
-      x, _ = common_layers.pad_to_same_length(
-          x, x, final_length_divisible_by=2, axis=1)
-      return x
-    x, _ = common_layers.pad_to_same_length(
-        x, x, final_length_divisible_by=2, axis=1)
-    x, _ = common_layers.pad_to_same_length(
-        x, x, final_length_divisible_by=2, axis=2)
-    return x
-
   def body(self, features):
     hparams = self.hparams
     filters = hparams.hidden_size
@@ -71,7 +53,7 @@ def body(self, features):
     for i in range(hparams.num_compress_steps):
       with tf.variable_scope("downstride%d" % i):
         layer_inputs.append(x)
-        x = self.make_even_size(x)
+        x = common_layers.make_even_size(x)
         if i < hparams.filter_double_steps:
           filters *= 2
         x = tf.layers.conv2d(x, filters, kernel2, activation=common_layers.belu,
@@ -196,13 +178,14 @@ def construct_latent_tower(self, images):
 
     with slim.arg_scope([slim.conv2d], reuse=False):
       stacked_images = tf.concat(images, 3)
-
+      stacked_images = common_layers.make_even_size(stacked_images)
       latent_enc1 = slim.conv2d(
           stacked_images,
           32, [3, 3],
           stride=2,
           scope="latent_conv1")
       latent_enc1 = slim.batch_norm(latent_enc1, scope="latent_bn1")
+      latent_enc1 = common_layers.make_even_size(latent_enc1)
 
       latent_enc2 = slim.conv2d(
           latent_enc1,
@@ -218,6 +201,7 @@ def construct_latent_tower(self, images):
           scope="latent_conv3")
       latent_enc3 = slim.batch_norm(latent_enc3, scope="latent_bn3")
 
+      latent_enc3 = common_layers.make_even_size(latent_enc3)
       latent_mean = slim.conv2d(
           latent_enc3,
           self.hparams.latent_channels, [3, 3],
@@ -348,6 +332,7 @@ def construct_model(self,
         # Predicted state is always fed back in
         state_action = tf.concat(axis=1, values=[action, current_state])
 
+        prev_image = common_layers.make_even_size(prev_image)
         enc0 = slim.layers.conv2d(
             prev_image,
             32, [5, 5],
@@ -362,6 +347,7 @@ def construct_model(self,
         hidden2, lstm_state2 = lstm_func(
             hidden1, lstm_state2, lstm_size[1], scope="state2")
         hidden2 = layer_norm(hidden2, scope="layer_norm3")
+        hidden2 = common_layers.make_even_size(hidden2)
         enc1 = slim.layers.conv2d(
             hidden2, hidden2.get_shape()[3], [3, 3], stride=2, scope="conv2")
 
@@ -371,16 +357,17 @@ def construct_model(self,
         hidden4, lstm_state4 = lstm_func(
             hidden3, lstm_state4, lstm_size[3], scope="state4")
         hidden4 = layer_norm(hidden4, scope="layer_norm5")
+        hidden4 = common_layers.make_even_size(hidden4)
         enc2 = slim.layers.conv2d(
             hidden4, hidden4.get_shape()[3], [3, 3], stride=2, scope="conv3")
 
         # Pass in state and action.
         smear = tf.reshape(
             state_action,
-            [-1, 1, 1, int(state_action.get_shape()[1])])
+            [-1, 1, 1, int(common_layers.shape_list(state_action)[1])])
+        enc2_shape = common_layers.shape_list(enc2)
         smear = tf.tile(
-            smear, [1, int(enc2.get_shape()[1]),
-                    int(enc2.get_shape()[2]), 1])
+            smear, [1, enc2_shape[1], enc2_shape[2], 1])
         if use_state:
           enc2 = tf.concat(axis=3, values=[enc2, smear])
 
@@ -404,6 +391,8 @@ def construct_model(self,
         enc4 = slim.layers.conv2d_transpose(
             hidden5, hidden5.get_shape()[3], 3, stride=2, scope="convt1")
 
+        enc1_shape = common_layers.shape_list(enc1)
+        enc4 = enc4[:, :enc1_shape[1], :enc1_shape[2], :]  # Cut to shape.
         hidden6, lstm_state6 = lstm_func(
             enc4, lstm_state6, lstm_size[5], scope="state6")  # 16x16
         hidden6 = layer_norm(hidden6, scope="layer_norm7")
@@ -412,6 +401,8 @@ def construct_model(self,
 
         enc5 = slim.layers.conv2d_transpose(
             hidden6, hidden6.get_shape()[3], 3, stride=2, scope="convt2")
+        enc0_shape = common_layers.shape_list(enc0)
+        enc5 = enc5[:, :enc0_shape[1], :enc0_shape[2], :]  # Cut to shape.
         hidden7, lstm_state7 = lstm_func(
             enc5, lstm_state7, lstm_size[6], scope="state7")  # 32x32
         hidden7 = layer_norm(hidden7, scope="layer_norm8")
@@ -604,29 +595,6 @@ def scheduled_sample(self,
     return tf.dynamic_stitch([ground_truth_idx, generated_idx],
                              [ground_truth_examps, generated_examps])
 
-  def init_state(self,
-                 inputs,
-                 state_shape,
-                 state_initializer=tf.zeros_initializer(),
-                 dtype=tf.float32):
-    """Helper function to create an initial state given inputs.
-
-    Args:
-      inputs: input Tensor, at least 2D, the first dimension being batch_size
-      state_shape: the shape of the state.
-      state_initializer: Initializer(shape, dtype) for state Tensor.
-      dtype: Optional dtype, needed when inputs is None.
-    Returns:
-       A tensors representing the initial state.
-    """
-
-    del state_initializer
-    del dtype
-
-    # recoded by @mbz
-    initial_state = tf.zeros([tf.shape(inputs)[0]] + state_shape)
-    return initial_state
-
   # TODO(mbz): use tf.distributions.kl_divergence instead.
   def kl_divergence(self, mu, log_sigma):
     """KL divergence of diagonal gaussian N(mu,exp(log_sigma)) and N(0,1).
@@ -668,10 +636,9 @@ def basic_conv_lstm_cell(self,
     Returns:
        a tuple of tensors representing output and the new state.
     """
-    spatial_size = [v.value for v in inputs.get_shape()[1:3]]
-
     if state is None:
-      state = self.init_state(inputs, spatial_size + [2 * num_channels])
+      inputs_shape = common_layers.shape_list(inputs)
+      state = tf.zeros(inputs_shape[:3] + [2 * num_channels])
     with tf.variable_scope(scope,
                            "BasicConvLstmCell",
                            [inputs, state],
@@ -767,13 +734,12 @@ def next_frame():
 @registry.register_hparams
 def next_frame_stochastic():
   """SV2P model."""
-  hparams = common_hparams.basic_params1()
+  hparams = next_frame()
   hparams.video_num_input_frames = 1
   hparams.video_num_target_frames = 4
   hparams.batch_size = 8
-  hparams.learning_rate_constant = 1e-3
-  hparams.learning_rate_schedule = "constant"
-  hparams.weight_decay = 0.0
+  hparams.target_modality = "video:raw"
+  hparams.input_modalities = "inputs:video:raw"
   hparams.add_hparam("stochastic_model", True)
   hparams.add_hparam("latent_channels", 1)
   hparams.add_hparam("latent_std_min", -5.0)

From 6713dd67046569344bbbd8f9f8a405ef3803d865 Mon Sep 17 00:00:00 2001
From: blazej0 <blazej0@users.noreply.github.com>
Date: Fri, 8 Jun 2018 05:16:50 +0200
Subject: [PATCH 0062/2720] Fixing initial frames and related issues. (#860)

---
 tensor2tensor/rl/collect.py                  |  6 ++--
 tensor2tensor/rl/envs/in_graph_batch_env.py  |  6 +++-
 tensor2tensor/rl/envs/py_func_batch_env.py   | 13 -------
 tensor2tensor/rl/envs/simulated_batch_env.py | 36 ++++----------------
 tensor2tensor/rl/envs/tf_atari_wrappers.py   | 28 +++++++++++----
 5 files changed, 36 insertions(+), 53 deletions(-)

diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 3335697ca..81bd03852 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -48,9 +48,11 @@ def define_collect(policy_factory, batch_env, hparams,
 
   should_reset_var = tf.Variable(True, trainable=False)
 
+  zeros_tensor = tf.zeros(len(batch_env))
+
   def group():
     return tf.group(batch_env.reset(tf.range(len(batch_env))),
-                    tf.assign(cumulative_rewards, tf.zeros(len(batch_env))))
+                    tf.assign(cumulative_rewards, zeros_tensor))
   reset_op = tf.cond(
       tf.logical_or(should_reset_var, tf.logical_or(eval_phase, on_simulated)),
       group, tf.no_op)
@@ -98,7 +100,7 @@ def step(index, scores_sum, scores_num):
         reset_env_op = batch_env.reset(agent_indices_to_reset)
         reset_cumulative_rewards_op = tf.scatter_update(
             cumulative_rewards, agent_indices_to_reset,
-            tf.zeros(tf.shape(agent_indices_to_reset)))
+            tf.gather(zeros_tensor, agent_indices_to_reset))
       with tf.control_dependencies([reset_env_op,
                                     reset_cumulative_rewards_op]):
         return [index + 1, scores_sum + scores_sum_delta,
diff --git a/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index ee0ae94d1..83ce0ca78 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -21,6 +21,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import tensorflow as tf
+
 
 class InGraphBatchEnv(object):
   """Abstract class for batch of environments inside the TensorFlow graph.
@@ -67,7 +69,9 @@ def reset(self, indices=None):
     Returns:
       Batch tensor of the new observations.
     """
-    raise NotImplementedError
+    return tf.cond(
+        tf.cast(tf.reduce_sum(indices+1), tf.bool),
+        lambda: self._reset_non_empty(indices), lambda: 0.0)
 
   @property
   def observ(self):
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index 92dca08ab..692a37298 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -95,19 +95,6 @@ def simulate(self, action):
       with tf.control_dependencies([self._observ.assign(observ)]):
         return tf.identity(reward), tf.identity(done)
 
-  def reset(self, indices=None):
-    """Reset the batch of environments.
-
-    Args:
-      indices: The batch indices of the environments to reset.
-
-    Returns:
-      Batch tensor of the new observations.
-    """
-    return tf.cond(
-        tf.cast(tf.shape(indices)[0], tf.bool),
-        lambda: self._reset_non_empty(indices), lambda: 0.0)
-
   def _reset_non_empty(self, indices):
     """Reset the batch of environments.
 
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 2ed6f2b65..c5287aced 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -39,24 +39,16 @@ def __init__(self, input_dataset, length):
     self.input_data_iterator = (
         input_dataset.batch(length).make_one_shot_iterator())
     self.length = length
-    initial_frames = self.get_initial_observations(length)
+    initial_frames = self.get_initial_observations()
     initial_shape = [length] + common_layers.shape_list(initial_frames)[1:]
     self._history_buff = tf.Variable(tf.zeros(initial_shape, tf.float32),
                                      trainable=False)
-    self._assigned = False
 
-  def get_initial_observations(self, n):
-    initial_frames = self.input_data_iterator.get_next()
-    return tf.cast(initial_frames[:n, ...], tf.float32)
+  def get_initial_observations(self):
+    return tf.cast(self.input_data_iterator.get_next(), tf.float32)
 
   def get_all_elements(self):
-    if self._assigned:
-      return self._history_buff.read_value()
-    assign = self._history_buff.assign(
-        self.get_initial_observations(self.length))
-    with tf.control_dependencies([assign]):
-      self._assigned = True
-      return tf.identity(self.initial_frames)
+    return self._history_buff.read_value()
 
   def move_by_one_element(self, element):
     last_removed = self.get_all_elements()[:, 1:, ...]
@@ -64,15 +56,12 @@ def move_by_one_element(self, element):
     moved = tf.concat([last_removed, element], axis=1)
     with tf.control_dependencies([moved]):
       with tf.control_dependencies([self._history_buff.assign(moved)]):
-        self._assigned = True
         return self._history_buff.read_value()
 
   def reset(self, indices):
-    number_of_indices = tf.size(indices)
-    initial_frames = self.get_initial_observations(number_of_indices)
+    initial_frames = tf.gather(self.get_initial_observations(), indices)
     scatter_op = tf.scatter_update(self._history_buff, indices, initial_frames)
     with tf.control_dependencies([scatter_op]):
-      self._assigned = True
       return self._history_buff.read_value()
 
 
@@ -179,19 +168,6 @@ def simulate(self, action):
              self.history_buffer.move_by_one_element(observ)]):
           return tf.identity(reward), tf.identity(done)
 
-  def reset(self, indices=None):
-    """Reset the batch of environments.
-
-    Args:
-      indices: The batch indices of the environments to reset.
-
-    Returns:
-      Batch tensor of the new observations.
-    """
-    return tf.cond(
-        tf.cast(tf.shape(indices)[0], tf.bool),
-        lambda: self._reset_non_empty(indices), lambda: 0.0)
-
   def _reset_non_empty(self, indices):
     """Reset the batch of environments.
 
@@ -209,4 +185,4 @@ def _reset_non_empty(self, indices):
   @property
   def observ(self):
     """Access the variable holding the current observation."""
-    return self._observ
+    return tf.identity(self._observ)
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 55e4bb807..719c8fd25 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -34,14 +34,17 @@ def __init__(self, batch_env):
   @property
   def observ(self):
     """Access the variable holding the current observation."""
-    return self._observ
+    return self._observ.read_value()
 
   def __len__(self):
     """Number of combined environments."""
     return self._length
 
-  def reset(self, indices=None):
-    return self._batch_env.reset(indices)
+  def _reset_non_empty(self, indices):
+    new_values = self._batch_env._reset_non_empty(indices)
+    assign_op = tf.scatter_update(self._observ, indices, new_values)
+    with tf.control_dependencies([assign_op]):
+      return tf.identity(new_values)
 
 
 class TransformWrapper(WrapperBase):
@@ -168,6 +171,14 @@ def not_done_step(a, _):
       with tf.control_dependencies([self._observ.assign(observation)]):
         return tf.identity(rewards[-1, ...]), tf.identity(dones[-1, ...])
 
+  def _reset_non_empty(self, indices):
+    new_values = tf.gather(self._batch_env._reset_non_empty(indices), indices)
+    inx = tf.concat([tf.ones(tf.size(tf.shape(new_values)),
+                             dtype=tf.int32)[:-1], [self.skip]], axis=0)
+    assign_op = tf.scatter_update(self._observ, indices,
+                                  tf.tile(new_values, inx))
+    with tf.control_dependencies([assign_op]):
+      return tf.identity(self.observ)
 
 class TimeLimitWrapper(WrapperBase):
   """Time limit wrapper."""
@@ -189,11 +200,14 @@ def simulate(self, action):
         with tf.control_dependencies([inc]):
           return tf.identity(reward), tf.identity(new_done)
 
-  def reset(self, indices=None):
+  def _reset_non_empty(self, indices):
     op_zero = tf.scatter_update(self._time_elapsed, indices,
-                                tf.zeros(tf.shape(indices), dtype=tf.int32))
-    with tf.control_dependencies([op_zero]):
-      return self._batch_env.reset(indices)
+                                tf.gather(tf.zeros((len(self),), tf.int32),
+                                          indices))
+    new_values = tf.gather(self._batch_env._reset_non_empty(indices), indices)
+    assign_op = tf.scatter_update(self._observ, indices, new_values)
+    with tf.control_dependencies([op_zero, assign_op]):
+      return tf.identity(self.observ)
 
 
 class MemoryWrapper(WrapperBase):

From b0258f6af6e28165403c6cacddb3aff7828e4cc7 Mon Sep 17 00:00:00 2001
From: Seppo Enarvi <seppo.git@marjaniemi.com>
Date: Fri, 8 Jun 2018 05:24:14 +0200
Subject: [PATCH 0063/2720] Improvements to sequence length computation in LSTM
 models (#859)

* Compute sequence lengths to avoid effects of padding.

* LSTM models compute sequence lengths for the dynamic_rnn functions. Assumes
  sequences end at the first 0.
* Unidirectional LSTM encoders reverse the input sequences using
  tf.reverse_sequence(), in order to take padding into account.
* Bidirectional LSTM encoders don't reverse the input sequences, since it makes
  no sense.
* Use LSTMCell instead of BasicLSTMCell. Oddly enough, this seems to give a
  consistent improvement.

* Fixed variable name.

* lstm_seq2seq_internal_bid_encoder() uses sequence lengths.

* Added a unit test for find_sequence_lengths().

* Convert lengths to integers.
---
 tensor2tensor/layers/common_layers.py |   2 +-
 tensor2tensor/models/lstm.py          | 155 +++++++++++++++++---------
 2 files changed, 105 insertions(+), 52 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index afbb0f3d9..fab035723 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1300,7 +1300,7 @@ def length_from_embedding(emb):
   Returns:
     a Tensor with shape [batch].
   """
-  return tf.reduce_sum(mask_from_embedding(emb), [1, 2, 3])
+  return tf.cast(tf.reduce_sum(mask_from_embedding(emb), [1, 2, 3]), tf.int32)
 
 
 def mask_leq(target_length, source_length):
diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 0b933b6ce..e985795e5 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -27,37 +27,69 @@
 import tensorflow as tf
 
 
-def lstm(inputs, hparams, train, name, initial_state=None,
-         sequence_length=None):
-  """Run LSTM cell on inputs, assuming they are [batch x time x size]."""
-
-  def dropout_lstm_cell():
-    return tf.contrib.rnn.DropoutWrapper(
-        tf.contrib.rnn.BasicLSTMCell(hparams.hidden_size),
-        input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
-
-  layers = [dropout_lstm_cell() for _ in range(hparams.num_hidden_layers)]
+def _dropout_lstm_cell(hparams, train):
+  return tf.contrib.rnn.DropoutWrapper(
+      tf.contrib.rnn.LSTMCell(hparams.hidden_size),
+      input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
+
+
+def lstm(inputs, sequence_length, hparams, train, name, initial_state=None):
+  """Adds a stack of LSTM layers on top of input.
+
+  Args:
+    inputs: The input `Tensor`, shaped `[batch_size, time_steps, hidden_size]`.
+    sequence_length: Lengths of the actual input sequence, excluding padding; a
+        `Tensor` shaped `[batch_size]`.
+    hparams: tf.contrib.training.HParams; hyperparameters.
+    train: bool; `True` when constructing training graph to enable dropout.
+    name: string; Create variable names under this scope.
+    initial_state: tuple of `LSTMStateTuple`s; the initial state of each layer.
+
+  Returns:
+    A tuple (outputs, states), where:
+      outputs: The output `Tensor`, shaped `[batch_size, time_steps,
+        hidden_size]`.
+      states: A tuple of `LSTMStateTuple`s; the final state of each layer.
+        Bidirectional LSTM returns a concatenation of last forward and backward
+        state, reduced to the original dimensionality.
+  """
+  layers = [_dropout_lstm_cell(hparams, train)
+            for _ in range(hparams.num_hidden_layers)]
   with tf.variable_scope(name):
     return tf.nn.dynamic_rnn(
         tf.contrib.rnn.MultiRNNCell(layers),
         inputs,
+        sequence_length,
         initial_state=initial_state,
-        sequence_length=sequence_length,
         dtype=tf.float32,
         time_major=False)
 
 
 def lstm_attention_decoder(inputs, hparams, train, name, initial_state,
-                           encoder_outputs, encoder_output_length=None,
-                           decoder_input_length=None):
-  """Run LSTM cell with attention on inputs of shape [batch x time x size]."""
-
-  def dropout_lstm_cell():
-    return tf.contrib.rnn.DropoutWrapper(
-        tf.nn.rnn_cell.BasicLSTMCell(hparams.hidden_size),
-        input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
-
-  layers = [dropout_lstm_cell() for _ in range(hparams.num_hidden_layers)]
+                           encoder_outputs, encoder_output_length,
+                           decoder_input_length):
+  """Run LSTM cell with attention on inputs of shape [batch x time x size].
+
+  Args:
+    inputs: The decoder input `Tensor`, shaped `[batch_size, decoder_steps,
+        hidden_size]`.
+    hparams: tf.contrib.training.HParams; hyperparameters.
+    train: bool; `True` when constructing training graph to enable dropout.
+    name: string; Create variable names under this scope.
+    initial_state: Tuple of `LSTMStateTuple`s; the initial state of each layer.
+    encoder_outputs: Encoder outputs; a `Tensor` shaped `[batch_size,
+        encoder_steps, hidden_size]`.
+    encoder_output_length: Lengths of the actual encoder outputs, excluding
+        padding; a `Tensor` shaped `[batch_size]`.
+    decoder_input_length: Lengths of the actual decoder inputs, excluding
+        padding; a `Tensor` shaped `[batch_size]`.
+
+  Returns:
+    The decoder output `Tensor`, shaped `[batch_size, decoder_steps,
+    hidden_size]`.
+  """
+  layers = [_dropout_lstm_cell(hparams, train)
+            for _ in range(hparams.num_hidden_layers)]
   if hparams.attention_mechanism == "luong":
     attention_mechanism_class = tf.contrib.seq2seq.LuongAttention
   elif hparams.attention_mechanism == "bahdanau":
@@ -81,36 +113,47 @@ def dropout_lstm_cell():
       cell_state=initial_state)
 
   with tf.variable_scope(name):
-    output, state = tf.nn.dynamic_rnn(
+    output, _ = tf.nn.dynamic_rnn(
         cell,
         inputs,
+        decoder_input_length,
         initial_state=initial_state,
-        sequence_length=decoder_input_length,
         dtype=tf.float32,
         time_major=False)
-
-    # For multi-head attention project output back to hidden size
+    # output is [batch_size, decoder_steps, attention_size], where
+    # attention_size is either hparams.hidden_size (when
+    # hparams.output_attention is 0) or hparams.attention_layer_size (when
+    # hparams.output_attention is 1) times the number of attention heads.
+    #
+    # For multi-head attention project output back to hidden size.
     if hparams.output_attention == 1 and hparams.num_heads > 1:
       output = tf.layers.dense(output, hparams.hidden_size)
 
-    return output, state
+    return output
 
 
 def lstm_seq2seq_internal(inputs, targets, hparams, train):
   """The basic LSTM seq2seq model, main step used for training."""
   with tf.variable_scope("lstm_seq2seq"):
     if inputs is not None:
+      inputs_length = common_layers.length_from_embedding(inputs)
       # Flatten inputs.
       inputs = common_layers.flatten4d3d(inputs)
+
       # LSTM encoder.
-      _, final_encoder_state = lstm(
-          tf.reverse(inputs, axis=[1]), hparams, train, "encoder")
+      inputs = tf.reverse_sequence(inputs, inputs_length, seq_axis=1)
+      _, final_encoder_state = lstm(inputs, inputs_length, hparams, train,
+                                    "encoder")
     else:
       final_encoder_state = None
+
     # LSTM decoder.
     shifted_targets = common_layers.shift_right(targets)
+    # Add 1 to account for the padding added to the left from shift_right
+    targets_length = common_layers.length_from_embedding(shifted_targets) + 1
     decoder_outputs, _ = lstm(
         common_layers.flatten4d3d(shifted_targets),
+        targets_length,
         hparams,
         train,
         "decoder",
@@ -127,41 +170,40 @@ def lstm_seq2seq_internal_attention(inputs, targets, hparams, train):
     inputs_length = common_layers.length_from_embedding(inputs)
     # Flatten inputs.
     inputs = common_layers.flatten4d3d(inputs)
+
     # LSTM encoder.
+    inputs = tf.reverse_sequence(inputs, inputs_length, seq_axis=1)
     encoder_outputs, final_encoder_state = lstm(
-        inputs, hparams, train, "encoder", sequence_length=inputs_length)
-    # LSTM decoder with attention
+        inputs, inputs_length, hparams, train, "encoder")
+
+    # LSTM decoder with attention.
     shifted_targets = common_layers.shift_right(targets)
     # Add 1 to account for the padding added to the left from shift_right
     targets_length = common_layers.length_from_embedding(shifted_targets) + 1
-    decoder_outputs, _ = lstm_attention_decoder(
+    decoder_outputs = lstm_attention_decoder(
         common_layers.flatten4d3d(shifted_targets), hparams, train, "decoder",
-        final_encoder_state, encoder_outputs,
-        encoder_output_length=inputs_length,
-        decoder_input_length=targets_length)
+        final_encoder_state, encoder_outputs, inputs_length, targets_length)
     return tf.expand_dims(decoder_outputs, axis=2)
 
 
-def lstm_bid_encoder(inputs, hparams, train, name):
+def lstm_bid_encoder(inputs, sequence_length, hparams, train, name):
   """Bidirectional LSTM for encoding inputs that are [batch x time x size]."""
 
-  def dropout_lstm_cell():
-    return tf.contrib.rnn.DropoutWrapper(
-        tf.contrib.rnn.BasicLSTMCell(hparams.hidden_size),
-        input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
-
   with tf.variable_scope(name):
     cell_fw = tf.contrib.rnn.MultiRNNCell(
-        [dropout_lstm_cell() for _ in range(hparams.num_hidden_layers)])
+        [_dropout_lstm_cell(hparams, train)
+         for _ in range(hparams.num_hidden_layers)])
 
     cell_bw = tf.contrib.rnn.MultiRNNCell(
-        [dropout_lstm_cell() for _ in range(hparams.num_hidden_layers)])
+        [_dropout_lstm_cell(hparams, train)
+         for _ in range(hparams.num_hidden_layers)])
 
     ((encoder_fw_outputs, encoder_bw_outputs),
      (encoder_fw_state, encoder_bw_state)) = tf.nn.bidirectional_dynamic_rnn(
-         cell_fw=cell_fw,
-         cell_bw=cell_bw,
-         inputs=inputs,
+         cell_fw,
+         cell_bw,
+         inputs,
+         sequence_length,
          dtype=tf.float32,
          time_major=False)
 
@@ -196,19 +238,24 @@ def lstm_seq2seq_internal_bid_encoder(inputs, targets, hparams, train):
   """The basic LSTM seq2seq model with bidirectional encoder."""
   with tf.variable_scope("lstm_seq2seq_bid_encoder"):
     if inputs is not None:
+      inputs_length = common_layers.length_from_embedding(inputs)
       # Flatten inputs.
       inputs = common_layers.flatten4d3d(inputs)
       # LSTM encoder.
       _, final_encoder_state = lstm_bid_encoder(
-          tf.reverse(inputs, axis=[1]), hparams, train, "encoder")
+          inputs, inputs_length, hparams, train, "encoder")
     else:
+      inputs_length = None
       final_encoder_state = None
     # LSTM decoder.
     shifted_targets = common_layers.shift_right(targets)
+    # Add 1 to account for the padding added to the left from shift_right
+    targets_length = common_layers.length_from_embedding(shifted_targets) + 1
     hparams_decoder = copy.copy(hparams)
     hparams_decoder.hidden_size = 2 * hparams.hidden_size
     decoder_outputs, _ = lstm(
         common_layers.flatten4d3d(shifted_targets),
+        targets_length,
         hparams_decoder,
         train,
         "decoder",
@@ -220,18 +267,22 @@ def lstm_seq2seq_internal_attention_bid_encoder(inputs, targets, hparams,
                                                 train):
   """LSTM seq2seq model with attention, main step used for training."""
   with tf.variable_scope("lstm_seq2seq_attention_bid_encoder"):
+    inputs_length = common_layers.length_from_embedding(inputs)
     # Flatten inputs.
     inputs = common_layers.flatten4d3d(inputs)
     # LSTM encoder.
     encoder_outputs, final_encoder_state = lstm_bid_encoder(
-        tf.reverse(inputs, axis=[1]), hparams, train, "encoder")
+        inputs, inputs_length, hparams, train, "encoder")
     # LSTM decoder with attention
     shifted_targets = common_layers.shift_right(targets)
+    # Add 1 to account for the padding added to the left from shift_right
+    targets_length = common_layers.length_from_embedding(shifted_targets) + 1
     hparams_decoder = copy.copy(hparams)
     hparams_decoder.hidden_size = 2 * hparams.hidden_size
-    decoder_outputs, _ = lstm_attention_decoder(
+    decoder_outputs = lstm_attention_decoder(
         common_layers.flatten4d3d(shifted_targets), hparams_decoder, train,
-        "decoder", final_encoder_state, encoder_outputs)
+        "decoder", final_encoder_state, encoder_outputs,
+        inputs_length, targets_length)
     return tf.expand_dims(decoder_outputs, axis=2)
 
 
@@ -244,11 +295,13 @@ def body(self, features):
       raise ValueError("LSTM models fail with orthogonal initializer.")
     train = self._hparams.mode == tf.estimator.ModeKeys.TRAIN
     inputs = features.get("inputs")
+    inputs_length = common_layers.length_from_embedding(inputs)
     # Flatten inputs.
     inputs = common_layers.flatten4d3d(inputs)
     # LSTM encoder.
-    encoder_output, _ = lstm(
-        tf.reverse(inputs, axis=[1]), self._hparams, train, "encoder")
+    inputs = tf.reverse_sequence(inputs, inputs_length, seq_axis=1)
+    encoder_output, _ = lstm(inputs, inputs_length, self._hparams, train,
+                             "encoder")
     return tf.expand_dims(encoder_output, axis=2)
 
 

From b547eeab2d60de69f6b1317085d7481b7a10e1cd Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 7 Jun 2018 20:28:11 -0700
Subject: [PATCH 0064/2720] internal merge of PR #860

PiperOrigin-RevId: 199739028
---
 tensor2tensor/layers/common_layers.py       |   2 +-
 tensor2tensor/models/lstm.py                | 155 +++++++-------------
 tensor2tensor/rl/collect.py                 |   6 +-
 tensor2tensor/rl/envs/in_graph_batch_env.py |   2 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py  |  25 +++-
 5 files changed, 75 insertions(+), 115 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index fab035723..afbb0f3d9 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1300,7 +1300,7 @@ def length_from_embedding(emb):
   Returns:
     a Tensor with shape [batch].
   """
-  return tf.cast(tf.reduce_sum(mask_from_embedding(emb), [1, 2, 3]), tf.int32)
+  return tf.reduce_sum(mask_from_embedding(emb), [1, 2, 3])
 
 
 def mask_leq(target_length, source_length):
diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index e985795e5..0b933b6ce 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -27,69 +27,37 @@
 import tensorflow as tf
 
 
-def _dropout_lstm_cell(hparams, train):
-  return tf.contrib.rnn.DropoutWrapper(
-      tf.contrib.rnn.LSTMCell(hparams.hidden_size),
-      input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
-
-
-def lstm(inputs, sequence_length, hparams, train, name, initial_state=None):
-  """Adds a stack of LSTM layers on top of input.
-
-  Args:
-    inputs: The input `Tensor`, shaped `[batch_size, time_steps, hidden_size]`.
-    sequence_length: Lengths of the actual input sequence, excluding padding; a
-        `Tensor` shaped `[batch_size]`.
-    hparams: tf.contrib.training.HParams; hyperparameters.
-    train: bool; `True` when constructing training graph to enable dropout.
-    name: string; Create variable names under this scope.
-    initial_state: tuple of `LSTMStateTuple`s; the initial state of each layer.
-
-  Returns:
-    A tuple (outputs, states), where:
-      outputs: The output `Tensor`, shaped `[batch_size, time_steps,
-        hidden_size]`.
-      states: A tuple of `LSTMStateTuple`s; the final state of each layer.
-        Bidirectional LSTM returns a concatenation of last forward and backward
-        state, reduced to the original dimensionality.
-  """
-  layers = [_dropout_lstm_cell(hparams, train)
-            for _ in range(hparams.num_hidden_layers)]
+def lstm(inputs, hparams, train, name, initial_state=None,
+         sequence_length=None):
+  """Run LSTM cell on inputs, assuming they are [batch x time x size]."""
+
+  def dropout_lstm_cell():
+    return tf.contrib.rnn.DropoutWrapper(
+        tf.contrib.rnn.BasicLSTMCell(hparams.hidden_size),
+        input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
+
+  layers = [dropout_lstm_cell() for _ in range(hparams.num_hidden_layers)]
   with tf.variable_scope(name):
     return tf.nn.dynamic_rnn(
         tf.contrib.rnn.MultiRNNCell(layers),
         inputs,
-        sequence_length,
         initial_state=initial_state,
+        sequence_length=sequence_length,
         dtype=tf.float32,
         time_major=False)
 
 
 def lstm_attention_decoder(inputs, hparams, train, name, initial_state,
-                           encoder_outputs, encoder_output_length,
-                           decoder_input_length):
-  """Run LSTM cell with attention on inputs of shape [batch x time x size].
-
-  Args:
-    inputs: The decoder input `Tensor`, shaped `[batch_size, decoder_steps,
-        hidden_size]`.
-    hparams: tf.contrib.training.HParams; hyperparameters.
-    train: bool; `True` when constructing training graph to enable dropout.
-    name: string; Create variable names under this scope.
-    initial_state: Tuple of `LSTMStateTuple`s; the initial state of each layer.
-    encoder_outputs: Encoder outputs; a `Tensor` shaped `[batch_size,
-        encoder_steps, hidden_size]`.
-    encoder_output_length: Lengths of the actual encoder outputs, excluding
-        padding; a `Tensor` shaped `[batch_size]`.
-    decoder_input_length: Lengths of the actual decoder inputs, excluding
-        padding; a `Tensor` shaped `[batch_size]`.
-
-  Returns:
-    The decoder output `Tensor`, shaped `[batch_size, decoder_steps,
-    hidden_size]`.
-  """
-  layers = [_dropout_lstm_cell(hparams, train)
-            for _ in range(hparams.num_hidden_layers)]
+                           encoder_outputs, encoder_output_length=None,
+                           decoder_input_length=None):
+  """Run LSTM cell with attention on inputs of shape [batch x time x size]."""
+
+  def dropout_lstm_cell():
+    return tf.contrib.rnn.DropoutWrapper(
+        tf.nn.rnn_cell.BasicLSTMCell(hparams.hidden_size),
+        input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
+
+  layers = [dropout_lstm_cell() for _ in range(hparams.num_hidden_layers)]
   if hparams.attention_mechanism == "luong":
     attention_mechanism_class = tf.contrib.seq2seq.LuongAttention
   elif hparams.attention_mechanism == "bahdanau":
@@ -113,47 +81,36 @@ def lstm_attention_decoder(inputs, hparams, train, name, initial_state,
       cell_state=initial_state)
 
   with tf.variable_scope(name):
-    output, _ = tf.nn.dynamic_rnn(
+    output, state = tf.nn.dynamic_rnn(
         cell,
         inputs,
-        decoder_input_length,
         initial_state=initial_state,
+        sequence_length=decoder_input_length,
         dtype=tf.float32,
         time_major=False)
-    # output is [batch_size, decoder_steps, attention_size], where
-    # attention_size is either hparams.hidden_size (when
-    # hparams.output_attention is 0) or hparams.attention_layer_size (when
-    # hparams.output_attention is 1) times the number of attention heads.
-    #
-    # For multi-head attention project output back to hidden size.
+
+    # For multi-head attention project output back to hidden size
     if hparams.output_attention == 1 and hparams.num_heads > 1:
       output = tf.layers.dense(output, hparams.hidden_size)
 
-    return output
+    return output, state
 
 
 def lstm_seq2seq_internal(inputs, targets, hparams, train):
   """The basic LSTM seq2seq model, main step used for training."""
   with tf.variable_scope("lstm_seq2seq"):
     if inputs is not None:
-      inputs_length = common_layers.length_from_embedding(inputs)
       # Flatten inputs.
       inputs = common_layers.flatten4d3d(inputs)
-
       # LSTM encoder.
-      inputs = tf.reverse_sequence(inputs, inputs_length, seq_axis=1)
-      _, final_encoder_state = lstm(inputs, inputs_length, hparams, train,
-                                    "encoder")
+      _, final_encoder_state = lstm(
+          tf.reverse(inputs, axis=[1]), hparams, train, "encoder")
     else:
       final_encoder_state = None
-
     # LSTM decoder.
     shifted_targets = common_layers.shift_right(targets)
-    # Add 1 to account for the padding added to the left from shift_right
-    targets_length = common_layers.length_from_embedding(shifted_targets) + 1
     decoder_outputs, _ = lstm(
         common_layers.flatten4d3d(shifted_targets),
-        targets_length,
         hparams,
         train,
         "decoder",
@@ -170,40 +127,41 @@ def lstm_seq2seq_internal_attention(inputs, targets, hparams, train):
     inputs_length = common_layers.length_from_embedding(inputs)
     # Flatten inputs.
     inputs = common_layers.flatten4d3d(inputs)
-
     # LSTM encoder.
-    inputs = tf.reverse_sequence(inputs, inputs_length, seq_axis=1)
     encoder_outputs, final_encoder_state = lstm(
-        inputs, inputs_length, hparams, train, "encoder")
-
-    # LSTM decoder with attention.
+        inputs, hparams, train, "encoder", sequence_length=inputs_length)
+    # LSTM decoder with attention
     shifted_targets = common_layers.shift_right(targets)
     # Add 1 to account for the padding added to the left from shift_right
     targets_length = common_layers.length_from_embedding(shifted_targets) + 1
-    decoder_outputs = lstm_attention_decoder(
+    decoder_outputs, _ = lstm_attention_decoder(
         common_layers.flatten4d3d(shifted_targets), hparams, train, "decoder",
-        final_encoder_state, encoder_outputs, inputs_length, targets_length)
+        final_encoder_state, encoder_outputs,
+        encoder_output_length=inputs_length,
+        decoder_input_length=targets_length)
     return tf.expand_dims(decoder_outputs, axis=2)
 
 
-def lstm_bid_encoder(inputs, sequence_length, hparams, train, name):
+def lstm_bid_encoder(inputs, hparams, train, name):
   """Bidirectional LSTM for encoding inputs that are [batch x time x size]."""
 
+  def dropout_lstm_cell():
+    return tf.contrib.rnn.DropoutWrapper(
+        tf.contrib.rnn.BasicLSTMCell(hparams.hidden_size),
+        input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
+
   with tf.variable_scope(name):
     cell_fw = tf.contrib.rnn.MultiRNNCell(
-        [_dropout_lstm_cell(hparams, train)
-         for _ in range(hparams.num_hidden_layers)])
+        [dropout_lstm_cell() for _ in range(hparams.num_hidden_layers)])
 
     cell_bw = tf.contrib.rnn.MultiRNNCell(
-        [_dropout_lstm_cell(hparams, train)
-         for _ in range(hparams.num_hidden_layers)])
+        [dropout_lstm_cell() for _ in range(hparams.num_hidden_layers)])
 
     ((encoder_fw_outputs, encoder_bw_outputs),
      (encoder_fw_state, encoder_bw_state)) = tf.nn.bidirectional_dynamic_rnn(
-         cell_fw,
-         cell_bw,
-         inputs,
-         sequence_length,
+         cell_fw=cell_fw,
+         cell_bw=cell_bw,
+         inputs=inputs,
          dtype=tf.float32,
          time_major=False)
 
@@ -238,24 +196,19 @@ def lstm_seq2seq_internal_bid_encoder(inputs, targets, hparams, train):
   """The basic LSTM seq2seq model with bidirectional encoder."""
   with tf.variable_scope("lstm_seq2seq_bid_encoder"):
     if inputs is not None:
-      inputs_length = common_layers.length_from_embedding(inputs)
       # Flatten inputs.
       inputs = common_layers.flatten4d3d(inputs)
       # LSTM encoder.
       _, final_encoder_state = lstm_bid_encoder(
-          inputs, inputs_length, hparams, train, "encoder")
+          tf.reverse(inputs, axis=[1]), hparams, train, "encoder")
     else:
-      inputs_length = None
       final_encoder_state = None
     # LSTM decoder.
     shifted_targets = common_layers.shift_right(targets)
-    # Add 1 to account for the padding added to the left from shift_right
-    targets_length = common_layers.length_from_embedding(shifted_targets) + 1
     hparams_decoder = copy.copy(hparams)
     hparams_decoder.hidden_size = 2 * hparams.hidden_size
     decoder_outputs, _ = lstm(
         common_layers.flatten4d3d(shifted_targets),
-        targets_length,
         hparams_decoder,
         train,
         "decoder",
@@ -267,22 +220,18 @@ def lstm_seq2seq_internal_attention_bid_encoder(inputs, targets, hparams,
                                                 train):
   """LSTM seq2seq model with attention, main step used for training."""
   with tf.variable_scope("lstm_seq2seq_attention_bid_encoder"):
-    inputs_length = common_layers.length_from_embedding(inputs)
     # Flatten inputs.
     inputs = common_layers.flatten4d3d(inputs)
     # LSTM encoder.
     encoder_outputs, final_encoder_state = lstm_bid_encoder(
-        inputs, inputs_length, hparams, train, "encoder")
+        tf.reverse(inputs, axis=[1]), hparams, train, "encoder")
     # LSTM decoder with attention
     shifted_targets = common_layers.shift_right(targets)
-    # Add 1 to account for the padding added to the left from shift_right
-    targets_length = common_layers.length_from_embedding(shifted_targets) + 1
     hparams_decoder = copy.copy(hparams)
     hparams_decoder.hidden_size = 2 * hparams.hidden_size
-    decoder_outputs = lstm_attention_decoder(
+    decoder_outputs, _ = lstm_attention_decoder(
         common_layers.flatten4d3d(shifted_targets), hparams_decoder, train,
-        "decoder", final_encoder_state, encoder_outputs,
-        inputs_length, targets_length)
+        "decoder", final_encoder_state, encoder_outputs)
     return tf.expand_dims(decoder_outputs, axis=2)
 
 
@@ -295,13 +244,11 @@ def body(self, features):
       raise ValueError("LSTM models fail with orthogonal initializer.")
     train = self._hparams.mode == tf.estimator.ModeKeys.TRAIN
     inputs = features.get("inputs")
-    inputs_length = common_layers.length_from_embedding(inputs)
     # Flatten inputs.
     inputs = common_layers.flatten4d3d(inputs)
     # LSTM encoder.
-    inputs = tf.reverse_sequence(inputs, inputs_length, seq_axis=1)
-    encoder_output, _ = lstm(inputs, inputs_length, self._hparams, train,
-                             "encoder")
+    encoder_output, _ = lstm(
+        tf.reverse(inputs, axis=[1]), self._hparams, train, "encoder")
     return tf.expand_dims(encoder_output, axis=2)
 
 
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 81bd03852..a3e807efd 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -51,8 +51,10 @@ def define_collect(policy_factory, batch_env, hparams,
   zeros_tensor = tf.zeros(len(batch_env))
 
   def group():
-    return tf.group(batch_env.reset(tf.range(len(batch_env))),
-                    tf.assign(cumulative_rewards, zeros_tensor))
+    return tf.group(
+        batch_env.reset(tf.range(len(batch_env))),
+        tf.assign(cumulative_rewards, zeros_tensor))
+
   reset_op = tf.cond(
       tf.logical_or(should_reset_var, tf.logical_or(eval_phase, on_simulated)),
       group, tf.no_op)
diff --git a/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index 83ce0ca78..5bc76f2f7 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -70,7 +70,7 @@ def reset(self, indices=None):
       Batch tensor of the new observations.
     """
     return tf.cond(
-        tf.cast(tf.reduce_sum(indices+1), tf.bool),
+        tf.cast(tf.reduce_sum(indices + 1), tf.bool),
         lambda: self._reset_non_empty(indices), lambda: 0.0)
 
   @property
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 719c8fd25..32bfcbde6 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -41,7 +41,9 @@ def __len__(self):
     return self._length
 
   def _reset_non_empty(self, indices):
+    # pylint: disable=protected-access
     new_values = self._batch_env._reset_non_empty(indices)
+    # pylint: enable=protected-access
     assign_op = tf.scatter_update(self._observ, indices, new_values)
     with tf.control_dependencies([assign_op]):
       return tf.identity(new_values)
@@ -172,14 +174,21 @@ def not_done_step(a, _):
         return tf.identity(rewards[-1, ...]), tf.identity(dones[-1, ...])
 
   def _reset_non_empty(self, indices):
+    # pylint: disable=protected-access
     new_values = tf.gather(self._batch_env._reset_non_empty(indices), indices)
-    inx = tf.concat([tf.ones(tf.size(tf.shape(new_values)),
-                             dtype=tf.int32)[:-1], [self.skip]], axis=0)
-    assign_op = tf.scatter_update(self._observ, indices,
-                                  tf.tile(new_values, inx))
+    # pylint: enable=protected-access
+    inx = tf.concat(
+        [
+            tf.ones(tf.size(tf.shape(new_values)), dtype=tf.int32)[:-1],
+            [self.skip]
+        ],
+        axis=0)
+    assign_op = tf.scatter_update(self._observ, indices, tf.tile(
+        new_values, inx))
     with tf.control_dependencies([assign_op]):
       return tf.identity(self.observ)
 
+
 class TimeLimitWrapper(WrapperBase):
   """Time limit wrapper."""
 
@@ -201,10 +210,12 @@ def simulate(self, action):
           return tf.identity(reward), tf.identity(new_done)
 
   def _reset_non_empty(self, indices):
-    op_zero = tf.scatter_update(self._time_elapsed, indices,
-                                tf.gather(tf.zeros((len(self),), tf.int32),
-                                          indices))
+    op_zero = tf.scatter_update(
+        self._time_elapsed, indices,
+        tf.gather(tf.zeros((len(self),), tf.int32), indices))
+    # pylint: disable=protected-access
     new_values = tf.gather(self._batch_env._reset_non_empty(indices), indices)
+    # pylint: enable=protected-access
     assign_op = tf.scatter_update(self._observ, indices, new_values)
     with tf.control_dependencies([op_zero, assign_op]):
       return tf.identity(self.observ)

From 4eec3f3d483ba1f9a4a4060bb13ca85f1753c45c Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 7 Jun 2018 20:31:54 -0700
Subject: [PATCH 0065/2720] internal merge of PR #859

PiperOrigin-RevId: 199739336
---
 tensor2tensor/layers/common_layers.py |   2 +-
 tensor2tensor/models/lstm.py          | 159 +++++++++++++++++---------
 2 files changed, 109 insertions(+), 52 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index afbb0f3d9..fab035723 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1300,7 +1300,7 @@ def length_from_embedding(emb):
   Returns:
     a Tensor with shape [batch].
   """
-  return tf.reduce_sum(mask_from_embedding(emb), [1, 2, 3])
+  return tf.cast(tf.reduce_sum(mask_from_embedding(emb), [1, 2, 3]), tf.int32)
 
 
 def mask_leq(target_length, source_length):
diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 0b933b6ce..beed64b37 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -27,37 +27,73 @@
 import tensorflow as tf
 
 
-def lstm(inputs, hparams, train, name, initial_state=None,
-         sequence_length=None):
-  """Run LSTM cell on inputs, assuming they are [batch x time x size]."""
-
-  def dropout_lstm_cell():
-    return tf.contrib.rnn.DropoutWrapper(
-        tf.contrib.rnn.BasicLSTMCell(hparams.hidden_size),
-        input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
-
-  layers = [dropout_lstm_cell() for _ in range(hparams.num_hidden_layers)]
+def _dropout_lstm_cell(hparams, train):
+  return tf.contrib.rnn.DropoutWrapper(
+      tf.contrib.rnn.LSTMCell(hparams.hidden_size),
+      input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
+
+
+def lstm(inputs, sequence_length, hparams, train, name, initial_state=None):
+  """Adds a stack of LSTM layers on top of input.
+
+  Args:
+    inputs: The input `Tensor`, shaped `[batch_size, time_steps, hidden_size]`.
+    sequence_length: Lengths of the actual input sequence, excluding padding; a
+        `Tensor` shaped `[batch_size]`.
+    hparams: tf.contrib.training.HParams; hyperparameters.
+    train: bool; `True` when constructing training graph to enable dropout.
+    name: string; Create variable names under this scope.
+    initial_state: tuple of `LSTMStateTuple`s; the initial state of each layer.
+
+  Returns:
+    A tuple (outputs, states), where:
+      outputs: The output `Tensor`, shaped `[batch_size, time_steps,
+        hidden_size]`.
+      states: A tuple of `LSTMStateTuple`s; the final state of each layer.
+        Bidirectional LSTM returns a concatenation of last forward and backward
+        state, reduced to the original dimensionality.
+  """
+  layers = [_dropout_lstm_cell(hparams, train)
+            for _ in range(hparams.num_hidden_layers)]
   with tf.variable_scope(name):
     return tf.nn.dynamic_rnn(
         tf.contrib.rnn.MultiRNNCell(layers),
         inputs,
+        sequence_length,
         initial_state=initial_state,
-        sequence_length=sequence_length,
         dtype=tf.float32,
         time_major=False)
 
 
 def lstm_attention_decoder(inputs, hparams, train, name, initial_state,
-                           encoder_outputs, encoder_output_length=None,
-                           decoder_input_length=None):
-  """Run LSTM cell with attention on inputs of shape [batch x time x size]."""
-
-  def dropout_lstm_cell():
-    return tf.contrib.rnn.DropoutWrapper(
-        tf.nn.rnn_cell.BasicLSTMCell(hparams.hidden_size),
-        input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
-
-  layers = [dropout_lstm_cell() for _ in range(hparams.num_hidden_layers)]
+                           encoder_outputs, encoder_output_length,
+                           decoder_input_length):
+  """Run LSTM cell with attention on inputs of shape [batch x time x size].
+
+  Args:
+    inputs: The decoder input `Tensor`, shaped `[batch_size, decoder_steps,
+        hidden_size]`.
+    hparams: tf.contrib.training.HParams; hyperparameters.
+    train: bool; `True` when constructing training graph to enable dropout.
+    name: string; Create variable names under this scope.
+    initial_state: Tuple of `LSTMStateTuple`s; the initial state of each layer.
+    encoder_outputs: Encoder outputs; a `Tensor` shaped `[batch_size,
+        encoder_steps, hidden_size]`.
+    encoder_output_length: Lengths of the actual encoder outputs, excluding
+        padding; a `Tensor` shaped `[batch_size]`.
+    decoder_input_length: Lengths of the actual decoder inputs, excluding
+        padding; a `Tensor` shaped `[batch_size]`.
+
+  Raises:
+    ValueError: If the hparams.attention_mechanism is anything other than
+        luong or bahdanau.
+
+  Returns:
+    The decoder output `Tensor`, shaped `[batch_size, decoder_steps,
+    hidden_size]`.
+  """
+  layers = [_dropout_lstm_cell(hparams, train)
+            for _ in range(hparams.num_hidden_layers)]
   if hparams.attention_mechanism == "luong":
     attention_mechanism_class = tf.contrib.seq2seq.LuongAttention
   elif hparams.attention_mechanism == "bahdanau":
@@ -81,36 +117,47 @@ def dropout_lstm_cell():
       cell_state=initial_state)
 
   with tf.variable_scope(name):
-    output, state = tf.nn.dynamic_rnn(
+    output, _ = tf.nn.dynamic_rnn(
         cell,
         inputs,
+        decoder_input_length,
         initial_state=initial_state,
-        sequence_length=decoder_input_length,
         dtype=tf.float32,
         time_major=False)
-
-    # For multi-head attention project output back to hidden size
+    # output is [batch_size, decoder_steps, attention_size], where
+    # attention_size is either hparams.hidden_size (when
+    # hparams.output_attention is 0) or hparams.attention_layer_size (when
+    # hparams.output_attention is 1) times the number of attention heads.
+    #
+    # For multi-head attention project output back to hidden size.
     if hparams.output_attention == 1 and hparams.num_heads > 1:
       output = tf.layers.dense(output, hparams.hidden_size)
 
-    return output, state
+    return output
 
 
 def lstm_seq2seq_internal(inputs, targets, hparams, train):
   """The basic LSTM seq2seq model, main step used for training."""
   with tf.variable_scope("lstm_seq2seq"):
     if inputs is not None:
+      inputs_length = common_layers.length_from_embedding(inputs)
       # Flatten inputs.
       inputs = common_layers.flatten4d3d(inputs)
+
       # LSTM encoder.
-      _, final_encoder_state = lstm(
-          tf.reverse(inputs, axis=[1]), hparams, train, "encoder")
+      inputs = tf.reverse_sequence(inputs, inputs_length, seq_axis=1)
+      _, final_encoder_state = lstm(inputs, inputs_length, hparams, train,
+                                    "encoder")
     else:
       final_encoder_state = None
+
     # LSTM decoder.
     shifted_targets = common_layers.shift_right(targets)
+    # Add 1 to account for the padding added to the left from shift_right
+    targets_length = common_layers.length_from_embedding(shifted_targets) + 1
     decoder_outputs, _ = lstm(
         common_layers.flatten4d3d(shifted_targets),
+        targets_length,
         hparams,
         train,
         "decoder",
@@ -127,41 +174,40 @@ def lstm_seq2seq_internal_attention(inputs, targets, hparams, train):
     inputs_length = common_layers.length_from_embedding(inputs)
     # Flatten inputs.
     inputs = common_layers.flatten4d3d(inputs)
+
     # LSTM encoder.
+    inputs = tf.reverse_sequence(inputs, inputs_length, seq_axis=1)
     encoder_outputs, final_encoder_state = lstm(
-        inputs, hparams, train, "encoder", sequence_length=inputs_length)
-    # LSTM decoder with attention
+        inputs, inputs_length, hparams, train, "encoder")
+
+    # LSTM decoder with attention.
     shifted_targets = common_layers.shift_right(targets)
     # Add 1 to account for the padding added to the left from shift_right
     targets_length = common_layers.length_from_embedding(shifted_targets) + 1
-    decoder_outputs, _ = lstm_attention_decoder(
+    decoder_outputs = lstm_attention_decoder(
         common_layers.flatten4d3d(shifted_targets), hparams, train, "decoder",
-        final_encoder_state, encoder_outputs,
-        encoder_output_length=inputs_length,
-        decoder_input_length=targets_length)
+        final_encoder_state, encoder_outputs, inputs_length, targets_length)
     return tf.expand_dims(decoder_outputs, axis=2)
 
 
-def lstm_bid_encoder(inputs, hparams, train, name):
+def lstm_bid_encoder(inputs, sequence_length, hparams, train, name):
   """Bidirectional LSTM for encoding inputs that are [batch x time x size]."""
 
-  def dropout_lstm_cell():
-    return tf.contrib.rnn.DropoutWrapper(
-        tf.contrib.rnn.BasicLSTMCell(hparams.hidden_size),
-        input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
-
   with tf.variable_scope(name):
     cell_fw = tf.contrib.rnn.MultiRNNCell(
-        [dropout_lstm_cell() for _ in range(hparams.num_hidden_layers)])
+        [_dropout_lstm_cell(hparams, train)
+         for _ in range(hparams.num_hidden_layers)])
 
     cell_bw = tf.contrib.rnn.MultiRNNCell(
-        [dropout_lstm_cell() for _ in range(hparams.num_hidden_layers)])
+        [_dropout_lstm_cell(hparams, train)
+         for _ in range(hparams.num_hidden_layers)])
 
     ((encoder_fw_outputs, encoder_bw_outputs),
      (encoder_fw_state, encoder_bw_state)) = tf.nn.bidirectional_dynamic_rnn(
-         cell_fw=cell_fw,
-         cell_bw=cell_bw,
-         inputs=inputs,
+         cell_fw,
+         cell_bw,
+         inputs,
+         sequence_length,
          dtype=tf.float32,
          time_major=False)
 
@@ -196,19 +242,24 @@ def lstm_seq2seq_internal_bid_encoder(inputs, targets, hparams, train):
   """The basic LSTM seq2seq model with bidirectional encoder."""
   with tf.variable_scope("lstm_seq2seq_bid_encoder"):
     if inputs is not None:
+      inputs_length = common_layers.length_from_embedding(inputs)
       # Flatten inputs.
       inputs = common_layers.flatten4d3d(inputs)
       # LSTM encoder.
       _, final_encoder_state = lstm_bid_encoder(
-          tf.reverse(inputs, axis=[1]), hparams, train, "encoder")
+          inputs, inputs_length, hparams, train, "encoder")
     else:
+      inputs_length = None
       final_encoder_state = None
     # LSTM decoder.
     shifted_targets = common_layers.shift_right(targets)
+    # Add 1 to account for the padding added to the left from shift_right
+    targets_length = common_layers.length_from_embedding(shifted_targets) + 1
     hparams_decoder = copy.copy(hparams)
     hparams_decoder.hidden_size = 2 * hparams.hidden_size
     decoder_outputs, _ = lstm(
         common_layers.flatten4d3d(shifted_targets),
+        targets_length,
         hparams_decoder,
         train,
         "decoder",
@@ -220,18 +271,22 @@ def lstm_seq2seq_internal_attention_bid_encoder(inputs, targets, hparams,
                                                 train):
   """LSTM seq2seq model with attention, main step used for training."""
   with tf.variable_scope("lstm_seq2seq_attention_bid_encoder"):
+    inputs_length = common_layers.length_from_embedding(inputs)
     # Flatten inputs.
     inputs = common_layers.flatten4d3d(inputs)
     # LSTM encoder.
     encoder_outputs, final_encoder_state = lstm_bid_encoder(
-        tf.reverse(inputs, axis=[1]), hparams, train, "encoder")
+        inputs, inputs_length, hparams, train, "encoder")
     # LSTM decoder with attention
     shifted_targets = common_layers.shift_right(targets)
+    # Add 1 to account for the padding added to the left from shift_right
+    targets_length = common_layers.length_from_embedding(shifted_targets) + 1
     hparams_decoder = copy.copy(hparams)
     hparams_decoder.hidden_size = 2 * hparams.hidden_size
-    decoder_outputs, _ = lstm_attention_decoder(
+    decoder_outputs = lstm_attention_decoder(
         common_layers.flatten4d3d(shifted_targets), hparams_decoder, train,
-        "decoder", final_encoder_state, encoder_outputs)
+        "decoder", final_encoder_state, encoder_outputs,
+        inputs_length, targets_length)
     return tf.expand_dims(decoder_outputs, axis=2)
 
 
@@ -244,11 +299,13 @@ def body(self, features):
       raise ValueError("LSTM models fail with orthogonal initializer.")
     train = self._hparams.mode == tf.estimator.ModeKeys.TRAIN
     inputs = features.get("inputs")
+    inputs_length = common_layers.length_from_embedding(inputs)
     # Flatten inputs.
     inputs = common_layers.flatten4d3d(inputs)
     # LSTM encoder.
-    encoder_output, _ = lstm(
-        tf.reverse(inputs, axis=[1]), self._hparams, train, "encoder")
+    inputs = tf.reverse_sequence(inputs, inputs_length, seq_axis=1)
+    encoder_output, _ = lstm(inputs, inputs_length, self._hparams, train,
+                             "encoder")
     return tf.expand_dims(encoder_output, axis=2)
 
 

From 60aa595233235d21d77ded917eb9b31b180fdd37 Mon Sep 17 00:00:00 2001
From: Jaehyeon Kim <jaywalnut310@users.noreply.github.com>
Date: Fri, 8 Jun 2018 12:49:45 +0900
Subject: [PATCH 0066/2720] Add Transformer's encoder-decoder cache (#827)

* fast enc-dec attention

* fix bugs for test

* caching splitted keys and values
---
 tensor2tensor/layers/common_attention.py | 68 ++++++++++++++++++------
 tensor2tensor/models/transformer.py      | 20 +++++--
 tensor2tensor/utils/decoding.py          | 50 ++++++++++++-----
 3 files changed, 105 insertions(+), 33 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 42f2eb745..815d0569b 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2604,6 +2604,28 @@ def masked_local_attention_2d(q,
     return output
 
 
+def compute_attention_component(antecedent,
+                       total_depth,
+                       filter_width=1,
+                       padding="VALID",
+                       name="c"):
+  """Computes attention compoenent (query, key or value).
+
+  Args:
+    antecedent: a Tensor with shape [batch, length, channels]
+    total_depth: an integer
+    filter_width: An integer specifying how wide you want the attention component to be.
+    padding: One of "VALID", "SAME" or "LEFT". Default is VALID: No padding.
+    name: a string specifying scope name.
+ 
+  Returns:
+    c : [batch, length, depth] tensor
+  """
+  if filter_width == 1:
+    return common_layers.dense(antecedent, total_depth, use_bias=False, name=name)
+  else:
+    return common_layers.conv1d(antecedent, total_depth, filter_width, padding, name=name)
+
 def compute_qkv(query_antecedent,
                 memory_antecedent,
                 total_key_depth,
@@ -2618,7 +2640,7 @@ def compute_qkv(query_antecedent,
     query_antecedent: a Tensor with shape [batch, length_q, channels]
     memory_antecedent: a Tensor with shape [batch, length_m, channels]
     total_key_depth: an integer
-    total_value_depth: and integer
+    total_value_depth: an integer
     q_filter_width: An integer specifying how wide you want the query to be.
     kv_filter_width: An integer specifying how wide you want the keys and values
     to be.
@@ -2630,15 +2652,11 @@ def compute_qkv(query_antecedent,
   """
   if memory_antecedent is None:
     memory_antecedent = query_antecedent
-  def _compute(inp, depth, filter_width, padding, name):
-    if filter_width == 1:
-      return common_layers.dense(inp, depth, use_bias=False, name=name)
-    return common_layers.conv1d(inp, depth, filter_width, padding, name=name)
-  q = _compute(
+  q = compute_attention_component(
       query_antecedent, total_key_depth, q_filter_width, q_padding, "q")
-  k = _compute(
+  k = compute_attention_component(
       memory_antecedent, total_key_depth, kv_filter_width, kv_padding, "k")
-  v = _compute(
+  v = compute_attention_component(
       memory_antecedent, total_value_depth, kv_filter_width, kv_padding, "v")
   return q, k, v
 
@@ -2663,7 +2681,7 @@ def multihead_attention(query_antecedent,
                         cache=None,
                         gap_size=0,
                         num_memory_blocks=2,
-                        name=None,
+                        name="multihead_attention",
                         save_weights_to=None,
                         make_image_summary=True,
                         dropout_broadcast_dims=None,
@@ -2747,12 +2765,12 @@ def multihead_attention(query_antecedent,
                      "attention heads (%d)." % (total_value_depth, num_heads))
   with tf.variable_scope(
       name,
-      default_name="multihead_attention",
       values=[query_antecedent, memory_antecedent]):
-    q, k, v = compute_qkv(query_antecedent, memory_antecedent, total_key_depth,
-                          total_value_depth, q_filter_width, kv_filter_width,
-                          q_padding, kv_padding)
 
+    if cache is None or memory_antecedent is None:
+      q, k, v =  compute_qkv(query_antecedent, memory_antecedent, total_key_depth,
+                            total_value_depth, q_filter_width, kv_filter_width,
+                            q_padding, kv_padding)
     if cache is not None:
       if attention_type != "dot_product":
         # TODO(petershaw): Support caching when using relative position
@@ -2763,12 +2781,24 @@ def multihead_attention(query_antecedent,
       if bias is None:
         raise ValueError("Bias required for caching. See function docstring "
                          "for details.")
-      k = cache["k"] = tf.concat([cache["k"], k], axis=1)
-      v = cache["v"] = tf.concat([cache["v"], v], axis=1)
+
+      if memory_antecedent is not None:
+        # Encoder-Decoder Attention Cache
+        q = compute_attention_component(query_antecedent, total_key_depth,
+                                        q_filter_width, q_padding, "q")
+        k = cache["k_encdec"]
+        v = cache["v_encdec"]
+      else:
+        k = split_heads(k, num_heads)
+        v = split_heads(v, num_heads)
+        k = cache["k"] = tf.concat([cache["k"], k], axis=2)
+        v = cache["v"] = tf.concat([cache["v"], v], axis=2)
 
     q = split_heads(q, num_heads)
-    k = split_heads(k, num_heads)
-    v = split_heads(v, num_heads)
+    if cache is None:
+      k = split_heads(k, num_heads)
+      v = split_heads(v, num_heads)
+
     key_depth_per_head = total_key_depth // num_heads
     q *= key_depth_per_head**-0.5
 
@@ -2808,6 +2838,10 @@ def multihead_attention(query_antecedent,
       x = dilated_self_attention_1d(q, k, v, block_length, block_width,
                                     gap_size, num_memory_blocks)
     x = combine_heads(x)
+
+    # Set last dim specifically.
+    x.set_shape(x.shape.as_list()[:-1] + [total_value_depth])
+
     x = common_layers.dense(
         x, output_depth, use_bias=False, name="output_transform")
     if additional_returned_value is not None:
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 1c7ca4553..32134db4c 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -477,14 +477,26 @@ def fast_decode(encoder_output,
 
   cache = {
       "layer_%d" % layer: {
-          "k": tf.zeros([batch_size, 0, key_channels]),
-          "v": tf.zeros([batch_size, 0, value_channels]),
+          "k": common_attention.split_heads(tf.zeros([batch_size, 0, key_channels]), hparams.num_heads),
+          "v": common_attention.split_heads(tf.zeros([batch_size, 0, value_channels]), hparams.num_heads),
           "f": tf.zeros([batch_size, 0, hparams.hidden_size]),
       } for layer in range(num_layers)
   }
 
   if encoder_output is not None:
-    cache["encoder_output"] = encoder_output
+    for layer in range(num_layers):
+      layer_name = "layer_%d" % layer
+      with tf.variable_scope("body/decoder/%s/encdec_attention/multihead_attention" % layer_name):
+        k_encdec = common_attention.compute_attention_component(
+              encoder_output, key_channels, name="k")
+        k_encdec = common_attention.split_heads(k_encdec, hparams.num_heads)
+        v_encdec = common_attention.compute_attention_component(
+              encoder_output, value_channels, name="v")
+        v_encdec = common_attention.split_heads(v_encdec, hparams.num_heads)
+      cache[layer_name]["k_encdec"] = k_encdec
+      cache[layer_name]["v_encdec"] = v_encdec
+
+    cache["encoder_output"] = tf.zeros_like(tf.placeholder(tf.float32, (None, 0, 0))) # Just a flag
     cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias
 
   if beam_size > 1:  # Beam Search
@@ -886,7 +898,6 @@ def transformer_decoder(decoder_input,
           x = common_layers.layer_postprocess(x, y, hparams)
         if encoder_output is not None:
           with tf.variable_scope("encdec_attention"):
-            # TODO(llion): Add caching.
             y = common_attention.multihead_attention(
                 common_layers.layer_preprocess(x, hparams),
                 encoder_output,
@@ -897,6 +908,7 @@ def transformer_decoder(decoder_input,
                 hparams.num_heads,
                 hparams.attention_dropout,
                 save_weights_to=save_weights_to,
+                cache=layer_cache,
                 make_image_summary=make_image_summary,
                 dropout_broadcast_dims=attention_dropout_broadcast_dims,
                 max_length=hparams.get("max_length"))
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 15996932c..98b98df16 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -19,6 +19,10 @@
 
 import operator
 import os
+import time
+
+# Dependency imports
+
 import numpy as np
 import six
 
@@ -38,7 +42,7 @@ def decode_hparams(overrides=""):
   """Hyperparameters for decoding."""
   hp = tf.contrib.training.HParams(
       save_images=False,
-      log_targets=True,
+      log_results=True,
       extra_length=100,
       batch_size=0,
       beam_size=4,
@@ -64,7 +68,7 @@ def log_decode_results(inputs,
                        save_images=False,
                        model_dir=None,
                        identity_output=False,
-                       log_targets=True):
+                       log_results=True):
   """Log inference results."""
   is_image = "image" in problem_name
   decoded_inputs = None
@@ -78,7 +82,8 @@ def log_decode_results(inputs,
     else:
       decoded_inputs = inputs_vocab.decode(_save_until_eos(inputs, is_image))
 
-    tf.logging.info("Inference results INPUT: %s" % decoded_inputs)
+    if log_results:
+      tf.logging.info("Inference results INPUT: %s" % decoded_inputs)
 
   decoded_targets = None
   decoded_outputs = None
@@ -88,11 +93,11 @@ def log_decode_results(inputs,
       decoded_targets = " ".join(map(str, targets.flatten()))
   else:
     decoded_outputs = targets_vocab.decode(_save_until_eos(outputs, is_image))
-    if targets is not None and log_targets:
+    if targets is not None and log_results:
       decoded_targets = targets_vocab.decode(_save_until_eos(targets, is_image))
-
-  tf.logging.info("Inference results OUTPUT: %s" % decoded_outputs)
-  if targets is not None and log_targets:
+  if log_results:
+    tf.logging.info("Inference results OUTPUT: %s" % decoded_outputs)
+  if targets is not None and log_results:
     tf.logging.info("Inference results TARGET: %s" % decoded_targets)
   return decoded_inputs, decoded_outputs, decoded_targets
 
@@ -181,7 +186,7 @@ def decode_from_dataset(estimator,
             model_dir=estimator.model_dir,
             identity_output=decode_hp.identity_output,
             targets=targets,
-            log_targets=decode_hp.log_targets)
+            log_results=decode_hp.log_results)
         decoded_outputs.append(decoded)
         if decode_hp.write_beam_scores:
           decoded_scores.append(score)
@@ -197,7 +202,7 @@ def decode_from_dataset(estimator,
           model_dir=estimator.model_dir,
           identity_output=decode_hp.identity_output,
           targets=targets,
-          log_targets=decode_hp.log_targets)
+          log_results=decode_hp.log_results)
       decoded_outputs.append(decoded)
 
     # Write out predictions if decode_to_file passed
@@ -258,7 +263,22 @@ def input_fn():
 
   decodes = []
   result_iter = estimator.predict(input_fn, checkpoint_path=checkpoint_path)
-  for result in result_iter:
+
+  start_time = time.time()
+  total_time_per_step = 0
+  total_cnt = 0
+
+  def timer(gen):
+    while True:
+      try:
+        start_time = time.time()
+        item = next(gen)
+        elapsed_time = time.time() - start_time
+        yield elapsed_time, item
+      except StopIteration:
+        break
+
+  for elapsed_time, result in timer(result_iter):
     if decode_hp.return_beams:
       beam_decodes = []
       beam_scores = []
@@ -271,7 +291,8 @@ def input_fn():
         score = scores and scores[k]
         _, decoded_outputs, _ = log_decode_results(result["inputs"], beam,
                                                    problem_name, None,
-                                                   inputs_vocab, targets_vocab)
+                                                   inputs_vocab, targets_vocab,
+                                                   log_results=decode_hp.log_results)
         beam_decodes.append(decoded_outputs)
         if decode_hp.write_beam_scores:
           beam_scores.append(score)
@@ -284,8 +305,13 @@ def input_fn():
     else:
       _, decoded_outputs, _ = log_decode_results(
           result["inputs"], result["outputs"], problem_name,
-          None, inputs_vocab, targets_vocab)
+          None, inputs_vocab, targets_vocab,
+          log_results=decode_hp.log_results)
       decodes.append(decoded_outputs)
+    total_time_per_step += elapsed_time
+    total_cnt += result["outputs"].shape[-1]
+  tf.logging.info("Elapsed Time: %5.5f" % (time.time() - start_time))
+  tf.logging.info("Averaged Single Token Generation Time: %5.7f" % (total_time_per_step / total_cnt))
 
   # Reversing the decoded inputs and outputs because they were reversed in
   # _decode_batch_input_fn

From 3c0c200e8e0da83444430caca91197dd3958b64d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 7 Jun 2018 20:55:20 -0700
Subject: [PATCH 0067/2720] internal merge of PR #827

PiperOrigin-RevId: 199742179
---
 tensor2tensor/layers/common_attention.py | 170 ++++++++++++++---------
 tensor2tensor/models/transformer.py      | 102 +++++++++-----
 tensor2tensor/utils/decoding.py          |  57 ++++----
 3 files changed, 198 insertions(+), 131 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 815d0569b..1967fa2e6 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -156,6 +156,7 @@ def decorator(x, *args, **kwargs):
       ),
       recompute_grad=True,
   )
+
   def memeff_attention_fn(*args, **kwargs):
     """Modify args/kwargs for compatibility with recompute_grad."""
     kwargs = kwargs.copy()
@@ -675,8 +676,8 @@ def add_positional_embedding_nd(x, max_length, name):
         tf.get_variable(
             name + "_%d" % i,
             shape,
-            initializer=tf.random_normal_initializer(0, depth**-0.5)) *
-        (depth**0.5))
+            initializer=tf.random_normal_initializer(0, depth**-0.5)) * (depth**
+                                                                         0.5))
     x += tf.slice(var, start, size)
   return x
 
@@ -832,8 +833,10 @@ def attention_bias_same_segment(query_segment_id, memory_segment_id):
   Returns:
     a `Tensor` with shape [batch, 1, query_length, memory_length].
   """
-  ret = tf.to_float(tf.not_equal(tf.expand_dims(query_segment_id, 2),
-                                 tf.expand_dims(memory_segment_id, 1))) * -1e9
+  ret = tf.to_float(
+      tf.not_equal(
+          tf.expand_dims(query_segment_id, 2),
+          tf.expand_dims(memory_segment_id, 1))) * -1e9
   return tf.expand_dims(ret, axis=1)
 
 
@@ -1236,9 +1239,8 @@ def grouped_attention_multihead(query_antecedent,
     q_requests = tf.one_hot(q_group, num_groups, axis=-1)
     m_requests = tf.to_float(tf.greater(m_pred_biased, 0.0))
     # include first memory position in all groups, to avoid division by zero.
-    m_requests = tf.maximum(m_requests,
-                            tf.reshape(
-                                tf.one_hot([0], length_kv), [1, length_kv, 1]))
+    m_requests = tf.maximum(
+        m_requests, tf.reshape(tf.one_hot([0], length_kv), [1, length_kv, 1]))
     q_group_size = tf.reduce_sum(q_requests, 1)
     m_group_size = tf.reduce_sum(m_requests, 1)
     q_group_target_size = tf.to_float(length_q) / tf.to_float(num_groups)
@@ -1247,12 +1249,12 @@ def grouped_attention_multihead(query_antecedent,
         tf.to_float(num_groups))
     capacity_q = tf.minimum(
         length_q,
-        tf.to_int32(
-            q_group_target_size * multiplicative_overhead + additive_overhead))
+        tf.to_int32(q_group_target_size * multiplicative_overhead +
+                    additive_overhead))
     capacity_m = tf.minimum(
         length_kv,
-        tf.to_int32(
-            m_group_target_size * multiplicative_overhead + additive_overhead))
+        tf.to_int32(m_group_target_size * multiplicative_overhead +
+                    additive_overhead))
     q_dispatcher = expert_utils.TruncatingDispatcher(q_requests, capacity_q)
     m_dispatcher = expert_utils.TruncatingDispatcher(m_requests, capacity_m)
     q_gates = q_dispatcher.gates()
@@ -1311,12 +1313,12 @@ def grouped_attention_multihead(query_antecedent,
     # decrease for groups that are too big.
     q_group_deviation = (q_group_size / q_group_target_size) - 1.0
     q_balance_loss = tf.reduce_sum(
-        tf.reduce_mean(q_pred_biased, axis=1) * q_group_deviation
-    ) / tf.to_float(batch)
+        tf.reduce_mean(q_pred_biased, axis=1) *
+        q_group_deviation) / tf.to_float(batch)
     m_group_deviation = (m_group_size / m_group_target_size) - 1.0
     m_balance_loss = tf.reduce_sum(
-        tf.reduce_mean(m_pred_biased, axis=1) * m_group_deviation
-    ) / tf.to_float(batch)
+        tf.reduce_mean(m_pred_biased, axis=1) *
+        m_group_deviation) / tf.to_float(batch)
 
     # The losses in this function only propagate back to variables
     # defined in this function, and the losses outside of this
@@ -1642,7 +1644,8 @@ def dot_product_self_attention_relative_v2(q,
     A Tensor.
   """
   with tf.variable_scope(
-      name, default_name="dot_product_self_attention_relative_v2",
+      name,
+      default_name="dot_product_self_attention_relative_v2",
       values=[q, k, v]):
 
     # This calculation only works for self attention.
@@ -1661,8 +1664,8 @@ def dot_product_self_attention_relative_v2(q,
     # [batch, num_heads, query_length, max_length]
     rel_logits = common_layers.dense(q, max_length, name="rel0")
     # [batch, num_heads, query_length, max_length]
-    rel_logits = tf.slice(
-        rel_logits, [0, 0, 0, max_length - length], [-1, -1, -1, -1])
+    rel_logits = tf.slice(rel_logits, [0, 0, 0, max_length - length],
+                          [-1, -1, -1, -1])
     rel_logits = _relative_position_to_absolute_position_masked(rel_logits)
     logits += rel_logits
 
@@ -1728,8 +1731,9 @@ def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
     v = tf.reshape(v, [batch, heads, num_blocks, block_length, depth_v])
     # attention shape: [batch, heads, num_blocks, block_length, block_length]
     attention = tf.matmul(q, k, transpose_b=True)
-    attention += tf.reshape(attention_bias_lower_triangle(block_length),
-                            [1, 1, 1, block_length, block_length])
+    attention += tf.reshape(
+        attention_bias_lower_triangle(block_length),
+        [1, 1, 1, block_length, block_length])
     attention = tf.nn.softmax(attention)
     # initial output shape: [batch, heads, num_blocks, block_length, depth_v]
     output = tf.matmul(attention, v)
@@ -1739,8 +1743,12 @@ def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
     return output
 
 
-def masked_local_attention_1d(q, k, v, block_length=128,
-                              make_image_summary=False, name=None):
+def masked_local_attention_1d(q,
+                              k,
+                              v,
+                              block_length=128,
+                              make_image_summary=False,
+                              name=None):
   """Attention to the source position and a neighborhood to the left of it.
 
   The sequence is divided into blocks of length block_size.
@@ -1817,14 +1825,13 @@ def local(x, depth):
       cur_block = tf.slice(x, [0, 0, 1, 0, 0], [-1, -1, -1, -1, -1])
       local_block = tf.concat([prev_block, cur_block], 3)
       return tf.reshape(local_block,
-                        [batch, heads, num_blocks - 1,
-                         block_length * 2, depth])
+                        [batch, heads, num_blocks - 1, block_length * 2, depth])
 
     local_k = local(k, depth_k)
     local_v = local(v, depth_v)
     tail_q = tf.slice(q, [0, 0, 1, 0, 0], [-1, -1, -1, -1, -1])
-    tail_q = tf.reshape(tail_q, [batch, heads, num_blocks - 1,
-                                 block_length, depth_k])
+    tail_q = tf.reshape(tail_q,
+                        [batch, heads, num_blocks - 1, block_length, depth_k])
     local_length = common_layers.shape_list(local_k)[3]
 
     # [batch, heads, num_blocks - 1, block_length, local_length]
@@ -1841,8 +1848,8 @@ def local(x, depth):
     # The naive way currently causes errors due to empty tensors.
     # output: [batch, heads, num_blocks-1, block_length, depth_v]
     output = tf.matmul(attention, local_v)
-    output = tf.reshape(output, [
-        batch, heads, (num_blocks-1)*block_length, depth_v])
+    output = tf.reshape(
+        output, [batch, heads, (num_blocks - 1) * block_length, depth_v])
     output = tf.concat([first_output, output], axis=2)
     output = tf.slice(output, [0, 0, 0, 0], [-1, -1, original_length, -1])
     output = tf.reshape(output, [batch, heads, original_length, depth_v])
@@ -2096,8 +2103,8 @@ def gather_dilated_memory_blocks(x,
   for block_id in range(num_memory_blocks):
     block_end_index = -(query_block_size + gap_size *
                         (block_id + 1) + memory_block_size * block_id) - 1
-    block_start_index = ((memory_block_size + gap_size) * (num_memory_blocks -
-                                                           (block_id + 1)))
+    block_start_index = (
+        (memory_block_size + gap_size) * (num_memory_blocks - (block_id + 1)))
     if direction != "left":
       [block_end_index,
        block_start_index] = [-block_start_index - 1, -block_end_index + 1]
@@ -2605,26 +2612,30 @@ def masked_local_attention_2d(q,
 
 
 def compute_attention_component(antecedent,
-                       total_depth,
-                       filter_width=1,
-                       padding="VALID",
-                       name="c"):
+                                total_depth,
+                                filter_width=1,
+                                padding="VALID",
+                                name="c"):
   """Computes attention compoenent (query, key or value).
 
   Args:
     antecedent: a Tensor with shape [batch, length, channels]
     total_depth: an integer
-    filter_width: An integer specifying how wide you want the attention component to be.
+    filter_width: An integer specifying how wide you want the attention
+      component to be.
     padding: One of "VALID", "SAME" or "LEFT". Default is VALID: No padding.
     name: a string specifying scope name.
- 
+
   Returns:
     c : [batch, length, depth] tensor
   """
   if filter_width == 1:
-    return common_layers.dense(antecedent, total_depth, use_bias=False, name=name)
+    return common_layers.dense(
+        antecedent, total_depth, use_bias=False, name=name)
   else:
-    return common_layers.conv1d(antecedent, total_depth, filter_width, padding, name=name)
+    return common_layers.conv1d(
+        antecedent, total_depth, filter_width, padding, name=name)
+
 
 def compute_qkv(query_antecedent,
                 memory_antecedent,
@@ -2652,12 +2663,12 @@ def compute_qkv(query_antecedent,
   """
   if memory_antecedent is None:
     memory_antecedent = query_antecedent
-  q = compute_attention_component(
-      query_antecedent, total_key_depth, q_filter_width, q_padding, "q")
-  k = compute_attention_component(
-      memory_antecedent, total_key_depth, kv_filter_width, kv_padding, "k")
-  v = compute_attention_component(
-      memory_antecedent, total_value_depth, kv_filter_width, kv_padding, "v")
+  q = compute_attention_component(query_antecedent, total_key_depth,
+                                  q_filter_width, q_padding, "q")
+  k = compute_attention_component(memory_antecedent, total_key_depth,
+                                  kv_filter_width, kv_padding, "k")
+  v = compute_attention_component(memory_antecedent, total_value_depth,
+                                  kv_filter_width, kv_padding, "v")
   return q, k, v
 
 
@@ -2763,14 +2774,12 @@ def multihead_attention(query_antecedent,
   if total_value_depth % num_heads != 0:
     raise ValueError("Value depth (%d) must be divisible by the number of "
                      "attention heads (%d)." % (total_value_depth, num_heads))
-  with tf.variable_scope(
-      name,
-      values=[query_antecedent, memory_antecedent]):
+  with tf.variable_scope(name, values=[query_antecedent, memory_antecedent]):
 
     if cache is None or memory_antecedent is None:
-      q, k, v =  compute_qkv(query_antecedent, memory_antecedent, total_key_depth,
-                            total_value_depth, q_filter_width, kv_filter_width,
-                            q_padding, kv_padding)
+      q, k, v = compute_qkv(query_antecedent, memory_antecedent,
+                            total_key_depth, total_value_depth, q_filter_width,
+                            kv_filter_width, q_padding, kv_padding)
     if cache is not None:
       if attention_type != "dot_product":
         # TODO(petershaw): Support caching when using relative position
@@ -2808,25 +2817,47 @@ def multihead_attention(query_antecedent,
       if isinstance(x, tuple):
         x, additional_returned_value = x  # Unpack
     elif attention_type == "dot_product":
-      x = dot_product_attention(q, k, v, bias, dropout_rate, image_shapes,
-                                save_weights_to=save_weights_to,
-                                make_image_summary=make_image_summary,
-                                dropout_broadcast_dims=dropout_broadcast_dims)
+      x = dot_product_attention(
+          q,
+          k,
+          v,
+          bias,
+          dropout_rate,
+          image_shapes,
+          save_weights_to=save_weights_to,
+          make_image_summary=make_image_summary,
+          dropout_broadcast_dims=dropout_broadcast_dims)
     elif attention_type == "dot_product_relative":
-      x = dot_product_attention_relative(q, k, v, bias, max_relative_position,
-                                         dropout_rate, image_shapes,
-                                         make_image_summary=make_image_summary)
+      x = dot_product_attention_relative(
+          q,
+          k,
+          v,
+          bias,
+          max_relative_position,
+          dropout_rate,
+          image_shapes,
+          make_image_summary=make_image_summary)
     elif attention_type == "dot_product_relative_v2":
       x = dot_product_self_attention_relative_v2(
-          q, k, v, bias, max_length, dropout_rate, image_shapes,
+          q,
+          k,
+          v,
+          bias,
+          max_length,
+          dropout_rate,
+          image_shapes,
           make_image_summary=make_image_summary,
           dropout_broadcast_dims=dropout_broadcast_dims)
     elif attention_type == "local_within_block_mask_right":
-      x = masked_within_block_local_attention_1d(q, k, v,
-                                                 block_length=block_length)
+      x = masked_within_block_local_attention_1d(
+          q, k, v, block_length=block_length)
     elif attention_type == "local_mask_right":
-      x = masked_local_attention_1d(q, k, v, block_length=block_length,
-                                    make_image_summary=make_image_summary)
+      x = masked_local_attention_1d(
+          q,
+          k,
+          v,
+          block_length=block_length,
+          make_image_summary=make_image_summary)
     elif attention_type == "local_unmasked":
       x = local_attention_1d(
           q, k, v, block_length=block_length, filter_width=block_width)
@@ -2890,8 +2921,8 @@ def multihead_attention_2d(query_antecedent,
       name,
       default_name="multihead_attention_2d",
       values=[query_antecedent, memory_antecedent]):
-    q, k, v = compute_qkv(query_antecedent, memory_antecedent,
-                          total_key_depth, total_value_depth)
+    q, k, v = compute_qkv(query_antecedent, memory_antecedent, total_key_depth,
+                          total_value_depth)
     # after splitting, shape is [batch, heads, h, w, depth]
     q = split_heads_2d(q, num_heads)
     k = split_heads_2d(k, num_heads)
@@ -2953,10 +2984,13 @@ def ffn_self_attention_layer(x,
     else:
       q = tf.expand_dims(
           common_layers.dense(
-              x, filter_depth, use_bias=False, name="q_transform"), axis=2)
+              x, filter_depth, use_bias=False, name="q_transform"),
+          axis=2)
       kv_combined = tf.expand_dims(
           common_layers.dense(
-              tf.concat([x, x], axis=1), filter_depth, use_bias=False,
+              tf.concat([x, x], axis=1),
+              filter_depth,
+              use_bias=False,
               name="kv_transform"),
           axis=2)
       k, v = tf.split(kv_combined, [x_shape[1], x_shape[1]], axis=1)
@@ -3980,8 +4014,8 @@ def forward_internal(x, wqkv, wo, attention_bias, norm_scale, norm_bias):
         y += tf.nn.conv1d(o, wo_split[h], 1, "SAME")
     return y
 
-  key = ("multihead_self_attention_memory_efficient %s %s" % (num_heads,
-                                                              epsilon))
+  key = (
+      "multihead_self_attention_memory_efficient %s %s" % (num_heads, epsilon))
   if not forget:
     forward_fn = forward_internal
   elif key in _function_cache:
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 32134db4c..8c029a675 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -78,8 +78,10 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
                                   1.0 - hparams.layer_prepostprocess_dropout)
 
     encoder_output = transformer_encoder(
-        encoder_input, self_attention_bias,
-        hparams, nonpadding=features_to_nonpadding(features, "inputs"),
+        encoder_input,
+        self_attention_bias,
+        hparams,
+        nonpadding=features_to_nonpadding(features, "inputs"),
         save_weights_to=self.attention_weights,
         losses=losses)
 
@@ -242,8 +244,8 @@ def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
       return self._beam_decode_slow(features, decode_length, beam_size,
                                     top_beams, alpha)
     with tf.variable_scope(self.name):
-      return self._fast_decode(
-          features, decode_length, beam_size, top_beams, alpha)
+      return self._fast_decode(features, decode_length, beam_size, top_beams,
+                               alpha)
 
   def _fast_decode(self,
                    features,
@@ -287,8 +289,9 @@ def _fast_decode(self,
       if target_modality.is_class_modality:
         decode_length = 1
       else:
-        decode_length = (common_layers.shape_list(inputs)[1] +
-                         features.get("decode_length", decode_length))
+        decode_length = (
+            common_layers.shape_list(inputs)[1] + features.get(
+                "decode_length", decode_length))
 
       # TODO(llion): Clean up this reshaping logic.
       inputs = tf.expand_dims(inputs, axis=1)
@@ -304,7 +307,10 @@ def _fast_decode(self,
         inputs = input_modality.bottom_sharded(inputs, dp)
       with tf.variable_scope("body"):
         encoder_output, encoder_decoder_attention_bias = dp(
-            self.encode, inputs, features["target_space_id"], hparams,
+            self.encode,
+            inputs,
+            features["target_space_id"],
+            hparams,
             features=features)
       encoder_output = encoder_output[0]
       encoder_decoder_attention_bias = encoder_decoder_attention_bias[0]
@@ -325,8 +331,8 @@ def _fast_decode(self,
       partial_targets = tf.to_int64(partial_targets)
       partial_targets_shape = common_layers.shape_list(partial_targets)
       partial_targets_length = partial_targets_shape[1]
-      decode_length = (partial_targets_length +
-                       features.get("decode_length", decode_length))
+      decode_length = (
+          partial_targets_length + features.get("decode_length", decode_length))
       batch_size = partial_targets_shape[0]
 
     if hparams.pos == "timing":
@@ -378,9 +384,13 @@ def symbols_to_logits_fn(ids, i, cache):
 
       with tf.variable_scope("body"):
         body_outputs = dp(
-            self.decode, targets, cache.get("encoder_output"),
+            self.decode,
+            targets,
+            cache.get("encoder_output"),
             cache.get("encoder_decoder_attention_bias"),
-            bias, hparams, cache,
+            bias,
+            hparams,
+            cache,
             nonpadding=features_to_nonpadding(features, "targets"))
 
       with tf.variable_scope(target_modality.name):
@@ -394,9 +404,12 @@ def symbols_to_logits_fn(ids, i, cache):
         # iteration in order to fill the corresponding parts of the cache.
         # This would require broader changes, though.
         vocab_size = tf.shape(ret)[1]
+
         def forced_logits():
-          return tf.one_hot(tf.tile(partial_targets[:, i], [beam_size]),
-                            vocab_size, 0.0, -1e9)
+          return tf.one_hot(
+              tf.tile(partial_targets[:, i], [beam_size]), vocab_size, 0.0,
+              -1e9)
+
         ret = tf.cond(
             tf.less(i, partial_targets_length), forced_logits, lambda: ret)
       return ret, cache
@@ -477,26 +490,33 @@ def fast_decode(encoder_output,
 
   cache = {
       "layer_%d" % layer: {
-          "k": common_attention.split_heads(tf.zeros([batch_size, 0, key_channels]), hparams.num_heads),
-          "v": common_attention.split_heads(tf.zeros([batch_size, 0, value_channels]), hparams.num_heads),
-          "f": tf.zeros([batch_size, 0, hparams.hidden_size]),
+          "k":
+              common_attention.split_heads(
+                  tf.zeros([batch_size, 0, key_channels]), hparams.num_heads),
+          "v":
+              common_attention.split_heads(
+                  tf.zeros([batch_size, 0, value_channels]), hparams.num_heads),
+          "f":
+              tf.zeros([batch_size, 0, hparams.hidden_size]),
       } for layer in range(num_layers)
   }
 
   if encoder_output is not None:
     for layer in range(num_layers):
       layer_name = "layer_%d" % layer
-      with tf.variable_scope("body/decoder/%s/encdec_attention/multihead_attention" % layer_name):
+      with tf.variable_scope(
+          "body/decoder/%s/encdec_attention/multihead_attention" % layer_name):
         k_encdec = common_attention.compute_attention_component(
-              encoder_output, key_channels, name="k")
+            encoder_output, key_channels, name="k")
         k_encdec = common_attention.split_heads(k_encdec, hparams.num_heads)
         v_encdec = common_attention.compute_attention_component(
-              encoder_output, value_channels, name="v")
+            encoder_output, value_channels, name="v")
         v_encdec = common_attention.split_heads(v_encdec, hparams.num_heads)
       cache[layer_name]["k_encdec"] = k_encdec
       cache[layer_name]["v_encdec"] = v_encdec
 
-    cache["encoder_output"] = tf.zeros_like(tf.placeholder(tf.float32, (None, 0, 0))) # Just a flag
+    cache["encoder_output"] = tf.zeros_like(
+        tf.placeholder(tf.float32, (None, 0, 0)))  # Just a flag
     cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias
 
   if beam_size > 1:  # Beam Search
@@ -629,7 +649,9 @@ def body(self, features):
     encoder_input = tf.nn.dropout(encoder_input,
                                   1.0 - hparams.layer_prepostprocess_dropout)
     encoder_output = transformer_encoder(
-        encoder_input, encoder_self_attention_bias, hparams,
+        encoder_input,
+        encoder_self_attention_bias,
+        hparams,
         nonpadding=features_to_nonpadding(features, "inputs"))
     encoder_output = tf.expand_dims(encoder_output, 2)
 
@@ -669,8 +691,8 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
     encoder_self_attention_bias = common_attention.attention_bias_same_segment(
         inputs_segmentation, inputs_segmentation)
     encoder_decoder_attention_bias = (
-        common_attention.attention_bias_same_segment(
-            targets_segmentation, inputs_segmentation))
+        common_attention.attention_bias_same_segment(targets_segmentation,
+                                                     inputs_segmentation))
   else:
     # Usual case - not a packed dataset.
     encoder_padding = common_attention.embedding_to_padding(encoder_input)
@@ -688,8 +710,8 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
       32,
       ishape_static[-1],
       name="target_space_embedding",
-      dtype=tf.bfloat16 if hparams.activation_dtype == "bfloat16"
-      else tf.float32)
+      dtype=tf.bfloat16
+      if hparams.activation_dtype == "bfloat16" else tf.float32)
   emb_target_space = tf.reshape(emb_target_space, [1, 1, -1])
   encoder_input += emb_target_space
   if hparams.pos == "timing":
@@ -798,8 +820,7 @@ def transformer_encoder(encoder_input,
     pad_remover = None
     if hparams.use_pad_remover and not common_layers.is_on_tpu():
       pad_remover = expert_utils.PadRemover(padding)
-    for layer in range(hparams.num_encoder_layers or
-                       hparams.num_hidden_layers):
+    for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):
       with tf.variable_scope("layer_%d" % layer):
         with tf.variable_scope("self_attention"):
           y = common_attention.multihead_attention(
@@ -820,8 +841,11 @@ def transformer_encoder(encoder_input,
           x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(
-              common_layers.layer_preprocess(x, hparams), hparams, pad_remover,
-              conv_padding="SAME", nonpadding_mask=nonpadding,
+              common_layers.layer_preprocess(x, hparams),
+              hparams,
+              pad_remover,
+              conv_padding="SAME",
+              nonpadding_mask=nonpadding,
               losses=losses)
           x = common_layers.layer_postprocess(x, y, hparams)
     # if normalization is done in layer_preprocess, then it should also be done
@@ -873,8 +897,7 @@ def transformer_decoder(decoder_input,
       common_layers.comma_separated_string_to_integer_list(
           getattr(hparams, "attention_dropout_broadcast_dims", "")))
   with tf.variable_scope(name):
-    for layer in range(hparams.num_decoder_layers or
-                       hparams.num_hidden_layers):
+    for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
       layer_name = "layer_%d" % layer
       layer_cache = cache[layer_name] if cache is not None else None
       with tf.variable_scope(layer_name):
@@ -1013,13 +1036,16 @@ def transformer_ffn_layer(x,
   elif ffn_layer == "sru":
     return common_layers.sru(x)
   elif ffn_layer == "local_moe_tpu":
-    overhead = (hparams.moe_overhead_train
-                if hparams.mode == tf.estimator.ModeKeys.TRAIN
-                else hparams.moe_overhead_eval)
+    overhead = (
+        hparams.moe_overhead_train
+        if hparams.mode == tf.estimator.ModeKeys.TRAIN else
+        hparams.moe_overhead_eval)
     ret, loss = expert_utils.local_moe_tpu(
-        x, hparams.filter_size // 2,
+        x,
+        hparams.filter_size // 2,
         hparams.hidden_size,
-        hparams.moe_num_experts, overhead=overhead,
+        hparams.moe_num_experts,
+        overhead=overhead,
         loss_coef=hparams.moe_loss_coef)
     if losses is None:
       raise ValueError(
@@ -1602,7 +1628,7 @@ def transformer_clean_big_tpu():
 def transformer_tpu_with_conv():
   """Cut down on the number of heads, and use convs instead."""
   hparams = transformer_tpu()
-  hparams.num_heads = 4   # Heads are expensive on TPUs.
+  hparams.num_heads = 4  # Heads are expensive on TPUs.
   hparams.ffn_layer = "conv_relu_conv"
   return hparams
 
@@ -1612,7 +1638,7 @@ def transformer_lm_tpu_0():
   """HParams for training languagemodel_lm1b8k on tpu.  92M Params."""
   hparams = transformer_clean_big()
   update_hparams_for_tpu(hparams)
-  hparams.num_heads = 4   # Heads are expensive on TPUs.
+  hparams.num_heads = 4  # Heads are expensive on TPUs.
   hparams.batch_size = 4096
   hparams.shared_embedding_and_softmax_weights = False
   hparams.layer_prepostprocess_dropout = 0.1
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 98b98df16..584b01623 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -21,8 +21,6 @@
 import os
 import time
 
-# Dependency imports
-
 import numpy as np
 import six
 
@@ -73,8 +71,8 @@ def log_decode_results(inputs,
   is_image = "image" in problem_name
   decoded_inputs = None
   if is_image and save_images:
-    save_path = os.path.join(model_dir, "%s_prediction_%d.jpg" %
-                             (problem_name, prediction_idx))
+    save_path = os.path.join(
+        model_dir, "%s_prediction_%d.jpg" % (problem_name, prediction_idx))
     show_and_save_image(inputs / 255., save_path)
   elif inputs_vocab:
     if identity_output:
@@ -139,8 +137,7 @@ def decode_from_dataset(estimator,
       decode_filename = decode_to_file + ("%.2d" % decode_hp.shard_id)
     else:
       decode_filename = decode_to_file
-    output_filepath = _decode_filename(decode_filename, problem_name,
-                                       decode_hp)
+    output_filepath = _decode_filename(decode_filename, problem_name, decode_hp)
     parts = output_filepath.split(".")
     parts[-1] = "targets"
     target_filepath = ".".join(parts)
@@ -211,8 +208,7 @@ def decode_from_dataset(estimator,
         beam_score_str = ""
         if decode_hp.write_beam_scores:
           beam_score_str = "\t%.2f" % decoded_scores[i]
-        output_file.write(
-            str(d_output) + beam_score_str + decode_hp.delimiter)
+        output_file.write(str(d_output) + beam_score_str + decode_hp.delimiter)
         target_file.write(str(d_target) + decode_hp.delimiter)
         input_file.write(str(d_input) + decode_hp.delimiter)
 
@@ -254,9 +250,9 @@ def decode_from_file(estimator,
   num_decode_batches = (len(sorted_inputs) - 1) // decode_hp.batch_size + 1
 
   def input_fn():
-    input_gen = _decode_batch_input_fn(
-        num_decode_batches, sorted_inputs, inputs_vocab,
-        decode_hp.batch_size, decode_hp.max_input_size)
+    input_gen = _decode_batch_input_fn(num_decode_batches, sorted_inputs,
+                                       inputs_vocab, decode_hp.batch_size,
+                                       decode_hp.max_input_size)
     gen_fn = make_input_fn_from_generator(input_gen)
     example = gen_fn()
     return _decode_input_tensor_to_features_dict(example, hparams)
@@ -289,29 +285,39 @@ def timer(gen):
       for k, beam in enumerate(output_beams):
         tf.logging.info("BEAM %d:" % k)
         score = scores and scores[k]
-        _, decoded_outputs, _ = log_decode_results(result["inputs"], beam,
-                                                   problem_name, None,
-                                                   inputs_vocab, targets_vocab,
-                                                   log_results=decode_hp.log_results)
+        _, decoded_outputs, _ = log_decode_results(
+            result["inputs"],
+            beam,
+            problem_name,
+            None,
+            inputs_vocab,
+            targets_vocab,
+            log_results=decode_hp.log_results)
         beam_decodes.append(decoded_outputs)
         if decode_hp.write_beam_scores:
           beam_scores.append(score)
       if decode_hp.write_beam_scores:
-        decodes.append("\t".join(
-            ["\t".join([d, "%.2f" % s]) for d, s
-             in zip(beam_decodes, beam_scores)]))
+        decodes.append("\t".join([
+            "\t".join([d, "%.2f" % s])
+            for d, s in zip(beam_decodes, beam_scores)
+        ]))
       else:
         decodes.append("\t".join(beam_decodes))
     else:
       _, decoded_outputs, _ = log_decode_results(
-          result["inputs"], result["outputs"], problem_name,
-          None, inputs_vocab, targets_vocab,
+          result["inputs"],
+          result["outputs"],
+          problem_name,
+          None,
+          inputs_vocab,
+          targets_vocab,
           log_results=decode_hp.log_results)
       decodes.append(decoded_outputs)
     total_time_per_step += elapsed_time
     total_cnt += result["outputs"].shape[-1]
   tf.logging.info("Elapsed Time: %5.5f" % (time.time() - start_time))
-  tf.logging.info("Averaged Single Token Generation Time: %5.7f" % (total_time_per_step / total_cnt))
+  tf.logging.info("Averaged Single Token Generation Time: %5.7f" %
+                  (total_time_per_step / total_cnt))
 
   # Reversing the decoded inputs and outputs because they were reversed in
   # _decode_batch_input_fn
@@ -400,8 +406,8 @@ def input_fn():
             targets_vocab.decode(_save_until_eos(result["outputs"], is_image)))
 
 
-def _decode_batch_input_fn(num_decode_batches, sorted_inputs,
-                           vocabulary, batch_size, max_input_size):
+def _decode_batch_input_fn(num_decode_batches, sorted_inputs, vocabulary,
+                           batch_size, max_input_size):
   """Generator to produce batches of inputs."""
   tf.logging.info(" batch %d" % num_decode_batches)
   # First reverse all the input sentences so that if you're going to get OOMs,
@@ -517,8 +523,9 @@ def show_and_save_image(img, save_path):
   try:
     import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top
   except ImportError as e:
-    tf.logging.warning("Showing and saving an image requires matplotlib to be "
-                       "installed: %s", e)
+    tf.logging.warning(
+        "Showing and saving an image requires matplotlib to be "
+        "installed: %s", e)
     raise NotImplementedError("Image display and save not implemented.")
   plt.imshow(img)
   plt.savefig(save_path)

From 71193b25e78f4650296d77fbda4b362a659f088b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 8 Jun 2018 13:00:52 -0700
Subject: [PATCH 0068/2720] Toy timeseries problem that derives from timeseries
 base class.

PiperOrigin-RevId: 199837579
---
 tensor2tensor/data_generators/all_problems.py |   2 +-
 tensor2tensor/data_generators/timeseries.py   | 117 +++++++++++++-----
 .../data_generators/timeseries_test.py        |  10 +-
 3 files changed, 93 insertions(+), 36 deletions(-)

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index caaf61069..ee2e91796 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -20,7 +20,6 @@
 import importlib
 import six
 
-
 modules = [
     "tensor2tensor.data_generators.algorithmic",
     "tensor2tensor.data_generators.algorithmic_math",
@@ -50,6 +49,7 @@
     "tensor2tensor.data_generators.snli",
     "tensor2tensor.data_generators.squad",
     "tensor2tensor.data_generators.subject_verb_agreement",
+    "tensor2tensor.data_generators.timeseries",
     "tensor2tensor.data_generators.translate_encs",
     "tensor2tensor.data_generators.translate_ende",
     "tensor2tensor.data_generators.translate_enet",
diff --git a/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index 6c670b45c..c36da4574 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -17,7 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
-import random
+import numpy as np
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -27,22 +27,8 @@
 import tensorflow as tf
 
 
-@registry.register_problem
-class TimeseriesToyProblem(problem.Problem):
-  """Base Problem for multi timeseries for datasets."""
-
-  def __init__(self,
-               was_reversed=False,
-               was_copy=False,
-               num_series=2,
-               num_train_shards=9,
-               num_eval_shards=1,
-               num_samples=100):
-    super(TimeseriesToyProblem, self).__init__(was_reversed, was_copy)
-    self._num_train_shards = num_train_shards
-    self._num_eval_shards = num_eval_shards
-    self._num_samples = num_samples
-    self._num_series = num_series
+class TimeseriesProblem(problem.Problem):
+  """Base Problem for multi timeseries datasets."""
 
   def feature_encoders(self, data_dir):
     del data_dir
@@ -62,12 +48,41 @@ def dataset_splits(self):
     # 10% evaluation data
     return [{
         "split": problem.DatasetSplit.TRAIN,
-        "shards": self._num_train_shards,
+        "shards": self.num_train_shards,
     }, {
         "split": problem.DatasetSplit.EVAL,
-        "shards": self._num_eval_shards,
+        "shards": self.num_eval_shards,
     }]
 
+  @property
+  def num_train_shards(self):
+    """Number of training shards."""
+    return 9
+
+  @property
+  def num_eval_shards(self):
+    """Number of eval shards."""
+    return 1
+
+  @property
+  def num_series(self):
+    """Number of timeseries."""
+    raise NotImplementedError()
+
+  @property
+  def num_input_timestamps(self):
+    """Number of timestamps to include in the input."""
+    raise NotImplementedError()
+
+  @property
+  def num_target_timestamps(self):
+    """Number of timestamps to include in the target."""
+    raise NotImplementedError()
+
+  def timeseries_dataset(self):
+    """Multi-timeseries data [ timestamps , self.num_series ] ."""
+    raise NotImplementedError()
+
   def eval_metrics(self):
     eval_metrics = [metrics.Metrics.RMSE]
     return eval_metrics
@@ -76,8 +91,10 @@ def preprocess_example(self, example, unused_mode, unused_hparams):
     # Time series are flat on disk, we un-flatten them back here.
     flat_inputs = example["inputs"]
     flat_targets = example["targets"]
-    example["inputs"] = tf.reshape(flat_inputs, [-1, self._num_series])
-    example["targets"] = tf.reshape(flat_targets, [-1, self._num_series])
+    example["inputs"] = tf.reshape(flat_inputs,
+                                   [self.num_input_timestamps, self.num_series])
+    example["targets"] = tf.reshape(
+        flat_targets, [self.num_target_timestamps, self.num_series])
     return example
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
@@ -85,14 +102,17 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
     del tmp_dir
     del dataset_split
 
-    series = [[float(i + n) for n in range(self._num_series)]
-              for i in range(10)]
-
-    # This generates _num_samples instances of each possible split of series;
-    # inputs & targets are of variable size.
-    for _ in range(self._num_samples):
-      split_index = random.randint(1, 9)
-      inputs, targets = series[:split_index], series[split_index:]
+    series = self.timeseries_dataset()
+    num_timestamps = len(series)
+
+    # Generate samples with num_input_timestamps for "inputs" and
+    # num_target_timestamps in the "targets".
+    for split_index in xrange(self.num_input_timestamps,
+                              num_timestamps - self.num_target_timestamps + 1):
+      inputs = series[split_index -
+                      self.num_input_timestamps:split_index, :].tolist()
+      targets = series[split_index:split_index +
+                       self.num_target_timestamps, :].tolist()
       # We need to flatten the lists on disk for tf,Example to work.
       flat_inputs = [item for sublist in inputs for item in sublist]
       flat_targets = [item for sublist in targets for item in sublist]
@@ -102,8 +122,8 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {"inputs": (registry.Modalities.REAL, self._num_series)}
-    p.target_modality = (registry.Modalities.REAL, self._num_series)
+    p.input_modality = {"inputs": (registry.Modalities.REAL, self.num_series)}
+    p.target_modality = (registry.Modalities.REAL, self.num_series)
     p.input_space_id = problem.SpaceID.REAL
     p.target_space_id = problem.SpaceID.REAL
 
@@ -140,3 +160,38 @@ def example_reading_spec(self):
     }
     data_items_to_decoders = None
     return (data_fields, data_items_to_decoders)
+
+
+@registry.register_problem
+class TimeseriesToyProblem(TimeseriesProblem):
+  """Timeseries problem with a toy dataset."""
+
+  @property
+  def num_train_shards(self):
+    """Number of training shards."""
+    return 1
+
+  @property
+  def num_eval_shards(self):
+    """Number of eval shards."""
+    return 1
+
+  @property
+  def num_series(self):
+    """Number of timeseries."""
+    return 2
+
+  @property
+  def num_input_timestamps(self):
+    """Number of timestamps to include in the input."""
+    return 2
+
+  @property
+  def num_target_timestamps(self):
+    """Number of timestamps to include in the target."""
+    return 2
+
+  def timeseries_dataset(self):
+    series = [[float(i + n) for n in range(self.num_series)] for i in range(10)]
+
+    return np.array(series)
diff --git a/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index a45c80895..9daabe80b 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -35,8 +35,7 @@ def setUpClass(cls):
     os.mkdir(cls.tmp_dir)
 
   def testTimeseriesToyProblem(self):
-    problem = timeseries.TimeseriesToyProblem(
-        num_train_shards=1, num_eval_shards=1, num_samples=4)
+    problem = timeseries.TimeseriesToyProblem()
     problem.generate_data(self.tmp_dir, self.tmp_dir)
 
     dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, self.tmp_dir)
@@ -47,16 +46,19 @@ def testTimeseriesToyProblem(self):
     with self.test_session() as sess:
       examples.append(sess.run(features))
       examples.append(sess.run(features))
+      examples.append(sess.run(features))
+      examples.append(sess.run(features))
+
       try:
         sess.run(features)
       except tf.errors.OutOfRangeError:
         exhausted = True
 
     self.assertTrue(exhausted)
-    self.assertEqual(2, len(examples))
+    self.assertEqual(4, len(examples))
 
     self.assertNotEqual(
-        list(examples[0]["inputs"]), list(examples[1]["inputs"]))
+        list(examples[0]["inputs"][0]), list(examples[1]["inputs"][0]))
 
 
 if __name__ == "__main__":

From 23bd23b9830059fbc349381b70d9429b5c40a139 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 8 Jun 2018 14:00:30 -0700
Subject: [PATCH 0069/2720] Correct some problems introduced by recent fast
 decoding change.

PiperOrigin-RevId: 199846706
---
 tensor2tensor/layers/common_attention.py | 6 +++---
 tensor2tensor/models/transformer.py      | 3 +--
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 1967fa2e6..535e84dd6 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1193,8 +1193,7 @@ def grouped_attention_multihead(query_antecedent,
                      "attention heads (%d)." % (total_value_depth, num_heads))
   depth_v = total_value_depth // num_heads
   with tf.variable_scope(
-      name,
-      default_name="multihead_attention_sparse",
+      name, default_name="multihead_attention_sparse",
       values=[query_antecedent, memory_antecedent]):
     q = common_layers.dense(
         query_antecedent, total_key_depth, use_bias=False, name="q_transform")
@@ -2774,7 +2773,8 @@ def multihead_attention(query_antecedent,
   if total_value_depth % num_heads != 0:
     raise ValueError("Value depth (%d) must be divisible by the number of "
                      "attention heads (%d)." % (total_value_depth, num_heads))
-  with tf.variable_scope(name, values=[query_antecedent, memory_antecedent]):
+  with tf.variable_scope(name, default_name="multihead_attention",
+                         values=[query_antecedent, memory_antecedent]):
 
     if cache is None or memory_antecedent is None:
       q, k, v = compute_qkv(query_antecedent, memory_antecedent,
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 8c029a675..b3928a4ea 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -515,8 +515,7 @@ def fast_decode(encoder_output,
       cache[layer_name]["k_encdec"] = k_encdec
       cache[layer_name]["v_encdec"] = v_encdec
 
-    cache["encoder_output"] = tf.zeros_like(
-        tf.placeholder(tf.float32, (None, 0, 0)))  # Just a flag
+    cache["encoder_output"] = encoder_output
     cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias
 
   if beam_size > 1:  # Beam Search

From 8bdb897ae025462d52970ef5b7ee43973525c973 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 8 Jun 2018 15:47:25 -0700
Subject: [PATCH 0070/2720] xrange->range

PiperOrigin-RevId: 199863241
---
 tensor2tensor/data_generators/timeseries.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index c36da4574..fd314723a 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -107,8 +107,8 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
 
     # Generate samples with num_input_timestamps for "inputs" and
     # num_target_timestamps in the "targets".
-    for split_index in xrange(self.num_input_timestamps,
-                              num_timestamps - self.num_target_timestamps + 1):
+    for split_index in range(self.num_input_timestamps,
+                             num_timestamps - self.num_target_timestamps + 1):
       inputs = series[split_index -
                       self.num_input_timestamps:split_index, :].tolist()
       targets = series[split_index:split_index +

From f34e33fb49d53051d891fa68c3fa55c264db6b64 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 8 Jun 2018 17:51:50 -0700
Subject: [PATCH 0071/2720] Following changes:

PiperOrigin-RevId: 199877255
---
 tensor2tensor/data_generators/gym_problems.py |   3 +-
 tensor2tensor/layers/modalities.py            |   2 -
 tensor2tensor/models/research/next_frame.py   | 217 ++++++++++--------
 3 files changed, 124 insertions(+), 98 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 722acaae5..9368a7d5e 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -147,7 +147,8 @@ def hparams(self, defaults, unused_model_hparams):
     }
     p.target_modality = {
         "targets": ("video", 256),
-        "target_reward": ("symbol:weights_all", self.num_rewards)
+        "target_reward": ("symbol:weights_all", self.num_rewards),
+        "target_action": ("symbol:weights_all", self.num_actions)
     }
     p.input_space_id = problem.SpaceID.IMAGE
     p.target_space_id = problem.SpaceID.IMAGE
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index c1c307651..797340123 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -561,8 +561,6 @@ def top(self, body_output, _):
 
   def loss(self, top_out, targets):
     top_out = tf.squeeze(top_out, axis=[-1])
-    assert(top_out.shape.as_list() == targets.shape.as_list()), \
-           "The dimensions doesn't match."
     loss = tf.square(top_out - tf.to_float(targets))
     return tf.reduce_sum(loss), tf.reduce_sum(tf.ones_like(loss))
 
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index fbec036c5..aabc6ea54 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -176,68 +176,75 @@ def construct_latent_tower(self, images):
     """
     sequence_length = len(images)
 
-    with slim.arg_scope([slim.conv2d], reuse=False):
-      stacked_images = tf.concat(images, 3)
-      stacked_images = common_layers.make_even_size(stacked_images)
-      latent_enc1 = slim.conv2d(
-          stacked_images,
-          32, [3, 3],
-          stride=2,
-          scope="latent_conv1")
-      latent_enc1 = slim.batch_norm(latent_enc1, scope="latent_bn1")
-      latent_enc1 = common_layers.make_even_size(latent_enc1)
-
-      latent_enc2 = slim.conv2d(
-          latent_enc1,
-          64, [3, 3],
-          stride=2,
-          scope="latent_conv2")
-      latent_enc2 = slim.batch_norm(latent_enc2, scope="latent_bn2")
-
-      latent_enc3 = slim.conv2d(
-          latent_enc2,
-          64, [3, 3],
-          stride=1,
-          scope="latent_conv3")
-      latent_enc3 = slim.batch_norm(latent_enc3, scope="latent_bn3")
-
-      latent_enc3 = common_layers.make_even_size(latent_enc3)
-      latent_mean = slim.conv2d(
-          latent_enc3,
-          self.hparams.latent_channels, [3, 3],
-          stride=2,
-          activation_fn=None,
-          scope="latent_mean")
-
-      latent_std = slim.conv2d(
-          latent_enc3,
-          self.hparams.latent_channels, [3, 3],
-          stride=2,
-          scope="latent_std")
-
-      latent_std += self.hparams.latent_std_min
+    with slim.arg_scope([slim.conv2d, slim.batch_norm], reuse=False):
+      images = tf.concat(images, 3)
+
+      x = images
+      x = common_layers.make_even_size(x)
+      x = slim.conv2d(x, 32, [3, 3], stride=2, scope="latent_conv1")
+      x = slim.batch_norm(x, scope="latent_bn1")
+      x = common_layers.make_even_size(x)
+      x = slim.conv2d(x, 64, [3, 3], stride=2, scope="latent_conv2")
+      x = slim.batch_norm(x, scope="latent_bn2")
+      x = slim.conv2d(x, 64, [3, 3], stride=1, scope="latent_conv3")
+      x = slim.batch_norm(x, scope="latent_bn3")
+
+      nc = self.hparams.latent_channels
+      mean = slim.conv2d(
+          x, nc, [3, 3], stride=2, activation_fn=None, scope="latent_mean")
+      std = slim.conv2d(x, nc, [3, 3], stride=2, scope="latent_std")
+      std += self.hparams.latent_std_min
 
     if self.hparams.multi_latent:
       # timestep x batch_size x latent_size
       samples = tf.random_normal(
-          [sequence_length-1] + latent_mean.shape, 0, 1,
+          [sequence_length-1] + mean.shape, 0, 1,
           dtype=tf.float32)
     else:
       # batch_size x latent_size
-      samples = tf.random_normal(tf.shape(latent_mean), 0, 1, dtype=tf.float32)
+      samples = tf.random_normal(tf.shape(mean), 0, 1, dtype=tf.float32)
 
     if self.hparams.mode == tf.estimator.ModeKeys.TRAIN:
-      return latent_mean, latent_std, samples
+      return mean, std, samples
     else:
       # No latent tower at inference time, just standard gaussian.
       return None, None, samples
 
+  def reward_prediction(self, inputs, reuse):
+    """Builds a reward prediction network."""
+    with slim.arg_scope([slim.conv2d, slim.batch_norm], reuse=reuse):
+      x = inputs
+      x = slim.conv2d(x, 32, [3, 3], scope="reward_conv1")
+      x = slim.batch_norm(x, scope="reward_bn1")
+      x = slim.conv2d(x, 16, [3, 3], scope="reward_conv2")
+      x = slim.batch_norm(x, scope="reward_bn2")
+      x = slim.conv2d(x, 1, [3, 3], scope="reward_conv3", activation_fn=None)
+    return x
+
+  def encode_to_shape(self, inputs, shape, reuse):
+    """Encode the given tensor to given image shape."""
+    with slim.arg_scope([slim.fully_connected], reuse=reuse):
+      w, h = shape[1].value, shape[2].value
+      x = inputs
+      x = tf.contrib.layers.flatten(x)
+      x = slim.fully_connected(x, w * h, scope="encoding_full")
+      x = tf.reshape(x, (-1, w, h, 1))
+      return x
+
+  def decode_to_shape(self, inputs, shape, reuse):
+    """Encode the given tensor to given image shape."""
+    with slim.arg_scope([slim.fully_connected], reuse=reuse):
+      x = inputs
+      x = tf.contrib.layers.flatten(x)
+      x = slim.fully_connected(x, shape[2].value, scope="decoding_full")
+      x = tf.expand_dims(x, axis=1)
+      return x
+
   def construct_model(self,
                       images,
                       actions,
-                      states,
+                      rewards,
                       k=-1,
-                      use_state=False,
                       num_masks=10,
                       cdna=True,
                       dna=False,
@@ -245,11 +252,13 @@ def construct_model(self,
     """Build convolutional lstm video predictor using CDNA, or DNA.
 
     Args:
-      images: tensor of ground truth image sequences
-      actions: tensor of action sequences
-      states: tensor of ground truth state sequences
+      images: list of tensors of ground truth image sequences
+              there should be a 4D image ?xWxHxC for each timestep
+      actions: list of action tensors
+               each action should be in the shape ?x1xZ
+      rewards: list of reward tensors
+               each reward should be in the shape ?x1xZ
       k: constant used for scheduled sampling. -1 to feed in own prediction.
-      use_state: True to include state and action in prediction
       num_masks: the number of different pixel motion predictions (and
                  the number of masks for each of those predictions)
       cdna: True to use Convoluational Dynamic Neural Advection (CDNA)
@@ -258,7 +267,9 @@ def construct_model(self,
                       feeding in own predictions
     Returns:
       gen_images: predicted future image frames
-      gen_states: predicted future states
+      gen_rewards: predicted future rewards
+      latent_mean: mean of approximated posterior
+      latent_std: std of approximated posterior
 
     Raises:
       ValueError: if more than one network option specified or more than 1 mask
@@ -273,12 +284,11 @@ def construct_model(self,
       raise ValueError("More than one, or no network option specified.")
 
     img_height, img_width, color_channels = self.hparams.problem.frame_shape
-    batch_size = common_layers.shape_list(images[0])[0]
+    batch_size = self.hparams.batch_size
     lstm_func = self.basic_conv_lstm_cell
 
-    # Generated robot states and images.
-    gen_states, gen_images = [], []
-    current_state = states[0]
+    # Predicted images and rewards.
+    gen_rewards, gen_images = [], []
 
     if k == -1:
       feedself = True
@@ -303,12 +313,12 @@ def construct_model(self,
       latent_mean, latent_std, samples = latent_tower_outputs
 
     # Main tower
-    timestep = 0
     layer_norm = tf.contrib.layers.layer_norm
 
-    for image, action in zip(images[:-1], actions[:-1]):
+    for timestep, image, action, reward in zip(
+        range(len(images)-1), images[:-1], actions[:-1], rewards[:-1]):
       # Reuse variables after the first timestep.
-      reuse = bool(gen_images)
+      reuse = timestep > 0
 
       done_warm_start = len(gen_images) > context_frames - 1
       with slim.arg_scope(
@@ -321,16 +331,17 @@ def construct_model(self,
         if feedself and done_warm_start:
           # Feed in generated image.
           prev_image = gen_images[-1]
+          prev_reward = gen_rewards[-1]
         elif done_warm_start:
           # Scheduled sampling
           prev_image = self.scheduled_sample(
-              image, gen_images[-1], self.hparams.batch_size, num_ground_truth)
+              image, gen_images[-1], batch_size, num_ground_truth)
+          prev_reward = self.scheduled_sample(
+              reward, gen_rewards[-1], batch_size, num_ground_truth)
         else:
           # Always feed in ground_truth
           prev_image = image
-
-        # Predicted state is always fed back in
-        state_action = tf.concat(axis=1, values=[action, current_state])
+          prev_reward = reward
 
         prev_image = common_layers.make_even_size(prev_image)
         enc0 = slim.layers.conv2d(
@@ -361,15 +372,10 @@ def construct_model(self,
         enc2 = slim.layers.conv2d(
             hidden4, hidden4.get_shape()[3], [3, 3], stride=2, scope="conv3")
 
-        # Pass in state and action.
-        smear = tf.reshape(
-            state_action,
-            [-1, 1, 1, int(common_layers.shape_list(state_action)[1])])
-        enc2_shape = common_layers.shape_list(enc2)
-        smear = tf.tile(
-            smear, [1, enc2_shape[1], enc2_shape[2], 1])
-        if use_state:
-          enc2 = tf.concat(axis=3, values=[enc2, smear])
+        # Pass in reward and action.
+        emb_action = self.encode_to_shape(action, enc2.get_shape(), reuse)
+        emb_reward = self.encode_to_shape(prev_reward, enc2.get_shape(), reuse)
+        enc2 = tf.concat(axis=3, values=[enc2, emb_action, emb_reward])
 
         # Setup latent
         if self.hparams.stochastic_model:
@@ -468,15 +474,12 @@ def construct_model(self,
           output += layer * mask
         gen_images.append(output)
 
-        current_state = slim.layers.fully_connected(
-            state_action,
-            int(current_state.get_shape()[1]),
-            scope="state_pred",
-            activation_fn=None)
-        gen_states.append(current_state)
-        timestep += 1
+        p_reward = self.reward_prediction(hidden5, reuse)
+        p_reward = self.decode_to_shape(p_reward, reward.shape, reuse)
 
-    return gen_images, gen_states, latent_mean, latent_std
+        gen_rewards.append(p_reward)
+
+    return gen_images, gen_rewards, latent_mean, latent_std
 
   def cdna_transformation(self,
                           prev_image,
@@ -662,48 +665,72 @@ def basic_conv_lstm_cell(self,
 
       return new_h, tf.concat(axis=3, values=[new_c, new_h])
 
+  def get_input_if_exists(self, features, key, batch_size, num_frames):
+    if key in features:
+      x = features[key]
+    else:
+      x = tf.zeros((batch_size, num_frames, 1, self.hparams.hidden_size))
+    return tf.unstack(x, axis=1)
+
   def body(self, features):
     hparams = self.hparams
+    batch_size = self.hparams.batch_size
 
     # Split inputs and targets time-wise into a list of frames.
     input_frames = tf.unstack(features["inputs"], axis=1)
     target_frames = tf.unstack(features["targets"], axis=1)
 
-    num_frames = (hparams.video_num_input_frames +
-                  hparams.video_num_target_frames)
-    batch_size = common_layers.shape_list(input_frames)[0]
-    fake_zeros = [tf.zeros((batch_size, 1), dtype=tf.float32)
-                  for _ in range(num_frames)]
+    # Get actions if exist otherwise use zeros
+    input_actions = self.get_input_if_exists(
+        features, "input_action", batch_size, hparams.video_num_input_frames)
+    target_actions = self.get_input_if_exists(
+        features, "target_action", batch_size, hparams.video_num_target_frames)
+
+    # Get rewards if exist otherwise use zeros
+    input_rewards = self.get_input_if_exists(
+        features, "input_reward", batch_size, hparams.video_num_input_frames)
+    target_rewards = self.get_input_if_exists(
+        features, "target_reward", batch_size, hparams.video_num_target_frames)
+
+    all_actions = input_actions + target_actions
+    all_rewards = input_rewards + target_rewards
+
     is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
-    gen_images, _, latent_mean, latent_std = self.construct_model(
+    gen_images, gen_rewards, latent_mean, latent_std = self.construct_model(
         images=input_frames + target_frames,
-        actions=fake_zeros,
-        states=fake_zeros,
+        actions=all_actions,
+        rewards=all_rewards,
         k=900.0 if is_training else -1.0,
-        use_state=False,
         num_masks=10,
         cdna=True,
         dna=False,
         context_frames=hparams.video_num_input_frames)
 
-    kl_loss = 0.0
     step_num = tf.train.get_or_create_global_step()
     beta = tf.cond(step_num > self.hparams.num_iterations_2nd_stage,
                    lambda: self.hparams.latent_loss_multiplier,
                    lambda: 0.0)
 
+    kl_loss = 0.0
     if is_training:
+      kl_loss = self.kl_divergence(latent_mean, latent_std)
+
       tf.summary.scalar("beta", beta)
       tf.summary.histogram("posterior_mean", latent_mean)
       tf.summary.histogram("posterior_std", latent_std)
-
-    if is_training:
-      kl_loss = self.kl_divergence(latent_mean, latent_std)
       tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
-    kl_loss *= beta
+
+    extra_loss = beta * kl_loss
 
     predictions = gen_images[hparams.video_num_input_frames-1:]
-    return predictions, kl_loss
+    reward_pred = tf.stack(
+        gen_rewards[hparams.video_num_input_frames-1:], axis=1)
+
+    return_targets = predictions
+    if "target_reward" in features:
+      return_targets = {"targets": predictions, "target_reward": reward_pred}
+
+    return return_targets, extra_loss
 
 
 @registry.register_hparams
@@ -735,8 +762,8 @@ def next_frame():
 def next_frame_stochastic():
   """SV2P model."""
   hparams = next_frame()
-  hparams.video_num_input_frames = 1
-  hparams.video_num_target_frames = 4
+  hparams.video_num_input_frames = 4
+  hparams.video_num_target_frames = 1
   hparams.batch_size = 8
   hparams.target_modality = "video:raw"
   hparams.input_modalities = "inputs:video:raw"

From 80a157b5d304973ed5cf73d2391d60950448c158 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 8 Jun 2018 20:13:52 -0700
Subject: [PATCH 0072/2720] Cleanup all the Slim reuse.

PiperOrigin-RevId: 199886069
---
 tensor2tensor/models/research/next_frame.py | 38 ++++++++-------------
 1 file changed, 14 insertions(+), 24 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index aabc6ea54..f7ec393a2 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -176,7 +176,7 @@ def construct_latent_tower(self, images):
     """
     sequence_length = len(images)
 
-    with slim.arg_scope([slim.conv2d, slim.batch_norm], reuse=False):
+    with tf.variable_scope("latent"):
       images = tf.concat(images, 3)
 
       x = images
@@ -210,9 +210,9 @@ def construct_latent_tower(self, images):
       # No latent tower at inference time, just standard gaussian.
       return None, None, samples
 
-  def reward_prediction(self, inputs, reuse):
+  def reward_prediction(self, inputs):
     """Builds a reward prediction network."""
-    with slim.arg_scope([slim.conv2d, slim.batch_norm], reuse=reuse):
+    with tf.variable_scope("reward_pred", reuse=tf.AUTO_REUSE):
       x = inputs
       x = slim.conv2d(x, 32, [3, 3], scope="reward_conv1")
       x = slim.batch_norm(x, scope="reward_bn1")
@@ -221,9 +221,9 @@ def reward_prediction(self, inputs, reuse):
       x = slim.conv2d(x, 1, [3, 3], scope="reward_conv3", activation_fn=None)
     return x
 
-  def encode_to_shape(self, inputs, shape, reuse):
+  def encode_to_shape(self, inputs, shape):
     """Encode the given tensor to given image shape."""
-    with slim.arg_scope([slim.fully_connected], reuse=reuse):
+    with tf.variable_scope("reward_enc", reuse=tf.AUTO_REUSE):
       w, h = shape[1].value, shape[2].value
       x = inputs
       x = tf.contrib.layers.flatten(x)
@@ -231,9 +231,9 @@ def encode_to_shape(self, inputs, shape, reuse):
       x = tf.reshape(x, (-1, w, h, 1))
       return x
 
-  def decode_to_shape(self, inputs, shape, reuse):
+  def decode_to_shape(self, inputs, shape):
     """Encode the given tensor to given image shape."""
-    with slim.arg_scope([slim.fully_connected], reuse=reuse):
+    with tf.variable_scope("reward_dec", reuse=tf.AUTO_REUSE):
       x = inputs
       x = tf.contrib.layers.flatten(x)
       x = slim.fully_connected(x, shape[2].value, scope="decoding_full")
@@ -317,17 +317,9 @@ def construct_model(self,
 
     for timestep, image, action, reward in zip(
         range(len(images)-1), images[:-1], actions[:-1], rewards[:-1]):
-      # Reuse variables after the first timestep.
-      reuse = timestep > 0
 
       done_warm_start = len(gen_images) > context_frames - 1
-      with slim.arg_scope(
-          [
-              lstm_func, slim.layers.conv2d, slim.layers.fully_connected,
-              layer_norm, slim.layers.conv2d_transpose
-          ],
-          reuse=reuse):
-
+      with tf.variable_scope("main", reuse=tf.AUTO_REUSE):
         if feedself and done_warm_start:
           # Feed in generated image.
           prev_image = gen_images[-1]
@@ -373,8 +365,8 @@ def construct_model(self,
             hidden4, hidden4.get_shape()[3], [3, 3], stride=2, scope="conv3")
 
         # Pass in reward and action.
-        emb_action = self.encode_to_shape(action, enc2.get_shape(), reuse)
-        emb_reward = self.encode_to_shape(prev_reward, enc2.get_shape(), reuse)
+        emb_action = self.encode_to_shape(action, enc2.get_shape())
+        emb_reward = self.encode_to_shape(prev_reward, enc2.get_shape())
         enc2 = tf.concat(axis=3, values=[enc2, emb_action, emb_reward])
 
         # Setup latent
@@ -474,8 +466,8 @@ def construct_model(self,
           output += layer * mask
         gen_images.append(output)
 
-        p_reward = self.reward_prediction(hidden5, reuse)
-        p_reward = self.decode_to_shape(p_reward, reward.shape, reuse)
+        p_reward = self.reward_prediction(hidden5)
+        p_reward = self.decode_to_shape(p_reward, reward.shape)
 
         gen_rewards.append(p_reward)
 
@@ -620,8 +612,7 @@ def basic_conv_lstm_cell(self,
                            num_channels,
                            filter_size=5,
                            forget_bias=1.0,
-                           scope=None,
-                           reuse=None):
+                           scope=None):
     """Basic LSTM recurrent network cell, with 2D convolution connctions.
 
     We add forget_bias (default: 1) to the biases of the forget gate in order to
@@ -635,7 +626,6 @@ def basic_conv_lstm_cell(self,
       filter_size: the shape of the each convolution filter.
       forget_bias: the initial value of the forget biases.
       scope: Optional scope for variable_scope.
-      reuse: whether or not the layer and the variables should be reused.
     Returns:
        a tuple of tensors representing output and the new state.
     """
@@ -645,7 +635,7 @@ def basic_conv_lstm_cell(self,
     with tf.variable_scope(scope,
                            "BasicConvLstmCell",
                            [inputs, state],
-                           reuse=reuse):
+                           reuse=tf.AUTO_REUSE):
       inputs.get_shape().assert_has_rank(4)
       state.get_shape().assert_has_rank(4)
       c, h = tf.split(axis=3, num_or_size_splits=2, value=state)

From 06bf2a0fb13703c221115ce090dad104b8fd058a Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sat, 9 Jun 2018 17:48:24 -0700
Subject: [PATCH 0073/2720] Switching to tf.contrib.rnn.ConvLSTMCell

PiperOrigin-RevId: 199934867
---
 tensor2tensor/models/research/next_frame.py | 96 ++++++---------------
 1 file changed, 28 insertions(+), 68 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index f7ec393a2..417b8ec4a 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -240,6 +240,17 @@ def decode_to_shape(self, inputs, shape):
       x = tf.expand_dims(x, axis=1)
       return x
 
+  def conv_lstm_2d(self, inputs, state, output_channels,
+                   kernel_size=5, scope=None):
+    input_shape = common_layers.shape_list(inputs)
+    cell = tf.contrib.rnn.ConvLSTMCell(
+        2, input_shape[1:], output_channels,
+        [kernel_size, kernel_size], name=scope)
+    if state is None:
+      state = cell.zero_state(self.hparams.batch_size, tf.float32)
+    outputs, new_state = cell(inputs, state)
+    return outputs, new_state
+
   def construct_model(self,
                       images,
                       actions,
@@ -285,7 +296,6 @@ def construct_model(self,
 
     img_height, img_width, color_channels = self.hparams.problem.frame_shape
     batch_size = self.hparams.batch_size
-    lstm_func = self.basic_conv_lstm_cell
 
     # Predicted images and rewards.
     gen_rewards, gen_images = [], []
@@ -303,9 +313,8 @@ def construct_model(self,
       feedself = False
 
     # LSTM state sizes and states.
-    lstm_size = np.int32(np.array([32, 32, 64, 64, 128, 64, 32]))
-    lstm_state1, lstm_state2, lstm_state3, lstm_state4 = None, None, None, None
-    lstm_state5, lstm_state6, lstm_state7 = None, None, None
+    lstm_size = np.array([32, 32, 64, 64, 128, 64, 32], dtype=np.int32)
+    lstm_state = [None] * 7
 
     # Latent tower
     if self.hparams.stochastic_model:
@@ -314,6 +323,7 @@ def construct_model(self,
 
     # Main tower
     layer_norm = tf.contrib.layers.layer_norm
+    lstm_func = self.conv_lstm_2d
 
     for timestep, image, action, reward in zip(
         range(len(images)-1), images[:-1], actions[:-1], rewards[:-1]):
@@ -344,21 +354,21 @@ def construct_model(self,
             normalizer_fn=layer_norm,
             normalizer_params={"scope": "layer_norm1"})
 
-        hidden1, lstm_state1 = lstm_func(
-            enc0, lstm_state1, lstm_size[0], scope="state1")
+        hidden1, lstm_state[0] = lstm_func(
+            enc0, lstm_state[0], lstm_size[0], scope="state1")
         hidden1 = layer_norm(hidden1, scope="layer_norm2")
-        hidden2, lstm_state2 = lstm_func(
-            hidden1, lstm_state2, lstm_size[1], scope="state2")
+        hidden2, lstm_state[1] = lstm_func(
+            hidden1, lstm_state[1], lstm_size[1], scope="state2")
         hidden2 = layer_norm(hidden2, scope="layer_norm3")
         hidden2 = common_layers.make_even_size(hidden2)
         enc1 = slim.layers.conv2d(
             hidden2, hidden2.get_shape()[3], [3, 3], stride=2, scope="conv2")
 
-        hidden3, lstm_state3 = lstm_func(
-            enc1, lstm_state3, lstm_size[2], scope="state3")
+        hidden3, lstm_state[2] = lstm_func(
+            enc1, lstm_state[2], lstm_size[2], scope="state3")
         hidden3 = layer_norm(hidden3, scope="layer_norm4")
-        hidden4, lstm_state4 = lstm_func(
-            hidden3, lstm_state4, lstm_size[3], scope="state4")
+        hidden4, lstm_state[3] = lstm_func(
+            hidden3, lstm_state[3], lstm_size[3], scope="state4")
         hidden4 = layer_norm(hidden4, scope="layer_norm5")
         hidden4 = common_layers.make_even_size(hidden4)
         enc2 = slim.layers.conv2d(
@@ -383,16 +393,16 @@ def construct_model(self,
         enc3 = slim.layers.conv2d(
             enc2, hidden4.get_shape()[3], [1, 1], stride=1, scope="conv4")
 
-        hidden5, lstm_state5 = lstm_func(
-            enc3, lstm_state5, lstm_size[4], scope="state5")  # last 8x8
+        hidden5, lstm_state[4] = lstm_func(
+            enc3, lstm_state[4], lstm_size[4], scope="state5")  # last 8x8
         hidden5 = layer_norm(hidden5, scope="layer_norm6")
         enc4 = slim.layers.conv2d_transpose(
             hidden5, hidden5.get_shape()[3], 3, stride=2, scope="convt1")
 
         enc1_shape = common_layers.shape_list(enc1)
         enc4 = enc4[:, :enc1_shape[1], :enc1_shape[2], :]  # Cut to shape.
-        hidden6, lstm_state6 = lstm_func(
-            enc4, lstm_state6, lstm_size[5], scope="state6")  # 16x16
+        hidden6, lstm_state[5] = lstm_func(
+            enc4, lstm_state[5], lstm_size[5], scope="state6")  # 16x16
         hidden6 = layer_norm(hidden6, scope="layer_norm7")
         # Skip connection.
         hidden6 = tf.concat(axis=3, values=[hidden6, enc1])  # both 16x16
@@ -401,8 +411,8 @@ def construct_model(self,
             hidden6, hidden6.get_shape()[3], 3, stride=2, scope="convt2")
         enc0_shape = common_layers.shape_list(enc0)
         enc5 = enc5[:, :enc0_shape[1], :enc0_shape[2], :]  # Cut to shape.
-        hidden7, lstm_state7 = lstm_func(
-            enc5, lstm_state7, lstm_size[6], scope="state7")  # 32x32
+        hidden7, lstm_state[6] = lstm_func(
+            enc5, lstm_state[6], lstm_size[6], scope="state7")  # 32x32
         hidden7 = layer_norm(hidden7, scope="layer_norm8")
 
         # Skip connection.
@@ -605,56 +615,6 @@ def kl_divergence(self, mu, log_sigma):
         1. + log_sigma - tf.square(mu) - tf.exp(log_sigma),
         axis=1)
 
-  @slim.add_arg_scope
-  def basic_conv_lstm_cell(self,
-                           inputs,
-                           state,
-                           num_channels,
-                           filter_size=5,
-                           forget_bias=1.0,
-                           scope=None):
-    """Basic LSTM recurrent network cell, with 2D convolution connctions.
-
-    We add forget_bias (default: 1) to the biases of the forget gate in order to
-    reduce the scale of forgetting in the beginning of the training.
-    It does not allow cell clipping, a projection layer, and does not
-    use peep-hole connections: it is the basic baseline.
-    Args:
-      inputs: input Tensor, 4D, batch x height x width x channels.
-      state: state Tensor, 4D, batch x height x width x channels.
-      num_channels: the number of output channels in the layer.
-      filter_size: the shape of the each convolution filter.
-      forget_bias: the initial value of the forget biases.
-      scope: Optional scope for variable_scope.
-    Returns:
-       a tuple of tensors representing output and the new state.
-    """
-    if state is None:
-      inputs_shape = common_layers.shape_list(inputs)
-      state = tf.zeros(inputs_shape[:3] + [2 * num_channels])
-    with tf.variable_scope(scope,
-                           "BasicConvLstmCell",
-                           [inputs, state],
-                           reuse=tf.AUTO_REUSE):
-      inputs.get_shape().assert_has_rank(4)
-      state.get_shape().assert_has_rank(4)
-      c, h = tf.split(axis=3, num_or_size_splits=2, value=state)
-      inputs_h = tf.concat(axis=3, values=[inputs, h])
-      # Parameters of gates are concatenated into one conv for efficiency.
-      i_j_f_o = slim.layers.conv2d(inputs_h,
-                                   4 * num_channels, [filter_size, filter_size],
-                                   stride=1,
-                                   activation_fn=None,
-                                   scope="Gates")
-
-      # i = input_gate, j = new_input, f = forget_gate, o = output_gate
-      i, j, f, o = tf.split(axis=3, num_or_size_splits=4, value=i_j_f_o)
-
-      new_c = c * tf.sigmoid(f + forget_bias) + tf.sigmoid(i) * tf.tanh(j)
-      new_h = tf.tanh(new_c) * tf.sigmoid(o)
-
-      return new_h, tf.concat(axis=3, values=[new_c, new_h])
-
   def get_input_if_exists(self, features, key, batch_size, num_frames):
     if key in features:
       x = features[key]

From 941c724a35634a87e531fabb5634b4cd040f5061 Mon Sep 17 00:00:00 2001
From: Tomasz Latkowski <tlatkowski@gmail.com>
Date: Sun, 10 Jun 2018 20:11:43 +0200
Subject: [PATCH 0074/2720] add style transfer in text as new problem

---
 tensor2tensor/data_generators/all_problems.py |   1 +
 .../data_generators/style_transfer.py         | 154 ++++++++++++++++++
 2 files changed, 155 insertions(+)
 create mode 100644 tensor2tensor/data_generators/style_transfer.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index ee2e91796..f3c4750b5 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -47,6 +47,7 @@
     "tensor2tensor.data_generators.problem_hparams",
     "tensor2tensor.data_generators.ptb",
     "tensor2tensor.data_generators.snli",
+    "tensor2tensor.data_generators.style_transfer",
     "tensor2tensor.data_generators.squad",
     "tensor2tensor.data_generators.subject_verb_agreement",
     "tensor2tensor.data_generators.timeseries",
diff --git a/tensor2tensor/data_generators/style_transfer.py b/tensor2tensor/data_generators/style_transfer.py
new file mode 100644
index 000000000..b7a6bf35b
--- /dev/null
+++ b/tensor2tensor/data_generators/style_transfer.py
@@ -0,0 +1,154 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 classes for text-based language style transfer problems.
+
+* StyleTransferProblem: abstract class for style transfer problems.
+* StyleTransferShakespeare: specific problem implementation that enriches language with Shakespeare-like style.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import tarfile
+
+import tensorflow as tf
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+
+logger = tf.logging
+
+"""
+Modern-Shakespeare corpus is consisted of:
+- 18,395 parallel sentences for training (train set),
+- 1,218 parallel sentences for evaluation (dev set),
+- 1,462 parallel sentence for testing (test set).
+"""
+
+_SHAKESPEARE_MODERN_TRAIN_DATASET = [[
+    "https://github.com/tlatkowski/st/raw/master/shakespeare.train.tgz",
+    ("train.original", "train.modern")
+]]
+
+_SHAKESPEARE_MODERN_DEV_DATASET = [[
+    "https://github.com/tlatkowski/st/raw/master/shakespeare.dev.tgz",
+    ("dev.original", "dev.modern")
+]]
+
+_TRAIN_SHARDS = 1
+_DEV_SHARDS = 1
+_SUBWORD_VOCAB_SIZE = 8000
+
+
+class StyleTransferProblem(text_problems.Text2TextProblem):
+  """Base class for transfering styles problems"""
+
+  @property
+  def target(self):
+    raise NotImplementedError()
+
+  @property
+  def source(self):
+    raise NotImplementedError()
+
+  def dataset_url(self, dataset_split):
+    raise NotImplementedError()
+
+  def vocab_data_files(self):
+    """Files to be passed to get_or_generate_vocab."""
+    return self.dataset_url(problem.DatasetSplit.TRAIN)
+
+  @property
+  def approx_vocab_size(self):
+    return _SUBWORD_VOCAB_SIZE
+
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each."""
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": _TRAIN_SHARDS,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": _DEV_SHARDS,
+    }]
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    dataset = self.dataset_url(dataset_split)
+
+    tag = "train" if dataset_split == problem.DatasetSplit.TRAIN else "dev"
+
+    url = dataset[0][0]
+    compressed_filename = os.path.basename(url)
+    compressed_filepath = os.path.join(tmp_dir, compressed_filename)
+    generator_utils.maybe_download(tmp_dir, compressed_filename, url)
+
+    mode = "r:gz" if compressed_filepath.endswith("gz") else "r"
+    with tarfile.open(compressed_filepath, mode) as corpus_tar:
+      corpus_tar.extractall(tmp_dir)
+
+    if self.vocab_type == text_problems.VocabType.SUBWORD:
+      generator_utils.get_or_generate_vocab(
+          data_dir, tmp_dir, self.vocab_filename, self.approx_vocab_size,
+          self.vocab_data_files())
+
+    source_file = os.path.join(tmp_dir, tag + ".modern")
+    target_file = os.path.join(tmp_dir, tag + ".original")
+    return text_problems.text2text_txt_iterator(source_file,
+                                                target_file)
+
+
+@registry.register_problem
+class StyleTransferShakespeareToModern(StyleTransferProblem):
+  """Transfering style from Shakespeare original English to modern one"""
+
+  @property
+  def target(self):
+    return '.modern'
+
+  @property
+  def source(self):
+    return '.original'
+
+  def dataset_url(self, dataset_split):
+    train = dataset_split == problem.DatasetSplit.TRAIN
+    return _SHAKESPEARE_MODERN_TRAIN_DATASET if train else _SHAKESPEARE_MODERN_DEV_DATASET
+
+
+@registry.register_problem
+class StyleTransferModernToShakespeare(StyleTransferProblem):
+  """Transfering style from modern English to Shakespeare original English"""
+
+  @property
+  def target(self):
+    return '.original'
+
+  @property
+  def source(self):
+    return '.modern'
+
+  def dataset_url(self, dataset_split):
+    train = dataset_split == problem.DatasetSplit.TRAIN
+    return _SHAKESPEARE_MODERN_TRAIN_DATASET if train else _SHAKESPEARE_MODERN_DEV_DATASET

From da8d5756333baec70441a2392ce81f82f58201e7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 11 Jun 2018 05:28:47 -0700
Subject: [PATCH 0075/2720] move gaussian initialiser multiplication into
 standard deviation.

PiperOrigin-RevId: 200036469
---
 tensor2tensor/layers/common_attention.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 535e84dd6..3a1990b13 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -3047,13 +3047,13 @@ def parameter_attention(x,
     k = tf.get_variable(
         "k",
         var_shape_k,
-        initializer=tf.random_normal_initializer(0, output_depth**-0.5)) * (
-            num_heads**0.5)
+        initializer=tf.random_normal_initializer(
+            0, output_depth**-0.5 * (num_heads**0.5)))
     v = tf.get_variable(
         "v",
         var_shape_v,
-        initializer=tf.random_normal_initializer(0, output_depth**-0.5)) * (
-            output_depth**0.5)
+        initializer=tf.random_normal_initializer(
+            0, output_depth**-0.5 * (output_depth**0.5)))
     batch_size = common_layers.shape_list(x)[0]
     length = common_layers.shape_list(x)[1]
     q = common_layers.dense(

From a9481bc34563f2f039650aa493a31c49155e59de Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 11 Jun 2018 12:32:10 -0700
Subject: [PATCH 0076/2720] Add WaterWorld T2T Problem

PiperOrigin-RevId: 200094339
---
 setup.py                                      |  1 +
 tensor2tensor/data_generators/gym_problems.py | 22 +++++++++
 .../data_generators/gym_problems_test.py      | 47 +++++++++++++++++++
 3 files changed, 70 insertions(+)
 create mode 100644 tensor2tensor/data_generators/gym_problems_test.py

diff --git a/setup.py b/setup.py
index fed6e8c20..0de3d2425 100644
--- a/setup.py
+++ b/setup.py
@@ -39,6 +39,7 @@
         'google-api-python-client',
         'gunicorn',
         'gym',
+        'gym_ple',
         'h5py',
         'numpy',
         'requests',
diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 9368a7d5e..f53c5cb65 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -21,6 +21,7 @@
 import math
 import os
 import gym
+import gym_ple       # pylint: disable=unused-import
 import numpy as np
 
 from tensor2tensor.data_generators import gym_utils
@@ -167,6 +168,27 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
              "reward": [int(reward - self.min_reward)]}
 
 
+@registry.register_problem
+class GymWaterWorldRandom5k(GymDiscreteProblem):
+  """WaterWorld game, random actions."""
+
+  @property
+  def env_name(self):
+    return "WaterWorldFast-v0"
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+  @property
+  def num_steps(self):
+    return 5000
+
+
 @registry.register_problem
 class GymPongRandom5k(GymDiscreteProblem):
   """Pong game, random actions."""
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
new file mode 100644
index 000000000..a0cc7bcb3
--- /dev/null
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -0,0 +1,47 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Gym generators tests."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import shutil
+
+from tensor2tensor.data_generators import gym_problems
+
+import tensorflow as tf
+
+
+class GymProblemsTest(tf.test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    cls.tmp_dir = tf.test.get_temp_dir()
+    shutil.rmtree(cls.tmp_dir)
+    os.mkdir(cls.tmp_dir)
+
+  def testGymAtariBoots(self):
+    problem = gym_problems.GymPongRandom5k()
+    self.assertEqual(5000, problem.num_steps)
+
+  def testPyGamesBoots(self):
+    problem = gym_problems.GymWaterWorldRandom5k()
+    self.assertEqual(5000, problem.num_steps)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 4da9c8ada1f1452789581fc4a4cd69c5ce18c145 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 11 Jun 2018 14:26:52 -0700
Subject: [PATCH 0077/2720] Corrections to make Transformer train on timeseries
 somewhat.

PiperOrigin-RevId: 200112812
---
 tensor2tensor/data_generators/timeseries.py      | 14 +++++++++++---
 tensor2tensor/data_generators/timeseries_test.py |  2 +-
 tensor2tensor/models/transformer.py              |  8 ++++++++
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index fd314723a..38d23eea8 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -87,14 +87,22 @@ def eval_metrics(self):
     eval_metrics = [metrics.Metrics.RMSE]
     return eval_metrics
 
+  @property
+  def normalizing_constant(self):
+    """Constant by which all data will be multiplied to be more normalized."""
+    return 1.0  # Adjust so that your loss is around 1 or 10 or 100, not 1e+9.
+
   def preprocess_example(self, example, unused_mode, unused_hparams):
     # Time series are flat on disk, we un-flatten them back here.
     flat_inputs = example["inputs"]
     flat_targets = example["targets"]
-    example["inputs"] = tf.reshape(flat_inputs,
-                                   [self.num_input_timestamps, self.num_series])
+    c = self.normalizing_constant
+    # Tensor2Tensor models expect [height, width, depth] examples, here we
+    # use height for time and set width to 1 and num_series is our depth.
+    example["inputs"] = tf.reshape(
+        flat_inputs, [self.num_input_timestamps, 1, self.num_series]) * c
     example["targets"] = tf.reshape(
-        flat_targets, [self.num_target_timestamps, self.num_series])
+        flat_targets, [self.num_target_timestamps, 1, self.num_series]) * c
     return example
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
diff --git a/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index 9daabe80b..ce0c4f911 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -58,7 +58,7 @@ def testTimeseriesToyProblem(self):
     self.assertEqual(4, len(examples))
 
     self.assertNotEqual(
-        list(examples[0]["inputs"][0]), list(examples[1]["inputs"][0]))
+        list(examples[0]["inputs"][0, 0]), list(examples[1]["inputs"][0, 0]))
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index b3928a4ea..aa3bda483 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1491,6 +1491,14 @@ def transformer_relative_big():
   return hparams
 
 
+@registry.register_hparams
+def transformer_timeseries():
+  hparams = transformer_small()
+  hparams.batch_size = 256
+  hparams.learning_rate_warmup_steps = 2000
+  return hparams
+
+
 def update_hparams_for_tpu(hparams):
   """Change hparams to be compatible with TPU training."""
 

From cd469410f5e3c7c1dd28cf642942b86d7d63d82c Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Mon, 11 Jun 2018 15:02:37 -0700
Subject: [PATCH 0078/2720] fix a bug in expert_utils.py

PiperOrigin-RevId: 200119008
---
 tensor2tensor/utils/expert_utils.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 86747480b..a60f3fa72 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -1478,15 +1478,16 @@ def _step(source_replica, target_replica, x_split, op="plus_eq"):
             assert op == "copy"
             x_split[target_device][shard] = tf.identity(source)
     center = parallelism.n // 2
+
     # accumulate everything towards the center.
-    for i in range(center, parallelism.n - 1)[::-1]:
+    for i in reversed(range(center, parallelism.n - 1)):
       _step(i + 1, i, x_split, op="plus_eq")
     for i in range(center):
       _step(i, i + 1, x_split, op="plus_eq")
     # copy everything away from the center.
     for i in range(center, parallelism.n - 1):
       _step(i, i + 1, x_split, op="copy")
-    for i in range(center)[::-1]:
+    for i in reversed(range(center)):
       _step(i + 1, i, x_split, op="copy")
     x_concat = parallelism(tf.concat, x_split, 0)
     y = parallelism(common_layers.reshape_like_all_dims, x_concat, x)

From c886c1f9e501792faae0c68f5fc1e34db1f1096d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 11 Jun 2018 15:46:25 -0700
Subject: [PATCH 0079/2720] Remove change requiring gym_ple and ple.

PiperOrigin-RevId: 200126981
---
 setup.py                                      |  1 -
 tensor2tensor/data_generators/gym_problems.py | 22 -------------------
 .../data_generators/gym_problems_test.py      |  4 ----
 3 files changed, 27 deletions(-)

diff --git a/setup.py b/setup.py
index 0de3d2425..fed6e8c20 100644
--- a/setup.py
+++ b/setup.py
@@ -39,7 +39,6 @@
         'google-api-python-client',
         'gunicorn',
         'gym',
-        'gym_ple',
         'h5py',
         'numpy',
         'requests',
diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index f53c5cb65..9368a7d5e 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -21,7 +21,6 @@
 import math
 import os
 import gym
-import gym_ple       # pylint: disable=unused-import
 import numpy as np
 
 from tensor2tensor.data_generators import gym_utils
@@ -168,27 +167,6 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
              "reward": [int(reward - self.min_reward)]}
 
 
-@registry.register_problem
-class GymWaterWorldRandom5k(GymDiscreteProblem):
-  """WaterWorld game, random actions."""
-
-  @property
-  def env_name(self):
-    return "WaterWorldFast-v0"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-  @property
-  def num_steps(self):
-    return 5000
-
-
 @registry.register_problem
 class GymPongRandom5k(GymDiscreteProblem):
   """Pong game, random actions."""
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index a0cc7bcb3..1bbd2a374 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -38,10 +38,6 @@ def testGymAtariBoots(self):
     problem = gym_problems.GymPongRandom5k()
     self.assertEqual(5000, problem.num_steps)
 
-  def testPyGamesBoots(self):
-    problem = gym_problems.GymWaterWorldRandom5k()
-    self.assertEqual(5000, problem.num_steps)
-
 
 if __name__ == "__main__":
   tf.test.main()

From 987f5de44aa50170836837022c9de5f912644ab1 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 11 Jun 2018 18:00:19 -0700
Subject: [PATCH 0080/2720] Better error messages in padded cross entropy and
 shape correction in reward prediction in stochastic next frame model.

PiperOrigin-RevId: 200145072
---
 tensor2tensor/layers/common_layers.py       | 7 +++++--
 tensor2tensor/models/research/next_frame.py | 1 +
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index fab035723..b7df15c30 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1826,15 +1826,18 @@ def padded_cross_entropy(logits,
         weights_fn=weights_fn,
         reduce_sum=reduce_sum)
   confidence = 1.0 - label_smoothing
-  vocab_size = shape_list(logits)[-1]
+  logits_shape = shape_list(logits)
+  vocab_size = logits_shape[-1]
   with tf.name_scope("padded_cross_entropy", values=[logits, labels]):
-    if len(logits.get_shape().as_list()) == 2:
+    if len(logits_shape) == 2:
       # Deal with the case where we did not insert extra dimensions due to
       # TPU issues.  No pad-to-same-length happens in this case.
       # TODO(noam): remove this logic once TPU can handle extra dimensions.
       labels = tf.reshape(labels, [-1])
     else:
       logits, labels = pad_with_zeros(logits, labels)
+    logits = tf.reshape(logits, shape_list(labels) + [vocab_size],
+                        name="padded_cross_entropy_size_check")
     logits = tf.cast(logits, tf.float32)
     xent = smoothing_cross_entropy(logits, labels, vocab_size, confidence,
                                    gaussian=gaussian)
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 417b8ec4a..4d6edf457 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -675,6 +675,7 @@ def body(self, features):
     predictions = gen_images[hparams.video_num_input_frames-1:]
     reward_pred = tf.stack(
         gen_rewards[hparams.video_num_input_frames-1:], axis=1)
+    reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove undeeded dimension.
 
     return_targets = predictions
     if "target_reward" in features:

From a12eff74c5f3fda0e21cd9d8064cb94f37993717 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 11 Jun 2018 18:15:19 -0700
Subject: [PATCH 0081/2720] Add support for non-causal decoder self-attention

PiperOrigin-RevId: 200146777
---
 tensor2tensor/models/transformer.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index aa3bda483..c605dc1ad 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -741,14 +741,21 @@ def transformer_prepare_decoder(targets, hparams, features=None):
     decoder_input: a Tensor, bottom of decoder stack
     decoder_self_attention_bias: a bias tensor for use in decoder self-attention
   """
-  if hparams.prepend_mode == "prepend_inputs_full_attention":
-    decoder_self_attention_bias = (
-        common_attention.attention_bias_prepend_inputs_full_attention(
-            common_attention.embedding_to_padding(targets)))
+  if hparams.causal_decoder_self_attention:
+    # Causal attention.
+    if hparams.prepend_mode == "prepend_inputs_full_attention":
+      decoder_self_attention_bias = (
+          common_attention.attention_bias_prepend_inputs_full_attention(
+              common_attention.embedding_to_padding(targets)))
+    else:
+      decoder_self_attention_bias = (
+          common_attention.attention_bias_lower_triangle(
+              common_layers.shape_list(targets)[1]))
   else:
+    # Full attention.
+    decoder_padding = common_attention.embedding_to_padding(targets)
     decoder_self_attention_bias = (
-        common_attention.attention_bias_lower_triangle(
-            common_layers.shape_list(targets)[1]))
+        common_attention.attention_bias_ignore_padding(decoder_padding))
 
   if features and "targets_segmentation" in features:
     # "Packed" dataset - keep the examples from seeing each other.
@@ -1103,6 +1110,7 @@ def transformer_base_v1():
   hparams.add_hparam("pos", "timing")  # timing, none
   hparams.add_hparam("nbr_decoder_problems", 1)
   hparams.add_hparam("proximity_bias", False)
+  hparams.add_hparam("causal_decoder_self_attention", True)
   hparams.add_hparam("use_pad_remover", True)
   hparams.add_hparam("self_attention_type", "dot_product")
   hparams.add_hparam("max_relative_position", 0)

From 6b0359732acb6c655cdd330a6285f02c5d6e0f88 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 11 Jun 2018 19:45:36 -0700
Subject: [PATCH 0082/2720] Internal change

PiperOrigin-RevId: 200153630
---
 tensor2tensor/models/transformer.py | 20 ++++++--------------
 1 file changed, 6 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index c605dc1ad..aa3bda483 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -741,21 +741,14 @@ def transformer_prepare_decoder(targets, hparams, features=None):
     decoder_input: a Tensor, bottom of decoder stack
     decoder_self_attention_bias: a bias tensor for use in decoder self-attention
   """
-  if hparams.causal_decoder_self_attention:
-    # Causal attention.
-    if hparams.prepend_mode == "prepend_inputs_full_attention":
-      decoder_self_attention_bias = (
-          common_attention.attention_bias_prepend_inputs_full_attention(
-              common_attention.embedding_to_padding(targets)))
-    else:
-      decoder_self_attention_bias = (
-          common_attention.attention_bias_lower_triangle(
-              common_layers.shape_list(targets)[1]))
+  if hparams.prepend_mode == "prepend_inputs_full_attention":
+    decoder_self_attention_bias = (
+        common_attention.attention_bias_prepend_inputs_full_attention(
+            common_attention.embedding_to_padding(targets)))
   else:
-    # Full attention.
-    decoder_padding = common_attention.embedding_to_padding(targets)
     decoder_self_attention_bias = (
-        common_attention.attention_bias_ignore_padding(decoder_padding))
+        common_attention.attention_bias_lower_triangle(
+            common_layers.shape_list(targets)[1]))
 
   if features and "targets_segmentation" in features:
     # "Packed" dataset - keep the examples from seeing each other.
@@ -1110,7 +1103,6 @@ def transformer_base_v1():
   hparams.add_hparam("pos", "timing")  # timing, none
   hparams.add_hparam("nbr_decoder_problems", 1)
   hparams.add_hparam("proximity_bias", False)
-  hparams.add_hparam("causal_decoder_self_attention", True)
   hparams.add_hparam("use_pad_remover", True)
   hparams.add_hparam("self_attention_type", "dot_product")
   hparams.add_hparam("max_relative_position", 0)

From 5045d571bc7e3d5fcaf014569c6b3212109206ea Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 12 Jun 2018 15:31:06 -0700
Subject: [PATCH 0083/2720] [TF:XLA] Eliminate more copies after HLO
 scheduling.

PiperOrigin-RevId: 200292049
---
 tensor2tensor/layers/common_layers.py | 24 ------------------------
 1 file changed, 24 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index b7df15c30..03b8f2353 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2798,12 +2798,6 @@ def grad_fn(inputs, variables, outputs, output_grads):
     # bfloat16. This is a hack to ensure that grad_vars are the right type.
     if grad_inputs[0].dtype == tf.bfloat16:
       grad_vars = [tf.cast(grad_var, tf.bfloat16) for grad_var in grad_vars]
-    if is_on_tpu():
-      # TODO(noam): remove this hack once XLA does the right thing.
-      # Force the gradients on the inputs to be computed before the variables
-      # are updated.  This saves memory by preventing XLA from making an extra
-      # copy of the variables.
-      grad_vars = force_dependency(grad_vars, grad_inputs)
     return grad_inputs, grad_vars
 
   @fn_with_custom_grad(grad_fn)
@@ -2815,24 +2809,6 @@ def fn_with_recompute(*args):
   return fn_with_recompute(*args)
 
 
-def force_dependency(xs, ys):
-  """Force all of xs to depend on all of ys, using a false data dependency.
-
-  XLA seems to ignore control dependencies.
-
-  Args:
-    xs: a list of tensors
-    ys: a list of tensors:
-  Returns:
-    a list of tensors of the same length as xs
-  """
-  def _first_element(x):
-    ndims = x.get_shape().ndims
-    return tf.reshape(tf.slice(x, [0] * ndims, [1] * ndims), [])
-  my_zero = tf.add_n([_first_element(y) for y  in ys if y is not None]) * 1e-30
-  return [x + my_zero for x in xs]
-
-
 def dense(x, units, **kwargs):
   """Identical to tf.layers.dense, Memory optimization on tpu."""
   fn = lambda x: tf.layers.dense(x, units, **kwargs)

From 0b23f5877169560db9f918e110a524cdb234280b Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 12 Jun 2018 17:02:22 -0700
Subject: [PATCH 0084/2720] Add WaterWorld T2T Problem

PiperOrigin-RevId: 200306422
---
 setup.py                                      |  1 +
 tensor2tensor/data_generators/gym_problems.py | 22 +++++++++++++++++++
 .../data_generators/gym_problems_test.py      |  4 ++++
 3 files changed, 27 insertions(+)

diff --git a/setup.py b/setup.py
index fed6e8c20..0de3d2425 100644
--- a/setup.py
+++ b/setup.py
@@ -39,6 +39,7 @@
         'google-api-python-client',
         'gunicorn',
         'gym',
+        'gym_ple',
         'h5py',
         'numpy',
         'requests',
diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 9368a7d5e..f53c5cb65 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -21,6 +21,7 @@
 import math
 import os
 import gym
+import gym_ple       # pylint: disable=unused-import
 import numpy as np
 
 from tensor2tensor.data_generators import gym_utils
@@ -167,6 +168,27 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
              "reward": [int(reward - self.min_reward)]}
 
 
+@registry.register_problem
+class GymWaterWorldRandom5k(GymDiscreteProblem):
+  """WaterWorld game, random actions."""
+
+  @property
+  def env_name(self):
+    return "WaterWorldFast-v0"
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+  @property
+  def num_steps(self):
+    return 5000
+
+
 @registry.register_problem
 class GymPongRandom5k(GymDiscreteProblem):
   """Pong game, random actions."""
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index 1bbd2a374..a0cc7bcb3 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -38,6 +38,10 @@ def testGymAtariBoots(self):
     problem = gym_problems.GymPongRandom5k()
     self.assertEqual(5000, problem.num_steps)
 
+  def testPyGamesBoots(self):
+    problem = gym_problems.GymWaterWorldRandom5k()
+    self.assertEqual(5000, problem.num_steps)
+
 
 if __name__ == "__main__":
   tf.test.main()

From 5e3b43eb94866addd446ecb0c31032e967007f5e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 12 Jun 2018 17:17:22 -0700
Subject: [PATCH 0085/2720] [TF:XLA] Rollback of: Eliminate more copies after
 HLO scheduling.

PiperOrigin-RevId: 200309129
---
 tensor2tensor/layers/common_layers.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 03b8f2353..b7df15c30 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2798,6 +2798,12 @@ def grad_fn(inputs, variables, outputs, output_grads):
     # bfloat16. This is a hack to ensure that grad_vars are the right type.
     if grad_inputs[0].dtype == tf.bfloat16:
       grad_vars = [tf.cast(grad_var, tf.bfloat16) for grad_var in grad_vars]
+    if is_on_tpu():
+      # TODO(noam): remove this hack once XLA does the right thing.
+      # Force the gradients on the inputs to be computed before the variables
+      # are updated.  This saves memory by preventing XLA from making an extra
+      # copy of the variables.
+      grad_vars = force_dependency(grad_vars, grad_inputs)
     return grad_inputs, grad_vars
 
   @fn_with_custom_grad(grad_fn)
@@ -2809,6 +2815,24 @@ def fn_with_recompute(*args):
   return fn_with_recompute(*args)
 
 
+def force_dependency(xs, ys):
+  """Force all of xs to depend on all of ys, using a false data dependency.
+
+  XLA seems to ignore control dependencies.
+
+  Args:
+    xs: a list of tensors
+    ys: a list of tensors:
+  Returns:
+    a list of tensors of the same length as xs
+  """
+  def _first_element(x):
+    ndims = x.get_shape().ndims
+    return tf.reshape(tf.slice(x, [0] * ndims, [1] * ndims), [])
+  my_zero = tf.add_n([_first_element(y) for y  in ys if y is not None]) * 1e-30
+  return [x + my_zero for x in xs]
+
+
 def dense(x, units, **kwargs):
   """Identical to tf.layers.dense, Memory optimization on tpu."""
   fn = lambda x: tf.layers.dense(x, units, **kwargs)

From f8ec31ed39fbdaba3195008580c45d07799f3c17 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 12 Jun 2018 17:47:14 -0700
Subject: [PATCH 0086/2720] Adding SSO version of BAIR dataset

PiperOrigin-RevId: 200312982
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 .../data_generators/bair_robot_pushing.py     | 124 ++++++++++++++++++
 2 files changed, 125 insertions(+)
 create mode 100644 tensor2tensor/data_generators/bair_robot_pushing.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index ee2e91796..4c2b49684 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -25,6 +25,7 @@
     "tensor2tensor.data_generators.algorithmic_math",
     "tensor2tensor.data_generators.audio",
     "tensor2tensor.data_generators.babi_qa",
+    "tensor2tensor.data_generators.bair_robot_pushing",
     "tensor2tensor.data_generators.celeba",
     "tensor2tensor.data_generators.cifar",
     "tensor2tensor.data_generators.cipher",
diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
new file mode 100644
index 000000000..0f40c1b3e
--- /dev/null
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -0,0 +1,124 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Berkeley (BAIR) robot pushing dataset.
+
+Self-Supervised Visual Planning with Temporal Skip Connections
+Frederik Ebert, Chelsea Finn, Alex X. Lee, and Sergey Levine.
+https://arxiv.org/abs/1710.05268
+
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import tarfile
+import numpy as np
+from PIL import Image
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import video_utils
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+DATA_URL = (
+    "http://rail.eecs.berkeley.edu/datasets/bair_robot_pushing_dataset_v0.tar")
+
+
+@registry.register_problem
+class BairRobotPushing(video_utils.VideoProblem):
+  """Berkeley (BAIR) robot pushing dataset."""
+
+  @property
+  def num_channels(self):
+    return 3
+
+  @property
+  def frame_height(self):
+    return 64
+
+  @property
+  def frame_width(self):
+    return 64
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  def parse_frames(self, filenames):
+    image_key = "{}/image_aux1/encoded"
+    action_key = "{}/action"
+    state_key = "{}/endeffector_pos"
+
+    for f in filenames:
+      print("Parsing ", f)
+      for serialized_example in tf.python_io.tf_record_iterator(f):
+        x = tf.train.Example()
+        x.ParseFromString(serialized_example)
+        # there are 4 features per frame
+        # main image, aux image, actions and states
+        nf = len(x.features.feature.keys()) // 4
+
+        for i in range(nf):
+          image_name = image_key.format(i)
+          action_name = action_key.format(i)
+          state_name = state_key.format(i)
+
+          byte_str = x.features.feature[image_name].bytes_list.value[0]
+          img = Image.frombytes(
+              "RGB", (self.frame_width, self.frame_height), byte_str)
+          arr = np.array(img.getdata())
+          frame = arr.reshape(
+              self.frame_width, self.frame_height, self.num_channels)
+
+          state = x.features.feature[state_name].float_list.value
+          action = x.features.feature[action_name].float_list.value
+
+          yield i, frame, state, action
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    path = generator_utils.maybe_download(
+        tmp_dir, os.path.basename(DATA_URL), DATA_URL)
+
+    tar = tarfile.open(path)
+    tar.extractall(tmp_dir)
+    tar.close()
+
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      base_dir = os.path.join(tmp_dir, "softmotion30_44k/train/*")
+    else:
+      base_dir = os.path.join(tmp_dir, "softmotion30_44k/test/*")
+
+    filenames = tf.gfile.Glob(base_dir)
+    for frame_number, frame, state, action in self.parse_frames(filenames):
+      yield {
+          "frame_number": [frame_number],
+          "frame": frame,
+          "state": state,
+          "action": action,
+      }
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    p.input_modality = {
+        # Pixels are in 0..255 range.
+        "inputs": ("video:raw", 256),
+    }
+    p.target_modality = {
+        "targets": ("video:raw", 256),
+    }

From 5128fcbd572c85de9dc7764e540e0d9a8823c301 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 12 Jun 2018 19:04:00 -0700
Subject: [PATCH 0087/2720] Adding OSS version of Google robot pushing dataset

PiperOrigin-RevId: 200320863
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 .../data_generators/google_robot_pushing.py   | 124 ++++++++++++++++++
 2 files changed, 125 insertions(+)
 create mode 100644 tensor2tensor/data_generators/google_robot_pushing.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 4c2b49684..7b8d8137f 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -33,6 +33,7 @@
     "tensor2tensor.data_generators.desc2code",
     "tensor2tensor.data_generators.fsns",
     "tensor2tensor.data_generators.gene_expression",
+    "tensor2tensor.data_generators.google_robot_pushing",
     "tensor2tensor.data_generators.gym_problems",
     "tensor2tensor.data_generators.ice_parsing",
     "tensor2tensor.data_generators.imagenet",
diff --git a/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
new file mode 100644
index 000000000..176cab263
--- /dev/null
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -0,0 +1,124 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Google robot pushing dataset.
+
+Unsupervised Learning for Physical Interaction through Video Prediction
+Chelsea Finn, Ian Goodfellow, Sergey Levine
+https://arxiv.org/abs/1605.07157
+
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import io
+import os
+import numpy as np
+from PIL import Image
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import video_utils
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+BASE_URL = "https://storage.googleapis.com/brain-robotics-data/push/"
+DATA_TRAIN = (264, "push_train/push_train.tfrecord-{:05d}-of-00264")
+DATA_TEST_SEEN = (5, "/push_testseen/push_testseen.tfrecord-{:05d}-of-00005")
+DATA_TEST_NOVEL = (5, "/push_testnovel/push_testnovel.tfrecord-{:05d}-of-00005")
+
+
+@registry.register_problem
+class GoogleRobotPushing(video_utils.VideoProblem):
+  """Google robot pushing dataset."""
+
+  @property
+  def num_channels(self):
+    return 3
+
+  @property
+  def frame_height(self):
+    return 64
+
+  @property
+  def frame_width(self):
+    return 64
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  def parse_frames(self, filename):
+    image_key = "move/{}/image/encoded"
+    action_key = "move/{}/commanded_pose/vec_pitch_yaw"
+    state_key = "move/{}/endeffector/vec_pitch_yaw"
+
+    for serialized_example in tf.python_io.tf_record_iterator(filename):
+      x = tf.train.Example()
+      x.ParseFromString(serialized_example)
+      # there are 6 features per frame
+      nf = len(x.features.feature.keys()) // 6
+
+      for i in range(nf):
+        image_name = image_key.format(i)
+        action_name = action_key.format(i)
+        state_name = state_key.format(i)
+
+        byte_str = x.features.feature[image_name].bytes_list.value[0]
+        img = Image.open(io.BytesIO(byte_str))
+        # The original images are much bigger than 64x64
+        img = img.resize((self.frame_width, self.frame_height),
+                         resample=Image.BILINEAR)
+        arr = np.array(img.getdata())
+        frame = arr.reshape(
+            self.frame_width, self.frame_height, self.num_channels)
+
+        state = x.features.feature[state_name].float_list.value
+        action = x.features.feature[action_name].float_list.value
+
+        yield i, frame, state, action
+
+  def get_urls(self, count, url_part):
+    template = os.path.join(BASE_URL, url_part)
+    return [template.format(i) for i in range(count)]
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      urls = self.get_urls(DATA_TRAIN[0], DATA_TRAIN[1])
+    else:
+      urls = self.get_urls(DATA_TEST_SEEN[0], DATA_TEST_SEEN[1])
+      urls += self.get_urls(DATA_TEST_NOVEL[0], DATA_TEST_NOVEL[1])
+
+    for url in urls:
+      path = generator_utils.maybe_download(tmp_dir, os.path.basename(url), url)
+      for frame_number, frame, state, action in self.parse_frames(path):
+        yield {
+            "frame_number": [frame_number],
+            "frame": frame,
+            "state": state,
+            "action": action,
+        }
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    p.input_modality = {
+        # Pixels are in 0..255 range.
+        "inputs": ("video:raw", 256),
+    }
+    p.target_modality = {
+        "targets": ("video:raw", 256),
+    }

From 5dac697d0634a3639de7a83986053d98cf183f2e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 13 Jun 2018 10:16:23 -0700
Subject: [PATCH 0088/2720] Add support for non-causal decoder self-attention

PiperOrigin-RevId: 200410759
---
 tensor2tensor/models/transformer.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index aa3bda483..c605dc1ad 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -741,14 +741,21 @@ def transformer_prepare_decoder(targets, hparams, features=None):
     decoder_input: a Tensor, bottom of decoder stack
     decoder_self_attention_bias: a bias tensor for use in decoder self-attention
   """
-  if hparams.prepend_mode == "prepend_inputs_full_attention":
-    decoder_self_attention_bias = (
-        common_attention.attention_bias_prepend_inputs_full_attention(
-            common_attention.embedding_to_padding(targets)))
+  if hparams.causal_decoder_self_attention:
+    # Causal attention.
+    if hparams.prepend_mode == "prepend_inputs_full_attention":
+      decoder_self_attention_bias = (
+          common_attention.attention_bias_prepend_inputs_full_attention(
+              common_attention.embedding_to_padding(targets)))
+    else:
+      decoder_self_attention_bias = (
+          common_attention.attention_bias_lower_triangle(
+              common_layers.shape_list(targets)[1]))
   else:
+    # Full attention.
+    decoder_padding = common_attention.embedding_to_padding(targets)
     decoder_self_attention_bias = (
-        common_attention.attention_bias_lower_triangle(
-            common_layers.shape_list(targets)[1]))
+        common_attention.attention_bias_ignore_padding(decoder_padding))
 
   if features and "targets_segmentation" in features:
     # "Packed" dataset - keep the examples from seeing each other.
@@ -1103,6 +1110,7 @@ def transformer_base_v1():
   hparams.add_hparam("pos", "timing")  # timing, none
   hparams.add_hparam("nbr_decoder_problems", 1)
   hparams.add_hparam("proximity_bias", False)
+  hparams.add_hparam("causal_decoder_self_attention", True)
   hparams.add_hparam("use_pad_remover", True)
   hparams.add_hparam("self_attention_type", "dot_product")
   hparams.add_hparam("max_relative_position", 0)

From f4731039345ad0ef5af69b61a451e0eb7d41ccf6 Mon Sep 17 00:00:00 2001
From: Shah Newaz Khan <shahnewazk@gmail.com>
Date: Wed, 13 Jun 2018 11:27:52 -0700
Subject: [PATCH 0089/2720] Added note to download original problem data when
 using the  flag for translation data-sets (#863)

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index acdbf4099..0d5aec2d7 100644
--- a/README.md
+++ b/README.md
@@ -153,7 +153,7 @@ There are a number of translation data-sets in T2T:
 
 You can get translations in the other direction by appending `_rev` to
 the problem name, e.g., for German-English use
-`--problem=translate_ende_wmt32k_rev`.
+`--problem=translate_ende_wmt32k_rev` (NOTE: You still need to download the orginal data with t2t-datagen `--problem=translate_ende_wmt32k`).
 
 For all translation problems, we suggest to try the Transformer model:
 `--model=transformer`. At first it is best to try the base setting,

From dc0e82d0fb2632fc6ef6bf0e50d194963e4da6d5 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 13 Jun 2018 12:33:17 -0700
Subject: [PATCH 0090/2720] Remove change requiring gym_ple and ple. Breaks
 Travis.

PiperOrigin-RevId: 200435130
---
 README.md                                     |  2 +-
 setup.py                                      |  1 -
 tensor2tensor/data_generators/gym_problems.py | 22 -------------------
 .../data_generators/gym_problems_test.py      |  4 ----
 4 files changed, 1 insertion(+), 28 deletions(-)

diff --git a/README.md b/README.md
index 0d5aec2d7..acdbf4099 100644
--- a/README.md
+++ b/README.md
@@ -153,7 +153,7 @@ There are a number of translation data-sets in T2T:
 
 You can get translations in the other direction by appending `_rev` to
 the problem name, e.g., for German-English use
-`--problem=translate_ende_wmt32k_rev` (NOTE: You still need to download the orginal data with t2t-datagen `--problem=translate_ende_wmt32k`).
+`--problem=translate_ende_wmt32k_rev`.
 
 For all translation problems, we suggest to try the Transformer model:
 `--model=transformer`. At first it is best to try the base setting,
diff --git a/setup.py b/setup.py
index 0de3d2425..fed6e8c20 100644
--- a/setup.py
+++ b/setup.py
@@ -39,7 +39,6 @@
         'google-api-python-client',
         'gunicorn',
         'gym',
-        'gym_ple',
         'h5py',
         'numpy',
         'requests',
diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index f53c5cb65..9368a7d5e 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -21,7 +21,6 @@
 import math
 import os
 import gym
-import gym_ple       # pylint: disable=unused-import
 import numpy as np
 
 from tensor2tensor.data_generators import gym_utils
@@ -168,27 +167,6 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
              "reward": [int(reward - self.min_reward)]}
 
 
-@registry.register_problem
-class GymWaterWorldRandom5k(GymDiscreteProblem):
-  """WaterWorld game, random actions."""
-
-  @property
-  def env_name(self):
-    return "WaterWorldFast-v0"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-  @property
-  def num_steps(self):
-    return 5000
-
-
 @registry.register_problem
 class GymPongRandom5k(GymDiscreteProblem):
   """Pong game, random actions."""
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index a0cc7bcb3..1bbd2a374 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -38,10 +38,6 @@ def testGymAtariBoots(self):
     problem = gym_problems.GymPongRandom5k()
     self.assertEqual(5000, problem.num_steps)
 
-  def testPyGamesBoots(self):
-    problem = gym_problems.GymWaterWorldRandom5k()
-    self.assertEqual(5000, problem.num_steps)
-
 
 if __name__ == "__main__":
   tf.test.main()

From 0d386398939684f26e55e4b8e800da9c3468576a Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 13 Jun 2018 13:26:17 -0700
Subject: [PATCH 0091/2720] Internal merge #863

PiperOrigin-RevId: 200443144
---
 README.md           | 4 +++-
 docs/walkthrough.md | 4 +++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index acdbf4099..1522a9542 100644
--- a/README.md
+++ b/README.md
@@ -153,7 +153,9 @@ There are a number of translation data-sets in T2T:
 
 You can get translations in the other direction by appending `_rev` to
 the problem name, e.g., for German-English use
-`--problem=translate_ende_wmt32k_rev`.
+`--problem=translate_ende_wmt32k_rev`
+(note that you still need to download the original data with t2t-datagen
+`--problem=translate_ende_wmt32k`).
 
 For all translation problems, we suggest to try the Transformer model:
 `--model=transformer`. At first it is best to try the base setting,
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index acdbf4099..1522a9542 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -153,7 +153,9 @@ There are a number of translation data-sets in T2T:
 
 You can get translations in the other direction by appending `_rev` to
 the problem name, e.g., for German-English use
-`--problem=translate_ende_wmt32k_rev`.
+`--problem=translate_ende_wmt32k_rev`
+(note that you still need to download the original data with t2t-datagen
+`--problem=translate_ende_wmt32k`).
 
 For all translation problems, we suggest to try the Transformer model:
 `--model=transformer`. At first it is best to try the base setting,

From 55631cc4141a7ca42918b530f9dbf277a80f662e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 13 Jun 2018 15:26:41 -0700
Subject: [PATCH 0092/2720] Add Human3.6M dataset.

PiperOrigin-RevId: 200464258
---
 tensor2tensor/data_generators/bair_robot_pushing.py   | 6 +++---
 tensor2tensor/data_generators/google_robot_pushing.py | 6 +++---
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index 0f40c1b3e..44309cf1f 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -41,7 +41,7 @@
 
 
 @registry.register_problem
-class BairRobotPushing(video_utils.VideoProblem):
+class VideoBairRobotPushing(video_utils.VideoProblem):
   """Berkeley (BAIR) robot pushing dataset."""
 
   @property
@@ -117,8 +117,8 @@ def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.input_modality = {
         # Pixels are in 0..255 range.
-        "inputs": ("video:raw", 256),
+        "inputs": ("video", 256),
     }
     p.target_modality = {
-        "targets": ("video:raw", 256),
+        "targets": ("video", 256),
     }
diff --git a/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index 176cab263..924671785 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -43,7 +43,7 @@
 
 
 @registry.register_problem
-class GoogleRobotPushing(video_utils.VideoProblem):
+class VideoGoogleRobotPushing(video_utils.VideoProblem):
   """Google robot pushing dataset."""
 
   @property
@@ -117,8 +117,8 @@ def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.input_modality = {
         # Pixels are in 0..255 range.
-        "inputs": ("video:raw", 256),
+        "inputs": ("video", 256),
     }
     p.target_modality = {
-        "targets": ("video:raw", 256),
+        "targets": ("video", 256),
     }

From 7c6d5908aad08c567e7b1098da5cf6656edab82e Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 13 Jun 2018 16:10:18 -0700
Subject: [PATCH 0093/2720] Correct the link to instruction for ImageNet data
 generation.

PiperOrigin-RevId: 200471134
---
 tensor2tensor/data_generators/imagenet.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 6b5037735..5735244ce 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -125,8 +125,8 @@ def num_classes(self):
   def generate_data(self, data_dir, tmp_dir, task_id=-1):
     # TODO(lukaszkaiser): find a better way than printing this.
     print("To generate the ImageNet dataset in the proper format, follow "
-          "instructions at https://github.com/tensorflow/models/blob/master"
-          "/inception/README.md#getting-started")
+          "instructions at https://github.com/tensorflow/models/tree/master"
+          "/research/inception/README.md#getting-started")
 
   def preprocess_example(self, example, mode, _):
     return imagenet_preprocess_example(example, mode)

From a8c501043bd488feeb8fcf0595b8476b4b31c0b1 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 13 Jun 2018 16:21:09 -0700
Subject: [PATCH 0094/2720] [TF:XLA] Rollback of: Eliminate more copies after
 HLO scheduling.

PiperOrigin-RevId: 200472722
---
 tensor2tensor/layers/common_layers.py | 24 ------------------------
 1 file changed, 24 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index b7df15c30..03b8f2353 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2798,12 +2798,6 @@ def grad_fn(inputs, variables, outputs, output_grads):
     # bfloat16. This is a hack to ensure that grad_vars are the right type.
     if grad_inputs[0].dtype == tf.bfloat16:
       grad_vars = [tf.cast(grad_var, tf.bfloat16) for grad_var in grad_vars]
-    if is_on_tpu():
-      # TODO(noam): remove this hack once XLA does the right thing.
-      # Force the gradients on the inputs to be computed before the variables
-      # are updated.  This saves memory by preventing XLA from making an extra
-      # copy of the variables.
-      grad_vars = force_dependency(grad_vars, grad_inputs)
     return grad_inputs, grad_vars
 
   @fn_with_custom_grad(grad_fn)
@@ -2815,24 +2809,6 @@ def fn_with_recompute(*args):
   return fn_with_recompute(*args)
 
 
-def force_dependency(xs, ys):
-  """Force all of xs to depend on all of ys, using a false data dependency.
-
-  XLA seems to ignore control dependencies.
-
-  Args:
-    xs: a list of tensors
-    ys: a list of tensors:
-  Returns:
-    a list of tensors of the same length as xs
-  """
-  def _first_element(x):
-    ndims = x.get_shape().ndims
-    return tf.reshape(tf.slice(x, [0] * ndims, [1] * ndims), [])
-  my_zero = tf.add_n([_first_element(y) for y  in ys if y is not None]) * 1e-30
-  return [x + my_zero for x in xs]
-
-
 def dense(x, units, **kwargs):
   """Identical to tf.layers.dense, Memory optimization on tpu."""
   fn = lambda x: tf.layers.dense(x, units, **kwargs)

From 751aac162fce5e4bb2561c59a389ce44f34f6719 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 13 Jun 2018 17:29:18 -0700
Subject: [PATCH 0095/2720] Move hparams function call to registry.hparams; add
 check for None return; fix next_frame_tpu

PiperOrigin-RevId: 200482340
---
 tensor2tensor/models/distillation.py        |  4 ++--
 tensor2tensor/models/research/next_frame.py |  1 +
 tensor2tensor/rl/model_rl_experiment.py     |  2 +-
 tensor2tensor/utils/registry.py             |  7 ++++++-
 tensor2tensor/utils/registry_test.py        | 17 +++++++++++++----
 tensor2tensor/utils/trainer_lib.py          |  2 +-
 tensor2tensor/utils/trainer_lib_test.py     |  2 +-
 7 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/models/distillation.py b/tensor2tensor/models/distillation.py
index 9ef341a77..1af4ee843 100644
--- a/tensor2tensor/models/distillation.py
+++ b/tensor2tensor/models/distillation.py
@@ -54,11 +54,11 @@ def __init__(self,
     elif hparams.distill_phase == "distill" and hparams.student_learning_rate:
       hparams.learning_rate = hparams.student_learning_rate
 
-    self.teacher_hparams = registry.hparams(hparams.teacher_hparams)()
+    self.teacher_hparams = registry.hparams(hparams.teacher_hparams)
     self.teacher_model = registry.model(
         hparams.teacher_model)(self.teacher_hparams, mode, problem_hparams,
                                data_parallelism, decode_hparams)
-    self.student_hparams = registry.hparams(hparams.student_hparams)()
+    self.student_hparams = registry.hparams(hparams.student_hparams)
     self.student_model = registry.model(
         hparams.student_model)(self.student_hparams, mode, problem_hparams,
                                data_parallelism, decode_hparams)
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 4d6edf457..1c80bd542 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -733,6 +733,7 @@ def next_frame_stochastic():
 def next_frame_tpu():
   hparams = next_frame()
   hparams.batch_size = 1
+  return hparams
 
 
 @registry.register_hparams
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 89ed60674..931ba0caf 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -781,7 +781,7 @@ def rl_modelrl_tiny_simulation_deterministic_starts():
 
 
 def create_loop_hparams():
-  hparams = registry.hparams(FLAGS.loop_hparams_set)()
+  hparams = registry.hparams(FLAGS.loop_hparams_set)
   hparams.parse(FLAGS.loop_hparams)
   return hparams
 
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index f73a92e5e..318abac96 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -165,12 +165,17 @@ def decorator(hp_fn, registration_name=None):
 
 
 def hparams(name):
+  """Retrieve registered hparams by name."""
   if name not in _HPARAMS:
     error_msg = "HParams set %s never registered. Sets registered:\n%s"
     raise LookupError(
         error_msg % (name,
                      display_list_by_prefix(list_hparams(), starting_spaces=4)))
-  return _HPARAMS[name]
+  hp = _HPARAMS[name]()
+  if hp is None:
+    raise TypeError("HParams %s is None. Make sure the registered function "
+                    "returns the HParams object." % name)
+  return hp
 
 
 def list_hparams():
diff --git a/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index bc593edf5..aa57e3482 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -109,13 +109,13 @@ def testHParamSet(self):
 
     @registry.register_hparams
     def my_hparams_set():
-      pass
+      return 3
 
     @registry.register_ranged_hparams
     def my_hparams_range(_):
       pass
 
-    self.assertTrue(registry.hparams("my_hparams_set") is my_hparams_set)
+    self.assertEqual(registry.hparams("my_hparams_set"), my_hparams_set())
     self.assertTrue(
         registry.ranged_hparams("my_hparams_range") is my_hparams_range)
 
@@ -123,13 +123,13 @@ def testNamedRegistration(self):
 
     @registry.register_hparams("a")
     def my_hparams_set():
-      pass
+      return 7
 
     @registry.register_ranged_hparams("a")
     def my_hparams_range(_):
       pass
 
-    self.assertTrue(registry.hparams("a") is my_hparams_set)
+    self.assertEqual(registry.hparams("a"), my_hparams_set())
     self.assertTrue(registry.ranged_hparams("a") is my_hparams_range)
 
   def testUnknownHparams(self):
@@ -138,6 +138,15 @@ def testUnknownHparams(self):
     with self.assertRaisesRegexp(LookupError, "never registered"):
       registry.ranged_hparams("not_registered")
 
+  def testNoneHparams(self):
+
+    @registry.register_hparams
+    def hp():
+      pass
+
+    with self.assertRaisesRegexp(TypeError, "is None"):
+      registry.hparams("hp")
+
   def testDuplicateRegistration(self):
 
     @registry.register_hparams
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index ad17b1da9..bbc6d15ce 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -69,7 +69,7 @@ def create_hparams(hparams_set,
                    data_dir=None,
                    problem_name=None):
   """Create HParams with data_dir and problem hparams, if kwargs provided."""
-  hparams = registry.hparams(hparams_set)()
+  hparams = registry.hparams(hparams_set)
   if data_dir:
     hparams.add_hparam("data_dir", data_dir)
   if problem_name:
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index e735c05b0..799168d44 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -67,7 +67,7 @@ def testExperiment(self):
         use_tpu=False)
     run_config = trainer_lib.create_run_config(
         model_dir=self.data_dir, num_gpus=0, use_tpu=False)
-    hparams = registry.hparams("transformer_tiny_tpu")()
+    hparams = registry.hparams("transformer_tiny_tpu")
     exp = exp_fn(run_config, hparams)
     exp.test()
 

From 7c077a78c716f4747992f18bdf48ed39abfbb69d Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Wed, 13 Jun 2018 17:46:18 -0700
Subject: [PATCH 0096/2720] Adding discretized mixture of logistics loss, an
 imagetransformerPlus model, a new cifar10 problem, and a bug fix for
 converting integer pixel intensities to [-1, 1].

PiperOrigin-RevId: 200484336
---
 tensor2tensor/data_generators/cifar.py        |  14 ++
 .../layers/common_image_attention.py          |   8 +-
 tensor2tensor/layers/common_layers.py         | 162 +++++++++++++++
 tensor2tensor/layers/modalities.py            |  13 +-
 tensor2tensor/models/image_transformer.py     | 195 ++++++++++++++++++
 tensor2tensor/utils/metrics.py                |  12 ++
 6 files changed, 397 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index 74532af17..dcdacb64a 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -27,6 +27,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import mnist
+from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -172,6 +173,19 @@ def preprocess_example(self, example, mode, unused_hparams):
     return example
 
 
+@registry.register_problem
+class ImageCifar10PlainGenDmol(ImageCifar10PlainGen):
+  """Discretized mixture of logistics problem."""
+
+  def dataset_filename(self):
+    return "image_cifar10_plain"  # Reuse CIFAR-10 plain data.
+
+  def eval_metrics(self):
+    return [
+        metrics.Metrics.DMOL_PERPLEXITY
+    ]
+
+
 @registry.register_problem
 class ImageCifar10Plain8(ImageCifar10):
   """CIFAR-10 rescaled to 8x8 for output: Conditional image generation."""
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index a34841a88..a30c23c60 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -621,10 +621,8 @@ def prepare_image(inputs, hparams, name=None):
   channels = hparams.num_channels
 
   hidden_size = hparams.hidden_size
-  # Only do lookup if the embeddings haven't been looked up already.
-  # if the last dimension is number of channels, then this is very likely the
-  # channel ids tensor. We have to make sure.
-  if inputs_shape[-1] == hparams.num_channels:
+  # Only do lookup if the modality is identity
+  if hparams.target_modality == "image:identity":
     inputs = tf.to_int32(inputs)
     x = get_channel_embeddings(channels, inputs, hidden_size, name=name)
   else:
@@ -675,5 +673,5 @@ def add_pos_signals(x, hparams, name="pos_emb"):
     else:
       assert hparams.pos == "emb"
       x = common_attention.add_positional_embedding_nd(
-          x, hparams.max_length, name=name)
+          x, hparams.max_length, name)
   return x
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 03b8f2353..17f8d3322 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -206,6 +206,17 @@ def convert_rgb_to_real(x):
     return x
 
 
+def convert_rgb_to_symmetric_real(x):
+  """Conversion of pixel values to real numbers."""
+  with tf.name_scope("rgb_to_real", values=[x]):
+    x = tf.to_float(x)
+    # Use the formula (value/127.5) - 1 to convert each channel value into a
+    # real number in the range -1 to 1. We use 127.5 instead of 128 because
+    # the intensities are in the range 0 to 255. This is used for dmol.
+    x = (x / 127.5) - 1
+    return x
+
+
 def convert_real_to_rgb(x):
   """Conversion of real numbers to pixel values."""
   with tf.name_scope("real_to_rgb", values=[x]):
@@ -1849,6 +1860,157 @@ def padded_cross_entropy(logits,
     return tf.reduce_sum(xent * weights), tf.reduce_sum(weights)
 
 
+def dml_loss(
+    pred, labels,
+    weights_fn=weights_nonzero,  # Unused
+    reduce_sum=True):
+  """Discretized mixture of logistics loss.
+
+  Args:
+    pred: a tensor of shape [batch, height, width, 10*num_mixtures]
+    labels: a [batch, height, width, channels] tensor of real pixel intensities
+    weights_fn: weights function
+    reduce_sum: A boolean, to return scalar mean loss instead of per position
+
+  Returns:
+    a pair of tensors of loss/sum of losses, denominator
+  """
+  del weights_fn  # Unused
+  real_labels = convert_rgb_to_real(labels)
+  dml_loss_value = discretized_mix_logistic_loss(real_labels, pred,
+                                                 sum_all=reduce_sum)
+  if reduce_sum:
+    return dml_loss_value, tf.reduce_sum(tf.ones(tf.shape(labels),
+                                                 tf.float32))
+  else:
+    return dml_loss_value/3., tf.ones(tf.shape(dml_loss_value),
+                                      tf.float32)
+
+
+def discretized_mix_logistic_loss(labels, pred, sum_all=True):
+  """Computes log likelihood for the discretized mixture of logistics loss.
+
+  Args:
+    labels: A [batch, height, width, channels] tensor of true pixel intensities
+      rescaled to [-1, 1]. The computation assumes channels is 3.
+    pred: A [batch, height, width, hparams.num_mixtures*10] tensor of floats
+      comprising one unnormalized mixture probability, three means
+      (one per channel), three standard deviations (one per channel),
+      and three coefficients which linearly parameterize dependence across
+      channels
+    sum_all: A boolean to return scalar mean loss or per position
+
+  Returns:
+    per_position_loss: A [batch, height, width, 3] tensor of the
+      conditional probability of each channel given all previous channels if
+      not sum_all else add all the losses (for eval)
+  """
+
+  # Extract mixture probabilities, means, scale, and coefficient parameters.
+  l_shape = shape_list(labels)
+  num_mix = shape_list(pred)[-1] // 10
+  logits = pred[:, :, :, :num_mix]  # unnormalized mixture probabilities
+  pred = tf.reshape(
+      pred[:, :, :, num_mix:],
+      [l_shape[0], l_shape[1], l_shape[2], l_shape[3], num_mix * 3])
+  means = pred[:, :, :, :, :num_mix]
+  log_scales = tf.maximum(pred[:, :, :, :, num_mix:2 * num_mix], -7.)
+  coeffs = tf.tanh(pred[:, :, :, :, 2 * num_mix:3 * num_mix])
+  labels = (  # tile labels across number of mixtures
+      tf.reshape(labels, [l_shape[0], l_shape[1], l_shape[2], l_shape[3], 1]) +
+      tf.zeros([l_shape[0], l_shape[1], l_shape[2], l_shape[3], num_mix])
+  )
+
+  # p(x) = sigmoid((x - means_i + 1/255.)/scale_i) -
+  #        sigmoid((x - means_i - 1/255.)/scale_i)
+  # for each channel i. The means are linearly parameterized.
+  means_0 = tf.reshape(means[:, :, :, 0, :],
+                       [l_shape[0], l_shape[1], l_shape[2], 1, num_mix])
+  means_1 = tf.reshape(means[:, :, :, 1, :] +
+                       coeffs[:, :, :, 0, :] * labels[:, :, :, 0, :],
+                       [l_shape[0], l_shape[1], l_shape[2], 1, num_mix])
+  means_2 = tf.reshape(means[:, :, :, 2, :] +
+                       coeffs[:, :, :, 1, :] * labels[:, :, :, 0, :] +
+                       coeffs[:, :, :, 2, :] * labels[:, :, :, 1, :],
+                       [l_shape[0], l_shape[1], l_shape[2], 1, num_mix])
+  means = tf.concat([means_0, means_1, means_2], 3)
+  centered_labels = labels - means
+  inv_stdv = tf.exp(-log_scales)
+  plus_in = inv_stdv * (centered_labels + 1. / 255.)
+  min_in = inv_stdv * (centered_labels - 1. / 255.)
+  cdf_plus = tf.nn.sigmoid(plus_in)
+  cdf_min = tf.nn.sigmoid(min_in)
+  # Compute log probability for edge case of 0 (before scaling), 255 (before
+  # scaling), and all other cases respectively.
+  prob_0 = plus_in - tf.nn.softplus(plus_in)
+  prob_255 = -tf.nn.softplus(min_in)
+  prob_event = cdf_plus - cdf_min
+
+  # Robustly select log-prob based on numerical edge-cases: (a) [-1, -1+eps);
+  # (b) (1-eps, 1]; (c) NaNs during `tf.gradients` of `tf.select`, which may
+  # cause `tf.log(0.)`; (d) p(x) < 1e-5.
+  mid_in = inv_stdv * centered_labels
+  log_prob_event_approx = (
+      mid_in - log_scales - 2. * tf.nn.softplus(mid_in) - np.log(127.5))
+  log_probs = tf.where(labels < -0.999, prob_0,
+                       tf.where(labels > 0.999, prob_255,
+                                tf.where(prob_event > 1e-5,
+                                         tf.log(tf.maximum(prob_event, 1e-12)),
+                                         log_prob_event_approx)))
+  # Sum over mixtures.
+  log_probs = tf.reduce_sum(log_probs, 3) + tf.nn.log_softmax(logits, axis=-1)
+  if sum_all:
+    output = -tf.reduce_sum(tf.reduce_logsumexp(log_probs, axis=-1))
+    return output
+  else:
+    output = -tf.reduce_logsumexp(log_probs, axis=-1, keep_dims=True)
+    return output
+
+
+def sample_from_discretized_mix_logistic(l, nr_mix, seed=None):
+  """Sampling from a discretized mixture of logistics using gumbel softmax.
+
+  Args:
+    l: output of body, of shape [batch, length, num_mixtures * 10]
+    nr_mix: Integer number of mixtures
+    seed: Random seed
+
+  Returns:
+    A tensor of shape [batch, length, 3] with real intensities scaled between
+    -1 and 1
+  """
+  ls_t = tf.shape(l)
+  ls = [ls_t[0], ls_t[1], ls_t[2], ls_t[3]]
+  xs = ls[:-1] + [3]
+  # unpack parameters
+  logit_probs = l[:, :, :, :nr_mix]
+  l = tf.reshape(l[:, :, :, nr_mix:], xs + [nr_mix * 3])
+  # sample mixture indicator from softmax using gumbel softmax
+  sel = tf.one_hot(tf.argmax(logit_probs - tf.log(-tf.log(tf.random_uniform(
+      tf.shape(logit_probs), minval=1e-5, maxval=1. - 1e-5, seed=seed))), 3),
+                   depth=nr_mix, dtype=tf.float32)
+  sel = tf.reshape(sel, xs[:-1] + [1, nr_mix])
+  # select logistic parameters
+  means = tf.reduce_sum(l[:, :, :, :, :nr_mix] * sel, 4)
+  log_scales = tf.maximum(tf.reduce_sum(
+      l[:, :, :, :, nr_mix:2 * nr_mix] * sel, 4), -7.)
+  coeffs = tf.reduce_sum(tf.nn.tanh(
+      l[:, :, :, :, 2 * nr_mix:3 * nr_mix]) * sel, 4)
+  # sample from logistic & clip to interval
+  # we don't actually round to the nearest 8bit value when sampling
+  u = tf.random_uniform(tf.shape(means), minval=1e-5, maxval=1. - 1e-5,
+                        seed=seed)
+  x = means + tf.exp(log_scales) * (tf.log(u) - tf.log(1. - u))
+  x0 = tf.clip_by_value(x[:, :, :, 0], -1., 1)
+  x1 = tf.clip_by_value(x[:, :, :, 1] + coeffs[:, :, :, 0] * x0, -1., 1)
+  x2 = tf.minimum(tf.maximum(
+      x[:, :, :, 2] + coeffs[:, :, :, 1] * x0 + coeffs[:, :, :, 2] * x1, -1.),
+                  1.)
+  return tf.concat([tf.reshape(x0, xs[:-1] + [1]),
+                    tf.reshape(x1, xs[:-1] + [1]),
+                    tf.reshape(x2, xs[:-1] + [1])], 3)
+
+
 def smoothing_cross_entropy(logits,
                             labels,
                             vocab_size,
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 797340123..f6de1d1f1 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -290,8 +290,10 @@ def bottom_compress(self, inputs, name="bottom"):
     """
     with tf.variable_scope(name):
       inputs = tf.to_float(inputs)
-      tf.summary.image("inputs", inputs, max_outputs=2)
-      inputs = common_layers.convert_rgb_to_real(inputs)
+      hp = self._model_hparams
+      if hp.mode != tf.estimator.ModeKeys.PREDICT:
+        tf.summary.image("inputs", inputs, max_outputs=2)
+      inputs = common_layers.convert_rgb_to_symmetric_real(inputs)
       ishape = common_layers.shape_list(inputs)
       inputs = tf.reshape(inputs, [-1, ishape[1], ishape[2] * ishape[3], 1])
       inputs.set_shape([None, None, None, 1])
@@ -334,6 +336,13 @@ def top(self, body_output, _):
       return x
 
 
+@registry.register_image_modality("image_channel_bottom_identity")
+class ImageChannelBottomIdentityModality(ImageChannelCompressModality):
+
+  def top(self, body_output, _):
+    return body_output
+
+
 @registry.register_image_modality("channel_embeddings_bottom")
 class ImageChannelEmbeddingsBottom(modality.Modality):
   """Modality for images using channel compression for generation."""
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 7c1ab1e6f..6bc0a4026 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -70,6 +70,59 @@ def body(self, features):
       return output
 
 
+@registry.register_model
+class ImagetransformerPlus(t2t_model.T2TModel):
+  """Imagetransformer with discretized mixture of logistics loss."""
+
+  def body(self, features):
+    hparams = copy.copy(self._hparams)
+    inputs = features["inputs"]
+    targets = features["targets"]
+    # Prepare decoder inputs and bias.
+    decoder_input, _, _ = cia.prepare_decoder(targets, hparams)
+    # Add class label to decoder input.
+    if not hparams.unconditional:
+      decoder_input += tf.reshape(
+          inputs,
+          [common_layers.shape_list(targets)[0], 1, 1, hparams.hidden_size])
+    decoder_output = cia.transformer_decoder_layers(
+        decoder_input,
+        None,
+        hparams.num_decoder_layers or hparams.num_hidden_layers,
+        hparams,
+        attention_type=hparams.dec_attention_type,
+        name="decoder")
+    # reshape it into [batch, height, width, depth]
+    decoder_output = tf.reshape(decoder_output, tf.shape(targets))
+    # there are 10 sets of parameters that you need to produce, location, scale,
+    # and coefficient parameter for each
+    output = tf.layers.dense(decoder_output, hparams.num_mixtures*10,
+                             use_bias=False, activation=None,
+                             name="output_mixtures_conv")
+    # TODO(avaswani) Figure out if we need residuals or layer norm
+    return output
+
+  def loss(self, pred, features):
+    return common_layers.dml_loss(pred, features["targets"])
+
+  def sample(self, features):
+    """Run the model and extract samples.
+
+    Args:
+      features: an map of string to `Tensor`.
+
+    Returns:
+       samples: an integer `Tensor`.
+       logits: a list of `Tensor`s, one per datashard.
+       losses: a dictionary: {loss-name (string): floating point `Scalar`}.
+    """
+    logits, losses = self(features)  # pylint: disable=not-callable
+
+    samples = common_layers.sample_from_discretized_mix_logistic(
+        logits, 10, seed=None)
+    return samples, logits, losses
+
+
 @registry.register_model
 class ImagetransformerMoe(t2t_model.T2TModel):
   """Conditional image generation with attention and MoE."""
@@ -175,6 +228,9 @@ def image_transformer_base():
 
   hparams.add_hparam("unconditional", False)  # unconditional generation
 
+  # parameters of discretized mixture of logistics loss from pixel cnn++
+  hparams.add_hparam("num_mixtures", 10)
+
   # These parameters are only used when ffn_layer=="local_moe_tpu"
   hparams.add_hparam("moe_overhead_train", 1.0)
   hparams.add_hparam("moe_overhead_eval", 2.0)
@@ -262,6 +318,145 @@ def imagetransformer_base_10l_8h_big_uncond_dr03_dan_64():
   return hparams
 
 
+@registry.register_hparams
+def imagetransformerpp_sep_channels_8l_8h():
+  """separate rgb embeddings."""
+  hparams = imagetransformer_base()
+  hparams.num_heads = 8
+  hparams.batch_size = 4
+  hparams.attention_key_channels = hparams.attention_value_channels = 0
+  hparams.hidden_size = 512
+  hparams.filter_size = 512
+  hparams.num_hidden_layers = 8
+  hparams.sampling_method = "random"
+  hparams.layer_preprocess_sequence = "n"
+  hparams.layer_postprocess_sequence = "da"
+  hparams.target_modality = "image:image_channel_bottom_identity"
+  hparams.summarize_grads = True
+  hparams.learning_rate = 0.1
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_8l_8h_big_cond_dr03_dan():
+  """big 1d model for conditional image generation.2.99 on cifar10."""
+  hparams = imagetransformerpp_sep_channels_8l_8h()
+  hparams.hidden_size = 512
+  hparams.num_heads = 8
+  hparams.filter_size = 2048
+  hparams.batch_size = 4
+  hparams.max_length = 3075
+  hparams.layer_prepostprocess_dropout = 0.3
+  hparams.layer_preprocess_sequence = "none"
+  hparams.layer_postprocess_sequence = "dan"
+  hparams.summarize_grads = True
+  hparams.learning_rate = 0.01
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_8l_8h_big_cond_dr03_dan_a():
+  hparams = imagetransformerpp_base_8l_8h_big_cond_dr03_dan()
+  hparams.learning_rate = 0.1
+  hparams.num_channels = 1
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_10l_8h_big_uncond_dr03_dan():
+  hparams = imagetransformerpp_base_8l_8h_big_cond_dr03_dan_a()
+  hparams.unconditional = True
+  hparams.num_decoder_layers = 10
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_a():
+  hparams = imagetransformerpp_base_10l_8h_big_uncond_dr03_dan()
+  hparams.learning_rate = 0.01
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_b():
+  hparams = imagetransformerpp_base_10l_8h_big_uncond_dr03_dan()
+  hparams.learning_rate = 0.1
+  hparams.hidden_size = 256
+  hparams.attention_key_channels = 512
+  hparams.attention_value_channels = 512
+  hparams.filter_size = 1024
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_e():
+  hparams = imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_b()
+  hparams.learning_rate_warmup_steps = 16000
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_f():
+  hparams = imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_b()
+  hparams.num_mixtures = 5
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_g():
+  hparams = imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_b()
+  hparams.filter_size = 512
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.learning_rate = 0.1
+  hparams.layer_preprocess_sequence = "none"
+  hparams.layer_postprocess_sequence = "dan"
+  hparams.pos = "emb"
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_k():
+  hparams = imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_g()
+  hparams.num_decoder_layers = 12
+  hparams.clip_grad_norm = 0.
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_l():
+  hparams = imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_g()
+  hparams.num_decoder_layers = 12
+  hparams.clip_grad_norm = 40.
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_m():
+  hparams = imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_k()
+  hparams.batch_size = 8
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_14l_8h_big_uncond_dr03_dan_p():
+  """Gets to 2.92 in just under 4 days on 8 p100s."""
+  hparams = imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_l()
+  hparams.num_decoder_layers = 14
+  hparams.batch_size = 8
+  hparams.layer_prepostprocess_dropout = 0.2
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_14l_8h_big_uncond_dr03_dan_eval():
+  """Gets to 2.92 in just under 4 days on 8 p100s."""
+  hparams = imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_l()
+  hparams.num_decoder_layers = 14
+  hparams.batch_size = 8
+  # hparams.layer_prepostprocess_dropout = 0.2
+  return hparams
+
+
 @registry.register_hparams
 def imagetransformer_base_8l_8h_big_cond_dr03_dan_128():
   hparams = imagetransformer_base_8l_8h_big_cond_dr03_dan()
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 1d9ac96b7..ed086b2b0 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -54,6 +54,7 @@ class Metrics(object):
   SIGMOID_CROSS_ENTROPY_ONE_HOT = "sigmoid_cross_entropy_one_hot"
   ROC_AUC = "roc_auc"
   IMAGE_SUMMARY = "image_summary"
+  DMOL_PERPLEXITY = "disc_mol_neg_log_perplexity"
   IMAGE_RMSE = "image_rmse"
 
 
@@ -232,6 +233,16 @@ def padded_neg_log_perplexity(predictions,
   return (-num, den)
 
 
+def dmol_neg_log_perplexity(predictions,
+                            labels,
+                            weights_fn=None):
+  """Average log-perplexity excluding padding 0s. No smoothing."""
+  del weights_fn  # Unused
+  num, den = common_layers.dml_loss(
+      predictions, labels, reduce_sum=False)
+  return (-num, den)
+
+
 def rounding_accuracy(predictions,
                       labels,
                       weights_fn=common_layers.weights_nonzero):
@@ -609,5 +620,6 @@ def metric_means():
     Metrics.SET_RECALL: set_recall,
     Metrics.ROC_AUC: roc_auc,
     Metrics.IMAGE_SUMMARY: image_summary,
+    Metrics.DMOL_PERPLEXITY: dmol_neg_log_perplexity,
     Metrics.IMAGE_RMSE: image_rmse,
 }

From 2fddb53f0b5f1fcc671f79f1a4c258d7b50899a1 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 13 Jun 2018 17:59:20 -0700
Subject: [PATCH 0097/2720] Corrections to video generation (also reduce
 logging bloat by having one Session).

PiperOrigin-RevId: 200485711
---
 tensor2tensor/bin/t2t_decoder.py              |  2 +
 .../data_generators/bair_robot_pushing.py     |  5 ++
 .../data_generators/generator_utils.py        |  2 +-
 .../data_generators/google_robot_pushing.py   |  5 ++
 tensor2tensor/data_generators/image_utils.py  |  4 +-
 tensor2tensor/data_generators/video_utils.py  | 55 ++++++++++++-------
 6 files changed, 49 insertions(+), 24 deletions(-)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 576941fe4..9de17a669 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -33,6 +33,7 @@
 
 import os
 from tensor2tensor.bin import t2t_trainer
+from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import registry
@@ -162,6 +163,7 @@ def main(_):
   trainer_lib.set_random_seed(FLAGS.random_seed)
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
 
+
   if FLAGS.score_file:
     filename = os.path.expanduser(FLAGS.score_file)
     if not tf.gfile.Exists(filename):
diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index 44309cf1f..498a714ed 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -60,6 +60,11 @@ def frame_width(self):
   def is_generate_per_split(self):
     return True
 
+  @property
+  def total_number_of_frames(self):
+    # TODO(mbz): correct this number to be the real total number of frames.
+    return 30 * 10 * 1000
+
   def parse_frames(self, filenames):
     image_key = "{}/image_aux1/encoded"
     action_key = "{}/action"
diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index af665d16b..9bd6cd199 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -154,7 +154,7 @@ def generate_files(generator, output_filenames,
   for case in generator:
     if case is None:
       continue
-    if counter > 0 and counter % 100000 == 0:
+    if counter % 100000 == 0:
       tf.logging.info("Generating case %d." % counter)
     counter += 1
     if max_cases and counter > max_cases:
diff --git a/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index 924671785..1f53fdd93 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -58,6 +58,11 @@ def frame_height(self):
   def frame_width(self):
     return 64
 
+  @property
+  def total_number_of_frames(self):
+    # TODO(mbz): correct this number to be the real total number of frames.
+    return 50 * 10 * 1000
+
   @property
   def is_generate_per_split(self):
     return True
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index a824913cd..e3088ac9a 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -158,9 +158,9 @@ def encode_images_as_png(images):
     for image in images:
       yield tf.image.encode_png(image).numpy()
   else:
-    (width, height, channels) = images[0].shape
+    (height, width, channels) = images[0].shape
     with tf.Graph().as_default():
-      image_t = tf.placeholder(dtype=tf.uint8, shape=(width, height, channels))
+      image_t = tf.placeholder(dtype=tf.uint8, shape=(height, width, channels))
       encoded_image_t = tf.image.encode_png(image_t)
       with tf.Session() as sess:
         for image in images:
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index d5a06435a..7be1efe1d 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -22,7 +22,6 @@
 import six
 
 from tensor2tensor.data_generators import generator_utils
-from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.utils import metrics
@@ -72,6 +71,12 @@ def frame_shape(self):
   @property
   def total_number_of_frames(self):
     """The total number of frames, needed for sharding."""
+    # It can also be a lower number -- we will switch shards every
+    # total_number_of_frames // num_shards time, so for example if
+    # you know that every video is 30 frames long and you have 100 shards
+    # then it's sufficient to set this to 30 * 100 so no shard-switching
+    # occurs during the generation of a video. For videos of variable length,
+    # just make this large so switching shards mid-video is very rare.
     raise NotImplementedError
 
   @property
@@ -244,25 +249,32 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
     Raises:
       ValueError: if the frame has a different number of channels than required.
     """
-    for features in self.generate_samples(data_dir, tmp_dir, dataset_split):
-      unencoded_frame = features.pop("frame")
-      height, width, channels = unencoded_frame.shape
-      if channels != self.num_channels:
-        raise ValueError("Generated frame has %d channels while the class "
-                         "assumes %d channels." % (channels, self.num_channels))
-      if height != self.frame_height:
-        raise ValueError("Generated frame has height %d while the class "
-                         "assumes height %d." % (height, self.frame_height))
-      if width != self.frame_width:
-        raise ValueError("Generated frame has width %d while the class "
-                         "assumes width %d." % (width, self.frame_width))
-      encoded_frame = six.next(
-          image_utils.encode_images_as_png([unencoded_frame]))
-      features["image/encoded"] = [encoded_frame]
-      features["image/format"] = ["png"]
-      features["image/height"] = [height]
-      features["image/width"] = [width]
-      yield features
+    with tf.Graph().as_default():
+      image_t = tf.placeholder(
+          dtype=tf.uint8,
+          shape=(self.frame_height, self.frame_width, self.num_channels))
+      encoded_image_t = tf.image.encode_png(image_t)
+      with tf.Session() as sess:
+        for features in self.generate_samples(data_dir, tmp_dir, dataset_split):
+          unencoded_frame = features.pop("frame")
+          height, width, channels = unencoded_frame.shape
+          if channels != self.num_channels:
+            raise ValueError("Generated frame has %d channels while the class "
+                             "assumes %d channels." % (channels,
+                                                       self.num_channels))
+          if height != self.frame_height:
+            raise ValueError("Generated frame has height %d while the class "
+                             "assumes height %d." % (height, self.frame_height))
+          if width != self.frame_width:
+            raise ValueError("Generated frame has width %d while the class "
+                             "assumes width %d." % (width, self.frame_width))
+          encoded_frame = sess.run(encoded_image_t, feed_dict={
+              image_t: unencoded_frame})
+          features["image/encoded"] = [encoded_frame]
+          features["image/format"] = ["png"]
+          features["image/height"] = [height]
+          features["image/width"] = [width]
+          yield features
 
   def generate_encoded_samples_debug(self, data_dir, tmp_dir, dataset_split):
     """Generate samples of the encoded frames and dump for debug if needed."""
@@ -303,7 +315,8 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
       for split, paths in split_paths:
         generator_utils.generate_files(
             self.generate_encoded_samples_debug(
-                data_dir, tmp_dir, split), paths)
+                data_dir, tmp_dir, split), paths,
+            cycle_every_n=self.total_number_of_frames // len(paths))
     else:
       generator_utils.generate_files(
           self.generate_encoded_samples_debug(

From 61ae5f5342be4f103b6f4e4c0e7f63e5e2726562 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 13 Jun 2018 21:13:35 -0700
Subject: [PATCH 0098/2720] Revert back to decoding output in certain cases.

PiperOrigin-RevId: 200501655
---
 tensor2tensor/utils/decoding.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 584b01623..1d79f8483 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -93,8 +93,7 @@ def log_decode_results(inputs,
     decoded_outputs = targets_vocab.decode(_save_until_eos(outputs, is_image))
     if targets is not None and log_results:
       decoded_targets = targets_vocab.decode(_save_until_eos(targets, is_image))
-  if log_results:
-    tf.logging.info("Inference results OUTPUT: %s" % decoded_outputs)
+  tf.logging.info("Inference results OUTPUT: %s" % decoded_outputs)
   if targets is not None and log_results:
     tf.logging.info("Inference results TARGET: %s" % decoded_targets)
   return decoded_inputs, decoded_outputs, decoded_targets

From 956dda1b31ceb8e2ce474f21813884e97d971168 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 13 Jun 2018 23:21:39 -0700
Subject: [PATCH 0099/2720] set a max for number of frames in Google robot
 pushing dataset.

PiperOrigin-RevId: 200510193
---
 tensor2tensor/data_generators/google_robot_pushing.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index 1f53fdd93..d2e8011ca 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -63,6 +63,10 @@ def total_number_of_frames(self):
     # TODO(mbz): correct this number to be the real total number of frames.
     return 50 * 10 * 1000
 
+  @property
+  def max_number_of_frames_per_video(self):
+    return 60
+
   @property
   def is_generate_per_split(self):
     return True
@@ -77,6 +81,8 @@ def parse_frames(self, filename):
       x.ParseFromString(serialized_example)
       # there are 6 features per frame
       nf = len(x.features.feature.keys()) // 6
+      # it seems features after 60 don't have any image
+      nf = min(nf, self.max_number_of_frames_per_video)
 
       for i in range(nf):
         image_name = image_key.format(i)

From 256f774d28f012d6e74fc369e5b58ff439a38a63 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 14 Jun 2018 08:53:07 -0700
Subject: [PATCH 0100/2720] Correct one name remaining after renaming
 bottleneck_size to bottleneck_bits.

PiperOrigin-RevId: 200563071
---
 tensor2tensor/models/research/autoencoders.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 5930e09e1..8cbfd1207 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -535,7 +535,7 @@ def autoencoder_discrete_cifar():
   """Discrete autoencoder model for compressing cifar."""
   hparams = autoencoder_ordered_discrete()
   hparams.bottleneck_noise = 0.0
-  hparams.bottleneck_size = 90
+  hparams.bottleneck_bits = 90
   hparams.unordered = True
   hparams.num_hidden_layers = 2
   hparams.hidden_size = 256

From 05080ec344475b0dabaa20fd758c8ff2b79b1995 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 14 Jun 2018 09:56:39 -0700
Subject: [PATCH 0101/2720] Rename ImagetransformerPlus.loss(self, pred,
 features) to ImagetransformerPlus.loss(self, *logits*, features) to keep lint
 happy.

PiperOrigin-RevId: 200572086
---
 tensor2tensor/models/image_transformer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 6bc0a4026..7833f46d9 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -102,8 +102,8 @@ def body(self, features):
     # TODO(avaswani) Figure out if we need residuals or layer norm
     return output
 
-  def loss(self, pred, features):
-    return common_layers.dml_loss(pred, features["targets"])
+  def loss(self, logits, features):
+    return common_layers.dml_loss(logits, features["targets"])
 
   def sample(self, features):
     """Run the model and extract samples.

From 863362028b435ce17b394cf952da4610ce7181c9 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Thu, 14 Jun 2018 12:52:24 -0700
Subject: [PATCH 0102/2720] Make transformer_nat work

PiperOrigin-RevId: 200603226
---
 .../models/research/transformer_nat.py        | 205 ++++++++++--------
 1 file changed, 113 insertions(+), 92 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index e9393fdc5..0a42b4b24 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -27,32 +27,30 @@
 from tensorflow.python.training import moving_averages
 
 
-def get_vq_bottleneck(bottleneck_size, hidden_size):
+def init_vq_bottleneck(bottleneck_size, hidden_size):
   """Get lookup table for VQ bottleneck."""
-  with tf.variable_scope("vq", reuse=tf.AUTO_REUSE):
-    means = tf.get_variable(
-        name="means",
-        shape=[bottleneck_size, hidden_size],
-        initializer=tf.uniform_unit_scaling_initializer())
-
-    ema_count = tf.get_variable(
-        name="ema_count",
-        shape=[bottleneck_size],
-        initializer=tf.constant_initializer(0),
+  means = tf.get_variable(
+      name="means",
+      shape=[bottleneck_size, hidden_size],
+      initializer=tf.uniform_unit_scaling_initializer())
+  ema_count = tf.get_variable(
+      name="ema_count",
+      shape=[bottleneck_size],
+      initializer=tf.constant_initializer(0),
+      trainable=False)
+  with tf.colocate_with(means):
+    ema_means = tf.get_variable(
+        name="ema_means",
+        initializer=means.initialized_value(),
         trainable=False)
 
-    with tf.colocate_with(means):
-      ema_means = tf.get_variable(
-          name="ema_means",
-          initializer=means.initialized_value(),
-          trainable=False)
-
   return means, ema_means, ema_count
 
 
-def vq_nearest_neighbor(x, means, hparams):
+def vq_nearest_neighbor(x, hparams):
   """Find the nearest element in means to elements in x."""
-  bottleneck_size = common_layers.shape_list(means)[0]
+  bottleneck_size = 2**hparams.bottleneck_bits
+  means = hparams.means
   x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keepdims=True)
   means_norm_sq = tf.reduce_sum(tf.square(means), axis=-1, keepdims=True)
   scalar_prod = tf.matmul(x, means, transpose_b=True)
@@ -60,60 +58,63 @@ def vq_nearest_neighbor(x, means, hparams):
   if hparams.bottleneck_kind == "em":
     x_means_idx = tf.multinomial(-dist, num_samples=hparams.num_samples)
     x_means_hot = tf.one_hot(
-        x_means_idx, depth=common_layers.shape_list(means)[0])
-    x_means_hot = tf.reduce_sum(x_means_hot, axis=1)
+        x_means_idx, depth=bottleneck_size)
+    x_means_hot = tf.reduce_mean(x_means_hot, axis=1)
   else:
     x_means_idx = tf.argmax(-dist, axis=-1)
-    x_means_hot = tf.one_hot(x_means_idx, bottleneck_size)
-  x_means_hot_flat = tf.reshape(x_means_hot, [-1, bottleneck_size])
-  x_means = tf.matmul(x_means_hot_flat, means)
+    x_means_hot = tf.one_hot(x_means_idx, depth=bottleneck_size)
+  x_means = tf.matmul(x_means_hot, means)
   e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
   return x_means_hot, e_loss
 
 
 def vq_discrete_bottleneck(x, hparams):
   """Simple vector quantized discrete bottleneck."""
+  tf.logging.info("Using EMA with beta = {}".format(hparams.beta))
   bottleneck_size = 2**hparams.bottleneck_bits
   x_shape = common_layers.shape_list(x)
-  means, ema_means, ema_count = get_vq_bottleneck(bottleneck_size,
-                                                  hparams.hidden_size)
   x = tf.reshape(x, [-1, hparams.hidden_size])
   x_means_hot, e_loss = vq_nearest_neighbor(
-      x, means, hparams)
+      x, hparams)
+  means, ema_means, ema_count = (hparams.means, hparams.ema_means,
+                                 hparams.ema_count)
 
   # Update the ema variables
   updated_ema_count = moving_averages.assign_moving_average(
       ema_count,
-      tf.reduce_sum(
-          tf.reshape(x_means_hot, shape=[-1, bottleneck_size]), axis=0),
+      tf.reduce_sum(x_means_hot, axis=0),
       hparams.decay,
       zero_debias=False)
 
   dw = tf.matmul(x_means_hot, x, transpose_a=True)
-  updated_ema_means = tf.identity(moving_averages.assign_moving_average(
-      ema_means, dw, hparams.decay, zero_debias=False))
+  updated_ema_means = moving_averages.assign_moving_average(
+      ema_means, dw, hparams.decay, zero_debias=False)
   n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
   updated_ema_count = (
       (updated_ema_count + hparams.epsilon) /
       (n + bottleneck_size * hparams.epsilon) * n)
-  updated_ema_means /= tf.expand_dims(updated_ema_count, axis=-1)
+  # pylint: disable=g-no-augmented-assignment
+  updated_ema_means = updated_ema_means / tf.expand_dims(
+      updated_ema_count, axis=-1)
+  # pylint: enable=g-no-augmented-assignment
   with tf.control_dependencies([e_loss]):
-    update_means = means.assign(updated_ema_means)
+    update_means = tf.assign(means, updated_ema_means)
     with tf.control_dependencies([update_means]):
       loss = hparams.beta * e_loss
 
-  d = tf.reshape(x_means_hot, x_shape[:-1] + [bottleneck_size])
-  return d, loss
+  discrete = tf.reshape(x_means_hot, x_shape[:-1] + [bottleneck_size])
+  return discrete, loss
 
 
 def vq_discrete_unbottleneck(x, hparams):
   """Simple undiscretization from vector quantized representation."""
   x_shape = common_layers.shape_list(x)
-  x = tf.to_float(x)
   bottleneck_size = 2**hparams.bottleneck_bits
-  means, _, _ = get_vq_bottleneck(bottleneck_size, hparams.hidden_size)
-  result = tf.matmul(tf.reshape(x, [-1, x_shape[-1]]), means)
-  return tf.reshape(result, x_shape[:-1] + [hparams.hidden_size])
+  means = hparams.means
+  x_flat = tf.reshape(x, [-1, bottleneck_size])
+  result = tf.matmul(x_flat, means)
+  result = tf.reshape(result, x_shape[:-1] + [hparams.hidden_size])
+  return result
 
 
 def residual_conv(x, repeat, k, hparams, name, reuse=None):
@@ -177,8 +178,7 @@ def encode(x, x_space, hparams, name):
 def decode_transformer(encoder_output, encoder_decoder_attention_bias, targets,
                        hparams, name):
   """Original Transformer decoder."""
-  orig_hparams = hparams
-  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+  with tf.variable_scope(name):
     targets = common_layers.flatten4d3d(targets)
 
     decoder_input, decoder_self_bias = (
@@ -195,22 +195,22 @@ def decode_transformer(encoder_output, encoder_decoder_attention_bias, targets,
     decoder_output = tf.reshape(
         decoder_output, [decoder_output_shape[0], -1, 1, hparams.hidden_size])
     # Expand since t2t expects 4d tensors.
-    hparams = orig_hparams
     return decoder_output
 
 
-def get_latent_pred_loss(latents_pred, latents_discrete, hparams):
+def get_latent_pred_loss(latents_pred, latents_discrete_hot, hparams):
   """Latent prediction and loss."""
   latents_logits = tf.layers.dense(
       latents_pred, 2**hparams.bottleneck_bits, name="extra_logits")
+  # loss = tf.losses.softmax_cross_entropy(onehot_labels=latents_discrete_hot,
+  #                                        logits=latents_logits)
   loss = tf.nn.softmax_cross_entropy_with_logits_v2(
-      labels=latents_discrete, logits=latents_logits)
+      labels=tf.stop_gradient(latents_discrete_hot), logits=latents_logits)
   return loss
 
 
 def ae_latent_sample_beam(latents_dense_in, inputs, ed, embed, hparams):
   """Sample from the latent space in the autoencoder."""
-
   def symbols_to_logits_fn(ids):
     """Go from ids to logits."""
     ids = tf.expand_dims(ids, axis=2)  # Ids start with added all-zeros.
@@ -232,9 +232,9 @@ def symbols_to_logits_fn(ids):
   ids, _ = beam_search.beam_search(
       symbols_to_logits_fn,
       initial_ids,
-      1,
-      length,
-      2**hparams.bottleneck_bits,
+      beam_size=1,
+      decode_length=length,
+      vocab_size=2**hparams.bottleneck_bits,
       alpha=0.0,
       eos_id=-1,
       stop_early=False)
@@ -245,59 +245,45 @@ def symbols_to_logits_fn(ids):
 
 def ae_transformer_internal(inputs, targets, target_space, hparams, cache=None):
   """Main step used for training."""
-  # Prepare.
-  if inputs is not None:
-    batch_size = common_layers.shape_list(inputs)[0]
-  else:
-    batch_size = common_layers.shape_list(targets)[0]
-  targets = tf.reshape(targets, [batch_size, -1, 1, hparams.hidden_size])
-
   # Encoder.
-  if inputs is not None:
-    inputs = common_layers.flatten4d3d(inputs)
-    inputs, ed = encode(inputs, target_space, hparams, "input_enc")
-    inputs_ex, ed_ex = inputs, ed
-  else:
-    ed, inputs_ex, ed_ex = None, None, None
+  inputs = common_layers.flatten4d3d(inputs)
+  inputs, ed = encode(inputs, target_space, hparams, "input_enc")
 
   # Autoencoding.
   losses = {"extra": tf.constant(0.0), "latent_pred": tf.constant(0.0)}
 
   max_targets_len_from_inputs = tf.concat([inputs, inputs], axis=1)
-
   targets, _ = common_layers.pad_to_same_length(
       targets,
       max_targets_len_from_inputs,
       final_length_divisible_by=2**hparams.num_compress_steps)
   targets_c = compress(targets, hparams, "compress")
-
   if hparams.mode != tf.estimator.ModeKeys.PREDICT:
     # Compress and bottleneck.
     latents_discrete_hot, extra_loss = vq_discrete_bottleneck(
         x=targets_c, hparams=hparams)
-    latents_dense = vq_discrete_unbottleneck(latents_discrete_hot, hparams)
+    latents_dense = vq_discrete_unbottleneck(
+        latents_discrete_hot, hparams=hparams)
+    latents_dense = targets_c + tf.stop_gradient(latents_dense - targets_c)
     latents_discrete = tf.argmax(latents_discrete_hot, axis=-1)
     tf.summary.histogram("codes", tf.reshape(latents_discrete[:, 0, :], [-1]))
-    pc = common_layers.inverse_exp_decay(hparams.startup_steps)
-    pc = pc if hparams.mode == tf.estimator.ModeKeys.TRAIN else 1.0
-    cond = tf.less(tf.random_uniform([batch_size]), pc)
-    latents_dense = tf.where(cond, latents_dense, targets_c)
-    losses["extra"] = extra_loss * tf.reduce_mean(tf.to_float(cond))
-
-    # Extra loss predicting latent code from input. Discrete only.
-    latents_pred = decode_transformer(inputs_ex, ed_ex, latents_dense, hparams,
+    losses["extra"] = extra_loss
+
+    # Extra loss predicting latent code from input.
+    latents_pred = decode_transformer(inputs, ed, latents_dense, hparams,
                                       "extra")
     latent_pred_loss = get_latent_pred_loss(latents_pred, latents_discrete_hot,
                                             hparams)
-    losses["latent_pred"] = tf.reduce_mean(latent_pred_loss * tf.to_float(cond))
+    losses["latent_pred"] = tf.reduce_mean(latent_pred_loss)
   else:
     latent_len = common_layers.shape_list(targets_c)[1]
     embed = functools.partial(vq_discrete_unbottleneck, hparams=hparams)
     latents_dense = tf.zeros_like(targets_c[:, :latent_len, :, :])
     if cache is None:
-      cache = ae_latent_sample_beam(latents_dense, inputs_ex, ed_ex, embed,
+      cache = ae_latent_sample_beam(latents_dense, inputs, ed, embed,
                                     hparams)
-    latents_dense = embed(tf.one_hot(cache, depth=2**hparams.bottleneck_bits))
+    cache_hot = tf.one_hot(cache, depth=2**hparams.bottleneck_bits)
+    latents_dense = embed(cache_hot)
 
   # Postprocess.
   d = latents_dense
@@ -311,12 +297,8 @@ def ae_transformer_internal(inputs, targets, target_space, hparams, cache=None):
     d = residual_conv(d, 1, (3, 1), hparams, "decompress_rc_%d" % j)
     d = decompress_step(d, hparams, i > 0, "decompress_%d" % j)
 
-  res = decode_transformer(inputs, ed, targets, hparams, "decoder")
-  # We'll start training the extra model of latents after mask_startup_steps.
-  nonlatent_steps = hparams.mask_startup_steps
-  latent_time = tf.less(nonlatent_steps,
-                        tf.to_int32(tf.train.get_global_step()))
-  losses["latent_pred"] *= tf.to_float(latent_time)
+  targets = d
+  res = decode_transformer(inputs, ed, d, hparams, "decoder")
   return res, losses, cache
 
 
@@ -324,6 +306,14 @@ def ae_transformer_internal(inputs, targets, target_space, hparams, cache=None):
 class TransformerNAT(t2t_model.T2TModel):
   """Nonautoregressive Transformer from https://arxiv.org/abs/1805.11063."""
 
+  def __init__(self, *args, **kwargs):
+    super(TransformerNAT, self).__init__(*args, **kwargs)
+    means, ema_means, ema_count = init_vq_bottleneck(
+        2**self._hparams.bottleneck_bits, self._hparams.hidden_size)
+    self._hparams.means = means
+    self._hparams.ema_means = ema_means
+    self._hparams.ema_count = ema_count
+
   @property
   def has_input(self):
     return self._problem_hparams.input_modality
@@ -338,11 +328,10 @@ def body(self, features):
       return res, loss
 
   def prepare_features_for_infer(self, features):
-    beam_batch_size = self._decode_hparams.beam_size
-    beam_batch_size *= self._decode_hparams.batch_size
-    inputs = tf.zeros([beam_batch_size, 1, 1, self._hparams.hidden_size])
+    batch_size = self._decode_hparams.batch_size
+    inputs = tf.zeros([batch_size, 1, 1, self._hparams.hidden_size])
     inputs = inputs if "inputs" in features else None
-    targets = tf.zeros([beam_batch_size, 1, 1, self._hparams.hidden_size])
+    targets = tf.zeros([batch_size, 1, 1, self._hparams.hidden_size])
     with tf.variable_scope("transformer_nat/body"):
       _, _, cache = ae_transformer_internal(
           inputs, targets, features["target_space_id"], self._hparams)
@@ -356,9 +345,29 @@ def infer(self,
             alpha=0.0,
             use_tpu=False):
     """Produce predictions from the model."""
-    infer_out = super(TransformerNAT, self).infer(
-        features, decode_length, beam_size, top_beams, alpha, use_tpu=use_tpu)
-    return infer_out["outputs"]
+    if not features:
+      features = {}
+    inputs_old = None
+    if "inputs" in features and len(features["inputs"].shape) < 4:
+      inputs_old = features["inputs"]
+      features["inputs"] = tf.expand_dims(features["inputs"], 2)
+
+    # Create an initial targets tensor.
+    if "partial_targets" in features:
+      initial_output = tf.convert_to_tensor(features["partial_targets"])
+    else:
+      batch_size = common_layers.shape_list(features["inputs"])[0]
+      length = common_layers.shape_list(features["inputs"])[1]
+      target_length = tf.to_int32(2.0 * tf.to_float(length))
+      initial_output = tf.zeros((batch_size, target_length, 1, 1),
+                                dtype=tf.int64)
+
+    features["targets"] = initial_output
+    logits, _ = self(features)  # pylint: disable=not-callable
+    samples = tf.argmax(logits, axis=-1)
+    if inputs_old is not None:  # Restore to not confuse Estimator.
+      features["inputs"] = inputs_old
+    return samples
 
 
 @registry.register_hparams
@@ -372,15 +381,14 @@ def transformer_nat_small():
   hparams.hidden_size = 384
   hparams.filter_size = 2048
   hparams.label_smoothing = 0.0
-  hparams.optimizer = "Adam"  # Can be unstable, maybe try Adam.
+  hparams.force_full_predict = True
+  hparams.optimizer = "Adam"
   hparams.optimizer_adam_epsilon = 1e-9
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.997  # Needs tuning, try 0.98 to 0.999.
-  hparams.add_hparam("bottleneck_kind", "em")
+  hparams.add_hparam("bottleneck_kind", "vq")
   hparams.add_hparam("bottleneck_bits", 12)
   hparams.add_hparam("num_compress_steps", 3)
-  hparams.add_hparam("startup_steps", 10000)
-  hparams.add_hparam("mask_startup_steps", 50000)
   hparams.add_hparam("beta", 0.25)
   hparams.add_hparam("epsilon", 1e-5)
   hparams.add_hparam("decay", 0.999)
@@ -397,3 +405,16 @@ def transformer_nat_base():
   hparams.filter_size = 4096
   hparams.num_hidden_layers = 6
   return hparams
+
+
+@registry.register_hparams
+def transformer_nat_big():
+  """Set of hyperparameters."""
+  hparams = transformer_nat_small()
+  hparams.batch_size = 2048
+  hparams.hidden_size = 1024
+  hparams.filter_size = 4096
+  hparams.num_hidden_layers = 6
+  hparams.num_heads = 16
+  hparams.layer_prepostprocess_dropout = 0.3
+  return hparams

From 8ba4bd5d05766f7fa15d5bc17c44855c56dc1360 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 14 Jun 2018 13:17:16 -0700
Subject: [PATCH 0103/2720] A number of small corrections to make model-based
 RL loop run with stochastic next-frame model.

PiperOrigin-RevId: 200607033
---
 tensor2tensor/data_generators/gym_problems.py | 19 ++++----
 tensor2tensor/data_generators/gym_utils.py    | 10 ----
 tensor2tensor/data_generators/video_utils.py  |  8 +++-
 tensor2tensor/layers/modalities.py            | 44 ++++++++---------
 tensor2tensor/models/research/next_frame.py   | 47 ++++++++++++++-----
 tensor2tensor/rl/envs/simulated_batch_env.py  |  5 +-
 tensor2tensor/rl/model_rl_experiment.py       | 20 ++++++++
 tensor2tensor/rl/rl_trainer_lib.py            |  2 +-
 8 files changed, 92 insertions(+), 63 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 9368a7d5e..28614f0c3 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -23,7 +23,8 @@
 import gym
 import numpy as np
 
-from tensor2tensor.data_generators import gym_utils
+# We need gym_utils for the game environments defined there.
+from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.layers import discretization
@@ -546,8 +547,7 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
             err = np.ndarray.astype(np.maximum(np.abs(
                 self.real_ob - observ, dtype=np.int) - 10, 0),
                                     np.uint8)
-            debug_im_np = np.concatenate([observ, self.real_ob, err], axis=1)
-            debug_im = gym_utils.encode_image_to_png(debug_im_np)
+            debug_im = np.concatenate([observ, self.real_ob, err], axis=1)
           if done:
             self.dones += 1
             self.sum_of_rewards += self.real_reward
@@ -569,14 +569,12 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
             self.sum_of_rewards += self.real_reward
         if FLAGS.autoencoder_path:
           if self.simulated_environment:
-            debug_im = gym_utils.encode_image_to_png(
-                self.autodecode(observ, sess))
+            debug_im = self.autodecode(observ, sess)
           else:
             orig_observ = observ
             observ = self.autoencode(observ, sess)
-            debug_im_np = np.concatenate([self.autodecode(observ, sess),
-                                          orig_observ], axis=1)
-            debug_im = gym_utils.encode_image_to_png(debug_im_np)
+            debug_im = np.concatenate([self.autodecode(observ, sess),
+                                       orig_observ], axis=1)
         ret_dict = {"frame": observ,
                     "image/format": ["png"],
                     "image/height": [self.frame_height],
@@ -585,7 +583,7 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
                     "done": [int(False)],
                     "reward": [int(reward) - self.min_reward]}
         if self.make_extra_debug_info:
-          ret_dict["image/encoded_debug"] = [debug_im]
+          ret_dict["image/debug"] = debug_im
         yield ret_dict
         pieces_generated += 1
 
@@ -622,8 +620,7 @@ def real_env(self):
   def restore_networks(self, sess):
     super(GymSimulatedDiscreteProblemWithAgent, self).restore_networks(sess)
     # TODO(blazej): adjust regexp for different models.
-    env_model_loader = tf.train.Saver(tf.global_variables(
-        "next_frame_basic.*"))
+    env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
     sess = tf.get_default_session()
 
     ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
diff --git a/tensor2tensor/data_generators/gym_utils.py b/tensor2tensor/data_generators/gym_utils.py
index 46abe9bfa..9d1b421d9 100644
--- a/tensor2tensor/data_generators/gym_utils.py
+++ b/tensor2tensor/data_generators/gym_utils.py
@@ -19,10 +19,6 @@
 
 import numpy as np
 
-import six
-
-from tensor2tensor.data_generators import image_utils
-
 
 # pylint: disable=method-hidden
 class WarmupWrapper(gym.Wrapper):
@@ -304,9 +300,3 @@ def wrapped_freeway_factory(warm_up_examples=0,
                       easy_freeway=False
                   ),
                   max_episode_steps=500)
-
-
-def encode_image_to_png(image):
-  encoded = six.next(
-      image_utils.encode_images_as_png([image]))
-  return encoded
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 7be1efe1d..c0f9ceea7 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -251,8 +251,7 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
     """
     with tf.Graph().as_default():
       image_t = tf.placeholder(
-          dtype=tf.uint8,
-          shape=(self.frame_height, self.frame_width, self.num_channels))
+          dtype=tf.uint8, shape=(None, None, None))
       encoded_image_t = tf.image.encode_png(image_t)
       with tf.Session() as sess:
         for features in self.generate_samples(data_dir, tmp_dir, dataset_split):
@@ -274,6 +273,11 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
           features["image/format"] = ["png"]
           features["image/height"] = [height]
           features["image/width"] = [width]
+          if "image/debug" in features:
+            unencoded_debug = features.pop("image/debug")
+            encoded_debug = sess.run(encoded_image_t, feed_dict={
+                image_t: unencoded_debug})
+            features["image/encoded_debug"] = encoded_debug
           yield features
 
   def generate_encoded_samples_debug(self, data_dir, tmp_dir, dataset_split):
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index f6de1d1f1..11aa106c3 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -549,31 +549,6 @@ def loss(self, top_out, targets):
         weights_fn=self.targets_weights_fn)
 
 
-@registry.register_video_modality("raw")
-class VideoModalityRaw(modality.Modality):
-  """Modality for raw videos, i.e., time-sequences of frames."""
-
-  def bottom(self, x):
-    common_layers.summarize_video(x, "inputs")
-    return common_layers.convert_rgb_to_real(x)
-
-  def targets_bottom(self, x):
-    common_layers.summarize_video(x, "targets_bottom")
-    return common_layers.convert_rgb_to_real(x)
-
-  def top(self, body_output, _):
-    frames = tf.stack(body_output, axis=1)
-    rgb_frames = common_layers.convert_real_to_rgb(frames)
-    common_layers.summarize_video(rgb_frames, "body_output")
-    # TODO(lukaszkaiser): remove the need for the last dimension of 1 in eval.
-    return tf.expand_dims(rgb_frames, axis=-1)
-
-  def loss(self, top_out, targets):
-    top_out = tf.squeeze(top_out, axis=[-1])
-    loss = tf.square(top_out - tf.to_float(targets))
-    return tf.reduce_sum(loss), tf.reduce_sum(tf.ones_like(loss))
-
-
 @registry.register_video_modality("embed")
 class VideoModalityEmbed(VideoModality):
   """Video Modality where bottom embeds pixels."""
@@ -661,6 +636,25 @@ def internal_loss(self, logits, targets):
     return tf.nn.relu((logits - targets)**2 - self.cutoff * self.cutoff)
 
 
+@registry.register_video_modality("l2raw")
+class VideoModalityL2Raw(VideoModalityL2):
+  """Modality with L2 loss and raw input (sequences of frames)."""
+
+  def bottom(self, x):
+    common_layers.summarize_video(x, "inputs")
+    return common_layers.convert_rgb_to_real(x)
+
+  def targets_bottom(self, x):
+    common_layers.summarize_video(x, "targets_bottom")
+    return common_layers.convert_rgb_to_real(x)
+
+  def top(self, body_output, _):
+    frames = tf.stack(body_output, axis=1)
+    rgb_frames = common_layers.convert_real_to_rgb(frames)
+    common_layers.summarize_video(rgb_frames, "body_output")
+    return tf.expand_dims(rgb_frames, axis=-1)
+
+
 @registry.register_class_label_modality("default")
 class ClassLabelModality(modality.Modality):
   """Used for label data."""
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 1c80bd542..74d62eea8 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -130,10 +130,17 @@ def logits_to_samples(logits):
       num_channels = self.hparams.problem.num_channels
     except AttributeError:
       num_channels = 1
-    features["targets"] = tf.zeros(
-        [self.hparams.batch_size, 1, 1, 1, num_channels], dtype=tf.int32)
+    if "inputs" in features:
+      inputs_shape = common_layers.shape_list(features["inputs"])
+      targets_shape = [inputs_shape[0], self.hparams.video_num_target_frames,
+                       inputs_shape[2], inputs_shape[3], num_channels]
+    else:
+      tf.logging.warn("Guessing targets shape as no inputs are given.")
+      targets_shape = [self.hparams.batch_size,
+                       self.hparams.video_num_target_frames, 1, 1, num_channels]
+    features["targets"] = tf.zeros(targets_shape, dtype=tf.int32)
     features["target_reward"] = tf.zeros(
-        [self.hparams.batch_size, 1, 1], dtype=tf.int32)
+        [targets_shape[0], 1, 1], dtype=tf.int32)
     logits, _ = self(features)  # pylint: disable=not-callable
     if isinstance(logits, dict):
       results = {}
@@ -151,6 +158,9 @@ def logits_to_samples(logits):
     return results
 
 
+_LARGE_STEP_NUMBER = 100000
+
+
 @registry.register_model
 class NextFrameStochastic(NextFrameBasic):
   """Stochastic next-frame model."""
@@ -247,7 +257,7 @@ def conv_lstm_2d(self, inputs, state, output_channels,
         2, input_shape[1:], output_channels,
         [kernel_size, kernel_size], name=scope)
     if state is None:
-      state = cell.zero_state(self.hparams.batch_size, tf.float32)
+      state = cell.zero_state(input_shape[0], tf.float32)
     outputs, new_state = cell(inputs, state)
     return outputs, new_state
 
@@ -295,7 +305,7 @@ def construct_model(self,
       raise ValueError("More than one, or no network option specified.")
 
     img_height, img_width, color_channels = self.hparams.problem.frame_shape
-    batch_size = self.hparams.batch_size
+    batch_size = common_layers.shape_list(images[0])[0]
 
     # Predicted images and rewards.
     gen_rewards, gen_images = [], []
@@ -305,7 +315,10 @@ def construct_model(self,
     else:
       # Scheduled sampling:
       # Calculate number of ground-truth frames to pass in.
-      iter_num = tf.train.get_or_create_global_step()
+      iter_num = tf.train.get_global_step()
+      # TODO(mbz): what should it be if it's undefined?
+      if iter_num is None:
+        iter_num = _LARGE_STEP_NUMBER
       num_ground_truth = tf.to_int32(
           tf.round(
               tf.to_float(batch_size) *
@@ -624,7 +637,7 @@ def get_input_if_exists(self, features, key, batch_size, num_frames):
 
   def body(self, features):
     hparams = self.hparams
-    batch_size = self.hparams.batch_size
+    batch_size = common_layers.shape_list(features["inputs"])[0]
 
     # Split inputs and targets time-wise into a list of frames.
     input_frames = tf.unstack(features["inputs"], axis=1)
@@ -656,7 +669,10 @@ def body(self, features):
         dna=False,
         context_frames=hparams.video_num_input_frames)
 
-    step_num = tf.train.get_or_create_global_step()
+    step_num = tf.train.get_global_step()
+    # TODO(mbz): what should it be if it's undefined?
+    if step_num is None:
+      step_num = _LARGE_STEP_NUMBER
     beta = tf.cond(step_num > self.hparams.num_iterations_2nd_stage,
                    lambda: self.hparams.latent_loss_multiplier,
                    lambda: 0.0)
@@ -716,12 +732,13 @@ def next_frame_stochastic():
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 1
   hparams.batch_size = 8
-  hparams.target_modality = "video:raw"
-  hparams.input_modalities = "inputs:video:raw"
+  hparams.target_modality = "video:l2raw"
+  hparams.input_modalities = "inputs:video:l2raw"
+  hparams.video_modality_loss_cutoff = 0.0
   hparams.add_hparam("stochastic_model", True)
   hparams.add_hparam("latent_channels", 1)
   hparams.add_hparam("latent_std_min", -5.0)
-  hparams.add_hparam("num_iterations_2nd_stage", 10000)
+  hparams.add_hparam("num_iterations_2nd_stage", 50000)
   hparams.add_hparam("latent_loss_multiplier", 1e-4)
   hparams.add_hparam("multi_latent", False)
   hparams.add_hparam("relu_shift", 1e-12)
@@ -729,6 +746,14 @@ def next_frame_stochastic():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_stochastic_cutoff():
+  """SV2P model with additional cutoff in L2 loss for environments like pong."""
+  hparams = next_frame_stochastic()
+  hparams.video_modality_loss_cutoff = 0.4
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_tpu():
   hparams = next_frame()
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index c5287aced..1c460f407 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -139,9 +139,8 @@ def simulate(self, action):
 
       observ = tf.to_float(tf.squeeze(model_output["targets"], axis=1))
 
-      reward = (tf.squeeze(model_output["target_reward"], axis=[1, 2, 3]) +
-                self._min_reward)
-      reward = tf.reshape(tf.to_float(reward), shape=(self.length,))
+      reward = tf.to_float(model_output["target_reward"])
+      reward = tf.reshape(reward, shape=(self.length,)) + self._min_reward
 
       if self._intrinsic_reward_scale:
         # Use the model's uncertainty about its prediction as an intrinsic
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 931ba0caf..7e55ec89a 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -490,6 +490,15 @@ def rl_modelrl_base():
   )
 
 
+@registry.register_hparams
+def rl_modelrl_base_stochastic():
+  """Base setting with a stochastic next-frame model."""
+  hparams = rl_modelrl_base()
+  hparams.generative_model = "next_frame_stochastic"
+  hparams.generative_model_params = "next_frame_stochastic_cutoff"
+  return hparams
+
+
 @registry.register_hparams
 def rl_modelrl_medium():
   """Small set for larger testing."""
@@ -532,6 +541,17 @@ def rl_modelrl_tiny():
       ).values())
 
 
+@registry.register_hparams
+def rl_modelrl_tiny_stochastic():
+  """Tiny setting with a stochastic next-frame model."""
+  hparams = rl_modelrl_tiny()
+  hparams.generative_model = "next_frame_stochastic"
+  hparams.generative_model_params = "next_frame_stochastic_cutoff"
+  hparams.true_env_generator_num_steps = 1000
+  hparams.simulated_env_generator_num_steps = 1000
+  return hparams
+
+
 @registry.register_hparams
 def rl_modelrl_l1_base():
   """Parameter set with L1 loss."""
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index ffb64dc7b..4f4b923c3 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -97,7 +97,7 @@ def train(hparams, environment_spec, event_dir=None, model_dir=None,
 
     if hparams.simulated_environment:
       env_model_loader = tf.train.Saver(
-          tf.global_variables("next_frame_basic.*"))
+          tf.global_variables("next_frame*"))
     else:
       env_model_loader = None
 

From f91e272ba39c262169485f2d8cb5c521d1c4d03a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 14 Jun 2018 13:54:21 -0700
Subject: [PATCH 0104/2720] Cleanup of transformer_sketch.

PiperOrigin-RevId: 200612905
---
 .../models/research/transformer_sketch.py     | 149 +++---------------
 1 file changed, 24 insertions(+), 125 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_sketch.py b/tensor2tensor/models/research/transformer_sketch.py
index bbe27d077..b5c6b2ac2 100644
--- a/tensor2tensor/models/research/transformer_sketch.py
+++ b/tensor2tensor/models/research/transformer_sketch.py
@@ -18,145 +18,44 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
-from tensor2tensor.models.research import transformer_vae
-from tensor2tensor.models.transformer import transformer_base
-from tensor2tensor.models.transformer import transformer_small
 from tensor2tensor.utils import registry
 
+import tensorflow as tf
+
 
 @registry.register_model
 class TransformerSketch(transformer.Transformer):
   """Transformer with strided convolutions."""
 
   def encode(self, inputs, target_space, hparams, features=None, losses=None):
-    """Add two layers strided convolutions ontop of encode."""
-    inputs = common_layers.conv_block(
-        inputs,
-        hparams.hidden_size, [((1, 1), (3, 3))],
-        first_relu=False,
-        padding="SAME",
-        force2d=True,
-        name="small_image_conv")
-
-    hparams.num_compress_steps = 2
-    compressed_inputs = transformer_vae.compress(inputs, None, is_2d=True,
-                                                 hparams=hparams,
-                                                 name="convolutions")
-
-    return super(TransformerSketch, self).encode(
-        compressed_inputs, target_space, hparams, features=features,
-        losses=losses)
+    """Add layers of strided convolutions on top of encoder."""
+    with tf.variable_scope("downstride"):
+      hparams = self.hparams
+      kernel, strides = (4, 4), (2, 2)
+      x = inputs
+      # Down-convolutions.
+      for i in range(hparams.num_compress_steps):
+        x = common_layers.make_even_size(x)
+        x = tf.layers.conv2d(
+            x, hparams.hidden_size, kernel, strides=strides,
+            padding="SAME", activation=common_layers.belu, name="conv_%d" % i)
+        x = common_layers.layer_norm(x)
+
+    encoder_output, encoder_decoder_attention_bias = super(
+        TransformerSketch, self).encode(
+            x, target_space, hparams, features=features, losses=losses)
+    return encoder_output, encoder_decoder_attention_bias
 
 
 @registry.register_hparams
 def transformer_sketch():
   """Basic transformer_sketch hparams."""
-  hparams = transformer_base()
-  hparams.batch_size = 2048
-  hparams.max_length = 784
-  hparams.clip_grad_norm = 5.
-  hparams.learning_rate_decay_scheme = "noam"
-  hparams.learning_rate = 0.2
-  hparams.learning_rate_warmup_steps = 10000
-  hparams.num_hidden_layers = 6
-  # hparams.initializer = "orthogonal"
-  hparams.sampling_method = "random"
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_sketch():
-  """Parameters based on base."""
-  hparams = transformer_base()
-  hparams.batch_size = 2048
-  hparams.max_length = 784
-  hparams.clip_grad_norm = 5.
-  hparams.learning_rate_decay_scheme = "noam"
-  hparams.learning_rate_warmup_steps = 8000
-  hparams.learning_rate = 0.2
-  hparams.num_hidden_layers = 6
-  hparams.initializer = "orthogonal"
-  hparams.sampling_method = "random"
-  return hparams
-
-
-@registry.register_hparams
-def transformer_small_sketch():
-  """Modified transformer_small."""
-  hparams = transformer_small()
-  hparams.batch_size = 2048
-  hparams.max_length = 784
-  hparams.clip_grad_norm = 5.
-  hparams.learning_rate_decay_scheme = "noam"
-  hparams.learning_rate = 0.1
-  hparams.initializer = "orthogonal"
+  hparams = transformer.transformer_small()
+  hparams.num_compress_steps = 4
+  hparams.batch_size = 32
+  hparams.clip_grad_norm = 2.
   hparams.sampling_method = "random"
-  hparams.learning_rate_warmup_steps = 10000
-  return hparams
-
-
-@registry.register_hparams
-def transformer_sketch_2layer():
-  hparams = transformer_sketch()
-  hparams.num_hidden_layers = 2
-  return hparams
-
-
-@registry.register_hparams
-def transformer_sketch_4layer():
-  hparams = transformer_sketch()
-  hparams.num_hidden_layers = 4
-  return hparams
-
-
-@registry.register_hparams
-def transformer_sketch_6layer():
-  hparams = transformer_sketch()
-  hparams.num_hidden_layers = 6
-  return hparams
-
-
-@registry.register_ranged_hparams("transformer_sketch_ranged")
-def transformer_sketch_ranged(rhp):
-  """Range of hparams for vizier."""
-  rhp.set_categorical("ffn_layer",
-                      ["conv_hidden_relu_with_sepconv", "conv_hidden_relu"])
-  rhp.set_discrete("batch_size", [1024, 2048, 4096])
-  rhp.set_discrete("num_hidden_layers", [2, 3, 4, 5, 6])
-  rhp.set_discrete("hidden_size", [32, 64, 128, 256, 512, 1024],
-                   scale=rhp.LOG_SCALE)
-  rhp.set_discrete("kernel_height", [1, 3, 5, 7])
-  rhp.set_discrete("kernel_width", [1, 3, 5, 7])
-  rhp.set_discrete("compress_steps", [0, 1, 2])
-  rhp.set_float("dropout", 0.0, 0.5)
-  rhp.set_float("weight_decay", 1e-4, .03, scale=rhp.LOG_SCALE)
-  rhp.set_float("label_smoothing", 0.0, 0.2)
-  rhp.set_float("clip_grad_norm", 0.01, 8.0, scale=rhp.LOG_SCALE)
-  rhp.set_float("learning_rate", 0.1, 1.0, scale=rhp.LOG_SCALE)
-  rhp.set_categorical("initializer",
-                      ["uniform", "orthogonal", "uniform_unit_scaling"])
-  rhp.set_float("initializer_gain", 0.5, 3.5)
-  rhp.set_categorical("learning_rate_decay_scheme",
-                      ["none", "sqrt", "noam", "exp"])
-  rhp.set_float("optimizer_adam_epsilon", 1e-7, 1e-2, scale=rhp.LOG_SCALE)
-  rhp.set_float("optimizer_adam_beta1", 0.8, 0.9)
-  rhp.set_float("optimizer_adam_beta2", 0.995, 0.999)
-  rhp.set_categorical("optimizer", [
-      "Adam", "Adagrad", "Momentum", "RMSProp", "SGD", "YellowFin"])
-
-
-@registry.register_hparams
-def transformer_opt():
-  """Parameters that work better."""
-  hparams = transformer_sketch()
-  hparams.batch_size = 1024
-  hparams.learning_rate = 0.28
-  hparams.num_hidden_layers = 3
-  hparams.dropout = 0.35
-  hparams.ffn_layer = "conv_hidden_relu_with_sepconv"
-  hparams.hidden_size = 128
-  hparams.initializer_gain = 2.6
-  hparams.weight_decay = 0.
   return hparams

From 0505d1ee8540d5ff0b3f1c7a0aa5b74126dfbdcd Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 14 Jun 2018 16:28:12 -0700
Subject: [PATCH 0105/2720] using tf.greater since step_num can now be a number
 too.

PiperOrigin-RevId: 200638706
---
 tensor2tensor/models/research/next_frame.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 74d62eea8..4ff04ebca 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -673,7 +673,7 @@ def body(self, features):
     # TODO(mbz): what should it be if it's undefined?
     if step_num is None:
       step_num = _LARGE_STEP_NUMBER
-    beta = tf.cond(step_num > self.hparams.num_iterations_2nd_stage,
+    beta = tf.cond(tf.greater(step_num, self.hparams.num_iterations_2nd_stage),
                    lambda: self.hparams.latent_loss_multiplier,
                    lambda: 0.0)
 

From 5c38266e96ab79231c42a71df48501f0531cfa0a Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 14 Jun 2018 16:34:06 -0700
Subject: [PATCH 0106/2720] Update ASR readme to use
 librispeech_train_full_test_clean

PiperOrigin-RevId: 200639752
---
 docs/tutorials/asr_with_transformer.md | 50 ++++++++++++++------------
 1 file changed, 28 insertions(+), 22 deletions(-)

diff --git a/docs/tutorials/asr_with_transformer.md b/docs/tutorials/asr_with_transformer.md
index c1210a9ae..9feb381b3 100644
--- a/docs/tutorials/asr_with_transformer.md
+++ b/docs/tutorials/asr_with_transformer.md
@@ -18,30 +18,17 @@ As the audio import in `t2t-datagen` uses `sox` to generate normalized
 waveforms, please install it as appropriate (e.g. `apt-get install sox`).
 
 ```
+# Generate both the full dataset and the small clean version, which we use for
+# evaluation.
 t2t-datagen --problem=librispeech --data_dir=$DATA_DIR --tmp_dir=$TMP_DIR
+t2t-datagen --problem=librispeech_clean --data_dir=$DATA_DIR --tmp_dir=$TMP_DIR
 ```
 
-You can also use smaller versions of the dataset by replacing `librispeech` with
-`librispeech_clean` or `librispeech_clean_small`
+The problem `librispeech_train_full_test_clean` will train on the full dataset
+but evaluate on the clean dataset.
 
-## Training on GPUs
-
-To train a model on GPU set up`OUT_DIR` and run the trainer:
-
-```
-t2t-trainer \
-  --model=transformer \
-  --hparams_set=transformer_librispeech \
-  --problem=librispeech \
-  --train_steps=120000 \
-  --eval_steps=3 \
-  --local_eval_frequency=100 \
-  --data_dir=$DATA_DIR \
-  --output_dir=$OUT_DIR
-```
-
-This model should achieve approximately 22% accuracy per sequence after
-approximately 80,000 steps.
+You can also use `librispeech_clean_small` which is a small version of the
+clean dataset.
 
 ## Training on Cloud TPUs
 
@@ -53,7 +40,7 @@ t2t-trainer \
   --model=transformer \
   --hparams_set=transformer_librispeech_tpu \
   --hparams=max_length=125550,max_input_seq_length=1550,max_target_seq_length=350,batch_size=16 \
-  --problem=librispeech \
+  --problem=librispeech_train_full_test_clean \
   --train_steps=210000 \
   --eval_steps=3 \
   --local_eval_frequency=100 \
@@ -71,7 +58,7 @@ t2t-trainer \
   --model=transformer \
   --hparams_set=transformer_librispeech_tpu \
   --hparams=max_length=295650,max_input_seq_length=3650,max_target_seq_length=650,batch_size=6 \
-  --problem=librispeech \
+  --problem=librispeech_train_full_test_clean \
   --train_steps=230000 \
   --eval_steps=3 \
   --local_eval_frequency=100 \
@@ -86,3 +73,22 @@ documentation](https://github.com/tensorflow/tensor2tensor/tree/master/docs/clou
 for Tensor2Tensor on Cloud TPUs, or the [official Google Cloud Platform
 documentation](https://cloud.google.com/tpu/docs/tutorials/transformer) for
 Cloud TPUs.
+
+## Training on GPUs
+
+To train a model on GPU set up`OUT_DIR` and run the trainer:
+
+```
+t2t-trainer \
+  --model=transformer \
+  --hparams_set=transformer_librispeech_tpu \
+  --problem=librispeech \
+  --train_steps=120000 \
+  --eval_steps=3 \
+  --local_eval_frequency=100 \
+  --data_dir=$DATA_DIR \
+  --output_dir=$OUT_DIR
+```
+
+This model should achieve approximately 22% accuracy per sequence after
+approximately 80,000 steps.

From 41c16e6d7af69c820575fa38e221af29ef9ffa95 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 14 Jun 2018 16:43:36 -0700
Subject: [PATCH 0107/2720] Lint fix

PiperOrigin-RevId: 200641075
---
 tensor2tensor/layers/modalities.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 11aa106c3..32d245433 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -644,7 +644,7 @@ def bottom(self, x):
     common_layers.summarize_video(x, "inputs")
     return common_layers.convert_rgb_to_real(x)
 
-  def targets_bottom(self, x):
+  def targets_bottom(self, x):  # pylint: disable=arguments-differ
     common_layers.summarize_video(x, "targets_bottom")
     return common_layers.convert_rgb_to_real(x)
 

From 33329655300824ef050d27eff466f23550d9a0e0 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 14 Jun 2018 17:01:13 -0700
Subject: [PATCH 0108/2720] Print which module had an ImportError

PiperOrigin-RevId: 200643604
---
 tensor2tensor/data_generators/all_problems.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 7b8d8137f..6ab99e7f5 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -88,6 +88,7 @@ def _handle_errors(errors):
   for module, err in errors:
     err_str = str(err)
     if err_str != _py_err_msg(module):
+      print("From module %s" % module)
       raise err
     if log_all:
       print("Did not import module: %s; Cause: %s" % (module, err_str))

From 3d7cd00f59abe963e1d9beed0566c919a4ea1195 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 14 Jun 2018 17:11:39 -0700
Subject: [PATCH 0109/2720] Add oauth2client explicitly to deps

PiperOrigin-RevId: 200645296
---
 setup.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/setup.py b/setup.py
index fed6e8c20..47ffba801 100644
--- a/setup.py
+++ b/setup.py
@@ -41,6 +41,7 @@
         'gym',
         'h5py',
         'numpy',
+        'oauth2client',
         'requests',
         'scipy',
         'sympy',

From ffff8ae1a1565471aca5923e8161e5693ff31780 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 14 Jun 2018 18:28:38 -0700
Subject: [PATCH 0110/2720] Fixing video decoding

PiperOrigin-RevId: 200653976
---
 tensor2tensor/models/research/next_frame.py |  5 +--
 tensor2tensor/utils/decoding.py             | 38 +++++++++++++++++++--
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 4ff04ebca..ed789081e 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -139,8 +139,9 @@ def logits_to_samples(logits):
       targets_shape = [self.hparams.batch_size,
                        self.hparams.video_num_target_frames, 1, 1, num_channels]
     features["targets"] = tf.zeros(targets_shape, dtype=tf.int32)
-    features["target_reward"] = tf.zeros(
-        [targets_shape[0], 1, 1], dtype=tf.int32)
+    if "target_reward" in self.hparams.problem_hparams.target_modality:
+      features["target_reward"] = tf.zeros(
+          [targets_shape[0], 1, 1], dtype=tf.int32)
     logits, _ = self(features)  # pylint: disable=not-callable
     if isinstance(logits, dict):
       results = {}
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 1d79f8483..bbc68fbe0 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -68,6 +68,22 @@ def log_decode_results(inputs,
                        identity_output=False,
                        log_results=True):
   """Log inference results."""
+
+  # TODO(lukaszkaiser) refactor this into feature_encoder
+  is_video = "video" in problem_name
+  if is_video:
+    def fix_and_save_video(vid, prefix):
+      save_path_template = os.path.join(
+          model_dir, "%s_%s_%d_{}.png" % (problem_name, prefix, prediction_idx))
+      # this is only required for predictions
+      if vid.shape[-1] == 1:
+        vid = np.squeeze(vid, axis=-1)
+      save_video(vid, save_path_template)
+    tf.logging.info("Saving video: {}".format(prediction_idx))
+    fix_and_save_video(inputs, "inputs")
+    fix_and_save_video(outputs, "outputs")
+    fix_and_save_video(targets, "targets")
+
   is_image = "image" in problem_name
   decoded_inputs = None
   if is_image and save_images:
@@ -80,7 +96,7 @@ def log_decode_results(inputs,
     else:
       decoded_inputs = inputs_vocab.decode(_save_until_eos(inputs, is_image))
 
-    if log_results:
+    if log_results and not is_video:
       tf.logging.info("Inference results INPUT: %s" % decoded_inputs)
 
   decoded_targets = None
@@ -93,8 +109,9 @@ def log_decode_results(inputs,
     decoded_outputs = targets_vocab.decode(_save_until_eos(outputs, is_image))
     if targets is not None and log_results:
       decoded_targets = targets_vocab.decode(_save_until_eos(targets, is_image))
-  tf.logging.info("Inference results OUTPUT: %s" % decoded_outputs)
-  if targets is not None and log_results:
+  if not is_video:
+    tf.logging.info("Inference results OUTPUT: %s" % decoded_outputs)
+  if targets is not None and log_results and not is_video:
     tf.logging.info("Inference results TARGET: %s" % decoded_targets)
   return decoded_inputs, decoded_outputs, decoded_targets
 
@@ -518,6 +535,21 @@ def _interactive_input_fn(hparams, decode_hp):
       yield features
 
 
+def save_video(video, save_path_template):
+  """Save frames of the videos into files."""
+  try:
+    from PIL import Image  # pylint: disable=g-import-not-at-top
+  except ImportError as e:
+    tf.logging.warning(
+        "Showing and saving an image requires PIL library to be "
+        "installed: %s", e)
+    raise NotImplementedError("Image display and save not implemented.")
+
+  for i, frame in enumerate(video):
+    save_path = save_path_template.format(i)
+    Image.fromarray(np.uint8(frame)).save(save_path)
+
+
 def show_and_save_image(img, save_path):
   try:
     import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top

From 3704b1f674c05d3d5b188e6004a91468e3d0c583 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 14 Jun 2018 21:23:33 -0700
Subject: [PATCH 0111/2720] Correct inference to run with RealModality for time
 series problems.

PiperOrigin-RevId: 200667028
---
 tensor2tensor/data_generators/text_encoder.py |  2 +-
 tensor2tensor/layers/modalities.py            |  4 ++
 tensor2tensor/models/transformer.py           |  3 ++
 tensor2tensor/utils/t2t_model.py              | 47 +++++++++++++++----
 4 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 4a37e69c4..8695aa816 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -1051,7 +1051,7 @@ def decode(self, ids, strip_extraneous=False):
       ValueError: if the ids are not of the appropriate size.
     """
     del strip_extraneous
-    return " ".join(ids)
+    return " ".join([str(i) for i in ids])
 
 
 def strip_ids(ids, ids_to_strip):
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 32d245433..bcfc9b80d 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -739,6 +739,10 @@ class RealModality(modality.Modality):
   * Top is a linear projection layer to vocab_size.
   """
 
+  @property
+  def top_is_pointwise(self):
+    return True
+
   def bottom(self, x):
     with tf.variable_scope("real"):
       return tf.layers.dense(tf.to_float(x), self._body_input_depth,
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index c605dc1ad..b96ae3f80 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -213,6 +213,9 @@ def _greedy_infer(self, features, decode_length):
     Raises:
       NotImplementedError: If there are multiple data shards.
     """
+    # For real-valued modalities use the slow decode path for now.
+    if self._target_modality_is_real:
+      return  super(Transformer, self)._greedy_infer(features, decode_length)
     with tf.variable_scope(self.name):
       return self._fast_decode(features, decode_length)
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 36fd72e8d..04a6dba8c 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -141,6 +141,12 @@ def _custom_getter(self):
     else:
       return None
 
+  @property
+  def _target_modality_is_real(self):
+    """Whether the target modality is real-valued."""
+    target_modality = self._problem_hparams.target_modality
+    return target_modality.name.startswith("real_")
+
   def call(self, inputs, **kwargs):
     del kwargs
     features = inputs
@@ -732,7 +738,11 @@ def _slow_greedy_infer(self, features, decode_length):
     def infer_step(recent_output, recent_logits, unused_loss):
       """Inference step."""
       if not tf.contrib.eager.in_eager_mode():
-        recent_output.set_shape([None, None, None, 1])
+        if self._target_modality_is_real:
+          dim = self._problem_hparams.target_modality.top_dimensionality
+          recent_output.set_shape([None, None, None, dim])
+        else:
+          recent_output.set_shape([None, None, None, 1])
       padded = tf.pad(recent_output, [[0, 0], [0, 1], [0, 0], [0, 0]])
       features["targets"] = padded
       # This is inefficient in that it generates samples at all timesteps,
@@ -745,10 +755,14 @@ def infer_step(recent_output, recent_logits, unused_loss):
       else:
         cur_sample = samples[:,
                              common_layers.shape_list(recent_output)[1], :, :]
-      cur_sample = tf.to_int64(tf.expand_dims(cur_sample, axis=1))
-      samples = tf.concat([recent_output, cur_sample], axis=1)
-      if not tf.contrib.eager.in_eager_mode():
-        samples.set_shape([None, None, None, 1])
+      if self._target_modality_is_real:
+        cur_sample = tf.expand_dims(cur_sample, axis=1)
+        samples = tf.concat([recent_output, cur_sample], axis=1)
+      else:
+        cur_sample = tf.to_int64(tf.expand_dims(cur_sample, axis=1))
+        samples = tf.concat([recent_output, cur_sample], axis=1)
+        if not tf.contrib.eager.in_eager_mode():
+          samples.set_shape([None, None, None, 1])
 
       # Assuming we have one shard for logits.
       logits = tf.concat([recent_logits, logits[:, -1:]], 1)
@@ -764,7 +778,11 @@ def infer_step(recent_output, recent_logits, unused_loss):
       batch_size = common_layers.shape_list(initial_output)[0]
     else:
       batch_size = common_layers.shape_list(features["inputs"])[0]
-      initial_output = tf.zeros((batch_size, 0, 1, 1), dtype=tf.int64)
+      if self._target_modality_is_real:
+        dim = self._problem_hparams.target_modality.top_dimensionality
+        initial_output = tf.zeros((batch_size, 0, 1, dim), dtype=tf.float32)
+      else:
+        initial_output = tf.zeros((batch_size, 0, 1, 1), dtype=tf.int64)
     # Hack: foldl complains when the output shape is less specified than the
     # input shape, so we confuse it about the input shape.
     initial_output = tf.slice(initial_output, [0, 0, 0, 0],
@@ -783,10 +801,17 @@ def infer_step(recent_output, recent_logits, unused_loss):
 
     # Initial values of result, logits and loss.
     result = initial_output
-    # tensor of shape [batch_size, time, 1, 1, vocab_size]
-    logits = tf.zeros((batch_size, 0, 1, 1, target_modality.top_dimensionality))
+    if self._target_modality_is_real:
+      logits = tf.zeros((batch_size, 0, 1, target_modality.top_dimensionality))
+      logits_shape_inv = [None, None, None, None]
+    else:
+      # tensor of shape [batch_size, time, 1, 1, vocab_size]
+      logits = tf.zeros((batch_size, 0, 1, 1,
+                         target_modality.top_dimensionality))
+      logits_shape_inv = [None, None, None, None, None]
     if not tf.contrib.eager.in_eager_mode():
-      logits.set_shape([None, None, None, None, None])
+      logits.set_shape(logits_shape_inv)
+
     loss = 0.0
 
     def while_exit_cond(result, logits, loss):  # pylint: disable=unused-argument
@@ -822,7 +847,7 @@ def fn_not_eos():
         infer_step, [result, logits, loss],
         shape_invariants=[
             tf.TensorShape([None, None, None, None]),
-            tf.TensorShape([None, None, None, None, None]),
+            tf.TensorShape(logits_shape_inv),
             tf.TensorShape([]),
         ],
         back_prop=False,
@@ -857,6 +882,8 @@ def sample(self, features):
        losses: a dictionary: {loss-name (string): floating point `Scalar`}.
     """
     logits, losses = self(features)  # pylint: disable=not-callable
+    if self._target_modality_is_real:
+      return logits, logits, losses  # Raw numbers returned from real modality.
     if self.hparams.sampling_method == "argmax":
       samples = tf.argmax(logits, axis=-1)
     else:

From 49f7f586595212b29c5ca699ff60b70d25bc20ba Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 15 Jun 2018 12:54:30 -0700
Subject: [PATCH 0112/2720] v1.6.4

PiperOrigin-RevId: 200759306
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 47ffba801..93892429f 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.6.3',
+    version='1.6.4',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',

From 9598fa2009b7f587a299a4ab36b8a2f01fced262 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Fri, 15 Jun 2018 13:11:23 -0700
Subject: [PATCH 0113/2720] Add word dropout to transformer_vae

PiperOrigin-RevId: 200761864
---
 tensor2tensor/models/research/transformer_vae.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index dcf19b5d6..5300968cf 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -360,7 +360,14 @@ def ae_transformer_internal(inputs,
     targets, _ = common_layers.pad_to_same_length(
         targets, max_targets_len_from_inputs,
         final_length_divisible_by=2**hparams.num_compress_steps)
-    targets_c = compress(targets, inputs, False, hparams, "compress")
+    if hparams.word_dropout:
+      mask = tf.random_uniform(shape=common_layers.shape_list(targets),
+                               minval=0.0, maxval=1.0)
+      targets_noisy = tf.where(mask > hparams.word_dropout, targets,
+                               tf.zeros_like(targets))
+    else:
+      targets_noisy = targets
+    targets_c = compress(targets_noisy, inputs, False, hparams, "compress")
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
       # Compress and bottleneck.
       latents_dense, latents_discrete, extra_loss, embed = hparams.bottleneck(
@@ -668,6 +675,7 @@ def transformer_ae_small():
   hparams.add_hparam("noise_dev", 0.5)
   hparams.add_hparam("d_mix", 0.5)
   hparams.add_hparam("logit_normalization", True)
+  hparams.add_hparam("word_dropout", 0.1)
   # Bottleneck kinds supported: dense, vae, semhash, gumbel-softmax, dvq.
   hparams.add_hparam("bottleneck_kind", "semhash")
   hparams.add_hparam("num_blocks", 1)

From 6fb153755420cba66f812d22572a5ed5d41c078b Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 15 Jun 2018 13:18:16 -0700
Subject: [PATCH 0114/2720] Lazy load PIL

PiperOrigin-RevId: 200762820
---
 tensor2tensor/data_generators/bair_robot_pushing.py   |  9 +++++++--
 tensor2tensor/data_generators/google_robot_pushing.py | 11 ++++++++---
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index 498a714ed..cd4cf3002 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -27,7 +27,6 @@
 import os
 import tarfile
 import numpy as np
-from PIL import Image
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -40,6 +39,12 @@
     "http://rail.eecs.berkeley.edu/datasets/bair_robot_pushing_dataset_v0.tar")
 
 
+# Lazy load PIL.Image
+def PIL_Image():  # pylint: disable=invalid-name
+  from PIL import Image  # pylint: disable=g-import-not-at-top
+  return Image
+
+
 @registry.register_problem
 class VideoBairRobotPushing(video_utils.VideoProblem):
   """Berkeley (BAIR) robot pushing dataset."""
@@ -85,7 +90,7 @@ def parse_frames(self, filenames):
           state_name = state_key.format(i)
 
           byte_str = x.features.feature[image_name].bytes_list.value[0]
-          img = Image.frombytes(
+          img = PIL_Image().frombytes(
               "RGB", (self.frame_width, self.frame_height), byte_str)
           arr = np.array(img.getdata())
           frame = arr.reshape(
diff --git a/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index d2e8011ca..75bafcedf 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -27,7 +27,6 @@
 import io
 import os
 import numpy as np
-from PIL import Image
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -42,6 +41,12 @@
 DATA_TEST_NOVEL = (5, "/push_testnovel/push_testnovel.tfrecord-{:05d}-of-00005")
 
 
+# Lazy load PIL.Image
+def PIL_Image():  # pylint: disable=invalid-name
+  from PIL import Image  # pylint: disable=g-import-not-at-top
+  return Image
+
+
 @registry.register_problem
 class VideoGoogleRobotPushing(video_utils.VideoProblem):
   """Google robot pushing dataset."""
@@ -90,10 +95,10 @@ def parse_frames(self, filename):
         state_name = state_key.format(i)
 
         byte_str = x.features.feature[image_name].bytes_list.value[0]
-        img = Image.open(io.BytesIO(byte_str))
+        img = PIL_Image().open(io.BytesIO(byte_str))
         # The original images are much bigger than 64x64
         img = img.resize((self.frame_width, self.frame_height),
-                         resample=Image.BILINEAR)
+                         resample=PIL_Image().BILINEAR)
         arr = np.array(img.getdata())
         frame = arr.reshape(
             self.frame_width, self.frame_height, self.num_channels)

From e96ca3a3457c63fc51880d92770bf24fdd1f19ad Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Fri, 15 Jun 2018 13:27:06 -0700
Subject: [PATCH 0115/2720] Bring back latent pred masking as it is somewhat
 unstable without it

PiperOrigin-RevId: 200764166
---
 tensor2tensor/models/research/transformer_nat.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 0a42b4b24..70e7075d0 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -299,6 +299,9 @@ def ae_transformer_internal(inputs, targets, target_space, hparams, cache=None):
 
   targets = d
   res = decode_transformer(inputs, ed, d, hparams, "decoder")
+  latent_time = tf.less(hparams.mask_startup_steps,
+                        tf.to_int32(tf.train.get_global_step()))
+  losses["latent_pred"] *= tf.to_float(latent_time)
   return res, losses, cache
 
 
@@ -385,7 +388,7 @@ def transformer_nat_small():
   hparams.optimizer = "Adam"
   hparams.optimizer_adam_epsilon = 1e-9
   hparams.optimizer_adam_beta1 = 0.9
-  hparams.optimizer_adam_beta2 = 0.997  # Needs tuning, try 0.98 to 0.999.
+  hparams.optimizer_adam_beta2 = 0.997
   hparams.add_hparam("bottleneck_kind", "vq")
   hparams.add_hparam("bottleneck_bits", 12)
   hparams.add_hparam("num_compress_steps", 3)
@@ -393,6 +396,7 @@ def transformer_nat_small():
   hparams.add_hparam("epsilon", 1e-5)
   hparams.add_hparam("decay", 0.999)
   hparams.add_hparam("num_samples", 10)
+  hparams.add_hparam("mask_startup_steps", 50000)
   return hparams
 
 

From 233fdf489634be333e2cd8d401e77304e5342590 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 15 Jun 2018 13:28:46 -0700
Subject: [PATCH 0116/2720] Adding tests for next_frame basic training step.

PiperOrigin-RevId: 200764433
---
 .../data_generators/video_generated.py        |   4 +-
 .../models/research/next_frame_test.py        | 104 ++++++++++++++++++
 2 files changed, 106 insertions(+), 2 deletions(-)
 create mode 100644 tensor2tensor/models/research/next_frame_test.py

diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 2ce9c318b..eef11cc38 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -95,10 +95,10 @@ def extra_reading_spec(self):
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.input_modality = {
-        "inputs": ("video:raw", 256),
+        "inputs": ("video", 256),
         "input_frame_number": ("symbol:identity", 1)
     }
-    p.target_modality = ("video:raw", 256)
+    p.target_modality = ("video", 256)
 
   @staticmethod
   def get_circle(x, y, z, c, s):
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
new file mode 100644
index 000000000..6284fbb0c
--- /dev/null
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -0,0 +1,104 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Basic tests for video prediction models."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import numpy as np
+
+from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
+from tensor2tensor.models.research import next_frame
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+class NextFrameTest(tf.test.TestCase):
+
+  def TestVideoModel(self,
+                     in_frames,
+                     out_frames,
+                     hparams,
+                     model,
+                     expected_last_dim):
+
+    x = np.random.random_integers(0, high=255, size=(8, in_frames, 64, 64, 3))
+    y = np.random.random_integers(0, high=255, size=(8, out_frames, 64, 64, 3))
+
+    hparams.video_num_input_frames = in_frames
+    hparams.video_num_target_frames = out_frames
+
+    problem = registry.problem("video_stochastic_shapes10k")
+    p_hparams = problem.get_hparams(hparams)
+    hparams.problem = problem
+    hparams.problem_hparams = p_hparams
+
+    with self.test_session() as session:
+      features = {
+          "inputs": tf.constant(x, dtype=tf.int32),
+          "targets": tf.constant(y, dtype=tf.int32),
+      }
+      model = model(
+          hparams, tf.estimator.ModeKeys.TRAIN)
+      logits, _ = model(features)
+      session.run(tf.global_variables_initializer())
+      res = session.run(logits)
+    expected_shape = y.shape + (expected_last_dim,)
+    self.assertEqual(res.shape, expected_shape)
+
+  def TestBasicModel(self, in_frames, out_frames):
+    self.TestVideoModel(
+        in_frames,
+        out_frames,
+        next_frame.next_frame(),
+        next_frame.NextFrameBasic,
+        256)
+
+  def testBasicModelSingleInputFrameSingleOutputFrames(self):
+    self.TestBasicModel(1, 1)
+
+  def testBasicModelSingleInputFrameMultiOutputFrames(self):
+    self.TestBasicModel(1, 6)
+
+  def testBasicModelMultiInputFrameSingleOutputFrames(self):
+    self.TestBasicModel(4, 1)
+
+  def testBasicModelMultiInputFrameMultiOutputFrames(self):
+    self.TestBasicModel(7, 5)
+
+  def TestStochasticModel(self, in_frames, out_frames):
+    self.TestVideoModel(
+        in_frames,
+        out_frames,
+        next_frame.next_frame_stochastic(),
+        next_frame.NextFrameStochastic,
+        1)
+
+  def testStochasticModelSingleInputFrameSingleOutputFrames(self):
+    self.TestStochasticModel(1, 1)
+
+  def testStochasticModelSingleInputFrameMultiOutputFrames(self):
+    self.TestStochasticModel(1, 6)
+
+  def testStochasticModelMultiInputFrameSingleOutputFrames(self):
+    self.TestStochasticModel(4, 1)
+
+  def testStochasticModelMultiInputFrameMultiOutputFrames(self):
+    self.TestStochasticModel(7, 5)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 76330b20e40b55bc571f0b2adca9eaee870548fd Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 15 Jun 2018 13:33:15 -0700
Subject: [PATCH 0117/2720] Update Travis to first check base dependencies

PiperOrigin-RevId: 200765178
---
 .travis.yml | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/.travis.yml b/.travis.yml
index 59065472f..ff0811c23 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -36,6 +36,10 @@ before_install:
   # - sudo apt-get install -qq tensorflow-model-server
 install:
   - pip install -q "tensorflow==$TF_VERSION"
+  # First ensure that the base dependencies are sufficient for a full import
+  - pip install -q .
+  - t2t-trainer --registry_help
+  # Then install the test dependencies
   - pip install -q .[tests]
   # Make sure to install the atari extras for gym
   - pip install "gym[atari]"

From 49e279eb6c871fbebc137d6f598758a275f521c3 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 15 Jun 2018 13:43:18 -0700
Subject: [PATCH 0118/2720] v1.6.5

PiperOrigin-RevId: 200766741
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 93892429f..aa1d3622c 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.6.4',
+    version='1.6.5',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',

From abd502722246fc38b009d8085bd9601ef09608b2 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 15 Jun 2018 15:57:20 -0700
Subject: [PATCH 0119/2720] Clean up and correct docs for DMOL loss and
 sampling.

PiperOrigin-RevId: 200788010
---
 tensor2tensor/layers/common_layers.py      | 200 +++++++++++++--------
 tensor2tensor/layers/common_layers_test.py |  65 +++++++
 tensor2tensor/models/image_transformer.py  |   2 +-
 3 files changed, 187 insertions(+), 80 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 17f8d3322..6afc4039c 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1887,64 +1887,104 @@ def dml_loss(
                                       tf.float32)
 
 
+def split_to_discretized_mix_logistic_params(inputs):
+  """Splits input tensor into parameters of discretized mixture logistic.
+
+  Args:
+    inputs: A [batch, height, width, num_mixtures*10] tensor of floats
+      comprising one unconstrained mixture probability, three means
+      (one per channel), three standard deviations (one per channel),
+      and three coefficients which linearly parameterize dependence across
+      channels.
+
+  Returns:
+    Tuple of unconstrained mixture probabilities, locations, scales, and
+    coefficient parameters of the distribution. The mixture probability has
+    shape [batch, height, width, num_mixtures]. Other parameters have shape
+    [batch, height, width, num_mixtures, 3].
+  """
+  batch, height, width, output_dim = shape_list(inputs)
+  num_mixtures = output_dim // 10
+  logits, locs, log_scales, coeffs = tf.split(
+      inputs,
+      num_or_size_splits=[num_mixtures, num_mixtures * 3,
+                          num_mixtures * 3, num_mixtures * 3],
+      axis=-1)
+  split_shape = [batch, height, width, num_mixtures, 3]
+  locs = tf.reshape(locs, split_shape)
+  log_scales = tf.reshape(log_scales, split_shape)
+  log_scales = tf.maximum(log_scales, -7.)
+  coeffs = tf.reshape(coeffs, split_shape)
+  coeffs = tf.tanh(coeffs)
+  return logits, locs, log_scales, coeffs
+
+
 def discretized_mix_logistic_loss(labels, pred, sum_all=True):
-  """Computes log likelihood for the discretized mixture of logistics loss.
+  """Computes negative log probability for the discretized mixture of logistics.
+
+  The distribution of a whole pixel is a mixture of 3-dimensional discretized
+  logistic distributions. The 3-D discretized logistic factorizes as 3 1-D
+  discretized logistic distributions, one for each channel. It defines
+
+  ```none
+  P(X = x)
+  = sum_{k=1}^K probs[k] * P(X = x | locs[k], scales[k])
+  = sum_{k=1}^K probs[k] * [
+      prod_{c=1}^3 DiscretizedLogistic(X[c] = x[c] | means[k][c], scales[k]) ]
+  ```
+
+  The means tensor is a linear combination of location parameters and previous
+  channels. The discretized logistic distribution assigns probability mass to an
+  event P(X=x) via logistic CDFs: P(X <= x + 0.5) - P(X > x - 0.5) for 1 < x <
+  254; P(X <= 0.5) for x = 0; and 1 - P(X > 245.5) for x = 255. Instead of
+  8-bit inputs, this implementation assumes the events are rescaled to [-1, 1].
 
   Args:
     labels: A [batch, height, width, channels] tensor of true pixel intensities
       rescaled to [-1, 1]. The computation assumes channels is 3.
-    pred: A [batch, height, width, hparams.num_mixtures*10] tensor of floats
-      comprising one unnormalized mixture probability, three means
+    pred: A [batch, height, width, num_mixtures*10] tensor of floats
+      comprising one unconstrained mixture probability, three means
       (one per channel), three standard deviations (one per channel),
       and three coefficients which linearly parameterize dependence across
-      channels
-    sum_all: A boolean to return scalar mean loss or per position
+      channels.
+    sum_all: A boolean to return scalar mean loss or per position.
 
   Returns:
-    per_position_loss: A [batch, height, width, 3] tensor of the
-      conditional probability of each channel given all previous channels if
-      not sum_all else add all the losses (for eval)
+    A [batch, height, width] tensor of the negative log conditional probability
+    of each pixel given all previous pixels if not sum_all else add all the
+    losses (for eval).
   """
 
-  # Extract mixture probabilities, means, scale, and coefficient parameters.
-  l_shape = shape_list(labels)
-  num_mix = shape_list(pred)[-1] // 10
-  logits = pred[:, :, :, :num_mix]  # unnormalized mixture probabilities
-  pred = tf.reshape(
-      pred[:, :, :, num_mix:],
-      [l_shape[0], l_shape[1], l_shape[2], l_shape[3], num_mix * 3])
-  means = pred[:, :, :, :, :num_mix]
-  log_scales = tf.maximum(pred[:, :, :, :, num_mix:2 * num_mix], -7.)
-  coeffs = tf.tanh(pred[:, :, :, :, 2 * num_mix:3 * num_mix])
-  labels = (  # tile labels across number of mixtures
-      tf.reshape(labels, [l_shape[0], l_shape[1], l_shape[2], l_shape[3], 1]) +
-      tf.zeros([l_shape[0], l_shape[1], l_shape[2], l_shape[3], num_mix])
-  )
+  logits, locs, log_scales, coeffs = split_to_discretized_mix_logistic_params(
+      pred)
+
+  # Tile labels to broadcast compute across the mixture dimension.
+  batch, height, width, num_mixtures = shape_list(logits)
+  labels = tf.tile(tf.reshape(labels, [batch, height, width, 1, 3]),
+                   [1, 1, 1, num_mixtures, 1])
 
   # p(x) = sigmoid((x - means_i + 1/255.)/scale_i) -
   #        sigmoid((x - means_i - 1/255.)/scale_i)
   # for each channel i. The means are linearly parameterized.
-  means_0 = tf.reshape(means[:, :, :, 0, :],
-                       [l_shape[0], l_shape[1], l_shape[2], 1, num_mix])
-  means_1 = tf.reshape(means[:, :, :, 1, :] +
-                       coeffs[:, :, :, 0, :] * labels[:, :, :, 0, :],
-                       [l_shape[0], l_shape[1], l_shape[2], 1, num_mix])
-  means_2 = tf.reshape(means[:, :, :, 2, :] +
-                       coeffs[:, :, :, 1, :] * labels[:, :, :, 0, :] +
-                       coeffs[:, :, :, 2, :] * labels[:, :, :, 1, :],
-                       [l_shape[0], l_shape[1], l_shape[2], 1, num_mix])
-  means = tf.concat([means_0, means_1, means_2], 3)
+  means_0 = locs[..., 0]
+  means_1 = locs[..., 1] + coeffs[..., 0] * labels[..., 0]
+  means_2 = (locs[..., 2] +
+             coeffs[..., 1] * labels[..., 0] +
+             coeffs[..., 2] * labels[..., 1])
+  means = tf.stack([means_0, means_1, means_2], axis=-1)
   centered_labels = labels - means
   inv_stdv = tf.exp(-log_scales)
   plus_in = inv_stdv * (centered_labels + 1. / 255.)
   min_in = inv_stdv * (centered_labels - 1. / 255.)
   cdf_plus = tf.nn.sigmoid(plus_in)
   cdf_min = tf.nn.sigmoid(min_in)
+
   # Compute log probability for edge case of 0 (before scaling), 255 (before
   # scaling), and all other cases respectively.
-  prob_0 = plus_in - tf.nn.softplus(plus_in)
-  prob_255 = -tf.nn.softplus(min_in)
-  prob_event = cdf_plus - cdf_min
+  log_prob_0 = plus_in - tf.nn.softplus(plus_in)
+  log_prob_255 = -tf.nn.softplus(min_in)
+  prob_event = tf.maximum(cdf_plus - cdf_min, 1e-12)
+  log_prob_event = tf.log(prob_event)
 
   # Robustly select log-prob based on numerical edge-cases: (a) [-1, -1+eps);
   # (b) (1-eps, 1]; (c) NaNs during `tf.gradients` of `tf.select`, which may
@@ -1952,63 +1992,65 @@ def discretized_mix_logistic_loss(labels, pred, sum_all=True):
   mid_in = inv_stdv * centered_labels
   log_prob_event_approx = (
       mid_in - log_scales - 2. * tf.nn.softplus(mid_in) - np.log(127.5))
-  log_probs = tf.where(labels < -0.999, prob_0,
-                       tf.where(labels > 0.999, prob_255,
-                                tf.where(prob_event > 1e-5,
-                                         tf.log(tf.maximum(prob_event, 1e-12)),
+  log_probs = tf.where(labels < -0.999, log_prob_0,
+                       tf.where(labels > 0.999, log_prob_255,
+                                tf.where(prob_event > 1e-5, log_prob_event,
                                          log_prob_event_approx)))
-  # Sum over mixtures.
-  log_probs = tf.reduce_sum(log_probs, 3) + tf.nn.log_softmax(logits, axis=-1)
+
+  # Sum over channels and compute log-probability of each mixture.
+  log_probs = tf.reduce_sum(log_probs, -1) + tf.nn.log_softmax(logits, axis=-1)
   if sum_all:
     output = -tf.reduce_sum(tf.reduce_logsumexp(log_probs, axis=-1))
     return output
   else:
-    output = -tf.reduce_logsumexp(log_probs, axis=-1, keep_dims=True)
+    output = -tf.reduce_logsumexp(log_probs, axis=-1)
     return output
 
 
-def sample_from_discretized_mix_logistic(l, nr_mix, seed=None):
-  """Sampling from a discretized mixture of logistics using gumbel softmax.
+def sample_from_discretized_mix_logistic(pred, seed=None):
+  """Sampling from a discretized mixture of logistics.
 
   Args:
-    l: output of body, of shape [batch, length, num_mixtures * 10]
-    nr_mix: Integer number of mixtures
-    seed: Random seed
+    pred: A [batch, height, width, num_mixtures*10] tensor of floats
+      comprising one unconstrained mixture probability, three means
+      (one per channel), three standard deviations (one per channel),
+      and three coefficients which linearly parameterize dependence across
+      channels.
+    seed: Random seed.
 
   Returns:
-    A tensor of shape [batch, length, 3] with real intensities scaled between
-    -1 and 1
+    A tensor of shape [batch, height, width, 3] with real intensities scaled
+    between -1 and 1.
   """
-  ls_t = tf.shape(l)
-  ls = [ls_t[0], ls_t[1], ls_t[2], ls_t[3]]
-  xs = ls[:-1] + [3]
-  # unpack parameters
-  logit_probs = l[:, :, :, :nr_mix]
-  l = tf.reshape(l[:, :, :, nr_mix:], xs + [nr_mix * 3])
-  # sample mixture indicator from softmax using gumbel softmax
-  sel = tf.one_hot(tf.argmax(logit_probs - tf.log(-tf.log(tf.random_uniform(
-      tf.shape(logit_probs), minval=1e-5, maxval=1. - 1e-5, seed=seed))), 3),
-                   depth=nr_mix, dtype=tf.float32)
-  sel = tf.reshape(sel, xs[:-1] + [1, nr_mix])
-  # select logistic parameters
-  means = tf.reduce_sum(l[:, :, :, :, :nr_mix] * sel, 4)
-  log_scales = tf.maximum(tf.reduce_sum(
-      l[:, :, :, :, nr_mix:2 * nr_mix] * sel, 4), -7.)
-  coeffs = tf.reduce_sum(tf.nn.tanh(
-      l[:, :, :, :, 2 * nr_mix:3 * nr_mix]) * sel, 4)
-  # sample from logistic & clip to interval
-  # we don't actually round to the nearest 8bit value when sampling
-  u = tf.random_uniform(tf.shape(means), minval=1e-5, maxval=1. - 1e-5,
-                        seed=seed)
-  x = means + tf.exp(log_scales) * (tf.log(u) - tf.log(1. - u))
-  x0 = tf.clip_by_value(x[:, :, :, 0], -1., 1)
-  x1 = tf.clip_by_value(x[:, :, :, 1] + coeffs[:, :, :, 0] * x0, -1., 1)
-  x2 = tf.minimum(tf.maximum(
-      x[:, :, :, 2] + coeffs[:, :, :, 1] * x0 + coeffs[:, :, :, 2] * x1, -1.),
-                  1.)
-  return tf.concat([tf.reshape(x0, xs[:-1] + [1]),
-                    tf.reshape(x1, xs[:-1] + [1]),
-                    tf.reshape(x2, xs[:-1] + [1])], 3)
+
+  logits, locs, log_scales, coeffs = split_to_discretized_mix_logistic_params(
+      pred)
+
+  # Sample mixture indicator given logits using the gumbel max trick.
+  num_mixtures = shape_list(logits)[-1]
+  gumbel_noise = -tf.log(-tf.log(tf.random_uniform(
+      tf.shape(logits), minval=1e-5, maxval=1. - 1e-5, seed=seed)))
+  sel = tf.one_hot(tf.argmax(logits + gumbel_noise, -1),
+                   depth=num_mixtures, dtype=tf.float32)
+
+  # Select mixture component's parameters.
+  sel = tf.expand_dims(sel, -1)
+  locs = tf.reduce_sum(locs * sel, 3)
+  log_scales = tf.reduce_sum(log_scales * sel, 3)
+  coeffs = tf.reduce_sum(coeffs * sel, 3)
+
+  # Sample from 3-D logistic & clip to interval. Note we don't round to the
+  # nearest 8-bit value when sampling.
+  uniform_noise = tf.random_uniform(
+      tf.shape(locs), minval=1e-5, maxval=1. - 1e-5, seed=seed)
+  logistic_noise = tf.log(uniform_noise) - tf.log(1. - uniform_noise)
+  x = locs + tf.exp(log_scales) * logistic_noise
+  x0 = x[..., 0]
+  x1 = x[..., 1] + coeffs[..., 0] * x0
+  x2 = x[..., 2] + coeffs[..., 1] * x0 + coeffs[..., 2] * x1
+  x = tf.stack([x0, x1, x2], axis=-1)
+  x = tf.clip_by_value(x, -1., 1.)
+  return x
 
 
 def smoothing_cross_entropy(logits,
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index 24f9448da..03096010d 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -484,6 +484,71 @@ def testPaddingCrossEntropyFactoredGrad(self):
     self.assertAllClose(actual_df, actual_df_factored)
     self.assertAllClose(actual_dw, actual_dw_factored)
 
+  def testDiscretizedMixLogisticLoss(self):
+    batch = 2
+    height = 4
+    width = 4
+    channels = 3
+    num_mixtures = 5
+    logits = tf.concat(  # assign all probability mass to first component
+        [tf.ones([batch, height, width, 1]) * 1e8,
+         tf.zeros([batch, height, width, num_mixtures - 1])],
+        axis=-1)
+    locs = tf.random_uniform([batch, height, width, num_mixtures * 3],
+                             minval=-.9, maxval=.9)
+    log_scales = tf.random_uniform([batch, height, width, num_mixtures * 3],
+                                   minval=-1., maxval=1.)
+    coeffs = tf.atanh(tf.zeros([batch, height, width, num_mixtures * 3]))
+    pred = tf.concat([logits, locs, log_scales, coeffs], axis=-1)
+
+    # Test labels that don't satisfy edge cases where 8-bit value is 0 or 255.
+    labels = tf.random_uniform([batch, height, width, channels],
+                               minval=-.9, maxval=.9)
+    locs_0 = locs[..., :3]
+    log_scales_0 = log_scales[..., :3]
+    centered_labels = labels - locs_0
+    inv_stdv = tf.exp(-log_scales_0)
+    plus_in = inv_stdv * (centered_labels + 1. / 255.)
+    min_in = inv_stdv * (centered_labels - 1. / 255.)
+    cdf_plus = tf.nn.sigmoid(plus_in)
+    cdf_min = tf.nn.sigmoid(min_in)
+    expected_loss = -tf.reduce_sum(tf.log(cdf_plus - cdf_min), axis=-1)
+
+    actual_loss = common_layers.discretized_mix_logistic_loss(
+        labels, pred, sum_all=False)
+    with self.test_session() as session:
+      actual_loss_val, expected_loss_val = session.run(
+          [actual_loss, expected_loss])
+    self.assertAllClose(actual_loss_val, expected_loss_val, rtol=1e-5)
+
+  def testSampleFromDiscretizedMixLogistic(self):
+    batch = 2
+    height = 4
+    width = 4
+    num_mixtures = 5
+    seed = 42
+    logits = tf.concat(  # assign all probability mass to first component
+        [tf.ones([batch, height, width, 1]) * 1e8,
+         tf.zeros([batch, height, width, num_mixtures - 1])],
+        axis=-1)
+    locs = tf.random_uniform([batch, height, width, num_mixtures * 3],
+                             minval=-.9, maxval=.9)
+    log_scales = tf.ones([batch, height, width, num_mixtures * 3]) * -1e8
+    coeffs = tf.atanh(tf.zeros([batch, height, width, num_mixtures * 3]))
+    pred = tf.concat([logits, locs, log_scales, coeffs], axis=-1)
+
+    locs_0 = locs[..., :3]
+    expected_sample = tf.clip_by_value(locs_0, -1., 1.)
+
+    actual_sample = common_layers.sample_from_discretized_mix_logistic(
+        pred, seed=seed)
+    with self.test_session() as session:
+      actual_sample_val, expected_sample_val = session.run(
+          [actual_sample, expected_sample])
+    # Use a low tolerance: samples numerically differ, as the actual
+    # implementation clips log-scales so they always contribute to sampling.
+    self.assertAllClose(actual_sample_val, expected_sample_val, atol=1e-2)
+
   def testFactoredTensorImplicitConversion(self):
     a = np.random.rand(3, 4, 5)
     b = np.random.rand(6, 5)
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 7833f46d9..bfd71eba0 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -119,7 +119,7 @@ def sample(self, features):
     logits, losses = self(features)  # pylint: disable=not-callable
 
     samples = common_layers.sample_from_discretized_mix_logistic(
-        logits, 10, seed=None)
+        logits, seed=None)
     return samples, logits, losses
 
 

From 0fed194ba5b5fc50b143977b906f21edb35d5584 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 15 Jun 2018 17:09:39 -0700
Subject: [PATCH 0120/2720] Enable general slow greedy inference on TPU for t2t
 models.

PiperOrigin-RevId: 200797559
---
 tensor2tensor/models/transformer.py      |   3 +-
 tensor2tensor/models/transformer_test.py |  37 +++++
 tensor2tensor/utils/t2t_model.py         | 197 ++++++++++++++++++++++-
 3 files changed, 228 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index b96ae3f80..e4778468b 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -194,12 +194,13 @@ def body(self, features):
     else:
       return ret
 
-  def _greedy_infer(self, features, decode_length):
+  def _greedy_infer(self, features, decode_length, use_tpu=False):
     """Fast version of greedy decoding.
 
     Args:
       features: an map of string to `Tensor`
       decode_length: an integer.  How many additional timesteps to decode.
+      use_tpu: A bool. Whether to build the inference graph for TPU.
 
     Returns:
       A dict of decoding results {
diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index d3d1ef686..bbfa3d742 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -240,6 +240,43 @@ def testTransformerWithEncoderDecoderAttentionLoss(self):
       res = session.run(extra_loss["attention_loss"])
     self.assertEqual(res.shape, ())
 
+  def testGreedySlowTPUVsNonTPU(self):
+    model, features = get_model(transformer.transformer_small())
+
+    decode_length = 3
+
+    out_logits, _ = model(features)
+    out_logits = tf.squeeze(out_logits, axis=[2, 3])
+    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
+        logits=tf.reshape(out_logits, [-1, VOCAB_SIZE]),
+        labels=tf.reshape(features["targets"], [-1]))
+    loss = tf.reduce_mean(loss)
+    apply_grad = tf.train.AdamOptimizer(0.001).minimize(loss)
+
+    with self.test_session():
+      tf.global_variables_initializer().run()
+      for _ in range(100):
+        apply_grad.run()
+
+    model.set_mode(tf.estimator.ModeKeys.PREDICT)
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+      slow_result_non_tpu = model._slow_greedy_infer(
+          features, decode_length)["outputs"]
+      slow_result_non_tpu = tf.squeeze(slow_result_non_tpu, axis=[2, 3])
+
+      slow_result_tpu = model._slow_greedy_infer_tpu(
+          features, decode_length)["outputs"]
+      slow_result_tpu = tf.squeeze(slow_result_tpu, axis=[2, 3])
+
+    with self.test_session():
+      slow_non_tpu_res = slow_result_non_tpu.eval()
+      slow_tpu_res = slow_result_tpu.eval()
+
+    self.assertEqual(slow_tpu_res.shape,
+                     (BATCH_SIZE, INPUT_LENGTH + decode_length))
+    self.assertAllClose(slow_tpu_res, slow_non_tpu_res)
+
 
 class TransformerScorerTest(tf.test.TestCase):
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 04a6dba8c..2331a3341 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -41,6 +41,7 @@
 import tensorflow as tf
 
 from tensorflow.python.layers import base
+from tensorflow.python.ops import inplace_ops
 from tensorflow.python.ops import variable_scope
 
 _no_problem_err_str = (
@@ -342,10 +343,23 @@ def _top_single(self, body_output, target_modality, features):
         logits = target_modality.top(body_output, features.get("targets"))
       else:
         # Take body outputs for the last position only, and targets too.
-        last_position_body_output = tf.expand_dims(
-            body_output[:, -1, :, :], axis=[1])
-        last_position_targets = tf.expand_dims(
-            features["targets"][:, -1:, :, :], axis=[1])
+        if "decode_loop_step" not in features:
+          last_position_body_output = tf.expand_dims(
+              body_output[:, -1, :, :], axis=[1])
+          last_position_targets = tf.expand_dims(
+              features["targets"][:, -1:, :, :], axis=[1])
+        else:
+          body_output_shape = body_output.shape.as_list()
+          last_position_body_output = tf.slice(
+              body_output,
+              [0, features["decode_loop_step"][0], 0, 0],
+              [body_output_shape[0], 1, body_output_shape[2],
+               body_output_shape[3]])
+          target_shape = features["targets"].shape.as_list()
+          last_position_targets = tf.slice(
+              features["targets"],
+              [0, features["decode_loop_step"][0], 0, 0],
+              [target_shape[0], 1, target_shape[2], target_shape[3]])
         logits = target_modality.top(last_position_body_output,
                                      last_position_targets)
     return logits
@@ -543,7 +557,6 @@ def infer(self,
           "losses": a dictionary: {loss-name (string): floating point `Scalar`
       }
     """
-    del use_tpu
     set_custom_getter_compose(self._custom_getter)
     with self._eager_var_store.as_default():
       # TODO(rsepassi): Make decoding work with real-valued model outputs
@@ -561,7 +574,7 @@ def infer(self,
           beam_size = 1  # No use to run beam-search for a single class.
       if beam_size == 1:
         log_info("Greedy Decoding")
-        results = self._greedy_infer(features, decode_length)
+        results = self._greedy_infer(features, decode_length, use_tpu)
       else:
         log_info("Beam Decoding with beam size %d" % beam_size)
         results = self._beam_decode(features, decode_length, beam_size,
@@ -674,7 +687,7 @@ def symbols_to_logits_fn(ids):
 
     return {"outputs": samples, "scores": scores}
 
-  def _greedy_infer(self, features, decode_length):
+  def _greedy_infer(self, features, decode_length, use_tpu=False):
     """A greedy inference method.
 
     Models should ideally implement a more efficient version of this function.
@@ -682,6 +695,29 @@ def _greedy_infer(self, features, decode_length):
     Args:
       features: an map of string to `Tensor`
       decode_length: an integer.  How many additional timesteps to decode.
+      use_tpu: A bool, whether to build the inference graph for TPU.
+
+    Returns:
+      A dict of decoding results {
+          "outputs": integer `Tensor` of decoded ids of shape
+              [batch_size, <= decode_length] if beam_size == 1 or
+              [batch_size, top_beams, <= decode_length]
+          "scores": None
+          "logits": `Tensor` of shape [batch_size, time, 1, 1, vocab_size].
+          "losses": a dictionary: {loss-name (string): floating point `Scalar`}
+      }
+    """
+    return (self._slow_greedy_infer_tpu(features, decode_length) if use_tpu else
+            self._slow_greedy_infer(features, decode_length))
+
+  def _slow_greedy_infer_tpu(self, features, decode_length):
+    """A slow greedy inference method on TPU.
+
+    Quadratic time in decode_length.
+
+    Args:
+      features: An map of string to `Tensor`.
+      decode_length: An integer, how many additional timesteps to decode.
 
     Returns:
       A dict of decoding results {
@@ -693,7 +729,152 @@ def _greedy_infer(self, features, decode_length):
           "losses": a dictionary: {loss-name (string): floating point `Scalar`}
       }
     """
-    return self._slow_greedy_infer(features, decode_length)
+    if not features:
+      features = {}
+    inputs_old = None
+    if "inputs" in features and len(features["inputs"].shape) < 4:
+      inputs_old = features["inputs"]
+      features["inputs"] = tf.expand_dims(features["inputs"], 2)
+    if not self.has_input:
+      # Prepare partial targets.
+      # In either features["inputs"] or features["targets"].
+      # We force the outputs to begin with these sequences.
+      partial_targets = features.get("inputs")
+      if partial_targets is None:
+        partial_targets = features["targets"]
+      features["partial_targets"] = tf.to_int64(partial_targets)
+    # Save the targets in a var and reassign it after the tf.while loop to avoid
+    # having targets being in a 'while' frame. This ensures targets when used
+    # in metric functions stays in the same frame as other vars.
+    targets_old = features.get("targets", None)
+
+    target_modality = self._problem_hparams.target_modality
+
+    def infer_step(i, recent_output, recent_logits, unused_loss):
+      """Inference step."""
+      if not tf.contrib.eager.in_eager_mode():
+        recent_output.set_shape([None, None, None, 1])
+      padded = tf.pad(recent_output, [[0, 0], [0, 1], [0, 0], [0, 0]])
+      features["targets"] = padded
+      # This is inefficient in that it generates samples at all timesteps,
+      # not just the last one, except if target_modality is pointwise.
+      features["decode_loop_step"] = i
+      samples, logits, losses = self.sample(features)
+      # Concatenate the already-generated recent_output with last timestep
+      # of the newly-generated samples.
+      if target_modality.top_is_pointwise:
+        cur_sample = samples[:, -1, :, :]
+      else:
+        cur_sample = samples[:, i, :, :]
+      samples = tf.transpose(recent_output, perm=[1, 0, 2, 3])
+      samples = inplace_ops.alias_inplace_update(
+          samples, i, tf.to_int64(cur_sample))
+      samples = tf.transpose(samples, perm=[1, 0, 2, 3])
+      if not tf.contrib.eager.in_eager_mode():
+        samples.set_shape([None, None, None, 1])
+
+      # Assuming we have one shard for logits.
+      recent_logits = tf.transpose(recent_logits, perm=[1, 0, 2, 3, 4])
+      recent_logits = inplace_ops.alias_inplace_update(
+          recent_logits, i, tf.squeeze(logits[:, -1:], axis=1))
+      logits = tf.transpose(recent_logits, perm=[1, 0, 2, 3, 4])
+      loss = sum([l for l in losses.values() if l is not None])
+      return i + 1, samples, logits, loss
+
+    # Create an initial output tensor. This will be passed
+    # to the infer_step, which adds one timestep at every iteration.
+    if "partial_targets" in features:
+      initial_output = tf.to_int64(features["partial_targets"])
+      while len(initial_output.get_shape().as_list()) < 4:
+        initial_output = tf.expand_dims(initial_output, 2)
+      batch_size = common_layers.shape_list(initial_output)[0]
+    else:
+      batch_size = common_layers.shape_list(features["inputs"])[0]
+      initial_output = tf.zeros((batch_size, 0, 1, 1), dtype=tf.int64)
+    # Hack: foldl complains when the output shape is less specified than the
+    # input shape, so we confuse it about the input shape.
+    initial_output = tf.slice(initial_output, [0, 0, 0, 0],
+                              common_layers.shape_list(initial_output))
+    target_modality = self._problem_hparams.target_modality
+    if target_modality.is_class_modality:
+      decode_length = 1
+    else:
+      if "partial_targets" in features:
+        prefix_length = common_layers.shape_list(
+            features["partial_targets"])[1]
+      else:
+        prefix_length = common_layers.shape_list(
+            features["inputs"])[1]
+      decode_length = prefix_length + decode_length
+
+    # Initial values of result, logits and loss.
+    result = tf.concat([initial_output,
+                        tf.zeros([batch_size, decode_length, 1, 1], tf.int64)],
+                       axis=1)
+    # tensor padded to [batch_size, decode_length, 1, 1, vocab_size]
+    logits = tf.zeros((batch_size, decode_length, 1, 1,
+                       target_modality.top_dimensionality))
+    if not tf.contrib.eager.in_eager_mode():
+      logits.set_shape([None, None, None, None, None])
+    loss = 0.0
+
+    def while_exit_cond(i, result, logits, loss):  # pylint: disable=unused-argument
+      """Exit the loop either if reach decode_length or EOS."""
+      not_overflow = i < decode_length
+
+      if self._problem_hparams.stop_at_eos:
+
+        def fn_not_eos():
+          # Check if the last predicted element is a EOS
+          return tf.reduce_any(
+              tf.not_equal(
+                  tf.squeeze(result[:, -1, :, :]), text_encoder.EOS_ID))
+
+        not_eos = tf.cond(
+            # We only check for early stopping if there is at least 1 element (
+            # otherwise not_eos will crash).
+            tf.not_equal(i, 0),
+            fn_not_eos,
+            lambda: True,
+        )
+
+        return tf.cond(
+            tf.equal(batch_size, 1),
+            # If batch_size == 1, we check EOS for early stopping.
+            lambda: tf.logical_and(not_overflow, not_eos),
+            # Else, just wait for max length
+            lambda: not_overflow)
+      return not_overflow
+
+    _, result, logits, loss = tf.while_loop(
+        while_exit_cond,
+        infer_step, [tf.constant(0), result, logits, loss],
+        shape_invariants=[
+            tf.TensorShape([]),
+            tf.TensorShape([batch_size, decode_length, 1, 1]),
+            tf.TensorShape([batch_size, decode_length, 1, 1,
+                            target_modality.top_dimensionality]),
+            tf.TensorShape([]),
+        ],
+        back_prop=False,
+        parallel_iterations=1)
+    if inputs_old is not None:  # Restore to not confuse Estimator.
+      features["inputs"] = inputs_old
+    # Reassign targets back to the previous value.
+    if targets_old is not None:
+      features["targets"] = targets_old
+    losses = {"training": loss}
+    if "partial_targets" in features:
+      partial_target_length = common_layers.shape_list(
+          features["partial_targets"])[1]
+      result = tf.slice(result, [0, partial_target_length, 0, 0],
+                        [-1, -1, -1, -1])
+    return {
+        "outputs": result,
+        "scores": None,
+        "logits": logits,
+        "losses": losses,
+    }
 
   def _slow_greedy_infer(self, features, decode_length):
     """A slow greedy inference method.

From 25b72af9300093c6021c3371717dafed9208e08d Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 16 Jun 2018 15:09:35 +0200
Subject: [PATCH 0121/2720] refactor of wrappers and collect

---
 tensor2tensor/models/research/rl.py          |  15 +-
 tensor2tensor/rl/collect.py                  | 169 ++++++++++++++-----
 tensor2tensor/rl/envs/py_func_batch_env.py   |  48 +-----
 tensor2tensor/rl/envs/simulated_batch_env.py |  15 +-
 tensor2tensor/rl/envs/utils.py               |  76 +++++++--
 tensor2tensor/rl/ppo.py                      |  40 ++---
 tensor2tensor/rl/rl_trainer_lib.py           |  65 +++----
 tensor2tensor/rl/rl_trainer_lib_test.py      |   6 +-
 8 files changed, 262 insertions(+), 172 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index c8d57792c..8281cc8db 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -22,6 +22,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
 from tensor2tensor.utils import registry
+from tensor2tensor.rl.envs import tf_atari_wrappers
 
 import tensorflow as tf
 
@@ -59,16 +60,26 @@ def ppo_base_v1():
 @registry.register_hparams
 def ppo_continuous_action_base():
   hparams = ppo_base_v1()
-  hparams.add_hparam("network", feed_forward_gaussian_fun)
+  hparams.add_hparam("policy_network", feed_forward_gaussian_fun)
+  hparams.add_hparam("policy_network_params", "basic_policy_parameters")
   return hparams
 
+@registry.register_hparams
+def basic_policy_parameters():
+  wrappers = None
+  return tf.contrib.training.HParams(wrappers=wrappers)
 
 @registry.register_hparams
 def ppo_discrete_action_base():
   hparams = ppo_base_v1()
-  hparams.add_hparam("network", feed_forward_categorical_fun)
+  hparams.add_hparam("policy_network", feed_forward_categorical_fun)
+  hparams.add_hparam("policy_network_params", "standard_atari_parameters")
   return hparams
 
+@registry.register_hparams
+def standard_atari_parameters():
+  wrappers = [[tf_atari_wrappers.MaxAndSkipWrapper, {"skip": 4}]]
+  return tf.contrib.training.HParams(wrappers=wrappers)
 
 @registry.register_hparams
 def ppo_atari_base():
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index a3e807efd..7edec7324 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -17,44 +17,104 @@
 from __future__ import division
 from __future__ import print_function
 
+
+
+import copy
 import tensorflow as tf
 
+from tensor2tensor.rl.envs import utils
 
-def define_collect(policy_factory, batch_env, hparams,
-                   eval_phase, policy_to_actions_lambda=None,
-                   scope="", preprocess=None, on_simulated=False):
-  """Collect trajectories."""
-  eval_phase = tf.convert_to_tensor(eval_phase)
-  on_simulated = tf.convert_to_tensor(on_simulated)
+from tensor2tensor.rl.envs.tf_atari_wrappers import WrapperBase
+from tensor2tensor.rl.envs.utils import get_policy
+from tensor2tensor.utils import registry
+
+def _rollout_metadata(batch_env):
   batch_env_shape = batch_env.observ.get_shape().as_list()
-  if preprocess is not None:
-    batch_env_shape = preprocess[1]
-  memory_shape = [hparams.epoch_length] + [batch_env_shape[0]]
-  memories_shapes_and_types = [
-      # observation
-      (memory_shape + batch_env_shape[1:], tf.float32),
-      (memory_shape, tf.float32),      # reward
-      (memory_shape, tf.bool),         # done
-      # action
-      (memory_shape + batch_env.action_shape, batch_env.action_dtype),
-      (memory_shape, tf.float32),      # pdf
-      (memory_shape, tf.float32),      # value function
+  batch_size =  [batch_env_shape[0]]
+  shapes_types_names = [
+    # observation
+    (batch_size + batch_env_shape[1:], tf.float32, "observation"),
+    (batch_size, tf.float32, "reward"),
+    (batch_size, tf.bool, "done"),
+    (batch_size + batch_env.action_shape, batch_env.action_dtype, "action"),
+    (batch_size, tf.float32, "pdf"),
+    (batch_size, tf.float32, "value_function"),
   ]
-  memory = [tf.Variable(tf.zeros(shape, dtype), trainable=False)
-            for (shape, dtype) in memories_shapes_and_types]
-  with tf.variable_scope(scope):
+  return shapes_types_names
+
+
+class _MemoryWrapper(WrapperBase):
+  """Memory wrapper."""
+
+  def __init__(self, batch_env):
+    super(_MemoryWrapper, self).__init__(batch_env)
+    infinity = 10000000
+    meta_data = list(zip(*_rollout_metadata(batch_env)))
+    shapes = meta_data[0][:4]
+    dtypes = meta_data[1][:4]
+    self.speculum = tf.FIFOQueue(infinity, shapes=shapes, dtypes=dtypes)
+    observs_shape = batch_env.observ.shape
+    observ_dtype = tf.float32
+    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
+                               trainable=False)
+
+  def simulate(self, action):
+    reward, done = self._batch_env.simulate(action)
+    # # image = tf.cast(self._batch_env.observ[0, ...], tf.uint8)
+    # image = self._batch_env.observ #possibly remove
+    with tf.control_dependencies([reward, done]):
+      assign = self._observ.assign(self._batch_env.observ)
+
+    with tf.control_dependencies([assign]):
+      enqueue_op = self.speculum.enqueue(
+          [self._observ, reward, done, action])
+      with tf.control_dependencies([enqueue_op]):
+        return tf.identity(reward), tf.identity(done)
+
+
+def define_collect(hparams, scope, eval_phase,
+                   collect_level=-1,
+                   policy_to_actions_lambda=None,
+                   on_simulated=False):
+  """Collect trajectories."""
+
+  with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
+    batch_env = utils.batch_env_factory(hparams)
+    policy_hparams = registry.hparams(hparams.policy_network_params)
+    wrappers = copy.copy(policy_hparams.wrappers) if policy_hparams.wrappers else []
+    #Put memory wrapper at the level you want to gather observations at
+    #Negative indices need to be shifted for insert to work correctly
+    collect_level = collect_level if collect_level>=0 else len(wrappers) + collect_level + 1
+    wrappers.insert(collect_level, [_MemoryWrapper, {}])
+    rollout_metadata = None
+    speculum = None
+    for w in wrappers:
+      batch_env = w[0](batch_env, **w[1])
+      if w[0] == _MemoryWrapper:
+        rollout_metadata = _rollout_metadata(batch_env)
+        speculum = batch_env.speculum
+
+    eval_phase = tf.convert_to_tensor(eval_phase)
+    on_simulated = tf.convert_to_tensor(on_simulated)
+
+    memory = [tf.get_variable("collect_memory_{}".format(name),
+                              shape=[hparams.epoch_length]+shape,
+                              dtype=dtype,
+                              initializer=tf.zeros_initializer(),
+                              trainable=False)
+              for (shape, dtype, name) in rollout_metadata]
+
+
     cumulative_rewards = tf.get_variable("cumulative_rewards", len(batch_env),
-                                         trainable=False)
+                                           trainable=False)
 
-  should_reset_var = tf.Variable(True, trainable=False)
+    should_reset_var = tf.Variable(True, trainable=False)
 
-  zeros_tensor = tf.zeros(len(batch_env))
+    zeros_tensor = tf.zeros(len(batch_env))
 
   def group():
-    return tf.group(
-        batch_env.reset(tf.range(len(batch_env))),
-        tf.assign(cumulative_rewards, zeros_tensor))
-
+    return tf.group(batch_env.reset(tf.range(len(batch_env))),
+                    tf.assign(cumulative_rewards, zeros_tensor))
   reset_op = tf.cond(
       tf.logical_or(should_reset_var, tf.logical_or(eval_phase, on_simulated)),
       group, tf.no_op)
@@ -71,24 +131,43 @@ def step(index, scores_sum, scores_num):
       # operation. We are waiting for tf.copy:
       # https://github.com/tensorflow/tensorflow/issues/11186
       obs_copy = batch_env.observ + 0
-      if preprocess is not None:
-        obs_copy = preprocess[0](obs_copy)
-      actor_critic = policy_factory(tf.expand_dims(obs_copy, 0))
-      policy = actor_critic.policy
-      if policy_to_actions_lambda:
-        action = policy_to_actions_lambda(policy)
-      else:
-        action = tf.cond(eval_phase,
-                         policy.mode,
-                         policy.sample)
-      postprocessed_action = actor_critic.action_postprocessing(action)
-      simulate_output = batch_env.simulate(postprocessed_action[0, ...])
-      pdf = policy.prob(action)[0]
-      with tf.control_dependencies(simulate_output):
-        reward, done = simulate_output
+
+      def env_step(arg1, arg2):
+        actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
+        policy = actor_critic.policy
+        if policy_to_actions_lambda:
+          action = policy_to_actions_lambda(policy)
+        else:
+          action = tf.cond(eval_phase,
+                           policy.mode,
+                           policy.sample)
+
+        postprocessed_action = actor_critic.action_postprocessing(action)
+        simulate_output = batch_env.simulate(postprocessed_action[0, ...])
+
+        pdf = policy.prob(action)[0]
+        value_function = actor_critic.value[0]
+        pdf = tf.reshape(pdf, shape=(hparams.num_agents, ))
+        value_function = tf.reshape(value_function, shape=(hparams.num_agents, ))
+
+        with tf.control_dependencies(simulate_output):
+          return tf.identity(pdf), tf.identity(value_function)
+
+      pdf, value_function = tf.while_loop(
+        lambda _1, _2: tf.equal(speculum.size(), 0),
+        env_step,
+        [tf.constant(0.0, shape=(hparams.num_agents,)),
+          tf.constant(0.0, shape=(hparams.num_agents,))],
+        parallel_iterations=1,
+        back_prop=False,
+        )
+
+      with tf.control_dependencies([pdf, value_function]):
+        obs, reward, done, action = speculum.dequeue()
+
         done = tf.reshape(done, (len(batch_env),))
-        to_save = [obs_copy, reward, done, action[0, ...], pdf,
-                   actor_critic.value[0]]
+        to_save = [obs, reward, done, action,
+                   pdf, value_function]
         save_ops = [tf.scatter_update(memory_slot, index, value)
                     for memory_slot, value in zip(memory, to_save)]
         cumulate_rewards_op = cumulative_rewards.assign_add(reward)
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index 692a37298..08ac236f1 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -20,12 +20,10 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-import gym
 
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
-
 import tensorflow as tf
-
+from tensor2tensor.rl.envs import utils
 
 class PyFuncBatchEnv(InGraphBatchEnv):
   """Batch of environments inside the TensorFlow graph.
@@ -42,10 +40,10 @@ def __init__(self, batch_env):
       batch_env: Batch environment.
     """
     self._batch_env = batch_env
-    observ_shape = self._parse_shape(self._batch_env.observation_space)
-    observ_dtype = self._parse_dtype(self._batch_env.observation_space)
-    self.action_shape = list(self._parse_shape(self._batch_env.action_space))
-    self.action_dtype = self._parse_dtype(self._batch_env.action_space)
+    observ_shape = utils.parse_shape(self._batch_env.observation_space)
+    observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
+    self.action_shape = list(utils.parse_shape(self._batch_env.action_space))
+    self.action_dtype = utils.parse_dtype(self._batch_env.action_space)
     with tf.variable_scope('env_temporary'):
       self._observ = tf.Variable(
           tf.zeros((len(self._batch_env),) + observ_shape, observ_dtype),
@@ -84,7 +82,7 @@ def simulate(self, action):
     with tf.name_scope('environment/simulate'):
       if action.dtype in (tf.float16, tf.float32, tf.float64):
         action = tf.check_numerics(action, 'action')
-      observ_dtype = self._parse_dtype(self._batch_env.observation_space)
+      observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
       observ, reward, done = tf.py_func(
           lambda a: self._batch_env.step(a)[:3], [action],
           [observ_dtype, tf.float32, tf.bool], name='step')
@@ -104,7 +102,7 @@ def _reset_non_empty(self, indices):
     Returns:
       Batch tensor of the new observations.
     """
-    observ_dtype = self._parse_dtype(self._batch_env.observation_space)
+    observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
     observ = tf.py_func(
         self._batch_env.reset, [indices], observ_dtype, name='reset')
     observ = tf.check_numerics(observ, 'observ')
@@ -119,34 +117,4 @@ def observ(self):
 
   def close(self):
     """Send close messages to the external process and join them."""
-    self._batch_env.close()
-
-  def _parse_shape(self, space):
-    """Get a tensor shape from a OpenAI Gym space.
-
-    Args:
-      space: Gym space.
-
-    Returns:
-      Shape tuple.
-    """
-    if isinstance(space, gym.spaces.Discrete):
-      return ()
-    if isinstance(space, gym.spaces.Box):
-      return space.shape
-    raise NotImplementedError()
-
-  def _parse_dtype(self, space):
-    """Get a tensor dtype from a OpenAI Gym space.
-
-    Args:
-      space: Gym space.
-
-    Returns:
-      TensorFlow data type.
-    """
-    if isinstance(space, gym.spaces.Discrete):
-      return tf.int32
-    if isinstance(space, gym.spaces.Box):
-      return tf.float32
-    raise NotImplementedError()
+    self._batch_env.close()
\ No newline at end of file
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 1c460f407..8f0920e19 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -22,6 +22,7 @@
 from __future__ import print_function
 from tensor2tensor.layers import common_layers
 from tensor2tensor.rl.envs import in_graph_batch_env
+from tensor2tensor.rl.envs import utils
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
@@ -90,7 +91,7 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   flags are held in according variables.
   """
 
-  def __init__(self, environment_lambda, length, problem,
+  def __init__(self, hparams, length, problem,
                simulation_random_starts=False, intrinsic_reward_scale=0.):
     """Batch of environments inside the TensorFlow graph."""
     self.length = length
@@ -98,16 +99,14 @@ def __init__(self, environment_lambda, length, problem,
     self._num_frames = problem.num_input_frames
     self._intrinsic_reward_scale = intrinsic_reward_scale
 
-    initialization_env = environment_lambda()
-    hparams = trainer_lib.create_hparams(
+    # initialization_env = environment_lambda()
+    model_hparams = trainer_lib.create_hparams(
         FLAGS.hparams_set, problem_name=FLAGS.problem)
-    hparams.force_full_predict = True
+    model_hparams.force_full_predict = True
     self._model = registry.model(FLAGS.model)(
-        hparams, tf.estimator.ModeKeys.PREDICT)
+      model_hparams, tf.estimator.ModeKeys.PREDICT)
 
-    self.action_space = initialization_env.action_space
-    self.action_shape = list(initialization_env.action_space.shape)
-    self.action_dtype = tf.int32
+    _, self.action_shape, self.action_dtype = utils.get_action_space(hparams)
 
     if simulation_random_starts:
       dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index de2af016a..dfc3de8e9 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -32,7 +32,6 @@
 import gym
 
 from tensor2tensor.rl.envs import batch_env
-
 from tensor2tensor.rl.envs import py_func_batch_env
 from tensor2tensor.rl.envs import simulated_batch_env
 
@@ -279,38 +278,87 @@ def _worker(self, constructor, conn):
     conn.close()
 
 
-def batch_env_factory(environment_lambda, hparams, num_agents, xvfb=False):
+def batch_env_factory(hparams, xvfb=False):
   """Factory of batch envs."""
-  wrappers = hparams.in_graph_wrappers if hasattr(
-      hparams, "in_graph_wrappers") else []
 
   if hparams.simulated_environment:
-    cur_batch_env = define_simulated_batch_env(
-        environment_lambda, num_agents, hparams.problem,
+    cur_batch_env = _define_simulated_batch_env(
+      hparams, hparams.num_agents, hparams.problem,
         hparams.simulation_random_starts,
         hparams.intrinsic_reward_scale)
   else:
-    cur_batch_env = define_batch_env(environment_lambda, num_agents, xvfb=xvfb)
-  for w in wrappers:
-    cur_batch_env = w[0](cur_batch_env, **w[1])
+    cur_batch_env = _define_batch_env(hparams, xvfb=xvfb)
   return cur_batch_env
 
 
-def define_batch_env(constructor, num_agents, xvfb=False):
+def _define_batch_env(hparams, xvfb=False):
   """Create environments and apply all desired wrappers."""
+  if isinstance(hparams.environment_spec, str):
+    environment_lambda = lambda: gym.make(hparams.environment_spec)
+  else:
+    environment_lambda = hparams.environment_spec
+
   with tf.variable_scope("environments"):
     envs = [
-        ExternalProcessEnv(constructor, xvfb)
-        for _ in range(num_agents)]
+        ExternalProcessEnv(environment_lambda, xvfb)
+        for _ in range(hparams.num_agents)]
     env = batch_env.BatchEnv(envs, blocking=False)
     env = py_func_batch_env.PyFuncBatchEnv(env)
     return env
 
 
-def define_simulated_batch_env(environment_lambda, num_agents, problem,
+def _define_simulated_batch_env(hparams, num_agents, problem,
                                simulation_random_starts=False,
                                intrinsic_reward_scale=0.):
   cur_batch_env = simulated_batch_env.SimulatedBatchEnv(
-      environment_lambda, num_agents, problem, simulation_random_starts,
+      hparams, num_agents, problem, simulation_random_starts,
       intrinsic_reward_scale)
   return cur_batch_env
+
+
+def get_action_space(hparams):
+  if isinstance(hparams.environment_spec, str):
+    environment_lambda = lambda: gym.make(hparams.environment_spec)
+  else:
+    environment_lambda = hparams.environment_spec
+
+  action_space = environment_lambda().action_space
+  action_shape = list(parse_shape(action_space))
+  action_dtype = parse_dtype(action_space)
+
+  return action_space, action_shape, action_dtype
+
+def get_policy(observations, hparams):
+  policy_network_lambda = hparams.policy_network
+  action_space, _, _ = get_action_space(hparams)
+  return policy_network_lambda(action_space, hparams, observations)
+
+def parse_shape(space):
+  """Get a tensor shape from a OpenAI Gym space.
+
+  Args:
+    space: Gym space.
+
+  Returns:
+    Shape tuple.
+  """
+  if isinstance(space, gym.spaces.Discrete):
+    return ()
+  if isinstance(space, gym.spaces.Box):
+    return space.shape
+  raise NotImplementedError()
+
+def parse_dtype(space):
+  """Get a tensor dtype from a OpenAI Gym space.
+
+  Args:
+    space: Gym space.
+
+  Returns:
+    TensorFlow data type.
+  """
+  if isinstance(space, gym.spaces.Discrete):
+    return tf.int32
+  if isinstance(space, gym.spaces.Box):
+    return tf.float32
+  raise NotImplementedError()
\ No newline at end of file
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index a62a8a920..750c02537 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -19,6 +19,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+from tensor2tensor.rl.envs.utils import get_policy
 
 import tensorflow as tf
 
@@ -29,25 +30,25 @@ def get_optimiser(config):
   return config.optimizer(learning_rate=config.learning_rate)
 
 
-def define_ppo_step(data_points, policy_factory, optimizer, config):
+def define_ppo_step(data_points, optimizer, hparams):
   """Define ppo step."""
   observation, action, discounted_reward, norm_advantage, old_pdf = data_points
-  new_policy_dist, new_value, _ = policy_factory(observation)
+  new_policy_dist, new_value, _ = get_policy(observation, hparams)
   new_pdf = new_policy_dist.prob(action)
 
   ratio = new_pdf / old_pdf
-  clipped_ratio = tf.clip_by_value(ratio, 1 - config.clipping_coef,
-                                   1 + config.clipping_coef)
+  clipped_ratio = tf.clip_by_value(ratio, 1 - hparams.clipping_coef,
+                                   1 + hparams.clipping_coef)
 
   surrogate_objective = tf.minimum(clipped_ratio * norm_advantage,
                                    ratio * norm_advantage)
   policy_loss = -tf.reduce_mean(surrogate_objective)
 
   value_error = new_value - discounted_reward
-  value_loss = config.value_loss_coef * tf.reduce_mean(value_error ** 2)
+  value_loss = hparams.value_loss_coef * tf.reduce_mean(value_error ** 2)
 
   entropy = new_policy_dist.entropy()
-  entropy_loss = -config.entropy_loss_coef * tf.reduce_mean(entropy)
+  entropy_loss = -hparams.entropy_loss_coef * tf.reduce_mean(entropy)
 
   losses = [policy_loss, value_loss, entropy_loss]
 
@@ -59,9 +60,9 @@ def define_ppo_step(data_points, policy_factory, optimizer, config):
   gradients_flat = sum([gradient[0] for gradient in gradients], ())
   gradients_variables_flat = sum([gradient[1] for gradient in gradients], ())
 
-  if config.max_gradients_norm:
+  if hparams.max_gradients_norm:
     gradients_flat, _ = tf.clip_by_global_norm(gradients_flat,
-                                               config.max_gradients_norm)
+                                               hparams.max_gradients_norm)
 
   optimize_op = optimizer.apply_gradients(zip(gradients_flat,
                                               gradients_variables_flat))
@@ -70,7 +71,7 @@ def define_ppo_step(data_points, policy_factory, optimizer, config):
     return [tf.identity(x) for x in losses + gradients_norms]
 
 
-def define_ppo_epoch(memory, policy_factory, config):
+def define_ppo_epoch(memory, hparams):
   """PPO epoch."""
   observation, reward, done, action, old_pdf, value = memory
 
@@ -78,14 +79,14 @@ def define_ppo_epoch(memory, policy_factory, config):
   observation = tf.stop_gradient(observation)
   action = tf.stop_gradient(action)
   reward = tf.stop_gradient(reward)
-  if hasattr(config, "rewards_preprocessing_fun"):
-    reward = config.rewards_preprocessing_fun(reward)
+  if hasattr(hparams, "rewards_preprocessing_fun"):
+    reward = hparams.rewards_preprocessing_fun(reward)
   done = tf.stop_gradient(done)
   value = tf.stop_gradient(value)
   old_pdf = tf.stop_gradient(old_pdf)
 
   advantage = calculate_generalized_advantage_estimator(
-      reward, value, done, config.gae_gamma, config.gae_lambda)
+      reward, value, done, hparams.gae_gamma, hparams.gae_lambda)
 
   discounted_reward = tf.stop_gradient(advantage + value)
 
@@ -96,23 +97,22 @@ def define_ppo_epoch(memory, policy_factory, config):
 
   add_lists_elementwise = lambda l1, l2: [x + y for x, y in zip(l1, l2)]
 
-  number_of_batches = (config.epoch_length * config.optimization_epochs
-                       / config.optimization_batch_size)
+  number_of_batches = (hparams.epoch_length * hparams.optimization_epochs
+                       / hparams.optimization_batch_size)
 
   dataset = tf.data.Dataset.from_tensor_slices(
       (observation, action, discounted_reward, advantage_normalized, old_pdf))
-  dataset = dataset.shuffle(buffer_size=config.epoch_length,
+  dataset = dataset.shuffle(buffer_size=hparams.epoch_length,
                             reshuffle_each_iteration=True)
-  dataset = dataset.repeat(config.optimization_epochs)
-  dataset = dataset.batch(config.optimization_batch_size)
+  dataset = dataset.repeat(hparams.optimization_epochs)
+  dataset = dataset.batch(hparams.optimization_batch_size)
   iterator = dataset.make_initializable_iterator()
-  optimizer = get_optimiser(config)
+  optimizer = get_optimiser(hparams)
 
   with tf.control_dependencies([iterator.initializer]):
     ppo_step_rets = tf.scan(
         lambda a, i: add_lists_elementwise(  # pylint: disable=g-long-lambda
-            a, define_ppo_step(iterator.get_next(), policy_factory, optimizer,
-                               config)),
+            a, define_ppo_step(iterator.get_next(), optimizer, hparams)),
         tf.range(number_of_batches),
         [0., 0., 0., 0., 0., 0.],
         parallel_iterations=1)
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 4f4b923c3..f431afe47 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -32,59 +32,42 @@
 import tensorflow as tf
 
 
-def define_train(hparams, environment_spec, event_dir):
+def define_train(hparams, event_dir):
   """Define the training setup."""
-  policy_lambda = hparams.network
-
-  if environment_spec == "stacked_pong":
-    environment_spec = lambda: gym.make("PongNoFrameskip-v4")
-    wrappers = hparams.in_graph_wrappers if hasattr(
-        hparams, "in_graph_wrappers") else []
-    wrappers.append((tf_atari_wrappers.MaxAndSkipWrapper, {"skip": 4}))
-    hparams.in_graph_wrappers = wrappers
-  if isinstance(environment_spec, str):
-    env_lambda = lambda: gym.make(environment_spec)
-  else:
-    env_lambda = environment_spec
-
-  batch_env = utils.batch_env_factory(
-      env_lambda, hparams, num_agents=hparams.num_agents)
-
-  policy_factory = functools.partial(
-      policy_lambda, batch_env.action_space, hparams)
 
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
     memory, collect_summary = collect.define_collect(
-        policy_factory, batch_env, hparams, eval_phase=False,
-        on_simulated=hparams.simulated_environment)
-    ppo_summary = ppo.define_ppo_epoch(memory, policy_factory, hparams)
+        hparams, "ppo_train", eval_phase=False, on_simulated=hparams.simulated_environment)
+    ppo_summary = ppo.define_ppo_epoch(memory, hparams)
     summary = tf.summary.merge([collect_summary, ppo_summary])
 
-  with tf.variable_scope("eval", reuse=tf.AUTO_REUSE):
-    eval_env_lambda = env_lambda
-    if event_dir and hparams.video_during_eval:
-      # Some environments reset environments automatically, when reached done
-      # state. For them we shall record only every second episode.
-      d = 2 if env_lambda().metadata.get("semantics.autoreset") else 1
-      eval_env_lambda = lambda: gym.wrappers.Monitor(  # pylint: disable=g-long-lambda
-          env_lambda(), event_dir, video_callable=lambda i: i % d == 0)
-      eval_env_lambda = (
-          lambda: utils.EvalVideoWrapper(eval_env_lambda()))
-    eval_batch_env = utils.batch_env_factory(
-        eval_env_lambda, hparams,
-        num_agents=hparams.num_eval_agents, xvfb=hparams.video_during_eval)
-
-    _, eval_summary = collect.define_collect(
-        policy_factory, eval_batch_env, hparams, eval_phase=True)
+  # with tf.variable_scope("eval", reuse=tf.AUTO_REUSE):
+  #   eval_env_lambda = env_lambda
+  #   if event_dir and hparams.video_during_eval:
+  #     # Some environments reset environments automatically, when reached done
+  #     # state. For them we shall record only every second episode.
+  #     d = 2 if env_lambda().metadata.get("semantics.autoreset") else 1
+  #     eval_env_lambda = lambda: gym.wrappers.Monitor(  # pylint: disable=g-long-lambda
+  #         env_lambda(), event_dir, video_callable=lambda i: i % d == 0)
+  #     eval_env_lambda = (
+  #         lambda: utils.EvalVideoWrapper(eval_env_lambda()))
+  #   eval_batch_env = utils.batch_env_factory(
+  #       eval_env_lambda, hparams,
+  #       num_agents=hparams.num_eval_agents, xvfb=hparams.video_during_eval)
+  #
+  #   _, eval_summary = collect_new.define_collect(
+  #       policy_factory, eval_batch_env, hparams, eval_phase=True)
+
+  #Fake for development
+  eval_summary = summary
   return summary, eval_summary
 
 
-def train(hparams, environment_spec, event_dir=None, model_dir=None,
+def train(hparams, event_dir=None, model_dir=None,
           restore_agent=True, epoch=0):
   """Train."""
   with tf.name_scope("rl_train"):
-    train_summary_op, eval_summary_op = define_train(hparams, environment_spec,
-                                                     event_dir)
+    train_summary_op, eval_summary_op = define_train(hparams, event_dir)
     if event_dir:
       summary_writer = tf.summary.FileWriter(
           event_dir, graph=tf.get_default_graph(), flush_secs=60)
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index 9e14c4e9a..396ce4292 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -31,13 +31,15 @@ def test_no_crash_pendulum(self):
     hparams = trainer_lib.create_hparams(
         "ppo_continuous_action_base",
         TrainTest.test_config)
-    rl_trainer_lib.train(hparams, "Pendulum-v0")
+    hparams.add_hparam("environment_spec", "Pendulum-v0")
+    rl_trainer_lib.train(hparams)
 
   def test_no_crash_cartpole(self):
     hparams = trainer_lib.create_hparams(
         "ppo_discrete_action_base",
         TrainTest.test_config)
-    rl_trainer_lib.train(hparams, "CartPole-v0")
+    hparams.add_hparam("environment_spec", "CartPole-v0")
+    rl_trainer_lib.train(hparams)
 
 
 if __name__ == "__main__":

From e78426487f4cd8b9998e5c030f067ae44264a779 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 16 Jun 2018 18:24:07 +0200
Subject: [PATCH 0122/2720] wrappers refactor part 2

---
 tensor2tensor/data_generators/gym_problems.py | 52 ++++++++-----------
 tensor2tensor/data_generators/video_utils.py  |  2 +-
 tensor2tensor/models/research/rl.py           | 10 +++-
 tensor2tensor/rl/collect.py                   |  2 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py    | 24 +--------
 5 files changed, 35 insertions(+), 55 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 28614f0c3..42d472ce2 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -332,14 +332,14 @@ def __init__(self, *args, **kwargs):
     self.environment_spec = lambda: gym.make(self.env_name)
     self._real_env = None
     self.real_env_problem = None
-    self.in_graph_wrappers = []
+    self._internal_memory_size = 10
+
     self.collect_hparams = rl.ppo_pong_base()
     if FLAGS.autoencoder_path:
       self.collect_hparams = rl.ppo_pong_ae_base()
     self.settable_num_steps = 50000
     self.simulated_environment = None
     self.eval_phase = False
-    self.warm_up = 10  # TODO(piotrm): This should be probably removed.
 
     # Debug info.
     self.dones = 0
@@ -427,10 +427,7 @@ def _setup(self):
       self.sum_of_rewards = 0.0
       self.successful_episode_reward_predictions = 0
 
-    in_graph_wrappers = self.in_graph_wrappers + [
-        (atari.MemoryWrapper, {}), (StackAndSkipWrapper, {"skip": 4})]
     env_hparams = tf.contrib.training.HParams(
-        in_graph_wrappers=in_graph_wrappers,
         problem=self.real_env_problem if self.real_env_problem else self,
         simulated_environment=self.simulated_environment)
     if self.simulated_environment:
@@ -439,15 +436,8 @@ def _setup(self):
       env_hparams.add_hparam("intrinsic_reward_scale",
                              self.intrinsic_reward_scale)
 
-    generator_batch_env = batch_env_factory(
-        self.environment_spec, env_hparams, num_agents=1, xvfb=False)
-
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      if FLAGS.agent_policy_path:
-        policy_lambda = self.collect_hparams.network
-      else:
-        # When no agent_policy_path is set, just generate random samples.
-        policy_lambda = rl.random_policy_fun
+    # generator_batch_env = batch_env_factory(
+    #     self.environment_spec, env_hparams, num_agents=1, xvfb=False)
 
     if FLAGS.autoencoder_path:
       # TODO(lukaszkaiser): remove hard-coded autoencoder params.
@@ -482,19 +472,17 @@ def preprocess_fn(x):
                      not self.simulated_environment)
     preprocess = (preprocess_fn, shape) if do_preprocess else None
 
-    def policy(x):
-      return policy_lambda(self.environment_spec().action_space,
-                           self.collect_hparams, x)
+    if not FLAGS.agent_policy_path:
+      self.collect_hparams.policy_network = rl.random_policy_fun
 
+    self.collect_hparams.add_hparam("environment_spec", self.environment_spec)
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      self.collect_hparams.epoch_length = 10
-      _, self.collect_trigger_op = collect.define_collect(
-          policy, generator_batch_env, self.collect_hparams,
-          eval_phase=self.eval_phase,
-          scope="define_collect", preprocess=preprocess)
-
-    self.avilable_data_size_op = atari.MemoryWrapper.singleton.speculum.size()
-    self.data_get_op = atari.MemoryWrapper.singleton.speculum.dequeue()
+      self.collect_hparams.epoch_length = self._internal_memory_size
+      self.collect_hparams.num_agents = 1 #it is possible to set more
+      self.collect_memory, self.collect_trigger_op = collect.define_collect(
+          self.collect_hparams, scope="gym_problems",
+          collect_level=0,
+          eval_phase=self.eval_phase)
 
   def restore_networks(self, sess):
     if FLAGS.agent_policy_path:
@@ -532,11 +520,17 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
         # built and self.collect_trigger_op is all that's used from it.
         FLAGS.autoencoder_path = None
       pieces_generated = 0
-      while pieces_generated < self.num_steps + self.warm_up:
-        avilable_data_size = sess.run(self.avilable_data_size_op)
-        if avilable_data_size < 1:
+      memory_index = 0
+      memory = None
+      while pieces_generated < self.num_steps:
+        if memory is None or memory_index>=self._internal_memory_size:
           sess.run(self.collect_trigger_op)
-        observ, reward, action, done = sess.run(self.data_get_op)
+          memory = sess.run(self.collect_memory)
+          memory_index = 0
+        data = [memory[i][memory_index][0] for i in range(4)]
+        memory_index += 1
+        observ, reward, done, action = data
+        observ = observ.astype(np.uint8) # TODO(piotrmilos). This should be probably done in collect
         debug_im = None
         if self.make_extra_debug_info:
           self.total_sim_reward += reward
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index c0f9ceea7..a6c45a2b0 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -277,7 +277,7 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
             unencoded_debug = features.pop("image/debug")
             encoded_debug = sess.run(encoded_image_t, feed_dict={
                 image_t: unencoded_debug})
-            features["image/encoded_debug"] = encoded_debug
+            features["image/encoded_debug"] = [encoded_debug]
           yield features
 
   def generate_encoded_samples_debug(self, data_dir, tmp_dir, dataset_split):
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 8281cc8db..160cf197a 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -81,6 +81,14 @@ def standard_atari_parameters():
   wrappers = [[tf_atari_wrappers.MaxAndSkipWrapper, {"skip": 4}]]
   return tf.contrib.training.HParams(wrappers=wrappers)
 
+@registry.register_hparams
+def discrete_random_action_base():
+  hparams = common_hparams.basic_params1()
+  hparams.add_hparam("policy_network", random_policy_fun)
+  hparams.add_hparam("policy_network_params", "standard_atari_parameters")
+  return hparams
+
+
 @registry.register_hparams
 def ppo_atari_base():
   """Atari base parameters."""
@@ -113,7 +121,7 @@ def ppo_pong_base():
   hparams.optimization_epochs = 2
   hparams.epochs_num = 1000
   hparams.num_eval_agents = 1
-  hparams.network = feed_forward_cnn_small_categorical_fun
+  hparams.policy_network = feed_forward_cnn_small_categorical_fun
   hparams.clipping_coef = 0.2
   hparams.optimization_batch_size = 4
   hparams.max_gradients_norm = 0.5
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 7edec7324..3f7c56cfe 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -67,7 +67,7 @@ def simulate(self, action):
 
     with tf.control_dependencies([assign]):
       enqueue_op = self.speculum.enqueue(
-          [self._observ, reward, done, action])
+          [self._observ.read_value(), reward, done, action])
       with tf.control_dependencies([enqueue_op]):
         return tf.identity(reward), tf.identity(done)
 
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 32bfcbde6..fe6eedc1b 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -219,26 +219,4 @@ def _reset_non_empty(self, indices):
     assign_op = tf.scatter_update(self._observ, indices, new_values)
     with tf.control_dependencies([op_zero, assign_op]):
       return tf.identity(self.observ)
-
-
-class MemoryWrapper(WrapperBase):
-  """Memory wrapper."""
-
-  def __init__(self, batch_env):
-    super(MemoryWrapper, self).__init__(batch_env)
-    MemoryWrapper.singleton = self
-    assert self._length == 1, "We support only one environment"
-    infinity = 10000000
-    self.speculum = tf.FIFOQueue(infinity, dtypes=[
-        tf.uint8, tf.float32, tf.int32, tf.bool])
-    self._observ = self._batch_env.observ
-
-  def simulate(self, action):
-    with tf.name_scope("environment/simulate"):  # Do we need this?
-      reward, done = self._batch_env.simulate(action)
-      image = tf.cast(self._batch_env.observ[0, ...], tf.uint8)
-      with tf.control_dependencies([reward, done]):
-        enqueue_op = self.speculum.enqueue(
-            [image, reward, action, done])
-        with tf.control_dependencies([enqueue_op]):
-          return tf.identity(reward), tf.identity(done)
+    
\ No newline at end of file

From eb118831ddef16f38a4e26c9885e8a9518aef0e7 Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Sat, 16 Jun 2018 10:46:22 -0700
Subject: [PATCH 0123/2720] Remove deprecated Experiment API code, and support
 SessionRunHooks on TPU.

PiperOrigin-RevId: 200847884
---
 tensor2tensor/utils/t2t_model.py   |  51 +++++--
 tensor2tensor/utils/trainer_lib.py | 234 ++++++++++++++++++++---------
 2 files changed, 198 insertions(+), 87 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 2331a3341..a99997d3c 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -117,6 +117,16 @@ def __init__(self,
     if self._problem_hparams:
       self._create_modalities(self._problem_hparams, self._hparams)
 
+  # Replace the two methods below in order to add custom SessionRunHooks to
+  # the training procedure.
+  @staticmethod
+  def train_hooks():
+    return []
+
+  @staticmethod
+  def eval_hooks():
+    return []
+
   @property
   def hparams(self):
     return self._hparams
@@ -270,8 +280,8 @@ def bottom(self, features):
     all_previous_modalities = []
 
     # Transform the input features
-    for key, input_modality in sorted(six.iteritems(
-        self._problem_hparams.input_modality)):
+    for key, input_modality in sorted(
+        six.iteritems(self._problem_hparams.input_modality)):
       if key not in features:
         tf.logging.warning("Missing feature %s - ignoring." % key)
         continue
@@ -519,8 +529,8 @@ def eval_autoregressive(self, features=None, decode_length=50):
 
   def _fill_problem_hparams_features(self, features):
     if features is not None:
-      for k, v in sorted(six.iteritems(
-          problem_hparams_to_features(self._problem_hparams))):
+      for k, v in sorted(
+          six.iteritems(problem_hparams_to_features(self._problem_hparams))):
         if k not in features:
           features[k] = tf.constant(v, name=k)
 
@@ -973,11 +983,9 @@ def infer_step(recent_output, recent_logits, unused_loss):
       decode_length = 1
     else:
       if "partial_targets" in features:
-        prefix_length = common_layers.shape_list(
-            features["partial_targets"])[1]
+        prefix_length = common_layers.shape_list(features["partial_targets"])[1]
       else:
-        prefix_length = common_layers.shape_list(
-            features["inputs"])[1]
+        prefix_length = common_layers.shape_list(features["inputs"])[1]
       decode_length = prefix_length + decode_length
 
     # Initial values of result, logits and loss.
@@ -1112,6 +1120,16 @@ def _to_single_features_dict(self, datashard_features):
         features[k].append(v)
     return features
 
+  @staticmethod
+  def get_train_hooks(model_name):
+    model_cls = registry.model(model_name)
+    return model_cls.train_hooks()
+
+  @staticmethod
+  def get_eval_hooks(model_name):
+    model_cls = registry.model(model_name)
+    return model_cls.eval_hooks()
+
   @staticmethod
   def make_estimator_model_fn(model_name,
                               hparams,
@@ -1273,8 +1291,8 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
           # the key is located in the center of metric_name: "metrics-%s/%s/%s"
           k = metric_name.split("/")[1]
           if k in logits:
-            eval_metrics[metric_name] = metric_fn(
-                logits[k], features, features[k])
+            eval_metrics[metric_name] = metric_fn(logits[k], features,
+                                                  features[k])
           else:
             # We do not make it an error because we sometimes run models that
             # predict only parts of the targets defined by the Problem class.
@@ -1283,8 +1301,8 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
             # like actions or rewards.
             tf.logging.warning("No key %s in logits for evaluation." % k)
         else:
-          eval_metrics[metric_name] = metric_fn(
-              logits, features, features["targets"])
+          eval_metrics[metric_name] = metric_fn(logits, features,
+                                                features["targets"])
       if isinstance(logits, dict):
         predictions = logits
       else:
@@ -1339,8 +1357,7 @@ def estimator_spec_predict(self, features, use_tpu=False):
 
     if use_tpu:
       return tf.contrib.tpu.TPUEstimatorSpec(
-          tf.estimator.ModeKeys.PREDICT,
-          predictions=predictions)
+          tf.estimator.ModeKeys.PREDICT, predictions=predictions)
     else:
       return tf.estimator.EstimatorSpec(
           tf.estimator.ModeKeys.PREDICT,
@@ -1684,8 +1701,10 @@ def _compose_custom_getters(getter_a, getter_b):
     return getter_b
   if not getter_b:
     return getter_a
+
   def getter_fn(getter, *args, **kwargs):
     return getter_b(functools.partial(getter_a, getter), *args, **kwargs)
+
   return getter_fn
 
 
@@ -1698,5 +1717,5 @@ def set_custom_getter_compose(custom_getter):
     custom_getter: a custom getter.
   """
   tf.get_variable_scope().set_custom_getter(
-      _compose_custom_getters(
-          tf.get_variable_scope().custom_getter, custom_getter))
+      _compose_custom_getters(tf.get_variable_scope().custom_getter,
+                              custom_getter))
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index bbc6d15ce..f9781c3ae 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -75,8 +75,7 @@ def create_hparams(hparams_set,
   if problem_name:
     add_problem_hparams(hparams, problem_name)
   if hparams_overrides_str:
-    tf.logging.info("Overriding hparams in %s with %s",
-                    hparams_set,
+    tf.logging.info("Overriding hparams in %s with %s", hparams_set,
                     hparams_overrides_str)
     hparams = hparams.parse(hparams_overrides_str)
   return hparams
@@ -185,8 +184,9 @@ def create_estimator(model_name,
 
   if use_tpu:
     problem = hparams.problem
-    batch_size = (problem.tpu_batch_size_per_shard(hparams) *
-                  run_config.tpu_config.num_shards)
+    batch_size = (
+        problem.tpu_batch_size_per_shard(hparams) *
+        run_config.tpu_config.num_shards)
     predict_batch_size = batch_size
     if decode_hparams and decode_hparams.batch_size:
       predict_batch_size = decode_hparams.batch_size
@@ -202,16 +202,18 @@ def create_estimator(model_name,
         model_fn=model_fn, model_dir=run_config.model_dir, config=run_config)
 
 
-def create_hooks(use_tfdbg=False, use_dbgprofile=False, dbgprofile_kwargs=None,
-                 use_validation_monitor=False, validation_monitor_kwargs=None,
-                 use_early_stopping=False, early_stopping_kwargs=None):
+def create_hooks(use_tfdbg=False,
+                 use_dbgprofile=False,
+                 dbgprofile_kwargs=None,
+                 use_early_stopping=False,
+                 early_stopping_kwargs=None):
   """Create train and eval hooks for Experiment."""
-  train_monitors = []
+  train_hooks = []
   eval_hooks = []
 
   if use_tfdbg:
     hook = debug.LocalCLIDebugHook()
-    train_monitors.append(hook)
+    train_hooks.append(hook)
     eval_hooks.append(hook)
 
   if use_dbgprofile:
@@ -220,22 +222,124 @@ def create_hooks(use_tfdbg=False, use_dbgprofile=False, dbgprofile_kwargs=None,
     tf.logging.info("Using ProfilerHook")
     defaults = dict(save_steps=10, show_dataflow=True, show_memory=True)
     defaults.update(dbgprofile_kwargs)
-    train_monitors.append(tf.train.ProfilerHook(**defaults))
-
-  if use_validation_monitor:
-    tf.logging.info("Using ValidationMonitor")
-    train_monitors.append(
-        tf.contrib.learn.monitors.ValidationMonitor(
-            hooks=eval_hooks, **validation_monitor_kwargs))
+    train_hooks.append(tf.train.ProfilerHook(**defaults))
 
   if use_early_stopping:
     tf.logging.info("Using EarlyStoppingHook")
     hook = metrics_hook.EarlyStoppingHook(**early_stopping_kwargs)
     # Adding to both training and eval so that eval aborts as well
-    train_monitors.append(hook)
+    train_hooks.append(hook)
     eval_hooks.append(hook)
 
-  return train_monitors, eval_hooks
+  return train_hooks, eval_hooks
+
+
+class T2TExperiment(object):
+  """Custom Experiment class for running distributed experiments."""
+
+  def __init__(self, estimator, hparams, train_spec, eval_spec):
+    self._train_spec = train_spec
+    self._eval_spec = eval_spec
+    self._hparams = hparams
+    self._estimator = estimator
+
+  def continuous_train_and_eval(self):
+    tf.estimator.train_and_evaluate(self._estimator, self._train_spec,
+                                    self._eval_spec)
+
+  def train_and_evaluate(self):
+    tf.logging.warning(
+        "Note that train_and_evaluate now behaves the same as"
+        " continuous_train_and_eval. tensor2tensor no longer supports"
+        " training and evaluation in the same graph."
+    )
+    self.continuous_train_and_eval()
+
+  def train(self):
+    self._estimator.train(
+        self._train_spec.input_fn,
+        hooks=self._train_spec.hooks,
+        max_steps=self._train_spec.max_steps)
+
+  def evaluate(self):
+    self._estimator.evaluate(
+        self._eval_spec.input_fn,
+        steps=self._eval_spec.steps,
+        hooks=self._eval_spec.hooks)
+
+  def continuous_eval(self):
+    """Evaluate until checkpoints stop being produced."""
+    last_ckpt = None
+    while True:
+      # Wait up to half an hour for a new checkpoint
+      last_ckpt = tf.contrib.training.wait_for_new_checkpoint(
+          self._hparams.model_dir,
+          last_ckpt,
+          seconds_to_sleep=60,
+          timeout=60 * 30)
+
+      if last_ckpt is None:
+        raise Exception("Eval timeout: no new checkpoints within 30mins")
+
+      self._estimator.evaluate(
+          self._eval_spec.input_fn,
+          steps=self._eval_spec.steps,
+          hooks=self._eval_spec.hooks)
+
+  def continuous_eval_on_train_data(self):
+    """Evaluate on train data until checkpoints stop being produced."""
+    last_ckpt = None
+    while True:
+      # Wait up to half an hour for a new checkpoint
+      last_ckpt = tf.contrib.training.wait_for_new_checkpoint(
+          self._hparams.model_dir,
+          last_ckpt,
+          seconds_to_sleep=60,
+          timeout=60 * 30)
+
+      if last_ckpt is None:
+        raise Exception("Eval timeout: no new checkpoints within 30mins")
+
+      self._estimator.evaluate(
+          self._train_spec.input_fn,
+          steps=self._eval_spec.steps,
+          hooks=self._eval_spec.hooks)
+
+  def test(self):
+    """Perform 1 step of train and 2 step of eval."""
+    self._estimator.train(
+        self._train_spec.input_fn,
+        hooks=self._train_spec.hooks,
+        max_steps=1)
+
+    self._estimator.evaluate(
+        self._eval_spec.input_fn,
+        steps=1,
+        hooks=self._eval_spec.hooks)
+
+  def run_std_server(self):
+    """Starts a TensorFlow server and joins the serving thread.
+
+    Typically used for parameter servers.
+
+    Raises:
+      ValueError: if not enough information is available in the estimator's
+        config to create a server.
+    """
+    config = self._estimator.config
+    if (not config.cluster_spec or not config.task_type or not config.master or
+        config.task_id is None):
+      raise ValueError("Could not start server; be sure to specify "
+                       "cluster_spec, task_type, master, and task in "
+                       "RunConfig or set the TF_CONFIG environment variable.")
+    server = tf.train.Server(
+        config.cluster_spec,
+        job_name=config.task_type,
+        task_index=config.task_id,
+        config=config.tf_config,
+        start=False)
+    server.start()
+    server.join()
 
 
 def create_experiment(run_config,
@@ -262,6 +366,7 @@ def create_experiment(run_config,
   hparams.add_hparam("data_dir", data_dir)
   hparams.add_hparam("train_steps", train_steps)
   hparams.add_hparam("eval_steps", eval_steps)
+  hparams.add_hparam("schedule", schedule)
   add_problem_hparams(hparams, problem_name)
 
   # Estimator
@@ -275,10 +380,10 @@ def create_experiment(run_config,
 
   # Input fns from Problem
   problem = hparams.problem
-  train_input_fn = problem.make_estimator_input_fn(
-      tf.estimator.ModeKeys.TRAIN, hparams)
-  eval_input_fn = problem.make_estimator_input_fn(
-      tf.estimator.ModeKeys.EVAL, hparams)
+  train_input_fn = problem.make_estimator_input_fn(tf.estimator.ModeKeys.TRAIN,
+                                                   hparams)
+  eval_input_fn = problem.make_estimator_input_fn(tf.estimator.ModeKeys.EVAL,
+                                                  hparams)
 
   # Export
   if export:
@@ -286,55 +391,42 @@ def create_experiment(run_config,
                     "See serving/export.py.")
 
   # Hooks
-  hooks_kwargs = {}
-  if not use_tpu:
-    dbgprofile_kwargs = {"output_dir": run_config.model_dir}
-    validation_monitor_kwargs = dict(
-        input_fn=eval_input_fn,
-        eval_steps=eval_steps,
-        every_n_steps=min_eval_frequency,
-        early_stopping_rounds=eval_early_stopping_steps,
-        early_stopping_metric=eval_early_stopping_metric,
-        early_stopping_metric_minimize=eval_early_stopping_metric_minimize)
-    early_stopping_kwargs = dict(
-        events_dir=os.path.join(run_config.model_dir, "eval_continuous"),
-        tag=eval_early_stopping_metric,
-        num_plateau_steps=eval_early_stopping_steps,
-        plateau_decrease=eval_early_stopping_metric_minimize,
-        plateau_delta=eval_early_stopping_metric_delta,
-        every_n_steps=min_eval_frequency)
-
-    # In-process eval (and possible early stopping)
-    if schedule == "continuous_train_and_eval" and min_eval_frequency:
-      tf.logging.warn("ValidationMonitor only works with "
-                      "--schedule=train_and_evaluate")
-    use_validation_monitor = (
-        schedule == "train_and_evaluate" and min_eval_frequency)
-    # Distributed early stopping
-    local_schedules = ["train_and_evaluate", "continuous_train_and_eval"]
-    use_early_stopping = (
-        schedule not in local_schedules and eval_early_stopping_steps)
-    train_monitors, eval_hooks = create_hooks(
-        use_tfdbg=use_tfdbg,
-        use_dbgprofile=use_dbgprofile,
-        dbgprofile_kwargs=dbgprofile_kwargs,
-        use_validation_monitor=use_validation_monitor,
-        use_early_stopping=use_early_stopping,
-        validation_monitor_kwargs=validation_monitor_kwargs,
-        early_stopping_kwargs=early_stopping_kwargs)
-    hooks_kwargs = {"train_monitors": train_monitors, "eval_hooks": eval_hooks}
-
-  # Experiment
-  return tf.contrib.learn.Experiment(
-      estimator=estimator,
-      train_input_fn=train_input_fn,
-      eval_input_fn=eval_input_fn,
-      train_steps=train_steps,
-      eval_steps=eval_steps,
-      min_eval_frequency=min_eval_frequency,
-      train_steps_per_iteration=min(min_eval_frequency, train_steps),
-      eval_delay_secs=0 if schedule == "evaluate" else 120,
-      **hooks_kwargs)
+  dbgprofile_kwargs = {"output_dir": run_config.model_dir}
+  early_stopping_kwargs = dict(
+      events_dir=os.path.join(run_config.model_dir, "eval_continuous"),
+      tag=eval_early_stopping_metric,
+      num_plateau_steps=eval_early_stopping_steps,
+      plateau_decrease=eval_early_stopping_metric_minimize,
+      plateau_delta=eval_early_stopping_metric_delta,
+      every_n_steps=min_eval_frequency)
+
+  # In-process eval (and possible early stopping)
+  if schedule == "continuous_train_and_eval" and min_eval_frequency:
+    tf.logging.warn("ValidationMonitor only works with "
+                    "--schedule=train_and_evaluate")
+  # Distributed early stopping
+  local_schedules = ["train_and_evaluate", "continuous_train_and_eval"]
+  use_early_stopping = (
+      schedule not in local_schedules and eval_early_stopping_steps)
+  train_hooks, eval_hooks = create_hooks(
+      use_tfdbg=use_tfdbg,
+      use_dbgprofile=use_dbgprofile,
+      dbgprofile_kwargs=dbgprofile_kwargs,
+      use_early_stopping=use_early_stopping,
+      early_stopping_kwargs=early_stopping_kwargs)
+  train_hooks += t2t_model.T2TModel.get_train_hooks(model_name)
+  eval_hooks += t2t_model.T2TModel.get_eval_hooks(model_name)
+
+  train_spec = tf.estimator.TrainSpec(
+      train_input_fn, max_steps=train_steps, hooks=train_hooks)
+  eval_spec = tf.estimator.EvalSpec(
+      eval_input_fn,
+      steps=eval_steps,
+      hooks=eval_hooks,
+      start_delay_secs=0 if hparams.schedule == "evaluate" else 120,
+      throttle_secs=600)
+
+  return T2TExperiment(estimator, hparams, train_spec, eval_spec)
 
 
 def create_experiment_fn(*args, **kwargs):

From 214a3cc6a4725f5bf8fe29606f7e632388edc376 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sun, 17 Jun 2018 20:58:38 -0700
Subject: [PATCH 0124/2720] preventing datasets from batching between videos.

PiperOrigin-RevId: 200934714
---
 .../data_generators/video_generated.py        | 15 +----
 tensor2tensor/data_generators/video_utils.py  | 61 ++++++++++++++++++-
 tensor2tensor/models/research/next_frame.py   |  7 ++-
 3 files changed, 66 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index eef11cc38..8c5ed756f 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -22,7 +22,6 @@
 
 import numpy as np
 
-from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
@@ -54,7 +53,8 @@ def frame_width(self):
 
   @property
   def total_number_of_frames(self):
-    return 10000
+    # 10k videos
+    return 10000 * self.video_length
 
   @property
   def video_length(self):
@@ -69,17 +69,6 @@ def eval_metrics(self):
                     metrics.Metrics.IMAGE_RMSE]
     return eval_metrics
 
-  @property
-  def dataset_splits(self):
-    """Splits of data to produce and number of output shards for each."""
-    return [{
-        "split": problem.DatasetSplit.TRAIN,
-        "shards": 1,
-    }, {
-        "split": problem.DatasetSplit.EVAL,
-        "shards": 1,
-    }]
-
   @property
   def extra_reading_spec(self):
     """Additional data fields to store on disk and their decoders."""
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index c0f9ceea7..21fec5358 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -100,6 +100,14 @@ def dataset_splits(self):
         "shards": 1,
     }]
 
+  @property
+  def only_keep_videos_from_0th_frame(self):
+    return True
+
+  @property
+  def use_not_breaking_batching(self):
+    return False
+
   def preprocess_example(self, example, mode, hparams):
     """Runtime preprocessing, e.g., resize example["frame"]."""
     return example
@@ -192,6 +200,52 @@ def features_from_batch(batched_prefeatures):
     # Batch and construct features.
     def _preprocess(example):
       return self.preprocess_example(example, mode, hparams)
+
+    def avoid_break_batching(dataset):
+      """Smart preprocessing to avoid break between videos!
+
+      Simple batching of images into videos may result into broken videos
+      with two parts from two different videos. This preprocessing avoids
+      this using the frame number.
+
+      Args:
+        dataset: raw not-batched dataset.
+
+      Returns:
+        batched not-broken videos.
+
+      """
+      def check_integrity_and_batch(*datasets):
+        """Checks whether a sequence of frames are from the same video.
+
+        Args:
+          *datasets: datasets each skipping 1 frame from the previous one.
+
+        Returns:
+          batched data and the integrity flag.
+        """
+        frame_numbers = [dataset["frame_number"][0] for dataset in datasets]
+
+        not_broken = tf.equal(
+            frame_numbers[-1] - frame_numbers[0], num_frames-1)
+        if self.only_keep_videos_from_0th_frame:
+          not_broken = tf.logical_and(not_broken, tf.equal(frame_numbers[0], 0))
+
+        features = {}
+        for key in datasets[0].keys():
+          values = [dataset[key] for dataset in datasets]
+          batch = tf.stack(values)
+          features[key] = batch
+        return features, not_broken
+
+      ds = [dataset.skip(i) for i in range(num_frames)]
+      dataset = tf.data.Dataset.zip(tuple(ds))
+      dataset = dataset.map(check_integrity_and_batch)
+      dataset = dataset.filter(lambda _, not_broken: not_broken)
+      dataset = dataset.map(lambda features, _: features)
+
+      return dataset
+
     preprocessed_dataset = dataset.map(_preprocess)
     num_frames = (hparams.video_num_input_frames +
                   hparams.video_num_target_frames)
@@ -199,8 +253,11 @@ def _preprocess(example):
     if self.random_skip:
       random_skip = tf.random_uniform([], maxval=num_frames, dtype=tf.int64)
       preprocessed_dataset = preprocessed_dataset.skip(random_skip)
-    batch_dataset = preprocessed_dataset.apply(
-        tf.contrib.data.batch_and_drop_remainder(num_frames))
+    if self.use_not_breaking_batching:
+      batch_dataset = avoid_break_batching(preprocessed_dataset)
+    else:
+      batch_dataset = preprocessed_dataset.apply(
+          tf.contrib.data.batch_and_drop_remainder(num_frames))
     dataset = batch_dataset.map(features_from_batch).shuffle(8)
     return dataset
 
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index ed789081e..73f19f325 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -658,10 +658,13 @@ def body(self, features):
 
     all_actions = input_actions + target_actions
     all_rewards = input_rewards + target_rewards
+    all_frames = input_frames + target_frames
+
+    tf.summary.image("full_video", tf.concat(all_frames, axis=1))
 
     is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
     gen_images, gen_rewards, latent_mean, latent_std = self.construct_model(
-        images=input_frames + target_frames,
+        images=all_frames,
         actions=all_actions,
         rewards=all_rewards,
         k=900.0 if is_training else -1.0,
@@ -730,7 +733,7 @@ def next_frame():
 def next_frame_stochastic():
   """SV2P model."""
   hparams = next_frame()
-  hparams.video_num_input_frames = 4
+  hparams.video_num_input_frames = 2
   hparams.video_num_target_frames = 1
   hparams.batch_size = 8
   hparams.target_modality = "video:l2raw"

From 173dbf74697814f2d32c1b1e7d611847afd1efbf Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 18 Jun 2018 11:25:14 +0200
Subject: [PATCH 0125/2720] Environment spec introduced

---
 tensor2tensor/data_generators/gym_problems.py |  3 --
 tensor2tensor/models/research/rl.py           | 33 ++++++++++---
 tensor2tensor/rl/collect.py                   |  6 +--
 tensor2tensor/rl/envs/simulated_batch_env.py  |  2 +-
 tensor2tensor/rl/envs/utils.py                | 49 ++++++++++++-------
 tensor2tensor/rl/rl_trainer_lib_test.py       | 12 +++--
 6 files changed, 68 insertions(+), 37 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 42d472ce2..3208e5adf 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -436,9 +436,6 @@ def _setup(self):
       env_hparams.add_hparam("intrinsic_reward_scale",
                              self.intrinsic_reward_scale)
 
-    # generator_batch_env = batch_env_factory(
-    #     self.environment_spec, env_hparams, num_agents=1, xvfb=False)
-
     if FLAGS.autoencoder_path:
       # TODO(lukaszkaiser): remove hard-coded autoencoder params.
       with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 160cf197a..65f845420 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -73,22 +73,14 @@ def basic_policy_parameters():
 def ppo_discrete_action_base():
   hparams = ppo_base_v1()
   hparams.add_hparam("policy_network", feed_forward_categorical_fun)
-  hparams.add_hparam("policy_network_params", "standard_atari_parameters")
   return hparams
 
-@registry.register_hparams
-def standard_atari_parameters():
-  wrappers = [[tf_atari_wrappers.MaxAndSkipWrapper, {"skip": 4}]]
-  return tf.contrib.training.HParams(wrappers=wrappers)
-
 @registry.register_hparams
 def discrete_random_action_base():
   hparams = common_hparams.basic_params1()
   hparams.add_hparam("policy_network", random_policy_fun)
-  hparams.add_hparam("policy_network_params", "standard_atari_parameters")
   return hparams
 
-
 @registry.register_hparams
 def ppo_atari_base():
   """Atari base parameters."""
@@ -127,6 +119,31 @@ def ppo_pong_base():
   hparams.max_gradients_norm = 0.5
   return hparams
 
+EnvironmentSpec = collections.namedtuple('EnvironmentSpec', 'env_lambda, wrappers, simulated_env')
+
+def standard_atari_env_spec(env):
+  """Parameters of environement specification"""
+  standard_wrappers = [[tf_atari_wrappers.MaxAndSkipWrapper, {"skip": 4}]]
+  env_lambda = None
+  if isinstance(env, str):
+    env_lambda = lambda: gym.make(env)
+  if callable(env):
+    env_lambda = env
+  assert env is not None, "Unknown specification of environment"
+
+  return EnvironmentSpec(env_lambda=env_lambda, wrappers=standard_wrappers, simulated_env=False)
+
+def simple_gym_spec(env):
+  """Parameters of environement specification"""
+  standard_wrappers = None
+  env_lambda = None
+  if isinstance(env, str):
+    env_lambda = lambda : gym.make(env)
+  if callable(env):
+    env_lambda = env
+  assert env is not None, "Unknown specification of environment"
+
+  return EnvironmentSpec(env_lambda=env_lambda, wrappers=standard_wrappers, simulated_env=False)
 
 @registry.register_hparams
 def ppo_pong_ae_base():
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 3f7c56cfe..757241bfe 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -60,8 +60,6 @@ def __init__(self, batch_env):
 
   def simulate(self, action):
     reward, done = self._batch_env.simulate(action)
-    # # image = tf.cast(self._batch_env.observ[0, ...], tf.uint8)
-    # image = self._batch_env.observ #possibly remove
     with tf.control_dependencies([reward, done]):
       assign = self._observ.assign(self._batch_env.observ)
 
@@ -80,8 +78,8 @@ def define_collect(hparams, scope, eval_phase,
 
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
     batch_env = utils.batch_env_factory(hparams)
-    policy_hparams = registry.hparams(hparams.policy_network_params)
-    wrappers = copy.copy(policy_hparams.wrappers) if policy_hparams.wrappers else []
+    environment_wrappers = hparams.environment_spec.wrappers
+    wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
     #Put memory wrapper at the level you want to gather observations at
     #Negative indices need to be shifted for insert to work correctly
     collect_level = collect_level if collect_level>=0 else len(wrappers) + collect_level + 1
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 8f0920e19..03a769c24 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -106,7 +106,7 @@ def __init__(self, hparams, length, problem,
     self._model = registry.model(FLAGS.model)(
       model_hparams, tf.estimator.ModeKeys.PREDICT)
 
-    _, self.action_shape, self.action_dtype = utils.get_action_space(hparams)
+    _, self.action_shape, self.action_dtype = utils.get_action_space(hparams.environment_spec)
 
     if simulation_random_starts:
       dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index dfc3de8e9..ffabb0464 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -281,27 +281,26 @@ def _worker(self, constructor, conn):
 def batch_env_factory(hparams, xvfb=False):
   """Factory of batch envs."""
 
-  if hparams.simulated_environment:
+  environment_spec = hparams.environment_spec
+
+  if environment_spec.simulated_env:
     cur_batch_env = _define_simulated_batch_env(
-      hparams, hparams.num_agents, hparams.problem,
+      environment_spec, hparams.num_agents, hparams.problem,
         hparams.simulation_random_starts,
         hparams.intrinsic_reward_scale)
   else:
-    cur_batch_env = _define_batch_env(hparams, xvfb=xvfb)
+    cur_batch_env = _define_batch_env(environment_spec, hparams.num_agents,
+                                      xvfb=xvfb)
   return cur_batch_env
 
 
-def _define_batch_env(hparams, xvfb=False):
+def _define_batch_env(environment_spec, num_agents, xvfb=False):
   """Create environments and apply all desired wrappers."""
-  if isinstance(hparams.environment_spec, str):
-    environment_lambda = lambda: gym.make(hparams.environment_spec)
-  else:
-    environment_lambda = hparams.environment_spec
 
   with tf.variable_scope("environments"):
     envs = [
-        ExternalProcessEnv(environment_lambda, xvfb)
-        for _ in range(hparams.num_agents)]
+        ExternalProcessEnv(environment_spec.env_lambda, xvfb)
+        for _ in range(num_agents)]
     env = batch_env.BatchEnv(envs, blocking=False)
     env = py_func_batch_env.PyFuncBatchEnv(env)
     return env
@@ -316,23 +315,37 @@ def _define_simulated_batch_env(hparams, num_agents, problem,
   return cur_batch_env
 
 
-def get_action_space(hparams):
-  if isinstance(hparams.environment_spec, str):
-    environment_lambda = lambda: gym.make(hparams.environment_spec)
-  else:
-    environment_lambda = hparams.environment_spec
-
-  action_space = environment_lambda().action_space
+def get_action_space(environment_spec):
+  """Get action spece associated with environment spec
+    
+  Args:
+     environment_spec:  EnvironmentSpec object
+     
+  Returns:
+    OpenAi Gym action spece
+  """
+  action_space = environment_spec.env_lambda().action_space
   action_shape = list(parse_shape(action_space))
   action_dtype = parse_dtype(action_space)
 
   return action_space, action_shape, action_dtype
 
+
 def get_policy(observations, hparams):
+  """Get policy network
+  
+  Args:
+    observations: Tensor with observations
+    hparams: parameters 
+    
+  Returns:
+    Tensor with policy and value function output
+  """
   policy_network_lambda = hparams.policy_network
-  action_space, _, _ = get_action_space(hparams)
+  action_space, _, _ = get_action_space(hparams.environment_spec)
   return policy_network_lambda(action_space, hparams, observations)
 
+
 def parse_shape(space):
   """Get a tensor shape from a OpenAI Gym space.
 
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index 396ce4292..ca400e34a 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -16,8 +16,10 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+from tensor2tensor.models.research.rl import simple_gym_spec, standard_atari_env_spec
 from tensor2tensor.rl import rl_trainer_lib
-from tensor2tensor.utils import trainer_lib
+from tensor2tensor.utils import trainer_lib, registry
 
 import tensorflow as tf
 
@@ -31,14 +33,18 @@ def test_no_crash_pendulum(self):
     hparams = trainer_lib.create_hparams(
         "ppo_continuous_action_base",
         TrainTest.test_config)
-    hparams.add_hparam("environment_spec", "Pendulum-v0")
+
+    hparams.add_hparam("environment_spec", simple_gym_spec("Pendulum-v0"))
+    hparams.add_hparam("environment_wrappers", None)
     rl_trainer_lib.train(hparams)
 
   def test_no_crash_cartpole(self):
     hparams = trainer_lib.create_hparams(
         "ppo_discrete_action_base",
         TrainTest.test_config)
-    hparams.add_hparam("environment_spec", "CartPole-v0")
+
+    hparams.add_hparam("environment_spec", standard_atari_env_spec("CartPole-v0"))
+
     rl_trainer_lib.train(hparams)
 
 

From 6086ca58fee4db89c9f2b12ebdbaf01bc365686b Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 18 Jun 2018 13:56:55 +0200
Subject: [PATCH 0126/2720] before big refactor of gym_problems

---
 tensor2tensor/data_generators/gym_problems.py |  5 +++--
 tensor2tensor/rl/envs/utils.py                |  8 ++++---
 tensor2tensor/rl/model_rl_experiment.py       | 21 ++++++++++++-------
 tensor2tensor/rl/rl_trainer_lib_test.py       |  2 --
 4 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 3208e5adf..5cc060e3f 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -30,6 +30,7 @@
 from tensor2tensor.layers import discretization
 from tensor2tensor.models.research import autoencoders
 from tensor2tensor.models.research import rl
+from tensor2tensor.models.research.rl import standard_atari_env_spec
 from tensor2tensor.rl import collect
 from tensor2tensor.rl.envs import tf_atari_wrappers as atari
 from tensor2tensor.rl.envs.tf_atari_wrappers import StackAndSkipWrapper
@@ -329,12 +330,13 @@ def __init__(self, *args, **kwargs):
     self.autoencoder_model = None
 
     # Defaults.
-    self.environment_spec = lambda: gym.make(self.env_name)
+    self.environment_spec = standard_atari_env_spec(self.env_name)
     self._real_env = None
     self.real_env_problem = None
     self._internal_memory_size = 10
 
     self.collect_hparams = rl.ppo_pong_base()
+    self.collect_hparams.add_hparam("environment_spec", self.environment_spec)
     if FLAGS.autoencoder_path:
       self.collect_hparams = rl.ppo_pong_ae_base()
     self.settable_num_steps = 50000
@@ -472,7 +474,6 @@ def preprocess_fn(x):
     if not FLAGS.agent_policy_path:
       self.collect_hparams.policy_network = rl.random_policy_fun
 
-    self.collect_hparams.add_hparam("environment_spec", self.environment_spec)
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       self.collect_hparams.epoch_length = self._internal_memory_size
       self.collect_hparams.num_agents = 1 #it is possible to set more
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index ffabb0464..c142333c4 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -284,12 +284,14 @@ def batch_env_factory(hparams, xvfb=False):
   environment_spec = hparams.environment_spec
 
   if environment_spec.simulated_env:
+    # TODO(piotrmilos): Consider passing only relevant paramters
     cur_batch_env = _define_simulated_batch_env(
-      environment_spec, hparams.num_agents, hparams.problem,
+      hparams, hparams.num_agents, hparams.problem,
         hparams.simulation_random_starts,
         hparams.intrinsic_reward_scale)
   else:
-    cur_batch_env = _define_batch_env(environment_spec, hparams.num_agents,
+
+    cur_batch_env = _define_batch_env(hparams.environment_spec, hparams.num_agents,
                                       xvfb=xvfb)
   return cur_batch_env
 
@@ -322,7 +324,7 @@ def get_action_space(environment_spec):
      environment_spec:  EnvironmentSpec object
      
   Returns:
-    OpenAi Gym action spece
+    OpenAi Gym action space
   """
   action_space = environment_spec.env_lambda().action_space
   action_shape = list(parse_shape(action_space))
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 7e55ec89a..4f7e328bb 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -30,9 +30,13 @@
 import math
 import os
 import time
+
+import copy
+
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.layers import discretization
+from tensor2tensor.models.research.rl import EnvironmentSpec
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.rl.envs.tf_atari_wrappers import StackAndSkipWrapper
 from tensor2tensor.rl.envs.tf_atari_wrappers import TimeLimitWrapper
@@ -139,7 +143,6 @@ def train_agent(problem_name, agent_model_dir,
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_epochs_num = hparams.ppo_epochs_num
   ppo_hparams.epochs_num = ppo_epochs_num
-  ppo_hparams.simulated_environment = True
   ppo_hparams.simulation_random_starts = hparams.simulation_random_starts
   ppo_hparams.intrinsic_reward_scale = hparams.intrinsic_reward_scale
   ppo_hparams.eval_every_epochs = 50
@@ -150,14 +153,17 @@ def train_agent(problem_name, agent_model_dir,
   ppo_hparams.world_model_dir = world_model_dir
   if hparams.ppo_learning_rate:
     ppo_hparams.learning_rate = hparams.ppo_learning_rate
+
+  environment_spec = gym_problem.environment_spec
   # 4x for the StackAndSkipWrapper minus one to always finish for reporting.
   ppo_time_limit = (ppo_hparams.epoch_length - 1) * 4
+  wrappers = copy.copy(environment_spec.wrappers)
+  wrappers.append([TimeLimitWrapper, {"timelimit": ppo_time_limit}])
+  ppo_hparams.add_hparam("environment_spec",
+                         EnvironmentSpec(env_lambda=environment_spec.env_lambda,
+                                         wrappers=wrappers,
+                                         simulated_env=True))
 
-  in_graph_wrappers = [
-      (TimeLimitWrapper, {"timelimit": ppo_time_limit}),
-      (StackAndSkipWrapper, {"skip": 4})]
-  in_graph_wrappers += gym_problem.in_graph_wrappers
-  ppo_hparams.add_hparam("in_graph_wrappers", in_graph_wrappers)
 
   with temporary_flags({
       "problem": problem_name,
@@ -167,8 +173,7 @@ def train_agent(problem_name, agent_model_dir,
       "data_dir": epoch_data_dir,
       "autoencoder_path": autoencoder_path,
   }):
-    rl_trainer_lib.train(ppo_hparams, gym_problem.env_name, event_dir,
-                         agent_model_dir, epoch=epoch)
+    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=epoch)
 
 
 def evaluate_world_model(simulated_problem_name, problem_name, hparams,
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index ca400e34a..fba44c31e 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -35,7 +35,6 @@ def test_no_crash_pendulum(self):
         TrainTest.test_config)
 
     hparams.add_hparam("environment_spec", simple_gym_spec("Pendulum-v0"))
-    hparams.add_hparam("environment_wrappers", None)
     rl_trainer_lib.train(hparams)
 
   def test_no_crash_cartpole(self):
@@ -44,7 +43,6 @@ def test_no_crash_cartpole(self):
         TrainTest.test_config)
 
     hparams.add_hparam("environment_spec", standard_atari_env_spec("CartPole-v0"))
-
     rl_trainer_lib.train(hparams)
 
 

From 34fca4c556795725227336490ef1e934314d49db Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 18 Jun 2018 16:38:26 +0200
Subject: [PATCH 0127/2720] refactor of gym_problems

---
 tensor2tensor/data_generators/gym_problems.py | 482 ++++++------------
 tensor2tensor/models/research/rl.py           |  10 +-
 tensor2tensor/rl/envs/simulated_batch_env.py  |   2 +-
 tensor2tensor/rl/model_rl_experiment.py       |  21 +-
 4 files changed, 174 insertions(+), 341 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 5cc060e3f..6304b48a8 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -62,6 +62,85 @@ class GymDiscreteProblem(video_utils.VideoProblem):
   def __init__(self, *args, **kwargs):
     super(GymDiscreteProblem, self).__init__(*args, **kwargs)
     self._env = None
+    self._env = None
+    self.debug_dump_frames_path = "debug_frames_env"
+    self.settable_num_steps = 5000
+
+    self.environment_spec = self.get_environment_spec()
+    self.eval_phase = False
+
+    # Debug info
+    self.make_extra_debug_info = True
+    self.dones = 0
+    self.real_reward = 0
+    self.total_sim_reward = 0.0
+    self.total_real_reward = 0.0
+    self.sum_of_rewards = 0.0
+    self.successful_episode_reward_predictions = 0
+    self.report_reward_statistics_every = 10
+
+  def _setup(self):
+    collect_hparams = rl.ppo_pong_base()
+    collect_hparams.add_hparam("environment_spec", self.environment_spec)
+
+    if not FLAGS.agent_policy_path:
+      collect_hparams.policy_network = rl.random_policy_fun
+
+    self._internal_memory_size = 10
+    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+      collect_hparams.epoch_length = self._internal_memory_size
+      collect_hparams.num_agents = 1 #TODO (piotrmilos). it is possible to set more
+      self.collect_memory, self.collect_trigger_op \
+        = collect.define_collect(collect_hparams, scope="gym_problems",
+                                 collect_level=0, eval_phase=self.eval_phase)
+
+  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
+    self._setup()
+    self.debug_dump_frames_path = os.path.join(
+        data_dir, self.debug_dump_frames_path)
+
+    with tf.Session() as sess:
+      sess.run(tf.global_variables_initializer())
+      self.restore_networks(sess)
+      pieces_generated = 0
+      memory_index = 0
+      memory = None
+      while pieces_generated < self.num_steps:
+        if memory is None or memory_index>=self._internal_memory_size:
+          sess.run(self.collect_trigger_op)
+          memory = sess.run(self.collect_memory)
+          memory_index = 0
+        data = [memory[i][memory_index][0] for i in range(4)]
+        memory_index += 1
+        observ, reward, done, action = data
+        observ = observ.astype(np.uint8) # TODO(piotrmilos). This should be probably done in collect
+        debug_im = None
+        if self.make_extra_debug_info:
+          debug_im = self.get_debug_image(data)
+
+        ret_dict = {"frame": observ,
+                    "image/format": ["png"],
+                    "image/height": [self.frame_height],
+                    "image/width": [self.frame_width],
+                    "action": [int(action)],
+                    "done": [int(False)],
+                    "reward": [int(reward) - self.min_reward]}
+        if self.make_extra_debug_info:
+          ret_dict["image/debug"] = debug_im
+        yield ret_dict
+        pieces_generated += 1
+
+  def get_debug_image(self, data):
+    raise NotImplemented()
+
+  def restore_networks(self, sess):
+    if FLAGS.agent_policy_path:
+      model_saver = tf.train.Saver(
+          tf.global_variables(".*network_parameters.*"))
+      ckpts = tf.train.get_checkpoint_state(FLAGS.agent_policy_path)
+      ckpt = ckpts.model_checkpoint_path
+      model_saver.restore(sess, ckpt)
+
 
   @property
   def num_input_frames(self):
@@ -88,6 +167,9 @@ def extra_reading_spec(self):
     }
     return data_fields, decoders
 
+  def get_environment_spec(self):
+    raise NotImplementedError()
+
   @property
   def is_generate_per_split(self):
     """Whether we have a train/test split or just hold out data."""
@@ -122,7 +204,7 @@ def num_rewards(self):
 
   @property
   def num_steps(self):
-    raise NotImplementedError()
+    return self.settable_num_steps
 
   @property
   def total_number_of_frames(self):
@@ -155,22 +237,9 @@ def hparams(self, defaults, unused_model_hparams):
     p.input_space_id = problem.SpaceID.IMAGE
     p.target_space_id = problem.SpaceID.IMAGE
 
-  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
-    next_observation = self.env.reset()
-    for _ in range(self.num_steps):
-      observation = next_observation
-      action = self.get_action(observation)
-      next_observation, reward, done, _ = self.env.step(action)
-      if done:
-        next_observation = self.env.reset()
-      yield {"frame": observation,
-             "action": [action],
-             "done": [done],
-             "reward": [int(reward - self.min_reward)]}
-
 
 @registry.register_problem
-class GymPongRandom5k(GymDiscreteProblem):
+class GymPongRandom(GymDiscreteProblem):
   """Pong game, random actions."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
@@ -199,22 +268,9 @@ def min_reward(self):
   def num_rewards(self):
     return 3
 
-  @property
-  def num_steps(self):
-    return 5000
-
 
 @registry.register_problem
-class GymPongRandom50k(GymPongRandom5k):
-  """Pong game, random actions."""
-
-  @property
-  def num_steps(self):
-    return 50000
-
-
-@registry.register_problem
-class GymWrappedPongRandom5k(GymDiscreteProblem):
+class GymWrappedPongRandom(GymDiscreteProblem):
   """Pong game, random actions."""
 
   @property
@@ -229,10 +285,6 @@ def min_reward(self):
   def num_rewards(self):
     return 3
 
-  @property
-  def num_steps(self):
-    return 5000
-
 
 @registry.register_problem
 class GymWrappedLongPongRandom(GymDiscreteProblem):
@@ -250,17 +302,13 @@ def min_reward(self):
   def num_rewards(self):
     return 3
 
-  @property
-  def num_steps(self):
-    return 5000
-
   @property
   def num_testing_steps(self):
     return 100
 
 
 @registry.register_problem
-class GymWrappedBreakoutRandom5k(GymDiscreteProblem):
+class GymWrappedBreakoutRandom(GymDiscreteProblem):
   """Pong game, random actions."""
 
   @property
@@ -275,22 +323,9 @@ def min_reward(self):
   def num_rewards(self):
     return 3
 
-  @property
-  def num_steps(self):
-    return 5000
-
 
 @registry.register_problem
-class GymWrappedPongRandom50k(GymPongRandom5k):
-  """Pong game, random actions."""
-
-  @property
-  def num_steps(self):
-    return 50000
-
-
-@registry.register_problem
-class GymFreewayRandom5k(GymDiscreteProblem):
+class GymFreewayRandom(GymDiscreteProblem):
   """Freeway game, random actions."""
 
   @property
@@ -305,290 +340,76 @@ def min_reward(self):
   def num_rewards(self):
     return 2
 
-  @property
-  def num_steps(self):
-    return 5000
-
-
-@registry.register_problem
-class GymFreewayRandom50k(GymFreewayRandom5k):
-  """Freeway game, random actions."""
-
-  @property
-  def num_steps(self):
-    return 50000
-
 
-class GymDiscreteProblemWithAgent(GymDiscreteProblem):
-  """Gym environment with discrete actions and rewards and an agent."""
+class GymRealDiscreteProblem(GymDiscreteProblem):
 
   def __init__(self, *args, **kwargs):
-    super(GymDiscreteProblemWithAgent, self).__init__(*args, **kwargs)
-    self._env = None
-    self.debug_dump_frames_path = "debug_frames_env"
-    self.make_extra_debug_info = True
-    self.autoencoder_model = None
-
-    # Defaults.
-    self.environment_spec = standard_atari_env_spec(self.env_name)
-    self._real_env = None
-    self.real_env_problem = None
-    self._internal_memory_size = 10
-
-    self.collect_hparams = rl.ppo_pong_base()
-    self.collect_hparams.add_hparam("environment_spec", self.environment_spec)
-    if FLAGS.autoencoder_path:
-      self.collect_hparams = rl.ppo_pong_ae_base()
-    self.settable_num_steps = 50000
-    self.simulated_environment = None
-    self.eval_phase = False
-
-    # Debug info.
-    self.dones = 0
-    self.real_reward = 0
-    self.total_sim_reward, self.total_real_reward = 0.0, 0.0
-    self.sum_of_rewards = 0.0
-    self.successful_episode_reward_predictions = 0
+    super(GymRealDiscreteProblem, self).__init__(*args, **kwargs)
+    self.make_extra_debug_info = False
 
-  @property
-  def real_env(self):
-    """Lazy caching environment construction."""
-    if self._real_env is None:
-      self._real_env = self.environment_spec()
-    return self._real_env
+  def get_debug_image(self):
+    #TODO(piotrmilos): possibly change this
+    raise NotImplemented()
 
-  @property
-  def num_steps(self):
-    return self.settable_num_steps
+  def get_environment_spec(self):
+    return standard_atari_env_spec(self.env_name)
 
-  @property
-  def raw_frame_height(self):
-    return self.env.observation_space.shape[0]
 
-  @property
-  def frame_height(self):
-    if FLAGS.autoencoder_path:
-      # TODO(lukaszkaiser): remove hard-coded autoencoder params.
-      return int(math.ceil(self.raw_frame_height / self.autoencoder_factor))
-    return self.raw_frame_height
-
-  @property
-  def autoencoder_factor(self):
-    """By how much to divide sizes when using autoencoders."""
-    hparams = autoencoders.autoencoder_discrete_pong()
-    return 2**hparams.num_hidden_layers
+class GymAEDiscreteProblem(GymDiscreteProblem):
+  pass
 
-  @property
-  def raw_frame_width(self):
-    return self.env.observation_space.shape[1]
+@registry.register_problem
+class GymSimulatedDiscreteProblem(GymDiscreteProblem):
+  """Simulated gym environment with discrete actions and rewards."""
 
-  @property
-  def frame_width(self):
-    if FLAGS.autoencoder_path:
-      # TODO(lukaszkaiser): remove hard-coded autoencoder params.
-      return int(math.ceil(self.raw_frame_width / self.autoencoder_factor))
-    return self.raw_frame_width
-
-  def setup_autoencoder(self):
-    if self.autoencoder_model is not None:
-      return
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      autoencoder_hparams = autoencoders.autoencoder_discrete_pong()
-      autoencoder_hparams.data_dir = "unused"
-      autoencoder_hparams.problem_hparams = self.get_hparams(
-          autoencoder_hparams)
-      autoencoder_hparams.problem = self
-      self.autoencoder_model = autoencoders.AutoencoderOrderedDiscrete(
-          autoencoder_hparams, tf.estimator.ModeKeys.EVAL)
-
-  def autoencode_tensor(self, x, batch_size=1):
-    if self.autoencoder_model is None:
-      return x
-    shape = [self.raw_frame_height, self.raw_frame_width, self.num_channels]
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      self.autoencoder_model.set_mode(tf.estimator.ModeKeys.EVAL)
-      autoencoded = self.autoencoder_model.encode(
-          tf.reshape(x, [batch_size, 1] + shape))
-    autoencoded = tf.reshape(
-        autoencoded, [batch_size] + self.frame_shape + [8])  # 8-bit groups.
-    if batch_size == 1:
-      autoencoded = tf.squeeze(autoencoded, axis=0)
-    return discretization.bit_to_int(autoencoded, 8)
+  def __init__(self, *args, **kwargs):
+    self.simulated_environment = True
+    self.make_extra_debug_info = True
+    self.debug_dump_frames_path = "debug_frames_sim"
+    self.intrinsic_reward_scale = 0.0
+    self.simulation_random_starts = True
+    super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
 
   def _setup(self):
+    super(GymSimulatedDiscreteProblem, self)._setup()
     if self.make_extra_debug_info:
       self.report_reward_statistics_every = 10
       self.dones = 0
       self.real_reward = 0
       # Slight weirdness to make sim env and real env aligned
-      if self.simulated_environment:
-        self.real_env.reset()
-        for _ in range(self.num_input_frames):
-          self.real_ob, _, _, _ = self.real_env.step(0)
+      self.real_env.reset()
+      for _ in range(self.num_input_frames):
+        self.real_ob, _, _, _ = self.real_env.step(0)
       self.total_sim_reward, self.total_real_reward = 0.0, 0.0
       self.sum_of_rewards = 0.0
       self.successful_episode_reward_predictions = 0
 
-    env_hparams = tf.contrib.training.HParams(
-        problem=self.real_env_problem if self.real_env_problem else self,
-        simulated_environment=self.simulated_environment)
-    if self.simulated_environment:
-      env_hparams.add_hparam("simulation_random_starts",
-                             self.simulation_random_starts)
-      env_hparams.add_hparam("intrinsic_reward_scale",
-                             self.intrinsic_reward_scale)
-
-    if FLAGS.autoencoder_path:
-      # TODO(lukaszkaiser): remove hard-coded autoencoder params.
-      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-        self.setup_autoencoder()
-        autoencoder_model = self.autoencoder_model
-        # Feeds for autoencoding.
-        shape = [self.raw_frame_height, self.raw_frame_width, self.num_channels]
-        self.autoencoder_feed = tf.placeholder(tf.int32, shape=shape)
-        self.autoencoder_result = self.autoencode_tensor(self.autoencoder_feed)
-        # Now for autodecoding.
-        shape = self.frame_shape
-        self.autodecoder_feed = tf.placeholder(tf.int32, shape=shape)
-        bottleneck = tf.reshape(
-            discretization.int_to_bit(self.autodecoder_feed, 8),
-            [1, 1, self.frame_height, self.frame_width, self.num_channels * 8])
-        autoencoder_model.set_mode(tf.estimator.ModeKeys.PREDICT)
-        self.autodecoder_result = autoencoder_model.decode(bottleneck)
-
-    def preprocess_fn(x):
-      shape = [self.raw_frame_height, self.raw_frame_width, self.num_channels]
-      # TODO(lukaszkaiser): we assume x comes from StackAndSkipWrapper skip=4.
-      xs = [tf.reshape(t, [1] + shape) for t in tf.split(x, 4, axis=-1)]
-      autoencoded = self.autoencode_tensor(tf.concat(xs, axis=0), batch_size=4)
-      encs = [tf.squeeze(t, axis=[0]) for t in tf.split(autoencoded, 4, axis=0)]
-      res = tf.to_float(tf.concat(encs, axis=-1))
-      return tf.expand_dims(res, axis=0)
-
-    # TODO(lukaszkaiser): x is from StackAndSkipWrapper thus 4*num_channels.
-    shape = [1, self.frame_height, self.frame_width, 4 * self.num_channels]
-    do_preprocess = (self.autoencoder_model is not None and
-                     not self.simulated_environment)
-    preprocess = (preprocess_fn, shape) if do_preprocess else None
-
-    if not FLAGS.agent_policy_path:
-      self.collect_hparams.policy_network = rl.random_policy_fun
-
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      self.collect_hparams.epoch_length = self._internal_memory_size
-      self.collect_hparams.num_agents = 1 #it is possible to set more
-      self.collect_memory, self.collect_trigger_op = collect.define_collect(
-          self.collect_hparams, scope="gym_problems",
-          collect_level=0,
-          eval_phase=self.eval_phase)
-
-  def restore_networks(self, sess):
-    if FLAGS.agent_policy_path:
-      model_saver = tf.train.Saver(
-          tf.global_variables(".*network_parameters.*"))
-      ckpts = tf.train.get_checkpoint_state(FLAGS.agent_policy_path)
-      ckpt = ckpts.model_checkpoint_path
-      model_saver.restore(sess, ckpt)
-    if FLAGS.autoencoder_path:
-      autoencoder_saver = tf.train.Saver(
-          tf.global_variables("autoencoder.*"))
-      ckpts = tf.train.get_checkpoint_state(FLAGS.autoencoder_path)
-      ckpt = ckpts.model_checkpoint_path
-      autoencoder_saver.restore(sess, ckpt)
-
-  def autoencode(self, image, sess):
-    return sess.run(self.autoencoder_result, {self.autoencoder_feed: image})
-
-  def autodecode(self, encoded, sess):
-    res = sess.run(self.autodecoder_result, {self.autodecoder_feed: encoded})
-    return res[0, 0, :self.raw_frame_height, :self.raw_frame_width, :]
-
-  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
-    self._setup()
-    self.debug_dump_frames_path = os.path.join(
-        data_dir, self.debug_dump_frames_path)
-
-    with tf.Session() as sess:
-      sess.run(tf.global_variables_initializer())
-      self.restore_networks(sess)
-      if FLAGS.only_use_ae_for_policy:
-        # If only the policy should use the autoencoder, then reset the flag so
-        # that other components here act as though there is no autoencoder and
-        # so write out full-resolution images. The policy graph was already
-        # built and self.collect_trigger_op is all that's used from it.
-        FLAGS.autoencoder_path = None
-      pieces_generated = 0
-      memory_index = 0
-      memory = None
-      while pieces_generated < self.num_steps:
-        if memory is None or memory_index>=self._internal_memory_size:
-          sess.run(self.collect_trigger_op)
-          memory = sess.run(self.collect_memory)
-          memory_index = 0
-        data = [memory[i][memory_index][0] for i in range(4)]
-        memory_index += 1
-        observ, reward, done, action = data
-        observ = observ.astype(np.uint8) # TODO(piotrmilos). This should be probably done in collect
-        debug_im = None
-        if self.make_extra_debug_info:
-          self.total_sim_reward += reward
-          if not self.simulated_environment:
-            self.real_ob = observ
-            self.real_reward = reward
-          if not FLAGS.autoencoder_path:
-            err = np.ndarray.astype(np.maximum(np.abs(
-                self.real_ob - observ, dtype=np.int) - 10, 0),
-                                    np.uint8)
-            debug_im = np.concatenate([observ, self.real_ob, err], axis=1)
-          if done:
-            self.dones += 1
-            self.sum_of_rewards += self.real_reward
-            if self.total_real_reward == self.total_sim_reward:
-              self.successful_episode_reward_predictions += 1
-
-            self.total_real_reward = 0.0
-            self.total_sim_reward = 0.0
-            self.real_reward = 0
-            if self.simulated_environment:
-              self.real_env.reset()
-              # Slight weirdness to make sim env and real env aligned
-              for _ in range(self.num_input_frames):
-                self.real_ob, _, _, _ = self.real_env.step(0)
-          else:
-            if self.simulated_environment:
-              self.real_ob, self.real_reward, _, _ = self.real_env.step(action)
-            self.total_real_reward += self.real_reward
-            self.sum_of_rewards += self.real_reward
-        if FLAGS.autoencoder_path:
-          if self.simulated_environment:
-            debug_im = self.autodecode(observ, sess)
-          else:
-            orig_observ = observ
-            observ = self.autoencode(observ, sess)
-            debug_im = np.concatenate([self.autodecode(observ, sess),
-                                       orig_observ], axis=1)
-        ret_dict = {"frame": observ,
-                    "image/format": ["png"],
-                    "image/height": [self.frame_height],
-                    "image/width": [self.frame_width],
-                    "action": [int(action)],
-                    "done": [int(False)],
-                    "reward": [int(reward) - self.min_reward]}
-        if self.make_extra_debug_info:
-          ret_dict["image/debug"] = debug_im
-        yield ret_dict
-        pieces_generated += 1
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgent(GymDiscreteProblemWithAgent):
-  """Simulated gym environment with discrete actions and rewards."""
+  def get_debug_image(self, data):
+    observ, reward, done, action = data
+    self.total_sim_reward += reward
+    err = np.ndarray.astype(np.maximum(np.abs(
+      self.real_ob - observ, dtype=np.int) - 10, 0),
+                            np.uint8)
+    debug_im = np.concatenate([observ, self.real_ob, err], axis=1)
+
+    if done:
+      self.dones += 1
+      self.sum_of_rewards += self.real_reward
+      if self.total_real_reward == self.total_sim_reward:
+        self.successful_episode_reward_predictions += 1
+
+      self.total_real_reward = 0.0
+      self.total_sim_reward = 0.0
+      self.real_reward = 0
+      # Slight weirdness to make sim env and real env aligned
+      for _ in range(self.num_input_frames):
+        self.real_ob, _, _, _ = self.real_env.step(0)
+    else:
+      self.real_ob, self.real_reward, _, _ = self.real_env.step(action)
+      self.total_real_reward += self.real_reward
+      self.sum_of_rewards += self.real_reward
 
-  def __init__(self, *args, **kwargs):
-    super(GymSimulatedDiscreteProblemWithAgent, self).__init__(*args, **kwargs)
-    self.simulated_environment = True
-    self.make_extra_debug_info = True
-    self.debug_dump_frames_path = "debug_frames_sim"
+    return debug_im
 
   @property
   def real_env(self):
@@ -609,8 +430,17 @@ def real_env(self):
           (TimeLimitWrapper, {"timelimit": timelimit}))
     return self._real_env
 
+  def get_environment_spec(self):
+    env_spec = standard_atari_env_spec(self.env_name)
+    env_spec.simulated_env = True
+    env_spec.add_hparam("simulation_random_starts",
+                           self.simulation_random_starts)
+    env_spec.add_hparam("intrinsic_reward_scale",
+                           self.intrinsic_reward_scale)
+    return env_spec
+
   def restore_networks(self, sess):
-    super(GymSimulatedDiscreteProblemWithAgent, self).restore_networks(sess)
+    super(GymSimulatedDiscreteProblem, self).restore_networks(sess)
     # TODO(blazej): adjust regexp for different models.
     env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
     sess = tf.get_default_session()
@@ -622,25 +452,25 @@ def restore_networks(self, sess):
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnPong(
-    GymSimulatedDiscreteProblemWithAgent, GymPongRandom5k):
+    GymSimulatedDiscreteProblem, GymPongRandom):
   pass
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnPong(
-    GymDiscreteProblemWithAgent, GymPongRandom5k):
+    GymRealDiscreteProblem, GymPongRandom):
   pass
 
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
-    GymSimulatedDiscreteProblemWithAgent, GymWrappedPongRandom5k):
+  GymSimulatedDiscreteProblem, GymWrappedPongRandom):
   pass
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedLongPong(
-    GymDiscreteProblemWithAgent, GymWrappedLongPongRandom):
+  GymRealDiscreteProblem, GymWrappedLongPongRandom):
   pass
 
 
@@ -652,13 +482,13 @@ class GymDiscreteProblemWithAgentOnWrappedLongPongAe(  # with autoencoder
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPong(
-    GymSimulatedDiscreteProblemWithAgent, GymWrappedLongPongRandom):
+    GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
   pass
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedBreakout(
-    GymDiscreteProblemWithAgent, GymWrappedBreakoutRandom5k):
+  GymRealDiscreteProblem, GymWrappedBreakoutRandom):
   pass
 
 
@@ -670,13 +500,13 @@ class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
-    GymSimulatedDiscreteProblemWithAgent, GymWrappedBreakoutRandom5k):
+    GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
   pass
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedPong(
-    GymDiscreteProblemWithAgent, GymWrappedPongRandom5k):
+  GymRealDiscreteProblem, GymWrappedPongRandom):
   """GymDiscreteProblemWithAgentOnWrappedPong."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
@@ -706,13 +536,13 @@ class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnFreeway(
-    GymSimulatedDiscreteProblemWithAgent, GymFreewayRandom5k):
+    GymSimulatedDiscreteProblem, GymFreewayRandom):
   pass
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnFreeway(
-    GymDiscreteProblemWithAgent, GymFreewayRandom5k):
+  GymRealDiscreteProblem, GymFreewayRandom):
   """Freeway with agent."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 65f845420..2a643e5af 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -119,7 +119,6 @@ def ppo_pong_base():
   hparams.max_gradients_norm = 0.5
   return hparams
 
-EnvironmentSpec = collections.namedtuple('EnvironmentSpec', 'env_lambda, wrappers, simulated_env')
 
 def standard_atari_env_spec(env):
   """Parameters of environement specification"""
@@ -131,7 +130,10 @@ def standard_atari_env_spec(env):
     env_lambda = env
   assert env is not None, "Unknown specification of environment"
 
-  return EnvironmentSpec(env_lambda=env_lambda, wrappers=standard_wrappers, simulated_env=False)
+  return tf.contrib.training.HParams(env_lambda=env_lambda,
+                                     wrappers=standard_wrappers,
+                                     simulated_env=False)
+
 
 def simple_gym_spec(env):
   """Parameters of environement specification"""
@@ -143,7 +145,9 @@ def simple_gym_spec(env):
     env_lambda = env
   assert env is not None, "Unknown specification of environment"
 
-  return EnvironmentSpec(env_lambda=env_lambda, wrappers=standard_wrappers, simulated_env=False)
+  return tf.contrib.training.HParams(env_lambda=env_lambda,
+                                     wrappers=standard_wrappers,
+                                     simulated_env=False)
 
 @registry.register_hparams
 def ppo_pong_ae_base():
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 03a769c24..9aac4c430 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -114,7 +114,7 @@ def __init__(self, hparams, length, problem,
       dataset = dataset.shuffle(buffer_size=100)
     else:
       dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
-                                shuffle_files=False, hparams=hparams).take(1)
+                                shuffle_files=True, hparams=hparams).take(1)
 
     dataset = dataset.map(lambda x: x["inputs"]).repeat()
     self.history_buffer = HistoryBuffer(dataset, self.length)
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 4f7e328bb..7d37f2d54 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -36,9 +36,7 @@
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.layers import discretization
-from tensor2tensor.models.research.rl import EnvironmentSpec
 from tensor2tensor.rl import rl_trainer_lib
-from tensor2tensor.rl.envs.tf_atari_wrappers import StackAndSkipWrapper
 from tensor2tensor.rl.envs.tf_atari_wrappers import TimeLimitWrapper
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
@@ -143,8 +141,6 @@ def train_agent(problem_name, agent_model_dir,
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_epochs_num = hparams.ppo_epochs_num
   ppo_hparams.epochs_num = ppo_epochs_num
-  ppo_hparams.simulation_random_starts = hparams.simulation_random_starts
-  ppo_hparams.intrinsic_reward_scale = hparams.intrinsic_reward_scale
   ppo_hparams.eval_every_epochs = 50
   ppo_hparams.save_models_every_epochs = ppo_epochs_num
   ppo_hparams.epoch_length = hparams.ppo_epoch_length
@@ -154,16 +150,19 @@ def train_agent(problem_name, agent_model_dir,
   if hparams.ppo_learning_rate:
     ppo_hparams.learning_rate = hparams.ppo_learning_rate
 
-  environment_spec = gym_problem.environment_spec
+  environment_spec = copy.copy(gym_problem.environment_spec)
+  environment_spec.simulated_env = True
+  environment_spec.add_hparam("simulation_random_starts",
+                              hparams.simulation_random_starts)
+  environment_spec.add_hparam("intrinsic_reward_scale",
+                              hparams.intrinsic_reward_scale)
+
   # 4x for the StackAndSkipWrapper minus one to always finish for reporting.
   ppo_time_limit = (ppo_hparams.epoch_length - 1) * 4
-  wrappers = copy.copy(environment_spec.wrappers)
-  wrappers.append([TimeLimitWrapper, {"timelimit": ppo_time_limit}])
-  ppo_hparams.add_hparam("environment_spec",
-                         EnvironmentSpec(env_lambda=environment_spec.env_lambda,
-                                         wrappers=wrappers,
-                                         simulated_env=True))
+  wrappers = environment_spec.wrappers + [[TimeLimitWrapper, {"timelimit": ppo_time_limit}]]
+  environment_spec.wrappers = wrappers
 
+  ppo_hparams.add_hparam("environment_spec", environment_spec)
 
   with temporary_flags({
       "problem": problem_name,

From afee9ad22fd46eed0f4c567738729b93d97f769d Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 18 Jun 2018 18:55:14 +0200
Subject: [PATCH 0128/2720] gym_problems refactor done

---
 tensor2tensor/data_generators/gym_problems.py | 75 ++++++++++++-------
 tensor2tensor/rl/envs/simulated_batch_env.py  | 20 ++---
 tensor2tensor/rl/envs/utils.py                |  6 +-
 tensor2tensor/rl/model_rl_experiment.py       | 17 +++--
 4 files changed, 70 insertions(+), 48 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 6304b48a8..50aa273de 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -238,6 +238,9 @@ def hparams(self, defaults, unused_model_hparams):
     p.target_space_id = problem.SpaceID.IMAGE
 
 
+class GymAEDiscreteProblem(GymDiscreteProblem):
+  pass
+
 @registry.register_problem
 class GymPongRandom(GymDiscreteProblem):
   """Pong game, random actions."""
@@ -355,10 +358,7 @@ def get_environment_spec(self):
     return standard_atari_env_spec(self.env_name)
 
 
-class GymAEDiscreteProblem(GymDiscreteProblem):
-  pass
 
-@registry.register_problem
 class GymSimulatedDiscreteProblem(GymDiscreteProblem):
   """Simulated gym environment with discrete actions and rewards."""
 
@@ -367,15 +367,12 @@ def __init__(self, *args, **kwargs):
     self.make_extra_debug_info = True
     self.debug_dump_frames_path = "debug_frames_sim"
     self.intrinsic_reward_scale = 0.0
-    self.simulation_random_starts = True
+    self.simulation_random_starts = False
     super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
 
   def _setup(self):
     super(GymSimulatedDiscreteProblem, self)._setup()
     if self.make_extra_debug_info:
-      self.report_reward_statistics_every = 10
-      self.dones = 0
-      self.real_reward = 0
       # Slight weirdness to make sim env and real env aligned
       self.real_env.reset()
       for _ in range(self.num_input_frames):
@@ -412,31 +409,34 @@ def get_debug_image(self, data):
     return debug_im
 
   @property
-  def real_env(self):
-    """Lazy caching environment construction."""
-    if self._real_env is None:
-      self._real_env = self.environment_spec()
-      if self.num_testing_steps is not None:
-        timelimit = self.num_testing_steps
-      else:
-        try:
-          # We assume that the real env is wrapped with TimeLimit.
-          history = self.num_input_frames
-          timelimit = self.real_env._max_episode_steps - history  # pylint: disable=protected-access
-        except:  # pylint: disable=bare-except
-          # If not, set some reasonable default.
-          timelimit = 100
-      self.in_graph_wrappers.append(
-          (TimeLimitWrapper, {"timelimit": timelimit}))
-    return self._real_env
+  def initial_frames_problem(self):
+    raise NotImplemented()
 
   def get_environment_spec(self):
     env_spec = standard_atari_env_spec(self.env_name)
+
+    #Set reasonable time limit (as we do not simulate done)
+    self.real_env = env_spec.env_lambda()
+    if self.num_testing_steps is not None:
+      timelimit = self.num_testing_steps
+    else:
+      try:
+        # We assume that the real env is wrapped with TimeLimit.
+        history = self.num_input_frames
+        timelimit = self.real_env._max_episode_steps - history  # pylint: disable=protected-access
+      except:  # pylint: disable=bare-except
+        # If not, set some reasonable default.
+        timelimit = 100
+
     env_spec.simulated_env = True
     env_spec.add_hparam("simulation_random_starts",
                            self.simulation_random_starts)
     env_spec.add_hparam("intrinsic_reward_scale",
                            self.intrinsic_reward_scale)
+    initial_frames_problem = registry.problem(self.initial_frames_problem)
+    env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
+    env_spec.wrappers.append([TimeLimitWrapper, {"timelimit": timelimit}])
+
     return env_spec
 
   def restore_networks(self, sess):
@@ -453,7 +453,10 @@ def restore_networks(self, sess):
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnPong(
     GymSimulatedDiscreteProblem, GymPongRandom):
-  pass
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_pong"
 
 
 @registry.register_problem
@@ -465,7 +468,10 @@ class GymDiscreteProblemWithAgentOnPong(
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
   GymSimulatedDiscreteProblem, GymWrappedPongRandom):
-  pass
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_pong"
 
 
 @registry.register_problem
@@ -483,7 +489,10 @@ class GymDiscreteProblemWithAgentOnWrappedLongPongAe(  # with autoencoder
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPong(
     GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
-  pass
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_long_pong"
 
 
 @registry.register_problem
@@ -501,7 +510,11 @@ class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
     GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
-  pass
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_breakout"
+
 
 
 @registry.register_problem
@@ -537,7 +550,11 @@ class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnFreeway(
     GymSimulatedDiscreteProblem, GymFreewayRandom):
-  pass
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_freeway"
+
 
 
 @registry.register_problem
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 9aac4c430..1e6425341 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -91,12 +91,14 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   flags are held in according variables.
   """
 
-  def __init__(self, hparams, length, problem,
-               simulation_random_starts=False, intrinsic_reward_scale=0.):
+  def __init__(self, hparams, length, simulation_random_starts=False,
+               intrinsic_reward_scale=0.):
     """Batch of environments inside the TensorFlow graph."""
     self.length = length
-    self._min_reward = problem.min_reward
-    self._num_frames = problem.num_input_frames
+    environment_spec = hparams.environment_spec
+    initial_frames_problem = environment_spec.initial_frames_problem
+    self._min_reward = initial_frames_problem.min_reward
+    self._num_frames = initial_frames_problem.num_input_frames
     self._intrinsic_reward_scale = intrinsic_reward_scale
 
     # initialization_env = environment_lambda()
@@ -106,21 +108,21 @@ def __init__(self, hparams, length, problem,
     self._model = registry.model(FLAGS.model)(
       model_hparams, tf.estimator.ModeKeys.PREDICT)
 
-    _, self.action_shape, self.action_dtype = utils.get_action_space(hparams.environment_spec)
+    _, self.action_shape, self.action_dtype = utils.get_action_space(environment_spec)
 
     if simulation_random_starts:
-      dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
+      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
                                 shuffle_files=True, hparams=hparams)
       dataset = dataset.shuffle(buffer_size=100)
     else:
-      dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
+      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
                                 shuffle_files=True, hparams=hparams).take(1)
 
     dataset = dataset.map(lambda x: x["inputs"]).repeat()
     self.history_buffer = HistoryBuffer(dataset, self.length)
 
-    shape = (self.length, problem.frame_height, problem.frame_width,
-             problem.num_channels)
+    shape = (self.length, initial_frames_problem.frame_height,
+             initial_frames_problem.frame_width, initial_frames_problem.num_channels)
     self._observ = tf.Variable(tf.zeros(shape, tf.float32), trainable=False)
 
   def __len__(self):
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index c142333c4..58f8b1ee9 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -286,7 +286,7 @@ def batch_env_factory(hparams, xvfb=False):
   if environment_spec.simulated_env:
     # TODO(piotrmilos): Consider passing only relevant paramters
     cur_batch_env = _define_simulated_batch_env(
-      hparams, hparams.num_agents, hparams.problem,
+        hparams, hparams.num_agents,
         hparams.simulation_random_starts,
         hparams.intrinsic_reward_scale)
   else:
@@ -308,11 +308,11 @@ def _define_batch_env(environment_spec, num_agents, xvfb=False):
     return env
 
 
-def _define_simulated_batch_env(hparams, num_agents, problem,
+def _define_simulated_batch_env(hparams, num_agents,
                                simulation_random_starts=False,
                                intrinsic_reward_scale=0.):
   cur_batch_env = simulated_batch_env.SimulatedBatchEnv(
-      hparams, num_agents, problem, simulation_random_starts,
+      hparams, num_agents, simulation_random_starts,
       intrinsic_reward_scale)
   return cur_batch_env
 
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 7d37f2d54..76d1383b6 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -145,7 +145,7 @@ def train_agent(problem_name, agent_model_dir,
   ppo_hparams.save_models_every_epochs = ppo_epochs_num
   ppo_hparams.epoch_length = hparams.ppo_epoch_length
   ppo_hparams.num_agents = hparams.ppo_num_agents
-  ppo_hparams.problem = gym_problem
+  # ppo_hparams.problem = gym_problem
   ppo_hparams.world_model_dir = world_model_dir
   if hparams.ppo_learning_rate:
     ppo_hparams.learning_rate = hparams.ppo_learning_rate
@@ -156,6 +156,9 @@ def train_agent(problem_name, agent_model_dir,
                               hparams.simulation_random_starts)
   environment_spec.add_hparam("intrinsic_reward_scale",
                               hparams.intrinsic_reward_scale)
+  environment_spec.add_hparam("initial_frames_problem",
+                              gym_problem)
+
 
   # 4x for the StackAndSkipWrapper minus one to always finish for reporting.
   ppo_time_limit = (ppo_hparams.epoch_length - 1) * 4
@@ -180,12 +183,12 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
                          autoencoder_path=None):
   """Generate simulated environment data and return reward accuracy."""
   gym_simulated_problem = registry.problem(simulated_problem_name)
-  gym_problem = registry.problem(problem_name)
+  # gym_problem = registry.problem(problem_name)
   sim_steps = hparams.simulated_env_generator_num_steps
   gym_simulated_problem.settable_num_steps = sim_steps
-  gym_simulated_problem.real_env_problem = gym_problem
-  gym_simulated_problem.simulation_random_starts = False
-  gym_simulated_problem.intrinsic_reward_scale = 0.
+  # gym_simulated_problem.real_env_problem = gym_problem
+  # gym_simulated_problem.simulation_random_starts = False
+  # gym_simulated_problem.intrinsic_reward_scale = 0.
   with temporary_flags({
       "problem": problem_name,
       "model": hparams.generative_model,
@@ -434,8 +437,6 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   return epoch_metrics[-1]
 
 
-
-
 def combine_training_data(problem, final_data_dir, old_data_dirs,
                           copy_last_eval_set=True):
   """Add training data from old_data_dirs into final_data_dir."""
@@ -542,6 +543,8 @@ def rl_modelrl_tiny():
           ppo_time_limit=5,
           ppo_epoch_length=5,
           ppo_num_agents=2,
+          # eval_world_model=True,
+          # simulation_random_starts=True
       ).values())
 
 

From 536aa1975d83d85746da069d7788329b621ebec2 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 18 Jun 2018 19:32:48 +0200
Subject: [PATCH 0129/2720] small bug fix

---
 tensor2tensor/data_generators/gym_problems_test.py | 2 +-
 tensor2tensor/rl/collect.py                        | 5 +----
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index 1bbd2a374..ccd87b6c8 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -35,7 +35,7 @@ def setUpClass(cls):
     os.mkdir(cls.tmp_dir)
 
   def testGymAtariBoots(self):
-    problem = gym_problems.GymPongRandom5k()
+    problem = gym_problems.GymPongRandom()
     self.assertEqual(5000, problem.num_steps)
 
 
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 757241bfe..5553099cf 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -18,15 +18,12 @@
 from __future__ import print_function
 
 
-
 import copy
 import tensorflow as tf
-
 from tensor2tensor.rl.envs import utils
-
 from tensor2tensor.rl.envs.tf_atari_wrappers import WrapperBase
 from tensor2tensor.rl.envs.utils import get_policy
-from tensor2tensor.utils import registry
+
 
 def _rollout_metadata(batch_env):
   batch_env_shape = batch_env.observ.get_shape().as_list()

From ee6a437c036f3056c8866278eb5a2c50491638ee Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 18 Jun 2018 11:12:37 -0700
Subject: [PATCH 0130/2720] Move problems import to problems.py (from
 all_problems.py)

PiperOrigin-RevId: 201026560
---
 tensor2tensor/data_generators/all_problems.py  | 18 ++++++++++--------
 tensor2tensor/problems.py                      |  5 ++++-
 .../all_problems_test.py => problems_test.py}  | 13 +++++--------
 3 files changed, 19 insertions(+), 17 deletions(-)
 rename tensor2tensor/{data_generators/all_problems_test.py => problems_test.py} (73%)

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 6ab99e7f5..dcb83832f 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -20,7 +20,7 @@
 import importlib
 import six
 
-modules = [
+MODULES = [
     "tensor2tensor.data_generators.algorithmic",
     "tensor2tensor.data_generators.algorithmic_math",
     "tensor2tensor.data_generators.audio",
@@ -67,6 +67,7 @@
     "tensor2tensor.data_generators.wikitext103",
     "tensor2tensor.data_generators.wsj_parsing",
 ]
+ALL_MODULES = list(MODULES)
 
 
 
@@ -94,10 +95,11 @@ def _handle_errors(errors):
       print("Did not import module: %s; Cause: %s" % (module, err_str))
 
 
-_errors = []
-for _module in modules:
-  try:
-    importlib.import_module(_module)
-  except ImportError as error:
-    _errors.append((_module, error))
-_handle_errors(_errors)
+def import_modules(modules):
+  errors = []
+  for module in modules:
+    try:
+      importlib.import_module(module)
+    except ImportError as error:
+      errors.append((module, error))
+  _handle_errors(errors)
diff --git a/tensor2tensor/problems.py b/tensor2tensor/problems.py
index c986c8eb5..3f349b908 100644
--- a/tensor2tensor/problems.py
+++ b/tensor2tensor/problems.py
@@ -17,7 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.data_generators import all_problems  # pylint: disable=unused-import
+from tensor2tensor.data_generators import all_problems
 from tensor2tensor.utils import registry
 
 
@@ -27,3 +27,6 @@ def problem(name):
 
 def available():
   return sorted(registry.list_problems())
+
+
+all_problems.import_modules(all_problems.ALL_MODULES)
diff --git a/tensor2tensor/data_generators/all_problems_test.py b/tensor2tensor/problems_test.py
similarity index 73%
rename from tensor2tensor/data_generators/all_problems_test.py
rename to tensor2tensor/problems_test.py
index 9d9fb3316..35db33b70 100644
--- a/tensor2tensor/data_generators/all_problems_test.py
+++ b/tensor2tensor/problems_test.py
@@ -12,23 +12,20 @@
 # 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 Tensor2Tensor's all_problems.py."""
-
+"""tensor2tensor.problems test."""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.data_generators import all_problems
+from tensor2tensor import problems
 
 import tensorflow as tf
 
 
-class AllProblemsTest(tf.test.TestCase):
+class ProblemsTest(tf.test.TestCase):
 
   def testImport(self):
-    """Make sure that importing all_problems doesn't break."""
-    self.assertIsNotNone(all_problems)
-
+    self.assertIsNotNone(problems)
 
-if __name__ == '__main__':
+if __name__ == "__main__":
   tf.test.main()

From 0c70614e8453d069b5231c11d3d776a15c5b32de Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 18 Jun 2018 20:16:29 +0200
Subject: [PATCH 0131/2720] test fixed

---
 tensor2tensor/data_generators/gym_problems_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index ccd87b6c8..9638acebf 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -35,7 +35,7 @@ def setUpClass(cls):
     os.mkdir(cls.tmp_dir)
 
   def testGymAtariBoots(self):
-    problem = gym_problems.GymPongRandom()
+    problem = gym_problems.GymSimulatedDiscreteProblemWithAgentOnPong()
     self.assertEqual(5000, problem.num_steps)
 
 

From b131f60663e2e23b500ba84dd827430c270273d9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 18 Jun 2018 11:52:18 -0700
Subject: [PATCH 0132/2720] Add a log indicating data is shuffled.

PiperOrigin-RevId: 201033717
---
 tensor2tensor/data_generators/generator_utils.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 9bd6cd199..64e09592f 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -481,6 +481,7 @@ def _shuffle_single(fname):
 
 
 def shuffle_dataset(filenames):
+  """Shuffles the dataset."""
   if outputs_exist(filenames):
     tf.logging.info("Skipping shuffle because output files exist")
     return
@@ -490,6 +491,7 @@ def shuffle_dataset(filenames):
     pool.map(_shuffle_single, filenames)
   else:
     _shuffle_single(filenames[0])
+  tf.logging.info("Data shuffled.")
 
 
 class SequencePacker(object):

From 2404b04ab37d380400d86989d650abaeca91ac86 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 18 Jun 2018 12:06:04 -0700
Subject: [PATCH 0133/2720] Set export_outputs for TPUEstimators

PiperOrigin-RevId: 201035937
---
 tensor2tensor/utils/t2t_model.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index a99997d3c..c29509c6e 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1355,18 +1355,20 @@ def estimator_spec_predict(self, features, use_tpu=False):
 
     _remove_summaries()
 
+    export_outputs = {
+        tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
+            tf.estimator.export.PredictOutput(export_out)
+    }
     if use_tpu:
       return tf.contrib.tpu.TPUEstimatorSpec(
-          tf.estimator.ModeKeys.PREDICT, predictions=predictions)
+          tf.estimator.ModeKeys.PREDICT,
+          predictions=predictions,
+          export_outputs=export_outputs)
     else:
       return tf.estimator.EstimatorSpec(
           tf.estimator.ModeKeys.PREDICT,
           predictions=predictions,
-          export_outputs={
-              tf.saved_model.signature_constants.
-              DEFAULT_SERVING_SIGNATURE_DEF_KEY:
-                  tf.estimator.export.PredictOutput(export_out)
-          })
+          export_outputs=export_outputs)
 
   def _normalize_body_output(self, body_out):
     if isinstance(body_out, tuple):

From d918b7e2964f25aaa9b7297f0bdc365e016303fd Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 18 Jun 2018 12:42:42 -0700
Subject: [PATCH 0134/2720] Remove recompute_grad workaround from
 common_layers.dense and batch_dense.

PiperOrigin-RevId: 201041076
---
 tensor2tensor/layers/common_layers.py | 37 ++++++---------------------
 1 file changed, 8 insertions(+), 29 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 6afc4039c..d1b76ef67 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3014,24 +3014,16 @@ def fn_with_recompute(*args):
 
 
 def dense(x, units, **kwargs):
-  """Identical to tf.layers.dense, Memory optimization on tpu."""
-  fn = lambda x: tf.layers.dense(x, units, **kwargs)
-  if is_on_tpu():
-    # TODO(noam): remove this hack once XLA does the right thing.
-    # Forces the gradients on the inputs to be computed before the variables
-    # are updated.  This saves memory by preventing XLA from making an extra
-    # copy of the variables.
-    return _recompute_grad(fn, [x])
-  else:
-    return fn(x)
+  """Identical to tf.layers.dense."""
+  return tf.layers.dense(x, units, **kwargs)
 
 
-def _batch_dense(inputs,
-                 units,
-                 activation=None,
-                 kernel_initializer=None,
-                 reuse=None,
-                 name=None):
+def batch_dense(inputs,
+                units,
+                activation=None,
+                kernel_initializer=None,
+                reuse=None,
+                name=None):
   """Multiply a batch of input matrices by a batch of parameter matrices.
 
   Each input matrix is multiplied by the corresponding parameter matrix.
@@ -3076,19 +3068,6 @@ def _batch_dense(inputs,
     return y
 
 
-def batch_dense(x, units, **kwargs):
-  """Identical to _batch_dense, Memory optimization on tpu."""
-  fn = lambda x: _batch_dense(x, units, **kwargs)
-  if is_on_tpu():
-    # TODO(noam): remove this hack once XLA does the right thing.
-    # Forces the gradients on the inputs to be computed before the variables
-    # are updated.  This saves memory by preventing XLA from making an extra
-    # copy of the variables.
-    return _recompute_grad(fn, [x])
-  else:
-    return fn(x)
-
-
 def mix(x1, x2, steps, is_training,
         min_prob=0.0, max_prob=1.0,
         mode="lin", simple=False, broadcast_last=False):

From 2a21962b31b30cee6928340da43d2191ecc94dce Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Mon, 18 Jun 2018 13:56:40 -0700
Subject: [PATCH 0135/2720] Add edge conditioned linear transformation for
 self-attention, with a faster algorithm. Add a GRU feedforward layer.

PiperOrigin-RevId: 201053232
---
 tensor2tensor/layers/common_attention.py      | 133 ++++++++--
 tensor2tensor/layers/common_layers.py         |  33 +++
 .../common_message_passing_attention.py       | 230 ++++++++++++++++++
 tensor2tensor/models/transformer.py           |  12 +-
 tensor2tensor/utils/metrics.py                |  11 +
 5 files changed, 401 insertions(+), 18 deletions(-)
 create mode 100644 tensor2tensor/layers/common_message_passing_attention.py

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 3a1990b13..c464ec2e5 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -682,6 +682,41 @@ def add_positional_embedding_nd(x, max_length, name):
   return x
 
 
+@expert_utils.add_name_scope()
+def make_edge_vectors(adjacency_matrix, num_edge_types, depth, name=None):
+  """Gets edge vectors for the edge types in the adjacency matrix.
+
+  Args:
+    adjacency_matrix: A [batch, num_nodes, num_nodes] tensor of ints.
+    num_edge_types: Number of different edge types
+    depth: Number of channels
+    name: a string
+  Returns:
+    A [batch, num_nodes, num_nodes, depth] vector of tensors
+  """
+  with tf.variable_scope(name, default_name="edge_vectors"):
+    att_adj_vectors_shape = [num_edge_types, depth]
+    adjacency_matrix_shape = common_layers.shape_list(adjacency_matrix)
+    adj_vectors = (
+        tf.get_variable(
+            "adj_vectors",
+            att_adj_vectors_shape,
+            initializer=tf.random_normal_initializer(0, depth**-0.5)) *
+        (depth**0.5))
+    # Avoiding gathers so that it works on TPUs
+    # adjacency_matrix_one_hot has shape
+    # [batch, num_nodes, num_nodes, num_edge_types]
+
+    adjacency_matrix_one_hot = tf.one_hot(adjacency_matrix, num_edge_types)
+
+    att_adj_vectors = tf.matmul(
+        tf.reshape(tf.to_float(adjacency_matrix_one_hot), [-1, num_edge_types]),
+        adj_vectors)
+    return tf.reshape(att_adj_vectors,
+                      [adjacency_matrix_shape[0], adjacency_matrix_shape[1],
+                       adjacency_matrix_shape[2], depth])
+
+
 class LshGating(object):
   """Class to split key/queries into separate buckets."""
 
@@ -1380,6 +1415,71 @@ def grouped_attention_multihead(query_antecedent,
     return o, extra_loss
 
 
+def graph_attention(q,
+                    k,
+                    v,
+                    bias,
+                    dropout_rate=0.0,
+                    image_shapes=None,
+                    name=None,
+                    make_image_summary=True,
+                    save_weights_to=None,
+                    dropout_broadcast_dims=None,
+                    adjacency_matrix=None,
+                    num_edge_types=5):
+  """graph attention.
+
+  Args:
+    q: a Tensor with shape [batch, heads, length_q, depth_k]
+    k: a Tensor with shape [batch, heads, length_kv, depth_k]
+    v: a Tensor with shape [batch, heads, length_kv, depth_v]
+    bias: bias Tensor (see attention_bias())
+    dropout_rate: a floating point number
+    image_shapes: optional tuple of integer scalars.
+      see comments for attention_image_summary()
+    name: an optional string
+    make_image_summary: True if you want an image summary.
+    save_weights_to: an optional dictionary to capture attention weights
+      for vizualization; the weights tensor will be appended there under
+      a string key created from the variable scope (including name).
+    dropout_broadcast_dims:  an optional list of integers less than 4
+      specifying in which dimensions to broadcast the dropout decisions.
+      saves memory.
+    adjacency_matrix: optional matrix of [batch, length, length] ids indicating
+      edge type
+    num_edge_types: an int indicating number of edge types
+  Returns:
+    A Tensor of shape [batch, length, depth(q)]
+  """
+  with tf.variable_scope(
+      name, default_name="dot_product_attention", values=[q, k, v]) as scope:
+    # [batch, num_heads, query_length, memory_length]
+    logits = tf.matmul(q, k, transpose_b=True)
+    if adjacency_matrix is not None:
+      key_head_depth = common_layers.shape_list(q)[-1]
+      adjacency_vectors = make_edge_vectors(
+          adjacency_matrix, num_edge_types, key_head_depth, name)
+      # zeroing out the vectors that have 0 entries in the adjacency
+      adjacency_vectors *= tf.to_float(
+          tf.expand_dims(adjacency_matrix, axis=-1))
+      # transposing q to be [batch, length_q, heads, depth_k]
+      # to allow for matmul with [batch, length_q, length_q, depth_k]
+      q_t = tf.transpose(q, [0, 2, 1, 3])
+      adj_logits = tf.matmul(q_t, adjacency_vectors, transpose_b=True)
+      logits += tf.transpose(adj_logits, [0, 2, 1, 3])
+    if bias is not None:
+      logits += bias
+    weights = tf.nn.softmax(logits, name="attention_weights")
+    if save_weights_to is not None:
+      save_weights_to[scope.name] = weights
+    # dropping out the attention links for each of the heads
+    weights = common_layers.dropout_with_broadcast_dims(
+        weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
+    if expert_utils.should_generate_summaries() and make_image_summary:
+      attention_image_summary(weights, image_shapes)
+    return tf.matmul(weights, v)
+
+
 def dot_product_attention(q,
                           k,
                           v,
@@ -1408,7 +1508,6 @@ def dot_product_attention(q,
     dropout_broadcast_dims:  an optional list of integers less than 4
       specifying in which dimensions to broadcast the dropout decisions.
       saves memory.
-
   Returns:
     A Tensor.
   """
@@ -2696,6 +2795,8 @@ def multihead_attention(query_antecedent,
                         make_image_summary=True,
                         dropout_broadcast_dims=None,
                         max_length=None,
+                        adjacency_matrix=None,
+                        num_edge_types=5,
                         **kwargs):
   """Multihead scaled-dot-product attention with input/output transformations.
 
@@ -2715,8 +2816,8 @@ def multihead_attention(query_antecedent,
                   see comments for attention_image_summary()
     attention_type: a string, either "dot_product", "dot_product_relative",
                     "local_mask_right", "local_unmasked", "masked_dilated_1d",
-                    "unmasked_dilated_1d" or any attention function with the
-                    signature (query, key, value, **kwargs)
+                    "unmasked_dilated_1d", graph, or any attention function
+                    with the signature (query, key, value, **kwargs)
     block_length: an integer - relevant for "local_mask_right"
     block_width: an integer - relevant for "local_unmasked"
     q_filter_width: An integer specifying how wide you want the query to be.
@@ -2744,6 +2845,9 @@ def multihead_attention(query_antecedent,
       specifying in which dimensions to broadcast the dropout decisions.
       saves memory.
     max_length: an integer - needed by relative attention
+    adjacency_matrix: an optional tensor of shape [batch, len_q, len_q]
+      containing edge vectors for attention
+    num_edge_types: number of edge types, an int
     **kwargs (dict): Parameters for the attention function
 
   Caching:
@@ -2817,16 +2921,17 @@ def multihead_attention(query_antecedent,
       if isinstance(x, tuple):
         x, additional_returned_value = x  # Unpack
     elif attention_type == "dot_product":
-      x = dot_product_attention(
-          q,
-          k,
-          v,
-          bias,
-          dropout_rate,
-          image_shapes,
-          save_weights_to=save_weights_to,
-          make_image_summary=make_image_summary,
-          dropout_broadcast_dims=dropout_broadcast_dims)
+      x = dot_product_attention(q, k, v, bias, dropout_rate, image_shapes,
+                                save_weights_to=save_weights_to,
+                                make_image_summary=make_image_summary,
+                                dropout_broadcast_dims=dropout_broadcast_dims)
+    elif attention_type == "edge_vector":
+      x = graph_attention(q, k, v, bias, dropout_rate, image_shapes,
+                          save_weights_to=save_weights_to,
+                          make_image_summary=make_image_summary,
+                          dropout_broadcast_dims=dropout_broadcast_dims,
+                          adjacency_matrix=adjacency_matrix,
+                          num_edge_types=num_edge_types)
     elif attention_type == "dot_product_relative":
       x = dot_product_attention_relative(
           q,
@@ -2905,7 +3010,7 @@ def multihead_attention_2d(query_antecedent,
     name: an optional string
 
   Returns:
-    A Tensor of shape [batch, h, w, depth_k]
+    A Tensor of shape [batch, h, w, output_depth]
 
   Raises:
     ValueError: if the key depth or value depth are not divisible by the
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index d1b76ef67..fd4c50c92 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1629,6 +1629,39 @@ def do_conv(args, name, bias_start, padding):
     return gate * x + (1 - gate) * candidate
 
 
+def gru_feedfwd(a_t, h_prev, filters, name=None):
+  """position-wise Feed-fwd GRU gates following the MPNN.
+
+  Args:
+    a_t: Tensor of shape [batch, length, depth] of current input
+    h_prev: Tensor of shape [batch, length, depth] of prev input
+    filters: an integer specifying number of dimensions of the filters
+    name: A string
+  Returns:
+    h_t: [batch, length, filters] hidden state
+  """
+
+  with tf.variable_scope(
+      name, default_name="GRU", values=[a_t, h_prev]):
+    # we use right matrix multiplication to handle batches
+    # W_z and W_r have shape 2d, d. U_z U_r have shape d,d
+    z_t = (tf.sigmoid(tpu_conv1d(a_t, filters, 1, padding="SAME",
+                                 name="W_z") +
+                      tpu_conv1d(h_prev, filters, 1, padding="SAME",
+                                 name="U_z")))
+    r_t = (tf.sigmoid(tpu_conv1d(a_t, filters, 1, padding="SAME",
+                                 name="W_r") +
+                      tpu_conv1d(h_prev, filters, 1, padding="SAME",
+                                 name="U_r")))
+    h_tilde = (tf.tanh(tpu_conv1d(a_t, filters, 1, padding="SAME",
+                                  name="W") +
+                       tpu_conv1d(r_t*h_prev, filters, 1, padding="SAME",
+                                  name="U")))
+    h_t = (1. - z_t)*h_prev + z_t * h_tilde
+
+  return h_t
+
+
 def conv_lstm(x,
               kernel_size,
               filters,
diff --git a/tensor2tensor/layers/common_message_passing_attention.py b/tensor2tensor/layers/common_message_passing_attention.py
new file mode 100644
index 000000000..3946662c0
--- /dev/null
+++ b/tensor2tensor/layers/common_message_passing_attention.py
@@ -0,0 +1,230 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Utilities for attention."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_layers
+
+import tensorflow as tf
+
+
+def compute_mpnn_qkv(node_states,
+                     total_key_depth,
+                     total_value_depth,
+                     num_edge_types,
+                     ignore_zero=True):
+  """Computes query, key and value for edge matrices.
+
+  Args:
+    node_states: a Tensor with shape [batch, num_nodes, channels]
+    total_key_depth: an integer
+    total_value_depth: an integer
+    num_edge_types: a integer specifying number of edge types
+    ignore_zero: If true, then edge type 0 will not be considered. Equivalent
+      to have a linear transformation of all 0's for edge type 0
+  Returns:
+    q: [batch, num_nodes, channels]
+    k: [batch, num_nodes * num_edge_types, channels]
+    v: [batch, num_nodes * num_edge_types, channels]
+  """
+  memory_antecedent = node_states
+  def _compute(inp, depth, filter_width, padding, name):
+    if filter_width == 1:
+      return common_layers.dense(inp, depth, use_bias=False, name=name)
+    else:
+      return common_layers.conv1d(inp, depth, filter_width, padding, name=name)
+  # For edge type 0, if ignore_zero, don't multiply with linear transformation,
+  # but just concat a bunch of 0's not only for efficiency but to make
+  # sure that it doesn't contribute anything to the terms
+  # TODO(avaswani): Better way to do this.
+  q = _compute(node_states, total_key_depth, 1, "VALID", "q_mpnn")
+  q_shape = common_layers.shape_list(q)
+  # k and v edge transforms have shape
+  # [batch, length, depth*nonignored_edge_types]
+  nonignored_edge_types = num_edge_types-int(ignore_zero)
+  k = _compute(memory_antecedent, total_key_depth*nonignored_edge_types, 1,
+               "VALID", "k_mpnn")
+  v = _compute(memory_antecedent, total_value_depth*nonignored_edge_types,
+               1, "VALID", "v_mpnn")
+  batch = q_shape[0]
+  length = q_shape[1]
+  k = tf.reshape(k,
+                 [batch, length, nonignored_edge_types, total_key_depth])
+  v = tf.reshape(v,
+                 [q_shape[0], q_shape[1], nonignored_edge_types,
+                  total_value_depth])
+  if ignore_zero:
+    k = tf.pad(k, [[0, 0], [0, 0], [1, 0], [0, 0]])
+    v = tf.pad(v, [[0, 0], [0, 0], [1, 0], [0, 0]])
+
+  k = tf.reshape(k,
+                 [q_shape[0], q_shape[1]*num_edge_types, total_key_depth])
+  v = tf.reshape(v,
+                 [q_shape[0], q_shape[1]*num_edge_types, total_value_depth])
+  return q, k, v
+
+
+def multihead_mpnn_attention(node_states,
+                             total_key_depth,
+                             total_value_depth,
+                             output_depth,
+                             num_heads,
+                             adjacency_matrix=None,
+                             num_edge_types=5,
+                             ignore_zero=True,
+                             name="mpnn_attention"):
+  """Multihead scaled-dot-product attention with input/output transformations.
+
+  Args:
+    node_states: A tensor of shape [batch, length, depth]
+    total_key_depth: An integer for key dimension
+    total_value_depth: An integer for value dimensions
+    output_depth: An intger for output dimemsions
+    num_heads: An integer
+    adjacency_matrix: An tensor of ints of shape [batch, length, length]
+    num_edge_types: An integer indicating number of edge bins
+    ignore_zero: A flag that says that edge type 0 should be ignored
+    name: A string
+
+  Returns:
+    The result of the attention transformation. The output shape is
+        [batch_size, length_q, output_depth]
+    unless the cache dict is provided in which case only the last memory
+    position is calculated and the output shape is [batch_size, 1, hidden_dim]
+    Optionaly returns an additional loss parameters (ex: load balance loss for
+    the experts) returned by the attention_type function.
+
+  Raises:
+    ValueError: if the key depth or value depth are not divisible by the
+      number of attention heads.
+  """
+  if total_key_depth % num_heads != 0:
+    raise ValueError("Key depth (%d) must be divisible by the number of "
+                     "attention heads (%d)." % (total_key_depth, num_heads))
+  if total_value_depth % num_heads != 0:
+    raise ValueError("Value depth (%d) must be divisible by the number of "
+                     "attention heads (%d)." % (total_value_depth, num_heads))
+  with tf.variable_scope(
+      name,
+      default_name="multihead_mpnn_attention",
+      values=[node_states]):
+    q, k, v = compute_mpnn_qkv(node_states,
+                               total_key_depth,
+                               total_value_depth,
+                               num_edge_types,
+                               ignore_zero=ignore_zero)
+    # reshaping k and v for head splitting
+    q_shape = tf.shape(q)
+    q = common_attention.split_heads(q, num_heads)
+    k = common_attention.split_heads(k, num_heads)
+    v = common_attention.split_heads(v, num_heads)
+    key_depth_per_head = total_key_depth // num_heads
+    q *= key_depth_per_head**-0.5
+    # make the heads dimension leading. We will loop over heads.
+    q = tf.transpose(q, [1, 0, 2, 3])
+    k = tf.transpose(k, [1, 0, 2, 3])
+    v = tf.transpose(v, [1, 0, 2, 3])
+    # putting edge as the dimension after batch for k and v
+    # k and v will be [heads, batch, num_edge_types, length, depth]
+    k = tf.reshape(k, [num_heads, q_shape[0], q_shape[1], num_edge_types,
+                       total_key_depth//num_heads])
+    k = tf.transpose(k, [0, 1, 3, 2, 4])
+
+    v = tf.reshape(v, [num_heads, q_shape[0], q_shape[1], num_edge_types,
+                       total_value_depth//num_heads])
+    v = tf.transpose(v, [0, 1, 3, 2, 4])
+
+    # doing attention separately for each head
+    head_outputs = []
+    for head_id in range(num_heads):
+      output = dot_product_mpnn_attention(q[head_id],
+                                          k[head_id],
+                                          v[head_id],
+                                          adjacency_matrix,
+                                          num_edge_types)
+      head_outputs.append(tf.expand_dims(output, axis=0))
+    # making x = [heads, batch, length, total_value_depth//num_heads]
+    x = tf.concat(head_outputs, axis=0)
+    x = tf.transpose(x, [1, 0, 2, 3])
+    # making x [batch, length, depth]
+    x = common_attention.combine_heads(x)
+    x = common_layers.dense(
+        x, output_depth, use_bias=False, name="output_transform")
+    return x
+
+
+def dot_product_mpnn_attention(q, k, v, adjacency_matrix, num_edge_types,
+                               ignore_zero=True, name=None):
+  """Dot product attention with edge vectors.
+
+  Args:
+    q: [batch, length, key_depth] tensor
+    k: [batch, num_edge_types, length, key_depth]
+    v: [batch, num_edge_types, length, depth]
+    adjacency_matrix: [batch, length, length] tensor of int edge types
+    num_edge_types: an int, specifying number of edge types
+    ignore_zero: A flag that says that edge type 0 should be ignored
+    name: optional string
+
+  Returns:
+    A tensor of shape [batch, length, depth(q)]
+  """
+  with tf.variable_scope(
+      name, default_name="dot_product_mpnn_attention",
+      values=[q, k, v, adjacency_matrix, num_edge_types]):
+    # Computing attention mask
+    # all edge logits will have shape [batch, edge_types, len, len]
+    all_edge_logits = tf.matmul(
+        tf.tile(tf.expand_dims(q, axis=1), [1, num_edge_types, 1, 1]),
+        k, transpose_b=True)
+    # adjacency_matrix_one_hot has shape [batch, len, len, num_edge_types]
+    adjacency_matrix_one_hot = tf.one_hot(adjacency_matrix, num_edge_types)
+    # making adjacency_matrix_one_hot [batch, edge_types, len, len]
+    adjacency_matrix_one_hot = tf.transpose(adjacency_matrix_one_hot,
+                                            [0, 3, 1, 2])
+    # getting dot products for q_i, k_j, and e_{ij}. This assumes that for
+    # edge type 0, the dot products are 0
+    all_edge_logits *= adjacency_matrix_one_hot
+    # logits will be [batch, length, length] after educing along
+    # axis 1 which has dimension num_edge_types.
+    logits = tf.reduce_sum(all_edge_logits, axis=1)
+    # ignoring edges if needed
+    bias = 0
+    if ignore_zero:
+      bias = tf.to_float(tf.equal(adjacency_matrix, 0)) * -1e9
+    logits += bias
+    # getting compatibilities
+    compatibility = tf.nn.softmax(logits)
+    common_attention.attention_image_summary(
+        tf.expand_dims(compatibility, axis=1), None)
+    # getting edge compatibilities ready to compute values.
+    # after tiling, edge_compatibility will be
+    # [batch, num_edge_types, length, length]
+    edge_compatibility = tf.tile(
+        tf.expand_dims(compatibility, axis=1), [1, num_edge_types, 1, 1])
+    # computing values
+    edge_compatibility *= adjacency_matrix_one_hot
+    # all edge values will be [batch, num_edge_types, length, depth]
+    # We also assumed that the linear transformations for edge_type 0 will
+    # all be zeros. That is [batch, 0] is a length*depth tensor of 0's
+    all_edge_values = tf.matmul(edge_compatibility, v)
+    # reducing along the num_edge_types dimension
+    output = tf.reduce_sum(all_edge_values, axis=1)
+    return output
+
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index e4778468b..4e9d6407e 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -967,7 +967,8 @@ def transformer_ffn_layer(x,
                           conv_padding="LEFT",
                           nonpadding_mask=None,
                           losses=None,
-                          cache=None):
+                          cache=None,
+                          readout_filter_size=0):
   """Feed-forward layer in the transformer.
 
   Args:
@@ -984,6 +985,8 @@ def transformer_ffn_layer(x,
     losses: optional list onto which to append extra training losses
     cache: dict, containing tensors which are the results of previous
         attentions, used for fast decoding.
+    readout_filter_size: if it's greater than 0, then it will be used instead of
+      filter_size
 
 
   Returns:
@@ -1020,7 +1023,7 @@ def transformer_ffn_layer(x,
   elif ffn_layer == "conv_relu_conv":
     return common_layers.conv_relu_conv(
         x,
-        hparams.filter_size,
+        readout_filter_size or hparams.filter_size,
         hparams.hidden_size,
         first_kernel_size=hparams.conv_first_kernel,
         second_kernel_size=1,
@@ -1032,12 +1035,13 @@ def transformer_ffn_layer(x,
     return common_attention.parameter_attention(
         x, hparams.parameter_attention_key_channels or hparams.hidden_size,
         hparams.parameter_attention_value_channels or hparams.hidden_size,
-        hparams.hidden_size, hparams.filter_size, hparams.num_heads,
+        hparams.hidden_size, readout_filter_size or hparams.filter_size,
+        hparams.num_heads,
         hparams.attention_dropout)
   elif ffn_layer == "conv_hidden_relu_with_sepconv":
     return common_layers.conv_hidden_relu(
         x,
-        hparams.filter_size,
+        readout_filter_size or hparams.filter_size,
         hparams.hidden_size,
         kernel_size=(3, 1),
         second_kernel_size=(31, 1),
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index ed086b2b0..9a644f9eb 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -55,6 +55,7 @@ class Metrics(object):
   ROC_AUC = "roc_auc"
   IMAGE_SUMMARY = "image_summary"
   DMOL_PERPLEXITY = "disc_mol_neg_log_perplexity"
+  ABS_ERR = "mean_absolute_error"
   IMAGE_RMSE = "image_rmse"
 
 
@@ -77,6 +78,15 @@ def padded_rmse(predictions, labels, weights_fn=common_layers.weights_all):
   return error_sqrt, tf.reduce_sum(weights)
 
 
+def abs_error(predictions, labels, weights_fn=None):
+  """Computes mean(abs(preds-target))."""
+  del weights_fn  # Unused
+  targets = tf.squeeze(labels, axis=[2, 3])
+  batch_abs_error = tf.abs(predictions - targets)
+  den = tf.ones(tf.shape(batch_abs_error), dtype=tf.float32)
+  return (batch_abs_error, den)
+
+
 def padded_log_poisson(predictions,
                        labels,
                        weights_fn=common_layers.weights_all):
@@ -621,5 +631,6 @@ def metric_means():
     Metrics.ROC_AUC: roc_auc,
     Metrics.IMAGE_SUMMARY: image_summary,
     Metrics.DMOL_PERPLEXITY: dmol_neg_log_perplexity,
+    Metrics.ABS_ERR: abs_error,
     Metrics.IMAGE_RMSE: image_rmse,
 }

From d8b9b3b821fd18177062becc29bf1c5061fe6b92 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 18 Jun 2018 14:13:28 -0700
Subject: [PATCH 0136/2720] Small datagen correction to only exclude parsing
 problems that need the flag.

PiperOrigin-RevId: 201056385
---
 tensor2tensor/bin/t2t_datagen.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index 30fc5e50b..91ff65922 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -149,7 +149,7 @@ def main(_):
     problems = [p for p in problems if "timit" not in p]
   # Remove parsing if paths are not given.
   if not FLAGS.parsing_path:
-    problems = [p for p in problems if "parsing" not in p]
+    problems = [p for p in problems if "parsing_english_ptb" not in p]
 
   if not problems:
     problems_str = "\n  * ".join(

From 6b9ff21e8778e066ba2e0bd8942dc376788ad6fb Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 18 Jun 2018 14:16:37 -0700
Subject: [PATCH 0137/2720] Update UniversalTransformer to attempt slow greedy
 infer on TPU if use_tpu=True

PiperOrigin-RevId: 201056940
---
 tensor2tensor/models/research/universal_transformer.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index c9ee89a69..19ec7c97d 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -215,12 +215,13 @@ def body(self, features):
 
     return decoder_output
 
-  def _greedy_infer(self, features, decode_length):
+  def _greedy_infer(self, features, decode_length, use_tpu=False):
     """Fast version of greedy decoding.
 
     Args:
       features: an map of string to `Tensor`
       decode_length: an integer.  How many additional timesteps to decode.
+      use_tpu: bool, whether to use the TPU codepath.
 
     Returns:
       A dict of decoding results {
@@ -234,8 +235,8 @@ def _greedy_infer(self, features, decode_length):
     Raises:
       NotImplementedError: If there are multiple data shards.
     """
-    # TODO(dehghani): Support fast decoding for Universal Transformer
-    return self._slow_greedy_infer(features, decode_length)
+    return (self._slow_greedy_infer_tpu(features, decode_length) if use_tpu else
+            self._slow_greedy_infer(features, decode_length))
 
   def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
     """Beam search decoding.

From 5545e991a96ae9a8b740f3cec3478fcf962721b7 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 18 Jun 2018 15:18:12 -0700
Subject: [PATCH 0138/2720] Lazy load inplace_ops to maintain older TF
 compatibility

PiperOrigin-RevId: 201068484
---
 tensor2tensor/utils/t2t_model.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index c29509c6e..09ae083e8 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -41,9 +41,9 @@
 import tensorflow as tf
 
 from tensorflow.python.layers import base
-from tensorflow.python.ops import inplace_ops
 from tensorflow.python.ops import variable_scope
 
+
 _no_problem_err_str = (
     "The default implementation of %s requires that the "
     "model be used with a Problem. If using a Problem, augment the "
@@ -53,6 +53,12 @@
     lambda method_name: _no_problem_err_str % (method_name, method_name))
 
 
+# Lazy load inplace_ops because moudle is only available in TF 1.8+
+def tf_inplace_ops():
+  from tensorflow.python.ops import inplace_ops  # pylint: disable=g-import-not-at-top
+  return inplace_ops
+
+
 class T2TModel(base.Layer):
   """Abstract base class for models.
 
@@ -777,7 +783,7 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
       else:
         cur_sample = samples[:, i, :, :]
       samples = tf.transpose(recent_output, perm=[1, 0, 2, 3])
-      samples = inplace_ops.alias_inplace_update(
+      samples = tf_inplace_ops().alias_inplace_update(
           samples, i, tf.to_int64(cur_sample))
       samples = tf.transpose(samples, perm=[1, 0, 2, 3])
       if not tf.contrib.eager.in_eager_mode():
@@ -785,7 +791,7 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
 
       # Assuming we have one shard for logits.
       recent_logits = tf.transpose(recent_logits, perm=[1, 0, 2, 3, 4])
-      recent_logits = inplace_ops.alias_inplace_update(
+      recent_logits = tf_inplace_ops().alias_inplace_update(
           recent_logits, i, tf.squeeze(logits[:, -1:], axis=1))
       logits = tf.transpose(recent_logits, perm=[1, 0, 2, 3, 4])
       loss = sum([l for l in losses.values() if l is not None])

From 889cafb605be2e6d0f0f69bd842929ba9d1f2f2d Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 18 Jun 2018 16:08:24 -0700
Subject: [PATCH 0139/2720] Disable slow greedy infer TPU test on TF < 1.8

PiperOrigin-RevId: 201076653
---
 tensor2tensor/models/transformer_test.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index bbfa3d742..149f6c373 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -241,6 +241,10 @@ def testTransformerWithEncoderDecoderAttentionLoss(self):
     self.assertEqual(res.shape, ())
 
   def testGreedySlowTPUVsNonTPU(self):
+    # Only works with TF 1.8+
+    major, minor, _ = [int(el) for el in tf.__version__.split(".")]
+    if major < 1 or (major == 1 and minor < 8):
+      return
     model, features = get_model(transformer.transformer_small())
 
     decode_length = 3

From ed4b81ab6160a836796d75ca4821a1e3a398635d Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 18 Jun 2018 16:13:19 -0700
Subject: [PATCH 0140/2720] Switch to the new video batching by default.

PiperOrigin-RevId: 201077397
---
 tensor2tensor/data_generators/video_utils.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 21fec5358..bf20c94ec 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -106,7 +106,7 @@ def only_keep_videos_from_0th_frame(self):
 
   @property
   def use_not_breaking_batching(self):
-    return False
+    return True
 
   def preprocess_example(self, example, mode, hparams):
     """Runtime preprocessing, e.g., resize example["frame"]."""
@@ -224,12 +224,15 @@ def check_integrity_and_batch(*datasets):
         Returns:
           batched data and the integrity flag.
         """
-        frame_numbers = [dataset["frame_number"][0] for dataset in datasets]
-
-        not_broken = tf.equal(
-            frame_numbers[-1] - frame_numbers[0], num_frames-1)
-        if self.only_keep_videos_from_0th_frame:
-          not_broken = tf.logical_and(not_broken, tf.equal(frame_numbers[0], 0))
+        not_broken = tf.constant(True)
+        if "frame_number" in datasets[0]:
+          frame_numbers = [dataset["frame_number"][0] for dataset in datasets]
+
+          not_broken = tf.equal(
+              frame_numbers[-1] - frame_numbers[0], num_frames-1)
+          if self.only_keep_videos_from_0th_frame:
+            not_broken = tf.logical_and(not_broken,
+                                        tf.equal(frame_numbers[0], 0))
 
         features = {}
         for key in datasets[0].keys():

From 79e5a6a1dc692f61621a124ec0fcc2cc1ae361ae Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 18 Jun 2018 16:54:22 -0700
Subject: [PATCH 0141/2720] Adding a script for video metric computation.

PiperOrigin-RevId: 201084099
---
 tensor2tensor/utils/compute_video_metrics.py |  43 ++++++
 tensor2tensor/utils/video_metrics.py         | 137 +++++++++++++++++++
 2 files changed, 180 insertions(+)
 create mode 100644 tensor2tensor/utils/compute_video_metrics.py
 create mode 100644 tensor2tensor/utils/video_metrics.py

diff --git a/tensor2tensor/utils/compute_video_metrics.py b/tensor2tensor/utils/compute_video_metrics.py
new file mode 100644
index 000000000..a13d6c397
--- /dev/null
+++ b/tensor2tensor/utils/compute_video_metrics.py
@@ -0,0 +1,43 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Computes and saves the metrics for video prediction and generation."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.bin import t2t_decoder
+from tensor2tensor.utils import video_metrics
+import tensorflow as tf
+
+
+FLAGS = tf.flags.FLAGS
+
+
+def main(_):
+  hparams = t2t_decoder.create_hparams()
+  problem = hparams.problem
+  frame_shape = [problem.frame_height,
+                 problem.frame_width,
+                 problem.num_channels]
+  video_metrics.compute_and_save_video_metrics(
+      FLAGS.output_dir,
+      FLAGS.problem,
+      hparams.video_num_target_frames,
+      frame_shape)
+
+
+if __name__ == "__main__":
+  tf.app.run(main)
diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
new file mode 100644
index 000000000..aabc76a0e
--- /dev/null
+++ b/tensor2tensor/utils/video_metrics.py
@@ -0,0 +1,137 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Computes the metrics for video prediction and generation."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import numpy as np
+import six
+import tensorflow as tf
+
+
+def load_image_map_function(filename, frame_shape):
+  image = tf.read_file(filename)
+  image = tf.image.decode_png(image)
+  image = tf.image.resize_images(image, frame_shape[0:2])
+  image.set_shape(frame_shape)
+  return image
+
+
+def load_videos(template, video_length, frame_shape):
+  """Loads videos from files.
+
+  Args:
+    template: template string for listing the image files.
+    video_length: length of the video.
+    frame_shape: shape of each frame.
+
+  Returns:
+    dataset: the tf dataset frame by frame.
+    dataset_len: number of the items which is the number of image files.
+
+  Raises:
+    ValueError: if no files found.
+  """
+  filenames = tf.gfile.Glob(template)
+  if not filenames:
+    raise ValueError("no files found.")
+  filenames = sorted(filenames)
+  dataset_len = len(filenames)
+  filenames = tf.constant(filenames)
+  dataset = tf.data.Dataset.from_tensor_slices(filenames)
+  dataset = dataset.apply(tf.contrib.data.map_and_batch(
+      lambda filename: load_image_map_function(filename, frame_shape),
+      video_length, drop_remainder=True))
+  return dataset, dataset_len
+
+
+def file_pattern(output_dir, problem_name, prefix):
+  return os.path.join(output_dir, "{}_{}*.png".format(problem_name, prefix))
+
+
+def get_target_and_output_filepatterns(output_dir, problem_name):
+  return (file_pattern(output_dir, problem_name, "outputs"),
+          file_pattern(output_dir, problem_name, "targets"))
+
+
+def get_zipped_dataset(output_files, target_files, video_length, frame_shape):
+  outputs, len_ = load_videos(output_files, video_length, frame_shape)
+  targets, len_ = load_videos(target_files, video_length, frame_shape)
+  zipped_dataset = tf.data.Dataset.zip((outputs, targets))
+  num_videos = len_ // video_length
+  return zipped_dataset, num_videos
+
+
+def save_results(results, output_dir, problem_name):
+  for name, array in six.iteritems(results):
+    output_filename = "{}_{}.npy".format(problem_name, name)
+    output_filename = os.path.join(output_dir, output_filename)
+    np.save(output_filename, array)
+
+
+def compute_metrics(output_video, target_video):
+  max_pixel_value = 255.0
+  psnr = tf.image.psnr(output_video, target_video, max_pixel_value)
+  ssim = tf.image.ssim(output_video, target_video, max_pixel_value)
+  return {"PSNR": psnr, "SSIM": ssim}
+
+
+def compute_video_metrics(output_dir, problem_name, video_length, frame_shape):
+  """Computes the average of all the metric over the whole dataset.
+
+  This function assumes that all the predicted and target frames
+  have been saved on the disk and sorting them by name will result
+  to consecutive frames saved in order.
+
+  Args:
+    output_dir: directory with all the saved frames.
+    problem_name: prefix of the saved frames usually name of the problem.
+    video_length: length of the videos.
+    frame_shape: shape of each frame in HxWxC format.
+
+  Returns:
+    Dictionary which contains the average of each metric per frame.
+  """
+  output_files, target_files = get_target_and_output_filepatterns(
+      output_dir, problem_name)
+  dataset, num_videos = get_zipped_dataset(
+      output_files, target_files, video_length, frame_shape)
+  output, target = dataset.make_one_shot_iterator().get_next()
+  metrics_dict = compute_metrics(output, target)
+  metrics_names, metrics = zip(*six.iteritems(metrics_dict))
+  means, update_ops = tf.metrics.mean_tensor(metrics)
+
+  with tf.Session() as sess:
+    sess.run(tf.local_variables_initializer())
+
+    # Compute mean over dataset
+    for i in range(num_videos):
+      print("Computing video: %d" % i)
+      sess.run(update_ops)
+    averaged_metrics = sess.run(means)
+
+    results = dict(zip(metrics_names, averaged_metrics))
+    return results
+
+
+def compute_and_save_video_metrics(
+    output_dir, problem_name, video_length, frame_shape):
+  results = compute_video_metrics(
+      output_dir, problem_name, video_length, frame_shape)
+  save_results(results, output_dir, problem_name)
+

From c3edd0c74152a618a71bd5a7a56ade749b337f23 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 18 Jun 2018 17:21:20 -0700
Subject: [PATCH 0142/2720] Add optional prefix to list_hparams

PiperOrigin-RevId: 201088315
---
 tensor2tensor/models/__init__.py | 2 +-
 tensor2tensor/utils/registry.py  | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 2aec38a24..0130ebd06 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -18,7 +18,7 @@
 from __future__ import print_function
 # pylint: disable=unused-import
 
-from tensor2tensor.layers import modalities
+from tensor2tensor.layers import modalities  # pylint: disable=g-import-not-at-top
 from tensor2tensor.models import basic
 from tensor2tensor.models import bytenet
 from tensor2tensor.models import distillation
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 318abac96..525c250f5 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -178,7 +178,9 @@ def hparams(name):
   return hp
 
 
-def list_hparams():
+def list_hparams(prefix=None):
+  if prefix:
+    return [name for name in _HPARAMS if name.startswith(prefix)]
   return list(_HPARAMS)
 
 

From e9b78321000d89e8e3ab30e0fc02b9706f09bde7 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 19 Jun 2018 10:26:20 +0200
Subject: [PATCH 0143/2720] small clean-ups

---
 tensor2tensor/data_generators/gym_problems.py | 218 ++++++------------
 tensor2tensor/rl/rl_trainer_lib.py            |  24 +-
 2 files changed, 78 insertions(+), 164 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 50aa273de..b308eef54 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -24,23 +24,14 @@
 import numpy as np
 
 # We need gym_utils for the game environments defined there.
-from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
-from tensor2tensor.layers import discretization
-from tensor2tensor.models.research import autoencoders
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research.rl import standard_atari_env_spec
 from tensor2tensor.rl import collect
-from tensor2tensor.rl.envs import tf_atari_wrappers as atari
-from tensor2tensor.rl.envs.tf_atari_wrappers import StackAndSkipWrapper
 from tensor2tensor.rl.envs.tf_atari_wrappers import TimeLimitWrapper
-from tensor2tensor.rl.envs.utils import batch_env_factory
-
-
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
-
 import tensorflow as tf
 
 
@@ -69,16 +60,6 @@ def __init__(self, *args, **kwargs):
     self.environment_spec = self.get_environment_spec()
     self.eval_phase = False
 
-    # Debug info
-    self.make_extra_debug_info = True
-    self.dones = 0
-    self.real_reward = 0
-    self.total_sim_reward = 0.0
-    self.total_real_reward = 0.0
-    self.sum_of_rewards = 0.0
-    self.successful_episode_reward_predictions = 0
-    self.report_reward_statistics_every = 10
-
   def _setup(self):
     collect_hparams = rl.ppo_pong_base()
     collect_hparams.add_hparam("environment_spec", self.environment_spec)
@@ -115,8 +96,6 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
         observ, reward, done, action = data
         observ = observ.astype(np.uint8) # TODO(piotrmilos). This should be probably done in collect
         debug_im = None
-        if self.make_extra_debug_info:
-          debug_im = self.get_debug_image(data)
 
         ret_dict = {"frame": observ,
                     "image/format": ["png"],
@@ -125,14 +104,10 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
                     "action": [int(action)],
                     "done": [int(False)],
                     "reward": [int(reward) - self.min_reward]}
-        if self.make_extra_debug_info:
-          ret_dict["image/debug"] = debug_im
+
         yield ret_dict
         pieces_generated += 1
 
-  def get_debug_image(self, data):
-    raise NotImplemented()
-
   def restore_networks(self, sess):
     if FLAGS.agent_policy_path:
       model_saver = tf.train.Saver(
@@ -141,7 +116,6 @@ def restore_networks(self, sess):
       ckpt = ckpts.model_checkpoint_path
       model_saver.restore(sess, ckpt)
 
-
   @property
   def num_input_frames(self):
     """Number of frames on input for real environment."""
@@ -241,6 +215,80 @@ def hparams(self, defaults, unused_model_hparams):
 class GymAEDiscreteProblem(GymDiscreteProblem):
   pass
 
+
+class GymRealDiscreteProblem(GymDiscreteProblem):
+
+  def __init__(self, *args, **kwargs):
+    super(GymRealDiscreteProblem, self).__init__(*args, **kwargs)
+    self.make_extra_debug_info = False
+
+  def get_environment_spec(self):
+    return standard_atari_env_spec(self.env_name)
+
+
+
+class GymSimulatedDiscreteProblem(GymDiscreteProblem):
+  """Simulated gym environment with discrete actions and rewards."""
+
+  def __init__(self, *args, **kwargs):
+    self.simulated_environment = True
+    self.debug_dump_frames_path = "debug_frames_sim"
+    self.intrinsic_reward_scale = 0.0
+    self.simulation_random_starts = False
+    super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
+
+  @property
+  def initial_frames_problem(self):
+    raise NotImplemented()
+
+  def get_environment_spec(self):
+    env_spec = standard_atari_env_spec(self.env_name)
+
+    #Set reasonable time limit (as we do not simulate done)
+    real_env = env_spec.env_lambda()
+    if self.num_testing_steps is not None:
+      timelimit = self.num_testing_steps
+    else:
+      try:
+        # We assume that the real env is wrapped with TimeLimit.
+        history = self.num_input_frames
+        timelimit = real_env._max_episode_steps - history  # pylint: disable=protected-access
+      except:  # pylint: disable=bare-except
+        # If not, set some reasonable default.
+        timelimit = 100
+
+    env_spec.simulated_env = True
+    env_spec.add_hparam("simulation_random_starts",
+                           self.simulation_random_starts)
+    env_spec.add_hparam("intrinsic_reward_scale",
+                           self.intrinsic_reward_scale)
+    initial_frames_problem = registry.problem(self.initial_frames_problem)
+    env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
+    env_spec.wrappers.append([TimeLimitWrapper, {"timelimit": timelimit}])
+
+    return env_spec
+
+  def restore_networks(self, sess):
+    super(GymSimulatedDiscreteProblem, self).restore_networks(sess)
+    # TODO(blazej): adjust regexp for different models.
+    env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
+    sess = tf.get_default_session()
+
+    ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
+    ckpt = ckpts.model_checkpoint_path
+    env_model_loader.restore(sess, ckpt)
+
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnPong(
+    GymSimulatedDiscreteProblem, GymPongRandom):
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_pong"
+
+
 @registry.register_problem
 class GymPongRandom(GymDiscreteProblem):
   """Pong game, random actions."""
@@ -343,122 +391,6 @@ def min_reward(self):
   def num_rewards(self):
     return 2
 
-
-class GymRealDiscreteProblem(GymDiscreteProblem):
-
-  def __init__(self, *args, **kwargs):
-    super(GymRealDiscreteProblem, self).__init__(*args, **kwargs)
-    self.make_extra_debug_info = False
-
-  def get_debug_image(self):
-    #TODO(piotrmilos): possibly change this
-    raise NotImplemented()
-
-  def get_environment_spec(self):
-    return standard_atari_env_spec(self.env_name)
-
-
-
-class GymSimulatedDiscreteProblem(GymDiscreteProblem):
-  """Simulated gym environment with discrete actions and rewards."""
-
-  def __init__(self, *args, **kwargs):
-    self.simulated_environment = True
-    self.make_extra_debug_info = True
-    self.debug_dump_frames_path = "debug_frames_sim"
-    self.intrinsic_reward_scale = 0.0
-    self.simulation_random_starts = False
-    super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
-
-  def _setup(self):
-    super(GymSimulatedDiscreteProblem, self)._setup()
-    if self.make_extra_debug_info:
-      # Slight weirdness to make sim env and real env aligned
-      self.real_env.reset()
-      for _ in range(self.num_input_frames):
-        self.real_ob, _, _, _ = self.real_env.step(0)
-      self.total_sim_reward, self.total_real_reward = 0.0, 0.0
-      self.sum_of_rewards = 0.0
-      self.successful_episode_reward_predictions = 0
-
-  def get_debug_image(self, data):
-    observ, reward, done, action = data
-    self.total_sim_reward += reward
-    err = np.ndarray.astype(np.maximum(np.abs(
-      self.real_ob - observ, dtype=np.int) - 10, 0),
-                            np.uint8)
-    debug_im = np.concatenate([observ, self.real_ob, err], axis=1)
-
-    if done:
-      self.dones += 1
-      self.sum_of_rewards += self.real_reward
-      if self.total_real_reward == self.total_sim_reward:
-        self.successful_episode_reward_predictions += 1
-
-      self.total_real_reward = 0.0
-      self.total_sim_reward = 0.0
-      self.real_reward = 0
-      # Slight weirdness to make sim env and real env aligned
-      for _ in range(self.num_input_frames):
-        self.real_ob, _, _, _ = self.real_env.step(0)
-    else:
-      self.real_ob, self.real_reward, _, _ = self.real_env.step(action)
-      self.total_real_reward += self.real_reward
-      self.sum_of_rewards += self.real_reward
-
-    return debug_im
-
-  @property
-  def initial_frames_problem(self):
-    raise NotImplemented()
-
-  def get_environment_spec(self):
-    env_spec = standard_atari_env_spec(self.env_name)
-
-    #Set reasonable time limit (as we do not simulate done)
-    self.real_env = env_spec.env_lambda()
-    if self.num_testing_steps is not None:
-      timelimit = self.num_testing_steps
-    else:
-      try:
-        # We assume that the real env is wrapped with TimeLimit.
-        history = self.num_input_frames
-        timelimit = self.real_env._max_episode_steps - history  # pylint: disable=protected-access
-      except:  # pylint: disable=bare-except
-        # If not, set some reasonable default.
-        timelimit = 100
-
-    env_spec.simulated_env = True
-    env_spec.add_hparam("simulation_random_starts",
-                           self.simulation_random_starts)
-    env_spec.add_hparam("intrinsic_reward_scale",
-                           self.intrinsic_reward_scale)
-    initial_frames_problem = registry.problem(self.initial_frames_problem)
-    env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
-    env_spec.wrappers.append([TimeLimitWrapper, {"timelimit": timelimit}])
-
-    return env_spec
-
-  def restore_networks(self, sess):
-    super(GymSimulatedDiscreteProblem, self).restore_networks(sess)
-    # TODO(blazej): adjust regexp for different models.
-    env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
-    sess = tf.get_default_session()
-
-    ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
-    ckpt = ckpts.model_checkpoint_path
-    env_model_loader.restore(sess, ckpt)
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnPong(
-    GymSimulatedDiscreteProblem, GymPongRandom):
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_pong"
-
-
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnPong(
     GymRealDiscreteProblem, GymPongRandom):
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index f431afe47..6ae473d20 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -41,26 +41,7 @@ def define_train(hparams, event_dir):
     ppo_summary = ppo.define_ppo_epoch(memory, hparams)
     summary = tf.summary.merge([collect_summary, ppo_summary])
 
-  # with tf.variable_scope("eval", reuse=tf.AUTO_REUSE):
-  #   eval_env_lambda = env_lambda
-  #   if event_dir and hparams.video_during_eval:
-  #     # Some environments reset environments automatically, when reached done
-  #     # state. For them we shall record only every second episode.
-  #     d = 2 if env_lambda().metadata.get("semantics.autoreset") else 1
-  #     eval_env_lambda = lambda: gym.wrappers.Monitor(  # pylint: disable=g-long-lambda
-  #         env_lambda(), event_dir, video_callable=lambda i: i % d == 0)
-  #     eval_env_lambda = (
-  #         lambda: utils.EvalVideoWrapper(eval_env_lambda()))
-  #   eval_batch_env = utils.batch_env_factory(
-  #       eval_env_lambda, hparams,
-  #       num_agents=hparams.num_eval_agents, xvfb=hparams.video_during_eval)
-  #
-  #   _, eval_summary = collect_new.define_collect(
-  #       policy_factory, eval_batch_env, hparams, eval_phase=True)
-
-  #Fake for development
-  eval_summary = summary
-  return summary, eval_summary
+  return summary, None
 
 
 def train(hparams, event_dir=None, model_dir=None,
@@ -106,7 +87,8 @@ def train(hparams, event_dir=None, model_dir=None,
           summary_writer.add_summary(summary, epoch_index)
         if (hparams.eval_every_epochs and
             epoch_index % hparams.eval_every_epochs == 0):
-          summary = sess.run(eval_summary_op)
+          print("Eval is to be implemented") #TODO(piotrmilos):implement
+          # summary = sess.run(eval_summary_op)
           if summary_writer and summary:
             summary_writer.add_summary(summary, epoch_index)
           else:

From bad7e0983d284eae65ff06bb547a4235e5075547 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 19 Jun 2018 11:12:06 +0200
Subject: [PATCH 0144/2720] small bug fix

---
 tensor2tensor/data_generators/gym_problems.py | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index b308eef54..44010e3c4 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -279,16 +279,6 @@ def restore_networks(self, sess):
     env_model_loader.restore(sess, ckpt)
 
 
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnPong(
-    GymSimulatedDiscreteProblem, GymPongRandom):
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_pong"
-
-
 @registry.register_problem
 class GymPongRandom(GymDiscreteProblem):
   """Pong game, random actions."""
@@ -375,6 +365,15 @@ def num_rewards(self):
     return 3
 
 
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnPong(
+    GymSimulatedDiscreteProblem, GymPongRandom):
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_pong"
+
+
 @registry.register_problem
 class GymFreewayRandom(GymDiscreteProblem):
   """Freeway game, random actions."""

From 022061e574c9c935f605dde40e00c51ab90fe0e8 Mon Sep 17 00:00:00 2001
From: Jose Fonollosa <jose.fonollosa@upc.edu>
Date: Tue, 19 Jun 2018 11:58:22 +0200
Subject: [PATCH 0145/2720] Fix CMVN variance normalization

Variance normalization should be performed dividing by the square root of the variance, not the variance
---
 tensor2tensor/data_generators/speech_recognition.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 9eed9fe50..74b314bfc 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -310,10 +310,10 @@ def preprocess_example(self, example, mode, hparams):
       assert fbank_size[0] == 1
 
       # This replaces CMVN estimation on data
-
+      var_epsilon = 1e-09
       mean = tf.reduce_mean(mel_fbanks, keepdims=True, axis=1)
-      variance = tf.reduce_mean((mel_fbanks-mean)**2, keepdims=True, axis=1)
-      mel_fbanks = (mel_fbanks - mean) / variance
+      variance = tf.reduce_mean(tf.square(mel_fbanks - mean), keepdims=True, axis=1)
+      mel_fbanks = (mel_fbanks - mean) * tf.rsqrt(variance + var_epsilon)
 
       # Later models like to flatten the two spatial dims. Instead, we add a
       # unit spatial dim and flatten the frequencies and channels.
@@ -377,13 +377,14 @@ def bottom(self, x):
               nonpadding_mask) * num_mel_bins * num_channels
 
           # This replaces CMVN estimation on data
+          var_epsilon = 1e-09
           mean = tf.reduce_sum(
               x, axis=[1], keepdims=True) / num_of_nonpadding_elements
           variance = (num_of_nonpadding_elements * mean**2. -
                       2. * mean * tf.reduce_sum(x, axis=[1], keepdims=True) +
                       tf.reduce_sum(x**2, axis=[1], keepdims=True)
                      ) / num_of_nonpadding_elements
-          x = (x - mean) / variance * tf.expand_dims(nonpadding_mask, -1)
+          x = (x - mean) * tf.rsqrt(variance + var_epsilon) * tf.expand_dims(nonpadding_mask, -1)
       else:
         x = inputs
 

From e455347334d3ae418d8587303c8e9db55aae30e8 Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Tue, 19 Jun 2018 10:19:03 -0700
Subject: [PATCH 0146/2720] Fix T2TExperiment to support more frequent evals

PiperOrigin-RevId: 201197206
---
 tensor2tensor/bin/t2t_trainer.py   |  1 +
 tensor2tensor/utils/flags.py       |  3 ++
 tensor2tensor/utils/trainer_lib.py | 66 +++++++++++++++++++++---------
 3 files changed, 51 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 998bd2584..949deec39 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -163,6 +163,7 @@ def create_experiment_fn():
       eval_steps=FLAGS.eval_steps,
       min_eval_frequency=FLAGS.local_eval_frequency,
       schedule=FLAGS.schedule,
+      eval_throttle_seconds=FLAGS.eval_throttle_seconds,
       export=FLAGS.export_saved_model,
       decode_hparams=decoding.decode_hparams(FLAGS.decode_hparams),
       use_tfdbg=FLAGS.tfdbg,
diff --git a/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index d9b753997..fc8d66674 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -87,6 +87,9 @@
 flags.DEFINE_integer("local_eval_frequency", 1000,
                      "Save checkpoints and run evaluation every N steps during "
                      "local training.")
+flags.DEFINE_integer("eval_throttle_seconds", 600,
+                     "Do not re-evaluate unless the last evaluation was started"
+                     " at least this many seconds ago.")
 flags.DEFINE_bool("locally_shard_to_cpu", False,
                   "Use CPU as a sharding device running locally. This allows "
                   "to test sharded model construction on a machine with 1 GPU.")
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index f9781c3ae..34362c9d7 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -205,6 +205,8 @@ def create_estimator(model_name,
 def create_hooks(use_tfdbg=False,
                  use_dbgprofile=False,
                  dbgprofile_kwargs=None,
+                 use_validation_monitor=False,
+                 validation_monitor_kwargs=None,
                  use_early_stopping=False,
                  early_stopping_kwargs=None):
   """Create train and eval hooks for Experiment."""
@@ -224,6 +226,12 @@ def create_hooks(use_tfdbg=False,
     defaults.update(dbgprofile_kwargs)
     train_hooks.append(tf.train.ProfilerHook(**defaults))
 
+  if use_validation_monitor:
+    tf.logging.info("Using ValidationMonitor")
+    train_hooks.append(
+        tf.contrib.learn.monitors.ValidationMonitor(
+            hooks=eval_hooks, **validation_monitor_kwargs))
+
   if use_early_stopping:
     tf.logging.info("Using EarlyStoppingHook")
     hook = metrics_hook.EarlyStoppingHook(**early_stopping_kwargs)
@@ -244,16 +252,17 @@ def __init__(self, estimator, hparams, train_spec, eval_spec):
     self._estimator = estimator
 
   def continuous_train_and_eval(self):
+
     tf.estimator.train_and_evaluate(self._estimator, self._train_spec,
                                     self._eval_spec)
+    return self.evaluate()
 
   def train_and_evaluate(self):
-    tf.logging.warning(
-        "Note that train_and_evaluate now behaves the same as"
-        " continuous_train_and_eval. tensor2tensor no longer supports"
-        " training and evaluation in the same graph."
-    )
-    self.continuous_train_and_eval()
+    if self._eval_spec is None:
+      tf.logging.warning("EvalSpec not provided. Estimator will not manage "
+                         "model evaluation. Assuming ValidationMonitor present "
+                         "in train_hooks.")
+      self.train()
 
   def train(self):
     self._estimator.train(
@@ -262,7 +271,7 @@ def train(self):
         max_steps=self._train_spec.max_steps)
 
   def evaluate(self):
-    self._estimator.evaluate(
+    return self._estimator.evaluate(
         self._eval_spec.input_fn,
         steps=self._eval_spec.steps,
         hooks=self._eval_spec.hooks)
@@ -307,15 +316,14 @@ def continuous_eval_on_train_data(self):
 
   def test(self):
     """Perform 1 step of train and 2 step of eval."""
+    if self._eval_spec is None:
+      return self.train_and_evaluate()
+
     self._estimator.train(
-        self._train_spec.input_fn,
-        hooks=self._train_spec.hooks,
-        max_steps=1)
+        self._train_spec.input_fn, hooks=self._train_spec.hooks, max_steps=1)
 
     self._estimator.evaluate(
-        self._eval_spec.input_fn,
-        steps=1,
-        hooks=self._eval_spec.hooks)
+        self._eval_spec.input_fn, steps=1, hooks=self._eval_spec.hooks)
 
   def run_std_server(self):
     """Starts a TensorFlow server and joins the serving thread.
@@ -350,6 +358,7 @@ def create_experiment(run_config,
                       train_steps,
                       eval_steps,
                       min_eval_frequency=2000,
+                      eval_throttle_seconds=600,
                       schedule="train_and_evaluate",
                       export=False,
                       decode_hparams=None,
@@ -391,6 +400,13 @@ def create_experiment(run_config,
                     "See serving/export.py.")
 
   # Hooks
+  validation_monitor_kwargs = dict(
+      input_fn=eval_input_fn,
+      eval_steps=eval_steps,
+      every_n_steps=min_eval_frequency,
+      early_stopping_rounds=eval_early_stopping_steps,
+      early_stopping_metric=eval_early_stopping_metric,
+      early_stopping_metric_minimize=eval_early_stopping_metric_minimize)
   dbgprofile_kwargs = {"output_dir": run_config.model_dir}
   early_stopping_kwargs = dict(
       events_dir=os.path.join(run_config.model_dir, "eval_continuous"),
@@ -404,6 +420,8 @@ def create_experiment(run_config,
   if schedule == "continuous_train_and_eval" and min_eval_frequency:
     tf.logging.warn("ValidationMonitor only works with "
                     "--schedule=train_and_evaluate")
+  use_validation_monitor = (
+      schedule == "train_and_evaluate" and min_eval_frequency)
   # Distributed early stopping
   local_schedules = ["train_and_evaluate", "continuous_train_and_eval"]
   use_early_stopping = (
@@ -412,19 +430,29 @@ def create_experiment(run_config,
       use_tfdbg=use_tfdbg,
       use_dbgprofile=use_dbgprofile,
       dbgprofile_kwargs=dbgprofile_kwargs,
+      use_validation_monitor=use_validation_monitor,
+      validation_monitor_kwargs=validation_monitor_kwargs,
       use_early_stopping=use_early_stopping,
       early_stopping_kwargs=early_stopping_kwargs)
   train_hooks += t2t_model.T2TModel.get_train_hooks(model_name)
   eval_hooks += t2t_model.T2TModel.get_eval_hooks(model_name)
 
+  train_hooks = tf.contrib.learn.monitors.replace_monitors_with_hooks(
+      train_hooks, estimator)
+  eval_hooks = tf.contrib.learn.monitors.replace_monitors_with_hooks(
+      eval_hooks, estimator)
+
   train_spec = tf.estimator.TrainSpec(
       train_input_fn, max_steps=train_steps, hooks=train_hooks)
-  eval_spec = tf.estimator.EvalSpec(
-      eval_input_fn,
-      steps=eval_steps,
-      hooks=eval_hooks,
-      start_delay_secs=0 if hparams.schedule == "evaluate" else 120,
-      throttle_secs=600)
+  if use_validation_monitor:
+    eval_spec = None
+  else:
+    eval_spec = tf.estimator.EvalSpec(
+        eval_input_fn,
+        steps=eval_steps,
+        hooks=eval_hooks,
+        start_delay_secs=0 if hparams.schedule == "evaluate" else 120,
+        throttle_secs=eval_throttle_seconds)
 
   return T2TExperiment(estimator, hparams, train_spec, eval_spec)
 

From 1b52078f6c00c3a1be570795eed507ad9caf6f53 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 19 Jun 2018 11:46:44 -0700
Subject: [PATCH 0147/2720] refactoring SV2P code.

PiperOrigin-RevId: 201214081
---
 tensor2tensor/models/research/next_frame.py | 419 ++++++++++----------
 1 file changed, 216 insertions(+), 203 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 73f19f325..c0dfe5221 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -262,15 +262,153 @@ def conv_lstm_2d(self, inputs, state, output_channels,
     outputs, new_state = cell(inputs, state)
     return outputs, new_state
 
+  def construct_predictive_tower(
+      self, input_image, input_reward, action, lstm_state, latent):
+    # Main tower
+    layer_norm = tf.contrib.layers.layer_norm
+    lstm_func = self.conv_lstm_2d
+    lstm_size = np.array([32, 32, 64, 64, 128, 64, 32], dtype=np.int32)
+    batch_size = common_layers.shape_list(input_image)[0]
+    # the number of different pixel motion predictions
+    # and the number of masks for each of those predictions
+    num_masks = self.hparams.num_masks
+
+    img_height, img_width, color_channels = self.hparams.problem.frame_shape
+
+    with tf.variable_scope("main", reuse=tf.AUTO_REUSE):
+      input_image = common_layers.make_even_size(input_image)
+      enc0 = slim.layers.conv2d(
+          input_image,
+          32, [5, 5],
+          stride=2,
+          scope="scale1_conv1",
+          normalizer_fn=layer_norm,
+          normalizer_params={"scope": "layer_norm1"})
+
+      hidden1, lstm_state[0] = lstm_func(
+          enc0, lstm_state[0], lstm_size[0], scope="state1")
+      hidden1 = layer_norm(hidden1, scope="layer_norm2")
+      hidden2, lstm_state[1] = lstm_func(
+          hidden1, lstm_state[1], lstm_size[1], scope="state2")
+      hidden2 = layer_norm(hidden2, scope="layer_norm3")
+      hidden2 = common_layers.make_even_size(hidden2)
+      enc1 = slim.layers.conv2d(
+          hidden2, hidden2.get_shape()[3], [3, 3], stride=2, scope="conv2")
+
+      hidden3, lstm_state[2] = lstm_func(
+          enc1, lstm_state[2], lstm_size[2], scope="state3")
+      hidden3 = layer_norm(hidden3, scope="layer_norm4")
+      hidden4, lstm_state[3] = lstm_func(
+          hidden3, lstm_state[3], lstm_size[3], scope="state4")
+      hidden4 = layer_norm(hidden4, scope="layer_norm5")
+      hidden4 = common_layers.make_even_size(hidden4)
+      enc2 = slim.layers.conv2d(
+          hidden4, hidden4.get_shape()[3], [3, 3], stride=2, scope="conv3")
+
+      # Pass in reward and action.
+      emb_action = self.encode_to_shape(action, enc2.get_shape())
+      emb_reward = self.encode_to_shape(input_reward, enc2.get_shape())
+      enc2 = tf.concat(axis=3, values=[enc2, emb_action, emb_reward])
+
+      if latent is not None:
+        with tf.control_dependencies([latent]):
+          enc2 = tf.concat([enc2, latent], 3)
+
+      enc3 = slim.layers.conv2d(
+          enc2, hidden4.get_shape()[3], [1, 1], stride=1, scope="conv4")
+
+      hidden5, lstm_state[4] = lstm_func(
+          enc3, lstm_state[4], lstm_size[4], scope="state5")  # last 8x8
+      hidden5 = layer_norm(hidden5, scope="layer_norm6")
+      enc4 = slim.layers.conv2d_transpose(
+          hidden5, hidden5.get_shape()[3], 3, stride=2, scope="convt1")
+
+      enc1_shape = common_layers.shape_list(enc1)
+      enc4 = enc4[:, :enc1_shape[1], :enc1_shape[2], :]  # Cut to shape.
+      hidden6, lstm_state[5] = lstm_func(
+          enc4, lstm_state[5], lstm_size[5], scope="state6")  # 16x16
+      hidden6 = layer_norm(hidden6, scope="layer_norm7")
+      # Skip connection.
+      hidden6 = tf.concat(axis=3, values=[hidden6, enc1])  # both 16x16
+
+      enc5 = slim.layers.conv2d_transpose(
+          hidden6, hidden6.get_shape()[3], 3, stride=2, scope="convt2")
+      enc0_shape = common_layers.shape_list(enc0)
+      enc5 = enc5[:, :enc0_shape[1], :enc0_shape[2], :]  # Cut to shape.
+      hidden7, lstm_state[6] = lstm_func(
+          enc5, lstm_state[6], lstm_size[6], scope="state7")  # 32x32
+      hidden7 = layer_norm(hidden7, scope="layer_norm8")
+
+      # Skip connection.
+      hidden7 = tf.concat(axis=3, values=[hidden7, enc0])  # both 32x32
+
+      enc6 = slim.layers.conv2d_transpose(
+          hidden7,
+          hidden7.get_shape()[3],
+          3,
+          stride=2,
+          scope="convt3",
+          activation_fn=None,
+          normalizer_fn=layer_norm,
+          normalizer_params={"scope": "layer_norm9"})
+
+      if self.hparams.model_options == "DNA":
+        # Using largest hidden state for predicting untied conv kernels.
+        enc7 = slim.layers.conv2d_transpose(
+            enc6,
+            self.hparams.dna_kernel_size**2,
+            1,
+            stride=1,
+            scope="convt4",
+            activation_fn=None)
+      else:
+        # Using largest hidden state for predicting a new image layer.
+        enc7 = slim.layers.conv2d_transpose(
+            enc6,
+            color_channels,
+            1,
+            stride=1,
+            scope="convt4",
+            activation_fn=None)
+        # This allows the network to also generate one image from scratch,
+        # which is useful when regions of the image become unoccluded.
+        transformed = [tf.nn.sigmoid(enc7)]
+
+      if self.hparams.model_options == "CDNA":
+        # cdna_input = tf.reshape(hidden5, [int(batch_size), -1])
+        cdna_input = tf.contrib.layers.flatten(hidden5)
+        transformed += self.cdna_transformation(
+            input_image, cdna_input, num_masks, int(color_channels))
+      elif self.hparams.model_options == "DNA":
+        # Only one mask is supported (more should be unnecessary).
+        if num_masks != 1:
+          raise ValueError("Only one mask is supported for DNA model.")
+        transformed = [self.dna_transformation(input_image, enc7)]
+
+      masks = slim.layers.conv2d_transpose(
+          enc6, num_masks + 1, 1,
+          stride=1, scope="convt7", activation_fn=None)
+      masks = tf.reshape(
+          tf.nn.softmax(tf.reshape(masks, [-1, num_masks + 1])),
+          [batch_size,
+           int(img_height),
+           int(img_width), num_masks + 1])
+      mask_list = tf.split(
+          axis=3, num_or_size_splits=num_masks + 1, value=masks)
+      output = mask_list[0] * input_image
+      for layer, mask in zip(transformed, mask_list[1:]):
+        output += layer * mask
+
+      p_reward = self.reward_prediction(hidden5)
+      p_reward = self.decode_to_shape(p_reward, input_reward.shape)
+
+      return output, p_reward, lstm_state
+
   def construct_model(self,
                       images,
                       actions,
                       rewards,
-                      k=-1,
-                      num_masks=10,
-                      cdna=True,
-                      dna=False,
-                      context_frames=2):
+                      k=-1):
     """Build convolutional lstm video predictor using CDNA, or DNA.
 
     Args:
@@ -281,12 +419,6 @@ def construct_model(self,
       rewards: list of reward tensors
                each reward should be in the shape ?x1xZ
       k: constant used for scheduled sampling. -1 to feed in own prediction.
-      num_masks: the number of different pixel motion predictions (and
-                 the number of masks for each of those predictions)
-      cdna: True to use Convoluational Dynamic Neural Advection (CDNA)
-      dna: True to use Dynamic Neural Advection (DNA)
-      context_frames: number of ground truth frames to pass in before
-                      feeding in own predictions
     Returns:
       gen_images: predicted future image frames
       gen_rewards: predicted future rewards
@@ -294,40 +426,20 @@ def construct_model(self,
       latent_std: std of approximated posterior
 
     Raises:
-      ValueError: if more than one network option specified or more than 1 mask
-      specified for DNA model.
+      ValueError: if more than 1 mask specified for DNA model.
     """
     # Each image is being used twice, in latent tower and main tower.
     # This is to make sure we are using the *same* image for both, ...
     # ... given how TF queues work.
     images = [tf.identity(image) for image in images]
 
-    if cdna + dna != 1:
-      raise ValueError("More than one, or no network option specified.")
-
-    img_height, img_width, color_channels = self.hparams.problem.frame_shape
     batch_size = common_layers.shape_list(images[0])[0]
+    context_frames = self.hparams.video_num_input_frames
 
     # Predicted images and rewards.
     gen_rewards, gen_images = [], []
 
-    if k == -1:
-      feedself = True
-    else:
-      # Scheduled sampling:
-      # Calculate number of ground-truth frames to pass in.
-      iter_num = tf.train.get_global_step()
-      # TODO(mbz): what should it be if it's undefined?
-      if iter_num is None:
-        iter_num = _LARGE_STEP_NUMBER
-      num_ground_truth = tf.to_int32(
-          tf.round(
-              tf.to_float(batch_size) *
-              (k / (k + tf.exp(tf.to_float(iter_num) / tf.to_float(k))))))
-      feedself = False
-
-    # LSTM state sizes and states.
-    lstm_size = np.array([32, 32, 64, 64, 128, 64, 32], dtype=np.int32)
+    # LSTM states.
     lstm_state = [None] * 7
 
     # Latent tower
@@ -335,165 +447,31 @@ def construct_model(self,
       latent_tower_outputs = self.construct_latent_tower(images)
       latent_mean, latent_std, samples = latent_tower_outputs
 
-    # Main tower
-    layer_norm = tf.contrib.layers.layer_norm
-    lstm_func = self.conv_lstm_2d
-
+    pred_image, pred_reward = None, None
     for timestep, image, action, reward in zip(
         range(len(images)-1), images[:-1], actions[:-1], rewards[:-1]):
-
-      done_warm_start = len(gen_images) > context_frames - 1
-      with tf.variable_scope("main", reuse=tf.AUTO_REUSE):
-        if feedself and done_warm_start:
-          # Feed in generated image.
-          prev_image = gen_images[-1]
-          prev_reward = gen_rewards[-1]
-        elif done_warm_start:
-          # Scheduled sampling
-          prev_image = self.scheduled_sample(
-              image, gen_images[-1], batch_size, num_ground_truth)
-          prev_reward = self.scheduled_sample(
-              reward, gen_rewards[-1], batch_size, num_ground_truth)
-        else:
-          # Always feed in ground_truth
-          prev_image = image
-          prev_reward = reward
-
-        prev_image = common_layers.make_even_size(prev_image)
-        enc0 = slim.layers.conv2d(
-            prev_image,
-            32, [5, 5],
-            stride=2,
-            scope="scale1_conv1",
-            normalizer_fn=layer_norm,
-            normalizer_params={"scope": "layer_norm1"})
-
-        hidden1, lstm_state[0] = lstm_func(
-            enc0, lstm_state[0], lstm_size[0], scope="state1")
-        hidden1 = layer_norm(hidden1, scope="layer_norm2")
-        hidden2, lstm_state[1] = lstm_func(
-            hidden1, lstm_state[1], lstm_size[1], scope="state2")
-        hidden2 = layer_norm(hidden2, scope="layer_norm3")
-        hidden2 = common_layers.make_even_size(hidden2)
-        enc1 = slim.layers.conv2d(
-            hidden2, hidden2.get_shape()[3], [3, 3], stride=2, scope="conv2")
-
-        hidden3, lstm_state[2] = lstm_func(
-            enc1, lstm_state[2], lstm_size[2], scope="state3")
-        hidden3 = layer_norm(hidden3, scope="layer_norm4")
-        hidden4, lstm_state[3] = lstm_func(
-            hidden3, lstm_state[3], lstm_size[3], scope="state4")
-        hidden4 = layer_norm(hidden4, scope="layer_norm5")
-        hidden4 = common_layers.make_even_size(hidden4)
-        enc2 = slim.layers.conv2d(
-            hidden4, hidden4.get_shape()[3], [3, 3], stride=2, scope="conv3")
-
-        # Pass in reward and action.
-        emb_action = self.encode_to_shape(action, enc2.get_shape())
-        emb_reward = self.encode_to_shape(prev_reward, enc2.get_shape())
-        enc2 = tf.concat(axis=3, values=[enc2, emb_action, emb_reward])
-
-        # Setup latent
-        if self.hparams.stochastic_model:
-          latent = samples
-          if self.hparams.multi_latent:
-            latent = samples[timestep]
-          if self.hparams.mode == tf.estimator.ModeKeys.TRAIN:
-            # TODO(mbz): put 1st stage of training back in if necessary
-            latent = latent_mean + tf.exp(latent_std / 2.0) * latent
-          with tf.control_dependencies([latent]):
-            enc2 = tf.concat([enc2, latent], 3)
-
-        enc3 = slim.layers.conv2d(
-            enc2, hidden4.get_shape()[3], [1, 1], stride=1, scope="conv4")
-
-        hidden5, lstm_state[4] = lstm_func(
-            enc3, lstm_state[4], lstm_size[4], scope="state5")  # last 8x8
-        hidden5 = layer_norm(hidden5, scope="layer_norm6")
-        enc4 = slim.layers.conv2d_transpose(
-            hidden5, hidden5.get_shape()[3], 3, stride=2, scope="convt1")
-
-        enc1_shape = common_layers.shape_list(enc1)
-        enc4 = enc4[:, :enc1_shape[1], :enc1_shape[2], :]  # Cut to shape.
-        hidden6, lstm_state[5] = lstm_func(
-            enc4, lstm_state[5], lstm_size[5], scope="state6")  # 16x16
-        hidden6 = layer_norm(hidden6, scope="layer_norm7")
-        # Skip connection.
-        hidden6 = tf.concat(axis=3, values=[hidden6, enc1])  # both 16x16
-
-        enc5 = slim.layers.conv2d_transpose(
-            hidden6, hidden6.get_shape()[3], 3, stride=2, scope="convt2")
-        enc0_shape = common_layers.shape_list(enc0)
-        enc5 = enc5[:, :enc0_shape[1], :enc0_shape[2], :]  # Cut to shape.
-        hidden7, lstm_state[6] = lstm_func(
-            enc5, lstm_state[6], lstm_size[6], scope="state7")  # 32x32
-        hidden7 = layer_norm(hidden7, scope="layer_norm8")
-
-        # Skip connection.
-        hidden7 = tf.concat(axis=3, values=[hidden7, enc0])  # both 32x32
-
-        enc6 = slim.layers.conv2d_transpose(
-            hidden7,
-            hidden7.get_shape()[3],
-            3,
-            stride=2,
-            scope="convt3",
-            activation_fn=None,
-            normalizer_fn=layer_norm,
-            normalizer_params={"scope": "layer_norm9"})
-
-        if dna:
-          # Using largest hidden state for predicting untied conv kernels.
-          enc7 = slim.layers.conv2d_transpose(
-              enc6,
-              self.hparams.dna_kernel_size**2,
-              1,
-              stride=1,
-              scope="convt4",
-              activation_fn=None)
-        else:
-          # Using largest hidden state for predicting a new image layer.
-          enc7 = slim.layers.conv2d_transpose(
-              enc6,
-              color_channels,
-              1,
-              stride=1,
-              scope="convt4",
-              activation_fn=None)
-          # This allows the network to also generate one image from scratch,
-          # which is useful when regions of the image become unoccluded.
-          transformed = [tf.nn.sigmoid(enc7)]
-
-        if cdna:
-          # cdna_input = tf.reshape(hidden5, [int(batch_size), -1])
-          cdna_input = tf.contrib.layers.flatten(hidden5)
-          transformed += self.cdna_transformation(
-              prev_image, cdna_input, num_masks, int(color_channels))
-        elif dna:
-          # Only one mask is supported (more should be unnecessary).
-          if num_masks != 1:
-            raise ValueError("Only one mask is supported for DNA model.")
-          transformed = [self.dna_transformation(prev_image, enc7)]
-
-        masks = slim.layers.conv2d_transpose(
-            enc6, num_masks + 1, 1,
-            stride=1, scope="convt7", activation_fn=None)
-        masks = tf.reshape(
-            tf.nn.softmax(tf.reshape(masks, [-1, num_masks + 1])),
-            [batch_size,
-             int(img_height),
-             int(img_width), num_masks + 1])
-        mask_list = tf.split(
-            axis=3, num_or_size_splits=num_masks + 1, value=masks)
-        output = mask_list[0] * prev_image
-        for layer, mask in zip(transformed, mask_list[1:]):
-          output += layer * mask
-        gen_images.append(output)
-
-        p_reward = self.reward_prediction(hidden5)
-        p_reward = self.decode_to_shape(p_reward, reward.shape)
-
-        gen_rewards.append(p_reward)
+      # Scheduled Sampling
+      done_warm_start = timestep > context_frames - 1
+      groundtruth_items = [image, reward]
+      generated_items = [pred_image, pred_reward]
+      input_image, input_reward = self.get_scheduled_sample_inputs(
+          done_warm_start, k, groundtruth_items, generated_items, batch_size)
+
+      # Setup latent
+      latent = None
+      if self.hparams.stochastic_model:
+        latent = samples
+        if self.hparams.multi_latent:
+          latent = samples[timestep]
+        if self.hparams.mode == tf.estimator.ModeKeys.TRAIN:
+          # TODO(mbz): put 1st stage of training back in if necessary
+          latent = latent_mean + tf.exp(latent_std / 2.0) * latent
+
+      # Prediction
+      pred_image, pred_reward, lstm_state = self.construct_predictive_tower(
+          input_image, input_reward, action, lstm_state, latent)
+      gen_images.append(pred_image)
+      gen_rewards.append(pred_reward)
 
     return gen_images, gen_rewards, latent_mean, latent_std
 
@@ -605,15 +583,49 @@ def scheduled_sample(self,
       New batch with num_ground_truth sampled from ground_truth_x and the rest
       from generated_x.
     """
-    idx = tf.random_shuffle(tf.range(int(batch_size)))
+    idx = tf.random_shuffle(tf.range(batch_size))
     ground_truth_idx = tf.gather(idx, tf.range(num_ground_truth))
-    generated_idx = tf.gather(idx, tf.range(num_ground_truth, int(batch_size)))
+    generated_idx = tf.gather(idx, tf.range(num_ground_truth, batch_size))
 
     ground_truth_examps = tf.gather(ground_truth_x, ground_truth_idx)
     generated_examps = tf.gather(generated_x, generated_idx)
     return tf.dynamic_stitch([ground_truth_idx, generated_idx],
                              [ground_truth_examps, generated_examps])
 
+  def get_scheduled_sample_inputs(
+      self, done_warm_start, k, groundtruth_items, generated_items, batch_size):
+
+    with tf.variable_scope("scheduled_sampling", reuse=tf.AUTO_REUSE):
+      if k < 0:
+        feedself = True
+      else:
+        # Scheduled sampling:
+        # Calculate number of ground-truth frames to pass in.
+        feedself = False
+        iter_num = tf.train.get_global_step()
+        # TODO(mbz): what should it be if it's undefined?
+        if iter_num is None:
+          iter_num = _LARGE_STEP_NUMBER
+        num_ground_truth = tf.to_int32(
+            tf.round(
+                tf.to_float(batch_size) *
+                (k / (k + tf.exp(tf.to_float(iter_num) / tf.to_float(k))))))
+
+      if feedself and done_warm_start:
+        # Feed in generated stuff.
+        output_items = generated_items
+      elif done_warm_start:
+        output_items = []
+        for item_gt, item_gen in zip(groundtruth_items, generated_items):
+          # Scheduled sampling
+          output_items.append(self.scheduled_sample(
+              item_gt, item_gen, batch_size, num_ground_truth))
+      else:
+        # Feed in ground_truth
+        output_items = groundtruth_items
+
+      return output_items
+
   # TODO(mbz): use tf.distributions.kl_divergence instead.
   def kl_divergence(self, mu, log_sigma):
     """KL divergence of diagonal gaussian N(mu,exp(log_sigma)) and N(0,1).
@@ -660,18 +672,13 @@ def body(self, features):
     all_rewards = input_rewards + target_rewards
     all_frames = input_frames + target_frames
 
-    tf.summary.image("full_video", tf.concat(all_frames, axis=1))
-
     is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
     gen_images, gen_rewards, latent_mean, latent_std = self.construct_model(
         images=all_frames,
         actions=all_actions,
         rewards=all_rewards,
-        k=900.0 if is_training else -1.0,
-        num_masks=10,
-        cdna=True,
-        dna=False,
-        context_frames=hparams.video_num_input_frames)
+        k=900.0 if is_training else -1.0
+    )
 
     step_num = tf.train.get_global_step()
     # TODO(mbz): what should it be if it's undefined?
@@ -697,6 +704,10 @@ def body(self, features):
         gen_rewards[hparams.video_num_input_frames-1:], axis=1)
     reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove undeeded dimension.
 
+    frames_gt = tf.concat(all_frames[hparams.video_num_input_frames:], axis=1)
+    frames_pd = tf.concat(predictions, axis=1)
+    tf.summary.image("full_video", tf.concat([frames_gt, frames_pd], axis=2))
+
     return_targets = predictions
     if "target_reward" in features:
       return_targets = {"targets": predictions, "target_reward": reward_pred}
@@ -733,16 +744,18 @@ def next_frame():
 def next_frame_stochastic():
   """SV2P model."""
   hparams = next_frame()
-  hparams.video_num_input_frames = 2
-  hparams.video_num_target_frames = 1
+  hparams.video_num_input_frames = 1
+  hparams.video_num_target_frames = 3
   hparams.batch_size = 8
   hparams.target_modality = "video:l2raw"
   hparams.input_modalities = "inputs:video:l2raw"
   hparams.video_modality_loss_cutoff = 0.0
   hparams.add_hparam("stochastic_model", True)
+  hparams.add_hparam("model_options", "CDNA")
+  hparams.add_hparam("num_masks", 10)
   hparams.add_hparam("latent_channels", 1)
   hparams.add_hparam("latent_std_min", -5.0)
-  hparams.add_hparam("num_iterations_2nd_stage", 50000)
+  hparams.add_hparam("num_iterations_2nd_stage", 10000)
   hparams.add_hparam("latent_loss_multiplier", 1e-4)
   hparams.add_hparam("multi_latent", False)
   hparams.add_hparam("relu_shift", 1e-12)

From 91adccf6e36dc9fdebb7fe988554f1bfc20720c9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 19 Jun 2018 12:49:00 -0700
Subject: [PATCH 0148/2720] Internal change

PiperOrigin-RevId: 201224558
---
 tensor2tensor/models/transformer_test.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 149f6c373..f97602a40 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -242,7 +242,9 @@ def testTransformerWithEncoderDecoderAttentionLoss(self):
 
   def testGreedySlowTPUVsNonTPU(self):
     # Only works with TF 1.8+
-    major, minor, _ = [int(el) for el in tf.__version__.split(".")]
+    # Version string can take the following form: "1.9.0-rc0"
+    major_str, minor_str, unused_rest = tf.__version__.split(".", 3)
+    major, minor = int(major_str), int(minor_str)
     if major < 1 or (major == 1 and minor < 8):
       return
     model, features = get_model(transformer.transformer_small())

From 724a854b2a8e53a2b1158c4f37c6783f96ad17a7 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 19 Jun 2018 12:49:20 -0700
Subject: [PATCH 0149/2720] Internal change

PiperOrigin-RevId: 201224601
---
 tensor2tensor/data_generators/program_search.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/program_search.py b/tensor2tensor/data_generators/program_search.py
index d1846dadc..775fa7c39 100644
--- a/tensor2tensor/data_generators/program_search.py
+++ b/tensor2tensor/data_generators/program_search.py
@@ -43,9 +43,9 @@ class ProgramSearchAlgolisp(text_problems.Text2TextProblem):
   DROPBOX = "https://www.dropbox.com"
   DATA_URLS = {
       problem.DatasetSplit.TRAIN: (
-          DROPBOX + "/s/wep81pcrar5fttl/metaset3.train.jsonl.gz?dl=1"),
+          DROPBOX + "/s/qhun6kml9yb2ui9/metaset3.train.jsonl.gz?dl=1"),
       problem.DatasetSplit.EVAL: (
-          DROPBOX + "/s/h3mn0abeiqy6foz/metaset3.dev.jsonl.gz?dl=1"),
+          DROPBOX + "/s/aajkw83j2ps8bzx/metaset3.dev.jsonl.gz?dl=1"),
       problem.DatasetSplit.TEST: (
           DROPBOX + "/s/f1x9ybkjpf371cp/metaset3.test.jsonl.gz?dl=1"),
   }

From 4d2f04d2bf49cdf7506f0bb870e75693c0b42303 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 19 Jun 2018 13:11:19 -0700
Subject: [PATCH 0150/2720] Align training with the original code so we can
 reproduce the results.

PiperOrigin-RevId: 201228083
---
 tensor2tensor/layers/modalities.py          |  8 ++++
 tensor2tensor/models/research/next_frame.py | 47 +++++++++------------
 2 files changed, 27 insertions(+), 28 deletions(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index bcfc9b80d..e425dadf3 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -654,6 +654,14 @@ def top(self, body_output, _):
     common_layers.summarize_video(rgb_frames, "body_output")
     return tf.expand_dims(rgb_frames, axis=-1)
 
+  def loss(self, top_out, targets):
+    prediction = top_out
+    prediction = tf.squeeze(prediction, axis=-1)
+    prediction = common_layers.convert_rgb_to_real(prediction)
+    groundtruth = common_layers.convert_rgb_to_real(targets)
+    loss = tf.losses.mean_squared_error(prediction, groundtruth)
+    return loss, tf.constant(1.0)
+
 
 @registry.register_class_label_modality("default")
 class ClassLabelModality(modality.Modality):
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index c0dfe5221..5e96f2d5c 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -185,8 +185,6 @@ def construct_latent_tower(self, images):
       latent_loss: loss of the latent twoer
       samples: random samples sampled from standard guassian
     """
-    sequence_length = len(images)
-
     with tf.variable_scope("latent"):
       images = tf.concat(images, 3)
 
@@ -206,20 +204,11 @@ def construct_latent_tower(self, images):
       std = slim.conv2d(x, nc, [3, 3], stride=2, scope="latent_std")
       std += self.hparams.latent_std_min
 
-    if self.hparams.multi_latent:
-      # timestep x batch_size x latent_size
-      samples = tf.random_normal(
-          [sequence_length-1] + mean.shape, 0, 1,
-          dtype=tf.float32)
-    else:
-      # batch_size x latent_size
-      samples = tf.random_normal(tf.shape(mean), 0, 1, dtype=tf.float32)
-
-    if self.hparams.mode == tf.estimator.ModeKeys.TRAIN:
-      return mean, std, samples
-    else:
       # No latent tower at inference time, just standard gaussian.
-      return None, None, samples
+      if self.hparams.mode != tf.estimator.ModeKeys.TRAIN:
+        return tf.zeros_like(mean), tf.zeros_like(std)
+
+      return mean, std
 
   def reward_prediction(self, inputs):
     """Builds a reward prediction network."""
@@ -404,6 +393,11 @@ def construct_predictive_tower(
 
       return output, p_reward, lstm_state
 
+  def get_guassian_latent(self, latent_mean, latent_std):
+    latent = tf.random_normal(tf.shape(latent_mean), 0, 1, dtype=tf.float32)
+    latent = latent_mean + tf.exp(latent_std / 2.0) * latent
+    return latent
+
   def construct_model(self,
                       images,
                       actions,
@@ -444,10 +438,9 @@ def construct_model(self,
 
     # Latent tower
     if self.hparams.stochastic_model:
-      latent_tower_outputs = self.construct_latent_tower(images)
-      latent_mean, latent_std, samples = latent_tower_outputs
+      latent_mean, latent_std = self.construct_latent_tower(images)
 
-    pred_image, pred_reward = None, None
+    pred_image, pred_reward, latent = None, None, None
     for timestep, image, action, reward in zip(
         range(len(images)-1), images[:-1], actions[:-1], rewards[:-1]):
       # Scheduled Sampling
@@ -457,15 +450,10 @@ def construct_model(self,
       input_image, input_reward = self.get_scheduled_sample_inputs(
           done_warm_start, k, groundtruth_items, generated_items, batch_size)
 
-      # Setup latent
-      latent = None
+      # Latent
       if self.hparams.stochastic_model:
-        latent = samples
-        if self.hparams.multi_latent:
-          latent = samples[timestep]
-        if self.hparams.mode == tf.estimator.ModeKeys.TRAIN:
-          # TODO(mbz): put 1st stage of training back in if necessary
-          latent = latent_mean + tf.exp(latent_std / 2.0) * latent
+        if timestep == 0 or self.hparams.multi_latent:
+          latent = self.get_guassian_latent(latent_mean, latent_std)
 
       # Prediction
       pred_image, pred_reward, lstm_state = self.construct_predictive_tower(
@@ -744,9 +732,12 @@ def next_frame():
 def next_frame_stochastic():
   """SV2P model."""
   hparams = next_frame()
+  hparams.optimizer = "TrueAdam"
+  hparams.learning_rate_schedule = "constant"
+  hparams.learning_rate_constant = 1e-3
   hparams.video_num_input_frames = 1
   hparams.video_num_target_frames = 3
-  hparams.batch_size = 8
+  hparams.batch_size = 16
   hparams.target_modality = "video:l2raw"
   hparams.input_modalities = "inputs:video:l2raw"
   hparams.video_modality_loss_cutoff = 0.0
@@ -756,7 +747,7 @@ def next_frame_stochastic():
   hparams.add_hparam("latent_channels", 1)
   hparams.add_hparam("latent_std_min", -5.0)
   hparams.add_hparam("num_iterations_2nd_stage", 10000)
-  hparams.add_hparam("latent_loss_multiplier", 1e-4)
+  hparams.add_hparam("latent_loss_multiplier", 1e-3)
   hparams.add_hparam("multi_latent", False)
   hparams.add_hparam("relu_shift", 1e-12)
   hparams.add_hparam("dna_kernel_size", 5)

From 8640c3f4a2329130c0e158f1a992c8d65059f4ef Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 19 Jun 2018 13:23:22 -0700
Subject: [PATCH 0151/2720] Adding 2Frames posterior model.

PiperOrigin-RevId: 201229975
---
 tensor2tensor/models/research/next_frame.py | 70 +++++++++++++++++----
 1 file changed, 59 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 5e96f2d5c..d897e54e2 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -186,7 +186,9 @@ def construct_latent_tower(self, images):
       samples: random samples sampled from standard guassian
     """
     with tf.variable_scope("latent"):
-      images = tf.concat(images, 3)
+      # this allows more predicted frames at inference time
+      latent_images = images[:self.hparams.latent_num_frames]
+      images = tf.concat(latent_images, 3)
 
       x = images
       x = common_layers.make_even_size(x)
@@ -422,11 +424,6 @@ def construct_model(self,
     Raises:
       ValueError: if more than 1 mask specified for DNA model.
     """
-    # Each image is being used twice, in latent tower and main tower.
-    # This is to make sure we are using the *same* image for both, ...
-    # ... given how TF queues work.
-    images = [tf.identity(image) for image in images]
-
     batch_size = common_layers.shape_list(images[0])[0]
     context_frames = self.hparams.video_num_input_frames
 
@@ -461,7 +458,7 @@ def construct_model(self,
       gen_images.append(pred_image)
       gen_rewards.append(pred_reward)
 
-    return gen_images, gen_rewards, latent_mean, latent_std
+    return gen_images, gen_rewards, [latent_mean], [latent_std]
 
   def cdna_transformation(self,
                           prev_image,
@@ -660,8 +657,14 @@ def body(self, features):
     all_rewards = input_rewards + target_rewards
     all_frames = input_frames + target_frames
 
+    # Each image is being used twice, in latent tower and main tower.
+    # This is to make sure we are using the *same* image for both, ...
+    # ... given how TF queues work.
+    # NOT sure if this is required at all. Doesn't hurt though! :)
+    all_frames = [tf.identity(frame) for frame in all_frames]
+
     is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
-    gen_images, gen_rewards, latent_mean, latent_std = self.construct_model(
+    gen_images, gen_rewards, latent_means, latent_stds = self.construct_model(
         images=all_frames,
         actions=all_actions,
         rewards=all_rewards,
@@ -678,11 +681,12 @@ def body(self, features):
 
     kl_loss = 0.0
     if is_training:
-      kl_loss = self.kl_divergence(latent_mean, latent_std)
+      for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
+        kl_loss += self.kl_divergence(mean, std)
+        tf.summary.histogram("posterior_mean_%d" % i, mean)
+        tf.summary.histogram("posterior_std_%d" % i, std)
 
       tf.summary.scalar("beta", beta)
-      tf.summary.histogram("posterior_mean", latent_mean)
-      tf.summary.histogram("posterior_std", latent_std)
       tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
 
     extra_loss = beta * kl_loss
@@ -703,6 +707,47 @@ def body(self, features):
     return return_targets, extra_loss
 
 
+@registry.register_model
+class NextFrameStochasticTwoFrames(NextFrameStochastic):
+  """Stochastic next-frame model with 2 frames posterior."""
+
+  def construct_model(self, images, actions, rewards, k=-1):
+    batch_size = common_layers.shape_list(images[0])[0]
+    context_frames = self.hparams.video_num_input_frames
+
+    # Predicted images and rewards.
+    gen_rewards, gen_images, latent_means, latent_stds = [], [], [], []
+
+    # LSTM states.
+    lstm_state = [None] * 7
+
+    pred_image, pred_reward, latent = None, None, None
+    for timestep, image, action, reward in zip(
+        range(len(images)-1), images[:-1], actions[:-1], rewards[:-1]):
+      # Scheduled Sampling
+      done_warm_start = timestep > context_frames - 1
+      groundtruth_items = [image, reward]
+      generated_items = [pred_image, pred_reward]
+      input_image, input_reward = self.get_scheduled_sample_inputs(
+          done_warm_start, k, groundtruth_items, generated_items, batch_size)
+
+      # Latent
+      # TODO(mbz): should we use input_image iunstead of image?
+      latent_images = [image, images[timestep+1]]
+      latent_mean, latent_std = self.construct_latent_tower(latent_images)
+      latent = self.get_guassian_latent(latent_mean, latent_std)
+      latent_means.append(latent_mean)
+      latent_stds.append(latent_std)
+
+      # Prediction
+      pred_image, pred_reward, lstm_state = self.construct_predictive_tower(
+          input_image, input_reward, action, lstm_state, latent)
+      gen_images.append(pred_image)
+      gen_rewards.append(pred_reward)
+
+    return gen_images, gen_rewards, latent_means, latent_stds
+
+
 @registry.register_hparams
 def next_frame():
   """Basic 2-frame conv model."""
@@ -751,6 +796,9 @@ def next_frame_stochastic():
   hparams.add_hparam("multi_latent", False)
   hparams.add_hparam("relu_shift", 1e-12)
   hparams.add_hparam("dna_kernel_size", 5)
+  hparams.add_hparam(
+      "latent_num_frames",  # use all frames by default.
+      hparams.video_num_input_frames + hparams.video_num_target_frames)
   return hparams
 
 

From 422f06b9f3d3a457dd35ab8bb99f9c5617220fb0 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 19 Jun 2018 15:01:13 -0700
Subject: [PATCH 0152/2720] Fixing Shapes dataset target modality so it
 decodes.

PiperOrigin-RevId: 201247103
---
 tensor2tensor/data_generators/video_generated.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 8c5ed756f..60088be97 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -87,7 +87,9 @@ def hparams(self, defaults, unused_model_hparams):
         "inputs": ("video", 256),
         "input_frame_number": ("symbol:identity", 1)
     }
-    p.target_modality = ("video", 256)
+    p.target_modality = {
+        "targets": ("video", 256),
+    }
 
   @staticmethod
   def get_circle(x, y, z, c, s):

From 2cd2f8998fd4aa9f909c1fa24823b96b437a95b8 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 19 Jun 2018 16:49:39 -0700
Subject: [PATCH 0153/2720] Add decode and continuous_decode schedules

PiperOrigin-RevId: 201265297
---
 tensor2tensor/bin/t2t_decoder.py   |  4 +-
 tensor2tensor/bin/t2t_trainer.py   |  2 +-
 tensor2tensor/utils/decoding.py    |  4 ++
 tensor2tensor/utils/trainer_lib.py | 88 +++++++++++++++++-------------
 4 files changed, 57 insertions(+), 41 deletions(-)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 9de17a669..19b05b591 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -72,8 +72,8 @@ def create_hparams():
 
 def create_decode_hparams():
   decode_hp = decoding.decode_hparams(FLAGS.decode_hparams)
-  decode_hp.add_hparam("shards", FLAGS.decode_shards)
-  decode_hp.add_hparam("shard_id", FLAGS.worker_id)
+  decode_hp.shards = FLAGS.decode_shards
+  decode_hp.shard_id = FLAGS.worker_id
   return decode_hp
 
 
diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 949deec39..426a32ed3 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -62,7 +62,7 @@
                      "See TensorFlow config.proto for details.")
 
 # To maintain compatibility with some internal libs, we guard against these flag
-# definitions possibly erring. Apologies for the ugliness.
+# definitions possibly erroring. Apologies for the ugliness.
 try:
   flags.DEFINE_string("master", "", "Address of TensorFlow master.")
   flags.DEFINE_string("output_dir", "", "Base output directory for run.")
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index bbc68fbe0..07f8ad6cc 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -51,6 +51,9 @@ def decode_hparams(overrides=""):
       identity_output=False,
       num_samples=-1,
       delimiter="\n",
+      decode_to_file=None,
+      shards=1,
+      shard_id=0,
       force_decode_length=False)
   hp.parse(overrides)
   return hp
@@ -148,6 +151,7 @@ def decode_from_dataset(estimator,
   predictions = estimator.predict(infer_input_fn)
 
   # Prepare output file writers if decode_to_file passed
+  decode_to_file = decode_to_file or decode_hp.decode_to_file
   if decode_to_file:
     if decode_hp.shards > 1:
       decode_filename = decode_to_file + ("%.2d" % decode_hp.shard_id)
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 34362c9d7..25f210518 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -22,6 +22,7 @@
 import random
 import numpy as np
 
+from tensor2tensor.utils import decoding
 from tensor2tensor.utils import devices
 from tensor2tensor.utils import metrics_hook
 from tensor2tensor.utils import registry
@@ -33,6 +34,24 @@
 from tensorflow.python import debug
 
 
+def next_checkpoint(model_dir, timeout_mins=120):
+  """Yields successive checkpoints from model_dir."""
+  last_ckpt = None
+  while True:
+    last_ckpt = tf.contrib.training.wait_for_new_checkpoint(
+        model_dir,
+        last_ckpt,
+        seconds_to_sleep=60,
+        timeout=60 * timeout_mins)
+
+    if last_ckpt is None:
+      tf.logging.info(
+          "Eval timeout: no new checkpoints within %dm" % timeout_mins)
+      break
+
+    yield last_ckpt
+
+
 def create_session_config(log_device_placement=False,
                           enable_graph_rewriter=False,
                           gpu_mem_fraction=0.95,
@@ -190,7 +209,7 @@ def create_estimator(model_name,
     predict_batch_size = batch_size
     if decode_hparams and decode_hparams.batch_size:
       predict_batch_size = decode_hparams.batch_size
-    return tf.contrib.tpu.TPUEstimator(
+    estimator = tf.contrib.tpu.TPUEstimator(
         model_fn=model_fn,
         model_dir=run_config.model_dir,
         config=run_config,
@@ -198,8 +217,9 @@ def create_estimator(model_name,
         eval_batch_size=batch_size if "eval" in schedule else None,
         predict_batch_size=predict_batch_size)
   else:
-    return tf.estimator.Estimator(
+    estimator = tf.estimator.Estimator(
         model_fn=model_fn, model_dir=run_config.model_dir, config=run_config)
+  return estimator
 
 
 def create_hooks(use_tfdbg=False,
@@ -245,14 +265,15 @@ def create_hooks(use_tfdbg=False,
 class T2TExperiment(object):
   """Custom Experiment class for running distributed experiments."""
 
-  def __init__(self, estimator, hparams, train_spec, eval_spec):
+  def __init__(self, estimator, hparams, train_spec, eval_spec,
+               decode_hparams=None):
     self._train_spec = train_spec
     self._eval_spec = eval_spec
     self._hparams = hparams
+    self._decode_hparams = decode_hparams
     self._estimator = estimator
 
   def continuous_train_and_eval(self):
-
     tf.estimator.train_and_evaluate(self._estimator, self._train_spec,
                                     self._eval_spec)
     return self.evaluate()
@@ -274,45 +295,25 @@ def evaluate(self):
     return self._estimator.evaluate(
         self._eval_spec.input_fn,
         steps=self._eval_spec.steps,
-        hooks=self._eval_spec.hooks)
+        hooks=self._eval_spec.hooks,
+        name="eval")
+
+  def evaluate_on_train_data(self):
+    self._estimator.evaluate(
+        self._train_spec.input_fn,
+        steps=self._eval_spec.steps,
+        hooks=self._eval_spec.hooks,
+        name="eval_train")
 
   def continuous_eval(self):
     """Evaluate until checkpoints stop being produced."""
-    last_ckpt = None
-    while True:
-      # Wait up to half an hour for a new checkpoint
-      last_ckpt = tf.contrib.training.wait_for_new_checkpoint(
-          self._hparams.model_dir,
-          last_ckpt,
-          seconds_to_sleep=60,
-          timeout=60 * 30)
-
-      if last_ckpt is None:
-        raise Exception("Eval timeout: no new checkpoints within 30mins")
-
-      self._estimator.evaluate(
-          self._eval_spec.input_fn,
-          steps=self._eval_spec.steps,
-          hooks=self._eval_spec.hooks)
+    for _ in next_checkpoint(self._hparams.model_dir):
+      self.evaluate()
 
   def continuous_eval_on_train_data(self):
     """Evaluate on train data until checkpoints stop being produced."""
-    last_ckpt = None
-    while True:
-      # Wait up to half an hour for a new checkpoint
-      last_ckpt = tf.contrib.training.wait_for_new_checkpoint(
-          self._hparams.model_dir,
-          last_ckpt,
-          seconds_to_sleep=60,
-          timeout=60 * 30)
-
-      if last_ckpt is None:
-        raise Exception("Eval timeout: no new checkpoints within 30mins")
-
-      self._estimator.evaluate(
-          self._train_spec.input_fn,
-          steps=self._eval_spec.steps,
-          hooks=self._eval_spec.hooks)
+    for _ in next_checkpoint(self._hparams.model_dir):
+      self.evaluate_on_train_data()
 
   def test(self):
     """Perform 1 step of train and 2 step of eval."""
@@ -349,6 +350,16 @@ def run_std_server(self):
     server.start()
     server.join()
 
+  def decode(self):
+    """Decodes from dataset."""
+    decoding.decode_from_dataset(self._estimator, self._hparams.problem.name,
+                                 self._hparams, self._decode_hparams)
+
+  def continuous_decode(self):
+    """Decode from dataset on new checkpoint."""
+    for _ in next_checkpoint(self._hparams.model_dir):
+      self.decode()
+
 
 def create_experiment(run_config,
                       hparams,
@@ -454,7 +465,8 @@ def create_experiment(run_config,
         start_delay_secs=0 if hparams.schedule == "evaluate" else 120,
         throttle_secs=eval_throttle_seconds)
 
-  return T2TExperiment(estimator, hparams, train_spec, eval_spec)
+  return T2TExperiment(estimator, hparams, train_spec, eval_spec,
+                       decode_hparams)
 
 
 def create_experiment_fn(*args, **kwargs):

From c3f236524314dda78a16268a68db98bf38ffe2a2 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 19 Jun 2018 16:53:06 -0700
Subject: [PATCH 0154/2720] Make GAN work by using loss inspired by sliced
 WGAN.

PiperOrigin-RevId: 201265831
---
 tensor2tensor/layers/common_layers.py |  52 ++++++++++
 tensor2tensor/layers/modalities.py    |   5 +-
 tensor2tensor/models/vanilla_gan.py   | 136 ++++++++++++--------------
 3 files changed, 118 insertions(+), 75 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index fd4c50c92..d16429768 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3332,3 +3332,55 @@ def make_even_size(x):
   x, _ = pad_to_same_length(x, x, final_length_divisible_by=2, axis=2)
   x.set_shape(new_shape)
   return x
+
+
+def sliced_gan_loss(input1, input2, discriminator, num_vecs,
+                    do_random_vecs=True, do_tanh=True):
+  """Loss inspired by the sliced WGAN paper: https://arxiv.org/abs/1804.01947.
+
+  Puts input1 and input2 through the provided discriminator to get logits.
+  Then, computes num_vecs random projections of the logits, sorts them on
+  the batch dimension and returns the L2 loss between the sorted vectors.
+  See the above-mentioned paper for the reasoning behind it.
+
+  Args:
+    input1: first discriminator inputs.
+    input2: second discriminator inputs.
+    discriminator: inputs -> logits function.
+    num_vecs: how many random vectors to use for projections.
+    do_random_vecs: whether to use random vectors or just tanh of the logits.
+    do_tanh: if true (default) we'll also just use tanh of the logits.
+
+  Returns:
+    The generator loss, i.e., the sliced approximation of the distance between
+    the projected distributions (warning: discriminator should maximize it).
+  """
+  with tf.variable_scope("sliced_gan"):
+    with tf.variable_scope("discriminator"):
+      logits1 = discriminator(input1)
+    with tf.variable_scope("discriminator", reuse=True):
+      logits2 = discriminator(input2)
+
+    if do_random_vecs:
+      random_vecs = tf.nn.l2_normalize(
+          tf.random_uniform([shape_list(logits1)[-1], num_vecs]), axis=0)
+
+    def get_sorted_projections(x):
+      """Make projections of x and sort them on the batch dimension."""
+      x = tf.reshape(x, [-1, shape_list(x)[-1]])
+      batch_size = shape_list(x)[0]
+      if do_random_vecs and do_tanh:
+        n = tf.nn.l2_normalize(x, axis=1)
+        proj = tf.concat([tf.matmul(n, random_vecs), tf.tanh(x)], axis=1)
+      elif do_random_vecs:
+        n = tf.nn.l2_normalize(x, axis=1)
+        proj = tf.matmul(n, random_vecs)
+      else:
+        proj = tf.tanh(x)
+      proj = tf.transpose(proj, [1, 0])  # [num_vecs, batch] after this.
+      values, _ = tf.nn.top_k(proj, k=batch_size, sorted=True)
+      return values
+
+    proj1 = get_sorted_projections(logits1)
+    proj2 = get_sorted_projections(logits2)
+    return tf.reduce_mean(tf.square(proj1 - proj2))
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index e425dadf3..11fa20abd 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -213,10 +213,9 @@ class ImageModality(modality.Modality):
 
   def bottom(self, x):
     with tf.variable_scope(self.name):
-      x = tf.to_float(x)
       if not tf.contrib.eager.in_eager_mode():
-        tf.summary.image("inputs", x, max_outputs=2)
-      return x
+        tf.summary.image("inputs", tf.cast(x, tf.uint8), max_outputs=2)
+      return tf.to_float(x)
 
   def targets_bottom(self, x):
     inputs = x
diff --git a/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index dcf32586f..c55c403d1 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -20,7 +20,9 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-from tensor2tensor.layers import common_hparams
+
+from tensor2tensor.layers import common_layers
+from tensor2tensor.models import basic
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -45,6 +47,10 @@ def deconv2d(
     return tf.reshape(tf.nn.bias_add(deconv, biases), deconv.get_shape())
 
 
+def reverse_gradient(x):
+  return -x + tf.stop_gradient(2 * x)
+
+
 class AbstractGAN(t2t_model.T2TModel):
   """Base class for all GANs."""
 
@@ -59,11 +65,11 @@ def discriminator(self, x, is_training, reuse=False):
     Returns:
       out_logit: the output logits (before sigmoid).
     """
-    hparams = self._hparams
+    hparams = self.hparams
     with tf.variable_scope(
         "discriminator", reuse=reuse,
         initializer=tf.random_normal_initializer(stddev=0.02)):
-      batch_size = hparams.batch_size
+      batch_size, height, width = common_layers.shape_list(x)[:3]
       # Mapping x from [bs, h, w, c] to [bs, 1]
       net = tf.layers.conv2d(x, 64, (4, 4), strides=(2, 2),
                              padding="SAME", name="d_conv1")
@@ -76,24 +82,22 @@ def discriminator(self, x, is_training, reuse=False):
         net = tf.layers.batch_normalization(net, training=is_training,
                                             momentum=0.999, name="d_bn2")
       net = lrelu(net)
-      size = hparams.height * hparams.width
+      size = height * width
       net = tf.reshape(net, [batch_size, size * 8])  # [bs, h * w * 8]
       net = tf.layers.dense(net, 1024, name="d_fc3")  # [bs, 1024]
       if hparams.discriminator_batchnorm:
         net = tf.layers.batch_normalization(net, training=is_training,
                                             momentum=0.999, name="d_bn3")
       net = lrelu(net)
-      out_logit = tf.layers.dense(net, 1, name="d_fc4")  # [bs, 1]
-      return out_logit
+      return net
 
-  def generator(self, z, is_training, reuse=False):
+  def generator(self, z, is_training, out_shape):
     """Generator outputting image in [0, 1]."""
-    hparams = self._hparams
-    height = hparams.height
-    width = hparams.width
+    hparams = self.hparams
+    height, width, c_dim = out_shape
     batch_size = hparams.batch_size
     with tf.variable_scope(
-        "generator", reuse=reuse,
+        "generator",
         initializer=tf.random_normal_initializer(stddev=0.02)):
       net = tf.layers.dense(z, 1024, name="g_fc1")
       net = tf.layers.batch_normalization(net, training=is_training,
@@ -110,10 +114,14 @@ def generator(self, z, is_training, reuse=False):
       net = tf.layers.batch_normalization(net, training=is_training,
                                           momentum=0.999, name="g_bn3")
       net = lrelu(net)
-      net = deconv2d(net, [batch_size, height, width, hparams.c_dim],
+      net = deconv2d(net, [batch_size, height, width, c_dim],
                      4, 4, 2, 2, name="g_dc4")
       out = tf.nn.sigmoid(net)
-      return out
+      return common_layers.convert_real_to_rgb(out)
+
+  def losses(self, inputs, generated):
+    """Return the losses dictionary."""
+    raise NotImplementedError
 
   def body(self, features):
     """Body of the model.
@@ -129,51 +137,24 @@ def body(self, features):
     is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
 
     # Input images.
-    inputs = features["inputs"]
+    inputs = tf.to_float(features["targets_raw"])
 
     # Noise vector.
-    z = tf.random_uniform(
-        shape=[self._hparams.batch_size, self._hparams.z_size],
-        minval=-1,
-        maxval=1,
-        name="z")
-
-    # Discriminator output for real images.
-    d_real_logits = self.discriminator(
-        inputs, is_training=is_training, reuse=False)
-
-    # Discriminator output for fake images.
-    g = self.generator(z, is_training=is_training, reuse=False)
-    d_fake_logits_g = self.discriminator(
-        g, is_training=is_training, reuse=True)
-    # Discriminator doesn't backprop to generator.
-    d_fake_logits_d = self.discriminator(
-        tf.stop_gradient(g), is_training=is_training, reuse=True)
-
-    # Loss on real and fake data.
-    d_loss_real = tf.reduce_mean(
-        tf.nn.sigmoid_cross_entropy_with_logits(
-            logits=d_real_logits, labels=tf.ones_like(d_real_logits)))
-    d_loss_fake_g = tf.reduce_mean(
-        tf.nn.sigmoid_cross_entropy_with_logits(
-            logits=d_fake_logits_g, labels=tf.zeros_like(d_fake_logits_g)))
-    d_loss_fake_d = tf.reduce_mean(
-        tf.nn.sigmoid_cross_entropy_with_logits(
-            logits=d_fake_logits_d, labels=tf.zeros_like(d_fake_logits_d)))
-    d_loss = d_loss_real + d_loss_fake_d
-
-    losses = {}  # All losses get added at the end.
-    losses["discriminator"] = d_loss
-    losses["generator"] = - d_loss_fake_g
-    # Include a dummy training loss to skip self.loss.
-    losses["training"] = tf.constant(0., dtype=tf.float32)
-
-    hparams = self._hparams
-    summary_g_image = tf.reshape(g[0, :], [1, hparams.height, hparams.width, 1])
+    z = tf.random_uniform([self.hparams.batch_size,
+                           self.hparams.bottleneck_bits],
+                          minval=-1, maxval=1, name="z")
+
+    # Generator output: fake images.
+    out_shape = common_layers.shape_list(inputs)[1:4]
+    g = self.generator(z, is_training, out_shape)
+
+    losses = self.losses(inputs, g)
+
+    summary_g_image = tf.reshape(
+        g[0, :], [1] + common_layers.shape_list(inputs)[1:])
     tf.summary.image("generated", summary_g_image, max_outputs=1)
 
-    if is_training:
-      # Returns an dummy output and the losses dictionary.
+    if is_training:  # Returns an dummy output and the losses dictionary.
       return tf.zeros_like(inputs), losses
     return tf.reshape(g, tf.shape(inputs)), losses
 
@@ -183,32 +164,43 @@ def top(self, body_output, features):
 
 
 @registry.register_model
-class VanillaGan(AbstractGAN):
-  """Simple GAN for demonstration."""
+class SlicedGan(AbstractGAN):
+  """Sliced GAN for demonstration."""
+
+  def losses(self, inputs, g):
+    """Losses in the sliced case."""
+    is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
+    def discriminate(x):
+      return self.discriminator(x, is_training=is_training, reuse=False)
+    generator_loss = common_layers.sliced_gan_loss(
+        inputs, reverse_gradient(g), discriminate, self.hparams.num_sliced_vecs)
+    return {"training": - generator_loss}
 
   def infer(self, *args, **kwargs):  # pylint: disable=arguments-differ
     del args, kwargs
-    with tf.variable_scope("body/vanilla_gan", reuse=tf.AUTO_REUSE):
-      z = tf.random_uniform(
-          shape=[1, self._hparams.random_sample_size],
-          minval=-1,
-          maxval=1,
-          name="z")
 
-      g_sample = self.generator(z, self._hparams)
+    try:
+      num_channels = self.hparams.problem.num_channels
+    except AttributeError:
+      num_channels = 1
+
+    with tf.variable_scope("body/vanilla_gan", reuse=tf.AUTO_REUSE):
+      hparams = self.hparams
+      z = tf.random_uniform([hparams.batch_size, hparams.bottleneck_bits],
+                            minval=-1, maxval=1, name="z")
+      out_shape = (hparams.sample_height, hparams.sample_width, num_channels)
+      g_sample = self.generator(z, False, out_shape)
       return g_sample
 
 
 @registry.register_hparams
-def vanilla_gan():
+def sliced_gan():
   """Basic parameters for a vanilla_gan."""
-  hparams = common_hparams.basic_params1()
-  hparams.label_smoothing = 0.0
+  hparams = basic.basic_autoencoder()
   hparams.hidden_size = 128
-  hparams.batch_size = 64
-  hparams.add_hparam("z_size", 64)
-  hparams.add_hparam("c_dim", 1)
-  hparams.add_hparam("height", 28)
-  hparams.add_hparam("width", 28)
-  hparams.add_hparam("discriminator_batchnorm", int(True))
+  hparams.batch_size = 128
+  hparams.weight_decay = 1e-6
+  hparams.bottleneck_bits = 128
+  hparams.add_hparam("discriminator_batchnorm", True)
+  hparams.add_hparam("num_sliced_vecs", 4096)
   return hparams

From ebc705261028796b28ff61d879bf1e5dcb707e21 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Tue, 19 Jun 2018 17:11:38 -0700
Subject: [PATCH 0155/2720] Fix bug with updating latents dense

PiperOrigin-RevId: 201268906
---
 tensor2tensor/layers/latent_layers.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 0ea5b30f4..10a7fc0af 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -422,6 +422,7 @@ def bottleneck_layer(targets_c, hparams):
       targets_c, hparams)
   latents_dense = discretization.parametrized_unbottleneck(
       latents_discrete_hot, hparams.hidden_size, hparams)
+  latents_dense = targets_c + tf.stop_gradient(latents_dense - targets_c)
   latents_discrete = tf.argmax(latents_discrete_hot, axis=-1)
 
   if DO_SUMMARIES:

From 639b2e186ab1d734e26f651f4fb86368d25e0cd3 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 19 Jun 2018 17:25:21 -0700
Subject: [PATCH 0156/2720] Adding Emily's model from the following paper:

PiperOrigin-RevId: 201270740
---
 tensor2tensor/layers/common_layers.py       |   6 +
 tensor2tensor/models/research/next_frame.py | 305 +++++++++++++++++++-
 2 files changed, 299 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index d16429768..b330998fb 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3384,3 +3384,9 @@ def get_sorted_projections(x):
     proj1 = get_sorted_projections(logits1)
     proj2 = get_sorted_projections(logits2)
     return tf.reduce_mean(tf.square(proj1 - proj2))
+
+
+def upscale(inputs, f, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR):
+  """Upscaling the image by a factor of f."""
+  height, width = shape_list(inputs)[1:3]
+  return tf.image.resize_images(inputs, (height * f, width * f), method)
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index d897e54e2..9f6093428 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -164,7 +164,11 @@ def logits_to_samples(logits):
 
 @registry.register_model
 class NextFrameStochastic(NextFrameBasic):
-  """Stochastic next-frame model."""
+  """ SV2P: Stochastic Variational Video Prediction.
+
+  based on the following papaer:
+  https://arxiv.org/abs/1710.11252
+  """
 
   def construct_latent_tower(self, images):
     """Builds convolutional latent tower for stochastic model.
@@ -403,8 +407,7 @@ def get_guassian_latent(self, latent_mean, latent_std):
   def construct_model(self,
                       images,
                       actions,
-                      rewards,
-                      k=-1):
+                      rewards):
     """Build convolutional lstm video predictor using CDNA, or DNA.
 
     Args:
@@ -414,7 +417,6 @@ def construct_model(self,
                each action should be in the shape ?x1xZ
       rewards: list of reward tensors
                each reward should be in the shape ?x1xZ
-      k: constant used for scheduled sampling. -1 to feed in own prediction.
     Returns:
       gen_images: predicted future image frames
       gen_rewards: predicted future rewards
@@ -445,7 +447,7 @@ def construct_model(self,
       groundtruth_items = [image, reward]
       generated_items = [pred_image, pred_reward]
       input_image, input_reward = self.get_scheduled_sample_inputs(
-          done_warm_start, k, groundtruth_items, generated_items, batch_size)
+          done_warm_start, groundtruth_items, generated_items, batch_size)
 
       # Latent
       if self.hparams.stochastic_model:
@@ -578,10 +580,10 @@ def scheduled_sample(self,
                              [ground_truth_examps, generated_examps])
 
   def get_scheduled_sample_inputs(
-      self, done_warm_start, k, groundtruth_items, generated_items, batch_size):
+      self, done_warm_start, groundtruth_items, generated_items, batch_size):
 
     with tf.variable_scope("scheduled_sampling", reuse=tf.AUTO_REUSE):
-      if k < 0:
+      if self.hparams.mode != tf.estimator.ModeKeys.TRAIN:
         feedself = True
       else:
         # Scheduled sampling:
@@ -591,6 +593,7 @@ def get_scheduled_sample_inputs(
         # TODO(mbz): what should it be if it's undefined?
         if iter_num is None:
           iter_num = _LARGE_STEP_NUMBER
+        k = self.hparams.scheduled_sampling_k
         num_ground_truth = tf.to_int32(
             tf.round(
                 tf.to_float(batch_size) *
@@ -660,7 +663,7 @@ def body(self, features):
     # Each image is being used twice, in latent tower and main tower.
     # This is to make sure we are using the *same* image for both, ...
     # ... given how TF queues work.
-    # NOT sure if this is required at all. Doesn't hurt though! :)
+    # NOT sure if this is required at all. Doesn"t hurt though! :)
     all_frames = [tf.identity(frame) for frame in all_frames]
 
     is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
@@ -668,11 +671,10 @@ def body(self, features):
         images=all_frames,
         actions=all_actions,
         rewards=all_rewards,
-        k=900.0 if is_training else -1.0
     )
 
     step_num = tf.train.get_global_step()
-    # TODO(mbz): what should it be if it's undefined?
+    # TODO(mbz): what should it be if it"s undefined?
     if step_num is None:
       step_num = _LARGE_STEP_NUMBER
     beta = tf.cond(tf.greater(step_num, self.hparams.num_iterations_2nd_stage),
@@ -711,7 +713,7 @@ def body(self, features):
 class NextFrameStochasticTwoFrames(NextFrameStochastic):
   """Stochastic next-frame model with 2 frames posterior."""
 
-  def construct_model(self, images, actions, rewards, k=-1):
+  def construct_model(self, images, actions, rewards):
     batch_size = common_layers.shape_list(images[0])[0]
     context_frames = self.hparams.video_num_input_frames
 
@@ -729,7 +731,7 @@ def construct_model(self, images, actions, rewards, k=-1):
       groundtruth_items = [image, reward]
       generated_items = [pred_image, pred_reward]
       input_image, input_reward = self.get_scheduled_sample_inputs(
-          done_warm_start, k, groundtruth_items, generated_items, batch_size)
+          done_warm_start, groundtruth_items, generated_items, batch_size)
 
       # Latent
       # TODO(mbz): should we use input_image iunstead of image?
@@ -748,6 +750,270 @@ def construct_model(self, images, actions, rewards, k=-1):
     return gen_images, gen_rewards, latent_means, latent_stds
 
 
+@registry.register_model
+class NextFrameStochasticEmily(NextFrameStochastic):
+  """Model architecture for video prediction model.
+
+     based on following paper:
+     "Stochastic Video Generation with a Learned Prior"
+     https://arxiv.org/pdf/1802.07687.pdf
+     by Emily Denton and Rob Fergus.
+
+     This code is a translation of the original code from PyTorch:
+     https://github.com/edenton/svg
+  """
+
+  def vgg_layer(self,
+                inputs,
+                nout,
+                kernel_size=3,
+                activation=tf.nn.leaky_relu,
+                padding="SAME",
+                scope=""):
+    """A layer of VGG network with batch norm.
+
+    Args:
+      inputs: image tensor
+      nout: number of output channels
+      kernel_size: size of the kernel
+      activation: activation function
+      padding: padding of the image
+      scope: slim scope of the op
+    Returns:
+      net: output of layer
+    """
+    net = slim.conv2d(inputs, nout, kernel_size=kernel_size, padding=padding,
+                      activation_fn=activation, scope=scope+"_conv")
+    net = slim.batch_norm(net, scope=scope+"_bn")
+    net = activation(net)
+    return net
+
+  def basic_lstm(self, inputs, state, num_units, scope=None):
+    input_shape = common_layers.shape_list(inputs)
+    cell = tf.contrib.rnn.BasicLSTMCell(num_units, name=scope)
+    if state is None:
+      state = cell.zero_state(input_shape[0], tf.float32)
+    outputs, new_state = cell(inputs, state)
+    return outputs, new_state
+
+  def encoder(self, inputs, nout):
+    """VGG based image encoder.
+
+    Args:
+      inputs: image tensor with size BSx64x64xC
+      nout: number of output channels
+    Returns:
+      net: encoded image with size BSxNout
+      skips: skip connection after each layer
+    """
+    vgg_layer = self.vgg_layer
+    net01 = inputs
+    # h1
+    net11 = slim.repeat(net01, 2, vgg_layer, 64, scope="h1")
+    net12 = slim.max_pool2d(net11, [2, 2], scope="h1_pool")
+    # h2
+    net21 = slim.repeat(net12, 2, vgg_layer, 128, scope="h2")
+    net22 = slim.max_pool2d(net21, [2, 2], scope="h2_pool")
+    # h3
+    net31 = slim.repeat(net22, 3, vgg_layer, 256, scope="h3")
+    net32 = slim.max_pool2d(net31, [2, 2], scope="h3_pool")
+    # h4
+    net41 = slim.repeat(net32, 3, vgg_layer, 512, scope="h4")
+    net42 = slim.max_pool2d(net41, [2, 2], scope="h4_pool")
+    # h5
+    net51 = slim.repeat(net42, 1, vgg_layer, nout, kernel_size=4,
+                        padding="VALID", activation=tf.tanh, scope="h5")
+    skips = [net11, net21, net31, net41]
+    return net51, skips
+
+  def decoder(self, inputs, skips, nout):
+    """VGG based image decoder.
+
+    Args:
+      inputs: image tensor with size BSxX
+      skips: skip connections from encoder
+      nout: number of output channels
+    Returns:
+      net: decoded image with size BSx64x64xNout
+      skips: skip connection after each layer
+    """
+    vgg_layer = self.vgg_layer
+    net = inputs
+    # d1
+    net = slim.conv2d_transpose(net, 512, kernel_size=4, padding="VALID",
+                                scope="d1_deconv", activation_fn=None)
+    net = slim.batch_norm(net, scope="d1_bn")
+    net = tf.nn.leaky_relu(net)
+    net = common_layers.upscale(net, 2)
+    # d2
+    net = tf.concat([net, skips[3]], axis=3)
+    net = slim.repeat(net, 2, vgg_layer, 512, scope="d2a")
+    net = slim.repeat(net, 1, vgg_layer, 256, scope="d2b")
+    net = common_layers.upscale(net, 2)
+    # d3
+    net = tf.concat([net, skips[2]], axis=3)
+    net = slim.repeat(net, 2, vgg_layer, 256, scope="d3a")
+    net = slim.repeat(net, 1, vgg_layer, 128, scope="d3b")
+    net = common_layers.upscale(net, 2)
+    # d4
+    net = tf.concat([net, skips[1]], axis=3)
+    net = slim.repeat(net, 1, vgg_layer, 128, scope="d4a")
+    net = slim.repeat(net, 1, vgg_layer, 64, scope="d4b")
+    net = common_layers.upscale(net, 2)
+    # d5
+    net = tf.concat([net, skips[0]], axis=3)
+    net = slim.repeat(net, 1, vgg_layer, 64, scope="d5")
+    net = slim.conv2d_transpose(net, nout, kernel_size=3, padding="SAME",
+                                scope="d6_deconv", activation_fn=tf.sigmoid)
+    return net
+
+  def stacked_lstm(self, inputs, states, hidden_size, output_size, nlayers):
+    """Stacked LSTM layers with FC layers as input and output embeddings.
+
+    Args:
+      inputs: input tensor
+      states: a list of internal lstm states for each layer
+      hidden_size: number of lstm units
+      output_size: size of the output
+      nlayers: number of lstm layers
+    Returns:
+      net: output of the network
+      skips: a list of updated lstm states for each layer
+    """
+    net = inputs
+    net = slim.layers.fully_connected(
+        net, hidden_size, activation_fn=None, scope="af1")
+    for i in xrange(nlayers):
+      net, states[i] = self.basic_lstm(
+          net, states[i], hidden_size, scope="alstm%d"%i)
+    net = slim.layers.fully_connected(
+        net, output_size, activation_fn=tf.tanh, scope="af2")
+    return net, states
+
+  def lstm_gaussian(self, inputs, states, hidden_size, output_size, nlayers):
+    """Stacked LSTM layers with FC layer as input and gaussian as output.
+
+    Args:
+      inputs: input tensor
+      states: a list of internal lstm states for each layer
+      hidden_size: number of lstm units
+      output_size: size of the output
+      nlayers: number of lstm layers
+    Returns:
+      mu: mean of the predicted gaussian
+      logvar: log(var) of the predicted gaussian
+      skips: a list of updated lstm states for each layer
+    """
+    net = inputs
+    net = slim.layers.fully_connected(net, hidden_size,
+                                      activation_fn=None, scope="bf1")
+    for i in xrange(nlayers):
+      net, states[i] = self.basic_lstm(
+          net, states[i], hidden_size, scope="blstm%d"%i)
+    mu = slim.layers.fully_connected(
+        net, output_size, activation_fn=None, scope="bf2mu")
+    logvar = slim.layers.fully_connected(
+        net, output_size, activation_fn=None, scope="bf2log")
+    return mu, logvar, states
+
+  def construct_model(self, images, actions, rewards):
+    """Builds the stochastic model.
+
+    The model first encodes all the images (x_t) in the sequence
+    using the encoder. Let"s call the output e_t. Then it predicts the
+    latent state of the next frame using a recurrent posterior network
+    z ~ q(z|e_{0:t}) = N(mu(e_{0:t}), sigma(e_{0:t})).
+    Another recurrent network predicts the embedding of the next frame
+    using the approximated posterior e_{t+1} = p(e_{t+1}|e_{0:t}, z)
+    Finally, the decoder decodes e_{t+1} into x_{t+1}.
+    Skip connections from encoder to decoder help with reconstruction.
+
+    Args:
+      images: tensor of ground truth image sequences
+      actions: NOT used list of action tensors
+      rewards: NOT used list of reward tensors
+
+    Returns:
+      gen_images: generated images
+      fakr_rewards: input rewards as reward prediction!
+      pred_mu: predited means of posterior
+      pred_logvar: predicted log(var) of posterior
+    """
+    # model does not support action conditioned and reward prediction
+    fakr_reward_prediction = rewards
+    del actions, rewards
+
+    z_dim = self.hparams.z_dim
+    g_dim = self.hparams.g_dim
+    rnn_size = self.hparams.rnn_size
+    posterior_rnn_layers = self.hparams.posterior_rnn_layers
+    predictor_rnn_layers = self.hparams.predictor_rnn_layers
+    context_frames = self.hparams.video_num_input_frames
+
+    seq_len = len(images)
+    batch_size, _, _, color_channels = common_layers.shape_list(images[0])
+
+    # LSTM initial sizesstates.
+    predictor_states = [None] * predictor_rnn_layers
+    posterior_states = [None] * posterior_rnn_layers
+
+    tf.logging.info(">>>> Encoding")
+    # Encoding:
+    enc_images, enc_skips = [], []
+    for i, image in enumerate(images):
+      with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
+        enc, skips = self.encoder(image, rnn_size)
+        enc = tf.contrib.layers.flatten(enc)
+        enc_images.append(enc)
+        enc_skips.append(skips)
+
+    tf.logging.info(">>>> Prediction")
+    # Prediction
+    pred_enc, pred_mu, pred_logvar = [], [], []
+    for i in range(1, seq_len):
+      with tf.variable_scope("prediction", reuse=tf.AUTO_REUSE):
+        # current encoding
+        h_current = enc_images[i-1]
+        # target encoding
+        h_target = enc_images[i]
+
+        z = tf.random_normal([batch_size, z_dim], 0, 1, dtype=tf.float32)
+        mu, logvar = tf.zeros_like(z), tf.zeros_like(z)
+
+        # Only use Posterior if it's training time
+        if self.hparams.mode == tf.estimator.ModeKeys.TRAIN:
+          mu, logvar, posterior_states = self.lstm_gaussian(
+              h_target, posterior_states, rnn_size, z_dim, posterior_rnn_layers)
+
+          # The original implementation has a multiplier of 0.5
+          # Removed here for simplicity i.e. replacing var with std
+          z = z * tf.exp(logvar) + mu
+
+        # Predict output encoding
+        h_pred, predictor_states = self.stacked_lstm(
+            tf.concat([h_current, z], axis=1),
+            predictor_states, rnn_size, g_dim, predictor_rnn_layers)
+
+        pred_enc.append(h_pred)
+        pred_mu.append(mu)
+        pred_logvar.append(logvar)
+
+    tf.logging.info(">>>> Decoding")
+    # Decoding
+    gen_images = []
+    for i in range(seq_len-1):
+      with tf.variable_scope("decoding", reuse=tf.AUTO_REUSE):
+        # use skip values of last available frame
+        skip_index = min(context_frames-1, i)
+
+        h_pred = tf.reshape(pred_enc[i], [batch_size, 1, 1, g_dim])
+        x_pred = self.decoder(h_pred, enc_skips[skip_index], color_channels)
+        gen_images.append(x_pred)
+
+    tf.logging.info(">>>> Done")
+    return gen_images, fakr_reward_prediction, pred_mu, pred_logvar
+
+
 @registry.register_hparams
 def next_frame():
   """Basic 2-frame conv model."""
@@ -796,12 +1062,27 @@ def next_frame_stochastic():
   hparams.add_hparam("multi_latent", False)
   hparams.add_hparam("relu_shift", 1e-12)
   hparams.add_hparam("dna_kernel_size", 5)
+  hparams.add_hparam("scheduled_sampling_k", 900.0)
   hparams.add_hparam(
       "latent_num_frames",  # use all frames by default.
       hparams.video_num_input_frames + hparams.video_num_target_frames)
   return hparams
 
 
+@registry.register_hparams
+def next_frame_stochastic_emily():
+  """Emily's model."""
+  hparams = next_frame_stochastic()
+  hparams.latent_loss_multiplier = 1e-4
+  hparams.learning_rate_constant = 0.002
+  hparams.add_hparam("z_dim", 10)
+  hparams.add_hparam("g_dim", 128)
+  hparams.add_hparam("rnn_size", 256)
+  hparams.add_hparam("posterior_rnn_layers", 1)
+  hparams.add_hparam("predictor_rnn_layers", 2)
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_stochastic_cutoff():
   """SV2P model with additional cutoff in L2 loss for environments like pong."""

From 2b3a192b0ca1bd28c082f394f55d946d5ff340c5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 19 Jun 2018 17:58:35 -0700
Subject: [PATCH 0157/2720] Add tests for 2frames and Emily's model

PiperOrigin-RevId: 201274828
---
 tensor2tensor/models/research/next_frame.py   |  6 +--
 .../models/research/next_frame_test.py        | 50 ++++++++-----------
 2 files changed, 23 insertions(+), 33 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 9f6093428..5e0c148bc 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -189,7 +189,7 @@ def construct_latent_tower(self, images):
       latent_loss: loss of the latent twoer
       samples: random samples sampled from standard guassian
     """
-    with tf.variable_scope("latent"):
+    with tf.variable_scope("latent", reuse=tf.AUTO_REUSE):
       # this allows more predicted frames at inference time
       latent_images = images[:self.hparams.latent_num_frames]
       images = tf.concat(latent_images, 3)
@@ -883,7 +883,7 @@ def stacked_lstm(self, inputs, states, hidden_size, output_size, nlayers):
     net = inputs
     net = slim.layers.fully_connected(
         net, hidden_size, activation_fn=None, scope="af1")
-    for i in xrange(nlayers):
+    for i in range(nlayers):
       net, states[i] = self.basic_lstm(
           net, states[i], hidden_size, scope="alstm%d"%i)
     net = slim.layers.fully_connected(
@@ -907,7 +907,7 @@ def lstm_gaussian(self, inputs, states, hidden_size, output_size, nlayers):
     net = inputs
     net = slim.layers.fully_connected(net, hidden_size,
                                       activation_fn=None, scope="bf1")
-    for i in xrange(nlayers):
+    for i in range(nlayers):
       net, states[i] = self.basic_lstm(
           net, states[i], hidden_size, scope="blstm%d"%i)
     mu = slim.layers.fully_connected(
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index 6284fbb0c..373f159a0 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -59,45 +59,35 @@ def TestVideoModel(self,
     expected_shape = y.shape + (expected_last_dim,)
     self.assertEqual(res.shape, expected_shape)
 
-  def TestBasicModel(self, in_frames, out_frames):
-    self.TestVideoModel(
-        in_frames,
-        out_frames,
+  def TestOnVariousInputOutputSizes(self, hparams, model, expected_last_dim):
+    self.TestVideoModel(1, 1, hparams, model, expected_last_dim)
+    self.TestVideoModel(1, 6, hparams, model, expected_last_dim)
+    self.TestVideoModel(4, 1, hparams, model, expected_last_dim)
+    self.TestVideoModel(7, 5, hparams, model, expected_last_dim)
+
+  def testBasic(self):
+    self.TestOnVariousInputOutputSizes(
         next_frame.next_frame(),
         next_frame.NextFrameBasic,
         256)
 
-  def testBasicModelSingleInputFrameSingleOutputFrames(self):
-    self.TestBasicModel(1, 1)
-
-  def testBasicModelSingleInputFrameMultiOutputFrames(self):
-    self.TestBasicModel(1, 6)
-
-  def testBasicModelMultiInputFrameSingleOutputFrames(self):
-    self.TestBasicModel(4, 1)
-
-  def testBasicModelMultiInputFrameMultiOutputFrames(self):
-    self.TestBasicModel(7, 5)
-
-  def TestStochasticModel(self, in_frames, out_frames):
-    self.TestVideoModel(
-        in_frames,
-        out_frames,
+  def testStochastic(self):
+    self.TestOnVariousInputOutputSizes(
         next_frame.next_frame_stochastic(),
         next_frame.NextFrameStochastic,
         1)
 
-  def testStochasticModelSingleInputFrameSingleOutputFrames(self):
-    self.TestStochasticModel(1, 1)
-
-  def testStochasticModelSingleInputFrameMultiOutputFrames(self):
-    self.TestStochasticModel(1, 6)
-
-  def testStochasticModelMultiInputFrameSingleOutputFrames(self):
-    self.TestStochasticModel(4, 1)
+  def testStochasticTwoFrames(self):
+    self.TestOnVariousInputOutputSizes(
+        next_frame.next_frame_stochastic(),
+        next_frame.NextFrameStochasticTwoFrames,
+        1)
 
-  def testStochasticModelMultiInputFrameMultiOutputFrames(self):
-    self.TestStochasticModel(7, 5)
+  def testStochasticEmily(self):
+    self.TestOnVariousInputOutputSizes(
+        next_frame.next_frame_stochastic_emily(),
+        next_frame.NextFrameStochasticEmily,
+        1)
 
 
 if __name__ == "__main__":

From 518c97c74fa625f4f3a8485ef9da146cc99adb6e Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 20 Jun 2018 00:35:21 -0700
Subject: [PATCH 0158/2720] Fix capability to strip reserved tokens while using
 SubwordTextEncoder

PiperOrigin-RevId: 201308529
---
 tensor2tensor/data_generators/text_encoder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 8695aa816..becc7737c 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -489,7 +489,7 @@ def __init__(self, filename=None):
     self.filename = filename
     if filename is not None:
       self._load_from_file(filename)
-    super(SubwordTextEncoder, self).__init__(num_reserved_ids=None)
+    super(SubwordTextEncoder, self).__init__()
 
   def encode(self, s):
     """Converts a native string to a list of subtoken ids.

From 68543a20dc9753963828e30386f567273c205d15 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Wed, 20 Jun 2018 10:36:01 +0200
Subject: [PATCH 0159/2720] bug fixes

---
 tensor2tensor/data_generators/gym_problems.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 44010e3c4..ae96dba82 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -24,6 +24,7 @@
 import numpy as np
 
 # We need gym_utils for the game environments defined there.
+from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.models.research import rl
@@ -59,6 +60,8 @@ def __init__(self, *args, **kwargs):
 
     self.environment_spec = self.get_environment_spec()
     self.eval_phase = False
+    self.sum_of_rewards = 0.0
+    self.dones = 0
 
   def _setup(self):
     collect_hparams = rl.ppo_pong_base()
@@ -95,7 +98,11 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
         memory_index += 1
         observ, reward, done, action = data
         observ = observ.astype(np.uint8) # TODO(piotrmilos). This should be probably done in collect
+
         debug_im = None
+        self.sum_of_rewards += reward
+        self.dones += int(done)
+
 
         ret_dict = {"frame": observ,
                     "image/format": ["png"],

From 90d58416aeda3dc41ade943e1489c6117a4a71f5 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Wed, 20 Jun 2018 10:36:27 +0200
Subject: [PATCH 0160/2720] removing circular imports

---
 tensor2tensor/rl/collect.py                  |   8 +-
 tensor2tensor/rl/envs/batch_env_factory.py   | 272 +++++++++++++++++++
 tensor2tensor/rl/envs/simulated_batch_env.py |   4 +-
 tensor2tensor/rl/envs/utils.py               | 252 -----------------
 4 files changed, 278 insertions(+), 258 deletions(-)
 create mode 100644 tensor2tensor/rl/envs/batch_env_factory.py

diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 5553099cf..9b18057ab 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -20,16 +20,16 @@
 
 import copy
 import tensorflow as tf
-from tensor2tensor.rl.envs import utils
+
+from tensor2tensor.rl.envs.batch_env_factory import batch_env_factory
 from tensor2tensor.rl.envs.tf_atari_wrappers import WrapperBase
 from tensor2tensor.rl.envs.utils import get_policy
 
 
 def _rollout_metadata(batch_env):
   batch_env_shape = batch_env.observ.get_shape().as_list()
-  batch_size =  [batch_env_shape[0]]
+  batch_size = [batch_env_shape[0]]
   shapes_types_names = [
-    # observation
     (batch_size + batch_env_shape[1:], tf.float32, "observation"),
     (batch_size, tf.float32, "reward"),
     (batch_size, tf.bool, "done"),
@@ -74,7 +74,7 @@ def define_collect(hparams, scope, eval_phase,
   """Collect trajectories."""
 
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
-    batch_env = utils.batch_env_factory(hparams)
+    batch_env = batch_env_factory(hparams)
     environment_wrappers = hparams.environment_spec.wrappers
     wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
     #Put memory wrapper at the level you want to gather observations at
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
new file mode 100644
index 000000000..a07961c2e
--- /dev/null
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -0,0 +1,272 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Utilities for creating batched environments."""
+
+# The code was based on Danijar Hafner's code from tf.agents:
+# https://github.com/tensorflow/agents/blob/master/agents/tools/wrappers.py
+# https://github.com/tensorflow/agents/blob/master/agents/scripts/utility.py
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.rl.envs import batch_env
+from tensor2tensor.rl.envs import py_func_batch_env
+from tensor2tensor.rl.envs import simulated_batch_env
+
+import tensorflow as tf
+import atexit
+import multiprocessing
+import os
+import random
+import signal
+import subprocess
+import sys
+import traceback
+
+
+def batch_env_factory(hparams, xvfb=False):
+  """Factory of batch envs."""
+
+  environment_spec = hparams.environment_spec
+
+  if environment_spec.simulated_env:
+    # TODO(piotrmilos): Consider passing only relevant paramters
+    cur_batch_env = _define_simulated_batch_env(
+        hparams, hparams.num_agents,
+        hparams.simulation_random_starts,
+        hparams.intrinsic_reward_scale)
+  else:
+
+    cur_batch_env = _define_batch_env(hparams.environment_spec, hparams.num_agents,
+                                      xvfb=xvfb)
+  return cur_batch_env
+
+
+def _define_batch_env(environment_spec, num_agents, xvfb=False):
+  """Create environments and apply all desired wrappers."""
+
+  with tf.variable_scope("environments"):
+    envs = [
+        ExternalProcessEnv(environment_spec.env_lambda, xvfb)
+        for _ in range(num_agents)]
+    env = batch_env.BatchEnv(envs, blocking=False)
+    env = py_func_batch_env.PyFuncBatchEnv(env)
+    return env
+
+
+def _define_simulated_batch_env(hparams, num_agents,
+                               simulation_random_starts=False,
+                               intrinsic_reward_scale=0.):
+  cur_batch_env = simulated_batch_env.SimulatedBatchEnv(
+      hparams, num_agents, simulation_random_starts,
+      intrinsic_reward_scale)
+  return cur_batch_env
+
+
+class ExternalProcessEnv(object):
+  """Step environment in a separate process for lock free parallelism."""
+
+  # Message types for communication via the pipe.
+  _ACCESS = 1
+  _CALL = 2
+  _RESULT = 3
+  _EXCEPTION = 4
+  _CLOSE = 5
+
+  def __init__(self, constructor, xvfb):
+    """Step environment in a separate process for lock free parallelism.
+
+    The environment will be created in the external process by calling the
+    specified callable. This can be an environment class, or a function
+    creating the environment and potentially wrapping it. The returned
+    environment should not access global variables.
+
+    Args:
+      constructor: Callable that creates and returns an OpenAI gym environment.
+      xvfb:  Frame buffer.
+
+    Attributes:
+      observation_space: The cached observation space of the environment.
+      action_space: The cached action space of the environment.
+    """
+    self._conn, conn = multiprocessing.Pipe()
+    if xvfb:
+      server_id = random.randint(10000, 99999)
+      auth_file_id = random.randint(10000, 99999999999)
+
+      xauthority_path = "/tmp/Xauthority_{}".format(auth_file_id)
+
+      command = "Xvfb :{} -screen 0 1400x900x24 -nolisten tcp -auth {}".format(
+          server_id, xauthority_path)
+      with open(os.devnull, "w") as devnull:
+        proc = subprocess.Popen(command.split(), shell=False, stdout=devnull,
+                                stderr=devnull)
+        atexit.register(lambda: os.kill(proc.pid, signal.SIGKILL))
+
+      def constructor_using_xvfb():
+        os.environ["DISPLAY"] = ":{}".format(server_id)
+        os.environ["XAUTHORITY"] = xauthority_path
+        return constructor()
+
+      self._process = multiprocessing.Process(
+          target=self._worker, args=(constructor_using_xvfb, conn))
+    else:
+      self._process = multiprocessing.Process(
+          target=self._worker, args=(constructor, conn))
+
+    atexit.register(self.close)
+    self._process.start()
+    self._observ_space = None
+    self._action_space = None
+
+  @property
+  def observation_space(self):
+    if not self._observ_space:
+      self._observ_space = self.__getattr__("observation_space")
+    return self._observ_space
+
+  @property
+  def action_space(self):
+    if not self._action_space:
+      self._action_space = self.__getattr__("action_space")
+    return self._action_space
+
+  def __getattr__(self, name):
+    """Request an attribute from the environment.
+
+    Note that this involves communication with the external process, so it can
+    be slow.
+
+    Args:
+      name: Attribute to access.
+
+    Returns:
+      Value of the attribute.
+    """
+    self._conn.send((self._ACCESS, name))
+    return self._receive()
+
+  def call(self, name, *args, **kwargs):
+    """Asynchronously call a method of the external environment.
+
+    Args:
+      name: Name of the method to call.
+      *args: Positional arguments to forward to the method.
+      **kwargs: Keyword arguments to forward to the method.
+
+    Returns:
+      Promise object that blocks and provides the return value when called.
+    """
+    payload = name, args, kwargs
+    self._conn.send((self._CALL, payload))
+    return self._receive
+
+  def close(self):
+    """Send a close message to the external process and join it."""
+    try:
+      self._conn.send((self._CLOSE, None))
+      self._conn.close()
+    except IOError:
+      # The connection was already closed.
+      pass
+    self._process.join()
+
+  def step(self, action, blocking=True):
+    """Step the environment.
+
+    Args:
+      action: The action to apply to the environment.
+      blocking: Whether to wait for the result.
+
+    Returns:
+      Transition tuple when blocking, otherwise callable that returns the
+      transition tuple.
+    """
+    promise = self.call("step", action)
+    if blocking:
+      return promise()
+    return promise
+
+  def reset(self, blocking=True):
+    """Reset the environment.
+
+    Args:
+      blocking: Whether to wait for the result.
+
+    Returns:
+      New observation when blocking, otherwise callable that returns the new
+      observation.
+    """
+    promise = self.call("reset")
+    if blocking:
+      return promise()
+    return promise
+
+  def _receive(self):
+    """Wait for a message from the worker process and return its payload.
+
+    Raises:
+      Exception: An exception was raised inside the worker process.
+      KeyError: The received message is of an unknown type.
+
+    Returns:
+      Payload object of the message.
+    """
+    message, payload = self._conn.recv()
+    # Re-raise exceptions in the main process.
+    if message == self._EXCEPTION:
+      stacktrace = payload
+      raise Exception(stacktrace)
+    if message == self._RESULT:
+      return payload
+    raise KeyError("Received message of unexpected type {}".format(message))
+
+  def _worker(self, constructor, conn):
+    """The process waits for actions and sends back environment results.
+
+    Args:
+      constructor: Constructor for the OpenAI Gym environment.
+      conn: Connection for communication to the main process.
+    """
+    try:
+      env = constructor()
+      while True:
+        try:
+          # Only block for short times to have keyboard exceptions be raised.
+          if not conn.poll(0.1):
+            continue
+          message, payload = conn.recv()
+        except (EOFError, KeyboardInterrupt):
+          break
+        if message == self._ACCESS:
+          name = payload
+          result = getattr(env, name)
+          conn.send((self._RESULT, result))
+          continue
+        if message == self._CALL:
+          name, args, kwargs = payload
+          result = getattr(env, name)(*args, **kwargs)
+          conn.send((self._RESULT, result))
+          continue
+        if message == self._CLOSE:
+          assert payload is None
+          env.close()
+          break
+        raise KeyError("Received message of unknown type {}".format(message))
+    except Exception:  # pylint: disable=broad-except
+      stacktrace = "".join(traceback.format_exception(*sys.exc_info()))  # pylint: disable=no-value-for-parameter
+      tf.logging.error("Error in environment process: {}".format(stacktrace))
+      conn.send((self._EXCEPTION, stacktrace))
+    conn.close()
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 1e6425341..ea9660bc1 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -22,7 +22,7 @@
 from __future__ import print_function
 from tensor2tensor.layers import common_layers
 from tensor2tensor.rl.envs import in_graph_batch_env
-from tensor2tensor.rl.envs import utils
+from tensor2tensor.rl.envs.utils import get_action_space
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
@@ -108,7 +108,7 @@ def __init__(self, hparams, length, simulation_random_starts=False,
     self._model = registry.model(FLAGS.model)(
       model_hparams, tf.estimator.ModeKeys.PREDICT)
 
-    _, self.action_shape, self.action_dtype = utils.get_action_space(environment_spec)
+    _, self.action_shape, self.action_dtype = get_action_space(environment_spec)
 
     if simulation_random_starts:
       dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index 58f8b1ee9..f304c9dcf 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -14,27 +14,11 @@
 # limitations under the License.
 """Utilities for using batched environments."""
 
-# The code was based on Danijar Hafner's code from tf.agents:
-# https://github.com/tensorflow/agents/blob/master/agents/tools/wrappers.py
-# https://github.com/tensorflow/agents/blob/master/agents/scripts/utility.py
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import atexit
-import multiprocessing
-import os
-import random
-import signal
-import subprocess
-import sys
-import traceback
 import gym
-
-from tensor2tensor.rl.envs import batch_env
-from tensor2tensor.rl.envs import py_func_batch_env
-from tensor2tensor.rl.envs import simulated_batch_env
-
 import tensorflow as tf
 
 
@@ -81,242 +65,6 @@ def _reset(self, **kwargs):
     return self._last_returned[0]
 
 
-class ExternalProcessEnv(object):
-  """Step environment in a separate process for lock free parallelism."""
-
-  # Message types for communication via the pipe.
-  _ACCESS = 1
-  _CALL = 2
-  _RESULT = 3
-  _EXCEPTION = 4
-  _CLOSE = 5
-
-  def __init__(self, constructor, xvfb):
-    """Step environment in a separate process for lock free parallelism.
-
-    The environment will be created in the external process by calling the
-    specified callable. This can be an environment class, or a function
-    creating the environment and potentially wrapping it. The returned
-    environment should not access global variables.
-
-    Args:
-      constructor: Callable that creates and returns an OpenAI gym environment.
-      xvfb:  Frame buffer.
-
-    Attributes:
-      observation_space: The cached observation space of the environment.
-      action_space: The cached action space of the environment.
-    """
-    self._conn, conn = multiprocessing.Pipe()
-    if xvfb:
-      server_id = random.randint(10000, 99999)
-      auth_file_id = random.randint(10000, 99999999999)
-
-      xauthority_path = "/tmp/Xauthority_{}".format(auth_file_id)
-
-      command = "Xvfb :{} -screen 0 1400x900x24 -nolisten tcp -auth {}".format(
-          server_id, xauthority_path)
-      with open(os.devnull, "w") as devnull:
-        proc = subprocess.Popen(command.split(), shell=False, stdout=devnull,
-                                stderr=devnull)
-        atexit.register(lambda: os.kill(proc.pid, signal.SIGKILL))
-
-      def constructor_using_xvfb():
-        os.environ["DISPLAY"] = ":{}".format(server_id)
-        os.environ["XAUTHORITY"] = xauthority_path
-        return constructor()
-
-      self._process = multiprocessing.Process(
-          target=self._worker, args=(constructor_using_xvfb, conn))
-    else:
-      self._process = multiprocessing.Process(
-          target=self._worker, args=(constructor, conn))
-
-    atexit.register(self.close)
-    self._process.start()
-    self._observ_space = None
-    self._action_space = None
-
-  @property
-  def observation_space(self):
-    if not self._observ_space:
-      self._observ_space = self.__getattr__("observation_space")
-    return self._observ_space
-
-  @property
-  def action_space(self):
-    if not self._action_space:
-      self._action_space = self.__getattr__("action_space")
-    return self._action_space
-
-  def __getattr__(self, name):
-    """Request an attribute from the environment.
-
-    Note that this involves communication with the external process, so it can
-    be slow.
-
-    Args:
-      name: Attribute to access.
-
-    Returns:
-      Value of the attribute.
-    """
-    self._conn.send((self._ACCESS, name))
-    return self._receive()
-
-  def call(self, name, *args, **kwargs):
-    """Asynchronously call a method of the external environment.
-
-    Args:
-      name: Name of the method to call.
-      *args: Positional arguments to forward to the method.
-      **kwargs: Keyword arguments to forward to the method.
-
-    Returns:
-      Promise object that blocks and provides the return value when called.
-    """
-    payload = name, args, kwargs
-    self._conn.send((self._CALL, payload))
-    return self._receive
-
-  def close(self):
-    """Send a close message to the external process and join it."""
-    try:
-      self._conn.send((self._CLOSE, None))
-      self._conn.close()
-    except IOError:
-      # The connection was already closed.
-      pass
-    self._process.join()
-
-  def step(self, action, blocking=True):
-    """Step the environment.
-
-    Args:
-      action: The action to apply to the environment.
-      blocking: Whether to wait for the result.
-
-    Returns:
-      Transition tuple when blocking, otherwise callable that returns the
-      transition tuple.
-    """
-    promise = self.call("step", action)
-    if blocking:
-      return promise()
-    return promise
-
-  def reset(self, blocking=True):
-    """Reset the environment.
-
-    Args:
-      blocking: Whether to wait for the result.
-
-    Returns:
-      New observation when blocking, otherwise callable that returns the new
-      observation.
-    """
-    promise = self.call("reset")
-    if blocking:
-      return promise()
-    return promise
-
-  def _receive(self):
-    """Wait for a message from the worker process and return its payload.
-
-    Raises:
-      Exception: An exception was raised inside the worker process.
-      KeyError: The received message is of an unknown type.
-
-    Returns:
-      Payload object of the message.
-    """
-    message, payload = self._conn.recv()
-    # Re-raise exceptions in the main process.
-    if message == self._EXCEPTION:
-      stacktrace = payload
-      raise Exception(stacktrace)
-    if message == self._RESULT:
-      return payload
-    raise KeyError("Received message of unexpected type {}".format(message))
-
-  def _worker(self, constructor, conn):
-    """The process waits for actions and sends back environment results.
-
-    Args:
-      constructor: Constructor for the OpenAI Gym environment.
-      conn: Connection for communication to the main process.
-    """
-    try:
-      env = constructor()
-      while True:
-        try:
-          # Only block for short times to have keyboard exceptions be raised.
-          if not conn.poll(0.1):
-            continue
-          message, payload = conn.recv()
-        except (EOFError, KeyboardInterrupt):
-          break
-        if message == self._ACCESS:
-          name = payload
-          result = getattr(env, name)
-          conn.send((self._RESULT, result))
-          continue
-        if message == self._CALL:
-          name, args, kwargs = payload
-          result = getattr(env, name)(*args, **kwargs)
-          conn.send((self._RESULT, result))
-          continue
-        if message == self._CLOSE:
-          assert payload is None
-          env.close()
-          break
-        raise KeyError("Received message of unknown type {}".format(message))
-    except Exception:  # pylint: disable=broad-except
-      stacktrace = "".join(traceback.format_exception(*sys.exc_info()))  # pylint: disable=no-value-for-parameter
-      tf.logging.error("Error in environment process: {}".format(stacktrace))
-      conn.send((self._EXCEPTION, stacktrace))
-    conn.close()
-
-
-def batch_env_factory(hparams, xvfb=False):
-  """Factory of batch envs."""
-
-  environment_spec = hparams.environment_spec
-
-  if environment_spec.simulated_env:
-    # TODO(piotrmilos): Consider passing only relevant paramters
-    cur_batch_env = _define_simulated_batch_env(
-        hparams, hparams.num_agents,
-        hparams.simulation_random_starts,
-        hparams.intrinsic_reward_scale)
-  else:
-
-    cur_batch_env = _define_batch_env(hparams.environment_spec, hparams.num_agents,
-                                      xvfb=xvfb)
-  return cur_batch_env
-
-
-def _define_batch_env(environment_spec, num_agents, xvfb=False):
-  """Create environments and apply all desired wrappers."""
-
-  with tf.variable_scope("environments"):
-    envs = [
-        ExternalProcessEnv(environment_spec.env_lambda, xvfb)
-        for _ in range(num_agents)]
-    env = batch_env.BatchEnv(envs, blocking=False)
-    env = py_func_batch_env.PyFuncBatchEnv(env)
-    return env
-
-
-def _define_simulated_batch_env(hparams, num_agents,
-                               simulation_random_starts=False,
-                               intrinsic_reward_scale=0.):
-  cur_batch_env = simulated_batch_env.SimulatedBatchEnv(
-      hparams, num_agents, simulation_random_starts,
-      intrinsic_reward_scale)
-  return cur_batch_env
-
-
 def get_action_space(environment_spec):
   """Get action spece associated with environment spec
     

From 5322dd6aaadbaae01eeafcd832630afeeb0beef5 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Wed, 20 Jun 2018 11:27:46 +0200
Subject: [PATCH 0161/2720] linter fixes

---
 tensor2tensor/data_generators/gym_problems.py | 25 +++++++------
 tensor2tensor/models/research/rl.py           |  6 ++--
 tensor2tensor/rl/collect.py                   | 36 +++++++++----------
 tensor2tensor/rl/envs/batch_env_factory.py    |  8 +++--
 tensor2tensor/rl/envs/py_func_batch_env.py    |  5 +--
 tensor2tensor/rl/envs/simulated_batch_env.py  | 17 +++++----
 tensor2tensor/rl/envs/utils.py                | 12 +++----
 tensor2tensor/rl/model_rl_experiment.py       |  4 +--
 tensor2tensor/rl/rl_trainer_lib.py            | 11 +++---
 tensor2tensor/rl/rl_trainer_lib_test.py       |  5 +--
 10 files changed, 67 insertions(+), 62 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index ae96dba82..b5dc8c78e 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -73,7 +73,8 @@ def _setup(self):
     self._internal_memory_size = 10
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       collect_hparams.epoch_length = self._internal_memory_size
-      collect_hparams.num_agents = 1 #TODO (piotrmilos). it is possible to set more
+      # TODO (piotrmilos). it is possible to set more then 1
+      collect_hparams.num_agents = 1
       self.collect_memory, self.collect_trigger_op \
         = collect.define_collect(collect_hparams, scope="gym_problems",
                                  collect_level=0, eval_phase=self.eval_phase)
@@ -90,20 +91,18 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
       memory_index = 0
       memory = None
       while pieces_generated < self.num_steps:
-        if memory is None or memory_index>=self._internal_memory_size:
+        if memory is None or memory_index >= self._internal_memory_size:
           sess.run(self.collect_trigger_op)
           memory = sess.run(self.collect_memory)
           memory_index = 0
         data = [memory[i][memory_index][0] for i in range(4)]
         memory_index += 1
         observ, reward, done, action = data
-        observ = observ.astype(np.uint8) # TODO(piotrmilos). This should be probably done in collect
+        observ = observ.astype(np.uint8)
 
-        debug_im = None
         self.sum_of_rewards += reward
         self.dones += int(done)
 
-
         ret_dict = {"frame": observ,
                     "image/format": ["png"],
                     "image/height": [self.frame_height],
@@ -246,7 +245,7 @@ def __init__(self, *args, **kwargs):
 
   @property
   def initial_frames_problem(self):
-    raise NotImplemented()
+    raise NotImplementedError()
 
   def get_environment_spec(self):
     env_spec = standard_atari_env_spec(self.env_name)
@@ -266,9 +265,9 @@ def get_environment_spec(self):
 
     env_spec.simulated_env = True
     env_spec.add_hparam("simulation_random_starts",
-                           self.simulation_random_starts)
+                        self.simulation_random_starts)
     env_spec.add_hparam("intrinsic_reward_scale",
-                           self.intrinsic_reward_scale)
+                        self.intrinsic_reward_scale)
     initial_frames_problem = registry.problem(self.initial_frames_problem)
     env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
     env_spec.wrappers.append([TimeLimitWrapper, {"timelimit": timelimit}])
@@ -405,7 +404,7 @@ class GymDiscreteProblemWithAgentOnPong(
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
-  GymSimulatedDiscreteProblem, GymWrappedPongRandom):
+    GymSimulatedDiscreteProblem, GymWrappedPongRandom):
 
   @property
   def initial_frames_problem(self):
@@ -414,7 +413,7 @@ def initial_frames_problem(self):
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedLongPong(
-  GymRealDiscreteProblem, GymWrappedLongPongRandom):
+    GymRealDiscreteProblem, GymWrappedLongPongRandom):
   pass
 
 
@@ -435,7 +434,7 @@ def initial_frames_problem(self):
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedBreakout(
-  GymRealDiscreteProblem, GymWrappedBreakoutRandom):
+    GymRealDiscreteProblem, GymWrappedBreakoutRandom):
   pass
 
 
@@ -457,7 +456,7 @@ def initial_frames_problem(self):
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedPong(
-  GymRealDiscreteProblem, GymWrappedPongRandom):
+    GymRealDiscreteProblem, GymWrappedPongRandom):
   """GymDiscreteProblemWithAgentOnWrappedPong."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
@@ -497,7 +496,7 @@ def initial_frames_problem(self):
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnFreeway(
-  GymRealDiscreteProblem, GymFreewayRandom):
+    GymRealDiscreteProblem, GymFreewayRandom):
   """Freeway with agent."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 2a643e5af..a7109455e 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -121,7 +121,7 @@ def ppo_pong_base():
 
 
 def standard_atari_env_spec(env):
-  """Parameters of environement specification"""
+  """Parameters of environment specification"""
   standard_wrappers = [[tf_atari_wrappers.MaxAndSkipWrapper, {"skip": 4}]]
   env_lambda = None
   if isinstance(env, str):
@@ -136,11 +136,11 @@ def standard_atari_env_spec(env):
 
 
 def simple_gym_spec(env):
-  """Parameters of environement specification"""
+  """Parameters of environment specification"""
   standard_wrappers = None
   env_lambda = None
   if isinstance(env, str):
-    env_lambda = lambda : gym.make(env)
+    env_lambda = lambda: gym.make(env)
   if callable(env):
     env_lambda = env
   assert env is not None, "Unknown specification of environment"
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 9b18057ab..64f150f7a 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -30,12 +30,12 @@ def _rollout_metadata(batch_env):
   batch_env_shape = batch_env.observ.get_shape().as_list()
   batch_size = [batch_env_shape[0]]
   shapes_types_names = [
-    (batch_size + batch_env_shape[1:], tf.float32, "observation"),
-    (batch_size, tf.float32, "reward"),
-    (batch_size, tf.bool, "done"),
-    (batch_size + batch_env.action_shape, batch_env.action_dtype, "action"),
-    (batch_size, tf.float32, "pdf"),
-    (batch_size, tf.float32, "value_function"),
+      (batch_size + batch_env_shape[1:], tf.float32, "observation"),
+      (batch_size, tf.float32, "reward"),
+      (batch_size, tf.bool, "done"),
+      (batch_size + batch_env.action_shape, batch_env.action_dtype, "action"),
+      (batch_size, tf.float32, "pdf"),
+      (batch_size, tf.float32, "value_function"),
   ]
   return shapes_types_names
 
@@ -79,7 +79,8 @@ def define_collect(hparams, scope, eval_phase,
     wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
     #Put memory wrapper at the level you want to gather observations at
     #Negative indices need to be shifted for insert to work correctly
-    collect_level = collect_level if collect_level>=0 else len(wrappers) + collect_level + 1
+    collect_level = collect_level if \
+      collect_level >= 0 else len(wrappers) + collect_level + 1
     wrappers.insert(collect_level, [_MemoryWrapper, {}])
     rollout_metadata = None
     speculum = None
@@ -101,7 +102,7 @@ def define_collect(hparams, scope, eval_phase,
 
 
     cumulative_rewards = tf.get_variable("cumulative_rewards", len(batch_env),
-                                           trainable=False)
+                                         trainable=False)
 
     should_reset_var = tf.Variable(True, trainable=False)
 
@@ -127,7 +128,7 @@ def step(index, scores_sum, scores_num):
       # https://github.com/tensorflow/tensorflow/issues/11186
       obs_copy = batch_env.observ + 0
 
-      def env_step(arg1, arg2):
+      def env_step(arg1, arg2): # pylint: disable=unused-argument
         actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
         policy = actor_critic.policy
         if policy_to_actions_lambda:
@@ -142,20 +143,19 @@ def env_step(arg1, arg2):
 
         pdf = policy.prob(action)[0]
         value_function = actor_critic.value[0]
-        pdf = tf.reshape(pdf, shape=(hparams.num_agents, ))
-        value_function = tf.reshape(value_function, shape=(hparams.num_agents, ))
+        pdf = tf.reshape(pdf, shape=(hparams.num_agents,))
+        value_function = tf.reshape(value_function, shape=(hparams.num_agents,))
 
         with tf.control_dependencies(simulate_output):
           return tf.identity(pdf), tf.identity(value_function)
 
       pdf, value_function = tf.while_loop(
-        lambda _1, _2: tf.equal(speculum.size(), 0),
-        env_step,
-        [tf.constant(0.0, shape=(hparams.num_agents,)),
-          tf.constant(0.0, shape=(hparams.num_agents,))],
-        parallel_iterations=1,
-        back_prop=False,
-        )
+          lambda _1, _2: tf.equal(speculum.size(), 0),
+          env_step,
+          [tf.constant(0.0, shape=(hparams.num_agents,)),
+           tf.constant(0.0, shape=(hparams.num_agents,))],
+          parallel_iterations=1,
+          back_prop=False,)
 
       with tf.control_dependencies([pdf, value_function]):
         obs, reward, done, action = speculum.dequeue()
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index a07961c2e..31f2e1b6b 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -17,6 +17,7 @@
 # The code was based on Danijar Hafner's code from tf.agents:
 # https://github.com/tensorflow/agents/blob/master/agents/tools/wrappers.py
 # https://github.com/tensorflow/agents/blob/master/agents/scripts/utility.py
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
@@ -49,7 +50,8 @@ def batch_env_factory(hparams, xvfb=False):
         hparams.intrinsic_reward_scale)
   else:
 
-    cur_batch_env = _define_batch_env(hparams.environment_spec, hparams.num_agents,
+    cur_batch_env = _define_batch_env(hparams.environment_spec,
+                                      hparams.num_agents,
                                       xvfb=xvfb)
   return cur_batch_env
 
@@ -67,8 +69,8 @@ def _define_batch_env(environment_spec, num_agents, xvfb=False):
 
 
 def _define_simulated_batch_env(hparams, num_agents,
-                               simulation_random_starts=False,
-                               intrinsic_reward_scale=0.):
+                                simulation_random_starts=False,
+                                intrinsic_reward_scale=0.):
   cur_batch_env = simulated_batch_env.SimulatedBatchEnv(
       hparams, num_agents, simulation_random_starts,
       intrinsic_reward_scale)
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index 08ac236f1..0cce5167c 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -22,8 +22,9 @@
 from __future__ import print_function
 
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
-import tensorflow as tf
 from tensor2tensor.rl.envs import utils
+import tensorflow as tf
+
 
 class PyFuncBatchEnv(InGraphBatchEnv):
   """Batch of environments inside the TensorFlow graph.
@@ -117,4 +118,4 @@ def observ(self):
 
   def close(self):
     """Send close messages to the external process and join them."""
-    self._batch_env.close()
\ No newline at end of file
+    self._batch_env.close()
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index ea9660bc1..5f1cfda6a 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -106,23 +106,28 @@ def __init__(self, hparams, length, simulation_random_starts=False,
         FLAGS.hparams_set, problem_name=FLAGS.problem)
     model_hparams.force_full_predict = True
     self._model = registry.model(FLAGS.model)(
-      model_hparams, tf.estimator.ModeKeys.PREDICT)
+        model_hparams, tf.estimator.ModeKeys.PREDICT)
 
     _, self.action_shape, self.action_dtype = get_action_space(environment_spec)
 
     if simulation_random_starts:
-      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
-                                shuffle_files=True, hparams=hparams)
+      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
+                                               FLAGS.data_dir,
+                                               shuffle_files=True,
+                                               hparams=hparams)
       dataset = dataset.shuffle(buffer_size=100)
     else:
-      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
-                                shuffle_files=True, hparams=hparams).take(1)
+      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
+                                               FLAGS.data_dir,
+                                               shuffle_files=True,
+                                               hparams=hparams).take(1)
 
     dataset = dataset.map(lambda x: x["inputs"]).repeat()
     self.history_buffer = HistoryBuffer(dataset, self.length)
 
     shape = (self.length, initial_frames_problem.frame_height,
-             initial_frames_problem.frame_width, initial_frames_problem.num_channels)
+             initial_frames_problem.frame_width,
+             initial_frames_problem.num_channels)
     self._observ = tf.Variable(tf.zeros(shape, tf.float32), trainable=False)
 
   def __len__(self):
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index f304c9dcf..ea6a8609c 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -67,10 +67,10 @@ def _reset(self, **kwargs):
 
 def get_action_space(environment_spec):
   """Get action spece associated with environment spec
-    
+
   Args:
      environment_spec:  EnvironmentSpec object
-     
+
   Returns:
     OpenAi Gym action space
   """
@@ -83,11 +83,11 @@ def get_action_space(environment_spec):
 
 def get_policy(observations, hparams):
   """Get policy network
-  
+
   Args:
     observations: Tensor with observations
-    hparams: parameters 
-    
+    hparams: parameters
+
   Returns:
     Tensor with policy and value function output
   """
@@ -124,4 +124,4 @@ def parse_dtype(space):
     return tf.int32
   if isinstance(space, gym.spaces.Box):
     return tf.float32
-  raise NotImplementedError()
\ No newline at end of file
+  raise NotImplementedError()
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 76d1383b6..8067f6aa6 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -159,10 +159,10 @@ def train_agent(problem_name, agent_model_dir,
   environment_spec.add_hparam("initial_frames_problem",
                               gym_problem)
 
-
   # 4x for the StackAndSkipWrapper minus one to always finish for reporting.
   ppo_time_limit = (ppo_hparams.epoch_length - 1) * 4
-  wrappers = environment_spec.wrappers + [[TimeLimitWrapper, {"timelimit": ppo_time_limit}]]
+  wrappers = environment_spec.wrappers + \
+             [[TimeLimitWrapper, {"timelimit": ppo_time_limit}]]
   environment_spec.wrappers = wrappers
 
   ppo_hparams.add_hparam("environment_spec", environment_spec)
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 6ae473d20..4a3a093d3 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -17,27 +17,24 @@
 from __future__ import division
 from __future__ import print_function
 
-import functools
 import os
-import gym
 
 from tensor2tensor import models  # pylint: disable=unused-import
 from tensor2tensor.models.research import rl  # pylint: disable=unused-import
 from tensor2tensor.rl import collect
 from tensor2tensor.rl import ppo
-from tensor2tensor.rl.envs import tf_atari_wrappers
-from tensor2tensor.rl.envs import utils
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
 
 
-def define_train(hparams, event_dir):
+def define_train(hparams, event_dir): # pylint: disable=unused-argument
   """Define the training setup."""
 
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
     memory, collect_summary = collect.define_collect(
-        hparams, "ppo_train", eval_phase=False, on_simulated=hparams.simulated_environment)
+        hparams, "ppo_train", eval_phase=False,
+        on_simulated=hparams.simulated_environment)
     ppo_summary = ppo.define_ppo_epoch(memory, hparams)
     summary = tf.summary.merge([collect_summary, ppo_summary])
 
@@ -48,7 +45,7 @@ def train(hparams, event_dir=None, model_dir=None,
           restore_agent=True, epoch=0):
   """Train."""
   with tf.name_scope("rl_train"):
-    train_summary_op, eval_summary_op = define_train(hparams, event_dir)
+    train_summary_op, eval_summary_op = define_train(hparams, event_dir) # pylint: disable=unused-variable
     if event_dir:
       summary_writer = tf.summary.FileWriter(
           event_dir, graph=tf.get_default_graph(), flush_secs=60)
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index fba44c31e..6066bd7fc 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -19,7 +19,7 @@
 
 from tensor2tensor.models.research.rl import simple_gym_spec, standard_atari_env_spec
 from tensor2tensor.rl import rl_trainer_lib
-from tensor2tensor.utils import trainer_lib, registry
+from tensor2tensor.utils import trainer_lib, registry # pylint: disable=unused-import
 
 import tensorflow as tf
 
@@ -42,7 +42,8 @@ def test_no_crash_cartpole(self):
         "ppo_discrete_action_base",
         TrainTest.test_config)
 
-    hparams.add_hparam("environment_spec", standard_atari_env_spec("CartPole-v0"))
+    hparams.add_hparam("environment_spec",
+                       standard_atari_env_spec("CartPole-v0"))
     rl_trainer_lib.train(hparams)
 
 

From 20754e50106d4cda3173456e9f9723d310aec413 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 20 Jun 2018 09:43:09 -0700
Subject: [PATCH 0162/2720] Correcting number of frames for RL envs like Pong

PiperOrigin-RevId: 201367273
---
 tensor2tensor/models/research/next_frame.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 5e0c148bc..0bbbd891c 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -1088,6 +1088,8 @@ def next_frame_stochastic_cutoff():
   """SV2P model with additional cutoff in L2 loss for environments like pong."""
   hparams = next_frame_stochastic()
   hparams.video_modality_loss_cutoff = 0.4
+  hparams.video_num_input_frames = 4
+  hparams.video_num_target_frames = 1
   return hparams
 
 

From f4fee3616118b68b6fa4c65f3aa523b51310cc3e Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 20 Jun 2018 13:27:57 -0700
Subject: [PATCH 0163/2720] Lint fix: losses signature

PiperOrigin-RevId: 201405468
---
 tensor2tensor/models/vanilla_gan.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index c55c403d1..e9fddc096 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -167,13 +167,14 @@ def top(self, body_output, features):
 class SlicedGan(AbstractGAN):
   """Sliced GAN for demonstration."""
 
-  def losses(self, inputs, g):
+  def losses(self, inputs, generated):
     """Losses in the sliced case."""
     is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
     def discriminate(x):
       return self.discriminator(x, is_training=is_training, reuse=False)
     generator_loss = common_layers.sliced_gan_loss(
-        inputs, reverse_gradient(g), discriminate, self.hparams.num_sliced_vecs)
+        inputs, reverse_gradient(generated), discriminate,
+        self.hparams.num_sliced_vecs)
     return {"training": - generator_loss}
 
   def infer(self, *args, **kwargs):  # pylint: disable=arguments-differ

From b9e947ac004d1efb495570af008beddb7e423fd0 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 20 Jun 2018 15:01:37 -0700
Subject: [PATCH 0164/2720] pad last eval batch to fixed batch size if
 applicable

PiperOrigin-RevId: 201422739
---
 tensor2tensor/data_generators/problem.py | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 5f7ca2808..ab8ea1cff 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -17,6 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 import collections
+import functools
 import os
 import random
 
@@ -812,9 +813,10 @@ def define_shapes(example):
       else:
         # On GPU, bucket by length
         dataset = dataset.filter(gpu_valid_size)
+        shard_multiplier = (config and config.data_parallelism.n) or 1
         batching_scheme = data_reader.hparams_to_batching_scheme(
             hparams,
-            shard_multiplier=(config and config.data_parallelism.n) or 1,
+            shard_multiplier=shard_multiplier,
             length_multiplier=self.get_hparams().batch_size_multiplier)
         if hparams.use_fixed_batch_size:
           # Here  batch_size really means examples per datashard.
@@ -825,18 +827,19 @@ def define_shapes(example):
             batching_scheme["batch_sizes"])
 
         if not is_training:
-
-          def _pad_batch(features):
-            if not config or config.data_parallelism.n <= 1:
-              return features
+          batch_multiple = shard_multiplier
+          if hparams.use_fixed_batch_size:
+            # Make sure the last batch has the same fixed size as the rest.
+            batch_multiple *= hparams.batch_size
+          if batch_multiple > 1:
             tf.logging.warn(
                 "Padding the batch to ensure that remainder eval batches have "
                 "a batch size divisible by the number of data shards. This may "
                 "lead to incorrect metrics for non-zero-padded features, e.g. "
                 "images. Use a single datashard (i.e. 1 GPU) in that case.")
-            return pad_batch(features, config.data_parallelism.n)
-
-          dataset = dataset.map(_pad_batch, num_parallel_calls=num_threads)
+            dataset = dataset.map(
+                functools.partial(pad_batch, batch_multiple=batch_multiple),
+                num_parallel_calls=num_threads)
 
     dataset = dataset.map(define_shapes, num_parallel_calls=num_threads)
 

From 7624602a58fcee2c6fc6899481a589f4adbb9dfa Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 20 Jun 2018 15:10:26 -0700
Subject: [PATCH 0165/2720] Ignore arguments-differ linter errors

PiperOrigin-RevId: 201424591
---
 pylintrc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pylintrc b/pylintrc
index 0d4e111ef..157d214e3 100644
--- a/pylintrc
+++ b/pylintrc
@@ -37,7 +37,7 @@ msg-template={msg_id}:{line:3} {obj}: {msg} [{symbol}]
 enable=indexing-exception,old-raise-syntax
 
 # List of checkers and warnings to disable.
-disable=design,similarities,no-self-use,attribute-defined-outside-init,locally-disabled,star-args,pointless-except,bad-option-value,global-statement,fixme,suppressed-message,useless-suppression,locally-enabled,file-ignored,multiple-imports,c-extension-no-member,trailing-newlines,unsubscriptable-object,misplaced-comparison-constant,no-member,abstract-method,no-else-return,missing-docstring,wrong-import-order,protected-access,inconsistent-return-statements,invalid-unary-operand-type,import-error,no-name-in-module
+disable=design,similarities,no-self-use,attribute-defined-outside-init,locally-disabled,star-args,pointless-except,bad-option-value,global-statement,fixme,suppressed-message,useless-suppression,locally-enabled,file-ignored,multiple-imports,c-extension-no-member,trailing-newlines,unsubscriptable-object,misplaced-comparison-constant,no-member,abstract-method,no-else-return,missing-docstring,wrong-import-order,protected-access,inconsistent-return-statements,invalid-unary-operand-type,import-error,no-name-in-module,arguments-differ
 
 [BASIC]
 

From e04b493a9fa73f269ba78d45dd164b123049acc8 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 20 Jun 2018 17:00:24 -0700
Subject: [PATCH 0166/2720] Sentiment IMDb with character level support.

PiperOrigin-RevId: 201442396
---
 tensor2tensor/data_generators/imdb.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensor2tensor/data_generators/imdb.py b/tensor2tensor/data_generators/imdb.py
index c914df22e..e6542e1c0 100644
--- a/tensor2tensor/data_generators/imdb.py
+++ b/tensor2tensor/data_generators/imdb.py
@@ -95,3 +95,11 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
           "inputs": doc,
           "label": int(label),
       }
+
+
+@registry.register_problem
+class SentimentIMDBCharacters(SentimentIMDB):
+  """IMDB Sentiment classification, character level."""
+
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER

From d7d9139d6a8926d5c556b87842a5c9559e0ef7b5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 20 Jun 2018 17:24:21 -0700
Subject: [PATCH 0167/2720] Internal change

PiperOrigin-RevId: 201445778
---
 tensor2tensor/utils/decoding.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 07f8ad6cc..4b6552dea 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -551,10 +551,12 @@ def save_video(video, save_path_template):
 
   for i, frame in enumerate(video):
     save_path = save_path_template.format(i)
-    Image.fromarray(np.uint8(frame)).save(save_path)
+    with tf.gfile.Open(save_path, "wb") as sp:
+      Image.fromarray(np.uint8(frame)).save(sp)
 
 
 def show_and_save_image(img, save_path):
+  """Shows an image using matplotlib and saves it."""
   try:
     import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top
   except ImportError as e:
@@ -563,7 +565,8 @@ def show_and_save_image(img, save_path):
         "installed: %s", e)
     raise NotImplementedError("Image display and save not implemented.")
   plt.imshow(img)
-  plt.savefig(save_path)
+  with tf.gfile.Open(save_path, "wb") as sp:
+    plt.savefig(sp)
 
 
 def _get_sorted_inputs(filename, num_shards=1, delimiter="\n"):

From af378f52eee11dbcde65ef46ed77d628a407de13 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Wed, 20 Jun 2018 17:28:07 -0700
Subject: [PATCH 0168/2720] Tentatively bring masking back as it makes it more
 stable

PiperOrigin-RevId: 201446222
---
 .../models/research/transformer_nat.py        | 26 +++++++++++++++----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 70e7075d0..078e96734 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -202,8 +202,6 @@ def get_latent_pred_loss(latents_pred, latents_discrete_hot, hparams):
   """Latent prediction and loss."""
   latents_logits = tf.layers.dense(
       latents_pred, 2**hparams.bottleneck_bits, name="extra_logits")
-  # loss = tf.losses.softmax_cross_entropy(onehot_labels=latents_discrete_hot,
-  #                                        logits=latents_logits)
   loss = tf.nn.softmax_cross_entropy_with_logits_v2(
       labels=tf.stop_gradient(latents_discrete_hot), logits=latents_logits)
   return loss
@@ -260,8 +258,13 @@ def ae_transformer_internal(inputs, targets, target_space, hparams, cache=None):
   targets_c = compress(targets, hparams, "compress")
   if hparams.mode != tf.estimator.ModeKeys.PREDICT:
     # Compress and bottleneck.
+    if hparams.denoising:
+      noise = tf.random_normal(
+          shape=common_layers.shape_list(targets_c), mean=0.0, stddev=0.1)
+    else:
+      noise = 0
     latents_discrete_hot, extra_loss = vq_discrete_bottleneck(
-        x=targets_c, hparams=hparams)
+        x=targets_c + noise, hparams=hparams)
     latents_dense = vq_discrete_unbottleneck(
         latents_discrete_hot, hparams=hparams)
     latents_dense = targets_c + tf.stop_gradient(latents_dense - targets_c)
@@ -297,8 +300,20 @@ def ae_transformer_internal(inputs, targets, target_space, hparams, cache=None):
     d = residual_conv(d, 1, (3, 1), hparams, "decompress_rc_%d" % j)
     d = decompress_step(d, hparams, i > 0, "decompress_%d" % j)
 
-  targets = d
-  res = decode_transformer(inputs, ed, d, hparams, "decoder")
+  masking = common_layers.inverse_lin_decay(hparams.mask_startup_steps)
+  masking *= common_layers.inverse_exp_decay(
+      hparams.mask_startup_steps // 4)  # Not much at start.
+  masking = tf.minimum(tf.maximum(masking, 0.0), 1.0)
+  if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+    masking = 1.0
+  mask = tf.less(masking,
+                 tf.random_uniform(common_layers.shape_list(targets)[:-1]))
+  mask = tf.expand_dims(tf.to_float(mask), 3)
+
+  # targets is always [batch, length, 1, depth]
+  targets = mask * targets + (1.0 - mask) * d
+
+  res = decode_transformer(inputs, ed, targets, hparams, "decoder")
   latent_time = tf.less(hparams.mask_startup_steps,
                         tf.to_int32(tf.train.get_global_step()))
   losses["latent_pred"] *= tf.to_float(latent_time)
@@ -397,6 +412,7 @@ def transformer_nat_small():
   hparams.add_hparam("decay", 0.999)
   hparams.add_hparam("num_samples", 10)
   hparams.add_hparam("mask_startup_steps", 50000)
+  hparams.add_hparam("denoising", False)
   return hparams
 
 

From c2048308ff151d1d4034b2672cfcd19f3fede65a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 20 Jun 2018 18:18:17 -0700
Subject: [PATCH 0169/2720] Add a GAN option to autoencoders, normalize ppo
 learning rate.

PiperOrigin-RevId: 201452757
---
 tensor2tensor/models/basic.py                 | 80 +++++++++++++++++--
 tensor2tensor/models/research/autoencoders.py |  7 +-
 tensor2tensor/models/research/rl.py           |  6 +-
 3 files changed, 82 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/models/basic.py b/tensor2tensor/models/basic.py
index 4d9c7b763..302396b39 100644
--- a/tensor2tensor/models/basic.py
+++ b/tensor2tensor/models/basic.py
@@ -25,6 +25,14 @@
 import tensorflow as tf
 
 
+def lrelu(input_, leak=0.2, name="lrelu"):
+  return tf.maximum(input_, leak * input_, name=name)
+
+
+def reverse_gradient(x):
+  return -x + tf.stop_gradient(2 * x)
+
+
 @registry.register_model
 class BasicFcRelu(t2t_model.T2TModel):
   """Basic fully-connected + ReLU model."""
@@ -59,8 +67,44 @@ def bottleneck(self, x):
         return tf.tanh(x) + noise * hparams.bottleneck_noise, 0.0
       return tf.tanh(x), 0.0
 
-  def unbottleneck(self, x, res_size):
-    with tf.variable_scope("unbottleneck"):
+  def discriminator(self, x, is_training):
+    """Discriminator architecture based on InfoGAN.
+
+    Args:
+      x: input images, shape [bs, h, w, channels]
+      is_training: boolean, are we in train or eval model.
+
+    Returns:
+      out_logit: the output logits (before sigmoid).
+    """
+    hparams = self.hparams
+    with tf.variable_scope(
+        "discriminator",
+        initializer=tf.random_normal_initializer(stddev=0.02)):
+      batch_size, height, width = common_layers.shape_list(x)[:3]
+      # Mapping x from [bs, h, w, c] to [bs, 1]
+      net = tf.layers.conv2d(x, 64, (4, 4), strides=(2, 2),
+                             padding="SAME", name="d_conv1")
+      # [bs, h/2, w/2, 64]
+      net = lrelu(net)
+      net = tf.layers.conv2d(net, 128, (4, 4), strides=(2, 2),
+                             padding="SAME", name="d_conv2")
+      # [bs, h/4, w/4, 128]
+      if hparams.discriminator_batchnorm:
+        net = tf.layers.batch_normalization(net, training=is_training,
+                                            momentum=0.999, name="d_bn2")
+      net = lrelu(net)
+      size = height * width
+      net = tf.reshape(net, [batch_size, size * 8])  # [bs, h * w * 8]
+      net = tf.layers.dense(net, 1024, name="d_fc3")  # [bs, 1024]
+      if hparams.discriminator_batchnorm:
+        net = tf.layers.batch_normalization(net, training=is_training,
+                                            momentum=0.999, name="d_bn3")
+      net = lrelu(net)
+      return net
+
+  def unbottleneck(self, x, res_size, reuse=None):
+    with tf.variable_scope("unbottleneck", reuse=reuse):
       x = tf.layers.dense(x, res_size, name="dense")
       return x
 
@@ -115,8 +159,13 @@ def body(self, features):
       self._cur_bottleneck_tensor = b
       b = self.unbottleneck(b, common_layers.shape_list(x)[-1])
       b = common_layers.mix(b, x, hparams.bottleneck_warmup_steps, is_training)
+      if hparams.add_gan_loss:
+        # Add a purely sampled batch on which we'll compute the GAN loss.
+        g = self.unbottleneck(self.sample(), common_layers.shape_list(x)[-1],
+                              reuse=True)
+        b = tf.concat([g, b], axis=0)
       # With probability bottleneck_max_prob use the bottleneck, otherwise x.
-      if hparams.bottleneck_max_prob < 1.0:
+      if hparams.bottleneck_max_prob < -1.0:
         x = tf.where(tf.less(tf.random_uniform([]),
                              hparams.bottleneck_max_prob), b, x)
       else:
@@ -135,11 +184,29 @@ def body(self, features):
       return x, {"bottleneck_loss": 0.0}
     # Cut to the right size and mix before returning.
     res = x[:, :shape[1], :shape[2], :]
+    # Add GAN loss if requested.
+    gan_loss = 0.0
+    if hparams.add_gan_loss:
+      # Split back if we added a purely sampled batch.
+      res_gan, res = tf.split(res, 2, axis=0)
+      num_channels = self.hparams.problem.num_channels
+      res_rgb = common_layers.convert_real_to_rgb(tf.nn.sigmoid(
+          tf.layers.dense(res_gan, num_channels, name="gan_rgb")))
+      tf.summary.image("gan", tf.cast(res_rgb, tf.uint8), max_outputs=1)
+      orig_rgb = tf.to_float(features["targets_raw"])
+      def discriminate(x):
+        return self.discriminator(x, is_training=is_training)
+      gan_loss = common_layers.sliced_gan_loss(
+          orig_rgb, reverse_gradient(res_rgb),
+          discriminate, self.hparams.num_sliced_vecs)
+      gan_loss *= common_layers.inverse_lin_decay(
+          hparams.bottleneck_warmup_steps)
+    # Mix the final result and return.
     res = common_layers.mix(res, features["targets"],
                             hparams.bottleneck_warmup_steps // 2, is_training)
-    return res, {"bottleneck_loss": b_loss}
+    return res, {"bottleneck_loss": b_loss, "gan_loss": - gan_loss}
 
-  def sample(self, features=None):
+  def sample(self, features=None, shape=None):
     del features
     hp = self.hparams
     div_x = 2**hp.num_hidden_layers
@@ -255,4 +322,7 @@ def basic_autoencoder():
   hparams.add_hparam("bottleneck_max_prob", 1.0)
   hparams.add_hparam("sample_height", 32)
   hparams.add_hparam("sample_width", 32)
+  hparams.add_hparam("discriminator_batchnorm", True)
+  hparams.add_hparam("num_sliced_vecs", 4096)
+  hparams.add_hparam("add_gan_loss", False)
   return hparams
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 8cbfd1207..2c7d9e62d 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -281,8 +281,9 @@ def bottleneck(self, x, bottleneck_bits=None):  # pylint: disable=arguments-diff
       self.hparams.bottleneck_bits = old_bottleneck_bits
     return res, loss
 
-  def unbottleneck(self, x, res_size):
-    return discretization.parametrized_unbottleneck(x, res_size, self.hparams)
+  def unbottleneck(self, x, res_size, reuse=None):
+    with tf.variable_scope("unbottleneck", reuse=reuse):
+      return discretization.parametrized_unbottleneck(x, res_size, self.hparams)
 
   def sample(self, features=None):
     del features
@@ -541,7 +542,7 @@ def autoencoder_discrete_cifar():
   hparams.hidden_size = 256
   hparams.num_residual_layers = 4
   hparams.batch_size = 32
-  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_constant = 1.0
   hparams.dropout = 0.1
   return hparams
 
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index c8d57792c..a3410a40f 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -74,7 +74,7 @@ def ppo_discrete_action_base():
 def ppo_atari_base():
   """Atari base parameters."""
   hparams = ppo_discrete_action_base()
-  hparams.learning_rate = 16e-5
+  hparams.learning_rate = 4e-4
   hparams.num_agents = 5
   hparams.epoch_length = 200
   hparams.gae_gamma = 0.985
@@ -92,7 +92,7 @@ def ppo_atari_base():
 def ppo_pong_base():
   """Pong base parameters."""
   hparams = ppo_discrete_action_base()
-  hparams.learning_rate = 8e-5
+  hparams.learning_rate = 2e-4
   hparams.num_agents = 8
   hparams.epoch_length = 200
   hparams.gae_gamma = 0.985
@@ -113,7 +113,7 @@ def ppo_pong_base():
 def ppo_pong_ae_base():
   """Pong autoencoder base parameters."""
   hparams = ppo_pong_base()
-  hparams.learning_rate = 4e-5
+  hparams.learning_rate = 2e-4
   hparams.network = dense_bitwise_categorical_fun
   return hparams
 

From f3880ba51315aef8edb18c28f843b3fa38d00264 Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Wed, 20 Jun 2018 21:32:20 -0700
Subject: [PATCH 0170/2720] Adding word dropout for robustness in latents and
 keeping default at 0.0

PiperOrigin-RevId: 201468482
---
 .../models/research/transformer_nat.py        | 22 ++++++++++++-
 .../models/research/transformer_vae.py        | 32 ++++++++++++++++++-
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 078e96734..2b0701829 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -255,7 +255,14 @@ def ae_transformer_internal(inputs, targets, target_space, hparams, cache=None):
       targets,
       max_targets_len_from_inputs,
       final_length_divisible_by=2**hparams.num_compress_steps)
-  targets_c = compress(targets, hparams, "compress")
+  if hparams.word_dropout:
+    mask = tf.random_uniform(shape=common_layers.shape_list(targets),
+                             minval=0.0, maxval=1.0)
+    targets_noisy = tf.where(mask > hparams.word_dropout, targets,
+                             tf.zeros_like(targets))
+  else:
+    targets_noisy = targets
+  targets_c = compress(targets_noisy, hparams, "compress")
   if hparams.mode != tf.estimator.ModeKeys.PREDICT:
     # Compress and bottleneck.
     if hparams.denoising:
@@ -412,6 +419,7 @@ def transformer_nat_small():
   hparams.add_hparam("decay", 0.999)
   hparams.add_hparam("num_samples", 10)
   hparams.add_hparam("mask_startup_steps", 50000)
+  hparams.add_hparam("word_dropout", 0.0)
   hparams.add_hparam("denoising", False)
   return hparams
 
@@ -427,6 +435,18 @@ def transformer_nat_base():
   return hparams
 
 
+@registry.register_hparams
+def transformer_nat_base_drop():
+  """Set of hyperparameters."""
+  hparams = transformer_nat_small()
+  hparams.batch_size = 2048
+  hparams.hidden_size = 512
+  hparams.filter_size = 4096
+  hparams.num_hidden_layers = 6
+  hparams.word_dropout = 0.2
+  return hparams
+
+
 @registry.register_hparams
 def transformer_nat_big():
   """Set of hyperparameters."""
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 5300968cf..7ed8f1f76 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -675,7 +675,7 @@ def transformer_ae_small():
   hparams.add_hparam("noise_dev", 0.5)
   hparams.add_hparam("d_mix", 0.5)
   hparams.add_hparam("logit_normalization", True)
-  hparams.add_hparam("word_dropout", 0.1)
+  hparams.add_hparam("word_dropout", 0.0)
   # Bottleneck kinds supported: dense, vae, semhash, gumbel-softmax, dvq.
   hparams.add_hparam("bottleneck_kind", "semhash")
   hparams.add_hparam("num_blocks", 1)
@@ -803,6 +803,36 @@ def transformer_ae_base():
   return hparams
 
 
+@registry.register_hparams
+def transformer_ae_base_noatt():
+  """Set of hyperparameters."""
+  hparams = transformer_ae_small()
+  hparams.batch_size = 2048
+  hparams.hidden_size = 512
+  hparams.filter_size = 4096
+  hparams.num_hidden_layers = 6
+  hparams.do_attend_decompress = False
+  hparams.do_attend_decompress = False
+  hparams.word_dropout = 0.0
+  hparams.z_size = 12
+  return hparams
+
+
+@registry.register_hparams
+def transformer_ae_base_drop_noatt():
+  """Set of hyperparameters."""
+  hparams = transformer_ae_small()
+  hparams.batch_size = 2048
+  hparams.hidden_size = 512
+  hparams.filter_size = 4096
+  hparams.num_hidden_layers = 6
+  hparams.do_attend_decompress = False
+  hparams.do_attend_decompress = False
+  hparams.word_dropout = 0.2
+  hparams.z_size = 12
+  return hparams
+
+
 @registry.register_hparams
 def transformer_ae_a3():
   """Set of hyperparameters."""

From 988b96f562e750ce6f9ccb568c776e1163cc1ff6 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Wed, 20 Jun 2018 21:50:14 -0700
Subject: [PATCH 0171/2720] Remove denoising from nat

PiperOrigin-RevId: 201469688
---
 tensor2tensor/models/research/transformer_nat.py | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 2b0701829..7d4ed16fb 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -265,13 +265,8 @@ def ae_transformer_internal(inputs, targets, target_space, hparams, cache=None):
   targets_c = compress(targets_noisy, hparams, "compress")
   if hparams.mode != tf.estimator.ModeKeys.PREDICT:
     # Compress and bottleneck.
-    if hparams.denoising:
-      noise = tf.random_normal(
-          shape=common_layers.shape_list(targets_c), mean=0.0, stddev=0.1)
-    else:
-      noise = 0
     latents_discrete_hot, extra_loss = vq_discrete_bottleneck(
-        x=targets_c + noise, hparams=hparams)
+        x=targets_c, hparams=hparams)
     latents_dense = vq_discrete_unbottleneck(
         latents_discrete_hot, hparams=hparams)
     latents_dense = targets_c + tf.stop_gradient(latents_dense - targets_c)
@@ -420,7 +415,6 @@ def transformer_nat_small():
   hparams.add_hparam("num_samples", 10)
   hparams.add_hparam("mask_startup_steps", 50000)
   hparams.add_hparam("word_dropout", 0.0)
-  hparams.add_hparam("denoising", False)
   return hparams
 
 

From 3b40fc622d953c91e233febab78ee00f850541d1 Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Wed, 20 Jun 2018 22:26:51 -0700
Subject: [PATCH 0172/2720] Internal change

PiperOrigin-RevId: 201472643
---
 tensor2tensor/utils/decoding.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 4b6552dea..07f8ad6cc 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -551,12 +551,10 @@ def save_video(video, save_path_template):
 
   for i, frame in enumerate(video):
     save_path = save_path_template.format(i)
-    with tf.gfile.Open(save_path, "wb") as sp:
-      Image.fromarray(np.uint8(frame)).save(sp)
+    Image.fromarray(np.uint8(frame)).save(save_path)
 
 
 def show_and_save_image(img, save_path):
-  """Shows an image using matplotlib and saves it."""
   try:
     import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top
   except ImportError as e:
@@ -565,8 +563,7 @@ def show_and_save_image(img, save_path):
         "installed: %s", e)
     raise NotImplementedError("Image display and save not implemented.")
   plt.imshow(img)
-  with tf.gfile.Open(save_path, "wb") as sp:
-    plt.savefig(sp)
+  plt.savefig(save_path)
 
 
 def _get_sorted_inputs(filename, num_shards=1, delimiter="\n"):

From 54484df7047f668babafb0b2296310bd01f8b0fe Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 20 Jun 2018 22:27:20 -0700
Subject: [PATCH 0173/2720] Another file-handling fix.

PiperOrigin-RevId: 201472675
---
 tensor2tensor/utils/video_metrics.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index aabc76a0e..a03ed2309 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -81,7 +81,8 @@ def save_results(results, output_dir, problem_name):
   for name, array in six.iteritems(results):
     output_filename = "{}_{}.npy".format(problem_name, name)
     output_filename = os.path.join(output_dir, output_filename)
-    np.save(output_filename, array)
+    with tf.gfile.Open(output_filename, "wb") as fname:
+      np.save(fname, array)
 
 
 def compute_metrics(output_video, target_video):

From cca0c89dc2c45d9c37736a3f7b1da906054f3f32 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 21 Jun 2018 09:32:05 -0700
Subject: [PATCH 0174/2720] Cloud ML Engine needs the default compute engine to
 be set

PiperOrigin-RevId: 201537521
---
 docs/cloud_mlengine.md | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/docs/cloud_mlengine.md b/docs/cloud_mlengine.md
index 25673901e..8b40947db 100644
--- a/docs/cloud_mlengine.md
+++ b/docs/cloud_mlengine.md
@@ -7,7 +7,10 @@ you can easily launch Tensor2Tensor on it, including for hyperparameter tuning.
 # Launch
 
 It's the same `t2t-trainer` you know and love with the addition of the
-`--cloud_mlengine` flag, which by default will launch on a 1-GPU machine.
+`--cloud_mlengine` flag, which by default will launch on a 1-GPU machine
+in the default compute region. See the [docs for `gcloud compute`]
+(https://cloud.google.com/compute/docs/gcloud-compute/#set_default_zone_and_region_in_your_local_client)
+to learn how to set the default compute region.
 
 ```
 # Note that both the data dir and output dir have to be on GCS
@@ -42,6 +45,7 @@ possibly with multiple GPUs. Multi-machine setups are not yet supported out of
 the box with the `--cloud_mlengine` flag, though multi-machine should in
 principle work just fine. Contributions/testers welcome.
 
+
 ## `--t2t_usr_dir`
 
 Launching on Cloud ML Engine works with `--t2t_usr_dir` as well as long as the

From 93d752da188c133552867d524be900434e6fd7cf Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 21 Jun 2018 10:22:07 -0700
Subject: [PATCH 0175/2720] Enable fast greedy inference on TPU for Transformer
 model.

PiperOrigin-RevId: 201545466
---
 tensor2tensor/layers/common_attention.py |  19 +-
 tensor2tensor/layers/common_layers.py    |  45 ++-
 tensor2tensor/models/transformer.py      | 363 ++++++++++++++++++++++-
 tensor2tensor/models/transformer_test.py |  80 ++++-
 4 files changed, 491 insertions(+), 16 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index c464ec2e5..bcd2b1799 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2904,8 +2904,23 @@ def multihead_attention(query_antecedent,
       else:
         k = split_heads(k, num_heads)
         v = split_heads(v, num_heads)
-        k = cache["k"] = tf.concat([cache["k"], k], axis=2)
-        v = cache["v"] = tf.concat([cache["v"], v], axis=2)
+        decode_loop_step = kwargs.get("decode_loop_step")
+        if decode_loop_step is None:
+          k = cache["k"] = tf.concat([cache["k"], k], axis=2)
+          v = cache["v"] = tf.concat([cache["v"], v], axis=2)
+        else:
+          # Inplace update is required for inference on TPU.
+          # Inplace_ops only supports inplace_update on the first dimension.
+          # TODO(shibow): explore updating the entire Tensor instead of using
+          # inplace_ops to avoid the transposes.
+          tmp_k = tf.transpose(cache["k"], perm=[2, 0, 1, 3])
+          tmp_k = common_layers.tf_inplace_ops().alias_inplace_update(
+              tmp_k, decode_loop_step, tf.squeeze(k, axis=2))
+          k = cache["k"] = tf.transpose(tmp_k, perm=[1, 2, 0, 3])
+          tmp_v = tf.transpose(cache["v"], perm=[2, 0, 1, 3])
+          tmp_v = common_layers.tf_inplace_ops().alias_inplace_update(
+              tmp_v, decode_loop_step, tf.squeeze(v, axis=2))
+          v = cache["v"] = tf.transpose(tmp_v, perm=[1, 2, 0, 3])
 
     q = split_heads(q, num_heads)
     if cache is None:
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index b330998fb..3c5f108f1 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -36,6 +36,12 @@
 allow_defun = False
 
 
+# Lazy load inplace_ops
+def tf_inplace_ops():
+  from tensorflow.python.ops import inplace_ops  # pylint: disable=g-import-not-at-top
+  return inplace_ops
+
+
 @function.Defun(
     python_grad_func=lambda x, dy: tf.convert_to_tensor(dy),
     shape_func=lambda op: [op.inputs[0].get_shape()])
@@ -1512,14 +1518,47 @@ def conv_relu_conv(inputs,
                    nonpadding_mask=None,
                    dropout=0.0,
                    name=None,
-                   cache=None):
-  """Hidden layer with RELU activation followed by linear projection."""
+                   cache=None,
+                   decode_loop_step=None):
+  """Hidden layer with RELU activation followed by linear projection.
+
+  Args:
+    inputs: A tensor.
+    filter_size: An integer.
+    output_size: An integer.
+    first_kernel_size: An integer.
+    second_kernel_size: An integer.
+    padding: A string.
+    nonpadding_mask: A tensor.
+    dropout: A float.
+    name: A string.
+    cache: A dict, containing Tensors which are the results of previous
+        attentions, used for fast decoding.
+    decode_loop_step: An integer, step number of the decoding loop.
+        Only used for inference on TPU. If it is not None, the function
+        will do inplace update for the cache instead of concatenating the
+        current result to the cache.
+
+  Returns:
+    A Tensor.
+  """
   with tf.variable_scope(name, "conv_relu_conv", [inputs]):
     inputs = maybe_zero_out_padding(
         inputs, first_kernel_size, nonpadding_mask)
 
     if cache:
-      inputs = cache["f"] = tf.concat([cache["f"], inputs], axis=1)
+      if decode_loop_step is None:
+        inputs = cache["f"] = tf.concat([cache["f"], inputs], axis=1)
+      else:
+        # Inplace update is required for inference on TPU.
+        # Inplace_ops only supports inplace_update on the first dimension.
+        # TODO(shibow): explore updating the entire Tensor instead of using
+        # inplace_ops to avoid the transposes.
+        tmp_f = tf.transpose(cache["f"], perm=[1, 0, 2])
+        tmp_f = tf_inplace_ops().alias_inplace_update(
+            tmp_f, decode_loop_step * tf.shape(inputs)[1],
+            tf.transpose(inputs, perm=[1, 0, 2]))
+        inputs = cache["f"] = tf.transpose(tmp_f, perm=[1, 0, 2])
       inputs = cache["f"] = inputs[:, -first_kernel_size:, :]
 
     h = tpu_conv1d(inputs, filter_size, first_kernel_size, padding=padding,
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 4e9d6407e..77496b5ff 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -94,6 +94,7 @@ def decode(self,
              decoder_self_attention_bias,
              hparams,
              cache=None,
+             decode_loop_step=None,
              nonpadding=None,
              losses=None):
     """Decode Transformer outputs from encoder representation.
@@ -110,6 +111,8 @@ def decode(self,
       hparams: hyperparameters for model.
       cache: dict, containing tensors which are the results of previous
           attentions, used for fast decoding.
+      decode_loop_step: An integer, step number of the decoding loop.
+          Only used for inference on TPU.
       nonpadding: optional Tensor with shape [batch_size, decoder_length]
       losses: optional list onto which to append extra training losses
 
@@ -126,6 +129,7 @@ def decode(self,
         encoder_decoder_attention_bias,
         hparams,
         cache=cache,
+        decode_loop_step=decode_loop_step,
         nonpadding=nonpadding,
         save_weights_to=self.attention_weights,
         losses=losses)
@@ -218,7 +222,8 @@ def _greedy_infer(self, features, decode_length, use_tpu=False):
     if self._target_modality_is_real:
       return  super(Transformer, self)._greedy_infer(features, decode_length)
     with tf.variable_scope(self.name):
-      return self._fast_decode(features, decode_length)
+      return (self._fast_decode_tpu(features, decode_length) if use_tpu else
+              self._fast_decode(features, decode_length))
 
   def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
     """Beam search decoding.
@@ -251,6 +256,213 @@ def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
       return self._fast_decode(features, decode_length, beam_size, top_beams,
                                alpha)
 
+  def _fast_decode_tpu(self,
+                       features,
+                       decode_length,
+                       beam_size=1):
+    """Fast decoding.
+
+    Implements only greedy decoding on TPU.
+
+    Args:
+      features: A map of string to model features.
+      decode_length: An integer, how many additional timesteps to decode.
+      beam_size: An integer, number of beams.
+
+    Returns:
+      A dict of decoding results {
+          "outputs": integer `Tensor` of decoded ids of shape
+              [batch_size, <= decode_length]
+          "scores": decoding log probs from the beam search,
+              None if using greedy decoding (beam_size=1)
+      }.
+
+    Raises:
+      NotImplementedError: If there are multiple data shards or beam_size > 1.
+    """
+    if self._num_datashards != 1:
+      raise NotImplementedError("Fast decoding only supports a single shard.")
+    dp = self._data_parallelism
+    hparams = self._hparams
+    target_modality = self._problem_hparams.target_modality
+
+    if self.has_input:
+      inputs = features["inputs"]
+      if target_modality.is_class_modality:
+        decode_length = 1
+      else:
+        decode_length = (
+            common_layers.shape_list(inputs)[1] + features.get(
+                "decode_length", decode_length))
+
+      # TODO(llion): Clean up this reshaping logic.
+      inputs = tf.expand_dims(inputs, axis=1)
+      if len(inputs.shape) < 5:
+        inputs = tf.expand_dims(inputs, axis=4)
+      s = common_layers.shape_list(inputs)
+      batch_size = s[0]
+      inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
+      # _shard_features called to ensure that the variable names match
+      inputs = self._shard_features({"inputs": inputs})["inputs"]
+      input_modality = self._problem_hparams.input_modality["inputs"]
+      with tf.variable_scope(input_modality.name):
+        inputs = input_modality.bottom_sharded(inputs, dp)
+      with tf.variable_scope("body"):
+        encoder_output, encoder_decoder_attention_bias = dp(
+            self.encode,
+            inputs,
+            features["target_space_id"],
+            hparams,
+            features=features)
+      encoder_output = encoder_output[0]
+      encoder_decoder_attention_bias = encoder_decoder_attention_bias[0]
+      partial_targets = None
+    else:
+      # The problem has no inputs.
+      encoder_output = None
+      encoder_decoder_attention_bias = None
+
+      # Prepare partial targets.
+      # In either features["inputs"] or features["targets"].
+      # We force the outputs to begin with these sequences.
+      partial_targets = features.get("inputs")
+      if partial_targets is None:
+        partial_targets = features["targets"]
+      assert partial_targets is not None
+      partial_targets = common_layers.expand_squeeze_to_nd(partial_targets, 2)
+      partial_targets = tf.to_int64(partial_targets)
+      partial_targets_shape = common_layers.shape_list(partial_targets)
+      partial_targets_length = partial_targets_shape[1]
+      decode_length = (
+          partial_targets_length + features.get("decode_length", decode_length))
+      batch_size = partial_targets_shape[0]
+
+    if hparams.pos == "timing":
+      timing_signal = common_attention.get_timing_signal_1d(
+          decode_length + 1, hparams.hidden_size)
+
+    def preprocess_targets(targets, i):
+      """Performs preprocessing steps on the targets to prepare for the decoder.
+
+      This includes:
+        - Embedding the ids.
+        - Flattening to 3D tensor.
+        - Optionally adding timing signals.
+
+      Args:
+        targets: A tensor, inputs ids to the decoder. [batch_size, 1].
+        i: An integer, Step number of the decoding loop.
+
+      Returns:
+        A tensor, processed targets [batch_size, 1, hidden_dim].
+      """
+      # _shard_features called to ensure that the variable names match
+      targets = self._shard_features({"targets": targets})["targets"]
+      with tf.variable_scope(target_modality.name):
+        targets = target_modality.targets_bottom_sharded(targets, dp)[0]
+      targets = common_layers.flatten4d3d(targets)
+
+      # TODO(llion): Explain! Is this even needed?
+      targets = tf.cond(
+          tf.equal(i, 0), lambda: tf.zeros_like(targets), lambda: targets)
+
+      if hparams.pos == "timing":
+        timing_signal_shape = timing_signal.shape.as_list()
+        targets += tf.slice(timing_signal, [0, i, 0],
+                            [timing_signal_shape[0], 1, timing_signal_shape[2]])
+      return targets
+
+    decoder_self_attention_bias = (
+        common_attention.attention_bias_lower_triangle(decode_length))
+    if hparams.proximity_bias:
+      decoder_self_attention_bias += common_attention.attention_bias_proximal(
+          decode_length)
+
+    def symbols_to_logits_tpu_fn(ids, i, cache):
+      """Go from ids to logits for next symbol on TPU.
+
+      Args:
+        ids: A tensor, symbol IDs.
+        i: An integer, step number of the decoding loop. Only used for inference
+            on TPU.
+        cache: A dict, containing tensors which are the results of previous
+            attentions, used for fast decoding.
+
+      Returns:
+        ret: A tensor, computed logits.
+        cache: A dict, containing tensors which are the results of previous
+            attentions, used for fast decoding.
+      """
+      ids = ids[:, -1:]
+      targets = tf.expand_dims(tf.expand_dims(ids, axis=2), axis=3)
+      targets = preprocess_targets(targets, i)
+
+      bias_shape = decoder_self_attention_bias.shape.as_list()
+      bias = tf.slice(decoder_self_attention_bias, [0, 0, i, 0],
+                      [bias_shape[0], bias_shape[1], 1, bias_shape[3]])
+
+      bias_padding = tf.fill([bias_shape[0], bias_shape[1], 1], -1e9)
+      tmp_bias = tf.transpose(bias, perm=[3, 0, 1, 2])
+      bias_index = i + 1
+      while_condition = lambda bias_index, _: tf.less(bias_index, decode_length)
+
+      def while_body(bias_index, tmp_bias):
+        tmp_bias = common_layers.tf_inplace_ops().alias_inplace_update(
+            tmp_bias, bias_index, bias_padding)
+        return bias_index + 1, tmp_bias
+
+      _, tmp_bias = tf.while_loop(
+          while_condition, while_body, (bias_index, tmp_bias))
+      bias = tf.transpose(tmp_bias, perm=[1, 2, 3, 0])
+
+      with tf.variable_scope("body"):
+        body_outputs = dp(
+            self.decode,
+            targets,
+            cache.get("encoder_output"),
+            cache.get("encoder_decoder_attention_bias"),
+            bias,
+            hparams,
+            cache,
+            i,
+            nonpadding=features_to_nonpadding(features, "targets"))
+
+      with tf.variable_scope(target_modality.name):
+        logits = target_modality.top_sharded(body_outputs, None, dp)[0]
+
+      ret = tf.squeeze(logits, axis=[1, 2, 3])
+      if partial_targets is not None:
+        # If the position is within the given partial targets, we alter the
+        # logits to always return those values.
+        # A faster approach would be to process the partial targets in one
+        # iteration in order to fill the corresponding parts of the cache.
+        # This would require broader changes, though.
+        vocab_size = tf.shape(ret)[1]
+
+        def forced_logits():
+          return tf.one_hot(
+              tf.tile(
+                  tf.slice(partial_targets, [0, i],
+                           [partial_targets.shape.as_list()[0], 1]),
+                  [beam_size]), vocab_size, 0.0, -1e9)
+
+        ret = tf.cond(
+            tf.less(i, partial_targets_length), forced_logits, lambda: ret)
+      return ret, cache
+
+    ret = fast_decode_tpu(
+        encoder_output=encoder_output,
+        encoder_decoder_attention_bias=encoder_decoder_attention_bias,
+        symbols_to_logits_fn=symbols_to_logits_tpu_fn,
+        hparams=hparams,
+        decode_length=decode_length,
+        beam_size=beam_size,
+        batch_size=batch_size,
+        force_decode_length=self._decode_hparams.force_decode_length)
+    if partial_targets is not None:
+      ret["outputs"] = ret["outputs"][:, partial_targets_length:]
+    return ret
+
   def _fast_decode(self,
                    features,
                    decode_length,
@@ -438,6 +650,140 @@ def forced_logits():
     return ret
 
 
+def fast_decode_tpu(encoder_output,
+                    encoder_decoder_attention_bias,
+                    symbols_to_logits_fn,
+                    hparams,
+                    decode_length,
+                    beam_size=1,
+                    eos_id=beam_search.EOS_ID,
+                    batch_size=None,
+                    force_decode_length=False):
+  """Given encoder output and a symbols to logits function, does fast decoding.
+
+  Implements only greedy decoding for TPU.
+
+  Args:
+    encoder_output: A tensor, output from encoder.
+    encoder_decoder_attention_bias: A tensor, bias for use in encoder-decoder
+        attention.
+    symbols_to_logits_fn: Incremental decoding, function mapping triple
+        `(ids, step, cache)` to symbol logits.
+    hparams: Run hyperparameters.
+    decode_length: An integer, how many additional timesteps to decode.
+    beam_size: An integer, number of beams.
+    eos_id: End-of-sequence symbol.
+    batch_size: An integer, must be passed if there is no input.
+    force_decode_length: A bool, whether to force the full decode length, or if
+        False, stop when all beams hit eos_id.
+
+  Returns:
+      A dict of decoding results {
+          "outputs": integer `Tensor` of decoded ids of shape
+              [batch_size, <= decode_length]
+          "scores": decoding log probs from the beam search,
+              None if using greedy decoding (beam_size=1)
+      }.
+
+  Raises:
+     NotImplementedError: If beam size > 1.
+  """
+  if encoder_output is not None:
+    batch_size = common_layers.shape_list(encoder_output)[0]
+
+  key_channels = hparams.attention_key_channels or hparams.hidden_size
+  value_channels = hparams.attention_value_channels or hparams.hidden_size
+  num_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
+
+  cache = {
+      "layer_%d" % layer: {
+          "k":
+          common_attention.split_heads(
+              tf.zeros([batch_size, decode_length, key_channels]),
+              hparams.num_heads),
+          "v":
+          common_attention.split_heads(
+              tf.zeros([batch_size, decode_length, value_channels]),
+              hparams.num_heads),
+          "f":
+          tf.zeros([batch_size, decode_length, hparams.hidden_size]),
+      } for layer in range(num_layers)
+  }
+
+  if encoder_output is not None:
+    for layer in range(num_layers):
+      layer_name = "layer_%d" % layer
+      with tf.variable_scope(
+          "body/decoder/%s/encdec_attention/multihead_attention" % layer_name):
+        k_encdec = common_attention.compute_attention_component(
+            encoder_output, key_channels, name="k")
+        k_encdec = common_attention.split_heads(k_encdec, hparams.num_heads)
+        v_encdec = common_attention.compute_attention_component(
+            encoder_output, value_channels, name="v")
+        v_encdec = common_attention.split_heads(v_encdec, hparams.num_heads)
+      cache[layer_name]["k_encdec"] = k_encdec
+      cache[layer_name]["v_encdec"] = v_encdec
+
+    cache["encoder_output"] = encoder_output
+    cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias
+
+  if beam_size > 1:  # Beam Search
+    raise NotImplementedError("Beam search inference on TPU is not supported")
+
+  # Greedy
+  def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
+    """One step of greedy decoding."""
+    logits, cache = symbols_to_logits_fn(next_id, i, cache)
+    log_probs = common_layers.log_prob_from_logits(logits)
+    temperature = (0.0 if hparams.sampling_method == "argmax" else
+                   hparams.sampling_temp)
+    next_id = common_layers.sample_with_temperature(logits, temperature)
+    hit_eos |= tf.equal(next_id, eos_id)
+
+    log_prob_indices = tf.stack(
+        [tf.range(tf.to_int64(batch_size)), next_id], axis=1)
+    log_prob += tf.gather_nd(log_probs, log_prob_indices)
+
+    next_id = tf.expand_dims(next_id, axis=1)
+    decoded_ids = tf.transpose(decoded_ids)
+    decoded_ids = common_layers.tf_inplace_ops().alias_inplace_update(
+        decoded_ids, i, tf.squeeze(next_id, axis=1))
+    decoded_ids = tf.transpose(decoded_ids)
+    return i + 1, hit_eos, next_id, decoded_ids, cache, log_prob
+
+  def is_not_finished(i, hit_eos, *_):
+    finished = i >= decode_length
+    if not force_decode_length:
+      finished |= tf.reduce_all(hit_eos)
+    return tf.logical_not(finished)
+
+  decoded_ids = tf.zeros([batch_size, decode_length], dtype=tf.int64)
+  hit_eos = tf.fill([batch_size], False)
+  next_id = tf.zeros([batch_size, 1], dtype=tf.int64)
+  initial_log_prob = tf.zeros([batch_size], dtype=tf.float32)
+
+  def compute_cache_shape_invariants(tensor):
+    return tf.TensorShape(tensor.shape.as_list())
+
+  _, _, _, decoded_ids, _, log_prob = tf.while_loop(
+      is_not_finished,
+      inner_loop, [
+          tf.constant(0), hit_eos, next_id, decoded_ids, cache,
+          initial_log_prob
+      ],
+      shape_invariants=[
+          tf.TensorShape([]),
+          tf.TensorShape([batch_size]),
+          tf.TensorShape([batch_size, 1]),
+          tf.TensorShape([batch_size, decode_length]),
+          nest.map_structure(compute_cache_shape_invariants, cache),
+          tf.TensorShape([batch_size]),
+      ])
+  scores = log_prob
+
+  return {"outputs": decoded_ids, "scores": scores}
+
+
 def fast_decode(encoder_output,
                 encoder_decoder_attention_bias,
                 symbols_to_logits_fn,
@@ -870,6 +1216,7 @@ def transformer_decoder(decoder_input,
                         encoder_decoder_attention_bias,
                         hparams,
                         cache=None,
+                        decode_loop_step=None,
                         name="decoder",
                         nonpadding=None,
                         save_weights_to=None,
@@ -887,6 +1234,8 @@ def transformer_decoder(decoder_input,
     hparams: hyperparameters for model
     cache: dict, containing tensors which are the results of previous
         attentions, used for fast decoding.
+    decode_loop_step: An integer, step number of the decoding loop.
+        Only used for inference on TPU.
     name: a string
     nonpadding: optional Tensor with shape [batch_size, encoder_length]
       indicating what positions are not padding.  This is used
@@ -927,7 +1276,8 @@ def transformer_decoder(decoder_input,
               cache=layer_cache,
               make_image_summary=make_image_summary,
               dropout_broadcast_dims=attention_dropout_broadcast_dims,
-              max_length=hparams.get("max_length"))
+              max_length=hparams.get("max_length"),
+              decode_loop_step=decode_loop_step)
           x = common_layers.layer_postprocess(x, y, hparams)
         if encoder_output is not None:
           with tf.variable_scope("encdec_attention"):
@@ -953,7 +1303,8 @@ def transformer_decoder(decoder_input,
               conv_padding="LEFT",
               nonpadding_mask=nonpadding,
               losses=losses,
-              cache=layer_cache)
+              cache=layer_cache,
+              decode_loop_step=decode_loop_step)
           x = common_layers.layer_postprocess(x, y, hparams)
     # if normalization is done in layer_preprocess, then it should also be done
     # on the output, since the output can grow very large, being the sum of
@@ -968,6 +1319,7 @@ def transformer_ffn_layer(x,
                           nonpadding_mask=None,
                           losses=None,
                           cache=None,
+                          decode_loop_step=None,
                           readout_filter_size=0):
   """Feed-forward layer in the transformer.
 
@@ -985,6 +1337,8 @@ def transformer_ffn_layer(x,
     losses: optional list onto which to append extra training losses
     cache: dict, containing tensors which are the results of previous
         attentions, used for fast decoding.
+    decode_loop_step: An integer, step number of the decoding loop.
+        Only used for inference on TPU.
     readout_filter_size: if it's greater than 0, then it will be used instead of
       filter_size
 
@@ -1030,7 +1384,8 @@ def transformer_ffn_layer(x,
         padding=conv_padding,
         nonpadding_mask=nonpadding_mask,
         dropout=hparams.relu_dropout,
-        cache=cache)
+        cache=cache,
+        decode_loop_step=decode_loop_step)
   elif ffn_layer == "parameter_attention":
     return common_attention.parameter_attention(
         x, hparams.parameter_attention_key_channels or hparams.hidden_size,
diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index f97602a40..7df6d8776 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -31,6 +31,12 @@
 VOCAB_SIZE = 10
 
 
+def tf_version_has_inplace_ops():
+  # Available in TF 1.8+
+  major, minor = [int(el) for el in tf.__version__.split(".")[:2]]
+  return major > 1 or (major == 1 and minor >= 8)
+
+
 def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
               has_input=True, model_cls=transformer.Transformer):
   if hparams is None:
@@ -240,13 +246,16 @@ def testTransformerWithEncoderDecoderAttentionLoss(self):
       res = session.run(extra_loss["attention_loss"])
     self.assertEqual(res.shape, ())
 
-  def testGreedySlowTPUVsNonTPU(self):
-    # Only works with TF 1.8+
-    # Version string can take the following form: "1.9.0-rc0"
-    major_str, minor_str, unused_rest = tf.__version__.split(".", 3)
-    major, minor = int(major_str), int(minor_str)
-    if major < 1 or (major == 1 and minor < 8):
-      return
+  def _create_greedy_infer_model(self, decode_length):
+    """Creates model for greedy inference testing.
+
+    Args:
+      decode_length: An integer, the decode length used for test.
+
+    Returns:
+      model: A t2t model.
+      features: An map of string to tensor.
+    """
     model, features = get_model(transformer.transformer_small())
 
     decode_length = 3
@@ -266,6 +275,16 @@ def testGreedySlowTPUVsNonTPU(self):
 
     model.set_mode(tf.estimator.ModeKeys.PREDICT)
 
+    return model, features
+
+  def testGreedySlowTPUVsNonTPU(self):
+    if not tf_version_has_inplace_ops():
+      return
+
+    decode_length = 3
+
+    model, features = self._create_greedy_infer_model(decode_length)
+
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
       slow_result_non_tpu = model._slow_greedy_infer(
           features, decode_length)["outputs"]
@@ -283,6 +302,53 @@ def testGreedySlowTPUVsNonTPU(self):
                      (BATCH_SIZE, INPUT_LENGTH + decode_length))
     self.assertAllClose(slow_tpu_res, slow_non_tpu_res)
 
+  def testGreedyFastTPUVsNonTPU(self):
+    if not tf_version_has_inplace_ops():
+      return
+
+    decode_length = 3
+
+    model, features = self._create_greedy_infer_model(decode_length)
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+      fast_result_non_tpu = model._greedy_infer(
+          features, decode_length, use_tpu=False)["outputs"]
+
+      fast_result_tpu = model._greedy_infer(
+          features, decode_length, use_tpu=True)["outputs"]
+
+    with self.test_session():
+      fast_non_tpu_res = fast_result_non_tpu.eval()
+      fast_tpu_res = fast_result_tpu.eval()
+
+    self.assertEqual(fast_tpu_res.shape,
+                     (BATCH_SIZE, INPUT_LENGTH + decode_length))
+    self.assertAllClose(fast_tpu_res, fast_non_tpu_res)
+
+  def testGreedyTPUSlowVsFast(self):
+    if not tf_version_has_inplace_ops():
+      return
+
+    decode_length = 3
+
+    model, features = self._create_greedy_infer_model(decode_length)
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+      slow_result = model._slow_greedy_infer_tpu(
+          features, decode_length)["outputs"]
+      slow_result = tf.squeeze(slow_result, axis=[2, 3])
+
+      fast_result = model._greedy_infer(
+          features, decode_length, use_tpu=True)["outputs"]
+
+    with self.test_session():
+      slow_res = slow_result.eval()
+      fast_res = fast_result.eval()
+
+    self.assertEqual(fast_res.shape,
+                     (BATCH_SIZE, INPUT_LENGTH + decode_length))
+    self.assertAllClose(fast_res, slow_res)
+
 
 class TransformerScorerTest(tf.test.TestCase):
 

From 830350d4ef1f6b659b42e9ffcb50361bdc1d1c47 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 21 Jun 2018 10:32:07 -0700
Subject: [PATCH 0176/2720] Use tf.gfile for saving images/videos.

PiperOrigin-RevId: 201547242
---
 tensor2tensor/utils/decoding.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 07f8ad6cc..4b6552dea 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -551,10 +551,12 @@ def save_video(video, save_path_template):
 
   for i, frame in enumerate(video):
     save_path = save_path_template.format(i)
-    Image.fromarray(np.uint8(frame)).save(save_path)
+    with tf.gfile.Open(save_path, "wb") as sp:
+      Image.fromarray(np.uint8(frame)).save(sp)
 
 
 def show_and_save_image(img, save_path):
+  """Shows an image using matplotlib and saves it."""
   try:
     import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top
   except ImportError as e:
@@ -563,7 +565,8 @@ def show_and_save_image(img, save_path):
         "installed: %s", e)
     raise NotImplementedError("Image display and save not implemented.")
   plt.imshow(img)
-  plt.savefig(save_path)
+  with tf.gfile.Open(save_path, "wb") as sp:
+    plt.savefig(sp)
 
 
 def _get_sorted_inputs(filename, num_shards=1, delimiter="\n"):

From e399e9f8e2ee93d46985ff6679a1ac5ec72d2c19 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Thu, 21 Jun 2018 12:00:53 -0700
Subject: [PATCH 0177/2720] Fix bug, change to reduce_mean for em
 discretization

PiperOrigin-RevId: 201562944
---
 tensor2tensor/layers/discretization.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index bcb534d6c..862187c28 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -706,7 +706,7 @@ def vq_nearest_neighbor(x, means, soft_em=False, num_samples=10):
     x_means_idx = tf.multinomial(-dist, num_samples=num_samples)
     x_means_hot = tf.one_hot(
         x_means_idx, depth=common_layers.shape_list(means)[0])
-    x_means_hot = tf.reduce_sum(x_means_hot, axis=1)
+    x_means_hot = tf.reduce_mean(x_means_hot, axis=1)
   else:
     x_means_idx = tf.argmax(-dist, axis=-1)
     x_means_hot = tf.one_hot(x_means_idx, bottleneck_size)

From 5e0d471b9c4abe24164f369c3ce058cad0c79560 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Thu, 21 Jun 2018 15:27:26 -0700
Subject: [PATCH 0178/2720] Add some more compatibility options to transformer.

PiperOrigin-RevId: 201597637
---
 tensor2tensor/layers/common_attention.py | 80 ++++++++++++++++++++----
 tensor2tensor/models/transformer.py      | 64 +++++++++++++------
 2 files changed, 113 insertions(+), 31 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index bcd2b1799..4881dd8f5 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -644,6 +644,27 @@ def add_timing_signal_nd(x, min_timescale=1.0, max_timescale=1.0e4):
   return x
 
 
+@expert_utils.add_name_scope()
+def add_positional_embedding(x, max_length, name, positions=None):
+  """Add positional embedding.
+
+  Args:
+    x: a Tensor with shape [batch, length, depth]
+    max_length: an integer.  static maximum size of any dimension.
+    name: a name for this layer.
+    positions: an optional tensor with shape [batch, length]
+
+  Returns:
+    a Tensor the same shape as x.
+  """
+  _, length, depth = common_layers.shape_list(x)
+  var = tf.get_variable(name, [max_length, depth])
+  if positions is not None:
+    return x + tf.gather(var, tf.to_int32(positions))
+  else:
+    return x + tf.expand_dims(tf.slice(var, [0, 0], [length, -1]), 0)
+
+
 @expert_utils.add_name_scope()
 def add_positional_embedding_nd(x, max_length, name):
   """Add n-dimensional positional embedding.
@@ -2713,7 +2734,8 @@ def compute_attention_component(antecedent,
                                 total_depth,
                                 filter_width=1,
                                 padding="VALID",
-                                name="c"):
+                                name="c",
+                                vars_3d_num_heads=None):
   """Computes attention compoenent (query, key or value).
 
   Args:
@@ -2723,10 +2745,25 @@ def compute_attention_component(antecedent,
       component to be.
     padding: One of "VALID", "SAME" or "LEFT". Default is VALID: No padding.
     name: a string specifying scope name.
+    vars_3d_num_heads: an optional integer (if we want to use 3d variables)
 
   Returns:
     c : [batch, length, depth] tensor
   """
+  if vars_3d_num_heads:
+    assert filter_width == 1
+    input_depth = antecedent.get_shape().as_list()[-1]
+    depth_per_head = total_depth // vars_3d_num_heads
+    initializer_stddev = input_depth ** -0.5
+    if "q" in name:
+      initializer_stddev *= depth_per_head ** -0.5
+    var = tf.get_variable(
+        name, [input_depth,
+               vars_3d_num_heads,
+               total_depth // vars_3d_num_heads],
+        initializer=tf.random_normal_initializer(stddev=initializer_stddev))
+    var = tf.reshape(var, [input_depth, total_depth])
+    return tf.tensordot(antecedent, var, axes=1)
   if filter_width == 1:
     return common_layers.dense(
         antecedent, total_depth, use_bias=False, name=name)
@@ -2742,7 +2779,8 @@ def compute_qkv(query_antecedent,
                 q_filter_width=1,
                 kv_filter_width=1,
                 q_padding="VALID",
-                kv_padding="VALID"):
+                kv_padding="VALID",
+                vars_3d_num_heads=None):
   """Computes query, key and value.
 
   Args:
@@ -2755,18 +2793,25 @@ def compute_qkv(query_antecedent,
     to be.
     q_padding: One of "VALID", "SAME" or "LEFT". Default is VALID: No padding.
     kv_padding: One of "VALID", "SAME" or "LEFT". Default is VALID: No padding.
+    vars_3d_num_heads: an optional (if we want to use 3d variables)
 
   Returns:
     q, k, v : [batch, length, depth] tensors
   """
   if memory_antecedent is None:
     memory_antecedent = query_antecedent
-  q = compute_attention_component(query_antecedent, total_key_depth,
-                                  q_filter_width, q_padding, "q")
-  k = compute_attention_component(memory_antecedent, total_key_depth,
-                                  kv_filter_width, kv_padding, "k")
-  v = compute_attention_component(memory_antecedent, total_value_depth,
-                                  kv_filter_width, kv_padding, "v")
+  q = compute_attention_component(
+      query_antecedent, total_key_depth,
+      q_filter_width, q_padding, "q",
+      vars_3d_num_heads=vars_3d_num_heads)
+  k = compute_attention_component(
+      memory_antecedent, total_key_depth,
+      kv_filter_width, kv_padding, "k",
+      vars_3d_num_heads=vars_3d_num_heads)
+  v = compute_attention_component(
+      memory_antecedent, total_value_depth,
+      kv_filter_width, kv_padding, "v",
+      vars_3d_num_heads=vars_3d_num_heads)
   return q, k, v
 
 
@@ -2797,6 +2842,7 @@ def multihead_attention(query_antecedent,
                         max_length=None,
                         adjacency_matrix=None,
                         num_edge_types=5,
+                        vars_3d=False,
                         **kwargs):
   """Multihead scaled-dot-product attention with input/output transformations.
 
@@ -2848,6 +2894,7 @@ def multihead_attention(query_antecedent,
     adjacency_matrix: an optional tensor of shape [batch, len_q, len_q]
       containing edge vectors for attention
     num_edge_types: number of edge types, an int
+    vars_3d: use 3-dimensional variables for input/output transformations
     **kwargs (dict): Parameters for the attention function
 
   Caching:
@@ -2877,13 +2924,15 @@ def multihead_attention(query_antecedent,
   if total_value_depth % num_heads != 0:
     raise ValueError("Value depth (%d) must be divisible by the number of "
                      "attention heads (%d)." % (total_value_depth, num_heads))
+  vars_3d_num_heads = num_heads if vars_3d else None
   with tf.variable_scope(name, default_name="multihead_attention",
                          values=[query_antecedent, memory_antecedent]):
 
     if cache is None or memory_antecedent is None:
       q, k, v = compute_qkv(query_antecedent, memory_antecedent,
                             total_key_depth, total_value_depth, q_filter_width,
-                            kv_filter_width, q_padding, kv_padding)
+                            kv_filter_width, q_padding, kv_padding,
+                            vars_3d_num_heads=vars_3d_num_heads)
     if cache is not None:
       if attention_type != "dot_product":
         # TODO(petershaw): Support caching when using relative position
@@ -2928,7 +2977,8 @@ def multihead_attention(query_antecedent,
       v = split_heads(v, num_heads)
 
     key_depth_per_head = total_key_depth // num_heads
-    q *= key_depth_per_head**-0.5
+    if not vars_3d:
+      q *= key_depth_per_head**-0.5
 
     additional_returned_value = None
     if callable(attention_type):  # Generic way to extend multihead_attention
@@ -2993,8 +3043,14 @@ def multihead_attention(query_antecedent,
     # Set last dim specifically.
     x.set_shape(x.shape.as_list()[:-1] + [total_value_depth])
 
-    x = common_layers.dense(
-        x, output_depth, use_bias=False, name="output_transform")
+    if vars_3d:
+      o_var = tf.get_variable(
+          "o", [num_heads, total_value_depth // num_heads, output_depth])
+      o_var = tf.reshape(o_var, [total_value_depth, output_depth])
+      x = tf.tensordot(x, o_var, axes=1)
+    else:
+      x = common_layers.dense(
+          x, output_depth, use_bias=False, name="output_transform")
     if additional_returned_value is not None:
       return x, additional_returned_value
     return x
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 77496b5ff..203e8709a 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -338,8 +338,14 @@ def _fast_decode_tpu(self,
       batch_size = partial_targets_shape[0]
 
     if hparams.pos == "timing":
-      timing_signal = common_attention.get_timing_signal_1d(
+      positional_encoding = common_attention.get_timing_signal_1d(
           decode_length + 1, hparams.hidden_size)
+    elif hparams.pos == "emb":
+      positional_encoding = common_attention.add_positional_embedding(
+          tf.zeros([1, decode_length + 1, hparams.hidden_size]),
+          hparams.max_length, "targets_positional_embedding", None)
+    else:
+      positional_encoding = None
 
     def preprocess_targets(targets, i):
       """Performs preprocessing steps on the targets to prepare for the decoder.
@@ -366,10 +372,11 @@ def preprocess_targets(targets, i):
       targets = tf.cond(
           tf.equal(i, 0), lambda: tf.zeros_like(targets), lambda: targets)
 
-      if hparams.pos == "timing":
-        timing_signal_shape = timing_signal.shape.as_list()
-        targets += tf.slice(timing_signal, [0, i, 0],
-                            [timing_signal_shape[0], 1, timing_signal_shape[2]])
+      if positional_encoding is not None:
+        positional_encoding_shape = positional_encoding.shape.as_list()
+        targets += tf.slice(
+            positional_encoding, [0, i, 0],
+            [positional_encoding_shape[0], 1, positional_encoding_shape[2]])
       return targets
 
     decoder_self_attention_bias = (
@@ -552,8 +559,14 @@ def _fast_decode(self,
       batch_size = partial_targets_shape[0]
 
     if hparams.pos == "timing":
-      timing_signal = common_attention.get_timing_signal_1d(
+      positional_encoding = common_attention.get_timing_signal_1d(
           decode_length + 1, hparams.hidden_size)
+    elif hparams.pos == "emb":
+      positional_encoding = common_attention.add_positional_embedding(
+          tf.zeros([1, decode_length + 1, hparams.hidden_size]),
+          hparams.max_length, "targets_positional_embedding", None)
+    else:
+      positional_encoding = None
 
     def preprocess_targets(targets, i):
       """Performs preprocessing steps on the targets to prepare for the decoder.
@@ -580,8 +593,8 @@ def preprocess_targets(targets, i):
       targets = tf.cond(
           tf.equal(i, 0), lambda: tf.zeros_like(targets), lambda: targets)
 
-      if hparams.pos == "timing":
-        targets += timing_signal[:, i:i + 1]
+      if positional_encoding is not None:
+        targets += positional_encoding[:, i:i + 1]
       return targets
 
     decoder_self_attention_bias = (
@@ -1053,22 +1066,27 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
   if hparams.proximity_bias:
     encoder_self_attention_bias += common_attention.attention_bias_proximal(
         common_layers.shape_list(inputs)[1])
-  # Append target_space_id embedding to inputs.
-  emb_target_space = common_layers.embedding(
-      target_space,
-      32,
-      ishape_static[-1],
-      name="target_space_embedding",
-      dtype=tf.bfloat16
-      if hparams.activation_dtype == "bfloat16" else tf.float32)
-  emb_target_space = tf.reshape(emb_target_space, [1, 1, -1])
-  encoder_input += emb_target_space
+  if hparams.get("use_target_space_embedding", True):
+    # Append target_space_id embedding to inputs.
+    emb_target_space = common_layers.embedding(
+        target_space,
+        32,
+        ishape_static[-1],
+        name="target_space_embedding",
+        dtype=tf.bfloat16
+        if hparams.activation_dtype == "bfloat16" else tf.float32)
+    emb_target_space = tf.reshape(emb_target_space, [1, 1, -1])
+    encoder_input += emb_target_space
   if hparams.pos == "timing":
     if inputs_position is not None:
       encoder_input = common_attention.add_timing_signal_1d_given_position(
           encoder_input, inputs_position)
     else:
       encoder_input = common_attention.add_timing_signal_1d(encoder_input)
+  elif hparams.pos == "emb":
+    encoder_input = common_attention.add_positional_embedding(
+        encoder_input, hparams.max_length, "inputs_positional_embedding",
+        inputs_position)
   if hparams.activation_dtype == "bfloat16":
     encoder_self_attention_bias = tf.cast(encoder_self_attention_bias,
                                           tf.bfloat16)
@@ -1125,6 +1143,11 @@ def transformer_prepare_decoder(targets, hparams, features=None):
           decoder_input, targets_position)
     else:
       decoder_input = common_attention.add_timing_signal_1d(decoder_input)
+  elif hparams.pos == "emb":
+    decoder_input = common_attention.add_positional_embedding(
+        decoder_input, hparams.max_length, "targets_positional_embedding",
+        targets_position)
+
   if hparams.activation_dtype == "bfloat16":
     decoder_self_attention_bias = tf.cast(decoder_self_attention_bias,
                                           tf.bfloat16)
@@ -1193,7 +1216,8 @@ def transformer_encoder(encoder_input,
               max_relative_position=hparams.max_relative_position,
               make_image_summary=make_image_summary,
               dropout_broadcast_dims=attention_dropout_broadcast_dims,
-              max_length=hparams.get("max_length"))
+              max_length=hparams.get("max_length"),
+              vars_3d=hparams.get("attention_variables_3d"))
           x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(
@@ -1478,6 +1502,8 @@ def transformer_base_v1():
   hparams.add_hparam("self_attention_type", "dot_product")
   hparams.add_hparam("max_relative_position", 0)
   hparams.add_hparam("conv_first_kernel", 3)
+  hparams.add_hparam("attention_variables_3d", False)
+  hparams.add_hparam("use_target_space_embedding", True)
   # These parameters are only used when ffn_layer=="local_moe_tpu"
   hparams.add_hparam("moe_overhead_train", 1.0)
   hparams.add_hparam("moe_overhead_eval", 2.0)

From 008a68451b3fed0edc50017389607d35a1e759e1 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 21 Jun 2018 16:03:20 -0700
Subject: [PATCH 0179/2720] Lint fixes

PiperOrigin-RevId: 201603381
---
 tensor2tensor/data_generators/gym_problems.py | 570 ++++++++++++------
 .../data_generators/gym_problems_test.py      |   2 +-
 tensor2tensor/data_generators/problem.py      |   2 +-
 tensor2tensor/data_generators/video_utils.py  |   2 +-
 tensor2tensor/models/basic.py                 |   2 +-
 tensor2tensor/models/research/rl.py           |  46 +-
 tensor2tensor/models/transformer_test.py      |   2 -
 tensor2tensor/models/vanilla_gan.py           |   2 +-
 tensor2tensor/rl/collect.py                   | 162 ++---
 tensor2tensor/rl/envs/batch_env_factory.py    | 274 ---------
 tensor2tensor/rl/envs/py_func_batch_env.py    |  45 +-
 tensor2tensor/rl/envs/simulated_batch_env.py  |  40 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py    |  24 +-
 tensor2tensor/rl/envs/utils.py                | 281 +++++++--
 tensor2tensor/rl/model_rl_experiment.py       |  43 +-
 tensor2tensor/rl/ppo.py                       |  40 +-
 tensor2tensor/rl/rl_trainer_lib.py            |  54 +-
 tensor2tensor/rl/rl_trainer_lib_test.py       |  13 +-
 18 files changed, 849 insertions(+), 755 deletions(-)
 delete mode 100644 tensor2tensor/rl/envs/batch_env_factory.py

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index b5dc8c78e..28614f0c3 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -27,12 +27,19 @@
 from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
+from tensor2tensor.layers import discretization
+from tensor2tensor.models.research import autoencoders
 from tensor2tensor.models.research import rl
-from tensor2tensor.models.research.rl import standard_atari_env_spec
 from tensor2tensor.rl import collect
+from tensor2tensor.rl.envs import tf_atari_wrappers as atari
+from tensor2tensor.rl.envs.tf_atari_wrappers import StackAndSkipWrapper
 from tensor2tensor.rl.envs.tf_atari_wrappers import TimeLimitWrapper
+from tensor2tensor.rl.envs.utils import batch_env_factory
+
+
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
+
 import tensorflow as tf
 
 
@@ -54,73 +61,6 @@ class GymDiscreteProblem(video_utils.VideoProblem):
   def __init__(self, *args, **kwargs):
     super(GymDiscreteProblem, self).__init__(*args, **kwargs)
     self._env = None
-    self._env = None
-    self.debug_dump_frames_path = "debug_frames_env"
-    self.settable_num_steps = 5000
-
-    self.environment_spec = self.get_environment_spec()
-    self.eval_phase = False
-    self.sum_of_rewards = 0.0
-    self.dones = 0
-
-  def _setup(self):
-    collect_hparams = rl.ppo_pong_base()
-    collect_hparams.add_hparam("environment_spec", self.environment_spec)
-
-    if not FLAGS.agent_policy_path:
-      collect_hparams.policy_network = rl.random_policy_fun
-
-    self._internal_memory_size = 10
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      collect_hparams.epoch_length = self._internal_memory_size
-      # TODO (piotrmilos). it is possible to set more then 1
-      collect_hparams.num_agents = 1
-      self.collect_memory, self.collect_trigger_op \
-        = collect.define_collect(collect_hparams, scope="gym_problems",
-                                 collect_level=0, eval_phase=self.eval_phase)
-
-  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
-    self._setup()
-    self.debug_dump_frames_path = os.path.join(
-        data_dir, self.debug_dump_frames_path)
-
-    with tf.Session() as sess:
-      sess.run(tf.global_variables_initializer())
-      self.restore_networks(sess)
-      pieces_generated = 0
-      memory_index = 0
-      memory = None
-      while pieces_generated < self.num_steps:
-        if memory is None or memory_index >= self._internal_memory_size:
-          sess.run(self.collect_trigger_op)
-          memory = sess.run(self.collect_memory)
-          memory_index = 0
-        data = [memory[i][memory_index][0] for i in range(4)]
-        memory_index += 1
-        observ, reward, done, action = data
-        observ = observ.astype(np.uint8)
-
-        self.sum_of_rewards += reward
-        self.dones += int(done)
-
-        ret_dict = {"frame": observ,
-                    "image/format": ["png"],
-                    "image/height": [self.frame_height],
-                    "image/width": [self.frame_width],
-                    "action": [int(action)],
-                    "done": [int(False)],
-                    "reward": [int(reward) - self.min_reward]}
-
-        yield ret_dict
-        pieces_generated += 1
-
-  def restore_networks(self, sess):
-    if FLAGS.agent_policy_path:
-      model_saver = tf.train.Saver(
-          tf.global_variables(".*network_parameters.*"))
-      ckpts = tf.train.get_checkpoint_state(FLAGS.agent_policy_path)
-      ckpt = ckpts.model_checkpoint_path
-      model_saver.restore(sess, ckpt)
 
   @property
   def num_input_frames(self):
@@ -147,9 +87,6 @@ def extra_reading_spec(self):
     }
     return data_fields, decoders
 
-  def get_environment_spec(self):
-    raise NotImplementedError()
-
   @property
   def is_generate_per_split(self):
     """Whether we have a train/test split or just hold out data."""
@@ -184,7 +121,7 @@ def num_rewards(self):
 
   @property
   def num_steps(self):
-    return self.settable_num_steps
+    raise NotImplementedError()
 
   @property
   def total_number_of_frames(self):
@@ -217,76 +154,22 @@ def hparams(self, defaults, unused_model_hparams):
     p.input_space_id = problem.SpaceID.IMAGE
     p.target_space_id = problem.SpaceID.IMAGE
 
-
-class GymAEDiscreteProblem(GymDiscreteProblem):
-  pass
-
-
-class GymRealDiscreteProblem(GymDiscreteProblem):
-
-  def __init__(self, *args, **kwargs):
-    super(GymRealDiscreteProblem, self).__init__(*args, **kwargs)
-    self.make_extra_debug_info = False
-
-  def get_environment_spec(self):
-    return standard_atari_env_spec(self.env_name)
-
-
-
-class GymSimulatedDiscreteProblem(GymDiscreteProblem):
-  """Simulated gym environment with discrete actions and rewards."""
-
-  def __init__(self, *args, **kwargs):
-    self.simulated_environment = True
-    self.debug_dump_frames_path = "debug_frames_sim"
-    self.intrinsic_reward_scale = 0.0
-    self.simulation_random_starts = False
-    super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
-
-  @property
-  def initial_frames_problem(self):
-    raise NotImplementedError()
-
-  def get_environment_spec(self):
-    env_spec = standard_atari_env_spec(self.env_name)
-
-    #Set reasonable time limit (as we do not simulate done)
-    real_env = env_spec.env_lambda()
-    if self.num_testing_steps is not None:
-      timelimit = self.num_testing_steps
-    else:
-      try:
-        # We assume that the real env is wrapped with TimeLimit.
-        history = self.num_input_frames
-        timelimit = real_env._max_episode_steps - history  # pylint: disable=protected-access
-      except:  # pylint: disable=bare-except
-        # If not, set some reasonable default.
-        timelimit = 100
-
-    env_spec.simulated_env = True
-    env_spec.add_hparam("simulation_random_starts",
-                        self.simulation_random_starts)
-    env_spec.add_hparam("intrinsic_reward_scale",
-                        self.intrinsic_reward_scale)
-    initial_frames_problem = registry.problem(self.initial_frames_problem)
-    env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
-    env_spec.wrappers.append([TimeLimitWrapper, {"timelimit": timelimit}])
-
-    return env_spec
-
-  def restore_networks(self, sess):
-    super(GymSimulatedDiscreteProblem, self).restore_networks(sess)
-    # TODO(blazej): adjust regexp for different models.
-    env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
-    sess = tf.get_default_session()
-
-    ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
-    ckpt = ckpts.model_checkpoint_path
-    env_model_loader.restore(sess, ckpt)
+  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
+    next_observation = self.env.reset()
+    for _ in range(self.num_steps):
+      observation = next_observation
+      action = self.get_action(observation)
+      next_observation, reward, done, _ = self.env.step(action)
+      if done:
+        next_observation = self.env.reset()
+      yield {"frame": observation,
+             "action": [action],
+             "done": [done],
+             "reward": [int(reward - self.min_reward)]}
 
 
 @registry.register_problem
-class GymPongRandom(GymDiscreteProblem):
+class GymPongRandom5k(GymDiscreteProblem):
   """Pong game, random actions."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
@@ -315,9 +198,22 @@ def min_reward(self):
   def num_rewards(self):
     return 3
 
+  @property
+  def num_steps(self):
+    return 5000
+
 
 @registry.register_problem
-class GymWrappedPongRandom(GymDiscreteProblem):
+class GymPongRandom50k(GymPongRandom5k):
+  """Pong game, random actions."""
+
+  @property
+  def num_steps(self):
+    return 50000
+
+
+@registry.register_problem
+class GymWrappedPongRandom5k(GymDiscreteProblem):
   """Pong game, random actions."""
 
   @property
@@ -332,6 +228,10 @@ def min_reward(self):
   def num_rewards(self):
     return 3
 
+  @property
+  def num_steps(self):
+    return 5000
+
 
 @registry.register_problem
 class GymWrappedLongPongRandom(GymDiscreteProblem):
@@ -349,13 +249,17 @@ def min_reward(self):
   def num_rewards(self):
     return 3
 
+  @property
+  def num_steps(self):
+    return 5000
+
   @property
   def num_testing_steps(self):
     return 100
 
 
 @registry.register_problem
-class GymWrappedBreakoutRandom(GymDiscreteProblem):
+class GymWrappedBreakoutRandom5k(GymDiscreteProblem):
   """Pong game, random actions."""
 
   @property
@@ -370,18 +274,22 @@ def min_reward(self):
   def num_rewards(self):
     return 3
 
+  @property
+  def num_steps(self):
+    return 5000
+
 
 @registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnPong(
-    GymSimulatedDiscreteProblem, GymPongRandom):
+class GymWrappedPongRandom50k(GymPongRandom5k):
+  """Pong game, random actions."""
 
   @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_pong"
+  def num_steps(self):
+    return 50000
 
 
 @registry.register_problem
-class GymFreewayRandom(GymDiscreteProblem):
+class GymFreewayRandom5k(GymDiscreteProblem):
   """Freeway game, random actions."""
 
   @property
@@ -396,24 +304,351 @@ def min_reward(self):
   def num_rewards(self):
     return 2
 
+  @property
+  def num_steps(self):
+    return 5000
+
+
+@registry.register_problem
+class GymFreewayRandom50k(GymFreewayRandom5k):
+  """Freeway game, random actions."""
+
+  @property
+  def num_steps(self):
+    return 50000
+
+
+class GymDiscreteProblemWithAgent(GymDiscreteProblem):
+  """Gym environment with discrete actions and rewards and an agent."""
+
+  def __init__(self, *args, **kwargs):
+    super(GymDiscreteProblemWithAgent, self).__init__(*args, **kwargs)
+    self._env = None
+    self.debug_dump_frames_path = "debug_frames_env"
+    self.make_extra_debug_info = True
+    self.autoencoder_model = None
+
+    # Defaults.
+    self.environment_spec = lambda: gym.make(self.env_name)
+    self._real_env = None
+    self.real_env_problem = None
+    self.in_graph_wrappers = []
+    self.collect_hparams = rl.ppo_pong_base()
+    if FLAGS.autoencoder_path:
+      self.collect_hparams = rl.ppo_pong_ae_base()
+    self.settable_num_steps = 50000
+    self.simulated_environment = None
+    self.eval_phase = False
+    self.warm_up = 10  # TODO(piotrm): This should be probably removed.
+
+    # Debug info.
+    self.dones = 0
+    self.real_reward = 0
+    self.total_sim_reward, self.total_real_reward = 0.0, 0.0
+    self.sum_of_rewards = 0.0
+    self.successful_episode_reward_predictions = 0
+
+  @property
+  def real_env(self):
+    """Lazy caching environment construction."""
+    if self._real_env is None:
+      self._real_env = self.environment_spec()
+    return self._real_env
+
+  @property
+  def num_steps(self):
+    return self.settable_num_steps
+
+  @property
+  def raw_frame_height(self):
+    return self.env.observation_space.shape[0]
+
+  @property
+  def frame_height(self):
+    if FLAGS.autoencoder_path:
+      # TODO(lukaszkaiser): remove hard-coded autoencoder params.
+      return int(math.ceil(self.raw_frame_height / self.autoencoder_factor))
+    return self.raw_frame_height
+
+  @property
+  def autoencoder_factor(self):
+    """By how much to divide sizes when using autoencoders."""
+    hparams = autoencoders.autoencoder_discrete_pong()
+    return 2**hparams.num_hidden_layers
+
+  @property
+  def raw_frame_width(self):
+    return self.env.observation_space.shape[1]
+
+  @property
+  def frame_width(self):
+    if FLAGS.autoencoder_path:
+      # TODO(lukaszkaiser): remove hard-coded autoencoder params.
+      return int(math.ceil(self.raw_frame_width / self.autoencoder_factor))
+    return self.raw_frame_width
+
+  def setup_autoencoder(self):
+    if self.autoencoder_model is not None:
+      return
+    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+      autoencoder_hparams = autoencoders.autoencoder_discrete_pong()
+      autoencoder_hparams.data_dir = "unused"
+      autoencoder_hparams.problem_hparams = self.get_hparams(
+          autoencoder_hparams)
+      autoencoder_hparams.problem = self
+      self.autoencoder_model = autoencoders.AutoencoderOrderedDiscrete(
+          autoencoder_hparams, tf.estimator.ModeKeys.EVAL)
+
+  def autoencode_tensor(self, x, batch_size=1):
+    if self.autoencoder_model is None:
+      return x
+    shape = [self.raw_frame_height, self.raw_frame_width, self.num_channels]
+    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+      self.autoencoder_model.set_mode(tf.estimator.ModeKeys.EVAL)
+      autoencoded = self.autoencoder_model.encode(
+          tf.reshape(x, [batch_size, 1] + shape))
+    autoencoded = tf.reshape(
+        autoencoded, [batch_size] + self.frame_shape + [8])  # 8-bit groups.
+    if batch_size == 1:
+      autoencoded = tf.squeeze(autoencoded, axis=0)
+    return discretization.bit_to_int(autoencoded, 8)
+
+  def _setup(self):
+    if self.make_extra_debug_info:
+      self.report_reward_statistics_every = 10
+      self.dones = 0
+      self.real_reward = 0
+      # Slight weirdness to make sim env and real env aligned
+      if self.simulated_environment:
+        self.real_env.reset()
+        for _ in range(self.num_input_frames):
+          self.real_ob, _, _, _ = self.real_env.step(0)
+      self.total_sim_reward, self.total_real_reward = 0.0, 0.0
+      self.sum_of_rewards = 0.0
+      self.successful_episode_reward_predictions = 0
+
+    in_graph_wrappers = self.in_graph_wrappers + [
+        (atari.MemoryWrapper, {}), (StackAndSkipWrapper, {"skip": 4})]
+    env_hparams = tf.contrib.training.HParams(
+        in_graph_wrappers=in_graph_wrappers,
+        problem=self.real_env_problem if self.real_env_problem else self,
+        simulated_environment=self.simulated_environment)
+    if self.simulated_environment:
+      env_hparams.add_hparam("simulation_random_starts",
+                             self.simulation_random_starts)
+      env_hparams.add_hparam("intrinsic_reward_scale",
+                             self.intrinsic_reward_scale)
+
+    generator_batch_env = batch_env_factory(
+        self.environment_spec, env_hparams, num_agents=1, xvfb=False)
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+      if FLAGS.agent_policy_path:
+        policy_lambda = self.collect_hparams.network
+      else:
+        # When no agent_policy_path is set, just generate random samples.
+        policy_lambda = rl.random_policy_fun
+
+    if FLAGS.autoencoder_path:
+      # TODO(lukaszkaiser): remove hard-coded autoencoder params.
+      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+        self.setup_autoencoder()
+        autoencoder_model = self.autoencoder_model
+        # Feeds for autoencoding.
+        shape = [self.raw_frame_height, self.raw_frame_width, self.num_channels]
+        self.autoencoder_feed = tf.placeholder(tf.int32, shape=shape)
+        self.autoencoder_result = self.autoencode_tensor(self.autoencoder_feed)
+        # Now for autodecoding.
+        shape = self.frame_shape
+        self.autodecoder_feed = tf.placeholder(tf.int32, shape=shape)
+        bottleneck = tf.reshape(
+            discretization.int_to_bit(self.autodecoder_feed, 8),
+            [1, 1, self.frame_height, self.frame_width, self.num_channels * 8])
+        autoencoder_model.set_mode(tf.estimator.ModeKeys.PREDICT)
+        self.autodecoder_result = autoencoder_model.decode(bottleneck)
+
+    def preprocess_fn(x):
+      shape = [self.raw_frame_height, self.raw_frame_width, self.num_channels]
+      # TODO(lukaszkaiser): we assume x comes from StackAndSkipWrapper skip=4.
+      xs = [tf.reshape(t, [1] + shape) for t in tf.split(x, 4, axis=-1)]
+      autoencoded = self.autoencode_tensor(tf.concat(xs, axis=0), batch_size=4)
+      encs = [tf.squeeze(t, axis=[0]) for t in tf.split(autoencoded, 4, axis=0)]
+      res = tf.to_float(tf.concat(encs, axis=-1))
+      return tf.expand_dims(res, axis=0)
+
+    # TODO(lukaszkaiser): x is from StackAndSkipWrapper thus 4*num_channels.
+    shape = [1, self.frame_height, self.frame_width, 4 * self.num_channels]
+    do_preprocess = (self.autoencoder_model is not None and
+                     not self.simulated_environment)
+    preprocess = (preprocess_fn, shape) if do_preprocess else None
+
+    def policy(x):
+      return policy_lambda(self.environment_spec().action_space,
+                           self.collect_hparams, x)
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+      self.collect_hparams.epoch_length = 10
+      _, self.collect_trigger_op = collect.define_collect(
+          policy, generator_batch_env, self.collect_hparams,
+          eval_phase=self.eval_phase,
+          scope="define_collect", preprocess=preprocess)
+
+    self.avilable_data_size_op = atari.MemoryWrapper.singleton.speculum.size()
+    self.data_get_op = atari.MemoryWrapper.singleton.speculum.dequeue()
+
+  def restore_networks(self, sess):
+    if FLAGS.agent_policy_path:
+      model_saver = tf.train.Saver(
+          tf.global_variables(".*network_parameters.*"))
+      ckpts = tf.train.get_checkpoint_state(FLAGS.agent_policy_path)
+      ckpt = ckpts.model_checkpoint_path
+      model_saver.restore(sess, ckpt)
+    if FLAGS.autoencoder_path:
+      autoencoder_saver = tf.train.Saver(
+          tf.global_variables("autoencoder.*"))
+      ckpts = tf.train.get_checkpoint_state(FLAGS.autoencoder_path)
+      ckpt = ckpts.model_checkpoint_path
+      autoencoder_saver.restore(sess, ckpt)
+
+  def autoencode(self, image, sess):
+    return sess.run(self.autoencoder_result, {self.autoencoder_feed: image})
+
+  def autodecode(self, encoded, sess):
+    res = sess.run(self.autodecoder_result, {self.autodecoder_feed: encoded})
+    return res[0, 0, :self.raw_frame_height, :self.raw_frame_width, :]
+
+  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
+    self._setup()
+    self.debug_dump_frames_path = os.path.join(
+        data_dir, self.debug_dump_frames_path)
+
+    with tf.Session() as sess:
+      sess.run(tf.global_variables_initializer())
+      self.restore_networks(sess)
+      if FLAGS.only_use_ae_for_policy:
+        # If only the policy should use the autoencoder, then reset the flag so
+        # that other components here act as though there is no autoencoder and
+        # so write out full-resolution images. The policy graph was already
+        # built and self.collect_trigger_op is all that's used from it.
+        FLAGS.autoencoder_path = None
+      pieces_generated = 0
+      while pieces_generated < self.num_steps + self.warm_up:
+        avilable_data_size = sess.run(self.avilable_data_size_op)
+        if avilable_data_size < 1:
+          sess.run(self.collect_trigger_op)
+        observ, reward, action, done = sess.run(self.data_get_op)
+        debug_im = None
+        if self.make_extra_debug_info:
+          self.total_sim_reward += reward
+          if not self.simulated_environment:
+            self.real_ob = observ
+            self.real_reward = reward
+          if not FLAGS.autoencoder_path:
+            err = np.ndarray.astype(np.maximum(np.abs(
+                self.real_ob - observ, dtype=np.int) - 10, 0),
+                                    np.uint8)
+            debug_im = np.concatenate([observ, self.real_ob, err], axis=1)
+          if done:
+            self.dones += 1
+            self.sum_of_rewards += self.real_reward
+            if self.total_real_reward == self.total_sim_reward:
+              self.successful_episode_reward_predictions += 1
+
+            self.total_real_reward = 0.0
+            self.total_sim_reward = 0.0
+            self.real_reward = 0
+            if self.simulated_environment:
+              self.real_env.reset()
+              # Slight weirdness to make sim env and real env aligned
+              for _ in range(self.num_input_frames):
+                self.real_ob, _, _, _ = self.real_env.step(0)
+          else:
+            if self.simulated_environment:
+              self.real_ob, self.real_reward, _, _ = self.real_env.step(action)
+            self.total_real_reward += self.real_reward
+            self.sum_of_rewards += self.real_reward
+        if FLAGS.autoencoder_path:
+          if self.simulated_environment:
+            debug_im = self.autodecode(observ, sess)
+          else:
+            orig_observ = observ
+            observ = self.autoencode(observ, sess)
+            debug_im = np.concatenate([self.autodecode(observ, sess),
+                                       orig_observ], axis=1)
+        ret_dict = {"frame": observ,
+                    "image/format": ["png"],
+                    "image/height": [self.frame_height],
+                    "image/width": [self.frame_width],
+                    "action": [int(action)],
+                    "done": [int(False)],
+                    "reward": [int(reward) - self.min_reward]}
+        if self.make_extra_debug_info:
+          ret_dict["image/debug"] = debug_im
+        yield ret_dict
+        pieces_generated += 1
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgent(GymDiscreteProblemWithAgent):
+  """Simulated gym environment with discrete actions and rewards."""
+
+  def __init__(self, *args, **kwargs):
+    super(GymSimulatedDiscreteProblemWithAgent, self).__init__(*args, **kwargs)
+    self.simulated_environment = True
+    self.make_extra_debug_info = True
+    self.debug_dump_frames_path = "debug_frames_sim"
+
+  @property
+  def real_env(self):
+    """Lazy caching environment construction."""
+    if self._real_env is None:
+      self._real_env = self.environment_spec()
+      if self.num_testing_steps is not None:
+        timelimit = self.num_testing_steps
+      else:
+        try:
+          # We assume that the real env is wrapped with TimeLimit.
+          history = self.num_input_frames
+          timelimit = self.real_env._max_episode_steps - history  # pylint: disable=protected-access
+        except:  # pylint: disable=bare-except
+          # If not, set some reasonable default.
+          timelimit = 100
+      self.in_graph_wrappers.append(
+          (TimeLimitWrapper, {"timelimit": timelimit}))
+    return self._real_env
+
+  def restore_networks(self, sess):
+    super(GymSimulatedDiscreteProblemWithAgent, self).restore_networks(sess)
+    # TODO(blazej): adjust regexp for different models.
+    env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
+    sess = tf.get_default_session()
+
+    ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
+    ckpt = ckpts.model_checkpoint_path
+    env_model_loader.restore(sess, ckpt)
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnPong(
+    GymSimulatedDiscreteProblemWithAgent, GymPongRandom5k):
+  pass
+
+
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnPong(
-    GymRealDiscreteProblem, GymPongRandom):
+    GymDiscreteProblemWithAgent, GymPongRandom5k):
   pass
 
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
-    GymSimulatedDiscreteProblem, GymWrappedPongRandom):
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_pong"
+    GymSimulatedDiscreteProblemWithAgent, GymWrappedPongRandom5k):
+  pass
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedLongPong(
-    GymRealDiscreteProblem, GymWrappedLongPongRandom):
+    GymDiscreteProblemWithAgent, GymWrappedLongPongRandom):
   pass
 
 
@@ -425,16 +660,13 @@ class GymDiscreteProblemWithAgentOnWrappedLongPongAe(  # with autoencoder
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPong(
-    GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_long_pong"
+    GymSimulatedDiscreteProblemWithAgent, GymWrappedLongPongRandom):
+  pass
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedBreakout(
-    GymRealDiscreteProblem, GymWrappedBreakoutRandom):
+    GymDiscreteProblemWithAgent, GymWrappedBreakoutRandom5k):
   pass
 
 
@@ -446,17 +678,13 @@ class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
-    GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_breakout"
-
+    GymSimulatedDiscreteProblemWithAgent, GymWrappedBreakoutRandom5k):
+  pass
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedPong(
-    GymRealDiscreteProblem, GymWrappedPongRandom):
+    GymDiscreteProblemWithAgent, GymWrappedPongRandom5k):
   """GymDiscreteProblemWithAgentOnWrappedPong."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
@@ -486,17 +714,13 @@ class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnFreeway(
-    GymSimulatedDiscreteProblem, GymFreewayRandom):
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_freeway"
-
+    GymSimulatedDiscreteProblemWithAgent, GymFreewayRandom5k):
+  pass
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnFreeway(
-    GymRealDiscreteProblem, GymFreewayRandom):
+    GymDiscreteProblemWithAgent, GymFreewayRandom5k):
   """Freeway with agent."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index 9638acebf..1bbd2a374 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -35,7 +35,7 @@ def setUpClass(cls):
     os.mkdir(cls.tmp_dir)
 
   def testGymAtariBoots(self):
-    problem = gym_problems.GymSimulatedDiscreteProblemWithAgentOnPong()
+    problem = gym_problems.GymPongRandom5k()
     self.assertEqual(5000, problem.num_steps)
 
 
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index ab8ea1cff..3186ceea4 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -813,7 +813,7 @@ def define_shapes(example):
       else:
         # On GPU, bucket by length
         dataset = dataset.filter(gpu_valid_size)
-        shard_multiplier = (config and config.data_parallelism.n) or 1
+        shard_multiplier = config.data_parallelism.n if config else 1
         batching_scheme = data_reader.hparams_to_batching_scheme(
             hparams,
             shard_multiplier=shard_multiplier,
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index b0bdb77d5..bf20c94ec 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -337,7 +337,7 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
             unencoded_debug = features.pop("image/debug")
             encoded_debug = sess.run(encoded_image_t, feed_dict={
                 image_t: unencoded_debug})
-            features["image/encoded_debug"] = [encoded_debug]
+            features["image/encoded_debug"] = encoded_debug
           yield features
 
   def generate_encoded_samples_debug(self, data_dir, tmp_dir, dataset_split):
diff --git a/tensor2tensor/models/basic.py b/tensor2tensor/models/basic.py
index 302396b39..bcaf8a872 100644
--- a/tensor2tensor/models/basic.py
+++ b/tensor2tensor/models/basic.py
@@ -207,7 +207,7 @@ def discriminate(x):
     return res, {"bottleneck_loss": b_loss, "gan_loss": - gan_loss}
 
   def sample(self, features=None, shape=None):
-    del features
+    del features, shape
     hp = self.hparams
     div_x = 2**hp.num_hidden_layers
     div_y = 1 if self.is1d else 2**hp.num_hidden_layers
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index bf21a5c55..a3410a40f 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -22,7 +22,6 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
 from tensor2tensor.utils import registry
-from tensor2tensor.rl.envs import tf_atari_wrappers
 
 import tensorflow as tf
 
@@ -60,26 +59,16 @@ def ppo_base_v1():
 @registry.register_hparams
 def ppo_continuous_action_base():
   hparams = ppo_base_v1()
-  hparams.add_hparam("policy_network", feed_forward_gaussian_fun)
-  hparams.add_hparam("policy_network_params", "basic_policy_parameters")
+  hparams.add_hparam("network", feed_forward_gaussian_fun)
   return hparams
 
-@registry.register_hparams
-def basic_policy_parameters():
-  wrappers = None
-  return tf.contrib.training.HParams(wrappers=wrappers)
 
 @registry.register_hparams
 def ppo_discrete_action_base():
   hparams = ppo_base_v1()
-  hparams.add_hparam("policy_network", feed_forward_categorical_fun)
+  hparams.add_hparam("network", feed_forward_categorical_fun)
   return hparams
 
-@registry.register_hparams
-def discrete_random_action_base():
-  hparams = common_hparams.basic_params1()
-  hparams.add_hparam("policy_network", random_policy_fun)
-  return hparams
 
 @registry.register_hparams
 def ppo_atari_base():
@@ -113,42 +102,13 @@ def ppo_pong_base():
   hparams.optimization_epochs = 2
   hparams.epochs_num = 1000
   hparams.num_eval_agents = 1
-  hparams.policy_network = feed_forward_cnn_small_categorical_fun
+  hparams.network = feed_forward_cnn_small_categorical_fun
   hparams.clipping_coef = 0.2
   hparams.optimization_batch_size = 4
   hparams.max_gradients_norm = 0.5
   return hparams
 
 
-def standard_atari_env_spec(env):
-  """Parameters of environment specification"""
-  standard_wrappers = [[tf_atari_wrappers.MaxAndSkipWrapper, {"skip": 4}]]
-  env_lambda = None
-  if isinstance(env, str):
-    env_lambda = lambda: gym.make(env)
-  if callable(env):
-    env_lambda = env
-  assert env is not None, "Unknown specification of environment"
-
-  return tf.contrib.training.HParams(env_lambda=env_lambda,
-                                     wrappers=standard_wrappers,
-                                     simulated_env=False)
-
-
-def simple_gym_spec(env):
-  """Parameters of environment specification"""
-  standard_wrappers = None
-  env_lambda = None
-  if isinstance(env, str):
-    env_lambda = lambda: gym.make(env)
-  if callable(env):
-    env_lambda = env
-  assert env is not None, "Unknown specification of environment"
-
-  return tf.contrib.training.HParams(env_lambda=env_lambda,
-                                     wrappers=standard_wrappers,
-                                     simulated_env=False)
-
 @registry.register_hparams
 def ppo_pong_ae_base():
   """Pong autoencoder base parameters."""
diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 7df6d8776..6325753a5 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -258,8 +258,6 @@ def _create_greedy_infer_model(self, decode_length):
     """
     model, features = get_model(transformer.transformer_small())
 
-    decode_length = 3
-
     out_logits, _ = model(features)
     out_logits = tf.squeeze(out_logits, axis=[2, 3])
     loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
diff --git a/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index e9fddc096..730d5b847 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -148,7 +148,7 @@ def body(self, features):
     out_shape = common_layers.shape_list(inputs)[1:4]
     g = self.generator(z, is_training, out_shape)
 
-    losses = self.losses(inputs, g)
+    losses = self.losses(inputs, g)  # pylint: disable=not-callable
 
     summary_g_image = tf.reshape(
         g[0, :], [1] + common_layers.shape_list(inputs)[1:])
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 64f150f7a..a3e807efd 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -17,100 +17,44 @@
 from __future__ import division
 from __future__ import print_function
 
-
-import copy
 import tensorflow as tf
 
-from tensor2tensor.rl.envs.batch_env_factory import batch_env_factory
-from tensor2tensor.rl.envs.tf_atari_wrappers import WrapperBase
-from tensor2tensor.rl.envs.utils import get_policy
-
 
-def _rollout_metadata(batch_env):
+def define_collect(policy_factory, batch_env, hparams,
+                   eval_phase, policy_to_actions_lambda=None,
+                   scope="", preprocess=None, on_simulated=False):
+  """Collect trajectories."""
+  eval_phase = tf.convert_to_tensor(eval_phase)
+  on_simulated = tf.convert_to_tensor(on_simulated)
   batch_env_shape = batch_env.observ.get_shape().as_list()
-  batch_size = [batch_env_shape[0]]
-  shapes_types_names = [
-      (batch_size + batch_env_shape[1:], tf.float32, "observation"),
-      (batch_size, tf.float32, "reward"),
-      (batch_size, tf.bool, "done"),
-      (batch_size + batch_env.action_shape, batch_env.action_dtype, "action"),
-      (batch_size, tf.float32, "pdf"),
-      (batch_size, tf.float32, "value_function"),
+  if preprocess is not None:
+    batch_env_shape = preprocess[1]
+  memory_shape = [hparams.epoch_length] + [batch_env_shape[0]]
+  memories_shapes_and_types = [
+      # observation
+      (memory_shape + batch_env_shape[1:], tf.float32),
+      (memory_shape, tf.float32),      # reward
+      (memory_shape, tf.bool),         # done
+      # action
+      (memory_shape + batch_env.action_shape, batch_env.action_dtype),
+      (memory_shape, tf.float32),      # pdf
+      (memory_shape, tf.float32),      # value function
   ]
-  return shapes_types_names
-
-
-class _MemoryWrapper(WrapperBase):
-  """Memory wrapper."""
-
-  def __init__(self, batch_env):
-    super(_MemoryWrapper, self).__init__(batch_env)
-    infinity = 10000000
-    meta_data = list(zip(*_rollout_metadata(batch_env)))
-    shapes = meta_data[0][:4]
-    dtypes = meta_data[1][:4]
-    self.speculum = tf.FIFOQueue(infinity, shapes=shapes, dtypes=dtypes)
-    observs_shape = batch_env.observ.shape
-    observ_dtype = tf.float32
-    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
-                               trainable=False)
-
-  def simulate(self, action):
-    reward, done = self._batch_env.simulate(action)
-    with tf.control_dependencies([reward, done]):
-      assign = self._observ.assign(self._batch_env.observ)
-
-    with tf.control_dependencies([assign]):
-      enqueue_op = self.speculum.enqueue(
-          [self._observ.read_value(), reward, done, action])
-      with tf.control_dependencies([enqueue_op]):
-        return tf.identity(reward), tf.identity(done)
-
-
-def define_collect(hparams, scope, eval_phase,
-                   collect_level=-1,
-                   policy_to_actions_lambda=None,
-                   on_simulated=False):
-  """Collect trajectories."""
-
-  with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
-    batch_env = batch_env_factory(hparams)
-    environment_wrappers = hparams.environment_spec.wrappers
-    wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
-    #Put memory wrapper at the level you want to gather observations at
-    #Negative indices need to be shifted for insert to work correctly
-    collect_level = collect_level if \
-      collect_level >= 0 else len(wrappers) + collect_level + 1
-    wrappers.insert(collect_level, [_MemoryWrapper, {}])
-    rollout_metadata = None
-    speculum = None
-    for w in wrappers:
-      batch_env = w[0](batch_env, **w[1])
-      if w[0] == _MemoryWrapper:
-        rollout_metadata = _rollout_metadata(batch_env)
-        speculum = batch_env.speculum
-
-    eval_phase = tf.convert_to_tensor(eval_phase)
-    on_simulated = tf.convert_to_tensor(on_simulated)
-
-    memory = [tf.get_variable("collect_memory_{}".format(name),
-                              shape=[hparams.epoch_length]+shape,
-                              dtype=dtype,
-                              initializer=tf.zeros_initializer(),
-                              trainable=False)
-              for (shape, dtype, name) in rollout_metadata]
-
-
+  memory = [tf.Variable(tf.zeros(shape, dtype), trainable=False)
+            for (shape, dtype) in memories_shapes_and_types]
+  with tf.variable_scope(scope):
     cumulative_rewards = tf.get_variable("cumulative_rewards", len(batch_env),
                                          trainable=False)
 
-    should_reset_var = tf.Variable(True, trainable=False)
+  should_reset_var = tf.Variable(True, trainable=False)
 
-    zeros_tensor = tf.zeros(len(batch_env))
+  zeros_tensor = tf.zeros(len(batch_env))
 
   def group():
-    return tf.group(batch_env.reset(tf.range(len(batch_env))),
-                    tf.assign(cumulative_rewards, zeros_tensor))
+    return tf.group(
+        batch_env.reset(tf.range(len(batch_env))),
+        tf.assign(cumulative_rewards, zeros_tensor))
+
   reset_op = tf.cond(
       tf.logical_or(should_reset_var, tf.logical_or(eval_phase, on_simulated)),
       group, tf.no_op)
@@ -127,42 +71,24 @@ def step(index, scores_sum, scores_num):
       # operation. We are waiting for tf.copy:
       # https://github.com/tensorflow/tensorflow/issues/11186
       obs_copy = batch_env.observ + 0
-
-      def env_step(arg1, arg2): # pylint: disable=unused-argument
-        actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
-        policy = actor_critic.policy
-        if policy_to_actions_lambda:
-          action = policy_to_actions_lambda(policy)
-        else:
-          action = tf.cond(eval_phase,
-                           policy.mode,
-                           policy.sample)
-
-        postprocessed_action = actor_critic.action_postprocessing(action)
-        simulate_output = batch_env.simulate(postprocessed_action[0, ...])
-
-        pdf = policy.prob(action)[0]
-        value_function = actor_critic.value[0]
-        pdf = tf.reshape(pdf, shape=(hparams.num_agents,))
-        value_function = tf.reshape(value_function, shape=(hparams.num_agents,))
-
-        with tf.control_dependencies(simulate_output):
-          return tf.identity(pdf), tf.identity(value_function)
-
-      pdf, value_function = tf.while_loop(
-          lambda _1, _2: tf.equal(speculum.size(), 0),
-          env_step,
-          [tf.constant(0.0, shape=(hparams.num_agents,)),
-           tf.constant(0.0, shape=(hparams.num_agents,))],
-          parallel_iterations=1,
-          back_prop=False,)
-
-      with tf.control_dependencies([pdf, value_function]):
-        obs, reward, done, action = speculum.dequeue()
-
+      if preprocess is not None:
+        obs_copy = preprocess[0](obs_copy)
+      actor_critic = policy_factory(tf.expand_dims(obs_copy, 0))
+      policy = actor_critic.policy
+      if policy_to_actions_lambda:
+        action = policy_to_actions_lambda(policy)
+      else:
+        action = tf.cond(eval_phase,
+                         policy.mode,
+                         policy.sample)
+      postprocessed_action = actor_critic.action_postprocessing(action)
+      simulate_output = batch_env.simulate(postprocessed_action[0, ...])
+      pdf = policy.prob(action)[0]
+      with tf.control_dependencies(simulate_output):
+        reward, done = simulate_output
         done = tf.reshape(done, (len(batch_env),))
-        to_save = [obs, reward, done, action,
-                   pdf, value_function]
+        to_save = [obs_copy, reward, done, action[0, ...], pdf,
+                   actor_critic.value[0]]
         save_ops = [tf.scatter_update(memory_slot, index, value)
                     for memory_slot, value in zip(memory, to_save)]
         cumulate_rewards_op = cumulative_rewards.assign_add(reward)
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
deleted file mode 100644
index 31f2e1b6b..000000000
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ /dev/null
@@ -1,274 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-"""Utilities for creating batched environments."""
-
-# The code was based on Danijar Hafner's code from tf.agents:
-# https://github.com/tensorflow/agents/blob/master/agents/tools/wrappers.py
-# https://github.com/tensorflow/agents/blob/master/agents/scripts/utility.py
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.rl.envs import batch_env
-from tensor2tensor.rl.envs import py_func_batch_env
-from tensor2tensor.rl.envs import simulated_batch_env
-
-import tensorflow as tf
-import atexit
-import multiprocessing
-import os
-import random
-import signal
-import subprocess
-import sys
-import traceback
-
-
-def batch_env_factory(hparams, xvfb=False):
-  """Factory of batch envs."""
-
-  environment_spec = hparams.environment_spec
-
-  if environment_spec.simulated_env:
-    # TODO(piotrmilos): Consider passing only relevant paramters
-    cur_batch_env = _define_simulated_batch_env(
-        hparams, hparams.num_agents,
-        hparams.simulation_random_starts,
-        hparams.intrinsic_reward_scale)
-  else:
-
-    cur_batch_env = _define_batch_env(hparams.environment_spec,
-                                      hparams.num_agents,
-                                      xvfb=xvfb)
-  return cur_batch_env
-
-
-def _define_batch_env(environment_spec, num_agents, xvfb=False):
-  """Create environments and apply all desired wrappers."""
-
-  with tf.variable_scope("environments"):
-    envs = [
-        ExternalProcessEnv(environment_spec.env_lambda, xvfb)
-        for _ in range(num_agents)]
-    env = batch_env.BatchEnv(envs, blocking=False)
-    env = py_func_batch_env.PyFuncBatchEnv(env)
-    return env
-
-
-def _define_simulated_batch_env(hparams, num_agents,
-                                simulation_random_starts=False,
-                                intrinsic_reward_scale=0.):
-  cur_batch_env = simulated_batch_env.SimulatedBatchEnv(
-      hparams, num_agents, simulation_random_starts,
-      intrinsic_reward_scale)
-  return cur_batch_env
-
-
-class ExternalProcessEnv(object):
-  """Step environment in a separate process for lock free parallelism."""
-
-  # Message types for communication via the pipe.
-  _ACCESS = 1
-  _CALL = 2
-  _RESULT = 3
-  _EXCEPTION = 4
-  _CLOSE = 5
-
-  def __init__(self, constructor, xvfb):
-    """Step environment in a separate process for lock free parallelism.
-
-    The environment will be created in the external process by calling the
-    specified callable. This can be an environment class, or a function
-    creating the environment and potentially wrapping it. The returned
-    environment should not access global variables.
-
-    Args:
-      constructor: Callable that creates and returns an OpenAI gym environment.
-      xvfb:  Frame buffer.
-
-    Attributes:
-      observation_space: The cached observation space of the environment.
-      action_space: The cached action space of the environment.
-    """
-    self._conn, conn = multiprocessing.Pipe()
-    if xvfb:
-      server_id = random.randint(10000, 99999)
-      auth_file_id = random.randint(10000, 99999999999)
-
-      xauthority_path = "/tmp/Xauthority_{}".format(auth_file_id)
-
-      command = "Xvfb :{} -screen 0 1400x900x24 -nolisten tcp -auth {}".format(
-          server_id, xauthority_path)
-      with open(os.devnull, "w") as devnull:
-        proc = subprocess.Popen(command.split(), shell=False, stdout=devnull,
-                                stderr=devnull)
-        atexit.register(lambda: os.kill(proc.pid, signal.SIGKILL))
-
-      def constructor_using_xvfb():
-        os.environ["DISPLAY"] = ":{}".format(server_id)
-        os.environ["XAUTHORITY"] = xauthority_path
-        return constructor()
-
-      self._process = multiprocessing.Process(
-          target=self._worker, args=(constructor_using_xvfb, conn))
-    else:
-      self._process = multiprocessing.Process(
-          target=self._worker, args=(constructor, conn))
-
-    atexit.register(self.close)
-    self._process.start()
-    self._observ_space = None
-    self._action_space = None
-
-  @property
-  def observation_space(self):
-    if not self._observ_space:
-      self._observ_space = self.__getattr__("observation_space")
-    return self._observ_space
-
-  @property
-  def action_space(self):
-    if not self._action_space:
-      self._action_space = self.__getattr__("action_space")
-    return self._action_space
-
-  def __getattr__(self, name):
-    """Request an attribute from the environment.
-
-    Note that this involves communication with the external process, so it can
-    be slow.
-
-    Args:
-      name: Attribute to access.
-
-    Returns:
-      Value of the attribute.
-    """
-    self._conn.send((self._ACCESS, name))
-    return self._receive()
-
-  def call(self, name, *args, **kwargs):
-    """Asynchronously call a method of the external environment.
-
-    Args:
-      name: Name of the method to call.
-      *args: Positional arguments to forward to the method.
-      **kwargs: Keyword arguments to forward to the method.
-
-    Returns:
-      Promise object that blocks and provides the return value when called.
-    """
-    payload = name, args, kwargs
-    self._conn.send((self._CALL, payload))
-    return self._receive
-
-  def close(self):
-    """Send a close message to the external process and join it."""
-    try:
-      self._conn.send((self._CLOSE, None))
-      self._conn.close()
-    except IOError:
-      # The connection was already closed.
-      pass
-    self._process.join()
-
-  def step(self, action, blocking=True):
-    """Step the environment.
-
-    Args:
-      action: The action to apply to the environment.
-      blocking: Whether to wait for the result.
-
-    Returns:
-      Transition tuple when blocking, otherwise callable that returns the
-      transition tuple.
-    """
-    promise = self.call("step", action)
-    if blocking:
-      return promise()
-    return promise
-
-  def reset(self, blocking=True):
-    """Reset the environment.
-
-    Args:
-      blocking: Whether to wait for the result.
-
-    Returns:
-      New observation when blocking, otherwise callable that returns the new
-      observation.
-    """
-    promise = self.call("reset")
-    if blocking:
-      return promise()
-    return promise
-
-  def _receive(self):
-    """Wait for a message from the worker process and return its payload.
-
-    Raises:
-      Exception: An exception was raised inside the worker process.
-      KeyError: The received message is of an unknown type.
-
-    Returns:
-      Payload object of the message.
-    """
-    message, payload = self._conn.recv()
-    # Re-raise exceptions in the main process.
-    if message == self._EXCEPTION:
-      stacktrace = payload
-      raise Exception(stacktrace)
-    if message == self._RESULT:
-      return payload
-    raise KeyError("Received message of unexpected type {}".format(message))
-
-  def _worker(self, constructor, conn):
-    """The process waits for actions and sends back environment results.
-
-    Args:
-      constructor: Constructor for the OpenAI Gym environment.
-      conn: Connection for communication to the main process.
-    """
-    try:
-      env = constructor()
-      while True:
-        try:
-          # Only block for short times to have keyboard exceptions be raised.
-          if not conn.poll(0.1):
-            continue
-          message, payload = conn.recv()
-        except (EOFError, KeyboardInterrupt):
-          break
-        if message == self._ACCESS:
-          name = payload
-          result = getattr(env, name)
-          conn.send((self._RESULT, result))
-          continue
-        if message == self._CALL:
-          name, args, kwargs = payload
-          result = getattr(env, name)(*args, **kwargs)
-          conn.send((self._RESULT, result))
-          continue
-        if message == self._CLOSE:
-          assert payload is None
-          env.close()
-          break
-        raise KeyError("Received message of unknown type {}".format(message))
-    except Exception:  # pylint: disable=broad-except
-      stacktrace = "".join(traceback.format_exception(*sys.exc_info()))  # pylint: disable=no-value-for-parameter
-      tf.logging.error("Error in environment process: {}".format(stacktrace))
-      conn.send((self._EXCEPTION, stacktrace))
-    conn.close()
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index 0cce5167c..692a37298 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -20,9 +20,10 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+import gym
 
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
-from tensor2tensor.rl.envs import utils
+
 import tensorflow as tf
 
 
@@ -41,10 +42,10 @@ def __init__(self, batch_env):
       batch_env: Batch environment.
     """
     self._batch_env = batch_env
-    observ_shape = utils.parse_shape(self._batch_env.observation_space)
-    observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
-    self.action_shape = list(utils.parse_shape(self._batch_env.action_space))
-    self.action_dtype = utils.parse_dtype(self._batch_env.action_space)
+    observ_shape = self._parse_shape(self._batch_env.observation_space)
+    observ_dtype = self._parse_dtype(self._batch_env.observation_space)
+    self.action_shape = list(self._parse_shape(self._batch_env.action_space))
+    self.action_dtype = self._parse_dtype(self._batch_env.action_space)
     with tf.variable_scope('env_temporary'):
       self._observ = tf.Variable(
           tf.zeros((len(self._batch_env),) + observ_shape, observ_dtype),
@@ -83,7 +84,7 @@ def simulate(self, action):
     with tf.name_scope('environment/simulate'):
       if action.dtype in (tf.float16, tf.float32, tf.float64):
         action = tf.check_numerics(action, 'action')
-      observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
+      observ_dtype = self._parse_dtype(self._batch_env.observation_space)
       observ, reward, done = tf.py_func(
           lambda a: self._batch_env.step(a)[:3], [action],
           [observ_dtype, tf.float32, tf.bool], name='step')
@@ -103,7 +104,7 @@ def _reset_non_empty(self, indices):
     Returns:
       Batch tensor of the new observations.
     """
-    observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
+    observ_dtype = self._parse_dtype(self._batch_env.observation_space)
     observ = tf.py_func(
         self._batch_env.reset, [indices], observ_dtype, name='reset')
     observ = tf.check_numerics(observ, 'observ')
@@ -119,3 +120,33 @@ def observ(self):
   def close(self):
     """Send close messages to the external process and join them."""
     self._batch_env.close()
+
+  def _parse_shape(self, space):
+    """Get a tensor shape from a OpenAI Gym space.
+
+    Args:
+      space: Gym space.
+
+    Returns:
+      Shape tuple.
+    """
+    if isinstance(space, gym.spaces.Discrete):
+      return ()
+    if isinstance(space, gym.spaces.Box):
+      return space.shape
+    raise NotImplementedError()
+
+  def _parse_dtype(self, space):
+    """Get a tensor dtype from a OpenAI Gym space.
+
+    Args:
+      space: Gym space.
+
+    Returns:
+      TensorFlow data type.
+    """
+    if isinstance(space, gym.spaces.Discrete):
+      return tf.int32
+    if isinstance(space, gym.spaces.Box):
+      return tf.float32
+    raise NotImplementedError()
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 5f1cfda6a..1c460f407 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -22,7 +22,6 @@
 from __future__ import print_function
 from tensor2tensor.layers import common_layers
 from tensor2tensor.rl.envs import in_graph_batch_env
-from tensor2tensor.rl.envs.utils import get_action_space
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
@@ -91,43 +90,38 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   flags are held in according variables.
   """
 
-  def __init__(self, hparams, length, simulation_random_starts=False,
-               intrinsic_reward_scale=0.):
+  def __init__(self, environment_lambda, length, problem,
+               simulation_random_starts=False, intrinsic_reward_scale=0.):
     """Batch of environments inside the TensorFlow graph."""
     self.length = length
-    environment_spec = hparams.environment_spec
-    initial_frames_problem = environment_spec.initial_frames_problem
-    self._min_reward = initial_frames_problem.min_reward
-    self._num_frames = initial_frames_problem.num_input_frames
+    self._min_reward = problem.min_reward
+    self._num_frames = problem.num_input_frames
     self._intrinsic_reward_scale = intrinsic_reward_scale
 
-    # initialization_env = environment_lambda()
-    model_hparams = trainer_lib.create_hparams(
+    initialization_env = environment_lambda()
+    hparams = trainer_lib.create_hparams(
         FLAGS.hparams_set, problem_name=FLAGS.problem)
-    model_hparams.force_full_predict = True
+    hparams.force_full_predict = True
     self._model = registry.model(FLAGS.model)(
-        model_hparams, tf.estimator.ModeKeys.PREDICT)
+        hparams, tf.estimator.ModeKeys.PREDICT)
 
-    _, self.action_shape, self.action_dtype = get_action_space(environment_spec)
+    self.action_space = initialization_env.action_space
+    self.action_shape = list(initialization_env.action_space.shape)
+    self.action_dtype = tf.int32
 
     if simulation_random_starts:
-      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
-                                               FLAGS.data_dir,
-                                               shuffle_files=True,
-                                               hparams=hparams)
+      dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
+                                shuffle_files=True, hparams=hparams)
       dataset = dataset.shuffle(buffer_size=100)
     else:
-      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
-                                               FLAGS.data_dir,
-                                               shuffle_files=True,
-                                               hparams=hparams).take(1)
+      dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
+                                shuffle_files=False, hparams=hparams).take(1)
 
     dataset = dataset.map(lambda x: x["inputs"]).repeat()
     self.history_buffer = HistoryBuffer(dataset, self.length)
 
-    shape = (self.length, initial_frames_problem.frame_height,
-             initial_frames_problem.frame_width,
-             initial_frames_problem.num_channels)
+    shape = (self.length, problem.frame_height, problem.frame_width,
+             problem.num_channels)
     self._observ = tf.Variable(tf.zeros(shape, tf.float32), trainable=False)
 
   def __len__(self):
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index fe6eedc1b..32bfcbde6 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -219,4 +219,26 @@ def _reset_non_empty(self, indices):
     assign_op = tf.scatter_update(self._observ, indices, new_values)
     with tf.control_dependencies([op_zero, assign_op]):
       return tf.identity(self.observ)
-    
\ No newline at end of file
+
+
+class MemoryWrapper(WrapperBase):
+  """Memory wrapper."""
+
+  def __init__(self, batch_env):
+    super(MemoryWrapper, self).__init__(batch_env)
+    MemoryWrapper.singleton = self
+    assert self._length == 1, "We support only one environment"
+    infinity = 10000000
+    self.speculum = tf.FIFOQueue(infinity, dtypes=[
+        tf.uint8, tf.float32, tf.int32, tf.bool])
+    self._observ = self._batch_env.observ
+
+  def simulate(self, action):
+    with tf.name_scope("environment/simulate"):  # Do we need this?
+      reward, done = self._batch_env.simulate(action)
+      image = tf.cast(self._batch_env.observ[0, ...], tf.uint8)
+      with tf.control_dependencies([reward, done]):
+        enqueue_op = self.speculum.enqueue(
+            [image, reward, action, done])
+        with tf.control_dependencies([enqueue_op]):
+          return tf.identity(reward), tf.identity(done)
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index ea6a8609c..de2af016a 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -14,11 +14,28 @@
 # limitations under the License.
 """Utilities for using batched environments."""
 
+# The code was based on Danijar Hafner's code from tf.agents:
+# https://github.com/tensorflow/agents/blob/master/agents/tools/wrappers.py
+# https://github.com/tensorflow/agents/blob/master/agents/scripts/utility.py
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import atexit
+import multiprocessing
+import os
+import random
+import signal
+import subprocess
+import sys
+import traceback
 import gym
+
+from tensor2tensor.rl.envs import batch_env
+
+from tensor2tensor.rl.envs import py_func_batch_env
+from tensor2tensor.rl.envs import simulated_batch_env
+
 import tensorflow as tf
 
 
@@ -65,63 +82,235 @@ def _reset(self, **kwargs):
     return self._last_returned[0]
 
 
-def get_action_space(environment_spec):
-  """Get action spece associated with environment spec
+class ExternalProcessEnv(object):
+  """Step environment in a separate process for lock free parallelism."""
 
-  Args:
-     environment_spec:  EnvironmentSpec object
+  # Message types for communication via the pipe.
+  _ACCESS = 1
+  _CALL = 2
+  _RESULT = 3
+  _EXCEPTION = 4
+  _CLOSE = 5
 
-  Returns:
-    OpenAi Gym action space
-  """
-  action_space = environment_spec.env_lambda().action_space
-  action_shape = list(parse_shape(action_space))
-  action_dtype = parse_dtype(action_space)
+  def __init__(self, constructor, xvfb):
+    """Step environment in a separate process for lock free parallelism.
 
-  return action_space, action_shape, action_dtype
+    The environment will be created in the external process by calling the
+    specified callable. This can be an environment class, or a function
+    creating the environment and potentially wrapping it. The returned
+    environment should not access global variables.
 
+    Args:
+      constructor: Callable that creates and returns an OpenAI gym environment.
+      xvfb:  Frame buffer.
 
-def get_policy(observations, hparams):
-  """Get policy network
+    Attributes:
+      observation_space: The cached observation space of the environment.
+      action_space: The cached action space of the environment.
+    """
+    self._conn, conn = multiprocessing.Pipe()
+    if xvfb:
+      server_id = random.randint(10000, 99999)
+      auth_file_id = random.randint(10000, 99999999999)
 
-  Args:
-    observations: Tensor with observations
-    hparams: parameters
+      xauthority_path = "/tmp/Xauthority_{}".format(auth_file_id)
 
-  Returns:
-    Tensor with policy and value function output
-  """
-  policy_network_lambda = hparams.policy_network
-  action_space, _, _ = get_action_space(hparams.environment_spec)
-  return policy_network_lambda(action_space, hparams, observations)
+      command = "Xvfb :{} -screen 0 1400x900x24 -nolisten tcp -auth {}".format(
+          server_id, xauthority_path)
+      with open(os.devnull, "w") as devnull:
+        proc = subprocess.Popen(command.split(), shell=False, stdout=devnull,
+                                stderr=devnull)
+        atexit.register(lambda: os.kill(proc.pid, signal.SIGKILL))
 
+      def constructor_using_xvfb():
+        os.environ["DISPLAY"] = ":{}".format(server_id)
+        os.environ["XAUTHORITY"] = xauthority_path
+        return constructor()
 
-def parse_shape(space):
-  """Get a tensor shape from a OpenAI Gym space.
+      self._process = multiprocessing.Process(
+          target=self._worker, args=(constructor_using_xvfb, conn))
+    else:
+      self._process = multiprocessing.Process(
+          target=self._worker, args=(constructor, conn))
 
-  Args:
-    space: Gym space.
+    atexit.register(self.close)
+    self._process.start()
+    self._observ_space = None
+    self._action_space = None
 
-  Returns:
-    Shape tuple.
-  """
-  if isinstance(space, gym.spaces.Discrete):
-    return ()
-  if isinstance(space, gym.spaces.Box):
-    return space.shape
-  raise NotImplementedError()
+  @property
+  def observation_space(self):
+    if not self._observ_space:
+      self._observ_space = self.__getattr__("observation_space")
+    return self._observ_space
 
-def parse_dtype(space):
-  """Get a tensor dtype from a OpenAI Gym space.
+  @property
+  def action_space(self):
+    if not self._action_space:
+      self._action_space = self.__getattr__("action_space")
+    return self._action_space
 
-  Args:
-    space: Gym space.
+  def __getattr__(self, name):
+    """Request an attribute from the environment.
 
-  Returns:
-    TensorFlow data type.
-  """
-  if isinstance(space, gym.spaces.Discrete):
-    return tf.int32
-  if isinstance(space, gym.spaces.Box):
-    return tf.float32
-  raise NotImplementedError()
+    Note that this involves communication with the external process, so it can
+    be slow.
+
+    Args:
+      name: Attribute to access.
+
+    Returns:
+      Value of the attribute.
+    """
+    self._conn.send((self._ACCESS, name))
+    return self._receive()
+
+  def call(self, name, *args, **kwargs):
+    """Asynchronously call a method of the external environment.
+
+    Args:
+      name: Name of the method to call.
+      *args: Positional arguments to forward to the method.
+      **kwargs: Keyword arguments to forward to the method.
+
+    Returns:
+      Promise object that blocks and provides the return value when called.
+    """
+    payload = name, args, kwargs
+    self._conn.send((self._CALL, payload))
+    return self._receive
+
+  def close(self):
+    """Send a close message to the external process and join it."""
+    try:
+      self._conn.send((self._CLOSE, None))
+      self._conn.close()
+    except IOError:
+      # The connection was already closed.
+      pass
+    self._process.join()
+
+  def step(self, action, blocking=True):
+    """Step the environment.
+
+    Args:
+      action: The action to apply to the environment.
+      blocking: Whether to wait for the result.
+
+    Returns:
+      Transition tuple when blocking, otherwise callable that returns the
+      transition tuple.
+    """
+    promise = self.call("step", action)
+    if blocking:
+      return promise()
+    return promise
+
+  def reset(self, blocking=True):
+    """Reset the environment.
+
+    Args:
+      blocking: Whether to wait for the result.
+
+    Returns:
+      New observation when blocking, otherwise callable that returns the new
+      observation.
+    """
+    promise = self.call("reset")
+    if blocking:
+      return promise()
+    return promise
+
+  def _receive(self):
+    """Wait for a message from the worker process and return its payload.
+
+    Raises:
+      Exception: An exception was raised inside the worker process.
+      KeyError: The received message is of an unknown type.
+
+    Returns:
+      Payload object of the message.
+    """
+    message, payload = self._conn.recv()
+    # Re-raise exceptions in the main process.
+    if message == self._EXCEPTION:
+      stacktrace = payload
+      raise Exception(stacktrace)
+    if message == self._RESULT:
+      return payload
+    raise KeyError("Received message of unexpected type {}".format(message))
+
+  def _worker(self, constructor, conn):
+    """The process waits for actions and sends back environment results.
+
+    Args:
+      constructor: Constructor for the OpenAI Gym environment.
+      conn: Connection for communication to the main process.
+    """
+    try:
+      env = constructor()
+      while True:
+        try:
+          # Only block for short times to have keyboard exceptions be raised.
+          if not conn.poll(0.1):
+            continue
+          message, payload = conn.recv()
+        except (EOFError, KeyboardInterrupt):
+          break
+        if message == self._ACCESS:
+          name = payload
+          result = getattr(env, name)
+          conn.send((self._RESULT, result))
+          continue
+        if message == self._CALL:
+          name, args, kwargs = payload
+          result = getattr(env, name)(*args, **kwargs)
+          conn.send((self._RESULT, result))
+          continue
+        if message == self._CLOSE:
+          assert payload is None
+          env.close()
+          break
+        raise KeyError("Received message of unknown type {}".format(message))
+    except Exception:  # pylint: disable=broad-except
+      stacktrace = "".join(traceback.format_exception(*sys.exc_info()))  # pylint: disable=no-value-for-parameter
+      tf.logging.error("Error in environment process: {}".format(stacktrace))
+      conn.send((self._EXCEPTION, stacktrace))
+    conn.close()
+
+
+def batch_env_factory(environment_lambda, hparams, num_agents, xvfb=False):
+  """Factory of batch envs."""
+  wrappers = hparams.in_graph_wrappers if hasattr(
+      hparams, "in_graph_wrappers") else []
+
+  if hparams.simulated_environment:
+    cur_batch_env = define_simulated_batch_env(
+        environment_lambda, num_agents, hparams.problem,
+        hparams.simulation_random_starts,
+        hparams.intrinsic_reward_scale)
+  else:
+    cur_batch_env = define_batch_env(environment_lambda, num_agents, xvfb=xvfb)
+  for w in wrappers:
+    cur_batch_env = w[0](cur_batch_env, **w[1])
+  return cur_batch_env
+
+
+def define_batch_env(constructor, num_agents, xvfb=False):
+  """Create environments and apply all desired wrappers."""
+  with tf.variable_scope("environments"):
+    envs = [
+        ExternalProcessEnv(constructor, xvfb)
+        for _ in range(num_agents)]
+    env = batch_env.BatchEnv(envs, blocking=False)
+    env = py_func_batch_env.PyFuncBatchEnv(env)
+    return env
+
+
+def define_simulated_batch_env(environment_lambda, num_agents, problem,
+                               simulation_random_starts=False,
+                               intrinsic_reward_scale=0.):
+  cur_batch_env = simulated_batch_env.SimulatedBatchEnv(
+      environment_lambda, num_agents, problem, simulation_random_starts,
+      intrinsic_reward_scale)
+  return cur_batch_env
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 8067f6aa6..7e55ec89a 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -30,13 +30,11 @@
 import math
 import os
 import time
-
-import copy
-
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.layers import discretization
 from tensor2tensor.rl import rl_trainer_lib
+from tensor2tensor.rl.envs.tf_atari_wrappers import StackAndSkipWrapper
 from tensor2tensor.rl.envs.tf_atari_wrappers import TimeLimitWrapper
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
@@ -141,31 +139,25 @@ def train_agent(problem_name, agent_model_dir,
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_epochs_num = hparams.ppo_epochs_num
   ppo_hparams.epochs_num = ppo_epochs_num
+  ppo_hparams.simulated_environment = True
+  ppo_hparams.simulation_random_starts = hparams.simulation_random_starts
+  ppo_hparams.intrinsic_reward_scale = hparams.intrinsic_reward_scale
   ppo_hparams.eval_every_epochs = 50
   ppo_hparams.save_models_every_epochs = ppo_epochs_num
   ppo_hparams.epoch_length = hparams.ppo_epoch_length
   ppo_hparams.num_agents = hparams.ppo_num_agents
-  # ppo_hparams.problem = gym_problem
+  ppo_hparams.problem = gym_problem
   ppo_hparams.world_model_dir = world_model_dir
   if hparams.ppo_learning_rate:
     ppo_hparams.learning_rate = hparams.ppo_learning_rate
-
-  environment_spec = copy.copy(gym_problem.environment_spec)
-  environment_spec.simulated_env = True
-  environment_spec.add_hparam("simulation_random_starts",
-                              hparams.simulation_random_starts)
-  environment_spec.add_hparam("intrinsic_reward_scale",
-                              hparams.intrinsic_reward_scale)
-  environment_spec.add_hparam("initial_frames_problem",
-                              gym_problem)
-
   # 4x for the StackAndSkipWrapper minus one to always finish for reporting.
   ppo_time_limit = (ppo_hparams.epoch_length - 1) * 4
-  wrappers = environment_spec.wrappers + \
-             [[TimeLimitWrapper, {"timelimit": ppo_time_limit}]]
-  environment_spec.wrappers = wrappers
 
-  ppo_hparams.add_hparam("environment_spec", environment_spec)
+  in_graph_wrappers = [
+      (TimeLimitWrapper, {"timelimit": ppo_time_limit}),
+      (StackAndSkipWrapper, {"skip": 4})]
+  in_graph_wrappers += gym_problem.in_graph_wrappers
+  ppo_hparams.add_hparam("in_graph_wrappers", in_graph_wrappers)
 
   with temporary_flags({
       "problem": problem_name,
@@ -175,7 +167,8 @@ def train_agent(problem_name, agent_model_dir,
       "data_dir": epoch_data_dir,
       "autoencoder_path": autoencoder_path,
   }):
-    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=epoch)
+    rl_trainer_lib.train(ppo_hparams, gym_problem.env_name, event_dir,
+                         agent_model_dir, epoch=epoch)
 
 
 def evaluate_world_model(simulated_problem_name, problem_name, hparams,
@@ -183,12 +176,12 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
                          autoencoder_path=None):
   """Generate simulated environment data and return reward accuracy."""
   gym_simulated_problem = registry.problem(simulated_problem_name)
-  # gym_problem = registry.problem(problem_name)
+  gym_problem = registry.problem(problem_name)
   sim_steps = hparams.simulated_env_generator_num_steps
   gym_simulated_problem.settable_num_steps = sim_steps
-  # gym_simulated_problem.real_env_problem = gym_problem
-  # gym_simulated_problem.simulation_random_starts = False
-  # gym_simulated_problem.intrinsic_reward_scale = 0.
+  gym_simulated_problem.real_env_problem = gym_problem
+  gym_simulated_problem.simulation_random_starts = False
+  gym_simulated_problem.intrinsic_reward_scale = 0.
   with temporary_flags({
       "problem": problem_name,
       "model": hparams.generative_model,
@@ -437,6 +430,8 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   return epoch_metrics[-1]
 
 
+
+
 def combine_training_data(problem, final_data_dir, old_data_dirs,
                           copy_last_eval_set=True):
   """Add training data from old_data_dirs into final_data_dir."""
@@ -543,8 +538,6 @@ def rl_modelrl_tiny():
           ppo_time_limit=5,
           ppo_epoch_length=5,
           ppo_num_agents=2,
-          # eval_world_model=True,
-          # simulation_random_starts=True
       ).values())
 
 
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index 750c02537..a62a8a920 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -19,7 +19,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-from tensor2tensor.rl.envs.utils import get_policy
 
 import tensorflow as tf
 
@@ -30,25 +29,25 @@ def get_optimiser(config):
   return config.optimizer(learning_rate=config.learning_rate)
 
 
-def define_ppo_step(data_points, optimizer, hparams):
+def define_ppo_step(data_points, policy_factory, optimizer, config):
   """Define ppo step."""
   observation, action, discounted_reward, norm_advantage, old_pdf = data_points
-  new_policy_dist, new_value, _ = get_policy(observation, hparams)
+  new_policy_dist, new_value, _ = policy_factory(observation)
   new_pdf = new_policy_dist.prob(action)
 
   ratio = new_pdf / old_pdf
-  clipped_ratio = tf.clip_by_value(ratio, 1 - hparams.clipping_coef,
-                                   1 + hparams.clipping_coef)
+  clipped_ratio = tf.clip_by_value(ratio, 1 - config.clipping_coef,
+                                   1 + config.clipping_coef)
 
   surrogate_objective = tf.minimum(clipped_ratio * norm_advantage,
                                    ratio * norm_advantage)
   policy_loss = -tf.reduce_mean(surrogate_objective)
 
   value_error = new_value - discounted_reward
-  value_loss = hparams.value_loss_coef * tf.reduce_mean(value_error ** 2)
+  value_loss = config.value_loss_coef * tf.reduce_mean(value_error ** 2)
 
   entropy = new_policy_dist.entropy()
-  entropy_loss = -hparams.entropy_loss_coef * tf.reduce_mean(entropy)
+  entropy_loss = -config.entropy_loss_coef * tf.reduce_mean(entropy)
 
   losses = [policy_loss, value_loss, entropy_loss]
 
@@ -60,9 +59,9 @@ def define_ppo_step(data_points, optimizer, hparams):
   gradients_flat = sum([gradient[0] for gradient in gradients], ())
   gradients_variables_flat = sum([gradient[1] for gradient in gradients], ())
 
-  if hparams.max_gradients_norm:
+  if config.max_gradients_norm:
     gradients_flat, _ = tf.clip_by_global_norm(gradients_flat,
-                                               hparams.max_gradients_norm)
+                                               config.max_gradients_norm)
 
   optimize_op = optimizer.apply_gradients(zip(gradients_flat,
                                               gradients_variables_flat))
@@ -71,7 +70,7 @@ def define_ppo_step(data_points, optimizer, hparams):
     return [tf.identity(x) for x in losses + gradients_norms]
 
 
-def define_ppo_epoch(memory, hparams):
+def define_ppo_epoch(memory, policy_factory, config):
   """PPO epoch."""
   observation, reward, done, action, old_pdf, value = memory
 
@@ -79,14 +78,14 @@ def define_ppo_epoch(memory, hparams):
   observation = tf.stop_gradient(observation)
   action = tf.stop_gradient(action)
   reward = tf.stop_gradient(reward)
-  if hasattr(hparams, "rewards_preprocessing_fun"):
-    reward = hparams.rewards_preprocessing_fun(reward)
+  if hasattr(config, "rewards_preprocessing_fun"):
+    reward = config.rewards_preprocessing_fun(reward)
   done = tf.stop_gradient(done)
   value = tf.stop_gradient(value)
   old_pdf = tf.stop_gradient(old_pdf)
 
   advantage = calculate_generalized_advantage_estimator(
-      reward, value, done, hparams.gae_gamma, hparams.gae_lambda)
+      reward, value, done, config.gae_gamma, config.gae_lambda)
 
   discounted_reward = tf.stop_gradient(advantage + value)
 
@@ -97,22 +96,23 @@ def define_ppo_epoch(memory, hparams):
 
   add_lists_elementwise = lambda l1, l2: [x + y for x, y in zip(l1, l2)]
 
-  number_of_batches = (hparams.epoch_length * hparams.optimization_epochs
-                       / hparams.optimization_batch_size)
+  number_of_batches = (config.epoch_length * config.optimization_epochs
+                       / config.optimization_batch_size)
 
   dataset = tf.data.Dataset.from_tensor_slices(
       (observation, action, discounted_reward, advantage_normalized, old_pdf))
-  dataset = dataset.shuffle(buffer_size=hparams.epoch_length,
+  dataset = dataset.shuffle(buffer_size=config.epoch_length,
                             reshuffle_each_iteration=True)
-  dataset = dataset.repeat(hparams.optimization_epochs)
-  dataset = dataset.batch(hparams.optimization_batch_size)
+  dataset = dataset.repeat(config.optimization_epochs)
+  dataset = dataset.batch(config.optimization_batch_size)
   iterator = dataset.make_initializable_iterator()
-  optimizer = get_optimiser(hparams)
+  optimizer = get_optimiser(config)
 
   with tf.control_dependencies([iterator.initializer]):
     ppo_step_rets = tf.scan(
         lambda a, i: add_lists_elementwise(  # pylint: disable=g-long-lambda
-            a, define_ppo_step(iterator.get_next(), optimizer, hparams)),
+            a, define_ppo_step(iterator.get_next(), policy_factory, optimizer,
+                               config)),
         tf.range(number_of_batches),
         [0., 0., 0., 0., 0., 0.],
         parallel_iterations=1)
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 4a3a093d3..4f4b923c3 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -17,35 +17,74 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
 import os
+import gym
 
 from tensor2tensor import models  # pylint: disable=unused-import
 from tensor2tensor.models.research import rl  # pylint: disable=unused-import
 from tensor2tensor.rl import collect
 from tensor2tensor.rl import ppo
+from tensor2tensor.rl.envs import tf_atari_wrappers
+from tensor2tensor.rl.envs import utils
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
 
 
-def define_train(hparams, event_dir): # pylint: disable=unused-argument
+def define_train(hparams, environment_spec, event_dir):
   """Define the training setup."""
+  policy_lambda = hparams.network
+
+  if environment_spec == "stacked_pong":
+    environment_spec = lambda: gym.make("PongNoFrameskip-v4")
+    wrappers = hparams.in_graph_wrappers if hasattr(
+        hparams, "in_graph_wrappers") else []
+    wrappers.append((tf_atari_wrappers.MaxAndSkipWrapper, {"skip": 4}))
+    hparams.in_graph_wrappers = wrappers
+  if isinstance(environment_spec, str):
+    env_lambda = lambda: gym.make(environment_spec)
+  else:
+    env_lambda = environment_spec
+
+  batch_env = utils.batch_env_factory(
+      env_lambda, hparams, num_agents=hparams.num_agents)
+
+  policy_factory = functools.partial(
+      policy_lambda, batch_env.action_space, hparams)
 
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
     memory, collect_summary = collect.define_collect(
-        hparams, "ppo_train", eval_phase=False,
+        policy_factory, batch_env, hparams, eval_phase=False,
         on_simulated=hparams.simulated_environment)
-    ppo_summary = ppo.define_ppo_epoch(memory, hparams)
+    ppo_summary = ppo.define_ppo_epoch(memory, policy_factory, hparams)
     summary = tf.summary.merge([collect_summary, ppo_summary])
 
-  return summary, None
+  with tf.variable_scope("eval", reuse=tf.AUTO_REUSE):
+    eval_env_lambda = env_lambda
+    if event_dir and hparams.video_during_eval:
+      # Some environments reset environments automatically, when reached done
+      # state. For them we shall record only every second episode.
+      d = 2 if env_lambda().metadata.get("semantics.autoreset") else 1
+      eval_env_lambda = lambda: gym.wrappers.Monitor(  # pylint: disable=g-long-lambda
+          env_lambda(), event_dir, video_callable=lambda i: i % d == 0)
+      eval_env_lambda = (
+          lambda: utils.EvalVideoWrapper(eval_env_lambda()))
+    eval_batch_env = utils.batch_env_factory(
+        eval_env_lambda, hparams,
+        num_agents=hparams.num_eval_agents, xvfb=hparams.video_during_eval)
+
+    _, eval_summary = collect.define_collect(
+        policy_factory, eval_batch_env, hparams, eval_phase=True)
+  return summary, eval_summary
 
 
-def train(hparams, event_dir=None, model_dir=None,
+def train(hparams, environment_spec, event_dir=None, model_dir=None,
           restore_agent=True, epoch=0):
   """Train."""
   with tf.name_scope("rl_train"):
-    train_summary_op, eval_summary_op = define_train(hparams, event_dir) # pylint: disable=unused-variable
+    train_summary_op, eval_summary_op = define_train(hparams, environment_spec,
+                                                     event_dir)
     if event_dir:
       summary_writer = tf.summary.FileWriter(
           event_dir, graph=tf.get_default_graph(), flush_secs=60)
@@ -84,8 +123,7 @@ def train(hparams, event_dir=None, model_dir=None,
           summary_writer.add_summary(summary, epoch_index)
         if (hparams.eval_every_epochs and
             epoch_index % hparams.eval_every_epochs == 0):
-          print("Eval is to be implemented") #TODO(piotrmilos):implement
-          # summary = sess.run(eval_summary_op)
+          summary = sess.run(eval_summary_op)
           if summary_writer and summary:
             summary_writer.add_summary(summary, epoch_index)
           else:
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index 6066bd7fc..9e14c4e9a 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -16,10 +16,8 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-from tensor2tensor.models.research.rl import simple_gym_spec, standard_atari_env_spec
 from tensor2tensor.rl import rl_trainer_lib
-from tensor2tensor.utils import trainer_lib, registry # pylint: disable=unused-import
+from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
 
@@ -33,18 +31,13 @@ def test_no_crash_pendulum(self):
     hparams = trainer_lib.create_hparams(
         "ppo_continuous_action_base",
         TrainTest.test_config)
-
-    hparams.add_hparam("environment_spec", simple_gym_spec("Pendulum-v0"))
-    rl_trainer_lib.train(hparams)
+    rl_trainer_lib.train(hparams, "Pendulum-v0")
 
   def test_no_crash_cartpole(self):
     hparams = trainer_lib.create_hparams(
         "ppo_discrete_action_base",
         TrainTest.test_config)
-
-    hparams.add_hparam("environment_spec",
-                       standard_atari_env_spec("CartPole-v0"))
-    rl_trainer_lib.train(hparams)
+    rl_trainer_lib.train(hparams, "CartPole-v0")
 
 
 if __name__ == "__main__":

From 1b26c5e62dc39c6d739a1f014535cf4dcfae12ee Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 21 Jun 2018 16:12:22 -0700
Subject: [PATCH 0180/2720] internal merge of PR #875

PiperOrigin-RevId: 201605027
---
 tensor2tensor/data_generators/gym_problems.py | 580 ++++++------------
 .../data_generators/gym_problems_test.py      |   4 +-
 tensor2tensor/data_generators/video_utils.py  |   2 +-
 tensor2tensor/models/research/rl.py           |  35 +-
 tensor2tensor/rl/collect.py                   | 164 +++--
 tensor2tensor/rl/envs/batch_env_factory.py    | 275 +++++++++
 tensor2tensor/rl/envs/py_func_batch_env.py    |  45 +-
 tensor2tensor/rl/envs/simulated_batch_env.py  |  40 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py    |  23 -
 tensor2tensor/rl/envs/utils.py                | 280 ++-------
 tensor2tensor/rl/model_rl_experiment.py       |  33 +-
 tensor2tensor/rl/ppo.py                       |  40 +-
 tensor2tensor/rl/rl_trainer_lib.py            |  54 +-
 tensor2tensor/rl/rl_trainer_lib_test.py       |  13 +-
 14 files changed, 744 insertions(+), 844 deletions(-)
 create mode 100644 tensor2tensor/rl/envs/batch_env_factory.py

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 28614f0c3..e2a2b7f3a 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -27,19 +27,11 @@
 from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
-from tensor2tensor.layers import discretization
-from tensor2tensor.models.research import autoencoders
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import collect
-from tensor2tensor.rl.envs import tf_atari_wrappers as atari
-from tensor2tensor.rl.envs.tf_atari_wrappers import StackAndSkipWrapper
-from tensor2tensor.rl.envs.tf_atari_wrappers import TimeLimitWrapper
-from tensor2tensor.rl.envs.utils import batch_env_factory
-
-
+from tensor2tensor.rl.envs import tf_atari_wrappers
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
-
 import tensorflow as tf
 
 
@@ -55,12 +47,94 @@
                      "still write out full-resolution frames.")
 
 
+def standard_atari_env_spec(env):
+  """Parameters of environment specification."""
+  standard_wrappers = [[tf_atari_wrappers.MaxAndSkipWrapper, {"skip": 4}]]
+  env_lambda = None
+  if isinstance(env, str):
+    env_lambda = lambda: gym.make(env)
+  if callable(env):
+    env_lambda = env
+  assert env is not None, "Unknown specification of environment"
+
+  return tf.contrib.training.HParams(env_lambda=env_lambda,
+                                     wrappers=standard_wrappers,
+                                     simulated_env=False)
+
+
 class GymDiscreteProblem(video_utils.VideoProblem):
   """Gym environment with discrete actions and rewards."""
 
   def __init__(self, *args, **kwargs):
     super(GymDiscreteProblem, self).__init__(*args, **kwargs)
     self._env = None
+    self._env = None
+    self.debug_dump_frames_path = "debug_frames_env"
+    self.settable_num_steps = 5000
+
+    self.environment_spec = self.get_environment_spec()
+    self.eval_phase = False
+    self.sum_of_rewards = 0.0
+    self.dones = 0
+
+  def _setup(self):
+    collect_hparams = rl.ppo_pong_base()
+    collect_hparams.add_hparam("environment_spec", self.environment_spec)
+
+    if not FLAGS.agent_policy_path:
+      collect_hparams.policy_network = rl.random_policy_fun
+
+    self._internal_memory_size = 10
+    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+      collect_hparams.epoch_length = self._internal_memory_size
+      # TODO(piotrmilos). it is possible to set more than 1.
+      collect_hparams.num_agents = 1
+      self.collect_memory, self.collect_trigger_op \
+        = collect.define_collect(collect_hparams, scope="gym_problems",
+                                 collect_level=0, eval_phase=self.eval_phase)
+
+  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
+    self._setup()
+    self.debug_dump_frames_path = os.path.join(
+        data_dir, self.debug_dump_frames_path)
+
+    with tf.Session() as sess:
+      sess.run(tf.global_variables_initializer())
+      self.restore_networks(sess)
+      pieces_generated = 0
+      memory_index = 0
+      memory = None
+      while pieces_generated < self.num_steps:
+        if memory is None or memory_index >= self._internal_memory_size:
+          sess.run(self.collect_trigger_op)
+          memory = sess.run(self.collect_memory)
+          memory_index = 0
+        data = [memory[i][memory_index][0] for i in range(4)]
+        memory_index += 1
+        observ, reward, done, action = data
+        observ = observ.astype(np.uint8)
+
+        self.sum_of_rewards += reward
+        self.dones += int(done)
+
+        ret_dict = {"frame": observ,
+                    "image/format": ["png"],
+                    "image/height": [self.frame_height],
+                    "image/width": [self.frame_width],
+                    "action": [int(action)],
+                    "done": [int(False)],
+                    "reward": [int(reward) - self.min_reward]}
+
+        yield ret_dict
+        pieces_generated += 1
+
+  def restore_networks(self, sess):
+    if FLAGS.agent_policy_path:
+      model_saver = tf.train.Saver(
+          tf.global_variables(".*network_parameters.*"))
+      ckpts = tf.train.get_checkpoint_state(FLAGS.agent_policy_path)
+      ckpt = ckpts.model_checkpoint_path
+      model_saver.restore(sess, ckpt)
 
   @property
   def num_input_frames(self):
@@ -87,6 +161,9 @@ def extra_reading_spec(self):
     }
     return data_fields, decoders
 
+  def get_environment_spec(self):
+    return standard_atari_env_spec(self.env_name)
+
   @property
   def is_generate_per_split(self):
     """Whether we have a train/test split or just hold out data."""
@@ -121,7 +198,7 @@ def num_rewards(self):
 
   @property
   def num_steps(self):
-    raise NotImplementedError()
+    return self.settable_num_steps
 
   @property
   def total_number_of_frames(self):
@@ -154,22 +231,73 @@ def hparams(self, defaults, unused_model_hparams):
     p.input_space_id = problem.SpaceID.IMAGE
     p.target_space_id = problem.SpaceID.IMAGE
 
-  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
-    next_observation = self.env.reset()
-    for _ in range(self.num_steps):
-      observation = next_observation
-      action = self.get_action(observation)
-      next_observation, reward, done, _ = self.env.step(action)
-      if done:
-        next_observation = self.env.reset()
-      yield {"frame": observation,
-             "action": [action],
-             "done": [done],
-             "reward": [int(reward - self.min_reward)]}
+
+class GymAEDiscreteProblem(GymDiscreteProblem):
+  pass
+
+
+class GymRealDiscreteProblem(GymDiscreteProblem):
+
+  def __init__(self, *args, **kwargs):
+    super(GymRealDiscreteProblem, self).__init__(*args, **kwargs)
+    self.make_extra_debug_info = False
+
+
+class GymSimulatedDiscreteProblem(GymDiscreteProblem):
+  """Simulated gym environment with discrete actions and rewards."""
+
+  def __init__(self, *args, **kwargs):
+    self.simulated_environment = True
+    self.debug_dump_frames_path = "debug_frames_sim"
+    self.intrinsic_reward_scale = 0.0
+    self.simulation_random_starts = False
+    super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
+
+  @property
+  def initial_frames_problem(self):
+    raise NotImplementedError()
+
+  def get_environment_spec(self):
+    env_spec = standard_atari_env_spec(self.env_name)
+
+    # Set reasonable time limit (as we do not simulate done).
+    real_env = env_spec.env_lambda()
+    if self.num_testing_steps is not None:
+      timelimit = self.num_testing_steps
+    else:
+      try:
+        # We assume that the real env is wrapped with TimeLimit.
+        history = self.num_input_frames
+        timelimit = real_env._max_episode_steps - history  # pylint: disable=protected-access
+      except:  # pylint: disable=bare-except
+        # If not, set some reasonable default.
+        timelimit = 100
+
+    env_spec.simulated_env = True
+    env_spec.add_hparam("simulation_random_starts",
+                        self.simulation_random_starts)
+    env_spec.add_hparam("intrinsic_reward_scale",
+                        self.intrinsic_reward_scale)
+    initial_frames_problem = registry.problem(self.initial_frames_problem)
+    env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
+    env_spec.wrappers.append(
+        [tf_atari_wrappers.TimeLimitWrapper, {"timelimit": timelimit}])
+
+    return env_spec
+
+  def restore_networks(self, sess):
+    super(GymSimulatedDiscreteProblem, self).restore_networks(sess)
+    # TODO(blazej): adjust regexp for different models.
+    env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
+    sess = tf.get_default_session()
+
+    ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
+    ckpt = ckpts.model_checkpoint_path
+    env_model_loader.restore(sess, ckpt)
 
 
 @registry.register_problem
-class GymPongRandom5k(GymDiscreteProblem):
+class GymPongRandom(GymDiscreteProblem):
   """Pong game, random actions."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
@@ -198,22 +326,9 @@ def min_reward(self):
   def num_rewards(self):
     return 3
 
-  @property
-  def num_steps(self):
-    return 5000
-
 
 @registry.register_problem
-class GymPongRandom50k(GymPongRandom5k):
-  """Pong game, random actions."""
-
-  @property
-  def num_steps(self):
-    return 50000
-
-
-@registry.register_problem
-class GymWrappedPongRandom5k(GymDiscreteProblem):
+class GymWrappedPongRandom(GymDiscreteProblem):
   """Pong game, random actions."""
 
   @property
@@ -228,10 +343,6 @@ def min_reward(self):
   def num_rewards(self):
     return 3
 
-  @property
-  def num_steps(self):
-    return 5000
-
 
 @registry.register_problem
 class GymWrappedLongPongRandom(GymDiscreteProblem):
@@ -249,17 +360,13 @@ def min_reward(self):
   def num_rewards(self):
     return 3
 
-  @property
-  def num_steps(self):
-    return 5000
-
   @property
   def num_testing_steps(self):
     return 100
 
 
 @registry.register_problem
-class GymWrappedBreakoutRandom5k(GymDiscreteProblem):
+class GymWrappedBreakoutRandom(GymDiscreteProblem):
   """Pong game, random actions."""
 
   @property
@@ -274,22 +381,18 @@ def min_reward(self):
   def num_rewards(self):
     return 3
 
-  @property
-  def num_steps(self):
-    return 5000
-
 
 @registry.register_problem
-class GymWrappedPongRandom50k(GymPongRandom5k):
-  """Pong game, random actions."""
+class GymSimulatedDiscreteProblemWithAgentOnPong(
+    GymSimulatedDiscreteProblem, GymPongRandom):
 
   @property
-  def num_steps(self):
-    return 50000
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_pong"
 
 
 @registry.register_problem
-class GymFreewayRandom5k(GymDiscreteProblem):
+class GymFreewayRandom(GymDiscreteProblem):
   """Freeway game, random actions."""
 
   @property
@@ -304,351 +407,25 @@ def min_reward(self):
   def num_rewards(self):
     return 2
 
-  @property
-  def num_steps(self):
-    return 5000
-
-
-@registry.register_problem
-class GymFreewayRandom50k(GymFreewayRandom5k):
-  """Freeway game, random actions."""
-
-  @property
-  def num_steps(self):
-    return 50000
-
-
-class GymDiscreteProblemWithAgent(GymDiscreteProblem):
-  """Gym environment with discrete actions and rewards and an agent."""
-
-  def __init__(self, *args, **kwargs):
-    super(GymDiscreteProblemWithAgent, self).__init__(*args, **kwargs)
-    self._env = None
-    self.debug_dump_frames_path = "debug_frames_env"
-    self.make_extra_debug_info = True
-    self.autoencoder_model = None
-
-    # Defaults.
-    self.environment_spec = lambda: gym.make(self.env_name)
-    self._real_env = None
-    self.real_env_problem = None
-    self.in_graph_wrappers = []
-    self.collect_hparams = rl.ppo_pong_base()
-    if FLAGS.autoencoder_path:
-      self.collect_hparams = rl.ppo_pong_ae_base()
-    self.settable_num_steps = 50000
-    self.simulated_environment = None
-    self.eval_phase = False
-    self.warm_up = 10  # TODO(piotrm): This should be probably removed.
-
-    # Debug info.
-    self.dones = 0
-    self.real_reward = 0
-    self.total_sim_reward, self.total_real_reward = 0.0, 0.0
-    self.sum_of_rewards = 0.0
-    self.successful_episode_reward_predictions = 0
-
-  @property
-  def real_env(self):
-    """Lazy caching environment construction."""
-    if self._real_env is None:
-      self._real_env = self.environment_spec()
-    return self._real_env
-
-  @property
-  def num_steps(self):
-    return self.settable_num_steps
-
-  @property
-  def raw_frame_height(self):
-    return self.env.observation_space.shape[0]
-
-  @property
-  def frame_height(self):
-    if FLAGS.autoencoder_path:
-      # TODO(lukaszkaiser): remove hard-coded autoencoder params.
-      return int(math.ceil(self.raw_frame_height / self.autoencoder_factor))
-    return self.raw_frame_height
-
-  @property
-  def autoencoder_factor(self):
-    """By how much to divide sizes when using autoencoders."""
-    hparams = autoencoders.autoencoder_discrete_pong()
-    return 2**hparams.num_hidden_layers
-
-  @property
-  def raw_frame_width(self):
-    return self.env.observation_space.shape[1]
-
-  @property
-  def frame_width(self):
-    if FLAGS.autoencoder_path:
-      # TODO(lukaszkaiser): remove hard-coded autoencoder params.
-      return int(math.ceil(self.raw_frame_width / self.autoencoder_factor))
-    return self.raw_frame_width
-
-  def setup_autoencoder(self):
-    if self.autoencoder_model is not None:
-      return
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      autoencoder_hparams = autoencoders.autoencoder_discrete_pong()
-      autoencoder_hparams.data_dir = "unused"
-      autoencoder_hparams.problem_hparams = self.get_hparams(
-          autoencoder_hparams)
-      autoencoder_hparams.problem = self
-      self.autoencoder_model = autoencoders.AutoencoderOrderedDiscrete(
-          autoencoder_hparams, tf.estimator.ModeKeys.EVAL)
-
-  def autoencode_tensor(self, x, batch_size=1):
-    if self.autoencoder_model is None:
-      return x
-    shape = [self.raw_frame_height, self.raw_frame_width, self.num_channels]
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      self.autoencoder_model.set_mode(tf.estimator.ModeKeys.EVAL)
-      autoencoded = self.autoencoder_model.encode(
-          tf.reshape(x, [batch_size, 1] + shape))
-    autoencoded = tf.reshape(
-        autoencoded, [batch_size] + self.frame_shape + [8])  # 8-bit groups.
-    if batch_size == 1:
-      autoencoded = tf.squeeze(autoencoded, axis=0)
-    return discretization.bit_to_int(autoencoded, 8)
-
-  def _setup(self):
-    if self.make_extra_debug_info:
-      self.report_reward_statistics_every = 10
-      self.dones = 0
-      self.real_reward = 0
-      # Slight weirdness to make sim env and real env aligned
-      if self.simulated_environment:
-        self.real_env.reset()
-        for _ in range(self.num_input_frames):
-          self.real_ob, _, _, _ = self.real_env.step(0)
-      self.total_sim_reward, self.total_real_reward = 0.0, 0.0
-      self.sum_of_rewards = 0.0
-      self.successful_episode_reward_predictions = 0
-
-    in_graph_wrappers = self.in_graph_wrappers + [
-        (atari.MemoryWrapper, {}), (StackAndSkipWrapper, {"skip": 4})]
-    env_hparams = tf.contrib.training.HParams(
-        in_graph_wrappers=in_graph_wrappers,
-        problem=self.real_env_problem if self.real_env_problem else self,
-        simulated_environment=self.simulated_environment)
-    if self.simulated_environment:
-      env_hparams.add_hparam("simulation_random_starts",
-                             self.simulation_random_starts)
-      env_hparams.add_hparam("intrinsic_reward_scale",
-                             self.intrinsic_reward_scale)
-
-    generator_batch_env = batch_env_factory(
-        self.environment_spec, env_hparams, num_agents=1, xvfb=False)
-
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      if FLAGS.agent_policy_path:
-        policy_lambda = self.collect_hparams.network
-      else:
-        # When no agent_policy_path is set, just generate random samples.
-        policy_lambda = rl.random_policy_fun
-
-    if FLAGS.autoencoder_path:
-      # TODO(lukaszkaiser): remove hard-coded autoencoder params.
-      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-        self.setup_autoencoder()
-        autoencoder_model = self.autoencoder_model
-        # Feeds for autoencoding.
-        shape = [self.raw_frame_height, self.raw_frame_width, self.num_channels]
-        self.autoencoder_feed = tf.placeholder(tf.int32, shape=shape)
-        self.autoencoder_result = self.autoencode_tensor(self.autoencoder_feed)
-        # Now for autodecoding.
-        shape = self.frame_shape
-        self.autodecoder_feed = tf.placeholder(tf.int32, shape=shape)
-        bottleneck = tf.reshape(
-            discretization.int_to_bit(self.autodecoder_feed, 8),
-            [1, 1, self.frame_height, self.frame_width, self.num_channels * 8])
-        autoencoder_model.set_mode(tf.estimator.ModeKeys.PREDICT)
-        self.autodecoder_result = autoencoder_model.decode(bottleneck)
-
-    def preprocess_fn(x):
-      shape = [self.raw_frame_height, self.raw_frame_width, self.num_channels]
-      # TODO(lukaszkaiser): we assume x comes from StackAndSkipWrapper skip=4.
-      xs = [tf.reshape(t, [1] + shape) for t in tf.split(x, 4, axis=-1)]
-      autoencoded = self.autoencode_tensor(tf.concat(xs, axis=0), batch_size=4)
-      encs = [tf.squeeze(t, axis=[0]) for t in tf.split(autoencoded, 4, axis=0)]
-      res = tf.to_float(tf.concat(encs, axis=-1))
-      return tf.expand_dims(res, axis=0)
-
-    # TODO(lukaszkaiser): x is from StackAndSkipWrapper thus 4*num_channels.
-    shape = [1, self.frame_height, self.frame_width, 4 * self.num_channels]
-    do_preprocess = (self.autoencoder_model is not None and
-                     not self.simulated_environment)
-    preprocess = (preprocess_fn, shape) if do_preprocess else None
-
-    def policy(x):
-      return policy_lambda(self.environment_spec().action_space,
-                           self.collect_hparams, x)
-
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      self.collect_hparams.epoch_length = 10
-      _, self.collect_trigger_op = collect.define_collect(
-          policy, generator_batch_env, self.collect_hparams,
-          eval_phase=self.eval_phase,
-          scope="define_collect", preprocess=preprocess)
-
-    self.avilable_data_size_op = atari.MemoryWrapper.singleton.speculum.size()
-    self.data_get_op = atari.MemoryWrapper.singleton.speculum.dequeue()
-
-  def restore_networks(self, sess):
-    if FLAGS.agent_policy_path:
-      model_saver = tf.train.Saver(
-          tf.global_variables(".*network_parameters.*"))
-      ckpts = tf.train.get_checkpoint_state(FLAGS.agent_policy_path)
-      ckpt = ckpts.model_checkpoint_path
-      model_saver.restore(sess, ckpt)
-    if FLAGS.autoencoder_path:
-      autoencoder_saver = tf.train.Saver(
-          tf.global_variables("autoencoder.*"))
-      ckpts = tf.train.get_checkpoint_state(FLAGS.autoencoder_path)
-      ckpt = ckpts.model_checkpoint_path
-      autoencoder_saver.restore(sess, ckpt)
-
-  def autoencode(self, image, sess):
-    return sess.run(self.autoencoder_result, {self.autoencoder_feed: image})
-
-  def autodecode(self, encoded, sess):
-    res = sess.run(self.autodecoder_result, {self.autodecoder_feed: encoded})
-    return res[0, 0, :self.raw_frame_height, :self.raw_frame_width, :]
-
-  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
-    self._setup()
-    self.debug_dump_frames_path = os.path.join(
-        data_dir, self.debug_dump_frames_path)
-
-    with tf.Session() as sess:
-      sess.run(tf.global_variables_initializer())
-      self.restore_networks(sess)
-      if FLAGS.only_use_ae_for_policy:
-        # If only the policy should use the autoencoder, then reset the flag so
-        # that other components here act as though there is no autoencoder and
-        # so write out full-resolution images. The policy graph was already
-        # built and self.collect_trigger_op is all that's used from it.
-        FLAGS.autoencoder_path = None
-      pieces_generated = 0
-      while pieces_generated < self.num_steps + self.warm_up:
-        avilable_data_size = sess.run(self.avilable_data_size_op)
-        if avilable_data_size < 1:
-          sess.run(self.collect_trigger_op)
-        observ, reward, action, done = sess.run(self.data_get_op)
-        debug_im = None
-        if self.make_extra_debug_info:
-          self.total_sim_reward += reward
-          if not self.simulated_environment:
-            self.real_ob = observ
-            self.real_reward = reward
-          if not FLAGS.autoencoder_path:
-            err = np.ndarray.astype(np.maximum(np.abs(
-                self.real_ob - observ, dtype=np.int) - 10, 0),
-                                    np.uint8)
-            debug_im = np.concatenate([observ, self.real_ob, err], axis=1)
-          if done:
-            self.dones += 1
-            self.sum_of_rewards += self.real_reward
-            if self.total_real_reward == self.total_sim_reward:
-              self.successful_episode_reward_predictions += 1
-
-            self.total_real_reward = 0.0
-            self.total_sim_reward = 0.0
-            self.real_reward = 0
-            if self.simulated_environment:
-              self.real_env.reset()
-              # Slight weirdness to make sim env and real env aligned
-              for _ in range(self.num_input_frames):
-                self.real_ob, _, _, _ = self.real_env.step(0)
-          else:
-            if self.simulated_environment:
-              self.real_ob, self.real_reward, _, _ = self.real_env.step(action)
-            self.total_real_reward += self.real_reward
-            self.sum_of_rewards += self.real_reward
-        if FLAGS.autoencoder_path:
-          if self.simulated_environment:
-            debug_im = self.autodecode(observ, sess)
-          else:
-            orig_observ = observ
-            observ = self.autoencode(observ, sess)
-            debug_im = np.concatenate([self.autodecode(observ, sess),
-                                       orig_observ], axis=1)
-        ret_dict = {"frame": observ,
-                    "image/format": ["png"],
-                    "image/height": [self.frame_height],
-                    "image/width": [self.frame_width],
-                    "action": [int(action)],
-                    "done": [int(False)],
-                    "reward": [int(reward) - self.min_reward]}
-        if self.make_extra_debug_info:
-          ret_dict["image/debug"] = debug_im
-        yield ret_dict
-        pieces_generated += 1
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgent(GymDiscreteProblemWithAgent):
-  """Simulated gym environment with discrete actions and rewards."""
-
-  def __init__(self, *args, **kwargs):
-    super(GymSimulatedDiscreteProblemWithAgent, self).__init__(*args, **kwargs)
-    self.simulated_environment = True
-    self.make_extra_debug_info = True
-    self.debug_dump_frames_path = "debug_frames_sim"
-
-  @property
-  def real_env(self):
-    """Lazy caching environment construction."""
-    if self._real_env is None:
-      self._real_env = self.environment_spec()
-      if self.num_testing_steps is not None:
-        timelimit = self.num_testing_steps
-      else:
-        try:
-          # We assume that the real env is wrapped with TimeLimit.
-          history = self.num_input_frames
-          timelimit = self.real_env._max_episode_steps - history  # pylint: disable=protected-access
-        except:  # pylint: disable=bare-except
-          # If not, set some reasonable default.
-          timelimit = 100
-      self.in_graph_wrappers.append(
-          (TimeLimitWrapper, {"timelimit": timelimit}))
-    return self._real_env
-
-  def restore_networks(self, sess):
-    super(GymSimulatedDiscreteProblemWithAgent, self).restore_networks(sess)
-    # TODO(blazej): adjust regexp for different models.
-    env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
-    sess = tf.get_default_session()
-
-    ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
-    ckpt = ckpts.model_checkpoint_path
-    env_model_loader.restore(sess, ckpt)
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnPong(
-    GymSimulatedDiscreteProblemWithAgent, GymPongRandom5k):
-  pass
-
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnPong(
-    GymDiscreteProblemWithAgent, GymPongRandom5k):
+    GymRealDiscreteProblem, GymPongRandom):
   pass
 
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
-    GymSimulatedDiscreteProblemWithAgent, GymWrappedPongRandom5k):
-  pass
+    GymSimulatedDiscreteProblem, GymWrappedPongRandom):
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_pong"
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedLongPong(
-    GymDiscreteProblemWithAgent, GymWrappedLongPongRandom):
+    GymRealDiscreteProblem, GymWrappedLongPongRandom):
   pass
 
 
@@ -660,13 +437,16 @@ class GymDiscreteProblemWithAgentOnWrappedLongPongAe(  # with autoencoder
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPong(
-    GymSimulatedDiscreteProblemWithAgent, GymWrappedLongPongRandom):
-  pass
+    GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_long_pong"
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedBreakout(
-    GymDiscreteProblemWithAgent, GymWrappedBreakoutRandom5k):
+    GymRealDiscreteProblem, GymWrappedBreakoutRandom):
   pass
 
 
@@ -678,13 +458,16 @@ class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
-    GymSimulatedDiscreteProblemWithAgent, GymWrappedBreakoutRandom5k):
-  pass
+    GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_breakout"
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedPong(
-    GymDiscreteProblemWithAgent, GymWrappedPongRandom5k):
+    GymRealDiscreteProblem, GymWrappedPongRandom):
   """GymDiscreteProblemWithAgentOnWrappedPong."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
@@ -714,13 +497,16 @@ class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnFreeway(
-    GymSimulatedDiscreteProblemWithAgent, GymFreewayRandom5k):
-  pass
+    GymSimulatedDiscreteProblem, GymFreewayRandom):
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_freeway"
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnFreeway(
-    GymDiscreteProblemWithAgent, GymFreewayRandom5k):
+    GymRealDiscreteProblem, GymFreewayRandom):
   """Freeway with agent."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index 1bbd2a374..f7aaa8e2e 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -35,8 +35,8 @@ def setUpClass(cls):
     os.mkdir(cls.tmp_dir)
 
   def testGymAtariBoots(self):
-    problem = gym_problems.GymPongRandom5k()
-    self.assertEqual(5000, problem.num_steps)
+    problem = gym_problems.GymPongRandom()
+    self.assertEqual(210, problem.frame_height)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index bf20c94ec..b0bdb77d5 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -337,7 +337,7 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
             unencoded_debug = features.pop("image/debug")
             encoded_debug = sess.run(encoded_image_t, feed_dict={
                 image_t: unencoded_debug})
-            features["image/encoded_debug"] = encoded_debug
+            features["image/encoded_debug"] = [encoded_debug]
           yield features
 
   def generate_encoded_samples_debug(self, data_dir, tmp_dir, dataset_split):
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index a3410a40f..f62e175d4 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -59,14 +59,28 @@ def ppo_base_v1():
 @registry.register_hparams
 def ppo_continuous_action_base():
   hparams = ppo_base_v1()
-  hparams.add_hparam("network", feed_forward_gaussian_fun)
+  hparams.add_hparam("policy_network", feed_forward_gaussian_fun)
+  hparams.add_hparam("policy_network_params", "basic_policy_parameters")
   return hparams
 
 
+@registry.register_hparams
+def basic_policy_parameters():
+  wrappers = None
+  return tf.contrib.training.HParams(wrappers=wrappers)
+
+
 @registry.register_hparams
 def ppo_discrete_action_base():
   hparams = ppo_base_v1()
-  hparams.add_hparam("network", feed_forward_categorical_fun)
+  hparams.add_hparam("policy_network", feed_forward_categorical_fun)
+  return hparams
+
+
+@registry.register_hparams
+def discrete_random_action_base():
+  hparams = common_hparams.basic_params1()
+  hparams.add_hparam("policy_network", random_policy_fun)
   return hparams
 
 
@@ -102,13 +116,28 @@ def ppo_pong_base():
   hparams.optimization_epochs = 2
   hparams.epochs_num = 1000
   hparams.num_eval_agents = 1
-  hparams.network = feed_forward_cnn_small_categorical_fun
+  hparams.policy_network = feed_forward_cnn_small_categorical_fun
   hparams.clipping_coef = 0.2
   hparams.optimization_batch_size = 4
   hparams.max_gradients_norm = 0.5
   return hparams
 
 
+def simple_gym_spec(env):
+  """Parameters of environment specification."""
+  standard_wrappers = None
+  env_lambda = None
+  if isinstance(env, str):
+    env_lambda = lambda: gym.make(env)
+  if callable(env):
+    env_lambda = env
+  assert env is not None, "Unknown specification of environment"
+
+  return tf.contrib.training.HParams(env_lambda=env_lambda,
+                                     wrappers=standard_wrappers,
+                                     simulated_env=False)
+
+
 @registry.register_hparams
 def ppo_pong_ae_base():
   """Pong autoencoder base parameters."""
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index a3e807efd..90a53721d 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -17,44 +17,101 @@
 from __future__ import division
 from __future__ import print_function
 
+
+import copy
+
+from tensor2tensor.rl.envs.batch_env_factory import batch_env_factory
+from tensor2tensor.rl.envs.tf_atari_wrappers import WrapperBase
+from tensor2tensor.rl.envs.utils import get_policy
+
 import tensorflow as tf
 
 
-def define_collect(policy_factory, batch_env, hparams,
-                   eval_phase, policy_to_actions_lambda=None,
-                   scope="", preprocess=None, on_simulated=False):
-  """Collect trajectories."""
-  eval_phase = tf.convert_to_tensor(eval_phase)
-  on_simulated = tf.convert_to_tensor(on_simulated)
+def _rollout_metadata(batch_env):
+  """Metadata for rollouts."""
   batch_env_shape = batch_env.observ.get_shape().as_list()
-  if preprocess is not None:
-    batch_env_shape = preprocess[1]
-  memory_shape = [hparams.epoch_length] + [batch_env_shape[0]]
-  memories_shapes_and_types = [
-      # observation
-      (memory_shape + batch_env_shape[1:], tf.float32),
-      (memory_shape, tf.float32),      # reward
-      (memory_shape, tf.bool),         # done
-      # action
-      (memory_shape + batch_env.action_shape, batch_env.action_dtype),
-      (memory_shape, tf.float32),      # pdf
-      (memory_shape, tf.float32),      # value function
+  batch_size = [batch_env_shape[0]]
+  shapes_types_names = [
+      (batch_size + batch_env_shape[1:], tf.float32, "observation"),
+      (batch_size, tf.float32, "reward"),
+      (batch_size, tf.bool, "done"),
+      (batch_size + batch_env.action_shape, batch_env.action_dtype, "action"),
+      (batch_size, tf.float32, "pdf"),
+      (batch_size, tf.float32, "value_function"),
   ]
-  memory = [tf.Variable(tf.zeros(shape, dtype), trainable=False)
-            for (shape, dtype) in memories_shapes_and_types]
-  with tf.variable_scope(scope):
+  return shapes_types_names
+
+
+class _MemoryWrapper(WrapperBase):
+  """Memory wrapper."""
+
+  def __init__(self, batch_env):
+    super(_MemoryWrapper, self).__init__(batch_env)
+    infinity = 10000000
+    meta_data = list(zip(*_rollout_metadata(batch_env)))
+    shapes = meta_data[0][:4]
+    dtypes = meta_data[1][:4]
+    self.speculum = tf.FIFOQueue(infinity, shapes=shapes, dtypes=dtypes)
+    observs_shape = batch_env.observ.shape
+    observ_dtype = tf.float32
+    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
+                               trainable=False)
+
+  def simulate(self, action):
+    reward, done = self._batch_env.simulate(action)
+    with tf.control_dependencies([reward, done]):
+      assign = self._observ.assign(self._batch_env.observ)
+
+    with tf.control_dependencies([assign]):
+      enqueue_op = self.speculum.enqueue(
+          [self._observ.read_value(), reward, done, action])
+      with tf.control_dependencies([enqueue_op]):
+        return tf.identity(reward), tf.identity(done)
+
+
+def define_collect(hparams, scope, eval_phase,
+                   collect_level=-1,
+                   policy_to_actions_lambda=None,
+                   on_simulated=False):
+  """Collect trajectories."""
+
+  with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
+    batch_env = batch_env_factory(hparams)
+    environment_wrappers = hparams.environment_spec.wrappers
+    wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
+    # Put memory wrapper at the level you want to gather observations at.
+    # Negative indices need to be shifted for insert to work correctly.
+    collect_level = collect_level if \
+      collect_level >= 0 else len(wrappers) + collect_level + 1
+    wrappers.insert(collect_level, [_MemoryWrapper, {}])
+    rollout_metadata = None
+    speculum = None
+    for w in wrappers:
+      batch_env = w[0](batch_env, **w[1])
+      if w[0] == _MemoryWrapper:
+        rollout_metadata = _rollout_metadata(batch_env)
+        speculum = batch_env.speculum
+
+    eval_phase = tf.convert_to_tensor(eval_phase)
+    on_simulated = tf.convert_to_tensor(on_simulated)
+
+    memory = [tf.get_variable("collect_memory_{}".format(name),
+                              shape=[hparams.epoch_length]+shape,
+                              dtype=dtype,
+                              initializer=tf.zeros_initializer(),
+                              trainable=False)
+              for (shape, dtype, name) in rollout_metadata]
+
     cumulative_rewards = tf.get_variable("cumulative_rewards", len(batch_env),
                                          trainable=False)
 
-  should_reset_var = tf.Variable(True, trainable=False)
+    should_reset_var = tf.Variable(True, trainable=False)
 
-  zeros_tensor = tf.zeros(len(batch_env))
+    zeros_tensor = tf.zeros(len(batch_env))
 
   def group():
-    return tf.group(
-        batch_env.reset(tf.range(len(batch_env))),
-        tf.assign(cumulative_rewards, zeros_tensor))
-
+    return tf.group(batch_env.reset(tf.range(len(batch_env))),
+                    tf.assign(cumulative_rewards, zeros_tensor))
   reset_op = tf.cond(
       tf.logical_or(should_reset_var, tf.logical_or(eval_phase, on_simulated)),
       group, tf.no_op)
@@ -71,24 +128,43 @@ def step(index, scores_sum, scores_num):
       # operation. We are waiting for tf.copy:
       # https://github.com/tensorflow/tensorflow/issues/11186
       obs_copy = batch_env.observ + 0
-      if preprocess is not None:
-        obs_copy = preprocess[0](obs_copy)
-      actor_critic = policy_factory(tf.expand_dims(obs_copy, 0))
-      policy = actor_critic.policy
-      if policy_to_actions_lambda:
-        action = policy_to_actions_lambda(policy)
-      else:
-        action = tf.cond(eval_phase,
-                         policy.mode,
-                         policy.sample)
-      postprocessed_action = actor_critic.action_postprocessing(action)
-      simulate_output = batch_env.simulate(postprocessed_action[0, ...])
-      pdf = policy.prob(action)[0]
-      with tf.control_dependencies(simulate_output):
-        reward, done = simulate_output
+
+      def env_step(arg1, arg2):  # pylint: disable=unused-argument
+        """Step of the environment."""
+        actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
+        policy = actor_critic.policy
+        if policy_to_actions_lambda:
+          action = policy_to_actions_lambda(policy)
+        else:
+          action = tf.cond(eval_phase,
+                           policy.mode,
+                           policy.sample)
+
+        postprocessed_action = actor_critic.action_postprocessing(action)
+        simulate_output = batch_env.simulate(postprocessed_action[0, ...])
+
+        pdf = policy.prob(action)[0]
+        value_function = actor_critic.value[0]
+        pdf = tf.reshape(pdf, shape=(hparams.num_agents,))
+        value_function = tf.reshape(value_function, shape=(hparams.num_agents,))
+
+        with tf.control_dependencies(simulate_output):
+          return tf.identity(pdf), tf.identity(value_function)
+
+      pdf, value_function = tf.while_loop(
+          lambda _1, _2: tf.equal(speculum.size(), 0),
+          env_step,
+          [tf.constant(0.0, shape=(hparams.num_agents,)),
+           tf.constant(0.0, shape=(hparams.num_agents,))],
+          parallel_iterations=1,
+          back_prop=False,)
+
+      with tf.control_dependencies([pdf, value_function]):
+        obs, reward, done, action = speculum.dequeue()
+
         done = tf.reshape(done, (len(batch_env),))
-        to_save = [obs_copy, reward, done, action[0, ...], pdf,
-                   actor_critic.value[0]]
+        to_save = [obs, reward, done, action,
+                   pdf, value_function]
         save_ops = [tf.scatter_update(memory_slot, index, value)
                     for memory_slot, value in zip(memory, to_save)]
         cumulate_rewards_op = cumulative_rewards.assign_add(reward)
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
new file mode 100644
index 000000000..70bb129a0
--- /dev/null
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -0,0 +1,275 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Utilities for creating batched environments."""
+
+# The code was based on Danijar Hafner's code from tf.agents:
+# https://github.com/tensorflow/agents/blob/master/agents/tools/wrappers.py
+# https://github.com/tensorflow/agents/blob/master/agents/scripts/utility.py
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import atexit
+import multiprocessing
+import os
+import random
+import signal
+import subprocess
+import sys
+import traceback
+
+from tensor2tensor.rl.envs import batch_env
+from tensor2tensor.rl.envs import py_func_batch_env
+from tensor2tensor.rl.envs import simulated_batch_env
+
+import tensorflow as tf
+
+
+def batch_env_factory(hparams, xvfb=False):
+  """Factory of batch envs."""
+
+  environment_spec = hparams.environment_spec
+
+  if environment_spec.simulated_env:
+    # TODO(piotrmilos): Consider passing only relevant parameters
+    cur_batch_env = _define_simulated_batch_env(
+        hparams, hparams.num_agents,
+        hparams.simulation_random_starts,
+        hparams.intrinsic_reward_scale)
+  else:
+
+    cur_batch_env = _define_batch_env(hparams.environment_spec,
+                                      hparams.num_agents,
+                                      xvfb=xvfb)
+  return cur_batch_env
+
+
+def _define_batch_env(environment_spec, num_agents, xvfb=False):
+  """Create environments and apply all desired wrappers."""
+
+  with tf.variable_scope("environments"):
+    envs = [
+        ExternalProcessEnv(environment_spec.env_lambda, xvfb)
+        for _ in range(num_agents)]
+    env = batch_env.BatchEnv(envs, blocking=False)
+    env = py_func_batch_env.PyFuncBatchEnv(env)
+    return env
+
+
+def _define_simulated_batch_env(hparams, num_agents,
+                                simulation_random_starts=False,
+                                intrinsic_reward_scale=0.):
+  cur_batch_env = simulated_batch_env.SimulatedBatchEnv(
+      hparams, num_agents, simulation_random_starts,
+      intrinsic_reward_scale)
+  return cur_batch_env
+
+
+class ExternalProcessEnv(object):
+  """Step environment in a separate process for lock free parallelism."""
+
+  # Message types for communication via the pipe.
+  _ACCESS = 1
+  _CALL = 2
+  _RESULT = 3
+  _EXCEPTION = 4
+  _CLOSE = 5
+
+  def __init__(self, constructor, xvfb):
+    """Step environment in a separate process for lock free parallelism.
+
+    The environment will be created in the external process by calling the
+    specified callable. This can be an environment class, or a function
+    creating the environment and potentially wrapping it. The returned
+    environment should not access global variables.
+
+    Args:
+      constructor: Callable that creates and returns an OpenAI gym environment.
+      xvfb:  Frame buffer.
+
+    Attributes:
+      observation_space: The cached observation space of the environment.
+      action_space: The cached action space of the environment.
+    """
+    self._conn, conn = multiprocessing.Pipe()
+    if xvfb:
+      server_id = random.randint(10000, 99999)
+      auth_file_id = random.randint(10000, 99999999999)
+
+      xauthority_path = "/tmp/Xauthority_{}".format(auth_file_id)
+
+      command = "Xvfb :{} -screen 0 1400x900x24 -nolisten tcp -auth {}".format(
+          server_id, xauthority_path)
+      with open(os.devnull, "w") as devnull:
+        proc = subprocess.Popen(command.split(), shell=False, stdout=devnull,
+                                stderr=devnull)
+        atexit.register(lambda: os.kill(proc.pid, signal.SIGKILL))
+
+      def constructor_using_xvfb():
+        os.environ["DISPLAY"] = ":{}".format(server_id)
+        os.environ["XAUTHORITY"] = xauthority_path
+        return constructor()
+
+      self._process = multiprocessing.Process(
+          target=self._worker, args=(constructor_using_xvfb, conn))
+    else:
+      self._process = multiprocessing.Process(
+          target=self._worker, args=(constructor, conn))
+
+    atexit.register(self.close)
+    self._process.start()
+    self._observ_space = None
+    self._action_space = None
+
+  @property
+  def observation_space(self):
+    if not self._observ_space:
+      self._observ_space = self.__getattr__("observation_space")
+    return self._observ_space
+
+  @property
+  def action_space(self):
+    if not self._action_space:
+      self._action_space = self.__getattr__("action_space")
+    return self._action_space
+
+  def __getattr__(self, name):
+    """Request an attribute from the environment.
+
+    Note that this involves communication with the external process, so it can
+    be slow.
+
+    Args:
+      name: Attribute to access.
+
+    Returns:
+      Value of the attribute.
+    """
+    self._conn.send((self._ACCESS, name))
+    return self._receive()
+
+  def call(self, name, *args, **kwargs):
+    """Asynchronously call a method of the external environment.
+
+    Args:
+      name: Name of the method to call.
+      *args: Positional arguments to forward to the method.
+      **kwargs: Keyword arguments to forward to the method.
+
+    Returns:
+      Promise object that blocks and provides the return value when called.
+    """
+    payload = name, args, kwargs
+    self._conn.send((self._CALL, payload))
+    return self._receive
+
+  def close(self):
+    """Send a close message to the external process and join it."""
+    try:
+      self._conn.send((self._CLOSE, None))
+      self._conn.close()
+    except IOError:
+      # The connection was already closed.
+      pass
+    self._process.join()
+
+  def step(self, action, blocking=True):
+    """Step the environment.
+
+    Args:
+      action: The action to apply to the environment.
+      blocking: Whether to wait for the result.
+
+    Returns:
+      Transition tuple when blocking, otherwise callable that returns the
+      transition tuple.
+    """
+    promise = self.call("step", action)
+    if blocking:
+      return promise()
+    return promise
+
+  def reset(self, blocking=True):
+    """Reset the environment.
+
+    Args:
+      blocking: Whether to wait for the result.
+
+    Returns:
+      New observation when blocking, otherwise callable that returns the new
+      observation.
+    """
+    promise = self.call("reset")
+    if blocking:
+      return promise()
+    return promise
+
+  def _receive(self):
+    """Wait for a message from the worker process and return its payload.
+
+    Raises:
+      Exception: An exception was raised inside the worker process.
+      KeyError: The received message is of an unknown type.
+
+    Returns:
+      Payload object of the message.
+    """
+    message, payload = self._conn.recv()
+    # Re-raise exceptions in the main process.
+    if message == self._EXCEPTION:
+      stacktrace = payload
+      raise Exception(stacktrace)
+    if message == self._RESULT:
+      return payload
+    raise KeyError("Received message of unexpected type {}".format(message))
+
+  def _worker(self, constructor, conn):
+    """The process waits for actions and sends back environment results.
+
+    Args:
+      constructor: Constructor for the OpenAI Gym environment.
+      conn: Connection for communication to the main process.
+    """
+    try:
+      env = constructor()
+      while True:
+        try:
+          # Only block for short times to have keyboard exceptions be raised.
+          if not conn.poll(0.1):
+            continue
+          message, payload = conn.recv()
+        except (EOFError, KeyboardInterrupt):
+          break
+        if message == self._ACCESS:
+          name = payload
+          result = getattr(env, name)
+          conn.send((self._RESULT, result))
+          continue
+        if message == self._CALL:
+          name, args, kwargs = payload
+          result = getattr(env, name)(*args, **kwargs)
+          conn.send((self._RESULT, result))
+          continue
+        if message == self._CLOSE:
+          assert payload is None
+          env.close()
+          break
+        raise KeyError("Received message of unknown type {}".format(message))
+    except Exception:  # pylint: disable=broad-except
+      stacktrace = "".join(traceback.format_exception(*sys.exc_info()))  # pylint: disable=no-value-for-parameter
+      tf.logging.error("Error in environment process: {}".format(stacktrace))
+      conn.send((self._EXCEPTION, stacktrace))
+    conn.close()
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index 692a37298..b5f909bdf 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -20,10 +20,9 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-import gym
 
+from tensor2tensor.rl.envs import utils
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
-
 import tensorflow as tf
 
 
@@ -42,10 +41,10 @@ def __init__(self, batch_env):
       batch_env: Batch environment.
     """
     self._batch_env = batch_env
-    observ_shape = self._parse_shape(self._batch_env.observation_space)
-    observ_dtype = self._parse_dtype(self._batch_env.observation_space)
-    self.action_shape = list(self._parse_shape(self._batch_env.action_space))
-    self.action_dtype = self._parse_dtype(self._batch_env.action_space)
+    observ_shape = utils.parse_shape(self._batch_env.observation_space)
+    observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
+    self.action_shape = list(utils.parse_shape(self._batch_env.action_space))
+    self.action_dtype = utils.parse_dtype(self._batch_env.action_space)
     with tf.variable_scope('env_temporary'):
       self._observ = tf.Variable(
           tf.zeros((len(self._batch_env),) + observ_shape, observ_dtype),
@@ -84,7 +83,7 @@ def simulate(self, action):
     with tf.name_scope('environment/simulate'):
       if action.dtype in (tf.float16, tf.float32, tf.float64):
         action = tf.check_numerics(action, 'action')
-      observ_dtype = self._parse_dtype(self._batch_env.observation_space)
+      observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
       observ, reward, done = tf.py_func(
           lambda a: self._batch_env.step(a)[:3], [action],
           [observ_dtype, tf.float32, tf.bool], name='step')
@@ -104,7 +103,7 @@ def _reset_non_empty(self, indices):
     Returns:
       Batch tensor of the new observations.
     """
-    observ_dtype = self._parse_dtype(self._batch_env.observation_space)
+    observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
     observ = tf.py_func(
         self._batch_env.reset, [indices], observ_dtype, name='reset')
     observ = tf.check_numerics(observ, 'observ')
@@ -120,33 +119,3 @@ def observ(self):
   def close(self):
     """Send close messages to the external process and join them."""
     self._batch_env.close()
-
-  def _parse_shape(self, space):
-    """Get a tensor shape from a OpenAI Gym space.
-
-    Args:
-      space: Gym space.
-
-    Returns:
-      Shape tuple.
-    """
-    if isinstance(space, gym.spaces.Discrete):
-      return ()
-    if isinstance(space, gym.spaces.Box):
-      return space.shape
-    raise NotImplementedError()
-
-  def _parse_dtype(self, space):
-    """Get a tensor dtype from a OpenAI Gym space.
-
-    Args:
-      space: Gym space.
-
-    Returns:
-      TensorFlow data type.
-    """
-    if isinstance(space, gym.spaces.Discrete):
-      return tf.int32
-    if isinstance(space, gym.spaces.Box):
-      return tf.float32
-    raise NotImplementedError()
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 1c460f407..5f1cfda6a 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -22,6 +22,7 @@
 from __future__ import print_function
 from tensor2tensor.layers import common_layers
 from tensor2tensor.rl.envs import in_graph_batch_env
+from tensor2tensor.rl.envs.utils import get_action_space
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
@@ -90,38 +91,43 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   flags are held in according variables.
   """
 
-  def __init__(self, environment_lambda, length, problem,
-               simulation_random_starts=False, intrinsic_reward_scale=0.):
+  def __init__(self, hparams, length, simulation_random_starts=False,
+               intrinsic_reward_scale=0.):
     """Batch of environments inside the TensorFlow graph."""
     self.length = length
-    self._min_reward = problem.min_reward
-    self._num_frames = problem.num_input_frames
+    environment_spec = hparams.environment_spec
+    initial_frames_problem = environment_spec.initial_frames_problem
+    self._min_reward = initial_frames_problem.min_reward
+    self._num_frames = initial_frames_problem.num_input_frames
     self._intrinsic_reward_scale = intrinsic_reward_scale
 
-    initialization_env = environment_lambda()
-    hparams = trainer_lib.create_hparams(
+    # initialization_env = environment_lambda()
+    model_hparams = trainer_lib.create_hparams(
         FLAGS.hparams_set, problem_name=FLAGS.problem)
-    hparams.force_full_predict = True
+    model_hparams.force_full_predict = True
     self._model = registry.model(FLAGS.model)(
-        hparams, tf.estimator.ModeKeys.PREDICT)
+        model_hparams, tf.estimator.ModeKeys.PREDICT)
 
-    self.action_space = initialization_env.action_space
-    self.action_shape = list(initialization_env.action_space.shape)
-    self.action_dtype = tf.int32
+    _, self.action_shape, self.action_dtype = get_action_space(environment_spec)
 
     if simulation_random_starts:
-      dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
-                                shuffle_files=True, hparams=hparams)
+      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
+                                               FLAGS.data_dir,
+                                               shuffle_files=True,
+                                               hparams=hparams)
       dataset = dataset.shuffle(buffer_size=100)
     else:
-      dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
-                                shuffle_files=False, hparams=hparams).take(1)
+      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
+                                               FLAGS.data_dir,
+                                               shuffle_files=True,
+                                               hparams=hparams).take(1)
 
     dataset = dataset.map(lambda x: x["inputs"]).repeat()
     self.history_buffer = HistoryBuffer(dataset, self.length)
 
-    shape = (self.length, problem.frame_height, problem.frame_width,
-             problem.num_channels)
+    shape = (self.length, initial_frames_problem.frame_height,
+             initial_frames_problem.frame_width,
+             initial_frames_problem.num_channels)
     self._observ = tf.Variable(tf.zeros(shape, tf.float32), trainable=False)
 
   def __len__(self):
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 32bfcbde6..9a1708cfb 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -219,26 +219,3 @@ def _reset_non_empty(self, indices):
     assign_op = tf.scatter_update(self._observ, indices, new_values)
     with tf.control_dependencies([op_zero, assign_op]):
       return tf.identity(self.observ)
-
-
-class MemoryWrapper(WrapperBase):
-  """Memory wrapper."""
-
-  def __init__(self, batch_env):
-    super(MemoryWrapper, self).__init__(batch_env)
-    MemoryWrapper.singleton = self
-    assert self._length == 1, "We support only one environment"
-    infinity = 10000000
-    self.speculum = tf.FIFOQueue(infinity, dtypes=[
-        tf.uint8, tf.float32, tf.int32, tf.bool])
-    self._observ = self._batch_env.observ
-
-  def simulate(self, action):
-    with tf.name_scope("environment/simulate"):  # Do we need this?
-      reward, done = self._batch_env.simulate(action)
-      image = tf.cast(self._batch_env.observ[0, ...], tf.uint8)
-      with tf.control_dependencies([reward, done]):
-        enqueue_op = self.speculum.enqueue(
-            [image, reward, action, done])
-        with tf.control_dependencies([enqueue_op]):
-          return tf.identity(reward), tf.identity(done)
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index de2af016a..b23dc358e 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -14,28 +14,11 @@
 # limitations under the License.
 """Utilities for using batched environments."""
 
-# The code was based on Danijar Hafner's code from tf.agents:
-# https://github.com/tensorflow/agents/blob/master/agents/tools/wrappers.py
-# https://github.com/tensorflow/agents/blob/master/agents/scripts/utility.py
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import atexit
-import multiprocessing
-import os
-import random
-import signal
-import subprocess
-import sys
-import traceback
 import gym
-
-from tensor2tensor.rl.envs import batch_env
-
-from tensor2tensor.rl.envs import py_func_batch_env
-from tensor2tensor.rl.envs import simulated_batch_env
-
 import tensorflow as tf
 
 
@@ -82,235 +65,64 @@ def _reset(self, **kwargs):
     return self._last_returned[0]
 
 
-class ExternalProcessEnv(object):
-  """Step environment in a separate process for lock free parallelism."""
-
-  # Message types for communication via the pipe.
-  _ACCESS = 1
-  _CALL = 2
-  _RESULT = 3
-  _EXCEPTION = 4
-  _CLOSE = 5
-
-  def __init__(self, constructor, xvfb):
-    """Step environment in a separate process for lock free parallelism.
-
-    The environment will be created in the external process by calling the
-    specified callable. This can be an environment class, or a function
-    creating the environment and potentially wrapping it. The returned
-    environment should not access global variables.
-
-    Args:
-      constructor: Callable that creates and returns an OpenAI gym environment.
-      xvfb:  Frame buffer.
-
-    Attributes:
-      observation_space: The cached observation space of the environment.
-      action_space: The cached action space of the environment.
-    """
-    self._conn, conn = multiprocessing.Pipe()
-    if xvfb:
-      server_id = random.randint(10000, 99999)
-      auth_file_id = random.randint(10000, 99999999999)
-
-      xauthority_path = "/tmp/Xauthority_{}".format(auth_file_id)
-
-      command = "Xvfb :{} -screen 0 1400x900x24 -nolisten tcp -auth {}".format(
-          server_id, xauthority_path)
-      with open(os.devnull, "w") as devnull:
-        proc = subprocess.Popen(command.split(), shell=False, stdout=devnull,
-                                stderr=devnull)
-        atexit.register(lambda: os.kill(proc.pid, signal.SIGKILL))
-
-      def constructor_using_xvfb():
-        os.environ["DISPLAY"] = ":{}".format(server_id)
-        os.environ["XAUTHORITY"] = xauthority_path
-        return constructor()
-
-      self._process = multiprocessing.Process(
-          target=self._worker, args=(constructor_using_xvfb, conn))
-    else:
-      self._process = multiprocessing.Process(
-          target=self._worker, args=(constructor, conn))
-
-    atexit.register(self.close)
-    self._process.start()
-    self._observ_space = None
-    self._action_space = None
-
-  @property
-  def observation_space(self):
-    if not self._observ_space:
-      self._observ_space = self.__getattr__("observation_space")
-    return self._observ_space
-
-  @property
-  def action_space(self):
-    if not self._action_space:
-      self._action_space = self.__getattr__("action_space")
-    return self._action_space
+def get_action_space(environment_spec):
+  """Get action spece associated with environment spec.
 
-  def __getattr__(self, name):
-    """Request an attribute from the environment.
+  Args:
+     environment_spec:  EnvironmentSpec object
 
-    Note that this involves communication with the external process, so it can
-    be slow.
-
-    Args:
-      name: Attribute to access.
-
-    Returns:
-      Value of the attribute.
-    """
-    self._conn.send((self._ACCESS, name))
-    return self._receive()
-
-  def call(self, name, *args, **kwargs):
-    """Asynchronously call a method of the external environment.
-
-    Args:
-      name: Name of the method to call.
-      *args: Positional arguments to forward to the method.
-      **kwargs: Keyword arguments to forward to the method.
-
-    Returns:
-      Promise object that blocks and provides the return value when called.
-    """
-    payload = name, args, kwargs
-    self._conn.send((self._CALL, payload))
-    return self._receive
-
-  def close(self):
-    """Send a close message to the external process and join it."""
-    try:
-      self._conn.send((self._CLOSE, None))
-      self._conn.close()
-    except IOError:
-      # The connection was already closed.
-      pass
-    self._process.join()
-
-  def step(self, action, blocking=True):
-    """Step the environment.
-
-    Args:
-      action: The action to apply to the environment.
-      blocking: Whether to wait for the result.
-
-    Returns:
-      Transition tuple when blocking, otherwise callable that returns the
-      transition tuple.
-    """
-    promise = self.call("step", action)
-    if blocking:
-      return promise()
-    return promise
-
-  def reset(self, blocking=True):
-    """Reset the environment.
-
-    Args:
-      blocking: Whether to wait for the result.
+  Returns:
+    OpenAi Gym action space
+  """
+  action_space = environment_spec.env_lambda().action_space
+  action_shape = list(parse_shape(action_space))
+  action_dtype = parse_dtype(action_space)
 
-    Returns:
-      New observation when blocking, otherwise callable that returns the new
-      observation.
-    """
-    promise = self.call("reset")
-    if blocking:
-      return promise()
-    return promise
+  return action_space, action_shape, action_dtype
 
-  def _receive(self):
-    """Wait for a message from the worker process and return its payload.
 
-    Raises:
-      Exception: An exception was raised inside the worker process.
-      KeyError: The received message is of an unknown type.
+def get_policy(observations, hparams):
+  """Get a policy network.
 
-    Returns:
-      Payload object of the message.
-    """
-    message, payload = self._conn.recv()
-    # Re-raise exceptions in the main process.
-    if message == self._EXCEPTION:
-      stacktrace = payload
-      raise Exception(stacktrace)
-    if message == self._RESULT:
-      return payload
-    raise KeyError("Received message of unexpected type {}".format(message))
+  Args:
+    observations: Tensor with observations
+    hparams: parameters
 
-  def _worker(self, constructor, conn):
-    """The process waits for actions and sends back environment results.
+  Returns:
+    Tensor with policy and value function output
+  """
+  policy_network_lambda = hparams.policy_network
+  action_space, _, _ = get_action_space(hparams.environment_spec)
+  return policy_network_lambda(action_space, hparams, observations)
 
-    Args:
-      constructor: Constructor for the OpenAI Gym environment.
-      conn: Connection for communication to the main process.
-    """
-    try:
-      env = constructor()
-      while True:
-        try:
-          # Only block for short times to have keyboard exceptions be raised.
-          if not conn.poll(0.1):
-            continue
-          message, payload = conn.recv()
-        except (EOFError, KeyboardInterrupt):
-          break
-        if message == self._ACCESS:
-          name = payload
-          result = getattr(env, name)
-          conn.send((self._RESULT, result))
-          continue
-        if message == self._CALL:
-          name, args, kwargs = payload
-          result = getattr(env, name)(*args, **kwargs)
-          conn.send((self._RESULT, result))
-          continue
-        if message == self._CLOSE:
-          assert payload is None
-          env.close()
-          break
-        raise KeyError("Received message of unknown type {}".format(message))
-    except Exception:  # pylint: disable=broad-except
-      stacktrace = "".join(traceback.format_exception(*sys.exc_info()))  # pylint: disable=no-value-for-parameter
-      tf.logging.error("Error in environment process: {}".format(stacktrace))
-      conn.send((self._EXCEPTION, stacktrace))
-    conn.close()
 
+def parse_shape(space):
+  """Get a tensor shape from a OpenAI Gym space.
 
-def batch_env_factory(environment_lambda, hparams, num_agents, xvfb=False):
-  """Factory of batch envs."""
-  wrappers = hparams.in_graph_wrappers if hasattr(
-      hparams, "in_graph_wrappers") else []
+  Args:
+    space: Gym space.
 
-  if hparams.simulated_environment:
-    cur_batch_env = define_simulated_batch_env(
-        environment_lambda, num_agents, hparams.problem,
-        hparams.simulation_random_starts,
-        hparams.intrinsic_reward_scale)
-  else:
-    cur_batch_env = define_batch_env(environment_lambda, num_agents, xvfb=xvfb)
-  for w in wrappers:
-    cur_batch_env = w[0](cur_batch_env, **w[1])
-  return cur_batch_env
+  Returns:
+    Shape tuple.
+  """
+  if isinstance(space, gym.spaces.Discrete):
+    return ()
+  if isinstance(space, gym.spaces.Box):
+    return space.shape
+  raise NotImplementedError()
 
 
-def define_batch_env(constructor, num_agents, xvfb=False):
-  """Create environments and apply all desired wrappers."""
-  with tf.variable_scope("environments"):
-    envs = [
-        ExternalProcessEnv(constructor, xvfb)
-        for _ in range(num_agents)]
-    env = batch_env.BatchEnv(envs, blocking=False)
-    env = py_func_batch_env.PyFuncBatchEnv(env)
-    return env
+def parse_dtype(space):
+  """Get a tensor dtype from a OpenAI Gym space.
 
+  Args:
+    space: Gym space.
 
-def define_simulated_batch_env(environment_lambda, num_agents, problem,
-                               simulation_random_starts=False,
-                               intrinsic_reward_scale=0.):
-  cur_batch_env = simulated_batch_env.SimulatedBatchEnv(
-      environment_lambda, num_agents, problem, simulation_random_starts,
-      intrinsic_reward_scale)
-  return cur_batch_env
+  Returns:
+    TensorFlow data type.
+  """
+  if isinstance(space, gym.spaces.Discrete):
+    return tf.int32
+  if isinstance(space, gym.spaces.Box):
+    return tf.float32
+  raise NotImplementedError()
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 7e55ec89a..4b412d364 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -26,15 +26,16 @@
 from __future__ import print_function
 
 import contextlib
+import copy
 import datetime
 import math
 import os
 import time
+
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.layers import discretization
 from tensor2tensor.rl import rl_trainer_lib
-from tensor2tensor.rl.envs.tf_atari_wrappers import StackAndSkipWrapper
 from tensor2tensor.rl.envs.tf_atari_wrappers import TimeLimitWrapper
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
@@ -139,25 +140,30 @@ def train_agent(problem_name, agent_model_dir,
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_epochs_num = hparams.ppo_epochs_num
   ppo_hparams.epochs_num = ppo_epochs_num
-  ppo_hparams.simulated_environment = True
-  ppo_hparams.simulation_random_starts = hparams.simulation_random_starts
-  ppo_hparams.intrinsic_reward_scale = hparams.intrinsic_reward_scale
   ppo_hparams.eval_every_epochs = 50
   ppo_hparams.save_models_every_epochs = ppo_epochs_num
   ppo_hparams.epoch_length = hparams.ppo_epoch_length
   ppo_hparams.num_agents = hparams.ppo_num_agents
-  ppo_hparams.problem = gym_problem
   ppo_hparams.world_model_dir = world_model_dir
   if hparams.ppo_learning_rate:
     ppo_hparams.learning_rate = hparams.ppo_learning_rate
+
+  environment_spec = copy.copy(gym_problem.environment_spec)
+  environment_spec.simulated_env = True
+  environment_spec.add_hparam("simulation_random_starts",
+                              hparams.simulation_random_starts)
+  environment_spec.add_hparam("intrinsic_reward_scale",
+                              hparams.intrinsic_reward_scale)
+  environment_spec.add_hparam("initial_frames_problem",
+                              gym_problem)
+
   # 4x for the StackAndSkipWrapper minus one to always finish for reporting.
   ppo_time_limit = (ppo_hparams.epoch_length - 1) * 4
+  wrappers = environment_spec.wrappers + \
+             [[TimeLimitWrapper, {"timelimit": ppo_time_limit}]]
+  environment_spec.wrappers = wrappers
 
-  in_graph_wrappers = [
-      (TimeLimitWrapper, {"timelimit": ppo_time_limit}),
-      (StackAndSkipWrapper, {"skip": 4})]
-  in_graph_wrappers += gym_problem.in_graph_wrappers
-  ppo_hparams.add_hparam("in_graph_wrappers", in_graph_wrappers)
+  ppo_hparams.add_hparam("environment_spec", environment_spec)
 
   with temporary_flags({
       "problem": problem_name,
@@ -167,8 +173,7 @@ def train_agent(problem_name, agent_model_dir,
       "data_dir": epoch_data_dir,
       "autoencoder_path": autoencoder_path,
   }):
-    rl_trainer_lib.train(ppo_hparams, gym_problem.env_name, event_dir,
-                         agent_model_dir, epoch=epoch)
+    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=epoch)
 
 
 def evaluate_world_model(simulated_problem_name, problem_name, hparams,
@@ -176,12 +181,8 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
                          autoencoder_path=None):
   """Generate simulated environment data and return reward accuracy."""
   gym_simulated_problem = registry.problem(simulated_problem_name)
-  gym_problem = registry.problem(problem_name)
   sim_steps = hparams.simulated_env_generator_num_steps
   gym_simulated_problem.settable_num_steps = sim_steps
-  gym_simulated_problem.real_env_problem = gym_problem
-  gym_simulated_problem.simulation_random_starts = False
-  gym_simulated_problem.intrinsic_reward_scale = 0.
   with temporary_flags({
       "problem": problem_name,
       "model": hparams.generative_model,
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index a62a8a920..750c02537 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -19,6 +19,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+from tensor2tensor.rl.envs.utils import get_policy
 
 import tensorflow as tf
 
@@ -29,25 +30,25 @@ def get_optimiser(config):
   return config.optimizer(learning_rate=config.learning_rate)
 
 
-def define_ppo_step(data_points, policy_factory, optimizer, config):
+def define_ppo_step(data_points, optimizer, hparams):
   """Define ppo step."""
   observation, action, discounted_reward, norm_advantage, old_pdf = data_points
-  new_policy_dist, new_value, _ = policy_factory(observation)
+  new_policy_dist, new_value, _ = get_policy(observation, hparams)
   new_pdf = new_policy_dist.prob(action)
 
   ratio = new_pdf / old_pdf
-  clipped_ratio = tf.clip_by_value(ratio, 1 - config.clipping_coef,
-                                   1 + config.clipping_coef)
+  clipped_ratio = tf.clip_by_value(ratio, 1 - hparams.clipping_coef,
+                                   1 + hparams.clipping_coef)
 
   surrogate_objective = tf.minimum(clipped_ratio * norm_advantage,
                                    ratio * norm_advantage)
   policy_loss = -tf.reduce_mean(surrogate_objective)
 
   value_error = new_value - discounted_reward
-  value_loss = config.value_loss_coef * tf.reduce_mean(value_error ** 2)
+  value_loss = hparams.value_loss_coef * tf.reduce_mean(value_error ** 2)
 
   entropy = new_policy_dist.entropy()
-  entropy_loss = -config.entropy_loss_coef * tf.reduce_mean(entropy)
+  entropy_loss = -hparams.entropy_loss_coef * tf.reduce_mean(entropy)
 
   losses = [policy_loss, value_loss, entropy_loss]
 
@@ -59,9 +60,9 @@ def define_ppo_step(data_points, policy_factory, optimizer, config):
   gradients_flat = sum([gradient[0] for gradient in gradients], ())
   gradients_variables_flat = sum([gradient[1] for gradient in gradients], ())
 
-  if config.max_gradients_norm:
+  if hparams.max_gradients_norm:
     gradients_flat, _ = tf.clip_by_global_norm(gradients_flat,
-                                               config.max_gradients_norm)
+                                               hparams.max_gradients_norm)
 
   optimize_op = optimizer.apply_gradients(zip(gradients_flat,
                                               gradients_variables_flat))
@@ -70,7 +71,7 @@ def define_ppo_step(data_points, policy_factory, optimizer, config):
     return [tf.identity(x) for x in losses + gradients_norms]
 
 
-def define_ppo_epoch(memory, policy_factory, config):
+def define_ppo_epoch(memory, hparams):
   """PPO epoch."""
   observation, reward, done, action, old_pdf, value = memory
 
@@ -78,14 +79,14 @@ def define_ppo_epoch(memory, policy_factory, config):
   observation = tf.stop_gradient(observation)
   action = tf.stop_gradient(action)
   reward = tf.stop_gradient(reward)
-  if hasattr(config, "rewards_preprocessing_fun"):
-    reward = config.rewards_preprocessing_fun(reward)
+  if hasattr(hparams, "rewards_preprocessing_fun"):
+    reward = hparams.rewards_preprocessing_fun(reward)
   done = tf.stop_gradient(done)
   value = tf.stop_gradient(value)
   old_pdf = tf.stop_gradient(old_pdf)
 
   advantage = calculate_generalized_advantage_estimator(
-      reward, value, done, config.gae_gamma, config.gae_lambda)
+      reward, value, done, hparams.gae_gamma, hparams.gae_lambda)
 
   discounted_reward = tf.stop_gradient(advantage + value)
 
@@ -96,23 +97,22 @@ def define_ppo_epoch(memory, policy_factory, config):
 
   add_lists_elementwise = lambda l1, l2: [x + y for x, y in zip(l1, l2)]
 
-  number_of_batches = (config.epoch_length * config.optimization_epochs
-                       / config.optimization_batch_size)
+  number_of_batches = (hparams.epoch_length * hparams.optimization_epochs
+                       / hparams.optimization_batch_size)
 
   dataset = tf.data.Dataset.from_tensor_slices(
       (observation, action, discounted_reward, advantage_normalized, old_pdf))
-  dataset = dataset.shuffle(buffer_size=config.epoch_length,
+  dataset = dataset.shuffle(buffer_size=hparams.epoch_length,
                             reshuffle_each_iteration=True)
-  dataset = dataset.repeat(config.optimization_epochs)
-  dataset = dataset.batch(config.optimization_batch_size)
+  dataset = dataset.repeat(hparams.optimization_epochs)
+  dataset = dataset.batch(hparams.optimization_batch_size)
   iterator = dataset.make_initializable_iterator()
-  optimizer = get_optimiser(config)
+  optimizer = get_optimiser(hparams)
 
   with tf.control_dependencies([iterator.initializer]):
     ppo_step_rets = tf.scan(
         lambda a, i: add_lists_elementwise(  # pylint: disable=g-long-lambda
-            a, define_ppo_step(iterator.get_next(), policy_factory, optimizer,
-                               config)),
+            a, define_ppo_step(iterator.get_next(), optimizer, hparams)),
         tf.range(number_of_batches),
         [0., 0., 0., 0., 0., 0.],
         parallel_iterations=1)
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 4f4b923c3..46ac189c0 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -17,74 +17,35 @@
 from __future__ import division
 from __future__ import print_function
 
-import functools
 import os
-import gym
 
 from tensor2tensor import models  # pylint: disable=unused-import
 from tensor2tensor.models.research import rl  # pylint: disable=unused-import
 from tensor2tensor.rl import collect
 from tensor2tensor.rl import ppo
-from tensor2tensor.rl.envs import tf_atari_wrappers
-from tensor2tensor.rl.envs import utils
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
 
 
-def define_train(hparams, environment_spec, event_dir):
+def define_train(hparams, event_dir):
   """Define the training setup."""
-  policy_lambda = hparams.network
-
-  if environment_spec == "stacked_pong":
-    environment_spec = lambda: gym.make("PongNoFrameskip-v4")
-    wrappers = hparams.in_graph_wrappers if hasattr(
-        hparams, "in_graph_wrappers") else []
-    wrappers.append((tf_atari_wrappers.MaxAndSkipWrapper, {"skip": 4}))
-    hparams.in_graph_wrappers = wrappers
-  if isinstance(environment_spec, str):
-    env_lambda = lambda: gym.make(environment_spec)
-  else:
-    env_lambda = environment_spec
-
-  batch_env = utils.batch_env_factory(
-      env_lambda, hparams, num_agents=hparams.num_agents)
-
-  policy_factory = functools.partial(
-      policy_lambda, batch_env.action_space, hparams)
-
+  del event_dir
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
     memory, collect_summary = collect.define_collect(
-        policy_factory, batch_env, hparams, eval_phase=False,
+        hparams, "ppo_train", eval_phase=False,
         on_simulated=hparams.simulated_environment)
-    ppo_summary = ppo.define_ppo_epoch(memory, policy_factory, hparams)
+    ppo_summary = ppo.define_ppo_epoch(memory, hparams)
     summary = tf.summary.merge([collect_summary, ppo_summary])
 
-  with tf.variable_scope("eval", reuse=tf.AUTO_REUSE):
-    eval_env_lambda = env_lambda
-    if event_dir and hparams.video_during_eval:
-      # Some environments reset environments automatically, when reached done
-      # state. For them we shall record only every second episode.
-      d = 2 if env_lambda().metadata.get("semantics.autoreset") else 1
-      eval_env_lambda = lambda: gym.wrappers.Monitor(  # pylint: disable=g-long-lambda
-          env_lambda(), event_dir, video_callable=lambda i: i % d == 0)
-      eval_env_lambda = (
-          lambda: utils.EvalVideoWrapper(eval_env_lambda()))
-    eval_batch_env = utils.batch_env_factory(
-        eval_env_lambda, hparams,
-        num_agents=hparams.num_eval_agents, xvfb=hparams.video_during_eval)
-
-    _, eval_summary = collect.define_collect(
-        policy_factory, eval_batch_env, hparams, eval_phase=True)
-  return summary, eval_summary
+  return summary, None
 
 
-def train(hparams, environment_spec, event_dir=None, model_dir=None,
+def train(hparams, event_dir=None, model_dir=None,
           restore_agent=True, epoch=0):
   """Train."""
   with tf.name_scope("rl_train"):
-    train_summary_op, eval_summary_op = define_train(hparams, environment_spec,
-                                                     event_dir)
+    train_summary_op, _ = define_train(hparams, event_dir)
     if event_dir:
       summary_writer = tf.summary.FileWriter(
           event_dir, graph=tf.get_default_graph(), flush_secs=60)
@@ -123,7 +84,6 @@ def train(hparams, environment_spec, event_dir=None, model_dir=None,
           summary_writer.add_summary(summary, epoch_index)
         if (hparams.eval_every_epochs and
             epoch_index % hparams.eval_every_epochs == 0):
-          summary = sess.run(eval_summary_op)
           if summary_writer and summary:
             summary_writer.add_summary(summary, epoch_index)
           else:
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index 9e14c4e9a..a919d88ae 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -16,7 +16,11 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+from tensor2tensor.data_generators.gym_problems import standard_atari_env_spec
+from tensor2tensor.models.research.rl import simple_gym_spec
 from tensor2tensor.rl import rl_trainer_lib
+from tensor2tensor.utils import registry  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -31,13 +35,18 @@ def test_no_crash_pendulum(self):
     hparams = trainer_lib.create_hparams(
         "ppo_continuous_action_base",
         TrainTest.test_config)
-    rl_trainer_lib.train(hparams, "Pendulum-v0")
+
+    hparams.add_hparam("environment_spec", simple_gym_spec("Pendulum-v0"))
+    rl_trainer_lib.train(hparams)
 
   def test_no_crash_cartpole(self):
     hparams = trainer_lib.create_hparams(
         "ppo_discrete_action_base",
         TrainTest.test_config)
-    rl_trainer_lib.train(hparams, "CartPole-v0")
+
+    hparams.add_hparam("environment_spec",
+                       standard_atari_env_spec("CartPole-v0"))
+    rl_trainer_lib.train(hparams)
 
 
 if __name__ == "__main__":

From c2a1a2d183f67c9e9ef661621974fb3da08cb056 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 21 Jun 2018 17:28:17 -0700
Subject: [PATCH 0181/2720] Rm unused decode_length from
 TransformerTest._create_greedy_infer_model

PiperOrigin-RevId: 201615206
---
 tensor2tensor/models/transformer_test.py | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 6325753a5..d4bb1253a 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -246,12 +246,9 @@ def testTransformerWithEncoderDecoderAttentionLoss(self):
       res = session.run(extra_loss["attention_loss"])
     self.assertEqual(res.shape, ())
 
-  def _create_greedy_infer_model(self, decode_length):
+  def _create_greedy_infer_model(self):
     """Creates model for greedy inference testing.
 
-    Args:
-      decode_length: An integer, the decode length used for test.
-
     Returns:
       model: A t2t model.
       features: An map of string to tensor.
@@ -281,7 +278,7 @@ def testGreedySlowTPUVsNonTPU(self):
 
     decode_length = 3
 
-    model, features = self._create_greedy_infer_model(decode_length)
+    model, features = self._create_greedy_infer_model()
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
       slow_result_non_tpu = model._slow_greedy_infer(
@@ -306,7 +303,7 @@ def testGreedyFastTPUVsNonTPU(self):
 
     decode_length = 3
 
-    model, features = self._create_greedy_infer_model(decode_length)
+    model, features = self._create_greedy_infer_model()
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
       fast_result_non_tpu = model._greedy_infer(
@@ -329,7 +326,7 @@ def testGreedyTPUSlowVsFast(self):
 
     decode_length = 3
 
-    model, features = self._create_greedy_infer_model(decode_length)
+    model, features = self._create_greedy_infer_model()
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
       slow_result = model._slow_greedy_infer_tpu(

From bda787ae7c07595de122fd5ef2afaf494cccdb22 Mon Sep 17 00:00:00 2001
From: Etienne Pot <epot@google.com>
Date: Thu, 21 Jun 2018 17:29:11 -0700
Subject: [PATCH 0182/2720] Fix transformer Hparams after
 causal_decoder_self_attention added to transformer_prepare_decoder

PiperOrigin-RevId: 201615304
---
 tensor2tensor/models/research/transformer_moe.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index 222fd3365..5d6f3d1c3 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -295,6 +295,7 @@ def transformer_moe_base():
   # Hparams used by transformer_prepare_decoder() function
   hparams.add_hparam("pos", "timing")  # timing, none
   hparams.add_hparam("proximity_bias", False)
+  hparams.add_hparam("causal_decoder_self_attention", True)
 
   hparams = common_attention.add_standard_attention_hparams(hparams)
 

From f8e55d1c9a06f04b3f1a13c9731c5b18ace0e56a Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 21 Jun 2018 17:44:24 -0700
Subject: [PATCH 0183/2720] Removing num_input_frames from Problem since it
 lives in Model now.

PiperOrigin-RevId: 201616980
---
 tensor2tensor/data_generators/gym_problems.py | 8 --------
 tensor2tensor/rl/envs/simulated_batch_env.py  | 5 +++--
 tensor2tensor/rl/model_rl_experiment.py       | 8 +++++++-
 3 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index e2a2b7f3a..f6558cf6a 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -68,7 +68,6 @@ class GymDiscreteProblem(video_utils.VideoProblem):
   def __init__(self, *args, **kwargs):
     super(GymDiscreteProblem, self).__init__(*args, **kwargs)
     self._env = None
-    self._env = None
     self.debug_dump_frames_path = "debug_frames_env"
     self.settable_num_steps = 5000
 
@@ -136,13 +135,6 @@ def restore_networks(self, sess):
       ckpt = ckpts.model_checkpoint_path
       model_saver.restore(sess, ckpt)
 
-  @property
-  def num_input_frames(self):
-    """Number of frames on input for real environment."""
-    # TODO(lukaszkaiser): This must be equal to hparams.video_num_input_frames,
-    # we should automate this to avoid bug in the future.
-    return 4
-
   def eval_metrics(self):
     eval_metrics = [metrics.Metrics.ACC, metrics.Metrics.ACC_PER_SEQ,
                     metrics.Metrics.IMAGE_RMSE]
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 5f1cfda6a..ed09018c9 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -91,14 +91,15 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   flags are held in according variables.
   """
 
-  def __init__(self, hparams, length, simulation_random_starts=False,
+  def __init__(self, hparams, length,
+               simulation_random_starts=False,
                intrinsic_reward_scale=0.):
     """Batch of environments inside the TensorFlow graph."""
     self.length = length
     environment_spec = hparams.environment_spec
     initial_frames_problem = environment_spec.initial_frames_problem
     self._min_reward = initial_frames_problem.min_reward
-    self._num_frames = initial_frames_problem.num_input_frames
+    self._num_frames = hparams.model_hparams.video_num_input_frames
     self._intrinsic_reward_scale = intrinsic_reward_scale
 
     # initialization_env = environment_lambda()
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 4b412d364..bba9e8da6 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -148,6 +148,10 @@ def train_agent(problem_name, agent_model_dir,
   if hparams.ppo_learning_rate:
     ppo_hparams.learning_rate = hparams.ppo_learning_rate
 
+  # Adding model hparams for model specific adjustments
+  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
+  ppo_hparams.add_hparam("model_hparams", model_hparams)
+
   environment_spec = copy.copy(gym_problem.environment_spec)
   environment_spec.simulated_env = True
   environment_spec.add_hparam("simulation_random_starts",
@@ -158,7 +162,9 @@ def train_agent(problem_name, agent_model_dir,
                               gym_problem)
 
   # 4x for the StackAndSkipWrapper minus one to always finish for reporting.
-  ppo_time_limit = (ppo_hparams.epoch_length - 1) * 4
+  ppo_time_limit = ppo_hparams.epoch_length - 1
+  ppo_time_limit *= model_hparams.video_num_input_frames
+
   wrappers = environment_spec.wrappers + \
              [[TimeLimitWrapper, {"timelimit": ppo_time_limit}]]
   environment_spec.wrappers = wrappers

From 57f0da164b14a6d17014c8d28447fc5cd6b46f91 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 21 Jun 2018 19:16:21 -0700
Subject: [PATCH 0184/2720] Fixing the simulated env num frames bug.

PiperOrigin-RevId: 201625851
---
 tensor2tensor/rl/envs/simulated_batch_env.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index ed09018c9..78a14d553 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -111,6 +111,11 @@ def __init__(self, hparams, length,
 
     _, self.action_shape, self.action_dtype = get_action_space(environment_spec)
 
+    # TODO(lukaszkaiser): do this in a more cleaner way
+    hparams.video_num_input_frames, hparams.video_num_target_frames = (
+        hparams.model_hparams.video_num_input_frames,
+        hparams.model_hparams.video_num_target_frames)
+
     if simulation_random_starts:
       dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
                                                FLAGS.data_dir,

From cc892360c2c9404631d7a99e474cc3ee05e6e699 Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Thu, 21 Jun 2018 22:01:12 -0700
Subject: [PATCH 0185/2720] Adding word dropout for robustness in latents and
 keeping default at 0.0

PiperOrigin-RevId: 201636872
---
 .../models/research/transformer_nat.py        | 22 +------------
 .../models/research/transformer_vae.py        | 32 +------------------
 2 files changed, 2 insertions(+), 52 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 7d4ed16fb..bf11926e1 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -255,14 +255,7 @@ def ae_transformer_internal(inputs, targets, target_space, hparams, cache=None):
       targets,
       max_targets_len_from_inputs,
       final_length_divisible_by=2**hparams.num_compress_steps)
-  if hparams.word_dropout:
-    mask = tf.random_uniform(shape=common_layers.shape_list(targets),
-                             minval=0.0, maxval=1.0)
-    targets_noisy = tf.where(mask > hparams.word_dropout, targets,
-                             tf.zeros_like(targets))
-  else:
-    targets_noisy = targets
-  targets_c = compress(targets_noisy, hparams, "compress")
+  targets_c = compress(targets, hparams, "compress")
   if hparams.mode != tf.estimator.ModeKeys.PREDICT:
     # Compress and bottleneck.
     latents_discrete_hot, extra_loss = vq_discrete_bottleneck(
@@ -414,7 +407,6 @@ def transformer_nat_small():
   hparams.add_hparam("decay", 0.999)
   hparams.add_hparam("num_samples", 10)
   hparams.add_hparam("mask_startup_steps", 50000)
-  hparams.add_hparam("word_dropout", 0.0)
   return hparams
 
 
@@ -429,18 +421,6 @@ def transformer_nat_base():
   return hparams
 
 
-@registry.register_hparams
-def transformer_nat_base_drop():
-  """Set of hyperparameters."""
-  hparams = transformer_nat_small()
-  hparams.batch_size = 2048
-  hparams.hidden_size = 512
-  hparams.filter_size = 4096
-  hparams.num_hidden_layers = 6
-  hparams.word_dropout = 0.2
-  return hparams
-
-
 @registry.register_hparams
 def transformer_nat_big():
   """Set of hyperparameters."""
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 7ed8f1f76..5300968cf 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -675,7 +675,7 @@ def transformer_ae_small():
   hparams.add_hparam("noise_dev", 0.5)
   hparams.add_hparam("d_mix", 0.5)
   hparams.add_hparam("logit_normalization", True)
-  hparams.add_hparam("word_dropout", 0.0)
+  hparams.add_hparam("word_dropout", 0.1)
   # Bottleneck kinds supported: dense, vae, semhash, gumbel-softmax, dvq.
   hparams.add_hparam("bottleneck_kind", "semhash")
   hparams.add_hparam("num_blocks", 1)
@@ -803,36 +803,6 @@ def transformer_ae_base():
   return hparams
 
 
-@registry.register_hparams
-def transformer_ae_base_noatt():
-  """Set of hyperparameters."""
-  hparams = transformer_ae_small()
-  hparams.batch_size = 2048
-  hparams.hidden_size = 512
-  hparams.filter_size = 4096
-  hparams.num_hidden_layers = 6
-  hparams.do_attend_decompress = False
-  hparams.do_attend_decompress = False
-  hparams.word_dropout = 0.0
-  hparams.z_size = 12
-  return hparams
-
-
-@registry.register_hparams
-def transformer_ae_base_drop_noatt():
-  """Set of hyperparameters."""
-  hparams = transformer_ae_small()
-  hparams.batch_size = 2048
-  hparams.hidden_size = 512
-  hparams.filter_size = 4096
-  hparams.num_hidden_layers = 6
-  hparams.do_attend_decompress = False
-  hparams.do_attend_decompress = False
-  hparams.word_dropout = 0.2
-  hparams.z_size = 12
-  return hparams
-
-
 @registry.register_hparams
 def transformer_ae_a3():
   """Set of hyperparameters."""

From 8a330d52d1b5c0530fa26e898cfb2a03f3ab859c Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Fri, 22 Jun 2018 13:43:41 +0200
Subject: [PATCH 0186/2720] crtical bug fix (data in memory not properly
 aligned)

---
 tensor2tensor/rl/collect.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 90a53721d..ba2a67e8f 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -58,15 +58,22 @@ def __init__(self, batch_env):
                                trainable=False)
 
   def simulate(self, action):
+
+    # There is subtlety here. We need to collect data
+    # obs, action = policy(obs), done, reward = env(abs, action)
+    # Thus we need to enqueue data before assigning new observation
+
     reward, done = self._batch_env.simulate(action)
+
     with tf.control_dependencies([reward, done]):
+      enqueue_op = self.speculum.enqueue(
+          [self._observ.read_value(), reward, done, action])
+
+    with tf.control_dependencies([enqueue_op]):
       assign = self._observ.assign(self._batch_env.observ)
 
     with tf.control_dependencies([assign]):
-      enqueue_op = self.speculum.enqueue(
-          [self._observ.read_value(), reward, done, action])
-      with tf.control_dependencies([enqueue_op]):
-        return tf.identity(reward), tf.identity(done)
+      return tf.identity(reward), tf.identity(done)
 
 
 def define_collect(hparams, scope, eval_phase,

From 022d85a37a4bada4db9423374a6b1120b9762d19 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 22 Jun 2018 11:18:21 -0700
Subject: [PATCH 0187/2720] Adding Stochastic RL test.

PiperOrigin-RevId: 201713329
---
 tensor2tensor/models/research/next_frame.py  | 37 +++++++++++++++-----
 tensor2tensor/rl/model_rl_experiment.py      |  4 +--
 tensor2tensor/rl/model_rl_experiment_test.py |  9 ++++-
 3 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 0bbbd891c..eab2becfe 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -18,7 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import numpy as np
 import six
 
 from tensor2tensor.layers import common_attention
@@ -170,6 +169,11 @@ class NextFrameStochastic(NextFrameBasic):
   https://arxiv.org/abs/1710.11252
   """
 
+  def tinyify(self, array):
+    if self.hparams.tiny_mode:
+      return [1 for _ in array]
+    return array
+
   def construct_latent_tower(self, images):
     """Builds convolutional latent tower for stochastic model.
 
@@ -189,6 +193,7 @@ def construct_latent_tower(self, images):
       latent_loss: loss of the latent twoer
       samples: random samples sampled from standard guassian
     """
+    conv_size = self.tinyify([32, 64, 64])
     with tf.variable_scope("latent", reuse=tf.AUTO_REUSE):
       # this allows more predicted frames at inference time
       latent_images = images[:self.hparams.latent_num_frames]
@@ -196,12 +201,12 @@ def construct_latent_tower(self, images):
 
       x = images
       x = common_layers.make_even_size(x)
-      x = slim.conv2d(x, 32, [3, 3], stride=2, scope="latent_conv1")
+      x = slim.conv2d(x, conv_size[0], [3, 3], stride=2, scope="latent_conv1")
       x = slim.batch_norm(x, scope="latent_bn1")
       x = common_layers.make_even_size(x)
-      x = slim.conv2d(x, 64, [3, 3], stride=2, scope="latent_conv2")
+      x = slim.conv2d(x, conv_size[1], [3, 3], stride=2, scope="latent_conv2")
       x = slim.batch_norm(x, scope="latent_bn2")
-      x = slim.conv2d(x, 64, [3, 3], stride=1, scope="latent_conv3")
+      x = slim.conv2d(x, conv_size[2], [3, 3], stride=1, scope="latent_conv3")
       x = slim.batch_norm(x, scope="latent_bn3")
 
       nc = self.hparams.latent_channels
@@ -218,11 +223,12 @@ def construct_latent_tower(self, images):
 
   def reward_prediction(self, inputs):
     """Builds a reward prediction network."""
+    conv_size = self.tinyify([32, 16])
     with tf.variable_scope("reward_pred", reuse=tf.AUTO_REUSE):
       x = inputs
-      x = slim.conv2d(x, 32, [3, 3], scope="reward_conv1")
+      x = slim.conv2d(x, conv_size[0], [3, 3], scope="reward_conv1")
       x = slim.batch_norm(x, scope="reward_bn1")
-      x = slim.conv2d(x, 16, [3, 3], scope="reward_conv2")
+      x = slim.conv2d(x, conv_size[1], [3, 3], scope="reward_conv2")
       x = slim.batch_norm(x, scope="reward_bn2")
       x = slim.conv2d(x, 1, [3, 3], scope="reward_conv3", activation_fn=None)
     return x
@@ -262,19 +268,21 @@ def construct_predictive_tower(
     # Main tower
     layer_norm = tf.contrib.layers.layer_norm
     lstm_func = self.conv_lstm_2d
-    lstm_size = np.array([32, 32, 64, 64, 128, 64, 32], dtype=np.int32)
     batch_size = common_layers.shape_list(input_image)[0]
     # the number of different pixel motion predictions
     # and the number of masks for each of those predictions
     num_masks = self.hparams.num_masks
 
+    lstm_size = self.tinyify([32, 32, 64, 64, 128, 64, 32])
+    conv_size = self.tinyify([32])
+
     img_height, img_width, color_channels = self.hparams.problem.frame_shape
 
     with tf.variable_scope("main", reuse=tf.AUTO_REUSE):
       input_image = common_layers.make_even_size(input_image)
       enc0 = slim.layers.conv2d(
           input_image,
-          32, [5, 5],
+          conv_size[0], [5, 5],
           stride=2,
           scope="scale1_conv1",
           normalizer_fn=layer_norm,
@@ -1066,6 +1074,7 @@ def next_frame_stochastic():
   hparams.add_hparam(
       "latent_num_frames",  # use all frames by default.
       hparams.video_num_input_frames + hparams.video_num_target_frames)
+  hparams.add_hparam("tiny_mode", False)
   return hparams
 
 
@@ -1093,6 +1102,18 @@ def next_frame_stochastic_cutoff():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_stochastic_tiny():
+  """SV2P model with additional cutoff in L2 loss for environments like pong."""
+  hparams = next_frame_stochastic()
+  hparams.tiny_mode = True
+  hparams.num_masks = 1
+  hparams.video_modality_loss_cutoff = 0.4
+  hparams.video_num_input_frames = 4
+  hparams.video_num_target_frames = 1
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_tpu():
   hparams = next_frame()
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index bba9e8da6..5c6c9f16c 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -553,9 +553,7 @@ def rl_modelrl_tiny_stochastic():
   """Tiny setting with a stochastic next-frame model."""
   hparams = rl_modelrl_tiny()
   hparams.generative_model = "next_frame_stochastic"
-  hparams.generative_model_params = "next_frame_stochastic_cutoff"
-  hparams.true_env_generator_num_steps = 1000
-  hparams.simulated_env_generator_num_steps = 1000
+  hparams.generative_model_params = "next_frame_stochastic_tiny"
   return hparams
 
 
diff --git a/tensor2tensor/rl/model_rl_experiment_test.py b/tensor2tensor/rl/model_rl_experiment_test.py
index 2fd327260..e2430983c 100644
--- a/tensor2tensor/rl/model_rl_experiment_test.py
+++ b/tensor2tensor/rl/model_rl_experiment_test.py
@@ -26,13 +26,20 @@
 
 class ModelRLExperimentTest(tf.test.TestCase):
 
-  def test_run(self):
+  def test_basic(self):
     FLAGS.output_dir = tf.test.get_temp_dir()
     FLAGS.loop_hparams_set = "rl_modelrl_tiny"
     FLAGS.loop_hparams = "generative_model_params=next_frame_tiny"
     FLAGS.schedule = "train"  # skip evaluation for world model training
     model_rl_experiment.main(None)
 
+  def test_stochastic(self):
+    FLAGS.output_dir = tf.test.get_temp_dir()
+    FLAGS.loop_hparams_set = "rl_modelrl_tiny_stochastic"
+    FLAGS.loop_hparams = "generative_model_params=next_frame_stochastic_tiny"
+    FLAGS.schedule = "train"  # skip evaluation for world model training
+    model_rl_experiment.main(None)
+
 
 if __name__ == "__main__":
   tf.test.main()

From 55434d4933e0b2eea13cb85ae8249ffeebe81a4d Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Fri, 22 Jun 2018 12:05:21 -0700
Subject: [PATCH 0188/2720] Don't use T2TExperiment with autotuners

PiperOrigin-RevId: 201720995
---
 tensor2tensor/bin/t2t_trainer.py   |  5 +-
 tensor2tensor/utils/trainer_lib.py | 77 ++++++++++++++++++++----------
 2 files changed, 55 insertions(+), 27 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 426a32ed3..8e3737d62 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -154,7 +154,7 @@ def create_hparams():
   return trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
 
 
-def create_experiment_fn():
+def create_experiment_fn(**kwargs):
   return trainer_lib.create_experiment_fn(
       model_name=FLAGS.model,
       problem_name=FLAGS.problem,
@@ -173,7 +173,8 @@ def create_experiment_fn():
       eval_early_stopping_metric_delta=FLAGS.eval_early_stopping_metric_delta,
       eval_early_stopping_metric_minimize=FLAGS.
       eval_early_stopping_metric_minimize,
-      use_tpu=FLAGS.use_tpu)
+      use_tpu=FLAGS.use_tpu,
+      **kwargs)
 
 
 def create_run_config(hp):
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 25f210518..3fd0ccc9f 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -39,10 +39,7 @@ def next_checkpoint(model_dir, timeout_mins=120):
   last_ckpt = None
   while True:
     last_ckpt = tf.contrib.training.wait_for_new_checkpoint(
-        model_dir,
-        last_ckpt,
-        seconds_to_sleep=60,
-        timeout=60 * timeout_mins)
+        model_dir, last_ckpt, seconds_to_sleep=60, timeout=60 * timeout_mins)
 
     if last_ckpt is None:
       tf.logging.info(
@@ -265,7 +262,11 @@ def create_hooks(use_tfdbg=False,
 class T2TExperiment(object):
   """Custom Experiment class for running distributed experiments."""
 
-  def __init__(self, estimator, hparams, train_spec, eval_spec,
+  def __init__(self,
+               estimator,
+               hparams,
+               train_spec,
+               eval_spec,
                decode_hparams=None):
     self._train_spec = train_spec
     self._eval_spec = eval_spec
@@ -273,6 +274,18 @@ def __init__(self, estimator, hparams, train_spec, eval_spec,
     self._decode_hparams = decode_hparams
     self._estimator = estimator
 
+  @property
+  def estimator(self):
+    return self._estimator
+
+  @property
+  def eval_steps(self):
+    return self._eval_spec.steps
+
+  @property
+  def train_steps(self):
+    return self._train_spec.max_steps
+
   def continuous_train_and_eval(self):
     tf.estimator.train_and_evaluate(self._estimator, self._train_spec,
                                     self._eval_spec)
@@ -361,25 +374,27 @@ def continuous_decode(self):
       self.decode()
 
 
-def create_experiment(run_config,
-                      hparams,
-                      model_name,
-                      problem_name,
-                      data_dir,
-                      train_steps,
-                      eval_steps,
-                      min_eval_frequency=2000,
-                      eval_throttle_seconds=600,
-                      schedule="train_and_evaluate",
-                      export=False,
-                      decode_hparams=None,
-                      use_tfdbg=False,
-                      use_dbgprofile=False,
-                      eval_early_stopping_steps=None,
-                      eval_early_stopping_metric=None,
-                      eval_early_stopping_metric_delta=None,
-                      eval_early_stopping_metric_minimize=True,
-                      use_tpu=False):
+def create_experiment(
+    run_config,
+    hparams,
+    model_name,
+    problem_name,
+    data_dir,
+    train_steps,
+    eval_steps,
+    min_eval_frequency=2000,
+    eval_throttle_seconds=600,
+    schedule="train_and_evaluate",
+    export=False,
+    decode_hparams=None,
+    use_tfdbg=False,
+    use_dbgprofile=False,
+    eval_early_stopping_steps=None,
+    eval_early_stopping_metric=None,
+    eval_early_stopping_metric_delta=None,
+    eval_early_stopping_metric_minimize=True,
+    autotune=False,
+    use_tpu=False):
   """Create Experiment."""
   # HParams
   hparams.add_hparam("model_dir", run_config.model_dir)
@@ -464,7 +479,19 @@ def create_experiment(run_config,
         hooks=eval_hooks,
         start_delay_secs=0 if hparams.schedule == "evaluate" else 120,
         throttle_secs=eval_throttle_seconds)
-
+  hooks_kwargs = {"train_hooks": train_hooks, "eval_hooks": eval_hooks}
+
+  if autotune:
+    return tf.contrib.learn.Experiment(
+        estimator=estimator,
+        train_input_fn=train_input_fn,
+        eval_input_fn=eval_input_fn,
+        train_steps=train_steps,
+        eval_steps=eval_steps,
+        min_eval_frequency=min_eval_frequency,
+        train_steps_per_iteration=min(min_eval_frequency, train_steps),
+        eval_delay_secs=0 if schedule == "evaluate" else 120,
+        **hooks_kwargs if not use_tpu else {})
   return T2TExperiment(estimator, hparams, train_spec, eval_spec,
                        decode_hparams)
 

From 268b92455aa4ea78b89cc4b63105e02c266e5751 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 22 Jun 2018 12:58:55 -0700
Subject: [PATCH 0189/2720] update decoding logging subdirectory

PiperOrigin-RevId: 201728454
---
 tensor2tensor/utils/decoding.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 4b6552dea..7e5d30a08 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -67,7 +67,7 @@ def log_decode_results(inputs,
                        targets_vocab,
                        targets=None,
                        save_images=False,
-                       model_dir=None,
+                       output_dir=None,
                        identity_output=False,
                        log_results=True):
   """Log inference results."""
@@ -77,7 +77,8 @@ def log_decode_results(inputs,
   if is_video:
     def fix_and_save_video(vid, prefix):
       save_path_template = os.path.join(
-          model_dir, "%s_%s_%d_{}.png" % (problem_name, prefix, prediction_idx))
+          output_dir,
+          "%s_%s_%d_{}.png" % (problem_name, prefix, prediction_idx))
       # this is only required for predictions
       if vid.shape[-1] == 1:
         vid = np.squeeze(vid, axis=-1)
@@ -91,7 +92,7 @@ def fix_and_save_video(vid, prefix):
   decoded_inputs = None
   if is_image and save_images:
     save_path = os.path.join(
-        model_dir, "%s_prediction_%d.jpg" % (problem_name, prediction_idx))
+        output_dir, "%s_prediction_%d.jpg" % (problem_name, prediction_idx))
     show_and_save_image(inputs / 255., save_path)
   elif inputs_vocab:
     if identity_output:
@@ -131,6 +132,10 @@ def decode_from_dataset(estimator,
   # We assume that worker_id corresponds to shard number.
   shard = decode_hp.shard_id if decode_hp.shards > 1 else None
 
+  # Setup the decode output directory for any artifacts that may be written out
+  output_dir = os.path.join(estimator.model_dir, "decode")
+  tf.gfile.MakeDirs(output_dir)
+
   # If decode_hp.batch_size is specified, use a fixed batch size
   if decode_hp.batch_size:
     hparams.batch_size = decode_hp.batch_size
@@ -200,7 +205,7 @@ def decode_from_dataset(estimator,
             inputs_vocab,
             targets_vocab,
             save_images=decode_hp.save_images,
-            model_dir=estimator.model_dir,
+            output_dir=output_dir,
             identity_output=decode_hp.identity_output,
             targets=targets,
             log_results=decode_hp.log_results)
@@ -216,7 +221,7 @@ def decode_from_dataset(estimator,
           inputs_vocab,
           targets_vocab,
           save_images=decode_hp.save_images,
-          model_dir=estimator.model_dir,
+          output_dir=output_dir,
           identity_output=decode_hp.identity_output,
           targets=targets,
           log_results=decode_hp.log_results)

From c3ea3784334200758fddf047aa2abbf7f06edb1f Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Fri, 22 Jun 2018 12:59:35 -0700
Subject: [PATCH 0190/2720] More fixes in latent layers

PiperOrigin-RevId: 201728524
---
 tensor2tensor/layers/latent_layers.py | 36 ++++++++++-----------------
 1 file changed, 13 insertions(+), 23 deletions(-)

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 10a7fc0af..86b9dfb7d 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -82,30 +82,20 @@ def multinomial_sample(x, vocab_size, temperature):
   return reshaped_samples
 
 
-def ae_latent_softmax(latents_pred, latents_discrete, hparams):
+def ae_latent_softmax(latents_pred, latents_discrete_hot, hparams):
   """Latent prediction and loss."""
   vocab_size = 2**hparams.bottleneck_bits
-  if hparams.num_decode_blocks < 2:
-    with tf.variable_scope("extra_logits"):
-      latents_logits = tf.layers.dense(latents_pred, vocab_size,
-                                       name="extra_logits")
-      if hparams.logit_normalization:
-        latents_logits *= tf.rsqrt(1e-8 +
-                                   tf.reduce_mean(tf.square(latents_logits)))
-
-      loss = None
-      if latents_discrete is not None:
-        if hparams.soft_em:
-          # latents_discrete is actually one-hot of multinomial samples
-          assert hparams.num_decode_blocks == 1
-          loss = tf.nn.softmax_cross_entropy_with_logits_v2(
-              labels=latents_discrete, logits=latents_logits)
-        else:
-          loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
-              labels=latents_discrete, logits=latents_logits)
-      sample = multinomial_sample(latents_logits, vocab_size,
-                                  hparams.sampling_temp)
-      return sample, loss
+  with tf.variable_scope("latent_logits"):
+    latents_logits = tf.layers.dense(latents_pred, vocab_size,
+                                     name="logits_dense")
+    if hparams.logit_normalization:
+      latents_logits *= tf.rsqrt(1e-8 +
+                                 tf.reduce_mean(tf.square(latents_logits)))
+    loss = tf.nn.softmax_cross_entropy_with_logits_v2(
+        labels=latents_discrete_hot, logits=latents_logits)
+    sample = multinomial_sample(latents_logits, vocab_size,
+                                hparams.sampling_temp)
+    return sample, loss
 
 
 def ae_latent_sample_beam(latents_dense_in, inputs, ed, embed, hparams):
@@ -427,7 +417,7 @@ def bottleneck_layer(targets_c, hparams):
 
   if DO_SUMMARIES:
     tf.summary.histogram("b0", tf.reshape(latents_discrete, [-1]))
-  return latents_dense, latents_discrete, extra_loss
+  return latents_dense, latents_discrete_hot, extra_loss
 
 
 def latent_prediction_model(inputs,

From a4fa55a3f128753d006d26ba8691eb97d14fbcfc Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Fri, 22 Jun 2018 14:29:09 -0700
Subject: [PATCH 0191/2720] add 3d vars hparams for transformer, do not create
 dummy vars during training/inference.

PiperOrigin-RevId: 201741899
---
 tensor2tensor/layers/common_attention.py | 21 +++++++++++++++------
 tensor2tensor/layers/common_hparams.py   |  2 ++
 tensor2tensor/models/transformer.py      |  6 ++++--
 tensor2tensor/utils/modality.py          |  2 ++
 tensor2tensor/utils/t2t_model.py         |  3 ++-
 5 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 4881dd8f5..66e77fbe7 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2801,16 +2801,25 @@ def compute_qkv(query_antecedent,
   if memory_antecedent is None:
     memory_antecedent = query_antecedent
   q = compute_attention_component(
-      query_antecedent, total_key_depth,
-      q_filter_width, q_padding, "q",
+      query_antecedent,
+      total_key_depth,
+      q_filter_width,
+      q_padding,
+      "q",
       vars_3d_num_heads=vars_3d_num_heads)
   k = compute_attention_component(
-      memory_antecedent, total_key_depth,
-      kv_filter_width, kv_padding, "k",
+      memory_antecedent,
+      total_key_depth,
+      kv_filter_width,
+      kv_padding,
+      "k",
       vars_3d_num_heads=vars_3d_num_heads)
   v = compute_attention_component(
-      memory_antecedent, total_value_depth,
-      kv_filter_width, kv_padding, "v",
+      memory_antecedent,
+      total_value_depth,
+      kv_filter_width,
+      kv_padding,
+      "v",
       vars_3d_num_heads=vars_3d_num_heads)
   return q, k, v
 
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index bc12b2629..4b6d62934 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -126,6 +126,8 @@ def basic_params1():
       # epsilon parameter to normalization function
       norm_epsilon=1e-6,
       symbol_modality_num_shards=1,
+      # pad vocabularies so that this value divides the vocabulary size.
+      vocab_divisor=1,
       # During training, we drop sequences whose inputs and targets are shorter
       # than min_length
       min_length=0,
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 203e8709a..d616dbfbd 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1301,7 +1301,8 @@ def transformer_decoder(decoder_input,
               make_image_summary=make_image_summary,
               dropout_broadcast_dims=attention_dropout_broadcast_dims,
               max_length=hparams.get("max_length"),
-              decode_loop_step=decode_loop_step)
+              decode_loop_step=decode_loop_step,
+              vars_3d=hparams.get("attention_variables_3d"))
           x = common_layers.layer_postprocess(x, y, hparams)
         if encoder_output is not None:
           with tf.variable_scope("encdec_attention"):
@@ -1318,7 +1319,8 @@ def transformer_decoder(decoder_input,
                 cache=layer_cache,
                 make_image_summary=make_image_summary,
                 dropout_broadcast_dims=attention_dropout_broadcast_dims,
-                max_length=hparams.get("max_length"))
+                max_length=hparams.get("max_length"),
+                vars_3d=hparams.get("attention_variables_3d"))
             x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index d38ee737b..815d6d5ed 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -52,6 +52,8 @@ class Modality(object):
 
   def __init__(self, model_hparams, vocab_size=None):
     self._model_hparams = model_hparams
+    if vocab_size is not None and hasattr(model_hparams, "vocab_divisor"):
+      vocab_size += (0 - vocab_size) % model_hparams.vocab_divisor
     self._vocab_size = vocab_size
 
   @property
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 09ae083e8..85b705eaf 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1181,7 +1181,8 @@ def estimator_model_fn(cls,
     Returns:
       TPUEstimatorSpec if use tpu else EstimatorSpec
     """
-    _create_dummy_vars()
+    if mode == tf.estimator.ModeKeys.TRAIN:
+      _create_dummy_vars()
     hparams = copy.deepcopy(hparams)
 
     # Instantiate model

From 3ea46fe7cb12ec46b95c10d10270f20c4d73a7a2 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 22 Jun 2018 14:30:22 -0700
Subject: [PATCH 0192/2720] Add postdecode hooks running

PiperOrigin-RevId: 201742130
---
 tensor2tensor/data_generators/problem.py | 11 ++++++
 tensor2tensor/utils/decoding.py          | 46 ++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 3186ceea4..d6cd7811e 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -605,6 +605,17 @@ def decode_example(self, serialized_example):
     decoded = decoder.decode(serialized_example, items=decode_items)
     return dict(zip(decode_items, decoded))
 
+  @property
+  def decode_hooks(self):
+    """List of functions to be run after full decodes have been produced.
+
+    Returns:
+      List of functions. Each function should expect a single argument, an
+      instance of decoding.DecodeHookArgs and optionally return a list of
+      tf.Summary.Value objects.
+    """
+    return []
+
   @property
   def has_inputs(self):
     return "inputs" in self.get_feature_encoders()
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 7e5d30a08..13c904bbc 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -17,6 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import operator
 import os
 import time
@@ -246,6 +247,13 @@ def decode_from_dataset(estimator,
     target_file.close()
     input_file.close()
 
+  run_postdecode_hooks(DecodeHookArgs(
+      estimator=estimator,
+      problem=problem,
+      output_dir=output_dir,
+      hparams=hparams,
+      decode_hparams=decode_hp))
+
   tf.logging.info("Completed inference on %d samples." % num_predictions)  # pylint: disable=undefined-loop-variable
 
 
@@ -694,3 +702,41 @@ def _decode_input_tensor_to_features_dict(feature_map, hparams):
       IMAGE_DECODE_LENGTH if input_is_image else tf.shape(x)[1] + 50)
   features["inputs"] = x
   return features
+
+
+def latest_checkpoint_step(ckpt_dir):
+  ckpt = tf.train.get_checkpoint_state(ckpt_dir)
+  if not ckpt:
+    return None
+  path = ckpt.model_checkpoint_path
+  step = int(path.split("-")[-1])
+  return step
+
+
+class DecodeHookArgs(collections.namedtuple(
+    "DecodeHookArgs",
+    ["estimator", "problem", "output_dir", "hparams", "decode_hparams"])):
+  pass
+
+
+def run_postdecode_hooks(decode_hook_args):
+  """Run hooks after decodes have run."""
+  hooks = decode_hook_args.problem.decode_hooks
+  if not hooks:
+    return
+  global_step = latest_checkpoint_step(decode_hook_args.estimator.model_dir)
+  if global_step is None:
+    tf.logging.info(
+        "Skipping decode hooks because no checkpoint yet available.")
+    return
+  tf.logging.info("Running decode hooks.")
+  summary_writer = tf.summary.FileWriter(decode_hook_args.output_dir)
+  for hook in hooks:
+    # Isolate each hook in case it creates TF ops
+    with tf.Graph():
+      summaries = hook(decode_hook_args)
+    if summaries:
+      summary = tf.Summary(value=list(summaries))
+      summary_writer.add_summary(summary, global_step)
+  summary_writer.close()
+  tf.logging.info("Decode hooks done.")

From 546b189dfd071119d8964717aa9b519172a3eca2 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 22 Jun 2018 15:05:44 -0700
Subject: [PATCH 0193/2720] Fix decode hooks runner

PiperOrigin-RevId: 201747870
---
 tensor2tensor/utils/decoding.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 13c904bbc..baeb6f675 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -733,7 +733,7 @@ def run_postdecode_hooks(decode_hook_args):
   summary_writer = tf.summary.FileWriter(decode_hook_args.output_dir)
   for hook in hooks:
     # Isolate each hook in case it creates TF ops
-    with tf.Graph():
+    with tf.Graph().as_default():
       summaries = hook(decode_hook_args)
     if summaries:
       summary = tf.Summary(value=list(summaries))

From cc92901fe8b634374df9a2f595d963cf1a44ef02 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 22 Jun 2018 15:35:52 -0700
Subject: [PATCH 0194/2720] Add video metrics hook to video problem.

PiperOrigin-RevId: 201752123
---
 tensor2tensor/data_generators/video_utils.py | 28 ++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index b0bdb77d5..e56e75ce3 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -26,6 +26,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
+from tensor2tensor.utils import video_metrics
 
 import tensorflow as tf
 
@@ -39,6 +40,29 @@ def resize_video_frames(images, size):
   return resized_images
 
 
+def summarize_video_metrics(hook_args):
+  """Computes video metrics summaries using the decoder output."""
+  problem_name = hook_args.problem.name
+  current_problem = hook_args.problem
+  hparams = hook_args.hparams
+  output_dir = hook_args.output_dir
+  frame_shape = [
+      current_problem.frame_height, current_problem.frame_width,
+      current_problem.num_channels
+  ]
+  metrics_graph = tf.Graph()
+  with metrics_graph.as_default():
+    metrics_results = video_metrics.compute_video_metrics(
+        output_dir, problem_name, hparams.video_num_target_frames, frame_shape)
+
+  summary_values = []
+  for name, array in six.iteritems(metrics_results):
+    for ind, val in enumerate(array):
+      tag = name + "_" + str(ind)
+      summary_values.append(tf.Summary.Value(tag=tag, simple_value=val))
+  return summary_values
+
+
 class VideoProblem(problem.Problem):
   """Base class for problems with videos."""
 
@@ -112,6 +136,10 @@ def preprocess_example(self, example, mode, hparams):
     """Runtime preprocessing, e.g., resize example["frame"]."""
     return example
 
+  @property
+  def decode_hooks(self):
+    return [summarize_video_metrics]
+
   @property
   def is_generate_per_split(self):
     """A single call to `generate_samples` generates for all `dataset_splits`.

From a754182abe4976fdee38ddb6553f6038d4867fa6 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Fri, 22 Jun 2018 15:44:13 -0700
Subject: [PATCH 0195/2720] add latent dropout, clean up hparams and unused
 code, add function descriptions

PiperOrigin-RevId: 201753189
---
 tensor2tensor/layers/latent_layers.py | 133 +++++++++++++++-----------
 1 file changed, 78 insertions(+), 55 deletions(-)

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 86b9dfb7d..900734473 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -30,18 +30,6 @@
 DO_SUMMARIES = True
 
 
-class Latent(object):
-  DISCRETE = "discrete"
-  DENSE = "dense"
-
-  @staticmethod
-  def get_choices():
-    return [
-        Latent.DISCRETE,
-        Latent.DENSE,
-    ]
-
-
 def add_learned_positional_embeddings(x, hparams):
   pos = tf.get_variable("pos",
                         [1, hparams.img_len*hparams.img_len,
@@ -72,9 +60,9 @@ def attend(x, source, hparams, name):
     return tf.reshape(res, xshape)
 
 
-def multinomial_sample(x, vocab_size, temperature):
+def multinomial_sample(x, vocab_size, sampling_method, temperature):
   """Multinomial sampling from a n-dimensional tensor."""
-  if temperature > 0:
+  if sampling_method == "random":
     samples = tf.multinomial(tf.reshape(x, [-1, vocab_size]) / temperature, 1)
   else:
     samples = tf.argmax(x, axis=-1)
@@ -94,6 +82,7 @@ def ae_latent_softmax(latents_pred, latents_discrete_hot, hparams):
     loss = tf.nn.softmax_cross_entropy_with_logits_v2(
         labels=latents_discrete_hot, logits=latents_logits)
     sample = multinomial_sample(latents_logits, vocab_size,
+                                hparams.sampling_method,
                                 hparams.sampling_temp)
     return sample, loss
 
@@ -110,10 +99,10 @@ def symbols_to_logits_fn(ids):
       latents_dense = embed(
           tf.one_hot(latents_discrete, depth=2**hparams.bottleneck_bits),
           hparams.hidden_size)
-      latents_pred = transformer_latent_decoder(inputs, ed, latents_dense,
-                                                hparams, "extra")
+      latents_pred = transformer_latent_decoder(
+          latents_dense, inputs, ed, hparams, name="latent_prediction")
       logits = tf.layers.dense(
-          latents_pred, 2**hparams.bottleneck_bits, name="extra_logits")
+          latents_pred, 2**hparams.bottleneck_bits, name="logits_dense")
       current_output_position = common_layers.shape_list(ids)[1] - 1
       logits = logits[:, current_output_position, :]
     return logits
@@ -290,7 +279,7 @@ def decompress_decoder(inputs, hparams,
 
 
 def decompress_decoder_2d(x, hparams, name):
-  """Dencoder that decompresses x to length height*width.
+  """Decoder that decompresses x to length height*width.
 
   Args:
     x: Tensor of shape [batch, compress_height, compress_width, hidden_dim].
@@ -325,11 +314,22 @@ def decompress_decoder_1d(x, hparams, name):
   return tf.squeeze(output, axis=2)
 
 
-def transformer_text_encoder(inputs, space_id,
-                             hparams, name="transformer_text_enc"):
-  """Transformer text encoder."""
+def transformer_text_encoder(
+    x, space_id, hparams, name="transformer_text_encoder"):
+  """Transformer text encoder over inputs with unmasked full attention.
+
+  Args:
+    x: Tensor of shape [batch, length, hidden_dim].
+    space_id: int, id.
+    hparams: Dict, hyperparameters.
+    name: string, variable scope.
+
+  Returns:
+    x: Tensor of shape [batch, length, hidden_dim].
+    ed: Tensor, bias for padded tokens in the input, shape [batch, length]
+  """
   with tf.variable_scope(name):
-    x = common_layers.flatten4d3d(inputs)
+    x = common_layers.flatten4d3d(x)
     (encoder_input, encoder_self_attention_bias,
      ed) = transformer.transformer_prepare_encoder(x, space_id, hparams)
     encoder_input = tf.nn.dropout(encoder_input, 1.0 - hparams.dropout)
@@ -337,16 +337,27 @@ def transformer_text_encoder(inputs, space_id,
         encoder_input, encoder_self_attention_bias, hparams), ed
 
 
-def transformer_image_decoder(encoder_output,
+def transformer_image_decoder(x,
+                              encoder_output,
                               ed_attention_bias,
-                              targets,
                               hparams,
                               name="transformer_dec"):
-  """Original Transformer decoder."""
+  """Transformer image decoder over inputs with local attention.
+
+  Args:
+    x: Tensor of shape [batch, height, width, hidden_dim].
+    encoder_output: Tensor, encoder output of shape [batch, length, hidden_dim].
+    ed_attention_bias: Tensor, bias for x.
+    hparams: Dict, hyperparameters.
+    name: string, variable scope.
+
+  Returns:
+    x: Tensor of shape [batch, height, width, hidden_dim].
+  """
   with tf.variable_scope(name):
-    batch_size = common_layers.shape_list(targets)[0]
+    batch_size = common_layers.shape_list(x)[0]
     # Reshape targets as b, 32, 32, 3*hidden size].
-    targets = tf.reshape(targets, [
+    targets = tf.reshape(x, [
         batch_size, hparams.img_len, hparams.img_len,
         hparams.num_channels*hparams.hidden_size])
 
@@ -369,24 +380,35 @@ def transformer_image_decoder(encoder_output,
     return decoder_output
 
 
-def transformer_latent_decoder(encoder_output,
+def transformer_latent_decoder(x,
+                               encoder_output,
                                ed_attention_bias,
-                               targets,
                                hparams,
                                name="transformer_latent_dec"):
-  """Original Transformer decoder."""
+  """Transformer decoder over latents using latent_attention_type.
+
+  Args:
+    x: Tensor of shape [batch, height, width, hidden_dim].
+    encoder_output: Tensor, encoder output of shape [batch, length, hidden_dim].
+    ed_attention_bias: Tensor, bias for x.
+    hparams: Dict, hyperparameters.
+    name: string, variable scope.
+
+  Returns:
+    x: Tensor of shape [batch, height, width, hidden_dim].
+  """
   with tf.variable_scope(name):
-    batch_size = common_layers.shape_list(targets)[0]
+    batch_size = common_layers.shape_list(x)[0]
     compress_ratio = 2**(hparams.num_compress_steps // 2)
     # Reshape targets as b, 32, 32, 3*hidden size].
-    targets = tf.reshape(targets, [
+    x = tf.reshape(x, [
         batch_size, hparams.img_len / compress_ratio,
         (hparams.img_len*hparams.num_latents) / compress_ratio,
         hparams.hidden_size
     ])
 
     # Prepare decoder inputs and bias.
-    decoder_input, _, _ = cia.prepare_decoder(targets, hparams)
+    decoder_input, _, _ = cia.prepare_decoder(x, hparams)
     # hparams.num_channels = 3
     decoder_output = cia.transformer_decoder_layers(
         decoder_input,
@@ -425,15 +447,13 @@ def latent_prediction_model(inputs,
                             latents_discrete,
                             latents_dense,
                             hparams,
-                            name="latent_pred"):
+                            name="latent_prediction"):
   """Transformer based latent prediction model."""
   with tf.variable_scope(name):
-    # latents_dense = discretization.parametrized_unbottleneck(
-    #     latents_discrete, hparams.hidden_size, hparams)
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
-      latents_pred = transformer_latent_decoder(inputs, ed_attention_bias,
-                                                tf.stop_gradient(latents_dense),
-                                                hparams, name + "_extra")
+      latents_pred = transformer_latent_decoder(
+          tf.stop_gradient(latents_dense), inputs, ed_attention_bias,
+          hparams, name)
       _, latent_pred_loss = ae_latent_softmax(
           latents_pred, tf.stop_gradient(latents_discrete), hparams)
   return latents_pred, latent_pred_loss
@@ -445,12 +465,13 @@ def transformer_autoencoder(inputs,
                             hparams,
                             cache=None,
                             predict_mask=1.0):
-  """AE Transformer, main step used for training."""
+  """Auto Encoder using transformer decoder and prior over latents."""
   # Define losses
   losses = {"extra": tf.constant(0.0), "latent_pred": tf.constant(0.0)}
 
   # Reshape image targets as 4d tensor.
   original_targets_shape = common_layers.shape_list(targets)
+  batch_size = original_targets_shape[0]
   if len(original_targets_shape) == 4:
     compress_fn = compress_encoder_2d
     decompress_fn = decompress_decoder_2d
@@ -472,15 +493,13 @@ def transformer_autoencoder(inputs,
   targets, _, _ = cia.maybe_reshape_4d_to_3d(targets)
 
   # Following code creates an exponentially decaying variable based on which
-  # we rescale the los values.
-  batch_size = common_layers.shape_list(targets_c)[0]
+  # we rescale the loss values.
   pc = common_layers.inverse_exp_decay(hparams.startup_steps)
   pc = pc if hparams.mode == tf.estimator.ModeKeys.TRAIN else 1.0
   cond = tf.less(tf.random_uniform([batch_size]), pc)
 
-  # TODO(lukaszkaiser): return extra losses batchwise, multiply before mean.
-  # Call bottleneck layer to get the latents.
-  # Returns embedded latents, discrete latents, loss and the embedding function.
+  # Call bottleneck layer, that takes encoder output and outputs the latents.
+  # Returns embedded latents, discrete latent codes, loss.
   if hparams.mode != tf.estimator.ModeKeys.PREDICT:
     latents_dense, latents_discrete, extra_loss = (
         bottleneck_layer(targets_c, hparams))
@@ -495,7 +514,12 @@ def transformer_autoencoder(inputs,
         hparams,
         name="latent_pred")
     latents_pred_loss = tf.reduce_mean(latents_pred_loss) * tf.to_float(cond)
-    # Assign latent loss
+    # Latent dropout.
+    latents_shape = common_layers.shape_list(latents_dense)
+    latents_dense = tf.nn.dropout(
+        latents_dense, 1 - hparams.latent_dropout,
+        noise_shape=[latents_shape[0], latents_shape[1], 1])
+    # Assign latent loss.
     losses["latent_pred"] = latents_pred_loss
     losses["extra_loss"] = extra_loss
   else:
@@ -524,7 +548,7 @@ def transformer_autoencoder(inputs,
   latents_decoder = decompress_fn(latents_decoder, hparams, name="decompress")
   # if we're operating in 2d space on images, then we're assuming that the
   # last dimension will not be a multiple of channels
-  latents_decoder = tf.reshape(
+  output = tf.reshape(
       latents_decoder,
       shape=[-1, hparams.img_len, hparams.img_len, hparams.hidden_size])
 
@@ -536,20 +560,19 @@ def transformer_autoencoder(inputs,
     mask = tf.less(masking, tf.random_uniform(
         common_layers.shape_list(targets)[:-1]))
     mask = tf.expand_dims(tf.to_float(mask), 2)
-    targets = mask * targets + (1.0 - mask) * latents_decoder
-  else:
-    targets = latents_decoder
+    output = mask * targets + (1.0 - mask) * output
+
   # reshape back to 4d here
-  targets = tf.reshape(targets, original_targets_shape)
+  output = tf.reshape(output, original_targets_shape)
   if hparams.decode_autoregressive:
     # Transformer decoder, that goes from inputs->targets
-    res = transformer_image_decoder(inputs, ed_attention_bias,
-                                    targets, hparams, "decoder")
+    decoder_output = transformer_image_decoder(
+        output, inputs, ed_attention_bias, hparams, "decoder")
   else:
-    res = targets
+    decoder_output = output
 
   # We'll start training the extra model of latents after mask_startup_steps.
   latent_time = tf.less(hparams.mask_startup_steps,
                         tf.to_int32(tf.train.get_global_step()))
   losses["latent_pred"] *= tf.to_float(latent_time)
-  return res, losses, cache
+  return decoder_output, losses, cache

From 01ce8acf20ba4cd71477aca700922315b7459125 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 22 Jun 2018 15:45:14 -0700
Subject: [PATCH 0196/2720] Move BasicAutoencoder out of basic to autoencoders
 and adjust callers.

PiperOrigin-RevId: 201753337
---
 tensor2tensor/models/basic.py                 | 273 +----------------
 tensor2tensor/models/basic_test.py            |  17 --
 tensor2tensor/models/research/autoencoders.py | 280 +++++++++++++++++-
 .../models/research/autoencoders_test.py      |   4 +
 tensor2tensor/models/vanilla_gan.py           |  15 +-
 5 files changed, 294 insertions(+), 295 deletions(-)

diff --git a/tensor2tensor/models/basic.py b/tensor2tensor/models/basic.py
index bcaf8a872..109190951 100644
--- a/tensor2tensor/models/basic.py
+++ b/tensor2tensor/models/basic.py
@@ -17,6 +17,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import registry
@@ -25,14 +26,6 @@
 import tensorflow as tf
 
 
-def lrelu(input_, leak=0.2, name="lrelu"):
-  return tf.maximum(input_, leak * input_, name=name)
-
-
-def reverse_gradient(x):
-  return -x + tf.stop_gradient(2 * x)
-
-
 @registry.register_model
 class BasicFcRelu(t2t_model.T2TModel):
   """Basic fully-connected + ReLU model."""
@@ -49,239 +42,6 @@ def body(self, features):
     return tf.expand_dims(tf.expand_dims(x, axis=1), axis=1)  # 4D For T2T.
 
 
-@registry.register_model
-class BasicAutoencoder(t2t_model.T2TModel):
-  """A basic autoencoder, try with image_mnist_rev or image_cifar10_rev."""
-
-  def __init__(self, *args, **kwargs):
-    super(BasicAutoencoder, self).__init__(*args, **kwargs)
-    self._cur_bottleneck_tensor = None
-    self.is1d = None
-
-  def bottleneck(self, x):
-    with tf.variable_scope("bottleneck"):
-      hparams = self.hparams
-      x = tf.layers.dense(x, hparams.bottleneck_bits, name="bottleneck")
-      if hparams.mode == tf.estimator.ModeKeys.TRAIN:
-        noise = 2.0 * tf.random_uniform(common_layers.shape_list(x)) - 1.0
-        return tf.tanh(x) + noise * hparams.bottleneck_noise, 0.0
-      return tf.tanh(x), 0.0
-
-  def discriminator(self, x, is_training):
-    """Discriminator architecture based on InfoGAN.
-
-    Args:
-      x: input images, shape [bs, h, w, channels]
-      is_training: boolean, are we in train or eval model.
-
-    Returns:
-      out_logit: the output logits (before sigmoid).
-    """
-    hparams = self.hparams
-    with tf.variable_scope(
-        "discriminator",
-        initializer=tf.random_normal_initializer(stddev=0.02)):
-      batch_size, height, width = common_layers.shape_list(x)[:3]
-      # Mapping x from [bs, h, w, c] to [bs, 1]
-      net = tf.layers.conv2d(x, 64, (4, 4), strides=(2, 2),
-                             padding="SAME", name="d_conv1")
-      # [bs, h/2, w/2, 64]
-      net = lrelu(net)
-      net = tf.layers.conv2d(net, 128, (4, 4), strides=(2, 2),
-                             padding="SAME", name="d_conv2")
-      # [bs, h/4, w/4, 128]
-      if hparams.discriminator_batchnorm:
-        net = tf.layers.batch_normalization(net, training=is_training,
-                                            momentum=0.999, name="d_bn2")
-      net = lrelu(net)
-      size = height * width
-      net = tf.reshape(net, [batch_size, size * 8])  # [bs, h * w * 8]
-      net = tf.layers.dense(net, 1024, name="d_fc3")  # [bs, 1024]
-      if hparams.discriminator_batchnorm:
-        net = tf.layers.batch_normalization(net, training=is_training,
-                                            momentum=0.999, name="d_bn3")
-      net = lrelu(net)
-      return net
-
-  def unbottleneck(self, x, res_size, reuse=None):
-    with tf.variable_scope("unbottleneck", reuse=reuse):
-      x = tf.layers.dense(x, res_size, name="dense")
-      return x
-
-  def make_even_size(self, x):
-    if not self.is1d:
-      return common_layers.make_even_size(x)
-    shape1 = x.get_shape().as_list()[1]
-    if shape1 is not None and shape1 % 2 == 0:
-      return x
-    x, _ = common_layers.pad_to_same_length(
-        x, x, final_length_divisible_by=2, axis=1)
-    return x
-
-  def encoder(self, x):
-    with tf.variable_scope("encoder"):
-      hparams = self.hparams
-      kernel, strides = self._get_kernel_and_strides()
-      # Down-convolutions.
-      for i in range(hparams.num_hidden_layers):
-        x = self.make_even_size(x)
-        x = tf.layers.conv2d(
-            x, hparams.hidden_size * 2**(i + 1), kernel, strides=strides,
-            padding="SAME", activation=common_layers.belu, name="conv_%d" % i)
-        x = common_layers.layer_norm(x)
-      return x
-
-  def decoder(self, x):
-    with tf.variable_scope("decoder"):
-      hparams = self.hparams
-      kernel, strides = self._get_kernel_and_strides()
-      # Up-convolutions.
-      for i in range(hparams.num_hidden_layers):
-        j = hparams.num_hidden_layers - i - 1
-        x = tf.layers.conv2d_transpose(
-            x, hparams.hidden_size * 2**j, kernel, strides=strides,
-            padding="SAME", activation=common_layers.belu, name="deconv_%d" % j)
-        x = common_layers.layer_norm(x)
-      return x
-
-  def body(self, features):
-    hparams = self.hparams
-    is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
-    if hparams.mode != tf.estimator.ModeKeys.PREDICT:
-      x = features["targets"]
-      shape = common_layers.shape_list(x)
-      is1d = shape[2] == 1
-      self.is1d = is1d
-      # Run encoder.
-      x = self.encoder(x)
-      # Bottleneck (mix during early training, not too important but stable).
-      b, b_loss = self.bottleneck(x)
-      self._cur_bottleneck_tensor = b
-      b = self.unbottleneck(b, common_layers.shape_list(x)[-1])
-      b = common_layers.mix(b, x, hparams.bottleneck_warmup_steps, is_training)
-      if hparams.add_gan_loss:
-        # Add a purely sampled batch on which we'll compute the GAN loss.
-        g = self.unbottleneck(self.sample(), common_layers.shape_list(x)[-1],
-                              reuse=True)
-        b = tf.concat([g, b], axis=0)
-      # With probability bottleneck_max_prob use the bottleneck, otherwise x.
-      if hparams.bottleneck_max_prob < -1.0:
-        x = tf.where(tf.less(tf.random_uniform([]),
-                             hparams.bottleneck_max_prob), b, x)
-      else:
-        x = b
-    else:
-      if self._cur_bottleneck_tensor is None:
-        b = self.sample()
-      else:
-        b = self._cur_bottleneck_tensor
-      res_size = self.hparams.hidden_size * 2**self.hparams.num_hidden_layers
-      res_size = min(res_size, hparams.max_hidden_size)
-      x = self.unbottleneck(b, res_size)
-    # Run decoder.
-    x = self.decoder(x)
-    if hparams.mode == tf.estimator.ModeKeys.PREDICT:
-      return x, {"bottleneck_loss": 0.0}
-    # Cut to the right size and mix before returning.
-    res = x[:, :shape[1], :shape[2], :]
-    # Add GAN loss if requested.
-    gan_loss = 0.0
-    if hparams.add_gan_loss:
-      # Split back if we added a purely sampled batch.
-      res_gan, res = tf.split(res, 2, axis=0)
-      num_channels = self.hparams.problem.num_channels
-      res_rgb = common_layers.convert_real_to_rgb(tf.nn.sigmoid(
-          tf.layers.dense(res_gan, num_channels, name="gan_rgb")))
-      tf.summary.image("gan", tf.cast(res_rgb, tf.uint8), max_outputs=1)
-      orig_rgb = tf.to_float(features["targets_raw"])
-      def discriminate(x):
-        return self.discriminator(x, is_training=is_training)
-      gan_loss = common_layers.sliced_gan_loss(
-          orig_rgb, reverse_gradient(res_rgb),
-          discriminate, self.hparams.num_sliced_vecs)
-      gan_loss *= common_layers.inverse_lin_decay(
-          hparams.bottleneck_warmup_steps)
-    # Mix the final result and return.
-    res = common_layers.mix(res, features["targets"],
-                            hparams.bottleneck_warmup_steps // 2, is_training)
-    return res, {"bottleneck_loss": b_loss, "gan_loss": - gan_loss}
-
-  def sample(self, features=None, shape=None):
-    del features, shape
-    hp = self.hparams
-    div_x = 2**hp.num_hidden_layers
-    div_y = 1 if self.is1d else 2**hp.num_hidden_layers
-    size = [hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
-            hp.bottleneck_bits]
-    # Sample in [-1, 1] as the bottleneck is under tanh.
-    return 2.0 * tf.random_uniform(size) - 1.0
-
-  def encode(self, x):
-    """Auto-encode x and return the bottleneck."""
-    features = {"targets": x}
-    self(features)  # pylint: disable=not-callable
-    res = tf.maximum(0.0, self._cur_bottleneck_tensor)  # Be 0/1 and not -1/1.
-    self._cur_bottleneck_tensor = None
-    return res
-
-  def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
-    """Produce predictions from the model by sampling."""
-    del args, kwargs
-    # Inputs and features preparation needed to handle edge cases.
-    if not features:
-      features = {}
-    inputs_old = None
-    if "inputs" in features and len(features["inputs"].shape) < 4:
-      inputs_old = features["inputs"]
-      features["inputs"] = tf.expand_dims(features["inputs"], 2)
-
-    # Sample and decode.
-    # TODO(lukaszkaiser): is this a universal enough way to get channels?
-    try:
-      num_channels = self.hparams.problem.num_channels
-    except AttributeError:
-      num_channels = 1
-    if "targets" not in features:
-      features["targets"] = tf.zeros(
-          [self.hparams.batch_size, 1, 1, num_channels],
-          dtype=tf.int32)
-    logits, _ = self(features)  # pylint: disable=not-callable
-    samples = tf.argmax(logits, axis=-1)
-
-    # Restore inputs to not confuse Estimator in edge cases.
-    if inputs_old is not None:
-      features["inputs"] = inputs_old
-
-    # Return samples.
-    return samples
-
-  def decode(self, bottleneck):
-    """Auto-decode from the bottleneck and return the result."""
-    # Get the shape from bottleneck and num channels.
-    shape = common_layers.shape_list(bottleneck)
-    try:
-      num_channels = self.hparams.problem.num_channels
-    except AttributeError:
-      num_channels = 1
-    dummy_targets = tf.zeros(shape[:-1] + [num_channels])
-    # Set the bottleneck to decode.
-    if len(shape) > 4:
-      bottleneck = tf.squeeze(bottleneck, axis=[1])
-    bottleneck = 2 * bottleneck - 1  # Be -1/1 instead of 0/1.
-    self._cur_bottleneck_tensor = bottleneck
-    # Run decoding.
-    res = self.infer({"targets": dummy_targets})
-    self._cur_bottleneck_tensor = None
-    return res
-
-  def _get_kernel_and_strides(self):
-    hparams = self.hparams
-    kernel = (hparams.kernel_height, hparams.kernel_width)
-    kernel = (hparams.kernel_height, 1) if self.is1d else kernel
-    strides = (2, 1) if self.is1d else (2, 2)
-    return (kernel, strides)
-
-
 @registry.register_hparams
 def basic_fc_small():
   """Small fully connected model."""
@@ -295,34 +55,3 @@ def basic_fc_small():
   hparams.weight_decay = 0.0
   hparams.dropout = 0.0
   return hparams
-
-
-@registry.register_hparams
-def basic_autoencoder():
-  """Basic autoencoder model."""
-  hparams = common_hparams.basic_params1()
-  hparams.optimizer = "Adam"
-  hparams.learning_rate_constant = 0.0002
-  hparams.learning_rate_warmup_steps = 500
-  hparams.learning_rate_schedule = "constant * linear_warmup"
-  hparams.label_smoothing = 0.0
-  hparams.batch_size = 128
-  hparams.hidden_size = 64
-  hparams.num_hidden_layers = 5
-  hparams.initializer = "uniform_unit_scaling"
-  hparams.initializer_gain = 1.0
-  hparams.weight_decay = 0.0
-  hparams.kernel_height = 4
-  hparams.kernel_width = 4
-  hparams.dropout = 0.1
-  hparams.add_hparam("max_hidden_size", 1024)
-  hparams.add_hparam("bottleneck_bits", 128)
-  hparams.add_hparam("bottleneck_noise", 0.1)
-  hparams.add_hparam("bottleneck_warmup_steps", 3000)
-  hparams.add_hparam("bottleneck_max_prob", 1.0)
-  hparams.add_hparam("sample_height", 32)
-  hparams.add_hparam("sample_width", 32)
-  hparams.add_hparam("discriminator_batchnorm", True)
-  hparams.add_hparam("num_sliced_vecs", 4096)
-  hparams.add_hparam("add_gan_loss", False)
-  return hparams
diff --git a/tensor2tensor/models/basic_test.py b/tensor2tensor/models/basic_test.py
index 7e9b8fd42..c50b103ca 100644
--- a/tensor2tensor/models/basic_test.py
+++ b/tensor2tensor/models/basic_test.py
@@ -44,23 +44,6 @@ def testBasicFcRelu(self):
       res = session.run(logits)
     self.assertEqual(res.shape, (1, 1, 1, 1, 10))
 
-  def testBasicAutoencoder(self):
-    x = np.random.random_integers(0, high=255, size=(1, 28, 28, 1))
-    y = np.random.random_integers(0, high=9, size=(1, 1))
-    hparams = trainer_lib.create_hparams(
-        "basic_autoencoder", problem_name="image_mnist_rev", data_dir=".")
-    with self.test_session() as session:
-      features = {
-          "targets": tf.constant(x, dtype=tf.int32),
-          "inputs": tf.constant(y, dtype=tf.int32),
-      }
-      tf.train.create_global_step()
-      model = basic.BasicAutoencoder(hparams, tf.estimator.ModeKeys.TRAIN)
-      logits, _ = model(features)
-      session.run(tf.global_variables_initializer())
-      res = session.run(logits)
-    self.assertEqual(res.shape, (1, 28, 28, 1, 256))
-
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 2c7d9e62d..2c63018fd 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -17,17 +17,260 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
 from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
-from tensor2tensor.models import basic
 from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
 
 import tensorflow as tf
 
 
+def lrelu(input_, leak=0.2, name="lrelu"):
+  return tf.maximum(input_, leak * input_, name=name)
+
+
+def reverse_gradient(x):
+  return -x + tf.stop_gradient(2 * x)
+
+
+@registry.register_model
+class AutoencoderBasic(t2t_model.T2TModel):
+  """A basic autoencoder, try with image_mnist_rev or image_cifar10_rev."""
+
+  def __init__(self, *args, **kwargs):
+    super(AutoencoderBasic, self).__init__(*args, **kwargs)
+    self._cur_bottleneck_tensor = None
+    self.is1d = None
+
+  def bottleneck(self, x):
+    with tf.variable_scope("bottleneck"):
+      hparams = self.hparams
+      x = tf.layers.dense(x, hparams.bottleneck_bits, name="bottleneck")
+      if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+        noise = 2.0 * tf.random_uniform(common_layers.shape_list(x)) - 1.0
+        return tf.tanh(x) + noise * hparams.bottleneck_noise, 0.0
+      return tf.tanh(x), 0.0
+
+  def discriminator(self, x, is_training):
+    """Discriminator architecture based on InfoGAN.
+
+    Args:
+      x: input images, shape [bs, h, w, channels]
+      is_training: boolean, are we in train or eval model.
+
+    Returns:
+      out_logit: the output logits (before sigmoid).
+    """
+    hparams = self.hparams
+    with tf.variable_scope(
+        "discriminator",
+        initializer=tf.random_normal_initializer(stddev=0.02)):
+      batch_size, height, width = common_layers.shape_list(x)[:3]
+      # Mapping x from [bs, h, w, c] to [bs, 1]
+      net = tf.layers.conv2d(x, 64, (4, 4), strides=(2, 2),
+                             padding="SAME", name="d_conv1")
+      # [bs, h/2, w/2, 64]
+      net = lrelu(net)
+      net = tf.layers.conv2d(net, 128, (4, 4), strides=(2, 2),
+                             padding="SAME", name="d_conv2")
+      # [bs, h/4, w/4, 128]
+      if hparams.discriminator_batchnorm:
+        net = tf.layers.batch_normalization(net, training=is_training,
+                                            momentum=0.999, name="d_bn2")
+      net = lrelu(net)
+      size = height * width
+      net = tf.reshape(net, [batch_size, size * 8])  # [bs, h * w * 8]
+      net = tf.layers.dense(net, 1024, name="d_fc3")  # [bs, 1024]
+      if hparams.discriminator_batchnorm:
+        net = tf.layers.batch_normalization(net, training=is_training,
+                                            momentum=0.999, name="d_bn3")
+      net = lrelu(net)
+      return net
+
+  def unbottleneck(self, x, res_size, reuse=None):
+    with tf.variable_scope("unbottleneck", reuse=reuse):
+      x = tf.layers.dense(x, res_size, name="dense")
+      return x
+
+  def make_even_size(self, x):
+    if not self.is1d:
+      return common_layers.make_even_size(x)
+    shape1 = x.get_shape().as_list()[1]
+    if shape1 is not None and shape1 % 2 == 0:
+      return x
+    x, _ = common_layers.pad_to_same_length(
+        x, x, final_length_divisible_by=2, axis=1)
+    return x
+
+  def encoder(self, x):
+    with tf.variable_scope("encoder"):
+      hparams = self.hparams
+      kernel, strides = self._get_kernel_and_strides()
+      # Down-convolutions.
+      for i in range(hparams.num_hidden_layers):
+        x = self.make_even_size(x)
+        x = tf.layers.conv2d(
+            x, hparams.hidden_size * 2**(i + 1), kernel, strides=strides,
+            padding="SAME", activation=common_layers.belu, name="conv_%d" % i)
+        x = common_layers.layer_norm(x)
+      return x
+
+  def decoder(self, x):
+    with tf.variable_scope("decoder"):
+      hparams = self.hparams
+      kernel, strides = self._get_kernel_and_strides()
+      # Up-convolutions.
+      for i in range(hparams.num_hidden_layers):
+        j = hparams.num_hidden_layers - i - 1
+        x = tf.layers.conv2d_transpose(
+            x, hparams.hidden_size * 2**j, kernel, strides=strides,
+            padding="SAME", activation=common_layers.belu, name="deconv_%d" % j)
+        x = common_layers.layer_norm(x)
+      return x
+
+  def body(self, features):
+    hparams = self.hparams
+    is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+    if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+      x = features["targets"]
+      shape = common_layers.shape_list(x)
+      is1d = shape[2] == 1
+      self.is1d = is1d
+      # Run encoder.
+      x = self.encoder(x)
+      # Bottleneck (mix during early training, not too important but stable).
+      b, b_loss = self.bottleneck(x)
+      self._cur_bottleneck_tensor = b
+      b = self.unbottleneck(b, common_layers.shape_list(x)[-1])
+      b = common_layers.mix(b, x, hparams.bottleneck_warmup_steps, is_training)
+      if hparams.add_gan_loss:
+        # Add a purely sampled batch on which we'll compute the GAN loss.
+        g = self.unbottleneck(self.sample(), common_layers.shape_list(x)[-1],
+                              reuse=True)
+        b = tf.concat([g, b], axis=0)
+      # With probability bottleneck_max_prob use the bottleneck, otherwise x.
+      if hparams.bottleneck_max_prob < -1.0:
+        x = tf.where(tf.less(tf.random_uniform([]),
+                             hparams.bottleneck_max_prob), b, x)
+      else:
+        x = b
+    else:
+      if self._cur_bottleneck_tensor is None:
+        b = self.sample()
+      else:
+        b = self._cur_bottleneck_tensor
+      res_size = self.hparams.hidden_size * 2**self.hparams.num_hidden_layers
+      res_size = min(res_size, hparams.max_hidden_size)
+      x = self.unbottleneck(b, res_size)
+    # Run decoder.
+    x = self.decoder(x)
+    if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+      return x, {"bottleneck_loss": 0.0}
+    # Cut to the right size and mix before returning.
+    res = x[:, :shape[1], :shape[2], :]
+    # Add GAN loss if requested.
+    gan_loss = 0.0
+    if hparams.add_gan_loss:
+      # Split back if we added a purely sampled batch.
+      res_gan, res = tf.split(res, 2, axis=0)
+      num_channels = self.hparams.problem.num_channels
+      res_rgb = common_layers.convert_real_to_rgb(tf.nn.sigmoid(
+          tf.layers.dense(res_gan, num_channels, name="gan_rgb")))
+      tf.summary.image("gan", tf.cast(res_rgb, tf.uint8), max_outputs=1)
+      orig_rgb = tf.to_float(features["targets_raw"])
+      def discriminate(x):
+        return self.discriminator(x, is_training=is_training)
+      gan_loss = common_layers.sliced_gan_loss(
+          orig_rgb, reverse_gradient(res_rgb),
+          discriminate, self.hparams.num_sliced_vecs)
+      gan_loss *= common_layers.inverse_lin_decay(
+          hparams.bottleneck_warmup_steps)
+    # Mix the final result and return.
+    res = common_layers.mix(res, features["targets"],
+                            hparams.bottleneck_warmup_steps // 2, is_training)
+    return res, {"bottleneck_loss": b_loss, "gan_loss": - gan_loss}
+
+  def sample(self, features=None, shape=None):
+    del features, shape
+    hp = self.hparams
+    div_x = 2**hp.num_hidden_layers
+    div_y = 1 if self.is1d else 2**hp.num_hidden_layers
+    size = [hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
+            hp.bottleneck_bits]
+    # Sample in [-1, 1] as the bottleneck is under tanh.
+    return 2.0 * tf.random_uniform(size) - 1.0
+
+  def encode(self, x):
+    """Auto-encode x and return the bottleneck."""
+    features = {"targets": x}
+    self(features)  # pylint: disable=not-callable
+    res = tf.maximum(0.0, self._cur_bottleneck_tensor)  # Be 0/1 and not -1/1.
+    self._cur_bottleneck_tensor = None
+    return res
+
+  def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
+    """Produce predictions from the model by sampling."""
+    del args, kwargs
+    # Inputs and features preparation needed to handle edge cases.
+    if not features:
+      features = {}
+    inputs_old = None
+    if "inputs" in features and len(features["inputs"].shape) < 4:
+      inputs_old = features["inputs"]
+      features["inputs"] = tf.expand_dims(features["inputs"], 2)
+
+    # Sample and decode.
+    # TODO(lukaszkaiser): is this a universal enough way to get channels?
+    try:
+      num_channels = self.hparams.problem.num_channels
+    except AttributeError:
+      num_channels = 1
+    if "targets" not in features:
+      features["targets"] = tf.zeros(
+          [self.hparams.batch_size, 1, 1, num_channels],
+          dtype=tf.int32)
+    logits, _ = self(features)  # pylint: disable=not-callable
+    samples = tf.argmax(logits, axis=-1)
+
+    # Restore inputs to not confuse Estimator in edge cases.
+    if inputs_old is not None:
+      features["inputs"] = inputs_old
+
+    # Return samples.
+    return samples
+
+  def decode(self, bottleneck):
+    """Auto-decode from the bottleneck and return the result."""
+    # Get the shape from bottleneck and num channels.
+    shape = common_layers.shape_list(bottleneck)
+    try:
+      num_channels = self.hparams.problem.num_channels
+    except AttributeError:
+      num_channels = 1
+    dummy_targets = tf.zeros(shape[:-1] + [num_channels])
+    # Set the bottleneck to decode.
+    if len(shape) > 4:
+      bottleneck = tf.squeeze(bottleneck, axis=[1])
+    bottleneck = 2 * bottleneck - 1  # Be -1/1 instead of 0/1.
+    self._cur_bottleneck_tensor = bottleneck
+    # Run decoding.
+    res = self.infer({"targets": dummy_targets})
+    self._cur_bottleneck_tensor = None
+    return res
+
+  def _get_kernel_and_strides(self):
+    hparams = self.hparams
+    kernel = (hparams.kernel_height, hparams.kernel_width)
+    kernel = (hparams.kernel_height, 1) if self.is1d else kernel
+    strides = (2, 1) if self.is1d else (2, 2)
+    return (kernel, strides)
+
+
 @registry.register_model
-class AutoencoderAutoregressive(basic.BasicAutoencoder):
+class AutoencoderAutoregressive(AutoencoderBasic):
   """Autoencoder with an autoregressive part."""
 
   def body(self, features):
@@ -422,10 +665,41 @@ def body(self, features):
     return res, losses
 
 
+@registry.register_hparams
+def autoencoder_basic():
+  """Basic autoencoder model."""
+  hparams = common_hparams.basic_params1()
+  hparams.optimizer = "Adam"
+  hparams.learning_rate_constant = 0.0002
+  hparams.learning_rate_warmup_steps = 500
+  hparams.learning_rate_schedule = "constant * linear_warmup"
+  hparams.label_smoothing = 0.0
+  hparams.batch_size = 128
+  hparams.hidden_size = 64
+  hparams.num_hidden_layers = 5
+  hparams.initializer = "uniform_unit_scaling"
+  hparams.initializer_gain = 1.0
+  hparams.weight_decay = 0.0
+  hparams.kernel_height = 4
+  hparams.kernel_width = 4
+  hparams.dropout = 0.1
+  hparams.add_hparam("max_hidden_size", 1024)
+  hparams.add_hparam("bottleneck_bits", 128)
+  hparams.add_hparam("bottleneck_noise", 0.1)
+  hparams.add_hparam("bottleneck_warmup_steps", 3000)
+  hparams.add_hparam("bottleneck_max_prob", 1.0)
+  hparams.add_hparam("sample_height", 32)
+  hparams.add_hparam("sample_width", 32)
+  hparams.add_hparam("discriminator_batchnorm", True)
+  hparams.add_hparam("num_sliced_vecs", 4096)
+  hparams.add_hparam("add_gan_loss", False)
+  return hparams
+
+
 @registry.register_hparams
 def autoencoder_autoregressive():
   """Autoregressive autoencoder model."""
-  hparams = basic.basic_autoencoder()
+  hparams = autoencoder_basic()
   hparams.add_hparam("autoregressive_forget_base", False)
   hparams.add_hparam("autoregressive_mode", "none")
   hparams.add_hparam("autoregressive_dropout", 0.4)
diff --git a/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index fc45aeede..dc1c79baa 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -52,6 +52,10 @@ def get_mnist_random_output(self, model_name, hparams_set=None,
   def mnist_output_shape(self):
     return (1, 28, 28, 1, 256)
 
+  def testAutoencoderBasic(self):
+    res = self.get_mnist_random_output("autoencoder_basic")
+    self.assertEqual(res.shape, self.mnist_output_shape)
+
   def testAutoencoderAutoregressive(self):
     res = self.get_mnist_random_output("autoencoder_autoregressive")
     self.assertEqual(res.shape, self.mnist_output_shape)
diff --git a/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index 730d5b847..4e4686bd5 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -21,8 +21,8 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
-from tensor2tensor.models import basic
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -197,10 +197,19 @@ def infer(self, *args, **kwargs):  # pylint: disable=arguments-differ
 @registry.register_hparams
 def sliced_gan():
   """Basic parameters for a vanilla_gan."""
-  hparams = basic.basic_autoencoder()
-  hparams.hidden_size = 128
+  hparams = common_hparams.basic_params1()
+  hparams.optimizer = "Adam"
+  hparams.learning_rate_constant = 0.0002
+  hparams.learning_rate_warmup_steps = 500
+  hparams.learning_rate_schedule = "constant * linear_warmup"
+  hparams.label_smoothing = 0.0
   hparams.batch_size = 128
+  hparams.hidden_size = 128
+  hparams.initializer = "uniform_unit_scaling"
+  hparams.initializer_gain = 1.0
   hparams.weight_decay = 1e-6
+  hparams.kernel_height = 4
+  hparams.kernel_width = 4
   hparams.bottleneck_bits = 128
   hparams.add_hparam("discriminator_batchnorm", True)
   hparams.add_hparam("num_sliced_vecs", 4096)

From e5fff475d3d7b12e7edbbd1a61bf702965e19281 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Fri, 22 Jun 2018 15:52:13 -0700
Subject: [PATCH 0197/2720] Internal change

PiperOrigin-RevId: 201754318
---
 tensor2tensor/utils/trainer_lib.py | 44 ++++++++++++++----------------
 1 file changed, 20 insertions(+), 24 deletions(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 3fd0ccc9f..ee8a1962d 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -262,37 +262,35 @@ def create_hooks(use_tfdbg=False,
 class T2TExperiment(object):
   """Custom Experiment class for running distributed experiments."""
 
-  def __init__(self,
-               estimator,
-               hparams,
-               train_spec,
-               eval_spec,
-               decode_hparams=None):
+  def __init__(self, estimator, hparams, train_spec, eval_spec,
+               use_validation_monitor, decode_hparams=None):
     self._train_spec = train_spec
     self._eval_spec = eval_spec
     self._hparams = hparams
     self._decode_hparams = decode_hparams
     self._estimator = estimator
+    self._use_validation_monitor = use_validation_monitor
 
   @property
   def estimator(self):
     return self._estimator
 
-  @property
-  def eval_steps(self):
-    return self._eval_spec.steps
-
   @property
   def train_steps(self):
     return self._train_spec.max_steps
 
-  def continuous_train_and_eval(self):
+  @property
+  def eval_steps(self):
+    return self._eval_spec.steps
+
+  def continuous_train_and_eval(self, continuous_eval_predicate_fn=None):
+    del continuous_eval_predicate_fn
     tf.estimator.train_and_evaluate(self._estimator, self._train_spec,
                                     self._eval_spec)
     return self.evaluate()
 
   def train_and_evaluate(self):
-    if self._eval_spec is None:
+    if self._use_validation_monitor:
       tf.logging.warning("EvalSpec not provided. Estimator will not manage "
                          "model evaluation. Assuming ValidationMonitor present "
                          "in train_hooks.")
@@ -330,7 +328,7 @@ def continuous_eval_on_train_data(self):
 
   def test(self):
     """Perform 1 step of train and 2 step of eval."""
-    if self._eval_spec is None:
+    if self._use_validation_monitor:
       return self.train_and_evaluate()
 
     self._estimator.train(
@@ -470,18 +468,15 @@ def create_experiment(
 
   train_spec = tf.estimator.TrainSpec(
       train_input_fn, max_steps=train_steps, hooks=train_hooks)
-  if use_validation_monitor:
-    eval_spec = None
-  else:
-    eval_spec = tf.estimator.EvalSpec(
-        eval_input_fn,
-        steps=eval_steps,
-        hooks=eval_hooks,
-        start_delay_secs=0 if hparams.schedule == "evaluate" else 120,
-        throttle_secs=eval_throttle_seconds)
-  hooks_kwargs = {"train_hooks": train_hooks, "eval_hooks": eval_hooks}
+  eval_spec = tf.estimator.EvalSpec(
+      eval_input_fn,
+      steps=eval_steps,
+      hooks=eval_hooks,
+      start_delay_secs=0 if hparams.schedule == "evaluate" else 120,
+      throttle_secs=eval_throttle_seconds)
 
   if autotune:
+    hooks_kwargs = {"train_monitors": train_hooks, "eval_hooks": eval_hooks}
     return tf.contrib.learn.Experiment(
         estimator=estimator,
         train_input_fn=train_input_fn,
@@ -493,7 +488,7 @@ def create_experiment(
         eval_delay_secs=0 if schedule == "evaluate" else 120,
         **hooks_kwargs if not use_tpu else {})
   return T2TExperiment(estimator, hparams, train_spec, eval_spec,
-                       decode_hparams)
+                       use_validation_monitor, decode_hparams)
 
 
 def create_experiment_fn(*args, **kwargs):
@@ -538,3 +533,4 @@ def restore_checkpoint(ckpt_dir, saver, sess, must_restore=False):
   saver.restore(sess, path)
   step = int(path.split("-")[-1])
   return step
+

From 60895e3ac21bd235d7339ec801117e6551557d2e Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Fri, 22 Jun 2018 16:20:21 -0700
Subject: [PATCH 0198/2720] Added an option to share the input embedding and
 target embedding without

PiperOrigin-RevId: 201758003
---
 tensor2tensor/layers/common_hparams.py | 11 +++++++----
 tensor2tensor/layers/modalities.py     |  6 ++++--
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 4b6d62934..b338bf2b8 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -154,12 +154,15 @@ def basic_params1():
       # during eval
       eval_run_autoregressive=False,
       # TODO(lukaszkaiser): these parameters should probably be set elsewhere.
-      # in SymbolModality, share the output embeddings and the softmax
-      # variables.
-      # You can also share the input embeddings with the output embeddings
+      # (SymbolModality) - If this flag is on, we try to share all of the input
+      # embeddings, the target embeddings and the softmax weights.
+      shared_embedding_and_softmax_weights=False,
+      # (SymbolModality) - If this flag is on, we try to share the input
+      # embeddings and the target embeddings.
+      # You can also share the input embeddings with the target embeddings
       # by using a problem_hparams that uses the same modality object for
       # the input_modality and target_modality.
-      shared_embedding_and_softmax_weights=False,
+      shared_embedding=False,
       # In SymbolModality, skip the top layer, assume we're providing logits.
       symbol_modality_skip_top=False,
       # For each feature for which you want to override the default input
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 11fa20abd..efbea07b5 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -114,12 +114,14 @@ def bottom_simple(self, x, name, reuse):
 
   def bottom(self, x):
     self._bottom_was_called = True
-    if self._model_hparams.shared_embedding_and_softmax_weights:
+    if (self._model_hparams.shared_embedding_and_softmax_weights or
+        self._model_hparams.get("shared_embedding")):
       return self.bottom_simple(x, "shared", reuse=None)
     return self.bottom_simple(x, "input_emb", reuse=None)
 
   def targets_bottom(self, x):
-    if self._model_hparams.shared_embedding_and_softmax_weights:
+    if (self._model_hparams.shared_embedding_and_softmax_weights or
+        self._model_hparams.get("shared_embedding")):
       try:
         return self.bottom_simple(x, "shared", reuse=True)
       except ValueError:

From 4eb36c1f146c7dc473dbfb93eadfebb86febc77a Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 22 Jun 2018 16:51:38 -0700
Subject: [PATCH 0199/2720] add problem and data generator for Mozilla Common
 Voice dataset

PiperOrigin-RevId: 201761875
---
 README.md                                     |   5 +-
 docs/walkthrough.md                           |   5 +-
 setup.py                                      |   9 +-
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/common_voice.py | 249 ++++++++++++++++++
 .../data_generators/speech_recognition.py     |   9 +-
 tensor2tensor/models/transformer.py           |  14 +
 7 files changed, 283 insertions(+), 9 deletions(-)
 create mode 100644 tensor2tensor/data_generators/common_voice.py

diff --git a/README.md b/README.md
index 1522a9542..10359f764 100644
--- a/README.md
+++ b/README.md
@@ -126,10 +126,13 @@ few steps (e.g., `--train_steps=2000`).
 
 For speech-to-text, we have these data-sets in T2T:
 
-* Librispeech (English speech to text): `--problem=librispeech` for
+* Librispeech (US English): `--problem=librispeech` for
     the whole set and `--problem=librispeech_clean` for a smaller
     but nicely filtered part.
 
+* Mozilla Common Voice (US English): `--problem=common_voice` for the whole set
+    `--problem=common_voice_clean` for a quality-checked subset.
+
 ### Summarization
 
 For summarizing longer text into shorter one we have these data-sets:
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 1522a9542..10359f764 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -126,10 +126,13 @@ few steps (e.g., `--train_steps=2000`).
 
 For speech-to-text, we have these data-sets in T2T:
 
-* Librispeech (English speech to text): `--problem=librispeech` for
+* Librispeech (US English): `--problem=librispeech` for
     the whole set and `--problem=librispeech_clean` for a smaller
     but nicely filtered part.
 
+* Mozilla Common Voice (US English): `--problem=common_voice` for the whole set
+    `--problem=common_voice_clean` for a quality-checked subset.
+
 ### Summarization
 
 For summarizing longer text into shorter one we have these data-sets:
diff --git a/setup.py b/setup.py
index aa1d3622c..6681f335e 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.6.5',
+    version='1.6.6',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',
@@ -46,16 +46,13 @@
         'scipy',
         'sympy',
         'six',
+        'tqdm',
     ],
     extras_require={
         'tensorflow': ['tensorflow>=1.5.0'],
         'tensorflow_gpu': ['tensorflow-gpu>=1.5.0'],
         'tests': [
-            'pytest',
-            'mock',
-            'pylint',
-            'jupyter',
-            'gsutil'
+            'pytest', 'mock', 'pylint', 'jupyter', 'gsutil'
             # Need atari extras for Travis tests, but because gym is already in
             # install_requires, pip skips the atari extras, so we instead do an
             # explicit pip install gym[atari] for the tests.
diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index dcb83832f..2cff8a1dd 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -30,6 +30,7 @@
     "tensor2tensor.data_generators.cifar",
     "tensor2tensor.data_generators.cipher",
     "tensor2tensor.data_generators.cnn_dailymail",
+    "tensor2tensor.data_generators.common_voice",
     "tensor2tensor.data_generators.desc2code",
     "tensor2tensor.data_generators.fsns",
     "tensor2tensor.data_generators.gene_expression",
diff --git a/tensor2tensor/data_generators/common_voice.py b/tensor2tensor/data_generators/common_voice.py
new file mode 100644
index 000000000..4bad08fdb
--- /dev/null
+++ b/tensor2tensor/data_generators/common_voice.py
@@ -0,0 +1,249 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Mozilla Common Voice dataset.
+
+Note: Generating the full set of examples can take upwards of 5 hours.
+As the Common Voice data are distributed in MP3 format, experimenters will need
+to have both SoX (http://sox.sourceforge.net) and on Linux, the libsox-fmt-mp3
+package installed. The original samples will be downsampled by the encoder.
+"""
+
+import csv
+import os
+import tarfile
+import tqdm  # pylint: disable=g-bad-import-order
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import speech_recognition
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+_COMMONVOICE_URL = "https://common-voice-data-download.s3.amazonaws.com/cv_corpus_v1.tar.gz"  # pylint: disable=line-too-long
+
+_COMMONVOICE_TRAIN_DATASETS = ["cv-valid-train", "cv-other-train"]
+_COMMONVOICE_DEV_DATASETS = ["cv-valid-dev", "cv-other-dev"]
+_COMMONVOICE_TEST_DATASETS = ["cv-valid-test", "cv-other-test"]
+
+
+def _collect_data(directory):
+  """Traverses directory collecting input and target files.
+
+  Args:
+   directory: base path to extracted audio and transcripts.
+  Returns:
+   list of (media_base, media_filepath, label) tuples
+  """
+  # Returns:
+  data_files = []
+  transcripts = [
+      filename for filename in os.listdir(directory)
+      if filename.endswith(".csv")
+  ]
+  for transcript in transcripts:
+    transcript_path = os.path.join(directory, transcript)
+    with open(transcript_path, "r") as transcript_file:
+      transcript_reader = csv.reader(transcript_file)
+      _ = transcript_reader.next()  # Skip headers.
+      for transcript_line in transcript_reader:
+        media_name, label = transcript_line[0:2]
+        filename = os.path.join(directory, media_name)
+        data_files.append((media_name, filename, label))
+  return data_files
+
+
+def _file_exists(path, filename):
+  """Checks if the filename exists under the path."""
+  return os.path.isfile(os.path.join(path, filename))
+
+
+def _is_relative(path, filename):
+  """Checks if the filename is relative, not absolute."""
+  return os.path.abspath(os.path.join(path, filename)).startswith(path)
+
+
+@registry.register_problem()
+class CommonVoice(speech_recognition.SpeechRecognitionProblem):
+  """Problem spec for Commonvoice using clean and noisy data."""
+
+  # Select only the clean data
+  TRAIN_DATASETS = _COMMONVOICE_TRAIN_DATASETS[:1]
+  DEV_DATASETS = _COMMONVOICE_DEV_DATASETS[:1]
+  TEST_DATASETS = _COMMONVOICE_TEST_DATASETS[:1]
+
+  @property
+  def num_shards(self):
+    return 100
+
+  @property
+  def use_subword_tokenizer(self):
+    return False
+
+  @property
+  def num_dev_shards(self):
+    return 1
+
+  @property
+  def num_test_shards(self):
+    return 1
+
+  @property
+  def use_train_shards_for_dev(self):
+    """If true, we only generate training data and hold out shards for dev."""
+    return False
+
+  def generator(self,
+                data_dir,
+                tmp_dir,
+                datasets,
+                eos_list=None,
+                start_from=0,
+                how_many=0):
+    del eos_list
+    i = 0
+
+    filename = os.path.basename(_COMMONVOICE_URL)
+    compressed_file = generator_utils.maybe_download(tmp_dir, filename,
+                                                     _COMMONVOICE_URL)
+
+    read_type = "r:gz" if filename.endswith(".tgz") else "r"
+    with tarfile.open(compressed_file, read_type) as corpus_tar:
+      # Create a subset of files that don't already exist.
+      #   tarfile.extractall errors when encountering an existing file
+      #   and tarfile.extract is extremely slow. For security, check that all
+      #   paths are relative.
+      members = [
+          f for f in corpus_tar if _is_relative(tmp_dir, f.name) and
+          not _file_exists(tmp_dir, f.name)
+      ]
+      corpus_tar.extractall(tmp_dir, members=members)
+
+    data_dir = os.path.join(tmp_dir, "cv_corpus_v1")
+    data_tuples = _collect_data(data_dir)
+    encoders = self.feature_encoders(None)
+    audio_encoder = encoders["waveforms"]
+    text_encoder = encoders["targets"]
+    for dataset in datasets:
+      data_tuples = (tup for tup in data_tuples if tup[0].startswith(dataset))
+      for utt_id, media_file, text_data in tqdm.tqdm(
+          sorted(data_tuples)[start_from:]):
+        if how_many > 0 and i == how_many:
+          return
+        i += 1
+        wav_data = audio_encoder.encode(media_file)
+        yield {
+            "waveforms": wav_data,
+            "waveform_lens": [len(wav_data)],
+            "targets": text_encoder.encode(text_data),
+            "raw_transcript": [text_data],
+            "utt_id": [utt_id],
+            "spk_id": ["unknown"],
+        }
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    train_paths = self.training_filepaths(
+        data_dir, self.num_shards, shuffled=False)
+    dev_paths = self.dev_filepaths(
+        data_dir, self.num_dev_shards, shuffled=False)
+    test_paths = self.test_filepaths(
+        data_dir, self.num_test_shards, shuffled=True)
+
+    generator_utils.generate_files(
+        self.generator(data_dir, tmp_dir, self.TEST_DATASETS), test_paths)
+
+    if self.use_train_shards_for_dev:
+      all_paths = train_paths + dev_paths
+      generator_utils.generate_files(
+          self.generator(data_dir, tmp_dir, self.TRAIN_DATASETS), all_paths)
+      generator_utils.shuffle_dataset(all_paths)
+    else:
+      generator_utils.generate_dataset_and_shuffle(
+          self.generator(data_dir, tmp_dir, self.TRAIN_DATASETS), train_paths,
+          self.generator(data_dir, tmp_dir, self.DEV_DATASETS), dev_paths)
+
+
+@registry.register_problem()
+class CommonVoiceTrainFullTestClean(CommonVoice):
+  """Problem to train on full set, but evaluate on clean data only."""
+
+  def training_filepaths(self, data_dir, num_shards, shuffled):
+    return CommonVoice.training_filepaths(self, data_dir, num_shards, shuffled)
+
+  def dev_filepaths(self, data_dir, num_shards, shuffled):
+    return CommonVoiceClean.dev_filepaths(self, data_dir, num_shards, shuffled)
+
+  def test_filepaths(self, data_dir, num_shards, shuffled):
+    return CommonVoiceClean.test_filepaths(self, data_dir, num_shards, shuffled)
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    raise Exception("Generate Commonvoice and Commonvoice_clean data.")
+
+  def filepattern(self, data_dir, mode, shard=None):
+    """Get filepattern for data files for mode.
+
+    Matches mode to a suffix.
+    * DatasetSplit.TRAIN: train
+    * DatasetSplit.EVAL: dev
+    * DatasetSplit.TEST: test
+    * tf.estimator.ModeKeys.PREDICT: dev
+
+    Args:
+      data_dir: str, data directory.
+      mode: DatasetSplit
+      shard: int, if provided, will only read data from the specified shard.
+
+    Returns:
+      filepattern str
+    """
+    shard_str = "-%05d" % shard if shard is not None else ""
+    if mode == problem.DatasetSplit.TRAIN:
+      path = os.path.join(data_dir, "common_voice")
+      suffix = "train"
+    elif mode in [problem.DatasetSplit.EVAL, tf.estimator.ModeKeys.PREDICT]:
+      path = os.path.join(data_dir, "common_voice_clean")
+      suffix = "dev"
+    else:
+      assert mode == problem.DatasetSplit.TEST
+      path = os.path.join(data_dir, "common_voice_clean")
+      suffix = "test"
+
+    return "%s-%s%s*" % (path, suffix, shard_str)
+
+
+@registry.register_problem()
+class CommonVoiceClean(CommonVoice):
+  """Problem spec for Common Voice using clean train and clean eval data."""
+
+  # Select only the "clean" data (crowdsourced quality control).
+  TRAIN_DATASETS = _COMMONVOICE_TRAIN_DATASETS[:1]
+  DEV_DATASETS = _COMMONVOICE_DEV_DATASETS[:1]
+  TEST_DATASETS = _COMMONVOICE_TEST_DATASETS[:1]
+
+
+@registry.register_problem()
+class CommonVoiceNoisy(CommonVoice):
+  """Problem spec for Common Voice using noisy train and noisy eval data."""
+
+  # Select only the "other" data.
+  TRAIN_DATASETS = _COMMONVOICE_TRAIN_DATASETS[1:]
+  DEV_DATASETS = _COMMONVOICE_DEV_DATASETS[1:]
+  TEST_DATASETS = _COMMONVOICE_TEST_DATASETS[1:]
+
+
+def set_common_voice_length_hparams(hparams):
+  hparams.max_length = 1650 * 80
+  hparams.max_input_seq_length = 1650
+  hparams.max_target_seq_length = 350
+  return hparams
diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 9eed9fe50..414d29333 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -179,7 +179,14 @@ def encode(self, s):
     # Make sure that the data is a single channel, 16bit, 16kHz wave.
     # TODO(chorowski): the directory may not be writable, this should fallback
     # to a temp path, and provide instructions for installing sox.
-    if not s.endswith(".wav"):
+    if s.endswith(".mp3"):
+      # TODO(dliebling) On Linux, check if libsox-fmt-mp3 is installed.
+      out_filepath = s[:-4] + ".wav"
+      call([
+          "sox", "--guard", s, "-r", "16k", "-b", "16", "-c", "1", out_filepath
+      ])
+      s = out_filepath
+    elif not s.endswith(".wav"):
       out_filepath = s + ".wav"
       if not os.path.exists(out_filepath):
         call(["sox", "-r", "16k", "-b", "16", "-c", "1", s, out_filepath])
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index d616dbfbd..2db47f24d 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -2136,6 +2136,20 @@ def transformer_librispeech_tpu():
   return transformer_librispeech_tpu_v2()
 
 
+@registry.register_hparams
+def transformer_common_voice():
+  """HParams for training ASR model on Mozilla Common Voice."""
+  return transformer_librispeech()
+
+
+@registry.register_hparams
+def transformer_common_voice_tpu():
+  """HParams for training ASR model on Mozilla Common Voice on TPU."""
+  hparams = transformer_librispeech_tpu()
+  hparams.batch_size = 8
+  return hparams
+
+
 @registry.register_hparams
 def transformer_supervised_attention():
   """HParams for supervised attention problems."""

From 41f79fc74ce38db60908985ce379be6bf13b3354 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 22 Jun 2018 17:46:58 -0700
Subject: [PATCH 0200/2720] internal merge of PR #885

PiperOrigin-RevId: 201767588
---
 tensor2tensor/data_generators/speech_recognition.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index a5c87af98..414d29333 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -317,10 +317,10 @@ def preprocess_example(self, example, mode, hparams):
       assert fbank_size[0] == 1
 
       # This replaces CMVN estimation on data
-      var_epsilon = 1e-09
+
       mean = tf.reduce_mean(mel_fbanks, keepdims=True, axis=1)
-      variance = tf.reduce_mean(tf.square(mel_fbanks - mean), keepdims=True, axis=1)
-      mel_fbanks = (mel_fbanks - mean) * tf.rsqrt(variance + var_epsilon)
+      variance = tf.reduce_mean((mel_fbanks-mean)**2, keepdims=True, axis=1)
+      mel_fbanks = (mel_fbanks - mean) / variance
 
       # Later models like to flatten the two spatial dims. Instead, we add a
       # unit spatial dim and flatten the frequencies and channels.
@@ -384,14 +384,13 @@ def bottom(self, x):
               nonpadding_mask) * num_mel_bins * num_channels
 
           # This replaces CMVN estimation on data
-          var_epsilon = 1e-09
           mean = tf.reduce_sum(
               x, axis=[1], keepdims=True) / num_of_nonpadding_elements
           variance = (num_of_nonpadding_elements * mean**2. -
                       2. * mean * tf.reduce_sum(x, axis=[1], keepdims=True) +
                       tf.reduce_sum(x**2, axis=[1], keepdims=True)
                      ) / num_of_nonpadding_elements
-          x = (x - mean) * tf.rsqrt(variance + var_epsilon) * tf.expand_dims(nonpadding_mask, -1)
+          x = (x - mean) / variance * tf.expand_dims(nonpadding_mask, -1)
       else:
         x = inputs
 

From 8b61a6e62d786658c57c8a4ddde11890dcc401be Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 22 Jun 2018 18:24:31 -0700
Subject: [PATCH 0201/2720] internal merge of PR #877

PiperOrigin-RevId: 201770879
---
 tensor2tensor/data_generators/speech_recognition.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 414d29333..9c0726a86 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -317,10 +317,11 @@ def preprocess_example(self, example, mode, hparams):
       assert fbank_size[0] == 1
 
       # This replaces CMVN estimation on data
-
+      var_epsilon = 1e-09
       mean = tf.reduce_mean(mel_fbanks, keepdims=True, axis=1)
-      variance = tf.reduce_mean((mel_fbanks-mean)**2, keepdims=True, axis=1)
-      mel_fbanks = (mel_fbanks - mean) / variance
+      variance = tf.reduce_mean(tf.square(mel_fbanks - mean),
+                                keepdims=True, axis=1)
+      mel_fbanks = (mel_fbanks - mean) * tf.rsqrt(variance + var_epsilon)
 
       # Later models like to flatten the two spatial dims. Instead, we add a
       # unit spatial dim and flatten the frequencies and channels.
@@ -384,13 +385,15 @@ def bottom(self, x):
               nonpadding_mask) * num_mel_bins * num_channels
 
           # This replaces CMVN estimation on data
+          var_epsilon = 1e-09
           mean = tf.reduce_sum(
               x, axis=[1], keepdims=True) / num_of_nonpadding_elements
           variance = (num_of_nonpadding_elements * mean**2. -
                       2. * mean * tf.reduce_sum(x, axis=[1], keepdims=True) +
                       tf.reduce_sum(x**2, axis=[1], keepdims=True)
                      ) / num_of_nonpadding_elements
-          x = (x - mean) / variance * tf.expand_dims(nonpadding_mask, -1)
+          x = (x - mean) * tf.rsqrt(variance + var_epsilon) * tf.expand_dims(
+              nonpadding_mask, -1)
       else:
         x = inputs
 

From 82b8d07ecf26c4eae32beda418a540b5d4586f15 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 22 Jun 2018 18:34:21 -0700
Subject: [PATCH 0202/2720] Added Gumbel-Softmax discrete bottleneck for latent
 transformer

PiperOrigin-RevId: 201771733
---
 tensor2tensor/layers/discretization.py      | 124 +++++++++++++++++++-
 tensor2tensor/layers/discretization_test.py |  10 ++
 2 files changed, 133 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 862187c28..1fd2e7fba 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -766,6 +766,122 @@ def vq_discrete_unbottleneck(x, hidden_size):
   return tf.reshape(result, x_shape[:-1] + [hidden_size])
 
 
+def gumbel_softmax_discrete_bottleneck(x,
+                                       bottleneck_bits,
+                                       beta=0.25,
+                                       decay=0.999,
+                                       epsilon=1e-5,
+                                       startup_steps=15000,
+                                       hard=False,
+                                       summary=True):
+  """VQ-VAE using Gumbel-Softmax.
+
+  Different from `gumbel_softmax()` function as
+  this function calculates the KL by using the discrete entropy
+  instead of taking the argmax, and it also uses an exponential moving average
+  to update the codebook while the `gumbel_softmax()` function includes no
+  codebook update.
+
+  Args:
+    x: A `float`-like `Tensor` containing the latent vectors to be compared to
+      the codebook, whose squared difference is used as the Gumbel-Softmax
+      logits.
+    bottleneck_bits: An `int` that sets the size of the bottleneck in `log_2`.
+    beta: Beta factor for commitment loss (Default: 0.25).
+    decay: Decay factor for exponential moving average (Default: 0.999).
+    epsilon: Small value to avoid dividing by zero in EMA update
+      (Default: 1e-5).
+    startup_steps: Number of steps for KL warmup (Default: 25000).
+    hard: When `True`, we use hard Gumbel-Softmax samples and force
+      discrete latents by taking the argmax. When `False`, we use soft samples,
+      which we treat as codebook weights (Default: False).
+    summary: When `True`, we save histogram summaries of the KL term (Default:
+      True).
+
+  Returns:
+    x_means_assignments: A `float`-like `Tensor` containing the codebook
+      assignments. When `hard == True`, this is one-hot, containing the arg-max
+      of the Gumbel-Softmax samples (and we use the straightthrough gradient).
+      Otherwise, it contains the Gumbel-Softmax samples exactly, which are
+      values from the `(K-1)`-simplex where `K` is the bottleneck size.
+    loss: The loss, which is the sum of the KL between the Gumbel-Softmax and
+      the uniform prior and the commitment loss multiplied by the beta factor.
+      We approximate the KL by using the entropy of a categorical distribution
+      instead of the Gumbel Softmax.
+
+  """
+  bottleneck_size = 2**bottleneck_bits
+  x_shape = common_layers.shape_list(x)
+  hidden_size = x_shape[-1]
+  means, ema_means, ema_count = get_vq_bottleneck(bottleneck_size, hidden_size)
+  x = tf.reshape(x, [-1, hidden_size])
+
+  bottleneck_size = common_layers.shape_list(means)[0]
+  x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keepdims=True)
+  means_norm_sq = tf.reduce_sum(tf.square(means), axis=-1, keepdims=True)
+  scalar_prod = tf.matmul(x, means, transpose_b=True)
+  dist = x_norm_sq + tf.transpose(means_norm_sq) - 2 * scalar_prod
+
+  class_probs = tf.nn.softmax(dist)
+  log_class_probs = tf.nn.log_softmax(dist)
+  gumbel_samples = gumbel_sample(common_layers.shape_list(dist))
+  gumbel_samples *= common_layers.inverse_exp_decay(startup_steps // 5) * 0.5
+  temperature = 1.2 - common_layers.inverse_lin_decay(startup_steps)
+
+  # 10% of the time keep reasonably high temperature to keep learning.
+  temperature = tf.cond(
+      tf.less(tf.random_uniform([]), 0.9), lambda: temperature,
+      lambda: tf.random_uniform([], minval=0.5, maxval=1.0))
+  gumbel_softmax_samples = tf.nn.softmax(
+      (log_class_probs + gumbel_samples) / temperature)
+
+  # Calculate KL between q and a uniform prior.
+  kl = tf.reduce_sum(class_probs * (log_class_probs -
+                                    tf.log(1.0/bottleneck_size)), -1)
+  if summary:
+    tf.summary.histogram("KL", tf.reshape(kl, [-1]))
+
+  # Straight-through gradient estimation when we're using hard assignments.
+  if hard:
+    x_means_idx = tf.reshape(tf.argmax(gumbel_softmax_samples, axis=-1), [-1])
+    x_means_hot = tf.one_hot(x_means_idx, bottleneck_size)
+    x_means_assignments = gumbel_softmax_samples + tf.stop_gradient(
+        x_means_hot - gumbel_softmax_samples)
+  else:
+    x_means_assignments = gumbel_softmax_samples
+  x_means_assignments_flat = tf.reshape(
+      x_means_assignments, [-1, bottleneck_size])
+  x_means = tf.matmul(x_means_assignments_flat, means)
+  commitment_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
+
+  # Update the ema variables.
+  updated_ema_count = moving_averages.assign_moving_average(
+      ema_count,
+      tf.reduce_sum(
+          tf.reshape(x_means_assignments, shape=[-1, bottleneck_size]), axis=0),
+      decay,
+      zero_debias=False)
+
+  dw = tf.matmul(x_means_assignments, x, transpose_a=True)
+  updated_ema_means = tf.identity(moving_averages.assign_moving_average(
+      ema_means, dw, decay, zero_debias=False))
+  n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
+  updated_ema_count = (
+      (updated_ema_count + epsilon) / (n + bottleneck_size * epsilon) * n)
+  updated_ema_means /= tf.expand_dims(updated_ema_count, axis=-1)
+  with tf.control_dependencies([commitment_loss]):
+    update_means = means.assign(updated_ema_means)
+    with tf.control_dependencies([update_means]):
+      loss = beta * commitment_loss
+
+  # Add KL loss.
+  loss += tf.reduce_mean(kl)
+
+  x_means_assignments = tf.reshape(
+      x_means_assignments, x_shape[:-1] + [bottleneck_size])
+  return x_means_assignments, loss
+
+
 def tanh_discrete_bottleneck(x, bottleneck_bits, bottleneck_noise,
                              discretize_warmup_steps, mode):
   """Simple discretization through tanh, flip bottleneck_noise many bits."""
@@ -844,6 +960,12 @@ def parametrized_bottleneck(x, hparams):
         hparams.vq_epsilon,
         soft_em=True,
         num_samples=hparams.vq_num_samples)
+  if hparams.bottleneck_kind == "gumbel_softmax":
+    return gumbel_softmax_discrete_bottleneck(x, hparams.bottleneck_bits,
+                                              hparams.vq_beta, hparams.vq_decay,
+                                              hparams.vq_epsilon,
+                                              hparams.startup_steps, hard=False,
+                                              summary=True)
 
   raise ValueError("Unsupported hparams.bottleneck_kind %s"
                    % hparams.bottleneck_kind)
@@ -856,7 +978,7 @@ def parametrized_unbottleneck(x, hidden_size, hparams):
   if hparams.bottleneck_kind == "isemhash":
     return isemhash_unbottleneck(
         x, hidden_size, hparams.isemhash_filter_size_multiplier)
-  if hparams.bottleneck_kind in ["vq", "em"]:
+  if hparams.bottleneck_kind in ["vq", "em", "gumbel_softmax"]:
     return vq_discrete_unbottleneck(x, hidden_size)
   raise ValueError("Unsupported hparams.bottleneck_kind %s"
                    % hparams.bottleneck_kind)
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index 8ad25a362..b017b6582 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -163,6 +163,16 @@ def testVQDiscreteUnbottlenck(self):
       x_means_eval = sess.run(x_means)
       self.assertEqual(np.shape(x_means_eval), (2, 3))
 
+  def testGumbleSoftmaxDiscreteBottleneck(self):
+    x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
+    tf.add_to_collection(tf.GraphKeys.GLOBAL_STEP, tf.constant(1))
+    x_means_hot, _ = discretization.gumbel_softmax_discrete_bottleneck(
+        x, bottleneck_bits=2)
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      x_means_hot_eval = sess.run(x_means_hot)
+      self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
+
 
 if __name__ == '__main__':
   tf.test.main()

From 36d322b67a3d75383df9878f9508f9253eb850c1 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sun, 24 Jun 2018 22:46:11 +0200
Subject: [PATCH 0203/2720] removing redundant parameter

---
 tensor2tensor/rl/collect.py        | 5 +++--
 tensor2tensor/rl/rl_trainer_lib.py | 7 ++++---
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index ba2a67e8f..dd65e8e47 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -78,8 +78,7 @@ def simulate(self, action):
 
 def define_collect(hparams, scope, eval_phase,
                    collect_level=-1,
-                   policy_to_actions_lambda=None,
-                   on_simulated=False):
+                   policy_to_actions_lambda=None):
   """Collect trajectories."""
 
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
@@ -100,6 +99,8 @@ def define_collect(hparams, scope, eval_phase,
         speculum = batch_env.speculum
 
     eval_phase = tf.convert_to_tensor(eval_phase)
+    on_simulated = hparams.environment_spec.simulated_env
+
     on_simulated = tf.convert_to_tensor(on_simulated)
 
     memory = [tf.get_variable("collect_memory_{}".format(name),
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 46ac189c0..a4b516c24 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -33,8 +33,7 @@ def define_train(hparams, event_dir):
   del event_dir
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
     memory, collect_summary = collect.define_collect(
-        hparams, "ppo_train", eval_phase=False,
-        on_simulated=hparams.simulated_environment)
+        hparams, "ppo_train", eval_phase=False)
     ppo_summary = ppo.define_ppo_epoch(memory, hparams)
     summary = tf.summary.merge([collect_summary, ppo_summary])
 
@@ -56,7 +55,9 @@ def train(hparams, event_dir=None, model_dir=None,
       summary_writer = None
       model_saver = None
 
-    if hparams.simulated_environment:
+    # TODO (piotr milos): This should be refactored, possibly with
+    # handlers for each type of env
+    if hparams.environment_spec.simulated_env:
       env_model_loader = tf.train.Saver(
           tf.global_variables("next_frame*"))
     else:

From a794e329ccf70f8c3b87fe91d484d53aa94db540 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sun, 24 Jun 2018 22:47:42 +0200
Subject: [PATCH 0204/2720] paremters clean-up

---
 tensor2tensor/rl/envs/simulated_batch_env.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 78a14d553..757a7fec1 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -99,7 +99,7 @@ def __init__(self, hparams, length,
     environment_spec = hparams.environment_spec
     initial_frames_problem = environment_spec.initial_frames_problem
     self._min_reward = initial_frames_problem.min_reward
-    self._num_frames = hparams.model_hparams.video_num_input_frames
+    self._num_frames = environment_spec.video_num_input_frames
     self._intrinsic_reward_scale = intrinsic_reward_scale
 
     # initialization_env = environment_lambda()
@@ -113,8 +113,8 @@ def __init__(self, hparams, length,
 
     # TODO(lukaszkaiser): do this in a more cleaner way
     hparams.video_num_input_frames, hparams.video_num_target_frames = (
-        hparams.model_hparams.video_num_input_frames,
-        hparams.model_hparams.video_num_target_frames)
+        hparams.environment_spec.video_num_input_frames,
+        hparams.environment_spec.video_num_target_frames)
 
     if simulation_random_starts:
       dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,

From 10393415ccda1da246062a2d85fc9f67d409b584 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sun, 24 Jun 2018 22:48:06 +0200
Subject: [PATCH 0205/2720] training test added

---
 tensor2tensor/rl/rl_trainer_lib_test.py | 35 ++++++++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index a919d88ae..e0bee9678 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -18,7 +18,7 @@
 from __future__ import print_function
 
 from tensor2tensor.data_generators.gym_problems import standard_atari_env_spec
-from tensor2tensor.models.research.rl import simple_gym_spec
+from tensor2tensor.models.research.rl import simple_gym_spec, feed_forward_cnn_small_categorical_fun
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.utils import registry  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
@@ -48,6 +48,39 @@ def test_no_crash_cartpole(self):
                        standard_atari_env_spec("CartPole-v0"))
     rl_trainer_lib.train(hparams)
 
+  # This test should sucessfully train pong.
+  # It should get train mean_score around 0 after 100 epoch
+  #
+  # This test should be run whenever ppo any bigger change
+  # is done on the ppo code
+  #
+  # To run the test change epochs_num=2 to epoch_num=200
+  # (it is set like that to meet travis timeouts
+  def test_train_pong(self):
+    hparams = tf.contrib.training.\
+      HParams(epochs_num=2,
+              eval_every_epochs=10,
+              num_agents=20,
+              optimization_epochs=3,
+              epoch_length=200,
+              entropy_loss_coef=0.003,
+              learning_rate=8e-05,
+              optimizer="Adam",
+              policy_network=feed_forward_cnn_small_categorical_fun,
+              gae_lambda=0.985,
+              num_eval_agents=1,
+              max_gradients_norm=0.5,
+              gae_gamma=0.985,
+              optimization_batch_size=4,
+              clipping_coef=0.2,
+              value_loss_coef=1,
+              save_models_every_epochs=False)
+
+    hparams.add_hparam("environment_spec",
+                       standard_atari_env_spec("PongNoFrameskip-v4"))
+    rl_trainer_lib.train(hparams)
+
+
 
 if __name__ == "__main__":
   tf.test.main()

From 8099236cae27f414550cab0c5ec984e00c2a5548 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sun, 24 Jun 2018 23:36:24 +0200
Subject: [PATCH 0206/2720] Reward per sequence reintroduced

---
 tensor2tensor/data_generators/gym_problems.py | 160 ++++++++++++++----
 tensor2tensor/rl/model_rl_experiment.py       |  20 +--
 2 files changed, 136 insertions(+), 44 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index f6558cf6a..9d0518024 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -73,21 +73,19 @@ def __init__(self, *args, **kwargs):
 
     self.environment_spec = self.get_environment_spec()
     self.eval_phase = False
-    self.sum_of_rewards = 0.0
-    self.dones = 0
 
   def _setup(self):
+    self._internal_memory_size = 10
+
     collect_hparams = rl.ppo_pong_base()
     collect_hparams.add_hparam("environment_spec", self.environment_spec)
+    collect_hparams.epoch_length = self._internal_memory_size
+    collect_hparams.num_agents = 1
 
     if not FLAGS.agent_policy_path:
       collect_hparams.policy_network = rl.random_policy_fun
 
-    self._internal_memory_size = 10
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      collect_hparams.epoch_length = self._internal_memory_size
-      # TODO(piotrmilos). it is possible to set more than 1.
-      collect_hparams.num_agents = 1
       self.collect_memory, self.collect_trigger_op \
         = collect.define_collect(collect_hparams, scope="gym_problems",
                                  collect_level=0, eval_phase=self.eval_phase)
@@ -110,13 +108,12 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
           memory_index = 0
         data = [memory[i][memory_index][0] for i in range(4)]
         memory_index += 1
-        observ, reward, done, action = data
-        observ = observ.astype(np.uint8)
+        observation, reward, done, action = data
+        observation = observation.astype(np.uint8)
 
-        self.sum_of_rewards += reward
-        self.dones += int(done)
+        debug_image = self.collect_statistics_and_generate_debug_image(*data)
 
-        ret_dict = {"frame": observ,
+        ret_dict = {"frame": observation,
                     "image/format": ["png"],
                     "image/height": [self.frame_height],
                     "image/width": [self.frame_width],
@@ -124,6 +121,9 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
                     "done": [int(False)],
                     "reward": [int(reward) - self.min_reward]}
 
+        if debug_image is not None:
+          ret_dict["image/debug"] = debug_image
+
         yield ret_dict
         pieces_generated += 1
 
@@ -176,6 +176,11 @@ def env(self):
   def num_actions(self):
     return self.env.action_space.n
 
+  def collect_statistics_and_generate_debug_image(self, observation,
+                                                  reward, done, action):
+    """This generates extra statistics and debug images"""
+    raise NotImplementedError()
+
   @property
   def frame_height(self):
     return self.env.observation_space.shape[0]
@@ -228,12 +233,53 @@ class GymAEDiscreteProblem(GymDiscreteProblem):
   pass
 
 
-class GymRealDiscreteProblem(GymDiscreteProblem):
+class BasicStatistics:
+  """Keeps basic statistics to calculate mean reward """
+
+  def __init__(self):
+    self.sum_of_rewards = 0.0
+    self.number_of_dones = 0
 
+
+
+class GymRealDiscreteProblem(GymDiscreteProblem):
   def __init__(self, *args, **kwargs):
     super(GymRealDiscreteProblem, self).__init__(*args, **kwargs)
+    self.statistics = BasicStatistics()
+
     self.make_extra_debug_info = False
 
+  def collect_statistics_and_generate_debug_image(self, observation,
+                                                  reward, done, action):
+    """Collects info required to calculate mean reward."""
+
+    self.statistics.sum_of_rewards += reward
+    self.statistics.number_of_dones += int(done)
+
+    debug_image = None
+
+    return debug_image
+
+
+class RewardPerSequenceStatistics(BasicStatistics):
+  """This encapsulates all pieces required to calculate
+  the correctness of rewards per sequence metric
+  """
+
+  def __init__(self):
+    super(RewardPerSequenceStatistics, self).__init__()
+
+    # data to calculate
+    # correctness of rewards per sequence metric
+    self.episode_sim_reward = 0.0
+    self.episode_real_reward = 0.0,
+    self.successful_episode_reward_predictions = 0
+    self.report_reward_statistics_every = 10
+
+    # auxiliary objects
+    self.real_env = None
+    self.real_ob = None
+
 
 class GymSimulatedDiscreteProblem(GymDiscreteProblem):
   """Simulated gym environment with discrete actions and rewards."""
@@ -243,17 +289,37 @@ def __init__(self, *args, **kwargs):
     self.debug_dump_frames_path = "debug_frames_sim"
     self.intrinsic_reward_scale = 0.0
     self.simulation_random_starts = False
+    self.statistics = RewardPerSequenceStatistics()
     super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
 
+  def _setup(self):
+    super(GymSimulatedDiscreteProblem, self)._setup()
+    self._reset_real_env()
+
   @property
   def initial_frames_problem(self):
     raise NotImplementedError()
 
+  @property
+  def num_input_frames(self):
+    """Number of frames on input for real environment."""
+    # TODO(lukaszkaiser): This must be equal to hparams.video_num_input_frames,
+    # we should automate this to avoid bug in the future.
+    return 4
+
+  @property
+  def video_num_target_frames(self):
+    """Number of frames on input for real environment."""
+    # TODO(piotrmilos): This must be equal to hparams.video_num_target_frames,
+    # we should automate this to avoid bug in the future.
+    return 1
+
   def get_environment_spec(self):
     env_spec = standard_atari_env_spec(self.env_name)
 
     # Set reasonable time limit (as we do not simulate done).
     real_env = env_spec.env_lambda()
+    self.statistics.real_env = real_env
     if self.num_testing_steps is not None:
       timelimit = self.num_testing_steps
     else:
@@ -273,10 +339,45 @@ def get_environment_spec(self):
     initial_frames_problem = registry.problem(self.initial_frames_problem)
     env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
     env_spec.wrappers.append(
-        [tf_atari_wrappers.TimeLimitWrapper, {"timelimit": timelimit}])
+      [tf_atari_wrappers.TimeLimitWrapper, {"timelimit": timelimit}])
+    env_spec.add_hparam("video_num_input_frames", self.num_input_frames)
+    env_spec.add_hparam("video_num_target_frames", self.video_num_target_frames)
 
     return env_spec
 
+  def _reset_real_env(self):
+    stat = self.statistics
+    stat.real_env.reset()
+    for _ in range(self.num_input_frames):
+      stat.real_ob, _, _, _ = stat.real_env.step(0)
+
+    stat.episode_sim_reward = 0.0
+    stat.episode_real_reward = 0.0
+
+  def collect_statistics_and_generate_debug_image(self, observation,
+                                                  reward, done, action):
+    stat = self.statistics
+
+    stat.sum_of_rewards += reward
+    stat.number_of_dones += int(done)
+    stat.episode_sim_reward += reward
+
+    ob = np.ndarray.astype(observation, np.int)
+    err = np.ndarray.astype(np.maximum(np.abs(
+      stat.real_ob - ob, dtype=np.int) - 10, 0),
+                            np.uint8)
+    debug_im = np.concatenate([observation, stat.real_ob, err], axis=1)
+
+    if done:
+      if stat.episode_sim_reward == stat.episode_real_reward:
+        stat.successful_episode_reward_predictions += 1
+      self._reset_real_env()
+    else:
+      stat.real_ob, real_reward, _, _ = stat.real_env.step(action)
+      stat.episode_real_reward += real_reward
+
+    return debug_im
+
   def restore_networks(self, sess):
     super(GymSimulatedDiscreteProblem, self).restore_networks(sess)
     # TODO(blazej): adjust regexp for different models.
@@ -376,8 +477,7 @@ def num_rewards(self):
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnPong(
-    GymSimulatedDiscreteProblem, GymPongRandom):
-
+  GymSimulatedDiscreteProblem, GymPongRandom):
   @property
   def initial_frames_problem(self):
     return "gym_discrete_problem_with_agent_on_pong"
@@ -402,14 +502,13 @@ def num_rewards(self):
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnPong(
-    GymRealDiscreteProblem, GymPongRandom):
+  GymRealDiscreteProblem, GymPongRandom):
   pass
 
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
-    GymSimulatedDiscreteProblem, GymWrappedPongRandom):
-
+  GymSimulatedDiscreteProblem, GymWrappedPongRandom):
   @property
   def initial_frames_problem(self):
     return "gym_discrete_problem_with_agent_on_wrapped_pong"
@@ -417,20 +516,19 @@ def initial_frames_problem(self):
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedLongPong(
-    GymRealDiscreteProblem, GymWrappedLongPongRandom):
+  GymRealDiscreteProblem, GymWrappedLongPongRandom):
   pass
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedLongPongAe(  # with autoencoder
-    GymDiscreteProblemWithAgentOnWrappedLongPong):
+  GymDiscreteProblemWithAgentOnWrappedLongPong):
   pass
 
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPong(
-    GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
-
+  GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
   @property
   def initial_frames_problem(self):
     return "gym_discrete_problem_with_agent_on_wrapped_long_pong"
@@ -438,20 +536,19 @@ def initial_frames_problem(self):
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedBreakout(
-    GymRealDiscreteProblem, GymWrappedBreakoutRandom):
+  GymRealDiscreteProblem, GymWrappedBreakoutRandom):
   pass
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
-    GymDiscreteProblemWithAgentOnWrappedBreakout):
+  GymDiscreteProblemWithAgentOnWrappedBreakout):
   pass
 
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
-    GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
-
+  GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
   @property
   def initial_frames_problem(self):
     return "gym_discrete_problem_with_agent_on_wrapped_breakout"
@@ -459,7 +556,7 @@ def initial_frames_problem(self):
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedPong(
-    GymRealDiscreteProblem, GymWrappedPongRandom):
+  GymRealDiscreteProblem, GymWrappedPongRandom):
   """GymDiscreteProblemWithAgentOnWrappedPong."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
@@ -483,14 +580,13 @@ def frame_width(self):
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
-    GymDiscreteProblemWithAgentOnWrappedPong):
+  GymDiscreteProblemWithAgentOnWrappedPong):
   pass
 
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnFreeway(
-    GymSimulatedDiscreteProblem, GymFreewayRandom):
-
+  GymSimulatedDiscreteProblem, GymFreewayRandom):
   @property
   def initial_frames_problem(self):
     return "gym_discrete_problem_with_agent_on_freeway"
@@ -498,7 +594,7 @@ def initial_frames_problem(self):
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnFreeway(
-    GymRealDiscreteProblem, GymFreewayRandom):
+  GymRealDiscreteProblem, GymFreewayRandom):
   """Freeway with agent."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
@@ -522,5 +618,5 @@ def frame_width(self):
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnFreewayAe(  # with autoencoder
-    GymDiscreteProblemWithAgentOnFreeway):
+  GymDiscreteProblemWithAgentOnFreeway):
   pass
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index bba9e8da6..0b1bb420a 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -102,7 +102,8 @@ def generate_real_env_data(problem_name, agent_policy_path, hparams, data_dir,
     gym_problem.settable_num_steps = hparams.true_env_generator_num_steps
     gym_problem.eval_phase = eval_phase
     gym_problem.generate_data(data_dir, tmp_dir)
-    mean_reward = gym_problem.sum_of_rewards / (1.0 + gym_problem.dones)
+    mean_reward = gym_problem.statistics.sum_of_rewards / \
+                  (1.0 + gym_problem.statistics.number_of_dones)
 
   return mean_reward
 
@@ -153,13 +154,8 @@ def train_agent(problem_name, agent_model_dir,
   ppo_hparams.add_hparam("model_hparams", model_hparams)
 
   environment_spec = copy.copy(gym_problem.environment_spec)
-  environment_spec.simulated_env = True
-  environment_spec.add_hparam("simulation_random_starts",
-                              hparams.simulation_random_starts)
-  environment_spec.add_hparam("intrinsic_reward_scale",
-                              hparams.intrinsic_reward_scale)
-  environment_spec.add_hparam("initial_frames_problem",
-                              gym_problem)
+  environment_spec.simulation_random_starts = hparams.simulation_random_starts
+  environment_spec.intrinsic_reward_scale = hparams.intrinsic_reward_scale
 
   # 4x for the StackAndSkipWrapper minus one to always finish for reporting.
   ppo_time_limit = ppo_hparams.epoch_length - 1
@@ -198,9 +194,9 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
       "autoencoder_path": autoencoder_path,
   }):
     gym_simulated_problem.generate_data(epoch_data_dir, tmp_dir)
-  n = max(1., gym_simulated_problem.dones)
+  n = max(1., gym_simulated_problem.statistics.number_of_dones)
   model_reward_accuracy = (
-      gym_simulated_problem.successful_episode_reward_predictions / float(n))
+      gym_simulated_problem.statistics.successful_episode_reward_predictions / float(n))
   return model_reward_accuracy
 
 
@@ -399,7 +395,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     if hparams.eval_world_model:
       log("Evaluating world model")
       model_reward_accuracy = evaluate_world_model(
-          simulated_problem_name, world_model_problem, hparams,
+        simulated_problem_name, world_model_problem, hparams,
           directories["world_model"],
           epoch_data_dir, directories["tmp"],
           autoencoder_path=autoencoder_model_dir)
@@ -411,7 +407,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     ppo_model_dir = directories["ppo"]
     if not hparams.ppo_continue_training:
       ppo_model_dir = ppo_event_dir
-    train_agent(world_model_problem, ppo_model_dir,
+    train_agent(simulated_problem_name, ppo_model_dir,
                 ppo_event_dir, directories["world_model"], epoch_data_dir,
                 hparams, autoencoder_path=autoencoder_model_dir, epoch=epoch)
 

From 951ea38e5fb5f33997bca8347a8d9e8bca2d9111 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 25 Jun 2018 10:16:49 +0200
Subject: [PATCH 0207/2720] renaming

---
 tensor2tensor/rl/model_rl_experiment.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 0b1bb420a..eb0635723 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -489,7 +489,7 @@ def rl_modelrl_base():
       game="wrapped_long_pong",
       # Whether to evaluate the world model in each iteration of the loop to get
       # the model_reward_accuracy metric.
-      eval_world_model=False,
+      eval_world_model=True,
   )
 
 

From 42bdbd43eca9b41ceec6981f4e4d88d3e359a3a7 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 25 Jun 2018 17:14:07 +0200
Subject: [PATCH 0208/2720] support for initialization

---
 tensor2tensor/rl/collect.py                | 30 +++++++++++++++-------
 tensor2tensor/rl/envs/py_func_batch_env.py |  3 +++
 tensor2tensor/rl/envs/tf_atari_wrappers.py |  8 ++++++
 3 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index dd65e8e47..6161fb2a2 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -79,10 +79,17 @@ def simulate(self, action):
 def define_collect(hparams, scope, eval_phase,
                    collect_level=-1,
                    policy_to_actions_lambda=None):
-  """Collect trajectories."""
-
+  """Collect trajectories.
+  Returns: memory - tensor with collected rollout
+           summaries - basic statistcs about the rollout
+           initialization_lambda - initializations to be done once 
+            tf.Session is created
+  """
+
+  to_initialize = []
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
     batch_env = batch_env_factory(hparams)
+    to_initialize.append(batch_env)
     environment_wrappers = hparams.environment_spec.wrappers
     wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
     # Put memory wrapper at the level you want to gather observations at.
@@ -94,14 +101,14 @@ def define_collect(hparams, scope, eval_phase,
     speculum = None
     for w in wrappers:
       batch_env = w[0](batch_env, **w[1])
+      to_initialize.append(batch_env)
       if w[0] == _MemoryWrapper:
         rollout_metadata = _rollout_metadata(batch_env)
         speculum = batch_env.speculum
 
-    eval_phase = tf.convert_to_tensor(eval_phase)
-    on_simulated = hparams.environment_spec.simulated_env
-
-    on_simulated = tf.convert_to_tensor(on_simulated)
+    def initialization_lambda(sess):
+      for batch_env in to_initialize:
+        batch_env.initialize(sess)
 
     memory = [tf.get_variable("collect_memory_{}".format(name),
                               shape=[hparams.epoch_length]+shape,
@@ -113,15 +120,20 @@ def define_collect(hparams, scope, eval_phase,
     cumulative_rewards = tf.get_variable("cumulative_rewards", len(batch_env),
                                          trainable=False)
 
+    eval_phase = tf.convert_to_tensor(eval_phase)
     should_reset_var = tf.Variable(True, trainable=False)
-
     zeros_tensor = tf.zeros(len(batch_env))
 
+  if "force_beginning_resets" in hparams:
+    force_beginning_resets = hparams.force_beginning_resets
+  else:
+    force_beginning_resets = False
+
   def group():
     return tf.group(batch_env.reset(tf.range(len(batch_env))),
                     tf.assign(cumulative_rewards, zeros_tensor))
   reset_op = tf.cond(
-      tf.logical_or(should_reset_var, tf.logical_or(eval_phase, on_simulated)),
+      tf.logical_or(should_reset_var, tf.convert_to_tensor(force_beginning_resets)),
       group, tf.no_op)
 
   with tf.control_dependencies([reset_op]):
@@ -217,4 +229,4 @@ def stop_condition(i, _, resets):
     summaries = tf.summary.merge(
         [mean_score_summary,
          tf.summary.scalar("episodes_finished_this_iter", scores_num)])
-    return memory, summaries
+    return memory, summaries, initialization_lambda
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index b5f909bdf..dc5a0fa79 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -61,6 +61,9 @@ def __getattr__(self, name):
     """
     return getattr(self._batch_env, name)
 
+  def initialize(self, sess):
+    pass
+
   def __len__(self):
     """Number of combined environments."""
     return len(self._batch_env)
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 9a1708cfb..e966836fd 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -31,6 +31,14 @@ def __init__(self, batch_env):
     self.action_shape = batch_env.action_shape
     self.action_dtype = batch_env.action_dtype
 
+  def initialize(self, sess):
+    """
+    Initializations to be run once the tf.Session is available 
+    
+    sess - tf.Session 
+    """
+    pass
+
   @property
   def observ(self):
     """Access the variable holding the current observation."""

From e7ca701503df1e8f76bd21399ee1a6bf033d55c1 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 25 Jun 2018 17:16:58 +0200
Subject: [PATCH 0209/2720] initalization and parameters refactor

---
 tensor2tensor/rl/envs/batch_env_factory.py   | 13 +++-------
 tensor2tensor/rl/envs/simulated_batch_env.py | 27 ++++++++++++--------
 2 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index 70bb129a0..e692ed9d3 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -46,9 +46,7 @@ def batch_env_factory(hparams, xvfb=False):
   if environment_spec.simulated_env:
     # TODO(piotrmilos): Consider passing only relevant parameters
     cur_batch_env = _define_simulated_batch_env(
-        hparams, hparams.num_agents,
-        hparams.simulation_random_starts,
-        hparams.intrinsic_reward_scale)
+        environment_spec, hparams.num_agents, hparams)
   else:
 
     cur_batch_env = _define_batch_env(hparams.environment_spec,
@@ -69,12 +67,9 @@ def _define_batch_env(environment_spec, num_agents, xvfb=False):
     return env
 
 
-def _define_simulated_batch_env(hparams, num_agents,
-                                simulation_random_starts=False,
-                                intrinsic_reward_scale=0.):
-  cur_batch_env = simulated_batch_env.SimulatedBatchEnv(
-      hparams, num_agents, simulation_random_starts,
-      intrinsic_reward_scale)
+def _define_simulated_batch_env(environment_spec, num_agents,
+                                other_hparms):
+  cur_batch_env = simulated_batch_env.SimulatedBatchEnv(environment_spec, num_agents, other_hparms)
   return cur_batch_env
 
 
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 757a7fec1..cba46464c 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -20,6 +20,9 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+import copy
+
 from tensor2tensor.layers import common_layers
 from tensor2tensor.rl.envs import in_graph_batch_env
 from tensor2tensor.rl.envs.utils import get_action_space
@@ -38,13 +41,16 @@ class HistoryBuffer(object):
 
   def __init__(self, input_dataset, length):
     self.input_data_iterator = (
-        input_dataset.batch(length).make_one_shot_iterator())
+        input_dataset.batch(length).make_initializable_iterator())
     self.length = length
     initial_frames = self.get_initial_observations()
     initial_shape = [length] + common_layers.shape_list(initial_frames)[1:]
     self._history_buff = tf.Variable(tf.zeros(initial_shape, tf.float32),
                                      trainable=False)
 
+  def initialize(self, sess):
+    sess.run(self.input_data_iterator.initializer)
+
   def get_initial_observations(self):
     return tf.cast(self.input_data_iterator.get_next(), tf.float32)
 
@@ -91,18 +97,14 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   flags are held in according variables.
   """
 
-  def __init__(self, hparams, length,
-               simulation_random_starts=False,
-               intrinsic_reward_scale=0.):
+  def __init__(self, environment_spec, length, other_hparams):
     """Batch of environments inside the TensorFlow graph."""
     self.length = length
-    environment_spec = hparams.environment_spec
     initial_frames_problem = environment_spec.initial_frames_problem
     self._min_reward = initial_frames_problem.min_reward
     self._num_frames = environment_spec.video_num_input_frames
-    self._intrinsic_reward_scale = intrinsic_reward_scale
+    self._intrinsic_reward_scale = environment_spec.intrinsic_reward_scale
 
-    # initialization_env = environment_lambda()
     model_hparams = trainer_lib.create_hparams(
         FLAGS.hparams_set, problem_name=FLAGS.problem)
     model_hparams.force_full_predict = True
@@ -112,20 +114,22 @@ def __init__(self, hparams, length,
     _, self.action_shape, self.action_dtype = get_action_space(environment_spec)
 
     # TODO(lukaszkaiser): do this in a more cleaner way
+    # remove other_hparams
+    hparams = copy.copy(other_hparams)
     hparams.video_num_input_frames, hparams.video_num_target_frames = (
         hparams.environment_spec.video_num_input_frames,
         hparams.environment_spec.video_num_target_frames)
 
-    if simulation_random_starts:
+    if environment_spec.simulation_random_starts:
       dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
                                                FLAGS.data_dir,
                                                shuffle_files=True,
                                                hparams=hparams)
       dataset = dataset.shuffle(buffer_size=100)
     else:
-      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
+      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.EVAL,
                                                FLAGS.data_dir,
-                                               shuffle_files=True,
+                                               shuffle_files=False,
                                                hparams=hparams).take(1)
 
     dataset = dataset.map(lambda x: x["inputs"]).repeat()
@@ -136,6 +140,9 @@ def __init__(self, hparams, length,
              initial_frames_problem.num_channels)
     self._observ = tf.Variable(tf.zeros(shape, tf.float32), trainable=False)
 
+  def initialize(self, sess):
+    self.history_buffer.initialize(sess)
+
   def __len__(self):
     """Number of combined environments."""
     return self.length

From f007c1f4ba283ffde281004b5d178b2390312a05 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 25 Jun 2018 17:26:46 +0200
Subject: [PATCH 0210/2720] initialization support in rl module

---
 tensor2tensor/rl/rl_trainer_lib.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index a4b516c24..f1935998a 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -32,19 +32,20 @@ def define_train(hparams, event_dir):
   """Define the training setup."""
   del event_dir
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-    memory, collect_summary = collect.define_collect(
+    memory, collect_summary, initialization\
+      = collect.define_collect(
         hparams, "ppo_train", eval_phase=False)
     ppo_summary = ppo.define_ppo_epoch(memory, hparams)
     summary = tf.summary.merge([collect_summary, ppo_summary])
 
-  return summary, None
+  return summary, None, initialization
 
 
 def train(hparams, event_dir=None, model_dir=None,
           restore_agent=True, epoch=0):
   """Train."""
   with tf.name_scope("rl_train"):
-    train_summary_op, _ = define_train(hparams, event_dir)
+    train_summary_op, _, initialization = define_train(hparams, event_dir)
     if event_dir:
       summary_writer = tf.summary.FileWriter(
           event_dir, graph=tf.get_default_graph(), flush_secs=60)
@@ -65,6 +66,7 @@ def train(hparams, event_dir=None, model_dir=None,
 
     with tf.Session() as sess:
       sess.run(tf.global_variables_initializer())
+      initialization(sess)
       if env_model_loader:
         trainer_lib.restore_checkpoint(
             hparams.world_model_dir, env_model_loader, sess, must_restore=True)

From 0d2fde0a4e079a020e0e062b2431fff103767ae1 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 25 Jun 2018 17:27:08 +0200
Subject: [PATCH 0211/2720] evaluation code reestablished

---
 tensor2tensor/data_generators/gym_problems.py | 101 ++++++++++++------
 tensor2tensor/rl/model_rl_experiment.py       |   9 +-
 2 files changed, 71 insertions(+), 39 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 9d0518024..8c0bc75c5 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -74,11 +74,14 @@ def __init__(self, *args, **kwargs):
     self.environment_spec = self.get_environment_spec()
     self.eval_phase = False
 
-  def _setup(self):
-    self._internal_memory_size = 10
+    self._internal_memory_size = 20
+    self._internal_memory_force_beginning_resets = False
 
+  def _setup(self):
     collect_hparams = rl.ppo_pong_base()
     collect_hparams.add_hparam("environment_spec", self.environment_spec)
+    collect_hparams.add_hparam("force_beginning_resets",
+                               self._internal_memory_force_beginning_resets)
     collect_hparams.epoch_length = self._internal_memory_size
     collect_hparams.num_agents = 1
 
@@ -86,7 +89,7 @@ def _setup(self):
       collect_hparams.policy_network = rl.random_policy_fun
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      self.collect_memory, self.collect_trigger_op \
+      self.collect_memory, self.collect_trigger_op, self.collect_init \
         = collect.define_collect(collect_hparams, scope="gym_problems",
                                  collect_level=0, eval_phase=self.eval_phase)
 
@@ -97,13 +100,13 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
 
     with tf.Session() as sess:
       sess.run(tf.global_variables_initializer())
+      self.collect_init(sess)
       self.restore_networks(sess)
       pieces_generated = 0
       memory_index = 0
       memory = None
       while pieces_generated < self.num_steps:
         if memory is None or memory_index >= self._internal_memory_size:
-          sess.run(self.collect_trigger_op)
           memory = sess.run(self.collect_memory)
           memory_index = 0
         data = [memory[i][memory_index][0] for i in range(4)]
@@ -111,7 +114,8 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
         observation, reward, done, action = data
         observation = observation.astype(np.uint8)
 
-        debug_image = self.collect_statistics_and_generate_debug_image(*data)
+        debug_image = self.collect_statistics_and_generate_debug_image(pieces_generated,
+                                                                       *data)
 
         ret_dict = {"frame": observation,
                     "image/format": ["png"],
@@ -176,8 +180,11 @@ def env(self):
   def num_actions(self):
     return self.env.action_space.n
 
-  def collect_statistics_and_generate_debug_image(self, observation,
-                                                  reward, done, action):
+  def collect_statistics_and_generate_debug_image(self, index,
+                                                  observation,
+                                                  reward,
+                                                  done,
+                                                  action):
     """This generates extra statistics and debug images"""
     raise NotImplementedError()
 
@@ -205,10 +212,6 @@ def total_number_of_frames(self):
   def min_reward(self):
     raise NotImplementedError()
 
-  @property
-  def num_testing_steps(self):
-    return None
-
   def get_action(self, observation=None):
     del observation
     return self.env.action_space.sample()
@@ -249,7 +252,7 @@ def __init__(self, *args, **kwargs):
 
     self.make_extra_debug_info = False
 
-  def collect_statistics_and_generate_debug_image(self, observation,
+  def collect_statistics_and_generate_debug_image(self, index, observation,
                                                   reward, done, action):
     """Collects info required to calculate mean reward."""
 
@@ -292,6 +295,17 @@ def __init__(self, *args, **kwargs):
     self.statistics = RewardPerSequenceStatistics()
     super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
 
+    # This is hackish way of introducing resets every
+    # self.num_testing_steps. It cannot be done easily
+    # using other ways as we do not control
+    # the amount of skips induced but wrappers
+    self._internal_memory_size = self.num_testing_steps
+    self._internal_memory_force_beginning_resets = True
+    env_spec = standard_atari_env_spec(self.env_name)
+    real_env = env_spec.env_lambda()
+    self.statistics.real_env = real_env
+
+
   def _setup(self):
     super(GymSimulatedDiscreteProblem, self)._setup()
     self._reset_real_env()
@@ -314,32 +328,20 @@ def video_num_target_frames(self):
     # we should automate this to avoid bug in the future.
     return 1
 
+  @property
+  def num_testing_steps(self):
+    return None
+
   def get_environment_spec(self):
     env_spec = standard_atari_env_spec(self.env_name)
-
-    # Set reasonable time limit (as we do not simulate done).
-    real_env = env_spec.env_lambda()
-    self.statistics.real_env = real_env
-    if self.num_testing_steps is not None:
-      timelimit = self.num_testing_steps
-    else:
-      try:
-        # We assume that the real env is wrapped with TimeLimit.
-        history = self.num_input_frames
-        timelimit = real_env._max_episode_steps - history  # pylint: disable=protected-access
-      except:  # pylint: disable=bare-except
-        # If not, set some reasonable default.
-        timelimit = 100
-
     env_spec.simulated_env = True
     env_spec.add_hparam("simulation_random_starts",
                         self.simulation_random_starts)
+
     env_spec.add_hparam("intrinsic_reward_scale",
                         self.intrinsic_reward_scale)
     initial_frames_problem = registry.problem(self.initial_frames_problem)
     env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
-    env_spec.wrappers.append(
-      [tf_atari_wrappers.TimeLimitWrapper, {"timelimit": timelimit}])
     env_spec.add_hparam("video_num_input_frames", self.num_input_frames)
     env_spec.add_hparam("video_num_target_frames", self.video_num_target_frames)
 
@@ -354,7 +356,8 @@ def _reset_real_env(self):
     stat.episode_sim_reward = 0.0
     stat.episode_real_reward = 0.0
 
-  def collect_statistics_and_generate_debug_image(self, observation,
+  def collect_statistics_and_generate_debug_image(self, index,
+                                                  observation,
                                                   reward, done, action):
     stat = self.statistics
 
@@ -368,7 +371,12 @@ def collect_statistics_and_generate_debug_image(self, observation,
                             np.uint8)
     debug_im = np.concatenate([observation, stat.real_ob, err], axis=1)
 
-    if done:
+    assert self._internal_memory_size==self.num_testing_steps and \
+           self._internal_memory_force_beginning_resets, \
+      "The collect memory should be set in force_beginning_resets mode" \
+      "for the code below to work properly"
+
+    if index%self._internal_memory_size == 0:
       if stat.episode_sim_reward == stat.episode_real_reward:
         stat.successful_episode_reward_predictions += 1
       self._reset_real_env()
@@ -381,6 +389,7 @@ def collect_statistics_and_generate_debug_image(self, observation,
   def restore_networks(self, sess):
     super(GymSimulatedDiscreteProblem, self).restore_networks(sess)
     # TODO(blazej): adjust regexp for different models.
+    # TODO(piotrmilos): move restoring networks to SimulatedBatchEnv.initialize
     env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
     sess = tf.get_default_session()
 
@@ -478,10 +487,16 @@ def num_rewards(self):
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnPong(
   GymSimulatedDiscreteProblem, GymPongRandom):
+
   @property
   def initial_frames_problem(self):
     return "gym_discrete_problem_with_agent_on_pong"
 
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
 
 @registry.register_problem
 class GymFreewayRandom(GymDiscreteProblem):
@@ -509,10 +524,16 @@ class GymDiscreteProblemWithAgentOnPong(
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
   GymSimulatedDiscreteProblem, GymWrappedPongRandom):
+
   @property
   def initial_frames_problem(self):
     return "gym_discrete_problem_with_agent_on_wrapped_pong"
 
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedLongPong(
@@ -529,10 +550,16 @@ class GymDiscreteProblemWithAgentOnWrappedLongPongAe(  # with autoencoder
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPong(
   GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
+
   @property
   def initial_frames_problem(self):
     return "gym_discrete_problem_with_agent_on_wrapped_long_pong"
 
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedBreakout(
@@ -549,10 +576,16 @@ class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
   GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
+
   @property
   def initial_frames_problem(self):
     return "gym_discrete_problem_with_agent_on_wrapped_breakout"
 
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedPong(
@@ -587,10 +620,16 @@ class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnFreeway(
   GymSimulatedDiscreteProblem, GymFreewayRandom):
+
   @property
   def initial_frames_problem(self):
     return "gym_discrete_problem_with_agent_on_freeway"
 
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnFreeway(
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index eb0635723..16e16e3fa 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -146,6 +146,7 @@ def train_agent(problem_name, agent_model_dir,
   ppo_hparams.epoch_length = hparams.ppo_epoch_length
   ppo_hparams.num_agents = hparams.ppo_num_agents
   ppo_hparams.world_model_dir = world_model_dir
+  ppo_hparams.add_hparam("force_beginning_resets", True)
   if hparams.ppo_learning_rate:
     ppo_hparams.learning_rate = hparams.ppo_learning_rate
 
@@ -157,14 +158,6 @@ def train_agent(problem_name, agent_model_dir,
   environment_spec.simulation_random_starts = hparams.simulation_random_starts
   environment_spec.intrinsic_reward_scale = hparams.intrinsic_reward_scale
 
-  # 4x for the StackAndSkipWrapper minus one to always finish for reporting.
-  ppo_time_limit = ppo_hparams.epoch_length - 1
-  ppo_time_limit *= model_hparams.video_num_input_frames
-
-  wrappers = environment_spec.wrappers + \
-             [[TimeLimitWrapper, {"timelimit": ppo_time_limit}]]
-  environment_spec.wrappers = wrappers
-
   ppo_hparams.add_hparam("environment_spec", environment_spec)
 
   with temporary_flags({

From 0bba246a7a0535d21c084e74400369a02ac832a9 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 25 Jun 2018 19:06:42 +0200
Subject: [PATCH 0212/2720] shorten the test (as it is running on cpu)

---
 tensor2tensor/rl/rl_trainer_lib_test.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index e0bee9678..39ff5f278 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -55,6 +55,7 @@ def test_no_crash_cartpole(self):
   # is done on the ppo code
   #
   # To run the test change epochs_num=2 to epoch_num=200
+  # and epoch_length=5 to epoch_length=200
   # (it is set like that to meet travis timeouts
   def test_train_pong(self):
     hparams = tf.contrib.training.\
@@ -62,7 +63,7 @@ def test_train_pong(self):
               eval_every_epochs=10,
               num_agents=20,
               optimization_epochs=3,
-              epoch_length=200,
+              epoch_length=5,
               entropy_loss_coef=0.003,
               learning_rate=8e-05,
               optimizer="Adam",

From 44b070448d14e3d3865a961564bc39be4715d3ae Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 25 Jun 2018 11:25:35 -0700
Subject: [PATCH 0213/2720] internal merge of PR #865

PiperOrigin-RevId: 201988159
---
 .../data_generators/style_transfer.py         | 33 ++++++++++---------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/data_generators/style_transfer.py b/tensor2tensor/data_generators/style_transfer.py
index b7a6bf35b..96af64bd8 100644
--- a/tensor2tensor/data_generators/style_transfer.py
+++ b/tensor2tensor/data_generators/style_transfer.py
@@ -12,12 +12,11 @@
 # 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 classes for text-based language style transfer problems.
+"""Base classes for text-based language style transfer problems.
 
 * StyleTransferProblem: abstract class for style transfer problems.
-* StyleTransferShakespeare: specific problem implementation that enriches language with Shakespeare-like style.
+* StyleTransferShakespeare: specific problem implementation that enriches
+  language with Shakespeare-like style.
 """
 
 from __future__ import absolute_import
@@ -27,13 +26,13 @@
 import os
 import tarfile
 
-import tensorflow as tf
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
+import tensorflow as tf
+
 logger = tf.logging
 
 """
@@ -59,7 +58,7 @@
 
 
 class StyleTransferProblem(text_problems.Text2TextProblem):
-  """Base class for transfering styles problems"""
+  """Base class for transferring styles problems"""
 
   @property
   def target(self):
@@ -122,33 +121,37 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
 
 @registry.register_problem
 class StyleTransferShakespeareToModern(StyleTransferProblem):
-  """Transfering style from Shakespeare original English to modern one"""
+  """Transferring style from Shakespeare original English to modern one"""
 
   @property
   def target(self):
-    return '.modern'
+    return ".modern"
 
   @property
   def source(self):
-    return '.original'
+    return ".original"
 
   def dataset_url(self, dataset_split):
     train = dataset_split == problem.DatasetSplit.TRAIN
-    return _SHAKESPEARE_MODERN_TRAIN_DATASET if train else _SHAKESPEARE_MODERN_DEV_DATASET
+    if train:
+      return _SHAKESPEARE_MODERN_TRAIN_DATASET
+    return _SHAKESPEARE_MODERN_DEV_DATASET
 
 
 @registry.register_problem
 class StyleTransferModernToShakespeare(StyleTransferProblem):
-  """Transfering style from modern English to Shakespeare original English"""
+  """Transferring style from modern English to Shakespeare original English"""
 
   @property
   def target(self):
-    return '.original'
+    return ".original"
 
   @property
   def source(self):
-    return '.modern'
+    return ".modern"
 
   def dataset_url(self, dataset_split):
     train = dataset_split == problem.DatasetSplit.TRAIN
-    return _SHAKESPEARE_MODERN_TRAIN_DATASET if train else _SHAKESPEARE_MODERN_DEV_DATASET
+    if train:
+      return _SHAKESPEARE_MODERN_TRAIN_DATASET
+    return _SHAKESPEARE_MODERN_DEV_DATASET

From 660d49fc7a20025c1a8ae4448ad92022f042f02b Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 25 Jun 2018 11:33:23 -0700
Subject: [PATCH 0214/2720] Add resize options (incl. dilation) to all
 MultiResolution problems

PiperOrigin-RevId: 201989631
---
 tensor2tensor/data_generators/celeba.py       | 22 ++++++---
 tensor2tensor/data_generators/celeba_test.py  | 49 +++++++++++++++++++
 tensor2tensor/data_generators/image_utils.py  | 43 +++++++++++++++-
 .../data_generators/image_utils_test.py       | 42 ++++++++++++++++
 tensor2tensor/data_generators/imagenet.py     | 23 +++++----
 .../data_generators/imagenet_test.py          | 48 ++++++++++++++++++
 tensor2tensor/data_generators/mscoco.py       | 18 +++++--
 tensor2tensor/data_generators/mscoco_test.py  | 48 ++++++++++++++++++
 8 files changed, 271 insertions(+), 22 deletions(-)
 create mode 100644 tensor2tensor/data_generators/celeba_test.py
 create mode 100644 tensor2tensor/data_generators/imagenet_test.py
 create mode 100644 tensor2tensor/data_generators/mscoco_test.py

diff --git a/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index 0cc1a6319..476e8f04c 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -157,22 +157,28 @@ def dataset_filename(self):
 
   def preprocess_example(self, example, mode, hparams):
     image = example["inputs"]
-    if hasattr(hparams, "resize_method"):
-      method = getattr(tf.image.ResizeMethod, hparams.resize_method)
-    else:  # default
-      method = tf.image.ResizeMethod.BICUBIC
+    # Get resize method. Include a default if not specified, or if it's not in
+    # TensorFlow's collection of pre-implemented resize methods.
+    resize_method = getattr(hparams, "resize_method", "BICUBIC")
+    resize_method = getattr(tf.image.ResizeMethod, resize_method, resize_method)
 
     # Remove boundaries in CelebA images. Remove 40 pixels each side
     # vertically and 20 pixels each side horizontally.
     image = tf.image.crop_to_bounding_box(image, 40, 20, 218 - 80, 178 - 40)
 
-    scaled_images = image_utils.make_multiscale(
-        image, hparams.resolutions,
-        resize_method=method, num_channels=self.num_channels)
+    highest_res = hparams.resolutions[-1]
+    if resize_method == "DILATED":
+      # Resize image so that dilated subsampling is properly divisible.
+      scaled_image = image_utils.resize_by_area(image, highest_res)
+      scaled_images = image_utils.make_multiscale_dilated(
+          scaled_image, hparams.resolutions, num_channels=self.num_channels)
+    else:
+      scaled_images = image_utils.make_multiscale(
+          image, hparams.resolutions,
+          resize_method=resize_method, num_channels=self.num_channels)
 
     # Pack tuple of scaled images into one tensor. We do this by enforcing the
     # columns to match for every resolution.
-    highest_res = hparams.resolutions[-1]
     example["inputs"] = image
     example["targets"] = tf.concat([
         tf.reshape(scaled_image,
diff --git a/tensor2tensor/data_generators/celeba_test.py b/tensor2tensor/data_generators/celeba_test.py
new file mode 100644
index 000000000..3c337d106
--- /dev/null
+++ b/tensor2tensor/data_generators/celeba_test.py
@@ -0,0 +1,49 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 CelebA."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators import celeba
+
+import tensorflow as tf
+
+from google3.testing.pybase import parameterized
+
+
+class CelebaTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.named_parameters(
+      ("Default", None),
+      ("Area", "AREA"),
+      ("Dilated", "DILATED"))
+  def testCelebaMultiResolutionPreprocessExample(self, resize_method):
+    example = {"inputs": tf.random_uniform([218, 178, 3], minval=-1.)}
+    mode = tf.estimator.ModeKeys.TRAIN
+    hparams = tf.contrib.training.HParams(resolutions=[8, 16, 32])
+    if resize_method is not None:
+      hparams.resize_method = resize_method
+
+    problem = celeba.ImageCelebaMultiResolution()
+    preprocessed_example = problem.preprocess_example(example, mode, hparams)
+    self.assertLen(preprocessed_example, 2)
+    self.assertEqual(preprocessed_example["inputs"].shape, (138, 138, 3))
+    self.assertEqual(preprocessed_example["targets"].shape, (42, 32, 3))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index e3088ac9a..9f4b6b0c9 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -22,6 +22,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -37,7 +38,18 @@ def resize_by_area(img, size):
 def make_multiscale(image, resolutions,
                     resize_method=tf.image.ResizeMethod.BICUBIC,
                     num_channels=3):
-  """Returns list of scaled images, one for each resolution."""
+  """Returns list of scaled images, one for each resolution.
+
+  Args:
+    image: Tensor of shape [height, height, num_channels].
+    resolutions: List of heights that image's height is resized to.
+    resize_method: tf.image.ResizeMethod.
+    num_channels: Number of channels in image.
+
+  Returns:
+    List of Tensors, one for each resolution with shape given by
+    [resolutions[i], resolutions[i], num_channels].
+  """
   scaled_images = []
   for height in resolutions:
     scaled_image = tf.image.resize_images(
@@ -51,6 +63,35 @@ def make_multiscale(image, resolutions,
   return scaled_images
 
 
+def make_multiscale_dilated(image, resolutions, num_channels=3):
+  """Returns list of scaled images, one for each resolution.
+
+  Resizes by skipping every nth pixel.
+
+  Args:
+    image: Tensor of shape [height, height, num_channels].
+    resolutions: List of heights that image's height is resized to. The function
+      assumes VALID padding, so the original image's height must be divisible
+      by each resolution's height to return the exact resolution size.
+    num_channels: Number of channels in image.
+
+  Returns:
+    List of Tensors, one for each resolution with shape given by
+    [resolutions[i], resolutions[i], num_channels] if resolutions properly
+    divide the original image's height; otherwise shape height and width is up
+    to valid skips.
+  """
+  image_height = common_layers.shape_list(image)[0]
+  scaled_images = []
+  for height in resolutions:
+    dilation_rate = image_height // height  # assuming height = width
+    scaled_image = image[::dilation_rate, ::dilation_rate]
+    scaled_image = tf.to_int64(scaled_image)
+    scaled_image.set_shape([None, None, num_channels])
+    scaled_images.append(scaled_image)
+  return scaled_images
+
+
 class ImageProblem(problem.Problem):
   """Base class for problems with images."""
 
diff --git a/tensor2tensor/data_generators/image_utils_test.py b/tensor2tensor/data_generators/image_utils_test.py
index c47992f12..c8f8ca47e 100644
--- a/tensor2tensor/data_generators/image_utils_test.py
+++ b/tensor2tensor/data_generators/image_utils_test.py
@@ -70,6 +70,48 @@ def testImageGenerator(self):
       decoded2 = sess.run(decoded_png_t, feed_dict={image_t: encoded_img2[0]})
       self.assertAllClose(decoded2, image2)
 
+  def testMakeMultiscaleDivisible(self):
+    image = tf.random_normal([256, 256, 3])
+    resolutions = [8, 16, 64, 256]
+    scaled_images = image_utils.make_multiscale(image, resolutions)
+    self.assertEqual(scaled_images[0].shape, (8, 8, 3))
+    self.assertEqual(scaled_images[1].shape, (16, 16, 3))
+    self.assertEqual(scaled_images[2].shape, (64, 64, 3))
+    self.assertEqual(scaled_images[3].shape, (256, 256, 3))
+
+  def testMakeMultiscaleIndivisible(self):
+    image = tf.random_normal([256, 256, 3])
+    resolutions = [255]
+    scaled_images = image_utils.make_multiscale(image, resolutions)
+    self.assertEqual(scaled_images[0].shape, (255, 255, 3))
+
+  def testMakeMultiscaleLarger(self):
+    image = tf.random_normal([256, 256, 3])
+    resolutions = [257]
+    scaled_images = image_utils.make_multiscale(image, resolutions)
+    self.assertEqual(scaled_images[0].shape, (257, 257, 3))
+
+  def testMakeMultiscaleDilatedDivisible(self):
+    image = tf.random_normal([256, 256, 3])
+    resolutions = [8, 16, 64, 256]
+    scaled_images = image_utils.make_multiscale_dilated(image, resolutions)
+    self.assertEqual(scaled_images[0].shape, (8, 8, 3))
+    self.assertEqual(scaled_images[1].shape, (16, 16, 3))
+    self.assertEqual(scaled_images[2].shape, (64, 64, 3))
+    self.assertEqual(scaled_images[3].shape, (256, 256, 3))
+
+  def testMakeMultiscaleDilatedIndivisible(self):
+    image = tf.random_normal([256, 256, 3])
+    resolutions = [255]
+    scaled_images = image_utils.make_multiscale_dilated(image, resolutions)
+    self.assertEqual(scaled_images[0].shape, (256, 256, 3))
+
+  def testMakeMultiscaleDilatedLarger(self):
+    image = tf.random_normal([256, 256, 3])
+    resolutions = [257]
+    with self.assertRaisesRegexp(ValueError, "strides.* must be non-zero"):
+      _ = image_utils.make_multiscale_dilated(image, resolutions)
+
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 5735244ce..7d569bbd8 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -240,21 +240,24 @@ def dev_shards(self):
 
   def preprocess_example(self, example, mode, hparams):
     image = example["inputs"]
+    # Get resize method. Include a default if not specified, or if it's not in
+    # TensorFlow's collection of pre-implemented resize methods.
+    resize_method = getattr(hparams, "resize_method", "BICUBIC")
+    resize_method = getattr(tf.image.ResizeMethod, resize_method, resize_method)
+
+    if resize_method == "DILATED":
+      scaled_images = image_utils.make_multiscale_dilated(
+          image, hparams.resolutions, num_channels=self.num_channels)
+    else:
+      scaled_images = image_utils.make_multiscale(
+          image, hparams.resolutions,
+          resize_method=resize_method, num_channels=self.num_channels)
 
-    if hasattr(hparams, "resize_method"):
-      method = getattr(tf.image.ResizeMethod, hparams.resize_method)
-    else:  # default
-      method = tf.image.ResizeMethod.BICUBIC
-
-    scaled_images = image_utils.make_multiscale(
-        image, hparams.resolutions,
-        resize_method=method, num_channels=self.num_channels)
-
-    highest_res = hparams.resolutions[-1]
     # Pack tuple of scaled images into one tensor. We do this by enforcing the
     # columns to match for every resolution.
     # TODO(avaswani, trandustin): We should create tuples because this will not
     # work if height*width of low res < width of high res
+    highest_res = hparams.resolutions[-1]
     example["inputs"] = tf.concat([
         tf.reshape(scaled_image,
                    [res**2 // highest_res, highest_res, self.num_channels])
diff --git a/tensor2tensor/data_generators/imagenet_test.py b/tensor2tensor/data_generators/imagenet_test.py
new file mode 100644
index 000000000..1e0007cee
--- /dev/null
+++ b/tensor2tensor/data_generators/imagenet_test.py
@@ -0,0 +1,48 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 ImageNet."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators import imagenet
+
+import tensorflow as tf
+
+from google3.testing.pybase import parameterized
+
+
+class ImagenetTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.named_parameters(
+      ("Default", None),
+      ("Area", "AREA"),
+      ("Dilated", "DILATED"))
+  def testImagenetMultiResolutionPreprocessExample(self, resize_method):
+    example = {"inputs": tf.random_uniform([64, 64, 3], minval=-1.)}
+    mode = tf.estimator.ModeKeys.TRAIN
+    hparams = tf.contrib.training.HParams(resolutions=[8, 16, 32])
+    if resize_method is not None:
+      hparams.resize_method = resize_method
+
+    problem = imagenet.ImageImagenetMultiResolutionGen()
+    preprocessed_example = problem.preprocess_example(example, mode, hparams)
+    self.assertLen(preprocessed_example, 1)
+    self.assertEqual(preprocessed_example["inputs"].shape, (42, 32, 3))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index acb8bc697..28a1ea935 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -225,12 +225,24 @@ def dataset_filename(self):
 
   def preprocess_example(self, example, mode, hparams):
     image = example["inputs"]
-    scaled_images = image_utils.make_multiscale(
-        image, hparams.resolutions, num_channels=self.num_channels)
+    # Get resize method. Include a default if not specified, or if it's not in
+    # TensorFlow's collection of pre-implemented resize methods.
+    resize_method = getattr(hparams, "resize_method", "BICUBIC")
+    resize_method = getattr(tf.image.ResizeMethod, resize_method, resize_method)
+
+    highest_res = hparams.resolutions[-1]
+    if resize_method == "DILATED":
+      # Resize image so that dilated subsampling is properly divisible.
+      scaled_image = image_utils.resize_by_area(image, highest_res)
+      scaled_images = image_utils.make_multiscale_dilated(
+          scaled_image, hparams.resolutions, num_channels=self.num_channels)
+    else:
+      scaled_images = image_utils.make_multiscale(
+          image, hparams.resolutions,
+          resize_method=resize_method, num_channels=self.num_channels)
 
     # Pack tuple of scaled images into one tensor. We do this by enforcing the
     # columns to match for every resolution.
-    highest_res = hparams.resolutions[-1]
     example["inputs"] = tf.concat([
         tf.reshape(scaled_image,
                    [res**2 // highest_res, highest_res, self.num_channels])
diff --git a/tensor2tensor/data_generators/mscoco_test.py b/tensor2tensor/data_generators/mscoco_test.py
new file mode 100644
index 000000000..0652feb22
--- /dev/null
+++ b/tensor2tensor/data_generators/mscoco_test.py
@@ -0,0 +1,48 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 MS COCO."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators import mscoco
+
+import tensorflow as tf
+
+from google3.testing.pybase import parameterized
+
+
+class MscocoTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.named_parameters(
+      ("Default", None),
+      ("Area", "AREA"),
+      ("Dilated", "DILATED"))
+  def testMsCocoMultiResolutionPreprocessExample(self, resize_method):
+    example = {"inputs": tf.random_uniform([400, 400, 3], minval=-1.)}
+    mode = tf.estimator.ModeKeys.TRAIN
+    hparams = tf.contrib.training.HParams(resolutions=[8, 16, 32])
+    if resize_method is not None:
+      hparams.resize_method = resize_method
+
+    problem = mscoco.ImageTextMsCocoMultiResolution()
+    preprocessed_example = problem.preprocess_example(example, mode, hparams)
+    self.assertLen(preprocessed_example, 1)
+    self.assertEqual(preprocessed_example["inputs"].shape, (42, 32, 3))
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 0f96e35818c7c68122981924aeddf1c911ee6f50 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 25 Jun 2018 12:02:46 -0700
Subject: [PATCH 0215/2720] Add resize options (incl. dilation) to all
 MultiResolution problems

PiperOrigin-RevId: 201994902
---
 tensor2tensor/data_generators/celeba.py       | 22 +++------
 tensor2tensor/data_generators/celeba_test.py  | 49 -------------------
 tensor2tensor/data_generators/image_utils.py  | 43 +---------------
 .../data_generators/image_utils_test.py       | 42 ----------------
 tensor2tensor/data_generators/imagenet.py     | 23 ++++-----
 .../data_generators/imagenet_test.py          | 48 ------------------
 tensor2tensor/data_generators/mscoco.py       | 18 ++-----
 tensor2tensor/data_generators/mscoco_test.py  | 48 ------------------
 8 files changed, 22 insertions(+), 271 deletions(-)
 delete mode 100644 tensor2tensor/data_generators/celeba_test.py
 delete mode 100644 tensor2tensor/data_generators/imagenet_test.py
 delete mode 100644 tensor2tensor/data_generators/mscoco_test.py

diff --git a/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index 476e8f04c..0cc1a6319 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -157,28 +157,22 @@ def dataset_filename(self):
 
   def preprocess_example(self, example, mode, hparams):
     image = example["inputs"]
-    # Get resize method. Include a default if not specified, or if it's not in
-    # TensorFlow's collection of pre-implemented resize methods.
-    resize_method = getattr(hparams, "resize_method", "BICUBIC")
-    resize_method = getattr(tf.image.ResizeMethod, resize_method, resize_method)
+    if hasattr(hparams, "resize_method"):
+      method = getattr(tf.image.ResizeMethod, hparams.resize_method)
+    else:  # default
+      method = tf.image.ResizeMethod.BICUBIC
 
     # Remove boundaries in CelebA images. Remove 40 pixels each side
     # vertically and 20 pixels each side horizontally.
     image = tf.image.crop_to_bounding_box(image, 40, 20, 218 - 80, 178 - 40)
 
-    highest_res = hparams.resolutions[-1]
-    if resize_method == "DILATED":
-      # Resize image so that dilated subsampling is properly divisible.
-      scaled_image = image_utils.resize_by_area(image, highest_res)
-      scaled_images = image_utils.make_multiscale_dilated(
-          scaled_image, hparams.resolutions, num_channels=self.num_channels)
-    else:
-      scaled_images = image_utils.make_multiscale(
-          image, hparams.resolutions,
-          resize_method=resize_method, num_channels=self.num_channels)
+    scaled_images = image_utils.make_multiscale(
+        image, hparams.resolutions,
+        resize_method=method, num_channels=self.num_channels)
 
     # Pack tuple of scaled images into one tensor. We do this by enforcing the
     # columns to match for every resolution.
+    highest_res = hparams.resolutions[-1]
     example["inputs"] = image
     example["targets"] = tf.concat([
         tf.reshape(scaled_image,
diff --git a/tensor2tensor/data_generators/celeba_test.py b/tensor2tensor/data_generators/celeba_test.py
deleted file mode 100644
index 3c337d106..000000000
--- a/tensor2tensor/data_generators/celeba_test.py
+++ /dev/null
@@ -1,49 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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 CelebA."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.data_generators import celeba
-
-import tensorflow as tf
-
-from google3.testing.pybase import parameterized
-
-
-class CelebaTest(parameterized.TestCase, tf.test.TestCase):
-
-  @parameterized.named_parameters(
-      ("Default", None),
-      ("Area", "AREA"),
-      ("Dilated", "DILATED"))
-  def testCelebaMultiResolutionPreprocessExample(self, resize_method):
-    example = {"inputs": tf.random_uniform([218, 178, 3], minval=-1.)}
-    mode = tf.estimator.ModeKeys.TRAIN
-    hparams = tf.contrib.training.HParams(resolutions=[8, 16, 32])
-    if resize_method is not None:
-      hparams.resize_method = resize_method
-
-    problem = celeba.ImageCelebaMultiResolution()
-    preprocessed_example = problem.preprocess_example(example, mode, hparams)
-    self.assertLen(preprocessed_example, 2)
-    self.assertEqual(preprocessed_example["inputs"].shape, (138, 138, 3))
-    self.assertEqual(preprocessed_example["targets"].shape, (42, 32, 3))
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 9f4b6b0c9..e3088ac9a 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -22,7 +22,6 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
-from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -38,18 +37,7 @@ def resize_by_area(img, size):
 def make_multiscale(image, resolutions,
                     resize_method=tf.image.ResizeMethod.BICUBIC,
                     num_channels=3):
-  """Returns list of scaled images, one for each resolution.
-
-  Args:
-    image: Tensor of shape [height, height, num_channels].
-    resolutions: List of heights that image's height is resized to.
-    resize_method: tf.image.ResizeMethod.
-    num_channels: Number of channels in image.
-
-  Returns:
-    List of Tensors, one for each resolution with shape given by
-    [resolutions[i], resolutions[i], num_channels].
-  """
+  """Returns list of scaled images, one for each resolution."""
   scaled_images = []
   for height in resolutions:
     scaled_image = tf.image.resize_images(
@@ -63,35 +51,6 @@ def make_multiscale(image, resolutions,
   return scaled_images
 
 
-def make_multiscale_dilated(image, resolutions, num_channels=3):
-  """Returns list of scaled images, one for each resolution.
-
-  Resizes by skipping every nth pixel.
-
-  Args:
-    image: Tensor of shape [height, height, num_channels].
-    resolutions: List of heights that image's height is resized to. The function
-      assumes VALID padding, so the original image's height must be divisible
-      by each resolution's height to return the exact resolution size.
-    num_channels: Number of channels in image.
-
-  Returns:
-    List of Tensors, one for each resolution with shape given by
-    [resolutions[i], resolutions[i], num_channels] if resolutions properly
-    divide the original image's height; otherwise shape height and width is up
-    to valid skips.
-  """
-  image_height = common_layers.shape_list(image)[0]
-  scaled_images = []
-  for height in resolutions:
-    dilation_rate = image_height // height  # assuming height = width
-    scaled_image = image[::dilation_rate, ::dilation_rate]
-    scaled_image = tf.to_int64(scaled_image)
-    scaled_image.set_shape([None, None, num_channels])
-    scaled_images.append(scaled_image)
-  return scaled_images
-
-
 class ImageProblem(problem.Problem):
   """Base class for problems with images."""
 
diff --git a/tensor2tensor/data_generators/image_utils_test.py b/tensor2tensor/data_generators/image_utils_test.py
index c8f8ca47e..c47992f12 100644
--- a/tensor2tensor/data_generators/image_utils_test.py
+++ b/tensor2tensor/data_generators/image_utils_test.py
@@ -70,48 +70,6 @@ def testImageGenerator(self):
       decoded2 = sess.run(decoded_png_t, feed_dict={image_t: encoded_img2[0]})
       self.assertAllClose(decoded2, image2)
 
-  def testMakeMultiscaleDivisible(self):
-    image = tf.random_normal([256, 256, 3])
-    resolutions = [8, 16, 64, 256]
-    scaled_images = image_utils.make_multiscale(image, resolutions)
-    self.assertEqual(scaled_images[0].shape, (8, 8, 3))
-    self.assertEqual(scaled_images[1].shape, (16, 16, 3))
-    self.assertEqual(scaled_images[2].shape, (64, 64, 3))
-    self.assertEqual(scaled_images[3].shape, (256, 256, 3))
-
-  def testMakeMultiscaleIndivisible(self):
-    image = tf.random_normal([256, 256, 3])
-    resolutions = [255]
-    scaled_images = image_utils.make_multiscale(image, resolutions)
-    self.assertEqual(scaled_images[0].shape, (255, 255, 3))
-
-  def testMakeMultiscaleLarger(self):
-    image = tf.random_normal([256, 256, 3])
-    resolutions = [257]
-    scaled_images = image_utils.make_multiscale(image, resolutions)
-    self.assertEqual(scaled_images[0].shape, (257, 257, 3))
-
-  def testMakeMultiscaleDilatedDivisible(self):
-    image = tf.random_normal([256, 256, 3])
-    resolutions = [8, 16, 64, 256]
-    scaled_images = image_utils.make_multiscale_dilated(image, resolutions)
-    self.assertEqual(scaled_images[0].shape, (8, 8, 3))
-    self.assertEqual(scaled_images[1].shape, (16, 16, 3))
-    self.assertEqual(scaled_images[2].shape, (64, 64, 3))
-    self.assertEqual(scaled_images[3].shape, (256, 256, 3))
-
-  def testMakeMultiscaleDilatedIndivisible(self):
-    image = tf.random_normal([256, 256, 3])
-    resolutions = [255]
-    scaled_images = image_utils.make_multiscale_dilated(image, resolutions)
-    self.assertEqual(scaled_images[0].shape, (256, 256, 3))
-
-  def testMakeMultiscaleDilatedLarger(self):
-    image = tf.random_normal([256, 256, 3])
-    resolutions = [257]
-    with self.assertRaisesRegexp(ValueError, "strides.* must be non-zero"):
-      _ = image_utils.make_multiscale_dilated(image, resolutions)
-
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 7d569bbd8..5735244ce 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -240,24 +240,21 @@ def dev_shards(self):
 
   def preprocess_example(self, example, mode, hparams):
     image = example["inputs"]
-    # Get resize method. Include a default if not specified, or if it's not in
-    # TensorFlow's collection of pre-implemented resize methods.
-    resize_method = getattr(hparams, "resize_method", "BICUBIC")
-    resize_method = getattr(tf.image.ResizeMethod, resize_method, resize_method)
-
-    if resize_method == "DILATED":
-      scaled_images = image_utils.make_multiscale_dilated(
-          image, hparams.resolutions, num_channels=self.num_channels)
-    else:
-      scaled_images = image_utils.make_multiscale(
-          image, hparams.resolutions,
-          resize_method=resize_method, num_channels=self.num_channels)
 
+    if hasattr(hparams, "resize_method"):
+      method = getattr(tf.image.ResizeMethod, hparams.resize_method)
+    else:  # default
+      method = tf.image.ResizeMethod.BICUBIC
+
+    scaled_images = image_utils.make_multiscale(
+        image, hparams.resolutions,
+        resize_method=method, num_channels=self.num_channels)
+
+    highest_res = hparams.resolutions[-1]
     # Pack tuple of scaled images into one tensor. We do this by enforcing the
     # columns to match for every resolution.
     # TODO(avaswani, trandustin): We should create tuples because this will not
     # work if height*width of low res < width of high res
-    highest_res = hparams.resolutions[-1]
     example["inputs"] = tf.concat([
         tf.reshape(scaled_image,
                    [res**2 // highest_res, highest_res, self.num_channels])
diff --git a/tensor2tensor/data_generators/imagenet_test.py b/tensor2tensor/data_generators/imagenet_test.py
deleted file mode 100644
index 1e0007cee..000000000
--- a/tensor2tensor/data_generators/imagenet_test.py
+++ /dev/null
@@ -1,48 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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 ImageNet."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.data_generators import imagenet
-
-import tensorflow as tf
-
-from google3.testing.pybase import parameterized
-
-
-class ImagenetTest(parameterized.TestCase, tf.test.TestCase):
-
-  @parameterized.named_parameters(
-      ("Default", None),
-      ("Area", "AREA"),
-      ("Dilated", "DILATED"))
-  def testImagenetMultiResolutionPreprocessExample(self, resize_method):
-    example = {"inputs": tf.random_uniform([64, 64, 3], minval=-1.)}
-    mode = tf.estimator.ModeKeys.TRAIN
-    hparams = tf.contrib.training.HParams(resolutions=[8, 16, 32])
-    if resize_method is not None:
-      hparams.resize_method = resize_method
-
-    problem = imagenet.ImageImagenetMultiResolutionGen()
-    preprocessed_example = problem.preprocess_example(example, mode, hparams)
-    self.assertLen(preprocessed_example, 1)
-    self.assertEqual(preprocessed_example["inputs"].shape, (42, 32, 3))
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index 28a1ea935..acb8bc697 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -225,24 +225,12 @@ def dataset_filename(self):
 
   def preprocess_example(self, example, mode, hparams):
     image = example["inputs"]
-    # Get resize method. Include a default if not specified, or if it's not in
-    # TensorFlow's collection of pre-implemented resize methods.
-    resize_method = getattr(hparams, "resize_method", "BICUBIC")
-    resize_method = getattr(tf.image.ResizeMethod, resize_method, resize_method)
-
-    highest_res = hparams.resolutions[-1]
-    if resize_method == "DILATED":
-      # Resize image so that dilated subsampling is properly divisible.
-      scaled_image = image_utils.resize_by_area(image, highest_res)
-      scaled_images = image_utils.make_multiscale_dilated(
-          scaled_image, hparams.resolutions, num_channels=self.num_channels)
-    else:
-      scaled_images = image_utils.make_multiscale(
-          image, hparams.resolutions,
-          resize_method=resize_method, num_channels=self.num_channels)
+    scaled_images = image_utils.make_multiscale(
+        image, hparams.resolutions, num_channels=self.num_channels)
 
     # Pack tuple of scaled images into one tensor. We do this by enforcing the
     # columns to match for every resolution.
+    highest_res = hparams.resolutions[-1]
     example["inputs"] = tf.concat([
         tf.reshape(scaled_image,
                    [res**2 // highest_res, highest_res, self.num_channels])
diff --git a/tensor2tensor/data_generators/mscoco_test.py b/tensor2tensor/data_generators/mscoco_test.py
deleted file mode 100644
index 0652feb22..000000000
--- a/tensor2tensor/data_generators/mscoco_test.py
+++ /dev/null
@@ -1,48 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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 MS COCO."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.data_generators import mscoco
-
-import tensorflow as tf
-
-from google3.testing.pybase import parameterized
-
-
-class MscocoTest(parameterized.TestCase, tf.test.TestCase):
-
-  @parameterized.named_parameters(
-      ("Default", None),
-      ("Area", "AREA"),
-      ("Dilated", "DILATED"))
-  def testMsCocoMultiResolutionPreprocessExample(self, resize_method):
-    example = {"inputs": tf.random_uniform([400, 400, 3], minval=-1.)}
-    mode = tf.estimator.ModeKeys.TRAIN
-    hparams = tf.contrib.training.HParams(resolutions=[8, 16, 32])
-    if resize_method is not None:
-      hparams.resize_method = resize_method
-
-    problem = mscoco.ImageTextMsCocoMultiResolution()
-    preprocessed_example = problem.preprocess_example(example, mode, hparams)
-    self.assertLen(preprocessed_example, 1)
-    self.assertEqual(preprocessed_example["inputs"].shape, (42, 32, 3))
-
-
-if __name__ == "__main__":
-  tf.test.main()

From ddd501897673c1721edf75a80369bc4d4576e509 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 25 Jun 2018 12:19:27 -0700
Subject: [PATCH 0216/2720] add multi label modality

PiperOrigin-RevId: 201997737
---
 tensor2tensor/layers/common_attention.py | 13 +++++++++++
 tensor2tensor/layers/modalities.py       | 28 ++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 66e77fbe7..8dbacf8ae 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -833,6 +833,19 @@ def embedding_to_padding(emb):
   return tf.to_float(tf.equal(emb_sum, 0.0))
 
 
+@expert_utils.add_name_scope()
+def padding_to_length(padding):
+  """Calculate the length of mask based on padding.
+
+  Args:
+    padding: a Tensor with shape [..., length].
+  Returns:
+    a Tensor with shape [...].
+  """
+  non_padding = 1.0 - padding
+  return tf.to_int64(tf.reduce_sum(non_padding, axis=-1))
+
+
 @expert_utils.add_name_scope()
 def attention_bias_local(length, max_backward, max_forward):
   """Create an bias tensor to be added to attention logits.
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index efbea07b5..a99a47110 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -705,6 +705,34 @@ def top(self, body_output, _):
       return tf.expand_dims(res, 3)
 
 
+@registry.register_class_label_modality("multi_label")
+class MultiLabelModality(ClassLabelModality):
+  """Used for multi label task."""
+
+  def targets_weights_fn(self):
+    """Target weight function for multi label, defaults to nonzero labels."""
+    weights_fn = common_layers.weights_nonzero
+    return weights_fn
+
+  def loss(self, top_out, targets):
+    """Average loss over the labels."""
+    logits = top_out
+    num_labels = tf.shape(targets)[1]
+    logits = tf.tile(logits, [1, num_labels, 1, 1])
+
+    xent, weights = common_layers.padded_cross_entropy(
+        logits,
+        targets,
+        self._model_hparams.label_smoothing,
+        weights_fn=self.targets_weights_fn,)
+    xent = tf.squeeze(xent, [2, 3])
+    weights = tf.squeeze(xent, [2, 3])
+    # average loss over all labels
+    loss = (tf.reduce_sum(xent, axis=1)
+            / (tf.reduce_sum(weights, axis=1) + 1e-8))
+    return tf.reduce_mean(loss)
+
+
 @registry.register_class_label_modality("onehot")
 class OneHotClassLabelModality(ClassLabelModality):
   """Used for one-hot encoded class labels."""

From 95948c502b9c4e58b9eba46f29920abeca94f225 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 25 Jun 2018 12:19:39 -0700
Subject: [PATCH 0217/2720] A starting setting of hyperparameters for text.

PiperOrigin-RevId: 201997759
---
 tensor2tensor/models/research/autoencoders.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 2c63018fd..8b5d9f197 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -782,6 +782,22 @@ def autoencoder_ordered_discrete():
   return hparams
 
 
+@registry.register_hparams
+def autoencoder_ordered_text():
+  """Ordered discrete autoencoder model for text."""
+  hparams = autoencoder_ordered_discrete()
+  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_warmup_steps = 2000
+  hparams.bottleneck_bits = 1024
+  hparams.batch_size = 2048
+  hparams.autoregressive_mode = "sru"
+  hparams.hidden_size = 256
+  hparams.max_hidden_size = 4096
+  hparams.bottleneck_warmup_steps = 10000
+  hparams.discretize_warmup_steps = 15000
+  return hparams
+
+
 @registry.register_hparams
 def autoencoder_ordered_discrete_vq():
   """Ordered discrete autoencoder model with VQ bottleneck."""

From a98443d372f3e96c54f77ee6c16c47186e870703 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 25 Jun 2018 12:25:15 -0700
Subject: [PATCH 0218/2720] Avoid potentially long chains of dataset concats

PiperOrigin-RevId: 201998472
---
 tensor2tensor/data_generators/problem.py | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index d6cd7811e..b90f77689 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -543,10 +543,11 @@ def dataset(self,
     data_files = sorted(tf.contrib.slim.parallel_reader.get_data_files(
         data_filepattern))
 
-    # Functions used in dataset transforms below
-    def _load_records_and_preprocess(filename):
-      # Load records from file with an 8MiB read buffer.
-      dataset = tf.data.TFRecordDataset(filename, buffer_size=8 * 1024 * 1024)
+    # Functions used in dataset transforms below. `filenames` can be either a
+    # `tf.string` tensor or `tf.data.Dataset` containing one or more filenames.
+    def _load_records_and_preprocess(filenames):
+      # Load records from file(s) with an 8MiB read buffer.
+      dataset = tf.data.TFRecordDataset(filenames, buffer_size=8 * 1024 * 1024)
       # Decode.
       dataset = dataset.map(self.decode_example, num_parallel_calls=num_threads)
       # Preprocess if requested.
@@ -566,17 +567,14 @@ def _load_records_and_preprocess(filename):
     if shuffle_files:
       random.shuffle(data_files)
 
+    dataset = tf.data.Dataset.from_tensor_slices(tf.constant(data_files))
     # Create data-set from files by parsing, pre-processing and interleaving.
     if shuffle_files:
-      dataset = tf.data.Dataset.from_tensor_slices(tf.constant(data_files))
       dataset = dataset.apply(
           tf.contrib.data.parallel_interleave(
               _load_records_and_preprocess, sloppy=True, cycle_length=8))
     else:
-      dataset = None
-      for f in data_files:
-        f_data = _load_records_and_preprocess(f)
-        dataset = f_data if dataset is None else dataset.concatenate(f_data)
+      dataset = _load_records_and_preprocess(dataset)
 
     dataset = dataset.map(
         self.maybe_reverse_and_copy, num_parallel_calls=num_threads)

From 90c5f41f92dd35b37d7e0003f81e82c8bae6648c Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Mon, 25 Jun 2018 13:02:20 -0700
Subject: [PATCH 0219/2720] Create a multihead graph attention function and
 move all graph functions from common attention.

PiperOrigin-RevId: 202004311
---
 tensor2tensor/layers/common_attention.py      |  77 ------
 .../common_message_passing_attention.py       | 226 +++++++++++++++++-
 2 files changed, 225 insertions(+), 78 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 8dbacf8ae..81fc7c44a 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1449,71 +1449,6 @@ def grouped_attention_multihead(query_antecedent,
     return o, extra_loss
 
 
-def graph_attention(q,
-                    k,
-                    v,
-                    bias,
-                    dropout_rate=0.0,
-                    image_shapes=None,
-                    name=None,
-                    make_image_summary=True,
-                    save_weights_to=None,
-                    dropout_broadcast_dims=None,
-                    adjacency_matrix=None,
-                    num_edge_types=5):
-  """graph attention.
-
-  Args:
-    q: a Tensor with shape [batch, heads, length_q, depth_k]
-    k: a Tensor with shape [batch, heads, length_kv, depth_k]
-    v: a Tensor with shape [batch, heads, length_kv, depth_v]
-    bias: bias Tensor (see attention_bias())
-    dropout_rate: a floating point number
-    image_shapes: optional tuple of integer scalars.
-      see comments for attention_image_summary()
-    name: an optional string
-    make_image_summary: True if you want an image summary.
-    save_weights_to: an optional dictionary to capture attention weights
-      for vizualization; the weights tensor will be appended there under
-      a string key created from the variable scope (including name).
-    dropout_broadcast_dims:  an optional list of integers less than 4
-      specifying in which dimensions to broadcast the dropout decisions.
-      saves memory.
-    adjacency_matrix: optional matrix of [batch, length, length] ids indicating
-      edge type
-    num_edge_types: an int indicating number of edge types
-  Returns:
-    A Tensor of shape [batch, length, depth(q)]
-  """
-  with tf.variable_scope(
-      name, default_name="dot_product_attention", values=[q, k, v]) as scope:
-    # [batch, num_heads, query_length, memory_length]
-    logits = tf.matmul(q, k, transpose_b=True)
-    if adjacency_matrix is not None:
-      key_head_depth = common_layers.shape_list(q)[-1]
-      adjacency_vectors = make_edge_vectors(
-          adjacency_matrix, num_edge_types, key_head_depth, name)
-      # zeroing out the vectors that have 0 entries in the adjacency
-      adjacency_vectors *= tf.to_float(
-          tf.expand_dims(adjacency_matrix, axis=-1))
-      # transposing q to be [batch, length_q, heads, depth_k]
-      # to allow for matmul with [batch, length_q, length_q, depth_k]
-      q_t = tf.transpose(q, [0, 2, 1, 3])
-      adj_logits = tf.matmul(q_t, adjacency_vectors, transpose_b=True)
-      logits += tf.transpose(adj_logits, [0, 2, 1, 3])
-    if bias is not None:
-      logits += bias
-    weights = tf.nn.softmax(logits, name="attention_weights")
-    if save_weights_to is not None:
-      save_weights_to[scope.name] = weights
-    # dropping out the attention links for each of the heads
-    weights = common_layers.dropout_with_broadcast_dims(
-        weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
-    if expert_utils.should_generate_summaries() and make_image_summary:
-      attention_image_summary(weights, image_shapes)
-    return tf.matmul(weights, v)
-
-
 def dot_product_attention(q,
                           k,
                           v,
@@ -2862,8 +2797,6 @@ def multihead_attention(query_antecedent,
                         make_image_summary=True,
                         dropout_broadcast_dims=None,
                         max_length=None,
-                        adjacency_matrix=None,
-                        num_edge_types=5,
                         vars_3d=False,
                         **kwargs):
   """Multihead scaled-dot-product attention with input/output transformations.
@@ -2913,9 +2846,6 @@ def multihead_attention(query_antecedent,
       specifying in which dimensions to broadcast the dropout decisions.
       saves memory.
     max_length: an integer - needed by relative attention
-    adjacency_matrix: an optional tensor of shape [batch, len_q, len_q]
-      containing edge vectors for attention
-    num_edge_types: number of edge types, an int
     vars_3d: use 3-dimensional variables for input/output transformations
     **kwargs (dict): Parameters for the attention function
 
@@ -3012,13 +2942,6 @@ def multihead_attention(query_antecedent,
                                 save_weights_to=save_weights_to,
                                 make_image_summary=make_image_summary,
                                 dropout_broadcast_dims=dropout_broadcast_dims)
-    elif attention_type == "edge_vector":
-      x = graph_attention(q, k, v, bias, dropout_rate, image_shapes,
-                          save_weights_to=save_weights_to,
-                          make_image_summary=make_image_summary,
-                          dropout_broadcast_dims=dropout_broadcast_dims,
-                          adjacency_matrix=adjacency_matrix,
-                          num_edge_types=num_edge_types)
     elif attention_type == "dot_product_relative":
       x = dot_product_attention_relative(
           q,
diff --git a/tensor2tensor/layers/common_message_passing_attention.py b/tensor2tensor/layers/common_message_passing_attention.py
index 3946662c0..a3fc38924 100644
--- a/tensor2tensor/layers/common_message_passing_attention.py
+++ b/tensor2tensor/layers/common_message_passing_attention.py
@@ -20,10 +20,234 @@
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import expert_utils
 
 import tensorflow as tf
 
 
+def multihead_graph_attention(query_antecedent,
+                              memory_antecedent,
+                              bias,
+                              total_key_depth,
+                              total_value_depth,
+                              output_depth,
+                              num_heads,
+                              dropout_rate,
+                              image_shapes=None,
+                              attention_type="edge_vector",
+                              name="multihead_graph_attention",
+                              save_weights_to=None,
+                              make_image_summary=True,
+                              dropout_broadcast_dims=None,
+                              adjacency_matrix=None,
+                              num_edge_types=5,
+                              vars_3d=False,
+                              **kwargs):
+  """Multihead scaled-dot-product attention with input/output transformations.
+
+  Args:
+    query_antecedent: a Tensor with shape [batch, length_q, channels]
+    memory_antecedent: a Tensor with shape [batch, length_m, channels] or None
+    bias: bias Tensor (see attention_bias())
+    total_key_depth: an integer
+    total_value_depth: an integer
+    output_depth: an integer
+    num_heads: an integer dividing total_key_depth and total_value_depth
+    dropout_rate: a floating point number
+    image_shapes: optional tuple of integer scalars.
+                  see comments for attention_image_summary()
+    attention_type: a string, either "dot_product", "dot_product_relative",
+                    "local_mask_right", "local_unmasked", "masked_dilated_1d",
+                    "unmasked_dilated_1d", graph, or any attention function
+                    with the signature (query, key, value, **kwargs)
+    name: an optional string.
+    save_weights_to: an optional dictionary to capture attention weights
+      for vizualization; the weights tensor will be appended there under
+      a string key created from the variable scope (including name).
+    make_image_summary: Whether to make an attention image summary.
+    dropout_broadcast_dims:  an optional list of integers less than 4
+      specifying in which dimensions to broadcast the dropout decisions.
+      saves memory.
+    adjacency_matrix: an optional tensor of shape [batch, len_q, len_q]
+      containing edge vectors for attention
+    num_edge_types: number of edge types, an int
+    vars_3d: use 3-dimensional variables for input/output transformations
+    **kwargs (dict): Parameters for the attention function
+
+  Caching:
+    WARNING: For decoder self-attention, i.e. when memory_antecedent == None,
+    the caching assumes that the bias contains future masking.
+
+    The caching works by saving all the previous key and value values so that
+    you are able to send just the last query location to this attention
+    function. I.e. if the cache dict is provided it assumes the query is of the
+    shape [batch_size, 1, hiddem_dim] rather than the full memory.
+
+  Returns:
+    The result of the attention transformation. The output shape is
+        [batch_size, length_q, hidden_dim]
+    unless the cache dict is provided in which case only the last memory
+    position is calculated and the output shape is [batch_size, 1, hidden_dim]
+    Optionally returns an additional loss parameters (ex: load balance loss for
+    the experts) returned by the attention_type function.
+
+  Raises:
+    ValueError: if the key depth or value depth are not divisible by the
+      number of attention heads.
+  """
+  if total_key_depth % num_heads != 0:
+    raise ValueError("Key depth (%d) must be divisible by the number of "
+                     "attention heads (%d)." % (total_key_depth, num_heads))
+  if total_value_depth % num_heads != 0:
+    raise ValueError("Value depth (%d) must be divisible by the number of "
+                     "attention heads (%d)." % (total_value_depth, num_heads))
+  vars_3d_num_heads = num_heads if vars_3d else None
+  with tf.variable_scope(name, default_name="multihead_attention",
+                         values=[query_antecedent, memory_antecedent]):
+
+    q, k, v = common_attention.compute_qkv(
+        query_antecedent, memory_antecedent, total_key_depth,
+        total_value_depth, vars_3d_num_heads=vars_3d_num_heads)
+    q = common_attention.split_heads(q, num_heads)
+    k = common_attention.split_heads(k, num_heads)
+    v = common_attention.split_heads(v, num_heads)
+
+    key_depth_per_head = total_key_depth // num_heads
+    if not vars_3d:
+      q *= key_depth_per_head**-0.5
+
+    additional_returned_value = None
+    if callable(attention_type):  # Generic way to extend multihead_attention
+      x = attention_type(q, k, v, **kwargs)
+      if isinstance(x, tuple):
+        x, additional_returned_value = x  # Unpack
+
+    elif attention_type == "edge_vector":
+      x = graph_attention(q, k, v, bias, dropout_rate, image_shapes,
+                          save_weights_to=save_weights_to,
+                          make_image_summary=make_image_summary,
+                          dropout_broadcast_dims=dropout_broadcast_dims,
+                          adjacency_matrix=adjacency_matrix,
+                          num_edge_types=num_edge_types)
+
+    x = common_attention.combine_heads(x)
+
+    # Set last dim specifically.
+    x.set_shape(x.shape.as_list()[:-1] + [total_value_depth])
+
+    if vars_3d:
+      o_var = tf.get_variable(
+          "o", [num_heads, total_value_depth // num_heads, output_depth])
+      o_var = tf.reshape(o_var, [total_value_depth, output_depth])
+      x = tf.tensordot(x, o_var, axes=1)
+    else:
+      x = common_layers.dense(
+          x, output_depth, use_bias=False, name="output_transform")
+    if additional_returned_value is not None:
+      return x, additional_returned_value
+    return x
+
+
+@expert_utils.add_name_scope()
+def make_edge_vectors(adjacency_matrix, num_edge_types, depth, name=None):
+  """Gets edge vectors for the edge types in the adjacency matrix.
+
+  Args:
+    adjacency_matrix: A [batch, num_nodes, num_nodes] tensor of ints.
+    num_edge_types: Number of different edge types
+    depth: Number of channels
+    name: a string
+  Returns:
+    A [batch, num_nodes, num_nodes, depth] vector of tensors
+  """
+  with tf.variable_scope(name, default_name="edge_vectors"):
+    att_adj_vectors_shape = [num_edge_types, depth]
+    adjacency_matrix_shape = common_layers.shape_list(adjacency_matrix)
+    adj_vectors = (
+        tf.get_variable(
+            "adj_vectors",
+            att_adj_vectors_shape,
+            initializer=tf.random_normal_initializer(0, depth**-0.5)) *
+        (depth**0.5))
+    # Avoiding gathers so that it works on TPUs
+    # adjacency_matrix_one_hot has shape
+    # [batch, num_nodes, num_nodes, num_edge_types]
+
+    adjacency_matrix_one_hot = tf.one_hot(adjacency_matrix, num_edge_types)
+
+    att_adj_vectors = tf.matmul(
+        tf.reshape(tf.to_float(adjacency_matrix_one_hot), [-1, num_edge_types]),
+        adj_vectors)
+    return tf.reshape(att_adj_vectors,
+                      [adjacency_matrix_shape[0], adjacency_matrix_shape[1],
+                       adjacency_matrix_shape[2], depth])
+
+
+def graph_attention(q,
+                    k,
+                    v,
+                    bias,
+                    dropout_rate=0.0,
+                    image_shapes=None,
+                    name=None,
+                    make_image_summary=True,
+                    save_weights_to=None,
+                    dropout_broadcast_dims=None,
+                    adjacency_matrix=None,
+                    num_edge_types=5):
+  """graph attention.
+
+  Args:
+    q: a Tensor with shape [batch, heads, length_q, depth_k]
+    k: a Tensor with shape [batch, heads, length_kv, depth_k]
+    v: a Tensor with shape [batch, heads, length_kv, depth_v]
+    bias: bias Tensor (see attention_bias())
+    dropout_rate: a floating point number
+    image_shapes: optional tuple of integer scalars.
+      see comments for attention_image_summary()
+    name: an optional string
+    make_image_summary: True if you want an image summary.
+    save_weights_to: an optional dictionary to capture attention weights
+      for vizualization; the weights tensor will be appended there under
+      a string key created from the variable scope (including name).
+    dropout_broadcast_dims:  an optional list of integers less than 4
+      specifying in which dimensions to broadcast the dropout decisions.
+      saves memory.
+    adjacency_matrix: optional matrix of [batch, length, length] ids indicating
+      edge type
+    num_edge_types: an int indicating number of edge types
+  Returns:
+    A Tensor of shape [batch, length, depth(q)]
+  """
+  with tf.variable_scope(
+      name, default_name="dot_product_attention", values=[q, k, v]) as scope:
+    # [batch, num_heads, query_length, memory_length]
+    logits = tf.matmul(q, k, transpose_b=True)
+    if adjacency_matrix is not None:
+      key_head_depth = common_layers.shape_list(q)[-1]
+      adjacency_vectors = make_edge_vectors(
+          adjacency_matrix, num_edge_types, key_head_depth, name)
+      # zeroing out the vectors that have 0 entries in the adjacency
+      adjacency_vectors *= tf.to_float(
+          tf.expand_dims(adjacency_matrix, axis=-1))
+      # transposing q to be [batch, length_q, heads, depth_k]
+      # to allow for matmul with [batch, length_q, length_q, depth_k]
+      q_t = tf.transpose(q, [0, 2, 1, 3])
+      adj_logits = tf.matmul(q_t, adjacency_vectors, transpose_b=True)
+      logits += tf.transpose(adj_logits, [0, 2, 1, 3])
+    if bias is not None:
+      logits += bias
+    weights = tf.nn.softmax(logits, name="attention_weights")
+    if save_weights_to is not None:
+      save_weights_to[scope.name] = weights
+    # dropping out the attention links for each of the heads
+    weights = common_layers.dropout_with_broadcast_dims(
+        weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
+    if common_layers.should_generate_summaries() and make_image_summary:
+      common_attention.attention_image_summary(weights, image_shapes)
+    return tf.matmul(weights, v)
+
+
 def compute_mpnn_qkv(node_states,
                      total_key_depth,
                      total_value_depth,
@@ -201,7 +425,7 @@ def dot_product_mpnn_attention(q, k, v, adjacency_matrix, num_edge_types,
     # getting dot products for q_i, k_j, and e_{ij}. This assumes that for
     # edge type 0, the dot products are 0
     all_edge_logits *= adjacency_matrix_one_hot
-    # logits will be [batch, length, length] after educing along
+    # logits will be [batch, length, length] after reducing along
     # axis 1 which has dimension num_edge_types.
     logits = tf.reduce_sum(all_edge_logits, axis=1)
     # ignoring edges if needed

From b2232412fb72528f6932caaab0299fdd4b96ff20 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 25 Jun 2018 13:41:27 -0700
Subject: [PATCH 0220/2720] Add resize options (incl. dilation) to all
 MultiResolution problems

PiperOrigin-RevId: 202011114
---
 setup.py                                      |  2 +-
 tensor2tensor/data_generators/celeba.py       | 22 +++++----
 tensor2tensor/data_generators/celeba_test.py  | 48 +++++++++++++++++++
 tensor2tensor/data_generators/image_utils.py  | 43 ++++++++++++++++-
 .../data_generators/image_utils_test.py       | 42 ++++++++++++++++
 tensor2tensor/data_generators/imagenet.py     | 23 +++++----
 .../data_generators/imagenet_test.py          | 47 ++++++++++++++++++
 tensor2tensor/data_generators/mscoco.py       | 18 +++++--
 tensor2tensor/data_generators/mscoco_test.py  | 47 ++++++++++++++++++
 9 files changed, 269 insertions(+), 23 deletions(-)
 create mode 100644 tensor2tensor/data_generators/celeba_test.py
 create mode 100644 tensor2tensor/data_generators/imagenet_test.py
 create mode 100644 tensor2tensor/data_generators/mscoco_test.py

diff --git a/setup.py b/setup.py
index 6681f335e..c7f70b239 100644
--- a/setup.py
+++ b/setup.py
@@ -52,7 +52,7 @@
         'tensorflow': ['tensorflow>=1.5.0'],
         'tensorflow_gpu': ['tensorflow-gpu>=1.5.0'],
         'tests': [
-            'pytest', 'mock', 'pylint', 'jupyter', 'gsutil'
+            'absl-py', 'pytest', 'mock', 'pylint', 'jupyter', 'gsutil'
             # Need atari extras for Travis tests, but because gym is already in
             # install_requires, pip skips the atari extras, so we instead do an
             # explicit pip install gym[atari] for the tests.
diff --git a/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index 0cc1a6319..476e8f04c 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -157,22 +157,28 @@ def dataset_filename(self):
 
   def preprocess_example(self, example, mode, hparams):
     image = example["inputs"]
-    if hasattr(hparams, "resize_method"):
-      method = getattr(tf.image.ResizeMethod, hparams.resize_method)
-    else:  # default
-      method = tf.image.ResizeMethod.BICUBIC
+    # Get resize method. Include a default if not specified, or if it's not in
+    # TensorFlow's collection of pre-implemented resize methods.
+    resize_method = getattr(hparams, "resize_method", "BICUBIC")
+    resize_method = getattr(tf.image.ResizeMethod, resize_method, resize_method)
 
     # Remove boundaries in CelebA images. Remove 40 pixels each side
     # vertically and 20 pixels each side horizontally.
     image = tf.image.crop_to_bounding_box(image, 40, 20, 218 - 80, 178 - 40)
 
-    scaled_images = image_utils.make_multiscale(
-        image, hparams.resolutions,
-        resize_method=method, num_channels=self.num_channels)
+    highest_res = hparams.resolutions[-1]
+    if resize_method == "DILATED":
+      # Resize image so that dilated subsampling is properly divisible.
+      scaled_image = image_utils.resize_by_area(image, highest_res)
+      scaled_images = image_utils.make_multiscale_dilated(
+          scaled_image, hparams.resolutions, num_channels=self.num_channels)
+    else:
+      scaled_images = image_utils.make_multiscale(
+          image, hparams.resolutions,
+          resize_method=resize_method, num_channels=self.num_channels)
 
     # Pack tuple of scaled images into one tensor. We do this by enforcing the
     # columns to match for every resolution.
-    highest_res = hparams.resolutions[-1]
     example["inputs"] = image
     example["targets"] = tf.concat([
         tf.reshape(scaled_image,
diff --git a/tensor2tensor/data_generators/celeba_test.py b/tensor2tensor/data_generators/celeba_test.py
new file mode 100644
index 000000000..f9fb2b999
--- /dev/null
+++ b/tensor2tensor/data_generators/celeba_test.py
@@ -0,0 +1,48 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 CelebA."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+from tensor2tensor.data_generators import celeba
+
+import tensorflow as tf
+
+
+class CelebaTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.named_parameters(
+      ("Default", None),
+      ("Area", "AREA"),
+      ("Dilated", "DILATED"))
+  def testCelebaMultiResolutionPreprocessExample(self, resize_method):
+    example = {"inputs": tf.random_uniform([218, 178, 3], minval=-1.)}
+    mode = tf.estimator.ModeKeys.TRAIN
+    hparams = tf.contrib.training.HParams(resolutions=[8, 16, 32])
+    if resize_method is not None:
+      hparams.resize_method = resize_method
+
+    problem = celeba.ImageCelebaMultiResolution()
+    preprocessed_example = problem.preprocess_example(example, mode, hparams)
+    self.assertLen(preprocessed_example, 2)
+    self.assertEqual(preprocessed_example["inputs"].shape, (138, 138, 3))
+    self.assertEqual(preprocessed_example["targets"].shape, (42, 32, 3))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index e3088ac9a..9f4b6b0c9 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -22,6 +22,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -37,7 +38,18 @@ def resize_by_area(img, size):
 def make_multiscale(image, resolutions,
                     resize_method=tf.image.ResizeMethod.BICUBIC,
                     num_channels=3):
-  """Returns list of scaled images, one for each resolution."""
+  """Returns list of scaled images, one for each resolution.
+
+  Args:
+    image: Tensor of shape [height, height, num_channels].
+    resolutions: List of heights that image's height is resized to.
+    resize_method: tf.image.ResizeMethod.
+    num_channels: Number of channels in image.
+
+  Returns:
+    List of Tensors, one for each resolution with shape given by
+    [resolutions[i], resolutions[i], num_channels].
+  """
   scaled_images = []
   for height in resolutions:
     scaled_image = tf.image.resize_images(
@@ -51,6 +63,35 @@ def make_multiscale(image, resolutions,
   return scaled_images
 
 
+def make_multiscale_dilated(image, resolutions, num_channels=3):
+  """Returns list of scaled images, one for each resolution.
+
+  Resizes by skipping every nth pixel.
+
+  Args:
+    image: Tensor of shape [height, height, num_channels].
+    resolutions: List of heights that image's height is resized to. The function
+      assumes VALID padding, so the original image's height must be divisible
+      by each resolution's height to return the exact resolution size.
+    num_channels: Number of channels in image.
+
+  Returns:
+    List of Tensors, one for each resolution with shape given by
+    [resolutions[i], resolutions[i], num_channels] if resolutions properly
+    divide the original image's height; otherwise shape height and width is up
+    to valid skips.
+  """
+  image_height = common_layers.shape_list(image)[0]
+  scaled_images = []
+  for height in resolutions:
+    dilation_rate = image_height // height  # assuming height = width
+    scaled_image = image[::dilation_rate, ::dilation_rate]
+    scaled_image = tf.to_int64(scaled_image)
+    scaled_image.set_shape([None, None, num_channels])
+    scaled_images.append(scaled_image)
+  return scaled_images
+
+
 class ImageProblem(problem.Problem):
   """Base class for problems with images."""
 
diff --git a/tensor2tensor/data_generators/image_utils_test.py b/tensor2tensor/data_generators/image_utils_test.py
index c47992f12..c8f8ca47e 100644
--- a/tensor2tensor/data_generators/image_utils_test.py
+++ b/tensor2tensor/data_generators/image_utils_test.py
@@ -70,6 +70,48 @@ def testImageGenerator(self):
       decoded2 = sess.run(decoded_png_t, feed_dict={image_t: encoded_img2[0]})
       self.assertAllClose(decoded2, image2)
 
+  def testMakeMultiscaleDivisible(self):
+    image = tf.random_normal([256, 256, 3])
+    resolutions = [8, 16, 64, 256]
+    scaled_images = image_utils.make_multiscale(image, resolutions)
+    self.assertEqual(scaled_images[0].shape, (8, 8, 3))
+    self.assertEqual(scaled_images[1].shape, (16, 16, 3))
+    self.assertEqual(scaled_images[2].shape, (64, 64, 3))
+    self.assertEqual(scaled_images[3].shape, (256, 256, 3))
+
+  def testMakeMultiscaleIndivisible(self):
+    image = tf.random_normal([256, 256, 3])
+    resolutions = [255]
+    scaled_images = image_utils.make_multiscale(image, resolutions)
+    self.assertEqual(scaled_images[0].shape, (255, 255, 3))
+
+  def testMakeMultiscaleLarger(self):
+    image = tf.random_normal([256, 256, 3])
+    resolutions = [257]
+    scaled_images = image_utils.make_multiscale(image, resolutions)
+    self.assertEqual(scaled_images[0].shape, (257, 257, 3))
+
+  def testMakeMultiscaleDilatedDivisible(self):
+    image = tf.random_normal([256, 256, 3])
+    resolutions = [8, 16, 64, 256]
+    scaled_images = image_utils.make_multiscale_dilated(image, resolutions)
+    self.assertEqual(scaled_images[0].shape, (8, 8, 3))
+    self.assertEqual(scaled_images[1].shape, (16, 16, 3))
+    self.assertEqual(scaled_images[2].shape, (64, 64, 3))
+    self.assertEqual(scaled_images[3].shape, (256, 256, 3))
+
+  def testMakeMultiscaleDilatedIndivisible(self):
+    image = tf.random_normal([256, 256, 3])
+    resolutions = [255]
+    scaled_images = image_utils.make_multiscale_dilated(image, resolutions)
+    self.assertEqual(scaled_images[0].shape, (256, 256, 3))
+
+  def testMakeMultiscaleDilatedLarger(self):
+    image = tf.random_normal([256, 256, 3])
+    resolutions = [257]
+    with self.assertRaisesRegexp(ValueError, "strides.* must be non-zero"):
+      _ = image_utils.make_multiscale_dilated(image, resolutions)
+
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 5735244ce..7d569bbd8 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -240,21 +240,24 @@ def dev_shards(self):
 
   def preprocess_example(self, example, mode, hparams):
     image = example["inputs"]
+    # Get resize method. Include a default if not specified, or if it's not in
+    # TensorFlow's collection of pre-implemented resize methods.
+    resize_method = getattr(hparams, "resize_method", "BICUBIC")
+    resize_method = getattr(tf.image.ResizeMethod, resize_method, resize_method)
+
+    if resize_method == "DILATED":
+      scaled_images = image_utils.make_multiscale_dilated(
+          image, hparams.resolutions, num_channels=self.num_channels)
+    else:
+      scaled_images = image_utils.make_multiscale(
+          image, hparams.resolutions,
+          resize_method=resize_method, num_channels=self.num_channels)
 
-    if hasattr(hparams, "resize_method"):
-      method = getattr(tf.image.ResizeMethod, hparams.resize_method)
-    else:  # default
-      method = tf.image.ResizeMethod.BICUBIC
-
-    scaled_images = image_utils.make_multiscale(
-        image, hparams.resolutions,
-        resize_method=method, num_channels=self.num_channels)
-
-    highest_res = hparams.resolutions[-1]
     # Pack tuple of scaled images into one tensor. We do this by enforcing the
     # columns to match for every resolution.
     # TODO(avaswani, trandustin): We should create tuples because this will not
     # work if height*width of low res < width of high res
+    highest_res = hparams.resolutions[-1]
     example["inputs"] = tf.concat([
         tf.reshape(scaled_image,
                    [res**2 // highest_res, highest_res, self.num_channels])
diff --git a/tensor2tensor/data_generators/imagenet_test.py b/tensor2tensor/data_generators/imagenet_test.py
new file mode 100644
index 000000000..eeb70f0a0
--- /dev/null
+++ b/tensor2tensor/data_generators/imagenet_test.py
@@ -0,0 +1,47 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 ImageNet."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+from tensor2tensor.data_generators import imagenet
+
+import tensorflow as tf
+
+
+class ImagenetTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.named_parameters(
+      ("Default", None),
+      ("Area", "AREA"),
+      ("Dilated", "DILATED"))
+  def testImagenetMultiResolutionPreprocessExample(self, resize_method):
+    example = {"inputs": tf.random_uniform([64, 64, 3], minval=-1.)}
+    mode = tf.estimator.ModeKeys.TRAIN
+    hparams = tf.contrib.training.HParams(resolutions=[8, 16, 32])
+    if resize_method is not None:
+      hparams.resize_method = resize_method
+
+    problem = imagenet.ImageImagenetMultiResolutionGen()
+    preprocessed_example = problem.preprocess_example(example, mode, hparams)
+    self.assertLen(preprocessed_example, 1)
+    self.assertEqual(preprocessed_example["inputs"].shape, (42, 32, 3))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index acb8bc697..28a1ea935 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -225,12 +225,24 @@ def dataset_filename(self):
 
   def preprocess_example(self, example, mode, hparams):
     image = example["inputs"]
-    scaled_images = image_utils.make_multiscale(
-        image, hparams.resolutions, num_channels=self.num_channels)
+    # Get resize method. Include a default if not specified, or if it's not in
+    # TensorFlow's collection of pre-implemented resize methods.
+    resize_method = getattr(hparams, "resize_method", "BICUBIC")
+    resize_method = getattr(tf.image.ResizeMethod, resize_method, resize_method)
+
+    highest_res = hparams.resolutions[-1]
+    if resize_method == "DILATED":
+      # Resize image so that dilated subsampling is properly divisible.
+      scaled_image = image_utils.resize_by_area(image, highest_res)
+      scaled_images = image_utils.make_multiscale_dilated(
+          scaled_image, hparams.resolutions, num_channels=self.num_channels)
+    else:
+      scaled_images = image_utils.make_multiscale(
+          image, hparams.resolutions,
+          resize_method=resize_method, num_channels=self.num_channels)
 
     # Pack tuple of scaled images into one tensor. We do this by enforcing the
     # columns to match for every resolution.
-    highest_res = hparams.resolutions[-1]
     example["inputs"] = tf.concat([
         tf.reshape(scaled_image,
                    [res**2 // highest_res, highest_res, self.num_channels])
diff --git a/tensor2tensor/data_generators/mscoco_test.py b/tensor2tensor/data_generators/mscoco_test.py
new file mode 100644
index 000000000..bb76041c7
--- /dev/null
+++ b/tensor2tensor/data_generators/mscoco_test.py
@@ -0,0 +1,47 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 MS COCO."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+from tensor2tensor.data_generators import mscoco
+
+import tensorflow as tf
+
+
+class MscocoTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.named_parameters(
+      ("Default", None),
+      ("Area", "AREA"),
+      ("Dilated", "DILATED"))
+  def testMsCocoMultiResolutionPreprocessExample(self, resize_method):
+    example = {"inputs": tf.random_uniform([400, 400, 3], minval=-1.)}
+    mode = tf.estimator.ModeKeys.TRAIN
+    hparams = tf.contrib.training.HParams(resolutions=[8, 16, 32])
+    if resize_method is not None:
+      hparams.resize_method = resize_method
+
+    problem = mscoco.ImageTextMsCocoMultiResolution()
+    preprocessed_example = problem.preprocess_example(example, mode, hparams)
+    self.assertLen(preprocessed_example, 1)
+    self.assertEqual(preprocessed_example["inputs"].shape, (42, 32, 3))
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 7e8101295d35f008f4b32cd23b6208f53bb651b9 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 25 Jun 2018 22:55:35 +0200
Subject: [PATCH 0221/2720] python 2.x new class style fix

---
 tensor2tensor/data_generators/gym_problems.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 8c0bc75c5..9213701b5 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -236,7 +236,7 @@ class GymAEDiscreteProblem(GymDiscreteProblem):
   pass
 
 
-class BasicStatistics:
+class BasicStatistics(object):
   """Keeps basic statistics to calculate mean reward """
 
   def __init__(self):

From 4731446f5afc665c7f4ae09836289c01118379df Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Mon, 25 Jun 2018 13:56:36 -0700
Subject: [PATCH 0222/2720] Fixed some documentation.

PiperOrigin-RevId: 202013837
---
 .../layers/common_message_passing_attention.py      | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/tensor2tensor/layers/common_message_passing_attention.py b/tensor2tensor/layers/common_message_passing_attention.py
index a3fc38924..9db0bcc04 100644
--- a/tensor2tensor/layers/common_message_passing_attention.py
+++ b/tensor2tensor/layers/common_message_passing_attention.py
@@ -74,22 +74,9 @@ def multihead_graph_attention(query_antecedent,
     vars_3d: use 3-dimensional variables for input/output transformations
     **kwargs (dict): Parameters for the attention function
 
-  Caching:
-    WARNING: For decoder self-attention, i.e. when memory_antecedent == None,
-    the caching assumes that the bias contains future masking.
-
-    The caching works by saving all the previous key and value values so that
-    you are able to send just the last query location to this attention
-    function. I.e. if the cache dict is provided it assumes the query is of the
-    shape [batch_size, 1, hiddem_dim] rather than the full memory.
-
   Returns:
     The result of the attention transformation. The output shape is
         [batch_size, length_q, hidden_dim]
-    unless the cache dict is provided in which case only the last memory
-    position is calculated and the output shape is [batch_size, 1, hidden_dim]
-    Optionally returns an additional loss parameters (ex: load balance loss for
-    the experts) returned by the attention_type function.
 
   Raises:
     ValueError: if the key depth or value depth are not divisible by the

From d424626d6ca6bb4ce140100107ee7f95aa17460a Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Mon, 25 Jun 2018 14:15:06 -0700
Subject: [PATCH 0223/2720] Bug fix in dml loss. Was doing the wrong
 normalization.

PiperOrigin-RevId: 202017352
---
 tensor2tensor/layers/common_layers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 3c5f108f1..a865a86a7 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1948,7 +1948,7 @@ def dml_loss(
     a pair of tensors of loss/sum of losses, denominator
   """
   del weights_fn  # Unused
-  real_labels = convert_rgb_to_real(labels)
+  real_labels = convert_rgb_to_symmetric_real(labels)
   dml_loss_value = discretized_mix_logistic_loss(real_labels, pred,
                                                  sum_all=reduce_sum)
   if reduce_sum:

From e020dc0b9b850cc1f54fb30f37a71cdbd55bcb69 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 25 Jun 2018 14:29:22 -0700
Subject: [PATCH 0224/2720] Video 2 gif converter

PiperOrigin-RevId: 202019907
---
 tensor2tensor/utils/video2gif.py | 92 ++++++++++++++++++++++++++++++++
 1 file changed, 92 insertions(+)
 create mode 100644 tensor2tensor/utils/video2gif.py

diff --git a/tensor2tensor/utils/video2gif.py b/tensor2tensor/utils/video2gif.py
new file mode 100644
index 000000000..e76acb215
--- /dev/null
+++ b/tensor2tensor/utils/video2gif.py
@@ -0,0 +1,92 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+r"""View the problem.
+
+This binary saves the videos in the problem(dataset) into gifs.
+
+The imagemagick package should be installed for conversion to gifs.
+
+Example usage to view dataset:
+
+  video2gif \
+      --data_dir ~/data \
+      --problem=gym_water_world_random5k \
+      --hparams_set=next_frame_stochastic \
+      --output_dir /usr/local/google/home/mbz/t2t_train/ww/ \
+      --data_dir /usr/local/google/home/mbz/temp/ \
+      --num_samples 10
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import sys
+import numpy as np
+from tensor2tensor.bin import t2t_trainer          # pylint: disable=unused-import
+from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
+from tensor2tensor.utils import decoding
+from tensor2tensor.utils import trainer_lib
+
+import tensorflow as tf
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+flags.DEFINE_integer("num_samples", -1, "Number of saved samples.")
+
+
+def create_gif(name):
+  cmd = "convert -delay 15 {0}* {0}.gif".format(name)
+  os.system(cmd)
+
+
+def main(_):
+  problem_name = FLAGS.problem
+  if "video" not in problem_name and "gym" not in problem_name:
+    print("This tool only works for video problems.")
+    return
+
+  mode = tf.estimator.ModeKeys.TRAIN
+  hparams = trainer_lib.create_hparams(
+      FLAGS.hparams_set,
+      FLAGS.hparams,
+      data_dir=os.path.expanduser(FLAGS.data_dir),
+      problem_name=problem_name)
+
+  dataset = hparams.problem.input_fn(mode, hparams)
+  features = dataset.make_one_shot_iterator().get_next()
+
+  tf.gfile.MakeDirs(FLAGS.output_dir)
+  base_template = os.path.join(FLAGS.output_dir, FLAGS.problem)
+  count = 0
+  with tf.train.MonitoredTrainingSession() as sess:
+    while not sess.should_stop():
+      # TODO(mbz): figure out what the second output is.
+      data, _ = sess.run(features)
+      video_batch = np.concatenate((data["inputs"], data["targets"]), axis=1)
+
+      for video in video_batch:
+        print("Saving {}/{}".format(count, FLAGS.num_samples))
+        name = "%s_%05d" % (base_template, count)
+        decoding.save_video(video, name + "_{}.png")
+        create_gif(name)
+        count += 1
+
+        if count == FLAGS.num_samples:
+          sys.exit(0)
+
+if __name__ == "__main__":
+  tf.app.run()

From b0f8951427af83b61eb8a7fb7a474261674e8d3d Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 26 Jun 2018 00:02:36 +0200
Subject: [PATCH 0225/2720] linter fixes

---
 tensor2tensor/rl/envs/batch_env_factory.py | 4 +++-
 tensor2tensor/rl/envs/tf_atari_wrappers.py | 6 +-----
 tensor2tensor/rl/model_rl_experiment.py    | 6 +++---
 tensor2tensor/rl/rl_trainer_lib.py         | 2 +-
 tensor2tensor/rl/rl_trainer_lib_test.py    | 4 ++--
 5 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index e692ed9d3..26ac6246a 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -69,7 +69,9 @@ def _define_batch_env(environment_spec, num_agents, xvfb=False):
 
 def _define_simulated_batch_env(environment_spec, num_agents,
                                 other_hparms):
-  cur_batch_env = simulated_batch_env.SimulatedBatchEnv(environment_spec, num_agents, other_hparms)
+  cur_batch_env = simulated_batch_env.SimulatedBatchEnv(environment_spec,
+                                                        num_agents,
+                                                        other_hparms)
   return cur_batch_env
 
 
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index e966836fd..6119aa657 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -32,11 +32,7 @@ def __init__(self, batch_env):
     self.action_dtype = batch_env.action_dtype
 
   def initialize(self, sess):
-    """
-    Initializations to be run once the tf.Session is available 
-    
-    sess - tf.Session 
-    """
+    """Initializations to be run once the tf.Session is available."""
     pass
 
   @property
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 16e16e3fa..3394771e7 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -36,7 +36,6 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.layers import discretization
 from tensor2tensor.rl import rl_trainer_lib
-from tensor2tensor.rl.envs.tf_atari_wrappers import TimeLimitWrapper
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
@@ -189,7 +188,8 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
     gym_simulated_problem.generate_data(epoch_data_dir, tmp_dir)
   n = max(1., gym_simulated_problem.statistics.number_of_dones)
   model_reward_accuracy = (
-      gym_simulated_problem.statistics.successful_episode_reward_predictions / float(n))
+      gym_simulated_problem.statistics.successful_episode_reward_predictions
+      / float(n))
   return model_reward_accuracy
 
 
@@ -388,7 +388,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     if hparams.eval_world_model:
       log("Evaluating world model")
       model_reward_accuracy = evaluate_world_model(
-        simulated_problem_name, world_model_problem, hparams,
+          simulated_problem_name, world_model_problem, hparams,
           directories["world_model"],
           epoch_data_dir, directories["tmp"],
           autoencoder_path=autoencoder_model_dir)
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index f1935998a..1d0e0b057 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -34,7 +34,7 @@ def define_train(hparams, event_dir):
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
     memory, collect_summary, initialization\
       = collect.define_collect(
-        hparams, "ppo_train", eval_phase=False)
+          hparams, "ppo_train", eval_phase=False)
     ppo_summary = ppo.define_ppo_epoch(memory, hparams)
     summary = tf.summary.merge([collect_summary, ppo_summary])
 
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index 39ff5f278..bfe143050 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -55,7 +55,7 @@ def test_no_crash_cartpole(self):
   # is done on the ppo code
   #
   # To run the test change epochs_num=2 to epoch_num=200
-  # and epoch_length=5 to epoch_length=200
+  # and epoch_length=4 to epoch_length=200
   # (it is set like that to meet travis timeouts
   def test_train_pong(self):
     hparams = tf.contrib.training.\
@@ -63,7 +63,7 @@ def test_train_pong(self):
               eval_every_epochs=10,
               num_agents=20,
               optimization_epochs=3,
-              epoch_length=5,
+              epoch_length=4,
               entropy_loss_coef=0.003,
               learning_rate=8e-05,
               optimizer="Adam",

From b7b82678378a3cb7f77b547990132348202fc2f9 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 25 Jun 2018 16:16:24 -0700
Subject: [PATCH 0226/2720] Make the dataset improvement compatible with older
 TF versions which we support in the open-source part.

PiperOrigin-RevId: 202038400
---
 tensor2tensor/data_generators/problem.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index b90f77689..5b0186d8e 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -574,7 +574,17 @@ def _load_records_and_preprocess(filenames):
           tf.contrib.data.parallel_interleave(
               _load_records_and_preprocess, sloppy=True, cycle_length=8))
     else:
-      dataset = _load_records_and_preprocess(dataset)
+      # TFRecordDataset can get filenames as dataset in TF 1.7+.
+      # TODO(lukaszkaiser): remove when we require TF 1.7+ in general.
+      major, minor = [int(el) for el in tf.__version__.split(".")[:2]]
+      filename_dataset_ok = major > 1 or (major == 1 and minor >= 7)
+      if filename_dataset_ok:  # We can just pass a Dataset of filenames.
+        dataset = _load_records_and_preprocess(dataset)
+      else:  # Go file-by-file (can be very slow).
+        dataset = None
+        for f in data_files:
+          f_data = _load_records_and_preprocess(f)
+          dataset = f_data if dataset is None else dataset.concatenate(f_data)
 
     dataset = dataset.map(
         self.maybe_reverse_and_copy, num_parallel_calls=num_threads)

From b90bca34a98ca5712dd5a09bda98de7aa45ffaed Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 25 Jun 2018 17:02:11 -0700
Subject: [PATCH 0227/2720] Separating the stochastic experiment test.

PiperOrigin-RevId: 202045656
---
 .../rl/model_rl_experiment_stochastic_test.py | 38 +++++++++++++++++++
 tensor2tensor/rl/model_rl_experiment_test.py  |  7 ----
 2 files changed, 38 insertions(+), 7 deletions(-)
 create mode 100644 tensor2tensor/rl/model_rl_experiment_stochastic_test.py

diff --git a/tensor2tensor/rl/model_rl_experiment_stochastic_test.py b/tensor2tensor/rl/model_rl_experiment_stochastic_test.py
new file mode 100644
index 000000000..51f207364
--- /dev/null
+++ b/tensor2tensor/rl/model_rl_experiment_stochastic_test.py
@@ -0,0 +1,38 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Tiny run of model_rl_experiment. Smoke test."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.rl import model_rl_experiment
+
+import tensorflow as tf
+
+FLAGS = tf.flags.FLAGS
+
+
+class ModelRLExperimentStochasticTest(tf.test.TestCase):
+
+  def test_stochastic(self):
+    FLAGS.output_dir = tf.test.get_temp_dir()
+    FLAGS.loop_hparams_set = "rl_modelrl_tiny_stochastic"
+    FLAGS.loop_hparams = "generative_model_params=next_frame_stochastic_tiny"
+    FLAGS.schedule = "train"  # skip evaluation for world model training
+    model_rl_experiment.main(None)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/rl/model_rl_experiment_test.py b/tensor2tensor/rl/model_rl_experiment_test.py
index e2430983c..90388d161 100644
--- a/tensor2tensor/rl/model_rl_experiment_test.py
+++ b/tensor2tensor/rl/model_rl_experiment_test.py
@@ -33,13 +33,6 @@ def test_basic(self):
     FLAGS.schedule = "train"  # skip evaluation for world model training
     model_rl_experiment.main(None)
 
-  def test_stochastic(self):
-    FLAGS.output_dir = tf.test.get_temp_dir()
-    FLAGS.loop_hparams_set = "rl_modelrl_tiny_stochastic"
-    FLAGS.loop_hparams = "generative_model_params=next_frame_stochastic_tiny"
-    FLAGS.schedule = "train"  # skip evaluation for world model training
-    model_rl_experiment.main(None)
-
 
 if __name__ == "__main__":
   tf.test.main()

From 77ff77f119b1c323f344b349ee983b0a34e21131 Mon Sep 17 00:00:00 2001
From: Sanyam Kapoor <sanyamkapoor@google.com>
Date: Mon, 25 Jun 2018 18:19:01 -0700
Subject: [PATCH 0228/2720] Setup skeleton for new problem and model

---
 .../gh_function_docstring_encoder.py          | 43 +++++++++++++++++++
 1 file changed, 43 insertions(+)
 create mode 100644 tensor2tensor/data_generators/gh_function_docstring_encoder.py

diff --git a/tensor2tensor/data_generators/gh_function_docstring_encoder.py b/tensor2tensor/data_generators/gh_function_docstring_encoder.py
new file mode 100644
index 000000000..bc5843d3a
--- /dev/null
+++ b/tensor2tensor/data_generators/gh_function_docstring_encoder.py
@@ -0,0 +1,43 @@
+import os
+from tensor2tensor.utils import t2t_model
+from tensor2tensor.utils import registry
+from tensor2tensor.data_generators import text_problems
+
+
+@registry.register_model
+class SimilarityTransformer(t2t_model.T2TModel):
+  # pylint: disable=abstract-method
+
+  """
+  This class defines the model to compute similarity scores between functions and
+  docstrings
+  """
+
+  def __init__(self, *args, **kwargs):
+    super(SimilarityTransformer, self).__init__(*args, **kwargs)
+
+
+  def body(self, features):
+    # TODO: need to fill this with Transformer encoder/decoder
+    # and loss calculation
+    raise NotImplementedError
+
+
+@registry.register_problem
+class GithubFunctionDocstring(text_problems.Text2TextProblem):
+  # pylint: disable=abstract-method
+
+  """This class defines the problem of finding similarity between Python function
+   and docstring"""
+
+  @property
+  def is_generate_per_split(self):
+    return False
+
+  def generate_samples(self, data_dir, _tmp_dir, dataset_split):  #pylint: disable=no-self-use
+    """This method returns the generator to return {"inputs": [text], "targets": [text]} dict"""
+
+    functions_file_path = os.path.join(data_dir, '{}.function'.format(dataset_split))
+    docstrings_file_path = os.path.join(data_dir, '{}.docstring'.format(dataset_split))
+
+    return text_problems.text2text_txt_iterator(functions_file_path, docstrings_file_path)

From db558a587185811f2e904ffdc8fd64a2f011f5e8 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 26 Jun 2018 00:12:21 -0700
Subject: [PATCH 0229/2720] Adding frame_number to Gym videos.

PiperOrigin-RevId: 202084013
---
 tensor2tensor/data_generators/gym_problems.py | 24 +++++++++++++------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index f6558cf6a..a4581773a 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -116,13 +116,16 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
         self.sum_of_rewards += reward
         self.dones += int(done)
 
-        ret_dict = {"frame": observ,
-                    "image/format": ["png"],
-                    "image/height": [self.frame_height],
-                    "image/width": [self.frame_width],
-                    "action": [int(action)],
-                    "done": [int(False)],
-                    "reward": [int(reward) - self.min_reward]}
+        ret_dict = {
+            "frame": observ,
+            "frame_number": [int(pieces_generated)],
+            "image/format": ["png"],
+            "image/height": [self.frame_height],
+            "image/width": [self.frame_width],
+            "action": [int(action)],
+            "done": [int(False)],
+            "reward": [int(reward) - self.min_reward]
+        }
 
         yield ret_dict
         pieces_generated += 1
@@ -144,10 +147,13 @@ def eval_metrics(self):
   def extra_reading_spec(self):
     """Additional data fields to store on disk and their decoders."""
     data_fields = {
+        "frame_number": tf.FixedLenFeature([1], tf.int64),
         "action": tf.FixedLenFeature([1], tf.int64),
         "reward": tf.FixedLenFeature([1], tf.int64)
     }
     decoders = {
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
+            tensor_key="frame_number"),
         "action": tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="action"),
         "reward": tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="reward"),
     }
@@ -204,6 +210,10 @@ def min_reward(self):
   def num_testing_steps(self):
     return None
 
+  @property
+  def only_keep_videos_from_0th_frame(self):
+    return False
+
   def get_action(self, observation=None):
     del observation
     return self.env.action_space.sample()

From f0dc6b7c4c938a0b4d850b67695ae15fb25ec929 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 26 Jun 2018 09:44:05 -0700
Subject: [PATCH 0230/2720] internal merge of PR #893

PiperOrigin-RevId: 202142951
---
 .../gh_function_docstring_encoder.py          | 52 +++++++++++--------
 1 file changed, 30 insertions(+), 22 deletions(-)

diff --git a/tensor2tensor/data_generators/gh_function_docstring_encoder.py b/tensor2tensor/data_generators/gh_function_docstring_encoder.py
index bc5843d3a..784c237f3 100644
--- a/tensor2tensor/data_generators/gh_function_docstring_encoder.py
+++ b/tensor2tensor/data_generators/gh_function_docstring_encoder.py
@@ -1,43 +1,51 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Github function to text similatrity problems."""
+
 import os
-from tensor2tensor.utils import t2t_model
-from tensor2tensor.utils import registry
+
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
 
 
 @registry.register_model
 class SimilarityTransformer(t2t_model.T2TModel):
-  # pylint: disable=abstract-method
-
-  """
-  This class defines the model to compute similarity scores between functions and
-  docstrings
-  """
-
-  def __init__(self, *args, **kwargs):
-    super(SimilarityTransformer, self).__init__(*args, **kwargs)
-
+  """Similarity scores between functions and docstrings."""
 
   def body(self, features):
-    # TODO: need to fill this with Transformer encoder/decoder
+    # TODO(sanyamkapoor): need to fill this with Transformer encoder/decoder
     # and loss calculation
     raise NotImplementedError
 
 
 @registry.register_problem
 class GithubFunctionDocstring(text_problems.Text2TextProblem):
-  # pylint: disable=abstract-method
-
-  """This class defines the problem of finding similarity between Python function
-   and docstring"""
+  """The problem of similarity between Python function and docstring."""
 
   @property
   def is_generate_per_split(self):
     return False
 
-  def generate_samples(self, data_dir, _tmp_dir, dataset_split):  #pylint: disable=no-self-use
-    """This method returns the generator to return {"inputs": [text], "targets": [text]} dict"""
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    """Returns the generator of {"inputs": [text], "targets": [text]} dict."""
 
-    functions_file_path = os.path.join(data_dir, '{}.function'.format(dataset_split))
-    docstrings_file_path = os.path.join(data_dir, '{}.docstring'.format(dataset_split))
+    functions_file_path = os.path.join(
+        data_dir, '{}.function'.format(dataset_split))
+    docstrings_file_path = os.path.join(
+        data_dir, '{}.docstring'.format(dataset_split))
 
-    return text_problems.text2text_txt_iterator(functions_file_path, docstrings_file_path)
+    return text_problems.text2text_txt_iterator(
+        functions_file_path, docstrings_file_path)

From d5b4385ab9df72c10b08e68141eaa4707ee9d3b3 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 26 Jun 2018 09:47:04 -0700
Subject: [PATCH 0231/2720] Allow to control the GAN loss mixing instead of
 just binary having it or not.

PiperOrigin-RevId: 202143522
---
 setup.py                                      |  2 +-
 tensor2tensor/models/research/autoencoders.py | 12 +++++++-----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/setup.py b/setup.py
index c7f70b239..6681f335e 100644
--- a/setup.py
+++ b/setup.py
@@ -52,7 +52,7 @@
         'tensorflow': ['tensorflow>=1.5.0'],
         'tensorflow_gpu': ['tensorflow-gpu>=1.5.0'],
         'tests': [
-            'absl-py', 'pytest', 'mock', 'pylint', 'jupyter', 'gsutil'
+            'pytest', 'mock', 'pylint', 'jupyter', 'gsutil'
             # Need atari extras for Travis tests, but because gym is already in
             # install_requires, pip skips the atari extras, so we instead do an
             # explicit pip install gym[atari] for the tests.
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 8b5d9f197..302d93c69 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -146,7 +146,7 @@ def body(self, features):
       self._cur_bottleneck_tensor = b
       b = self.unbottleneck(b, common_layers.shape_list(x)[-1])
       b = common_layers.mix(b, x, hparams.bottleneck_warmup_steps, is_training)
-      if hparams.add_gan_loss:
+      if hparams.gan_loss_factor != 0.0:
         # Add a purely sampled batch on which we'll compute the GAN loss.
         g = self.unbottleneck(self.sample(), common_layers.shape_list(x)[-1],
                               reuse=True)
@@ -173,7 +173,7 @@ def body(self, features):
     res = x[:, :shape[1], :shape[2], :]
     # Add GAN loss if requested.
     gan_loss = 0.0
-    if hparams.add_gan_loss:
+    if hparams.gan_loss_factor != 0.0:
       # Split back if we added a purely sampled batch.
       res_gan, res = tf.split(res, 2, axis=0)
       num_channels = self.hparams.problem.num_channels
@@ -186,8 +186,7 @@ def discriminate(x):
       gan_loss = common_layers.sliced_gan_loss(
           orig_rgb, reverse_gradient(res_rgb),
           discriminate, self.hparams.num_sliced_vecs)
-      gan_loss *= common_layers.inverse_lin_decay(
-          hparams.bottleneck_warmup_steps)
+      gan_loss *= hparams.gan_loss_factor
     # Mix the final result and return.
     res = common_layers.mix(res, features["targets"],
                             hparams.bottleneck_warmup_steps // 2, is_training)
@@ -692,7 +691,7 @@ def autoencoder_basic():
   hparams.add_hparam("sample_width", 32)
   hparams.add_hparam("discriminator_batchnorm", True)
   hparams.add_hparam("num_sliced_vecs", 4096)
-  hparams.add_hparam("add_gan_loss", False)
+  hparams.add_hparam("gan_loss_factor", 0.0)
   return hparams
 
 
@@ -778,6 +777,9 @@ def autoencoder_ordered_discrete():
   """Ordered discrete autoencoder model."""
   hparams = autoencoder_residual_discrete()
   hparams.bottleneck_noise = 1.0
+  hparams.gan_loss_factor = 0.0
+  hparams.dropout = 0.1
+  hparams.residual_dropout = 0.3
   hparams.add_hparam("unordered", False)
   return hparams
 

From 50b1103bd8063255616c82ddabcb234b1b96554d Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 26 Jun 2018 19:05:20 +0200
Subject: [PATCH 0232/2720] RewardPerSequenceStatistics implementation

---
 tensor2tensor/data_generators/gym_problems.py | 44 ++++++++++++++++---
 tensor2tensor/data_generators/video_utils.py  |  2 +-
 tensor2tensor/rl/envs/simulated_batch_env.py  | 15 ++++---
 tensor2tensor/rl/model_rl_experiment.py       | 10 ++---
 4 files changed, 51 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 9213701b5..69d73f8c3 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -25,6 +25,8 @@
 
 # We need gym_utils for the game environments defined there.
 from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
+
+from tensorflow.contrib.training import HParams
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.models.research import rl
@@ -76,6 +78,7 @@ def __init__(self, *args, **kwargs):
 
     self._internal_memory_size = 20
     self._internal_memory_force_beginning_resets = False
+    self._session = None
 
   def _setup(self):
     collect_hparams = rl.ppo_pong_base()
@@ -89,18 +92,24 @@ def _setup(self):
       collect_hparams.policy_network = rl.random_policy_fun
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      self.collect_memory, self.collect_trigger_op, self.collect_init \
+      self.collect_memory, self.collect_trigger_op, collect_init \
         = collect.define_collect(collect_hparams, scope="gym_problems",
                                  collect_level=0, eval_phase=self.eval_phase)
 
+    self._session = tf.Session()
+    collect_init(self._session)
+    self._session.run(tf.global_variables_initializer())
+
+  @property
+  def random_skip(self):
+    return False
+
   def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
     self._setup()
     self.debug_dump_frames_path = os.path.join(
         data_dir, self.debug_dump_frames_path)
 
-    with tf.Session() as sess:
-      sess.run(tf.global_variables_initializer())
-      self.collect_init(sess)
+    with self._session as sess:
       self.restore_networks(sess)
       pieces_generated = 0
       memory_index = 0
@@ -112,6 +121,7 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
         data = [memory[i][memory_index][0] for i in range(4)]
         memory_index += 1
         observation, reward, done, action = data
+        print("Data:{}".format(data[1:]))
         observation = observation.astype(np.uint8)
 
         debug_image = self.collect_statistics_and_generate_debug_image(pieces_generated,
@@ -305,9 +315,29 @@ def __init__(self, *args, **kwargs):
     real_env = env_spec.env_lambda()
     self.statistics.real_env = real_env
 
-
   def _setup(self):
     super(GymSimulatedDiscreteProblem, self)._setup()
+
+    environment_spec = self.environment_spec
+    hparams = HParams(video_num_input_frames=
+                      environment_spec.video_num_input_frames,
+                      video_num_target_frames=
+                      environment_spec.video_num_target_frames,
+                      environment_spec=environment_spec)
+
+    initial_frames_problem = environment_spec.initial_frames_problem
+    # initial_frames_problem.random_skip = False
+    dataset = initial_frames_problem.dataset(
+      tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
+      shuffle_files=False, hparams=hparams)
+    dataset = dataset.map(lambda x: x["input_action"]).take(1)
+    input_data_iterator = (
+      dataset.batch(1).make_initializable_iterator())
+    self._session.run(input_data_iterator.initializer)
+
+    res = self._session.run(input_data_iterator.get_next())
+    print("Actions:{}".format(res))
+    self._initial_action = res[0, :, 0]
     self._reset_real_env()
 
   @property
@@ -350,8 +380,8 @@ def get_environment_spec(self):
   def _reset_real_env(self):
     stat = self.statistics
     stat.real_env.reset()
-    for _ in range(self.num_input_frames):
-      stat.real_ob, _, _, _ = stat.real_env.step(0)
+    for a in self._initial_action:
+      stat.real_ob, _, _, _ = stat.real_env.step(a)
 
     stat.episode_sim_reward = 0.0
     stat.episode_real_reward = 0.0
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index e56e75ce3..a42d33cd9 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -289,7 +289,7 @@ def check_integrity_and_batch(*datasets):
     else:
       batch_dataset = preprocessed_dataset.apply(
           tf.contrib.data.batch_and_drop_remainder(num_frames))
-    dataset = batch_dataset.map(features_from_batch).shuffle(8)
+    dataset = batch_dataset.map(features_from_batch)  #.shuffle(8)
     return dataset
 
   def eval_metrics(self):
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index cba46464c..e1ae418a7 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -23,6 +23,8 @@
 
 import copy
 
+from tensorflow.contrib.training import HParams
+
 from tensor2tensor.layers import common_layers
 from tensor2tensor.rl.envs import in_graph_batch_env
 from tensor2tensor.rl.envs.utils import get_action_space
@@ -113,12 +115,11 @@ def __init__(self, environment_spec, length, other_hparams):
 
     _, self.action_shape, self.action_dtype = get_action_space(environment_spec)
 
-    # TODO(lukaszkaiser): do this in a more cleaner way
-    # remove other_hparams
-    hparams = copy.copy(other_hparams)
-    hparams.video_num_input_frames, hparams.video_num_target_frames = (
-        hparams.environment_spec.video_num_input_frames,
-        hparams.environment_spec.video_num_target_frames)
+    hparams = HParams(video_num_input_frames=
+                      environment_spec.video_num_input_frames,
+                      video_num_target_frames=
+                      environment_spec.video_num_target_frames,
+                      environment_spec=environment_spec)
 
     if environment_spec.simulation_random_starts:
       dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
@@ -127,7 +128,7 @@ def __init__(self, environment_spec, length, other_hparams):
                                                hparams=hparams)
       dataset = dataset.shuffle(buffer_size=100)
     else:
-      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.EVAL,
+      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
                                                FLAGS.data_dir,
                                                shuffle_files=False,
                                                hparams=hparams).take(1)
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index ef3b0a364..7e69ddc1d 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -378,11 +378,11 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
                         autoencoder_model_dir, epoch_data_dir)
 
     # Train world model
-    log("Training world model")
-    train_world_model(world_model_problem, epoch_data_dir,
-                      directories["world_model"], hparams, epoch,
-                      use_autoencoder=using_autoencoder)
-
+    # log("Training world model")
+    # train_world_model(world_model_problem, epoch_data_dir,
+    #                   directories["world_model"], hparams, epoch,
+    #                   use_autoencoder=using_autoencoder)
+    #
     # Evaluate world model
     model_reward_accuracy = 0.
     if hparams.eval_world_model:

From c57dcb407210985a858a10e84218ba90a43ac439 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 26 Jun 2018 19:17:45 +0200
Subject: [PATCH 0233/2720] Minor clean-ups

---
 tensor2tensor/data_generators/gym_problems.py |  1 -
 tensor2tensor/rl/model_rl_experiment.py       | 10 +++++-----
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 05e49b092..c24888943 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -349,7 +349,6 @@ def _setup(self):
     self._session.run(input_data_iterator.initializer)
 
     res = self._session.run(input_data_iterator.get_next())
-    print("Actions:{}".format(res))
     self._initial_action = res[0, :, 0]
     self._reset_real_env()
 
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 7e69ddc1d..ef3b0a364 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -378,11 +378,11 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
                         autoencoder_model_dir, epoch_data_dir)
 
     # Train world model
-    # log("Training world model")
-    # train_world_model(world_model_problem, epoch_data_dir,
-    #                   directories["world_model"], hparams, epoch,
-    #                   use_autoencoder=using_autoencoder)
-    #
+    log("Training world model")
+    train_world_model(world_model_problem, epoch_data_dir,
+                      directories["world_model"], hparams, epoch,
+                      use_autoencoder=using_autoencoder)
+
     # Evaluate world model
     model_reward_accuracy = 0.
     if hparams.eval_world_model:

From f7a1f04552df9ad29520b66ace7dace4e2da2223 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 26 Jun 2018 19:18:53 +0200
Subject: [PATCH 0234/2720] minor clean-ups 2

---
 tensor2tensor/data_generators/gym_problems.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index c24888943..5938f413c 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -121,7 +121,6 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
         data = [memory[i][memory_index][0] for i in range(4)]
         memory_index += 1
         observation, reward, done, action = data
-        print("Data:{}".format(data[1:]))
         observation = observation.astype(np.uint8)
 
         debug_image = self.collect_statistics_and_generate_debug_image(pieces_generated,

From 647f4b2ee9780036ab6fb070a35783eea6c2c49f Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 26 Jun 2018 10:24:40 -0700
Subject: [PATCH 0235/2720] Exclude model_experiment tests from travis for now.

PiperOrigin-RevId: 202150199
---
 .travis.yml | 2 ++
 setup.py    | 2 +-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/.travis.yml b/.travis.yml
index ff0811c23..dc65b9b66 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -57,6 +57,7 @@ script:
   #   * trainer_lib_test
   #   * visualization_test
   #   * model_rl_experiment_test
+  #   * model_rl_experiment_test_stochastic
   # algorithmic_math_test: flaky
   # universal_transformer_test: requires new feature in tf.foldl (rm with TF 1.9)
   - pytest
@@ -67,6 +68,7 @@ script:
     --ignore=tensor2tensor/data_generators/algorithmic_math_test.py
     --ignore=tensor2tensor/models/research/universal_transformer_test.py
     --ignore=tensor2tensor/rl/model_rl_experiment_test.py
+    --ignore=tensor2tensor/rl/model_rl_experiment_test_stochastic.py
   - pytest tensor2tensor/utils/registry_test.py
   - pytest tensor2tensor/utils/trainer_lib_test.py
   - pytest tensor2tensor/visualization/visualization_test.py
diff --git a/setup.py b/setup.py
index 6681f335e..c7f70b239 100644
--- a/setup.py
+++ b/setup.py
@@ -52,7 +52,7 @@
         'tensorflow': ['tensorflow>=1.5.0'],
         'tensorflow_gpu': ['tensorflow-gpu>=1.5.0'],
         'tests': [
-            'pytest', 'mock', 'pylint', 'jupyter', 'gsutil'
+            'absl-py', 'pytest', 'mock', 'pylint', 'jupyter', 'gsutil'
             # Need atari extras for Travis tests, but because gym is already in
             # install_requires, pip skips the atari extras, so we instead do an
             # explicit pip install gym[atari] for the tests.

From f204e897e15c512336da1f6eda645a4899b92c92 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 26 Jun 2018 11:16:18 -0700
Subject: [PATCH 0236/2720] Typo in travis config.

PiperOrigin-RevId: 202160180
---
 .travis.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index dc65b9b66..2fb36ae47 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -57,7 +57,7 @@ script:
   #   * trainer_lib_test
   #   * visualization_test
   #   * model_rl_experiment_test
-  #   * model_rl_experiment_test_stochastic
+  #   * model_rl_experiment_stochastic_test
   # algorithmic_math_test: flaky
   # universal_transformer_test: requires new feature in tf.foldl (rm with TF 1.9)
   - pytest
@@ -68,7 +68,7 @@ script:
     --ignore=tensor2tensor/data_generators/algorithmic_math_test.py
     --ignore=tensor2tensor/models/research/universal_transformer_test.py
     --ignore=tensor2tensor/rl/model_rl_experiment_test.py
-    --ignore=tensor2tensor/rl/model_rl_experiment_test_stochastic.py
+    --ignore=tensor2tensor/rl/model_rl_experiment_stochastic_test.py
   - pytest tensor2tensor/utils/registry_test.py
   - pytest tensor2tensor/utils/trainer_lib_test.py
   - pytest tensor2tensor/visualization/visualization_test.py

From 03871e6fd08ccf65dd3ac16c4003a04f60ece6ad Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 26 Jun 2018 11:38:49 -0700
Subject: [PATCH 0237/2720] internal merge of PR #890

PiperOrigin-RevId: 202164726
---
 tensor2tensor/data_generators/gym_problems.py | 72 ++++++++++---------
 tensor2tensor/data_generators/video_utils.py  |  2 +-
 tensor2tensor/rl/collect.py                   | 11 +--
 tensor2tensor/rl/envs/simulated_batch_env.py  |  6 +-
 tensor2tensor/rl/model_rl_experiment.py       |  2 -
 tensor2tensor/rl/rl_trainer_lib.py            |  2 +-
 tensor2tensor/rl/rl_trainer_lib_test.py       |  7 +-
 7 files changed, 48 insertions(+), 54 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 5938f413c..e85d78a27 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -26,7 +26,6 @@
 # We need gym_utils for the game environments defined there.
 from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
 
-from tensorflow.contrib.training import HParams
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.models.research import rl
@@ -34,8 +33,11 @@
 from tensor2tensor.rl.envs import tf_atari_wrappers
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
+
 import tensorflow as tf
 
+from tensorflow.contrib.training import HParams
+
 
 flags = tf.flags
 FLAGS = flags.FLAGS
@@ -120,11 +122,11 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
           memory_index = 0
         data = [memory[i][memory_index][0] for i in range(4)]
         memory_index += 1
-        observation, reward, done, action = data
+        observation, reward, _, action = data
         observation = observation.astype(np.uint8)
 
-        debug_image = self.collect_statistics_and_generate_debug_image(pieces_generated,
-                                                                       *data)
+        debug_image = self.collect_statistics_and_generate_debug_image(
+            pieces_generated, *data)
         ret_dict = {
             "frame": observation,
             "frame_number": [int(pieces_generated)],
@@ -199,7 +201,7 @@ def collect_statistics_and_generate_debug_image(self, index,
                                                   reward,
                                                   done,
                                                   action):
-    """This generates extra statistics and debug images"""
+    """This generates extra statistics and debug images."""
     raise NotImplementedError()
 
   @property
@@ -266,8 +268,9 @@ def __init__(self):
     self.number_of_dones = 0
 
 
-
 class GymRealDiscreteProblem(GymDiscreteProblem):
+  """Discrete problem."""
+
   def __init__(self, *args, **kwargs):
     super(GymRealDiscreteProblem, self).__init__(*args, **kwargs)
     self.statistics = BasicStatistics()
@@ -340,11 +343,11 @@ def _setup(self):
     initial_frames_problem = environment_spec.initial_frames_problem
     # initial_frames_problem.random_skip = False
     dataset = initial_frames_problem.dataset(
-      tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
-      shuffle_files=False, hparams=hparams)
+        tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
+        shuffle_files=False, hparams=hparams)
     dataset = dataset.map(lambda x: x["input_action"]).take(1)
     input_data_iterator = (
-      dataset.batch(1).make_initializable_iterator())
+        dataset.batch(1).make_initializable_iterator())
     self._session.run(input_data_iterator.initializer)
 
     res = self._session.run(input_data_iterator.get_next())
@@ -408,16 +411,15 @@ def collect_statistics_and_generate_debug_image(self, index,
 
     ob = np.ndarray.astype(observation, np.int)
     err = np.ndarray.astype(np.maximum(np.abs(
-      stat.real_ob - ob, dtype=np.int) - 10, 0),
-                            np.uint8)
+        stat.real_ob - ob, dtype=np.int) - 10, 0), np.uint8)
     debug_im = np.concatenate([observation, stat.real_ob, err], axis=1)
 
-    assert self._internal_memory_size==self.num_testing_steps and \
-           self._internal_memory_force_beginning_resets, \
-      "The collect memory should be set in force_beginning_resets mode" \
-      "for the code below to work properly"
+    assert (self._internal_memory_size == self.num_testing_steps and
+            self._internal_memory_force_beginning_resets), (
+                "The collect memory should be set in force_beginning_resets "
+                "mode for the code below to work properly.")
 
-    if index%self._internal_memory_size == 0:
+    if index % self._internal_memory_size == 0:
       if stat.episode_sim_reward == stat.episode_real_reward:
         stat.successful_episode_reward_predictions += 1
       self._reset_real_env()
@@ -527,7 +529,8 @@ def num_rewards(self):
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnPong(
-  GymSimulatedDiscreteProblem, GymPongRandom):
+    GymSimulatedDiscreteProblem, GymPongRandom):
+  """Simulated pong."""
 
   @property
   def initial_frames_problem(self):
@@ -538,7 +541,6 @@ def num_testing_steps(self):
     return 100
 
 
-
 @registry.register_problem
 class GymFreewayRandom(GymDiscreteProblem):
   """Freeway game, random actions."""
@@ -558,13 +560,14 @@ def num_rewards(self):
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnPong(
-  GymRealDiscreteProblem, GymPongRandom):
+    GymRealDiscreteProblem, GymPongRandom):
   pass
 
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
-  GymSimulatedDiscreteProblem, GymWrappedPongRandom):
+    GymSimulatedDiscreteProblem, GymWrappedPongRandom):
+  """Similated pong."""
 
   @property
   def initial_frames_problem(self):
@@ -575,22 +578,22 @@ def num_testing_steps(self):
     return 100
 
 
-
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedLongPong(
-  GymRealDiscreteProblem, GymWrappedLongPongRandom):
+    GymRealDiscreteProblem, GymWrappedLongPongRandom):
   pass
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedLongPongAe(  # with autoencoder
-  GymDiscreteProblemWithAgentOnWrappedLongPong):
+    GymDiscreteProblemWithAgentOnWrappedLongPong):
   pass
 
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPong(
-  GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
+    GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
+  """Similated pong."""
 
   @property
   def initial_frames_problem(self):
@@ -601,22 +604,22 @@ def num_testing_steps(self):
     return 100
 
 
-
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedBreakout(
-  GymRealDiscreteProblem, GymWrappedBreakoutRandom):
+    GymRealDiscreteProblem, GymWrappedBreakoutRandom):
   pass
 
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
-  GymDiscreteProblemWithAgentOnWrappedBreakout):
+    GymDiscreteProblemWithAgentOnWrappedBreakout):
   pass
 
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
-  GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
+    GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
+  """Similated breakout."""
 
   @property
   def initial_frames_problem(self):
@@ -627,10 +630,9 @@ def num_testing_steps(self):
     return 100
 
 
-
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedPong(
-  GymRealDiscreteProblem, GymWrappedPongRandom):
+    GymRealDiscreteProblem, GymWrappedPongRandom):
   """GymDiscreteProblemWithAgentOnWrappedPong."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
@@ -654,13 +656,14 @@ def frame_width(self):
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
-  GymDiscreteProblemWithAgentOnWrappedPong):
+    GymDiscreteProblemWithAgentOnWrappedPong):
   pass
 
 
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnFreeway(
-  GymSimulatedDiscreteProblem, GymFreewayRandom):
+    GymSimulatedDiscreteProblem, GymFreewayRandom):
+  """Similated freeway."""
 
   @property
   def initial_frames_problem(self):
@@ -671,10 +674,9 @@ def num_testing_steps(self):
     return 100
 
 
-
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnFreeway(
-  GymRealDiscreteProblem, GymFreewayRandom):
+    GymRealDiscreteProblem, GymFreewayRandom):
   """Freeway with agent."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
@@ -698,5 +700,5 @@ def frame_width(self):
 
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnFreewayAe(  # with autoencoder
-  GymDiscreteProblemWithAgentOnFreeway):
+    GymDiscreteProblemWithAgentOnFreeway):
   pass
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index a42d33cd9..6d3c47a6b 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -289,7 +289,7 @@ def check_integrity_and_batch(*datasets):
     else:
       batch_dataset = preprocessed_dataset.apply(
           tf.contrib.data.batch_and_drop_remainder(num_frames))
-    dataset = batch_dataset.map(features_from_batch)  #.shuffle(8)
+    dataset = batch_dataset.map(features_from_batch)  # shuffle(8)
     return dataset
 
   def eval_metrics(self):
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 6161fb2a2..20d77dc45 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -79,13 +79,7 @@ def simulate(self, action):
 def define_collect(hparams, scope, eval_phase,
                    collect_level=-1,
                    policy_to_actions_lambda=None):
-  """Collect trajectories.
-  Returns: memory - tensor with collected rollout
-           summaries - basic statistcs about the rollout
-           initialization_lambda - initializations to be done once 
-            tf.Session is created
-  """
-
+  """Collect trajectories."""
   to_initialize = []
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
     batch_env = batch_env_factory(hparams)
@@ -133,7 +127,8 @@ def group():
     return tf.group(batch_env.reset(tf.range(len(batch_env))),
                     tf.assign(cumulative_rewards, zeros_tensor))
   reset_op = tf.cond(
-      tf.logical_or(should_reset_var, tf.convert_to_tensor(force_beginning_resets)),
+      tf.logical_or(should_reset_var, tf.convert_to_tensor(
+          force_beginning_resets)),
       group, tf.no_op)
 
   with tf.control_dependencies([reset_op]):
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index e1ae418a7..c531bc275 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -21,10 +21,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import copy
-
-from tensorflow.contrib.training import HParams
-
 from tensor2tensor.layers import common_layers
 from tensor2tensor.rl.envs import in_graph_batch_env
 from tensor2tensor.rl.envs.utils import get_action_space
@@ -33,6 +29,8 @@
 
 import tensorflow as tf
 
+from tensorflow.contrib.training import HParams
+
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index ef3b0a364..e04bb5028 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -426,8 +426,6 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   return epoch_metrics[-1]
 
 
-
-
 def combine_training_data(problem, final_data_dir, old_data_dirs,
                           copy_last_eval_set=True):
   """Add training data from old_data_dirs into final_data_dir."""
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 1d0e0b057..ec8618832 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -56,7 +56,7 @@ def train(hparams, event_dir=None, model_dir=None,
       summary_writer = None
       model_saver = None
 
-    # TODO (piotr milos): This should be refactored, possibly with
+    # TODO(piotrmilos): This should be refactored, possibly with
     # handlers for each type of env
     if hparams.environment_spec.simulated_env:
       env_model_loader = tf.train.Saver(
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index bfe143050..4081e4aa9 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -18,7 +18,8 @@
 from __future__ import print_function
 
 from tensor2tensor.data_generators.gym_problems import standard_atari_env_spec
-from tensor2tensor.models.research.rl import simple_gym_spec, feed_forward_cnn_small_categorical_fun
+from tensor2tensor.models.research.rl import feed_forward_cnn_small_categorical_fun
+from tensor2tensor.models.research.rl import simple_gym_spec
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.utils import registry  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
@@ -79,8 +80,8 @@ def test_train_pong(self):
 
     hparams.add_hparam("environment_spec",
                        standard_atari_env_spec("PongNoFrameskip-v4"))
-    rl_trainer_lib.train(hparams)
-
+    # TODO(lukaszkaiser): enable tests with Atari.
+    # rl_trainer_lib.train(hparams)
 
 
 if __name__ == "__main__":

From 8b693b5daa89f9c000582ddb009cd894159d2966 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Tue, 26 Jun 2018 12:05:19 -0700
Subject: [PATCH 0238/2720] Improve error message for unrecognized vocab type
 for text problems.

PiperOrigin-RevId: 202169649
---
 tensor2tensor/data_generators/imdb.py          | 1 +
 tensor2tensor/data_generators/text_problems.py | 3 ++-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/imdb.py b/tensor2tensor/data_generators/imdb.py
index e6542e1c0..acd615653 100644
--- a/tensor2tensor/data_generators/imdb.py
+++ b/tensor2tensor/data_generators/imdb.py
@@ -101,5 +101,6 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
 class SentimentIMDBCharacters(SentimentIMDB):
   """IMDB Sentiment classification, character level."""
 
+  @property
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 79e185190..956f1d385 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -226,7 +226,8 @@ def get_or_create_vocab(self, data_dir, tmp_dir, force_get=False):
       encoder = text_encoder.TokenTextEncoder(vocab_filename,
                                               replace_oov=self.oov_token)
     else:
-      raise ValueError("Unrecognized VocabType")
+      raise ValueError(
+          "Unrecognized VocabType: %s" % str(self.vocab_type))
     return encoder
 
   def _maybe_pack_examples(self, generator):

From ca23d52161d3d7005d595df0027dd507789da250 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 26 Jun 2018 14:00:14 -0700
Subject: [PATCH 0239/2720] Adding comments to the tensor2tensor graph neural
 network layer. This is the first of two CLs in which this file will be
 commented in detail.

PiperOrigin-RevId: 202189335
---
 .../common_message_passing_attention.py       | 434 ++++++++++++------
 1 file changed, 293 insertions(+), 141 deletions(-)

diff --git a/tensor2tensor/layers/common_message_passing_attention.py b/tensor2tensor/layers/common_message_passing_attention.py
index 9db0bcc04..69c9287e7 100644
--- a/tensor2tensor/layers/common_message_passing_attention.py
+++ b/tensor2tensor/layers/common_message_passing_attention.py
@@ -17,7 +17,6 @@
 from __future__ import division
 from __future__ import print_function
 
-
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import expert_utils
@@ -89,12 +88,17 @@ def multihead_graph_attention(query_antecedent,
     raise ValueError("Value depth (%d) must be divisible by the number of "
                      "attention heads (%d)." % (total_value_depth, num_heads))
   vars_3d_num_heads = num_heads if vars_3d else None
-  with tf.variable_scope(name, default_name="multihead_attention",
-                         values=[query_antecedent, memory_antecedent]):
+  with tf.variable_scope(
+      name,
+      default_name="multihead_attention",
+      values=[query_antecedent, memory_antecedent]):
 
     q, k, v = common_attention.compute_qkv(
-        query_antecedent, memory_antecedent, total_key_depth,
-        total_value_depth, vars_3d_num_heads=vars_3d_num_heads)
+        query_antecedent,
+        memory_antecedent,
+        total_key_depth,
+        total_value_depth,
+        vars_3d_num_heads=vars_3d_num_heads)
     q = common_attention.split_heads(q, num_heads)
     k = common_attention.split_heads(k, num_heads)
     v = common_attention.split_heads(v, num_heads)
@@ -110,12 +114,18 @@ def multihead_graph_attention(query_antecedent,
         x, additional_returned_value = x  # Unpack
 
     elif attention_type == "edge_vector":
-      x = graph_attention(q, k, v, bias, dropout_rate, image_shapes,
-                          save_weights_to=save_weights_to,
-                          make_image_summary=make_image_summary,
-                          dropout_broadcast_dims=dropout_broadcast_dims,
-                          adjacency_matrix=adjacency_matrix,
-                          num_edge_types=num_edge_types)
+      x = graph_attention(
+          q,
+          k,
+          v,
+          bias,
+          dropout_rate,
+          image_shapes,
+          save_weights_to=save_weights_to,
+          make_image_summary=make_image_summary,
+          dropout_broadcast_dims=dropout_broadcast_dims,
+          adjacency_matrix=adjacency_matrix,
+          num_edge_types=num_edge_types)
 
     x = common_attention.combine_heads(x)
 
@@ -165,9 +175,10 @@ def make_edge_vectors(adjacency_matrix, num_edge_types, depth, name=None):
     att_adj_vectors = tf.matmul(
         tf.reshape(tf.to_float(adjacency_matrix_one_hot), [-1, num_edge_types]),
         adj_vectors)
-    return tf.reshape(att_adj_vectors,
-                      [adjacency_matrix_shape[0], adjacency_matrix_shape[1],
-                       adjacency_matrix_shape[2], depth])
+    return tf.reshape(att_adj_vectors, [
+        adjacency_matrix_shape[0], adjacency_matrix_shape[1],
+        adjacency_matrix_shape[2], depth
+    ])
 
 
 def graph_attention(q,
@@ -212,8 +223,8 @@ def graph_attention(q,
     logits = tf.matmul(q, k, transpose_b=True)
     if adjacency_matrix is not None:
       key_head_depth = common_layers.shape_list(q)[-1]
-      adjacency_vectors = make_edge_vectors(
-          adjacency_matrix, num_edge_types, key_head_depth, name)
+      adjacency_vectors = make_edge_vectors(adjacency_matrix, num_edge_types,
+                                            key_head_depth, name)
       # zeroing out the vectors that have 0 entries in the adjacency
       adjacency_vectors *= tf.to_float(
           tf.expand_dims(adjacency_matrix, axis=-1))
@@ -242,52 +253,97 @@ def compute_mpnn_qkv(node_states,
                      ignore_zero=True):
   """Computes query, key and value for edge matrices.
 
+  Let B be the number of batches.
+  Let N be the number of nodes in the graph.
+  Let D be the size of the node hidden states.
+  Let K be the size of the attention keys/queries (total_key_depth).
+  Let V be the size of the attention values (total_value_depth).
+  Let T be the total number of edge types (num_edge_types).
+
+  Computes the queries, keys, and values for attention.
+  * For each node N_i in the graph, a query Q_i of size K is computed. This
+    query is used to determine the relative weights to give to each of the
+    node's incoming edges.
+  * For each node N_j and edge type t, a key K_jt of size K is computed. When an
+    edge of type t goes from node N_j to any other node, K_jt is the key that is
+    in the attention process.
+  * For each node N_j and edge type t, a value V_jt of size V is computed. When
+    an edge of type t goes from node N_j to node N_i, Attention(Q_i, K_jt)
+    produces a weight w_ijt. The message sent along this edge is w_ijt * V_jt.
+
   Args:
-    node_states: a Tensor with shape [batch, num_nodes, channels]
-    total_key_depth: an integer
-    total_value_depth: an integer
-    num_edge_types: a integer specifying number of edge types
+    node_states: A Tensor with shape [B, N, D].
+    total_key_depth: an integer (K).
+    total_value_depth: an integer (V).
+    num_edge_types: a integer specifying number of edge types (T).
     ignore_zero: If true, then edge type 0 will not be considered. Equivalent
-      to have a linear transformation of all 0's for edge type 0
+      to having a linear transformation of all 0's for edge type 0. All queries,
+      keys, and values for edge type 0 will be all 0's.
   Returns:
-    q: [batch, num_nodes, channels]
-    k: [batch, num_nodes * num_edge_types, channels]
-    v: [batch, num_nodes * num_edge_types, channels]
+    q: The attention queries for each destination node (shape [B, N, K]).
+    k: The attention keys for each node and edge type (shape [B, N*T, K]).
+    v: The attention values for each node and edge type (shape [B, N*T, V]).
   """
-  memory_antecedent = node_states
-  def _compute(inp, depth, filter_width, padding, name):
-    if filter_width == 1:
-      return common_layers.dense(inp, depth, use_bias=False, name=name)
-    else:
-      return common_layers.conv1d(inp, depth, filter_width, padding, name=name)
-  # For edge type 0, if ignore_zero, don't multiply with linear transformation,
-  # but just concat a bunch of 0's not only for efficiency but to make
-  # sure that it doesn't contribute anything to the terms
-  # TODO(avaswani): Better way to do this.
-  q = _compute(node_states, total_key_depth, 1, "VALID", "q_mpnn")
-  q_shape = common_layers.shape_list(q)
-  # k and v edge transforms have shape
-  # [batch, length, depth*nonignored_edge_types]
-  nonignored_edge_types = num_edge_types-int(ignore_zero)
-  k = _compute(memory_antecedent, total_key_depth*nonignored_edge_types, 1,
-               "VALID", "k_mpnn")
-  v = _compute(memory_antecedent, total_value_depth*nonignored_edge_types,
-               1, "VALID", "v_mpnn")
-  batch = q_shape[0]
-  length = q_shape[1]
-  k = tf.reshape(k,
-                 [batch, length, nonignored_edge_types, total_key_depth])
-  v = tf.reshape(v,
-                 [q_shape[0], q_shape[1], nonignored_edge_types,
-                  total_value_depth])
+
+  # node_states is initially a tensor with shape [B, N, D]. The call to dense
+  # creates a D x K kernel that serves as a fully-connected layer.
+  #
+  # For each possible batch b and node n in the first two dimensions of
+  # node_states, the corresponding size-D vector (the third dimension of
+  # node_states) is the hidden state for node n in batch b. Each of these size-D
+  # vectors is multiplied by the kernel to produce an attention query of size K.
+  # The result is a tensor of size [B, N, K] containing the attention queries
+  # for each node in each batch.
+  q = common_layers.dense(
+      node_states, total_key_depth, use_bias=False, name="q_mpnn")
+
+  q_shape = common_layers.shape_list(q)  # As above, q_shape = [B, N, K].
+
+  # T (or T-1 if ignore_zero).
+  nonignored_edge_types = num_edge_types - int(ignore_zero)
+
+  # Creates the attention keys in a manner similar to the process of creating
+  # the attention queries. One key is created for each type of outgoing edge the
+  # corresponding node might have, meaning k will have shape [B, N, K*T].
+  k = common_layers.dense(
+      node_states,
+      total_key_depth * nonignored_edge_types,
+      use_bias=False,
+      name="k_mpnn")
+
+  # The values over which self-attention is performed. They are created in
+  # a manner largely identical to that of the keys.
+  v = common_layers.dense(
+      node_states,
+      total_value_depth * nonignored_edge_types,
+      use_bias=False,
+      name="v_mpnn")
+
+  batch = q_shape[0]  # B.
+  length = q_shape[1]  # N.
+
+  # Making the fourth dimension explicit by separating the vectors of size
+  # K*T (in k) and V*T (in v) into two-dimensional matrices with shape [K, T]
+  # (in k) and [V, T] in v.
+  #
+  # This reshape is only necessary when ignore_zero is True (for the padding
+  # step that follows).
+  k = tf.reshape(k, [batch, length, nonignored_edge_types, total_key_depth])
+  v = tf.reshape(
+      v, [q_shape[0], q_shape[1], nonignored_edge_types, total_value_depth])
+
+  # If we previously ignored edge type 0, then we need to pad the keys and
+  # values to take this additional edge type into account. To do so, we
+  # pad the third dimension of k and v (which has size T-1 if ignore_zero is
+  # True) to size T with zeroes.
   if ignore_zero:
     k = tf.pad(k, [[0, 0], [0, 0], [1, 0], [0, 0]])
     v = tf.pad(v, [[0, 0], [0, 0], [1, 0], [0, 0]])
 
-  k = tf.reshape(k,
-                 [q_shape[0], q_shape[1]*num_edge_types, total_key_depth])
+  # Flatten out the fourth dimension.
+  k = tf.reshape(k, [q_shape[0], q_shape[1] * num_edge_types, total_key_depth])
   v = tf.reshape(v,
-                 [q_shape[0], q_shape[1]*num_edge_types, total_value_depth])
+                 [q_shape[0], q_shape[1] * num_edge_types, total_value_depth])
   return q, k, v
 
 
@@ -302,24 +358,34 @@ def multihead_mpnn_attention(node_states,
                              name="mpnn_attention"):
   """Multihead scaled-dot-product attention with input/output transformations.
 
+  Let B be the number of batches.
+  Let N be the number of nodes in the graph.
+  Let D be the size of the node hidden states.
+  Let K be the size of the attention keys/queries (total_key_depth).
+  Let V be the size of the attention values (total_value_depth).
+  Let O be the size of the attention output (output_depth).
+  Let H be the number of heads (num_heads).
+  Let T be the total number of edge types (num_edge_types).
+
+  The key and value depths are split across all of the heads. For example, if
+  the key depth is 6 and there are three heads, then the key for each head has
+  depth 2.
+
   Args:
-    node_states: A tensor of shape [batch, length, depth]
-    total_key_depth: An integer for key dimension
-    total_value_depth: An integer for value dimensions
-    output_depth: An intger for output dimemsions
-    num_heads: An integer
-    adjacency_matrix: An tensor of ints of shape [batch, length, length]
-    num_edge_types: An integer indicating number of edge bins
-    ignore_zero: A flag that says that edge type 0 should be ignored
-    name: A string
+    node_states: A Tensor with shape [B, N, D]
+    total_key_depth: An integer (K).
+    total_value_depth: An integer (V).
+    output_depth: An integer (O).
+    num_heads: An integer (H).
+    adjacency_matrix: An Tensor of ints with shape [B, N, N]. If there is an
+      edge from node j to node i in batch b, then adjacency_matrix[b, i, j]
+      contains the type of that edge as an integer. Otherwise, it contains 0.
+    num_edge_types: An integer indicating number of edge types (T).
+    ignore_zero: A flag that says that edge type 0 should be ignored.
+    name: A string.
 
   Returns:
-    The result of the attention transformation. The output shape is
-        [batch_size, length_q, output_depth]
-    unless the cache dict is provided in which case only the last memory
-    position is calculated and the output shape is [batch_size, 1, hidden_dim]
-    Optionaly returns an additional loss parameters (ex: load balance loss for
-    the experts) returned by the attention_type function.
+    The result of the attention transformation. The output shape is [B, N, O].
 
   Raises:
     ValueError: if the key depth or value depth are not divisible by the
@@ -332,110 +398,196 @@ def multihead_mpnn_attention(node_states,
     raise ValueError("Value depth (%d) must be divisible by the number of "
                      "attention heads (%d)." % (total_value_depth, num_heads))
   with tf.variable_scope(
-      name,
-      default_name="multihead_mpnn_attention",
-      values=[node_states]):
-    q, k, v = compute_mpnn_qkv(node_states,
-                               total_key_depth,
-                               total_value_depth,
-                               num_edge_types,
-                               ignore_zero=ignore_zero)
-    # reshaping k and v for head splitting
-    q_shape = tf.shape(q)
-    q = common_attention.split_heads(q, num_heads)
-    k = common_attention.split_heads(k, num_heads)
-    v = common_attention.split_heads(v, num_heads)
+      name, default_name="multihead_mpnn_attention", values=[node_states]):
+    # Create the query for each node's incoming edges.
+    # Create the keys/values for each node for each possible outgoing edge type.
+    q, k, v = compute_mpnn_qkv(
+        node_states,
+        total_key_depth,
+        total_value_depth,
+        num_edge_types,
+        ignore_zero=ignore_zero)
+
+    q_shape = tf.shape(q)  # As above, q_shape is [B, N, K].
+
+    # Divides each query/key/value into separate heads. Specifically, the
+    # query/key/value for each (batch, node) pair (i.e., the third dimensions
+    # of q, k, and v) are broken into H separate pieces. These pieces are used
+    # as the separate attention heads. The resulting tensors have shape
+    # [B, H, N, ?/H], where ? = K, K*T or V*T as appropriate.
+    q = common_attention.split_heads(q, num_heads)  # Shape [B, H, N, K/H].
+    k = common_attention.split_heads(k, num_heads)  # Shape [B, H, N, K*T/H].
+    v = common_attention.split_heads(v, num_heads)  # Shape [B, H, N, V*T/H].
     key_depth_per_head = total_key_depth // num_heads
+
+    # Ensures that the logits don't have too large of a magnitude.
     q *= key_depth_per_head**-0.5
-    # make the heads dimension leading. We will loop over heads.
-    q = tf.transpose(q, [1, 0, 2, 3])
-    k = tf.transpose(k, [1, 0, 2, 3])
-    v = tf.transpose(v, [1, 0, 2, 3])
-    # putting edge as the dimension after batch for k and v
-    # k and v will be [heads, batch, num_edge_types, length, depth]
-    k = tf.reshape(k, [num_heads, q_shape[0], q_shape[1], num_edge_types,
-                       total_key_depth//num_heads])
-    k = tf.transpose(k, [0, 1, 3, 2, 4])
-
-    v = tf.reshape(v, [num_heads, q_shape[0], q_shape[1], num_edge_types,
-                       total_value_depth//num_heads])
-    v = tf.transpose(v, [0, 1, 3, 2, 4])
-
-    # doing attention separately for each head
+
+    # Rearrange the dimensions so that the head is first. This will make
+    # subsequent steps easier (we loop over the head).
+    q = tf.transpose(q, [1, 0, 2, 3])  # Shape [H, B, N, K/H].
+    k = tf.transpose(k, [1, 0, 2, 3])  # Shape [H, B, N, K*T/H].
+    v = tf.transpose(v, [1, 0, 2, 3])  # Shape [H, B, N, V*T/H].
+
+    # Split the keys and values into separate per-edge-type keys and values.
+    k = tf.reshape(k, [
+        num_heads, q_shape[0], q_shape[1], num_edge_types,
+        total_key_depth // num_heads
+    ])  # Shape [H, B, N, T, K/H].
+    k = tf.transpose(k, [0, 1, 3, 2, 4])  # Shape [H, B, T, N, K/H].
+
+    v = tf.reshape(v, [
+        num_heads, q_shape[0], q_shape[1], num_edge_types,
+        total_value_depth // num_heads
+    ])  # Shape [H, B, N, T, V/H].
+    v = tf.transpose(v, [0, 1, 3, 2, 4])  # Shape [H, B, T, N, V/H].
+
+    # Perform attention for each head and combine the results into a list.
+    # head_outputs stores a list of tensors, each with shape [1, B, N, V/H].
+    # The last dimension contains the values computed for each attention head.
+    # Each value was determined by computing attention over all of the
+    # incoming edges for node n, weighting the incoming values accordingly,
+    # and adding those weighted values together.
     head_outputs = []
     for head_id in range(num_heads):
-      output = dot_product_mpnn_attention(q[head_id],
-                                          k[head_id],
-                                          v[head_id],
-                                          adjacency_matrix,
-                                          num_edge_types)
+      output = dot_product_mpnn_attention(q[head_id], k[head_id], v[head_id],
+                                          adjacency_matrix, num_edge_types)
+
+      # Store this result in the list of attention results for each head.
+      # The call to expand_dims gives output shape [1, B, N, V/H], which will
+      # come in handy when we combine the heads together.
       head_outputs.append(tf.expand_dims(output, axis=0))
-    # making x = [heads, batch, length, total_value_depth//num_heads]
-    x = tf.concat(head_outputs, axis=0)
-    x = tf.transpose(x, [1, 0, 2, 3])
-    # making x [batch, length, depth]
-    x = common_attention.combine_heads(x)
+
+    # Combine the heads together into one tensor and rearrange the dimensions.
+    x = tf.concat(head_outputs, axis=0)  # Shape [H, B, N, V/H].
+    x = tf.transpose(x, [1, 0, 2, 3])  # Shape [B, H, N, V/H].
+
+    # Concatenate the values produced by each head together into one vector.
+    x = common_attention.combine_heads(x)  # Shape [B, N, V].
+
+    # A fully-connected linear layer to convert from the value vectors of size V
+    # to output vectors of length O (the appropriate output length).
     x = common_layers.dense(
         x, output_depth, use_bias=False, name="output_transform")
     return x
 
 
-def dot_product_mpnn_attention(q, k, v, adjacency_matrix, num_edge_types,
-                               ignore_zero=True, name=None):
+def dot_product_mpnn_attention(q,
+                               k,
+                               v,
+                               adjacency_matrix,
+                               num_edge_types,
+                               ignore_zero=True,
+                               name=None):
   """Dot product attention with edge vectors.
 
+  Let B be the number of batches.
+  Let N be the number of nodes in the graph.
+  Let K be the size of the attention keys/queries.
+  Let V be the size of the attention values.
+  Let T be the total number of edge types (num_edge_types).
+
   Args:
-    q: [batch, length, key_depth] tensor
-    k: [batch, num_edge_types, length, key_depth]
-    v: [batch, num_edge_types, length, depth]
-    adjacency_matrix: [batch, length, length] tensor of int edge types
-    num_edge_types: an int, specifying number of edge types
-    ignore_zero: A flag that says that edge type 0 should be ignored
-    name: optional string
+    q: The query Tensor of shape [B, N, K].
+    k: The key Tensor of shape [B, T, N, K].
+    v: The value Tensor of shape [B, T, N, V].
+    adjacency_matrix: A Tensor of shape [B, N, N]. An entry at indices b, i, j
+     is the integer edge type of the edge from node j to node i in batch b.
+    num_edge_types: An integer specifying number of edge types (T).
+    ignore_zero: A flag that says that edge type 0 should be ignored.
+    name: A string.
 
   Returns:
-    A tensor of shape [batch, length, depth(q)]
+    A Tensor of shape [B, N, V] storing the result of computing attention
+    weights using the queries and keys and combining the values according to
+    those weights.
   """
+  # TODO(jfrankle): Consider ways to handle graphs that have multiple edges
+  # between the same nodes (with only one edge of each type. adjacency_matrix
+  # will need to be converted to shape [B, T, N, N].
   with tf.variable_scope(
-      name, default_name="dot_product_mpnn_attention",
+      name,
+      default_name="dot_product_mpnn_attention",
       values=[q, k, v, adjacency_matrix, num_edge_types]):
-    # Computing attention mask
-    # all edge logits will have shape [batch, edge_types, len, len]
+    # Computes the raw dot-product attention values between each query and
+    # the corresponding keys it needs to consider.
+    #
+    # This operation takes the dot product of (the query for
+    # each node) and (the key for each node for each possible edge type),
+    # creating an N x N matrix for each edge type. The entry at index (i, j)
+    # is the dot-product for the edge from node i to node j of the appropriate
+    # type. These dot products will eventually become attention weights
+    # specifying how much node i weights an edge of that type coming from node
+    # j.
     all_edge_logits = tf.matmul(
         tf.tile(tf.expand_dims(q, axis=1), [1, num_edge_types, 1, 1]),
-        k, transpose_b=True)
-    # adjacency_matrix_one_hot has shape [batch, len, len, num_edge_types]
+        k,
+        transpose_b=True)
+
+    # The adjacency matrix assumes there is only one directed edge (i <- j) for
+    # each pair of nodes. If such an edge exists, it contains the integer
+    # type of that edge at position (i, j) of the adjacency matrix.
+    #
+    # adjacency_matrix_one_hot has shape [B, N, N, T]. If there is an edge
+    # from node j to node i of type t, then index t of the last dimension is
+    # 1 for entry (i, j) of the second and third dimensions.
     adjacency_matrix_one_hot = tf.one_hot(adjacency_matrix, num_edge_types)
-    # making adjacency_matrix_one_hot [batch, edge_types, len, len]
+
+    # Rearranging the dimensions to match the shape of all_edge_logits.
     adjacency_matrix_one_hot = tf.transpose(adjacency_matrix_one_hot,
                                             [0, 3, 1, 2])
-    # getting dot products for q_i, k_j, and e_{ij}. This assumes that for
-    # edge type 0, the dot products are 0
+
+    # Element-wise multiplies all_edge_logits and adjacency_matrix_one_hot.
+    #
+    # In other words: all_edge_logits contains N x N matrices of query-key
+    # products. This element-wise multiplication zeroes out entries that do not
+    # correspond to actual edges in the graph of the appropriate edge type.
+    # all_edge_logits retains shape [B, T, N, N].
     all_edge_logits *= adjacency_matrix_one_hot
-    # logits will be [batch, length, length] after reducing along
-    # axis 1 which has dimension num_edge_types.
+
+    # Since there can only be one edge from node A to node B, we can collapse
+    # the T different adjacency matrices containing key-query pairs into one
+    # adjacency matrix. logits is [B, N, N].
     logits = tf.reduce_sum(all_edge_logits, axis=1)
-    # ignoring edges if needed
+
+    # If we do not have any special treatment for edge type 0, add a large,
+    # negative bias to each location without an edge so that the softmax of
+    # entries with the value 0 become a small negative number instead.
+    #
+    # TODO(avaswani): Better explanation of the rationale behind ignore_zero
+    # here and throughout.
     bias = 0
     if ignore_zero:
       bias = tf.to_float(tf.equal(adjacency_matrix, 0)) * -1e9
     logits += bias
-    # getting compatibilities
-    compatibility = tf.nn.softmax(logits)
+
+    # Turn the raw key-query products into a probability distribution (or,
+    # in terms of attention, weights). The softmax is computed across the
+    # last dimension of logits.
+    compatibility = tf.nn.softmax(logits)  # Shape [B, N, N].
+
+    # Computes a summary showing the attention matrix as an image. Does not do
+    # any work toward actually performing attention.
     common_attention.attention_image_summary(
         tf.expand_dims(compatibility, axis=1), None)
-    # getting edge compatibilities ready to compute values.
-    # after tiling, edge_compatibility will be
-    # [batch, num_edge_types, length, length]
+
+    # Repeats the attention matrix T times for each batch, producing
+    # a tensor with shape [B, T, N, N] where the [N, N] component is T
+    # repeats of the values found in compatibility.
     edge_compatibility = tf.tile(
         tf.expand_dims(compatibility, axis=1), [1, num_edge_types, 1, 1])
-    # computing values
-    edge_compatibility *= adjacency_matrix_one_hot
-    # all edge values will be [batch, num_edge_types, length, depth]
-    # We also assumed that the linear transformations for edge_type 0 will
-    # all be zeros. That is [batch, 0] is a length*depth tensor of 0's
-    all_edge_values = tf.matmul(edge_compatibility, v)
-    # reducing along the num_edge_types dimension
-    output = tf.reduce_sum(all_edge_values, axis=1)
-    return output
 
+    # Zeroes out the entries in edge_compatibility that do not correspond to
+    # actual edges.
+    edge_compatibility *= adjacency_matrix_one_hot  # Shape [B, T, N, N].
+
+    # Computes the incoming value vectors for each node by weighting them
+    # according to the attention weights. These values are still segregated by
+    # edge type.
+    all_edge_values = tf.matmul(edge_compatibility, v)  # Shape = [B, T, N, V].
+
+    # Combines the weighted value vectors together across edge types into a
+    # single N x V matrix for each batch.
+    output = tf.reduce_sum(all_edge_values, axis=1)  # Shape [B, N, V].
+
+    return output

From 84f8d37e47ba4b216914e05fc60525e78c164362 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 26 Jun 2018 14:15:14 -0700
Subject: [PATCH 0240/2720] Mark unused argument (nit).

PiperOrigin-RevId: 202192084
---
 tensor2tensor/rl/envs/simulated_batch_env.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index c531bc275..28a41a6c5 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -99,6 +99,7 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
 
   def __init__(self, environment_spec, length, other_hparams):
     """Batch of environments inside the TensorFlow graph."""
+    del other_hparams
     self.length = length
     initial_frames_problem = environment_spec.initial_frames_problem
     self._min_reward = initial_frames_problem.min_reward

From 699c65e6db9c97e4cbd246aa4c0e59f5a266e954 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 26 Jun 2018 15:08:20 -0700
Subject: [PATCH 0241/2720] internal change, just adding comments.

PiperOrigin-RevId: 202202011
---
 tensor2tensor/layers/common_attention.py | 5 +++--
 tensor2tensor/layers/common_layers.py    | 5 +++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 81fc7c44a..d3399b8c6 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2912,8 +2912,9 @@ def multihead_attention(query_antecedent,
         else:
           # Inplace update is required for inference on TPU.
           # Inplace_ops only supports inplace_update on the first dimension.
-          # TODO(shibow): explore updating the entire Tensor instead of using
-          # inplace_ops to avoid the transposes.
+          # The performance of current implementation is better than updating
+          # the tensor by adding the result of matmul(one_hot,
+          # update_in_current_step)
           tmp_k = tf.transpose(cache["k"], perm=[2, 0, 1, 3])
           tmp_k = common_layers.tf_inplace_ops().alias_inplace_update(
               tmp_k, decode_loop_step, tf.squeeze(k, axis=2))
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index a865a86a7..bf9b5aa7f 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1552,8 +1552,9 @@ def conv_relu_conv(inputs,
       else:
         # Inplace update is required for inference on TPU.
         # Inplace_ops only supports inplace_update on the first dimension.
-        # TODO(shibow): explore updating the entire Tensor instead of using
-        # inplace_ops to avoid the transposes.
+        # The performance of current implementation is better than updating
+        # the tensor by adding the result of matmul(one_hot,
+        # update_in_current_step)
         tmp_f = tf.transpose(cache["f"], perm=[1, 0, 2])
         tmp_f = tf_inplace_ops().alias_inplace_update(
             tmp_f, decode_loop_step * tf.shape(inputs)[1],

From dc6372c03a2004654f1de74c66bdf21dc07a84ea Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 26 Jun 2018 23:33:01 +0200
Subject: [PATCH 0242/2720] saving eval frames

---
 tensor2tensor/rl/model_rl_experiment.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index e04bb5028..6cb436354 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -190,6 +190,9 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
   model_reward_accuracy = (
       gym_simulated_problem.statistics.successful_episode_reward_predictions
       / float(n))
+  old_path = os.path.join(epoch_data_dir, "debug_frames_env")
+  new_path = os.path.join(epoch_data_dir, "debug_frames_env_eval")
+  tf.gfile.Rename(old_path, new_path)
   return model_reward_accuracy
 
 

From e0c972342ffea87317f742475f04dce18c765a53 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 26 Jun 2018 23:32:48 +0200
Subject: [PATCH 0243/2720] align fix

---
 tensor2tensor/data_generators/gym_problems.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index e85d78a27..6e7f041c9 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -341,7 +341,6 @@ def _setup(self):
                       environment_spec=environment_spec)
 
     initial_frames_problem = environment_spec.initial_frames_problem
-    # initial_frames_problem.random_skip = False
     dataset = initial_frames_problem.dataset(
         tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
         shuffle_files=False, hparams=hparams)
@@ -351,7 +350,7 @@ def _setup(self):
     self._session.run(input_data_iterator.initializer)
 
     res = self._session.run(input_data_iterator.get_next())
-    self._initial_action = res[0, :, 0]
+    self._initial_action = res[0, :, 0][:-1]
     self._reset_real_env()
 
   @property
@@ -406,7 +405,6 @@ def collect_statistics_and_generate_debug_image(self, index,
     stat = self.statistics
 
     stat.sum_of_rewards += reward
-    stat.number_of_dones += int(done)
     stat.episode_sim_reward += reward
 
     ob = np.ndarray.astype(observation, np.int)
@@ -419,9 +417,10 @@ def collect_statistics_and_generate_debug_image(self, index,
                 "The collect memory should be set in force_beginning_resets "
                 "mode for the code below to work properly.")
 
-    if index % self._internal_memory_size == 0:
+    if (index+1) % self._internal_memory_size == 0:
       if stat.episode_sim_reward == stat.episode_real_reward:
         stat.successful_episode_reward_predictions += 1
+        stat.number_of_dones += 1
       self._reset_real_env()
     else:
       stat.real_ob, real_reward, _, _ = stat.real_env.step(action)

From f45ce6f0fce3dc0a8c6e0011838664b7aa792ea2 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Wed, 27 Jun 2018 01:18:27 +0200
Subject: [PATCH 0244/2720] Putting StackAndSkipWrapper as defaul

---
 tensor2tensor/data_generators/gym_problems.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 6e7f041c9..a8dc37fde 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -53,7 +53,7 @@
 
 def standard_atari_env_spec(env):
   """Parameters of environment specification."""
-  standard_wrappers = [[tf_atari_wrappers.MaxAndSkipWrapper, {"skip": 4}]]
+  standard_wrappers = [[tf_atari_wrappers.StackAndSkipWrapper, {"skip": 4}]]
   env_lambda = None
   if isinstance(env, str):
     env_lambda = lambda: gym.make(env)

From 4de9951fbef6abe21f2c5bd6e9424e37e4ad98cd Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 26 Jun 2018 17:11:34 -0700
Subject: [PATCH 0245/2720] Fixing video2gif order bug.

PiperOrigin-RevId: 202222867
---
 tensor2tensor/utils/video2gif.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/video2gif.py b/tensor2tensor/utils/video2gif.py
index e76acb215..4a16c43a8 100644
--- a/tensor2tensor/utils/video2gif.py
+++ b/tensor2tensor/utils/video2gif.py
@@ -81,7 +81,7 @@ def main(_):
       for video in video_batch:
         print("Saving {}/{}".format(count, FLAGS.num_samples))
         name = "%s_%05d" % (base_template, count)
-        decoding.save_video(video, name + "_{}.png")
+        decoding.save_video(video, name + "_{:05d}.png")
         create_gif(name)
         count += 1
 

From f6a92ac76e85dc779b9339a4d3730f87fe041e37 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Tue, 26 Jun 2018 22:22:12 -0700
Subject: [PATCH 0246/2720] support additional options in transformer decoding

PiperOrigin-RevId: 202252975
---
 tensor2tensor/layers/common_attention.py | 21 +++++++++++++--------
 tensor2tensor/models/transformer.py      | 16 ++++++++++------
 2 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index d3399b8c6..ee62b0f88 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -659,10 +659,14 @@ def add_positional_embedding(x, max_length, name, positions=None):
   """
   _, length, depth = common_layers.shape_list(x)
   var = tf.get_variable(name, [max_length, depth])
-  if positions is not None:
-    return x + tf.gather(var, tf.to_int32(positions))
+  if positions is None:
+    sliced = tf.cond(
+        tf.less(length, max_length),
+        lambda: tf.slice(var, [0, 0], [length, -1]),
+        lambda: tf.pad(var, [[0, length - max_length], [0, 0]]))
+    return x + tf.expand_dims(sliced, 0)
   else:
-    return x + tf.expand_dims(tf.slice(var, [0, 0], [length, -1]), 0)
+    return x + tf.gather(var, tf.to_int32(positions))
 
 
 @expert_utils.add_name_scope()
@@ -2683,7 +2687,7 @@ def compute_attention_component(antecedent,
                                 filter_width=1,
                                 padding="VALID",
                                 name="c",
-                                vars_3d_num_heads=None):
+                                vars_3d_num_heads=0):
   """Computes attention compoenent (query, key or value).
 
   Args:
@@ -2698,7 +2702,7 @@ def compute_attention_component(antecedent,
   Returns:
     c : [batch, length, depth] tensor
   """
-  if vars_3d_num_heads:
+  if vars_3d_num_heads > 0:
     assert filter_width == 1
     input_depth = antecedent.get_shape().as_list()[-1]
     depth_per_head = total_depth // vars_3d_num_heads
@@ -2728,7 +2732,7 @@ def compute_qkv(query_antecedent,
                 kv_filter_width=1,
                 q_padding="VALID",
                 kv_padding="VALID",
-                vars_3d_num_heads=None):
+                vars_3d_num_heads=0):
   """Computes query, key and value.
 
   Args:
@@ -2876,7 +2880,7 @@ def multihead_attention(query_antecedent,
   if total_value_depth % num_heads != 0:
     raise ValueError("Value depth (%d) must be divisible by the number of "
                      "attention heads (%d)." % (total_value_depth, num_heads))
-  vars_3d_num_heads = num_heads if vars_3d else None
+  vars_3d_num_heads = num_heads if vars_3d else 0
   with tf.variable_scope(name, default_name="multihead_attention",
                          values=[query_antecedent, memory_antecedent]):
 
@@ -2899,7 +2903,8 @@ def multihead_attention(query_antecedent,
       if memory_antecedent is not None:
         # Encoder-Decoder Attention Cache
         q = compute_attention_component(query_antecedent, total_key_depth,
-                                        q_filter_width, q_padding, "q")
+                                        q_filter_width, q_padding, "q",
+                                        vars_3d_num_heads=vars_3d_num_heads)
         k = cache["k_encdec"]
         v = cache["v_encdec"]
       else:
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 2db47f24d..c76897eee 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -343,7 +343,7 @@ def _fast_decode_tpu(self,
     elif hparams.pos == "emb":
       positional_encoding = common_attention.add_positional_embedding(
           tf.zeros([1, decode_length + 1, hparams.hidden_size]),
-          hparams.max_length, "targets_positional_embedding", None)
+          hparams.max_length, "body/targets_positional_embedding", None)
     else:
       positional_encoding = None
 
@@ -563,8 +563,8 @@ def _fast_decode(self,
           decode_length + 1, hparams.hidden_size)
     elif hparams.pos == "emb":
       positional_encoding = common_attention.add_positional_embedding(
-          tf.zeros([1, decode_length + 1, hparams.hidden_size]),
-          hparams.max_length, "targets_positional_embedding", None)
+          tf.zeros([1, decode_length, hparams.hidden_size]),
+          hparams.max_length, "body/targets_positional_embedding", None)
     else:
       positional_encoding = None
 
@@ -850,6 +850,8 @@ def fast_decode(encoder_output,
   key_channels = hparams.attention_key_channels or hparams.hidden_size
   value_channels = hparams.attention_value_channels or hparams.hidden_size
   num_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
+  vars_3d_num_heads = (
+      hparams.num_heads if hparams.get("attention_variables_3d") else 0)
 
   cache = {
       "layer_%d" % layer: {
@@ -870,10 +872,12 @@ def fast_decode(encoder_output,
       with tf.variable_scope(
           "body/decoder/%s/encdec_attention/multihead_attention" % layer_name):
         k_encdec = common_attention.compute_attention_component(
-            encoder_output, key_channels, name="k")
+            encoder_output, key_channels, name="k",
+            vars_3d_num_heads=vars_3d_num_heads)
         k_encdec = common_attention.split_heads(k_encdec, hparams.num_heads)
         v_encdec = common_attention.compute_attention_component(
-            encoder_output, value_channels, name="v")
+            encoder_output, value_channels, name="v",
+            vars_3d_num_heads=vars_3d_num_heads)
         v_encdec = common_attention.split_heads(v_encdec, hparams.num_heads)
       cache[layer_name]["k_encdec"] = k_encdec
       cache[layer_name]["v_encdec"] = v_encdec
@@ -1023,7 +1027,7 @@ def body(self, features):
 def features_to_nonpadding(features, inputs_or_targets="inputs"):
   key = inputs_or_targets + "_segmentation"
   if features and key in features:
-    return tf.minimum(features[key], 1.0)
+    return tf.minimum(tf.to_float(features[key]), 1.0)
   return None
 
 

From 30b573eadc79bc7d8826c008f12d47e4e167faf5 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 26 Jun 2018 22:26:34 -0700
Subject: [PATCH 0247/2720] Fixing the frame_number in Gym

PiperOrigin-RevId: 202253336
---
 tensor2tensor/data_generators/gym_problems.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index e85d78a27..189b90223 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -114,6 +114,7 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
     with self._session as sess:
       self.restore_networks(sess)
       pieces_generated = 0
+      frame_counter = 0
       memory_index = 0
       memory = None
       while pieces_generated < self.num_steps:
@@ -122,19 +123,19 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
           memory_index = 0
         data = [memory[i][memory_index][0] for i in range(4)]
         memory_index += 1
-        observation, reward, _, action = data
+        observation, reward, done, action = data
         observation = observation.astype(np.uint8)
 
         debug_image = self.collect_statistics_and_generate_debug_image(
             pieces_generated, *data)
         ret_dict = {
             "frame": observation,
-            "frame_number": [int(pieces_generated)],
+            "frame_number": [int(frame_counter)],
             "image/format": ["png"],
             "image/height": [self.frame_height],
             "image/width": [self.frame_width],
             "action": [int(action)],
-            "done": [int(False)],
+            "done": [int(done)],
             "reward": [int(reward) - self.min_reward]
         }
 
@@ -143,6 +144,9 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
 
         yield ret_dict
         pieces_generated += 1
+        frame_counter += 1
+        if done:
+          frame_counter = 0
 
   def restore_networks(self, sess):
     if FLAGS.agent_policy_path:
@@ -202,7 +206,7 @@ def collect_statistics_and_generate_debug_image(self, index,
                                                   done,
                                                   action):
     """This generates extra statistics and debug images."""
-    raise NotImplementedError()
+    return None
 
   @property
   def frame_height(self):
@@ -341,7 +345,6 @@ def _setup(self):
                       environment_spec=environment_spec)
 
     initial_frames_problem = environment_spec.initial_frames_problem
-    # initial_frames_problem.random_skip = False
     dataset = initial_frames_problem.dataset(
         tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
         shuffle_files=False, hparams=hparams)

From 3f4f6e89f6ddeadec070ac399dd396de965b224f Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Wed, 27 Jun 2018 13:43:27 +0200
Subject: [PATCH 0248/2720] stats calculation bugfix

---
 tensor2tensor/data_generators/gym_problems.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 54cc5d716..4bf57da8a 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -400,9 +400,6 @@ def _reset_real_env(self):
     for a in self._initial_action:
       stat.real_ob, _, _, _ = stat.real_env.step(a)
 
-    stat.episode_sim_reward = 0.0
-    stat.episode_real_reward = 0.0
-
   def collect_statistics_and_generate_debug_image(self, index,
                                                   observation,
                                                   reward, done, action):
@@ -422,9 +419,13 @@ def collect_statistics_and_generate_debug_image(self, index,
                 "mode for the code below to work properly.")
 
     if (index+1) % self._internal_memory_size == 0:
+
       if stat.episode_sim_reward == stat.episode_real_reward:
         stat.successful_episode_reward_predictions += 1
-        stat.number_of_dones += 1
+        stat.episode_sim_reward = 0.0
+        stat.episode_real_reward = 0.0
+
+      stat.number_of_dones += 1
       self._reset_real_env()
     else:
       stat.real_ob, real_reward, _, _ = stat.real_env.step(action)

From 54f5b9e67903f5a3048f60794e896b06e9168feb Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Wed, 27 Jun 2018 14:49:07 +0200
Subject: [PATCH 0249/2720] better parameter passing to ppo

---
 tensor2tensor/data_generators/gym_problems.py |  2 +-
 tensor2tensor/rl/model_rl_experiment.py       | 16 +++++++++-------
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 4bf57da8a..cb4ba778d 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -304,7 +304,7 @@ def __init__(self):
     # data to calculate
     # correctness of rewards per sequence metric
     self.episode_sim_reward = 0.0
-    self.episode_real_reward = 0.0,
+    self.episode_real_reward = 0.0
     self.successful_episode_reward_predictions = 0
     self.report_reward_statistics_every = 10
 
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 6cb436354..920be540e 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -138,16 +138,19 @@ def train_agent(problem_name, agent_model_dir,
   """Train the PPO agent in the simulated environment."""
   gym_problem = registry.problem(problem_name)
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
+  ppo_params_names = ["epochs_num", "epoch_length",
+                      "learning_rate", "num_agents",
+                      "optimization_epochs"]
+
+  for param_name in ppo_params_names:
+    ppo_param_name = "ppo_"+ param_name
+    if ppo_param_name in hparams:
+      ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
+
   ppo_epochs_num = hparams.ppo_epochs_num
-  ppo_hparams.epochs_num = ppo_epochs_num
-  ppo_hparams.eval_every_epochs = 50
   ppo_hparams.save_models_every_epochs = ppo_epochs_num
-  ppo_hparams.epoch_length = hparams.ppo_epoch_length
-  ppo_hparams.num_agents = hparams.ppo_num_agents
   ppo_hparams.world_model_dir = world_model_dir
   ppo_hparams.add_hparam("force_beginning_resets", True)
-  if hparams.ppo_learning_rate:
-    ppo_hparams.learning_rate = hparams.ppo_learning_rate
 
   # Adding model hparams for model specific adjustments
   model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
@@ -476,7 +479,6 @@ def rl_modelrl_base():
       # though it is not necessary.
       ppo_epoch_length=60,
       ppo_num_agents=16,
-      ppo_learning_rate=0.,
       # Whether the PPO agent should be restored from the previous iteration, or
       # should start fresh each time.
       ppo_continue_training=True,

From 35e3a63a972fdebaa8ed0e79e0783231ec6200cd Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Wed, 27 Jun 2018 08:37:17 -0700
Subject: [PATCH 0250/2720] Add Cleverhans integration to T2T

PiperOrigin-RevId: 202314334
---
 setup.py                                 |   4 +
 tensor2tensor/bin/t2t_attack.py          | 134 +++++++++++++++++++++++
 tensor2tensor/data_generators/problem.py |   3 +-
 tensor2tensor/models/shake_shake.py      |  26 ++++-
 tensor2tensor/utils/adv_attack_utils.py  |  60 ++++++++++
 tensor2tensor/utils/registry.py          | 115 ++++++++++++++++---
 tensor2tensor/utils/t2t_model.py         |  62 ++++++-----
 7 files changed, 358 insertions(+), 46 deletions(-)
 create mode 100644 tensor2tensor/bin/t2t_attack.py
 create mode 100644 tensor2tensor/utils/adv_attack_utils.py

diff --git a/setup.py b/setup.py
index c7f70b239..569e60c81 100644
--- a/setup.py
+++ b/setup.py
@@ -33,6 +33,7 @@
     ],
     install_requires=[
         'bz2file',
+        'cleverhans',
         'flask',
         'future',
         'gevent',
@@ -66,5 +67,8 @@
         'License :: OSI Approved :: Apache Software License',
         'Topic :: Scientific/Engineering :: Artificial Intelligence',
     ],
+    dependency_links=[
+        'git+https://github.com/tensorflow/cleverhans.git#egg=cleverhans'
+    ],
     keywords='tensorflow machine learning',
 )
diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
new file mode 100644
index 000000000..88f54c5b0
--- /dev/null
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -0,0 +1,134 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Adversarially attack a model."""
+
+import os
+
+from tensor2tensor.bin import t2t_trainer
+from tensor2tensor.data_generators import problem as problem_lib
+from tensor2tensor.utils import adv_attack_utils
+from tensor2tensor.utils import cloud_mlengine
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
+from tensor2tensor.utils import trainer_lib
+from tensor2tensor.utils import usr_dir
+
+import tensorflow as tf
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+# See flags.py for additional command-line flags.
+flags.DEFINE_string("attack_params_set", None,
+                    "Which attack parameters to use.")
+flags.DEFINE_boolean(
+    "ignore_incorrect", False, "Ignore examples that are "
+    "incorrectly classified to begin with.")
+
+
+def create_attack_params():
+  return registry.attack_params(FLAGS.attack_params_set)
+
+
+def create_attack(attack):
+  return registry.attacks(attack)
+
+
+def main(argv):
+  tf.logging.set_verbosity(tf.logging.INFO)
+  trainer_lib.set_random_seed(FLAGS.random_seed)
+  usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
+  t2t_trainer.maybe_log_registry_and_exit()
+
+
+  if FLAGS.cloud_mlengine:
+    cloud_mlengine.launch()
+    return
+
+  if FLAGS.generate_data:
+    t2t_trainer.generate_data()
+
+  if cloud_mlengine.job_dir():
+    FLAGS.output_dir = cloud_mlengine.job_dir()
+
+  if argv:
+    t2t_trainer.set_hparams_from_args(argv[1:])
+  hparams = t2t_trainer.create_hparams()
+  trainer_lib.add_problem_hparams(hparams, FLAGS.problem)
+  attack_params = create_attack_params()
+  attack_params.add_hparam("eps", 0.0)
+
+  config = t2t_trainer.create_run_config(hparams)
+  params = {"batch_size": hparams.batch_size}
+
+  # add "_rev" as a hack to avoid image standardization
+  problem = registry.problem(FLAGS.problem + "_rev")
+  input_fn = problem.make_estimator_input_fn(
+      tf.estimator.ModeKeys.EVAL, hparams)
+  features, _ = input_fn(params, config).make_one_shot_iterator().get_next()
+  inputs, labels = features["targets"], features["inputs"]
+  inputs = tf.to_float(inputs)
+  labels = tf.squeeze(labels)
+
+  sess = tf.Session()
+
+  model_fn = t2t_model.T2TModel.make_estimator_model_fn(
+      FLAGS.model, hparams, use_tpu=FLAGS.use_tpu)
+  ch_model = adv_attack_utils.T2TAttackModel(model_fn, params, config)
+
+  acc_mask = None
+  probs = ch_model.get_probs(inputs)
+  if FLAGS.ignore_incorrect:
+    preds = tf.argmax(probs, -1)
+    preds = tf.squeeze(preds)
+    acc_mask = tf.to_float(tf.equal(labels, preds))
+  one_hot_labels = tf.one_hot(labels, probs.shape[-1])
+
+  attack = create_attack(attack_params.attack)(ch_model, sess=sess)
+
+  # Restore weights
+  saver = tf.train.Saver()
+  checkpoint_path = os.path.expanduser(FLAGS.output_dir or
+                                       FLAGS.checkpoint_path)
+  saver.restore(sess, tf.train.latest_checkpoint(checkpoint_path))
+
+  # reuse variables
+  tf.get_variable_scope().reuse_variables()
+
+  def compute_accuracy(x, labels, mask):
+    preds = ch_model.get_probs(x)
+    preds = tf.squeeze(preds)
+    preds = tf.argmax(preds, -1, output_type=labels.dtype)
+    _, acc_update_op = tf.metrics.accuracy(
+        labels=labels, predictions=preds, weights=mask)
+    sess.run(tf.initialize_local_variables())
+    for _ in range(FLAGS.eval_steps):
+      acc = sess.run(acc_update_op)
+    return acc
+
+  acc = compute_accuracy(inputs, labels, acc_mask)
+  epsilon_acc_pairs = [(0.0, acc)]
+  for epsilon in attack_params.attack_epsilons:
+    attack_params.eps = epsilon
+    adv_x = attack.generate(inputs, y=one_hot_labels, **attack_params.values())
+    acc = compute_accuracy(adv_x, labels, acc_mask)
+    epsilon_acc_pairs.append((epsilon, acc))
+
+  for epsilon, acc in epsilon_acc_pairs:
+    tf.logging.info("Accuracy @ eps=%f: %f" % (epsilon, acc))
+
+
+if __name__ == "__main__":
+  tf.app.run()
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 5b0186d8e..d2d100768 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -812,8 +812,7 @@ def define_shapes(example):
         # on TPU, we use params["batch_size"], which specifies the number of
         # examples across all datashards
         batch_size = params["batch_size"]
-        dataset = dataset.apply(
-            tf.contrib.data.batch_and_drop_remainder(batch_size))
+        dataset = dataset.batch(batch_size, drop_remainder=True)
       else:
         num_shards = config.data_parallelism.n if config else 1
         batch_size = hparams.batch_size * num_shards
diff --git a/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index 964290d11..77d953380 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -33,15 +33,15 @@ def shake_shake_skip_connection(x, output_filters, stride, is_training):
   stride_spec = [1, stride, stride, 1]
   # Skip path 1.
   path1 = tf.nn.avg_pool(x, [1, 1, 1, 1], stride_spec, "VALID")
-  path1 = tf.layers.conv2d(path1, int(output_filters / 2), (1, 1),
-                           padding="SAME", name="path1_conv")
+  path1 = tf.layers.conv2d(
+      path1, int(output_filters / 2), (1, 1), padding="SAME", name="path1_conv")
 
   # Skip path 2.
   pad_arr = [[0, 0], [0, 1], [0, 1], [0, 0]]  # First pad with 0's then crop.
   path2 = tf.pad(x, pad_arr)[:, 1:, 1:, :]
   path2 = tf.nn.avg_pool(path2, [1, 1, 1, 1], stride_spec, "VALID")
-  path2 = tf.layers.conv2d(path2, int(output_filters / 2), (1, 1),
-                           padding="SAME", name="path2_conv")
+  path2 = tf.layers.conv2d(
+      path2, int(output_filters / 2), (1, 1), padding="SAME", name="path2_conv")
 
   # Concat and apply BN.
   final_path = tf.concat(values=[path1, path2], axis=-1)
@@ -55,8 +55,12 @@ def shake_shake_branch(x, output_filters, stride, rand_forward, rand_backward,
   """Building a 2 branching convnet."""
   is_training = hparams.mode == tf.contrib.learn.ModeKeys.TRAIN
   x = tf.nn.relu(x)
-  x = tf.layers.conv2d(x, output_filters, (3, 3), strides=(stride, stride),
-                       padding="SAME", name="conv1")
+  x = tf.layers.conv2d(
+      x,
+      output_filters, (3, 3),
+      strides=(stride, stride),
+      padding="SAME",
+      name="conv1")
   x = tf.layers.batch_normalization(x, training=is_training, name="bn1")
   x = tf.nn.relu(x)
   x = tf.layers.conv2d(x, output_filters, (3, 3), padding="SAME", name="conv2")
@@ -205,3 +209,13 @@ def shakeshake_tpu():
   hparams.learning_rate_cosine_cycle_steps = 180000
   hparams.learning_rate = 0.6
   return hparams
+
+
+@registry.register_attack_params
+def shake_shake_fgsm():
+  aparams = tf.contrib.training.HParams()
+  aparams.attack = "fgsm"
+  aparams.attack_epsilons = [(i+1) * 0.1 for i in range(12)]
+  aparams.add_hparam("clip_min", 0.0)
+  aparams.add_hparam("clip_max", 255.0)
+  return aparams
diff --git a/tensor2tensor/utils/adv_attack_utils.py b/tensor2tensor/utils/adv_attack_utils.py
new file mode 100644
index 000000000..d543f7615
--- /dev/null
+++ b/tensor2tensor/utils/adv_attack_utils.py
@@ -0,0 +1,60 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Utilities to assist in performing adversarial attack using Cleverhans."""
+
+from cleverhans import attacks
+from cleverhans import model
+
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+@registry.register_attack
+def fgsm():
+  return attacks.FastGradientMethod
+
+
+@registry.register_attack
+def madry():
+  return attacks.MadryEtAl
+
+
+class T2TAttackModel(model.Model):
+  """Wrapper of Cleverhans Model object."""
+
+  def __init__(self, model_fn, params, config):
+    self._model_fn = model_fn
+    self._params = params
+    self._config = config
+    self._logits_dict = {}
+
+  def get_logits(self, x):
+    if x.name in self._logits_dict:
+      return self._logits_dict[x.name]
+
+    x = tf.map_fn(tf.image.per_image_standardization, x)
+
+    logits = self._model_fn(
+        {
+            "inputs": x
+        },
+        None,
+        "attack",
+        params=self._params,
+        config=self._config)
+    self._logits_dict[x.name] = logits
+
+    return tf.squeeze(logits)
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 525c250f5..2a759c1bf 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -51,6 +51,8 @@ class MyModel(T2TModel):
 _MODELS = {}
 _HPARAMS = {}
 _RANGED_HPARAMS = {}
+_ATTACKS = {}
+_ATTACK_PARAMS = {}
 _PROBLEMS = {}
 
 
@@ -85,7 +87,8 @@ def _convert_camel_to_snake(name):
 
 
 def _reset():
-  for ctr in [_MODELS, _HPARAMS, _RANGED_HPARAMS] + list(_MODALITIES.values()):
+  for ctr in [_MODELS, _HPARAMS, _RANGED_HPARAMS, _ATTACK_PARAMS] + list(
+      _MODALITIES.values()):
     ctr.clear()
 
 
@@ -135,8 +138,8 @@ def decorator(model_cls, registration_name=None):
 
 def model(name):
   if name not in _MODELS:
-    raise LookupError("Model %s never registered.  Available models:\n %s" % (
-        name, "\n".join(list_models())))
+    raise LookupError("Model %s never registered.  Available models:\n %s" %
+                      (name, "\n".join(list_models())))
 
   return _MODELS[name]
 
@@ -279,12 +282,91 @@ def list_problems():
   return sorted(list(_PROBLEMS))
 
 
+def register_attack(name=None):
+  """Register an attack HParams set. Same behaviour as register_hparams."""
+
+  def decorator(attack_fn, registration_name=None):
+    """Registers & returns attack_fn with registration_name or default name."""
+    attack_name = registration_name or default_name(attack_fn)
+    if attack_name in _ATTACKS and not tf.contrib.eager.in_eager_mode():
+      raise LookupError("Attack %s already registered." % attack_name)
+    _ATTACKS[attack_name] = attack_fn
+    return attack_fn
+
+  # Handle if decorator was used without parens
+  if callable(name):
+    attack_fn = name
+    return decorator(attack_fn, registration_name=default_name(attack_fn))
+
+  return lambda attack_fn: decorator(attack_fn, name)
+
+
+def attacks(name):
+  """Retrieve registered attack by name."""
+  if name not in _ATTACKS:
+    error_msg = "Attack %s never registered. Sets registered:\n%s"
+    raise LookupError(
+        error_msg % (name,
+                     display_list_by_prefix(list_attacks(), starting_spaces=4)))
+  attack = _ATTACKS[name]()
+  if attack is None:
+    raise TypeError(
+        "Attack %s is None. Make sure the registered function returns a "
+        "`cleverhans.attack.Attack` object." % name)
+  return attack
+
+
+def list_attacks(prefix=None):
+  if prefix:
+    return [name for name in _ATTACKS if name.startswith(prefix)]
+  return list(_ATTACKS)
+
+
+def register_attack_params(name=None):
+  """Register an attack HParams set. Same behaviour as register_hparams."""
+
+  def decorator(ap_fn, registration_name=None):
+    """Registers & returns ap_fn with registration_name or default name."""
+    ap_name = registration_name or default_name(ap_fn)
+    if ap_name in _ATTACK_PARAMS and not tf.contrib.eager.in_eager_mode():
+      raise LookupError("Attack HParams set %s already registered." % ap_name)
+    _ATTACK_PARAMS[ap_name] = ap_fn
+    return ap_fn
+
+  # Handle if decorator was used without parens
+  if callable(name):
+    ap_fn = name
+    return decorator(ap_fn, registration_name=default_name(ap_fn))
+
+  return lambda ap_fn: decorator(ap_fn, name)
+
+
+def attack_params(name):
+  """Retrieve registered aparams by name."""
+  if name not in _ATTACK_PARAMS:
+    error_msg = "Attack HParams set %s never registered. Sets registered:\n%s"
+    raise LookupError(
+        error_msg %
+        (name, display_list_by_prefix(list_attack_params(), starting_spaces=4)))
+  ap = _ATTACK_PARAMS[name]()
+  if ap is None:
+    raise TypeError("Attack HParams %s is None. Make sure the registered "
+                    "function returns the HParams object." % name)
+  return ap
+
+
+def list_attack_params(prefix=None):
+  if prefix:
+    return [name for name in _ATTACK_PARAMS if name.startswith(prefix)]
+  return list(_ATTACK_PARAMS)
+
+
 def _internal_get_modality(name, mod_collection, collection_str):
   if name is None:
     name = "default"
   if name not in mod_collection:
-    raise LookupError("%s modality %s never registered." % (collection_str,
-                                                            name))
+    raise LookupError(
+        "%s modality %s never registered." % (collection_str, name))
   return mod_collection[name]
 
 
@@ -330,8 +412,8 @@ def decorator(mod_cls, registration_name=None):
     """Registers & returns mod_cls with registration_name or default name."""
     mod_name = registration_name or default_name(mod_cls)
     if mod_name in mod_collection and not tf.contrib.eager.in_eager_mode():
-      raise LookupError("%s modality %s already registered." % (collection_str,
-                                                                mod_name))
+      raise LookupError(
+          "%s modality %s already registered." % (collection_str, mod_name))
     mod_collection[mod_name] = mod_cls
     return mod_cls
 
@@ -390,8 +472,8 @@ def list_modalities():
   for modality_type, modalities in six.iteritems(_MODALITIES):
     all_modalities.extend([
         "%s:%s" % (mtype, modality)
-        for mtype, modality in zip([modality_type] * len(modalities),
-                                   modalities)
+        for mtype, modality in zip([modality_type] *
+                                   len(modalities), modalities)
     ])
   return all_modalities
 
@@ -471,15 +553,22 @@ def help_string():
 
   Problems:
 %s
-  """
+
+  Attacks:
+%s
+
+  Attack HParams:
+%s
+"""
   m, hp, rhp, mod, probs = [
-      display_list_by_prefix(entries, starting_spaces=4)
-      for entries in [
+      display_list_by_prefix(entries, starting_spaces=4) for entries in [
           list_models(),
           list_hparams(),
           list_ranged_hparams(),
           list_modalities(),
-          list_problems()
+          list_problems(),
+          list_attacks(),
+          list_attack_params()
       ]
   ]
   return help_str % (m, hp, rhp, mod, probs)
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 85b705eaf..8f98ce619 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -43,7 +43,6 @@
 from tensorflow.python.layers import base
 from tensorflow.python.ops import variable_scope
 
-
 _no_problem_err_str = (
     "The default implementation of %s requires that the "
     "model be used with a Problem. If using a Problem, augment the "
@@ -71,7 +70,8 @@ def __init__(self,
                mode=tf.estimator.ModeKeys.TRAIN,
                problem_hparams=None,
                data_parallelism=None,
-               decode_hparams=None):
+               decode_hparams=None,
+               **kwargs):
     """Create a T2TModel.
 
     Args:
@@ -85,6 +85,7 @@ def __init__(self,
         specifies devices for data parallelism.
       decode_hparams: a hyperparameter object with decoding parameters.
         See decoding.decode_hparams.
+      **kwargs: arguments to pass to base.Layer constructor.
 
     Returns:
       a T2TModel
@@ -93,7 +94,7 @@ def __init__(self,
     default_name = registry.default_name(type(self))
     name = self.REGISTERED_NAME or default_name
     super(T2TModel, self).__init__(
-        trainable=mode == tf.estimator.ModeKeys.TRAIN, name=name)
+        trainable=mode == tf.estimator.ModeKeys.TRAIN, name=name, **kwargs)
 
     if not problem_hparams and hasattr(hparams, "problem_hparams"):
       problem_hparams = hparams.problem_hparams
@@ -272,7 +273,8 @@ def model_fn(self, features):
       else:
         logits = self.top(output, features)
         losses["training"] = 0.0
-        if self._hparams.mode != tf.estimator.ModeKeys.PREDICT:
+        if (self._hparams.mode != tf.estimator.ModeKeys.PREDICT and
+            self._hparams.mode != "attack"):
           losses["training"] = self.loss(logits, features)
 
       return logits, losses
@@ -367,14 +369,13 @@ def _top_single(self, body_output, target_modality, features):
         else:
           body_output_shape = body_output.shape.as_list()
           last_position_body_output = tf.slice(
-              body_output,
-              [0, features["decode_loop_step"][0], 0, 0],
-              [body_output_shape[0], 1, body_output_shape[2],
-               body_output_shape[3]])
+              body_output, [0, features["decode_loop_step"][0], 0, 0], [
+                  body_output_shape[0], 1, body_output_shape[2],
+                  body_output_shape[3]
+              ])
           target_shape = features["targets"].shape.as_list()
           last_position_targets = tf.slice(
-              features["targets"],
-              [0, features["decode_loop_step"][0], 0, 0],
+              features["targets"], [0, features["decode_loop_step"][0], 0, 0],
               [target_shape[0], 1, target_shape[2], target_shape[3]])
         logits = target_modality.top(last_position_body_output,
                                      last_position_targets)
@@ -723,8 +724,8 @@ def _greedy_infer(self, features, decode_length, use_tpu=False):
           "losses": a dictionary: {loss-name (string): floating point `Scalar`}
       }
     """
-    return (self._slow_greedy_infer_tpu(features, decode_length) if use_tpu else
-            self._slow_greedy_infer(features, decode_length))
+    return (self._slow_greedy_infer_tpu(features, decode_length)
+            if use_tpu else self._slow_greedy_infer(features, decode_length))
 
   def _slow_greedy_infer_tpu(self, features, decode_length):
     """A slow greedy inference method on TPU.
@@ -783,8 +784,8 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
       else:
         cur_sample = samples[:, i, :, :]
       samples = tf.transpose(recent_output, perm=[1, 0, 2, 3])
-      samples = tf_inplace_ops().alias_inplace_update(
-          samples, i, tf.to_int64(cur_sample))
+      samples = tf_inplace_ops().alias_inplace_update(samples, i,
+                                                      tf.to_int64(cur_sample))
       samples = tf.transpose(samples, perm=[1, 0, 2, 3])
       if not tf.contrib.eager.in_eager_mode():
         samples.set_shape([None, None, None, 1])
@@ -816,17 +817,16 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
       decode_length = 1
     else:
       if "partial_targets" in features:
-        prefix_length = common_layers.shape_list(
-            features["partial_targets"])[1]
+        prefix_length = common_layers.shape_list(features["partial_targets"])[1]
       else:
-        prefix_length = common_layers.shape_list(
-            features["inputs"])[1]
+        prefix_length = common_layers.shape_list(features["inputs"])[1]
       decode_length = prefix_length + decode_length
 
     # Initial values of result, logits and loss.
-    result = tf.concat([initial_output,
-                        tf.zeros([batch_size, decode_length, 1, 1], tf.int64)],
-                       axis=1)
+    result = tf.concat(
+        [initial_output,
+         tf.zeros([batch_size, decode_length, 1, 1], tf.int64)],
+        axis=1)
     # tensor padded to [batch_size, decode_length, 1, 1, vocab_size]
     logits = tf.zeros((batch_size, decode_length, 1, 1,
                        target_modality.top_dimensionality))
@@ -868,8 +868,10 @@ def fn_not_eos():
         shape_invariants=[
             tf.TensorShape([]),
             tf.TensorShape([batch_size, decode_length, 1, 1]),
-            tf.TensorShape([batch_size, decode_length, 1, 1,
-                            target_modality.top_dimensionality]),
+            tf.TensorShape([
+                batch_size, decode_length, 1, 1,
+                target_modality.top_dimensionality
+            ]),
             tf.TensorShape([]),
         ],
         back_prop=False,
@@ -1143,7 +1145,11 @@ def make_estimator_model_fn(model_name,
                               use_tpu=False):
     model_cls = registry.model(model_name)
 
-    def wrapping_model_fn(features, labels, mode, params=None, config=None):
+    def wrapping_model_fn(features,
+                          labels,
+                          mode,
+                          params=None,
+                          config=None):
       return model_cls.estimator_model_fn(
           hparams,
           features,
@@ -1189,11 +1195,13 @@ def estimator_model_fn(cls,
     data_parallelism = None
     if not use_tpu and config:
       data_parallelism = config.data_parallelism
+    reuse = tf.get_variable_scope().reuse
     model = cls(
         hparams,
         mode,
         data_parallelism=data_parallelism,
-        decode_hparams=decode_hparams)
+        decode_hparams=decode_hparams,
+        _reuse=reuse)
 
     # PREDICT mode
     if mode == tf.estimator.ModeKeys.PREDICT:
@@ -1228,6 +1236,10 @@ def estimator_model_fn(cls,
 
     assert "training" in losses_dict
 
+    # Attack mode
+    if mode == "attack":
+      return logits
+
     # Summarize losses
     with tf.name_scope("losses"):
       for loss_name, loss_val in sorted(losses_dict.items()):

From 0c2100043da504d26da8618e4d136e0706d664e3 Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Wed, 27 Jun 2018 09:08:50 -0700
Subject: [PATCH 0251/2720] Fix Travis CI and add multiple passes over data

PiperOrigin-RevId: 202318863
---
 setup.py                        | 1 -
 tensor2tensor/bin/t2t_attack.py | 3 ++-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index 569e60c81..2a509074e 100644
--- a/setup.py
+++ b/setup.py
@@ -33,7 +33,6 @@
     ],
     install_requires=[
         'bz2file',
-        'cleverhans',
         'flask',
         'future',
         'gevent',
diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 88f54c5b0..15d7f30f4 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -77,7 +77,8 @@ def main(argv):
   problem = registry.problem(FLAGS.problem + "_rev")
   input_fn = problem.make_estimator_input_fn(
       tf.estimator.ModeKeys.EVAL, hparams)
-  features, _ = input_fn(params, config).make_one_shot_iterator().get_next()
+  dataset = input_fn(params, config).repeat()
+  features, _ = dataset.make_one_shot_iterator().get_next()
   inputs, labels = features["targets"], features["inputs"]
   inputs = tf.to_float(inputs)
   labels = tf.squeeze(labels)

From ef874abe3446a75d15a536292127f324233ce9a6 Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Wed, 27 Jun 2018 09:28:19 -0700
Subject: [PATCH 0252/2720] Fix help string

PiperOrigin-RevId: 202321399
---
 tensor2tensor/utils/registry.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 2a759c1bf..59f6711ee 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -560,7 +560,7 @@ def help_string():
   Attack HParams:
 %s
 """
-  m, hp, rhp, mod, probs = [
+  m, hp, rhp, mod, probs, atks, ap = [
       display_list_by_prefix(entries, starting_spaces=4) for entries in [
           list_models(),
           list_hparams(),
@@ -571,4 +571,4 @@ def help_string():
           list_attack_params()
       ]
   ]
-  return help_str % (m, hp, rhp, mod, probs)
+  return help_str % (m, hp, rhp, mod, probs, atks, ap)

From 73be3d2dd9f5242d1c844fb4122ae45a4cd99981 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 27 Jun 2018 10:13:50 -0700
Subject: [PATCH 0253/2720] Adding LSUN Bedroom dataset.

PiperOrigin-RevId: 202328415
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/image_lsun.py   | 101 ++++++++++++++++++
 2 files changed, 102 insertions(+)
 create mode 100644 tensor2tensor/data_generators/image_lsun.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 47336c8de..afae0033a 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -38,6 +38,7 @@
     "tensor2tensor.data_generators.gym_problems",
     "tensor2tensor.data_generators.ice_parsing",
     "tensor2tensor.data_generators.imagenet",
+    "tensor2tensor.data_generators.image_lsun",
     "tensor2tensor.data_generators.imdb",
     "tensor2tensor.data_generators.lambada",
     "tensor2tensor.data_generators.librispeech",
diff --git a/tensor2tensor/data_generators/image_lsun.py b/tensor2tensor/data_generators/image_lsun.py
new file mode 100644
index 000000000..ce49d02c7
--- /dev/null
+++ b/tensor2tensor/data_generators/image_lsun.py
@@ -0,0 +1,101 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""LSUN datasets (bedrooms only for now)."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import io
+import os
+import zipfile
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import image_utils
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+
+_LSUN_URL = "http://lsun.cs.princeton.edu/htbin/download.cgi?tag=latest&category=%s&set=%s"
+_LSUN_DATA_FILENAME = "lsun-%s-%s.zip"
+
+
+def pil_image():
+  import PIL  # pylint: disable=g-import-not-at-top
+  return PIL.Image
+
+
+def _get_lsun(directory, category, split_name):
+  """Downloads all lsun files to directory unless they are there."""
+  generator_utils.maybe_download(directory,
+                                 _LSUN_DATA_FILENAME % (category, split_name),
+                                 _LSUN_URL % (category, split_name))
+
+
+@registry.register_problem
+class LsunBedrooms(image_utils.ImageProblem):
+  """LSUN Bedrooms."""
+
+  @property
+  def num_channels(self):
+    """Number of color channels."""
+    return 3
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    """Generates LSUN bedrooms dataset and writes it in data_dir."""
+    generator_utils.generate_dataset_and_shuffle(
+        self.read_and_convert_to_png(tmp_dir, "train"),
+        self.training_filepaths(data_dir, 100, shuffled=False),
+        self.read_and_convert_to_png(tmp_dir, "val"),
+        self.dev_filepaths(data_dir, 1, shuffled=False))
+
+  def read_and_convert_to_png(self, tmp_dir, split_name):
+    """Downloads the datasets, extracts from zip and yields in PNG format."""
+    category = "bedroom"
+    _get_lsun(tmp_dir, category, split_name)
+    filename = _LSUN_DATA_FILENAME % (category, split_name)
+    data_path = os.path.join(tmp_dir, filename)
+    print("Extracting zip file.")
+    zip_ref = zipfile.ZipFile(data_path, "r")
+    zip_ref.extractall(tmp_dir)
+    zip_ref.close()
+
+    print("Opening database.")
+    data_file = os.path.join(tmp_dir,
+                             "%s_%s_lmdb/data.mdb" % (category, split_name))
+
+    filename_queue = tf.train.string_input_producer([data_file], num_epochs=1)
+    reader = tf.LMDBReader()
+    _, webp_image_tensor = reader.read(filename_queue)
+
+    object_count = 0
+    with tf.train.MonitoredTrainingSession() as session:
+      while True:
+        webp_image = session.run(webp_image_tensor)
+        object_count += 1
+        if object_count % 1000 == 0:
+          print("Extracted %d objects." % object_count)
+        # Unfortunately Tensorflow doesn't support reading or parsing
+        # WebP images, so we have to do it via Image PIL library.
+        image = pil_image().open(io.BytesIO(webp_image))
+        buf = io.BytesIO()
+        width, height = image.size
+        image.save(buf, "PNG")
+        yield {
+            "image/encoded": [buf.getvalue()],
+            "image/format": ["png"],
+            "image/class/label": [0],
+            "image/height": [height],
+            "image/width": [width]
+        }

From 23341c3ef8762c3bf6833ee87bd6c8d5c87ccbbe Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Wed, 27 Jun 2018 10:40:34 -0700
Subject: [PATCH 0254/2720] Clean up unused imports and args

PiperOrigin-RevId: 202333134
---
 tensor2tensor/bin/t2t_attack.py               |  2 +-
 tensor2tensor/data_generators/gym_problems.py | 98 +++++++++----------
 2 files changed, 49 insertions(+), 51 deletions(-)

diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 15d7f30f4..b15caea65 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -17,7 +17,7 @@
 import os
 
 from tensor2tensor.bin import t2t_trainer
-from tensor2tensor.data_generators import problem as problem_lib
+from tensor2tensor.data_generators import problem as problem_lib  # pylint: disable=unused-import
 from tensor2tensor.utils import adv_attack_utils
 from tensor2tensor.utils import cloud_mlengine
 from tensor2tensor.utils import registry
diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 189b90223..4ced44fa9 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -38,17 +38,16 @@
 
 from tensorflow.contrib.training import HParams
 
-
 flags = tf.flags
 FLAGS = flags.FLAGS
 
-
 flags.DEFINE_string("agent_policy_path", None, "File with model for agent.")
 flags.DEFINE_string("autoencoder_path", None,
                     "File with model for autoencoder.")
-flags.DEFINE_boolean("only_use_ae_for_policy", False,
-                     "Whether to only use the autoencoder for the policy and "
-                     "still write out full-resolution frames.")
+flags.DEFINE_boolean(
+    "only_use_ae_for_policy", False,
+    "Whether to only use the autoencoder for the policy and "
+    "still write out full-resolution frames.")
 
 
 def standard_atari_env_spec(env):
@@ -61,9 +60,8 @@ def standard_atari_env_spec(env):
     env_lambda = env
   assert env is not None, "Unknown specification of environment"
 
-  return tf.contrib.training.HParams(env_lambda=env_lambda,
-                                     wrappers=standard_wrappers,
-                                     simulated_env=False)
+  return tf.contrib.training.HParams(
+      env_lambda=env_lambda, wrappers=standard_wrappers, simulated_env=False)
 
 
 class GymDiscreteProblem(video_utils.VideoProblem):
@@ -108,8 +106,8 @@ def random_skip(self):
 
   def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
     self._setup()
-    self.debug_dump_frames_path = os.path.join(
-        data_dir, self.debug_dump_frames_path)
+    self.debug_dump_frames_path = os.path.join(data_dir,
+                                               self.debug_dump_frames_path)
 
     with self._session as sess:
       self.restore_networks(sess)
@@ -157,8 +155,10 @@ def restore_networks(self, sess):
       model_saver.restore(sess, ckpt)
 
   def eval_metrics(self):
-    eval_metrics = [metrics.Metrics.ACC, metrics.Metrics.ACC_PER_SEQ,
-                    metrics.Metrics.IMAGE_RMSE]
+    eval_metrics = [
+        metrics.Metrics.ACC, metrics.Metrics.ACC_PER_SEQ,
+        metrics.Metrics.IMAGE_RMSE
+    ]
     return eval_metrics
 
   @property
@@ -170,10 +170,12 @@ def extra_reading_spec(self):
         "reward": tf.FixedLenFeature([1], tf.int64)
     }
     decoders = {
-        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
-            tensor_key="frame_number"),
-        "action": tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="action"),
-        "reward": tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="reward"),
+        "frame_number":
+            tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="frame_number"),
+        "action":
+            tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="action"),
+        "reward":
+            tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="reward"),
     }
     return data_fields, decoders
 
@@ -200,13 +202,12 @@ def env(self):
   def num_actions(self):
     return self.env.action_space.n
 
-  def collect_statistics_and_generate_debug_image(self, index,
-                                                  observation,
-                                                  reward,
-                                                  done,
-                                                  action):
+  # pylint: disable=unused-argument
+  def collect_statistics_and_generate_debug_image(self, index, observation,
+                                                  reward, done, action):
     """This generates extra statistics and debug images."""
     return None
+  # pylint: enable=unused-argument
 
   @property
   def frame_height(self):
@@ -338,19 +339,19 @@ def _setup(self):
     super(GymSimulatedDiscreteProblem, self)._setup()
 
     environment_spec = self.environment_spec
-    hparams = HParams(video_num_input_frames=
-                      environment_spec.video_num_input_frames,
-                      video_num_target_frames=
-                      environment_spec.video_num_target_frames,
-                      environment_spec=environment_spec)
+    hparams = HParams(
+        video_num_input_frames=environment_spec.video_num_input_frames,
+        video_num_target_frames=environment_spec.video_num_target_frames,
+        environment_spec=environment_spec)
 
     initial_frames_problem = environment_spec.initial_frames_problem
     dataset = initial_frames_problem.dataset(
-        tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir,
-        shuffle_files=False, hparams=hparams)
+        tf.estimator.ModeKeys.TRAIN,
+        FLAGS.data_dir,
+        shuffle_files=False,
+        hparams=hparams)
     dataset = dataset.map(lambda x: x["input_action"]).take(1)
-    input_data_iterator = (
-        dataset.batch(1).make_initializable_iterator())
+    input_data_iterator = (dataset.batch(1).make_initializable_iterator())
     self._session.run(input_data_iterator.initializer)
 
     res = self._session.run(input_data_iterator.get_next())
@@ -385,8 +386,7 @@ def get_environment_spec(self):
     env_spec.add_hparam("simulation_random_starts",
                         self.simulation_random_starts)
 
-    env_spec.add_hparam("intrinsic_reward_scale",
-                        self.intrinsic_reward_scale)
+    env_spec.add_hparam("intrinsic_reward_scale", self.intrinsic_reward_scale)
     initial_frames_problem = registry.problem(self.initial_frames_problem)
     env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
     env_spec.add_hparam("video_num_input_frames", self.num_input_frames)
@@ -403,8 +403,7 @@ def _reset_real_env(self):
     stat.episode_sim_reward = 0.0
     stat.episode_real_reward = 0.0
 
-  def collect_statistics_and_generate_debug_image(self, index,
-                                                  observation,
+  def collect_statistics_and_generate_debug_image(self, index, observation,
                                                   reward, done, action):
     stat = self.statistics
 
@@ -413,8 +412,8 @@ def collect_statistics_and_generate_debug_image(self, index,
     stat.episode_sim_reward += reward
 
     ob = np.ndarray.astype(observation, np.int)
-    err = np.ndarray.astype(np.maximum(np.abs(
-        stat.real_ob - ob, dtype=np.int) - 10, 0), np.uint8)
+    err = np.ndarray.astype(
+        np.maximum(np.abs(stat.real_ob - ob, dtype=np.int) - 10, 0), np.uint8)
     debug_im = np.concatenate([observation, stat.real_ob, err], axis=1)
 
     assert (self._internal_memory_size == self.num_testing_steps and
@@ -531,8 +530,8 @@ def num_rewards(self):
 
 
 @registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnPong(
-    GymSimulatedDiscreteProblem, GymPongRandom):
+class GymSimulatedDiscreteProblemWithAgentOnPong(GymSimulatedDiscreteProblem,
+                                                 GymPongRandom):
   """Simulated pong."""
 
   @property
@@ -562,8 +561,7 @@ def num_rewards(self):
 
 
 @registry.register_problem
-class GymDiscreteProblemWithAgentOnPong(
-    GymRealDiscreteProblem, GymPongRandom):
+class GymDiscreteProblemWithAgentOnPong(GymRealDiscreteProblem, GymPongRandom):
   pass
 
 
@@ -582,8 +580,8 @@ def num_testing_steps(self):
 
 
 @registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedLongPong(
-    GymRealDiscreteProblem, GymWrappedLongPongRandom):
+class GymDiscreteProblemWithAgentOnWrappedLongPong(GymRealDiscreteProblem,
+                                                   GymWrappedLongPongRandom):
   pass
 
 
@@ -608,8 +606,8 @@ def num_testing_steps(self):
 
 
 @registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedBreakout(
-    GymRealDiscreteProblem, GymWrappedBreakoutRandom):
+class GymDiscreteProblemWithAgentOnWrappedBreakout(GymRealDiscreteProblem,
+                                                   GymWrappedBreakoutRandom):
   pass
 
 
@@ -634,8 +632,8 @@ def num_testing_steps(self):
 
 
 @registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedPong(
-    GymRealDiscreteProblem, GymWrappedPongRandom):
+class GymDiscreteProblemWithAgentOnWrappedPong(GymRealDiscreteProblem,
+                                               GymWrappedPongRandom):
   """GymDiscreteProblemWithAgentOnWrappedPong."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading
@@ -664,8 +662,8 @@ class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
 
 
 @registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnFreeway(
-    GymSimulatedDiscreteProblem, GymFreewayRandom):
+class GymSimulatedDiscreteProblemWithAgentOnFreeway(GymSimulatedDiscreteProblem,
+                                                    GymFreewayRandom):
   """Similated freeway."""
 
   @property
@@ -678,8 +676,8 @@ def num_testing_steps(self):
 
 
 @registry.register_problem
-class GymDiscreteProblemWithAgentOnFreeway(
-    GymRealDiscreteProblem, GymFreewayRandom):
+class GymDiscreteProblemWithAgentOnFreeway(GymRealDiscreteProblem,
+                                           GymFreewayRandom):
   """Freeway with agent."""
 
   # Hard-coding num_actions, frame_height, frame_width to avoid loading

From c4c0d47c75d29040bb5a3d41e8a42e91062f9b5c Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Wed, 27 Jun 2018 11:19:06 -0700
Subject: [PATCH 0255/2720] MultiProblem class for multitask learning.

PiperOrigin-RevId: 202339795
---
 .../data_generators/multi_problem.py          | 94 +++++++++++++++++++
 1 file changed, 94 insertions(+)
 create mode 100644 tensor2tensor/data_generators/multi_problem.py

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
new file mode 100644
index 000000000..f708ca475
--- /dev/null
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -0,0 +1,94 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 class for combining multiple problems for multitask learning."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+import tensorflow as tf
+
+
+class MultiProblem(problem.Problem):
+  """MultiProblem base class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(MultiProblem, self).__init__(was_reversed, was_copy)
+    self.task_list = []
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    assert len(self.task_list) > 1
+
+    for task in self.task_list:
+      assert task.vocab_type == text_problems.VocabType.CHARACTER
+      task.generate_data(data_dir, tmp_dir, task_id)
+
+  def add_task_id(self, task_id, serialized_example):
+    """Convert example to code switching mode by adding a task id."""
+    serialized_example["targets"] = tf.concat(serialized_example["inputs"],
+                                              [task_id],
+                                              serialized_example["targets"], 0)
+    del serialized_example["inputs"]
+
+  def filepattern(self, data_dir, mode, shard=None):
+    return [task.filepattern(data_dir, mode, shard) for task in self.task_list]
+
+  def dataset(self,
+              mode,
+              data_dir=None,
+              num_threads=None,
+              output_buffer_size=None,
+              shuffle_files=None,
+              hparams=None,
+              preprocess=True,
+              dataset_split=None,
+              shard=None,
+              partition_id=0,
+              num_partitions=1,
+              max_records=-1):
+
+    datasets = []
+    is_training = mode == tf.estimator.ModeKeys.TRAIN
+
+    for task in self.task_list:
+      task_dataset = task.dataset(mode, data_dir, num_threads,
+                                  output_buffer_size, shuffle_files,
+                                  hparams, preprocess, dataset_split,
+                                  shard, partition_id, num_partitions,
+                                  max_records).repeat()
+      task_dataset = task_dataset.map(
+          # pylint: disable=cell-var-from-loop
+          lambda x: self.add_task_id(task.task_id, x),
+          num_parallel_threads=num_threads)
+      datasets.append(task_dataset)
+
+    def flatten_zip(zipped):
+      flattened = tf.data.Dataset.from_tensors(zipped[0])
+      for ex in zipped[1:]:
+        flattened.concatenate(tf.data.Dataset.from_tensors(ex))
+
+      return flattened
+
+    if is_training:
+      single_mtl_dataset = tf.data.Dataset.zip(datasets).flat_map(
+          flatten_zip)
+    else:
+      single_mtl_dataset = datasets[0]
+      for data in datasets[1:]:
+        single_mtl_dataset.concatenate(data)
+
+    return single_mtl_dataset

From 46d690025f660f83e53e581573ed4c922acecffc Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Wed, 27 Jun 2018 11:44:48 -0700
Subject: [PATCH 0256/2720] Add ImageSummary to TPU and fix sliced_gan_loss on
 TPU

PiperOrigin-RevId: 202344570
---
 tensor2tensor/layers/common_layers.py         | 328 +++++++++++-------
 tensor2tensor/layers/modalities.py            | 134 ++++---
 tensor2tensor/models/research/autoencoders.py | 197 +++++++----
 tensor2tensor/utils/t2t_model.py              |  41 ++-
 4 files changed, 434 insertions(+), 266 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index bf9b5aa7f..45857d5e9 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -99,7 +99,8 @@ def dropout_with_broadcast_dims(x, keep_prob, broadcast_dims=None, **kwargs):
     # Allow dimensions like "-1" as well.
     broadcast_dims = [dim + ndims if dim < 0 else dim for dim in broadcast_dims]
     kwargs["noise_shape"] = [
-        1 if i in broadcast_dims else shape[i] for i in range(ndims)]
+        1 if i in broadcast_dims else shape[i] for i in range(ndims)
+    ]
   return tf.nn.dropout(x, keep_prob, **kwargs)
 
 
@@ -299,7 +300,8 @@ def cumsum(x, axis=0, exclusive=False):
   ret = tf.tensordot(x, mask, axes=[[axis], [0]])
   if axis != rank - 1:
     ret = tf.transpose(
-        ret, list(range(axis)) + [rank - 1] + list(range(axis, rank - 1)))
+        ret,
+        list(range(axis)) + [rank - 1] + list(range(axis, rank - 1)))
   return ret
 
 
@@ -602,10 +604,10 @@ def tpu_conv1d(inputs, filters, kernel_size, padding="SAME", name="tpu_conv1d"):
   for i in range(kernel_size):
     shifted = tf.slice(padded, [0, i, 0], tf.shape(inputs)) if i else inputs
     shifted.set_shape(inputs.get_shape())
-    results.append(dense(
-        shifted, filters, use_bias=(i == 0), name=name + "_%d" % i))
+    results.append(
+        dense(shifted, filters, use_bias=(i == 0), name=name + "_%d" % i))
   ret = tf.add_n(results)
-  ret *= kernel_size ** -0.5
+  ret *= kernel_size**-0.5
   return ret
 
 
@@ -1456,8 +1458,7 @@ def maybe_zero_out_padding(inputs, kernel_size, nonpadding_mask):
   Returns:
     a Tensor with the same shape as inputs
   """
-  if (kernel_size != 1 and
-      kernel_size != (1, 1) and
+  if (kernel_size != 1 and kernel_size != (1, 1) and
       nonpadding_mask is not None):
     while nonpadding_mask.get_shape().ndims < inputs.get_shape().ndims:
       nonpadding_mask = tf.expand_dims(nonpadding_mask, -1)
@@ -1543,8 +1544,7 @@ def conv_relu_conv(inputs,
     A Tensor.
   """
   with tf.variable_scope(name, "conv_relu_conv", [inputs]):
-    inputs = maybe_zero_out_padding(
-        inputs, first_kernel_size, nonpadding_mask)
+    inputs = maybe_zero_out_padding(inputs, first_kernel_size, nonpadding_mask)
 
     if cache:
       if decode_loop_step is None:
@@ -1557,13 +1557,14 @@ def conv_relu_conv(inputs,
         # update_in_current_step)
         tmp_f = tf.transpose(cache["f"], perm=[1, 0, 2])
         tmp_f = tf_inplace_ops().alias_inplace_update(
-            tmp_f, decode_loop_step * tf.shape(inputs)[1],
+            tmp_f,
+            decode_loop_step * tf.shape(inputs)[1],
             tf.transpose(inputs, perm=[1, 0, 2]))
         inputs = cache["f"] = tf.transpose(tmp_f, perm=[1, 0, 2])
       inputs = cache["f"] = inputs[:, -first_kernel_size:, :]
 
-    h = tpu_conv1d(inputs, filter_size, first_kernel_size, padding=padding,
-                   name="conv1")
+    h = tpu_conv1d(
+        inputs, filter_size, first_kernel_size, padding=padding, name="conv1")
 
     if cache:
       h = h[:, -1:, :]
@@ -1572,8 +1573,8 @@ def conv_relu_conv(inputs,
     if dropout != 0.0:
       h = tf.nn.dropout(h, 1.0 - dropout)
     h = maybe_zero_out_padding(h, second_kernel_size, nonpadding_mask)
-    return tpu_conv1d(h, output_size, second_kernel_size, padding=padding,
-                      name="conv2")
+    return tpu_conv1d(
+        h, output_size, second_kernel_size, padding=padding, name="conv2")
 
 
 def sepconv_relu_sepconv(inputs,
@@ -1587,16 +1588,19 @@ def sepconv_relu_sepconv(inputs,
                          name=None):
   """Hidden layer with RELU activation followed by linear projection."""
   with tf.variable_scope(name, "sepconv_relu_sepconv", [inputs]):
-    inputs = maybe_zero_out_padding(
-        inputs, first_kernel_size, nonpadding_mask)
+    inputs = maybe_zero_out_padding(inputs, first_kernel_size, nonpadding_mask)
     if inputs.get_shape().ndims == 3:
       is_3d = True
       inputs = tf.expand_dims(inputs, 2)
     else:
       is_3d = False
     h = separable_conv(
-        inputs, filter_size, first_kernel_size, activation=tf.nn.relu,
-        padding=padding, name="conv1")
+        inputs,
+        filter_size,
+        first_kernel_size,
+        activation=tf.nn.relu,
+        padding=padding,
+        name="conv1")
     if dropout != 0.0:
       h = tf.nn.dropout(h, 1.0 - dropout)
     h = maybe_zero_out_padding(h, second_kernel_size, nonpadding_mask)
@@ -1681,23 +1685,22 @@ def gru_feedfwd(a_t, h_prev, filters, name=None):
     h_t: [batch, length, filters] hidden state
   """
 
-  with tf.variable_scope(
-      name, default_name="GRU", values=[a_t, h_prev]):
+  with tf.variable_scope(name, default_name="GRU", values=[a_t, h_prev]):
     # we use right matrix multiplication to handle batches
     # W_z and W_r have shape 2d, d. U_z U_r have shape d,d
-    z_t = (tf.sigmoid(tpu_conv1d(a_t, filters, 1, padding="SAME",
-                                 name="W_z") +
-                      tpu_conv1d(h_prev, filters, 1, padding="SAME",
-                                 name="U_z")))
-    r_t = (tf.sigmoid(tpu_conv1d(a_t, filters, 1, padding="SAME",
-                                 name="W_r") +
-                      tpu_conv1d(h_prev, filters, 1, padding="SAME",
-                                 name="U_r")))
-    h_tilde = (tf.tanh(tpu_conv1d(a_t, filters, 1, padding="SAME",
-                                  name="W") +
-                       tpu_conv1d(r_t*h_prev, filters, 1, padding="SAME",
-                                  name="U")))
-    h_t = (1. - z_t)*h_prev + z_t * h_tilde
+    z_t = (
+        tf.sigmoid(
+            tpu_conv1d(a_t, filters, 1, padding="SAME", name="W_z") +
+            tpu_conv1d(h_prev, filters, 1, padding="SAME", name="U_z")))
+    r_t = (
+        tf.sigmoid(
+            tpu_conv1d(a_t, filters, 1, padding="SAME", name="W_r") +
+            tpu_conv1d(h_prev, filters, 1, padding="SAME", name="U_r")))
+    h_tilde = (
+        tf.tanh(
+            tpu_conv1d(a_t, filters, 1, padding="SAME", name="W") +
+            tpu_conv1d(r_t * h_prev, filters, 1, padding="SAME", name="U")))
+    h_t = (1. - z_t) * h_prev + z_t * h_tilde
 
   return h_t
 
@@ -1920,11 +1923,13 @@ def padded_cross_entropy(logits,
       labels = tf.reshape(labels, [-1])
     else:
       logits, labels = pad_with_zeros(logits, labels)
-    logits = tf.reshape(logits, shape_list(labels) + [vocab_size],
-                        name="padded_cross_entropy_size_check")
+    logits = tf.reshape(
+        logits,
+        shape_list(labels) + [vocab_size],
+        name="padded_cross_entropy_size_check")
     logits = tf.cast(logits, tf.float32)
-    xent = smoothing_cross_entropy(logits, labels, vocab_size, confidence,
-                                   gaussian=gaussian)
+    xent = smoothing_cross_entropy(
+        logits, labels, vocab_size, confidence, gaussian=gaussian)
     weights = weights_fn(labels)
     if cutoff > 0.0:
       xent = tf.nn.relu(xent - cutoff)
@@ -1934,7 +1939,8 @@ def padded_cross_entropy(logits,
 
 
 def dml_loss(
-    pred, labels,
+    pred,
+    labels,
     weights_fn=weights_nonzero,  # Unused
     reduce_sum=True):
   """Discretized mixture of logistics loss.
@@ -1950,14 +1956,12 @@ def dml_loss(
   """
   del weights_fn  # Unused
   real_labels = convert_rgb_to_symmetric_real(labels)
-  dml_loss_value = discretized_mix_logistic_loss(real_labels, pred,
-                                                 sum_all=reduce_sum)
+  dml_loss_value = discretized_mix_logistic_loss(
+      real_labels, pred, sum_all=reduce_sum)
   if reduce_sum:
-    return dml_loss_value, tf.reduce_sum(tf.ones(tf.shape(labels),
-                                                 tf.float32))
+    return dml_loss_value, tf.reduce_sum(tf.ones(tf.shape(labels), tf.float32))
   else:
-    return dml_loss_value/3., tf.ones(tf.shape(dml_loss_value),
-                                      tf.float32)
+    return dml_loss_value / 3., tf.ones(tf.shape(dml_loss_value), tf.float32)
 
 
 def split_to_discretized_mix_logistic_params(inputs):
@@ -1980,8 +1984,9 @@ def split_to_discretized_mix_logistic_params(inputs):
   num_mixtures = output_dim // 10
   logits, locs, log_scales, coeffs = tf.split(
       inputs,
-      num_or_size_splits=[num_mixtures, num_mixtures * 3,
-                          num_mixtures * 3, num_mixtures * 3],
+      num_or_size_splits=[
+          num_mixtures, num_mixtures * 3, num_mixtures * 3, num_mixtures * 3
+      ],
       axis=-1)
   split_shape = [batch, height, width, num_mixtures, 3]
   locs = tf.reshape(locs, split_shape)
@@ -2033,17 +2038,18 @@ def discretized_mix_logistic_loss(labels, pred, sum_all=True):
 
   # Tile labels to broadcast compute across the mixture dimension.
   batch, height, width, num_mixtures = shape_list(logits)
-  labels = tf.tile(tf.reshape(labels, [batch, height, width, 1, 3]),
-                   [1, 1, 1, num_mixtures, 1])
+  labels = tf.tile(
+      tf.reshape(labels, [batch, height, width, 1, 3]),
+      [1, 1, 1, num_mixtures, 1])
 
   # p(x) = sigmoid((x - means_i + 1/255.)/scale_i) -
   #        sigmoid((x - means_i - 1/255.)/scale_i)
   # for each channel i. The means are linearly parameterized.
   means_0 = locs[..., 0]
   means_1 = locs[..., 1] + coeffs[..., 0] * labels[..., 0]
-  means_2 = (locs[..., 2] +
-             coeffs[..., 1] * labels[..., 0] +
-             coeffs[..., 2] * labels[..., 1])
+  means_2 = (
+      locs[..., 2] + coeffs[..., 1] * labels[..., 0] +
+      coeffs[..., 2] * labels[..., 1])
   means = tf.stack([means_0, means_1, means_2], axis=-1)
   centered_labels = labels - means
   inv_stdv = tf.exp(-log_scales)
@@ -2065,10 +2071,11 @@ def discretized_mix_logistic_loss(labels, pred, sum_all=True):
   mid_in = inv_stdv * centered_labels
   log_prob_event_approx = (
       mid_in - log_scales - 2. * tf.nn.softplus(mid_in) - np.log(127.5))
-  log_probs = tf.where(labels < -0.999, log_prob_0,
-                       tf.where(labels > 0.999, log_prob_255,
-                                tf.where(prob_event > 1e-5, log_prob_event,
-                                         log_prob_event_approx)))
+  log_probs = tf.where(
+      labels < -0.999, log_prob_0,
+      tf.where(
+          labels > 0.999, log_prob_255,
+          tf.where(prob_event > 1e-5, log_prob_event, log_prob_event_approx)))
 
   # Sum over channels and compute log-probability of each mixture.
   log_probs = tf.reduce_sum(log_probs, -1) + tf.nn.log_softmax(logits, axis=-1)
@@ -2101,10 +2108,13 @@ def sample_from_discretized_mix_logistic(pred, seed=None):
 
   # Sample mixture indicator given logits using the gumbel max trick.
   num_mixtures = shape_list(logits)[-1]
-  gumbel_noise = -tf.log(-tf.log(tf.random_uniform(
-      tf.shape(logits), minval=1e-5, maxval=1. - 1e-5, seed=seed)))
-  sel = tf.one_hot(tf.argmax(logits + gumbel_noise, -1),
-                   depth=num_mixtures, dtype=tf.float32)
+  gumbel_noise = -tf.log(-tf.log(
+      tf.random_uniform(
+          tf.shape(logits), minval=1e-5, maxval=1. - 1e-5, seed=seed)))
+  sel = tf.one_hot(
+      tf.argmax(logits + gumbel_noise, -1),
+      depth=num_mixtures,
+      dtype=tf.float32)
 
   # Select mixture component's parameters.
   sel = tf.expand_dims(sel, -1)
@@ -2254,16 +2264,19 @@ def gated_linear_unit_layer(x, name=None):
   Returns:
     x: A tensor
   """
-  with tf.variable_scope(
-      name, default_name="glu_layer", values=[x]):
+  with tf.variable_scope(name, default_name="glu_layer", values=[x]):
     depth = shape_list(x)[-1]
     x = tf.layers.dense(x, depth * 2, activation=None)
     x, gating_x = tf.split(x, 2, axis=-1)
     return x * tf.nn.sigmoid(gating_x)
 
 
-def sru_with_scan(x, num_layers=2,
-                  activation=None, initial_state=None, name=None, reuse=None):
+def sru_with_scan(x,
+                  num_layers=2,
+                  activation=None,
+                  initial_state=None,
+                  name=None,
+                  reuse=None):
   """SRU cell as in https://arxiv.org/abs/1709.02755.
 
   This implementation uses tf.scan and can incur overhead, see the full SRU
@@ -2291,22 +2304,26 @@ def sru_with_scan(x, num_layers=2,
     x = tf.reshape(x, [x_shape[0], -1, x_shape[-1]])
     x = tf.transpose(x, [1, 0, 2])  # Scan assumes time on axis 0.
     initial_state = initial_state or tf.zeros([x_shape[0], x_shape[-1]])
+
     # SRU state manipulation function.
     def next_state(cur_state, args_tup):
       cur_x_times_one_minus_f, cur_f = args_tup
       return cur_f * cur_state + cur_x_times_one_minus_f
+
     # Calculate SRU on each layer.
     for i in range(num_layers):
       # The parallel part of the SRU.
       x_orig = x
-      x, f, r = tf.split(tf.layers.dense(x, 3 * x_shape[-1],
-                                         name="kernel_%d" % i), 3, axis=-1)
+      x, f, r = tf.split(
+          tf.layers.dense(x, 3 * x_shape[-1], name="kernel_%d" % i), 3, axis=-1)
       f, r = tf.sigmoid(f), tf.sigmoid(r)
       x_times_one_minus_f = x * (1.0 - f)  # Compute in parallel for speed.
       # Calculate states.
-      c_states = tf.scan(next_state, (x_times_one_minus_f, f),
-                         initializer=initial_state,
-                         parallel_iterations=2, name="scan_%d" % i)
+      c_states = tf.scan(
+          next_state, (x_times_one_minus_f, f),
+          initializer=initial_state,
+          parallel_iterations=2,
+          name="scan_%d" % i)
       # Final output.
       if activation is not None:
         c_states = activation(c_states)
@@ -2338,8 +2355,12 @@ def __call__(self, inputs_t, state_t):
     return outputs_t, state_next
 
 
-def sru(x, num_layers=2,
-        activation=None, initial_state=None, name=None, reuse=None):
+def sru(x,
+        num_layers=2,
+        activation=None,
+        initial_state=None,
+        name=None,
+        reuse=None):
   """SRU cell as in https://arxiv.org/abs/1709.02755.
 
   As defined in the paper:
@@ -2385,8 +2406,8 @@ def sru(x, num_layers=2,
     for i in range(num_layers):
       # The parallel part of the SRU.
       x_orig = x
-      x, f, r = tf.split(tf.layers.dense(x, 3 * x_shape[-1],
-                                         name="kernel_%d" % i), 3, axis=-1)
+      x, f, r = tf.split(
+          tf.layers.dense(x, 3 * x_shape[-1], name="kernel_%d" % i), 3, axis=-1)
       f, r = tf.sigmoid(f), tf.sigmoid(r)
       x_times_one_minus_f = x * (1.0 - f)  # Compute in parallel for speed.
       # Calculate states.
@@ -2622,8 +2643,8 @@ def to_tensor(self):
     product = tf.matmul(flat_a, self.b, transpose_b=True)
     product_shape = a_shape[:-1] + [result_dim]
     product = tf.reshape(product, product_shape)
-    product.set_shape(
-        self.a.get_shape().as_list()[:-1] + [self.b.get_shape()[0]])
+    product.set_shape(self.a.get_shape().as_list()[:-1] +
+                      [self.b.get_shape()[0]])
     return product
 
 
@@ -3011,9 +3032,9 @@ def ones_matrix_band_part(rows, cols, num_lower, num_upper, out_shape=None):
       band = band.reshape(out_shape)
     band = tf.constant(band, tf.float32)
   else:
-    band = tf.matrix_band_part(tf.ones([rows, cols]),
-                               tf.cast(num_lower, tf.int64),
-                               tf.cast(num_upper, tf.int64))
+    band = tf.matrix_band_part(
+        tf.ones([rows, cols]), tf.cast(num_lower, tf.int64),
+        tf.cast(num_upper, tf.int64))
     if out_shape:
       band = tf.reshape(band, out_shape)
 
@@ -3127,56 +3148,76 @@ def batch_dense(inputs,
   if not isinstance(batch, int) or not isinstance(input_units, int):
     raise ValueError("inputs must have static dimensions 0 and 2")
   with tf.variable_scope(
-      name, default_name="batch_dense", values=[inputs],
-      reuse=reuse, dtype=inputs.dtype):
+      name,
+      default_name="batch_dense",
+      values=[inputs],
+      reuse=reuse,
+      dtype=inputs.dtype):
     if kernel_initializer is None:
       kernel_initializer = tf.random_normal_initializer(
           stddev=input_units**-0.5)
     w = tf.get_variable(
         "w", [batch, input_units, units],
-        initializer=kernel_initializer, dtype=inputs.dtype)
+        initializer=kernel_initializer,
+        dtype=inputs.dtype)
     y = tf.matmul(inputs, w)
     if activation is not None:
       y = activation(y)
     return y
 
 
-def mix(x1, x2, steps, is_training,
-        min_prob=0.0, max_prob=1.0,
-        mode="lin", simple=False, broadcast_last=False):
+def mix(x1,
+        x2,
+        steps,
+        is_training,
+        min_prob=0.0,
+        max_prob=1.0,
+        mode="lin",
+        simple=False,
+        broadcast_last=False):
   """Mix starting with x2, mixing mixing, going towards x1."""
-  if not is_training:
-    if max_prob >= 1.0:
-      return x1
-    alpha_shape = shape_list(x1)
-    if broadcast_last:
-      alpha_shape = alpha_shape[:-1] + [1]
-    alpha = tf.random_uniform(alpha_shape)
-    alpha = tf.to_float(tf.less(alpha, max_prob))
-    return alpha * x1 + (1.0 - alpha) * x2
-
-  def get_res():
-    """Create the result. Separate function to speed it up later (see below)."""
-    if mode == "lin":
-      alpha_p = inverse_lin_decay(steps)
+  with tf.name_scope("mix"):
+    if not is_training:
+      if max_prob >= 1.0:
+        return x1
+      alpha_shape = shape_list(x1)
+      if broadcast_last:
+        alpha_shape = alpha_shape[:-1] + [1]
+      alpha = tf.random_uniform(alpha_shape)
+      alpha = tf.to_float(tf.less(alpha, max_prob))
+      return alpha * x1 + (1.0 - alpha) * x2
+
+    def get_res():
+      """Create the result.
+
+      Separate function to speed it up later (see below).
+
+      Returns:
+        Tensor of mixed inputs.
+      """
+      if mode == "lin":
+        alpha_p = inverse_lin_decay(steps)
+      else:
+        alpha_p = inverse_exp_decay(steps)
+      alpha_p = alpha_p * (max_prob - min_prob) + min_prob
+      if simple:
+        return alpha_p * x1 + (1.0 - alpha_p) * x2
+      alpha_shape = shape_list(x1)
+      if broadcast_last:
+        alpha_shape = alpha_shape[:-1] + [1]
+      alpha = tf.random_uniform(alpha_shape)
+      alpha = tf.to_float(tf.less(alpha, alpha_p))
+      return alpha * x1 + (1.0 - alpha) * x2
+
+    if max_prob < 1.0:
+      return get_res()
+
+    # Prevent sampling after steps is passed to speed it up.
+    if is_on_tpu():
+      return get_res()
     else:
-      alpha_p = inverse_exp_decay(steps)
-    alpha_p = alpha_p * (max_prob - min_prob) + min_prob
-    if simple:
-      return alpha_p * x1 + (1.0 - alpha_p) * x2
-    alpha_shape = shape_list(x1)
-    if broadcast_last:
-      alpha_shape = alpha_shape[:-1] + [1]
-    alpha = tf.random_uniform(alpha_shape)
-    alpha = tf.to_float(tf.less(alpha, alpha_p))
-    return alpha * x1 + (1.0 - alpha) * x2
-
-  if max_prob < 1.0:
-    return get_res()
-
-  # Prevent sampling after steps is passed to speed it up.
-  return tf.cond(tf.less(tf.train.get_global_step(), steps),
-                 get_res, lambda: x1)
+      return tf.cond(
+          tf.less(tf.train.get_global_step(), steps), get_res, lambda: x1)
 
 
 def brelu(x):
@@ -3213,7 +3254,8 @@ def argmax_with_score(logits, axis=None):
   flat_predictions = tf.reshape(predictions, [prefix_size])
   flat_indices = tf.stack(
       [tf.range(tf.to_int64(prefix_size)),
-       tf.to_int64(flat_predictions)], axis=1)
+       tf.to_int64(flat_predictions)],
+      axis=1)
   flat_scores = tf.gather_nd(flat_logits, flat_indices)
 
   # Unflatten
@@ -3268,8 +3310,11 @@ def index_last_dim_with_indices(x, indices):
   flat_indices = tf.reshape(indices, [list_product(x_shape[:-1])])
 
   idx = tf.stack(
-      [tf.range(tf.to_int64(shape_list(flat_indices)[0])),
-       tf.to_int64(flat_indices)], axis=1)
+      [
+          tf.range(tf.to_int64(shape_list(flat_indices)[0])),
+          tf.to_int64(flat_indices)
+      ],
+      axis=1)
   flat_x_idx = tf.gather_nd(flat_x, idx)
 
   x_idx = tf.reshape(flat_x_idx, x_shape[:-1])
@@ -3311,12 +3356,14 @@ def summarize_video(video, prefix, max_outputs=1):
     return
   if video.get_shape().as_list()[1] is None:
     tf.summary.image(
-        "%s_last_frame" % prefix, tf.cast(video[:, -1, :, :, :], tf.uint8),
+        "%s_last_frame" % prefix,
+        tf.cast(video[:, -1, :, :, :], tf.uint8),
         max_outputs=max_outputs)
   else:
     for k in range(video_shape[1]):
       tf.summary.image(
-          "%s_frame_%d" % (prefix, k), tf.cast(video[:, k, :, :, :], tf.uint8),
+          "%s_frame_%d" % (prefix, k),
+          tf.cast(video[:, k, :, :, :], tf.uint8),
           max_outputs=max_outputs)
 
 
@@ -3328,9 +3375,10 @@ def time_to_channels(embedded_video):
                      "[batch, time, height, width, channels] but got one "
                      "of shape: %s" % str(video_shape))
   transposed = tf.transpose(embedded_video, [0, 2, 3, 1, 4])
-  return tf.reshape(transposed,
-                    [video_shape[0], video_shape[2], video_shape[3],
-                     video_shape[1] * video_shape[4]])
+  return tf.reshape(transposed, [
+      video_shape[0], video_shape[2], video_shape[3],
+      video_shape[1] * video_shape[4]
+  ])
 
 
 def cast_like(x, y):
@@ -3343,8 +3391,8 @@ def cast_like(x, y):
 
   cast_x = tf.cast(x, y.dtype)
   if cast_x.device != x.device:
-    tf.logging.warning("Cast for %s may induce copy from '%s' to '%s'",
-                       x.name, x.device, cast_x.device)
+    tf.logging.warning("Cast for %s may induce copy from '%s' to '%s'", x.name,
+                       x.device, cast_x.device)
   return cast_x
 
 
@@ -3374,8 +3422,12 @@ def make_even_size(x):
   return x
 
 
-def sliced_gan_loss(input1, input2, discriminator, num_vecs,
-                    do_random_vecs=True, do_tanh=True):
+def sliced_gan_loss(input1,
+                    input2,
+                    discriminator,
+                    num_vecs,
+                    do_random_vecs=True,
+                    do_tanh=True):
   """Loss inspired by the sliced WGAN paper: https://arxiv.org/abs/1804.01947.
 
   Puts input1 and input2 through the provided discriminator to get logits.
@@ -3418,7 +3470,19 @@ def get_sorted_projections(x):
       else:
         proj = tf.tanh(x)
       proj = tf.transpose(proj, [1, 0])  # [num_vecs, batch] after this.
-      values, _ = tf.nn.top_k(proj, k=batch_size, sorted=True)
+
+      if is_on_tpu():
+        proj_dtype = proj.dtype
+        proj = tf.cast(proj, tf.bfloat16)
+
+        # Currently TPU only supports 1-D top_k calls.
+        map_fn = lambda x: tf.nn.top_k(x, k=batch_size, sorted=True)[0]
+        values = tf.map_fn(map_fn, proj)
+
+        values = tf.cast(values, proj_dtype)
+      else:
+        values, _ = tf.nn.top_k(proj, k=batch_size, sorted=True)
+
       return values
 
     proj1 = get_sorted_projections(logits1)
@@ -3430,3 +3494,13 @@ def upscale(inputs, f, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR):
   """Upscaling the image by a factor of f."""
   height, width = shape_list(inputs)[1:3]
   return tf.image.resize_images(inputs, (height * f, width * f), method)
+
+
+def tpu_safe_image_summary(image):
+  if is_on_tpu():
+    # We only support float32 images at the moment due to casting complications.
+    if image.dtype != tf.float32:
+      image = tf.to_float(image)
+  else:
+    image = tf.cast(image, tf.uint8)
+  return image
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index a99a47110..e93849394 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -59,8 +59,7 @@ def targets_weights_fn(self):
           # autoregressively predicting the inputs portion, while the
           # evaluation is only done on the output
           hp.prepend_mode != "prepend_inputs_masked_attention" or
-          hp.mode != tf.estimator.ModeKeys.TRAIN
-      ):
+          hp.mode != tf.estimator.ModeKeys.TRAIN):
         weights_fn = common_layers.weights_prepend_inputs_to_targets
 
     return weights_fn
@@ -165,8 +164,8 @@ def top(self, body_output, _):
           # TODO(noam): remove this once TPU is more forgiving of extra dims.
           return logits
         else:
-          return tf.reshape(
-              logits, body_output_shape[:-1] + [1, self._vocab_size])
+          return tf.reshape(logits,
+                            body_output_shape[:-1] + [1, self._vocab_size])
 
 
 @registry.register_symbol_modality("weights_all")
@@ -216,15 +215,18 @@ class ImageModality(modality.Modality):
   def bottom(self, x):
     with tf.variable_scope(self.name):
       if not tf.contrib.eager.in_eager_mode():
-        tf.summary.image("inputs", tf.cast(x, tf.uint8), max_outputs=2)
+        tf.summary.image(
+            "inputs", common_layers.tpu_safe_image_summary(x), max_outputs=2)
       return tf.to_float(x)
 
   def targets_bottom(self, x):
     inputs = x
     with tf.variable_scope(self.name):
       if not tf.contrib.eager.in_eager_mode():
-        tf.summary.image("targets_bottom",
-                         tf.cast(inputs, tf.uint8), max_outputs=1)
+        tf.summary.image(
+            "targets_bottom",
+            common_layers.tpu_safe_image_summary(inputs),
+            max_outputs=1)
       inputs_shape = common_layers.shape_list(inputs)
       if len(inputs_shape) != 4:
         raise ValueError("Assuming images given as int tensors in the format "
@@ -239,8 +241,10 @@ def targets_bottom(self, x):
       # Let's now merge all channels that were embedded into a single vector.
       merged_size = self.PIXEL_EMBEDDING_SIZE * inputs_shape[3]
       embedded = tf.reshape(embedded, inputs_shape[:3] + [merged_size])
-      merged = tf.layers.dense(embedded, self._body_input_depth,
-                               name="merge_pixel_embedded_channels")
+      merged = tf.layers.dense(
+          embedded,
+          self._body_input_depth,
+          name="merge_pixel_embedded_channels")
       return merged
 
   def top(self, body_output, _):
@@ -253,8 +257,11 @@ def top(self, body_output, _):
       res = tf.layers.dense(body_output, self.top_dimensionality * num_channels)
       res = tf.reshape(res, reshape_shape)
       if not tf.get_variable_scope().reuse:
-        res_argmax = tf.cast(tf.argmax(res, axis=-1), tf.uint8)
-        tf.summary.image("result", res_argmax, max_outputs=1)
+        res_argmax = tf.argmax(res, axis=-1)
+        tf.summary.image(
+            "result",
+            common_layers.tpu_safe_image_summary(res_argmax),
+            max_outputs=1)
       return res
 
   def loss(self, top_out, targets):
@@ -293,18 +300,22 @@ def bottom_compress(self, inputs, name="bottom"):
       inputs = tf.to_float(inputs)
       hp = self._model_hparams
       if hp.mode != tf.estimator.ModeKeys.PREDICT:
-        tf.summary.image("inputs", inputs, max_outputs=2)
+        tf.summary.image(
+            "inputs",
+            common_layers.tpu_safe_image_summary(inputs),
+            max_outputs=2)
       inputs = common_layers.convert_rgb_to_symmetric_real(inputs)
       ishape = common_layers.shape_list(inputs)
       inputs = tf.reshape(inputs, [-1, ishape[1], ishape[2] * ishape[3], 1])
       inputs.set_shape([None, None, None, 1])
       # We compress RGB intensities for each pixel using a conv.
-      x = tf.layers.conv2d(inputs,
-                           self._body_input_depth, (1, self.num_channels),
-                           padding="VALID",
-                           strides=(1, self.num_channels),
-                           activation=tf.nn.relu,
-                           name="conv_input")
+      x = tf.layers.conv2d(
+          inputs,
+          self._body_input_depth, (1, self.num_channels),
+          padding="VALID",
+          strides=(1, self.num_channels),
+          activation=tf.nn.relu,
+          name="conv_input")
       x.set_shape([None, None, None, self._body_input_depth])
       return x
 
@@ -322,16 +333,15 @@ def top(self, body_output, _):
       batch = common_layers.shape_list(body_output)[0]
       x = tf.layers.conv2d(
           body_output,
-          hidden_dim*channels, (1, 1),
+          hidden_dim * channels, (1, 1),
           strides=(1, 1),
           padding="VALID",
           activation=tf.nn.relu,
           name="decompress_conv")
       x = tf.reshape(x, [batch, img_len, img_len * channels, hidden_dim])
       x = common_layers.layer_preprocess(x, self._model_hparams)
-      x = tf.layers.dense(x, 256,
-                          use_bias=True, activation=None,
-                          name="output_conv")
+      x = tf.layers.dense(
+          x, 256, use_bias=True, activation=None, name="output_conv")
       x = tf.reshape(x,
                      [-1, img_len, img_len, channels, self.top_dimensionality])
       return x
@@ -348,7 +358,10 @@ def top(self, body_output, _):
 class ImageChannelEmbeddingsBottom(modality.Modality):
   """Modality for images using channel compression for generation."""
 
-  def get_channel_embeddings(self, io_depth, targets, hidden_size,
+  def get_channel_embeddings(self,
+                             io_depth,
+                             targets,
+                             hidden_size,
                              name="channel"):
     """Get separate embedding for each of the channels."""
     targets_split = tf.split(targets, io_depth, axis=3)
@@ -371,18 +384,17 @@ def targets_bottom(self, x):
     io_depth = self._model_hparams.num_channels
     tshape = common_layers.shape_list(inputs)
     hidden_size = self._model_hparams.hidden_size
-    target_embeddings = self.get_channel_embeddings(
-        io_depth, inputs, hidden_size, "input_bottom")
+    target_embeddings = self.get_channel_embeddings(io_depth, inputs,
+                                                    hidden_size, "input_bottom")
     return tf.reshape(target_embeddings,
-                      [tshape[0], tshape[1], tshape[2]*io_depth, hidden_size])
+                      [tshape[0], tshape[1], tshape[2] * io_depth, hidden_size])
 
   def top(self, body_output, _):
     with tf.variable_scope(self.name):
       img_len = self._model_hparams.img_len
       channels = self._model_hparams.num_channels
-      x = tf.layers.dense(body_output, 256,
-                          use_bias=True, activation=None,
-                          name="output_conv")
+      x = tf.layers.dense(
+          body_output, 256, use_bias=True, activation=None, name="output_conv")
       x = tf.reshape(x,
                      [-1, img_len, img_len, channels, self.top_dimensionality])
       return x
@@ -495,10 +507,10 @@ def bottom(self, x):
       inputs = tf.reshape(inputs, inputs_shape)
       # Concatenate the time dimension on channels for image models to work.
       transposed = tf.transpose(inputs, [0, 2, 3, 1, 4])
-      return tf.reshape(
-          transposed,
-          [inputs_shape[0], inputs_shape[2], inputs_shape[3],
-           inputs_shape[1] * inputs_shape[4]])
+      return tf.reshape(transposed, [
+          inputs_shape[0], inputs_shape[2], inputs_shape[3],
+          inputs_shape[1] * inputs_shape[4]
+      ])
 
   def targets_bottom(self, x, summary_prefix="targets_bottom"):  # pylint: disable=arguments-differ
     inputs = x
@@ -516,8 +528,10 @@ def targets_bottom(self, x, summary_prefix="targets_bottom"):  # pylint: disable
       merged_size = self.PIXEL_EMBEDDING_SIZE * inputs_shape[4]
       embedded = tf.reshape(embedded, inputs_shape[:4] + [merged_size])
       transposed = common_layers.time_to_channels(embedded)
-      return tf.layers.dense(transposed, self._body_input_depth,
-                             name="merge_pixel_embedded_frames")
+      return tf.layers.dense(
+          transposed,
+          self._body_input_depth,
+          name="merge_pixel_embedded_frames")
 
   def top(self, body_output, _):
     num_channels = self._model_hparams.problem.num_channels
@@ -526,14 +540,16 @@ def top(self, body_output, _):
       body_output_shape = common_layers.shape_list(body_output)
       reshape_shape = body_output_shape[:3]
       reshape_shape.extend([num_channels, num_frames, self.top_dimensionality])
-      res = tf.layers.dense(
-          body_output, self.top_dimensionality * num_channels * num_frames)
+      res = tf.layers.dense(body_output,
+                            self.top_dimensionality * num_channels * num_frames)
       res = tf.reshape(res, reshape_shape)
       res = tf.transpose(res, [0, 4, 1, 2, 3, 5])
       if not tf.get_variable_scope().reuse:
-        res_argmax = tf.cast(tf.argmax(res[:, -1, :, :, :, :], axis=-1),
-                             tf.uint8)
-        tf.summary.image("result", res_argmax, max_outputs=1)
+        res_argmax = tf.argmax(res[:, -1, :, :, :, :], axis=-1)
+        tf.summary.image(
+            "result",
+            common_layers.tpu_safe_image_summary(res_argmax),
+            max_outputs=1)
       return res
 
   def loss(self, top_out, targets):
@@ -569,12 +585,14 @@ def bottom(self, x):
       common_layers.summarize_video(inputs, "bottom")
       # Embed bitwise.
       assert self.top_dimensionality == 256
-      embedded = discretization.int_to_bit_embed(
-          inputs, 8, self.PIXEL_EMBEDDING_SIZE)
+      embedded = discretization.int_to_bit_embed(inputs, 8,
+                                                 self.PIXEL_EMBEDDING_SIZE)
       # Transpose and project.
       transposed = common_layers.time_to_channels(embedded)
-      return tf.layers.dense(transposed, self._body_input_depth,
-                             name="merge_pixel_embedded_frames")
+      return tf.layers.dense(
+          transposed,
+          self._body_input_depth,
+          name="merge_pixel_embedded_frames")
 
   def targets_bottom(self, x):  # pylint: disable=arguments-differ
     inputs = x
@@ -582,12 +600,14 @@ def targets_bottom(self, x):  # pylint: disable=arguments-differ
       common_layers.summarize_video(inputs, "targets_bottom")
       # Embed bitwise.
       assert self.top_dimensionality == 256
-      embedded = discretization.int_to_bit_embed(
-          inputs, 8, self.PIXEL_EMBEDDING_SIZE)
+      embedded = discretization.int_to_bit_embed(inputs, 8,
+                                                 self.PIXEL_EMBEDDING_SIZE)
       # Transpose and project.
       transposed = common_layers.time_to_channels(embedded)
-      return tf.layers.dense(transposed, self._body_input_depth,
-                             name="merge_pixel_embedded_frames")
+      return tf.layers.dense(
+          transposed,
+          self._body_input_depth,
+          name="merge_pixel_embedded_frames")
 
 
 @registry.register_video_modality("l1")
@@ -603,8 +623,11 @@ def top(self, body_output, _):
       res = tf.reshape(res, body_output_shape[:3] + [num_channels, num_frames])
       res = tf.transpose(res, [0, 4, 1, 2, 3])  # Move frames next to batch.
       if not tf.get_variable_scope().reuse:
-        res_argmax = tf.cast(res[:, -1, :, :, :], tf.uint8)
-        tf.summary.image("result", res_argmax, max_outputs=1)
+        res_argmax = res[:, -1, :, :, :]
+        tf.summary.image(
+            "result",
+            common_layers.tpu_safe_image_summary(res_argmax),
+            max_outputs=1)
       return tf.expand_dims(res, axis=-1)  # Add an axis like in perplexity.
 
   @property
@@ -724,12 +747,13 @@ def loss(self, top_out, targets):
         logits,
         targets,
         self._model_hparams.label_smoothing,
-        weights_fn=self.targets_weights_fn,)
+        weights_fn=self.targets_weights_fn,
+    )
     xent = tf.squeeze(xent, [2, 3])
     weights = tf.squeeze(xent, [2, 3])
     # average loss over all labels
-    loss = (tf.reduce_sum(xent, axis=1)
-            / (tf.reduce_sum(weights, axis=1) + 1e-8))
+    loss = (
+        tf.reduce_sum(xent, axis=1) / (tf.reduce_sum(weights, axis=1) + 1e-8))
     return tf.reduce_mean(loss)
 
 
@@ -782,8 +806,8 @@ def top_is_pointwise(self):
 
   def bottom(self, x):
     with tf.variable_scope("real"):
-      return tf.layers.dense(tf.to_float(x), self._body_input_depth,
-                             name="bottom")
+      return tf.layers.dense(
+          tf.to_float(x), self._body_input_depth, name="bottom")
 
   def top(self, body_output, _):
     with tf.variable_scope("real"):
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 302d93c69..5eaaf9c47 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -66,27 +66,26 @@ def discriminator(self, x, is_training):
     """
     hparams = self.hparams
     with tf.variable_scope(
-        "discriminator",
-        initializer=tf.random_normal_initializer(stddev=0.02)):
+        "discriminator", initializer=tf.random_normal_initializer(stddev=0.02)):
       batch_size, height, width = common_layers.shape_list(x)[:3]
       # Mapping x from [bs, h, w, c] to [bs, 1]
-      net = tf.layers.conv2d(x, 64, (4, 4), strides=(2, 2),
-                             padding="SAME", name="d_conv1")
+      net = tf.layers.conv2d(
+          x, 64, (4, 4), strides=(2, 2), padding="SAME", name="d_conv1")
       # [bs, h/2, w/2, 64]
       net = lrelu(net)
-      net = tf.layers.conv2d(net, 128, (4, 4), strides=(2, 2),
-                             padding="SAME", name="d_conv2")
+      net = tf.layers.conv2d(
+          net, 128, (4, 4), strides=(2, 2), padding="SAME", name="d_conv2")
       # [bs, h/4, w/4, 128]
       if hparams.discriminator_batchnorm:
-        net = tf.layers.batch_normalization(net, training=is_training,
-                                            momentum=0.999, name="d_bn2")
+        net = tf.layers.batch_normalization(
+            net, training=is_training, momentum=0.999, name="d_bn2")
       net = lrelu(net)
       size = height * width
       net = tf.reshape(net, [batch_size, size * 8])  # [bs, h * w * 8]
       net = tf.layers.dense(net, 1024, name="d_fc3")  # [bs, 1024]
       if hparams.discriminator_batchnorm:
-        net = tf.layers.batch_normalization(net, training=is_training,
-                                            momentum=0.999, name="d_bn3")
+        net = tf.layers.batch_normalization(
+            net, training=is_training, momentum=0.999, name="d_bn3")
       net = lrelu(net)
       return net
 
@@ -113,8 +112,13 @@ def encoder(self, x):
       for i in range(hparams.num_hidden_layers):
         x = self.make_even_size(x)
         x = tf.layers.conv2d(
-            x, hparams.hidden_size * 2**(i + 1), kernel, strides=strides,
-            padding="SAME", activation=common_layers.belu, name="conv_%d" % i)
+            x,
+            hparams.hidden_size * 2**(i + 1),
+            kernel,
+            strides=strides,
+            padding="SAME",
+            activation=common_layers.belu,
+            name="conv_%d" % i)
         x = common_layers.layer_norm(x)
       return x
 
@@ -126,8 +130,13 @@ def decoder(self, x):
       for i in range(hparams.num_hidden_layers):
         j = hparams.num_hidden_layers - i - 1
         x = tf.layers.conv2d_transpose(
-            x, hparams.hidden_size * 2**j, kernel, strides=strides,
-            padding="SAME", activation=common_layers.belu, name="deconv_%d" % j)
+            x,
+            hparams.hidden_size * 2**j,
+            kernel,
+            strides=strides,
+            padding="SAME",
+            activation=common_layers.belu,
+            name="deconv_%d" % j)
         x = common_layers.layer_norm(x)
       return x
 
@@ -148,13 +157,13 @@ def body(self, features):
       b = common_layers.mix(b, x, hparams.bottleneck_warmup_steps, is_training)
       if hparams.gan_loss_factor != 0.0:
         # Add a purely sampled batch on which we'll compute the GAN loss.
-        g = self.unbottleneck(self.sample(), common_layers.shape_list(x)[-1],
-                              reuse=True)
+        g = self.unbottleneck(
+            self.sample(), common_layers.shape_list(x)[-1], reuse=True)
         b = tf.concat([g, b], axis=0)
       # With probability bottleneck_max_prob use the bottleneck, otherwise x.
       if hparams.bottleneck_max_prob < -1.0:
-        x = tf.where(tf.less(tf.random_uniform([]),
-                             hparams.bottleneck_max_prob), b, x)
+        x = tf.where(
+            tf.less(tf.random_uniform([]), hparams.bottleneck_max_prob), b, x)
       else:
         x = b
     else:
@@ -177,28 +186,34 @@ def body(self, features):
       # Split back if we added a purely sampled batch.
       res_gan, res = tf.split(res, 2, axis=0)
       num_channels = self.hparams.problem.num_channels
-      res_rgb = common_layers.convert_real_to_rgb(tf.nn.sigmoid(
-          tf.layers.dense(res_gan, num_channels, name="gan_rgb")))
-      tf.summary.image("gan", tf.cast(res_rgb, tf.uint8), max_outputs=1)
+      res_rgb = common_layers.convert_real_to_rgb(
+          tf.nn.sigmoid(tf.layers.dense(res_gan, num_channels, name="gan_rgb")))
+      tf.summary.image(
+          "gan", common_layers.tpu_safe_image_summary(res_rgb), max_outputs=1)
       orig_rgb = tf.to_float(features["targets_raw"])
+
       def discriminate(x):
         return self.discriminator(x, is_training=is_training)
-      gan_loss = common_layers.sliced_gan_loss(
-          orig_rgb, reverse_gradient(res_rgb),
-          discriminate, self.hparams.num_sliced_vecs)
+
+      gan_loss = common_layers.sliced_gan_loss(orig_rgb,
+                                               reverse_gradient(res_rgb),
+                                               discriminate,
+                                               self.hparams.num_sliced_vecs)
       gan_loss *= hparams.gan_loss_factor
     # Mix the final result and return.
     res = common_layers.mix(res, features["targets"],
                             hparams.bottleneck_warmup_steps // 2, is_training)
-    return res, {"bottleneck_loss": b_loss, "gan_loss": - gan_loss}
+    return res, {"bottleneck_loss": b_loss, "gan_loss": -gan_loss}
 
   def sample(self, features=None, shape=None):
     del features, shape
     hp = self.hparams
     div_x = 2**hp.num_hidden_layers
     div_y = 1 if self.is1d else 2**hp.num_hidden_layers
-    size = [hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
-            hp.bottleneck_bits]
+    size = [
+        hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
+        hp.bottleneck_bits
+    ]
     # Sample in [-1, 1] as the bottleneck is under tanh.
     return 2.0 * tf.random_uniform(size) - 1.0
 
@@ -229,8 +244,7 @@ def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
       num_channels = 1
     if "targets" not in features:
       features["targets"] = tf.zeros(
-          [self.hparams.batch_size, 1, 1, num_channels],
-          dtype=tf.int32)
+          [self.hparams.batch_size, 1, 1, num_channels], dtype=tf.int32)
     logits, _ = self(features)  # pylint: disable=not-callable
     samples = tf.argmax(logits, axis=-1)
 
@@ -294,9 +308,12 @@ def body(self, features):
       assert hparams.mode == tf.estimator.ModeKeys.PREDICT
       features["targets"] = tf.zeros_like(basic_result)
     targets_dropout = common_layers.mix(
-        features["targets"], tf.zeros_like(basic_result),
-        hparams.bottleneck_warmup_steps, is_training,
-        max_prob=1.0 - hparams.autoregressive_dropout, broadcast_last=True)
+        features["targets"],
+        tf.zeros_like(basic_result),
+        hparams.bottleneck_warmup_steps,
+        is_training,
+        max_prob=1.0 - hparams.autoregressive_dropout,
+        broadcast_last=True)
     # Sometimes it's useful to look at non-autoregressive evals.
     if (hparams.mode == tf.estimator.ModeKeys.EVAL and
         hparams.autoregressive_eval_pure_autoencoder):
@@ -311,24 +328,36 @@ def body(self, features):
       concat1d = common_layers.shift_right_3d(concat1d)
     # The autoregressive part depends on the mode.
     if hparams.autoregressive_mode == "conv3":
-      res = common_layers.conv1d(concat1d, shape[3], 3, padding="LEFT",
-                                 activation=common_layers.belu,
-                                 name="autoregressive_conv3")
+      res = common_layers.conv1d(
+          concat1d,
+          shape[3],
+          3,
+          padding="LEFT",
+          activation=common_layers.belu,
+          name="autoregressive_conv3")
       return tf.reshape(res, shape), losses
     if hparams.autoregressive_mode == "conv5":
-      res = common_layers.conv1d(concat1d, shape[3], 5, padding="LEFT",
-                                 activation=common_layers.belu,
-                                 name="autoregressive_conv5")
+      res = common_layers.conv1d(
+          concat1d,
+          shape[3],
+          5,
+          padding="LEFT",
+          activation=common_layers.belu,
+          name="autoregressive_conv5")
       return tf.reshape(res, shape), losses
     if hparams.autoregressive_mode == "sru":
-      res = common_layers.conv1d(concat1d, shape[3], 3, padding="LEFT",
-                                 activation=common_layers.belu,
-                                 name="autoregressive_sru_conv3")
+      res = common_layers.conv1d(
+          concat1d,
+          shape[3],
+          3,
+          padding="LEFT",
+          activation=common_layers.belu,
+          name="autoregressive_sru_conv3")
       res = common_layers.sru(res)
       return tf.reshape(res, shape), losses
 
-    raise ValueError("Unsupported autoregressive mode: %s"
-                     % hparams.autoregressive_mode)
+    raise ValueError(
+        "Unsupported autoregressive mode: %s" % hparams.autoregressive_mode)
 
   def infer(self, features, *args, **kwargs):
     """Produce predictions from the model by sampling."""
@@ -347,11 +376,9 @@ def infer(self, features, *args, **kwargs):
       num_channels = 1
     if "targets" not in features:
       features["targets"] = tf.zeros(
-          [self.hparams.batch_size, 1, 1, num_channels],
-          dtype=tf.int32)
+          [self.hparams.batch_size, 1, 1, num_channels], dtype=tf.int32)
     logits, _ = self(features)  # pylint: disable=not-callable
-    samples = common_layers.sample_with_temperature(
-        logits, 0.0)
+    samples = common_layers.sample_with_temperature(logits, 0.0)
     shape = common_layers.shape_list(samples)
 
     # Sample again if requested for the autoregressive part.
@@ -371,8 +398,8 @@ def infer(self, features, *args, **kwargs):
         samples = common_layers.sample_with_temperature(
             logits, self.hparams.sampling_temp)
         samples1d = tf.reshape(samples, [shape[0], -1, shape[3]])
-        samples1d = tf.concat([old_samples1d[:, :i, :], samples1d[:, i:, :]],
-                              axis=1)
+        samples1d = tf.concat(
+            [old_samples1d[:, :i, :], samples1d[:, i:, :]], axis=1)
         samples = tf.reshape(samples1d, shape)
 
     # Restore inputs to not confuse Estimator in edge cases.
@@ -394,9 +421,12 @@ def dropout(self, x):
     # return tf.nn.dropout(x, 1.0 - self.hparams.dropout)
     is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
     return common_layers.mix(
-        tf.zeros_like(x), x,
-        self.hparams.bottleneck_warmup_steps, is_training,
-        max_prob=self.hparams.dropout, broadcast_last=True)
+        tf.zeros_like(x),
+        x,
+        self.hparams.bottleneck_warmup_steps,
+        is_training,
+        max_prob=self.hparams.dropout,
+        broadcast_last=True)
 
   def encoder(self, x):
     with tf.variable_scope("encoder"):
@@ -411,7 +441,10 @@ def encoder(self, x):
         residual_conv = tf.layers.separable_conv2d
       # Input embedding with a non-zero bias for uniform inputs.
       x = tf.layers.dense(
-          x, hparams.hidden_size, name="embed", activation=common_layers.belu,
+          x,
+          hparams.hidden_size,
+          name="embed",
+          activation=common_layers.belu,
           bias_initializer=tf.random_normal_initializer(stddev=0.01))
       x = common_attention.add_timing_signal_nd(x)
       # Down-convolutions.
@@ -422,8 +455,13 @@ def encoder(self, x):
           filters = hparams.hidden_size * 2**(i + 1)
           filters = min(filters, hparams.max_hidden_size)
           x = tf.layers.conv2d(
-              x, filters, kernel, strides=strides,
-              padding="SAME", activation=common_layers.belu, name="strided")
+              x,
+              filters,
+              kernel,
+              strides=strides,
+              padding="SAME",
+              activation=common_layers.belu,
+              name="strided")
           y = x
           for r in range(hparams.num_residual_layers):
             residual_filters = filters
@@ -431,8 +469,11 @@ def encoder(self, x):
               residual_filters = int(
                   filters * hparams.residual_filter_multiplier)
             y = residual_conv(
-                y, residual_filters, residual_kernel,
-                padding="SAME", activation=common_layers.belu,
+                y,
+                residual_filters,
+                residual_kernel,
+                padding="SAME",
+                activation=common_layers.belu,
                 name="residual_%d" % r)
           x += tf.nn.dropout(y, 1.0 - hparams.residual_dropout)
           x = common_layers.layer_norm(x)
@@ -458,8 +499,13 @@ def decoder(self, x):
           j = hparams.num_hidden_layers - i - 1
           filters = hparams.hidden_size * 2**j
           x = tf.layers.conv2d_transpose(
-              x, filters, kernel, strides=strides,
-              padding="SAME", activation=common_layers.belu, name="strided")
+              x,
+              filters,
+              kernel,
+              strides=strides,
+              padding="SAME",
+              activation=common_layers.belu,
+              name="strided")
           y = x
           for r in range(hparams.num_residual_layers):
             residual_filters = filters
@@ -467,8 +513,11 @@ def decoder(self, x):
               residual_filters = int(
                   filters * hparams.residual_filter_multiplier)
             y = residual_conv(
-                y, residual_filters, residual_kernel,
-                padding="SAME", activation=common_layers.belu,
+                y,
+                residual_filters,
+                residual_kernel,
+                padding="SAME",
+                activation=common_layers.belu,
                 name="residual_%d" % r)
           x += tf.nn.dropout(y, 1.0 - hparams.residual_dropout)
           x = common_layers.layer_norm(x)
@@ -497,8 +546,10 @@ def sample(self, features=None):
     hp = self.hparams
     div_x = 2**hp.num_hidden_layers
     div_y = 1 if self.is1d else 2**hp.num_hidden_layers
-    size = [hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
-            hp.bottleneck_bits]
+    size = [
+        hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
+        hp.bottleneck_bits
+    ]
     rand = tf.random_uniform(size)
     return 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
 
@@ -532,8 +583,10 @@ def sample(self, features=None):
     hp = self.hparams
     div_x = 2**hp.num_hidden_layers
     div_y = 1 if self.is1d else 2**hp.num_hidden_layers
-    size = [hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
-            hp.bottleneck_bits]
+    size = [
+        hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
+        hp.bottleneck_bits
+    ]
     rand = tf.random_uniform(size)
     res = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
     # If you want to set some first bits to a fixed value, do this:
@@ -612,8 +665,10 @@ def full_stack(self, b, x_size, bottleneck_bits, losses, is_training, i):
     losses["stack%d_loss" % i] = self.stack_loss(b, b_pred, "step%d" % i)
     b_shape = common_layers.shape_list(b)
     if is_training:
-      b1 = tf.cond(tf.less(tf.random_uniform([]), 0.5),
-                   lambda: b, lambda: b1)
+      condition = tf.less(tf.random_uniform([]), 0.5)
+      condition = tf.reshape(condition, [1] * len(b.shape))
+      condition = tf.tile(condition, b.shape)
+      b1 = tf.where(condition, b, b1)
     return tf.reshape(b1, b_shape)
 
   def body(self, features):
@@ -637,14 +692,14 @@ def body(self, features):
       # Bottleneck (mix during early training, not too important but stable).
       b, b_loss = self.bottleneck(x)
       losses = {"bottleneck0_loss": b_loss}
-      b = self.full_stack(b, 2 * x_size, 2 * hparams.bottleneck_bits,
-                          losses, is_training, num_stacks - 1)
+      b = self.full_stack(b, 2 * x_size, 2 * hparams.bottleneck_bits, losses,
+                          is_training, num_stacks - 1)
       b = self.unbottleneck(b, x_size)
       b = common_layers.mix(b, x, hparams.bottleneck_warmup_steps, is_training)
       # With probability bottleneck_max_prob use the bottleneck, otherwise x.
       if hparams.bottleneck_max_prob < 1.0:
-        x = tf.where(tf.less(tf.random_uniform([]),
-                             hparams.bottleneck_max_prob), b, x)
+        x = tf.where(
+            tf.less(tf.random_uniform([]), hparams.bottleneck_max_prob), b, x)
       else:
         x = b
     else:
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 8f98ce619..63dd73e9d 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1145,11 +1145,7 @@ def make_estimator_model_fn(model_name,
                               use_tpu=False):
     model_cls = registry.model(model_name)
 
-    def wrapping_model_fn(features,
-                          labels,
-                          mode,
-                          params=None,
-                          config=None):
+    def wrapping_model_fn(features, labels, mode, params=None, config=None):
       return model_cls.estimator_model_fn(
           hparams,
           features,
@@ -1539,19 +1535,32 @@ def _create_host_call(model_dir):
   gs_t = tf.reshape(tf.to_int32(tf.train.get_global_step()), [1])
   summary_kwargs = collections.OrderedDict()
   for t in summaries:
-    if t.op.type != "ScalarSummary":
+    if t.op.type not in ["ScalarSummary", "ImageSummary"]:
+      tf.logging.warn("Ignoring unsupported tf.Summary type %s" % t.op.type)
       continue
 
     name = t.op.name
     tensor = t.op.inputs[1]
-    assert tensor.shape.is_compatible_with([])
-    if tensor.dtype == tf.int64:
-      tensor = tf.to_int32(tensor)
-    summary_kwargs[name] = tf.reshape(tensor, [1])
+    if t.op.type == "ScalarSummary":
+      assert tensor.shape.is_compatible_with([])
+      if tensor.dtype == tf.int64:
+        tensor = tf.to_int32(tensor)
+      summary_kwargs["ScalarSummary" + name] = tf.reshape(tensor, [1])
+    elif t.op.type == "ImageSummary":
+      # TODO(aidangomez): as we move to support more types, update
+      # common_layers.tpu_safe_image_summary
+      if tensor.dtype != tf.float32:
+        tf.logging.warn(
+            "Currently T2T on TPU only supports ImageSummary of "
+            "tf.float32-type Tensors. Skipping Tensor "
+            "%s with dtype %s..." % (tensor.name, tensor.dtype))
+        continue
+      # tensor = tf.to_float(tensor)
+      summary_kwargs["ImageSummary" + name] = tensor
   summary_kwargs["global_step"] = gs_t
 
   def host_call_fn(**kwargs):
-    """Training host call. Creates scalar summaries for training metrics.
+    """Training host call. Creates summaries for training metrics.
 
     Args:
       **kwargs: Dict of {str: Tensor} , with `Tensor` of shape `[batch]`. Must
@@ -1563,9 +1572,15 @@ def host_call_fn(**kwargs):
     gs = tf.to_int64(kwargs.pop("global_step")[0])
     with tf.contrib.summary.create_file_writer(model_dir).as_default():
       with tf.contrib.summary.always_record_summaries():
+        # We need to use tf.contrib.summary in order to feed the `step`.
         for name, value in sorted(six.iteritems(kwargs)):
-          tf.contrib.summary.scalar(
-              name, tf.reduce_mean(tf.to_float(value)), step=gs)
+          if name.startswith("ScalarSummary"):
+            name = name.strip("ScalarSummary")
+            tf.contrib.summary.scalar(
+                name, tf.reduce_mean(tf.to_float(value)), step=gs)
+          elif name.startswith("ImageSummary"):
+            name = name.strip("ImageSummary")
+            tf.contrib.summary.image(name, value, step=gs)
 
         return tf.contrib.summary.all_summary_ops()
 

From 96d6e35d76d5bed83cdd5e4730450a3a34a37582 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?B=C5=82a=C5=BCej=20O?= <blazej.osinski@codilime.com>
Date: Wed, 27 Jun 2018 22:35:02 +0200
Subject: [PATCH 0257/2720] Small fix to StackAndSkipWrapper resets.

---
 tensor2tensor/rl/envs/tf_atari_wrappers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 6119aa657..570a26915 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -190,7 +190,7 @@ def _reset_non_empty(self, indices):
     assign_op = tf.scatter_update(self._observ, indices, tf.tile(
         new_values, inx))
     with tf.control_dependencies([assign_op]):
-      return tf.identity(self.observ)
+      return tf.gather(self.observ, indices)
 
 
 class TimeLimitWrapper(WrapperBase):

From ea576658c608d8b805bbe64c1c85814a96b879b9 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 27 Jun 2018 16:39:21 -0700
Subject: [PATCH 0258/2720] Raise exceptions when trying to decode from a
 packed dataset.

PiperOrigin-RevId: 202391946
---
 tensor2tensor/models/transformer.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index c76897eee..9a28aeef2 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -282,6 +282,11 @@ def _fast_decode_tpu(self,
     """
     if self._num_datashards != 1:
       raise NotImplementedError("Fast decoding only supports a single shard.")
+    if "targets_segmentation" in features:
+      raise NotImplementedError(
+          "Decoding not supported on packed datasets "
+          " If you want to decode from a dataset, use the non-packed version"
+          " of the dataset when decoding.")
     dp = self._data_parallelism
     hparams = self._hparams
     target_modality = self._problem_hparams.target_modality
@@ -506,7 +511,11 @@ def _fast_decode(self,
     dp = self._data_parallelism
     hparams = self._hparams
     target_modality = self._problem_hparams.target_modality
-
+    if "targets_segmentation" in features:
+      raise NotImplementedError(
+          "Decoding not supported on packed datasets "
+          " If you want to decode from a dataset, use the non-packed version"
+          " of the dataset when decoding.")
     if self.has_input:
       inputs = features["inputs"]
       if target_modality.is_class_modality:

From c4d81d73a7c63fd2a363c705b17c6733333bb306 Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Wed, 27 Jun 2018 17:32:36 -0700
Subject: [PATCH 0259/2720] Revert support for ImageSummaries due to immense
 slowdown

PiperOrigin-RevId: 202398898
---
 tensor2tensor/utils/t2t_model.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 63dd73e9d..d3ba09e44 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1535,7 +1535,9 @@ def _create_host_call(model_dir):
   gs_t = tf.reshape(tf.to_int32(tf.train.get_global_step()), [1])
   summary_kwargs = collections.OrderedDict()
   for t in summaries:
-    if t.op.type not in ["ScalarSummary", "ImageSummary"]:
+    # TODO(aidangomez): enable ImageSummary support when we have a faster method
+    # see @shibow's comment in cl/202344570
+    if t.op.type not in ["ScalarSummary"]:
       tf.logging.warn("Ignoring unsupported tf.Summary type %s" % t.op.type)
       continue
 

From 82726e2708f7dd88abdd7c484c5cb7bb86cf7ede Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 27 Jun 2018 22:27:35 -0700
Subject: [PATCH 0260/2720] Fix case where `ids` is empty.

PiperOrigin-RevId: 202423317
---
 tensor2tensor/data_generators/text_encoder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index becc7737c..7d590babd 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -1057,6 +1057,6 @@ def decode(self, ids, strip_extraneous=False):
 def strip_ids(ids, ids_to_strip):
   """Strip ids_to_strip from the end ids."""
   ids = list(ids)
-  while ids[-1] in ids_to_strip:
+  while ids and ids[-1] in ids_to_strip:
     ids.pop()
   return ids

From 9222b5ec6da198da6ed969adafc7dcfbb06a1456 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Thu, 28 Jun 2018 09:43:25 +0200
Subject: [PATCH 0261/2720] removing StackAndSkipWrapper

---
 tensor2tensor/data_generators/gym_problems.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 194f302ae..766ac9ab6 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -52,7 +52,7 @@
 
 def standard_atari_env_spec(env):
   """Parameters of environment specification."""
-  standard_wrappers = [[tf_atari_wrappers.StackAndSkipWrapper, {"skip": 4}]]
+  standard_wrappers = [[tf_atari_wrappers.MaxAndSkipWrapper, {"skip": 4}]]
   env_lambda = None
   if isinstance(env, str):
     env_lambda = lambda: gym.make(env)

From cb0292566e903563d7aa990cd2166699873b7b96 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?B=C5=82a=C5=BCej=20O?= <blazej.osinski@codilime.com>
Date: Thu, 28 Jun 2018 11:39:02 +0200
Subject: [PATCH 0262/2720] Uniforming semantics of _reset_non_empty.

---
 tensor2tensor/rl/envs/simulated_batch_env.py | 2 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py   | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 28a41a6c5..abb713dae 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -198,7 +198,7 @@ def _reset_non_empty(self, indices):
     with tf.control_dependencies([self.history_buffer.reset(indices)]):
       with tf.control_dependencies([self._observ.assign(
           self.history_buffer.get_all_elements()[:, -1, ...])]):
-        return tf.identity(self._observ.read_value())
+        return tf.gather(self._observ.read_value(), indices)
 
   @property
   def observ(self):
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 570a26915..03a64d9ed 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -179,7 +179,7 @@ def not_done_step(a, _):
 
   def _reset_non_empty(self, indices):
     # pylint: disable=protected-access
-    new_values = tf.gather(self._batch_env._reset_non_empty(indices), indices)
+    new_values = self._batch_env._reset_non_empty(indices)
     # pylint: enable=protected-access
     inx = tf.concat(
         [
@@ -218,8 +218,8 @@ def _reset_non_empty(self, indices):
         self._time_elapsed, indices,
         tf.gather(tf.zeros((len(self),), tf.int32), indices))
     # pylint: disable=protected-access
-    new_values = tf.gather(self._batch_env._reset_non_empty(indices), indices)
+    new_values = self._batch_env._reset_non_empty(indices)
     # pylint: enable=protected-access
     assign_op = tf.scatter_update(self._observ, indices, new_values)
     with tf.control_dependencies([op_zero, assign_op]):
-      return tf.identity(self.observ)
+      return tf.gather(self.observ, indices)

From 74d2124dedd9260358d7f52252619520b5d39093 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?B=C5=82a=C5=BCej=20O?= <blazej.osinski@codilime.com>
Date: Thu, 28 Jun 2018 12:02:36 +0200
Subject: [PATCH 0263/2720] Re-enabling StackAndSkipWrapper.

---
 tensor2tensor/data_generators/gym_problems.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 766ac9ab6..194f302ae 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -52,7 +52,7 @@
 
 def standard_atari_env_spec(env):
   """Parameters of environment specification."""
-  standard_wrappers = [[tf_atari_wrappers.MaxAndSkipWrapper, {"skip": 4}]]
+  standard_wrappers = [[tf_atari_wrappers.StackAndSkipWrapper, {"skip": 4}]]
   env_lambda = None
   if isinstance(env, str):
     env_lambda = lambda: gym.make(env)

From 9486e574f2cd8d44e24279970850345490509045 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 28 Jun 2018 06:23:46 -0700
Subject: [PATCH 0264/2720] Add option to use weighted sum of transforms for
 MPNN.

PiperOrigin-RevId: 202468482
---
 .../common_message_passing_attention.py       | 136 +++++++++++++-----
 1 file changed, 98 insertions(+), 38 deletions(-)

diff --git a/tensor2tensor/layers/common_message_passing_attention.py b/tensor2tensor/layers/common_message_passing_attention.py
index 69c9287e7..a13210082 100644
--- a/tensor2tensor/layers/common_message_passing_attention.py
+++ b/tensor2tensor/layers/common_message_passing_attention.py
@@ -40,6 +40,7 @@ def multihead_graph_attention(query_antecedent,
                               dropout_broadcast_dims=None,
                               adjacency_matrix=None,
                               num_edge_types=5,
+                              ignore_zero=True,
                               vars_3d=False,
                               **kwargs):
   """Multihead scaled-dot-product attention with input/output transformations.
@@ -70,6 +71,7 @@ def multihead_graph_attention(query_antecedent,
     adjacency_matrix: an optional tensor of shape [batch, len_q, len_q]
       containing edge vectors for attention
     num_edge_types: number of edge types, an int
+    ignore_zero: A flag that says that edge type 0 should be ignored.
     vars_3d: use 3-dimensional variables for input/output transformations
     **kwargs (dict): Parameters for the attention function
 
@@ -125,7 +127,8 @@ def multihead_graph_attention(query_antecedent,
           make_image_summary=make_image_summary,
           dropout_broadcast_dims=dropout_broadcast_dims,
           adjacency_matrix=adjacency_matrix,
-          num_edge_types=num_edge_types)
+          num_edge_types=num_edge_types,
+          ignore_zero=ignore_zero)
 
     x = common_attention.combine_heads(x)
 
@@ -146,13 +149,18 @@ def multihead_graph_attention(query_antecedent,
 
 
 @expert_utils.add_name_scope()
-def make_edge_vectors(adjacency_matrix, num_edge_types, depth, name=None):
+def make_edge_vectors(adjacency_matrix,
+                      num_edge_types,
+                      depth,
+                      ignore_zero,
+                      name=None):
   """Gets edge vectors for the edge types in the adjacency matrix.
 
   Args:
     adjacency_matrix: A [batch, num_nodes, num_nodes] tensor of ints.
     num_edge_types: Number of different edge types
     depth: Number of channels
+    ignore_zero: Whether to use zeros vector for edge type 0.
     name: a string
   Returns:
     A [batch, num_nodes, num_nodes, depth] vector of tensors
@@ -169,16 +177,22 @@ def make_edge_vectors(adjacency_matrix, num_edge_types, depth, name=None):
     # Avoiding gathers so that it works on TPUs
     # adjacency_matrix_one_hot has shape
     # [batch, num_nodes, num_nodes, num_edge_types]
-
     adjacency_matrix_one_hot = tf.one_hot(adjacency_matrix, num_edge_types)
 
     att_adj_vectors = tf.matmul(
         tf.reshape(tf.to_float(adjacency_matrix_one_hot), [-1, num_edge_types]),
         adj_vectors)
-    return tf.reshape(att_adj_vectors, [
+    # Reshape to be [batch, num_nodes, num_nodes, depth].
+    att_adj_vectors = tf.reshape(att_adj_vectors, [
         adjacency_matrix_shape[0], adjacency_matrix_shape[1],
         adjacency_matrix_shape[2], depth
     ])
+    if ignore_zero:
+      # Set vectors for edge type 0 to be all zeros.
+      mask = tf.not_equal(tf.expand_dims(adjacency_matrix, axis=-1), 0)
+      return att_adj_vectors * tf.to_float(mask)
+    else:
+      return att_adj_vectors
 
 
 def graph_attention(q,
@@ -192,7 +206,8 @@ def graph_attention(q,
                     save_weights_to=None,
                     dropout_broadcast_dims=None,
                     adjacency_matrix=None,
-                    num_edge_types=5):
+                    num_edge_types=5,
+                    ignore_zero=True):
   """graph attention.
 
   Args:
@@ -214,6 +229,7 @@ def graph_attention(q,
     adjacency_matrix: optional matrix of [batch, length, length] ids indicating
       edge type
     num_edge_types: an int indicating number of edge types
+    ignore_zero: A flag that says that edge type 0 should be ignored.
   Returns:
     A Tensor of shape [batch, length, depth(q)]
   """
@@ -223,11 +239,12 @@ def graph_attention(q,
     logits = tf.matmul(q, k, transpose_b=True)
     if adjacency_matrix is not None:
       key_head_depth = common_layers.shape_list(q)[-1]
-      adjacency_vectors = make_edge_vectors(adjacency_matrix, num_edge_types,
-                                            key_head_depth, name)
-      # zeroing out the vectors that have 0 entries in the adjacency
-      adjacency_vectors *= tf.to_float(
-          tf.expand_dims(adjacency_matrix, axis=-1))
+      adjacency_vectors = make_edge_vectors(
+          adjacency_matrix,
+          num_edge_types,
+          key_head_depth,
+          ignore_zero=ignore_zero,
+          name=name)
       # transposing q to be [batch, length_q, heads, depth_k]
       # to allow for matmul with [batch, length_q, length_q, depth_k]
       q_t = tf.transpose(q, [0, 2, 1, 3])
@@ -249,7 +266,7 @@ def graph_attention(q,
 def compute_mpnn_qkv(node_states,
                      total_key_depth,
                      total_value_depth,
-                     num_edge_types,
+                     num_transforms,
                      ignore_zero=True):
   """Computes query, key and value for edge matrices.
 
@@ -258,7 +275,7 @@ def compute_mpnn_qkv(node_states,
   Let D be the size of the node hidden states.
   Let K be the size of the attention keys/queries (total_key_depth).
   Let V be the size of the attention values (total_value_depth).
-  Let T be the total number of edge types (num_edge_types).
+  Let T be the total number of transforms (num_transforms).
 
   Computes the queries, keys, and values for attention.
   * For each node N_i in the graph, a query Q_i of size K is computed. This
@@ -275,7 +292,8 @@ def compute_mpnn_qkv(node_states,
     node_states: A Tensor with shape [B, N, D].
     total_key_depth: an integer (K).
     total_value_depth: an integer (V).
-    num_edge_types: a integer specifying number of edge types (T).
+    num_transforms: a integer specifying number of transforms (T). This is
+      typically the number of edge types.
     ignore_zero: If true, then edge type 0 will not be considered. Equivalent
       to having a linear transformation of all 0's for edge type 0. All queries,
       keys, and values for edge type 0 will be all 0's.
@@ -300,14 +318,14 @@ def compute_mpnn_qkv(node_states,
   q_shape = common_layers.shape_list(q)  # As above, q_shape = [B, N, K].
 
   # T (or T-1 if ignore_zero).
-  nonignored_edge_types = num_edge_types - int(ignore_zero)
+  nonignored_transforms = num_transforms - int(ignore_zero)
 
   # Creates the attention keys in a manner similar to the process of creating
   # the attention queries. One key is created for each type of outgoing edge the
   # corresponding node might have, meaning k will have shape [B, N, K*T].
   k = common_layers.dense(
       node_states,
-      total_key_depth * nonignored_edge_types,
+      total_key_depth * nonignored_transforms,
       use_bias=False,
       name="k_mpnn")
 
@@ -315,7 +333,7 @@ def compute_mpnn_qkv(node_states,
   # a manner largely identical to that of the keys.
   v = common_layers.dense(
       node_states,
-      total_value_depth * nonignored_edge_types,
+      total_value_depth * nonignored_transforms,
       use_bias=False,
       name="v_mpnn")
 
@@ -328,9 +346,9 @@ def compute_mpnn_qkv(node_states,
   #
   # This reshape is only necessary when ignore_zero is True (for the padding
   # step that follows).
-  k = tf.reshape(k, [batch, length, nonignored_edge_types, total_key_depth])
+  k = tf.reshape(k, [batch, length, nonignored_transforms, total_key_depth])
   v = tf.reshape(
-      v, [q_shape[0], q_shape[1], nonignored_edge_types, total_value_depth])
+      v, [q_shape[0], q_shape[1], nonignored_transforms, total_value_depth])
 
   # If we previously ignored edge type 0, then we need to pad the keys and
   # values to take this additional edge type into account. To do so, we
@@ -341,9 +359,9 @@ def compute_mpnn_qkv(node_states,
     v = tf.pad(v, [[0, 0], [0, 0], [1, 0], [0, 0]])
 
   # Flatten out the fourth dimension.
-  k = tf.reshape(k, [q_shape[0], q_shape[1] * num_edge_types, total_key_depth])
+  k = tf.reshape(k, [q_shape[0], q_shape[1] * num_transforms, total_key_depth])
   v = tf.reshape(v,
-                 [q_shape[0], q_shape[1] * num_edge_types, total_value_depth])
+                 [q_shape[0], q_shape[1] * num_transforms, total_value_depth])
   return q, k, v
 
 
@@ -354,6 +372,8 @@ def multihead_mpnn_attention(node_states,
                              num_heads,
                              adjacency_matrix=None,
                              num_edge_types=5,
+                             num_transforms=None,
+                             use_weighted_sum=False,
                              ignore_zero=True,
                              name="mpnn_attention"):
   """Multihead scaled-dot-product attention with input/output transformations.
@@ -365,7 +385,7 @@ def multihead_mpnn_attention(node_states,
   Let V be the size of the attention values (total_value_depth).
   Let O be the size of the attention output (output_depth).
   Let H be the number of heads (num_heads).
-  Let T be the total number of edge types (num_edge_types).
+  Let T be the total number of transforms (num_transforms).
 
   The key and value depths are split across all of the heads. For example, if
   the key depth is 6 and there are three heads, then the key for each head has
@@ -380,7 +400,11 @@ def multihead_mpnn_attention(node_states,
     adjacency_matrix: An Tensor of ints with shape [B, N, N]. If there is an
       edge from node j to node i in batch b, then adjacency_matrix[b, i, j]
       contains the type of that edge as an integer. Otherwise, it contains 0.
-    num_edge_types: An integer indicating number of edge types (T).
+    num_edge_types: An integer indicating number of edge types.
+    num_transforms: An integer indicating number of transforms (T). If None,
+      then num_transforms will be equal to num_edge_types.
+    use_weighted_sum: If False, will only use a single transform per edge type.
+      Otherwise, use a learned weighted sum of transforms per edge type.
     ignore_zero: A flag that says that edge type 0 should be ignored.
     name: A string.
 
@@ -399,13 +423,17 @@ def multihead_mpnn_attention(node_states,
                      "attention heads (%d)." % (total_value_depth, num_heads))
   with tf.variable_scope(
       name, default_name="multihead_mpnn_attention", values=[node_states]):
+    # If not explicitly set, use num_transforms set to num_edge_types.
+    num_transforms = (
+        num_edge_types if num_transforms is None else num_transforms)
+
     # Create the query for each node's incoming edges.
     # Create the keys/values for each node for each possible outgoing edge type.
     q, k, v = compute_mpnn_qkv(
         node_states,
         total_key_depth,
         total_value_depth,
-        num_edge_types,
+        num_transforms,
         ignore_zero=ignore_zero)
 
     q_shape = tf.shape(q)  # As above, q_shape is [B, N, K].
@@ -431,13 +459,13 @@ def multihead_mpnn_attention(node_states,
 
     # Split the keys and values into separate per-edge-type keys and values.
     k = tf.reshape(k, [
-        num_heads, q_shape[0], q_shape[1], num_edge_types,
+        num_heads, q_shape[0], q_shape[1], num_transforms,
         total_key_depth // num_heads
     ])  # Shape [H, B, N, T, K/H].
     k = tf.transpose(k, [0, 1, 3, 2, 4])  # Shape [H, B, T, N, K/H].
 
     v = tf.reshape(v, [
-        num_heads, q_shape[0], q_shape[1], num_edge_types,
+        num_heads, q_shape[0], q_shape[1], num_transforms,
         total_value_depth // num_heads
     ])  # Shape [H, B, N, T, V/H].
     v = tf.transpose(v, [0, 1, 3, 2, 4])  # Shape [H, B, T, N, V/H].
@@ -450,8 +478,14 @@ def multihead_mpnn_attention(node_states,
     # and adding those weighted values together.
     head_outputs = []
     for head_id in range(num_heads):
-      output = dot_product_mpnn_attention(q[head_id], k[head_id], v[head_id],
-                                          adjacency_matrix, num_edge_types)
+      output = dot_product_mpnn_attention(
+          q[head_id],
+          k[head_id],
+          v[head_id],
+          adjacency_matrix,
+          num_edge_types,
+          num_transforms=num_transforms,
+          use_weighted_sum=use_weighted_sum)
 
       # Store this result in the list of attention results for each head.
       # The call to expand_dims gives output shape [1, B, N, V/H], which will
@@ -477,6 +511,8 @@ def dot_product_mpnn_attention(q,
                                v,
                                adjacency_matrix,
                                num_edge_types,
+                               num_transforms=None,
+                               use_weighted_sum=False,
                                ignore_zero=True,
                                name=None):
   """Dot product attention with edge vectors.
@@ -485,7 +521,7 @@ def dot_product_mpnn_attention(q,
   Let N be the number of nodes in the graph.
   Let K be the size of the attention keys/queries.
   Let V be the size of the attention values.
-  Let T be the total number of edge types (num_edge_types).
+  Let T be the total number of transforms (num_transforms).
 
   Args:
     q: The query Tensor of shape [B, N, K].
@@ -493,7 +529,11 @@ def dot_product_mpnn_attention(q,
     v: The value Tensor of shape [B, T, N, V].
     adjacency_matrix: A Tensor of shape [B, N, N]. An entry at indices b, i, j
      is the integer edge type of the edge from node j to node i in batch b.
-    num_edge_types: An integer specifying number of edge types (T).
+    num_edge_types: An integer specifying number of edge types.
+    num_transforms: An integer indicating number of transforms (T). If None,
+      then num_transforms will be equal to num_edge_types.
+    use_weighted_sum: If False, will only use a single transform per edge type.
+      Otherwise, use a learned weighted sum of transforms per edge type.
     ignore_zero: A flag that says that edge type 0 should be ignored.
     name: A string.
 
@@ -501,6 +541,10 @@ def dot_product_mpnn_attention(q,
     A Tensor of shape [B, N, V] storing the result of computing attention
     weights using the queries and keys and combining the values according to
     those weights.
+
+  Raises:
+    ValueError: if num_transforms doesn't equal num_edge_types and not using
+      weighted sum.
   """
   # TODO(jfrankle): Consider ways to handle graphs that have multiple edges
   # between the same nodes (with only one edge of each type. adjacency_matrix
@@ -509,6 +553,14 @@ def dot_product_mpnn_attention(q,
       name,
       default_name="dot_product_mpnn_attention",
       values=[q, k, v, adjacency_matrix, num_edge_types]):
+    # If not explicitly set, use num_transforms set to num_edge_types.
+    num_transforms = (
+        num_edge_types if num_transforms is None else num_transforms)
+
+    if not use_weighted_sum and num_transforms != num_edge_types:
+      raise ValueError("num_transforms must equal num_edge_types unless "
+                       "use_weighted_sum is True")
+
     # Computes the raw dot-product attention values between each query and
     # the corresponding keys it needs to consider.
     #
@@ -528,22 +580,30 @@ def dot_product_mpnn_attention(q,
     # each pair of nodes. If such an edge exists, it contains the integer
     # type of that edge at position (i, j) of the adjacency matrix.
     #
-    # adjacency_matrix_one_hot has shape [B, N, N, T]. If there is an edge
-    # from node j to node i of type t, then index t of the last dimension is
-    # 1 for entry (i, j) of the second and third dimensions.
-    adjacency_matrix_one_hot = tf.one_hot(adjacency_matrix, num_edge_types)
+    # Construct edge_vectors of shape [B, N, N, T].
+    if use_weighted_sum:
+      # Use dense representation for edge vectors.
+      edge_vectors = make_edge_vectors(
+          adjacency_matrix,
+          num_edge_types,
+          num_transforms,
+          ignore_zero=ignore_zero)
+    else:
+      # Generate one-hot vectors based on edge types.
+      # If there is an edge from node j to node i of type t, then index t of the
+      # last dimension is 1 for entry (i, j) of the second and third dimensions.
+      edge_vectors = tf.one_hot(adjacency_matrix, num_transforms)
 
     # Rearranging the dimensions to match the shape of all_edge_logits.
-    adjacency_matrix_one_hot = tf.transpose(adjacency_matrix_one_hot,
-                                            [0, 3, 1, 2])
+    edge_vectors = tf.transpose(edge_vectors, [0, 3, 1, 2])
 
-    # Element-wise multiplies all_edge_logits and adjacency_matrix_one_hot.
+    # Element-wise multiplies all_edge_logits and edge_vectors.
     #
     # In other words: all_edge_logits contains N x N matrices of query-key
     # products. This element-wise multiplication zeroes out entries that do not
     # correspond to actual edges in the graph of the appropriate edge type.
     # all_edge_logits retains shape [B, T, N, N].
-    all_edge_logits *= adjacency_matrix_one_hot
+    all_edge_logits *= edge_vectors
 
     # Since there can only be one edge from node A to node B, we can collapse
     # the T different adjacency matrices containing key-query pairs into one
@@ -579,7 +639,7 @@ def dot_product_mpnn_attention(q,
 
     # Zeroes out the entries in edge_compatibility that do not correspond to
     # actual edges.
-    edge_compatibility *= adjacency_matrix_one_hot  # Shape [B, T, N, N].
+    edge_compatibility *= edge_vectors  # Shape [B, T, N, N].
 
     # Computes the incoming value vectors for each node by weighting them
     # according to the attention weights. These values are still segregated by

From 1ef2c24834c4e5c9d9550c1dc47c6c1cdd066559 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 28 Jun 2018 11:03:33 -0700
Subject: [PATCH 0265/2720] Clean up and test dml_loss. Change arg order in
 DMOL.

PiperOrigin-RevId: 202508603
---
 tensor2tensor/layers/common_layers.py      | 64 ++++++++++++----------
 tensor2tensor/layers/common_layers_test.py | 33 ++++++++++-
 2 files changed, 67 insertions(+), 30 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 45857d5e9..fe493b2d5 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1938,30 +1938,44 @@ def padded_cross_entropy(logits,
     return tf.reduce_sum(xent * weights), tf.reduce_sum(weights)
 
 
-def dml_loss(
-    pred,
-    labels,
-    weights_fn=weights_nonzero,  # Unused
-    reduce_sum=True):
+def _weights_one_third(labels):
+  """Returns Tensor of shape [batch, height, width]. Each element is 1/3."""
+  return tf.ones(tf.shape(labels)[:-1]) / 3.
+
+
+def dml_loss(pred,
+             labels,
+             weights_fn=_weights_one_third,
+             reduce_sum=True):
   """Discretized mixture of logistics loss.
 
   Args:
-    pred: a tensor of shape [batch, height, width, 10*num_mixtures]
-    labels: a [batch, height, width, channels] tensor of real pixel intensities
-    weights_fn: weights function
-    reduce_sum: A boolean, to return scalar mean loss instead of per position
+    pred: A [batch, height, width, num_mixtures*10] tensor of floats
+      comprising one unconstrained mixture probability, three means
+      (one per channel), three standard deviations (one per channel),
+      and three coefficients which linearly parameterize dependence across
+      channels.
+    labels: A [batch, height, width, channels] tensor of 8-bit pixel
+      intensities. The computation assumes channels is 3.
+    weights_fn: A function of labels, returning a Tensor of shape
+      [batch, height, width] which weights each loss term. Default is to scale
+      each loss term by 1/3 so that they capture the average across channels.
+    reduce_sum: A boolean, to return scalar loss instead of per position.
 
   Returns:
-    a pair of tensors of loss/sum of losses, denominator
+    Tuple of loss tensors for numerator and denominator, each a scalar if
+    reduce_sum else of shape [batch, height, width]. The sum of their divisions
+    is the number of nats for each pixel in labels.
   """
-  del weights_fn  # Unused
   real_labels = convert_rgb_to_symmetric_real(labels)
-  dml_loss_value = discretized_mix_logistic_loss(
-      real_labels, pred, sum_all=reduce_sum)
+  dml_loss_value = discretized_mix_logistic_loss(pred=pred, labels=real_labels)
+  weights = weights_fn(labels)
+  loss_num = weights * dml_loss_value
+  loss_den = weights_nonzero(weights)
   if reduce_sum:
-    return dml_loss_value, tf.reduce_sum(tf.ones(tf.shape(labels), tf.float32))
-  else:
-    return dml_loss_value / 3., tf.ones(tf.shape(dml_loss_value), tf.float32)
+    loss_num = tf.reduce_sum(loss_num)
+    loss_den = tf.reduce_sum(loss_den)
+  return loss_num, loss_den
 
 
 def split_to_discretized_mix_logistic_params(inputs):
@@ -1997,7 +2011,7 @@ def split_to_discretized_mix_logistic_params(inputs):
   return logits, locs, log_scales, coeffs
 
 
-def discretized_mix_logistic_loss(labels, pred, sum_all=True):
+def discretized_mix_logistic_loss(pred, labels):
   """Computes negative log probability for the discretized mixture of logistics.
 
   The distribution of a whole pixel is a mixture of 3-dimensional discretized
@@ -2018,19 +2032,17 @@ def discretized_mix_logistic_loss(labels, pred, sum_all=True):
   8-bit inputs, this implementation assumes the events are rescaled to [-1, 1].
 
   Args:
-    labels: A [batch, height, width, channels] tensor of true pixel intensities
-      rescaled to [-1, 1]. The computation assumes channels is 3.
     pred: A [batch, height, width, num_mixtures*10] tensor of floats
       comprising one unconstrained mixture probability, three means
       (one per channel), three standard deviations (one per channel),
       and three coefficients which linearly parameterize dependence across
       channels.
-    sum_all: A boolean to return scalar mean loss or per position.
+    labels: A [batch, height, width, channels] tensor of true pixel intensities
+      rescaled to [-1, 1]. The computation assumes channels is 3.
 
   Returns:
     A [batch, height, width] tensor of the negative log conditional probability
-    of each pixel given all previous pixels if not sum_all else add all the
-    losses (for eval).
+    of each pixel given all previous pixels.
   """
 
   logits, locs, log_scales, coeffs = split_to_discretized_mix_logistic_params(
@@ -2079,12 +2091,8 @@ def discretized_mix_logistic_loss(labels, pred, sum_all=True):
 
   # Sum over channels and compute log-probability of each mixture.
   log_probs = tf.reduce_sum(log_probs, -1) + tf.nn.log_softmax(logits, axis=-1)
-  if sum_all:
-    output = -tf.reduce_sum(tf.reduce_logsumexp(log_probs, axis=-1))
-    return output
-  else:
-    output = -tf.reduce_logsumexp(log_probs, axis=-1)
-    return output
+  output = -tf.reduce_logsumexp(log_probs, axis=-1)
+  return output
 
 
 def sample_from_discretized_mix_logistic(pred, seed=None):
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index 03096010d..bc86cfdff 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -17,13 +17,16 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+from absl.testing import parameterized
 import numpy as np
+
 from tensor2tensor.layers import common_layers
 
 import tensorflow as tf
 
 
-class CommonLayersTest(tf.test.TestCase):
+class CommonLayersTest(parameterized.TestCase, tf.test.TestCase):
 
   def testIndexLastDimWithIndices(self):
     x = np.array([[2., 3., 4., 5.],
@@ -484,6 +487,32 @@ def testPaddingCrossEntropyFactoredGrad(self):
     self.assertAllClose(actual_df, actual_df_factored)
     self.assertAllClose(actual_dw, actual_dw_factored)
 
+  @parameterized.parameters(
+      (2, 4, 4, 5, True),
+      (2, 4, 4, 5, False),
+      (1, 16, 16, 1, True),
+      (1, 16, 16, 1, False),
+  )
+  def testDmlLoss(self, batch, height, width, num_mixtures, reduce_sum):
+    channels = 3
+    pred = tf.random_normal([batch, height, width, num_mixtures * 10])
+    labels = tf.random_uniform([batch, height, width, channels],
+                               minval=0, maxval=256, dtype=tf.int32)
+    actual_loss_num, actual_loss_den = common_layers.dml_loss(
+        pred=pred, labels=labels, reduce_sum=reduce_sum)
+    actual_loss = actual_loss_num / actual_loss_den
+
+    real_labels = common_layers.convert_rgb_to_symmetric_real(labels)
+    expected_loss = common_layers.discretized_mix_logistic_loss(
+        pred=pred, labels=real_labels) / channels
+    if reduce_sum:
+      expected_loss = tf.reduce_mean(expected_loss)
+
+    with self.test_session() as sess:
+      actual_loss_val, expected_loss_val = sess.run(
+          [actual_loss, expected_loss])
+    self.assertAllClose(actual_loss_val, expected_loss_val)
+
   def testDiscretizedMixLogisticLoss(self):
     batch = 2
     height = 4
@@ -515,7 +544,7 @@ def testDiscretizedMixLogisticLoss(self):
     expected_loss = -tf.reduce_sum(tf.log(cdf_plus - cdf_min), axis=-1)
 
     actual_loss = common_layers.discretized_mix_logistic_loss(
-        labels, pred, sum_all=False)
+        pred=pred, labels=labels)
     with self.test_session() as session:
       actual_loss_val, expected_loss_val = session.run(
           [actual_loss, expected_loss])

From a70d0f7c2a475975b5bc625d54f7b919589b83eb Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 28 Jun 2018 11:36:43 -0700
Subject: [PATCH 0266/2720] more helpful logging

PiperOrigin-RevId: 202515771
---
 tensor2tensor/data_generators/generator_utils.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 64e09592f..f8dd1f12a 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -145,7 +145,8 @@ def generate_files(generator, output_filenames,
       switching to the next shard; by default set to 1, switch every case.
   """
   if outputs_exist(output_filenames):
-    tf.logging.info("Skipping generator because outputs files exist")
+    tf.logging.info("Skipping generator because outputs files exists at {}"
+                    .format(output_filenames))
     return
   tmp_filenames = [fname + ".incomplete" for fname in output_filenames]
   num_shards = len(output_filenames)

From d4277056dab2e40cd02ec9fa24779427b7454f17 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 28 Jun 2018 14:49:38 -0700
Subject: [PATCH 0267/2720] + Fixing the action/reward scope bug

PiperOrigin-RevId: 202547981
---
 tensor2tensor/models/research/next_frame.py | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index eab2becfe..c81264711 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -223,19 +223,20 @@ def construct_latent_tower(self, images):
 
   def reward_prediction(self, inputs):
     """Builds a reward prediction network."""
-    conv_size = self.tinyify([32, 16])
+    conv_size = self.tinyify([32, 16, 1])
     with tf.variable_scope("reward_pred", reuse=tf.AUTO_REUSE):
       x = inputs
+      x = slim.batch_norm(x, scope="reward_bn0")
       x = slim.conv2d(x, conv_size[0], [3, 3], scope="reward_conv1")
       x = slim.batch_norm(x, scope="reward_bn1")
       x = slim.conv2d(x, conv_size[1], [3, 3], scope="reward_conv2")
       x = slim.batch_norm(x, scope="reward_bn2")
-      x = slim.conv2d(x, 1, [3, 3], scope="reward_conv3", activation_fn=None)
+      x = slim.conv2d(x, conv_size[2], [3, 3], scope="reward_conv3")
     return x
 
-  def encode_to_shape(self, inputs, shape):
+  def encode_to_shape(self, inputs, shape, scope):
     """Encode the given tensor to given image shape."""
-    with tf.variable_scope("reward_enc", reuse=tf.AUTO_REUSE):
+    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
       w, h = shape[1].value, shape[2].value
       x = inputs
       x = tf.contrib.layers.flatten(x)
@@ -243,9 +244,9 @@ def encode_to_shape(self, inputs, shape):
       x = tf.reshape(x, (-1, w, h, 1))
       return x
 
-  def decode_to_shape(self, inputs, shape):
+  def decode_to_shape(self, inputs, shape, scope):
     """Encode the given tensor to given image shape."""
-    with tf.variable_scope("reward_dec", reuse=tf.AUTO_REUSE):
+    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
       x = inputs
       x = tf.contrib.layers.flatten(x)
       x = slim.fully_connected(x, shape[2].value, scope="decoding_full")
@@ -309,8 +310,9 @@ def construct_predictive_tower(
           hidden4, hidden4.get_shape()[3], [3, 3], stride=2, scope="conv3")
 
       # Pass in reward and action.
-      emb_action = self.encode_to_shape(action, enc2.get_shape())
-      emb_reward = self.encode_to_shape(input_reward, enc2.get_shape())
+      emb_action = self.encode_to_shape(action, enc2.get_shape(), "action_enc")
+      emb_reward = self.encode_to_shape(
+          input_reward, enc2.get_shape(), "reward_enc")
       enc2 = tf.concat(axis=3, values=[enc2, emb_action, emb_reward])
 
       if latent is not None:
@@ -403,7 +405,8 @@ def construct_predictive_tower(
         output += layer * mask
 
       p_reward = self.reward_prediction(hidden5)
-      p_reward = self.decode_to_shape(p_reward, input_reward.shape)
+      p_reward = self.decode_to_shape(
+          p_reward, input_reward.shape, "reward_dec")
 
       return output, p_reward, lstm_state
 

From ea05fd5935af7b17cd16cf62bd5407e52d1fedba Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Thu, 28 Jun 2018 16:00:13 -0700
Subject: [PATCH 0268/2720] Add relative attention to masked local 1d attention
 and added attention dropout in masked 1d.

PiperOrigin-RevId: 202560871
---
 tensor2tensor/layers/common_attention.py      | 248 +++++++++++++++++-
 .../layers/common_image_attention.py          |  10 +-
 tensor2tensor/layers/common_layers.py         |   2 +-
 tensor2tensor/models/image_transformer.py     |  27 +-
 4 files changed, 271 insertions(+), 16 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index ee62b0f88..7752a95cf 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1814,11 +1814,42 @@ def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
     return output
 
 
+def _local_unmasked_relative_to_absolute(x):
+  """Converts tensor from relative to aboslute indexing for local attention.
+
+  Args:
+    x: a Tensor of shape [batch (or batch*num_blocks), heads,
+                          length, 2 * length - 1]
+
+  Returns:
+    A Tensor of shape [batch (or batch*num_blocks), heads, length, length-1]
+  """
+  x_shape = common_layers.shape_list(x)
+  batch = x_shape[0]
+  heads = x_shape[1]
+  length = x_shape[2]
+  # Concat columns of pad to shift from relative to absolute indexing.
+  col_pad = tf.zeros((batch, heads, length, 1))
+  x = tf.concat([x, col_pad], axis=3)
+
+  # Concat extra elements so to add up to shape (len+1, 2*len-2).
+  flat_x = tf.reshape(x, [batch, heads, length * 2 * length])
+  flat_pad = tf.zeros((batch, heads, length-1))
+  flat_x_padded = tf.concat([flat_x, flat_pad], axis=2)
+
+  # Reshape and slice out the padded elements.
+  final_x = tf.reshape(flat_x_padded, [batch, heads, length+1, 2*length-1])
+  final_x = final_x[:, :, :, length-1:]
+  final_x = final_x[:, :, :length, :]
+  return final_x
+
+
 def masked_local_attention_1d(q,
                               k,
                               v,
                               block_length=128,
                               make_image_summary=False,
+                              dropout_rate=0.,
                               name=None):
   """Attention to the source position and a neighborhood to the left of it.
 
@@ -1836,6 +1867,7 @@ def masked_local_attention_1d(q,
     v: a Tensor with shape [batch, heads, length, depth_v]
     block_length: an integer
     make_image_summary: a boolean, whether to make an attention image summary.
+    dropout_rate: Dropout rate for attention dropout
     name: an optional string
 
   Returns:
@@ -1843,6 +1875,7 @@ def masked_local_attention_1d(q,
   """
   with tf.variable_scope(
       name, default_name="local_attention_1d", values=[q, k, v]):
+
     batch = common_layers.shape_list(q)[0]
     heads = common_layers.shape_list(q)[1]
     length = common_layers.shape_list(q)[2]
@@ -1850,7 +1883,6 @@ def masked_local_attention_1d(q,
       const = tf.contrib.util.constant_value(block_length)
       if const is not None:
         block_length = int(const)
-
     # If (length < 2 * block_length), then we use only one block.
     if isinstance(length, int) and isinstance(block_length, int):
       block_length = length if length < block_length * 2 else block_length
@@ -1876,11 +1908,13 @@ def masked_local_attention_1d(q,
     first_q = tf.slice(q, [0, 0, 0, 0], [-1, -1, block_length, -1])
     first_k = tf.slice(k, [0, 0, 0, 0], [-1, -1, block_length, -1])
     first_v = tf.slice(v, [0, 0, 0, 0], [-1, -1, block_length, -1])
+
     first_output = dot_product_attention(
         first_q,
         first_k,
         first_v,
         attention_bias_lower_triangle(block_length),
+        dropout_rate=dropout_rate,
         make_image_summary=make_image_summary,
         name="fist_block")
 
@@ -1889,17 +1923,10 @@ def masked_local_attention_1d(q,
     k = tf.reshape(k, [batch, heads, num_blocks, block_length, depth_k])
     v = tf.reshape(v, [batch, heads, num_blocks, block_length, depth_v])
 
-    def local(x, depth):
-      """Create a local version of the keys or values."""
-      prev_block = tf.slice(x, [0, 0, 0, 0, 0],
-                            [-1, -1, num_blocks - 1, -1, -1])
-      cur_block = tf.slice(x, [0, 0, 1, 0, 0], [-1, -1, -1, -1, -1])
-      local_block = tf.concat([prev_block, cur_block], 3)
-      return tf.reshape(local_block,
-                        [batch, heads, num_blocks - 1, block_length * 2, depth])
-
-    local_k = local(k, depth_k)
-    local_v = local(v, depth_v)
+    local_k = _make_local_block(k, depth_k, batch, heads, num_blocks,
+                                block_length)
+    local_v = _make_local_block(v, depth_v, batch, heads, num_blocks,
+                                block_length)
     tail_q = tf.slice(q, [0, 0, 1, 0, 0], [-1, -1, -1, -1, -1])
     tail_q = tf.reshape(tail_q,
                         [batch, heads, num_blocks - 1, block_length, depth_k])
@@ -1909,12 +1936,16 @@ def local(x, depth):
     attention = tf.matmul(tail_q, local_k, transpose_b=True)
 
     # make sure source_pos <= target_pos
-    good_part = common_layers.ones_matrix_band_part(block_length, local_length,
+    good_part = common_layers.ones_matrix_band_part(block_length,
+                                                    local_length,
                                                     -1, block_length)
     mask = (1.0 - good_part) * -1e9
     mask = common_layers.cast_like(mask, attention)
     attention += tf.reshape(mask, [1, 1, 1, block_length, local_length])
     attention = tf.nn.softmax(attention)
+    attention = common_layers.dropout_with_broadcast_dims(
+        attention, 1.0 - dropout_rate,
+        broadcast_dims=None)
     # TODO(noam): figure out how to show a summary for the remaining blocks.
     # The naive way currently causes errors due to empty tensors.
     # output: [batch, heads, num_blocks-1, block_length, depth_v]
@@ -1927,6 +1958,190 @@ def local(x, depth):
     return output
 
 
+def _make_local_block(x, depth, batch, heads, num_blocks, block_length):
+  """Helper function to create a local version of the keys or values for 1d."""
+  prev_block = tf.slice(x, [0, 0, 0, 0, 0],
+                        [-1, -1, num_blocks - 1, -1, -1])
+  cur_block = tf.slice(x, [0, 0, 1, 0, 0], [-1, -1, -1, -1, -1])
+  local_block = tf.concat([prev_block, cur_block], 3)
+  return tf.reshape(local_block,
+                    [batch, heads, num_blocks - 1, block_length * 2, depth])
+
+
+def masked_rel_local_attention_1d(q,
+                                  k,
+                                  v,
+                                  block_length=128,
+                                  make_image_summary=False,
+                                  dropout_rate=0.,
+                                  share_rel_embed=False,
+                                  name=None):
+  """Masked local 1d attention with relative positions.
+
+  The sequence is divided into blocks of length block_size.
+  Attention for a given query position can only see memory positions
+  less than or equal to the query position, in the corresponding block
+  and the previous block.
+
+  If mask_right is True, then a target position cannot see greater source
+  positions.
+
+  Args:
+    q: a Tensor with shape [batch, heads, length, depth_k]
+    k: a Tensor with shape [batch, heads, length, depth_k]
+    v: a Tensor with shape [batch, heads, length, depth_v]
+    block_length: an integer
+    make_image_summary: a boolean, whether to make an attention image summary.
+    dropout_rate: Dropout rate for attention dropout
+    share_rel_embed: Boolean for sharing relative embeddings
+    name: an optional string
+
+  Returns:
+    a Tensor of shape [batch, heads, length, depth_v]
+  """
+  with tf.variable_scope(
+      name, default_name="local_attention_1d", values=[q, k, v]):
+
+    default_block_length = block_length
+    batch = common_layers.shape_list(q)[0]
+    heads = common_layers.shape_list(q)[1]
+    length = common_layers.shape_list(q)[2]
+    # If (length < 2 * block_length), then we use only one block.
+    block_length = length if length < block_length * 2 else block_length
+    depth_k = common_layers.shape_list(k)[3]
+    depth_v = common_layers.shape_list(v)[3]
+    original_length = length
+    padding_size = tf.mod(-length, block_length)
+    length += padding_size
+    padding = [[0, 0], [0, 0], [0, padding_size], [0, 0]]
+    q = tf.pad(q, padding)
+    k = tf.pad(k, padding)
+    v = tf.pad(v, padding)
+
+    num_blocks = length // block_length
+    # compute attention for the first query block.
+    first_q = tf.slice(q, [0, 0, 0, 0], [-1, -1, block_length, -1])
+    first_k = tf.slice(k, [0, 0, 0, 0], [-1, -1, block_length, -1])
+    first_v = tf.slice(v, [0, 0, 0, 0], [-1, -1, block_length, -1])
+    # Relative embeddings will be used later as well.
+    # TODO(avaswani,annahuang): check why 2*bl was breaking for music
+    # We only multiply with the needed embeddings as we slice them out.
+    rel_embed_length = 4 * default_block_length
+    # Relative embeddings can be shared or unshared
+    first_logits = tf.matmul(first_q, first_k, transpose_b=True)
+    if share_rel_embed:
+      relative_embeddings = (
+          tf.get_variable(name="relative_embeddings",
+                          shape=(rel_embed_length, depth_k)))
+      masked_relative_embeddings = tf.slice(
+          relative_embeddings,
+          [rel_embed_length - block_length, 0], [-1, -1])
+      first_relative_logits = tf.einsum(
+          "bhld,md->bhlm", first_q, masked_relative_embeddings)
+    else:
+      relative_embeddings = (
+          tf.get_variable(name="relative_embeddings",
+                          shape=(heads, rel_embed_length, depth_k)))
+      masked_relative_embeddings = tf.slice(
+          relative_embeddings,
+          [0, rel_embed_length - block_length, 0], [-1, -1, -1])
+      first_relative_logits = tf.einsum(
+          "bhld,hmd->bhlm", first_q, masked_relative_embeddings)
+    first_logits += (
+        _relative_position_to_absolute_position_masked(first_relative_logits))
+    # adding a mask
+    first_logits += (
+        common_layers.cast_like(attention_bias_lower_triangle(block_length),
+                                first_logits))
+    first_att = tf.nn.softmax(first_logits,
+                              name="first_attention_weights")
+    # dropping out the attention links for each of the heads
+    first_att = common_layers.dropout_with_broadcast_dims(
+        first_att, 1.0 - dropout_rate,
+        broadcast_dims=None)
+    # only call image summary for the first block
+    if common_layers.should_generate_summaries() and make_image_summary:
+      attention_image_summary(first_att, None)
+    first_output = tf.matmul(first_att, first_v)
+
+    # compute attention for all subsequent query blocks.
+    q = tf.reshape(q, [batch, heads, num_blocks, block_length, depth_k])
+    k = tf.reshape(k, [batch, heads, num_blocks, block_length, depth_k])
+    v = tf.reshape(v, [batch, heads, num_blocks, block_length, depth_v])
+    local_k = _make_local_block(k, depth_k, batch, heads, num_blocks,
+                                block_length)
+    local_v = _make_local_block(v, depth_v, batch, heads, num_blocks,
+                                block_length)
+    tail_q = tf.slice(q, [0, 0, 1, 0, 0], [-1, -1, -1, -1, -1])
+    tail_q = tf.reshape(tail_q,
+                        [batch, heads, num_blocks - 1, block_length, depth_k])
+    local_length = common_layers.shape_list(local_k)[3]
+
+    # collapsing num blocks and batch size so that we can reuse
+    # functions
+    def _reshape_for_relative(x):
+      x_shape = common_layers.shape_list(x)
+      # [batch, num_blocks, heads, length, depth]
+      x = tf.transpose(x, [0, 2, 1, 3, 4])
+      x = tf.reshape(x, [batch*x_shape[2], heads, x_shape[3],
+                         x_shape[4]])
+      return x
+    rel_tail_q = _reshape_for_relative(tail_q)
+    rel_k = _reshape_for_relative(local_k)
+    rel_v = _reshape_for_relative(local_v)
+    # Computing relative logits separately for the masked and unmasked parts
+    # because the reshaping logic is different for both
+    if share_rel_embed:
+      used_relative_embeddings = tf.slice(
+          relative_embeddings,
+          [rel_embed_length - 2*block_length, 0], [-1, -1])
+      rel_logits = tf.einsum(
+          "bhld,md->bhlm", rel_tail_q, used_relative_embeddings)
+      masked_rel_logits = tf.slice(rel_logits, [0, 0, 0, block_length],
+                                   [-1, -1, -1, -1])
+    else:
+      used_relative_embeddings = tf.slice(
+          relative_embeddings,
+          [0, rel_embed_length - 2*block_length, 0], [-1, -1, -1])
+      rel_logits = tf.einsum(
+          "bhld,hmd->bhlm", rel_tail_q, used_relative_embeddings)
+      masked_rel_logits = tf.slice(rel_logits, [0, 0, 0, block_length],
+                                   [-1, -1, -1, -1])
+    masked_rel_logits = _relative_position_to_absolute_position_masked(
+        masked_rel_logits)
+    unmasked_rel_logits = tf.slice(rel_logits, [0, 0, 0, 0],
+                                   [-1, -1, -1, 2*block_length-1])
+    unmasked_rel_logits = _local_unmasked_relative_to_absolute(
+        unmasked_rel_logits)
+    all_rel_logits = tf.concat([unmasked_rel_logits, masked_rel_logits],
+                               axis=3)
+    all_logits = (
+        tf.matmul(rel_tail_q, rel_k, transpose_b=True) + all_rel_logits)
+    # make sure source_pos <= target_pos
+    good_part = common_layers.ones_matrix_band_part(block_length,
+                                                    local_length,
+                                                    -1, block_length)
+    mask = (1.0 - good_part) * -1e9
+    mask = common_layers.cast_like(mask, all_logits)
+    all_logits += tf.reshape(mask, [1, 1, block_length, local_length])
+    attention = tf.nn.softmax(all_logits)
+    attention = common_layers.dropout_with_broadcast_dims(
+        attention, 1.0 - dropout_rate,
+        broadcast_dims=None)
+    output = tf.matmul(attention, rel_v)
+    # bring to [batch, heads, num_blocks-1, block_length, depth]
+    output = tf.reshape(output,
+                        [batch, num_blocks-1, heads, block_length, depth_v])
+    output = tf.transpose(output, [0, 2, 1, 3, 4])
+
+    output = tf.reshape(
+        output, [batch, heads, (num_blocks - 1) * block_length, depth_v])
+    output = tf.concat([first_output, output], axis=2)
+    output = tf.slice(output, [0, 0, 0, 0], [-1, -1, original_length, -1])
+    output = tf.reshape(output, [batch, heads, original_length, depth_v])
+    return output
+
+
 def local_attention_1d(q, k, v, block_length=128, filter_width=100, name=None):
   """strided block local self-attention.
 
@@ -2784,6 +2999,7 @@ def multihead_attention(query_antecedent,
                         output_depth,
                         num_heads,
                         dropout_rate,
+                        shared_rel=False,
                         max_relative_position=None,
                         image_shapes=None,
                         attention_type="dot_product",
@@ -2814,6 +3030,7 @@ def multihead_attention(query_antecedent,
     output_depth: an integer
     num_heads: an integer dividing total_key_depth and total_value_depth
     dropout_rate: a floating point number
+    shared_rel: boolean to share relative embeddings
     max_relative_position: Maximum distance between inputs to generate
                            unique relation embeddings for. Only relevant
                            when using "dot_product_relative" attention.
@@ -2972,6 +3189,11 @@ def multihead_attention(query_antecedent,
     elif attention_type == "local_within_block_mask_right":
       x = masked_within_block_local_attention_1d(
           q, k, v, block_length=block_length)
+    elif attention_type == "rel_local_mask_right":
+      x = masked_rel_local_attention_1d(q, k, v, block_length=block_length,
+                                        make_image_summary=make_image_summary,
+                                        dropout_rate=dropout_rate,
+                                        share_rel_embed=shared_rel)
     elif attention_type == "local_mask_right":
       x = masked_local_attention_1d(
           q,
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index a30c23c60..4a54f0db9 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -33,6 +33,7 @@ class AttentionType(object):
   MOE_LOCAL_1D = "moe_local1d"
   LOCAL_BLOCK = "local_block"
   NON_CAUSAL_1D = "local_1d_noncausal"
+  RELATIVE_LOCAL_1D = "rel_local_1d"
 
   @staticmethod
   def get_choices():
@@ -45,6 +46,7 @@ def get_choices():
         AttentionType.LOCAL_BLOCK,
         AttentionType.DILATED,
         AttentionType.NON_CAUSAL_1D,
+        AttentionType.RELATIVE_LOCAL_1D,
     ]
 
 
@@ -126,6 +128,7 @@ def local_attention_1d(x,
         hparams.num_heads,
         hparams.attention_dropout,
         attention_type=attention_type,
+        shared_rel=hparams.shared_rel,
         block_width=hparams.block_width,
         block_length=hparams.block_length,
         q_padding=q_padding,
@@ -302,6 +305,11 @@ def transformer_decoder_layers(inputs,
                                hparams,
                                attention_type="local_mask_right",
                                q_padding="LEFT", kv_padding="LEFT")
+      elif attention_type == AttentionType.RELATIVE_LOCAL_1D:
+        y = local_attention_1d(common_layers.layer_preprocess(x, hparams),
+                               hparams,
+                               attention_type="rel_local_mask_right",
+                               q_padding="LEFT", kv_padding="LEFT")
       elif attention_type == AttentionType.NON_CAUSAL_1D:
         y = local_attention_1d(common_layers.layer_preprocess(x, hparams),
                                hparams,
@@ -608,7 +616,7 @@ def prepare_decoder(targets, hparams):
     x = tf.reshape(x, [targets_shape[0],
                        x_shape[1], x_shape[2], hparams.hidden_size])
     x = add_pos_signals(x, hparams, "dec_pos")
-  x = common_layers.cast_like(x, targets)
+  # x = common_layers.cast_like(x, targets)
   return x, x_shape[1], x_shape[2]
 
 
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index fe493b2d5..97aab456c 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -220,7 +220,7 @@ def convert_rgb_to_symmetric_real(x):
     # Use the formula (value/127.5) - 1 to convert each channel value into a
     # real number in the range -1 to 1. We use 127.5 instead of 128 because
     # the intensities are in the range 0 to 255. This is used for dmol.
-    x = (x / 127.5) - 1
+    x = (x / 128) - 1
     return x
 
 
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index bfd71eba0..9e8d84e6e 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -171,7 +171,7 @@ def image_transformer_base():
   """Set of hyperparameters."""
   hparams = common_hparams.basic_params1()
   hparams.hidden_size = 512
-  hparams.batch_size = 1
+  hparams.batch_size = 4
   hparams.max_length = 3075
   hparams.dropout = 0.0
   hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
@@ -236,6 +236,9 @@ def image_transformer_base():
   hparams.add_hparam("moe_overhead_eval", 2.0)
   hparams.moe_num_experts = 8
   hparams.moe_loss_coef = 1e-3
+
+  # These parameters are for relative attention
+  hparams.add_hparam("shared_rel", False)  # share relative embeddings
   return hparams
 
 
@@ -245,6 +248,13 @@ def imagetransformer_base():
   return hparams
 
 
+@registry.register_hparams
+def imagetransformer_base_rel():
+  hparams = imagetransformer_base()
+  hparams.dec_attention_type = cia.AttentionType.RELATIVE_LOCAL_1D
+  return hparams
+
+
 @registry.register_hparams
 def imagetransformer_sep_channels():
   """separate rgb embeddings."""
@@ -437,6 +447,21 @@ def imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_m():
   return hparams
 
 
+@registry.register_hparams
+def imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_m_rel():
+  hparams = imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_k()
+  hparams.batch_size = 8
+  hparams.dec_attention_type = cia.AttentionType.RELATIVE_LOCAL_1D
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_m_relsh():
+  hparams = imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_m_rel()
+  hparams.shared_rel = True
+  return hparams
+
+
 @registry.register_hparams
 def imagetransformerpp_base_14l_8h_big_uncond_dr03_dan_p():
   """Gets to 2.92 in just under 4 days on 8 p100s."""

From 867f7b34bc97048db544025ec6a59eb17c246d1b Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 28 Jun 2018 16:49:16 -0700
Subject: [PATCH 0269/2720] Revert two minor changes from relative attention
 commit.

PiperOrigin-RevId: 202568883
---
 tensor2tensor/layers/common_image_attention.py | 2 +-
 tensor2tensor/layers/common_layers.py          | 7 +++----
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 4a54f0db9..22f8bbad1 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -616,7 +616,7 @@ def prepare_decoder(targets, hparams):
     x = tf.reshape(x, [targets_shape[0],
                        x_shape[1], x_shape[2], hparams.hidden_size])
     x = add_pos_signals(x, hparams, "dec_pos")
-  # x = common_layers.cast_like(x, targets)
+  x = common_layers.cast_like(x, targets)
   return x, x_shape[1], x_shape[2]
 
 
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 97aab456c..709920ddc 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -217,10 +217,9 @@ def convert_rgb_to_symmetric_real(x):
   """Conversion of pixel values to real numbers."""
   with tf.name_scope("rgb_to_real", values=[x]):
     x = tf.to_float(x)
-    # Use the formula (value/127.5) - 1 to convert each channel value into a
-    # real number in the range -1 to 1. We use 127.5 instead of 128 because
-    # the intensities are in the range 0 to 255. This is used for dmol.
-    x = (x / 128) - 1
+    # Convert each pixel intensity in [0, 1, 2, ..., 255] into a real number in
+    # the range [-1, 1].
+    x = (x / 127.5) - 1
     return x
 
 

From 5717e591ebe91487a8cba3891124d150d143cc2c Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Thu, 28 Jun 2018 17:11:48 -0700
Subject: [PATCH 0270/2720] Internal change

PiperOrigin-RevId: 202572428
---
 tensor2tensor/data_generators/all_problems.py |  1 +
 tensor2tensor/data_generators/imdb.py         |  6 ++-
 tensor2tensor/data_generators/lm1b.py         |  4 ++
 tensor2tensor/data_generators/lm1b_imdb.py    | 39 +++++++++++++++++++
 .../data_generators/multi_problem.py          | 33 +++++++++-------
 tensor2tensor/data_generators/problem.py      |  2 +
 6 files changed, 69 insertions(+), 16 deletions(-)
 create mode 100644 tensor2tensor/data_generators/lm1b_imdb.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index afae0033a..ef75308a5 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -43,6 +43,7 @@
     "tensor2tensor.data_generators.lambada",
     "tensor2tensor.data_generators.librispeech",
     "tensor2tensor.data_generators.lm1b",
+    "tensor2tensor.data_generators.lm1b_imdb",
     "tensor2tensor.data_generators.mnist",
     "tensor2tensor.data_generators.mscoco",
     "tensor2tensor.data_generators.multinli",
diff --git a/tensor2tensor/data_generators/imdb.py b/tensor2tensor/data_generators/imdb.py
index acd615653..1feaf40d1 100644
--- a/tensor2tensor/data_generators/imdb.py
+++ b/tensor2tensor/data_generators/imdb.py
@@ -99,8 +99,12 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
 
 @registry.register_problem
 class SentimentIMDBCharacters(SentimentIMDB):
-  """IMDB Sentiment classification, character level."""
+  """IMDB sentiment classification, character level."""
 
   @property
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
+
+  @property
+  def task_id(self):
+    return problem.SpaceID.EN_CHR_SENT
diff --git a/tensor2tensor/data_generators/lm1b.py b/tensor2tensor/data_generators/lm1b.py
index 84a03fef7..d16bd109b 100644
--- a/tensor2tensor/data_generators/lm1b.py
+++ b/tensor2tensor/data_generators/lm1b.py
@@ -175,6 +175,10 @@ class LanguagemodelLm1bCharacters(LanguagemodelLm1b32k):
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
 
+  @property
+  def task_id(self):
+    return problem.SpaceID.EN_CHR
+
 
 @registry.register_problem
 class LanguagemodelLm1bCharactersPacked(LanguagemodelLm1bCharacters):
diff --git a/tensor2tensor/data_generators/lm1b_imdb.py b/tensor2tensor/data_generators/lm1b_imdb.py
new file mode 100644
index 000000000..b4b8ba9aa
--- /dev/null
+++ b/tensor2tensor/data_generators/lm1b_imdb.py
@@ -0,0 +1,39 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generators for LM1B and IMDb combined data-set."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators import imdb
+from tensor2tensor.data_generators import lm1b
+from tensor2tensor.data_generators import multi_problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+
+
+@registry.register_problem
+class LanguagemodelLm1bSentimentIMDB(multi_problem.MultiProblem):
+  """LM1b and IMDb mixed problem class for multitask learning."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelLm1bSentimentIMDB, self).__init__(was_reversed, was_copy)
+    self.task_list.append(lm1b.LanguagemodelLm1bCharacters())
+    self.task_list.append(imdb.SentimentIMDBCharacters())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index f708ca475..e1cfdad44 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -37,14 +37,19 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
       assert task.vocab_type == text_problems.VocabType.CHARACTER
       task.generate_data(data_dir, tmp_dir, task_id)
 
-  def add_task_id(self, task_id, serialized_example):
+  def add_task_id(self, task, example):
     """Convert example to code switching mode by adding a task id."""
-    serialized_example["targets"] = tf.concat(serialized_example["inputs"],
-                                              [task_id],
-                                              serialized_example["targets"], 0)
-    del serialized_example["inputs"]
+    if task.has_inputs:
+      inputs = example.pop("inputs")
+      concat_list = [inputs, [task.task_id], example["targets"]]
+    else:
+      concat_list = [[task.task_id], example["targets"]]
+
+    example["targets"] = tf.concat(concat_list, 0)
+    return example
 
   def filepattern(self, data_dir, mode, shard=None):
+    print("Generating multi problem filepattern")
     return [task.filepattern(data_dir, mode, shard) for task in self.task_list]
 
   def dataset(self,
@@ -70,21 +75,19 @@ def dataset(self,
                                   hparams, preprocess, dataset_split,
                                   shard, partition_id, num_partitions,
                                   max_records).repeat()
-      task_dataset = task_dataset.map(
-          # pylint: disable=cell-var-from-loop
-          lambda x: self.add_task_id(task.task_id, x),
-          num_parallel_threads=num_threads)
+      # pylint: disable=cell-var-from-loop
+      task_dataset = task_dataset.map(lambda x: self.add_task_id(task, x))
       datasets.append(task_dataset)
 
-    def flatten_zip(zipped):
-      flattened = tf.data.Dataset.from_tensors(zipped[0])
-      for ex in zipped[1:]:
-        flattened.concatenate(tf.data.Dataset.from_tensors(ex))
+    self._hparams = self.task_list[0].get_hparams()
 
-      return flattened
+    # TODO(urvashik): make this independent of the number of tasks
+    def flatten_zip(d0, d1):
+      return tf.data.Dataset.from_tensors(d0).concatenate(
+          tf.data.Dataset.from_tensors(d1))
 
     if is_training:
-      single_mtl_dataset = tf.data.Dataset.zip(datasets).flat_map(
+      single_mtl_dataset = tf.data.Dataset.zip(tuple(datasets)).flat_map(
           flatten_zip)
     else:
       single_mtl_dataset = datasets[0]
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index d2d100768..7bce7e1c6 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -100,6 +100,8 @@ class SpaceID(object):
   STROKES = 29
   # Pickled Python
   PICKLED_PYTHON = 30
+  # English characters sentiment
+  EN_CHR_SENT = 31
 
 
 def default_model_hparams():

From 8ee32845fa74f316a15b45063e4b756ba564e4eb Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 28 Jun 2018 22:59:21 -0700
Subject: [PATCH 0271/2720] add multilabel match3 metrics

PiperOrigin-RevId: 202601314
---
 tensor2tensor/utils/metrics.py      | 41 +++++++++++++++++++++++++++++
 tensor2tensor/utils/metrics_test.py | 23 ++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 9a644f9eb..fac0fabe6 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -37,6 +37,7 @@ class Metrics(object):
   ACC = "accuracy"
   ACC_TOP5 = "accuracy_top5"
   ACC_PER_SEQ = "accuracy_per_sequence"
+  ACC_MULTILABEL_MATCH3 = "accuracy_multilabel_match3"
   NEG_LOG_PERPLEXITY = "neg_log_perplexity"
   APPROX_BLEU = "approx_bleu_score"
   RMSE = "rmse"
@@ -280,6 +281,45 @@ def padded_accuracy(predictions,
     return tf.to_float(tf.equal(outputs, padded_labels)), weights
 
 
+def multilabel_accuracy_matchk(predictions,
+                               labels,
+                               k,
+                               weights_fn=common_layers.weights_nonzero):
+  """Used to evaluate the VQA accuracy.
+
+  Let n be the times that predictions appear in labels, then final score
+  is min(n/k, 1).
+  Refer to https://arxiv.org/pdf/1505.00468.pdf.
+
+  Args:
+    predictions: A tensor with shape [batch_size, 1, 1, 1, vocab_size].
+    labels: A tensor with shape [batch_size, length, 1, 1].
+    k: A tensor constant.
+    weights_fn: weight function.
+  Returns:
+    scores: min(n/k, 1).
+    weights: 1 if labels contains non-zero label else 0.
+
+  """
+  predictions = tf.to_int32(tf.argmax(predictions, axis=-1))
+  length = tf.shape(labels)[1]
+  predictions = tf.tile(predictions, [1, length, 1, 1])
+  scores = tf.to_float(tf.equal(predictions, labels))
+  scores = tf.reduce_sum(scores, axis=[1, 2, 3])
+  scores = tf.minimum(scores / tf.to_float(k), 1)
+
+  weights = weights_fn(labels)
+  weights = tf.reduce_sum(weights, axis=[1, 2, 3])
+  weights = tf.to_float(tf.greater(weights, 0.))
+
+  return scores, weights
+
+
+def multilabel_accuracy_match3(predictions, labels,
+                               weights_fn=common_layers.weights_nonzero):
+  return multilabel_accuracy_matchk(predictions, labels, 3, weights_fn)
+
+
 def set_precision(predictions, labels,
                   weights_fn=common_layers.weights_nonzero):
   """Precision of set predictions.
@@ -613,6 +653,7 @@ def metric_means():
     Metrics.ACC: padded_accuracy,
     Metrics.ACC_TOP5: padded_accuracy_top5,
     Metrics.ACC_PER_SEQ: padded_sequence_accuracy,
+    Metrics.ACC_MULTILABEL_MATCH3: multilabel_accuracy_match3,
     Metrics.NEG_LOG_PERPLEXITY: padded_neg_log_perplexity,
     Metrics.APPROX_BLEU: bleu_hook.bleu_score,
     Metrics.RMSE: padded_rmse,
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 3c19a7bf5..d3c9d9cb9 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -19,6 +19,7 @@
 import numpy as np
 from tensor2tensor.utils import metrics
 
+
 import tensorflow as tf
 
 
@@ -218,6 +219,28 @@ def testRocAuc(self):
       s = session.run(score)
     self.assertAlmostEqual(s, 0.750, places=3)
 
+  def testMultilabelMatch3(self):
+    predictions = np.random.randint(1, 5, size=(100, 1, 1, 1))
+    targets = np.random.randint(1, 5, size=(100, 10, 1, 1))
+    weights = np.random.randint(0, 2, size=(100, 1, 1, 1))
+    targets *= weights
+
+    predictions_repeat = np.repeat(predictions, 10, axis=1)
+    expected = (predictions_repeat == targets).astype(float)
+    expected = np.sum(expected, axis=(1, 2, 3))
+    expected = np.minimum(expected / 3.0, 1.)
+    expected = np.sum(expected * weights[:, 0, 0, 0]) / np.sum(weights)
+    with self.test_session() as session:
+      scores, weights_ = metrics.multilabel_accuracy_match3(
+          tf.one_hot(predictions, depth=5, dtype=tf.float32),
+          tf.constant(targets, dtype=tf.int32))
+      a, a_op = tf.metrics.mean(scores, weights_)
+      session.run(tf.local_variables_initializer())
+      session.run(tf.global_variables_initializer())
+      _ = session.run(a_op)
+      actual = session.run(a)
+    self.assertAlmostEqual(actual, expected)
+
 
 if __name__ == '__main__':
   tf.test.main()

From ab4c897cc09f7a9b80c2bb2cd129f1ae9376a923 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 29 Jun 2018 13:03:53 -0700
Subject: [PATCH 0272/2720] Added Image prefix to LsunBedrooms dataset.

PiperOrigin-RevId: 202694743
---
 tensor2tensor/data_generators/image_lsun.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/image_lsun.py b/tensor2tensor/data_generators/image_lsun.py
index ce49d02c7..373f8aa98 100644
--- a/tensor2tensor/data_generators/image_lsun.py
+++ b/tensor2tensor/data_generators/image_lsun.py
@@ -44,7 +44,7 @@ def _get_lsun(directory, category, split_name):
 
 
 @registry.register_problem
-class LsunBedrooms(image_utils.ImageProblem):
+class ImageLsunBedrooms(image_utils.ImageProblem):
   """LSUN Bedrooms."""
 
   @property

From 9fc6f55ad25cdb407d9e4fbc3c64931b57d33fc4 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Fri, 29 Jun 2018 13:15:05 -0700
Subject: [PATCH 0273/2720] Add hparams for vizier, remove obsolete ones and
 include for relative

PiperOrigin-RevId: 202696307
---
 tensor2tensor/models/image_transformer.py | 162 ++++------------------
 1 file changed, 28 insertions(+), 134 deletions(-)

diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 9e8d84e6e..91f5190ab 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -289,15 +289,6 @@ def imagetransformer_sep_channels_8l_multipos3():
   return hparams
 
 
-@registry.register_hparams
-def imagetransformer_sep_output_channels_8l():
-  """separate rgb embeddings."""
-  hparams = imagetransformer_sep_channels_8l()
-  hparams.sep_rgb_embed = True
-  hparams.sampling_method = "random"
-  return hparams
-
-
 @registry.register_hparams
 def imagetransformer_base_8l_8h_big_cond_dr03_dan():
   """big 1d model for conditional image generation.2.99 on cifar10."""
@@ -625,14 +616,6 @@ def imagetransformer_base_14l_8h_big_uncond():
   return hparams
 
 
-@registry.register_hparams
-def imagetransformer_base_14l_8h_big_uncond_dr01():
-  """big 1d model for conditional image generation."""
-  hparams = imagetransformer_base_14l_8h_big_uncond()
-  hparams.layer_prepostprocess_dropout = 0.1
-  return hparams
-
-
 @registry.register_hparams
 def imagetransformer_sep_channels_12l_16h_imagenet_large():
   """separate rgb embeddings."""
@@ -682,7 +665,7 @@ def imagetransformer_sep_output_channels_8l_local_and_global_att():
 @registry.register_hparams
 def imagetransformer_base_10l_16h_big_uncond_dr01_imgnet():
   """big 1d model for conditional image generation."""
-  hparams = imagetransformer_base_14l_8h_big_uncond_dr01()
+  hparams = imagetransformer_base_14l_8h_big_dr01()
   # num_hidden_layers
   hparams.num_decoder_layers = 10
   hparams.num_heads = 16
@@ -696,7 +679,7 @@ def imagetransformer_base_10l_16h_big_uncond_dr01_imgnet():
 @registry.register_hparams
 def imagetransformer_base_10l_16h_big_dr01_imgnet():
   """big 1d model for conditional image generation."""
-  hparams = imagetransformer_base_14l_8h_big_uncond_dr01()
+  hparams = imagetransformer_base_14l_8h_big_dr01()
   # num_hidden_layers
   hparams.num_decoder_layers = 10
   hparams.num_heads = 16
@@ -722,63 +705,6 @@ def imagetransformer_sep_channels_8l_8h():
   return hparams
 
 
-@registry.register_hparams
-def imagetransformer_sep_channels_10l_8h():
-  """separate rgb embeddings."""
-  hparams = imagetransformer_sep_channels_8l_8h()
-  hparams.num_hidden_layers = 8
-  hparams.learning_rate_warmup_steps = 16000
-  hparams.sampling_method = "random"
-  return hparams
-
-
-@registry.register_hparams
-def imagetransformer_sep_channels_12l_8h():
-  """separate rgb embeddings."""
-  hparams = imagetransformer_sep_channels_8l_8h()
-  hparams.num_hidden_layers = 12
-  hparams.batch_size = 2
-  hparams.learning_rate_warmup_steps = 16000
-  hparams.sampling_method = "random"
-  return hparams
-
-
-@registry.register_hparams
-def imagetransformer_sep_channels_12l_8h_nda():
-  """separate rgb embeddings."""
-  hparams = imagetransformer_sep_channels_8l_8h()
-  hparams.num_hidden_layers = 12
-  hparams.batch_size = 2
-  hparams.learning_rate_warmup_steps = 16000
-  hparams.sampling_method = "random"
-  hparams.layer_preprocess_sequence = "n"
-  hparams.layer_postprocess_sequence = "da"
-  return hparams
-
-
-@registry.register_hparams
-def imagetransformer_sep_channels_12l_8h_4k():
-  """separate rgb embeddings."""
-  hparams = imagetransformer_sep_channels_8l_8h()
-  hparams.num_hidden_layers = 12
-  hparams.batch_size = 2
-  hparams.learning_rate_warmup_steps = 4000
-  hparams.sampling_method = "random"
-  return hparams
-
-
-@registry.register_hparams
-def imagetransformer_sep_channels_12l_8h_sep_rgb():
-  """separate rgb embeddings."""
-  hparams = imagetransformer_sep_channels_8l_8h()
-  hparams.num_hidden_layers = 12
-  hparams.batch_size = 2
-  hparams.learning_rate_warmup_steps = 16000
-  hparams.sep_rgb_embed = True
-  hparams.sampling_method = "random"
-  return hparams
-
-
 @registry.register_hparams
 def imagetransformer_sep_channels_8l_8h_local_and_global_att():
   """separate rgb embeddings."""
@@ -794,28 +720,10 @@ def imagetransformer_sep_channels_8l_8h_local_and_global_att():
   return hparams
 
 
-@registry.register_hparams
-def imagetransformer_sep_channels_8l_self_att_ffn():
-  """separate rgb embeddings."""
-  hparams = imagetransformer_sep_channels_8l()
-  hparams.num_parts = 4
-  hparams.ffn_layer = "self_attention_ffn"
-  hparams.share_kv = True
-  return hparams
-
-
-@registry.register_hparams
-def imagetransformer_sep_channels_8l_glu_ffn():
-  """separate rgb embeddings."""
-  hparams = imagetransformer_sep_channels_8l()
-  hparams.ffn_layer = "glu_ffn"
-  return hparams
-
-
 @registry.register_hparams
 def imagetransformer_bas8l_8h_big_uncond_dr03_imgnet():
   """big 1d model for conditional image generation."""
-  hparams = imagetransformer_base_14l_8h_big_uncond_dr01()
+  hparams = imagetransformer_base_14l_8h_big_dr01()
   # num_hidden_layers
   hparams.num_decoder_layers = 8
   hparams.num_heads = 8
@@ -836,7 +744,8 @@ def imagetransformer_tiny():
 
 @registry.register_hparams
 def imagetransformer_tiny_tpu():
-  hparams = imagetransformer_base()
+  hparams = imagetransformer_tiny()
+  update_hparams_for_tpu(hparams)
   hparams.num_hidden_layers = 2
   hparams.hidden_size = 16
   hparams.batch_size = 2
@@ -877,7 +786,6 @@ def imagetransformer_moe_tiny():
 
 
 def update_hparams_for_tpu(hparams):
-  hparams.use_pad_remover = False  # where op not supported
   hparams.optimizer = "TrueAdam"
   hparams.batch_size = 4
 
@@ -1001,6 +909,26 @@ def imagetransformer_b12l_4h_b256_uncond_dr03_tpu():
   return hparams
 
 
+@registry.register_hparams
+def imagetransformer_b12l_4h_b256_uncond_dr03_rel_tpu():
+  """works very well on 4x4."""
+  hparams = imagetransformer_bas8l_8h_big_uncond_dr03_imgnet()
+  hparams.dec_attention_type = cia.RELATIVE_LOCAL_1D
+  return hparams
+
+
+@registry.register_ranged_hparams
+def imagetransformer_cifar_tpu_range(rhp):
+  """Range of hyperparameters for vizier."""
+  # After starting from base, set intervals for some parameters.
+  rhp.set_float("learning_rate", 0.01, 1.0, scale=rhp.LOG_SCALE)
+  rhp.set_discrete("num_decoder_layers", [8, 10, 12, 14, 16])
+  rhp.set_discrete("hidden_size", [256, 512, 1024])
+  rhp.set_discrete("block_length", [128, 256, 512])
+  rhp.set_categorical("dec_attention_type", [
+      cia.RELATIVE_LOCAL_1D, cia.LOCAL_1D])
+
+
 @registry.register_hparams
 def imagetransformer_b12l_4h_b128_h512_uncond_dr03_tpu():
   """TPU related big model."""
@@ -1023,33 +951,17 @@ def imagetransformer_b12l_4h_b128_h512_uncond_dr03_tpu():
 @registry.register_hparams
 def imagetransformer_b12l_4h_b128_h512_uncond_dr03_im():
   """TPU related imagenet model."""
-  hparams = imagetransformer_bas8l_8h_big_uncond_dr03_imgnet()
+  hparams = imagetransformer_b12l_4h_b256_uncond_dr03_tpu()
   update_hparams_for_tpu(hparams)
   hparams.batch_size = 4
-  hparams.num_heads = 4   # heads are expensive on tpu
-  hparams.num_decoder_layers = 12
-  hparams.block_length = 128
-  hparams.hidden_size = 512
-  hparams.filter_size = 2048
-  hparams.learning_rate = 0.2
-  hparams.learning_rate_warmup_steps = 6000
-  hparams.layer_preprocess_sequence = "none"
-  hparams.layer_postprocess_sequence = "dan"
   hparams.layer_prepostprocess_dropout = 0.1
   return hparams
 
 
 @registry.register_hparams
-def imagetransformer_b12l_4h_small_uncond_dr03_tpu():
+def imagetransformer_b12l_4h_uncond_dr03_tpu():
   """TPU related small model."""
-  hparams = imagetransformer_bas8l_8h_big_uncond_dr03_imgnet()
-  update_hparams_for_tpu(hparams)
-  hparams.batch_size = 4
-  hparams.num_heads = 4   # heads are expensive on tpu
-  hparams.num_decoder_layers = 8
-  hparams.block_length = 256
-  hparams.hidden_size = 512
-  hparams.filter_size = 2048
+  hparams = imagetransformer_b12l_4h_b256_uncond_dr03_tpu()
   hparams.learning_rate = 0.2
   hparams.learning_rate_warmup_steps = 4000
   hparams.layer_preprocess_sequence = "none"
@@ -1092,24 +1004,6 @@ def imagetransformer_b12l_8h_b256_uncond_dr03_tpu():
   return hparams
 
 
-@registry.register_hparams
-def imagetransformer_b12l_4h_b256_uncond_dr03_lr025_tpu():
-  hparams = imagetransformer_b12l_4h_b256_uncond_dr03_tpu()
-  update_hparams_for_tpu(hparams)
-  hparams.learning_rate = 0.25
-  hparams.learning_rate_warmup_steps = 10000
-  return hparams
-
-
-@registry.register_hparams
-def imagetransformer_b10l_4h_big_uncond_dr03_lr05_tpu():
-  hparams = imagetransformer_b10l_4h_big_uncond_dr03_lr025_tpu()
-  update_hparams_for_tpu(hparams)
-  hparams.learning_rate = 0.5
-  hparams.learning_rate_warmup_steps = 16000
-  return hparams
-
-
 @registry.register_hparams
 def imagetransformer_b10l_4h_big_uncond_dr01_tpu():
   """big 1d model for conditional image generation."""

From 3acf6f906e26614f684507f6e4c603a02aee88e1 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 29 Jun 2018 13:34:45 -0700
Subject: [PATCH 0274/2720] vqa attention test

PiperOrigin-RevId: 202699226
---
 tensor2tensor/layers/modalities.py | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index e93849394..ab2b46966 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -723,7 +723,7 @@ def top(self, body_output, _):
     """
     with tf.variable_scope(self.name):
       x = body_output
-      x = tf.reduce_mean(x, axis=[1, 2], keep_dims=True)
+      x = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
       res = tf.layers.dense(x, self._vocab_size)
       return tf.expand_dims(res, 3)
 
@@ -732,29 +732,33 @@ def top(self, body_output, _):
 class MultiLabelModality(ClassLabelModality):
   """Used for multi label task."""
 
+  @property
   def targets_weights_fn(self):
     """Target weight function for multi label, defaults to nonzero labels."""
-    weights_fn = common_layers.weights_nonzero
-    return weights_fn
+    return common_layers.weights_nonzero
 
   def loss(self, top_out, targets):
     """Average loss over the labels."""
     logits = top_out
     num_labels = tf.shape(targets)[1]
-    logits = tf.tile(logits, [1, num_labels, 1, 1])
+    logits = tf.tile(logits, [1, num_labels, 1, 1, 1])
 
     xent, weights = common_layers.padded_cross_entropy(
         logits,
         targets,
         self._model_hparams.label_smoothing,
         weights_fn=self.targets_weights_fn,
+        reduce_sum=False,
     )
     xent = tf.squeeze(xent, [2, 3])
-    weights = tf.squeeze(xent, [2, 3])
+    weights = tf.squeeze(weights, [2, 3])
     # average loss over all labels
-    loss = (
-        tf.reduce_sum(xent, axis=1) / (tf.reduce_sum(weights, axis=1) + 1e-8))
-    return tf.reduce_mean(loss)
+    loss = tf.reduce_sum(xent, axis=1)
+    weights = tf.reduce_sum(weights, axis=1)
+    loss /= (weights + 1e-8)
+    weights = tf.to_float(tf.greater(weights, 0.))
+
+    return tf.reduce_sum(loss*weights), tf.reduce_sum(weights)
 
 
 @registry.register_class_label_modality("onehot")

From 573e128c9c7fd8c97142709639af016001f35a56 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 29 Jun 2018 13:42:50 -0700
Subject: [PATCH 0275/2720] more saving filenames issues

PiperOrigin-RevId: 202700295
---
 tensor2tensor/utils/decoding.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index baeb6f675..b132e1b83 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -79,7 +79,7 @@ def log_decode_results(inputs,
     def fix_and_save_video(vid, prefix):
       save_path_template = os.path.join(
           output_dir,
-          "%s_%s_%d_{}.png" % (problem_name, prefix, prediction_idx))
+          "%s_%s_%05d_{:05d}.png" % (problem_name, prefix, prediction_idx))
       # this is only required for predictions
       if vid.shape[-1] == 1:
         vid = np.squeeze(vid, axis=-1)

From d9d596cf53c24f269c7a72431d145ec865a53738 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 29 Jun 2018 14:19:33 -0700
Subject: [PATCH 0276/2720] Clean up docstrings and inline comments in
 common_attention.

PiperOrigin-RevId: 202706161
---
 tensor2tensor/layers/common_attention.py | 265 +++++++++++++----------
 1 file changed, 146 insertions(+), 119 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 7752a95cf..ae0193bde 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -830,8 +830,10 @@ def embedding_to_padding(emb):
 
   Args:
     emb: a Tensor with shape [..., depth].
+
   Returns:
-    a float Tensor with shape [...].
+    a float Tensor with shape [...]. Each element is 1 if its corresponding
+    embedding vector is all zero, and is 0 otherwise.
   """
   emb_sum = tf.reduce_sum(tf.abs(emb), axis=-1)
   return tf.to_float(tf.equal(emb_sum, 0.0))
@@ -994,39 +996,37 @@ def attention_bias_proximal(length):
 
 
 @expert_utils.add_name_scope()
-def attention_bias_batch(
-    batch_coordinates_q,
-    batch_coordinates_k=None,
-    condition_fn=None,
-):
+def attention_bias_batch(batch_coordinates_q,
+                         batch_coordinates_k=None,
+                         condition_fn=None):
   """Generate a mask to prevent the batch to attend to each others.
 
   Args:
-    batch_coordinates_q (tf.Tensor): int32 of shape [length_q, 1] containing the
+    batch_coordinates_q: Int-like Tensor of shape [length_q, 1] containing the
       coordinates of the batches
-    batch_coordinates_k (tf.Tensor): int32 of shape [length_k, 1] containing the
-      coordinates of the batches. If None, do self attention (q and k identical)
-    condition_fn (fct): A function defining which type of mask build
+    batch_coordinates_k: Int-like Tensor of shape [length_k, 1] containing the
+      coordinates of the batches. If None, do self-attention.
+    condition_fn: Callable defining the attention mask.
 
   Returns:
-    tf.Tensor: float32 mask of shape [length_q, length_k] containing either 0 or
-      -infinity (-1e9)
+    Float-like Tensor of shape [length_q, length_k] containing either 0 or
+    -infinity (-1e9).
   """
   if batch_coordinates_k is None:
     batch_coordinates_k = batch_coordinates_q
 
-  # Convert to float first because of b/25387198
+  # Convert to float first because of b/25387198.
   def to_float(bc):
     bc = tf.squeeze(bc, 1)
     bc = tf.to_float(bc)
     return bc
 
+  # Broadcast to create [length_q, length_k] mask.
   bc_v = tf.expand_dims(to_float(batch_coordinates_q), 1)
   bc_h = tf.expand_dims(to_float(batch_coordinates_k), 0)
-  bias_batch = bc_h - bc_v  # Broadcast to create [length_q, length_k] mask.
-  # Threshold non zeros to 1.0.
+  bias_batch = bc_h - bc_v
   bias_batch = condition_fn(bias_batch)
-  bias_batch *= -1e9  # Set non zeros to -infinity
+  bias_batch *= -1e9
   return bias_batch
 
 
@@ -1463,14 +1463,16 @@ def dot_product_attention(q,
                           make_image_summary=True,
                           save_weights_to=None,
                           dropout_broadcast_dims=None):
-  """dot-product attention.
+  """Dot-product attention.
 
   Args:
-    q: a Tensor with shape [batch, heads, length_q, depth_k]
-    k: a Tensor with shape [batch, heads, length_kv, depth_k]
-    v: a Tensor with shape [batch, heads, length_kv, depth_v]
+    q: Tensor with shape [..., length_q, depth_k].
+    k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
+      match with q.
+    v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
+      match with q.
     bias: bias Tensor (see attention_bias())
-    dropout_rate: a floating point number
+    dropout_rate: a float.
     image_shapes: optional tuple of integer scalars.
       see comments for attention_image_summary()
     name: an optional string
@@ -1478,16 +1480,15 @@ def dot_product_attention(q,
     save_weights_to: an optional dictionary to capture attention weights
       for visualization; the weights tensor will be appended there under
       a string key created from the variable scope (including name).
-    dropout_broadcast_dims:  an optional list of integers less than 4
-      specifying in which dimensions to broadcast the dropout decisions.
-      saves memory.
+    dropout_broadcast_dims: an optional list of integers less than rank of q.
+      Specifies in which dimensions to broadcast the dropout decisions.
+
   Returns:
-    A Tensor.
+    Tensor with shape [..., length_q, depth_v].
   """
   with tf.variable_scope(
       name, default_name="dot_product_attention", values=[q, k, v]) as scope:
-    # [batch, num_heads, query_length, memory_length]
-    logits = tf.matmul(q, k, transpose_b=True)
+    logits = tf.matmul(q, k, transpose_b=True)  # [..., length_q, length_kv]
     if bias is not None:
       bias = common_layers.cast_like(bias, logits)
       logits += bias
@@ -1495,7 +1496,7 @@ def dot_product_attention(q,
     if save_weights_to is not None:
       save_weights_to[scope.name] = weights
       save_weights_to[scope.name + "/logits"] = logits
-    # dropping out the attention links for each of the heads
+    # Drop out attention links for each head.
     weights = common_layers.dropout_with_broadcast_dims(
         weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
     if common_layers.should_generate_summaries() and make_image_summary:
@@ -1763,9 +1764,9 @@ def dot_product_self_attention_relative_v2(q,
 def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
   """Attention to the source and a neighborhood to the left within a block.
 
-  The sequence is divided into blocks of length block_size.
-  Attention for a given query position can only see memory positions
-  less than or equal to the query position in the corresponding block
+  The sequence is divided into blocks of length block_length. Attention for a
+  given query position can only see memory positions less than or equal to the
+  query position in the corresponding block.
 
   Args:
     q: a Tensor with shape [batch, heads, length, depth_k]
@@ -1786,6 +1787,7 @@ def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
       if const is not None:
         block_length = int(const)
 
+    # Pad query, key, value to ensure multiple of block length.
     depth_k = common_layers.shape_list(k)[3]
     depth_v = common_layers.shape_list(v)[3]
     original_length = length
@@ -1795,20 +1797,23 @@ def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
     q = tf.pad(q, padding)
     k = tf.pad(k, padding)
     v = tf.pad(v, padding)
+
+    # Compute attention for all subsequent query blocks.
     num_blocks = tf.div(length, block_length)
-    # compute attention for all subsequent query blocks.
     q = tf.reshape(q, [batch, heads, num_blocks, block_length, depth_k])
     k = tf.reshape(k, [batch, heads, num_blocks, block_length, depth_k])
     v = tf.reshape(v, [batch, heads, num_blocks, block_length, depth_v])
-    # attention shape: [batch, heads, num_blocks, block_length, block_length]
+    # [batch, heads, num_blocks, block_length, block_length]
     attention = tf.matmul(q, k, transpose_b=True)
     attention += tf.reshape(
         attention_bias_lower_triangle(block_length),
         [1, 1, 1, block_length, block_length])
     attention = tf.nn.softmax(attention)
-    # initial output shape: [batch, heads, num_blocks, block_length, depth_v]
+    # [batch, heads, num_blocks, block_length, depth_v]
     output = tf.matmul(attention, v)
     output = tf.reshape(output, [batch, heads, -1, depth_v])
+
+    # Remove the padding if introduced.
     output = tf.slice(output, [0, 0, 0, 0], [-1, -1, original_length, -1])
     output.set_shape(v_shape)
     return output
@@ -1853,13 +1858,9 @@ def masked_local_attention_1d(q,
                               name=None):
   """Attention to the source position and a neighborhood to the left of it.
 
-  The sequence is divided into blocks of length block_size.
-  Attention for a given query position can only see memory positions
-  less than or equal to the query position, in the corresponding block
-  and the previous block.
-
-  If mask_right is True, then a target position cannot see greater source
-  positions.
+  The sequence is divided into blocks of length block_length. Attention for a
+  given query position can only see memory positions less than or equal to the
+  query position, in the corresponding block and the previous block.
 
   Args:
     q: a Tensor with shape [batch, heads, length, depth_k]
@@ -1889,6 +1890,8 @@ def masked_local_attention_1d(q,
     else:
       block_length = tf.where(
           tf.less(length, block_length * 2), length, block_length)
+
+    # Pad query, key, value to ensure multiple of block length.
     depth_k = common_layers.shape_list(k)[3]
     depth_v = common_layers.shape_list(v)[3]
     original_length = length
@@ -1904,7 +1907,7 @@ def masked_local_attention_1d(q,
     else:
       num_blocks = tf.div(length, block_length)
 
-    # compute attention for the first query block.
+    # Compute attention for the first query block.
     first_q = tf.slice(q, [0, 0, 0, 0], [-1, -1, block_length, -1])
     first_k = tf.slice(k, [0, 0, 0, 0], [-1, -1, block_length, -1])
     first_v = tf.slice(v, [0, 0, 0, 0], [-1, -1, block_length, -1])
@@ -1916,9 +1919,9 @@ def masked_local_attention_1d(q,
         attention_bias_lower_triangle(block_length),
         dropout_rate=dropout_rate,
         make_image_summary=make_image_summary,
-        name="fist_block")
+        name="first_block")
 
-    # compute attention for all subsequent query blocks.
+    # Compute attention for all subsequent query blocks.
     q = tf.reshape(q, [batch, heads, num_blocks, block_length, depth_k])
     k = tf.reshape(k, [batch, heads, num_blocks, block_length, depth_k])
     v = tf.reshape(v, [batch, heads, num_blocks, block_length, depth_v])
@@ -1953,6 +1956,8 @@ def masked_local_attention_1d(q,
     output = tf.reshape(
         output, [batch, heads, (num_blocks - 1) * block_length, depth_v])
     output = tf.concat([first_output, output], axis=2)
+
+    # Remove the padding if introduced.
     output = tf.slice(output, [0, 0, 0, 0], [-1, -1, original_length, -1])
     output = tf.reshape(output, [batch, heads, original_length, depth_v])
     return output
@@ -2143,7 +2148,11 @@ def _reshape_for_relative(x):
 
 
 def local_attention_1d(q, k, v, block_length=128, filter_width=100, name=None):
-  """strided block local self-attention.
+  """Strided block local self-attention.
+
+  The sequence is divided into blocks of length block_length. Attention for a
+  given query position can see all memory positions in the corresponding block
+  and filter_width many positions to the left of the block.
 
   Args:
     q: a Tensor with shape [batch, heads, length, depth_k]
@@ -2164,7 +2173,7 @@ def local_attention_1d(q, k, v, block_length=128, filter_width=100, name=None):
     num_heads = common_layers.shape_list(q)[1]
     original_length = common_layers.shape_list(q)[2]
 
-    # making sure q is a multiple of d
+    # Pad query, key, value to ensure multiple of corresponding lengths.
     def pad_to_multiple(x, pad_length):
       x_length = common_layers.shape_list(x)[2]
       return tf.pad(x, [[0, 0], [0, 0], [0, -x_length % pad_length], [0, 0]])
@@ -2176,24 +2185,21 @@ def pad_l_and_r(x, pad_length):
     k = pad_to_multiple(k, block_length)
     v = pad_to_multiple(v, block_length)
 
-    # Setting up q blocks
+    # Set up query blocks.
     new_q_shape = common_layers.shape_list(q)
-    # Setting up q blocks
     q = tf.reshape(q, [
         new_q_shape[0], new_q_shape[1], new_q_shape[2] // block_length,
         block_length, new_q_shape[3]
     ])
 
-    # Setting up k and v values
+    # Set up key and value blocks.
+    # Get gather indices.
     k = pad_l_and_r(k, filter_width)
     v = pad_l_and_r(v, filter_width)
-
     length = common_layers.shape_list(k)[2]
     full_filter_width = block_length + 2 * filter_width
-    # getting gather indices
     indices = tf.range(0, length, delta=1, name="index_range")
-    # making indices [1, length, 1] to appy convs
-    indices = tf.reshape(indices, [1, -1, 1])
+    indices = tf.reshape(indices, [1, -1, 1])  # [1, length, 1] for convs
     kernel = tf.expand_dims(tf.eye(full_filter_width), axis=1)
     gather_indices = tf.nn.conv1d(
         tf.cast(indices, tf.float32),
@@ -2204,11 +2210,10 @@ def pad_l_and_r(x, pad_length):
 
     gather_indices = tf.squeeze(tf.cast(gather_indices, tf.int32), axis=0)
 
-    # [length, batch, heads, dim]
+    # Reshape keys and values to [length, batch, heads, dim] for gather. Then
+    # reshape to [batch, heads, blocks, block_length + filter_width, dim].
     k_t = tf.transpose(k, [2, 0, 1, 3])
     k_new = tf.gather(k_t, gather_indices)
-
-    # [batch, heads, blocks, block_length, dim]
     k_new = tf.transpose(k_new, [2, 3, 0, 1, 4])
 
     attention_bias = tf.expand_dims(embedding_to_padding(k_new) * -1e9, axis=-2)
@@ -2226,13 +2231,25 @@ def pad_l_and_r(x, pad_length):
         name="local_1d",
         make_image_summary=False)
     output = tf.reshape(output, [batch_size, num_heads, -1, depth_v])
-    # Remove the padding if introduced
+
+    # Remove the padding if introduced.
     output = tf.slice(output, [0, 0, 0, 0], [-1, -1, original_length, -1])
     output.set_shape(v_shape)
     return output
 
 
 def reshape_by_blocks(x, x_shape, memory_block_size):
+  """Reshapes input by splitting its length over blocks of memory_block_size.
+
+  Args:
+    x: a Tensor with shape [batch, heads, length, depth]
+    x_shape: tf.TensorShape of x.
+    memory_block_size: Integer which divides length.
+
+  Returns:
+    Tensor with shape
+    [batch, heads, length // memory_block_size, memory_block_size, depth].
+  """
   x = tf.reshape(x, [
       x_shape[0], x_shape[1], x_shape[2] // memory_block_size,
       memory_block_size, x_shape[3]
@@ -2248,7 +2265,7 @@ def dilated_self_attention_1d(q,
                               gap_size=2,
                               num_memory_blocks=2,
                               name=None):
-  """dilated self-attention.
+  """Dilated self-attention.
 
   Args:
     q: a Tensor with shape [batch, heads, length, depth_k]
@@ -2273,7 +2290,7 @@ def dilated_self_attention_1d(q,
     num_heads = v_shape[1]
     original_length = common_layers.shape_list(q)[2]
 
-    # making sure q is a multiple of query block size
+    # Pad query, key, value to ensure multiple of corresponding lengths.
     def pad_to_multiple(x, pad_length):
       x_length = common_layers.shape_list(x)[2]
       return tf.pad(x, [[0, 0], [0, 0], [0, -x_length % pad_length], [0, 0]])
@@ -2284,26 +2301,25 @@ def pad_l_and_r(x, pad_length):
     q = pad_to_multiple(q, query_block_size)
     v = pad_to_multiple(v, query_block_size)
     k = pad_to_multiple(k, query_block_size)
-
     q.set_shape(v_list_shape)
     v.set_shape(v_list_shape)
     k.set_shape(v_list_shape)
-    # Setting up q blocks
+
+    # Set up query blocks.
     new_q_shape = common_layers.shape_list(q)
-    # Setting up q blocks
     q = reshape_by_blocks(q, new_q_shape, query_block_size)
     self_k_part = reshape_by_blocks(k, new_q_shape, query_block_size)
     self_v_part = reshape_by_blocks(v, new_q_shape, query_block_size)
 
-    # Setting up k and v windows
+    # Set up key and value windows.
     k_v_padding = (gap_size + memory_block_size) * num_memory_blocks
     k = pad_l_and_r(k, k_v_padding)
     v = pad_l_and_r(v, k_v_padding)
-    # getting gather indices
+
+    # Get gather indices.
     index_length = (new_q_shape[2] - query_block_size + memory_block_size)
     indices = tf.range(0, index_length, delta=1, name="index_range")
-    # making indices [1, length, 1] to appy convs
-    indices = tf.reshape(indices, [1, -1, 1])
+    indices = tf.reshape(indices, [1, -1, 1])  # [1, length, 1] for convs
     kernel = tf.expand_dims(tf.eye(memory_block_size), axis=1)
     gather_indices = tf.nn.conv1d(
         tf.cast(indices, tf.float32),
@@ -2314,7 +2330,7 @@ def pad_l_and_r(x, pad_length):
 
     gather_indices = tf.squeeze(tf.cast(gather_indices, tf.int32), axis=0)
 
-    # get left and right memory blocks for each query
+    # Get left and right memory blocks for each query.
     # [length, batch, heads, dim]
     k_t = tf.transpose(k, [2, 0, 1, 3])
     v_t = tf.transpose(v, [2, 0, 1, 3])
@@ -2356,7 +2372,8 @@ def pad_l_and_r(x, pad_length):
         name="dilated_1d",
         make_image_summary=False)
     output = tf.reshape(output, [batch_size, num_heads, -1, depth_v])
-    # Remove the padding if introduced
+
+    # Remove the padding if introduced.
     output = tf.slice(output, [0, 0, 0, 0], [-1, -1, original_length, -1])
     output.set_shape(v_list_shape)
     return output
@@ -2372,18 +2389,18 @@ def gather_dilated_memory_blocks(x,
   """Gathers blocks with gaps in between.
 
   Args:
-    x: A tensor of shape [length, batch, heads, depth]
-    num_memory_blocks:     num_memory_blocks: how many memory blocks to look
-      in "direction". Each will be separated by gap_size.
+    x: Tensor of shape [length, batch, heads, depth]
+    num_memory_blocks: how many memory blocks to look in "direction". Each will
+      be separated by gap_size.
     gap_size: an integer indicating the gap size
     query_block_size: an integer indicating size of query block
     memory_block_size: an integer indicating the size of a memory block.
     gather_indices: The indices to gather from.
     direction: left or right
+
   Returns:
-    a tensor of shape [batch, heads, blocks, block_length, depth]
+    Tensor of shape [batch, heads, blocks, block_length, depth]
   """
-
   gathered_blocks = []
   # gathering memory blocks
   for block_id in range(num_memory_blocks):
@@ -2414,7 +2431,7 @@ def masked_dilated_self_attention_1d(q,
                                      gap_size=2,
                                      num_memory_blocks=2,
                                      name=None):
-  """dilated self-attention. TODO(avaswani): Try it and write a paper on it.
+  """Dilated self-attention. TODO(avaswani): Try it and write a paper on it.
 
   Args:
     q: a Tensor with shape [batch, heads, length, depth_k]
@@ -2439,7 +2456,7 @@ def masked_dilated_self_attention_1d(q,
     num_heads = v_shape[1]
     original_length = common_layers.shape_list(q)[2]
 
-    # making sure q is a multiple of query block size
+    # Pad query, key, value to ensure multiple of corresponding lengths.
     def pad_to_multiple(x, pad_length):
       x_length = common_layers.shape_list(x)[2]
       return tf.pad(x, [[0, 0], [0, 0], [0, -x_length % pad_length], [0, 0]])
@@ -2453,23 +2470,23 @@ def pad_l(x, left_pad_length):
     q.set_shape(v_list_shape)
     v.set_shape(v_list_shape)
     k.set_shape(v_list_shape)
-    # Setting up q blocks
-    new_q_shape = common_layers.shape_list(q)
 
-    # Setting up q blocks
+    # Set up query blocks.
+    new_q_shape = common_layers.shape_list(q)
     q = reshape_by_blocks(q, new_q_shape, query_block_size)
+
+    # Set up key and value windows.
     self_k_part = reshape_by_blocks(k, new_q_shape, query_block_size)
     self_v_part = reshape_by_blocks(v, new_q_shape, query_block_size)
-    # Setting up k and v windows
     k_v_padding = (gap_size + memory_block_size) * num_memory_blocks
     k = pad_l(k, k_v_padding)
     v = pad_l(v, k_v_padding)
-    # getting gather indices
+
+    # Get gather indices.
     index_length = (new_q_shape[2] - query_block_size + memory_block_size)
 
     indices = tf.range(0, index_length, delta=1, name="index_range")
-    # making indices [1, length, 1] to appy convs
-    indices = tf.reshape(indices, [1, -1, 1])
+    indices = tf.reshape(indices, [1, -1, 1])  # [1, length, 1] for convs
     kernel = tf.expand_dims(tf.eye(memory_block_size), axis=1)
     gather_indices = tf.nn.conv1d(
         tf.cast(indices, tf.float32),
@@ -2479,7 +2496,7 @@ def pad_l(x, left_pad_length):
         name="gather_conv")
     gather_indices = tf.squeeze(tf.cast(gather_indices, tf.int32), axis=0)
 
-    # get left and right memory blocks for each query
+    # Get left and right memory blocks for each query.
     # [length, batch, heads, dim]
     k_t = tf.transpose(k, [2, 0, 1, 3])
     v_t = tf.transpose(v, [2, 0, 1, 3])
@@ -2491,7 +2508,7 @@ def pad_l(x, left_pad_length):
         v_t, num_memory_blocks, gap_size, query_block_size, memory_block_size,
         gather_indices)
 
-    # combine memory windows
+    # Combine memory windows.
     block_q_shape = common_layers.shape_list(q)
     masked_attention_bias = tf.tile(
         tf.expand_dims(attention_bias_lower_triangle(query_block_size), axis=0),
@@ -2514,7 +2531,8 @@ def pad_l(x, left_pad_length):
         name="dilated_1d",
         make_image_summary=False)
     output = tf.reshape(output, [batch_size, num_heads, -1, depth_v])
-    # Remove the padding if introduced
+
+    # Remove the padding if introduced.
     output = tf.slice(output, [0, 0, 0, 0], [-1, -1, original_length, -1])
     output.set_shape(v_list_shape)
     return output
@@ -2526,7 +2544,13 @@ def local_attention_2d(q,
                        query_shape=(8, 16),
                        memory_flange=(8, 16),
                        name=None):
-  """strided block local self-attention.
+  """Strided block local self-attention.
+
+  The 2-D sequence is divided into 2-D blocks of shape query_shape. Attention
+  for a given query position can only see memory positions less than or equal to
+  the query position. The memory positions are the corresponding block with
+  memory_flange many positions to add to the height and width of the block
+  (namely, left, top, and right).
 
   Args:
     q: a Tensor with shape [batch, heads, h, w, depth_k]
@@ -2545,21 +2569,21 @@ def local_attention_2d(q,
     q_shape = q.get_shape().as_list()
     v_shape = common_layers.shape_list(v)
 
+    # Pad query, key, value to ensure multiple of corresponding lengths.
     q = pad_to_multiple_2d(q, query_shape)
     k = pad_to_multiple_2d(k, query_shape)
     v = pad_to_multiple_2d(v, query_shape)
     padded_q_shape = common_layers.shape_list(q)
-    # Setting up k and v values
     paddings = [[0, 0], [0, 0], [memory_flange[0], memory_flange[1]],
                 [memory_flange[0], memory_flange[1]], [0, 0]]
     k = tf.pad(k, paddings)
     v = tf.pad(v, paddings)
 
-    # Setting up q blocks
+    # Set up query blocks.
     q_indices = gather_indices_2d(q, query_shape, query_shape)
     q_new = gather_blocks_2d(q, q_indices)
 
-    # Setting up k and v blocks
+    # Set up key and value blocks.
     memory_shape = (query_shape[0] + 2 * memory_flange[0],
                     query_shape[1] + 2 * memory_flange[1])
     k_and_v_indices = gather_indices_2d(k, memory_shape, query_shape)
@@ -2577,9 +2601,10 @@ def local_attention_2d(q,
         dropout_rate=0.,
         name="local_2d",
         make_image_summary=False)
-    # putting the representations back in the right place
+    # Put representations back into original shapes.
     output = scatter_blocks_2d(output, q_indices, padded_q_shape)
-    # Remove the padding if introduced
+
+    # Remove the padding if introduced.
     output = tf.slice(output, [0, 0, 0, 0, 0],
                       [-1, -1, v_shape[2], v_shape[3], -1])
     output.set_shape(q_shape)
@@ -2693,7 +2718,7 @@ def make_2d_block_raster_mask(query_shape, memory_flange):
           tf.concat(mask_pieces, axis=1)
       ],
       axis=1)
-  # 0. is visible location, 1.0 is masked.
+  # 0.0 is visible location, 1.0 is masked.
   return 1. - final_mask
 
 
@@ -2814,7 +2839,7 @@ def masked_local_attention_2d(q,
                               query_shape=(8, 16),
                               memory_flange=(8, 16),
                               name=None):
-  """strided block local self-attention.
+  """Strided block local self-attention.
 
     Each position in a query block can attend to all the generated queries in
     the query block, which are generated in raster scan, and positions that are
@@ -2842,12 +2867,15 @@ def masked_local_attention_2d(q,
     q_shape = q.get_shape().as_list()
     v_shape = common_layers.shape_list(v)
 
+    # Pad query to ensure multiple of corresponding lengths.
     q = pad_to_multiple_2d(q, query_shape)
     padded_q_shape = common_layers.shape_list(q)
-    # Setting up q blocks
+
+    # Set up query blocks.
     q_indices = gather_indices_2d(q, query_shape, query_shape)
     q_new = gather_blocks_2d(q, q_indices)
-    # Setting up k and v blocks
+
+    # Set up key and value blocks.
     k_flange, k_center = get_memory_region(k, query_shape, memory_flange,
                                            q_indices)
     v_flange, v_center = get_memory_region(v, query_shape, memory_flange,
@@ -2858,7 +2886,8 @@ def masked_local_attention_2d(q,
     else:
       k_new = k_center
       v_new = v_center
-    # Getting the masks ready
+
+    # Set up the masks.
     query_elements = np.prod(query_shape)
     padding_mask = None
     if k_flange is not None:
@@ -2875,7 +2904,7 @@ def masked_local_attention_2d(q,
         center_attention_bias,
         [v_center_shape[0], v_center_shape[1], v_center_shape[2], 1, 1])
     if padding_mask is not None:
-      # Combining the mask for padding and visible region
+      # Combine the mask for padding and visible region.
       attention_bias = tf.concat([padding_mask, center_attention_bias], axis=4)
     else:
       attention_bias = center_attention_bias
@@ -2888,9 +2917,11 @@ def masked_local_attention_2d(q,
         dropout_rate=0.,
         name="masked_local_2d",
         make_image_summary=False)
-    # putting the representations back in the right place
+
+    # Put representations back into original shapes.
     output = scatter_blocks_2d(output, q_indices, padded_q_shape)
-    # Remove the padding if introduced
+
+    # Remove the padding if introduced.
     output = tf.slice(output, [0, 0, 0, 0, 0],
                       [-1, -1, v_shape[2], v_shape[3], -1])
     output.set_shape(q_shape)
@@ -3077,7 +3108,7 @@ def multihead_attention(query_antecedent,
     The caching works by saving all the previous key and value values so that
     you are able to send just the last query location to this attention
     function. I.e. if the cache dict is provided it assumes the query is of the
-    shape [batch_size, 1, hiddem_dim] rather than the full memory.
+    shape [batch_size, 1, hidden_dim] rather than the full memory.
 
   Returns:
     The result of the attention transformation. The output shape is
@@ -3303,10 +3334,9 @@ def ffn_self_attention_layer(x,
   We use self-attention to do feedforward computations. We apply this function
   positionwise where for each position, we linearly transform the output to have
   depth filter_depth, and break up the result depth-wise into num_parts
-  contiguous parts.  The parts self-attend, we concatenate the results
-  depth-wise, and we linearly transform to a depth of output_depth. The
-  goal is to get multiplicative interactions between components of a
-  representation.
+  contiguous parts. The parts self-attend, we concatenate the results
+  depth-wise, and we linearly transform to a depth of output_depth. The goal is
+  to get multiplicative interactions between components of a representation.
 
   Args:
     x: a Tensor with shape [batch, length, channels]
@@ -3318,9 +3348,8 @@ def ffn_self_attention_layer(x,
     name: an optional string
 
   Returns:
-    A Tensor.
+    A Tensor with shape [batch, length, output_depth].
   """
-
   with tf.variable_scope(
       name, default_name="feedforward_self_attention", values=[x]):
     x_shape = common_layers.shape_list(x)
@@ -3369,8 +3398,8 @@ def parameter_attention(x,
   """Attention over parameters.
 
   We use the same multi-headed attention as in the other layers, but the memory
-  keys and values are model parameters.  There are no linear transformation
-  on the keys or values.
+  keys and values are model parameters. There are no linear transformation on
+  the keys or values.
 
   We are also a bit more careful about memory usage, since the number of
   memory positions may be very large.
@@ -3386,7 +3415,7 @@ def parameter_attention(x,
     name: an optional string
 
   Returns:
-    A Tensor.
+    A Tensor with shape [batch, length_q, output_depth].
   """
   with tf.variable_scope(name, default_name="parameter_attention", values=[x]):
     head_size_k = total_key_depth // num_heads
@@ -3451,15 +3480,13 @@ def coordinate_tensor(shape, axis):
   return tf.zeros(shape, dtype=tf.int32) + tf.reshape(r, r_shape)
 
 
-def self_attention_expert(
-    x,
-    batch_coordinate,
-    mask_right=True,
-    split_batch=False,
-    attention_num_head=1,
-    attention_kq_size=None,
-    attention_v_size=None,
-):
+def self_attention_expert(x,
+                          batch_coordinate,
+                          mask_right=True,
+                          split_batch=False,
+                          attention_num_head=1,
+                          attention_kq_size=None,
+                          attention_v_size=None):
   """Implementing attention that runs inside each expert.
 
   Args:
@@ -4283,7 +4310,7 @@ def construct_bias_vectors(t, axis):
 
 
 def scaled_dot_product_attention_simple(q, k, v, bias, name=None):
-  """scaled dot-product attention.  One head.  One spatial dimension.
+  """Scaled dot-product attention. One head. One spatial dimension.
 
   Args:
     q: a Tensor with shape [batch, length_q, depth_k]

From 847ccf7ff70bc09aada863bacccc5bea54f58bc2 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 29 Jun 2018 23:47:35 -0700
Subject: [PATCH 0277/2720] add support for multiple decoding

PiperOrigin-RevId: 202756437
---
 tensor2tensor/data_generators/video_utils.py |  6 +--
 tensor2tensor/utils/compute_video_metrics.py | 11 +++-
 tensor2tensor/utils/decoding.py              | 53 +++++++++++++++-----
 tensor2tensor/utils/video_metrics.py         | 45 ++++++++++++++---
 4 files changed, 91 insertions(+), 24 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 6d3c47a6b..e29be7887 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -45,15 +45,15 @@ def summarize_video_metrics(hook_args):
   problem_name = hook_args.problem.name
   current_problem = hook_args.problem
   hparams = hook_args.hparams
-  output_dir = hook_args.output_dir
+  output_dirs = hook_args.output_dirs
   frame_shape = [
       current_problem.frame_height, current_problem.frame_width,
       current_problem.num_channels
   ]
   metrics_graph = tf.Graph()
   with metrics_graph.as_default():
-    metrics_results = video_metrics.compute_video_metrics(
-        output_dir, problem_name, hparams.video_num_target_frames, frame_shape)
+    metrics_results, _ = video_metrics.compute_video_metrics(
+        output_dirs, problem_name, hparams.video_num_target_frames, frame_shape)
 
   summary_values = []
   for name, array in six.iteritems(metrics_results):
diff --git a/tensor2tensor/utils/compute_video_metrics.py b/tensor2tensor/utils/compute_video_metrics.py
index a13d6c397..45b077004 100644
--- a/tensor2tensor/utils/compute_video_metrics.py
+++ b/tensor2tensor/utils/compute_video_metrics.py
@@ -18,6 +18,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import os
+
 from tensor2tensor.bin import t2t_decoder
 from tensor2tensor.utils import video_metrics
 import tensorflow as tf
@@ -32,8 +34,15 @@ def main(_):
   frame_shape = [problem.frame_height,
                  problem.frame_width,
                  problem.num_channels]
+  decode_hp = t2t_decoder.create_decode_hparams()
+
+  output_dirs = [
+      os.path.join(FLAGS.output_dir, "decode_%05d" % decode_id)
+      for decode_id in range(decode_hp.num_decodes)
+  ]
+
   video_metrics.compute_and_save_video_metrics(
-      FLAGS.output_dir,
+      output_dirs,
       FLAGS.problem,
       hparams.video_num_target_frames,
       frame_shape)
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index b132e1b83..7944c5df5 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -55,6 +55,7 @@ def decode_hparams(overrides=""):
       decode_to_file=None,
       shards=1,
       shard_id=0,
+      num_decodes=1,
       force_decode_length=False)
   hp.parse(overrides)
   return hp
@@ -74,7 +75,7 @@ def log_decode_results(inputs,
   """Log inference results."""
 
   # TODO(lukaszkaiser) refactor this into feature_encoder
-  is_video = "video" in problem_name
+  is_video = "video" in problem_name or "gym" in problem_name
   if is_video:
     def fix_and_save_video(vid, prefix):
       save_path_template = os.path.join(
@@ -133,7 +134,7 @@ def decode_from_dataset(estimator,
   # We assume that worker_id corresponds to shard number.
   shard = decode_hp.shard_id if decode_hp.shards > 1 else None
 
-  # Setup the decode output directory for any artifacts that may be written out
+  # Setup decode output directory for any artifacts that may be written out
   output_dir = os.path.join(estimator.model_dir, "decode")
   tf.gfile.MakeDirs(output_dir)
 
@@ -153,6 +154,39 @@ def decode_from_dataset(estimator,
   infer_input_fn = problem.make_estimator_input_fn(
       tf.estimator.ModeKeys.PREDICT, hparams, dataset_kwargs=dataset_kwargs)
 
+  output_dirs = []
+  for decode_id in range(decode_hp.num_decodes):
+    tf.logging.info("Decoding {}".format(decode_id))
+
+    output_dir = os.path.join(estimator.model_dir, "decode_%05d" % decode_id)
+    tf.gfile.MakeDirs(output_dir)
+    output_dirs.append(output_dir)
+
+    decode_once(estimator,
+                problem_name,
+                hparams,
+                infer_input_fn,
+                decode_hp,
+                decode_to_file,
+                output_dir)
+
+  run_postdecode_hooks(DecodeHookArgs(
+      estimator=estimator,
+      problem=problem,
+      output_dirs=output_dirs,
+      hparams=hparams,
+      decode_hparams=decode_hp))
+
+
+def decode_once(estimator,
+                problem_name,
+                hparams,
+                infer_input_fn,
+                decode_hp,
+                decode_to_file,
+                output_dir):
+  """Decodes once."""
+
   # Get the predictions as an iterable
   predictions = estimator.predict(infer_input_fn)
 
@@ -247,15 +281,6 @@ def decode_from_dataset(estimator,
     target_file.close()
     input_file.close()
 
-  run_postdecode_hooks(DecodeHookArgs(
-      estimator=estimator,
-      problem=problem,
-      output_dir=output_dir,
-      hparams=hparams,
-      decode_hparams=decode_hp))
-
-  tf.logging.info("Completed inference on %d samples." % num_predictions)  # pylint: disable=undefined-loop-variable
-
 
 def decode_from_file(estimator,
                      filename,
@@ -715,7 +740,7 @@ def latest_checkpoint_step(ckpt_dir):
 
 class DecodeHookArgs(collections.namedtuple(
     "DecodeHookArgs",
-    ["estimator", "problem", "output_dir", "hparams", "decode_hparams"])):
+    ["estimator", "problem", "output_dirs", "hparams", "decode_hparams"])):
   pass
 
 
@@ -730,7 +755,9 @@ def run_postdecode_hooks(decode_hook_args):
         "Skipping decode hooks because no checkpoint yet available.")
     return
   tf.logging.info("Running decode hooks.")
-  summary_writer = tf.summary.FileWriter(decode_hook_args.output_dir)
+  parent_dir = os.path.join(decode_hook_args.output_dirs[0], os.pardir)
+  final_dir = os.path.join(parent_dir, "decode")
+  summary_writer = tf.summary.FileWriter(final_dir)
   for hook in hooks:
     # Isolate each hook in case it creates TF ops
     with tf.Graph().as_default():
diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index a03ed2309..54a4265be 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -92,8 +92,9 @@ def compute_metrics(output_video, target_video):
   return {"PSNR": psnr, "SSIM": ssim}
 
 
-def compute_video_metrics(output_dir, problem_name, video_length, frame_shape):
-  """Computes the average of all the metric over the whole dataset.
+def compute_one_decoding_video_metrics(
+    output_dir, problem_name, video_length, frame_shape):
+  """Computes the average of all the metric for one decoding.
 
   This function assumes that all the predicted and target frames
   have been saved on the disk and sorting them by name will result
@@ -130,9 +131,39 @@ def compute_video_metrics(output_dir, problem_name, video_length, frame_shape):
     return results
 
 
-def compute_and_save_video_metrics(
-    output_dir, problem_name, video_length, frame_shape):
-  results = compute_video_metrics(
-      output_dir, problem_name, video_length, frame_shape)
-  save_results(results, output_dir, problem_name)
+def compute_all_metrics_statistics(all_results):
+  """Computes statistics of metrics across multiple decodings."""
+  statistics = {}
+  for key in all_results[0].keys():
+    values = [result[key] for result in all_results]
+    values = np.vstack(values)
+    statistics[key + "_MEAN"] = np.mean(values, axis=0)
+    statistics[key + "_STD"] = np.std(values, axis=0)
+    statistics[key + "_MIN"] = np.min(values, axis=0)
+    statistics[key + "_MAX"] = np.max(values, axis=0)
+  return statistics
+
 
+def compute_video_metrics(output_dirs, problem_name, video_length, frame_shape):
+  all_results = [
+      compute_one_decoding_video_metrics(
+          output_dir, problem_name, video_length, frame_shape)
+      for output_dir in output_dirs
+  ]
+  statistics = compute_all_metrics_statistics(all_results)
+  return statistics, all_results
+
+
+def compute_and_save_video_metrics(
+    output_dirs, problem_name, video_length, frame_shape):
+  """Compute and saves the video metrics."""
+  statistics, all_results = compute_video_metrics(
+      output_dirs, problem_name, video_length, frame_shape)
+  for results, output_dir in zip(all_results, output_dirs):
+    save_results(results, output_dir, problem_name)
+
+  parent_dir = os.path.join(output_dirs[0], os.pardir)
+  final_dir = os.path.join(parent_dir, "decode")
+  tf.gfile.MakeDirs(parent_dir)
+
+  save_results(statistics, final_dir, problem_name)

From de3aaab899703f3abc94ac428c75f8ee71c671bc Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Sat, 30 Jun 2018 00:20:13 -0700
Subject: [PATCH 0278/2720] Adding separated target losses in TB

PiperOrigin-RevId: 202758263
---
 tensor2tensor/utils/t2t_model.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index d3ba09e44..979d81ae6 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -434,6 +434,12 @@ def loss(self, logits, features):
       losses = {}
       for k, v in six.iteritems(logits):
         losses[k] = self._loss_single(v, target_modality[k], features[k])
+
+        n, d = losses[k]
+        tf.summary.scalar(k + "_loss", n / d)
+        tf.summary.scalar(k + "_loss_num", n)
+        tf.summary.scalar(k + "_loss_den", d)
+
       return tf.add_n([n / d for n, d in losses.values()])
     else:
       if self._problem_hparams:

From 7a8971f170f9db2c756baf31c6ccbe7e670b4569 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sat, 30 Jun 2018 23:19:37 -0700
Subject: [PATCH 0279/2720] restore hook and test

PiperOrigin-RevId: 202824653
---
 tensor2tensor/utils/restore_hook.py | 86 +++++++++++++++++++++++++++++
 1 file changed, 86 insertions(+)
 create mode 100644 tensor2tensor/utils/restore_hook.py

diff --git a/tensor2tensor/utils/restore_hook.py b/tensor2tensor/utils/restore_hook.py
new file mode 100644
index 000000000..13d0da9e2
--- /dev/null
+++ b/tensor2tensor/utils/restore_hook.py
@@ -0,0 +1,86 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Restore hooks."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import tarfile
+
+from tensor2tensor.data_generators import generator_utils
+import tensorflow as tf
+
+
+class RestoreHook(tf.train.SessionRunHook):
+  """Restore variables from a checkpoint path."""
+
+  def __init__(self, checkpoint_path, new_model_scope="", old_model_scope="",
+               include=None, exclude=None):
+    self._checkpoint_path = checkpoint_path
+    self._new_model_scope = new_model_scope
+    self._old_model_scope = old_model_scope
+    self._include = include
+    self._exclude = exclude
+
+  def begin(self):
+    """Load variables from checkpoint.
+
+    New model variables have the following name foramt:
+    new_model_scope/old_model_scope/xxx/xxx:0 To find the map of
+    name to variable, need to strip the new_model_scope and then
+    match the old_model_scope and remove the suffix :0.
+
+    """
+    variables_to_restore = tf.contrib.framework.get_variables_to_restore(
+        include=self._include, exclude=self._exclude)
+    # remove new_model_scope from variable name prefix
+    assignment_map = {variable.name[len(self._new_model_scope):]: variable
+                      for variable in variables_to_restore
+                      if variable.name.startswith(self._new_model_scope)}
+    # remove :0 from variable name suffix
+    assignment_map = {name.split(":")[0]: variable
+                      for name, variable in assignment_map.iteritems()
+                      if name.startswith(self._old_model_scope)}
+    self._assignment_map = assignment_map
+
+    tf.logging.info("restoring variables from checkpoint %s"%(
+        self._checkpoint_path))
+    tf.train.init_from_checkpoint(self._checkpoint_path, self._assignment_map)
+
+
+class RestoreResnetHook(RestoreHook):
+  """Restore Resnet models given scopes."""
+
+  _RESNET_URL = "http://download.tensorflow.org/models/{}_2017_04_14.tar.gz"
+
+  def __init__(self, new_model_scope="", include=None, exclude=None,
+               old_model_scope="resnet_v2_152/", model_dir="/tmp"):
+    model_name = old_model_scope[:-1]
+    checkpoint_path = self.get_model(model_name, model_dir)
+    super(RestoreResnetHook, self).__init__(
+        checkpoint_path, new_model_scope, old_model_scope, include, exclude)
+
+  def get_model(self, model_name, model_dir):
+    """Download the model given model name and extract it to a directory."""
+    resnet_url = self._RESNET_URL.format(model_name)
+    model_filename = "{}.tar.gz".format(model_name)
+    ckpt_filename = "{}.ckpt".format(model_name)
+
+    path = generator_utils.maybe_download(model_dir, model_filename, resnet_url)
+    with tarfile.open(path, "r:gz") as modeltar:
+      modeltar.extractall(model_dir)
+    return os.path.join(model_dir, ckpt_filename)

From acc697e99c706fe2c18adc86c8149fc2c38c64cc Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Mon, 2 Jul 2018 09:26:36 -0700
Subject: [PATCH 0280/2720] Fix MultiProblem.

PiperOrigin-RevId: 202962533
---
 tensor2tensor/data_generators/multi_problem.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index e1cfdad44..92ccf2a89 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -20,6 +20,7 @@
 
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.layers import discretization
 import tensorflow as tf
 
 
@@ -39,6 +40,12 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
   def add_task_id(self, task, example):
     """Convert example to code switching mode by adding a task id."""
+    if hasattr(task, "class_labels"):
+      # TODO(urvashik): handle the case where num_labels > 9
+      example["targets"] = tf.cast(discretization.int_to_bit(
+          example["targets"], 1, base=10) + 50, tf.int64)
+      example["targets"] = tf.squeeze(example["targets"], axis=[-1])
+
     if task.has_inputs:
       inputs = example.pop("inputs")
       concat_list = [inputs, [task.task_id], example["targets"]]
@@ -52,6 +59,14 @@ def filepattern(self, data_dir, mode, shard=None):
     print("Generating multi problem filepattern")
     return [task.filepattern(data_dir, mode, shard) for task in self.task_list]
 
+  def get_hparams(self, model_hparams=None):
+    if self._hparams is not None:
+      return self._hparams
+
+    self._hparams = self.task_list[0].get_hparams()
+
+    return self._hparams
+
   def dataset(self,
               mode,
               data_dir=None,
@@ -79,7 +94,7 @@ def dataset(self,
       task_dataset = task_dataset.map(lambda x: self.add_task_id(task, x))
       datasets.append(task_dataset)
 
-    self._hparams = self.task_list[0].get_hparams()
+    self.get_hparams()
 
     # TODO(urvashik): make this independent of the number of tasks
     def flatten_zip(d0, d1):

From bd05d62822ea17894f61ba525e29e43265f05712 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 2 Jul 2018 12:00:47 -0700
Subject: [PATCH 0281/2720] Adding a warning for when there is no frame_number
 but the flag is set.

PiperOrigin-RevId: 202989866
---
 tensor2tensor/data_generators/video_utils.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index e29be7887..a82059158 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -261,6 +261,9 @@ def check_integrity_and_batch(*datasets):
           if self.only_keep_videos_from_0th_frame:
             not_broken = tf.logical_and(not_broken,
                                         tf.equal(frame_numbers[0], 0))
+        else:
+          tf.logging.warning("use_not_breaking_batching is True but "
+                             "no frame_number is in the dataset.")
 
         features = {}
         for key in datasets[0].keys():

From 870996e07ca2889a4990a68a0c9c7d06eff99d3b Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 2 Jul 2018 12:04:36 -0700
Subject: [PATCH 0282/2720] fixing parameters of BAIR dataset.

PiperOrigin-RevId: 202990593
---
 .../data_generators/bair_robot_pushing.py     | 49 ++++++++++++++-----
 1 file changed, 37 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index cd4cf3002..784010542 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -67,8 +67,43 @@ def is_generate_per_split(self):
 
   @property
   def total_number_of_frames(self):
-    # TODO(mbz): correct this number to be the real total number of frames.
-    return 30 * 10 * 1000
+    return 1305600
+
+  @property
+  def random_skip(self):
+    return False
+
+  def eval_metrics(self):
+    return []
+
+  @property
+  def only_keep_videos_from_0th_frame(self):
+    return True
+
+  @property
+  def use_not_breaking_batching(self):
+    return True
+
+  @property
+  def extra_reading_spec(self):
+    """Additional data fields to store on disk and their decoders."""
+    data_fields = {
+        "frame_number": tf.FixedLenFeature([1], tf.int64),
+    }
+    decoders = {
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
+            tensor_key="frame_number"),
+    }
+    return data_fields, decoders
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    p.input_modality = {
+        "inputs": ("video", 256),
+    }
+    p.target_modality = {
+        "targets": ("video", 256),
+    }
 
   def parse_frames(self, filenames):
     image_key = "{}/image_aux1/encoded"
@@ -122,13 +157,3 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
           "state": state,
           "action": action,
       }
-
-  def hparams(self, defaults, unused_model_hparams):
-    p = defaults
-    p.input_modality = {
-        # Pixels are in 0..255 range.
-        "inputs": ("video", 256),
-    }
-    p.target_modality = {
-        "targets": ("video", 256),
-    }

From 786cc2586bd4eb6b5e9bcf4e1d8334db4fe5717d Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Mon, 2 Jul 2018 12:48:25 -0700
Subject: [PATCH 0283/2720] Return correct dtype from cumsum.

PiperOrigin-RevId: 202996883
---
 tensor2tensor/layers/common_layers.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 709920ddc..d76d18d46 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -294,8 +294,9 @@ def cumsum(x, axis=0, exclusive=False):
   length = x_shape[axis]
   my_range = tf.range(length)
   comparator = tf.less if exclusive else tf.less_equal
-  mask = tf.to_float(
-      comparator(tf.expand_dims(my_range, 1), tf.expand_dims(my_range, 0)))
+  mask = tf.cast(
+      comparator(tf.expand_dims(my_range, 1), tf.expand_dims(my_range, 0)),
+      x.dtype)
   ret = tf.tensordot(x, mask, axes=[[axis], [0]])
   if axis != rank - 1:
     ret = tf.transpose(

From 5d01bf67a1c063b7cf978919545e7cc71c0ca164 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 2 Jul 2018 13:50:11 -0700
Subject: [PATCH 0284/2720] restore hook test

PiperOrigin-RevId: 203006293
---
 tensor2tensor/utils/restore_hook.py | 32 +++--------------------------
 1 file changed, 3 insertions(+), 29 deletions(-)

diff --git a/tensor2tensor/utils/restore_hook.py b/tensor2tensor/utils/restore_hook.py
index 13d0da9e2..523c19a9d 100644
--- a/tensor2tensor/utils/restore_hook.py
+++ b/tensor2tensor/utils/restore_hook.py
@@ -18,17 +18,15 @@
 from __future__ import division
 from __future__ import print_function
 
-import os
-import tarfile
+import six
 
-from tensor2tensor.data_generators import generator_utils
 import tensorflow as tf
 
 
 class RestoreHook(tf.train.SessionRunHook):
   """Restore variables from a checkpoint path."""
 
-  def __init__(self, checkpoint_path, new_model_scope="", old_model_scope="",
+  def __init__(self, checkpoint_path="", new_model_scope="", old_model_scope="",
                include=None, exclude=None):
     self._checkpoint_path = checkpoint_path
     self._new_model_scope = new_model_scope
@@ -53,34 +51,10 @@ def begin(self):
                       if variable.name.startswith(self._new_model_scope)}
     # remove :0 from variable name suffix
     assignment_map = {name.split(":")[0]: variable
-                      for name, variable in assignment_map.iteritems()
+                      for name, variable in six.iteritems(assignment_map)
                       if name.startswith(self._old_model_scope)}
     self._assignment_map = assignment_map
 
     tf.logging.info("restoring variables from checkpoint %s"%(
         self._checkpoint_path))
     tf.train.init_from_checkpoint(self._checkpoint_path, self._assignment_map)
-
-
-class RestoreResnetHook(RestoreHook):
-  """Restore Resnet models given scopes."""
-
-  _RESNET_URL = "http://download.tensorflow.org/models/{}_2017_04_14.tar.gz"
-
-  def __init__(self, new_model_scope="", include=None, exclude=None,
-               old_model_scope="resnet_v2_152/", model_dir="/tmp"):
-    model_name = old_model_scope[:-1]
-    checkpoint_path = self.get_model(model_name, model_dir)
-    super(RestoreResnetHook, self).__init__(
-        checkpoint_path, new_model_scope, old_model_scope, include, exclude)
-
-  def get_model(self, model_name, model_dir):
-    """Download the model given model name and extract it to a directory."""
-    resnet_url = self._RESNET_URL.format(model_name)
-    model_filename = "{}.tar.gz".format(model_name)
-    ckpt_filename = "{}.ckpt".format(model_name)
-
-    path = generator_utils.maybe_download(model_dir, model_filename, resnet_url)
-    with tarfile.open(path, "r:gz") as modeltar:
-      modeltar.extractall(model_dir)
-    return os.path.join(model_dir, ckpt_filename)

From 91c34073981f3f1dd77164671b2108adaddd20ee Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Mon, 2 Jul 2018 13:51:01 -0700
Subject: [PATCH 0285/2720] set tpu computation_shape parameter to [1,1,1] in
 experimental

PiperOrigin-RevId: 203006418
---
 tensor2tensor/utils/trainer_lib.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index ee8a1962d..d38e2a348 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -159,7 +159,8 @@ def create_run_config(master="",
         iterations_per_loop=iterations_per_loop,
         num_shards=num_shards,
         per_host_input_for_training=True,
-        initial_infeed_sleep_secs=tpu_infeed_sleep_secs)
+        initial_infeed_sleep_secs=tpu_infeed_sleep_secs,
+        computation_shape=[1, 1, 1] if no_data_parallelism else None)
     run_config_args["tpu_config"] = tpu_config
 
   config = run_config_cls(**run_config_args)
@@ -533,4 +534,3 @@ def restore_checkpoint(ckpt_dir, saver, sess, must_restore=False):
   saver.restore(sess, path)
   step = int(path.split("-")[-1])
   return step
-

From 5cacecf8207aeef83d60e52b4b47291a973b414a Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Mon, 2 Jul 2018 14:13:38 -0700
Subject: [PATCH 0286/2720] small bug fixes for rel attention

PiperOrigin-RevId: 203010511
---
 tensor2tensor/models/image_transformer.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 91f5190ab..6abd0d626 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -912,8 +912,9 @@ def imagetransformer_b12l_4h_b256_uncond_dr03_tpu():
 @registry.register_hparams
 def imagetransformer_b12l_4h_b256_uncond_dr03_rel_tpu():
   """works very well on 4x4."""
-  hparams = imagetransformer_bas8l_8h_big_uncond_dr03_imgnet()
-  hparams.dec_attention_type = cia.RELATIVE_LOCAL_1D
+  hparams = imagetransformer_b12l_4h_b256_uncond_dr03_tpu()
+  hparams.shared_rel = True
+  hparams.dec_attention_type = cia.AttentionType.RELATIVE_LOCAL_1D
   return hparams
 
 
@@ -926,7 +927,7 @@ def imagetransformer_cifar_tpu_range(rhp):
   rhp.set_discrete("hidden_size", [256, 512, 1024])
   rhp.set_discrete("block_length", [128, 256, 512])
   rhp.set_categorical("dec_attention_type", [
-      cia.RELATIVE_LOCAL_1D, cia.LOCAL_1D])
+      cia.AttentionType.RELATIVE_LOCAL_1D, cia.AttentionType.LOCAL_1D])
 
 
 @registry.register_hparams

From ca5898adaf8eeb66a9aa2874b98de4c58392c44a Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 2 Jul 2018 14:33:38 -0700
Subject: [PATCH 0287/2720] Combine Imagetransformer{,Plus} as one model under
 one hparam.

PiperOrigin-RevId: 203013801
---
 .../layers/common_image_attention.py          |  97 ++++++++++++----
 .../layers/common_image_attention_test.py     | 108 ++++++++++++++++++
 tensor2tensor/models/image_transformer.py     | 100 +++++++++-------
 .../models/image_transformer_test.py          |  26 +++--
 4 files changed, 259 insertions(+), 72 deletions(-)
 create mode 100644 tensor2tensor/layers/common_image_attention_test.py

diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 22f8bbad1..9454e2c13 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -50,6 +50,19 @@ def get_choices():
     ]
 
 
+class DistributionType(object):
+  """Types of distributions used in cia."""
+  CAT = "cat"
+  DMOL = "dmol"
+
+  @staticmethod
+  def get_choices():
+    return [
+        DistributionType.CAT,
+        DistributionType.DMOL,
+    ]
+
+
 def maybe_reshape_4d_to_3d(x):
   """Reshape input from 4D to 3D if necessary."""
   x_shape = common_layers.shape_list(x)
@@ -513,34 +526,58 @@ def transformer_layers_sharded(dp,
 
 
 def postprocess_image(x, rows, cols, hparams):
-  """Postprocessing after decoding."""
+  """Postprocessing after decoding.
+
+  Args:
+    x: Tensor of shape [batch, ...], where ... can be any rank such that the
+      number of elements in x is batch * rows * cols * hparams.hidden_size.
+    rows: Integer representing number of rows in a 2-D data point.
+    cols: Integer representing number of columns in a 2-D data point.
+    hparams: tf.contrib.training.HParams set.
+
+  Returns:
+    Tensor of shape [batch, rows, cols, depth], where depth is
+    hparams.num_mixtures * 10 if hparams.likelihood is DMOL, otherwise 256. In
+    the special case of inference and block raster scan order, it is a Tensor
+    of shape [batch, num_blocks_rows, num_block_cols, block_length, block_width,
+    depth].
+  """
   batch = common_layers.shape_list(x)[0]
-  channels = 256
   x = tf.reshape(x, [batch, rows, cols, hparams.hidden_size])
-  targets = tf.layers.dense(x, 256, use_bias=True, activation=None,
-                            name="output_conv")
+  likelihood = getattr(hparams, "likelihood", DistributionType.CAT)
+  if likelihood == DistributionType.DMOL:
+    depth = hparams.num_mixtures * 10
+    targets = tf.layers.dense(x,
+                              depth,
+                              use_bias=False,
+                              activation=None,
+                              name="output_conv")
+  else:
+    depth = 256
+    targets = tf.layers.dense(x,
+                              depth,
+                              use_bias=True,
+                              activation=None,
+                              name="output_conv")
   if (hparams.mode == tf.contrib.learn.ModeKeys.INFER and
       hparams.block_raster_scan):
     y = targets
-    y = tf.reshape(y, [batch, -1, hparams.img_len*3, channels])
     yshape = common_layers.shape_list(y)
     block_length = hparams.query_shape[0]
     block_width = hparams.query_shape[1]
 
     # Break into block row wise.
     y = tf.reshape(y,
-                   [batch, yshape[1] // block_length,
-                    block_length,
-                    yshape[2], channels])
+                   [batch, yshape[1] // block_length, block_length,
+                    yshape[2], depth])
     yshape = common_layers.shape_list(y)
     # Break into blocks width wise.
     y_blocks = tf.reshape(y,
                           [batch, yshape[1], yshape[2],
-                           yshape[3] // block_width,
-                           block_width, channels])
+                           yshape[3] // block_width, block_width, depth])
 
-    # Reshape targets as [batch_size, num_blocks_rows, num_block_cols,
-    # block_length, block_width, channels]
+    # Reshape targets as [batch, num_blocks_rows, num_block_cols, block_length,
+    # block_width, depth].
     targets = tf.transpose(y_blocks, [0, 1, 3, 2, 4, 5])
 
   return targets
@@ -641,18 +678,36 @@ def prepare_image(inputs, hparams, name=None):
 
 
 def create_output(decoder_output, rows, cols, targets, hparams):
-  """Create output from decoder output and vars."""
+  """Creates output from decoder output and vars.
+
+  Args:
+    decoder_output: Tensor of shape [batch, ...], where ... can be any rank such
+      that the number of elements is batch * rows * cols * hparams.hidden_size.
+    rows: Integer representing number of rows in a 2-D data point.
+    cols: Integer representing number of columns in a 2-D data point.
+    targets: Tensor of shape [batch, hparams.img_len, hparams.img_len,
+      hparams.num_channels].
+    hparams: tf.contrib.training.HParams set.
+
+  Returns:
+    Tensor of shape [batch, hparams.img_len, hparams.img_len,
+    hparams.num_mixtures * 10] if hparams.likelihood is DMOL, otherwise
+    [batch, hparams.img_len, hparams.img_len, hparams.num_channels, 256].
+    In the special case of predict mode, it is a Tensor of rank 5.
+  """
   decoded_image = postprocess_image(decoder_output, rows, cols, hparams)
-  targets_shape = common_layers.shape_list(targets)
+  depth = common_layers.shape_list(decoded_image)[-1]
+  batch, height, width, channels = common_layers.shape_list(targets)
+  likelihood = getattr(hparams, "likelihood", DistributionType.CAT)
   if hparams.mode == tf.estimator.ModeKeys.PREDICT:
-    # Hardcoding that the number of intensity values is 256.
-    y = tf.reshape(decoded_image, [targets_shape[0], -1, 1, 1, 256])
-    output = y[:, :targets_shape[1], :, :, :]
+    y = tf.reshape(decoded_image, [batch, -1, 1, 1, depth])
+    output = y[:, :height, :, :, :]
+  elif likelihood == DistributionType.CAT:
+    # Unpack the cols dimension of the Categorical.
+    output = tf.reshape(decoded_image,
+                        [batch, height, width, channels, depth])
   else:
-    output = tf.reshape(decoded_image, [
-        targets_shape[0], targets_shape[1], targets_shape[2],
-        targets_shape[3], 256
-    ])
+    output = decoded_image
   return output
 
 
diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
new file mode 100644
index 000000000..6d57e7413
--- /dev/null
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -0,0 +1,108 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 common image attention utilities."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+from tensor2tensor.layers import common_image_attention
+
+import tensorflow as tf
+
+
+class CommonImageAttentionTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(
+      (common_image_attention.DistributionType.DMOL, 5, 50),
+      (common_image_attention.DistributionType.CAT, None, 256),
+  )
+  def testPostProcessImageTrainMode(self, likelihood, num_mixtures, depth):
+    batch = 1
+    rows = 8
+    cols = 24
+    hparams = tf.contrib.training.HParams(
+        hidden_size=2,
+        likelihood=likelihood,
+        mode=tf.estimator.ModeKeys.TRAIN,
+        num_mixtures=num_mixtures,
+    )
+    inputs = tf.random_uniform([batch, rows, cols, hparams.hidden_size],
+                               minval=-1., maxval=1.)
+    outputs = common_image_attention.postprocess_image(
+        inputs, rows, cols, hparams)
+    self.assertEqual(outputs.shape, (batch, rows, cols, depth))
+
+  @parameterized.parameters(
+      (common_image_attention.DistributionType.DMOL, 5, 50),
+      (common_image_attention.DistributionType.CAT, None, 256),
+  )
+  def testPostProcessImageInferMode(self, likelihood, num_mixtures, depth):
+    batch = 1
+    rows = 8
+    cols = 24
+    block_length = 4
+    block_width = 2
+    hparams = tf.contrib.training.HParams(
+        block_raster_scan=True,
+        hidden_size=2,
+        likelihood=likelihood,
+        mode=tf.contrib.learn.ModeKeys.INFER,
+        num_mixtures=num_mixtures,
+        query_shape=[block_length, block_width],
+    )
+    inputs = tf.random_uniform([batch, rows, cols, hparams.hidden_size],
+                               minval=-1., maxval=1.)
+    outputs = common_image_attention.postprocess_image(
+        inputs, rows, cols, hparams)
+    num_blocks_rows = rows // block_length
+    num_blocks_cols = cols // block_width
+    self.assertEqual(outputs.shape,
+                     (batch, num_blocks_rows, num_blocks_cols,
+                      block_length, block_width, depth))
+
+  @parameterized.parameters(
+      (common_image_attention.DistributionType.DMOL, 5, 50),
+      (common_image_attention.DistributionType.CAT, None, 256),
+  )
+  def testCreateOutputTrainMode(self, likelihood, num_mixtures, depth):
+    batch = 1
+    height = 8
+    width = 8
+    channels = 3
+    rows = height
+    if likelihood == common_image_attention.DistributionType.CAT:
+      cols = channels * width
+    else:
+      cols = width
+    hparams = tf.contrib.training.HParams(
+        hidden_size=2,
+        likelihood=likelihood,
+        mode=tf.estimator.ModeKeys.TRAIN,
+        num_mixtures=num_mixtures,
+    )
+    decoder_output = tf.random_normal([batch, rows, cols, hparams.hidden_size])
+    targets = tf.random_uniform([batch, height, width, channels],
+                                minval=-1., maxval=1.)
+    output = common_image_attention.create_output(
+        decoder_output, rows, cols, targets, hparams)
+    if hparams.likelihood == common_image_attention.DistributionType.CAT:
+      self.assertEqual(output.shape, (batch, height, width, channels, depth))
+    else:
+      self.assertEqual(output.shape, (batch, height, width, depth))
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 6abd0d626..e59431d32 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -35,14 +35,26 @@
 
 @registry.register_model
 class Imagetransformer(t2t_model.T2TModel):
-  """Conditional image generation with attention. See file docstring."""
+  """Conditional image generation with attention. See file docstring.
+
+  The model admits either a Categorical or discretized mixture of logistic
+  distributions (DMOL) as the likelihood. When using DMOL for training, double
+  check that the evaluation metrics also use it.
+  """
 
   def body(self, features):
     hparams = copy.copy(self._hparams)
     inputs = features["inputs"]
     targets = features["targets"]
-    if not (tf.get_variable_scope().reuse or
-            hparams.mode == tf.contrib.learn.ModeKeys.INFER):
+    if (hparams.likelihood == cia.DistributionType.DMOL and
+        (hparams.target_modality != "image:image_channel_bottom_identity" or
+         hparams.num_channels != 1)):
+      raise ValueError("When using DMOL for the likelihood, target_modality "
+                       "must be image:image_channel_bottom_identity and "
+                       "num_channels must be 1.")
+    if (not tf.get_variable_scope().reuse and
+        hparams.mode != tf.contrib.learn.ModeKeys.INFER and
+        hparams.target_modality != "image:image_channel_bottom_identity"):
       tf.summary.image("targets", tf.to_float(targets), max_outputs=1)
 
     # Extra losses list if we want to use moe.
@@ -69,41 +81,11 @@ def body(self, features):
     else:
       return output
 
-
-@registry.register_model
-class ImagetransformerPlus(t2t_model.T2TModel):
-  """Imagetransformer with discretized mixture of logistics loss."""
-
-  def body(self, features):
-    hparams = copy.copy(self._hparams)
-    inputs = features["inputs"]
-    targets = features["targets"]
-    # Prepare decoder inputs and bias.
-    decoder_input, _, _ = cia.prepare_decoder(targets, hparams)
-    # Add class label to decoder input.
-    if not hparams.unconditional:
-      decoder_input += tf.reshape(
-          inputs,
-          [common_layers.shape_list(targets)[0], 1, 1, hparams.hidden_size])
-    decoder_output = cia.transformer_decoder_layers(
-        decoder_input,
-        None,
-        hparams.num_decoder_layers or hparams.num_hidden_layers,
-        hparams,
-        attention_type=hparams.dec_attention_type,
-        name="decoder")
-    # reshape it into [batch, height, width, depth]
-    decoder_output = tf.reshape(decoder_output, tf.shape(targets))
-    # there are 10 sets of parameters that you need to produce, location, scale,
-    # and coefficient parameter for each
-    output = tf.layers.dense(decoder_output, hparams.num_mixtures*10,
-                             use_bias=False, activation=None,
-                             name="output_mixtures_conv")
-    # TODO(avaswani) Figure out if we need residuals or layer norm
-    return output
-
   def loss(self, logits, features):
-    return common_layers.dml_loss(logits, features["targets"])
+    if self._hparams.likelihood == cia.DistributionType.DMOL:
+      return common_layers.dml_loss(logits, features["targets"])
+
+    return super(Imagetransformer, self).loss(logits, features)
 
   def sample(self, features):
     """Run the model and extract samples.
@@ -116,11 +98,32 @@ def sample(self, features):
        logits: a list of `Tensor`s, one per datashard.
        losses: a dictionary: {loss-name (string): floating point `Scalar`}.
     """
-    logits, losses = self(features)  # pylint: disable=not-callable
+    if self._hparams.likelihood == cia.DistributionType.DMOL:
+      logits, losses = self(features)  # pylint: disable=not-callable
+      samples = common_layers.sample_from_discretized_mix_logistic(
+          logits, seed=None)
+      return samples, logits, losses
+
+    return super(Imagetransformer, self).sample(features)
 
-    samples = common_layers.sample_from_discretized_mix_logistic(
-        logits, seed=None)
-    return samples, logits, losses
+  def _slow_greedy_infer(self, features, decode_length):
+    """A slow greedy inference method.
+
+    Quadratic time in decode_length.
+
+    Args:
+      features: an map of string to `Tensor`
+      decode_length: an integer.  How many additional timesteps to decode.
+
+    Returns:
+       samples: an integer `Tensor`.
+       logits: `Tensor` of shape [batch_size, time, 1, 1, vocab_size].
+       losses: a dictionary: {loss-name (string): floating point `Scalar`}
+    """
+    if self._hparams.likelihood == cia.DistributionType.DMOL:
+      raise NotImplementedError("Decoding is not currently available for DMOL.")
+    return super(Imagetransformer, self)._slow_greedy_infer(features,
+                                                            decode_length)
 
 
 @registry.register_model
@@ -226,6 +229,7 @@ def image_transformer_base():
   hparams.add_hparam("q_filter_width", 1)
   hparams.add_hparam("kv_filter_width", 1)
 
+  hparams.add_hparam("likelihood", cia.DistributionType.CAT)
   hparams.add_hparam("unconditional", False)  # unconditional generation
 
   # parameters of discretized mixture of logistics loss from pixel cnn++
@@ -323,6 +327,9 @@ def imagetransformer_base_10l_8h_big_uncond_dr03_dan_64():
 def imagetransformerpp_sep_channels_8l_8h():
   """separate rgb embeddings."""
   hparams = imagetransformer_base()
+  hparams.likelihood = cia.DistributionType.DMOL
+  hparams.num_channels = 1
+  hparams.target_modality = "image:image_channel_bottom_identity"
   hparams.num_heads = 8
   hparams.batch_size = 4
   hparams.attention_key_channels = hparams.attention_value_channels = 0
@@ -332,7 +339,6 @@ def imagetransformerpp_sep_channels_8l_8h():
   hparams.sampling_method = "random"
   hparams.layer_preprocess_sequence = "n"
   hparams.layer_postprocess_sequence = "da"
-  hparams.target_modality = "image:image_channel_bottom_identity"
   hparams.summarize_grads = True
   hparams.learning_rate = 0.1
   return hparams
@@ -359,7 +365,6 @@ def imagetransformerpp_base_8l_8h_big_cond_dr03_dan():
 def imagetransformerpp_base_8l_8h_big_cond_dr03_dan_a():
   hparams = imagetransformerpp_base_8l_8h_big_cond_dr03_dan()
   hparams.learning_rate = 0.1
-  hparams.num_channels = 1
   return hparams
 
 
@@ -742,6 +747,15 @@ def imagetransformer_tiny():
   return hparams
 
 
+@registry.register_hparams
+def imagetransformerpp_tiny():
+  hparams = imagetransformer_base()
+  hparams.likelihood = cia.DistributionType.DMOL
+  hparams.num_channels = 1
+  hparams.target_modality = "image:image_channel_bottom_identity"
+  return hparams
+
+
 @registry.register_hparams
 def imagetransformer_tiny_tpu():
   hparams = imagetransformer_tiny()
diff --git a/tensor2tensor/models/image_transformer_test.py b/tensor2tensor/models/image_transformer_test.py
index ad0f45f58..8f4911f0d 100644
--- a/tensor2tensor/models/image_transformer_test.py
+++ b/tensor2tensor/models/image_transformer_test.py
@@ -17,21 +17,31 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+from absl.testing import parameterized
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.layers import common_image_attention
 from tensor2tensor.models import image_transformer
 
 import tensorflow as tf
 
 
-class ImagetransformerTest(tf.test.TestCase):
+class ImagetransformerTest(parameterized.TestCase, tf.test.TestCase):
 
-  def _test_imagetransformer(self, net):
+  @parameterized.named_parameters(
+      ("ImageTransformerCat",
+       image_transformer.Imagetransformer,
+       image_transformer.imagetransformer_tiny()),
+      ("ImageTransformerDmol",
+       image_transformer.Imagetransformer,
+       image_transformer.imagetransformerpp_tiny()),
+  )
+  def testImagetransformer(self, net, hparams):
     batch_size = 3
     size = 7
     vocab_size = 256
-    hparams = image_transformer.imagetransformer_base()
     p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
     inputs = -1 + np.random.random_integers(
         vocab_size, size=(batch_size, 1, 1, 1))
@@ -47,11 +57,11 @@ def _test_imagetransformer(self, net):
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
-    self.assertEqual(res.shape, (batch_size, size, size, 3, vocab_size))
-
-  def testImagetransformer(self):
-    self._test_imagetransformer(image_transformer.Imagetransformer)
-
+    if hparams.likelihood == common_image_attention.DistributionType.CAT:
+      expected = (batch_size, size, size, 3, vocab_size)
+    else:
+      expected = (batch_size, size, size, hparams.num_mixtures * 10)
+    self.assertEqual(res.shape, expected)
 
 if __name__ == "__main__":
   tf.test.main()

From 72a4876ecf1c3763bb5a482d73c7c73efb356e11 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Mon, 2 Jul 2018 15:12:49 -0700
Subject: [PATCH 0288/2720] support different hparams.activation_dtype in a few
 more places.

PiperOrigin-RevId: 203020676
---
 tensor2tensor/layers/common_attention.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index ae0193bde..a654206e7 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -658,7 +658,7 @@ def add_positional_embedding(x, max_length, name, positions=None):
     a Tensor the same shape as x.
   """
   _, length, depth = common_layers.shape_list(x)
-  var = tf.get_variable(name, [max_length, depth])
+  var = tf.cast(tf.get_variable(name, [max_length, depth]), x.dtype)
   if positions is None:
     sliced = tf.cond(
         tf.less(length, max_length),
@@ -2960,6 +2960,7 @@ def compute_attention_component(antecedent,
                vars_3d_num_heads,
                total_depth // vars_3d_num_heads],
         initializer=tf.random_normal_initializer(stddev=initializer_stddev))
+    var = tf.cast(var, antecedent.dtype)
     var = tf.reshape(var, [input_depth, total_depth])
     return tf.tensordot(antecedent, var, axes=1)
   if filter_width == 1:
@@ -3250,6 +3251,7 @@ def multihead_attention(query_antecedent,
     if vars_3d:
       o_var = tf.get_variable(
           "o", [num_heads, total_value_depth // num_heads, output_depth])
+      o_var = tf.cast(o_var, x.dtype)
       o_var = tf.reshape(o_var, [total_value_depth, output_depth])
       x = tf.tensordot(x, o_var, axes=1)
     else:

From 93f24d4c241c8ff0bbcb6807361a9161de3289f9 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Mon, 2 Jul 2018 15:29:03 -0700
Subject: [PATCH 0289/2720] respect hparams.use_fixed_batch_size in

PiperOrigin-RevId: 203023248
---
 tensor2tensor/data_generators/problem.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 7bce7e1c6..9bf267426 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -257,7 +257,7 @@ def tpu_batch_size_per_shard(self, model_hparams):
     Returns:
       an integer
     """
-    if self.batch_size_means_tokens:
+    if self.batch_size_means_tokens and not model_hparams.use_fixed_batch_size:
       return model_hparams.batch_size // self.max_length(model_hparams)
     else:
       return model_hparams.batch_size

From 7bb67a18e1e4a0cddd1d61c65c937f14c1c124e3 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Mon, 2 Jul 2018 22:20:24 -0700
Subject: [PATCH 0290/2720] fix test failure

PiperOrigin-RevId: 203066059
---
 tensor2tensor/bin/t2t_trainer.py   | 5 ++++-
 tensor2tensor/utils/trainer_lib.py | 5 +++--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 8e3737d62..985ab0075 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -190,6 +190,8 @@ def create_run_config(hp):
   if save_ckpt_secs:
     save_ckpt_steps = None
   assert FLAGS.output_dir or FLAGS.checkpoint_path
+  computation_shape = None
+
   # the various custom getters we have written do not play well together yet.
   # TODO(noam): ask rsepassi for help here.
   daisy_chain_variables = (
@@ -226,7 +228,8 @@ def create_run_config(hp):
       tpu_infeed_sleep_secs=FLAGS.tpu_infeed_sleep_secs,
       inter_op_parallelism_threads=FLAGS.inter_op_parallelism_threads,
       log_step_count_steps=FLAGS.log_step_count_steps,
-      intra_op_parallelism_threads=FLAGS.intra_op_parallelism_threads)
+      intra_op_parallelism_threads=FLAGS.intra_op_parallelism_threads,
+      computation_shape=computation_shape)
 
 
 def generate_data():
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index d38e2a348..7288b6026 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -126,7 +126,8 @@ def create_run_config(master="",
                       use_tpu=False,
                       inter_op_parallelism_threads=0,
                       log_step_count_steps=100,
-                      intra_op_parallelism_threads=0):
+                      intra_op_parallelism_threads=0,
+                      computation_shape=None):
   """Create RunConfig, TPUConfig, and Parallelism object."""
   session_config = create_session_config(
       log_device_placement=log_device_placement,
@@ -160,7 +161,7 @@ def create_run_config(master="",
         num_shards=num_shards,
         per_host_input_for_training=True,
         initial_infeed_sleep_secs=tpu_infeed_sleep_secs,
-        computation_shape=[1, 1, 1] if no_data_parallelism else None)
+        computation_shape=computation_shape)
     run_config_args["tpu_config"] = tpu_config
 
   config = run_config_cls(**run_config_args)

From 3b5d95c6e2327017e06fe721d2810579f7917320 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Tue, 3 Jul 2018 11:32:03 -0700
Subject: [PATCH 0291/2720] Allow passing of extra kwargs to TPUConfig
 initializer.

PiperOrigin-RevId: 203156633
---
 tensor2tensor/bin/t2t_trainer.py   | 4 ++--
 tensor2tensor/utils/trainer_lib.py | 8 ++++++--
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 985ab0075..0958dba3b 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -190,7 +190,7 @@ def create_run_config(hp):
   if save_ckpt_secs:
     save_ckpt_steps = None
   assert FLAGS.output_dir or FLAGS.checkpoint_path
-  computation_shape = None
+  tpu_config_extra_kwargs = {}
 
   # the various custom getters we have written do not play well together yet.
   # TODO(noam): ask rsepassi for help here.
@@ -229,7 +229,7 @@ def create_run_config(hp):
       inter_op_parallelism_threads=FLAGS.inter_op_parallelism_threads,
       log_step_count_steps=FLAGS.log_step_count_steps,
       intra_op_parallelism_threads=FLAGS.intra_op_parallelism_threads,
-      computation_shape=computation_shape)
+      tpu_config_extra_kwargs=tpu_config_extra_kwargs)
 
 
 def generate_data():
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 7288b6026..30b00df68 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -127,7 +127,7 @@ def create_run_config(master="",
                       inter_op_parallelism_threads=0,
                       log_step_count_steps=100,
                       intra_op_parallelism_threads=0,
-                      computation_shape=None):
+                      tpu_config_extra_kwargs=None):
   """Create RunConfig, TPUConfig, and Parallelism object."""
   session_config = create_session_config(
       log_device_placement=log_device_placement,
@@ -155,13 +155,15 @@ def create_run_config(master="",
 
   # If using TPU, use TPU RunConfig, add TPUConfig, and add additional args
   if use_tpu:
+    if tpu_config_extra_kwargs is None:
+      tpu_config_extra_kwargs = {}
     run_config_cls = tf.contrib.tpu.RunConfig
     tpu_config = tf.contrib.tpu.TPUConfig(
         iterations_per_loop=iterations_per_loop,
         num_shards=num_shards,
         per_host_input_for_training=True,
         initial_infeed_sleep_secs=tpu_infeed_sleep_secs,
-        computation_shape=computation_shape)
+        **tpu_config_extra_kwargs)
     run_config_args["tpu_config"] = tpu_config
 
   config = run_config_cls(**run_config_args)
@@ -205,6 +207,8 @@ def create_estimator(model_name,
     batch_size = (
         problem.tpu_batch_size_per_shard(hparams) *
         run_config.tpu_config.num_shards)
+    if getattr(hparams, "mtf_mode", False):
+      batch_size = problem.tpu_batch_size_per_shard(hparams)
     predict_batch_size = batch_size
     if decode_hparams and decode_hparams.batch_size:
       predict_batch_size = decode_hparams.batch_size

From 84bd6b23ba87e206c9ef3020d2e5e1f195807750 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 3 Jul 2018 12:03:35 -0700
Subject: [PATCH 0292/2720] add restore_hook to vqa_attention_baseline

PiperOrigin-RevId: 203161768
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index a654206e7..d1a6d6ee0 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -849,7 +849,7 @@ def padding_to_length(padding):
     a Tensor with shape [...].
   """
   non_padding = 1.0 - padding
-  return tf.to_int64(tf.reduce_sum(non_padding, axis=-1))
+  return tf.to_int32(tf.reduce_sum(non_padding, axis=-1))
 
 
 @expert_utils.add_name_scope()

From 342e214dea360a7f472fc82f3dd0775d7e224c52 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Tue, 3 Jul 2018 14:31:42 -0700
Subject: [PATCH 0293/2720] Clean up docstrings and inline comments in
 latent_layers.

PiperOrigin-RevId: 203184230
---
 tensor2tensor/layers/latent_layers.py | 181 +++++++++++++++++---------
 1 file changed, 117 insertions(+), 64 deletions(-)

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 900734473..bc13c7f76 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -61,7 +61,17 @@ def attend(x, source, hparams, name):
 
 
 def multinomial_sample(x, vocab_size, sampling_method, temperature):
-  """Multinomial sampling from a n-dimensional tensor."""
+  """Multinomial sampling from a n-dimensional tensor.
+
+  Args:
+    x: Tensor of shape [..., vocab_size]. Parameterizes logits of multinomial.
+    vocab_size: Number of classes in multinomial distribution.
+    sampling_method: String, "random" or otherwise deterministic.
+    temperature: Positive float.
+
+  Returns:
+    Tensor of shape [...].
+  """
   if sampling_method == "random":
     samples = tf.multinomial(tf.reshape(x, [-1, vocab_size]) / temperature, 1)
   else:
@@ -71,7 +81,17 @@ def multinomial_sample(x, vocab_size, sampling_method, temperature):
 
 
 def ae_latent_softmax(latents_pred, latents_discrete_hot, hparams):
-  """Latent prediction and loss."""
+  """Latent prediction and loss.
+
+  Args:
+    latents_pred: Tensor of shape [..., depth].
+    latents_discrete_hot: Tensor of shape [..., vocab_size].
+    hparams: tf.contrib.training.HParams.
+
+  Returns:
+    sample: Tensor of shape [...], a sample from a multinomial distribution.
+    loss: Tensor of shape [...], the softmax cross-entropy.
+  """
   vocab_size = 2**hparams.bottleneck_bits
   with tf.variable_scope("latent_logits"):
     latents_logits = tf.layers.dense(latents_pred, vocab_size,
@@ -88,7 +108,22 @@ def ae_latent_softmax(latents_pred, latents_discrete_hot, hparams):
 
 
 def ae_latent_sample_beam(latents_dense_in, inputs, ed, embed, hparams):
-  """Sample from the latent space in the autoencoder."""
+  """Samples from the latent space in the autoencoder.
+
+  Args:
+    latents_dense_in: Tensor of shape [batch, length_q, ...]. Only the shape of
+      its first two dimensions are used.
+    inputs: Tensor of shape [batch, length_kv, hparams.hidden_size]. Encodings
+      to attend to in decoder.
+    ed: Tensor which broadcasts with shape [batch, hparams.num_heads, length_q,
+      length_kv]. Encoder-decoder attention bias.
+    embed: Callable which embeds discrete latent hot-vectors and a hidden size
+      and returns dense vectors.
+    hparams: tf.contrib.training.HParams.
+
+  Returns:
+    Tensor of shape [batch, length].
+  """
 
   def symbols_to_logits_fn(ids):
     """Go from ids to logits."""
@@ -132,11 +167,11 @@ def residual_block_layer(inputs, hparams):
     dropout, add and normalize according to hparams.layer_postprocess_sequence.
 
   Args:
-    inputs: Tensor of shape [batch_size, height, width, hidden_dim].
-    hparams: Dict, hyperparameters.
+    inputs: Tensor of shape [batch, height, width, hparams.hidden_size].
+    hparams: tf.contrib.training.HParams.
 
   Returns:
-    x: Tensor of shape [batch_size, height, width, hidden_dim]
+    Tensor of shape [batch, height, width, hparams.hidden_size].
   """
   kernel = (hparams.res_kernel_size, hparams.res_kernel_size)
   x = inputs
@@ -160,25 +195,26 @@ def residual_block_layer(inputs, hparams):
   return x
 
 
-def compress_encoder(inputs, hparams,
+def compress_encoder(inputs,
+                     hparams,
                      strides=(2, 2),
                      kernel=(3, 3),
                      name="compress"):
-  """Encoder that compresses inputs to length/2**num_compress_steps.
+  """Encoder that compresses 2-D inputs by 2**num_compress_steps.
 
   Args:
-    inputs: Tensor of shape [batch, height, width, hidden_dim].
-    hparams: Dict, hyperparameters.
+    inputs: Tensor of shape [batch, height, width, channels].
+    hparams: tf.contrib.training.HParams.
     strides: Tuple, strides for conv block.
     kernel: Tuple, kernel window size for conv block.
     name: string, variable scope.
 
   Returns:
-    x: Tensor of shape [batch, height*width/2**(compress_steps), hidden_dim].
+    Tensor of shape [batch, (height*width) / 2**(hparams.num_compress_steps),
+    hparams.hidden_size].
   """
   with tf.variable_scope(name):
     x = inputs
-    # Compress conv layers with strides and kernels as passed to the function.
     for i in range(hparams.num_compress_steps // 2):
       with tf.variable_scope("compress_conv_%d" % i):
         y = common_layers.conv_block(
@@ -191,7 +227,6 @@ def compress_encoder(inputs, hparams,
         y = tf.nn.dropout(y, 1.0 - hparams.dropout)
         x = y
 
-    # Residual blocks.
     x = residual_block_layer(x, hparams)
 
     # If using multiple copies of latents, blow up the hidden size and then
@@ -205,15 +240,16 @@ def compress_encoder(inputs, hparams,
 
 
 def compress_encoder_2d(x, hparams, name):
-  """Encoder that compresses inputs to height*width/2**num_compress_steps.
+  """Encoder that compresses 2-D inputs by 2**num_compress_steps.
 
   Args:
-    x: Tensor of shape [batch, height, width, hidden_dim].
-    hparams: Dict, hyperparameters.
+    x: Tensor of shape [batch, height, width, channels].
+    hparams: tf.contrib.training.HParams.
     name: string, variable scope.
 
   Returns:
-    x: Tensor of shape [batch, height*width/2**(compress_steps), hidden_dim].
+    Tensor of shape [batch, (height*width) / 2**hparams.num_compress_steps,
+    hparams.hidden_size].
   """
   return compress_encoder(x, hparams,
                           strides=(2, 2),
@@ -222,15 +258,16 @@ def compress_encoder_2d(x, hparams, name):
 
 
 def compress_encoder_1d(x, hparams, name):
-  """Encoder that compresses inputs to length/2**num_compress_steps.
+  """Encoder that compresses 1-D inputs by 2**num_compress_steps.
 
   Args:
-    x: Tensor of shape [batch, length, hidden_dim].
-    hparams: Dict, hyperparameters.
+    x: Tensor of shape [batch, length, channels].
+    hparams: tf.contrib.training.HParams.
     name: string, variable scope.
 
   Returns:
-    x: Tensor of shape [batch, length/2**(compress_steps), hidden_dim].
+    Tensor of shape [batch, length / 2**hparams.num_compress_steps,
+    hparams.hidden_size].
   """
   x = tf.expand_dims(x, axis=2)
   return compress_encoder(x, hparams,
@@ -239,30 +276,27 @@ def compress_encoder_1d(x, hparams, name):
                           name=name)
 
 
-def decompress_decoder(inputs, hparams,
+def decompress_decoder(inputs,
+                       hparams,
                        strides=(2, 2),
                        kernel=(3, 3),
                        name="decompress"):
-  """Encoder that compresses inputs to length/2**num_compress_steps.
+  """Decoder that decompresses 2-D inputs by 2**num_compress_steps.
 
   Args:
-    inputs: Tensor of shape [batch, compress_height, compress_width, hidden_dim]
-    hparams: Dict, hyperparameters.
+    inputs: Tensor of shape [batch, compress_height, compress_width, channels].
+    hparams: tf.contrib.training.HParams.
     strides: Tuple, strides for conv block.
     kernel: Tuple, kernel window size for conv block.
     name: string, variable scope.
 
   Returns:
-    x: Tensor of shape [batch, height, width, hidden_dim].
+    Tensor of shape [batch, height, width, hparams.hidden_size].
   """
   with tf.variable_scope(name):
     x = inputs
-    # Reshape?
     x = tf.layers.dense(x, hparams.hidden_size, name=name + "_dense")
-    # Residual blocks.
     x = residual_block_layer(x, hparams)
-
-    # Decompress conv layers with strides and kernels as passed to the function.
     for i in range(hparams.num_compress_steps // 2):
       j = hparams.num_compress_steps // 2 - i - 1
       with tf.variable_scope(name + "_%d" % j):
@@ -279,15 +313,15 @@ def decompress_decoder(inputs, hparams,
 
 
 def decompress_decoder_2d(x, hparams, name):
-  """Decoder that decompresses x to length height*width.
+  """Decoder that decompresses 2-D inputs by 2**num_compress_steps.
 
   Args:
-    x: Tensor of shape [batch, compress_height, compress_width, hidden_dim].
-    hparams: Dict, hyperparameters.
+    x: Tensor of shape [batch, compress_height, compress_width, channels].
+    hparams: tf.contrib.training.HParams.
     name: string, variable scope.
 
   Returns:
-    x: Tensor of shape [batch, height, width, hidden_dim].
+    Tensor of shape [batch, height, width, hparams.hidden_size].
   """
   return decompress_decoder(x, hparams,
                             strides=(2, 2),
@@ -296,15 +330,15 @@ def decompress_decoder_2d(x, hparams, name):
 
 
 def decompress_decoder_1d(x, hparams, name):
-  """Dencoder that decompresses x to original target length.
+  """Decoder that decompresses 1-D inputs by 2**num_compress_steps.
 
   Args:
-    x: Tensor of shape [batch, compress_length, hidden_dim].
-    hparams: Dict, hyperparameters.
+    x: Tensor of shape [batch, compress_length, channels].
+    hparams: tf.contrib.training.HParams.
     name: string, variable scope.
 
   Returns:
-    x: Tensor of shape [batch, length, hidden_dim].
+    Tensor of shape [batch, length, hparams.hidden_size].
   """
   x = tf.expand_dims(x, axis=2)
   output = decompress_decoder(x, hparams,
@@ -314,19 +348,22 @@ def decompress_decoder_1d(x, hparams, name):
   return tf.squeeze(output, axis=2)
 
 
-def transformer_text_encoder(
-    x, space_id, hparams, name="transformer_text_encoder"):
+def transformer_text_encoder(x,
+                             space_id,
+                             hparams,
+                             name="transformer_text_encoder"):
   """Transformer text encoder over inputs with unmasked full attention.
 
   Args:
-    x: Tensor of shape [batch, length, hidden_dim].
+    x: Tensor of shape [batch, length, 1, hparams.hidden_size].
     space_id: int, id.
-    hparams: Dict, hyperparameters.
+    hparams: tf.contrib.training.HParams.
     name: string, variable scope.
 
   Returns:
-    x: Tensor of shape [batch, length, hidden_dim].
-    ed: Tensor, bias for padded tokens in the input, shape [batch, length]
+    encoder_output: Tensor of shape [batch, length, hparams.hidden_size].
+    ed: Tensor of shape [batch, 1, 1, length]. Encoder-decoder attention bias
+      for any padded tokens.
   """
   with tf.variable_scope(name):
     x = common_layers.flatten4d3d(x)
@@ -345,24 +382,23 @@ def transformer_image_decoder(x,
   """Transformer image decoder over inputs with local attention.
 
   Args:
-    x: Tensor of shape [batch, height, width, hidden_dim].
-    encoder_output: Tensor, encoder output of shape [batch, length, hidden_dim].
-    ed_attention_bias: Tensor, bias for x.
-    hparams: Dict, hyperparameters.
+    x: Tensor of shape [batch, ...], and whose size is batch * height * width *
+      hparams.num_channels * hparams.hidden_size.
+    encoder_output: Tensor of shape [batch, length_kv, hparams.hidden_size].
+    ed_attention_bias: Tensor which broadcasts with shape [batch,
+      hparams.num_heads, length_q, length_kv]. Encoder-decoder attention bias.
+    hparams: tf.contrib.training.HParams.
     name: string, variable scope.
 
   Returns:
-    x: Tensor of shape [batch, height, width, hidden_dim].
+    Tensor of shape [batch, height, width * hparams.num_channels,
+    hparams.hidden_size].
   """
   with tf.variable_scope(name):
     batch_size = common_layers.shape_list(x)[0]
-    # Reshape targets as b, 32, 32, 3*hidden size].
     targets = tf.reshape(x, [
         batch_size, hparams.img_len, hparams.img_len,
         hparams.num_channels*hparams.hidden_size])
-
-    # Prepare decoder inputs and bias. This also shifts targets and adds 2D
-    # position embeddings to target.
     decoder_input, _, _ = cia.prepare_decoder(targets, hparams)
     decoder_output = cia.transformer_decoder_layers(
         decoder_input,
@@ -388,28 +424,27 @@ def transformer_latent_decoder(x,
   """Transformer decoder over latents using latent_attention_type.
 
   Args:
-    x: Tensor of shape [batch, height, width, hidden_dim].
-    encoder_output: Tensor, encoder output of shape [batch, length, hidden_dim].
-    ed_attention_bias: Tensor, bias for x.
-    hparams: Dict, hyperparameters.
+    x: Tensor of shape [batch, ...], and whose size is batch * length_q *
+      hparams.hidden_size. Here, length_q is the latent length, which is
+      height * width * hparams.num_latents / (2**hparams.num_compress_steps).
+    encoder_output: Tensor of shape [batch, length_kv, hparams.hidden_size].
+    ed_attention_bias: Tensor which broadcasts with shape [batch,
+      hparams.num_heads, length_q, length_kv]. Encoder-decoder attention bias.
+    hparams: tf.contrib.training.HParams.
     name: string, variable scope.
 
   Returns:
-    x: Tensor of shape [batch, height, width, hidden_dim].
+    Tensor of shape [batch, length_q, hparams.hidden_size].
   """
   with tf.variable_scope(name):
     batch_size = common_layers.shape_list(x)[0]
     compress_ratio = 2**(hparams.num_compress_steps // 2)
-    # Reshape targets as b, 32, 32, 3*hidden size].
     x = tf.reshape(x, [
         batch_size, hparams.img_len / compress_ratio,
         (hparams.img_len*hparams.num_latents) / compress_ratio,
         hparams.hidden_size
     ])
-
-    # Prepare decoder inputs and bias.
     decoder_input, _, _ = cia.prepare_decoder(x, hparams)
-    # hparams.num_channels = 3
     decoder_output = cia.transformer_decoder_layers(
         decoder_input,
         encoder_output,
@@ -448,7 +483,25 @@ def latent_prediction_model(inputs,
                             latents_dense,
                             hparams,
                             name="latent_prediction"):
-  """Transformer based latent prediction model."""
+  """Transformer-based latent prediction model.
+
+  It is an autoregressive decoder over latents_discrete given inputs.
+
+  Args:
+    inputs: Tensor of shape [batch, length_kv, hparams.hidden_size]. Inputs to
+      attend to for the decoder on latents.
+    ed_attention_bias: Tensor which broadcasts with shape [batch,
+      hparams.num_heads, length_q, length_kv]. Encoder-decoder attention bias.
+    latents_discrete: Tensor of shape [batch, length_q, vocab_size].
+      One-hot latents to compute log-probability of given inputs.
+    latents_dense: Tensor of shape [batch, length_q, hparams.hidden_size].
+    hparams: tf.contrib.training.HParams.
+    name: string, variable scope.
+
+  Returns:
+    latents_pred: Tensor of shape [batch, length_q, hparams.hidden_size].
+    latents_pred_loss: Tensor of shape [batch, length_q].
+  """
   with tf.variable_scope(name):
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
       latents_pred = transformer_latent_decoder(
@@ -465,7 +518,7 @@ def transformer_autoencoder(inputs,
                             hparams,
                             cache=None,
                             predict_mask=1.0):
-  """Auto Encoder using transformer decoder and prior over latents."""
+  """Auto-encoder using transformer decoder and prior over latents."""
   # Define losses
   losses = {"extra": tf.constant(0.0), "latent_pred": tf.constant(0.0)}
 

From 1585cc66df47e9a6e4a4da259fffe202e698777e Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Mon, 9 Jul 2018 10:39:06 -0700
Subject: [PATCH 0294/2720] Adding Stanford Sentiment Treebank (binary
 classification).

PiperOrigin-RevId: 203789847
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/sst_binary.py   | 122 ++++++++++++++++++
 2 files changed, 123 insertions(+)
 create mode 100644 tensor2tensor/data_generators/sst_binary.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index ef75308a5..3aa15c111 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -54,6 +54,7 @@
     "tensor2tensor.data_generators.snli",
     "tensor2tensor.data_generators.style_transfer",
     "tensor2tensor.data_generators.squad",
+    "tensor2tensor.data_generators.sst_binary",
     "tensor2tensor.data_generators.subject_verb_agreement",
     "tensor2tensor.data_generators.timeseries",
     "tensor2tensor.data_generators.translate_encs",
diff --git a/tensor2tensor/data_generators/sst_binary.py b/tensor2tensor/data_generators/sst_binary.py
new file mode 100644
index 000000000..7b9dd585c
--- /dev/null
+++ b/tensor2tensor/data_generators/sst_binary.py
@@ -0,0 +1,122 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Stanford Sentiment Treebank Binary Classification Problem."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import zipfile
+import six
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+EOS = text_encoder.EOS
+
+
+@registry.register_problem
+class SentimentSSTBinary(text_problems.Text2ClassProblem):
+  """Stanford Sentiment Treebank binary classification problems."""
+
+  # Link to data from GLUE: https://gluebenchmark.com/tasks
+  _SST2_URL = ("https://firebasestorage.googleapis.com/v0/b/"
+               "mtl-sentence-representations.appspot.com/o/"
+               "data%2FSST-2.zip?alt=media&token=aabc5f6b-e466-"
+               "44a2-b9b4-cf6337f84ac8")
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 10,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  @property
+  def approx_vocab_size(self):
+    return 2**13  # 8k vocab suffices for this small dataset.
+
+  @property
+  def vocab_filename(self):
+    return "vocab.sst_binary.%d" % self.approx_vocab_size
+
+  @property
+  def num_classes(self):
+    return 2
+
+  def class_labels(self, data_dir):
+    del data_dir
+    # Note this binary classification is different from usual MNLI.
+    return ["neg", "pos"]
+
+  def _maybe_download_corpora(self, tmp_dir):
+    sst_binary_filename = "SST-2.zip"
+    sst_binary_finalpath = os.path.join(tmp_dir, "SST-2")
+    if not tf.gfile.Exists(sst_binary_finalpath):
+      zip_filepath = generator_utils.maybe_download(
+          tmp_dir, sst_binary_filename, self._SST2_URL)
+      zip_ref = zipfile.ZipFile(zip_filepath, "r")
+      zip_ref.extractall(tmp_dir)
+      zip_ref.close()
+
+    return sst_binary_finalpath
+
+  def example_generator(self, filename):
+    for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
+      if idx == 0: continue  # skip header
+      if six.PY2:
+        line = unicode(line.strip(), "utf-8")
+      else:
+        line = line.strip().decode("utf-8")
+      sent, label = line.split("\t")
+      yield {
+          "inputs": sent,
+          "label": int(label)
+      }
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    sst_binary_dir = self._maybe_download_corpora(tmp_dir)
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      filesplit = "train.tsv"
+    else:
+      filesplit = "dev.tsv"
+
+    filename = os.path.join(sst_binary_dir, filesplit)
+    for example in self.example_generator(filename):
+      yield example
+
+
+@registry.register_problem
+class SentimentSSTBinaryCharacters(SentimentSSTBinary):
+  """Binary Stanford Sentiment Treebank problems, character level"""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  @property
+  def task_id(self):
+    return problem.SpaceID.EN_CHR_SENT

From 30e85d1b955a6e91721de5723e58e9266a3ec22f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 9 Jul 2018 10:56:56 -0700
Subject: [PATCH 0295/2720] Added synthetic data generator for timeseries.
 Added a new timeseries problem with 10 series and 100k samples that is
 generated using the synthetic data generator.

PiperOrigin-RevId: 203793037
---
 tensor2tensor/data_generators/timeseries.py   | 119 ++++++++++++++++++
 .../timeseries_data_generator.py              |  62 +++++++++
 .../timeseries_data_generator_test.py         | 110 ++++++++++++++++
 .../data_generators/timeseries_test.py        |   6 +
 4 files changed, 297 insertions(+)
 create mode 100644 tensor2tensor/data_generators/timeseries_data_generator.py
 create mode 100644 tensor2tensor/data_generators/timeseries_data_generator_test.py

diff --git a/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index 38d23eea8..f1367b573 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -21,6 +21,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import timeseries_data_generator
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -203,3 +204,121 @@ def timeseries_dataset(self):
     series = [[float(i + n) for n in range(self.num_series)] for i in range(10)]
 
     return np.array(series)
+
+
+@registry.register_problem
+class TimeseriesSyntheticDataSeries10Samples100k(TimeseriesProblem):
+  """10 synthetic timeseries with 100K samples/timestamps."""
+
+  @property
+  def num_train_shards(self):
+    """Number of training shards."""
+    return 9
+
+  @property
+  def num_eval_shards(self):
+    """Number of eval shards."""
+    return 1
+
+  @property
+  def num_series(self):
+    """Number of timeseries."""
+    return 10
+
+  @property
+  def num_input_timestamps(self):
+    """Number of timestamps to include in the input."""
+    return 500
+
+  @property
+  def num_target_timestamps(self):
+    """Number of timestamps to include in the target."""
+    return 100
+
+  @property
+  def normalizing_constant(self):
+    return 0.01
+
+  @property
+  def timeseries_params(self):
+    """Parameters for each timeseries."""
+    timeseries_params = [{
+        "m": 0.006,
+        "b": 300.0,
+        "A": 50.0,
+        "freqcoeff": 1500.0,
+        "rndA": 15.0,
+        "fn": np.sin
+    }, {
+        "m": 0.000,
+        "b": 500.0,
+        "A": 35.0,
+        "freqcoeff": 3500.0,
+        "rndA": 25.0,
+        "fn": np.cos
+    }, {
+        "m": -0.003,
+        "b": 800.0,
+        "A": 65.0,
+        "freqcoeff": 2500.0,
+        "rndA": 5.0,
+        "fn": np.sin
+    }, {
+        "m": 0.009,
+        "b": 600.0,
+        "A": 20.0,
+        "freqcoeff": 1000.0,
+        "rndA": 1.0,
+        "fn": np.cos
+    }, {
+        "m": 0.002,
+        "b": 700.0,
+        "A": 40.0,
+        "freqcoeff": 2000.0,
+        "rndA": 35.0,
+        "fn": np.sin
+    }, {
+        "m": -0.008,
+        "b": 1000.0,
+        "A": 70.0,
+        "freqcoeff": 3000.0,
+        "rndA": 25.0,
+        "fn": np.cos
+    }, {
+        "m": 0.000,
+        "b": 100.0,
+        "A": 25.0,
+        "freqcoeff": 1500.0,
+        "rndA": 10.0,
+        "fn": np.sin
+    }, {
+        "m": 0.004,
+        "b": 1500.0,
+        "A": 54.0,
+        "freqcoeff": 900.0,
+        "rndA": 55.0,
+        "fn": np.cos
+    }, {
+        "m": 0.005,
+        "b": 2000.0,
+        "A": 32.0,
+        "freqcoeff": 1100.0,
+        "rndA": 43.0,
+        "fn": np.sin
+    }, {
+        "m": 0.010,
+        "b": 2500.0,
+        "A": 43.0,
+        "freqcoeff": 1900.0,
+        "rndA": 53.0,
+        "fn": np.cos
+    }]
+
+    return timeseries_params
+
+  def timeseries_dataset(self):
+    series = np.array(
+        timeseries_data_generator.generate_data(100000, self.timeseries_params))
+
+    series = series.transpose()
+    return series
diff --git a/tensor2tensor/data_generators/timeseries_data_generator.py b/tensor2tensor/data_generators/timeseries_data_generator.py
new file mode 100644
index 000000000..db21dc529
--- /dev/null
+++ b/tensor2tensor/data_generators/timeseries_data_generator.py
@@ -0,0 +1,62 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generator for the timeseries problem."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+
+def generate_data(timeseries_length, timeseries_params):
+  """Generates synthetic timeseries using input parameters.
+
+  Each generated timeseries has timeseries_length data points.
+  Parameters for each timeseries are specified by timeseries_params.
+
+  Args:
+    timeseries_length: Number of data points to generate for each timeseries.
+    timeseries_params: Parameters used to generate the timeseries. The following
+      parameters need to be specified for each timeseries:
+      m = Slope of the timeseries used to compute the timeseries trend.
+      b = y-intercept of the timeseries used to compute the timeseries trend.
+      A = Timeseries amplitude used to compute timeseries period.
+      freqcoeff = Frequency coefficient used to compute timeseries period.
+      rndA = Random amplitude used to inject noise into the timeseries.
+      fn = Base timeseries function (np.cos or np.sin).
+      Example params for two timeseries.
+      [{"m": 0.006, "b": 300.0, "A":50.0, "freqcoeff":1500.0, "rndA":15.0,
+      "fn": np.sin},
+      {"m": 0.000, "b": 500.0, "A":35.0, "freqcoeff":3500.0, "rndA":25.0,
+      "fn": np.cos}]
+
+  Returns:
+    Multi-timeseries (list of list).
+  """
+  x = range(timeseries_length)
+
+  multi_timeseries = []
+  for p in timeseries_params:
+    # Trend
+    y1 = [p["m"] * i + p["b"] for i in x]
+    # Period
+    y2 = [p["A"] * p["fn"](i / p["freqcoeff"]) for i in x]
+    # Noise
+    y3 = np.random.normal(0, p["rndA"], timeseries_length).tolist()
+    # Sum of Trend, Period and Noise. Replace negative values with zero.
+    y = [max(a + b + c, 0) for a, b, c in zip(y1, y2, y3)]
+    multi_timeseries.append(y)
+
+  return multi_timeseries
diff --git a/tensor2tensor/data_generators/timeseries_data_generator_test.py b/tensor2tensor/data_generators/timeseries_data_generator_test.py
new file mode 100644
index 000000000..85eb2a2f3
--- /dev/null
+++ b/tensor2tensor/data_generators/timeseries_data_generator_test.py
@@ -0,0 +1,110 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Timeseries data generator tests."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensor2tensor.data_generators import timeseries_data_generator
+
+import tensorflow as tf
+
+
+class TimeseriesDataGeneratorTest(tf.test.TestCase):
+
+  def testGenerateData(self):
+    timeseries_params = [{
+        "m": 0.006,
+        "b": 300.0,
+        "A": 50.0,
+        "freqcoeff": 1500.0,
+        "rndA": 15.0,
+        "fn": np.sin
+    }, {
+        "m": 0.000,
+        "b": 500.0,
+        "A": 35.0,
+        "freqcoeff": 3500.0,
+        "rndA": 25.0,
+        "fn": np.cos
+    }, {
+        "m": -0.003,
+        "b": 800.0,
+        "A": 65.0,
+        "freqcoeff": 2500.0,
+        "rndA": 5.0,
+        "fn": np.sin
+    }, {
+        "m": 0.009,
+        "b": 600.0,
+        "A": 20.0,
+        "freqcoeff": 1000.0,
+        "rndA": 1.0,
+        "fn": np.cos
+    }, {
+        "m": 0.002,
+        "b": 700.0,
+        "A": 40.0,
+        "freqcoeff": 2000.0,
+        "rndA": 35.0,
+        "fn": np.sin
+    }, {
+        "m": -0.008,
+        "b": 1000.0,
+        "A": 70.0,
+        "freqcoeff": 3000.0,
+        "rndA": 25.0,
+        "fn": np.cos
+    }, {
+        "m": 0.000,
+        "b": 100.0,
+        "A": 25.0,
+        "freqcoeff": 1500.0,
+        "rndA": 10.0,
+        "fn": np.sin
+    }, {
+        "m": 0.004,
+        "b": 1500.0,
+        "A": 54.0,
+        "freqcoeff": 900.0,
+        "rndA": 55.0,
+        "fn": np.cos
+    }, {
+        "m": 0.005,
+        "b": 2000.0,
+        "A": 32.0,
+        "freqcoeff": 1100.0,
+        "rndA": 43.0,
+        "fn": np.sin
+    }, {
+        "m": 0.010,
+        "b": 2500.0,
+        "A": 43.0,
+        "freqcoeff": 1900.0,
+        "rndA": 53.0,
+        "fn": np.cos
+    }]
+    multi_timeseries = timeseries_data_generator.generate_data(
+        20, timeseries_params)
+
+    self.assertEqual(10, len(multi_timeseries))
+    self.assertEqual(20, len(multi_timeseries[0]))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index ce0c4f911..ccfc34565 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -60,6 +60,12 @@ def testTimeseriesToyProblem(self):
     self.assertNotEqual(
         list(examples[0]["inputs"][0, 0]), list(examples[1]["inputs"][0, 0]))
 
+  def testTimeseriesSyntheticData10Series100kSamples(self):
+    problem = timeseries.TimeseriesSyntheticDataSeries10Samples100k()
+    self.assertEqual(10, problem.num_series)
+    self.assertEqual(500, problem.num_input_timestamps)
+    self.assertEqual(100, problem.num_target_timestamps)
+
 
 if __name__ == "__main__":
   tf.test.main()

From 7acadc6fe0fbe435a9a6cc52cd2d229a026a5ce1 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Mon, 9 Jul 2018 11:16:33 -0700
Subject: [PATCH 0296/2720] Adding WNLI data.

PiperOrigin-RevId: 203797175
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/problem.py      |   6 +
 .../data_generators/text_problems.py          |  41 ++++++
 tensor2tensor/data_generators/wnli.py         | 134 ++++++++++++++++++
 4 files changed, 182 insertions(+)
 create mode 100644 tensor2tensor/data_generators/wnli.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 3aa15c111..039a9986d 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -71,6 +71,7 @@
     "tensor2tensor.data_generators.wikisum.wikisum",
     "tensor2tensor.data_generators.wikitext103",
     "tensor2tensor.data_generators.wsj_parsing",
+    "tensor2tensor.data_generators.wnli",
 ]
 ALL_MODULES = list(MODULES)
 
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 9bf267426..dae28a4c3 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -102,6 +102,12 @@ class SpaceID(object):
   PICKLED_PYTHON = 30
   # English characters sentiment
   EN_CHR_SENT = 31
+  # English Premise Hypothesis pair
+  EN_PR_HYP = 32
+  # English NLI
+  EN_NLI = 33
+  # COLA
+  COLA = 34
 
 
 def default_model_hparams():
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 956f1d385..b50a475c5 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -499,6 +499,47 @@ def example_reading_spec(self):
     return (data_fields, data_items_to_decoders)
 
 
+class TextConcat2ClassProblem(Text2ClassProblem):
+  """Base class for text classification problems with multiple inputs.
+
+  For problems where there are multiple input sentences and we wish to concat
+  these inputs with a special delimiter. See, for example, NLI tasks.
+  """
+
+  @property
+  def concat_token(self):
+    raise NotImplementedError()
+
+  @property
+  def concat_id(self):
+    raise NotImplementedError()
+
+  @property
+  def additional_reserved_tokens(self):
+    return [self.concat_token]
+
+  def generate_text_for_vocab(self, data_dir, tmp_dir):
+    for i, sample in enumerate(
+        self.generate_samples(data_dir, tmp_dir, problem.DatasetSplit.TRAIN)):
+      for inp in sample["inputs"]:
+        yield inp
+        if self.max_samples_for_vocab and (i + 1) >= self.max_samples_for_vocab:
+          break
+
+  def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
+    generator = self.generate_samples(data_dir, tmp_dir, dataset_split)
+    encoder = self.get_or_create_vocab(data_dir, tmp_dir)
+    for sample in generator:
+      inputs = []
+      for idx, inp in enumerate(sample["inputs"]):
+        inputs += encoder.encode(inp)
+        inputs.append(text_encoder.EOS_ID)
+        if idx < len(sample["inputs"])-1:
+          inputs.append(self.concat_id)
+      label = sample["label"]
+      yield {"inputs": inputs, "targets": [label]}
+
+
 def txt_line_iterator(txt_path):
   """Iterate through lines of file."""
   with tf.gfile.Open(txt_path) as f:
diff --git a/tensor2tensor/data_generators/wnli.py b/tensor2tensor/data_generators/wnli.py
new file mode 100644
index 000000000..fc124b8f9
--- /dev/null
+++ b/tensor2tensor/data_generators/wnli.py
@@ -0,0 +1,134 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generators for the Winograd NLI dataset."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import zipfile
+import six
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+EOS = text_encoder.EOS
+
+
+@registry.register_problem
+class WinogradNLI(text_problems.TextConcat2ClassProblem):
+  """Winograd NLI classification problems."""
+
+  # Link to data from GLUE: https://gluebenchmark.com/tasks
+  _WNLI_URL = ("https://firebasestorage.googleapis.com/v0/b/"
+               "mtl-sentence-representations.appspot.com/o/"
+               "data%2FWNLI.zip?alt=media&token=068ad0a0-ded7-"
+               "4bd7-99a5-5e00222e0faf")
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 1,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  @property
+  def approx_vocab_size(self):
+    return 2**13  # 8k vocab suffices for this small dataset.
+
+  @property
+  def vocab_filename(self):
+    return "vocab.wnli.%d" % self.approx_vocab_size
+
+  @property
+  def num_classes(self):
+    return 2
+
+  @property
+  def concat_token(self):
+    return "<EN-PR-HYP>"
+
+  @property
+  def concat_id(self):
+    if self.vocab_type == text_problems.VocabType.CHARACTER:
+      return problem.SpaceID.EN_PR_HYP
+    return 2
+
+  def class_labels(self, data_dir):
+    del data_dir
+    # Note this binary classification is different from usual MNLI.
+    return ["contradiction", "entailment"]
+
+  def _maybe_download_corpora(self, tmp_dir):
+    wnli_filename = "WNLI.zip"
+    wnli_finalpath = os.path.join(tmp_dir, "WNLI")
+    if not tf.gfile.Exists(wnli_finalpath):
+      zip_filepath = generator_utils.maybe_download(
+          tmp_dir, wnli_filename, self._WNLI_URL)
+      zip_ref = zipfile.ZipFile(zip_filepath, "r")
+      zip_ref.extractall(tmp_dir)
+      zip_ref.close()
+
+    return wnli_finalpath
+
+  def example_generator(self, filename):
+    for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
+      if idx == 0: continue  # skip header
+      if six.PY2:
+        line = unicode(line.strip(), "utf-8")
+      else:
+        line = line.strip().decode("utf-8")
+      _, s1, s2, l = line.split("\t")
+      # inputs = " ".join([s1 EOS, s2])
+      inputs = [s1, s2]
+      yield {
+          "inputs": inputs,
+          "label": int(l)
+      }
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    wnli_dir = self._maybe_download_corpora(tmp_dir)
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      filesplit = "train.tsv"
+    else:
+      filesplit = "dev.tsv"
+
+    filename = os.path.join(wnli_dir, filesplit)
+    for example in self.example_generator(filename):
+      yield example
+
+
+@registry.register_problem
+class WinogradNLICharacters(WinogradNLI):
+  """Winograd NLI classification problems, character level"""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  @property
+  def task_id(self):
+    return problem.SpaceID.EN_NLI

From 94265ab7d33c084727f132357dc558c75c010ac3 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 9 Jul 2018 11:19:07 -0700
Subject: [PATCH 0297/2720] Add a generic L2 modality.

PiperOrigin-RevId: 203797672
---
 tensor2tensor/layers/modalities.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index ab2b46966..49da0c4cb 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -797,6 +797,18 @@ def top(self, body_output, _):
     return body_output
 
 
+@registry.register_generic_modality("l2_loss")
+class GenericL2LossModality(IdentityModality):
+  """Generic modality with L2 as Loss."""
+
+  def targets_bottom(self, x):
+    return tf.to_float(x)
+
+  def loss(self, body_output, targets):
+    loss = tf.square(body_output - tf.to_float(targets))
+    return tf.reduce_mean(loss), tf.constant(1.0)
+
+
 class RealModality(modality.Modality):
   """Base class for real (i.e. float) vectors.
 

From 4d30048add95c3a35101d33fce148d17937e58e1 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 9 Jul 2018 11:20:01 -0700
Subject: [PATCH 0298/2720] enable video metric calculation from in-memory
 decodings

PiperOrigin-RevId: 203797878
---
 tensor2tensor/utils/video_metrics.py | 96 +++++++++++++++++++++-------
 1 file changed, 72 insertions(+), 24 deletions(-)

diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index 54a4265be..43b473d6e 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -69,12 +69,14 @@ def get_target_and_output_filepatterns(output_dir, problem_name):
           file_pattern(output_dir, problem_name, "targets"))
 
 
-def get_zipped_dataset(output_files, target_files, video_length, frame_shape):
+def get_zipped_dataset_from_png_files(
+    output_files, target_files, video_length, frame_shape):
   outputs, len_ = load_videos(output_files, video_length, frame_shape)
   targets, len_ = load_videos(target_files, video_length, frame_shape)
   zipped_dataset = tf.data.Dataset.zip((outputs, targets))
   num_videos = len_ // video_length
-  return zipped_dataset, num_videos
+  iterator = zipped_dataset.make_one_shot_iterator()
+  return iterator, None, num_videos
 
 
 def save_results(results, output_dir, problem_name):
@@ -87,39 +89,58 @@ def save_results(results, output_dir, problem_name):
 
 def compute_metrics(output_video, target_video):
   max_pixel_value = 255.0
+  output_video = tf.to_float(output_video)
+  target_video = tf.to_float(target_video)
   psnr = tf.image.psnr(output_video, target_video, max_pixel_value)
   ssim = tf.image.ssim(output_video, target_video, max_pixel_value)
   return {"PSNR": psnr, "SSIM": ssim}
 
 
-def compute_one_decoding_video_metrics(
-    output_dir, problem_name, video_length, frame_shape):
-  """Computes the average of all the metric for one decoding.
+def stack_data_given_key(predictions, key):
+  x = [p[key] for p in predictions]
+  x = np.stack(x, axis=0)
+  return x
 
-  This function assumes that all the predicted and target frames
-  have been saved on the disk and sorting them by name will result
-  to consecutive frames saved in order.
+
+def get_zipped_dataset_from_predictions(predictions):
+  """Creates dataset from in-memory predictions."""
+  targets = stack_data_given_key(predictions, "targets")
+  outputs = stack_data_given_key(predictions, "outputs")
+  num_videos = len(targets)
+
+  targets_placeholder = tf.placeholder(targets.dtype, targets.shape)
+  outputs_placeholder = tf.placeholder(outputs.dtype, outputs.shape)
+  dataset = tf.data.Dataset.from_tensor_slices(
+      (targets_placeholder, outputs_placeholder))
+  iterator = dataset.make_initializable_iterator()
+  feed_dict = {targets_placeholder: targets,
+               outputs_placeholder: outputs}
+
+  return iterator, feed_dict, num_videos
+
+
+def compute_one_decoding_video_metrics(iterator, feed_dict, num_videos):
+  """Computes the average of all the metric for one decoding.
 
   Args:
-    output_dir: directory with all the saved frames.
-    problem_name: prefix of the saved frames usually name of the problem.
-    video_length: length of the videos.
-    frame_shape: shape of each frame in HxWxC format.
+    iterator: dataset iterator.
+    feed_dict: feed dict to initialize iterator.
+    num_videos: number of videos.
 
   Returns:
     Dictionary which contains the average of each metric per frame.
   """
-  output_files, target_files = get_target_and_output_filepatterns(
-      output_dir, problem_name)
-  dataset, num_videos = get_zipped_dataset(
-      output_files, target_files, video_length, frame_shape)
-  output, target = dataset.make_one_shot_iterator().get_next()
+  output, target = iterator.get_next()
+
   metrics_dict = compute_metrics(output, target)
   metrics_names, metrics = zip(*six.iteritems(metrics_dict))
   means, update_ops = tf.metrics.mean_tensor(metrics)
 
   with tf.Session() as sess:
     sess.run(tf.local_variables_initializer())
+    initalizer = iterator._initializer  # pylint: disable=protected-access
+    if initalizer is not None:
+      sess.run(initalizer, feed_dict=feed_dict)
 
     # Compute mean over dataset
     for i in range(num_videos):
@@ -144,12 +165,39 @@ def compute_all_metrics_statistics(all_results):
   return statistics
 
 
-def compute_video_metrics(output_dirs, problem_name, video_length, frame_shape):
-  all_results = [
-      compute_one_decoding_video_metrics(
-          output_dir, problem_name, video_length, frame_shape)
-      for output_dir in output_dirs
-  ]
+def compute_video_metrics_from_predictions(predictions):
+  all_results = []
+  for prediction in predictions:
+    args = get_zipped_dataset_from_predictions(prediction)
+    all_results.append(compute_one_decoding_video_metrics(*args))
+  statistics = compute_all_metrics_statistics(all_results)
+  return statistics
+
+
+def compute_video_metrics_from_png_files(
+    output_dirs, problem_name, video_length, frame_shape):
+  """Computes the average of all the metric for one decoding.
+
+  This function assumes that all the predicted and target frames
+  have been saved on the disk and sorting them by name will result
+  to consecutive frames saved in order.
+
+  Args:
+    output_dirs: directory with all the saved frames.
+    problem_name: prefix of the saved frames usually name of the problem.
+    video_length: length of the videos.
+    frame_shape: shape of each frame in HxWxC format.
+
+  Returns:
+    Dictionary which contains the average of each metric per frame.
+  """
+  all_results = []
+  for output_dir in output_dirs:
+    output_files, target_files = get_target_and_output_filepatterns(
+        output_dir, problem_name)
+    args = get_zipped_dataset_from_png_files(
+        output_files, target_files, video_length, frame_shape)
+    all_results.append(compute_one_decoding_video_metrics(*args))
   statistics = compute_all_metrics_statistics(all_results)
   return statistics, all_results
 
@@ -157,7 +205,7 @@ def compute_video_metrics(output_dirs, problem_name, video_length, frame_shape):
 def compute_and_save_video_metrics(
     output_dirs, problem_name, video_length, frame_shape):
   """Compute and saves the video metrics."""
-  statistics, all_results = compute_video_metrics(
+  statistics, all_results = compute_video_metrics_from_png_files(
       output_dirs, problem_name, video_length, frame_shape)
   for results, output_dir in zip(all_results, output_dirs):
     save_results(results, output_dir, problem_name)

From d59ba339a2f0804c0bad9df3ecfe177bcb43a33c Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Mon, 9 Jul 2018 11:34:55 -0700
Subject: [PATCH 0299/2720] Adding the Corpus of Linguistic Acceptability.

PiperOrigin-RevId: 203800690
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/cola.py         | 121 ++++++++++++++++++
 tensor2tensor/data_generators/problem.py      |   1 +
 3 files changed, 123 insertions(+)
 create mode 100644 tensor2tensor/data_generators/cola.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 039a9986d..b065c25f5 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -30,6 +30,7 @@
     "tensor2tensor.data_generators.cifar",
     "tensor2tensor.data_generators.cipher",
     "tensor2tensor.data_generators.cnn_dailymail",
+    "tensor2tensor.data_generators.cola",
     "tensor2tensor.data_generators.common_voice",
     "tensor2tensor.data_generators.desc2code",
     "tensor2tensor.data_generators.fsns",
diff --git a/tensor2tensor/data_generators/cola.py b/tensor2tensor/data_generators/cola.py
new file mode 100644
index 000000000..47e47c86b
--- /dev/null
+++ b/tensor2tensor/data_generators/cola.py
@@ -0,0 +1,121 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generators for the Corpus of Liguistic Acceptability."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import zipfile
+import six
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+EOS = text_encoder.EOS
+
+
+@registry.register_problem
+class Cola(text_problems.Text2ClassProblem):
+  """Corpus of Linguistic Acceptability classification problems."""
+
+  # Link to data from GLUE: https://gluebenchmark.com/tasks
+  _COLA_URL = ("https://firebasestorage.googleapis.com/v0/b/"
+               "mtl-sentence-representations.appspot.com/o/"
+               "data%2FCoLA.zip?alt=media&token=46d5e637-3411-"
+               "4188-bc44-5809b5bfb5f4")
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 10,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  @property
+  def approx_vocab_size(self):
+    return 2**13  # 8k vocab suffices for this small dataset.
+
+  @property
+  def vocab_filename(self):
+    return "vocab.cola.%d" % self.approx_vocab_size
+
+  @property
+  def num_classes(self):
+    return 2
+
+  def class_labels(self, data_dir):
+    del data_dir
+    # Note this binary classification is different from usual MNLI.
+    return ["unacceptable", "acceptable"]
+
+  def _maybe_download_corpora(self, tmp_dir):
+    cola_filename = "CoLA.zip"
+    cola_finalpath = os.path.join(tmp_dir, "CoLA")
+    if not tf.gfile.Exists(cola_finalpath):
+      zip_filepath = generator_utils.maybe_download(
+          tmp_dir, cola_filename, self._COLA_URL)
+      zip_ref = zipfile.ZipFile(zip_filepath, "r")
+      zip_ref.extractall(tmp_dir)
+      zip_ref.close()
+
+    return cola_finalpath
+
+  def example_generator(self, filename):
+    for line in tf.gfile.Open(filename, "rb"):
+      if six.PY2:
+        line = unicode(line.strip(), "utf-8")
+      else:
+        line = line.strip().decode("utf-8")
+      _, label, _, sent = line.split("\t")
+      yield {
+          "inputs": sent,
+          "label": int(label)
+      }
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    cola_dir = self._maybe_download_corpora(tmp_dir)
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      filesplit = "train.tsv"
+    else:
+      filesplit = "dev.tsv"
+
+    filename = os.path.join(cola_dir, filesplit)
+    for example in self.example_generator(filename):
+      yield example
+
+
+@registry.register_problem
+class ColaCharacters(Cola):
+  """Corpus of Linguistic Acceptability problems, character level"""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  @property
+  def task_id(self):
+    return problem.SpaceID.EN_COLA
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index dae28a4c3..2d83b25f1 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -1118,3 +1118,4 @@ def skip_random_fraction(dataset, data_file):
   # replicas reading the same data in lock-step.
   num_skip = random.randint(0, _file_num_records_cached(data_file))
   return dataset.skip(num_skip)
+

From dc1557df06b5af2e91cf97a977efc80c24d9e045 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 9 Jul 2018 14:17:35 -0700
Subject: [PATCH 0300/2720] WikisumWeb data gen command fixes

PiperOrigin-RevId: 203827908
---
 tensor2tensor/data_generators/wikisum/README.md          | 2 +-
 .../data_generators/wikisum/get_references_web.py        | 9 +++++----
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/wikisum/README.md b/tensor2tensor/data_generators/wikisum/README.md
index 72037fb1a..f71d550da 100644
--- a/tensor2tensor/data_generators/wikisum/README.md
+++ b/tensor2tensor/data_generators/wikisum/README.md
@@ -174,7 +174,7 @@ python -m tensor2tensor.data_generators.wikisum.parallel_launch \
   --name=wikisum-web-refs \
   --log_dir=$BUCKET/logs \
   --setup_command="pip3 install tensorflow tensor2tensor aiohttp cchardet aiodns bs4 -U -q --user" \
-  --command_prefix="python3 wikisum/get_references_web.py --out_dir=$BUCKET/wiki_references --shard_id"
+  --command_prefix="python3 -m tensor2tensor.data_generators.wikisum.get_references_web --out_dir=$BUCKET/wiki_references --shard_id"
 
 # Generate vocabulary file
 python -m tensor2tensor.data_generators.wikisum.generate_vocab \
diff --git a/tensor2tensor/data_generators/wikisum/get_references_web.py b/tensor2tensor/data_generators/wikisum/get_references_web.py
index 05ddda100..44b4ac94e 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web.py
@@ -49,10 +49,11 @@
 flags = tf.flags
 FLAGS = flags.FLAGS
 
-flags.DEFINE_string("command",
-                    "python3 wikisum/get_references_web_single_group.py",
-                    "Command to run get_references_web_single_group, without "
-                    "flags.")
+flags.DEFINE_string(
+    "command",
+    "python3 -m "
+    "tensor2tensor.data_generators.wikisum.get_references_web_single_group",
+    "Command to run get_references_web_single_group, without flags.")
 
 
 def main(_):

From 36e9e7f88dc6ff245ed1c98eb3c541964b33f025 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 9 Jul 2018 14:36:17 -0700
Subject: [PATCH 0301/2720] Make fashion mnist have 1 channel

PiperOrigin-RevId: 203831190
---
 tensor2tensor/data_generators/mnist.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/data_generators/mnist.py b/tensor2tensor/data_generators/mnist.py
index 5db518461..382a38e30 100644
--- a/tensor2tensor/data_generators/mnist.py
+++ b/tensor2tensor/data_generators/mnist.py
@@ -228,6 +228,10 @@ class ImageFashionMnist(image_utils.Image2ClassProblem):
   def is_small(self):
     return True
 
+  @property
+  def num_channels(self):
+    return 1
+
   @property
   def num_classes(self):
     return 10

From 06123f9d0ebc4e0cacd6abed8ddceeead0ddb760 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 9 Jul 2018 15:40:10 -0700
Subject: [PATCH 0302/2720] Set shape in image_fashion_mnist

PiperOrigin-RevId: 203843511
---
 tensor2tensor/data_generators/mnist.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tensor2tensor/data_generators/mnist.py b/tensor2tensor/data_generators/mnist.py
index 382a38e30..7f83a8680 100644
--- a/tensor2tensor/data_generators/mnist.py
+++ b/tensor2tensor/data_generators/mnist.py
@@ -244,6 +244,12 @@ def class_labels(self):
   def train_shards(self):
     return 10
 
+  def preprocess_example(self, example, mode, unused_hparams):
+    image = example["inputs"]
+    image.set_shape([_MNIST_IMAGE_SIZE, _MNIST_IMAGE_SIZE, 1])
+    example["inputs"] = image
+    return example
+
   def generator(self, data_dir, tmp_dir, is_training):
     if is_training:
       return fashion_mnist_generator(tmp_dir, True, 60000)

From aa1b7ea08649ac053fd781a77d1157e2725572af Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 9 Jul 2018 17:29:52 -0700
Subject: [PATCH 0303/2720] Fix metrics_test and disable an autoencoder test

PiperOrigin-RevId: 203860731
---
 .../models/research/autoencoders_test.py       | 18 ++++++++++--------
 tensor2tensor/utils/metrics_test.py            |  6 +++---
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index dc1c79baa..9f1ce548a 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -34,16 +34,16 @@ def get_mnist_random_output(self, model_name, hparams_set=None,
     hparams_set = hparams_set or model_name
     x = np.random.random_integers(0, high=255, size=(1, 28, 28, 1))
     y = np.random.random_integers(0, high=9, size=(1, 1))
+    features = {
+        "targets": tf.constant(x, dtype=tf.int32),
+        "inputs": tf.constant(y, dtype=tf.int32),
+    }
     hparams = trainer_lib.create_hparams(
         hparams_set, problem_name="image_mnist_rev", data_dir=".")
+    model = registry.model(model_name)(hparams, mode)
+    tf.train.create_global_step()
+    logits, _ = model(features)
     with self.test_session() as session:
-      features = {
-          "targets": tf.constant(x, dtype=tf.int32),
-          "inputs": tf.constant(y, dtype=tf.int32),
-      }
-      tf.train.create_global_step()
-      model = registry.model(model_name)(hparams, mode)
-      logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
     return res
@@ -81,7 +81,9 @@ def testAutoencoderOrderedDiscreteVQ(self):
         "autoencoder_ordered_discrete", "autoencoder_ordered_discrete_vq")
     self.assertEqual(res.shape, self.mnist_output_shape)
 
-  def testAutoencoderStacked(self):
+  # TODO(lukaszkaiser): Re-enable test by conserving lost shape information
+  # in autoencoder_stacked.
+  def x_testAutoencoderStacked(self):
     res = self.get_mnist_random_output("autoencoder_stacked")
     self.assertEqual(res.shape, self.mnist_output_shape)
 
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index d3c9d9cb9..953bc51ae 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -16,14 +16,14 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
 import numpy as np
 from tensor2tensor.utils import metrics
 
-
 import tensorflow as tf
 
 
-class CommonLayersTest(tf.test.TestCase):
+class MetricsTest(tf.test.TestCase):
 
   def testAccuracyMetric(self):
     predictions = np.random.randint(1, 5, size=(12, 12, 12, 1))
@@ -239,7 +239,7 @@ def testMultilabelMatch3(self):
       session.run(tf.global_variables_initializer())
       _ = session.run(a_op)
       actual = session.run(a)
-    self.assertAlmostEqual(actual, expected)
+    self.assertAlmostEqual(actual, expected, places=6)
 
 
 if __name__ == '__main__':

From 4642c62c70ef995424747483a25316992bb3c549 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 10 Jul 2018 08:00:25 -0700
Subject: [PATCH 0304/2720] Remove unnecessary parts of bias inplace update for
 TPU based

PiperOrigin-RevId: 203946997
---
 tensor2tensor/models/transformer.py | 14 --------------
 1 file changed, 14 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 9a28aeef2..c6963b369 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -413,20 +413,6 @@ def symbols_to_logits_tpu_fn(ids, i, cache):
       bias = tf.slice(decoder_self_attention_bias, [0, 0, i, 0],
                       [bias_shape[0], bias_shape[1], 1, bias_shape[3]])
 
-      bias_padding = tf.fill([bias_shape[0], bias_shape[1], 1], -1e9)
-      tmp_bias = tf.transpose(bias, perm=[3, 0, 1, 2])
-      bias_index = i + 1
-      while_condition = lambda bias_index, _: tf.less(bias_index, decode_length)
-
-      def while_body(bias_index, tmp_bias):
-        tmp_bias = common_layers.tf_inplace_ops().alias_inplace_update(
-            tmp_bias, bias_index, bias_padding)
-        return bias_index + 1, tmp_bias
-
-      _, tmp_bias = tf.while_loop(
-          while_condition, while_body, (bias_index, tmp_bias))
-      bias = tf.transpose(tmp_bias, perm=[1, 2, 3, 0])
-
       with tf.variable_scope("body"):
         body_outputs = dp(
             self.decode,

From 6cd6a2a0fc0814ecd2ce8ea0fc326154499ea03e Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 10 Jul 2018 11:31:41 -0700
Subject: [PATCH 0305/2720] naming lint fix

PiperOrigin-RevId: 203982613
---
 tensor2tensor/models/research/autoencoders_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index 9f1ce548a..82ac3fdab 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -83,7 +83,7 @@ def testAutoencoderOrderedDiscreteVQ(self):
 
   # TODO(lukaszkaiser): Re-enable test by conserving lost shape information
   # in autoencoder_stacked.
-  def x_testAutoencoderStacked(self):
+  def xtestAutoencoderStacked(self):
     res = self.get_mnist_random_output("autoencoder_stacked")
     self.assertEqual(res.shape, self.mnist_output_shape)
 

From 0efa0a6be54ae178a5ec39052c2d7dee6b5b6374 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 10 Jul 2018 11:46:50 -0700
Subject: [PATCH 0306/2720] Division by zero bug fix.

PiperOrigin-RevId: 203985260
---
 tensor2tensor/utils/bleu_hook.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/bleu_hook.py b/tensor2tensor/utils/bleu_hook.py
index 56bff62b0..e03c5b77a 100644
--- a/tensor2tensor/utils/bleu_hook.py
+++ b/tensor2tensor/utils/bleu_hook.py
@@ -115,7 +115,12 @@ def compute_bleu(reference_corpus,
 
   if use_bp:
     ratio = translation_length / reference_length
-    bp = math.exp(1 - 1. / ratio) if ratio < 1.0 else 1.0
+    if ratio <= 0.0:
+      bp = 0.0
+    elif ratio >= 1.0:
+      bp = 1.0
+    else:
+      bp = math.exp(1 - 1. / ratio)
   bleu = geo_mean * bp
   return np.float32(bleu)
 

From 79deef4b334d03d371c9d5033cfb99979b304ad7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 10 Jul 2018 12:19:31 -0700
Subject: [PATCH 0307/2720] xavier initializer

PiperOrigin-RevId: 203990559
---
 tensor2tensor/utils/optimize.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index a64a2869a..970914e22 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -234,5 +234,7 @@ def get_variable_initializer(hparams):
   elif hparams.initializer == "uniform_unit_scaling":
     return tf.variance_scaling_initializer(
         hparams.initializer_gain, mode="fan_avg", distribution="uniform")
+  elif hparams.initializer == "xavier":
+    return tf.contrib.layers.xavier_initializer()
   else:
     raise ValueError("Unrecognized initializer: %s" % hparams.initializer)

From 2ffe24dc386eff86962319adf291df4f8d3f0b16 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 10 Jul 2018 12:20:55 -0700
Subject: [PATCH 0308/2720] In memory decoding.

PiperOrigin-RevId: 203990815
---
 tensor2tensor/bin/t2t_decoder.py             |  1 +
 tensor2tensor/data_generators/video_utils.py | 10 ++++-
 tensor2tensor/utils/decoding.py              | 46 +++++++++++++-------
 3 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 19b05b591..8d0693150 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -60,6 +60,7 @@
 flags.DEFINE_integer("decode_shards", 1, "Number of decoding replicas.")
 flags.DEFINE_string("score_file", "", "File to score. Each line in the file "
                     "must be in the format input \t target.")
+flags.DEFINE_bool("decode_in_memory", False, "Decode in memory.")
 
 
 def create_hparams():
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index a82059158..c2f4abcf8 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -46,14 +46,20 @@ def summarize_video_metrics(hook_args):
   current_problem = hook_args.problem
   hparams = hook_args.hparams
   output_dirs = hook_args.output_dirs
+  predictions = hook_args.predictions
   frame_shape = [
       current_problem.frame_height, current_problem.frame_width,
       current_problem.num_channels
   ]
   metrics_graph = tf.Graph()
   with metrics_graph.as_default():
-    metrics_results, _ = video_metrics.compute_video_metrics(
-        output_dirs, problem_name, hparams.video_num_target_frames, frame_shape)
+    if predictions:
+      metrics_results = video_metrics.compute_video_metrics_from_predictions(
+          predictions)
+    else:
+      metrics_results, _ = video_metrics.compute_video_metrics(
+          output_dirs, problem_name,
+          hparams.video_num_target_frames, frame_shape)
 
   summary_values = []
   for name, array in six.iteritems(metrics_results):
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 7944c5df5..362557ac8 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -154,28 +154,37 @@ def decode_from_dataset(estimator,
   infer_input_fn = problem.make_estimator_input_fn(
       tf.estimator.ModeKeys.PREDICT, hparams, dataset_kwargs=dataset_kwargs)
 
-  output_dirs = []
+  predictions, output_dirs = [], []
   for decode_id in range(decode_hp.num_decodes):
     tf.logging.info("Decoding {}".format(decode_id))
 
-    output_dir = os.path.join(estimator.model_dir, "decode_%05d" % decode_id)
-    tf.gfile.MakeDirs(output_dir)
-    output_dirs.append(output_dir)
-
-    decode_once(estimator,
-                problem_name,
-                hparams,
-                infer_input_fn,
-                decode_hp,
-                decode_to_file,
-                output_dir)
+    # Create decode directory if not in-memory decoding.
+    if not FLAGS.decode_in_memory:
+      output_dir = os.path.join(estimator.model_dir, "decode_%05d" % decode_id)
+      tf.gfile.MakeDirs(output_dir)
+      output_dirs.append(output_dir)
+
+    result = decode_once(estimator,
+                         problem_name,
+                         hparams,
+                         infer_input_fn,
+                         decode_hp,
+                         decode_to_file,
+                         output_dir,
+                         log_results=not FLAGS.decode_in_memory)
+
+    if FLAGS.decode_in_memory:
+      output_dirs = [output_dir]
+      predictions.append(result)
 
   run_postdecode_hooks(DecodeHookArgs(
       estimator=estimator,
       problem=problem,
       output_dirs=output_dirs,
       hparams=hparams,
-      decode_hparams=decode_hp))
+      decode_hparams=decode_hp,
+      predictions=predictions
+  ))
 
 
 def decode_once(estimator,
@@ -184,12 +193,16 @@ def decode_once(estimator,
                 infer_input_fn,
                 decode_hp,
                 decode_to_file,
-                output_dir):
+                output_dir,
+                log_results=True):
   """Decodes once."""
 
   # Get the predictions as an iterable
   predictions = estimator.predict(infer_input_fn)
 
+  if not log_results:
+    return list(predictions)
+
   # Prepare output file writers if decode_to_file passed
   decode_to_file = decode_to_file or decode_hp.decode_to_file
   if decode_to_file:
@@ -215,6 +228,7 @@ def decode_once(estimator,
   inputs_vocab_key = "inputs" if has_input else "targets"
   inputs_vocab = problem_hparams.vocabulary[inputs_vocab_key]
   targets_vocab = problem_hparams.vocabulary["targets"]
+
   for num_predictions, prediction in enumerate(predictions):
     num_predictions += 1
     inputs = prediction["inputs"]
@@ -740,7 +754,8 @@ def latest_checkpoint_step(ckpt_dir):
 
 class DecodeHookArgs(collections.namedtuple(
     "DecodeHookArgs",
-    ["estimator", "problem", "output_dirs", "hparams", "decode_hparams"])):
+    ["estimator", "problem", "output_dirs", "hparams",
+     "decode_hparams", "predictions"])):
   pass
 
 
@@ -758,6 +773,7 @@ def run_postdecode_hooks(decode_hook_args):
   parent_dir = os.path.join(decode_hook_args.output_dirs[0], os.pardir)
   final_dir = os.path.join(parent_dir, "decode")
   summary_writer = tf.summary.FileWriter(final_dir)
+
   for hook in hooks:
     # Isolate each hook in case it creates TF ops
     with tf.Graph().as_default():

From ff5003c1e24d08191c5de10d441c269d5a137a2f Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Tue, 10 Jul 2018 13:07:15 -0700
Subject: [PATCH 0309/2720] Fix for decoding text classification problems.

PiperOrigin-RevId: 203997895
---
 tensor2tensor/utils/decoding.py | 41 ++++++++++++++++++++++-----------
 1 file changed, 28 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 362557ac8..8e17e22d9 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -29,6 +29,8 @@
 
 from tensor2tensor.data_generators import problem as problem_lib
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
 import tensorflow as tf
 
 FLAGS = tf.flags.FLAGS
@@ -91,6 +93,10 @@ def fix_and_save_video(vid, prefix):
     fix_and_save_video(targets, "targets")
 
   is_image = "image" in problem_name
+  is_text2class = isinstance(registry.problem(problem_name),
+                             text_problems.Text2ClassProblem)
+  skip_eos_postprocess = is_image or is_text2class
+
   decoded_inputs = None
   if is_image and save_images:
     save_path = os.path.join(
@@ -100,7 +106,8 @@ def fix_and_save_video(vid, prefix):
     if identity_output:
       decoded_inputs = " ".join(map(str, inputs.flatten()))
     else:
-      decoded_inputs = inputs_vocab.decode(_save_until_eos(inputs, is_image))
+      decoded_inputs = inputs_vocab.decode(_save_until_eos(
+          inputs, skip_eos_postprocess))
 
     if log_results and not is_video:
       tf.logging.info("Inference results INPUT: %s" % decoded_inputs)
@@ -112,9 +119,11 @@ def fix_and_save_video(vid, prefix):
     if targets is not None:
       decoded_targets = " ".join(map(str, targets.flatten()))
   else:
-    decoded_outputs = targets_vocab.decode(_save_until_eos(outputs, is_image))
+    decoded_outputs = targets_vocab.decode(_save_until_eos(
+        outputs, skip_eos_postprocess))
     if targets is not None and log_results:
-      decoded_targets = targets_vocab.decode(_save_until_eos(targets, is_image))
+      decoded_targets = targets_vocab.decode(_save_until_eos(
+          targets, skip_eos_postprocess))
   if not is_video:
     tf.logging.info("Inference results OUTPUT: %s" % decoded_outputs)
   if targets is not None and log_results and not is_video:
@@ -446,6 +455,11 @@ def input_fn():
 def decode_interactively(estimator, hparams, decode_hp, checkpoint_path=None):
   """Interactive decoding."""
 
+  is_image = "image" in hparams.problem.name
+  is_text2class = isinstance(hparams.problem,
+                             text_problems.Text2ClassProblem)
+  skip_eos_postprocess = is_image or is_text2class
+
   def input_fn():
     gen_fn = make_input_fn_from_generator(
         _interactive_input_fn(hparams, decode_hp))
@@ -455,7 +469,6 @@ def input_fn():
 
   result_iter = estimator.predict(input_fn, checkpoint_path=checkpoint_path)
   for result in result_iter:
-    is_image = False  # TODO(lukaszkaiser): find out from problem id / class.
     targets_vocab = hparams.problem_hparams.vocabulary["targets"]
 
     if decode_hp.return_beams:
@@ -465,7 +478,8 @@ def input_fn():
         scores = np.split(result["scores"], decode_hp.beam_size, axis=0)
       for k, beam in enumerate(beams):
         tf.logging.info("BEAM %d:" % k)
-        beam_string = targets_vocab.decode(_save_until_eos(beam, is_image))
+        beam_string = targets_vocab.decode(_save_until_eos(
+            beam, skip_eos_postprocess))
         if scores is not None:
           tf.logging.info("\"%s\"\tScore:%f" % (beam_string, scores[k]))
         else:
@@ -475,7 +489,8 @@ def input_fn():
         tf.logging.info(" ".join(map(str, result["outputs"].flatten())))
       else:
         tf.logging.info(
-            targets_vocab.decode(_save_until_eos(result["outputs"], is_image)))
+            targets_vocab.decode(_save_until_eos(
+                result["outputs"], skip_eos_postprocess)))
 
 
 def _decode_batch_input_fn(num_decode_batches, sorted_inputs, vocabulary,
@@ -659,17 +674,17 @@ def _get_sorted_inputs(filename, num_shards=1, delimiter="\n"):
   return sorted_inputs, sorted_keys
 
 
-def _save_until_eos(hyp, is_image):
+def _save_until_eos(ids, skip=False):
   """Strips everything after the first <EOS> token, which is normally 1."""
-  hyp = hyp.flatten()
-  if is_image:
-    return hyp
+  ids = ids.flatten()
+  if skip:
+    return ids
   try:
-    index = list(hyp).index(text_encoder.EOS_ID)
-    return hyp[0:index]
+    index = list(ids).index(text_encoder.EOS_ID)
+    return ids[0:index]
   except ValueError:
     # No EOS_ID: return the array as-is.
-    return hyp
+    return ids
 
 
 def _interactive_input_tensor_to_features_dict(feature_map, hparams):

From f81261c8ae2e6c7f835dc9440497dc0ca1083913 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 10 Jul 2018 17:28:07 -0700
Subject: [PATCH 0310/2720] Add unique vs non-unique option for sorting task.

PiperOrigin-RevId: 204042250
---
 tensor2tensor/data_generators/algorithmic.py | 39 ++++++++++++++++----
 1 file changed, 32 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index 93f3b27a9..50fd887f7 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -23,6 +23,7 @@
 from tensor2tensor.data_generators import generator_utils as utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
 
@@ -453,7 +454,7 @@ class AlgorithmicSortProblem(AlgorithmicProblem):
 
   @property
   def num_symbols(self):
-    return 10
+    return max(self.train_length, self.dev_length)
 
   @property
   def train_length(self):
@@ -461,14 +462,19 @@ def train_length(self):
 
   @property
   def dev_length(self):
-    return 10
+    return self.train_length * 2
+
+  @property
+  def unique(self):
+    """Unique numbers wo/ replacement or w/ replacement in sorting task."""
+    return False
 
   def generator(self, nbr_symbols, max_length, nbr_cases):
     """Generating for sorting task on sequence of symbols.
 
     The length of the sequence is drawn uniformly at random from [1, max_length]
-    and then symbols are drawn uniformly at random from [0, nbr_symbols) until
-    nbr_cases sequences have been produced.
+    and then symbols are drawn (uniquely w/ or w/o replacement) uniformly at
+    random from [0, nbr_symbols) until nbr_cases sequences have been produced.
 
     Args:
       nbr_symbols: number of symbols to use in each sequence.
@@ -480,6 +486,25 @@ def generator(self, nbr_symbols, max_length, nbr_cases):
       target-list is input-list sorted.
     """
     for _ in range(nbr_cases):
-      l = np.random.randint(max_length) + 1
-      inputs = list(np.random.randint(nbr_symbols, size=l))
-      yield {"inputs": inputs, "targets": list(sorted(inputs))}
+      # Sample the sequence length.
+      length = np.random.randint(max_length) + 1
+
+      if self.unique:
+        # Sample our inputs w/o replacement.
+        inputs = np.arange(nbr_symbols)
+        np.random.shuffle(inputs)
+
+        # Truncate to the desired length.
+        inputs = inputs[:length]
+        inputs = list(inputs)
+      else:
+        inputs = list(np.random.randint(nbr_symbols, size=length))
+
+      # Targets are simply the sorted inputs.
+      targets = list(sorted(inputs))
+
+      yield {"inputs": inputs, "targets": targets}
+
+  def eval_metrics(self):
+    defaults = super(AlgorithmicSortProblem, self).eval_metrics()
+    return defaults + [metrics.Metrics.EDIT_DISTANCE]

From 4eb4f4cabb7e548f502396cf4e9e48a237dd76c4 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Wed, 11 Jul 2018 02:33:32 +0200
Subject: [PATCH 0311/2720] RL fixes and autoencoder reimplementation (#908)

* autoencoder wrapper

* autoencoder wrapper v2

* read agent's policy

* Autoencoder rl_modelrl_ae_tiny experiment.

* linter cleanup and configuration bugfix

* bug masked

* gym_problems specifications cleanup

* removing unused argument

* bug fix in eval

* eval_metrics added to summaries

* fixed done bug

* removing unused (and possibly bugged) wrappers

* Cleanup and review

* autoencoder restoring restored

* force_beginning_resets result reporting

* statistics collection bug fixed

* restore bug fix

* StackAndSkipWrapper bugfix

* linter fixes

* more linter fixes

* more linter fixes
---
 tensor2tensor/data_generators/all_problems.py |   2 +-
 tensor2tensor/data_generators/gym_problems.py | 421 ++++++------------
 .../data_generators/gym_problems_specs.py     | 277 ++++++++++++
 .../data_generators/gym_problems_test.py      |   4 +-
 tensor2tensor/models/research/rl.py           |   1 -
 tensor2tensor/rl/collect.py                   |  56 ++-
 tensor2tensor/rl/envs/batch_env.py            |   3 +
 tensor2tensor/rl/envs/batch_env_factory.py    |   9 +-
 tensor2tensor/rl/envs/in_graph_batch_env.py   |   2 +-
 tensor2tensor/rl/envs/py_func_batch_env.py    |   2 +-
 tensor2tensor/rl/envs/simulated_batch_env.py  |   6 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py    | 163 +++----
 tensor2tensor/rl/model_rl_experiment.py       | 111 +++--
 tensor2tensor/rl/rl_trainer_lib.py            |   5 +-
 tensor2tensor/rl/rl_trainer_lib_test.py       |  16 +-
 15 files changed, 625 insertions(+), 453 deletions(-)
 create mode 100644 tensor2tensor/data_generators/gym_problems_specs.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index b065c25f5..36641436a 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -36,7 +36,7 @@
     "tensor2tensor.data_generators.fsns",
     "tensor2tensor.data_generators.gene_expression",
     "tensor2tensor.data_generators.google_robot_pushing",
-    "tensor2tensor.data_generators.gym_problems",
+    "tensor2tensor.data_generators.gym_problems_specs",
     "tensor2tensor.data_generators.ice_parsing",
     "tensor2tensor.data_generators.imagenet",
     "tensor2tensor.data_generators.image_lsun",
diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 194f302ae..a2edecefb 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -23,31 +23,23 @@
 import gym
 import numpy as np
 
-# We need gym_utils for the game environments defined there.
-from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
-
+import tensorflow as tf
+from tensorflow.contrib.training import HParams
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
-from tensor2tensor.models.research import rl
+from tensor2tensor.models.research import rl, autoencoders
 from tensor2tensor.rl import collect
 from tensor2tensor.rl.envs import tf_atari_wrappers
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
-
-from tensorflow.contrib.training import HParams
-
 flags = tf.flags
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string("agent_policy_path", None, "File with model for agent.")
+
 flags.DEFINE_string("autoencoder_path", None,
                     "File with model for autoencoder.")
-flags.DEFINE_boolean(
-    "only_use_ae_for_policy", False,
-    "Whether to only use the autoencoder for the policy and "
-    "still write out full-resolution frames.")
 
 
 def standard_atari_env_spec(env):
@@ -63,24 +55,42 @@ def standard_atari_env_spec(env):
   return tf.contrib.training.HParams(
       env_lambda=env_lambda, wrappers=standard_wrappers, simulated_env=False)
 
+def standard_atari_ae_env_spec(env):
+  """Parameters of environment specification."""
+  standard_wrappers = [[tf_atari_wrappers.StackAndSkipWrapper, {"skip": 4}],
+                       [tf_atari_wrappers.AutoencoderWrapper, {}]]
+  env_lambda = None
+  if isinstance(env, str):
+    env_lambda = lambda: gym.make(env)
+  if callable(env):
+    env_lambda = env
+  assert env is not None, "Unknown specification of environment"
+
+  return tf.contrib.training.HParams(env_lambda=env_lambda,
+                                     wrappers=standard_wrappers,
+                                     simulated_env=False)
+
 
 class GymDiscreteProblem(video_utils.VideoProblem):
   """Gym environment with discrete actions and rewards."""
 
   def __init__(self, *args, **kwargs):
     super(GymDiscreteProblem, self).__init__(*args, **kwargs)
+    #TODO(piotrmilos): Check if self._env is used.
     self._env = None
     self.debug_dump_frames_path = "debug_frames_env"
     self.settable_num_steps = 5000
 
     self.environment_spec = self.get_environment_spec()
-    self.eval_phase = False
+    self.settable_eval_phase = False
 
     self._internal_memory_size = 20
     self._internal_memory_force_beginning_resets = False
     self._session = None
 
   def _setup(self):
+    #TODO(piotrmilos):this should be consistent with
+    # ppo_params in model_rl_experiment
     collect_hparams = rl.ppo_pong_base()
     collect_hparams.add_hparam("environment_spec", self.environment_spec)
     collect_hparams.add_hparam("force_beginning_resets",
@@ -91,14 +101,21 @@ def _setup(self):
     if not FLAGS.agent_policy_path:
       collect_hparams.policy_network = rl.random_policy_fun
 
+    policy_to_actions_lambda = None
+    if self.settable_eval_phase:
+      policy_to_actions_lambda = lambda policy: policy.mode()
+
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       self.collect_memory, self.collect_trigger_op, collect_init \
         = collect.define_collect(collect_hparams, scope="gym_problems",
-                                 collect_level=0, eval_phase=self.eval_phase)
+                                 eval_phase=False, collect_level=0,
+                                 policy_to_actions_lambda
+                                 =policy_to_actions_lambda)
 
     self._session = tf.Session()
     collect_init(self._session)
     self._session.run(tf.global_variables_initializer())
+    self.restore_networks(self._session)
 
   @property
   def random_skip(self):
@@ -110,18 +127,21 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
                                                self.debug_dump_frames_path)
 
     with self._session as sess:
-      self.restore_networks(sess)
-      pieces_generated = 0
       frame_counter = 0
       memory_index = 0
       memory = None
-      while pieces_generated < self.num_steps:
+      pieces_generated = 0
+
+      # TODO(piotrmilos): self.settable_eval_phase possibly violates sematics
+      # of VideoProblem
+      while pieces_generated < self.num_steps or self.settable_eval_phase:
         if memory is None or memory_index >= self._internal_memory_size:
           memory = sess.run(self.collect_memory)
           memory_index = 0
         data = [memory[i][memory_index][0] for i in range(4)]
         memory_index += 1
         observation, reward, done, action = data
+        #TODO(piotrmilos): cleanup types management
         observation = observation.astype(np.uint8)
 
         debug_image = self.collect_statistics_and_generate_debug_image(
@@ -134,13 +154,17 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
             "image/width": [self.frame_width],
             "action": [int(action)],
             "done": [int(done)],
-            "reward": [int(reward) - self.min_reward]
+            "reward": [int(reward - self.min_reward)]
         }
 
         if debug_image is not None:
           ret_dict["image/debug"] = debug_image
 
         yield ret_dict
+
+        if done and self.settable_eval_phase:
+          return
+
         pieces_generated += 1
         frame_counter += 1
         if done:
@@ -164,6 +188,8 @@ def eval_metrics(self):
   @property
   def extra_reading_spec(self):
     """Additional data fields to store on disk and their decoders."""
+
+    #TODO(piotrmilos): shouldn't done be included here?
     data_fields = {
         "frame_number": tf.FixedLenFeature([1], tf.int64),
         "action": tf.FixedLenFeature([1], tf.int64),
@@ -194,6 +220,7 @@ def env_name(self):
 
   @property
   def env(self):
+    #TODO(piotrmilos): possibly remove
     if self._env is None:
       self._env = gym.make(self.env_name)
     return self._env
@@ -241,10 +268,6 @@ def num_testing_steps(self):
   def only_keep_videos_from_0th_frame(self):
     return False
 
-  def get_action(self, observation=None):
-    del observation
-    return self.env.action_space.sample()
-
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.input_modality = {
@@ -261,18 +284,17 @@ def hparams(self, defaults, unused_model_hparams):
     p.target_space_id = problem.SpaceID.IMAGE
 
 
-class GymAEDiscreteProblem(GymDiscreteProblem):
-  pass
-
-
 class BasicStatistics(object):
   """Keeps basic statistics to calculate mean reward """
 
   def __init__(self):
     self.sum_of_rewards = 0.0
     self.number_of_dones = 0
+    self.sum_of_rewards_current_episode = 0.0
+    self.last_done = False
 
 
+#TODO(piotrmilos): merge with the superclass
 class GymRealDiscreteProblem(GymDiscreteProblem):
   """Discrete problem."""
 
@@ -286,14 +308,63 @@ def collect_statistics_and_generate_debug_image(self, index, observation,
                                                   reward, done, action):
     """Collects info required to calculate mean reward."""
 
-    self.statistics.sum_of_rewards += reward
-    self.statistics.number_of_dones += int(done)
+    self.statistics.sum_of_rewards_current_episode += reward
+    # we ignore consecutive dones as they are artefacts of skip wrappers
+    if done and not self.statistics.last_done:
+      self.statistics.number_of_dones += int(done)
+      self.statistics.sum_of_rewards +=\
+        self.statistics.sum_of_rewards_current_episode
+      self.statistics.sum_of_rewards_current_episode = 0.0
 
-    debug_image = None
+    self.statistics.last_done = done
 
+    debug_image = None
     return debug_image
 
 
+class GymDiscreteProblemWithAutoencoder(GymRealDiscreteProblem):
+  def get_environment_spec(self):
+    return standard_atari_ae_env_spec(self.env_name)
+
+
+  def restore_networks(self, sess):
+    super(GymDiscreteProblemWithAutoencoder, self).restore_networks(sess)
+    if FLAGS.autoencoder_path:
+      autoencoder_saver = tf.train.Saver(
+          tf.global_variables("autoencoder.*"))
+      ckpts = tf.train.get_checkpoint_state(FLAGS.autoencoder_path)
+      ckpt = ckpts.model_checkpoint_path
+      autoencoder_saver.restore(sess, ckpt)
+
+class GymDiscreteProblemAutoencoded(GymRealDiscreteProblem):
+
+  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
+    raise RuntimeError("GymDiscreteProblemAutoencoded can be used only"
+                       " for reading encoded frames")
+
+  def get_environment_spec(self):
+    return standard_atari_ae_env_spec(self.env_name)
+
+  @property
+  def autoencoder_factor(self):
+    """By how much to divide sizes when using autoencoders."""
+    hparams = autoencoders.autoencoder_discrete_pong()
+    return 2**hparams.num_hidden_layers
+
+  @property
+  def frame_height(self):
+    height = self.env.observation_space.shape[0]
+    ae_height = int(math.ceil(height / self.autoencoder_factor))
+    return ae_height
+
+  @property
+  def frame_width(self):
+    width = self.env.observation_space.shape[1]
+    return int(math.ceil(width / self.autoencoder_factor))
+
+
+
+
 class RewardPerSequenceStatistics(BasicStatistics):
   """This encapsulates all pieces required to calculate
   the correctness of rewards per sequence metric
@@ -318,12 +389,9 @@ class GymSimulatedDiscreteProblem(GymDiscreteProblem):
   """Simulated gym environment with discrete actions and rewards."""
 
   def __init__(self, *args, **kwargs):
-    self.simulated_environment = True
-    self.debug_dump_frames_path = "debug_frames_sim"
-    self.intrinsic_reward_scale = 0.0
-    self.simulation_random_starts = False
-    self.statistics = RewardPerSequenceStatistics()
     super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
+    self.statistics = RewardPerSequenceStatistics()
+    self.debug_dump_frames_path = "debug_frames_sim"
 
     # This is hackish way of introducing resets every
     # self.num_testing_steps. It cannot be done easily
@@ -355,7 +423,7 @@ def _setup(self):
     self._session.run(input_data_iterator.initializer)
 
     res = self._session.run(input_data_iterator.get_next())
-    self._initial_action = res[0, :, 0][:-1]
+    self._initial_actions = res[0, :, 0][:-1]
     self._reset_real_env()
 
   @property
@@ -383,10 +451,8 @@ def num_testing_steps(self):
   def get_environment_spec(self):
     env_spec = standard_atari_env_spec(self.env_name)
     env_spec.simulated_env = True
-    env_spec.add_hparam("simulation_random_starts",
-                        self.simulation_random_starts)
-
-    env_spec.add_hparam("intrinsic_reward_scale", self.intrinsic_reward_scale)
+    env_spec.add_hparam("simulation_random_starts", False)
+    env_spec.add_hparam("intrinsic_reward_scale", 0.0)
     initial_frames_problem = registry.problem(self.initial_frames_problem)
     env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
     env_spec.add_hparam("video_num_input_frames", self.num_input_frames)
@@ -397,7 +463,7 @@ def get_environment_spec(self):
   def _reset_real_env(self):
     stat = self.statistics
     stat.real_env.reset()
-    for a in self._initial_action:
+    for a in self._initial_actions:
       stat.real_ob, _, _, _ = stat.real_env.step(a)
 
   def collect_statistics_and_generate_debug_image(self, index,
@@ -405,6 +471,8 @@ def collect_statistics_and_generate_debug_image(self, index,
                                                   reward, done, action):
     stat = self.statistics
 
+    # TODO(piotrmilos): possibly make the same behaviour as
+    # in the BasicStatistics
     stat.sum_of_rewards += reward
     stat.episode_sim_reward += reward
 
@@ -438,270 +506,39 @@ def restore_networks(self, sess):
     # TODO(blazej): adjust regexp for different models.
     # TODO(piotrmilos): move restoring networks to SimulatedBatchEnv.initialize
     env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
-    sess = tf.get_default_session()
-
     ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
     ckpt = ckpts.model_checkpoint_path
     env_model_loader.restore(sess, ckpt)
 
 
-@registry.register_problem
-class GymPongRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  # Hard-coding num_actions, frame_height, frame_width to avoid loading
-  # libale.so file.
-  @property
-  def num_actions(self):
-    return 6
-
-  @property
-  def frame_height(self):
-    return 210
-
-  @property
-  def frame_width(self):
-    return 160
-
-  @property
-  def env_name(self):
-    return "PongDeterministic-v4"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-
-@registry.register_problem
-class GymWrappedPongRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  @property
-  def env_name(self):
-    return "T2TPongWarmUp20RewSkip200Steps-v1"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-
-@registry.register_problem
-class GymWrappedLongPongRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  @property
-  def env_name(self):
-    return "T2TPongWarmUp20RewSkip2000Steps-v1"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymWrappedBreakoutRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  @property
-  def env_name(self):
-    return "T2TBreakoutWarmUp20RewSkip500Steps-v1"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnPong(GymSimulatedDiscreteProblem,
-                                                 GymPongRandom):
-  """Simulated pong."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_pong"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymFreewayRandom(GymDiscreteProblem):
-  """Freeway game, random actions."""
-
-  @property
-  def env_name(self):
-    return "FreewayDeterministic-v4"
-
-  @property
-  def min_reward(self):
-    return 0
-
-  @property
-  def num_rewards(self):
-    return 2
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnPong(GymRealDiscreteProblem, GymPongRandom):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
-    GymSimulatedDiscreteProblem, GymWrappedPongRandom):
-  """Similated pong."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_pong"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedLongPong(GymRealDiscreteProblem,
-                                                   GymWrappedLongPongRandom):
-  pass
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedLongPongAe(  # with autoencoder
-    GymDiscreteProblemWithAgentOnWrappedLongPong):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPong(
-    GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
-  """Similated pong."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_long_pong"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedBreakout(GymRealDiscreteProblem,
-                                                   GymWrappedBreakoutRandom):
-  pass
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
-    GymDiscreteProblemWithAgentOnWrappedBreakout):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
-    GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
-  """Similated breakout."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_breakout"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedPong(GymRealDiscreteProblem,
-                                               GymWrappedPongRandom):
-  """GymDiscreteProblemWithAgentOnWrappedPong."""
-
-  # Hard-coding num_actions, frame_height, frame_width to avoid loading
-  # libale.so file.
-  @property
-  def num_actions(self):
-    return 6
-
-  @property
-  def frame_height(self):
-    if not FLAGS.autoencoder_path:
-      return 210
-    return int(math.ceil(210 / self.autoencoder_factor))
-
-  @property
-  def frame_width(self):
-    if not FLAGS.autoencoder_path:
-      return 160
-    return int(math.ceil(160 / self.autoencoder_factor))
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
-    GymDiscreteProblemWithAgentOnWrappedPong):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnFreeway(GymSimulatedDiscreteProblem,
-                                                    GymFreewayRandom):
-  """Similated freeway."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_freeway"
-
-  @property
-  def num_testing_steps(self):
-    return 100
+class GymSimulatedDiscreteProblemAutoencoded(GymSimulatedDiscreteProblem):
+  def get_environment_spec(self):
+    env_spec = standard_atari_env_spec(self.env_name)
+    env_spec.wrappers = [[tf_atari_wrappers.IntToBitWrapper, {}]]
+    env_spec.simulated_env = True
+    env_spec.add_hparam("simulation_random_starts", False)
 
+    env_spec.add_hparam("intrinsic_reward_scale", 0.0)
+    initial_frames_problem = registry.problem(self.initial_frames_problem)
+    env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
+    env_spec.add_hparam("video_num_input_frames", self.num_input_frames)
+    env_spec.add_hparam("video_num_target_frames", self.video_num_target_frames)
 
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnFreeway(GymRealDiscreteProblem,
-                                           GymFreewayRandom):
-  """Freeway with agent."""
+    return env_spec
 
-  # Hard-coding num_actions, frame_height, frame_width to avoid loading
-  # libale.so file.
   @property
-  def num_actions(self):
-    return 3
+  def autoencoder_factor(self):
+    """By how much to divide sizes when using autoencoders."""
+    hparams = autoencoders.autoencoder_discrete_pong()
+    return 2**hparams.num_hidden_layers
 
   @property
   def frame_height(self):
-    if not FLAGS.autoencoder_path:
-      return 210
-    return int(math.ceil(210 / self.autoencoder_factor))
+    height = self.env.observation_space.shape[0]
+    ae_height = int(math.ceil(height / self.autoencoder_factor))
+    return ae_height
 
   @property
   def frame_width(self):
-    if not FLAGS.autoencoder_path:
-      return 160
-    return int(math.ceil(160 / self.autoencoder_factor))
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnFreewayAe(  # with autoencoder
-    GymDiscreteProblemWithAgentOnFreeway):
-  pass
+    width = self.env.observation_space.shape[1]
+    return int(math.ceil(width / self.autoencoder_factor))
diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
new file mode 100644
index 000000000..a04a9a5d6
--- /dev/null
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -0,0 +1,277 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Definitions of data generators for gym problems."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+# We need gym_utils for the game environments defined there.
+from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
+from tensor2tensor.data_generators.gym_problems import GymDiscreteProblem,\
+  GymSimulatedDiscreteProblem, GymRealDiscreteProblem, \
+  GymDiscreteProblemWithAutoencoder, GymDiscreteProblemAutoencoded, \
+  GymSimulatedDiscreteProblemAutoencoded
+from tensor2tensor.utils import registry
+
+
+
+@registry.register_problem
+class GymPongRandom(GymDiscreteProblem):
+  """Pong game, random actions."""
+
+  # Hard-coding num_actions, frame_height, frame_width to avoid loading
+  # libale.so file.
+  @property
+  def num_actions(self):
+    return 6
+
+  @property
+  def frame_height(self):
+    return 210
+
+  @property
+  def frame_width(self):
+    return 160
+
+  @property
+  def env_name(self):
+    return "PongDeterministic-v4"
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+
+@registry.register_problem
+class GymWrappedPongRandom(GymDiscreteProblem):
+  """Pong game, random actions."""
+
+  @property
+  def env_name(self):
+    return "T2TPongWarmUp20RewSkip200Steps-v1"
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+
+@registry.register_problem
+class GymWrappedLongPongRandom(GymDiscreteProblem):
+  """Pong game, random actions."""
+
+  @property
+  def env_name(self):
+    return "T2TPongWarmUp20RewSkip2000Steps-v1"
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymWrappedBreakoutRandom(GymDiscreteProblem):
+  """Pong game, random actions."""
+
+  @property
+  def env_name(self):
+    return "T2TBreakoutWarmUp20RewSkip500Steps-v1"
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnPong(GymSimulatedDiscreteProblem,
+                                                 GymPongRandom):
+  """Simulated pong."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_pong"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymFreewayRandom(GymDiscreteProblem):
+  """Freeway game, random actions."""
+
+  @property
+  def env_name(self):
+    return "FreewayDeterministic-v4"
+
+  @property
+  def min_reward(self):
+    return 0
+
+  @property
+  def num_rewards(self):
+    return 2
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnPong(GymRealDiscreteProblem, GymPongRandom):
+  pass
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
+    GymSimulatedDiscreteProblem, GymWrappedPongRandom):
+  """Similated pong."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_pong"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedLongPong(GymRealDiscreteProblem,
+                                                   GymWrappedLongPongRandom):
+  pass
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedLongPongWithAutoencoder(
+    GymDiscreteProblemWithAutoencoder, GymWrappedLongPongRandom):
+  pass
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedLongPongAutoencoded(
+    GymDiscreteProblemAutoencoded, GymWrappedLongPongRandom):
+  pass
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPong(
+    GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
+  """Simulated pong."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_long_pong"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPongAutoencoded(
+    GymSimulatedDiscreteProblemAutoencoded, GymWrappedLongPongRandom):
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_long_pong_autoencoded"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedBreakout(GymRealDiscreteProblem,
+                                                   GymWrappedBreakoutRandom):
+  pass
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
+    GymDiscreteProblemWithAgentOnWrappedBreakout):
+  pass
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
+    GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
+  """Similated breakout."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_breakout"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedPong(GymRealDiscreteProblem,
+                                               GymWrappedPongRandom):
+  """GymDiscreteProblemWithAgentOnWrappedPong."""
+
+  # Hard-coding num_actions, frame_height, frame_width to avoid loading
+  # libale.so file.
+  @property
+  def num_actions(self):
+    return 6
+
+  @property
+  def frame_height(self):
+    return 210
+
+  @property
+  def frame_width(self):
+    return 160
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
+    GymDiscreteProblemWithAgentOnWrappedPong):
+  pass
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnFreeway(GymSimulatedDiscreteProblem,
+                                                    GymFreewayRandom):
+  """Similated freeway."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_freeway"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index f7aaa8e2e..061a53dc9 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -21,7 +21,7 @@
 import os
 import shutil
 
-from tensor2tensor.data_generators import gym_problems
+from tensor2tensor.data_generators import gym_problems_specs
 
 import tensorflow as tf
 
@@ -35,7 +35,7 @@ def setUpClass(cls):
     os.mkdir(cls.tmp_dir)
 
   def testGymAtariBoots(self):
-    problem = gym_problems.GymPongRandom()
+    problem = gym_problems_specs.GymPongRandom()
     self.assertEqual(210, problem.frame_height)
 
 
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index f62e175d4..de71cde6f 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -50,7 +50,6 @@ def ppo_base_v1():
   hparams.add_hparam("save_models_every_epochs", 30)
   hparams.add_hparam("optimization_batch_size", 50)
   hparams.add_hparam("max_gradients_norm", 0.5)
-  hparams.add_hparam("simulated_environment", False)
   hparams.add_hparam("simulation_random_starts", False)
   hparams.add_hparam("intrinsic_reward_scale", 0.)
   return hparams
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 20d77dc45..9d6a26a36 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -32,6 +32,7 @@ def _rollout_metadata(batch_env):
   batch_env_shape = batch_env.observ.get_shape().as_list()
   batch_size = [batch_env_shape[0]]
   shapes_types_names = [
+      #TODO(piotrmilos): possibly retrive the observation type for batch_env
       (batch_size + batch_env_shape[1:], tf.float32, "observation"),
       (batch_size, tf.float32, "reward"),
       (batch_size, tf.bool, "done"),
@@ -49,10 +50,13 @@ def __init__(self, batch_env):
     super(_MemoryWrapper, self).__init__(batch_env)
     infinity = 10000000
     meta_data = list(zip(*_rollout_metadata(batch_env)))
+    #In memory wrapper we do not collect pdfs neither value_function
+    #thus we only need the first 4 entries of meta_data
     shapes = meta_data[0][:4]
     dtypes = meta_data[1][:4]
     self.speculum = tf.FIFOQueue(infinity, shapes=shapes, dtypes=dtypes)
     observs_shape = batch_env.observ.shape
+    # TODO(piotrmilos): possibly retrive the observation type for batch_env
     observ_dtype = tf.float32
     self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
                                trainable=False)
@@ -79,7 +83,14 @@ def simulate(self, action):
 def define_collect(hparams, scope, eval_phase,
                    collect_level=-1,
                    policy_to_actions_lambda=None):
-  """Collect trajectories."""
+  """ Collect trajectories.
+      Returns memory (observtions, rewards, dones, actions,
+      pdfs, values_functions)
+      containing a rollout of enviroment from collect_level of nested wrapper
+      structure. Note that pdfs and values_functions are meaningful only if
+      collect_level==-1.
+  """
+
   to_initialize = []
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
     batch_env = batch_env_factory(hparams)
@@ -122,13 +133,13 @@ def initialization_lambda(sess):
     force_beginning_resets = hparams.force_beginning_resets
   else:
     force_beginning_resets = False
+  force_beginning_resets = tf.convert_to_tensor(force_beginning_resets)
 
   def group():
     return tf.group(batch_env.reset(tf.range(len(batch_env))),
                     tf.assign(cumulative_rewards, zeros_tensor))
   reset_op = tf.cond(
-      tf.logical_or(should_reset_var, tf.convert_to_tensor(
-          force_beginning_resets)),
+      tf.logical_or(should_reset_var, force_beginning_resets),
       group, tf.no_op)
 
   with tf.control_dependencies([reset_op]):
@@ -144,7 +155,7 @@ def step(index, scores_sum, scores_num):
       # https://github.com/tensorflow/tensorflow/issues/11186
       obs_copy = batch_env.observ + 0
 
-      def env_step(arg1, arg2):  # pylint: disable=unused-argument
+      def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
         """Step of the environment."""
         actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
         policy = actor_critic.policy
@@ -156,35 +167,42 @@ def env_step(arg1, arg2):  # pylint: disable=unused-argument
                            policy.sample)
 
         postprocessed_action = actor_critic.action_postprocessing(action)
-        simulate_output = batch_env.simulate(postprocessed_action[0, ...])
+        reward, done = batch_env.simulate(postprocessed_action[0, ...])
 
         pdf = policy.prob(action)[0]
         value_function = actor_critic.value[0]
         pdf = tf.reshape(pdf, shape=(hparams.num_agents,))
         value_function = tf.reshape(value_function, shape=(hparams.num_agents,))
+        done = tf.reshape(done, shape=(hparams.num_agents,))
 
-        with tf.control_dependencies(simulate_output):
-          return tf.identity(pdf), tf.identity(value_function)
+        with tf.control_dependencies([reward, done]):
+          return tf.identity(pdf), tf.identity(value_function), \
+                 tf.identity(done)
 
-      pdf, value_function = tf.while_loop(
-          lambda _1, _2: tf.equal(speculum.size(), 0),
+      # TODO(piotrmilos): while_body is executed at most once,
+      # thus should be replaced with tf.cond
+      pdf, value_function, top_level_done = tf.while_loop(
+          lambda _1, _2, _3: tf.equal(speculum.size(), 0),
           env_step,
           [tf.constant(0.0, shape=(hparams.num_agents,)),
-           tf.constant(0.0, shape=(hparams.num_agents,))],
+           tf.constant(0.0, shape=(hparams.num_agents,)),
+           tf.constant(False, shape=(hparams.num_agents,))],
           parallel_iterations=1,
           back_prop=False,)
 
       with tf.control_dependencies([pdf, value_function]):
         obs, reward, done, action = speculum.dequeue()
 
-        done = tf.reshape(done, (len(batch_env),))
         to_save = [obs, reward, done, action,
                    pdf, value_function]
         save_ops = [tf.scatter_update(memory_slot, index, value)
                     for memory_slot, value in zip(memory, to_save)]
         cumulate_rewards_op = cumulative_rewards.assign_add(reward)
-        agent_indices_to_reset = tf.where(done)[:, 0]
+
+
+        agent_indices_to_reset = tf.where(top_level_done)[:, 0]
       with tf.control_dependencies([cumulate_rewards_op]):
+        # TODO (piotrmilos): possibly we need cumulative_rewards.read_value()
         scores_sum_delta = tf.reduce_sum(
             tf.gather(cumulative_rewards, agent_indices_to_reset))
         scores_num_delta = tf.count_nonzero(done, dtype=tf.int32)
@@ -211,6 +229,20 @@ def stop_condition(i, _, resets):
         init,
         parallel_iterations=1,
         back_prop=False)
+
+  # We handle force_beginning_resets differently. We assume that all envs are
+  # reseted at the end of episod (though it happens at the beginning of the
+  # next one
+  scores_num = tf.cond(force_beginning_resets,
+                       lambda: scores_num + len(batch_env),
+                       lambda: scores_num)
+
+  with tf.control_dependencies([scores_sum]):
+    scores_sum = tf.cond(force_beginning_resets,
+                         lambda: scores_sum + tf.reduce_sum
+                         (cumulative_rewards.read_value()),
+                         lambda: scores_sum)
+
   mean_score = tf.cond(tf.greater(scores_num, 0),
                        lambda: scores_sum / tf.cast(scores_num, tf.float32),
                        lambda: 0.)
diff --git a/tensor2tensor/rl/envs/batch_env.py b/tensor2tensor/rl/envs/batch_env.py
index f3a72844c..9a0ad136c 100644
--- a/tensor2tensor/rl/envs/batch_env.py
+++ b/tensor2tensor/rl/envs/batch_env.py
@@ -95,6 +95,8 @@ def step(self, actions):
           for env, action in zip(self._envs, actions)]
       transitions = [transition() for transition in transitions]
     observs, rewards, dones, infos = zip(*transitions)
+
+    # TODO(piotrmilos): Do we really want cast to float32
     observ = np.stack(observs).astype(np.float32)
     reward = np.stack(rewards).astype(np.float32)
     done = np.stack(dones)
@@ -118,6 +120,7 @@ def reset(self, indices=None):
       observs = [self._envs[index].reset(blocking=False) for index in indices]
       observs = [observ() for observ in observs]
     observ = np.stack(observs)
+    #TODO(piotrmilos): Do we really want this?
     observ = observ.astype(np.float32)
     return observ
 
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index 26ac6246a..8f610f42a 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -46,9 +46,8 @@ def batch_env_factory(hparams, xvfb=False):
   if environment_spec.simulated_env:
     # TODO(piotrmilos): Consider passing only relevant parameters
     cur_batch_env = _define_simulated_batch_env(
-        environment_spec, hparams.num_agents, hparams)
+        environment_spec, hparams.num_agents)
   else:
-
     cur_batch_env = _define_batch_env(hparams.environment_spec,
                                       hparams.num_agents,
                                       xvfb=xvfb)
@@ -67,11 +66,9 @@ def _define_batch_env(environment_spec, num_agents, xvfb=False):
     return env
 
 
-def _define_simulated_batch_env(environment_spec, num_agents,
-                                other_hparms):
+def _define_simulated_batch_env(environment_spec, num_agents):
   cur_batch_env = simulated_batch_env.SimulatedBatchEnv(environment_spec,
-                                                        num_agents,
-                                                        other_hparms)
+                                                        num_agents)
   return cur_batch_env
 
 
diff --git a/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index 5bc76f2f7..56e13878d 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -76,7 +76,7 @@ def reset(self, indices=None):
   @property
   def observ(self):
     """Access the variable holding the current observation."""
-    return self._observ
+    return self._observ.read_value()
 
   def close(self):
     """Send close messages to the external process and join them."""
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index dc5a0fa79..76916d991 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -117,7 +117,7 @@ def _reset_non_empty(self, indices):
   @property
   def observ(self):
     """Access the variable holding the current observation."""
-    return self._observ
+    return self._observ.read_value()
 
   def close(self):
     """Send close messages to the external process and join them."""
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index abb713dae..f6515becc 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -97,9 +97,9 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   flags are held in according variables.
   """
 
-  def __init__(self, environment_spec, length, other_hparams):
+  def __init__(self, environment_spec, length):
     """Batch of environments inside the TensorFlow graph."""
-    del other_hparams
+
     self.length = length
     initial_frames_problem = environment_spec.initial_frames_problem
     self._min_reward = initial_frames_problem.min_reward
@@ -203,4 +203,4 @@ def _reset_non_empty(self, indices):
   @property
   def observ(self):
     """Access the variable holding the current observation."""
-    return tf.identity(self._observ)
+    return self._observ.read_value()
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 03a64d9ed..0e2a39722 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -17,7 +17,12 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+from tensor2tensor.layers import discretization
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
+from tensor2tensor.models.research import autoencoders
+import math
+
 
 import tensorflow as tf
 
@@ -53,66 +58,15 @@ def _reset_non_empty(self, indices):
       return tf.identity(new_values)
 
 
-class TransformWrapper(WrapperBase):
-  """Transform wrapper."""
-
-  def __init__(self, batch_env, transform_observation=None,
-               transform_reward=tf.identity, transform_done=tf.identity):
-    super(TransformWrapper, self).__init__(batch_env)
-    if transform_observation is not None:
-      _, observ_shape, observ_dtype = transform_observation  # pylint: disable=unpacking-non-sequence
-      self._observ = tf.Variable(
-          tf.zeros(len(self) + observ_shape, observ_dtype), trainable=False)
-    else:
-      self._observ = self._batch_env.observ
-
-    self.transform_observation = transform_observation
-    self.transform_reward = transform_reward
-    self.transform_done = transform_done
-
-  def simulate(self, action):
-    with tf.name_scope("environment/simulate"):  # Do we need this?
-      reward, done = self._batch_env.simulate(action)
-      with tf.control_dependencies([reward]):
-        if self.transform_observation:
-          observ = self.transform_observation[0](self._batch_env.observ)
-          assign_op = self._observ.assign(observ)
-        else:
-          assign_op = tf.no_op()  # TODO(lukaszkaiser): looks as if it's broken.
-        with tf.control_dependencies([assign_op]):
-          return self.transform_reward(reward), self.transform_done(done)
-
-
-class WarpFrameWrapper(TransformWrapper):
-  """Wrap frames."""
-
-  def __init__(self, batch_env):
-    """Warp frames to 84x84 as done in the Nature paper and later work."""
-
-    dims = [84, 84]
-    nature_transform = lambda o: tf.image.rgb_to_grayscale(  # pylint: disable=g-long-lambda
-        tf.image.resize_images(o, dims))
-
-    super(WarpFrameWrapper, self).__init__(batch_env, transform_observation=(
-        nature_transform, dims, tf.float32))
-
-
-class ShiftRewardWrapper(TransformWrapper):
-  """Shift the reward."""
-
-  def __init__(self, batch_env, add_value):
-    shift_reward = lambda r: tf.add(r, add_value)
-    super(ShiftRewardWrapper, self).__init__(
-        batch_env, transform_reward=shift_reward)
-
-
 class MaxAndSkipWrapper(WrapperBase):
-  """Max and skip wrapper."""
+  """ Max and skip wrapper.
+      The wrapper works under assumptions that issuing an action
+      to an environment with done=True has not effect.
+  """
 
   def __init__(self, batch_env, skip=4):
     super(MaxAndSkipWrapper, self).__init__(batch_env)
     self.skip = skip
-    self._observ = None
     observs_shape = batch_env.observ.shape
     observ_dtype = tf.float32
     self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
@@ -126,6 +80,7 @@ def simulate(self, action):
       def not_done_step(a, _):
         reward, done = self._batch_env.simulate(action)
         with tf.control_dependencies([reward, done]):
+          # TODO(piotrmilos): possibly ignore envs with done
           r0 = tf.maximum(a[0], self._batch_env.observ)
           r1 = tf.add(a[1], reward)
           r2 = tf.logical_or(a[2], done)
@@ -142,7 +97,10 @@ def not_done_step(a, _):
 
 
 class StackAndSkipWrapper(WrapperBase):
-  """Stack and skip wrapper."""
+  """ Stack and skip wrapper.
+      The wrapper works under assumptions that issuing an action
+      to an environment with done=True has not effect.
+  """
 
   def __init__(self, batch_env, skip=4):
     super(StackAndSkipWrapper, self).__init__(batch_env)
@@ -162,7 +120,7 @@ def simulate(self, action):
       def not_done_step(a, _):
         reward, done = self._batch_env.simulate(action)
         with tf.control_dependencies([reward, done]):
-          r0 = self._batch_env.observ
+          r0 = self._batch_env.observ + 0
           r1 = tf.add(a[1], reward)
           r2 = tf.logical_or(a[2], done)
           return (r0, r1, r2)
@@ -193,33 +151,82 @@ def _reset_non_empty(self, indices):
       return tf.gather(self.observ, indices)
 
 
-class TimeLimitWrapper(WrapperBase):
-  """Time limit wrapper."""
+class AutoencoderWrapper(WrapperBase):
+  """ Transforms the observations taking the bottleneck
+      state of an autoencoder"""
 
-  # TODO(lukaszkaiser): Check if TimeLimitWrapper does what it's supposed to do.
-  def __init__(self, batch_env, timelimit=100):
-    super(TimeLimitWrapper, self).__init__(batch_env)
-    self.timelimit = timelimit
-    self._time_elapsed = tf.Variable(tf.zeros((len(self),), tf.int32),
-                                     trainable=False)
+  def __init__(self, batch_env):
+    super(AutoencoderWrapper, self).__init__(batch_env)
+    batch_size, height, width, _ = self._batch_env.observ.get_shape().as_list()
+    ae_height = int(math.ceil(height / self.autoencoder_factor))
+    ae_width = int(math.ceil(width / self.autoencoder_factor))
+    ae_channels = 24 #TODO (piotrmilos): make it better
+    observ_shape = (batch_size, ae_height, ae_width, ae_channels)
+    self._observ = self._observ = tf.Variable(
+        tf.zeros(observ_shape, tf.float32), trainable=False)
+    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+      autoencoder_hparams = autoencoders.autoencoder_discrete_pong()
+      self.autoencoder_model = autoencoders.AutoencoderOrderedDiscrete(
+          autoencoder_hparams, tf.estimator.ModeKeys.EVAL)
+    self.autoencoder_model.set_mode(tf.estimator.ModeKeys.EVAL)
+
+  @property
+  def autoencoder_factor(self):
+    """By how much to divide sizes when using autoencoders."""
+    hparams = autoencoders.autoencoder_discrete_pong()
+    return 2**hparams.num_hidden_layers
 
   def simulate(self, action):
-    with tf.name_scope("environment/simulate"):
-      reward, done = self._batch_env.simulate(action)
-      with tf.control_dependencies([reward, done]):
-        new_done = tf.logical_or(done, self._time_elapsed > self.timelimit)
-        inc = self._time_elapsed.assign_add(tf.ones_like(self._time_elapsed))
+    reward, done = self._batch_env.simulate(action)
+    with tf.control_dependencies([reward, done]):
+      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+        ret = self.autoencoder_model.encode(self._batch_env.observ)
+        assign_op = self._observ.assign(ret)
+        with tf.control_dependencies([assign_op]):
+          return tf.identity(reward), tf.identity(done)
+
+  def _reset_non_empty(self, indices):
+    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+      new_values = self._batch_env._reset_non_empty(indices)
+      ret = self.autoencoder_model.encode(new_values)
+      assign_op = tf.scatter_update(self._observ, indices, ret)
+      with tf.control_dependencies([assign_op]):
+        return tf.gather(self.observ, indices)
+
 
-        with tf.control_dependencies([inc]):
-          return tf.identity(reward), tf.identity(new_done)
+class IntToBitWrapper(WrapperBase):
+  """Unpacks the observations from integer values to bit values"""
+
+  def __init__(self, batch_env):
+    super(IntToBitWrapper, self).__init__(batch_env)
+    batch_size, height, width, channels = \
+      self._batch_env.observ.get_shape().as_list()
+    #We treat each channel as 8-bit integer to be expanded to 8 channels
+    self.observ_shape = (height, width, channels*8)
+    self._observ = self._observ = tf.Variable(
+        tf.zeros((batch_size,) + self.observ_shape, tf.float32),
+        trainable=False)
+
+  def simulate(self, action):
+    action = tf.Print(action, [action], message="action=", summarize=200)
+
+    # action = tf.zeros_like(action) #Temporary hacked bugfix
+    reward, done = self._batch_env.simulate(action)
+    with tf.control_dependencies([reward, done]):
+      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+        unpacked = discretization.int_to_bit(self._batch_env.observ, 8)
+        unpacked = tf.reshape(unpacked, (-1,)+self.observ_shape)
+        assign_op = self._observ.assign(unpacked)
+        with tf.control_dependencies([assign_op]):
+          return tf.identity(reward), tf.identity(done)
 
   def _reset_non_empty(self, indices):
-    op_zero = tf.scatter_update(
-        self._time_elapsed, indices,
-        tf.gather(tf.zeros((len(self),), tf.int32), indices))
     # pylint: disable=protected-access
     new_values = self._batch_env._reset_non_empty(indices)
+    new_values_unpacked = discretization.int_to_bit(new_values, 8)
+    new_values_unpacked = tf.reshape(new_values_unpacked, (-1,)
+                                     +self.observ_shape)
     # pylint: enable=protected-access
-    assign_op = tf.scatter_update(self._observ, indices, new_values)
-    with tf.control_dependencies([op_zero, assign_op]):
-      return tf.gather(self.observ, indices)
+    assign_op = tf.scatter_update(self._observ, indices, new_values_unpacked)
+    with tf.control_dependencies([assign_op]):
+      return tf.identity(new_values_unpacked)
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 920be540e..ae1b57435 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -95,14 +95,15 @@ def generate_real_env_data(problem_name, agent_policy_path, hparams, data_dir,
       "problem": problem_name,
       "agent_policy_path": agent_policy_path,
       "autoencoder_path": autoencoder_path,
-      "only_use_ae_for_policy": True,
   }):
     gym_problem = registry.problem(problem_name)
     gym_problem.settable_num_steps = hparams.true_env_generator_num_steps
-    gym_problem.eval_phase = eval_phase
+    gym_problem.settable_eval_phase = eval_phase
     gym_problem.generate_data(data_dir, tmp_dir)
-    mean_reward = gym_problem.statistics.sum_of_rewards / \
-                  (1.0 + gym_problem.statistics.number_of_dones)
+    mean_reward = None
+    if gym_problem.statistics.number_of_dones:
+      mean_reward = (gym_problem.statistics.sum_of_rewards /
+                     gym_problem.statistics.number_of_dones)
 
   return mean_reward
 
@@ -133,8 +134,7 @@ def train_autoencoder(problem_name, data_dir, output_dir, hparams, epoch):
 
 
 def train_agent(problem_name, agent_model_dir,
-                event_dir, world_model_dir, epoch_data_dir, hparams,
-                autoencoder_path=None, epoch=0):
+                event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0):
   """Train the PPO agent in the simulated environment."""
   gym_problem = registry.problem(problem_name)
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
@@ -168,14 +168,12 @@ def train_agent(problem_name, agent_model_dir,
       "hparams_set": hparams.generative_model_params,
       "output_dir": world_model_dir,
       "data_dir": epoch_data_dir,
-      "autoencoder_path": autoencoder_path,
   }):
     rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=epoch)
 
 
 def evaluate_world_model(simulated_problem_name, problem_name, hparams,
-                         world_model_dir, epoch_data_dir, tmp_dir,
-                         autoencoder_path=None):
+                         world_model_dir, epoch_data_dir, tmp_dir):
   """Generate simulated environment data and return reward accuracy."""
   gym_simulated_problem = registry.problem(simulated_problem_name)
   sim_steps = hparams.simulated_env_generator_num_steps
@@ -186,21 +184,19 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
       "hparams_set": hparams.generative_model_params,
       "data_dir": epoch_data_dir,
       "output_dir": world_model_dir,
-      "autoencoder_path": autoencoder_path,
   }):
     gym_simulated_problem.generate_data(epoch_data_dir, tmp_dir)
   n = max(1., gym_simulated_problem.statistics.number_of_dones)
   model_reward_accuracy = (
       gym_simulated_problem.statistics.successful_episode_reward_predictions
       / float(n))
-  old_path = os.path.join(epoch_data_dir, "debug_frames_env")
-  new_path = os.path.join(epoch_data_dir, "debug_frames_env_eval")
+  old_path = os.path.join(epoch_data_dir, "debug_frames_sim")
+  new_path = os.path.join(epoch_data_dir, "debug_frames_sim_eval")
   tf.gfile.Rename(old_path, new_path)
   return model_reward_accuracy
 
 
-def train_world_model(problem_name, data_dir, output_dir, hparams, epoch,
-                      use_autoencoder=False):
+def train_world_model(problem_name, data_dir, output_dir, hparams, epoch):
   """Train the world model on problem_name."""
   train_steps = hparams.model_train_steps * (epoch + 2)
   with temporary_flags({
@@ -211,9 +207,6 @@ def train_world_model(problem_name, data_dir, output_dir, hparams, epoch,
       "hparams_set": hparams.generative_model_params,
       "eval_steps": 100,
       "train_steps": train_steps,
-      # Hack: If training on autoencoded frames, autoencoder_path needs to be
-      # set so that the problem reports the right sizes for frames.
-      "autoencoder_path": "dummy" if use_autoencoder else None,
   }):
     t2t_trainer.main([])
 
@@ -250,13 +243,15 @@ def generator():
       while True:
         try:
           pngs_np, examples_np = sess.run([pngs, examples])
-          rewards_np = [list(el) for el in examples_np["reward"]]
-          actions_np = [list(el) for el in examples_np["action"]]
-          pngs_np = [el for el in pngs_np]
-          for action, reward, png in zip(actions_np, rewards_np, pngs_np):
+          rewards = examples_np["reward"].tolist()
+          actions = examples_np["action"].tolist()
+          frame_numbers = examples_np["frame_number"].tolist()
+          for action, reward, frame_number, png in \
+                  zip(actions, rewards, frame_numbers, pngs_np):
             yield {
                 "action": action,
                 "reward": reward,
+                "frame_number": frame_number,
                 "image/encoded": [png],
                 "image/format": ["png"],
                 "image/height": [encoded_frame_height],
@@ -332,17 +327,22 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   directories = setup_directories(output_dir, subdirectories)
 
   # Problems
-  problem_name = "gym_discrete_problem_with_agent_on_%s" % hparams.game
-  ae_problem_name = problem_name + "_ae"
-  simulated_problem_name = (
-      "gym_simulated_discrete_problem_with_agent_on_%s" % hparams.game)
-  world_model_problem = ae_problem_name if using_autoencoder else problem_name
-  check_problems([problem_name, world_model_problem, simulated_problem_name])
+  if using_autoencoder:
+    problem_name = \
+      "gym_discrete_problem_with_agent_on_%s_with_autoencoder" % hparams.game
+    world_model_problem = \
+      "gym_discrete_problem_with_agent_on_%s_autoencoded" % hparams.game
+    simulated_problem_name = \
+      "gym_simulated_discrete_problem_with_agent_on_%s_autoencoded" \
+      % hparams.game
+  else:
+    problem_name = "gym_discrete_problem_with_agent_on_%s" % hparams.game
+    world_model_problem = problem_name
+    simulated_problem_name = "gym_simulated_discrete_problem_with_agent_on_%s"\
+                             % hparams.game
 
   # Autoencoder model dir
-  autoencoder_model_dir = (FLAGS.autoencoder_path or
-                           directories.get("autoencoder"))
-  FLAGS.autoencoder_path = None
+  autoencoder_model_dir = directories.get("autoencoder")
 
   # Timing log function
   log_relative_time = make_relative_timing_fn()
@@ -357,7 +357,18 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   tf.logging.info("Generating real environment data with random policy")
   mean_reward = generate_real_env_data(
       problem_name, None, hparams, data_dir, directories["tmp"])
-  tf.logging.info("Mean reward (random): %.4f", mean_reward)
+  tf.logging.info("Mean reward (random): {}".format(mean_reward))
+
+
+  eval_metrics_event_dir = os.path.join(directories["world_model"],
+                                        "eval_metrics_event_dir")
+  eval_metrics_writer = tf.summary.FileWriter(eval_metrics_event_dir)
+  model_reward_accuracy_summary = tf.Summary()
+  model_reward_accuracy_summary.value.add(tag='model_reward_accuracy',
+                                          simple_value=None)
+  mean_reward_summary = tf.Summary()
+  mean_reward_summary.value.add(tag='mean_reward',
+                                simple_value=None)
 
   for epoch in range(hparams.epochs):
     is_final_epoch = (epoch + 1) == hparams.epochs
@@ -386,8 +397,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     # Train world model
     log("Training world model")
     train_world_model(world_model_problem, epoch_data_dir,
-                      directories["world_model"], hparams, epoch,
-                      use_autoencoder=using_autoencoder)
+                      directories["world_model"], hparams, epoch)
 
     # Evaluate world model
     model_reward_accuracy = 0.
@@ -396,33 +406,49 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       model_reward_accuracy = evaluate_world_model(
           simulated_problem_name, world_model_problem, hparams,
           directories["world_model"],
-          epoch_data_dir, directories["tmp"],
-          autoencoder_path=autoencoder_model_dir)
+          epoch_data_dir, directories["tmp"])
       log("World model reward accuracy: %.4f", model_reward_accuracy)
 
     # Train PPO
     log("Training PPO")
-    ppo_event_dir = os.path.join(directories["ppo"], str(epoch))
+    ppo_event_dir = os.path.join(directories["world_model"],
+                                 "ppo_summaries", str(epoch))
     ppo_model_dir = directories["ppo"]
     if not hparams.ppo_continue_training:
       ppo_model_dir = ppo_event_dir
     train_agent(simulated_problem_name, ppo_model_dir,
                 ppo_event_dir, directories["world_model"], epoch_data_dir,
-                hparams, autoencoder_path=autoencoder_model_dir, epoch=epoch)
+                hparams, epoch=epoch)
 
     # Collect data from the real environment.
     log("Generating real environment data")
-    if is_final_epoch:
-      epoch_data_dir = os.path.join(epoch_data_dir, "final_eval")
+    eval_data_dir = os.path.join(epoch_data_dir, "eval")
     mean_reward = generate_real_env_data(
-        problem_name, ppo_model_dir, hparams, epoch_data_dir,
+        problem_name, ppo_model_dir, hparams, eval_data_dir,
         directories["tmp"], autoencoder_path=autoencoder_model_dir,
-        eval_phase=is_final_epoch)
-    log("Mean reward during generation: %.4f", mean_reward)
+        eval_phase=True)
+    log("Mean eval reward: {}".format(mean_reward))
+
+    if not is_final_epoch:
+      generation_mean_reward = generate_real_env_data(
+          problem_name, ppo_model_dir, hparams, epoch_data_dir,
+          directories["tmp"], autoencoder_path=autoencoder_model_dir,
+          eval_phase=False)
+      log("Mean reward during generation: {}".format(generation_mean_reward))
 
     # Report metrics.
     eval_metrics = {"model_reward_accuracy": model_reward_accuracy,
                     "mean_reward": mean_reward}
+
+    model_reward_accuracy_summary.value[0].simple_value \
+      = model_reward_accuracy
+
+    mean_reward_summary.value[0].simple_value \
+      = mean_reward
+
+    eval_metrics_writer.add_summary(model_reward_accuracy_summary, epoch)
+    eval_metrics_writer.add_summary(mean_reward_summary, epoch)
+
     epoch_metrics.append(eval_metrics)
     log("Eval metrics: %s", str(eval_metrics))
     if report_fn:
@@ -672,6 +698,7 @@ def rl_modelrl_ae_tiny():
   hparams.ppo_params = "ppo_pong_ae_base"
   hparams.generative_model_params = "next_frame_ae"
   hparams.autoencoder_train_steps = 2
+  hparams.eval_world_model = False
   return hparams
 
 
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index ec8618832..3d908495a 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -28,9 +28,8 @@
 import tensorflow as tf
 
 
-def define_train(hparams, event_dir):
+def define_train(hparams):
   """Define the training setup."""
-  del event_dir
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
     memory, collect_summary, initialization\
       = collect.define_collect(
@@ -45,7 +44,7 @@ def train(hparams, event_dir=None, model_dir=None,
           restore_agent=True, epoch=0):
   """Train."""
   with tf.name_scope("rl_train"):
-    train_summary_op, _, initialization = define_train(hparams, event_dir)
+    train_summary_op, _, initialization = define_train(hparams)
     if event_dir:
       summary_writer = tf.summary.FileWriter(
           event_dir, graph=tf.get_default_graph(), flush_secs=60)
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index 4081e4aa9..7021f416f 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -50,21 +50,15 @@ def test_no_crash_cartpole(self):
     rl_trainer_lib.train(hparams)
 
   # This test should sucessfully train pong.
-  # It should get train mean_score around 0 after 100 epoch
-  #
-  # This test should be run whenever ppo any bigger change
-  # is done on the ppo code
-  #
-  # To run the test change epochs_num=2 to epoch_num=200
-  # and epoch_length=4 to epoch_length=200
-  # (it is set like that to meet travis timeouts
+  # It should get train mean_score around 0 after 200 epoch
+  # By default the test is disabled to avoid travis timeouts
   def test_train_pong(self):
     hparams = tf.contrib.training.\
-      HParams(epochs_num=2,
+      HParams(epochs_num=300,
               eval_every_epochs=10,
-              num_agents=20,
+              num_agents=10,
               optimization_epochs=3,
-              epoch_length=4,
+              epoch_length=200,
               entropy_loss_coef=0.003,
               learning_rate=8e-05,
               optimizer="Adam",

From be120c324508d541d3f655e285e75c87c80b3f3f Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 11 Jul 2018 00:48:49 -0700
Subject: [PATCH 0312/2720] Clean up docstrings in common_layers.

PiperOrigin-RevId: 204080961
---
 tensor2tensor/data_generators/all_problems.py |   2 +-
 tensor2tensor/data_generators/gym_problems.py | 421 ++++++++++++------
 .../data_generators/gym_problems_specs.py     | 277 ------------
 .../data_generators/gym_problems_test.py      |   4 +-
 tensor2tensor/layers/common_layers.py         |  53 +--
 tensor2tensor/models/research/rl.py           |   1 +
 tensor2tensor/rl/collect.py                   |  56 +--
 tensor2tensor/rl/envs/batch_env.py            |   3 -
 tensor2tensor/rl/envs/batch_env_factory.py    |   9 +-
 tensor2tensor/rl/envs/in_graph_batch_env.py   |   2 +-
 tensor2tensor/rl/envs/py_func_batch_env.py    |   2 +-
 tensor2tensor/rl/envs/simulated_batch_env.py  |   6 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py    | 163 ++++---
 tensor2tensor/rl/model_rl_experiment.py       | 111 ++---
 tensor2tensor/rl/rl_trainer_lib.py            |   5 +-
 tensor2tensor/rl/rl_trainer_lib_test.py       |  16 +-
 16 files changed, 481 insertions(+), 650 deletions(-)
 delete mode 100644 tensor2tensor/data_generators/gym_problems_specs.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 36641436a..b065c25f5 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -36,7 +36,7 @@
     "tensor2tensor.data_generators.fsns",
     "tensor2tensor.data_generators.gene_expression",
     "tensor2tensor.data_generators.google_robot_pushing",
-    "tensor2tensor.data_generators.gym_problems_specs",
+    "tensor2tensor.data_generators.gym_problems",
     "tensor2tensor.data_generators.ice_parsing",
     "tensor2tensor.data_generators.imagenet",
     "tensor2tensor.data_generators.image_lsun",
diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index a2edecefb..194f302ae 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -23,23 +23,31 @@
 import gym
 import numpy as np
 
-import tensorflow as tf
-from tensorflow.contrib.training import HParams
+# We need gym_utils for the game environments defined there.
+from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
+
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
-from tensor2tensor.models.research import rl, autoencoders
+from tensor2tensor.models.research import rl
 from tensor2tensor.rl import collect
 from tensor2tensor.rl.envs import tf_atari_wrappers
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
+import tensorflow as tf
+
+from tensorflow.contrib.training import HParams
+
 flags = tf.flags
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string("agent_policy_path", None, "File with model for agent.")
-
 flags.DEFINE_string("autoencoder_path", None,
                     "File with model for autoencoder.")
+flags.DEFINE_boolean(
+    "only_use_ae_for_policy", False,
+    "Whether to only use the autoencoder for the policy and "
+    "still write out full-resolution frames.")
 
 
 def standard_atari_env_spec(env):
@@ -55,42 +63,24 @@ def standard_atari_env_spec(env):
   return tf.contrib.training.HParams(
       env_lambda=env_lambda, wrappers=standard_wrappers, simulated_env=False)
 
-def standard_atari_ae_env_spec(env):
-  """Parameters of environment specification."""
-  standard_wrappers = [[tf_atari_wrappers.StackAndSkipWrapper, {"skip": 4}],
-                       [tf_atari_wrappers.AutoencoderWrapper, {}]]
-  env_lambda = None
-  if isinstance(env, str):
-    env_lambda = lambda: gym.make(env)
-  if callable(env):
-    env_lambda = env
-  assert env is not None, "Unknown specification of environment"
-
-  return tf.contrib.training.HParams(env_lambda=env_lambda,
-                                     wrappers=standard_wrappers,
-                                     simulated_env=False)
-
 
 class GymDiscreteProblem(video_utils.VideoProblem):
   """Gym environment with discrete actions and rewards."""
 
   def __init__(self, *args, **kwargs):
     super(GymDiscreteProblem, self).__init__(*args, **kwargs)
-    #TODO(piotrmilos): Check if self._env is used.
     self._env = None
     self.debug_dump_frames_path = "debug_frames_env"
     self.settable_num_steps = 5000
 
     self.environment_spec = self.get_environment_spec()
-    self.settable_eval_phase = False
+    self.eval_phase = False
 
     self._internal_memory_size = 20
     self._internal_memory_force_beginning_resets = False
     self._session = None
 
   def _setup(self):
-    #TODO(piotrmilos):this should be consistent with
-    # ppo_params in model_rl_experiment
     collect_hparams = rl.ppo_pong_base()
     collect_hparams.add_hparam("environment_spec", self.environment_spec)
     collect_hparams.add_hparam("force_beginning_resets",
@@ -101,21 +91,14 @@ def _setup(self):
     if not FLAGS.agent_policy_path:
       collect_hparams.policy_network = rl.random_policy_fun
 
-    policy_to_actions_lambda = None
-    if self.settable_eval_phase:
-      policy_to_actions_lambda = lambda policy: policy.mode()
-
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       self.collect_memory, self.collect_trigger_op, collect_init \
         = collect.define_collect(collect_hparams, scope="gym_problems",
-                                 eval_phase=False, collect_level=0,
-                                 policy_to_actions_lambda
-                                 =policy_to_actions_lambda)
+                                 collect_level=0, eval_phase=self.eval_phase)
 
     self._session = tf.Session()
     collect_init(self._session)
     self._session.run(tf.global_variables_initializer())
-    self.restore_networks(self._session)
 
   @property
   def random_skip(self):
@@ -127,21 +110,18 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
                                                self.debug_dump_frames_path)
 
     with self._session as sess:
+      self.restore_networks(sess)
+      pieces_generated = 0
       frame_counter = 0
       memory_index = 0
       memory = None
-      pieces_generated = 0
-
-      # TODO(piotrmilos): self.settable_eval_phase possibly violates sematics
-      # of VideoProblem
-      while pieces_generated < self.num_steps or self.settable_eval_phase:
+      while pieces_generated < self.num_steps:
         if memory is None or memory_index >= self._internal_memory_size:
           memory = sess.run(self.collect_memory)
           memory_index = 0
         data = [memory[i][memory_index][0] for i in range(4)]
         memory_index += 1
         observation, reward, done, action = data
-        #TODO(piotrmilos): cleanup types management
         observation = observation.astype(np.uint8)
 
         debug_image = self.collect_statistics_and_generate_debug_image(
@@ -154,17 +134,13 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
             "image/width": [self.frame_width],
             "action": [int(action)],
             "done": [int(done)],
-            "reward": [int(reward - self.min_reward)]
+            "reward": [int(reward) - self.min_reward]
         }
 
         if debug_image is not None:
           ret_dict["image/debug"] = debug_image
 
         yield ret_dict
-
-        if done and self.settable_eval_phase:
-          return
-
         pieces_generated += 1
         frame_counter += 1
         if done:
@@ -188,8 +164,6 @@ def eval_metrics(self):
   @property
   def extra_reading_spec(self):
     """Additional data fields to store on disk and their decoders."""
-
-    #TODO(piotrmilos): shouldn't done be included here?
     data_fields = {
         "frame_number": tf.FixedLenFeature([1], tf.int64),
         "action": tf.FixedLenFeature([1], tf.int64),
@@ -220,7 +194,6 @@ def env_name(self):
 
   @property
   def env(self):
-    #TODO(piotrmilos): possibly remove
     if self._env is None:
       self._env = gym.make(self.env_name)
     return self._env
@@ -268,6 +241,10 @@ def num_testing_steps(self):
   def only_keep_videos_from_0th_frame(self):
     return False
 
+  def get_action(self, observation=None):
+    del observation
+    return self.env.action_space.sample()
+
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.input_modality = {
@@ -284,17 +261,18 @@ def hparams(self, defaults, unused_model_hparams):
     p.target_space_id = problem.SpaceID.IMAGE
 
 
+class GymAEDiscreteProblem(GymDiscreteProblem):
+  pass
+
+
 class BasicStatistics(object):
   """Keeps basic statistics to calculate mean reward """
 
   def __init__(self):
     self.sum_of_rewards = 0.0
     self.number_of_dones = 0
-    self.sum_of_rewards_current_episode = 0.0
-    self.last_done = False
 
 
-#TODO(piotrmilos): merge with the superclass
 class GymRealDiscreteProblem(GymDiscreteProblem):
   """Discrete problem."""
 
@@ -308,61 +286,12 @@ def collect_statistics_and_generate_debug_image(self, index, observation,
                                                   reward, done, action):
     """Collects info required to calculate mean reward."""
 
-    self.statistics.sum_of_rewards_current_episode += reward
-    # we ignore consecutive dones as they are artefacts of skip wrappers
-    if done and not self.statistics.last_done:
-      self.statistics.number_of_dones += int(done)
-      self.statistics.sum_of_rewards +=\
-        self.statistics.sum_of_rewards_current_episode
-      self.statistics.sum_of_rewards_current_episode = 0.0
-
-    self.statistics.last_done = done
+    self.statistics.sum_of_rewards += reward
+    self.statistics.number_of_dones += int(done)
 
     debug_image = None
-    return debug_image
-
-
-class GymDiscreteProblemWithAutoencoder(GymRealDiscreteProblem):
-  def get_environment_spec(self):
-    return standard_atari_ae_env_spec(self.env_name)
-
-
-  def restore_networks(self, sess):
-    super(GymDiscreteProblemWithAutoencoder, self).restore_networks(sess)
-    if FLAGS.autoencoder_path:
-      autoencoder_saver = tf.train.Saver(
-          tf.global_variables("autoencoder.*"))
-      ckpts = tf.train.get_checkpoint_state(FLAGS.autoencoder_path)
-      ckpt = ckpts.model_checkpoint_path
-      autoencoder_saver.restore(sess, ckpt)
-
-class GymDiscreteProblemAutoencoded(GymRealDiscreteProblem):
-
-  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
-    raise RuntimeError("GymDiscreteProblemAutoencoded can be used only"
-                       " for reading encoded frames")
-
-  def get_environment_spec(self):
-    return standard_atari_ae_env_spec(self.env_name)
-
-  @property
-  def autoencoder_factor(self):
-    """By how much to divide sizes when using autoencoders."""
-    hparams = autoencoders.autoencoder_discrete_pong()
-    return 2**hparams.num_hidden_layers
-
-  @property
-  def frame_height(self):
-    height = self.env.observation_space.shape[0]
-    ae_height = int(math.ceil(height / self.autoencoder_factor))
-    return ae_height
-
-  @property
-  def frame_width(self):
-    width = self.env.observation_space.shape[1]
-    return int(math.ceil(width / self.autoencoder_factor))
-
 
+    return debug_image
 
 
 class RewardPerSequenceStatistics(BasicStatistics):
@@ -389,9 +318,12 @@ class GymSimulatedDiscreteProblem(GymDiscreteProblem):
   """Simulated gym environment with discrete actions and rewards."""
 
   def __init__(self, *args, **kwargs):
-    super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
-    self.statistics = RewardPerSequenceStatistics()
+    self.simulated_environment = True
     self.debug_dump_frames_path = "debug_frames_sim"
+    self.intrinsic_reward_scale = 0.0
+    self.simulation_random_starts = False
+    self.statistics = RewardPerSequenceStatistics()
+    super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
 
     # This is hackish way of introducing resets every
     # self.num_testing_steps. It cannot be done easily
@@ -423,7 +355,7 @@ def _setup(self):
     self._session.run(input_data_iterator.initializer)
 
     res = self._session.run(input_data_iterator.get_next())
-    self._initial_actions = res[0, :, 0][:-1]
+    self._initial_action = res[0, :, 0][:-1]
     self._reset_real_env()
 
   @property
@@ -451,8 +383,10 @@ def num_testing_steps(self):
   def get_environment_spec(self):
     env_spec = standard_atari_env_spec(self.env_name)
     env_spec.simulated_env = True
-    env_spec.add_hparam("simulation_random_starts", False)
-    env_spec.add_hparam("intrinsic_reward_scale", 0.0)
+    env_spec.add_hparam("simulation_random_starts",
+                        self.simulation_random_starts)
+
+    env_spec.add_hparam("intrinsic_reward_scale", self.intrinsic_reward_scale)
     initial_frames_problem = registry.problem(self.initial_frames_problem)
     env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
     env_spec.add_hparam("video_num_input_frames", self.num_input_frames)
@@ -463,7 +397,7 @@ def get_environment_spec(self):
   def _reset_real_env(self):
     stat = self.statistics
     stat.real_env.reset()
-    for a in self._initial_actions:
+    for a in self._initial_action:
       stat.real_ob, _, _, _ = stat.real_env.step(a)
 
   def collect_statistics_and_generate_debug_image(self, index,
@@ -471,8 +405,6 @@ def collect_statistics_and_generate_debug_image(self, index,
                                                   reward, done, action):
     stat = self.statistics
 
-    # TODO(piotrmilos): possibly make the same behaviour as
-    # in the BasicStatistics
     stat.sum_of_rewards += reward
     stat.episode_sim_reward += reward
 
@@ -506,39 +438,270 @@ def restore_networks(self, sess):
     # TODO(blazej): adjust regexp for different models.
     # TODO(piotrmilos): move restoring networks to SimulatedBatchEnv.initialize
     env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
+    sess = tf.get_default_session()
+
     ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
     ckpt = ckpts.model_checkpoint_path
     env_model_loader.restore(sess, ckpt)
 
 
-class GymSimulatedDiscreteProblemAutoencoded(GymSimulatedDiscreteProblem):
-  def get_environment_spec(self):
-    env_spec = standard_atari_env_spec(self.env_name)
-    env_spec.wrappers = [[tf_atari_wrappers.IntToBitWrapper, {}]]
-    env_spec.simulated_env = True
-    env_spec.add_hparam("simulation_random_starts", False)
+@registry.register_problem
+class GymPongRandom(GymDiscreteProblem):
+  """Pong game, random actions."""
 
-    env_spec.add_hparam("intrinsic_reward_scale", 0.0)
-    initial_frames_problem = registry.problem(self.initial_frames_problem)
-    env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
-    env_spec.add_hparam("video_num_input_frames", self.num_input_frames)
-    env_spec.add_hparam("video_num_target_frames", self.video_num_target_frames)
+  # Hard-coding num_actions, frame_height, frame_width to avoid loading
+  # libale.so file.
+  @property
+  def num_actions(self):
+    return 6
+
+  @property
+  def frame_height(self):
+    return 210
+
+  @property
+  def frame_width(self):
+    return 160
+
+  @property
+  def env_name(self):
+    return "PongDeterministic-v4"
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+
+@registry.register_problem
+class GymWrappedPongRandom(GymDiscreteProblem):
+  """Pong game, random actions."""
+
+  @property
+  def env_name(self):
+    return "T2TPongWarmUp20RewSkip200Steps-v1"
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+
+@registry.register_problem
+class GymWrappedLongPongRandom(GymDiscreteProblem):
+  """Pong game, random actions."""
+
+  @property
+  def env_name(self):
+    return "T2TPongWarmUp20RewSkip2000Steps-v1"
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymWrappedBreakoutRandom(GymDiscreteProblem):
+  """Pong game, random actions."""
+
+  @property
+  def env_name(self):
+    return "T2TBreakoutWarmUp20RewSkip500Steps-v1"
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnPong(GymSimulatedDiscreteProblem,
+                                                 GymPongRandom):
+  """Simulated pong."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_pong"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymFreewayRandom(GymDiscreteProblem):
+  """Freeway game, random actions."""
+
+  @property
+  def env_name(self):
+    return "FreewayDeterministic-v4"
+
+  @property
+  def min_reward(self):
+    return 0
+
+  @property
+  def num_rewards(self):
+    return 2
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnPong(GymRealDiscreteProblem, GymPongRandom):
+  pass
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
+    GymSimulatedDiscreteProblem, GymWrappedPongRandom):
+  """Similated pong."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_pong"
+
+  @property
+  def num_testing_steps(self):
+    return 100
 
-    return env_spec
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedLongPong(GymRealDiscreteProblem,
+                                                   GymWrappedLongPongRandom):
+  pass
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedLongPongAe(  # with autoencoder
+    GymDiscreteProblemWithAgentOnWrappedLongPong):
+  pass
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPong(
+    GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
+  """Similated pong."""
 
   @property
-  def autoencoder_factor(self):
-    """By how much to divide sizes when using autoencoders."""
-    hparams = autoencoders.autoencoder_discrete_pong()
-    return 2**hparams.num_hidden_layers
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_long_pong"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedBreakout(GymRealDiscreteProblem,
+                                                   GymWrappedBreakoutRandom):
+  pass
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
+    GymDiscreteProblemWithAgentOnWrappedBreakout):
+  pass
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
+    GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
+  """Similated breakout."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_breakout"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedPong(GymRealDiscreteProblem,
+                                               GymWrappedPongRandom):
+  """GymDiscreteProblemWithAgentOnWrappedPong."""
+
+  # Hard-coding num_actions, frame_height, frame_width to avoid loading
+  # libale.so file.
+  @property
+  def num_actions(self):
+    return 6
+
+  @property
+  def frame_height(self):
+    if not FLAGS.autoencoder_path:
+      return 210
+    return int(math.ceil(210 / self.autoencoder_factor))
+
+  @property
+  def frame_width(self):
+    if not FLAGS.autoencoder_path:
+      return 160
+    return int(math.ceil(160 / self.autoencoder_factor))
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
+    GymDiscreteProblemWithAgentOnWrappedPong):
+  pass
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnFreeway(GymSimulatedDiscreteProblem,
+                                                    GymFreewayRandom):
+  """Similated freeway."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_freeway"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnFreeway(GymRealDiscreteProblem,
+                                           GymFreewayRandom):
+  """Freeway with agent."""
+
+  # Hard-coding num_actions, frame_height, frame_width to avoid loading
+  # libale.so file.
+  @property
+  def num_actions(self):
+    return 3
 
   @property
   def frame_height(self):
-    height = self.env.observation_space.shape[0]
-    ae_height = int(math.ceil(height / self.autoencoder_factor))
-    return ae_height
+    if not FLAGS.autoencoder_path:
+      return 210
+    return int(math.ceil(210 / self.autoencoder_factor))
 
   @property
   def frame_width(self):
-    width = self.env.observation_space.shape[1]
-    return int(math.ceil(width / self.autoencoder_factor))
+    if not FLAGS.autoencoder_path:
+      return 160
+    return int(math.ceil(160 / self.autoencoder_factor))
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnFreewayAe(  # with autoencoder
+    GymDiscreteProblemWithAgentOnFreeway):
+  pass
diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
deleted file mode 100644
index a04a9a5d6..000000000
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ /dev/null
@@ -1,277 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-"""Definitions of data generators for gym problems."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-
-# We need gym_utils for the game environments defined there.
-from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
-from tensor2tensor.data_generators.gym_problems import GymDiscreteProblem,\
-  GymSimulatedDiscreteProblem, GymRealDiscreteProblem, \
-  GymDiscreteProblemWithAutoencoder, GymDiscreteProblemAutoencoded, \
-  GymSimulatedDiscreteProblemAutoencoded
-from tensor2tensor.utils import registry
-
-
-
-@registry.register_problem
-class GymPongRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  # Hard-coding num_actions, frame_height, frame_width to avoid loading
-  # libale.so file.
-  @property
-  def num_actions(self):
-    return 6
-
-  @property
-  def frame_height(self):
-    return 210
-
-  @property
-  def frame_width(self):
-    return 160
-
-  @property
-  def env_name(self):
-    return "PongDeterministic-v4"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-
-@registry.register_problem
-class GymWrappedPongRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  @property
-  def env_name(self):
-    return "T2TPongWarmUp20RewSkip200Steps-v1"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-
-@registry.register_problem
-class GymWrappedLongPongRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  @property
-  def env_name(self):
-    return "T2TPongWarmUp20RewSkip2000Steps-v1"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymWrappedBreakoutRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  @property
-  def env_name(self):
-    return "T2TBreakoutWarmUp20RewSkip500Steps-v1"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnPong(GymSimulatedDiscreteProblem,
-                                                 GymPongRandom):
-  """Simulated pong."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_pong"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymFreewayRandom(GymDiscreteProblem):
-  """Freeway game, random actions."""
-
-  @property
-  def env_name(self):
-    return "FreewayDeterministic-v4"
-
-  @property
-  def min_reward(self):
-    return 0
-
-  @property
-  def num_rewards(self):
-    return 2
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnPong(GymRealDiscreteProblem, GymPongRandom):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
-    GymSimulatedDiscreteProblem, GymWrappedPongRandom):
-  """Similated pong."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_pong"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedLongPong(GymRealDiscreteProblem,
-                                                   GymWrappedLongPongRandom):
-  pass
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedLongPongWithAutoencoder(
-    GymDiscreteProblemWithAutoencoder, GymWrappedLongPongRandom):
-  pass
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedLongPongAutoencoded(
-    GymDiscreteProblemAutoencoded, GymWrappedLongPongRandom):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPong(
-    GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
-  """Simulated pong."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_long_pong"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPongAutoencoded(
-    GymSimulatedDiscreteProblemAutoencoded, GymWrappedLongPongRandom):
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_long_pong_autoencoded"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedBreakout(GymRealDiscreteProblem,
-                                                   GymWrappedBreakoutRandom):
-  pass
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
-    GymDiscreteProblemWithAgentOnWrappedBreakout):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
-    GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
-  """Similated breakout."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_breakout"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedPong(GymRealDiscreteProblem,
-                                               GymWrappedPongRandom):
-  """GymDiscreteProblemWithAgentOnWrappedPong."""
-
-  # Hard-coding num_actions, frame_height, frame_width to avoid loading
-  # libale.so file.
-  @property
-  def num_actions(self):
-    return 6
-
-  @property
-  def frame_height(self):
-    return 210
-
-  @property
-  def frame_width(self):
-    return 160
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
-    GymDiscreteProblemWithAgentOnWrappedPong):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnFreeway(GymSimulatedDiscreteProblem,
-                                                    GymFreewayRandom):
-  """Similated freeway."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_freeway"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index 061a53dc9..f7aaa8e2e 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -21,7 +21,7 @@
 import os
 import shutil
 
-from tensor2tensor.data_generators import gym_problems_specs
+from tensor2tensor.data_generators import gym_problems
 
 import tensorflow as tf
 
@@ -35,7 +35,7 @@ def setUpClass(cls):
     os.mkdir(cls.tmp_dir)
 
   def testGymAtariBoots(self):
-    problem = gym_problems_specs.GymPongRandom()
+    problem = gym_problems.GymPongRandom()
     self.assertEqual(210, problem.frame_height)
 
 
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index d76d18d46..0bb726742 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -89,8 +89,9 @@ def dropout_with_broadcast_dims(x, keep_prob, broadcast_dims=None, **kwargs):
     broadcast_dims: an optional list of integers
       the dimensions along which to broadcast the keep/drop flags.
     **kwargs: keyword arguments to tf.nn.dropout other than "noise_shape".
+
   Returns:
-    A Tensor with the same size and shape as x.
+    Tensor of the same shape as x.
   """
   assert "noise_shape" not in kwargs
   if broadcast_dims:
@@ -284,8 +285,9 @@ def cumsum(x, axis=0, exclusive=False):
     x: a Tensor
     axis: an integer
     exclusive: a boolean
+
   Returns:
-    a Tensor with the same shape as x
+    Tensor of the same shape as x.
   """
   if not is_on_tpu():
     return tf.cumsum(x, axis=axis, exclusive=exclusive)
@@ -311,8 +313,9 @@ def dropout_no_scaling(x, keep_prob):
   Args:
     x: a Tensor
     keep_prob: a floating point number
+
   Returns:
-    a Tensor of the same size and shape as x
+    Tensor of the same shape as x.
   """
   if keep_prob == 1.0:
     return x
@@ -1456,7 +1459,7 @@ def maybe_zero_out_padding(inputs, kernel_size, nonpadding_mask):
     nonpadding_mask: a Tensor with shape [batch, length]
 
   Returns:
-    a Tensor with the same shape as inputs
+    Tensor of the same shape as inputs.
   """
   if (kernel_size != 1 and kernel_size != (1, 1) and
       nonpadding_mask is not None):
@@ -2152,8 +2155,8 @@ def smoothing_cross_entropy(logits,
   """Cross entropy with label smoothing to limit over-confidence.
 
   Args:
-    logits: Tensor of size [batch_size, ?, ?, ?, vocab_size]
-    labels: Tensor of size [batch_size, ?, ?, ?]
+    logits: Tensor of shape [batch_size, ?, ?, ?, vocab_size].
+    labels: Tensor of shape [batch_size, ?, ?, ?].
     vocab_size: Tensor representing the size of the vocabulary.
     confidence: Used to determine on and off values for label smoothing.
       If `gaussian` is true, `confidence` is the variance to the Gaussian
@@ -2161,7 +2164,7 @@ def smoothing_cross_entropy(logits,
     gaussian: Uses a Gaussian distribution for label smoothing
 
   Returns:
-
+    Tensor of shape [batch_size, ?, ?, ?].
   """
   with tf.name_scope("smoothing_cross_entropy", values=[logits, labels]):
     # Low confidence is given to all non-true labels, uniformly.
@@ -2200,15 +2203,15 @@ def global_pool_1d(inputs, pooling_type="MAX", mask=None):
   to get a representation of a set.
 
   Args:
-    inputs: A tensor of dimensions batch_size x sequence_length x input_dims
+    inputs: A tensor of shape [batch_size, sequence_length, input_dims]
       containing the sequences of input vectors.
     pooling_type: the pooling type to use, MAX or AVR
-    mask: A tensor of dimensions batch_size x sequence_length containing a
+    mask: A tensor of shape [batch_size, sequence_length] containing a
       mask for the inputs with 1's for existing elements, and 0's elsewhere.
 
   Returns:
-    output: A tensor of dimensions batch_size x input_dims
-      dimension containing the sequences of transformed vectors.
+    A tensor of shape [batch_size, input_dims] containing the sequences of
+    transformed vectors.
   """
   with tf.name_scope("global_pool", values=[inputs]):
     if mask is not None:
@@ -2238,13 +2241,13 @@ def running_global_pool_1d(inputs, pooling_type="MAX"):
   Currently only supports maximum. Equivalent to using a lower triangle bias.
 
   Args:
-    inputs: A tensor of dimensions batch_size x sequence_length x input_dims
+    inputs: A tensor of shape [batch_size, sequence_length, input_dims]
       containing the sequences of input vectors.
     pooling_type: Pooling type to use. Currently only supports 'MAX'.
 
   Returns:
-    output: A tensor of dimensions batch_size x sequence_length x input_dims
-      dimension containing the running 'totals'.
+    A tensor of shape [batch_size, sequence_length, input_dims] containing the
+    running 'totals'.
   """
   del pooling_type
   with tf.name_scope("running_global_pool", values=[inputs]):
@@ -2270,7 +2273,7 @@ def gated_linear_unit_layer(x, name=None):
     name: A string
 
   Returns:
-    x: A tensor
+    A tensor of the same shape as x.
   """
   with tf.variable_scope(name, default_name="glu_layer", values=[x]):
     depth = shape_list(x)[-1]
@@ -2447,17 +2450,17 @@ def linear_set_layer(layer_size,
 
   Args:
     layer_size: Dimension to transform the input vectors to.
-    inputs: A tensor of dimensions batch_size x sequence_length x input_dims
+    inputs: A tensor of shape [batch_size, sequence_length, input_dims]
       containing the sequences of input vectors.
-    context: A tensor of dimensions batch_size x context_dims
-      containing a global statistic about the set.
+    context: A tensor of shape [batch_size, context_dims] containing a global
+      statistic about the set.
     activation_fn: The activation function to use.
     dropout: Dropout probability.
     name: name.
 
   Returns:
-    output: A tensor of dimensions batch_size x sequence_length x output_dims
-      dimension containing the sequences of transformed vectors.
+    Tensor of shape [batch_size, sequence_length, output_dims] containing the
+    sequences of transformed vectors.
   """
   with tf.variable_scope(
       name, default_name="linear_set_layer", values=[inputs]):
@@ -2497,9 +2500,9 @@ def ravanbakhsh_set_layer(layer_size,
 
   Args:
     layer_size: Dimension to transform the input vectors to.
-    inputs: A tensor of dimensions batch_size x sequence_length x vector
+    inputs: A tensor of shape [batch_size, sequence_length, vector]
       containing the sequences of input vectors.
-    mask: A tensor of dimensions batch_size x sequence_length containing a
+    mask: A tensor of shape [batch_size, sequence_length] containing a
       mask for the inputs with 1's for existing elements, and 0's elsewhere.
     sequential: If true, will use a running global pool so each element will
       only depend on those before it. Set true if this layer is being used in
@@ -2509,8 +2512,8 @@ def ravanbakhsh_set_layer(layer_size,
     name: name.
 
   Returns:
-    output: A tensor of dimensions batch_size x sequence_length x vector
-      dimension containing the sequences of transformed vectors.
+    Tensor of shape [batch_size, sequence_length, vector] containing the
+    sequences of transformed vectors.
   """
   del dropout
   with tf.variable_scope(name, "ravanbakhsh_set_layer", [inputs]):
@@ -2587,7 +2590,7 @@ def underlying_variable(t):
     t: a Tensor
 
   Returns:
-    a tf.Varaible object.
+    tf.Variable.
   """
   t = underlying_variable_ref(t)
   assert t is not None
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index de71cde6f..f62e175d4 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -50,6 +50,7 @@ def ppo_base_v1():
   hparams.add_hparam("save_models_every_epochs", 30)
   hparams.add_hparam("optimization_batch_size", 50)
   hparams.add_hparam("max_gradients_norm", 0.5)
+  hparams.add_hparam("simulated_environment", False)
   hparams.add_hparam("simulation_random_starts", False)
   hparams.add_hparam("intrinsic_reward_scale", 0.)
   return hparams
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 9d6a26a36..20d77dc45 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -32,7 +32,6 @@ def _rollout_metadata(batch_env):
   batch_env_shape = batch_env.observ.get_shape().as_list()
   batch_size = [batch_env_shape[0]]
   shapes_types_names = [
-      #TODO(piotrmilos): possibly retrive the observation type for batch_env
       (batch_size + batch_env_shape[1:], tf.float32, "observation"),
       (batch_size, tf.float32, "reward"),
       (batch_size, tf.bool, "done"),
@@ -50,13 +49,10 @@ def __init__(self, batch_env):
     super(_MemoryWrapper, self).__init__(batch_env)
     infinity = 10000000
     meta_data = list(zip(*_rollout_metadata(batch_env)))
-    #In memory wrapper we do not collect pdfs neither value_function
-    #thus we only need the first 4 entries of meta_data
     shapes = meta_data[0][:4]
     dtypes = meta_data[1][:4]
     self.speculum = tf.FIFOQueue(infinity, shapes=shapes, dtypes=dtypes)
     observs_shape = batch_env.observ.shape
-    # TODO(piotrmilos): possibly retrive the observation type for batch_env
     observ_dtype = tf.float32
     self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
                                trainable=False)
@@ -83,14 +79,7 @@ def simulate(self, action):
 def define_collect(hparams, scope, eval_phase,
                    collect_level=-1,
                    policy_to_actions_lambda=None):
-  """ Collect trajectories.
-      Returns memory (observtions, rewards, dones, actions,
-      pdfs, values_functions)
-      containing a rollout of enviroment from collect_level of nested wrapper
-      structure. Note that pdfs and values_functions are meaningful only if
-      collect_level==-1.
-  """
-
+  """Collect trajectories."""
   to_initialize = []
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
     batch_env = batch_env_factory(hparams)
@@ -133,13 +122,13 @@ def initialization_lambda(sess):
     force_beginning_resets = hparams.force_beginning_resets
   else:
     force_beginning_resets = False
-  force_beginning_resets = tf.convert_to_tensor(force_beginning_resets)
 
   def group():
     return tf.group(batch_env.reset(tf.range(len(batch_env))),
                     tf.assign(cumulative_rewards, zeros_tensor))
   reset_op = tf.cond(
-      tf.logical_or(should_reset_var, force_beginning_resets),
+      tf.logical_or(should_reset_var, tf.convert_to_tensor(
+          force_beginning_resets)),
       group, tf.no_op)
 
   with tf.control_dependencies([reset_op]):
@@ -155,7 +144,7 @@ def step(index, scores_sum, scores_num):
       # https://github.com/tensorflow/tensorflow/issues/11186
       obs_copy = batch_env.observ + 0
 
-      def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
+      def env_step(arg1, arg2):  # pylint: disable=unused-argument
         """Step of the environment."""
         actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
         policy = actor_critic.policy
@@ -167,42 +156,35 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
                            policy.sample)
 
         postprocessed_action = actor_critic.action_postprocessing(action)
-        reward, done = batch_env.simulate(postprocessed_action[0, ...])
+        simulate_output = batch_env.simulate(postprocessed_action[0, ...])
 
         pdf = policy.prob(action)[0]
         value_function = actor_critic.value[0]
         pdf = tf.reshape(pdf, shape=(hparams.num_agents,))
         value_function = tf.reshape(value_function, shape=(hparams.num_agents,))
-        done = tf.reshape(done, shape=(hparams.num_agents,))
 
-        with tf.control_dependencies([reward, done]):
-          return tf.identity(pdf), tf.identity(value_function), \
-                 tf.identity(done)
+        with tf.control_dependencies(simulate_output):
+          return tf.identity(pdf), tf.identity(value_function)
 
-      # TODO(piotrmilos): while_body is executed at most once,
-      # thus should be replaced with tf.cond
-      pdf, value_function, top_level_done = tf.while_loop(
-          lambda _1, _2, _3: tf.equal(speculum.size(), 0),
+      pdf, value_function = tf.while_loop(
+          lambda _1, _2: tf.equal(speculum.size(), 0),
           env_step,
           [tf.constant(0.0, shape=(hparams.num_agents,)),
-           tf.constant(0.0, shape=(hparams.num_agents,)),
-           tf.constant(False, shape=(hparams.num_agents,))],
+           tf.constant(0.0, shape=(hparams.num_agents,))],
           parallel_iterations=1,
           back_prop=False,)
 
       with tf.control_dependencies([pdf, value_function]):
         obs, reward, done, action = speculum.dequeue()
 
+        done = tf.reshape(done, (len(batch_env),))
         to_save = [obs, reward, done, action,
                    pdf, value_function]
         save_ops = [tf.scatter_update(memory_slot, index, value)
                     for memory_slot, value in zip(memory, to_save)]
         cumulate_rewards_op = cumulative_rewards.assign_add(reward)
-
-
-        agent_indices_to_reset = tf.where(top_level_done)[:, 0]
+        agent_indices_to_reset = tf.where(done)[:, 0]
       with tf.control_dependencies([cumulate_rewards_op]):
-        # TODO (piotrmilos): possibly we need cumulative_rewards.read_value()
         scores_sum_delta = tf.reduce_sum(
             tf.gather(cumulative_rewards, agent_indices_to_reset))
         scores_num_delta = tf.count_nonzero(done, dtype=tf.int32)
@@ -229,20 +211,6 @@ def stop_condition(i, _, resets):
         init,
         parallel_iterations=1,
         back_prop=False)
-
-  # We handle force_beginning_resets differently. We assume that all envs are
-  # reseted at the end of episod (though it happens at the beginning of the
-  # next one
-  scores_num = tf.cond(force_beginning_resets,
-                       lambda: scores_num + len(batch_env),
-                       lambda: scores_num)
-
-  with tf.control_dependencies([scores_sum]):
-    scores_sum = tf.cond(force_beginning_resets,
-                         lambda: scores_sum + tf.reduce_sum
-                         (cumulative_rewards.read_value()),
-                         lambda: scores_sum)
-
   mean_score = tf.cond(tf.greater(scores_num, 0),
                        lambda: scores_sum / tf.cast(scores_num, tf.float32),
                        lambda: 0.)
diff --git a/tensor2tensor/rl/envs/batch_env.py b/tensor2tensor/rl/envs/batch_env.py
index 9a0ad136c..f3a72844c 100644
--- a/tensor2tensor/rl/envs/batch_env.py
+++ b/tensor2tensor/rl/envs/batch_env.py
@@ -95,8 +95,6 @@ def step(self, actions):
           for env, action in zip(self._envs, actions)]
       transitions = [transition() for transition in transitions]
     observs, rewards, dones, infos = zip(*transitions)
-
-    # TODO(piotrmilos): Do we really want cast to float32
     observ = np.stack(observs).astype(np.float32)
     reward = np.stack(rewards).astype(np.float32)
     done = np.stack(dones)
@@ -120,7 +118,6 @@ def reset(self, indices=None):
       observs = [self._envs[index].reset(blocking=False) for index in indices]
       observs = [observ() for observ in observs]
     observ = np.stack(observs)
-    #TODO(piotrmilos): Do we really want this?
     observ = observ.astype(np.float32)
     return observ
 
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index 8f610f42a..26ac6246a 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -46,8 +46,9 @@ def batch_env_factory(hparams, xvfb=False):
   if environment_spec.simulated_env:
     # TODO(piotrmilos): Consider passing only relevant parameters
     cur_batch_env = _define_simulated_batch_env(
-        environment_spec, hparams.num_agents)
+        environment_spec, hparams.num_agents, hparams)
   else:
+
     cur_batch_env = _define_batch_env(hparams.environment_spec,
                                       hparams.num_agents,
                                       xvfb=xvfb)
@@ -66,9 +67,11 @@ def _define_batch_env(environment_spec, num_agents, xvfb=False):
     return env
 
 
-def _define_simulated_batch_env(environment_spec, num_agents):
+def _define_simulated_batch_env(environment_spec, num_agents,
+                                other_hparms):
   cur_batch_env = simulated_batch_env.SimulatedBatchEnv(environment_spec,
-                                                        num_agents)
+                                                        num_agents,
+                                                        other_hparms)
   return cur_batch_env
 
 
diff --git a/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index 56e13878d..5bc76f2f7 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -76,7 +76,7 @@ def reset(self, indices=None):
   @property
   def observ(self):
     """Access the variable holding the current observation."""
-    return self._observ.read_value()
+    return self._observ
 
   def close(self):
     """Send close messages to the external process and join them."""
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index 76916d991..dc5a0fa79 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -117,7 +117,7 @@ def _reset_non_empty(self, indices):
   @property
   def observ(self):
     """Access the variable holding the current observation."""
-    return self._observ.read_value()
+    return self._observ
 
   def close(self):
     """Send close messages to the external process and join them."""
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index f6515becc..abb713dae 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -97,9 +97,9 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   flags are held in according variables.
   """
 
-  def __init__(self, environment_spec, length):
+  def __init__(self, environment_spec, length, other_hparams):
     """Batch of environments inside the TensorFlow graph."""
-
+    del other_hparams
     self.length = length
     initial_frames_problem = environment_spec.initial_frames_problem
     self._min_reward = initial_frames_problem.min_reward
@@ -203,4 +203,4 @@ def _reset_non_empty(self, indices):
   @property
   def observ(self):
     """Access the variable holding the current observation."""
-    return self._observ.read_value()
+    return tf.identity(self._observ)
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 0e2a39722..03a64d9ed 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -17,12 +17,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
-from tensor2tensor.layers import discretization
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
-from tensor2tensor.models.research import autoencoders
-import math
-
 
 import tensorflow as tf
 
@@ -58,15 +53,66 @@ def _reset_non_empty(self, indices):
       return tf.identity(new_values)
 
 
+class TransformWrapper(WrapperBase):
+  """Transform wrapper."""
+
+  def __init__(self, batch_env, transform_observation=None,
+               transform_reward=tf.identity, transform_done=tf.identity):
+    super(TransformWrapper, self).__init__(batch_env)
+    if transform_observation is not None:
+      _, observ_shape, observ_dtype = transform_observation  # pylint: disable=unpacking-non-sequence
+      self._observ = tf.Variable(
+          tf.zeros(len(self) + observ_shape, observ_dtype), trainable=False)
+    else:
+      self._observ = self._batch_env.observ
+
+    self.transform_observation = transform_observation
+    self.transform_reward = transform_reward
+    self.transform_done = transform_done
+
+  def simulate(self, action):
+    with tf.name_scope("environment/simulate"):  # Do we need this?
+      reward, done = self._batch_env.simulate(action)
+      with tf.control_dependencies([reward]):
+        if self.transform_observation:
+          observ = self.transform_observation[0](self._batch_env.observ)
+          assign_op = self._observ.assign(observ)
+        else:
+          assign_op = tf.no_op()  # TODO(lukaszkaiser): looks as if it's broken.
+        with tf.control_dependencies([assign_op]):
+          return self.transform_reward(reward), self.transform_done(done)
+
+
+class WarpFrameWrapper(TransformWrapper):
+  """Wrap frames."""
+
+  def __init__(self, batch_env):
+    """Warp frames to 84x84 as done in the Nature paper and later work."""
+
+    dims = [84, 84]
+    nature_transform = lambda o: tf.image.rgb_to_grayscale(  # pylint: disable=g-long-lambda
+        tf.image.resize_images(o, dims))
+
+    super(WarpFrameWrapper, self).__init__(batch_env, transform_observation=(
+        nature_transform, dims, tf.float32))
+
+
+class ShiftRewardWrapper(TransformWrapper):
+  """Shift the reward."""
+
+  def __init__(self, batch_env, add_value):
+    shift_reward = lambda r: tf.add(r, add_value)
+    super(ShiftRewardWrapper, self).__init__(
+        batch_env, transform_reward=shift_reward)
+
+
 class MaxAndSkipWrapper(WrapperBase):
-  """ Max and skip wrapper.
-      The wrapper works under assumptions that issuing an action
-      to an environment with done=True has not effect.
-  """
+  """Max and skip wrapper."""
 
   def __init__(self, batch_env, skip=4):
     super(MaxAndSkipWrapper, self).__init__(batch_env)
     self.skip = skip
+    self._observ = None
     observs_shape = batch_env.observ.shape
     observ_dtype = tf.float32
     self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
@@ -80,7 +126,6 @@ def simulate(self, action):
       def not_done_step(a, _):
         reward, done = self._batch_env.simulate(action)
         with tf.control_dependencies([reward, done]):
-          # TODO(piotrmilos): possibly ignore envs with done
           r0 = tf.maximum(a[0], self._batch_env.observ)
           r1 = tf.add(a[1], reward)
           r2 = tf.logical_or(a[2], done)
@@ -97,10 +142,7 @@ def not_done_step(a, _):
 
 
 class StackAndSkipWrapper(WrapperBase):
-  """ Stack and skip wrapper.
-      The wrapper works under assumptions that issuing an action
-      to an environment with done=True has not effect.
-  """
+  """Stack and skip wrapper."""
 
   def __init__(self, batch_env, skip=4):
     super(StackAndSkipWrapper, self).__init__(batch_env)
@@ -120,7 +162,7 @@ def simulate(self, action):
       def not_done_step(a, _):
         reward, done = self._batch_env.simulate(action)
         with tf.control_dependencies([reward, done]):
-          r0 = self._batch_env.observ + 0
+          r0 = self._batch_env.observ
           r1 = tf.add(a[1], reward)
           r2 = tf.logical_or(a[2], done)
           return (r0, r1, r2)
@@ -151,82 +193,33 @@ def _reset_non_empty(self, indices):
       return tf.gather(self.observ, indices)
 
 
-class AutoencoderWrapper(WrapperBase):
-  """ Transforms the observations taking the bottleneck
-      state of an autoencoder"""
+class TimeLimitWrapper(WrapperBase):
+  """Time limit wrapper."""
 
-  def __init__(self, batch_env):
-    super(AutoencoderWrapper, self).__init__(batch_env)
-    batch_size, height, width, _ = self._batch_env.observ.get_shape().as_list()
-    ae_height = int(math.ceil(height / self.autoencoder_factor))
-    ae_width = int(math.ceil(width / self.autoencoder_factor))
-    ae_channels = 24 #TODO (piotrmilos): make it better
-    observ_shape = (batch_size, ae_height, ae_width, ae_channels)
-    self._observ = self._observ = tf.Variable(
-        tf.zeros(observ_shape, tf.float32), trainable=False)
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      autoencoder_hparams = autoencoders.autoencoder_discrete_pong()
-      self.autoencoder_model = autoencoders.AutoencoderOrderedDiscrete(
-          autoencoder_hparams, tf.estimator.ModeKeys.EVAL)
-    self.autoencoder_model.set_mode(tf.estimator.ModeKeys.EVAL)
-
-  @property
-  def autoencoder_factor(self):
-    """By how much to divide sizes when using autoencoders."""
-    hparams = autoencoders.autoencoder_discrete_pong()
-    return 2**hparams.num_hidden_layers
+  # TODO(lukaszkaiser): Check if TimeLimitWrapper does what it's supposed to do.
+  def __init__(self, batch_env, timelimit=100):
+    super(TimeLimitWrapper, self).__init__(batch_env)
+    self.timelimit = timelimit
+    self._time_elapsed = tf.Variable(tf.zeros((len(self),), tf.int32),
+                                     trainable=False)
 
   def simulate(self, action):
-    reward, done = self._batch_env.simulate(action)
-    with tf.control_dependencies([reward, done]):
-      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-        ret = self.autoencoder_model.encode(self._batch_env.observ)
-        assign_op = self._observ.assign(ret)
-        with tf.control_dependencies([assign_op]):
-          return tf.identity(reward), tf.identity(done)
-
-  def _reset_non_empty(self, indices):
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      new_values = self._batch_env._reset_non_empty(indices)
-      ret = self.autoencoder_model.encode(new_values)
-      assign_op = tf.scatter_update(self._observ, indices, ret)
-      with tf.control_dependencies([assign_op]):
-        return tf.gather(self.observ, indices)
-
+    with tf.name_scope("environment/simulate"):
+      reward, done = self._batch_env.simulate(action)
+      with tf.control_dependencies([reward, done]):
+        new_done = tf.logical_or(done, self._time_elapsed > self.timelimit)
+        inc = self._time_elapsed.assign_add(tf.ones_like(self._time_elapsed))
 
-class IntToBitWrapper(WrapperBase):
-  """Unpacks the observations from integer values to bit values"""
-
-  def __init__(self, batch_env):
-    super(IntToBitWrapper, self).__init__(batch_env)
-    batch_size, height, width, channels = \
-      self._batch_env.observ.get_shape().as_list()
-    #We treat each channel as 8-bit integer to be expanded to 8 channels
-    self.observ_shape = (height, width, channels*8)
-    self._observ = self._observ = tf.Variable(
-        tf.zeros((batch_size,) + self.observ_shape, tf.float32),
-        trainable=False)
-
-  def simulate(self, action):
-    action = tf.Print(action, [action], message="action=", summarize=200)
-
-    # action = tf.zeros_like(action) #Temporary hacked bugfix
-    reward, done = self._batch_env.simulate(action)
-    with tf.control_dependencies([reward, done]):
-      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-        unpacked = discretization.int_to_bit(self._batch_env.observ, 8)
-        unpacked = tf.reshape(unpacked, (-1,)+self.observ_shape)
-        assign_op = self._observ.assign(unpacked)
-        with tf.control_dependencies([assign_op]):
-          return tf.identity(reward), tf.identity(done)
+        with tf.control_dependencies([inc]):
+          return tf.identity(reward), tf.identity(new_done)
 
   def _reset_non_empty(self, indices):
+    op_zero = tf.scatter_update(
+        self._time_elapsed, indices,
+        tf.gather(tf.zeros((len(self),), tf.int32), indices))
     # pylint: disable=protected-access
     new_values = self._batch_env._reset_non_empty(indices)
-    new_values_unpacked = discretization.int_to_bit(new_values, 8)
-    new_values_unpacked = tf.reshape(new_values_unpacked, (-1,)
-                                     +self.observ_shape)
     # pylint: enable=protected-access
-    assign_op = tf.scatter_update(self._observ, indices, new_values_unpacked)
-    with tf.control_dependencies([assign_op]):
-      return tf.identity(new_values_unpacked)
+    assign_op = tf.scatter_update(self._observ, indices, new_values)
+    with tf.control_dependencies([op_zero, assign_op]):
+      return tf.gather(self.observ, indices)
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index ae1b57435..920be540e 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -95,15 +95,14 @@ def generate_real_env_data(problem_name, agent_policy_path, hparams, data_dir,
       "problem": problem_name,
       "agent_policy_path": agent_policy_path,
       "autoencoder_path": autoencoder_path,
+      "only_use_ae_for_policy": True,
   }):
     gym_problem = registry.problem(problem_name)
     gym_problem.settable_num_steps = hparams.true_env_generator_num_steps
-    gym_problem.settable_eval_phase = eval_phase
+    gym_problem.eval_phase = eval_phase
     gym_problem.generate_data(data_dir, tmp_dir)
-    mean_reward = None
-    if gym_problem.statistics.number_of_dones:
-      mean_reward = (gym_problem.statistics.sum_of_rewards /
-                     gym_problem.statistics.number_of_dones)
+    mean_reward = gym_problem.statistics.sum_of_rewards / \
+                  (1.0 + gym_problem.statistics.number_of_dones)
 
   return mean_reward
 
@@ -134,7 +133,8 @@ def train_autoencoder(problem_name, data_dir, output_dir, hparams, epoch):
 
 
 def train_agent(problem_name, agent_model_dir,
-                event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0):
+                event_dir, world_model_dir, epoch_data_dir, hparams,
+                autoencoder_path=None, epoch=0):
   """Train the PPO agent in the simulated environment."""
   gym_problem = registry.problem(problem_name)
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
@@ -168,12 +168,14 @@ def train_agent(problem_name, agent_model_dir,
       "hparams_set": hparams.generative_model_params,
       "output_dir": world_model_dir,
       "data_dir": epoch_data_dir,
+      "autoencoder_path": autoencoder_path,
   }):
     rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=epoch)
 
 
 def evaluate_world_model(simulated_problem_name, problem_name, hparams,
-                         world_model_dir, epoch_data_dir, tmp_dir):
+                         world_model_dir, epoch_data_dir, tmp_dir,
+                         autoencoder_path=None):
   """Generate simulated environment data and return reward accuracy."""
   gym_simulated_problem = registry.problem(simulated_problem_name)
   sim_steps = hparams.simulated_env_generator_num_steps
@@ -184,19 +186,21 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
       "hparams_set": hparams.generative_model_params,
       "data_dir": epoch_data_dir,
       "output_dir": world_model_dir,
+      "autoencoder_path": autoencoder_path,
   }):
     gym_simulated_problem.generate_data(epoch_data_dir, tmp_dir)
   n = max(1., gym_simulated_problem.statistics.number_of_dones)
   model_reward_accuracy = (
       gym_simulated_problem.statistics.successful_episode_reward_predictions
       / float(n))
-  old_path = os.path.join(epoch_data_dir, "debug_frames_sim")
-  new_path = os.path.join(epoch_data_dir, "debug_frames_sim_eval")
+  old_path = os.path.join(epoch_data_dir, "debug_frames_env")
+  new_path = os.path.join(epoch_data_dir, "debug_frames_env_eval")
   tf.gfile.Rename(old_path, new_path)
   return model_reward_accuracy
 
 
-def train_world_model(problem_name, data_dir, output_dir, hparams, epoch):
+def train_world_model(problem_name, data_dir, output_dir, hparams, epoch,
+                      use_autoencoder=False):
   """Train the world model on problem_name."""
   train_steps = hparams.model_train_steps * (epoch + 2)
   with temporary_flags({
@@ -207,6 +211,9 @@ def train_world_model(problem_name, data_dir, output_dir, hparams, epoch):
       "hparams_set": hparams.generative_model_params,
       "eval_steps": 100,
       "train_steps": train_steps,
+      # Hack: If training on autoencoded frames, autoencoder_path needs to be
+      # set so that the problem reports the right sizes for frames.
+      "autoencoder_path": "dummy" if use_autoencoder else None,
   }):
     t2t_trainer.main([])
 
@@ -243,15 +250,13 @@ def generator():
       while True:
         try:
           pngs_np, examples_np = sess.run([pngs, examples])
-          rewards = examples_np["reward"].tolist()
-          actions = examples_np["action"].tolist()
-          frame_numbers = examples_np["frame_number"].tolist()
-          for action, reward, frame_number, png in \
-                  zip(actions, rewards, frame_numbers, pngs_np):
+          rewards_np = [list(el) for el in examples_np["reward"]]
+          actions_np = [list(el) for el in examples_np["action"]]
+          pngs_np = [el for el in pngs_np]
+          for action, reward, png in zip(actions_np, rewards_np, pngs_np):
             yield {
                 "action": action,
                 "reward": reward,
-                "frame_number": frame_number,
                 "image/encoded": [png],
                 "image/format": ["png"],
                 "image/height": [encoded_frame_height],
@@ -327,22 +332,17 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   directories = setup_directories(output_dir, subdirectories)
 
   # Problems
-  if using_autoencoder:
-    problem_name = \
-      "gym_discrete_problem_with_agent_on_%s_with_autoencoder" % hparams.game
-    world_model_problem = \
-      "gym_discrete_problem_with_agent_on_%s_autoencoded" % hparams.game
-    simulated_problem_name = \
-      "gym_simulated_discrete_problem_with_agent_on_%s_autoencoded" \
-      % hparams.game
-  else:
-    problem_name = "gym_discrete_problem_with_agent_on_%s" % hparams.game
-    world_model_problem = problem_name
-    simulated_problem_name = "gym_simulated_discrete_problem_with_agent_on_%s"\
-                             % hparams.game
+  problem_name = "gym_discrete_problem_with_agent_on_%s" % hparams.game
+  ae_problem_name = problem_name + "_ae"
+  simulated_problem_name = (
+      "gym_simulated_discrete_problem_with_agent_on_%s" % hparams.game)
+  world_model_problem = ae_problem_name if using_autoencoder else problem_name
+  check_problems([problem_name, world_model_problem, simulated_problem_name])
 
   # Autoencoder model dir
-  autoencoder_model_dir = directories.get("autoencoder")
+  autoencoder_model_dir = (FLAGS.autoencoder_path or
+                           directories.get("autoencoder"))
+  FLAGS.autoencoder_path = None
 
   # Timing log function
   log_relative_time = make_relative_timing_fn()
@@ -357,18 +357,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   tf.logging.info("Generating real environment data with random policy")
   mean_reward = generate_real_env_data(
       problem_name, None, hparams, data_dir, directories["tmp"])
-  tf.logging.info("Mean reward (random): {}".format(mean_reward))
-
-
-  eval_metrics_event_dir = os.path.join(directories["world_model"],
-                                        "eval_metrics_event_dir")
-  eval_metrics_writer = tf.summary.FileWriter(eval_metrics_event_dir)
-  model_reward_accuracy_summary = tf.Summary()
-  model_reward_accuracy_summary.value.add(tag='model_reward_accuracy',
-                                          simple_value=None)
-  mean_reward_summary = tf.Summary()
-  mean_reward_summary.value.add(tag='mean_reward',
-                                simple_value=None)
+  tf.logging.info("Mean reward (random): %.4f", mean_reward)
 
   for epoch in range(hparams.epochs):
     is_final_epoch = (epoch + 1) == hparams.epochs
@@ -397,7 +386,8 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     # Train world model
     log("Training world model")
     train_world_model(world_model_problem, epoch_data_dir,
-                      directories["world_model"], hparams, epoch)
+                      directories["world_model"], hparams, epoch,
+                      use_autoencoder=using_autoencoder)
 
     # Evaluate world model
     model_reward_accuracy = 0.
@@ -406,49 +396,33 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       model_reward_accuracy = evaluate_world_model(
           simulated_problem_name, world_model_problem, hparams,
           directories["world_model"],
-          epoch_data_dir, directories["tmp"])
+          epoch_data_dir, directories["tmp"],
+          autoencoder_path=autoencoder_model_dir)
       log("World model reward accuracy: %.4f", model_reward_accuracy)
 
     # Train PPO
     log("Training PPO")
-    ppo_event_dir = os.path.join(directories["world_model"],
-                                 "ppo_summaries", str(epoch))
+    ppo_event_dir = os.path.join(directories["ppo"], str(epoch))
     ppo_model_dir = directories["ppo"]
     if not hparams.ppo_continue_training:
       ppo_model_dir = ppo_event_dir
     train_agent(simulated_problem_name, ppo_model_dir,
                 ppo_event_dir, directories["world_model"], epoch_data_dir,
-                hparams, epoch=epoch)
+                hparams, autoencoder_path=autoencoder_model_dir, epoch=epoch)
 
     # Collect data from the real environment.
     log("Generating real environment data")
-    eval_data_dir = os.path.join(epoch_data_dir, "eval")
+    if is_final_epoch:
+      epoch_data_dir = os.path.join(epoch_data_dir, "final_eval")
     mean_reward = generate_real_env_data(
-        problem_name, ppo_model_dir, hparams, eval_data_dir,
+        problem_name, ppo_model_dir, hparams, epoch_data_dir,
         directories["tmp"], autoencoder_path=autoencoder_model_dir,
-        eval_phase=True)
-    log("Mean eval reward: {}".format(mean_reward))
-
-    if not is_final_epoch:
-      generation_mean_reward = generate_real_env_data(
-          problem_name, ppo_model_dir, hparams, epoch_data_dir,
-          directories["tmp"], autoencoder_path=autoencoder_model_dir,
-          eval_phase=False)
-      log("Mean reward during generation: {}".format(generation_mean_reward))
+        eval_phase=is_final_epoch)
+    log("Mean reward during generation: %.4f", mean_reward)
 
     # Report metrics.
     eval_metrics = {"model_reward_accuracy": model_reward_accuracy,
                     "mean_reward": mean_reward}
-
-    model_reward_accuracy_summary.value[0].simple_value \
-      = model_reward_accuracy
-
-    mean_reward_summary.value[0].simple_value \
-      = mean_reward
-
-    eval_metrics_writer.add_summary(model_reward_accuracy_summary, epoch)
-    eval_metrics_writer.add_summary(mean_reward_summary, epoch)
-
     epoch_metrics.append(eval_metrics)
     log("Eval metrics: %s", str(eval_metrics))
     if report_fn:
@@ -698,7 +672,6 @@ def rl_modelrl_ae_tiny():
   hparams.ppo_params = "ppo_pong_ae_base"
   hparams.generative_model_params = "next_frame_ae"
   hparams.autoencoder_train_steps = 2
-  hparams.eval_world_model = False
   return hparams
 
 
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 3d908495a..ec8618832 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -28,8 +28,9 @@
 import tensorflow as tf
 
 
-def define_train(hparams):
+def define_train(hparams, event_dir):
   """Define the training setup."""
+  del event_dir
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
     memory, collect_summary, initialization\
       = collect.define_collect(
@@ -44,7 +45,7 @@ def train(hparams, event_dir=None, model_dir=None,
           restore_agent=True, epoch=0):
   """Train."""
   with tf.name_scope("rl_train"):
-    train_summary_op, _, initialization = define_train(hparams)
+    train_summary_op, _, initialization = define_train(hparams, event_dir)
     if event_dir:
       summary_writer = tf.summary.FileWriter(
           event_dir, graph=tf.get_default_graph(), flush_secs=60)
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index 7021f416f..4081e4aa9 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -50,15 +50,21 @@ def test_no_crash_cartpole(self):
     rl_trainer_lib.train(hparams)
 
   # This test should sucessfully train pong.
-  # It should get train mean_score around 0 after 200 epoch
-  # By default the test is disabled to avoid travis timeouts
+  # It should get train mean_score around 0 after 100 epoch
+  #
+  # This test should be run whenever ppo any bigger change
+  # is done on the ppo code
+  #
+  # To run the test change epochs_num=2 to epoch_num=200
+  # and epoch_length=4 to epoch_length=200
+  # (it is set like that to meet travis timeouts
   def test_train_pong(self):
     hparams = tf.contrib.training.\
-      HParams(epochs_num=300,
+      HParams(epochs_num=2,
               eval_every_epochs=10,
-              num_agents=10,
+              num_agents=20,
               optimization_epochs=3,
-              epoch_length=200,
+              epoch_length=4,
               entropy_loss_coef=0.003,
               learning_rate=8e-05,
               optimizer="Adam",

From 13cb58891cb8014546792a6006bc1e1c55c84856 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 11 Jul 2018 01:36:12 -0700
Subject: [PATCH 0313/2720] Clean up implementations in latent_layers.

PiperOrigin-RevId: 204085905
---
 tensor2tensor/layers/latent_layers.py | 100 +++++++++++---------------
 1 file changed, 43 insertions(+), 57 deletions(-)

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index bc13c7f76..24cc150ef 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -30,16 +30,6 @@
 DO_SUMMARIES = True
 
 
-def add_learned_positional_embeddings(x, hparams):
-  pos = tf.get_variable("pos",
-                        [1, hparams.img_len*hparams.img_len,
-                         1, hparams.hidden_size])
-  pos = pos[:, :common_layers.shape_list(x)[1], :, :]
-  x = tf.expand_dims(x, axis=2)
-  x += pos
-  return x
-
-
 def attend(x, source, hparams, name):
   """Attend function."""
   with tf.variable_scope(name):
@@ -101,7 +91,8 @@ def ae_latent_softmax(latents_pred, latents_discrete_hot, hparams):
                                  tf.reduce_mean(tf.square(latents_logits)))
     loss = tf.nn.softmax_cross_entropy_with_logits_v2(
         labels=latents_discrete_hot, logits=latents_logits)
-    sample = multinomial_sample(latents_logits, vocab_size,
+    sample = multinomial_sample(latents_logits,
+                                vocab_size,
                                 hparams.sampling_method,
                                 hparams.sampling_temp)
     return sample, loss
@@ -232,9 +223,11 @@ def compress_encoder(inputs,
     # If using multiple copies of latents, blow up the hidden size and then
     # reshape to increase by num_latents.
     shape_x = common_layers.shape_list(x)
-    x = tf.layers.dense(x, hparams.num_latents*hparams.hidden_size,
+    x = tf.layers.dense(x,
+                        hparams.num_latents * hparams.hidden_size,
                         name=name + "_dense")
-    new_shape = [shape_x[0], shape_x[1] * shape_x[2]*hparams.num_latents,
+    new_shape = [shape_x[0],
+                 shape_x[1] * shape_x[2] * hparams.num_latents,
                  hparams.hidden_size]
     return tf.reshape(x, new_shape)
 
@@ -370,8 +363,9 @@ def transformer_text_encoder(x,
     (encoder_input, encoder_self_attention_bias,
      ed) = transformer.transformer_prepare_encoder(x, space_id, hparams)
     encoder_input = tf.nn.dropout(encoder_input, 1.0 - hparams.dropout)
-    return transformer.transformer_encoder(
-        encoder_input, encoder_self_attention_bias, hparams), ed
+    encoder_output = transformer.transformer_encoder(
+        encoder_input, encoder_self_attention_bias, hparams)
+    return encoder_output, ed
 
 
 def transformer_image_decoder(x,
@@ -396,9 +390,10 @@ def transformer_image_decoder(x,
   """
   with tf.variable_scope(name):
     batch_size = common_layers.shape_list(x)[0]
-    targets = tf.reshape(x, [
-        batch_size, hparams.img_len, hparams.img_len,
-        hparams.num_channels*hparams.hidden_size])
+    targets = tf.reshape(x, [batch_size,
+                             hparams.img_len,
+                             hparams.img_len,
+                             hparams.num_channels * hparams.hidden_size])
     decoder_input, _, _ = cia.prepare_decoder(targets, hparams)
     decoder_output = cia.transformer_decoder_layers(
         decoder_input,
@@ -408,11 +403,11 @@ def transformer_image_decoder(x,
         attention_type=hparams.dec_attention_type,
         encoder_decoder_attention_bias=ed_attention_bias,
         name="decoder")
-    decoder_output_shape = common_layers.shape_list(decoder_output)
-    decoder_output = tf.reshape(decoder_output, [
-        decoder_output_shape[0],
-        hparams.img_len, hparams.img_len*hparams.num_channels,
-        hparams.hidden_size])
+    decoder_output = tf.reshape(decoder_output,
+                                [batch_size,
+                                 hparams.img_len,
+                                 hparams.img_len * hparams.num_channels,
+                                 hparams.hidden_size])
     return decoder_output
 
 
@@ -438,12 +433,11 @@ def transformer_latent_decoder(x,
   """
   with tf.variable_scope(name):
     batch_size = common_layers.shape_list(x)[0]
-    compress_ratio = 2**(hparams.num_compress_steps // 2)
-    x = tf.reshape(x, [
-        batch_size, hparams.img_len / compress_ratio,
-        (hparams.img_len*hparams.num_latents) / compress_ratio,
-        hparams.hidden_size
-    ])
+    compressed_img_len = hparams.img_len / 2**(hparams.num_compress_steps // 2)
+    x = tf.reshape(x, [batch_size,
+                       compressed_img_len,
+                       compressed_img_len * hparams.num_latents,
+                       hparams.hidden_size])
     decoder_input, _, _ = cia.prepare_decoder(x, hparams)
     decoder_output = cia.transformer_decoder_layers(
         decoder_input,
@@ -453,13 +447,10 @@ def transformer_latent_decoder(x,
         attention_type=hparams.latent_attention_type,
         encoder_decoder_attention_bias=ed_attention_bias,
         name="decoder")
-    decoder_output_shape = common_layers.shape_list(decoder_output)
-    decoder_output = tf.reshape(decoder_output, [
-        decoder_output_shape[0],
-        (hparams.img_len * hparams.img_len *
-         hparams.num_latents) / (2**hparams.num_compress_steps),
-        hparams.hidden_size
-    ])
+    decoder_output = tf.reshape(decoder_output,
+                                [batch_size,
+                                 compressed_img_len**2 * hparams.num_latents,
+                                 hparams.hidden_size])
     return decoder_output
 
 
@@ -519,8 +510,7 @@ def transformer_autoencoder(inputs,
                             cache=None,
                             predict_mask=1.0):
   """Auto-encoder using transformer decoder and prior over latents."""
-  # Define losses
-  losses = {"extra": tf.constant(0.0), "latent_pred": tf.constant(0.0)}
+  losses = {"extra": 0., "latent_pred": 0.}
 
   # Reshape image targets as 4d tensor.
   original_targets_shape = common_layers.shape_list(targets)
@@ -532,10 +522,8 @@ def transformer_autoencoder(inputs,
     compress_fn = compress_encoder_1d
     decompress_fn = decompress_decoder_1d
 
-  # Encoder decoder attention bias.
-  ed_attention_bias = None
-
   # Input Encoder if present.
+  ed_attention_bias = None
   if inputs is not None:
     inputs = common_layers.flatten4d3d(inputs)
     inputs, ed_attention_bias = transformer_text_encoder(
@@ -558,7 +546,6 @@ def transformer_autoencoder(inputs,
         bottleneck_layer(targets_c, hparams))
     extra_loss = tf.reduce_mean(extra_loss) * tf.to_float(cond)
 
-    # Call the autoregressive latent prediction model.
     _, latents_pred_loss = latent_prediction_model(
         inputs,
         ed_attention_bias,
@@ -567,14 +554,19 @@ def transformer_autoencoder(inputs,
         hparams,
         name="latent_pred")
     latents_pred_loss = tf.reduce_mean(latents_pred_loss) * tf.to_float(cond)
-    # Latent dropout.
+
     latents_shape = common_layers.shape_list(latents_dense)
     latents_dense = tf.nn.dropout(
         latents_dense, 1 - hparams.latent_dropout,
         noise_shape=[latents_shape[0], latents_shape[1], 1])
-    # Assign latent loss.
-    losses["latent_pred"] = latents_pred_loss
+
     losses["extra_loss"] = extra_loss
+    losses["latent_pred"] = latents_pred_loss
+
+    # We'll start training the extra model of latents after mask_startup_steps.
+    latent_time = tf.less(hparams.mask_startup_steps,
+                          tf.to_int32(tf.train.get_global_step()))
+    losses["latent_pred"] *= tf.to_float(latent_time)
   else:
     latent_len = (
         hparams.img_len * hparams.img_len * hparams.num_latents) / 2**(
@@ -591,13 +583,13 @@ def transformer_autoencoder(inputs,
 
   latents_decoder = latents_dense
   if len(original_targets_shape) == 4:
-    cmp_img_len = hparams.img_len / (2**(hparams.num_compress_steps // 2))
-    latents_decoder = tf.reshape(
-        latents_decoder,
-        [batch_size, cmp_img_len, cmp_img_len,
-         hparams.num_latents*hparams.hidden_size])
+    compressed_img_len = hparams.img_len / 2**(hparams.num_compress_steps // 2)
+    latents_decoder = tf.reshape(latents_decoder,
+                                 [batch_size,
+                                  compressed_img_len,
+                                  compressed_img_len,
+                                  hparams.num_latents * hparams.hidden_size])
 
-  # Decompress either using 1D or 2D upconvs.
   latents_decoder = decompress_fn(latents_decoder, hparams, name="decompress")
   # if we're operating in 2d space on images, then we're assuming that the
   # last dimension will not be a multiple of channels
@@ -606,7 +598,6 @@ def transformer_autoencoder(inputs,
       shape=[-1, hparams.img_len, hparams.img_len, hparams.hidden_size])
 
   if hparams.use_gold_targets:
-    latents_decoder, _, _ = cia.maybe_reshape_4d_to_3d(latents_decoder)
     masking = common_layers.inverse_exp_decay(hparams.mask_startup_steps)
     if hparams.mode == tf.estimator.ModeKeys.PREDICT:
       masking = predict_mask
@@ -623,9 +614,4 @@ def transformer_autoencoder(inputs,
         output, inputs, ed_attention_bias, hparams, "decoder")
   else:
     decoder_output = output
-
-  # We'll start training the extra model of latents after mask_startup_steps.
-  latent_time = tf.less(hparams.mask_startup_steps,
-                        tf.to_int32(tf.train.get_global_step()))
-  losses["latent_pred"] *= tf.to_float(latent_time)
   return decoder_output, losses, cache

From 42f595306fd52df91685e6c34e524c9641b6ae2b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 11 Jul 2018 10:25:10 -0700
Subject: [PATCH 0314/2720] Add Text2textTmpdirTokens to be used with
 user-supplied vocab.

PiperOrigin-RevId: 204146334
---
 .../data_generators/text_problems.py          | 42 +++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index b50a475c5..b67af6702 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -652,6 +652,48 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
     return text2text_txt_iterator(inputs_file, targets_file)
 
 
+@registry.register_problem
+class Text2textTmpdirTokens(Text2textTmpdir):
+  """Allows training a token-based variant of Text2textTmpdir.
+
+  Put your training and evaluation data into the following files in tmp_dir,
+  with 1 record per line along with a vocabulary file with 1 token per line
+  (you can leave out PAD, EOS, and UNK as those will be automatically added)
+
+  * inputs.train.txt
+  * targets.train.txt
+  * inputs.eval.txt
+  * targets.eval.txt
+  * vocab.txt
+  """
+
+  @property
+  def vocab_type(self):
+    return VocabType.TOKEN
+
+  @property
+  def oov_token(self):
+    return "<UNK>"
+
+  def _generate_vocab(self, tmp_dir):
+    vocab_list = [self.oov_token]
+    user_vocab_file = os.path.join(tmp_dir, "vocab.txt")
+    with tf.gfile.GFile(user_vocab_file, "r") as vocab_file:
+      for line in vocab_file:
+        token = line.strip()
+        vocab_list.append(token)
+    token_encoder = text_encoder.TokenTextEncoder(None, vocab_list=vocab_list)
+    return token_encoder
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    vocab_filepath = os.path.join(data_dir, self.vocab_filename)
+    if not tf.gfile.Exists(vocab_filepath):
+      token_encoder = self._generate_vocab(tmp_dir)
+      token_encoder.store_to_file(vocab_filepath)
+    super(Text2textTmpdirTokens, self).generate_samples(data_dir, tmp_dir,
+                                                        dataset_split)
+
+
 class ChoppedTextProblem(Text2SelfProblem):
   """Tokenize and chop text files into fixed-length language-modeling examples.
 

From 4d6ff2f152c5f732a6a061ace13b789e81cdfbda Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Wed, 11 Jul 2018 11:58:55 -0700
Subject: [PATCH 0315/2720] Update the vocabulary size for the Stanford
 Sentiment problem.

PiperOrigin-RevId: 204163583
---
 tensor2tensor/data_generators/sst_binary.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/sst_binary.py b/tensor2tensor/data_generators/sst_binary.py
index 7b9dd585c..ffd934d1f 100644
--- a/tensor2tensor/data_generators/sst_binary.py
+++ b/tensor2tensor/data_generators/sst_binary.py
@@ -57,7 +57,7 @@ def dataset_splits(self):
 
   @property
   def approx_vocab_size(self):
-    return 2**13  # 8k vocab suffices for this small dataset.
+    return 2**14
 
   @property
   def vocab_filename(self):

From 2db7dfce5ec7d4ad2b8a4c7ea8226b05abb8b14b Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Wed, 11 Jul 2018 11:59:15 -0700
Subject: [PATCH 0316/2720] Clean up MultiNLI.

PiperOrigin-RevId: 204163634
---
 tensor2tensor/data_generators/multinli.py | 205 +++++++++-------------
 tensor2tensor/data_generators/problem.py  |   2 +
 2 files changed, 86 insertions(+), 121 deletions(-)

diff --git a/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index 6c50cf9c7..bed35c014 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -12,164 +12,127 @@
 # 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.
-"""Data generators for MultiNLI (https://www.nyu.edu/projects/bowman/multinli/).
-"""
+"""Data generators for MultiNLI."""
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import json
 import os
 import zipfile
+import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
-from tensor2tensor.utils import metrics
+from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
-
 import tensorflow as tf
 
-EOS = text_encoder.EOS_ID
+EOS = text_encoder.EOS
 
 
-class MultinliProblem(problem.Problem):
-  """Base class for MultiNLI classification problems."""
+@registry.register_problem
+class MultiNLI(text_problems.TextConcat2ClassProblem):
+  """MultiNLI classification problems."""
 
-  _ZIP = 'multinli_1.0.zip'
-  _URL = 'https://www.nyu.edu/projects/bowman/multinli/' + _ZIP
-  _LABEL_DICT = {'contradiction': 0,
-                 'entailment': 1,
-                 'neutral': 2}
-  _LABELS = {'contradiction', 'entailment', 'neutral'}
+  # Link to data from GLUE: https://gluebenchmark.com/tasks
+  _MNLI_URL = ("https://firebasestorage.googleapis.com/v0/b/"
+               "mtl-sentence-representations.appspot.com/o/"
+               "data%2FMNLI.zip?alt=media&token=50329ea1-e339-"
+               "40e2-809c-10c40afff3ce")
 
   @property
-  def num_shards(self):
-    return 10
+  def is_generate_per_split(self):
+    return True
 
   @property
-  def vocab_file(self):
-    if self._matched:
-      return 'multinli_matched.vocab'
-    else:
-      return 'multinli_mismatched.vocab'
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 100,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
 
   @property
-  def targeted_vocab_size(self):
-    return 2**14
+  def approx_vocab_size(self):
+    return 2**15
 
   @property
-  def _matched(self):
-    raise NotImplementedError()
+  def vocab_filename(self):
+    return "vocab.mnli.%d" % self.approx_vocab_size
 
   @property
-  def _train_file(self):
-    return 'multinli_1.0/multinli_1.0_train.jsonl'
+  def num_classes(self):
+    return 3
 
   @property
-  def _dev_file(self):
-    if self._matched:
-      return 'multinli_1.0/multinli_1.0_dev_matched.jsonl'
-    else:
-      return 'multinli_1.0/multinli_1.0_dev_mismatched.jsonl'
+  def concat_token(self):
+    return "<EN-PR-HYP>"
 
-  def _examples(self, data_dir, tmp_dir, train):
-    del data_dir
-    file_path = generator_utils.maybe_download(tmp_dir, self._ZIP, self._URL)
-    zip_ref = zipfile.ZipFile(file_path, 'r')
-    zip_ref.extractall(tmp_dir)
-    zip_ref.close()
-
-    data_file = self._train_file if train else self._dev_file
-    examples = []
-    with tf.gfile.GFile(os.path.join(tmp_dir, data_file), mode='r') as f:
-      for line in f:
-        record = json.loads(line)
-        try:
-          label_str = record['gold_label'].encode('ascii')
-          if label_str != '-':
-            label = self._LABEL_DICT[label_str]
-            sentence1 = record['sentence1'].encode('ascii')
-            sentence2 = record['sentence2'].encode('ascii')
-            examples.append({'sentence1': sentence1,
-                             'sentence2': sentence2,
-                             'label': label})
-        except UnicodeEncodeError:
-          pass
-
-    return examples
-
-  def _inputs_and_targets(self, encoder, examples):
-    for e in examples:
-      enc_s1 = encoder.encode(e['sentence1'])
-      enc_s2 = encoder.encode(e['sentence2'])
+  @property
+  def concat_id(self):
+    if self.vocab_type == text_problems.VocabType.CHARACTER:
+      return problem.SpaceID.EN_PR_HYP
+    return 2
 
+  def class_labels(self, data_dir):
+    del data_dir
+    # Note this binary classification is different from usual MNLI.
+    return ["contradiction", "entailment", "neutral"]
+
+  def _maybe_download_corpora(self, tmp_dir):
+    mnli_filename = "MNLI.zip"
+    mnli_finalpath = os.path.join(tmp_dir, "MNLI")
+    if not tf.gfile.Exists(mnli_finalpath):
+      zip_filepath = generator_utils.maybe_download(
+          tmp_dir, mnli_filename, self._MNLI_URL)
+      zip_ref = zipfile.ZipFile(zip_filepath, "r")
+      zip_ref.extractall(tmp_dir)
+      zip_ref.close()
+
+    return mnli_finalpath
+
+  def example_generator(self, filename):
+    label_list = self.class_labels(data_dir=None)
+    for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
+      if idx == 0: continue  # skip header
+      if six.PY2:
+        line = unicode(line.strip(), "utf-8")
+      else:
+        line = line.strip().decode("utf-8")
+      split_line = line.split("\t")
+      # Works for both splits even though dev has some extra human labels.
+      s1, s2 = split_line[8:10]
+      l = label_list.index(split_line[-1])
+      inputs = [s1, s2]
       yield {
-          'inputs': enc_s1 + [EOS] + enc_s2 + [EOS],
-          'targets': [e['label']]
+          "inputs": inputs,
+          "label": l
       }
 
-  def generate_data(self, data_dir, tmp_dir, task_id=-1):
-    train_paths = self.training_filepaths(
-        data_dir, self.num_shards, shuffled=False)
-    dev_paths = self.dev_filepaths(data_dir, 1, shuffled=False)
-
-    train_examples = self._examples(data_dir, tmp_dir, train=True)
-    dev_examples = self._examples(data_dir, tmp_dir, train=False)
-
-    encoder = generator_utils.get_or_generate_vocab_inner(
-        data_dir, self.vocab_file, self.targeted_vocab_size,
-        (e['sentence1'] + ' ' + e['sentence2']
-         for e in train_examples + dev_examples)
-        )
-
-    generator_utils.generate_dataset_and_shuffle(
-        self._inputs_and_targets(encoder, train_examples), train_paths,
-        self._inputs_and_targets(encoder, dev_examples), dev_paths)
-
-  def hparams(self, defaults, unused_model_hparams):
-    p = defaults
-    source_vocab_size = self._encoders['inputs'].vocab_size
-    p.input_modality = {
-        'inputs': (registry.Modalities.SYMBOL, source_vocab_size)
-    }
-    p.target_modality = (registry.Modalities.CLASS_LABEL, 3)
-    p.input_space_id = problem.SpaceID.EN_TOK
-    p.target_space_id = problem.SpaceID.GENERIC
-
-  def feature_encoders(self, data_dir):
-    vocab_filename = os.path.join(data_dir, self.vocab_file)
-    encoder = text_encoder.SubwordTextEncoder(vocab_filename)
-    return {
-        'inputs': encoder,
-        'targets': text_encoder.ClassLabelEncoder(self._LABELS),
-    }
-
-  def example_reading_spec(self):
-    data_fields = {
-        'inputs': tf.VarLenFeature(tf.int64),
-        'targets': tf.FixedLenFeature([1], tf.int64),
-    }
-    data_items_to_decoders = None
-    return (data_fields, data_items_to_decoders)
-
-  def eval_metrics(self):
-    return [metrics.Metrics.ACC]
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    mnli_dir = self._maybe_download_corpora(tmp_dir)
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      filesplit = ["train.tsv"]
+    else:
+      filesplit = ["dev_matched.tsv", "dev_mismatched.tsv"]
+
+    for fs in filesplit:
+      filename = os.path.join(mnli_dir, fs)
+      for example in self.example_generator(filename):
+        yield example
 
 
 @registry.register_problem
-class MultinliMatched(MultinliProblem):
-  """MultiNLI with matched dev set."""
+class MultiNLICharacters(MultiNLI):
+  """MultiNLI classification problems, character level"""
 
   @property
-  def _matched(self):
-    return True
-
-
-@registry.register_problem
-class MultinliMismatched(MultinliProblem):
-  """MultiNLI with mismatched dev set."""
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
 
   @property
-  def _matched(self):
-    return False
+  def task_id(self):
+    return problem.SpaceID.THREE_CL_NLI
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 2d83b25f1..c62496141 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -108,6 +108,8 @@ class SpaceID(object):
   EN_NLI = 33
   # COLA
   COLA = 34
+  # 3 class NLI
+  THREE_CL_NLI = 37
 
 
 def default_model_hparams():

From 40cccaefa79849d3d2c63011d8360005f0d8998c Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 11 Jul 2018 14:28:29 -0700
Subject: [PATCH 0317/2720] internal merge of PR #908

PiperOrigin-RevId: 204189647
---
 tensor2tensor/data_generators/all_problems.py |   2 +-
 tensor2tensor/data_generators/gym_problems.py | 416 ++++++------------
 .../data_generators/gym_problems_specs.py     | 285 ++++++++++++
 .../data_generators/gym_problems_test.py      |   4 +-
 tensor2tensor/models/research/rl.py           |   1 -
 tensor2tensor/rl/collect.py                   |  71 ++-
 tensor2tensor/rl/envs/batch_env.py            |   3 +
 tensor2tensor/rl/envs/batch_env_factory.py    |   9 +-
 tensor2tensor/rl/envs/in_graph_batch_env.py   |   2 +-
 tensor2tensor/rl/envs/py_func_batch_env.py    |   2 +-
 tensor2tensor/rl/envs/simulated_batch_env.py  |   6 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py    | 163 +++----
 tensor2tensor/rl/model_rl_experiment.py       | 113 +++--
 tensor2tensor/rl/rl_trainer_lib.py            |   5 +-
 tensor2tensor/rl/rl_trainer_lib_test.py       |  18 +-
 15 files changed, 648 insertions(+), 452 deletions(-)
 create mode 100644 tensor2tensor/data_generators/gym_problems_specs.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index b065c25f5..36641436a 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -36,7 +36,7 @@
     "tensor2tensor.data_generators.fsns",
     "tensor2tensor.data_generators.gene_expression",
     "tensor2tensor.data_generators.google_robot_pushing",
-    "tensor2tensor.data_generators.gym_problems",
+    "tensor2tensor.data_generators.gym_problems_specs",
     "tensor2tensor.data_generators.ice_parsing",
     "tensor2tensor.data_generators.imagenet",
     "tensor2tensor.data_generators.image_lsun",
diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 194f302ae..32b4ff9e4 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -23,11 +23,9 @@
 import gym
 import numpy as np
 
-# We need gym_utils for the game environments defined there.
-from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
-
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
+from tensor2tensor.models.research import autoencoders
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import collect
 from tensor2tensor.rl.envs import tf_atari_wrappers
@@ -42,12 +40,9 @@
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string("agent_policy_path", None, "File with model for agent.")
+
 flags.DEFINE_string("autoencoder_path", None,
                     "File with model for autoencoder.")
-flags.DEFINE_boolean(
-    "only_use_ae_for_policy", False,
-    "Whether to only use the autoencoder for the policy and "
-    "still write out full-resolution frames.")
 
 
 def standard_atari_env_spec(env):
@@ -64,23 +59,42 @@ def standard_atari_env_spec(env):
       env_lambda=env_lambda, wrappers=standard_wrappers, simulated_env=False)
 
 
+def standard_atari_ae_env_spec(env):
+  """Parameters of environment specification."""
+  standard_wrappers = [[tf_atari_wrappers.StackAndSkipWrapper, {"skip": 4}],
+                       [tf_atari_wrappers.AutoencoderWrapper, {}]]
+  env_lambda = None
+  if isinstance(env, str):
+    env_lambda = lambda: gym.make(env)
+  if callable(env):
+    env_lambda = env
+  assert env is not None, "Unknown specification of environment"
+
+  return tf.contrib.training.HParams(env_lambda=env_lambda,
+                                     wrappers=standard_wrappers,
+                                     simulated_env=False)
+
+
 class GymDiscreteProblem(video_utils.VideoProblem):
   """Gym environment with discrete actions and rewards."""
 
   def __init__(self, *args, **kwargs):
     super(GymDiscreteProblem, self).__init__(*args, **kwargs)
+    # TODO(piotrmilos): Check if self._env is used.
     self._env = None
     self.debug_dump_frames_path = "debug_frames_env"
     self.settable_num_steps = 5000
 
     self.environment_spec = self.get_environment_spec()
-    self.eval_phase = False
+    self.settable_eval_phase = False
 
     self._internal_memory_size = 20
     self._internal_memory_force_beginning_resets = False
     self._session = None
 
   def _setup(self):
+    # TODO(piotrmilos):this should be consistent with
+    # ppo_params in model_rl_experiment
     collect_hparams = rl.ppo_pong_base()
     collect_hparams.add_hparam("environment_spec", self.environment_spec)
     collect_hparams.add_hparam("force_beginning_resets",
@@ -91,14 +105,21 @@ def _setup(self):
     if not FLAGS.agent_policy_path:
       collect_hparams.policy_network = rl.random_policy_fun
 
+    policy_to_actions_lambda = None
+    if self.settable_eval_phase:
+      policy_to_actions_lambda = lambda policy: policy.mode()
+
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       self.collect_memory, self.collect_trigger_op, collect_init \
         = collect.define_collect(collect_hparams, scope="gym_problems",
-                                 collect_level=0, eval_phase=self.eval_phase)
+                                 eval_phase=False, collect_level=0,
+                                 policy_to_actions_lambda
+                                 =policy_to_actions_lambda)
 
     self._session = tf.Session()
     collect_init(self._session)
     self._session.run(tf.global_variables_initializer())
+    self.restore_networks(self._session)
 
   @property
   def random_skip(self):
@@ -110,18 +131,21 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
                                                self.debug_dump_frames_path)
 
     with self._session as sess:
-      self.restore_networks(sess)
-      pieces_generated = 0
       frame_counter = 0
       memory_index = 0
       memory = None
-      while pieces_generated < self.num_steps:
+      pieces_generated = 0
+
+      # TODO(piotrmilos): self.settable_eval_phase possibly violates sematics
+      # of VideoProblem
+      while pieces_generated < self.num_steps or self.settable_eval_phase:
         if memory is None or memory_index >= self._internal_memory_size:
           memory = sess.run(self.collect_memory)
           memory_index = 0
         data = [memory[i][memory_index][0] for i in range(4)]
         memory_index += 1
         observation, reward, done, action = data
+        # TODO(piotrmilos): cleanup types management
         observation = observation.astype(np.uint8)
 
         debug_image = self.collect_statistics_and_generate_debug_image(
@@ -134,13 +158,17 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
             "image/width": [self.frame_width],
             "action": [int(action)],
             "done": [int(done)],
-            "reward": [int(reward) - self.min_reward]
+            "reward": [int(reward - self.min_reward)]
         }
 
         if debug_image is not None:
           ret_dict["image/debug"] = debug_image
 
         yield ret_dict
+
+        if done and self.settable_eval_phase:
+          return
+
         pieces_generated += 1
         frame_counter += 1
         if done:
@@ -164,6 +192,8 @@ def eval_metrics(self):
   @property
   def extra_reading_spec(self):
     """Additional data fields to store on disk and their decoders."""
+
+    # TODO(piotrmilos): shouldn't done be included here?
     data_fields = {
         "frame_number": tf.FixedLenFeature([1], tf.int64),
         "action": tf.FixedLenFeature([1], tf.int64),
@@ -194,6 +224,7 @@ def env_name(self):
 
   @property
   def env(self):
+    # TODO(piotrmilos): possibly remove
     if self._env is None:
       self._env = gym.make(self.env_name)
     return self._env
@@ -241,10 +272,6 @@ def num_testing_steps(self):
   def only_keep_videos_from_0th_frame(self):
     return False
 
-  def get_action(self, observation=None):
-    del observation
-    return self.env.action_space.sample()
-
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.input_modality = {
@@ -261,18 +288,17 @@ def hparams(self, defaults, unused_model_hparams):
     p.target_space_id = problem.SpaceID.IMAGE
 
 
-class GymAEDiscreteProblem(GymDiscreteProblem):
-  pass
-
-
 class BasicStatistics(object):
   """Keeps basic statistics to calculate mean reward """
 
   def __init__(self):
     self.sum_of_rewards = 0.0
     self.number_of_dones = 0
+    self.sum_of_rewards_current_episode = 0.0
+    self.last_done = False
 
 
+# TODO(piotrmilos): merge with the superclass
 class GymRealDiscreteProblem(GymDiscreteProblem):
   """Discrete problem."""
 
@@ -286,14 +312,64 @@ def collect_statistics_and_generate_debug_image(self, index, observation,
                                                   reward, done, action):
     """Collects info required to calculate mean reward."""
 
-    self.statistics.sum_of_rewards += reward
-    self.statistics.number_of_dones += int(done)
+    self.statistics.sum_of_rewards_current_episode += reward
+    # we ignore consecutive dones as they are artefacts of skip wrappers
+    if done and not self.statistics.last_done:
+      self.statistics.number_of_dones += int(done)
+      self.statistics.sum_of_rewards +=\
+        self.statistics.sum_of_rewards_current_episode
+      self.statistics.sum_of_rewards_current_episode = 0.0
 
-    debug_image = None
+    self.statistics.last_done = done
 
+    debug_image = None
     return debug_image
 
 
+class GymDiscreteProblemWithAutoencoder(GymRealDiscreteProblem):
+  """Gym discrete problem with autoencoder."""
+
+  def get_environment_spec(self):
+    return standard_atari_ae_env_spec(self.env_name)
+
+  def restore_networks(self, sess):
+    super(GymDiscreteProblemWithAutoencoder, self).restore_networks(sess)
+    if FLAGS.autoencoder_path:
+      autoencoder_saver = tf.train.Saver(
+          tf.global_variables("autoencoder.*"))
+      ckpts = tf.train.get_checkpoint_state(FLAGS.autoencoder_path)
+      ckpt = ckpts.model_checkpoint_path
+      autoencoder_saver.restore(sess, ckpt)
+
+
+class GymDiscreteProblemAutoencoded(GymRealDiscreteProblem):
+  """Gym discrete problem with frames already autoencoded."""
+
+  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
+    raise RuntimeError("GymDiscreteProblemAutoencoded can be used only"
+                       " for reading encoded frames")
+
+  def get_environment_spec(self):
+    return standard_atari_ae_env_spec(self.env_name)
+
+  @property
+  def autoencoder_factor(self):
+    """By how much to divide sizes when using autoencoders."""
+    hparams = autoencoders.autoencoder_discrete_pong()
+    return 2**hparams.num_hidden_layers
+
+  @property
+  def frame_height(self):
+    height = self.env.observation_space.shape[0]
+    ae_height = int(math.ceil(height / self.autoencoder_factor))
+    return ae_height
+
+  @property
+  def frame_width(self):
+    width = self.env.observation_space.shape[1]
+    return int(math.ceil(width / self.autoencoder_factor))
+
+
 class RewardPerSequenceStatistics(BasicStatistics):
   """This encapsulates all pieces required to calculate
   the correctness of rewards per sequence metric
@@ -318,12 +394,9 @@ class GymSimulatedDiscreteProblem(GymDiscreteProblem):
   """Simulated gym environment with discrete actions and rewards."""
 
   def __init__(self, *args, **kwargs):
-    self.simulated_environment = True
-    self.debug_dump_frames_path = "debug_frames_sim"
-    self.intrinsic_reward_scale = 0.0
-    self.simulation_random_starts = False
-    self.statistics = RewardPerSequenceStatistics()
     super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
+    self.statistics = RewardPerSequenceStatistics()
+    self.debug_dump_frames_path = "debug_frames_sim"
 
     # This is hackish way of introducing resets every
     # self.num_testing_steps. It cannot be done easily
@@ -355,7 +428,7 @@ def _setup(self):
     self._session.run(input_data_iterator.initializer)
 
     res = self._session.run(input_data_iterator.get_next())
-    self._initial_action = res[0, :, 0][:-1]
+    self._initial_actions = res[0, :, 0][:-1]
     self._reset_real_env()
 
   @property
@@ -383,10 +456,8 @@ def num_testing_steps(self):
   def get_environment_spec(self):
     env_spec = standard_atari_env_spec(self.env_name)
     env_spec.simulated_env = True
-    env_spec.add_hparam("simulation_random_starts",
-                        self.simulation_random_starts)
-
-    env_spec.add_hparam("intrinsic_reward_scale", self.intrinsic_reward_scale)
+    env_spec.add_hparam("simulation_random_starts", False)
+    env_spec.add_hparam("intrinsic_reward_scale", 0.0)
     initial_frames_problem = registry.problem(self.initial_frames_problem)
     env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
     env_spec.add_hparam("video_num_input_frames", self.num_input_frames)
@@ -397,7 +468,7 @@ def get_environment_spec(self):
   def _reset_real_env(self):
     stat = self.statistics
     stat.real_env.reset()
-    for a in self._initial_action:
+    for a in self._initial_actions:
       stat.real_ob, _, _, _ = stat.real_env.step(a)
 
   def collect_statistics_and_generate_debug_image(self, index,
@@ -405,6 +476,8 @@ def collect_statistics_and_generate_debug_image(self, index,
                                                   reward, done, action):
     stat = self.statistics
 
+    # TODO(piotrmilos): possibly make the same behaviour as
+    # in the BasicStatistics
     stat.sum_of_rewards += reward
     stat.episode_sim_reward += reward
 
@@ -438,270 +511,41 @@ def restore_networks(self, sess):
     # TODO(blazej): adjust regexp for different models.
     # TODO(piotrmilos): move restoring networks to SimulatedBatchEnv.initialize
     env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
-    sess = tf.get_default_session()
-
     ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
     ckpt = ckpts.model_checkpoint_path
     env_model_loader.restore(sess, ckpt)
 
 
-@registry.register_problem
-class GymPongRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  # Hard-coding num_actions, frame_height, frame_width to avoid loading
-  # libale.so file.
-  @property
-  def num_actions(self):
-    return 6
-
-  @property
-  def frame_height(self):
-    return 210
-
-  @property
-  def frame_width(self):
-    return 160
-
-  @property
-  def env_name(self):
-    return "PongDeterministic-v4"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-
-@registry.register_problem
-class GymWrappedPongRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  @property
-  def env_name(self):
-    return "T2TPongWarmUp20RewSkip200Steps-v1"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-
-@registry.register_problem
-class GymWrappedLongPongRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  @property
-  def env_name(self):
-    return "T2TPongWarmUp20RewSkip2000Steps-v1"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
+class GymSimulatedDiscreteProblemAutoencoded(GymSimulatedDiscreteProblem):
+  """Gym simulated discrete problem with frames already autoencoded."""
 
-@registry.register_problem
-class GymWrappedBreakoutRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  @property
-  def env_name(self):
-    return "T2TBreakoutWarmUp20RewSkip500Steps-v1"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnPong(GymSimulatedDiscreteProblem,
-                                                 GymPongRandom):
-  """Simulated pong."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_pong"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymFreewayRandom(GymDiscreteProblem):
-  """Freeway game, random actions."""
-
-  @property
-  def env_name(self):
-    return "FreewayDeterministic-v4"
-
-  @property
-  def min_reward(self):
-    return 0
-
-  @property
-  def num_rewards(self):
-    return 2
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnPong(GymRealDiscreteProblem, GymPongRandom):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
-    GymSimulatedDiscreteProblem, GymWrappedPongRandom):
-  """Similated pong."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_pong"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedLongPong(GymRealDiscreteProblem,
-                                                   GymWrappedLongPongRandom):
-  pass
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedLongPongAe(  # with autoencoder
-    GymDiscreteProblemWithAgentOnWrappedLongPong):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPong(
-    GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
-  """Similated pong."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_long_pong"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedBreakout(GymRealDiscreteProblem,
-                                                   GymWrappedBreakoutRandom):
-  pass
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
-    GymDiscreteProblemWithAgentOnWrappedBreakout):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
-    GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
-  """Similated breakout."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_breakout"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedPong(GymRealDiscreteProblem,
-                                               GymWrappedPongRandom):
-  """GymDiscreteProblemWithAgentOnWrappedPong."""
-
-  # Hard-coding num_actions, frame_height, frame_width to avoid loading
-  # libale.so file.
-  @property
-  def num_actions(self):
-    return 6
-
-  @property
-  def frame_height(self):
-    if not FLAGS.autoencoder_path:
-      return 210
-    return int(math.ceil(210 / self.autoencoder_factor))
-
-  @property
-  def frame_width(self):
-    if not FLAGS.autoencoder_path:
-      return 160
-    return int(math.ceil(160 / self.autoencoder_factor))
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
-    GymDiscreteProblemWithAgentOnWrappedPong):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnFreeway(GymSimulatedDiscreteProblem,
-                                                    GymFreewayRandom):
-  """Similated freeway."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_freeway"
-
-  @property
-  def num_testing_steps(self):
-    return 100
+  def get_environment_spec(self):
+    env_spec = standard_atari_env_spec(self.env_name)
+    env_spec.wrappers = [[tf_atari_wrappers.IntToBitWrapper, {}]]
+    env_spec.simulated_env = True
+    env_spec.add_hparam("simulation_random_starts", False)
 
+    env_spec.add_hparam("intrinsic_reward_scale", 0.0)
+    initial_frames_problem = registry.problem(self.initial_frames_problem)
+    env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
+    env_spec.add_hparam("video_num_input_frames", self.num_input_frames)
+    env_spec.add_hparam("video_num_target_frames", self.video_num_target_frames)
 
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnFreeway(GymRealDiscreteProblem,
-                                           GymFreewayRandom):
-  """Freeway with agent."""
+    return env_spec
 
-  # Hard-coding num_actions, frame_height, frame_width to avoid loading
-  # libale.so file.
   @property
-  def num_actions(self):
-    return 3
+  def autoencoder_factor(self):
+    """By how much to divide sizes when using autoencoders."""
+    hparams = autoencoders.autoencoder_discrete_pong()
+    return 2**hparams.num_hidden_layers
 
   @property
   def frame_height(self):
-    if not FLAGS.autoencoder_path:
-      return 210
-    return int(math.ceil(210 / self.autoencoder_factor))
+    height = self.env.observation_space.shape[0]
+    ae_height = int(math.ceil(height / self.autoencoder_factor))
+    return ae_height
 
   @property
   def frame_width(self):
-    if not FLAGS.autoencoder_path:
-      return 160
-    return int(math.ceil(160 / self.autoencoder_factor))
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnFreewayAe(  # with autoencoder
-    GymDiscreteProblemWithAgentOnFreeway):
-  pass
+    width = self.env.observation_space.shape[1]
+    return int(math.ceil(width / self.autoencoder_factor))
diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
new file mode 100644
index 000000000..db0383aea
--- /dev/null
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -0,0 +1,285 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Definitions of data generators for gym problems."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+# We need gym_utils for the game environments defined there.
+from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
+# pylint: disable=g-multiple-import
+from tensor2tensor.data_generators.gym_problems import GymDiscreteProblem,\
+  GymSimulatedDiscreteProblem, GymRealDiscreteProblem, \
+  GymDiscreteProblemWithAutoencoder, GymDiscreteProblemAutoencoded, \
+  GymSimulatedDiscreteProblemAutoencoded
+# pylint: enable=g-multiple-import
+from tensor2tensor.utils import registry
+
+
+@registry.register_problem
+class GymPongRandom(GymDiscreteProblem):
+  """Pong game, random actions."""
+
+  # Hard-coding num_actions, frame_height, frame_width to avoid loading
+  # libale.so file.
+  @property
+  def num_actions(self):
+    return 6
+
+  @property
+  def frame_height(self):
+    return 210
+
+  @property
+  def frame_width(self):
+    return 160
+
+  @property
+  def env_name(self):
+    return "PongDeterministic-v4"
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+
+@registry.register_problem
+class GymWrappedPongRandom(GymDiscreteProblem):
+  """Pong game, random actions."""
+
+  @property
+  def env_name(self):
+    return "T2TPongWarmUp20RewSkip200Steps-v1"
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+
+@registry.register_problem
+class GymWrappedLongPongRandom(GymDiscreteProblem):
+  """Pong game, random actions."""
+
+  @property
+  def env_name(self):
+    return "T2TPongWarmUp20RewSkip2000Steps-v1"
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymWrappedBreakoutRandom(GymDiscreteProblem):
+  """Pong game, random actions."""
+
+  @property
+  def env_name(self):
+    return "T2TBreakoutWarmUp20RewSkip500Steps-v1"
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnPong(GymSimulatedDiscreteProblem,
+                                                 GymPongRandom):
+  """Simulated pong."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_pong"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymFreewayRandom(GymDiscreteProblem):
+  """Freeway game, random actions."""
+
+  @property
+  def env_name(self):
+    return "FreewayDeterministic-v4"
+
+  @property
+  def min_reward(self):
+    return 0
+
+  @property
+  def num_rewards(self):
+    return 2
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnPong(GymRealDiscreteProblem, GymPongRandom):
+  pass
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnFreeway(GymRealDiscreteProblem,
+                                           GymFreewayRandom):
+  pass
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
+    GymSimulatedDiscreteProblem, GymWrappedPongRandom):
+  """Similated pong."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_pong"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedLongPong(GymRealDiscreteProblem,
+                                                   GymWrappedLongPongRandom):
+  pass
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedLongPongWithAutoencoder(
+    GymDiscreteProblemWithAutoencoder, GymWrappedLongPongRandom):
+  pass
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedLongPongAutoencoded(
+    GymDiscreteProblemAutoencoded, GymWrappedLongPongRandom):
+  pass
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPong(
+    GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
+  """Simulated pong."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_long_pong"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPongAutoencoded(
+    GymSimulatedDiscreteProblemAutoencoded, GymWrappedLongPongRandom):
+  """GymSimulatedDiscreteProblemWithAgentOnWrappedLongPongAutoencoded."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_long_pong_autoencoded"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedBreakout(GymRealDiscreteProblem,
+                                                   GymWrappedBreakoutRandom):
+  pass
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
+    GymDiscreteProblemWithAgentOnWrappedBreakout):
+  pass
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
+    GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
+  """Similated breakout."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_breakout"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedPong(GymRealDiscreteProblem,
+                                               GymWrappedPongRandom):
+  """GymDiscreteProblemWithAgentOnWrappedPong."""
+
+  # Hard-coding num_actions, frame_height, frame_width to avoid loading
+  # libale.so file.
+  @property
+  def num_actions(self):
+    return 6
+
+  @property
+  def frame_height(self):
+    return 210
+
+  @property
+  def frame_width(self):
+    return 160
+
+
+@registry.register_problem
+class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
+    GymDiscreteProblemWithAgentOnWrappedPong):
+  pass
+
+
+@registry.register_problem
+class GymSimulatedDiscreteProblemWithAgentOnFreeway(GymSimulatedDiscreteProblem,
+                                                    GymFreewayRandom):
+  """Similated freeway."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_freeway"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index f7aaa8e2e..061a53dc9 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -21,7 +21,7 @@
 import os
 import shutil
 
-from tensor2tensor.data_generators import gym_problems
+from tensor2tensor.data_generators import gym_problems_specs
 
 import tensorflow as tf
 
@@ -35,7 +35,7 @@ def setUpClass(cls):
     os.mkdir(cls.tmp_dir)
 
   def testGymAtariBoots(self):
-    problem = gym_problems.GymPongRandom()
+    problem = gym_problems_specs.GymPongRandom()
     self.assertEqual(210, problem.frame_height)
 
 
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index f62e175d4..de71cde6f 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -50,7 +50,6 @@ def ppo_base_v1():
   hparams.add_hparam("save_models_every_epochs", 30)
   hparams.add_hparam("optimization_batch_size", 50)
   hparams.add_hparam("max_gradients_norm", 0.5)
-  hparams.add_hparam("simulated_environment", False)
   hparams.add_hparam("simulation_random_starts", False)
   hparams.add_hparam("intrinsic_reward_scale", 0.)
   return hparams
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 20d77dc45..f2ce0301c 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -32,6 +32,7 @@ def _rollout_metadata(batch_env):
   batch_env_shape = batch_env.observ.get_shape().as_list()
   batch_size = [batch_env_shape[0]]
   shapes_types_names = [
+      # TODO(piotrmilos): possibly retrieve the observation type for batch_env
       (batch_size + batch_env_shape[1:], tf.float32, "observation"),
       (batch_size, tf.float32, "reward"),
       (batch_size, tf.bool, "done"),
@@ -49,10 +50,13 @@ def __init__(self, batch_env):
     super(_MemoryWrapper, self).__init__(batch_env)
     infinity = 10000000
     meta_data = list(zip(*_rollout_metadata(batch_env)))
+    # In memory wrapper we do not collect pdfs neither value_function
+    # thus we only need the first 4 entries of meta_data
     shapes = meta_data[0][:4]
     dtypes = meta_data[1][:4]
     self.speculum = tf.FIFOQueue(infinity, shapes=shapes, dtypes=dtypes)
     observs_shape = batch_env.observ.shape
+    # TODO(piotrmilos): possibly retrieve the observation type for batch_env
     observ_dtype = tf.float32
     self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
                                trainable=False)
@@ -79,7 +83,23 @@ def simulate(self, action):
 def define_collect(hparams, scope, eval_phase,
                    collect_level=-1,
                    policy_to_actions_lambda=None):
-  """Collect trajectories."""
+  """Collect trajectories.
+
+  Args:
+    hparams: HParams.
+    scope: var scope.
+    eval_phase: bool, is eval phase.
+    collect_level: int, which level to collect observations.
+    policy_to_actions_lambda: lambda.
+
+  Returns:
+    Returns memory (observtions, rewards, dones, actions,
+    pdfs, values_functions)
+    containing a rollout of environment from collect_level of nested wrapper
+    structure. Note that pdfs and values_functions are meaningful only if
+    collect_level==-1.
+  """
+
   to_initialize = []
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
     batch_env = batch_env_factory(hparams)
@@ -122,14 +142,13 @@ def initialization_lambda(sess):
     force_beginning_resets = hparams.force_beginning_resets
   else:
     force_beginning_resets = False
+  force_beginning_resets = tf.convert_to_tensor(force_beginning_resets)
 
   def group():
     return tf.group(batch_env.reset(tf.range(len(batch_env))),
                     tf.assign(cumulative_rewards, zeros_tensor))
   reset_op = tf.cond(
-      tf.logical_or(should_reset_var, tf.convert_to_tensor(
-          force_beginning_resets)),
-      group, tf.no_op)
+      tf.logical_or(should_reset_var, force_beginning_resets), group, tf.no_op)
 
   with tf.control_dependencies([reset_op]):
     reset_once_op = tf.assign(should_reset_var, False)
@@ -144,7 +163,7 @@ def step(index, scores_sum, scores_num):
       # https://github.com/tensorflow/tensorflow/issues/11186
       obs_copy = batch_env.observ + 0
 
-      def env_step(arg1, arg2):  # pylint: disable=unused-argument
+      def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
         """Step of the environment."""
         actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
         policy = actor_critic.policy
@@ -156,35 +175,44 @@ def env_step(arg1, arg2):  # pylint: disable=unused-argument
                            policy.sample)
 
         postprocessed_action = actor_critic.action_postprocessing(action)
-        simulate_output = batch_env.simulate(postprocessed_action[0, ...])
+        reward, done = batch_env.simulate(postprocessed_action[0, ...])
 
         pdf = policy.prob(action)[0]
         value_function = actor_critic.value[0]
         pdf = tf.reshape(pdf, shape=(hparams.num_agents,))
         value_function = tf.reshape(value_function, shape=(hparams.num_agents,))
+        done = tf.reshape(done, shape=(hparams.num_agents,))
 
-        with tf.control_dependencies(simulate_output):
-          return tf.identity(pdf), tf.identity(value_function)
+        with tf.control_dependencies([reward, done]):
+          return tf.identity(pdf), tf.identity(value_function), \
+                 tf.identity(done)
 
-      pdf, value_function = tf.while_loop(
-          lambda _1, _2: tf.equal(speculum.size(), 0),
+      # TODO(piotrmilos): while_body is executed at most once,
+      # thus should be replaced with tf.cond
+      pdf, value_function, top_level_done = tf.while_loop(
+          lambda _1, _2, _3: tf.equal(speculum.size(), 0),
           env_step,
-          [tf.constant(0.0, shape=(hparams.num_agents,)),
-           tf.constant(0.0, shape=(hparams.num_agents,))],
+          [
+              tf.constant(0.0, shape=(hparams.num_agents,)),
+              tf.constant(0.0, shape=(hparams.num_agents,)),
+              tf.constant(False, shape=(hparams.num_agents,))
+          ],
           parallel_iterations=1,
-          back_prop=False,)
+          back_prop=False,
+      )
 
       with tf.control_dependencies([pdf, value_function]):
         obs, reward, done, action = speculum.dequeue()
 
-        done = tf.reshape(done, (len(batch_env),))
         to_save = [obs, reward, done, action,
                    pdf, value_function]
         save_ops = [tf.scatter_update(memory_slot, index, value)
                     for memory_slot, value in zip(memory, to_save)]
         cumulate_rewards_op = cumulative_rewards.assign_add(reward)
-        agent_indices_to_reset = tf.where(done)[:, 0]
+
+        agent_indices_to_reset = tf.where(top_level_done)[:, 0]
       with tf.control_dependencies([cumulate_rewards_op]):
+        # TODO(piotrmilos): possibly we need cumulative_rewards.read_value()
         scores_sum_delta = tf.reduce_sum(
             tf.gather(cumulative_rewards, agent_indices_to_reset))
         scores_num_delta = tf.count_nonzero(done, dtype=tf.int32)
@@ -211,6 +239,19 @@ def stop_condition(i, _, resets):
         init,
         parallel_iterations=1,
         back_prop=False)
+
+  # We handle force_beginning_resets differently. We assume that all envs are
+  # reseted at the end of episod (though it happens at the beginning of the
+  # next one
+  scores_num = tf.cond(force_beginning_resets,
+                       lambda: scores_num + len(batch_env), lambda: scores_num)
+
+  with tf.control_dependencies([scores_sum]):
+    scores_sum = tf.cond(
+        force_beginning_resets,
+        lambda: scores_sum + tf.reduce_sum(cumulative_rewards.read_value()),
+        lambda: scores_sum)
+
   mean_score = tf.cond(tf.greater(scores_num, 0),
                        lambda: scores_sum / tf.cast(scores_num, tf.float32),
                        lambda: 0.)
diff --git a/tensor2tensor/rl/envs/batch_env.py b/tensor2tensor/rl/envs/batch_env.py
index f3a72844c..bf433066f 100644
--- a/tensor2tensor/rl/envs/batch_env.py
+++ b/tensor2tensor/rl/envs/batch_env.py
@@ -95,6 +95,8 @@ def step(self, actions):
           for env, action in zip(self._envs, actions)]
       transitions = [transition() for transition in transitions]
     observs, rewards, dones, infos = zip(*transitions)
+
+    # TODO(piotrmilos): Do we really want cast to float32
     observ = np.stack(observs).astype(np.float32)
     reward = np.stack(rewards).astype(np.float32)
     done = np.stack(dones)
@@ -118,6 +120,7 @@ def reset(self, indices=None):
       observs = [self._envs[index].reset(blocking=False) for index in indices]
       observs = [observ() for observ in observs]
     observ = np.stack(observs)
+    # TODO(piotrmilos): Do we really want this?
     observ = observ.astype(np.float32)
     return observ
 
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index 26ac6246a..8f610f42a 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -46,9 +46,8 @@ def batch_env_factory(hparams, xvfb=False):
   if environment_spec.simulated_env:
     # TODO(piotrmilos): Consider passing only relevant parameters
     cur_batch_env = _define_simulated_batch_env(
-        environment_spec, hparams.num_agents, hparams)
+        environment_spec, hparams.num_agents)
   else:
-
     cur_batch_env = _define_batch_env(hparams.environment_spec,
                                       hparams.num_agents,
                                       xvfb=xvfb)
@@ -67,11 +66,9 @@ def _define_batch_env(environment_spec, num_agents, xvfb=False):
     return env
 
 
-def _define_simulated_batch_env(environment_spec, num_agents,
-                                other_hparms):
+def _define_simulated_batch_env(environment_spec, num_agents):
   cur_batch_env = simulated_batch_env.SimulatedBatchEnv(environment_spec,
-                                                        num_agents,
-                                                        other_hparms)
+                                                        num_agents)
   return cur_batch_env
 
 
diff --git a/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index 5bc76f2f7..56e13878d 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -76,7 +76,7 @@ def reset(self, indices=None):
   @property
   def observ(self):
     """Access the variable holding the current observation."""
-    return self._observ
+    return self._observ.read_value()
 
   def close(self):
     """Send close messages to the external process and join them."""
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index dc5a0fa79..76916d991 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -117,7 +117,7 @@ def _reset_non_empty(self, indices):
   @property
   def observ(self):
     """Access the variable holding the current observation."""
-    return self._observ
+    return self._observ.read_value()
 
   def close(self):
     """Send close messages to the external process and join them."""
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index abb713dae..f6515becc 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -97,9 +97,9 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   flags are held in according variables.
   """
 
-  def __init__(self, environment_spec, length, other_hparams):
+  def __init__(self, environment_spec, length):
     """Batch of environments inside the TensorFlow graph."""
-    del other_hparams
+
     self.length = length
     initial_frames_problem = environment_spec.initial_frames_problem
     self._min_reward = initial_frames_problem.min_reward
@@ -203,4 +203,4 @@ def _reset_non_empty(self, indices):
   @property
   def observ(self):
     """Access the variable holding the current observation."""
-    return tf.identity(self._observ)
+    return self._observ.read_value()
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 03a64d9ed..f01020abc 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -17,6 +17,11 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+import math
+
+from tensor2tensor.layers import discretization
+from tensor2tensor.models.research import autoencoders
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
 
 import tensorflow as tf
@@ -53,66 +58,15 @@ def _reset_non_empty(self, indices):
       return tf.identity(new_values)
 
 
-class TransformWrapper(WrapperBase):
-  """Transform wrapper."""
-
-  def __init__(self, batch_env, transform_observation=None,
-               transform_reward=tf.identity, transform_done=tf.identity):
-    super(TransformWrapper, self).__init__(batch_env)
-    if transform_observation is not None:
-      _, observ_shape, observ_dtype = transform_observation  # pylint: disable=unpacking-non-sequence
-      self._observ = tf.Variable(
-          tf.zeros(len(self) + observ_shape, observ_dtype), trainable=False)
-    else:
-      self._observ = self._batch_env.observ
-
-    self.transform_observation = transform_observation
-    self.transform_reward = transform_reward
-    self.transform_done = transform_done
-
-  def simulate(self, action):
-    with tf.name_scope("environment/simulate"):  # Do we need this?
-      reward, done = self._batch_env.simulate(action)
-      with tf.control_dependencies([reward]):
-        if self.transform_observation:
-          observ = self.transform_observation[0](self._batch_env.observ)
-          assign_op = self._observ.assign(observ)
-        else:
-          assign_op = tf.no_op()  # TODO(lukaszkaiser): looks as if it's broken.
-        with tf.control_dependencies([assign_op]):
-          return self.transform_reward(reward), self.transform_done(done)
-
-
-class WarpFrameWrapper(TransformWrapper):
-  """Wrap frames."""
-
-  def __init__(self, batch_env):
-    """Warp frames to 84x84 as done in the Nature paper and later work."""
-
-    dims = [84, 84]
-    nature_transform = lambda o: tf.image.rgb_to_grayscale(  # pylint: disable=g-long-lambda
-        tf.image.resize_images(o, dims))
-
-    super(WarpFrameWrapper, self).__init__(batch_env, transform_observation=(
-        nature_transform, dims, tf.float32))
-
-
-class ShiftRewardWrapper(TransformWrapper):
-  """Shift the reward."""
-
-  def __init__(self, batch_env, add_value):
-    shift_reward = lambda r: tf.add(r, add_value)
-    super(ShiftRewardWrapper, self).__init__(
-        batch_env, transform_reward=shift_reward)
-
-
 class MaxAndSkipWrapper(WrapperBase):
-  """Max and skip wrapper."""
+  """ Max and skip wrapper.
+      The wrapper works under assumptions that issuing an action
+      to an environment with done=True has not effect.
+  """
 
   def __init__(self, batch_env, skip=4):
     super(MaxAndSkipWrapper, self).__init__(batch_env)
     self.skip = skip
-    self._observ = None
     observs_shape = batch_env.observ.shape
     observ_dtype = tf.float32
     self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
@@ -126,6 +80,7 @@ def simulate(self, action):
       def not_done_step(a, _):
         reward, done = self._batch_env.simulate(action)
         with tf.control_dependencies([reward, done]):
+          # TODO(piotrmilos): possibly ignore envs with done
           r0 = tf.maximum(a[0], self._batch_env.observ)
           r1 = tf.add(a[1], reward)
           r2 = tf.logical_or(a[2], done)
@@ -142,7 +97,10 @@ def not_done_step(a, _):
 
 
 class StackAndSkipWrapper(WrapperBase):
-  """Stack and skip wrapper."""
+  """ Stack and skip wrapper.
+      The wrapper works under assumptions that issuing an action
+      to an environment with done=True has not effect.
+  """
 
   def __init__(self, batch_env, skip=4):
     super(StackAndSkipWrapper, self).__init__(batch_env)
@@ -162,7 +120,7 @@ def simulate(self, action):
       def not_done_step(a, _):
         reward, done = self._batch_env.simulate(action)
         with tf.control_dependencies([reward, done]):
-          r0 = self._batch_env.observ
+          r0 = self._batch_env.observ + 0
           r1 = tf.add(a[1], reward)
           r2 = tf.logical_or(a[2], done)
           return (r0, r1, r2)
@@ -193,33 +151,82 @@ def _reset_non_empty(self, indices):
       return tf.gather(self.observ, indices)
 
 
-class TimeLimitWrapper(WrapperBase):
-  """Time limit wrapper."""
+class AutoencoderWrapper(WrapperBase):
+  """ Transforms the observations taking the bottleneck
+      state of an autoencoder"""
 
-  # TODO(lukaszkaiser): Check if TimeLimitWrapper does what it's supposed to do.
-  def __init__(self, batch_env, timelimit=100):
-    super(TimeLimitWrapper, self).__init__(batch_env)
-    self.timelimit = timelimit
-    self._time_elapsed = tf.Variable(tf.zeros((len(self),), tf.int32),
-                                     trainable=False)
+  def __init__(self, batch_env):
+    super(AutoencoderWrapper, self).__init__(batch_env)
+    batch_size, height, width, _ = self._batch_env.observ.get_shape().as_list()
+    ae_height = int(math.ceil(height / self.autoencoder_factor))
+    ae_width = int(math.ceil(width / self.autoencoder_factor))
+    ae_channels = 24  # TODO(piotrmilos): make it better
+    observ_shape = (batch_size, ae_height, ae_width, ae_channels)
+    self._observ = self._observ = tf.Variable(
+        tf.zeros(observ_shape, tf.float32), trainable=False)
+    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+      autoencoder_hparams = autoencoders.autoencoder_discrete_pong()
+      self.autoencoder_model = autoencoders.AutoencoderOrderedDiscrete(
+          autoencoder_hparams, tf.estimator.ModeKeys.EVAL)
+    self.autoencoder_model.set_mode(tf.estimator.ModeKeys.EVAL)
+
+  @property
+  def autoencoder_factor(self):
+    """By how much to divide sizes when using autoencoders."""
+    hparams = autoencoders.autoencoder_discrete_pong()
+    return 2**hparams.num_hidden_layers
 
   def simulate(self, action):
-    with tf.name_scope("environment/simulate"):
-      reward, done = self._batch_env.simulate(action)
-      with tf.control_dependencies([reward, done]):
-        new_done = tf.logical_or(done, self._time_elapsed > self.timelimit)
-        inc = self._time_elapsed.assign_add(tf.ones_like(self._time_elapsed))
+    reward, done = self._batch_env.simulate(action)
+    with tf.control_dependencies([reward, done]):
+      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+        ret = self.autoencoder_model.encode(self._batch_env.observ)
+        assign_op = self._observ.assign(ret)
+        with tf.control_dependencies([assign_op]):
+          return tf.identity(reward), tf.identity(done)
+
+  def _reset_non_empty(self, indices):
+    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+      new_values = self._batch_env._reset_non_empty(indices)  # pylint: disable=protected-access
+      ret = self.autoencoder_model.encode(new_values)
+      assign_op = tf.scatter_update(self._observ, indices, ret)
+      with tf.control_dependencies([assign_op]):
+        return tf.gather(self.observ, indices)
+
 
-        with tf.control_dependencies([inc]):
-          return tf.identity(reward), tf.identity(new_done)
+class IntToBitWrapper(WrapperBase):
+  """Unpacks the observations from integer values to bit values"""
+
+  def __init__(self, batch_env):
+    super(IntToBitWrapper, self).__init__(batch_env)
+    batch_size, height, width, channels = \
+      self._batch_env.observ.get_shape().as_list()
+    # We treat each channel as 8-bit integer to be expanded to 8 channels
+    self.observ_shape = (height, width, channels*8)
+    self._observ = self._observ = tf.Variable(
+        tf.zeros((batch_size,) + self.observ_shape, tf.float32),
+        trainable=False)
+
+  def simulate(self, action):
+    action = tf.Print(action, [action], message="action=", summarize=200)
+
+    # action = tf.zeros_like(action) #Temporary hacked bugfix
+    reward, done = self._batch_env.simulate(action)
+    with tf.control_dependencies([reward, done]):
+      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+        unpacked = discretization.int_to_bit(self._batch_env.observ, 8)
+        unpacked = tf.reshape(unpacked, (-1,)+self.observ_shape)
+        assign_op = self._observ.assign(unpacked)
+        with tf.control_dependencies([assign_op]):
+          return tf.identity(reward), tf.identity(done)
 
   def _reset_non_empty(self, indices):
-    op_zero = tf.scatter_update(
-        self._time_elapsed, indices,
-        tf.gather(tf.zeros((len(self),), tf.int32), indices))
     # pylint: disable=protected-access
     new_values = self._batch_env._reset_non_empty(indices)
+    new_values_unpacked = discretization.int_to_bit(new_values, 8)
+    new_values_unpacked = tf.reshape(new_values_unpacked, (-1,)
+                                     +self.observ_shape)
     # pylint: enable=protected-access
-    assign_op = tf.scatter_update(self._observ, indices, new_values)
-    with tf.control_dependencies([op_zero, assign_op]):
-      return tf.gather(self.observ, indices)
+    assign_op = tf.scatter_update(self._observ, indices, new_values_unpacked)
+    with tf.control_dependencies([assign_op]):
+      return tf.identity(new_values_unpacked)
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 920be540e..248aa9876 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -95,14 +95,15 @@ def generate_real_env_data(problem_name, agent_policy_path, hparams, data_dir,
       "problem": problem_name,
       "agent_policy_path": agent_policy_path,
       "autoencoder_path": autoencoder_path,
-      "only_use_ae_for_policy": True,
   }):
     gym_problem = registry.problem(problem_name)
     gym_problem.settable_num_steps = hparams.true_env_generator_num_steps
-    gym_problem.eval_phase = eval_phase
+    gym_problem.settable_eval_phase = eval_phase
     gym_problem.generate_data(data_dir, tmp_dir)
-    mean_reward = gym_problem.statistics.sum_of_rewards / \
-                  (1.0 + gym_problem.statistics.number_of_dones)
+    mean_reward = None
+    if gym_problem.statistics.number_of_dones:
+      mean_reward = (gym_problem.statistics.sum_of_rewards /
+                     gym_problem.statistics.number_of_dones)
 
   return mean_reward
 
@@ -133,8 +134,7 @@ def train_autoencoder(problem_name, data_dir, output_dir, hparams, epoch):
 
 
 def train_agent(problem_name, agent_model_dir,
-                event_dir, world_model_dir, epoch_data_dir, hparams,
-                autoencoder_path=None, epoch=0):
+                event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0):
   """Train the PPO agent in the simulated environment."""
   gym_problem = registry.problem(problem_name)
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
@@ -168,14 +168,12 @@ def train_agent(problem_name, agent_model_dir,
       "hparams_set": hparams.generative_model_params,
       "output_dir": world_model_dir,
       "data_dir": epoch_data_dir,
-      "autoencoder_path": autoencoder_path,
   }):
     rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=epoch)
 
 
 def evaluate_world_model(simulated_problem_name, problem_name, hparams,
-                         world_model_dir, epoch_data_dir, tmp_dir,
-                         autoencoder_path=None):
+                         world_model_dir, epoch_data_dir, tmp_dir):
   """Generate simulated environment data and return reward accuracy."""
   gym_simulated_problem = registry.problem(simulated_problem_name)
   sim_steps = hparams.simulated_env_generator_num_steps
@@ -186,21 +184,20 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
       "hparams_set": hparams.generative_model_params,
       "data_dir": epoch_data_dir,
       "output_dir": world_model_dir,
-      "autoencoder_path": autoencoder_path,
   }):
     gym_simulated_problem.generate_data(epoch_data_dir, tmp_dir)
   n = max(1., gym_simulated_problem.statistics.number_of_dones)
   model_reward_accuracy = (
       gym_simulated_problem.statistics.successful_episode_reward_predictions
       / float(n))
-  old_path = os.path.join(epoch_data_dir, "debug_frames_env")
-  new_path = os.path.join(epoch_data_dir, "debug_frames_env_eval")
-  tf.gfile.Rename(old_path, new_path)
+  old_path = os.path.join(epoch_data_dir, "debug_frames_sim")
+  new_path = os.path.join(epoch_data_dir, "debug_frames_sim_eval")
+  if not tf.gfile.Exists(new_path):
+    tf.gfile.Rename(old_path, new_path)
   return model_reward_accuracy
 
 
-def train_world_model(problem_name, data_dir, output_dir, hparams, epoch,
-                      use_autoencoder=False):
+def train_world_model(problem_name, data_dir, output_dir, hparams, epoch):
   """Train the world model on problem_name."""
   train_steps = hparams.model_train_steps * (epoch + 2)
   with temporary_flags({
@@ -211,9 +208,6 @@ def train_world_model(problem_name, data_dir, output_dir, hparams, epoch,
       "hparams_set": hparams.generative_model_params,
       "eval_steps": 100,
       "train_steps": train_steps,
-      # Hack: If training on autoencoded frames, autoencoder_path needs to be
-      # set so that the problem reports the right sizes for frames.
-      "autoencoder_path": "dummy" if use_autoencoder else None,
   }):
     t2t_trainer.main([])
 
@@ -250,13 +244,15 @@ def generator():
       while True:
         try:
           pngs_np, examples_np = sess.run([pngs, examples])
-          rewards_np = [list(el) for el in examples_np["reward"]]
-          actions_np = [list(el) for el in examples_np["action"]]
-          pngs_np = [el for el in pngs_np]
-          for action, reward, png in zip(actions_np, rewards_np, pngs_np):
+          rewards = examples_np["reward"].tolist()
+          actions = examples_np["action"].tolist()
+          frame_numbers = examples_np["frame_number"].tolist()
+          for action, reward, frame_number, png in \
+                  zip(actions, rewards, frame_numbers, pngs_np):
             yield {
                 "action": action,
                 "reward": reward,
+                "frame_number": frame_number,
                 "image/encoded": [png],
                 "image/format": ["png"],
                 "image/height": [encoded_frame_height],
@@ -332,17 +328,22 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   directories = setup_directories(output_dir, subdirectories)
 
   # Problems
-  problem_name = "gym_discrete_problem_with_agent_on_%s" % hparams.game
-  ae_problem_name = problem_name + "_ae"
-  simulated_problem_name = (
-      "gym_simulated_discrete_problem_with_agent_on_%s" % hparams.game)
-  world_model_problem = ae_problem_name if using_autoencoder else problem_name
-  check_problems([problem_name, world_model_problem, simulated_problem_name])
+  if using_autoencoder:
+    problem_name = (
+        "gym_discrete_problem_with_agent_on_%s_with_autoencoder" % hparams.game)
+    world_model_problem = (
+        "gym_discrete_problem_with_agent_on_%s_autoencoded" % hparams.game)
+    simulated_problem_name = (
+        "gym_simulated_discrete_problem_with_agent_on_%s_autoencoded"
+        % hparams.game)
+  else:
+    problem_name = ("gym_discrete_problem_with_agent_on_%s" % hparams.game)
+    world_model_problem = problem_name
+    simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
+                              % hparams.game)
 
   # Autoencoder model dir
-  autoencoder_model_dir = (FLAGS.autoencoder_path or
-                           directories.get("autoencoder"))
-  FLAGS.autoencoder_path = None
+  autoencoder_model_dir = directories.get("autoencoder")
 
   # Timing log function
   log_relative_time = make_relative_timing_fn()
@@ -357,7 +358,17 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   tf.logging.info("Generating real environment data with random policy")
   mean_reward = generate_real_env_data(
       problem_name, None, hparams, data_dir, directories["tmp"])
-  tf.logging.info("Mean reward (random): %.4f", mean_reward)
+  tf.logging.info("Mean reward (random): {}".format(mean_reward))
+
+  eval_metrics_event_dir = os.path.join(directories["world_model"],
+                                        "eval_metrics_event_dir")
+  eval_metrics_writer = tf.summary.FileWriter(eval_metrics_event_dir)
+  model_reward_accuracy_summary = tf.Summary()
+  model_reward_accuracy_summary.value.add(tag="model_reward_accuracy",
+                                          simple_value=None)
+  mean_reward_summary = tf.Summary()
+  mean_reward_summary.value.add(tag="mean_reward",
+                                simple_value=None)
 
   for epoch in range(hparams.epochs):
     is_final_epoch = (epoch + 1) == hparams.epochs
@@ -386,8 +397,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     # Train world model
     log("Training world model")
     train_world_model(world_model_problem, epoch_data_dir,
-                      directories["world_model"], hparams, epoch,
-                      use_autoencoder=using_autoencoder)
+                      directories["world_model"], hparams, epoch)
 
     # Evaluate world model
     model_reward_accuracy = 0.
@@ -396,33 +406,49 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       model_reward_accuracy = evaluate_world_model(
           simulated_problem_name, world_model_problem, hparams,
           directories["world_model"],
-          epoch_data_dir, directories["tmp"],
-          autoencoder_path=autoencoder_model_dir)
+          epoch_data_dir, directories["tmp"])
       log("World model reward accuracy: %.4f", model_reward_accuracy)
 
     # Train PPO
     log("Training PPO")
-    ppo_event_dir = os.path.join(directories["ppo"], str(epoch))
+    ppo_event_dir = os.path.join(directories["world_model"],
+                                 "ppo_summaries", str(epoch))
     ppo_model_dir = directories["ppo"]
     if not hparams.ppo_continue_training:
       ppo_model_dir = ppo_event_dir
     train_agent(simulated_problem_name, ppo_model_dir,
                 ppo_event_dir, directories["world_model"], epoch_data_dir,
-                hparams, autoencoder_path=autoencoder_model_dir, epoch=epoch)
+                hparams, epoch=epoch)
 
     # Collect data from the real environment.
     log("Generating real environment data")
-    if is_final_epoch:
-      epoch_data_dir = os.path.join(epoch_data_dir, "final_eval")
+    eval_data_dir = os.path.join(epoch_data_dir, "eval")
     mean_reward = generate_real_env_data(
-        problem_name, ppo_model_dir, hparams, epoch_data_dir,
+        problem_name, ppo_model_dir, hparams, eval_data_dir,
         directories["tmp"], autoencoder_path=autoencoder_model_dir,
-        eval_phase=is_final_epoch)
-    log("Mean reward during generation: %.4f", mean_reward)
+        eval_phase=True)
+    log("Mean eval reward: {}".format(mean_reward))
+
+    if not is_final_epoch:
+      generation_mean_reward = generate_real_env_data(
+          problem_name, ppo_model_dir, hparams, epoch_data_dir,
+          directories["tmp"], autoencoder_path=autoencoder_model_dir,
+          eval_phase=False)
+      log("Mean reward during generation: {}".format(generation_mean_reward))
 
     # Report metrics.
     eval_metrics = {"model_reward_accuracy": model_reward_accuracy,
                     "mean_reward": mean_reward}
+
+    model_reward_accuracy_summary.value[0].simple_value \
+      = model_reward_accuracy
+
+    mean_reward_summary.value[0].simple_value \
+      = mean_reward
+
+    eval_metrics_writer.add_summary(model_reward_accuracy_summary, epoch)
+    eval_metrics_writer.add_summary(mean_reward_summary, epoch)
+
     epoch_metrics.append(eval_metrics)
     log("Eval metrics: %s", str(eval_metrics))
     if report_fn:
@@ -672,6 +698,7 @@ def rl_modelrl_ae_tiny():
   hparams.ppo_params = "ppo_pong_ae_base"
   hparams.generative_model_params = "next_frame_ae"
   hparams.autoencoder_train_steps = 2
+  hparams.eval_world_model = False
   return hparams
 
 
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index ec8618832..3d908495a 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -28,9 +28,8 @@
 import tensorflow as tf
 
 
-def define_train(hparams, event_dir):
+def define_train(hparams):
   """Define the training setup."""
-  del event_dir
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
     memory, collect_summary, initialization\
       = collect.define_collect(
@@ -45,7 +44,7 @@ def train(hparams, event_dir=None, model_dir=None,
           restore_agent=True, epoch=0):
   """Train."""
   with tf.name_scope("rl_train"):
-    train_summary_op, _, initialization = define_train(hparams, event_dir)
+    train_summary_op, _, initialization = define_train(hparams)
     if event_dir:
       summary_writer = tf.summary.FileWriter(
           event_dir, graph=tf.get_default_graph(), flush_secs=60)
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index 4081e4aa9..963b86dec 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -49,22 +49,16 @@ def test_no_crash_cartpole(self):
                        standard_atari_env_spec("CartPole-v0"))
     rl_trainer_lib.train(hparams)
 
-  # This test should sucessfully train pong.
-  # It should get train mean_score around 0 after 100 epoch
-  #
-  # This test should be run whenever ppo any bigger change
-  # is done on the ppo code
-  #
-  # To run the test change epochs_num=2 to epoch_num=200
-  # and epoch_length=4 to epoch_length=200
-  # (it is set like that to meet travis timeouts
+  # This test should successfully train pong.
+  # It should get train mean_score around 0 after 200 epoch
+  # By default the test is disabled to avoid travis timeouts
   def test_train_pong(self):
     hparams = tf.contrib.training.\
-      HParams(epochs_num=2,
+      HParams(epochs_num=300,
               eval_every_epochs=10,
-              num_agents=20,
+              num_agents=10,
               optimization_epochs=3,
-              epoch_length=4,
+              epoch_length=200,
               entropy_loss_coef=0.003,
               learning_rate=8e-05,
               optimizer="Adam",

From 512f60db847661c5bf2193e4bfe24ab29305fb8c Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Wed, 11 Jul 2018 14:49:20 -0700
Subject: [PATCH 0318/2720] Add Recognizing Textual Entailment data.

PiperOrigin-RevId: 204193774
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/rte.py          | 135 ++++++++++++++++++
 tensor2tensor/data_generators/wnli.py         |   3 +-
 3 files changed, 137 insertions(+), 2 deletions(-)
 create mode 100644 tensor2tensor/data_generators/rte.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 36641436a..4c7436f8c 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -52,6 +52,7 @@
     "tensor2tensor.data_generators.ocr",
     "tensor2tensor.data_generators.problem_hparams",
     "tensor2tensor.data_generators.ptb",
+    "tensor2tensor.data_generators.rte",
     "tensor2tensor.data_generators.snli",
     "tensor2tensor.data_generators.style_transfer",
     "tensor2tensor.data_generators.squad",
diff --git a/tensor2tensor/data_generators/rte.py b/tensor2tensor/data_generators/rte.py
new file mode 100644
index 000000000..474c30b2f
--- /dev/null
+++ b/tensor2tensor/data_generators/rte.py
@@ -0,0 +1,135 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generators for the Recognizing Textual Entailment dataset."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import zipfile
+import six
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+EOS = text_encoder.EOS
+
+
+@registry.register_problem
+class RTE(text_problems.TextConcat2ClassProblem):
+  """Recognizing Textual Entailment classification problems."""
+
+  # Link to data from GLUE: https://gluebenchmark.com/tasks
+  _RTE_URL = ("https://firebasestorage.googleapis.com/v0/b/"
+              "mtl-sentence-representations.appspot.com/o/"
+              "data%2FRTE.zip?alt=media&token=5efa7e85-a0bb-"
+              "4f19-8ea2-9e1840f077fb")
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 1,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  @property
+  def approx_vocab_size(self):
+    return 2**13  # 8k vocab suffices for this small dataset.
+
+  @property
+  def vocab_filename(self):
+    return "vocab.rte.%d" % self.approx_vocab_size
+
+  @property
+  def num_classes(self):
+    return 2
+
+  @property
+  def concat_token(self):
+    return "<EN-PR-HYP>"
+
+  @property
+  def concat_id(self):
+    if self.vocab_type == text_problems.VocabType.CHARACTER:
+      return problem.SpaceID.EN_PR_HYP
+    return 2
+
+  def class_labels(self, data_dir):
+    del data_dir
+    # Note this binary classification is different from usual MNLI.
+    return ["not_entailment", "entailment"]
+
+  def _maybe_download_corpora(self, tmp_dir):
+    rte_filename = "RTE.zip"
+    rte_finalpath = os.path.join(tmp_dir, "RTE")
+    if not tf.gfile.Exists(rte_finalpath):
+      zip_filepath = generator_utils.maybe_download(
+          tmp_dir, rte_filename, self._RTE_URL)
+      zip_ref = zipfile.ZipFile(zip_filepath, "r")
+      zip_ref.extractall(tmp_dir)
+      zip_ref.close()
+
+    return rte_finalpath
+
+  def example_generator(self, filename):
+    label_list = self.class_labels(data_dir=None)
+    for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
+      if idx == 0: continue  # skip header
+      if six.PY2:
+        line = unicode(line.strip(), "utf-8")
+      else:
+        line = line.strip().decode("utf-8")
+      _, s1, s2, l = line.split("\t")
+      inputs = [s1, s2]
+      l = label_list.index(l)
+      yield {
+          "inputs": inputs,
+          "label": l
+      }
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    rte_dir = self._maybe_download_corpora(tmp_dir)
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      filesplit = "train.tsv"
+    else:
+      filesplit = "dev.tsv"
+
+    filename = os.path.join(rte_dir, filesplit)
+    for example in self.example_generator(filename):
+      yield example
+
+
+@registry.register_problem
+class RTECharacters(RTE):
+  """Recognizing Textual Entailment classification problems, character level"""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  @property
+  def task_id(self):
+    return problem.SpaceID.EN_NLI
diff --git a/tensor2tensor/data_generators/wnli.py b/tensor2tensor/data_generators/wnli.py
index fc124b8f9..b4945d1e5 100644
--- a/tensor2tensor/data_generators/wnli.py
+++ b/tensor2tensor/data_generators/wnli.py
@@ -80,7 +80,7 @@ def concat_id(self):
   def class_labels(self, data_dir):
     del data_dir
     # Note this binary classification is different from usual MNLI.
-    return ["contradiction", "entailment"]
+    return ["not_entailment", "entailment"]
 
   def _maybe_download_corpora(self, tmp_dir):
     wnli_filename = "WNLI.zip"
@@ -102,7 +102,6 @@ def example_generator(self, filename):
       else:
         line = line.strip().decode("utf-8")
       _, s1, s2, l = line.split("\t")
-      # inputs = " ".join([s1 EOS, s2])
       inputs = [s1, s2]
       yield {
           "inputs": inputs,

From 6f09793a019c42639febbda0a43950f6f0b4c33f Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 11 Jul 2018 14:54:15 -0700
Subject: [PATCH 0319/2720] outputting "extra" predictions if exists.

PiperOrigin-RevId: 204194624
---
 tensor2tensor/utils/t2t_model.py | 17 ++++++-----------
 1 file changed, 6 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 979d81ae6..3be734b0f 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1337,7 +1337,7 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
   def estimator_spec_predict(self, features, use_tpu=False):
     """Construct EstimatorSpec for PREDICT mode."""
     decode_hparams = self._decode_hparams
-    infer_out = self.infer(
+    predictions = self.infer(
         features,
         beam_size=decode_hparams.beam_size,
         top_beams=(decode_hparams.beam_size
@@ -1345,24 +1345,19 @@ def estimator_spec_predict(self, features, use_tpu=False):
         alpha=decode_hparams.alpha,
         decode_length=decode_hparams.extra_length,
         use_tpu=use_tpu)
-    if isinstance(infer_out, dict):
-      outputs = infer_out["outputs"]
-      scores = infer_out["scores"]
-    else:
-      outputs = infer_out
-      scores = None
+    if not isinstance(predictions, dict):
+      predictions = {"outputs": predictions}
 
     inputs = features.get("inputs")
     if inputs is None:
       inputs = features["targets"]
 
-    predictions = {
-        "outputs": outputs,
-        "scores": scores,
+    predictions.update({
         "inputs": inputs,
         "targets": features.get("infer_targets"),
         "batch_prediction_key": features.get("batch_prediction_key"),
-    }
+        })
+
     _del_dict_nones(predictions)
 
     export_out = {"outputs": predictions["outputs"]}

From a5481fa8bcfd35ed43ce123529c2084a3be23d9f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 11 Jul 2018 15:13:27 -0700
Subject: [PATCH 0320/2720] Add Imagenet 32x32 generator to Tensor2Tensor.

PiperOrigin-RevId: 204198183
---
 tensor2tensor/data_generators/imagenet.py | 34 +++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 7d569bbd8..e9db5011d 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -186,6 +186,40 @@ def preprocess_example(self, example, mode, _):
     return example
 
 
+@registry.register_problem
+class ImageImagenet32Gen(ImageImagenet):
+  """Imagenet 32 from the pixen cnn paper"""
+
+  @property
+  def train_shards(self):
+    return 1024
+
+  @property
+  def dev_shards(self):
+    return 10
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    generator_utils.generate_dataset_and_shuffle(
+        self.generator(data_dir, tmp_dir, True),
+        self.training_filepaths(data_dir, self.train_shards, shuffled=True),
+        self.generator(data_dir, tmp_dir, False),
+        self.dev_filepaths(data_dir, self.dev_shards, shuffled=True))
+
+  def generator(self, data_dir, tmp_dir, is_training):
+    if is_training:
+      return imagenet_pixelrnn_generator(
+          tmp_dir, int(True), size=_IMAGENET_SMALL_IMAGE_SIZE)
+    else:
+      return imagenet_pixelrnn_generator(
+          tmp_dir, int(is_training), size=_IMAGENET_SMALL_IMAGE_SIZE)
+
+  def preprocess_example(self, example, mode, unused_hparams):
+    example["inputs"].set_shape([_IMAGENET_SMALL_IMAGE_SIZE,
+                                 _IMAGENET_SMALL_IMAGE_SIZE, 3])
+    example["inputs"] = tf.to_int64(example["inputs"])
+    return example
+
+
 @registry.register_problem
 class ImageImagenet64Gen(ImageImagenet):
   """Imagenet 64 from the pixen cnn paper"""

From 9f193ac2e8716902337a7e151d409a5b37e579ee Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Wed, 11 Jul 2018 15:14:57 -0700
Subject: [PATCH 0321/2720] Fix ResNet when running on cifar

PiperOrigin-RevId: 204198477
---
 tensor2tensor/models/resnet.py | 54 ++++++++++++++++++----------------
 1 file changed, 28 insertions(+), 26 deletions(-)

diff --git a/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index c2db3e064..a6996f773 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -312,7 +312,8 @@ def resnet_v2(inputs,
               layers,
               filters,
               data_format="channels_first",
-              is_training=False):
+              is_training=False,
+              is_cifar=False):
   """Resnet model.
 
   Args:
@@ -327,27 +328,11 @@ def resnet_v2(inputs,
     data_format: `str`, "channels_first" `[batch, channels, height,
         width]` or "channels_last" `[batch, height, width, channels]`.
     is_training: bool, build in training mode or not.
+    is_cifar: bool, whether the data is CIFAR or not.
 
   Returns:
     Pre-logit activations.
   """
-  inputs = conv2d_fixed_padding(
-      inputs=inputs,
-      filters=filters[0],
-      kernel_size=7,
-      strides=2,
-      data_format=data_format)
-  inputs = tf.identity(inputs, "initial_conv")
-  inputs = batch_norm_relu(inputs, is_training, data_format=data_format)
-
-  inputs = tf.layers.max_pooling2d(
-      inputs=inputs,
-      pool_size=3,
-      strides=2,
-      padding="SAME",
-      data_format=data_format)
-  inputs = tf.identity(inputs, "initial_max_pool")
-
   inputs = block_layer(
       inputs=inputs,
       filters=filters[1],
@@ -375,7 +360,7 @@ def resnet_v2(inputs,
       is_training=is_training,
       name="block_layer3",
       data_format=data_format)
-  if filters[4]:
+  if not is_cifar:
     inputs = block_layer(
         inputs=inputs,
         filters=filters[4],
@@ -400,6 +385,7 @@ def body(self, features):
         "bottleneck": bottleneck_block,
     }
     assert hp.block_fn in block_fns
+    is_training = hp.mode == tf.estimator.ModeKeys.TRAIN
 
     inputs = features["inputs"]
 
@@ -410,13 +396,32 @@ def body(self, features):
       inputs = tf.transpose(inputs, [0, 3, 1, 2])
       data_format = "channels_first"
 
+    inputs = conv2d_fixed_padding(
+        inputs=inputs,
+        filters=hp.filter_sizes[0],
+        kernel_size=7,
+        strides=1 if hp.is_cifar else 2,
+        data_format=data_format)
+    inputs = tf.identity(inputs, "initial_conv")
+    inputs = batch_norm_relu(inputs, is_training, data_format=data_format)
+
+    if not hp.is_cifar:
+      inputs = tf.layers.max_pooling2d(
+          inputs=inputs,
+          pool_size=3,
+          strides=2,
+          padding="SAME",
+          data_format=data_format)
+      inputs = tf.identity(inputs, "initial_max_pool")
+
     out = resnet_v2(
         inputs,
         block_fns[hp.block_fn],
         hp.layer_sizes,
         hp.filter_sizes,
         data_format,
-        is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
+        is_training=is_training,
+        is_cifar=hp.is_cifar)
 
     if hp.use_nchw:
       out = tf.transpose(out, [0, 2, 3, 1])
@@ -458,6 +463,7 @@ def resnet_base():
   hparams.add_hparam("filter_sizes", [64, 64, 128, 256, 512])
   hparams.add_hparam("block_fn", "bottleneck")
   hparams.add_hparam("use_nchw", True)
+  hparams.add_hparam("is_cifar", False)
 
   # Variable init
   hparams.initializer = "normal_unit_scaling"
@@ -516,13 +522,9 @@ def resnet_cifar_15():
   """Set of hyperparameters."""
   hp = resnet_base()
   hp.block_fn = "residual"
+  hp.is_cifar = True
   hp.layer_sizes = [2, 2, 2]
-  hp.filter_sizes = [16, 16, 32, 64, None]
-
-  hp.learning_rate = 0.1 * 128. * 8. / 256.
-  hp.learning_rate_decay_scheme = "piecewise"
-  hp.add_hparam("learning_rate_boundaries", [40000, 60000, 80000])
-  hp.add_hparam("learning_rate_multiples", [0.1, 0.01, 0.001])
+  hp.filter_sizes = [16, 32, 64, 128]
 
   return hp
 

From bae5ff79cae83fa03ccb96884690762629cb3d16 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Wed, 11 Jul 2018 15:17:07 -0700
Subject: [PATCH 0322/2720] Add Question-Answering NLI data.

PiperOrigin-RevId: 204198855
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/problem.py      |   2 +
 tensor2tensor/data_generators/qnli.py         | 135 ++++++++++++++++++
 3 files changed, 138 insertions(+)
 create mode 100644 tensor2tensor/data_generators/qnli.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 4c7436f8c..bed452483 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -52,6 +52,7 @@
     "tensor2tensor.data_generators.ocr",
     "tensor2tensor.data_generators.problem_hparams",
     "tensor2tensor.data_generators.ptb",
+    "tensor2tensor.data_generators.qnli",
     "tensor2tensor.data_generators.rte",
     "tensor2tensor.data_generators.snli",
     "tensor2tensor.data_generators.style_transfer",
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index c62496141..a6578d721 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -108,6 +108,8 @@ class SpaceID(object):
   EN_NLI = 33
   # COLA
   COLA = 34
+  # Enligh Question Context pair
+  EN_Q_CONT = 35
   # 3 class NLI
   THREE_CL_NLI = 37
 
diff --git a/tensor2tensor/data_generators/qnli.py b/tensor2tensor/data_generators/qnli.py
new file mode 100644
index 000000000..ed7ed604b
--- /dev/null
+++ b/tensor2tensor/data_generators/qnli.py
@@ -0,0 +1,135 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generators for the Question-Answering NLI dataset."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import zipfile
+import six
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+EOS = text_encoder.EOS
+
+
+@registry.register_problem
+class QuestionNLI(text_problems.TextConcat2ClassProblem):
+  """Question Answering NLI classification problems."""
+
+  # Link to data from GLUE: https://gluebenchmark.com/tasks
+  _QNLI_URL = ("https://firebasestorage.googleapis.com/v0/b/"
+               "mtl-sentence-representations.appspot.com/o/"
+               "data%2FQNLI.zip?alt=media&token=c24cad61-f2df-"
+               "4f04-9ab6-aa576fa829d0")
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 100,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  @property
+  def approx_vocab_size(self):
+    return 2**15
+
+  @property
+  def vocab_filename(self):
+    return "vocab.qnli.%d" % self.approx_vocab_size
+
+  @property
+  def num_classes(self):
+    return 2
+
+  @property
+  def concat_token(self):
+    return "<EN-Q-CONT>"
+
+  @property
+  def concat_id(self):
+    if self.vocab_type == text_problems.VocabType.CHARACTER:
+      return problem.SpaceID.EN_Q_CONT
+    return 2
+
+  def class_labels(self, data_dir):
+    del data_dir
+    # Note this binary classification is different from usual MNLI.
+    return ["not_entailment", "entailment"]
+
+  def _maybe_download_corpora(self, tmp_dir):
+    qnli_filename = "QNLI.zip"
+    qnli_finalpath = os.path.join(tmp_dir, "QNLI")
+    if not tf.gfile.Exists(qnli_finalpath):
+      zip_filepath = generator_utils.maybe_download(
+          tmp_dir, qnli_filename, self._QNLI_URL)
+      zip_ref = zipfile.ZipFile(zip_filepath, "r")
+      zip_ref.extractall(tmp_dir)
+      zip_ref.close()
+
+    return qnli_finalpath
+
+  def example_generator(self, filename):
+    label_list = self.class_labels(data_dir=None)
+    for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
+      if idx == 0: continue  # skip header
+      if six.PY2:
+        line = unicode(line.strip(), "utf-8")
+      else:
+        line = line.strip().decode("utf-8")
+      _, s1, s2, l = line.split("\t")
+      inputs = [s1, s2]
+      l = label_list.index(l)
+      yield {
+          "inputs": inputs,
+          "label": l
+      }
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    qnli_dir = self._maybe_download_corpora(tmp_dir)
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      filesplit = "train.tsv"
+    else:
+      filesplit = "dev.tsv"
+
+    filename = os.path.join(qnli_dir, filesplit)
+    for example in self.example_generator(filename):
+      yield example
+
+
+@registry.register_problem
+class QuestionNLICharacters(QuestionNLI):
+  """Question-Answering NLI classification problems, character level"""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  @property
+  def task_id(self):
+    return problem.SpaceID.EN_NLI

From dd5d1e567c537cfe468ba89f97303ad5f3db95cd Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Wed, 11 Jul 2018 15:35:06 -0700
Subject: [PATCH 0323/2720] Add Quora Question Pairs data.

PiperOrigin-RevId: 204201783
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/problem.py      |   2 +
 tensor2tensor/data_generators/quora_qpairs.py | 141 ++++++++++++++++++
 3 files changed, 144 insertions(+)
 create mode 100644 tensor2tensor/data_generators/quora_qpairs.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index bed452483..62b448655 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -53,6 +53,7 @@
     "tensor2tensor.data_generators.problem_hparams",
     "tensor2tensor.data_generators.ptb",
     "tensor2tensor.data_generators.qnli",
+    "tensor2tensor.data_generators.quora_qpairs",
     "tensor2tensor.data_generators.rte",
     "tensor2tensor.data_generators.snli",
     "tensor2tensor.data_generators.style_transfer",
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index a6578d721..a21ab5e58 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -110,6 +110,8 @@ class SpaceID(object):
   COLA = 34
   # Enligh Question Context pair
   EN_Q_CONT = 35
+  # English similarity task
+  EN_SIM = 36
   # 3 class NLI
   THREE_CL_NLI = 37
 
diff --git a/tensor2tensor/data_generators/quora_qpairs.py b/tensor2tensor/data_generators/quora_qpairs.py
new file mode 100644
index 000000000..4aff774a8
--- /dev/null
+++ b/tensor2tensor/data_generators/quora_qpairs.py
@@ -0,0 +1,141 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generators for the Quora Question Pairs dataset."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import zipfile
+import six
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+EOS = text_encoder.EOS
+
+
+@registry.register_problem
+class QuoraQuestionPairs(text_problems.TextConcat2ClassProblem):
+  """Quora duplicate question pairs binary classification problems."""
+
+  # Link to data from GLUE: https://gluebenchmark.com/tasks
+  _QQP_URL = ("https://firebasestorage.googleapis.com/v0/b/"
+              "mtl-sentence-representations.appspot.com/o/"
+              "data%2FQQP.zip?alt=media&token=700c6acf-160d-"
+              "4d89-81d1-de4191d02cb5")
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 100,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  @property
+  def approx_vocab_size(self):
+    return 2**15
+
+  @property
+  def vocab_filename(self):
+    return "vocab.qqp.%d" % self.approx_vocab_size
+
+  @property
+  def num_classes(self):
+    return 2
+
+  @property
+  def concat_token(self):
+    return "<SENT_SEP>"
+
+  @property
+  def concat_id(self):
+    if self.vocab_type == text_problems.VocabType.CHARACTER:
+      return problem.SpaceID.EN_CHR
+    return 2
+
+  def class_labels(self, data_dir):
+    del data_dir
+    return ["not_duplicate", "duplicate"]
+
+  def _maybe_download_corpora(self, tmp_dir):
+    qqp_filename = "QQP.zip"
+    qqp_finalpath = os.path.join(tmp_dir, "QQP")
+    if not tf.gfile.Exists(qqp_finalpath):
+      zip_filepath = generator_utils.maybe_download(
+          tmp_dir, qqp_filename, self._QQP_URL)
+      zip_ref = zipfile.ZipFile(zip_filepath, "r")
+      zip_ref.extractall(tmp_dir)
+      zip_ref.close()
+
+    return qqp_finalpath
+
+  def example_generator(self, filename):
+    skipped = 0
+    for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
+      if idx == 0: continue  # skip header
+      if six.PY2:
+        line = unicode(line.strip(), "utf-8")
+      else:
+        line = line.strip().decode("utf-8")
+      split_line = line.split("\t")
+      if len(split_line) < 6:
+        skipped += 1
+        tf.logging.info("Skipping %d" % skipped)
+        continue
+      s1, s2, l = split_line[3:]
+      # A neat data augmentation trick from Radford et al. (2018)
+      # https://blog.openai.com/language-unsupervised/
+      inputs = [[s1, s2], [s2, s1]]
+      for inp in inputs:
+        yield {
+            "inputs": inp,
+            "label": int(l)
+        }
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    qqp_dir = self._maybe_download_corpora(tmp_dir)
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      filesplit = "train.tsv"
+    else:
+      filesplit = "dev.tsv"
+
+    filename = os.path.join(qqp_dir, filesplit)
+    for example in self.example_generator(filename):
+      yield example
+
+
+@registry.register_problem
+class QuoraQuestionPairsCharacters(QuoraQuestionPairs):
+  """Quora duplicate question pairs classification problems, character level"""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  @property
+  def task_id(self):
+    return problem.SpaceID.EN_SIM

From d24f7ccd258545d8acdc5430030c1ccab5cd362d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 11 Jul 2018 16:42:04 -0700
Subject: [PATCH 0324/2720] multilabel metrics include label length equals 0

PiperOrigin-RevId: 204212959
---
 tensor2tensor/utils/metrics.py      | 13 ++++++-------
 tensor2tensor/utils/metrics_test.py |  2 +-
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index fac0fabe6..41b2b4c4b 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -298,19 +298,18 @@ def multilabel_accuracy_matchk(predictions,
     weights_fn: weight function.
   Returns:
     scores: min(n/k, 1).
-    weights: 1 if labels contains non-zero label else 0.
+    weights: returns all ones.
 
   """
   predictions = tf.to_int32(tf.argmax(predictions, axis=-1))
-  length = tf.shape(labels)[1]
-  predictions = tf.tile(predictions, [1, length, 1, 1])
   scores = tf.to_float(tf.equal(predictions, labels))
+  # those label == 0 do not count
+  weights = weights_fn(labels)
+  scores *= weights
   scores = tf.reduce_sum(scores, axis=[1, 2, 3])
   scores = tf.minimum(scores / tf.to_float(k), 1)
-
-  weights = weights_fn(labels)
-  weights = tf.reduce_sum(weights, axis=[1, 2, 3])
-  weights = tf.to_float(tf.greater(weights, 0.))
+  # every sample count
+  weights = tf.ones(tf.shape(scores), dtype=tf.float32)
 
   return scores, weights
 
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 953bc51ae..9ad3bd2a9 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -229,7 +229,7 @@ def testMultilabelMatch3(self):
     expected = (predictions_repeat == targets).astype(float)
     expected = np.sum(expected, axis=(1, 2, 3))
     expected = np.minimum(expected / 3.0, 1.)
-    expected = np.sum(expected * weights[:, 0, 0, 0]) / np.sum(weights)
+    expected = np.sum(expected * weights[:, 0, 0, 0]) / weights.shape[0]
     with self.test_session() as session:
       scores, weights_ = metrics.multilabel_accuracy_match3(
           tf.one_hot(predictions, depth=5, dtype=tf.float32),

From 85d273f3ed3c298628b14270ac85840365d47225 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 12 Jul 2018 09:52:14 -0700
Subject: [PATCH 0325/2720] Add support for TF 1.9 and drop support for TF 1.5

PiperOrigin-RevId: 204315619
---
 .travis.yml                                     | 8 ++++----
 setup.py                                        | 4 ++--
 tensor2tensor/utils/multistep_optimizer_test.py | 5 -----
 3 files changed, 6 insertions(+), 11 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 2fb36ae47..19fa25160 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -7,25 +7,25 @@ env:
     - T2T_PROBLEM=algorithmic_reverse_binary40_test
     - T2T_DATA_DIR=/tmp/t2t-data
     - T2T_TRAIN_DIR=/tmp/t2t-train
-    - TF_LATEST="1.8.*"
+    - TF_LATEST="1.9.*"
     # This is necessary to have gsutil work with Python 2.7
     - BOTO_CONFIG=/dev/null
   matrix:
     # We test against the last 4 versions of TensorFlow
     # If updating, also update TF_LATEST above
-    - TF_VERSION="1.5.*"
     - TF_VERSION="1.6.*"
     - TF_VERSION="1.7.*"
     - TF_VERSION="1.8.*"
+    - TF_VERSION="1.9.*"
 matrix:
   exclude:
     # We test against all versions in Python 2 but only the latest in Python 3
-    - python: "3.6"
-      env: TF_VERSION="1.5.*"
     - python: "3.6"
       env: TF_VERSION="1.6.*"
     - python: "3.6"
       env: TF_VERSION="1.7.*"
+    - python: "3.6"
+      env: TF_VERSION="1.8.*"
 before_install:
   # Disabled TensorFlow Serving install until bug fixed. See "Export and query"
   # section below.
diff --git a/setup.py b/setup.py
index 2a509074e..b1cc9754c 100644
--- a/setup.py
+++ b/setup.py
@@ -49,8 +49,8 @@
         'tqdm',
     ],
     extras_require={
-        'tensorflow': ['tensorflow>=1.5.0'],
-        'tensorflow_gpu': ['tensorflow-gpu>=1.5.0'],
+        'tensorflow': ['tensorflow>=1.6.0'],
+        'tensorflow_gpu': ['tensorflow-gpu>=1.6.0'],
         'tests': [
             'absl-py', 'pytest', 'mock', 'pylint', 'jupyter', 'gsutil'
             # Need atari extras for Travis tests, but because gym is already in
diff --git a/tensor2tensor/utils/multistep_optimizer_test.py b/tensor2tensor/utils/multistep_optimizer_test.py
index 50affc12b..584bd95c3 100644
--- a/tensor2tensor/utils/multistep_optimizer_test.py
+++ b/tensor2tensor/utils/multistep_optimizer_test.py
@@ -25,11 +25,6 @@
 class MultistepAdamOptimizerTest(tf.test.TestCase):
 
   def testMultistep(self):
-    ver = tf.__version__.split('.')
-    # TODO(rsepassi): Remove version check once 1.5 is not tested anymore
-    if int(ver[0]) <= 1 and int(ver[1]) < 6:
-      # MultistepAdamOptimizer requires TF >= 1.6
-      return
     dtype = tf.float32
     beta1 = 0.2
     beta2 = 0.99

From a33fc0be68dcd29326300dec483ec6c90f94bbfe Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 12 Jul 2018 10:46:09 -0700
Subject: [PATCH 0326/2720] Comment out disabled test (lint)

PiperOrigin-RevId: 204325408
---
 tensor2tensor/models/research/autoencoders_test.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index 82ac3fdab..b5717b9e6 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -83,9 +83,9 @@ def testAutoencoderOrderedDiscreteVQ(self):
 
   # TODO(lukaszkaiser): Re-enable test by conserving lost shape information
   # in autoencoder_stacked.
-  def xtestAutoencoderStacked(self):
-    res = self.get_mnist_random_output("autoencoder_stacked")
-    self.assertEqual(res.shape, self.mnist_output_shape)
+  # def testAutoencoderStacked(self):
+  #  res = self.get_mnist_random_output("autoencoder_stacked")
+  #  self.assertEqual(res.shape, self.mnist_output_shape)
 
 if __name__ == "__main__":
   tf.test.main()

From 31d867b6314965b12847978e526a096619647c10 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Thu, 12 Jul 2018 11:47:01 -0700
Subject: [PATCH 0327/2720] Add linear decay learning rate option.

PiperOrigin-RevId: 204336641
---
 tensor2tensor/utils/learning_rate.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index e5e05c90e..ec15b696b 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -28,6 +28,9 @@ def learning_rate_factor(name, step_num, hparams):
     return hparams.learning_rate_constant
   elif name == "linear_warmup":
     return tf.minimum(1.0, step_num / hparams.learning_rate_warmup_steps)
+  elif name == "linear_decay":
+    ret = (hparams.train_steps - step_num) / hparams.learning_rate_decay_steps
+    return tf.minimum(1.0, tf.maximum(0.0, ret))
   elif name == "rsqrt_decay":
     return tf.rsqrt(tf.maximum(step_num, hparams.learning_rate_warmup_steps))
   elif name == "rsqrt_hidden_size":

From da6c3e9eb9480e14546674efdbed9494a653a46b Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Thu, 12 Jul 2018 12:50:05 -0700
Subject: [PATCH 0328/2720] no public changes

PiperOrigin-RevId: 204347115
---
 tensor2tensor/data_generators/problem.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index a21ab5e58..9ce9e63d2 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -723,6 +723,14 @@ def _dataset_partition(self, mode, config):
       # Reset in the case when using TPU but alternating TRAIN and EVAL.
       self._next_partition_id = 0
       return 0, 1
+    # BEGIN GOOGLE-INTERNAL
+    # make mesh-tensorflow on TPU work with patch CL/202825176
+    # TODO(ylc): fix this once TPU estimator changes are checked in.
+    if getattr(config.tpu_config, "symmetric_sharding_enabled", False):
+      tf.logging.info("symmetric_sharding_enabled")
+      self._next_partition_id = 0
+      return 0, 1
+    # END GOOOGLE-INTERNAL
     if config.tpu_config.per_host_input_for_training:
       num_partitions = max(config.tpu_config.num_shards // 8, 1)
     else:

From 8c1a703f6ec644703734a6ea233ed5bdbc4d9997 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 12 Jul 2018 17:00:40 -0700
Subject: [PATCH 0329/2720] Make model_rl_experiment robust to restarts

PiperOrigin-RevId: 204390799
---
 tensor2tensor/data_generators/gym_problems.py | 72 ++++++++++++++++---
 tensor2tensor/rl/model_rl_experiment.py       | 20 +++---
 2 files changed, 73 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 32b4ff9e4..470658da8 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import json
 import math
 import os
 import gym
@@ -91,6 +92,7 @@ def __init__(self, *args, **kwargs):
     self._internal_memory_size = 20
     self._internal_memory_force_beginning_resets = False
     self._session = None
+    self.statistics = BasicStatistics()
 
   def _setup(self):
     # TODO(piotrmilos):this should be consistent with
@@ -110,11 +112,13 @@ def _setup(self):
       policy_to_actions_lambda = lambda policy: policy.mode()
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      self.collect_memory, self.collect_trigger_op, collect_init \
-        = collect.define_collect(collect_hparams, scope="gym_problems",
-                                 eval_phase=False, collect_level=0,
-                                 policy_to_actions_lambda
-                                 =policy_to_actions_lambda)
+      self.collect_memory, self.collect_trigger_op, collect_init = (
+          collect.define_collect(
+              collect_hparams,
+              scope="gym_problems",
+              eval_phase=False,
+              collect_level=0,
+              policy_to_actions_lambda=policy_to_actions_lambda))
 
     self._session = tf.Session()
     collect_init(self._session)
@@ -287,6 +291,16 @@ def hparams(self, defaults, unused_model_hparams):
     p.input_space_id = problem.SpaceID.IMAGE
     p.target_space_id = problem.SpaceID.IMAGE
 
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    super(GymDiscreteProblem, self).generate_data(data_dir, tmp_dir, task_id)
+    # Save stats to file, or restore if data was already generated.
+    stats_file = os.path.join(data_dir,
+                              "%s.stats.json" % self.dataset_filename())
+    if tf.gfile.Exists(stats_file):
+      self.statistics.update_from_file(stats_file)
+    else:
+      self.statistics.save_to_file(stats_file)
+
 
 class BasicStatistics(object):
   """Keeps basic statistics to calculate mean reward """
@@ -297,6 +311,35 @@ def __init__(self):
     self.sum_of_rewards_current_episode = 0.0
     self.last_done = False
 
+  def update_from_dict(self, stats_dict):
+    keys = set(self.to_dict().keys())
+    for k, v in stats_dict.items():
+      if k not in keys:
+        raise ValueError("Key %s not a property of %s" %
+                         (k, type(self).__name__))
+      setattr(self, k, v)
+    return self
+
+  def to_dict(self):
+    # Cast the values to base types as some are numpy types.
+    keys_and_types = [
+        ("sum_of_rewards", float),
+        ("number_of_dones", int),
+        ("sum_of_rewards_current_episode", float),
+        ("last_done", bool),
+    ]
+    stats_dict = dict([(k, t(getattr(self, k))) for k, t in keys_and_types])
+    return stats_dict
+
+  def save_to_file(self, fname):
+    with tf.gfile.Open(fname, "w") as f:
+      f.write(json.dumps(self.to_dict()))
+
+  def update_from_file(self, fname):
+    with tf.gfile.Open(fname) as f:
+      self.update_from_dict(json.loads(f.read()))
+      return self
+
 
 # TODO(piotrmilos): merge with the superclass
 class GymRealDiscreteProblem(GymDiscreteProblem):
@@ -316,8 +359,8 @@ def collect_statistics_and_generate_debug_image(self, index, observation,
     # we ignore consecutive dones as they are artefacts of skip wrappers
     if done and not self.statistics.last_done:
       self.statistics.number_of_dones += int(done)
-      self.statistics.sum_of_rewards +=\
-        self.statistics.sum_of_rewards_current_episode
+      self.statistics.sum_of_rewards += (
+          self.statistics.sum_of_rewards_current_episode)
       self.statistics.sum_of_rewards_current_episode = 0.0
 
     self.statistics.last_done = done
@@ -389,13 +432,24 @@ def __init__(self):
     self.real_env = None
     self.real_ob = None
 
+  def to_dict(self):
+    stats_dict = super(RewardPerSequenceStatistics, self).to_dict()
+    keys_and_types = [
+        ("episode_sim_reward", float),
+        ("episode_real_reward", float),
+        ("successful_episode_reward_predictions", int),
+        ("report_reward_statistics_every", int),
+    ]
+    additional = dict([(k, t(getattr(self, k))) for k, t in keys_and_types])
+    stats_dict.update(additional)
+    return stats_dict
+
 
 class GymSimulatedDiscreteProblem(GymDiscreteProblem):
   """Simulated gym environment with discrete actions and rewards."""
 
   def __init__(self, *args, **kwargs):
     super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
-    self.statistics = RewardPerSequenceStatistics()
     self.debug_dump_frames_path = "debug_frames_sim"
 
     # This is hackish way of introducing resets every
@@ -406,6 +460,8 @@ def __init__(self, *args, **kwargs):
     self._internal_memory_force_beginning_resets = True
     env_spec = standard_atari_env_spec(self.env_name)
     real_env = env_spec.env_lambda()
+
+    self.statistics = RewardPerSequenceStatistics()
     self.statistics.real_env = real_env
 
   def _setup(self):
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 248aa9876..a4318fe13 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -436,25 +436,23 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
           eval_phase=False)
       log("Mean reward during generation: {}".format(generation_mean_reward))
 
-    # Report metrics.
-    eval_metrics = {"model_reward_accuracy": model_reward_accuracy,
-                    "mean_reward": mean_reward}
-
-    model_reward_accuracy_summary.value[0].simple_value \
-      = model_reward_accuracy
-
-    mean_reward_summary.value[0].simple_value \
-      = mean_reward
-
+    # Summarize metrics
+    assert model_reward_accuracy is not None
+    assert mean_reward is not None
+    model_reward_accuracy_summary.value[0].simple_value = model_reward_accuracy
+    mean_reward_summary.value[0].simple_value = mean_reward
     eval_metrics_writer.add_summary(model_reward_accuracy_summary, epoch)
     eval_metrics_writer.add_summary(mean_reward_summary, epoch)
 
+    # Report metrics
+    eval_metrics = {"model_reward_accuracy": model_reward_accuracy,
+                    "mean_reward": mean_reward}
     epoch_metrics.append(eval_metrics)
     log("Eval metrics: %s", str(eval_metrics))
     if report_fn:
       report_fn(eval_metrics[report_metric], epoch)
 
-  # Report the evaluation metrics from the final epoch
+  # Return the evaluation metrics from the final epoch
   return epoch_metrics[-1]
 
 

From 9a0e884403f9fdb5e77bb007c3750cb8812a25be Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 13 Jul 2018 10:23:11 -0700
Subject: [PATCH 0330/2720] Escape reserved tokens when constructing
 SubwordTextEncoder. Fixes bug with spaces in decoding.

PiperOrigin-RevId: 204490623
---
 tensor2tensor/data_generators/text_encoder.py      |  6 +++++-
 tensor2tensor/data_generators/text_encoder_test.py | 12 ++++--------
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 7d590babd..38ca98c84 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -850,7 +850,11 @@ def build_from_token_counts(self,
       # Reinitialize to the candidate vocabulary.
       new_subtoken_strings = [subtoken for _, subtoken in new_subtoken_strings]
       if reserved_tokens:
-        new_subtoken_strings = reserved_tokens + new_subtoken_strings
+        escaped_reserved_tokens = [
+            _escape_token(native_to_unicode(t), self._alphabet)
+            for t in reserved_tokens
+        ]
+        new_subtoken_strings = escaped_reserved_tokens + new_subtoken_strings
 
       self._init_subtokens_from_list(new_subtoken_strings)
       tf.logging.info("vocab_size = %d" % self.vocab_size)
diff --git a/tensor2tensor/data_generators/text_encoder_test.py b/tensor2tensor/data_generators/text_encoder_test.py
index f7c7ed4e2..ead1de54f 100644
--- a/tensor2tensor/data_generators/text_encoder_test.py
+++ b/tensor2tensor/data_generators/text_encoder_test.py
@@ -222,10 +222,8 @@ def test_custom_reserved_tokens(self):
         10, token_counts, 2, 10, reserved_tokens=reserved_tokens)
 
     # Make sure that reserved tokens appear in the right places.
-    start_id = encoder._subtoken_string_to_id[start_symbol]
-    end_id = encoder._subtoken_string_to_id[end_symbol]
-    self.assertEqual(start_id, 2)
-    self.assertEqual(end_id, 3)
+    self.assertEqual(encoder.decode([2]), start_symbol)
+    self.assertEqual(encoder.decode([3]), end_symbol)
 
     # Make sure that we haven't messed up the ability to reconstruct.
     reconstructed_corpus = encoder.decode(encoder.encode(corpus))
@@ -354,10 +352,8 @@ def gen():
         gen(), 10, reserved_tokens=reserved_tokens)
 
     # Make sure that reserved tokens appear in the right places.
-    start_id = encoder._subtoken_string_to_id[start_symbol]
-    end_id = encoder._subtoken_string_to_id[end_symbol]
-    self.assertEqual(start_id, 2)
-    self.assertEqual(end_id, 3)
+    self.assertEqual(encoder.decode([2]), start_symbol)
+    self.assertEqual(encoder.decode([3]), end_symbol)
 
     self.assertEqual("hi%s" % start_symbol,
                      encoder.decode(encoder.encode("hi") + [2]))

From 0416dfca6969821029ae0a3aed4d947ecdfb99dd Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Fri, 13 Jul 2018 10:53:19 -0700
Subject: [PATCH 0331/2720] Task specific ids for MultiProblem.

PiperOrigin-RevId: 204495861
---
 tensor2tensor/data_generators/cola.py         |  2 +-
 tensor2tensor/data_generators/imdb.py         |  2 +-
 tensor2tensor/data_generators/lm1b.py         |  2 +-
 .../data_generators/multi_problem.py          | 14 +++++++-----
 tensor2tensor/data_generators/multinli.py     |  4 ++--
 tensor2tensor/data_generators/problem.py      | 22 +++++++++++++------
 tensor2tensor/data_generators/qnli.py         |  4 ++--
 tensor2tensor/data_generators/quora_qpairs.py |  4 ++--
 tensor2tensor/data_generators/rte.py          |  4 ++--
 tensor2tensor/data_generators/sst_binary.py   |  2 +-
 tensor2tensor/data_generators/wnli.py         |  4 ++--
 11 files changed, 37 insertions(+), 27 deletions(-)

diff --git a/tensor2tensor/data_generators/cola.py b/tensor2tensor/data_generators/cola.py
index 47e47c86b..e8e033f7a 100644
--- a/tensor2tensor/data_generators/cola.py
+++ b/tensor2tensor/data_generators/cola.py
@@ -118,4 +118,4 @@ def vocab_type(self):
 
   @property
   def task_id(self):
-    return problem.SpaceID.EN_COLA
+    return problem.TaskID.COLA
diff --git a/tensor2tensor/data_generators/imdb.py b/tensor2tensor/data_generators/imdb.py
index 1feaf40d1..6ea5715a8 100644
--- a/tensor2tensor/data_generators/imdb.py
+++ b/tensor2tensor/data_generators/imdb.py
@@ -107,4 +107,4 @@ def vocab_type(self):
 
   @property
   def task_id(self):
-    return problem.SpaceID.EN_CHR_SENT
+    return problem.TaskID.EN_CHR_SENT
diff --git a/tensor2tensor/data_generators/lm1b.py b/tensor2tensor/data_generators/lm1b.py
index d16bd109b..cebd94c95 100644
--- a/tensor2tensor/data_generators/lm1b.py
+++ b/tensor2tensor/data_generators/lm1b.py
@@ -177,7 +177,7 @@ def vocab_type(self):
 
   @property
   def task_id(self):
-    return problem.SpaceID.EN_CHR
+    return problem.TaskID.EN_CHR
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 92ccf2a89..183bfb331 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -67,6 +67,13 @@ def get_hparams(self, model_hparams=None):
 
     return self._hparams
 
+  def flatten_zip(self, *args):
+    flattened = tf.data.Dataset.from_tensors(args[0])
+    for ex in args[1:]:
+      flattened.concatenate(tf.data.Dataset.from_tensors(ex))
+
+    return flattened
+
   def dataset(self,
               mode,
               data_dir=None,
@@ -96,14 +103,9 @@ def dataset(self,
 
     self.get_hparams()
 
-    # TODO(urvashik): make this independent of the number of tasks
-    def flatten_zip(d0, d1):
-      return tf.data.Dataset.from_tensors(d0).concatenate(
-          tf.data.Dataset.from_tensors(d1))
-
     if is_training:
       single_mtl_dataset = tf.data.Dataset.zip(tuple(datasets)).flat_map(
-          flatten_zip)
+          self.flatten_zip)
     else:
       single_mtl_dataset = datasets[0]
       for data in datasets[1:]:
diff --git a/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index bed35c014..9ab08f9fa 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -74,7 +74,7 @@ def concat_token(self):
   @property
   def concat_id(self):
     if self.vocab_type == text_problems.VocabType.CHARACTER:
-      return problem.SpaceID.EN_PR_HYP
+      return problem.TaskID.EN_PR_HYP
     return 2
 
   def class_labels(self, data_dir):
@@ -135,4 +135,4 @@ def vocab_type(self):
 
   @property
   def task_id(self):
-    return problem.SpaceID.THREE_CL_NLI
+    return problem.TaskID.THREE_CL_NLI
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 9ce9e63d2..7de6c4ff1 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -100,20 +100,28 @@ class SpaceID(object):
   STROKES = 29
   # Pickled Python
   PICKLED_PYTHON = 30
+
+
+class TaskID(object):
+  """Problem specific task ids. Add more as needed."""
+  # English characters
+  EN_CHR = 0
   # English characters sentiment
-  EN_CHR_SENT = 31
+  EN_CHR_SENT = 1
   # English Premise Hypothesis pair
-  EN_PR_HYP = 32
+  EN_PR_HYP = 2
   # English NLI
-  EN_NLI = 33
+  EN_NLI = 3
   # COLA
-  COLA = 34
+  COLA = 4
   # Enligh Question Context pair
-  EN_Q_CONT = 35
+  EN_Q_CONT = 5
   # English similarity task
-  EN_SIM = 36
+  EN_SIM = 6
+  # English sentence pair
+  EN_SENT_PAIR = 7
   # 3 class NLI
-  THREE_CL_NLI = 37
+  THREE_CL_NLI = 8
 
 
 def default_model_hparams():
diff --git a/tensor2tensor/data_generators/qnli.py b/tensor2tensor/data_generators/qnli.py
index ed7ed604b..101f1dba2 100644
--- a/tensor2tensor/data_generators/qnli.py
+++ b/tensor2tensor/data_generators/qnli.py
@@ -74,7 +74,7 @@ def concat_token(self):
   @property
   def concat_id(self):
     if self.vocab_type == text_problems.VocabType.CHARACTER:
-      return problem.SpaceID.EN_Q_CONT
+      return problem.TaskID.EN_Q_CONT
     return 2
 
   def class_labels(self, data_dir):
@@ -132,4 +132,4 @@ def vocab_type(self):
 
   @property
   def task_id(self):
-    return problem.SpaceID.EN_NLI
+    return problem.TaskID.EN_NLI
diff --git a/tensor2tensor/data_generators/quora_qpairs.py b/tensor2tensor/data_generators/quora_qpairs.py
index 4aff774a8..489fa8650 100644
--- a/tensor2tensor/data_generators/quora_qpairs.py
+++ b/tensor2tensor/data_generators/quora_qpairs.py
@@ -74,7 +74,7 @@ def concat_token(self):
   @property
   def concat_id(self):
     if self.vocab_type == text_problems.VocabType.CHARACTER:
-      return problem.SpaceID.EN_CHR
+      return problem.TaskID.EN_SENT_PAIR
     return 2
 
   def class_labels(self, data_dir):
@@ -138,4 +138,4 @@ def vocab_type(self):
 
   @property
   def task_id(self):
-    return problem.SpaceID.EN_SIM
+    return problem.TaskID.EN_SIM
diff --git a/tensor2tensor/data_generators/rte.py b/tensor2tensor/data_generators/rte.py
index 474c30b2f..bc36a45ab 100644
--- a/tensor2tensor/data_generators/rte.py
+++ b/tensor2tensor/data_generators/rte.py
@@ -74,7 +74,7 @@ def concat_token(self):
   @property
   def concat_id(self):
     if self.vocab_type == text_problems.VocabType.CHARACTER:
-      return problem.SpaceID.EN_PR_HYP
+      return problem.TaskID.EN_PR_HYP
     return 2
 
   def class_labels(self, data_dir):
@@ -132,4 +132,4 @@ def vocab_type(self):
 
   @property
   def task_id(self):
-    return problem.SpaceID.EN_NLI
+    return problem.TaskID.EN_NLI
diff --git a/tensor2tensor/data_generators/sst_binary.py b/tensor2tensor/data_generators/sst_binary.py
index ffd934d1f..ce5fcc586 100644
--- a/tensor2tensor/data_generators/sst_binary.py
+++ b/tensor2tensor/data_generators/sst_binary.py
@@ -119,4 +119,4 @@ def vocab_type(self):
 
   @property
   def task_id(self):
-    return problem.SpaceID.EN_CHR_SENT
+    return problem.TaskID.EN_CHR_SENT
diff --git a/tensor2tensor/data_generators/wnli.py b/tensor2tensor/data_generators/wnli.py
index b4945d1e5..0a6a4caf2 100644
--- a/tensor2tensor/data_generators/wnli.py
+++ b/tensor2tensor/data_generators/wnli.py
@@ -74,7 +74,7 @@ def concat_token(self):
   @property
   def concat_id(self):
     if self.vocab_type == text_problems.VocabType.CHARACTER:
-      return problem.SpaceID.EN_PR_HYP
+      return problem.TaskID.EN_PR_HYP
     return 2
 
   def class_labels(self, data_dir):
@@ -130,4 +130,4 @@ def vocab_type(self):
 
   @property
   def task_id(self):
-    return problem.SpaceID.EN_NLI
+    return problem.TaskID.EN_NLI

From 0386c4974298cea770e0ca0bfffc1b58c34451fa Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 13 Jul 2018 14:58:01 -0700
Subject: [PATCH 0332/2720] Separating the reward model.

PiperOrigin-RevId: 204533573
---
 tensor2tensor/models/research/next_frame.py | 175 +++++++++++++-------
 1 file changed, 116 insertions(+), 59 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index c81264711..8b338654f 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -221,18 +221,102 @@ def construct_latent_tower(self, images):
 
       return mean, std
 
-  def reward_prediction(self, inputs):
+  def bottom_part_tower(self, input_image, input_reward, action, latent,
+                        lstm_state, lstm_size, conv_size):
+    """The bottom part of predictive towers.
+
+    With the current (early) design, the main prediction tower and
+    the reward prediction tower share the same arcitecture. TF Scope can be
+    adjusted as required to either share or not share the weights between
+    the two towers.
+
+    Args:
+      input_image: the current image.
+      input_reward: the current reward.
+      action: the action taken by the agent.
+      latent: the latent vector.
+      lstm_state: the current internal states of conv lstms.
+      lstm_size: the size of lstms.
+      conv_size: the size of convolutions.
+
+    Returns:
+      - the output of the partial network.
+      - intermidate outputs for skip connections.
+    """
+    layer_norm = tf.contrib.layers.layer_norm
+    lstm_func = self.conv_lstm_2d
+
+    input_image = common_layers.make_even_size(input_image)
+    enc0 = slim.layers.conv2d(
+        input_image,
+        conv_size[0], [5, 5],
+        stride=2,
+        scope="scale1_conv1",
+        normalizer_fn=layer_norm,
+        normalizer_params={"scope": "layer_norm1"})
+
+    hidden1, lstm_state[0] = lstm_func(
+        enc0, lstm_state[0], lstm_size[0], scope="state1")
+    hidden1 = layer_norm(hidden1, scope="layer_norm2")
+    hidden2, lstm_state[1] = lstm_func(
+        hidden1, lstm_state[1], lstm_size[1], scope="state2")
+    hidden2 = layer_norm(hidden2, scope="layer_norm3")
+    hidden2 = common_layers.make_even_size(hidden2)
+    enc1 = slim.layers.conv2d(
+        hidden2, hidden2.get_shape()[3], [3, 3], stride=2, scope="conv2")
+
+    hidden3, lstm_state[2] = lstm_func(
+        enc1, lstm_state[2], lstm_size[2], scope="state3")
+    hidden3 = layer_norm(hidden3, scope="layer_norm4")
+    hidden4, lstm_state[3] = lstm_func(
+        hidden3, lstm_state[3], lstm_size[3], scope="state4")
+    hidden4 = layer_norm(hidden4, scope="layer_norm5")
+    hidden4 = common_layers.make_even_size(hidden4)
+    enc2 = slim.layers.conv2d(
+        hidden4, hidden4.get_shape()[3], [3, 3], stride=2, scope="conv3")
+
+    # Pass in reward and action.
+    emb_action = self.encode_to_shape(action, enc2.get_shape(), "action_enc")
+    emb_reward = self.encode_to_shape(
+        input_reward, enc2.get_shape(), "reward_enc")
+    enc2 = tf.concat(axis=3, values=[enc2, emb_action, emb_reward])
+
+    if latent is not None:
+      with tf.control_dependencies([latent]):
+        enc2 = tf.concat([enc2, latent], 3)
+
+    enc3 = slim.layers.conv2d(
+        enc2, hidden4.get_shape()[3], [1, 1], stride=1, scope="conv4")
+
+    hidden5, lstm_state[4] = lstm_func(
+        enc3, lstm_state[4], lstm_size[4], scope="state5")  # last 8x8
+    hidden5 = layer_norm(hidden5, scope="layer_norm6")
+
+    return hidden5, (enc0, enc1)
+
+  def reward_prediction(
+      self, input_image, input_reward, action, lstm_state, latent):
     """Builds a reward prediction network."""
-    conv_size = self.tinyify([32, 16, 1])
+    conv_size = self.tinyify([32, 32, 16, 4])
+    lstm_size = self.tinyify([32, 64, 128, 64, 32])
+
     with tf.variable_scope("reward_pred", reuse=tf.AUTO_REUSE):
-      x = inputs
+      hidden5, _ = self.bottom_part_tower(
+          input_image, input_reward, action, latent,
+          lstm_state, lstm_size, conv_size)
+
+      x = hidden5
       x = slim.batch_norm(x, scope="reward_bn0")
-      x = slim.conv2d(x, conv_size[0], [3, 3], scope="reward_conv1")
+      x = slim.conv2d(x, conv_size[1], [3, 3], scope="reward_conv1")
       x = slim.batch_norm(x, scope="reward_bn1")
-      x = slim.conv2d(x, conv_size[1], [3, 3], scope="reward_conv2")
+      x = slim.conv2d(x, conv_size[2], [3, 3], scope="reward_conv2")
       x = slim.batch_norm(x, scope="reward_bn2")
-      x = slim.conv2d(x, conv_size[2], [3, 3], scope="reward_conv3")
-    return x
+      x = slim.conv2d(x, conv_size[3], [3, 3], scope="reward_conv3")
+
+      pred_reward = self.decode_to_shape(
+          x, input_reward.shape, "reward_dec")
+
+      return pred_reward, lstm_state
 
   def encode_to_shape(self, inputs, shape, scope):
     """Encode the given tensor to given image shape."""
@@ -280,51 +364,11 @@ def construct_predictive_tower(
     img_height, img_width, color_channels = self.hparams.problem.frame_shape
 
     with tf.variable_scope("main", reuse=tf.AUTO_REUSE):
-      input_image = common_layers.make_even_size(input_image)
-      enc0 = slim.layers.conv2d(
-          input_image,
-          conv_size[0], [5, 5],
-          stride=2,
-          scope="scale1_conv1",
-          normalizer_fn=layer_norm,
-          normalizer_params={"scope": "layer_norm1"})
-
-      hidden1, lstm_state[0] = lstm_func(
-          enc0, lstm_state[0], lstm_size[0], scope="state1")
-      hidden1 = layer_norm(hidden1, scope="layer_norm2")
-      hidden2, lstm_state[1] = lstm_func(
-          hidden1, lstm_state[1], lstm_size[1], scope="state2")
-      hidden2 = layer_norm(hidden2, scope="layer_norm3")
-      hidden2 = common_layers.make_even_size(hidden2)
-      enc1 = slim.layers.conv2d(
-          hidden2, hidden2.get_shape()[3], [3, 3], stride=2, scope="conv2")
-
-      hidden3, lstm_state[2] = lstm_func(
-          enc1, lstm_state[2], lstm_size[2], scope="state3")
-      hidden3 = layer_norm(hidden3, scope="layer_norm4")
-      hidden4, lstm_state[3] = lstm_func(
-          hidden3, lstm_state[3], lstm_size[3], scope="state4")
-      hidden4 = layer_norm(hidden4, scope="layer_norm5")
-      hidden4 = common_layers.make_even_size(hidden4)
-      enc2 = slim.layers.conv2d(
-          hidden4, hidden4.get_shape()[3], [3, 3], stride=2, scope="conv3")
-
-      # Pass in reward and action.
-      emb_action = self.encode_to_shape(action, enc2.get_shape(), "action_enc")
-      emb_reward = self.encode_to_shape(
-          input_reward, enc2.get_shape(), "reward_enc")
-      enc2 = tf.concat(axis=3, values=[enc2, emb_action, emb_reward])
-
-      if latent is not None:
-        with tf.control_dependencies([latent]):
-          enc2 = tf.concat([enc2, latent], 3)
-
-      enc3 = slim.layers.conv2d(
-          enc2, hidden4.get_shape()[3], [1, 1], stride=1, scope="conv4")
-
-      hidden5, lstm_state[4] = lstm_func(
-          enc3, lstm_state[4], lstm_size[4], scope="state5")  # last 8x8
-      hidden5 = layer_norm(hidden5, scope="layer_norm6")
+      hidden5, skips = self.bottom_part_tower(
+          input_image, input_reward, action, latent,
+          lstm_state, lstm_size, conv_size)
+      enc0, enc1 = skips
+
       enc4 = slim.layers.conv2d_transpose(
           hidden5, hidden5.get_shape()[3], 3, stride=2, scope="convt1")
 
@@ -404,11 +448,7 @@ def construct_predictive_tower(
       for layer, mask in zip(transformed, mask_list[1:]):
         output += layer * mask
 
-      p_reward = self.reward_prediction(hidden5)
-      p_reward = self.decode_to_shape(
-          p_reward, input_reward.shape, "reward_dec")
-
-      return output, p_reward, lstm_state
+      return output, lstm_state
 
   def get_guassian_latent(self, latent_mean, latent_std):
     latent = tf.random_normal(tf.shape(latent_mean), 0, 1, dtype=tf.float32)
@@ -445,6 +485,7 @@ def construct_model(self,
 
     # LSTM states.
     lstm_state = [None] * 7
+    reward_lstm_state = [None] * 5
 
     # Latent tower
     if self.hparams.stochastic_model:
@@ -466,8 +507,15 @@ def construct_model(self,
           latent = self.get_guassian_latent(latent_mean, latent_std)
 
       # Prediction
-      pred_image, pred_reward, lstm_state = self.construct_predictive_tower(
+      pred_image, lstm_state = self.construct_predictive_tower(
           input_image, input_reward, action, lstm_state, latent)
+
+      if self.hparams.reward_prediction:
+        pred_reward, reward_lstm_state = self.reward_prediction(
+            input_image, input_reward, action, reward_lstm_state, latent)
+      else:
+        pred_reward = input_reward
+
       gen_images.append(pred_image)
       gen_rewards.append(pred_reward)
 
@@ -733,6 +781,7 @@ def construct_model(self, images, actions, rewards):
 
     # LSTM states.
     lstm_state = [None] * 7
+    reward_lstm_state = [None] * 5
 
     pred_image, pred_reward, latent = None, None, None
     for timestep, image, action, reward in zip(
@@ -753,8 +802,15 @@ def construct_model(self, images, actions, rewards):
       latent_stds.append(latent_std)
 
       # Prediction
-      pred_image, pred_reward, lstm_state = self.construct_predictive_tower(
+      pred_image, lstm_state = self.construct_predictive_tower(
           input_image, input_reward, action, lstm_state, latent)
+
+      if self.hparams.reward_prediction:
+        pred_reward, reward_lstm_state = self.reward_prediction(
+            input_image, input_reward, action, reward_lstm_state, latent)
+      else:
+        pred_reward = input_reward
+
       gen_images.append(pred_image)
       gen_rewards.append(pred_reward)
 
@@ -1064,6 +1120,7 @@ def next_frame_stochastic():
   hparams.input_modalities = "inputs:video:l2raw"
   hparams.video_modality_loss_cutoff = 0.0
   hparams.add_hparam("stochastic_model", True)
+  hparams.add_hparam("reward_prediction", True)
   hparams.add_hparam("model_options", "CDNA")
   hparams.add_hparam("num_masks", 10)
   hparams.add_hparam("latent_channels", 1)

From 0ca95698d4a5ef9922f91b7265979559cd5f444c Mon Sep 17 00:00:00 2001
From: Stefan Schweter <stefan@schweter.it>
Date: Sat, 14 Jul 2018 21:23:10 +0200
Subject: [PATCH 0333/2720] readme: add universal transformer paper

---
 README.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/README.md b/README.md
index 10359f764..6fe27fc4a 100644
--- a/README.md
+++ b/README.md
@@ -402,5 +402,6 @@ T2T](https://research.googleblog.com/2017/06/accelerating-deep-learning-research
 * [Self-Attention with Relative Position Representations](https://arxiv.org/abs/1803.02155)
 * [Fast Decoding in Sequence Models using Discrete Latent Variables](https://arxiv.org/abs/1803.03382)
 * [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://arxiv.org/abs/1804.04235)
+* [Universal Transformers](https://arxiv.org/abs/1807.03819)
 
 *Note: This is not an official Google product.*

From f70e4ce623d206e4df35e4f7021c1a4cef5466e3 Mon Sep 17 00:00:00 2001
From: Stefan Schweter <stefan@schweter.it>
Date: Sat, 14 Jul 2018 21:23:27 +0200
Subject: [PATCH 0334/2720] doc: add universal transformer paper

---
 docs/walkthrough.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 10359f764..6fe27fc4a 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -402,5 +402,6 @@ T2T](https://research.googleblog.com/2017/06/accelerating-deep-learning-research
 * [Self-Attention with Relative Position Representations](https://arxiv.org/abs/1803.02155)
 * [Fast Decoding in Sequence Models using Discrete Latent Variables](https://arxiv.org/abs/1803.03382)
 * [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://arxiv.org/abs/1804.04235)
+* [Universal Transformers](https://arxiv.org/abs/1807.03819)
 
 *Note: This is not an official Google product.*

From 21d5f79212380bfc74ccf4ba7869158b3788d9ad Mon Sep 17 00:00:00 2001
From: Stefan Schweter <stefan@schweter.it>
Date: Sat, 14 Jul 2018 21:23:44 +0200
Subject: [PATCH 0335/2720] models: add universal transformer paper

---
 tensor2tensor/models/research/universal_transformer.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 19ec7c97d..9558122a3 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 """Universal Transformers.
 
+Universal Transformer is described in https://arxiv.org/abs/1807.03819.
 
 Universal Transformer is recurrent in depth while employing self-attention
 to combine information from different parts of sequences.

From 8e735aa52517041147bba08c1cf54aa700e2c522 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 16 Jul 2018 13:32:36 -0700
Subject: [PATCH 0336/2720] modify generator_utils.py and translate.py, so that
 translate problem has a consistent vocab generator API as other text2text
 problems.

PiperOrigin-RevId: 204797793
---
 .../data_generators/generator_utils.py        | 92 ++++++++++---------
 tensor2tensor/data_generators/translate.py    | 10 +-
 2 files changed, 51 insertions(+), 51 deletions(-)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index f8dd1f12a..04716a94c 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -343,52 +343,54 @@ def get_or_generate_vocab(data_dir, tmp_dir, vocab_filename, vocab_size,
                           sources, file_byte_budget=1e6):
   """Generate a vocabulary from the datasets in sources."""
 
-  def generate():
-    """Generate lines for vocabulary generation."""
-    tf.logging.info("Generating vocab from: %s", str(sources))
-    for source in sources:
-      url = source[0]
-      filename = os.path.basename(url)
-      compressed_file = maybe_download(tmp_dir, filename, url)
-
-      for lang_file in source[1]:
-        tf.logging.info("Reading file: %s" % lang_file)
-        filepath = os.path.join(tmp_dir, lang_file)
-
-        # Extract from tar if needed.
-        if not tf.gfile.Exists(filepath):
-          read_type = "r:gz" if filename.endswith("tgz") else "r"
-          with tarfile.open(compressed_file, read_type) as corpus_tar:
-            corpus_tar.extractall(tmp_dir)
-
-        # For some datasets a second extraction is necessary.
-        if lang_file.endswith(".gz"):
-          new_filepath = os.path.join(tmp_dir, lang_file[:-3])
-          if tf.gfile.Exists(new_filepath):
-            tf.logging.info(
-                "Subdirectory %s already exists, skipping unpacking" % filepath)
-          else:
-            tf.logging.info("Unpacking subdirectory %s" % filepath)
-            gunzip_file(filepath, new_filepath)
-          filepath = new_filepath
-
-        with tf.gfile.GFile(filepath, mode="r") as source_file:
-          file_byte_budget_ = file_byte_budget
-          counter = 0
-          countermax = int(source_file.size() / file_byte_budget_ / 2)
-          for line in source_file:
-            if counter < countermax:
-              counter += 1
-            else:
-              if file_byte_budget_ <= 0:
-                break
-              line = line.strip()
-              file_byte_budget_ -= len(line)
-              counter = 0
-              yield line
-
+  vocab_generator = generate_lines_for_vocab(tmp_dir, sources, file_byte_budget)
   return get_or_generate_vocab_inner(data_dir, vocab_filename, vocab_size,
-                                     generate())
+                                     vocab_generator)
+
+
+def generate_lines_for_vocab(tmp_dir, sources, file_byte_budget=1e6):
+  """Generate lines for vocabulary generation."""
+  tf.logging.info("Generating vocab from: %s", str(sources))
+  for source in sources:
+    url = source[0]
+    filename = os.path.basename(url)
+    compressed_file = maybe_download(tmp_dir, filename, url)
+
+    for lang_file in source[1]:
+      tf.logging.info("Reading file: %s" % lang_file)
+      filepath = os.path.join(tmp_dir, lang_file)
+
+      # Extract from tar if needed.
+      if not tf.gfile.Exists(filepath):
+        read_type = "r:gz" if filename.endswith("tgz") else "r"
+        with tarfile.open(compressed_file, read_type) as corpus_tar:
+          corpus_tar.extractall(tmp_dir)
+
+      # For some datasets a second extraction is necessary.
+      if lang_file.endswith(".gz"):
+        new_filepath = os.path.join(tmp_dir, lang_file[:-3])
+        if tf.gfile.Exists(new_filepath):
+          tf.logging.info(
+              "Subdirectory %s already exists, skipping unpacking" % filepath)
+        else:
+          tf.logging.info("Unpacking subdirectory %s" % filepath)
+          gunzip_file(filepath, new_filepath)
+        filepath = new_filepath
+
+      with tf.gfile.GFile(filepath, mode="r") as source_file:
+        file_byte_budget_ = file_byte_budget
+        counter = 0
+        countermax = int(source_file.size() / file_byte_budget_ / 2)
+        for line in source_file:
+          if counter < countermax:
+            counter += 1
+          else:
+            if file_byte_budget_ <= 0:
+              break
+            line = line.strip()
+            file_byte_budget_ -= len(line)
+            counter = 0
+            yield line
 
 
 def get_or_generate_tabbed_vocab(data_dir, tmp_dir, source_filename,
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 28984b6fa..cf3272604 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -53,15 +53,13 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
     tag = "train" if dataset_split == problem.DatasetSplit.TRAIN else "dev"
     data_path = compile_data(tmp_dir, datasets, "%s-compiled-%s" % (self.name,
                                                                     tag))
-
-    if self.vocab_type == text_problems.VocabType.SUBWORD:
-      generator_utils.get_or_generate_vocab(
-          data_dir, tmp_dir, self.vocab_filename, self.approx_vocab_size,
-          self.vocab_data_files())
-
     return text_problems.text2text_txt_iterator(data_path + ".lang1",
                                                 data_path + ".lang2")
 
+  def generate_text_for_vocab(self, data_dir, tmp_dir):
+    return generator_utils.generate_lines_for_vocab(tmp_dir,
+                                                    self.vocab_data_files())
+
 
 def _preprocess_sgm(line, is_sgm):
   """Preprocessing to strip tags in SGM files."""

From 95151619156d117b6cad6a91b670b644a6e180d0 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 16 Jul 2018 14:11:57 -0700
Subject: [PATCH 0337/2720] Internal change

PiperOrigin-RevId: 204804923
---
 tensor2tensor/bin/t2t_decoder.py | 1 +
 tensor2tensor/utils/decoding.py  | 7 ++++---
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 8d0693150..8df00b163 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -75,6 +75,7 @@ def create_decode_hparams():
   decode_hp = decoding.decode_hparams(FLAGS.decode_hparams)
   decode_hp.shards = FLAGS.decode_shards
   decode_hp.shard_id = FLAGS.worker_id
+  decode_hp.decode_in_memory = FLAGS.decode_in_memory
   return decode_hp
 
 
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 8e17e22d9..2ee8ff753 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -55,6 +55,7 @@ def decode_hparams(overrides=""):
       num_samples=-1,
       delimiter="\n",
       decode_to_file=None,
+      decode_in_memory=False,
       shards=1,
       shard_id=0,
       num_decodes=1,
@@ -168,7 +169,7 @@ def decode_from_dataset(estimator,
     tf.logging.info("Decoding {}".format(decode_id))
 
     # Create decode directory if not in-memory decoding.
-    if not FLAGS.decode_in_memory:
+    if not decode_hp.decode_in_memory:
       output_dir = os.path.join(estimator.model_dir, "decode_%05d" % decode_id)
       tf.gfile.MakeDirs(output_dir)
       output_dirs.append(output_dir)
@@ -180,9 +181,9 @@ def decode_from_dataset(estimator,
                          decode_hp,
                          decode_to_file,
                          output_dir,
-                         log_results=not FLAGS.decode_in_memory)
+                         log_results=not decode_hp.decode_in_memory)
 
-    if FLAGS.decode_in_memory:
+    if decode_hp.decode_in_memory:
       output_dirs = [output_dir]
       predictions.append(result)
 

From 962d932d4db4e7662454b104a57446c3ad777b5a Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 16 Jul 2018 17:18:24 -0700
Subject: [PATCH 0338/2720] Fix docstring for Transformer body().

PiperOrigin-RevId: 204834704
---
 tensor2tensor/models/transformer.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index c6963b369..9a61c7a4c 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -53,8 +53,8 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
     """Encode transformer inputs.
 
     Args:
-      inputs: Transformer inputs [batch_size, input_length, input_height,
-        hidden_dim] which will be flattened along the two spatial dimensions.
+      inputs: Transformer inputs [batch_size, input_length, 1, hidden_dim] which
+        will be flattened along the two spatial dimensions.
       target_space: scalar, target space ID.
       hparams: hyperparameters for model.
       features: optionally pass the entire features dictionary as well.
@@ -148,9 +148,10 @@ def body(self, features):
 
     Args:
       features: Map of features to the model. Should contain the following:
-          "inputs": Transformer inputs [batch_size, input_length, hidden_dim]
+          "inputs": Transformer inputs.
+              [batch_size, input_length, 1, hidden_dim].
           "targets": Target decoder outputs.
-              [batch_size, decoder_length, hidden_dim]
+              [batch_size, decoder_length, 1, hidden_dim]
           "target_space_id": A scalar int from data_generators.problem.SpaceID.
 
     Returns:

From 1ab149ca44290a3046f10c47bec79d8b53ccd051 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 16 Jul 2018 23:00:29 -0700
Subject: [PATCH 0339/2720] Removing Slim :( and switching to tf.layers and
 tf.contrib.layers.

PiperOrigin-RevId: 204862278
---
 tensor2tensor/models/research/next_frame.py | 260 +++++++++++---------
 1 file changed, 139 insertions(+), 121 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 8b338654f..d58f96505 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -27,7 +27,8 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow as tf
-slim = tf.contrib.slim
+tfl = tf.layers
+tfcl = tf.contrib.layers
 
 
 @registry.register_model
@@ -169,6 +170,10 @@ class NextFrameStochastic(NextFrameBasic):
   https://arxiv.org/abs/1710.11252
   """
 
+  @property
+  def is_training(self):
+    return self.hparams.mode == tf.estimator.ModeKeys.TRAIN
+
   def tinyify(self, array):
     if self.hparams.tiny_mode:
       return [1 for _ in array]
@@ -201,18 +206,25 @@ def construct_latent_tower(self, images):
 
       x = images
       x = common_layers.make_even_size(x)
-      x = slim.conv2d(x, conv_size[0], [3, 3], stride=2, scope="latent_conv1")
-      x = slim.batch_norm(x, scope="latent_bn1")
+      x = tfl.conv2d(x, conv_size[0], [3, 3], strides=(2, 2),
+                     padding="SAME", activation=tf.nn.relu, name="latent_conv1")
+      x = tfl.batch_normalization(x,
+                                  training=self.is_training, name="latent_bn1")
       x = common_layers.make_even_size(x)
-      x = slim.conv2d(x, conv_size[1], [3, 3], stride=2, scope="latent_conv2")
-      x = slim.batch_norm(x, scope="latent_bn2")
-      x = slim.conv2d(x, conv_size[2], [3, 3], stride=1, scope="latent_conv3")
-      x = slim.batch_norm(x, scope="latent_bn3")
+      x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
+                     padding="SAME", activation=tf.nn.relu, name="latent_conv2")
+      x = tfl.batch_normalization(x,
+                                  training=self.is_training, name="latent_bn2")
+      x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(1, 1),
+                     padding="SAME", activation=tf.nn.relu, name="latent_conv3")
+      x = tfl.batch_normalization(x,
+                                  training=self.is_training, name="latent_bn3")
 
       nc = self.hparams.latent_channels
-      mean = slim.conv2d(
-          x, nc, [3, 3], stride=2, activation_fn=None, scope="latent_mean")
-      std = slim.conv2d(x, nc, [3, 3], stride=2, scope="latent_std")
+      mean = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
+                        padding="SAME", activation=None, name="latent_mean")
+      std = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
+                       padding="SAME", activation=tf.nn.relu, name="latent_std")
       std += self.hparams.latent_std_min
 
       # No latent tower at inference time, just standard gaussian.
@@ -243,37 +255,37 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
       - the output of the partial network.
       - intermidate outputs for skip connections.
     """
-    layer_norm = tf.contrib.layers.layer_norm
     lstm_func = self.conv_lstm_2d
 
     input_image = common_layers.make_even_size(input_image)
-    enc0 = slim.layers.conv2d(
+    enc0 = tfl.conv2d(
         input_image,
         conv_size[0], [5, 5],
-        stride=2,
-        scope="scale1_conv1",
-        normalizer_fn=layer_norm,
-        normalizer_params={"scope": "layer_norm1"})
+        strides=(2, 2),
+        activation=tf.nn.relu,
+        padding="SAME",
+        name="scale1_conv1")
+    enc0 = tfcl.layer_norm(enc0, scope="layer_norm1")
 
     hidden1, lstm_state[0] = lstm_func(
-        enc0, lstm_state[0], lstm_size[0], scope="state1")
-    hidden1 = layer_norm(hidden1, scope="layer_norm2")
+        enc0, lstm_state[0], lstm_size[0], name="state1")
+    hidden1 = tfcl.layer_norm(hidden1, scope="layer_norm2")
     hidden2, lstm_state[1] = lstm_func(
-        hidden1, lstm_state[1], lstm_size[1], scope="state2")
-    hidden2 = layer_norm(hidden2, scope="layer_norm3")
+        hidden1, lstm_state[1], lstm_size[1], name="state2")
+    hidden2 = tfcl.layer_norm(hidden2, scope="layer_norm3")
     hidden2 = common_layers.make_even_size(hidden2)
-    enc1 = slim.layers.conv2d(
-        hidden2, hidden2.get_shape()[3], [3, 3], stride=2, scope="conv2")
+    enc1 = tfl.conv2d(hidden2, hidden2.get_shape()[3], [3, 3], strides=(2, 2),
+                      padding="SAME", activation=tf.nn.relu, name="conv2")
 
     hidden3, lstm_state[2] = lstm_func(
-        enc1, lstm_state[2], lstm_size[2], scope="state3")
-    hidden3 = layer_norm(hidden3, scope="layer_norm4")
+        enc1, lstm_state[2], lstm_size[2], name="state3")
+    hidden3 = tfcl.layer_norm(hidden3, scope="layer_norm4")
     hidden4, lstm_state[3] = lstm_func(
-        hidden3, lstm_state[3], lstm_size[3], scope="state4")
-    hidden4 = layer_norm(hidden4, scope="layer_norm5")
+        hidden3, lstm_state[3], lstm_size[3], name="state4")
+    hidden4 = tfcl.layer_norm(hidden4, scope="layer_norm5")
     hidden4 = common_layers.make_even_size(hidden4)
-    enc2 = slim.layers.conv2d(
-        hidden4, hidden4.get_shape()[3], [3, 3], stride=2, scope="conv3")
+    enc2 = tfl.conv2d(hidden4, hidden4.get_shape()[3], [3, 3], strides=(2, 2),
+                      padding="SAME", activation=tf.nn.relu, name="conv3")
 
     # Pass in reward and action.
     emb_action = self.encode_to_shape(action, enc2.get_shape(), "action_enc")
@@ -285,12 +297,12 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
       with tf.control_dependencies([latent]):
         enc2 = tf.concat([enc2, latent], 3)
 
-    enc3 = slim.layers.conv2d(
-        enc2, hidden4.get_shape()[3], [1, 1], stride=1, scope="conv4")
+    enc3 = tfl.conv2d(enc2, hidden4.get_shape()[3], [1, 1], strides=(1, 1),
+                      padding="SAME", activation=tf.nn.relu, name="conv4")
 
     hidden5, lstm_state[4] = lstm_func(
-        enc3, lstm_state[4], lstm_size[4], scope="state5")  # last 8x8
-    hidden5 = layer_norm(hidden5, scope="layer_norm6")
+        enc3, lstm_state[4], lstm_size[4], name="state5")  # last 8x8
+    hidden5 = tfcl.layer_norm(hidden5, scope="layer_norm6")
 
     return hidden5, (enc0, enc1)
 
@@ -306,12 +318,18 @@ def reward_prediction(
           lstm_state, lstm_size, conv_size)
 
       x = hidden5
-      x = slim.batch_norm(x, scope="reward_bn0")
-      x = slim.conv2d(x, conv_size[1], [3, 3], scope="reward_conv1")
-      x = slim.batch_norm(x, scope="reward_bn1")
-      x = slim.conv2d(x, conv_size[2], [3, 3], scope="reward_conv2")
-      x = slim.batch_norm(x, scope="reward_bn2")
-      x = slim.conv2d(x, conv_size[3], [3, 3], scope="reward_conv3")
+      x = tfl.batch_normalization(x,
+                                  training=self.is_training, name="reward_bn0")
+      x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
+                     padding="SAME", activation=tf.nn.relu, name="reward_conv1")
+      x = tfl.batch_normalization(x,
+                                  training=self.is_training, name="reward_bn1")
+      x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(2, 2),
+                     padding="SAME", activation=tf.nn.relu, name="reward_conv2")
+      x = tfl.batch_normalization(x,
+                                  training=self.is_training, name="reward_bn2")
+      x = tfl.conv2d(x, conv_size[3], [3, 3], strides=(2, 2),
+                     padding="SAME", activation=tf.nn.relu, name="reward_conv3")
 
       pred_reward = self.decode_to_shape(
           x, input_reward.shape, "reward_dec")
@@ -323,8 +341,8 @@ def encode_to_shape(self, inputs, shape, scope):
     with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
       w, h = shape[1].value, shape[2].value
       x = inputs
-      x = tf.contrib.layers.flatten(x)
-      x = slim.fully_connected(x, w * h, scope="encoding_full")
+      x = tfcl.flatten(x)
+      x = tfl.dense(x, w * h, activation=tf.nn.relu, name="enc_dense")
       x = tf.reshape(x, (-1, w, h, 1))
       return x
 
@@ -332,17 +350,17 @@ def decode_to_shape(self, inputs, shape, scope):
     """Encode the given tensor to given image shape."""
     with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
       x = inputs
-      x = tf.contrib.layers.flatten(x)
-      x = slim.fully_connected(x, shape[2].value, scope="decoding_full")
+      x = tfcl.flatten(x)
+      x = tfl.dense(x, shape[2].value, activation=tf.nn.relu, name="dec_dense")
       x = tf.expand_dims(x, axis=1)
       return x
 
   def conv_lstm_2d(self, inputs, state, output_channels,
-                   kernel_size=5, scope=None):
+                   kernel_size=5, name=None):
     input_shape = common_layers.shape_list(inputs)
     cell = tf.contrib.rnn.ConvLSTMCell(
         2, input_shape[1:], output_channels,
-        [kernel_size, kernel_size], name=scope)
+        [kernel_size, kernel_size], name=name)
     if state is None:
       state = cell.zero_state(input_shape[0], tf.float32)
     outputs, new_state = cell(inputs, state)
@@ -351,7 +369,6 @@ def conv_lstm_2d(self, inputs, state, output_channels,
   def construct_predictive_tower(
       self, input_image, input_reward, action, lstm_state, latent):
     # Main tower
-    layer_norm = tf.contrib.layers.layer_norm
     lstm_func = self.conv_lstm_2d
     batch_size = common_layers.shape_list(input_image)[0]
     # the number of different pixel motion predictions
@@ -369,63 +386,66 @@ def construct_predictive_tower(
           lstm_state, lstm_size, conv_size)
       enc0, enc1 = skips
 
-      enc4 = slim.layers.conv2d_transpose(
-          hidden5, hidden5.get_shape()[3], 3, stride=2, scope="convt1")
+      enc4 = tfl.conv2d_transpose(
+          hidden5, hidden5.get_shape()[3], 3, strides=2, name="convt1")
 
       enc1_shape = common_layers.shape_list(enc1)
       enc4 = enc4[:, :enc1_shape[1], :enc1_shape[2], :]  # Cut to shape.
       hidden6, lstm_state[5] = lstm_func(
-          enc4, lstm_state[5], lstm_size[5], scope="state6")  # 16x16
-      hidden6 = layer_norm(hidden6, scope="layer_norm7")
+          enc4, lstm_state[5], lstm_size[5], name="state6")  # 16x16
+      hidden6 = tfcl.layer_norm(hidden6, scope="layer_norm7")
       # Skip connection.
       hidden6 = tf.concat(axis=3, values=[hidden6, enc1])  # both 16x16
 
-      enc5 = slim.layers.conv2d_transpose(
-          hidden6, hidden6.get_shape()[3], 3, stride=2, scope="convt2")
+      enc5 = tfl.conv2d_transpose(
+          hidden6, hidden6.get_shape()[3], [3, 3], strides=(2, 2),
+          padding="SAME", activation=tf.nn.relu, name="convt2")
       enc0_shape = common_layers.shape_list(enc0)
       enc5 = enc5[:, :enc0_shape[1], :enc0_shape[2], :]  # Cut to shape.
       hidden7, lstm_state[6] = lstm_func(
-          enc5, lstm_state[6], lstm_size[6], scope="state7")  # 32x32
-      hidden7 = layer_norm(hidden7, scope="layer_norm8")
+          enc5, lstm_state[6], lstm_size[6], name="state7")  # 32x32
+      hidden7 = tfcl.layer_norm(hidden7, scope="layer_norm8")
 
       # Skip connection.
       hidden7 = tf.concat(axis=3, values=[hidden7, enc0])  # both 32x32
 
-      enc6 = slim.layers.conv2d_transpose(
+      enc6 = tfl.conv2d_transpose(
           hidden7,
           hidden7.get_shape()[3],
-          3,
-          stride=2,
-          scope="convt3",
-          activation_fn=None,
-          normalizer_fn=layer_norm,
-          normalizer_params={"scope": "layer_norm9"})
+          [3, 3],
+          strides=(2, 2),
+          padding="SAME",
+          name="convt3",
+          activation=None)
+      enc6 = tfcl.layer_norm(enc6, scope="layer_norm9")
 
       if self.hparams.model_options == "DNA":
         # Using largest hidden state for predicting untied conv kernels.
-        enc7 = slim.layers.conv2d_transpose(
+        enc7 = tfl.conv2d_transpose(
             enc6,
             self.hparams.dna_kernel_size**2,
-            1,
-            stride=1,
-            scope="convt4",
-            activation_fn=None)
+            [1, 1],
+            strides=(1, 1),
+            padding="SAME",
+            name="convt4",
+            activation=None)
       else:
         # Using largest hidden state for predicting a new image layer.
-        enc7 = slim.layers.conv2d_transpose(
+        enc7 = tfl.conv2d_transpose(
             enc6,
             color_channels,
-            1,
-            stride=1,
-            scope="convt4",
-            activation_fn=None)
+            [1, 1],
+            strides=(1, 1),
+            padding="SAME",
+            name="convt4",
+            activation=None)
         # This allows the network to also generate one image from scratch,
         # which is useful when regions of the image become unoccluded.
         transformed = [tf.nn.sigmoid(enc7)]
 
       if self.hparams.model_options == "CDNA":
         # cdna_input = tf.reshape(hidden5, [int(batch_size), -1])
-        cdna_input = tf.contrib.layers.flatten(hidden5)
+        cdna_input = tfcl.flatten(hidden5)
         transformed += self.cdna_transformation(
             input_image, cdna_input, num_masks, int(color_channels))
       elif self.hparams.model_options == "DNA":
@@ -434,9 +454,9 @@ def construct_predictive_tower(
           raise ValueError("Only one mask is supported for DNA model.")
         transformed = [self.dna_transformation(input_image, enc7)]
 
-      masks = slim.layers.conv2d_transpose(
-          enc6, num_masks + 1, 1,
-          stride=1, scope="convt7", activation_fn=None)
+      masks = tfl.conv2d_transpose(
+          enc6, num_masks + 1, [1, 1], strides=(1, 1),
+          name="convt7", padding="SAME", activation=None)
       masks = tf.reshape(
           tf.nn.softmax(tf.reshape(masks, [-1, num_masks + 1])),
           [batch_size,
@@ -541,12 +561,12 @@ def cdna_transformation(self,
     width = int(prev_image.get_shape()[2])
 
     # Predict kernels using linear function of last hidden layer.
-    cdna_kerns = slim.layers.fully_connected(
+    cdna_kerns = tfl.dense(
         cdna_input,
         self.hparams.dna_kernel_size *
         self.hparams.dna_kernel_size * num_masks,
-        scope="cdna_params",
-        activation_fn=None)
+        name="cdna_params",
+        activation=None)
 
     # Reshape and normalize.
     cdna_kerns = tf.reshape(
@@ -725,7 +745,6 @@ def body(self, features):
     # NOT sure if this is required at all. Doesn"t hurt though! :)
     all_frames = [tf.identity(frame) for frame in all_frames]
 
-    is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
     gen_images, gen_rewards, latent_means, latent_stds = self.construct_model(
         images=all_frames,
         actions=all_actions,
@@ -741,7 +760,7 @@ def body(self, features):
                    lambda: 0.0)
 
     kl_loss = 0.0
-    if is_training:
+    if self.is_training:
       for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
         kl_loss += self.kl_divergence(mean, std)
         tf.summary.histogram("posterior_mean_%d" % i, mean)
@@ -836,7 +855,7 @@ def vgg_layer(self,
                 kernel_size=3,
                 activation=tf.nn.leaky_relu,
                 padding="SAME",
-                scope=""):
+                scope=None):
     """A layer of VGG network with batch norm.
 
     Args:
@@ -845,19 +864,21 @@ def vgg_layer(self,
       kernel_size: size of the kernel
       activation: activation function
       padding: padding of the image
-      scope: slim scope of the op
+      scope: variable scope of the op
     Returns:
       net: output of layer
     """
-    net = slim.conv2d(inputs, nout, kernel_size=kernel_size, padding=padding,
-                      activation_fn=activation, scope=scope+"_conv")
-    net = slim.batch_norm(net, scope=scope+"_bn")
-    net = activation(net)
+    with tf.variable_scope(scope):
+      net = tfl.conv2d(inputs, nout, kernel_size=kernel_size, padding=padding,
+                       activation=None, name="conv")
+      net = tfl.batch_normalization(net,
+                                    training=self.is_training, name="bn")
+      net = activation(net)
     return net
 
-  def basic_lstm(self, inputs, state, num_units, scope=None):
+  def basic_lstm(self, inputs, state, num_units, name=None):
     input_shape = common_layers.shape_list(inputs)
-    cell = tf.contrib.rnn.BasicLSTMCell(num_units, name=scope)
+    cell = tf.contrib.rnn.BasicLSTMCell(num_units, name=name)
     if state is None:
       state = cell.zero_state(input_shape[0], tf.float32)
     outputs, new_state = cell(inputs, state)
@@ -876,19 +897,19 @@ def encoder(self, inputs, nout):
     vgg_layer = self.vgg_layer
     net01 = inputs
     # h1
-    net11 = slim.repeat(net01, 2, vgg_layer, 64, scope="h1")
-    net12 = slim.max_pool2d(net11, [2, 2], scope="h1_pool")
+    net11 = tfcl.repeat(net01, 2, vgg_layer, 64, scope="h1")
+    net12 = tfl.max_pooling2d(net11, [2, 2], strides=(2, 2), name="h1_pool")
     # h2
-    net21 = slim.repeat(net12, 2, vgg_layer, 128, scope="h2")
-    net22 = slim.max_pool2d(net21, [2, 2], scope="h2_pool")
+    net21 = tfcl.repeat(net12, 2, vgg_layer, 128, scope="h2")
+    net22 = tfl.max_pooling2d(net21, [2, 2], strides=(2, 2), name="h2_pool")
     # h3
-    net31 = slim.repeat(net22, 3, vgg_layer, 256, scope="h3")
-    net32 = slim.max_pool2d(net31, [2, 2], scope="h3_pool")
+    net31 = tfcl.repeat(net22, 3, vgg_layer, 256, scope="h3")
+    net32 = tfl.max_pooling2d(net31, [2, 2], strides=(2, 2), name="h3_pool")
     # h4
-    net41 = slim.repeat(net32, 3, vgg_layer, 512, scope="h4")
-    net42 = slim.max_pool2d(net41, [2, 2], scope="h4_pool")
+    net41 = tfcl.repeat(net32, 3, vgg_layer, 512, scope="h4")
+    net42 = tfl.max_pooling2d(net41, [2, 2], strides=(2, 2), name="h4_pool")
     # h5
-    net51 = slim.repeat(net42, 1, vgg_layer, nout, kernel_size=4,
+    net51 = tfcl.repeat(net42, 1, vgg_layer, nout, kernel_size=4,
                         padding="VALID", activation=tf.tanh, scope="h5")
     skips = [net11, net21, net31, net41]
     return net51, skips
@@ -907,31 +928,31 @@ def decoder(self, inputs, skips, nout):
     vgg_layer = self.vgg_layer
     net = inputs
     # d1
-    net = slim.conv2d_transpose(net, 512, kernel_size=4, padding="VALID",
-                                scope="d1_deconv", activation_fn=None)
-    net = slim.batch_norm(net, scope="d1_bn")
+    net = tfl.conv2d_transpose(net, 512, kernel_size=4, padding="VALID",
+                               name="d1_deconv", activation=None)
+    net = tfl.batch_normalization(net, training=self.is_training, name="d1_bn")
     net = tf.nn.leaky_relu(net)
     net = common_layers.upscale(net, 2)
     # d2
     net = tf.concat([net, skips[3]], axis=3)
-    net = slim.repeat(net, 2, vgg_layer, 512, scope="d2a")
-    net = slim.repeat(net, 1, vgg_layer, 256, scope="d2b")
+    net = tfcl.repeat(net, 2, vgg_layer, 512, scope="d2a")
+    net = tfcl.repeat(net, 1, vgg_layer, 256, scope="d2b")
     net = common_layers.upscale(net, 2)
     # d3
     net = tf.concat([net, skips[2]], axis=3)
-    net = slim.repeat(net, 2, vgg_layer, 256, scope="d3a")
-    net = slim.repeat(net, 1, vgg_layer, 128, scope="d3b")
+    net = tfcl.repeat(net, 2, vgg_layer, 256, scope="d3a")
+    net = tfcl.repeat(net, 1, vgg_layer, 128, scope="d3b")
     net = common_layers.upscale(net, 2)
     # d4
     net = tf.concat([net, skips[1]], axis=3)
-    net = slim.repeat(net, 1, vgg_layer, 128, scope="d4a")
-    net = slim.repeat(net, 1, vgg_layer, 64, scope="d4b")
+    net = tfcl.repeat(net, 1, vgg_layer, 128, scope="d4a")
+    net = tfcl.repeat(net, 1, vgg_layer, 64, scope="d4b")
     net = common_layers.upscale(net, 2)
     # d5
     net = tf.concat([net, skips[0]], axis=3)
-    net = slim.repeat(net, 1, vgg_layer, 64, scope="d5")
-    net = slim.conv2d_transpose(net, nout, kernel_size=3, padding="SAME",
-                                scope="d6_deconv", activation_fn=tf.sigmoid)
+    net = tfcl.repeat(net, 1, vgg_layer, 64, scope="d5")
+    net = tfl.conv2d_transpose(net, nout, kernel_size=3, padding="SAME",
+                               name="d6_deconv", activation=tf.sigmoid)
     return net
 
   def stacked_lstm(self, inputs, states, hidden_size, output_size, nlayers):
@@ -948,13 +969,13 @@ def stacked_lstm(self, inputs, states, hidden_size, output_size, nlayers):
       skips: a list of updated lstm states for each layer
     """
     net = inputs
-    net = slim.layers.fully_connected(
-        net, hidden_size, activation_fn=None, scope="af1")
+    net = tfl.dense(
+        net, hidden_size, activation=None, name="af1")
     for i in range(nlayers):
       net, states[i] = self.basic_lstm(
-          net, states[i], hidden_size, scope="alstm%d"%i)
-    net = slim.layers.fully_connected(
-        net, output_size, activation_fn=tf.tanh, scope="af2")
+          net, states[i], hidden_size, name="alstm%d"%i)
+    net = tfl.dense(
+        net, output_size, activation=tf.nn.tanh, name="af2")
     return net, states
 
   def lstm_gaussian(self, inputs, states, hidden_size, output_size, nlayers):
@@ -972,15 +993,12 @@ def lstm_gaussian(self, inputs, states, hidden_size, output_size, nlayers):
       skips: a list of updated lstm states for each layer
     """
     net = inputs
-    net = slim.layers.fully_connected(net, hidden_size,
-                                      activation_fn=None, scope="bf1")
+    net = tfl.dense(net, hidden_size, activation=None, name="bf1")
     for i in range(nlayers):
       net, states[i] = self.basic_lstm(
-          net, states[i], hidden_size, scope="blstm%d"%i)
-    mu = slim.layers.fully_connected(
-        net, output_size, activation_fn=None, scope="bf2mu")
-    logvar = slim.layers.fully_connected(
-        net, output_size, activation_fn=None, scope="bf2log")
+          net, states[i], hidden_size, name="blstm%d"%i)
+    mu = tfl.dense(net, output_size, activation=None, name="bf2mu")
+    logvar = tfl.dense(net, output_size, activation=None, name="bf2log")
     return mu, logvar, states
 
   def construct_model(self, images, actions, rewards):
@@ -1030,7 +1048,7 @@ def construct_model(self, images, actions, rewards):
     for i, image in enumerate(images):
       with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
         enc, skips = self.encoder(image, rnn_size)
-        enc = tf.contrib.layers.flatten(enc)
+        enc = tfcl.flatten(enc)
         enc_images.append(enc)
         enc_skips.append(skips)
 

From 9b3f1d7c1fd456ac24d7b40fd35dd08cc1cb4c81 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Tue, 17 Jul 2018 02:37:15 -0700
Subject: [PATCH 0340/2720] Add compression using self attention baseline

PiperOrigin-RevId: 204881406
---
 tensor2tensor/layers/latent_layers.py | 39 ++++++++++++++++++---------
 1 file changed, 26 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 24cc150ef..00c16a358 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -30,17 +30,13 @@
 DO_SUMMARIES = True
 
 
-def attend(x, source, hparams, name):
+def compress_self_attention_layer(x, hparams, name):
   """Attend function."""
   with tf.variable_scope(name):
-    # x = tf.squeeze(x, axis=2)
     x, xshape, _ = cia.maybe_reshape_4d_to_3d(x)
-    if len(source.get_shape()) > 3:
-      source = tf.squeeze(source, axis=2)
-    source = common_attention.add_timing_signal_1d(source)
     y = common_attention.multihead_attention(
         common_layers.layer_preprocess(x, hparams),
-        source,
+        None,
         None,
         hparams.attention_key_channels or hparams.hidden_size,
         hparams.attention_value_channels or hparams.hidden_size,
@@ -70,19 +66,19 @@ def multinomial_sample(x, vocab_size, sampling_method, temperature):
   return reshaped_samples
 
 
-def ae_latent_softmax(latents_pred, latents_discrete_hot, hparams):
+def ae_latent_softmax(latents_pred, latents_discrete_hot, vocab_size, hparams):
   """Latent prediction and loss.
 
   Args:
     latents_pred: Tensor of shape [..., depth].
     latents_discrete_hot: Tensor of shape [..., vocab_size].
+    vocab_size: an int representing the vocab size.
     hparams: tf.contrib.training.HParams.
 
   Returns:
     sample: Tensor of shape [...], a sample from a multinomial distribution.
     loss: Tensor of shape [...], the softmax cross-entropy.
   """
-  vocab_size = 2**hparams.bottleneck_bits
   with tf.variable_scope("latent_logits"):
     latents_logits = tf.layers.dense(latents_pred, vocab_size,
                                      name="logits_dense")
@@ -91,6 +87,9 @@ def ae_latent_softmax(latents_pred, latents_discrete_hot, hparams):
                                  tf.reduce_mean(tf.square(latents_logits)))
     loss = tf.nn.softmax_cross_entropy_with_logits_v2(
         labels=latents_discrete_hot, logits=latents_logits)
+
+    # TODO(trandustin): tease this out from ae_latent_softmax.
+    # we use just the loss portion to anchor prior / encoder on text.
     sample = multinomial_sample(latents_logits,
                                 vocab_size,
                                 hparams.sampling_method,
@@ -216,6 +215,10 @@ def compress_encoder(inputs,
             padding="SAME",
             name="compress_conv_%d" % i)
         y = tf.nn.dropout(y, 1.0 - hparams.dropout)
+        if hparams.do_compress_attend:
+          y = compress_self_attention_layer(
+              x, hparams, name="compress_selfatt_%d" % i)
+          y += x
         x = y
 
     x = residual_block_layer(x, hparams)
@@ -293,6 +296,10 @@ def decompress_decoder(inputs,
     for i in range(hparams.num_compress_steps // 2):
       j = hparams.num_compress_steps // 2 - i - 1
       with tf.variable_scope(name + "_%d" % j):
+        if hparams.do_decompress_attend:
+          y = compress_self_attention_layer(
+              x, hparams, name="decompress_selfatt")
+          x += y
         y = tf.layers.conv2d_transpose(
             x,
             hparams.hidden_size,
@@ -464,7 +471,7 @@ def bottleneck_layer(targets_c, hparams):
   latents_discrete = tf.argmax(latents_discrete_hot, axis=-1)
 
   if DO_SUMMARIES:
-    tf.summary.histogram("b0", tf.reshape(latents_discrete, [-1]))
+    tf.summary.histogram("discrete_latents", tf.reshape(latents_discrete, [-1]))
   return latents_dense, latents_discrete_hot, extra_loss
 
 
@@ -473,6 +480,7 @@ def latent_prediction_model(inputs,
                             latents_discrete,
                             latents_dense,
                             hparams,
+                            vocab_size=None,
                             name="latent_prediction"):
   """Transformer-based latent prediction model.
 
@@ -487,6 +495,7 @@ def latent_prediction_model(inputs,
       One-hot latents to compute log-probability of given inputs.
     latents_dense: Tensor of shape [batch, length_q, hparams.hidden_size].
     hparams: tf.contrib.training.HParams.
+    vocab_size: int, if given else None.
     name: string, variable scope.
 
   Returns:
@@ -495,11 +504,15 @@ def latent_prediction_model(inputs,
   """
   with tf.variable_scope(name):
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
-      latents_pred = transformer_latent_decoder(
-          tf.stop_gradient(latents_dense), inputs, ed_attention_bias,
-          hparams, name)
+      latents_pred = transformer_latent_decoder(tf.stop_gradient(latents_dense),
+                                                inputs,
+                                                ed_attention_bias,
+                                                hparams,
+                                                name)
+      vocab_size = (2**hparams.bottleneck_bits
+                    if vocab_size is None else vocab_size)
       _, latent_pred_loss = ae_latent_softmax(
-          latents_pred, tf.stop_gradient(latents_discrete), hparams)
+          latents_pred, tf.stop_gradient(latents_discrete), vocab_size, hparams)
   return latents_pred, latent_pred_loss
 
 

From aceb40f3356230b370f92c8124cc3d6f3e8fd7c9 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 17 Jul 2018 12:29:59 -0700
Subject: [PATCH 0341/2720] Fixed dataset function to return deterministic
 order when shuffle_files=False

PiperOrigin-RevId: 204955345
---
 tensor2tensor/data_generators/algorithmic.py  | 35 +++++++-
 tensor2tensor/data_generators/problem.py      | 20 +++--
 tensor2tensor/data_generators/problem_test.py | 90 +++++++++++++++++++
 tensor2tensor/data_generators/video_utils.py  |  3 +-
 tensor2tensor/utils/trainer_lib_test.py       | 43 +++------
 5 files changed, 150 insertions(+), 41 deletions(-)
 create mode 100644 tensor2tensor/data_generators/problem_test.py

diff --git a/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index 50fd887f7..a40378263 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -16,15 +16,17 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-import numpy as np
 
+import os
+import shutil
+import numpy as np
 from six.moves import range  # pylint: disable=redefined-builtin
-
 from tensor2tensor.data_generators import generator_utils as utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
+import tensorflow as tf
 
 
 class AlgorithmicProblem(problem.Problem):
@@ -508,3 +510,32 @@ def generator(self, nbr_symbols, max_length, nbr_cases):
   def eval_metrics(self):
     defaults = super(AlgorithmicSortProblem, self).eval_metrics()
     return defaults + [metrics.Metrics.EDIT_DISTANCE]
+
+
+@registry.register_problem
+class TinyAlgo(AlgorithmicIdentityBinary40):
+  """A small algorthmic problem for testing."""
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    """Ganerate data for this problem."""
+
+    del tmp_dir, task_id
+    identity_problem = AlgorithmicIdentityBinary40()
+    utils.generate_files(
+        identity_problem.generator(self.num_symbols, 40, 100000),
+        self.training_filepaths(data_dir, 1, shuffled=True), 100)
+    utils.generate_files(
+        identity_problem.generator(self.num_symbols, 400, 10000),
+        self.dev_filepaths(data_dir, 1, shuffled=True), 100)
+
+  @classmethod
+  def setUpForTest(cls):
+    """Setup directories and files required to run the problem."""
+
+    tmp_dir = tf.test.get_temp_dir()
+    shutil.rmtree(tmp_dir)
+    os.mkdir(tmp_dir)
+    cls.data_dir = tmp_dir
+
+    # Generate a small test dataset
+    cls().generate_data(TinyAlgo.data_dir, None)
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 7de6c4ff1..2850fbc4d 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -345,7 +345,7 @@ def eval_metrics(self):
   # END SUBCLASS INTERFACE
   # ============================================================================
 
-  def preprocess(self, dataset, mode, hparams):
+  def preprocess(self, dataset, mode, hparams, interleave=True):
     """Runtime preprocessing on the whole dataset.
 
     Return a tf.data.Datset -- the preprocessed version of the given one.
@@ -355,6 +355,9 @@ def preprocess(self, dataset, mode, hparams):
       dataset: the Dataset of already decoded but not yet preprocessed features.
       mode: tf.estimator.ModeKeys
       hparams: HParams, model hyperparameters
+      interleave: bool, whether to use parallel_interleave, which is faster
+        but will alter the order of samples non-deterministically, or flat_map,
+        which is slower but will preserve the sample order.
 
     Returns:
       a Dataset
@@ -365,10 +368,12 @@ def _preprocess(example):
         examples = tf.data.Dataset.from_tensors(examples)
       return examples
 
-    is_training = mode == tf.estimator.ModeKeys.TRAIN
-    dataset = dataset.apply(
-        tf.contrib.data.parallel_interleave(
-            _preprocess, sloppy=is_training, cycle_length=8))
+    if interleave:
+      dataset = dataset.apply(
+          tf.contrib.data.parallel_interleave(
+              _preprocess, sloppy=True, cycle_length=8))
+    else:
+      dataset = dataset.flat_map(_preprocess)
 
     return dataset
 
@@ -568,6 +573,7 @@ def dataset(self,
     # Functions used in dataset transforms below. `filenames` can be either a
     # `tf.string` tensor or `tf.data.Dataset` containing one or more filenames.
     def _load_records_and_preprocess(filenames):
+      """Reads files from a string tensor or a dataset of filenames."""
       # Load records from file(s) with an 8MiB read buffer.
       dataset = tf.data.TFRecordDataset(filenames, buffer_size=8 * 1024 * 1024)
       # Decode.
@@ -575,7 +581,8 @@ def _load_records_and_preprocess(filenames):
       # Preprocess if requested.
       # Note that preprocessing should happen per-file as order may matter.
       if preprocess:
-        dataset = self.preprocess(dataset, mode, hparams)
+        dataset = self.preprocess(dataset, mode, hparams,
+                                  interleave=shuffle_files)
       return dataset
 
     if len(data_files) < num_partitions:
@@ -1140,4 +1147,3 @@ def skip_random_fraction(dataset, data_file):
   # replicas reading the same data in lock-step.
   num_skip = random.randint(0, _file_num_records_cached(data_file))
   return dataset.skip(num_skip)
-
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
new file mode 100644
index 000000000..f9c7c7f45
--- /dev/null
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -0,0 +1,90 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Test for common problem functionalities."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.data_generators import algorithmic
+import tensorflow as tf
+
+
+def assert_tensors_equal(sess, t1, t2, n):
+  """Compute tensors `n` times and ensure that they are equal."""
+
+  for _ in range(n):
+
+    v1, v2 = sess.run([t1, t2])
+
+    if v1.shape != v2.shape:
+      return False
+
+    if not np.all(v1 == v2):
+      return False
+
+  return True
+
+
+class ProblemTest(tf.test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    algorithmic.TinyAlgo.setUpForTest()
+
+  def testNoShuffleDeterministic(self):
+    problem = algorithmic.TinyAlgo()
+    dataset = problem.dataset(mode=tf.estimator.ModeKeys.TRAIN,
+                              data_dir=algorithmic.TinyAlgo.data_dir,
+                              shuffle_files=False)
+
+    tensor1 = dataset.make_one_shot_iterator().get_next()["targets"]
+    tensor2 = dataset.make_one_shot_iterator().get_next()["targets"]
+
+    with tf.Session() as sess:
+      self.assertTrue(assert_tensors_equal(sess, tensor1, tensor2, 20))
+
+  def testNoShufflePreprocess(self):
+
+    problem = algorithmic.TinyAlgo()
+    dataset1 = problem.dataset(mode=tf.estimator.ModeKeys.TRAIN,
+                               data_dir=algorithmic.TinyAlgo.data_dir,
+                               shuffle_files=False, preprocess=False)
+    dataset2 = problem.dataset(mode=tf.estimator.ModeKeys.TRAIN,
+                               data_dir=algorithmic.TinyAlgo.data_dir,
+                               shuffle_files=False, preprocess=True)
+
+    tensor1 = dataset1.make_one_shot_iterator().get_next()["targets"]
+    tensor2 = dataset2.make_one_shot_iterator().get_next()["targets"]
+
+    with tf.Session() as sess:
+      self.assertTrue(assert_tensors_equal(sess, tensor1, tensor2, 20))
+
+  def testNoShuffleFail(self):
+    problem = algorithmic.TinyAlgo()
+    dataset = problem.dataset(mode=tf.estimator.ModeKeys.TRAIN,
+                              data_dir=algorithmic.TinyAlgo.data_dir,
+                              shuffle_files=True)
+
+    tensor1 = dataset.make_one_shot_iterator().get_next()["targets"]
+    tensor2 = dataset.make_one_shot_iterator().get_next()["targets"]
+
+    with tf.Session() as sess:
+      self.assertFalse(assert_tensors_equal(sess, tensor1, tensor2, 20))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index c2f4abcf8..eef40b56f 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -185,7 +185,8 @@ def example_reading_spec(self):
 
     return data_fields, data_items_to_decoders
 
-  def preprocess(self, dataset, mode, hparams):
+  def preprocess(self, dataset, mode, hparams, interleave=True):
+    del interleave
     def split_on_batch(x):
       """Split x on batch dimension into x[:size, ...] and x[size:, ...]."""
       length = len(x.get_shape())
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index 799168d44..a0c5e98fa 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -18,55 +18,32 @@
 from __future__ import division
 from __future__ import print_function
 
-import os
-import shutil
 from tensor2tensor import models  # pylint: disable=unused-import
 from tensor2tensor.data_generators import algorithmic
-from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem as problem_lib
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
-
 import tensorflow as tf
 
 
-@registry.register_problem
-class TinyAlgo(algorithmic.AlgorithmicIdentityBinary40):
-
-  def generate_data(self, data_dir, tmp_dir, task_id=-1):
-    del tmp_dir, task_id
-    identity_problem = algorithmic.AlgorithmicIdentityBinary40()
-    generator_utils.generate_files(
-        identity_problem.generator(self.num_symbols, 40, 100000),
-        self.training_filepaths(data_dir, 1, shuffled=True), 100)
-    generator_utils.generate_files(
-        identity_problem.generator(self.num_symbols, 400, 10000),
-        self.dev_filepaths(data_dir, 1, shuffled=True), 100)
-
-
 class TrainerLibTest(tf.test.TestCase):
 
   @classmethod
   def setUpClass(cls):
-    tmp_dir = tf.test.get_temp_dir()
-    shutil.rmtree(tmp_dir)
-    os.mkdir(tmp_dir)
-    cls.data_dir = tmp_dir
-
-    # Generate a small test dataset
-    registry.problem("tiny_algo").generate_data(cls.data_dir, None)
+    algorithmic.TinyAlgo.setUpForTest()
 
   def testExperiment(self):
     exp_fn = trainer_lib.create_experiment_fn(
         "transformer",
         "tiny_algo",
-        self.data_dir,
+        algorithmic.TinyAlgo.data_dir,
         train_steps=1,
         eval_steps=1,
         min_eval_frequency=1,
         use_tpu=False)
     run_config = trainer_lib.create_run_config(
-        model_dir=self.data_dir, num_gpus=0, use_tpu=False)
+        model_dir=algorithmic.TinyAlgo.data_dir, num_gpus=0,
+        use_tpu=False)
     hparams = registry.hparams("transformer_tiny_tpu")
     exp = exp_fn(run_config, hparams)
     exp.test()
@@ -74,11 +51,13 @@ def testExperiment(self):
   def testModel(self):
     # HParams
     hparams = trainer_lib.create_hparams(
-        "transformer_tiny", data_dir=self.data_dir, problem_name="tiny_algo")
+        "transformer_tiny", data_dir=algorithmic.TinyAlgo.data_dir,
+        problem_name="tiny_algo")
 
     # Dataset
     problem = hparams.problem
-    dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, self.data_dir)
+    dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN,
+                              algorithmic.TinyAlgo.data_dir)
     dataset = dataset.repeat(None).padded_batch(10, dataset.output_shapes)
     features = dataset.make_one_shot_iterator().get_next()
     features = problem_lib.standardize_shapes(features)
@@ -101,7 +80,8 @@ def testModel(self):
   def testMultipleTargetModalities(self):
     # HParams
     hparams = trainer_lib.create_hparams(
-        "transformer_tiny", data_dir=self.data_dir, problem_name="tiny_algo")
+        "transformer_tiny", data_dir=algorithmic.TinyAlgo.data_dir,
+        problem_name="tiny_algo")
     tm = hparams.problem.get_hparams().target_modality
     hparams.problem.get_hparams().target_modality = {
         "targets": tm,
@@ -111,7 +91,8 @@ def testMultipleTargetModalities(self):
 
     # Dataset
     problem = hparams.problem
-    dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, self.data_dir)
+    dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN,
+                              algorithmic.TinyAlgo.data_dir)
     dataset = dataset.repeat(None).padded_batch(10, dataset.output_shapes)
     features = dataset.make_one_shot_iterator().get_next()
     features = problem_lib.standardize_shapes(features)

From 79ed80acfab7bc1932e667866d7c09b6e023737d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 17 Jul 2018 12:38:07 -0700
Subject: [PATCH 0342/2720] Update Similarity Transformer model body

PiperOrigin-RevId: 204956830
---
 .../data_generators/function_docstring.py     | 69 ++++++++++++++
 .../gh_function_docstring_encoder.py          | 51 -----------
 .../models/research/similarity_transformer.py | 89 +++++++++++++++++++
 3 files changed, 158 insertions(+), 51 deletions(-)
 create mode 100644 tensor2tensor/data_generators/function_docstring.py
 delete mode 100644 tensor2tensor/data_generators/gh_function_docstring_encoder.py
 create mode 100644 tensor2tensor/models/research/similarity_transformer.py

diff --git a/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
new file mode 100644
index 000000000..6f8cb1f2d
--- /dev/null
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -0,0 +1,69 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Github function/text similatrity problems."""
+import csv
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import translate
+from tensor2tensor.utils import metrics
+from tensor2tensor.utils import registry
+
+
+# There are 10 splits of the data as CSV files.
+_DATA_BASE_URL = 'https://storage.googleapis.com/kubeflow-examples/t2t-code-search/data'
+_GITHUB_FUNCTION_DOCSTRING_FILES = [
+    [
+        '{}/pairs-0000{}-of-00010.csv'.format(_DATA_BASE_URL, i),
+        'pairs-0000{}-of-00010.csv'.format(i),
+    ]
+    for i in range(10)
+]
+
+
+@registry.register_problem
+class GithubFunctionDocstring(translate.TranslateProblem):
+  """This class defines the problem of finding similarity between Python
+  function and docstring"""
+
+  @property
+  def is_generate_per_split(self):
+    return False
+
+  @property
+  def approx_vocab_size(self):
+    return 2**13
+
+  def source_data_files(self, dataset_split):  # pylint: disable=no-self-use,unused-argument
+    # TODO(sanyamkapoor): separate train/eval data set.
+    return _GITHUB_FUNCTION_DOCSTRING_FILES
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):  # pylint: disable=no-self-use,unused-argument
+    """Returns a generator to return {"inputs": [text], "targets": [text]}."""
+
+    pair_csv_files = [
+        generator_utils.maybe_download(data_dir, filename, uri)
+        for uri, filename in self.source_data_files(dataset_split)
+    ]
+
+    for pairs_file in pair_csv_files:
+      with open(pairs_file, 'r') as csv_file:
+        pairs_reader = csv.reader(csv_file)
+        for row in pairs_reader:
+          function_tokens, docstring_tokens = row[-2:]
+          yield {'inputs': docstring_tokens, 'targets': function_tokens}
+
+  def eval_metrics(self):  # pylint: disable=no-self-use
+    return [
+        metrics.Metrics.ACC
+    ]
diff --git a/tensor2tensor/data_generators/gh_function_docstring_encoder.py b/tensor2tensor/data_generators/gh_function_docstring_encoder.py
deleted file mode 100644
index 784c237f3..000000000
--- a/tensor2tensor/data_generators/gh_function_docstring_encoder.py
+++ /dev/null
@@ -1,51 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-"""Github function to text similatrity problems."""
-
-import os
-
-from tensor2tensor.data_generators import text_problems
-from tensor2tensor.utils import registry
-from tensor2tensor.utils import t2t_model
-
-
-@registry.register_model
-class SimilarityTransformer(t2t_model.T2TModel):
-  """Similarity scores between functions and docstrings."""
-
-  def body(self, features):
-    # TODO(sanyamkapoor): need to fill this with Transformer encoder/decoder
-    # and loss calculation
-    raise NotImplementedError
-
-
-@registry.register_problem
-class GithubFunctionDocstring(text_problems.Text2TextProblem):
-  """The problem of similarity between Python function and docstring."""
-
-  @property
-  def is_generate_per_split(self):
-    return False
-
-  def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    """Returns the generator of {"inputs": [text], "targets": [text]} dict."""
-
-    functions_file_path = os.path.join(
-        data_dir, '{}.function'.format(dataset_split))
-    docstrings_file_path = os.path.join(
-        data_dir, '{}.docstring'.format(dataset_split))
-
-    return text_problems.text2text_txt_iterator(
-        functions_file_path, docstrings_file_path)
diff --git a/tensor2tensor/models/research/similarity_transformer.py b/tensor2tensor/models/research/similarity_transformer.py
new file mode 100644
index 000000000..0f477258f
--- /dev/null
+++ b/tensor2tensor/models/research/similarity_transformer.py
@@ -0,0 +1,89 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Using Transformer Networks for String similarities."""
+from tensor2tensor.data_generators import problem
+from tensor2tensor.layers import common_layers
+from tensor2tensor.models import transformer
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
+import tensorflow as tf
+
+
+@registry.register_model
+class SimilarityTransformer(t2t_model.T2TModel):
+  """
+  This class defines the model to compute similarity scores between functions
+  and docstrings
+  """
+
+  def top(self, body_output, features):  # pylint: disable=no-self-use,unused-argument
+    return body_output
+
+  def body(self, features):
+    """Body of the Similarity Transformer Network."""
+
+    with tf.variable_scope('string_embedding'):
+      string_embedding = self.encode(features, 'inputs')
+
+    if 'targets' in features:
+      with tf.variable_scope('code_embedding'):
+        code_embedding = self.encode(features, 'targets')
+
+      string_embedding_norm = tf.nn.l2_normalize(string_embedding, axis=1)
+      code_embedding_norm = tf.nn.l2_normalize(code_embedding, axis=1)
+
+      # All-vs-All cosine distance matrix, reshaped as row-major.
+      cosine_dist = 1.0 - tf.matmul(string_embedding_norm, code_embedding_norm,
+                                    transpose_b=True)
+      cosine_dist_flat = tf.reshape(cosine_dist, [-1, 1])
+
+      # Positive samples on the diagonal, reshaped as row-major.
+      label_matrix = tf.eye(tf.shape(cosine_dist)[0], dtype=tf.int32)
+      label_matrix_flat = tf.reshape(label_matrix, [-1])
+
+      logits = tf.concat([1.0 - cosine_dist_flat, cosine_dist_flat], axis=1)
+      labels = tf.one_hot(label_matrix_flat, 2)
+
+      loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels,
+                                                     logits=logits)
+
+      return string_embedding, {'training': loss}
+
+    return string_embedding
+
+  def encode(self, features, input_key):
+    hparams = self._hparams
+    inputs = common_layers.flatten4d3d(features[input_key])
+
+    (encoder_input, encoder_self_attention_bias, _) = (
+        transformer.transformer_prepare_encoder(inputs, problem.SpaceID.EN_TOK,
+                                                self._hparams))
+
+    encoder_input = tf.nn.dropout(encoder_input,
+                                  1.0 - hparams.layer_prepostprocess_dropout)
+    encoder_output = transformer.transformer_encoder(
+        encoder_input,
+        encoder_self_attention_bias,
+        self._hparams,
+        nonpadding=transformer.features_to_nonpadding(features, input_key))
+    encoder_output = tf.expand_dims(encoder_output, 2)
+
+    encoder_output = tf.reduce_mean(tf.squeeze(encoder_output, axis=2), axis=1)
+
+    return encoder_output
+
+  def infer(self, features=None, **kwargs):  # pylint: disable=no-self-use,unused-argument
+    predictions, _ = self(features)
+    return predictions

From 1d29da2d1ee8397b5ace402189b546abc34790a7 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 17 Jul 2018 17:25:14 -0700
Subject: [PATCH 0343/2720] Corrections to data generation for github and style
 in the similarity transformer.

PiperOrigin-RevId: 205003792
---
 tensor2tensor/data_generators/algorithmic.py  |  2 +-
 tensor2tensor/data_generators/all_problems.py |  1 +
 .../data_generators/function_docstring.py     | 21 ++++++++++++-------
 tensor2tensor/data_generators/problem_test.py |  2 +-
 .../models/research/similarity_transformer.py | 16 +++++++-------
 tensor2tensor/utils/trainer_lib_test.py       |  2 +-
 6 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index a40378263..aea9a5e90 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -529,7 +529,7 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
         self.dev_filepaths(data_dir, 1, shuffled=True), 100)
 
   @classmethod
-  def setUpForTest(cls):
+  def setup_for_test(cls):
     """Setup directories and files required to run the problem."""
 
     tmp_dir = tf.test.get_temp_dir()
diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 62b448655..419c02d48 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -34,6 +34,7 @@
     "tensor2tensor.data_generators.common_voice",
     "tensor2tensor.data_generators.desc2code",
     "tensor2tensor.data_generators.fsns",
+    "tensor2tensor.data_generators.function_docstring",
     "tensor2tensor.data_generators.gene_expression",
     "tensor2tensor.data_generators.google_robot_pushing",
     "tensor2tensor.data_generators.gym_problems_specs",
diff --git a/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
index 6f8cb1f2d..3058cccce 100644
--- a/tensor2tensor/data_generators/function_docstring.py
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -21,11 +21,11 @@
 
 
 # There are 10 splits of the data as CSV files.
-_DATA_BASE_URL = 'https://storage.googleapis.com/kubeflow-examples/t2t-code-search/data'
+_DATA_BASE_URL = "https://storage.googleapis.com/kubeflow-examples/t2t-code-search/data"
 _GITHUB_FUNCTION_DOCSTRING_FILES = [
     [
-        '{}/pairs-0000{}-of-00010.csv'.format(_DATA_BASE_URL, i),
-        'pairs-0000{}-of-00010.csv'.format(i),
+        "{}/pairs-0000{}-of-00010.csv".format(_DATA_BASE_URL, i),
+        "pairs-0000{}-of-00010.csv".format(i),
     ]
     for i in range(10)
 ]
@@ -44,11 +44,11 @@ def is_generate_per_split(self):
   def approx_vocab_size(self):
     return 2**13
 
-  def source_data_files(self, dataset_split):  # pylint: disable=no-self-use,unused-argument
+  def source_data_files(self, dataset_split):
     # TODO(sanyamkapoor): separate train/eval data set.
     return _GITHUB_FUNCTION_DOCSTRING_FILES
 
-  def generate_samples(self, data_dir, tmp_dir, dataset_split):  # pylint: disable=no-self-use,unused-argument
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
     """Returns a generator to return {"inputs": [text], "targets": [text]}."""
 
     pair_csv_files = [
@@ -57,13 +57,18 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):  # pylint: disable
     ]
 
     for pairs_file in pair_csv_files:
-      with open(pairs_file, 'r') as csv_file:
+      with open(pairs_file, "r") as csv_file:
         pairs_reader = csv.reader(csv_file)
         for row in pairs_reader:
           function_tokens, docstring_tokens = row[-2:]
-          yield {'inputs': docstring_tokens, 'targets': function_tokens}
+          yield {"inputs": docstring_tokens, "targets": function_tokens}
 
-  def eval_metrics(self):  # pylint: disable=no-self-use
+  def generate_text_for_vocab(self, data_dir, tmp_dir):
+    for sample in self.generate_samples(data_dir, tmp_dir, None):
+      yield sample["inputs"]
+      yield sample["targets"]
+
+  def eval_metrics(self):
     return [
         metrics.Metrics.ACC
     ]
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index f9c7c7f45..5e76db9e8 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -43,7 +43,7 @@ class ProblemTest(tf.test.TestCase):
 
   @classmethod
   def setUpClass(cls):
-    algorithmic.TinyAlgo.setUpForTest()
+    algorithmic.TinyAlgo.setup_for_test()
 
   def testNoShuffleDeterministic(self):
     problem = algorithmic.TinyAlgo()
diff --git a/tensor2tensor/models/research/similarity_transformer.py b/tensor2tensor/models/research/similarity_transformer.py
index 0f477258f..ed915c3f4 100644
--- a/tensor2tensor/models/research/similarity_transformer.py
+++ b/tensor2tensor/models/research/similarity_transformer.py
@@ -28,18 +28,18 @@ class SimilarityTransformer(t2t_model.T2TModel):
   and docstrings
   """
 
-  def top(self, body_output, features):  # pylint: disable=no-self-use,unused-argument
+  def top(self, body_output, features):
     return body_output
 
   def body(self, features):
     """Body of the Similarity Transformer Network."""
 
-    with tf.variable_scope('string_embedding'):
-      string_embedding = self.encode(features, 'inputs')
+    with tf.variable_scope("string_embedding"):
+      string_embedding = self.encode(features, "inputs")
 
-    if 'targets' in features:
-      with tf.variable_scope('code_embedding'):
-        code_embedding = self.encode(features, 'targets')
+    if "targets" in features:
+      with tf.variable_scope("code_embedding"):
+        code_embedding = self.encode(features, "targets")
 
       string_embedding_norm = tf.nn.l2_normalize(string_embedding, axis=1)
       code_embedding_norm = tf.nn.l2_normalize(code_embedding, axis=1)
@@ -59,7 +59,7 @@ def body(self, features):
       loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels,
                                                      logits=logits)
 
-      return string_embedding, {'training': loss}
+      return string_embedding, {"training": loss}
 
     return string_embedding
 
@@ -84,6 +84,6 @@ def encode(self, features, input_key):
 
     return encoder_output
 
-  def infer(self, features=None, **kwargs):  # pylint: disable=no-self-use,unused-argument
+  def infer(self, features=None, **kwargs):
     predictions, _ = self(features)
     return predictions
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index a0c5e98fa..aa9d57976 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -30,7 +30,7 @@ class TrainerLibTest(tf.test.TestCase):
 
   @classmethod
   def setUpClass(cls):
-    algorithmic.TinyAlgo.setUpForTest()
+    algorithmic.TinyAlgo.setup_for_test()
 
   def testExperiment(self):
     exp_fn = trainer_lib.create_experiment_fn(

From 47a099472f718c0c3c3a2872aa0d33659fe31787 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 18 Jul 2018 09:49:24 -0700
Subject: [PATCH 0344/2720] Fixing the formula for RMSE metric to incorporate
 mean of the squared

PiperOrigin-RevId: 205092596
---
 tensor2tensor/utils/metrics.py      |  2 +-
 tensor2tensor/utils/metrics_test.py | 12 ++++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 41b2b4c4b..683dc5e48 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -75,7 +75,7 @@ def padded_rmse(predictions, labels, weights_fn=common_layers.weights_all):
   predictions, labels = common_layers.pad_with_zeros(predictions, labels)
   weights = weights_fn(labels)
   error = tf.pow(predictions - labels, 2)
-  error_sqrt = tf.sqrt(tf.reduce_sum(error * weights))
+  error_sqrt = tf.sqrt(tf.reduce_mean(error * weights))
   return error_sqrt, tf.reduce_sum(weights)
 
 
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 9ad3bd2a9..3e504d881 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -69,6 +69,18 @@ def testSequenceAccuracyMetric(self):
       actual = session.run(a)
     self.assertEqual(actual, expected)
 
+  def testRMSEMetric(self):
+    predictions = np.full((10, 1), 1)  # All 1's
+    targets = np.full((10, 1), 3)  # All 3's
+    expected = np.sqrt(np.mean((predictions - targets)**2))  # RMSE = 2.0
+    with self.test_session() as session:
+      rmse, _ = metrics.padded_rmse(
+          tf.constant(predictions, dtype=tf.int32),
+          tf.constant(targets, dtype=tf.int32))
+      session.run(tf.global_variables_initializer())
+      actual = session.run(rmse)
+    self.assertEqual(actual, expected)
+
   def testSequenceEditDistanceMetric(self):
     predictions = np.array([[3, 4, 5, 1, 0, 0],
                             [2, 1, 3, 4, 0, 0],

From 6b1438d95f8cb6b4cc9aeb872c0a22ff2b3d0bcd Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 18 Jul 2018 09:52:08 -0700
Subject: [PATCH 0345/2720] no public changes

PiperOrigin-RevId: 205092926
---
 tensor2tensor/bin/t2t_trainer.py         |  3 ++-
 tensor2tensor/data_generators/problem.py | 12 ++++++------
 tensor2tensor/utils/trainer_lib.py       | 16 +++++++++-------
 3 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 0958dba3b..8d57476c0 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -30,8 +30,9 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
-
 import tensorflow as tf
+from tensorflow.contrib.tpu.python.tpu import tpu_config
+
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 2850fbc4d..ca7b8f9c5 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -27,6 +27,7 @@
 from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import metrics
 import tensorflow as tf
+from tensorflow.contrib.tpu.python.tpu import tpu_config
 
 
 
@@ -738,15 +739,14 @@ def _dataset_partition(self, mode, config):
       # Reset in the case when using TPU but alternating TRAIN and EVAL.
       self._next_partition_id = 0
       return 0, 1
+    phift = config.tpu_config.per_host_input_for_training
     # BEGIN GOOGLE-INTERNAL
-    # make mesh-tensorflow on TPU work with patch CL/202825176
-    # TODO(ylc): fix this once TPU estimator changes are checked in.
-    if getattr(config.tpu_config, "symmetric_sharding_enabled", False):
-      tf.logging.info("symmetric_sharding_enabled")
-      self._next_partition_id = 0
+    # This is the mesh-tensorflow case.  Still requires patch of cl/204685944
+    if (hasattr(tpu_config.InputPipelineConfig, "BROADCAST") and
+        phift == tpu_config.InputPipelineConfig.BROADCAST):
       return 0, 1
     # END GOOOGLE-INTERNAL
-    if config.tpu_config.per_host_input_for_training:
+    if phift:
       num_partitions = max(config.tpu_config.num_shards // 8, 1)
     else:
       num_partitions = config.tpu_config.num_shards
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 30b00df68..09439dd40 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -155,15 +155,17 @@ def create_run_config(master="",
 
   # If using TPU, use TPU RunConfig, add TPUConfig, and add additional args
   if use_tpu:
-    if tpu_config_extra_kwargs is None:
-      tpu_config_extra_kwargs = {}
+    tpu_config_kwargs = {
+        "iterations_per_loop": iterations_per_loop,
+        "num_shards": num_shards,
+        "per_host_input_for_training": True,
+        "initial_infeed_sleep_secs": tpu_infeed_sleep_secs,
+    }
+    if tpu_config_extra_kwargs is not None:
+      tpu_config_kwargs.update(tpu_config_extra_kwargs)
     run_config_cls = tf.contrib.tpu.RunConfig
     tpu_config = tf.contrib.tpu.TPUConfig(
-        iterations_per_loop=iterations_per_loop,
-        num_shards=num_shards,
-        per_host_input_for_training=True,
-        initial_infeed_sleep_secs=tpu_infeed_sleep_secs,
-        **tpu_config_extra_kwargs)
+        **tpu_config_kwargs)
     run_config_args["tpu_config"] = tpu_config
 
   config = run_config_cls(**run_config_args)

From c647fc0bb6d8e4bef979efa3360b1cf866ef971b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 18 Jul 2018 11:24:06 -0700
Subject: [PATCH 0346/2720] Create new token encoder and text2text problem for
 token transformer with copying mechanism.

PiperOrigin-RevId: 205108126
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 .../data_generators/pointer_generator_word.py | 197 ++++++++++++++++++
 2 files changed, 198 insertions(+)
 create mode 100644 tensor2tensor/data_generators/pointer_generator_word.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 419c02d48..b3e85bf77 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -51,6 +51,7 @@
     "tensor2tensor.data_generators.multinli",
     "tensor2tensor.data_generators.program_search",
     "tensor2tensor.data_generators.ocr",
+    "tensor2tensor.data_generators.pointer_generator_word",
     "tensor2tensor.data_generators.problem_hparams",
     "tensor2tensor.data_generators.ptb",
     "tensor2tensor.data_generators.qnli",
diff --git a/tensor2tensor/data_generators/pointer_generator_word.py b/tensor2tensor/data_generators/pointer_generator_word.py
new file mode 100644
index 000000000..7e17c06c1
--- /dev/null
+++ b/tensor2tensor/data_generators/pointer_generator_word.py
@@ -0,0 +1,197 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generator for pointer-generator for word transformer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+
+@registry.register_problem
+class Text2textCopyableTokens(text_problems.Text2textTmpdirTokens):
+  """Allows training a variant of Text2textTmpdirTokens that supports copying.
+
+  Handling the case where the input contains OOV tokens. Store a temporary vocab
+  ID for source OOV, so that the decoder can directly copy from the input.
+  Uses TokenTextEncoderOov as the vocab encoder.
+  """
+
+  def get_or_create_vocab(self, data_dir, tmp_dir, force_get=False):
+    vocab_filename = os.path.join(data_dir, self.vocab_filename)
+    encoder = text_encoder.TokenTextEncoderOov(
+        vocab_filename, replace_oov=self.oov_token)
+    return encoder
+
+  def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
+    generator = self.generate_samples(data_dir, tmp_dir, dataset_split)
+    encoder = self.get_or_create_vocab(data_dir, tmp_dir)
+    return self.text2text_generate_encoded_oovs(
+        generator, encoder, has_inputs=self.has_inputs)
+
+  def text2text_generate_encoded_oovs(self,
+                                      sample_generator,
+                                      vocab,
+                                      targets_vocab=None,
+                                      has_inputs=True):
+    """Encode Text2Text samples from the generator with the vocab."""
+    targets_vocab = targets_vocab or vocab
+    for sample in sample_generator:
+      if has_inputs:
+        (sample["inputs"], sample["inputs_extend"], source_oovs,
+         _) = vocab.encode(sample["inputs"])
+        sample["inputs"].append(text_encoder.EOS_ID)
+        sample["inputs_extend"].append(text_encoder.EOS_ID)
+      # need to pass the source OOV tokens to the target encoder
+      sample["targets"], sample["targets_extend"] = targets_vocab.encode_target(
+          sample["targets"], source_oovs)
+      sample["targets"].append(text_encoder.EOS_ID)
+      sample["targets_extend"].append(text_encoder.EOS_ID)
+      yield sample
+
+  def example_reading_spec(self):
+    data_fields = {
+        "inputs": tf.VarLenFeature(tf.int64),
+        "inputs_extend": tf.VarLenFeature(tf.int64),
+        "targets": tf.VarLenFeature(tf.int64),
+        "targets_extend": tf.VarLenFeature(tf.int64)
+    }
+    data_items_to_decoders = None
+    return (data_fields, data_items_to_decoders)
+
+
+class TokenTextEncoderOov(text_encoder.TokenTextEncoder):
+  """Encoder based on a user-supplied vocabulary (file or list).
+
+  This encoder extends over TokenTextEncoder by additionally assigning distinct
+  temporary IDs to OOV tokens appearing in the source sequence. This facilitates
+  decoding with the pointer-generator mechanism using word-based tokenization.
+
+  NOTE: TokenTextEncoderOov does not conform to the TextEncoder API; it changes
+  the signature of encode and decode.
+  """
+
+  def encode(self, s):
+    """Converts a space-separated string of tokens to lists of ids.
+
+    Also store temporary vocabulary IDs for source OOV tokens. OOVs are
+    represented by their temporary OOV number. E.g., if the vocabulary size
+    is 50k and the source has 3 OOVs, then these temporary OOV numbers will
+    be 50000, 50001, 50002.
+
+    Args:
+      s: human-readable string to be converted.
+
+    Returns:
+      ids: list of integers
+      ids_extend: list of integers including extended temporary vocab IDs for
+      source OOVs.
+      oovs: A dict storing source OOV words, used for the decoder to copy. The
+      key is OOV word, and the value is the order they appear in the source,
+      starting from 0.
+      source_oov_id_to_token: a list of source OOV tokens, in the same order as
+      they appear in the source.
+    """
+    sentence = s
+    tokens = sentence.strip().split()
+    ids = []
+    ids_extend = []
+    oovs = {}
+    for t in tokens:
+      if t in self._token_to_id:
+        ids.append(self._token_to_id[t])
+        ids_extend.append(self._token_to_id[t])
+      else:
+        next_oov_id = len(oovs)
+        oov_num = oovs.get(t, next_oov_id)
+        if oov_num == next_oov_id:
+          oovs[t] = oov_num
+        ids_extend.append(self.vocab_size + oov_num)
+        ids.append(self._token_to_id[self._replace_oov])
+    source_oov_id_to_token = [""] * len(oovs)
+    for oov in oovs:
+      source_oov_id_to_token[oovs[oov]] = oov
+    if self._reverse:
+      return ids[::-1], ids_extend[::-1], oovs, source_oov_id_to_token
+    else:
+      return ids, ids_extend, oovs, source_oov_id_to_token
+
+  def encode_target(self, target, source_oovs):
+    """Converts a space-separated string of tokens to lists of ids.
+
+    Also store a version of extened vocabulary IDs.
+    For target OOVs that are in the source, encode them using the temporary
+    vocab IDs.
+    For target OOVs not in the source, encode them as <UNK>
+
+    Args:
+      target: target string
+      source_oovs: source OOV words stored in dict, key is the word, value is
+      the order in which they appear in the source starting from 0
+
+    Returns:
+      ids: list of integers
+      ids_extend: list of integers including extended vocabulary IDs.
+    """
+    tokens = target.strip().split()
+    ids = []
+    ids_extend = []
+    for t in tokens:
+      if t in self._token_to_id:
+        i = self._token_to_id[t]
+        ids.append(i)
+        ids_extend.append(i)
+      else:
+        ids.append(self._token_to_id[self._replace_oov])
+        if t in source_oovs:
+          vocab_idx = self.vocab_size + source_oovs[t]
+          ids_extend.append(vocab_idx)
+        else:
+          ids_extend.append(self._token_to_id[self._replace_oov])
+    if self._reverse:
+      return ids[::-1], ids_extend[::-1]
+    else:
+      return ids, ids_extend
+
+  def decode_oov(self, ids, source_oov):
+    return " ".join(self.decode_list_oov(ids, source_oov))
+
+  def decode_list_oov(self, ids, source_oov_id_to_token):
+    """decode ids back to tokens, considering OOVs temporary IDs.
+
+    Args:
+      ids: vocab ids. Could possibly include source temporary OOV ID starting
+      from vocab_size.
+      source_oov_id_to_token: a list of source OOV tokens, with the order the
+      same as they appear in the source.
+
+    Returns:
+      decoded tokens, possibly including source OOV tokens.
+
+    """
+    seq = reversed(ids) if self._reverse else ids
+    tokens = []
+    for cur_id in seq:
+      if cur_id in self._id_to_token:
+        tokens.append(self._id_to_token[cur_id])
+      else:
+        tokens.append(source_oov_id_to_token[cur_id - self.vocab_size])
+    return tokens

From 702cc9d89fa5abe10fa9ef33168e4df85897b928 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 18 Jul 2018 12:07:26 -0700
Subject: [PATCH 0347/2720] Fixing the old renamed function in video decoding.

PiperOrigin-RevId: 205115133
---
 tensor2tensor/data_generators/video_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index eef40b56f..c341d1e82 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -57,7 +57,7 @@ def summarize_video_metrics(hook_args):
       metrics_results = video_metrics.compute_video_metrics_from_predictions(
           predictions)
     else:
-      metrics_results, _ = video_metrics.compute_video_metrics(
+      metrics_results, _ = video_metrics.compute_video_metrics_from_png_files(
           output_dirs, problem_name,
           hparams.video_num_target_frames, frame_shape)
 

From afc9ba896b65385a97a526768d2dd30f73dd53e6 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 18 Jul 2018 12:16:59 -0700
Subject: [PATCH 0348/2720] Internal change

PiperOrigin-RevId: 205116622
---
 .../get_references_web_single_group.py        | 59 +-------------
 tensor2tensor/data_generators/wikisum/html.py | 76 +++++++++++++++++++
 .../data_generators/wikisum/wikisum.py        |  6 +-
 3 files changed, 81 insertions(+), 60 deletions(-)
 create mode 100644 tensor2tensor/data_generators/wikisum/html.py

diff --git a/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py b/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
index 5198fec5a..d6678a421 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
@@ -30,9 +30,9 @@
 
 import asyncio
 import aiohttp
-import bs4
 import tensorflow as tf
 
+from tensor2tensor.data_generators.wikisum import html
 from tensor2tensor.data_generators.wikisum import utils
 
 
@@ -101,63 +101,8 @@ def shard(items, num_shards):
   return sharded
 
 
-def soup_strings(soup):
-  paragraph_tags = set(["caption", "details", "h1", "h2", "h3", "h4", "h5",
-                        "h6", "li", "p", "td", "div", "span"])
-
-  skip_children = None
-  for descendant in soup.descendants:
-    # If we've treated a tag as a contiguous paragraph, don't re-emit the
-    # children (see below).
-    if skip_children is not None:
-      try:
-        in_skip = descendant in skip_children
-      except RecursionError:
-        # Possible for this check to hit a nasty infinite recursion because of
-        # BeautifulSoup __eq__ checks.
-        in_skip = True
-      if in_skip:
-        continue
-      else:
-        skip_children = None
-
-    # Treat some tags as contigous paragraphs, regardless of other tags nested
-    # inside (like <a> or <b>).
-    if isinstance(descendant, bs4.Tag):
-      if descendant.name in paragraph_tags:
-        if descendant.find_all(paragraph_tags):
-          # If there are nested paragraph tags, don't treat it as a single
-          # contiguous tag.
-          continue
-        skip_children = list(descendant.descendants)
-        text = " ".join(descendant.get_text(" ", strip=True).split())
-        if text:
-          yield text
-        continue
-
-    if (isinstance(descendant, bs4.Comment) or
-        not isinstance(descendant, bs4.NavigableString)):
-      continue
-
-    text = " ".join(descendant.strip().split())
-    if text:
-      yield text
-
-
 def mp_get_text(url, html):
-  return url, get_text_from_html(html)
-
-
-def get_text_from_html(html):
-  try:
-    soup = bs4.BeautifulSoup(html, 'html.parser')
-  except:
-    # Some docs don't parse
-    return ""
-  # Remove script and style tags
-  for s in soup(["script", "style"]):
-    s.decompose()
-  return "\n".join([s for s in soup_strings(soup)])
+  return url, html.get_text_from_html(html)
 
 
 def encode(s):
diff --git a/tensor2tensor/data_generators/wikisum/html.py b/tensor2tensor/data_generators/wikisum/html.py
new file mode 100644
index 000000000..d27cf28fa
--- /dev/null
+++ b/tensor2tensor/data_generators/wikisum/html.py
@@ -0,0 +1,76 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Utils to parse HTML content into plaintext."""
+
+import bs4
+
+
+def get_text_from_html(html):
+  """Returns a plaintext representation of HTML content."""
+
+  try:
+    soup = bs4.BeautifulSoup(html, "html.parser")
+  except:
+    # Some docs don't parse
+    return ""
+  # Remove script and style tags
+  for s in soup(["script", "style"]):
+    s.decompose()
+  return "\n".join([s for s in _soup_strings(soup)])
+
+
+def _soup_strings(soup):
+  paragraph_tags = set([
+      "caption", "details", "h1", "h2", "h3", "h4", "h5", "h6", "li", "p", "td",
+      "div", "span"
+  ])
+
+  skip_children = None
+  for descendant in soup.descendants:
+    # If we've treated a tag as a contiguous paragraph, don't re-emit the
+    # children (see below).
+    if skip_children is not None:
+      try:
+        in_skip = descendant in skip_children
+      except RecursionError:
+        # Possible for this check to hit a nasty infinite recursion because of
+        # BeautifulSoup __eq__ checks.
+        in_skip = True
+      if in_skip:
+        continue
+      else:
+        skip_children = None
+
+    # Treat some tags as contiguous paragraphs, regardless of other tags nested
+    # inside (like <a> or <b>).
+    if isinstance(descendant, bs4.Tag):
+      if descendant.name in paragraph_tags:
+        if descendant.find_all(paragraph_tags):
+          # If there are nested paragraph tags, don't treat it as a single
+          # contiguous tag.
+          continue
+        skip_children = list(descendant.descendants)
+        text = " ".join(descendant.get_text(" ", strip=True).split())
+        if text:
+          yield text
+        continue
+
+    if (isinstance(descendant, bs4.Comment) or
+        not isinstance(descendant, bs4.NavigableString)):
+      continue
+
+    text = " ".join(descendant.strip().split())
+    if text:
+      yield text
diff --git a/tensor2tensor/data_generators/wikisum/wikisum.py b/tensor2tensor/data_generators/wikisum/wikisum.py
index 2609707ff..9fe1e7971 100644
--- a/tensor2tensor/data_generators/wikisum/wikisum.py
+++ b/tensor2tensor/data_generators/wikisum/wikisum.py
@@ -341,7 +341,7 @@ def _tokens_to_score(tokens):
   return {t for t in tokens if re.search("[a-z0-9]", t)}
 
 
-def _rank_reference_paragraphs(wiki_title, references_content):
+def rank_reference_paragraphs(wiki_title, references_content):
   """Rank and return reference paragraphs by tf-idf score on title tokens."""
   title_tokens = _tokens_to_score(set(
       tokenizer.encode(text_encoder.native_to_unicode(wiki_title))))
@@ -428,8 +428,8 @@ def example_generator():
 
         # Rank reference paragraphs with TFIDF
         wiki_title = _normalize_text(wiki.title)
-        ranked_paragraphs = _rank_reference_paragraphs(wiki_title,
-                                                       wiki_ref_content)
+        ranked_paragraphs = rank_reference_paragraphs(wiki_title,
+                                                      wiki_ref_content)
 
         # Construct inputs from Wiki title and references
         inputs = []

From 5124b5178cc8794bc6e9080e2c4bb262a797d236 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 18 Jul 2018 13:12:23 -0700
Subject: [PATCH 0349/2720] Corrections to make PPO lr tuning and similarity
 transformer run.

PiperOrigin-RevId: 205124670
---
 tensor2tensor/models/__init__.py                        | 1 +
 tensor2tensor/models/research/similarity_transformer.py | 1 +
 tensor2tensor/rl/model_rl_experiment.py                 | 1 +
 3 files changed, 3 insertions(+)

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 0130ebd06..abe3f1f22 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -45,6 +45,7 @@
 from tensor2tensor.models.research import multimodel
 from tensor2tensor.models.research import next_frame
 from tensor2tensor.models.research import rl
+from tensor2tensor.models.research import similarity_transformer
 from tensor2tensor.models.research import super_lm
 from tensor2tensor.models.research import transformer_moe
 from tensor2tensor.models.research import transformer_nat
diff --git a/tensor2tensor/models/research/similarity_transformer.py b/tensor2tensor/models/research/similarity_transformer.py
index ed915c3f4..f60e261f9 100644
--- a/tensor2tensor/models/research/similarity_transformer.py
+++ b/tensor2tensor/models/research/similarity_transformer.py
@@ -85,5 +85,6 @@ def encode(self, features, input_key):
     return encoder_output
 
   def infer(self, features=None, **kwargs):
+    del kwargs
     predictions, _ = self(features)
     return predictions
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index a4318fe13..cb9c95ffe 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -503,6 +503,7 @@ def rl_modelrl_base():
       # though it is not necessary.
       ppo_epoch_length=60,
       ppo_num_agents=16,
+      ppo_learning_rate=0.,  # Will be changed, just so it exists.
       # Whether the PPO agent should be restored from the previous iteration, or
       # should start fresh each time.
       ppo_continue_training=True,

From b7b448da238e2dcc6755b8d7caecf8f6bf93b9b8 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 18 Jul 2018 13:29:27 -0700
Subject: [PATCH 0350/2720] Fixed typo.

PiperOrigin-RevId: 205127272
---
 tensor2tensor/models/distillation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/distillation.py b/tensor2tensor/models/distillation.py
index 1af4ee843..839501a95 100644
--- a/tensor2tensor/models/distillation.py
+++ b/tensor2tensor/models/distillation.py
@@ -28,7 +28,7 @@
 class Distillation(t2t_model.T2TModel):
   """Distillation from a teacher to student network.
 
-  First, a teacher is train on a task; Second, a student is trained to perform
+  First, a teacher is trained on a task; Second, a student is trained to perform
   the task while matching the teacher's softened outputs. For more details, see
   the paper below.
 

From 46724a0e04d76edb62fdf1ae14fa7a3ba17c65ef Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 18 Jul 2018 13:40:56 -0700
Subject: [PATCH 0351/2720] Increasing the number of preprocessing threads to
 num_cores.

PiperOrigin-RevId: 205129038
---
 tensor2tensor/data_generators/problem.py    | 9 ++++++++-
 tensor2tensor/models/research/next_frame.py | 1 +
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index ca7b8f9c5..021c0ef41 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -18,6 +18,7 @@
 from __future__ import print_function
 import collections
 import functools
+import multiprocessing
 import os
 import random
 
@@ -170,6 +171,12 @@ def _file_num_records_cached(filename):
 _file_num_records_cache = {}
 
 
+def cpu_count():
+  """Return the number of available cores."""
+  num_available_cores = multiprocessing.cpu_count()
+  return num_available_cores
+
+
 class Problem(object):
   """Problem base class. Specifies a T2T problem.
 
@@ -785,7 +792,7 @@ def input_fn(self,
     if config and config.use_tpu:
       num_threads = 64
     else:
-      num_threads = 4 if is_training else 1
+      num_threads = cpu_count() if is_training else 1
 
     max_length = self.max_length(hparams)
 
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index d58f96505..24c568c71 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -1184,6 +1184,7 @@ def next_frame_stochastic_cutoff():
 def next_frame_stochastic_tiny():
   """SV2P model with additional cutoff in L2 loss for environments like pong."""
   hparams = next_frame_stochastic()
+  hparams.batch_size = 2
   hparams.tiny_mode = True
   hparams.num_masks = 1
   hparams.video_modality_loss_cutoff = 0.4

From e38969ea254d62fb3670509807997dfbaac1d0fd Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Wed, 18 Jul 2018 14:42:22 -0700
Subject: [PATCH 0352/2720] Move reconstruction into body and use custom
 training loss

PiperOrigin-RevId: 205139940
---
 tensor2tensor/models/research/autoencoders.py | 36 +++++++++++++------
 1 file changed, 26 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 5eaaf9c47..9fbd76680 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -180,30 +180,46 @@ def body(self, features):
       return x, {"bottleneck_loss": 0.0}
     # Cut to the right size and mix before returning.
     res = x[:, :shape[1], :shape[2], :]
+
+    num_channels = self.hparams.problem.num_channels
+    reconstr = tf.layers.dense(res, num_channels)
+    reconstr = tf.nn.sigmoid(reconstr)
+    reconstr = 256. * reconstr - 0.5
     # Add GAN loss if requested.
     gan_loss = 0.0
     if hparams.gan_loss_factor != 0.0:
       # Split back if we added a purely sampled batch.
-      res_gan, res = tf.split(res, 2, axis=0)
-      num_channels = self.hparams.problem.num_channels
-      res_rgb = common_layers.convert_real_to_rgb(
-          tf.nn.sigmoid(tf.layers.dense(res_gan, num_channels, name="gan_rgb")))
+      reconstr_gan, reconstr = tf.split(reconstr, 2, axis=0)
       tf.summary.image(
-          "gan", common_layers.tpu_safe_image_summary(res_rgb), max_outputs=1)
+          "gan",
+          common_layers.tpu_safe_image_summary(reconstr_gan),
+          max_outputs=1)
       orig_rgb = tf.to_float(features["targets_raw"])
 
       def discriminate(x):
         return self.discriminator(x, is_training=is_training)
 
       gan_loss = common_layers.sliced_gan_loss(orig_rgb,
-                                               reverse_gradient(res_rgb),
+                                               reverse_gradient(reconstr_gan),
                                                discriminate,
                                                self.hparams.num_sliced_vecs)
       gan_loss *= hparams.gan_loss_factor
-    # Mix the final result and return.
-    res = common_layers.mix(res, features["targets"],
-                            hparams.bottleneck_warmup_steps // 2, is_training)
-    return res, {"bottleneck_loss": b_loss, "gan_loss": -gan_loss}
+
+    tf.summary.image(
+        "ae", common_layers.tpu_safe_image_summary(reconstr), max_outputs=1)
+
+    # Project to correct vocab_size/channels
+    training = tf.reduce_mean(
+        tf.square(tf.to_float(features["targets_raw"]) - reconstr))
+
+    outputs = tf.round(reconstr)
+    outputs = tf.one_hot(tf.to_int32(outputs), 256)
+
+    return outputs, {
+        "training": training,
+        "b_loss": b_loss,
+        "gan_loss": -gan_loss
+    }
 
   def sample(self, features=None, shape=None):
     del features, shape

From 9b20e7d32b4f89bfdf02ec3f85f70c628bf49425 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 18 Jul 2018 16:09:23 -0700
Subject: [PATCH 0353/2720] Added TEST shard to the dataset split. Data
 generator will now create

PiperOrigin-RevId: 205154358
---
 tensor2tensor/data_generators/timeseries.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index f1367b573..80d2da6bf 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -45,14 +45,16 @@ def is_generate_per_split(self):
 
   @property
   def dataset_splits(self):
-    """Splits of data to produce and number of output shards for each."""
-    # 10% evaluation data
+    """Splits of data to produce and number the output shards for each."""
     return [{
         "split": problem.DatasetSplit.TRAIN,
         "shards": self.num_train_shards,
     }, {
         "split": problem.DatasetSplit.EVAL,
         "shards": self.num_eval_shards,
+    }, {
+        "split": problem.DatasetSplit.TEST,
+        "shards": self.num_test_shards,
     }]
 
   @property
@@ -65,6 +67,11 @@ def num_eval_shards(self):
     """Number of eval shards."""
     return 1
 
+  @property
+  def num_test_shards(self):
+    """Number of test shards."""
+    return 1
+
   @property
   def num_series(self):
     """Number of timeseries."""
@@ -185,6 +192,11 @@ def num_eval_shards(self):
     """Number of eval shards."""
     return 1
 
+  @property
+  def num_test_shards(self):
+    """Number of eval shards."""
+    return 0
+
   @property
   def num_series(self):
     """Number of timeseries."""

From 58b51414a746f854c74aec9cd2296c55b4bf4521 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 19 Jul 2018 11:47:54 -0700
Subject: [PATCH 0354/2720] Change number of input timestamps from 500 to 250
 for initial tests.

PiperOrigin-RevId: 205276381
---
 tensor2tensor/data_generators/timeseries.py      | 2 +-
 tensor2tensor/data_generators/timeseries_test.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index 80d2da6bf..71a9fb5e0 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -240,7 +240,7 @@ def num_series(self):
   @property
   def num_input_timestamps(self):
     """Number of timestamps to include in the input."""
-    return 500
+    return 250
 
   @property
   def num_target_timestamps(self):
diff --git a/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index ccfc34565..038ddedcc 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -63,7 +63,7 @@ def testTimeseriesToyProblem(self):
   def testTimeseriesSyntheticData10Series100kSamples(self):
     problem = timeseries.TimeseriesSyntheticDataSeries10Samples100k()
     self.assertEqual(10, problem.num_series)
-    self.assertEqual(500, problem.num_input_timestamps)
+    self.assertEqual(250, problem.num_input_timestamps)
     self.assertEqual(100, problem.num_target_timestamps)
 
 

From 78084b542a4b74883e926f6502c93bc176fc40d0 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 20 Jul 2018 03:28:59 -0700
Subject: [PATCH 0355/2720] Internal change

PiperOrigin-RevId: 205376281
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/celebahq.py     | 119 ++++++++++++++++++
 tensor2tensor/data_generators/problem.py      |  53 +++++---
 tensor2tensor/models/image_transformer.py     |  77 +++++++++---
 4 files changed, 214 insertions(+), 36 deletions(-)
 create mode 100644 tensor2tensor/data_generators/celebahq.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index b3e85bf77..099abb74a 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -27,6 +27,7 @@
     "tensor2tensor.data_generators.babi_qa",
     "tensor2tensor.data_generators.bair_robot_pushing",
     "tensor2tensor.data_generators.celeba",
+    "tensor2tensor.data_generators.celebahq",
     "tensor2tensor.data_generators.cifar",
     "tensor2tensor.data_generators.cipher",
     "tensor2tensor.data_generators.cnn_dailymail",
diff --git a/tensor2tensor/data_generators/celebahq.py b/tensor2tensor/data_generators/celebahq.py
new file mode 100644
index 000000000..24e43e493
--- /dev/null
+++ b/tensor2tensor/data_generators/celebahq.py
@@ -0,0 +1,119 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""CelebA-HQ."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+from tensor2tensor.data_generators import image_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.utils import metrics
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+@registry.register_problem
+class ImageCelebahq128(image_utils.ImageProblem):
+  """CelebA-HQ dataset, downsampled as 128x128."""
+
+  def dataset_filename(self):
+    return "image_celebahq-128"
+
+  def example_reading_spec(self):
+    data_fields = {
+        "image/encoded": tf.FixedLenFeature((), tf.string),
+        "image/format": tf.FixedLenFeature((), tf.string, default_value="png"),
+    }
+    _, data_items_to_decoders = super(
+        ImageCelebahq128, self).example_reading_spec()
+    return data_fields, data_items_to_decoders
+
+  def filepattern(self, data_dir, mode, shard=None):
+    """Get filepattern for data files for mode.
+
+    Args:
+      data_dir: str, data directory.
+      mode: DatasetSplit
+      shard: int, if provided, will only read data from the specified shard.
+
+    Returns:
+      filepattern str
+    """
+    path = os.path.join(data_dir, self.dataset_filename())
+    if shard is not None:
+      shard_str = "%05d" % shard
+    elif mode == problem.DatasetSplit.TRAIN:
+      # Use the first 90 shards.
+      shard_str = "000[0-8]"
+    else:
+      assert mode in [problem.DatasetSplit.EVAL,
+                      tf.estimator.ModeKeys.PREDICT,
+                      problem.DatasetSplit.TEST]
+      # Use the last 10 shards.
+      shard_str = "0009"
+
+    return "%s-%s*" % (path, shard_str)
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    raise NotImplementedError("Data preprocessing for CelebA-HQ is not "
+                              "currently available. Please follow the steps "
+                              "in https://github.com/tkarras/progressive_growin"
+                              "g_of_gans.")
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    p.batch_size_multiplier = 1
+    p.input_modality = {"inputs": ("image:identity", 256)}
+    p.input_space_id = 1
+
+  def preprocess_example(self, example, mode, hparams):
+    del mode, hparams  # unused
+    example["inputs"].set_shape((128, 128, 3))
+    return example
+
+
+@registry.register_problem
+class ImageCelebahq128Dmol(ImageCelebahq128):
+  """CelebA-HQ dataset with discretized mixture of logistics for evaluation."""
+
+  def eval_metrics(self):
+    return [
+        metrics.Metrics.DMOL_PERPLEXITY
+    ]
+
+
+@registry.register_problem
+class ImageCelebahq256(ImageCelebahq128):
+  """CelebA-HQ dataset, downsampled as 256x256."""
+
+  def dataset_filename(self):
+    return "image_celebahq-256"
+
+  def preprocess_example(self, example, mode, hparams):
+    del mode, hparams  # unused
+    example["inputs"].set_shape((256, 256, 3))
+    return example
+
+
+@registry.register_problem
+class ImageCelebahq256Dmol(ImageCelebahq256):
+  """CelebA-HQ dataset with discretized mixture of logistics for evaluation."""
+
+  def eval_metrics(self):
+    return [
+        metrics.Metrics.DMOL_PERPLEXITY
+    ]
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 021c0ef41..e6253120f 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -487,18 +487,24 @@ def maybe_reverse_features(self, feature_map):
     """Reverse features between inputs and targets if the problem is '_rev'."""
     if not self._was_reversed:
       return
-    inputs, targets = feature_map["inputs"], feature_map["targets"]
-    feature_map["inputs"], feature_map["targets"] = targets, inputs
-    if "inputs_segmentation" in feature_map:
-      inputs_seg = feature_map["inputs_segmentation"]
-      targets_seg = feature_map["targets_segmentation"]
-      feature_map["inputs_segmentation"] = targets_seg
+    inputs = feature_map.pop("inputs", None)
+    targets = feature_map.pop("targets", None)
+    inputs_seg = feature_map.pop("inputs_segmentation", None)
+    targets_seg = feature_map.pop("targets_segmentation", None)
+    inputs_pos = feature_map.pop("inputs_position", None)
+    targets_pos = feature_map.pop("targets_position", None)
+    if inputs is not None:
+      feature_map["targets"] = inputs
+    if targets is not None:
+      feature_map["inputs"] = targets
+    if inputs_seg is not None:
       feature_map["targets_segmentation"] = inputs_seg
-    if "inputs_position" in feature_map:
-      inputs_pos = feature_map["inputs_position"]
-      targets_pos = feature_map["targets_position"]
-      feature_map["inputs_position"] = targets_pos
+    if targets_seg is not None:
+      feature_map["inputs_segmentation"] = targets_seg
+    if inputs_pos is not None:
       feature_map["targets_position"] = inputs_pos
+    if targets_pos is not None:
+      feature_map["inputs_position"] = targets_pos
 
   def maybe_copy_features(self, feature_map):
     if not self._was_copy:
@@ -1009,22 +1015,33 @@ def _reverse_problem_hparams(p_hparams):
   p = p_hparams
 
   # Swap modalities.
-  input_modality = p.input_modality["inputs"]
+  input_modality = p.input_modality.get("inputs")
   target_modality = p.target_modality
-  p.input_modality["inputs"] = target_modality
   p.target_modality = input_modality
+  if target_modality is not None:
+    p.input_modality["inputs"] = target_modality
+  else:
+    p.input_modality = {}
 
   # Swap vocabularies.
-  input_vocabulary = p.vocabulary["inputs"]
-  target_vocabulary = p.vocabulary["targets"]
-  p.vocabulary["inputs"] = target_vocabulary
-  p.vocabulary["targets"] = input_vocabulary
+  input_vocabulary = p.vocabulary.pop("inputs", None)
+  target_vocabulary = p.vocabulary.pop("targets", None)
+  if input_vocabulary is not None:
+    p.vocabulary["targets"] = input_vocabulary
+  if target_vocabulary is not None:
+    p.vocabulary["inputs"] = target_vocabulary
 
   # Swap input/target space ids.
   input_space_id = p.input_space_id
   target_space_id = p.target_space_id
-  p.input_space_id = target_space_id
-  p.target_space_id = input_space_id
+  if input_space_id is not None:
+    p.target_space_id = input_space_id
+  else:
+    p.target_space_id = SpaceID.GENERIC
+  if target_space_id is not None:
+    p.input_space_id = target_space_id
+  else:
+    p.input_space_id = SpaceID.GENERIC
 
   # Mark that p was reversed.
   p.was_reversed = True
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index e59431d32..aa6d08fb2 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -44,7 +44,6 @@ class Imagetransformer(t2t_model.T2TModel):
 
   def body(self, features):
     hparams = copy.copy(self._hparams)
-    inputs = features["inputs"]
     targets = features["targets"]
     if (hparams.likelihood == cia.DistributionType.DMOL and
         (hparams.target_modality != "image:image_channel_bottom_identity" or
@@ -63,6 +62,7 @@ def body(self, features):
     decoder_input, rows, cols = cia.prepare_decoder(targets, hparams)
     # Add class label to decoder input.
     if not hparams.unconditional:
+      inputs = features["inputs"]
       decoder_input += tf.reshape(
           inputs,
           [common_layers.shape_list(targets)[0], 1, 1, hparams.hidden_size])
@@ -394,20 +394,6 @@ def imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_b():
   return hparams
 
 
-@registry.register_hparams
-def imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_e():
-  hparams = imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_b()
-  hparams.learning_rate_warmup_steps = 16000
-  return hparams
-
-
-@registry.register_hparams
-def imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_f():
-  hparams = imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_b()
-  hparams.num_mixtures = 5
-  return hparams
-
-
 @registry.register_hparams
 def imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_g():
   hparams = imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_b()
@@ -424,7 +410,6 @@ def imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_g():
 def imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_k():
   hparams = imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_g()
   hparams.num_decoder_layers = 12
-  hparams.clip_grad_norm = 0.
   return hparams
 
 
@@ -468,6 +453,60 @@ def imagetransformerpp_base_14l_8h_big_uncond_dr03_dan_p():
   return hparams
 
 
+@registry.register_hparams
+def imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_m_bs1():
+  """For 128x128."""
+  # TODO(trandustin): why are these running? max_length and img_len not set
+  # 256x256 was also training without setting max_length
+  hparams = imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_m()
+  hparams.batch_size = 1
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_14l_8h_big_uncond_dr03_dan_p_bs1():
+  """For 128x128."""
+  hparams = imagetransformerpp_base_14l_8h_big_uncond_dr03_dan_p()
+  hparams.batch_size = 1
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_5l_8h_big_uncond_dr00_dan_g_bs1():
+  """For 256x256."""
+  hparams = imagetransformerpp_base_10l_8h_big_uncond_dr03_dan_g()
+  # TODO(trandustin): I forgot to set this in the runs! Maybe it's not used in
+  # image transformer training implementation?
+  # hparams.img_len = 256
+  hparams.max_length = 66000  # allow for 256x256
+  hparams.batch_size = 1
+  hparams.num_decoder_layers = 5
+  hparams.hidden_size = 128
+  hparams.filter_size = 128
+  hparams.attention_key_channels = 64
+  hparams.attention_value_channels = 64
+  hparams.layer_prepostprocess_dropout = 0.0
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_5l_8h_dr00_dan_g_bs1_adafactor():
+  """For 256x256."""
+  hparams = imagetransformerpp_base_5l_8h_big_uncond_dr00_dan_g_bs1()
+  # Use Adafactor which uses less memory than Adam, and its recommendations.
+  hparams.optimizer = "Adafactor"
+  hparams.learning_rate_schedule = "rsqrt_decay"
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformerpp_base_6l_8h_dr00_dan_g_bs1_adafactor():
+  """For 256x256."""
+  hparams = imagetransformerpp_base_5l_8h_dr00_dan_g_bs1_adafactor()
+  hparams.num_decoder_layers = 6
+  return hparams
+
+
 @registry.register_hparams
 def imagetransformerpp_base_14l_8h_big_uncond_dr03_dan_eval():
   """Gets to 2.92 in just under 4 days on 8 p100s."""
@@ -741,15 +780,17 @@ def imagetransformer_bas8l_8h_big_uncond_dr03_imgnet():
 @registry.register_hparams
 def imagetransformer_tiny():
   hparams = imagetransformer_base()
-  hparams.num_hidden_layers = 2
+  hparams.num_decoder_layers = 2
   hparams.hidden_size = 64
   hparams.batch_size = 1
+  hparams.unconditional = True
+  hparams.max_length = 66000  # allow for 256x256
   return hparams
 
 
 @registry.register_hparams
 def imagetransformerpp_tiny():
-  hparams = imagetransformer_base()
+  hparams = imagetransformer_tiny()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
   hparams.target_modality = "image:image_channel_bottom_identity"

From 1ba41509006f93a1cc41d3664a81182b16ec0285 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 20 Jul 2018 10:52:00 -0700
Subject: [PATCH 0356/2720] Internal change

PiperOrigin-RevId: 205420980
---
 tensor2tensor/utils/trainer_lib.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 09439dd40..64e5fa7f9 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -314,8 +314,7 @@ def evaluate(self):
     return self._estimator.evaluate(
         self._eval_spec.input_fn,
         steps=self._eval_spec.steps,
-        hooks=self._eval_spec.hooks,
-        name="eval")
+        hooks=self._eval_spec.hooks)
 
   def evaluate_on_train_data(self):
     self._estimator.evaluate(

From cb0ae2e960b65278e5891ece38a7f902bbe7548b Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 20 Jul 2018 11:00:58 -0700
Subject: [PATCH 0357/2720] Correct import error handling to account for more
 deeply nested modules (e.g. wikisum.wikisum)

PiperOrigin-RevId: 205422617
---
 tensor2tensor/data_generators/all_problems.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 099abb74a..0188a84ff 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -86,7 +86,11 @@
 
 def _py_err_msg(module):
   if six.PY2:
-    msg = "No module named %s" % module.split(".")[-1]
+    # Py2 error will reference the module relative to the current module
+    shared_module = "data_generators."
+    start_idx = module.index(shared_module) + len(shared_module)
+    err_name = module[start_idx:]
+    msg = "No module named %s" % err_name
   else:
     msg = "No module named '%s'" % module
   return msg

From 3989fa10ece85653ff343e58c34b5a88c3062921 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 20 Jul 2018 13:36:26 -0700
Subject: [PATCH 0358/2720] fix Text2textTmpdirTokens problem.

PiperOrigin-RevId: 205445630
---
 tensor2tensor/data_generators/pointer_generator_word.py | 2 +-
 tensor2tensor/data_generators/text_problems.py          | 5 +++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/pointer_generator_word.py b/tensor2tensor/data_generators/pointer_generator_word.py
index 7e17c06c1..b2759113c 100644
--- a/tensor2tensor/data_generators/pointer_generator_word.py
+++ b/tensor2tensor/data_generators/pointer_generator_word.py
@@ -37,7 +37,7 @@ class Text2textCopyableTokens(text_problems.Text2textTmpdirTokens):
 
   def get_or_create_vocab(self, data_dir, tmp_dir, force_get=False):
     vocab_filename = os.path.join(data_dir, self.vocab_filename)
-    encoder = text_encoder.TokenTextEncoderOov(
+    encoder = TokenTextEncoderOov(
         vocab_filename, replace_oov=self.oov_token)
     return encoder
 
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index b67af6702..a17e47f3f 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -690,8 +690,9 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
     if not tf.gfile.Exists(vocab_filepath):
       token_encoder = self._generate_vocab(tmp_dir)
       token_encoder.store_to_file(vocab_filepath)
-    super(Text2textTmpdirTokens, self).generate_samples(data_dir, tmp_dir,
-                                                        dataset_split)
+    return super(Text2textTmpdirTokens, self).generate_samples(data_dir,
+                                                               tmp_dir,
+                                                               dataset_split)
 
 
 class ChoppedTextProblem(Text2SelfProblem):

From 9e899a2f3fd389cad9fe25d47f4914eb568b95d5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 20 Jul 2018 14:13:22 -0700
Subject: [PATCH 0359/2720] Added Gumbel-Softmax VQ-VAE discrete bottleneck
 with or without IAF for transformer_ae.

PiperOrigin-RevId: 205451561
---
 tensor2tensor/layers/discretization.py        | 304 ++++++++++++++++--
 tensor2tensor/layers/discretization_test.py   |  10 +-
 .../models/research/transformer_vae.py        |  91 +++++-
 3 files changed, 359 insertions(+), 46 deletions(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 1fd2e7fba..da90b07af 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -19,6 +19,7 @@
 
 from functools import partial
 
+from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 
 import tensorflow as tf
@@ -51,15 +52,16 @@ def slice_hidden(x, hidden_size, num_blocks):
   """Slice encoder hidden state into block_dim.
 
   Args:
-    x: Encoder hidden state of shape [-1, hidden_size].
+    x: Encoder hidden state of shape [..., 1, hidden_size].
     hidden_size: Dimension of the latent space.
     num_blocks: Number of blocks in DVQ.
 
   Returns:
-    Sliced states of shape [-1, num_blocks, block_dim].
+    Sliced states of shape [..., num_blocks, block_dim].
   """
-  block_dim = int(hidden_size // num_blocks)
-  x_sliced = tf.reshape(x, shape=[-1, num_blocks, block_dim])
+  block_dim = hidden_size // num_blocks
+  x_shape = common_layers.shape_list(x)
+  x_sliced = tf.reshape(x, shape=(x_shape[:-2] + [num_blocks, block_dim]))
   return x_sliced
 
 
@@ -68,22 +70,25 @@ def nearest_neighbor(x,
                      block_v_size,
                      random_top_k=1,
                      soft_em=False,
-                     num_samples=1):
+                     num_samples=1,
+                     summary=True):
   """Find the nearest element in means to elements in x.
 
   Args:
     x: Batch of encoder continuous latent states sliced/projected into shape
-      [-1, num_blocks, block_dim].
-    means: Embedding table of shpae [num_blocks, block_v_size, block_dim].
+      [batch_size, latent_dim, num_blocks, block_dim].
+    means: Embedding table of shape [num_blocks, block_v_size, block_dim].
     block_v_size: Number of table entries per block.
     random_top_k: Noisy top-k if this is bigger than 1 (Default: 1).
     soft_em: If True then use soft EM rather than hard EM (Default: False).
     num_samples: Number of samples to take in soft EM (Default: 1).
+    summary: If True then record summary histogram of entropies (Default: True).
 
   Returns:
     Tensor with nearest element in mean encoded in one-hot notation
     and distances.
   """
+  x = tf.reshape(x, [-1] + common_layers.shape_list(x)[2:])
   x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keep_dims=True)
   means_norm_sq = tf.reduce_sum(tf.square(means), axis=-1, keep_dims=True)
   scalar_prod = tf.matmul(
@@ -102,8 +107,14 @@ def nearest_neighbor(x,
         ],
         axis=1)
     nearest_hot = tf.one_hot(nearest_idx, depth=block_v_size)
+    neg_q_entropy = tf.reduce_sum(
+        nearest_hot * tf.expand_dims(tf.nn.log_softmax(-dist), 2), axis=2)
+    neg_q_entropy = tf.reduce_mean(neg_q_entropy)
     nearest_hot = tf.reduce_mean(nearest_hot, axis=-2)
+    if summary:
+      tf.summary.histogram("neg_q_entropy", tf.reshape(neg_q_entropy, [-1]))
   else:
+    neg_q_entropy = 0.
     if random_top_k > 1:
       _, top_k_idx = tf.nn.top_k(-dist, k=random_top_k)
       nearest_idx = tf.gather(
@@ -114,45 +125,74 @@ def nearest_neighbor(x,
     else:
       nearest_idx = tf.argmax(-dist, axis=-1)
     nearest_hot = tf.one_hot(nearest_idx, block_v_size)
-  return nearest_hot
+  return nearest_hot, neg_q_entropy
 
 
 def embedding_lookup(x,
                      means,
                      num_blocks,
                      block_v_size,
+                     bottleneck_kind="dvq",
                      random_top_k=1,
                      soft_em=False,
-                     num_samples=1):
+                     num_samples=1,
+                     do_hard_gumbel_softmax=False,
+                     do_iaf=False,
+                     approximate_gs_entropy=False):
   """Compute nearest neighbors and loss for training the embeddings via DVQ.
 
   Args:
     x: Batch of encoder continuous latent states sliced/projected into shape
-      [-1, num_blocks, block_dim].
+      [batch_size, latent_dim, num_blocks, block_dim].
     means: Embedding table of shape [num_blocks, block_v_size, block_dim].
     num_blocks: Number of blocks in DVQ.
     block_v_size: Number of table entries per block.
+    bottleneck_kind: Discrete bottleneck type (Default: "dvq").
     random_top_k: Noisy top-k if this is bigger than 1 (Default: 1).
     soft_em: If True then use soft EM rather than hard EM (Default: False).
     num_samples: Number of samples to use for soft EM (Default: 1).
+    do_hard_gumbel_softmax: Boolean determining whether we take hard or soft
+      Gumbel-Softmax samples for gumbel-softmax-dvq bottleneck (Default: False).
+    do_iaf: Boolean determining whether we use inverse autoregressive flows for
+      gumbel-softmax-dvq bottleneck (Default: False).
+    approximate_gs_entropy: If True, use the categorical density instead of the
+      Gumbel-Softmax density when calculating the sample entropy (Default:
+      False).
 
   Returns:
-    The nearest neighbor in one hot form, the nearest neighbor itself, the
-    commitment loss, embedding training loss and distances.
+    x_means_hot: The nearest neighbor in one hot form, with shape
+      [batch_size * latent_dim, num_blocks, block_v_size].
+    x_means: The nearest neighbor itself, with shape [batch_size * latent_dim,
+      num_blocks, block_dim].
+    q_loss: Commitment loss.
+    e_loss: Embedding training loss.
+    neg_q_entropy: Negative entropy of variational approximation (0 if q is
+      deterministic).
   """
-  x_means_hot = nearest_neighbor(
-      x,
-      means,
-      block_v_size,
-      random_top_k,
-      soft_em=soft_em,
-      num_samples=num_samples)
+  if bottleneck_kind == "gumbel-softmax-dvq":
+    x_means_hot, neg_q_entropy = gumbel_softmax_nearest_neighbor_dvq(
+        x,
+        means,
+        block_v_size,
+        hard=do_hard_gumbel_softmax,
+        num_samples=num_samples,
+        do_iaf=do_iaf,
+        approximate_gs_entropy=approximate_gs_entropy)
+  else:
+    x_means_hot, neg_q_entropy = nearest_neighbor(
+        x,
+        means,
+        block_v_size,
+        random_top_k,
+        soft_em=soft_em,
+        num_samples=num_samples)
   x_means_hot_flat = tf.reshape(x_means_hot, [-1, num_blocks, block_v_size])
   x_means = tf.matmul(tf.transpose(x_means_hot_flat, perm=[1, 0, 2]), means)
   x_means = tf.transpose(x_means, [1, 0, 2])
+  x = tf.reshape(x, [-1] + common_layers.shape_list(x)[2:])
   q_loss = tf.reduce_mean(tf.square((tf.stop_gradient(x) - x_means)))
   e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
-  return x_means_hot, x_means, q_loss, e_loss
+  return x_means_hot, x_means, q_loss, e_loss, neg_q_entropy
 
 
 def bit_to_int(x_bit, num_bits, base=2):
@@ -249,7 +289,7 @@ def embed(x,
     elif bottleneck_kind == "gumbel-softmax":
       hot = tf.one_hot(x, 2**z_size)
       h1 = tf.layers.dense(hot, hidden_size, name="dae_dense")
-    elif bottleneck_kind == "dvq":
+    elif bottleneck_kind in ["dvq", "gumbel-softmax-dvq"]:
       if block_v_size is None:
         raise ValueError("Bottleneck kind is dvq but block_v_size is None.")
 
@@ -439,7 +479,10 @@ def discrete_bottleneck(x,
                         ema=True,
                         ema_count=None,
                         ema_means=None,
-                        summary=True):
+                        summary=True,
+                        do_hard_gumbel_softmax=False,
+                        do_iaf=False,
+                        approximate_gs_entropy=False):
   """Discretization bottleneck for latent variables.
 
   Args:
@@ -454,7 +497,7 @@ def discrete_bottleneck(x,
     startup_steps: Number of steps after which latent predictor is trained
       (Default: 50000).
     bottleneck_kind: Kind of discretization bottleneck to use; one of dvq,
-      semhash, gumbel-softmax (Default: dvq).
+      semhash, gumbel-softmax, gumbel-softmax-dvq (Default: dvq).
     num_blocks: Number of blocks to use for decomposed vector
       quantization (Default: 2).
     num_residuals: Number of residual units used to compute nearest
@@ -480,6 +523,13 @@ def discrete_bottleneck(x,
       examples in a batch it was the closest to (Default: None).
     ema_means: Exponentially averaged version of the embeddings (Default: None).
     summary: If True, then write summaries (Default: True).
+    do_hard_gumbel_softmax: Boolean determining hard or soft Gumbel-Softmax
+      samples (Default: False).
+    do_iaf: Boolean determining whether we do inverse autoregresive flows for
+      Gumbel-Softmax DVQ bottleneck (Default: False).
+    approximate_gs_entropy: If true, we approximate the Gumbel-Softmax density
+      as a categorical distribution when calculating the sample entropy
+      (Default: False).
 
   Returns:
     Embedding to pass to the decoder, discrete latent, loss, and the embedding
@@ -490,7 +540,7 @@ def discrete_bottleneck(x,
     ema_count or ema_means is None if we are using ema, or unknown args.
   """
   block_v_size = None
-  if bottleneck_kind == "dvq":
+  if bottleneck_kind in ["dvq", "gumbel-softmax-dvq"]:
     # Define the dvq parameters
     assert means is not None
 
@@ -509,7 +559,7 @@ def discrete_bottleneck(x,
     block_v_size = 2**(z_size_per_residual / num_blocks)
     block_v_size = int(block_v_size)
 
-    # Set the reshape method corresponding to projections or slices
+    # Set the reshape method corresponding to projections or slices.
     if reshape_method == "slice":
       reshape_fn = partial(
           slice_hidden, hidden_size=hidden_size, num_blocks=num_blocks)
@@ -525,7 +575,7 @@ def discrete_bottleneck(x,
     else:
       raise ValueError("Unknown reshape_method")
 
-    # Check if the ema settings make sense
+    # Check if the ema settings make sense.
     if ema:
       if ema_count is None:
         raise ValueError("ema_count is None but ema is True")
@@ -534,6 +584,7 @@ def discrete_bottleneck(x,
 
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     l = tf.constant(0.0)
+    neg_q_entropy = tf.constant(0.0)
     if bottleneck_kind == "dense":
       c = tf.layers.dense(x, z_size, name="vcc")
       h1 = tf.layers.dense(c, filter_size, name="vch1")
@@ -569,16 +620,20 @@ def discrete_bottleneck(x,
                                  kl_warmup_steps, summary)
       c = tf.argmax(hot, axis=-1)
       h1 = tf.layers.dense(hot, hidden_size, name="dae_dense")
-    elif bottleneck_kind == "dvq":
+    elif bottleneck_kind in ["dvq", "gumbel-softmax-dvq"]:
       x_reshaped = reshape_fn(x)
-      x_res = x_reshaped
+      x_res = tf.reshape(x_reshaped,
+                         [-1] + common_layers.shape_list(x_reshaped)[2:])
       x_means_hot = []
       x_means = 0
       l = 0
       for i in range(num_residuals):
-        x_means_hot_res, x_means_res, q_loss_res, e_loss_res = embedding_lookup(
-            x_res, means[i], num_blocks, block_v_size, random_top_k, soft_em,
-            num_samples)
+        x_means_hot_res, x_means_res, q_loss_res, e_loss_res, neg_q_entropy = (
+            embedding_lookup(x_reshaped, means[i], num_blocks, block_v_size,
+                             bottleneck_kind, random_top_k, soft_em,
+                             num_samples, do_hard_gumbel_softmax,
+                             do_iaf=do_iaf,
+                             approximate_gs_entropy=approximate_gs_entropy))
         # Update the ema variables
         if ema:
           tf.logging.info("Using EMA with beta = {}".format(beta))
@@ -663,7 +718,7 @@ def discrete_bottleneck(x,
         num_residuals=num_residuals,
         block_v_size=block_v_size,
         means=means)
-    return res, c, l, embed_fn
+    return res, c, l, embed_fn, neg_q_entropy
 
 
 # New API for discretization bottlenecks:
@@ -766,6 +821,191 @@ def vq_discrete_unbottleneck(x, hidden_size):
   return tf.reshape(result, x_shape[:-1] + [hidden_size])
 
 
+def gumbel_softmax_nearest_neighbor_dvq(x,
+                                        means,
+                                        block_v_size,
+                                        hard=False,
+                                        temperature_init=1.2,
+                                        num_samples=1,
+                                        startup_steps=15000,
+                                        summary=True,
+                                        do_iaf=False,
+                                        approximate_gs_entropy=False):
+  """Sample from Gumbel-Softmax and compute neighbors and losses.
+
+  Args:
+    x: A `float`-like `Tensor` of shape [batch_size, latent_dim, num_blocks,
+      block_dim] containing the latent vectors to be compared to the codebook.
+    means: Embedding table of shape [num_blocks, block_v_size, block_dim].
+    block_v_size: Number of discrete codes per block.
+    hard: Determines whether we take hard or soft Gumbel-Softmax samples
+      (Default: False).
+    temperature_init: Initial temperature used for Gumbel-Softmax samples,
+      after it which it decays to 0 (Default: 1.2).
+    num_samples: Number of samples drawn for each latent (Default: 1).
+    startup_steps: Number of steps it takes to decay temperature to 0 (Default:
+      15000).
+    summary: When `True`, we save histogram summaries of the KL term (Default:
+      True).
+    do_iaf: When `True`, we perform inverse autoregressive flow with
+      Gumbel-Softmax sample (Default: False).
+    approximate_gs_entropy: When `True`, we approximate Gumbel-Softmax
+      density as categorical when calculating sample entropy (Default: False).
+
+  Returns:
+    x_means_assignments: A `float`-like `Tensor` containing the codebook
+      assignments, averaged over samples, with shape [batch_size * latent_dim,
+      num_blocks, block_v_size].
+    neg_q_entropy: The negative entropy of the variational distribution,
+      averaged over samples.
+  """
+  batch_size, latent_dim, num_blocks, block_dim = common_layers.shape_list(x)
+
+  # Combine latent_dim and batch_size for computing distances.
+  x = tf.reshape(x, [-1, num_blocks, block_dim])
+
+  # Compute distances using (x - means)**2 = x**2 + means**2 - 2*x*means.
+  x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keepdims=True)
+  means_norm_sq = tf.reduce_sum(tf.square(means), axis=-1, keepdims=True)
+  means_norm_sq = tf.transpose(means_norm_sq, perm=[2, 0, 1])
+  scalar_prod = tf.matmul(
+      tf.transpose(x, perm=[1, 0, 2]), tf.transpose(means, perm=[0, 2, 1]))
+  scalar_prod = tf.transpose(scalar_prod, perm=[1, 0, 2])
+  dist = x_norm_sq + means_norm_sq - 2 * scalar_prod
+
+  # IAF requires latents to have their own dimension, so reshape dist from
+  # [batch_size * latent_dim, num_blocks, block_v_size] to
+  # [batch_size * num_blocks, latent_dim, block_v_size].
+  dist = tf.reshape(dist, [batch_size, latent_dim, num_blocks, -1])
+  dist = tf.reshape(tf.transpose(dist, perm=[0, 2, 1, 3]),
+                    [-1, latent_dim, block_v_size])
+  log_class_probs = tf.nn.log_softmax(-dist)
+
+  sample_shape = [num_samples] + common_layers.shape_list(dist)
+  gumbel_samples = gumbel_sample(sample_shape)
+
+  # Temperature decays linearly.
+  temperature = temperature_init - common_layers.inverse_lin_decay(
+      startup_steps)
+
+  # 10% of the time keep reasonably high temperature to keep learning.
+  temperature = tf.cond(
+      tf.less(tf.random_uniform([]), 0.9), lambda: temperature,
+      lambda: tf.random_uniform([], minval=0.5, maxval=1.0))
+
+  gumbel_softmax_samples = tf.nn.softmax(
+      (tf.expand_dims(log_class_probs, 0) + gumbel_samples) / temperature)
+  q_samples = tf.clip_by_value(gumbel_softmax_samples, 1e-6, 1-1e-6)
+
+  if approximate_gs_entropy:
+    q_dist = tf.contrib.distributions.Multinomial(total_count=1.0, logits=-dist)
+  else:
+    q_dist = tf.contrib.distributions.RelaxedOneHotCategorical(
+        temperature, logits=-dist)
+
+  # Take mean over samples to approximate entropy.
+  neg_q_entropy = tf.reduce_mean(q_dist.log_prob(q_samples), 0)
+
+  if summary:
+    tf.summary.histogram("neg_q_entropy", tf.reshape(neg_q_entropy, [-1]))
+  neg_q_entropy = tf.reduce_mean(neg_q_entropy)
+
+  if do_iaf:
+    q_samples = tf.reshape(q_samples, [-1, latent_dim, block_v_size])
+
+    # Shift samples so log_pi[:, i, :] is only a function of
+    # q_samples[:, :i, :]. We do this by adding a first row of zeros to the
+    # latents, shifting the other rows down by one, and removing the last row.
+
+    top_latent = tf.zeros([batch_size * num_blocks, 1, block_v_size])
+    shifted_samples = tf.concat([top_latent, q_samples[:, :-1, :]], axis=1)
+
+    d_k = 64
+    d_v = 64
+    query_projection = tf.get_variable(
+        "query_projection", [block_v_size, d_k], dtype=tf.float32)
+    keys_projection = tf.get_variable(
+        "keys_projection", [block_v_size, d_k], dtype=tf.float32)
+    values_projection = tf.get_variable(
+        "values_projection", [block_v_size, d_v], dtype=tf.float32)
+    query = tf.reduce_sum(tf.expand_dims(shifted_samples, -1) *
+                          tf.reshape(query_projection,
+                                     [1, 1, block_v_size, d_k]), 2)
+    keys = tf.reduce_sum(tf.expand_dims(shifted_samples, -1) *
+                         tf.reshape(keys_projection,
+                                    [1, 1, block_v_size, d_k]), 2)
+    values = tf.reduce_sum(tf.expand_dims(shifted_samples, -1) *
+                           tf.reshape(values_projection,
+                                      [1, 1, block_v_size, d_v]), 2)
+
+    # Masked self-attention with a single head.
+    # TODO(vafa): Add support for multiple heads
+    attention_output = common_attention.masked_local_attention_1d(
+        q=tf.expand_dims(query, 1),
+        k=tf.expand_dims(keys, 1),
+        v=tf.expand_dims(values, 1),
+        block_length=1)
+    attention_output = tf.reshape(
+        attention_output, [-1] + common_layers.shape_list(attention_output)[2:])
+
+    ffn_output = common_layers.conv_relu_conv(
+        attention_output,
+        filter_size=64,
+        output_size=block_v_size,
+        first_kernel_size=3,
+        second_kernel_size=1,
+        padding="LEFT",
+        nonpadding_mask=None,
+        dropout=0.,
+        cache=None,
+        decode_loop_step=None)
+
+    log_pi = tf.nn.log_softmax(ffn_output)
+
+    # Flow 1: Adding log_pi to q_samples and dividing by the temperature.
+    # Note that we drop last dimension of q_samples for centered-softmax, which
+    # we can do without recalculating probabilities because the last dimension
+    # of log_pi and q_samples are deterministic given the other dimensions.
+    # Flow 2: Centered-softmax.
+
+    chained_bijectors = tf.contrib.distributions.bijectors.Chain(
+        [tf.contrib.distributions.bijectors.SoftmaxCentered(),
+         tf.contrib.distributions.bijectors.Affine(
+             shift=log_pi[:, :, :-1],
+             scale_identity_multiplier=1./temperature)])
+    q_samples = chained_bijectors.forward(q_samples[:, :, :-1])
+    neg_q_entropy += tf.reduce_mean(
+        chained_bijectors.inverse_log_det_jacobian(q_samples, event_ndims=1))
+
+    q_samples = tf.reshape(
+        q_samples,
+        [num_samples, batch_size * num_blocks, latent_dim, block_v_size])
+
+  if hard:
+    x_means_idx = tf.argmax(q_samples, -1)
+
+    # Take average of one-hot vectors over samples.
+    x_means_hot = tf.reduce_mean(tf.one_hot(x_means_idx, block_v_size), 0)
+    x_means_assignments = (tf.reduce_mean(q_samples, 0) +
+                           tf.stop_gradient(x_means_hot - tf.reduce_mean(
+                               q_samples, 0)))
+  else:
+    x_means_assignments = tf.reduce_mean(gumbel_softmax_samples, 0)
+
+  # Reshape assignments to [batch_size * latent_dim, num_blocks,
+  # block_v_size]. We have to transpose between reshapes to make sure the
+  # dimensions have the correct interpretation.
+  x_means_assignments = tf.reshape(
+      x_means_assignments,
+      [batch_size, num_blocks, latent_dim, block_v_size])
+  x_means_assignments = tf.transpose(x_means_assignments, [0, 2, 1, 3])
+  x_means_assignments = tf.reshape(
+      x_means_assignments,
+      [batch_size * latent_dim, num_blocks, block_v_size])
+
+  return x_means_assignments, neg_q_entropy
+
+
 def gumbel_softmax_discrete_bottleneck(x,
                                        bottleneck_bits,
                                        beta=0.25,
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index b017b6582..384fc5008 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -90,7 +90,7 @@ def testSliceHiddenZeros(self):
     with self.test_session() as sess:
       tf.global_variables_initializer().run()
       x_sliced_eval = sess.run(x_sliced)
-      self.assertEqual(np.shape(x_sliced_eval), (1, num_blocks, block_dim))
+      self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
       self.assertTrue(np.all(x_sliced_eval == 0))
 
   def testSliceHiddenOnes(self):
@@ -102,16 +102,16 @@ def testSliceHiddenOnes(self):
     with self.test_session() as sess:
       tf.global_variables_initializer().run()
       x_sliced_eval = sess.run(x_sliced)
-      self.assertEqual(np.shape(x_sliced_eval), (1, num_blocks, block_dim))
+      self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
       self.assertTrue(np.all(x_sliced_eval == 1))
 
   def testNearestNeighbors(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
-    x = tf.expand_dims(x, axis=0)
+    x = tf.reshape(x, [1, 1, 2, 3])
     means = tf.constant(
         [[1, 0, 0], [0, 1, 0], [0, 0, 1], [9, 9, 9]], dtype=tf.float32)
     means = tf.stack([means, means], axis=0)
-    x_means_hot = discretization.nearest_neighbor(x, means, block_v_size=4)
+    x_means_hot, _ = discretization.nearest_neighbor(x, means, block_v_size=4)
     x_means_hot_test = np.array([[0, 1, 0, 0], [1, 0, 0, 0]])
     x_means_hot_test = np.expand_dims(x_means_hot_test, axis=0)
     with self.test_session() as sess:
@@ -163,7 +163,7 @@ def testVQDiscreteUnbottlenck(self):
       x_means_eval = sess.run(x_means)
       self.assertEqual(np.shape(x_means_eval), (2, 3))
 
-  def testGumbleSoftmaxDiscreteBottleneck(self):
+  def testGumbelSoftmaxDiscreteBottleneck(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
     tf.add_to_collection(tf.GraphKeys.GLOBAL_STEP, tf.constant(1))
     x_means_hot, _ = discretization.gumbel_softmax_discrete_bottleneck(
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 5300968cf..8c27730fe 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -346,7 +346,8 @@ def ae_transformer_internal(inputs,
     ed, inputs_ex, ed_ex = None, None, None
 
   # Autoencoding.
-  losses = {"extra": tf.constant(0.0), "latent_pred": tf.constant(0.0)}
+  losses = {"extra": tf.constant(0.0), "latent_pred": tf.constant(0.0),
+            "neg_q_entropy": tf.constant(0.0)}
   if hparams.do_ae:
     # flatten here
     original_targets_shape = tf.shape(targets)
@@ -370,11 +371,12 @@ def ae_transformer_internal(inputs,
     targets_c = compress(targets_noisy, inputs, False, hparams, "compress")
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
       # Compress and bottleneck.
-      latents_dense, latents_discrete, extra_loss, embed = hparams.bottleneck(
-          x=targets_c,
-          filter_size=hparams.compress_filter_size,
-          name="vc",
-          mode=hparams.mode)
+      latents_dense, latents_discrete, extra_loss, embed, neg_q_entropy = (
+          hparams.bottleneck(
+              x=targets_c,
+              filter_size=hparams.compress_filter_size,
+              name="vc",
+              mode=hparams.mode))
       if _DO_SUMMARIES:
         tf.summary.histogram("b0", tf.reshape(latents_discrete[:, 0, :], [-1]))
       pc = common_layers.inverse_exp_decay(hparams.startup_steps)
@@ -392,7 +394,8 @@ def ae_transformer_internal(inputs,
         _, latent_pred_loss = ae_latent_softmax(
             latents_pred, tf.stop_gradient(latents_discrete), hparams)
         losses["latent_pred"] = tf.reduce_mean(
-            latent_pred_loss * tf.to_float(cond))
+            latent_pred_loss * tf.to_float(cond)) * hparams.prior_scale
+        losses["neg_q_entropy"] = neg_q_entropy * hparams.entropy_scale
       else:
         inputs_c = decode_transformer(inputs, ed, targets_c, hparams, "dec_c")
         losses["latent_pred"] = tf.reduce_mean((inputs_c - targets_c)**2) * 20
@@ -513,9 +516,12 @@ def __init__(self, *args, **kwargs):
         softmax_k=self._hparams.softmax_k,
         kl_warmup_steps=self._hparams.kl_warmup_steps,
         ema=self._hparams.ema,
-        summary=_DO_SUMMARIES)
+        summary=_DO_SUMMARIES,
+        do_hard_gumbel_softmax=self._hparams.do_hard_gumbel_softmax,
+        do_iaf=self._hparams.do_iaf,
+        approximate_gs_entropy=self._hparams.approximate_gs_entropy)
     # Set the discretization bottleneck specific things here
-    if self._hparams.bottleneck_kind == "dvq":
+    if self._hparams.bottleneck_kind in ["dvq", "gumbel-softmax-dvq"]:
       z_size_per_residual = self._hparams.z_size / self._hparams.num_residuals
       block_dim = int(self._hparams.hidden_size // self._hparams.num_blocks)
       block_v_size = 2**(z_size_per_residual / self._hparams.num_blocks)
@@ -714,6 +720,11 @@ def transformer_ae_small():
   hparams.add_hparam("soft_em", False)
   hparams.add_hparam("num_samples", 10)
   hparams.add_hparam("inv_temp", 1.0)
+  hparams.add_hparam("entropy_scale", 0.0)
+  hparams.add_hparam("prior_scale", 1.0)
+  hparams.add_hparam("do_hard_gumbel_softmax", False)
+  hparams.add_hparam("do_iaf", False)
+  hparams.add_hparam("approximate_gs_entropy", False)
   hparams.kl_warmup_steps = 150000
   hparams.force_full_predict = True
 
@@ -840,3 +851,65 @@ def transformer_ae_base_tpu():
   transformer.update_hparams_for_tpu(hparams)
   hparams.batch_size = 512
   return hparams
+
+
+@registry.register_hparams
+def transformer_ae_base_noatt():
+  """Set of hyperparameters."""
+  hparams = transformer_ae_base()
+  hparams.reshape_method = "slice"
+  hparams.bottleneck_kind = "dvq"
+  hparams.hidden_size = 512
+  hparams.num_blocks = 1
+  hparams.num_decode_blocks = 1
+  hparams.z_size = 12
+  hparams.do_attend_decompress = False
+  return hparams
+
+
+@registry.register_hparams
+def transformer_ae_base_ablation_1():
+  hparams = transformer_ae_base_noatt()
+  hparams.soft_em = True
+  return hparams
+
+
+@registry.register_hparams
+def transformer_ae_base_ablation_2():
+  hparams = transformer_ae_base_ablation_1()
+  hparams.entropy_scale = 0.1
+  return hparams
+
+
+@registry.register_hparams
+def transformer_ae_base_ablation_3():
+  hparams = transformer_ae_base_ablation_2()
+  hparams.prior_scale = 0.1
+  hparams.entropy_scale = 0.1
+  return hparams
+
+
+@registry.register_hparams
+def transformer_ae_base_ablation_4():
+  hparams = transformer_ae_base_ablation_3()
+  hparams.entropy_scale = 0.0
+  hparams.prior_scale = 1.0
+  hparams.bottleneck_kind = "gumbel-softmax-dvq"
+  hparams.do_hard_gumbel_softmax = True
+  hparams.approximate_gs_entropy = True
+  return hparams
+
+
+@registry.register_hparams
+def transformer_ae_base_ablation_5():
+  hparams = transformer_ae_base_ablation_4()
+  hparams.do_hard_gumbel_softmax = False
+  return hparams
+
+
+@registry.register_hparams
+def transformer_ae_base_iaf():
+  hparams = transformer_ae_base_ablation_5()
+  hparams.do_iaf = True
+  hparams.num_samples = 1
+  return hparams

From a0addd067cef1ceda5dd8936c4d8cf5a05156b25 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Fri, 20 Jul 2018 14:43:00 -0700
Subject: [PATCH 0360/2720] Turn off label smoothing for eval.

PiperOrigin-RevId: 205456357
---
 tensor2tensor/utils/t2t_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 3be734b0f..e59ec2373 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -472,7 +472,7 @@ def set_mode(self, mode):
     # When not in training mode, set all forms of dropout to zero.
     if mode != tf.estimator.ModeKeys.TRAIN:
       for key in hparams.values():
-        if key.endswith("dropout"):
+        if key.endswith("dropout") or key == "label_smoothing":
           log_info("Setting hparams.%s to 0.0", key)
           setattr(hparams, key, 0.0)
     self._hparams = hparams

From b9a88027413f858f396c31fb64bc07d86a8e6fa3 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 20 Jul 2018 18:12:29 -0700
Subject: [PATCH 0361/2720] Fixing the frame-reward shift bug.

PiperOrigin-RevId: 205482096
---
 tensor2tensor/rl/collect.py | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index f2ce0301c..69776f8b1 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -55,11 +55,6 @@ def __init__(self, batch_env):
     shapes = meta_data[0][:4]
     dtypes = meta_data[1][:4]
     self.speculum = tf.FIFOQueue(infinity, shapes=shapes, dtypes=dtypes)
-    observs_shape = batch_env.observ.shape
-    # TODO(piotrmilos): possibly retrieve the observation type for batch_env
-    observ_dtype = tf.float32
-    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
-                               trainable=False)
 
   def simulate(self, action):
 
@@ -71,12 +66,9 @@ def simulate(self, action):
 
     with tf.control_dependencies([reward, done]):
       enqueue_op = self.speculum.enqueue(
-          [self._observ.read_value(), reward, done, action])
+          [self._batch_env.observ, reward, done, action])
 
     with tf.control_dependencies([enqueue_op]):
-      assign = self._observ.assign(self._batch_env.observ)
-
-    with tf.control_dependencies([assign]):
       return tf.identity(reward), tf.identity(done)
 
 

From d26796145ea37c02c75675d839182994da5434d2 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 20 Jul 2018 18:33:50 -0700
Subject: [PATCH 0362/2720] Log the number of ground_truth frames fed at every
 step.

PiperOrigin-RevId: 205483551
---
 tensor2tensor/models/research/next_frame.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 24c568c71..e3cd9070a 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -677,6 +677,7 @@ def get_scheduled_sample_inputs(
             tf.round(
                 tf.to_float(batch_size) *
                 (k / (k + tf.exp(tf.to_float(iter_num) / tf.to_float(k))))))
+        tf.summary.scalar("num_ground_truth", num_ground_truth)
 
       if feedself and done_warm_start:
         # Feed in generated stuff.

From 33b6e4382e19fa22eb7594c05f492bf1144d652f Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Sun, 22 Jul 2018 13:53:02 -0700
Subject: [PATCH 0363/2720] Internal change

PiperOrigin-RevId: 205589211
---
 .../models/research/transformer_vae.py          | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 8c27730fe..d2b12c0a9 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -190,10 +190,10 @@ def decode_transformer(encoder_output,
             [common_layers.shape_list(targets)[0], 1, 1, hparams.hidden_size])
       decoder_output = cia.transformer_decoder_layers(
           decoder_input,
-          None,
-          bias,
-          hparams.num_decoder_layers or hparams.num_hidden_layers,
-          hparams,
+          encoder_output=None,
+          num_layers=hparams.num_decoder_layers or hparams.num_hidden_layers,
+          hparams=hparams,
+          self_attention_bias=bias,
           attention_type=hparams.dec_attention_type,
           name="decoder")
     decoder_output_shape = common_layers.shape_list(decoder_output)
@@ -573,7 +573,10 @@ def __init__(self, *args, **kwargs):
           for i in range(self._hparams.num_residuals):
             ema_means_i = tf.get_variable(
                 "ema_means_{}".format(i),
-                initializer=means.initialized_value()[i],
+                [self._hparams.num_blocks, block_v_size, block_dim],
+                initializer=(lambda shape, dtype=None, partition_info=None,  # pylint: disable=g-long-lambda
+                                    verify_shape=None:
+                             means.initialized_value()[i]),
                 trainable=False)
             ema_means.append(ema_means_i)
 
@@ -770,6 +773,9 @@ def imagetransformer_ae_cifar():
   hparams.pos = "timing"  # timing, none
   hparams.nbr_decoder_problems = 1
   hparams.num_output_layers = 3
+  # TODO(trandustin): semhash doesn't work if filter_size != hidden_size. For
+  # now, set default to dvq.
+  hparams.bottleneck_kind = "dvq"
   hparams.add_hparam("block_size", 1)
 
   # dilated attention based flags
@@ -789,6 +795,7 @@ def imagetransformer_ae_cifar():
   hparams.sep_rgb_embed = False
   hparams.add_hparam("dec_attention_type", cia.AttentionType.LOCAL_1D)
   hparams.add_hparam("block_raster_scan", False)
+  hparams.add_hparam("shared_rel", False)
 
   # multipos attention params
   hparams.add_hparam("q_filter_width", 1)

From bc6eb56b89dcd84f1ee25bb276ddb44fc20e9aec Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Sun, 22 Jul 2018 17:57:59 -0700
Subject: [PATCH 0364/2720] Extend TransformerAE to work with >1000 latent
 variables.

PiperOrigin-RevId: 205599118
---
 .../models/research/transformer_vae.py         | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index d2b12c0a9..0e9a72e91 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -431,8 +431,10 @@ def bn_inputs():
         latents_dense = embed(cache)
     # Postprocess.
     d = latents_dense
-    pos = tf.get_variable("pos", [1, 1000, 1, hparams.hidden_size])
-    pos = pos[:, :common_layers.shape_list(latents_dense)[1] + 1, :, :]
+    latent_len = common_layers.shape_list(latents_dense)[1]
+    if isinstance(latent_len, tf.Tensor):
+      latent_len = hparams.max_length
+    pos = tf.get_variable("pos", [1, latent_len + 1, 1, hparams.hidden_size])
     latents_dense = tf.pad(latents_dense,
                            [[0, 0], [1, 0], [0, 0], [0, 0]]) + pos
 
@@ -810,6 +812,18 @@ def imagetransformer_ae_cifar():
   return hparams
 
 
+def imagetransformer_ae_imagenet():
+  """For 64x64 ImageNet. ~56M trainable variables."""
+  hparams = imagetransformer_ae_cifar()
+  hparams.max_length = int(64 * 64 * 3)
+  hparams.img_len = 64
+  hparams.num_heads = 4  # Heads are expensive on TPUs.
+  # Reduce architecture from 32x32 CIFAR-10 in order to fit in memory.
+  hparams.num_decoder_layers = 8
+  hparams.num_compress_steps = 2
+  return hparams
+
+
 @registry.register_hparams
 def transformer_ae_base():
   """Set of hyperparameters."""

From 3f43cd235708ca6445a98fd67e49f5dbb96d56c5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 23 Jul 2018 09:13:48 -0700
Subject: [PATCH 0365/2720] Add test-split to CelebA dataset creation and sort
 filenames to ensure correct split.

PiperOrigin-RevId: 205671298
---
 tensor2tensor/data_generators/celeba.py | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index 476e8f04c..cf068f040 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -114,7 +114,7 @@ def process_attrs(raw_data):
     img_landmarks, _ = process_landmarks(landmarks_raw)
     img_attrs, _ = process_attrs(attr_raw)
 
-    image_files = tf.gfile.Glob(unzipped_folder + "/*.jpg")
+    image_files = list(sorted(tf.gfile.Glob(unzipped_folder + "/*.jpg")))
     for filename in image_files[start_from:start_from + how_many]:
       img_name = os.path.basename(filename)
       landmarks = img_landmarks[img_name]
@@ -137,12 +137,25 @@ def train_shards(self):
   def dev_shards(self):
     return 10
 
+  @property
+  def test_shards(self):
+    return 10
+
   def generate_data(self, data_dir, tmp_dir, task_id=-1):
-    generator_utils.generate_dataset_and_shuffle(
-        self.generator(tmp_dir, 162770),  # train
-        self.training_filepaths(data_dir, self.train_shards, shuffled=False),
-        self.generator(tmp_dir, 19867, 162770),  # dev
-        self.dev_filepaths(data_dir, self.dev_shards, shuffled=False))
+    train_gen = self.generator(tmp_dir, 162770)
+    train_paths = self.training_filepaths(
+        data_dir, self.train_shards, shuffled=False)
+    generator_utils.generate_files(train_gen, train_paths)
+
+    dev_gen = self.generator(tmp_dir, 19867, 162770)
+    dev_paths = self.dev_filepaths(data_dir, self.dev_shards, shuffled=False)
+    generator_utils.generate_files(dev_gen, dev_paths)
+
+    test_gen = self.generator(tmp_dir, 19962, 162770+19867)
+    test_paths = self.test_filepaths(data_dir, self.test_shards, shuffled=False)
+    generator_utils.generate_files(test_gen, test_paths)
+
+    generator_utils.shuffle_dataset(train_paths + dev_paths + test_paths)
 
 
 @registry.register_problem

From f6992c34ff8152e4dbbfc27a08f80a7e6af4b555 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 23 Jul 2018 10:11:11 -0700
Subject: [PATCH 0366/2720] Allow option to set both
 save_checkpoints_{steps,secs} to None.

PiperOrigin-RevId: 205680280
---
 tensor2tensor/utils/trainer_lib.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 64e5fa7f9..d00056b36 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -143,6 +143,7 @@ def create_run_config(master="",
       "session_config": session_config,
       "save_summary_steps": 100,
       "save_checkpoints_steps": save_checkpoints_steps,
+      "save_checkpoints_secs": save_checkpoints_secs,
       "keep_checkpoint_max": keep_checkpoint_max,
       "keep_checkpoint_every_n_hours": keep_checkpoint_every_n_hours,
       "tf_random_seed": random_seed,
@@ -150,7 +151,6 @@ def create_run_config(master="",
   }
   if save_checkpoints_secs:
     del run_config_args["save_checkpoints_steps"]
-    run_config_args["save_checkpoints_secs"] = save_checkpoints_secs
   run_config_cls = tf.contrib.learn.RunConfig
 
   # If using TPU, use TPU RunConfig, add TPUConfig, and add additional args

From 2ee73083ce8a630c19f95fc713ec2b7d1409f1e7 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 23 Jul 2018 10:33:17 -0700
Subject: [PATCH 0367/2720] Clean up docs, logic, and var names in
 discrete_bottleneck.

PiperOrigin-RevId: 205684438
---
 tensor2tensor/layers/discretization.py        | 475 ++++++++++--------
 .../models/research/transformer_vae.py        |  49 +-
 2 files changed, 285 insertions(+), 239 deletions(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index da90b07af..5426e7974 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -75,8 +75,8 @@ def nearest_neighbor(x,
   """Find the nearest element in means to elements in x.
 
   Args:
-    x: Batch of encoder continuous latent states sliced/projected into shape
-      [batch_size, latent_dim, num_blocks, block_dim].
+    x: Continuous encodings of shape [batch_size, latent_dim, num_blocks,
+      block_dim].
     means: Embedding table of shape [num_blocks, block_v_size, block_dim].
     block_v_size: Number of table entries per block.
     random_top_k: Noisy top-k if this is bigger than 1 (Default: 1).
@@ -137,37 +137,40 @@ def embedding_lookup(x,
                      soft_em=False,
                      num_samples=1,
                      do_hard_gumbel_softmax=False,
+                     temperature_warmup_steps=150000,
                      do_iaf=False,
                      approximate_gs_entropy=False):
   """Compute nearest neighbors and loss for training the embeddings via DVQ.
 
   Args:
-    x: Batch of encoder continuous latent states sliced/projected into shape
-      [batch_size, latent_dim, num_blocks, block_dim].
+    x: Continuous encodings of shape [batch_size, latent_dim, num_blocks,
+      block_dim].
     means: Embedding table of shape [num_blocks, block_v_size, block_dim].
     num_blocks: Number of blocks in DVQ.
     block_v_size: Number of table entries per block.
-    bottleneck_kind: Discrete bottleneck type (Default: "dvq").
-    random_top_k: Noisy top-k if this is bigger than 1 (Default: 1).
-    soft_em: If True then use soft EM rather than hard EM (Default: False).
-    num_samples: Number of samples to use for soft EM (Default: 1).
-    do_hard_gumbel_softmax: Boolean determining whether we take hard or soft
-      Gumbel-Softmax samples for gumbel-softmax-dvq bottleneck (Default: False).
-    do_iaf: Boolean determining whether we use inverse autoregressive flows for
-      gumbel-softmax-dvq bottleneck (Default: False).
-    approximate_gs_entropy: If True, use the categorical density instead of the
-      Gumbel-Softmax density when calculating the sample entropy (Default:
-      False).
+    bottleneck_kind: Discrete bottleneck type.
+    random_top_k: Noisy top-k if this is bigger than 1.
+    soft_em: If True then use soft EM rather than hard EM.
+    num_samples: Number of samples to use for soft EM.
+    do_hard_gumbel_softmax: Whether to use hard or soft Gumbel-Softmax samples
+      for gumbel-softmax-dvq bottleneck.
+    temperature_warmup_steps: Number of steps it takes to decay temperature to
+      0. Used only if bottleneck_kind is gumbel-softmax-dvq.
+    do_iaf: Whether to apply inverse autoregressive flows for gumbel-softmax-dvq
+      bottleneck.
+    approximate_gs_entropy: Whether to approximate the Gumbel-Softmax density
+      as a categorical distribution when calculating the sample entropy. Used
+      only if bottleneck_kind is gumbel-softmax-dvq.
 
   Returns:
     x_means_hot: The nearest neighbor in one hot form, with shape
       [batch_size * latent_dim, num_blocks, block_v_size].
     x_means: The nearest neighbor itself, with shape [batch_size * latent_dim,
       num_blocks, block_dim].
-    q_loss: Commitment loss.
-    e_loss: Embedding training loss.
-    neg_q_entropy: Negative entropy of variational approximation (0 if q is
-      deterministic).
+    q_loss: Scalar Tensor representing codebook loss.
+    e_loss: Scalar Tensor representing commitment loss.
+    neg_q_entropy: Scalar Tensor representing negative entropy of variational
+      approximation (0 if it is deterministic).
   """
   if bottleneck_kind == "gumbel-softmax-dvq":
     x_means_hot, neg_q_entropy = gumbel_softmax_nearest_neighbor_dvq(
@@ -176,6 +179,7 @@ def embedding_lookup(x,
         block_v_size,
         hard=do_hard_gumbel_softmax,
         num_samples=num_samples,
+        temperature_warmup_steps=temperature_warmup_steps,
         do_iaf=do_iaf,
         approximate_gs_entropy=approximate_gs_entropy)
   else:
@@ -250,22 +254,21 @@ def embed(x,
           hidden_size,
           z_size,
           filter_size,
-          name,
           bottleneck_kind="dvq",
           soft_em=False,
           num_blocks=2,
           num_residuals=1,
           block_v_size=None,
-          means=None):
+          means=None,
+          name=None):
   """Embedding function that takes discrete latent and returns embedding.
 
   Args:
     x: Input to the discretization bottleneck.
     hidden_size: Dimension of the latent state.
-    z_size: Number of bits used to produce discrete code; discrete codes range
-      from 1 to 2**z_size.
-    filter_size: Filter size to be used for the embedding function.
-    name: Name for the bottleneck scope.
+    z_size: Number of bits, where discrete codes range from 1 to 2**z_size.
+    filter_size: Dimension to project embedding by. Used only if bottleneck_kind
+      is semhash.
     bottleneck_kind: Kind of discretization bottleneck to use; one of dvq,
       semhash, gumbel-softmax (Default: dvq).
     soft_em: If True then it uses a multi-sample version of EM (Default: False).
@@ -273,6 +276,7 @@ def embed(x,
     num_residuals: Number of residuals (Default: 1).
     block_v_size: Number of embedding entries per block (Default: None).
     means: The embedding table for dvq (Default: None).
+    name: Name for the bottleneck scope.
 
   Returns:
     Continuous embedding to be passed on to the decoder.
@@ -280,7 +284,7 @@ def embed(x,
   Raises:
     ValueError: For unknown or missing arguments.
   """
-  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+  with tf.variable_scope(name, default_name="embed", reuse=tf.AUTO_REUSE):
     if bottleneck_kind == "semhash":
       c = int_to_bit(x, z_size)
       h1a = tf.layers.dense(c, filter_size, name="vch1a")
@@ -334,19 +338,18 @@ def embed(x,
     return h1
 
 
-def vae(x, name, z_size):
+def vae(x, z_size, name=None):
   """Simple variational autoencoder without discretization.
 
   Args:
     x: Input to the discretization bottleneck.
+    z_size: Number of bits, where discrete codes range from 1 to 2**z_size.
     name: Name for the bottleneck scope.
-    z_size: Number of bits used to produce discrete code; discrete codes range
-      from 1 to 2**z_size.
 
   Returns:
     Embedding function, latent, loss, mu and log_simga.
   """
-  with tf.variable_scope(name):
+  with tf.variable_scope(name, default_name="vae"):
     mu = tf.layers.dense(x, z_size, name="mu")
     log_sigma = tf.layers.dense(x, z_size, name="log_sigma")
     shape = common_layers.shape_list(x)
@@ -391,29 +394,28 @@ def gumbel_sample(shape):
 
 
 def gumbel_softmax(x,
-                   name,
                    z_size,
                    mode,
                    softmax_k=0,
-                   kl_warmup_steps=150000,
-                   summary=True):
+                   temperature_warmup_steps=150000,
+                   summary=True,
+                   name=None):
   """Gumbel softmax discretization bottleneck.
 
   Args:
     x: Input to the discretization bottleneck.
+    z_size: Number of bits, where discrete codes range from 1 to 2**z_size.
+    mode: tf.estimator.ModeKeys.
+    softmax_k: If > 0 then do top-k softmax.
+    temperature_warmup_steps: Number of steps it takes to decay temperature to
+      0.
+    summary: Whether to write summaries.
     name: Name for the bottleneck scope.
-    z_size: Number of bits used to produce discrete code; discrete codes range
-      from 1 to 2**z_size.
-    mode: Mode represents whether we are training or testing for bottlenecks
-      that differ in behavior (Default: None).
-    softmax_k: If > 1 then do top-k softmax (Default: 0).
-    kl_warmup_steps: Number of steps for kl warmup (Default: 150000).
-    summary: If True, then write summaries (Default: True).
 
   Returns:
-    Embedding function, discrete code and loss.
+    Embedding function, discrete code, and loss.
   """
-  with tf.variable_scope(name):
+  with tf.variable_scope(name, default_name="gumbel_softmax"):
     m = tf.layers.dense(x, 2**z_size, name="mask")
     if softmax_k > 0:
       m, kl = top_k_softmax(m, softmax_k)
@@ -422,7 +424,7 @@ def gumbel_softmax(x,
 
     # Gumbel-softmax sample.
     gumbel_samples = gumbel_sample(common_layers.shape_list(m))
-    steps = kl_warmup_steps
+    steps = temperature_warmup_steps
     gumbel_samples *= common_layers.inverse_exp_decay(steps // 5) * 0.5
     temperature = 1.2 - common_layers.inverse_lin_decay(steps)
 
@@ -453,188 +455,187 @@ def gumbel_softmax(x,
     return m, ret, d_dev * 5.0 + tf.reduce_mean(kl) * 0.002
 
 
-def discrete_bottleneck(x,
+def discrete_bottleneck(inputs,
                         hidden_size,
                         z_size,
                         filter_size,
-                        name,
                         mode=None,
-                        startup_steps=50000,
                         bottleneck_kind="dvq",
                         num_blocks=2,
                         num_residuals=1,
                         reshape_method="slice",
                         projection_tensors=None,
-                        means=None,
                         beta=0.25,
-                        noise_dev=1.,
+                        ema=True,
+                        means=None,
+                        ema_count=None,
+                        ema_means=None,
+                        epsilon=1e-5,
                         decay=0.999,
-                        discrete_mix=0.5,
                         random_top_k=1,
                         soft_em=False,
                         num_samples=1,
-                        epsilon=1e-5,
                         softmax_k=0,
-                        kl_warmup_steps=150000,
-                        ema=True,
-                        ema_count=None,
-                        ema_means=None,
-                        summary=True,
+                        temperature_warmup_steps=150000,
                         do_hard_gumbel_softmax=False,
                         do_iaf=False,
-                        approximate_gs_entropy=False):
-  """Discretization bottleneck for latent variables.
+                        approximate_gs_entropy=False,
+                        discrete_mix=0.5,
+                        noise_dev=1.,
+                        startup_steps=50000,
+                        summary=True,
+                        name=None):
+  """Discretization bottleneck.
 
   Args:
-    x: Input to the discretization bottleneck.
-    hidden_size: Dimension of the latent state.
-    z_size: Number of bits used to produce discrete code; discrete codes range
-      from 1 to 2**z_size.
-    filter_size: Filter size to be used for the embedding function.
-    name: Name for the bottleneck scope.
-    mode: Mode represents whether we are training or testing for bottlenecks
-      that differ in behavior (Default: None).
-    startup_steps: Number of steps after which latent predictor is trained
-      (Default: 50000).
-    bottleneck_kind: Kind of discretization bottleneck to use; one of dvq,
-      semhash, gumbel-softmax, gumbel-softmax-dvq (Default: dvq).
-    num_blocks: Number of blocks to use for decomposed vector
-      quantization (Default: 2).
+    inputs: Input to the bottleneck, a Tensor of shape [..., channels].
+    hidden_size: Dimension of the dense output.
+    z_size: Number of bits, where discrete codes range from 1 to 2**z_size.
+    filter_size: Filter size in the embedding function.
+    mode: tf.estimator.ModeKeys.
+    bottleneck_kind: Kind of discretization bottleneck. One of dense, dvq
+      (decomposed vector quantization), gumbel-softmax, gumbel-softmax-dvq,
+      semhash, or vae.
+    num_blocks: Number of blocks. Used only if bottleneck_kind is DVQ.
     num_residuals: Number of residual units used to compute nearest
-      neighbors (Default: 1).
-    reshape_method: Method to reshape for DVQ (Default: slice).
+      neighbors. Used only if bottleneck_kind is DVQ.
+    reshape_method: Method to reshape. Used only if bottleneck_kind is DVQ.
     projection_tensors: If the reshape method is project, then these are the
-      tensors used to project (Default: None).
-    means: The embedding table for dvq (Default: None).
-    beta: Beta factor for the DVQ loss (Default: 0.25).
-    noise_dev: Stddev for noise added for semhash (Default: 0).
-    decay: Decay factor for the exponential moving average (Default: 0.999).
-    discrete_mix: Factor for mixing discrete and non-discrete input for semhash
-      (Default: 0.5).
-    random_top_k: Noisy top-k for DVQ (Default: 1).
-    soft_em: If True then use soft EM rather than hard EM (Default: False).
-    num_samples: Number of samples for soft EM (Default: 1).
-    epsilon: Epsilon parameter for DVQ (Default: 1e-5).
-    softmax_k: If > 1 then do top-k softmax (Default: 0).
-    kl_warmup_steps: Number of steps for kl warmup (Default: 150000).
-    ema: If True update embeddings using exponential moving averages (Default:
-      True).
+      tensors used to project.
+    beta: Scale factor for codebook loss and EMA. Used only if bottleneck_kind
+      is DVQ.
+    ema: Whether to update embeddings using exponential moving averages. Used
+      only if bottleneck_kind is DVQ.
+    means: The embedding table. Used only if ema is True.
     ema_count: Table of counts for each embedding corresponding to how many
-      examples in a batch it was the closest to (Default: None).
-    ema_means: Exponentially averaged version of the embeddings (Default: None).
-    summary: If True, then write summaries (Default: True).
-    do_hard_gumbel_softmax: Boolean determining hard or soft Gumbel-Softmax
-      samples (Default: False).
-    do_iaf: Boolean determining whether we do inverse autoregresive flows for
-      Gumbel-Softmax DVQ bottleneck (Default: False).
-    approximate_gs_entropy: If true, we approximate the Gumbel-Softmax density
-      as a categorical distribution when calculating the sample entropy
-      (Default: False).
+      examples in a batch it was the closest to. Used only if ema is True.
+    ema_means: Exponentially averaged version of the embeddings. Used only if
+      ema is True.
+    epsilon: Small value to avoid dividing by zero in EMA update. Used only if
+      ema is True.
+    decay: Decay factor for the exponential moving average. Used only if ema is
+      True.
+    random_top_k: Noisy top-k. Used only if bottleneck_kind is DVQ.
+    soft_em: Whether to use soft EM or hard EM. Used only if bottleneck_kind is
+      DVQ.
+    num_samples: Number of samples for soft EM. Used only if soft_em is True.
+    softmax_k: If > 0 then do top-k softmax. Used only if bottleneck_kind
+      is gumbel-softmax.
+    temperature_warmup_steps: Number of steps it takes to decay temperature to
+      0. Used only if bottleneck_kind is gumbel-softmax or gumbel-softmax-dvq.
+    do_hard_gumbel_softmax: Whether to use hard or soft Gumbel-Softmax
+      samples. Used only if bottleneck_kind is gumbel-softmax-dvq.
+    do_iaf: Whether to apply inverse autoregresive flows. Used only if
+      bottleneck_kind is gumbel-softmax-dvq.
+    approximate_gs_entropy: Whether to approximate the Gumbel-Softmax density
+      as a categorical distribution when calculating the sample entropy. Used
+      only if bottleneck_kind is gumbel-softmax-dvq.
+    discrete_mix: Factor for mixing discrete and non-discrete input. Used only
+      if bottleneck_kind is semhash.
+    noise_dev: Noise stddev. Used only if bottleneck_kind is semhash.
+    startup_steps: Number of steps after which latent predictor is trained. Used
+      only if bottleneck_kind is semhash.
+    summary: Whether to write summaries.
+    name: Name for the bottleneck scope.
 
   Returns:
-    Embedding to pass to the decoder, discrete latent, loss, and the embedding
-    function.
+    outputs_dense: Tensor of shape [..., output_dim]. The output dimension is
+      hidden_size if bottleneck_kind is gumbel-softmax, DVQ; filter_size if
+      bottleneck_kind is dense, semhash, vae. If bottleneck_kind is DVQ,
+      outputs_dense represents the codebook (means) indexed by outputs_discrete.
+    outputs_discrete: Tensor of shape [...]. Discrete codes, each an index in
+      [0, 2**z_size). It uses the hot representation if soft_em is True.
+    extra_loss: Scalar Tensor. Sum of codebook and commitment losses if
+      bottleneck_kind is DVQ; else zero.
+    embed_fn: Function embed with arguments partially filled in.
+    neg_q_entropy: Scalar Tensor representing negative entropy of variational
+      approximation (0 if it is deterministic).
 
   Raises:
     ValueError: If projection_tensors is None for reshape_method project, or
-    ema_count or ema_means is None if we are using ema, or unknown args.
+    ema_count or ema_means is None if ema is True, or unknown args.
   """
-  block_v_size = None
   if bottleneck_kind in ["dvq", "gumbel-softmax-dvq"]:
-    # Define the dvq parameters
     assert means is not None
-
-    # Check block dimensions add up
     if hidden_size % num_blocks != 0:
       raise ValueError("num_blocks does not divide hidden size")
 
     if z_size % num_residuals != 0:
       raise ValueError("num_residuals does not divide embedding table size")
-
     z_size_per_residual = int(z_size / num_residuals)
 
     if z_size_per_residual % num_blocks != 0:
       raise ValueError("num_blocks does not divide embedding table size")
+    block_v_size = 2**int(z_size_per_residual / num_blocks)
 
-    block_v_size = 2**(z_size_per_residual / num_blocks)
-    block_v_size = int(block_v_size)
-
-    # Set the reshape method corresponding to projections or slices.
-    if reshape_method == "slice":
-      reshape_fn = partial(
-          slice_hidden, hidden_size=hidden_size, num_blocks=num_blocks)
-    elif reshape_method == "project":
-      if projection_tensors is None:
-        raise ValueError(
-            "Projection tensors is None for reshape_method project")
-      reshape_fn = partial(
-          project_hidden,
-          projection_tensors=projection_tensors,
-          hidden_size=hidden_size,
-          num_blocks=num_blocks)
-    else:
-      raise ValueError("Unknown reshape_method")
-
-    # Check if the ema settings make sense.
     if ema:
       if ema_count is None:
         raise ValueError("ema_count is None but ema is True")
       if ema_means is None:
         raise ValueError("ema_means is None but ema is True")
+  else:
+    block_v_size = None
+
+  with tf.variable_scope(name,
+                         default_name="discrete_bottleneck",
+                         reuse=tf.AUTO_REUSE):
+    embed_fn = partial(
+        embed,
+        hidden_size=hidden_size,
+        z_size=z_size,
+        filter_size=filter_size,
+        bottleneck_kind=bottleneck_kind,
+        soft_em=soft_em,
+        num_blocks=num_blocks,
+        num_residuals=num_residuals,
+        block_v_size=block_v_size,
+        means=means,
+        name=name)
 
-  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
-    l = tf.constant(0.0)
-    neg_q_entropy = tf.constant(0.0)
     if bottleneck_kind == "dense":
-      c = tf.layers.dense(x, z_size, name="vcc")
-      h1 = tf.layers.dense(c, filter_size, name="vch1")
-    elif bottleneck_kind == "vae":
-      c, l, _, _ = vae(x, z_size, "vae")
-      h1 = tf.layers.dense(c, filter_size, name="vch1")
-    elif bottleneck_kind == "semhash":
-      c = tf.layers.dense(x, z_size, name="vcc")
-      y_clean = common_layers.saturating_sigmoid(c)
-      if summary:
-        tf.summary.histogram("y_clean", tf.reshape(y_clean, [-1]))
-      if noise_dev > 0 and mode == tf.estimator.ModeKeys.TRAIN:
-        noise = tf.truncated_normal(
-            common_layers.shape_list(c), mean=0.0, stddev=noise_dev)
-        y = common_layers.saturating_sigmoid(c + noise)
-      else:
-        y = y_clean
-      d = tf.to_float(tf.less(0.5, y))
-      y_discrete = tf.stop_gradient(d) + y - tf.stop_gradient(y)
-      pd = common_layers.inverse_exp_decay(startup_steps * 2)
-      pd *= discrete_mix
-      pd = pd if mode == tf.estimator.ModeKeys.TRAIN else 1.0
-      c = tf.where(
-          tf.less(tf.random_uniform([common_layers.shape_list(y)[0]]), pd),
-          y_discrete, y)
-      h1a = tf.layers.dense(c, filter_size, name="vch1a")
-      h1b = tf.layers.dense(1.0 - c, filter_size, name="vch1b")
-      h1 = h1a + h1b
-      dx = tf.to_int32(tf.stop_gradient(d))
-      c = bit_to_int(dx, z_size)
-    elif bottleneck_kind == "gumbel-softmax":
-      _, hot, l = gumbel_softmax(x, name, z_size, mode, softmax_k,
-                                 kl_warmup_steps, summary)
-      c = tf.argmax(hot, axis=-1)
-      h1 = tf.layers.dense(hot, hidden_size, name="dae_dense")
+      # Note discrete output is continuous here.
+      outputs_discrete = tf.layers.dense(inputs, z_size, name="vcc")
+      outputs_dense = tf.layers.dense(outputs_discrete,
+                                      filter_size,
+                                      name="vch1")
+      extra_loss = tf.constant(0.0)
+      neg_q_entropy = tf.constant(0.0)
     elif bottleneck_kind in ["dvq", "gumbel-softmax-dvq"]:
-      x_reshaped = reshape_fn(x)
+      if reshape_method == "slice":
+        x_reshaped = slice_hidden(inputs,
+                                  hidden_size=hidden_size,
+                                  num_blocks=num_blocks)
+      elif reshape_method == "project":
+        if projection_tensors is None:
+          raise ValueError(
+              "Projection tensors is None for reshape_method project")
+        x_reshaped = project_hidden(inputs,
+                                    projection_tensors=projection_tensors,
+                                    hidden_size=hidden_size,
+                                    num_blocks=num_blocks)
+      else:
+        raise ValueError("Unknown reshape_method")
+
       x_res = tf.reshape(x_reshaped,
                          [-1] + common_layers.shape_list(x_reshaped)[2:])
       x_means_hot = []
       x_means = 0
-      l = 0
+      extra_loss = 0
       for i in range(num_residuals):
         x_means_hot_res, x_means_res, q_loss_res, e_loss_res, neg_q_entropy = (
-            embedding_lookup(x_reshaped, means[i], num_blocks, block_v_size,
-                             bottleneck_kind, random_top_k, soft_em,
-                             num_samples, do_hard_gumbel_softmax,
+            embedding_lookup(x_reshaped,
+                             means=means[i],
+                             num_blocks=num_blocks,
+                             block_v_size=block_v_size,
+                             bottleneck_kind=bottleneck_kind,
+                             random_top_k=random_top_k,
+                             soft_em=soft_em,
+                             num_samples=num_samples,
+                             temperature_warmup_steps=temperature_warmup_steps,
+                             do_hard_gumbel_softmax=do_hard_gumbel_softmax,
                              do_iaf=do_iaf,
                              approximate_gs_entropy=approximate_gs_entropy))
-        # Update the ema variables
+        # Update the EMA variables.
         if ema:
           tf.logging.info("Using EMA with beta = {}".format(beta))
           updated_ema_count_res = moving_averages.assign_moving_average(
@@ -646,9 +647,8 @@ def discrete_bottleneck(x,
               decay,
               zero_debias=False)
 
-          dw = tf.matmul(
-              tf.transpose(x_means_hot_res, perm=[1, 2, 0]),
-              tf.transpose(x_res, perm=[1, 0, 2]))
+          dw = tf.matmul(tf.transpose(x_means_hot_res, perm=[1, 2, 0]),
+                         tf.transpose(x_res, perm=[1, 0, 2]))
 
           updated_ema_means_res = moving_averages.assign_moving_average(
               ema_means[i], dw, decay, zero_debias=False)
@@ -663,20 +663,20 @@ def discrete_bottleneck(x,
           with tf.control_dependencies([e_loss_res]):
             update_means_res = tf.assign(means[i], updated_ema_means_res)
             with tf.control_dependencies([update_means_res]):
-              l += beta * e_loss_res
+              extra_loss += beta * e_loss_res
         else:
-          l += q_loss_res + beta * e_loss_res
+          extra_loss += q_loss_res + beta * e_loss_res
 
-        # Update the residuals
+        # Update the residuals.
         x_res -= x_means_res
         x_means += x_means_res
         x_means_hot.append(x_means_hot_res)
 
-      # Get the discrete latent representation
+      # Get the discrete latent representation.
       x_means_hot = tf.stack(x_means_hot, axis=1)
       x_means_idx = tf.argmax(x_means_hot, axis=-1)
 
-      # Get the binary representation
+      # Get the binary representation.
       x_means_bits = int_to_bit(
           x_means_idx,
           num_bits=int(z_size / (num_residuals * num_blocks)),
@@ -685,40 +685,80 @@ def discrete_bottleneck(x,
       new_shape = shape[:-2]
       new_shape[-1] = z_size
       x_means_bits = tf.reshape(x_means_bits, shape=new_shape)
-      c = bit_to_int(tf.to_int32(x_means_bits), num_bits=z_size, base=2)
+      outputs_discrete = bit_to_int(tf.to_int32(x_means_bits),
+                                    num_bits=z_size,
+                                    base=2)
 
-      # Adjust shape of c
-      shape_x = common_layers.shape_list(x)
-      new_shape = shape_x[:-1]
-      c = tf.reshape(c, new_shape)
+      # Adjust shape of discrete outputs.
+      inputs_shape = common_layers.shape_list(inputs)
+      outputs_discrete = tf.reshape(outputs_discrete, inputs_shape[:-1])
 
-      # If we are doing soft EM then c is x_means_hot
+      # If we're using soft EM then set discretes to the hot representation.
       if soft_em:
-        c = x_means_hot
-        new_shape.append(block_v_size)
-        c = tf.reshape(c, new_shape)
+        outputs_discrete = x_means_hot
+        outputs_discrete = tf.reshape(outputs_discrete,
+                                      inputs_shape[:-1] + [block_v_size])
 
-      x_means = tf.reshape(x_means, shape_x)
-      x_reshaped = tf.reshape(x_reshaped, shape_x)
-      h1 = x_reshaped + tf.stop_gradient(x_means - x_reshaped)
+      # Reshape assuming hidden_size == inputs_shape[:-1].
+      x_means = tf.reshape(x_means, inputs_shape)
+      outputs_dense = inputs + tf.stop_gradient(x_means - inputs)
+    elif bottleneck_kind == "gumbel-softmax":
+      _, outputs_hot, extra_loss = gumbel_softmax(
+          inputs,
+          z_size=z_size,
+          mode=mode,
+          softmax_k=softmax_k,
+          temperature_warmup_steps=temperature_warmup_steps,
+          summary=summary,
+          name=name)
+      outputs_discrete = tf.argmax(outputs_hot, axis=-1)
+      outputs_dense = tf.layers.dense(outputs_hot,
+                                      hidden_size,
+                                      name="dae_dense")
+      neg_q_entropy = tf.constant(0.0)
+    elif bottleneck_kind == "semhash":
+      outputs_discrete = tf.layers.dense(inputs, z_size, name="vcc")
+      y_clean = common_layers.saturating_sigmoid(outputs_discrete)
+      if summary:
+        tf.summary.histogram("y_clean", tf.reshape(y_clean, [-1]))
+      if noise_dev > 0 and mode == tf.estimator.ModeKeys.TRAIN:
+        noise = tf.truncated_normal(common_layers.shape_list(outputs_discrete),
+                                    mean=0.0,
+                                    stddev=noise_dev)
+        y = common_layers.saturating_sigmoid(outputs_discrete + noise)
+      else:
+        y = y_clean
+      d = tf.to_float(tf.less(0.5, y))
+      y_discrete = tf.stop_gradient(d) + y - tf.stop_gradient(y)
+      pd = common_layers.inverse_exp_decay(startup_steps * 2)
+      pd *= discrete_mix
+      pd = pd if mode == tf.estimator.ModeKeys.TRAIN else 1.0
+      c = tf.where(
+          tf.less(tf.random_uniform([common_layers.shape_list(y)[0]]), pd),
+          y_discrete, y)
+      outputs_dense_a = tf.layers.dense(c,
+                                        filter_size,
+                                        name="vch1a")
+      outputs_dense_b = tf.layers.dense(1.0 - c,
+                                        filter_size,
+                                        name="vch1b")
+      outputs_dense = outputs_dense_a + outputs_dense_b
+      dx = tf.to_int32(tf.stop_gradient(d))
+      outputs_discrete = bit_to_int(dx, z_size)
+      extra_loss = tf.constant(0.0)
+      neg_q_entropy = tf.constant(0.0)
+    elif bottleneck_kind == "vae":
+      outputs_discrete, extra_loss, _, _ = vae(inputs,
+                                               z_size,
+                                               name="vae")
+      outputs_dense = tf.layers.dense(outputs_discrete,
+                                      filter_size,
+                                      name="vch1")
+      neg_q_entropy = tf.constant(0.0)
     else:
       raise ValueError("Unknown discretization method.")
 
-    res = h1
-
-    embed_fn = partial(
-        embed,
-        hidden_size=hidden_size,
-        z_size=z_size,
-        filter_size=filter_size,
-        name=name,
-        bottleneck_kind=bottleneck_kind,
-        soft_em=soft_em,
-        num_blocks=num_blocks,
-        num_residuals=num_residuals,
-        block_v_size=block_v_size,
-        means=means)
-    return res, c, l, embed_fn, neg_q_entropy
+  return outputs_dense, outputs_discrete, extra_loss, embed_fn, neg_q_entropy
 
 
 # New API for discretization bottlenecks:
@@ -827,7 +867,7 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
                                         hard=False,
                                         temperature_init=1.2,
                                         num_samples=1,
-                                        startup_steps=15000,
+                                        temperature_warmup_steps=150000,
                                         summary=True,
                                         do_iaf=False,
                                         approximate_gs_entropy=False):
@@ -843,8 +883,8 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
     temperature_init: Initial temperature used for Gumbel-Softmax samples,
       after it which it decays to 0 (Default: 1.2).
     num_samples: Number of samples drawn for each latent (Default: 1).
-    startup_steps: Number of steps it takes to decay temperature to 0 (Default:
-      15000).
+    temperature_warmup_steps: Number of steps it takes to decay temperature to 0
+      (Default: 150000).
     summary: When `True`, we save histogram summaries of the KL term (Default:
       True).
     do_iaf: When `True`, we perform inverse autoregressive flow with
@@ -886,7 +926,7 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
 
   # Temperature decays linearly.
   temperature = temperature_init - common_layers.inverse_lin_decay(
-      startup_steps)
+      temperature_warmup_steps)
 
   # 10% of the time keep reasonably high temperature to keep learning.
   temperature = tf.cond(
@@ -1011,7 +1051,7 @@ def gumbel_softmax_discrete_bottleneck(x,
                                        beta=0.25,
                                        decay=0.999,
                                        epsilon=1e-5,
-                                       startup_steps=15000,
+                                       temperature_warmup_steps=150000,
                                        hard=False,
                                        summary=True):
   """VQ-VAE using Gumbel-Softmax.
@@ -1031,7 +1071,8 @@ def gumbel_softmax_discrete_bottleneck(x,
     decay: Decay factor for exponential moving average (Default: 0.999).
     epsilon: Small value to avoid dividing by zero in EMA update
       (Default: 1e-5).
-    startup_steps: Number of steps for KL warmup (Default: 25000).
+    temperature_warmup_steps: Number of steps it takes to decay temperature to 0
+      (Default: 150000).
     hard: When `True`, we use hard Gumbel-Softmax samples and force
       discrete latents by taking the argmax. When `False`, we use soft samples,
       which we treat as codebook weights (Default: False).
@@ -1065,8 +1106,9 @@ def gumbel_softmax_discrete_bottleneck(x,
   class_probs = tf.nn.softmax(dist)
   log_class_probs = tf.nn.log_softmax(dist)
   gumbel_samples = gumbel_sample(common_layers.shape_list(dist))
-  gumbel_samples *= common_layers.inverse_exp_decay(startup_steps // 5) * 0.5
-  temperature = 1.2 - common_layers.inverse_lin_decay(startup_steps)
+  steps = temperature_warmup_steps
+  gumbel_samples *= common_layers.inverse_exp_decay(steps // 5) * 0.5
+  temperature = 1.2 - common_layers.inverse_lin_decay(steps)
 
   # 10% of the time keep reasonably high temperature to keep learning.
   temperature = tf.cond(
@@ -1201,10 +1243,13 @@ def parametrized_bottleneck(x, hparams):
         soft_em=True,
         num_samples=hparams.vq_num_samples)
   if hparams.bottleneck_kind == "gumbel_softmax":
-    return gumbel_softmax_discrete_bottleneck(x, hparams.bottleneck_bits,
-                                              hparams.vq_beta, hparams.vq_decay,
+    return gumbel_softmax_discrete_bottleneck(x,
+                                              hparams.bottleneck_bits,
+                                              hparams.vq_beta,
+                                              hparams.vq_decay,
                                               hparams.vq_epsilon,
-                                              hparams.startup_steps, hard=False,
+                                              hparams.temperature_warmup_steps,
+                                              hard=False,
                                               summary=True)
 
   raise ValueError("Unsupported hparams.bottleneck_kind %s"
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 0e9a72e91..d2e253279 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -372,11 +372,10 @@ def ae_transformer_internal(inputs,
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
       # Compress and bottleneck.
       latents_dense, latents_discrete, extra_loss, embed, neg_q_entropy = (
-          hparams.bottleneck(
-              x=targets_c,
-              filter_size=hparams.compress_filter_size,
-              name="vc",
-              mode=hparams.mode))
+          hparams.bottleneck(inputs=targets_c,
+                             filter_size=hparams.compress_filter_size,
+                             mode=hparams.mode,
+                             name="vc"))
       if _DO_SUMMARIES:
         tf.summary.histogram("b0", tf.reshape(latents_discrete[:, 0, :], [-1]))
       pc = common_layers.inverse_exp_decay(hparams.startup_steps)
@@ -401,11 +400,11 @@ def ae_transformer_internal(inputs,
         losses["latent_pred"] = tf.reduce_mean((inputs_c - targets_c)**2) * 20
         def bn_inputs():
           with tf.variable_scope(tf.get_variable_scope(), reuse=True):
-            bn, _, _, _ = hparams.bottleneck(
-                x=inputs_c,
+            bn, _, _, _, _ = hparams.bottleneck(
+                inputs=inputs_c,
                 filter_size=hparams.compress_filter_size,
-                name="vc",
-                mode=hparams.mode)
+                mode=hparams.mode,
+                name="vc")
           return bn
         inputs_c = bn_inputs
         ptc = 1.0 - common_layers.inverse_lin_decay(200000) * 0.5
@@ -415,15 +414,17 @@ def bn_inputs():
     else:
       if hparams.bottleneck_kind in ["dense", "vae"]:
         inputs_c = decode_transformer(inputs, ed, targets_c, hparams, "dec_c")
-        latents_dense, _, _, _ = hparams.bottleneck(
-            x=inputs_c,
+        latents_dense, _, _, _, _ = hparams.bottleneck(
+            inputs=inputs_c,
             filter_size=hparams.compress_filter_size,
-            name="vc",
-            mode=hparams.mode)
+            mode=hparams.mode,
+            name="vc")
       else:
         latent_len = common_layers.shape_list(targets_c)[1]
-        _, _, _, embed = hparams.bottleneck(
-            x=targets_c, filter_size=hparams.compress_filter_size, name="vc")
+        _, _, _, embed, _ = hparams.bottleneck(
+            inputs=targets_c,
+            filter_size=hparams.compress_filter_size,
+            name="vc")
         latents_dense = tf.zeros_like(targets_c[:, :latent_len, :, :])
         if cache is None:
           cache = ae_latent_sample(
@@ -502,26 +503,26 @@ def __init__(self, *args, **kwargs):
         hidden_size=self._hparams.hidden_size,
         z_size=self._hparams.z_size,
         filter_size=self._hparams.filter_size,
-        startup_steps=self.hparams.startup_steps,
         bottleneck_kind=self._hparams.bottleneck_kind,
         num_blocks=self._hparams.num_blocks,
         num_residuals=self.hparams.num_residuals,
         reshape_method=self._hparams.reshape_method,
         beta=self._hparams.beta,
-        noise_dev=self._hparams.noise_dev,
+        ema=self._hparams.ema,
+        epsilon=self._hparams.epsilon,
         decay=self._hparams.decay,
-        discrete_mix=self._hparams.d_mix,
         random_top_k=self._hparams.random_top_k,
         soft_em=self.hparams.soft_em,
         num_samples=self.hparams.num_samples,
-        epsilon=self._hparams.epsilon,
         softmax_k=self._hparams.softmax_k,
-        kl_warmup_steps=self._hparams.kl_warmup_steps,
-        ema=self._hparams.ema,
-        summary=_DO_SUMMARIES,
+        temperature_warmup_steps=self._hparams.temperature_warmup_steps,
         do_hard_gumbel_softmax=self._hparams.do_hard_gumbel_softmax,
         do_iaf=self._hparams.do_iaf,
-        approximate_gs_entropy=self._hparams.approximate_gs_entropy)
+        approximate_gs_entropy=self._hparams.approximate_gs_entropy,
+        discrete_mix=self._hparams.d_mix,
+        noise_dev=self._hparams.noise_dev,
+        startup_steps=self.hparams.startup_steps,
+        summary=_DO_SUMMARIES)
     # Set the discretization bottleneck specific things here
     if self._hparams.bottleneck_kind in ["dvq", "gumbel-softmax-dvq"]:
       z_size_per_residual = self._hparams.z_size / self._hparams.num_residuals
@@ -730,7 +731,7 @@ def transformer_ae_small():
   hparams.add_hparam("do_hard_gumbel_softmax", False)
   hparams.add_hparam("do_iaf", False)
   hparams.add_hparam("approximate_gs_entropy", False)
-  hparams.kl_warmup_steps = 150000
+  hparams.add_hparam("temperature_warmup_steps", 150000)
   hparams.force_full_predict = True
 
   # task params

From a14d88703ddba366e38419db8fd523c544df82ab Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 23 Jul 2018 10:42:56 -0700
Subject: [PATCH 0368/2720] Update README.md with image generation.

PiperOrigin-RevId: 205686148
---
 README.md           | 20 ++++++++++++++++++++
 docs/walkthrough.md | 20 ++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/README.md b/README.md
index 6fe27fc4a..c96925c4d 100644
--- a/README.md
+++ b/README.md
@@ -47,6 +47,7 @@ pip install tensor2tensor && t2t-trainer \
 
 * [Suggested Datasets and Models](#suggested-datasets-and-models)
   * [Image Classification](#image-classification)
+  * [Image Generation](#image-generation)
   * [Language Modeling](#language-modeling)
   * [Sentiment Analysis](#sentiment-analysis)
   * [Speech Recognition](#speech-recognition)
@@ -97,6 +98,25 @@ For CIFAR and MNIST, we suggest to try the shake-shake model:
 This setting trained for `--train_steps=700000` should yield
 close to 97% accuracy on CIFAR-10.
 
+### Image Generation
+
+For (un)conditional image generation, we have a number of standard data-sets:
+
+* CelebA: `--problem=img2img_celeba` for image-to-image translation, namely,
+    superresolution from 8x8 to 32x32.
+* CelebA-HQ: `--problem=image_celeba256_rev` for a downsampled 256x256.
+* CIFAR-10: `--problem=image_cifar10_plain_gen_rev` for class-conditional
+    32x32 generation.
+* LSUN Bedrooms: `--problem=image_lsun_bedrooms_rev`
+* MS-COCO: `--problem=image_text_ms_coco_rev` for text-to-image generation.
+* Small ImageNet (a large data-set): `--problem=image_imagenet32_gen_rev` for
+    32x32 or `--problem=image_imagenet64_gen_rev` for 64x64.
+
+We suggest to use the Image Transformer, i.e., `--model=imagetransformer`, or
+variational auto-encoder, i.e., `--model=transformer_ae`. For CIFAR-10, using
+`--hparams_set=imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_m` yields 2.90
+bits per dimension.
+
 ### Language Modeling
 
 For language modeling, we have these data-sets in T2T:
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 6fe27fc4a..c96925c4d 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -47,6 +47,7 @@ pip install tensor2tensor && t2t-trainer \
 
 * [Suggested Datasets and Models](#suggested-datasets-and-models)
   * [Image Classification](#image-classification)
+  * [Image Generation](#image-generation)
   * [Language Modeling](#language-modeling)
   * [Sentiment Analysis](#sentiment-analysis)
   * [Speech Recognition](#speech-recognition)
@@ -97,6 +98,25 @@ For CIFAR and MNIST, we suggest to try the shake-shake model:
 This setting trained for `--train_steps=700000` should yield
 close to 97% accuracy on CIFAR-10.
 
+### Image Generation
+
+For (un)conditional image generation, we have a number of standard data-sets:
+
+* CelebA: `--problem=img2img_celeba` for image-to-image translation, namely,
+    superresolution from 8x8 to 32x32.
+* CelebA-HQ: `--problem=image_celeba256_rev` for a downsampled 256x256.
+* CIFAR-10: `--problem=image_cifar10_plain_gen_rev` for class-conditional
+    32x32 generation.
+* LSUN Bedrooms: `--problem=image_lsun_bedrooms_rev`
+* MS-COCO: `--problem=image_text_ms_coco_rev` for text-to-image generation.
+* Small ImageNet (a large data-set): `--problem=image_imagenet32_gen_rev` for
+    32x32 or `--problem=image_imagenet64_gen_rev` for 64x64.
+
+We suggest to use the Image Transformer, i.e., `--model=imagetransformer`, or
+variational auto-encoder, i.e., `--model=transformer_ae`. For CIFAR-10, using
+`--hparams_set=imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_m` yields 2.90
+bits per dimension.
+
 ### Language Modeling
 
 For language modeling, we have these data-sets in T2T:

From 08eef676aa76b9d18929270a1a2b213e086b7fa0 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 23 Jul 2018 11:18:57 -0700
Subject: [PATCH 0369/2720] Adding a simple One-Hot Symbol modality.

PiperOrigin-RevId: 205692939
---
 tensor2tensor/layers/modalities.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 49da0c4cb..c1efbdfda 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -177,6 +177,26 @@ def targets_weights_fn(self):
     return common_layers.weights_all
 
 
+@registry.register_symbol_modality("one_hot")
+class SymbolModalityOneHot(SymbolModality):
+  """Simple SymbolModality with one hot as embeddings."""
+
+  def bottom(self, x):
+    return tf.one_hot(x, self._vocab_size)
+
+  def targets_bottom(self, x):
+    return tf.one_hot(x, self._vocab_size)
+
+  def top(self, body_output, _):
+    return body_output
+
+  def loss(self, top_out, targets):
+    labels = tf.one_hot(targets, self._vocab_size)
+    loss = tf.nn.softmax_cross_entropy_with_logits(
+        logits=top_out, labels=labels)
+    return tf.reduce_mean(loss), tf.constant(1.0)
+
+
 @registry.register_symbol_modality("ctc")
 class CTCSymbolModality(SymbolModality):
   """SymbolModality that uses CTC loss."""

From e84c39aee20f9682982b7a523ba1e0a1bb8043b1 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 23 Jul 2018 11:50:12 -0700
Subject: [PATCH 0370/2720] outputting "extra" predictions if exists.

PiperOrigin-RevId: 205698529
---
 tensor2tensor/utils/t2t_model.py | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index e59ec2373..e504b9866 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1337,7 +1337,7 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
   def estimator_spec_predict(self, features, use_tpu=False):
     """Construct EstimatorSpec for PREDICT mode."""
     decode_hparams = self._decode_hparams
-    predictions = self.infer(
+    infer_out = self.infer(
         features,
         beam_size=decode_hparams.beam_size,
         top_beams=(decode_hparams.beam_size
@@ -1345,19 +1345,24 @@ def estimator_spec_predict(self, features, use_tpu=False):
         alpha=decode_hparams.alpha,
         decode_length=decode_hparams.extra_length,
         use_tpu=use_tpu)
-    if not isinstance(predictions, dict):
-      predictions = {"outputs": predictions}
+    if isinstance(infer_out, dict):
+      outputs = infer_out["outputs"]
+      scores = infer_out["scores"]
+    else:
+      outputs = infer_out
+      scores = None
 
     inputs = features.get("inputs")
     if inputs is None:
       inputs = features["targets"]
 
-    predictions.update({
+    predictions = {
+        "outputs": outputs,
+        "scores": scores,
         "inputs": inputs,
         "targets": features.get("infer_targets"),
         "batch_prediction_key": features.get("batch_prediction_key"),
-        })
-
+    }
     _del_dict_nones(predictions)
 
     export_out = {"outputs": predictions["outputs"]}

From 73bd9377514a04edd9a669f2a32039602fe623ea Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 23 Jul 2018 12:02:01 -0700
Subject: [PATCH 0371/2720] Internal change

PiperOrigin-RevId: 205700378
---
 tensor2tensor/data_generators/all_problems.py | 17 +++++------------
 1 file changed, 5 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 0188a84ff..5a1ffedce 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -18,7 +18,7 @@
 from __future__ import print_function
 
 import importlib
-import six
+import re
 
 MODULES = [
     "tensor2tensor.data_generators.algorithmic",
@@ -84,16 +84,9 @@
 
 
 
-def _py_err_msg(module):
-  if six.PY2:
-    # Py2 error will reference the module relative to the current module
-    shared_module = "data_generators."
-    start_idx = module.index(shared_module) + len(shared_module)
-    err_name = module[start_idx:]
-    msg = "No module named %s" % err_name
-  else:
-    msg = "No module named '%s'" % module
-  return msg
+def _is_import_err_msg(err_str, module):
+  module_pattern = "(.)?".join(["(%s)?" % m for m in module.split(".")])
+  return re.match("^No module named (')?%s(')?$" % module_pattern, err_str)
 
 
 def _handle_errors(errors):
@@ -105,7 +98,7 @@ def _handle_errors(errors):
   print(err_msg.format(num_missing=len(errors)))
   for module, err in errors:
     err_str = str(err)
-    if err_str != _py_err_msg(module):
+    if not _is_import_err_msg(err_str, module):
       print("From module %s" % module)
       raise err
     if log_all:

From bfcea1525cf9a65ba851a6006f670285f6e30a51 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 23 Jul 2018 16:51:28 -0700
Subject: [PATCH 0372/2720] Stochastic Adversarial Video Prediction: Part 1 of
 2

PiperOrigin-RevId: 205747907
---
 tensor2tensor/layers/modalities.py            |  22 +-
 tensor2tensor/models/research/next_frame.py   | 329 +++++++++++++++---
 .../models/research/next_frame_test.py        |  20 +-
 3 files changed, 324 insertions(+), 47 deletions(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index c1efbdfda..e1ae5401f 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -684,6 +684,13 @@ def internal_loss(self, logits, targets):
 class VideoModalityL2Raw(VideoModalityL2):
   """Modality with L2 loss and raw input (sequences of frames)."""
 
+  def convert_rgb_to_real(self, prediction, targets):
+    """Convert prediction and target from rgb to real."""
+    prediction = tf.squeeze(prediction, axis=-1)
+    prediction = common_layers.convert_rgb_to_real(prediction)
+    targets = common_layers.convert_rgb_to_real(targets)
+    return prediction, targets
+
   def bottom(self, x):
     common_layers.summarize_video(x, "inputs")
     return common_layers.convert_rgb_to_real(x)
@@ -699,14 +706,21 @@ def top(self, body_output, _):
     return tf.expand_dims(rgb_frames, axis=-1)
 
   def loss(self, top_out, targets):
-    prediction = top_out
-    prediction = tf.squeeze(prediction, axis=-1)
-    prediction = common_layers.convert_rgb_to_real(prediction)
-    groundtruth = common_layers.convert_rgb_to_real(targets)
+    prediction, groundtruth = self.convert_rgb_to_real(top_out, targets)
     loss = tf.losses.mean_squared_error(prediction, groundtruth)
     return loss, tf.constant(1.0)
 
 
+@registry.register_video_modality("l1raw")
+class VideoModalityL1Raw(VideoModalityL2Raw):
+  """Modality with L1 loss and raw input (sequences of frames)."""
+
+  def loss(self, top_out, targets):
+    prediction, groundtruth = self.convert_rgb_to_real(top_out, targets)
+    loss = tf.losses.absolute_difference(prediction, groundtruth)
+    return loss, tf.constant(1.0)
+
+
 @registry.register_class_label_modality("default")
 class ClassLabelModality(modality.Modality):
   """Used for label data."""
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index e3cd9070a..ffc4da732 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -27,10 +27,20 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow as tf
+from tensorflow_models.slim.nets.cyclegan import cyclegan_upsample
 tfl = tf.layers
 tfcl = tf.contrib.layers
 
 
+def basic_lstm(inputs, state, num_units, name=None):
+  input_shape = common_layers.shape_list(inputs)
+  cell = tf.contrib.rnn.BasicLSTMCell(num_units, name=name)
+  if state is None:
+    state = cell.zero_state(input_shape[0], tf.float32)
+  outputs, new_state = cell(inputs, state)
+  return outputs, new_state
+
+
 @registry.register_model
 class NextFrameBasic(t2t_model.T2TModel):
   """Basic next-frame model, may take actions and predict rewards too."""
@@ -179,6 +189,24 @@ def tinyify(self, array):
       return [1 for _ in array]
     return array
 
+  def tile_and_concat(self, image, latent, concat_latent=True):
+    """Tile latent and concatenate to image across depth."""
+    if not concat_latent:
+      return image
+    height, width = image.shape[1:3].as_list()
+    tf.logging.info("Height")
+    tf.logging.info(height)
+    tf.logging.info("Width")
+    tf.logging.info(width)
+    _, latent_dims = latent.shape.as_list()
+
+    height_multiples = height // latent_dims
+    pad = height - (height_multiples * latent_dims)
+    latent = tf.reshape(latent, (-1, latent_dims, 1, 1))
+    latent = tf.tile(latent, (1, height_multiples, width, 1))
+    latent = tf.pad(latent, [[0, 0], [pad // 2, pad // 2], [0, 0], [0, 0]])
+    return tf.concat([image, latent], axis=-1)
+
   def construct_latent_tower(self, images):
     """Builds convolutional latent tower for stochastic model.
 
@@ -234,7 +262,7 @@ def construct_latent_tower(self, images):
       return mean, std
 
   def bottom_part_tower(self, input_image, input_reward, action, latent,
-                        lstm_state, lstm_size, conv_size):
+                        lstm_state, lstm_size, conv_size, concat_latent=False):
     """The bottom part of predictive towers.
 
     With the current (early) design, the main prediction tower and
@@ -250,6 +278,7 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
       lstm_state: the current internal states of conv lstms.
       lstm_size: the size of lstms.
       conv_size: the size of convolutions.
+      concat_latent: whether or not to concatenate the latent at every step.
 
     Returns:
       - the output of the partial network.
@@ -258,8 +287,11 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
     lstm_func = self.conv_lstm_2d
 
     input_image = common_layers.make_even_size(input_image)
+    concat_input_image = self.tile_and_concat(
+        input_image, latent, concat_latent=concat_latent)
+
     enc0 = tfl.conv2d(
-        input_image,
+        concat_input_image,
         conv_size[0], [5, 5],
         strides=(2, 2),
         activation=tf.nn.relu,
@@ -269,6 +301,7 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
 
     hidden1, lstm_state[0] = lstm_func(
         enc0, lstm_state[0], lstm_size[0], name="state1")
+    hidden1 = self.tile_and_concat(hidden1, latent, concat_latent=concat_latent)
     hidden1 = tfcl.layer_norm(hidden1, scope="layer_norm2")
     hidden2, lstm_state[1] = lstm_func(
         hidden1, lstm_state[1], lstm_size[1], name="state2")
@@ -276,12 +309,15 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
     hidden2 = common_layers.make_even_size(hidden2)
     enc1 = tfl.conv2d(hidden2, hidden2.get_shape()[3], [3, 3], strides=(2, 2),
                       padding="SAME", activation=tf.nn.relu, name="conv2")
+    enc1 = self.tile_and_concat(enc1, latent, concat_latent=concat_latent)
 
     hidden3, lstm_state[2] = lstm_func(
         enc1, lstm_state[2], lstm_size[2], name="state3")
+    hidden3 = self.tile_and_concat(hidden3, latent, concat_latent=concat_latent)
     hidden3 = tfcl.layer_norm(hidden3, scope="layer_norm4")
     hidden4, lstm_state[3] = lstm_func(
         hidden3, lstm_state[3], lstm_size[3], name="state4")
+    hidden4 = self.tile_and_concat(hidden4, latent, concat_latent=concat_latent)
     hidden4 = tfcl.layer_norm(hidden4, scope="layer_norm5")
     hidden4 = common_layers.make_even_size(hidden4)
     enc2 = tfl.conv2d(hidden4, hidden4.get_shape()[3], [3, 3], strides=(2, 2),
@@ -293,7 +329,7 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
         input_reward, enc2.get_shape(), "reward_enc")
     enc2 = tf.concat(axis=3, values=[enc2, emb_action, emb_reward])
 
-    if latent is not None:
+    if latent is not None and not concat_latent:
       with tf.control_dependencies([latent]):
         enc2 = tf.concat([enc2, latent], 3)
 
@@ -303,7 +339,7 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
     hidden5, lstm_state[4] = lstm_func(
         enc3, lstm_state[4], lstm_size[4], name="state5")  # last 8x8
     hidden5 = tfcl.layer_norm(hidden5, scope="layer_norm6")
-
+    hidden5 = self.tile_and_concat(hidden5, latent, concat_latent=concat_latent)
     return hidden5, (enc0, enc1)
 
   def reward_prediction(
@@ -367,13 +403,15 @@ def conv_lstm_2d(self, inputs, state, output_channels,
     return outputs, new_state
 
   def construct_predictive_tower(
-      self, input_image, input_reward, action, lstm_state, latent):
+      self, input_image, input_reward, action, lstm_state, latent,
+      concat_latent=False):
     # Main tower
     lstm_func = self.conv_lstm_2d
     batch_size = common_layers.shape_list(input_image)[0]
     # the number of different pixel motion predictions
     # and the number of masks for each of those predictions
     num_masks = self.hparams.num_masks
+    upsample_method = self.hparams.upsample_method
 
     lstm_size = self.tinyify([32, 32, 64, 64, 128, 64, 32])
     conv_size = self.tinyify([32])
@@ -383,25 +421,35 @@ def construct_predictive_tower(
     with tf.variable_scope("main", reuse=tf.AUTO_REUSE):
       hidden5, skips = self.bottom_part_tower(
           input_image, input_reward, action, latent,
-          lstm_state, lstm_size, conv_size)
+          lstm_state, lstm_size, conv_size, concat_latent=concat_latent)
       enc0, enc1 = skips
 
-      enc4 = tfl.conv2d_transpose(
-          hidden5, hidden5.get_shape()[3], 3, strides=2, name="convt1")
+      with tf.variable_scope("upsample1", reuse=tf.AUTO_REUSE):
+        enc4 = cyclegan_upsample(
+            hidden5, num_outputs=hidden5.shape.as_list()[-1],
+            stride=[2, 2], method=upsample_method)
 
       enc1_shape = common_layers.shape_list(enc1)
       enc4 = enc4[:, :enc1_shape[1], :enc1_shape[2], :]  # Cut to shape.
+      enc4 = self.tile_and_concat(enc4, latent, concat_latent=concat_latent)
+
       hidden6, lstm_state[5] = lstm_func(
           enc4, lstm_state[5], lstm_size[5], name="state6")  # 16x16
+      hidden6 = self.tile_and_concat(
+          hidden6, latent, concat_latent=concat_latent)
       hidden6 = tfcl.layer_norm(hidden6, scope="layer_norm7")
       # Skip connection.
       hidden6 = tf.concat(axis=3, values=[hidden6, enc1])  # both 16x16
 
-      enc5 = tfl.conv2d_transpose(
-          hidden6, hidden6.get_shape()[3], [3, 3], strides=(2, 2),
-          padding="SAME", activation=tf.nn.relu, name="convt2")
+      with tf.variable_scope("upsample2", reuse=tf.AUTO_REUSE):
+        enc5 = cyclegan_upsample(
+            hidden6, num_outputs=hidden6.shape.as_list()[-1],
+            stride=[2, 2], method=upsample_method)
+
       enc0_shape = common_layers.shape_list(enc0)
       enc5 = enc5[:, :enc0_shape[1], :enc0_shape[2], :]  # Cut to shape.
+      enc5 = self.tile_and_concat(enc5, latent, concat_latent=concat_latent)
+
       hidden7, lstm_state[6] = lstm_func(
           enc5, lstm_state[6], lstm_size[6], name="state7")  # 32x32
       hidden7 = tfcl.layer_norm(hidden7, scope="layer_norm8")
@@ -409,15 +457,12 @@ def construct_predictive_tower(
       # Skip connection.
       hidden7 = tf.concat(axis=3, values=[hidden7, enc0])  # both 32x32
 
-      enc6 = tfl.conv2d_transpose(
-          hidden7,
-          hidden7.get_shape()[3],
-          [3, 3],
-          strides=(2, 2),
-          padding="SAME",
-          name="convt3",
-          activation=None)
+      with tf.variable_scope("upsample3", reuse=tf.AUTO_REUSE):
+        enc6 = cyclegan_upsample(
+            hidden7, num_outputs=hidden7.shape.as_list()[-1],
+            stride=[2, 2], method=upsample_method)
       enc6 = tfcl.layer_norm(enc6, scope="layer_norm9")
+      enc6 = self.tile_and_concat(enc6, latent, concat_latent=concat_latent)
 
       if self.hparams.model_options == "DNA":
         # Using largest hidden state for predicting untied conv kernels.
@@ -454,9 +499,9 @@ def construct_predictive_tower(
           raise ValueError("Only one mask is supported for DNA model.")
         transformed = [self.dna_transformation(input_image, enc7)]
 
-      masks = tfl.conv2d_transpose(
-          enc6, num_masks + 1, [1, 1], strides=(1, 1),
-          name="convt7", padding="SAME", activation=None)
+      masks = tfl.conv2d(
+          enc6, filters=num_masks + 1, kernel_size=[1, 1],
+          strides=(1, 1), name="convt7", padding="SAME")
       masks = tf.reshape(
           tf.nn.softmax(tf.reshape(masks, [-1, num_masks + 1])),
           [batch_size,
@@ -470,7 +515,7 @@ def construct_predictive_tower(
 
       return output, lstm_state
 
-  def get_guassian_latent(self, latent_mean, latent_std):
+  def get_gaussian_latent(self, latent_mean, latent_std):
     latent = tf.random_normal(tf.shape(latent_mean), 0, 1, dtype=tf.float32)
     latent = latent_mean + tf.exp(latent_std / 2.0) * latent
     return latent
@@ -524,7 +569,7 @@ def construct_model(self,
       # Latent
       if self.hparams.stochastic_model:
         if timestep == 0 or self.hparams.multi_latent:
-          latent = self.get_guassian_latent(latent_mean, latent_std)
+          latent = self.get_gaussian_latent(latent_mean, latent_std)
 
       # Prediction
       pred_image, lstm_state = self.construct_predictive_tower(
@@ -704,7 +749,7 @@ def kl_divergence(self, mu, log_sigma):
     Returns:
       the KL loss.
     """
-
+    # TODO(mechcoder): Sum across all but the first dimension.
     return -.5 * tf.reduce_sum(
         1. + log_sigma - tf.square(mu) - tf.exp(log_sigma),
         axis=1)
@@ -756,9 +801,35 @@ def body(self, features):
     # TODO(mbz): what should it be if it"s undefined?
     if step_num is None:
       step_num = _LARGE_STEP_NUMBER
-    beta = tf.cond(tf.greater(step_num, self.hparams.num_iterations_2nd_stage),
-                   lambda: self.hparams.latent_loss_multiplier,
-                   lambda: 0.0)
+
+    schedule = self.hparams.latent_loss_multiplier_schedule
+    second_stage = self.hparams.num_iterations_2nd_stage
+    # TODO(mechcoder): Add log_annealing schedule.
+    if schedule == "constant":
+      beta = tf.cond(tf.greater(step_num, second_stage),
+                     lambda: self.hparams.latent_loss_multiplier,
+                     lambda: 0.0)
+    elif schedule == "linear_anneal":
+      # Linearly anneal beta from 0.0 to self.hparams.latent_loss_multiplier.
+      # between self.hparams.num_iterations_2nd_stage to anneal_end.
+      # beta = latent_loss * (1 - (global_step - 2nd_stage) / (anneal_end - 2nd_stage))  # pylint:disable=line-too-long
+      anneal_end = self.hparams.anneal_end
+      latent_multiplier = self.hparams.latent_loss_multiplier
+      if anneal_end < second_stage:
+        raise ValueError("Expected hparams.num_iterations_2nd_stage < "
+                         "hparams.anneal_end %d, got %d." %
+                         (second_stage, anneal_end))
+
+      def anneal_loss(step_num):
+        step_num = tf.cast(step_num, dtype=tf.float32)
+        fraction = (float(anneal_end) - step_num) / (anneal_end - second_stage)
+        return self.hparams.latent_loss_multiplier * (1 - fraction)
+
+      beta = tf.case(
+          pred_fn_pairs={
+              tf.less(step_num, second_stage): lambda: 0.0,
+              tf.greater(step_num, anneal_end): lambda: latent_multiplier},
+          default=lambda: anneal_loss(step_num))
 
     kl_loss = 0.0
     if self.is_training:
@@ -770,8 +841,7 @@ def body(self, features):
       tf.summary.scalar("beta", beta)
       tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
 
-    extra_loss = beta * kl_loss
-
+    extra_loss = beta * tf.reduce_mean(kl_loss)
     predictions = gen_images[hparams.video_num_input_frames-1:]
     reward_pred = tf.stack(
         gen_rewards[hparams.video_num_input_frames-1:], axis=1)
@@ -817,7 +887,7 @@ def construct_model(self, images, actions, rewards):
       # TODO(mbz): should we use input_image iunstead of image?
       latent_images = [image, images[timestep+1]]
       latent_mean, latent_std = self.construct_latent_tower(latent_images)
-      latent = self.get_guassian_latent(latent_mean, latent_std)
+      latent = self.get_gaussian_latent(latent_mean, latent_std)
       latent_means.append(latent_mean)
       latent_stds.append(latent_std)
 
@@ -877,14 +947,6 @@ def vgg_layer(self,
       net = activation(net)
     return net
 
-  def basic_lstm(self, inputs, state, num_units, name=None):
-    input_shape = common_layers.shape_list(inputs)
-    cell = tf.contrib.rnn.BasicLSTMCell(num_units, name=name)
-    if state is None:
-      state = cell.zero_state(input_shape[0], tf.float32)
-    outputs, new_state = cell(inputs, state)
-    return outputs, new_state
-
   def encoder(self, inputs, nout):
     """VGG based image encoder.
 
@@ -973,7 +1035,7 @@ def stacked_lstm(self, inputs, states, hidden_size, output_size, nlayers):
     net = tfl.dense(
         net, hidden_size, activation=None, name="af1")
     for i in range(nlayers):
-      net, states[i] = self.basic_lstm(
+      net, states[i] = basic_lstm(
           net, states[i], hidden_size, name="alstm%d"%i)
     net = tfl.dense(
         net, output_size, activation=tf.nn.tanh, name="af2")
@@ -996,7 +1058,7 @@ def lstm_gaussian(self, inputs, states, hidden_size, output_size, nlayers):
     net = inputs
     net = tfl.dense(net, hidden_size, activation=None, name="bf1")
     for i in range(nlayers):
-      net, states[i] = self.basic_lstm(
+      net, states[i] = basic_lstm(
           net, states[i], hidden_size, name="blstm%d"%i)
     mu = tfl.dense(net, output_size, activation=None, name="bf2mu")
     logvar = tfl.dense(net, output_size, activation=None, name="bf2log")
@@ -1146,6 +1208,7 @@ def next_frame_stochastic():
   hparams.add_hparam("latent_std_min", -5.0)
   hparams.add_hparam("num_iterations_2nd_stage", 10000)
   hparams.add_hparam("latent_loss_multiplier", 1e-3)
+  hparams.add_hparam("latent_loss_multiplier_schedule", "constant")
   hparams.add_hparam("multi_latent", False)
   hparams.add_hparam("relu_shift", 1e-12)
   hparams.add_hparam("dna_kernel_size", 5)
@@ -1154,6 +1217,8 @@ def next_frame_stochastic():
       "latent_num_frames",  # use all frames by default.
       hparams.video_num_input_frames + hparams.video_num_target_frames)
   hparams.add_hparam("tiny_mode", False)
+  hparams.add_hparam("anneal_end", 100000)
+  hparams.add_hparam("upsample_method", "conv2d_transpose")
   return hparams
 
 
@@ -1171,6 +1236,19 @@ def next_frame_stochastic_emily():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_savp():
+  """SVAP model."""
+  hparams = next_frame_stochastic()
+  hparams.add_hparam("z_dim", 8)
+  hparams.target_modality = "video:l1raw"
+  hparams.input_modalities = "inputs:video:l1raw"
+  hparams.latent_loss_multiplier_schedule = "linear_anneal"
+  hparams.anneal_end = 100000
+  hparams.upsample_method = "bilinear_upsample_conv"
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_stochastic_cutoff():
   """SV2P model with additional cutoff in L2 loss for environments like pong."""
@@ -1294,3 +1372,172 @@ def next_frame_ae_range(rhp):
   rhp.set_float("learning_rate_constant", 1., 2.)
   rhp.set_float("initializer_gain", 0.8, 1.5)
   rhp.set_int("filter_double_steps", 2, 3)
+
+
+@registry.register_model
+class NextFrameSavp(NextFrameStochastic):
+  """Stochastic Adversarial Video Prediction."""
+
+  def encoder(self, inputs, n_layers=3):
+    """COnvnet that encodes inputs into mean and std of a gaussian.
+
+    Args:
+     inputs: 5-D Tensor, shape (batch_size, num_frames, width, height, channels)
+     n_layers: Number of layers.
+
+    Returns:
+     z_mu: Mean of the latent gaussians.
+     z_log_var: log(var) of the latent gaussians.
+
+    Raises:
+      ValueError: If inputs is not a 5-D tensor or not float32.
+    """
+    latent_dims = self.hparams.z_dim
+
+    shape_as_list = inputs.shape.as_list()
+    if len(shape_as_list) != 5:
+      raise ValueError("Expected inputs to be a 5-D, got %d" %
+                       len(shape_as_list))
+    if inputs.dtype != tf.float32:
+      raise ValueError("Expected dtype tf.float32, got %s" % inputs.dtype)
+
+    # Flatten (N,T,W,H,C) into (NT,W,H,C)
+    batch_size, _ = shape_as_list[:2]
+    inputs = tf.reshape(inputs, [-1] + list(inputs.shape)[2:])
+    n_filters = 64
+    rectified = None
+
+    # Applies 3 layer conv-net with padding, instance normalization
+    # and leaky relu as per the encoder in
+    # https://github.com/alexlee-gk/video_prediction
+    padding = [[0, 0], [1, 1], [1, 1], [0, 0]]
+    for i in range(n_layers):
+      with tf.variable_scope("layer_%d" % (i + 1)):
+        n_filters *= 2**i
+        if i:
+          padded = tf.pad(rectified, padding)
+        else:
+          padded = tf.pad(inputs, padding)
+        convolved = tf.layers.conv2d(padded, filters=n_filters, kernel_size=4,
+                                     strides=2, padding="VALID")
+        normalized = tf.contrib.layers.instance_norm(convolved)
+        rectified = tf.nn.leaky_relu(normalized, alpha=0.2)
+
+    # Mean pooling across all spatial dimensions.
+    pooled = tf.nn.avg_pool(
+        rectified, [1] + rectified.shape[1:3].as_list() + [1],
+        strides=[1, 1, 1, 1], padding="VALID")
+    squeezed = tf.squeeze(pooled, [1, 2])
+
+    # Down-project and output the mean and log of the standard deviation of
+    # the latents.
+    with tf.variable_scope("z_mu"):
+      z_mu = tf.layers.dense(squeezed, latent_dims)
+    with tf.variable_scope("z_log_sigma_sq"):
+      z_log_var = tf.layers.dense(squeezed, latent_dims)
+      z_log_var = tf.clip_by_value(z_log_var, -10, 10)
+
+    # Reshape to (batch_size X num_frames X latent_dims)
+    z_mu = tf.reshape(z_mu, (batch_size, -1, latent_dims))
+    z_log_var = tf.reshape(
+        z_log_var, (batch_size, -1, latent_dims))
+    return z_mu, z_log_var
+
+  def construct_model(self, images, actions, rewards):
+    """Model that takes in images and returns predictions.
+
+    Args:
+      images: list of 4-D Tensors indexed by time.
+              (batch_size, width, height, channels)
+      actions: list of action tensors
+               each action should be in the shape ?x1xZ
+      rewards: list of reward tensors
+               each reward should be in the shape ?x1xZ
+
+    Returns:
+      video: list of 4-D predicted frames.
+      all_rewards: predicted rewards.
+      latent_means: list of gaussian means conditioned on the input at
+                    every frame.
+      latent_stds: list of gaussian stds conditioned on the input at
+                   every frame.
+    """
+    latent_dims = self.hparams.z_dim
+    context_frames = self.hparams.video_num_input_frames
+    seq_len = len(images)
+    input_shape = common_layers.shape_list(images[0])
+    batch_size = input_shape[0]
+
+    # Model does not support reward-conditioned frame generation.
+    fake_rewards = rewards[:-1]
+
+    # Concatenate x_{t-1} and x_{t} along depth and encode it to
+    # produce the mean and standard deviation of z_{t-1}
+    image_pairs = tf.concat([images[:seq_len - 1],
+                             images[1:seq_len]], axis=-1)
+
+    z_mu, z_log_sigma_sq = self.encoder(image_pairs)
+    # Unstack z_mu and z_log_sigma_sq along the time dimension.
+    z_mu = tf.unstack(z_mu, axis=0)
+    z_log_sigma_sq = tf.unstack(z_log_sigma_sq, axis=0)
+    iterable = zip(images[:-1], actions[:-1], fake_rewards,
+                   z_mu, z_log_sigma_sq)
+
+    # Initialize LSTM State
+    lstm_state = [None] * 7
+    gen_cond_video, gen_prior_video, all_rewards, latent_means, latent_stds = \
+      [], [], [], [], []
+    pred_image, prior_latent_state, cond_latent_state = None, None, None
+    train_mode = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
+
+    with tf.variable_scope("prediction", reuse=tf.AUTO_REUSE):
+
+      for step, (image, action, reward, mu, log_sigma_sq) in enumerate(iterable):  # pylint:disable=line-too-long
+        # Sample latents using a gaussian centered at conditional mu and std.
+        latent = self.get_gaussian_latent(mu, log_sigma_sq)
+
+        # Sample prior latents from isotropic normal distribution.
+        prior_latent = tf.random_normal(tf.shape(latent), dtype=tf.float32)
+
+        # LSTM that encodes correlations between conditional latents.
+        # Pg 22 in https://arxiv.org/pdf/1804.01523.pdf
+        enc_cond_latent, cond_latent_state = basic_lstm(
+            latent, cond_latent_state, latent_dims, name="cond_latent")
+
+        # LSTM that encodes correlations between prior latents.
+        enc_prior_latent, prior_latent_state = basic_lstm(
+            prior_latent, prior_latent_state, latent_dims, name="prior_latent")
+
+        # Scheduled Sampling
+        done_warm_start = step > context_frames - 1
+        groundtruth_items = [image]
+        generated_items = [pred_image]
+        input_image = self.get_scheduled_sample_inputs(
+            done_warm_start, groundtruth_items, generated_items, batch_size)[0]
+
+        all_latents = tf.concat([enc_cond_latent, enc_prior_latent], axis=0)
+        all_image = tf.concat([input_image, input_image], axis=0)
+        all_action = tf.concat([action, action], axis=0)
+        all_rewards = tf.concat([reward, reward], axis=0)
+
+        all_pred_images, lstm_state = self.construct_predictive_tower(
+            all_image, all_rewards, all_action, lstm_state, all_latents,
+            concat_latent=True)
+
+        cond_pred_images, prior_pred_images = \
+          all_pred_images[:batch_size], all_pred_images[batch_size:]
+
+        if train_mode:
+          pred_image = cond_pred_images
+        else:
+          pred_image = prior_pred_images
+
+        gen_cond_video.append(cond_pred_images)
+        gen_prior_video.append(prior_pred_images)
+        latent_means.append(mu)
+        latent_stds.append(log_sigma_sq)
+
+    if train_mode:
+      return gen_cond_video, fake_rewards, latent_means, latent_stds
+    else:
+      return gen_prior_video, fake_rewards, latent_means, latent_stds
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index 373f159a0..39d7b7d57 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -33,14 +33,15 @@ def TestVideoModel(self,
                      out_frames,
                      hparams,
                      model,
-                     expected_last_dim):
+                     expected_last_dim,
+                     upsample_method="conv2d_transpose"):
 
     x = np.random.random_integers(0, high=255, size=(8, in_frames, 64, 64, 3))
     y = np.random.random_integers(0, high=255, size=(8, out_frames, 64, 64, 3))
 
     hparams.video_num_input_frames = in_frames
     hparams.video_num_target_frames = out_frames
-
+    hparams.upsample_method = upsample_method
     problem = registry.problem("video_stochastic_shapes10k")
     p_hparams = problem.get_hparams(hparams)
     hparams.problem = problem
@@ -65,6 +66,12 @@ def TestOnVariousInputOutputSizes(self, hparams, model, expected_last_dim):
     self.TestVideoModel(4, 1, hparams, model, expected_last_dim)
     self.TestVideoModel(7, 5, hparams, model, expected_last_dim)
 
+  def TestOnVariousUpSampleLayers(self, hparams, model, expected_last_dim):
+    self.TestVideoModel(4, 1, hparams, model, expected_last_dim,
+                        upsample_method="bilinear_upsample_conv")
+    self.TestVideoModel(4, 1, hparams, model, expected_last_dim,
+                        upsample_method="nn_upsample_conv")
+
   def testBasic(self):
     self.TestOnVariousInputOutputSizes(
         next_frame.next_frame(),
@@ -89,6 +96,15 @@ def testStochasticEmily(self):
         next_frame.NextFrameStochasticEmily,
         1)
 
+  def testStochasticSavp(self):
+    self.TestOnVariousInputOutputSizes(
+        next_frame.next_frame_savp(),
+        next_frame.NextFrameSavp,
+        1)
+    self.TestOnVariousUpSampleLayers(
+        next_frame.next_frame_savp(),
+        next_frame.NextFrameSavp,
+        1)
 
 if __name__ == "__main__":
   tf.test.main()

From 70b46c3ae8196d050549e87c7f6aa90574e5afb8 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 23 Jul 2018 17:49:07 -0700
Subject: [PATCH 0373/2720] Allowing extra outputs while removing all
 non-Tensors.

PiperOrigin-RevId: 205755617
---
 tensor2tensor/utils/t2t_model.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index e504b9866..72d60e3c9 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1363,7 +1363,8 @@ def estimator_spec_predict(self, features, use_tpu=False):
         "targets": features.get("infer_targets"),
         "batch_prediction_key": features.get("batch_prediction_key"),
     }
-    _del_dict_nones(predictions)
+
+    _del_dict_non_tensors(predictions)
 
     export_out = {"outputs": predictions["outputs"]}
     if "scores" in predictions:
@@ -1595,9 +1596,9 @@ def host_call_fn(**kwargs):
   return (host_call_fn, summary_kwargs)
 
 
-def _del_dict_nones(d):
+def _del_dict_non_tensors(d):
   for k in list(d.keys()):
-    if d[k] is None:
+    if not isinstance(d[k], tf.Tensor):
       del d[k]
 
 

From c29de87c6951adb1ef7b5986e70f898e0d335a19 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Mon, 23 Jul 2018 20:29:42 -0700
Subject: [PATCH 0374/2720] Fixing data concatenation in MultiProblem eval.

PiperOrigin-RevId: 205769079
---
 tensor2tensor/data_generators/multi_problem.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 183bfb331..5a0611553 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -109,6 +109,6 @@ def dataset(self,
     else:
       single_mtl_dataset = datasets[0]
       for data in datasets[1:]:
-        single_mtl_dataset.concatenate(data)
+        single_mtl_dataset = single_mtl_dataset.concatenate(data)
 
     return single_mtl_dataset

From 4a5db4ac52f89dc237eced8722a059fa2ab06d70 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 23 Jul 2018 20:30:47 -0700
Subject: [PATCH 0375/2720] Travis fixes

PiperOrigin-RevId: 205769187
---
 tensor2tensor/bin/t2t_trainer.py              | 1 -
 tensor2tensor/data_generators/wikisum/html.py | 7 ++++---
 tensor2tensor/models/research/next_frame.py   | 7 ++++++-
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 8d57476c0..5f7e36caa 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -31,7 +31,6 @@
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
 import tensorflow as tf
-from tensorflow.contrib.tpu.python.tpu import tpu_config
 
 
 flags = tf.flags
diff --git a/tensor2tensor/data_generators/wikisum/html.py b/tensor2tensor/data_generators/wikisum/html.py
index d27cf28fa..32c2b82d5 100644
--- a/tensor2tensor/data_generators/wikisum/html.py
+++ b/tensor2tensor/data_generators/wikisum/html.py
@@ -22,7 +22,7 @@ def get_text_from_html(html):
 
   try:
     soup = bs4.BeautifulSoup(html, "html.parser")
-  except:
+  except:  # pylint: disable=bare-except
     # Some docs don't parse
     return ""
   # Remove script and style tags
@@ -32,6 +32,7 @@ def get_text_from_html(html):
 
 
 def _soup_strings(soup):
+  """Return text strings in soup."""
   paragraph_tags = set([
       "caption", "details", "h1", "h2", "h3", "h4", "h5", "h6", "li", "p", "td",
       "div", "span"
@@ -43,8 +44,8 @@ def _soup_strings(soup):
     # children (see below).
     if skip_children is not None:
       try:
-        in_skip = descendant in skip_children
-      except RecursionError:
+        in_skip = descendant in skip_children  # pylint: disable=unsupported-membership-test
+      except RecursionError:  # pylint: disable=undefined-variable
         # Possible for this check to hit a nasty infinite recursion because of
         # BeautifulSoup __eq__ checks.
         in_skip = True
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index ffc4da732..9f42ed645 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -27,7 +27,12 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow as tf
-from tensorflow_models.slim.nets.cyclegan import cyclegan_upsample
+
+try:
+  from tensorflow_models.slim.nets.cyclegan import cyclegan_upsample  # pylint: disable=g-import-not-at-top
+except ImportError:
+  pass
+
 tfl = tf.layers
 tfcl = tf.contrib.layers
 

From 6c8f435db09bdcc8509c36d6f2d22f3a0004b7cd Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 24 Jul 2018 10:43:33 -0700
Subject: [PATCH 0376/2720] summarize variables

PiperOrigin-RevId: 205853262
---
 tensor2tensor/utils/optimize.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 970914e22..c405c8f3f 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -33,6 +33,8 @@ def optimize(loss, learning_rate, hparams, use_tpu=False):
   loss = weight_decay_and_noise(loss, hparams, learning_rate)
   loss = tf.identity(loss, name="total_loss")
   log_variable_sizes(verbose=hparams.summarize_vars)
+  if hparams.summarize_vars:
+    summarize_variables()
   diet_vars = [
       v for v in tf.global_variables() if v.dtype == dtypes.float16_ref
   ]
@@ -216,6 +218,24 @@ def log_variable_sizes(var_list=None, tag=None, verbose=False):
   tf.logging.info("%s Total size: %d", tag, total_size)
 
 
+def summarize_variables(var_list=None, tag=None):
+  """Summarize the variables.
+
+  Args:
+    var_list: a list of variables; defaults to trainable_variables.
+    tag: name scope of the summary; defaults to training_variables/.
+  """
+  if var_list is None:
+    var_list = tf.trainable_variables()
+  if tag is None:
+    tag = "training_variables/"
+
+  name_to_var = {v.name: v for v in var_list}
+  for v_name in list(name_to_var):
+    v = name_to_var[v_name]
+    tf.summary.histogram(tag + v_name, v)
+
+
 def get_variable_initializer(hparams):
   """Get variable initializer from hparams."""
   if not hparams.initializer:

From 2fefad913de198c4bf9ceb4f7bddb879521d1a85 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Tue, 24 Jul 2018 14:19:59 -0700
Subject: [PATCH 0377/2720] Add MSR Paraphrase Corpus.

PiperOrigin-RevId: 205890541
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/mrpc.py         | 145 ++++++++++++++++++
 2 files changed, 146 insertions(+)
 create mode 100644 tensor2tensor/data_generators/mrpc.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 5a1ffedce..fe816b683 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -48,6 +48,7 @@
     "tensor2tensor.data_generators.lm1b",
     "tensor2tensor.data_generators.lm1b_imdb",
     "tensor2tensor.data_generators.mnist",
+    "tensor2tensor.data_generators.mrpc",
     "tensor2tensor.data_generators.mscoco",
     "tensor2tensor.data_generators.multinli",
     "tensor2tensor.data_generators.program_search",
diff --git a/tensor2tensor/data_generators/mrpc.py b/tensor2tensor/data_generators/mrpc.py
new file mode 100644
index 000000000..3c441d285
--- /dev/null
+++ b/tensor2tensor/data_generators/mrpc.py
@@ -0,0 +1,145 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generators for the MSR Paraphrase Corpus."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import six
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+EOS = text_encoder.EOS
+
+
+@registry.register_problem
+class MSRParaphraseCorpus(text_problems.TextConcat2ClassProblem):
+  """MSR Paraphrase Identification problems."""
+
+  # Link to data from GLUE: https://gluebenchmark.com/tasks
+  DEV_IDS = ("https://firebasestorage.googleapis.com/v0/b/"
+             "mtl-sentence-representations.appspot.com/o/"
+             "data%2FWNLI.zip?alt=media&token=068ad0a0-ded7-"
+             "4bd7-99a5-5e00222e0faf")
+  MRPC_TRAIN = ("https://s3.amazonaws.com/senteval/senteval_data/"
+                "msr_paraphrase_train.txt")
+  MRPC_TEST = ("https://s3.amazonaws.com/senteval/senteval_data/"
+               "msr_paraphrase_test.txt")
+  DATA_DIR = "MRPC"
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 10,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  @property
+  def approx_vocab_size(self):
+    return 2**13  # 8k vocab suffices for this small dataset.
+
+  @property
+  def vocab_filename(self):
+    return "vocab.mrpc.%d" % self.approx_vocab_size
+
+  @property
+  def num_classes(self):
+    return 2
+
+  @property
+  def concat_token(self):
+    return "<SENT_SEP>"
+
+  @property
+  def concat_id(self):
+    if self.vocab_type == text_problems.VocabType.CHARACTER:
+      return problem.TaskID.EN_SENT_PAIR
+    return 2
+
+  def class_labels(self, data_dir):
+    del data_dir
+    return ["not_paraphrase", "paraphrase"]
+
+  def _maybe_download_corpora(self, tmp_dir):
+    mrpc_dir = os.path.join(tmp_dir, self.DATA_DIR)
+    tf.gfile.MakeDirs(mrpc_dir)
+    mrpc_train_finalpath = os.path.join(mrpc_dir, "msr_paraphrase_train.txt")
+    mrpc_test_finalpath = os.path.join(mrpc_dir, "msr_paraphrase_test.txt")
+    mrpc_dev_ids_finalpath = os.path.join(mrpc_dir, "dev_ids.tsv")
+
+    def download_file(tdir, filepath, url):
+      if not tf.gfile.Exists(filepath):
+        generator_utils.maybe_download(tdir, filepath, url)
+
+    download_file(mrpc_dir, mrpc_train_finalpath, self.MRPC_TRAIN)
+    download_file(mrpc_dir, mrpc_test_finalpath, self.MRPC_TEST)
+    download_file(mrpc_dir, mrpc_dev_ids_finalpath, self.DEV_IDS)
+
+    return mrpc_dir
+
+  def example_generator(self, filename, dev_ids):
+    for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
+      if idx == 0: continue  # skip header
+      if six.PY2:
+        line = unicode(line.strip(), "utf-8")
+      else:
+        line = line.strip().decode("utf-8")
+      l, id1, id2, s1, s2 = line.split("\t")
+      if dev_ids and [id1, id2] not in dev_ids:
+        continue
+      inputs = [[s1, s2], [s2, s1]]
+      for inp in inputs:
+        yield {
+            "inputs": inp,
+            "label": int(l)
+        }
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    mrpc_dir = self._maybe_download_corpora(tmp_dir)
+    filesplit = "msr_paraphrase_train.txt"
+    dev_ids = []
+    if dataset_split != problem.DatasetSplit.TRAIN:
+      for row in tf.gfile.Open(os.path.join(mrpc_dir, "dev_ids.tsv")):
+        dev_ids.append(row.strip().split("\t"))
+
+    filename = os.path.join(mrpc_dir, filesplit)
+    for example in self.example_generator(filename, dev_ids):
+      yield example
+
+
+@registry.register_problem
+class MSRParaphraseCorpusCharacters(MSRParaphraseCorpus):
+  """MSR Paraphrase Identification problems, character level"""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  @property
+  def task_id(self):
+    return problem.TaskID.EN_SIM

From a89306da94ab84584e3b3bd3fe98e2f110286392 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 24 Jul 2018 14:24:56 -0700
Subject: [PATCH 0378/2720] Allow tile_and_concat to handle dynamic input.

PiperOrigin-RevId: 205891344
---
 tensor2tensor/models/research/next_frame.py   | 43 +++++++++++++------
 .../models/research/next_frame_test.py        | 30 +++++++++++++
 2 files changed, 60 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 9f42ed645..bca4ef8de 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -194,16 +194,25 @@ def tinyify(self, array):
       return [1 for _ in array]
     return array
 
-  def tile_and_concat(self, image, latent, concat_latent=True):
-    """Tile latent and concatenate to image across depth."""
+  @staticmethod
+  def tile_and_concat(image, latent, concat_latent=True):
+    """Tile latent and concatenate to image across depth.
+
+    Args:
+      image: 4-D Tensor, (batch_size X height X width X channels)
+      latent: 2-D Tensor, (batch_size X latent_dims)
+      concat_latent: If set to False, the image is returned as is.
+
+    Returns:
+      concat_latent: 4-D Tensor, (batch_size X height X width X channels+1)
+        latent tiled and concatenated to the image across the channels.
+    """
     if not concat_latent:
       return image
-    height, width = image.shape[1:3].as_list()
-    tf.logging.info("Height")
-    tf.logging.info(height)
-    tf.logging.info("Width")
-    tf.logging.info(width)
-    _, latent_dims = latent.shape.as_list()
+    image_shape = common_layers.shape_list(image)
+    latent_shape = common_layers.shape_list(latent)
+    height, width = image_shape[1], image_shape[2]
+    latent_dims = latent_shape[1]
 
     height_multiples = height // latent_dims
     pad = height - (height_multiples * latent_dims)
@@ -397,13 +406,19 @@ def decode_to_shape(self, inputs, shape, scope):
       return x
 
   def conv_lstm_2d(self, inputs, state, output_channels,
-                   kernel_size=5, name=None):
+                   kernel_size=5, name=None, spatial_dims=None):
     input_shape = common_layers.shape_list(inputs)
+    batch_size, input_channels = input_shape[0], input_shape[-1]
+    if spatial_dims is None:
+      input_shape = input_shape[1:]
+    else:
+      input_shape = spatial_dims + [input_channels]
+
     cell = tf.contrib.rnn.ConvLSTMCell(
-        2, input_shape[1:], output_channels,
+        2, input_shape, output_channels,
         [kernel_size, kernel_size], name=name)
     if state is None:
-      state = cell.zero_state(input_shape[0], tf.float32)
+      state = cell.zero_state(batch_size, tf.float32)
     outputs, new_state = cell(inputs, state)
     return outputs, new_state
 
@@ -439,7 +454,8 @@ def construct_predictive_tower(
       enc4 = self.tile_and_concat(enc4, latent, concat_latent=concat_latent)
 
       hidden6, lstm_state[5] = lstm_func(
-          enc4, lstm_state[5], lstm_size[5], name="state6")  # 16x16
+          enc4, lstm_state[5], lstm_size[5], name="state6",
+          spatial_dims=enc1_shape[1:-1])  # 16x16
       hidden6 = self.tile_and_concat(
           hidden6, latent, concat_latent=concat_latent)
       hidden6 = tfcl.layer_norm(hidden6, scope="layer_norm7")
@@ -456,7 +472,8 @@ def construct_predictive_tower(
       enc5 = self.tile_and_concat(enc5, latent, concat_latent=concat_latent)
 
       hidden7, lstm_state[6] = lstm_func(
-          enc5, lstm_state[6], lstm_size[6], name="state7")  # 32x32
+          enc5, lstm_state[6], lstm_size[6], name="state7",
+          spatial_dims=enc0_shape[1:-1])  # 32x32
       hidden7 = tfcl.layer_norm(hidden7, scope="layer_norm8")
 
       # Skip connection.
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index 39d7b7d57..45dfed9c3 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -106,5 +106,35 @@ def testStochasticSavp(self):
         next_frame.NextFrameSavp,
         1)
 
+  def testDynamicTileAndConcat(self):
+    with tf.Graph().as_default():
+      # image = (1 X 4 X 4 X 1)
+      image = [[1, 2, 3, 4],
+               [2, 4, 5, 6],
+               [7, 8, 9, 10],
+               [7, 9, 10, 1]]
+      image = tf.expand_dims(tf.expand_dims(image, axis=0), axis=-1)
+      image_t = tf.cast(tf.convert_to_tensor(image), dtype=tf.float32)
+
+      # latent = (1 X 2)
+      latent = np.array([[90, 100]])
+      latent_t = tf.cast(tf.convert_to_tensor(latent), dtype=tf.float32)
+
+      with tf.Session() as session:
+        tiled = next_frame.NextFrameStochastic.tile_and_concat(
+            image_t, latent_t)
+        tiled_np = session.run(tiled)
+        tiled_latent = tiled_np[0, :, :, -1]
+        self.assertAllEqual(tiled_np.shape, (1, 4, 4, 2))
+
+        self.assertAllEqual(tiled_np[:, :, :, :1], image)
+        self.assertAllEqual(
+            tiled_latent,
+            [[90, 90, 90, 90],
+             [100, 100, 100, 100],
+             [90, 90, 90, 90],
+             [100, 100, 100, 100]])
+
+
 if __name__ == "__main__":
   tf.test.main()

From a61bb785f45e0707d440d9f7d55045d0266b8e0d Mon Sep 17 00:00:00 2001
From: cbockman <c.bockman@gmail.com>
Date: Tue, 24 Jul 2018 11:53:50 -1000
Subject: [PATCH 0379/2720] minor docs fix to adam (#936)

Removed note about re-scaling the lr with Adam.

I believe this is a legacy of the old way LR+Adam was handled and that the docs weren't updated when the code was updated.

Pushing a PR because this is confusing to read now.  :)  (Unless I'm wrong!)
---
 tensor2tensor/utils/optimize.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index c405c8f3f..7249a76d8 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -80,7 +80,7 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
     tf.logging.info("Using optimizer %s", optimizer_name)
 
     if optimizer_name == "Adam":
-      # We change the default epsilon for Adam and re-scale lr.
+      # We change the default epsilon for Adam.
       # Using LazyAdam as it's much faster for large vocabulary embeddings.
       self._opt = tf.contrib.opt.LazyAdamOptimizer(
           lr,

From d5b748caac8fcba6adffaa2a9c95c807ed6dbaea Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 24 Jul 2018 16:39:14 -0700
Subject: [PATCH 0380/2720] add l2 normalization and zero_add

PiperOrigin-RevId: 205913392
---
 tensor2tensor/layers/common_layers.py | 35 ++++++++++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 0bb726742..a48784f7b 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -695,6 +695,23 @@ def noam_norm(x, epsilon=1.0, name=None):
         tf.to_float(shape[-1])))
 
 
+def l2_norm(x, filters=None, epsilon=1e-6, name=None, reuse=None):
+  """Layer normalization with l2 norm."""
+  if filters is None:
+    filters = shape_list(x)[-1]
+  with tf.variable_scope(
+      name, default_name="l2_norm", values=[x], reuse=reuse):
+    scale = tf.get_variable(
+        "l2_norm_scale", [filters], initializer=tf.ones_initializer())
+    bias = tf.get_variable(
+        "l2_norm_bias", [filters], initializer=tf.zeros_initializer())
+    epsilon, scale, bias = [cast_like(t, x) for t in [epsilon, scale, bias]]
+    mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
+    l2norm = tf.reduce_sum(tf.square(x - mean), axis=[-1], keepdims=True)
+    norm_x = (x - mean) * tf.rsqrt(l2norm + epsilon)
+    return norm_x * scale + bias
+
+
 def apply_norm(x, norm_type, depth, epsilon):
   """Apply Normalization."""
   if norm_type == "layer":
@@ -705,10 +722,21 @@ def apply_norm(x, norm_type, depth, epsilon):
     return tf.layers.batch_normalization(x, epsilon=epsilon)
   if norm_type == "noam":
     return noam_norm(x, epsilon)
+  if norm_type == "l2":
+    return l2_norm(x, filters=depth, epsilon=epsilon)
   if norm_type == "none":
     return x
   raise ValueError("Parameter normalizer_fn must be one of: 'layer', 'batch',"
-                   "'noam', 'none'.")
+                   "'noam', 'lr', 'none'.")
+
+
+def zero_add(previous_value, x, name=None, reuse=False):
+  """Resnet connection with zero initialization."""
+  with tf.name_scope(
+      name, default_name="zero_add", values=[previous_value, x], reuse=reuse):
+    gamma = tf.get_variable(
+        "gamma", [None], initializer=tf.zeros_initializer())
+    return previous_value + gamma*x
 
 
 def layer_prepostprocess(previous_value,
@@ -728,6 +756,7 @@ def layer_prepostprocess(previous_value,
     a: add previous_value
     n: apply normalization
     d: apply dropout
+    z: zero add
 
   For example, if sequence=="dna", then the output is
     previous_value + normalize(dropout(x))
@@ -755,6 +784,8 @@ def layer_prepostprocess(previous_value,
     for c in sequence:
       if c == "a":
         x += previous_value
+      elif c == "z":
+        x = zero_add(previous_value, x)
       elif c == "n":
         x = apply_norm(x, norm_type, depth, epsilon)
       else:
@@ -787,6 +818,8 @@ def layer_preprocess(layer_input, hparams):
   """
   assert "a" not in hparams.layer_preprocess_sequence, (
       "No residual connections allowed in hparams.layer_preprocess_sequence")
+  assert "z" not in hparams.layer_preprocess_sequence, (
+      "No residual connections allowed in hparams.layer_preprocess_sequence")
   return layer_prepostprocess(
       None,
       layer_input,

From 55ad83230d946a8027b8bee57a537de78896a3e4 Mon Sep 17 00:00:00 2001
From: lhlmgr <lhlmgr@users.noreply.github.com>
Date: Wed, 25 Jul 2018 02:04:39 +0200
Subject: [PATCH 0381/2720] Update cifar.py (#935)

With `Python 3.6.5` the following exception will be thrown:

```
  File "/home/lhlmgr/anaconda3/lib/python3.6/site-packages/tensor2tensor/data_generators/cifar.py", line 146, in cifar_generator
    data = cPickle.load(f)
UnicodeDecodeError: 'ascii' codec can't decode byte 0xff in position 0: ordinal not in range(128)
```
Adding the the parameter `encoding='latin1'` fixes this error for `cifar10` and `cifar100`
---
 tensor2tensor/data_generators/cifar.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index dcdacb64a..efd36bb84 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -92,7 +92,7 @@ def cifar_generator(cifar_version, tmp_dir, training, how_many, start_from=0):
   for filename in data_files:
     path = os.path.join(tmp_dir, prefix, filename)
     with tf.gfile.Open(path, "rb") as f:
-      data = cPickle.load(f)
+      data = cPickle.load(f, encoding='latin1')
     images = data["data"]
     num_images = images.shape[0]
     images = images.reshape((num_images, 3, image_size, image_size))

From 27645c21b5507fa2b8e43bdf68b2533c383e4fdf Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 24 Jul 2018 17:17:19 -0700
Subject: [PATCH 0382/2720] internal merge of PR #935

PiperOrigin-RevId: 205918738
---
 tensor2tensor/data_generators/cifar.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index efd36bb84..f19deb84c 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -21,6 +21,7 @@
 import os
 import tarfile
 import numpy as np
+import six
 
 from six.moves import cPickle
 
@@ -92,7 +93,10 @@ def cifar_generator(cifar_version, tmp_dir, training, how_many, start_from=0):
   for filename in data_files:
     path = os.path.join(tmp_dir, prefix, filename)
     with tf.gfile.Open(path, "rb") as f:
-      data = cPickle.load(f, encoding='latin1')
+      if six.PY2:
+        data = cPickle.load(f)
+      else:
+        data = cPickle.load(f, encoding="latin1")
     images = data["data"]
     num_images = images.shape[0]
     images = images.reshape((num_images, 3, image_size, image_size))

From c49d05d390d645e3ed09b0bfb7771c0438cbc1b9 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 24 Jul 2018 17:32:37 -0700
Subject: [PATCH 0383/2720] More RL experiments.

PiperOrigin-RevId: 205920441
---
 tensor2tensor/rl/collect.py             | 22 ++++++++++++++++++----
 tensor2tensor/rl/model_rl_experiment.py |  4 ++--
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 69776f8b1..515bb8fce 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -48,6 +48,7 @@ class _MemoryWrapper(WrapperBase):
 
   def __init__(self, batch_env):
     super(_MemoryWrapper, self).__init__(batch_env)
+    self._is_simple = False  # TODO(lukaszkaiser): why do we need it? mbz?
     infinity = 10000000
     meta_data = list(zip(*_rollout_metadata(batch_env)))
     # In memory wrapper we do not collect pdfs neither value_function
@@ -55,6 +56,11 @@ def __init__(self, batch_env):
     shapes = meta_data[0][:4]
     dtypes = meta_data[1][:4]
     self.speculum = tf.FIFOQueue(infinity, shapes=shapes, dtypes=dtypes)
+    observs_shape = batch_env.observ.shape
+    # TODO(piotrmilos): possibly retrieve the observation type for batch_env
+    observ_dtype = tf.float32
+    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
+                               trainable=False)
 
   def simulate(self, action):
 
@@ -64,11 +70,19 @@ def simulate(self, action):
 
     reward, done = self._batch_env.simulate(action)
 
-    with tf.control_dependencies([reward, done]):
-      enqueue_op = self.speculum.enqueue(
-          [self._batch_env.observ, reward, done, action])
+    if self._is_simple:
+      with tf.control_dependencies([reward, done]):
+        enqueue_op = self.speculum.enqueue(
+            [self._batch_env.observ, reward, done, action])
+    else:
+      with tf.control_dependencies([reward, done]):
+        enqueue_op = self.speculum.enqueue(
+            [self._observ.read_value(), reward, done, action])
 
-    with tf.control_dependencies([enqueue_op]):
+      with tf.control_dependencies([enqueue_op]):
+        assign = self._observ.assign(self._batch_env.observ)
+
+    with tf.control_dependencies([assign]):
       return tf.identity(reward), tf.identity(done)
 
 
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index cb9c95ffe..19f6dca77 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -490,7 +490,7 @@ def rl_modelrl_base():
       generative_model_params="next_frame",
       ppo_params="ppo_pong_base",
       autoencoder_train_steps=0,
-      model_train_steps=50000,
+      model_train_steps=100000,
       simulated_env_generator_num_steps=2000,
       simulation_random_starts=True,
       intrinsic_reward_scale=0.,
@@ -503,7 +503,7 @@ def rl_modelrl_base():
       # though it is not necessary.
       ppo_epoch_length=60,
       ppo_num_agents=16,
-      ppo_learning_rate=0.,  # Will be changed, just so it exists.
+      ppo_learning_rate=2e-4,  # Will be changed, just so it exists.
       # Whether the PPO agent should be restored from the previous iteration, or
       # should start fresh each time.
       ppo_continue_training=True,

From 1c327cdcaa83e95add639cbec5f1a713bdd2ad46 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 25 Jul 2018 16:52:27 -0700
Subject: [PATCH 0384/2720] Copy cyclegan_upsample from tensorflow/models to
 common_layers, since that repo isn't pip installable and can't be linked to
 T2T.

PiperOrigin-RevId: 206078743
---
 tensor2tensor/layers/common_layers.py       | 65 +++++++++++++++++++++
 tensor2tensor/layers/common_layers_test.py  | 59 +++++++++++++++++++
 tensor2tensor/models/research/next_frame.py | 11 +---
 3 files changed, 127 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index a48784f7b..23e11d2b3 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3548,3 +3548,68 @@ def tpu_safe_image_summary(image):
   else:
     image = tf.cast(image, tf.uint8)
   return image
+
+
+# This has been (shamefully) copied from
+# GitHub tensorflow/models/blob/master/research/slim/nets/cyclegan.py
+#
+# tensorflow/models cannot be pip installed, and even if it were we don't want
+# to depend on all the models in it.
+#
+# Therefore copying and forgoing any more bugfixes into it is the most
+# expedient way to use this function.
+def cyclegan_upsample(net, num_outputs, stride, method="conv2d_transpose"):
+  """Upsamples the given inputs.
+
+  Args:
+    net: A Tensor of size [batch_size, height, width, filters].
+    num_outputs: The number of output filters.
+    stride: A list of 2 scalars or a 1x2 Tensor indicating the scale,
+      relative to the inputs, of the output dimensions. For example, if kernel
+      size is [2, 3], then the output height and width will be twice and three
+      times the input size.
+    method: The upsampling method: 'nn_upsample_conv',
+      'bilinear_upsample_conv', or 'conv2d_transpose'.
+
+  Returns:
+    A Tensor which was upsampled using the specified method.
+
+  Raises:
+    ValueError: if `method` is not recognized.
+  """
+
+  with tf.variable_scope("upconv"):
+    net_shape = tf.shape(net)
+    height = net_shape[1]
+    width = net_shape[2]
+
+    # Reflection pad by 1 in spatial dimensions (axes 1, 2 = h, w) to make a
+    # 3x3 "valid" convolution produce an output with the same dimension as the
+    # input.
+    spatial_pad_1 = np.array([[0, 0], [1, 1], [1, 1], [0, 0]])
+
+    if method == "nn_upsample_conv":
+      net = tf.image.resize_nearest_neighbor(
+          net, [stride[0] * height, stride[1] * width])
+      net = tf.pad(net, spatial_pad_1, "REFLECT")
+      net = tf.contrib.layers.conv2d(net, num_outputs, kernel_size=[3, 3],
+                                     padding="valid")
+    elif method == "bilinear_upsample_conv":
+      net = tf.image.resize_bilinear(
+          net, [stride[0] * height, stride[1] * width])
+      net = tf.pad(net, spatial_pad_1, "REFLECT")
+      net = tf.contrib.layers.conv2d(net, num_outputs, kernel_size=[3, 3],
+                                     padding="valid")
+    elif method == "conv2d_transpose":
+      # This corrects 1 pixel offset for images with even width and height.
+      # conv2d is left aligned and conv2d_transpose is right aligned for even
+      # sized images (while doing "SAME" padding).
+      # Note: This doesn"t reflect actual model in paper.
+      net = tf.contrib.layers.conv2d_transpose(
+          net, num_outputs, kernel_size=[3, 3], stride=stride,
+          padding="valid")
+      net = net[:, 1:, 1:, :]
+    else:
+      raise ValueError("Unknown method: [%s]" % method)
+
+    return net
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index bc86cfdff..d24193769 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -627,6 +627,65 @@ def testConvHiddenReluMemoryEfficient(self):
     self.assertAllClose(dnorm_bias, dnorm_bias_f)
     self.assertAllClose(dx, dx_f)
 
+  def testCycleGANUpsampleNnUpsampleConv(self):
+    batch = 8
+    height = 32
+    width = 32
+    num_channels = 3
+    output_filters = 10
+    stride = [2, 3]  # we want height to be x2 and width to be x3
+    random_input = np.random.rand(batch, height, width, num_channels)
+
+    # nn_upsample_conv gives exactly the shapes we'd expect.
+    upsampled_output = common_layers.cyclegan_upsample(
+        random_input, output_filters, stride, "nn_upsample_conv")
+    upsampled_output_shape = common_layers.shape_list(upsampled_output)
+    self.assertEqual(batch, upsampled_output_shape[0])
+    self.assertEqual(height * stride[0], upsampled_output_shape[1])
+    self.assertEqual(width * stride[1], upsampled_output_shape[2])
+    self.assertEqual(output_filters, upsampled_output_shape[3])
+
+  def testCycleGANUpsampleBilinearUpsampleConv(self):
+    batch = 8
+    height = 32
+    width = 32
+    num_channels = 3
+    output_filters = 10
+    stride = [2, 3]  # we want height to be x2 and width to be x3
+    random_input = np.random.rand(batch, height, width, num_channels)
+
+    # bilinear_upsample_conv gives exactly the shapes we'd expect.
+    upsampled_output = common_layers.cyclegan_upsample(
+        random_input, output_filters, stride, "bilinear_upsample_conv")
+    upsampled_output_shape = common_layers.shape_list(upsampled_output)
+    self.assertEqual(batch, upsampled_output_shape[0])
+    self.assertEqual(height * stride[0], upsampled_output_shape[1])
+    self.assertEqual(width * stride[1], upsampled_output_shape[2])
+    self.assertEqual(output_filters, upsampled_output_shape[3])
+
+  def testCycleGANUpsampleConv2dTranspose(self):
+    batch = 8
+    height = 32
+    width = 32
+    num_channels = 3
+    output_filters = 10
+    stride = [2, 3]  # we want height to be x2 and width to be x3
+    random_input = np.random.rand(batch, height, width, num_channels)
+
+    # conv2d_transpose is a little tricky.
+    # height_new = (height_old - 1) * stride + kernel - 2*padding - correction
+    # here kernel = 3, padding = 0, correction = 1
+    upsampled_height = (height - 1) * stride[0] + 3 - 2*0 - 1
+    upsampled_width = (width - 1) * stride[1] + 3 - 2*0 - 1
+    upsampled_output = common_layers.cyclegan_upsample(random_input,
+                                                       output_filters, stride,
+                                                       "conv2d_transpose")
+    upsampled_output_shape = common_layers.shape_list(upsampled_output)
+    self.assertEqual(batch, upsampled_output_shape[0])
+    self.assertEqual(upsampled_height, upsampled_output_shape[1])
+    self.assertEqual(upsampled_width, upsampled_output_shape[2])
+    self.assertEqual(output_filters, upsampled_output_shape[3])
+
 
 class FnWithCustomGradTest(tf.test.TestCase):
 
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index bca4ef8de..bce98c37a 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -28,11 +28,6 @@
 
 import tensorflow as tf
 
-try:
-  from tensorflow_models.slim.nets.cyclegan import cyclegan_upsample  # pylint: disable=g-import-not-at-top
-except ImportError:
-  pass
-
 tfl = tf.layers
 tfcl = tf.contrib.layers
 
@@ -445,7 +440,7 @@ def construct_predictive_tower(
       enc0, enc1 = skips
 
       with tf.variable_scope("upsample1", reuse=tf.AUTO_REUSE):
-        enc4 = cyclegan_upsample(
+        enc4 = common_layers.cyclegan_upsample(
             hidden5, num_outputs=hidden5.shape.as_list()[-1],
             stride=[2, 2], method=upsample_method)
 
@@ -463,7 +458,7 @@ def construct_predictive_tower(
       hidden6 = tf.concat(axis=3, values=[hidden6, enc1])  # both 16x16
 
       with tf.variable_scope("upsample2", reuse=tf.AUTO_REUSE):
-        enc5 = cyclegan_upsample(
+        enc5 = common_layers.cyclegan_upsample(
             hidden6, num_outputs=hidden6.shape.as_list()[-1],
             stride=[2, 2], method=upsample_method)
 
@@ -480,7 +475,7 @@ def construct_predictive_tower(
       hidden7 = tf.concat(axis=3, values=[hidden7, enc0])  # both 32x32
 
       with tf.variable_scope("upsample3", reuse=tf.AUTO_REUSE):
-        enc6 = cyclegan_upsample(
+        enc6 = common_layers.cyclegan_upsample(
             hidden7, num_outputs=hidden7.shape.as_list()[-1],
             stride=[2, 2], method=upsample_method)
       enc6 = tfcl.layer_norm(enc6, scope="layer_norm9")

From 7047c94abf1f9653774b5c478b8d9e6732c95225 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 25 Jul 2018 17:13:45 -0700
Subject: [PATCH 0385/2720] zero add

PiperOrigin-RevId: 206081798
---
 tensor2tensor/layers/common_layers.py | 25 ++++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 23e11d2b3..3fc4d335f 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -730,12 +730,27 @@ def apply_norm(x, norm_type, depth, epsilon):
                    "'noam', 'lr', 'none'.")
 
 
-def zero_add(previous_value, x, name=None, reuse=False):
-  """Resnet connection with zero initialization."""
-  with tf.name_scope(
-      name, default_name="zero_add", values=[previous_value, x], reuse=reuse):
+def zero_add(previous_value, x, name=None, reuse=None):
+  """Resnet connection with zero initialization.
+
+  Another type of resnet connection which returns previous_value + gamma * x.
+  gamma is a trainable scalar and initialized with zero. It is useful when a
+  module is plugged into a trained model and we want to make sure it matches the
+  original model's performance.
+
+  Args:
+    previous_value:  A tensor.
+    x: A tensor.
+    name: name of variable scope; defaults to zero_add.
+    reuse: reuse scope.
+
+  Returns:
+    previous_value + gamma * x.
+  """
+  with tf.variable_scope(
+      name, default_name="zero_add", reuse=reuse):
     gamma = tf.get_variable(
-        "gamma", [None], initializer=tf.zeros_initializer())
+        "gamma", (), initializer=tf.zeros_initializer())
     return previous_value + gamma*x
 
 

From b7f9acb5735cf836beff2a49565a8932ec3c31bf Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 25 Jul 2018 17:42:17 -0700
Subject: [PATCH 0386/2720] Add decoding option to pass through the features
 dictionary to predictions

PiperOrigin-RevId: 206085079
---
 tensor2tensor/utils/decoding.py  | 3 ++-
 tensor2tensor/utils/t2t_model.py | 5 ++++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 2ee8ff753..3f1691ebc 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -59,7 +59,8 @@ def decode_hparams(overrides=""):
       shards=1,
       shard_id=0,
       num_decodes=1,
-      force_decode_length=False)
+      force_decode_length=False,
+      pass_through_features=False)
   hp.parse(overrides)
   return hp
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 72d60e3c9..8fd1b37b9 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1361,8 +1361,11 @@ def estimator_spec_predict(self, features, use_tpu=False):
         "scores": scores,
         "inputs": inputs,
         "targets": features.get("infer_targets"),
-        "batch_prediction_key": features.get("batch_prediction_key"),
     }
+    if decode_hparams.pass_through_features:
+      for k in features:
+        if k not in list(predictions.keys()) + ["infer_targets"]:
+          predictions[k] = features[k]
 
     _del_dict_non_tensors(predictions)
 

From bc5187788a7af2af273f98fb41b10853ac663ba0 Mon Sep 17 00:00:00 2001
From: Tomasz Latkowski <13836101+tlatkowski@users.noreply.github.com>
Date: Thu, 26 Jul 2018 05:02:06 +0200
Subject: [PATCH 0387/2720] fixed bug with data_gen in style transfer problem
 (#951)

* fixed bug with data_gen in style transfer problems

* renamed style transfer problem
---
 .../data_generators/style_transfer.py         | 48 ++++++-----
 .../data_generators/style_transfer_test.py    | 79 +++++++++++++++++++
 2 files changed, 108 insertions(+), 19 deletions(-)
 create mode 100644 tensor2tensor/data_generators/style_transfer_test.py

diff --git a/tensor2tensor/data_generators/style_transfer.py b/tensor2tensor/data_generators/style_transfer.py
index 96af64bd8..ded5af9ba 100644
--- a/tensor2tensor/data_generators/style_transfer.py
+++ b/tensor2tensor/data_generators/style_transfer.py
@@ -57,7 +57,7 @@
 _SUBWORD_VOCAB_SIZE = 8000
 
 
-class StyleTransferProblem(text_problems.Text2TextProblem):
+class StyleTransferProblemShakespeare(text_problems.Text2TextProblem):
   """Base class for transferring styles problems"""
 
   @property
@@ -69,7 +69,10 @@ def source(self):
     raise NotImplementedError()
 
   def dataset_url(self, dataset_split):
-    raise NotImplementedError()
+    train = dataset_split == problem.DatasetSplit.TRAIN
+    if train:
+      return _SHAKESPEARE_MODERN_TRAIN_DATASET
+    return _SHAKESPEARE_MODERN_DEV_DATASET
 
   def vocab_data_files(self):
     """Files to be passed to get_or_generate_vocab."""
@@ -97,8 +100,6 @@ def is_generate_per_split(self):
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
     dataset = self.dataset_url(dataset_split)
 
-    tag = "train" if dataset_split == problem.DatasetSplit.TRAIN else "dev"
-
     url = dataset[0][0]
     compressed_filename = os.path.basename(url)
     compressed_filepath = os.path.join(tmp_dir, compressed_filename)
@@ -113,14 +114,19 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
           data_dir, tmp_dir, self.vocab_filename, self.approx_vocab_size,
           self.vocab_data_files())
 
-    source_file = os.path.join(tmp_dir, tag + ".modern")
-    target_file = os.path.join(tmp_dir, tag + ".original")
+    source_file, target_file = self.source_target_paths(dataset_split, tmp_dir)
     return text_problems.text2text_txt_iterator(source_file,
                                                 target_file)
 
+  def source_target_paths(self, dataset_split, tmp_dir):
+    tag = "train" if dataset_split == problem.DatasetSplit.TRAIN else "dev"
+    source_path = os.path.join(tmp_dir, tag + self.source)
+    target_path = os.path.join(tmp_dir, tag + self.target)
+    return source_path, target_path
+
 
 @registry.register_problem
-class StyleTransferShakespeareToModern(StyleTransferProblem):
+class StyleTransferShakespeareToModern(StyleTransferProblemShakespeare):
   """Transferring style from Shakespeare original English to modern one"""
 
   @property
@@ -131,15 +137,9 @@ def target(self):
   def source(self):
     return ".original"
 
-  def dataset_url(self, dataset_split):
-    train = dataset_split == problem.DatasetSplit.TRAIN
-    if train:
-      return _SHAKESPEARE_MODERN_TRAIN_DATASET
-    return _SHAKESPEARE_MODERN_DEV_DATASET
-
 
 @registry.register_problem
-class StyleTransferModernToShakespeare(StyleTransferProblem):
+class StyleTransferModernToShakespeare(StyleTransferProblemShakespeare):
   """Transferring style from modern English to Shakespeare original English"""
 
   @property
@@ -150,8 +150,18 @@ def target(self):
   def source(self):
     return ".modern"
 
-  def dataset_url(self, dataset_split):
-    train = dataset_split == problem.DatasetSplit.TRAIN
-    if train:
-      return _SHAKESPEARE_MODERN_TRAIN_DATASET
-    return _SHAKESPEARE_MODERN_DEV_DATASET
+
+@registry.register_problem
+class StyleTransferShakespeareToModernCharacters(StyleTransferShakespeareToModern):
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+
+@registry.register_problem
+class StyleTransferModernToShakespeareCharacters(StyleTransferModernToShakespeare):
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
\ No newline at end of file
diff --git a/tensor2tensor/data_generators/style_transfer_test.py b/tensor2tensor/data_generators/style_transfer_test.py
new file mode 100644
index 000000000..77f5c8b81
--- /dev/null
+++ b/tensor2tensor/data_generators/style_transfer_test.py
@@ -0,0 +1,79 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.data_generators.style_transfer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators import style_transfer
+from tensor2tensor.data_generators import problem
+
+import tensorflow as tf
+
+
+class StyleTransferProblemShakespeareTest(tf.test.TestCase):
+
+  def testSourceAndTargetPathsTrainModern2Shakespeare(self):
+    tmp_dir = 'tmp_dir'
+    modern_to_shakespeare_data_gen = style_transfer.StyleTransferModernToShakespeare()
+    actual_source, actual_target = modern_to_shakespeare_data_gen.source_target_paths(
+      problem.DatasetSplit.TRAIN, tmp_dir)
+
+    expected_source = '{}/train.modern'.format(tmp_dir)
+    expected_target = '{}/train.original'.format(tmp_dir)
+
+    self.assertEqual(actual_source, expected_source)
+    self.assertEqual(actual_target, expected_target)
+
+  def testSourceAndTargetPathsTrainShakespeare2Modern(self):
+    tmp_dir = 'tmp_dir'
+    shakespeare_to_modern_data_gen = style_transfer.StyleTransferShakespeareToModern()
+    actual_source, actual_target = shakespeare_to_modern_data_gen.source_target_paths(
+      problem.DatasetSplit.TRAIN, tmp_dir)
+
+    expected_source = '{}/train.original'.format(tmp_dir)
+    expected_target = '{}/train.modern'.format(tmp_dir)
+
+    self.assertEqual(actual_source, expected_source)
+    self.assertEqual(actual_target, expected_target)
+
+  def testSourceAndTargetPathsDevModern2Shakespeare(self):
+    tmp_dir = 'tmp_dir'
+    modern_to_shakespeare_data_gen = style_transfer.StyleTransferModernToShakespeare()
+    actual_source, actual_target = modern_to_shakespeare_data_gen.source_target_paths(
+      problem.DatasetSplit.EVAL, tmp_dir)
+
+    expected_source = '{}/dev.modern'.format(tmp_dir)
+    expected_target = '{}/dev.original'.format(tmp_dir)
+
+    self.assertEqual(actual_source, expected_source)
+    self.assertEqual(actual_target, expected_target)
+
+  def testSourceAndTargetPathsDevShakespeare2Modern(self):
+    tmp_dir = 'tmp_dir'
+    shakespeare_to_modern_data_gen = style_transfer.StyleTransferShakespeareToModern()
+    actual_source, actual_target = shakespeare_to_modern_data_gen.source_target_paths(
+      problem.DatasetSplit.EVAL, tmp_dir)
+
+    expected_source = '{}/dev.original'.format(tmp_dir)
+    expected_target = '{}/dev.modern'.format(tmp_dir)
+
+    self.assertEqual(actual_source, expected_source)
+    self.assertEqual(actual_target, expected_target)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 5b91af6a02b952cda369743829b252f7270bcb39 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 25 Jul 2018 20:10:26 -0700
Subject: [PATCH 0388/2720] internal merge of PR #951

PiperOrigin-RevId: 206097452
---
 .../data_generators/style_transfer.py         |  8 ++-
 .../data_generators/style_transfer_test.py    | 61 +++++++++++--------
 2 files changed, 39 insertions(+), 30 deletions(-)

diff --git a/tensor2tensor/data_generators/style_transfer.py b/tensor2tensor/data_generators/style_transfer.py
index ded5af9ba..b3096c7a5 100644
--- a/tensor2tensor/data_generators/style_transfer.py
+++ b/tensor2tensor/data_generators/style_transfer.py
@@ -152,7 +152,8 @@ def source(self):
 
 
 @registry.register_problem
-class StyleTransferShakespeareToModernCharacters(StyleTransferShakespeareToModern):
+class StyleTransferShakespeareToModernCharacters(
+    StyleTransferShakespeareToModern):
 
   @property
   def vocab_type(self):
@@ -160,8 +161,9 @@ def vocab_type(self):
 
 
 @registry.register_problem
-class StyleTransferModernToShakespeareCharacters(StyleTransferModernToShakespeare):
+class StyleTransferModernToShakespeareCharacters(
+    StyleTransferModernToShakespeare):
 
   @property
   def vocab_type(self):
-    return text_problems.VocabType.CHARACTER
\ No newline at end of file
+    return text_problems.VocabType.CHARACTER
diff --git a/tensor2tensor/data_generators/style_transfer_test.py b/tensor2tensor/data_generators/style_transfer_test.py
index 77f5c8b81..6e4e1aa94 100644
--- a/tensor2tensor/data_generators/style_transfer_test.py
+++ b/tensor2tensor/data_generators/style_transfer_test.py
@@ -18,58 +18,65 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.data_generators import style_transfer
 from tensor2tensor.data_generators import problem
-
+from tensor2tensor.data_generators import style_transfer
 import tensorflow as tf
 
 
 class StyleTransferProblemShakespeareTest(tf.test.TestCase):
 
   def testSourceAndTargetPathsTrainModern2Shakespeare(self):
-    tmp_dir = 'tmp_dir'
-    modern_to_shakespeare_data_gen = style_transfer.StyleTransferModernToShakespeare()
-    actual_source, actual_target = modern_to_shakespeare_data_gen.source_target_paths(
-      problem.DatasetSplit.TRAIN, tmp_dir)
+    tmp_dir = "tmp_dir"
+    modern_to_shakespeare_data_gen = (
+        style_transfer.StyleTransferModernToShakespeare())
+    actual_source, actual_target = (
+        modern_to_shakespeare_data_gen.source_target_paths(
+            problem.DatasetSplit.TRAIN, tmp_dir))
 
-    expected_source = '{}/train.modern'.format(tmp_dir)
-    expected_target = '{}/train.original'.format(tmp_dir)
+    expected_source = "{}/train.modern".format(tmp_dir)
+    expected_target = "{}/train.original".format(tmp_dir)
 
     self.assertEqual(actual_source, expected_source)
     self.assertEqual(actual_target, expected_target)
 
   def testSourceAndTargetPathsTrainShakespeare2Modern(self):
-    tmp_dir = 'tmp_dir'
-    shakespeare_to_modern_data_gen = style_transfer.StyleTransferShakespeareToModern()
-    actual_source, actual_target = shakespeare_to_modern_data_gen.source_target_paths(
-      problem.DatasetSplit.TRAIN, tmp_dir)
+    tmp_dir = "tmp_dir"
+    shakespeare_to_modern_data_gen = (
+        style_transfer.StyleTransferShakespeareToModern())
+    actual_source, actual_target = (
+        shakespeare_to_modern_data_gen.source_target_paths(
+            problem.DatasetSplit.TRAIN, tmp_dir))
 
-    expected_source = '{}/train.original'.format(tmp_dir)
-    expected_target = '{}/train.modern'.format(tmp_dir)
+    expected_source = "{}/train.original".format(tmp_dir)
+    expected_target = "{}/train.modern".format(tmp_dir)
 
     self.assertEqual(actual_source, expected_source)
     self.assertEqual(actual_target, expected_target)
 
   def testSourceAndTargetPathsDevModern2Shakespeare(self):
-    tmp_dir = 'tmp_dir'
-    modern_to_shakespeare_data_gen = style_transfer.StyleTransferModernToShakespeare()
-    actual_source, actual_target = modern_to_shakespeare_data_gen.source_target_paths(
-      problem.DatasetSplit.EVAL, tmp_dir)
+    tmp_dir = "tmp_dir"
+    modern_to_shakespeare_data_gen = (
+        style_transfer.StyleTransferModernToShakespeare())
+    actual_source, actual_target = (
+        modern_to_shakespeare_data_gen.source_target_paths(
+            problem.DatasetSplit.EVAL, tmp_dir))
 
-    expected_source = '{}/dev.modern'.format(tmp_dir)
-    expected_target = '{}/dev.original'.format(tmp_dir)
+    expected_source = "{}/dev.modern".format(tmp_dir)
+    expected_target = "{}/dev.original".format(tmp_dir)
 
     self.assertEqual(actual_source, expected_source)
     self.assertEqual(actual_target, expected_target)
 
   def testSourceAndTargetPathsDevShakespeare2Modern(self):
-    tmp_dir = 'tmp_dir'
-    shakespeare_to_modern_data_gen = style_transfer.StyleTransferShakespeareToModern()
-    actual_source, actual_target = shakespeare_to_modern_data_gen.source_target_paths(
-      problem.DatasetSplit.EVAL, tmp_dir)
-
-    expected_source = '{}/dev.original'.format(tmp_dir)
-    expected_target = '{}/dev.modern'.format(tmp_dir)
+    tmp_dir = "tmp_dir"
+    shakespeare_to_modern_data_gen = (
+        style_transfer.StyleTransferShakespeareToModern())
+    actual_source, actual_target = (
+        shakespeare_to_modern_data_gen.source_target_paths(
+            problem.DatasetSplit.EVAL, tmp_dir))
+
+    expected_source = "{}/dev.original".format(tmp_dir)
+    expected_target = "{}/dev.modern".format(tmp_dir)
 
     self.assertEqual(actual_source, expected_source)
     self.assertEqual(actual_target, expected_target)

From c43866713438da0c206718b223ae3098183b3b34 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 26 Jul 2018 10:34:37 -0700
Subject: [PATCH 0389/2720] Always pass-through features on prediction

PiperOrigin-RevId: 206181730
---
 tensor2tensor/utils/decoding.py  |  3 +--
 tensor2tensor/utils/t2t_model.py | 13 +++++++++----
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 3f1691ebc..2ee8ff753 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -59,8 +59,7 @@ def decode_hparams(overrides=""):
       shards=1,
       shard_id=0,
       num_decodes=1,
-      force_decode_length=False,
-      pass_through_features=False)
+      force_decode_length=False)
   hp.parse(overrides)
   return hp
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 8fd1b37b9..1d50c6397 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1362,10 +1362,15 @@ def estimator_spec_predict(self, features, use_tpu=False):
         "inputs": inputs,
         "targets": features.get("infer_targets"),
     }
-    if decode_hparams.pass_through_features:
-      for k in features:
-        if k not in list(predictions.keys()) + ["infer_targets"]:
-          predictions[k] = features[k]
+
+    # Pass through remaining features
+    for name, feature in features.items():
+      if name not in list(predictions.keys()) + ["infer_targets"]:
+        if not feature.shape.as_list():
+          # All features must have a batch dimension
+          batch_size = common_layers.shape_list(outputs)[0]
+          feature = tf.tile(tf.expand_dims(feature, 0), [batch_size])
+        predictions[name] = feature
 
     _del_dict_non_tensors(predictions)
 

From cec250db49fae8f941a86f1c9685bcd13735c4c4 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 26 Jul 2018 10:40:58 -0700
Subject: [PATCH 0390/2720] Use session in cyclegan_upsample tests.

PiperOrigin-RevId: 206182864
---
 tensor2tensor/layers/common_layers_test.py | 33 ++++++++++++----------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index d24193769..c7f308952 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -639,11 +639,12 @@ def testCycleGANUpsampleNnUpsampleConv(self):
     # nn_upsample_conv gives exactly the shapes we'd expect.
     upsampled_output = common_layers.cyclegan_upsample(
         random_input, output_filters, stride, "nn_upsample_conv")
-    upsampled_output_shape = common_layers.shape_list(upsampled_output)
-    self.assertEqual(batch, upsampled_output_shape[0])
-    self.assertEqual(height * stride[0], upsampled_output_shape[1])
-    self.assertEqual(width * stride[1], upsampled_output_shape[2])
-    self.assertEqual(output_filters, upsampled_output_shape[3])
+    upsampled_output_shape = tf.shape(upsampled_output)
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      self.assertAllEqual(
+          [batch, height * stride[0], width * stride[1], output_filters],
+          session.run(upsampled_output_shape))
 
   def testCycleGANUpsampleBilinearUpsampleConv(self):
     batch = 8
@@ -657,11 +658,12 @@ def testCycleGANUpsampleBilinearUpsampleConv(self):
     # bilinear_upsample_conv gives exactly the shapes we'd expect.
     upsampled_output = common_layers.cyclegan_upsample(
         random_input, output_filters, stride, "bilinear_upsample_conv")
-    upsampled_output_shape = common_layers.shape_list(upsampled_output)
-    self.assertEqual(batch, upsampled_output_shape[0])
-    self.assertEqual(height * stride[0], upsampled_output_shape[1])
-    self.assertEqual(width * stride[1], upsampled_output_shape[2])
-    self.assertEqual(output_filters, upsampled_output_shape[3])
+    upsampled_output_shape = tf.shape(upsampled_output)
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      self.assertAllEqual(
+          [batch, height * stride[0], width * stride[1], output_filters],
+          session.run(upsampled_output_shape))
 
   def testCycleGANUpsampleConv2dTranspose(self):
     batch = 8
@@ -680,11 +682,12 @@ def testCycleGANUpsampleConv2dTranspose(self):
     upsampled_output = common_layers.cyclegan_upsample(random_input,
                                                        output_filters, stride,
                                                        "conv2d_transpose")
-    upsampled_output_shape = common_layers.shape_list(upsampled_output)
-    self.assertEqual(batch, upsampled_output_shape[0])
-    self.assertEqual(upsampled_height, upsampled_output_shape[1])
-    self.assertEqual(upsampled_width, upsampled_output_shape[2])
-    self.assertEqual(output_filters, upsampled_output_shape[3])
+    upsampled_output_shape = tf.shape(upsampled_output)
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      self.assertAllEqual(
+          [batch, upsampled_height, upsampled_width, output_filters],
+          session.run(upsampled_output_shape))
 
 
 class FnWithCustomGradTest(tf.test.TestCase):

From 342fa97074d0a2b54cd070e79643af0cef5a2cc4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 26 Jul 2018 11:42:01 -0700
Subject: [PATCH 0391/2720] Truncate targets after prepending inputs, when
 prepend_mode set

PiperOrigin-RevId: 206194497
---
 tensor2tensor/data_generators/problem.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index e6253120f..b871e365a 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -139,14 +139,14 @@ def preprocess_example_common(example, hparams, mode):
   """Preprocessing steps common to all models."""
   if hparams.max_input_seq_length > 0:
     example["inputs"] = example["inputs"][:hparams.max_input_seq_length]
-  if hparams.max_target_seq_length > 0:
-    example["targets"] = example["targets"][:hparams.max_target_seq_length]
   if hparams.prepend_mode != "none":
     if mode == tf.estimator.ModeKeys.PREDICT:
       example["partial_targets"] = tf.concat([example["inputs"], [0]], 0)
     else:
       example["targets"] = tf.concat(
           [example["inputs"], [0], example["targets"]], 0)
+  if hparams.max_target_seq_length > 0:
+    example["targets"] = example["targets"][:hparams.max_target_seq_length]
   if hparams.split_to_length:
     example["targets"] = tf.reshape(example["targets"],
                                     [-1, hparams.split_to_length, 1, 1])

From ee65f4707d77e47c9d6bb71d58a6fd27f8f94e03 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Thu, 26 Jul 2018 18:37:59 -0700
Subject: [PATCH 0392/2720] Internal change

PiperOrigin-RevId: 206256396
---
 tensor2tensor/data_generators/all_problems.py |  1 +
 tensor2tensor/data_generators/lm1b_mnli.py    | 39 +++++++++++++++++++
 .../data_generators/multi_problem.py          | 16 ++++++--
 tensor2tensor/data_generators/problem.py      | 18 ++++-----
 tensor2tensor/layers/common_layers.py         | 25 ++++++++++++
 tensor2tensor/utils/metrics.py                |  8 ++++
 tensor2tensor/utils/t2t_model.py              |  6 ++-
 7 files changed, 99 insertions(+), 14 deletions(-)
 create mode 100644 tensor2tensor/data_generators/lm1b_mnli.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index fe816b683..3621862a8 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -47,6 +47,7 @@
     "tensor2tensor.data_generators.librispeech",
     "tensor2tensor.data_generators.lm1b",
     "tensor2tensor.data_generators.lm1b_imdb",
+    "tensor2tensor.data_generators.lm1b_mnli",
     "tensor2tensor.data_generators.mnist",
     "tensor2tensor.data_generators.mrpc",
     "tensor2tensor.data_generators.mscoco",
diff --git a/tensor2tensor/data_generators/lm1b_mnli.py b/tensor2tensor/data_generators/lm1b_mnli.py
new file mode 100644
index 000000000..a0d8e076d
--- /dev/null
+++ b/tensor2tensor/data_generators/lm1b_mnli.py
@@ -0,0 +1,39 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generators for LM1B and MNLI combined datasets."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators import lm1b
+from tensor2tensor.data_generators import multi_problem
+from tensor2tensor.data_generators import multinli
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+
+
+@registry.register_problem
+class LanguagemodelLm1bMultiNLI(multi_problem.MultiProblem):
+  """LM1b and MNLI mixed problem class for multitask learning."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelLm1bMultiNLI, self).__init__(was_reversed, was_copy)
+    self.task_list.append(lm1b.LanguagemodelLm1bCharacters())
+    self.task_list.append(multinli.MultiNLICharacters())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 5a0611553..f14bfaf4d 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -21,6 +21,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.layers import discretization
+from tensor2tensor.utils import metrics
 import tensorflow as tf
 
 
@@ -63,14 +64,14 @@ def get_hparams(self, model_hparams=None):
     if self._hparams is not None:
       return self._hparams
 
-    self._hparams = self.task_list[0].get_hparams()
+    self._hparams = self.task_list[0].get_hparams(model_hparams)
 
     return self._hparams
 
   def flatten_zip(self, *args):
     flattened = tf.data.Dataset.from_tensors(args[0])
     for ex in args[1:]:
-      flattened.concatenate(tf.data.Dataset.from_tensors(ex))
+      flattened = flattened.concatenate(tf.data.Dataset.from_tensors(ex))
 
     return flattened
 
@@ -96,7 +97,9 @@ def dataset(self,
                                   output_buffer_size, shuffle_files,
                                   hparams, preprocess, dataset_split,
                                   shard, partition_id, num_partitions,
-                                  max_records).repeat()
+                                  max_records)
+      if is_training:
+        task_dataset = task_dataset.repeat()
       # pylint: disable=cell-var-from-loop
       task_dataset = task_dataset.map(lambda x: self.add_task_id(task, x))
       datasets.append(task_dataset)
@@ -108,7 +111,12 @@ def dataset(self,
           self.flatten_zip)
     else:
       single_mtl_dataset = datasets[0]
-      for data in datasets[1:]:
+      for data in datasets[0:]:
         single_mtl_dataset = single_mtl_dataset.concatenate(data)
 
     return single_mtl_dataset
+
+  def eval_metrics(self):
+    return [
+        metrics.Metrics.ACC, metrics.Metrics.NEG_LOG_PERPLEXITY
+    ]
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index b871e365a..73d49ccf5 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -107,23 +107,23 @@ class SpaceID(object):
 class TaskID(object):
   """Problem specific task ids. Add more as needed."""
   # English characters
-  EN_CHR = 0
+  EN_CHR = 2
   # English characters sentiment
-  EN_CHR_SENT = 1
+  EN_CHR_SENT = 3
   # English Premise Hypothesis pair
-  EN_PR_HYP = 2
+  EN_PR_HYP = 4
   # English NLI
-  EN_NLI = 3
+  EN_NLI = 5
   # COLA
-  COLA = 4
+  COLA = 6
   # Enligh Question Context pair
-  EN_Q_CONT = 5
+  EN_Q_CONT = 7
   # English similarity task
-  EN_SIM = 6
+  EN_SIM = 8
   # English sentence pair
-  EN_SENT_PAIR = 7
+  EN_SENT_PAIR = 9
   # 3 class NLI
-  THREE_CL_NLI = 8
+  THREE_CL_NLI = 10
 
 
 def default_model_hparams():
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 3fc4d335f..a6e600867 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1890,6 +1890,31 @@ def weights_prepend_inputs_to_targets(labels):
   return tf.to_float(tf.not_equal(past_first_zero * nonzero, 0))
 
 
+def weights_multi_problem(labels, taskid=-1):
+  """Assign weight 1.0 to only the "targets" portion of the labels.
+
+  Weight 1.0 is assigned to all labels past the taskid.
+
+  Args:
+    labels: A Tensor of int32s.
+    taskid: an int32 representing the task id for a problem.
+
+  Returns:
+    A Tensor of floats.
+
+  Raises:
+    ValueError: The Task ID must be valid.
+  """
+  if taskid < 0:
+    raise ValueError("Task ID must be non-negative.")
+
+  past_taskid = tf.cumsum(tf.to_float(tf.equal(labels, taskid)), axis=1)
+  # Additionally zero out the task id location
+  past_taskid *= tf.to_float(tf.not_equal(labels, taskid))
+  non_taskid = tf.to_float(labels)
+  return tf.to_float(tf.not_equal(past_taskid * non_taskid, 0))
+
+
 def weights_all(labels):
   """Assign weight 1.0 to all labels."""
   return tf.ones_like(labels, dtype=tf.float32)
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 683dc5e48..f5ab0033e 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -566,10 +566,15 @@ def image_wrapped_metric_fn(predictions,
 
     return image_wrapped_metric_fn
 
+  def weights_fn_for_mp(problem_task_id):
+    return lambda x: common_layers.weights_multi_problem(x, problem_task_id)
+
   eval_metrics = dict()
   for problem_instance in problems:
     problem_name = problem_instance.name
     metrics = problem_instance.eval_metrics()
+    if hasattr(model_hparams.problem, "task_list"):
+      metrics = model_hparams.problem.eval_metrics()
     if not all([m in METRICS_FNS for m in metrics]):
       error_str = ("Unrecognized metric. Problem %s specified metrics "
                    "%s. Recognized metrics are %s.")
@@ -585,6 +590,9 @@ def image_wrapped_metric_fn(predictions,
       if isinstance(modality, tuple):
         modality = registry.create_modality(modality, model_hparams)
       weights_fn = modality.targets_weights_fn
+      if hasattr(model_hparams.problem, "task_list"):
+        ptid = problem_instance.task_id  # pylint: disable=cell-var-from-loop
+        weights_fn = weights_fn_for_mp(ptid)
 
       for metric in metrics:
         metric_fn = METRICS_FNS[metric]
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 1d50c6397..ee21160ce 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1305,7 +1305,11 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
             eval_metrics=(eval_metrics_fn, [logits, labels]),
             loss=loss)
     else:
-      eval_metrics_fns = metrics.create_evaluation_metrics([problem], hparams)
+      task_list = [problem]
+      if hasattr(problem, "task_list"):
+        task_list = problem.task_list
+
+      eval_metrics_fns = metrics.create_evaluation_metrics(task_list, hparams)
       eval_metrics = {}
       for metric_name, metric_fn in six.iteritems(eval_metrics_fns):
         if isinstance(logits, dict):

From dc03545fdd0b4408a3d5bd0d96baf652ac9f4702 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 27 Jul 2018 09:43:14 -0700
Subject: [PATCH 0393/2720] Fix KL axis.

PiperOrigin-RevId: 206329450
---
 tensor2tensor/models/research/next_frame.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index bce98c37a..fa061c028 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -766,10 +766,9 @@ def kl_divergence(self, mu, log_sigma):
     Returns:
       the KL loss.
     """
-    # TODO(mechcoder): Sum across all but the first dimension.
-    return -.5 * tf.reduce_sum(
-        1. + log_sigma - tf.square(mu) - tf.exp(log_sigma),
-        axis=1)
+    batch_size = common_layers.shape_list(mu)[0]
+    kl = -.5 * tf.reduce_sum(1. + log_sigma - tf.square(mu) - tf.exp(log_sigma))
+    return kl / tf.to_float(batch_size)
 
   def get_input_if_exists(self, features, key, batch_size, num_frames):
     if key in features:

From 937c1ae6049e9f1038ec1ef5affb8f257e512736 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 27 Jul 2018 10:23:14 -0700
Subject: [PATCH 0394/2720] Internal change

PiperOrigin-RevId: 206335166
---
 tensor2tensor/bin/t2t_trainer.py      |  5 ++++
 tensor2tensor/layers/common_layers.py |  3 ++
 tensor2tensor/utils/optimize.py       | 13 ++++----
 tensor2tensor/utils/t2t_model.py      | 43 ++++++++++++++++-----------
 tensor2tensor/utils/trainer_lib.py    | 12 +++++---
 5 files changed, 50 insertions(+), 26 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 5f7e36caa..1883452d9 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -48,6 +48,10 @@
 flags.DEFINE_integer("iterations_per_loop", 100,
                      "Number of iterations in a TPU training loop.")
 flags.DEFINE_bool("use_tpu", False, "Whether to use TPU.")
+flags.DEFINE_bool("xla_compile", False, "Whether to use XLA to compile. When "
+                  "this is set to True, computation will be constructed to "
+                  "optimize for XLA as if use_tpu=True but run on CPU/GPU "
+                  "instead of TPU.")
 flags.DEFINE_integer("tpu_infeed_sleep_secs", None,
                      "How long to sleep the infeed thread.")
 flags.DEFINE_bool("generate_data", False, "Generate data before training?")
@@ -174,6 +178,7 @@ def create_experiment_fn(**kwargs):
       eval_early_stopping_metric_minimize=FLAGS.
       eval_early_stopping_metric_minimize,
       use_tpu=FLAGS.use_tpu,
+      xla_compile=FLAGS.xla_compile,
       **kwargs)
 
 
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index a6e600867..a30b5ff0d 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3412,6 +3412,9 @@ def should_generate_summaries():
   Returns:
     a boolean
   """
+  if is_on_tpu():
+    # Summaries don't work well with TPU and XLA.
+    return False
   if "while/" in tf.contrib.framework.get_name_scope():
     # Summaries don't work well within tf.while_loop()
     return False
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 7249a76d8..bdbedbd17 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -44,9 +44,11 @@ def optimize(loss, learning_rate, hparams, use_tpu=False):
   if use_tpu:
     opt = tf.contrib.tpu.CrossShardOptimizer(opt)
 
-  tf.summary.scalar("learning_rate", learning_rate)
-  opt_summaries = ["loss"]
-  if hparams.summarize_grads:
+  opt_summaries = []
+  if common_layers.should_generate_summaries():
+    tf.summary.scalar("learning_rate", learning_rate)
+    opt_summaries = ["loss"]
+  if hparams.summarize_grads and common_layers.should_generate_summaries():
     tf.logging.info("Summarizing gradients")
     opt_summaries.extend(["gradients", "gradient_norm", "global_gradient_norm"])
 
@@ -136,7 +138,7 @@ def weight_decay_and_noise(loss, hparams, learning_rate, var_list=None):
   noise_vars = [v for v in var_list if "/body/" in v.name]
 
   weight_decay_loss = weight_decay(hparams.weight_decay, decay_vars)
-  if hparams.weight_decay:
+  if hparams.weight_decay and common_layers.should_generate_summaries():
     tf.summary.scalar("losses/weight_decay", weight_decay_loss)
   weight_noise_ops = weight_noise(hparams.weight_noise, learning_rate,
                                   noise_vars)
@@ -161,7 +163,8 @@ def weight_noise(noise_rate, learning_rate, var_list):
   for v in var_list:
     with tf.device(v._ref().device):  # pylint: disable=protected-access
       scale = noise_rate * learning_rate * 0.001
-      tf.summary.scalar("weight_noise_scale", scale)
+      if common_layers.should_generate_summaries():
+        tf.summary.scalar("weight_noise_scale", scale)
       noise = tf.truncated_normal(v.shape) * scale
       noise_op = v.assign_add(noise)
       noise_ops.append(noise_op)
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index ee21160ce..acbb3854e 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -194,7 +194,8 @@ def body_sharded(self, sharded_features):
 
   def model_fn_sharded(self, sharded_features):
     dp = self._data_parallelism
-    summarize_features(sharded_features, num_shards=dp.n)
+    if common_layers.should_generate_summaries():
+      summarize_features(sharded_features, num_shards=dp.n)
     datashard_to_features = self._to_features_per_datashard(sharded_features)
     if self.use_body_sharded:
       # MoE models override body_sharded
@@ -436,9 +437,10 @@ def loss(self, logits, features):
         losses[k] = self._loss_single(v, target_modality[k], features[k])
 
         n, d = losses[k]
-        tf.summary.scalar(k + "_loss", n / d)
-        tf.summary.scalar(k + "_loss_num", n)
-        tf.summary.scalar(k + "_loss_den", d)
+        if common_layers.should_generate_summaries():
+          tf.summary.scalar(k + "_loss", n / d)
+          tf.summary.scalar(k + "_loss_num", n)
+          tf.summary.scalar(k + "_loss_den", d)
 
       return tf.add_n([n / d for n, d in losses.values()])
     else:
@@ -453,15 +455,14 @@ def loss(self, logits, features):
         target_modality = target_modality["targets"]
       return self._loss_single(logits, target_modality, features["targets"])
 
-  def optimize(self, loss, num_async_replicas=1):
+  def optimize(self, loss, num_async_replicas=1, use_tpu=False):
     """Return a training op minimizing loss."""
     lr = learning_rate.learning_rate_schedule(self.hparams)
     if num_async_replicas > 1:
       log_info("Dividing learning rate by num_async_replicas: %d",
                num_async_replicas)
     lr /= math.sqrt(float(num_async_replicas))
-    train_op = optimize.optimize(
-        loss, lr, self.hparams, use_tpu=common_layers.is_on_tpu())
+    train_op = optimize.optimize(loss, lr, self.hparams, use_tpu=use_tpu)
     return train_op
 
   def set_mode(self, mode):
@@ -1148,7 +1149,8 @@ def get_eval_hooks(model_name):
   def make_estimator_model_fn(model_name,
                               hparams,
                               decode_hparams=None,
-                              use_tpu=False):
+                              use_tpu=False,
+                              xla_compile=False):
     model_cls = registry.model(model_name)
 
     def wrapping_model_fn(features, labels, mode, params=None, config=None):
@@ -1160,7 +1162,8 @@ def wrapping_model_fn(features, labels, mode, params=None, config=None):
           config=config,
           params=params,
           decode_hparams=decode_hparams,
-          use_tpu=use_tpu)
+          use_tpu=use_tpu,
+          xla_compile=xla_compile)
 
     return wrapping_model_fn
 
@@ -1173,7 +1176,8 @@ def estimator_model_fn(cls,
                          config=None,
                          params=None,
                          decode_hparams=None,
-                         use_tpu=False):
+                         use_tpu=False,
+                         xla_compile=False):
     """Model fn for Estimator.
 
     Args:
@@ -1185,6 +1189,7 @@ def estimator_model_fn(cls,
       params: dict, may include batch_size
       decode_hparams: HParams, used when mode == PREDICT.
       use_tpu: bool, whether using TPU
+      xla_compile: bool, whether to use XLA to compile graph, unimplemented.
 
     Returns:
       TPUEstimatorSpec if use tpu else EstimatorSpec
@@ -1216,7 +1221,7 @@ def estimator_model_fn(cls,
       logits, losses_dict = model(features)  # pylint: disable=not-callable
 
     # Set known shapes
-    if use_tpu:
+    if use_tpu or xla_compile:
       if isinstance(logits, dict):
         for k, v in sorted(six.iteritems(logits)):
           if "scalar/" in k:
@@ -1243,9 +1248,10 @@ def estimator_model_fn(cls,
       return logits
 
     # Summarize losses
-    with tf.name_scope("losses"):
-      for loss_name, loss_val in sorted(losses_dict.items()):
-        tf.summary.scalar(loss_name, loss_val)
+    if common_layers.should_generate_summaries():
+      with tf.name_scope("losses"):
+        for loss_name, loss_val in sorted(losses_dict.items()):
+          tf.summary.scalar(loss_name, loss_val)
 
     # Accumulate losses
     loss = sum(losses_dict[key] for key in sorted(losses_dict.keys()))
@@ -1260,11 +1266,14 @@ def estimator_model_fn(cls,
     num_async_replicas = (1 if (use_tpu or not config) else
                           config.t2t_device_info["num_async_replicas"])
     return model.estimator_spec_train(
-        loss, num_async_replicas=num_async_replicas)
+        loss,
+        num_async_replicas=num_async_replicas,
+        use_tpu=use_tpu and not xla_compile)
 
-  def estimator_spec_train(self, loss, num_async_replicas=1):
+  def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
     """Construct EstimatorSpec for TRAIN mode."""
-    train_op = self.optimize(loss, num_async_replicas=num_async_replicas)
+    train_op = self.optimize(
+        loss, num_async_replicas=num_async_replicas, use_tpu=use_tpu)
 
     if common_layers.is_on_tpu():
       host_call = _create_host_call(self.hparams.model_dir)
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index d00056b36..ac52d9e3e 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -199,10 +199,12 @@ def create_estimator(model_name,
                      run_config,
                      schedule="train_and_evaluate",
                      decode_hparams=None,
-                     use_tpu=False):
+                     use_tpu=False,
+                     xla_compile=False):
   """Create a T2T Estimator."""
   model_fn = t2t_model.T2TModel.make_estimator_model_fn(
-      model_name, hparams, decode_hparams=decode_hparams, use_tpu=use_tpu)
+      model_name, hparams, decode_hparams=decode_hparams, use_tpu=use_tpu,
+      xla_compile=xla_compile)
 
   if use_tpu:
     problem = hparams.problem
@@ -399,7 +401,8 @@ def create_experiment(
     eval_early_stopping_metric_delta=None,
     eval_early_stopping_metric_minimize=True,
     autotune=False,
-    use_tpu=False):
+    use_tpu=False,
+    xla_compile=False):
   """Create Experiment."""
   # HParams
   hparams.add_hparam("model_dir", run_config.model_dir)
@@ -416,7 +419,8 @@ def create_experiment(
       run_config,
       schedule=schedule,
       decode_hparams=decode_hparams,
-      use_tpu=use_tpu)
+      use_tpu=use_tpu,
+      xla_compile=xla_compile)
 
   # Input fns from Problem
   problem = hparams.problem

From f8c992587b14d0ce0b5bf7f4f6ae2ec4a65bd263 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 27 Jul 2018 13:56:14 -0700
Subject: [PATCH 0395/2720] Shuffling videos. If
 only_keep_videos_from_0th_frame is False (default for gym envs), the videos
 in the dataset are HIGHLY correlated (they only differ in one frame) which
 makes the training highly unstable. This shuffling is not the best solution,
 but rather a hacky solution to this issue. Ideally, we want a full shuffling
 on the whole dataset.

PiperOrigin-RevId: 206369208
---
 tensor2tensor/data_generators/video_utils.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index c341d1e82..ceebe7efa 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -299,7 +299,8 @@ def check_integrity_and_batch(*datasets):
     else:
       batch_dataset = preprocessed_dataset.apply(
           tf.contrib.data.batch_and_drop_remainder(num_frames))
-    dataset = batch_dataset.map(features_from_batch)  # shuffle(8)
+    dataset = batch_dataset.map(features_from_batch)
+    dataset = dataset.shuffle(256)
     return dataset
 
   def eval_metrics(self):

From 4503c368bb9b80beb79201182cac191f7670bb91 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 27 Jul 2018 15:23:57 -0700
Subject: [PATCH 0396/2720] Warm start.

PiperOrigin-RevId: 206383295
---
 tensor2tensor/bin/t2t_trainer.py   | 1 +
 tensor2tensor/utils/flags.py       | 1 +
 tensor2tensor/utils/trainer_lib.py | 8 +++++++-
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 1883452d9..91ebfee67 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -205,6 +205,7 @@ def create_run_config(hp):
       hp.weight_dtype == "float32")
   return trainer_lib.create_run_config(
       model_dir=os.path.expanduser(FLAGS.output_dir),
+      warm_start_from=FLAGS.warm_start_from,
       master=FLAGS.master,
       iterations_per_loop=FLAGS.iterations_per_loop,
       num_shards=FLAGS.tpu_num_shards,
diff --git a/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index fc8d66674..3ae12988a 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -82,6 +82,7 @@
                      "is max(iterations_per_loop, local_eval_frequency).")
 flags.DEFINE_bool("log_device_placement", False,
                   "Whether to log device placement.")
+flags.DEFINE_string("warm_start_from", None, "Warm start from checkpoint.")
 
 # Distributed training flags
 flags.DEFINE_integer("local_eval_frequency", 1000,
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index ac52d9e3e..94e818735 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -99,6 +99,7 @@ def create_hparams(hparams_set,
 
 def create_run_config(master="",
                       model_dir=None,
+                      warm_start_from=None,
                       iterations_per_loop=1000,
                       num_shards=8,
                       log_device_placement=False,
@@ -169,6 +170,7 @@ def create_run_config(master="",
     run_config_args["tpu_config"] = tpu_config
 
   config = run_config_cls(**run_config_args)
+  config.warm_start_from = warm_start_from
 
   # If not using TPU, add device info for data_parallelism
   config.use_tpu = use_tpu
@@ -225,7 +227,11 @@ def create_estimator(model_name,
         predict_batch_size=predict_batch_size)
   else:
     estimator = tf.estimator.Estimator(
-        model_fn=model_fn, model_dir=run_config.model_dir, config=run_config)
+        model_fn=model_fn,
+        model_dir=run_config.model_dir,
+        config=run_config,
+        warm_start_from=run_config.warm_start_from
+    )
   return estimator
 
 

From a39b078307ad9ffcbb8b37f34af8c10185ab7d9f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 27 Jul 2018 15:41:50 -0700
Subject: [PATCH 0397/2720] Fix division by zero.

PiperOrigin-RevId: 206385795
---
 tensor2tensor/utils/bleu_hook.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/utils/bleu_hook.py b/tensor2tensor/utils/bleu_hook.py
index e03c5b77a..73ef1081e 100644
--- a/tensor2tensor/utils/bleu_hook.py
+++ b/tensor2tensor/utils/bleu_hook.py
@@ -114,13 +114,16 @@ def compute_bleu(reference_corpus,
     geo_mean = math.exp(p_log_sum/max_order)
 
   if use_bp:
-    ratio = translation_length / reference_length
-    if ratio <= 0.0:
-      bp = 0.0
-    elif ratio >= 1.0:
+    if not reference_length:
       bp = 1.0
     else:
-      bp = math.exp(1 - 1. / ratio)
+      ratio = translation_length / reference_length
+      if ratio <= 0.0:
+        bp = 0.0
+      elif ratio >= 1.0:
+        bp = 1.0
+      else:
+        bp = math.exp(1 - 1. / ratio)
   bleu = geo_mean * bp
   return np.float32(bleu)
 

From d4163c1ec7fd8631f3e51d968445649953f17929 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 27 Jul 2018 17:45:59 -0700
Subject: [PATCH 0398/2720] Add support for warm starting from a subset of a
 pretrained model

PiperOrigin-RevId: 206400383
---
 tensor2tensor/layers/common_hparams.py |  5 +++++
 tensor2tensor/utils/t2t_model.py       | 16 ++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index b338bf2b8..f202cec8e 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -248,6 +248,11 @@ def basic_params1():
       #   roundoff.  Initial experiments show that model quality is similar
       #   to baseline for about 3M training steps, but worse thereafter.
       weight_dtype="float32",
+      # Directory containing a checkpoint for a pretrained model. This will only
+      # be used if a new run is being started. Parameters not found in the
+      # pretrained model will be randomly initialized. Superfluous parameters in
+      # the pretrained model will be ignored.
+      pretrained_model_dir="",
   )
 
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index acbb3854e..e6a7aa72c 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1275,6 +1275,22 @@ def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
     train_op = self.optimize(
         loss, num_async_replicas=num_async_replicas, use_tpu=use_tpu)
 
+    # TODO(mitchellstern): Add support for partitioned variables?
+    if (tf.train.latest_checkpoint(self._hparams.model_dir) is None and
+        self._hparams.pretrained_model_dir):
+      pretrained_model_dir = self._hparams.pretrained_model_dir
+      reader = tf.contrib.framework.load_checkpoint(pretrained_model_dir)
+      variable_map = {}
+      for var in tf.contrib.framework.get_trainable_variables():
+        var_name = var.name.split(":")[0]
+        if reader.has_tensor(var_name):
+          tf.logging.info("Loading variable from checkpoint: %s", var_name)
+          variable_map[var_name] = var
+        else:
+          tf.logging.info(
+              "Cannot find variable in checkpoint, skipping: %s", var_name)
+      tf.train.init_from_checkpoint(pretrained_model_dir, variable_map)
+
     if common_layers.is_on_tpu():
       host_call = _create_host_call(self.hparams.model_dir)
       _remove_summaries()

From daa34ff7ae1e6ff19394221b4fff0fff793f9084 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 27 Jul 2018 18:34:34 -0700
Subject: [PATCH 0399/2720] Add support for start-of-string token to
 Transformer fast_decode

PiperOrigin-RevId: 206404291
---
 tensor2tensor/models/transformer.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 9a61c7a4c..7966125d1 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -665,6 +665,7 @@ def fast_decode_tpu(encoder_output,
                     hparams,
                     decode_length,
                     beam_size=1,
+                    sos_id=0,
                     eos_id=beam_search.EOS_ID,
                     batch_size=None,
                     force_decode_length=False):
@@ -681,6 +682,7 @@ def fast_decode_tpu(encoder_output,
     hparams: Run hyperparameters.
     decode_length: An integer, how many additional timesteps to decode.
     beam_size: An integer, number of beams.
+    sos_id: Start-of-sequence symbol.
     eos_id: End-of-sequence symbol.
     batch_size: An integer, must be passed if there is no input.
     force_decode_length: A bool, whether to force the full decode length, or if
@@ -768,7 +770,7 @@ def is_not_finished(i, hit_eos, *_):
 
   decoded_ids = tf.zeros([batch_size, decode_length], dtype=tf.int64)
   hit_eos = tf.fill([batch_size], False)
-  next_id = tf.zeros([batch_size, 1], dtype=tf.int64)
+  next_id = sos_id * tf.ones([batch_size, 1], dtype=tf.int64)
   initial_log_prob = tf.zeros([batch_size], dtype=tf.float32)
 
   def compute_cache_shape_invariants(tensor):
@@ -802,6 +804,7 @@ def fast_decode(encoder_output,
                 beam_size=1,
                 top_beams=1,
                 alpha=1.0,
+                sos_id=0,
                 eos_id=beam_search.EOS_ID,
                 batch_size=None,
                 force_decode_length=False):
@@ -823,6 +826,7 @@ def fast_decode(encoder_output,
     top_beams: an integer. How many of the beams to return.
     alpha: Float that controls the length penalty. larger the alpha, stronger
       the preference for longer translations.
+    sos_id: End-of-sequence symbol in beam search.
     eos_id: End-of-sequence symbol in beam search.
     batch_size: an integer scalar - must be passed if there is no input
     force_decode_length: bool, whether to force the full decode length, or if
@@ -882,7 +886,7 @@ def fast_decode(encoder_output,
     cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias
 
   if beam_size > 1:  # Beam Search
-    initial_ids = tf.zeros([batch_size], dtype=tf.int32)
+    initial_ids = sos_id * tf.ones([batch_size], dtype=tf.int32)
     decoded_ids, scores = beam_search.beam_search(
         symbols_to_logits_fn,
         initial_ids,
@@ -927,7 +931,7 @@ def is_not_finished(i, hit_eos, *_):
 
     decoded_ids = tf.zeros([batch_size, 0], dtype=tf.int64)
     hit_eos = tf.fill([batch_size], False)
-    next_id = tf.zeros([batch_size, 1], dtype=tf.int64)
+    next_id = sos_id * tf.ones([batch_size, 1], dtype=tf.int64)
     initial_log_prob = tf.zeros([batch_size], dtype=tf.float32)
     _, _, _, decoded_ids, _, log_prob = tf.while_loop(
         is_not_finished,

From 27796faa97965b11142f32903870620a3193feb5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 27 Jul 2018 19:48:52 -0700
Subject: [PATCH 0400/2720] internal

PiperOrigin-RevId: 206408185
---
 tensor2tensor/utils/decoding.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 2ee8ff753..f298a5bf2 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -48,6 +48,7 @@ def decode_hparams(overrides=""):
       batch_size=0,
       beam_size=4,
       alpha=0.6,
+      eos_penalty=0.0,
       return_beams=False,
       write_beam_scores=False,
       max_input_size=-1,

From fc76cc27371be1f868384a9e74b1a1a743d588f5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 27 Jul 2018 20:03:19 -0700
Subject: [PATCH 0401/2720] Add fast block parallel transformer model

PiperOrigin-RevId: 206408714
---
 tensor2tensor/models/image_transformer_2d.py | 334 +++++++++++++++++++
 tensor2tensor/utils/decoding.py              |   3 +
 2 files changed, 337 insertions(+)

diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 6bc9f7ecc..a6cde618b 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -24,6 +24,7 @@
 from __future__ import print_function
 
 import copy
+import numpy as np
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
@@ -97,6 +98,269 @@ def body(self, features):
     return output
 
 
+@registry.register_model
+class Img2imgTransformerBlockParallel(t2t_model.T2TModel):
+  """Image-to-image transformer predicting blocks of the output in parallel."""
+
+  def body(self, features):
+    assert self._hparams.block_size > 0
+    assert not common_layers.is_on_tpu()
+
+    hparams = copy.copy(self._hparams)
+    targets = features["targets"]
+    inputs = features["inputs"]
+    if not (tf.get_variable_scope().reuse or
+            hparams.mode == tf.contrib.learn.ModeKeys.INFER):
+      tf.summary.image("inputs", inputs, max_outputs=1)
+      tf.summary.image("targets", targets, max_outputs=1)
+
+    encoder_input = cia.prepare_encoder(inputs, hparams)
+    encoder_output = cia.transformer_encoder_layers(
+        encoder_input,
+        hparams.num_encoder_layers,
+        hparams,
+        attention_type=hparams.enc_attention_type,
+        name="encoder")
+    decoder_input, rows, cols = cia.prepare_decoder(
+        targets, hparams)
+    decoder_output = cia.transformer_decoder_layers(
+        decoder_input,
+        encoder_output,
+        hparams.num_decoder_layers,
+        hparams,
+        attention_type=hparams.dec_attention_type,
+        name="decoder")
+
+    assert not isinstance(decoder_output, tuple)
+    assert len(decoder_output.shape) == 4
+
+    relu_dropout_broadcast_dims = (
+        common_layers.comma_separated_string_to_integer_list(
+            getattr(self._hparams, "relu_dropout_broadcast_dims", "")))
+
+    with tf.variable_scope("block_size_%d" % self._hparams.block_size):
+      tf.logging.info("Using block_size %d", self._hparams.block_size)
+      block_output = common_layers.dense_relu_dense(
+          decoder_output,
+          self._hparams.block_size * self._hparams.filter_size,
+          self._hparams.block_size * self._hparams.hidden_size,
+          dropout=self._hparams.relu_dropout,
+          dropout_broadcast_dims=relu_dropout_broadcast_dims)
+
+    batch_size, rows, cols = common_layers.shape_list(decoder_output)[:3]
+    decoder_output = tf.reshape(decoder_output, [
+        batch_size,
+        rows,
+        cols,
+        1,
+        self._hparams.hidden_size
+    ])
+    block_output = tf.reshape(block_output, [
+        batch_size,
+        rows,
+        cols,
+        self._hparams.block_size,
+        self._hparams.hidden_size
+    ])
+
+    block_output = common_layers.layer_postprocess(
+        decoder_output, block_output, self._hparams)
+
+    return block_output
+
+  def top(self, body_output, features):
+    assert self._hparams.block_size > 0
+
+    train_or_eval = (
+        self._hparams.mode == tf.estimator.ModeKeys.TRAIN or
+        self._hparams.mode == tf.estimator.ModeKeys.EVAL)
+
+    if train_or_eval:
+      if self._hparams.mode == tf.estimator.ModeKeys.TRAIN:
+        features["block_index"] = tf.random_uniform(
+            shape=[], minval=0, maxval=self._hparams.block_size, dtype=tf.int64)
+      else:
+        features["block_index"] = 0
+      body_output = body_output[:, :, :, features["block_index"], :]
+
+    decoded_image = tf.layers.dense(
+        body_output, 256, use_bias=True, activation=None, name="output_conv")
+
+    assert len(features["targets"].shape) == 4
+    targets_shape = common_layers.shape_list(features["targets"])
+
+    if train_or_eval:
+      output = tf.reshape(decoded_image, targets_shape + [256])
+    else:
+      output = tf.reshape(decoded_image, [
+          targets_shape[0], -1, self._hparams.block_size, 1, 256])
+      output = output[:, :targets_shape[1], :, :, :]
+
+    return output
+
+  def loss(self, logits, features):
+    assert self._hparams.block_size > 0
+
+    if self._hparams.mode == tf.estimator.ModeKeys.PREDICT:
+      return 0.0
+
+    def shift_left_2d(x, k):
+      return tf.pad(x, [[0, 0], [0, k]])[:, k:]
+
+    def shift_left_4d_raster_scan(x, k):
+      batch_size = common_layers.shape_list(x)[0]
+      return tf.reshape(
+          shift_left_2d(tf.reshape(x, [batch_size, -1]), k), tf.shape(x))
+
+    targets = features["targets"]
+    assert len(targets.shape) == 4
+
+    targets = tf.stack([
+        shift_left_4d_raster_scan(targets, i)
+        for i in range(self._hparams.block_size)
+    ], axis=4)
+
+    if (self._hparams.mode == tf.estimator.ModeKeys.TRAIN or
+        self._hparams.mode == tf.estimator.ModeKeys.EVAL):
+      assert "block_index" in features
+      targets = targets[:, :, :, :, features["block_index"]]
+
+    features["targets"] = targets
+
+    loss = super(Img2imgTransformerBlockParallel, self).loss(logits, features)
+
+    if self._hparams.mode == tf.estimator.ModeKeys.TRAIN:
+      k = features["block_index"]
+      loss_num, loss_den = loss
+      loss_val = loss_num / loss_den
+      for i in range(self._hparams.block_size):
+        # Hack: if you report a loss of NaN, TensorBoard will plot a point at
+        # the previous value without a connecting line. This is used here to
+        # separate out the training losses by block index.
+        one_or_nan = tf.cond(tf.equal(k, i), lambda: 1.0, lambda: float("nan"))
+        tf.summary.scalar(
+            "block_index_%d" % i, one_or_nan * loss_val, family="losses")
+
+    return loss
+
+  def _greedy_infer(self, features, decode_length, use_tpu=False):
+    assert not use_tpu
+    return self._slow_greedy_infer_guess_and_check(features, decode_length)
+
+  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
+    raise NotImplementedError
+
+  def _slow_greedy_infer_guess_and_check(self, features, decode_length):
+    assert self._hparams.block_size > 0
+    assert self._hparams.force_full_predict
+    assert self._hparams.sampling_method == "argmax"
+    assert self._decode_hparams.batch_size == 1
+    assert self._decode_hparams.block_size > 0
+    assert self._decode_hparams.block_size <= self._hparams.block_size
+    assert (
+        (self._decode_hparams.guess_and_check_top_k > 0) +
+        (self._decode_hparams.guess_and_check_epsilon >= 0) == 1)
+
+    inputs_old = features["inputs"]
+    assert "targets" not in features
+
+    assert len(features["inputs"].shape) in [3, 4]
+    if len(features["inputs"].shape) < 4:
+      features["inputs"] = tf.expand_dims(features["inputs"], 2)
+
+    block_size = self._decode_hparams.block_size
+    decode_length += tf.shape(features["inputs"])[1]
+
+    def while_exit_cond(result, length):  # pylint: disable=unused-argument
+      return length < decode_length
+
+    def infer_step(result, length):
+      """Inference step."""
+
+      def print_info(samples, result, length, new_length):
+        tf.logging.info(
+            "length=%s new_length=%s length_diff=%s samples-result=%s",
+            length,
+            new_length,
+            new_length - length,
+            np.array_str(
+                samples[0, -block_size-1:-1, 0, 0] -
+                result[0, -block_size:, 0, 0]
+            ).replace("\n", ""),
+        )
+
+      features["targets"] = tf.pad(result, [[0, 0], [0, 1], [0, 0], [0, 0]])
+      samples, logits, losses = self.sample(features)  # pylint: disable=unused-variable
+
+      _, top_k_indices = tf.nn.top_k(
+          logits[:, :-1, :1, :, :],
+          k=self._decode_hparams.guess_and_check_top_k)
+      in_top_k = tf.reduce_any(
+          tf.equal(tf.to_int64(top_k_indices), tf.expand_dims(result, 4)),
+          axis=4)
+
+      within_epsilon = tf.less_equal(
+          tf.abs(result - samples[:, :-1, :1, :]),
+          self._decode_hparams.guess_and_check_epsilon)
+
+      if self._decode_hparams.guess_and_check_top_k:
+        tf.logging.info(
+            "Using guess_and_check_top_k=%s",
+            self._decode_hparams.guess_and_check_top_k)
+        correct = in_top_k
+      else:
+        tf.logging.info(
+            "Using guess_and_check_epsilon=%s",
+            self._decode_hparams.guess_and_check_epsilon)
+        correct = within_epsilon
+
+      correct_cumsum = tf.cumsum(tf.to_int32(correct), axis=1)
+      perfect_cumsum = 1 + tf.range(tf.shape(correct)[1])
+      for axis in [0, 2, 3]:
+        perfect_cumsum = tf.expand_dims(perfect_cumsum, axis=axis)
+
+      new_length = tf.reduce_sum(
+          tf.to_int32(tf.equal(correct_cumsum, perfect_cumsum)), axis=1)
+      new_length = tf.squeeze(new_length, axis=[0, 1, 2])
+      new_length = tf.minimum(new_length, decode_length)
+
+      new_result = tf.concat([
+          result[:, :new_length, :, :],
+          tf.reshape(
+              samples[:, new_length, :block_size, :], [1, block_size, 1, 1])
+      ], axis=1)
+
+      with tf.control_dependencies([
+          tf.py_func(print_info, [samples, result, length, new_length], [])
+      ]):
+        new_result = tf.identity(new_result)
+
+      return new_result, new_length
+
+    result = tf.zeros((1, 0, 1, 1), dtype=tf.int64)
+    length = tf.squeeze(tf.zeros(1, dtype=tf.int32))
+
+    result, length = tf.while_loop(
+        while_exit_cond,
+        infer_step,
+        [result, length],
+        shape_invariants=[
+            tf.TensorShape([1, None, 1, 1]),
+            tf.TensorShape([]),
+        ],
+        back_prop=False,
+        parallel_iterations=1)
+
+    result = result[:, :length, :, :]
+
+    features["inputs"] = inputs_old
+
+    return {
+        "outputs": result,
+        "scores": None,
+    }
+
+
 @registry.register_hparams
 def image_transformer2d_base():
   """Set of hyperparameters."""
@@ -414,6 +678,76 @@ def img2img_transformer_b3():
   return hparams
 
 
+@registry.register_hparams
+def img2img_transformer_b3_bs1():
+  hparams = img2img_transformer_b3()
+  hparams.block_size = 1
+  return hparams
+
+
+@registry.register_hparams
+def img2img_transformer_b3_bs2():
+  hparams = img2img_transformer_b3()
+  hparams.block_size = 2
+  return hparams
+
+
+@registry.register_hparams
+def img2img_transformer_b3_bs3():
+  hparams = img2img_transformer_b3()
+  hparams.block_size = 3
+  return hparams
+
+
+@registry.register_hparams
+def img2img_transformer_b3_bs4():
+  hparams = img2img_transformer_b3()
+  hparams.block_size = 4
+  return hparams
+
+
+@registry.register_hparams
+def img2img_transformer_b3_bs5():
+  hparams = img2img_transformer_b3()
+  hparams.block_size = 5
+  return hparams
+
+
+@registry.register_hparams
+def img2img_transformer_b3_bs6():
+  hparams = img2img_transformer_b3()
+  hparams.block_size = 6
+  return hparams
+
+
+@registry.register_hparams
+def img2img_transformer_b3_bs7():
+  hparams = img2img_transformer_b3()
+  hparams.block_size = 7
+  return hparams
+
+
+@registry.register_hparams
+def img2img_transformer_b3_bs8():
+  hparams = img2img_transformer_b3()
+  hparams.block_size = 8
+  return hparams
+
+
+@registry.register_hparams
+def img2img_transformer_b3_bs9():
+  hparams = img2img_transformer_b3()
+  hparams.block_size = 9
+  return hparams
+
+
+@registry.register_hparams
+def img2img_transformer_b3_bs10():
+  hparams = img2img_transformer_b3()
+  hparams.block_size = 10
+  return hparams
+
+
 @registry.register_hparams
 def img2img_transformer_dilated():
   """Try dilated."""
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index f298a5bf2..7f53dea3e 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -49,6 +49,9 @@ def decode_hparams(overrides=""):
       beam_size=4,
       alpha=0.6,
       eos_penalty=0.0,
+      block_size=0,
+      guess_and_check_top_k=0,
+      guess_and_check_epsilon=-1,
       return_beams=False,
       write_beam_scores=False,
       max_input_size=-1,

From 814c4a13e1b03e3db444b31be2e6337c02c11bc7 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Sat, 28 Jul 2018 16:06:09 -0700
Subject: [PATCH 0402/2720] Switching to tf.scan for SV2P.

PiperOrigin-RevId: 206456825
---
 tensor2tensor/layers/modalities.py          |   4 +-
 tensor2tensor/models/research/next_frame.py | 159 ++++++++++++--------
 2 files changed, 100 insertions(+), 63 deletions(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index e1ae5401f..64046e0ac 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -700,7 +700,9 @@ def targets_bottom(self, x):  # pylint: disable=arguments-differ
     return common_layers.convert_rgb_to_real(x)
 
   def top(self, body_output, _):
-    frames = tf.stack(body_output, axis=1)
+    frames = body_output
+    if isinstance(body_output, list):
+      frames = tf.stack(body_output, axis=1)
     rgb_frames = common_layers.convert_real_to_rgb(frames)
     common_layers.summarize_video(rgb_frames, "body_output")
     return tf.expand_dims(rgb_frames, axis=-1)
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index fa061c028..a336651a5 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -238,24 +238,24 @@ def construct_latent_tower(self, images):
     conv_size = self.tinyify([32, 64, 64])
     with tf.variable_scope("latent", reuse=tf.AUTO_REUSE):
       # this allows more predicted frames at inference time
-      latent_images = images[:self.hparams.latent_num_frames]
-      images = tf.concat(latent_images, 3)
+      images = tf.unstack(images[:self.hparams.latent_num_frames], axis=0)
+      images = tf.concat(images, 3)
 
       x = images
       x = common_layers.make_even_size(x)
       x = tfl.conv2d(x, conv_size[0], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="latent_conv1")
-      x = tfl.batch_normalization(x,
-                                  training=self.is_training, name="latent_bn1")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="latent_bn1")
       x = common_layers.make_even_size(x)
       x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="latent_conv2")
-      x = tfl.batch_normalization(x,
-                                  training=self.is_training, name="latent_bn2")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="latent_bn2")
       x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(1, 1),
                      padding="SAME", activation=tf.nn.relu, name="latent_conv3")
-      x = tfl.batch_normalization(x,
-                                  training=self.is_training, name="latent_bn3")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="latent_bn3")
 
       nc = self.hparams.latent_channels
       mean = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
@@ -363,16 +363,16 @@ def reward_prediction(
           lstm_state, lstm_size, conv_size)
 
       x = hidden5
-      x = tfl.batch_normalization(x,
-                                  training=self.is_training, name="reward_bn0")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="reward_bn0")
       x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="reward_conv1")
-      x = tfl.batch_normalization(x,
-                                  training=self.is_training, name="reward_bn1")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="reward_bn1")
       x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="reward_conv2")
-      x = tfl.batch_normalization(x,
-                                  training=self.is_training, name="reward_bn2")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="reward_bn2")
       x = tfl.conv2d(x, conv_size[3], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="reward_conv3")
 
@@ -559,47 +559,56 @@ def construct_model(self,
     Raises:
       ValueError: if more than 1 mask specified for DNA model.
     """
-    batch_size = common_layers.shape_list(images[0])[0]
     context_frames = self.hparams.video_num_input_frames
 
-    # Predicted images and rewards.
-    gen_rewards, gen_images = [], []
+    def process_single_frame(prev_outputs, inputs):
+      """Process a single frame of the video."""
+      cur_image, cur_reward, action = inputs
+      time_step, prev_image, prev_reward, lstm_states = prev_outputs[:4]
 
-    # LSTM states.
-    lstm_state = [None] * 7
-    reward_lstm_state = [None] * 5
+      # TODO(mbz): No scheduled sampling for now!
+      input_image, input_reward = tf.cond(
+          tf.greater(time_step, context_frames),
+          lambda: (prev_image, prev_reward),
+          lambda: (cur_image, cur_reward))
+
+      # Prediction
+      pred_image, lstm_states = self.construct_predictive_tower(
+          input_image, input_reward, action, lstm_states, latent)
+
+      if self.hparams.reward_prediction:
+        reward_lstm_states = prev_outputs[4]
+        pred_reward, reward_lstm_states = self.reward_prediction(
+            input_image, input_reward, action, reward_lstm_states, latent)
+      else:
+        pred_reward = input_reward
+
+      time_step += 1
+      outputs = (time_step, pred_image, pred_reward, lstm_states)
+      if self.hparams.reward_prediction:
+        outputs += (reward_lstm_states,)
+
+      return outputs
 
     # Latent tower
+    latent = None
     if self.hparams.stochastic_model:
       latent_mean, latent_std = self.construct_latent_tower(images)
+      latent = self.get_gaussian_latent(latent_mean, latent_std)
 
-    pred_image, pred_reward, latent = None, None, None
-    for timestep, image, action, reward in zip(
-        range(len(images)-1), images[:-1], actions[:-1], rewards[:-1]):
-      # Scheduled Sampling
-      done_warm_start = timestep > context_frames - 1
-      groundtruth_items = [image, reward]
-      generated_items = [pred_image, pred_reward]
-      input_image, input_reward = self.get_scheduled_sample_inputs(
-          done_warm_start, groundtruth_items, generated_items, batch_size)
-
-      # Latent
-      if self.hparams.stochastic_model:
-        if timestep == 0 or self.hparams.multi_latent:
-          latent = self.get_gaussian_latent(latent_mean, latent_std)
+    # Initialize all the variables
+    lstm_states, reward_lstm_states = [None] * 7, [None] * 5
+    inputs = images[0], rewards[0], actions[0]
+    prev_outputs = (tf.constant(0), images[0], rewards[0], lstm_states)
+    if self.hparams.reward_prediction:
+      prev_outputs += (reward_lstm_states,)
 
-      # Prediction
-      pred_image, lstm_state = self.construct_predictive_tower(
-          input_image, input_reward, action, lstm_state, latent)
+    initializers = process_single_frame(prev_outputs, inputs)
 
-      if self.hparams.reward_prediction:
-        pred_reward, reward_lstm_state = self.reward_prediction(
-            input_image, input_reward, action, reward_lstm_state, latent)
-      else:
-        pred_reward = input_reward
+    inputs = (images[1:], actions[1:], rewards[1:])
 
-      gen_images.append(pred_image)
-      gen_rewards.append(pred_reward)
+    outputs = tf.scan(process_single_frame, inputs, initializers)
+    gen_images, gen_rewards = outputs[1:3]
 
     return gen_images, gen_rewards, [latent_mean], [latent_std]
 
@@ -775,15 +784,19 @@ def get_input_if_exists(self, features, key, batch_size, num_frames):
       x = features[key]
     else:
       x = tf.zeros((batch_size, num_frames, 1, self.hparams.hidden_size))
-    return tf.unstack(x, axis=1)
+    return self.swap_time_and_batch_axes(x)
+
+  def swap_time_and_batch_axes(self, x):
+    transposed_axes = tf.concat([[1, 0], tf.range(2, tf.rank(x))], axis=0)
+    return tf.transpose(x, transposed_axes)
 
   def body(self, features):
     hparams = self.hparams
     batch_size = common_layers.shape_list(features["inputs"])[0]
 
-    # Split inputs and targets time-wise into a list of frames.
-    input_frames = tf.unstack(features["inputs"], axis=1)
-    target_frames = tf.unstack(features["targets"], axis=1)
+    # Swap time and batch axes.
+    input_frames = self.swap_time_and_batch_axes(features["inputs"])
+    target_frames = self.swap_time_and_batch_axes(features["targets"])
 
     # Get actions if exist otherwise use zeros
     input_actions = self.get_input_if_exists(
@@ -797,15 +810,15 @@ def body(self, features):
     target_rewards = self.get_input_if_exists(
         features, "target_reward", batch_size, hparams.video_num_target_frames)
 
-    all_actions = input_actions + target_actions
-    all_rewards = input_rewards + target_rewards
-    all_frames = input_frames + target_frames
+    all_actions = tf.concat([input_actions, target_actions], axis=0)
+    all_rewards = tf.concat([input_rewards, target_rewards], axis=0)
+    all_frames = tf.concat([input_frames, target_frames], axis=0)
 
     # Each image is being used twice, in latent tower and main tower.
     # This is to make sure we are using the *same* image for both, ...
     # ... given how TF queues work.
     # NOT sure if this is required at all. Doesn"t hurt though! :)
-    all_frames = [tf.identity(frame) for frame in all_frames]
+    all_frames = tf.identity(all_frames)
 
     gen_images, gen_rewards, latent_means, latent_stds = self.construct_model(
         images=all_frames,
@@ -857,15 +870,29 @@ def anneal_loss(step_num):
       tf.summary.scalar("beta", beta)
       tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
 
-    extra_loss = beta * tf.reduce_mean(kl_loss)
+    extra_loss = beta * kl_loss
+
+    # Ignore the predictions from the input frames.
+    # This is NOT the same as original paper/implementation.
     predictions = gen_images[hparams.video_num_input_frames-1:]
-    reward_pred = tf.stack(
-        gen_rewards[hparams.video_num_input_frames-1:], axis=1)
+    reward_pred = gen_rewards[hparams.video_num_input_frames-1:]
     reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove undeeded dimension.
 
-    frames_gt = tf.concat(all_frames[hparams.video_num_input_frames:], axis=1)
-    frames_pd = tf.concat(predictions, axis=1)
-    tf.summary.image("full_video", tf.concat([frames_gt, frames_pd], axis=2))
+    # TODO(mbz): clean this up!
+    def fix_video_dims_and_concat_on_x_axis(x):
+      x = tf.transpose(x, [1, 3, 4, 0, 2])
+      x = tf.reshape(x, [batch_size, 64, 3, -1])
+      x = tf.transpose(x, [0, 3, 1, 2])
+      return x
+
+    frames_gd = fix_video_dims_and_concat_on_x_axis(target_frames)
+    frames_pd = fix_video_dims_and_concat_on_x_axis(predictions)
+    side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
+    tf.summary.image("full_video", side_by_side_video)
+
+    # Swap back time and batch axes.
+    predictions = self.swap_time_and_batch_axes(predictions)
+    reward_pred = self.swap_time_and_batch_axes(reward_pred)
 
     return_targets = predictions
     if "target_reward" in features:
@@ -879,6 +906,10 @@ class NextFrameStochasticTwoFrames(NextFrameStochastic):
   """Stochastic next-frame model with 2 frames posterior."""
 
   def construct_model(self, images, actions, rewards):
+    images = tf.unstack(images, axis=0)
+    actions = tf.unstack(actions, axis=0)
+    rewards = tf.unstack(rewards, axis=0)
+
     batch_size = common_layers.shape_list(images[0])[0]
     context_frames = self.hparams.video_num_input_frames
 
@@ -920,6 +951,9 @@ def construct_model(self, images, actions, rewards):
       gen_images.append(pred_image)
       gen_rewards.append(pred_reward)
 
+    gen_images = tf.stack(gen_images, axis=0)
+    gen_rewards = tf.stack(gen_rewards, axis=0)
+
     return gen_images, gen_rewards, latent_means, latent_stds
 
 
@@ -1104,7 +1138,7 @@ def construct_model(self, images, actions, rewards):
       pred_logvar: predicted log(var) of posterior
     """
     # model does not support action conditioned and reward prediction
-    fakr_reward_prediction = rewards
+    fake_reward_prediction = rewards
     del actions, rewards
 
     z_dim = self.hparams.z_dim
@@ -1114,8 +1148,7 @@ def construct_model(self, images, actions, rewards):
     predictor_rnn_layers = self.hparams.predictor_rnn_layers
     context_frames = self.hparams.video_num_input_frames
 
-    seq_len = len(images)
-    batch_size, _, _, color_channels = common_layers.shape_list(images[0])
+    seq_len, batch_size, _, _, color_channels = common_layers.shape_list(images)
 
     # LSTM initial sizesstates.
     predictor_states = [None] * predictor_rnn_layers
@@ -1124,6 +1157,7 @@ def construct_model(self, images, actions, rewards):
     tf.logging.info(">>>> Encoding")
     # Encoding:
     enc_images, enc_skips = [], []
+    images = tf.unstack(images, axis=0)
     for i, image in enumerate(images):
       with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
         enc, skips = self.encoder(image, rnn_size)
@@ -1175,7 +1209,8 @@ def construct_model(self, images, actions, rewards):
         gen_images.append(x_pred)
 
     tf.logging.info(">>>> Done")
-    return gen_images, fakr_reward_prediction, pred_mu, pred_logvar
+    gen_images = tf.stack(gen_images, axis=0)
+    return gen_images, fake_reward_prediction, pred_mu, pred_logvar
 
 
 @registry.register_hparams

From 9e6d737f066596ebaf1ba981bc5fa384e9b1bbc7 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Sun, 29 Jul 2018 22:32:28 -0700
Subject: [PATCH 0403/2720] fixing the hierarchy of video metrics in TB.

PiperOrigin-RevId: 206528847
---
 tensor2tensor/data_generators/video_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index ceebe7efa..11471d089 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -64,7 +64,7 @@ def summarize_video_metrics(hook_args):
   summary_values = []
   for name, array in six.iteritems(metrics_results):
     for ind, val in enumerate(array):
-      tag = name + "_" + str(ind)
+      tag = "metric_{}/{}".format(name, ind)
       summary_values.append(tf.Summary.Value(tag=tag, simple_value=val))
   return summary_values
 

From 17aaf775cb0f186ebb27bdbfd6883ebdbbae07a2 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 30 Jul 2018 10:18:19 -0700
Subject: [PATCH 0404/2720] Switching to tf.scan for SV2P.

PiperOrigin-RevId: 206603029
---
 tensor2tensor/layers/modalities.py          |   4 +-
 tensor2tensor/models/research/next_frame.py | 159 ++++++++------------
 2 files changed, 63 insertions(+), 100 deletions(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 64046e0ac..e1ae5401f 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -700,9 +700,7 @@ def targets_bottom(self, x):  # pylint: disable=arguments-differ
     return common_layers.convert_rgb_to_real(x)
 
   def top(self, body_output, _):
-    frames = body_output
-    if isinstance(body_output, list):
-      frames = tf.stack(body_output, axis=1)
+    frames = tf.stack(body_output, axis=1)
     rgb_frames = common_layers.convert_real_to_rgb(frames)
     common_layers.summarize_video(rgb_frames, "body_output")
     return tf.expand_dims(rgb_frames, axis=-1)
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index a336651a5..fa061c028 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -238,24 +238,24 @@ def construct_latent_tower(self, images):
     conv_size = self.tinyify([32, 64, 64])
     with tf.variable_scope("latent", reuse=tf.AUTO_REUSE):
       # this allows more predicted frames at inference time
-      images = tf.unstack(images[:self.hparams.latent_num_frames], axis=0)
-      images = tf.concat(images, 3)
+      latent_images = images[:self.hparams.latent_num_frames]
+      images = tf.concat(latent_images, 3)
 
       x = images
       x = common_layers.make_even_size(x)
       x = tfl.conv2d(x, conv_size[0], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="latent_conv1")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="latent_bn1")
+      x = tfl.batch_normalization(x,
+                                  training=self.is_training, name="latent_bn1")
       x = common_layers.make_even_size(x)
       x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="latent_conv2")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="latent_bn2")
+      x = tfl.batch_normalization(x,
+                                  training=self.is_training, name="latent_bn2")
       x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(1, 1),
                      padding="SAME", activation=tf.nn.relu, name="latent_conv3")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="latent_bn3")
+      x = tfl.batch_normalization(x,
+                                  training=self.is_training, name="latent_bn3")
 
       nc = self.hparams.latent_channels
       mean = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
@@ -363,16 +363,16 @@ def reward_prediction(
           lstm_state, lstm_size, conv_size)
 
       x = hidden5
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="reward_bn0")
+      x = tfl.batch_normalization(x,
+                                  training=self.is_training, name="reward_bn0")
       x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="reward_conv1")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="reward_bn1")
+      x = tfl.batch_normalization(x,
+                                  training=self.is_training, name="reward_bn1")
       x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="reward_conv2")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="reward_bn2")
+      x = tfl.batch_normalization(x,
+                                  training=self.is_training, name="reward_bn2")
       x = tfl.conv2d(x, conv_size[3], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="reward_conv3")
 
@@ -559,56 +559,47 @@ def construct_model(self,
     Raises:
       ValueError: if more than 1 mask specified for DNA model.
     """
+    batch_size = common_layers.shape_list(images[0])[0]
     context_frames = self.hparams.video_num_input_frames
 
-    def process_single_frame(prev_outputs, inputs):
-      """Process a single frame of the video."""
-      cur_image, cur_reward, action = inputs
-      time_step, prev_image, prev_reward, lstm_states = prev_outputs[:4]
-
-      # TODO(mbz): No scheduled sampling for now!
-      input_image, input_reward = tf.cond(
-          tf.greater(time_step, context_frames),
-          lambda: (prev_image, prev_reward),
-          lambda: (cur_image, cur_reward))
-
-      # Prediction
-      pred_image, lstm_states = self.construct_predictive_tower(
-          input_image, input_reward, action, lstm_states, latent)
-
-      if self.hparams.reward_prediction:
-        reward_lstm_states = prev_outputs[4]
-        pred_reward, reward_lstm_states = self.reward_prediction(
-            input_image, input_reward, action, reward_lstm_states, latent)
-      else:
-        pred_reward = input_reward
-
-      time_step += 1
-      outputs = (time_step, pred_image, pred_reward, lstm_states)
-      if self.hparams.reward_prediction:
-        outputs += (reward_lstm_states,)
+    # Predicted images and rewards.
+    gen_rewards, gen_images = [], []
 
-      return outputs
+    # LSTM states.
+    lstm_state = [None] * 7
+    reward_lstm_state = [None] * 5
 
     # Latent tower
-    latent = None
     if self.hparams.stochastic_model:
       latent_mean, latent_std = self.construct_latent_tower(images)
-      latent = self.get_gaussian_latent(latent_mean, latent_std)
 
-    # Initialize all the variables
-    lstm_states, reward_lstm_states = [None] * 7, [None] * 5
-    inputs = images[0], rewards[0], actions[0]
-    prev_outputs = (tf.constant(0), images[0], rewards[0], lstm_states)
-    if self.hparams.reward_prediction:
-      prev_outputs += (reward_lstm_states,)
+    pred_image, pred_reward, latent = None, None, None
+    for timestep, image, action, reward in zip(
+        range(len(images)-1), images[:-1], actions[:-1], rewards[:-1]):
+      # Scheduled Sampling
+      done_warm_start = timestep > context_frames - 1
+      groundtruth_items = [image, reward]
+      generated_items = [pred_image, pred_reward]
+      input_image, input_reward = self.get_scheduled_sample_inputs(
+          done_warm_start, groundtruth_items, generated_items, batch_size)
+
+      # Latent
+      if self.hparams.stochastic_model:
+        if timestep == 0 or self.hparams.multi_latent:
+          latent = self.get_gaussian_latent(latent_mean, latent_std)
 
-    initializers = process_single_frame(prev_outputs, inputs)
+      # Prediction
+      pred_image, lstm_state = self.construct_predictive_tower(
+          input_image, input_reward, action, lstm_state, latent)
 
-    inputs = (images[1:], actions[1:], rewards[1:])
+      if self.hparams.reward_prediction:
+        pred_reward, reward_lstm_state = self.reward_prediction(
+            input_image, input_reward, action, reward_lstm_state, latent)
+      else:
+        pred_reward = input_reward
 
-    outputs = tf.scan(process_single_frame, inputs, initializers)
-    gen_images, gen_rewards = outputs[1:3]
+      gen_images.append(pred_image)
+      gen_rewards.append(pred_reward)
 
     return gen_images, gen_rewards, [latent_mean], [latent_std]
 
@@ -784,19 +775,15 @@ def get_input_if_exists(self, features, key, batch_size, num_frames):
       x = features[key]
     else:
       x = tf.zeros((batch_size, num_frames, 1, self.hparams.hidden_size))
-    return self.swap_time_and_batch_axes(x)
-
-  def swap_time_and_batch_axes(self, x):
-    transposed_axes = tf.concat([[1, 0], tf.range(2, tf.rank(x))], axis=0)
-    return tf.transpose(x, transposed_axes)
+    return tf.unstack(x, axis=1)
 
   def body(self, features):
     hparams = self.hparams
     batch_size = common_layers.shape_list(features["inputs"])[0]
 
-    # Swap time and batch axes.
-    input_frames = self.swap_time_and_batch_axes(features["inputs"])
-    target_frames = self.swap_time_and_batch_axes(features["targets"])
+    # Split inputs and targets time-wise into a list of frames.
+    input_frames = tf.unstack(features["inputs"], axis=1)
+    target_frames = tf.unstack(features["targets"], axis=1)
 
     # Get actions if exist otherwise use zeros
     input_actions = self.get_input_if_exists(
@@ -810,15 +797,15 @@ def body(self, features):
     target_rewards = self.get_input_if_exists(
         features, "target_reward", batch_size, hparams.video_num_target_frames)
 
-    all_actions = tf.concat([input_actions, target_actions], axis=0)
-    all_rewards = tf.concat([input_rewards, target_rewards], axis=0)
-    all_frames = tf.concat([input_frames, target_frames], axis=0)
+    all_actions = input_actions + target_actions
+    all_rewards = input_rewards + target_rewards
+    all_frames = input_frames + target_frames
 
     # Each image is being used twice, in latent tower and main tower.
     # This is to make sure we are using the *same* image for both, ...
     # ... given how TF queues work.
     # NOT sure if this is required at all. Doesn"t hurt though! :)
-    all_frames = tf.identity(all_frames)
+    all_frames = [tf.identity(frame) for frame in all_frames]
 
     gen_images, gen_rewards, latent_means, latent_stds = self.construct_model(
         images=all_frames,
@@ -870,29 +857,15 @@ def anneal_loss(step_num):
       tf.summary.scalar("beta", beta)
       tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
 
-    extra_loss = beta * kl_loss
-
-    # Ignore the predictions from the input frames.
-    # This is NOT the same as original paper/implementation.
+    extra_loss = beta * tf.reduce_mean(kl_loss)
     predictions = gen_images[hparams.video_num_input_frames-1:]
-    reward_pred = gen_rewards[hparams.video_num_input_frames-1:]
+    reward_pred = tf.stack(
+        gen_rewards[hparams.video_num_input_frames-1:], axis=1)
     reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove undeeded dimension.
 
-    # TODO(mbz): clean this up!
-    def fix_video_dims_and_concat_on_x_axis(x):
-      x = tf.transpose(x, [1, 3, 4, 0, 2])
-      x = tf.reshape(x, [batch_size, 64, 3, -1])
-      x = tf.transpose(x, [0, 3, 1, 2])
-      return x
-
-    frames_gd = fix_video_dims_and_concat_on_x_axis(target_frames)
-    frames_pd = fix_video_dims_and_concat_on_x_axis(predictions)
-    side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
-    tf.summary.image("full_video", side_by_side_video)
-
-    # Swap back time and batch axes.
-    predictions = self.swap_time_and_batch_axes(predictions)
-    reward_pred = self.swap_time_and_batch_axes(reward_pred)
+    frames_gt = tf.concat(all_frames[hparams.video_num_input_frames:], axis=1)
+    frames_pd = tf.concat(predictions, axis=1)
+    tf.summary.image("full_video", tf.concat([frames_gt, frames_pd], axis=2))
 
     return_targets = predictions
     if "target_reward" in features:
@@ -906,10 +879,6 @@ class NextFrameStochasticTwoFrames(NextFrameStochastic):
   """Stochastic next-frame model with 2 frames posterior."""
 
   def construct_model(self, images, actions, rewards):
-    images = tf.unstack(images, axis=0)
-    actions = tf.unstack(actions, axis=0)
-    rewards = tf.unstack(rewards, axis=0)
-
     batch_size = common_layers.shape_list(images[0])[0]
     context_frames = self.hparams.video_num_input_frames
 
@@ -951,9 +920,6 @@ def construct_model(self, images, actions, rewards):
       gen_images.append(pred_image)
       gen_rewards.append(pred_reward)
 
-    gen_images = tf.stack(gen_images, axis=0)
-    gen_rewards = tf.stack(gen_rewards, axis=0)
-
     return gen_images, gen_rewards, latent_means, latent_stds
 
 
@@ -1138,7 +1104,7 @@ def construct_model(self, images, actions, rewards):
       pred_logvar: predicted log(var) of posterior
     """
     # model does not support action conditioned and reward prediction
-    fake_reward_prediction = rewards
+    fakr_reward_prediction = rewards
     del actions, rewards
 
     z_dim = self.hparams.z_dim
@@ -1148,7 +1114,8 @@ def construct_model(self, images, actions, rewards):
     predictor_rnn_layers = self.hparams.predictor_rnn_layers
     context_frames = self.hparams.video_num_input_frames
 
-    seq_len, batch_size, _, _, color_channels = common_layers.shape_list(images)
+    seq_len = len(images)
+    batch_size, _, _, color_channels = common_layers.shape_list(images[0])
 
     # LSTM initial sizesstates.
     predictor_states = [None] * predictor_rnn_layers
@@ -1157,7 +1124,6 @@ def construct_model(self, images, actions, rewards):
     tf.logging.info(">>>> Encoding")
     # Encoding:
     enc_images, enc_skips = [], []
-    images = tf.unstack(images, axis=0)
     for i, image in enumerate(images):
       with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
         enc, skips = self.encoder(image, rnn_size)
@@ -1209,8 +1175,7 @@ def construct_model(self, images, actions, rewards):
         gen_images.append(x_pred)
 
     tf.logging.info(">>>> Done")
-    gen_images = tf.stack(gen_images, axis=0)
-    return gen_images, fake_reward_prediction, pred_mu, pred_logvar
+    return gen_images, fakr_reward_prediction, pred_mu, pred_logvar
 
 
 @registry.register_hparams

From cce566273cfa9d91484ff82393b16c28648ebebd Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 30 Jul 2018 11:37:00 -0700
Subject: [PATCH 0405/2720] part of the rolled back changes.

PiperOrigin-RevId: 206618354
---
 tensor2tensor/models/research/next_frame.py | 28 ++++++++++-----------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index fa061c028..f6dc8b758 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -245,17 +245,17 @@ def construct_latent_tower(self, images):
       x = common_layers.make_even_size(x)
       x = tfl.conv2d(x, conv_size[0], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="latent_conv1")
-      x = tfl.batch_normalization(x,
-                                  training=self.is_training, name="latent_bn1")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="latent_bn1")
       x = common_layers.make_even_size(x)
       x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="latent_conv2")
-      x = tfl.batch_normalization(x,
-                                  training=self.is_training, name="latent_bn2")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="latent_bn2")
       x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(1, 1),
                      padding="SAME", activation=tf.nn.relu, name="latent_conv3")
-      x = tfl.batch_normalization(x,
-                                  training=self.is_training, name="latent_bn3")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="latent_bn3")
 
       nc = self.hparams.latent_channels
       mean = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
@@ -363,16 +363,16 @@ def reward_prediction(
           lstm_state, lstm_size, conv_size)
 
       x = hidden5
-      x = tfl.batch_normalization(x,
-                                  training=self.is_training, name="reward_bn0")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="reward_bn0")
       x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="reward_conv1")
-      x = tfl.batch_normalization(x,
-                                  training=self.is_training, name="reward_bn1")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="reward_bn1")
       x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="reward_conv2")
-      x = tfl.batch_normalization(x,
-                                  training=self.is_training, name="reward_bn2")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="reward_bn2")
       x = tfl.conv2d(x, conv_size[3], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="reward_conv3")
 
@@ -1104,7 +1104,7 @@ def construct_model(self, images, actions, rewards):
       pred_logvar: predicted log(var) of posterior
     """
     # model does not support action conditioned and reward prediction
-    fakr_reward_prediction = rewards
+    fake_reward_prediction = rewards
     del actions, rewards
 
     z_dim = self.hparams.z_dim
@@ -1175,7 +1175,7 @@ def construct_model(self, images, actions, rewards):
         gen_images.append(x_pred)
 
     tf.logging.info(">>>> Done")
-    return gen_images, fakr_reward_prediction, pred_mu, pred_logvar
+    return gen_images, fake_reward_prediction, pred_mu, pred_logvar
 
 
 @registry.register_hparams

From 44b5e4b37c3a6f4216d53364a57adf636300e459 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 30 Jul 2018 12:08:12 -0700
Subject: [PATCH 0406/2720] separating next_frame_params to reduce the
 filesize.

PiperOrigin-RevId: 206623734
---
 tensor2tensor/models/research/next_frame.py   | 214 +---------------
 .../models/research/next_frame_params.py      | 234 ++++++++++++++++++
 2 files changed, 235 insertions(+), 213 deletions(-)
 create mode 100644 tensor2tensor/models/research/next_frame_params.py

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index f6dc8b758..2e77cfb9b 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -21,8 +21,8 @@
 import six
 
 from tensor2tensor.layers import common_attention
-from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.models.research import next_frame_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -1178,218 +1178,6 @@ def construct_model(self, images, actions, rewards):
     return gen_images, fake_reward_prediction, pred_mu, pred_logvar
 
 
-@registry.register_hparams
-def next_frame():
-  """Basic 2-frame conv model."""
-  hparams = common_hparams.basic_params1()
-  hparams.video_num_input_frames = 4
-  hparams.video_num_target_frames = 1
-  hparams.hidden_size = 64
-  hparams.batch_size = 4
-  hparams.num_hidden_layers = 2
-  hparams.optimizer = "Adafactor"
-  hparams.learning_rate_constant = 1.5
-  hparams.learning_rate_warmup_steps = 1500
-  hparams.learning_rate_schedule = "linear_warmup * constant * rsqrt_decay"
-  hparams.label_smoothing = 0.0
-  hparams.initializer = "uniform_unit_scaling"
-  hparams.initializer_gain = 1.3
-  hparams.weight_decay = 0.0
-  hparams.clip_grad_norm = 1.0
-  hparams.dropout = 0.5
-  hparams.add_hparam("num_compress_steps", 6)
-  hparams.add_hparam("filter_double_steps", 2)
-  hparams.add_hparam("video_modality_loss_cutoff", 0.02)
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_stochastic():
-  """SV2P model."""
-  hparams = next_frame()
-  hparams.optimizer = "TrueAdam"
-  hparams.learning_rate_schedule = "constant"
-  hparams.learning_rate_constant = 1e-3
-  hparams.video_num_input_frames = 1
-  hparams.video_num_target_frames = 3
-  hparams.batch_size = 16
-  hparams.target_modality = "video:l2raw"
-  hparams.input_modalities = "inputs:video:l2raw"
-  hparams.video_modality_loss_cutoff = 0.0
-  hparams.add_hparam("stochastic_model", True)
-  hparams.add_hparam("reward_prediction", True)
-  hparams.add_hparam("model_options", "CDNA")
-  hparams.add_hparam("num_masks", 10)
-  hparams.add_hparam("latent_channels", 1)
-  hparams.add_hparam("latent_std_min", -5.0)
-  hparams.add_hparam("num_iterations_2nd_stage", 10000)
-  hparams.add_hparam("latent_loss_multiplier", 1e-3)
-  hparams.add_hparam("latent_loss_multiplier_schedule", "constant")
-  hparams.add_hparam("multi_latent", False)
-  hparams.add_hparam("relu_shift", 1e-12)
-  hparams.add_hparam("dna_kernel_size", 5)
-  hparams.add_hparam("scheduled_sampling_k", 900.0)
-  hparams.add_hparam(
-      "latent_num_frames",  # use all frames by default.
-      hparams.video_num_input_frames + hparams.video_num_target_frames)
-  hparams.add_hparam("tiny_mode", False)
-  hparams.add_hparam("anneal_end", 100000)
-  hparams.add_hparam("upsample_method", "conv2d_transpose")
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_stochastic_emily():
-  """Emily's model."""
-  hparams = next_frame_stochastic()
-  hparams.latent_loss_multiplier = 1e-4
-  hparams.learning_rate_constant = 0.002
-  hparams.add_hparam("z_dim", 10)
-  hparams.add_hparam("g_dim", 128)
-  hparams.add_hparam("rnn_size", 256)
-  hparams.add_hparam("posterior_rnn_layers", 1)
-  hparams.add_hparam("predictor_rnn_layers", 2)
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_savp():
-  """SVAP model."""
-  hparams = next_frame_stochastic()
-  hparams.add_hparam("z_dim", 8)
-  hparams.target_modality = "video:l1raw"
-  hparams.input_modalities = "inputs:video:l1raw"
-  hparams.latent_loss_multiplier_schedule = "linear_anneal"
-  hparams.anneal_end = 100000
-  hparams.upsample_method = "bilinear_upsample_conv"
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_stochastic_cutoff():
-  """SV2P model with additional cutoff in L2 loss for environments like pong."""
-  hparams = next_frame_stochastic()
-  hparams.video_modality_loss_cutoff = 0.4
-  hparams.video_num_input_frames = 4
-  hparams.video_num_target_frames = 1
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_stochastic_tiny():
-  """SV2P model with additional cutoff in L2 loss for environments like pong."""
-  hparams = next_frame_stochastic()
-  hparams.batch_size = 2
-  hparams.tiny_mode = True
-  hparams.num_masks = 1
-  hparams.video_modality_loss_cutoff = 0.4
-  hparams.video_num_input_frames = 4
-  hparams.video_num_target_frames = 1
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_tpu():
-  hparams = next_frame()
-  hparams.batch_size = 1
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_ae():
-  """Conv autoencoder."""
-  hparams = next_frame()
-  hparams.input_modalities = "inputs:video:bitwise"
-  hparams.hidden_size = 256
-  hparams.batch_size = 8
-  hparams.num_hidden_layers = 4
-  hparams.num_compress_steps = 4
-  hparams.dropout = 0.4
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_small():
-  """Small conv model."""
-  hparams = next_frame()
-  hparams.hidden_size = 32
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_tiny():
-  """Tiny for testing."""
-  hparams = next_frame()
-  hparams.hidden_size = 32
-  hparams.num_hidden_layers = 1
-  hparams.num_compress_steps = 2
-  hparams.filter_double_steps = 1
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_l1():
-  """Basic conv model with L1 modality."""
-  hparams = next_frame()
-  hparams.target_modality = "video:l1"
-  hparams.video_modality_loss_cutoff = 2.4
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_l2():
-  """Basic conv model with L2 modality."""
-  hparams = next_frame()
-  hparams.target_modality = "video:l2"
-  hparams.video_modality_loss_cutoff = 2.4
-  return hparams
-
-
-@registry.register_ranged_hparams
-def next_frame_base_range(rhp):
-  """Basic tuning grid."""
-  rhp.set_float("dropout", 0.2, 0.6)
-  rhp.set_discrete("hidden_size", [64, 128, 256])
-  rhp.set_int("num_compress_steps", 5, 8)
-  rhp.set_discrete("batch_size", [4, 8, 16, 32])
-  rhp.set_int("num_hidden_layers", 1, 3)
-  rhp.set_int("filter_double_steps", 1, 6)
-  rhp.set_float("learning_rate_constant", 1., 4.)
-  rhp.set_int("learning_rate_warmup_steps", 500, 3000)
-  rhp.set_float("initializer_gain", 0.8, 1.8)
-
-
-@registry.register_ranged_hparams
-def next_frame_doubling_range(rhp):
-  """Filter doubling and dropout tuning grid."""
-  rhp.set_float("dropout", 0.2, 0.6)
-  rhp.set_int("filter_double_steps", 2, 5)
-
-
-@registry.register_ranged_hparams
-def next_frame_clipgrad_range(rhp):
-  """Filter doubling and dropout tuning grid."""
-  rhp.set_float("dropout", 0.3, 0.4)
-  rhp.set_float("clip_grad_norm", 0.5, 10.0)
-
-
-@registry.register_ranged_hparams
-def next_frame_xent_cutoff_range(rhp):
-  """Cross-entropy tuning grid."""
-  rhp.set_float("video_modality_loss_cutoff", 0.005, 0.05)
-
-
-@registry.register_ranged_hparams
-def next_frame_ae_range(rhp):
-  """Autoencoder world model tuning grid."""
-  rhp.set_float("dropout", 0.3, 0.5)
-  rhp.set_int("num_compress_steps", 1, 3)
-  rhp.set_int("num_hidden_layers", 2, 6)
-  rhp.set_float("learning_rate_constant", 1., 2.)
-  rhp.set_float("initializer_gain", 0.8, 1.5)
-  rhp.set_int("filter_double_steps", 2, 3)
-
-
 @registry.register_model
 class NextFrameSavp(NextFrameStochastic):
   """Stochastic Adversarial Video Prediction."""
diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
new file mode 100644
index 000000000..5217d7824
--- /dev/null
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -0,0 +1,234 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Param sets for next frame prediction models."""
+
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_hparams
+from tensor2tensor.utils import registry
+
+
+@registry.register_hparams
+def next_frame():
+  """Basic 2-frame conv model."""
+  hparams = common_hparams.basic_params1()
+  hparams.video_num_input_frames = 4
+  hparams.video_num_target_frames = 1
+  hparams.hidden_size = 64
+  hparams.batch_size = 4
+  hparams.num_hidden_layers = 2
+  hparams.optimizer = "Adafactor"
+  hparams.learning_rate_constant = 1.5
+  hparams.learning_rate_warmup_steps = 1500
+  hparams.learning_rate_schedule = "linear_warmup * constant * rsqrt_decay"
+  hparams.label_smoothing = 0.0
+  hparams.initializer = "uniform_unit_scaling"
+  hparams.initializer_gain = 1.3
+  hparams.weight_decay = 0.0
+  hparams.clip_grad_norm = 1.0
+  hparams.dropout = 0.5
+  hparams.add_hparam("num_compress_steps", 6)
+  hparams.add_hparam("filter_double_steps", 2)
+  hparams.add_hparam("video_modality_loss_cutoff", 0.02)
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_stochastic():
+  """SV2P model."""
+  hparams = next_frame()
+  hparams.optimizer = "TrueAdam"
+  hparams.learning_rate_schedule = "constant"
+  hparams.learning_rate_constant = 1e-3
+  hparams.video_num_input_frames = 1
+  hparams.video_num_target_frames = 3
+  hparams.batch_size = 16
+  hparams.target_modality = "video:l2raw"
+  hparams.input_modalities = "inputs:video:l2raw"
+  hparams.video_modality_loss_cutoff = 0.0
+  hparams.add_hparam("stochastic_model", True)
+  hparams.add_hparam("reward_prediction", True)
+  hparams.add_hparam("model_options", "CDNA")
+  hparams.add_hparam("num_masks", 10)
+  hparams.add_hparam("latent_channels", 1)
+  hparams.add_hparam("latent_std_min", -5.0)
+  hparams.add_hparam("num_iterations_2nd_stage", 10000)
+  hparams.add_hparam("latent_loss_multiplier", 1e-3)
+  hparams.add_hparam("latent_loss_multiplier_schedule", "constant")
+  hparams.add_hparam("multi_latent", False)
+  hparams.add_hparam("relu_shift", 1e-12)
+  hparams.add_hparam("dna_kernel_size", 5)
+  hparams.add_hparam("scheduled_sampling_k", 900.0)
+  hparams.add_hparam(
+      "latent_num_frames",  # use all frames by default.
+      hparams.video_num_input_frames + hparams.video_num_target_frames)
+  hparams.add_hparam("tiny_mode", False)
+  hparams.add_hparam("anneal_end", 100000)
+  hparams.add_hparam("upsample_method", "conv2d_transpose")
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_stochastic_emily():
+  """Emily's model."""
+  hparams = next_frame_stochastic()
+  hparams.latent_loss_multiplier = 1e-4
+  hparams.learning_rate_constant = 0.002
+  hparams.add_hparam("z_dim", 10)
+  hparams.add_hparam("g_dim", 128)
+  hparams.add_hparam("rnn_size", 256)
+  hparams.add_hparam("posterior_rnn_layers", 1)
+  hparams.add_hparam("predictor_rnn_layers", 2)
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_savp():
+  """SVAP model."""
+  hparams = next_frame_stochastic()
+  hparams.add_hparam("z_dim", 8)
+  hparams.target_modality = "video:l1raw"
+  hparams.input_modalities = "inputs:video:l1raw"
+  hparams.latent_loss_multiplier_schedule = "linear_anneal"
+  hparams.anneal_end = 100000
+  hparams.upsample_method = "bilinear_upsample_conv"
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_stochastic_cutoff():
+  """SV2P model with additional cutoff in L2 loss for environments like pong."""
+  hparams = next_frame_stochastic()
+  hparams.video_modality_loss_cutoff = 0.4
+  hparams.video_num_input_frames = 4
+  hparams.video_num_target_frames = 1
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_stochastic_tiny():
+  """SV2P model with additional cutoff in L2 loss for environments like pong."""
+  hparams = next_frame_stochastic()
+  hparams.batch_size = 2
+  hparams.tiny_mode = True
+  hparams.num_masks = 1
+  hparams.video_modality_loss_cutoff = 0.4
+  hparams.video_num_input_frames = 4
+  hparams.video_num_target_frames = 1
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_tpu():
+  hparams = next_frame()
+  hparams.batch_size = 1
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_ae():
+  """Conv autoencoder."""
+  hparams = next_frame()
+  hparams.input_modalities = "inputs:video:bitwise"
+  hparams.hidden_size = 256
+  hparams.batch_size = 8
+  hparams.num_hidden_layers = 4
+  hparams.num_compress_steps = 4
+  hparams.dropout = 0.4
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_small():
+  """Small conv model."""
+  hparams = next_frame()
+  hparams.hidden_size = 32
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_tiny():
+  """Tiny for testing."""
+  hparams = next_frame()
+  hparams.hidden_size = 32
+  hparams.num_hidden_layers = 1
+  hparams.num_compress_steps = 2
+  hparams.filter_double_steps = 1
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_l1():
+  """Basic conv model with L1 modality."""
+  hparams = next_frame()
+  hparams.target_modality = "video:l1"
+  hparams.video_modality_loss_cutoff = 2.4
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_l2():
+  """Basic conv model with L2 modality."""
+  hparams = next_frame()
+  hparams.target_modality = "video:l2"
+  hparams.video_modality_loss_cutoff = 2.4
+  return hparams
+
+
+@registry.register_ranged_hparams
+def next_frame_base_range(rhp):
+  """Basic tuning grid."""
+  rhp.set_float("dropout", 0.2, 0.6)
+  rhp.set_discrete("hidden_size", [64, 128, 256])
+  rhp.set_int("num_compress_steps", 5, 8)
+  rhp.set_discrete("batch_size", [4, 8, 16, 32])
+  rhp.set_int("num_hidden_layers", 1, 3)
+  rhp.set_int("filter_double_steps", 1, 6)
+  rhp.set_float("learning_rate_constant", 1., 4.)
+  rhp.set_int("learning_rate_warmup_steps", 500, 3000)
+  rhp.set_float("initializer_gain", 0.8, 1.8)
+
+
+@registry.register_ranged_hparams
+def next_frame_doubling_range(rhp):
+  """Filter doubling and dropout tuning grid."""
+  rhp.set_float("dropout", 0.2, 0.6)
+  rhp.set_int("filter_double_steps", 2, 5)
+
+
+@registry.register_ranged_hparams
+def next_frame_clipgrad_range(rhp):
+  """Filter doubling and dropout tuning grid."""
+  rhp.set_float("dropout", 0.3, 0.4)
+  rhp.set_float("clip_grad_norm", 0.5, 10.0)
+
+
+@registry.register_ranged_hparams
+def next_frame_xent_cutoff_range(rhp):
+  """Cross-entropy tuning grid."""
+  rhp.set_float("video_modality_loss_cutoff", 0.005, 0.05)
+
+
+@registry.register_ranged_hparams
+def next_frame_ae_range(rhp):
+  """Autoencoder world model tuning grid."""
+  rhp.set_float("dropout", 0.3, 0.5)
+  rhp.set_int("num_compress_steps", 1, 3)
+  rhp.set_int("num_hidden_layers", 2, 6)
+  rhp.set_float("learning_rate_constant", 1., 2.)
+  rhp.set_float("initializer_gain", 0.8, 1.5)
+  rhp.set_int("filter_double_steps", 2, 3)
+

From a595dc5d3737e5064017eef948482bb78a0af2dd Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 30 Jul 2018 12:48:19 -0700
Subject: [PATCH 0407/2720] Add a param to toggle wiki text normalization

PiperOrigin-RevId: 206629880
---
 tensor2tensor/data_generators/wikisum/wikisum.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/wikisum/wikisum.py b/tensor2tensor/data_generators/wikisum/wikisum.py
index 9fe1e7971..b7ec85231 100644
--- a/tensor2tensor/data_generators/wikisum/wikisum.py
+++ b/tensor2tensor/data_generators/wikisum/wikisum.py
@@ -341,23 +341,25 @@ def _tokens_to_score(tokens):
   return {t for t in tokens if re.search("[a-z0-9]", t)}
 
 
-def rank_reference_paragraphs(wiki_title, references_content):
+def rank_reference_paragraphs(wiki_title, references_content, normalize=True):
   """Rank and return reference paragraphs by tf-idf score on title tokens."""
-  title_tokens = _tokens_to_score(set(
-      tokenizer.encode(text_encoder.native_to_unicode(wiki_title))))
+  normalized_title = _normalize_text(wiki_title)
+  title_tokens = _tokens_to_score(
+      set(tokenizer.encode(text_encoder.native_to_unicode(normalized_title))))
   ref_paragraph_info = []
   doc_counts = collections.defaultdict(int)
   for ref in references_content:
     for paragraph in ref.split("\n"):
-      paragraph = _normalize_text(paragraph)
-      if cc_utils.filter_paragraph(paragraph):
+      normalized_paragraph = _normalize_text(paragraph)
+      if cc_utils.filter_paragraph(normalized_paragraph):
         # Skip paragraph
         continue
-      counts = _token_counts(paragraph, title_tokens)
+      counts = _token_counts(normalized_paragraph, title_tokens)
       for token in title_tokens:
         if counts[token]:
           doc_counts[token] += 1
-      info = {"content": paragraph, "counts": counts}
+      content = normalized_paragraph if normalize else paragraph
+      info = {"content": content, "counts": counts}
       ref_paragraph_info.append(info)
 
   for info in ref_paragraph_info:

From aa60bee802f3579c5cafe47345bc042e1a4feefc Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 30 Jul 2018 16:48:59 -0700
Subject: [PATCH 0408/2720] Enable RL experiments on all Atari games

PiperOrigin-RevId: 206670359
---
 tensor2tensor/data_generators/gym_problems.py | 14 +++-
 .../data_generators/gym_problems_specs.py     | 67 +++++++++++++++++++
 tensor2tensor/data_generators/gym_utils.py    | 23 +++++++
 tensor2tensor/rl/model_rl_experiment.py       | 10 ++-
 4 files changed, 110 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 470658da8..09f981661 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -131,8 +131,18 @@ def random_skip(self):
 
   def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
     self._setup()
-    self.debug_dump_frames_path = os.path.join(data_dir,
-                                               self.debug_dump_frames_path)
+
+    # We only want to save frames for eval and simulated experience, not the
+    # frames used for world model training.
+    base_dir = os.path.basename(os.path.dirname(data_dir + "/"))
+    if (base_dir == "eval" or self.debug_dump_frames_path in [
+        "debug_frames_sim_eval", "debug_frames_sim"
+    ]):
+      self.debug_dump_frames_path = os.path.join(data_dir,
+                                                 self.debug_dump_frames_path)
+    else:
+      # Disable frame saving
+      self.debug_dump_frames_path = ""
 
     with self._session as sess:
       frame_counter = 0
diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index db0383aea..2632cb730 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import gym
 
 # We need gym_utils for the game environments defined there.
 from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
@@ -29,6 +30,25 @@
 # pylint: enable=g-multiple-import
 from tensor2tensor.utils import registry
 
+ATARI_GAMES = [
+    "air_raid", "alien", "amidar", "assault", "asterix", "asteroids",
+    "atlantis", "bank_heist", "battle_zone", "beam_rider", "berzerk", "bowling",
+    "boxing", "breakout", "carnival", "centipede", "chopper_command",
+    "crazy_climber", "demon_attack", "double_dunk", "elevator_action", "enduro",
+    "fishing_derby", "freeway", "frostbite", "gopher", "gravitar", "hero",
+    "ice_hockey", "jamesbond", "journey_escape", "kangaroo", "krull",
+    "kung_fu_master", "montezuma_revenge", "ms_pacman", "name_this_game",
+    "phoenix", "pitfall", "pong", "pooyan", "private_eye", "qbert", "riverraid",
+    "road_runner", "robotank", "seaquest", "skiing", "solaris",
+    "space_invaders", "star_gunner", "tennis", "time_pilot", "tutankham",
+    "up_n_down", "venture", "video_pinball", "wizard_of_wor", "yars_revenge",
+    "zaxxon"
+]
+# Removed because XDeterministic-v4 did not exist:
+# * adventure
+# * defender
+# * kaboom
+
 
 @registry.register_problem
 class GymPongRandom(GymDiscreteProblem):
@@ -283,3 +303,50 @@ def initial_frames_problem(self):
   def num_testing_steps(self):
     return 100
 
+
+class GymClippedRewardRandom(GymDiscreteProblem):
+  """Base class for clipped reward games."""
+
+  @property
+  def env_name(self):
+    raise NotImplementedError
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+
+def dynamically_create_gym_clipped_reward_problem(game_name):
+  """Dynamically create env wrapper and Problems for game."""
+  # e.g. game_name == bank_heist
+  assert game_name in ATARI_GAMES
+  camel_game_name = "".join(
+      [w[0].upper() + w[1:] for w in game_name.split("_")])
+  env_name = "%sDeterministic-v4" % camel_game_name
+  wrapped_env_name = "T2T%s" % env_name
+
+  # Register an environment that does the reward clipping
+  gym.envs.register(
+      id=wrapped_env_name,
+      entry_point=lambda: gym_utils.wrapped_factory(  # pylint: disable=g-long-lambda
+          env=env_name, reward_clipping=True))
+
+  # Create and register the Random and WithAgent Problem classes
+  problem_cls = type(camel_game_name + "Random", (GymClippedRewardRandom,),
+                     {"env_name": wrapped_env_name})
+  with_agent_cls = type("GymDiscreteProblemWithAgentOn%s" % camel_game_name,
+                        (GymRealDiscreteProblem, problem_cls), {})
+  registry.register_problem(with_agent_cls)
+
+  # Create and register the simulated Problem
+  simulated_cls = type(
+      "GymSimulatedDiscreteProblemWithAgentOn%s" % camel_game_name,
+      (GymSimulatedDiscreteProblem, problem_cls), {
+          "initial_frames_problem": with_agent_cls.name,
+          "num_testing_steps": 100
+      })
+  registry.register_problem(simulated_cls)
diff --git a/tensor2tensor/data_generators/gym_utils.py b/tensor2tensor/data_generators/gym_utils.py
index 9d1b421d9..529caa1d6 100644
--- a/tensor2tensor/data_generators/gym_utils.py
+++ b/tensor2tensor/data_generators/gym_utils.py
@@ -300,3 +300,26 @@ def wrapped_freeway_factory(warm_up_examples=0,
                       easy_freeway=False
                   ),
                   max_episode_steps=500)
+
+
+class DefaultGymWrapper(gym.Wrapper):
+  """Warmup wrapper."""
+
+  def __init__(self, env, reward_clipping=True):
+    super(DefaultGymWrapper, self).__init__(env)
+    self.reward_clipping = reward_clipping
+
+  def step(self, action):
+    ob, rew, done, info = self.env.step(action)
+
+    if self.reward_clipping:
+      rew = np.sign(rew)
+
+    return ob, rew, done, info
+
+
+def wrapped_factory(env, reward_clipping):
+  """Wrapped games."""
+  env = gym.make(env)
+  env = DefaultGymWrapper(env, reward_clipping)
+  return env
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 19f6dca77..a207e6161 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -34,6 +34,7 @@
 
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import gym_problems_specs
 from tensor2tensor.layers import discretization
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.utils import registry
@@ -341,6 +342,11 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     world_model_problem = problem_name
     simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
                               % hparams.game)
+    if problem_name not in registry.list_problems():
+      tf.logging.info("Game Problem %s not found; dynamically registering",
+                      problem_name)
+      gym_problems_specs.dynamically_create_gym_clipped_reward_problem(
+          hparams.game)
 
   # Autoencoder model dir
   autoencoder_model_dir = directories.get("autoencoder")
@@ -494,14 +500,14 @@ def rl_modelrl_base():
       simulated_env_generator_num_steps=2000,
       simulation_random_starts=True,
       intrinsic_reward_scale=0.,
-      ppo_epochs_num=200,  # This should be enough to see something
+      ppo_epochs_num=400,  # This should be enough to see something
       # Our simulated envs do not know how to reset.
       # You should set ppo_time_limit to the value you believe that
       # the simulated env produces a reasonable output.
       ppo_time_limit=200,  # TODO(blazej): this param is unused
       # It makes sense to have ppo_time_limit=ppo_epoch_length,
       # though it is not necessary.
-      ppo_epoch_length=60,
+      ppo_epoch_length=30,
       ppo_num_agents=16,
       ppo_learning_rate=2e-4,  # Will be changed, just so it exists.
       # Whether the PPO agent should be restored from the previous iteration, or

From cfb4f6297a514245584cf4f88d9345a166d45950 Mon Sep 17 00:00:00 2001
From: Alexander Ku <alexku@google.com>
Date: Mon, 30 Jul 2018 17:14:57 -0700
Subject: [PATCH 0409/2720] Auxiliary loss transformer.

PiperOrigin-RevId: 206674021
---
 .../models/research/transformer_aux.py        | 174 ++++++++++++++++++
 .../models/research/transformer_aux_test.py   | 115 ++++++++++++
 2 files changed, 289 insertions(+)
 create mode 100644 tensor2tensor/models/research/transformer_aux.py
 create mode 100644 tensor2tensor/models/research/transformer_aux_test.py

diff --git a/tensor2tensor/models/research/transformer_aux.py b/tensor2tensor/models/research/transformer_aux.py
new file mode 100644
index 000000000..1e096a32f
--- /dev/null
+++ b/tensor2tensor/models/research/transformer_aux.py
@@ -0,0 +1,174 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Transformer with auxiliary losses from https://arxiv.org/abs/1803.00144."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_layers
+from tensor2tensor.models import transformer
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+def shift_and_pad(tensor, shift, axis=0):
+  """Shifts and pads with zero along an axis.
+
+  Example:
+    shift_and_pad([1, 2, 3, 4], 2)  --> [0, 0, 1, 2]
+    shift_and_pad([1, 2, 3, 4], -2) --> [3, 4, 0, 0]
+
+  Args:
+    tensor: Tensor; to be shifted and padded.
+    shift: int; number of positions to shift by.
+    axis: int; along which axis to shift and pad.
+
+  Returns:
+    A Tensor with the same shape as the input tensor.
+  """
+  shape = tensor.shape
+  rank = len(shape)
+  assert 0 <= abs(axis) < rank
+
+  length = int(shape[axis])
+  assert 0 <= abs(shift) < length
+
+  paddings = [(0, 0)] * rank
+  begin = [0] * rank
+  size = [-1] * rank
+
+  if shift > 0:
+    paddings[axis] = (shift, 0)
+    size[axis] = length - shift
+  elif shift < 0:
+    paddings[axis] = (0, -shift)
+    begin[axis] = -shift
+
+  ret = tf.pad(tf.slice(tensor, begin, size), paddings)
+
+  return ret
+
+
+@registry.register_model
+class TransformerAux(transformer.Transformer):
+  """Attention net. See file docstring."""
+
+  def _extract_shift_values(self):
+    """Parses the shift string.
+
+    The hparams should contain the key shift_values, which maps to a
+    comma-separated string of integers. These integers specify the number of
+    timesteps to predict/reconstruct to compute auxiliary losses.
+
+    For instance, "-4,2,6" means to reconstruct the target 4 steps before and
+    predict the targets 2 steps and 6 steps ahead.
+
+    Returns:
+      List of int != 0 shift values to compute the auxiliary losses.
+    """
+    shift_values_str = self._hparams.get("shift_values", "")
+    shift_values = [int(x) for x in shift_values_str.split(",")]
+
+    tf.logging.info(
+        "Computing auxiliary losses for the following shifts: %s",
+        shift_values)
+
+    return shift_values
+
+  def auxiliary_loss(self, body_output, features, shift):
+    """Auxiliary predict loss.
+
+    Args:
+      body_output: Tensor with shape [batch_size, decoder_length, hidden_dim].
+      features: Map of features to the model. Must contain the following:
+          "targets": Target decoder outputs.
+              [batch_size, decoder_length, 1, hidden_dim]
+      shift: int != 0, amount to shift/pad the target sequence.
+        If shift > 0, it represents the number of previous timesteps to
+        reconstruct; if shift < 0, it represents the number of future timesteps
+        to predict.
+
+    Returns:
+      A 2-tuple of the numerator and denominator of the cross-entropy loss.
+
+    Raises:
+      ValueError: if features does not contain a targets_raw tensor.
+    """
+    assert isinstance(shift, int) and shift != 0
+    name = "reconst_%d" % shift if shift > 0 else "predict_%d" % abs(shift)
+
+    if features and "targets_raw" in features:
+      targets = features["targets_raw"]
+      targets = common_layers.flatten4d3d(targets)
+    else:
+      raise ValueError(
+          "Feature map must contain a targets_raw tensor.")
+
+    with tf.variable_scope(name):
+      logits = self.top(body_output, features)
+      labels = shift_and_pad(targets, shift, axis=1)
+      return common_layers.padded_cross_entropy(
+          logits,
+          labels,
+          self._hparams.label_smoothing)
+
+  def body(self, features):
+    """Transformer main model_fn.
+
+    Args:
+      features: Map of features to the model. Should contain the following:
+          "inputs": Transformer inputs.
+              [batch_size, input_length, 1, hidden_dim].
+          "targets": Target decoder outputs.
+              [batch_size, target_length, 1, hidden_dim]
+          "target_space_id": A scalar int from data_generators.problem.SpaceID.
+
+    Returns:
+      A 2-tuple containing:
+          Logit tensor. [batch_size, decoder_length, vocab_size]
+          Map of keys to loss tensors. Should contain the following:
+              "training": Training loss (shift == 0).
+              "auxiliary": Auxiliary loss (shift != 0).
+    """
+    output = super(TransformerAux, self).body(features)
+    output, losses = self._normalize_body_output(output)
+
+    aux = 0.0
+    for shift in self._extract_shift_values():
+      loss_num, loss_den = self.auxiliary_loss(output, features, shift)
+      aux += loss_num / loss_den
+    losses["auxiliary"] = aux
+
+    return output, losses
+
+
+@registry.register_hparams
+def transformer_aux_base():
+  """Set of hyperparameters."""
+  hparams = transformer.transformer_base()
+  hparams.shared_embedding_and_softmax_weights = False
+  hparams.add_hparam("shift_values", "1,2,3,4")
+  return hparams
+
+
+@registry.register_hparams
+def transformer_aux_tiny():
+  """Set of hyperparameters."""
+  hparams = transformer.transformer_tiny()
+  hparams.shared_embedding_and_softmax_weights = False
+  hparams.add_hparam("shift_values", "1,2")
+  return hparams
diff --git a/tensor2tensor/models/research/transformer_aux_test.py b/tensor2tensor/models/research/transformer_aux_test.py
new file mode 100644
index 000000000..3b3531b38
--- /dev/null
+++ b/tensor2tensor/models/research/transformer_aux_test.py
@@ -0,0 +1,115 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.models.research.transformer_aux."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+
+import numpy as np
+from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.models.research import transformer_aux
+import tensorflow as tf
+
+
+class TransformerAuxTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(
+      dict(
+          tensor=np.array(
+              [1, 2, 3, 4]
+          ),
+          shift=0,
+          axis=0,
+          target=np.array(
+              [1, 2, 3, 4]
+          ),
+      ),
+      dict(
+          tensor=np.array(
+              [1, 2, 3, 4]
+          ),
+          shift=2,
+          axis=0,
+          target=np.array(
+              [0, 0, 1, 2]
+          ),
+      ),
+      dict(
+          tensor=np.array(
+              [1, 2, 3, 4]
+          ),
+          shift=-2,
+          axis=0,
+          target=np.array(
+              [3, 4, 0, 0]
+          ),
+      ),
+      dict(
+          tensor=np.array(
+              [[1, 2, 3, 4],
+               [5, 6, 7, 8]]
+          ),
+          shift=2,
+          axis=1,
+          target=np.array(
+              [[0, 0, 1, 2],
+               [0, 0, 5, 6]]
+          ),
+      ),
+  )
+  def test_shift_and_pad(self, tensor, shift, axis, target):
+    with self.test_session() as session:
+      output = transformer_aux.shift_and_pad(tensor, shift, axis)
+      output_val = session.run(output)
+      self.assertAllEqual(output_val, target)
+
+  def test_transformer_aux_body(self):
+    batch_size = 3
+    input_length = 5
+    target_length = 16
+    vocab_size = 9
+    hparams = transformer_aux.transformer_aux_tiny()
+    hparams.shift_values = "-5,1,2,3"
+    p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
+    hparams.problem_hparams = p_hparams
+    inputs = -1 + np.random.random_integers(
+        vocab_size, size=(batch_size, input_length, 1, 1))
+    targets = -1 + np.random.random_integers(
+        vocab_size, size=(batch_size, target_length, 1, 1))
+    features = {
+        "inputs": tf.constant(inputs, dtype=tf.int32),
+        "targets": tf.constant(targets, dtype=tf.int32),
+        "target_space_id": tf.constant(1, dtype=tf.int32),
+    }
+    tf.train.create_global_step()
+    model = transformer_aux.TransformerAux(hparams, tf.estimator.ModeKeys.TRAIN,
+                                           p_hparams)
+    logits, losses = model(features)
+
+    self.assertIn("training", losses)
+    self.assertIn("auxiliary", losses)
+
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      logits_val = session.run(logits)
+      self.assertEqual(logits_val.shape,
+                       (batch_size, target_length, 1, 1, vocab_size))
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 920349269d21d80f6993375bdfd60d8c4f244b17 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 30 Jul 2018 18:24:59 -0700
Subject: [PATCH 0410/2720] vqa attention add self attention

PiperOrigin-RevId: 206683543
---
 tensor2tensor/layers/vqa_layers.py | 346 +++++++++++++++++++++++++++++
 1 file changed, 346 insertions(+)
 create mode 100644 tensor2tensor/layers/vqa_layers.py

diff --git a/tensor2tensor/layers/vqa_layers.py b/tensor2tensor/layers/vqa_layers.py
new file mode 100644
index 000000000..74044429d
--- /dev/null
+++ b/tensor2tensor/layers/vqa_layers.py
@@ -0,0 +1,346 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Some customization of common_attention."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_layers
+
+import tensorflow as tf
+
+# pylint: disable=unused-import
+from tensorflow.contrib import slim
+from tensorflow.contrib.slim.python.slim.nets.resnet_v1 import resnet_v1_152
+from tensorflow.contrib.slim.python.slim.nets.resnet_v2 import resnet_v2_152
+
+
+def summarize_tensors(tensor_dict, tag=None):
+  """Summarize the tensors.
+
+  Args:
+    tensor_dict: a dictionary of tensors.
+    tag: name scope of the summary; defaults to tensors/.
+  """
+  if tag is None:
+    tag = "tensors/"
+
+  for t_name in list(tensor_dict):
+    t = tensor_dict[t_name]
+    tf.summary.histogram(tag + t_name, t)
+
+
+def image_embedding(images,
+                    model_fn=resnet_v1_152,
+                    trainable=True,
+                    is_training=True,
+                    weight_decay=0.0001,
+                    batch_norm_decay=0.997,
+                    batch_norm_epsilon=1e-5,
+                    batch_norm_scale=True,
+                    add_summaries=False,
+                    reuse=False):
+  """Extract image features from pretrained resnet model."""
+
+  is_resnet_training = trainable and is_training
+
+  batch_norm_params = {
+      "is_training": is_resnet_training,
+      "trainable": trainable,
+      "decay": batch_norm_decay,
+      "epsilon": batch_norm_epsilon,
+      "scale": batch_norm_scale,
+  }
+
+  if trainable:
+    weights_regularizer = tf.contrib.layers.l2_regularizer(weight_decay)
+  else:
+    weights_regularizer = None
+
+  with tf.variable_scope(model_fn.__name__, [images], reuse=reuse) as scope:
+    with slim.arg_scope(
+        [slim.conv2d],
+        weights_regularizer=weights_regularizer,
+        trainable=trainable):
+      with slim.arg_scope(
+          [slim.conv2d],
+          weights_initializer=slim.variance_scaling_initializer(),
+          activation_fn=tf.nn.relu,
+          normalizer_fn=slim.batch_norm,
+          normalizer_params=batch_norm_params):
+        with slim.arg_scope([slim.batch_norm],
+                            is_training=is_resnet_training,
+                            trainable=trainable):
+          with slim.arg_scope([slim.max_pool2d], padding="SAME"):
+            net, end_points = model_fn(
+                images, num_classes=None, global_pool=False,
+                is_training=is_resnet_training,
+                reuse=reuse, scope=scope)
+
+  if add_summaries:
+    for v in end_points.values():
+      tf.contrib.layers.summaries.summarize_activation(v)
+
+  return net
+
+
+def multihead_attention(query_antecedent,
+                        memory_antecedent,
+                        bias,
+                        total_key_depth,
+                        total_value_depth,
+                        output_depth,
+                        num_heads,
+                        dropout_rate,
+                        shared_rel=False,
+                        max_relative_position=None,
+                        image_shapes=None,
+                        attention_type="dot_product",
+                        block_length=128,
+                        block_width=128,
+                        q_filter_width=1,
+                        kv_filter_width=1,
+                        q_padding="VALID",
+                        kv_padding="VALID",
+                        cache=None,
+                        gap_size=0,
+                        num_memory_blocks=2,
+                        name="multihead_attention",
+                        save_weights_to=None,
+                        make_image_summary=True,
+                        dropout_broadcast_dims=None,
+                        max_length=None,
+                        vars_3d=False,
+                        scale_dotproduct=True,
+                        **kwargs):
+  """Multihead scaled-dot-product attention with input/output transformations.
+
+  Args:
+    query_antecedent: a Tensor with shape [batch, length_q, channels]
+    memory_antecedent: a Tensor with shape [batch, length_m, channels] or None
+    bias: bias Tensor (see attention_bias())
+    total_key_depth: an integer
+    total_value_depth: an integer
+    output_depth: an integer
+    num_heads: an integer dividing total_key_depth and total_value_depth
+    dropout_rate: a floating point number
+    shared_rel: boolean to share relative embeddings
+    max_relative_position: Maximum distance between inputs to generate
+                           unique relation embeddings for. Only relevant
+                           when using "dot_product_relative" attention.
+    image_shapes: optional tuple of integer scalars.
+                  see comments for attention_image_summary()
+    attention_type: a string, either "dot_product", "dot_product_relative",
+                    "local_mask_right", "local_unmasked", "masked_dilated_1d",
+                    "unmasked_dilated_1d", graph, or any attention function
+                    with the signature (query, key, value, **kwargs)
+    block_length: an integer - relevant for "local_mask_right"
+    block_width: an integer - relevant for "local_unmasked"
+    q_filter_width: An integer specifying how wide you want the query to be.
+    kv_filter_width: An integer specifying how wide you want the keys and values
+                     to be.
+    q_padding: One of "VALID", "SAME" or "LEFT". Default is VALID: No padding.
+               kv_padding: One of "VALID", "SAME" or "LEFT". Default is "VALID":
+               no padding.
+    cache: dict containing Tensors which are the results of previous
+           attentions, used for fast decoding. Expects the dict to contrain two
+           keys ('k' and 'v'), for the initial call the values for these keys
+           should be empty Tensors of the appropriate shape.
+               'k' [batch_size, 0, key_channels]
+               'v' [batch_size, 0, value_channels]
+    gap_size: Integer option for dilated attention to indicate spacing between
+              memory blocks.
+    num_memory_blocks: Integer option to indicate how many memory blocks to look
+                       at.
+    name: an optional string.
+    save_weights_to: an optional dictionary to capture attention weights
+      for vizualization; the weights tensor will be appended there under
+      a string key created from the variable scope (including name).
+    make_image_summary: Whether to make an attention image summary.
+    dropout_broadcast_dims:  an optional list of integers less than 4
+      specifying in which dimensions to broadcast the dropout decisions.
+      saves memory.
+    max_length: an integer - needed by relative attention
+    vars_3d: use 3-dimensional variables for input/output transformations
+    scale_dotproduct: whether to normalize the attention product.
+    **kwargs (dict): Parameters for the attention function
+
+  Caching:
+    WARNING: For decoder self-attention, i.e. when memory_antecedent == None,
+    the caching assumes that the bias contains future masking.
+
+    The caching works by saving all the previous key and value values so that
+    you are able to send just the last query location to this attention
+    function. I.e. if the cache dict is provided it assumes the query is of the
+    shape [batch_size, 1, hidden_dim] rather than the full memory.
+
+  Returns:
+    The result of the attention transformation. The output shape is
+        [batch_size, length_q, hidden_dim]
+    unless the cache dict is provided in which case only the last memory
+    position is calculated and the output shape is [batch_size, 1, hidden_dim]
+    Optionally returns an additional loss parameters (ex: load balance loss for
+    the experts) returned by the attention_type function.
+
+  Raises:
+    ValueError: if the key depth or value depth are not divisible by the
+      number of attention heads.
+  """
+  if total_key_depth % num_heads != 0:
+    raise ValueError("Key depth (%d) must be divisible by the number of "
+                     "attention heads (%d)." % (total_key_depth, num_heads))
+  if total_value_depth % num_heads != 0:
+    raise ValueError("Value depth (%d) must be divisible by the number of "
+                     "attention heads (%d)." % (total_value_depth, num_heads))
+  vars_3d_num_heads = num_heads if vars_3d else 0
+  with tf.variable_scope(name, default_name="multihead_attention",
+                         values=[query_antecedent, memory_antecedent]):
+
+    if cache is None or memory_antecedent is None:
+      q, k, v = common_attention.compute_qkv(
+          query_antecedent, memory_antecedent,
+          total_key_depth, total_value_depth, q_filter_width,
+          kv_filter_width, q_padding, kv_padding,
+          vars_3d_num_heads=vars_3d_num_heads)
+    if cache is not None:
+      if attention_type != "dot_product":
+        # TODO(petershaw): Support caching when using relative position
+        # representations, i.e. "dot_product_relative" attention.
+        raise NotImplementedError(
+            "Caching is not guaranteed to work with attention types other than"
+            " dot_product.")
+      if bias is None:
+        raise ValueError("Bias required for caching. See function docstring "
+                         "for details.")
+
+      if memory_antecedent is not None:
+        # Encoder-Decoder Attention Cache
+        q = common_attention.compute_attention_component(
+            query_antecedent, total_key_depth,
+            q_filter_width, q_padding, "q",
+            vars_3d_num_heads=vars_3d_num_heads)
+        k = cache["k_encdec"]
+        v = cache["v_encdec"]
+      else:
+        k = common_attention.split_heads(k, num_heads)
+        v = common_attention.split_heads(v, num_heads)
+        decode_loop_step = kwargs.get("decode_loop_step")
+        if decode_loop_step is None:
+          k = cache["k"] = tf.concat([cache["k"], k], axis=2)
+          v = cache["v"] = tf.concat([cache["v"], v], axis=2)
+        else:
+          # Inplace update is required for inference on TPU.
+          # Inplace_ops only supports inplace_update on the first dimension.
+          # The performance of current implementation is better than updating
+          # the tensor by adding the result of matmul(one_hot,
+          # update_in_current_step)
+          tmp_k = tf.transpose(cache["k"], perm=[2, 0, 1, 3])
+          tmp_k = common_layers.tf_inplace_ops().alias_inplace_update(
+              tmp_k, decode_loop_step, tf.squeeze(k, axis=2))
+          k = cache["k"] = tf.transpose(tmp_k, perm=[1, 2, 0, 3])
+          tmp_v = tf.transpose(cache["v"], perm=[2, 0, 1, 3])
+          tmp_v = common_layers.tf_inplace_ops().alias_inplace_update(
+              tmp_v, decode_loop_step, tf.squeeze(v, axis=2))
+          v = cache["v"] = tf.transpose(tmp_v, perm=[1, 2, 0, 3])
+
+    q = common_attention.split_heads(q, num_heads)
+    if cache is None:
+      k = common_attention.split_heads(k, num_heads)
+      v = common_attention.split_heads(v, num_heads)
+
+    key_depth_per_head = total_key_depth // num_heads
+    if not vars_3d:
+      if scale_dotproduct:
+        q *= key_depth_per_head**-0.5
+
+    additional_returned_value = None
+    if callable(attention_type):  # Generic way to extend multihead_attention
+      x = attention_type(q, k, v, **kwargs)
+      if isinstance(x, tuple):
+        x, additional_returned_value = x  # Unpack
+    elif attention_type == "dot_product":
+      x = common_attention.dot_product_attention(
+          q, k, v, bias, dropout_rate, image_shapes,
+          save_weights_to=save_weights_to,
+          make_image_summary=make_image_summary,
+          dropout_broadcast_dims=dropout_broadcast_dims)
+    elif attention_type == "dot_product_relative":
+      x = common_attention.dot_product_attention_relative(
+          q,
+          k,
+          v,
+          bias,
+          max_relative_position,
+          dropout_rate,
+          image_shapes,
+          make_image_summary=make_image_summary)
+    elif attention_type == "dot_product_relative_v2":
+      x = common_attention.dot_product_self_attention_relative_v2(
+          q,
+          k,
+          v,
+          bias,
+          max_length,
+          dropout_rate,
+          image_shapes,
+          make_image_summary=make_image_summary,
+          dropout_broadcast_dims=dropout_broadcast_dims)
+    elif attention_type == "local_within_block_mask_right":
+      x = common_attention.masked_within_block_local_attention_1d(
+          q, k, v, block_length=block_length)
+    elif attention_type == "rel_local_mask_right":
+      x = common_attention.masked_rel_local_attention_1d(
+          q, k, v, block_length=block_length,
+          make_image_summary=make_image_summary,
+          dropout_rate=dropout_rate,
+          share_rel_embed=shared_rel)
+    elif attention_type == "local_mask_right":
+      x = common_attention.masked_local_attention_1d(
+          q,
+          k,
+          v,
+          block_length=block_length,
+          make_image_summary=make_image_summary)
+    elif attention_type == "local_unmasked":
+      x = common_attention.local_attention_1d(
+          q, k, v, block_length=block_length, filter_width=block_width)
+    elif attention_type == "masked_dilated_1d":
+      x = common_attention.masked_dilated_self_attention_1d(
+          q, k, v, block_length, block_width,
+          gap_size, num_memory_blocks)
+    else:
+      assert attention_type == "unmasked_dilated_1d"
+      x = common_attention.dilated_self_attention_1d(
+          q, k, v, block_length, block_width,
+          gap_size, num_memory_blocks)
+    x = common_attention.combine_heads(x)
+
+    # Set last dim specifically.
+    x.set_shape(x.shape.as_list()[:-1] + [total_value_depth])
+
+    if vars_3d:
+      o_var = tf.get_variable(
+          "o", [num_heads, total_value_depth // num_heads, output_depth])
+      o_var = tf.cast(o_var, x.dtype)
+      o_var = tf.reshape(o_var, [total_value_depth, output_depth])
+      x = tf.tensordot(x, o_var, axes=1)
+    else:
+      x = common_layers.dense(
+          x, output_depth, use_bias=False, name="output_transform")
+    if additional_returned_value is not None:
+      return x, additional_returned_value
+    return x

From 7f791e60e6bb5e7d15e66431064fe087f0b5fec1 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Mon, 30 Jul 2018 18:29:25 -0700
Subject: [PATCH 0411/2720] Adding support for relative task ids in
 MultiProblem.

PiperOrigin-RevId: 206683973
---
 tensor2tensor/data_generators/cola.py         |  3 +--
 tensor2tensor/data_generators/imdb.py         |  3 +--
 tensor2tensor/data_generators/lm1b.py         |  3 +--
 tensor2tensor/data_generators/mrpc.py         |  3 +--
 .../data_generators/multi_problem.py          | 25 +++++++++++++++++--
 tensor2tensor/data_generators/multinli.py     |  3 +--
 tensor2tensor/data_generators/problem.py      | 10 ++++++++
 tensor2tensor/data_generators/qnli.py         |  3 +--
 tensor2tensor/data_generators/quora_qpairs.py |  3 +--
 tensor2tensor/data_generators/rte.py          |  3 +--
 tensor2tensor/data_generators/sst_binary.py   |  3 +--
 tensor2tensor/data_generators/wnli.py         |  3 +--
 12 files changed, 43 insertions(+), 22 deletions(-)

diff --git a/tensor2tensor/data_generators/cola.py b/tensor2tensor/data_generators/cola.py
index e8e033f7a..0028c4f26 100644
--- a/tensor2tensor/data_generators/cola.py
+++ b/tensor2tensor/data_generators/cola.py
@@ -116,6 +116,5 @@ class ColaCharacters(Cola):
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
 
-  @property
-  def task_id(self):
+  def global_task_id(self):
     return problem.TaskID.COLA
diff --git a/tensor2tensor/data_generators/imdb.py b/tensor2tensor/data_generators/imdb.py
index 6ea5715a8..cbb37d3df 100644
--- a/tensor2tensor/data_generators/imdb.py
+++ b/tensor2tensor/data_generators/imdb.py
@@ -105,6 +105,5 @@ class SentimentIMDBCharacters(SentimentIMDB):
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
 
-  @property
-  def task_id(self):
+  def global_task_id(self):
     return problem.TaskID.EN_CHR_SENT
diff --git a/tensor2tensor/data_generators/lm1b.py b/tensor2tensor/data_generators/lm1b.py
index cebd94c95..dca331c6c 100644
--- a/tensor2tensor/data_generators/lm1b.py
+++ b/tensor2tensor/data_generators/lm1b.py
@@ -175,8 +175,7 @@ class LanguagemodelLm1bCharacters(LanguagemodelLm1b32k):
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
 
-  @property
-  def task_id(self):
+  def global_task_id(self):
     return problem.TaskID.EN_CHR
 
 
diff --git a/tensor2tensor/data_generators/mrpc.py b/tensor2tensor/data_generators/mrpc.py
index 3c441d285..083970be3 100644
--- a/tensor2tensor/data_generators/mrpc.py
+++ b/tensor2tensor/data_generators/mrpc.py
@@ -140,6 +140,5 @@ class MSRParaphraseCorpusCharacters(MSRParaphraseCorpus):
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
 
-  @property
-  def task_id(self):
+  def global_task_id(self):
     return problem.TaskID.EN_SIM
diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index f14bfaf4d..0fb37679c 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -19,6 +19,7 @@
 from __future__ import print_function
 
 from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.layers import discretization
 from tensor2tensor.utils import metrics
@@ -92,12 +93,15 @@ def dataset(self,
     datasets = []
     is_training = mode == tf.estimator.ModeKeys.TRAIN
 
-    for task in self.task_list:
+    for idx, task in enumerate(self.task_list):
       task_dataset = task.dataset(mode, data_dir, num_threads,
                                   output_buffer_size, shuffle_files,
                                   hparams, preprocess, dataset_split,
                                   shard, partition_id, num_partitions,
                                   max_records)
+      if idx == 0:
+        self.update_task_ids(data_dir)
+
       if is_training:
         task_dataset = task_dataset.repeat()
       # pylint: disable=cell-var-from-loop
@@ -111,7 +115,7 @@ def dataset(self,
           self.flatten_zip)
     else:
       single_mtl_dataset = datasets[0]
-      for data in datasets[0:]:
+      for data in datasets[1:]:
         single_mtl_dataset = single_mtl_dataset.concatenate(data)
 
     return single_mtl_dataset
@@ -120,3 +124,20 @@ def eval_metrics(self):
     return [
         metrics.Metrics.ACC, metrics.Metrics.NEG_LOG_PERPLEXITY
     ]
+
+  def update_task_ids(self, data_dir):
+    primary_task = self.task_list[0]
+    if primary_task.has_inputs:
+      raise ValueError("Only support language models as primary problem which "
+                       "supplies the vocabulary and the hparams.")
+
+    encoder = primary_task.feature_encoders(data_dir=data_dir)["targets"]
+
+    id_offset = encoder.vocab_size + text_encoder.NUM_RESERVED_TOKENS
+    if hasattr(primary_task, "additional_reserved_tokens"):
+      id_offset += len(primary_task.additional_reserved_tokens)
+
+    for idx, _ in enumerate(self.task_list):
+      # protect against the ord mapping of chars with the 2x multiplier.
+      self.task_list[idx].set_task_id(idx + 2 * id_offset)
+      print(self.task_list[idx].task_id)
diff --git a/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index 9ab08f9fa..7ecdad85e 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -133,6 +133,5 @@ class MultiNLICharacters(MultiNLI):
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
 
-  @property
-  def task_id(self):
+  def global_task_id(self):
     return problem.TaskID.THREE_CL_NLI
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 73d49ccf5..bd1362335 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -349,6 +349,15 @@ def eval_metrics(self):
         metrics.Metrics.ACC_PER_SEQ, metrics.Metrics.NEG_LOG_PERPLEXITY
     ]
 
+  @property
+  def task_id(self):
+    if self._task_id == -1 and hasattr(self, "global_task_id"):
+      self._task_id = self.global_task_id()
+    return self._task_id
+
+  def set_task_id(self, new_task_id):
+    self._task_id = new_task_id
+
   # ============================================================================
   # END SUBCLASS INTERFACE
   # ============================================================================
@@ -449,6 +458,7 @@ def __init__(self, was_reversed=False, was_copy=False):
     self._encoders = None
     self._hparams = None
     self._feature_info = None
+    self._task_id = -1
 
   def get_feature_encoders(self, data_dir=None):
     if self._encoders is None:
diff --git a/tensor2tensor/data_generators/qnli.py b/tensor2tensor/data_generators/qnli.py
index 101f1dba2..93773adcb 100644
--- a/tensor2tensor/data_generators/qnli.py
+++ b/tensor2tensor/data_generators/qnli.py
@@ -130,6 +130,5 @@ class QuestionNLICharacters(QuestionNLI):
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
 
-  @property
-  def task_id(self):
+  def global_task_id(self):
     return problem.TaskID.EN_NLI
diff --git a/tensor2tensor/data_generators/quora_qpairs.py b/tensor2tensor/data_generators/quora_qpairs.py
index 489fa8650..7cfbb5fd6 100644
--- a/tensor2tensor/data_generators/quora_qpairs.py
+++ b/tensor2tensor/data_generators/quora_qpairs.py
@@ -136,6 +136,5 @@ class QuoraQuestionPairsCharacters(QuoraQuestionPairs):
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
 
-  @property
-  def task_id(self):
+  def global_task_id(self):
     return problem.TaskID.EN_SIM
diff --git a/tensor2tensor/data_generators/rte.py b/tensor2tensor/data_generators/rte.py
index bc36a45ab..d62624316 100644
--- a/tensor2tensor/data_generators/rte.py
+++ b/tensor2tensor/data_generators/rte.py
@@ -130,6 +130,5 @@ class RTECharacters(RTE):
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
 
-  @property
-  def task_id(self):
+  def global_task_id(self):
     return problem.TaskID.EN_NLI
diff --git a/tensor2tensor/data_generators/sst_binary.py b/tensor2tensor/data_generators/sst_binary.py
index ce5fcc586..17b1a3ef5 100644
--- a/tensor2tensor/data_generators/sst_binary.py
+++ b/tensor2tensor/data_generators/sst_binary.py
@@ -117,6 +117,5 @@ class SentimentSSTBinaryCharacters(SentimentSSTBinary):
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
 
-  @property
-  def task_id(self):
+  def global_task_id(self):
     return problem.TaskID.EN_CHR_SENT
diff --git a/tensor2tensor/data_generators/wnli.py b/tensor2tensor/data_generators/wnli.py
index 0a6a4caf2..e44459fc6 100644
--- a/tensor2tensor/data_generators/wnli.py
+++ b/tensor2tensor/data_generators/wnli.py
@@ -128,6 +128,5 @@ class WinogradNLICharacters(WinogradNLI):
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
 
-  @property
-  def task_id(self):
+  def global_task_id(self):
     return problem.TaskID.EN_NLI

From 22d42f489adce5389f5e10197e63c2d6595f399d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 30 Jul 2018 18:51:23 -0700
Subject: [PATCH 0412/2720] Add documentation and update sample generation.

PiperOrigin-RevId: 206686342
---
 .../data_generators/function_docstring.py     | 89 +++++++++++++------
 .../models/research/similarity_transformer.py | 33 +++----
 2 files changed, 77 insertions(+), 45 deletions(-)

diff --git a/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
index 3058cccce..c44319342 100644
--- a/tensor2tensor/data_generators/function_docstring.py
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -14,27 +14,42 @@
 # limitations under the License.
 """Github function/text similatrity problems."""
 import csv
+import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import translate
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
+import tensorflow as tf
 
-
-# There are 10 splits of the data as CSV files.
-_DATA_BASE_URL = "https://storage.googleapis.com/kubeflow-examples/t2t-code-search/data"
-_GITHUB_FUNCTION_DOCSTRING_FILES = [
-    [
-        "{}/pairs-0000{}-of-00010.csv".format(_DATA_BASE_URL, i),
-        "pairs-0000{}-of-00010.csv".format(i),
-    ]
-    for i in range(10)
-]
+# pylint: disable=g-import-not-at-top
+if six.PY2:
+  from StringIO import StringIO
+else:
+  from io import StringIO
+# pylint: enable=g-import-not-at-top
 
 
 @registry.register_problem
 class GithubFunctionDocstring(translate.TranslateProblem):
-  """This class defines the problem of finding similarity between Python
-  function and docstring"""
+  """Function and Docstring similarity Problem.
+
+  This problem contains the data consisting of function
+  and docstring pairs as CSV files. The files are structured
+  such that they contain two columns without headers containing
+  the docstring tokens and function tokens. The delimiter is
+  ",".
+  """
+
+  @property
+  def base_url(self):
+    return "gs://kubeflow-examples/t2t-code-search/raw_data"
+
+  @property
+  def pair_files_list(self):
+    return [
+        "func-doc-pairs-000{:02}-of-00100.csv".format(i)
+        for i in range(100)
+    ]
 
   @property
   def is_generate_per_split(self):
@@ -44,29 +59,45 @@ def is_generate_per_split(self):
   def approx_vocab_size(self):
     return 2**13
 
-  def source_data_files(self, dataset_split):
-    # TODO(sanyamkapoor): separate train/eval data set.
-    return _GITHUB_FUNCTION_DOCSTRING_FILES
+  def source_data_files(self, _):
+    # TODO(sanyamkapoor): Manually separate train/eval data set.
+    return self.pair_files_list
+
+  @property
+  def max_samples_for_vocab(self):
+    # FIXME(sanyamkapoor): This exists to handle memory explosion.
+    return int(3.5e5)
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    """Returns a generator to return {"inputs": [text], "targets": [text]}."""
+    """A generator to return data samples.Returns the data generator to return.
+
+
+    Args:
+      data_dir: A string representing the data directory.
+      tmp_dir: A string representing the temporary directory and is
+              used to download files if not already available.
+      dataset_split: Train, Test or Eval.
+
+    Yields:
+      Each element yielded is of a Python dict of the form
+        {"inputs": "STRING", "targets": "STRING"}
+    """
 
-    pair_csv_files = [
-        generator_utils.maybe_download(data_dir, filename, uri)
-        for uri, filename in self.source_data_files(dataset_split)
+    csv_file_names = self.source_data_files(dataset_split)
+    csv_files = [
+        generator_utils.maybe_download(tmp_dir, filename,
+                                       "{}/{}".format(self.base_url,
+                                                      filename))
+        for filename in csv_file_names
     ]
 
-    for pairs_file in pair_csv_files:
+    for pairs_file in csv_files:
+      tf.logging.debug("Reading {}".format(pairs_file))
       with open(pairs_file, "r") as csv_file:
-        pairs_reader = csv.reader(csv_file)
-        for row in pairs_reader:
-          function_tokens, docstring_tokens = row[-2:]
-          yield {"inputs": docstring_tokens, "targets": function_tokens}
-
-  def generate_text_for_vocab(self, data_dir, tmp_dir):
-    for sample in self.generate_samples(data_dir, tmp_dir, None):
-      yield sample["inputs"]
-      yield sample["targets"]
+        for line in csv_file:
+          reader = csv.reader(StringIO(line))
+          for docstring_tokens, function_tokens in reader:
+            yield {"inputs": docstring_tokens, "targets": function_tokens}
 
   def eval_metrics(self):
     return [
diff --git a/tensor2tensor/models/research/similarity_transformer.py b/tensor2tensor/models/research/similarity_transformer.py
index f60e261f9..08575d0d6 100644
--- a/tensor2tensor/models/research/similarity_transformer.py
+++ b/tensor2tensor/models/research/similarity_transformer.py
@@ -23,23 +23,25 @@
 
 @registry.register_model
 class SimilarityTransformer(t2t_model.T2TModel):
-  """
-  This class defines the model to compute similarity scores between functions
-  and docstrings
+  """Transformer Model for Similarity between two strings.
+
+  This model defines the architecture using two transformer
+  networks, each of which embed a string and the loss is
+  calculated as a Binary Cross-Entropy loss. Normalized
+  Dot Product is used as the distance measure between two
+  string embeddings.
   """
 
-  def top(self, body_output, features):
+  def top(self, body_output, _):
     return body_output
 
   def body(self, features):
-    """Body of the Similarity Transformer Network."""
-
-    with tf.variable_scope("string_embedding"):
-      string_embedding = self.encode(features, "inputs")
+    with tf.variable_scope('string_embedding'):
+      string_embedding = self.encode(features, 'inputs')
 
-    if "targets" in features:
-      with tf.variable_scope("code_embedding"):
-        code_embedding = self.encode(features, "targets")
+    if 'targets' in features:
+      with tf.variable_scope('code_embedding'):
+        code_embedding = self.encode(features, 'targets')
 
       string_embedding_norm = tf.nn.l2_normalize(string_embedding, axis=1)
       code_embedding_norm = tf.nn.l2_normalize(code_embedding, axis=1)
@@ -59,7 +61,7 @@ def body(self, features):
       loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels,
                                                      logits=logits)
 
-      return string_embedding, {"training": loss}
+      return string_embedding, {'training': loss}
 
     return string_embedding
 
@@ -69,18 +71,17 @@ def encode(self, features, input_key):
 
     (encoder_input, encoder_self_attention_bias, _) = (
         transformer.transformer_prepare_encoder(inputs, problem.SpaceID.EN_TOK,
-                                                self._hparams))
+                                                hparams))
 
     encoder_input = tf.nn.dropout(encoder_input,
                                   1.0 - hparams.layer_prepostprocess_dropout)
     encoder_output = transformer.transformer_encoder(
         encoder_input,
         encoder_self_attention_bias,
-        self._hparams,
+        hparams,
         nonpadding=transformer.features_to_nonpadding(features, input_key))
-    encoder_output = tf.expand_dims(encoder_output, 2)
 
-    encoder_output = tf.reduce_mean(tf.squeeze(encoder_output, axis=2), axis=1)
+    encoder_output = tf.reduce_mean(encoder_output, axis=1)
 
     return encoder_output
 

From 38d5cf75e8dc561521305eb7e97d5a10a5444911 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 30 Jul 2018 19:05:28 -0700
Subject: [PATCH 0413/2720] Improvements to model based RL: FullPong, action
 embeding concatenations, longer PPO in the final epoch.

PiperOrigin-RevId: 206688010
---
 .../data_generators/gym_problems_specs.py     | 36 +++++++++----------
 tensor2tensor/data_generators/gym_utils.py    |  5 ++-
 tensor2tensor/models/research/next_frame.py   |  5 ++-
 .../models/research/next_frame_params.py      |  1 +
 tensor2tensor/rl/model_rl_experiment.py       |  9 +++--
 5 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index 2632cb730..6c961c2a5 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -99,12 +99,12 @@ def num_rewards(self):
 
 
 @registry.register_problem
-class GymWrappedLongPongRandom(GymDiscreteProblem):
+class GymWrappedFullPongRandom(GymDiscreteProblem):
   """Pong game, random actions."""
 
   @property
   def env_name(self):
-    return "T2TPongWarmUp20RewSkip2000Steps-v1"
+    return "T2TPongWarmUp20RewSkipFull-v1"
 
   @property
   def min_reward(self):
@@ -181,7 +181,7 @@ class GymDiscreteProblemWithAgentOnFreeway(GymRealDiscreteProblem,
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
     GymSimulatedDiscreteProblem, GymWrappedPongRandom):
-  """Similated pong."""
+  """Simulated pong."""
 
   @property
   def initial_frames_problem(self):
@@ -193,31 +193,31 @@ def num_testing_steps(self):
 
 
 @registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedLongPong(GymRealDiscreteProblem,
-                                                   GymWrappedLongPongRandom):
+class GymDiscreteProblemWithAgentOnWrappedFullPong(GymRealDiscreteProblem,
+                                                   GymWrappedFullPongRandom):
   pass
 
 
 @registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedLongPongWithAutoencoder(
-    GymDiscreteProblemWithAutoencoder, GymWrappedLongPongRandom):
+class GymDiscreteProblemWithAgentOnWrappedFullPongWithAutoencoder(
+    GymDiscreteProblemWithAutoencoder, GymWrappedFullPongRandom):
   pass
 
 
 @registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedLongPongAutoencoded(
-    GymDiscreteProblemAutoencoded, GymWrappedLongPongRandom):
+class GymDiscreteProblemWithAgentOnWrappedFullPongAutoencoded(
+    GymDiscreteProblemAutoencoded, GymWrappedFullPongRandom):
   pass
 
 
 @registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPong(
-    GymSimulatedDiscreteProblem, GymWrappedLongPongRandom):
+class GymSimulatedDiscreteProblemWithAgentOnWrappedFullPong(
+    GymSimulatedDiscreteProblem, GymWrappedFullPongRandom):
   """Simulated pong."""
 
   @property
   def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_long_pong"
+    return "gym_discrete_problem_with_agent_on_wrapped_full_pong"
 
   @property
   def num_testing_steps(self):
@@ -225,13 +225,13 @@ def num_testing_steps(self):
 
 
 @registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedLongPongAutoencoded(
-    GymSimulatedDiscreteProblemAutoencoded, GymWrappedLongPongRandom):
-  """GymSimulatedDiscreteProblemWithAgentOnWrappedLongPongAutoencoded."""
+class GymSimulatedDiscreteProblemWithAgentOnWrappedFullPongAutoencoded(
+    GymSimulatedDiscreteProblemAutoencoded, GymWrappedFullPongRandom):
+  """GymSimulatedDiscreteProblemWithAgentOnWrappedFullPongAutoencoded."""
 
   @property
   def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_long_pong_autoencoded"
+    return "gym_discrete_problem_with_agent_on_wrapped_full_pong_autoencoded"
 
   @property
   def num_testing_steps(self):
@@ -253,7 +253,7 @@ class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
     GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
-  """Similated breakout."""
+  """Simulated breakout."""
 
   @property
   def initial_frames_problem(self):
@@ -293,7 +293,7 @@ class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnFreeway(GymSimulatedDiscreteProblem,
                                                     GymFreewayRandom):
-  """Similated freeway."""
+  """Simulated freeway."""
 
   @property
   def initial_frames_problem(self):
diff --git a/tensor2tensor/data_generators/gym_utils.py b/tensor2tensor/data_generators/gym_utils.py
index 529caa1d6..af14451ae 100644
--- a/tensor2tensor/data_generators/gym_utils.py
+++ b/tensor2tensor/data_generators/gym_utils.py
@@ -126,10 +126,9 @@ def wrapped_pong_factory(warm_up_examples=0, action_space_reduction=False,
                   max_episode_steps=200)
 
 
-gym.envs.register(id="T2TPongWarmUp20RewSkip2000Steps-v1",
+gym.envs.register(id="T2TPongWarmUp20RewSkipFull-v1",
                   entry_point=lambda: wrapped_pong_factory(  # pylint: disable=g-long-lambda
-                      warm_up_examples=20, reward_skip_steps=15),
-                  max_episode_steps=2000)
+                      warm_up_examples=20, reward_skip_steps=15))
 
 
 class BreakoutWrapper(WarmupWrapper):
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 2e77cfb9b..5148c71eb 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -77,7 +77,10 @@ def body(self, features):
       action_mask = tf.layers.dense(action, filters, name="action_mask")
       zeros_mask = tf.zeros(common_layers.shape_list(x)[:-1] + [filters],
                             dtype=tf.float32)
-      x *= action_mask + zeros_mask
+      if hparams.concatenate_actions:
+        x = tf.concat([x, action_mask + zeros_mask], axis=-1)
+      else:
+        x *= action_mask + zeros_mask
 
     # Run a stack of convolutions.
     for i in range(hparams.num_hidden_layers):
diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index 5217d7824..13dcedbef 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -43,6 +43,7 @@ def next_frame():
   hparams.add_hparam("num_compress_steps", 6)
   hparams.add_hparam("filter_double_steps", 2)
   hparams.add_hparam("video_modality_loss_cutoff", 0.02)
+  hparams.add_hparam("concatenate_actions", True)
   return hparams
 
 
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index a207e6161..79e01fc57 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -135,7 +135,8 @@ def train_autoencoder(problem_name, data_dir, output_dir, hparams, epoch):
 
 
 def train_agent(problem_name, agent_model_dir,
-                event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0):
+                event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
+                is_final_epoch=False):
   """Train the PPO agent in the simulated environment."""
   gym_problem = registry.problem(problem_name)
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
@@ -149,6 +150,8 @@ def train_agent(problem_name, agent_model_dir,
       ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
 
   ppo_epochs_num = hparams.ppo_epochs_num
+  if is_final_epoch:
+    ppo_epochs_num *= 2
   ppo_hparams.save_models_every_epochs = ppo_epochs_num
   ppo_hparams.world_model_dir = world_model_dir
   ppo_hparams.add_hparam("force_beginning_resets", True)
@@ -424,7 +427,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       ppo_model_dir = ppo_event_dir
     train_agent(simulated_problem_name, ppo_model_dir,
                 ppo_event_dir, directories["world_model"], epoch_data_dir,
-                hparams, epoch=epoch)
+                hparams, epoch=epoch, is_final_epoch=is_final_epoch)
 
     # Collect data from the real environment.
     log("Generating real environment data")
@@ -513,7 +516,7 @@ def rl_modelrl_base():
       # Whether the PPO agent should be restored from the previous iteration, or
       # should start fresh each time.
       ppo_continue_training=True,
-      game="wrapped_long_pong",
+      game="wrapped_full_pong",
       # Whether to evaluate the world model in each iteration of the loop to get
       # the model_reward_accuracy metric.
       eval_world_model=True,

From e980a940c69a71fb15ac28182ca30859069e346b Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Mon, 30 Jul 2018 19:45:25 -0700
Subject: [PATCH 0414/2720] wmt_enfr dataset should not use vocabulary based on
 "small" dataset - it does

PiperOrigin-RevId: 206691360
---
 tensor2tensor/data_generators/translate.py      | 1 +
 tensor2tensor/data_generators/translate_enfr.py | 6 ++++--
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index cf3272604..b8bc020cc 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -87,6 +87,7 @@ def compile_data(tmp_dir, datasets, filename):
   if tf.gfile.Exists(lang1_fname) and tf.gfile.Exists(lang2_fname):
     tf.logging.info("Skipping compile data, found files:\n%s\n%s", lang1_fname,
                     lang2_fname)
+    return filename
   with tf.gfile.GFile(lang1_fname, mode="w") as lang1_resfile:
     with tf.gfile.GFile(lang2_fname, mode="w") as lang2_resfile:
       for dataset in datasets:
diff --git a/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index 93d2e3f42..e190f9982 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -86,7 +86,8 @@ def approx_vocab_size(self):
 
   @property
   def vocab_filename(self):
-    return "vocab.enfr.%d" % self.approx_vocab_size
+    return "vocab.enfr.%s.%d" % (
+        "small" if self.use_small_dataset else "large", self.approx_vocab_size)
 
   @property
   def use_small_dataset(self):
@@ -101,7 +102,8 @@ def source_data_files(self, dataset_split):
     return datasets
 
   def vocab_data_files(self):
-    return _ENFR_TRAIN_SMALL_DATA
+    return (_ENFR_TRAIN_SMALL_DATA if self.use_small_dataset
+            else _ENFR_TRAIN_LARGE_DATA)
 
 
 @registry.register_problem

From d9593cfc4680f5a12f4e76ae71b0b93f955707c8 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 30 Jul 2018 22:21:10 -0700
Subject: [PATCH 0415/2720] exp learning rate

PiperOrigin-RevId: 206704854
---
 tensor2tensor/utils/learning_rate.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index ec15b696b..65d7ba872 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -33,6 +33,18 @@ def learning_rate_factor(name, step_num, hparams):
     return tf.minimum(1.0, tf.maximum(0.0, ret))
   elif name == "rsqrt_decay":
     return tf.rsqrt(tf.maximum(step_num, hparams.learning_rate_warmup_steps))
+  elif name == "rsqrt_normalized_decay":
+    scale = tf.sqrt(hparams.learning_rate_warmup_steps)
+    return scale * tf.rsqrt(tf.maximum(
+        step_num, hparams.learning_rate_warmup_steps))
+  elif name == "exp_decay":
+    decay_steps = hparams.learning_rate_decay_steps
+    warmup_steps = hparams.learning_rate_warmup_steps
+    p = (step_num - warmup_steps) / decay_steps
+    p = tf.maximum(p, 0.)
+    if hparams.learning_rate_decay_staircase:
+      p = tf.floor(p)
+    return tf.pow(hparams.learning_rate_decay_rate, p)
   elif name == "rsqrt_hidden_size":
     return hparams.hidden_size ** -0.5
   elif name == "legacy":

From 0c1b7ea2ac9858417ff3332aba3ae6fe0dc6394a Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Tue, 31 Jul 2018 08:29:57 -0700
Subject: [PATCH 0416/2720] Add targeted dropout method.

PiperOrigin-RevId: 206765502
---
 tensor2tensor/layers/common_layers.py | 74 ++++++++++++++++++++-------
 1 file changed, 55 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index a30b5ff0d..29840c6a0 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -699,8 +699,7 @@ def l2_norm(x, filters=None, epsilon=1e-6, name=None, reuse=None):
   """Layer normalization with l2 norm."""
   if filters is None:
     filters = shape_list(x)[-1]
-  with tf.variable_scope(
-      name, default_name="l2_norm", values=[x], reuse=reuse):
+  with tf.variable_scope(name, default_name="l2_norm", values=[x], reuse=reuse):
     scale = tf.get_variable(
         "l2_norm_scale", [filters], initializer=tf.ones_initializer())
     bias = tf.get_variable(
@@ -747,11 +746,9 @@ def zero_add(previous_value, x, name=None, reuse=None):
   Returns:
     previous_value + gamma * x.
   """
-  with tf.variable_scope(
-      name, default_name="zero_add", reuse=reuse):
-    gamma = tf.get_variable(
-        "gamma", (), initializer=tf.zeros_initializer())
-    return previous_value + gamma*x
+  with tf.variable_scope(name, default_name="zero_add", reuse=reuse):
+    gamma = tf.get_variable("gamma", (), initializer=tf.zeros_initializer())
+    return previous_value + gamma * x
 
 
 def layer_prepostprocess(previous_value,
@@ -2019,10 +2016,7 @@ def _weights_one_third(labels):
   return tf.ones(tf.shape(labels)[:-1]) / 3.
 
 
-def dml_loss(pred,
-             labels,
-             weights_fn=_weights_one_third,
-             reduce_sum=True):
+def dml_loss(pred, labels, weights_fn=_weights_one_third, reduce_sum=True):
   """Discretized mixture of logistics loss.
 
   Args:
@@ -3635,24 +3629,66 @@ def cyclegan_upsample(net, num_outputs, stride, method="conv2d_transpose"):
       net = tf.image.resize_nearest_neighbor(
           net, [stride[0] * height, stride[1] * width])
       net = tf.pad(net, spatial_pad_1, "REFLECT")
-      net = tf.contrib.layers.conv2d(net, num_outputs, kernel_size=[3, 3],
-                                     padding="valid")
+      net = tf.contrib.layers.conv2d(
+          net, num_outputs, kernel_size=[3, 3], padding="valid")
     elif method == "bilinear_upsample_conv":
-      net = tf.image.resize_bilinear(
-          net, [stride[0] * height, stride[1] * width])
+      net = tf.image.resize_bilinear(net,
+                                     [stride[0] * height, stride[1] * width])
       net = tf.pad(net, spatial_pad_1, "REFLECT")
-      net = tf.contrib.layers.conv2d(net, num_outputs, kernel_size=[3, 3],
-                                     padding="valid")
+      net = tf.contrib.layers.conv2d(
+          net, num_outputs, kernel_size=[3, 3], padding="valid")
     elif method == "conv2d_transpose":
       # This corrects 1 pixel offset for images with even width and height.
       # conv2d is left aligned and conv2d_transpose is right aligned for even
       # sized images (while doing "SAME" padding).
       # Note: This doesn"t reflect actual model in paper.
       net = tf.contrib.layers.conv2d_transpose(
-          net, num_outputs, kernel_size=[3, 3], stride=stride,
-          padding="valid")
+          net, num_outputs, kernel_size=[3, 3], stride=stride, padding="valid")
       net = net[:, 1:, 1:, :]
     else:
       raise ValueError("Unknown method: [%s]" % method)
 
     return net
+
+
+def targeted_dropout(inputs,
+                     k,
+                     keep_prob,
+                     targeting_fn,
+                     is_training,
+                     do_prune=False):
+  """Applies targeted dropout.
+
+  Applies dropout at a rate of `1 - keep_prob` to only those elements of `x`
+  marked by `targeting_fn`. See below and paper for more detail:
+
+  "Targeted Dropout for Posthoc Pruning" Aidan N. Gomez, Ivan Zhang,
+    Kevin Swersky, Yarin Gal, and Geoffrey E. Hinton.
+
+  Args:
+    inputs: Tensor, inputs to apply targeted dropout to.
+    k: Scalar Tensor or python scalar, sets the number of elements to target in
+      `inputs`. Must be within `[0, tf.shape(x)[-1]]` and compatible with
+      second argument of `targeting_fn`.
+    keep_prob: Scalar Tensor, passed as `tf.nn.dropout`'s `keep_prob` argument.
+    targeting_fn: callable `fn(inputs, k) -> Boolean Tensor`, produces a
+      boolean mask the same shape as `inputs` where True indicates an element
+      will be dropped, and False not.
+    is_training: bool, indicates whether currently training.
+    do_prune: bool, indicates whether to prune the `k * (1 - keep_prob)`
+      elements of `inputs` expected to be dropped each forwards pass.
+  Returns:
+    Tensor, same shape and dtype as `inputs`.
+  """
+  if not is_training and do_prune:
+    k = tf.round(k * (1 - keep_prob))
+
+  mask = targeting_fn(inputs, k)
+  mask = tf.cast(mask, inputs.dtype)
+
+  if is_training:
+    return inputs * (1 - mask) + tf.nn.dropout(inputs, keep_prob) * mask
+  elif do_prune:
+    return inputs * (1 - mask)
+  else:
+    return inputs

From 6a595f1b3dc1d47c758a2c1674a5bda6b953ad66 Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Tue, 31 Jul 2018 09:41:15 -0700
Subject: [PATCH 0417/2720] Change word dropout default to 0.

PiperOrigin-RevId: 206776309
---
 tensor2tensor/models/research/transformer_vae.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index d2e253279..768ed4f16 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -687,7 +687,7 @@ def transformer_ae_small():
   hparams.add_hparam("noise_dev", 0.5)
   hparams.add_hparam("d_mix", 0.5)
   hparams.add_hparam("logit_normalization", True)
-  hparams.add_hparam("word_dropout", 0.1)
+  hparams.add_hparam("word_dropout", 0.0)
   # Bottleneck kinds supported: dense, vae, semhash, gumbel-softmax, dvq.
   hparams.add_hparam("bottleneck_kind", "semhash")
   hparams.add_hparam("num_blocks", 1)

From e6bd25aa8e15d2504c4374b03d7879c9ca5ac9ca Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 31 Jul 2018 10:58:12 -0700
Subject: [PATCH 0418/2720] Fixing the latent_num_frames bug.

PiperOrigin-RevId: 206790323
---
 tensor2tensor/models/research/next_frame.py        | 7 ++++++-
 tensor2tensor/models/research/next_frame_params.py | 4 +---
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 5148c71eb..2f2e3c680 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -241,7 +241,12 @@ def construct_latent_tower(self, images):
     conv_size = self.tinyify([32, 64, 64])
     with tf.variable_scope("latent", reuse=tf.AUTO_REUSE):
       # this allows more predicted frames at inference time
-      latent_images = images[:self.hparams.latent_num_frames]
+      latent_num_frames = self.hparams.latent_num_frames
+      if latent_num_frames == 0:  # use all frames by default.
+        latent_num_frames = (self.hparams.video_num_input_frames +
+                             self.hparams.video_num_target_frames)
+      tf.logging.info("Creating latent tower with %d frames."%latent_num_frames)
+      latent_images = images[:latent_num_frames]
       images = tf.concat(latent_images, 3)
 
       x = images
diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index 13dcedbef..663f683ed 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -73,9 +73,7 @@ def next_frame_stochastic():
   hparams.add_hparam("relu_shift", 1e-12)
   hparams.add_hparam("dna_kernel_size", 5)
   hparams.add_hparam("scheduled_sampling_k", 900.0)
-  hparams.add_hparam(
-      "latent_num_frames",  # use all frames by default.
-      hparams.video_num_input_frames + hparams.video_num_target_frames)
+  hparams.add_hparam("latent_num_frames", 0)  # 0 means use all frames.
   hparams.add_hparam("tiny_mode", False)
   hparams.add_hparam("anneal_end", 100000)
   hparams.add_hparam("upsample_method", "conv2d_transpose")

From 87597694dfb4fbabe5b0594fcdc27657366d911c Mon Sep 17 00:00:00 2001
From: Keyon Vafa <vafa@google.com>
Date: Tue, 31 Jul 2018 11:08:37 -0700
Subject: [PATCH 0419/2720] Changed prior loss and reconstruction loss scaling
 to 1/batch_size

PiperOrigin-RevId: 206792617
---
 tensor2tensor/layers/discretization.py        | 56 ++++++++++++++-----
 .../models/research/transformer_vae.py        | 26 +++++++++
 2 files changed, 69 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 5426e7974..42918af22 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -71,6 +71,7 @@ def nearest_neighbor(x,
                      random_top_k=1,
                      soft_em=False,
                      num_samples=1,
+                     sum_over_latents=False,
                      summary=True):
   """Find the nearest element in means to elements in x.
 
@@ -79,10 +80,12 @@ def nearest_neighbor(x,
       block_dim].
     means: Embedding table of shape [num_blocks, block_v_size, block_dim].
     block_v_size: Number of table entries per block.
-    random_top_k: Noisy top-k if this is bigger than 1 (Default: 1).
-    soft_em: If True then use soft EM rather than hard EM (Default: False).
-    num_samples: Number of samples to take in soft EM (Default: 1).
-    summary: If True then record summary histogram of entropies (Default: True).
+    random_top_k: Noisy top-k if this is bigger than 1.
+    soft_em: If True then use soft EM rather than hard EM.
+    num_samples: Number of samples to take in soft EM.
+    sum_over_latents: Whether to sum over non-batch dimensions when calculating
+      negative entropy loss. Used only when doing soft EM.
+    summary: If True then record summary histogram of entropies.
 
   Returns:
     Tensor with nearest element in mean encoded in one-hot notation
@@ -109,7 +112,9 @@ def nearest_neighbor(x,
     nearest_hot = tf.one_hot(nearest_idx, depth=block_v_size)
     neg_q_entropy = tf.reduce_sum(
         nearest_hot * tf.expand_dims(tf.nn.log_softmax(-dist), 2), axis=2)
-    neg_q_entropy = tf.reduce_mean(neg_q_entropy)
+    if sum_over_latents:
+      neg_q_entropy = tf.reduce_sum(neg_q_entropy, [1, 2])
+    neg_q_entropy = tf.reduce_mean(neg_q_entropy, axis=0)
     nearest_hot = tf.reduce_mean(nearest_hot, axis=-2)
     if summary:
       tf.summary.histogram("neg_q_entropy", tf.reshape(neg_q_entropy, [-1]))
@@ -139,7 +144,8 @@ def embedding_lookup(x,
                      do_hard_gumbel_softmax=False,
                      temperature_warmup_steps=150000,
                      do_iaf=False,
-                     approximate_gs_entropy=False):
+                     approximate_gs_entropy=False,
+                     sum_over_latents=False):
   """Compute nearest neighbors and loss for training the embeddings via DVQ.
 
   Args:
@@ -161,6 +167,9 @@ def embedding_lookup(x,
     approximate_gs_entropy: Whether to approximate the Gumbel-Softmax density
       as a categorical distribution when calculating the sample entropy. Used
       only if bottleneck_kind is gumbel-softmax-dvq.
+    sum_over_latents: Whether to sum over non-batch dimensions when calculating
+      negative entropy loss. Used only if soft EM or when bottleneck_kind is
+      gumbel-softmax-dvq.
 
   Returns:
     x_means_hot: The nearest neighbor in one hot form, with shape
@@ -181,7 +190,8 @@ def embedding_lookup(x,
         num_samples=num_samples,
         temperature_warmup_steps=temperature_warmup_steps,
         do_iaf=do_iaf,
-        approximate_gs_entropy=approximate_gs_entropy)
+        approximate_gs_entropy=approximate_gs_entropy,
+        sum_over_latents=sum_over_latents)
   else:
     x_means_hot, neg_q_entropy = nearest_neighbor(
         x,
@@ -189,11 +199,15 @@ def embedding_lookup(x,
         block_v_size,
         random_top_k,
         soft_em=soft_em,
-        num_samples=num_samples)
+        num_samples=num_samples,
+        sum_over_latents=sum_over_latents)
   x_means_hot_flat = tf.reshape(x_means_hot, [-1, num_blocks, block_v_size])
   x_means = tf.matmul(tf.transpose(x_means_hot_flat, perm=[1, 0, 2]), means)
   x_means = tf.transpose(x_means, [1, 0, 2])
   x = tf.reshape(x, [-1] + common_layers.shape_list(x)[2:])
+
+  # Currently, we use the mean scaling for the commitment loss, as opposed to
+  # summing across all non-batch dimensions.
   q_loss = tf.reduce_mean(tf.square((tf.stop_gradient(x) - x_means)))
   e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
   return x_means_hot, x_means, q_loss, e_loss, neg_q_entropy
@@ -480,6 +494,7 @@ def discrete_bottleneck(inputs,
                         do_hard_gumbel_softmax=False,
                         do_iaf=False,
                         approximate_gs_entropy=False,
+                        sum_over_latents=False,
                         discrete_mix=0.5,
                         noise_dev=1.,
                         startup_steps=50000,
@@ -530,6 +545,9 @@ def discrete_bottleneck(inputs,
     approximate_gs_entropy: Whether to approximate the Gumbel-Softmax density
       as a categorical distribution when calculating the sample entropy. Used
       only if bottleneck_kind is gumbel-softmax-dvq.
+    sum_over_latents: Whether to sum over all non-batch dimensions before
+      taking mean of entropy loss term. Used only if bottleneck kind is DVQ
+      or gumbel-softmax-dvq.
     discrete_mix: Factor for mixing discrete and non-discrete input. Used only
       if bottleneck_kind is semhash.
     noise_dev: Noise stddev. Used only if bottleneck_kind is semhash.
@@ -634,7 +652,8 @@ def discrete_bottleneck(inputs,
                              temperature_warmup_steps=temperature_warmup_steps,
                              do_hard_gumbel_softmax=do_hard_gumbel_softmax,
                              do_iaf=do_iaf,
-                             approximate_gs_entropy=approximate_gs_entropy))
+                             approximate_gs_entropy=approximate_gs_entropy,
+                             sum_over_latents=sum_over_latents))
         # Update the EMA variables.
         if ema:
           tf.logging.info("Using EMA with beta = {}".format(beta))
@@ -870,7 +889,8 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
                                         temperature_warmup_steps=150000,
                                         summary=True,
                                         do_iaf=False,
-                                        approximate_gs_entropy=False):
+                                        approximate_gs_entropy=False,
+                                        sum_over_latents=False):
   """Sample from Gumbel-Softmax and compute neighbors and losses.
 
   Args:
@@ -891,6 +911,8 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
       Gumbel-Softmax sample (Default: False).
     approximate_gs_entropy: When `True`, we approximate Gumbel-Softmax
       density as categorical when calculating sample entropy (Default: False).
+    sum_over_latents: Whether to sum over non-batch dimensions when calculating
+      negative entropy loss.
 
   Returns:
     x_means_assignments: A `float`-like `Tensor` containing the codebook
@@ -945,9 +967,12 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
 
   # Take mean over samples to approximate entropy.
   neg_q_entropy = tf.reduce_mean(q_dist.log_prob(q_samples), 0)
-
   if summary:
     tf.summary.histogram("neg_q_entropy", tf.reshape(neg_q_entropy, [-1]))
+  if sum_over_latents:
+    neg_q_entropy = tf.reshape(neg_q_entropy,
+                               [batch_size, num_blocks, latent_dim])
+    neg_q_entropy = tf.reduce_sum(neg_q_entropy, [1, 2])
   neg_q_entropy = tf.reduce_mean(neg_q_entropy)
 
   if do_iaf:
@@ -1014,8 +1039,13 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
              shift=log_pi[:, :, :-1],
              scale_identity_multiplier=1./temperature)])
     q_samples = chained_bijectors.forward(q_samples[:, :, :-1])
-    neg_q_entropy += tf.reduce_mean(
-        chained_bijectors.inverse_log_det_jacobian(q_samples, event_ndims=1))
+    log_det = chained_bijectors.inverse_log_det_jacobian(
+        q_samples, event_ndims=1)
+    log_det = tf.reshape(log_det,
+                         [num_samples, batch_size, num_blocks, latent_dim])
+    if sum_over_latents:
+      log_det = tf.reduce_sum(log_det, axis=[2, 3])
+    neg_q_entropy += tf.reduce_mean(log_det)
 
     q_samples = tf.reshape(
         q_samples,
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 768ed4f16..362744f19 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -392,6 +392,15 @@ def ae_transformer_internal(inputs,
             task="translate")
         _, latent_pred_loss = ae_latent_softmax(
             latents_pred, tf.stop_gradient(latents_discrete), hparams)
+
+        # Scale by latent dimension for summary so we can compare across
+        # batches.
+        if _DO_SUMMARIES:
+          tf.summary.scalar("latent_pred_loss_mean",
+                            tf.reduce_mean(latent_pred_loss))
+        if hparams.sum_over_latents:
+          latent_pred_loss = tf.reduce_sum(latent_pred_loss, [1, 2])
+
         losses["latent_pred"] = tf.reduce_mean(
             latent_pred_loss * tf.to_float(cond)) * hparams.prior_scale
         losses["neg_q_entropy"] = neg_q_entropy * hparams.entropy_scale
@@ -594,6 +603,22 @@ def __init__(self, *args, **kwargs):
   def has_input(self):
     return self._problem_hparams.input_modality
 
+  def loss(self, logits, features):
+    """Computes cross-entropy loss and scales by 1/batch_size."""
+    labels = features["targets"]
+    logits_shape = common_layers.shape_list(logits)
+    vocab_size = logits_shape[-1]
+    with tf.name_scope("padded_cross_entropy", values=[logits, labels]):
+      logits, labels = common_layers.pad_with_zeros(logits, labels)
+      logits = tf.reshape(
+          logits,
+          common_layers.shape_list(labels) + [vocab_size],
+          name="padded_cross_entropy_size_check")
+      logits = tf.cast(logits, tf.float32)
+      xent = common_layers.smoothing_cross_entropy(
+          logits, labels, vocab_size, confidence=1.0, gaussian=False)
+      return tf.reduce_sum(xent) / tf.cast(logits_shape[0], tf.float32)
+
   def body(self, features):
     inputs = features["inputs"] if "inputs" in features else None
     if self._hparams.drop_inputs:
@@ -732,6 +757,7 @@ def transformer_ae_small():
   hparams.add_hparam("do_iaf", False)
   hparams.add_hparam("approximate_gs_entropy", False)
   hparams.add_hparam("temperature_warmup_steps", 150000)
+  hparams.add_hparam("sum_over_latents", False)
   hparams.force_full_predict = True
 
   # task params

From cfc2b29113611ee15c1be6f4c6604ef5df7b42e9 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 31 Jul 2018 11:49:03 -0700
Subject: [PATCH 0420/2720] restore hook log num of variables

PiperOrigin-RevId: 206800389
---
 tensor2tensor/utils/restore_hook.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/restore_hook.py b/tensor2tensor/utils/restore_hook.py
index 523c19a9d..4d6fdd3c4 100644
--- a/tensor2tensor/utils/restore_hook.py
+++ b/tensor2tensor/utils/restore_hook.py
@@ -55,6 +55,6 @@ def begin(self):
                       if name.startswith(self._old_model_scope)}
     self._assignment_map = assignment_map
 
-    tf.logging.info("restoring variables from checkpoint %s"%(
-        self._checkpoint_path))
+    tf.logging.info("restoring %d variables from checkpoint %s"%(
+        len(assignment_map), self._checkpoint_path))
     tf.train.init_from_checkpoint(self._checkpoint_path, self._assignment_map)

From edc7c6e35f1ecfc0a9111e9e61679f4bd68fbf7b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 31 Jul 2018 14:02:01 -0700
Subject: [PATCH 0421/2720] Implementing GGNN with unit test.

PiperOrigin-RevId: 206822733
---
 .../common_message_passing_attention.py       | 205 +++++++++++++-----
 1 file changed, 155 insertions(+), 50 deletions(-)

diff --git a/tensor2tensor/layers/common_message_passing_attention.py b/tensor2tensor/layers/common_message_passing_attention.py
index a13210082..c716b1d39 100644
--- a/tensor2tensor/layers/common_message_passing_attention.py
+++ b/tensor2tensor/layers/common_message_passing_attention.py
@@ -263,6 +263,71 @@ def graph_attention(q,
     return tf.matmul(weights, v)
 
 
+def _compute_edge_transforms(node_states,
+                             depth,
+                             num_transforms,
+                             ignore_zero=True,
+                             name="transform"):
+  """Helper function that computes transformation for keys and values.
+
+  Let B be the number of batches.
+  Let N be the number of nodes in the graph.
+  Let D be the size of the node hidden states.
+  Let K be the size of the attention keys/queries (total_key_depth).
+  Let V be the size of the attention values (total_value_depth).
+  Let T be the total number of transforms (num_transforms).
+
+  Computes the transforms for keys or values for attention.
+  * For each node N_j and edge type t, a key K_jt of size K is computed. When an
+    edge of type t goes from node N_j to any other node, K_jt is the key that is
+    in the attention process.
+  * For each node N_j and edge type t, a value V_jt of size V is computed. When
+    an edge of type t goes from node N_j to node N_i, Attention(Q_i, K_jt)
+    produces a weight w_ijt. The message sent along this edge is w_ijt * V_jt.
+
+  Args:
+    node_states: A tensor of shape [B, L, D]
+    depth: An integer (K or V)
+    num_transforms: An integer (T),
+    ignore_zero: A boolean to ignore 0 edge
+    name: A name for the function
+
+  Returns:
+    x: A The attention keys or values for each node and edge type
+      (shape [B, N*T, K or V])
+  """
+  node_shapes = common_layers.shape_list(node_states)
+  nonignored_transforms = num_transforms - int(ignore_zero)
+  x = common_layers.dense(
+      node_states,
+      depth * nonignored_transforms,
+      use_bias=False,
+      name=name)
+
+  batch = node_shapes[0]  # B.
+  length = node_shapes[1]  # N.
+
+  # Making the fourth dimension explicit by separating the vectors of size
+  # K*T (in k) and V*T (in v) into two-dimensional matrices with shape [K, T]
+  # (in k) and [V, T] in v.
+  #
+  # This reshape is only necessary when ignore_zero is True (for the padding
+  # step that follows).
+  x = tf.reshape(x, [batch, length, nonignored_transforms, depth])
+
+  # If we previously ignored edge type 0, then we need to pad the keys and
+  # values to take this additional edge type into account. To do so, we
+  # pad the third dimension of k and v (which has size T-1 if ignore_zero is
+  # True) to size T with zeroes.
+  if ignore_zero:
+    x = tf.pad(x, [[0, 0], [0, 0], [1, 0], [0, 0]])
+
+  # Flatten out the fourth dimension.
+  x = tf.reshape(x, [batch, length * num_transforms, depth])
+
+  return x
+
+
 def compute_mpnn_qkv(node_states,
                      total_key_depth,
                      total_value_depth,
@@ -315,53 +380,20 @@ def compute_mpnn_qkv(node_states,
   q = common_layers.dense(
       node_states, total_key_depth, use_bias=False, name="q_mpnn")
 
-  q_shape = common_layers.shape_list(q)  # As above, q_shape = [B, N, K].
-
-  # T (or T-1 if ignore_zero).
-  nonignored_transforms = num_transforms - int(ignore_zero)
-
   # Creates the attention keys in a manner similar to the process of creating
   # the attention queries. One key is created for each type of outgoing edge the
   # corresponding node might have, meaning k will have shape [B, N, K*T].
-  k = common_layers.dense(
-      node_states,
-      total_key_depth * nonignored_transforms,
-      use_bias=False,
-      name="k_mpnn")
-
-  # The values over which self-attention is performed. They are created in
-  # a manner largely identical to that of the keys.
-  v = common_layers.dense(
-      node_states,
-      total_value_depth * nonignored_transforms,
-      use_bias=False,
-      name="v_mpnn")
-
-  batch = q_shape[0]  # B.
-  length = q_shape[1]  # N.
+  k = _compute_edge_transforms(node_states,
+                               total_key_depth,
+                               num_transforms,
+                               ignore_zero=ignore_zero,
+                               name="k_mpnn")
+  v = _compute_edge_transforms(node_states,
+                               total_value_depth,
+                               num_transforms,
+                               ignore_zero=ignore_zero,
+                               name="v_mpnn")
 
-  # Making the fourth dimension explicit by separating the vectors of size
-  # K*T (in k) and V*T (in v) into two-dimensional matrices with shape [K, T]
-  # (in k) and [V, T] in v.
-  #
-  # This reshape is only necessary when ignore_zero is True (for the padding
-  # step that follows).
-  k = tf.reshape(k, [batch, length, nonignored_transforms, total_key_depth])
-  v = tf.reshape(
-      v, [q_shape[0], q_shape[1], nonignored_transforms, total_value_depth])
-
-  # If we previously ignored edge type 0, then we need to pad the keys and
-  # values to take this additional edge type into account. To do so, we
-  # pad the third dimension of k and v (which has size T-1 if ignore_zero is
-  # True) to size T with zeroes.
-  if ignore_zero:
-    k = tf.pad(k, [[0, 0], [0, 0], [1, 0], [0, 0]])
-    v = tf.pad(v, [[0, 0], [0, 0], [1, 0], [0, 0]])
-
-  # Flatten out the fourth dimension.
-  k = tf.reshape(k, [q_shape[0], q_shape[1] * num_transforms, total_key_depth])
-  v = tf.reshape(v,
-                 [q_shape[0], q_shape[1] * num_transforms, total_value_depth])
   return q, k, v
 
 
@@ -641,13 +673,86 @@ def dot_product_mpnn_attention(q,
     # actual edges.
     edge_compatibility *= edge_vectors  # Shape [B, T, N, N].
 
-    # Computes the incoming value vectors for each node by weighting them
-    # according to the attention weights. These values are still segregated by
-    # edge type.
-    all_edge_values = tf.matmul(edge_compatibility, v)  # Shape = [B, T, N, V].
+    output = compute_values(edge_compatibility, v)
+    return output
+
+
+def ggnn_fast_dense(node_states,
+                    adjacency_matrix,
+                    num_edge_types,
+                    total_value_depth,
+                    ignore_zero=True,
+                    name=None):
+  """ggnn version of the MPNN from Gilmer et al.
 
-    # Combines the weighted value vectors together across edge types into a
-    # single N x V matrix for each batch.
-    output = tf.reduce_sum(all_edge_values, axis=1)  # Shape [B, N, V].
+  Let B be the number of batches.
+  Let D be the size of the node hidden states.
+  Let K be the size of the attention keys/queries.
+  Let V be the size of the output of the ggnn
+
+  Args:
+    node_states: The value Tensor of shape [B, T, N, D].
+    adjacency_matrix: A Tensor of shape [B, N, N]. An entry at indices b, i, j
+     is the integer edge type of the edge from node j to node i in batch b.
+    num_edge_types: An integer specifying number of edge types.
+    total_value_depth: An integer (V)
+    ignore_zero: A boolean to ignore edge type 0.
+    name: A string.
 
+  Returns:
+    A Tensor of shape [B, N, V] storing the result of computing attention
+    weights using the queries and keys and combining the values according to
+    those weights.
+
+  Raises:
+    ValueError: if num_transforms doesn't equal num_edge_types and not using
+      weighted sum.
+  """
+  # between the same nodes (with only one edge of each type. adjacency_matrix
+  # will need to be converted to shape [B, T, N, N].
+  with tf.variable_scope(
+      name,
+      default_name="ggnn_fast_dense",
+      values=[node_states, adjacency_matrix, num_edge_types]):
+    nodes_shape = common_layers.shape_list(node_states)
+    v = _compute_edge_transforms(node_states,
+                                 total_value_depth,
+                                 num_edge_types,
+                                 ignore_zero=ignore_zero,
+                                 name="v_mpnn")
+    v = tf.reshape(v, [nodes_shape[0], nodes_shape[1], num_edge_types,
+                       total_value_depth
+                      ])  # Shape [B, N, T, V].
+    v = tf.transpose(v, [0, 2, 1, 3])  # Shape [B, T, N, V].
+
+    # Generate one-hot vectors based on edge types.
+    # If there is an edge from node j to node i of type t, then index t of the
+    # last dimension is 1 for entry (i, j) of the second and third dimensions.
+    edge_vectors = tf.one_hot(adjacency_matrix, num_edge_types)
+
+    # Rearranging the dimensions to match the shape of all_edge_logits.
+    edge_vectors = tf.transpose(edge_vectors, [0, 3, 1, 2])
+    output = compute_values(edge_vectors, v)
     return output
+
+
+def compute_values(edge_compatibility, v):
+  """Compute values. If edge compatibilities is just adjacency, we get ggnn.
+
+  Args:
+    edge_compatibility: A tensor of shape [batch, num_transforms, length, depth]
+    v: A tensor of shape [batch, num_transforms, length, depth]
+
+  Returns:
+    output: A [batch, length, depth] tensor
+  """
+
+  # Computes the incoming value vectors for each node by weighting them
+  # according to the attention weights. These values are still segregated by
+  # edge type.
+  all_edge_values = tf.matmul(edge_compatibility, v)  # Shape = [B, T, N, V].
+
+  # Combines the weighted value vectors together across edge types into a
+  # single N x V matrix for each batch.
+  output = tf.reduce_sum(all_edge_values, axis=1)  # Shape [B, N, V].
+  return output

From 816cf149aaf3bc0bc96ab48d6e28f2ed377821d1 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 31 Jul 2018 15:55:36 -0700
Subject: [PATCH 0422/2720] Internal change

PiperOrigin-RevId: 206843035
---
 tensor2tensor/models/research/next_frame_test.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index 45dfed9c3..77aedf738 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -21,6 +21,7 @@
 
 from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
 from tensor2tensor.models.research import next_frame
+from tensor2tensor.models.research import next_frame_params
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -74,35 +75,35 @@ def TestOnVariousUpSampleLayers(self, hparams, model, expected_last_dim):
 
   def testBasic(self):
     self.TestOnVariousInputOutputSizes(
-        next_frame.next_frame(),
+        next_frame_params.next_frame(),
         next_frame.NextFrameBasic,
         256)
 
   def testStochastic(self):
     self.TestOnVariousInputOutputSizes(
-        next_frame.next_frame_stochastic(),
+        next_frame_params.next_frame_stochastic(),
         next_frame.NextFrameStochastic,
         1)
 
   def testStochasticTwoFrames(self):
     self.TestOnVariousInputOutputSizes(
-        next_frame.next_frame_stochastic(),
+        next_frame_params.next_frame_stochastic(),
         next_frame.NextFrameStochasticTwoFrames,
         1)
 
   def testStochasticEmily(self):
     self.TestOnVariousInputOutputSizes(
-        next_frame.next_frame_stochastic_emily(),
+        next_frame_params.next_frame_stochastic_emily(),
         next_frame.NextFrameStochasticEmily,
         1)
 
   def testStochasticSavp(self):
     self.TestOnVariousInputOutputSizes(
-        next_frame.next_frame_savp(),
+        next_frame_params.next_frame_savp(),
         next_frame.NextFrameSavp,
         1)
     self.TestOnVariousUpSampleLayers(
-        next_frame.next_frame_savp(),
+        next_frame_params.next_frame_savp(),
         next_frame.NextFrameSavp,
         1)
 

From 096b186e026f1ab3dc06f0ed7a024bbc0037c9a7 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 31 Jul 2018 16:00:22 -0700
Subject: [PATCH 0423/2720] Move to GA grpc API in serving_utils

PiperOrigin-RevId: 206843788
---
 tensor2tensor/serving/serving_utils.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index c7d44bf47..c11689d93 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -20,14 +20,14 @@
 
 import base64
 from googleapiclient import discovery
-from grpc.beta import implementations
+import grpc
 
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.utils import cloud_tpu as cloud
 import tensorflow as tf
 from tensorflow_serving.apis import predict_pb2
-from tensorflow_serving.apis import prediction_service_pb2
+from tensorflow_serving.apis import prediction_service_pb2_grpc
 
 
 
@@ -40,10 +40,8 @@ def _make_example(input_ids, feature_name="inputs"):
 
 
 def _create_stub(server):
-  host, port = server.split(":")
-  channel = implementations.insecure_channel(host, int(port))
-  # TODO(bgb): Migrate to GA API.
-  return prediction_service_pb2.beta_create_PredictionService_stub(channel)
+  channel = grpc.insecure_channel(server)
+  return prediction_service_pb2_grpc.PredictionServiceStub(channel)
 
 
 def _encode(inputs, encoder, add_eos=True):

From b2d749b3a66e875da43f807a078d94bf2dc5c176 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 31 Jul 2018 16:10:59 -0700
Subject: [PATCH 0424/2720] Fix 2/2 of failing travis tests.

PiperOrigin-RevId: 206845448
---
 tensor2tensor/models/research/transformer_vae.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 362744f19..4877544c7 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -443,8 +443,10 @@ def bn_inputs():
     d = latents_dense
     latent_len = common_layers.shape_list(latents_dense)[1]
     if isinstance(latent_len, tf.Tensor):
-      latent_len = hparams.max_length
+      # TODO(trandustin): Fix this in a better manner.
+      latent_len = max(1000, hparams.max_length)
     pos = tf.get_variable("pos", [1, latent_len + 1, 1, hparams.hidden_size])
+    pos = pos[:, :common_layers.shape_list(latents_dense)[1] + 1, :, :]
     latents_dense = tf.pad(latents_dense,
                            [[0, 0], [1, 0], [0, 0], [0, 0]]) + pos
 

From 246690f4cbfcbc38ee53d60b1841f32dcb6081ac Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 31 Jul 2018 16:31:32 -0700
Subject: [PATCH 0425/2720] Disable multiprocessing shuffle - more headache
 than it's worth

PiperOrigin-RevId: 206848733
---
 tensor2tensor/data_generators/generator_utils.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 04716a94c..93b9f489a 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -19,7 +19,6 @@
 from __future__ import print_function
 
 import gzip
-import multiprocessing as mp
 import os
 import random
 import stat
@@ -489,11 +488,8 @@ def shuffle_dataset(filenames):
     tf.logging.info("Skipping shuffle because output files exist")
     return
   tf.logging.info("Shuffling data...")
-  if len(filenames) > 1:
-    pool = mp.Pool(min(len(filenames), 20))
-    pool.map(_shuffle_single, filenames)
-  else:
-    _shuffle_single(filenames[0])
+  for filename in filenames:
+    _shuffle_single(filename)
   tf.logging.info("Data shuffled.")
 
 

From 9a07e1000601be4ffdd52ff00f328c11b4f6f01e Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 31 Jul 2018 16:35:05 -0700
Subject: [PATCH 0426/2720] Disable pylint check not-context-manager because
 this seems to be finicky with tf.contrib.slim.arg_scope

PiperOrigin-RevId: 206849355
---
 pylintrc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pylintrc b/pylintrc
index 157d214e3..47e9c7ab3 100644
--- a/pylintrc
+++ b/pylintrc
@@ -37,7 +37,7 @@ msg-template={msg_id}:{line:3} {obj}: {msg} [{symbol}]
 enable=indexing-exception,old-raise-syntax
 
 # List of checkers and warnings to disable.
-disable=design,similarities,no-self-use,attribute-defined-outside-init,locally-disabled,star-args,pointless-except,bad-option-value,global-statement,fixme,suppressed-message,useless-suppression,locally-enabled,file-ignored,multiple-imports,c-extension-no-member,trailing-newlines,unsubscriptable-object,misplaced-comparison-constant,no-member,abstract-method,no-else-return,missing-docstring,wrong-import-order,protected-access,inconsistent-return-statements,invalid-unary-operand-type,import-error,no-name-in-module,arguments-differ
+disable=design,similarities,no-self-use,attribute-defined-outside-init,locally-disabled,star-args,pointless-except,bad-option-value,global-statement,fixme,suppressed-message,useless-suppression,locally-enabled,file-ignored,multiple-imports,c-extension-no-member,trailing-newlines,unsubscriptable-object,misplaced-comparison-constant,no-member,abstract-method,no-else-return,missing-docstring,wrong-import-order,protected-access,inconsistent-return-statements,invalid-unary-operand-type,import-error,no-name-in-module,arguments-differ,not-context-manager
 
 [BASIC]
 

From 1ec9635eb2489f598b8658c1559e8db6938447ed Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 31 Jul 2018 16:47:58 -0700
Subject: [PATCH 0427/2720] add a body_scope to Transformer's decoding methods
 for models that don't fully

PiperOrigin-RevId: 206851213
---
 tensor2tensor/models/transformer.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 7966125d1..5d3733e88 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -668,7 +668,8 @@ def fast_decode_tpu(encoder_output,
                     sos_id=0,
                     eos_id=beam_search.EOS_ID,
                     batch_size=None,
-                    force_decode_length=False):
+                    force_decode_length=False,
+                    scope_prefix="body/"):
   """Given encoder output and a symbols to logits function, does fast decoding.
 
   Implements only greedy decoding for TPU.
@@ -687,6 +688,7 @@ def fast_decode_tpu(encoder_output,
     batch_size: An integer, must be passed if there is no input.
     force_decode_length: A bool, whether to force the full decode length, or if
         False, stop when all beams hit eos_id.
+    scope_prefix: str, prefix for decoder layer variable scopes.
 
   Returns:
       A dict of decoding results {
@@ -725,7 +727,8 @@ def fast_decode_tpu(encoder_output,
     for layer in range(num_layers):
       layer_name = "layer_%d" % layer
       with tf.variable_scope(
-          "body/decoder/%s/encdec_attention/multihead_attention" % layer_name):
+          "%sdecoder/%s/encdec_attention/multihead_attention" % (scope_prefix,
+                                                                 layer_name)):
         k_encdec = common_attention.compute_attention_component(
             encoder_output, key_channels, name="k")
         k_encdec = common_attention.split_heads(k_encdec, hparams.num_heads)
@@ -807,7 +810,8 @@ def fast_decode(encoder_output,
                 sos_id=0,
                 eos_id=beam_search.EOS_ID,
                 batch_size=None,
-                force_decode_length=False):
+                force_decode_length=False,
+                scope_prefix="body/"):
   """Given encoder output and a symbols to logits function, does fast decoding.
 
   Implements both greedy and beam search decoding, uses beam search iff
@@ -831,6 +835,7 @@ def fast_decode(encoder_output,
     batch_size: an integer scalar - must be passed if there is no input
     force_decode_length: bool, whether to force the full decode length, or if
       False, stop when all beams hit eos_id.
+    scope_prefix: str, prefix for decoder layer variable scopes.
 
   Returns:
       A dict of decoding results {
@@ -870,7 +875,8 @@ def fast_decode(encoder_output,
     for layer in range(num_layers):
       layer_name = "layer_%d" % layer
       with tf.variable_scope(
-          "body/decoder/%s/encdec_attention/multihead_attention" % layer_name):
+          "%sdecoder/%s/encdec_attention/multihead_attention" % (scope_prefix,
+                                                                 layer_name)):
         k_encdec = common_attention.compute_attention_component(
             encoder_output, key_channels, name="k",
             vars_3d_num_heads=vars_3d_num_heads)

From 97cb147985eba2e171809d926ca5696d694aa27f Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Tue, 31 Jul 2018 17:22:16 -0700
Subject: [PATCH 0428/2720] Relative memory efficient unmasked self-attention.
 Added a shape test as well.

PiperOrigin-RevId: 206855956
---
 tensor2tensor/layers/common_attention.py      | 177 +++++++++++++++++-
 tensor2tensor/layers/common_attention_test.py |  14 ++
 2 files changed, 181 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index d1a6d6ee0..e3341d9ac 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1693,7 +1693,6 @@ def dot_product_self_attention_relative_v2(q,
   """Calculate relative position-aware dot-product self-attention.
 
   Only works for masked self-attention (no looking forward).
-  TODO(noam): extend to unmasked self-attention
 
   The attention calculation is augmented with learned representations for the
   relative position between each element in q and each element in k and v.
@@ -1761,6 +1760,147 @@ def dot_product_self_attention_relative_v2(q,
     return ret
 
 
+def _absolute_position_to_relative_position_unmasked(x):
+  """Helper function for dot_product_unmasked_self_attention_relative_v2.
+
+  Rearrange an attention logits or weights Tensor.
+
+  The dimensions of the input represent:
+  [batch, heads, query_position, memory_position]
+
+  The dimensions of the output represent:
+  [batch, heads, query_position, memory_position - query_position + length - 1]
+
+  Only works with unmasked_attention.
+
+  Args:
+    x: a Tensor with shape [batch, heads, length, length]
+
+  Returns:
+    a Tensor with shape [batch, heads, length, 2*length-1]
+  """
+  batch, heads, length, _ = common_layers.shape_list(x)
+  # padd along column
+  x = tf.pad(x, [[0, 0], [0, 0], [0, 0], [0, length-1]])
+  x_flat = tf.reshape(x, [batch, heads, length**2 + length*(length -1)])
+  # add 0's in the beginning that will skew the elements after reshape
+  x_flat = tf.pad(x_flat, [[0, 0], [0, 0], [length, 0]])
+  x = tf.reshape(x_flat, [batch, heads, length, 2*length])
+  x = tf.slice(x, [0, 0, 0, 1], [batch, heads, length,
+                                 2*length -1])
+  return x
+
+
+def dot_product_unmasked_self_attention_relative_v2(
+    q, k, v, bias, max_length=None, dropout_rate=0.0, image_shapes=None,
+    name=None, make_image_summary=True, dropout_broadcast_dims=None):
+  """Calculate relative position-aware dot-product self-attention.
+
+  The attention calculation is augmented with learned representations for the
+  relative position between each element in q and each element in k and v.
+
+  Args:
+    q: a Tensor with shape [batch, heads, length, depth].
+    k: a Tensor with shape [batch, heads, length, depth].
+    v: a Tensor with shape [batch, heads, length, depth].
+    bias: bias Tensor.
+    max_length: an integer - changing this invalidates checkpoints
+    dropout_rate: a floating point number.
+    image_shapes: optional tuple of integer scalars.
+    name: an optional string.
+    make_image_summary: Whether to make an attention image summary.
+    dropout_broadcast_dims:  an optional list of integers less than 4
+      specifying in which dimensions to broadcast the dropout decisions.
+      saves memory.
+
+  Returns:
+    A Tensor.
+  """
+  with tf.variable_scope(
+      name,
+      default_name="dot_product_unmasked_self_attention_relative_v2",
+      values=[q, k, v]):
+
+    # This calculation only works for self attention.
+    # q, k and v must therefore have the same shape.
+    q.get_shape().assert_is_compatible_with(k.get_shape())
+    q.get_shape().assert_is_compatible_with(v.get_shape())
+
+    # Use separate embeddings suitable for keys and values.
+    length = common_layers.shape_list(q)[2]
+    assert max_length is not None
+    k_shape = common_layers.shape_list(k)
+    depth_k = k_shape[-1]
+    initializer_stddev = depth_k**-0.5
+    # TODO(avaswani): Add option for unshared relative embeddings
+    key_relative_embeddings = (
+        tf.get_variable(name="key_relative_embeddings",
+                        shape=(2*max_length-1, depth_k),
+                        initializer=tf.random_normal_initializer(
+                            stddev=initializer_stddev)))
+    # [batch, num_heads, query_length, memory_length]
+    logits = tf.matmul(q, k, transpose_b=True)
+    # slice out the right band of rel embeddings to save computation
+    # First pad the relative embeddings with zeros if the sequence length
+    # is longer than max length
+    padded_key_relative_embeddings = tf.pad(key_relative_embeddings,
+                                            [[tf.maximum(
+                                                length-max_length, 0),
+                                              tf.maximum(
+                                                  length-max_length, 0)],
+                                             [0, 0]])
+    used_key_relative_embeddings = tf.slice(padded_key_relative_embeddings,
+                                            [tf.maximum(
+                                                max_length-length,
+                                                length-max_length),
+                                             0],
+                                            [2*length -1, -1])
+    unmasked_rel_logits = tf.einsum("bhld, md -> bhlm", q,
+                                    used_key_relative_embeddings)
+    unmasked_rel_logits = _relative_position_to_absolute_position_unmasked(
+        unmasked_rel_logits)
+    logits += unmasked_rel_logits
+
+    if bias is not None:
+      logits += bias
+    weights = tf.nn.softmax(logits, name="attention_weights")
+    # dropping out the attention links for each of the heads
+    weights = common_layers.dropout_with_broadcast_dims(
+        weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
+    # relative_weights.set_shape([None, None, None, max_length])
+    if common_layers.should_generate_summaries() and make_image_summary:
+      attention_image_summary(weights, image_shapes)
+    ret = tf.matmul(weights, v)
+    # getting the contribution of the relative embeddings to the values
+    # [batch, num_heads, query_length, 2*memory_length-1]
+    relative_weights = _absolute_position_to_relative_position_unmasked(
+        weights)
+    depth_v = common_layers.shape_list(v)[3]
+    initializer_stddev = depth_v**-0.5
+    value_relative_embeddings = (
+        tf.get_variable(name="value_relative_embeddings",
+                        shape=(2*max_length-1, depth_v),
+                        initializer=tf.random_normal_initializer(
+                            stddev=initializer_stddev)))
+    # slice out the right band of rel embeddings to save computation
+    padded_value_relative_embeddings = tf.pad(value_relative_embeddings,
+                                              [[tf.maximum(
+                                                  length-max_length, 0),
+                                                tf.maximum(
+                                                    length-max_length, 0)],
+                                               [0, 0]])
+    used_value_relative_embeddings = tf.slice(padded_value_relative_embeddings,
+                                              [tf.maximum(
+                                                  max_length-length,
+                                                  length-max_length),
+                                               0],
+                                              [2*length -1, -1])
+
+    ret += tf.einsum("bhlm, md -> bhld", relative_weights,
+                     used_value_relative_embeddings)
+    return ret
+
+
 def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
   """Attention to the source and a neighborhood to the left within a block.
 
@@ -1819,7 +1959,7 @@ def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
     return output
 
 
-def _local_unmasked_relative_to_absolute(x):
+def _relative_position_to_absolute_position_unmasked(x):
   """Converts tensor from relative to aboslute indexing for local attention.
 
   Args:
@@ -1837,7 +1977,7 @@ def _local_unmasked_relative_to_absolute(x):
   col_pad = tf.zeros((batch, heads, length, 1))
   x = tf.concat([x, col_pad], axis=3)
 
-  # Concat extra elements so to add up to shape (len+1, 2*len-2).
+  # Concat extra elements so to add up to shape (len+1, 2*len-1).
   flat_x = tf.reshape(x, [batch, heads, length * 2 * length])
   flat_pad = tf.zeros((batch, heads, length-1))
   flat_x_padded = tf.concat([flat_x, flat_pad], axis=2)
@@ -2034,10 +2174,13 @@ def masked_rel_local_attention_1d(q,
     rel_embed_length = 4 * default_block_length
     # Relative embeddings can be shared or unshared
     first_logits = tf.matmul(first_q, first_k, transpose_b=True)
+    initializer_stddev = depth_k**-0.5
     if share_rel_embed:
       relative_embeddings = (
           tf.get_variable(name="relative_embeddings",
-                          shape=(rel_embed_length, depth_k)))
+                          shape=(rel_embed_length, depth_k),
+                          initializer=tf.random_normal_initializer(
+                              stddev=initializer_stddev)))
       masked_relative_embeddings = tf.slice(
           relative_embeddings,
           [rel_embed_length - block_length, 0], [-1, -1])
@@ -2046,7 +2189,9 @@ def masked_rel_local_attention_1d(q,
     else:
       relative_embeddings = (
           tf.get_variable(name="relative_embeddings",
-                          shape=(heads, rel_embed_length, depth_k)))
+                          shape=(heads, rel_embed_length, depth_k),
+                          initializer=tf.random_normal_initializer(
+                              stddev=initializer_stddev)))
       masked_relative_embeddings = tf.slice(
           relative_embeddings,
           [0, rel_embed_length - block_length, 0], [-1, -1, -1])
@@ -2102,21 +2247,20 @@ def _reshape_for_relative(x):
           [rel_embed_length - 2*block_length, 0], [-1, -1])
       rel_logits = tf.einsum(
           "bhld,md->bhlm", rel_tail_q, used_relative_embeddings)
-      masked_rel_logits = tf.slice(rel_logits, [0, 0, 0, block_length],
-                                   [-1, -1, -1, -1])
     else:
       used_relative_embeddings = tf.slice(
           relative_embeddings,
           [0, rel_embed_length - 2*block_length, 0], [-1, -1, -1])
       rel_logits = tf.einsum(
           "bhld,hmd->bhlm", rel_tail_q, used_relative_embeddings)
-      masked_rel_logits = tf.slice(rel_logits, [0, 0, 0, block_length],
-                                   [-1, -1, -1, -1])
+
+    masked_rel_logits = tf.slice(rel_logits, [0, 0, 0, block_length],
+                                 [-1, -1, -1, -1])
     masked_rel_logits = _relative_position_to_absolute_position_masked(
         masked_rel_logits)
     unmasked_rel_logits = tf.slice(rel_logits, [0, 0, 0, 0],
                                    [-1, -1, -1, 2*block_length-1])
-    unmasked_rel_logits = _local_unmasked_relative_to_absolute(
+    unmasked_rel_logits = _relative_position_to_absolute_position_unmasked(
         unmasked_rel_logits)
     all_rel_logits = tf.concat([unmasked_rel_logits, masked_rel_logits],
                                axis=3)
@@ -2134,6 +2278,7 @@ def _reshape_for_relative(x):
         attention, 1.0 - dropout_rate,
         broadcast_dims=None)
     output = tf.matmul(attention, rel_v)
+
     # bring to [batch, heads, num_blocks-1, block_length, depth]
     output = tf.reshape(output,
                         [batch, num_blocks-1, heads, block_length, depth_v])
@@ -2144,6 +2289,7 @@ def _reshape_for_relative(x):
     output = tf.concat([first_output, output], axis=2)
     output = tf.slice(output, [0, 0, 0, 0], [-1, -1, original_length, -1])
     output = tf.reshape(output, [batch, heads, original_length, depth_v])
+
     return output
 
 
@@ -3207,6 +3353,17 @@ def multihead_attention(query_antecedent,
           dropout_rate,
           image_shapes,
           make_image_summary=make_image_summary)
+    elif attention_type == "dot_product_unmasked_relative_v2":
+      x = dot_product_unmasked_self_attention_relative_v2(
+          q,
+          k,
+          v,
+          bias,
+          max_length,
+          dropout_rate,
+          image_shapes,
+          make_image_summary=make_image_summary,
+          dropout_broadcast_dims=dropout_broadcast_dims)
     elif attention_type == "dot_product_relative_v2":
       x = dot_product_self_attention_relative_v2(
           q,
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 825cde343..cf47eff8b 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -398,6 +398,20 @@ def testDotProductAttentionRelative(self):
       res = session.run(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
+  def testDotProductUnMaskedAttentionRelativeV2(self):
+    x = np.random.rand(5, 7, 12, 32)
+    y = np.random.rand(5, 7, 12, 32)
+    with self.test_session() as session:
+      a = common_attention.dot_product_unmasked_self_attention_relative_v2(
+          tf.constant(x, dtype=tf.float32),
+          tf.constant(y, dtype=tf.float32),
+          tf.constant(y, dtype=tf.float32),
+          None,
+          35)
+      session.run(tf.global_variables_initializer())
+      res = session.run(a)
+    self.assertEqual(res.shape, (5, 7, 12, 32))
+
   def testBiasBatchCoordinates(self):
     """Testing the batch coordinates mask."""
     q = tf.constant([0, 0, 1, 1, 1, 1, 2, 2, 2], dtype=tf.int32)

From 9bea9b479da5b1de6f9f7bb7159914299995e28c Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 31 Jul 2018 17:25:22 -0700
Subject: [PATCH 0429/2720] Move GithubFunctionDocstring problem to
 Text2TextProblem

PiperOrigin-RevId: 206856332
---
 tensor2tensor/data_generators/function_docstring.py | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
index c44319342..6ddd9994b 100644
--- a/tensor2tensor/data_generators/function_docstring.py
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -16,7 +16,7 @@
 import csv
 import six
 from tensor2tensor.data_generators import generator_utils
-from tensor2tensor.data_generators import translate
+from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 import tensorflow as tf
@@ -30,7 +30,7 @@
 
 
 @registry.register_problem
-class GithubFunctionDocstring(translate.TranslateProblem):
+class GithubFunctionDocstring(text_problems.Text2TextProblem):
   """Function and Docstring similarity Problem.
 
   This problem contains the data consisting of function
@@ -59,10 +59,6 @@ def is_generate_per_split(self):
   def approx_vocab_size(self):
     return 2**13
 
-  def source_data_files(self, _):
-    # TODO(sanyamkapoor): Manually separate train/eval data set.
-    return self.pair_files_list
-
   @property
   def max_samples_for_vocab(self):
     # FIXME(sanyamkapoor): This exists to handle memory explosion.
@@ -83,7 +79,8 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
         {"inputs": "STRING", "targets": "STRING"}
     """
 
-    csv_file_names = self.source_data_files(dataset_split)
+    # TODO(sanyamkapoor): Manually separate train/eval data set.
+    csv_file_names = self.pair_files_list
     csv_files = [
         generator_utils.maybe_download(tmp_dir, filename,
                                        "{}/{}".format(self.base_url,

From 354c0bccd5089eba9d89b07fa3c09815f6dfe43c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 31 Jul 2018 17:44:07 -0700
Subject: [PATCH 0430/2720] Better StringIO import and dataset format.

PiperOrigin-RevId: 206858458
---
 .../data_generators/function_docstring.py     | 20 +++++++------------
 1 file changed, 7 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
index 6ddd9994b..7d72d5898 100644
--- a/tensor2tensor/data_generators/function_docstring.py
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -14,20 +14,13 @@
 # limitations under the License.
 """Github function/text similatrity problems."""
 import csv
-import six
+from six import StringIO
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 import tensorflow as tf
 
-# pylint: disable=g-import-not-at-top
-if six.PY2:
-  from StringIO import StringIO
-else:
-  from io import StringIO
-# pylint: enable=g-import-not-at-top
-
 
 @registry.register_problem
 class GithubFunctionDocstring(text_problems.Text2TextProblem):
@@ -47,7 +40,10 @@ def base_url(self):
   @property
   def pair_files_list(self):
     return [
-        "func-doc-pairs-000{:02}-of-00100.csv".format(i)
+        [
+            "{}/func-doc-pairs-000{:02}-of-00100.csv".format(self.base_url, i),
+            ("func-doc-pairs-000{:02}-of-00100.csv".format(i),)
+        ]
         for i in range(100)
     ]
 
@@ -82,10 +78,8 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
     # TODO(sanyamkapoor): Manually separate train/eval data set.
     csv_file_names = self.pair_files_list
     csv_files = [
-        generator_utils.maybe_download(tmp_dir, filename,
-                                       "{}/{}".format(self.base_url,
-                                                      filename))
-        for filename in csv_file_names
+        generator_utils.maybe_download(tmp_dir, file_list[0], uri)
+        for uri, file_list in csv_file_names
     ]
 
     for pairs_file in csv_files:

From c8e3cc1bdf687ca0c0d04ada0efe6ec1e7441e00 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 31 Jul 2018 19:43:19 -0700
Subject: [PATCH 0431/2720] putting tf.scan back in. fixed.

PiperOrigin-RevId: 206869200
---
 tensor2tensor/layers/modalities.py          |   4 +-
 tensor2tensor/models/research/next_frame.py | 140 +++++++++++++-------
 2 files changed, 97 insertions(+), 47 deletions(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index e1ae5401f..64046e0ac 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -700,7 +700,9 @@ def targets_bottom(self, x):  # pylint: disable=arguments-differ
     return common_layers.convert_rgb_to_real(x)
 
   def top(self, body_output, _):
-    frames = tf.stack(body_output, axis=1)
+    frames = body_output
+    if isinstance(body_output, list):
+      frames = tf.stack(body_output, axis=1)
     rgb_frames = common_layers.convert_real_to_rgb(frames)
     common_layers.summarize_video(rgb_frames, "body_output")
     return tf.expand_dims(rgb_frames, axis=-1)
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 2f2e3c680..5ca02e35e 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -246,7 +246,7 @@ def construct_latent_tower(self, images):
         latent_num_frames = (self.hparams.video_num_input_frames +
                              self.hparams.video_num_target_frames)
       tf.logging.info("Creating latent tower with %d frames."%latent_num_frames)
-      latent_images = images[:latent_num_frames]
+      latent_images = tf.unstack(images[:latent_num_frames], axis=0)
       images = tf.concat(latent_images, 3)
 
       x = images
@@ -567,47 +567,61 @@ def construct_model(self,
     Raises:
       ValueError: if more than 1 mask specified for DNA model.
     """
-    batch_size = common_layers.shape_list(images[0])[0]
     context_frames = self.hparams.video_num_input_frames
 
-    # Predicted images and rewards.
-    gen_rewards, gen_images = [], []
+    def process_single_frame(prev_outputs, inputs):
+      """Process a single frame of the video."""
+      cur_image, cur_reward, action = inputs
+      time_step, prev_image, prev_reward, lstm_states = prev_outputs[:4]
 
-    # LSTM states.
-    lstm_state = [None] * 7
-    reward_lstm_state = [None] * 5
+      # TODO(mbz): No scheduled sampling for now!
+      input_image, input_reward = tf.cond(
+          tf.greater(time_step, context_frames),
+          lambda: (prev_image, prev_reward),
+          lambda: (cur_image, cur_reward))
+
+      # Prediction
+      pred_image, lstm_states = self.construct_predictive_tower(
+          input_image, input_reward, action, lstm_states, latent)
+
+      if self.hparams.reward_prediction:
+        reward_lstm_states = prev_outputs[4]
+        pred_reward, reward_lstm_states = self.reward_prediction(
+            input_image, input_reward, action, reward_lstm_states, latent)
+      else:
+        pred_reward = input_reward
+
+      time_step += 1
+      outputs = (time_step, pred_image, pred_reward, lstm_states)
+      if self.hparams.reward_prediction:
+        outputs += (reward_lstm_states,)
+
+      return outputs
 
     # Latent tower
+    latent = None
     if self.hparams.stochastic_model:
       latent_mean, latent_std = self.construct_latent_tower(images)
+      latent = self.get_gaussian_latent(latent_mean, latent_std)
 
-    pred_image, pred_reward, latent = None, None, None
-    for timestep, image, action, reward in zip(
-        range(len(images)-1), images[:-1], actions[:-1], rewards[:-1]):
-      # Scheduled Sampling
-      done_warm_start = timestep > context_frames - 1
-      groundtruth_items = [image, reward]
-      generated_items = [pred_image, pred_reward]
-      input_image, input_reward = self.get_scheduled_sample_inputs(
-          done_warm_start, groundtruth_items, generated_items, batch_size)
+    # HACK: Do first step outside to initialize all the variables
+    lstm_states, reward_lstm_states = [None] * 7, [None] * 5
+    inputs = images[0], rewards[0], actions[0]
+    prev_outputs = (tf.constant(0), images[0], rewards[0], lstm_states)
+    if self.hparams.reward_prediction:
+      prev_outputs += (reward_lstm_states,)
 
-      # Latent
-      if self.hparams.stochastic_model:
-        if timestep == 0 or self.hparams.multi_latent:
-          latent = self.get_gaussian_latent(latent_mean, latent_std)
+    initializers = process_single_frame(prev_outputs, inputs)
+    first_gen_images = tf.expand_dims(initializers[1], axis=0)
+    first_gen_rewards = tf.expand_dims(initializers[2], axis=0)
 
-      # Prediction
-      pred_image, lstm_state = self.construct_predictive_tower(
-          input_image, input_reward, action, lstm_state, latent)
+    inputs = (images[1:-1], actions[1:-1], rewards[1:-1])
 
-      if self.hparams.reward_prediction:
-        pred_reward, reward_lstm_state = self.reward_prediction(
-            input_image, input_reward, action, reward_lstm_state, latent)
-      else:
-        pred_reward = input_reward
+    outputs = tf.scan(process_single_frame, inputs, initializers)
+    gen_images, gen_rewards = outputs[1:3]
 
-      gen_images.append(pred_image)
-      gen_rewards.append(pred_reward)
+    gen_images = tf.concat((first_gen_images, gen_images), axis=0)
+    gen_rewards = tf.concat((first_gen_rewards, gen_rewards), axis=0)
 
     return gen_images, gen_rewards, [latent_mean], [latent_std]
 
@@ -783,15 +797,19 @@ def get_input_if_exists(self, features, key, batch_size, num_frames):
       x = features[key]
     else:
       x = tf.zeros((batch_size, num_frames, 1, self.hparams.hidden_size))
-    return tf.unstack(x, axis=1)
+    return self.swap_time_and_batch_axes(x)
+
+  def swap_time_and_batch_axes(self, x):
+    transposed_axes = tf.concat([[1, 0], tf.range(2, tf.rank(x))], axis=0)
+    return tf.transpose(x, transposed_axes)
 
   def body(self, features):
     hparams = self.hparams
     batch_size = common_layers.shape_list(features["inputs"])[0]
 
-    # Split inputs and targets time-wise into a list of frames.
-    input_frames = tf.unstack(features["inputs"], axis=1)
-    target_frames = tf.unstack(features["targets"], axis=1)
+    # Swap time and batch axes.
+    input_frames = self.swap_time_and_batch_axes(features["inputs"])
+    target_frames = self.swap_time_and_batch_axes(features["targets"])
 
     # Get actions if exist otherwise use zeros
     input_actions = self.get_input_if_exists(
@@ -805,15 +823,15 @@ def body(self, features):
     target_rewards = self.get_input_if_exists(
         features, "target_reward", batch_size, hparams.video_num_target_frames)
 
-    all_actions = input_actions + target_actions
-    all_rewards = input_rewards + target_rewards
-    all_frames = input_frames + target_frames
+    all_actions = tf.concat([input_actions, target_actions], axis=0)
+    all_rewards = tf.concat([input_rewards, target_rewards], axis=0)
+    all_frames = tf.concat([input_frames, target_frames], axis=0)
 
     # Each image is being used twice, in latent tower and main tower.
     # This is to make sure we are using the *same* image for both, ...
     # ... given how TF queues work.
     # NOT sure if this is required at all. Doesn"t hurt though! :)
-    all_frames = [tf.identity(frame) for frame in all_frames]
+    all_frames = tf.identity(all_frames)
 
     gen_images, gen_rewards, latent_means, latent_stds = self.construct_model(
         images=all_frames,
@@ -865,15 +883,29 @@ def anneal_loss(step_num):
       tf.summary.scalar("beta", beta)
       tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
 
-    extra_loss = beta * tf.reduce_mean(kl_loss)
+    extra_loss = beta * kl_loss
+
+    # Ignore the predictions from the input frames.
+    # This is NOT the same as original paper/implementation.
     predictions = gen_images[hparams.video_num_input_frames-1:]
-    reward_pred = tf.stack(
-        gen_rewards[hparams.video_num_input_frames-1:], axis=1)
+    reward_pred = gen_rewards[hparams.video_num_input_frames-1:]
     reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove undeeded dimension.
 
-    frames_gt = tf.concat(all_frames[hparams.video_num_input_frames:], axis=1)
-    frames_pd = tf.concat(predictions, axis=1)
-    tf.summary.image("full_video", tf.concat([frames_gt, frames_pd], axis=2))
+    # TODO(mbz): clean this up!
+    def fix_video_dims_and_concat_on_x_axis(x):
+      x = tf.transpose(x, [1, 3, 4, 0, 2])
+      x = tf.reshape(x, [batch_size, 64, 3, -1])
+      x = tf.transpose(x, [0, 3, 1, 2])
+      return x
+
+    frames_gd = fix_video_dims_and_concat_on_x_axis(target_frames)
+    frames_pd = fix_video_dims_and_concat_on_x_axis(predictions)
+    side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
+    tf.summary.image("full_video", side_by_side_video)
+
+    # Swap back time and batch axes.
+    predictions = self.swap_time_and_batch_axes(predictions)
+    reward_pred = self.swap_time_and_batch_axes(reward_pred)
 
     return_targets = predictions
     if "target_reward" in features:
@@ -887,6 +919,10 @@ class NextFrameStochasticTwoFrames(NextFrameStochastic):
   """Stochastic next-frame model with 2 frames posterior."""
 
   def construct_model(self, images, actions, rewards):
+    images = tf.unstack(images, axis=0)
+    actions = tf.unstack(actions, axis=0)
+    rewards = tf.unstack(rewards, axis=0)
+
     batch_size = common_layers.shape_list(images[0])[0]
     context_frames = self.hparams.video_num_input_frames
 
@@ -928,6 +964,9 @@ def construct_model(self, images, actions, rewards):
       gen_images.append(pred_image)
       gen_rewards.append(pred_reward)
 
+    gen_images = tf.stack(gen_images, axis=0)
+    gen_rewards = tf.stack(gen_rewards, axis=0)
+
     return gen_images, gen_rewards, latent_means, latent_stds
 
 
@@ -1122,8 +1161,7 @@ def construct_model(self, images, actions, rewards):
     predictor_rnn_layers = self.hparams.predictor_rnn_layers
     context_frames = self.hparams.video_num_input_frames
 
-    seq_len = len(images)
-    batch_size, _, _, color_channels = common_layers.shape_list(images[0])
+    seq_len, batch_size, _, _, color_channels = common_layers.shape_list(images)
 
     # LSTM initial sizesstates.
     predictor_states = [None] * predictor_rnn_layers
@@ -1132,6 +1170,7 @@ def construct_model(self, images, actions, rewards):
     tf.logging.info(">>>> Encoding")
     # Encoding:
     enc_images, enc_skips = [], []
+    images = tf.unstack(images, axis=0)
     for i, image in enumerate(images):
       with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
         enc, skips = self.encoder(image, rnn_size)
@@ -1183,6 +1222,7 @@ def construct_model(self, images, actions, rewards):
         gen_images.append(x_pred)
 
     tf.logging.info(">>>> Done")
+    gen_images = tf.stack(gen_images, axis=0)
     return gen_images, fake_reward_prediction, pred_mu, pred_logvar
 
 
@@ -1274,6 +1314,10 @@ def construct_model(self, images, actions, rewards):
       latent_stds: list of gaussian stds conditioned on the input at
                    every frame.
     """
+    images = tf.unstack(images, axis=0)
+    actions = tf.unstack(actions, axis=0)
+    rewards = tf.unstack(rewards, axis=0)
+
     latent_dims = self.hparams.z_dim
     context_frames = self.hparams.video_num_input_frames
     seq_len = len(images)
@@ -1349,6 +1393,10 @@ def construct_model(self, images, actions, rewards):
         latent_means.append(mu)
         latent_stds.append(log_sigma_sq)
 
+    gen_cond_video = tf.stack(gen_cond_video, axis=0)
+    gen_prior_video = tf.stack(gen_prior_video, axis=0)
+    fake_rewards = tf.stack(fake_rewards, axis=0)
+
     if train_mode:
       return gen_cond_video, fake_rewards, latent_means, latent_stds
     else:

From cd74ce010a5dd18a576d57f3769c03de04d2954e Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 31 Jul 2018 20:30:41 -0700
Subject: [PATCH 0432/2720] + Adding probability based scheduled sampling (SS).

PiperOrigin-RevId: 206872932
---
 tensor2tensor/models/research/next_frame.py   | 133 +++++++++++++-----
 .../models/research/next_frame_params.py      |   3 +
 .../models/research/next_frame_test.py        |  43 ++++++
 3 files changed, 143 insertions(+), 36 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 5ca02e35e..b5ef66b84 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from functools import partial
 import six
 
 from tensor2tensor.layers import common_attention
@@ -569,16 +570,19 @@ def construct_model(self,
     """
     context_frames = self.hparams.video_num_input_frames
 
+    batch_size = common_layers.shape_list(images)[1]
+    ss_func = self.get_scheduled_sample_func(batch_size)
+
     def process_single_frame(prev_outputs, inputs):
       """Process a single frame of the video."""
       cur_image, cur_reward, action = inputs
       time_step, prev_image, prev_reward, lstm_states = prev_outputs[:4]
 
-      # TODO(mbz): No scheduled sampling for now!
-      input_image, input_reward = tf.cond(
-          tf.greater(time_step, context_frames),
-          lambda: (prev_image, prev_reward),
-          lambda: (cur_image, cur_reward))
+      generated_items = [prev_image, prev_reward]
+      groundtruth_items = [cur_image, cur_reward]
+      done_warm_start = tf.greater(time_step, context_frames - 1)
+      input_image, input_reward = self.get_scheduled_sample_inputs(
+          done_warm_start, groundtruth_items, generated_items, ss_func)
 
       # Prediction
       pred_image, lstm_states = self.construct_predictive_tower(
@@ -717,22 +721,23 @@ def dna_transformation(self,
                             [4])
     return tf.reduce_sum(kernel * inputs, [3], keep_dims=False)
 
-  def scheduled_sample(self,
-                       ground_truth_x,
-                       generated_x,
-                       batch_size,
-                       num_ground_truth):
+  @staticmethod
+  def scheduled_sample_count(ground_truth_x,
+                             generated_x,
+                             batch_size,
+                             scheduled_sample_var):
     """Sample batch with specified mix of groundtruth and generated data points.
 
     Args:
       ground_truth_x: tensor of ground-truth data points.
       generated_x: tensor of generated data points.
       batch_size: batch size
-      num_ground_truth: number of ground-truth examples to include in batch.
+      scheduled_sample_var: number of ground-truth examples to include in batch.
     Returns:
       New batch with num_ground_truth sampled from ground_truth_x and the rest
       from generated_x.
     """
+    num_ground_truth = scheduled_sample_var
     idx = tf.random_shuffle(tf.range(batch_size))
     ground_truth_idx = tf.gather(idx, tf.range(num_ground_truth))
     generated_idx = tf.gather(idx, tf.range(num_ground_truth, batch_size))
@@ -742,41 +747,91 @@ def scheduled_sample(self,
     return tf.dynamic_stitch([ground_truth_idx, generated_idx],
                              [ground_truth_examps, generated_examps])
 
-  def get_scheduled_sample_inputs(
-      self, done_warm_start, groundtruth_items, generated_items, batch_size):
+  @staticmethod
+  def scheduled_sample_prob(ground_truth_x,
+                            generated_x,
+                            batch_size,
+                            scheduled_sample_var):
+    """Probability based scheduled sampling.
 
-    with tf.variable_scope("scheduled_sampling", reuse=tf.AUTO_REUSE):
-      if self.hparams.mode != tf.estimator.ModeKeys.TRAIN:
-        feedself = True
+    Args:
+      ground_truth_x: tensor of ground-truth data points.
+      generated_x: tensor of generated data points.
+      batch_size: batch size
+      scheduled_sample_var: probability of choosing from ground_truth.
+    Returns:
+      New batch with randomly selected data points.
+    """
+    probability_threshold = scheduled_sample_var
+    probability_of_generated = tf.random_uniform([batch_size])
+    array_ind = tf.to_int32(probability_of_generated > probability_threshold)
+    indices = tf.range(batch_size) + array_ind * batch_size
+    xy = tf.concat([ground_truth_x, generated_x], axis=0)
+    output = tf.gather(xy, indices)
+    return output
+
+  def get_scheduled_sample_func(self, batch_size):
+    """Creates a function for scheduled sampling based on given hparams."""
+    with tf.variable_scope("scheduled_sampling_func", reuse=False):
+      iter_num = tf.train.get_global_step()
+      # TODO(lukaszkaiser): figure out why iter_num can be None.
+      if iter_num is None:
+        iter_num = _LARGE_STEP_NUMBER
+
+      if self.hparams.scheduled_sampling_mode == "prob":
+        decay_steps = self.hparams.scheduled_sampling_decay_steps
+        probability = tf.train.polynomial_decay(
+            1.0, iter_num, decay_steps, 0.0)
+        scheduled_sampling_func = NextFrameStochastic.scheduled_sample_prob
+        scheduled_sampling_func_var = probability
       else:
-        # Scheduled sampling:
         # Calculate number of ground-truth frames to pass in.
-        feedself = False
-        iter_num = tf.train.get_global_step()
-        # TODO(mbz): what should it be if it's undefined?
-        if iter_num is None:
-          iter_num = _LARGE_STEP_NUMBER
         k = self.hparams.scheduled_sampling_k
         num_ground_truth = tf.to_int32(
             tf.round(
                 tf.to_float(batch_size) *
                 (k / (k + tf.exp(tf.to_float(iter_num) / tf.to_float(k))))))
-        tf.summary.scalar("num_ground_truth", num_ground_truth)
+        scheduled_sampling_func = NextFrameStochastic.scheduled_sample_count
+        scheduled_sampling_func_var = num_ground_truth
+
+      tf.summary.scalar("scheduled_sampling_var", scheduled_sampling_func_var)
+      partial_func = partial(scheduled_sampling_func,
+                             batch_size=batch_size,
+                             scheduled_sample_var=scheduled_sampling_func_var)
+      return partial_func
+
+  def get_scheduled_sample_inputs(self,
+                                  done_warm_start,
+                                  groundtruth_items,
+                                  generated_items,
+                                  scheduled_sampling_func):
+    """Scheduled sampling.
 
-      if feedself and done_warm_start:
-        # Feed in generated stuff.
-        output_items = generated_items
-      elif done_warm_start:
+    Args:
+      done_warm_start: whether we are done with warm start or not.
+      groundtruth_items: list of ground truth items.
+      generated_items: list of generated items.
+      scheduled_sampling_func: scheduled sampling function to choose between
+        groundtruth items and generated items.
+
+    Returns:
+      A mix list of ground truth and generated items.
+    """
+    def sample():
+      """Calculate the scheduled sampling params based on iteration number."""
+      with tf.variable_scope("scheduled_sampling", reuse=tf.AUTO_REUSE):
         output_items = []
         for item_gt, item_gen in zip(groundtruth_items, generated_items):
-          # Scheduled sampling
-          output_items.append(self.scheduled_sample(
-              item_gt, item_gen, batch_size, num_ground_truth))
-      else:
-        # Feed in ground_truth
-        output_items = groundtruth_items
+          output_items.append(scheduled_sampling_func(item_gt, item_gen))
+        return output_items
+
+    cases = {
+        tf.logical_not(self.is_training): lambda: generated_items,
+        tf.logical_not(done_warm_start): lambda: groundtruth_items,
+    }
+    output_items = tf.case(cases, default=sample)
 
-      return output_items
+    return output_items
 
   # TODO(mbz): use tf.distributions.kl_divergence instead.
   def kl_divergence(self, mu, log_sigma):
@@ -933,6 +988,9 @@ def construct_model(self, images, actions, rewards):
     lstm_state = [None] * 7
     reward_lstm_state = [None] * 5
 
+    # Create scheduled sampling function
+    ss_func = self.get_scheduled_sample_func(batch_size)
+
     pred_image, pred_reward, latent = None, None, None
     for timestep, image, action, reward in zip(
         range(len(images)-1), images[:-1], actions[:-1], rewards[:-1]):
@@ -941,7 +999,7 @@ def construct_model(self, images, actions, rewards):
       groundtruth_items = [image, reward]
       generated_items = [pred_image, pred_reward]
       input_image, input_reward = self.get_scheduled_sample_inputs(
-          done_warm_start, groundtruth_items, generated_items, batch_size)
+          done_warm_start, groundtruth_items, generated_items, ss_func)
 
       # Latent
       # TODO(mbz): should we use input_image iunstead of image?
@@ -1346,6 +1404,9 @@ def construct_model(self, images, actions, rewards):
     pred_image, prior_latent_state, cond_latent_state = None, None, None
     train_mode = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
 
+    # Create scheduled sampling function
+    ss_func = self.get_scheduled_sample_func(batch_size)
+
     with tf.variable_scope("prediction", reuse=tf.AUTO_REUSE):
 
       for step, (image, action, reward, mu, log_sigma_sq) in enumerate(iterable):  # pylint:disable=line-too-long
@@ -1369,7 +1430,7 @@ def construct_model(self, images, actions, rewards):
         groundtruth_items = [image]
         generated_items = [pred_image]
         input_image = self.get_scheduled_sample_inputs(
-            done_warm_start, groundtruth_items, generated_items, batch_size)[0]
+            done_warm_start, groundtruth_items, generated_items, ss_func)[0]
 
         all_latents = tf.concat([enc_cond_latent, enc_prior_latent], axis=0)
         all_image = tf.concat([input_image, input_image], axis=0)
diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index 663f683ed..44ecf5bec 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -72,6 +72,9 @@ def next_frame_stochastic():
   hparams.add_hparam("multi_latent", False)
   hparams.add_hparam("relu_shift", 1e-12)
   hparams.add_hparam("dna_kernel_size", 5)
+  # Scheduled sampling method. Choose between prob or count.
+  hparams.add_hparam("scheduled_sampling_mode", "prob")
+  hparams.add_hparam("scheduled_sampling_decay_steps", 10000)
   hparams.add_hparam("scheduled_sampling_k", 900.0)
   hparams.add_hparam("latent_num_frames", 0)  # 0 means use all frames.
   hparams.add_hparam("tiny_mode", False)
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index 77aedf738..69ba0d899 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -107,6 +107,49 @@ def testStochasticSavp(self):
         next_frame.NextFrameSavp,
         1)
 
+  @staticmethod
+  def run_scheduled_sample_func(func, var, batch_size):
+    ground_truth_x = list(range(1, batch_size+1))
+    generated_x = [-x for x in ground_truth_x]
+    ground_truth_x = tf.convert_to_tensor(ground_truth_x)
+    generated_x = tf.convert_to_tensor(generated_x)
+    ss_out = func(ground_truth_x, generated_x, batch_size, var)
+    with tf.Session() as session:
+      output = session.run([ground_truth_x, generated_x, ss_out])
+    return output
+
+  def testScheduledSampleProbStart(self):
+    ground_truth_x, _, ss_out = NextFrameTest.run_scheduled_sample_func(
+        next_frame.NextFrameStochastic.scheduled_sample_prob, 1.0, 10)
+    self.assertAllEqual(ground_truth_x, ss_out)
+
+  def testScheduledSampleProbMid(self):
+    _, _, ss_out = NextFrameTest.run_scheduled_sample_func(
+        next_frame.NextFrameStochastic.scheduled_sample_prob, 0.5, 1000)
+    positive_count = np.sum(ss_out > 0)
+    self.assertAlmostEqual(positive_count / 1000.0, 0.5, places=2)
+
+  def testScheduledSampleProbEnd(self):
+    _, generated_x, ss_out = NextFrameTest.run_scheduled_sample_func(
+        next_frame.NextFrameStochastic.scheduled_sample_prob, 0.0, 10)
+    self.assertAllEqual(generated_x, ss_out)
+
+  def testScheduledSampleCountStart(self):
+    ground_truth_x, _, ss_out = NextFrameTest.run_scheduled_sample_func(
+        next_frame.NextFrameStochastic.scheduled_sample_count, 10, 10)
+    self.assertAllEqual(ground_truth_x, ss_out)
+
+  def testScheduledSampleCountMid(self):
+    _, _, ss_out = NextFrameTest.run_scheduled_sample_func(
+        next_frame.NextFrameStochastic.scheduled_sample_count, 5, 10)
+    positive_count = np.sum(ss_out > 0)
+    self.assertEqual(positive_count, 5)
+
+  def testScheduledSampleCountEnd(self):
+    _, generated_x, ss_out = NextFrameTest.run_scheduled_sample_func(
+        next_frame.NextFrameStochastic.scheduled_sample_count, 0, 10)
+    self.assertAllEqual(generated_x, ss_out)
+
   def testDynamicTileAndConcat(self):
     with tf.Graph().as_default():
       # image = (1 X 4 X 4 X 1)

From aa95e4c8a1f25204ba695494fc785db042494402 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 31 Jul 2018 21:13:14 -0700
Subject: [PATCH 0433/2720] data generator for wsj parsing.

PiperOrigin-RevId: 206876325
---
 tensor2tensor/data_generators/wsj_parsing.py | 48 ++++++++++++++++++--
 1 file changed, 45 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/wsj_parsing.py b/tensor2tensor/data_generators/wsj_parsing.py
index 30d91572b..3db8b7060 100644
--- a/tensor2tensor/data_generators/wsj_parsing.py
+++ b/tensor2tensor/data_generators/wsj_parsing.py
@@ -14,9 +14,14 @@
 # limitations under the License.
 """Data generators for parsing data-sets."""
 
-# import os
-# from tensor2tensor.data_generators import generator_utils
-
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
 import tensorflow as tf
 
 
@@ -26,6 +31,43 @@
 FLAGS = tf.flags.FLAGS
 
 
+@registry.register_problem
+class WsjParsing(text_problems.Text2textTmpdir):
+  """Generate vocabulary and training data for parsing.
+  """
+
+  # These files are used for vocab generation
+  TRAIN_FILES = ("wsj.train.text.txt", "wsj.train.tags.txt")
+
+  # These files are used for generating encoded samples
+  TRAIN_FILES_TREE = "wsjTrain.trees"
+  EVAL_FILES_TREE = "wsjEval.trees"
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    del data_dir
+    is_training = dataset_split == problem.DatasetSplit.TRAIN
+    tree_file = self.TRAIN_FILES_TREE if is_training else self.EVAL_FILES_TREE
+    tree_file_path = os.path.join(tmp_dir, tree_file)
+    with tf.gfile.GFile(tree_file_path, mode="r") as cur_tree_file:
+      for line in cur_tree_file:
+        (words, tags) = words_and_tags_from_wsj_tree(line)
+        yield {"inputs": words, "targets": tags}
+
+  def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
+    generator = self.generate_samples(data_dir, tmp_dir, dataset_split)
+    encoder = self.get_or_create_vocab(data_dir, tmp_dir)
+    return text_problems.text2text_generate_encoded(generator, encoder,
+                                                    has_inputs=self.has_inputs)
+
+  def generate_text_for_vocab(self, data_dir, tmp_dir):
+    files = [os.path.join(tmp_dir, f) for f in self.TRAIN_FILES]
+    inputs_file, targets_file = files
+    for sample in text_problems.text2text_txt_iterator(inputs_file,
+                                                       targets_file):
+      yield sample["inputs"]
+      yield sample["targets"]
+
+
 def words_and_tags_from_wsj_tree(tree_string):
   """Generates linearized trees and tokens from the wsj tree format.
 

From e88adbee773a2b88a0533ee02c87697d68a1bdb6 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 1 Aug 2018 08:03:28 -0700
Subject: [PATCH 0434/2720] Remove inter-example padding from packed datasets. 
 Previously,

PiperOrigin-RevId: 206929744
---
 tensor2tensor/data_generators/text_problems.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index a17e47f3f..16afc53a1 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -172,6 +172,15 @@ def packed_length(self):
     """
     return None
 
+  @property
+  def packed_spacing(self):
+    """If this is a packed dataset, how much padding to insert between examples.
+
+    Returns:
+      int
+    """
+    return 0
+
   # END: Subclass interface
 
   @property
@@ -238,6 +247,7 @@ def _maybe_pack_examples(self, generator):
         generator,
         self.has_inputs,
         self.packed_length,
+        spacing=self.packed_spacing,
         chop_long_sequences=not self.has_inputs)
 
   def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):

From 502e3fe1b4f4eafa8a9b668189d5112b637c59f9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 1 Aug 2018 09:48:15 -0700
Subject: [PATCH 0435/2720] Travis fixes for next_frame_test.py
 testDynamicTileAndConcat

PiperOrigin-RevId: 206943401
---
 tensor2tensor/models/research/next_frame_test.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index 69ba0d899..55bbeacc3 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -157,8 +157,8 @@ def testDynamicTileAndConcat(self):
                [2, 4, 5, 6],
                [7, 8, 9, 10],
                [7, 9, 10, 1]]
-      image = tf.expand_dims(tf.expand_dims(image, axis=0), axis=-1)
-      image_t = tf.cast(tf.convert_to_tensor(image), dtype=tf.float32)
+      image_t = tf.expand_dims(tf.expand_dims(image, axis=0), axis=-1)
+      image_t = tf.cast(image_t, dtype=tf.float32)
 
       # latent = (1 X 2)
       latent = np.array([[90, 100]])
@@ -167,11 +167,11 @@ def testDynamicTileAndConcat(self):
       with tf.Session() as session:
         tiled = next_frame.NextFrameStochastic.tile_and_concat(
             image_t, latent_t)
-        tiled_np = session.run(tiled)
+        tiled_np, image_np = session.run([tiled, image_t])
         tiled_latent = tiled_np[0, :, :, -1]
         self.assertAllEqual(tiled_np.shape, (1, 4, 4, 2))
 
-        self.assertAllEqual(tiled_np[:, :, :, :1], image)
+        self.assertAllEqual(tiled_np[:, :, :, :1], image_np)
         self.assertAllEqual(
             tiled_latent,
             [[90, 90, 90, 90],

From 93745d175936cf9c3f77682653d5aba18ed14141 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 1 Aug 2018 09:48:58 -0700
Subject: [PATCH 0436/2720] More Travis fixes, on and before TF 1.7 the way the
 test was written failed.

PiperOrigin-RevId: 206943496
---
 tensor2tensor/layers/common_layers_test.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index c7f308952..fd617dd10 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -634,7 +634,8 @@ def testCycleGANUpsampleNnUpsampleConv(self):
     num_channels = 3
     output_filters = 10
     stride = [2, 3]  # we want height to be x2 and width to be x3
-    random_input = np.random.rand(batch, height, width, num_channels)
+    random_input = np.random.rand(batch, height, width, num_channels).astype(
+        np.float32)
 
     # nn_upsample_conv gives exactly the shapes we'd expect.
     upsampled_output = common_layers.cyclegan_upsample(
@@ -653,7 +654,8 @@ def testCycleGANUpsampleBilinearUpsampleConv(self):
     num_channels = 3
     output_filters = 10
     stride = [2, 3]  # we want height to be x2 and width to be x3
-    random_input = np.random.rand(batch, height, width, num_channels)
+    random_input = np.random.rand(batch, height, width, num_channels).astype(
+        np.float32)
 
     # bilinear_upsample_conv gives exactly the shapes we'd expect.
     upsampled_output = common_layers.cyclegan_upsample(
@@ -672,7 +674,8 @@ def testCycleGANUpsampleConv2dTranspose(self):
     num_channels = 3
     output_filters = 10
     stride = [2, 3]  # we want height to be x2 and width to be x3
-    random_input = np.random.rand(batch, height, width, num_channels)
+    random_input = np.random.rand(batch, height, width, num_channels).astype(
+        np.float32)
 
     # conv2d_transpose is a little tricky.
     # height_new = (height_old - 1) * stride + kernel - 2*padding - correction

From eb5d4bad2b34689f5f0a53daa85b370f968b5d65 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 1 Aug 2018 09:59:00 -0700
Subject: [PATCH 0437/2720] Camelcase method name to fix lint

PiperOrigin-RevId: 206945092
---
 tensor2tensor/models/research/next_frame_test.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index 55bbeacc3..e4d2f4178 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -108,7 +108,7 @@ def testStochasticSavp(self):
         1)
 
   @staticmethod
-  def run_scheduled_sample_func(func, var, batch_size):
+  def runScheduledSampleFunc(func, var, batch_size):
     ground_truth_x = list(range(1, batch_size+1))
     generated_x = [-x for x in ground_truth_x]
     ground_truth_x = tf.convert_to_tensor(ground_truth_x)
@@ -119,34 +119,34 @@ def run_scheduled_sample_func(func, var, batch_size):
     return output
 
   def testScheduledSampleProbStart(self):
-    ground_truth_x, _, ss_out = NextFrameTest.run_scheduled_sample_func(
+    ground_truth_x, _, ss_out = NextFrameTest.runScheduledSampleFunc(
         next_frame.NextFrameStochastic.scheduled_sample_prob, 1.0, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
   def testScheduledSampleProbMid(self):
-    _, _, ss_out = NextFrameTest.run_scheduled_sample_func(
+    _, _, ss_out = NextFrameTest.runScheduledSampleFunc(
         next_frame.NextFrameStochastic.scheduled_sample_prob, 0.5, 1000)
     positive_count = np.sum(ss_out > 0)
     self.assertAlmostEqual(positive_count / 1000.0, 0.5, places=2)
 
   def testScheduledSampleProbEnd(self):
-    _, generated_x, ss_out = NextFrameTest.run_scheduled_sample_func(
+    _, generated_x, ss_out = NextFrameTest.runScheduledSampleFunc(
         next_frame.NextFrameStochastic.scheduled_sample_prob, 0.0, 10)
     self.assertAllEqual(generated_x, ss_out)
 
   def testScheduledSampleCountStart(self):
-    ground_truth_x, _, ss_out = NextFrameTest.run_scheduled_sample_func(
+    ground_truth_x, _, ss_out = NextFrameTest.runScheduledSampleFunc(
         next_frame.NextFrameStochastic.scheduled_sample_count, 10, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
   def testScheduledSampleCountMid(self):
-    _, _, ss_out = NextFrameTest.run_scheduled_sample_func(
+    _, _, ss_out = NextFrameTest.runScheduledSampleFunc(
         next_frame.NextFrameStochastic.scheduled_sample_count, 5, 10)
     positive_count = np.sum(ss_out > 0)
     self.assertEqual(positive_count, 5)
 
   def testScheduledSampleCountEnd(self):
-    _, generated_x, ss_out = NextFrameTest.run_scheduled_sample_func(
+    _, generated_x, ss_out = NextFrameTest.runScheduledSampleFunc(
         next_frame.NextFrameStochastic.scheduled_sample_count, 0, 10)
     self.assertAllEqual(generated_x, ss_out)
 

From e5a495545024c839d47a339e48245b5f7c3a7b5e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 1 Aug 2018 11:03:59 -0700
Subject: [PATCH 0438/2720] Allow additional hooks to create_experiment().

PiperOrigin-RevId: 206956986
---
 tensor2tensor/utils/trainer_lib.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 94e818735..6ac4a5fe3 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -408,7 +408,9 @@ def create_experiment(
     eval_early_stopping_metric_minimize=True,
     autotune=False,
     use_tpu=False,
-    xla_compile=False):
+    xla_compile=False,
+    additional_train_hooks=None,
+    additional_eval_hooks=None):
   """Create Experiment."""
   # HParams
   hparams.add_hparam("model_dir", run_config.model_dir)
@@ -477,6 +479,10 @@ def create_experiment(
       early_stopping_kwargs=early_stopping_kwargs)
   train_hooks += t2t_model.T2TModel.get_train_hooks(model_name)
   eval_hooks += t2t_model.T2TModel.get_eval_hooks(model_name)
+  if additional_train_hooks:
+    train_hooks += additional_train_hooks
+  if additional_eval_hooks:
+    eval_hooks += additional_eval_hooks
 
   train_hooks = tf.contrib.learn.monitors.replace_monitors_with_hooks(
       train_hooks, estimator)

From 19ea73e22c1bb1d5f33d187e13112f4dc3b4d08f Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 1 Aug 2018 12:11:00 -0700
Subject: [PATCH 0439/2720] Decouple the SAVP model from the rest of
 next_frame.py

PiperOrigin-RevId: 206969462
---
 tensor2tensor/models/__init__.py              |   1 +
 tensor2tensor/models/research/next_frame.py   | 180 ---------------
 .../models/research/next_frame_params.py      |   2 +-
 .../models/research/next_frame_savp.py        | 209 ++++++++++++++++++
 .../models/research/next_frame_savp_test.py   |  88 ++++++++
 .../models/research/next_frame_test.py        |  16 --
 6 files changed, 299 insertions(+), 197 deletions(-)
 create mode 100644 tensor2tensor/models/research/next_frame_savp.py
 create mode 100644 tensor2tensor/models/research/next_frame_savp_test.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index abe3f1f22..e5f49e04a 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -44,6 +44,7 @@
 from tensor2tensor.models.research import lm_experiments
 from tensor2tensor.models.research import multimodel
 from tensor2tensor.models.research import next_frame
+from tensor2tensor.models.research import next_frame_savp
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research import similarity_transformer
 from tensor2tensor.models.research import super_lm
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index b5ef66b84..3af91c2a1 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -1282,183 +1282,3 @@ def construct_model(self, images, actions, rewards):
     tf.logging.info(">>>> Done")
     gen_images = tf.stack(gen_images, axis=0)
     return gen_images, fake_reward_prediction, pred_mu, pred_logvar
-
-
-@registry.register_model
-class NextFrameSavp(NextFrameStochastic):
-  """Stochastic Adversarial Video Prediction."""
-
-  def encoder(self, inputs, n_layers=3):
-    """COnvnet that encodes inputs into mean and std of a gaussian.
-
-    Args:
-     inputs: 5-D Tensor, shape (batch_size, num_frames, width, height, channels)
-     n_layers: Number of layers.
-
-    Returns:
-     z_mu: Mean of the latent gaussians.
-     z_log_var: log(var) of the latent gaussians.
-
-    Raises:
-      ValueError: If inputs is not a 5-D tensor or not float32.
-    """
-    latent_dims = self.hparams.z_dim
-
-    shape_as_list = inputs.shape.as_list()
-    if len(shape_as_list) != 5:
-      raise ValueError("Expected inputs to be a 5-D, got %d" %
-                       len(shape_as_list))
-    if inputs.dtype != tf.float32:
-      raise ValueError("Expected dtype tf.float32, got %s" % inputs.dtype)
-
-    # Flatten (N,T,W,H,C) into (NT,W,H,C)
-    batch_size, _ = shape_as_list[:2]
-    inputs = tf.reshape(inputs, [-1] + list(inputs.shape)[2:])
-    n_filters = 64
-    rectified = None
-
-    # Applies 3 layer conv-net with padding, instance normalization
-    # and leaky relu as per the encoder in
-    # https://github.com/alexlee-gk/video_prediction
-    padding = [[0, 0], [1, 1], [1, 1], [0, 0]]
-    for i in range(n_layers):
-      with tf.variable_scope("layer_%d" % (i + 1)):
-        n_filters *= 2**i
-        if i:
-          padded = tf.pad(rectified, padding)
-        else:
-          padded = tf.pad(inputs, padding)
-        convolved = tf.layers.conv2d(padded, filters=n_filters, kernel_size=4,
-                                     strides=2, padding="VALID")
-        normalized = tf.contrib.layers.instance_norm(convolved)
-        rectified = tf.nn.leaky_relu(normalized, alpha=0.2)
-
-    # Mean pooling across all spatial dimensions.
-    pooled = tf.nn.avg_pool(
-        rectified, [1] + rectified.shape[1:3].as_list() + [1],
-        strides=[1, 1, 1, 1], padding="VALID")
-    squeezed = tf.squeeze(pooled, [1, 2])
-
-    # Down-project and output the mean and log of the standard deviation of
-    # the latents.
-    with tf.variable_scope("z_mu"):
-      z_mu = tf.layers.dense(squeezed, latent_dims)
-    with tf.variable_scope("z_log_sigma_sq"):
-      z_log_var = tf.layers.dense(squeezed, latent_dims)
-      z_log_var = tf.clip_by_value(z_log_var, -10, 10)
-
-    # Reshape to (batch_size X num_frames X latent_dims)
-    z_mu = tf.reshape(z_mu, (batch_size, -1, latent_dims))
-    z_log_var = tf.reshape(
-        z_log_var, (batch_size, -1, latent_dims))
-    return z_mu, z_log_var
-
-  def construct_model(self, images, actions, rewards):
-    """Model that takes in images and returns predictions.
-
-    Args:
-      images: list of 4-D Tensors indexed by time.
-              (batch_size, width, height, channels)
-      actions: list of action tensors
-               each action should be in the shape ?x1xZ
-      rewards: list of reward tensors
-               each reward should be in the shape ?x1xZ
-
-    Returns:
-      video: list of 4-D predicted frames.
-      all_rewards: predicted rewards.
-      latent_means: list of gaussian means conditioned on the input at
-                    every frame.
-      latent_stds: list of gaussian stds conditioned on the input at
-                   every frame.
-    """
-    images = tf.unstack(images, axis=0)
-    actions = tf.unstack(actions, axis=0)
-    rewards = tf.unstack(rewards, axis=0)
-
-    latent_dims = self.hparams.z_dim
-    context_frames = self.hparams.video_num_input_frames
-    seq_len = len(images)
-    input_shape = common_layers.shape_list(images[0])
-    batch_size = input_shape[0]
-
-    # Model does not support reward-conditioned frame generation.
-    fake_rewards = rewards[:-1]
-
-    # Concatenate x_{t-1} and x_{t} along depth and encode it to
-    # produce the mean and standard deviation of z_{t-1}
-    image_pairs = tf.concat([images[:seq_len - 1],
-                             images[1:seq_len]], axis=-1)
-
-    z_mu, z_log_sigma_sq = self.encoder(image_pairs)
-    # Unstack z_mu and z_log_sigma_sq along the time dimension.
-    z_mu = tf.unstack(z_mu, axis=0)
-    z_log_sigma_sq = tf.unstack(z_log_sigma_sq, axis=0)
-    iterable = zip(images[:-1], actions[:-1], fake_rewards,
-                   z_mu, z_log_sigma_sq)
-
-    # Initialize LSTM State
-    lstm_state = [None] * 7
-    gen_cond_video, gen_prior_video, all_rewards, latent_means, latent_stds = \
-      [], [], [], [], []
-    pred_image, prior_latent_state, cond_latent_state = None, None, None
-    train_mode = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
-
-    # Create scheduled sampling function
-    ss_func = self.get_scheduled_sample_func(batch_size)
-
-    with tf.variable_scope("prediction", reuse=tf.AUTO_REUSE):
-
-      for step, (image, action, reward, mu, log_sigma_sq) in enumerate(iterable):  # pylint:disable=line-too-long
-        # Sample latents using a gaussian centered at conditional mu and std.
-        latent = self.get_gaussian_latent(mu, log_sigma_sq)
-
-        # Sample prior latents from isotropic normal distribution.
-        prior_latent = tf.random_normal(tf.shape(latent), dtype=tf.float32)
-
-        # LSTM that encodes correlations between conditional latents.
-        # Pg 22 in https://arxiv.org/pdf/1804.01523.pdf
-        enc_cond_latent, cond_latent_state = basic_lstm(
-            latent, cond_latent_state, latent_dims, name="cond_latent")
-
-        # LSTM that encodes correlations between prior latents.
-        enc_prior_latent, prior_latent_state = basic_lstm(
-            prior_latent, prior_latent_state, latent_dims, name="prior_latent")
-
-        # Scheduled Sampling
-        done_warm_start = step > context_frames - 1
-        groundtruth_items = [image]
-        generated_items = [pred_image]
-        input_image = self.get_scheduled_sample_inputs(
-            done_warm_start, groundtruth_items, generated_items, ss_func)[0]
-
-        all_latents = tf.concat([enc_cond_latent, enc_prior_latent], axis=0)
-        all_image = tf.concat([input_image, input_image], axis=0)
-        all_action = tf.concat([action, action], axis=0)
-        all_rewards = tf.concat([reward, reward], axis=0)
-
-        all_pred_images, lstm_state = self.construct_predictive_tower(
-            all_image, all_rewards, all_action, lstm_state, all_latents,
-            concat_latent=True)
-
-        cond_pred_images, prior_pred_images = \
-          all_pred_images[:batch_size], all_pred_images[batch_size:]
-
-        if train_mode:
-          pred_image = cond_pred_images
-        else:
-          pred_image = prior_pred_images
-
-        gen_cond_video.append(cond_pred_images)
-        gen_prior_video.append(prior_pred_images)
-        latent_means.append(mu)
-        latent_stds.append(log_sigma_sq)
-
-    gen_cond_video = tf.stack(gen_cond_video, axis=0)
-    gen_prior_video = tf.stack(gen_prior_video, axis=0)
-    fake_rewards = tf.stack(fake_rewards, axis=0)
-
-    if train_mode:
-      return gen_cond_video, fake_rewards, latent_means, latent_stds
-    else:
-      return gen_prior_video, fake_rewards, latent_means, latent_stds
diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index 44ecf5bec..94aea9b94 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -99,7 +99,7 @@ def next_frame_stochastic_emily():
 
 @registry.register_hparams
 def next_frame_savp():
-  """SVAP model."""
+  """SAVP model."""
   hparams = next_frame_stochastic()
   hparams.add_hparam("z_dim", 8)
   hparams.target_modality = "video:l1raw"
diff --git a/tensor2tensor/models/research/next_frame_savp.py b/tensor2tensor/models/research/next_frame_savp.py
new file mode 100644
index 000000000..73b65c599
--- /dev/null
+++ b/tensor2tensor/models/research/next_frame_savp.py
@@ -0,0 +1,209 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Stochastic Adversarial Video Prediction model.
+
+Reference: https://arxiv.org/abs/1804.01523
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_layers
+from tensor2tensor.models.research import next_frame
+from tensor2tensor.models.research import next_frame_params  # pylint: disable=unused-import
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+
+@registry.register_model
+class NextFrameSAVP(next_frame.NextFrameStochastic):
+  """Stochastic Adversarial Video Prediction."""
+
+  def encoder(self, inputs, n_layers=3):
+    """COnvnet that encodes inputs into mean and std of a gaussian.
+
+    Args:
+     inputs: 5-D Tensor, shape (batch_size, num_frames, width, height, channels)
+     n_layers: Number of layers.
+
+    Returns:
+     z_mu: Mean of the latent gaussians.
+     z_log_var: log(var) of the latent gaussians.
+
+    Raises:
+      ValueError: If inputs is not a 5-D tensor or not float32.
+    """
+    latent_dims = self.hparams.z_dim
+
+    shape_as_list = inputs.shape.as_list()
+    if len(shape_as_list) != 5:
+      raise ValueError("Expected inputs to be a 5-D, got %d" %
+                       len(shape_as_list))
+    if inputs.dtype != tf.float32:
+      raise ValueError("Expected dtype tf.float32, got %s" % inputs.dtype)
+
+    # Flatten (N,T,W,H,C) into (NT,W,H,C)
+    batch_size, _ = shape_as_list[:2]
+    inputs = tf.reshape(inputs, [-1] + list(inputs.shape)[2:])
+    n_filters = 64
+    rectified = None
+
+    # Applies 3 layer conv-net with padding, instance normalization
+    # and leaky relu as per the encoder in
+    # https://github.com/alexlee-gk/video_prediction
+    padding = [[0, 0], [1, 1], [1, 1], [0, 0]]
+    for i in range(n_layers):
+      with tf.variable_scope("layer_%d" % (i + 1)):
+        n_filters *= 2**i
+        if i:
+          padded = tf.pad(rectified, padding)
+        else:
+          padded = tf.pad(inputs, padding)
+        convolved = tf.layers.conv2d(padded, filters=n_filters, kernel_size=4,
+                                     strides=2, padding="VALID")
+        normalized = tf.contrib.layers.instance_norm(convolved)
+        rectified = tf.nn.leaky_relu(normalized, alpha=0.2)
+
+    # Mean pooling across all spatial dimensions.
+    pooled = tf.nn.avg_pool(
+        rectified, [1] + rectified.shape[1:3].as_list() + [1],
+        strides=[1, 1, 1, 1], padding="VALID")
+    squeezed = tf.squeeze(pooled, [1, 2])
+
+    # Down-project and output the mean and log of the standard deviation of
+    # the latents.
+    with tf.variable_scope("z_mu"):
+      z_mu = tf.layers.dense(squeezed, latent_dims)
+    with tf.variable_scope("z_log_sigma_sq"):
+      z_log_var = tf.layers.dense(squeezed, latent_dims)
+      z_log_var = tf.clip_by_value(z_log_var, -10, 10)
+
+    # Reshape to (batch_size X num_frames X latent_dims)
+    z_mu = tf.reshape(z_mu, (batch_size, -1, latent_dims))
+    z_log_var = tf.reshape(
+        z_log_var, (batch_size, -1, latent_dims))
+    return z_mu, z_log_var
+
+  def construct_model(self, images, actions, rewards):
+    """Model that takes in images and returns predictions.
+
+    Args:
+      images: list of 4-D Tensors indexed by time.
+              (batch_size, width, height, channels)
+      actions: list of action tensors
+               each action should be in the shape ?x1xZ
+      rewards: list of reward tensors
+               each reward should be in the shape ?x1xZ
+
+    Returns:
+      video: list of 4-D predicted frames.
+      all_rewards: predicted rewards.
+      latent_means: list of gaussian means conditioned on the input at
+                    every frame.
+      latent_stds: list of gaussian stds conditioned on the input at
+                   every frame.
+    """
+    images = tf.unstack(images, axis=0)
+    actions = tf.unstack(actions, axis=0)
+    rewards = tf.unstack(rewards, axis=0)
+
+    latent_dims = self.hparams.z_dim
+    context_frames = self.hparams.video_num_input_frames
+    seq_len = len(images)
+    input_shape = common_layers.shape_list(images[0])
+    batch_size = input_shape[0]
+
+    # Model does not support reward-conditioned frame generation.
+    fake_rewards = rewards[:-1]
+
+    # Concatenate x_{t-1} and x_{t} along depth and encode it to
+    # produce the mean and standard deviation of z_{t-1}
+    image_pairs = tf.concat([images[:seq_len - 1],
+                             images[1:seq_len]], axis=-1)
+
+    z_mu, z_log_sigma_sq = self.encoder(image_pairs)
+    # Unstack z_mu and z_log_sigma_sq along the time dimension.
+    z_mu = tf.unstack(z_mu, axis=0)
+    z_log_sigma_sq = tf.unstack(z_log_sigma_sq, axis=0)
+    iterable = zip(images[:-1], actions[:-1], fake_rewards,
+                   z_mu, z_log_sigma_sq)
+
+    # Initialize LSTM State
+    lstm_state = [None] * 7
+    gen_cond_video, gen_prior_video, all_rewards, latent_means, latent_stds = \
+      [], [], [], [], []
+    pred_image = tf.zeros_like(images[0])
+    prior_latent_state, cond_latent_state = None, None
+    train_mode = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
+
+    # Create scheduled sampling function
+    ss_func = self.get_scheduled_sample_func(batch_size)
+
+    with tf.variable_scope("prediction", reuse=tf.AUTO_REUSE):
+
+      for step, (image, action, reward, mu, log_sigma_sq) in enumerate(iterable):  # pylint:disable=line-too-long
+        # Sample latents using a gaussian centered at conditional mu and std.
+        latent = self.get_gaussian_latent(mu, log_sigma_sq)
+
+        # Sample prior latents from isotropic normal distribution.
+        prior_latent = tf.random_normal(tf.shape(latent), dtype=tf.float32)
+
+        # LSTM that encodes correlations between conditional latents.
+        # Pg 22 in https://arxiv.org/pdf/1804.01523.pdf
+        enc_cond_latent, cond_latent_state = next_frame.basic_lstm(
+            latent, cond_latent_state, latent_dims, name="cond_latent")
+
+        # LSTM that encodes correlations between prior latents.
+        enc_prior_latent, prior_latent_state = next_frame.basic_lstm(
+            prior_latent, prior_latent_state, latent_dims, name="prior_latent")
+
+        # Scheduled Sampling
+        done_warm_start = step > context_frames - 1
+        groundtruth_items = [image]
+        generated_items = [pred_image]
+        input_image = self.get_scheduled_sample_inputs(
+            done_warm_start, groundtruth_items, generated_items, ss_func)
+
+        all_latents = tf.concat([enc_cond_latent, enc_prior_latent], axis=0)
+        all_image = tf.concat([input_image, input_image], axis=0)
+        all_action = tf.concat([action, action], axis=0)
+        all_rewards = tf.concat([reward, reward], axis=0)
+
+        all_pred_images, lstm_state = self.construct_predictive_tower(
+            all_image, all_rewards, all_action, lstm_state, all_latents,
+            concat_latent=True)
+
+        cond_pred_images, prior_pred_images = \
+          all_pred_images[:batch_size], all_pred_images[batch_size:]
+
+        if train_mode:
+          pred_image = cond_pred_images
+        else:
+          pred_image = prior_pred_images
+
+        gen_cond_video.append(cond_pred_images)
+        gen_prior_video.append(prior_pred_images)
+        latent_means.append(mu)
+        latent_stds.append(log_sigma_sq)
+
+    gen_cond_video = tf.stack(gen_cond_video, axis=0)
+    gen_prior_video = tf.stack(gen_prior_video, axis=0)
+    fake_rewards = tf.stack(fake_rewards, axis=0)
+
+    if train_mode:
+      return gen_cond_video, fake_rewards, latent_means, latent_stds
+    else:
+      return gen_prior_video, fake_rewards, latent_means, latent_stds
diff --git a/tensor2tensor/models/research/next_frame_savp_test.py b/tensor2tensor/models/research/next_frame_savp_test.py
new file mode 100644
index 000000000..9facd903d
--- /dev/null
+++ b/tensor2tensor/models/research/next_frame_savp_test.py
@@ -0,0 +1,88 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for the SAVP model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import numpy as np
+
+from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
+from tensor2tensor.models.research import next_frame_params
+from tensor2tensor.models.research import next_frame_savp
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+class NextFrameSAVPTest(tf.test.TestCase):
+
+  def TestVideoModel(self,
+                     in_frames,
+                     out_frames,
+                     hparams,
+                     model,
+                     expected_last_dim,
+                     upsample_method="conv2d_transpose"):
+
+    x = np.random.random_integers(0, high=255, size=(8, in_frames, 64, 64, 3))
+    y = np.random.random_integers(0, high=255, size=(8, out_frames, 64, 64, 3))
+
+    hparams.video_num_input_frames = in_frames
+    hparams.video_num_target_frames = out_frames
+    hparams.upsample_method = upsample_method
+    problem = registry.problem("video_stochastic_shapes10k")
+    p_hparams = problem.get_hparams(hparams)
+    hparams.problem = problem
+    hparams.problem_hparams = p_hparams
+
+    with self.test_session() as session:
+      features = {
+          "inputs": tf.constant(x, dtype=tf.int32),
+          "targets": tf.constant(y, dtype=tf.int32),
+      }
+      model = model(
+          hparams, tf.estimator.ModeKeys.TRAIN)
+      logits, _ = model(features)
+      session.run(tf.global_variables_initializer())
+      res = session.run(logits)
+    expected_shape = y.shape + (expected_last_dim,)
+    self.assertEqual(res.shape, expected_shape)
+
+  def TestOnVariousInputOutputSizes(self, hparams, model, expected_last_dim):
+    self.TestVideoModel(1, 1, hparams, model, expected_last_dim)
+    self.TestVideoModel(1, 6, hparams, model, expected_last_dim)
+    self.TestVideoModel(4, 1, hparams, model, expected_last_dim)
+    self.TestVideoModel(7, 5, hparams, model, expected_last_dim)
+
+  def TestOnVariousUpSampleLayers(self, hparams, model, expected_last_dim):
+    self.TestVideoModel(4, 1, hparams, model, expected_last_dim,
+                        upsample_method="bilinear_upsample_conv")
+    self.TestVideoModel(4, 1, hparams, model, expected_last_dim,
+                        upsample_method="nn_upsample_conv")
+
+  def testStochasticSavp(self):
+    self.TestOnVariousInputOutputSizes(
+        next_frame_params.next_frame_savp(),
+        next_frame_savp.NextFrameSAVP,
+        1)
+    self.TestOnVariousUpSampleLayers(
+        next_frame_params.next_frame_savp(),
+        next_frame_savp.NextFrameSAVP,
+        1)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index e4d2f4178..fd9d668d6 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -67,12 +67,6 @@ def TestOnVariousInputOutputSizes(self, hparams, model, expected_last_dim):
     self.TestVideoModel(4, 1, hparams, model, expected_last_dim)
     self.TestVideoModel(7, 5, hparams, model, expected_last_dim)
 
-  def TestOnVariousUpSampleLayers(self, hparams, model, expected_last_dim):
-    self.TestVideoModel(4, 1, hparams, model, expected_last_dim,
-                        upsample_method="bilinear_upsample_conv")
-    self.TestVideoModel(4, 1, hparams, model, expected_last_dim,
-                        upsample_method="nn_upsample_conv")
-
   def testBasic(self):
     self.TestOnVariousInputOutputSizes(
         next_frame_params.next_frame(),
@@ -97,16 +91,6 @@ def testStochasticEmily(self):
         next_frame.NextFrameStochasticEmily,
         1)
 
-  def testStochasticSavp(self):
-    self.TestOnVariousInputOutputSizes(
-        next_frame_params.next_frame_savp(),
-        next_frame.NextFrameSavp,
-        1)
-    self.TestOnVariousUpSampleLayers(
-        next_frame_params.next_frame_savp(),
-        next_frame.NextFrameSavp,
-        1)
-
   @staticmethod
   def runScheduledSampleFunc(func, var, batch_size):
     ground_truth_x = list(range(1, batch_size+1))

From 123c3267806fd277506ab8acbf899ad3a96c1f9b Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 1 Aug 2018 16:59:01 -0700
Subject: [PATCH 0440/2720] Cutting next_frame fat!

PiperOrigin-RevId: 207019530
---
 tensor2tensor/layers/common_layers.py         |  17 +
 tensor2tensor/layers/common_video.py          | 263 ++++++++++++++
 tensor2tensor/layers/common_video_test.py     |  97 +++++
 tensor2tensor/models/research/next_frame.py   | 336 +++---------------
 .../models/research/next_frame_params.py      |   2 +-
 .../models/research/next_frame_savp.py        |   5 +-
 .../models/research/next_frame_test.py        |  72 ----
 7 files changed, 428 insertions(+), 364 deletions(-)
 create mode 100644 tensor2tensor/layers/common_video.py
 create mode 100644 tensor2tensor/layers/common_video_test.py

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 29840c6a0..1ee456f96 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3692,3 +3692,20 @@ def targeted_dropout(inputs,
     return inputs * (1 - mask)
   else:
     return inputs
+
+
+# TODO(mbz): use tf.distributions.kl_divergence instead.
+def kl_divergence(mu, log_sigma):
+  """KL divergence of diagonal gaussian N(mu,exp(log_sigma)) and N(0,1).
+
+  Args:
+    mu: mu parameter of the distribution.
+    log_sigma: log(sigma) parameter of the distribution.
+  Returns:
+    the KL loss.
+  """
+  batch_size = shape_list(mu)[0]
+  kl = -.5 * tf.reduce_sum(1. + log_sigma - tf.square(mu) - tf.exp(log_sigma))
+  return kl / tf.to_float(batch_size)
+
+
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
new file mode 100644
index 000000000..6a15401fa
--- /dev/null
+++ b/tensor2tensor/layers/common_video.py
@@ -0,0 +1,263 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Layers common to multiple models."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_layers
+import tensorflow as tf
+
+tfl = tf.layers
+
+
+def encode_to_shape(inputs, shape, scope):
+  """Encode the given tensor to given image shape."""
+  with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
+    w, h = shape[1].value, shape[2].value
+    x = inputs
+    x = tf.contrib.layers.flatten(x)
+    x = tfl.dense(x, w * h, activation=tf.nn.relu, name="enc_dense")
+    x = tf.reshape(x, (-1, w, h, 1))
+    return x
+
+
+def decode_to_shape(inputs, shape, scope):
+  """Encode the given tensor to given image shape."""
+  with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
+    x = inputs
+    x = tf.contrib.layers.flatten(x)
+    x = tfl.dense(x, shape[2].value, activation=tf.nn.relu, name="dec_dense")
+    x = tf.expand_dims(x, axis=1)
+    return x
+
+
+def basic_lstm(inputs, state, num_units, name=None):
+  """Basic LSTM."""
+  input_shape = common_layers.shape_list(inputs)
+  cell = tf.contrib.rnn.BasicLSTMCell(num_units, name=name)
+  if state is None:
+    state = cell.zero_state(input_shape[0], tf.float32)
+  outputs, new_state = cell(inputs, state)
+  return outputs, new_state
+
+
+def conv_lstm_2d(inputs, state, output_channels,
+                 kernel_size=5, name=None, spatial_dims=None):
+  """2D Convolutional LSTM."""
+  input_shape = common_layers.shape_list(inputs)
+  batch_size, input_channels = input_shape[0], input_shape[-1]
+  if spatial_dims is None:
+    input_shape = input_shape[1:]
+  else:
+    input_shape = spatial_dims + [input_channels]
+
+  cell = tf.contrib.rnn.ConvLSTMCell(
+      2, input_shape, output_channels,
+      [kernel_size, kernel_size], name=name)
+  if state is None:
+    state = cell.zero_state(batch_size, tf.float32)
+  outputs, new_state = cell(inputs, state)
+  return outputs, new_state
+
+
+def scheduled_sample_count(ground_truth_x,
+                           generated_x,
+                           batch_size,
+                           scheduled_sample_var):
+  """Sample batch with specified mix of groundtruth and generated data points.
+
+  Args:
+    ground_truth_x: tensor of ground-truth data points.
+    generated_x: tensor of generated data points.
+    batch_size: batch size
+    scheduled_sample_var: number of ground-truth examples to include in batch.
+  Returns:
+    New batch with num_ground_truth sampled from ground_truth_x and the rest
+    from generated_x.
+  """
+  num_ground_truth = scheduled_sample_var
+  idx = tf.random_shuffle(tf.range(batch_size))
+  ground_truth_idx = tf.gather(idx, tf.range(num_ground_truth))
+  generated_idx = tf.gather(idx, tf.range(num_ground_truth, batch_size))
+
+  ground_truth_examps = tf.gather(ground_truth_x, ground_truth_idx)
+  generated_examps = tf.gather(generated_x, generated_idx)
+  return tf.dynamic_stitch([ground_truth_idx, generated_idx],
+                           [ground_truth_examps, generated_examps])
+
+
+def scheduled_sample_prob(ground_truth_x,
+                          generated_x,
+                          batch_size,
+                          scheduled_sample_var):
+  """Probability based scheduled sampling.
+
+  Args:
+    ground_truth_x: tensor of ground-truth data points.
+    generated_x: tensor of generated data points.
+    batch_size: batch size
+    scheduled_sample_var: probability of choosing from ground_truth.
+  Returns:
+    New batch with randomly selected data points.
+  """
+  probability_threshold = scheduled_sample_var
+  probability_of_generated = tf.random_uniform([batch_size])
+  array_ind = tf.to_int32(probability_of_generated > probability_threshold)
+  indices = tf.range(batch_size) + array_ind * batch_size
+  xy = tf.concat([ground_truth_x, generated_x], axis=0)
+  output = tf.gather(xy, indices)
+  return output
+
+
+def dna_transformation(prev_image, dna_input, dna_kernel_size, relu_shift):
+  """Apply dynamic neural advection to previous image.
+
+  Args:
+    prev_image: previous image to be transformed.
+    dna_input: hidden lyaer to be used for computing DNA transformation.
+    dna_kernel_size: dna kernel size.
+    relu_shift: shift for ReLU function.
+  Returns:
+    List of images transformed by the predicted CDNA kernels.
+  """
+  # Construct translated images.
+  prev_image_pad = tf.pad(prev_image, [[0, 0], [2, 2], [2, 2], [0, 0]])
+  image_height = int(prev_image.get_shape()[1])
+  image_width = int(prev_image.get_shape()[2])
+
+  inputs = []
+  for xkern in range(dna_kernel_size):
+    for ykern in range(dna_kernel_size):
+      inputs.append(
+          tf.expand_dims(
+              tf.slice(prev_image_pad, [0, xkern, ykern, 0],
+                       [-1, image_height, image_width, -1]), [3]))
+  inputs = tf.concat(axis=3, values=inputs)
+
+  # Normalize channels to 1.
+  kernel = tf.nn.relu(dna_input - relu_shift) + relu_shift
+  kernel = tf.expand_dims(
+      kernel / tf.reduce_sum(kernel, [3], keep_dims=True), [4])
+  return tf.reduce_sum(kernel * inputs, [3], keep_dims=False)
+
+
+def cdna_transformation(prev_image, cdna_input, num_masks, color_channels,
+                        dna_kernel_size, relu_shift):
+  """Apply convolutional dynamic neural advection to previous image.
+
+  Args:
+    prev_image: previous image to be transformed.
+    cdna_input: hidden lyaer to be used for computing CDNA kernels.
+    num_masks: number of masks and hence the number of CDNA transformations.
+    color_channels: the number of color channels in the images.
+    dna_kernel_size: dna kernel size.
+    relu_shift: shift for ReLU function.
+  Returns:
+    List of images transformed by the predicted CDNA kernels.
+  """
+  batch_size = tf.shape(cdna_input)[0]
+  height = int(prev_image.get_shape()[1])
+  width = int(prev_image.get_shape()[2])
+
+  # Predict kernels using linear function of last hidden layer.
+  cdna_kerns = tfl.dense(
+      cdna_input, dna_kernel_size * dna_kernel_size * num_masks,
+      name="cdna_params",
+      activation=None)
+
+  # Reshape and normalize.
+  cdna_kerns = tf.reshape(
+      cdna_kerns, [batch_size, dna_kernel_size, dna_kernel_size, 1, num_masks])
+  cdna_kerns = (tf.nn.relu(cdna_kerns - relu_shift) + relu_shift)
+  norm_factor = tf.reduce_sum(cdna_kerns, [1, 2, 3], keep_dims=True)
+  cdna_kerns /= norm_factor
+
+  # Treat the color channel dimension as the batch dimension since the same
+  # transformation is applied to each color channel.
+  # Treat the batch dimension as the channel dimension so that
+  # depthwise_conv2d can apply a different transformation to each sample.
+  cdna_kerns = tf.transpose(cdna_kerns, [1, 2, 0, 4, 3])
+  cdna_kerns = tf.reshape(
+      cdna_kerns, [dna_kernel_size, dna_kernel_size, batch_size, num_masks])
+  # Swap the batch and channel dimensions.
+  prev_image = tf.transpose(prev_image, [3, 1, 2, 0])
+
+  # Transform image.
+  transformed = tf.nn.depthwise_conv2d(
+      prev_image, cdna_kerns, [1, 1, 1, 1], "SAME")
+
+  # Transpose the dimensions to where they belong.
+  transformed = tf.reshape(
+      transformed, [color_channels, height, width, batch_size, num_masks])
+  transformed = tf.transpose(transformed, [3, 1, 2, 0, 4])
+  transformed = tf.unstack(transformed, axis=-1)
+  return transformed
+
+
+def vgg_layer(inputs,
+              nout,
+              kernel_size=3,
+              activation=tf.nn.leaky_relu,
+              padding="SAME",
+              is_training=False,
+              scope=None):
+  """A layer of VGG network with batch norm.
+
+  Args:
+    inputs: image tensor
+    nout: number of output channels
+    kernel_size: size of the kernel
+    activation: activation function
+    padding: padding of the image
+    is_training: whether it is training mode or not
+    scope: variable scope of the op
+  Returns:
+    net: output of layer
+  """
+  with tf.variable_scope(scope):
+    net = tfl.conv2d(inputs, nout, kernel_size=kernel_size, padding=padding,
+                     activation=None, name="conv")
+    net = tfl.batch_normalization(net, training=is_training, name="bn")
+    net = activation(net)
+  return net
+
+
+def tile_and_concat(image, latent, concat_latent=True):
+  """Tile latent and concatenate to image across depth.
+
+  Args:
+    image: 4-D Tensor, (batch_size X height X width X channels)
+    latent: 2-D Tensor, (batch_size X latent_dims)
+    concat_latent: If set to False, the image is returned as is.
+
+  Returns:
+    concat_latent: 4-D Tensor, (batch_size X height X width X channels+1)
+      latent tiled and concatenated to the image across the channels.
+  """
+  if not concat_latent:
+    return image
+  image_shape = common_layers.shape_list(image)
+  latent_shape = common_layers.shape_list(latent)
+  height, width = image_shape[1], image_shape[2]
+  latent_dims = latent_shape[1]
+
+  height_multiples = height // latent_dims
+  pad = height - (height_multiples * latent_dims)
+  latent = tf.reshape(latent, (-1, latent_dims, 1, 1))
+  latent = tf.tile(latent, (1, height_multiples, width, 1))
+  latent = tf.pad(latent, [[0, 0], [pad // 2, pad // 2], [0, 0], [0, 0]])
+  return tf.concat([image, latent], axis=-1)
+
diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
new file mode 100644
index 000000000..4a4f939b5
--- /dev/null
+++ b/tensor2tensor/layers/common_video_test.py
@@ -0,0 +1,97 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import numpy as np
+
+from tensor2tensor.layers import common_video
+
+import tensorflow as tf
+
+
+class CommonVideoTest(tf.test.TestCase):
+
+  @staticmethod
+  def runScheduledSampleFunc(func, var, batch_size):
+    ground_truth_x = list(range(1, batch_size+1))
+    generated_x = [-x for x in ground_truth_x]
+    ground_truth_x = tf.convert_to_tensor(ground_truth_x)
+    generated_x = tf.convert_to_tensor(generated_x)
+    ss_out = func(ground_truth_x, generated_x, batch_size, var)
+    with tf.Session() as session:
+      output = session.run([ground_truth_x, generated_x, ss_out])
+    return output
+
+  def testScheduledSampleProbStart(self):
+    ground_truth_x, _, ss_out = CommonVideoTest.runScheduledSampleFunc(
+        common_video.NextFrameStochastic.scheduled_sample_prob, 1.0, 10)
+    self.assertAllEqual(ground_truth_x, ss_out)
+
+  def testScheduledSampleProbMid(self):
+    _, _, ss_out = CommonVideoTest.runScheduledSampleFunc(
+        common_video.NextFrameStochastic.scheduled_sample_prob, 0.5, 1000)
+    positive_count = np.sum(ss_out > 0)
+    self.assertAlmostEqual(positive_count / 1000.0, 0.5, places=2)
+
+  def testScheduledSampleProbEnd(self):
+    _, generated_x, ss_out = CommonVideoTest.runScheduledSampleFunc(
+        common_video.NextFrameStochastic.scheduled_sample_prob, 0.0, 10)
+    self.assertAllEqual(generated_x, ss_out)
+
+  def testScheduledSampleCountStart(self):
+    ground_truth_x, _, ss_out = CommonVideoTest.runScheduledSampleFunc(
+        common_video.NextFrameStochastic.scheduled_sample_count, 10, 10)
+    self.assertAllEqual(ground_truth_x, ss_out)
+
+  def testScheduledSampleCountMid(self):
+    _, _, ss_out = CommonVideoTest.runScheduledSampleFunc(
+        common_video.NextFrameStochastic.scheduled_sample_count, 5, 10)
+    positive_count = np.sum(ss_out > 0)
+    self.assertEqual(positive_count, 5)
+
+  def testScheduledSampleCountEnd(self):
+    _, generated_x, ss_out = CommonVideoTest.runScheduledSampleFunc(
+        common_video.NextFrameStochastic.scheduled_sample_count, 0, 10)
+    self.assertAllEqual(generated_x, ss_out)
+
+  def testDynamicTileAndConcat(self):
+    with tf.Graph().as_default():
+      # image = (1 X 4 X 4 X 1)
+      image = [[1, 2, 3, 4],
+               [2, 4, 5, 6],
+               [7, 8, 9, 10],
+               [7, 9, 10, 1]]
+      image_t = tf.expand_dims(tf.expand_dims(image, axis=0), axis=-1)
+      image_t = tf.cast(image_t, dtype=tf.float32)
+
+      # latent = (1 X 2)
+      latent = np.array([[90, 100]])
+      latent_t = tf.cast(tf.convert_to_tensor(latent), dtype=tf.float32)
+
+      with tf.Session() as session:
+        tiled = common_video.tile_and_concat(
+            image_t, latent_t)
+        tiled_np, image_np = session.run([tiled, image_t])
+        tiled_latent = tiled_np[0, :, :, -1]
+        self.assertAllEqual(tiled_np.shape, (1, 4, 4, 2))
+
+        self.assertAllEqual(tiled_np[:, :, :, :1], image_np)
+        self.assertAllEqual(
+            tiled_latent,
+            [[90, 90, 90, 90],
+             [100, 100, 100, 100],
+             [90, 90, 90, 90],
+             [100, 100, 100, 100]])
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 3af91c2a1..28b338423 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -23,6 +23,7 @@
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import next_frame_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
@@ -33,15 +34,6 @@
 tfcl = tf.contrib.layers
 
 
-def basic_lstm(inputs, state, num_units, name=None):
-  input_shape = common_layers.shape_list(inputs)
-  cell = tf.contrib.rnn.BasicLSTMCell(num_units, name=name)
-  if state is None:
-    state = cell.zero_state(input_shape[0], tf.float32)
-  outputs, new_state = cell(inputs, state)
-  return outputs, new_state
-
-
 @registry.register_model
 class NextFrameBasic(t2t_model.T2TModel):
   """Basic next-frame model, may take actions and predict rewards too."""
@@ -193,33 +185,6 @@ def tinyify(self, array):
       return [1 for _ in array]
     return array
 
-  @staticmethod
-  def tile_and_concat(image, latent, concat_latent=True):
-    """Tile latent and concatenate to image across depth.
-
-    Args:
-      image: 4-D Tensor, (batch_size X height X width X channels)
-      latent: 2-D Tensor, (batch_size X latent_dims)
-      concat_latent: If set to False, the image is returned as is.
-
-    Returns:
-      concat_latent: 4-D Tensor, (batch_size X height X width X channels+1)
-        latent tiled and concatenated to the image across the channels.
-    """
-    if not concat_latent:
-      return image
-    image_shape = common_layers.shape_list(image)
-    latent_shape = common_layers.shape_list(latent)
-    height, width = image_shape[1], image_shape[2]
-    latent_dims = latent_shape[1]
-
-    height_multiples = height // latent_dims
-    pad = height - (height_multiples * latent_dims)
-    latent = tf.reshape(latent, (-1, latent_dims, 1, 1))
-    latent = tf.tile(latent, (1, height_multiples, width, 1))
-    latent = tf.pad(latent, [[0, 0], [pad // 2, pad // 2], [0, 0], [0, 0]])
-    return tf.concat([image, latent], axis=-1)
-
   def construct_latent_tower(self, images):
     """Builds convolutional latent tower for stochastic model.
 
@@ -302,10 +267,11 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
       - the output of the partial network.
       - intermidate outputs for skip connections.
     """
-    lstm_func = self.conv_lstm_2d
+    lstm_func = common_video.conv_lstm_2d
+    tile_and_concat = common_video.tile_and_concat
 
     input_image = common_layers.make_even_size(input_image)
-    concat_input_image = self.tile_and_concat(
+    concat_input_image = tile_and_concat(
         input_image, latent, concat_latent=concat_latent)
 
     enc0 = tfl.conv2d(
@@ -319,7 +285,7 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
 
     hidden1, lstm_state[0] = lstm_func(
         enc0, lstm_state[0], lstm_size[0], name="state1")
-    hidden1 = self.tile_and_concat(hidden1, latent, concat_latent=concat_latent)
+    hidden1 = tile_and_concat(hidden1, latent, concat_latent=concat_latent)
     hidden1 = tfcl.layer_norm(hidden1, scope="layer_norm2")
     hidden2, lstm_state[1] = lstm_func(
         hidden1, lstm_state[1], lstm_size[1], name="state2")
@@ -327,23 +293,24 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
     hidden2 = common_layers.make_even_size(hidden2)
     enc1 = tfl.conv2d(hidden2, hidden2.get_shape()[3], [3, 3], strides=(2, 2),
                       padding="SAME", activation=tf.nn.relu, name="conv2")
-    enc1 = self.tile_and_concat(enc1, latent, concat_latent=concat_latent)
+    enc1 = tile_and_concat(enc1, latent, concat_latent=concat_latent)
 
     hidden3, lstm_state[2] = lstm_func(
         enc1, lstm_state[2], lstm_size[2], name="state3")
-    hidden3 = self.tile_and_concat(hidden3, latent, concat_latent=concat_latent)
+    hidden3 = tile_and_concat(hidden3, latent, concat_latent=concat_latent)
     hidden3 = tfcl.layer_norm(hidden3, scope="layer_norm4")
     hidden4, lstm_state[3] = lstm_func(
         hidden3, lstm_state[3], lstm_size[3], name="state4")
-    hidden4 = self.tile_and_concat(hidden4, latent, concat_latent=concat_latent)
+    hidden4 = tile_and_concat(hidden4, latent, concat_latent=concat_latent)
     hidden4 = tfcl.layer_norm(hidden4, scope="layer_norm5")
     hidden4 = common_layers.make_even_size(hidden4)
     enc2 = tfl.conv2d(hidden4, hidden4.get_shape()[3], [3, 3], strides=(2, 2),
                       padding="SAME", activation=tf.nn.relu, name="conv3")
 
     # Pass in reward and action.
-    emb_action = self.encode_to_shape(action, enc2.get_shape(), "action_enc")
-    emb_reward = self.encode_to_shape(
+    emb_action = common_video.encode_to_shape(
+        action, enc2.get_shape(), "action_enc")
+    emb_reward = common_video.encode_to_shape(
         input_reward, enc2.get_shape(), "reward_enc")
     enc2 = tf.concat(axis=3, values=[enc2, emb_action, emb_reward])
 
@@ -357,7 +324,7 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
     hidden5, lstm_state[4] = lstm_func(
         enc3, lstm_state[4], lstm_size[4], name="state5")  # last 8x8
     hidden5 = tfcl.layer_norm(hidden5, scope="layer_norm6")
-    hidden5 = self.tile_and_concat(hidden5, latent, concat_latent=concat_latent)
+    hidden5 = tile_and_concat(hidden5, latent, concat_latent=concat_latent)
     return hidden5, (enc0, enc1)
 
   def reward_prediction(
@@ -385,57 +352,22 @@ def reward_prediction(
       x = tfl.conv2d(x, conv_size[3], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="reward_conv3")
 
-      pred_reward = self.decode_to_shape(
+      pred_reward = common_video.decode_to_shape(
           x, input_reward.shape, "reward_dec")
 
       return pred_reward, lstm_state
 
-  def encode_to_shape(self, inputs, shape, scope):
-    """Encode the given tensor to given image shape."""
-    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
-      w, h = shape[1].value, shape[2].value
-      x = inputs
-      x = tfcl.flatten(x)
-      x = tfl.dense(x, w * h, activation=tf.nn.relu, name="enc_dense")
-      x = tf.reshape(x, (-1, w, h, 1))
-      return x
-
-  def decode_to_shape(self, inputs, shape, scope):
-    """Encode the given tensor to given image shape."""
-    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
-      x = inputs
-      x = tfcl.flatten(x)
-      x = tfl.dense(x, shape[2].value, activation=tf.nn.relu, name="dec_dense")
-      x = tf.expand_dims(x, axis=1)
-      return x
-
-  def conv_lstm_2d(self, inputs, state, output_channels,
-                   kernel_size=5, name=None, spatial_dims=None):
-    input_shape = common_layers.shape_list(inputs)
-    batch_size, input_channels = input_shape[0], input_shape[-1]
-    if spatial_dims is None:
-      input_shape = input_shape[1:]
-    else:
-      input_shape = spatial_dims + [input_channels]
-
-    cell = tf.contrib.rnn.ConvLSTMCell(
-        2, input_shape, output_channels,
-        [kernel_size, kernel_size], name=name)
-    if state is None:
-      state = cell.zero_state(batch_size, tf.float32)
-    outputs, new_state = cell(inputs, state)
-    return outputs, new_state
-
   def construct_predictive_tower(
       self, input_image, input_reward, action, lstm_state, latent,
       concat_latent=False):
     # Main tower
-    lstm_func = self.conv_lstm_2d
+    lstm_func = common_video.conv_lstm_2d
     batch_size = common_layers.shape_list(input_image)[0]
     # the number of different pixel motion predictions
     # and the number of masks for each of those predictions
     num_masks = self.hparams.num_masks
     upsample_method = self.hparams.upsample_method
+    tile_and_concat = common_video.tile_and_concat
 
     lstm_size = self.tinyify([32, 32, 64, 64, 128, 64, 32])
     conv_size = self.tinyify([32])
@@ -455,13 +387,12 @@ def construct_predictive_tower(
 
       enc1_shape = common_layers.shape_list(enc1)
       enc4 = enc4[:, :enc1_shape[1], :enc1_shape[2], :]  # Cut to shape.
-      enc4 = self.tile_and_concat(enc4, latent, concat_latent=concat_latent)
+      enc4 = tile_and_concat(enc4, latent, concat_latent=concat_latent)
 
       hidden6, lstm_state[5] = lstm_func(
           enc4, lstm_state[5], lstm_size[5], name="state6",
           spatial_dims=enc1_shape[1:-1])  # 16x16
-      hidden6 = self.tile_and_concat(
-          hidden6, latent, concat_latent=concat_latent)
+      hidden6 = tile_and_concat(hidden6, latent, concat_latent=concat_latent)
       hidden6 = tfcl.layer_norm(hidden6, scope="layer_norm7")
       # Skip connection.
       hidden6 = tf.concat(axis=3, values=[hidden6, enc1])  # both 16x16
@@ -473,7 +404,7 @@ def construct_predictive_tower(
 
       enc0_shape = common_layers.shape_list(enc0)
       enc5 = enc5[:, :enc0_shape[1], :enc0_shape[2], :]  # Cut to shape.
-      enc5 = self.tile_and_concat(enc5, latent, concat_latent=concat_latent)
+      enc5 = tile_and_concat(enc5, latent, concat_latent=concat_latent)
 
       hidden7, lstm_state[6] = lstm_func(
           enc5, lstm_state[6], lstm_size[6], name="state7",
@@ -488,7 +419,7 @@ def construct_predictive_tower(
             hidden7, num_outputs=hidden7.shape.as_list()[-1],
             stride=[2, 2], method=upsample_method)
       enc6 = tfcl.layer_norm(enc6, scope="layer_norm9")
-      enc6 = self.tile_and_concat(enc6, latent, concat_latent=concat_latent)
+      enc6 = tile_and_concat(enc6, latent, concat_latent=concat_latent)
 
       if self.hparams.model_options == "DNA":
         # Using largest hidden state for predicting untied conv kernels.
@@ -517,13 +448,17 @@ def construct_predictive_tower(
       if self.hparams.model_options == "CDNA":
         # cdna_input = tf.reshape(hidden5, [int(batch_size), -1])
         cdna_input = tfcl.flatten(hidden5)
-        transformed += self.cdna_transformation(
-            input_image, cdna_input, num_masks, int(color_channels))
+        transformed += common_video.cdna_transformation(
+            input_image, cdna_input, num_masks, int(color_channels),
+            self.hparams.dna_kernel_size, self.hparams.relu_shift)
       elif self.hparams.model_options == "DNA":
         # Only one mask is supported (more should be unnecessary).
         if num_masks != 1:
           raise ValueError("Only one mask is supported for DNA model.")
-        transformed = [self.dna_transformation(input_image, enc7)]
+        transformed = [
+            common_video.dna_transformation(
+                input_image, enc7,
+                self.hparams.dna_kernel_size, self.hparams.relu_shift)]
 
       masks = tfl.conv2d(
           enc6, filters=num_masks + 1, kernel_size=[1, 1],
@@ -629,147 +564,6 @@ def process_single_frame(prev_outputs, inputs):
 
     return gen_images, gen_rewards, [latent_mean], [latent_std]
 
-  def cdna_transformation(self,
-                          prev_image,
-                          cdna_input,
-                          num_masks,
-                          color_channels):
-    """Apply convolutional dynamic neural advection to previous image.
-
-    Args:
-      prev_image: previous image to be transformed.
-      cdna_input: hidden lyaer to be used for computing CDNA kernels.
-      num_masks: number of masks and hence the number of CDNA transformations.
-      color_channels: the number of color channels in the images.
-    Returns:
-      List of images transformed by the predicted CDNA kernels.
-    """
-    batch_size = tf.shape(cdna_input)[0]
-    height = int(prev_image.get_shape()[1])
-    width = int(prev_image.get_shape()[2])
-
-    # Predict kernels using linear function of last hidden layer.
-    cdna_kerns = tfl.dense(
-        cdna_input,
-        self.hparams.dna_kernel_size *
-        self.hparams.dna_kernel_size * num_masks,
-        name="cdna_params",
-        activation=None)
-
-    # Reshape and normalize.
-    cdna_kerns = tf.reshape(
-        cdna_kerns, [batch_size, self.hparams.dna_kernel_size,
-                     self.hparams.dna_kernel_size, 1, num_masks])
-    cdna_kerns = (tf.nn.relu(cdna_kerns - self.hparams.relu_shift)
-                  + self.hparams.relu_shift)
-    norm_factor = tf.reduce_sum(cdna_kerns, [1, 2, 3], keep_dims=True)
-    cdna_kerns /= norm_factor
-
-    # Treat the color channel dimension as the batch dimension since the same
-    # transformation is applied to each color channel.
-    # Treat the batch dimension as the channel dimension so that
-    # depthwise_conv2d can apply a different transformation to each sample.
-    cdna_kerns = tf.transpose(cdna_kerns, [1, 2, 0, 4, 3])
-    cdna_kerns = tf.reshape(cdna_kerns,
-                            [self.hparams.dna_kernel_size,
-                             self.hparams.dna_kernel_size,
-                             batch_size,
-                             num_masks])
-    # Swap the batch and channel dimensions.
-    prev_image = tf.transpose(prev_image, [3, 1, 2, 0])
-
-    # Transform image.
-    transformed = tf.nn.depthwise_conv2d(prev_image, cdna_kerns, [1, 1, 1, 1],
-                                         "SAME")
-
-    # Transpose the dimensions to where they belong.
-    transformed = tf.reshape(
-        transformed, [color_channels, height, width, batch_size, num_masks])
-    transformed = tf.transpose(transformed, [3, 1, 2, 0, 4])
-    transformed = tf.unstack(transformed, axis=-1)
-    return transformed
-
-  def dna_transformation(self,
-                         prev_image,
-                         dna_input):
-    """Apply dynamic neural advection to previous image.
-
-    Args:
-      prev_image: previous image to be transformed.
-      dna_input: hidden lyaer to be used for computing DNA transformation.
-    Returns:
-      List of images transformed by the predicted CDNA kernels.
-    """
-    # Construct translated images.
-    prev_image_pad = tf.pad(prev_image, [[0, 0], [2, 2], [2, 2], [0, 0]])
-    image_height = int(prev_image.get_shape()[1])
-    image_width = int(prev_image.get_shape()[2])
-
-    inputs = []
-    for xkern in range(self.hparams.dna_kernel_size):
-      for ykern in range(self.hparams.dna_kernel_size):
-        inputs.append(
-            tf.expand_dims(
-                tf.slice(prev_image_pad, [0, xkern, ykern, 0],
-                         [-1, image_height, image_width, -1]), [3]))
-    inputs = tf.concat(axis=3, values=inputs)
-
-    # Normalize channels to 1.
-    kernel = (tf.nn.relu(dna_input -self.hparams.relu_shift)
-              + self.hparams.relu_shift)
-    kernel = tf.expand_dims(kernel / tf.reduce_sum(kernel, [3], keep_dims=True),
-                            [4])
-    return tf.reduce_sum(kernel * inputs, [3], keep_dims=False)
-
-  @staticmethod
-  def scheduled_sample_count(ground_truth_x,
-                             generated_x,
-                             batch_size,
-                             scheduled_sample_var):
-    """Sample batch with specified mix of groundtruth and generated data points.
-
-    Args:
-      ground_truth_x: tensor of ground-truth data points.
-      generated_x: tensor of generated data points.
-      batch_size: batch size
-      scheduled_sample_var: number of ground-truth examples to include in batch.
-    Returns:
-      New batch with num_ground_truth sampled from ground_truth_x and the rest
-      from generated_x.
-    """
-    num_ground_truth = scheduled_sample_var
-    idx = tf.random_shuffle(tf.range(batch_size))
-    ground_truth_idx = tf.gather(idx, tf.range(num_ground_truth))
-    generated_idx = tf.gather(idx, tf.range(num_ground_truth, batch_size))
-
-    ground_truth_examps = tf.gather(ground_truth_x, ground_truth_idx)
-    generated_examps = tf.gather(generated_x, generated_idx)
-    return tf.dynamic_stitch([ground_truth_idx, generated_idx],
-                             [ground_truth_examps, generated_examps])
-
-  @staticmethod
-  def scheduled_sample_prob(ground_truth_x,
-                            generated_x,
-                            batch_size,
-                            scheduled_sample_var):
-    """Probability based scheduled sampling.
-
-    Args:
-      ground_truth_x: tensor of ground-truth data points.
-      generated_x: tensor of generated data points.
-      batch_size: batch size
-      scheduled_sample_var: probability of choosing from ground_truth.
-    Returns:
-      New batch with randomly selected data points.
-    """
-    probability_threshold = scheduled_sample_var
-    probability_of_generated = tf.random_uniform([batch_size])
-    array_ind = tf.to_int32(probability_of_generated > probability_threshold)
-    indices = tf.range(batch_size) + array_ind * batch_size
-    xy = tf.concat([ground_truth_x, generated_x], axis=0)
-    output = tf.gather(xy, indices)
-    return output
-
   def get_scheduled_sample_func(self, batch_size):
     """Creates a function for scheduled sampling based on given hparams."""
     with tf.variable_scope("scheduled_sampling_func", reuse=False):
@@ -782,7 +576,7 @@ def get_scheduled_sample_func(self, batch_size):
         decay_steps = self.hparams.scheduled_sampling_decay_steps
         probability = tf.train.polynomial_decay(
             1.0, iter_num, decay_steps, 0.0)
-        scheduled_sampling_func = NextFrameStochastic.scheduled_sample_prob
+        scheduled_sampling_func = common_video.scheduled_sample_prob
         scheduled_sampling_func_var = probability
       else:
         # Calculate number of ground-truth frames to pass in.
@@ -791,7 +585,7 @@ def get_scheduled_sample_func(self, batch_size):
             tf.round(
                 tf.to_float(batch_size) *
                 (k / (k + tf.exp(tf.to_float(iter_num) / tf.to_float(k))))))
-        scheduled_sampling_func = NextFrameStochastic.scheduled_sample_count
+        scheduled_sampling_func = common_video.scheduled_sample_count
         scheduled_sampling_func_var = num_ground_truth
 
       tf.summary.scalar("scheduled_sampling_var", scheduled_sampling_func_var)
@@ -826,27 +620,13 @@ def sample():
         return output_items
 
     cases = {
-        tf.logical_not(self.is_training): lambda: generated_items,
         tf.logical_not(done_warm_start): lambda: groundtruth_items,
+        tf.logical_not(self.is_training): lambda: generated_items,
     }
-    output_items = tf.case(cases, default=sample)
+    output_items = tf.case(cases, default=sample, strict=True)
 
     return output_items
 
-  # TODO(mbz): use tf.distributions.kl_divergence instead.
-  def kl_divergence(self, mu, log_sigma):
-    """KL divergence of diagonal gaussian N(mu,exp(log_sigma)) and N(0,1).
-
-    Args:
-      mu: mu parameter of the distribution.
-      log_sigma: log(sigma) parameter of the distribution.
-    Returns:
-      the KL loss.
-    """
-    batch_size = common_layers.shape_list(mu)[0]
-    kl = -.5 * tf.reduce_sum(1. + log_sigma - tf.square(mu) - tf.exp(log_sigma))
-    return kl / tf.to_float(batch_size)
-
   def get_input_if_exists(self, features, key, batch_size, num_frames):
     if key in features:
       x = features[key]
@@ -931,7 +711,7 @@ def anneal_loss(step_num):
     kl_loss = 0.0
     if self.is_training:
       for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
-        kl_loss += self.kl_divergence(mean, std)
+        kl_loss += common_layers.kl_divergence(mean, std)
         tf.summary.histogram("posterior_mean_%d" % i, mean)
         tf.summary.histogram("posterior_std_%d" % i, std)
 
@@ -991,7 +771,7 @@ def construct_model(self, images, actions, rewards):
     # Create scheduled sampling function
     ss_func = self.get_scheduled_sample_func(batch_size)
 
-    pred_image, pred_reward, latent = None, None, None
+    pred_image, pred_reward, latent = images[0], rewards[0], None
     for timestep, image, action, reward in zip(
         range(len(images)-1), images[:-1], actions[:-1], rewards[:-1]):
       # Scheduled Sampling
@@ -1041,33 +821,6 @@ class NextFrameStochasticEmily(NextFrameStochastic):
      https://github.com/edenton/svg
   """
 
-  def vgg_layer(self,
-                inputs,
-                nout,
-                kernel_size=3,
-                activation=tf.nn.leaky_relu,
-                padding="SAME",
-                scope=None):
-    """A layer of VGG network with batch norm.
-
-    Args:
-      inputs: image tensor
-      nout: number of output channels
-      kernel_size: size of the kernel
-      activation: activation function
-      padding: padding of the image
-      scope: variable scope of the op
-    Returns:
-      net: output of layer
-    """
-    with tf.variable_scope(scope):
-      net = tfl.conv2d(inputs, nout, kernel_size=kernel_size, padding=padding,
-                       activation=None, name="conv")
-      net = tfl.batch_normalization(net,
-                                    training=self.is_training, name="bn")
-      net = activation(net)
-    return net
-
   def encoder(self, inputs, nout):
     """VGG based image encoder.
 
@@ -1078,23 +831,28 @@ def encoder(self, inputs, nout):
       net: encoded image with size BSxNout
       skips: skip connection after each layer
     """
-    vgg_layer = self.vgg_layer
+    vgg_layer = common_video.vgg_layer
     net01 = inputs
     # h1
-    net11 = tfcl.repeat(net01, 2, vgg_layer, 64, scope="h1")
+    net11 = tfcl.repeat(net01, 2, vgg_layer, 64,
+                        scope="h1", is_training=self.is_training)
     net12 = tfl.max_pooling2d(net11, [2, 2], strides=(2, 2), name="h1_pool")
     # h2
-    net21 = tfcl.repeat(net12, 2, vgg_layer, 128, scope="h2")
+    net21 = tfcl.repeat(net12, 2, vgg_layer, 128,
+                        scope="h2", is_training=self.is_training)
     net22 = tfl.max_pooling2d(net21, [2, 2], strides=(2, 2), name="h2_pool")
     # h3
-    net31 = tfcl.repeat(net22, 3, vgg_layer, 256, scope="h3")
+    net31 = tfcl.repeat(net22, 3, vgg_layer, 256,
+                        scope="h3", is_training=self.is_training)
     net32 = tfl.max_pooling2d(net31, [2, 2], strides=(2, 2), name="h3_pool")
     # h4
-    net41 = tfcl.repeat(net32, 3, vgg_layer, 512, scope="h4")
+    net41 = tfcl.repeat(net32, 3, vgg_layer, 512,
+                        scope="h4", is_training=self.is_training)
     net42 = tfl.max_pooling2d(net41, [2, 2], strides=(2, 2), name="h4_pool")
     # h5
-    net51 = tfcl.repeat(net42, 1, vgg_layer, nout, kernel_size=4,
-                        padding="VALID", activation=tf.tanh, scope="h5")
+    net51 = tfcl.repeat(net42, 1, vgg_layer, nout,
+                        kernel_size=4, padding="VALID", activation=tf.tanh,
+                        scope="h5", is_training=self.is_training)
     skips = [net11, net21, net31, net41]
     return net51, skips
 
@@ -1109,7 +867,7 @@ def decoder(self, inputs, skips, nout):
       net: decoded image with size BSx64x64xNout
       skips: skip connection after each layer
     """
-    vgg_layer = self.vgg_layer
+    vgg_layer = common_video.vgg_layer
     net = inputs
     # d1
     net = tfl.conv2d_transpose(net, 512, kernel_size=4, padding="VALID",
@@ -1156,7 +914,7 @@ def stacked_lstm(self, inputs, states, hidden_size, output_size, nlayers):
     net = tfl.dense(
         net, hidden_size, activation=None, name="af1")
     for i in range(nlayers):
-      net, states[i] = basic_lstm(
+      net, states[i] = common_video.basic_lstm(
           net, states[i], hidden_size, name="alstm%d"%i)
     net = tfl.dense(
         net, output_size, activation=tf.nn.tanh, name="af2")
@@ -1179,7 +937,7 @@ def lstm_gaussian(self, inputs, states, hidden_size, output_size, nlayers):
     net = inputs
     net = tfl.dense(net, hidden_size, activation=None, name="bf1")
     for i in range(nlayers):
-      net, states[i] = basic_lstm(
+      net, states[i] = common_video.basic_lstm(
           net, states[i], hidden_size, name="blstm%d"%i)
     mu = tfl.dense(net, output_size, activation=None, name="bf2mu")
     logvar = tfl.dense(net, output_size, activation=None, name="bf2log")
diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index 94aea9b94..3737fbb27 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -73,7 +73,7 @@ def next_frame_stochastic():
   hparams.add_hparam("relu_shift", 1e-12)
   hparams.add_hparam("dna_kernel_size", 5)
   # Scheduled sampling method. Choose between prob or count.
-  hparams.add_hparam("scheduled_sampling_mode", "prob")
+  hparams.add_hparam("scheduled_sampling_mode", "count")
   hparams.add_hparam("scheduled_sampling_decay_steps", 10000)
   hparams.add_hparam("scheduled_sampling_k", 900.0)
   hparams.add_hparam("latent_num_frames", 0)  # 0 means use all frames.
diff --git a/tensor2tensor/models/research/next_frame_savp.py b/tensor2tensor/models/research/next_frame_savp.py
index 73b65c599..8681ec32e 100644
--- a/tensor2tensor/models/research/next_frame_savp.py
+++ b/tensor2tensor/models/research/next_frame_savp.py
@@ -22,6 +22,7 @@
 from __future__ import print_function
 
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import next_frame
 from tensor2tensor.models.research import next_frame_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
@@ -163,11 +164,11 @@ def construct_model(self, images, actions, rewards):
 
         # LSTM that encodes correlations between conditional latents.
         # Pg 22 in https://arxiv.org/pdf/1804.01523.pdf
-        enc_cond_latent, cond_latent_state = next_frame.basic_lstm(
+        enc_cond_latent, cond_latent_state = common_video.basic_lstm(
             latent, cond_latent_state, latent_dims, name="cond_latent")
 
         # LSTM that encodes correlations between prior latents.
-        enc_prior_latent, prior_latent_state = next_frame.basic_lstm(
+        enc_prior_latent, prior_latent_state = common_video.basic_lstm(
             prior_latent, prior_latent_state, latent_dims, name="prior_latent")
 
         # Scheduled Sampling
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index fd9d668d6..d31559c8e 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -91,78 +91,6 @@ def testStochasticEmily(self):
         next_frame.NextFrameStochasticEmily,
         1)
 
-  @staticmethod
-  def runScheduledSampleFunc(func, var, batch_size):
-    ground_truth_x = list(range(1, batch_size+1))
-    generated_x = [-x for x in ground_truth_x]
-    ground_truth_x = tf.convert_to_tensor(ground_truth_x)
-    generated_x = tf.convert_to_tensor(generated_x)
-    ss_out = func(ground_truth_x, generated_x, batch_size, var)
-    with tf.Session() as session:
-      output = session.run([ground_truth_x, generated_x, ss_out])
-    return output
-
-  def testScheduledSampleProbStart(self):
-    ground_truth_x, _, ss_out = NextFrameTest.runScheduledSampleFunc(
-        next_frame.NextFrameStochastic.scheduled_sample_prob, 1.0, 10)
-    self.assertAllEqual(ground_truth_x, ss_out)
-
-  def testScheduledSampleProbMid(self):
-    _, _, ss_out = NextFrameTest.runScheduledSampleFunc(
-        next_frame.NextFrameStochastic.scheduled_sample_prob, 0.5, 1000)
-    positive_count = np.sum(ss_out > 0)
-    self.assertAlmostEqual(positive_count / 1000.0, 0.5, places=2)
-
-  def testScheduledSampleProbEnd(self):
-    _, generated_x, ss_out = NextFrameTest.runScheduledSampleFunc(
-        next_frame.NextFrameStochastic.scheduled_sample_prob, 0.0, 10)
-    self.assertAllEqual(generated_x, ss_out)
-
-  def testScheduledSampleCountStart(self):
-    ground_truth_x, _, ss_out = NextFrameTest.runScheduledSampleFunc(
-        next_frame.NextFrameStochastic.scheduled_sample_count, 10, 10)
-    self.assertAllEqual(ground_truth_x, ss_out)
-
-  def testScheduledSampleCountMid(self):
-    _, _, ss_out = NextFrameTest.runScheduledSampleFunc(
-        next_frame.NextFrameStochastic.scheduled_sample_count, 5, 10)
-    positive_count = np.sum(ss_out > 0)
-    self.assertEqual(positive_count, 5)
-
-  def testScheduledSampleCountEnd(self):
-    _, generated_x, ss_out = NextFrameTest.runScheduledSampleFunc(
-        next_frame.NextFrameStochastic.scheduled_sample_count, 0, 10)
-    self.assertAllEqual(generated_x, ss_out)
-
-  def testDynamicTileAndConcat(self):
-    with tf.Graph().as_default():
-      # image = (1 X 4 X 4 X 1)
-      image = [[1, 2, 3, 4],
-               [2, 4, 5, 6],
-               [7, 8, 9, 10],
-               [7, 9, 10, 1]]
-      image_t = tf.expand_dims(tf.expand_dims(image, axis=0), axis=-1)
-      image_t = tf.cast(image_t, dtype=tf.float32)
-
-      # latent = (1 X 2)
-      latent = np.array([[90, 100]])
-      latent_t = tf.cast(tf.convert_to_tensor(latent), dtype=tf.float32)
-
-      with tf.Session() as session:
-        tiled = next_frame.NextFrameStochastic.tile_and_concat(
-            image_t, latent_t)
-        tiled_np, image_np = session.run([tiled, image_t])
-        tiled_latent = tiled_np[0, :, :, -1]
-        self.assertAllEqual(tiled_np.shape, (1, 4, 4, 2))
-
-        self.assertAllEqual(tiled_np[:, :, :, :1], image_np)
-        self.assertAllEqual(
-            tiled_latent,
-            [[90, 90, 90, 90],
-             [100, 100, 100, 100],
-             [90, 90, 90, 90],
-             [100, 100, 100, 100]])
-
 
 if __name__ == "__main__":
   tf.test.main()

From 7a5f20b475b391c102acc22d710d1aa722f828b9 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 1 Aug 2018 17:15:02 -0700
Subject: [PATCH 0441/2720] Test against last 3 versions of TF plus tf-nightly;
 only test models/research against tf-nightly

PiperOrigin-RevId: 207022097
---
 .travis.yml                              | 27 ++++++++++++++----------
 setup.py                                 |  4 ++--
 tensor2tensor/data_generators/problem.py | 12 +----------
 tensor2tensor/serving/README.md          |  1 -
 4 files changed, 19 insertions(+), 25 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 19fa25160..6211fdd48 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -11,21 +11,18 @@ env:
     # This is necessary to have gsutil work with Python 2.7
     - BOTO_CONFIG=/dev/null
   matrix:
-    # We test against the last 4 versions of TensorFlow
+    # We test against recent versions of TensorFlow and tf-nightly.
     # If updating, also update TF_LATEST above
-    - TF_VERSION="1.6.*"
-    - TF_VERSION="1.7.*"
     - TF_VERSION="1.8.*"
     - TF_VERSION="1.9.*"
+    - TF_VERSION="tf-nightly"
 matrix:
   exclude:
     # We test against all versions in Python 2 but only the latest in Python 3
-    - python: "3.6"
-      env: TF_VERSION="1.6.*"
-    - python: "3.6"
-      env: TF_VERSION="1.7.*"
     - python: "3.6"
       env: TF_VERSION="1.8.*"
+    - python: "3.6"
+      env: TF_VERSION="tf-nightly"
 before_install:
   # Disabled TensorFlow Serving install until bug fixed. See "Export and query"
   # section below.
@@ -36,6 +33,12 @@ before_install:
   # - sudo apt-get install -qq tensorflow-model-server
 install:
   - pip install -q "tensorflow==$TF_VERSION"
+  - if [[ "$TF_VERSION" == "tf-nightly"  ]];
+    then
+      pip install tf-nightly;
+    else
+      pip install -q "tensorflow==$TF_VERSION";
+    fi
   # First ensure that the base dependencies are sufficient for a full import
   - pip install -q .
   - t2t-trainer --registry_help
@@ -58,6 +61,7 @@ script:
   #   * visualization_test
   #   * model_rl_experiment_test
   #   * model_rl_experiment_stochastic_test
+  #   * models/research
   # algorithmic_math_test: flaky
   # universal_transformer_test: requires new feature in tf.foldl (rm with TF 1.9)
   - pytest
@@ -69,9 +73,14 @@ script:
     --ignore=tensor2tensor/models/research/universal_transformer_test.py
     --ignore=tensor2tensor/rl/model_rl_experiment_test.py
     --ignore=tensor2tensor/rl/model_rl_experiment_stochastic_test.py
+    --ignore=tensor2tensor/models/research
   - pytest tensor2tensor/utils/registry_test.py
   - pytest tensor2tensor/utils/trainer_lib_test.py
   - pytest tensor2tensor/visualization/visualization_test.py
+  - if [[ "$TF_VERSION" == "$TF_LATEST"  ]] || [[ "$TF_VERSION" == "tf-nightly"  ]];
+    then
+      pytest tensor2tensor/models/research;
+    fi
 
   # Run installed scripts
   - t2t-datagen 2>&1 | grep translate && echo passed
@@ -88,11 +97,7 @@ script:
   # Do some things only on Python 2 and the latest TF version
   - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
         pylint -j 2 tensor2tensor;
-    fi
-  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
         pytest tensor2tensor/rl/model_rl_experiment_test.py;
-    fi
-  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
         jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/hello_t2t.ipynb;
     fi
 
diff --git a/setup.py b/setup.py
index b1cc9754c..409486722 100644
--- a/setup.py
+++ b/setup.py
@@ -49,8 +49,8 @@
         'tqdm',
     ],
     extras_require={
-        'tensorflow': ['tensorflow>=1.6.0'],
-        'tensorflow_gpu': ['tensorflow-gpu>=1.6.0'],
+        'tensorflow': ['tensorflow>=1.8.0'],
+        'tensorflow_gpu': ['tensorflow-gpu>=1.8.0'],
         'tests': [
             'absl-py', 'pytest', 'mock', 'pylint', 'jupyter', 'gsutil'
             # Need atari extras for Travis tests, but because gym is already in
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index bd1362335..9124a6bee 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -627,17 +627,7 @@ def _load_records_and_preprocess(filenames):
           tf.contrib.data.parallel_interleave(
               _load_records_and_preprocess, sloppy=True, cycle_length=8))
     else:
-      # TFRecordDataset can get filenames as dataset in TF 1.7+.
-      # TODO(lukaszkaiser): remove when we require TF 1.7+ in general.
-      major, minor = [int(el) for el in tf.__version__.split(".")[:2]]
-      filename_dataset_ok = major > 1 or (major == 1 and minor >= 7)
-      if filename_dataset_ok:  # We can just pass a Dataset of filenames.
-        dataset = _load_records_and_preprocess(dataset)
-      else:  # Go file-by-file (can be very slow).
-        dataset = None
-        for f in data_files:
-          f_data = _load_records_and_preprocess(f)
-          dataset = f_data if dataset is None else dataset.concatenate(f_data)
+      dataset = _load_records_and_preprocess(dataset)
 
     dataset = dataset.map(
         self.maybe_reverse_and_copy, num_parallel_calls=num_threads)
diff --git a/tensor2tensor/serving/README.md b/tensor2tensor/serving/README.md
index 8bb35da27..156a2503d 100644
--- a/tensor2tensor/serving/README.md
+++ b/tensor2tensor/serving/README.md
@@ -82,7 +82,6 @@ This step only needs to be performed once.
 VERSION=v0
 gcloud ml-engine versions create $VERSION \
   --model $MODEL_NAME \
-  --runtime-version 1.6 \
   --origin $ORIGIN
 ```
 

From 60bdc2bfdecc878ed02c94e5b542edb51314a6df Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 1 Aug 2018 18:06:58 -0700
Subject: [PATCH 0442/2720] fix travis.yml for tf-nightly

PiperOrigin-RevId: 207028087
---
 .travis.yml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/.travis.yml b/.travis.yml
index 6211fdd48..992537d65 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -32,7 +32,6 @@ before_install:
   - sudo apt-get install -qq libhdf5-dev
   # - sudo apt-get install -qq tensorflow-model-server
 install:
-  - pip install -q "tensorflow==$TF_VERSION"
   - if [[ "$TF_VERSION" == "tf-nightly"  ]];
     then
       pip install tf-nightly;

From 16f7d6f1e5ae078cd488ca5160f0837e1459b2fd Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 1 Aug 2018 18:13:15 -0700
Subject: [PATCH 0443/2720] Enable MLEngine path to use Cloud TPUs

PiperOrigin-RevId: 207028905
---
 tensor2tensor/utils/cloud_mlengine.py | 16 ++++++++--------
 tensor2tensor/utils/trainer_lib.py    |  6 ++++++
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index f9a4080a8..88e87778d 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -91,10 +91,8 @@ def flags_as_args():
   return args
 
 
-def get_default_master_type(num_gpus=1, use_tpu=False):
+def get_default_master_type(num_gpus=1):
   """Returns master_type for trainingInput."""
-  if use_tpu:
-    return "cloud_tpu"
   gpus_to_master_map = {
       0: "standard",
       1: "standard_p100",
@@ -120,9 +118,13 @@ def configure_job():
       "jobDir": FLAGS.output_dir,
       "scaleTier": "CUSTOM",
       "masterType": FLAGS.cloud_mlengine_master_type or get_default_master_type(
-          num_gpus=FLAGS.worker_gpu,
-          use_tpu=FLAGS.use_tpu)
+          num_gpus=FLAGS.worker_gpu)
   }
+  if FLAGS.use_tpu:
+    training_input["masterType"] = (FLAGS.cloud_mlengine_master_type or
+                                    "standard")
+    training_input["workerType"] = "cloud_tpu"
+    training_input["workerCount"] = 1
   if FLAGS.hparams_range:
     tf.logging.info("Configuring hyperparameter tuning.")
     training_input["hyperparameters"] = configure_autotune(
@@ -277,9 +279,7 @@ def validate_flags():
   if FLAGS.worker_gpu:
     assert FLAGS.worker_gpu in [1, 4, 8]
   if FLAGS.cloud_mlengine_master_type:
-    if FLAGS.use_tpu:
-      assert FLAGS.cloud_mlengine_master_type == "cloud_tpu"
-    elif FLAGS.worker_gpu:
+    if FLAGS.worker_gpu:
       if FLAGS.worker_gpu == 1:
         assert FLAGS.cloud_mlengine_master_type in ["standard_gpu",
                                                     "standard_p100"]
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 6ac4a5fe3..05a31ac89 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -168,6 +168,12 @@ def create_run_config(master="",
     tpu_config = tf.contrib.tpu.TPUConfig(
         **tpu_config_kwargs)
     run_config_args["tpu_config"] = tpu_config
+    if not master and "KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS" in os.environ:
+      # If running on TPU but no master is set and the KUBE env var is present
+      # then we're running on ML Engine. Set the master.
+      run_config_args["master"] = os.environ[
+          "KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS"]
+      run_config_args["evaluation_master"] = run_config_args["master"]
 
   config = run_config_cls(**run_config_args)
   config.warm_start_from = warm_start_from

From a3e8d934d147897d439b88db87106a476aa8f441 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 1 Aug 2018 18:24:12 -0700
Subject: [PATCH 0444/2720] Split up some Travis commands

PiperOrigin-RevId: 207030088
---
 .travis.yml | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/.travis.yml b/.travis.yml
index 992537d65..b45cbcb6a 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -94,9 +94,14 @@ script:
   - t2t-decoder --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR --model=transformer --hparams_set=transformer_tiny --output_dir=$T2T_TRAIN_DIR --decode_hparams='num_samples=10'
 
   # Do some things only on Python 2 and the latest TF version
+  # Each should be in a separate block to get proper errors on Travis.
   - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
         pylint -j 2 tensor2tensor;
+    fi
+  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
         pytest tensor2tensor/rl/model_rl_experiment_test.py;
+    fi
+  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
         jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/hello_t2t.ipynb;
     fi
 

From e5150c1445be17979d359719ea11ce4925ad12a2 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 1 Aug 2018 18:37:55 -0700
Subject: [PATCH 0445/2720] breaking down next_frame.py

PiperOrigin-RevId: 207031577
---
 tensor2tensor/layers/common_video.py          |   6 +
 tensor2tensor/layers/common_video_test.py     |  16 +-
 tensor2tensor/models/research/next_frame.py   | 877 ------------------
 .../models/research/next_frame_emily.py       | 263 ++++++
 .../models/research/next_frame_savp.py        |   4 +-
 .../models/research/next_frame_savp_test.py   |  88 --
 .../models/research/next_frame_sv2p.py        | 671 ++++++++++++++
 .../models/research/next_frame_test.py        |  25 +-
 8 files changed, 974 insertions(+), 976 deletions(-)
 create mode 100644 tensor2tensor/models/research/next_frame_emily.py
 delete mode 100644 tensor2tensor/models/research/next_frame_savp_test.py
 create mode 100644 tensor2tensor/models/research/next_frame_sv2p.py

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 6a15401fa..4579a2709 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -23,6 +23,12 @@
 tfl = tf.layers
 
 
+def swap_time_and_batch_axes(inputs):
+  """Swaps time and batch axis (the first two axis)."""
+  transposed_axes = tf.concat([[1, 0], tf.range(2, tf.rank(inputs))], axis=0)
+  return tf.transpose(inputs, transposed_axes)
+
+
 def encode_to_shape(inputs, shape, scope):
   """Encode the given tensor to given image shape."""
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index 4a4f939b5..3eb32ea5e 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -37,34 +37,34 @@ def runScheduledSampleFunc(func, var, batch_size):
 
   def testScheduledSampleProbStart(self):
     ground_truth_x, _, ss_out = CommonVideoTest.runScheduledSampleFunc(
-        common_video.NextFrameStochastic.scheduled_sample_prob, 1.0, 10)
+        common_video.scheduled_sample_prob, 1.0, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
   def testScheduledSampleProbMid(self):
     _, _, ss_out = CommonVideoTest.runScheduledSampleFunc(
-        common_video.NextFrameStochastic.scheduled_sample_prob, 0.5, 1000)
+        common_video.scheduled_sample_prob, 0.5, 1000)
     positive_count = np.sum(ss_out > 0)
     self.assertAlmostEqual(positive_count / 1000.0, 0.5, places=2)
 
   def testScheduledSampleProbEnd(self):
     _, generated_x, ss_out = CommonVideoTest.runScheduledSampleFunc(
-        common_video.NextFrameStochastic.scheduled_sample_prob, 0.0, 10)
+        common_video.scheduled_sample_prob, 0.0, 10)
     self.assertAllEqual(generated_x, ss_out)
 
   def testScheduledSampleCountStart(self):
     ground_truth_x, _, ss_out = CommonVideoTest.runScheduledSampleFunc(
-        common_video.NextFrameStochastic.scheduled_sample_count, 10, 10)
+        common_video.scheduled_sample_count, 10, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
   def testScheduledSampleCountMid(self):
     _, _, ss_out = CommonVideoTest.runScheduledSampleFunc(
-        common_video.NextFrameStochastic.scheduled_sample_count, 5, 10)
+        common_video.scheduled_sample_count, 5, 10)
     positive_count = np.sum(ss_out > 0)
     self.assertEqual(positive_count, 5)
 
   def testScheduledSampleCountEnd(self):
     _, generated_x, ss_out = CommonVideoTest.runScheduledSampleFunc(
-        common_video.NextFrameStochastic.scheduled_sample_count, 0, 10)
+        common_video.scheduled_sample_count, 0, 10)
     self.assertAllEqual(generated_x, ss_out)
 
   def testDynamicTileAndConcat(self):
@@ -95,3 +95,7 @@ def testDynamicTileAndConcat(self):
              [100, 100, 100, 100],
              [90, 90, 90, 90],
              [100, 100, 100, 100]])
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index 28b338423..a49f86fa4 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -18,12 +18,10 @@
 from __future__ import division
 from __future__ import print_function
 
-from functools import partial
 import six
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
-from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import next_frame_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
@@ -165,878 +163,3 @@ def logits_to_samples(logits):
     return results
 
 
-_LARGE_STEP_NUMBER = 100000
-
-
-@registry.register_model
-class NextFrameStochastic(NextFrameBasic):
-  """ SV2P: Stochastic Variational Video Prediction.
-
-  based on the following papaer:
-  https://arxiv.org/abs/1710.11252
-  """
-
-  @property
-  def is_training(self):
-    return self.hparams.mode == tf.estimator.ModeKeys.TRAIN
-
-  def tinyify(self, array):
-    if self.hparams.tiny_mode:
-      return [1 for _ in array]
-    return array
-
-  def construct_latent_tower(self, images):
-    """Builds convolutional latent tower for stochastic model.
-
-    At training time this tower generates a latent distribution (mean and std)
-    conditioned on the entire video. This latent variable will be fed to the
-    main tower as an extra variable to be used for future frames prediction.
-    At inference time, the tower is disabled and only returns latents sampled
-    from N(0,1).
-    If the multi_latent flag is on, a different latent for every timestep would
-    be generated.
-
-    Args:
-      images: tensor of ground truth image sequences
-    Returns:
-      latent_mean: predicted latent mean
-      latent_std: predicted latent standard deviation
-      latent_loss: loss of the latent twoer
-      samples: random samples sampled from standard guassian
-    """
-    conv_size = self.tinyify([32, 64, 64])
-    with tf.variable_scope("latent", reuse=tf.AUTO_REUSE):
-      # this allows more predicted frames at inference time
-      latent_num_frames = self.hparams.latent_num_frames
-      if latent_num_frames == 0:  # use all frames by default.
-        latent_num_frames = (self.hparams.video_num_input_frames +
-                             self.hparams.video_num_target_frames)
-      tf.logging.info("Creating latent tower with %d frames."%latent_num_frames)
-      latent_images = tf.unstack(images[:latent_num_frames], axis=0)
-      images = tf.concat(latent_images, 3)
-
-      x = images
-      x = common_layers.make_even_size(x)
-      x = tfl.conv2d(x, conv_size[0], [3, 3], strides=(2, 2),
-                     padding="SAME", activation=tf.nn.relu, name="latent_conv1")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="latent_bn1")
-      x = common_layers.make_even_size(x)
-      x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
-                     padding="SAME", activation=tf.nn.relu, name="latent_conv2")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="latent_bn2")
-      x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(1, 1),
-                     padding="SAME", activation=tf.nn.relu, name="latent_conv3")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="latent_bn3")
-
-      nc = self.hparams.latent_channels
-      mean = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
-                        padding="SAME", activation=None, name="latent_mean")
-      std = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
-                       padding="SAME", activation=tf.nn.relu, name="latent_std")
-      std += self.hparams.latent_std_min
-
-      # No latent tower at inference time, just standard gaussian.
-      if self.hparams.mode != tf.estimator.ModeKeys.TRAIN:
-        return tf.zeros_like(mean), tf.zeros_like(std)
-
-      return mean, std
-
-  def bottom_part_tower(self, input_image, input_reward, action, latent,
-                        lstm_state, lstm_size, conv_size, concat_latent=False):
-    """The bottom part of predictive towers.
-
-    With the current (early) design, the main prediction tower and
-    the reward prediction tower share the same arcitecture. TF Scope can be
-    adjusted as required to either share or not share the weights between
-    the two towers.
-
-    Args:
-      input_image: the current image.
-      input_reward: the current reward.
-      action: the action taken by the agent.
-      latent: the latent vector.
-      lstm_state: the current internal states of conv lstms.
-      lstm_size: the size of lstms.
-      conv_size: the size of convolutions.
-      concat_latent: whether or not to concatenate the latent at every step.
-
-    Returns:
-      - the output of the partial network.
-      - intermidate outputs for skip connections.
-    """
-    lstm_func = common_video.conv_lstm_2d
-    tile_and_concat = common_video.tile_and_concat
-
-    input_image = common_layers.make_even_size(input_image)
-    concat_input_image = tile_and_concat(
-        input_image, latent, concat_latent=concat_latent)
-
-    enc0 = tfl.conv2d(
-        concat_input_image,
-        conv_size[0], [5, 5],
-        strides=(2, 2),
-        activation=tf.nn.relu,
-        padding="SAME",
-        name="scale1_conv1")
-    enc0 = tfcl.layer_norm(enc0, scope="layer_norm1")
-
-    hidden1, lstm_state[0] = lstm_func(
-        enc0, lstm_state[0], lstm_size[0], name="state1")
-    hidden1 = tile_and_concat(hidden1, latent, concat_latent=concat_latent)
-    hidden1 = tfcl.layer_norm(hidden1, scope="layer_norm2")
-    hidden2, lstm_state[1] = lstm_func(
-        hidden1, lstm_state[1], lstm_size[1], name="state2")
-    hidden2 = tfcl.layer_norm(hidden2, scope="layer_norm3")
-    hidden2 = common_layers.make_even_size(hidden2)
-    enc1 = tfl.conv2d(hidden2, hidden2.get_shape()[3], [3, 3], strides=(2, 2),
-                      padding="SAME", activation=tf.nn.relu, name="conv2")
-    enc1 = tile_and_concat(enc1, latent, concat_latent=concat_latent)
-
-    hidden3, lstm_state[2] = lstm_func(
-        enc1, lstm_state[2], lstm_size[2], name="state3")
-    hidden3 = tile_and_concat(hidden3, latent, concat_latent=concat_latent)
-    hidden3 = tfcl.layer_norm(hidden3, scope="layer_norm4")
-    hidden4, lstm_state[3] = lstm_func(
-        hidden3, lstm_state[3], lstm_size[3], name="state4")
-    hidden4 = tile_and_concat(hidden4, latent, concat_latent=concat_latent)
-    hidden4 = tfcl.layer_norm(hidden4, scope="layer_norm5")
-    hidden4 = common_layers.make_even_size(hidden4)
-    enc2 = tfl.conv2d(hidden4, hidden4.get_shape()[3], [3, 3], strides=(2, 2),
-                      padding="SAME", activation=tf.nn.relu, name="conv3")
-
-    # Pass in reward and action.
-    emb_action = common_video.encode_to_shape(
-        action, enc2.get_shape(), "action_enc")
-    emb_reward = common_video.encode_to_shape(
-        input_reward, enc2.get_shape(), "reward_enc")
-    enc2 = tf.concat(axis=3, values=[enc2, emb_action, emb_reward])
-
-    if latent is not None and not concat_latent:
-      with tf.control_dependencies([latent]):
-        enc2 = tf.concat([enc2, latent], 3)
-
-    enc3 = tfl.conv2d(enc2, hidden4.get_shape()[3], [1, 1], strides=(1, 1),
-                      padding="SAME", activation=tf.nn.relu, name="conv4")
-
-    hidden5, lstm_state[4] = lstm_func(
-        enc3, lstm_state[4], lstm_size[4], name="state5")  # last 8x8
-    hidden5 = tfcl.layer_norm(hidden5, scope="layer_norm6")
-    hidden5 = tile_and_concat(hidden5, latent, concat_latent=concat_latent)
-    return hidden5, (enc0, enc1)
-
-  def reward_prediction(
-      self, input_image, input_reward, action, lstm_state, latent):
-    """Builds a reward prediction network."""
-    conv_size = self.tinyify([32, 32, 16, 4])
-    lstm_size = self.tinyify([32, 64, 128, 64, 32])
-
-    with tf.variable_scope("reward_pred", reuse=tf.AUTO_REUSE):
-      hidden5, _ = self.bottom_part_tower(
-          input_image, input_reward, action, latent,
-          lstm_state, lstm_size, conv_size)
-
-      x = hidden5
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="reward_bn0")
-      x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
-                     padding="SAME", activation=tf.nn.relu, name="reward_conv1")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="reward_bn1")
-      x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(2, 2),
-                     padding="SAME", activation=tf.nn.relu, name="reward_conv2")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="reward_bn2")
-      x = tfl.conv2d(x, conv_size[3], [3, 3], strides=(2, 2),
-                     padding="SAME", activation=tf.nn.relu, name="reward_conv3")
-
-      pred_reward = common_video.decode_to_shape(
-          x, input_reward.shape, "reward_dec")
-
-      return pred_reward, lstm_state
-
-  def construct_predictive_tower(
-      self, input_image, input_reward, action, lstm_state, latent,
-      concat_latent=False):
-    # Main tower
-    lstm_func = common_video.conv_lstm_2d
-    batch_size = common_layers.shape_list(input_image)[0]
-    # the number of different pixel motion predictions
-    # and the number of masks for each of those predictions
-    num_masks = self.hparams.num_masks
-    upsample_method = self.hparams.upsample_method
-    tile_and_concat = common_video.tile_and_concat
-
-    lstm_size = self.tinyify([32, 32, 64, 64, 128, 64, 32])
-    conv_size = self.tinyify([32])
-
-    img_height, img_width, color_channels = self.hparams.problem.frame_shape
-
-    with tf.variable_scope("main", reuse=tf.AUTO_REUSE):
-      hidden5, skips = self.bottom_part_tower(
-          input_image, input_reward, action, latent,
-          lstm_state, lstm_size, conv_size, concat_latent=concat_latent)
-      enc0, enc1 = skips
-
-      with tf.variable_scope("upsample1", reuse=tf.AUTO_REUSE):
-        enc4 = common_layers.cyclegan_upsample(
-            hidden5, num_outputs=hidden5.shape.as_list()[-1],
-            stride=[2, 2], method=upsample_method)
-
-      enc1_shape = common_layers.shape_list(enc1)
-      enc4 = enc4[:, :enc1_shape[1], :enc1_shape[2], :]  # Cut to shape.
-      enc4 = tile_and_concat(enc4, latent, concat_latent=concat_latent)
-
-      hidden6, lstm_state[5] = lstm_func(
-          enc4, lstm_state[5], lstm_size[5], name="state6",
-          spatial_dims=enc1_shape[1:-1])  # 16x16
-      hidden6 = tile_and_concat(hidden6, latent, concat_latent=concat_latent)
-      hidden6 = tfcl.layer_norm(hidden6, scope="layer_norm7")
-      # Skip connection.
-      hidden6 = tf.concat(axis=3, values=[hidden6, enc1])  # both 16x16
-
-      with tf.variable_scope("upsample2", reuse=tf.AUTO_REUSE):
-        enc5 = common_layers.cyclegan_upsample(
-            hidden6, num_outputs=hidden6.shape.as_list()[-1],
-            stride=[2, 2], method=upsample_method)
-
-      enc0_shape = common_layers.shape_list(enc0)
-      enc5 = enc5[:, :enc0_shape[1], :enc0_shape[2], :]  # Cut to shape.
-      enc5 = tile_and_concat(enc5, latent, concat_latent=concat_latent)
-
-      hidden7, lstm_state[6] = lstm_func(
-          enc5, lstm_state[6], lstm_size[6], name="state7",
-          spatial_dims=enc0_shape[1:-1])  # 32x32
-      hidden7 = tfcl.layer_norm(hidden7, scope="layer_norm8")
-
-      # Skip connection.
-      hidden7 = tf.concat(axis=3, values=[hidden7, enc0])  # both 32x32
-
-      with tf.variable_scope("upsample3", reuse=tf.AUTO_REUSE):
-        enc6 = common_layers.cyclegan_upsample(
-            hidden7, num_outputs=hidden7.shape.as_list()[-1],
-            stride=[2, 2], method=upsample_method)
-      enc6 = tfcl.layer_norm(enc6, scope="layer_norm9")
-      enc6 = tile_and_concat(enc6, latent, concat_latent=concat_latent)
-
-      if self.hparams.model_options == "DNA":
-        # Using largest hidden state for predicting untied conv kernels.
-        enc7 = tfl.conv2d_transpose(
-            enc6,
-            self.hparams.dna_kernel_size**2,
-            [1, 1],
-            strides=(1, 1),
-            padding="SAME",
-            name="convt4",
-            activation=None)
-      else:
-        # Using largest hidden state for predicting a new image layer.
-        enc7 = tfl.conv2d_transpose(
-            enc6,
-            color_channels,
-            [1, 1],
-            strides=(1, 1),
-            padding="SAME",
-            name="convt4",
-            activation=None)
-        # This allows the network to also generate one image from scratch,
-        # which is useful when regions of the image become unoccluded.
-        transformed = [tf.nn.sigmoid(enc7)]
-
-      if self.hparams.model_options == "CDNA":
-        # cdna_input = tf.reshape(hidden5, [int(batch_size), -1])
-        cdna_input = tfcl.flatten(hidden5)
-        transformed += common_video.cdna_transformation(
-            input_image, cdna_input, num_masks, int(color_channels),
-            self.hparams.dna_kernel_size, self.hparams.relu_shift)
-      elif self.hparams.model_options == "DNA":
-        # Only one mask is supported (more should be unnecessary).
-        if num_masks != 1:
-          raise ValueError("Only one mask is supported for DNA model.")
-        transformed = [
-            common_video.dna_transformation(
-                input_image, enc7,
-                self.hparams.dna_kernel_size, self.hparams.relu_shift)]
-
-      masks = tfl.conv2d(
-          enc6, filters=num_masks + 1, kernel_size=[1, 1],
-          strides=(1, 1), name="convt7", padding="SAME")
-      masks = tf.reshape(
-          tf.nn.softmax(tf.reshape(masks, [-1, num_masks + 1])),
-          [batch_size,
-           int(img_height),
-           int(img_width), num_masks + 1])
-      mask_list = tf.split(
-          axis=3, num_or_size_splits=num_masks + 1, value=masks)
-      output = mask_list[0] * input_image
-      for layer, mask in zip(transformed, mask_list[1:]):
-        output += layer * mask
-
-      return output, lstm_state
-
-  def get_gaussian_latent(self, latent_mean, latent_std):
-    latent = tf.random_normal(tf.shape(latent_mean), 0, 1, dtype=tf.float32)
-    latent = latent_mean + tf.exp(latent_std / 2.0) * latent
-    return latent
-
-  def construct_model(self,
-                      images,
-                      actions,
-                      rewards):
-    """Build convolutional lstm video predictor using CDNA, or DNA.
-
-    Args:
-      images: list of tensors of ground truth image sequences
-              there should be a 4D image ?xWxHxC for each timestep
-      actions: list of action tensors
-               each action should be in the shape ?x1xZ
-      rewards: list of reward tensors
-               each reward should be in the shape ?x1xZ
-    Returns:
-      gen_images: predicted future image frames
-      gen_rewards: predicted future rewards
-      latent_mean: mean of approximated posterior
-      latent_std: std of approximated posterior
-
-    Raises:
-      ValueError: if more than 1 mask specified for DNA model.
-    """
-    context_frames = self.hparams.video_num_input_frames
-
-    batch_size = common_layers.shape_list(images)[1]
-    ss_func = self.get_scheduled_sample_func(batch_size)
-
-    def process_single_frame(prev_outputs, inputs):
-      """Process a single frame of the video."""
-      cur_image, cur_reward, action = inputs
-      time_step, prev_image, prev_reward, lstm_states = prev_outputs[:4]
-
-      generated_items = [prev_image, prev_reward]
-      groundtruth_items = [cur_image, cur_reward]
-      done_warm_start = tf.greater(time_step, context_frames - 1)
-      input_image, input_reward = self.get_scheduled_sample_inputs(
-          done_warm_start, groundtruth_items, generated_items, ss_func)
-
-      # Prediction
-      pred_image, lstm_states = self.construct_predictive_tower(
-          input_image, input_reward, action, lstm_states, latent)
-
-      if self.hparams.reward_prediction:
-        reward_lstm_states = prev_outputs[4]
-        pred_reward, reward_lstm_states = self.reward_prediction(
-            input_image, input_reward, action, reward_lstm_states, latent)
-      else:
-        pred_reward = input_reward
-
-      time_step += 1
-      outputs = (time_step, pred_image, pred_reward, lstm_states)
-      if self.hparams.reward_prediction:
-        outputs += (reward_lstm_states,)
-
-      return outputs
-
-    # Latent tower
-    latent = None
-    if self.hparams.stochastic_model:
-      latent_mean, latent_std = self.construct_latent_tower(images)
-      latent = self.get_gaussian_latent(latent_mean, latent_std)
-
-    # HACK: Do first step outside to initialize all the variables
-    lstm_states, reward_lstm_states = [None] * 7, [None] * 5
-    inputs = images[0], rewards[0], actions[0]
-    prev_outputs = (tf.constant(0), images[0], rewards[0], lstm_states)
-    if self.hparams.reward_prediction:
-      prev_outputs += (reward_lstm_states,)
-
-    initializers = process_single_frame(prev_outputs, inputs)
-    first_gen_images = tf.expand_dims(initializers[1], axis=0)
-    first_gen_rewards = tf.expand_dims(initializers[2], axis=0)
-
-    inputs = (images[1:-1], actions[1:-1], rewards[1:-1])
-
-    outputs = tf.scan(process_single_frame, inputs, initializers)
-    gen_images, gen_rewards = outputs[1:3]
-
-    gen_images = tf.concat((first_gen_images, gen_images), axis=0)
-    gen_rewards = tf.concat((first_gen_rewards, gen_rewards), axis=0)
-
-    return gen_images, gen_rewards, [latent_mean], [latent_std]
-
-  def get_scheduled_sample_func(self, batch_size):
-    """Creates a function for scheduled sampling based on given hparams."""
-    with tf.variable_scope("scheduled_sampling_func", reuse=False):
-      iter_num = tf.train.get_global_step()
-      # TODO(lukaszkaiser): figure out why iter_num can be None.
-      if iter_num is None:
-        iter_num = _LARGE_STEP_NUMBER
-
-      if self.hparams.scheduled_sampling_mode == "prob":
-        decay_steps = self.hparams.scheduled_sampling_decay_steps
-        probability = tf.train.polynomial_decay(
-            1.0, iter_num, decay_steps, 0.0)
-        scheduled_sampling_func = common_video.scheduled_sample_prob
-        scheduled_sampling_func_var = probability
-      else:
-        # Calculate number of ground-truth frames to pass in.
-        k = self.hparams.scheduled_sampling_k
-        num_ground_truth = tf.to_int32(
-            tf.round(
-                tf.to_float(batch_size) *
-                (k / (k + tf.exp(tf.to_float(iter_num) / tf.to_float(k))))))
-        scheduled_sampling_func = common_video.scheduled_sample_count
-        scheduled_sampling_func_var = num_ground_truth
-
-      tf.summary.scalar("scheduled_sampling_var", scheduled_sampling_func_var)
-      partial_func = partial(scheduled_sampling_func,
-                             batch_size=batch_size,
-                             scheduled_sample_var=scheduled_sampling_func_var)
-      return partial_func
-
-  def get_scheduled_sample_inputs(self,
-                                  done_warm_start,
-                                  groundtruth_items,
-                                  generated_items,
-                                  scheduled_sampling_func):
-    """Scheduled sampling.
-
-    Args:
-      done_warm_start: whether we are done with warm start or not.
-      groundtruth_items: list of ground truth items.
-      generated_items: list of generated items.
-      scheduled_sampling_func: scheduled sampling function to choose between
-        groundtruth items and generated items.
-
-    Returns:
-      A mix list of ground truth and generated items.
-    """
-    def sample():
-      """Calculate the scheduled sampling params based on iteration number."""
-      with tf.variable_scope("scheduled_sampling", reuse=tf.AUTO_REUSE):
-        output_items = []
-        for item_gt, item_gen in zip(groundtruth_items, generated_items):
-          output_items.append(scheduled_sampling_func(item_gt, item_gen))
-        return output_items
-
-    cases = {
-        tf.logical_not(done_warm_start): lambda: groundtruth_items,
-        tf.logical_not(self.is_training): lambda: generated_items,
-    }
-    output_items = tf.case(cases, default=sample, strict=True)
-
-    return output_items
-
-  def get_input_if_exists(self, features, key, batch_size, num_frames):
-    if key in features:
-      x = features[key]
-    else:
-      x = tf.zeros((batch_size, num_frames, 1, self.hparams.hidden_size))
-    return self.swap_time_and_batch_axes(x)
-
-  def swap_time_and_batch_axes(self, x):
-    transposed_axes = tf.concat([[1, 0], tf.range(2, tf.rank(x))], axis=0)
-    return tf.transpose(x, transposed_axes)
-
-  def body(self, features):
-    hparams = self.hparams
-    batch_size = common_layers.shape_list(features["inputs"])[0]
-
-    # Swap time and batch axes.
-    input_frames = self.swap_time_and_batch_axes(features["inputs"])
-    target_frames = self.swap_time_and_batch_axes(features["targets"])
-
-    # Get actions if exist otherwise use zeros
-    input_actions = self.get_input_if_exists(
-        features, "input_action", batch_size, hparams.video_num_input_frames)
-    target_actions = self.get_input_if_exists(
-        features, "target_action", batch_size, hparams.video_num_target_frames)
-
-    # Get rewards if exist otherwise use zeros
-    input_rewards = self.get_input_if_exists(
-        features, "input_reward", batch_size, hparams.video_num_input_frames)
-    target_rewards = self.get_input_if_exists(
-        features, "target_reward", batch_size, hparams.video_num_target_frames)
-
-    all_actions = tf.concat([input_actions, target_actions], axis=0)
-    all_rewards = tf.concat([input_rewards, target_rewards], axis=0)
-    all_frames = tf.concat([input_frames, target_frames], axis=0)
-
-    # Each image is being used twice, in latent tower and main tower.
-    # This is to make sure we are using the *same* image for both, ...
-    # ... given how TF queues work.
-    # NOT sure if this is required at all. Doesn"t hurt though! :)
-    all_frames = tf.identity(all_frames)
-
-    gen_images, gen_rewards, latent_means, latent_stds = self.construct_model(
-        images=all_frames,
-        actions=all_actions,
-        rewards=all_rewards,
-    )
-
-    step_num = tf.train.get_global_step()
-    # TODO(mbz): what should it be if it"s undefined?
-    if step_num is None:
-      step_num = _LARGE_STEP_NUMBER
-
-    schedule = self.hparams.latent_loss_multiplier_schedule
-    second_stage = self.hparams.num_iterations_2nd_stage
-    # TODO(mechcoder): Add log_annealing schedule.
-    if schedule == "constant":
-      beta = tf.cond(tf.greater(step_num, second_stage),
-                     lambda: self.hparams.latent_loss_multiplier,
-                     lambda: 0.0)
-    elif schedule == "linear_anneal":
-      # Linearly anneal beta from 0.0 to self.hparams.latent_loss_multiplier.
-      # between self.hparams.num_iterations_2nd_stage to anneal_end.
-      # beta = latent_loss * (1 - (global_step - 2nd_stage) / (anneal_end - 2nd_stage))  # pylint:disable=line-too-long
-      anneal_end = self.hparams.anneal_end
-      latent_multiplier = self.hparams.latent_loss_multiplier
-      if anneal_end < second_stage:
-        raise ValueError("Expected hparams.num_iterations_2nd_stage < "
-                         "hparams.anneal_end %d, got %d." %
-                         (second_stage, anneal_end))
-
-      def anneal_loss(step_num):
-        step_num = tf.cast(step_num, dtype=tf.float32)
-        fraction = (float(anneal_end) - step_num) / (anneal_end - second_stage)
-        return self.hparams.latent_loss_multiplier * (1 - fraction)
-
-      beta = tf.case(
-          pred_fn_pairs={
-              tf.less(step_num, second_stage): lambda: 0.0,
-              tf.greater(step_num, anneal_end): lambda: latent_multiplier},
-          default=lambda: anneal_loss(step_num))
-
-    kl_loss = 0.0
-    if self.is_training:
-      for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
-        kl_loss += common_layers.kl_divergence(mean, std)
-        tf.summary.histogram("posterior_mean_%d" % i, mean)
-        tf.summary.histogram("posterior_std_%d" % i, std)
-
-      tf.summary.scalar("beta", beta)
-      tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
-
-    extra_loss = beta * kl_loss
-
-    # Ignore the predictions from the input frames.
-    # This is NOT the same as original paper/implementation.
-    predictions = gen_images[hparams.video_num_input_frames-1:]
-    reward_pred = gen_rewards[hparams.video_num_input_frames-1:]
-    reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove undeeded dimension.
-
-    # TODO(mbz): clean this up!
-    def fix_video_dims_and_concat_on_x_axis(x):
-      x = tf.transpose(x, [1, 3, 4, 0, 2])
-      x = tf.reshape(x, [batch_size, 64, 3, -1])
-      x = tf.transpose(x, [0, 3, 1, 2])
-      return x
-
-    frames_gd = fix_video_dims_and_concat_on_x_axis(target_frames)
-    frames_pd = fix_video_dims_and_concat_on_x_axis(predictions)
-    side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
-    tf.summary.image("full_video", side_by_side_video)
-
-    # Swap back time and batch axes.
-    predictions = self.swap_time_and_batch_axes(predictions)
-    reward_pred = self.swap_time_and_batch_axes(reward_pred)
-
-    return_targets = predictions
-    if "target_reward" in features:
-      return_targets = {"targets": predictions, "target_reward": reward_pred}
-
-    return return_targets, extra_loss
-
-
-@registry.register_model
-class NextFrameStochasticTwoFrames(NextFrameStochastic):
-  """Stochastic next-frame model with 2 frames posterior."""
-
-  def construct_model(self, images, actions, rewards):
-    images = tf.unstack(images, axis=0)
-    actions = tf.unstack(actions, axis=0)
-    rewards = tf.unstack(rewards, axis=0)
-
-    batch_size = common_layers.shape_list(images[0])[0]
-    context_frames = self.hparams.video_num_input_frames
-
-    # Predicted images and rewards.
-    gen_rewards, gen_images, latent_means, latent_stds = [], [], [], []
-
-    # LSTM states.
-    lstm_state = [None] * 7
-    reward_lstm_state = [None] * 5
-
-    # Create scheduled sampling function
-    ss_func = self.get_scheduled_sample_func(batch_size)
-
-    pred_image, pred_reward, latent = images[0], rewards[0], None
-    for timestep, image, action, reward in zip(
-        range(len(images)-1), images[:-1], actions[:-1], rewards[:-1]):
-      # Scheduled Sampling
-      done_warm_start = timestep > context_frames - 1
-      groundtruth_items = [image, reward]
-      generated_items = [pred_image, pred_reward]
-      input_image, input_reward = self.get_scheduled_sample_inputs(
-          done_warm_start, groundtruth_items, generated_items, ss_func)
-
-      # Latent
-      # TODO(mbz): should we use input_image iunstead of image?
-      latent_images = [image, images[timestep+1]]
-      latent_mean, latent_std = self.construct_latent_tower(latent_images)
-      latent = self.get_gaussian_latent(latent_mean, latent_std)
-      latent_means.append(latent_mean)
-      latent_stds.append(latent_std)
-
-      # Prediction
-      pred_image, lstm_state = self.construct_predictive_tower(
-          input_image, input_reward, action, lstm_state, latent)
-
-      if self.hparams.reward_prediction:
-        pred_reward, reward_lstm_state = self.reward_prediction(
-            input_image, input_reward, action, reward_lstm_state, latent)
-      else:
-        pred_reward = input_reward
-
-      gen_images.append(pred_image)
-      gen_rewards.append(pred_reward)
-
-    gen_images = tf.stack(gen_images, axis=0)
-    gen_rewards = tf.stack(gen_rewards, axis=0)
-
-    return gen_images, gen_rewards, latent_means, latent_stds
-
-
-@registry.register_model
-class NextFrameStochasticEmily(NextFrameStochastic):
-  """Model architecture for video prediction model.
-
-     based on following paper:
-     "Stochastic Video Generation with a Learned Prior"
-     https://arxiv.org/pdf/1802.07687.pdf
-     by Emily Denton and Rob Fergus.
-
-     This code is a translation of the original code from PyTorch:
-     https://github.com/edenton/svg
-  """
-
-  def encoder(self, inputs, nout):
-    """VGG based image encoder.
-
-    Args:
-      inputs: image tensor with size BSx64x64xC
-      nout: number of output channels
-    Returns:
-      net: encoded image with size BSxNout
-      skips: skip connection after each layer
-    """
-    vgg_layer = common_video.vgg_layer
-    net01 = inputs
-    # h1
-    net11 = tfcl.repeat(net01, 2, vgg_layer, 64,
-                        scope="h1", is_training=self.is_training)
-    net12 = tfl.max_pooling2d(net11, [2, 2], strides=(2, 2), name="h1_pool")
-    # h2
-    net21 = tfcl.repeat(net12, 2, vgg_layer, 128,
-                        scope="h2", is_training=self.is_training)
-    net22 = tfl.max_pooling2d(net21, [2, 2], strides=(2, 2), name="h2_pool")
-    # h3
-    net31 = tfcl.repeat(net22, 3, vgg_layer, 256,
-                        scope="h3", is_training=self.is_training)
-    net32 = tfl.max_pooling2d(net31, [2, 2], strides=(2, 2), name="h3_pool")
-    # h4
-    net41 = tfcl.repeat(net32, 3, vgg_layer, 512,
-                        scope="h4", is_training=self.is_training)
-    net42 = tfl.max_pooling2d(net41, [2, 2], strides=(2, 2), name="h4_pool")
-    # h5
-    net51 = tfcl.repeat(net42, 1, vgg_layer, nout,
-                        kernel_size=4, padding="VALID", activation=tf.tanh,
-                        scope="h5", is_training=self.is_training)
-    skips = [net11, net21, net31, net41]
-    return net51, skips
-
-  def decoder(self, inputs, skips, nout):
-    """VGG based image decoder.
-
-    Args:
-      inputs: image tensor with size BSxX
-      skips: skip connections from encoder
-      nout: number of output channels
-    Returns:
-      net: decoded image with size BSx64x64xNout
-      skips: skip connection after each layer
-    """
-    vgg_layer = common_video.vgg_layer
-    net = inputs
-    # d1
-    net = tfl.conv2d_transpose(net, 512, kernel_size=4, padding="VALID",
-                               name="d1_deconv", activation=None)
-    net = tfl.batch_normalization(net, training=self.is_training, name="d1_bn")
-    net = tf.nn.leaky_relu(net)
-    net = common_layers.upscale(net, 2)
-    # d2
-    net = tf.concat([net, skips[3]], axis=3)
-    net = tfcl.repeat(net, 2, vgg_layer, 512, scope="d2a")
-    net = tfcl.repeat(net, 1, vgg_layer, 256, scope="d2b")
-    net = common_layers.upscale(net, 2)
-    # d3
-    net = tf.concat([net, skips[2]], axis=3)
-    net = tfcl.repeat(net, 2, vgg_layer, 256, scope="d3a")
-    net = tfcl.repeat(net, 1, vgg_layer, 128, scope="d3b")
-    net = common_layers.upscale(net, 2)
-    # d4
-    net = tf.concat([net, skips[1]], axis=3)
-    net = tfcl.repeat(net, 1, vgg_layer, 128, scope="d4a")
-    net = tfcl.repeat(net, 1, vgg_layer, 64, scope="d4b")
-    net = common_layers.upscale(net, 2)
-    # d5
-    net = tf.concat([net, skips[0]], axis=3)
-    net = tfcl.repeat(net, 1, vgg_layer, 64, scope="d5")
-    net = tfl.conv2d_transpose(net, nout, kernel_size=3, padding="SAME",
-                               name="d6_deconv", activation=tf.sigmoid)
-    return net
-
-  def stacked_lstm(self, inputs, states, hidden_size, output_size, nlayers):
-    """Stacked LSTM layers with FC layers as input and output embeddings.
-
-    Args:
-      inputs: input tensor
-      states: a list of internal lstm states for each layer
-      hidden_size: number of lstm units
-      output_size: size of the output
-      nlayers: number of lstm layers
-    Returns:
-      net: output of the network
-      skips: a list of updated lstm states for each layer
-    """
-    net = inputs
-    net = tfl.dense(
-        net, hidden_size, activation=None, name="af1")
-    for i in range(nlayers):
-      net, states[i] = common_video.basic_lstm(
-          net, states[i], hidden_size, name="alstm%d"%i)
-    net = tfl.dense(
-        net, output_size, activation=tf.nn.tanh, name="af2")
-    return net, states
-
-  def lstm_gaussian(self, inputs, states, hidden_size, output_size, nlayers):
-    """Stacked LSTM layers with FC layer as input and gaussian as output.
-
-    Args:
-      inputs: input tensor
-      states: a list of internal lstm states for each layer
-      hidden_size: number of lstm units
-      output_size: size of the output
-      nlayers: number of lstm layers
-    Returns:
-      mu: mean of the predicted gaussian
-      logvar: log(var) of the predicted gaussian
-      skips: a list of updated lstm states for each layer
-    """
-    net = inputs
-    net = tfl.dense(net, hidden_size, activation=None, name="bf1")
-    for i in range(nlayers):
-      net, states[i] = common_video.basic_lstm(
-          net, states[i], hidden_size, name="blstm%d"%i)
-    mu = tfl.dense(net, output_size, activation=None, name="bf2mu")
-    logvar = tfl.dense(net, output_size, activation=None, name="bf2log")
-    return mu, logvar, states
-
-  def construct_model(self, images, actions, rewards):
-    """Builds the stochastic model.
-
-    The model first encodes all the images (x_t) in the sequence
-    using the encoder. Let"s call the output e_t. Then it predicts the
-    latent state of the next frame using a recurrent posterior network
-    z ~ q(z|e_{0:t}) = N(mu(e_{0:t}), sigma(e_{0:t})).
-    Another recurrent network predicts the embedding of the next frame
-    using the approximated posterior e_{t+1} = p(e_{t+1}|e_{0:t}, z)
-    Finally, the decoder decodes e_{t+1} into x_{t+1}.
-    Skip connections from encoder to decoder help with reconstruction.
-
-    Args:
-      images: tensor of ground truth image sequences
-      actions: NOT used list of action tensors
-      rewards: NOT used list of reward tensors
-
-    Returns:
-      gen_images: generated images
-      fakr_rewards: input rewards as reward prediction!
-      pred_mu: predited means of posterior
-      pred_logvar: predicted log(var) of posterior
-    """
-    # model does not support action conditioned and reward prediction
-    fake_reward_prediction = rewards
-    del actions, rewards
-
-    z_dim = self.hparams.z_dim
-    g_dim = self.hparams.g_dim
-    rnn_size = self.hparams.rnn_size
-    posterior_rnn_layers = self.hparams.posterior_rnn_layers
-    predictor_rnn_layers = self.hparams.predictor_rnn_layers
-    context_frames = self.hparams.video_num_input_frames
-
-    seq_len, batch_size, _, _, color_channels = common_layers.shape_list(images)
-
-    # LSTM initial sizesstates.
-    predictor_states = [None] * predictor_rnn_layers
-    posterior_states = [None] * posterior_rnn_layers
-
-    tf.logging.info(">>>> Encoding")
-    # Encoding:
-    enc_images, enc_skips = [], []
-    images = tf.unstack(images, axis=0)
-    for i, image in enumerate(images):
-      with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
-        enc, skips = self.encoder(image, rnn_size)
-        enc = tfcl.flatten(enc)
-        enc_images.append(enc)
-        enc_skips.append(skips)
-
-    tf.logging.info(">>>> Prediction")
-    # Prediction
-    pred_enc, pred_mu, pred_logvar = [], [], []
-    for i in range(1, seq_len):
-      with tf.variable_scope("prediction", reuse=tf.AUTO_REUSE):
-        # current encoding
-        h_current = enc_images[i-1]
-        # target encoding
-        h_target = enc_images[i]
-
-        z = tf.random_normal([batch_size, z_dim], 0, 1, dtype=tf.float32)
-        mu, logvar = tf.zeros_like(z), tf.zeros_like(z)
-
-        # Only use Posterior if it's training time
-        if self.hparams.mode == tf.estimator.ModeKeys.TRAIN:
-          mu, logvar, posterior_states = self.lstm_gaussian(
-              h_target, posterior_states, rnn_size, z_dim, posterior_rnn_layers)
-
-          # The original implementation has a multiplier of 0.5
-          # Removed here for simplicity i.e. replacing var with std
-          z = z * tf.exp(logvar) + mu
-
-        # Predict output encoding
-        h_pred, predictor_states = self.stacked_lstm(
-            tf.concat([h_current, z], axis=1),
-            predictor_states, rnn_size, g_dim, predictor_rnn_layers)
-
-        pred_enc.append(h_pred)
-        pred_mu.append(mu)
-        pred_logvar.append(logvar)
-
-    tf.logging.info(">>>> Decoding")
-    # Decoding
-    gen_images = []
-    for i in range(seq_len-1):
-      with tf.variable_scope("decoding", reuse=tf.AUTO_REUSE):
-        # use skip values of last available frame
-        skip_index = min(context_frames-1, i)
-
-        h_pred = tf.reshape(pred_enc[i], [batch_size, 1, 1, g_dim])
-        x_pred = self.decoder(h_pred, enc_skips[skip_index], color_channels)
-        gen_images.append(x_pred)
-
-    tf.logging.info(">>>> Done")
-    gen_images = tf.stack(gen_images, axis=0)
-    return gen_images, fake_reward_prediction, pred_mu, pred_logvar
diff --git a/tensor2tensor/models/research/next_frame_emily.py b/tensor2tensor/models/research/next_frame_emily.py
new file mode 100644
index 000000000..ff2275ede
--- /dev/null
+++ b/tensor2tensor/models/research/next_frame_emily.py
@@ -0,0 +1,263 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Model architecture for video prediction model.
+
+   based on following paper:
+   "Stochastic Video Generation with a Learned Prior"
+   https://arxiv.org/pdf/1802.07687.pdf
+   by Emily Denton and Rob Fergus.
+
+   This code is a translation of the original code from PyTorch:
+   https://github.com/edenton/svg
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
+from tensor2tensor.models.research import next_frame_params  # pylint: disable=unused-import
+from tensor2tensor.models.research import next_frame_sv2p
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+tfl = tf.layers
+tfcl = tf.contrib.layers
+
+
+@registry.register_model
+class NextFrameStochasticEmily(next_frame_sv2p.NextFrameStochastic):
+  """Stochastic Variational Video Prediction Without Learned Prior."""
+
+  def encoder(self, inputs, nout):
+    """VGG based image encoder.
+
+    Args:
+      inputs: image tensor with size BSx64x64xC
+      nout: number of output channels
+    Returns:
+      net: encoded image with size BSxNout
+      skips: skip connection after each layer
+    """
+    vgg_layer = common_video.vgg_layer
+    net01 = inputs
+    # h1
+    net11 = tfcl.repeat(net01, 2, vgg_layer, 64,
+                        scope="h1", is_training=self.is_training)
+    net12 = tfl.max_pooling2d(net11, [2, 2], strides=(2, 2), name="h1_pool")
+    # h2
+    net21 = tfcl.repeat(net12, 2, vgg_layer, 128,
+                        scope="h2", is_training=self.is_training)
+    net22 = tfl.max_pooling2d(net21, [2, 2], strides=(2, 2), name="h2_pool")
+    # h3
+    net31 = tfcl.repeat(net22, 3, vgg_layer, 256,
+                        scope="h3", is_training=self.is_training)
+    net32 = tfl.max_pooling2d(net31, [2, 2], strides=(2, 2), name="h3_pool")
+    # h4
+    net41 = tfcl.repeat(net32, 3, vgg_layer, 512,
+                        scope="h4", is_training=self.is_training)
+    net42 = tfl.max_pooling2d(net41, [2, 2], strides=(2, 2), name="h4_pool")
+    # h5
+    net51 = tfcl.repeat(net42, 1, vgg_layer, nout,
+                        kernel_size=4, padding="VALID", activation=tf.tanh,
+                        scope="h5", is_training=self.is_training)
+    skips = [net11, net21, net31, net41]
+    return net51, skips
+
+  def decoder(self, inputs, skips, nout):
+    """VGG based image decoder.
+
+    Args:
+      inputs: image tensor with size BSxX
+      skips: skip connections from encoder
+      nout: number of output channels
+    Returns:
+      net: decoded image with size BSx64x64xNout
+      skips: skip connection after each layer
+    """
+    vgg_layer = common_video.vgg_layer
+    net = inputs
+    # d1
+    net = tfl.conv2d_transpose(net, 512, kernel_size=4, padding="VALID",
+                               name="d1_deconv", activation=None)
+    net = tfl.batch_normalization(net, training=self.is_training, name="d1_bn")
+    net = tf.nn.leaky_relu(net)
+    net = common_layers.upscale(net, 2)
+    # d2
+    net = tf.concat([net, skips[3]], axis=3)
+    net = tfcl.repeat(net, 2, vgg_layer, 512, scope="d2a")
+    net = tfcl.repeat(net, 1, vgg_layer, 256, scope="d2b")
+    net = common_layers.upscale(net, 2)
+    # d3
+    net = tf.concat([net, skips[2]], axis=3)
+    net = tfcl.repeat(net, 2, vgg_layer, 256, scope="d3a")
+    net = tfcl.repeat(net, 1, vgg_layer, 128, scope="d3b")
+    net = common_layers.upscale(net, 2)
+    # d4
+    net = tf.concat([net, skips[1]], axis=3)
+    net = tfcl.repeat(net, 1, vgg_layer, 128, scope="d4a")
+    net = tfcl.repeat(net, 1, vgg_layer, 64, scope="d4b")
+    net = common_layers.upscale(net, 2)
+    # d5
+    net = tf.concat([net, skips[0]], axis=3)
+    net = tfcl.repeat(net, 1, vgg_layer, 64, scope="d5")
+    net = tfl.conv2d_transpose(net, nout, kernel_size=3, padding="SAME",
+                               name="d6_deconv", activation=tf.sigmoid)
+    return net
+
+  def stacked_lstm(self, inputs, states, hidden_size, output_size, nlayers):
+    """Stacked LSTM layers with FC layers as input and output embeddings.
+
+    Args:
+      inputs: input tensor
+      states: a list of internal lstm states for each layer
+      hidden_size: number of lstm units
+      output_size: size of the output
+      nlayers: number of lstm layers
+    Returns:
+      net: output of the network
+      skips: a list of updated lstm states for each layer
+    """
+    net = inputs
+    net = tfl.dense(
+        net, hidden_size, activation=None, name="af1")
+    for i in range(nlayers):
+      net, states[i] = common_video.basic_lstm(
+          net, states[i], hidden_size, name="alstm%d"%i)
+    net = tfl.dense(
+        net, output_size, activation=tf.nn.tanh, name="af2")
+    return net, states
+
+  def lstm_gaussian(self, inputs, states, hidden_size, output_size, nlayers):
+    """Stacked LSTM layers with FC layer as input and gaussian as output.
+
+    Args:
+      inputs: input tensor
+      states: a list of internal lstm states for each layer
+      hidden_size: number of lstm units
+      output_size: size of the output
+      nlayers: number of lstm layers
+    Returns:
+      mu: mean of the predicted gaussian
+      logvar: log(var) of the predicted gaussian
+      skips: a list of updated lstm states for each layer
+    """
+    net = inputs
+    net = tfl.dense(net, hidden_size, activation=None, name="bf1")
+    for i in range(nlayers):
+      net, states[i] = common_video.basic_lstm(
+          net, states[i], hidden_size, name="blstm%d"%i)
+    mu = tfl.dense(net, output_size, activation=None, name="bf2mu")
+    logvar = tfl.dense(net, output_size, activation=None, name="bf2log")
+    return mu, logvar, states
+
+  def construct_model(self, images, actions, rewards):
+    """Builds the stochastic model.
+
+    The model first encodes all the images (x_t) in the sequence
+    using the encoder. Let"s call the output e_t. Then it predicts the
+    latent state of the next frame using a recurrent posterior network
+    z ~ q(z|e_{0:t}) = N(mu(e_{0:t}), sigma(e_{0:t})).
+    Another recurrent network predicts the embedding of the next frame
+    using the approximated posterior e_{t+1} = p(e_{t+1}|e_{0:t}, z)
+    Finally, the decoder decodes e_{t+1} into x_{t+1}.
+    Skip connections from encoder to decoder help with reconstruction.
+
+    Args:
+      images: tensor of ground truth image sequences
+      actions: NOT used list of action tensors
+      rewards: NOT used list of reward tensors
+
+    Returns:
+      gen_images: generated images
+      fakr_rewards: input rewards as reward prediction!
+      pred_mu: predited means of posterior
+      pred_logvar: predicted log(var) of posterior
+    """
+    # model does not support action conditioned and reward prediction
+    fake_reward_prediction = rewards
+    del actions, rewards
+
+    z_dim = self.hparams.z_dim
+    g_dim = self.hparams.g_dim
+    rnn_size = self.hparams.rnn_size
+    posterior_rnn_layers = self.hparams.posterior_rnn_layers
+    predictor_rnn_layers = self.hparams.predictor_rnn_layers
+    context_frames = self.hparams.video_num_input_frames
+
+    seq_len, batch_size, _, _, color_channels = common_layers.shape_list(images)
+
+    # LSTM initial sizesstates.
+    predictor_states = [None] * predictor_rnn_layers
+    posterior_states = [None] * posterior_rnn_layers
+
+    tf.logging.info(">>>> Encoding")
+    # Encoding:
+    enc_images, enc_skips = [], []
+    images = tf.unstack(images, axis=0)
+    for i, image in enumerate(images):
+      with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
+        enc, skips = self.encoder(image, rnn_size)
+        enc = tfcl.flatten(enc)
+        enc_images.append(enc)
+        enc_skips.append(skips)
+
+    tf.logging.info(">>>> Prediction")
+    # Prediction
+    pred_enc, pred_mu, pred_logvar = [], [], []
+    for i in range(1, seq_len):
+      with tf.variable_scope("prediction", reuse=tf.AUTO_REUSE):
+        # current encoding
+        h_current = enc_images[i-1]
+        # target encoding
+        h_target = enc_images[i]
+
+        z = tf.random_normal([batch_size, z_dim], 0, 1, dtype=tf.float32)
+        mu, logvar = tf.zeros_like(z), tf.zeros_like(z)
+
+        # Only use Posterior if it's training time
+        if self.hparams.mode == tf.estimator.ModeKeys.TRAIN:
+          mu, logvar, posterior_states = self.lstm_gaussian(
+              h_target, posterior_states, rnn_size, z_dim, posterior_rnn_layers)
+
+          # The original implementation has a multiplier of 0.5
+          # Removed here for simplicity i.e. replacing var with std
+          z = z * tf.exp(logvar) + mu
+
+        # Predict output encoding
+        h_pred, predictor_states = self.stacked_lstm(
+            tf.concat([h_current, z], axis=1),
+            predictor_states, rnn_size, g_dim, predictor_rnn_layers)
+
+        pred_enc.append(h_pred)
+        pred_mu.append(mu)
+        pred_logvar.append(logvar)
+
+    tf.logging.info(">>>> Decoding")
+    # Decoding
+    gen_images = []
+    for i in range(seq_len-1):
+      with tf.variable_scope("decoding", reuse=tf.AUTO_REUSE):
+        # use skip values of last available frame
+        skip_index = min(context_frames-1, i)
+
+        h_pred = tf.reshape(pred_enc[i], [batch_size, 1, 1, g_dim])
+        x_pred = self.decoder(h_pred, enc_skips[skip_index], color_channels)
+        gen_images.append(x_pred)
+
+    tf.logging.info(">>>> Done")
+    gen_images = tf.stack(gen_images, axis=0)
+    return gen_images, fake_reward_prediction, pred_mu, pred_logvar
diff --git a/tensor2tensor/models/research/next_frame_savp.py b/tensor2tensor/models/research/next_frame_savp.py
index 8681ec32e..ad3044d91 100644
--- a/tensor2tensor/models/research/next_frame_savp.py
+++ b/tensor2tensor/models/research/next_frame_savp.py
@@ -23,14 +23,14 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
-from tensor2tensor.models.research import next_frame
 from tensor2tensor.models.research import next_frame_params  # pylint: disable=unused-import
+from tensor2tensor.models.research import next_frame_sv2p
 from tensor2tensor.utils import registry
 import tensorflow as tf
 
 
 @registry.register_model
-class NextFrameSAVP(next_frame.NextFrameStochastic):
+class NextFrameSAVP(next_frame_sv2p.NextFrameStochastic):
   """Stochastic Adversarial Video Prediction."""
 
   def encoder(self, inputs, n_layers=3):
diff --git a/tensor2tensor/models/research/next_frame_savp_test.py b/tensor2tensor/models/research/next_frame_savp_test.py
deleted file mode 100644
index 9facd903d..000000000
--- a/tensor2tensor/models/research/next_frame_savp_test.py
+++ /dev/null
@@ -1,88 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tests for the SAVP model."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-import numpy as np
-
-from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
-from tensor2tensor.models.research import next_frame_params
-from tensor2tensor.models.research import next_frame_savp
-from tensor2tensor.utils import registry
-
-import tensorflow as tf
-
-
-class NextFrameSAVPTest(tf.test.TestCase):
-
-  def TestVideoModel(self,
-                     in_frames,
-                     out_frames,
-                     hparams,
-                     model,
-                     expected_last_dim,
-                     upsample_method="conv2d_transpose"):
-
-    x = np.random.random_integers(0, high=255, size=(8, in_frames, 64, 64, 3))
-    y = np.random.random_integers(0, high=255, size=(8, out_frames, 64, 64, 3))
-
-    hparams.video_num_input_frames = in_frames
-    hparams.video_num_target_frames = out_frames
-    hparams.upsample_method = upsample_method
-    problem = registry.problem("video_stochastic_shapes10k")
-    p_hparams = problem.get_hparams(hparams)
-    hparams.problem = problem
-    hparams.problem_hparams = p_hparams
-
-    with self.test_session() as session:
-      features = {
-          "inputs": tf.constant(x, dtype=tf.int32),
-          "targets": tf.constant(y, dtype=tf.int32),
-      }
-      model = model(
-          hparams, tf.estimator.ModeKeys.TRAIN)
-      logits, _ = model(features)
-      session.run(tf.global_variables_initializer())
-      res = session.run(logits)
-    expected_shape = y.shape + (expected_last_dim,)
-    self.assertEqual(res.shape, expected_shape)
-
-  def TestOnVariousInputOutputSizes(self, hparams, model, expected_last_dim):
-    self.TestVideoModel(1, 1, hparams, model, expected_last_dim)
-    self.TestVideoModel(1, 6, hparams, model, expected_last_dim)
-    self.TestVideoModel(4, 1, hparams, model, expected_last_dim)
-    self.TestVideoModel(7, 5, hparams, model, expected_last_dim)
-
-  def TestOnVariousUpSampleLayers(self, hparams, model, expected_last_dim):
-    self.TestVideoModel(4, 1, hparams, model, expected_last_dim,
-                        upsample_method="bilinear_upsample_conv")
-    self.TestVideoModel(4, 1, hparams, model, expected_last_dim,
-                        upsample_method="nn_upsample_conv")
-
-  def testStochasticSavp(self):
-    self.TestOnVariousInputOutputSizes(
-        next_frame_params.next_frame_savp(),
-        next_frame_savp.NextFrameSAVP,
-        1)
-    self.TestOnVariousUpSampleLayers(
-        next_frame_params.next_frame_savp(),
-        next_frame_savp.NextFrameSAVP,
-        1)
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
new file mode 100644
index 000000000..57b5e8246
--- /dev/null
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -0,0 +1,671 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""SV2P: Stochastic Variational Video Prediction.
+
+   based on the following paper:
+   https://arxiv.org/abs/1710.11252
+   by Mohammad Babaeizadeh, Chelsea Finn, Dumitru Erhan,
+      Roy H. Campbell and Sergey Levine
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from functools import partial
+
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
+from tensor2tensor.models.research import next_frame
+from tensor2tensor.models.research import next_frame_params  # pylint: disable=unused-import
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+tfl = tf.layers
+tfcl = tf.contrib.layers
+
+_LARGE_STEP_NUMBER = 100000
+
+
+@registry.register_model
+class NextFrameStochastic(next_frame.NextFrameBasic):
+  """Stochastic Variational Video Prediction."""
+
+  @property
+  def is_training(self):
+    return self.hparams.mode == tf.estimator.ModeKeys.TRAIN
+
+  def tinyify(self, array):
+    if self.hparams.tiny_mode:
+      return [1 for _ in array]
+    return array
+
+  def get_gaussian_latent(self, latent_mean, latent_std):
+    latent = tf.random_normal(tf.shape(latent_mean), 0, 1, dtype=tf.float32)
+    latent = latent_mean + tf.exp(latent_std / 2.0) * latent
+    return latent
+
+  def get_scheduled_sample_func(self, batch_size):
+    """Creates a function for scheduled sampling based on given hparams."""
+    with tf.variable_scope("scheduled_sampling_func", reuse=False):
+      iter_num = tf.train.get_global_step()
+      # TODO(lukaszkaiser): figure out why iter_num can be None.
+      if iter_num is None:
+        iter_num = _LARGE_STEP_NUMBER
+
+      if self.hparams.scheduled_sampling_mode == "prob":
+        decay_steps = self.hparams.scheduled_sampling_decay_steps
+        probability = tf.train.polynomial_decay(
+            1.0, iter_num, decay_steps, 0.0)
+        scheduled_sampling_func = common_video.scheduled_sample_prob
+        scheduled_sampling_func_var = probability
+      else:
+        # Calculate number of ground-truth frames to pass in.
+        k = self.hparams.scheduled_sampling_k
+        num_ground_truth = tf.to_int32(
+            tf.round(
+                tf.to_float(batch_size) *
+                (k / (k + tf.exp(tf.to_float(iter_num) / tf.to_float(k))))))
+        scheduled_sampling_func = common_video.scheduled_sample_count
+        scheduled_sampling_func_var = num_ground_truth
+
+      tf.summary.scalar("scheduled_sampling_var", scheduled_sampling_func_var)
+      partial_func = partial(scheduled_sampling_func,
+                             batch_size=batch_size,
+                             scheduled_sample_var=scheduled_sampling_func_var)
+      return partial_func
+
+  def get_scheduled_sample_inputs(self,
+                                  done_warm_start,
+                                  groundtruth_items,
+                                  generated_items,
+                                  scheduled_sampling_func):
+    """Scheduled sampling.
+
+    Args:
+      done_warm_start: whether we are done with warm start or not.
+      groundtruth_items: list of ground truth items.
+      generated_items: list of generated items.
+      scheduled_sampling_func: scheduled sampling function to choose between
+        groundtruth items and generated items.
+
+    Returns:
+      A mix list of ground truth and generated items.
+    """
+    def sample():
+      """Calculate the scheduled sampling params based on iteration number."""
+      with tf.variable_scope("scheduled_sampling", reuse=tf.AUTO_REUSE):
+        output_items = []
+        for item_gt, item_gen in zip(groundtruth_items, generated_items):
+          output_items.append(scheduled_sampling_func(item_gt, item_gen))
+        return output_items
+
+    cases = {
+        tf.logical_not(done_warm_start): lambda: groundtruth_items,
+        tf.logical_not(self.is_training): lambda: generated_items,
+    }
+    output_items = tf.case(cases, default=sample, strict=True)
+
+    return output_items
+
+  def get_input_if_exists(self, features, key, batch_size, num_frames):
+    if key in features:
+      x = features[key]
+    else:
+      x = tf.zeros((batch_size, num_frames, 1, self.hparams.hidden_size))
+    return common_video.swap_time_and_batch_axes(x)
+
+  def bottom_part_tower(self, input_image, input_reward, action, latent,
+                        lstm_state, lstm_size, conv_size, concat_latent=False):
+    """The bottom part of predictive towers.
+
+    With the current (early) design, the main prediction tower and
+    the reward prediction tower share the same arcitecture. TF Scope can be
+    adjusted as required to either share or not share the weights between
+    the two towers.
+
+    Args:
+      input_image: the current image.
+      input_reward: the current reward.
+      action: the action taken by the agent.
+      latent: the latent vector.
+      lstm_state: the current internal states of conv lstms.
+      lstm_size: the size of lstms.
+      conv_size: the size of convolutions.
+      concat_latent: whether or not to concatenate the latent at every step.
+
+    Returns:
+      - the output of the partial network.
+      - intermidate outputs for skip connections.
+    """
+    lstm_func = common_video.conv_lstm_2d
+    tile_and_concat = common_video.tile_and_concat
+
+    input_image = common_layers.make_even_size(input_image)
+    concat_input_image = tile_and_concat(
+        input_image, latent, concat_latent=concat_latent)
+
+    enc0 = tfl.conv2d(
+        concat_input_image,
+        conv_size[0], [5, 5],
+        strides=(2, 2),
+        activation=tf.nn.relu,
+        padding="SAME",
+        name="scale1_conv1")
+    enc0 = tfcl.layer_norm(enc0, scope="layer_norm1")
+
+    hidden1, lstm_state[0] = lstm_func(
+        enc0, lstm_state[0], lstm_size[0], name="state1")
+    hidden1 = tile_and_concat(hidden1, latent, concat_latent=concat_latent)
+    hidden1 = tfcl.layer_norm(hidden1, scope="layer_norm2")
+    hidden2, lstm_state[1] = lstm_func(
+        hidden1, lstm_state[1], lstm_size[1], name="state2")
+    hidden2 = tfcl.layer_norm(hidden2, scope="layer_norm3")
+    hidden2 = common_layers.make_even_size(hidden2)
+    enc1 = tfl.conv2d(hidden2, hidden2.get_shape()[3], [3, 3], strides=(2, 2),
+                      padding="SAME", activation=tf.nn.relu, name="conv2")
+    enc1 = tile_and_concat(enc1, latent, concat_latent=concat_latent)
+
+    hidden3, lstm_state[2] = lstm_func(
+        enc1, lstm_state[2], lstm_size[2], name="state3")
+    hidden3 = tile_and_concat(hidden3, latent, concat_latent=concat_latent)
+    hidden3 = tfcl.layer_norm(hidden3, scope="layer_norm4")
+    hidden4, lstm_state[3] = lstm_func(
+        hidden3, lstm_state[3], lstm_size[3], name="state4")
+    hidden4 = tile_and_concat(hidden4, latent, concat_latent=concat_latent)
+    hidden4 = tfcl.layer_norm(hidden4, scope="layer_norm5")
+    hidden4 = common_layers.make_even_size(hidden4)
+    enc2 = tfl.conv2d(hidden4, hidden4.get_shape()[3], [3, 3], strides=(2, 2),
+                      padding="SAME", activation=tf.nn.relu, name="conv3")
+
+    # Pass in reward and action.
+    emb_action = common_video.encode_to_shape(
+        action, enc2.get_shape(), "action_enc")
+    emb_reward = common_video.encode_to_shape(
+        input_reward, enc2.get_shape(), "reward_enc")
+    enc2 = tf.concat(axis=3, values=[enc2, emb_action, emb_reward])
+
+    if latent is not None and not concat_latent:
+      with tf.control_dependencies([latent]):
+        enc2 = tf.concat([enc2, latent], 3)
+
+    enc3 = tfl.conv2d(enc2, hidden4.get_shape()[3], [1, 1], strides=(1, 1),
+                      padding="SAME", activation=tf.nn.relu, name="conv4")
+
+    hidden5, lstm_state[4] = lstm_func(
+        enc3, lstm_state[4], lstm_size[4], name="state5")  # last 8x8
+    hidden5 = tfcl.layer_norm(hidden5, scope="layer_norm6")
+    hidden5 = tile_and_concat(hidden5, latent, concat_latent=concat_latent)
+    return hidden5, (enc0, enc1)
+
+  def construct_latent_tower(self, images):
+    """Builds convolutional latent tower for stochastic model.
+
+    At training time this tower generates a latent distribution (mean and std)
+    conditioned on the entire video. This latent variable will be fed to the
+    main tower as an extra variable to be used for future frames prediction.
+    At inference time, the tower is disabled and only returns latents sampled
+    from N(0,1).
+    If the multi_latent flag is on, a different latent for every timestep would
+    be generated.
+
+    Args:
+      images: tensor of ground truth image sequences
+    Returns:
+      latent_mean: predicted latent mean
+      latent_std: predicted latent standard deviation
+      latent_loss: loss of the latent twoer
+      samples: random samples sampled from standard guassian
+    """
+    conv_size = self.tinyify([32, 64, 64])
+    with tf.variable_scope("latent", reuse=tf.AUTO_REUSE):
+      # this allows more predicted frames at inference time
+      latent_num_frames = self.hparams.latent_num_frames
+      if latent_num_frames == 0:  # use all frames by default.
+        latent_num_frames = (self.hparams.video_num_input_frames +
+                             self.hparams.video_num_target_frames)
+      tf.logging.info("Creating latent tower with %d frames."%latent_num_frames)
+      latent_images = tf.unstack(images[:latent_num_frames], axis=0)
+      images = tf.concat(latent_images, 3)
+
+      x = images
+      x = common_layers.make_even_size(x)
+      x = tfl.conv2d(x, conv_size[0], [3, 3], strides=(2, 2),
+                     padding="SAME", activation=tf.nn.relu, name="latent_conv1")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="latent_bn1")
+      x = common_layers.make_even_size(x)
+      x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
+                     padding="SAME", activation=tf.nn.relu, name="latent_conv2")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="latent_bn2")
+      x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(1, 1),
+                     padding="SAME", activation=tf.nn.relu, name="latent_conv3")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="latent_bn3")
+
+      nc = self.hparams.latent_channels
+      mean = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
+                        padding="SAME", activation=None, name="latent_mean")
+      std = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
+                       padding="SAME", activation=tf.nn.relu, name="latent_std")
+      std += self.hparams.latent_std_min
+
+      # No latent tower at inference time, just standard gaussian.
+      if self.hparams.mode != tf.estimator.ModeKeys.TRAIN:
+        return tf.zeros_like(mean), tf.zeros_like(std)
+
+      return mean, std
+
+  def reward_prediction(
+      self, input_image, input_reward, action, lstm_state, latent):
+    """Builds a reward prediction network."""
+    conv_size = self.tinyify([32, 32, 16, 4])
+    lstm_size = self.tinyify([32, 64, 128, 64, 32])
+
+    with tf.variable_scope("reward_pred", reuse=tf.AUTO_REUSE):
+      hidden5, _ = self.bottom_part_tower(
+          input_image, input_reward, action, latent,
+          lstm_state, lstm_size, conv_size)
+
+      x = hidden5
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="reward_bn0")
+      x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
+                     padding="SAME", activation=tf.nn.relu, name="reward_conv1")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="reward_bn1")
+      x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(2, 2),
+                     padding="SAME", activation=tf.nn.relu, name="reward_conv2")
+      x = tfcl.batch_norm(x, updates_collections=None,
+                          is_training=self.is_training, scope="reward_bn2")
+      x = tfl.conv2d(x, conv_size[3], [3, 3], strides=(2, 2),
+                     padding="SAME", activation=tf.nn.relu, name="reward_conv3")
+
+      pred_reward = common_video.decode_to_shape(
+          x, input_reward.shape, "reward_dec")
+
+      return pred_reward, lstm_state
+
+  def construct_predictive_tower(
+      self, input_image, input_reward, action, lstm_state, latent,
+      concat_latent=False):
+    # Main tower
+    lstm_func = common_video.conv_lstm_2d
+    batch_size = common_layers.shape_list(input_image)[0]
+    # the number of different pixel motion predictions
+    # and the number of masks for each of those predictions
+    num_masks = self.hparams.num_masks
+    upsample_method = self.hparams.upsample_method
+    tile_and_concat = common_video.tile_and_concat
+
+    lstm_size = self.tinyify([32, 32, 64, 64, 128, 64, 32])
+    conv_size = self.tinyify([32])
+
+    img_height, img_width, color_channels = self.hparams.problem.frame_shape
+
+    with tf.variable_scope("main", reuse=tf.AUTO_REUSE):
+      hidden5, skips = self.bottom_part_tower(
+          input_image, input_reward, action, latent,
+          lstm_state, lstm_size, conv_size, concat_latent=concat_latent)
+      enc0, enc1 = skips
+
+      with tf.variable_scope("upsample1", reuse=tf.AUTO_REUSE):
+        enc4 = common_layers.cyclegan_upsample(
+            hidden5, num_outputs=hidden5.shape.as_list()[-1],
+            stride=[2, 2], method=upsample_method)
+
+      enc1_shape = common_layers.shape_list(enc1)
+      enc4 = enc4[:, :enc1_shape[1], :enc1_shape[2], :]  # Cut to shape.
+      enc4 = tile_and_concat(enc4, latent, concat_latent=concat_latent)
+
+      hidden6, lstm_state[5] = lstm_func(
+          enc4, lstm_state[5], lstm_size[5], name="state6",
+          spatial_dims=enc1_shape[1:-1])  # 16x16
+      hidden6 = tile_and_concat(hidden6, latent, concat_latent=concat_latent)
+      hidden6 = tfcl.layer_norm(hidden6, scope="layer_norm7")
+      # Skip connection.
+      hidden6 = tf.concat(axis=3, values=[hidden6, enc1])  # both 16x16
+
+      with tf.variable_scope("upsample2", reuse=tf.AUTO_REUSE):
+        enc5 = common_layers.cyclegan_upsample(
+            hidden6, num_outputs=hidden6.shape.as_list()[-1],
+            stride=[2, 2], method=upsample_method)
+
+      enc0_shape = common_layers.shape_list(enc0)
+      enc5 = enc5[:, :enc0_shape[1], :enc0_shape[2], :]  # Cut to shape.
+      enc5 = tile_and_concat(enc5, latent, concat_latent=concat_latent)
+
+      hidden7, lstm_state[6] = lstm_func(
+          enc5, lstm_state[6], lstm_size[6], name="state7",
+          spatial_dims=enc0_shape[1:-1])  # 32x32
+      hidden7 = tfcl.layer_norm(hidden7, scope="layer_norm8")
+
+      # Skip connection.
+      hidden7 = tf.concat(axis=3, values=[hidden7, enc0])  # both 32x32
+
+      with tf.variable_scope("upsample3", reuse=tf.AUTO_REUSE):
+        enc6 = common_layers.cyclegan_upsample(
+            hidden7, num_outputs=hidden7.shape.as_list()[-1],
+            stride=[2, 2], method=upsample_method)
+      enc6 = tfcl.layer_norm(enc6, scope="layer_norm9")
+      enc6 = tile_and_concat(enc6, latent, concat_latent=concat_latent)
+
+      if self.hparams.model_options == "DNA":
+        # Using largest hidden state for predicting untied conv kernels.
+        enc7 = tfl.conv2d_transpose(
+            enc6,
+            self.hparams.dna_kernel_size**2,
+            [1, 1],
+            strides=(1, 1),
+            padding="SAME",
+            name="convt4",
+            activation=None)
+      else:
+        # Using largest hidden state for predicting a new image layer.
+        enc7 = tfl.conv2d_transpose(
+            enc6,
+            color_channels,
+            [1, 1],
+            strides=(1, 1),
+            padding="SAME",
+            name="convt4",
+            activation=None)
+        # This allows the network to also generate one image from scratch,
+        # which is useful when regions of the image become unoccluded.
+        transformed = [tf.nn.sigmoid(enc7)]
+
+      if self.hparams.model_options == "CDNA":
+        # cdna_input = tf.reshape(hidden5, [int(batch_size), -1])
+        cdna_input = tfcl.flatten(hidden5)
+        transformed += common_video.cdna_transformation(
+            input_image, cdna_input, num_masks, int(color_channels),
+            self.hparams.dna_kernel_size, self.hparams.relu_shift)
+      elif self.hparams.model_options == "DNA":
+        # Only one mask is supported (more should be unnecessary).
+        if num_masks != 1:
+          raise ValueError("Only one mask is supported for DNA model.")
+        transformed = [
+            common_video.dna_transformation(
+                input_image, enc7,
+                self.hparams.dna_kernel_size, self.hparams.relu_shift)]
+
+      masks = tfl.conv2d(
+          enc6, filters=num_masks + 1, kernel_size=[1, 1],
+          strides=(1, 1), name="convt7", padding="SAME")
+      masks = tf.reshape(
+          tf.nn.softmax(tf.reshape(masks, [-1, num_masks + 1])),
+          [batch_size,
+           int(img_height),
+           int(img_width), num_masks + 1])
+      mask_list = tf.split(
+          axis=3, num_or_size_splits=num_masks + 1, value=masks)
+      output = mask_list[0] * input_image
+      for layer, mask in zip(transformed, mask_list[1:]):
+        output += layer * mask
+
+      return output, lstm_state
+
+  def construct_model(self,
+                      images,
+                      actions,
+                      rewards):
+    """Build convolutional lstm video predictor using CDNA, or DNA.
+
+    Args:
+      images: list of tensors of ground truth image sequences
+              there should be a 4D image ?xWxHxC for each timestep
+      actions: list of action tensors
+               each action should be in the shape ?x1xZ
+      rewards: list of reward tensors
+               each reward should be in the shape ?x1xZ
+    Returns:
+      gen_images: predicted future image frames
+      gen_rewards: predicted future rewards
+      latent_mean: mean of approximated posterior
+      latent_std: std of approximated posterior
+
+    Raises:
+      ValueError: if more than 1 mask specified for DNA model.
+    """
+    context_frames = self.hparams.video_num_input_frames
+
+    batch_size = common_layers.shape_list(images)[1]
+    ss_func = self.get_scheduled_sample_func(batch_size)
+
+    def process_single_frame(prev_outputs, inputs):
+      """Process a single frame of the video."""
+      cur_image, cur_reward, action = inputs
+      time_step, prev_image, prev_reward, lstm_states = prev_outputs[:4]
+
+      generated_items = [prev_image, prev_reward]
+      groundtruth_items = [cur_image, cur_reward]
+      done_warm_start = tf.greater(time_step, context_frames - 1)
+      input_image, input_reward = self.get_scheduled_sample_inputs(
+          done_warm_start, groundtruth_items, generated_items, ss_func)
+
+      # Prediction
+      pred_image, lstm_states = self.construct_predictive_tower(
+          input_image, input_reward, action, lstm_states, latent)
+
+      if self.hparams.reward_prediction:
+        reward_lstm_states = prev_outputs[4]
+        pred_reward, reward_lstm_states = self.reward_prediction(
+            input_image, input_reward, action, reward_lstm_states, latent)
+      else:
+        pred_reward = input_reward
+
+      time_step += 1
+      outputs = (time_step, pred_image, pred_reward, lstm_states)
+      if self.hparams.reward_prediction:
+        outputs += (reward_lstm_states,)
+
+      return outputs
+
+    # Latent tower
+    latent = None
+    if self.hparams.stochastic_model:
+      latent_mean, latent_std = self.construct_latent_tower(images)
+      latent = self.get_gaussian_latent(latent_mean, latent_std)
+
+    # HACK: Do first step outside to initialize all the variables
+    lstm_states, reward_lstm_states = [None] * 7, [None] * 5
+    inputs = images[0], rewards[0], actions[0]
+    prev_outputs = (tf.constant(0), images[0], rewards[0], lstm_states)
+    if self.hparams.reward_prediction:
+      prev_outputs += (reward_lstm_states,)
+
+    initializers = process_single_frame(prev_outputs, inputs)
+    first_gen_images = tf.expand_dims(initializers[1], axis=0)
+    first_gen_rewards = tf.expand_dims(initializers[2], axis=0)
+
+    inputs = (images[1:-1], actions[1:-1], rewards[1:-1])
+
+    outputs = tf.scan(process_single_frame, inputs, initializers)
+    gen_images, gen_rewards = outputs[1:3]
+
+    gen_images = tf.concat((first_gen_images, gen_images), axis=0)
+    gen_rewards = tf.concat((first_gen_rewards, gen_rewards), axis=0)
+
+    return gen_images, gen_rewards, [latent_mean], [latent_std]
+
+  def body(self, features):
+    hparams = self.hparams
+    batch_size = common_layers.shape_list(features["inputs"])[0]
+
+    # Swap time and batch axes.
+    input_frames = common_video.swap_time_and_batch_axes(features["inputs"])
+    target_frames = common_video.swap_time_and_batch_axes(features["targets"])
+
+    # Get actions if exist otherwise use zeros
+    input_actions = self.get_input_if_exists(
+        features, "input_action", batch_size, hparams.video_num_input_frames)
+    target_actions = self.get_input_if_exists(
+        features, "target_action", batch_size, hparams.video_num_target_frames)
+
+    # Get rewards if exist otherwise use zeros
+    input_rewards = self.get_input_if_exists(
+        features, "input_reward", batch_size, hparams.video_num_input_frames)
+    target_rewards = self.get_input_if_exists(
+        features, "target_reward", batch_size, hparams.video_num_target_frames)
+
+    all_actions = tf.concat([input_actions, target_actions], axis=0)
+    all_rewards = tf.concat([input_rewards, target_rewards], axis=0)
+    all_frames = tf.concat([input_frames, target_frames], axis=0)
+
+    # Each image is being used twice, in latent tower and main tower.
+    # This is to make sure we are using the *same* image for both, ...
+    # ... given how TF queues work.
+    # NOT sure if this is required at all. Doesn"t hurt though! :)
+    all_frames = tf.identity(all_frames)
+
+    gen_images, gen_rewards, latent_means, latent_stds = self.construct_model(
+        images=all_frames,
+        actions=all_actions,
+        rewards=all_rewards,
+    )
+
+    step_num = tf.train.get_global_step()
+    # TODO(mbz): what should it be if it"s undefined?
+    if step_num is None:
+      step_num = _LARGE_STEP_NUMBER
+
+    schedule = self.hparams.latent_loss_multiplier_schedule
+    second_stage = self.hparams.num_iterations_2nd_stage
+    # TODO(mechcoder): Add log_annealing schedule.
+    if schedule == "constant":
+      beta = tf.cond(tf.greater(step_num, second_stage),
+                     lambda: self.hparams.latent_loss_multiplier,
+                     lambda: 0.0)
+    elif schedule == "linear_anneal":
+      # Linearly anneal beta from 0.0 to self.hparams.latent_loss_multiplier.
+      # between self.hparams.num_iterations_2nd_stage to anneal_end.
+      # beta = latent_loss * (1 - (global_step - 2nd_stage) / (anneal_end - 2nd_stage))  # pylint:disable=line-too-long
+      anneal_end = self.hparams.anneal_end
+      latent_multiplier = self.hparams.latent_loss_multiplier
+      if anneal_end < second_stage:
+        raise ValueError("Expected hparams.num_iterations_2nd_stage < "
+                         "hparams.anneal_end %d, got %d." %
+                         (second_stage, anneal_end))
+
+      def anneal_loss(step_num):
+        step_num = tf.cast(step_num, dtype=tf.float32)
+        fraction = (float(anneal_end) - step_num) / (anneal_end - second_stage)
+        return self.hparams.latent_loss_multiplier * (1 - fraction)
+
+      beta = tf.case(
+          pred_fn_pairs={
+              tf.less(step_num, second_stage): lambda: 0.0,
+              tf.greater(step_num, anneal_end): lambda: latent_multiplier},
+          default=lambda: anneal_loss(step_num))
+
+    kl_loss = 0.0
+    if self.is_training:
+      for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
+        kl_loss += common_layers.kl_divergence(mean, std)
+        tf.summary.histogram("posterior_mean_%d" % i, mean)
+        tf.summary.histogram("posterior_std_%d" % i, std)
+
+      tf.summary.scalar("beta", beta)
+      tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
+
+    extra_loss = beta * kl_loss
+
+    # Ignore the predictions from the input frames.
+    # This is NOT the same as original paper/implementation.
+    predictions = gen_images[hparams.video_num_input_frames-1:]
+    reward_pred = gen_rewards[hparams.video_num_input_frames-1:]
+    reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove undeeded dimension.
+
+    # TODO(mbz): clean this up!
+    def fix_video_dims_and_concat_on_x_axis(x):
+      x = tf.transpose(x, [1, 3, 4, 0, 2])
+      x = tf.reshape(x, [batch_size, 64, 3, -1])
+      x = tf.transpose(x, [0, 3, 1, 2])
+      return x
+
+    frames_gd = fix_video_dims_and_concat_on_x_axis(target_frames)
+    frames_pd = fix_video_dims_and_concat_on_x_axis(predictions)
+    side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
+    tf.summary.image("full_video", side_by_side_video)
+
+    # Swap back time and batch axes.
+    predictions = common_video.swap_time_and_batch_axes(predictions)
+    reward_pred = common_video.swap_time_and_batch_axes(reward_pred)
+
+    return_targets = predictions
+    if "target_reward" in features:
+      return_targets = {"targets": predictions, "target_reward": reward_pred}
+
+    return return_targets, extra_loss
+
+
+@registry.register_model
+class NextFrameStochasticTwoFrames(NextFrameStochastic):
+  """Stochastic next-frame model with 2 frames posterior."""
+
+  def construct_model(self, images, actions, rewards):
+    images = tf.unstack(images, axis=0)
+    actions = tf.unstack(actions, axis=0)
+    rewards = tf.unstack(rewards, axis=0)
+
+    batch_size = common_layers.shape_list(images[0])[0]
+    context_frames = self.hparams.video_num_input_frames
+
+    # Predicted images and rewards.
+    gen_rewards, gen_images, latent_means, latent_stds = [], [], [], []
+
+    # LSTM states.
+    lstm_state = [None] * 7
+    reward_lstm_state = [None] * 5
+
+    # Create scheduled sampling function
+    ss_func = self.get_scheduled_sample_func(batch_size)
+
+    pred_image, pred_reward, latent = images[0], rewards[0], None
+    for timestep, image, action, reward in zip(
+        range(len(images)-1), images[:-1], actions[:-1], rewards[:-1]):
+      # Scheduled Sampling
+      done_warm_start = timestep > context_frames - 1
+      groundtruth_items = [image, reward]
+      generated_items = [pred_image, pred_reward]
+      input_image, input_reward = self.get_scheduled_sample_inputs(
+          done_warm_start, groundtruth_items, generated_items, ss_func)
+
+      # Latent
+      # TODO(mbz): should we use input_image iunstead of image?
+      latent_images = [image, images[timestep+1]]
+      latent_mean, latent_std = self.construct_latent_tower(latent_images)
+      latent = self.get_gaussian_latent(latent_mean, latent_std)
+      latent_means.append(latent_mean)
+      latent_stds.append(latent_std)
+
+      # Prediction
+      pred_image, lstm_state = self.construct_predictive_tower(
+          input_image, input_reward, action, lstm_state, latent)
+
+      if self.hparams.reward_prediction:
+        pred_reward, reward_lstm_state = self.reward_prediction(
+            input_image, input_reward, action, reward_lstm_state, latent)
+      else:
+        pred_reward = input_reward
+
+      gen_images.append(pred_image)
+      gen_rewards.append(pred_reward)
+
+    gen_images = tf.stack(gen_images, axis=0)
+    gen_rewards = tf.stack(gen_rewards, axis=0)
+
+    return gen_images, gen_rewards, latent_means, latent_stds
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index d31559c8e..72e518c02 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -21,7 +21,10 @@
 
 from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
 from tensor2tensor.models.research import next_frame
+from tensor2tensor.models.research import next_frame_emily
 from tensor2tensor.models.research import next_frame_params
+from tensor2tensor.models.research import next_frame_savp
+from tensor2tensor.models.research import next_frame_sv2p
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -67,6 +70,12 @@ def TestOnVariousInputOutputSizes(self, hparams, model, expected_last_dim):
     self.TestVideoModel(4, 1, hparams, model, expected_last_dim)
     self.TestVideoModel(7, 5, hparams, model, expected_last_dim)
 
+  def TestOnVariousUpSampleLayers(self, hparams, model, expected_last_dim):
+    self.TestVideoModel(4, 1, hparams, model, expected_last_dim,
+                        upsample_method="bilinear_upsample_conv")
+    self.TestVideoModel(4, 1, hparams, model, expected_last_dim,
+                        upsample_method="nn_upsample_conv")
+
   def testBasic(self):
     self.TestOnVariousInputOutputSizes(
         next_frame_params.next_frame(),
@@ -76,19 +85,29 @@ def testBasic(self):
   def testStochastic(self):
     self.TestOnVariousInputOutputSizes(
         next_frame_params.next_frame_stochastic(),
-        next_frame.NextFrameStochastic,
+        next_frame_sv2p.NextFrameStochastic,
         1)
 
   def testStochasticTwoFrames(self):
     self.TestOnVariousInputOutputSizes(
         next_frame_params.next_frame_stochastic(),
-        next_frame.NextFrameStochasticTwoFrames,
+        next_frame_sv2p.NextFrameStochasticTwoFrames,
         1)
 
   def testStochasticEmily(self):
     self.TestOnVariousInputOutputSizes(
         next_frame_params.next_frame_stochastic_emily(),
-        next_frame.NextFrameStochasticEmily,
+        next_frame_emily.NextFrameStochasticEmily,
+        1)
+
+  def testStochasticSavp(self):
+    self.TestOnVariousInputOutputSizes(
+        next_frame_params.next_frame_savp(),
+        next_frame_savp.NextFrameSAVP,
+        1)
+    self.TestOnVariousUpSampleLayers(
+        next_frame_params.next_frame_savp(),
+        next_frame_savp.NextFrameSAVP,
         1)
 
 

From e99ae9135e214875f83df70f1c366c8a86736bb9 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 1 Aug 2018 19:00:27 -0700
Subject: [PATCH 0446/2720] Internal change

PiperOrigin-RevId: 207033295
---
 tensor2tensor/bin/t2t_trainer.py      |  5 ----
 tensor2tensor/layers/common_layers.py |  3 --
 tensor2tensor/utils/optimize.py       | 13 ++++----
 tensor2tensor/utils/t2t_model.py      | 43 +++++++++++----------------
 tensor2tensor/utils/trainer_lib.py    | 10 ++-----
 5 files changed, 25 insertions(+), 49 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 91ebfee67..9dc07dced 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -48,10 +48,6 @@
 flags.DEFINE_integer("iterations_per_loop", 100,
                      "Number of iterations in a TPU training loop.")
 flags.DEFINE_bool("use_tpu", False, "Whether to use TPU.")
-flags.DEFINE_bool("xla_compile", False, "Whether to use XLA to compile. When "
-                  "this is set to True, computation will be constructed to "
-                  "optimize for XLA as if use_tpu=True but run on CPU/GPU "
-                  "instead of TPU.")
 flags.DEFINE_integer("tpu_infeed_sleep_secs", None,
                      "How long to sleep the infeed thread.")
 flags.DEFINE_bool("generate_data", False, "Generate data before training?")
@@ -178,7 +174,6 @@ def create_experiment_fn(**kwargs):
       eval_early_stopping_metric_minimize=FLAGS.
       eval_early_stopping_metric_minimize,
       use_tpu=FLAGS.use_tpu,
-      xla_compile=FLAGS.xla_compile,
       **kwargs)
 
 
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 1ee456f96..c601f57d2 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3406,9 +3406,6 @@ def should_generate_summaries():
   Returns:
     a boolean
   """
-  if is_on_tpu():
-    # Summaries don't work well with TPU and XLA.
-    return False
   if "while/" in tf.contrib.framework.get_name_scope():
     # Summaries don't work well within tf.while_loop()
     return False
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index bdbedbd17..7249a76d8 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -44,11 +44,9 @@ def optimize(loss, learning_rate, hparams, use_tpu=False):
   if use_tpu:
     opt = tf.contrib.tpu.CrossShardOptimizer(opt)
 
-  opt_summaries = []
-  if common_layers.should_generate_summaries():
-    tf.summary.scalar("learning_rate", learning_rate)
-    opt_summaries = ["loss"]
-  if hparams.summarize_grads and common_layers.should_generate_summaries():
+  tf.summary.scalar("learning_rate", learning_rate)
+  opt_summaries = ["loss"]
+  if hparams.summarize_grads:
     tf.logging.info("Summarizing gradients")
     opt_summaries.extend(["gradients", "gradient_norm", "global_gradient_norm"])
 
@@ -138,7 +136,7 @@ def weight_decay_and_noise(loss, hparams, learning_rate, var_list=None):
   noise_vars = [v for v in var_list if "/body/" in v.name]
 
   weight_decay_loss = weight_decay(hparams.weight_decay, decay_vars)
-  if hparams.weight_decay and common_layers.should_generate_summaries():
+  if hparams.weight_decay:
     tf.summary.scalar("losses/weight_decay", weight_decay_loss)
   weight_noise_ops = weight_noise(hparams.weight_noise, learning_rate,
                                   noise_vars)
@@ -163,8 +161,7 @@ def weight_noise(noise_rate, learning_rate, var_list):
   for v in var_list:
     with tf.device(v._ref().device):  # pylint: disable=protected-access
       scale = noise_rate * learning_rate * 0.001
-      if common_layers.should_generate_summaries():
-        tf.summary.scalar("weight_noise_scale", scale)
+      tf.summary.scalar("weight_noise_scale", scale)
       noise = tf.truncated_normal(v.shape) * scale
       noise_op = v.assign_add(noise)
       noise_ops.append(noise_op)
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index e6a7aa72c..640119301 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -194,8 +194,7 @@ def body_sharded(self, sharded_features):
 
   def model_fn_sharded(self, sharded_features):
     dp = self._data_parallelism
-    if common_layers.should_generate_summaries():
-      summarize_features(sharded_features, num_shards=dp.n)
+    summarize_features(sharded_features, num_shards=dp.n)
     datashard_to_features = self._to_features_per_datashard(sharded_features)
     if self.use_body_sharded:
       # MoE models override body_sharded
@@ -437,10 +436,9 @@ def loss(self, logits, features):
         losses[k] = self._loss_single(v, target_modality[k], features[k])
 
         n, d = losses[k]
-        if common_layers.should_generate_summaries():
-          tf.summary.scalar(k + "_loss", n / d)
-          tf.summary.scalar(k + "_loss_num", n)
-          tf.summary.scalar(k + "_loss_den", d)
+        tf.summary.scalar(k + "_loss", n / d)
+        tf.summary.scalar(k + "_loss_num", n)
+        tf.summary.scalar(k + "_loss_den", d)
 
       return tf.add_n([n / d for n, d in losses.values()])
     else:
@@ -455,14 +453,15 @@ def loss(self, logits, features):
         target_modality = target_modality["targets"]
       return self._loss_single(logits, target_modality, features["targets"])
 
-  def optimize(self, loss, num_async_replicas=1, use_tpu=False):
+  def optimize(self, loss, num_async_replicas=1):
     """Return a training op minimizing loss."""
     lr = learning_rate.learning_rate_schedule(self.hparams)
     if num_async_replicas > 1:
       log_info("Dividing learning rate by num_async_replicas: %d",
                num_async_replicas)
     lr /= math.sqrt(float(num_async_replicas))
-    train_op = optimize.optimize(loss, lr, self.hparams, use_tpu=use_tpu)
+    train_op = optimize.optimize(
+        loss, lr, self.hparams, use_tpu=common_layers.is_on_tpu())
     return train_op
 
   def set_mode(self, mode):
@@ -1149,8 +1148,7 @@ def get_eval_hooks(model_name):
   def make_estimator_model_fn(model_name,
                               hparams,
                               decode_hparams=None,
-                              use_tpu=False,
-                              xla_compile=False):
+                              use_tpu=False):
     model_cls = registry.model(model_name)
 
     def wrapping_model_fn(features, labels, mode, params=None, config=None):
@@ -1162,8 +1160,7 @@ def wrapping_model_fn(features, labels, mode, params=None, config=None):
           config=config,
           params=params,
           decode_hparams=decode_hparams,
-          use_tpu=use_tpu,
-          xla_compile=xla_compile)
+          use_tpu=use_tpu)
 
     return wrapping_model_fn
 
@@ -1176,8 +1173,7 @@ def estimator_model_fn(cls,
                          config=None,
                          params=None,
                          decode_hparams=None,
-                         use_tpu=False,
-                         xla_compile=False):
+                         use_tpu=False):
     """Model fn for Estimator.
 
     Args:
@@ -1189,7 +1185,6 @@ def estimator_model_fn(cls,
       params: dict, may include batch_size
       decode_hparams: HParams, used when mode == PREDICT.
       use_tpu: bool, whether using TPU
-      xla_compile: bool, whether to use XLA to compile graph, unimplemented.
 
     Returns:
       TPUEstimatorSpec if use tpu else EstimatorSpec
@@ -1221,7 +1216,7 @@ def estimator_model_fn(cls,
       logits, losses_dict = model(features)  # pylint: disable=not-callable
 
     # Set known shapes
-    if use_tpu or xla_compile:
+    if use_tpu:
       if isinstance(logits, dict):
         for k, v in sorted(six.iteritems(logits)):
           if "scalar/" in k:
@@ -1248,10 +1243,9 @@ def estimator_model_fn(cls,
       return logits
 
     # Summarize losses
-    if common_layers.should_generate_summaries():
-      with tf.name_scope("losses"):
-        for loss_name, loss_val in sorted(losses_dict.items()):
-          tf.summary.scalar(loss_name, loss_val)
+    with tf.name_scope("losses"):
+      for loss_name, loss_val in sorted(losses_dict.items()):
+        tf.summary.scalar(loss_name, loss_val)
 
     # Accumulate losses
     loss = sum(losses_dict[key] for key in sorted(losses_dict.keys()))
@@ -1266,14 +1260,11 @@ def estimator_model_fn(cls,
     num_async_replicas = (1 if (use_tpu or not config) else
                           config.t2t_device_info["num_async_replicas"])
     return model.estimator_spec_train(
-        loss,
-        num_async_replicas=num_async_replicas,
-        use_tpu=use_tpu and not xla_compile)
+        loss, num_async_replicas=num_async_replicas)
 
-  def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
+  def estimator_spec_train(self, loss, num_async_replicas=1):
     """Construct EstimatorSpec for TRAIN mode."""
-    train_op = self.optimize(
-        loss, num_async_replicas=num_async_replicas, use_tpu=use_tpu)
+    train_op = self.optimize(loss, num_async_replicas=num_async_replicas)
 
     # TODO(mitchellstern): Add support for partitioned variables?
     if (tf.train.latest_checkpoint(self._hparams.model_dir) is None and
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 05a31ac89..c5ea9df24 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -207,12 +207,10 @@ def create_estimator(model_name,
                      run_config,
                      schedule="train_and_evaluate",
                      decode_hparams=None,
-                     use_tpu=False,
-                     xla_compile=False):
+                     use_tpu=False):
   """Create a T2T Estimator."""
   model_fn = t2t_model.T2TModel.make_estimator_model_fn(
-      model_name, hparams, decode_hparams=decode_hparams, use_tpu=use_tpu,
-      xla_compile=xla_compile)
+      model_name, hparams, decode_hparams=decode_hparams, use_tpu=use_tpu)
 
   if use_tpu:
     problem = hparams.problem
@@ -414,7 +412,6 @@ def create_experiment(
     eval_early_stopping_metric_minimize=True,
     autotune=False,
     use_tpu=False,
-    xla_compile=False,
     additional_train_hooks=None,
     additional_eval_hooks=None):
   """Create Experiment."""
@@ -433,8 +430,7 @@ def create_experiment(
       run_config,
       schedule=schedule,
       decode_hparams=decode_hparams,
-      use_tpu=use_tpu,
-      xla_compile=xla_compile)
+      use_tpu=use_tpu)
 
   # Input fns from Problem
   problem = hparams.problem

From 1a754af96bfdfd6015c7763374c0d5910a93dbf2 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 1 Aug 2018 20:36:57 -0700
Subject: [PATCH 0447/2720] fix the scheduled sampling order bug.

PiperOrigin-RevId: 207040593
---
 tensor2tensor/models/research/next_frame_sv2p.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 57b5e8246..b92937cda 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -112,10 +112,10 @@ def sample():
           output_items.append(scheduled_sampling_func(item_gt, item_gen))
         return output_items
 
-    cases = {
-        tf.logical_not(done_warm_start): lambda: groundtruth_items,
-        tf.logical_not(self.is_training): lambda: generated_items,
-    }
+    cases = [
+        (tf.logical_not(done_warm_start), lambda: groundtruth_items),
+        (tf.logical_not(self.is_training), lambda: generated_items),
+    ]
     output_items = tf.case(cases, default=sample, strict=True)
 
     return output_items

From 9786320b8a5d5122c43e191553f70b3b45d73efc Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 1 Aug 2018 20:46:23 -0700
Subject: [PATCH 0448/2720] Fix SAVP test.

PiperOrigin-RevId: 207041279
---
 tensor2tensor/models/research/next_frame_savp.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/next_frame_savp.py b/tensor2tensor/models/research/next_frame_savp.py
index ad3044d91..1e4883291 100644
--- a/tensor2tensor/models/research/next_frame_savp.py
+++ b/tensor2tensor/models/research/next_frame_savp.py
@@ -175,7 +175,7 @@ def construct_model(self, images, actions, rewards):
         done_warm_start = step > context_frames - 1
         groundtruth_items = [image]
         generated_items = [pred_image]
-        input_image = self.get_scheduled_sample_inputs(
+        input_image, = self.get_scheduled_sample_inputs(
             done_warm_start, groundtruth_items, generated_items, ss_func)
 
         all_latents = tf.concat([enc_cond_latent, enc_prior_latent], axis=0)

From 993d6dffa6bbb05cc56ca3e5390ea19f4e35013e Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 1 Aug 2018 20:52:07 -0700
Subject: [PATCH 0449/2720] Safer inputs for initial prediction in SV2P models,
 to highlight scheduled sampling bugs.

PiperOrigin-RevId: 207041676
---
 tensor2tensor/models/research/next_frame_sv2p.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index b92937cda..83c6aeb96 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -483,7 +483,10 @@ def process_single_frame(prev_outputs, inputs):
     # HACK: Do first step outside to initialize all the variables
     lstm_states, reward_lstm_states = [None] * 7, [None] * 5
     inputs = images[0], rewards[0], actions[0]
-    prev_outputs = (tf.constant(0), images[0], rewards[0], lstm_states)
+    prev_outputs = (tf.constant(0),
+                    tf.zeros_like(images[0]),
+                    tf.zeros_like(rewards[0]),
+                    lstm_states)
     if self.hparams.reward_prediction:
       prev_outputs += (reward_lstm_states,)
 
@@ -634,7 +637,9 @@ def construct_model(self, images, actions, rewards):
     # Create scheduled sampling function
     ss_func = self.get_scheduled_sample_func(batch_size)
 
-    pred_image, pred_reward, latent = images[0], rewards[0], None
+    pred_image = tf.zeros_like(images[0])
+    pred_reward = tf.zeros_like(rewards[0])
+    latent = None
     for timestep, image, action, reward in zip(
         range(len(images)-1), images[:-1], actions[:-1], rewards[:-1]):
       # Scheduled Sampling

From affd67e12457394fd142520b4900b430a7f40bfe Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 1 Aug 2018 21:09:44 -0700
Subject: [PATCH 0450/2720] Fix model registration.

PiperOrigin-RevId: 207043189
---
 tensor2tensor/models/__init__.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index e5f49e04a..311956c50 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -44,7 +44,9 @@
 from tensor2tensor.models.research import lm_experiments
 from tensor2tensor.models.research import multimodel
 from tensor2tensor.models.research import next_frame
+from tensor2tensor.models.research import next_frame_emily
 from tensor2tensor.models.research import next_frame_savp
+from tensor2tensor.models.research import next_frame_sv2p
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research import similarity_transformer
 from tensor2tensor.models.research import super_lm

From 5c7fe30550d5dc576b72373c94aeba6d8f9058fa Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 1 Aug 2018 21:20:32 -0700
Subject: [PATCH 0451/2720] Update method name to conform to linter

PiperOrigin-RevId: 207043896
---
 tensor2tensor/layers/common_video_test.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index 3eb32ea5e..efa031957 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -25,7 +25,7 @@
 class CommonVideoTest(tf.test.TestCase):
 
   @staticmethod
-  def runScheduledSampleFunc(func, var, batch_size):
+  def RunScheduledSampleFunc(func, var, batch_size):
     ground_truth_x = list(range(1, batch_size+1))
     generated_x = [-x for x in ground_truth_x]
     ground_truth_x = tf.convert_to_tensor(ground_truth_x)
@@ -36,34 +36,34 @@ def runScheduledSampleFunc(func, var, batch_size):
     return output
 
   def testScheduledSampleProbStart(self):
-    ground_truth_x, _, ss_out = CommonVideoTest.runScheduledSampleFunc(
+    ground_truth_x, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
         common_video.scheduled_sample_prob, 1.0, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
   def testScheduledSampleProbMid(self):
-    _, _, ss_out = CommonVideoTest.runScheduledSampleFunc(
+    _, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
         common_video.scheduled_sample_prob, 0.5, 1000)
     positive_count = np.sum(ss_out > 0)
     self.assertAlmostEqual(positive_count / 1000.0, 0.5, places=2)
 
   def testScheduledSampleProbEnd(self):
-    _, generated_x, ss_out = CommonVideoTest.runScheduledSampleFunc(
+    _, generated_x, ss_out = CommonVideoTest.RunScheduledSampleFunc(
         common_video.scheduled_sample_prob, 0.0, 10)
     self.assertAllEqual(generated_x, ss_out)
 
   def testScheduledSampleCountStart(self):
-    ground_truth_x, _, ss_out = CommonVideoTest.runScheduledSampleFunc(
+    ground_truth_x, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
         common_video.scheduled_sample_count, 10, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
   def testScheduledSampleCountMid(self):
-    _, _, ss_out = CommonVideoTest.runScheduledSampleFunc(
+    _, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
         common_video.scheduled_sample_count, 5, 10)
     positive_count = np.sum(ss_out > 0)
     self.assertEqual(positive_count, 5)
 
   def testScheduledSampleCountEnd(self):
-    _, generated_x, ss_out = CommonVideoTest.runScheduledSampleFunc(
+    _, generated_x, ss_out = CommonVideoTest.RunScheduledSampleFunc(
         common_video.scheduled_sample_count, 0, 10)
     self.assertAllEqual(generated_x, ss_out)
 

From 7e2e9f1851bfb66fbb791a1c1d546256c967e3a8 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 2 Aug 2018 09:53:59 -0700
Subject: [PATCH 0452/2720] fix action/reward bug.

PiperOrigin-RevId: 207116631
---
 .../models/research/next_frame_sv2p.py        |   2 +-
 .../models/research/next_frame_test.py        | 121 +++++++++++++++---
 2 files changed, 101 insertions(+), 22 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 83c6aeb96..53507ec28 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -494,7 +494,7 @@ def process_single_frame(prev_outputs, inputs):
     first_gen_images = tf.expand_dims(initializers[1], axis=0)
     first_gen_rewards = tf.expand_dims(initializers[2], axis=0)
 
-    inputs = (images[1:-1], actions[1:-1], rewards[1:-1])
+    inputs = (images[1:-1], rewards[1:-1], actions[1:-1])
 
     outputs = tf.scan(process_single_frame, inputs, initializers)
     gen_images, gen_rewards = outputs[1:3]
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index 72e518c02..adf02bc14 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -30,8 +30,68 @@
 import tensorflow as tf
 
 
+def fill_hparams(hparams, in_frames, out_frames):
+  hparams.video_num_input_frames = in_frames
+  hparams.video_num_target_frames = out_frames
+  problem = registry.problem("video_stochastic_shapes10k")
+  p_hparams = problem.get_hparams(hparams)
+  hparams.problem = problem
+  hparams.problem_hparams = p_hparams
+  return hparams
+
+
+def full_modalities(hparams):
+  hparams.problem_hparams.input_modality = {
+      "inputs": ("video:l2raw", 256),
+      "input_reward": ("symbol:one_hot", 3),
+      "input_action": ("symbol:one_hot", 5)
+  }
+  hparams.problem_hparams.target_modality = {
+      "targets": ("video:l2raw", 256),
+      "target_reward": ("symbol:one_hot", 3),
+      "target_action": ("symbol:one_hot", 5),
+  }
+  return hparams
+
+
+def create_basic_features(in_frames, out_frames):
+  x = np.random.randint(0, 256, size=(8, in_frames, 64, 64, 3))
+  y = np.random.randint(0, 256, size=(8, out_frames, 64, 64, 3))
+  features = {
+      "inputs": tf.constant(x, dtype=tf.int32),
+      "targets": tf.constant(y, dtype=tf.int32),
+  }
+  return features
+
+
+def create_full_features(in_frames, out_frames):
+  features = create_basic_features(in_frames, out_frames)
+  x = np.random.randint(0, 5, size=(8, in_frames, 1))
+  y = np.random.randint(0, 5, size=(8, out_frames, 1))
+  features["input_action"] = tf.constant(x, dtype=tf.int32)
+  features["target_action"] = tf.constant(y, dtype=tf.int32)
+  x = np.random.randint(0, 5, size=(8, in_frames, 1))
+  y = np.random.randint(0, 5, size=(8, out_frames, 1))
+  features["input_reward"] = tf.constant(x, dtype=tf.int32)
+  features["target_reward"] = tf.constant(y, dtype=tf.int32)
+  return features
+
+
+def get_tensor_shape(tensor):
+  return tuple([d.value for d in tensor.shape])
+
+
 class NextFrameTest(tf.test.TestCase):
 
+  def RunModel(self, model, hparams, features):
+    with self.test_session() as session:
+      model = model(
+          hparams, tf.estimator.ModeKeys.TRAIN)
+      logits, _ = model(features)
+      session.run(tf.global_variables_initializer())
+      res = session.run(logits)
+    return res
+
   def TestVideoModel(self,
                      in_frames,
                      out_frames,
@@ -39,30 +99,36 @@ def TestVideoModel(self,
                      model,
                      expected_last_dim,
                      upsample_method="conv2d_transpose"):
+    hparams = fill_hparams(hparams, in_frames, out_frames)
+    hparams.upsample_method = upsample_method
 
-    x = np.random.random_integers(0, high=255, size=(8, in_frames, 64, 64, 3))
-    y = np.random.random_integers(0, high=255, size=(8, out_frames, 64, 64, 3))
+    features = create_basic_features(in_frames, out_frames)
+    output = self.RunModel(model, hparams, features)
 
-    hparams.video_num_input_frames = in_frames
-    hparams.video_num_target_frames = out_frames
-    hparams.upsample_method = upsample_method
-    problem = registry.problem("video_stochastic_shapes10k")
-    p_hparams = problem.get_hparams(hparams)
-    hparams.problem = problem
-    hparams.problem_hparams = p_hparams
+    targets = features["targets"]
+    expected_shape = get_tensor_shape(targets) + (expected_last_dim,)
+    self.assertEqual(output.shape, expected_shape)
 
-    with self.test_session() as session:
-      features = {
-          "inputs": tf.constant(x, dtype=tf.int32),
-          "targets": tf.constant(y, dtype=tf.int32),
-      }
-      model = model(
-          hparams, tf.estimator.ModeKeys.TRAIN)
-      logits, _ = model(features)
-      session.run(tf.global_variables_initializer())
-      res = session.run(logits)
-    expected_shape = y.shape + (expected_last_dim,)
-    self.assertEqual(res.shape, expected_shape)
+  def TestVideoModelWithActionAndRewards(self,
+                                         in_frames,
+                                         out_frames,
+                                         hparams,
+                                         model,
+                                         expected_last_dim):
+    hparams = fill_hparams(hparams, in_frames, out_frames)
+    hparams = full_modalities(hparams)
+
+    features = create_full_features(in_frames, out_frames)
+
+    res = self.RunModel(model, hparams, features)
+
+    output, targets = res["targets"], features["targets"]
+    expected_shape = get_tensor_shape(targets) + (expected_last_dim,)
+    self.assertEqual(output.shape, expected_shape)
+
+    output, targets = res["target_reward"], features["target_reward"]
+    expected_shape = get_tensor_shape(targets)[:2] + (3,)
+    self.assertEqual(output.shape, expected_shape)
 
   def TestOnVariousInputOutputSizes(self, hparams, model, expected_last_dim):
     self.TestVideoModel(1, 1, hparams, model, expected_last_dim)
@@ -70,6 +136,13 @@ def TestOnVariousInputOutputSizes(self, hparams, model, expected_last_dim):
     self.TestVideoModel(4, 1, hparams, model, expected_last_dim)
     self.TestVideoModel(7, 5, hparams, model, expected_last_dim)
 
+  def TestWithActionAndRewards(self, hparams, model, expected_last_dim):
+    test_func = self.TestVideoModelWithActionAndRewards
+    test_func(1, 1, hparams, model, expected_last_dim)
+    test_func(1, 6, hparams, model, expected_last_dim)
+    test_func(4, 1, hparams, model, expected_last_dim)
+    test_func(7, 5, hparams, model, expected_last_dim)
+
   def TestOnVariousUpSampleLayers(self, hparams, model, expected_last_dim):
     self.TestVideoModel(4, 1, hparams, model, expected_last_dim,
                         upsample_method="bilinear_upsample_conv")
@@ -88,6 +161,12 @@ def testStochastic(self):
         next_frame_sv2p.NextFrameStochastic,
         1)
 
+  def testStochasticWithActionsAndRewards(self):
+    self.TestWithActionAndRewards(
+        next_frame_params.next_frame_stochastic(),
+        next_frame_sv2p.NextFrameStochastic,
+        1)
+
   def testStochasticTwoFrames(self):
     self.TestOnVariousInputOutputSizes(
         next_frame_params.next_frame_stochastic(),

From abab5b7b8ce93d3d6e55648a3d4b7641460ebe14 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 2 Aug 2018 11:50:57 -0700
Subject: [PATCH 0453/2720] adding more tests for action only (no rewards)
 scenarios.

PiperOrigin-RevId: 207139070
---
 .../models/research/next_frame_test.py        | 43 ++++++++++++++++++-
 1 file changed, 42 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index adf02bc14..cc57e7668 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -40,6 +40,18 @@ def fill_hparams(hparams, in_frames, out_frames):
   return hparams
 
 
+def action_modalities(hparams):
+  hparams.problem_hparams.input_modality = {
+      "inputs": ("video:l2raw", 256),
+      "input_action": ("symbol:one_hot", 5)
+  }
+  hparams.problem_hparams.target_modality = {
+      "targets": ("video:l2raw", 256),
+      "target_action": ("symbol:one_hot", 5),
+  }
+  return hparams
+
+
 def full_modalities(hparams):
   hparams.problem_hparams.input_modality = {
       "inputs": ("video:l2raw", 256),
@@ -64,12 +76,17 @@ def create_basic_features(in_frames, out_frames):
   return features
 
 
-def create_full_features(in_frames, out_frames):
+def create_action_features(in_frames, out_frames):
   features = create_basic_features(in_frames, out_frames)
   x = np.random.randint(0, 5, size=(8, in_frames, 1))
   y = np.random.randint(0, 5, size=(8, out_frames, 1))
   features["input_action"] = tf.constant(x, dtype=tf.int32)
   features["target_action"] = tf.constant(y, dtype=tf.int32)
+  return features
+
+
+def create_full_features(in_frames, out_frames):
+  features = create_basic_features(in_frames, out_frames)
   x = np.random.randint(0, 5, size=(8, in_frames, 1))
   y = np.random.randint(0, 5, size=(8, out_frames, 1))
   features["input_reward"] = tf.constant(x, dtype=tf.int32)
@@ -109,6 +126,23 @@ def TestVideoModel(self,
     expected_shape = get_tensor_shape(targets) + (expected_last_dim,)
     self.assertEqual(output.shape, expected_shape)
 
+  def TestVideoModelWithActions(self,
+                                in_frames,
+                                out_frames,
+                                hparams,
+                                model,
+                                expected_last_dim):
+    hparams = fill_hparams(hparams, in_frames, out_frames)
+    hparams = action_modalities(hparams)
+    hparams.reward_prediction = False
+
+    features = create_action_features(in_frames, out_frames)
+    output = self.RunModel(model, hparams, features)
+
+    targets = features["targets"]
+    expected_shape = get_tensor_shape(targets) + (expected_last_dim,)
+    self.assertEqual(output.shape, expected_shape)
+
   def TestVideoModelWithActionAndRewards(self,
                                          in_frames,
                                          out_frames,
@@ -136,6 +170,13 @@ def TestOnVariousInputOutputSizes(self, hparams, model, expected_last_dim):
     self.TestVideoModel(4, 1, hparams, model, expected_last_dim)
     self.TestVideoModel(7, 5, hparams, model, expected_last_dim)
 
+  def TestWithActions(self, hparams, model, expected_last_dim):
+    test_func = self.TestVideoModelWithActionAndRewards
+    test_func(1, 1, hparams, model, expected_last_dim)
+    test_func(1, 6, hparams, model, expected_last_dim)
+    test_func(4, 1, hparams, model, expected_last_dim)
+    test_func(7, 5, hparams, model, expected_last_dim)
+
   def TestWithActionAndRewards(self, hparams, model, expected_last_dim):
     test_func = self.TestVideoModelWithActionAndRewards
     test_func(1, 1, hparams, model, expected_last_dim)

From dbed734d14065a300df13ba0f4aed9e3d7291330 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 2 Aug 2018 13:09:13 -0700
Subject: [PATCH 0454/2720] minor refactorings.

PiperOrigin-RevId: 207151820
---
 .../models/research/next_frame_sv2p.py        | 83 ++++++++++---------
 1 file changed, 44 insertions(+), 39 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 53507ec28..dc16bd1b8 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -57,14 +57,52 @@ def get_gaussian_latent(self, latent_mean, latent_std):
     latent = latent_mean + tf.exp(latent_std / 2.0) * latent
     return latent
 
+  def get_iteration_num(self):
+    step_num = tf.train.get_global_step()
+    # TODO(lukaszkaiser): what should it be if it"s undefined?
+    if step_num is None:
+      step_num = 1000000
+    return step_num
+
+  def get_beta(self):
+    """Get KL multiplier (beta) based on the schedule."""
+    step_num = self.get_iteration_num()
+    schedule = self.hparams.latent_loss_multiplier_schedule
+    second_stage = self.hparams.num_iterations_2nd_stage
+    # TODO(mechcoder): Add log_annealing schedule.
+    if schedule == "constant":
+      beta = tf.cond(tf.greater(step_num, second_stage),
+                     lambda: self.hparams.latent_loss_multiplier,
+                     lambda: 0.0)
+    elif schedule == "linear_anneal":
+      # Linearly anneal beta from 0.0 to self.hparams.latent_loss_multiplier.
+      # between self.hparams.num_iterations_2nd_stage to anneal_end.
+      # beta = latent_loss * (1 - (global_step - 2nd_stage) / (anneal_end - 2nd_stage))  # pylint:disable=line-too-long
+      anneal_end = self.hparams.anneal_end
+      latent_multiplier = self.hparams.latent_loss_multiplier
+      if anneal_end < second_stage:
+        raise ValueError("Expected hparams.num_iterations_2nd_stage < "
+                         "hparams.anneal_end %d, got %d." %
+                         (second_stage, anneal_end))
+
+      def anneal_loss(step_num):
+        step_num = tf.cast(step_num, dtype=tf.float32)
+        fraction = (float(anneal_end) - step_num) / (anneal_end - second_stage)
+        return self.hparams.latent_loss_multiplier * (1 - fraction)
+
+      beta = tf.case(
+          pred_fn_pairs={
+              tf.less(step_num, second_stage): lambda: 0.0,
+              tf.greater(step_num, anneal_end): lambda: latent_multiplier},
+          default=lambda: anneal_loss(step_num))
+    else:
+      raise ValueError("Unknown beta schedule.")
+    return beta
+
   def get_scheduled_sample_func(self, batch_size):
     """Creates a function for scheduled sampling based on given hparams."""
     with tf.variable_scope("scheduled_sampling_func", reuse=False):
-      iter_num = tf.train.get_global_step()
-      # TODO(lukaszkaiser): figure out why iter_num can be None.
-      if iter_num is None:
-        iter_num = _LARGE_STEP_NUMBER
-
+      iter_num = self.get_iteration_num()
       if self.hparams.scheduled_sampling_mode == "prob":
         decay_steps = self.hparams.scheduled_sampling_decay_steps
         probability = tf.train.polynomial_decay(
@@ -540,40 +578,7 @@ def body(self, features):
         rewards=all_rewards,
     )
 
-    step_num = tf.train.get_global_step()
-    # TODO(mbz): what should it be if it"s undefined?
-    if step_num is None:
-      step_num = _LARGE_STEP_NUMBER
-
-    schedule = self.hparams.latent_loss_multiplier_schedule
-    second_stage = self.hparams.num_iterations_2nd_stage
-    # TODO(mechcoder): Add log_annealing schedule.
-    if schedule == "constant":
-      beta = tf.cond(tf.greater(step_num, second_stage),
-                     lambda: self.hparams.latent_loss_multiplier,
-                     lambda: 0.0)
-    elif schedule == "linear_anneal":
-      # Linearly anneal beta from 0.0 to self.hparams.latent_loss_multiplier.
-      # between self.hparams.num_iterations_2nd_stage to anneal_end.
-      # beta = latent_loss * (1 - (global_step - 2nd_stage) / (anneal_end - 2nd_stage))  # pylint:disable=line-too-long
-      anneal_end = self.hparams.anneal_end
-      latent_multiplier = self.hparams.latent_loss_multiplier
-      if anneal_end < second_stage:
-        raise ValueError("Expected hparams.num_iterations_2nd_stage < "
-                         "hparams.anneal_end %d, got %d." %
-                         (second_stage, anneal_end))
-
-      def anneal_loss(step_num):
-        step_num = tf.cast(step_num, dtype=tf.float32)
-        fraction = (float(anneal_end) - step_num) / (anneal_end - second_stage)
-        return self.hparams.latent_loss_multiplier * (1 - fraction)
-
-      beta = tf.case(
-          pred_fn_pairs={
-              tf.less(step_num, second_stage): lambda: 0.0,
-              tf.greater(step_num, anneal_end): lambda: latent_multiplier},
-          default=lambda: anneal_loss(step_num))
-
+    beta = self.get_beta()
     kl_loss = 0.0
     if self.is_training:
       for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):

From 27bc4e88d7a621fef19b3885f4aa1a0b6122046f Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 2 Aug 2018 13:14:17 -0700
Subject: [PATCH 0455/2720] make tests fast to run.

PiperOrigin-RevId: 207152631
---
 tensor2tensor/models/research/next_frame_params.py | 2 +-
 tensor2tensor/models/research/next_frame_test.py   | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index 3737fbb27..603a6c4d1 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -44,6 +44,7 @@ def next_frame():
   hparams.add_hparam("filter_double_steps", 2)
   hparams.add_hparam("video_modality_loss_cutoff", 0.02)
   hparams.add_hparam("concatenate_actions", True)
+  hparams.add_hparam("tiny_mode", False)
   return hparams
 
 
@@ -77,7 +78,6 @@ def next_frame_stochastic():
   hparams.add_hparam("scheduled_sampling_decay_steps", 10000)
   hparams.add_hparam("scheduled_sampling_k", 900.0)
   hparams.add_hparam("latent_num_frames", 0)  # 0 means use all frames.
-  hparams.add_hparam("tiny_mode", False)
   hparams.add_hparam("anneal_end", 100000)
   hparams.add_hparam("upsample_method", "conv2d_transpose")
   return hparams
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index cc57e7668..26622bde9 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -37,6 +37,7 @@ def fill_hparams(hparams, in_frames, out_frames):
   p_hparams = problem.get_hparams(hparams)
   hparams.problem = problem
   hparams.problem_hparams = p_hparams
+  hparams.tiny_mode = True
   return hparams
 
 

From bb05a6db7132ebff06ef555b0fec525937116d68 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 2 Aug 2018 13:43:02 -0700
Subject: [PATCH 0456/2720] conv-only reward model based on the current frame.

PiperOrigin-RevId: 207157290
---
 tensor2tensor/layers/common_video.py          |  4 +-
 .../models/research/next_frame_params.py      |  1 +
 .../models/research/next_frame_sv2p.py        | 52 +++++++++----------
 3 files changed, 27 insertions(+), 30 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 4579a2709..8b41d221b 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -35,7 +35,7 @@ def encode_to_shape(inputs, shape, scope):
     w, h = shape[1].value, shape[2].value
     x = inputs
     x = tf.contrib.layers.flatten(x)
-    x = tfl.dense(x, w * h, activation=tf.nn.relu, name="enc_dense")
+    x = tfl.dense(x, w * h, activation=None, name="enc_dense")
     x = tf.reshape(x, (-1, w, h, 1))
     return x
 
@@ -45,7 +45,7 @@ def decode_to_shape(inputs, shape, scope):
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
     x = inputs
     x = tf.contrib.layers.flatten(x)
-    x = tfl.dense(x, shape[2].value, activation=tf.nn.relu, name="dec_dense")
+    x = tfl.dense(x, shape[2].value, activation=None, name="dec_dense")
     x = tf.expand_dims(x, axis=1)
     return x
 
diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index 603a6c4d1..ca76b6d65 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -63,6 +63,7 @@ def next_frame_stochastic():
   hparams.video_modality_loss_cutoff = 0.0
   hparams.add_hparam("stochastic_model", True)
   hparams.add_hparam("reward_prediction", True)
+  hparams.add_hparam("reward_prediction_stop_gradient", False)
   hparams.add_hparam("model_options", "CDNA")
   hparams.add_hparam("num_masks", 10)
   hparams.add_hparam("latent_channels", 1)
diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index dc16bd1b8..2d4e70c24 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -228,16 +228,21 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
     enc2 = tfl.conv2d(hidden4, hidden4.get_shape()[3], [3, 3], strides=(2, 2),
                       padding="SAME", activation=tf.nn.relu, name="conv3")
 
-    # Pass in reward and action.
-    emb_action = common_video.encode_to_shape(
-        action, enc2.get_shape(), "action_enc")
-    emb_reward = common_video.encode_to_shape(
-        input_reward, enc2.get_shape(), "reward_enc")
-    enc2 = tf.concat(axis=3, values=[enc2, emb_action, emb_reward])
+    # Pass in action if exists.
+    if action is not None:
+      emb_action = common_video.encode_to_shape(
+          action, enc2.get_shape(), "action_enc")
+      enc2 = tf.concat(values=[enc2, emb_action], axis=3)
+
+    # Pass in reward if exists.
+    if input_reward is not None:
+      emb_reward = common_video.encode_to_shape(
+          input_reward, enc2.get_shape(), "reward_enc")
+      enc2 = tf.concat(values=[enc2, emb_reward], axis=3)
 
     if latent is not None and not concat_latent:
       with tf.control_dependencies([latent]):
-        enc2 = tf.concat([enc2, latent], 3)
+        enc2 = tf.concat([enc2, latent], axis=3)
 
     enc3 = tfl.conv2d(enc2, hidden4.get_shape()[3], [1, 1], strides=(1, 1),
                       padding="SAME", activation=tf.nn.relu, name="conv4")
@@ -307,18 +312,12 @@ def construct_latent_tower(self, images):
 
       return mean, std
 
-  def reward_prediction(
-      self, input_image, input_reward, action, lstm_state, latent):
+  def reward_prediction(self, input_image, input_reward, action, latent):
     """Builds a reward prediction network."""
     conv_size = self.tinyify([32, 32, 16, 4])
-    lstm_size = self.tinyify([32, 64, 128, 64, 32])
 
     with tf.variable_scope("reward_pred", reuse=tf.AUTO_REUSE):
-      hidden5, _ = self.bottom_part_tower(
-          input_image, input_reward, action, latent,
-          lstm_state, lstm_size, conv_size)
-
-      x = hidden5
+      x = input_image
       x = tfcl.batch_norm(x, updates_collections=None,
                           is_training=self.is_training, scope="reward_bn0")
       x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
@@ -335,7 +334,7 @@ def reward_prediction(
       pred_reward = common_video.decode_to_shape(
           x, input_reward.shape, "reward_dec")
 
-      return pred_reward, lstm_state
+      return pred_reward
 
   def construct_predictive_tower(
       self, input_image, input_reward, action, lstm_state, latent,
@@ -486,7 +485,7 @@ def construct_model(self,
     def process_single_frame(prev_outputs, inputs):
       """Process a single frame of the video."""
       cur_image, cur_reward, action = inputs
-      time_step, prev_image, prev_reward, lstm_states = prev_outputs[:4]
+      time_step, prev_image, prev_reward, lstm_states = prev_outputs
 
       generated_items = [prev_image, prev_reward]
       groundtruth_items = [cur_image, cur_reward]
@@ -499,16 +498,16 @@ def process_single_frame(prev_outputs, inputs):
           input_image, input_reward, action, lstm_states, latent)
 
       if self.hparams.reward_prediction:
-        reward_lstm_states = prev_outputs[4]
-        pred_reward, reward_lstm_states = self.reward_prediction(
-            input_image, input_reward, action, reward_lstm_states, latent)
+        reward_input_image = pred_image
+        if self.hparams.reward_prediction_stop_gradient:
+          reward_input_image = tf.stop_gradient(reward_input_image)
+        pred_reward = self.reward_prediction(
+            reward_input_image, input_reward, action, latent)
       else:
         pred_reward = input_reward
 
       time_step += 1
       outputs = (time_step, pred_image, pred_reward, lstm_states)
-      if self.hparams.reward_prediction:
-        outputs += (reward_lstm_states,)
 
       return outputs
 
@@ -519,14 +518,12 @@ def process_single_frame(prev_outputs, inputs):
       latent = self.get_gaussian_latent(latent_mean, latent_std)
 
     # HACK: Do first step outside to initialize all the variables
-    lstm_states, reward_lstm_states = [None] * 7, [None] * 5
+    lstm_states = [None] * 7
     inputs = images[0], rewards[0], actions[0]
     prev_outputs = (tf.constant(0),
                     tf.zeros_like(images[0]),
                     tf.zeros_like(rewards[0]),
                     lstm_states)
-    if self.hparams.reward_prediction:
-      prev_outputs += (reward_lstm_states,)
 
     initializers = process_single_frame(prev_outputs, inputs)
     first_gen_images = tf.expand_dims(initializers[1], axis=0)
@@ -637,7 +634,6 @@ def construct_model(self, images, actions, rewards):
 
     # LSTM states.
     lstm_state = [None] * 7
-    reward_lstm_state = [None] * 5
 
     # Create scheduled sampling function
     ss_func = self.get_scheduled_sample_func(batch_size)
@@ -667,8 +663,8 @@ def construct_model(self, images, actions, rewards):
           input_image, input_reward, action, lstm_state, latent)
 
       if self.hparams.reward_prediction:
-        pred_reward, reward_lstm_state = self.reward_prediction(
-            input_image, input_reward, action, reward_lstm_state, latent)
+        pred_reward = self.reward_prediction(
+            pred_image, input_reward, action, latent)
       else:
         pred_reward = input_reward
 

From 4dc1ed266aada025fc901a630980048fe9b50860 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 2 Aug 2018 13:43:40 -0700
Subject: [PATCH 0457/2720] Internal change

PiperOrigin-RevId: 207157414
---
 tensor2tensor/layers/common_attention.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index e3341d9ac..d951b119b 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1855,7 +1855,7 @@ def dot_product_unmasked_self_attention_relative_v2(
                                                 length-max_length),
                                              0],
                                             [2*length -1, -1])
-    unmasked_rel_logits = tf.einsum("bhld, md -> bhlm", q,
+    unmasked_rel_logits = tf.einsum("bhld,md->bhlm", q,
                                     used_key_relative_embeddings)
     unmasked_rel_logits = _relative_position_to_absolute_position_unmasked(
         unmasked_rel_logits)
@@ -1896,7 +1896,7 @@ def dot_product_unmasked_self_attention_relative_v2(
                                                0],
                                               [2*length -1, -1])
 
-    ret += tf.einsum("bhlm, md -> bhld", relative_weights,
+    ret += tf.einsum("bhlm,md->bhld", relative_weights,
                      used_value_relative_embeddings)
     return ret
 

From ecbcc1013b9694ec566c751ea95937ae2fb16da6 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 2 Aug 2018 15:13:37 -0700
Subject: [PATCH 0458/2720] Fix lint on next_frame_sv2p.py

PiperOrigin-RevId: 207174029
---
 tensor2tensor/models/research/next_frame_sv2p.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 2d4e70c24..d93bfc49b 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -314,6 +314,9 @@ def construct_latent_tower(self, images):
 
   def reward_prediction(self, input_image, input_reward, action, latent):
     """Builds a reward prediction network."""
+    del action
+    del latent
+
     conv_size = self.tinyify([32, 32, 16, 4])
 
     with tf.variable_scope("reward_pred", reuse=tf.AUTO_REUSE):

From c3dfd0e235cb5ad706592cf8daf36f140a90e76c Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 2 Aug 2018 16:14:30 -0700
Subject: [PATCH 0459/2720] v 1.7.0

PiperOrigin-RevId: 207183874
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 409486722..33f2b8526 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.6.6',
+    version='1.7.0',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',

From e7d5e6c111ca4d5b97bae366218a4f229819d489 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 2 Aug 2018 16:57:14 -0700
Subject: [PATCH 0460/2720] resizing frames.

PiperOrigin-RevId: 207190244
---
 tensor2tensor/data_generators/video_utils.py       | 10 +++++-----
 tensor2tensor/models/research/next_frame_params.py |  1 +
 tensor2tensor/models/research/next_frame_sv2p.py   |  5 ++---
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 11471d089..46bc80acf 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -140,6 +140,11 @@ def use_not_breaking_batching(self):
 
   def preprocess_example(self, example, mode, hparams):
     """Runtime preprocessing, e.g., resize example["frame"]."""
+    if hparams.preprocess_resize_frames is not None:
+      example["frame"] = tf.image.resize_images(
+          example["frame"],
+          hparams.preprocess_resize_frames,
+          tf.image.ResizeMethod.BILINEAR)
     return example
 
   @property
@@ -219,11 +224,6 @@ def features_from_batch(batched_prefeatures):
       for k, v in six.iteritems(batched_prefeatures):
         if k == "frame":  # We rename past frames to inputs and targets.
           s1, s2 = split_on_batch(v)
-          # Reshape just to make sure shapes are right and set.
-          s1 = tf.reshape(
-              s1, [hparams.video_num_input_frames] + self.frame_shape)
-          s2 = tf.reshape(
-              s2, [hparams.video_num_target_frames] + self.frame_shape)
           features["inputs"] = s1
           features["targets"] = s2
         else:
diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index ca76b6d65..c36c30765 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -43,6 +43,7 @@ def next_frame():
   hparams.add_hparam("num_compress_steps", 6)
   hparams.add_hparam("filter_double_steps", 2)
   hparams.add_hparam("video_modality_loss_cutoff", 0.02)
+  hparams.add_hparam("preprocess_resize_frames", None)
   hparams.add_hparam("concatenate_actions", True)
   hparams.add_hparam("tiny_mode", False)
   return hparams
diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index d93bfc49b..43f0a2ff6 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -344,7 +344,8 @@ def construct_predictive_tower(
       concat_latent=False):
     # Main tower
     lstm_func = common_video.conv_lstm_2d
-    batch_size = common_layers.shape_list(input_image)[0]
+    frame_shape = common_layers.shape_list(input_image)
+    batch_size, img_height, img_width, color_channels = frame_shape
     # the number of different pixel motion predictions
     # and the number of masks for each of those predictions
     num_masks = self.hparams.num_masks
@@ -354,8 +355,6 @@ def construct_predictive_tower(
     lstm_size = self.tinyify([32, 32, 64, 64, 128, 64, 32])
     conv_size = self.tinyify([32])
 
-    img_height, img_width, color_channels = self.hparams.problem.frame_shape
-
     with tf.variable_scope("main", reuse=tf.AUTO_REUSE):
       hidden5, skips = self.bottom_part_tower(
           input_image, input_reward, action, latent,

From 1d6a3bb9df9af99d97f64a19bc8701af515e515b Mon Sep 17 00:00:00 2001
From: Keyon Vafa <vafa@google.com>
Date: Thu, 2 Aug 2018 18:14:55 -0700
Subject: [PATCH 0461/2720] change transformer_vae.loss() so it uses the
 sum_over_latents flag.

PiperOrigin-RevId: 207200375
---
 tensor2tensor/models/research/transformer_vae.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 4877544c7..b0f77a9bd 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -619,7 +619,11 @@ def loss(self, logits, features):
       logits = tf.cast(logits, tf.float32)
       xent = common_layers.smoothing_cross_entropy(
           logits, labels, vocab_size, confidence=1.0, gaussian=False)
-      return tf.reduce_sum(xent) / tf.cast(logits_shape[0], tf.float32)
+      if self._hparams.sum_over_latents:
+        recon_loss = tf.reduce_sum(xent) / tf.cast(logits_shape[0], tf.float32)
+      else:
+        recon_loss = tf.reduce_mean(xent)
+      return recon_loss
 
   def body(self, features):
     inputs = features["inputs"] if "inputs" in features else None

From 7b532bdf98cf0f7f9d31fda88a9891708b57f1fe Mon Sep 17 00:00:00 2001
From: Keyon Vafa <vafa@google.com>
Date: Thu, 2 Aug 2018 19:31:40 -0700
Subject: [PATCH 0462/2720] Include transformer layers when using IAF

PiperOrigin-RevId: 207207082
---
 tensor2tensor/layers/discretization.py        | 175 ++++++++++--------
 .../models/research/transformer_vae.py        |   6 +-
 2 files changed, 98 insertions(+), 83 deletions(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 42918af22..c55f64b69 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -19,7 +19,8 @@
 
 from functools import partial
 
-from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_hparams
+from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
 
 import tensorflow as tf
@@ -143,7 +144,7 @@ def embedding_lookup(x,
                      num_samples=1,
                      do_hard_gumbel_softmax=False,
                      temperature_warmup_steps=150000,
-                     do_iaf=False,
+                     num_flows=0,
                      approximate_gs_entropy=False,
                      sum_over_latents=False):
   """Compute nearest neighbors and loss for training the embeddings via DVQ.
@@ -162,7 +163,7 @@ def embedding_lookup(x,
       for gumbel-softmax-dvq bottleneck.
     temperature_warmup_steps: Number of steps it takes to decay temperature to
       0. Used only if bottleneck_kind is gumbel-softmax-dvq.
-    do_iaf: Whether to apply inverse autoregressive flows for gumbel-softmax-dvq
+    num_flows: Number of inverse autoregressive flows for gumbel-softmax-dvq
       bottleneck.
     approximate_gs_entropy: Whether to approximate the Gumbel-Softmax density
       as a categorical distribution when calculating the sample entropy. Used
@@ -189,7 +190,7 @@ def embedding_lookup(x,
         hard=do_hard_gumbel_softmax,
         num_samples=num_samples,
         temperature_warmup_steps=temperature_warmup_steps,
-        do_iaf=do_iaf,
+        num_flows=num_flows,
         approximate_gs_entropy=approximate_gs_entropy,
         sum_over_latents=sum_over_latents)
   else:
@@ -492,7 +493,7 @@ def discrete_bottleneck(inputs,
                         softmax_k=0,
                         temperature_warmup_steps=150000,
                         do_hard_gumbel_softmax=False,
-                        do_iaf=False,
+                        num_flows=0,
                         approximate_gs_entropy=False,
                         sum_over_latents=False,
                         discrete_mix=0.5,
@@ -540,7 +541,7 @@ def discrete_bottleneck(inputs,
       0. Used only if bottleneck_kind is gumbel-softmax or gumbel-softmax-dvq.
     do_hard_gumbel_softmax: Whether to use hard or soft Gumbel-Softmax
       samples. Used only if bottleneck_kind is gumbel-softmax-dvq.
-    do_iaf: Whether to apply inverse autoregresive flows. Used only if
+    num_flows: Number of inverse autoregresive flows. Used only if
       bottleneck_kind is gumbel-softmax-dvq.
     approximate_gs_entropy: Whether to approximate the Gumbel-Softmax density
       as a categorical distribution when calculating the sample entropy. Used
@@ -651,7 +652,7 @@ def discrete_bottleneck(inputs,
                              num_samples=num_samples,
                              temperature_warmup_steps=temperature_warmup_steps,
                              do_hard_gumbel_softmax=do_hard_gumbel_softmax,
-                             do_iaf=do_iaf,
+                             num_flows=num_flows,
                              approximate_gs_entropy=approximate_gs_entropy,
                              sum_over_latents=sum_over_latents))
         # Update the EMA variables.
@@ -888,7 +889,7 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
                                         num_samples=1,
                                         temperature_warmup_steps=150000,
                                         summary=True,
-                                        do_iaf=False,
+                                        num_flows=0,
                                         approximate_gs_entropy=False,
                                         sum_over_latents=False):
   """Sample from Gumbel-Softmax and compute neighbors and losses.
@@ -907,8 +908,8 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
       (Default: 150000).
     summary: When `True`, we save histogram summaries of the KL term (Default:
       True).
-    do_iaf: When `True`, we perform inverse autoregressive flow with
-      Gumbel-Softmax sample (Default: False).
+    num_flows: Number of inverse autoregressive flows with Gumbel-Softmax
+      samples.
     approximate_gs_entropy: When `True`, we approximate Gumbel-Softmax
       density as categorical when calculating sample entropy (Default: False).
     sum_over_latents: Whether to sum over non-batch dimensions when calculating
@@ -975,77 +976,52 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
     neg_q_entropy = tf.reduce_sum(neg_q_entropy, [1, 2])
   neg_q_entropy = tf.reduce_mean(neg_q_entropy)
 
-  if do_iaf:
+  if num_flows > 0:
+    hparams = iaf_hparams(hidden_size=512, filter_size=4096)
     q_samples = tf.reshape(q_samples, [-1, latent_dim, block_v_size])
-
-    # Shift samples so log_pi[:, i, :] is only a function of
-    # q_samples[:, :i, :]. We do this by adding a first row of zeros to the
-    # latents, shifting the other rows down by one, and removing the last row.
-
-    top_latent = tf.zeros([batch_size * num_blocks, 1, block_v_size])
-    shifted_samples = tf.concat([top_latent, q_samples[:, :-1, :]], axis=1)
-
-    d_k = 64
-    d_v = 64
-    query_projection = tf.get_variable(
-        "query_projection", [block_v_size, d_k], dtype=tf.float32)
-    keys_projection = tf.get_variable(
-        "keys_projection", [block_v_size, d_k], dtype=tf.float32)
-    values_projection = tf.get_variable(
-        "values_projection", [block_v_size, d_v], dtype=tf.float32)
-    query = tf.reduce_sum(tf.expand_dims(shifted_samples, -1) *
-                          tf.reshape(query_projection,
-                                     [1, 1, block_v_size, d_k]), 2)
-    keys = tf.reduce_sum(tf.expand_dims(shifted_samples, -1) *
-                         tf.reshape(keys_projection,
-                                    [1, 1, block_v_size, d_k]), 2)
-    values = tf.reduce_sum(tf.expand_dims(shifted_samples, -1) *
-                           tf.reshape(values_projection,
-                                      [1, 1, block_v_size, d_v]), 2)
-
-    # Masked self-attention with a single head.
-    # TODO(vafa): Add support for multiple heads
-    attention_output = common_attention.masked_local_attention_1d(
-        q=tf.expand_dims(query, 1),
-        k=tf.expand_dims(keys, 1),
-        v=tf.expand_dims(values, 1),
-        block_length=1)
-    attention_output = tf.reshape(
-        attention_output, [-1] + common_layers.shape_list(attention_output)[2:])
-
-    ffn_output = common_layers.conv_relu_conv(
-        attention_output,
-        filter_size=64,
-        output_size=block_v_size,
-        first_kernel_size=3,
-        second_kernel_size=1,
-        padding="LEFT",
-        nonpadding_mask=None,
-        dropout=0.,
-        cache=None,
-        decode_loop_step=None)
-
-    log_pi = tf.nn.log_softmax(ffn_output)
-
-    # Flow 1: Adding log_pi to q_samples and dividing by the temperature.
-    # Note that we drop last dimension of q_samples for centered-softmax, which
-    # we can do without recalculating probabilities because the last dimension
-    # of log_pi and q_samples are deterministic given the other dimensions.
-    # Flow 2: Centered-softmax.
-
-    chained_bijectors = tf.contrib.distributions.bijectors.Chain(
-        [tf.contrib.distributions.bijectors.SoftmaxCentered(),
-         tf.contrib.distributions.bijectors.Affine(
-             shift=log_pi[:, :, :-1],
-             scale_identity_multiplier=1./temperature)])
-    q_samples = chained_bijectors.forward(q_samples[:, :, :-1])
-    log_det = chained_bijectors.inverse_log_det_jacobian(
-        q_samples, event_ndims=1)
-    log_det = tf.reshape(log_det,
-                         [num_samples, batch_size, num_blocks, latent_dim])
-    if sum_over_latents:
-      log_det = tf.reduce_sum(log_det, axis=[2, 3])
-    neg_q_entropy += tf.reduce_mean(log_det)
+    for flow in range(num_flows):
+      shifted_samples = tf.pad(q_samples, [[0, 0], [1, 0], [0, 0]])[:, :-1, :]
+
+      # Project samples from  [batch_size, latent_size, block_v_size] to
+      # [batch_size, latent_size, hidden_size].
+      shifted_samples = common_layers.dense(
+          shifted_samples, hparams.hidden_size)
+      # TODO(vafa): Include masking as a flag.
+      mask = True
+      if mask:
+        attention_type = cia.AttentionType.LOCAL_1D
+      else:
+        attention_type = cia.AttentionType.GLOBAL
+      ffn_output = cia.transformer_decoder_layers(
+          inputs=shifted_samples,
+          encoder_output=None,
+          num_layers=6,
+          hparams=hparams,
+          attention_type=attention_type,
+          name="transformer_" + str(flow))
+
+      # Project samples back to [batch_size, latent_size, block_v_size].
+      ffn_output = common_layers.dense(ffn_output, block_v_size)
+      log_pi = tf.nn.log_softmax(ffn_output)
+
+      # Flow 1: Adding log_pi to q_samples and dividing by the temperature.
+      # Note that we drop the last dimension of q_samples for centered-softmax,
+      # which we can do without recalculating probabilities because the last
+      # dimension of log_pi and q_samples are deterministic given the others.
+      # Flow 2: Centered-softmax.
+      chained_bijectors = tf.contrib.distributions.bijectors.Chain(
+          [tf.contrib.distributions.bijectors.SoftmaxCentered(),
+           tf.contrib.distributions.bijectors.Affine(
+               shift=log_pi[:, :, :-1],
+               scale_identity_multiplier=1./temperature)])
+      q_samples = chained_bijectors.forward(q_samples[:, :, :-1])
+      log_det = chained_bijectors.inverse_log_det_jacobian(
+          q_samples, event_ndims=1)
+      log_det = tf.reshape(log_det,
+                           [num_samples, batch_size, num_blocks, latent_dim])
+      if sum_over_latents:
+        log_det = tf.reduce_sum(log_det, axis=[2, 3])
+      neg_q_entropy += tf.reduce_mean(log_det)
 
     q_samples = tf.reshape(
         q_samples,
@@ -1297,3 +1273,42 @@ def parametrized_unbottleneck(x, hidden_size, hparams):
     return vq_discrete_unbottleneck(x, hidden_size)
   raise ValueError("Unsupported hparams.bottleneck_kind %s"
                    % hparams.bottleneck_kind)
+
+
+def iaf_hparams(hidden_size=512, filter_size=4096):
+  """Create hyperpameters for inverse autoregressive flows.
+
+  Args:
+    hidden_size: Width of attention layers and neural network output layer.
+    filter_size: Hidden layer width for neural network.
+
+  Returns:
+    hparams: Hyperpameters with basic presets for inverse autoregressive flows.
+  """
+  hparams = common_hparams.basic_params1()
+
+  # Attention hyperparameters.
+  hparams.hidden_size = hidden_size
+  hparams.add_hparam("attention_key_channels", None)
+  hparams.add_hparam("attention_value_channels", None)
+  hparams.add_hparam("num_heads", 4)
+  hparams.add_hparam("attention_dropout", 0.1)
+  hparams.add_hparam("shared_rel", False)
+  hparams.add_hparam("block_width", 1)
+  hparams.add_hparam("block_length", 1)
+  hparams.add_hparam("q_filter_width", 1)
+  hparams.add_hparam("kv_filter_width", 1)
+
+  # Preprocessing and postprocesing hyperparameters.
+  hparams.layer_preprocess_sequence = "n"
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.norm_type = "layer"
+  hparams.norm_epsilon = 1e-06
+  hparams.layer_prepostprocess_dropout_broadcast_dims = ""
+  hparams.layer_postprocess_sequence = "da"
+
+  # Feedforward neural network hyperparameters.
+  hparams.add_hparam("filter_size", filter_size)
+  hparams.add_hparam("ffn_layer", "conv_hidden_relu")
+  hparams.add_hparam("relu_dropout", 0.1)
+  return hparams
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index b0f77a9bd..98ecb4250 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -528,7 +528,7 @@ def __init__(self, *args, **kwargs):
         softmax_k=self._hparams.softmax_k,
         temperature_warmup_steps=self._hparams.temperature_warmup_steps,
         do_hard_gumbel_softmax=self._hparams.do_hard_gumbel_softmax,
-        do_iaf=self._hparams.do_iaf,
+        num_flows=self._hparams.num_flows,
         approximate_gs_entropy=self._hparams.approximate_gs_entropy,
         discrete_mix=self._hparams.d_mix,
         noise_dev=self._hparams.noise_dev,
@@ -760,7 +760,7 @@ def transformer_ae_small():
   hparams.add_hparam("entropy_scale", 0.0)
   hparams.add_hparam("prior_scale", 1.0)
   hparams.add_hparam("do_hard_gumbel_softmax", False)
-  hparams.add_hparam("do_iaf", False)
+  hparams.add_hparam("num_flows", 0)
   hparams.add_hparam("approximate_gs_entropy", False)
   hparams.add_hparam("temperature_warmup_steps", 150000)
   hparams.add_hparam("sum_over_latents", False)
@@ -964,6 +964,6 @@ def transformer_ae_base_ablation_5():
 @registry.register_hparams
 def transformer_ae_base_iaf():
   hparams = transformer_ae_base_ablation_5()
-  hparams.do_iaf = True
+  hparams.num_flows = 1
   hparams.num_samples = 1
   return hparams

From 7967b446c5a0f7e7ea2e6b39150abf850abec463 Mon Sep 17 00:00:00 2001
From: Christopher Beitel <cwbeitel@users.noreply.github.com>
Date: Thu, 2 Aug 2018 22:14:56 -0700
Subject: [PATCH 0463/2720] img2img_allen_brain problem defs (#872)

* img2img_allen_brain problem defs

only outstanding issue blocking this should be that the test in allen_brain_test.py that uses tf.eager breaks many/all other t2t tests not compatible with eager mode.

squashed these:

commit 76c8cd3f105c07919f84d70b02b1fbd38b5376fe
Author: cwbeitel <me@cwbeitel-github.com>
Date:   Fri Jun 15 18:35:13 2018 +0000

    fix mock_raw_data

commit f0dc3a0fcb6691c5b8854f419485966cb2f7b9c2
Merge: 3e17473 4b7360d
Author: cwbeitel <me@cwbeitel-github.com>
Date:   Thu Jun 14 22:54:33 2018 +0000

    Merge branch 'allen-deps' of https://github.com/cwbeitel/tensor2tensor into allen-deps

commit 3e174736ec6628fe0b8097c1a3b0c2b3646ddee9
Author: cwbeitel <me@cwbeitel-github.com>
Date:   Fri Jun 8 17:04:39 2018 +0000

    cleanup

commit 3e8ef9d61382afd795a26a16ffb254974a80dc8e
Author: cwbeitel <me@cwbeitel-github.com>
Date:   Wed Jun 6 23:07:28 2018 +0000

    removed image summary writing from metrics

    training was 10x slower than CIFAR img2img problem, profiled and found 91% of time was being sp
ent writing ImageSummary's, removed that from metrics and saw the speed increase as expected.

commit 481455649596724da4ab7878eff7d84ed495af44
Author: cwbeitel <me@cwbeitel-github.com>
Date:   Tue May 29 16:00:17 2018 +0000

    basic img2img prob. defs that maybe get data

commit 320f06c05697e35a108e74e7d3a0fbd29bebee3b
Author: cwbeitel <me@cwbeitel-github.com>
Date:   Thu May 24 18:53:02 2018 +0000

    test without accessing network, follow rcfile

    * made test that accesses remote API optional, off by default
    * included test that runs sub-image extractor on mocked raw input image within mocked directory
 structure mirroring what would be produced by running actual data downloader
    * refactoring according to rcfile

commit 503123b09e1d784f9ca441c4b3f723d99baba7fb
Author: cwbeitel <me@cwbeitel-github.com>
Date:   Wed May 23 22:09:31 2018 +0000

    add [allen] deps to travis setup

commit 239188b8c3e28d6fb0fb31c0d9167291874dc2f1
Author: cwbeitel <me@cwbeitel-github.com>
Date:   Wed May 23 21:39:13 2018 +0000

    data download and sub-image gen utility

    * if additional [allen] deps not present suggest running `pip install tensor2tensor[allen]`
    * obtains a list of section dataset ID's from Allen Institute API
    * downloads specified number of raw images from specified number of section datasets
    * produces specified number of sub-images of specified size from raw images.
    * simple e2e test that runs the data downloader and sub-image generator with default parameters

commit af34cc660c90663b662eaa2bc7a5a66cfadf8762
Author: cwbeitel <me@cwbeitel-github.com>
Date:   Wed May 23 17:47:35 2018 +0000

    deps. for allen brain img probs. via extras

* run allen tests separately on ci

* travis debug

* include previously excluded setup.py that defines [allen] step

* smaller e2e test, maybe addresses travis OOM

* linter fixes

* Remove allensdk dependency; various improvements

* Simplify allen_brain_utils* to download smaller static list of images, avoiding allensdk dependency.

* Simplify _generator given input is already in [0, 255]

* _generator ignores likely non-tissue/background regions based on max intensity

* Fractional in-painting step added to pre-processor, depends on inpaint_fraction property, randomly positioned square region

* Random crop to target dimensions instead of downsampling leaving full resolution of target intact

* Re-using same generated data for different problems that share examples (differing in terms of how these are pre-processed)

* cleanup
---
 .travis.yml                                   |   8 +-
 setup.py                                      |   1 +
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/allen_brain.py  | 433 ++++++++++++++++++
 .../data_generators/allen_brain_test.py       | 174 +++++++
 .../data_generators/allen_brain_utils.py      | 108 +++++
 .../data_generators/allen_brain_utils_test.py |  83 ++++
 7 files changed, 807 insertions(+), 1 deletion(-)
 create mode 100644 tensor2tensor/data_generators/allen_brain.py
 create mode 100644 tensor2tensor/data_generators/allen_brain_test.py
 create mode 100644 tensor2tensor/data_generators/allen_brain_utils.py
 create mode 100644 tensor2tensor/data_generators/allen_brain_utils_test.py

diff --git a/.travis.yml b/.travis.yml
index b45cbcb6a..e4c209350 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -42,7 +42,7 @@ install:
   - pip install -q .
   - t2t-trainer --registry_help
   # Then install the test dependencies
-  - pip install -q .[tests]
+  - pip install -q .[tests,allen]
   # Make sure to install the atari extras for gym
   - pip install "gym[atari]"
   # Make sure we have the latest version of numpy - avoid problems we were
@@ -59,6 +59,8 @@ script:
   #   * trainer_lib_test
   #   * visualization_test
   #   * model_rl_experiment_test
+  #   * allen_brain_test
+  #   * allen_brain_utils_test
   #   * model_rl_experiment_stochastic_test
   #   * models/research
   # algorithmic_math_test: flaky
@@ -71,11 +73,15 @@ script:
     --ignore=tensor2tensor/data_generators/algorithmic_math_test.py
     --ignore=tensor2tensor/models/research/universal_transformer_test.py
     --ignore=tensor2tensor/rl/model_rl_experiment_test.py
+    --ignore=tensor2tensor/data_generators/allen_brain_test.py
+    --ignore=tensor2tensor/data_generators/allen_brain_utils_test.py
     --ignore=tensor2tensor/rl/model_rl_experiment_stochastic_test.py
     --ignore=tensor2tensor/models/research
   - pytest tensor2tensor/utils/registry_test.py
   - pytest tensor2tensor/utils/trainer_lib_test.py
   - pytest tensor2tensor/visualization/visualization_test.py
+  - pytest tensor2tensor/data_generators/allen_brain_test.py
+  - pytest tensor2tensor/data_generators/allen_brain_utils_test.py
   - if [[ "$TF_VERSION" == "$TF_LATEST"  ]] || [[ "$TF_VERSION" == "tf-nightly"  ]];
     then
       pytest tensor2tensor/models/research;
diff --git a/setup.py b/setup.py
index 33f2b8526..e9740ad00 100644
--- a/setup.py
+++ b/setup.py
@@ -58,6 +58,7 @@
             # explicit pip install gym[atari] for the tests.
             # 'gym[atari]',
         ],
+        'allen': ['Pillow==5.1.0', 'pandas==0.23.0']
     },
     classifiers=[
         'Development Status :: 4 - Beta',
diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 3621862a8..792cd1896 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -23,6 +23,7 @@
 MODULES = [
     "tensor2tensor.data_generators.algorithmic",
     "tensor2tensor.data_generators.algorithmic_math",
+    "tensor2tensor.data_generators.allen_brain",
     "tensor2tensor.data_generators.audio",
     "tensor2tensor.data_generators.babi_qa",
     "tensor2tensor.data_generators.bair_robot_pushing",
diff --git a/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
new file mode 100644
index 000000000..59722d3fe
--- /dev/null
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -0,0 +1,433 @@
+# coding=utf-8
+# 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.
+
+"""Problem definitions for Allen Brain Atlas problems.
+
+Notes:
+
+  * TODO(cwbeitel): Want to be able to increase up-sampling ratio and/or
+    in-paint fraction over the course of training. This could be done by
+    defining a range of problems or perhaps more aptly with an hparam
+    that is dialed up depending on training performance.
+
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from io import BytesIO
+import math
+import numpy as np
+import os
+import requests
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import image_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import metrics
+
+from tensor2tensor.data_generators.allen_brain_utils import try_importing_pil_image
+
+import tensorflow as tf
+
+_BASE_EXAMPLE_IMAGE_SIZE = 64
+
+
+# A 100 image random subset of non-failed acquisitions of Mouse imaging
+# products from Allen Brain Institute (api.brain-map.org) dataset. The
+# full set (or a desired subset) of image IDs can be obtained following
+# the steps described here: http://help.brain-map.org/display/api,
+# e.g. https://gist.github.com/cwbeitel/5dffe90eb561637e35cdf6aa4ee3e704
+_IMAGE_IDS = [
+    '74887117', '71894997', '69443979', '79853548', '101371232', '77857182',
+    '70446772', '68994990', '69141561', '70942310', '70942316', '68298378',
+    '69690156', '74364867', '77874134', '75925043', '73854431', '69206601',
+    '71771457', '101311379', '74777533', '70960269', '71604493', '102216720',
+    '74776437', '75488723', '79815814', '77857132', '77857138', '74952778',
+    '69068486', '648167', '75703410', '74486118', '77857098', '637407',
+    '67849516', '69785503', '71547630', '69068504', '69184074', '74853078',
+    '74890694', '74890698', '75488687', '71138602', '71652378', '68079764',
+    '70619061', '68280153', '73527042', '69764608', '68399025', '244297',
+    '69902658', '68234159', '71495521', '74488395', '73923026', '68280155',
+    '75488747', '69589140', '71342189', '75119214', '79455452', '71774294',
+    '74364957', '68031779', '71389422', '67937572', '69912671', '73854471',
+    '75008183', '101371376', '75703290', '69533924', '79853544', '77343882',
+    '74887133', '332587', '69758622', '69618413', '77929999', '244293',
+    '334792', '75825136', '75008103', '70196678', '71883965', '74486130',
+    '74693566', '76107119', '76043858', '70252433', '68928364', '74806345',
+    '67848661', '75900326', '71773690', '75008171']
+
+
+def _get_case_file_paths(tmp_dir, case, training_fraction=0.95):
+  """Obtain a list of image paths corresponding to training or eval case.
+
+  Args:
+    tmp_dir: str, the root path to which raw images were written, at the
+      top level having meta/ and raw/ subdirs.
+    size: int, the size of sub-images to consider (`size`x`size`).
+    case: bool, whether obtaining file paths for training (true) or eval
+      (false).
+    training_fraction: float, the fraction of the sub-image path list to
+      consider as the basis for training examples.
+
+  Returns:
+    list: A list of file paths.
+  """
+
+  paths = tf.gfile.Glob("%s/*.jpg" % tmp_dir)
+
+  if not paths:
+    raise ValueError("Search of tmp_dir (%s) " % tmp_dir,
+                     "for subimage paths yielded an empty list, ",
+                     "can't proceed with returning training/eval split.")
+
+  split_index = int(math.floor(len(paths)*training_fraction))
+
+  if split_index >= len(paths):
+    raise ValueError("For a path list of size %s "
+                     "and a training_fraction of %s "
+                     "the resulting split_index of the paths list, "
+                     "%s, would leave no elements for the eval "
+                     "condition." % (len(paths),
+                                     training_fraction,
+                                     split_index))
+
+  if case:
+    return paths[:split_index]
+  else:
+    return paths[split_index:]
+
+
+def maybe_download_image_dataset(image_ids, target_dir):
+  """Download a set of images from api.brain-map.org to `target_dir`.
+
+  Args:
+    image_ids: list, a list of image ids.
+    target_dir: str, a directory to which to download the images.
+  """
+
+  tf.gfile.MakeDirs(target_dir)
+
+  num_images = len(image_ids)
+
+  for i, image_id in enumerate(image_ids):
+
+    destination = os.path.join(target_dir, "%s.jpg" % i)
+    tmp_destination = "%s.temp" % destination
+
+    source_url = ("http://api.brain-map.org/api/v2/"
+                  "section_image_download/%s" % image_id)
+
+    if tf.gfile.Exists(destination):
+      tf.logging.info("Image with ID already present, "
+                      "skipping download (%s of %s)." % (
+                          i+1, num_images
+                      ))
+      continue
+
+    tf.logging.info("Downloading image with id %s (%s of %s)" % (
+        image_id, i+1, num_images
+    ))
+
+    response = requests.get(source_url, stream=True)
+
+    response.raise_for_status()
+
+    with open(tmp_destination, "w") as f:
+      for block in response.iter_content(1024):
+        f.write(block)
+
+    tf.gfile.Rename(tmp_destination, destination)
+
+
+def random_square_mask(shape, fraction):
+  """Create a numpy array with specified shape and masked fraction.
+
+  Args:
+    shape: tuple, shape of the mask to create.
+    fraction: float, fraction of the mask area to populate with `mask_scalar`.
+    mask_scalar: float, the scalar to apply to the otherwise 1-valued mask.
+
+  Returns:
+    numpy.array: A numpy array storing the mask.
+  """
+
+  mask = np.ones(shape)
+
+  patch_area = shape[0]*shape[1]*fraction
+  patch_dim = np.int(math.floor(math.sqrt(patch_area)))
+  if patch_area == 0 or patch_dim == 0:
+    return mask
+
+  x = np.random.randint(shape[0] - patch_dim)
+  y = np.random.randint(shape[1] - patch_dim)
+
+  mask[x:(x + patch_dim), y:(y + patch_dim), :] = 0
+
+  return mask
+
+
+def _generator(tmp_dir, training, size=_BASE_EXAMPLE_IMAGE_SIZE,
+               training_fraction=0.95):
+  """Base problem example generator for Allen Brain Atlas problems.
+
+  Args:
+
+    tmp_dir: str, a directory where raw example input data has been stored.
+    training: bool, whether the mode of operation is training (or,
+      alternatively, evaluation), determining whether examples in tmp_dir
+      prefixed with train or dev will be used.
+    size: int, the image size to add to the example annotation.
+
+  Yields:
+    A dictionary representing the images with the following fields:
+      * image/encoded: The string encoding the image as JPEG.
+      * image/format: The string "jpeg" indicating the image format.
+      * image/height: The integer indicating the image height.
+      * image/width: The integer indicating the image height.
+
+  """
+
+  maybe_download_image_dataset(_IMAGE_IDS, tmp_dir)
+
+  image_files = _get_case_file_paths(tmp_dir=tmp_dir,
+                                     case=training,
+                                     training_fraction=training_fraction)
+
+  image_obj = try_importing_pil_image()
+
+  tf.logging.info("Loaded case file paths (n=%s)" % len(image_files))
+  height = size
+  width = size
+
+  for input_path in image_files:
+
+    img = image_obj.open(input_path)
+    img = np.float32(img)
+    shape = np.shape(img)
+
+    for h_index in range(0, int(math.floor(shape[0]/size))):
+
+      h_offset = h_index * size
+      h_end = h_offset + size - 1
+
+      for v_index in range(0, int(math.floor(shape[1]/size))):
+
+        v_offset = v_index * size
+        v_end = v_offset + size - 1
+
+        # Extract a sub-image tile.
+        # pylint: disable=invalid-sequence-index
+        subimage = np.uint8(img[h_offset:h_end, v_offset:v_end])
+
+        # Filter images that are likely background (not tissue).
+        if np.amax(subimage) < 230:
+          continue
+
+        subimage = image_obj.fromarray(subimage)
+        buff = BytesIO()
+        subimage.save(buff, format="JPEG")
+        subimage_encoded = buff.getvalue()
+
+        yield {
+            "image/encoded": [subimage_encoded],
+            "image/format": ["jpeg"],
+            "image/height": [height],
+            "image/width": [width]
+        }
+
+
+@registry.register_problem
+class Img2imgAllenBrain(problem.Problem):
+  """Allen Brain Atlas histology dataset.
+
+  See also: http://help.brain-map.org/
+
+  Notes:
+
+    * 64px to 64px identity mapping problem, no in-painting.
+
+  """
+
+  @property
+  def train_shards(self):
+    return 100
+
+  @property
+  def dev_shards(self):
+    return 10
+
+  @property
+  def training_fraction(self):
+    return 0.95
+
+  @property
+  def num_channels(self):
+    """Number of color channels."""
+    return 3
+
+  @property
+  def input_dim(self):
+    """The x and y dimension of the input image."""
+    # By default, there is no input image, only a target.
+    return 64
+
+  @property
+  def output_dim(self):
+    """The x and y dimension of the target image."""
+    return 64
+
+  @property
+  def inpaint_fraction(self):
+    """The fraction of the input image to be in-painted."""
+    # By default, no in-painting is performed.
+    return None
+
+  def preprocess_example(self, example, mode, hparams):
+
+    # Crop to target shape instead of down-sampling target, leaving target
+    # of maximum available resolution.
+    target_shape = (self.output_dim, self.output_dim, self.num_channels)
+    example["targets"] = tf.random_crop(example["targets"], target_shape)
+
+    example["inputs"] = image_utils.resize_by_area(example["targets"],
+                                                   self.input_dim)
+
+    if self.inpaint_fraction is not None and self.inpaint_fraction > 0:
+
+      mask = random_square_mask((self.input_dim,
+                                 self.input_dim,
+                                 self.num_channels),
+                                self.inpaint_fraction)
+
+      example["inputs"] = tf.multiply(
+          tf.convert_to_tensor(mask, dtype=tf.int64),
+          example["inputs"])
+
+      if self.input_dim is None:
+        raise ValueError("Cannot train in-painting for examples with "
+                         "only targets (i.e. input_dim is None, "
+                         "implying there are only targets to be "
+                         "generated).")
+
+    return example
+
+  def feature_encoders(self, data_dir):
+    del data_dir
+    return {
+        "inputs": text_encoder.ImageEncoder(channels=self.num_channels),
+        "targets": text_encoder.ImageEncoder(channels=self.num_channels)
+    }
+
+  def example_reading_spec(self):
+    data_fields = {
+        "image/encoded": tf.FixedLenFeature((), tf.string),
+        "image/format": tf.FixedLenFeature((), tf.string),
+    }
+
+    data_items_to_decoders = {
+        "targets":
+            tf.contrib.slim.tfexample_decoder.Image(
+                image_key="image/encoded",
+                format_key="image/format",
+                channels=self.num_channels),
+    }
+
+    return data_fields, data_items_to_decoders
+
+  def eval_metrics(self):
+    eval_metrics = [
+        metrics.Metrics.ACC,
+        metrics.Metrics.ACC_PER_SEQ,
+        metrics.Metrics.NEG_LOG_PERPLEXITY
+    ]
+    return eval_metrics
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    generator_utils.generate_dataset_and_shuffle(
+        self.generator(tmp_dir, True),
+        self.training_filepaths(data_dir, self.train_shards, shuffled=True),
+        self.generator(tmp_dir, False),
+        self.dev_filepaths(data_dir, self.dev_shards, shuffled=True))
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    p.input_modality = {"inputs": ("image:identity", 256)}
+    p.target_modality = ("image:identity", 256)
+    p.batch_size_multiplier = 256
+    p.input_space_id = problem.SpaceID.IMAGE
+    p.target_space_id = problem.SpaceID.IMAGE
+
+  def generator(self, tmp_dir, is_training):
+    if is_training:
+      return _generator(tmp_dir, True, size=_BASE_EXAMPLE_IMAGE_SIZE,
+                        training_fraction=self.training_fraction)
+    else:
+      return _generator(tmp_dir, False, size=_BASE_EXAMPLE_IMAGE_SIZE,
+                        training_fraction=self.training_fraction)
+
+
+@registry.register_problem
+class Img2imgAllenBrainDim48to64(Img2imgAllenBrain):
+  """48px to 64px resolution up-sampling problem."""
+
+  def dataset_filename(self):
+    return "img2img_allen_brain"  # Reuse base problem data
+
+  @property
+  def input_dim(self):
+    return 48
+
+  @property
+  def output_dim(self):
+    return 64
+
+
+@registry.register_problem
+class Img2imgAllenBrainDim8to32(Img2imgAllenBrain):
+  """8px to 32px resolution up-sampling problem."""
+
+  def dataset_filename(self):
+    return "img2img_allen_brain"  # Reuse base problem data
+
+  @property
+  def input_dim(self):
+    return 8
+
+  @property
+  def output_dim(self):
+    return 32
+
+
+@registry.register_problem
+class Img2imgAllenBrainDim16to16Paint1(Img2imgAllenBrain):
+  """In-painting problem (1%) with no resolution upsampling."""
+
+  def dataset_filename(self):
+    return "img2img_allen_brain"  # Reuse base problem data
+
+  @property
+  def input_dim(self):
+    return 16
+
+  @property
+  def output_dim(self):
+    return 16
+
+  @property
+  def inpaint_fraction(self):
+    return 0.01
+
diff --git a/tensor2tensor/data_generators/allen_brain_test.py b/tensor2tensor/data_generators/allen_brain_test.py
new file mode 100644
index 000000000..98413f1a5
--- /dev/null
+++ b/tensor2tensor/data_generators/allen_brain_test.py
@@ -0,0 +1,174 @@
+# coding=utf-8
+# 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 of the Allen Brain Atlas problems."""
+
+import tensorflow as tf
+from tensorflow.contrib.eager.python import tfe
+
+from tensor2tensor.data_generators import allen_brain
+from tensor2tensor.data_generators.allen_brain import _generator
+from tensor2tensor.data_generators.allen_brain_utils import mock_raw_data
+from tensor2tensor.data_generators.allen_brain_utils import TemporaryDirectory
+from tensor2tensor.models import image_transformer_2d
+
+tfe.enable_eager_execution()
+Modes = tf.estimator.ModeKeys
+
+
+class TestAllenBrain(tf.test.TestCase):
+  """Tests that are common to all Allen Brain Atlas problems."""
+
+  def setUp(self):
+
+    self.all_problems = [
+        #allen_brain.Img2imgAllenBrain,
+        #allen_brain.Img2imgAllenBrainDim48to64,
+        #allen_brain.Img2imgAllenBrainDim8to32,
+        #allen_brain.Img2imgAllenBrainDim16to32,
+        allen_brain.Img2imgAllenBrainDim16to16Paint1
+    ]
+
+  def test_generator_produces_examples(self):
+    """Basic test that the generator produces examples with expected keys."""
+
+    for is_training in [True, False]:
+      with TemporaryDirectory() as tmp_dir:
+        mock_raw_data(tmp_dir, raw_dim=256, num_images=100)
+        for example in _generator(tmp_dir, is_training):
+          for key in ["image/encoded", "image/format",
+                      "image/height", "image/width"]:
+            self.assertTrue(key in example.keys())
+
+  def test_generate_data_produces_examples_of_correct_shape(self):
+    """Test examples have correct input and output shapes.
+
+    Notes:
+
+      * Loops over all AllenBrainImage2image* problems.
+
+    """
+
+    with TemporaryDirectory() as tmp_dir:
+      mock_raw_data(tmp_dir, raw_dim=256, num_images=100)
+      with TemporaryDirectory() as data_dir:
+        for problem_obj in self.all_problems:
+          problem_object = problem_obj()
+
+          problem_object.generate_data(data_dir, tmp_dir)
+
+          for mode in [Modes.TRAIN, Modes.EVAL]:
+
+            dataset = problem_object.dataset(mode, data_dir)
+            example = tfe.Iterator(dataset).next()
+
+            num_channels = problem_object.num_channels
+
+            # Check that the input tensor has the right shape
+            input_dim = problem_object.input_dim
+            self.assertEqual(example["inputs"].numpy().shape,
+                             (input_dim, input_dim, num_channels))
+
+            # Check that the targets tensor has the right shape
+            output_dim = problem_object.output_dim
+            self.assertEqual(example["targets"].numpy().shape,
+                             (output_dim, output_dim, num_channels))
+
+  def test_transformer2d_single_step_e2e(self):
+    """Minimal end-to-end test of training and eval on allen_brain_image2image.
+
+    Notes:
+
+      * Runs problem generate_data
+
+      * Runs a single step of training
+
+      * Runs model in eval mode to obtain a prediction and confirms the
+        resulting shape.
+
+        * TODO: Running this in predict mode crashes in my environment.
+          Separately have seen predict mode not produce the right shape
+          output tensors, as if .infer is still a wip.
+
+    """
+
+    problem_object = allen_brain.Img2imgAllenBrainDim8to32()
+
+    with TemporaryDirectory() as tmp_dir:
+
+      mock_raw_data(tmp_dir, raw_dim=256, num_images=100)
+
+      with TemporaryDirectory() as data_dir:
+
+        problem_object.generate_data(data_dir, tmp_dir)
+
+        input_xy_dim = problem_object.input_dim
+        target_xy_dim = problem_object.output_dim
+        num_channels = problem_object.num_channels
+
+        hparams = image_transformer_2d.img2img_transformer2d_tiny()
+        hparams.data_dir = data_dir
+
+        p_hparams = problem_object.get_hparams(hparams)
+
+        model = image_transformer_2d.Img2imgTransformer(
+            hparams, tf.estimator.ModeKeys.TRAIN, p_hparams
+        )
+
+        @tfe.implicit_value_and_gradients
+        def loss_fn(features):
+          _, losses = model(features)
+          return losses["training"]
+
+        batch_size = 1
+        train_dataset = problem_object.dataset(Modes.TRAIN, data_dir)
+        train_dataset = train_dataset.repeat(None).batch(batch_size)
+
+        optimizer = tf.train.AdamOptimizer()
+
+        example = tfe.Iterator(train_dataset).next()
+        example["targets"] = tf.reshape(example["targets"],
+                                        [batch_size,
+                                         target_xy_dim,
+                                         target_xy_dim,
+                                         num_channels])
+        _, gv = loss_fn(example)
+        optimizer.apply_gradients(gv)
+
+        model.set_mode(Modes.EVAL)
+        dataset = problem_object.dataset(Modes.EVAL, data_dir)
+
+        example = tfe.Iterator(dataset).next()
+        example["inputs"] = tf.reshape(example["inputs"],
+                                       [1,
+                                        input_xy_dim,
+                                        input_xy_dim,
+                                        num_channels])
+        example["targets"] = tf.reshape(example["targets"],
+                                        [1,
+                                         target_xy_dim,
+                                         target_xy_dim,
+                                         num_channels])
+
+        predictions, _ = model(example)
+
+        self.assertEqual(predictions.numpy().shape,
+                         (1,
+                          target_xy_dim,
+                          target_xy_dim,
+                          num_channels,
+                          256))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/data_generators/allen_brain_utils.py b/tensor2tensor/data_generators/allen_brain_utils.py
new file mode 100644
index 000000000..e85ffc621
--- /dev/null
+++ b/tensor2tensor/data_generators/allen_brain_utils.py
@@ -0,0 +1,108 @@
+# coding=utf-8
+# 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.
+
+"""Utils. for Allen Brain Atlas dataset, download and subimages."""
+
+import os
+import shutil
+import tempfile
+
+import numpy as np
+import tensorflow as tf
+
+
+def try_importing_pil_image():
+  """Import a PIL Image object if the function is called."""
+  try:
+    from PIL import Image
+  except ImportError:
+    tf.logging.error("Can't import Image from PIL (Pillow). Please install it, "
+                     "such as by running `pip install Pillow`.")
+    exit(1)
+
+  return Image
+
+
+def mock_raw_image(x_dim=1024, y_dim=1024, num_channels=3,
+                   output_path=None, write_image=True):
+  """Generate random `x_dim` by `y_dim`, optionally to `output_path`.
+
+  Args:
+    output_path: str, path to which to write image.
+    x_dim: int, the x dimension of generated raw image.
+    y_dim: int, the x dimension of generated raw image.
+    return_raw_image: bool, whether to return the generated image (as a
+      numpy array).
+
+  Returns:
+    numpy.array: The random `x_dim` by `y_dim` image (i.e. array).
+
+  """
+
+  rand_shape = (x_dim, y_dim, num_channels)
+  tf.logging.debug(rand_shape)
+
+  if num_channels != 3:
+    raise NotImplementedError("mock_raw_image for channels != 3 not yet "
+                              "implemented.")
+
+  img = np.random.random(rand_shape)
+  img = np.uint8(img*255)
+
+  if write_image:
+    if not isinstance(output_path, str):
+      raise ValueError("Output path must be of type str if write_image=True, "
+                       "saw %s." % output_path)
+
+    image_obj = try_importing_pil_image()
+    pil_img = image_obj.fromarray(img, mode="RGB")
+    with tf.gfile.Open(output_path, "w") as f:
+      pil_img.save(f, "jpeg")
+
+  return img
+
+
+def mock_raw_data(tmp_dir, raw_dim=1024, num_channels=3, num_images=1):
+  """Mock a raw data download directory with meta and raw subdirs.
+
+  Notes:
+
+    * This utility is shared by tests in both allen_brain_utils and
+      allen_brain so kept here instead of in one of *_test.
+
+  Args:
+    tmp_dir: str, temporary dir in which to mock data.
+    raw_dim: int, the x and y dimension of generated raw imgs.
+
+  """
+
+  tf.gfile.MakeDirs(tmp_dir)
+
+  for image_id in range(0, num_images):
+
+    raw_image_path = os.path.join(tmp_dir, "%s.jpg" % image_id)
+
+    mock_raw_image(x_dim=raw_dim, y_dim=raw_dim,
+                   num_channels=num_channels,
+                   output_path=raw_image_path)
+
+
+class TemporaryDirectory(object):
+  """For py2 support of `with tempfile.TemporaryDirectory() as name:`"""
+
+  def __enter__(self):
+    self.name = tempfile.mkdtemp()
+    return self.name
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    shutil.rmtree(self.name)
diff --git a/tensor2tensor/data_generators/allen_brain_utils_test.py b/tensor2tensor/data_generators/allen_brain_utils_test.py
new file mode 100644
index 000000000..441235f7c
--- /dev/null
+++ b/tensor2tensor/data_generators/allen_brain_utils_test.py
@@ -0,0 +1,83 @@
+# coding=utf-8
+# 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 of utilities supporting Allen Brain Atlas problems."""
+
+import os
+
+import tensorflow as tf
+
+from tensor2tensor.data_generators.allen_brain_utils import mock_raw_data
+from tensor2tensor.data_generators.allen_brain_utils import mock_raw_image
+from tensor2tensor.data_generators.allen_brain_utils import TemporaryDirectory
+
+
+class TestTemporaryDirectory(tf.test.TestCase):
+  """Tests of py2/py3 tmpdir context pattern compatibility class."""
+
+  def test_makes_tmpdir(self):
+    """Test that a tmpdir is created."""
+    with TemporaryDirectory() as tmp_dir:
+
+      # Within the temporary context the tmpdir has been created
+      self.assertTrue(tf.gfile.Exists(tmp_dir))
+
+    # The tmpdir no longer exists outside of the temporary context
+    self.assertFalse(tf.gfile.Exists(tmp_dir))
+
+
+class TestImageMock(tf.test.TestCase):
+  """Tests of image mocking utility."""
+
+  def test_image_mock_produces_expected_shape(self):
+    """Test that the image mocking utility produces expected shape output."""
+
+    with TemporaryDirectory() as tmp_dir:
+
+      cases = [
+          {
+              "x_dim": 8,
+              "y_dim": 8,
+              "num_channels": 3,
+              "output_path": "/foo",
+              "write_image": True
+          }
+      ]
+
+      for cid, case in enumerate(cases):
+        output_path = os.path.join(tmp_dir, "dummy%s.jpg" % cid)
+        img = mock_raw_image(x_dim=case["x_dim"],
+                             y_dim=case["y_dim"],
+                             num_channels=case["num_channels"],
+                             output_path=output_path,
+                             write_image=case["write_image"])
+
+        self.assertEqual(img.shape, (case["x_dim"], case["y_dim"],
+                                     case["num_channels"]))
+        if case["write_image"]:
+          self.assertTrue(tf.gfile.Exists(output_path))
+
+
+class TestMockRawData(tf.test.TestCase):
+  """Tests of raw data mocking utility."""
+
+  def test_runs(self):
+    """Test that data mocking utility runs for cases expected to succeed."""
+
+    with TemporaryDirectory() as tmp_dir:
+
+      mock_raw_data(tmp_dir, raw_dim=256, num_channels=3, num_images=40)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From ac8f6e3bea7b9a7d070a961ee9c8d2975d42a924 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 3 Aug 2018 10:52:36 -0700
Subject: [PATCH 0464/2720] internal merge of PR #872

PiperOrigin-RevId: 207290554
---
 .travis.yml                                   |   3 -
 setup.py                                      |   2 +-
 tensor2tensor/data_generators/allen_brain.py  |  66 +++++----
 .../data_generators/allen_brain_test.py       | 139 ++++++++++++++++--
 .../data_generators/allen_brain_utils.py      | 108 --------------
 .../data_generators/allen_brain_utils_test.py |  83 -----------
 6 files changed, 164 insertions(+), 237 deletions(-)
 delete mode 100644 tensor2tensor/data_generators/allen_brain_utils.py
 delete mode 100644 tensor2tensor/data_generators/allen_brain_utils_test.py

diff --git a/.travis.yml b/.travis.yml
index e4c209350..b162ae9f9 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -60,7 +60,6 @@ script:
   #   * visualization_test
   #   * model_rl_experiment_test
   #   * allen_brain_test
-  #   * allen_brain_utils_test
   #   * model_rl_experiment_stochastic_test
   #   * models/research
   # algorithmic_math_test: flaky
@@ -74,14 +73,12 @@ script:
     --ignore=tensor2tensor/models/research/universal_transformer_test.py
     --ignore=tensor2tensor/rl/model_rl_experiment_test.py
     --ignore=tensor2tensor/data_generators/allen_brain_test.py
-    --ignore=tensor2tensor/data_generators/allen_brain_utils_test.py
     --ignore=tensor2tensor/rl/model_rl_experiment_stochastic_test.py
     --ignore=tensor2tensor/models/research
   - pytest tensor2tensor/utils/registry_test.py
   - pytest tensor2tensor/utils/trainer_lib_test.py
   - pytest tensor2tensor/visualization/visualization_test.py
   - pytest tensor2tensor/data_generators/allen_brain_test.py
-  - pytest tensor2tensor/data_generators/allen_brain_utils_test.py
   - if [[ "$TF_VERSION" == "$TF_LATEST"  ]] || [[ "$TF_VERSION" == "tf-nightly"  ]];
     then
       pytest tensor2tensor/models/research;
diff --git a/setup.py b/setup.py
index e9740ad00..ff69570f3 100644
--- a/setup.py
+++ b/setup.py
@@ -58,7 +58,7 @@
             # explicit pip install gym[atari] for the tests.
             # 'gym[atari]',
         ],
-        'allen': ['Pillow==5.1.0', 'pandas==0.23.0']
+        'allen': ['Pillow==5.1.0', 'pandas==0.23.0'],
     },
     classifiers=[
         'Development Status :: 4 - Beta',
diff --git a/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index 59722d3fe..6ee1f5f8d 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -1,16 +1,17 @@
 # coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
 # 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
+#     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.
-
 """Problem definitions for Allen Brain Atlas problems.
 
 Notes:
@@ -28,18 +29,17 @@
 
 from io import BytesIO
 import math
-import numpy as np
 import os
+
+import numpy as np
 import requests
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
-from tensor2tensor.utils import registry
 from tensor2tensor.utils import metrics
-
-from tensor2tensor.data_generators.allen_brain_utils import try_importing_pil_image
+from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -52,23 +52,28 @@
 # the steps described here: http://help.brain-map.org/display/api,
 # e.g. https://gist.github.com/cwbeitel/5dffe90eb561637e35cdf6aa4ee3e704
 _IMAGE_IDS = [
-    '74887117', '71894997', '69443979', '79853548', '101371232', '77857182',
-    '70446772', '68994990', '69141561', '70942310', '70942316', '68298378',
-    '69690156', '74364867', '77874134', '75925043', '73854431', '69206601',
-    '71771457', '101311379', '74777533', '70960269', '71604493', '102216720',
-    '74776437', '75488723', '79815814', '77857132', '77857138', '74952778',
-    '69068486', '648167', '75703410', '74486118', '77857098', '637407',
-    '67849516', '69785503', '71547630', '69068504', '69184074', '74853078',
-    '74890694', '74890698', '75488687', '71138602', '71652378', '68079764',
-    '70619061', '68280153', '73527042', '69764608', '68399025', '244297',
-    '69902658', '68234159', '71495521', '74488395', '73923026', '68280155',
-    '75488747', '69589140', '71342189', '75119214', '79455452', '71774294',
-    '74364957', '68031779', '71389422', '67937572', '69912671', '73854471',
-    '75008183', '101371376', '75703290', '69533924', '79853544', '77343882',
-    '74887133', '332587', '69758622', '69618413', '77929999', '244293',
-    '334792', '75825136', '75008103', '70196678', '71883965', '74486130',
-    '74693566', '76107119', '76043858', '70252433', '68928364', '74806345',
-    '67848661', '75900326', '71773690', '75008171']
+    "74887117", "71894997", "69443979", "79853548", "101371232", "77857182",
+    "70446772", "68994990", "69141561", "70942310", "70942316", "68298378",
+    "69690156", "74364867", "77874134", "75925043", "73854431", "69206601",
+    "71771457", "101311379", "74777533", "70960269", "71604493", "102216720",
+    "74776437", "75488723", "79815814", "77857132", "77857138", "74952778",
+    "69068486", "648167", "75703410", "74486118", "77857098", "637407",
+    "67849516", "69785503", "71547630", "69068504", "69184074", "74853078",
+    "74890694", "74890698", "75488687", "71138602", "71652378", "68079764",
+    "70619061", "68280153", "73527042", "69764608", "68399025", "244297",
+    "69902658", "68234159", "71495521", "74488395", "73923026", "68280155",
+    "75488747", "69589140", "71342189", "75119214", "79455452", "71774294",
+    "74364957", "68031779", "71389422", "67937572", "69912671", "73854471",
+    "75008183", "101371376", "75703290", "69533924", "79853544", "77343882",
+    "74887133", "332587", "69758622", "69618413", "77929999", "244293",
+    "334792", "75825136", "75008103", "70196678", "71883965", "74486130",
+    "74693566", "76107119", "76043858", "70252433", "68928364", "74806345",
+    "67848661", "75900326", "71773690", "75008171"]
+
+
+def PIL_Image():  # pylint: disable=invalid-name
+  from PIL import Image  # pylint: disable=g-import-not-at-top
+  return Image
 
 
 def _get_case_file_paths(tmp_dir, case, training_fraction=0.95):
@@ -77,7 +82,6 @@ def _get_case_file_paths(tmp_dir, case, training_fraction=0.95):
   Args:
     tmp_dir: str, the root path to which raw images were written, at the
       top level having meta/ and raw/ subdirs.
-    size: int, the size of sub-images to consider (`size`x`size`).
     case: bool, whether obtaining file paths for training (true) or eval
       (false).
     training_fraction: float, the fraction of the sub-image path list to
@@ -85,6 +89,10 @@ def _get_case_file_paths(tmp_dir, case, training_fraction=0.95):
 
   Returns:
     list: A list of file paths.
+
+  Raises:
+    ValueError: if images not found in tmp_dir, or if training_fraction would
+      leave no examples for eval.
   """
 
   paths = tf.gfile.Glob("%s/*.jpg" % tmp_dir)
@@ -146,7 +154,7 @@ def maybe_download_image_dataset(image_ids, target_dir):
 
     response.raise_for_status()
 
-    with open(tmp_destination, "w") as f:
+    with tf.gfile.Open(tmp_destination, "w") as f:
       for block in response.iter_content(1024):
         f.write(block)
 
@@ -159,7 +167,6 @@ def random_square_mask(shape, fraction):
   Args:
     shape: tuple, shape of the mask to create.
     fraction: float, fraction of the mask area to populate with `mask_scalar`.
-    mask_scalar: float, the scalar to apply to the otherwise 1-valued mask.
 
   Returns:
     numpy.array: A numpy array storing the mask.
@@ -191,6 +198,8 @@ def _generator(tmp_dir, training, size=_BASE_EXAMPLE_IMAGE_SIZE,
       alternatively, evaluation), determining whether examples in tmp_dir
       prefixed with train or dev will be used.
     size: int, the image size to add to the example annotation.
+    training_fraction: float, the fraction of the sub-image path list to
+      consider as the basis for training examples.
 
   Yields:
     A dictionary representing the images with the following fields:
@@ -207,7 +216,7 @@ def _generator(tmp_dir, training, size=_BASE_EXAMPLE_IMAGE_SIZE,
                                      case=training,
                                      training_fraction=training_fraction)
 
-  image_obj = try_importing_pil_image()
+  image_obj = PIL_Image()
 
   tf.logging.info("Loaded case file paths (n=%s)" % len(image_files))
   height = size
@@ -230,8 +239,7 @@ def _generator(tmp_dir, training, size=_BASE_EXAMPLE_IMAGE_SIZE,
         v_end = v_offset + size - 1
 
         # Extract a sub-image tile.
-        # pylint: disable=invalid-sequence-index
-        subimage = np.uint8(img[h_offset:h_end, v_offset:v_end])
+        subimage = np.uint8(img[h_offset:h_end, v_offset:v_end])  # pylint: disable=invalid-sequence-index
 
         # Filter images that are likely background (not tissue).
         if np.amax(subimage) < 230:
diff --git a/tensor2tensor/data_generators/allen_brain_test.py b/tensor2tensor/data_generators/allen_brain_test.py
index 98413f1a5..739220aaa 100644
--- a/tensor2tensor/data_generators/allen_brain_test.py
+++ b/tensor2tensor/data_generators/allen_brain_test.py
@@ -1,29 +1,103 @@
 # coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
 # 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
+#     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 of the Allen Brain Atlas problems."""
 
-import tensorflow as tf
-from tensorflow.contrib.eager.python import tfe
+import os
+import shutil
+import tempfile
+
+import numpy as np
 
 from tensor2tensor.data_generators import allen_brain
-from tensor2tensor.data_generators.allen_brain import _generator
-from tensor2tensor.data_generators.allen_brain_utils import mock_raw_data
-from tensor2tensor.data_generators.allen_brain_utils import TemporaryDirectory
 from tensor2tensor.models import image_transformer_2d
 
+import tensorflow as tf
+
+tfe = tf.contrib.eager
 tfe.enable_eager_execution()
-Modes = tf.estimator.ModeKeys
+Modes = tf.estimator.ModeKeys  # pylint: disable=invalid-name
+
+
+def mock_raw_image(x_dim=1024, y_dim=1024, num_channels=3,
+                   output_path=None, write_image=True):
+  """Generate random `x_dim` by `y_dim`, optionally to `output_path`.
+
+  Args:
+    x_dim: int, the x dimension of generated raw image.
+    y_dim: int, the x dimension of generated raw image.
+    num_channels: int, number of channels in image.
+    output_path: str, path to which to write image.
+    write_image: bool, whether to write the image to output_path.
+
+  Returns:
+    numpy.array: The random `x_dim` by `y_dim` image (i.e. array).
+  """
+
+  rand_shape = (x_dim, y_dim, num_channels)
+
+  if num_channels != 3:
+    raise NotImplementedError("mock_raw_image for channels != 3 not yet "
+                              "implemented.")
+
+  img = np.random.random(rand_shape)
+  img = np.uint8(img*255)
+
+  if write_image:
+    image_obj = allen_brain.PIL_Image()
+    pil_img = image_obj.fromarray(img, mode="RGB")
+    with tf.gfile.Open(output_path, "w") as f:
+      pil_img.save(f, "jpeg")
+
+  return img
+
+
+def mock_raw_data(tmp_dir, raw_dim=1024, num_channels=3, num_images=1):
+  """Mock a raw data download directory with meta and raw subdirs.
+
+  Notes:
+
+    * This utility is shared by tests in both allen_brain_utils and
+      allen_brain so kept here instead of in one of *_test.
+
+  Args:
+    tmp_dir: str, temporary dir in which to mock data.
+    raw_dim: int, the x and y dimension of generated raw imgs.
+    num_channels: int, number of channels in image.
+    num_images: int, number of images to mock.
+  """
+
+  tf.gfile.MakeDirs(tmp_dir)
+
+  for image_id in range(num_images):
+
+    raw_image_path = os.path.join(tmp_dir, "%s.jpg" % image_id)
+
+    mock_raw_image(x_dim=raw_dim, y_dim=raw_dim,
+                   num_channels=num_channels,
+                   output_path=raw_image_path)
+
+
+class TemporaryDirectory(object):
+  """For py2 support of `with tempfile.TemporaryDirectory() as name:`"""
+
+  def __enter__(self):
+    self.name = tempfile.mkdtemp()
+    return self.name
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    shutil.rmtree(self.name)
 
 
 class TestAllenBrain(tf.test.TestCase):
@@ -32,10 +106,6 @@ class TestAllenBrain(tf.test.TestCase):
   def setUp(self):
 
     self.all_problems = [
-        #allen_brain.Img2imgAllenBrain,
-        #allen_brain.Img2imgAllenBrainDim48to64,
-        #allen_brain.Img2imgAllenBrainDim8to32,
-        #allen_brain.Img2imgAllenBrainDim16to32,
         allen_brain.Img2imgAllenBrainDim16to16Paint1
     ]
 
@@ -45,7 +115,7 @@ def test_generator_produces_examples(self):
     for is_training in [True, False]:
       with TemporaryDirectory() as tmp_dir:
         mock_raw_data(tmp_dir, raw_dim=256, num_images=100)
-        for example in _generator(tmp_dir, is_training):
+        for example in allen_brain._generator(tmp_dir, is_training):
           for key in ["image/encoded", "image/format",
                       "image/height", "image/width"]:
             self.assertTrue(key in example.keys())
@@ -170,5 +240,48 @@ def loss_fn(features):
                           256))
 
 
+class TestImageMock(tf.test.TestCase):
+  """Tests of image mocking utility."""
+
+  def test_image_mock_produces_expected_shape(self):
+    """Test that the image mocking utility produces expected shape output."""
+
+    with TemporaryDirectory() as tmp_dir:
+
+      cases = [
+          {
+              "x_dim": 8,
+              "y_dim": 8,
+              "num_channels": 3,
+              "output_path": "/foo",
+              "write_image": True
+          }
+      ]
+
+      for cid, case in enumerate(cases):
+        output_path = os.path.join(tmp_dir, "dummy%s.jpg" % cid)
+        img = mock_raw_image(x_dim=case["x_dim"],
+                             y_dim=case["y_dim"],
+                             num_channels=case["num_channels"],
+                             output_path=output_path,
+                             write_image=case["write_image"])
+
+        self.assertEqual(img.shape, (case["x_dim"], case["y_dim"],
+                                     case["num_channels"]))
+        if case["write_image"]:
+          self.assertTrue(tf.gfile.Exists(output_path))
+
+
+class TestMockRawData(tf.test.TestCase):
+  """Tests of raw data mocking utility."""
+
+  def test_runs(self):
+    """Test that data mocking utility runs for cases expected to succeed."""
+
+    with TemporaryDirectory() as tmp_dir:
+
+      mock_raw_data(tmp_dir, raw_dim=256, num_channels=3, num_images=40)
+
+
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/data_generators/allen_brain_utils.py b/tensor2tensor/data_generators/allen_brain_utils.py
deleted file mode 100644
index e85ffc621..000000000
--- a/tensor2tensor/data_generators/allen_brain_utils.py
+++ /dev/null
@@ -1,108 +0,0 @@
-# coding=utf-8
-# 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.
-
-"""Utils. for Allen Brain Atlas dataset, download and subimages."""
-
-import os
-import shutil
-import tempfile
-
-import numpy as np
-import tensorflow as tf
-
-
-def try_importing_pil_image():
-  """Import a PIL Image object if the function is called."""
-  try:
-    from PIL import Image
-  except ImportError:
-    tf.logging.error("Can't import Image from PIL (Pillow). Please install it, "
-                     "such as by running `pip install Pillow`.")
-    exit(1)
-
-  return Image
-
-
-def mock_raw_image(x_dim=1024, y_dim=1024, num_channels=3,
-                   output_path=None, write_image=True):
-  """Generate random `x_dim` by `y_dim`, optionally to `output_path`.
-
-  Args:
-    output_path: str, path to which to write image.
-    x_dim: int, the x dimension of generated raw image.
-    y_dim: int, the x dimension of generated raw image.
-    return_raw_image: bool, whether to return the generated image (as a
-      numpy array).
-
-  Returns:
-    numpy.array: The random `x_dim` by `y_dim` image (i.e. array).
-
-  """
-
-  rand_shape = (x_dim, y_dim, num_channels)
-  tf.logging.debug(rand_shape)
-
-  if num_channels != 3:
-    raise NotImplementedError("mock_raw_image for channels != 3 not yet "
-                              "implemented.")
-
-  img = np.random.random(rand_shape)
-  img = np.uint8(img*255)
-
-  if write_image:
-    if not isinstance(output_path, str):
-      raise ValueError("Output path must be of type str if write_image=True, "
-                       "saw %s." % output_path)
-
-    image_obj = try_importing_pil_image()
-    pil_img = image_obj.fromarray(img, mode="RGB")
-    with tf.gfile.Open(output_path, "w") as f:
-      pil_img.save(f, "jpeg")
-
-  return img
-
-
-def mock_raw_data(tmp_dir, raw_dim=1024, num_channels=3, num_images=1):
-  """Mock a raw data download directory with meta and raw subdirs.
-
-  Notes:
-
-    * This utility is shared by tests in both allen_brain_utils and
-      allen_brain so kept here instead of in one of *_test.
-
-  Args:
-    tmp_dir: str, temporary dir in which to mock data.
-    raw_dim: int, the x and y dimension of generated raw imgs.
-
-  """
-
-  tf.gfile.MakeDirs(tmp_dir)
-
-  for image_id in range(0, num_images):
-
-    raw_image_path = os.path.join(tmp_dir, "%s.jpg" % image_id)
-
-    mock_raw_image(x_dim=raw_dim, y_dim=raw_dim,
-                   num_channels=num_channels,
-                   output_path=raw_image_path)
-
-
-class TemporaryDirectory(object):
-  """For py2 support of `with tempfile.TemporaryDirectory() as name:`"""
-
-  def __enter__(self):
-    self.name = tempfile.mkdtemp()
-    return self.name
-
-  def __exit__(self, exc_type, exc_value, traceback):
-    shutil.rmtree(self.name)
diff --git a/tensor2tensor/data_generators/allen_brain_utils_test.py b/tensor2tensor/data_generators/allen_brain_utils_test.py
deleted file mode 100644
index 441235f7c..000000000
--- a/tensor2tensor/data_generators/allen_brain_utils_test.py
+++ /dev/null
@@ -1,83 +0,0 @@
-# coding=utf-8
-# 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 of utilities supporting Allen Brain Atlas problems."""
-
-import os
-
-import tensorflow as tf
-
-from tensor2tensor.data_generators.allen_brain_utils import mock_raw_data
-from tensor2tensor.data_generators.allen_brain_utils import mock_raw_image
-from tensor2tensor.data_generators.allen_brain_utils import TemporaryDirectory
-
-
-class TestTemporaryDirectory(tf.test.TestCase):
-  """Tests of py2/py3 tmpdir context pattern compatibility class."""
-
-  def test_makes_tmpdir(self):
-    """Test that a tmpdir is created."""
-    with TemporaryDirectory() as tmp_dir:
-
-      # Within the temporary context the tmpdir has been created
-      self.assertTrue(tf.gfile.Exists(tmp_dir))
-
-    # The tmpdir no longer exists outside of the temporary context
-    self.assertFalse(tf.gfile.Exists(tmp_dir))
-
-
-class TestImageMock(tf.test.TestCase):
-  """Tests of image mocking utility."""
-
-  def test_image_mock_produces_expected_shape(self):
-    """Test that the image mocking utility produces expected shape output."""
-
-    with TemporaryDirectory() as tmp_dir:
-
-      cases = [
-          {
-              "x_dim": 8,
-              "y_dim": 8,
-              "num_channels": 3,
-              "output_path": "/foo",
-              "write_image": True
-          }
-      ]
-
-      for cid, case in enumerate(cases):
-        output_path = os.path.join(tmp_dir, "dummy%s.jpg" % cid)
-        img = mock_raw_image(x_dim=case["x_dim"],
-                             y_dim=case["y_dim"],
-                             num_channels=case["num_channels"],
-                             output_path=output_path,
-                             write_image=case["write_image"])
-
-        self.assertEqual(img.shape, (case["x_dim"], case["y_dim"],
-                                     case["num_channels"]))
-        if case["write_image"]:
-          self.assertTrue(tf.gfile.Exists(output_path))
-
-
-class TestMockRawData(tf.test.TestCase):
-  """Tests of raw data mocking utility."""
-
-  def test_runs(self):
-    """Test that data mocking utility runs for cases expected to succeed."""
-
-    with TemporaryDirectory() as tmp_dir:
-
-      mock_raw_data(tmp_dir, raw_dim=256, num_channels=3, num_images=40)
-
-
-if __name__ == "__main__":
-  tf.test.main()

From 927f78c820f52fc8a5c3bfbfefd5e739c4e2bad7 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 3 Aug 2018 11:17:42 -0700
Subject: [PATCH 0465/2720] Do not require preprocess_resize_frames in hparams
 for VideoProblem

PiperOrigin-RevId: 207295010
---
 tensor2tensor/data_generators/video_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 46bc80acf..f8e7ef26e 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -140,7 +140,7 @@ def use_not_breaking_batching(self):
 
   def preprocess_example(self, example, mode, hparams):
     """Runtime preprocessing, e.g., resize example["frame"]."""
-    if hparams.preprocess_resize_frames is not None:
+    if getattr(hparams, "preprocess_resize_frames", None) is not None:
       example["frame"] = tf.image.resize_images(
           example["frame"],
           hparams.preprocess_resize_frames,

From 6edbd6b7d14ae81e359b328a47d6c6187d0d5844 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 3 Aug 2018 11:29:38 -0700
Subject: [PATCH 0466/2720] Remove flaky testNoShuffleFail test that was
 testing non-guaranteed behavior

PiperOrigin-RevId: 207297174
---
 tensor2tensor/data_generators/problem_test.py | 12 ------------
 1 file changed, 12 deletions(-)

diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 5e76db9e8..26e370bd8 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -73,18 +73,6 @@ def testNoShufflePreprocess(self):
     with tf.Session() as sess:
       self.assertTrue(assert_tensors_equal(sess, tensor1, tensor2, 20))
 
-  def testNoShuffleFail(self):
-    problem = algorithmic.TinyAlgo()
-    dataset = problem.dataset(mode=tf.estimator.ModeKeys.TRAIN,
-                              data_dir=algorithmic.TinyAlgo.data_dir,
-                              shuffle_files=True)
-
-    tensor1 = dataset.make_one_shot_iterator().get_next()["targets"]
-    tensor2 = dataset.make_one_shot_iterator().get_next()["targets"]
-
-    with tf.Session() as sess:
-      self.assertFalse(assert_tensors_equal(sess, tensor1, tensor2, 20))
-
 
 if __name__ == "__main__":
   tf.test.main()

From a2a08662f8cdbac21424e9a0af6218d376a00141 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 3 Aug 2018 13:32:12 -0700
Subject: [PATCH 0467/2720] Stochastic Adversarial Video Prediction (Part 2 of
 3). Implements an adversarial loss component via hparam="use_gan=True".
 Latent states z_{t} sampled from ~ N(0.0, 1.0) are encoded using a LSTM and
 concatenated along with X_{t-1} to predict the frames at X_{t}. The predicted
 frames are then stacked and discriminated from the real video using a 3-D
 discriminator.

PiperOrigin-RevId: 207315716
---
 tensor2tensor/layers/common_layers.py         |  49 ++++-
 tensor2tensor/layers/common_layers_test.py    |  19 ++
 .../models/research/next_frame_params.py      |   6 +
 .../models/research/next_frame_savp.py        | 202 ++++++++++++++++++
 .../models/research/next_frame_sv2p.py        |  27 ++-
 .../models/research/next_frame_test.py        |  14 ++
 tensor2tensor/utils/update_ops_hook.py        |  29 +++
 7 files changed, 333 insertions(+), 13 deletions(-)
 create mode 100644 tensor2tensor/utils/update_ops_hook.py

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index c601f57d2..851a0222d 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -711,6 +711,46 @@ def l2_norm(x, filters=None, epsilon=1e-6, name=None, reuse=None):
     return norm_x * scale + bias
 
 
+def apply_spectral_norm(x):
+  """Normalizes x using the spectral norm.
+
+  The implementation follows Algorithm 1 of
+  https://arxiv.org/abs/1802.05957. If x is not a 2-D Tensor, then it is
+  reshaped such that the number of channels (last-dimension) is the same.
+
+  Args:
+    x: Tensor with the last dimension equal to the number of filters.
+
+  Returns:
+    x: Tensor with the same shape as x normalized by the spectral norm.
+    assign_op: Op to be run after every step to update the vector "u".
+  """
+  weights_shape = shape_list(x)
+  other, num_filters = tf.reduce_prod(weights_shape[:-1]), weights_shape[-1]
+
+  # Reshape into a 2-D matrix with outer size num_filters.
+  weights_2d = tf.reshape(x, (other, num_filters))
+
+  # v = Wu / ||W u||
+  with tf.variable_scope("u", reuse=tf.AUTO_REUSE):
+    u = tf.get_variable(
+        "u", [num_filters, 1], initializer=tf.truncated_normal_initializer(),
+        trainable=False)
+  v = tf.nn.l2_normalize(tf.matmul(weights_2d, u))
+
+  # u_new = vW / ||v W||
+  u_new = tf.nn.l2_normalize(tf.matmul(tf.transpose(v), weights_2d))
+
+  # s = v*W*u
+  spectral_norm = tf.squeeze(tf.matmul(
+      tf.transpose(v),
+      tf.matmul(weights_2d, tf.transpose(u_new))))
+
+  # set u equal to u_new in the next iteration.
+  assign_op = tf.assign(u, tf.transpose(u_new))
+  return tf.divide(x, spectral_norm), assign_op
+
+
 def apply_norm(x, norm_type, depth, epsilon):
   """Apply Normalization."""
   if norm_type == "layer":
@@ -3505,7 +3545,8 @@ def sliced_gan_loss(input1,
                     discriminator,
                     num_vecs,
                     do_random_vecs=True,
-                    do_tanh=True):
+                    do_tanh=True,
+                    return_logits=False):
   """Loss inspired by the sliced WGAN paper: https://arxiv.org/abs/1804.01947.
 
   Puts input1 and input2 through the provided discriminator to get logits.
@@ -3520,6 +3561,7 @@ def sliced_gan_loss(input1,
     num_vecs: how many random vectors to use for projections.
     do_random_vecs: whether to use random vectors or just tanh of the logits.
     do_tanh: if true (default) we'll also just use tanh of the logits.
+    return_logits: Whether or not to return the logits.
 
   Returns:
     The generator loss, i.e., the sliced approximation of the distance between
@@ -3565,7 +3607,10 @@ def get_sorted_projections(x):
 
     proj1 = get_sorted_projections(logits1)
     proj2 = get_sorted_projections(logits2)
-    return tf.reduce_mean(tf.square(proj1 - proj2))
+    dist = tf.reduce_mean(tf.square(proj1 - proj2))
+    if return_logits:
+      return dist, logits1, logits2
+    return dist
 
 
 def upscale(inputs, f, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR):
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index fd617dd10..b510e84e5 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -692,6 +692,25 @@ def testCycleGANUpsampleConv2dTranspose(self):
           [batch, upsampled_height, upsampled_width, output_filters],
           session.run(upsampled_output_shape))
 
+  def testSpectralNorm(self):
+    # Test that after 20 calls to apply_spectral_norm, the spectral
+    # norm of the normalized matrix is close to 1.0
+    with tf.Graph().as_default():
+      weights = tf.get_variable("w", dtype=tf.float32, shape=[2, 3, 50, 100])
+      weights = tf.multiply(weights, 10.0)
+      normed_weight, assign_op = common_layers.apply_spectral_norm(weights)
+
+      with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+
+        for _ in range(20):
+          sess.run(assign_op)
+          normed_weight, assign_op = common_layers.apply_spectral_norm(
+              weights)
+        normed_weight = sess.run(normed_weight).reshape(-1, 100)
+        _, s, _ = np.linalg.svd(normed_weight)
+        self.assertTrue(np.allclose(s[0], 1.0, rtol=0.1))
+
 
 class FnWithCustomGradTest(tf.test.TestCase):
 
diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index c36c30765..5b52ea944 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -104,6 +104,12 @@ def next_frame_savp():
   """SAVP model."""
   hparams = next_frame_stochastic()
   hparams.add_hparam("z_dim", 8)
+  hparams.add_hparam("num_discriminator_filters", 32)
+  hparams.add_hparam("use_vae", True)
+  hparams.add_hparam("use_gan", False)
+  hparams.add_hparam("use_spectral_norm", True)
+  hparams.add_hparam("gan_loss", "cross_entropy")
+  hparams.add_hparam("gan_loss_multiplier", 0.01)
   hparams.target_modality = "video:l1raw"
   hparams.input_modalities = "inputs:video:l1raw"
   hparams.latent_loss_multiplier_schedule = "linear_anneal"
diff --git a/tensor2tensor/models/research/next_frame_savp.py b/tensor2tensor/models/research/next_frame_savp.py
index 1e4883291..1b4678995 100644
--- a/tensor2tensor/models/research/next_frame_savp.py
+++ b/tensor2tensor/models/research/next_frame_savp.py
@@ -20,14 +20,18 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+import numbers
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import next_frame_params  # pylint: disable=unused-import
 from tensor2tensor.models.research import next_frame_sv2p
 from tensor2tensor.utils import registry
+from tensor2tensor.utils import update_ops_hook
 import tensorflow as tf
 
+gan_losses = tf.contrib.gan.losses.wargs
+
 
 @registry.register_model
 class NextFrameSAVP(next_frame_sv2p.NextFrameStochastic):
@@ -98,6 +102,195 @@ def encoder(self, inputs, n_layers=3):
         z_log_var, (batch_size, -1, latent_dims))
     return z_mu, z_log_var
 
+  def discriminator(self, frames):
+    """3-D SNGAN discriminator.
+
+    Args:
+      frames: a list of batch-major tensors indexed by time.
+
+    Returns:
+      logits: 1-D Tensor with shape=batch_size.
+              Positive logits imply that the discriminator thinks that it
+              belongs to the true class.
+    """
+    ndf = self.hparams.num_discriminator_filters
+    frames = tf.stack(frames)
+
+    # Switch from time-major axis to batch-major axis.
+    frames = common_video.swap_time_and_batch_axes(frames)
+
+    # 3-D Conv-net mapping inputs to activations.
+    num_outputs = [ndf, ndf*2, ndf*2, ndf*4, ndf*4, ndf*8, ndf*8]
+    kernel_sizes = [3, 4, 3, 4, 3, 4, 3]
+    strides = [1, [1, 2, 2], 1, [1, 2, 2], 1, 2, 1]
+    names = ["video_sn_conv0_0", "video_sn_conv0_1", "video_sn_conv1_0",
+             "video_sn_conv1_1", "video_sn_conv2_0", "video_sn_conv2_1",
+             "video_sn_conv3_0"]
+    iterable = zip(num_outputs, kernel_sizes, strides, names)
+    activations = frames
+    for num_filters, kernel_size, stride, name in iterable:
+      activations = self.pad_conv3d_lrelu(activations, num_filters, kernel_size,
+                                          stride, name)
+
+    # Flatten and apply fully-connected layer.
+    num_fc_dimensions = tf.reduce_prod(
+        common_layers.shape_list(activations)[1:])
+    activations = tf.reshape(activations, (-1, num_fc_dimensions))
+    return tf.squeeze(tf.layers.dense(activations, 1))
+
+  def d_step(self, true_frames, gen_frames):
+    """Performs the discriminator step in computing the GAN loss.
+
+    Applies stop-gradient to the generated frames while computing the
+    discriminator loss to make sure that the gradients are not back-propagated
+    to the generator. This makes sure that only the discriminator is updated.
+
+    Args:
+      true_frames: True outputs
+      gen_frames: Generated frames.
+    Returns:
+      d_loss: Loss component due to the discriminator.
+    """
+    hparam_to_disc_loss = {
+        "least_squares": gan_losses.least_squares_discriminator_loss,
+        "cross_entropy": gan_losses.modified_discriminator_loss,
+        "wasserstein": gan_losses.wasserstein_discriminator_loss}
+
+    # Concat across batch-axis.
+    _, batch_size, _, _, _ = common_layers.shape_list(true_frames)
+    all_frames = tf.concat(
+        [true_frames, tf.stop_gradient(gen_frames)], axis=1)
+
+    all_logits = self.discriminator(all_frames)
+    true_logits, fake_logits_stop = \
+      all_logits[:batch_size], all_logits[batch_size:]
+    mean_true_logits = tf.reduce_mean(true_logits)
+    tf.summary.scalar("mean_true_logits", mean_true_logits)
+
+    mean_fake_logits_stop = tf.reduce_mean(fake_logits_stop)
+    tf.summary.scalar("mean_fake_logits_stop", mean_fake_logits_stop)
+
+    discriminator_loss_func = hparam_to_disc_loss[self.hparams.gan_loss]
+    gan_d_loss = discriminator_loss_func(
+        discriminator_real_outputs=true_logits,
+        discriminator_gen_outputs=fake_logits_stop,
+        add_summaries=True)
+    return gan_d_loss, true_logits, fake_logits_stop
+
+  def g_step(self, gen_frames, fake_logits_stop):
+    """Performs the generator step in computing the GAN loss.
+
+    Args:
+      gen_frames: Generated frames
+      fake_logits_stop: Logits corresponding to the generated frames as per
+                        the discriminator. Assumed to have a stop-gradient term.
+    Returns:
+      gan_g_loss_pos_d: Loss.
+      gan_g_loss_neg_d: -gan_g_loss_pos_d but with a stop gradient on generator.
+    """
+    hparam_to_gen_loss = {
+        "least_squares": gan_losses.least_squares_generator_loss,
+        "cross_entropy": gan_losses.modified_generator_loss,
+        "wasserstein": gan_losses.wasserstein_generator_loss
+    }
+
+    fake_logits = self.discriminator(gen_frames)
+    mean_fake_logits = tf.reduce_mean(fake_logits)
+    tf.summary.scalar("mean_fake_logits", mean_fake_logits)
+
+    # Generator loss.
+    # Using gan_g_loss_pos_d updates the discriminator as well.
+    # To avoid this add gan_g_loss_neg_d = -gan_g_loss_pos_d
+    # but with stop gradient on the generator.
+    # This makes sure that the net gradient on the discriminator is zero and
+    # net-gradient on the generator is just due to the gan_g_loss_pos_d.
+    generator_loss_func = hparam_to_gen_loss[self.hparams.gan_loss]
+    gan_g_loss_pos_d = generator_loss_func(
+        discriminator_gen_outputs=fake_logits, add_summaries=True)
+    gan_g_loss_neg_d = -generator_loss_func(
+        discriminator_gen_outputs=fake_logits_stop, add_summaries=True)
+    return gan_g_loss_pos_d, gan_g_loss_neg_d
+
+  def get_gan_loss(self, true_frames, gen_frames):
+    """Get the discriminator + generator loss at every step.
+
+    This performs an 1:1 update of the discriminator and generator at every
+    step.
+
+    Args:
+      true_frames: 5-D Tensor of shape (num_steps, batch_size, H, W, C)
+                   Assumed to be ground truth.
+      gen_frames: 5-D Tensor of shape (num_steps, batch_size, H, W, C)
+                  Assumed to be fake.
+    Returns:
+      loss: 0-D Tensor, with d_loss + g_loss
+    """
+    # D - STEP
+    with tf.variable_scope("gan_discriminator", reuse=tf.AUTO_REUSE):
+      gan_d_loss, _, fake_logits_stop = self.d_step(
+          true_frames, gen_frames)
+
+    # G - STEP
+    with tf.variable_scope("gan_discriminator", reuse=True):
+      gan_g_loss_pos_d, gan_g_loss_neg_d = self.g_step(
+          gen_frames, fake_logits_stop)
+    gan_g_loss = gan_g_loss_pos_d + gan_g_loss_neg_d
+    gan_loss = gan_g_loss + gan_d_loss
+    tf.summary.scalar("gan_loss", gan_g_loss_pos_d + gan_d_loss)
+    return self.hparams.gan_loss_multiplier * gan_loss
+
+  def get_extra_loss(self, latent_means=None, latent_stds=None,
+                     true_frames=None, gen_frames=None, beta=1.0):
+    if not self.is_training:
+      return 0.0
+    if self.hparams.use_vae:
+      return super(NextFrameSAVP, self).get_extra_loss(
+          latent_means=latent_means, latent_stds=latent_stds, beta=beta)
+    elif self.hparams.use_gan:
+      # Strip out the first context_frames for the true_frames
+      # Strip out the first context_frames - 1 for the gen_frames
+      context_frames = self.hparams.video_num_input_frames
+      true_frames = tf.stack(
+          tf.unstack(true_frames, axis=0)[context_frames:])
+      gen_frames = tf.stack(
+          tf.unstack(gen_frames, axis=0)[context_frames-1:])
+      return self.get_gan_loss(true_frames, gen_frames)
+
+  def pad_conv3d_lrelu(self, activations, n_filters, kernel_size, strides,
+                       scope):
+    """Pad, apply 3-D convolution and leaky relu."""
+    padding = [[0, 0], [1, 1], [1, 1], [1, 1], [0, 0]]
+
+    # tf.nn.conv3d accepts a list of 5 values for strides
+    # with first and last value equal to 1
+    if isinstance(strides, numbers.Integral):
+      strides = [strides] * 3
+    strides = [1] + strides + [1]
+
+    # Filter_shape = [K, K, K, num_input, num_output]
+    filter_shape = (
+        [kernel_size]*3 + activations.shape[-1:].as_list() + [n_filters])
+
+    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
+      conv_filter = tf.get_variable(
+          "conv_filter", shape=filter_shape,
+          initializer=tf.truncated_normal_initializer(stddev=0.02))
+
+      if self.hparams.use_spectral_norm:
+        conv_filter, assign_op = common_layers.apply_spectral_norm(conv_filter)
+        if self.is_training:
+          tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, assign_op)
+
+      padded = tf.pad(activations, padding)
+      convolved = tf.nn.conv3d(
+          padded, conv_filter, strides=strides, padding="VALID")
+      rectified = tf.nn.leaky_relu(convolved, alpha=0.2)
+    return rectified
+
+  @staticmethod
+  def train_hooks():
+    return [update_ops_hook.UpdateOpsHook()]
+
   def construct_model(self, images, actions, rewards):
     """Model that takes in images and returns predictions.
 
@@ -116,7 +309,16 @@ def construct_model(self, images, actions, rewards):
                     every frame.
       latent_stds: list of gaussian stds conditioned on the input at
                    every frame.
+
+    Raises:
+      ValueError: If not exactly one of self.hparams.vae or self.hparams.gan
+                  is set to True.
     """
+    if self.hparams.use_vae and self.hparams.use_gan:
+      raise ValueError("VAE + GAN variant not implemented")
+    if not self.hparams.use_vae and not self.hparams.use_gan:
+      raise ValueError("Set at least one of use_vae or use_gan to be True")
+
     images = tf.unstack(images, axis=0)
     actions = tf.unstack(actions, axis=0)
     rewards = tf.unstack(rewards, axis=0)
diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 43f0a2ff6..045fd7c7a 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -541,6 +541,18 @@ def process_single_frame(prev_outputs, inputs):
 
     return gen_images, gen_rewards, [latent_mean], [latent_std]
 
+  def get_extra_loss(self, latent_means=None, latent_stds=None,
+                     true_frames=None, gen_frames=None, beta=1.0):
+    """Losses in addition to the default modality losses."""
+    kl_loss = 0.0
+    if self.is_training:
+      for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
+        kl_loss += common_layers.kl_divergence(mean, std)
+        tf.summary.histogram("posterior_mean_%d" % i, mean)
+        tf.summary.histogram("posterior_std_%d" % i, std)
+      tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
+    return beta * kl_loss
+
   def body(self, features):
     hparams = self.hparams
     batch_size = common_layers.shape_list(features["inputs"])[0]
@@ -578,17 +590,10 @@ def body(self, features):
     )
 
     beta = self.get_beta()
-    kl_loss = 0.0
-    if self.is_training:
-      for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
-        kl_loss += common_layers.kl_divergence(mean, std)
-        tf.summary.histogram("posterior_mean_%d" % i, mean)
-        tf.summary.histogram("posterior_std_%d" % i, std)
-
-      tf.summary.scalar("beta", beta)
-      tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
-
-    extra_loss = beta * kl_loss
+    extra_loss = self.get_extra_loss(
+        latent_means=latent_means,
+        latent_stds=latent_stds, beta=beta, true_frames=all_frames,
+        gen_frames=gen_images)
 
     # Ignore the predictions from the input frames.
     # This is NOT the same as original paper/implementation.
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index 26622bde9..b7bd2c13a 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -231,6 +231,20 @@ def testStochasticSavp(self):
         next_frame_savp.NextFrameSAVP,
         1)
 
+  def testStochasticSavpGAN(self):
+    hparams = next_frame_params.next_frame_savp()
+    hparams.use_gan = True
+    hparams.use_vae = False
+    self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
+
+  def testStochasticInvalidVAEGANCombinations(self):
+    hparams = next_frame_params.next_frame_savp()
+    for use_vae, use_gan in [[True, True], [False, False]]:
+      hparams.use_gan = use_gan
+      hparams.use_vae = use_vae
+      self.assertRaises(ValueError, self.TestVideoModel,
+                        7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
+
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/utils/update_ops_hook.py b/tensor2tensor/utils/update_ops_hook.py
new file mode 100644
index 000000000..184272eaf
--- /dev/null
+++ b/tensor2tensor/utils/update_ops_hook.py
@@ -0,0 +1,29 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Hook to run tf.GraphKeys.UPDATE_OPS."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+
+class UpdateOpsHook(tf.train.SessionRunHook):
+  """Hook to run assign_ops."""
+
+  def before_run(self, run_context):
+    del run_context
+    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
+    return tf.train.SessionRunArgs(update_ops)

From 20cf69bc336f0a981f268e0fef25459db462a319 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Fri, 3 Aug 2018 13:49:29 -0700
Subject: [PATCH 0468/2720] Update hparams and docs to reflect best config

PiperOrigin-RevId: 207318450
---
 README.md                                 | 10 ++--
 docs/walkthrough.md                       | 10 ++--
 tensor2tensor/models/image_transformer.py | 58 ++++++++++++++++++++++-
 3 files changed, 71 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index c96925c4d..814e58404 100644
--- a/README.md
+++ b/README.md
@@ -113,9 +113,13 @@ For (un)conditional image generation, we have a number of standard data-sets:
     32x32 or `--problem=image_imagenet64_gen_rev` for 64x64.
 
 We suggest to use the Image Transformer, i.e., `--model=imagetransformer`, or
-variational auto-encoder, i.e., `--model=transformer_ae`. For CIFAR-10, using
-`--hparams_set=imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_m` yields 2.90
-bits per dimension.
+the Image Transformer Plus, i.e., `--model=imagetransformerpp` that uses
+discretized mixture of logistics, or variational auto-encoder, i.e.,
+`--model=transformer_ae`.
+For CIFAR-10, using `--hparams_set=imagetransformer_cifar10_base` or
+`--hparams_set=imagetransformer_cifar10_base_dmol` yields 2.90 bits per
+dimension. For Imagenet-32, using
+`--hparams_set=imagetransformer_imagenet32_base` yields 3.77 bits per dimension.
 
 ### Language Modeling
 
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index c96925c4d..814e58404 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -113,9 +113,13 @@ For (un)conditional image generation, we have a number of standard data-sets:
     32x32 or `--problem=image_imagenet64_gen_rev` for 64x64.
 
 We suggest to use the Image Transformer, i.e., `--model=imagetransformer`, or
-variational auto-encoder, i.e., `--model=transformer_ae`. For CIFAR-10, using
-`--hparams_set=imagetransformerpp_base_12l_8h_big_uncond_dr03_dan_m` yields 2.90
-bits per dimension.
+the Image Transformer Plus, i.e., `--model=imagetransformerpp` that uses
+discretized mixture of logistics, or variational auto-encoder, i.e.,
+`--model=transformer_ae`.
+For CIFAR-10, using `--hparams_set=imagetransformer_cifar10_base` or
+`--hparams_set=imagetransformer_cifar10_base_dmol` yields 2.90 bits per
+dimension. For Imagenet-32, using
+`--hparams_set=imagetransformer_imagenet32_base` yields 3.77 bits per dimension.
 
 ### Language Modeling
 
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index aa6d08fb2..0c894bc27 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -252,8 +252,64 @@ def imagetransformer_base():
   return hparams
 
 
+@registry.register_hparams
+def imagetransformer_cifar10_base():
+  """Best config for 2.90 bits/dim on CIFAR10 using cross entropy."""
+  hparams = image_transformer_base()
+  hparams.batch_size = 4
+  hparams.num_heads = 4
+  hparams.num_decoder_layers = 12
+  hparams.block_length = 256
+  hparams.hidden_size = 512
+  hparams.filter_size = 2048
+  hparams.learning_rate = 0.5
+  hparams.learning_rate_warmup_steps = 4000
+  hparams.layer_preprocess_sequence = "none"
+  hparams.layer_postprocess_sequence = "dan"
+  hparams.layer_prepostprocess_dropout = 0.3
+  hparams.unconditional = True
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformer_cifar10_base_dmol():
+  """Best config for 2.90 bits/dim on CIFAR10 using DMOL."""
+  hparams = image_transformer_base()
+  hparams.likelihood = cia.DistributionType.DMOL
+  hparams.num_channels = 1
+  hparams.target_modality = "image:image_channel_bottom_identity"
+  hparams.num_heads = 8
+  hparams.batch_size = 8
+  hparams.sampling_method = "random"
+  hparams.layer_preprocess_sequence = "n"
+  hparams.layer_postprocess_sequence = "da"
+  hparams.summarize_grads = True
+  hparams.hidden_size = 256
+  hparams.filter_size = 512
+  hparams.attention_key_channels = 512
+  hparams.attention_value_channels = 512
+  hparams.num_decoder_layers = 12
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.learning_rate = 0.1
+  hparams.layer_preprocess_sequence = "none"
+  hparams.layer_postprocess_sequence = "dan"
+  hparams.pos = "emb"
+  hparams.unconditional = True
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformer_imagenet32_base():
+  """Best config for ImageNet-32 with 3.77 bits/dim using cross entropy."""
+  hparams = imagetransformer_cifar10_base()
+  hparams.batch_size = 4
+  hparams.layer_prepostprocess_dropout = 0.1
+  return hparams
+
+
 @registry.register_hparams
 def imagetransformer_base_rel():
+  """Base with relative attention."""
   hparams = imagetransformer_base()
   hparams.dec_attention_type = cia.AttentionType.RELATIVE_LOCAL_1D
   return hparams
@@ -880,7 +936,7 @@ def imagetransformer_b10l_4h_big_uncond_dr03_tpu():
   hparams.num_heads = 4   # heads are expensive on tpu
   hparams.num_decoder_layers = 10
   hparams.block_length = 128
-  hparams.hidden_size = 256
+  hparams.hidden_size = 512
   hparams.filter_size = 1024
   hparams.learning_rate = 0.2
   hparams.layer_preprocess_sequence = "none"

From 098c0a8fc24083e84ff9675cba59c2e8c60012e8 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Fri, 3 Aug 2018 15:59:37 -0700
Subject: [PATCH 0469/2720] MultiProblems with subwords.

PiperOrigin-RevId: 207338544
---
 tensor2tensor/data_generators/lm1b_mnli.py    | 15 +++
 tensor2tensor/data_generators/mrpc.py         | 10 --
 .../data_generators/multi_problem.py          | 97 ++++++++++++++-----
 tensor2tensor/data_generators/multinli.py     | 19 ++--
 tensor2tensor/data_generators/qnli.py         | 10 --
 tensor2tensor/data_generators/quora_qpairs.py | 10 --
 tensor2tensor/data_generators/rte.py          | 10 --
 .../data_generators/text_problems.py          | 15 +--
 tensor2tensor/data_generators/wnli.py         | 10 --
 9 files changed, 100 insertions(+), 96 deletions(-)

diff --git a/tensor2tensor/data_generators/lm1b_mnli.py b/tensor2tensor/data_generators/lm1b_mnli.py
index a0d8e076d..a0835723f 100644
--- a/tensor2tensor/data_generators/lm1b_mnli.py
+++ b/tensor2tensor/data_generators/lm1b_mnli.py
@@ -25,6 +25,21 @@
 from tensor2tensor.utils import registry
 
 
+@registry.register_problem
+class LanguagemodelLm1bMultiNLISubwords(multi_problem.MultiProblem):
+  """LM1b and MNLI mixed problem class for multitask learning."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelLm1bMultiNLISubwords, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(lm1b.LanguagemodelLm1b32k())
+    self.task_list.append(multinli.MultiNLISharedVocab())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
+
+
 @registry.register_problem
 class LanguagemodelLm1bMultiNLI(multi_problem.MultiProblem):
   """LM1b and MNLI mixed problem class for multitask learning."""
diff --git a/tensor2tensor/data_generators/mrpc.py b/tensor2tensor/data_generators/mrpc.py
index 083970be3..f6e2b504b 100644
--- a/tensor2tensor/data_generators/mrpc.py
+++ b/tensor2tensor/data_generators/mrpc.py
@@ -71,16 +71,6 @@ def vocab_filename(self):
   def num_classes(self):
     return 2
 
-  @property
-  def concat_token(self):
-    return "<SENT_SEP>"
-
-  @property
-  def concat_id(self):
-    if self.vocab_type == text_problems.VocabType.CHARACTER:
-      return problem.TaskID.EN_SENT_PAIR
-    return 2
-
   def class_labels(self, data_dir):
     del data_dir
     return ["not_paraphrase", "paraphrase"]
diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 0fb37679c..b22dbc8d2 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -23,6 +23,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.layers import discretization
 from tensor2tensor.utils import metrics
+from tensor2tensor.utils import registry
 import tensorflow as tf
 
 
@@ -37,16 +38,20 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
     assert len(self.task_list) > 1
 
     for task in self.task_list:
-      assert task.vocab_type == text_problems.VocabType.CHARACTER
       task.generate_data(data_dir, tmp_dir, task_id)
 
-  def add_task_id(self, task, example):
+  def add_task_id(self, task, example, encoder):
     """Convert example to code switching mode by adding a task id."""
     if hasattr(task, "class_labels"):
-      # TODO(urvashik): handle the case where num_labels > 9
-      example["targets"] = tf.cast(discretization.int_to_bit(
-          example["targets"], 1, base=10) + 50, tf.int64)
-      example["targets"] = tf.squeeze(example["targets"], axis=[-1])
+      if self.vocab_type == text_problems.VocabType.CHARACTER:
+        # TODO(urvashik): handle the case where num_labels > 9
+        example["targets"] = tf.cast(discretization.int_to_bit(
+            example["targets"], 1, base=10) + 50, tf.int64)
+        example["targets"] = tf.squeeze(example["targets"], axis=[-1])
+      elif self.vocab_type == text_problems.VocabType.SUBWORD:
+        offset = encoder.vocab_size + len(self.task_list)
+        # An additional +1 because of 0-indexing
+        example["targets"] = offset + example["targets"] + 1
 
     if task.has_inputs:
       inputs = example.pop("inputs")
@@ -66,10 +71,30 @@ def get_hparams(self, model_hparams=None):
       return self._hparams
 
     self._hparams = self.task_list[0].get_hparams(model_hparams)
+    # increase the vocab size in order to account for task ids
+    vocab_size_inc = len(self.task_list)
+    vocab_size_inc += self.get_max_num_classes()
+    vocab_size = self._hparams.vocabulary["targets"].vocab_size
+    self._hparams.target_modality = (registry.Modalities.SYMBOL,
+                                     vocab_size + vocab_size_inc)
 
     return self._hparams
 
   def flatten_zip(self, *args):
+    """A list of examples to a dataset containing mixed examples.
+
+    Given a list of `n` dataset examples, flatten them by converting
+    each element into a dataset and concatenating them to convert into a
+    single dataset.
+
+    Args:
+      *args: A list containing one example each from `n` different datasets.
+
+    Returns:
+      flattened: A new dataset containing the examples from the list as part
+        of a single dataset.
+    """
+
     flattened = tf.data.Dataset.from_tensors(args[0])
     for ex in args[1:]:
       flattened = flattened.concatenate(tf.data.Dataset.from_tensors(ex))
@@ -90,33 +115,38 @@ def dataset(self,
               num_partitions=1,
               max_records=-1):
 
+    # A list of datasets corresponding to the tasks in the task_list object
+    # that need to be mixed.
     datasets = []
     is_training = mode == tf.estimator.ModeKeys.TRAIN
 
+    primary_task = self.task_list[0]
+    if primary_task.has_inputs:
+      raise ValueError("Only support language models as primary problem which "
+                       "supplies the vocabulary and the hparams.")
+    enc = primary_task.feature_encoders(data_dir=data_dir)["targets"]
+
     for idx, task in enumerate(self.task_list):
       task_dataset = task.dataset(mode, data_dir, num_threads,
                                   output_buffer_size, shuffle_files,
                                   hparams, preprocess, dataset_split,
                                   shard, partition_id, num_partitions,
                                   max_records)
+
       if idx == 0:
-        self.update_task_ids(data_dir)
+        self.update_task_ids(enc)
 
       if is_training:
         task_dataset = task_dataset.repeat()
       # pylint: disable=cell-var-from-loop
-      task_dataset = task_dataset.map(lambda x: self.add_task_id(task, x))
+      task_dataset = task_dataset.map(lambda x: self.add_task_id(task, x, enc))
       datasets.append(task_dataset)
 
+    # Setup the problem hparams by setting them to the LM task hparams.
     self.get_hparams()
 
-    if is_training:
-      single_mtl_dataset = tf.data.Dataset.zip(tuple(datasets)).flat_map(
-          self.flatten_zip)
-    else:
-      single_mtl_dataset = datasets[0]
-      for data in datasets[1:]:
-        single_mtl_dataset = single_mtl_dataset.concatenate(data)
+    single_mtl_dataset = tf.data.Dataset.zip(tuple(datasets)).flat_map(
+        self.flatten_zip)
 
     return single_mtl_dataset
 
@@ -125,19 +155,40 @@ def eval_metrics(self):
         metrics.Metrics.ACC, metrics.Metrics.NEG_LOG_PERPLEXITY
     ]
 
-  def update_task_ids(self, data_dir):
-    primary_task = self.task_list[0]
-    if primary_task.has_inputs:
-      raise ValueError("Only support language models as primary problem which "
-                       "supplies the vocabulary and the hparams.")
+  def update_task_ids(self, encoder):
+    """Generate task_ids for each problem.
 
-    encoder = primary_task.feature_encoders(data_dir=data_dir)["targets"]
+    These ids correspond to the index of the task in the task_list.
 
+    Args:
+      encoder: this provides the size of the vocab which is used to compute
+        the index offset.
+    """
+    primary_task = self.task_list[0]
     id_offset = encoder.vocab_size + text_encoder.NUM_RESERVED_TOKENS
     if hasattr(primary_task, "additional_reserved_tokens"):
       id_offset += len(primary_task.additional_reserved_tokens)
 
     for idx, _ in enumerate(self.task_list):
-      # protect against the ord mapping of chars with the 2x multiplier.
-      self.task_list[idx].set_task_id(idx + 2 * id_offset)
+      # Subtract one to get actual indices in the context of 0-indexing
+      self.task_list[idx].set_task_id(idx + id_offset - 1)
       print(self.task_list[idx].task_id)
+
+  def get_max_num_classes(self):
+    """Compute the maximum number of classes any subtask has.
+
+    This is useful for modifying the size of the softmax to include the output
+    labels for the classification tasks. Currently, labels from different tasks
+    are overloaded.
+
+    Returns:
+      num: Highest number of output classes in any text classification sub-task
+        within this MultiProblem.
+    """
+    num = 0
+    for task in self.task_list:
+      if hasattr(task, "num_classes"):
+        if num < task.num_classes:
+          num = task.num_classes
+
+    return num
diff --git a/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index 7ecdad85e..be879be2f 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -67,16 +67,6 @@ def vocab_filename(self):
   def num_classes(self):
     return 3
 
-  @property
-  def concat_token(self):
-    return "<EN-PR-HYP>"
-
-  @property
-  def concat_id(self):
-    if self.vocab_type == text_problems.VocabType.CHARACTER:
-      return problem.TaskID.EN_PR_HYP
-    return 2
-
   def class_labels(self, data_dir):
     del data_dir
     # Note this binary classification is different from usual MNLI.
@@ -135,3 +125,12 @@ def vocab_type(self):
 
   def global_task_id(self):
     return problem.TaskID.THREE_CL_NLI
+
+
+@registry.register_problem
+class MultiNLISharedVocab(MultiNLI):
+  """MultiNLI classification problems with the LM1b vocabulary"""
+
+  @property
+  def vocab_filename(self):
+    return "vocab.lm1b.en.%d" % 2**15
diff --git a/tensor2tensor/data_generators/qnli.py b/tensor2tensor/data_generators/qnli.py
index 93773adcb..a075e8fb3 100644
--- a/tensor2tensor/data_generators/qnli.py
+++ b/tensor2tensor/data_generators/qnli.py
@@ -67,16 +67,6 @@ def vocab_filename(self):
   def num_classes(self):
     return 2
 
-  @property
-  def concat_token(self):
-    return "<EN-Q-CONT>"
-
-  @property
-  def concat_id(self):
-    if self.vocab_type == text_problems.VocabType.CHARACTER:
-      return problem.TaskID.EN_Q_CONT
-    return 2
-
   def class_labels(self, data_dir):
     del data_dir
     # Note this binary classification is different from usual MNLI.
diff --git a/tensor2tensor/data_generators/quora_qpairs.py b/tensor2tensor/data_generators/quora_qpairs.py
index 7cfbb5fd6..89736d7f3 100644
--- a/tensor2tensor/data_generators/quora_qpairs.py
+++ b/tensor2tensor/data_generators/quora_qpairs.py
@@ -67,16 +67,6 @@ def vocab_filename(self):
   def num_classes(self):
     return 2
 
-  @property
-  def concat_token(self):
-    return "<SENT_SEP>"
-
-  @property
-  def concat_id(self):
-    if self.vocab_type == text_problems.VocabType.CHARACTER:
-      return problem.TaskID.EN_SENT_PAIR
-    return 2
-
   def class_labels(self, data_dir):
     del data_dir
     return ["not_duplicate", "duplicate"]
diff --git a/tensor2tensor/data_generators/rte.py b/tensor2tensor/data_generators/rte.py
index d62624316..0ea78e144 100644
--- a/tensor2tensor/data_generators/rte.py
+++ b/tensor2tensor/data_generators/rte.py
@@ -67,16 +67,6 @@ def vocab_filename(self):
   def num_classes(self):
     return 2
 
-  @property
-  def concat_token(self):
-    return "<EN-PR-HYP>"
-
-  @property
-  def concat_id(self):
-    if self.vocab_type == text_problems.VocabType.CHARACTER:
-      return problem.TaskID.EN_PR_HYP
-    return 2
-
   def class_labels(self, data_dir):
     del data_dir
     # Note this binary classification is different from usual MNLI.
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 16afc53a1..4e2448f12 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -515,18 +515,7 @@ class TextConcat2ClassProblem(Text2ClassProblem):
   For problems where there are multiple input sentences and we wish to concat
   these inputs with a special delimiter. See, for example, NLI tasks.
   """
-
-  @property
-  def concat_token(self):
-    raise NotImplementedError()
-
-  @property
-  def concat_id(self):
-    raise NotImplementedError()
-
-  @property
-  def additional_reserved_tokens(self):
-    return [self.concat_token]
+  CONCAT_TOKEN = "$"
 
   def generate_text_for_vocab(self, data_dir, tmp_dir):
     for i, sample in enumerate(
@@ -545,7 +534,7 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
         inputs += encoder.encode(inp)
         inputs.append(text_encoder.EOS_ID)
         if idx < len(sample["inputs"])-1:
-          inputs.append(self.concat_id)
+          inputs.append(encoder.encode(self.CONCAT_TOKEN)[0])
       label = sample["label"]
       yield {"inputs": inputs, "targets": [label]}
 
diff --git a/tensor2tensor/data_generators/wnli.py b/tensor2tensor/data_generators/wnli.py
index e44459fc6..a00ff0754 100644
--- a/tensor2tensor/data_generators/wnli.py
+++ b/tensor2tensor/data_generators/wnli.py
@@ -67,16 +67,6 @@ def vocab_filename(self):
   def num_classes(self):
     return 2
 
-  @property
-  def concat_token(self):
-    return "<EN-PR-HYP>"
-
-  @property
-  def concat_id(self):
-    if self.vocab_type == text_problems.VocabType.CHARACTER:
-      return problem.TaskID.EN_PR_HYP
-    return 2
-
   def class_labels(self, data_dir):
     del data_dir
     # Note this binary classification is different from usual MNLI.

From 1952055dfe5042dde6a4fcb411f844fb71ecad24 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Fri, 3 Aug 2018 16:29:35 -0700
Subject: [PATCH 0470/2720] whitelist of games for model based.

PiperOrigin-RevId: 207342890
---
 .../data_generators/gym_problems_specs.py     | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index 6c961c2a5..e957b67ab 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -49,6 +49,25 @@
 # * defender
 # * kaboom
 
+# Subset of games with promissing results on model based training.
+ATARI_WHITELIST_GAMES = [
+    "amidar",
+    "bank_heist",
+    "berzerk",
+    "boxing",
+    "breakout",
+    "crazy_climber",
+    "freeway",
+    "frostbite",
+    "gopher",
+    "hero",
+    "kung_fu_master",
+    "pong",
+    "road_runner",
+    "seaquest",
+    "wrapped_full_pong",  # TODO(blazej): check if we get equally good results
+]                         # on vanilla pong.
+
 
 @registry.register_problem
 class GymPongRandom(GymDiscreteProblem):

From ab2eec23661106b49862a5224479ddc64c6d4efe Mon Sep 17 00:00:00 2001
From: Xu Song <eson.org@gmail.com>
Date: Sat, 4 Aug 2018 07:36:31 +0800
Subject: [PATCH 0471/2720] update to wmt18 dataset (#904)

---
 tensor2tensor/data_generators/translate_encs.py |  6 +++---
 tensor2tensor/data_generators/translate_ende.py |  6 +++---
 tensor2tensor/data_generators/translate_enzh.py | 15 +++++++--------
 3 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/data_generators/translate_encs.py b/tensor2tensor/data_generators/translate_encs.py
index bc4e5df3b..be7341d53 100644
--- a/tensor2tensor/data_generators/translate_encs.py
+++ b/tensor2tensor/data_generators/translate_encs.py
@@ -35,9 +35,9 @@
       "11234/1-1458/data-plaintext-format.tar"),
      ("tsv", 3, 2, "data.plaintext-format/*train.gz")],
     [
-        "http://data.statmt.org/wmt17/translation-task/training-parallel-nc-v12.tgz",  # pylint: disable=line-too-long
-        ("training/news-commentary-v12.cs-en.en",
-         "training/news-commentary-v12.cs-en.cs")
+        "http://data.statmt.org/wmt18/translation-task/training-parallel-nc-v13.tgz",  # pylint: disable=line-too-long
+        ("training/news-commentary-v13.cs-en.en",
+         "training/news-commentary-v13.cs-en.cs")
     ],
     [
         "http://www.statmt.org/wmt13/training-parallel-commoncrawl.tgz",
diff --git a/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index 70bd53528..095d5fcff 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -31,9 +31,9 @@
 
 _ENDE_TRAIN_DATASETS = [
     [
-        "http://data.statmt.org/wmt17/translation-task/training-parallel-nc-v12.tgz",  # pylint: disable=line-too-long
-        ("training/news-commentary-v12.de-en.en",
-         "training/news-commentary-v12.de-en.de")
+        "http://data.statmt.org/wmt18/translation-task/training-parallel-nc-v13.tgz",  # pylint: disable=line-too-long
+        ("training/news-commentary-v13.de-en.en",
+         "training/news-commentary-v13.de-en.de")
     ],
     [
         "http://www.statmt.org/wmt13/training-parallel-commoncrawl.tgz",
diff --git a/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index 1944b0fff..6503b8cfc 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -36,23 +36,22 @@
 # End-of-sentence marker.
 EOS = text_encoder.EOS_ID
 
-# This is far from being the real WMT17 task - only toyset here
+# This is far from being the real WMT18 task - only toyset here
 # you need to register to get UN data and CWT data. Also, by convention,
 # this is EN to ZH - use translate_enzh_wmt8k_rev for ZH to EN task
 #
-# News Commentary, around 220k lines
-# This dataset is only a small fraction of full WMT17 task
+# News Commentary, around 252k lines
+# This dataset is only a small fraction of full WMT18 task
 _NC_TRAIN_DATASETS = [[
-    "http://data.statmt.org/wmt17/translation-task/training-parallel-nc-v12"
-    ".tgz", [
-        "training/news-commentary-v12.zh-en.en",
-        "training/news-commentary-v12.zh-en.zh"
+    "http://data.statmt.org/wmt18/translation-task/training-parallel-nc-v13.tgz", [
+        "training/news-commentary-v13.zh-en.en",
+        "training/news-commentary-v13.zh-en.zh"
     ]
 ]]
 
 # Test set from News Commentary. 2000 lines
 _NC_TEST_DATASETS = [[
-    "http://data.statmt.org/wmt17/translation-task/dev.tgz",
+    "http://data.statmt.org/wmt18/translation-task/dev.tgz",
     ("dev/newsdev2017-enzh-src.en.sgm", "dev/newsdev2017-enzh-ref.zh.sgm")
 ]]
 

From 68d383e6210f9be27904167bd3bd51dc27858a38 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 3 Aug 2018 16:55:13 -0700
Subject: [PATCH 0472/2720] internal merge of PR #904

PiperOrigin-RevId: 207347039
---
 tensor2tensor/data_generators/translate_enzh.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index 6503b8cfc..68ba28daa 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -42,8 +42,9 @@
 #
 # News Commentary, around 252k lines
 # This dataset is only a small fraction of full WMT18 task
+_STAT_MT_URL = "http://data.statmt.org/wmt18/translation-task/"
 _NC_TRAIN_DATASETS = [[
-    "http://data.statmt.org/wmt18/translation-task/training-parallel-nc-v13.tgz", [
+    _STAT_MT_URL + "training-parallel-nc-v13.tgz", [
         "training/news-commentary-v13.zh-en.en",
         "training/news-commentary-v13.zh-en.zh"
     ]
@@ -51,7 +52,7 @@
 
 # Test set from News Commentary. 2000 lines
 _NC_TEST_DATASETS = [[
-    "http://data.statmt.org/wmt18/translation-task/dev.tgz",
+    _STAT_MT_URL + "dev.tgz",
     ("dev/newsdev2017-enzh-src.en.sgm", "dev/newsdev2017-enzh-ref.zh.sgm")
 ]]
 

From a2a8b76ba035303e74f0934320b1548ea81c6842 Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Fri, 3 Aug 2018 18:05:35 -0700
Subject: [PATCH 0473/2720] Support text for discrete autoencoder

PiperOrigin-RevId: 207354213
---
 tensor2tensor/data_generators/image_utils.py  |   5 +
 tensor2tensor/layers/discretization.py        | 132 +++++++++++++-----
 tensor2tensor/layers/discretization_test.py   |   8 +-
 tensor2tensor/models/research/autoencoders.py |  74 ++++++----
 4 files changed, 155 insertions(+), 64 deletions(-)

diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 9f4b6b0c9..98ec9da81 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -100,6 +100,11 @@ def num_channels(self):
     """Number of color channels."""
     return 3
 
+  @property
+  def vocab_size(self):
+    """Number of pixel values."""
+    return 256
+
   def example_reading_spec(self):
     data_fields = {
         "image/encoded": tf.FixedLenFeature((), tf.string),
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index c55f64b69..4de85375a 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -673,8 +673,8 @@ def discrete_bottleneck(inputs,
           updated_ema_means_res = moving_averages.assign_moving_average(
               ema_means[i], dw, decay, zero_debias=False)
           n = tf.reduce_sum(updated_ema_count_res, axis=-1, keep_dims=True)
-          updated_ema_count_res = ((updated_ema_count_res + epsilon) /
-                                   (n + 2**z_size * epsilon) * n)
+          updated_ema_count_res = (
+              (updated_ema_count_res + epsilon) / (n + 2**z_size * epsilon) * n)
           # pylint: disable=g-no-augmented-assignment
           updated_ema_means_res = updated_ema_means_res / tf.expand_dims(
               updated_ema_count_res, axis=-1)
@@ -787,17 +787,17 @@ def discrete_bottleneck(inputs,
 # * The [method]_unbottleneck function moves from discretized state to dense.
 
 
-def get_vq_bottleneck(bottleneck_size, hidden_size):
+def get_vq_codebook(codebook_size, hidden_size):
   """Get lookup table for VQ bottleneck."""
   with tf.variable_scope("vq", reuse=tf.AUTO_REUSE):
     means = tf.get_variable(
         name="means",
-        shape=[bottleneck_size, hidden_size],
+        shape=[codebook_size, hidden_size],
         initializer=tf.uniform_unit_scaling_initializer())
 
     ema_count = tf.get_variable(
         name="ema_count",
-        shape=[bottleneck_size],
+        shape=[codebook_size],
         initializer=tf.constant_initializer(0),
         trainable=False)
 
@@ -840,9 +840,27 @@ def vq_discrete_bottleneck(x,
                            num_samples=10):
   """Simple vector quantized discrete bottleneck."""
   bottleneck_size = 2**bottleneck_bits
+  return vq_body(
+      x,
+      bottleneck_size,
+      beta=beta,
+      decay=decay,
+      epsilon=epsilon,
+      soft_em=soft_em,
+      num_samples=num_samples)
+
+
+def vq_body(x,
+            codebook_size,
+            beta=0.25,
+            decay=0.999,
+            epsilon=1e-5,
+            soft_em=False,
+            num_samples=10):
+  """Discretize each x into one of codebook_size codes."""
   x_shape = common_layers.shape_list(x)
   hidden_size = x_shape[-1]
-  means, ema_means, ema_count = get_vq_bottleneck(bottleneck_size, hidden_size)
+  means, ema_means, ema_count = get_vq_codebook(codebook_size, hidden_size)
   x = tf.reshape(x, [-1, hidden_size])
   x_means_hot, e_loss = vq_nearest_neighbor(
       x, means, soft_em=soft_em, num_samples=num_samples)
@@ -850,33 +868,67 @@ def vq_discrete_bottleneck(x,
   # Update the ema variables
   updated_ema_count = moving_averages.assign_moving_average(
       ema_count,
-      tf.reduce_sum(
-          tf.reshape(x_means_hot, shape=[-1, bottleneck_size]), axis=0),
+      tf.reduce_sum(tf.reshape(x_means_hot, shape=[-1, codebook_size]), axis=0),
       decay,
       zero_debias=False)
 
   dw = tf.matmul(x_means_hot, x, transpose_a=True)
-  updated_ema_means = tf.identity(moving_averages.assign_moving_average(
-      ema_means, dw, decay, zero_debias=False))
+  updated_ema_means = tf.identity(
+      moving_averages.assign_moving_average(
+          ema_means, dw, decay, zero_debias=False))
   n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
   updated_ema_count = (
-      (updated_ema_count + epsilon) / (n + bottleneck_size * epsilon) * n)
+      (updated_ema_count + epsilon) / (n + codebook_size * epsilon) * n)
   updated_ema_means /= tf.expand_dims(updated_ema_count, axis=-1)
   with tf.control_dependencies([e_loss]):
     update_means = means.assign(updated_ema_means)
     with tf.control_dependencies([update_means]):
       loss = beta * e_loss
 
-  d = tf.reshape(x_means_hot, x_shape[:-1] + [bottleneck_size])
+  d = tf.reshape(x_means_hot, x_shape[:-1] + [codebook_size])
   return d, loss
 
 
+def vq_loss(x,
+            one_hot_targets,
+            codebook_size,
+            beta=0.25,
+            decay=0.999,
+            epsilon=1e-5,
+            soft_em=False,
+            num_samples=10):
+  """Simple vector quantized discrete bottleneck."""
+  x_shape = common_layers.shape_list(x)
+  target_shape = common_layers.shape_list(one_hot_targets)
+  hidden_size = x_shape[-1]
+  means, _, _ = get_vq_codebook(codebook_size, hidden_size)
+  x = tf.reshape(x, [-1, hidden_size])
+
+  one_hot_targets = tf.reshape(one_hot_targets, [-1, target_shape[-1]])
+  target_means = tf.matmul(one_hot_targets, means)
+  targets_loss = tf.reduce_sum((x - target_means)**2, -1)
+  targets_loss = tf.reduce_mean(targets_loss)
+
+  discrete_x, code_loss = vq_body(
+      x,
+      codebook_size,
+      beta=beta,
+      decay=decay,
+      epsilon=epsilon,
+      soft_em=soft_em,
+      num_samples=num_samples)
+
+  discrete_x = tf.reshape(discrete_x, x_shape[:-1] + [codebook_size])
+  target_means = tf.reshape(target_means, target_shape[:-1] + [hidden_size])
+  return discrete_x, target_means, code_loss, targets_loss
+
+
 def vq_discrete_unbottleneck(x, hidden_size):
   """Simple undiscretization from vector quantized representation."""
   x_shape = common_layers.shape_list(x)
   x = tf.to_float(x)
   bottleneck_size = common_layers.shape_list(x)[-1]
-  means, _, _ = get_vq_bottleneck(bottleneck_size, hidden_size)
+  means, _, _ = get_vq_codebook(bottleneck_size, hidden_size)
   result = tf.matmul(tf.reshape(x, [-1, x_shape[-1]]), means)
   return tf.reshape(result, x_shape[:-1] + [hidden_size])
 
@@ -1100,7 +1152,7 @@ def gumbel_softmax_discrete_bottleneck(x,
   bottleneck_size = 2**bottleneck_bits
   x_shape = common_layers.shape_list(x)
   hidden_size = x_shape[-1]
-  means, ema_means, ema_count = get_vq_bottleneck(bottleneck_size, hidden_size)
+  means, ema_means, ema_count = get_vq_codebook(bottleneck_size, hidden_size)
   x = tf.reshape(x, [-1, hidden_size])
 
   bottleneck_size = common_layers.shape_list(means)[0]
@@ -1124,8 +1176,8 @@ def gumbel_softmax_discrete_bottleneck(x,
       (log_class_probs + gumbel_samples) / temperature)
 
   # Calculate KL between q and a uniform prior.
-  kl = tf.reduce_sum(class_probs * (log_class_probs -
-                                    tf.log(1.0/bottleneck_size)), -1)
+  kl = tf.reduce_sum(
+      class_probs * (log_class_probs - tf.log(1.0 / bottleneck_size)), -1)
   if summary:
     tf.summary.histogram("KL", tf.reshape(kl, [-1]))
 
@@ -1137,8 +1189,8 @@ def gumbel_softmax_discrete_bottleneck(x,
         x_means_hot - gumbel_softmax_samples)
   else:
     x_means_assignments = gumbel_softmax_samples
-  x_means_assignments_flat = tf.reshape(
-      x_means_assignments, [-1, bottleneck_size])
+  x_means_assignments_flat = tf.reshape(x_means_assignments,
+                                        [-1, bottleneck_size])
   x_means = tf.matmul(x_means_assignments_flat, means)
   commitment_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
 
@@ -1151,8 +1203,9 @@ def gumbel_softmax_discrete_bottleneck(x,
       zero_debias=False)
 
   dw = tf.matmul(x_means_assignments, x, transpose_a=True)
-  updated_ema_means = tf.identity(moving_averages.assign_moving_average(
-      ema_means, dw, decay, zero_debias=False))
+  updated_ema_means = tf.identity(
+      moving_averages.assign_moving_average(
+          ema_means, dw, decay, zero_debias=False))
   n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
   updated_ema_count = (
       (updated_ema_count + epsilon) / (n + bottleneck_size * epsilon) * n)
@@ -1165,16 +1218,16 @@ def gumbel_softmax_discrete_bottleneck(x,
   # Add KL loss.
   loss += tf.reduce_mean(kl)
 
-  x_means_assignments = tf.reshape(
-      x_means_assignments, x_shape[:-1] + [bottleneck_size])
+  x_means_assignments = tf.reshape(x_means_assignments,
+                                   x_shape[:-1] + [bottleneck_size])
   return x_means_assignments, loss
 
 
 def tanh_discrete_bottleneck(x, bottleneck_bits, bottleneck_noise,
                              discretize_warmup_steps, mode):
   """Simple discretization through tanh, flip bottleneck_noise many bits."""
-  x = tf.tanh(tf.layers.dense(x, bottleneck_bits,
-                              name="tanh_discrete_bottleneck"))
+  x = tf.tanh(
+      tf.layers.dense(x, bottleneck_bits, name="tanh_discrete_bottleneck"))
   d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
   if mode == tf.estimator.ModeKeys.TRAIN:
     noise = tf.random_uniform(common_layers.shape_list(x))
@@ -1191,9 +1244,13 @@ def tanh_discrete_unbottleneck(x, hidden_size):
   return x
 
 
-def isemhash_bottleneck(x, bottleneck_bits, bottleneck_noise,
-                        discretize_warmup_steps, mode,
-                        isemhash_noise_dev=0.5, isemhash_mix_prob=0.5):
+def isemhash_bottleneck(x,
+                        bottleneck_bits,
+                        bottleneck_noise,
+                        discretize_warmup_steps,
+                        mode,
+                        isemhash_noise_dev=0.5,
+                        isemhash_mix_prob=0.5):
   """Improved semantic hashing bottleneck."""
   with tf.variable_scope("isemhash_bottleneck"):
     x = tf.layers.dense(x, bottleneck_bits, name="dense")
@@ -1208,9 +1265,12 @@ def isemhash_bottleneck(x, bottleneck_bits, bottleneck_noise,
       noise = tf.random_uniform(common_layers.shape_list(x))
       noise = 2.0 * tf.to_float(tf.less(bottleneck_noise, noise)) - 1.0
       d *= noise
-      d = common_layers.mix(d, 2.0 * y - 1.0, discretize_warmup_steps,
-                            mode == tf.estimator.ModeKeys.TRAIN,
-                            max_prob=isemhash_mix_prob)
+      d = common_layers.mix(
+          d,
+          2.0 * y - 1.0,
+          discretize_warmup_steps,
+          mode == tf.estimator.ModeKeys.TRAIN,
+          max_prob=isemhash_mix_prob)
     return d, 0.0
 
 
@@ -1258,8 +1318,8 @@ def parametrized_bottleneck(x, hparams):
                                               hard=False,
                                               summary=True)
 
-  raise ValueError("Unsupported hparams.bottleneck_kind %s"
-                   % hparams.bottleneck_kind)
+  raise ValueError(
+      "Unsupported hparams.bottleneck_kind %s" % hparams.bottleneck_kind)
 
 
 def parametrized_unbottleneck(x, hidden_size, hparams):
@@ -1267,12 +1327,12 @@ def parametrized_unbottleneck(x, hidden_size, hparams):
   if hparams.bottleneck_kind == "tanh_discrete":
     return tanh_discrete_unbottleneck(x, hidden_size)
   if hparams.bottleneck_kind == "isemhash":
-    return isemhash_unbottleneck(
-        x, hidden_size, hparams.isemhash_filter_size_multiplier)
+    return isemhash_unbottleneck(x, hidden_size,
+                                 hparams.isemhash_filter_size_multiplier)
   if hparams.bottleneck_kind in ["vq", "em", "gumbel_softmax"]:
     return vq_discrete_unbottleneck(x, hidden_size)
-  raise ValueError("Unsupported hparams.bottleneck_kind %s"
-                   % hparams.bottleneck_kind)
+  raise ValueError(
+      "Unsupported hparams.bottleneck_kind %s" % hparams.bottleneck_kind)
 
 
 def iaf_hparams(hidden_size=512, filter_size=4096):
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index 384fc5008..3970d14d8 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -124,11 +124,11 @@ def testGetVQBottleneck(self):
     bottleneck_bits = 2
     bottleneck_size = 2**bottleneck_bits
     hidden_size = 3
-    means, _, ema_count = discretization.get_vq_bottleneck(bottleneck_size,
-                                                           hidden_size)
+    means, _, ema_count = discretization.get_vq_codebook(
+        bottleneck_size, hidden_size)
     assign_op = means.assign(tf.zeros(shape=[bottleneck_size, hidden_size]))
-    means_new, _, _ = discretization.get_vq_bottleneck(bottleneck_size,
-                                                       hidden_size)
+    means_new, _, _ = discretization.get_vq_codebook(bottleneck_size,
+                                                     hidden_size)
     with self.test_session() as sess:
       tf.global_variables_initializer().run()
       sess.run(assign_op)
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 9fbd76680..6b9da76a4 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -18,6 +18,8 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.data_generators import image_utils
+from tensor2tensor.data_generators import text_problems
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
@@ -145,6 +147,7 @@ def body(self, features):
     is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
       x = features["targets"]
+      labels = features["targets_raw"]
       shape = common_layers.shape_list(x)
       is1d = shape[2] == 1
       self.is1d = is1d
@@ -181,42 +184,62 @@ def body(self, features):
     # Cut to the right size and mix before returning.
     res = x[:, :shape[1], :shape[2], :]
 
-    num_channels = self.hparams.problem.num_channels
-    reconstr = tf.layers.dense(res, num_channels)
-    reconstr = tf.nn.sigmoid(reconstr)
-    reconstr = 256. * reconstr - 0.5
+    is_image = isinstance(self.hparams.problem, image_utils.ImageProblem)
+    if is_image:
+      vocab_size = self.hparams.problem.vocab_size
+
+      res = tf.layers.dense(
+          res, self.hparams.problem.num_channels * self.hparams.hidden_size)
+      output_shape = common_layers.shape_list(res)[:-1] + [
+          self.hparams.problem.num_channels, self.hparams.hidden_size
+      ]
+      res = tf.reshape(res, output_shape)
+    elif isinstance(self.hparams.problem, text_problems.Text2TextProblem):
+      vocab_size = self._problem_hparams.target_modality.top_dimensionality
+      res = tf.layers.dense(res, self.hparams.hidden_size)
+    else:
+      raise Exception("Unsupported problem type: %s" % self.hparams.problem)
+
+    one_hot_labels = tf.one_hot(labels, vocab_size)
+    code_loss_gan = 0.0
+    if hparams.gan_loss_factor != 0.0:
+      res_gan, res = tf.split(res, 2, axis=0)
+      with tf.variable_scope("vq"):
+        reconstr_gan, _, code_loss_gan, _ = discretization.vq_loss(
+            res, one_hot_labels, vocab_size)
+
+    with tf.variable_scope("vq", reuse=tf.AUTO_REUSE):
+      reconstr, target_codes, code_loss, targets_loss = discretization.vq_loss(
+          res, one_hot_labels, vocab_size)
+
     # Add GAN loss if requested.
     gan_loss = 0.0
     if hparams.gan_loss_factor != 0.0:
-      # Split back if we added a purely sampled batch.
-      reconstr_gan, reconstr = tf.split(reconstr, 2, axis=0)
-      tf.summary.image(
-          "gan",
-          common_layers.tpu_safe_image_summary(reconstr_gan),
-          max_outputs=1)
-      orig_rgb = tf.to_float(features["targets_raw"])
+      if is_image:
+        tf.summary.image(
+            "gan",
+            common_layers.tpu_safe_image_summary(tf.argmax(reconstr_gan, -1)),
+            max_outputs=1)
 
       def discriminate(x):
         return self.discriminator(x, is_training=is_training)
 
-      gan_loss = common_layers.sliced_gan_loss(orig_rgb,
-                                               reverse_gradient(reconstr_gan),
+      gan_loss = common_layers.sliced_gan_loss(target_codes,
+                                               reverse_gradient(res_gan),
                                                discriminate,
                                                self.hparams.num_sliced_vecs)
       gan_loss *= hparams.gan_loss_factor
 
-    tf.summary.image(
-        "ae", common_layers.tpu_safe_image_summary(reconstr), max_outputs=1)
-
-    # Project to correct vocab_size/channels
-    training = tf.reduce_mean(
-        tf.square(tf.to_float(features["targets_raw"]) - reconstr))
-
-    outputs = tf.round(reconstr)
-    outputs = tf.one_hot(tf.to_int32(outputs), 256)
+    if is_image:
+      tf.summary.image(
+          "ae",
+          common_layers.tpu_safe_image_summary(tf.argmax(reconstr, -1)),
+          max_outputs=1)
 
-    return outputs, {
-        "training": training,
+    return reconstr, {
+        "training": targets_loss,
+        "code_loss": code_loss,
+        "code_loss_gan": code_loss_gan,
         "b_loss": b_loss,
         "gan_loss": -gan_loss
     }
@@ -318,6 +341,7 @@ def body(self, features):
         basic1d = hparams.sampled_basic1d_tensor
       else:
         hparams.sampled_basic1d_tensor = basic1d
+    print(common_layers.shape_list(features["targets"]))
     # Prepare inputs for autoregressive modes.
     if common_layers.shape_list(features["targets"])[1] == 1:
       # This happens on the first step of predicitions.
@@ -868,6 +892,8 @@ def autoencoder_ordered_text():
   hparams.max_hidden_size = 4096
   hparams.bottleneck_warmup_steps = 10000
   hparams.discretize_warmup_steps = 15000
+  hparams.target_modality = "symbol:identity"
+  hparams.input_modalities = "symbol:identity"
   return hparams
 
 

From 044600688d83640adbb7b3c21ec91fb258217da2 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 3 Aug 2018 18:51:31 -0700
Subject: [PATCH 0474/2720] Double PPO epoch length in the last iteration,
 adjust variable names to not collide.

PiperOrigin-RevId: 207357509
---
 tensor2tensor/rl/collect.py             | 13 +++++++------
 tensor2tensor/rl/model_rl_experiment.py |  1 +
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 515bb8fce..3c7441a21 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -130,12 +130,13 @@ def initialization_lambda(sess):
       for batch_env in to_initialize:
         batch_env.initialize(sess)
 
-    memory = [tf.get_variable("collect_memory_{}".format(name),
-                              shape=[hparams.epoch_length]+shape,
-                              dtype=dtype,
-                              initializer=tf.zeros_initializer(),
-                              trainable=False)
-              for (shape, dtype, name) in rollout_metadata]
+    memory = [
+        tf.get_variable("collect_memory_%d_%s" % (hparams.epoch_length, name),
+                        shape=[hparams.epoch_length] + shape,
+                        dtype=dtype,
+                        initializer=tf.zeros_initializer(),
+                        trainable=False)
+        for (shape, dtype, name) in rollout_metadata]
 
     cumulative_rewards = tf.get_variable("cumulative_rewards", len(batch_env),
                                          trainable=False)
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 79e01fc57..5a563231a 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -152,6 +152,7 @@ def train_agent(problem_name, agent_model_dir,
   ppo_epochs_num = hparams.ppo_epochs_num
   if is_final_epoch:
     ppo_epochs_num *= 2
+    ppo_hparams.epoch_length *= 2
   ppo_hparams.save_models_every_epochs = ppo_epochs_num
   ppo_hparams.world_model_dir = world_model_dir
   ppo_hparams.add_hparam("force_beginning_resets", True)

From 4bf60607c87717f68fa5516b2cb74d23477d9ca4 Mon Sep 17 00:00:00 2001
From: cclauss <cclauss@bluewin.ch>
Date: Sun, 5 Aug 2018 02:14:48 +0200
Subject: [PATCH 0475/2720] Travis CI: Python syntax errors or undefined names
 (#942)

Use [flake8](http://flake8.pycqa.org) to look for Python syntax errors or undefined names on both Python 2 and Python 3.

__E901,E999,F821,F822,F823__ are the "showstopper" flake8 issues that can halt the runtime with a SyntaxError, NameError, etc.  Most other flake8 issues are merely "style violations" -- useful for readability but they do not effect runtime safety.  This PR therefore recommends a flake8 run of these tests on the entire codebase.
* F821: undefined name `name`
* F822: undefined name `name` in `__all__`
* F823: local variable `name` referenced before assignment
* E901: SyntaxError or IndentationError
* E999: SyntaxError -- failed to compile a file into an Abstract Syntax Tree
---
 .travis.yml | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/.travis.yml b/.travis.yml
index b162ae9f9..65c1bdb9d 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -48,6 +48,11 @@ install:
   # Make sure we have the latest version of numpy - avoid problems we were
   # seeing with Python 3
   - pip install -q -U numpy
+before_script:
+  # stop the build if there are Python syntax errors or undefined names
+  if [[ "$TF_VERSION" == "1.9.*" ]]; then
+    pip install flake8; flake8 . --count --select=E901,E999,F821,F822,F823 --show-source --statistics;
+  fi
 script:
   # Check import
   - python -c "from tensor2tensor.models import transformer; print(transformer.Transformer.__name__)"

From d34041ff910964c7a78de8786425856168e3a580 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 4 Aug 2018 17:17:39 -0700
Subject: [PATCH 0476/2720] internal merge of PR #942

PiperOrigin-RevId: 207416830
---
 .travis.yml | 7 ++++++-
 setup.py    | 2 +-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 65c1bdb9d..6f5f86bad 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -51,7 +51,12 @@ install:
 before_script:
   # stop the build if there are Python syntax errors or undefined names
   if [[ "$TF_VERSION" == "1.9.*" ]]; then
-    pip install flake8; flake8 . --count --select=E901,E999,F821,F822,F823 --show-source --statistics;
+    # * F821: undefined name `name`
+    # * F822: undefined name `name` in `__all__`
+    # * F823: local variable `name` referenced before assignment
+    # * E901: SyntaxError or IndentationError
+    # * E999: SyntaxError -- failed to compile a file into an Abstract Syntax Tree
+    flake8 . --count --select=F821,F822,F823,E901,E999 --show-source --statistics;
   fi
 script:
   # Check import
diff --git a/setup.py b/setup.py
index ff69570f3..2bf760637 100644
--- a/setup.py
+++ b/setup.py
@@ -52,7 +52,7 @@
         'tensorflow': ['tensorflow>=1.8.0'],
         'tensorflow_gpu': ['tensorflow-gpu>=1.8.0'],
         'tests': [
-            'absl-py', 'pytest', 'mock', 'pylint', 'jupyter', 'gsutil'
+            'absl-py', 'pytest', 'mock', 'pylint', 'jupyter', 'gsutil', 'flake8'
             # Need atari extras for Travis tests, but because gym is already in
             # install_requires, pip skips the atari extras, so we instead do an
             # explicit pip install gym[atari] for the tests.

From cb23cf354a3f25138ce02d97538c1eafdc92d7bb Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Sun, 5 Aug 2018 19:46:34 -0700
Subject: [PATCH 0477/2720] Add multigraph support to mpnn

PiperOrigin-RevId: 207483286
---
 .../common_message_passing_attention.py       | 107 ++++++------------
 1 file changed, 32 insertions(+), 75 deletions(-)

diff --git a/tensor2tensor/layers/common_message_passing_attention.py b/tensor2tensor/layers/common_message_passing_attention.py
index c716b1d39..af005c35d 100644
--- a/tensor2tensor/layers/common_message_passing_attention.py
+++ b/tensor2tensor/layers/common_message_passing_attention.py
@@ -40,7 +40,6 @@ def multihead_graph_attention(query_antecedent,
                               dropout_broadcast_dims=None,
                               adjacency_matrix=None,
                               num_edge_types=5,
-                              ignore_zero=True,
                               vars_3d=False,
                               **kwargs):
   """Multihead scaled-dot-product attention with input/output transformations.
@@ -71,7 +70,6 @@ def multihead_graph_attention(query_antecedent,
     adjacency_matrix: an optional tensor of shape [batch, len_q, len_q]
       containing edge vectors for attention
     num_edge_types: number of edge types, an int
-    ignore_zero: A flag that says that edge type 0 should be ignored.
     vars_3d: use 3-dimensional variables for input/output transformations
     **kwargs (dict): Parameters for the attention function
 
@@ -127,8 +125,7 @@ def multihead_graph_attention(query_antecedent,
           make_image_summary=make_image_summary,
           dropout_broadcast_dims=dropout_broadcast_dims,
           adjacency_matrix=adjacency_matrix,
-          num_edge_types=num_edge_types,
-          ignore_zero=ignore_zero)
+          num_edge_types=num_edge_types)
 
     x = common_attention.combine_heads(x)
 
@@ -152,7 +149,6 @@ def multihead_graph_attention(query_antecedent,
 def make_edge_vectors(adjacency_matrix,
                       num_edge_types,
                       depth,
-                      ignore_zero,
                       name=None):
   """Gets edge vectors for the edge types in the adjacency matrix.
 
@@ -160,8 +156,7 @@ def make_edge_vectors(adjacency_matrix,
     adjacency_matrix: A [batch, num_nodes, num_nodes] tensor of ints.
     num_edge_types: Number of different edge types
     depth: Number of channels
-    ignore_zero: Whether to use zeros vector for edge type 0.
-    name: a string
+    name: A optional string name for scoping
   Returns:
     A [batch, num_nodes, num_nodes, depth] vector of tensors
   """
@@ -174,25 +169,16 @@ def make_edge_vectors(adjacency_matrix,
             att_adj_vectors_shape,
             initializer=tf.random_normal_initializer(0, depth**-0.5)) *
         (depth**0.5))
-    # Avoiding gathers so that it works on TPUs
-    # adjacency_matrix_one_hot has shape
-    # [batch, num_nodes, num_nodes, num_edge_types]
-    adjacency_matrix_one_hot = tf.one_hot(adjacency_matrix, num_edge_types)
 
     att_adj_vectors = tf.matmul(
-        tf.reshape(tf.to_float(adjacency_matrix_one_hot), [-1, num_edge_types]),
+        tf.reshape(tf.to_float(adjacency_matrix), [-1, num_edge_types]),
         adj_vectors)
     # Reshape to be [batch, num_nodes, num_nodes, depth].
     att_adj_vectors = tf.reshape(att_adj_vectors, [
         adjacency_matrix_shape[0], adjacency_matrix_shape[1],
         adjacency_matrix_shape[2], depth
     ])
-    if ignore_zero:
-      # Set vectors for edge type 0 to be all zeros.
-      mask = tf.not_equal(tf.expand_dims(adjacency_matrix, axis=-1), 0)
-      return att_adj_vectors * tf.to_float(mask)
-    else:
-      return att_adj_vectors
+    return att_adj_vectors
 
 
 def graph_attention(q,
@@ -206,8 +192,7 @@ def graph_attention(q,
                     save_weights_to=None,
                     dropout_broadcast_dims=None,
                     adjacency_matrix=None,
-                    num_edge_types=5,
-                    ignore_zero=True):
+                    num_edge_types=5):
   """graph attention.
 
   Args:
@@ -229,7 +214,6 @@ def graph_attention(q,
     adjacency_matrix: optional matrix of [batch, length, length] ids indicating
       edge type
     num_edge_types: an int indicating number of edge types
-    ignore_zero: A flag that says that edge type 0 should be ignored.
   Returns:
     A Tensor of shape [batch, length, depth(q)]
   """
@@ -243,13 +227,13 @@ def graph_attention(q,
           adjacency_matrix,
           num_edge_types,
           key_head_depth,
-          ignore_zero=ignore_zero,
           name=name)
       # transposing q to be [batch, length_q, heads, depth_k]
       # to allow for matmul with [batch, length_q, length_q, depth_k]
       q_t = tf.transpose(q, [0, 2, 1, 3])
       adj_logits = tf.matmul(q_t, adjacency_vectors, transpose_b=True)
       logits += tf.transpose(adj_logits, [0, 2, 1, 3])
+      # [batch, depth, num_nodes, num_nodes]
     if bias is not None:
       logits += bias
     weights = tf.nn.softmax(logits, name="attention_weights")
@@ -266,7 +250,6 @@ def graph_attention(q,
 def _compute_edge_transforms(node_states,
                              depth,
                              num_transforms,
-                             ignore_zero=True,
                              name="transform"):
   """Helper function that computes transformation for keys and values.
 
@@ -289,7 +272,6 @@ def _compute_edge_transforms(node_states,
     node_states: A tensor of shape [B, L, D]
     depth: An integer (K or V)
     num_transforms: An integer (T),
-    ignore_zero: A boolean to ignore 0 edge
     name: A name for the function
 
   Returns:
@@ -297,10 +279,9 @@ def _compute_edge_transforms(node_states,
       (shape [B, N*T, K or V])
   """
   node_shapes = common_layers.shape_list(node_states)
-  nonignored_transforms = num_transforms - int(ignore_zero)
   x = common_layers.dense(
       node_states,
-      depth * nonignored_transforms,
+      depth * num_transforms,
       use_bias=False,
       name=name)
 
@@ -311,16 +292,7 @@ def _compute_edge_transforms(node_states,
   # K*T (in k) and V*T (in v) into two-dimensional matrices with shape [K, T]
   # (in k) and [V, T] in v.
   #
-  # This reshape is only necessary when ignore_zero is True (for the padding
-  # step that follows).
-  x = tf.reshape(x, [batch, length, nonignored_transforms, depth])
-
-  # If we previously ignored edge type 0, then we need to pad the keys and
-  # values to take this additional edge type into account. To do so, we
-  # pad the third dimension of k and v (which has size T-1 if ignore_zero is
-  # True) to size T with zeroes.
-  if ignore_zero:
-    x = tf.pad(x, [[0, 0], [0, 0], [1, 0], [0, 0]])
+  x = tf.reshape(x, [batch, length, num_transforms, depth])
 
   # Flatten out the fourth dimension.
   x = tf.reshape(x, [batch, length * num_transforms, depth])
@@ -331,8 +303,7 @@ def _compute_edge_transforms(node_states,
 def compute_mpnn_qkv(node_states,
                      total_key_depth,
                      total_value_depth,
-                     num_transforms,
-                     ignore_zero=True):
+                     num_transforms):
   """Computes query, key and value for edge matrices.
 
   Let B be the number of batches.
@@ -359,9 +330,6 @@ def compute_mpnn_qkv(node_states,
     total_value_depth: an integer (V).
     num_transforms: a integer specifying number of transforms (T). This is
       typically the number of edge types.
-    ignore_zero: If true, then edge type 0 will not be considered. Equivalent
-      to having a linear transformation of all 0's for edge type 0. All queries,
-      keys, and values for edge type 0 will be all 0's.
   Returns:
     q: The attention queries for each destination node (shape [B, N, K]).
     k: The attention keys for each node and edge type (shape [B, N*T, K]).
@@ -386,12 +354,10 @@ def compute_mpnn_qkv(node_states,
   k = _compute_edge_transforms(node_states,
                                total_key_depth,
                                num_transforms,
-                               ignore_zero=ignore_zero,
                                name="k_mpnn")
   v = _compute_edge_transforms(node_states,
                                total_value_depth,
                                num_transforms,
-                               ignore_zero=ignore_zero,
                                name="v_mpnn")
 
   return q, k, v
@@ -406,7 +372,6 @@ def multihead_mpnn_attention(node_states,
                              num_edge_types=5,
                              num_transforms=None,
                              use_weighted_sum=False,
-                             ignore_zero=True,
                              name="mpnn_attention"):
   """Multihead scaled-dot-product attention with input/output transformations.
 
@@ -429,7 +394,7 @@ def multihead_mpnn_attention(node_states,
     total_value_depth: An integer (V).
     output_depth: An integer (O).
     num_heads: An integer (H).
-    adjacency_matrix: An Tensor of ints with shape [B, N, N]. If there is an
+    adjacency_matrix: An Tensor of ints with shape [B, T, N, N]. If there is an
       edge from node j to node i in batch b, then adjacency_matrix[b, i, j]
       contains the type of that edge as an integer. Otherwise, it contains 0.
     num_edge_types: An integer indicating number of edge types.
@@ -437,7 +402,6 @@ def multihead_mpnn_attention(node_states,
       then num_transforms will be equal to num_edge_types.
     use_weighted_sum: If False, will only use a single transform per edge type.
       Otherwise, use a learned weighted sum of transforms per edge type.
-    ignore_zero: A flag that says that edge type 0 should be ignored.
     name: A string.
 
   Returns:
@@ -465,8 +429,7 @@ def multihead_mpnn_attention(node_states,
         node_states,
         total_key_depth,
         total_value_depth,
-        num_transforms,
-        ignore_zero=ignore_zero)
+        num_transforms)
 
     q_shape = tf.shape(q)  # As above, q_shape is [B, N, K].
 
@@ -545,7 +508,6 @@ def dot_product_mpnn_attention(q,
                                num_edge_types,
                                num_transforms=None,
                                use_weighted_sum=False,
-                               ignore_zero=True,
                                name=None):
   """Dot product attention with edge vectors.
 
@@ -559,14 +521,15 @@ def dot_product_mpnn_attention(q,
     q: The query Tensor of shape [B, N, K].
     k: The key Tensor of shape [B, T, N, K].
     v: The value Tensor of shape [B, T, N, V].
-    adjacency_matrix: A Tensor of shape [B, N, N]. An entry at indices b, i, j
-     is the integer edge type of the edge from node j to node i in batch b.
+    adjacency_matrix: A Tensor of shape [B, N, N, T]. An entry at
+      indices b, i, j, k is the indicator of the edge
+      from node j to node i in batch b. A standard adjacency matrix will only
+      have one edge type while a mutigraph will have multiple edge types.
     num_edge_types: An integer specifying number of edge types.
     num_transforms: An integer indicating number of transforms (T). If None,
       then num_transforms will be equal to num_edge_types.
     use_weighted_sum: If False, will only use a single transform per edge type.
       Otherwise, use a learned weighted sum of transforms per edge type.
-    ignore_zero: A flag that says that edge type 0 should be ignored.
     name: A string.
 
   Returns:
@@ -618,8 +581,7 @@ def dot_product_mpnn_attention(q,
       edge_vectors = make_edge_vectors(
           adjacency_matrix,
           num_edge_types,
-          num_transforms,
-          ignore_zero=ignore_zero)
+          num_transforms)
     else:
       # Generate one-hot vectors based on edge types.
       # If there is an edge from node j to node i of type t, then index t of the
@@ -640,17 +602,16 @@ def dot_product_mpnn_attention(q,
     # Since there can only be one edge from node A to node B, we can collapse
     # the T different adjacency matrices containing key-query pairs into one
     # adjacency matrix. logits is [B, N, N].
+    # TODO(dbieber): Use a reshape instead of reduce sum to attend over all
+    # edges instead of over all neighboring nodes to handle the multigraph case.
     logits = tf.reduce_sum(all_edge_logits, axis=1)
 
-    # If we do not have any special treatment for edge type 0, add a large,
-    # negative bias to each location without an edge so that the softmax of
-    # entries with the value 0 become a small negative number instead.
-    #
-    # TODO(avaswani): Better explanation of the rationale behind ignore_zero
-    # here and throughout.
+    # For pairs of nodes with no edges between them, add a large negative bias
+    # to each location without an edge so that the softmax of entries with the
+    # value 0 become a small negative number instead.
     bias = 0
-    if ignore_zero:
-      bias = tf.to_float(tf.equal(adjacency_matrix, 0)) * -1e9
+    bias = tf.to_float(tf.equal(
+        tf.reduce_sum(adjacency_matrix, axis=-1), 0)) * -1e9
     logits += bias
 
     # Turn the raw key-query products into a probability distribution (or,
@@ -681,22 +642,23 @@ def ggnn_fast_dense(node_states,
                     adjacency_matrix,
                     num_edge_types,
                     total_value_depth,
-                    ignore_zero=True,
                     name=None):
   """ggnn version of the MPNN from Gilmer et al.
 
   Let B be the number of batches.
   Let D be the size of the node hidden states.
   Let K be the size of the attention keys/queries.
-  Let V be the size of the output of the ggnn
+  Let V be the size of the output of the ggnn.
+  Let T be the number of transforms / edge types.
 
   Args:
     node_states: The value Tensor of shape [B, T, N, D].
-    adjacency_matrix: A Tensor of shape [B, N, N]. An entry at indices b, i, j
-     is the integer edge type of the edge from node j to node i in batch b.
+    adjacency_matrix: A Tensor of shape [B, N, N, T]. An entry at
+      indices b, i, j, k is the indicator of the edge from node j to node i in
+      batch b. A standard adjacency matrix will only have values of one, while a
+      mutigraph may have larger integer values.
     num_edge_types: An integer specifying number of edge types.
     total_value_depth: An integer (V)
-    ignore_zero: A boolean to ignore edge type 0.
     name: A string.
 
   Returns:
@@ -718,20 +680,14 @@ def ggnn_fast_dense(node_states,
     v = _compute_edge_transforms(node_states,
                                  total_value_depth,
                                  num_edge_types,
-                                 ignore_zero=ignore_zero,
                                  name="v_mpnn")
     v = tf.reshape(v, [nodes_shape[0], nodes_shape[1], num_edge_types,
                        total_value_depth
                       ])  # Shape [B, N, T, V].
     v = tf.transpose(v, [0, 2, 1, 3])  # Shape [B, T, N, V].
 
-    # Generate one-hot vectors based on edge types.
-    # If there is an edge from node j to node i of type t, then index t of the
-    # last dimension is 1 for entry (i, j) of the second and third dimensions.
-    edge_vectors = tf.one_hot(adjacency_matrix, num_edge_types)
-
     # Rearranging the dimensions to match the shape of all_edge_logits.
-    edge_vectors = tf.transpose(edge_vectors, [0, 3, 1, 2])
+    edge_vectors = tf.transpose(adjacency_matrix, [0, 3, 1, 2])
     output = compute_values(edge_vectors, v)
     return output
 
@@ -750,7 +706,8 @@ def compute_values(edge_compatibility, v):
   # Computes the incoming value vectors for each node by weighting them
   # according to the attention weights. These values are still segregated by
   # edge type.
-  all_edge_values = tf.matmul(edge_compatibility, v)  # Shape = [B, T, N, V].
+  # Shape = [B, T, N, V].
+  all_edge_values = tf.matmul(tf.to_float(edge_compatibility), v)
 
   # Combines the weighted value vectors together across edge types into a
   # single N x V matrix for each batch.

From d0f6fe98b3f8b82f8cc80248ce684ad375875c4d Mon Sep 17 00:00:00 2001
From: Xu Song <eson.org@gmail.com>
Date: Mon, 6 Aug 2018 11:33:13 +0800
Subject: [PATCH 0478/2720] remove duplicate lines (#972)

---
 tensor2tensor/data_generators/translate_enzh.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index 68ba28daa..500614b9f 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -33,9 +33,6 @@
 # End-of-sentence marker.
 EOS = text_encoder.EOS_ID
 
-# End-of-sentence marker.
-EOS = text_encoder.EOS_ID
-
 # This is far from being the real WMT18 task - only toyset here
 # you need to register to get UN data and CWT data. Also, by convention,
 # this is EN to ZH - use translate_enzh_wmt8k_rev for ZH to EN task

From a06e8b6492f484c06334bc56926dce07dbfb7a18 Mon Sep 17 00:00:00 2001
From: cclauss <cclauss@bluewin.ch>
Date: Mon, 6 Aug 2018 19:49:18 +0200
Subject: [PATCH 0479/2720] from six.moves import range for Python 3 (#971)

__xrange()__ was removed in Python 3 in favor of __range()__ so convert all instances of __xrange()__ to __six.moves.range()__ for compatibility across Python 2 and Python 3.
---
 tensor2tensor/models/research/universal_transformer_util.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 231cf9a60..fb358f2e8 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -48,6 +48,8 @@
 import copy
 import functools
 
+from six.moves import range  # pylint: disable=redefined-builtin
+
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
@@ -227,7 +229,7 @@ def add_vanilla_transformer_layer(x, num_layers):
       # and add position timing signal at the beginning of each step, so for
       # the vanilla transformer, we need to add timing signal here.
       x = common_attention.add_timing_signal_1d(x)
-    for layer in xrange(num_layers):
+    for layer in range(num_layers):
       with tf.variable_scope("layer_%d" % layer):
         x = ffn_unit(attention_unit(x))
     return x

From efe0305ac0934a9fce88921bb83d2d813aa9bb9e Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sun, 5 Aug 2018 20:39:29 -0700
Subject: [PATCH 0480/2720] internal merge of PR #972

PiperOrigin-RevId: 207486471
---
 tensor2tensor/models/research/universal_transformer_util.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index fb358f2e8..231cf9a60 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -48,8 +48,6 @@
 import copy
 import functools
 
-from six.moves import range  # pylint: disable=redefined-builtin
-
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
@@ -229,7 +227,7 @@ def add_vanilla_transformer_layer(x, num_layers):
       # and add position timing signal at the beginning of each step, so for
       # the vanilla transformer, we need to add timing signal here.
       x = common_attention.add_timing_signal_1d(x)
-    for layer in range(num_layers):
+    for layer in xrange(num_layers):
       with tf.variable_scope("layer_%d" % layer):
         x = ffn_unit(attention_unit(x))
     return x

From 3508d42f6a24c87da1daa5090762eb2db6557314 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 6 Aug 2018 11:21:59 -0700
Subject: [PATCH 0481/2720] internal merge of PR #971

PiperOrigin-RevId: 207577076
---
 tensor2tensor/models/research/universal_transformer_util.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 231cf9a60..fb358f2e8 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -48,6 +48,8 @@
 import copy
 import functools
 
+from six.moves import range  # pylint: disable=redefined-builtin
+
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
@@ -227,7 +229,7 @@ def add_vanilla_transformer_layer(x, num_layers):
       # and add position timing signal at the beginning of each step, so for
       # the vanilla transformer, we need to add timing signal here.
       x = common_attention.add_timing_signal_1d(x)
-    for layer in xrange(num_layers):
+    for layer in range(num_layers):
       with tf.variable_scope("layer_%d" % layer):
         x = ffn_unit(attention_unit(x))
     return x

From c88edeae7c5578b487f234ca8842d3580a082d99 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 6 Aug 2018 15:00:10 -0700
Subject: [PATCH 0482/2720] adding more info into TB.

PiperOrigin-RevId: 207614353
---
 tensor2tensor/models/research/next_frame_sv2p.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 045fd7c7a..910993303 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -97,6 +97,8 @@ def anneal_loss(step_num):
           default=lambda: anneal_loss(step_num))
     else:
       raise ValueError("Unknown beta schedule.")
+
+    tf.summary.scalar("beta", beta)
     return beta
 
   def get_scheduled_sample_func(self, batch_size):
@@ -589,6 +591,12 @@ def body(self, features):
         rewards=all_rewards,
     )
 
+    tf.summary.histogram("input_action", tf.argmax(input_actions, axis=3))
+    tf.summary.histogram("target_action", tf.argmax(target_actions, axis=3))
+    tf.summary.histogram("input_reward", tf.argmax(input_rewards, axis=3))
+    tf.summary.histogram("target_reward", tf.argmax(target_rewards, axis=3))
+    tf.summary.histogram("gen_rewards", tf.argmax(gen_rewards, axis=3))
+
     beta = self.get_beta()
     extra_loss = self.get_extra_loss(
         latent_means=latent_means,

From 857c61e817d93fc29e89dd62a13089bcdca87ced Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 6 Aug 2018 17:26:09 -0700
Subject: [PATCH 0483/2720] revert back summaries as they only work with
 one_hot modalities.

PiperOrigin-RevId: 207637270
---
 tensor2tensor/models/research/next_frame_sv2p.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 910993303..c3cefa93b 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -591,12 +591,6 @@ def body(self, features):
         rewards=all_rewards,
     )
 
-    tf.summary.histogram("input_action", tf.argmax(input_actions, axis=3))
-    tf.summary.histogram("target_action", tf.argmax(target_actions, axis=3))
-    tf.summary.histogram("input_reward", tf.argmax(input_rewards, axis=3))
-    tf.summary.histogram("target_reward", tf.argmax(target_rewards, axis=3))
-    tf.summary.histogram("gen_rewards", tf.argmax(gen_rewards, axis=3))
-
     beta = self.get_beta()
     extra_loss = self.get_extra_loss(
         latent_means=latent_means,

From df95a6fe9f5b80e24689616236d78fec84d54311 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 6 Aug 2018 17:53:39 -0700
Subject: [PATCH 0484/2720] rsqrt_normalized_decay

PiperOrigin-RevId: 207640487
---
 tensor2tensor/utils/learning_rate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index 65d7ba872..e9b65b83a 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -34,7 +34,7 @@ def learning_rate_factor(name, step_num, hparams):
   elif name == "rsqrt_decay":
     return tf.rsqrt(tf.maximum(step_num, hparams.learning_rate_warmup_steps))
   elif name == "rsqrt_normalized_decay":
-    scale = tf.sqrt(hparams.learning_rate_warmup_steps)
+    scale = tf.sqrt(tf.to_float(hparams.learning_rate_warmup_steps))
     return scale * tf.rsqrt(tf.maximum(
         step_num, hparams.learning_rate_warmup_steps))
   elif name == "exp_decay":

From 1d5814b74856f6f783bcdcb46eaa97cd36d51d16 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Mon, 6 Aug 2018 22:40:19 -0700
Subject: [PATCH 0485/2720] Move mesh-tensorflow into open-source directory.

PiperOrigin-RevId: 207663652
---
 .../mesh_tensorflow/mesh_tensorflow.py        | 3234 +++++++++++++++++
 tensor2tensor/mesh_tensorflow/mnist.py        |  217 ++
 .../mesh_tensorflow/mnist_dataset.py          |  131 +
 .../mesh_tensorflow/mtf_beam_search.py        |  573 +++
 .../mesh_tensorflow/mtf_image_transformer.py  |  271 ++
 .../mtf_image_transformer_test.py             |   95 +
 tensor2tensor/mesh_tensorflow/mtf_layers.py   |  851 +++++
 .../mesh_tensorflow/mtf_layers_test.py        |  292 ++
 tensor2tensor/mesh_tensorflow/mtf_model.py    |  287 ++
 tensor2tensor/mesh_tensorflow/mtf_optimize.py |  269 ++
 tensor2tensor/mesh_tensorflow/mtf_toy.py      |  178 +
 .../mesh_tensorflow/mtf_toy_model_tpu.py      |  219 ++
 .../mesh_tensorflow/mtf_transformer.py        |  895 +++++
 .../mesh_tensorflow/mtf_transformer_compat.py |  926 +++++
 .../mesh_tensorflow/mtf_transformer_test.py   |  155 +
 tensor2tensor/mesh_tensorflow/mtf_utils.py    |   29 +
 .../mesh_tensorflow/placement_mesh_impl.py    |  492 +++
 .../mesh_tensorflow/simd_mesh_impl.py         |  342 ++
 .../mesh_tensorflow/tpu_variables.py          |  199 +
 19 files changed, 9655 insertions(+)
 create mode 100644 tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mnist.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mnist_dataset.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_beam_search.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_layers.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_layers_test.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_model.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_optimize.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_toy.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_transformer.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_transformer_compat.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_transformer_test.py
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_utils.py
 create mode 100644 tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
 create mode 100644 tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
 create mode 100644 tensor2tensor/mesh_tensorflow/tpu_variables.py

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
new file mode 100644
index 000000000..f83e0d28b
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -0,0 +1,3234 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Mesh-TensorFlow."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+from functools import reduce  # pylint: disable=redefined-builtin; for py3
+from operator import mul
+import re
+import google3
+from past.builtins import xrange
+import six
+
+from tensor2tensor.mesh_tensorflow import mtf_utils
+import tensorflow as tf
+
+
+Dimension = collections.namedtuple("Dimension", ["name", "size"])
+
+
+def convert_to_dimension(d):
+  """Convert something to a Dimension.
+
+  Args:
+    d: either a Dimension or a pair of (string, int) or None
+  Returns:
+    a Dimension, or None
+  Raises:
+    ValueError: if d cannot be converted to a Dimension
+  """
+  if d is None:
+    return None
+  if isinstance(d, Dimension):
+    return d
+  tf.logging.info("d = %s" % d)
+  name, size = d
+  if isinstance(name, str) and isinstance(size, int):
+    return Dimension(name, size)
+  else:
+    raise ValueError("could not convert %s to Dimension" % d)
+
+
+class TensorShape(object):
+  """Shape of a Tensor."""
+
+  def __init__(self, dims):
+    self._dims = tuple(dims)
+    # verify no repeated dims
+    if len(set([d.name for d in dims])) != len(dims):
+      raise ValueError("Shape must not have repeated dimensions %s" % dims)
+
+  @property
+  def dims(self):
+    return list(self._dims)
+
+  @property
+  def ndims(self):
+    return len(self._dims)
+
+  def __repr__(self):
+    return self.to_string
+
+  def __eq__(self, other):
+    return self.dims == other.dims
+
+  def __ne__(self, other):
+    return self.dims != other.dims
+
+  def __add__(self, other):
+    if isinstance(other, TensorShape):
+      other = other.dims
+    if isinstance(other, Dimension):
+      other = [other]
+    return TensorShape(self.dims + other)
+
+  def __sub__(self, other):
+    if other is None:
+      return self
+    if isinstance(other, TensorShape):
+      other = other.dims
+    if isinstance(other, Dimension):
+      other = [other]
+    return TensorShape([d for d in self.dims if d not in other])
+
+  def __len__(self):
+    return len(self._dims)
+
+  def __getitem__(self, key):
+    return self._dims[key]
+
+  def __iter__(self):
+    return iter(self._dims)
+
+  @property
+  def to_tf_shape(self):
+    return [d.size for d in self.dims]
+
+  @property
+  def size(self):
+    return list_product(self.to_tf_shape)
+
+  @property
+  def to_string(self):
+    return "TensorShape[%s]" % ", ".join(
+        ["%s=%d" % (d.name, d.size) for d in self.dims])
+
+  @property
+  def cumprod(self):
+    """cumulative product (exclusive) of dimension sizes."""
+    return _cumprod(self.to_tf_shape)[::-1]
+
+  def cumprod_to_tensor_axis(self, cumprod):
+    """Tensor axis i such that self.cumprod[i] == cumprod, or None."""
+    try:
+      return self.cumprod.index(cumprod)
+    except ValueError:
+      return None
+
+  @property
+  def dimension_names(self):
+    return [d.name for d in self.dims]
+
+  def rename_dimension(self, old_name, new_name):
+    """Returns a copy where one dimension is renamed."""
+    if old_name not in self.dimension_names:
+      raise ValueError("Shape %s does not have dimension named %s"
+                       % (self, old_name))
+    return TensorShape(
+        [Dimension(new_name, d.size) if d.name == old_name else d
+         for d in self.dims])
+
+  def resize_dimension(self, name, new_size):
+    """Returns a copy where one dimension has a different size."""
+    if name not in self.dimension_names:
+      raise ValueError("Shape %s does not have dimension named %s"
+                       % (self, name))
+    return TensorShape(
+        [Dimension(name, new_size) if d.name == name else d
+         for d in self.dims])
+
+
+def convert_to_tensor_shape(x):
+  if x is None:
+    return None
+  if isinstance(x, TensorShape):
+    return x
+  return TensorShape([convert_to_dimension(d) for d in x])
+
+
+class TensorLayout(object):
+  """Mapping from tensor dimension to mesh dimension.
+
+  Represented as a list of optional integers with length tensor.ndims.
+  Each item is either a unique integer inicating the mesh dimension over
+  which that tensor dimension is split, or None, indicating that this
+  tensor dimension is not split.
+  """
+
+  def __init__(self, tensor_axis_to_mesh_axis):
+    self._tensor_axis_to_mesh_axis = tensor_axis_to_mesh_axis
+
+  def __eq__(self, other):
+    return self.tensor_axis_to_mesh_axis == other.tensor_axis_to_mesh_axis
+
+  def __ne__(self, other):
+    return self.tensor_axis_to_mesh_axis != other.tensor_axis_to_mesh_axis
+
+  def __repr__(self):
+    return "TensorLayout%s" % self.tensor_axis_to_mesh_axis
+
+  @property
+  def tensor_axis_to_mesh_axis(self):
+    return self._tensor_axis_to_mesh_axis
+
+  @property
+  def is_fully_replicated(self):
+    return (self.tensor_axis_to_mesh_axis ==
+            [None] * len(self.tensor_axis_to_mesh_axis))
+
+
+class Graph(object):
+  """Distributed-TF graph."""
+
+  def __init__(self):
+    self._operations = []
+    self._tensors = []
+    self._trainable_variables = []
+    self._all_variables = []
+
+  def __repr__(self):
+    return self.to_string
+
+  @property
+  def operations(self):
+    return self._operations
+
+  @property
+  def tensors(self):
+    return self._tensors
+
+  @property
+  def trainable_variables(self):
+    return self._trainable_variables
+
+  @property
+  def all_variables(self):
+    return self._all_variables
+
+  @property
+  def to_string(self):
+    return "\n".join([op.to_string for op in self.operations])
+
+
+class Lowering(object):
+  """Lowering of a Graph from mesh-tensorflow to tensorflow."""
+
+  def __init__(self, graph, mesh_to_impl):
+    """Create a Lowering of a graph.
+
+    layout is a dictionary whose keys are the meshes in the graph
+    and whose values are themselves dictionaries mapping tensor-dimension
+    names to mesh dimensions (integers).
+
+    Args:
+      graph: a Graph
+      mesh_to_impl: {Mesh: MeshImpl}
+    """
+    # tf.logging.info("LOWERING GRAPH:\n%s" % graph.to_string)
+    self.mesh_to_impl = mesh_to_impl   # {Mesh: MeshImpl}
+    self.graph = graph
+    self._counters = []
+    self.tensors = {}                  # {Tensor: Mesh.LaidOutTensor}
+    self.operations = {}               # {Operation: tf.Operation}
+    self.variables = {}                # {Variable: LaidOutVariable}
+    for op in graph.operations:
+      # tf.logging.info("Lowering operation %s" % op.to_string)
+      with tf.name_scope(op.name):
+        op.lower(self)
+      for out in op.outputs:
+        self.add_counter(
+            "output/%s" % type(op).__name__, self.laid_out_size(out))
+        self.add_counter("output_unique/%s" % type(op).__name__, out.size)
+    log_variable_sizes(
+        graph.trainable_variables, "Trainable Variables", verbose=True)
+    tf.logging.info("Counters:\n" + pretty_print_counters(self._counters))
+
+  def mesh_impl(self, m):
+    if not isinstance(m, Mesh):
+      m = m.mesh
+    return self.mesh_to_impl[m]
+
+  def outfeed(self, x):
+    """Turn a Tensor into a tf.Tensor.
+
+    Args:
+      x: a Tensor
+    Returns:
+      a tf.Tensor
+    """
+    mesh_impl = self.mesh_impl(x)
+    return mesh_impl.outfeed(x, self.tensors[x].to_laid_out_tensor())
+
+  def lowered_operation(self, op):
+    return self.operations[op]
+
+  def copy_masters_to_slices(self):
+    return tf.group(
+        [v.copy_master_to_slices for _, v in six.iteritems(self.variables)])
+
+  def copy_slices_to_masters(self):
+    return tf.group(
+        [v.copy_slices_to_master for _, v in six.iteritems(self.variables)])
+
+  def tensor_layout(self, t):
+    return self.mesh_impl(t).tensor_layout(t)
+
+  def add_counter(self, key, value):
+    assert isinstance(value, int)
+    self._counters.append((key, value))
+
+  @property
+  def counters(self):
+    return self._counters
+
+  def laid_out_size(self, tensor):
+    """Total size of all slices.
+
+    Args:
+      tensor: a Tensor
+
+    Returns:
+      an integer
+    """
+    return self.mesh_impl(tensor).laid_out_size(tensor.shape)
+
+
+class Mesh(object):
+  """A placeholder with no functionality.
+
+  A Graph is built with each tensor assigned to a mesh.  The mesh does not
+  know its shape or its implementation.
+
+  A Lowering asssigns a MeshImpl to each mesh.
+  """
+
+  def __init__(self, graph, name):
+    self._graph = graph
+    self._name = name
+
+  @property
+  def graph(self):
+    return self._graph
+
+
+class MeshImpl(object):
+  """Implementation of a mesh.
+
+  Knows its shape, its underlying devices, and its layout
+  (mapping from TensorDim to mesh-dimension).
+
+  Subclasses will include PlacementMeshImpl and SimdMeshImpl
+  """
+
+  def __init__(self, shape, layout):
+    """Create a mesh.
+
+    Args:
+      shape: a list of ints
+      layout: dict from string to int
+    """
+    self._shape = shape
+    self._layout = layout
+
+  @property
+  def shape(self):
+    return self._shape
+
+  @property
+  def ndims(self):
+    return len(self._shape)
+
+  @property
+  def layout(self):
+    return self._layout
+
+  @property
+  def size(self):
+    return list_product(self._shape)
+
+  def tensor_dimension_to_mesh_axis(self, tensor_dimension):
+    """Mesh axis associated with tensor dimension (or None).
+
+    Args:
+      tensor_dimension: a Dimension
+    Returns:
+      an integer or None
+    """
+    return self.layout.get(tensor_dimension.name)
+
+  def tensor_layout(self, arg):
+    """Compute TensorLayout given a mesh and a TensorShape.
+
+    Args:
+      arg: a Tensor or TensorShape
+    Returns:
+      a TensorLayout
+    """
+    if isinstance(arg, Tensor):
+      arg = arg.shape
+    return TensorLayout(
+        [self.tensor_dimension_to_mesh_axis(d) for d in arg.dims])
+
+  def mesh_axis_to_tensor_axis(self, tensor_shape):
+    """Reverse-mapping of a tensor layout.
+
+    Args:
+      tensor_shape: a TensorShape
+    Returns:
+      a list of length self.ndims, where each element is either an integer
+        index of a tensor axis, or None
+    """
+    layout = self.tensor_layout(tensor_shape)
+    return [layout.tensor_axis_to_mesh_axis.index(mesh_axis)
+            if mesh_axis in layout.tensor_axis_to_mesh_axis else None
+            for mesh_axis in xrange(self.ndims)]
+
+  def mesh_axis_to_cumprod(self, tensor_shape):
+    """For each mesh axis, give the product of previous tensor axes.
+
+    Args:
+      tensor_shape: a TensorShape
+    Returns:
+      a list with length self.ndims where each element is an integer or None.
+    """
+    ma2ta = self.mesh_axis_to_tensor_axis(tensor_shape)
+    ta2cumprod = tensor_shape.cumprod
+    return [None if ta is None else ta2cumprod[ta] for ta in ma2ta]
+
+  def slice_shape(self, tensor_shape):
+    """Shape of each slice of the tensor.
+
+    Args:
+      tensor_shape: a TensorShape
+    Returns:
+      a list of integers with length tensor_shape.ndims
+    Raises:
+      ValueError: if a tensor dimension is not divisible by the corresponding
+        mesh dimension.
+    """
+    tensor_layout = self.tensor_layout(tensor_shape)
+    ret = []
+    for dim_size, mesh_dim in zip(
+        tensor_shape.to_tf_shape, tensor_layout.tensor_axis_to_mesh_axis):
+      if mesh_dim is None:
+        ret.append(dim_size)
+      else:
+        if dim_size % self.shape[mesh_dim] != 0:
+          raise ValueError(
+              "Tensor dimension size not divisible by mesh dimension size:"
+              " tensor_shape=%s tensor_layout=%s"
+              % (tensor_shape, tensor_layout))
+        ret.append(dim_size // self.shape[mesh_dim])
+    return ret
+
+  def slice_begin(self, tensor_shape, pnum):
+    """Begin position for the tensor slice for the given processor.
+
+    Args:
+      tensor_shape: a TensorShape
+      pnum: an integer <= self.size
+    Returns:
+      a list of integers with length tensor_shape.ndims
+    """
+    tensor_layout = self.tensor_layout(tensor_shape)
+    coordinates = pnum_to_processor_coordinates(self.shape, pnum)
+    ret = []
+    for dim_size, mesh_axis in zip(
+        tensor_shape.to_tf_shape, tensor_layout.tensor_axis_to_mesh_axis):
+      if mesh_axis is None:
+        ret.append(0)
+      else:
+        ret.append(dim_size // self.shape[mesh_axis] * coordinates[mesh_axis])
+    return ret
+
+  def laid_out_size(self, tensor_shape):
+    """Total size of all slices.
+
+    Args:
+      tensor_shape: a TensorShape
+
+    Returns:
+      an integer
+    """
+    return list_product(self.slice_shape(tensor_shape)) * self.size
+
+  def slicewise(self, fn, *inputs):
+    """Execute a function in parallel on all slices.
+
+    Args:
+      fn: a function from tf.Tensors to tf.Tensor or a tuple of tf.Tensors.
+      *inputs: a list of inputs.  Each input is either a LaidOutTensor or
+        is convertible to a tf.Tensor.
+    Returns:
+      a LaidOutTensor, or a tuple of LaidOutTensors if fn returns a tuple.
+    """
+    raise NotImplementedError("Slicewise not implemented")
+
+  def Print(self, x, data, message, **kwargs):  # pylint: disable=invalid-name
+    """call tf.Print.
+
+    Args:
+      x: a LaidOutTensor
+      data: a list of LaidOutTensor
+      message: a string
+      **kwargs: keyword arguments to tf.print
+    Returns:
+      a LaidOutTensor
+    """
+    del data, message, kwargs
+    tf.logging.warning("Warning - mtf.Print not implemented for this mesh type")
+    return x
+
+  def allreduce(self, x, mesh_axes, reduction_fn_string):
+    """Grouped allreduce, (summed across the given dimensions).
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axes: a list of integers - the mesh dimensions to be reduced
+      reduction_fn_string: "SUM" or "MAX"
+    Returns:
+      a LaidOutTensor
+    """
+    raise NotImplementedError("Allreduce not implemented")
+
+  def allsplit(self, x, mesh_axis, split_axis):
+    """Inverse of allconcat - split each slice and keep only one piece of it.
+
+    The number of ways to split is the number of processors in the group.
+    The part that is kept corrseponds to the processor's index in the group.
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axis: an integer the mesh axis along which to split
+      split_axis: an integer (the Tensor axis along which to split)
+    Returns:
+      a LaidOutTensor
+    """
+    num_splits = self.shape[mesh_axis]
+    def my_fn(x, coordinate):
+      slice_begin = [
+          dimsize // num_splits * coordinate if i == split_axis
+          else 0 for i, dimsize in enumerate(x.shape.as_list())]
+      slice_size = [
+          dimsize // num_splits if i == split_axis
+          else dimsize for i, dimsize in enumerate(x.shape.as_list())]
+      return tf.slice(x, slice_begin, slice_size)
+    return self.slicewise(my_fn, x, self.laid_out_pcoord(mesh_axis))
+
+  def allconcat(self, x, mesh_axis, concat_axis):
+    """Grouped allconcat (like MPI allgather followed by concat).
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axis: an integer - the mesh axis along which to group
+      concat_axis: an integer (the Tensor axis along which to concatenate)
+    Returns:
+      a LaidOutTensor
+    """
+    raise NotImplementedError("Allconcat not implemented")
+
+  def alltoall(self, x, mesh_axis, split_axis, concat_axis):
+    """Grouped alltoall (like MPI alltoall with splitting and concatenation).
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axis: an integer the mesh axis along which to group
+      split_axis: an integer (the Tensor axis along which to split)
+      concat_axis: an integer (the Tensor axis along which to concatenate)
+    Returns:
+      a LaidOutTensor
+    """
+    raise NotImplementedError("Alltoall not implemented")
+
+  def laid_out_pnum(self):
+    """Returns a LaidOutTensor containing the processor number.
+
+    Returns:
+      a LaidOutTensor where each slice is an integer scalar
+    """
+    raise NotImplementedError("laid_out_pnum not implemented")
+
+  def laid_out_pcoord(self, mesh_axis):
+    """Returns a LaidOutTensor containing the processor coordinate.
+
+    Args:
+      mesh_axis: an integer
+    Returns:
+      a LaidOutTensor where each slice is an integer scalar
+    """
+    divisor = list_product(self.shape[mesh_axis + 1:])
+    modulus = self.shape[mesh_axis]
+    def my_fn(pnum):
+      return (pnum // divisor) % modulus
+    return self.slicewise(my_fn, self.laid_out_pnum())
+
+  def broadcast_impl(self, old_slices, old_shape, new_shape):
+    """Implementation of a broadcast operation.
+
+    Args:
+      old_slices: a LaidOutTensor
+      old_shape: a TensorShape
+      new_shape: a TensorShape
+    Returns:
+      a LaidOutTensor
+    """
+    new_slice_shape = self.slice_shape(new_shape)
+    def tf_fn(x):
+      return (tf.zeros(new_slice_shape, dtype=x.dtype) +
+              _expand_dims(x, old_shape, new_shape))
+    return self.slicewise(tf_fn, old_slices)
+
+  def make_slices(self, tf_tensor, tensor_shape):
+    """Turn a single tf.Tensor into a list of slices, one for each processor.
+
+    Args:
+      tf_tensor: a tf.Tensor
+      tensor_shape: a TensorShape
+
+    Returns:
+      a list of tf.tensor with length self.size
+    """
+    tensor_layout = self.tensor_layout(tensor_shape)
+    slice_shape = self.slice_shape(tensor_shape)
+    def my_fn(pnum):
+      if tensor_layout.is_fully_replicated:
+        return tf_tensor
+      else:
+        slice_begin = self.slice_begin(tensor_shape, pnum)
+        return tf.slice(tf_tensor, slice_begin, slice_shape)
+
+    return parallel([tf_tensor.device] * self.size, my_fn,
+                    list(xrange(self.size)))
+
+  def combine_slices(self, slices, tensor_shape, device=None):
+    """Turn a set of slices into a single tensor.
+
+    Args:
+      slices: a list of tf.Tensor with length self.size
+      tensor_shape: a TensorShape
+      device: an optional device string.
+        if absent, we use the devices of the slices.
+
+    Returns:
+      a tf.Tensor
+    """
+    if tensor_shape.ndims == 0:
+      return slices[0]
+
+    ret = slices[:]
+    for unused_mesh_axis, (mesh_axis_size, tensor_axis) in enumerate(
+        zip(self.shape, self.mesh_axis_to_tensor_axis(tensor_shape))):
+      slice_size = len(ret) // mesh_axis_size
+      if tensor_axis is None:
+        ret = ret[:slice_size]
+      else:
+        if device:
+          devices = [device] * slice_size
+        else:
+          devices = [ret[i].device for i in xrange(slice_size)]
+        concat_inputs = [[ret[i + slice_size * j]
+                          for j in xrange(mesh_axis_size)]
+                         for i in xrange(slice_size)]
+        ret = parallel(
+            devices, tf.concat, concat_inputs,
+            axis=[tensor_axis] * len(devices))
+    assert len(ret) == 1
+    return ret[0]
+
+  def outfeed(self, x, laid_out_x):
+    """Turn a Tensor into a tf.Tensor.
+
+    Args:
+      x: a Tensor
+      laid_out_x: a LaidOutTensor
+    Returns:
+      a tf.Tensor
+    """
+    raise NotImplementedError("Outfeed not implemented")
+
+  def infeed(self, x, tf_x):
+    """Infeed a tf.Tensor, producing a LaidOutTensor.
+
+    Args:
+      x: a Tensor
+      tf_x: a tf.Tensor
+    Returns:
+      a LaidOutTensor
+    """
+    raise NotImplementedError("Infeed not implemented")
+
+  @property
+  def supports_control_dependencies(self):
+    return True
+
+
+class LazyAllreduceSum(object):
+  """Represents a LaidOutTensor with a lazy allreduce.
+
+  The purpose of delaying allreduce is that it saves bandwidth to first add
+  and then allreduce, as opposed to the other way around.
+  """
+
+  def __init__(self,
+               mesh_impl,
+               laid_out_input,
+               mesh_axes,
+               add_counter_fn=None):
+    """Create a LazyAllreduceSum.
+
+    Args:
+      mesh_impl: a mesh_impl
+      laid_out_input: a LaidOutTensor
+      mesh_axes: a list of mesh axes
+      add_counter_fn: a function taking no arguments which calls
+        lowering.add_counter if and when the allreduce executes.
+    Returns:
+      a LazyAllreduceSum
+    """
+    self.mesh_impl = mesh_impl
+    self.laid_out_input = laid_out_input
+    self.mesh_axes = mesh_axes
+    self._add_counter_fn = add_counter_fn
+    self._reduced = None
+
+  def to_laid_out_tensor(self):
+    if not self._reduced:
+      self._reduced = self.mesh_impl.allreduce(
+          self.laid_out_input, self.mesh_axes, "SUM")
+      if self._add_counter_fn:
+        self._add_counter_fn()
+    return self._reduced
+
+  def __add__(self, other):
+    """Add to another LazyAllreduceSum.
+
+    Args:
+      other: a LazyAllreduceSum or a LaidOutTensor
+    Returns:
+      a LazyAllreduceSum or a LaidOutTensor
+    """
+    if (isinstance(other, LazyAllreduceSum) and
+        self.mesh_impl == other.mesh_impl and
+        self.mesh_axes == other.mesh_axes):
+      return LazyAllreduceSum(
+          self.mesh_impl,
+          self.mesh_impl.slicewise(
+              tf.add, self.laid_out_input, other.laid_out_input),
+          self.mesh_axes,
+          add_counter_fn=self._add_counter_fn)
+    else:
+      return self.mesh_impl.slicewise(
+          tf.add, self.to_laid_out_tensor(), other.to_laid_out_tensor())
+
+
+def convert_args_to_laid_out_tensors(xs):
+  """Convert list elements to laid-out-tensors when possible.
+
+  Args:
+    xs: a list
+  Returns:
+    a list
+  """
+  ret = []
+  for x in xs:
+    try:
+      ret.append(x.to_laid_out_tensor())
+    except AttributeError:
+      ret.append(x)
+  return ret
+
+
+class Tensor(object):
+  """A Distributed Tensor."""
+
+  def __init__(self, operation, shape, dtype, name=None):
+    if not isinstance(shape, TensorShape):
+      raise ValueError("shape must be a TensorShape got %s" % shape.to_string)
+    if not isinstance(dtype, tf.DType):
+      raise ValueError("dtype must be a tf.DType got %s" % dtype)
+    self._mesh = operation.mesh
+    self._operation = operation
+    self._shape = shape
+    self._dtype = dtype
+    if name is None:
+      name = self.operation.name
+    self._name = name
+    self._mesh.graph.tensors.append(self)
+
+  @property
+  def shape(self):
+    return self._shape
+
+  @property
+  def size(self):
+    return self.shape.size
+
+  @property
+  def mesh(self):
+    return self._mesh
+
+  @property
+  def graph(self):
+    return self._mesh.graph
+
+  @property
+  def operation(self):
+    return self._operation
+
+  @property
+  def dtype(self):
+    return self._dtype
+
+  @property
+  def name(self):
+    return self._name
+
+  def __repr__(self):
+    return self.to_string
+
+  def __add__(self, other):
+    return add(self, other)
+
+  def __radd__(self, other):
+    return add(self, other)
+
+  def __sub__(self, other):
+    return sub(self, other)
+
+  def __rsub__(self, other):
+    return sub(other, self)
+
+  def __mul__(self, other):
+    return multiply(self, other)
+
+  def __rmul__(self, other):
+    return multiply(self, other)
+
+  def __neg__(self):
+    return negative(self)
+
+  def __truediv__(self, other):
+    return divide(self, other)
+
+  def __rtruediv__(self, other):
+    return divide(other, self)
+
+  def __floordiv__(self, other):
+    return floordiv(self, other)
+
+  def __rfloordiv__(self, other):
+    return floordiv(other, self)
+
+  def __mod__(self, other):
+    return mod(self, other)
+
+  def __rmod__(self, other):
+    return mod(other, self)
+
+  @property
+  def to_string(self):
+    return "Tensor[%s, %s, %s]" % (self.name, self.shape.to_string, self.dtype)
+
+
+class Operation(object):
+  """A Distributed Operation."""
+
+  def __init__(self, inputs, mesh=None, name=None):
+    if mesh is None:
+      if not inputs:
+        raise ValueError("mesh must be specified if no inputs")
+      mesh = inputs[0].mesh
+    self._inputs = inputs
+    self._outputs = []
+    self._mesh = mesh
+    assert name is not None
+    scope_name = tf.get_variable_scope().name
+    if scope_name:
+      name = scope_name + "/" + name
+    self._name = name
+    mesh.graph.operations.append(self)
+
+  @property
+  def graph(self):
+    return self._mesh.graph
+
+  @property
+  def mesh(self):
+    return self._mesh
+
+  @property
+  def name(self):
+    return self._name
+
+  @property
+  def inputs(self):
+    return self._inputs[:]
+
+  @property
+  def outputs(self):
+    return self._outputs[:]
+
+  @property
+  def to_string(self):
+    return "%s[Inputs=(%s) Outputs=(%s)]" % (
+        type(self).__name__,
+        ", ".join([t.to_string for t in self.inputs]),
+        ", ".join([t.to_string for t in self.outputs]))
+
+  @property
+  def has_gradient(self):
+    return (
+        [t for t in self.inputs if t.dtype.is_floating] and
+        [t for t in self.outputs if t.dtype.is_floating])
+
+  def gradient(self, unused_grad_ys):
+    raise NotImplementedError("Gradient not implemented")
+
+  def lower(self, lowering):
+    raise NotImplementedError("Lower not implemented")
+
+
+class SlicewiseOperation(Operation):
+  """Apply any tensorflow function slice-wise.
+
+  Calls the Tensorflow function on each slice of the inputs to produce the
+  corresponding slice of the outputs.  Gradients are computed through
+  tensorflow.
+
+  The user must specify "splittable_dims": a list of Dimensions which can
+  be split while still keeping this computation valid.  For example, for
+  component-wise functions, all the dimensions are splittable, but if the
+  function is a reduction, the reduced dimensions are not splittable.
+  """
+
+  def __init__(self,
+               tf_fn,
+               inputs,
+               output_shape,
+               output_dtype,
+               splittable_dims,
+               grad_function=None,
+               name=None):
+    """Create a SlicewiseOperation.
+
+    grad_function is a python function taking this operation and a gradients
+    Tensor and producing input gradients tensors.
+    e.g.
+    def _square_grad(op, dy):
+      return [dy * op.inputs[0] * 2]
+
+    Args:
+      tf_fn: a function taking n tf.Tensors and returning a tf.Tensor
+      inputs: a list of n Tensors
+      output_shape: a TensorShape
+      output_dtype: a dtype
+      splittable_dims: a list of Dimensions which are ok to split
+      grad_function: an optional python function
+      name: an optional string
+    """
+    super(SlicewiseOperation, self).__init__(inputs, name=name or "slicewise")
+    self._tf_fn = tf_fn
+    self._outputs = [Tensor(self, output_shape, output_dtype)]
+    self._splittable_dims = splittable_dims
+    self._grad_function = grad_function
+
+  def gradient(self, grad_ys):
+    if self._grad_function is not None:
+      return self._grad_function(self, grad_ys[0])
+    return GenericGradOperation(self, grad_ys).outputs
+
+  def lower(self, lowering):
+    # Check that only splittable dims are split
+    for t in self.inputs + self.outputs:
+      layout = lowering.tensor_layout(t)
+      for d, mesh_axis in zip(t.shape.dims, layout.tensor_axis_to_mesh_axis):
+        if (mesh_axis is not None and d not in self._splittable_dims):
+          raise ValueError("dimension %s is not declared as splittable" % d)
+    lowering.tensors[self.outputs[0]] = lowering.mesh_impl(self).slicewise(
+        self._tf_fn, *[lowering.tensors[x] for x in self.inputs])
+
+
+def slicewise(tf_fn,
+              xs,
+              output_shape=None,
+              output_dtype=None,
+              splittable_dims=None,
+              grad_function=None,
+              name=None):
+  """Slice-wise call to any tensorflow function.
+
+  The output shape and dtype default to those of the first input.
+  splittable_dims is a list of Dimensions which can be split while keeping the
+  computation valid.
+
+  Args:
+    tf_fn: a function taking n tf.Tensors and returning a tf.Tensor
+    xs: a list of n Tensors
+    output_shape: a TensorShape
+    output_dtype: a dtype
+    splittable_dims: a list of Dimensions which are ok to split
+    grad_function: an optional gradients function
+    name: an optional string
+
+  Returns:
+    a Tensor
+  """
+  return SlicewiseOperation(
+      tf_fn,
+      xs,
+      convert_to_tensor_shape(output_shape) or xs[0].shape,
+      output_dtype or xs[0].dtype,
+      splittable_dims,
+      grad_function,
+      name=name).outputs[0]
+
+
+def cwise(tf_fn, xs, output_dtype=None, grad_function=None, name=None):
+  """Component-wise operation with no broadcasting.
+
+  Args:
+    tf_fn: a component-wise function taking n tf.Tensor inputs and producing
+      a tf.Tensor output
+    xs: n Tensors
+    output_dtype: an optional dtype
+    grad_function: an optional python function
+    name: an optional string
+
+  Returns:
+    a Tensor
+  """
+  return slicewise(
+      tf_fn, xs, output_dtype=output_dtype, splittable_dims=xs[0].shape.dims,
+      grad_function=grad_function, name=name or "cwise")
+
+
+def square(x, name="square"):
+  return cwise(
+      tf.square, [x], name=name,
+      grad_function=lambda op, dy: [dy * op.inputs[0] * 2])
+
+
+def sqrt(x, name="sqrt"):
+  return cwise(
+      tf.sqrt, [x], name=name,
+      grad_function=lambda op, dy: [dy * 0.5 / op.outputs[0]])
+
+
+def _rsqrt_grad(op, dy):
+  return [dy * -0.5 * op.outputs[0] * op.outputs[0] * op.outputs[0]]
+
+
+def rsqrt(x, name="rsqrt"):
+  return cwise(
+      tf.rsqrt, [x], name=name, grad_function=_rsqrt_grad)
+
+
+def log(x, name="log"):
+  return cwise(
+      tf.log, [x], name=name,
+      grad_function=lambda op, dy: [dy / op.inputs[0]])
+
+
+def exp(x, name="exp"):
+  return cwise(tf.exp, [x], name=name,
+               grad_function=lambda op, dy: [dy * op.outputs[0]])
+
+
+def pow(x, y):  # pylint: disable=redefined-builtin
+  return exp(log(x) * y)
+
+
+def negative(x, name="negative"):
+  return cwise(tf.negative, [x], name=name,
+               grad_function=lambda op, dy: [negative(dy)])
+
+
+def logical_not(x, name="logical_not"):
+  return cwise(tf.logical_not, [x], name=name)
+
+
+def reciprocal(x, name="reciprocal"):
+  return cwise(
+      tf.reciprocal, [x], name=name,
+      grad_function=lambda op, dy: [negative(dy * square(op.outputs[0]))])
+
+
+def _relu_grad(op, dy):
+  return [dy * cast(greater(op.inputs[0], 0), op.inputs[0].dtype)]
+
+
+def relu(x, name="relu"):
+  return cwise(tf.nn.relu, [x], name=name, grad_function=_relu_grad)
+
+
+def cast(x, dtype, name="cast"):
+  if dtype == x.dtype:
+    return x
+  return cwise(
+      lambda x: tf.cast(x, dtype), [x], output_dtype=dtype, name=name,
+      grad_function=lambda op, dy: [cast(dy, op.inputs[0].dtype)])
+
+
+def to_float(x, name="to_float"):
+  return cast(x, tf.float32, name=name)
+
+
+def to_int32(x, name="to_int32"):
+  return cast(x, tf.int32, name=name)
+
+
+class GenericGradOperation(Operation):
+  """Gradients that follow regular TF.
+
+  Calling tf.gradients multiple times seems really slow in python.
+  TODO(noam): can we speed this up using functions or some other method?
+  """
+
+  def __init__(self, forward_op, grad_ys, name=None):
+    # tf.logging.info("forward inp %s, operations %s, grad_ys: %s",
+    #                 forward_op.inputs, forward_op.outputs, grad_ys)
+    super(GenericGradOperation, self).__init__(
+        forward_op.inputs + forward_op.outputs + grad_ys,
+        name=name or "generic_grad")
+    self._grad_ys = grad_ys
+    self._forward_op = forward_op
+    self._outputs = [Tensor(self, x.shape, x.dtype) for x in forward_op.inputs]
+
+  def lower(self, lowering):
+    # lists of lists of tf.Tensor
+    all_ys = transpose_list_of_lists(
+        [lowering.tensors[y].tensor_list for y in self._forward_op.outputs])
+    all_xs = transpose_list_of_lists(
+        [lowering.tensors[x].tensor_list for x in self._forward_op.inputs])
+    all_grad_ys = transpose_list_of_lists(
+        [lowering.tensors[dy].tensor_list for dy in self._grad_ys])
+    all_grad_xs = [tf.gradients(ys=ys, xs=xs, grad_ys=grad_ys) for
+                   ys, xs, grad_ys in zip(all_ys, all_xs, all_grad_ys)]
+    grad_xs = transpose_list_of_lists(all_grad_xs)
+    for out, grad_x in zip(self.outputs, grad_xs):
+      lowering.tensors[out] = (
+          lowering.mesh_impl(self).LaidOutTensor.from_tensor_list(grad_x))
+
+
+class ScalarMultiplyOperation(Operation):
+  """Multiply by a tf Scalar (no backprop to scalar)."""
+
+  def __init__(self, x, scalar, name=None):
+    super(ScalarMultiplyOperation, self).__init__(
+        [x], name=name or "scalar_mul")
+    self._outputs = [Tensor(self, x.shape, x.dtype)]
+    self._scalar = scalar
+
+  def gradient(self, grad_ys):
+    dy = grad_ys[0]
+    return [dy * self._scalar]
+
+  def lower(self, lowering):
+    lowering.tensors[self.outputs[0]] = lowering.mesh_impl(self).slicewise(
+        lambda x: x * self._scalar, lowering.tensors[self.inputs[0]])
+
+
+class ScalarAddOperation(Operation):
+  """Add a tf Scalar (no backprop to scalar)."""
+
+  def __init__(self, x, scalar, name=None):
+    super(ScalarAddOperation, self).__init__([x], name=name or "scalar_add")
+    self._outputs = [Tensor(self, x.shape, x.dtype)]
+    self._scalar = scalar
+
+  def gradient(self, grad_ys):
+    return grad_ys
+
+  def lower(self, lowering):
+    lowering.tensors[self.outputs[0]] = lowering.mesh_impl(self).slicewise(
+        lambda x: x + self._scalar, lowering.tensors[self.inputs[0]])
+
+
+class BinaryOpWithBroadcasting(Operation):
+  """Binary operation with broadcasting."""
+
+  def __init__(self, tf_fn, x1, x2, output_shape, output_dtype, name=None):
+    super(BinaryOpWithBroadcasting, self).__init__(
+        [x1, x2], name=name or "binary_op")
+    assert isinstance(output_dtype, tf.DType)
+    self._outputs = [Tensor(self, output_shape, output_dtype)]
+    self._tf_fn = tf_fn
+
+  def gradient(self, unused_grad_ys):
+    raise ValueError("Gradient not implememnted")
+
+  def lower(self, lowering):
+    x1 = self.inputs[0]
+    x2 = self.inputs[1]
+    output = self.outputs[0]
+    laid_out_x1 = lowering.tensors[x1]
+    laid_out_x2 = lowering.tensors[x2]
+    mesh_impl = lowering.mesh_impl(self)
+    if x1.shape != output.shape:
+      laid_out_x1 = mesh_impl.slicewise(
+          _expand_dims, laid_out_x1, x1.shape, output.shape)
+    if x2.shape != output.shape:
+      laid_out_x2 = mesh_impl.slicewise(
+          _expand_dims, laid_out_x2, x2.shape, output.shape)
+    lowering.tensors[self.outputs[0]] = mesh_impl.slicewise(
+        self._tf_fn, laid_out_x1, laid_out_x2)
+
+
+def binary_arguments_to_tensors(x1, x2):
+  """Convert argument of a binary operation to Tensors.
+
+  Args:
+    x1: a Tensor or something convertible to a tf Scalar
+    x2: a Tensor or something convertible to a tf Scalar
+
+  Returns:
+    new_x1: a Tensor
+    new_x2: a Tensor
+
+  Raises:
+    ValueError: on failure
+  """
+  if not isinstance(x1, Tensor) and not isinstance(x2, Tensor):
+    raise ValueError("at least one of x1 and x2 must be an mtf Tensor")
+  elif isinstance(x1, Tensor) and isinstance(x2, Tensor):
+    return x1, x2
+  elif isinstance(x1, Tensor):
+    return x1, infeed(x1.mesh, tf.convert_to_tensor(x2, dtype=x1.dtype),
+                      TensorShape([]))
+  else:
+    return infeed(x2.mesh, tf.convert_to_tensor(x1, dtype=x2.dtype),
+                  TensorShape([])), x2
+
+
+def binary_op_with_broadcasting(
+    tf_fn, x1, x2, output_shape=None, output_dtype=None):
+  x1, x2 = binary_arguments_to_tensors(x1, x2)
+  output_shape = _infer_binary_broadcast_shape(x1.shape, x2.shape, output_shape)
+  output_dtype = output_dtype or x1.dtype
+  assert isinstance(output_dtype, tf.DType)
+  return BinaryOpWithBroadcasting(
+      tf_fn, x1, x2, convert_to_tensor_shape(output_shape),
+      output_dtype).outputs[0]
+
+
+def maximum(x1, x2, output_shape=None):
+  return binary_op_with_broadcasting(
+      tf.maximum, x1, x2, output_shape=output_shape)
+
+
+def minimum(x1, x2, output_shape=None):
+  return binary_op_with_broadcasting(
+      tf.minimum, x1, x2, output_shape=output_shape)
+
+
+def less(x1, x2, output_shape=None):
+  return binary_op_with_broadcasting(
+      tf.less, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
+
+
+def greater(x1, x2, output_shape=None):
+  return binary_op_with_broadcasting(
+      tf.greater, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
+
+
+def equal(x1, x2, output_shape=None):
+  return binary_op_with_broadcasting(
+      tf.equal, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
+
+
+def not_equal(x1, x2, output_shape=None):
+  return binary_op_with_broadcasting(
+      tf.not_equal, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
+
+
+def logical_and(x1, x2, output_shape=None):
+  return binary_op_with_broadcasting(
+      tf.logical_and, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
+
+
+def logical_or(x1, x2, output_shape=None):
+  return binary_op_with_broadcasting(
+      tf.logical_or, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
+
+
+def floordiv(x1, x2, output_shape=None):
+  output_dtype = x1.dtype if isinstance(x1, Tensor) else x2.dtype
+  return binary_op_with_broadcasting(
+      tf.floordiv, x1, x2, output_dtype=output_dtype, output_shape=output_shape)
+
+
+def mod(x1, x2, output_shape=None):
+  output_dtype = x1.dtype if isinstance(x1, Tensor) else x2.dtype
+  return binary_op_with_broadcasting(
+      tf.mod, x1, x2, output_dtype=output_dtype, output_shape=output_shape)
+
+
+class AddOperation(BinaryOpWithBroadcasting):
+  """Binary addition with broadcasting."""
+
+  def __init__(self, x1, x2, output_shape, name=None):
+    super(AddOperation, self).__init__(
+        tf.add, x1, x2, output_shape, x1.dtype, name=name or "add")
+    if x1.dtype != x2.dtype:
+      raise ValueError("Dtypes must be equal.")
+
+  def gradient(self, grad_ys):
+    dy = grad_ys[0]
+    return [reduce_sum(dy, output_shape=self.inputs[0].shape),
+            reduce_sum(dy, output_shape=self.inputs[1].shape)]
+
+
+class BroadcastOperation(Operation):
+  """Broadcast - output dims are a superset of input dims, in any order."""
+
+  def __init__(self, x, output_shape, name=None):
+    super(BroadcastOperation, self).__init__([x], name=name or "broadcast")
+    self._outputs = [Tensor(self, output_shape, x.dtype)]
+
+  def gradient(self, grad_ys):
+    return [reduce_sum(grad_ys[0], output_shape=self.inputs[0].shape)]
+
+  def lower(self, lowering):
+    ret = lowering.mesh_impl(self).broadcast_impl(
+        lowering.tensors[self.inputs[0]], self.inputs[0].shape,
+        self.outputs[0].shape)
+    lowering.tensors[self.outputs[0]] = ret
+
+
+def broadcast(x, new_shape):
+  return BroadcastOperation(x, new_shape).outputs[0]
+
+
+def _reduce_helper(input_shape,
+                   output_shape,
+                   mesh_layout,
+                   reduction_fn_string="SUM"):
+  """Returns slicewise function and reduced mesh dimensions.
+
+  Args:
+    input_shape: a TensorShape
+    output_shape: a TensorShape
+    mesh_layout: a dict (string -> int)
+    reduction_fn_string: "SUM" or "MAX"
+  Returns:
+    reduce_slice_fn: a function from tf.Tensor to tf.Tensor
+    reduced_mesh_axes: a list of integers
+  """
+  reduce_dims_indices = [
+      i for i, d in enumerate(input_shape.dims) if d not in output_shape.dims]
+  reduced_input_shape = TensorShape([
+      d for d in input_shape.dims if d in output_shape.dims])
+  perm = [reduced_input_shape.dims.index(d) for d in output_shape.dims]
+  def reduce_slice_fn(xslice):
+    ret = xslice
+    if reduce_dims_indices:
+      ret = reduction_fn(reduction_fn_string)(xslice, reduce_dims_indices)
+    if perm != list(xrange(len(perm))):
+      ret = tf.transpose(ret, perm)
+    return ret
+  reduced_mesh_axes = []
+  for i in reduce_dims_indices:
+    mesh_axis = mesh_layout.get(input_shape.dims[i].name, None)
+    if mesh_axis is not None:
+      reduced_mesh_axes.append(mesh_axis)
+  return reduce_slice_fn, reduced_mesh_axes
+
+
+class ReduceOperation(Operation):
+  """Reduction - output dims are a subset of input dims, in any order."""
+
+  def __init__(self, x, output_shape, reduction_fn_string, name=None):
+    super(ReduceOperation, self).__init__([x], name=name or "reduce")
+    self._outputs = [Tensor(self, output_shape, x.dtype)]
+    self._reduction_fn_string = reduction_fn_string
+
+  def gradient(self, grad_ys):
+    if self._reduction_fn_string == "SUM":
+      return [broadcast(grad_ys[0], self.inputs[0].shape)]
+    else:
+      raise ValueError("Gradients to other reductions not implemented")
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    slicewise_fn, reduced_mesh_axes = _reduce_helper(
+        self.inputs[0].shape, self.outputs[0].shape,
+        mesh_impl.layout, self._reduction_fn_string)
+    y = mesh_impl.slicewise(slicewise_fn, lowering.tensors[self.inputs[0]])
+    if reduced_mesh_axes:
+      def add_counter_fn():
+        lowering.add_counter("allreduce/%s/reduce_op" % reduced_mesh_axes,
+                             lowering.laid_out_size(self.outputs[0]))
+      if self._reduction_fn_string == "SUM":
+        y = LazyAllreduceSum(
+            mesh_impl, y, reduced_mesh_axes, add_counter_fn=add_counter_fn)
+      else:
+        y = mesh_impl.allreduce(
+            y, reduced_mesh_axes, self._reduction_fn_string)
+        add_counter_fn()
+    lowering.tensors[self.outputs[0]] = y
+
+
+class ConcatOperation(Operation):
+  """tf.concat.
+
+  All inputs have the same shape, except for the size of the dimension named
+  dim_name.
+  """
+
+  def __init__(self, xs, concat_dim_name, name=None):
+    super(ConcatOperation, self).__init__(xs, name=name or "concat")
+    # verify that the shapes are all compatible
+    dim_names = [dim.name for dim in xs[0].shape.dims]
+    self._concat_dim_name = concat_dim_name
+
+    if concat_dim_name not in dim_names:
+      raise ValueError("xs[0] does not contain a dimension named dim_name")
+    self._axis = dim_names.index(concat_dim_name)
+
+    should_be_equal = [
+        x.shape.resize_dimension(concat_dim_name, 0) for x in xs]
+    if not all(s == should_be_equal[0] for s in should_be_equal):
+      raise ValueError("shapes are not compatible %s" % xs)
+
+    self._input_sizes = [x.shape.dims[self._axis].size for x in xs]
+    output_size = sum(self._input_sizes)
+    self._outputs = [
+        Tensor(self, xs[0].shape.resize_dimension(concat_dim_name, output_size),
+               xs[0].dtype)]
+
+  def gradient(self, grad_ys):
+    dy = grad_ys[0]
+    return split(dy, self.outputs[0].shape.dims[self._axis], self._input_sizes)
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    if self._concat_dim_name in mesh_impl.layout:
+      raise ValueError("can't concat along split axis")
+    def slicewise_fn(*args):
+      return tf.concat(args, axis=self._axis, name="concat")
+    y = mesh_impl.slicewise(
+        slicewise_fn, *[lowering.tensors[x] for x in self._inputs])
+    lowering.tensors[self.outputs[0]] = y
+
+
+def concat(xs, concat_dim_name, name=None):
+  """Like tf.concat.
+
+  All inputs must have equal shape except for the sizes in the concatenated
+  dimension.  The dimension names should be the same, even that of the
+  concatenated dimension.
+
+  Args:
+    xs: a list of Tensors
+    concat_dim_name: a string
+    name: an optional string
+  Returns:
+    a Tensor
+  """
+  return ConcatOperation(xs, concat_dim_name, name).outputs[0]
+
+
+class SplitOperation(Operation):
+  """like tf.split.
+
+  TODO(noam, nikip): this code has never been run.  Run it and test it.
+  """
+
+  def __init__(self, x, split_dim, num_or_size_splits, name=None):
+    super(SplitOperation, self).__init__([x], name=name or "concat")
+
+    self._split_dim = split_dim
+    if split_dim not in x.shape.dims:
+      raise ValueError("%s does not contain dimension %s" % (x, split_dim))
+    self._axis = x.shape.dims.index(split_dim)
+
+    if isinstance(num_or_size_splits, list):
+      self._output_sizes = num_or_size_splits
+      if sum(num_or_size_splits) != split_dim.size:
+        raise ValueError(
+            "Sizes do not add up %s %s" % (num_or_size_splits, split_dim))
+    else:
+      assert isinstance(num_or_size_splits, int)
+      assert split_dim.size % num_or_size_splits == 0
+      self._output_sizes = (
+          [split_dim.size / num_or_size_splits] * num_or_size_splits)
+
+    self._outputs = [
+        Tensor(self, x.shape.resize_dimension(split_dim.name, output_size),
+               x.dtype) for output_size in self._output_sizes]
+
+  def gradient(self, grad_ys):
+    return concat(grad_ys, self._split_dim.name)
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    if self._split_dim.name in mesh_impl.layout:
+      raise ValueError("can't split along split axis")
+    def slicewise_fn(x):
+      # Since we return a tuple of tf.Tensor, slicewise will collate the
+      # outputs and return a tuple of LaidOutTensors.
+      return tuple(tf.split(x, self._output_sizes, axis=self._axis))
+    values = mesh_impl.slicewise(
+        slicewise_fn, lowering.tensors[self.inputs[0]])
+    for t, v in zip(self._outputs, values):
+      lowering.tensors[t] = v
+
+
+def split(x, split_dim, num_or_size_splits, name=None):
+  """Like tf.split.
+
+  Args:
+    x: a Tensor
+    split_dim: a Dimension in x.shape.dims
+    num_or_size_splits: either an integer dividing split_dim.size
+       or a list of integers adding up to split_dim.size
+    name: an optional string
+  Returns:
+    a list of Tensors.
+  """
+  return SplitOperation(x, split_dim, num_or_size_splits, name=name).outputs
+
+
+class StackOperation(Operation):
+  """Like tf.stack."""
+
+  def __init__(self, xs, dim_name, axis, name=None):
+    super(StackOperation, self).__init__(xs, name=name or "stack")
+    self._axis = axis
+    self._new_dim = Dimension(dim_name, len(xs))
+    input_shape = xs[0].shape
+    for x in xs:
+      if x.shape != xs[0].shape:
+        raise ValueError(
+            "inputs to stack must have the same shape, got %s" % xs)
+    output_shape = TensorShape(
+        input_shape.dims[:axis] + [self._new_dim]+ input_shape.dims[axis:])
+    self._outputs = [Tensor(self, output_shape, xs[0].dtype)]
+
+  def gradient(self, grad_ys):
+    return unstack(grad_ys[0], self._new_dim)
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    if self._new_dim in mesh_impl.layout:
+      raise ValueError("can't stack along split axis")
+    inputs = [lowering.tensors[t] for t in self._inputs]
+    def slicewise_fn(*args):
+      return tf.stack(args, axis=self._axis)
+    ret = mesh_impl.slicewise(slicewise_fn, *inputs)
+    lowering.tensors[self.outputs[0]] = ret
+
+
+def stack(xs, dim_name, axis, name=None):
+  """Stack multiple Tensors to make a new dimension.
+
+  Args:
+    xs: a list of Tensors with identical shapes.
+    dim_name: a string (name of the new dimension)
+    axis: an integer (index of the new dimension in the output shape)
+    name: an optional string
+
+  Returns:
+    a Tensor
+  """
+  ret = StackOperation(xs, dim_name, axis, name).outputs[0]
+  return ret
+
+
+class UnstackOperation(Operation):
+  """Split into multiple Tensors, eliminating a dimension."""
+
+  def __init__(self, x, dim, name=None):
+    super(UnstackOperation, self).__init__([x], name=name or "unstack")
+    self._dim = dim
+    self._axis = x.shape.dims.index(dim)
+    output_shape = x.shape - dim
+    self._outputs = [
+        Tensor(self, output_shape, x.dtype) for _ in xrange(dim.size)]
+
+  def gradient(self, grad_ys):
+    return [stack(grad_ys, self._dim.name, self._axis)]
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    if self._dim in mesh_impl.layout:
+      raise ValueError("can't unstack along split axis")
+    def slicewise_fn(x):
+      return tuple(tf.unstack(x, num=self._dim.size, axis=self._axis))
+    output_values = mesh_impl.slicewise(
+        slicewise_fn, lowering.tensors[self._inputs[0]])
+    for t, v in zip(self.outputs, list(output_values)):
+      lowering.tensors[t] = v
+
+
+def unstack(x, dim, name=None):
+  """Split into multiple Tensors, eliminating a dimension.
+
+  Args:
+    x: a Tensor
+    dim: a Dimension
+    name: an optional string
+
+  Returns:
+    a list of dim.size Tensors, each with shape (x.shape - dim)
+  """
+  return UnstackOperation(x, dim, name).outputs
+
+
+def _einsum_helper(input_shapes, output_shape, mesh_layout):
+  """Returns slicewise function and reduced mesh dimensions.
+
+  Assumes the output shape contains no new dimensions.
+
+  Args:
+    input_shapes: a list of TensorShapes
+    output_shape: a TensorShape
+    mesh_layout: a dict (string -> int)
+  Returns:
+    einsum_slice_fn: a function from tf.Tensors to tf.Tensor
+    reduced_mesh_axes: a list of integers
+  """
+  input_shape_set = set(sum([s.dims for s in input_shapes], []))
+  total_num_dims = len(input_shape_set)
+  # list of input shapes that contain all dimensions.
+  full_shapes = [
+      s for s in input_shapes + [output_shape] if s.ndims == total_num_dims]
+  full_shape = (
+      full_shapes[0] if full_shapes else TensorShape(list(input_shape_set)))
+  reduce_slice_fn, reduced_mesh_axes = _reduce_helper(
+      full_shape, output_shape, mesh_layout)
+  def einsum_slice_fn_naive(*slices):
+    # naive einsum implementation where we broadcst all inputs to the full
+    # shape, multiply componentwise, then reduce.
+    return reduce_slice_fn(reduce(tf.multiply, [
+        _expand_dims(x, input_shape, full_shape)
+        for x, input_shape in zip(slices, input_shapes)]))
+  if full_shapes:
+    # it is not wasteful of space to broadcast fully and then reduce.
+    # this helps to avoid some inefficient GPU implementations.
+    einsum_slice_fn = einsum_slice_fn_naive
+  else:
+    # call tf.einsum
+    equation = _einsum_equation(input_shapes, output_shape)
+    def einsum_slice_fn(*slices):
+      if slices[0].dtype.is_floating:
+        return tf.einsum(equation, *slices)
+      else:
+        return einsum_slice_fn_naive(*slices)
+  return einsum_slice_fn, reduced_mesh_axes
+
+
+class EinsumOperation(Operation):
+  """Einstein summation (matmul, etc).
+
+  The equation follows the dimensions in the input and output shapes.
+
+  Every dimension must occur in at least two of the input/output Tensors.
+  i.e. no new dimensions in the output, and no reduction of dimensions that
+  occur in only one input.
+  """
+
+  def __init__(self, inputs, output_shape, name=None):
+    super(EinsumOperation, self).__init__(inputs, name=name or "einsum")
+    if not inputs:
+      raise ValueError("Einsum needs at least one input")
+    for x in inputs:
+      if x.dtype != inputs[0].dtype:
+        raise ValueError("Input dtypes must be equal")
+    self._outputs = [Tensor(self, output_shape, inputs[0].dtype)]
+
+  def gradient(self, grad_ys):
+    dy = grad_ys[0]
+    xs = self.inputs
+    return [
+        einsum([dy] + [xs[j] for j in xrange(len(xs)) if j != i], xs[i].shape)
+        for i in xrange(len(self.inputs))]
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    xs = self.inputs
+    input_shape_set = set(sum([x.shape.dims for x in xs], []))
+    output_shape = self.outputs[0].shape
+    intersection_shape = TensorShape(
+        [d for d in output_shape.dims if d in input_shape_set])
+    einsum_slice_fn, reduced_mesh_axes = _einsum_helper(
+        [x.shape for x in self.inputs], intersection_shape, mesh_impl.layout)
+    y = mesh_impl.slicewise(
+        einsum_slice_fn, *[lowering.tensors[x] for x in self.inputs])
+    if reduced_mesh_axes:
+      def add_counter_fn():
+        lowering.add_counter(
+            "allreduce/%s/einsum_op" % reduced_mesh_axes,
+            mesh_impl.laid_out_size(intersection_shape))
+      y = LazyAllreduceSum(
+          mesh_impl, y, reduced_mesh_axes, add_counter_fn=add_counter_fn)
+    # broadcast from intersection_shape to output_shape
+    if intersection_shape != output_shape:
+      y = mesh_impl.broadcast_impl(y, intersection_shape, output_shape)
+    lowering.tensors[self.outputs[0]] = y
+    computation_shape = TensorShape(list(input_shape_set))
+    lowering.add_counter("einsum", mesh_impl.laid_out_size(computation_shape))
+    lowering.add_counter("einsum_unique", computation_shape.size)
+
+
+class SliceOperation(Operation):
+  """tf.slice.
+
+  We support the slice operation along one axis. Similar to tf.slice, specify
+  the begin and size values for the slice_dim.
+  """
+
+  def __init__(self, x, begin, size, slice_dim_name, name=None):
+    super(SliceOperation, self).__init__([x], name=name or "slice")
+    dim_names = x.shape.dimension_names
+    self._axis = axis = dim_names.index(slice_dim_name)
+    self._begin = begin
+    self._slice_dim = Dimension(slice_dim_name, size)
+    input_shape = self._inputs[0].shape
+    output_shape = TensorShape(
+        input_shape.dims[:axis] + [self._slice_dim] + input_shape.dims[axis+1:])
+    self._outputs = [Tensor(self, output_shape, x.dtype)]
+
+  def gradient(self, grad_ys):
+    actual_size = self._inputs[0].shape.dims[self._axis].size
+    return [
+        pad(grad_ys[0],
+            [self._begin, actual_size - self._slice_dim.size - self._begin],
+            self._slice_dim.name)]
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    if self._slice_dim in mesh_impl.layout:
+      raise ValueError("can't slice along split axis")
+    inputs = self._inputs[0]
+    ndims = self._inputs[0].shape.ndims
+    axis = self._axis
+    begin = [0] * axis + [self._begin] + [0] * (ndims - axis - 1)
+    size = self._outputs[0].shape.to_tf_shape
+
+    def slicewise_fn(x, begin, size):
+      return tf.slice(x, begin, size, name="slice")
+    y = mesh_impl.slicewise(
+        slicewise_fn, lowering.tensors[inputs], begin, size)
+    lowering.tensors[self.outputs[0]] = y
+
+
+class PadOperation(Operation):
+  """tf.pad.
+
+  Similar to tf.pad but we only pad along one axis given by pad_dim_name
+  with values specified by paddings. paddings is a list of two
+  values, giving the padding value before and after pad_dim.
+  """
+
+  def __init__(self, x, paddings, pad_dim_name, name=None):
+    super(PadOperation, self).__init__([x], name=name or "pad")
+    assert len(paddings) == 2
+    input_shape = self._inputs[0].shape
+    dim_names = [dim.name for dim in x.shape.dims]
+    if pad_dim_name not in dim_names:
+      raise ValueError("Padding dim name %s not found in input." % pad_dim_name)
+    self._paddings = paddings
+    self._axis = axis = dim_names.index(pad_dim_name)
+    output_size = input_shape.dims[axis].size + sum(paddings)
+    self._output_dim = Dimension(pad_dim_name, output_size)
+    output_shape = TensorShape(
+        input_shape.dims[:axis] +
+        [self._output_dim] + input_shape.dims[axis+1:])
+    self._outputs = [Tensor(self, output_shape, x.dtype)]
+
+  def gradient(self, grad_ys):
+    # slice_dim = self._inputs[0].shape.dims[self._axis]
+    slice_dim_name = self._output_dim.name
+    slice_size = self._inputs[0].shape.dims[self._axis].size
+    return [slice(grad_ys[0], self._paddings[0], slice_size, slice_dim_name)]
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    if self._output_dim in mesh_impl.layout:
+      raise ValueError("can't pad along split axis")
+    inputs = self._inputs[0]
+    ndims = self._inputs[0].shape.ndims
+    axis = self._axis
+    paddings = [[0, 0]] * axis + [self._paddings] + [[0, 0]]* (ndims - axis - 1)
+
+    def slicewise_fn(x, paddings):
+      return tf.pad(x, paddings, name="pad")
+    y = mesh_impl.slicewise(
+        slicewise_fn, lowering.tensors[inputs], paddings)
+    lowering.tensors[self.outputs[0]] = y
+
+
+class OneHotOperation(Operation):
+  """one_hot.
+  """
+
+  def __init__(self, indices, output_dim, on_value, off_value, dtype,
+               name=None):
+    super(OneHotOperation, self).__init__([indices], name=name or "one_hot")
+    if not indices.dtype.is_integer:
+      raise ValueError("indices requires an integer dtype got %s" % indices)
+    self._output_dim = output_dim
+    self._on_value = on_value
+    self._off_value = off_value
+    self._dtype = dtype
+    output_shape = TensorShape(indices.shape.dims + [output_dim])
+    self._outputs = [Tensor(self, output_shape, dtype)]
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    indices = self.inputs[0]
+    output_shape = self.outputs[0].shape
+    output_slice_shape = mesh_impl.slice_shape(output_shape)
+    mesh_axis = mesh_impl.tensor_dimension_to_mesh_axis(self._output_dim)
+    depth = output_slice_shape[-1]
+    if mesh_axis is None:
+      offset = 0
+    else:
+      offset = mesh_impl.slicewise(
+          tf.multiply, mesh_impl.laid_out_pcoord(mesh_axis), depth)
+
+    def slicewise_fn(indices_slice, offset):
+      return tf.one_hot(indices_slice - offset,
+                        depth,
+                        on_value=tf.cast(self._on_value, self._dtype),
+                        off_value=tf.cast(self._off_value, self._dtype),
+                        dtype=self._dtype)
+    y = mesh_impl.slicewise(
+        slicewise_fn, lowering.tensors[indices], offset)
+    lowering.tensors[self.outputs[0]] = y
+
+
+class InfeedOperation(Operation):
+  """Infeed a tf.Tensor onto a mesh."""
+
+  def __init__(self, mesh, tf_tensor, shape, name=None):
+    super(InfeedOperation, self).__init__([], mesh=mesh, name=name or "infeed")
+    self._outputs = [Tensor(self, shape, tf_tensor.dtype)]
+    self._tf_tensor = tf_tensor
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    lowering.tensors[self.outputs[0]] = mesh_impl.infeed(
+        self.outputs[0], self._tf_tensor)
+
+
+def anonymous_shape(shape):
+  shape = convert_to_tensor_shape(shape)
+  return TensorShape([Dimension("_anonymous_%i" % i, d.size)
+                      for i, d in enumerate(shape)])
+
+
+def anonymize(x):
+  return reshape(x, anonymous_shape(x.shape))
+
+
+def infeed(mesh, tf_tensor, shape=None, name=None):
+  tf_tensor = tf.convert_to_tensor(tf_tensor)
+  if shape is None:
+    shape = TensorShape([])
+    assert not tf_tensor.shape.as_list()
+  return InfeedOperation(
+      mesh, tf_tensor, convert_to_tensor_shape(shape), name=name).outputs[0]
+
+
+def infeed_fully_replicated(mesh, tf_tensor, shape, name=None):
+  return reshape(infeed(mesh, tf_tensor, anonymous_shape(shape), name), shape)
+
+
+class Variable(Operation):
+  """Variable."""
+
+  def __init__(self, mesh, name, shape, dtype, initializer,
+               trainable, **kwargs):
+    super(Variable, self).__init__([], mesh, name="name_will_be_set_later")
+    self._trainable = trainable
+    with tf.device("cpu:0"), mtf_utils.outside_all_rewrites():
+      self.master = tf.get_variable(
+          name, shape.to_tf_shape, dtype=dtype, initializer=initializer,
+          **kwargs)
+    self._name = self.master.name[:self.master.name.find(":")]
+    self._outputs = [Tensor(self, shape, dtype)]
+    self.graph.all_variables.append(self)
+    if trainable:
+      self.graph.trainable_variables.append(self)
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    with mtf_utils.outside_all_rewrites():
+      sv = mesh_impl.LaidOutVariable(self, mesh_impl)
+    lowering.variables[self] = sv
+    lowering.tensors[self.outputs[0]] = sv.laid_out_tensor
+    if self._trainable:
+      lowering.add_counter("variables/trainable", self.outputs[0].size)
+    else:
+      lowering.add_counter("variables/untrainable", self.outputs[0].size)
+
+  @property
+  def value(self):
+    return self.outputs[0]
+
+  @property
+  def shape(self):
+    return self.value.shape
+
+  @property
+  def dtype(self):
+    return self.value.dtype
+
+
+def get_variable(mesh, name, shape, dtype=tf.float32,
+                 initializer=None, trainable=True,
+                 activation_dtype=None, **kwargs):
+  ret = Variable(
+      mesh, name, convert_to_tensor_shape(shape), dtype, initializer,
+      trainable, **kwargs).outputs[0]
+  if activation_dtype and activation_dtype != dtype:
+    ret = cast(ret, activation_dtype)
+  return ret
+
+
+class Assign(Operation):
+  """Assign to a variable."""
+
+  def __init__(self, var, new_val, name=None):
+    super(Assign, self).__init__([new_val], var.mesh, name=name or "assign")
+    self._var = var
+    self._outputs = []
+
+  def lower(self, lowering):
+    lowering.operations[self] = lowering.variables[self._var].assign_to_slices(
+        lowering.tensors[self.inputs[0]].to_laid_out_tensor().all_slices)
+
+
+def assign(var, new_val):
+  """Assign a new value to a variable.
+
+  Args:
+    var: either a Variable operation or its output Tensor.
+    new_val: a Tensor
+  Returns:
+    an Operation
+  Raises:
+    ValueError: if var is not a Variable and var.operation is not a Variable
+  """
+  if isinstance(var, Tensor):
+    var = var.operation
+  if not isinstance(var, Variable):
+    raise ValueError("var must be a mtf.Variable or its output Tensor.")
+  return Assign(var, new_val)
+
+
+class Depend(Operation):
+  """Control dependency."""
+
+  def __init__(self, x, dependencies, name=None):
+    super(Depend, self).__init__([x], x.mesh, name=name or "depend")
+    for d in dependencies:
+      if not isinstance(d, Operation):
+        raise ValueError("dependencies must be mtf.Operations. got %s" % d)
+    self._dependencies = dependencies
+    self._outputs = [Tensor(self, x.shape, x.dtype)]
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    if not mesh_impl.supports_control_dependencies:
+      raise ValueError("Mesh does not suppport control dependencies.")
+    with tf.control_dependencies(
+        [lowering.operations[d] for d in self._dependencies]):
+      lowering.tensors[self.outputs[0]] = mesh_impl.slicewise(
+          tf.identity,
+          lowering.tensors[self.inputs[0]])
+
+  def gradient(self, grad_ys):
+    return grad_ys
+
+
+def depend(x, dependencies):
+  """Identity of Tensor x that dependes on operations dependencies.
+
+  Args:
+    x: a Tensor
+    dependencies: a list of Operations
+  Returns:
+    an tensor
+  """
+  return Depend(x, dependencies).outputs[0]
+
+
+class Constant(Operation):
+  """A tensor where every element is the same constant value."""
+
+  def __init__(self, mesh, value, shape, dtype, name=None):
+    super(Constant, self).__init__([], mesh, name=name or "constant")
+    self._outputs = [Tensor(self, shape, dtype)]
+    self._value = value
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    slice_shape = mesh_impl.slice_shape(self.outputs[0].shape)
+    def tf_fn():
+      return tf.constant(value=self._value,
+                         dtype=self.outputs[0].dtype,
+                         shape=slice_shape)
+    lowering.tensors[self.outputs[0]] = mesh_impl.slicewise(tf_fn)
+
+
+def constant(mesh, value, shape=None, dtype=tf.float32):
+  shape = convert_to_tensor_shape(shape)
+  return Constant(mesh, value,
+                  shape if shape is not None else TensorShape([]),
+                  dtype).outputs[0]
+
+
+def zeros(mesh, shape, dtype=tf.float32):
+  return constant(mesh, 0, shape=convert_to_tensor_shape(shape), dtype=dtype)
+
+
+def zeros_like(t):
+  return zeros(t.mesh, t.shape, dtype=t.dtype)
+
+
+class StopGradient(Operation):
+  """Similar to tf.stop_gradient."""
+
+  def __init__(self, x, name=None):
+    super(StopGradient, self).__init__(
+        [x], x.mesh, name=name or "stop_gradient")
+    self._outputs = [Tensor(self, x.shape, x.dtype)]
+
+  def lower(self, lowering):
+    lowering.tensors[self.outputs[0]] = lowering.tensors[self.inputs[0]]
+
+  @property
+  def has_gradient(self):
+    return False
+
+
+def stop_gradient(x):
+  return StopGradient(x).outputs[0]
+
+
+class PrintOperation(Operation):
+  """Similar to tf.stop_gradient."""
+
+  def __init__(self, x, data, message, name=None, **kwargs):
+    super(PrintOperation, self).__init__(
+        [x], x.mesh, name=name or "Print")
+    self._outputs = [Tensor(self, x.shape, x.dtype)]
+    self._data = data
+    self._message = message
+    self._kwargs = kwargs
+
+  def lower(self, lowering):
+    lowering.tensors[self.outputs[0]] = lowering.mesh_impl(self).Print(
+        lowering.tensors[self.inputs[0]],
+        [lowering.tensors[d] for d in self._data], self._message,
+        **self._kwargs)
+
+  def gradient(self, grad_ys):
+    return grad_ys
+
+
+def Print(x, data, message, **kwargs):  # pylint: disable=invalid-name
+  """Call tf.Print.
+
+  Args:
+    x: a Tensor.
+    data: a list of Tensor
+    message: a string
+    **kwargs: keyword arguments to tf.Print
+  Returns:
+    a Tensor which is identical in value to x
+  """
+  return PrintOperation(x, data, message, **kwargs).outputs[0]
+
+
+class ReshapeOperation(Operation):
+  """Similar to tf.stop_gradient."""
+
+  def __init__(self, x, new_shape, name=None):
+    super(ReshapeOperation, self).__init__([x], x.mesh, name=name or "reshape")
+    self._outputs = [Tensor(self, new_shape, x.dtype)]
+
+  def lower(self, lowering):
+    """Lower the ReshapeOperation.
+
+    Reshaping can require collective communication between processors.
+    We haven't yet implemented all possible reshapes.  We try to handle the
+    common cases here - otherwise we raise a NotImplementedError.
+
+    Args:
+      lowering: a Lowering
+    Raises:
+      NotImplementedError: if we haven't covered this case
+    """
+    old_shape = self.inputs[0].shape
+    new_shape = self.outputs[0].shape
+    mesh_impl = lowering.mesh_impl(self)
+    slices = lowering.tensors[self.inputs[0]]
+
+    mesh_axis_to_cumprod_old = mesh_impl.mesh_axis_to_cumprod(old_shape)
+    mesh_axis_to_cumprod_new = mesh_impl.mesh_axis_to_cumprod(new_shape)
+    # Figure out what needs to be done for different mesh-axes
+    mesh_axes_allsplit = []
+    mesh_axes_allconcat = []
+    mesh_axes_alltoall = []
+    for mesh_axis, (old_cumprod, new_cumprod) in enumerate(
+        zip(mesh_axis_to_cumprod_old, mesh_axis_to_cumprod_new)):
+      if new_cumprod != old_cumprod:
+        if old_cumprod is None:
+          # split in new layout but not in old layout - we need an allsplit
+          mesh_axes_allsplit.append(mesh_axis)
+        elif new_cumprod is None:
+          # split in old layout but not in new layout - we need an allconcat
+          mesh_axes_allconcat.append(mesh_axis)
+        else:
+          # split differently in old and new layouts - we need an alltoall
+          mesh_axes_alltoall.append(mesh_axis)
+
+    laid_out_size = mesh_impl.laid_out_size(old_shape)
+
+    for mesh_axis in mesh_axes_allsplit:
+      tensor_axis = old_shape.cumprod_to_tensor_axis(
+          mesh_axis_to_cumprod_new[mesh_axis])
+      if tensor_axis is None:
+        # TODO(noam): try to handle this case
+        raise NotImplementedError(
+            "Try first reshaping to insert a new tf dimension,"
+            " then changing layout.")
+      slices = mesh_impl.allsplit(slices, mesh_axis, tensor_axis)
+      laid_out_size //= mesh_impl.shape[mesh_axis]
+    for mesh_axis in mesh_axes_alltoall:
+      split_tensor_axis = old_shape.cumprod_to_tensor_axis(
+          mesh_axis_to_cumprod_new[mesh_axis])
+      if split_tensor_axis is None:
+        # TODO(noam): try to handle this case
+        raise NotImplementedError(
+            "Try first reshaping to insert a new tf dimension,"
+            " then changing layout.")
+      concat_tensor_axis = old_shape.cumprod_to_tensor_axis(
+          mesh_axis_to_cumprod_old[mesh_axis])
+      assert concat_tensor_axis is not None
+      slices = mesh_impl.alltoall(
+          slices, mesh_axis, split_tensor_axis, concat_tensor_axis)
+      lowering.add_counter(
+          "alltoall/%s/reshape_op" % mesh_axis, laid_out_size)
+
+    for mesh_axis in mesh_axes_allconcat:
+      tensor_axis = old_shape.cumprod_to_tensor_axis(
+          mesh_axis_to_cumprod_old[mesh_axis])
+      assert tensor_axis is not None
+      slices = mesh_impl.allconcat(slices, mesh_axis, tensor_axis)
+      laid_out_size *= mesh_impl.shape[mesh_axis]
+      lowering.add_counter(
+          "allconcat/%s/reshape_op" % mesh_axis, laid_out_size)
+    # now reshape the slices
+    old_slice_shape = mesh_impl.slice_shape(old_shape)
+    new_slice_shape = mesh_impl.slice_shape(new_shape)
+    if new_slice_shape != old_slice_shape:
+      def reshape_fn(x):
+        return tf.reshape(x, new_slice_shape)
+      slices = mesh_impl.slicewise(reshape_fn, slices)
+    lowering.tensors[self.outputs[0]] = slices
+
+  def gradient(self, grad_ys):
+    return [reshape(grad_ys[0], self.inputs[0].shape)]
+
+
+def reshape(x, new_shape):
+  return ReshapeOperation(x, convert_to_tensor_shape(new_shape)).outputs[0]
+
+
+def rename_dimension(x, old_name, new_name):
+  """Reshape a Tensor, renaming one dimension.
+
+  Args:
+    x: a Tensor
+    old_name: a string
+    new_name: a string
+
+  Returns:
+    a Tensor
+  """
+  return reshape(x, x.shape.rename_dimension(old_name, new_name))
+
+
+def einsum(xs, output_shape=None, name=None):
+  """Einstein summation.
+
+  If output_shape is not specified and there are two inputs, reduce over
+  all common dimensions and default the output shape to the unique dimensions
+  of the first input followed by the unique dimensions of the second input.
+
+  Args:
+    xs: a list of Tensors
+    output_shape: an optional TensorShape.
+    name: an optional string
+  Returns:
+    a Tensor
+  Raises:
+    ValueError: if the output shape cannot be inferred
+  """
+  output_shape = convert_to_tensor_shape(output_shape)
+  if output_shape is None:
+    if len(xs) == 2:
+      output_shape = TensorShape(
+          [d for d in xs[0].shape.dims if d not in xs[1].shape.dims] +
+          [d for d in xs[1].shape.dims if d not in xs[0].shape.dims])
+    else:
+      raise ValueError("could not infer einsum output_shape for inputs %s" %
+                       [x.to_string for x in xs])
+  return EinsumOperation(xs, output_shape, name=name).outputs[0]
+
+
+def matmul(a, b, output_shape=None, name=None):
+  return einsum([a, b], output_shape=output_shape, name=name)
+
+
+def _reduction_output_shape(x, output_shape, reduced_dim):
+  """Helper function to reduce_sum, etc."""
+  if output_shape is None:
+    if reduced_dim is None:
+      return TensorShape([])
+    else:
+      if reduced_dim not in x.shape.dims:
+        raise ValueError(
+            "reduced_dim=%s not in x.shape.dims=%s" % (reduced_dim, x.shape))
+      return x.shape - reduced_dim
+  elif reduced_dim is not None:
+    raise ValueError("do not specify both reduced_dim and output_shape")
+  else:
+    return output_shape
+
+
+def reduce_sum(x,
+               disable_positional_args=None,
+               output_shape=None,
+               reduced_dim=None,
+               name=None):
+  """Reduction on 1 or more axes.
+
+  If reduced_dim is present, then only that dimension is reduced out.
+  Alternatively, specify output_shape.
+  Do not specify both reduced_dim and output_shape.
+  If neither is specified, then all dimensions are reduced out.
+
+  Args:
+    x: a Tensor
+    disable_positional_args: None
+    output_shape: an optional TensorShape.  Must be a subsequence of x.shape.
+    reduced_dim: a mtf.Dimension
+    name: an optional string
+  Returns:
+    a Tensor
+  """
+  output_shape = convert_to_tensor_shape(output_shape)
+  reduced_dim = convert_to_dimension(reduced_dim)
+  assert disable_positional_args is None
+  output_shape = _reduction_output_shape(x, output_shape, reduced_dim)
+  if output_shape == x.shape:
+    return x
+  return ReduceOperation(x, output_shape, "SUM", name=name).outputs[0]
+
+
+def reduce_mean(x,
+                disable_positional_args=None,
+                output_shape=None,
+                reduced_dim=None,
+                name=None):
+  """Reduction on 1 or more axes.
+
+  If reduced_dim is present, then only that dimension is reduced out.
+  Alternatively, specify output_shape.
+  Do not specify both reduced_dim and output_shape.
+  If neither is specified, then all dimensions are reduced out.
+
+  Args:
+    x: a Tensor
+    disable_positional_args: None
+    output_shape: an optional TensorShape. Must be a subsequence of x.shape.
+    reduced_dim: a mtf.Dimension
+    name: an optional string
+
+  Returns:
+    a Tensor
+  """
+  output_shape = convert_to_tensor_shape(output_shape)
+  reduced_dim = convert_to_dimension(reduced_dim)
+  assert disable_positional_args is None
+  output_shape = _reduction_output_shape(x, output_shape, reduced_dim)
+  with tf.variable_scope(name, default_name="reduce_mean"):
+    if output_shape == x.shape:
+      return x
+    return reduce_sum(
+        x, output_shape=output_shape) * (output_shape.size / x.shape.size)
+
+
+def reduce_max(x,
+               disable_positional_args=None,
+               output_shape=None,
+               reduced_dim=None,
+               name=None):
+  """Reduction on 1 or more axes.
+
+  Args:
+    x: a Tensor
+    disable_positional_args: None
+    output_shape: an optional TensorShape.  Must be a subsequence of x.shape.
+    reduced_dim: an optional Dimension
+    name: an optional string
+  Returns:
+    a Tensor
+  """
+  output_shape = convert_to_tensor_shape(output_shape)
+  reduced_dim = convert_to_dimension(reduced_dim)
+  assert disable_positional_args is None
+  output_shape = _reduction_output_shape(x, output_shape, reduced_dim)
+  if output_shape is None:
+    output_shape = TensorShape([])
+  if output_shape == x.shape:
+    return x
+  return ReduceOperation(
+      x, output_shape, "MAX", name=name or "reduce_max").outputs[0]
+
+
+def reduce_min(x,
+               disable_positional_args=None,
+               output_shape=None,
+               reduced_dim=None,
+               name=None):
+  """Reduction on 1 or more axes.
+
+  Args:
+    x: a Tensor
+    disable_positional_args: None
+    output_shape: an optional TensorShape.  Must be a subsequence of x.shape.
+    reduced_dim: an optional Dimension
+    name: an optional string
+  Returns:
+    a Tensor
+  """
+  output_shape = convert_to_tensor_shape(output_shape)
+  reduced_dim = convert_to_dimension(reduced_dim)
+  assert disable_positional_args is None
+  output_shape = _reduction_output_shape(x, output_shape, reduced_dim)
+  if output_shape is None:
+    output_shape = TensorShape([])
+  if output_shape == x.shape:
+    return x
+  return ReduceOperation(
+      x, output_shape, "MIN", name=name or "reduce_min").outputs[0]
+
+
+def reduce_all(x,
+               disable_positional_args=None,
+               output_shape=None,
+               reduced_dim=None,
+               name=None):
+  output_shape = convert_to_tensor_shape(output_shape)
+  reduced_dim = convert_to_dimension(reduced_dim)
+  return cast(reduce_min(to_float(x),
+                         disable_positional_args=disable_positional_args,
+                         output_shape=output_shape,
+                         reduced_dim=reduced_dim,
+                         name=name or "reduce_all"), tf.bool)
+
+
+def reduce_any(x,
+               disable_positional_args=None,
+               output_shape=None,
+               reduced_dim=None,
+               name=None):
+  output_shape = convert_to_tensor_shape(output_shape)
+  reduced_dim = convert_to_dimension(reduced_dim)
+  return cast(reduce_max(to_float(x),
+                         disable_positional_args=disable_positional_args,
+                         output_shape=output_shape,
+                         reduced_dim=reduced_dim,
+                         name=name or "reduce_any"), tf.bool)
+
+
+def top_1(x, reduced_dim, dtype=tf.int32, name=None):
+  """Argmax and Max.
+
+  Args:
+    x: a Tensor
+    reduced_dim: a Dimension in x.shape.dims
+    dtype: a tf.dtype (for the output)
+    name: an optional string
+  Returns:
+    indices: a Tensor with given dtype
+    values: optional Tensor equal to mtf.reduce_max(x, reduced_dim=reduced_dim)
+  """
+  reduced_dim = convert_to_dimension(reduced_dim)
+  with tf.name_scope(name, default_name="top_1"):
+    max_val = reduce_max(x, reduced_dim=reduced_dim)
+    is_max = to_float(equal(x, max_val))
+    pos = range(x.mesh, reduced_dim, tf.float32)
+    ret = reduce_max(is_max * pos, reduced_dim=reduced_dim)
+    ret = cast(ret, dtype)
+    return ret, max_val
+
+
+def argmax(x, reduced_dim, dtype=tf.int32, name=None):
+  reduced_dim = convert_to_dimension(reduced_dim)
+  return top_1(x, reduced_dim, dtype, name)[0]
+
+
+def top_k(x, reduced_dim, new_dim, dtype=tf.int32, name=None):
+  """Like tf.top_k.
+
+  This operation returns two tensors with the same shape.  The output shape
+  is identical to the shape of x, except that reduced_dim is replaced by
+  new_dim.
+
+  Args:
+    x: a Tensor
+    reduced_dim: a Dimension in x.shape.dims.
+    new_dim: a Dimension.  The size determines k.
+    dtype: optional dtype for indices.
+    name: optional string.
+  Returns:
+    indices: a Tensor with given dtype.
+    values: a Tensor with same type as x.
+  """
+  reduced_dim = convert_to_dimension(reduced_dim)
+  new_dim = convert_to_dimension(new_dim)
+  indices = []
+  values = []
+  k = new_dim.size
+  with tf.name_scope(name, default_name="top_k"):
+    for i in xrange(k):
+      max_index, max_val = top_1(x, reduced_dim, dtype)
+      indices.append(max_index)
+      values.append(max_val)
+      if i + 1 < k:
+        x += one_hot(max_index, reduced_dim, on_value=-1e9)
+  axis = x.shape.dims.index(reduced_dim)
+  return stack(indices, new_dim.name, axis), stack(values, new_dim.name, axis)
+
+
+def sample_with_temperature(x, dim, temperature=1.0, dtype=tf.int32, name=None):
+  dim = convert_to_dimension(dim)
+  with tf.name_scope(name, default_name="sample_with_temperature"):
+    if temperature != 0.0:
+      # gumbel trick
+      g = -log(-log(random_uniform(x.mesh, x.shape, dtype=x.dtype)))
+      x += g * temperature
+    return argmax(x, dim, dtype, name)
+
+
+def add(x1, x2, output_shape=None, name=None):
+  """Binary addition with broadcsting.
+
+  Args:
+    x1: a Tensor
+    x2: a Tensor
+    output_shape: an optional TensorShape
+    name: an optional string
+  Returns:
+    a Tensor
+  """
+  output_shape = convert_to_tensor_shape(output_shape)
+  if not isinstance(x2, Tensor):
+    return ScalarAddOperation(x1, x2).outputs[0]
+  with tf.name_scope(name, default_name="add"):
+    x1, x2 = binary_arguments_to_tensors(x1, x2)
+    return AddOperation(
+        x1, x2, output_shape=_infer_binary_broadcast_shape(
+            x1.shape, x2.shape, output_shape)).outputs[0]
+
+
+def sub(x1, x2, output_shape=None, name=None):
+  """Binary subtraction with broadcsting.
+
+  Args:
+    x1: a Tensor
+    x2: a Tensor
+    output_shape: an optional TensorShape
+    name: an optional string
+  Returns:
+    a Tensor
+  """
+  output_shape = convert_to_tensor_shape(output_shape)
+  if not isinstance(x2, Tensor):
+    return ScalarAddOperation(x1, -x2).outputs[0]
+  with tf.name_scope(name, default_name="sub"):
+    x1, x2 = binary_arguments_to_tensors(x1, x2)
+    return add(x1, negative(x2), output_shape=output_shape)
+
+
+def multiply(x1, x2, output_shape=None, name=None):
+  """Binary multiplication with broadcsting.
+
+  Args:
+    x1: a Tensor
+    x2: a Tensor
+    output_shape: an optional TensorShape
+    name: an optional string
+  Returns:
+    a Tensor
+  """
+  if not isinstance(x2, Tensor):
+    return ScalarMultiplyOperation(x1, x2).outputs[0]
+  with tf.name_scope(name, default_name="mul"):
+    x1, x2 = binary_arguments_to_tensors(x1, x2)
+    return einsum(
+        [x1, x2],
+        output_shape=_infer_binary_broadcast_shape(
+            x1.shape, x2.shape, output_shape))
+
+
+def divide(x1, x2, output_shape=None, name=None):
+  """Binary division with broadcsting.
+
+  Args:
+    x1: a Tensor
+    x2: a Tensor
+    output_shape: an optional TensorShape
+    name: an optional string
+  Returns:
+    a Tensor
+  """
+  output_shape = convert_to_tensor_shape(output_shape)
+  if not isinstance(x2, Tensor):
+    return ScalarMultiplyOperation(x1, 1.0 / x2).outputs[0]
+  with tf.name_scope(name, default_name="divide"):
+    x1, x2 = binary_arguments_to_tensors(x1, x2)
+    return multiply(x1, reciprocal(x2), output_shape=output_shape)
+
+
+def slice(x, begin, size, slice_dim_name, name=None):  # pylint: disable=redefined-builtin
+  """Slice operation.
+
+  Args:
+    x: a list of Tensors
+    begin: integer, where to begin slicing from along the axis
+    size: integer, size to slice from axis.
+    slice_dim_name: string, dimension name of slicing axis.
+    name: an optional string
+  Returns:
+    a Tensor with shape extended by output_shape for the last axis.
+  """
+  return SliceOperation(
+      x, begin, size, slice_dim_name, name=name).outputs[0]
+
+
+def pad(x, paddings, dim_name, name=None):
+  """Slice operation.
+
+  Args:
+    x: a list of Tensors
+    paddings: list of integers of size 2, padding size before and after for dim.
+    dim_name: string, name for the padding dim
+    name: an optional string
+  Returns:
+    a Tensor with shape extended by output_shape for the last axis.
+  """
+  return PadOperation(
+      x, paddings, dim_name, name=name).outputs[0]
+
+
+def one_hot(indices, output_dim, on_value=1.0,
+            off_value=0.0, dtype=tf.float32, name=None):
+  """One hot operation.
+
+  Args:
+    indices: a Tensor
+    output_dim: a Dimension
+    on_value: Value taken when indices are on at a location, default 1
+    off_value: Value taken when indices are off at a location, default 0
+    dtype: a tf.DType
+    name: an optional string
+  Returns:
+    a Tensor with shape extended by output_dim for the last axis.
+  """
+  return OneHotOperation(
+      indices, output_dim, on_value, off_value, dtype, name=name).outputs[0]
+
+
+def gather(weights, indices, dim, output_shape=None):
+  """Shorthand for einsum([one_hot(indices, dim)], weights).
+
+  Args:
+    weights: a Tensor
+    indices: a Tensor with integer type
+    dim: a Dimension
+    output_shape: an optional mtf.TensorShape
+  Returns:
+    a Tensor
+  """
+  dim = convert_to_dimension(dim)
+  output_shape = convert_to_tensor_shape(output_shape)
+  if weights.dtype == tf.bool:
+    return cast(gather(to_float(weights), indices, dim, output_shape), tf.bool)
+  return einsum([one_hot(indices, dim, dtype=weights.dtype), weights],
+                output_shape=output_shape)
+
+
+def gradients(ys, xs, grad_ys=None):
+  """Compute gradients in dtf.
+
+  Args:
+    ys: a list of Tensors
+    xs: a list of Tensors
+    grad_ys: an optional list of Tensors
+
+  Returns:
+    grad_xs: a list of Tensors
+  """
+  graph = ys[0].graph
+  if not grad_ys:
+    grad_ys = [Constant(y.mesh, 1.0, y.shape, y.dtype).outputs[0] for y in ys]
+  # figure out what Tensors are downstream of xs
+  downstream = set(xs)
+  for op in graph.operations:
+    if op.has_gradient:
+      if set(op.inputs) & downstream:
+        downstream |= set(op.outputs)
+  tensor_to_gradient = dict(zip(ys, grad_ys))
+  for op in graph.operations[::-1]:
+    grad_outputs = [tensor_to_gradient.get(out) for out in op.outputs]
+    if op.has_gradient and any(grad_outputs) and (set(op.inputs) & downstream):
+      with tf.variable_scope(op.name + "/gradients"):
+        input_grads = op.gradient(grad_outputs)
+        for inp, grad in zip(op.inputs, input_grads):
+          if inp in downstream and grad is not None:
+            if inp in tensor_to_gradient:
+              tensor_to_gradient[inp] += grad
+            else:
+              tensor_to_gradient[inp] = grad
+  return [tensor_to_gradient.get(x, None) for x in xs]
+
+
+def _infer_binary_broadcast_shape(shape1, shape2, given_output_shape=None):
+  """Infer shape of the output of a binary op with broadcasting.
+
+  If the output shape is not given with given_output_shape, then we check
+  to see if one of the shapes is a subsequence of the other one, and we
+  return the one that is the supersequence.  Otherwise, we list the dimensions
+  of shape1, followed by all new dimensions in shape2.
+
+  Args:
+    shape1: a TensorShape
+    shape2: a TensorShape
+    given_output_shape: an optional TensorShape
+  Returns:
+    a TensorShape
+  """
+  shape1 = convert_to_tensor_shape(shape1)
+  shape2 = convert_to_tensor_shape(shape2)
+  given_output_shape = convert_to_tensor_shape(given_output_shape)
+  if given_output_shape is not None:
+    return given_output_shape
+  if is_subsequence(shape1.dims, shape2.dims):
+    return shape2
+  if is_subsequence(shape2.dims, shape1.dims):
+    return shape1
+  return TensorShape(
+      shape1.dims + [d for d in shape2.dims if d not in shape1.dims])
+
+
+def _expand_dims(x, input_shape, output_shape):
+  """Expand dimensions and transpose if necessary.
+
+  Args:
+    x: a tf.Tensor
+    input_shape: a TensorShape
+    output_shape: a TensorShape whose dimensions are a superset of
+      those in input_shape
+
+  Returns:
+    a tf.Tensor
+  """
+  verify_no_new_dims([output_shape], input_shape)
+  if input_shape == output_shape or input_shape.ndims == 0:
+    return x
+  perm = [input_shape.dims.index(d) for d in output_shape.dims
+          if d in input_shape.dims]
+  x = tf.transpose(x, perm)
+  for i, d in enumerate(output_shape.dims):
+    if d not in input_shape.dims:
+      x = tf.expand_dims(x, i)
+  return x
+
+
+def _einsum_equation(input_shapes, output_shape):
+  """Turn shapes into an einsum equation.
+
+  e.g. "ij,jk->ik"
+
+  Args:
+    input_shapes: a list of TensorShapes
+    output_shape: a TensorShape
+  Returns:
+    a string
+  """
+  ret = []
+  next_letter = ord("a")
+  dim_to_letter = {}
+  for shape_num, shape in enumerate(input_shapes + [output_shape]):
+    if shape_num == len(input_shapes):
+      ret.append("->")
+    elif shape_num > 0:
+      ret.append(",")
+    for d in shape.dims:
+      if d not in dim_to_letter:
+        dim_to_letter[d] = chr(next_letter)
+        next_letter += 1
+      ret.append(dim_to_letter[d])
+  ret = "".join(ret)
+  return ret
+
+
+def is_subsequence(short_seq, long_seq):
+  """Is short_seq a subsequence of long_seq."""
+  if not short_seq:
+    return True
+  pos = 0
+  for x in long_seq:
+    if pos == len(short_seq):
+      return True
+    if short_seq[pos] == x:
+      pos += 1
+  if pos == len(short_seq):
+    return True
+  return False
+
+
+def verify_no_new_dims(input_shapes, output_shape):
+  """Verifies that all dimensions in the output are in at least one input.
+
+  Args:
+    input_shapes: a list of TensorShapes
+    output_shape: a TensorShape
+  Raises:
+    ValueError: if there are new dimensions in the output.
+  """
+  all_input_dims = set(sum([s.dims for s in input_shapes], []))
+  all_output_dims = set(output_shape.dims)
+  if not all_output_dims.issubset(all_input_dims):
+    raise ValueError(
+        "No new dimensions allowed in output"
+        " input_shapes = %s output_shape= %s"
+        % ([s.dims for s in input_shapes], output_shape.dims))
+
+
+def pnum_to_processor_coordinates(mesh_shape, pnum):
+  """Coordinates of a processor in the mesh.
+
+  Args:
+    mesh_shape: a list of integers
+    pnum: an integer less than len(mesh_shape)
+
+  Returns:
+    a list of integers with length len(mesh_shape)
+  """
+  ret = []
+  for dimsize in mesh_shape[::-1]:
+    ret.append(pnum % dimsize)
+    pnum //= dimsize
+  return ret[::-1]
+
+
+def processor_coordinates_to_pnum(mesh_shape, coord):
+  """Inverse of pnum_to_processor_coordinates.
+
+  Args:
+    mesh_shape: a list of integers
+    coord: a list of integers with length len(mesh_shape)
+
+  Returns:
+    an integer less than len(mesh_shape)
+  """
+  ret = 0
+  multiplier = 1
+  for c, d in zip(coord[::-1], mesh_shape[::-1]):
+    ret += multiplier * c
+    multiplier *= d
+  return ret
+
+
+def pnum_to_group(mesh_shape, group_dims, pnum):
+  """Group number for grouped allreduce.
+
+  Args:
+    mesh_shape: a list of integers
+    group_dims: a list of integers (the dimensions reduced over)
+    pnum: an integer
+
+  Returns:
+    an integer
+  """
+  coord = pnum_to_processor_coordinates(mesh_shape, pnum)
+  remaining_shape = [d for i, d in enumerate(mesh_shape) if i not in group_dims]
+  remaining_coord = [d for i, d in enumerate(coord) if i not in group_dims]
+  return processor_coordinates_to_pnum(remaining_shape, remaining_coord)
+
+
+def processor_groups(mesh_shape, group_dims):
+  """Groups of processors which differ only in the given dimensions.
+
+  Args:
+    mesh_shape: a list of integers
+    group_dims: a list of integers
+
+  Returns:
+    a list of lists of integers (processor numbers)
+  """
+  group_numbers = [
+      pnum_to_group(mesh_shape, group_dims, pnum)
+      for pnum in xrange(list_product(mesh_shape))]
+  ret = []
+  for pnum, g in enumerate(group_numbers):
+    while len(ret) <= g:
+      ret.append([])
+    ret[g].append(pnum)
+  return ret
+
+
+def list_product(l):
+  return reduce(mul, l, 1)
+
+
+def log_softmax(x, reduced_dim, name=None):
+  """log(softmax(x)).
+
+  Args:
+    x: a Tensor whose shape contains vocab_dim
+    reduced_dim: a Dimension
+    name: an optional string
+
+  Returns:
+    a Tensor with the same shape as x
+  """
+  reduced_dim = convert_to_dimension(reduced_dim)
+  with tf.variable_scope(name, default_name="log_softmax"):
+    reduced_shape = x.shape - reduced_dim
+    max_logit = reduce_max(stop_gradient(x), output_shape=reduced_shape)
+    x -= max_logit
+    exp_x = exp(x)
+    sum_exp_x = reduce_sum(exp_x, output_shape=reduced_shape)
+    log_denom = log(sum_exp_x)
+    return x - log_denom
+
+
+def softmax(x, reduced_dim, name=None):
+  with tf.variable_scope(name, default_name="softmax"):
+    return exp(log_softmax(x, reduced_dim))
+
+
+def range(mesh, dim, dtype, name=None):  # pylint: disable=redefined-builtin
+  """Create a 1d mesh tensor with a range from [0, dim.size).
+
+  Args:
+    mesh: a Mesh
+    dim: a Dimension
+    dtype: a tf.DType
+    name: an optional string
+
+  Returns:
+    a Tensor
+  """
+  dim = convert_to_dimension(dim)
+  with tf.variable_scope(name, default_name="range"):
+    return infeed(
+        mesh, tf.range(dim.size, dtype=dtype), shape=TensorShape([dim]))
+
+
+def pretty_print_counters(counters):
+  """print counters hierarchically.
+
+  Each counter is a pair of a string and a number.
+  The string can have slashes, meaning that the number also counts towards
+  each prefix.  e.g.  "parameters/trainable" counts towards both "parameters"
+  and "parameters/trainable".
+
+  Args:
+    counters: a list of (string, number) pairs
+
+  Returns:
+    a string
+  """
+  totals = collections.defaultdict(int)
+  for (name, val) in counters:
+    prefixes = [name[:i] for i in xrange(len(name)) if name[i] == "/"] + [name]
+    for p in prefixes:
+      totals[p] += val
+  parts = []
+  for name, val in sorted(six.iteritems(totals)):
+    parts.append(" " * name.count("/") + "%s: %.3g" % (name, val))
+  return "\n".join(parts)
+
+
+def parse_mesh_shape(mesh_shape):
+  """Parase a string to a list of integers.
+
+  All non-digits are taken as delimeters
+
+  Args:
+    mesh_shape: a string or a list of integers
+  Returns:
+    a list of integers
+  """
+  if isinstance(mesh_shape, list):
+    return mesh_shape
+  return [int(x) for x in re.sub("[^0-9]", " ", mesh_shape).split()]
+
+
+def parse_layout(layout_string):
+  r"""Parase a string specifying a layout.
+
+  The layout_string is a list of name, integer pairs.
+  Each pair is separated by a colon, and the delimeters between paris are any
+  of " ,.;"
+
+  e.g. "batch:0 vocab:1 filter_size:1 heads:1"
+
+  Args:
+    layout_string: a string or a dictionary
+  Returns:
+    a dictionary from string to int
+  """
+  if isinstance(layout_string, dict):
+    return layout_string
+  ret = {}
+  for s in re.sub("[,.;]", " ", layout_string).split():
+    dim_name, mesh_axis = s.split(":")
+    ret[dim_name] = int(mesh_axis)
+  return ret
+
+
+def parallel(devices, fn, *args, **kwargs):
+  """Call a function once on each device.
+
+  Args:
+    devices: a list of n devices
+    fn: a function
+    *args: arguments, each of which is a list of length n
+    **kwargs: keyword-args, each of which is a list of length n
+  Returns:
+    a list of length n
+  Raises:
+    ValueError: if the arguments are not all lists of length n
+  """
+  if not isinstance(devices, list):
+    raise ValueError("devices must be a list")
+  for x in list(args) + list(six.itervalues(kwargs)):
+    if not isinstance(x, list) or len(x) != len(devices):
+      raise ValueError(
+          "Argument not a list with same length as devices "
+          "arg=%s devices=%s %s %s" % (x, devices, len(x), len(devices)))
+  ret = []
+  for i, device in enumerate(devices):
+    with tf.device(device):
+      with tf.variable_scope("parallel_%d" % i):
+        my_args = [x[i] for x in args]
+        my_kwargs = {k: v[i] for k, v in six.iteritems(kwargs)}
+        ret.append(fn(*my_args, **my_kwargs))
+  return ret
+
+
+def transpose_list_of_lists(lol):
+  """Transpose a list of equally-sized python lists.
+
+  Args:
+    lol: a list of lists
+  Returns:
+    a list of lists
+  Raises:
+    ValueError: if list is empty
+  """
+  if not lol:
+    raise ValueError("cannot transpose the empty list")
+  return [list(x) for x in zip(*lol)]
+
+
+def binary_reduction_fn(reduction_fn_string):
+  if reduction_fn_string == "SUM":
+    return tf.add
+  elif reduction_fn_string == "MAX":
+    return tf.maximum
+  elif reduction_fn_string == "MIN":
+    return tf.minimum
+  else:
+    raise ValueError("Unknown reduction_fn_string %s" % reduction_fn_string)
+
+
+def reduction_fn(reduction_fn_string):
+  if reduction_fn_string == "SUM":
+    return tf.reduce_sum
+  elif reduction_fn_string == "MAX":
+    return tf.reduce_max
+  elif reduction_fn_string == "MIN":
+    return tf.reduce_min
+  else:
+    raise ValueError("Unknown reduction_fn_string %s" % reduction_fn_string)
+
+
+class MtfCheckpointSaverListener(tf.train.CheckpointSaverListener):
+  """Copy slices to masters before saving."""
+
+  def __init__(self, lowering):
+    self._op = lowering.copy_slices_to_masters()
+
+  def begin(self):
+    # You can add ops to the graph here.
+    tf.logging.info("Starting the session.")
+
+  def before_save(self, session, global_step_value):
+    # assigns
+    tf.logging.info("Before Save.")
+    session.run(self._op)
+    tf.logging.info("About to write a checkpoint")
+
+  def after_save(self, session, global_step_value):
+    tf.logging.info("Done writing checkpoint.")
+
+  def end(self, session, global_step_value):
+    tf.logging.info("Done with the session.")
+
+
+class MtfRestoreHook(tf.train.SessionRunHook):
+  """Copy masters to slices after restoring."""
+
+  def __init__(self, lowering):
+    self._lowering = lowering
+
+  def begin(self):
+    self._op = self._lowering.copy_masters_to_slices()
+
+  def after_create_session(self, session, coord):
+    session.run(self._op)
+
+
+class RandomOperation(Operation):
+  """Random operation such as tf.random_uniform."""
+
+  def __init__(self, mesh, shape, tf_fn, **kwargs):
+    super(RandomOperation, self).__init__(
+        [], mesh=mesh, name=kwargs.get("name", "random"))
+    self._tf_fn = tf_fn
+    self._kwargs = kwargs
+    self._outputs = [Tensor(self, shape, kwargs.get("dtype", tf.float32))]
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    output_shape = self.outputs[0].shape
+    lowering.tensors[self.outputs[0]] = (
+        mesh_impl.random(output_shape, self._tf_fn, self._kwargs))
+
+
+def random_uniform(mesh, shape, **kwargs):
+  """Random uniform.
+
+  Args:
+    mesh: a Mesh
+    shape: a TensorShape
+    **kwargs: keyword args for tf.random_uniform, except seed
+
+  Returns:
+    a Tensor
+  """
+  shape = convert_to_tensor_shape(shape)
+  return RandomOperation(mesh, shape, tf.random_uniform, **kwargs).outputs[0]
+
+
+def dropout(x, keep_prob, noise_shape=None, name=None):
+  """Dropout layer.
+
+  Args:
+    x: a Tensor
+    keep_prob: a float between 0.0 and 1.0
+    noise_shape: an optional TensorShape (a subset of x.shape)
+    name: an optional string
+
+  Returns:
+    a Tensor
+  """
+  noise_shape = convert_to_tensor_shape(noise_shape)
+  if noise_shape is None:
+    noise_shape = x.shape
+  with tf.variable_scope(name, default_name="dropout"):
+    if keep_prob == 1.0:
+      return x
+    noise = cast(less(random_uniform(
+        x.mesh, noise_shape, dtype=x.dtype), keep_prob), x.dtype)
+    noise /= keep_prob
+    return x * noise
+
+
+def _cumprod(l):
+  """Cumulative product of a list.
+
+  Args:
+    l: a list of integers
+  Returns:
+    a list with one more element (starting with 1)
+  """
+  ret = [1]
+  for item in l:
+    ret.append(ret[-1] * item)
+  return ret
+
+
+def log_variable_sizes(var_list, tag, verbose=True):
+  """Log the sizes and shapes of variables, and the total size.
+
+  Args:
+    var_list: a list of variables; defaults to trainable_variables
+    tag: a string; defaults to "Trainable Variables"
+    verbose: bool, if True, log every weight; otherwise, log total size only.
+  """
+  if not var_list:
+    return
+
+  name_to_var = {v.name: v for v in var_list}
+  total_size = 0
+  for v_name in sorted(list(name_to_var)):
+    v = name_to_var[v_name]
+    v_size = v.shape.size
+    if verbose:
+      tf.logging.info("Weight    %s\tshape    %s\tsize    %d",
+                      v.name.ljust(80),
+                      str(v.shape).ljust(30), v_size)
+    total_size += v_size
+  tf.logging.info("%s Total size: %d", tag, total_size)
+
+
+class WhileLoopOperation(Operation):
+  """While loop."""
+
+  def __init__(self, cond_fn, body_fn, inputs,
+               tf_kwargs=None, name="while_loop"):
+    super(WhileLoopOperation, self).__init__(
+        inputs, mesh=inputs[0].mesh, name=name)
+    self._cond_fn = cond_fn
+    self._body_fn = body_fn
+    self._tf_kwargs = tf_kwargs or {}
+    assert not self._tf_kwargs.get("back_prop", False)
+    ops = self.graph.operations
+    before = len(ops)
+    def make_placeholders(name):
+      return [Tensor(self, t.shape, t.dtype, name="%s_%d" % (name, i))
+              for i, t in enumerate(inputs)]
+    self._cond_inputs = make_placeholders("cond_input")
+    self._cond_output = self._cond_fn(*self._cond_inputs)
+    self._cond_ops = ops[before:]
+    del ops[before:]
+    self._body_inputs = make_placeholders("body_input")
+    self._body_outputs = self._body_fn(*self._body_inputs)
+    for (i, (inp, body_out)) in enumerate(zip(inputs, self._body_outputs)):
+      if inp.shape != body_out.shape:
+        raise ValueError(
+            "shape mismatch i=%d inp=%s body_out=%s" % (i, inp, body_out))
+    self._body_ops = ops[before:]
+    del ops[before:]
+    self._outputs = make_placeholders("output")
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    def tf_cond_fn(*tf_inputs):
+      for tf_inp, mtf_inp in zip(tf_inputs, self._cond_inputs):
+        lowering.tensors[mtf_inp] = mesh_impl.LaidOutTensor(tf_inp)
+      for op in self._cond_ops:
+        with tf.name_scope(op.name):
+          op.lower(lowering)
+      lowered_output = lowering.tensors[self._cond_output]
+      ret = lowered_output.to_laid_out_tensor().tensor_list[0]
+      return ret
+
+    def tf_body_fn(*tf_inputs):
+      for tf_inp, mtf_inp in zip(tf_inputs, self._body_inputs):
+        lowering.tensors[mtf_inp] = mesh_impl.LaidOutTensor(tf_inp)
+      for op in self._body_ops:
+        with tf.name_scope(op.name):
+          op.lower(lowering)
+      return [
+          lowering.tensors[mtf_out].to_laid_out_tensor().tensor_list
+          for mtf_out in self._body_outputs]
+
+    lowered_inputs = [
+        lowering.tensors[t].to_laid_out_tensor().tensor_list
+        for t in self.inputs]
+
+    tf_outs = tf.while_loop(tf_cond_fn,
+                            tf_body_fn,
+                            lowered_inputs,
+                            back_prop=False,
+                            **self._tf_kwargs)
+    for tf_out, mtf_out in zip(tf_outs, self._outputs):
+      lowering.tensors[mtf_out] = mesh_impl.LaidOutTensor(tf_out)
+
+
+def while_loop(cond_fn, body_fn, inputs, num_loop_vars=None, **kwargs):
+  """While Loop.
+
+  num_loop_vars is a hack for the multi-gpu setup.  In this case, loops
+  are generally slow, as all loop variables are placed on device.  By setting
+  num_loop_vars=k, then all of the loop variables except for the first k
+  are handled as mtf Variables instead of loop variables, using explicit
+  updates and control dependencies.  In this case, we only return the
+  first num_loop_vars outputs.  Do not use this option on TPU, since it
+  is unnecessary and also produces incorrect results, since xla does not
+  respect control dependencies.
+
+  Args:
+    cond_fn: a function from n Tensors to scalar boolean Tensor
+    body_fn: a function from n Tensors to n Tensors
+    inputs: a list of n Tensors
+    num_loop_vars: an optional integer.
+    **kwargs: additional kwargs passed to tf.while_loop
+
+  Returns:
+    a list of n Tensors.
+  """
+  if num_loop_vars is None:
+    return WhileLoopOperation(cond_fn, body_fn, inputs, kwargs).outputs
+  # Turn all loop vars except for the first ones into non-loop vars.
+  # see comments in docstring.
+  assert num_loop_vars > 0
+  extra_inputs = inputs[num_loop_vars:]
+  my_vars = tuple([get_variable(
+      x.mesh, "loop_var_%d" % i,
+      x.shape, initializer=tf.zeros_initializer(),
+      dtype=x.dtype,
+      collections=[tf.GraphKeys.LOCAL_VARIABLES])
+                   for i, x in enumerate(extra_inputs)])
+  first_input = depend(
+      inputs[0], [assign(var, x) for var, x in zip(my_vars, extra_inputs)])
+  inputs = [first_input] + inputs[1:num_loop_vars]
+  def my_cond_fn(*inputs):
+    return cond_fn(*(inputs + my_vars))
+  def my_body_fn(*inputs):
+    outputs = tuple(body_fn(*(inputs + my_vars)))
+    extra_outputs = outputs[num_loop_vars:]
+    first_output = depend(
+        outputs[0], [assign(var, x) for var, x in zip(my_vars, extra_outputs)])
+    outputs = (first_output,) + outputs[1:num_loop_vars]
+    return outputs
+  return WhileLoopOperation(
+      my_cond_fn, my_body_fn, inputs, kwargs).outputs
+
+
+def where(condition, if_true, if_false):
+  dtype = if_true.dtype
+  return (
+      if_true * cast(condition, dtype) +
+      if_false * cast(logical_not(condition), dtype))
diff --git a/tensor2tensor/mesh_tensorflow/mnist.py b/tensor2tensor/mesh_tensorflow/mnist.py
new file mode 100644
index 000000000..0a0763872
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mnist.py
@@ -0,0 +1,217 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Mnist using mesh-tensrflow and tf.Estimator."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import google3
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+from tensor2tensor.mesh_tensorflow import mnist_dataset as dataset
+from tensor2tensor.mesh_tensorflow import mtf_layers
+from tensor2tensor.mesh_tensorflow import mtf_optimize
+from tensor2tensor.mesh_tensorflow import placement_mesh_impl
+import tensorflow as tf
+
+
+tf.flags.DEFINE_string("data_dir", "/tmp/mnist_data",
+                       "Path to directory containing the MNIST dataset")
+tf.flags.DEFINE_string("model_dir", "/tmp/mnist_model", "Estimator model_dir")
+tf.flags.DEFINE_integer("batch_size", 200,
+                        "Mini-batch size for the training. Note that this "
+                        "is the global batch size and not the per-shard batch.")
+tf.flags.DEFINE_integer("hidden_size", 512, "Size of each hidden layer.")
+tf.flags.DEFINE_integer("train_epochs", 40, "Total number of training epochs.")
+tf.flags.DEFINE_integer("epochs_between_evals", 1,
+                        "# of epochs between evaluations.")
+tf.flags.DEFINE_integer("eval_steps", 0,
+                        "Total number of evaluation steps. If `0`, evaluation "
+                        "after training is skipped.")
+tf.flags.DEFINE_string("mesh_shape", "2;2", "mesh shape")
+tf.flags.DEFINE_string("layout", "batch:0;hidden1:1", "computation layout")
+
+FLAGS = tf.flags.FLAGS
+
+
+def mnist_model(image, labels, mesh):
+  """The model.
+
+  Args:
+    image: tf.Tensor with shape [batch, 28*28]
+    labels: a tf.Tensor with shape [batch] and dtype tf.int32
+    mesh: a mtf.Mesh
+
+  Returns:
+    logits: a tf.Tensor with shape [batch, 10]
+    loss: a mtf.Tensor with shape []
+  """
+  batch_dim = mtf.Dimension("batch", FLAGS.batch_size)
+  rows_dim = mtf.Dimension("rows", 28)
+  cols_dim = mtf.Dimension("cols", 28)
+  classes_dim = mtf.Dimension("classes", 10)
+  hidden_dim1 = mtf.Dimension("hidden1", FLAGS.hidden_size)
+  hidden_dim2 = mtf.Dimension("hidden2", FLAGS.hidden_size)
+
+  x = mtf.infeed(mesh, tf.reshape(image, [-1, 28, 28]),
+                 mtf.TensorShape([batch_dim, rows_dim, cols_dim]))
+  h1 = mtf_layers.dense(
+      x, hidden_dim1, reduced_dims=[rows_dim, cols_dim],
+      activation=mtf.relu, name="hidden1")
+  h2 = mtf_layers.dense(
+      h1, hidden_dim2, activation=mtf.relu, name="hidden2")
+  logits = mtf_layers.dense(h2, classes_dim, name="logits")
+  if labels is None:
+    loss = None
+  else:
+    labels = mtf.infeed(mesh, labels, mtf.TensorShape([batch_dim]))
+    loss = mtf_layers.softmax_cross_entropy_with_logits(
+        logits, mtf.one_hot(labels, classes_dim), classes_dim)
+    loss = mtf.reduce_mean(loss)
+  return logits, loss
+
+
+def model_fn(features, labels, mode, params):
+  """The model_fn argument for creating an Estimator."""
+  tf.logging.info("features = %s labels = %s mode = %s params=%s" %
+                  (features, labels, mode, params))
+  global_step = tf.train.get_global_step()
+  graph = mtf.Graph()
+  mesh = mtf.Mesh(graph, "my_mesh")
+  logits, loss = mnist_model(features, labels, mesh)
+  mesh_shape = mtf.parse_mesh_shape(FLAGS.mesh_shape)
+  mesh_size = mtf.list_product(mesh_shape)
+  mesh_devices = [""] * mesh_size
+  mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+      mesh_shape, mtf.parse_layout(FLAGS.layout), mesh_devices)
+
+  if mode == tf.estimator.ModeKeys.TRAIN:
+    var_grads = mtf.gradients(
+        [loss], [v.outputs[0] for v in graph.trainable_variables])
+    optimizer = mtf_optimize.AdafactorOptimizer()
+    update_ops = []
+    for grad, var in zip(var_grads, graph.trainable_variables):
+      update_ops.extend(optimizer.apply_grad(grad, var))
+
+  lowering = mtf.Lowering(graph, {mesh: mesh_impl})
+  restore_hook = mtf.MtfRestoreHook(lowering)
+
+  tf_logits = lowering.outfeed(logits)
+  if mode != tf.estimator.ModeKeys.PREDICT:
+    tf_loss = lowering.outfeed(loss)
+    tf.summary.scalar("loss", tf_loss)
+
+  if mode == tf.estimator.ModeKeys.TRAIN:
+    tf_update_ops = [lowering.lowered_operation(op) for op in update_ops]
+    tf_update_ops.append(tf.assign_add(global_step, 1))
+    train_op = tf.group(tf_update_ops)
+    saver = tf.train.Saver(
+        tf.global_variables(),
+        sharded=True,
+        max_to_keep=10,
+        keep_checkpoint_every_n_hours=2,
+        defer_build=False, save_relative_paths=True)
+    tf.add_to_collection(tf.GraphKeys.SAVERS, saver)
+    saver_listener = mtf.MtfCheckpointSaverListener(lowering)
+    saver_hook = tf.train.CheckpointSaverHook(
+        FLAGS.model_dir,
+        save_steps=1000,
+        saver=saver,
+        listeners=[saver_listener])
+
+    accuracy = tf.metrics.accuracy(
+        labels=labels, predictions=tf.argmax(tf_logits, axis=1))
+
+    # Name tensors to be logged with LoggingTensorHook.
+    tf.identity(tf_loss, "cross_entropy")
+    tf.identity(accuracy[1], name="train_accuracy")
+
+    # Save accuracy scalar to Tensorboard output.
+    tf.summary.scalar("train_accuracy", accuracy[1])
+
+    # restore_hook must come before saver_hook
+    return tf.estimator.EstimatorSpec(
+        tf.estimator.ModeKeys.TRAIN, loss=tf_loss, train_op=train_op,
+        training_chief_hooks=[restore_hook, saver_hook])
+
+  if mode == tf.estimator.ModeKeys.PREDICT:
+    predictions = {
+        "classes": tf.argmax(tf_logits, axis=1),
+        "probabilities": tf.nn.softmax(tf_logits),
+    }
+    return tf.estimator.EstimatorSpec(
+        mode=tf.estimator.ModeKeys.PREDICT,
+        predictions=predictions,
+        prediction_hooks=[restore_hook],
+        export_outputs={
+            "classify": tf.estimator.export.PredictOutput(predictions)
+        })
+  if mode == tf.estimator.ModeKeys.EVAL:
+    return tf.estimator.EstimatorSpec(
+        mode=tf.estimator.ModeKeys.EVAL,
+        loss=tf_loss,
+        evaluation_hooks=[restore_hook],
+        eval_metric_ops={
+            "accuracy":
+            tf.metrics.accuracy(
+                labels=labels, predictions=tf.argmax(tf_logits, axis=1)),
+        })
+
+
+def run_mnist():
+  """Run MNIST training and eval loop."""
+  mnist_classifier = tf.estimator.Estimator(
+      model_fn=model_fn,
+      model_dir=FLAGS.model_dir,
+      params={
+          "model_dir": FLAGS.model_dir,
+          "mesh_shape": "2.2",
+          "layout": "batch.0;hidden1.1",
+          "batch_size": FLAGS.batch_size,
+          "hidden_size": 512,
+      })
+
+  # Set up training and evaluation input functions.
+  def train_input_fn():
+    """Prepare data for training."""
+
+    # When choosing shuffle buffer sizes, larger sizes result in better
+    # randomness, while smaller sizes use less memory. MNIST is a small
+    # enough dataset that we can easily shuffle the full epoch.
+    ds = dataset.train(FLAGS.data_dir)
+    ds = ds.cache().shuffle(buffer_size=50000).batch(FLAGS.batch_size)
+
+    # Iterate through the dataset a set number (`epochs_between_evals`) of times
+    # during each training session.
+    ds = ds.repeat(FLAGS.epochs_between_evals)
+    return ds
+
+  def eval_input_fn():
+    return dataset.test(FLAGS.data_dir).batch(
+        FLAGS.batch_size).make_one_shot_iterator().get_next()
+
+  # Train and evaluate model.
+  for _ in range(FLAGS.train_epochs // FLAGS.epochs_between_evals):
+    mnist_classifier.train(input_fn=train_input_fn, hooks=None)
+    eval_results = mnist_classifier.evaluate(input_fn=eval_input_fn)
+    print("\nEvaluation results:\n\t%s\n" % eval_results)
+
+
+def main(_):
+  run_mnist()
+
+
+if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
+  tf.app.run()
diff --git a/tensor2tensor/mesh_tensorflow/mnist_dataset.py b/tensor2tensor/mesh_tensorflow/mnist_dataset.py
new file mode 100644
index 000000000..99bb0e94c
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mnist_dataset.py
@@ -0,0 +1,131 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+#  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.
+"""tf.data.Dataset interface to the MNIST dataset."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gzip
+import os
+import shutil
+import tempfile
+
+import google3
+import numpy as np
+from six.moves import urllib
+import tensorflow as tf
+
+
+def read32(bytestream):
+  """Read 4 bytes from bytestream as an unsigned 32-bit integer."""
+  dt = np.dtype(np.uint32).newbyteorder('>')
+  return np.frombuffer(bytestream.read(4), dtype=dt)[0]
+
+
+def check_image_file_header(filename):
+  """Validate that filename corresponds to images for the MNIST dataset."""
+  with tf.gfile.Open(filename, 'rb') as f:
+    magic = read32(f)
+    read32(f)  # num_images, unused
+    rows = read32(f)
+    cols = read32(f)
+    if magic != 2051:
+      raise ValueError('Invalid magic number %d in MNIST file %s' % (magic,
+                                                                     f.name))
+    if rows != 28 or cols != 28:
+      raise ValueError(
+          'Invalid MNIST file %s: Expected 28x28 images, found %dx%d' %
+          (f.name, rows, cols))
+
+
+def check_labels_file_header(filename):
+  """Validate that filename corresponds to labels for the MNIST dataset."""
+  with tf.gfile.Open(filename, 'rb') as f:
+    magic = read32(f)
+    read32(f)  # num_items, unused
+    if magic != 2049:
+      raise ValueError('Invalid magic number %d in MNIST file %s' % (magic,
+                                                                     f.name))
+
+
+def download(directory, filename):
+  """Download (and unzip) a file from the MNIST dataset if not already done."""
+  filepath = os.path.join(directory, filename)
+  if tf.gfile.Exists(filepath):
+    return filepath
+  if not tf.gfile.Exists(directory):
+    tf.gfile.MakeDirs(directory)
+  url = 'http://yann.lecun.com/exdb/mnist/' + filename + '.gz'
+  _, zipped_filepath = tempfile.mkstemp(suffix='.gz')
+  print('Downloading %s to %s' % (url, zipped_filepath))
+  urllib.request.urlretrieve(url, zipped_filepath)
+  with gzip.open(zipped_filepath, 'rb') as f_in, \
+      tf.gfile.Open(filepath, 'wb') as f_out:
+    shutil.copyfileobj(f_in, f_out)
+  os.remove(zipped_filepath)
+  return filepath
+
+
+def dataset(directory, images_file, labels_file):
+  """Download and parse MNIST dataset."""
+
+  images_file = download(directory, images_file)
+  labels_file = download(directory, labels_file)
+
+  check_image_file_header(images_file)
+  check_labels_file_header(labels_file)
+
+  def decode_image(image):
+    # Normalize from [0, 255] to [0.0, 1.0]
+    image = tf.decode_raw(image, tf.uint8)
+    image = tf.cast(image, tf.float32)
+    image = tf.reshape(image, [784])
+    return image / 255.0
+
+  def decode_label(label):
+    label = tf.decode_raw(label, tf.uint8)  # tf.string -> [tf.uint8]
+    label = tf.reshape(label, [])  # label is a scalar
+    return tf.to_int32(label)
+
+  images = tf.data.FixedLengthRecordDataset(
+      images_file, 28 * 28, header_bytes=16).map(decode_image)
+  labels = tf.data.FixedLengthRecordDataset(
+      labels_file, 1, header_bytes=8).map(decode_label)
+  return tf.data.Dataset.zip((images, labels))
+
+
+def train(directory):
+  """tf.data.Dataset object for MNIST training data."""
+  return dataset(directory, 'train-images-idx3-ubyte',
+                 'train-labels-idx1-ubyte')
+
+
+def test(directory):
+  """tf.data.Dataset object for MNIST test data."""
+  return dataset(directory, 't10k-images-idx3-ubyte', 't10k-labels-idx1-ubyte')
diff --git a/tensor2tensor/mesh_tensorflow/mtf_beam_search.py b/tensor2tensor/mesh_tensorflow/mtf_beam_search.py
new file mode 100644
index 000000000..f0f42b7a0
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_beam_search.py
@@ -0,0 +1,573 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Implementation of beam search with penalties."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+import tensorflow as tf
+
+# Assuming EOS_ID is 1
+EOS_ID = 1
+# Default value for INF
+INF = 1. * 1e7
+
+
+def _concat_equal_sizes(xs, dim, new_dim_name):
+  axis = xs[0].shape.dims.index(dim)
+  ret = mtf.stack(xs, "tmp_concat", axis)
+  new_shape = mtf.TensorShape(
+      xs[0].shape.dims[:axis]
+      + [mtf.Dimension(new_dim_name, dim.size * len(xs))]
+      + xs[0].shape.dims[axis + 1:])
+  return mtf.reshape(ret, new_shape)
+
+
+def _expand_to_beam_size(tensor, beam_size):
+  """Tiles a given tensor by beam_size.
+
+  Args:
+    tensor: tensor to tile [batch_size, ...]
+    beam_size: How much to tile the tensor by.
+
+  Returns:
+    Tiled tensor [batch_size, beam_size, ...]
+  """
+  tensor = tf.expand_dims(tensor, axis=1)
+  tile_dims = [1] * tensor.shape.ndims
+  tile_dims[1] = beam_size
+
+  return tf.tile(tensor, tile_dims)
+
+
+def get_state_shape_invariants(tensor):
+  """Returns the shape of the tensor but sets middle dims to None."""
+  shape = tensor.shape.as_list()
+  for i in range(1, len(shape) - 1):
+    shape[i] = None
+  return tf.TensorShape(shape)
+
+
+def compute_batch_indices(batch_size, beam_size):
+  """Computes the i'th coordinate that contains the batch index for gathers.
+
+  Batch pos is a tensor like [[0,0,0,0,],[1,1,1,1],..]. It says which
+  batch the beam item is in. This will create the i of the i,j coordinate
+  needed for the gather.
+
+  Args:
+    batch_size: Batch size
+    beam_size: Size of the beam.
+  Returns:
+    batch_pos: [batch_size, beam_size] tensor of ids
+  """
+  batch_pos = tf.range(batch_size * beam_size) // beam_size
+  batch_pos = tf.reshape(batch_pos, [batch_size, beam_size])
+  return batch_pos
+
+
+def compute_topk_scores_and_seq(sequences, scores, scores_to_gather, flags,
+                                beam_dim, prefix="default",
+                                states=None):
+  """Given sequences and scores, will gather the top k=beam size sequences.
+
+  This function is used to grow alive, and finished. It takes sequences,
+  scores, and flags, and returns the top k from sequences, scores_to_gather,
+  and flags based on the values in scores.
+
+  This method permits easy introspection using tfdbg.  It adds three named ops
+  that are prefixed by `prefix`:
+    - _topk_seq: the tensor for topk_seq returned by this method.
+    - _topk_flags: the tensor for topk_finished_flags returned by this method.
+    - _topk_scores: the tensor for tokp_gathered_scores returned by this method.
+
+  Args:
+    sequences: Tensor of sequences that we need to gather from.
+      [batch_size, beam_size, seq_length]
+    scores: Tensor of scores for each sequence in sequences.
+      [batch_size, beam_size]. We will use these to compute the topk.
+    scores_to_gather: Tensor of scores for each sequence in sequences.
+      [batch_size, beam_size]. We will return the gathered scores from here.
+      Scores to gather is different from scores because for grow_alive, we will
+      need to return log_probs, while for grow_finished, we will need to return
+      the length penalized scores.
+    flags: Tensor of bools for sequences that say whether a sequence has reached
+      EOS or not
+    beam_dim: mtf.Dimension
+    prefix: an optional string
+    states: an optional list of mtf.Tensor
+  Returns:
+    Tuple of
+    (topk_seq [batch_size, beam_size, decode_length],
+     topk_gathered_scores [batch_size, beam_size],
+     topk_finished_flags[batch_size, beam_size],
+     topk_gathered_states)
+  """
+  unused_batch_dim, old_beam_dim, unused_length_dim = sequences.shape.dims
+  topk_indices, _ = mtf.top_k(scores, old_beam_dim, beam_dim)
+
+  # Gather up the highest scoring sequences.
+  # For each operation added, give it
+  # a concrete name to simplify observing these operations with tfdbg.
+  # Clients can capture these tensors by watching these node names.
+  def gather(tensor, name):
+    with tf.name_scope(prefix + name):
+      output_shape = mtf.TensorShape(
+          [beam_dim if d == old_beam_dim else d for d in tensor.shape.dims])
+      return mtf.gather(
+          tensor, topk_indices, old_beam_dim, output_shape=output_shape)
+  topk_seq = gather(sequences, "_seq")
+  topk_flags = gather(flags, "_flags")
+  topk_gathered_scores = gather(scores_to_gather, "_scores")
+  if states is None:
+    topk_gathered_states = None
+  else:
+    topk_gathered_states = [gather(state, "_topk_states") for state in states]
+  return topk_seq, topk_gathered_scores, topk_flags, topk_gathered_states
+
+
+def beam_search(logits_fn,
+                initial_ids,
+                alpha,
+                states=None,
+                eos_id=EOS_ID,
+                stop_early=True,
+                decode_length=None,
+                use_tpu=True):
+  """Beam search with length penalties.
+
+  Requires a function that can take the currently decoded symbols and return
+  the logits for the next symbol. The implementation is inspired by
+  https://arxiv.org/abs/1609.08144.
+
+  When running, the beam search steps can be visualized by using tfdbg to watch
+  the operations generating the output ids for each beam step.  These operations
+  have the pattern:
+    (alive|finished)_topk_(seq,scores)
+
+  Operations marked `alive` represent the new beam sequences that will be
+  processed in the next step.  Operations marked `finished` represent the
+  completed beam sequences, which may be padded with 0s if no beams finished.
+
+  Operations marked `seq` store the full beam sequence for the time step.
+  Operations marked `scores` store the sequence's final log scores.
+
+  The beam search steps will be processed sequentially in order, so when
+  capturing observed from these operations, tensors, clients can make
+  assumptions about which step is being recorded.
+
+  WARNING: Assumes 2nd dimension of tensors in `states` and not invariant, this
+  means that the shape of the 2nd dimension of these tensors will not be
+  available (i.e. set to None) inside logits_fn.
+
+  Args:
+    logits_fn: Interface to the model, to provide logits.
+        Shoud take:
+          step_num - mtf Scalar
+          ids - mtf Tensor with shape [batch, beam, length]
+        Should return:
+          logits - [batch, beam, vocab_size]
+    initial_ids: a mtf.Tensor with shape [batch_dim, beam_dim, length_dim])
+    alpha: alpha for length penalty.
+    states: list of mtf.Tensor
+    eos_id: ID for end of sentence.
+    stop_early: a boolean - stop once best sequence is provably determined.
+    decode_length: a mtf Scalar of dtype tf.int32 - maximum length of decodes
+    use_tpu: a boolean
+  Returns:
+    Tuple of
+    (decoded beams [batch, beam, length]
+     decoding probabilities [batch, beam_size])
+  """
+  batch_dim, beam_dim, length_dim = initial_ids.shape.dims
+  mesh = initial_ids.mesh
+
+  batch_by_beam = mtf.TensorShape([batch_dim, beam_dim])
+  initial_log_probs = mtf.broadcast(
+      mtf.one_hot(
+          mtf.constant(mesh, 0, dtype=tf.int32),
+          beam_dim,
+          on_value=0.0,
+          off_value=-INF),
+      batch_by_beam)
+
+  length_scalar = mtf.constant(mesh, length_dim.size, dtype=tf.int32)
+  if decode_length is None:
+    decode_length = length_scalar
+  else:
+    decode_length = mtf.minimum(decode_length, length_scalar)
+
+  alive_log_probs = initial_log_probs
+  alive_seq = initial_ids
+
+  # Finished will keep track of all the sequences that have finished so far
+  # Finished log probs will be negative infinity in the beginning
+  # finished_flags will keep track of booleans
+  finished_seq = initial_ids
+  finished_scores = mtf.constant(mesh, -INF, batch_by_beam)
+
+  # Setting the scores of the initial to negative infinity.
+  finished_flags = mtf.constant(mesh, False, batch_by_beam, tf.bool)
+
+  def grow_finished(finished_seq, finished_scores, finished_flags, curr_seq,
+                    curr_scores, curr_finished):
+    """Given sequences and scores, will gather the top k=beam size sequences.
+
+    Args:
+      finished_seq: Current finished sequences.
+        [batch, beam, length]
+      finished_scores: scores for each of these sequences.
+        [batch, beam]
+      finished_flags: finished bools for each of these sequences.
+        [batch, beam]
+      curr_seq: current topk sequence that has been grown by one position.
+        [batch, beam, length]
+      curr_scores: scores for each of these sequences. [batch, beam]
+      curr_finished: Finished flags for each of these sequences.
+        [batch, beam]
+    Returns:
+      Tuple of
+        (Topk sequences based on scores,
+         log probs of these sequences,
+         Finished flags of these sequences,
+         None (no states))
+    """
+
+    # Set the scores of the unfinished seq in curr_seq to large negative
+    # values
+    curr_scores += (1. - mtf.to_float(curr_finished)) * -INF
+    unused_batch_dim, beam_dim, unused_length_dim = finished_seq.shape.dims
+    # concatenating the sequences and scores along beam axis
+    def _my_concat(a, b):
+      a = mtf.rename_dimension(a, "beam", "triple_beam")
+      b = mtf.rename_dimension(b, "double_beam", "triple_beam")
+      return mtf.concat([a, b], "triple_beam")
+
+    curr_finished_seq = _my_concat(finished_seq, curr_seq)
+    curr_finished_scores = _my_concat(finished_scores, curr_scores)
+    curr_finished_flags = _my_concat(finished_flags, curr_finished)
+    return compute_topk_scores_and_seq(
+        curr_finished_seq, curr_finished_scores, curr_finished_scores,
+        curr_finished_flags, beam_dim, "grow_finished", states=None)
+
+  def grow_alive(curr_seq, curr_scores, curr_log_probs, curr_finished, states):
+    """Given sequences and scores, will gather the top k=beam size sequences.
+
+    Args:
+      curr_seq: current topk sequence that has been grown by one position.
+        [batch, beam, length]
+      curr_scores: scores for each of these sequences. [batch_size, beam_size]
+      curr_log_probs: log probs for each of these sequences.
+        [batch, beam]
+      curr_finished: Finished flags for each of these sequences.
+        [batch, beam]
+      states: list of mtf.Tensor
+    Returns:
+      Tuple of
+        (Topk sequences based on scores,
+         log probs of these sequences,
+         Finished flags of these sequences)
+    """
+    # Set the scores of the finished seq in curr_seq to large negative
+    # values
+    curr_scores += mtf.to_float(curr_finished) * -INF
+    return compute_topk_scores_and_seq(curr_seq, curr_scores, curr_log_probs,
+                                       curr_finished, beam_dim,
+                                       "grow_alive", states)
+
+  def grow_topk(i, alive_seq, alive_log_probs, states=None):
+    r"""Inner beam search loop.
+
+    This function takes the current alive sequences, and grows them to topk
+    sequences where k = 2*beam. We use 2*beam because, we could have beam_size
+    number of sequences that might hit <EOS> and there will be no alive
+    sequences to continue. With 2*beam_size, this will not happen. This relies
+    on the assumption the vocab size is > beam size. If this is true, we'll
+    have at least beam_size non <EOS> extensions if we extract the next top
+    2*beam words.
+    Length penalty is given by = (5+len(decode)/6) ^ -\alpha. Pls refer to
+    https://arxiv.org/abs/1609.08144.
+
+    Args:
+      i: loop index
+      alive_seq: Topk sequences decoded so far [batch, beam, length]
+      alive_log_probs: probabilities of these sequences. [batch, beam]
+      states: optional list of mtf.Tensor
+    Returns:
+      Tuple of
+        (Topk sequences extended by the next word,
+         The log probs of these sequences,
+         The scores with length penalty of these sequences,
+         Flags indicating which of these sequences have finished decoding,
+         list of transformed decoding states)
+    """
+    logits, new_states = logits_fn(i, alive_seq, states)
+    batch_dim, beam_dim, vocab_dim = logits.shape.dims
+
+    # Convert logits to normalized log probs
+    candidate_log_probs = mtf.log_softmax(logits, vocab_dim)
+
+    # Multiply the probabilities by the current probabilities of the beam.
+    # (batch_size, beam_size, vocab_size) + (batch_size, beam_size, 1)
+    log_probs = candidate_log_probs + alive_log_probs
+
+    length_penalty = mtf.pow(((5. + mtf.to_float(i + 1)) / 6.), alpha)
+
+    curr_scores = log_probs / length_penalty
+
+    # scores have shape [batch, beam, vocab]
+    beam_and_vocab_dim = mtf.Dimension(
+        "beam_and_vocab", beam_dim.size * vocab_dim.size)
+    flat_shape = mtf.TensorShape([batch_dim, beam_and_vocab_dim])
+    double_beam = mtf.Dimension("double_beam", beam_dim.size * 2)
+    # Flatten out (beam_size, vocab_size) probs in to a list of possibilities
+    flat_curr_scores = mtf.reshape(curr_scores, flat_shape)
+
+    top_ids, top_scores = mtf.top_k(
+        flat_curr_scores, reduced_dim=beam_and_vocab_dim, new_dim=double_beam)
+
+    # Recovering the log probs because we will need to send them back
+    top_log_probs = top_scores * length_penalty
+
+    # Work out what beam the top probs are in.
+    top_beam_index = top_ids // vocab_dim.size
+    top_ids %= vocab_dim.size  # Unflatten the ids
+
+    def my_gather(tensor):
+      return mtf.gather(
+          tensor, top_beam_index, beam_dim,
+          output_shape=mtf.TensorShape(
+              [double_beam if d == beam_dim else d for d in tensor.shape.dims]))
+
+    # Gather up the most probable 2*beams both for the ids and finished_in_alive
+    # bools
+    top_seq = my_gather(alive_seq)
+
+    if states:
+      states = [my_gather(state) for state in new_states]
+
+    # Append the most probable alive
+    top_seq += top_ids * mtf.one_hot(i, length_dim, dtype=tf.int32)
+    top_finished = mtf.equal(top_ids, eos_id)
+
+    return top_seq, top_log_probs, top_scores, top_finished, states
+
+  def inner_loop(i, alive_seq, alive_log_probs, finished_seq, finished_scores,
+                 finished_flags, *states):
+    """Inner beam search loop.
+
+    There are three groups of tensors, alive, finished, and topk.
+    The alive group contains information about the current alive sequences
+    The topk group contains information about alive + topk current decoded words
+    the finished group contains information about finished sentences, that is,
+    the ones that have decoded to <EOS>. These are what we return.
+    The general beam search algorithm is as follows:
+    While we haven't terminated (pls look at termination condition)
+      1. Grow the current alive to get beam*2 topk sequences
+      2. Among the topk, keep the top beam_size ones that haven't reached EOS
+      into alive
+      3. Among the topk, keep the top beam_size ones have reached EOS into
+      finished
+    Repeat
+    To make things simple with using fixed size tensors, we will end
+    up inserting unfinished sequences into finished in the beginning. To stop
+    that we add -ve INF to the score of the unfinished sequence so that when a
+    true finished sequence does appear, it will have a higher score than all the
+    unfinished ones.
+
+    Args:
+      i: loop index
+      alive_seq: Topk sequences decoded so far [batch_size, beam_size, i+1]
+      alive_log_probs: probabilities of the beams. [batch_size, beam_size]
+      finished_seq: Current finished sequences.
+        [batch_size, beam_size, i+1]
+      finished_scores: scores for each of these sequences.
+        [batch_size, beam_size]
+      finished_flags: finished bools for each of these sequences.
+        [batch_size, beam_size]
+      *states: mtf Tensors
+
+    Returns:
+      Tuple of
+        (Incremented loop index
+         New alive sequences,
+         Log probs of the alive sequences,
+         New finished sequences,
+         Scores of the new finished sequences,
+         Flags indicating which sequence in finished as reached EOS,
+         dict of final decoding states)
+    """
+
+    # Each inner loop, we carry out three steps:
+    # 1. Get the current topk items.
+    # 2. Extract the ones that have finished and haven't finished
+    # 3. Recompute the contents of finished based on scores.
+    (top2k_seq, top2k_log_probs, top2k_scores, top2k_finished,
+     top2k_states) = grow_topk(i, alive_seq, alive_log_probs, states)
+    alive_seq, alive_log_probs, _, states = grow_alive(
+        top2k_seq, top2k_scores, top2k_log_probs, top2k_finished, top2k_states)
+    finished_seq, finished_scores, finished_flags, _ = grow_finished(
+        finished_seq, finished_scores, finished_flags, top2k_seq, top2k_scores,
+        top2k_finished)
+    return (i + 1, alive_seq, alive_log_probs, finished_seq, finished_scores,
+            finished_flags) + tuple(states)
+
+  def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
+                   finished_scores, finished_in_finished, *unused_states):
+    """Checking termination condition.
+
+    We terminate when we decoded up to decode_length or the lowest scoring item
+    in finished has a greater score that the highest prob item in alive divided
+    by the max length penalty
+
+    Args:
+      i: loop index
+      alive_log_probs: probabilities of the beams. [batch_size, beam_size]
+      finished_scores: scores for each of these sequences.
+        [batch_size, beam_size]
+      finished_in_finished: finished bools for each of these sequences.
+        [batch_size, beam_size]
+
+    Returns:
+      Bool.
+    """
+    # TODO(noam): support a different decode length...
+    # decode_length = mtf.constant(mesh, length_dim.size, dtype=tf.int32)
+
+    # del alive_log_probs, finished_scores, finished_in_finished
+    # return mtf.less(i, length_dim.size)
+    if not stop_early:
+      return mtf.less(i, decode_length)
+    max_length_penalty = mtf.pow(
+        ((5. + mtf.to_float(decode_length)) / 6.), alpha)
+    # The best possible score of the most likely alive sequence.
+    lower_bound_alive_scores = mtf.gather(
+        alive_log_probs, mtf.constant(mesh, 0, dtype=tf.int32),
+        beam_dim) / max_length_penalty
+
+    # Now to compute the lowest score of a finished sequence in finished
+    # If the sequence isn't finished, we multiply it's score by 0. since
+    # scores are all -ve, taking the min will give us the score of the lowest
+    # finished item.
+    lowest_score_of_finished_in_finished = mtf.reduce_min(
+        finished_scores * mtf.to_float(finished_in_finished),
+        reduced_dim=beam_dim)
+
+    # If none of the sequences have finished, then the min will be 0 and
+    # we have to replace it by -ve INF if it is. The score of any seq in alive
+    # will be much higher than -ve INF and the termination condition will not
+    # be met.
+    lowest_score_of_finished_in_finished += (
+        (1. - mtf.to_float(mtf.reduce_any(
+            finished_in_finished, reduced_dim=beam_dim))) * -INF)
+
+    bound_is_met = mtf.reduce_all(
+        mtf.greater(lowest_score_of_finished_in_finished,
+                    lower_bound_alive_scores))
+    return mtf.logical_and(
+        mtf.less(i, decode_length), mtf.logical_not(bound_is_met))
+
+  initial_step_num = mtf.constant(mesh, 0, dtype=tf.int32)
+  while_loop_inputs = [
+      initial_step_num, alive_seq, alive_log_probs, finished_seq,
+      finished_scores, finished_flags] + states
+
+  (_, alive_seq, alive_log_probs, finished_seq, finished_scores,
+   finished_flags) = mtf.while_loop(
+       _is_finished, inner_loop, while_loop_inputs,
+       num_loop_vars=None if use_tpu else 6)[:6]
+
+  # Accounting for corner case: It's possible that no sequence in alive for a
+  # particular batch item ever reached EOS. In that case, we should just copy
+  # the contents of alive for that batch item. tf.reduce_any(finished_flags, 1)
+  # if 0, means that no sequence for that batch index had reached EOS. We need
+  # to do the same for the scores as well.
+  finished_seq = mtf.where(
+      mtf.reduce_any(finished_flags, reduced_dim=beam_dim),
+      finished_seq, alive_seq)
+  finished_scores = mtf.where(
+      mtf.reduce_any(finished_flags, reduced_dim=beam_dim),
+      finished_scores, alive_log_probs)
+  return finished_seq, finished_scores
+
+
+def greedy_decode(logits_fn,
+                  initial_ids,
+                  temperature=0.0,
+                  initial_states=None,
+                  eos_id=EOS_ID,
+                  forced_ids=None,
+                  use_tpu=True):
+  """Greedy decoding.
+
+  Args:
+    logits_fn: Interface to the model, to provide logits.
+        Shoud take:
+          step_num - mtf Scalar
+          ids - mtf Tensor with shape [..., length]
+          states - list of mtf.Tensor
+        Should return:
+          logits - [batch, vocab_size]
+          new_states - list of mtf.Tensor
+    initial_ids: mtf.Tensor with shape [..., length], containing zeros.
+    temperature: a float between 0.0 (argmax) and 1.0 (random)
+    initial_states: list of mtf.Tensor
+    eos_id: ID for end of sentence.
+    forced_ids: optional mtf.Tensor with shape [..., length]
+    use_tpu: a boolean
+  Returns:
+    Tensor with shape [..., length]
+  """
+  length_dim = initial_ids.shape.dims[-1]
+  mesh = initial_ids.mesh
+  num_steps = mtf.constant(mesh, length_dim.size, dtype=tf.int32)
+  def cond_fn(step_num, prev_ids, *unused_states):
+    """Should we run another loop iteration."""
+    overflow = mtf.equal(step_num, num_steps)
+    has_eos = mtf.reduce_any(
+        mtf.equal(prev_ids, eos_id), reduced_dim=length_dim)
+    all_has_eos = mtf.reduce_all(has_eos)
+    return mtf.logical_not(mtf.logical_or(overflow, all_has_eos))
+  def body_fn(step_num, ids, *states):
+    """Body function for greedy decoding.
+
+    Args:
+      step_num: a mtf.Tensor
+      ids: a mtf.Tensor
+      *states: additional mtf.Tensors
+    Returns:
+      new_step_num, new_ids, *new_states
+    """
+    logits, new_states = logits_fn(step_num, ids, states)
+    vocab_dim = logits.shape.dims[-1]
+    new_ids = mtf.sample_with_temperature(
+        logits, vocab_dim, temperature)
+    if forced_ids is not None:
+      # force the new ids to equal the partial targets where specified
+      # (positions where partial_targets contain nonzero values)
+      forced = mtf.gather(forced_ids, step_num, length_dim)
+      new_ids = forced + new_ids * mtf.to_int32(mtf.equal(forced, 0))
+    ids += new_ids * mtf.one_hot(step_num, length_dim, dtype=tf.int32)
+    new_step_num = step_num + 1
+    return [new_step_num, ids] + new_states
+  initial_step_num = mtf.constant(mesh, 0, dtype=tf.int32)
+  while_loop_inputs = [initial_step_num, initial_ids] + initial_states
+  final_step_num, mtf_samples = mtf.while_loop(
+      cond_fn, body_fn, while_loop_inputs,
+      num_loop_vars=None if use_tpu else 2)[:2]
+  mtf_samples = mtf.Print(mtf_samples, [final_step_num], "output_length")
+  return mtf_samples
diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
new file mode 100644
index 000000000..736ce646d
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
@@ -0,0 +1,271 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Image Transformer model with model and data parallelism using MTF.
+
+Integration of Mesh tensorflow with Image Transformer to do model parallelism.
+Currently, this supports unconditional image generation. Specify a particular
+architecture layout in the hparams that specifies how different dimensions are
+split or replicated along the mesh dimensions.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import copy
+from tensor2tensor.layers import common_hparams
+from tensor2tensor.layers import common_layers
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+from tensor2tensor.mesh_tensorflow import mtf_layers
+from tensor2tensor.mesh_tensorflow import mtf_model
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+
+@registry.register_model
+class MtfImageTransformer(mtf_model.MtfModel):
+  """Transformer in mesh_tensorflow."""
+
+  def set_activation_type(self):
+    hparams = self._hparams
+    if hparams.activation_dtype == "float32":
+      activation_dtype = tf.float32
+    elif hparams.activation_dtype == "float16":
+      activation_dtype = tf.float16
+    elif hparams.activation_dtype == "bfloat16":
+      activation_dtype = tf.bfloat16
+    else:
+      raise ValueError(
+          "unknown hparams.activation_dtype %s" % hparams.activation_dtype)
+    return activation_dtype
+
+  def mtf_model_fn(self, features, mesh):
+    features = copy.copy(features)
+    tf.logging.info("features = %s" % features)
+    hparams = self._hparams
+    activation_dtype = self.set_activation_type()
+
+    # We assume fixed vocab size for targets
+    targets_vocab_size = self._problem_hparams.target_modality._vocab_size  # pylint: disable=protected-access
+    targets = tf.to_int32(features["targets"])
+
+    # Image preprocessing, reshape into a 1D sequence and shift right.
+    length = hparams.img_len*hparams.img_len*hparams.num_channels
+    targets = tf.reshape(targets, [hparams.batch_size, length])
+    shifted_targets = common_layers.shift_right_2d(targets)
+
+    # Declare all the dimensions
+    model_dim = mtf.Dimension("model", hparams.hidden_size)
+    batch_dim = mtf.Dimension("batch", hparams.batch_size)
+    length_dim = mtf.Dimension("length", length)
+    filter_dim = mtf.Dimension("filter_size", hparams.filter_size)
+    kv_channels = mtf.Dimension("kv_channels", hparams.d_kv)
+    heads = mtf.Dimension("heads", hparams.num_heads)
+
+    def infeed_to_batch_by_length(x, name):
+      return mtf.infeed(
+          mesh, x, mtf.TensorShape([batch_dim, length_dim]), name=name)
+
+    def layer_prepostprocess_dropout(x):
+      return mtf.dropout(
+          x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
+          noise_shape=mtf.TensorShape([batch_dim, model_dim]))
+
+    targets = infeed_to_batch_by_length(targets, "targets")
+    shifted_targets = infeed_to_batch_by_length(
+        shifted_targets, "shifted_targets")
+
+    extra_losses = []
+
+    # TODO(nikip): Verify conditional.
+    if self.has_input and not hparams.unconditional:
+      vocab_size = hparams.num_classes
+      inputs_vocab_dim = mtf.Dimension("vocab", vocab_size)
+      inputs = tf.squeeze(tf.to_int32(features["inputs"]), [2, 3])
+      inputs = infeed_to_batch_by_length(inputs, "inputs")
+
+      # Input embeddings
+      inputs, _ = mtf_layers.embedding(
+          inputs, inputs_vocab_dim, model_dim,
+          activation_dtype=activation_dtype,
+          name="inputs_embedding")
+
+    # Create targets content and position embeddings.
+    targets_position = mtf.range(mesh, length_dim, dtype=tf.int32)
+    targets_vocab_size = 256 * hparams.num_channels
+    targets_vocab_dim = mtf.Dimension("vocab", targets_vocab_size)
+    outputs_vocab_dim = mtf.Dimension("output_vocab", 256)
+
+    # Create embedding var for targets and positions and do a gather.
+    targets_embedding_var = mtf.get_variable(
+        mesh, "targets_embedding",
+        mtf.TensorShape([targets_vocab_dim, model_dim]),
+        initializer=tf.random_normal_initializer(),
+        activation_dtype=activation_dtype)
+
+    positional_embedding_var = mtf.get_variable(
+        mesh, "positional_embedding",
+        mtf.TensorShape([targets_vocab_dim, model_dim]),
+        initializer=tf.random_normal_initializer(),
+        activation_dtype=activation_dtype)
+    x = (mtf.gather(targets_embedding_var, shifted_targets, targets_vocab_dim) +
+         mtf.gather(
+             positional_embedding_var, targets_position, targets_vocab_dim))
+
+    # Image Transformer Decoder
+    # [ self attention - ffn - residual + dropout] x n
+    for layer in range(hparams.num_decoder_layers):
+      layer_name = "decoder_layer_%d" % layer
+      with tf.variable_scope(layer_name):
+        # Self attention layer
+        x += layer_prepostprocess_dropout(
+            mtf_layers.masked_local_attention_1d(
+                mtf_layers.layer_norm(x, model_dim, name="layer_norm_self_att"),
+                None,
+                kv_channels,
+                heads,
+                name="self_att"))
+        # ffn layer
+        x += layer_prepostprocess_dropout(mtf_layers.dense_relu_dense(
+            mtf_layers.layer_norm(x, model_dim, name="layer_norm_ffn"),
+            filter_dim, hparams.dropout, dropout_broadcast_dims=[length_dim]))
+
+    x = mtf_layers.layer_norm(x, model_dim, name="decoder_final_layer_norm")
+
+    # Calculate the logits and loss.
+    logits = mtf_layers.dense(x, outputs_vocab_dim, name="logits")
+    soft_targets = mtf.one_hot(
+        targets, outputs_vocab_dim, dtype=activation_dtype)
+    loss = mtf_layers.softmax_cross_entropy_with_logits(
+        logits, soft_targets, outputs_vocab_dim)
+
+    loss = mtf.reduce_mean(loss)
+    for l in extra_losses:
+      loss += l
+    return logits, loss
+
+
+@registry.register_hparams
+def mtf_image_transformer_base():
+  """Set of hyperparameters."""
+  hparams = common_hparams.basic_params1()
+  hparams.no_data_parallelism = True
+  hparams.use_fixed_batch_size = True
+  hparams.batch_size = 1
+  hparams.max_length = 256
+  hparams.hidden_size = 256
+  hparams.label_smoothing = 0.0
+  # 8-way model-parallelism
+  hparams.add_hparam("mesh_shape", "8")
+  hparams.add_hparam("layout", "vocab:0;filter_size:0;heads:0")
+  hparams.add_hparam("num_heads", 8)
+  hparams.add_hparam("filter_size", 512)
+  hparams.add_hparam("num_encoder_layers", 0)
+  hparams.add_hparam("num_decoder_layers", 6)
+  hparams.add_hparam("attention_key_size", 256)
+  hparams.add_hparam("attention_value_size", 256)
+  # Share weights between input and target embeddings
+  hparams.shared_embedding = True
+
+  # mixture of experts hparams
+  hparams.add_hparam("ffn_layer", "dense_relu_dense")
+  hparams.add_hparam("moe_overhead_train", 1.0)
+  hparams.add_hparam("moe_overhead_eval", 2.0)
+  hparams.moe_num_experts = 16
+  hparams.moe_loss_coef = 1e-3
+
+  hparams.shared_embedding_and_softmax_weights = True
+  hparams.optimizer = "Adafactor"
+  hparams.learning_rate_schedule = "rsqrt_decay"
+  hparams.learning_rate_warmup_steps = 10000
+  hparams.add_hparam("d_kv", 32)
+
+  # Image related hparams
+  hparams.add_hparam("img_len", 32)
+  hparams.add_hparam("num_channels", 3)
+  hparams.add_hparam("unconditional", True)
+  hparams.add_hparam("block_length", 128)
+  return hparams
+
+
+@registry.register_hparams
+def mtf_image_transformer_tiny():
+  """Catch bugs locally..."""
+  hparams = mtf_image_transformer_base()
+  hparams.hidden_size = 128
+  hparams.filter_size = 256
+  hparams.batch_size = 4
+  hparams.num_encoder_layers = 1
+  hparams.num_decoder_layers = 1
+  hparams.num_heads = 4
+  hparams.attention_key_size = 128
+  hparams.attention_value_size = 128
+  # data parallelism and model-parallelism
+  hparams.mesh_shape = "2.2"
+  hparams.layout = "batch:0;filter_size:1"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_image_transformer_single():
+  """Small single parameters."""
+  hparams = mtf_image_transformer_tiny()
+  hparams.mesh_shape = ""
+  hparams.layout = ""
+  hparams.hidden_size = 32
+  hparams.filter_size = 32
+  hparams.batch_size = 1
+  hparams.num_encoder_layers = 1
+  hparams.num_decoder_layers = 1
+  hparams.num_heads = 2
+  hparams.attention_key_size = 32
+  hparams.attention_value_size = 32
+  return hparams
+
+
+@registry.register_hparams
+def mtf_image_transformer_base_single():
+  """Small single parameters."""
+  hparams = mtf_image_transformer_base()
+  hparams.num_decoder_layers = 6
+  hparams.filter_size = 256
+  hparams.block_length = 128
+  hparams.mesh_shape = ""
+  hparams.layout = ""
+  return hparams
+
+
+@registry.register_hparams
+def mtf_image_transformer_tiny_moe():
+  hparams = mtf_image_transformer_tiny()
+  hparams.mesh_shape = "4"
+  hparams.layout = "batch:0,experts:0"
+  hparams.ffn_layer = "moe"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_image_transformer_tiny_8gpu():
+  hparams = mtf_image_transformer_tiny()
+  hparams.mesh_shape = "8"
+  hparams.layout = "vocab:0;filter_size:0;heads:0"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_image_transformer_length_sharded():
+  hparams = mtf_image_transformer_tiny()
+  hparams.mesh_shape = "2"
+  hparams.layout = "length:0"
+  return hparams
diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py b/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
new file mode 100644
index 000000000..655e50194
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
@@ -0,0 +1,95 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 Image Transformer on Mesh TensorFlow."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+from tensor2tensor.mesh_tensorflow import mtf_image_transformer
+from tensor2tensor.mesh_tensorflow import placement_mesh_impl
+
+import tensorflow as tf
+
+# Constants shared between all functions.
+BATCH_SIZE = 8
+INPUT_LENGTH = 8
+IMG_LENGTH = 8
+VOCAB_SIZE = 256
+
+
+def get_model(hparams=None,
+              mode=tf.estimator.ModeKeys.TRAIN,
+              model_cls=mtf_image_transformer.MtfImageTransformer):
+  if hparams is None:
+    hparams = mtf_image_transformer.mtf_image_transformer_single()
+  hparams.max_length = INPUT_LENGTH
+  hparams.batch_size = BATCH_SIZE
+  hparams.img_len = IMG_LENGTH
+  hparams.num_channels = 1
+
+  p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE, VOCAB_SIZE)
+  p_hparams.input_modality = {}
+  hparams.problem_hparams = p_hparams
+
+  targets = -1 + np.random.random_integers(
+      VOCAB_SIZE, size=(BATCH_SIZE, IMG_LENGTH*IMG_LENGTH, 1, 1))
+  features = {
+      "targets": tf.constant(targets, dtype=tf.int32, name="targets"),
+  }
+
+  return model_cls(hparams, mode, p_hparams), features, hparams
+
+
+def get_placement_mesh(hparams):
+  graph = mtf.Graph()
+  mesh = mtf.Mesh(graph, "my_mesh")
+  mesh_shape = mtf.parse_mesh_shape(hparams.mesh_shape)
+  mesh_size = mtf.list_product(mesh_shape)
+
+  mesh_devices = [""] * mesh_size
+  mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+      mesh_shape, mtf.parse_layout(hparams.layout), mesh_devices)
+  return mesh, mesh_impl
+
+
+class MtfImageTransformerTest(tf.test.TestCase):
+
+  def testMtfImageTransformer(self):
+    hparams = mtf_image_transformer.mtf_image_transformer_single()
+
+    model, features, hparams = get_model(hparams)
+    hparams.mesh_shape = ""
+    hparams.layout = ""
+    mesh, mesh_impl = get_placement_mesh(hparams)
+
+    logits, _ = model.mtf_model_fn(features, mesh)
+    lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
+    tf_group = lowering.copy_masters_to_slices()
+    tf_logits = lowering.outfeed(logits)
+
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      session.run(tf_group)
+      res = session.run(tf_logits)
+    self.assertEqual(res.shape, (BATCH_SIZE, IMG_LENGTH*IMG_LENGTH, VOCAB_SIZE))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers.py b/tensor2tensor/mesh_tensorflow/mtf_layers.py
new file mode 100644
index 000000000..09a2dea93
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers.py
@@ -0,0 +1,851 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Layers for mesh tensorflow."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import functools
+from tensor2tensor.layers import common_layers
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+import tensorflow as tf
+
+
+def dense(x, output_dim, reduced_dims=None, expert_dims=None,
+          use_bias=True, activation=None, name=None):
+  """Dense layer doing (kernel*x + bias) computation.
+
+  Args:
+    x: a mtf.Tensor of shape [..., reduced_dims].
+    output_dim: a mtf.Dimension
+    reduced_dims: an optional list of mtf.Dimensions of x to be reduced. If
+      omitted, we reduce the last dimension.
+    expert_dims: an optional list of mtf.Dimension which represent different
+      experts. Different experts get different weights.
+    use_bias: a boolean, whether to add bias.
+    activation: an optional function from mtf.Tensor to mtf.Tensor
+    name: a string. variable scope.
+
+  Returns:
+    a mtf.Tensor of shape [..., output_dim].
+  """
+  if expert_dims is None:
+    expert_dims = []
+  if reduced_dims is None:
+    reduced_dims = x.shape.dims[-1:]
+  w_shape = mtf.TensorShape(expert_dims + reduced_dims + [output_dim])
+  output_shape = mtf.TensorShape(
+      [d for d in x.shape.dims if d not in reduced_dims] + [output_dim])
+  with tf.variable_scope(name, default_name="dense"):
+    stddev = mtf.list_product(d.size for d in reduced_dims) ** -0.5
+    w = mtf.get_variable(
+        x.mesh,
+        "kernel",
+        w_shape,
+        initializer=tf.random_normal_initializer(stddev=stddev),
+        activation_dtype=x.dtype)
+    y = mtf.matmul(x, w, output_shape=output_shape)
+    if use_bias:
+      b = mtf.get_variable(
+          x.mesh,
+          "bias",
+          mtf.TensorShape(expert_dims + [output_dim]),
+          initializer=tf.zeros_initializer(),
+          activation_dtype=x.dtype)
+      y += b
+    if activation is not None:
+      y = activation(y)
+    return y
+
+
+def layer_norm(x, dim, epsilon=1e-6, name="layer_prepostprocess"):
+  """Layer normalization over dimension dim.
+
+  Args:
+    x: a mtf.Tensor whose shape contains dim.
+    dim: a mtf.Dimension
+    epsilon: a floating point number
+    name: a string. variable scope.
+
+  Returns:
+    a mtf.Tensor with same shape as x.
+  """
+  with tf.variable_scope(name + "/layer_norm"):
+    scale = mtf.get_variable(
+        x.mesh,
+        "layer_norm_scale",
+        mtf.TensorShape([dim]),
+        initializer=tf.ones_initializer(),
+        activation_dtype=x.dtype)
+    bias = mtf.get_variable(
+        x.mesh,
+        "layer_norm_bias",
+        mtf.TensorShape([dim]),
+        initializer=tf.zeros_initializer(),
+        activation_dtype=x.dtype)
+    reduced_shape = x.shape - dim
+    mean = mtf.reduce_mean(x, output_shape=reduced_shape)
+    variance = mtf.reduce_mean(mtf.square(x - mean), output_shape=reduced_shape)
+    norm_x = (x - mean) * mtf.rsqrt(variance + epsilon)
+    return norm_x * scale + bias
+
+
+def softmax_cross_entropy_with_logits(logits, targets, vocab_dim):
+  """Per-example softmax loss.
+
+  Args:
+    logits: a mtf.Tensor whose shape contains vocab_dim
+    targets: a mtf.Tensor with the same shape as logits
+    vocab_dim: a mtf.Dimension
+
+  Returns:
+    a mtf.Tensor whose shape is equal to logits.shape - vocab_dim
+
+  Raises:
+    ValueError: if the shapes do not match.
+  """
+  if logits.shape != targets.shape:
+    raise ValueError(
+        "logits shape must equal targets shape"
+        "logits=%s targets=%s" % (logits.to_string, targets.to_string))
+  if vocab_dim not in logits.shape.dims:
+    raise ValueError("vocab_dim must be in logits.shape.dims")
+  log_softmax = mtf.log_softmax(logits, vocab_dim)
+  return mtf.negative(
+      mtf.reduce_sum(log_softmax * targets, reduced_dim=vocab_dim))
+
+
+def weights_nonzero(targets, dtype=tf.float32):
+  def my_fn(x):
+    return tf.cast(tf.not_equal(x, 0), dtype)
+  return mtf.cwise(my_fn, [targets], output_dtype=dtype, name="weights_nonzero")
+
+
+def dense_relu_dense(x,
+                     hidden_channels,
+                     dropout=0.0,
+                     dropout_broadcast_dims=None,
+                     name=None):
+  """Hidden layer with ReLU activation followed by linear projection.
+
+  The output has the same number of channels as the input.
+
+  Args:
+    x: a mtf.Tensor
+    hidden_channels: a mtf.Dimension - channels in the hidden layer
+    dropout: an optional float
+    dropout_broadcast_dims: an optional list of mtf.Dimension
+    name: an optional string
+
+  Returns:
+    a mtf.Tensor with the same shape as x.
+  """
+  with tf.variable_scope(name, default_name="dense_relu_dense"):
+    io_channels = x.shape.dims[-1]
+    stddev = (hidden_channels.size * io_channels.size) ** -0.25
+    io = mtf.Dimension("io", 2)
+    w = mtf.get_variable(
+        x.mesh,
+        "kernel",
+        mtf.TensorShape([io, io_channels, hidden_channels]),
+        initializer=tf.random_normal_initializer(stddev=stddev),
+        activation_dtype=x.dtype)
+    wi, wo = mtf.unstack(w, io)
+    h = mtf.relu(mtf.einsum([x, wi]))
+    if dropout != 0.0:
+      h = mtf.dropout(h, 1.0 - dropout,
+                      noise_shape=h.shape - dropout_broadcast_dims)
+    return mtf.einsum([h, wo])
+
+
+def masked_local_attention_1d(query_antecedent,
+                              memory_antecedent,
+                              kv_channels,
+                              heads,
+                              block_length=128,
+                              name=None):
+  """Attention to the source position and a neighborhood to the left of it.
+
+  The sequence is divided into blocks of length block_size.
+  Attention for a given query position can only see memory positions
+  less than or equal to the query position, in the corresponding block
+  and the previous block.
+
+  Args:
+    query_antecedent: a mtf.Tensor with shape [batch, query_length, io_channels]
+    memory_antecedent: a mtf.Tensor with shape
+      [batch, memory_length, io_channels] (optional). Currently, memory_length
+      must have the same size as query_length, but a different name.
+    kv_channels: a mtf.Dimension (the size of the key and value vectors)
+    heads: a mtf.Dimension (the number of heads)
+    block_length: an integer, representing receptive fields for attention.
+    name: an optional string.
+
+  Returns:
+    a Tensor of shape [batch, query_length, io_channels]
+
+  Raises:
+    ValueError: if channels or depth don't match.
+  """
+  with tf.variable_scope(
+      name, default_name="multihead_attention",
+      values=[query_antecedent, memory_antecedent]):
+
+    batch, query_length, io_channels = query_antecedent.shape.dims
+    q_var, k_var, v_var, o_var = multihead_attention_vars(
+        query_antecedent.mesh, heads, io_channels, kv_channels,
+        query_antecedent.dtype)
+
+    if memory_antecedent is None:
+      memory_antecedent = rename_length_to_memory_length(
+          query_antecedent, query_length.name)
+    memory_batch, memory_length, memory_channels = memory_antecedent.shape.dims
+    if memory_batch != batch:
+      raise ValueError("memory batch must equal query batch")
+    if memory_channels != io_channels:
+      raise ValueError("memory channels must equal query channels")
+
+    # Get query q, keys k and values v.
+    q = mtf.einsum(
+        [query_antecedent, q_var],
+        mtf.TensorShape([batch, heads, query_length, kv_channels]))
+    k = mtf.einsum(
+        [memory_antecedent, k_var],
+        mtf.TensorShape([batch, heads, memory_length, kv_channels]))
+    v = mtf.einsum(
+        [memory_antecedent, v_var],
+        mtf.TensorShape([batch, heads, memory_length, kv_channels]))
+
+    # Let's assume for now we don't have padding and the block length equally
+    # divides the memory length.
+    block_length = (query_length.size
+                    if query_length.size < block_length * 2 else block_length)
+    blength = mtf.Dimension("block_length", block_length)
+    mlength = mtf.Dimension("mem_block_length", block_length)
+    num_blocks = mtf.Dimension("num_blocks", query_length.size // block_length)
+
+    q = mtf.reshape(
+        q, mtf.TensorShape([batch, heads, num_blocks, blength, kv_channels]))
+    k = mtf.reshape(
+        k, mtf.TensorShape([batch, heads, num_blocks, mlength, kv_channels]))
+    v = mtf.reshape(
+        v, mtf.TensorShape([batch, heads, num_blocks, mlength, kv_channels]))
+
+    # compute attention for the first query block.
+    def first_block_attention():
+      """Compute attention for the first block."""
+      first_q = mtf.slice(q, 0, 1, num_blocks.name)
+      first_k = mtf.slice(k, 0, 1, num_blocks.name)
+      first_v = mtf.slice(v, 0, 1, num_blocks.name)
+      block = first_q.shape.dims[2]
+
+      first_logits = mtf.einsum(
+          [first_q, first_k],
+          mtf.TensorShape([batch, heads, block, blength, mlength]))
+      weights = mtf.softmax(first_logits, mlength)
+      first_output = mtf.einsum(
+          [weights, first_v],
+          mtf.TensorShape([batch, heads, block, blength, kv_channels]))
+      return first_output
+
+    # Attention for first block, since query_length = key_length.
+    first_output = first_block_attention()
+
+    # Concatenate two adjacent blocks to compute the overlapping memory block.
+    def local(x):
+      """Helper function to get memory blocks."""
+      prev_block = mtf.slice(x, 0, num_blocks.size-1, num_blocks.name)
+      cur_block = mtf.slice(x, 1, num_blocks.size-1, num_blocks.name)
+      local_block = mtf.concat([prev_block, cur_block], mlength.name)
+      return local_block
+
+    local_k = local(k)
+    local_v = local(v)
+    mblocks = local_k.shape.dims[2]
+    mlength = local_k.shape.dims[3]
+    # Calculate the causal mask to avoid peeking into the future. We compute
+    # this once and reuse it for all blocks since the block_size is known.
+    mask = attention_bias_local_block(query_antecedent.mesh,
+                                      blength, mlength)
+
+    # Remove the first block from q since we already computed that.
+    tail_q = mtf.slice(q, 1, num_blocks.size-1, num_blocks.name)
+
+    # Compatibility between q and k for rest of the blocks.
+    # Shape [batch, heads, num_blocks - 1, block_length, local_length]
+    attention = mtf.einsum(
+        [tail_q, local_k],
+        mtf.TensorShape([batch, heads, mblocks, blength, mlength]))
+    attention += mask
+    attention = mtf.softmax(attention, mlength)
+
+    # Run attention for rest of the blocks.
+    # Shape [batch, heads, num_blocks-1, block_length, kv_channels]
+    output = mtf.einsum(
+        [attention, local_v],
+        mtf.TensorShape([batch, heads, mblocks, blength, kv_channels]))
+    # Now concatenate the first and rest of the blocks.
+    final_output = mtf.concat([first_output, output], num_blocks.name)
+    final_output = mtf.reshape(final_output, mtf.TensorShape(
+        [batch, heads, query_length, kv_channels]))
+    return mtf.einsum([final_output, o_var],
+                      mtf.TensorShape([batch, query_length, io_channels]))
+
+
+def rename_length_to_memory_length(
+    x, length_name="length", memory_length_name="memory_length"):
+  return mtf.rename_dimension(x, length_name, memory_length_name)
+
+
+def multihead_attention_vars(
+    mesh, heads, io_channels, kv_channels, activation_dtype):
+  """Create Parameters for Multihead Attention.
+
+  Args:
+    mesh: a Mesh
+    heads: a Dimension
+    io_channels: a Dimension
+    kv_channels: a Dimension
+    activation_dtype: a tf.dtype
+
+  Returns:
+    q_var: a Tensor with shape [heads, io_channels, kv_channels]
+    k_var: a Tensor with shape [heads, io_channels, kv_channels]
+    v_var: a Tensor with shape [heads, io_channels, kv_channels]
+    o_var: a Tensor with shape [heads, io_channels, kv_channels]
+  """
+  qkvo = mtf.Dimension("qkvo", 4)
+  qk_stddev = (io_channels.size ** -0.5) * (kv_channels.size ** -0.25)
+  v_stddev = io_channels.size ** -0.5
+  o_stddev = (io_channels.size * heads.size) ** -0.5
+  def qkvo_initializer(shape,
+                       dtype=None,
+                       partition_info=None,
+                       verify_shape=None):
+    del partition_info, verify_shape
+    return tf.random_normal(shape, dtype=dtype) * tf.reshape(
+        [qk_stddev, qk_stddev, v_stddev, o_stddev], [4, 1, 1, 1])
+  var = mtf.get_variable(
+      mesh, "qkvo", mtf.TensorShape([qkvo, heads, io_channels, kv_channels]),
+      initializer=qkvo_initializer, activation_dtype=activation_dtype)
+  q_var, k_var, v_var, o_var = mtf.unstack(var, qkvo)
+  return q_var, k_var, v_var, o_var
+
+
+def dot_product_attention(q,
+                          k,
+                          v,
+                          mask,
+                          dropout=0.0,
+                          dropout_broadcast_dims=None):
+  """Dot-product attention.
+
+  Args:
+    q: Tensor with shape [...., length_q, depth_k]. Typically leading dimensions
+      are [batch, heads].
+    k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
+      match with q.
+    v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
+      match with q.
+    mask: mask Tensor (see attention_mask())
+    dropout: a float.
+    dropout_broadcast_dims: an optional list of mtf.Dimension
+
+  Returns:
+    Tensor with shape [..., length_q, depth_v].
+  """
+  length_kv = k.shape.dims[-2]
+  logits_shape = mtf.TensorShape(q.shape.dims[:-1] + [length_kv])
+  logits = mtf.einsum([q, k], logits_shape)
+  if mask is not None:
+    logits += mask
+  weights = mtf.softmax(logits, length_kv)
+  if dropout != 0.0:
+    weights = mtf.dropout(
+        weights, 1.0 - dropout,
+        noise_shape=weights.shape - dropout_broadcast_dims)
+  depth_v = v.shape.dims[-1]
+  outputs_shape = mtf.TensorShape(q.shape.dims[:-1] + [depth_v])
+  outputs = mtf.einsum([weights, v], outputs_shape)
+  return outputs
+
+
+def multihead_attention(query_antecedent,
+                        memory_antecedent,
+                        mask,
+                        kv_channels,
+                        heads,
+                        dropout=0.0,
+                        dropout_broadcast_dims=None,
+                        name="multihead_attention"):
+  """Multihead scaled-dot-product attention with input/output transformations.
+
+  In order to use only one variable containing the four weight matrices
+  packed together, we insist that the query and memory antecedents have the
+  same dimensionality (io_channels) and that the keys and values have the
+  same dimensionality (kv_channels).
+
+  Args:
+    query_antecedent: a mtf.Tensor with shape [batch, query_length, io_channels]
+    memory_antecedent: a mtf.Tensor with shape
+      [batch, memory_length, io_channels] (optional)
+    mask: mask Tensor (see attention_mask())
+    kv_channels: a mtf.Dimension (the size of the key and value vectors)
+    heads: a mtf.Dimension (the number of heads)
+    dropout: a floating point value
+    dropout_broadcast_dims: an optional list of mtf.Dimension
+    name: an optional string.
+
+  Returns:
+    A mtf.Tensor with shape [batch, query_length, io_channels]
+
+  Raises:
+    ValueError: if the dimensions do not match.
+  """
+  batch, query_length, io_channels = query_antecedent.shape.dims
+  with tf.variable_scope(name,
+                         default_name="multihead_attention",
+                         values=[query_antecedent, memory_antecedent]):
+    q_var, k_var, v_var, o_var = multihead_attention_vars(
+        query_antecedent.mesh, heads, io_channels, kv_channels,
+        query_antecedent.dtype)
+    if memory_antecedent is None:
+      memory_antecedent = rename_length_to_memory_length(
+          query_antecedent, query_length.name)
+    memory_batch, memory_length, memory_channels = memory_antecedent.shape.dims
+    if memory_batch != batch:
+      raise ValueError("memory batch must equal query batch")
+    if memory_channels != io_channels:
+      raise ValueError("memory channels must equal query channels")
+    q = mtf.einsum(
+        [query_antecedent, q_var],
+        mtf.TensorShape([batch, heads, query_length, kv_channels]))
+    k = mtf.einsum(
+        [memory_antecedent, k_var],
+        mtf.TensorShape([batch, heads, memory_length, kv_channels]))
+    v = mtf.einsum(
+        [memory_antecedent, v_var],
+        mtf.TensorShape([batch, heads, memory_length, kv_channels]))
+    o = dot_product_attention(
+        q, k, v, mask, dropout, dropout_broadcast_dims)
+    return mtf.einsum(
+        [o, o_var], mtf.TensorShape([batch, query_length, io_channels]))
+
+
+def multihead_self_attention_incremental(query_antecedent,
+                                         prev_k,
+                                         prev_v,
+                                         step_num,
+                                         name="multihead_attention"):
+  """Incremental self-attention (one decode step).
+
+  In order to use only one variable containing the four weight matrices
+  packed together, we insist that the query and memory antecedents have the
+  same dimensionality (io_channels) and that the keys and values have the
+  same dimensionality (kv_channels).
+
+  Args:
+    query_antecedent: a mtf.Tensor with shape [batch..., io_channels]
+    prev_k: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels]
+    prev_v: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels]
+    step_num: mtf Scalar with dtype tf.int32
+    name: an optional string.
+
+  Returns:
+    y: A mtf.Tensor with shape [batch..., io_channels]
+    new_k: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels]
+    new_v: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels]
+
+  Raises:
+    ValueError: if the dimensions do not match.
+  """
+  batch_dims = query_antecedent.shape.dims[:-1]
+  io_channels = query_antecedent.shape.dims[-1]
+  heads, memory_length, kv_channels = prev_k.shape.dims[-3:]
+  with tf.variable_scope(name, default_name="multihead_attention"):
+    q_var, k_var, v_var, o_var = multihead_attention_vars(
+        query_antecedent.mesh, heads, io_channels, kv_channels,
+        query_antecedent.dtype)
+    memory_antecedent = query_antecedent
+    q = mtf.einsum(
+        [query_antecedent, q_var],
+        mtf.TensorShape(batch_dims + [heads, kv_channels]))
+    k = mtf.einsum(
+        [memory_antecedent, k_var],
+        mtf.TensorShape(batch_dims + [heads, kv_channels]))
+    v = mtf.einsum(
+        [memory_antecedent, v_var],
+        mtf.TensorShape(batch_dims + [heads, kv_channels]))
+    k = prev_k + mtf.multiply(
+        k, mtf.one_hot(step_num, memory_length), output_shape=prev_k.shape)
+    v = prev_v + mtf.multiply(
+        v, mtf.one_hot(step_num, memory_length), output_shape=prev_v.shape)
+
+    mask = mtf.to_float(mtf.greater(mtf.range(
+        query_antecedent.mesh, memory_length, dtype=tf.int32), step_num)
+                       ) * -1e9
+    o = dot_product_attention(q, k, v, mask)
+    y = mtf.einsum([o, o_var], query_antecedent.shape)
+    return y, k, v
+
+
+def multihead_encdec_attention_incremental(query_antecedent,
+                                           q_var, o_var, k, v,
+                                           mask,
+                                           name="multihead_attention"):
+  """Incremental attention over encoder (one decode step).
+
+  In order to use only one variable containing the four weight matrices
+  packed together, we insist that the query and memory antecedents have the
+  same dimensionality (io_channels) and that the keys and values have the
+  same dimensionality (kv_channels).
+
+  memory_dims is a subset of query_dims
+
+  Args:
+    query_antecedent: a mtf.Tensor with shape query_dims + [io_channels]
+    q_var: a mtf.Tensor with shape [heads, io_channels, kv_channels]
+    o_var: a mtf.Tensor with shape [heads, io_channels, kv_channels]
+    k: memory_dims + [heads, memory_length, kv_channels]
+    v: memory_dims + [heads, memory_length, kv_channels]
+    mask: mask Tensor (see attention_mask())
+    name: an optional string.
+
+  Returns:
+    A mtf.Tensor with shape [batch, qlen, io_channels]
+  """
+  heads, _, kv_channels = k.shape.dims[-3:]
+  query_dims = query_antecedent.shape.dims[:-1]
+  with tf.variable_scope(name, default_name="multihead_attention"):
+    q = mtf.einsum(
+        [query_antecedent, q_var],
+        mtf.TensorShape(query_dims + [heads, kv_channels]))
+    o = dot_product_attention(q, k, v, mask)
+    return mtf.einsum([o, o_var], query_antecedent.shape)
+
+
+def attention_mask_ignore_padding(inputs, dtype=tf.float32):
+  """Bias for encoder-decoder attention.
+
+  Args:
+    inputs: a mtf.Tensor with shape [..., length_dim]
+    dtype: a tf.dtype
+
+  Returns:
+    a mtf.Tensor with shape [..., memory_length_dim]
+  """
+  inputs = rename_length_to_memory_length(inputs)
+  return mtf.cast(mtf.equal(inputs, 0), dtype) * -1e9
+
+
+def attention_mask_autoregressive(query_pos, dtype=tf.float32):
+  """Bias for self-attention where attention to the right is disallowed.
+
+  Args:
+    query_pos: a mtf.Tensor with shape [..., length_dim]
+    dtype: a tf.dtype
+
+  Returns:
+    a mtf.Tensor with shape [..., length_dim, memory_length_dim]
+  """
+  memory_pos = rename_length_to_memory_length(query_pos)
+  return mtf.cast(mtf.less(query_pos, memory_pos), dtype) * -1e9
+
+
+def attention_mask_same_segment(
+    query_segment, memory_segment=None, dtype=tf.float32):
+  """Bias for attention where attention between segments is disallowed.
+
+  Args:
+    query_segment: a mtf.Tensor with shape [..., length_dim]
+    memory_segment: a mtf.Tensor with shape [..., memory_length_dim]
+    dtype: a tf.dtype
+
+  Returns:
+    a mtf.Tensor with shape [..., length_dim, memory_length_dim]
+  """
+  memory_segment = rename_length_to_memory_length(
+      memory_segment or query_segment)
+  return mtf.cast(mtf.not_equal(query_segment, memory_segment), dtype) * -1e9
+
+
+def attention_bias_local_block(mesh, block_length, memory_length,
+                               dtype=tf.int32):
+  """Bias for attention for local blocks where attention to right is disallowed.
+
+  Args:
+    mesh: a MeshTensorflow object
+    block_length: a mtf.Dimension
+    memory_length: a mtf.Dimension
+    dtype: a tf.dtype
+
+  Returns:
+    a mtf.Tensor with shape [rows, cols]
+  """
+  mask = mtf.cast(mtf.less(mtf.range(mesh, block_length, dtype=dtype),
+                           mtf.range(mesh, memory_length, dtype=dtype)),
+                  dtype=dtype)
+  mask = mtf.cast(mask, dtype=tf.float32)  * -1e9
+  return mask
+
+
+def moe_v0(inputs,
+           hidden_dim,
+           output_dim,
+           experts_dim,
+           loss_coef=1e-3,
+           overhead=1.0):
+  """Local mixture of experts that works well on TPU.
+
+  See https://arxiv.org/abs/1701.06538
+
+  There are num_experts expert networks, each containing a relu-activated
+  hidden layer of size hidden_size, followed by an output projection.
+
+  The number of parameters is thus:
+    num_experts * (input_size * hidden_size + hidden_size * output_size)
+
+  The input is 3d: [batch, length, depth], consisting of the representations
+  of all positions in a batch of sequences.
+
+  Each position of each sequence is sent to 0-2 experts.  The expert
+  choices and the combination weights are determined by a learned gating
+  function.
+
+  This function returns a small auxiliary loss that should be added to the
+  training loss of the model.  This loss helps to balance expert usage.
+  Without the loss, it is very likely that a few experts will be trained and
+  the rest will starve.
+
+  Several hacks are necessary to get around current TPU limitations:
+
+  - To ensure static shapes, we enforce (by truncation/padding)
+    that each sequence send the same number of elements to each expert.
+
+    It would make more sense to enforce this equality over the entire batch,
+    as opposed to on individual sequences.  This would allow more freedom
+    for individual sequences to be unbalanced.  Unfortunately, that would
+    slow down our hacked-up gather-by-matmul implementation.
+
+    TODO(noam): There is no real reason for a single sequence to be the unit
+      of equal allocation.  Reshaping the inputs would allow us to pick a
+      different unit of equal allocation.
+
+  TODO(noam): Factor this code better.  We want to be able to substitute
+  different code for the experts themselves.  We also want to integrate this
+  gating/dispatching logic into multi-device mixtures-of-experts.
+
+  Args:
+    inputs: a mtf.Tensor with shape [batch_dim, length_dim, input_dim]
+    hidden_dim: a mtf.Dimension
+    output_dim: a mtf.Dimension
+    experts_dim: a mtf.Dimension
+    loss_coef: a float scalar
+    overhead: multiplicative factor of how much spare capacity to assign
+
+  Returns:
+    outputs: a Tensor with shape [batch_dim, length_dim, output_dim]
+    loss: a mtf scalar
+  """
+  batch_dim, length_dim, input_dim = inputs.shape.dims
+
+  # Each sequence sends expert_capacity positions to each expert.
+  expert_capacity = min(
+      length_dim.size,
+      int((length_dim.size * 2 * overhead) / experts_dim.size))
+  expert_capacity_dim = mtf.Dimension("expert_capacity", expert_capacity)
+
+  experts_dim_unsplit = mtf.Dimension("expert_unsplit", experts_dim.size)
+  batch_dim_unsplit = mtf.Dimension("batch_unsplit", batch_dim.size)
+
+  # This is the learned gating function.
+  # shape = [batch_dim, length_dim, experts_dim_unsplit]
+  gates = mtf.softmax(dense(inputs, experts_dim_unsplit), experts_dim_unsplit)
+
+  assignment_shape = mtf.TensorShape(
+      [batch_dim, length_dim, experts_dim_unsplit, expert_capacity_dim])
+
+  backward_assignment = mtf.slicewise(
+      functools.partial(
+          _truncated_top_2_gating, expert_capacity=expert_capacity),
+      [gates],
+      output_shape=assignment_shape,
+      splittable_dims=[batch_dim],
+      name="backward_assignment")
+
+  forward_assignment = mtf.cast(
+      mtf.cast(backward_assignment, tf.bool), inputs.dtype)
+
+  # put num_experts dimension first to make split easier in alltoall
+  expert_inputs = mtf.einsum([inputs, forward_assignment], mtf.TensorShape(
+      [experts_dim_unsplit, batch_dim, expert_capacity_dim, input_dim]))
+
+  expert_inputs = mtf.reshape(expert_inputs, mtf.TensorShape(
+      [experts_dim, batch_dim_unsplit, expert_capacity_dim, input_dim]))
+
+  # Now feed the expert inputs through the experts.
+  h = dense(expert_inputs, hidden_dim, expert_dims=[experts_dim],
+            activation=mtf.relu, name="x0")
+  expert_output = dense(h, output_dim, expert_dims=[experts_dim], name="x1")
+
+  expert_output = mtf.reshape(expert_output, mtf.TensorShape(
+      [experts_dim_unsplit, batch_dim, expert_capacity_dim, input_dim]))
+
+  output = mtf.einsum([expert_output, backward_assignment], mtf.TensorShape(
+      [batch_dim, length_dim, output_dim]))
+
+  importance = mtf.reduce_sum(backward_assignment, output_shape=mtf.TensorShape(
+      [batch_dim, experts_dim_unsplit]))
+
+  loss = cv_squared(importance) * loss_coef
+  return output, loss
+
+
+def cv_squared(x):
+  """The squared coefficient of variation of a sample.
+
+  Useful as a loss to encourage a positive distribution to be more uniform.
+  Epsilons added for numerical stability.
+  Returns 0 for an empty Tensor.
+
+  Args:
+    x: a mtf.Tensor
+
+  Returns:
+    a mtf Scalar
+  """
+  epsilon = 1e-10
+  mean = mtf.reduce_mean(x)
+  variance = mtf.reduce_mean(mtf.square(x - mean))
+  return variance / (mtf.square(mean) + epsilon)
+
+
+def _truncated_top_2_gating(gates, expert_capacity, show_summaries=False):
+  """Compute gating for mixture-of-experts in TensorFlow.
+
+  gates is usually the output of a softmax function.
+  The return value is a dense representation of the mapping between
+  the input positions in the positions in the batches sent to the experts.
+
+  TODO(noam): this function contains code factored out of
+  expert_utils.local_moe_tpu.  Move this function to that file and
+  call it from both places.
+
+  Args:
+    gates: a Tensor with shape [batch, length, num_experts]
+    expert_capacity: an integer
+    show_summaries: a boolean
+
+  Returns:
+    a Tensor with shape [batch, length, num_experts, expert_capacity]
+  """
+  def _to_float(x):
+    return tf.cast(x, gates.dtype)
+  batch = tf.shape(gates)[0]
+  length = tf.shape(gates)[1]
+  num_experts = tf.shape(gates)[2]
+  expert_capacity_f = _to_float(expert_capacity)
+  # Find the top expert for each position.
+  gate_1, index_1 = common_layers.top_1_tpu(gates)
+  # [batch, length, num_experts]
+  mask_1 = tf.one_hot(index_1, num_experts, dtype=gates.dtype)
+  # [batch, length, num_experts]
+  # This is the position within the expert's mini-batch for this sequence
+  position_in_expert_1 = common_layers.cumsum(
+      mask_1, axis=1, exclusive=True) * mask_1
+  # Remove the elements that don't fit.
+  mask_1 *= _to_float(tf.less(position_in_expert_1, expert_capacity_f))
+  # [batch, 1, num_experts]
+  # How many examples in this sequence go to this expert
+  mask_1_count = tf.reduce_sum(mask_1, axis=1, keepdims=True)
+  # [batch, length] - mostly ones, but zeros where something didn't fit
+  mask_1_flat = tf.reduce_sum(mask_1, axis=2)
+  position_in_expert_1 = tf.reduce_sum(position_in_expert_1, axis=2)
+  # Weight assigned to first expert.
+  gate_1 *= mask_1_flat
+
+  # Pick a second-place expert for each position.
+  # We first mask out the experts that we expect to be over-capacity
+  space_remaining = expert_capacity_f - mask_1_count
+  use_rate = (mask_1_count + 1.0) / _to_float(length)
+  # At what point in the sequence do we expect the expert to be full.
+  expected_exhaustion_pos = space_remaining / use_rate
+  # A Tensor with shape [batch, length, num_experts] representing a boolean
+  #   - whether we expect that the expert will already be full.
+  expected_exhausted = _to_float(tf.greater(
+      tf.reshape(_to_float(tf.range(length)), [1, length, 1]),
+      expected_exhaustion_pos))
+  masked_gates = gates - mask_1 - expected_exhausted
+  # This section is similar to the section above.
+  gate_2, index_2 = common_layers.top_1_tpu(masked_gates)
+  # [batch, length, num_experts]
+  mask_2 = tf.one_hot(index_2, num_experts, dtype=gates.dtype)
+  position_in_expert_2 = (
+      common_layers.cumsum(mask_2, axis=1, exclusive=True) + mask_1_count)
+  position_in_expert_2 *= mask_2
+  mask_2 *= _to_float(tf.less(position_in_expert_2, expert_capacity_f))
+  mask_2_count = tf.reduce_sum(mask_2, axis=1, keepdims=True)
+  mask_2_flat = tf.reduce_sum(mask_2, axis=2)
+  position_in_expert_2 = tf.reduce_sum(position_in_expert_2, axis=2)
+  gate_2 *= mask_2_flat
+
+  # What fraction didn't fit - show summaries
+  if show_summaries:
+    miss_rate_1 = (
+        1.0 - tf.reduce_sum(mask_1_count) / _to_float(batch * length))
+    miss_rate_2 = (
+        1.0 - tf.reduce_sum(mask_2_count) / _to_float(batch * length))
+    tf.summary.scalar("miss_rate_1", miss_rate_1)
+    tf.summary.scalar("miss_rate_2", miss_rate_2)
+
+  # renormalize the two gate values to add up to 1
+  denom = gate_1 + gate_2 + 1e-9
+  gate_1 /= denom
+  gate_2 /= denom
+
+  # inputs: [batch, length, input_size]
+  # forward_assignment: [batch, length, num_experts * expert_capacity]
+  # expert_inputs: [batch, num_experts * expert_capacity, input_size]
+
+  segment_ids_forward_1 = (
+      (index_1 * expert_capacity) +
+      tf.to_int32(position_in_expert_1) +
+      tf.to_int32(1.0 - mask_1_flat) * (num_experts * expert_capacity))
+
+  segment_ids_forward_2 = (
+      (index_2 * expert_capacity) +
+      tf.to_int32(position_in_expert_2) +
+      tf.to_int32(1.0 - mask_2_flat) * (num_experts * expert_capacity))
+
+  # Gather and scatter are painfully slow on TPU.
+  # We will use one_hot and matmul instead.
+
+  # [batch, length, num_experts * expert_capacity]
+  one_hot_1 = tf.one_hot(
+      segment_ids_forward_1, num_experts * expert_capacity, dtype=gates.dtype)
+  one_hot_2 = tf.one_hot(
+      segment_ids_forward_2, num_experts * expert_capacity, dtype=gates.dtype)
+
+  # expert_output: [batch, num_experts * expert_capacity, output_size]
+  # backward_assignment: [batch, length, num_experts * expert_capacity]
+  # output: [batch, length, output_size]
+  assignment = (
+      one_hot_1 * tf.cast(tf.expand_dims(gate_1, 2), gates.dtype) +
+      one_hot_2 * tf.cast(tf.expand_dims(gate_2, 2), gates.dtype))
+
+  assignment = tf.reshape(
+      assignment, [batch, length, num_experts, expert_capacity])
+
+  return assignment
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers_test.py b/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
new file mode 100644
index 000000000..d2592c7c3
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
@@ -0,0 +1,292 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 Mesh TensorFlow layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+
+from tensor2tensor.layers import common_layers
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+from tensor2tensor.mesh_tensorflow import mtf_layers
+from tensor2tensor.mesh_tensorflow import placement_mesh_impl
+
+import tensorflow as tf
+
+
+class MtfLayersTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(
+      (4, True),
+      (8, False),
+  )
+  def testDense(self, units, use_bias):
+    batch = 2
+    channels = 3
+    inputs = tf.random_normal([batch, channels])
+
+    graph = mtf.Graph()
+    mesh = mtf.Mesh(graph, "my_mesh")
+    batch_dim = mtf.Dimension("batch", batch)
+    channels_dim = mtf.Dimension("channels", channels)
+    depth_dim = mtf.Dimension("depth", units)
+
+    mtf_inputs = mtf.infeed(mesh, inputs,
+                            shape=mtf.TensorShape([batch_dim, channels_dim]))
+    mtf_outputs = mtf_layers.dense(mtf_inputs,
+                                   output_dim=depth_dim,
+                                   reduced_dims=[channels_dim],
+                                   activation=mtf.relu,
+                                   use_bias=use_bias)
+    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+        shape=[1], layout={}, devices=[""])
+    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
+    actual_outputs = lowering.outfeed(mtf_outputs)
+
+    expected_outputs = tf.keras.layers.Dense(units=units,
+                                             activation=tf.nn.relu,
+                                             use_bias=use_bias)(inputs)
+    tf_group = lowering.copy_masters_to_slices()
+    init = tf.global_variables_initializer()
+    with self.test_session() as sess:
+      sess.run(init)
+      sess.run(tf_group)
+      actual, expected = sess.run([actual_outputs, expected_outputs])
+
+    self.assertEqual(actual.shape, expected.shape)
+
+  def testLayerNorm(self):
+    batch = 2
+    channels = 3
+    inputs = tf.random_normal([batch, channels])
+
+    graph = mtf.Graph()
+    mesh = mtf.Mesh(graph, "my_mesh")
+    batch_dim = mtf.Dimension("batch", batch)
+    channels_dim = mtf.Dimension("channels", channels)
+
+    mtf_inputs = mtf.infeed(mesh, inputs,
+                            shape=mtf.TensorShape([batch_dim, channels_dim]))
+    mtf_outputs = mtf_layers.layer_norm(mtf_inputs,
+                                        dim=channels_dim)
+    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+        shape=[1], layout={}, devices=[""])
+    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
+    actual_outputs = lowering.outfeed(mtf_outputs)
+
+    expected_outputs = common_layers.layer_norm(inputs)
+    tf_group = lowering.copy_masters_to_slices()
+    init = tf.global_variables_initializer()
+    with self.test_session() as sess:
+      sess.run(init)
+      sess.run(tf_group)
+      actual, expected = sess.run([actual_outputs, expected_outputs])
+
+    self.assertEqual(actual.shape, expected.shape)
+
+  def testWeightsNonzero(self):
+    inputs = tf.constant([[3, 1, 0], [1, 0, 0]])
+
+    graph = mtf.Graph()
+    mesh = mtf.Mesh(graph, "my_mesh")
+    batch_dim = mtf.Dimension("batch", inputs.shape.as_list()[0])
+    channels_dim = mtf.Dimension("channels", inputs.shape.as_list()[1])
+
+    mtf_inputs = mtf.infeed(mesh, inputs,
+                            shape=mtf.TensorShape([batch_dim, channels_dim]))
+    mtf_outputs = mtf_layers.weights_nonzero(mtf_inputs)
+    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+        shape=[1], layout={}, devices=[""])
+    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
+    actual_outputs = lowering.outfeed(mtf_outputs)
+
+    expected_outputs = common_layers.weights_nonzero(inputs)
+    tf_group = lowering.copy_masters_to_slices()
+    with self.test_session() as sess:
+      sess.run(tf_group)
+      actual, expected = sess.run([actual_outputs, expected_outputs])
+
+    self.assertAllEqual(actual, expected)
+
+  def testDenseReluDense(self):
+    batch = 2
+    channels = 3
+    hidden = 5
+    inputs = tf.random_normal([batch, channels])
+
+    graph = mtf.Graph()
+    mesh = mtf.Mesh(graph, "my_mesh")
+    batch_dim = mtf.Dimension("batch", batch)
+    channels_dim = mtf.Dimension("channels", channels)
+    hidden_dim = mtf.Dimension("hidden", hidden)
+
+    mtf_inputs = mtf.infeed(mesh, inputs,
+                            shape=mtf.TensorShape([batch_dim, channels_dim]))
+    mtf_outputs = mtf_layers.dense_relu_dense(mtf_inputs,
+                                              hidden_channels=hidden_dim)
+    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+        shape=[1], layout={}, devices=[""])
+    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
+    actual_outputs = lowering.outfeed(mtf_outputs)
+
+    tf_group = lowering.copy_masters_to_slices()
+    init = tf.global_variables_initializer()
+    with self.test_session() as sess:
+      sess.run(init)
+      sess.run(tf_group)
+      actual = sess.run(actual_outputs)
+
+    self.assertEqual(actual.shape, inputs.shape)
+
+  @parameterized.parameters(
+      (4, 2),
+  )
+  def testMaskedLocalAttention1D(self, kv_channels, heads):
+    batch = 2
+    length_q = 16
+    length_m = 16
+    channels = 3
+    query = tf.random_normal([batch, length_q, channels])
+    memory = tf.random_normal([batch, length_m, channels])
+
+    graph = mtf.Graph()
+    mesh = mtf.Mesh(graph, "my_mesh")
+    batch_dim = mtf.Dimension("batch", batch)
+    length_q_dim = mtf.Dimension("length_q", length_q)
+    length_m_dim = mtf.Dimension("length_m", length_m)
+    channels_dim = mtf.Dimension("channels", channels)
+    kv_channels_dim = mtf.Dimension("kv_channels", kv_channels)
+    heads_dim = mtf.Dimension("heads", heads)
+
+    mtf_query = mtf.infeed(
+        mesh, query,
+        shape=mtf.TensorShape([batch_dim, length_q_dim, channels_dim]))
+    mtf_memory = mtf.infeed(
+        mesh, memory,
+        shape=mtf.TensorShape([batch_dim, length_m_dim, channels_dim]))
+    mtf_outputs = mtf_layers.masked_local_attention_1d(
+        mtf_query,
+        mtf_memory,
+        kv_channels=kv_channels_dim,
+        heads=heads_dim,
+        block_length=2)
+    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+        shape=[1], layout={}, devices=[""])
+    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
+    actual_outputs = lowering.outfeed(mtf_outputs)
+
+    tf_group = lowering.copy_masters_to_slices()
+    init = tf.global_variables_initializer()
+    with self.test_session() as sess:
+      sess.run(init)
+      sess.run(tf_group)
+      actual = sess.run(actual_outputs)
+
+    self.assertEqual(actual.shape, (batch, length_q, channels))
+
+  @parameterized.parameters(
+      (2, 4, 5, 7, 3, 1),
+  )
+  def testDotProductAttention(
+      self, batch, heads, length_q, length_kv, depth_k, depth_v):
+    query = tf.random_normal([batch, heads, length_q, depth_k])
+    key = tf.random_normal([batch, heads, length_kv, depth_k])
+    value = tf.random_normal([batch, heads, length_kv, depth_v])
+
+    graph = mtf.Graph()
+    mesh = mtf.Mesh(graph, "my_mesh")
+    batch_dim = mtf.Dimension("batch", batch)
+    heads_dim = mtf.Dimension("heads", heads)
+    length_q_dim = mtf.Dimension("length_q", length_q)
+    length_kv_dim = mtf.Dimension("length_kv", length_kv)
+    depth_k_dim = mtf.Dimension("depth_k", depth_k)
+    depth_v_dim = mtf.Dimension("depth_v", depth_v)
+
+    mtf_query = mtf.infeed(
+        mesh, query,
+        shape=mtf.TensorShape(
+            [batch_dim, heads_dim, length_q_dim, depth_k_dim]))
+    mtf_key = mtf.infeed(
+        mesh, key,
+        shape=mtf.TensorShape(
+            [batch_dim, heads_dim, length_kv_dim, depth_k_dim]))
+    mtf_value = mtf.infeed(
+        mesh, value,
+        shape=mtf.TensorShape(
+            [batch_dim, heads_dim, length_kv_dim, depth_v_dim]))
+    mtf_outputs = mtf_layers.dot_product_attention(
+        mtf_query,
+        mtf_key,
+        mtf_value,
+        mask=None)
+    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+        shape=[1], layout={}, devices=[""])
+    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
+    actual_outputs = lowering.outfeed(mtf_outputs)
+
+    tf_group = lowering.copy_masters_to_slices()
+    init = tf.global_variables_initializer()
+    with self.test_session() as sess:
+      sess.run(init)
+      sess.run(tf_group)
+      actual = sess.run(actual_outputs)
+
+    self.assertEqual(actual.shape, (batch, heads, length_q, depth_v))
+
+  @parameterized.parameters(
+      (16, 4),
+      (32, 8),
+  )
+  def testMultiheadAttention(self, kv_channels, heads):
+    batch = 2
+    length = 8
+    channels = 3
+    query = tf.random_normal([batch, length, channels])
+
+    graph = mtf.Graph()
+    mesh = mtf.Mesh(graph, "my_mesh")
+    batch_dim = mtf.Dimension("batch", batch)
+    length_dim = mtf.Dimension("length", length)
+    channels_dim = mtf.Dimension("channels", channels)
+    kv_channels_dim = mtf.Dimension("kv_channels", kv_channels)
+    heads_dim = mtf.Dimension("heads", heads)
+
+    mtf_query = mtf.infeed(
+        mesh, query,
+        shape=mtf.TensorShape([batch_dim, length_dim, channels_dim]))
+    mtf_outputs = mtf_layers.multihead_attention(
+        mtf_query,
+        memory_antecedent=None,
+        mask=None,
+        kv_channels=kv_channels_dim,
+        heads=heads_dim)
+    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+        shape=[1], layout={}, devices=[""])
+    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
+    actual_outputs = lowering.outfeed(mtf_outputs)
+
+    tf_group = lowering.copy_masters_to_slices()
+    init = tf.global_variables_initializer()
+    with self.test_session() as sess:
+      sess.run(init)
+      sess.run(tf_group)
+      actual = sess.run(actual_outputs)
+
+    self.assertEqual(actual.shape, query.shape)
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/mesh_tensorflow/mtf_model.py b/tensor2tensor/mesh_tensorflow/mtf_model.py
new file mode 100644
index 000000000..59f975d83
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_model.py
@@ -0,0 +1,287 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Mesh-Tensorflow Model in tensor2tensor."""
+
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+import copy
+import six
+
+
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+from tensor2tensor.mesh_tensorflow import mtf_optimize
+from tensor2tensor.mesh_tensorflow import mtf_utils
+from tensor2tensor.mesh_tensorflow import placement_mesh_impl
+from tensor2tensor.mesh_tensorflow import simd_mesh_impl
+from tensor2tensor.utils import learning_rate
+from tensor2tensor.utils import metrics
+from tensor2tensor.utils import t2t_model
+
+import tensorflow as tf
+
+from tensorflow.contrib.tpu.python.tpu import tpu_estimator
+
+
+class MtfModel(t2t_model.T2TModel):
+  """Toy model to test mesh_tensorflow."""
+
+  @classmethod
+  def estimator_model_fn(cls,
+                         hparams,
+                         features,
+                         labels,
+                         mode,
+                         config=None,
+                         params=None,
+                         decode_hparams=None,
+                         use_tpu=False,
+                         xla_compile=False):
+    hparams = copy.deepcopy(hparams)
+    hparams.use_tpu = use_tpu
+    # merge decode_hparams into hparams if present
+    if mode == tf.estimator.ModeKeys.PREDICT and decode_hparams is not None:
+      for k, v in six.iteritems(decode_hparams.values()):
+        if hasattr(hparams, k) and getattr(hparams, k) != v:
+          tf.logging.warning("Overriding hparams.%s with %s from decode_hparams"
+                             % (k, v))
+        setattr(hparams, k, v)
+
+    # Instantiate model
+    data_parallelism = None
+    if not use_tpu and config:
+      data_parallelism = config.data_parallelism
+    model = cls(
+        hparams,
+        mode,
+        data_parallelism=data_parallelism,
+        decode_hparams=decode_hparams)
+
+    global_step = tf.train.get_global_step()
+    graph = mtf.Graph()
+    mesh = mtf.Mesh(graph, "my_mesh")
+
+    mesh_shape = mtf.parse_mesh_shape(hparams.mesh_shape)
+    mesh_size = mtf.list_product(mesh_shape)
+    if use_tpu:
+      mesh_devices = [""] * mesh_size
+      mesh_impl = simd_mesh_impl.SimdMeshImpl(
+          mesh_shape, mtf.parse_layout(hparams.layout), mesh_devices,
+          params["context"].device_assignment)
+    else:
+      if len(data_parallelism.ps_devices) == 1:
+        mesh_devices = [""] * mesh_size
+      else:
+        assert len(data_parallelism.ps_devices) == mesh_size
+        mesh_devices = data_parallelism.ps_devices
+      mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+          mesh_shape, mtf.parse_layout(hparams.layout), mesh_devices)
+
+    # PREDICT mode
+    if mode == tf.estimator.ModeKeys.PREDICT:
+      return model.estimator_spec_predict(features, mesh, mesh_impl, use_tpu)
+
+    logits, loss = model.mtf_model_fn(features, mesh)
+    if use_tpu and logits is not None:
+      logits = mtf.anonymize(logits)
+
+    # TRAIN mode
+    if mode == tf.estimator.ModeKeys.TRAIN:
+      var_grads = mtf.gradients(
+          [loss], [v.outputs[0] for v in graph.trainable_variables])
+      lr = learning_rate.learning_rate_schedule(hparams)
+      mtf_lr = mtf.infeed(
+          mesh, tf.convert_to_tensor(lr, dtype=tf.float32), mtf.TensorShape([]))
+      optimizer = mtf_optimize.make_optimizer(hparams, mtf_lr)
+      update_ops = []
+      for grad, var in zip(var_grads, graph.trainable_variables):
+        update_ops.extend(optimizer.apply_grad(grad, var))
+
+    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
+
+    tf_loss = lowering.outfeed(loss)
+    tf_loss = tf.to_float(tf_loss)
+    if logits and mode != tf.estimator.ModeKeys.TRAIN:
+      tf_logits = lowering.outfeed(logits)
+
+    if mode == tf.estimator.ModeKeys.TRAIN:
+      tf_update_ops = [lowering.lowered_operation(op) for op in update_ops]
+      tf_update_ops.append(tf.assign_add(global_step, 1))
+      # tf.logging.info("tf_update_ops: {}".format(tf_update_ops))
+      train_op = tf.group(tf_update_ops)
+
+    with mtf_utils.outside_all_rewrites():
+      # Copy master variables to slices. Must be called first.
+      restore_hook = mtf.MtfRestoreHook(lowering)
+      saver = tf.train.Saver(
+          tf.global_variables(),
+          sharded=True,
+          max_to_keep=10,
+          keep_checkpoint_every_n_hours=2,
+          defer_build=False,
+          save_relative_paths=True)
+      tf.add_to_collection(tf.GraphKeys.SAVERS, saver)
+      saver_listener = mtf.MtfCheckpointSaverListener(lowering)
+      saver_hook = tf.train.CheckpointSaverHook(
+          hparams.model_dir,
+          save_steps=1000,
+          saver=saver,
+          listeners=[saver_listener])
+
+    # EVAL mode
+    if mode == tf.estimator.ModeKeys.EVAL:
+      tf_logits = lowering.outfeed(logits)
+      return model.estimator_spec_eval(features, tf_logits, labels, tf_loss,
+                                       restore_hook, use_tpu)
+
+    if use_tpu:
+      _remove_summaries()
+      return tpu_estimator.TPUEstimatorSpec(
+          mode=tf.estimator.ModeKeys.TRAIN,
+          loss=tf_loss,
+          train_op=train_op,
+          training_hooks=[restore_hook, saver_hook])
+    else:
+      return tf.estimator.EstimatorSpec(
+          tf.estimator.ModeKeys.TRAIN, loss=tf_loss, train_op=train_op,
+          training_chief_hooks=[restore_hook, saver_hook])
+
+  def estimator_spec_eval(
+      self, features, logits, labels, loss, restore_hook, use_tpu):
+    """Construct EstimatorSpec for EVAL mode."""
+    hparams = self.hparams
+    problem = hparams.problem
+    if logits.get_shape().ndims == 3:
+      logits = tf.expand_dims(tf.expand_dims(logits, 2), 3)
+    eval_metrics_fns = metrics.create_evaluation_metrics([problem], hparams)
+
+    if use_tpu:
+      def metric_fn(tf_logits, labels):
+        with tf.device("cpu:0"), mtf_utils.outside_all_rewrites():
+          eval_metrics = {}
+          for metric_name, metric_fn in six.iteritems(eval_metrics_fns):
+            if metric_name.split("/")[-1] not in t2t_model.TPU_METRIC_BLACKLIST:
+              eval_metrics[metric_name] = metric_fn(
+                  tf_logits, None, tf.identity(labels))
+          return eval_metrics
+      return tpu_estimator.TPUEstimatorSpec(
+          tf.estimator.ModeKeys.EVAL,
+          evaluation_hooks=[restore_hook],
+          loss=loss,
+          eval_metrics=(metric_fn, [logits, labels]))
+    else:
+      eval_metrics = {}
+      predictions = {"predictions": logits}
+      for metric_name, metric_fn in six.iteritems(eval_metrics_fns):
+        eval_metrics[metric_name] = metric_fn(logits, features,
+                                              features["targets"])
+
+      return tf.estimator.EstimatorSpec(
+          tf.estimator.ModeKeys.EVAL,
+          predictions=predictions,
+          eval_metric_ops=eval_metrics,
+          evaluation_hooks=[restore_hook],
+          loss=loss)
+
+  def estimator_spec_predict(self, features, mesh, mesh_impl, use_tpu):
+    mtf_samples = self.sample(features, mesh)
+    lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
+    outputs = lowering.outfeed(mtf_samples)
+    if self.has_input:
+      ndims = len(outputs.shape.as_list())
+      actual_batch_size = tf.shape(features["inputs"])[0]
+      outputs = tf.slice(
+          outputs, [0] * ndims, [actual_batch_size] + [-1] * (ndims - 1))
+    predictions = {
+        "outputs": outputs,
+        "targets": features.get("infer_targets", features.get("inputs")),
+        "inputs": features.get("inputs"),
+    }
+    if use_tpu:
+      _remove_summaries()
+      return tpu_estimator.TPUEstimatorSpec(
+          mode=tf.estimator.ModeKeys.PREDICT,
+          predictions=predictions,
+          prediction_hooks=[mtf.MtfRestoreHook(lowering)])
+    else:
+      return tf.estimator.EstimatorSpec(
+          tf.estimator.ModeKeys.PREDICT,
+          predictions=predictions,
+          prediction_hooks=[mtf.MtfRestoreHook(lowering)])
+
+  def sample(self, features, mesh):
+    """Sample from the model."""
+    raise NotImplementedError("TODO(noam): write generic slow mtf sample.")
+
+  def mtf_model_fn(self, features, mesh):
+    raise NotImplementedError("Not implemented")
+
+
+def _remove_summaries():
+  g = tf.get_default_graph()
+  key = tf.GraphKeys.SUMMARIES
+  del g.get_collection_ref(key)[:]
+  assert not g.get_collection(key)
+
+
+def _create_host_call(model_dir):
+  """Construct a host_call writing scalar summaries.
+
+  Args:
+    model_dir: String containing path to train
+
+  Returns:
+    (fn, args) Pair to be called by TPUEstimator as the host_call.
+  """
+  graph = tf.get_default_graph()
+  summaries = graph.get_collection(tf.GraphKeys.SUMMARIES)
+
+  gs_t = tf.reshape(tf.to_int32(tf.train.get_global_step()), [1])
+  summary_kwargs = collections.OrderedDict()
+  for t in summaries:
+    if t.op.type != "ScalarSummary":
+      continue
+
+    name = t.op.name
+    tensor = t.op.inputs[1]
+    assert tensor.shape.is_compatible_with([])
+    if tensor.dtype == tf.int64:
+      tensor = tf.to_int32(tensor)
+    summary_kwargs[name] = tf.reshape(tensor, [1])
+  summary_kwargs["global_step"] = gs_t
+
+  def host_call_fn(**kwargs):
+    """Training host call. Creates scalar summaries for training metrics.
+
+    Args:
+      **kwargs: Dict of {str: Tensor} , with `Tensor` of shape `[batch]`. Must
+        contain key "global_step" with value of current global_step Tensor.
+
+    Returns:
+      List of summary ops to run on the CPU host.
+    """
+    gs = tf.to_int64(kwargs.pop("global_step")[0])
+    with tf.contrib.summary.create_file_writer(model_dir).as_default():
+      with tf.contrib.summary.always_record_summaries():
+        for name, value in sorted(six.iteritems(kwargs)):
+          tf.contrib.summary.scalar(
+              name, tf.reduce_mean(tf.to_float(value)), step=gs)
+
+        return tf.contrib.summary.all_summary_ops()
+
+  return (host_call_fn, summary_kwargs)
diff --git a/tensor2tensor/mesh_tensorflow/mtf_optimize.py b/tensor2tensor/mesh_tensorflow/mtf_optimize.py
new file mode 100644
index 000000000..df822fb27
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_optimize.py
@@ -0,0 +1,269 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Mesh-Tensorflow Optimizers."""
+
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+import tensorflow as tf
+
+
+def make_optimizer(hparams, lr):
+  if hparams.optimizer == "SGD":
+    return SgdOptimizer(lr)
+  elif hparams.optimizer == "Adafactor":
+    return adafactor_optimizer_from_hparams(hparams, lr)
+  else:
+    raise ValueError("Unknown Optimizer")
+
+
+class Optimizer(object):
+  """Base optmizer class."""
+
+  def apply_grad(self, grad, var):
+    raise ValueError("Apply_Grad not implemented %s %s" % (grad, var))
+
+
+class SgdOptimizer(Optimizer):
+  """oOptimizer implementing SGD."""
+
+  def __init__(self, lr):
+    self._lr = lr
+
+  @property
+  def lr(self):
+    return self._lr
+
+  def apply_grad(self, grad, var):
+    return [mtf.assign(var, var.outputs[0] - (grad * self.lr))]
+
+
+class AdafactorOptimizer(Optimizer):
+  """Adafactor."""
+
+  def __init__(self,
+               multiply_by_parameter_scale=True,
+               learning_rate=None,
+               decay_rate=None,
+               beta1=0.0,
+               clipping_threshold=1.0,
+               factored=True,
+               epsilon1=1e-30,
+               epsilon2=1e-3):
+    """Construct a new Adafactor optimizer.
+
+    See class comment.
+
+    Args:
+      multiply_by_parameter_scale: a boolean
+      learning_rate: an optional Scalar.
+      decay_rate: an optional Scalar.
+      beta1: a float value between 0 and 1
+      clipping_threshold: an optional float >= 1
+      factored: a boolean - whether to use factored second-moment estimator
+        for 2d variables
+      epsilon1: Regularization constant for squared gradient.
+      epsilon2: Regularization constant for parameter scale.
+
+    Raises:
+      ValueError: if absolute_update_scale and relative_update_scale_fn are both
+        present or both absent.
+    """
+    self._multiply_by_parameter_scale = multiply_by_parameter_scale
+    if learning_rate is None:
+      learning_rate = self._learning_rate_default(multiply_by_parameter_scale)
+    self._learning_rate = learning_rate
+    if decay_rate is None:
+      decay_rate = self._decay_rate_default()
+    self._decay_rate = decay_rate
+    self._beta1 = beta1
+    self._clipping_threshold = clipping_threshold
+    self._factored = factored
+    self._epsilon1 = epsilon1
+    self._epsilon2 = epsilon2
+
+  def _factored_dims(self, shape):
+    """Should we use a factored second moment estimator.
+
+    Based on the shape of the variable.
+    If we factor the accumulator, then this function returns a list of two
+    mtf.Dimensions to reduce over.  We always pick the two largest dimensions.
+    If there are not two dimensions of size >=128, then we do not factor.
+
+    Args:
+      shape: a TensorShape
+    Returns:
+      either a list of 2 Dimensions or None
+    """
+    if not self._factored or shape.ndims < 2:
+      return None
+    sorted_dims = sorted(shape.dims, key=lambda d: -d.size)
+    if sorted_dims[1].size < 128:
+      return None
+    return sorted_dims[:2]
+
+  def _parameter_scale(self, var):
+    """Estimate the scale of the parameters from the current values.
+
+    We include a minimum value of 0.001 to give it a chance to escape 0
+    if it was zero-initialized.
+
+    Instead of using the value, we could impute the scale from the shape,
+    as initializers do.
+
+    Args:
+      var: a variable or Tensor.
+    Returns:
+      a Scalar
+    """
+    return mtf.maximum(reduce_rms(var), self._epsilon2)
+
+  def apply_grad(self, grad, var):
+    # create slots
+    factored_dims = self._factored_dims(var.shape)
+    if factored_dims:
+      d0, d1 = factored_dims
+      vr_shape = var.shape - d0
+      vc_shape = var.shape - d1
+      vr = mtf.get_variable(
+          var.mesh, var.name + "_slot_vr", vr_shape,
+          initializer=tf.zeros_initializer(), trainable=False)
+      vc = mtf.get_variable(
+          var.mesh, var.name + "_slot_vc", vc_shape,
+          initializer=tf.zeros_initializer(), trainable=False)
+    else:
+      v = mtf.get_variable(
+          var.mesh, var.name + "_slot_v", var.shape,
+          initializer=tf.zeros_initializer(), trainable=False)
+    if self._beta1:
+      m = mtf.get_variable(
+          var.mesh, var.name + "_slot_m", var.shape,
+          iniitalizer=tf.zeros_initializer(), trainable=False)
+
+    with tf.variable_scope(var.name + "/adafactor"):
+      grad_squared = mtf.square(grad) + self._epsilon1
+      decay_rate = self._decay_rate
+      old_val = var.value
+      if self._multiply_by_parameter_scale:
+        update_scale = self._parameter_scale(old_val) * self._learning_rate
+      else:
+        update_scale = self._learning_rate
+      mixing_rate = 1.0 - decay_rate
+      updates = []
+      if factored_dims:
+        grad_squared_row_mean = mtf.reduce_mean(
+            grad_squared, output_shape=vr_shape)
+        grad_squared_col_mean = mtf.reduce_mean(
+            grad_squared, output_shape=vc_shape)
+        new_vr = vr * decay_rate + grad_squared_row_mean * mixing_rate
+        new_vc = vc * decay_rate + grad_squared_col_mean * mixing_rate
+        vr_update = mtf.assign(vr, new_vr)
+        vc_update = mtf.assign(vc, new_vc)
+        updates.extend([vr_update, vc_update])
+        long_term_mean = mtf.reduce_mean(new_vr, reduced_dim=d1)
+        r_factor = mtf.rsqrt(new_vr / long_term_mean)
+        c_factor = mtf.rsqrt(new_vc)
+        x = grad * r_factor * c_factor
+      else:
+        new_v = v * decay_rate + grad_squared * mixing_rate
+        v_update = mtf.assign(v, new_v)
+        updates.append(v_update)
+        x = grad * mtf.rsqrt(new_v)
+      if self._clipping_threshold is not None:
+        clipping_denom = mtf.maximum(
+            1.0, reduce_rms(x) / self._clipping_threshold)
+        x /= clipping_denom
+      subtrahend = x * update_scale
+      if self._beta1:
+        new_m = self._beta1 * m.value + (1.0 - self._beta1) * subtrahend
+        subtrahend = new_m
+        updates.append(mtf.assign(m, new_m))
+      new_val = old_val - subtrahend
+      var_update = mtf.assign(var, new_val)
+      updates.append(var_update)
+      return updates
+
+  def _decay_rate_default(self):
+    return adafactor_decay_rate_pow(0.8)
+
+  def _learning_rate_default(self, multiply_by_parameter_scale):
+    learning_rate = tf.minimum(tf.rsqrt(step_num() + 1.0), 0.01)
+    if not multiply_by_parameter_scale:
+      learning_rate *= 0.05
+    return learning_rate
+
+
+def adafactor_decay_rate_adam(beta2):
+  """Second-moment decay rate like Adam, subsuming the correction factor.
+
+  Args:
+    beta2: a float between 0 and 1
+  Returns:
+    a scalar
+  """
+  t = tf.to_float(tf.train.get_or_create_global_step()) + 1.0
+  decay = beta2 * (1.0 - tf.pow(beta2, t - 1.0)) / (1.0 - tf.pow(beta2, t))
+  return decay
+
+
+def adafactor_decay_rate_pow(exponent):
+  """Second moment decay rate where memory-length grows as step_num^exponent.
+
+  Args:
+    exponent: a float between 0 and 1
+  Returns:
+    a scalar
+  """
+  return 1.0 - tf.pow((step_num() + 1.0), -exponent)
+
+
+def step_num():
+  return tf.to_float(tf.train.get_or_create_global_step())
+
+
+def adafactor_optimizer_from_hparams(hparams, lr):
+  """Create an Adafactor optimizer based on model hparams.
+
+  Args:
+    hparams: model hyperparameters
+    lr: learning rate scalar.
+  Returns:
+    an AdafactorOptimizer
+  Raises:
+    ValueError: on illegal values
+  """
+  if hparams.optimizer_adafactor_decay_type == "Adam":
+    decay_rate = adafactor_decay_rate_adam(
+        hparams.optimizer_adafactor_beta2)
+  elif hparams.optimizer_adafactor_decay_type == "pow":
+    decay_rate = adafactor_decay_rate_pow(
+        hparams.optimizer_adafactor_memory_exponent)
+  else:
+    raise ValueError("unknown optimizer_adafactor_decay_type")
+  return AdafactorOptimizer(
+      multiply_by_parameter_scale=(
+          hparams.optimizer_adafactor_multiply_by_parameter_scale),
+      learning_rate=lr,
+      decay_rate=decay_rate,
+      beta1=hparams.optimizer_adafactor_beta1,
+      clipping_threshold=hparams.optimizer_adafactor_clipping_threshold,
+      factored=hparams.optimizer_adafactor_factored)
+
+
+def reduce_rms(x):
+  return mtf.sqrt(mtf.reduce_mean(mtf.square(x)))
diff --git a/tensor2tensor/mesh_tensorflow/mtf_toy.py b/tensor2tensor/mesh_tensorflow/mtf_toy.py
new file mode 100644
index 000000000..811d47086
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_toy.py
@@ -0,0 +1,178 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Toy model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_hparams
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+from tensor2tensor.mesh_tensorflow import mtf_layers
+from tensor2tensor.mesh_tensorflow import mtf_model
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+from tensorflow.contrib.tpu.python.ops import tpu_ops
+
+
+@registry.register_model
+class MtfToy(mtf_model.MtfModel):
+  """Toy model to test mesh_tensorflow."""
+
+  def mtf_model_fn(self, features, mesh):
+    hparams = self._hparams
+    # tf_x = tf.random_uniform([hparams.batch_size, hparams.io_size])
+    tf_x = tf.matmul(
+        tf.reshape(
+            tf.lin_space(0., 1.0, hparams.batch_size), [hparams.batch_size, 1]),
+        tf.reshape(
+            tf.lin_space(0., 1.0, hparams.io_size), [1, hparams.io_size]))
+    batch_dim = mtf.Dimension("batch", hparams.batch_size)
+
+    hidden_dim = mtf.Dimension("hidden", hparams.hidden_size)
+    io_dim = mtf.Dimension("io", hparams.io_size)
+    x = mtf.infeed_fully_replicated(
+        mesh, tf_x, mtf.TensorShape([batch_dim, io_dim]))
+    h = mtf_layers.dense(x, hidden_dim, name="layer1", use_bias=False)
+    y = mtf_layers.dense(h, io_dim, name="layer2", use_bias=False)
+
+    loss = mtf.reduce_sum(mtf.square(y - x))
+    return None, loss
+
+
+@registry.register_model
+class MtfSimple(mtf_model.MtfModel):
+  """Toy model to test mesh_tensorflow."""
+
+  def mtf_model_fn(self, features, mesh):
+    hparams = self._hparams
+    # tf_x = tf.random_uniform([hparams.batch_size, hparams.io_size])
+    tf_x = tf.matmul(
+        tf.reshape(
+            tf.lin_space(0., 1.0, hparams.batch_size), [hparams.batch_size, 1]),
+        tf.reshape(
+            tf.lin_space(0., 1.0, hparams.io_size), [1, hparams.io_size]))
+    batch_dim = mtf.Dimension("batch", hparams.batch_size)
+    hidden_dim = mtf.Dimension("hidden", hparams.hidden_size)
+    io_dim = mtf.Dimension("io", hparams.io_size)
+
+    x = mtf.infeed_fully_replicated(
+        mesh, tf_x, mtf.TensorShape([batch_dim, io_dim]))
+    h = mtf_layers.dense(x, hidden_dim, name="layer1", use_bias=False)
+    y = mtf_layers.dense(h, io_dim, name="layer2", use_bias=False)
+    loss = mtf.reduce_sum(mtf.square(y - x))
+    return None, loss
+
+
+@registry.register_model
+class MtfToyNormal(mtf_model.MtfModel):
+  """Toy model to test mesh_tensorflow."""
+
+  def mtf_model_fn(self, features, mesh):
+    hparams = self._hparams
+    hparams.batch_size = 10
+    hparams.io_size = 4
+    hparams.hidden_size = 2
+    tf_x = tf.matmul(
+        tf.reshape(
+            tf.lin_space(0., 1.0, hparams.batch_size), [hparams.batch_size, 1]),
+        tf.reshape(
+            tf.lin_space(0., 1.0, hparams.io_size), [1, hparams.io_size]))
+    # tf_x = tf.random_uniform([hparams.batch_size, hparams.io_size])
+
+    hidden_1_variable = tf.get_variable(
+        "a",
+        shape=[hparams.io_size, hparams.hidden_size],
+        initializer=tf.random_normal_initializer())
+    hidden_2_variable = tf.get_variable(
+        "b",
+        shape=[hparams.hidden_size, hparams.io_size],
+        initializer=tf.random_normal_initializer())
+
+    hidden_layer_1 = tf.matmul(tf_x, hidden_1_variable)
+    hidden_layer_2 = tf.matmul(hidden_layer_1, hidden_2_variable)
+    hidden_layer_2 = tpu_ops.cross_replica_sum(hidden_layer_2)
+    loss = tf.reduce_mean(tf.square(hidden_layer_2 - tf_x))
+    return None, loss
+
+
+def set_sgd_optimizer(hparams):
+  hparams.optimizer = "SGD"
+  hparams.learning_rate_schedule = "constant"
+  hparams.learning_rate_constant = 0.01
+
+
+def set_adafactor_optimizer(hparams):
+  hparams.optimizer = "Adafactor"
+  hparams.learning_rate_schedule = "rsqrt_decay"
+  hparams.optimizer_adafactor_factored = True
+  hparams.learning_rate_warmup_steps = 1000
+
+
+@registry.register_hparams
+def mtf_toy_base():
+  """Set of hyperparameters."""
+  hparams = common_hparams.basic_params1()
+  hparams.no_data_parallelism = True
+  hparams.use_fixed_batch_size = True
+  hparams.add_hparam("mtf_mode", True)
+  hparams.batch_size = 64
+  set_adafactor_optimizer(hparams)
+  hparams.add_hparam("io_size", 32)
+  hparams.hidden_size = 32
+  hparams.add_hparam("mesh_shape", "4.2")
+  hparams.add_hparam("layout", "batch:0;hidden:1")
+  return hparams
+
+
+@registry.register_hparams
+def mtf_toy_data_parallel():
+  """Set of hyperparameters."""
+  hparams = mtf_toy_base()
+  hparams.add_hparam("layout", "batch:0")
+  return hparams
+
+
+@registry.register_hparams
+def mtf_toy_model_parallel():
+  """Set of hyperparameters."""
+  hparams = mtf_toy_base()
+  hparams.add_hparam("layout", "hidden:0")
+  return hparams
+
+
+@registry.register_hparams
+def mtf_toy_data_parallel_m2():
+  """Set of hyperparameters."""
+  hparams = mtf_toy_data_parallel()
+  hparams.mesh_shape = "2"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_toy_model_parallel_m2():
+  """Set of hyperparameters."""
+  hparams = mtf_toy_model_parallel()
+  hparams.mesh_shape = "2"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_toy_m32():
+  """Set of hyperparameters."""
+  hparams = mtf_toy_base()
+  hparams.mesh_shape = "8;4"
+  return hparams
diff --git a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
new file mode 100644
index 000000000..6c4725fce
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
@@ -0,0 +1,219 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""A toy model using mesh-tensrflow."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import google3
+import numpy
+
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+from tensor2tensor.mesh_tensorflow import mtf_layers
+from tensor2tensor.mesh_tensorflow import mtf_optimize
+from tensor2tensor.mesh_tensorflow import mtf_utils
+from tensor2tensor.mesh_tensorflow.simd_mesh_impl import SimdMeshImpl
+import tensorflow as tf
+
+from tensorflow.contrib.tpu.python.tpu import tpu_config
+from tensorflow.contrib.tpu.python.tpu import tpu_estimator
+from tensorflow.python.data.ops.dataset_ops import Dataset
+from tensorflow.python.estimator import estimator as estimator_lib
+from tensorflow.python.platform import flags
+from tensorflow.python.platform import tf_logging as logging
+
+
+FLAGS = flags.FLAGS
+
+tf.flags.DEFINE_integer('batch_size', 64, 'Training batch size.')
+tf.flags.DEFINE_integer('io_size', 2, 'Number of channels per feature.')
+tf.flags.DEFINE_integer('hidden_size', 2, 'Size of each hidden layer.')
+tf.flags.DEFINE_string('mesh_shape', '2;1', 'mesh shape')
+tf.flags.DEFINE_string('layout', 'batch:0', 'computation layout')
+tf.flags.DEFINE_integer('iterations', 100,
+                        'Number of iterations per training loop.')
+tf.flags.DEFINE_integer('train_steps', 10000, 'max steps')
+tf.flags.DEFINE_integer('steps_per_checkpoint', 200, 'steps_per_checkpoint')
+tf.flags.DEFINE_integer('num_shards', 2, 'Number of shards.')
+tf.flags.DEFINE_string('master', 'local',
+                       'BNS name of the TensorFlow master to use.')
+tf.flags.DEFINE_string('evaluation_master', 'local',
+                       'BNS name of the TensorFlow master to use.')
+tf.flags.DEFINE_string(
+    'model_dir',
+    default='',
+    help='The directory where the model will be stored.')
+
+
+class ToyModelInput(object):
+  """Wrapper class that acts as the input_fn to TPUEstimator."""
+
+  def __init__(self):
+    self._num_examples = 10000  # 10k
+    self._images = numpy.random.uniform(
+        0, 1.0, [self._num_examples, FLAGS.io_size]).astype(numpy.float32)
+    self._labels = self._images
+    logging.info('init ToyModelInput()')
+
+  def __call__(self, params):
+    """Input function which provides a single batch for train or eval."""
+    # Retrieves the batch size for the current shard. The # of shards is
+    # computed according to the input pipeline deployment. See
+    # `tf.contrib.tpu.RunConfig` for details.
+    batch_size = params['batch_size']
+    logging.info('call ToyModelInput() with batch size {}'.format(batch_size))
+
+    dataset = Dataset.from_tensor_slices((self._images, self._labels))
+    dataset = dataset.repeat()
+
+    dataset = dataset.prefetch(batch_size)
+    dataset = dataset.apply(
+        tf.contrib.data.batch_and_drop_remainder(batch_size))
+
+    dataset = dataset.prefetch(2)  # Prefetch overlaps in-feed with training
+    return dataset
+
+
+def toy_model(features, mesh):
+  """A toy model implemented by mesh tensorlfow."""
+  batch_dim = mtf.Dimension('batch', FLAGS.batch_size)
+  hidden_dim = mtf.Dimension('hidden', FLAGS.hidden_size)
+  io_dim = mtf.Dimension('io', FLAGS.io_size)
+
+  x = mtf.infeed(mesh, features, mtf.TensorShape([batch_dim, io_dim]))
+  h = mtf_layers.dense(x, hidden_dim, name='layer1', use_bias=False)
+  y = mtf_layers.dense(h, io_dim, name='layer2', use_bias=False)
+
+  loss = mtf.reduce_sum(mtf.square(y - x))
+  return y, loss
+
+
+def model_fn(features, labels, mode, params):
+  """A model is called by TpuEstimator."""
+  del labels
+  global_step = tf.train.get_global_step()
+  graph = mtf.Graph()
+  mesh = mtf.Mesh(graph, 'my_mesh')
+  mesh_shape = mtf.parse_mesh_shape(FLAGS.mesh_shape)
+  mesh_size = mtf.list_product(mesh_shape)
+  mesh_devices = [''] * mesh_size
+  mesh_impl = SimdMeshImpl(mesh_shape, mtf.parse_layout(FLAGS.layout),
+                           mesh_devices, params['context'].device_assignment)
+  with mtf_utils.outside_all_rewrites():
+    logits, loss = toy_model(features, mesh)
+
+  # TRAIN mode
+  if mode == tf.estimator.ModeKeys.TRAIN:
+    var_grads = mtf.gradients([loss],
+                              [v.outputs[0] for v in graph.trainable_variables])
+    optimizer = mtf_optimize.AdafactorOptimizer()
+    update_ops = []
+    for grad, var in zip(var_grads, graph.trainable_variables):
+      update_ops.extend(optimizer.apply_grad(grad, var))
+
+  lowering = mtf.Lowering(graph, {mesh: mesh_impl})
+
+  tf_loss = lowering.outfeed(loss)
+  tf_logits = lowering.outfeed(logits)
+
+  if mode == tf.estimator.ModeKeys.TRAIN:
+    tf_update_ops = [lowering.lowered_operation(op) for op in update_ops]
+    tf_update_ops.append(tf.assign_add(global_step, 1))
+    tf.logging.info('tf_update_ops: {}'.format(tf_update_ops))
+    train_op = tf.group(tf_update_ops)
+
+  with mtf_utils.outside_all_rewrites():
+    # Copy master variables to slices. Must be called first.
+    restore_hook = mtf.MtfRestoreHook(lowering)
+    if mode == tf.estimator.ModeKeys.TRAIN:
+      saver = tf.train.Saver(
+          tf.global_variables(),
+          sharded=True,
+          max_to_keep=10,
+          keep_checkpoint_every_n_hours=2,
+          defer_build=False,
+          save_relative_paths=True)
+      tf.add_to_collection(tf.GraphKeys.SAVERS, saver)
+      saver_listener = mtf.MtfCheckpointSaverListener(lowering)
+      saver_hook = tf.train.CheckpointSaverHook(
+          FLAGS.model_dir,
+          save_steps=1000,
+          saver=saver,
+          listeners=[saver_listener])
+
+      return tpu_estimator.TPUEstimatorSpec(
+          tf.estimator.ModeKeys.TRAIN,
+          loss=tf_loss,
+          train_op=train_op,
+          training_hooks=[restore_hook, saver_hook])
+    elif mode == tf.estimator.ModeKeys.EVAL:
+
+      def metric_fn(tf_logits):
+        mean_logitss = tf.metrics.mean(tf_logits)
+        return {'mean_logitss': mean_logitss}
+
+      eval_metrics = (metric_fn, [tf_logits])
+
+      return tpu_estimator.TPUEstimatorSpec(
+          tf.estimator.ModeKeys.EVAL,
+          evaluation_hooks=[restore_hook],
+          loss=tf_loss,
+          eval_metrics=eval_metrics)
+
+
+def run_toy_model_tpu():
+  """Run a toy model on TPU."""
+  iterations_per_loop = FLAGS.iterations
+  config = tpu_config.RunConfig(
+      master=FLAGS.master,
+      evaluation_master=FLAGS.evaluation_master,
+      model_dir=FLAGS.model_dir,
+      save_checkpoints_steps=None,  # Disable the default saver
+      save_checkpoints_secs=None,  # Disable the default saver
+      log_step_count_steps=iterations_per_loop,
+      tpu_config=tpu_config.TPUConfig(
+          num_shards=FLAGS.num_shards,
+          iterations_per_loop=iterations_per_loop,
+          num_cores_per_replica=1,
+          per_host_input_for_training=tpu_config.InputPipelineConfig.BROADCAST))
+  classifier = tpu_estimator.TPUEstimator(
+      use_tpu=True,
+      model_fn=model_fn,
+      config=config,
+      train_batch_size=FLAGS.batch_size,
+      eval_batch_size=FLAGS.batch_size)
+  current_step = estimator_lib._load_global_step_from_checkpoint_dir(FLAGS.model_dir)  # pylint: disable=protected-access,line-too-long
+  logging.info('Current step %d', current_step)
+  while current_step < FLAGS.train_steps:
+    next_checkpoint = min(current_step + FLAGS.steps_per_checkpoint,
+                          FLAGS.train_steps)
+    classifier.train(input_fn=ToyModelInput(), max_steps=next_checkpoint)
+    current_step = next_checkpoint
+
+    tf.logging.info('Starting to evaluate.')
+    eval_results = classifier.evaluate(
+        input_fn=ToyModelInput(),
+        steps=156)  # since we have 10000 examples and batch_size = 64 per host
+    logging.info('Eval results: %s', eval_results)
+  # classifier.train(input_fn=ToyModelInput(), max_steps=FLAGS.train_steps)
+
+
+def main(_):
+  run_toy_model_tpu()
+
+
+if __name__ == '__main__':
+  tf.app.run()
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
new file mode 100644
index 000000000..550af3391
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
@@ -0,0 +1,895 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Transformer model."""
+
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import copy
+from tensor2tensor.layers import common_hparams
+from tensor2tensor.layers import common_layers
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+from tensor2tensor.mesh_tensorflow import mtf_beam_search
+from tensor2tensor.mesh_tensorflow import mtf_layers
+from tensor2tensor.mesh_tensorflow import mtf_model
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+
+@registry.register_model
+class MtfTransformer(mtf_model.MtfModel):
+  """Transformer in mesh_tensorflow."""
+
+  @property
+  def batch_dim(self):
+    return mtf.Dimension("batch", self._hparams.batch_size)
+
+  @property
+  def inputs_vocab_dim(self):
+    assert self.has_input
+    return mtf.Dimension("vocab", self._inputs_vocab_size)
+
+  @property
+  def targets_vocab_dim(self):
+    return mtf.Dimension("vocab", self._targets_vocab_size)
+
+  @property
+  def model_dim(self):
+    return mtf.Dimension("d_model", self._hparams.d_model)
+
+  @property
+  def max_length_dim(self):
+    return mtf.Dimension("max_length", self._hparams.max_length)
+
+  @property
+  def length_dim(self):
+    return mtf.Dimension("length", self._hparams.max_length)
+
+  @property
+  def memory_length_dim(self):
+    return mtf.Dimension("memory_length", self._hparams.max_length)
+
+  @property
+  def heads_dim(self):
+    return mtf.Dimension("heads", self._hparams.num_heads)
+
+  @property
+  def kv_dim(self):
+    return mtf.Dimension("d_kv", self._hparams.d_kv)
+
+  @property
+  def feedforward_dim(self):
+    return mtf.Dimension("d_ff", self._hparams.d_ff)
+
+  @property
+  def experts_dim(self):
+    return mtf.Dimension("experts", self._hparams.moe_num_experts)
+
+  @property
+  def activation_dtype(self):
+    if self._hparams.activation_dtype == "float32":
+      return tf.float32
+    elif self._hparams.activation_dtype == "bfloat16":
+      return tf.bfloat16
+    else:
+      raise ValueError(
+          "unknown hparams.activation_dtype %s"
+          % self._hparams.activation_dtype)
+
+  def _infeed_to_batch_by_length(self, x, name, mesh, hparams):
+    x = tf.reshape(x, [self.batch_dim.size, self.length_dim.size])
+    return mtf.infeed_fully_replicated(
+        mesh, x, mtf.TensorShape([self.batch_dim, self.length_dim]), name=name)
+
+  def _embedding_and_softmax_vars(self, mesh):
+    hparams = self._hparams
+    targets_embedding_var = mtf.get_variable(
+        mesh, "targets_embedding",
+        mtf.TensorShape([self.targets_vocab_dim, self.model_dim]),
+        initializer=tf.random_normal_initializer(),
+        activation_dtype=self.activation_dtype)
+    if self.has_input:
+      if hparams.shared_embedding:
+        inputs_embedding_var = targets_embedding_var
+      else:
+        inputs_embedding_var = mtf.get_variable(
+            mesh, "inputs_embedding",
+            mtf.TensorShape([self.inputs_vocab_dim, self.model_dim]),
+            initializer=tf.random_normal_initializer(),
+            activation_dtype=self.activation_dtype)
+    else:
+      inputs_embedding_var = None
+    if hparams.shared_embedding_and_softmax_weights:
+      softmax_var = targets_embedding_var * (self.model_dim.size ** -0.5)
+    else:
+      softmax_var = mtf.get_variable(
+          mesh,
+          "softmax",
+          mtf.TensorShape([self.targets_vocab_dim, self.model_dim]),
+          initializer=tf.random_normal_initializer(
+              stddev=self.model_dim.size**-0.5),
+          activation_dtype=self.activation_dtype)
+    positional_embedding_var = mtf.get_variable(
+        mesh, "positional_embedding",
+        mtf.TensorShape([self.max_length_dim, self.model_dim]),
+        initializer=tf.random_normal_initializer(),
+        activation_dtype=self.activation_dtype)
+    return (inputs_embedding_var, targets_embedding_var,
+            softmax_var, positional_embedding_var)
+
+  def _mtf_model_fn(self, features, mesh):
+    features = copy.copy(features)
+    hparams = self._hparams
+    targets = tf.to_int32(features["targets"])
+    if len(targets.get_shape()) > 2:
+      tf.logging.info("targets = %s" % targets)
+      targets = tf.squeeze(targets, [2, 3])
+    # pad targets to max_length
+    def pad_to_max_length(x):
+      extra_length = hparams.max_length - tf.shape(x)[1]
+      x = tf.pad(x, [[0, 0], [0, extra_length]])
+      x = tf.reshape(x, [hparams.batch_size, hparams.max_length])
+      return x
+    targets = pad_to_max_length(targets)
+    for key in ["targets_segmentation", "targets_position",
+                "inputs_segmentation", "inputs_position"]:
+      if key in features:
+        features[key] = pad_to_max_length(features[key])
+    shifted_targets = common_layers.shift_right_2d(targets)
+
+    targets = self._infeed_to_batch_by_length(targets, "targets", mesh, hparams)
+    shifted_targets = self._infeed_to_batch_by_length(
+        shifted_targets, "shifted_targets", mesh, hparams)
+
+    if "targets_segmentation" in features:
+      # "Packed" dataset - keep the examples from seeing each other.
+      targets_segmentation = self._infeed_to_batch_by_length(
+          features["targets_segmentation"], "targets_segmentation",
+          mesh, hparams)
+      targets_position = self._infeed_to_batch_by_length(
+          features["targets_position"], "targets_position",
+          mesh, hparams)
+      decoder_self_attention_mask = (
+          mtf_layers.attention_mask_autoregressive(
+              targets_position, dtype=self.activation_dtype) +
+          mtf_layers.attention_mask_same_segment(
+              targets_segmentation, dtype=self.activation_dtype))
+    else:
+      targets_position = mtf.range(mesh, self.length_dim, dtype=tf.int32)
+      decoder_self_attention_mask = mtf_layers.attention_mask_autoregressive(
+          targets_position, dtype=self.activation_dtype)
+
+    def layer_prepostprocess_dropout(x):
+      return mtf.dropout(
+          x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
+          noise_shape=mtf.TensorShape([self.batch_dim, self.model_dim]))
+
+    extra_losses = []
+    (inputs_embedding_var,
+     targets_embedding_var,
+     softmax_var,
+     positional_embedding_var) = self._embedding_and_softmax_vars(mesh)
+    if self.has_input:
+      inputs = tf.squeeze(tf.to_int32(features["inputs"]), [2, 3])
+      inputs = pad_to_max_length(inputs)
+      inputs = self._infeed_to_batch_by_length(inputs, "inputs", mesh, hparams)
+      if "inputs_segmentation" in features:
+        # "Packed" dataset - keep the examples from seeing each other.
+        inputs_segmentation = self._infeed_to_batch_by_length(
+            features["inputs_segmentation"], "inputs_segmentation",
+            mesh, hparams)
+        inputs_position = self._infeed_to_batch_by_length(
+            features["inputs_position"], "inputs_position",
+            mesh, hparams)
+        encoder_self_attention_mask = (
+            mtf_layers.attention_mask_same_segment(
+                inputs_segmentation, dtype=self.activation_dtype))
+        encoder_decoder_attention_mask = (
+            mtf_layers.attention_mask_same_segment(
+                targets_segmentation, inputs_segmentation,
+                dtype=self.activation_dtype))
+      else:
+        inputs_position = mtf.range(mesh, self.length_dim, dtype=tf.int32)
+        encoder_self_attention_mask = (
+            mtf_layers.attention_mask_ignore_padding(
+                inputs, dtype=self.activation_dtype))
+        encoder_decoder_attention_mask = encoder_self_attention_mask
+
+      x = (mtf.gather(inputs_embedding_var, inputs, self.inputs_vocab_dim) +
+           mtf.gather(positional_embedding_var, inputs_position,
+                      self.max_length_dim))
+      x = layer_prepostprocess_dropout(x)
+      with tf.variable_scope("encoder"):
+        x = self._layer_stack(x,
+                              hparams.num_encoder_layers,
+                              self_attention_mask=encoder_self_attention_mask,
+                              losses=extra_losses)
+      encoder_output = mtf.rename_dimension(
+          x, self.length_dim.name, self.memory_length_dim.name)
+    else:
+      encoder_output = None
+      encoder_decoder_attention_mask = None
+
+    # DECODER
+    x = (mtf.gather(
+        targets_embedding_var, shifted_targets, self.targets_vocab_dim) +
+         mtf.gather(
+             positional_embedding_var, targets_position, self.max_length_dim))
+    x = layer_prepostprocess_dropout(x)
+
+    # Decoder
+    with tf.variable_scope("decoder"):
+      x = self._layer_stack(
+          x,
+          hparams.num_decoder_layers,
+          encoder_output=encoder_output,
+          self_attention_mask=decoder_self_attention_mask,
+          encdec_attention_mask=encoder_decoder_attention_mask,
+          losses=extra_losses)
+    logits = mtf.matmul(x, softmax_var)
+    off_value = hparams.label_smoothing / self._targets_vocab_size
+    on_value = 1.0 - hparams.label_smoothing + off_value
+    soft_targets = mtf.one_hot(
+        targets, self.targets_vocab_dim, on_value=on_value, off_value=off_value,
+        dtype=self.activation_dtype)
+    loss = mtf_layers.softmax_cross_entropy_with_logits(
+        logits, soft_targets, self.targets_vocab_dim)
+    weights = mtf_layers.weights_nonzero(
+        targets, dtype=self.activation_dtype)
+    loss = mtf.reduce_mean(loss * weights)
+    for l in extra_losses:
+      loss += l
+    return logits, loss
+
+  def mtf_model_fn(self, features, mesh):
+    with tf.variable_scope("transformer"):
+      return self._mtf_model_fn(features, mesh)
+
+  @property
+  def _targets_vocab_size(self):
+    targets_vocab_size = self._problem_hparams.target_modality._vocab_size  # pylint: disable=protected-access
+    targets_vocab_size += (-targets_vocab_size) % self._hparams.vocab_divisor
+    return targets_vocab_size
+
+  @property
+  def _inputs_vocab_size(self):
+    if not self.has_input:
+      return None
+    inputs_vocab_size = self._problem_hparams.input_modality[   # pylint: disable=protected-access
+        "inputs"]._vocab_size
+    inputs_vocab_size += (-inputs_vocab_size) % self._hparams.vocab_divisor
+    return inputs_vocab_size
+
+  def _feedforward_layer(self, x, losses=None):
+    """Feed-forward layer.
+
+    Args:
+      x: a mtf.Tensor with shape [batch_dim, length_dim, model_dim]
+      losses: a list to be appended-to
+    Returns:
+      a mtf.Tensor with shape [batch_dim, length_dim, model_dim]
+    Raises:
+      ValueError: if hparams make no sense
+    """
+    hparams = self._hparams
+    feedforward_layer = hparams.feedforward_layer
+    if feedforward_layer == "dense_relu_dense":
+      return mtf_layers.dense_relu_dense(
+          x, self.feedforward_dim, dropout=hparams.relu_dropout,
+          dropout_broadcast_dims=[self.length_dim])
+    elif feedforward_layer == "moe":
+      overhead = (
+          hparams.moe_overhead_train
+          if hparams.mode == tf.estimator.ModeKeys.TRAIN else
+          hparams.moe_overhead_eval)
+      output, loss = mtf_layers.moe_v0(
+          x,
+          self.feedforward_dim,
+          self.model_dim,
+          self.experts_dim,
+          loss_coef=hparams.moe_loss_coef,
+          overhead=overhead)
+      if losses is not None:
+        losses.append(loss)
+        return output
+    else:
+      raise ValueError(
+          "hparams.feedforward_layer not recognized %s" % feedforward_layer)
+
+  def _layer_stack(self,
+                   x,
+                   num_layers,
+                   encoder_output=None,
+                   self_attention_mask=None,
+                   encdec_attention_mask=None,
+                   losses=None):
+    """Encoder or decoder stack.
+
+    Args:
+      x: a mtf.Tensor with shape [batch_dim, length_dim, model_dim]
+      num_layers: an integer
+      encoder_output: an optional mtf.Tensor with shape
+        [batch_dim, encoder_length_dim, model_dim]
+      self_attention_mask: an optional mtf.Tensor with shape
+        [batch, length_dim, memory_length_dim] containing values 0 or -inf.
+      encdec_attention_mask: an optional mtf.Tensor with shape
+        [batch, length_dim, encoder_length_dim] containing values 0 or -inf.
+      losses: a list to be appended-to
+    Returns:
+      a mtf.Tensor with shape [batch_dim, length_dim, model_dim]
+    Raises:
+      ValueError: if hparams make no sense
+    """
+    hparams = self._hparams
+
+    def layer_prepostprocess_dropout(x):
+      return mtf.dropout(
+          x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
+          noise_shape=mtf.TensorShape([self.batch_dim, self.model_dim]))
+    num_layer_norms = num_layers * (2 if encoder_output is None else 3) + 1
+    layer_norms_dim = mtf.Dimension("layer_norms", num_layer_norms)
+    layer_norm_combined_var = mtf.get_variable(
+        x.mesh,
+        "layer_norm_scale",
+        mtf.TensorShape([layer_norms_dim, self.model_dim]),
+        initializer=tf.ones_initializer(),
+        activation_dtype=x.dtype)
+    layer_norm_vars = mtf.unstack(layer_norm_combined_var, layer_norms_dim)
+    def normalize(x):
+      scale = layer_norm_vars.pop(0)
+      variance = mtf.reduce_mean(mtf.square(x), reduced_dim=self.model_dim)
+      return x * mtf.rsqrt(variance + hparams.norm_epsilon) * scale
+
+    for layer in range(num_layers):
+      with tf.variable_scope("layer_%d" % layer):
+        # Self attention layer
+        x += layer_prepostprocess_dropout(
+            mtf_layers.multihead_attention(
+                normalize(x), None,
+                self_attention_mask, self.kv_dim, self.heads_dim,
+                dropout=hparams.attention_dropout,
+                dropout_broadcast_dims=[self.length_dim],
+                name="self_attention"))
+        if encoder_output is not None:
+          # Encoder-Decoder attention layer
+          x += layer_prepostprocess_dropout(
+              mtf_layers.multihead_attention(
+                  normalize(x), encoder_output,
+                  encdec_attention_mask, self.kv_dim, self.heads_dim,
+                  dropout=hparams.attention_dropout,
+                  dropout_broadcast_dims=[self.length_dim],
+                  name="encdec_attention"))
+        # ffn layer
+        x += layer_prepostprocess_dropout(
+            self._feedforward_layer(normalize(x), losses=losses))
+    x = normalize(x)
+    assert not layer_norm_vars
+    return x
+
+  def sample(self, features, mesh):
+    with tf.variable_scope("transformer"):
+      return self._sample(features, mesh)
+
+  def _sample(self, features, mesh):
+    hparams = self._hparams
+    (inputs_embedding_var,
+     targets_embedding_var,
+     softmax_var,
+     positional_embedding_var) = self._embedding_and_softmax_vars(mesh)
+    if self.has_input:
+      inputs = features["inputs"]
+      while len(inputs.shape.as_list()) > 2:
+        inputs = tf.squeeze(inputs, axis=2)
+      actual_batch_size = tf.shape(inputs)[0]
+      actual_length = tf.shape(inputs)[1]
+      inputs = tf.pad(
+          inputs, [[0, hparams.batch_size - actual_batch_size],
+                   [0, hparams.max_length - actual_length]])
+      inputs = self._infeed_to_batch_by_length(
+          inputs, "inputs", mesh, hparams)
+      x = (mtf.gather(inputs_embedding_var, inputs, self.inputs_vocab_dim) +
+           mtf.reshape(positional_embedding_var,
+                       mtf.TensorShape([self.length_dim, self.model_dim])))
+      encoder_attention_mask = (
+          mtf_layers.attention_mask_ignore_padding(
+              inputs, dtype=self.activation_dtype))
+      with tf.variable_scope("encoder"):
+        x = self._layer_stack(x,
+                              hparams.num_encoder_layers,
+                              self_attention_mask=encoder_attention_mask)
+      encoder_output = mtf.rename_dimension(
+          x, self.length_dim.name, self.memory_length_dim.name)
+      encdec_tensors = []
+      for layer_num in xrange(hparams.num_decoder_layers):
+        with tf.variable_scope("decoder/layer_%d/encdec_attention" % layer_num):
+          q_var, k_var, v_var, o_var = mtf_layers.multihead_attention_vars(
+              mesh, self.heads_dim, self.model_dim,
+              self.kv_dim, self.activation_dtype)
+          k = mtf.einsum(
+              [encoder_output, k_var],
+              mtf.TensorShape(
+                  [self.batch_dim, self.heads_dim,
+                   self.memory_length_dim, self.kv_dim]))
+          v = mtf.einsum(
+              [encoder_output, v_var],
+              mtf.TensorShape(
+                  [self.batch_dim, self.heads_dim,
+                   self.memory_length_dim, self.kv_dim]))
+        encdec_tensors.append((q_var, o_var, k, v))
+      partial_targets = None
+    else:
+      encdec_tensors = None
+      encoder_output = None
+      encoder_attention_mask = None
+      # Prepare partial targets.
+      # In either features["inputs"] or features["targets"].
+      # We force the outputs to begin with these sequences.
+      partial_targets = features.get("inputs", None)
+      if partial_targets is None:
+        partial_targets = features.get("targets", None)
+      if partial_targets is not None:
+        partial_targets = common_layers.expand_squeeze_to_nd(partial_targets, 2)
+        partial_targets = tf.to_int32(partial_targets)
+        partial_targets_batch = tf.shape(partial_targets)[0]
+        partial_targets_length = tf.shape(partial_targets)[1]
+        partial_targets = tf.pad(
+            partial_targets, [[0, hparams.batch_size - partial_targets_batch],
+                              [0, hparams.max_length - partial_targets_length]])
+        partial_targets = self._infeed_to_batch_by_length(
+            partial_targets, "partial_targets", mesh, hparams)
+
+    if hparams.beam_size == 1:
+      ids_shape = mtf.TensorShape([self.batch_dim, self.length_dim])
+      kv_shape = mtf.TensorShape([self.batch_dim, self.heads_dim,
+                                  self.memory_length_dim, self.kv_dim])
+    else:
+      beam_dim = mtf.Dimension("beam", hparams.beam_size)
+      ids_shape = mtf.TensorShape([self.batch_dim, beam_dim, self.length_dim])
+      kv_shape = mtf.TensorShape([self.batch_dim, beam_dim, self.heads_dim,
+                                  self.memory_length_dim, self.kv_dim])
+
+    initial_ids = mtf.constant(mesh, 0, ids_shape, dtype=tf.int32)
+    initial_kv_states = (
+        [mtf.zeros(mesh, kv_shape, dtype=self.activation_dtype)]
+        * (2 * hparams.num_decoder_layers))
+    def logits_fn(step_num, ids, states):
+      """Produce logits for this step, and new states."""
+      self_attention_k = states[:hparams.num_decoder_layers]
+      self_attention_v = states[hparams.num_decoder_layers:]
+      ids_this_step = mtf.gather(ids, step_num - 1, self.length_dim)
+      x = (mtf.gather(targets_embedding_var, ids_this_step,
+                      self.targets_vocab_dim) +
+           mtf.gather(positional_embedding_var, step_num, self.max_length_dim))
+      with tf.variable_scope("decoder"):
+        x, new_self_attention_k, new_self_attention_v = (
+            self._decoder_layer_stack_incremental(
+                x,
+                step_num,
+                encdec_tensors,
+                self_attention_k,
+                self_attention_v,
+                encdec_attention_mask=encoder_attention_mask))
+      logits = mtf.matmul(x, softmax_var)
+      return logits, new_self_attention_k + new_self_attention_v
+
+    if hparams.beam_size == 1:
+      temperature = (0.0 if hparams.sampling_method == "argmax"
+                     else hparams.sampling_temp)
+      return mtf_beam_search.greedy_decode(
+          logits_fn,
+          initial_ids,
+          temperature=temperature,
+          initial_states=initial_kv_states,
+          forced_ids=partial_targets,
+          use_tpu=hparams.use_tpu)
+    else:
+      if self.has_input:
+        input_length = mtf.reduce_sum(
+            mtf.to_float(mtf.cast(inputs, tf.bool)),
+            reduced_dim=self.length_dim)
+        max_input_length = mtf.reduce_max(input_length)
+        decode_length = mtf.cast(
+            max_input_length * hparams.decode_length_multiplier
+            + hparams.decode_length_constant, tf.int32)
+      else:
+        decode_length = None
+      beams, unused_scores = mtf_beam_search.beam_search(
+          logits_fn,
+          initial_ids,
+          hparams.alpha,
+          states=initial_kv_states,
+          decode_length=decode_length,
+          use_tpu=hparams.use_tpu)
+      return mtf.gather(beams, mtf.constant(mesh, 0, dtype=tf.int32), beam_dim)
+
+  def _decoder_layer_stack_incremental(self,
+                                       x,
+                                       step_num,
+                                       encdec_tensors,
+                                       self_attention_k,
+                                       self_attention_v,
+                                       encdec_attention_mask=None):
+    """Decoder layer stack during inference.
+
+    We are processing only one position at a time.
+
+    The self-attention keys and values have already been computed for
+    previous positions.  In addition to the decoder output, we need to
+    produce the updated self-attention keys and values.
+
+    If there is an encoder, then additional Tensors are supplied in
+    encdec_tensors, which give us the keys and values for encoder-decoder
+    attention as well as the weight matrices q_var and o_var.
+
+    Args:
+      x: a mtf.Tensor with shape [batch_dim, model_dim]
+      step_num: an mtf integer Scalar
+      encdec_tensors: an optional list of num_layers tuples, each of the form
+        (q_var, o_var, k, v)
+      self_attention_k: an optional list of num_layers Tensors each with shape
+        [batch, heads, memory_length, kv_channels]
+      self_attention_v: an optional list of num_layers Tensors each with shape
+        [batch, heads, memory_length, kv_channels]
+      encdec_attention_mask: an optional mtf.Tensor with shape
+        [batch, length_dim, encoder_length_dim] containing values 0 or -inf.
+
+    Returns:
+      y: a mtf.Tensor with shape [batch_dim, model_dim]
+      new_self_attention_k: a list of num_layers mtf.Tensors, with the same
+        shapes as the elements of self_attention_k
+      new_self_attention_v: a list of num_layers mtf.Tensors, with the same
+        shapes as the elements of self_attention_v
+
+    Raises:
+      ValueError: if hparams make no sense
+    """
+    hparams = self._hparams
+    num_layers = hparams.num_decoder_layers
+    num_layer_norms = num_layers * (2 if encdec_tensors is None else 3) + 1
+    layer_norms_dim = mtf.Dimension("layer_norms", num_layer_norms)
+    layer_norm_combined_var = mtf.get_variable(
+        x.mesh,
+        "layer_norm_scale",
+        mtf.TensorShape([layer_norms_dim, self.model_dim]),
+        initializer=tf.ones_initializer(),
+        activation_dtype=x.dtype)
+    layer_norm_vars = mtf.unstack(layer_norm_combined_var, layer_norms_dim)
+    def normalize(x):
+      scale = layer_norm_vars.pop(0)
+      variance = mtf.reduce_mean(mtf.square(x), reduced_dim=self.model_dim)
+      return x * mtf.rsqrt(variance + hparams.norm_epsilon) * scale
+
+    new_self_attention_k = []
+    new_self_attention_v = []
+    for layer in range(num_layers):
+      with tf.variable_scope("layer_%d" % layer):
+        # Self attention layer
+        y, new_k, new_v = mtf_layers.multihead_self_attention_incremental(
+            normalize(x),
+            prev_k=self_attention_k[layer],
+            prev_v=self_attention_v[layer],
+            step_num=step_num,
+            name="self_attention")
+        new_self_attention_k.append(new_k)
+        new_self_attention_v.append(new_v)
+        x += y
+        if encdec_tensors is not None:
+          # Encoder-Decoder attention layer
+          q_var, o_var, k, v = encdec_tensors[layer]
+          x += mtf_layers.multihead_encdec_attention_incremental(
+              normalize(x),
+              q_var, o_var, k, v,
+              encdec_attention_mask,
+              name="encdec_attention")
+        # ffn layer
+        x += self._feedforward_layer(normalize(x), hparams)
+    x = normalize(x)
+    assert not layer_norm_vars
+    return x, new_self_attention_k, new_self_attention_v
+
+
+@registry.register_hparams
+def mtf_transformer_base():
+  """Set of hyperparameters."""
+  hparams = common_hparams.basic_params1()
+  hparams.no_data_parallelism = True
+  hparams.use_fixed_batch_size = True
+  hparams.add_hparam("mtf_mode", True)
+  hparams.batch_size = 64
+  hparams.max_length = 256
+  hparams.add_hparam("d_model", 512)
+  hparams.add_hparam("d_kv", 128)
+  hparams.label_smoothing = 0.1
+  # 8-way model-parallelism
+  hparams.add_hparam("mesh_shape", "8")
+  hparams.add_hparam("layout", "vocab:0;d_ff:0;heads:0")
+  hparams.add_hparam("num_heads", 8)
+  hparams.add_hparam("d_ff", 2048)
+  hparams.add_hparam("num_encoder_layers", 6)
+  hparams.add_hparam("num_decoder_layers", 6)
+  hparams.add_hparam("attention_dropout", 0.1)
+  hparams.add_hparam("relu_dropout", 0.1)
+  hparams.layer_prepostprocess_dropout = 0.1
+
+  # round up vocab sizes to be a multiple of this value
+  hparams.vocab_divisor = 128
+
+  # mixture of experts hparams
+  hparams.add_hparam("feedforward_layer", "dense_relu_dense")
+  hparams.add_hparam("moe_overhead_train", 1.0)
+  hparams.add_hparam("moe_overhead_eval", 2.0)
+  hparams.moe_num_experts = 16
+  hparams.moe_loss_coef = 1e-3
+
+  # Use targets_embedding_var * rsqrt(d_model) as softmax_var
+  hparams.shared_embedding_and_softmax_weights = True
+  # Reuse targets_embedding_var as inputs_embedding_var
+  hparams.shared_embedding = True
+  hparams.optimizer = "Adafactor"
+  hparams.learning_rate_schedule = "rsqrt_decay*linear_decay"
+  hparams.learning_rate_warmup_steps = 10000
+  hparams.activation_dtype = "float32"
+
+  # These parameters make Transformer model compatible with MtfTransformer
+  # Do not override these, as mtf_transformer does not support other options.
+  hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
+  hparams.target_modality = "symbol:identity"
+  hparams.input_modalities = "inputs:symbol:identity"
+
+  # Parameters for computing the maximum decode length in beam search.
+  # Maximum decode length is:
+  #    min(max_length,
+  #        decode_length_multiplier * input_length + decode_length_constant)
+  hparams.add_hparam("decode_length_multiplier", 1.5)
+  hparams.add_hparam("decode_length_constant", 10.0)
+
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_shared_embedding():
+  hparams = mtf_transformer_base()
+  hparams.shared_embedding_and_softmax_weights = False
+  hparams.shared_embedding = True
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_no_share():
+  hparams = mtf_transformer_base()
+  hparams.shared_embedding_and_softmax_weights = False
+  hparams.shared_embedding = False
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_tiny():
+  """Catch bugs locally..."""
+  hparams = mtf_transformer_base()
+  hparams.d_model = 128
+  hparams.d_ff = 512
+  hparams.batch_size = 4
+  hparams.num_encoder_layers = 2
+  hparams.num_decoder_layers = 2
+  hparams.num_heads = 4
+  # data parallelism and model-parallelism
+  hparams.mesh_shape = "2.2"
+  hparams.layout = "batch:0;vocab:1;d_ff:1;heads:1"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_single():
+  hparams = mtf_transformer_tiny()
+  hparams.mesh_shape = ""
+  hparams.layout = ""
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_tiny_moe():
+  hparams = mtf_transformer_tiny()
+  hparams.mesh_shape = "4"
+  hparams.layout = "batch:0,experts:0"
+  hparams.feedforward_layer = "moe"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_tiny_8gpu():
+  hparams = mtf_transformer_tiny()
+  hparams.mesh_shape = "8"
+  hparams.layout = "vocab:0;d_ff:0;heads:0"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_length_sharded():
+  hparams = mtf_transformer_tiny()
+  hparams.mesh_shape = "2"
+  hparams.layout = "length:0"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_lm():
+  """Set of hyperparameters."""
+  hparams = mtf_transformer_base()
+  hparams.label_smoothing = 0.0
+  return hparams
+
+
+def mtf_transformer_paper_lm(sz):
+  """Config for language-model experiments."""
+  n = 2 ** sz
+  hparams = mtf_transformer_base()
+  hparams.label_smoothing = 0.0
+  hparams.batch_size = 128
+  hparams.d_model = 1024
+  hparams.d_ff = int(8192 * n)
+  hparams.num_heads = int(8 * n)
+  hparams.shared_embedding_and_softmax_weights = False
+  hparams.learning_rate_decay_steps = 27300  # one epoch for lm1b
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_lm_m1():
+  hparams = mtf_transformer_paper_lm(-1)
+  hparams.mesh_shape = "32"
+  hparams.layout = "batch:0"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_lm_0():
+  hparams = mtf_transformer_paper_lm(0)
+  hparams.mesh_shape = "32"
+  hparams.layout = "batch:0"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_lm_1():
+  hparams = mtf_transformer_paper_lm(1)
+  hparams.mesh_shape = "4;8"
+  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_lm_2():
+  hparams = mtf_transformer_paper_lm(2)
+  hparams.mesh_shape = "4;8"
+  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_lm_3():
+  hparams = mtf_transformer_paper_lm(3)
+  hparams.mesh_shape = "8;16"
+  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_lm_4():
+  hparams = mtf_transformer_paper_lm(4)
+  hparams.mesh_shape = "8;16"
+  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_lm_5():
+  hparams = mtf_transformer_paper_lm(5)
+  hparams.mesh_shape = "8;16"
+  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  return hparams
+
+
+def mtf_transformer_paper_tr(sz):
+  """Config for translation experiments.
+
+  translate_enfr_wmt32k_packed - tokens=2.385e9
+  batch:128 * length:256 = 32768 tokens/step
+  steps/epoch = 72800
+  Let's run for 3 epochs (218K steps)
+  Leaerning-rate-decay for last epoch
+
+  Args:
+    sz: an integer
+  Returns:
+    hyperparameters
+  """
+  n = 2 ** sz
+  hparams = mtf_transformer_base()
+  hparams.label_smoothing = 0.1
+  hparams.batch_size = 128
+  hparams.d_model = 1024
+  hparams.d_ff = int(4096 * n)
+  hparams.num_heads = int(8 * n)
+  hparams.shared_embedding_and_softmax_weights = False
+  hparams.learning_rate_decay_steps = 51400  # one epoch for enfr
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_tr_m1():
+  hparams = mtf_transformer_paper_tr(-1)
+  hparams.mesh_shape = "32"
+  hparams.layout = "batch:0"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_tr_0():
+  hparams = mtf_transformer_paper_tr(0)
+  hparams.mesh_shape = "32"
+  hparams.layout = "batch:0"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_tr_1():
+  hparams = mtf_transformer_paper_tr(1)
+  hparams.mesh_shape = "4;8"
+  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_tr_2():
+  hparams = mtf_transformer_paper_tr(2)
+  hparams.mesh_shape = "4;8"
+  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_tr_3():
+  hparams = mtf_transformer_paper_tr(3)
+  hparams.mesh_shape = "8;16"
+  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_tr_4():
+  return mtf_transformer_paper_tr(4)
+
+
+@registry.register_hparams
+def mtf_transformer_paper_tr_5():
+  return mtf_transformer_paper_tr(5)
+
+
+@registry.register_hparams
+def mtf_transformer_lm_moe():
+  """Mixture of experts language model."""
+  hparams = mtf_transformer_base()
+  hparams.label_smoothing = 0.0
+  hparams.batch_size = 128
+  hparams.d_model = 1024
+  hparams.d_ff = 4096
+  hparams.attention_key_channels = 1024
+  hparams.attention_value_channels = 1024
+  hparams.shared_embedding_and_softmax_weights = False
+  hparams.num_heads = 8
+  hparams.layout = "batch:0,experts:0"
+  hparams.feedforward_layer = "moe"
+  return hparams
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer_compat.py b/tensor2tensor/mesh_tensorflow/mtf_transformer_compat.py
new file mode 100644
index 000000000..2e48f9de7
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer_compat.py
@@ -0,0 +1,926 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Temporary hack for decoding mtf_transformer models.
+
+This is a transformer implementation in regular TensorFlow which is
+checkpoint-compatible with MtfTransformer for eval/inference.
+
+The purpose of this model is to run inference on MtfTransformer models.
+We are working on native decoding in MtfTransformer which will be faster and
+cleaner.
+
+TODO(noam): Remove once we can decode in mtf.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from six.moves import range  # pylint: disable=redefined-builtin
+
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import beam_search
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
+
+import tensorflow as tf
+
+from tensorflow.python.util import nest
+
+
+@registry.register_model
+class MtfTransformerCompat(t2t_model.T2TModel):
+  """Attention net.  See file docstring."""
+
+  def __init__(self, *args, **kwargs):
+    with tf.variable_scope("transformer"):
+      self._top_scope = tf.get_variable_scope()
+    kwargs["_scope"] = "transformer"
+    super(MtfTransformerCompat, self).__init__(*args, **kwargs)
+    self._name = "transformer"
+    self._base_name = "transformer"
+
+  @property
+  def _targets_vocab_size(self):
+    targets_vocab_size = self._problem_hparams.vocabulary["targets"].vocab_size
+    targets_vocab_size += (-targets_vocab_size) % self._hparams.vocab_divisor
+    return targets_vocab_size
+
+  @property
+  def _inputs_vocab_size(self):
+    if not self.has_input:
+      return None
+    inputs_vocab_size = self._problem_hparams.vocabulary["inputs"].vocab_size
+    inputs_vocab_size += (-inputs_vocab_size) % self._hparams.vocab_divisor
+    return inputs_vocab_size
+
+  @property
+  def _embedding_and_softmax_var_names(self):
+    """Figure out the variable names for the embedding and softmax variables.
+
+    Equality between the returned names means that we should share the
+    variables.
+
+    Returns:
+      inputs_embedding_name: a string or None
+      targets_embedding_name: a string
+      softmax_var_name: a string
+    Raises:
+      ValueError: if we try to share embeddings with different vocab sizes.
+    """
+    hparams = self._hparams
+    inputs_embedding_name = "input_emb"
+    targets_embedding_name = "target_emb"
+    softmax_var_name = "softmax"
+    if (self.has_input and
+        (hparams.shared_embedding or
+         hparams.shared_embedding_and_softmax_weights) and
+        self._inputs_vocab_size != self._targets_vocab_size):
+      raise ValueError(
+          "hparams.shared_embedding_and_softmax_weights "
+          " or hparams.shared_embedding require "
+          "that input and target vocabulary sizes be equal %s vs %s"
+          % (self._inputs_vocab_size, self._targets_vocab_size))
+    if hparams.shared_embedding_and_softmax_weights:
+      inputs_embedding_name = "shared"
+      targets_embedding_name = "shared"
+      softmax_var_name = "shared"
+    elif hparams.shared_embedding:
+      inputs_embedding_name = "shared"
+      targets_embedding_name = "shared"
+    targets_embedding_name = (
+        "symbol_modality_%d_%d/%s/weights_0" %
+        (self._targets_vocab_size, hparams.d_model, targets_embedding_name))
+    softmax_var_name = (
+        "symbol_modality_%d_%d/%s/weights_0" %
+        (self._targets_vocab_size, hparams.d_model, softmax_var_name))
+    if self.has_input:
+      inputs_embedding_name = (
+          "symbol_modality_%d_%d/%s/weights_0" %
+          (self._inputs_vocab_size, hparams.d_model, inputs_embedding_name))
+    else:
+      inputs_embedding_name = None
+    return inputs_embedding_name, targets_embedding_name, softmax_var_name
+
+  @property
+  def _get_targets_emb_var(self):
+    with tf.variable_scope(self._top_scope, reuse=tf.AUTO_REUSE):
+      return tf.get_variable(
+          "targets_embedding",
+          [self._targets_vocab_size, self._hparams.d_model])
+
+  @property
+  def _get_inputs_emb_var(self):
+    if self._hparams.shared_embedding:
+      return self._get_targets_emb_var
+    with tf.variable_scope(self._top_scope, reuse=tf.AUTO_REUSE):
+      return tf.get_variable(
+          "inputs_embedding",
+          [self._inputs_vocab_size, self._hparams.d_model])
+
+  @property
+  def _get_softmax_var(self):
+    if self._hparams.shared_embedding_and_softmax_weights:
+      return self._get_targets_emb_var * (self._hparams.d_model ** -0.5)
+    with tf.variable_scope(self._top_scope, reuse=tf.AUTO_REUSE):
+      return tf.get_variable(
+          "softmax",
+          [self._targets_vocab_size, self._hparams.d_model])
+
+  def encode(self, inputs, hparams, features=None, losses=None):
+    """Encode transformer inputs.
+
+    Args:
+      inputs: Transformer inputs [batch_size, input_length, input_height,
+        hidden_dim] which will be flattened along the two spatial dimensions.
+      hparams: hyperparameters for model.
+      features: optionally pass the entire features dictionary as well.
+        This is needed now for "packed" datasets.
+      losses: optional list onto which to append extra training losses
+
+    Returns:
+      Tuple of:
+          encoder_output: Encoder representation.
+              [batch_size, input_length, hidden_dim]
+          encoder_decoder_attention_bias: Bias and mask weights for
+              encoder-decoder attention. [batch_size, input_length]
+    """
+    tf.logging.info("Encode inputs = %s" % inputs)
+    encoder_input, self_attention_bias, encoder_decoder_attention_bias = (
+        transformer_prepare_encoder(
+            self._get_inputs_emb_var, inputs, hparams, features=features))
+
+    encoder_input = tf.nn.dropout(encoder_input,
+                                  1.0 - hparams.layer_prepostprocess_dropout)
+
+    encoder_output = transformer_encoder(
+        encoder_input,
+        self_attention_bias,
+        hparams,
+        losses=losses)
+
+    return encoder_output, encoder_decoder_attention_bias
+
+  def decode(self,
+             decoder_input,
+             encoder_output,
+             encoder_decoder_attention_bias,
+             decoder_self_attention_bias,
+             hparams,
+             cache=None,
+             losses=None):
+    """Decode Transformer outputs from encoder representation.
+
+    Args:
+      decoder_input: inputs to bottom of the model.
+          [batch_size, decoder_length, hidden_dim]
+      encoder_output: Encoder representation.
+          [batch_size, input_length, hidden_dim]
+      encoder_decoder_attention_bias: Bias and mask weights for
+          encoder-decoder attention. [batch_size, input_length]
+      decoder_self_attention_bias: Bias and mask weights for decoder
+          self-attention. [batch_size, decoder_length]
+      hparams: hyperparameters for model.
+      cache: dict, containing tensors which are the results of previous
+          attentions, used for fast decoding.
+      losses: optional list onto which to append extra training losses
+
+    Returns:
+      Final decoder representation. [batch_size, decoder_length, hidden_dim]
+    """
+    decoder_input = tf.nn.dropout(decoder_input,
+                                  1.0 - hparams.layer_prepostprocess_dropout)
+
+    decoder_output = transformer_decoder(
+        decoder_input,
+        encoder_output,
+        decoder_self_attention_bias,
+        encoder_decoder_attention_bias,
+        hparams,
+        cache=cache,
+        losses=losses)
+
+    ret = tf.tensordot(decoder_output, self._get_softmax_var, axes=[[-1], [1]])
+    ret = tf.expand_dims(tf.expand_dims(ret, 2), 3)
+    return ret
+
+  def body(self, features):
+    """Transformer main model_fn.
+
+    Args:
+      features: Map of features to the model. Should contain the following:
+          "inputs": Transformer inputs [batch_size, input_length, hidden_dim]
+          "targets": Target decoder outputs.
+              [batch_size, decoder_length, hidden_dim]
+
+    Returns:
+      Final decoder representation. [batch_size, decoder_length, hidden_dim]
+    """
+    with tf.variable_scope(self._top_scope):
+      hparams = self._hparams
+      losses = []
+
+      if self.has_input:
+        inputs = tf.squeeze(features["inputs_raw"], (2, 3))
+        encoder_output, encoder_decoder_attention_bias = self.encode(
+            inputs, hparams, features=features, losses=losses)
+      else:
+        encoder_output, encoder_decoder_attention_bias = (None, None)
+
+      targets = tf.squeeze(features["targets_raw"], (2, 3))
+      decoder_input, decoder_self_attention_bias = transformer_prepare_decoder(
+          self._get_targets_emb_var, targets, hparams, features=features)
+
+      decoder_output = self.decode(
+          decoder_input,
+          encoder_output,
+          encoder_decoder_attention_bias,
+          decoder_self_attention_bias,
+          hparams,
+          losses=losses)
+
+      if losses:
+        return decoder_output, {"extra_loss": tf.add_n(losses)}
+      else:
+        return decoder_output
+
+  def _greedy_infer(self, features, decode_length, use_tpu=False):
+    """Fast version of greedy decoding.
+
+    Args:
+      features: an map of string to `Tensor`
+      decode_length: an integer.  How many additional timesteps to decode.
+      use_tpu: a boolean
+
+    Returns:
+      A dict of decoding results {
+          "outputs": integer `Tensor` of decoded ids of shape
+              [batch_size, <= decode_length] if beam_size == 1 or
+              [batch_size, top_beams, <= decode_length]
+          "scores": decoding log probs from the beam search,
+              None if using greedy decoding (beam_size=1)
+      }
+
+    Raises:
+      NotImplementedError: If there are multiple data shards.
+    """
+    with tf.variable_scope(self.name):
+      return  self._fast_decode(features, decode_length)
+
+  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
+    """Beam search decoding.
+
+    Args:
+      features: an map of string to `Tensor`
+      decode_length: an integer.  How many additional timesteps to decode.
+      beam_size: number of beams.
+      top_beams: an integer. How many of the beams to return.
+      alpha: Float that controls the length penalty. larger the alpha, stronger
+        the preference for longer translations.
+
+    Returns:
+      A dict of decoding results {
+          "outputs": integer `Tensor` of decoded ids of shape
+              [batch_size, <= decode_length] if beam_size == 1 or
+              [batch_size, top_beams, <= decode_length]
+          "scores": decoding log probs from the beam search,
+              None if using greedy decoding (beam_size=1)
+      }
+    """
+    with tf.variable_scope(self.name):
+      return self._fast_decode(features, decode_length, beam_size, top_beams,
+                               alpha)
+
+  def _fast_decode(self,
+                   features,
+                   decode_length,
+                   beam_size=1,
+                   top_beams=1,
+                   alpha=1.0):
+    """Fast decoding.
+
+    Implements both greedy and beam search decoding, uses beam search iff
+    beam_size > 1, otherwise beam search related arguments are ignored.
+
+    Args:
+      features: a map of string to model  features.
+      decode_length: an integer.  How many additional timesteps to decode.
+      beam_size: number of beams.
+      top_beams: an integer. How many of the beams to return.
+      alpha: Float that controls the length penalty. larger the alpha, stronger
+        the preference for longer translations.
+
+    Returns:
+      A dict of decoding results {
+          "outputs": integer `Tensor` of decoded ids of shape
+              [batch_size, <= decode_length] if beam_size == 1 or
+              [batch_size, top_beams, <= decode_length]
+          "scores": decoding log probs from the beam search,
+              None if using greedy decoding (beam_size=1)
+      }
+
+    Raises:
+      NotImplementedError: If there are multiple data shards.
+    """
+    if self._num_datashards != 1:
+      raise NotImplementedError("Fast decoding only supports a single shard.")
+    dp = self._data_parallelism
+    hparams = self._hparams
+    target_modality = self._problem_hparams.target_modality
+    if "targets_segmentation" in features:
+      raise NotImplementedError(
+          "Decoding not supported on packed datasets "
+          " If you want to decode from a dataset, use the non-packed version"
+          " of the dataset when decoding.")
+    if self.has_input:
+      inputs = features["inputs"]
+      if target_modality.is_class_modality:
+        decode_length = 1
+      else:
+        decode_length = (
+            common_layers.shape_list(inputs)[1] + features.get(
+                "decode_length", decode_length))
+
+      # TODO(llion): Clean up this reshaping logic.
+      inputs = tf.expand_dims(inputs, axis=1)
+      if len(inputs.shape) < 5:
+        inputs = tf.expand_dims(inputs, axis=4)
+      s = common_layers.shape_list(inputs)
+      batch_size = s[0]
+      inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
+      inputs = tf.squeeze(inputs, (2, 3))
+      # _shard_features called to ensure that the variable names match
+      inputs = self._shard_features({"inputs": inputs})["inputs"]
+
+      # input_modality = self._problem_hparams.input_modality["inputs"]
+      # with tf.variable_scope(input_modality.name):
+      #   inputs = input_modality.bottom_sharded(inputs, dp)
+      encoder_output, encoder_decoder_attention_bias = dp(
+          self.encode,
+          inputs,
+          hparams,
+          features=features)
+      encoder_output = encoder_output[0]
+      encoder_decoder_attention_bias = encoder_decoder_attention_bias[0]
+      partial_targets = None
+    else:
+      # The problem has no inputs.
+      encoder_output = None
+      encoder_decoder_attention_bias = None
+
+      # Prepare partial targets.
+      # In either features["inputs"] or features["targets"].
+      # We force the outputs to begin with these sequences.
+      partial_targets = features.get("inputs")
+      if partial_targets is None:
+        partial_targets = features["targets"]
+      assert partial_targets is not None
+      partial_targets = common_layers.expand_squeeze_to_nd(partial_targets, 2)
+      partial_targets = tf.to_int64(partial_targets)
+      partial_targets_shape = common_layers.shape_list(partial_targets)
+      partial_targets_length = partial_targets_shape[1]
+      decode_length = (
+          partial_targets_length + features.get("decode_length", decode_length))
+      batch_size = partial_targets_shape[0]
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+      positional_encoding = common_attention.add_positional_embedding(
+          tf.zeros([1, decode_length, hparams.d_model]),
+          hparams.max_length, "positional_embedding", None)
+
+    def preprocess_targets(targets, i):
+      """Performs preprocessing steps on the targets to prepare for the decoder.
+
+      This includes:
+        - Embedding the ids.
+        - Flattening to 3D tensor.
+        - Optionally adding timing signals.
+
+      Args:
+        targets: inputs ids to the decoder. [batch_size, 1]
+        i: scalar, Step number of the decoding loop.
+
+      Returns:
+        Processed targets [batch_size, 1, hidden_dim]
+      """
+      targets_emb_var = self._get_targets_emb_var
+      targets = tf.gather(targets_emb_var, targets)
+      tf.logging.info("targets = %s" % targets)
+      targets = tf.squeeze(targets, (2, 3))
+      if positional_encoding is not None:
+        targets += positional_encoding[:, i:i + 1]
+      return targets
+
+    def symbols_to_logits_fn(ids, i, cache):
+      """Go from ids to logits for next symbol."""
+      ids = ids[:, -1:]
+      targets = tf.expand_dims(tf.expand_dims(ids, axis=2), axis=3)
+      targets = preprocess_targets(targets, i)
+
+      bias = None  # decoder_self_attention_bias[:, :, i:i + 1, :i + 1]
+
+      body_outputs = dp(
+          self.decode,
+          targets,
+          cache.get("encoder_output"),
+          cache.get("encoder_decoder_attention_bias"),
+          bias,
+          hparams,
+          cache)
+
+      logits = body_outputs[0]
+      # with tf.variable_scope(target_modality.name):
+      #   logits = target_modality.top_sharded(body_outputs, None, dp)[0]
+
+      ret = tf.squeeze(logits, axis=[1, 2, 3])
+      if partial_targets is not None:
+        # If the position is within the given partial targets, we alter the
+        # logits to always return those values.
+        # A faster approach would be to process the partial targets in one
+        # iteration in order to fill the corresponding parts of the cache.
+        # This would require broader changes, though.
+        vocab_size = tf.shape(ret)[1]
+
+        def forced_logits():
+          return tf.one_hot(
+              tf.tile(partial_targets[:, i], [beam_size]), vocab_size, 0.0,
+              -1e9)
+
+        ret = tf.cond(
+            tf.less(i, partial_targets_length), forced_logits, lambda: ret)
+      return ret, cache
+
+    ret = fast_decode(
+        encoder_output=encoder_output,
+        encoder_decoder_attention_bias=encoder_decoder_attention_bias,
+        symbols_to_logits_fn=symbols_to_logits_fn,
+        hparams=hparams,
+        decode_length=decode_length,
+        vocab_size=target_modality.top_dimensionality,
+        beam_size=beam_size,
+        top_beams=top_beams,
+        alpha=alpha,
+        batch_size=batch_size,
+        force_decode_length=self._decode_hparams.force_decode_length)
+    if partial_targets is not None:
+      if beam_size <= 1 or top_beams <= 1:
+        ret["outputs"] = ret["outputs"][:, partial_targets_length:]
+      else:
+        ret["outputs"] = ret["outputs"][:, :, partial_targets_length:]
+    return ret
+
+
+def fast_decode(encoder_output,
+                encoder_decoder_attention_bias,
+                symbols_to_logits_fn,
+                hparams,
+                decode_length,
+                vocab_size,
+                beam_size=1,
+                top_beams=1,
+                alpha=1.0,
+                eos_id=beam_search.EOS_ID,
+                batch_size=None,
+                force_decode_length=False):
+  """Given encoder output and a symbols to logits function, does fast decoding.
+
+  Implements both greedy and beam search decoding, uses beam search iff
+  beam_size > 1, otherwise beam search related arguments are ignored.
+
+  Args:
+    encoder_output: Output from encoder.
+    encoder_decoder_attention_bias: a bias tensor for use in encoder-decoder
+      attention
+    symbols_to_logits_fn: Incremental decoding; function mapping triple
+      `(ids, step, cache)` to symbol logits.
+    hparams: run hyperparameters
+    decode_length: an integer.  How many additional timesteps to decode.
+    vocab_size: Output vocabulary size.
+    beam_size: number of beams.
+    top_beams: an integer. How many of the beams to return.
+    alpha: Float that controls the length penalty. larger the alpha, stronger
+      the preference for longer translations.
+    eos_id: End-of-sequence symbol in beam search.
+    batch_size: an integer scalar - must be passed if there is no input
+    force_decode_length: bool, whether to force the full decode length, or if
+      False, stop when all beams hit eos_id.
+
+  Returns:
+      A dict of decoding results {
+          "outputs": integer `Tensor` of decoded ids of shape
+              [batch_size, <= decode_length] if top_beams == 1 or
+              [batch_size, top_beams, <= decode_length] otherwise
+          "scores": decoding log probs from the beam search,
+              None if using greedy decoding (beam_size=1)
+      }
+
+    Raises:
+      NotImplementedError: If beam size > 1 with partial targets.
+  """
+  if encoder_output is not None:
+    batch_size = common_layers.shape_list(encoder_output)[0]
+
+  num_layers = hparams.num_decoder_layers
+  cache = {
+      "layer_%d" % layer: {
+          "k": tf.zeros([batch_size, hparams.num_heads,
+                         0, hparams.d_kv]),
+          "v": tf.zeros([batch_size, hparams.num_heads,
+                         0, hparams.d_kv]),
+      } for layer in range(num_layers)
+  }
+
+  if encoder_output is not None:
+    for layer in range(num_layers):
+      layer_name = "layer_%d" % layer
+      with tf.variable_scope("decoder/%s" % layer_name):
+        k_encdec, v_encdec = multihead_attention_compat(
+            None,
+            encoder_output,
+            None,
+            hparams.d_kv,
+            hparams.num_heads,
+            name="encdec_attention")
+      cache[layer_name]["k_encdec"] = k_encdec
+      cache[layer_name]["v_encdec"] = v_encdec
+
+    cache["encoder_output"] = encoder_output
+    cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias
+
+  if beam_size > 1:  # Beam Search
+    initial_ids = tf.zeros([batch_size], dtype=tf.int32)
+    decoded_ids, scores = beam_search.beam_search(
+        symbols_to_logits_fn,
+        initial_ids,
+        beam_size,
+        decode_length,
+        vocab_size,
+        alpha,
+        states=cache,
+        eos_id=eos_id,
+        stop_early=(top_beams == 1))
+
+    if top_beams == 1:
+      decoded_ids = decoded_ids[:, 0, 1:]
+      scores = scores[:, 0]
+    else:
+      decoded_ids = decoded_ids[:, :top_beams, 1:]
+      scores = scores[:, :top_beams]
+  else:  # Greedy
+
+    def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
+      """One step of greedy decoding."""
+      logits, cache = symbols_to_logits_fn(next_id, i, cache)
+      log_probs = common_layers.log_prob_from_logits(logits)
+      temperature = (0.0 if hparams.sampling_method == "argmax" else
+                     hparams.sampling_temp)
+      next_id = common_layers.sample_with_temperature(logits, temperature)
+      hit_eos |= tf.equal(next_id, eos_id)
+
+      log_prob_indices = tf.stack(
+          [tf.range(tf.to_int64(batch_size)), next_id], axis=1)
+      log_prob += tf.gather_nd(log_probs, log_prob_indices)
+
+      next_id = tf.expand_dims(next_id, axis=1)
+      decoded_ids = tf.concat([decoded_ids, next_id], axis=1)
+      return i + 1, hit_eos, next_id, decoded_ids, cache, log_prob
+
+    def is_not_finished(i, hit_eos, *_):
+      finished = i >= decode_length
+      if not force_decode_length:
+        finished |= tf.reduce_all(hit_eos)
+      return tf.logical_not(finished)
+
+    decoded_ids = tf.zeros([batch_size, 0], dtype=tf.int64)
+    hit_eos = tf.fill([batch_size], False)
+    next_id = tf.zeros([batch_size, 1], dtype=tf.int64)
+    initial_log_prob = tf.zeros([batch_size], dtype=tf.float32)
+    _, _, _, decoded_ids, _, log_prob = tf.while_loop(
+        is_not_finished,
+        inner_loop, [
+            tf.constant(0), hit_eos, next_id, decoded_ids, cache,
+            initial_log_prob
+        ],
+        shape_invariants=[
+            tf.TensorShape([]),
+            tf.TensorShape([None]),
+            tf.TensorShape([None, None]),
+            tf.TensorShape([None, None]),
+            nest.map_structure(beam_search.get_state_shape_invariants, cache),
+            tf.TensorShape([None]),
+        ])
+    scores = log_prob
+
+  return {"outputs": decoded_ids, "scores": scores}
+
+
+def transformer_prepare_encoder(
+    inputs_emb_var, inputs, hparams, features=None):
+  """Prepare one shard of the model for the encoder.
+
+  Args:
+    inputs_emb_var: a Tensor
+    inputs: a Tensor.
+    hparams: run hyperparameters
+    features: optionally pass the entire features dictionary as well.
+      This is needed now for "packed" datasets.
+
+  Returns:
+    encoder_input: a Tensor, bottom of encoder stack
+    encoder_self_attention_bias: a bias tensor for use in encoder self-attention
+    encoder_decoder_attention_bias: a bias tensor for use in encoder-decoder
+      attention
+  """
+  encoder_input = tf.gather(inputs_emb_var, inputs)
+
+  if features and "inputs_segmentation" in features:
+    # Packed dataset.  Keep the examples from seeing each other.
+    inputs_segmentation = features["inputs_segmentation"]
+    inputs_position = features["inputs_position"]
+    targets_segmentation = features["targets_segmentation"]
+    encoder_self_attention_bias = common_attention.attention_bias_same_segment(
+        inputs_segmentation, inputs_segmentation)
+    encoder_decoder_attention_bias = (
+        common_attention.attention_bias_same_segment(targets_segmentation,
+                                                     inputs_segmentation))
+  else:
+    # Usual case - not a packed dataset.
+    encoder_padding = tf.to_float(tf.equal(inputs, 0))
+    ignore_padding = common_attention.attention_bias_ignore_padding(
+        encoder_padding)
+    encoder_self_attention_bias = ignore_padding
+    encoder_decoder_attention_bias = ignore_padding
+    inputs_position = None
+  with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+    encoder_input = common_attention.add_positional_embedding(
+        encoder_input, hparams.max_length, "positional_embedding",
+        inputs_position)
+  if hparams.activation_dtype == "bfloat16":
+    encoder_self_attention_bias = tf.cast(encoder_self_attention_bias,
+                                          tf.bfloat16)
+    encoder_decoder_attention_bias = tf.cast(encoder_decoder_attention_bias,
+                                             tf.bfloat16)
+  return (encoder_input, encoder_self_attention_bias,
+          encoder_decoder_attention_bias)
+
+
+def transformer_prepare_decoder(
+    targets_emb_var, targets, hparams, features=None):
+  """Prepare one shard of the model for the decoder.
+
+  Args:
+    targets_emb_var: a Tensor
+    targets: a Tensor.
+    hparams: run hyperparameters
+    features: optionally pass the entire features dictionary as well.
+      This is needed now for "packed" datasets.
+
+  Returns:
+    decoder_input: a Tensor, bottom of decoder stack
+    decoder_self_attention_bias: a bias tensor for use in decoder self-attention
+  """
+  decoder_self_attention_bias = (
+      common_attention.attention_bias_lower_triangle(
+          common_layers.shape_list(targets)[1]))
+
+  if features and "targets_segmentation" in features:
+    # "Packed" dataset - keep the examples from seeing each other.
+    targets_segmentation = features["targets_segmentation"]
+    targets_position = features["targets_position"]
+    decoder_self_attention_bias += common_attention.attention_bias_same_segment(
+        targets_segmentation, targets_segmentation)
+  else:
+    targets_position = None
+  decoder_input = tf.gather(
+      targets_emb_var, common_layers.shift_right_2d(targets))
+  with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+    decoder_input = common_attention.add_positional_embedding(
+        decoder_input, hparams.max_length, "positional_embedding",
+        targets_position)
+
+  if hparams.activation_dtype == "bfloat16":
+    decoder_self_attention_bias = tf.cast(decoder_self_attention_bias,
+                                          tf.bfloat16)
+  return (decoder_input, decoder_self_attention_bias)
+
+
+def transformer_encoder(encoder_input,
+                        encoder_self_attention_bias,
+                        hparams,
+                        name="encoder",
+                        losses=None):
+  """A stack of transformer layers.
+
+  Args:
+    encoder_input: a Tensor
+    encoder_self_attention_bias: bias Tensor for self-attention
+       (see common_attention.attention_bias())
+    hparams: hyperparameters for model
+    name: a string
+    losses: optional list onto which to append extra training losses
+
+  Returns:
+    y: a Tensors
+  """
+  x = encoder_input
+  with tf.variable_scope(name):
+    num_layer_norms = hparams.num_encoder_layers * 2 + 1
+    layer_norm_combined_var = tf.get_variable(
+        "layer_norm_scale", [num_layer_norms, hparams.d_model])
+    layer_norm_vars = tf.unstack(layer_norm_combined_var, num_layer_norms)
+    def normalize(x):
+      scale = layer_norm_vars.pop(0)
+      variance = tf.reduce_mean(tf.square(x), -1, keep_dims=True)
+      return x * tf.rsqrt(variance + hparams.norm_epsilon) * scale
+    for layer in range(hparams.num_encoder_layers):
+      with tf.variable_scope("layer_%d" % layer):
+        x += multihead_attention_compat(
+            normalize(x),
+            None,
+            encoder_self_attention_bias,
+            kv_channels=hparams.d_kv,
+            heads=hparams.num_heads,
+            name="self_attention")
+        x += transformer_feedforward_layer(normalize(x), hparams, losses=losses)
+    x = normalize(x)
+    return x
+
+
+def transformer_decoder(decoder_input,
+                        encoder_output,
+                        decoder_self_attention_bias,
+                        encoder_decoder_attention_bias,
+                        hparams,
+                        cache=None,
+                        name="decoder",
+                        losses=None):
+  """A stack of transformer layers.
+
+  Args:
+    decoder_input: a Tensor
+    encoder_output: a Tensor
+    decoder_self_attention_bias: bias Tensor for self-attention
+      (see common_attention.attention_bias())
+    encoder_decoder_attention_bias: bias Tensor for encoder-decoder attention
+      (see common_attention.attention_bias())
+    hparams: hyperparameters for model
+    cache: dict, containing tensors which are the results of previous
+        attentions, used for fast decoding.
+    name: a string
+    losses: optional list onto which to append extra training losses
+
+  Returns:
+    y: a Tensors
+  """
+  x = decoder_input
+  with tf.variable_scope(name):
+    num_layer_norms = (
+        hparams.num_decoder_layers * (2 if encoder_output is None else 3) + 1)
+    layer_norm_combined_var = tf.get_variable(
+        "layer_norm_scale", [num_layer_norms, hparams.d_model])
+    layer_norm_vars = tf.unstack(layer_norm_combined_var, num_layer_norms)
+    def normalize(x):
+      scale = layer_norm_vars.pop(0)
+      variance = tf.reduce_mean(tf.square(x), -1, keep_dims=True)
+      return x * tf.rsqrt(variance + hparams.norm_epsilon) * scale
+    for layer in range(hparams.num_decoder_layers):
+      layer_name = "layer_%d" % layer
+      layer_cache = cache[layer_name] if cache is not None else None
+      with tf.variable_scope(layer_name):
+        x += multihead_attention_compat(
+            normalize(x),
+            None,
+            decoder_self_attention_bias,
+            kv_channels=hparams.d_kv,
+            heads=hparams.num_heads,
+            cache=layer_cache,
+            name="self_attention")
+        if encoder_output is not None:
+          x += multihead_attention_compat(
+              normalize(x),
+              encoder_output,
+              encoder_decoder_attention_bias,
+              kv_channels=hparams.d_kv,
+              heads=hparams.num_heads,
+              cache=layer_cache,
+              name="encdec_attention")
+        x += transformer_feedforward_layer(normalize(x), hparams, losses=losses)
+    x = normalize(x)
+    return x
+
+
+def transformer_feedforward_layer(x, hparams, losses=None):
+  """Feed-forward layer in the transformer.
+
+  Args:
+    x: a Tensor of shape [batch_size, length, hparams.d_model]
+    hparams: hyperparameters for model
+    losses: an optional list
+
+  Returns:
+    a Tensor of shape [batch_size, length, hparams.d_model]
+
+  Raises:
+    ValueError: If losses arg is None, but layer generates extra losses.
+  """
+  del losses
+  feedforward_layer = hparams.feedforward_layer
+  if feedforward_layer == "dense_relu_dense":
+    return dense_relu_dense_compat(x, hparams.d_ff)
+  else:
+    raise ValueError("Unknown hparams.feedforward_layer = %s"
+                     % hparams.feedforward_layer)
+
+
+def dense_relu_dense_compat(x, filter_depth, name=None):
+  """Hidden layer with RELU activation followed by linear projection.
+
+  Args:
+    x: a Tensor
+    filter_depth: integer
+    name: an optional string
+
+  Returns:
+    a tf.Tensor
+  """
+  with tf.variable_scope(name, default_name="dense_relu_dense"):
+    io_channels = x.shape.as_list()[-1]
+    w = tf.get_variable("kernel", [2, io_channels, filter_depth])
+    wi, wo = tf.unstack(w, num=2, axis=0)
+    h = tf.nn.relu(tf.tensordot(x, wi, axes=[[-1], [0]]))
+    return tf.tensordot(h, wo, axes=[[-1], [1]])
+
+
+def multihead_attention_compat(query_antecedent,
+                               memory_antecedent,
+                               mask,
+                               kv_channels,
+                               heads,
+                               cache=None,
+                               name="multihead_attention"):
+  """Multihead scaled-dot-product attention with input/output transformations.
+
+  In order to use only one variable containing the four weight matrices
+  packed together, we insist that the query and memory antecedents have the
+  same dimensionality (io_channels) and that the keys and values have the
+  same dimensionality (kv_channels).
+
+  Args:
+    query_antecedent: a Tensor with shape [batch, query_length, io_channels]
+    memory_antecedent: a Tensor with shape
+      [batch, memory_length, io_channels] (optional)
+    mask: mask Tensor (see attention_mask())
+    kv_channels: integer
+    heads: integer
+    cache: an optional dict
+    name: an optional string.
+
+  Returns:
+    A Tensor with shape [batch, qlen, io_channels]
+
+  Raises:
+    ValueError: if the dimensions do not match.
+  """
+  memory_or_query_antecedent = (
+      memory_antecedent if memory_antecedent is not None
+      else query_antecedent)
+  io_channels = memory_or_query_antecedent.shape.as_list()[-1]
+  with tf.variable_scope(name,
+                         default_name="multihead_attention",
+                         values=[query_antecedent, memory_antecedent],
+                         reuse=tf.AUTO_REUSE):
+    var = tf.get_variable("qkvo", [4, heads, io_channels, kv_channels])
+    q_var, k_var, v_var, o_var = tf.unstack(var, num=4, axis=0)
+    if cache is None or memory_antecedent is None:
+      k = tf.einsum("bmi,hik->bhmk", memory_or_query_antecedent, k_var)
+      v = tf.einsum("bmi,hiv->bhmv", memory_or_query_antecedent, v_var)
+      if query_antecedent is None:
+        # we are computing the cache.
+        return k, v
+      q = tf.einsum("bqi,hik->bhqk", query_antecedent, q_var)
+    if cache is not None:
+      if memory_antecedent is not None:
+        q = tf.einsum("bqi,hik->bhqk", query_antecedent, q_var)
+        k = cache["k_encdec"]
+        v = cache["v_encdec"]
+      else:
+        k = cache["k"] = tf.concat([cache["k"], k], axis=2)
+        v = cache["v"] = tf.concat([cache["v"], v], axis=2)
+    logits = tf.einsum("bhqk,bhmk->bhqm", q, k)
+    if mask is not None:
+      logits += mask
+    weights = tf.nn.softmax(logits)
+    o = tf.einsum("bhqm,bhmv->bhqv", weights, v)
+    return tf.einsum("bhqv,hiv->bqi", o, o_var)
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py b/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py
new file mode 100644
index 000000000..fc95edc63
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py
@@ -0,0 +1,155 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 Transformer on Mesh TensorFlow."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+from tensor2tensor.mesh_tensorflow import mtf_transformer
+from tensor2tensor.mesh_tensorflow import placement_mesh_impl
+
+import tensorflow as tf
+
+# Constants shared between all functions.
+BATCH_SIZE = 2
+INPUT_LENGTH = 6
+TARGET_LENGTH = 6
+VOCAB_SIZE = 128
+
+
+def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
+              has_input=True, model_cls=mtf_transformer.MtfTransformer):
+  if hparams is None:
+    hparams = mtf_transformer.mtf_transformer_single()
+  hparams.max_length = INPUT_LENGTH
+  hparams.batch_size = BATCH_SIZE
+
+  p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE, VOCAB_SIZE)
+  if not has_input:
+    p_hparams.input_modality = {}
+  hparams.problem_hparams = p_hparams
+
+  inputs = -1 + np.random.random_integers(
+      VOCAB_SIZE, size=(BATCH_SIZE, INPUT_LENGTH, 1, 1))
+  targets = -1 + np.random.random_integers(
+      VOCAB_SIZE, size=(BATCH_SIZE, TARGET_LENGTH, 1, 1))
+  features = {
+      "targets": tf.constant(targets, dtype=tf.int32, name="targets"),
+      "target_space_id": tf.constant(1, dtype=tf.int32)
+  }
+  if has_input:
+    features["inputs"] = tf.constant(inputs, dtype=tf.int32, name="inputs")
+
+  return model_cls(hparams, mode, p_hparams), features, hparams
+
+
+def get_placement_mesh(hparams):
+  graph = mtf.Graph()
+  mesh = mtf.Mesh(graph, "my_mesh")
+  mesh_shape = mtf.parse_mesh_shape(hparams.mesh_shape)
+  mesh_size = mtf.list_product(mesh_shape)
+
+  mesh_devices = [""] * mesh_size
+  mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+      mesh_shape, mtf.parse_layout(hparams.layout), mesh_devices)
+  return mesh, mesh_impl
+
+
+class MtfTransformerTest(tf.test.TestCase):
+
+  def testMtfTransformer(self):
+    hparams = mtf_transformer.mtf_transformer_single()
+
+    model, features, hparams = get_model(hparams)
+    hparams.mesh_shape = ""
+    hparams.layout = ""
+    mesh, mesh_impl = get_placement_mesh(hparams)
+
+    logits, _ = model.mtf_model_fn(features, mesh)
+    lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
+    tf_group = lowering.copy_masters_to_slices()
+    tf_logits = lowering.outfeed(logits)
+
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      session.run(tf_group)
+      res = session.run(tf_logits)
+    self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, VOCAB_SIZE))
+
+  def testMtfTransformerDataParallel(self):
+    hparams = mtf_transformer.mtf_transformer_single()
+
+    model, features, hparams = get_model(hparams)
+    hparams.mesh_shape = "2"
+    hparams.layout = "batch:0"
+    mesh, mesh_impl = get_placement_mesh(hparams)
+
+    logits, _ = model.mtf_model_fn(features, mesh)
+    lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
+    tf_group = lowering.copy_masters_to_slices()
+    tf_logits = lowering.outfeed(logits)
+
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      session.run(tf_group)
+      res = session.run(tf_logits)
+    self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, VOCAB_SIZE))
+
+  def testMtfTransformerModelParallel(self):
+    hparams = mtf_transformer.mtf_transformer_single()
+
+    model, features, hparams = get_model(hparams)
+    hparams.mesh_shape = "2"
+    hparams.layout = "length:0"
+    mesh, mesh_impl = get_placement_mesh(hparams)
+
+    logits, _ = model.mtf_model_fn(features, mesh)
+    lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
+    tf_group = lowering.copy_masters_to_slices()
+    tf_logits = lowering.outfeed(logits)
+
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      session.run(tf_group)
+      res = session.run(tf_logits)
+    self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, VOCAB_SIZE))
+
+  def testMtfTransformerDataModelParallel(self):
+    hparams = mtf_transformer.mtf_transformer_single()
+
+    model, features, hparams = get_model(hparams)
+    hparams.mesh_shape = "2.2"
+    hparams.layout = "batch:0;vocab:1;d_ff:1;heads:1"
+    mesh, mesh_impl = get_placement_mesh(hparams)
+
+    logits, _ = model.mtf_model_fn(features, mesh)
+    lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
+    tf_group = lowering.copy_masters_to_slices()
+    tf_logits = lowering.outfeed(logits)
+
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      session.run(tf_group)
+      res = session.run(tf_logits)
+    self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, VOCAB_SIZE))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/mesh_tensorflow/mtf_utils.py b/tensor2tensor/mesh_tensorflow/mtf_utils.py
new file mode 100644
index 000000000..70cee4923
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_utils.py
@@ -0,0 +1,29 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Common utilities for mesh tensorflow."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import contextlib
+
+from tensorflow.python.framework import ops
+
+
+@contextlib.contextmanager
+def outside_all_rewrites():
+  with ops.control_dependencies(None):
+    yield
diff --git a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
new file mode 100644
index 000000000..32fe65f69
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
@@ -0,0 +1,492 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Placement Mesh Implementation (for CPU/GPU clusters)."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+import google3
+from past.builtins import xrange
+
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+import tensorflow as tf
+
+
+class PlacementMeshImpl(mtf.MeshImpl):
+  """Mesh implemented using explicit device placement."""
+
+  def __init__(self, shape, layout, devices):
+    super(PlacementMeshImpl, self).__init__(shape, layout)
+    self._devices = devices
+
+  class LaidOutTensor(object):
+    """One Slice for each processor."""
+
+    def __init__(self, tensor_list):
+      self._tensor_list = tensor_list
+
+    def __repr__(self):
+      return "[" + ",".join([str(t) for t in self._tensor_list]) + "]"
+
+    @property
+    def tensor_list(self):
+      return self._tensor_list
+
+    @classmethod
+    def from_tensor_list(cls, tensor_list):
+      return cls(tensor_list)
+
+    @property
+    def all_slices(self):
+      return self._tensor_list
+
+    def to_laid_out_tensor(self):
+      return self
+
+  class LaidOutVariable(object):
+    """Maintains slice-variables and copy operations."""
+
+    def __init__(self, variable, mesh_impl):
+      """Create a LaidOutVariable.
+
+      Args:
+        variable: a Variable (Operation)
+        mesh_impl: a MeshImpl
+      """
+      self._variable = variable
+      self._mesh_impl = mesh_impl
+      shape = variable.outputs[0].shape
+      dtype = variable.outputs[0].dtype
+      slice_shape = mesh_impl.slice_shape(shape)
+      base_name = variable.name
+      slices = []
+      for pnum in xrange(mesh_impl.size):
+        with tf.device(mesh_impl.devices[pnum]):
+          slices.append(tf.get_variable(
+              base_name + "_slice_%d" % pnum,
+              slice_shape,
+              dtype=dtype, collections=[]))
+      self._laid_out_tensor = mesh_impl.LaidOutTensor(slices)
+      self._copy_master_to_slices = self.assign_to_slices(
+          mesh_impl.make_slices(variable.master, shape))
+      self._copy_slices_to_master = tf.assign(
+          variable.master,
+          mesh_impl.combine_slices(self._laid_out_tensor.all_slices, shape))
+
+    def assign_to_slices(self, slices):
+      """Assign to the slice variables.
+
+      Args:
+        slices: a list of tf.Tensor
+
+      Returns:
+        a tf.operation
+      """
+      return tf.group(mtf.parallel(
+          self._mesh_impl.devices, tf.assign,
+          self.laid_out_tensor.all_slices, slices))
+
+    @property
+    def laid_out_tensor(self):
+      return self._laid_out_tensor
+
+    @property
+    def copy_master_to_slices(self):
+      return self._copy_master_to_slices
+
+    @property
+    def copy_slices_to_master(self):
+      return self._copy_slices_to_master
+
+  def slicewise(self, fn, *inputs):
+    """Execute a function in parallel on all slices.
+
+    Args:
+      fn: a function from tf.Tensors to tf.Tensor or a tuple of tf.Tensors.
+      *inputs: a list of inputs.  Each input is either a LaidOutTensor or
+        is convertible to a tf.Tensor.
+    Returns:
+      a LaidOutTensor, or a tuple of LaidOutTensors if fn returns a tuple.
+    """
+    if fn == tf.add:
+      assert len(inputs) == 2
+      if isinstance(inputs[0], mtf.LazyAllreduceSum):
+        # sum of LazyAllreduceSum (keep delaying the allreduce)
+        return inputs[0] + inputs[1]
+    # convert all inputs to LaidOutTensor where possible
+    inputs = mtf.convert_args_to_laid_out_tensors(inputs)
+    inputs = [x.tensor_list if isinstance(x, self.LaidOutTensor)
+              else [x] * len(self.devices) for x in inputs]
+    ret = mtf.parallel(self.devices, fn, *inputs)
+    if isinstance(ret[0], tuple):
+      ret = mtf.transpose_list_of_lists(ret)
+      return tuple([self.LaidOutTensor(t) for t in ret])
+    else:
+      return self.LaidOutTensor(ret)
+
+  def Print(self, x, data, message, **kwargs):  # pylint: disable=invalid-name
+    """call tf.Print.
+
+    Args:
+      x: a LaidOutTensor
+      data: a list of LaidOutTensor
+      message: a string
+      **kwargs: keyword arguments to tf.print
+    Returns:
+      a LaidOutTensor
+    """
+    tf.logging.info("PlacementMeshImpl::Print")
+    new_slices = x.tensor_list[:]
+    with tf.device(self._devices[0]):
+      new_slices[0] = tf.Print(
+          new_slices[0], [t for d in data for t in d.tensor_list],
+          message, **kwargs)
+    return self.LaidOutTensor(new_slices)
+
+  def allreduce(self, x, mesh_axes, reduction_fn_string):
+    """Grouped allreduce, (across the given dimensions).
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axes: a list of integers - the mesh dimensions to be reduced
+      reduction_fn_string: "SUM" or "MAX"
+    Returns:
+      a LaidOutTensor
+    """
+    return self._collective_with_groups(
+        x, mesh_axes, functools.partial(
+            allreduce_ring, reduction_fn_string=reduction_fn_string))
+
+  def allconcat(self, x, mesh_axis, concat_axis):
+    """Grouped allconcat (like MPI allgather followed by concat).
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axis: an integer - the mesh axis along which to group
+      concat_axis: an integer (the Tensor axis along which to concatenate)
+    Returns:
+      a LaidOutTensor
+    """
+    return self._collective_with_groups(
+        x, [mesh_axis],
+        functools.partial(allconcat_ring, concat_axis=concat_axis))
+
+  def alltoall(self, x, mesh_axis, split_axis, concat_axis):
+    """Grouped alltoall.
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axis: an integer the mesh axis along which to group
+      split_axis: an integer (the Tensor axis along which to split)
+      concat_axis: an integer (the Tensor axis along which to concatenate)
+    Returns:
+      a LaidOutTensor
+    """
+    return self._collective_with_groups(
+        x, [mesh_axis],
+        functools.partial(
+            alltoall_ring, split_axis=split_axis, concat_axis=concat_axis))
+
+  def _collective_with_groups(self, x, mesh_axes, collective):
+    """Grouped collective, (across the given dimensions).
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axes: a list of integers - the mesh dimensions to be reduced
+      collective: fn from list(tf.Tensor), list(device) -> list(tf.Tensor)
+    Returns:
+      a LaidOutTensor
+    """
+    if not mesh_axes:
+      return x
+    x = x.to_laid_out_tensor()
+    if len(mesh_axes) == self.ndims:
+      return self.LaidOutTensor(collective(x.tensor_list, self._devices))
+    else:
+      groups = mtf.processor_groups(self.shape, mesh_axes)
+      ret = [None] * self.size
+      for g in groups:
+        inputs = [x.tensor_list[pnum] for pnum in g]
+        devices = [self._devices[pnum] for pnum in g]
+        reduced = collective(inputs, devices)
+        for pnum, y in zip(g, reduced):
+          ret[pnum] = y
+      return self.LaidOutTensor(ret)
+
+  def random(self, shape, tf_fn, kwargs):
+    """Call a random tf operation (e.g. random_uniform).
+
+    Args:
+      shape: a TensorShape
+      tf_fn: a function such as tf.random_uniform
+      kwargs: kwargs to pass to tf_fn, except for seed
+
+    Returns:
+      a LaidOutTensor
+    """
+    slice_shape = self.slice_shape(shape)
+    var_scope = tf.get_variable_scope().name
+    def my_fn(pnum):
+      # seeds are necessary to make sure that slices that should have the
+      # same values actually do have the same values.
+      seed = hash("%s%s" % (var_scope, self.slice_begin(shape, pnum)))
+      return tf_fn(slice_shape, seed=seed, **kwargs)
+    return self.slicewise(my_fn, self.laid_out_pnum())
+
+  def laid_out_pnum(self):
+    """Returns a LaidOutTensor containing the processor number."""
+    return self.LaidOutTensor(list(range(self.size)))
+
+  @property
+  def devices(self):
+    return self._devices
+
+  def outfeed(self, x, laid_out_x):
+    """Turn a Tensor into a tf.Tensor.
+
+    Args:
+      x: a Tensor
+      laid_out_x: a LaidOutTensor
+    Returns:
+      a tf.Tensor
+    """
+    return self.combine_slices(laid_out_x.all_slices, x.shape)
+
+  def infeed(self, x, tf_x):
+    """Infeed a tf.Tensor, producing a LaidOutTensor.
+
+    Args:
+      x: a Tensor
+      tf_x: a tf.Tensor
+    Returns:
+      a LaidOutTensor
+    """
+    return self.LaidOutTensor(self.make_slices(tf_x, x.shape))
+
+
+def allreduce_ring_single_shard(xs, devices, reduction_fn_string="SUM"):
+  """Compute the reduction of all Tensors and put the result everywhere.
+
+  Performance-optimized for a ring of devices.
+
+  Args:
+    xs: a list of n tf.Tensors
+    devices: a list of strings
+    reduction_fn_string: "SUM" or "MAX"
+
+  Returns:
+    a list of n Tensors
+  Raises:
+    ValueError: if devices is not a list of n strings
+  """
+  n = len(xs)
+  binary_reduction = mtf.binary_reduction_fn(reduction_fn_string)
+  assert len(devices) == n, "devices must be a list of length len(xs)"
+  if n == 1:
+    return xs
+  result = [None] * n
+  if n % 2 == 0:
+    left_center = n // 2 - 1
+    right_center = left_center + 1
+  else:
+    left_center = n // 2
+    right_center = left_center
+  left_sum = xs[0]
+  for i in xrange(1, left_center + 1):
+    with tf.device(devices[i]):
+      left_sum = binary_reduction(left_sum, xs[i])
+  right_sum = xs[n-1]
+  for i in reversed(xrange(left_center + 1, n - 1)):
+    with tf.device(devices[i]):
+      right_sum = binary_reduction(xs[i], right_sum)
+  with tf.device(devices[left_center]):
+    result[left_center] = binary_reduction(left_sum, right_sum)
+  if n % 2 == 0:
+    with tf.device(devices[right_center]):
+      result[right_center] = binary_reduction(left_sum, right_sum)
+  for i in reversed(xrange(left_center)):
+    with tf.device(devices[i]):
+      result[i] = tf.identity(result[i + 1])
+  for i in xrange(right_center + 1, n):
+    with tf.device(devices[i]):
+      result[i] = tf.identity(result[i - 1])
+  return result
+
+
+def allreduce_ring(xs, devices, reduction_fn_string="SUM"):
+  """Compute the reduction of all Tensors and put the result everywhere.
+
+  Performance-optimized for a ring of devices.
+
+  Args:
+    xs: a list of n tf.Tensors
+    devices: a list of strings
+    reduction_fn_string: "SUM" or "MAX"
+
+  Returns:
+    a list of n Tensors
+  Raises:
+    ValueError: if devices is not a list of n strings
+  """
+  n = len(xs)
+  if len(devices) != n:
+    raise ValueError("devices must be a list of length len(xs)")
+  if n == 1:
+    return xs
+  shape = xs[0].shape.as_list()
+  # tf.logging.info("allreduce_ring shape = %s" % shape)
+  size = None if None in shape else mtf.list_product(shape)
+  if size is None or size < 1024 or size % n != 0:
+    return allreduce_ring_single_shard(xs, devices, reduction_fn_string)
+
+  def _circular_shift(l, n):
+    n %= len(l)
+    return l[-n:] + l[:-n]
+  def _flatten_and_split(x):
+    return tf.split(tf.reshape(x, [size]), n)
+  def _concat_and_reshape(xs):
+    return tf.reshape(tf.concat(xs, 0), shape)
+
+  # [device, shard]
+  x_split = mtf.parallel(devices, _flatten_and_split, xs)
+  x_split_t = mtf.transpose_list_of_lists(x_split)
+
+  y_split_t = []
+  for shard in xrange(n):
+    shard_xs = _circular_shift(x_split_t[shard], shard)
+    shard_devices = _circular_shift(devices, shard)
+    shard_ys = allreduce_ring_single_shard(
+        shard_xs, shard_devices, reduction_fn_string)
+    y_split_t.append(_circular_shift(shard_ys, -shard))
+  y_split = mtf.transpose_list_of_lists(y_split_t)
+  ys = mtf.parallel(devices, _concat_and_reshape, y_split)
+  return ys
+
+
+def allconcat_ring(xs, devices, concat_axis):
+  """Concatenate all Tensors everywhere.
+
+  Performance-optimized for a ring of devices.
+
+  Args:
+    xs: a list of n tf.Tensors
+    devices: a list of n strings
+    concat_axis: an integer
+
+  Returns:
+    a list of n Tensors
+  """
+  n = len(xs)
+  if n == 1:
+    return xs
+  # [target, source]
+  parts = [[xs[target] if target == source else None for source in xrange(n)]
+           for target in xrange(n)]
+  for distance in xrange(1, n // 2 + 1):
+    for target in xrange(n):
+      source = (target + distance) % n
+      if parts[target][source] is None:
+        with tf.device(devices[target]):
+          parts[target][source] = tf.identity(parts[(target + 1) % n][source])
+      source = (target - distance) % n
+      if parts[target][source] is None:
+        with tf.device(devices[target]):
+          parts[target][source] = tf.identity(parts[(target - 1) % n][source])
+  return mtf.parallel(devices, tf.concat, parts, axis=[concat_axis] * n)
+
+
+def alltoall_pointtwise(xs, devices, split_axis, concat_axis):
+  """MPI alltoall operation.
+
+  Implementation of alltoall using pointwise communication.
+
+  Args:
+    xs: a list of n tf.Tensors
+    devices: a list of n strings
+    split_axis: an integer
+    concat_axis: an integer
+
+  Returns:
+    a list of n Tensors
+  """
+  n = len(xs)
+  if n == 1:
+    return xs
+  # [target, source]
+  parts = mtf.transpose_list_of_lists(
+      mtf.parallel(devices, tf.split, xs, [n] * n, axis=[split_axis] * n))
+  return mtf.parallel(devices, tf.concat, parts, axis=[concat_axis] * n)
+
+
+def alltoall_ring(xs, devices, split_axis, concat_axis):
+  """MPI alltoall operation.
+
+  Performance-optimized for a ring of devices.
+
+  Args:
+    xs: a list of n tf.Tensors
+    devices: a list of n strings
+    split_axis: an integer
+    concat_axis: an integer
+
+  Returns:
+    a list of n Tensors
+  """
+  n = len(xs)
+  if n == 1:
+    return xs
+  # set up
+  # [target, source]
+  parts = [[None] * n for i in xrange(n)]
+  def my_split(x, size_splits):
+    total_size = tf.shape(x)[split_axis]
+    part_size = total_size // sum(size_splits)
+    return tf.split(x, [s * part_size for s in size_splits], axis=split_axis)
+  forward_message_size = (n - 1) // 2
+  backward_message_size = (n - 1) - forward_message_size
+  forward_messages = [None] * n
+  backward_messages = [None] * n
+  for i in xrange(n):
+    with tf.device(devices[i]):
+      if i >= backward_message_size:
+        a, b, c, d = my_split(
+            xs[i], [i - backward_message_size,
+                    backward_message_size, 1, n - i - 1])
+        backward_messages[i] = b
+        parts[i][i] = c
+        forward_messages[i] = tf.concat([d, a], axis=split_axis)
+      else:
+        a, b, c, d = my_split(
+            xs[i], [i, 1, forward_message_size, backward_message_size - i])
+        backward_messages[i] = tf.concat([d, a], axis=split_axis)
+        parts[i][i] = b
+        forward_messages[i] = c
+  for step in xrange(1, max(forward_message_size, backward_message_size) + 1):
+    new_forward_messages = [None] * n
+    new_backward_messages = [None] * n
+    for i in xrange(n):
+      with tf.device(devices[i]):
+        if forward_message_size > 0:
+          parts[i][(i - step) % n], new_forward_messages[i] = my_split(
+              forward_messages[(i - 1) % n], [1, forward_message_size - 1])
+        if backward_message_size > 0:
+          new_backward_messages[i], parts[i][(i + step) % n] = my_split(
+              backward_messages[(i + 1) % n], [backward_message_size - 1, 1])
+    forward_message_size -= 1
+    backward_message_size -= 1
+    forward_messages = new_forward_messages
+    backward_messages = new_backward_messages
+  return mtf.parallel(devices, tf.concat, parts, axis=[concat_axis] * n)
diff --git a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
new file mode 100644
index 000000000..012527b6c
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
@@ -0,0 +1,342 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""SIMD Mesh implementation (for TPU/XLA)."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import google3
+
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+from tensor2tensor.mesh_tensorflow import mtf_utils
+from tensor2tensor.mesh_tensorflow import tpu_variables
+
+import tensorflow as tf
+
+from tensorflow.contrib.tpu.python.ops import tpu_ops
+from tensorflow.python.framework import ops
+
+
+class SimdMeshImpl(mtf.MeshImpl):
+  """Mesh implementation for TPU using SIMD and MPI operations."""
+
+  def __init__(self, shape, layout, devices, device_assignment):
+    super(SimdMeshImpl, self).__init__(shape, layout)
+    self._devices = devices
+    self._device_assignment = device_assignment
+    tf.logging.info("SimdMeshImpl init: {0} {1}".format(shape, layout))
+    self._pnum_tensor = None
+
+  @property
+  def pnum_tensor(self):
+    if self._pnum_tensor is not None:
+      return self._pnum_tensor
+    with mtf_utils.outside_all_rewrites():
+      tf.logging.info("Create pnum_tensor")
+      self._pnum_tensor = tpu_ops.tpu_replicated_input(
+          list(range(self.size)), name="pnum_constants")
+      return self._pnum_tensor
+
+  class LaidOutTensor(object):
+    """One Slice."""
+
+    def __init__(self, tensor_list):
+      assert isinstance(tensor_list, list)
+      self._tensor_list = tensor_list
+
+    def __repr__(self):
+      return "[" + ",".join([str(t) for t in self._tensor_list]) + "]"
+
+    @property
+    def tensor_list(self):
+      return self._tensor_list
+
+    @property
+    def one_slice(self):
+      return self._tensor_list[0]
+
+    @classmethod
+    def from_tensor_list(cls, tensor_list):
+      return cls(tensor_list)
+
+    @property
+    def all_slices(self):
+      return self._tensor_list
+
+    def to_laid_out_tensor(self):
+      return self
+
+  class LaidOutVariable(object):
+    """Maintains slice-variables and copy operations."""
+
+    def __init__(self, variable, mesh_impl):
+      """Create a LaidOutVariable.
+
+      Args:
+        variable: a Variable (Operation)
+        mesh_impl: a MeshImpl
+      """
+      self._variable = variable
+      self._mesh_impl = mesh_impl
+      shape = variable.outputs[0].shape
+      dtype = variable.outputs[0].dtype
+      slice_shape = mesh_impl.slice_shape(shape)
+      base_name = variable.name
+      slices = []
+      for pnum in xrange(mesh_impl.size):
+        slice_var_name = base_name + "_slice_%d" % pnum
+        tpu_device = mesh_impl.device_assignment.tpu_device(replica=pnum)
+        # The initializer is unimportant, since the slice variables will be
+        # overwritten.  zeros_initializer() is here to avoid the default
+        # initialization which adds lots of useless operations to the TF graph.
+        with ops.device(tpu_device):
+          slices.append(
+              tf.get_variable(
+                  slice_var_name,
+                  slice_shape,
+                  dtype=dtype,
+                  collections=[],
+                  initializer=tf.zeros_initializer()))
+      self._laid_out_tensor = mesh_impl.LaidOutTensor(
+          [tpu_variables.ReplicatedVariable(base_name, slices)])
+      with tf.device("cpu:0"), mtf_utils.outside_all_rewrites():
+        self._copy_master_to_slices = self.assign_to_slices(
+            mesh_impl.make_slices(variable.master, shape),
+            assign_to_tensor_list=slices)
+        self._copy_slices_to_master = tf.assign(
+            variable.master,
+            mesh_impl.combine_slices(slices, shape, device="cpu:0"))
+
+    def assign_to_slices(self, slice_values, assign_to_tensor_list=None):
+      """Assign to the slice variables.
+
+      Args:
+        slice_values: a list of tf.Tensor
+        assign_to_tensor_list: an optional list of tf.Variable
+
+      Returns:
+        a tf.operation
+      """
+      if assign_to_tensor_list is None:
+        assign_to_tensor_list = self._laid_out_tensor.all_slices
+      # Handle both N -> 1 and N -> N cases.
+      num_slices = min(
+          len(assign_to_tensor_list), len(slice_values))
+      devices = [""] * num_slices
+      return tf.group(
+          mtf.parallel(devices, tf.assign, assign_to_tensor_list[:num_slices],
+                       slice_values[:num_slices]))
+
+    @property
+    def laid_out_tensor(self):
+      return self._laid_out_tensor
+
+    @property
+    def copy_master_to_slices(self):
+      return self._copy_master_to_slices
+
+    @property
+    def copy_slices_to_master(self):
+      return self._copy_slices_to_master
+
+  def laid_out_pnum(self):
+    """Returns a LaidOutTensor containing the processor number.
+
+    Returns:
+      a LaidOutTensor where each slice is an integer scalar
+    """
+    return self.LaidOutTensor([self.pnum_tensor])
+
+  def allreduce(self, x, mesh_axes, reduction_fn_string):
+    """Grouped allreduce, (summed across the given dimensions).
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axes: a list of integers
+      reduction_fn_string: "SUM"
+    Returns:
+      a LaidOutTensor
+    Raises:
+      ValueError: if the reduction is not yet implemented.
+    """
+    if not mesh_axes:
+      return x
+    x = x.to_laid_out_tensor()
+    if reduction_fn_string == "SUM":
+      partitioning = [
+          mtf.pnum_to_group(self.shape, mesh_axes, pnum)
+          for pnum in xrange(self.size)]
+      return self.LaidOutTensor(
+          [tpu_ops.cross_replica_sum(x.one_slice, partitioning)])
+    else:
+      for axis in mesh_axes:
+        x = self.allconcat(x, axis, 0, stack=True)
+        x = self.LaidOutTensor(
+            [mtf.reduction_fn(reduction_fn_string)(x.one_slice, 0)])
+      return x
+
+  def allconcat(self, x, mesh_axis, concat_axis, stack=False):
+    """Grouped allconcat (like MPI allgather followed by concat).
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axis: an integer - the mesh axis along which to group
+      concat_axis: an integer (the Tensor axis along which to concatenate)
+      stack: a boolean - whether to stack instead of concat
+    Returns:
+      a LaidOutTensor
+    """
+    x = x.to_laid_out_tensor()
+    coord = self.laid_out_pcoord(mesh_axis)
+    t = x.one_slice
+    old_shape = t.shape.as_list()
+    num_parts = self.shape[mesh_axis]
+    t = tf.expand_dims(t, concat_axis)
+    t *= tf.reshape(
+        tf.one_hot(coord.one_slice, num_parts, dtype=t.dtype),
+        [num_parts if i == concat_axis else 1
+         for i in xrange(len(old_shape) + 1)])
+    if not stack:
+      new_shape = old_shape[:]
+      new_shape[concat_axis] *= num_parts
+      t = tf.reshape(t, new_shape)
+    return self.allreduce(self.LaidOutTensor([t]), [mesh_axis], "SUM")
+
+  def alltoall(self, x, mesh_axis, split_axis, concat_axis):
+    """Grouped alltoall (like MPI alltoall with splitting and concatenation).
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axis: an integer the mesh axis along which to group
+      split_axis: an integer (the Tensor axis along which to split)
+      concat_axis: an integer (the Tensor axis along which to concatenate)
+    Returns:
+      a LaidOutTensor
+    """
+    x = x.to_laid_out_tensor()
+    x = self.allconcat(x, mesh_axis, concat_axis)
+    x = self.allsplit(x, mesh_axis, split_axis)
+    return x
+
+  def slice(self, tf_tensor, tensor_shape):
+    """"Slice out the correspoding part of tensor given the pnum variable."""
+    tensor_layout = self.tensor_layout(tensor_shape)
+
+    if tensor_layout.is_fully_replicated:
+      return self.LaidOutTensor([tf_tensor])
+    else:
+      slice_shape = self.slice_shape(tensor_shape)
+      slice_begins = [
+          self.slice_begin(tensor_shape, pnum) for pnum in xrange(self.size)
+      ]
+      slice_begins_tensor = tf.stack(slice_begins)
+      # slice on source device
+      selected_slice_begin = tf.gather(slice_begins_tensor, self.pnum_tensor)
+      return self.LaidOutTensor(
+          [tf.slice(tf_tensor, selected_slice_begin, slice_shape)])
+
+  def slicewise(self, fn, *inputs):
+    """Execute a function in parallel on all slices.
+
+    Args:
+      fn: a function from tf.Tensors to tf.Tensor or a tuple of tf.Tensors.
+      *inputs: a list of inputs.  Each input is either a LaidOutTensor or
+        is convertible to a tf.Tensor.
+    Returns:
+      a LaidOutTensor, or a tuple of LaidOutTensors if fn returns a tuple.
+    """
+    if fn == tf.add:
+      assert len(inputs) == 2
+      if isinstance(inputs[0], mtf.LazyAllreduceSum):
+        # sum of LazyAllreduceSum (keep delaying the allreduce)
+        return inputs[0] + inputs[1]
+    # convert all inputs to LaidOutTensor where possible
+    inputs = mtf.convert_args_to_laid_out_tensors(inputs)
+    ret = fn(*[x.one_slice if isinstance(x, self.LaidOutTensor)
+               else x for x in inputs])
+    if isinstance(ret, tuple):
+      return tuple([self.LaidOutTensor([t]) for t in ret])
+    else:
+      return self.LaidOutTensor([ret])
+
+  @property
+  def device_assignment(self):
+    return self._device_assignment
+
+  @property
+  def devices(self):
+    return self._devices
+
+  def random(self, shape, tf_fn, kwargs):
+    """Call a random tf operation (e.g. random_uniform).
+
+    Args:
+      shape: a TensorShape
+      tf_fn: a function such as tf.random_uniform
+      kwargs: kwargs to pass to tf_fn, except for seed
+
+    Returns:
+      a LaidOutTensor
+    """
+    # TODO(noam): can we make things better with stateless_random?
+    slice_shape = self.slice_shape(shape)
+    x = tf_fn(slice_shape, **kwargs)
+    # TPU does not have seeds enabled.  Sync up the
+    # random choices by zeroing out all but the first core per group of
+    # identical slices, then allreducing by group.
+    layout = self.tensor_layout(shape)
+    # we need to sync across these axes.
+    mesh_axes = [i for i in xrange(self.ndims)
+                 if i not in layout.tensor_axis_to_mesh_axis]
+    multiplier = 1.0
+    for axis in mesh_axes:
+      multiplier *= tf.cast(
+          tf.equal(self.laid_out_pcoord(axis).one_slice, 0), x.dtype)
+    x *= multiplier
+    x = self.LaidOutTensor([x])
+    x = self.allreduce(x, mesh_axes, "SUM")
+    return x
+
+  def outfeed(self, x, laid_out_x):
+    """Turn a Tensor into a tf.Tensor.
+
+    Args:
+      x: a Tensor
+      laid_out_x: a LaidOutTensor
+    Returns:
+      a tf.Tensor
+    """
+    tensor_layout = self.tensor_layout(x.shape)
+    if not tensor_layout.is_fully_replicated:
+      raise NotImplementedError(
+          "SimdMeshImpl only supports outfeed of fully-replicated Tensors."
+          " Try reshaping to new dimension names.")
+    return laid_out_x.one_slice
+
+  def infeed(self, x, tf_x):
+    """Infeed a tf.Tensor, producing a LaidOutTensor.
+
+    Args:
+      x: a Tensor
+      tf_x: a tf.Tensor
+    Returns:
+      a LaidOutTensor
+    """
+    return self.slice(tf_x, x.shape)
+
+  @property
+  def supports_control_dependencies(self):
+    return False
diff --git a/tensor2tensor/mesh_tensorflow/tpu_variables.py b/tensor2tensor/mesh_tensorflow/tpu_variables.py
new file mode 100644
index 000000000..0bc3e2c04
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/tpu_variables.py
@@ -0,0 +1,199 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Distributed variable implementation for TPUs."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import contextlib
+
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import gen_resource_variable_ops
+
+
+@contextlib.contextmanager
+def _handle_graph(handle):
+  with handle.graph.as_default():
+    yield
+
+
+def _enclosing_tpu_context():
+  # pylint: disable=protected-access
+  context = ops.get_default_graph()._get_control_flow_context()
+  # pylint: enable=protected-access
+  while context is not None and not isinstance(
+      context, control_flow_ops.XLAControlFlowContext):
+    context = context.outer_context
+  return context
+
+
+class ReplicatedVariable(object):
+  """A replicated variable for use on TPUs.
+
+  When accessed inside a tpu.replicate() context, this variable acts as if it
+  is a single variable whose handle is a replicated input to the computation.
+
+  Outside a tpu.replicate() context currently this object has pretty murky
+  semantics, especially with respect to things such as
+  * initialization
+  * colocation.
+
+  TODO(phawkins): merge this with the TPU DistributionStrategy code.
+  """
+
+  def __init__(self, name, variables):
+    self._name = name
+    self._primary_var = variables[0]
+    self._vars = variables
+    self._cached_value = None
+    self._dtype = variables[0].dtype
+
+  @property
+  def handle(self):
+    tpu_context = _enclosing_tpu_context()
+    if tpu_context is None:
+      return self._primary_var.handle
+
+    return tpu_context.get_replicated_var_handle(self)
+
+  @contextlib.contextmanager
+  def _assign_dependencies(self):
+    """Makes assignments depend on the cached value, if any.
+
+    This prevents undefined behavior with reads not ordered wrt writes.
+
+    Yields:
+      None.
+    """
+    if self._cached_value is not None:
+      with ops.control_dependencies([self._cached_value]):
+        yield
+    else:
+      yield
+
+  @property
+  def initializer(self):
+    return control_flow_ops.group([v.initializer for v in self._vars])
+
+  @property
+  def graph(self):
+    return self._primary_var.graph
+
+  @property
+  def _shared_name(self):
+    return self._common_name
+
+  @property
+  def _unique_id(self):
+    return self._primary_var._unique_id  # pylint: disable=protected-access
+
+  @property
+  def name(self):
+    return self._name
+
+  @property
+  def dtype(self):
+    return self._primary_var.dtype
+
+  @property
+  def shape(self):
+    return self._primary_var.shape
+
+  def get_shape(self):
+    return self._primary_var.get_shape()
+
+  def to_proto(self, export_scope=None):
+    return self._primary_var.to_proto(export_scope=export_scope)
+
+  @property
+  def constraint(self):
+    return None
+
+  @property
+  def op(self):
+    return self.get().op
+
+  def _read_variable_op(self):
+    if _enclosing_tpu_context() is None:
+      return self._primary_var.read_value()
+    v = gen_resource_variable_ops.read_variable_op(self.handle, self._dtype)
+    return v
+
+  def read_value(self):
+    return self._read_variable_op()
+
+  def assign(self, value, use_locking=None, name=None, read_value=False):
+    del use_locking
+    with _handle_graph(self.handle), self._assign_dependencies():
+      value_tensor = ops.convert_to_tensor(value, dtype=self.dtype)
+      assign_op = gen_resource_variable_ops.assign_variable_op(
+          self.handle, value_tensor, name=name)
+    if read_value:
+      return self._read_variable_op()
+    return assign_op
+
+  def assign_add(self, delta, use_locking=None, name=None, read_value=True):
+    del use_locking
+    with _handle_graph(self.handle), self._assign_dependencies():
+      assign_add_op = gen_resource_variable_ops.assign_add_variable_op(
+          self.handle,
+          ops.convert_to_tensor(delta, dtype=self.dtype),
+          name=name)
+    if read_value:
+      return self._read_variable_op()
+    return assign_add_op
+
+  def assign_sub(self, delta, use_locking=None, name=None, read_value=True):
+    del use_locking
+    with _handle_graph(self.handle), self._assign_dependencies():
+      assign_sub_op = gen_resource_variable_ops.assign_sub_variable_op(
+          self.handle,
+          ops.convert_to_tensor(delta, dtype=self.dtype),
+          name=name)
+    if read_value:
+      return self._read_variable_op()
+    return assign_sub_op
+
+  def get(self):
+    return self._primary_var
+
+  def _should_act_as_resource_variable(self):
+    """Pass resource_variable_ops.is_resource_variable check."""
+    pass
+
+  def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
+    """Converts a variable to a tensor."""
+    # pylint: disable=protected-access
+    if _enclosing_tpu_context() is None:
+      return self._primary_var._dense_var_to_tensor(dtype, name, as_ref)
+    # pylint: enable=protected-access
+    if dtype is not None and dtype != self.dtype:
+      return NotImplemented
+    if as_ref:
+      return self.handle
+    else:
+      return self.read_value()
+
+
+# Register a conversion function which reads the value of the variable,
+# allowing instances of the class to be used as tensors.
+def _tensor_conversion(var, dtype=None, name=None, as_ref=False):
+  return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref)  # pylint: disable=protected-access
+
+
+ops.register_tensor_conversion_function(ReplicatedVariable, _tensor_conversion)
+ops.register_dense_tensor_like_type(ReplicatedVariable)

From 9d4ab4c3109d7276f61c42906321678937b073c0 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 7 Aug 2018 09:26:32 -0700
Subject: [PATCH 0486/2720] We've decided against using flake8. Main reason
 being pylint should catch all these errors and is used internally as well as
 externally. And flake8 seems to have corner cases on valid code.

PiperOrigin-RevId: 207727711
---
 .travis.yml | 10 ----------
 setup.py    |  2 +-
 2 files changed, 1 insertion(+), 11 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 6f5f86bad..b162ae9f9 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -48,16 +48,6 @@ install:
   # Make sure we have the latest version of numpy - avoid problems we were
   # seeing with Python 3
   - pip install -q -U numpy
-before_script:
-  # stop the build if there are Python syntax errors or undefined names
-  if [[ "$TF_VERSION" == "1.9.*" ]]; then
-    # * F821: undefined name `name`
-    # * F822: undefined name `name` in `__all__`
-    # * F823: local variable `name` referenced before assignment
-    # * E901: SyntaxError or IndentationError
-    # * E999: SyntaxError -- failed to compile a file into an Abstract Syntax Tree
-    flake8 . --count --select=F821,F822,F823,E901,E999 --show-source --statistics;
-  fi
 script:
   # Check import
   - python -c "from tensor2tensor.models import transformer; print(transformer.Transformer.__name__)"
diff --git a/setup.py b/setup.py
index 2bf760637..ff69570f3 100644
--- a/setup.py
+++ b/setup.py
@@ -52,7 +52,7 @@
         'tensorflow': ['tensorflow>=1.8.0'],
         'tensorflow_gpu': ['tensorflow-gpu>=1.8.0'],
         'tests': [
-            'absl-py', 'pytest', 'mock', 'pylint', 'jupyter', 'gsutil', 'flake8'
+            'absl-py', 'pytest', 'mock', 'pylint', 'jupyter', 'gsutil'
             # Need atari extras for Travis tests, but because gym is already in
             # install_requires, pip skips the atari extras, so we instead do an
             # explicit pip install gym[atari] for the tests.

From 55f301340e7eb2c8ea0f5239e23adab1a8c3de47 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 7 Aug 2018 09:29:37 -0700
Subject: [PATCH 0487/2720] Internal change

PiperOrigin-RevId: 207728093
---
 tensor2tensor/mesh_tensorflow/mesh_tensorflow.py     | 1 -
 tensor2tensor/mesh_tensorflow/mnist.py               | 2 +-
 tensor2tensor/mesh_tensorflow/mnist_dataset.py       | 1 -
 tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py   | 1 -
 tensor2tensor/mesh_tensorflow/placement_mesh_impl.py | 1 -
 tensor2tensor/mesh_tensorflow/simd_mesh_impl.py      | 2 --
 6 files changed, 1 insertion(+), 7 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index f83e0d28b..a389d4dbd 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -21,7 +21,6 @@
 from functools import reduce  # pylint: disable=redefined-builtin; for py3
 from operator import mul
 import re
-import google3
 from past.builtins import xrange
 import six
 
diff --git a/tensor2tensor/mesh_tensorflow/mnist.py b/tensor2tensor/mesh_tensorflow/mnist.py
index 0a0763872..31c261695 100644
--- a/tensor2tensor/mesh_tensorflow/mnist.py
+++ b/tensor2tensor/mesh_tensorflow/mnist.py
@@ -17,7 +17,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-import google3
+
 from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
 from tensor2tensor.mesh_tensorflow import mnist_dataset as dataset
 from tensor2tensor.mesh_tensorflow import mtf_layers
diff --git a/tensor2tensor/mesh_tensorflow/mnist_dataset.py b/tensor2tensor/mesh_tensorflow/mnist_dataset.py
index 99bb0e94c..15db689b6 100644
--- a/tensor2tensor/mesh_tensorflow/mnist_dataset.py
+++ b/tensor2tensor/mesh_tensorflow/mnist_dataset.py
@@ -36,7 +36,6 @@
 import shutil
 import tempfile
 
-import google3
 import numpy as np
 from six.moves import urllib
 import tensorflow as tf
diff --git a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
index 6c4725fce..8dfea0f7c 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
@@ -18,7 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import google3
 import numpy
 
 from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
diff --git a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
index 32fe65f69..f6a4d4a31 100644
--- a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
@@ -18,7 +18,6 @@
 from __future__ import print_function
 
 import functools
-import google3
 from past.builtins import xrange
 
 from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
diff --git a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
index 012527b6c..fbd9ae66c 100644
--- a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
@@ -18,8 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import google3
-
 from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
 from tensor2tensor.mesh_tensorflow import mtf_utils
 from tensor2tensor.mesh_tensorflow import tpu_variables

From 5f1fdbf9844254a9d1ce4d35d1af25038c00a79c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 7 Aug 2018 09:32:37 -0700
Subject: [PATCH 0488/2720] Small changes to decode_from_dataset.

PiperOrigin-RevId: 207728656
---
 tensor2tensor/utils/decoding.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 7f53dea3e..12eb1aa96 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -141,7 +141,8 @@ def decode_from_dataset(estimator,
                         hparams,
                         decode_hp,
                         decode_to_file=None,
-                        dataset_split=None):
+                        dataset_split=None,
+                        checkpoint_path=None):
   """Perform decoding from dataset."""
   tf.logging.info("Performing local inference from dataset for %s.",
                   str(problem_name))
@@ -185,7 +186,8 @@ def decode_from_dataset(estimator,
                          decode_hp,
                          decode_to_file,
                          output_dir,
-                         log_results=not decode_hp.decode_in_memory)
+                         log_results=not decode_hp.decode_in_memory,
+                         checkpoint_path=checkpoint_path)
 
     if decode_hp.decode_in_memory:
       output_dirs = [output_dir]
@@ -199,6 +201,7 @@ def decode_from_dataset(estimator,
       decode_hparams=decode_hp,
       predictions=predictions
   ))
+  return predictions
 
 
 def decode_once(estimator,
@@ -208,11 +211,13 @@ def decode_once(estimator,
                 decode_hp,
                 decode_to_file,
                 output_dir,
-                log_results=True):
+                log_results=True,
+                checkpoint_path=None):
   """Decodes once."""
 
   # Get the predictions as an iterable
-  predictions = estimator.predict(infer_input_fn)
+  predictions = estimator.predict(infer_input_fn,
+                                  checkpoint_path=checkpoint_path)
 
   if not log_results:
     return list(predictions)

From 6f4f9133e9fee011955810291bee8125a14ba055 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 7 Aug 2018 09:33:12 -0700
Subject: [PATCH 0489/2720] Fix argument passing to common_layers.conv1d

PiperOrigin-RevId: 207728760
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index d951b119b..6393de99b 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -3114,7 +3114,7 @@ def compute_attention_component(antecedent,
         antecedent, total_depth, use_bias=False, name=name)
   else:
     return common_layers.conv1d(
-        antecedent, total_depth, filter_width, padding, name=name)
+        antecedent, total_depth, filter_width, padding=padding, name=name)
 
 
 def compute_qkv(query_antecedent,

From 553e65aa896711e7c8475dce508a34cb1ccb64db Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 7 Aug 2018 11:38:10 -0700
Subject: [PATCH 0490/2720] Wrap t2t_datagen problem imports in try/except to
 support bare build target.

PiperOrigin-RevId: 207752289
---
 tensor2tensor/bin/t2t_datagen.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index 91ff65922..ec9fc2162 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -36,16 +36,22 @@
 import numpy as np
 
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
-from tensor2tensor.data_generators import algorithmic_math
-from tensor2tensor.data_generators import audio
 from tensor2tensor.data_generators import generator_utils
-from tensor2tensor.data_generators import snli
-from tensor2tensor.data_generators import wsj_parsing
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import usr_dir
 
 import tensorflow as tf
 
+try:
+  # pylint: disable=g-import-not-at-top
+  from tensor2tensor.data_generators import algorithmic_math
+  from tensor2tensor.data_generators import audio
+  from tensor2tensor.data_generators import snli
+  from tensor2tensor.data_generators import wsj_parsing
+  # pylint: enable=g-import-not-at-top
+except ImportError:
+  pass
+
 
 flags = tf.flags
 FLAGS = flags.FLAGS
@@ -145,10 +151,10 @@ def main(_):
     problems = []
 
   # Remove TIMIT if paths are not given.
-  if not FLAGS.timit_paths:
+  if getattr(FLAGS, "timit_paths", None):
     problems = [p for p in problems if "timit" not in p]
   # Remove parsing if paths are not given.
-  if not FLAGS.parsing_path:
+  if getattr(FLAGS, "parsing_path", None):
     problems = [p for p in problems if "parsing_english_ptb" not in p]
 
   if not problems:

From 77c3da7fd50b4eba665eb11fc4b7d0a4c34f4c55 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 7 Aug 2018 14:10:38 -0700
Subject: [PATCH 0491/2720] Add apache license to notebooks.

PiperOrigin-RevId: 207776977
---
 tensor2tensor/notebooks/asr_transformer.ipynb | 36 +++++++++++++++++++
 .../TransformerVisualization.ipynb            | 32 +++++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/tensor2tensor/notebooks/asr_transformer.ipynb b/tensor2tensor/notebooks/asr_transformer.ipynb
index 2328160eb..4e011b47f 100644
--- a/tensor2tensor/notebooks/asr_transformer.ipynb
+++ b/tensor2tensor/notebooks/asr_transformer.ipynb
@@ -1,5 +1,37 @@
 {
   "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "cellView": "form",
+        "colab": {
+          "autoexec": {
+            "startup": false,
+            "wait_interval": 0
+          }
+        },
+        "colab_type": "code",
+        "id": "6uNrFWq5BRba"
+      },
+      "outputs": [],
+      "source": [
+        "#@title\n",
+        "# Copyright 2018 Google LLC.\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",
+        "# https://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."
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": 0,
@@ -360,6 +392,10 @@
       "default_view": {},
       "name": "ASR with Transformer example notebook",
       "provenance": [
+        {
+          "file_id": "/piper/depot/google3/third_party/py/tensor2tensor/notebooks/asr_transformer.ipynb",
+          "timestamp": 1533672794744
+        },
         {
           "file_id": "notebooks/SR_with_Transformer_example_notebook.ipynb",
           "timestamp": 1525703542020
diff --git a/tensor2tensor/visualization/TransformerVisualization.ipynb b/tensor2tensor/visualization/TransformerVisualization.ipynb
index 91ae49ea1..dea6b3c6b 100644
--- a/tensor2tensor/visualization/TransformerVisualization.ipynb
+++ b/tensor2tensor/visualization/TransformerVisualization.ipynb
@@ -1,5 +1,37 @@
 {
  "cells": [
+  {
+    "cell_type": "code",
+    "execution_count": 0,
+    "metadata": {
+      "cellView": "form",
+      "colab": {
+        "autoexec": {
+          "startup": false,
+          "wait_interval": 0
+        }
+      },
+      "colab_type": "code",
+      "id": "6uNrFWq5BRba"
+    },
+    "outputs": [],
+    "source": [
+      "#@title\n",
+      "# Copyright 2018 Google LLC.\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",
+      "# https://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."
+    ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

From c1aed144c999c014fe3843655ab626f610e24f90 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Tue, 7 Aug 2018 14:46:31 -0700
Subject: [PATCH 0492/2720] Introducing pixel noise video modality.

PiperOrigin-RevId: 207783266
---
 tensor2tensor/layers/modalities.py            | 23 +++++++++++++++++++
 .../models/research/next_frame_params.py      |  9 ++++++++
 2 files changed, 32 insertions(+)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 64046e0ac..d404495ab 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -630,6 +630,29 @@ def targets_bottom(self, x):  # pylint: disable=arguments-differ
           name="merge_pixel_embedded_frames")
 
 
+@registry.register_video_modality("pixel_noise")
+class VideoModalityPixelNoise(VideoModality):
+  """Video modality that introduces pixel noise on input during training."""
+
+  def bottom(self, x):
+    inputs = x
+    if self._model_hparams.mode == tf.estimator.ModeKeys.TRAIN:
+      background = tf.contrib.distributions.percentile(inputs, 50.,
+                                                       axis=[0, 1, 2, 3])
+      input_shape = common_layers.shape_list(inputs)
+      input_size = tf.reduce_prod(input_shape[:-1])
+      input_mask = tf.multinomial(
+          tf.log([[self.input_noise, 1.-self.input_noise]]), input_size)
+      input_mask = tf.reshape(tf.cast(input_mask, tf.int32),
+                              input_shape[:-1]+[1])
+      inputs = inputs * input_mask + background * (1 - input_mask)
+    return super(VideoModalityPixelNoise, self).bottom(inputs)
+
+  @property
+  def input_noise(self):
+    return getattr(self._model_hparams, "video_modality_input_noise", 0.25)
+
+
 @registry.register_video_modality("l1")
 class VideoModalityL1(VideoModality):
   """Video modality that predicts a scalar per channel with an L1 loss."""
diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index 5b52ea944..85bb29d7f 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -49,6 +49,15 @@ def next_frame():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_pixel_noise():
+  """Basic 2-frame conv model with pixel noise."""
+  hparams = next_frame()
+  hparams.add_hparam("video_modality_input_noise", 0.25)
+  hparams.input_modalities = "inputs:video:pixel_noise"
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_stochastic():
   """SV2P model."""

From 0132483525b010577c5a809cff87b9d9905fd420 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 7 Aug 2018 15:31:19 -0700
Subject: [PATCH 0493/2720] Internal change

PiperOrigin-RevId: 207791391
---
 .travis.yml                               |   3 +
 tensor2tensor/notebooks/t2t_problem.ipynb | 570 ++++++++++++++++++++++
 2 files changed, 573 insertions(+)
 create mode 100644 tensor2tensor/notebooks/t2t_problem.ipynb

diff --git a/.travis.yml b/.travis.yml
index b162ae9f9..398c15703 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -107,6 +107,9 @@ script:
   - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
         jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/hello_t2t.ipynb;
     fi
+  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
+        jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/t2t_problem.ipynb;
+    fi
 
   # Export and query (on Python 2 only)
   # Bug: https://github.com/tensorflow/serving/issues/819
diff --git a/tensor2tensor/notebooks/t2t_problem.ipynb b/tensor2tensor/notebooks/t2t_problem.ipynb
new file mode 100644
index 000000000..1eddf9e6b
--- /dev/null
+++ b/tensor2tensor/notebooks/t2t_problem.ipynb
@@ -0,0 +1,570 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "Wd48fv-zDMe6"
+      },
+      "source": [
+        "# Welcome to the [Tensor2Tensor](https://github.com/tensorflow/tensor2tensor) Dataset Colab!\n",
+        "\n",
+        "Tensor2Tensor, or T2T for short, is a library of deep learning models and datasets designed to make deep learning more accessible and [accelerate ML research](https://research.googleblog.com/2017/06/accelerating-deep-learning-research.html).\n",
+        "\n",
+        "**This colab shows you how to add your own dataset to T2T so that you can train one of the several preexisting models on your newly added dataset!**\n",
+        "\n",
+        "For a tutorial that covers all the broader aspects of T2T using existing datasets and models, please see this [IPython notebook](https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/hello_t2t.ipynb)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "cellView": "form",
+        "colab": {
+          "autoexec": {
+            "startup": false,
+            "wait_interval": 0
+          }
+        },
+        "colab_type": "code",
+        "id": "FesA0dakI2kh"
+      },
+      "outputs": [],
+      "source": [
+        "#@title\n",
+        "# Copyright 2018 Google LLC.\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",
+        "# https://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."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "toc",
+        "id": "av8U13aqyEdf"
+      },
+      "source": [
+        "\u003e[Welcome to the Tensor2Tensor Dataset Colab!](#scrollTo=Wd48fv-zDMe6)\n",
+        "\n",
+        "\u003e\u003e[Installation \u0026 Setup](#scrollTo=Urn4QmNfI3hw)\n",
+        "\n",
+        "\u003e\u003e[Define the Problem](#scrollTo=LUoP57gOjlk9)\n",
+        "\n",
+        "\u003e\u003e\u003e[Run t2t_datagen](#scrollTo=Q1xBmlrFLSPX)\n",
+        "\n",
+        "\u003e\u003e[Viewing the generated data.](#scrollTo=MCqJhdnYgiG-)\n",
+        "\n",
+        "\u003e\u003e\u003e[tf.python_io.tf_record_iterator](#scrollTo=uNpohcPXKsLN)\n",
+        "\n",
+        "\u003e\u003e\u003e[Using tf.data.Dataset](#scrollTo=6o_1BHGQC5w5)\n",
+        "\n",
+        "\u003e\u003e[Terminology](#scrollTo=xRtfC0sHBlSo)\n",
+        "\n",
+        "\u003e\u003e\u003e[Problem](#scrollTo=xRtfC0sHBlSo)\n",
+        "\n",
+        "\u003e\u003e\u003e[Modalities](#scrollTo=xRtfC0sHBlSo)\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "Urn4QmNfI3hw"
+      },
+      "source": [
+        "## Installation \u0026 Setup\n",
+        "\n",
+        "\n",
+        "We'll install T2T and TensorFlow.\n",
+        "\n",
+        "We also need to setup the directories where T2T will:\n",
+        "\n",
+        "*   Generate the dataset and write the TFRecords file representing the training and the eval set, vocabulary files etc `DATA_DIR`\n",
+        "*   Run the training, keep the graph and the checkpoint files `OUTPUT_DIR` and\n",
+        "*   Use as a scratch directory to download your dataset from a URL, unzip it, etc. `TMP_DIR`"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "cellView": "form",
+        "colab": {
+          "autoexec": {
+            "startup": false,
+            "wait_interval": 0
+          }
+        },
+        "colab_type": "code",
+        "id": "IBWBeE39JYaR"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Run for installation.\n",
+        "\n",
+        "! pip install -q -U tensor2tensor\n",
+        "! pip install -q tensorflow"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "cellView": "form",
+        "colab": {
+          "autoexec": {
+            "startup": false,
+            "wait_interval": 0
+          }
+        },
+        "colab_type": "code",
+        "id": "sbTULiroLs2w"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Run this only once - Sets up TF Eager execution.\n",
+        "\n",
+        "import tensorflow as tf\n",
+        "\n",
+        "# Enable Eager execution - useful for seeing the generated data.\n",
+        "tf.enable_eager_execution()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "cellView": "form",
+        "colab": {
+          "autoexec": {
+            "startup": false,
+            "wait_interval": 0
+          }
+        },
+        "colab_type": "code",
+        "id": "A8JljOzDYF-Z"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Setting a random seed.\n",
+        "\n",
+        "from tensor2tensor.utils import trainer_lib\n",
+        "\n",
+        "# Set a seed so that we have deterministic outputs.\n",
+        "RANDOM_SEED = 301\n",
+        "trainer_lib.set_random_seed(RANDOM_SEED)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "cellView": "form",
+        "colab": {
+          "autoexec": {
+            "startup": false,
+            "wait_interval": 0
+          }
+        },
+        "colab_type": "code",
+        "id": "ioW-V1qpqSCE"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Run for setting up directories.\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "# Setup and create directories.\n",
+        "DATA_DIR = os.path.expanduser(\"/tmp/t2t/data\")\n",
+        "OUTPUT_DIR = os.path.expanduser(\"/tmp/t2t/output\")\n",
+        "TMP_DIR = os.path.expanduser(\"/tmp/t2t/tmp\")\n",
+        "\n",
+        "# Create them.\n",
+        "tf.gfile.MakeDirs(DATA_DIR)\n",
+        "tf.gfile.MakeDirs(OUTPUT_DIR)\n",
+        "tf.gfile.MakeDirs(TMP_DIR)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "LUoP57gOjlk9"
+      },
+      "source": [
+        "## Define the `Problem`\n",
+        "\n",
+        "To simplify our setting our input text sampled randomly from [a, z] - each sentence has between [3, 20] words with each word being [1, 8] characters in length.\n",
+        "\n",
+        "Example input: \"olrkpi z cldv xqcxisg cutzllf doteq\" -- this will be generated by `sample_sentence()`\n",
+        "\n",
+        "Our output will be the input words sorted according to length.\n",
+        "\n",
+        "Example output: \"z cldv doteq olrkpi xqcxisg cutzllf\" -- this will be processed by `target_sentence()`\n",
+        "\n",
+        "Let's dive right into our first problem -- we'll explain as we go on.\n",
+        "\n",
+        "Take some time to read each line along with its comments -- or skip them and come back later to clarify your understanding."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "cellView": "form",
+        "colab": {
+          "autoexec": {
+            "startup": false,
+            "wait_interval": 0
+          }
+        },
+        "colab_type": "code",
+        "id": "pDDiPxqg9UF-"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Define `sample_sentence()` and `target_sentence(input_sentence)`\n",
+        "import random\n",
+        "import string\n",
+        "\n",
+        "def sample_sentence():\n",
+        "    # Our sentence has between 3 and 20 words\n",
+        "    num_words = random.randint(3, 20)\n",
+        "    words = []\n",
+        "    for i in range(num_words):\n",
+        "        # Our words have between 1 and 8 characters.\n",
+        "        num_chars = random.randint(1, 8)\n",
+        "        chars = []\n",
+        "        for j in range(num_chars):\n",
+        "            chars.append(random.choice(string.ascii_lowercase))\n",
+        "        words.append(\"\".join(chars))\n",
+        "    return \" \".join(words)\n",
+        "\n",
+        "def target_sentence(input_sentence):\n",
+        "    words = input_sentence.split(\" \")\n",
+        "    return \" \".join(sorted(words, key=lambda x: len(x)))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {
+          "autoexec": {
+            "startup": false,
+            "wait_interval": 0
+          }
+        },
+        "colab_type": "code",
+        "id": "KcT_x4ma-Uaq"
+      },
+      "outputs": [],
+      "source": [
+        "# `Problem` is the base class for any dataset that we want to add to T2T -- it\n",
+        "# unifies the specification of the problem for generating training data,\n",
+        "# training, evaluation and inference.\n",
+        "#\n",
+        "# All its methods (except `generate_data`) have reasonable default\n",
+        "# implementations.\n",
+        "#\n",
+        "# A sub-class must implement `generate_data(data_dir, tmp_dir)` -- this method\n",
+        "# is called by t2t-trainer or t2t-datagen to actually generate TFRecord dataset\n",
+        "# files on disk.\n",
+        "from tensor2tensor.data_generators import problem\n",
+        "\n",
+        "# Certain categories of problems are very common, like where either the input or\n",
+        "# output is text, for such problems we define an (abstract) sub-class of\n",
+        "# `Problem` called `Text2TextProblem` -- this implements `generate_data` in\n",
+        "# terms of another function `generate_samples`. Sub-classes must override\n",
+        "# `generate_samples` and `is_generate_per_split`.\n",
+        "from tensor2tensor.data_generators import text_problems\n",
+        "\n",
+        "# Every non-abstract problem sub-class (as well as models and hyperparameter\n",
+        "# sets) must be registered with T2T so that T2T knows about it and can look it\n",
+        "# up when you specify your problem on the commandline to t2t-trainer or\n",
+        "# t2t-datagen.\n",
+        "#\n",
+        "# One uses:\n",
+        "# `register_problem` for a new Problem sub-class.\n",
+        "# `register_model` for a new T2TModel sub-class.\n",
+        "# `register_hparams` for a new hyperparameter set. All hyperparameter sets\n",
+        "# typically extend `common_hparams.basic_params1` (directly or indirectly).\n",
+        "from tensor2tensor.utils import registry\n",
+        "\n",
+        "\n",
+        "# By default, when you register a problem (or model or hyperparameter set) the\n",
+        "# name with which it gets registered is the 'snake case' version -- so here\n",
+        "# the Problem class `SortWordsAccordingToLengthRandom` will be registered with\n",
+        "# the name `sort_words_according_to_length_random`.\n",
+        "#\n",
+        "# One can override this default by actually assigning a name as follows:\n",
+        "# `@registry.register_problem(\"my_awesome_problem\")`\n",
+        "#\n",
+        "# The registered name is specified to the t2t-trainer or t2t-datagen using the\n",
+        "# commandline flag `--problem`.\n",
+        "@registry.register_problem\n",
+        "\n",
+        "# We inherit from `Text2TextProblem` which takes care of a lot of details\n",
+        "# regarding reading and writing the data to disk, what vocabulary type one\n",
+        "# should use, its size etc -- so that we need not worry about them, one can,\n",
+        "# of course, override those.\n",
+        "class SortWordsAccordingToLengthRandom(text_problems.Text2TextProblem):\n",
+        "  \"\"\"Sort words on length in randomly generated text.\"\"\"\n",
+        "\n",
+        "  # START: Methods we should override.\n",
+        "\n",
+        "  # The methods that need to be overriden from `Text2TextProblem` are:\n",
+        "  # `is_generate_per_split` and\n",
+        "  # `generate_samples`.\n",
+        "\n",
+        "  @property\n",
+        "  def is_generate_per_split(self):\n",
+        "    # If we have pre-existing data splits for (train, eval, test) then we set\n",
+        "    # this to True, which will have generate_samples be called for each of the\n",
+        "    # dataset_splits.\n",
+        "    #\n",
+        "    # If we do not have pre-existing data splits, we set this to False, which\n",
+        "    # will have generate_samples be called just once and the Problem will\n",
+        "    # automatically partition the data into dataset_splits.\n",
+        "    return False\n",
+        "\n",
+        "  def generate_samples(self, data_dir, tmp_dir, dataset_split):\n",
+        "    # Here we are generating the data in-situ using the `sample_sentence`\n",
+        "    # function, otherwise we would have downloaded the data and put it in\n",
+        "    # `tmp_dir` -- and read it from that location.\n",
+        "    del tmp_dir\n",
+        "\n",
+        "    # Unused here, is used in `Text2TextProblem.generate_data`.\n",
+        "    del data_dir\n",
+        "\n",
+        "    # This would have been useful if `self.is_generate_per_split()` was True.\n",
+        "    # In that case we would have checked if we were generating a training,\n",
+        "    # evaluation or test sample. This is of type `problem.DatasetSplit`.\n",
+        "    del dataset_split\n",
+        "\n",
+        "    # Just an arbitrary limit to our number of examples, this can be set higher.\n",
+        "    MAX_EXAMPLES = 10\n",
+        "\n",
+        "    for i in range(MAX_EXAMPLES):\n",
+        "      sentence_input = sample_sentence()\n",
+        "      sentence_target = target_sentence(sentence_input)\n",
+        "      yield {\n",
+        "          \"inputs\"  : sentence_input,\n",
+        "          \"targets\" : sentence_target,\n",
+        "      }\n",
+        "\n",
+        "  # END: Methods we should override.\n",
+        "\n",
+        "  # START: Overridable methods.\n",
+        "\n",
+        "  @property\n",
+        "  def vocab_type(self):\n",
+        "    # We can use different types of vocabularies, `VocabType.CHARACTER`,\n",
+        "    # `VocabType.SUBWORD` and `VocabType.TOKEN`.\n",
+        "    #\n",
+        "    # SUBWORD and CHARACTER are fully invertible -- but SUBWORD provides a good\n",
+        "    # tradeoff between CHARACTER and TOKEN.\n",
+        "    return text_problems.VocabType.SUBWORD\n",
+        "\n",
+        "  @property\n",
+        "  def approx_vocab_size(self):\n",
+        "    # Approximate vocab size to generate. Only for VocabType.SUBWORD.\n",
+        "    return 2**13  # ~8k\n",
+        "\n",
+        "  @property\n",
+        "  def dataset_splits(self):\n",
+        "    # Since we are responsible for generating the dataset splits, we override\n",
+        "    # `Text2TextProblem.dataset_splits` to specify that we intend to keep\n",
+        "    # 80% data for training and 10% for evaluation and testing each.\n",
+        "    return [{\n",
+        "        \"split\": problem.DatasetSplit.TRAIN,\n",
+        "        \"shards\": 8,\n",
+        "    }, {\n",
+        "        \"split\": problem.DatasetSplit.EVAL,\n",
+        "        \"shards\": 1,\n",
+        "    }, {\n",
+        "        \"split\": problem.DatasetSplit.TEST,\n",
+        "        \"shards\": 1,\n",
+        "    }]\n",
+        "\n",
+        " # END: Overridable methods."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "HwxQpOKhrolK"
+      },
+      "source": [
+        "That's it!\n",
+        "\n",
+        "To use this with `t2t-trainer` or `t2t-datagen`, save it to a directory, add an `__init__.py` that imports it, and then specify that directory with `--t2t_usr_dir`.\n",
+        "\n",
+        "i.e. as follows:\n",
+        "\n",
+        "```\n",
+        "$ t2t-datagen \\\n",
+        "  --problem=sort_words_according_to_length_random \\\n",
+        "  --data_dir=/tmp/t2t/data \\\n",
+        "  --tmp_dir=/tmp/t2t/tmp \\\n",
+        "  --t2t_usr_dir=/tmp/t2t/usr\n",
+        "\n",
+        "```\n",
+        "\n",
+        "However, we'll generate the data from the colab itself as well -- this is what `t2t-datagen` essentially does."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "Q1xBmlrFLSPX"
+      },
+      "source": [
+        "## Generate the data.\n",
+        "\n",
+        "We will now generate the data by calling `Problem.generate_data()` and inspect it."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {
+          "autoexec": {
+            "startup": false,
+            "wait_interval": 0
+          }
+        },
+        "colab_type": "code",
+        "id": "atYWRpM1FgaJ"
+      },
+      "outputs": [],
+      "source": [
+        "sort_len_problem = SortWordsAccordingToLengthRandom()\n",
+        "\n",
+        "sort_len_problem.generate_data(DATA_DIR, TMP_DIR)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "MCqJhdnYgiG-"
+      },
+      "source": [
+        "## Viewing the generated data.\n",
+        "\n",
+        "`tf.data.Dataset` is the recommended API for inputting data into a TensorFlow graph and the `Problem.dataset()` method returns a `tf.data.Dataset` object.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {
+          "autoexec": {
+            "startup": false,
+            "wait_interval": 0
+          }
+        },
+        "colab_type": "code",
+        "id": "PZczDWnOQDp2"
+      },
+      "outputs": [],
+      "source": [
+        "tfe = tf.contrib.eager\n",
+        "\n",
+        "Modes = tf.estimator.ModeKeys\n",
+        "\n",
+        "# We can iterate over our examples by making an iterator and calling next on it.\n",
+        "eager_iterator = tfe.Iterator(sort_len_problem.dataset(Modes.EVAL, DATA_DIR))\n",
+        "example = eager_iterator.next()\n",
+        "\n",
+        "input_tensor = example[\"inputs\"]\n",
+        "target_tensor = example[\"targets\"]\n",
+        "\n",
+        "# The tensors are actually encoded using the generated vocabulary file -- you\n",
+        "# can inspect the actual vocab file in DATA_DIR.\n",
+        "print(\"Tensor Input: \" + str(input_tensor))\n",
+        "print(\"Tensor Target: \" + str(target_tensor))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {
+          "autoexec": {
+            "startup": false,
+            "wait_interval": 0
+          }
+        },
+        "colab_type": "code",
+        "id": "1DtfzgqivAxl"
+      },
+      "outputs": [],
+      "source": [
+        "\n",
+        "# We use the encoders to decode the tensors to the actual input text.\n",
+        "input_encoder = sort_len_problem.get_feature_encoders(\n",
+        "    data_dir=DATA_DIR)[\"inputs\"]\n",
+        "target_encoder = sort_len_problem.get_feature_encoders(\n",
+        "    data_dir=DATA_DIR)[\"targets\"]\n",
+        "\n",
+        "input_decoded = input_encoder.decode(input_tensor.numpy())\n",
+        "target_decoded = target_encoder.decode(target_tensor.numpy())\n",
+        "\n",
+        "print(\"Decoded Input: \" + input_decoded)\n",
+        "print(\"Decoded Target: \" + target_decoded)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "xRtfC0sHBlSo"
+      },
+      "source": [
+        "## To be continued ...\n",
+        "\n",
+        "Stay tuned for additions to this notebook for adding problems with non-text modalities like Images, Audio and Video!"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [],
+      "default_view": {},
+      "name": "t2t_problem.ipynb",
+      "provenance": [
+        {
+          "file_id": "1FwspR4PzEZAiQCGziob5oov-8DyEXSnw",
+          "timestamp": 1533664607636
+        }
+      ],
+      "version": "0.3.2",
+      "views": {}
+    },
+    "kernelspec": {
+      "display_name": "Python 2",
+      "name": "python2"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

From 2674c371dcc2b42bebfd60bdda5eeadf260fc236 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Tue, 7 Aug 2018 16:31:22 -0700
Subject: [PATCH 0494/2720] Data mixing changes for MultiProblem.

PiperOrigin-RevId: 207801056
---
 .../data_generators/multi_problem.py          | 34 +++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index b22dbc8d2..b59a686cd 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -30,6 +30,8 @@
 class MultiProblem(problem.Problem):
   """MultiProblem base class."""
 
+  _ADDED_EVAL_COUNT = 20000
+
   def __init__(self, was_reversed=False, was_copy=False):
     super(MultiProblem, self).__init__(was_reversed, was_copy)
     self.task_list = []
@@ -138,15 +140,43 @@ def dataset(self,
 
       if is_training:
         task_dataset = task_dataset.repeat()
+
       # pylint: disable=cell-var-from-loop
       task_dataset = task_dataset.map(lambda x: self.add_task_id(task, x, enc))
+
+      if not is_training:
+        pad_data = tf.data.Dataset.from_tensor_slices({
+            "targets": tf.zeros([self._ADDED_EVAL_COUNT, 1], dtype=tf.int64),
+            "batch_prediction_key": tf.zeros(
+                [self._ADDED_EVAL_COUNT, 1], dtype=tf.int64),
+        })
+        task_dataset = task_dataset.concatenate(pad_data)
+
       datasets.append(task_dataset)
 
     # Setup the problem hparams by setting them to the LM task hparams.
     self.get_hparams()
 
-    single_mtl_dataset = tf.data.Dataset.zip(tuple(datasets)).flat_map(
-        self.flatten_zip)
+    if is_training:
+      dataset_iterators = [d.make_one_shot_iterator() for d in datasets]
+
+      def get_next_from_dataset(dataset_iter):
+        return dataset_iter.get_next()
+
+      def mix_data(example):
+        del example
+        return tf.data.Dataset.from_tensors(tf.cond(
+            tf.less(tf.random_uniform([]), 0.5),
+            lambda d=dataset_iterators[0]: get_next_from_dataset(d),
+            lambda d=dataset_iterators[1]: get_next_from_dataset(d)
+        ))
+
+      single_mtl_dataset = tf.data.Dataset.from_tensors(tf.zeros([1])).repeat()
+      single_mtl_dataset = single_mtl_dataset.flat_map(mix_data)
+
+    else:
+      single_mtl_dataset = tf.data.Dataset.zip(tuple(datasets)).flat_map(
+          self.flatten_zip)
 
     return single_mtl_dataset
 

From 2fd91d34b8e6d79599c0612e446175174e838b9d Mon Sep 17 00:00:00 2001
From: repoloper <41405759+repoloper@users.noreply.github.com>
Date: Tue, 7 Aug 2018 17:04:21 -0700
Subject: [PATCH 0495/2720] typo fixed from tf.logging.into to tf.logging.info

---
 tensor2tensor/utils/get_rouge.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/get_rouge.py b/tensor2tensor/utils/get_rouge.py
index 474bb2393..65dc883a6 100644
--- a/tensor2tensor/utils/get_rouge.py
+++ b/tensor2tensor/utils/get_rouge.py
@@ -46,7 +46,7 @@ def prep_data(decode_dir, target_dir):
         write_to_file(os.path.join(decode_dir, "rouge.%06d.txt" % (i+1)), d)
         write_to_file(os.path.join(target_dir, "rouge.A.%06d.txt" % (i+1)), t)
         if (i+1 % 1000) == 0:
-          tf.logging.into("Written %d examples to file" % i)
+          tf.logging.info("Written %d examples to file" % i)
 
 
 def main(_):

From 0c4a475c10bd14c05434bad2b2768bf3f4e4e699 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 7 Aug 2018 19:01:57 -0700
Subject: [PATCH 0496/2720] minor cleanup

PiperOrigin-RevId: 207819482
---
 tensor2tensor/notebooks/asr_transformer.ipynb | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/tensor2tensor/notebooks/asr_transformer.ipynb b/tensor2tensor/notebooks/asr_transformer.ipynb
index 4e011b47f..85593bf72 100644
--- a/tensor2tensor/notebooks/asr_transformer.ipynb
+++ b/tensor2tensor/notebooks/asr_transformer.ipynb
@@ -392,10 +392,6 @@
       "default_view": {},
       "name": "ASR with Transformer example notebook",
       "provenance": [
-        {
-          "file_id": "/piper/depot/google3/third_party/py/tensor2tensor/notebooks/asr_transformer.ipynb",
-          "timestamp": 1533672794744
-        },
         {
           "file_id": "notebooks/SR_with_Transformer_example_notebook.ipynb",
           "timestamp": 1525703542020

From ebbcd3b3ded83d6f1457f2f6808a23ec9b3552f1 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Tue, 7 Aug 2018 22:55:56 -0700
Subject: [PATCH 0497/2720] Mesh-Tensorflow - Changed a bunch of names and
 interfaces.  Rename TensorShape to Shape.  Mesh shapes now have named
 dimensions, and are instances of class Shape.  Rename "infeed", "outfeed" to
 "import_tf_tensor", "export_to_tf_tensor".  Delete unused models.  Add
 missing dropout in mtf_transformer, and update some configs.

PiperOrigin-RevId: 207836648
---
 .../mesh_tensorflow/mesh_tensorflow.py        | 521 +++++-----
 tensor2tensor/mesh_tensorflow/mnist.py        |  22 +-
 .../mesh_tensorflow/mtf_beam_search.py        |  61 +-
 .../mesh_tensorflow/mtf_image_transformer.py  |  18 +-
 .../mtf_image_transformer_test.py             |   9 +-
 tensor2tensor/mesh_tensorflow/mtf_layers.py   |  70 +-
 .../mesh_tensorflow/mtf_layers_test.py        |  68 +-
 tensor2tensor/mesh_tensorflow/mtf_model.py    |  26 +-
 tensor2tensor/mesh_tensorflow/mtf_optimize.py |   2 +-
 tensor2tensor/mesh_tensorflow/mtf_toy.py      | 178 ----
 .../mesh_tensorflow/mtf_toy_model_tpu.py      |  28 +-
 .../mesh_tensorflow/mtf_transformer.py        | 180 ++--
 .../mesh_tensorflow/mtf_transformer_compat.py | 926 ------------------
 .../mesh_tensorflow/mtf_transformer_test.py   |  27 +-
 .../mesh_tensorflow/placement_mesh_impl.py    |   8 +-
 .../mesh_tensorflow/simd_mesh_impl.py         |  16 +-
 16 files changed, 516 insertions(+), 1644 deletions(-)
 delete mode 100644 tensor2tensor/mesh_tensorflow/mtf_toy.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/mtf_transformer_compat.py

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index a389d4dbd..12f076534 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -45,21 +45,20 @@ def convert_to_dimension(d):
     return None
   if isinstance(d, Dimension):
     return d
-  tf.logging.info("d = %s" % d)
   name, size = d
   if isinstance(name, str) and isinstance(size, int):
     return Dimension(name, size)
   else:
-    raise ValueError("could not convert %s to Dimension" % d)
+    raise ValueError("could not convert %s to Dimension" % (d,))
 
 
-class TensorShape(object):
+class Shape(object):
   """Shape of a Tensor."""
 
   def __init__(self, dims):
     self._dims = tuple(dims)
     # verify no repeated dims
-    if len(set([d.name for d in dims])) != len(dims):
+    if len(set(dims)) != len(dims):
       raise ValueError("Shape must not have repeated dimensions %s" % dims)
 
   @property
@@ -80,20 +79,20 @@ def __ne__(self, other):
     return self.dims != other.dims
 
   def __add__(self, other):
-    if isinstance(other, TensorShape):
+    if isinstance(other, Shape):
       other = other.dims
     if isinstance(other, Dimension):
       other = [other]
-    return TensorShape(self.dims + other)
+    return Shape(self.dims + other)
 
   def __sub__(self, other):
     if other is None:
       return self
-    if isinstance(other, TensorShape):
+    if isinstance(other, Shape):
       other = other.dims
     if isinstance(other, Dimension):
       other = [other]
-    return TensorShape([d for d in self.dims if d not in other])
+    return Shape([d for d in self.dims if d not in other])
 
   def __len__(self):
     return len(self._dims)
@@ -105,22 +104,22 @@ def __iter__(self):
     return iter(self._dims)
 
   @property
-  def to_tf_shape(self):
+  def to_integer_list(self):
     return [d.size for d in self.dims]
 
   @property
   def size(self):
-    return list_product(self.to_tf_shape)
+    return list_product(self.to_integer_list)
 
   @property
   def to_string(self):
-    return "TensorShape[%s]" % ", ".join(
+    return "Shape[%s]" % ", ".join(
         ["%s=%d" % (d.name, d.size) for d in self.dims])
 
   @property
   def cumprod(self):
     """cumulative product (exclusive) of dimension sizes."""
-    return _cumprod(self.to_tf_shape)[::-1]
+    return _cumprod(self.to_integer_list)[::-1]
 
   def cumprod_to_tensor_axis(self, cumprod):
     """Tensor axis i such that self.cumprod[i] == cumprod, or None."""
@@ -138,7 +137,7 @@ def rename_dimension(self, old_name, new_name):
     if old_name not in self.dimension_names:
       raise ValueError("Shape %s does not have dimension named %s"
                        % (self, old_name))
-    return TensorShape(
+    return Shape(
         [Dimension(new_name, d.size) if d.name == old_name else d
          for d in self.dims])
 
@@ -147,30 +146,100 @@ def resize_dimension(self, name, new_size):
     if name not in self.dimension_names:
       raise ValueError("Shape %s does not have dimension named %s"
                        % (self, name))
-    return TensorShape(
+    return Shape(
         [Dimension(name, new_size) if d.name == name else d
          for d in self.dims])
 
 
-def convert_to_tensor_shape(x):
+def convert_to_shape(x):
   if x is None:
     return None
-  if isinstance(x, TensorShape):
+  if isinstance(x, Shape):
     return x
-  return TensorShape([convert_to_dimension(d) for d in x])
+  if isinstance(x, str):
+    x = _parse_string_to_list_of_pairs(x, seconds_to_int=True)
+  return Shape([convert_to_dimension(d) for d in x])
+
+
+class ComputationLayout(object):
+  """Represents layout of a computation.
+
+  Consists of a set of pairs of strings (tensor_dim_name, mesh_dim_name)
+  """
+
+  def __init__(self, pairs):
+    if isinstance(pairs, str):
+      pairs = _parse_string_to_list_of_pairs(pairs)
+    self._pairs = set(pairs)
+
+  def __repr__(self):
+    return "ComputationLayout%s" % self._pairs
+
+  def tensor_dimension_to_mesh_axis(self, tensor_dimension, mesh_shape):
+    """Mesh axis associated with tensor dimension (or None).
+
+    Args:
+      tensor_dimension: a Dimension
+      mesh_shape: a Shape
+    Returns:
+      an integer or None
+    Raises:
+      ValueError: if one Tensor dimension maps to two mesh dimensions.
+    """
+    val = [i for i, mesh_dimension in enumerate(mesh_shape)
+           if (tensor_dimension.name, mesh_dimension.name) in self._pairs]
+    if len(val) > 1:
+      raise ValueError(
+          "Tensor dimension maps to multiple mesh dimensions"
+          " tensor_dimension=%s mesh_shape=%s layout=%s"
+          % (tensor_dimension, mesh_shape, self._pairs))
+    return val[0] if val else None
+
+  def tensor_layout(self, tensor_shape, mesh_shape):
+    """Compute TensorLayout given a tensor shape and a mesh shape.
+
+    Args:
+      tensor_shape: a Shape
+      mesh_shape: a Shape
+    Returns:
+      a TensorLayout
+    Raises:
+      ValueError: if two tensor dimensions map to the same mesh dimension.
+    """
+    ret = [self.tensor_dimension_to_mesh_axis(d, mesh_shape)
+           for d in tensor_shape]
+    not_nones = [a for a in ret if a is not None]
+    if len(not_nones) != len(set(not_nones)):
+      raise ValueError(
+          "Two tensor dimensions may not map to the same mesh dimesnsion:"
+          " layout=%s tensor_shape=%s mesh_shape=%s " %
+          (self, tensor_shape, mesh_shape))
+    return TensorLayout(ret)
+
+
+def convert_to_computation_layout(x):
+  if isinstance(x, ComputationLayout):
+    return x
+  else:
+    return ComputationLayout(x)
 
 
 class TensorLayout(object):
   """Mapping from tensor dimension to mesh dimension.
 
-  Represented as a list of optional integers with length tensor.ndims.
-  Each item is either a unique integer inicating the mesh dimension over
+  Represented as a tuple of optional integers with length tensor.ndims.
+  Each item is either a unique integer inicating the mesh axis over
   which that tensor dimension is split, or None, indicating that this
   tensor dimension is not split.
   """
 
   def __init__(self, tensor_axis_to_mesh_axis):
-    self._tensor_axis_to_mesh_axis = tensor_axis_to_mesh_axis
+    """Create a TensorLayout.
+
+    Args:
+      tensor_axis_to_mesh_axis: a sequence of optional integers.
+    """
+    self._tensor_axis_to_mesh_axis = tuple(tensor_axis_to_mesh_axis)
 
   def __eq__(self, other):
     return self.tensor_axis_to_mesh_axis == other.tensor_axis_to_mesh_axis
@@ -179,16 +248,40 @@ def __ne__(self, other):
     return self.tensor_axis_to_mesh_axis != other.tensor_axis_to_mesh_axis
 
   def __repr__(self):
-    return "TensorLayout%s" % self.tensor_axis_to_mesh_axis
+    return "TensorLayout%s" % (self.tensor_axis_to_mesh_axis,)
+
+  def __len__(self):
+    return len(self._tensor_axis_to_mesh_axis)
+
+  def __getitem__(self, key):
+    return self._tensor_axis_to_mesh_axis[key]
+
+  def __iter__(self):
+    return iter(self._tensor_axis_to_mesh_axis)
 
   @property
   def tensor_axis_to_mesh_axis(self):
+    """Convert to a tuple of optional integers."""
     return self._tensor_axis_to_mesh_axis
 
   @property
   def is_fully_replicated(self):
-    return (self.tensor_axis_to_mesh_axis ==
-            [None] * len(self.tensor_axis_to_mesh_axis))
+    """Do all tensor dimensions map to None."""
+    return self.tensor_axis_to_mesh_axis == (None,) * len(self)
+
+  def mesh_axis_to_tensor_axis(self, mesh_ndims):
+    """For each mesh axis, which Tensor axis maps to it.
+
+    Args:
+      mesh_ndims: an integer
+
+    Returns:
+      a tuple of optional integers, with length mesh_ndims
+    """
+    return tuple(
+        [self._tensor_axis_to_mesh_axis.index(mesh_axis)
+         if mesh_axis in self._tensor_axis_to_mesh_axis else None
+         for mesh_axis in xrange(mesh_ndims)])
 
 
 class Graph(object):
@@ -262,7 +355,7 @@ def mesh_impl(self, m):
       m = m.mesh
     return self.mesh_to_impl[m]
 
-  def outfeed(self, x):
+  def export_to_tf_tensor(self, x):
     """Turn a Tensor into a tf.Tensor.
 
     Args:
@@ -271,7 +364,8 @@ def outfeed(self, x):
       a tf.Tensor
     """
     mesh_impl = self.mesh_impl(x)
-    return mesh_impl.outfeed(x, self.tensors[x].to_laid_out_tensor())
+    return mesh_impl.export_to_tf_tensor(
+        x, self.tensors[x].to_laid_out_tensor())
 
   def lowered_operation(self, op):
     return self.operations[op]
@@ -284,8 +378,8 @@ def copy_slices_to_masters(self):
     return tf.group(
         [v.copy_slices_to_master for _, v in six.iteritems(self.variables)])
 
-  def tensor_layout(self, t):
-    return self.mesh_impl(t).tensor_layout(t)
+  # def tensor_layout(self, t):
+  #   return self.mesh_impl(t).tensor_layout(t)
 
   def add_counter(self, key, value):
     assert isinstance(value, int)
@@ -339,10 +433,10 @@ def __init__(self, shape, layout):
 
     Args:
       shape: a list of ints
-      layout: dict from string to int
+      layout: a ComputationLayout
     """
-    self._shape = shape
-    self._layout = layout
+    self._shape = convert_to_shape(shape)
+    self._layout = convert_to_computation_layout(layout)
 
   @property
   def shape(self):
@@ -358,7 +452,7 @@ def layout(self):
 
   @property
   def size(self):
-    return list_product(self._shape)
+    return self.shape.size
 
   def tensor_dimension_to_mesh_axis(self, tensor_dimension):
     """Mesh axis associated with tensor dimension (or None).
@@ -368,44 +462,31 @@ def tensor_dimension_to_mesh_axis(self, tensor_dimension):
     Returns:
       an integer or None
     """
-    return self.layout.get(tensor_dimension.name)
+    return self.layout.tensor_dimension_to_mesh_axis(
+        tensor_dimension, self.shape)
 
   def tensor_layout(self, arg):
-    """Compute TensorLayout given a mesh and a TensorShape.
+    """Compute TensorLayout for a Tensor or a Shape.
 
     Args:
-      arg: a Tensor or TensorShape
+      arg: a Tensor or Shape
     Returns:
       a TensorLayout
     """
     if isinstance(arg, Tensor):
       arg = arg.shape
-    return TensorLayout(
-        [self.tensor_dimension_to_mesh_axis(d) for d in arg.dims])
-
-  def mesh_axis_to_tensor_axis(self, tensor_shape):
-    """Reverse-mapping of a tensor layout.
-
-    Args:
-      tensor_shape: a TensorShape
-    Returns:
-      a list of length self.ndims, where each element is either an integer
-        index of a tensor axis, or None
-    """
-    layout = self.tensor_layout(tensor_shape)
-    return [layout.tensor_axis_to_mesh_axis.index(mesh_axis)
-            if mesh_axis in layout.tensor_axis_to_mesh_axis else None
-            for mesh_axis in xrange(self.ndims)]
+    return self.layout.tensor_layout(arg, self.shape)
 
   def mesh_axis_to_cumprod(self, tensor_shape):
     """For each mesh axis, give the product of previous tensor axes.
 
     Args:
-      tensor_shape: a TensorShape
+      tensor_shape: a Shape
     Returns:
       a list with length self.ndims where each element is an integer or None.
     """
-    ma2ta = self.mesh_axis_to_tensor_axis(tensor_shape)
+    tensor_layout = self.tensor_layout(tensor_shape)
+    ma2ta = tensor_layout.mesh_axis_to_tensor_axis(self.ndims)
     ta2cumprod = tensor_shape.cumprod
     return [None if ta is None else ta2cumprod[ta] for ta in ma2ta]
 
@@ -413,7 +494,7 @@ def slice_shape(self, tensor_shape):
     """Shape of each slice of the tensor.
 
     Args:
-      tensor_shape: a TensorShape
+      tensor_shape: a Shape
     Returns:
       a list of integers with length tensor_shape.ndims
     Raises:
@@ -422,24 +503,25 @@ def slice_shape(self, tensor_shape):
     """
     tensor_layout = self.tensor_layout(tensor_shape)
     ret = []
-    for dim_size, mesh_dim in zip(
-        tensor_shape.to_tf_shape, tensor_layout.tensor_axis_to_mesh_axis):
-      if mesh_dim is None:
-        ret.append(dim_size)
+    for tensor_dim, mesh_axis in zip(
+        tensor_shape, tensor_layout.tensor_axis_to_mesh_axis):
+      if mesh_axis is None:
+        ret.append(tensor_dim.size)
       else:
-        if dim_size % self.shape[mesh_dim] != 0:
+        mesh_dim = self.shape[mesh_axis]
+        if tensor_dim.size % mesh_dim.size != 0:
           raise ValueError(
               "Tensor dimension size not divisible by mesh dimension size:"
               " tensor_shape=%s tensor_layout=%s"
               % (tensor_shape, tensor_layout))
-        ret.append(dim_size // self.shape[mesh_dim])
+        ret.append(tensor_dim.size // mesh_dim.size)
     return ret
 
   def slice_begin(self, tensor_shape, pnum):
     """Begin position for the tensor slice for the given processor.
 
     Args:
-      tensor_shape: a TensorShape
+      tensor_shape: a Shape
       pnum: an integer <= self.size
     Returns:
       a list of integers with length tensor_shape.ndims
@@ -448,18 +530,19 @@ def slice_begin(self, tensor_shape, pnum):
     coordinates = pnum_to_processor_coordinates(self.shape, pnum)
     ret = []
     for dim_size, mesh_axis in zip(
-        tensor_shape.to_tf_shape, tensor_layout.tensor_axis_to_mesh_axis):
+        tensor_shape.to_integer_list, tensor_layout.tensor_axis_to_mesh_axis):
       if mesh_axis is None:
         ret.append(0)
       else:
-        ret.append(dim_size // self.shape[mesh_axis] * coordinates[mesh_axis])
+        ret.append(
+            dim_size // self.shape[mesh_axis].size * coordinates[mesh_axis])
     return ret
 
   def laid_out_size(self, tensor_shape):
     """Total size of all slices.
 
     Args:
-      tensor_shape: a TensorShape
+      tensor_shape: a Shape
 
     Returns:
       an integer
@@ -518,7 +601,7 @@ def allsplit(self, x, mesh_axis, split_axis):
     Returns:
       a LaidOutTensor
     """
-    num_splits = self.shape[mesh_axis]
+    num_splits = self.shape[mesh_axis].size
     def my_fn(x, coordinate):
       slice_begin = [
           dimsize // num_splits * coordinate if i == split_axis
@@ -570,8 +653,8 @@ def laid_out_pcoord(self, mesh_axis):
     Returns:
       a LaidOutTensor where each slice is an integer scalar
     """
-    divisor = list_product(self.shape[mesh_axis + 1:])
-    modulus = self.shape[mesh_axis]
+    divisor = list_product(self.shape.to_integer_list[mesh_axis + 1:])
+    modulus = self.shape[mesh_axis].size
     def my_fn(pnum):
       return (pnum // divisor) % modulus
     return self.slicewise(my_fn, self.laid_out_pnum())
@@ -581,8 +664,8 @@ def broadcast_impl(self, old_slices, old_shape, new_shape):
 
     Args:
       old_slices: a LaidOutTensor
-      old_shape: a TensorShape
-      new_shape: a TensorShape
+      old_shape: a Shape
+      new_shape: a Shape
     Returns:
       a LaidOutTensor
     """
@@ -597,7 +680,7 @@ def make_slices(self, tf_tensor, tensor_shape):
 
     Args:
       tf_tensor: a tf.Tensor
-      tensor_shape: a TensorShape
+      tensor_shape: a Shape
 
     Returns:
       a list of tf.tensor with length self.size
@@ -619,7 +702,7 @@ def combine_slices(self, slices, tensor_shape, device=None):
 
     Args:
       slices: a list of tf.Tensor with length self.size
-      tensor_shape: a TensorShape
+      tensor_shape: a Shape
       device: an optional device string.
         if absent, we use the devices of the slices.
 
@@ -630,9 +713,10 @@ def combine_slices(self, slices, tensor_shape, device=None):
       return slices[0]
 
     ret = slices[:]
-    for unused_mesh_axis, (mesh_axis_size, tensor_axis) in enumerate(
-        zip(self.shape, self.mesh_axis_to_tensor_axis(tensor_shape))):
-      slice_size = len(ret) // mesh_axis_size
+    tensor_layout = self.tensor_layout(tensor_shape)
+    for mesh_dim, tensor_axis in zip(
+        self.shape, tensor_layout.mesh_axis_to_tensor_axis(self.ndims)):
+      slice_size = len(ret) // mesh_dim.size
       if tensor_axis is None:
         ret = ret[:slice_size]
       else:
@@ -641,7 +725,7 @@ def combine_slices(self, slices, tensor_shape, device=None):
         else:
           devices = [ret[i].device for i in xrange(slice_size)]
         concat_inputs = [[ret[i + slice_size * j]
-                          for j in xrange(mesh_axis_size)]
+                          for j in xrange(mesh_dim.size)]
                          for i in xrange(slice_size)]
         ret = parallel(
             devices, tf.concat, concat_inputs,
@@ -649,7 +733,7 @@ def combine_slices(self, slices, tensor_shape, device=None):
     assert len(ret) == 1
     return ret[0]
 
-  def outfeed(self, x, laid_out_x):
+  def export_to_tf_tensor(self, x, laid_out_x):
     """Turn a Tensor into a tf.Tensor.
 
     Args:
@@ -658,10 +742,10 @@ def outfeed(self, x, laid_out_x):
     Returns:
       a tf.Tensor
     """
-    raise NotImplementedError("Outfeed not implemented")
+    raise NotImplementedError("export_to_tf_tensor not implemented")
 
-  def infeed(self, x, tf_x):
-    """Infeed a tf.Tensor, producing a LaidOutTensor.
+  def import_tf_tensor(self, x, tf_x):
+    """Import a tf.Tensor, producing a LaidOutTensor.
 
     Args:
       x: a Tensor
@@ -669,7 +753,7 @@ def infeed(self, x, tf_x):
     Returns:
       a LaidOutTensor
     """
-    raise NotImplementedError("Infeed not implemented")
+    raise NotImplementedError("Import not implemented")
 
   @property
   def supports_control_dependencies(self):
@@ -756,8 +840,8 @@ class Tensor(object):
   """A Distributed Tensor."""
 
   def __init__(self, operation, shape, dtype, name=None):
-    if not isinstance(shape, TensorShape):
-      raise ValueError("shape must be a TensorShape got %s" % shape.to_string)
+    if not isinstance(shape, Shape):
+      raise ValueError("shape must be a Shape got %s" % shape.to_string)
     if not isinstance(dtype, tf.DType):
       raise ValueError("dtype must be a tf.DType got %s" % dtype)
     self._mesh = operation.mesh
@@ -934,7 +1018,7 @@ def _square_grad(op, dy):
     Args:
       tf_fn: a function taking n tf.Tensors and returning a tf.Tensor
       inputs: a list of n Tensors
-      output_shape: a TensorShape
+      output_shape: a Shape
       output_dtype: a dtype
       splittable_dims: a list of Dimensions which are ok to split
       grad_function: an optional python function
@@ -953,12 +1037,13 @@ def gradient(self, grad_ys):
 
   def lower(self, lowering):
     # Check that only splittable dims are split
+    mesh_impl = lowering.mesh_impl(self)
     for t in self.inputs + self.outputs:
-      layout = lowering.tensor_layout(t)
+      layout = mesh_impl.tensor_layout(t)
       for d, mesh_axis in zip(t.shape.dims, layout.tensor_axis_to_mesh_axis):
         if (mesh_axis is not None and d not in self._splittable_dims):
           raise ValueError("dimension %s is not declared as splittable" % d)
-    lowering.tensors[self.outputs[0]] = lowering.mesh_impl(self).slicewise(
+    lowering.tensors[self.outputs[0]] = mesh_impl.slicewise(
         self._tf_fn, *[lowering.tensors[x] for x in self.inputs])
 
 
@@ -978,7 +1063,7 @@ def slicewise(tf_fn,
   Args:
     tf_fn: a function taking n tf.Tensors and returning a tf.Tensor
     xs: a list of n Tensors
-    output_shape: a TensorShape
+    output_shape: a Shape
     output_dtype: a dtype
     splittable_dims: a list of Dimensions which are ok to split
     grad_function: an optional gradients function
@@ -990,7 +1075,7 @@ def slicewise(tf_fn,
   return SlicewiseOperation(
       tf_fn,
       xs,
-      convert_to_tensor_shape(output_shape) or xs[0].shape,
+      convert_to_shape(output_shape) or xs[0].shape,
       output_dtype or xs[0].dtype,
       splittable_dims,
       grad_function,
@@ -1207,11 +1292,11 @@ def binary_arguments_to_tensors(x1, x2):
   elif isinstance(x1, Tensor) and isinstance(x2, Tensor):
     return x1, x2
   elif isinstance(x1, Tensor):
-    return x1, infeed(x1.mesh, tf.convert_to_tensor(x2, dtype=x1.dtype),
-                      TensorShape([]))
+    return x1, import_tf_tensor(
+        x1.mesh, tf.convert_to_tensor(x2, dtype=x1.dtype), Shape([]))
   else:
-    return infeed(x2.mesh, tf.convert_to_tensor(x1, dtype=x2.dtype),
-                  TensorShape([])), x2
+    return import_tf_tensor(x2.mesh, tf.convert_to_tensor(x1, dtype=x2.dtype),
+                            Shape([])), x2
 
 
 def binary_op_with_broadcasting(
@@ -1221,7 +1306,7 @@ def binary_op_with_broadcasting(
   output_dtype = output_dtype or x1.dtype
   assert isinstance(output_dtype, tf.DType)
   return BinaryOpWithBroadcasting(
-      tf_fn, x1, x2, convert_to_tensor_shape(output_shape),
+      tf_fn, x1, x2, convert_to_shape(output_shape),
       output_dtype).outputs[0]
 
 
@@ -1315,14 +1400,14 @@ def broadcast(x, new_shape):
 
 def _reduce_helper(input_shape,
                    output_shape,
-                   mesh_layout,
+                   input_tensor_layout,
                    reduction_fn_string="SUM"):
   """Returns slicewise function and reduced mesh dimensions.
 
   Args:
-    input_shape: a TensorShape
-    output_shape: a TensorShape
-    mesh_layout: a dict (string -> int)
+    input_shape: a Shape
+    output_shape: a Shape
+    input_tensor_layout: a TensorLayout
     reduction_fn_string: "SUM" or "MAX"
   Returns:
     reduce_slice_fn: a function from tf.Tensor to tf.Tensor
@@ -1330,7 +1415,7 @@ def _reduce_helper(input_shape,
   """
   reduce_dims_indices = [
       i for i, d in enumerate(input_shape.dims) if d not in output_shape.dims]
-  reduced_input_shape = TensorShape([
+  reduced_input_shape = Shape([
       d for d in input_shape.dims if d in output_shape.dims])
   perm = [reduced_input_shape.dims.index(d) for d in output_shape.dims]
   def reduce_slice_fn(xslice):
@@ -1342,7 +1427,7 @@ def reduce_slice_fn(xslice):
     return ret
   reduced_mesh_axes = []
   for i in reduce_dims_indices:
-    mesh_axis = mesh_layout.get(input_shape.dims[i].name, None)
+    mesh_axis = input_tensor_layout[i]
     if mesh_axis is not None:
       reduced_mesh_axes.append(mesh_axis)
   return reduce_slice_fn, reduced_mesh_axes
@@ -1366,7 +1451,8 @@ def lower(self, lowering):
     mesh_impl = lowering.mesh_impl(self)
     slicewise_fn, reduced_mesh_axes = _reduce_helper(
         self.inputs[0].shape, self.outputs[0].shape,
-        mesh_impl.layout, self._reduction_fn_string)
+        mesh_impl.tensor_layout(self.inputs[0]),
+        self._reduction_fn_string)
     y = mesh_impl.slicewise(slicewise_fn, lowering.tensors[self.inputs[0]])
     if reduced_mesh_axes:
       def add_counter_fn():
@@ -1416,7 +1502,8 @@ def gradient(self, grad_ys):
 
   def lower(self, lowering):
     mesh_impl = lowering.mesh_impl(self)
-    if self._concat_dim_name in mesh_impl.layout:
+    if mesh_impl.tensor_dimension_to_mesh_axis(
+        Dimension(self._concat_dim_name, 0)) is not None:
       raise ValueError("can't concat along split axis")
     def slicewise_fn(*args):
       return tf.concat(args, axis=self._axis, name="concat")
@@ -1476,7 +1563,7 @@ def gradient(self, grad_ys):
 
   def lower(self, lowering):
     mesh_impl = lowering.mesh_impl(self)
-    if self._split_dim.name in mesh_impl.layout:
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._split_dim) is not None:
       raise ValueError("can't split along split axis")
     def slicewise_fn(x):
       # Since we return a tuple of tf.Tensor, slicewise will collate the
@@ -1515,7 +1602,7 @@ def __init__(self, xs, dim_name, axis, name=None):
       if x.shape != xs[0].shape:
         raise ValueError(
             "inputs to stack must have the same shape, got %s" % xs)
-    output_shape = TensorShape(
+    output_shape = Shape(
         input_shape.dims[:axis] + [self._new_dim]+ input_shape.dims[axis:])
     self._outputs = [Tensor(self, output_shape, xs[0].dtype)]
 
@@ -1524,7 +1611,7 @@ def gradient(self, grad_ys):
 
   def lower(self, lowering):
     mesh_impl = lowering.mesh_impl(self)
-    if self._new_dim in mesh_impl.layout:
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._new_dim) is not None:
       raise ValueError("can't stack along split axis")
     inputs = [lowering.tensors[t] for t in self._inputs]
     def slicewise_fn(*args):
@@ -1565,7 +1652,7 @@ def gradient(self, grad_ys):
 
   def lower(self, lowering):
     mesh_impl = lowering.mesh_impl(self)
-    if self._dim in mesh_impl.layout:
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._dim) is not None:
       raise ValueError("can't unstack along split axis")
     def slicewise_fn(x):
       return tuple(tf.unstack(x, num=self._dim.size, axis=self._axis))
@@ -1589,15 +1676,15 @@ def unstack(x, dim, name=None):
   return UnstackOperation(x, dim, name).outputs
 
 
-def _einsum_helper(input_shapes, output_shape, mesh_layout):
+def _einsum_helper(input_shapes, output_shape, mesh_impl):
   """Returns slicewise function and reduced mesh dimensions.
 
   Assumes the output shape contains no new dimensions.
 
   Args:
-    input_shapes: a list of TensorShapes
-    output_shape: a TensorShape
-    mesh_layout: a dict (string -> int)
+    input_shapes: a list of Shapes
+    output_shape: a Shape
+    mesh_impl: a MeshImpl
   Returns:
     einsum_slice_fn: a function from tf.Tensors to tf.Tensor
     reduced_mesh_axes: a list of integers
@@ -1608,9 +1695,9 @@ def _einsum_helper(input_shapes, output_shape, mesh_layout):
   full_shapes = [
       s for s in input_shapes + [output_shape] if s.ndims == total_num_dims]
   full_shape = (
-      full_shapes[0] if full_shapes else TensorShape(list(input_shape_set)))
+      full_shapes[0] if full_shapes else Shape(list(input_shape_set)))
   reduce_slice_fn, reduced_mesh_axes = _reduce_helper(
-      full_shape, output_shape, mesh_layout)
+      full_shape, output_shape, mesh_impl.tensor_layout(full_shape))
   def einsum_slice_fn_naive(*slices):
     # naive einsum implementation where we broadcst all inputs to the full
     # shape, multiply componentwise, then reduce.
@@ -1663,10 +1750,10 @@ def lower(self, lowering):
     xs = self.inputs
     input_shape_set = set(sum([x.shape.dims for x in xs], []))
     output_shape = self.outputs[0].shape
-    intersection_shape = TensorShape(
+    intersection_shape = Shape(
         [d for d in output_shape.dims if d in input_shape_set])
     einsum_slice_fn, reduced_mesh_axes = _einsum_helper(
-        [x.shape for x in self.inputs], intersection_shape, mesh_impl.layout)
+        [x.shape for x in self.inputs], intersection_shape, mesh_impl)
     y = mesh_impl.slicewise(
         einsum_slice_fn, *[lowering.tensors[x] for x in self.inputs])
     if reduced_mesh_axes:
@@ -1680,7 +1767,7 @@ def add_counter_fn():
     if intersection_shape != output_shape:
       y = mesh_impl.broadcast_impl(y, intersection_shape, output_shape)
     lowering.tensors[self.outputs[0]] = y
-    computation_shape = TensorShape(list(input_shape_set))
+    computation_shape = Shape(list(input_shape_set))
     lowering.add_counter("einsum", mesh_impl.laid_out_size(computation_shape))
     lowering.add_counter("einsum_unique", computation_shape.size)
 
@@ -1699,7 +1786,7 @@ def __init__(self, x, begin, size, slice_dim_name, name=None):
     self._begin = begin
     self._slice_dim = Dimension(slice_dim_name, size)
     input_shape = self._inputs[0].shape
-    output_shape = TensorShape(
+    output_shape = Shape(
         input_shape.dims[:axis] + [self._slice_dim] + input_shape.dims[axis+1:])
     self._outputs = [Tensor(self, output_shape, x.dtype)]
 
@@ -1712,13 +1799,13 @@ def gradient(self, grad_ys):
 
   def lower(self, lowering):
     mesh_impl = lowering.mesh_impl(self)
-    if self._slice_dim in mesh_impl.layout:
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._slice_dim) is not None:
       raise ValueError("can't slice along split axis")
     inputs = self._inputs[0]
     ndims = self._inputs[0].shape.ndims
     axis = self._axis
     begin = [0] * axis + [self._begin] + [0] * (ndims - axis - 1)
-    size = self._outputs[0].shape.to_tf_shape
+    size = self._outputs[0].shape.to_integer_list
 
     def slicewise_fn(x, begin, size):
       return tf.slice(x, begin, size, name="slice")
@@ -1746,7 +1833,7 @@ def __init__(self, x, paddings, pad_dim_name, name=None):
     self._axis = axis = dim_names.index(pad_dim_name)
     output_size = input_shape.dims[axis].size + sum(paddings)
     self._output_dim = Dimension(pad_dim_name, output_size)
-    output_shape = TensorShape(
+    output_shape = Shape(
         input_shape.dims[:axis] +
         [self._output_dim] + input_shape.dims[axis+1:])
     self._outputs = [Tensor(self, output_shape, x.dtype)]
@@ -1759,7 +1846,7 @@ def gradient(self, grad_ys):
 
   def lower(self, lowering):
     mesh_impl = lowering.mesh_impl(self)
-    if self._output_dim in mesh_impl.layout:
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._output_dim) is not None:
       raise ValueError("can't pad along split axis")
     inputs = self._inputs[0]
     ndims = self._inputs[0].shape.ndims
@@ -1786,7 +1873,7 @@ def __init__(self, indices, output_dim, on_value, off_value, dtype,
     self._on_value = on_value
     self._off_value = off_value
     self._dtype = dtype
-    output_shape = TensorShape(indices.shape.dims + [output_dim])
+    output_shape = Shape(indices.shape.dims + [output_dim])
     self._outputs = [Tensor(self, output_shape, dtype)]
 
   def lower(self, lowering):
@@ -1813,41 +1900,42 @@ def slicewise_fn(indices_slice, offset):
     lowering.tensors[self.outputs[0]] = y
 
 
-class InfeedOperation(Operation):
-  """Infeed a tf.Tensor onto a mesh."""
+class ImportOperation(Operation):
+  """Import a tf.Tensor onto a mesh."""
 
   def __init__(self, mesh, tf_tensor, shape, name=None):
-    super(InfeedOperation, self).__init__([], mesh=mesh, name=name or "infeed")
+    super(ImportOperation, self).__init__([], mesh=mesh, name=name or "import")
     self._outputs = [Tensor(self, shape, tf_tensor.dtype)]
     self._tf_tensor = tf_tensor
 
   def lower(self, lowering):
     mesh_impl = lowering.mesh_impl(self)
-    lowering.tensors[self.outputs[0]] = mesh_impl.infeed(
+    lowering.tensors[self.outputs[0]] = mesh_impl.import_tf_tensor(
         self.outputs[0], self._tf_tensor)
 
 
 def anonymous_shape(shape):
-  shape = convert_to_tensor_shape(shape)
-  return TensorShape([Dimension("_anonymous_%i" % i, d.size)
-                      for i, d in enumerate(shape)])
+  shape = convert_to_shape(shape)
+  return Shape([Dimension("_anonymous_%i" % i, d.size)
+                for i, d in enumerate(shape)])
 
 
 def anonymize(x):
   return reshape(x, anonymous_shape(x.shape))
 
 
-def infeed(mesh, tf_tensor, shape=None, name=None):
+def import_tf_tensor(mesh, tf_tensor, shape=None, name=None):
   tf_tensor = tf.convert_to_tensor(tf_tensor)
   if shape is None:
-    shape = TensorShape([])
+    shape = Shape([])
     assert not tf_tensor.shape.as_list()
-  return InfeedOperation(
-      mesh, tf_tensor, convert_to_tensor_shape(shape), name=name).outputs[0]
+  return ImportOperation(
+      mesh, tf_tensor, convert_to_shape(shape), name=name).outputs[0]
 
 
-def infeed_fully_replicated(mesh, tf_tensor, shape, name=None):
-  return reshape(infeed(mesh, tf_tensor, anonymous_shape(shape), name), shape)
+def import_fully_replicated(mesh, tf_tensor, shape, name=None):
+  return reshape(import_tf_tensor(
+      mesh, tf_tensor, anonymous_shape(shape), name), shape)
 
 
 class Variable(Operation):
@@ -1859,7 +1947,7 @@ def __init__(self, mesh, name, shape, dtype, initializer,
     self._trainable = trainable
     with tf.device("cpu:0"), mtf_utils.outside_all_rewrites():
       self.master = tf.get_variable(
-          name, shape.to_tf_shape, dtype=dtype, initializer=initializer,
+          name, shape.to_integer_list, dtype=dtype, initializer=initializer,
           **kwargs)
     self._name = self.master.name[:self.master.name.find(":")]
     self._outputs = [Tensor(self, shape, dtype)]
@@ -1895,7 +1983,7 @@ def get_variable(mesh, name, shape, dtype=tf.float32,
                  initializer=None, trainable=True,
                  activation_dtype=None, **kwargs):
   ret = Variable(
-      mesh, name, convert_to_tensor_shape(shape), dtype, initializer,
+      mesh, name, convert_to_shape(shape), dtype, initializer,
       trainable, **kwargs).outputs[0]
   if activation_dtype and activation_dtype != dtype:
     ret = cast(ret, activation_dtype)
@@ -1989,14 +2077,14 @@ def tf_fn():
 
 
 def constant(mesh, value, shape=None, dtype=tf.float32):
-  shape = convert_to_tensor_shape(shape)
+  shape = convert_to_shape(shape)
   return Constant(mesh, value,
-                  shape if shape is not None else TensorShape([]),
+                  shape if shape is not None else Shape([]),
                   dtype).outputs[0]
 
 
 def zeros(mesh, shape, dtype=tf.float32):
-  return constant(mesh, 0, shape=convert_to_tensor_shape(shape), dtype=dtype)
+  return constant(mesh, 0, shape=convert_to_shape(shape), dtype=dtype)
 
 
 def zeros_like(t):
@@ -2112,7 +2200,7 @@ def lower(self, lowering):
             "Try first reshaping to insert a new tf dimension,"
             " then changing layout.")
       slices = mesh_impl.allsplit(slices, mesh_axis, tensor_axis)
-      laid_out_size //= mesh_impl.shape[mesh_axis]
+      laid_out_size //= mesh_impl.shape[mesh_axis].size
     for mesh_axis in mesh_axes_alltoall:
       split_tensor_axis = old_shape.cumprod_to_tensor_axis(
           mesh_axis_to_cumprod_new[mesh_axis])
@@ -2134,7 +2222,7 @@ def lower(self, lowering):
           mesh_axis_to_cumprod_old[mesh_axis])
       assert tensor_axis is not None
       slices = mesh_impl.allconcat(slices, mesh_axis, tensor_axis)
-      laid_out_size *= mesh_impl.shape[mesh_axis]
+      laid_out_size *= mesh_impl.shape[mesh_axis].size
       lowering.add_counter(
           "allconcat/%s/reshape_op" % mesh_axis, laid_out_size)
     # now reshape the slices
@@ -2151,7 +2239,7 @@ def gradient(self, grad_ys):
 
 
 def reshape(x, new_shape):
-  return ReshapeOperation(x, convert_to_tensor_shape(new_shape)).outputs[0]
+  return ReshapeOperation(x, convert_to_shape(new_shape)).outputs[0]
 
 
 def rename_dimension(x, old_name, new_name):
@@ -2177,17 +2265,17 @@ def einsum(xs, output_shape=None, name=None):
 
   Args:
     xs: a list of Tensors
-    output_shape: an optional TensorShape.
+    output_shape: an optional Shape.
     name: an optional string
   Returns:
     a Tensor
   Raises:
     ValueError: if the output shape cannot be inferred
   """
-  output_shape = convert_to_tensor_shape(output_shape)
+  output_shape = convert_to_shape(output_shape)
   if output_shape is None:
     if len(xs) == 2:
-      output_shape = TensorShape(
+      output_shape = Shape(
           [d for d in xs[0].shape.dims if d not in xs[1].shape.dims] +
           [d for d in xs[1].shape.dims if d not in xs[0].shape.dims])
     else:
@@ -2204,7 +2292,7 @@ def _reduction_output_shape(x, output_shape, reduced_dim):
   """Helper function to reduce_sum, etc."""
   if output_shape is None:
     if reduced_dim is None:
-      return TensorShape([])
+      return Shape([])
     else:
       if reduced_dim not in x.shape.dims:
         raise ValueError(
@@ -2231,13 +2319,13 @@ def reduce_sum(x,
   Args:
     x: a Tensor
     disable_positional_args: None
-    output_shape: an optional TensorShape.  Must be a subsequence of x.shape.
+    output_shape: an optional Shape.  Must be a subsequence of x.shape.
     reduced_dim: a mtf.Dimension
     name: an optional string
   Returns:
     a Tensor
   """
-  output_shape = convert_to_tensor_shape(output_shape)
+  output_shape = convert_to_shape(output_shape)
   reduced_dim = convert_to_dimension(reduced_dim)
   assert disable_positional_args is None
   output_shape = _reduction_output_shape(x, output_shape, reduced_dim)
@@ -2261,14 +2349,14 @@ def reduce_mean(x,
   Args:
     x: a Tensor
     disable_positional_args: None
-    output_shape: an optional TensorShape. Must be a subsequence of x.shape.
+    output_shape: an optional Shape. Must be a subsequence of x.shape.
     reduced_dim: a mtf.Dimension
     name: an optional string
 
   Returns:
     a Tensor
   """
-  output_shape = convert_to_tensor_shape(output_shape)
+  output_shape = convert_to_shape(output_shape)
   reduced_dim = convert_to_dimension(reduced_dim)
   assert disable_positional_args is None
   output_shape = _reduction_output_shape(x, output_shape, reduced_dim)
@@ -2289,18 +2377,18 @@ def reduce_max(x,
   Args:
     x: a Tensor
     disable_positional_args: None
-    output_shape: an optional TensorShape.  Must be a subsequence of x.shape.
+    output_shape: an optional Shape.  Must be a subsequence of x.shape.
     reduced_dim: an optional Dimension
     name: an optional string
   Returns:
     a Tensor
   """
-  output_shape = convert_to_tensor_shape(output_shape)
+  output_shape = convert_to_shape(output_shape)
   reduced_dim = convert_to_dimension(reduced_dim)
   assert disable_positional_args is None
   output_shape = _reduction_output_shape(x, output_shape, reduced_dim)
   if output_shape is None:
-    output_shape = TensorShape([])
+    output_shape = Shape([])
   if output_shape == x.shape:
     return x
   return ReduceOperation(
@@ -2317,18 +2405,18 @@ def reduce_min(x,
   Args:
     x: a Tensor
     disable_positional_args: None
-    output_shape: an optional TensorShape.  Must be a subsequence of x.shape.
+    output_shape: an optional Shape.  Must be a subsequence of x.shape.
     reduced_dim: an optional Dimension
     name: an optional string
   Returns:
     a Tensor
   """
-  output_shape = convert_to_tensor_shape(output_shape)
+  output_shape = convert_to_shape(output_shape)
   reduced_dim = convert_to_dimension(reduced_dim)
   assert disable_positional_args is None
   output_shape = _reduction_output_shape(x, output_shape, reduced_dim)
   if output_shape is None:
-    output_shape = TensorShape([])
+    output_shape = Shape([])
   if output_shape == x.shape:
     return x
   return ReduceOperation(
@@ -2340,7 +2428,7 @@ def reduce_all(x,
                output_shape=None,
                reduced_dim=None,
                name=None):
-  output_shape = convert_to_tensor_shape(output_shape)
+  output_shape = convert_to_shape(output_shape)
   reduced_dim = convert_to_dimension(reduced_dim)
   return cast(reduce_min(to_float(x),
                          disable_positional_args=disable_positional_args,
@@ -2354,7 +2442,7 @@ def reduce_any(x,
                output_shape=None,
                reduced_dim=None,
                name=None):
-  output_shape = convert_to_tensor_shape(output_shape)
+  output_shape = convert_to_shape(output_shape)
   reduced_dim = convert_to_dimension(reduced_dim)
   return cast(reduce_max(to_float(x),
                          disable_positional_args=disable_positional_args,
@@ -2439,12 +2527,12 @@ def add(x1, x2, output_shape=None, name=None):
   Args:
     x1: a Tensor
     x2: a Tensor
-    output_shape: an optional TensorShape
+    output_shape: an optional Shape
     name: an optional string
   Returns:
     a Tensor
   """
-  output_shape = convert_to_tensor_shape(output_shape)
+  output_shape = convert_to_shape(output_shape)
   if not isinstance(x2, Tensor):
     return ScalarAddOperation(x1, x2).outputs[0]
   with tf.name_scope(name, default_name="add"):
@@ -2460,12 +2548,12 @@ def sub(x1, x2, output_shape=None, name=None):
   Args:
     x1: a Tensor
     x2: a Tensor
-    output_shape: an optional TensorShape
+    output_shape: an optional Shape
     name: an optional string
   Returns:
     a Tensor
   """
-  output_shape = convert_to_tensor_shape(output_shape)
+  output_shape = convert_to_shape(output_shape)
   if not isinstance(x2, Tensor):
     return ScalarAddOperation(x1, -x2).outputs[0]
   with tf.name_scope(name, default_name="sub"):
@@ -2479,7 +2567,7 @@ def multiply(x1, x2, output_shape=None, name=None):
   Args:
     x1: a Tensor
     x2: a Tensor
-    output_shape: an optional TensorShape
+    output_shape: an optional Shape
     name: an optional string
   Returns:
     a Tensor
@@ -2500,12 +2588,12 @@ def divide(x1, x2, output_shape=None, name=None):
   Args:
     x1: a Tensor
     x2: a Tensor
-    output_shape: an optional TensorShape
+    output_shape: an optional Shape
     name: an optional string
   Returns:
     a Tensor
   """
-  output_shape = convert_to_tensor_shape(output_shape)
+  output_shape = convert_to_shape(output_shape)
   if not isinstance(x2, Tensor):
     return ScalarMultiplyOperation(x1, 1.0 / x2).outputs[0]
   with tf.name_scope(name, default_name="divide"):
@@ -2569,12 +2657,12 @@ def gather(weights, indices, dim, output_shape=None):
     weights: a Tensor
     indices: a Tensor with integer type
     dim: a Dimension
-    output_shape: an optional mtf.TensorShape
+    output_shape: an optional mtf.Shape
   Returns:
     a Tensor
   """
   dim = convert_to_dimension(dim)
-  output_shape = convert_to_tensor_shape(output_shape)
+  output_shape = convert_to_shape(output_shape)
   if weights.dtype == tf.bool:
     return cast(gather(to_float(weights), indices, dim, output_shape), tf.bool)
   return einsum([one_hot(indices, dim, dtype=weights.dtype), weights],
@@ -2625,22 +2713,22 @@ def _infer_binary_broadcast_shape(shape1, shape2, given_output_shape=None):
   of shape1, followed by all new dimensions in shape2.
 
   Args:
-    shape1: a TensorShape
-    shape2: a TensorShape
-    given_output_shape: an optional TensorShape
+    shape1: a Shape
+    shape2: a Shape
+    given_output_shape: an optional Shape
   Returns:
-    a TensorShape
+    a Shape
   """
-  shape1 = convert_to_tensor_shape(shape1)
-  shape2 = convert_to_tensor_shape(shape2)
-  given_output_shape = convert_to_tensor_shape(given_output_shape)
+  shape1 = convert_to_shape(shape1)
+  shape2 = convert_to_shape(shape2)
+  given_output_shape = convert_to_shape(given_output_shape)
   if given_output_shape is not None:
     return given_output_shape
   if is_subsequence(shape1.dims, shape2.dims):
     return shape2
   if is_subsequence(shape2.dims, shape1.dims):
     return shape1
-  return TensorShape(
+  return Shape(
       shape1.dims + [d for d in shape2.dims if d not in shape1.dims])
 
 
@@ -2649,8 +2737,8 @@ def _expand_dims(x, input_shape, output_shape):
 
   Args:
     x: a tf.Tensor
-    input_shape: a TensorShape
-    output_shape: a TensorShape whose dimensions are a superset of
+    input_shape: a Shape
+    output_shape: a Shape whose dimensions are a superset of
       those in input_shape
 
   Returns:
@@ -2674,8 +2762,8 @@ def _einsum_equation(input_shapes, output_shape):
   e.g. "ij,jk->ik"
 
   Args:
-    input_shapes: a list of TensorShapes
-    output_shape: a TensorShape
+    input_shapes: a list of Shapes
+    output_shape: a Shape
   Returns:
     a string
   """
@@ -2715,8 +2803,8 @@ def verify_no_new_dims(input_shapes, output_shape):
   """Verifies that all dimensions in the output are in at least one input.
 
   Args:
-    input_shapes: a list of TensorShapes
-    output_shape: a TensorShape
+    input_shapes: a list of Shapes
+    output_shape: a Shape
   Raises:
     ValueError: if there are new dimensions in the output.
   """
@@ -2733,14 +2821,14 @@ def pnum_to_processor_coordinates(mesh_shape, pnum):
   """Coordinates of a processor in the mesh.
 
   Args:
-    mesh_shape: a list of integers
+    mesh_shape: a Shape
     pnum: an integer less than len(mesh_shape)
 
   Returns:
     a list of integers with length len(mesh_shape)
   """
   ret = []
-  for dimsize in mesh_shape[::-1]:
+  for dimsize in mesh_shape.to_integer_list[::-1]:
     ret.append(pnum % dimsize)
     pnum //= dimsize
   return ret[::-1]
@@ -2750,7 +2838,7 @@ def processor_coordinates_to_pnum(mesh_shape, coord):
   """Inverse of pnum_to_processor_coordinates.
 
   Args:
-    mesh_shape: a list of integers
+    mesh_shape: a Shape
     coord: a list of integers with length len(mesh_shape)
 
   Returns:
@@ -2758,7 +2846,7 @@ def processor_coordinates_to_pnum(mesh_shape, coord):
   """
   ret = 0
   multiplier = 1
-  for c, d in zip(coord[::-1], mesh_shape[::-1]):
+  for c, d in zip(coord[::-1], mesh_shape.to_integer_list[::-1]):
     ret += multiplier * c
     multiplier *= d
   return ret
@@ -2768,7 +2856,7 @@ def pnum_to_group(mesh_shape, group_dims, pnum):
   """Group number for grouped allreduce.
 
   Args:
-    mesh_shape: a list of integers
+    mesh_shape: a Shape
     group_dims: a list of integers (the dimensions reduced over)
     pnum: an integer
 
@@ -2776,7 +2864,8 @@ def pnum_to_group(mesh_shape, group_dims, pnum):
     an integer
   """
   coord = pnum_to_processor_coordinates(mesh_shape, pnum)
-  remaining_shape = [d for i, d in enumerate(mesh_shape) if i not in group_dims]
+  remaining_shape = Shape(
+      [d for i, d in enumerate(mesh_shape) if i not in group_dims])
   remaining_coord = [d for i, d in enumerate(coord) if i not in group_dims]
   return processor_coordinates_to_pnum(remaining_shape, remaining_coord)
 
@@ -2785,7 +2874,7 @@ def processor_groups(mesh_shape, group_dims):
   """Groups of processors which differ only in the given dimensions.
 
   Args:
-    mesh_shape: a list of integers
+    mesh_shape: a Shape
     group_dims: a list of integers
 
   Returns:
@@ -2793,7 +2882,7 @@ def processor_groups(mesh_shape, group_dims):
   """
   group_numbers = [
       pnum_to_group(mesh_shape, group_dims, pnum)
-      for pnum in xrange(list_product(mesh_shape))]
+      for pnum in xrange(mesh_shape.size)]
   ret = []
   for pnum, g in enumerate(group_numbers):
     while len(ret) <= g:
@@ -2847,8 +2936,8 @@ def range(mesh, dim, dtype, name=None):  # pylint: disable=redefined-builtin
   """
   dim = convert_to_dimension(dim)
   with tf.variable_scope(name, default_name="range"):
-    return infeed(
-        mesh, tf.range(dim.size, dtype=dtype), shape=TensorShape([dim]))
+    return import_tf_tensor(
+        mesh, tf.range(dim.size, dtype=dtype), shape=Shape([dim]))
 
 
 def pretty_print_counters(counters):
@@ -2876,41 +2965,33 @@ def pretty_print_counters(counters):
   return "\n".join(parts)
 
 
-def parse_mesh_shape(mesh_shape):
-  """Parase a string to a list of integers.
-
-  All non-digits are taken as delimeters
-
-  Args:
-    mesh_shape: a string or a list of integers
-  Returns:
-    a list of integers
-  """
-  if isinstance(mesh_shape, list):
-    return mesh_shape
-  return [int(x) for x in re.sub("[^0-9]", " ", mesh_shape).split()]
+def _parse_string_to_list_of_pairs(s, seconds_to_int=False):
+  r"""Parase a string into a list of pairs.
 
+  If seconds_to_int, then the second elements are integers, otherwise, they
+  are strings.
 
-def parse_layout(layout_string):
-  r"""Parase a string specifying a layout.
+  In the input string, each pair is separated by a colon, and the delimeters
+  between paris are any of " ,.;"
 
-  The layout_string is a list of name, integer pairs.
-  Each pair is separated by a colon, and the delimeters between paris are any
-  of " ,.;"
-
-  e.g. "batch:0 vocab:1 filter_size:1 heads:1"
+  e.g. "rows:32,cols:32"
 
   Args:
-    layout_string: a string or a dictionary
+    s: a string
+    seconds_to_int: a boolean
   Returns:
-    a dictionary from string to int
+    a list of pairs
+  Raises:
+    ValueError: on badly formatted string
   """
-  if isinstance(layout_string, dict):
-    return layout_string
-  ret = {}
-  for s in re.sub("[,.;]", " ", layout_string).split():
-    dim_name, mesh_axis = s.split(":")
-    ret[dim_name] = int(mesh_axis)
+  ret = []
+  for p in [s.split(":") for s in re.sub("[,.;]", " ", s).split()]:
+    if len(p) != 2:
+      raise ValueError("bad input to _parse_string_to_list_of_pairs %s" % s)
+    if seconds_to_int:
+      ret.append((p[0], int(p[1])))
+    else:
+      ret.append(tuple(p))
   return ret
 
 
@@ -3039,13 +3120,13 @@ def random_uniform(mesh, shape, **kwargs):
 
   Args:
     mesh: a Mesh
-    shape: a TensorShape
+    shape: a Shape
     **kwargs: keyword args for tf.random_uniform, except seed
 
   Returns:
     a Tensor
   """
-  shape = convert_to_tensor_shape(shape)
+  shape = convert_to_shape(shape)
   return RandomOperation(mesh, shape, tf.random_uniform, **kwargs).outputs[0]
 
 
@@ -3055,13 +3136,13 @@ def dropout(x, keep_prob, noise_shape=None, name=None):
   Args:
     x: a Tensor
     keep_prob: a float between 0.0 and 1.0
-    noise_shape: an optional TensorShape (a subset of x.shape)
+    noise_shape: an optional Shape (a subset of x.shape)
     name: an optional string
 
   Returns:
     a Tensor
   """
-  noise_shape = convert_to_tensor_shape(noise_shape)
+  noise_shape = convert_to_shape(noise_shape)
   if noise_shape is None:
     noise_shape = x.shape
   with tf.variable_scope(name, default_name="dropout"):
diff --git a/tensor2tensor/mesh_tensorflow/mnist.py b/tensor2tensor/mesh_tensorflow/mnist.py
index 31c261695..60f81fe1b 100644
--- a/tensor2tensor/mesh_tensorflow/mnist.py
+++ b/tensor2tensor/mesh_tensorflow/mnist.py
@@ -39,8 +39,9 @@
 tf.flags.DEFINE_integer("eval_steps", 0,
                         "Total number of evaluation steps. If `0`, evaluation "
                         "after training is skipped.")
-tf.flags.DEFINE_string("mesh_shape", "2;2", "mesh shape")
-tf.flags.DEFINE_string("layout", "batch:0;hidden1:1", "computation layout")
+tf.flags.DEFINE_string("mesh_shape", "rows:2;cols:2", "mesh shape")
+tf.flags.DEFINE_string("layout", "batch:rows;hidden1:cols",
+                       "computation layout")
 
 FLAGS = tf.flags.FLAGS
 
@@ -64,8 +65,8 @@ def mnist_model(image, labels, mesh):
   hidden_dim1 = mtf.Dimension("hidden1", FLAGS.hidden_size)
   hidden_dim2 = mtf.Dimension("hidden2", FLAGS.hidden_size)
 
-  x = mtf.infeed(mesh, tf.reshape(image, [-1, 28, 28]),
-                 mtf.TensorShape([batch_dim, rows_dim, cols_dim]))
+  x = mtf.import_tf_tensor(mesh, tf.reshape(image, [-1, 28, 28]),
+                           mtf.Shape([batch_dim, rows_dim, cols_dim]))
   h1 = mtf_layers.dense(
       x, hidden_dim1, reduced_dims=[rows_dim, cols_dim],
       activation=mtf.relu, name="hidden1")
@@ -75,7 +76,7 @@ def mnist_model(image, labels, mesh):
   if labels is None:
     loss = None
   else:
-    labels = mtf.infeed(mesh, labels, mtf.TensorShape([batch_dim]))
+    labels = mtf.import_tf_tensor(mesh, labels, mtf.Shape([batch_dim]))
     loss = mtf_layers.softmax_cross_entropy_with_logits(
         logits, mtf.one_hot(labels, classes_dim), classes_dim)
     loss = mtf.reduce_mean(loss)
@@ -90,11 +91,12 @@ def model_fn(features, labels, mode, params):
   graph = mtf.Graph()
   mesh = mtf.Mesh(graph, "my_mesh")
   logits, loss = mnist_model(features, labels, mesh)
-  mesh_shape = mtf.parse_mesh_shape(FLAGS.mesh_shape)
-  mesh_size = mtf.list_product(mesh_shape)
+  mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
+  computation_layout = mtf.ComputationLayout(FLAGS.layout)
+  mesh_size = mesh_shape.size
   mesh_devices = [""] * mesh_size
   mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-      mesh_shape, mtf.parse_layout(FLAGS.layout), mesh_devices)
+      mesh_shape, computation_layout, mesh_devices)
 
   if mode == tf.estimator.ModeKeys.TRAIN:
     var_grads = mtf.gradients(
@@ -107,9 +109,9 @@ def model_fn(features, labels, mode, params):
   lowering = mtf.Lowering(graph, {mesh: mesh_impl})
   restore_hook = mtf.MtfRestoreHook(lowering)
 
-  tf_logits = lowering.outfeed(logits)
+  tf_logits = lowering.export_to_tf_tensor(logits)
   if mode != tf.estimator.ModeKeys.PREDICT:
-    tf_loss = lowering.outfeed(loss)
+    tf_loss = lowering.export_to_tf_tensor(loss)
     tf.summary.scalar("loss", tf_loss)
 
   if mode == tf.estimator.ModeKeys.TRAIN:
diff --git a/tensor2tensor/mesh_tensorflow/mtf_beam_search.py b/tensor2tensor/mesh_tensorflow/mtf_beam_search.py
index f0f42b7a0..3fb8592ba 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_beam_search.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_beam_search.py
@@ -27,59 +27,6 @@
 INF = 1. * 1e7
 
 
-def _concat_equal_sizes(xs, dim, new_dim_name):
-  axis = xs[0].shape.dims.index(dim)
-  ret = mtf.stack(xs, "tmp_concat", axis)
-  new_shape = mtf.TensorShape(
-      xs[0].shape.dims[:axis]
-      + [mtf.Dimension(new_dim_name, dim.size * len(xs))]
-      + xs[0].shape.dims[axis + 1:])
-  return mtf.reshape(ret, new_shape)
-
-
-def _expand_to_beam_size(tensor, beam_size):
-  """Tiles a given tensor by beam_size.
-
-  Args:
-    tensor: tensor to tile [batch_size, ...]
-    beam_size: How much to tile the tensor by.
-
-  Returns:
-    Tiled tensor [batch_size, beam_size, ...]
-  """
-  tensor = tf.expand_dims(tensor, axis=1)
-  tile_dims = [1] * tensor.shape.ndims
-  tile_dims[1] = beam_size
-
-  return tf.tile(tensor, tile_dims)
-
-
-def get_state_shape_invariants(tensor):
-  """Returns the shape of the tensor but sets middle dims to None."""
-  shape = tensor.shape.as_list()
-  for i in range(1, len(shape) - 1):
-    shape[i] = None
-  return tf.TensorShape(shape)
-
-
-def compute_batch_indices(batch_size, beam_size):
-  """Computes the i'th coordinate that contains the batch index for gathers.
-
-  Batch pos is a tensor like [[0,0,0,0,],[1,1,1,1],..]. It says which
-  batch the beam item is in. This will create the i of the i,j coordinate
-  needed for the gather.
-
-  Args:
-    batch_size: Batch size
-    beam_size: Size of the beam.
-  Returns:
-    batch_pos: [batch_size, beam_size] tensor of ids
-  """
-  batch_pos = tf.range(batch_size * beam_size) // beam_size
-  batch_pos = tf.reshape(batch_pos, [batch_size, beam_size])
-  return batch_pos
-
-
 def compute_topk_scores_and_seq(sequences, scores, scores_to_gather, flags,
                                 beam_dim, prefix="default",
                                 states=None):
@@ -126,7 +73,7 @@ def compute_topk_scores_and_seq(sequences, scores, scores_to_gather, flags,
   # Clients can capture these tensors by watching these node names.
   def gather(tensor, name):
     with tf.name_scope(prefix + name):
-      output_shape = mtf.TensorShape(
+      output_shape = mtf.Shape(
           [beam_dim if d == old_beam_dim else d for d in tensor.shape.dims])
       return mtf.gather(
           tensor, topk_indices, old_beam_dim, output_shape=output_shape)
@@ -196,7 +143,7 @@ def beam_search(logits_fn,
   batch_dim, beam_dim, length_dim = initial_ids.shape.dims
   mesh = initial_ids.mesh
 
-  batch_by_beam = mtf.TensorShape([batch_dim, beam_dim])
+  batch_by_beam = mtf.Shape([batch_dim, beam_dim])
   initial_log_probs = mtf.broadcast(
       mtf.one_hot(
           mtf.constant(mesh, 0, dtype=tf.int32),
@@ -332,7 +279,7 @@ def grow_topk(i, alive_seq, alive_log_probs, states=None):
     # scores have shape [batch, beam, vocab]
     beam_and_vocab_dim = mtf.Dimension(
         "beam_and_vocab", beam_dim.size * vocab_dim.size)
-    flat_shape = mtf.TensorShape([batch_dim, beam_and_vocab_dim])
+    flat_shape = mtf.Shape([batch_dim, beam_and_vocab_dim])
     double_beam = mtf.Dimension("double_beam", beam_dim.size * 2)
     # Flatten out (beam_size, vocab_size) probs in to a list of possibilities
     flat_curr_scores = mtf.reshape(curr_scores, flat_shape)
@@ -350,7 +297,7 @@ def grow_topk(i, alive_seq, alive_log_probs, states=None):
     def my_gather(tensor):
       return mtf.gather(
           tensor, top_beam_index, beam_dim,
-          output_shape=mtf.TensorShape(
+          output_shape=mtf.Shape(
               [double_beam if d == beam_dim else d for d in tensor.shape.dims]))
 
     # Gather up the most probable 2*beams both for the ids and finished_in_alive
diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
index 736ce646d..0b80c4c7a 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
@@ -73,17 +73,17 @@ def mtf_model_fn(self, features, mesh):
     kv_channels = mtf.Dimension("kv_channels", hparams.d_kv)
     heads = mtf.Dimension("heads", hparams.num_heads)
 
-    def infeed_to_batch_by_length(x, name):
-      return mtf.infeed(
-          mesh, x, mtf.TensorShape([batch_dim, length_dim]), name=name)
+    def import_to_batch_by_length(x, name):
+      return mtf.import_tf_tensor(
+          mesh, x, mtf.Shape([batch_dim, length_dim]), name=name)
 
     def layer_prepostprocess_dropout(x):
       return mtf.dropout(
           x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
-          noise_shape=mtf.TensorShape([batch_dim, model_dim]))
+          noise_shape=mtf.Shape([batch_dim, model_dim]))
 
-    targets = infeed_to_batch_by_length(targets, "targets")
-    shifted_targets = infeed_to_batch_by_length(
+    targets = import_to_batch_by_length(targets, "targets")
+    shifted_targets = import_to_batch_by_length(
         shifted_targets, "shifted_targets")
 
     extra_losses = []
@@ -93,7 +93,7 @@ def layer_prepostprocess_dropout(x):
       vocab_size = hparams.num_classes
       inputs_vocab_dim = mtf.Dimension("vocab", vocab_size)
       inputs = tf.squeeze(tf.to_int32(features["inputs"]), [2, 3])
-      inputs = infeed_to_batch_by_length(inputs, "inputs")
+      inputs = import_to_batch_by_length(inputs, "inputs")
 
       # Input embeddings
       inputs, _ = mtf_layers.embedding(
@@ -110,13 +110,13 @@ def layer_prepostprocess_dropout(x):
     # Create embedding var for targets and positions and do a gather.
     targets_embedding_var = mtf.get_variable(
         mesh, "targets_embedding",
-        mtf.TensorShape([targets_vocab_dim, model_dim]),
+        mtf.Shape([targets_vocab_dim, model_dim]),
         initializer=tf.random_normal_initializer(),
         activation_dtype=activation_dtype)
 
     positional_embedding_var = mtf.get_variable(
         mesh, "positional_embedding",
-        mtf.TensorShape([targets_vocab_dim, model_dim]),
+        mtf.Shape([targets_vocab_dim, model_dim]),
         initializer=tf.random_normal_initializer(),
         activation_dtype=activation_dtype)
     x = (mtf.gather(targets_embedding_var, shifted_targets, targets_vocab_dim) +
diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py b/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
index 655e50194..770b53af5 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
@@ -60,12 +60,11 @@ def get_model(hparams=None,
 def get_placement_mesh(hparams):
   graph = mtf.Graph()
   mesh = mtf.Mesh(graph, "my_mesh")
-  mesh_shape = mtf.parse_mesh_shape(hparams.mesh_shape)
-  mesh_size = mtf.list_product(mesh_shape)
+  mesh_shape = mtf.convert_to_shape(hparams.mesh_shape)
 
-  mesh_devices = [""] * mesh_size
+  mesh_devices = [""] * mesh_shape.size
   mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-      mesh_shape, mtf.parse_layout(hparams.layout), mesh_devices)
+      mesh_shape, hparams.layout, mesh_devices)
   return mesh, mesh_impl
 
 
@@ -82,7 +81,7 @@ def testMtfImageTransformer(self):
     logits, _ = model.mtf_model_fn(features, mesh)
     lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
     tf_group = lowering.copy_masters_to_slices()
-    tf_logits = lowering.outfeed(logits)
+    tf_logits = lowering.export_to_tf_tensor(logits)
 
     with self.test_session() as session:
       session.run(tf.global_variables_initializer())
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers.py b/tensor2tensor/mesh_tensorflow/mtf_layers.py
index 09a2dea93..bb23a800c 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers.py
@@ -45,8 +45,8 @@ def dense(x, output_dim, reduced_dims=None, expert_dims=None,
     expert_dims = []
   if reduced_dims is None:
     reduced_dims = x.shape.dims[-1:]
-  w_shape = mtf.TensorShape(expert_dims + reduced_dims + [output_dim])
-  output_shape = mtf.TensorShape(
+  w_shape = mtf.Shape(expert_dims + reduced_dims + [output_dim])
+  output_shape = mtf.Shape(
       [d for d in x.shape.dims if d not in reduced_dims] + [output_dim])
   with tf.variable_scope(name, default_name="dense"):
     stddev = mtf.list_product(d.size for d in reduced_dims) ** -0.5
@@ -61,7 +61,7 @@ def dense(x, output_dim, reduced_dims=None, expert_dims=None,
       b = mtf.get_variable(
           x.mesh,
           "bias",
-          mtf.TensorShape(expert_dims + [output_dim]),
+          mtf.Shape(expert_dims + [output_dim]),
           initializer=tf.zeros_initializer(),
           activation_dtype=x.dtype)
       y += b
@@ -86,13 +86,13 @@ def layer_norm(x, dim, epsilon=1e-6, name="layer_prepostprocess"):
     scale = mtf.get_variable(
         x.mesh,
         "layer_norm_scale",
-        mtf.TensorShape([dim]),
+        mtf.Shape([dim]),
         initializer=tf.ones_initializer(),
         activation_dtype=x.dtype)
     bias = mtf.get_variable(
         x.mesh,
         "layer_norm_bias",
-        mtf.TensorShape([dim]),
+        mtf.Shape([dim]),
         initializer=tf.zeros_initializer(),
         activation_dtype=x.dtype)
     reduced_shape = x.shape - dim
@@ -159,7 +159,7 @@ def dense_relu_dense(x,
     w = mtf.get_variable(
         x.mesh,
         "kernel",
-        mtf.TensorShape([io, io_channels, hidden_channels]),
+        mtf.Shape([io, io_channels, hidden_channels]),
         initializer=tf.random_normal_initializer(stddev=stddev),
         activation_dtype=x.dtype)
     wi, wo = mtf.unstack(w, io)
@@ -220,13 +220,13 @@ def masked_local_attention_1d(query_antecedent,
     # Get query q, keys k and values v.
     q = mtf.einsum(
         [query_antecedent, q_var],
-        mtf.TensorShape([batch, heads, query_length, kv_channels]))
+        mtf.Shape([batch, heads, query_length, kv_channels]))
     k = mtf.einsum(
         [memory_antecedent, k_var],
-        mtf.TensorShape([batch, heads, memory_length, kv_channels]))
+        mtf.Shape([batch, heads, memory_length, kv_channels]))
     v = mtf.einsum(
         [memory_antecedent, v_var],
-        mtf.TensorShape([batch, heads, memory_length, kv_channels]))
+        mtf.Shape([batch, heads, memory_length, kv_channels]))
 
     # Let's assume for now we don't have padding and the block length equally
     # divides the memory length.
@@ -237,11 +237,11 @@ def masked_local_attention_1d(query_antecedent,
     num_blocks = mtf.Dimension("num_blocks", query_length.size // block_length)
 
     q = mtf.reshape(
-        q, mtf.TensorShape([batch, heads, num_blocks, blength, kv_channels]))
+        q, mtf.Shape([batch, heads, num_blocks, blength, kv_channels]))
     k = mtf.reshape(
-        k, mtf.TensorShape([batch, heads, num_blocks, mlength, kv_channels]))
+        k, mtf.Shape([batch, heads, num_blocks, mlength, kv_channels]))
     v = mtf.reshape(
-        v, mtf.TensorShape([batch, heads, num_blocks, mlength, kv_channels]))
+        v, mtf.Shape([batch, heads, num_blocks, mlength, kv_channels]))
 
     # compute attention for the first query block.
     def first_block_attention():
@@ -253,11 +253,11 @@ def first_block_attention():
 
       first_logits = mtf.einsum(
           [first_q, first_k],
-          mtf.TensorShape([batch, heads, block, blength, mlength]))
+          mtf.Shape([batch, heads, block, blength, mlength]))
       weights = mtf.softmax(first_logits, mlength)
       first_output = mtf.einsum(
           [weights, first_v],
-          mtf.TensorShape([batch, heads, block, blength, kv_channels]))
+          mtf.Shape([batch, heads, block, blength, kv_channels]))
       return first_output
 
     # Attention for first block, since query_length = key_length.
@@ -287,7 +287,7 @@ def local(x):
     # Shape [batch, heads, num_blocks - 1, block_length, local_length]
     attention = mtf.einsum(
         [tail_q, local_k],
-        mtf.TensorShape([batch, heads, mblocks, blength, mlength]))
+        mtf.Shape([batch, heads, mblocks, blength, mlength]))
     attention += mask
     attention = mtf.softmax(attention, mlength)
 
@@ -295,13 +295,13 @@ def local(x):
     # Shape [batch, heads, num_blocks-1, block_length, kv_channels]
     output = mtf.einsum(
         [attention, local_v],
-        mtf.TensorShape([batch, heads, mblocks, blength, kv_channels]))
+        mtf.Shape([batch, heads, mblocks, blength, kv_channels]))
     # Now concatenate the first and rest of the blocks.
     final_output = mtf.concat([first_output, output], num_blocks.name)
-    final_output = mtf.reshape(final_output, mtf.TensorShape(
+    final_output = mtf.reshape(final_output, mtf.Shape(
         [batch, heads, query_length, kv_channels]))
     return mtf.einsum([final_output, o_var],
-                      mtf.TensorShape([batch, query_length, io_channels]))
+                      mtf.Shape([batch, query_length, io_channels]))
 
 
 def rename_length_to_memory_length(
@@ -338,7 +338,7 @@ def qkvo_initializer(shape,
     return tf.random_normal(shape, dtype=dtype) * tf.reshape(
         [qk_stddev, qk_stddev, v_stddev, o_stddev], [4, 1, 1, 1])
   var = mtf.get_variable(
-      mesh, "qkvo", mtf.TensorShape([qkvo, heads, io_channels, kv_channels]),
+      mesh, "qkvo", mtf.Shape([qkvo, heads, io_channels, kv_channels]),
       initializer=qkvo_initializer, activation_dtype=activation_dtype)
   q_var, k_var, v_var, o_var = mtf.unstack(var, qkvo)
   return q_var, k_var, v_var, o_var
@@ -367,7 +367,7 @@ def dot_product_attention(q,
     Tensor with shape [..., length_q, depth_v].
   """
   length_kv = k.shape.dims[-2]
-  logits_shape = mtf.TensorShape(q.shape.dims[:-1] + [length_kv])
+  logits_shape = mtf.Shape(q.shape.dims[:-1] + [length_kv])
   logits = mtf.einsum([q, k], logits_shape)
   if mask is not None:
     logits += mask
@@ -377,7 +377,7 @@ def dot_product_attention(q,
         weights, 1.0 - dropout,
         noise_shape=weights.shape - dropout_broadcast_dims)
   depth_v = v.shape.dims[-1]
-  outputs_shape = mtf.TensorShape(q.shape.dims[:-1] + [depth_v])
+  outputs_shape = mtf.Shape(q.shape.dims[:-1] + [depth_v])
   outputs = mtf.einsum([weights, v], outputs_shape)
   return outputs
 
@@ -431,17 +431,17 @@ def multihead_attention(query_antecedent,
       raise ValueError("memory channels must equal query channels")
     q = mtf.einsum(
         [query_antecedent, q_var],
-        mtf.TensorShape([batch, heads, query_length, kv_channels]))
+        mtf.Shape([batch, heads, query_length, kv_channels]))
     k = mtf.einsum(
         [memory_antecedent, k_var],
-        mtf.TensorShape([batch, heads, memory_length, kv_channels]))
+        mtf.Shape([batch, heads, memory_length, kv_channels]))
     v = mtf.einsum(
         [memory_antecedent, v_var],
-        mtf.TensorShape([batch, heads, memory_length, kv_channels]))
+        mtf.Shape([batch, heads, memory_length, kv_channels]))
     o = dot_product_attention(
         q, k, v, mask, dropout, dropout_broadcast_dims)
     return mtf.einsum(
-        [o, o_var], mtf.TensorShape([batch, query_length, io_channels]))
+        [o, o_var], mtf.Shape([batch, query_length, io_channels]))
 
 
 def multihead_self_attention_incremental(query_antecedent,
@@ -481,13 +481,13 @@ def multihead_self_attention_incremental(query_antecedent,
     memory_antecedent = query_antecedent
     q = mtf.einsum(
         [query_antecedent, q_var],
-        mtf.TensorShape(batch_dims + [heads, kv_channels]))
+        mtf.Shape(batch_dims + [heads, kv_channels]))
     k = mtf.einsum(
         [memory_antecedent, k_var],
-        mtf.TensorShape(batch_dims + [heads, kv_channels]))
+        mtf.Shape(batch_dims + [heads, kv_channels]))
     v = mtf.einsum(
         [memory_antecedent, v_var],
-        mtf.TensorShape(batch_dims + [heads, kv_channels]))
+        mtf.Shape(batch_dims + [heads, kv_channels]))
     k = prev_k + mtf.multiply(
         k, mtf.one_hot(step_num, memory_length), output_shape=prev_k.shape)
     v = prev_v + mtf.multiply(
@@ -531,7 +531,7 @@ def multihead_encdec_attention_incremental(query_antecedent,
   with tf.variable_scope(name, default_name="multihead_attention"):
     q = mtf.einsum(
         [query_antecedent, q_var],
-        mtf.TensorShape(query_dims + [heads, kv_channels]))
+        mtf.Shape(query_dims + [heads, kv_channels]))
     o = dot_product_attention(q, k, v, mask)
     return mtf.einsum([o, o_var], query_antecedent.shape)
 
@@ -674,7 +674,7 @@ def moe_v0(inputs,
   # shape = [batch_dim, length_dim, experts_dim_unsplit]
   gates = mtf.softmax(dense(inputs, experts_dim_unsplit), experts_dim_unsplit)
 
-  assignment_shape = mtf.TensorShape(
+  assignment_shape = mtf.Shape(
       [batch_dim, length_dim, experts_dim_unsplit, expert_capacity_dim])
 
   backward_assignment = mtf.slicewise(
@@ -689,10 +689,10 @@ def moe_v0(inputs,
       mtf.cast(backward_assignment, tf.bool), inputs.dtype)
 
   # put num_experts dimension first to make split easier in alltoall
-  expert_inputs = mtf.einsum([inputs, forward_assignment], mtf.TensorShape(
+  expert_inputs = mtf.einsum([inputs, forward_assignment], mtf.Shape(
       [experts_dim_unsplit, batch_dim, expert_capacity_dim, input_dim]))
 
-  expert_inputs = mtf.reshape(expert_inputs, mtf.TensorShape(
+  expert_inputs = mtf.reshape(expert_inputs, mtf.Shape(
       [experts_dim, batch_dim_unsplit, expert_capacity_dim, input_dim]))
 
   # Now feed the expert inputs through the experts.
@@ -700,13 +700,13 @@ def moe_v0(inputs,
             activation=mtf.relu, name="x0")
   expert_output = dense(h, output_dim, expert_dims=[experts_dim], name="x1")
 
-  expert_output = mtf.reshape(expert_output, mtf.TensorShape(
+  expert_output = mtf.reshape(expert_output, mtf.Shape(
       [experts_dim_unsplit, batch_dim, expert_capacity_dim, input_dim]))
 
-  output = mtf.einsum([expert_output, backward_assignment], mtf.TensorShape(
+  output = mtf.einsum([expert_output, backward_assignment], mtf.Shape(
       [batch_dim, length_dim, output_dim]))
 
-  importance = mtf.reduce_sum(backward_assignment, output_shape=mtf.TensorShape(
+  importance = mtf.reduce_sum(backward_assignment, output_shape=mtf.Shape(
       [batch_dim, experts_dim_unsplit]))
 
   loss = cv_squared(importance) * loss_coef
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers_test.py b/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
index d2592c7c3..1c37d4c8d 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
@@ -45,17 +45,17 @@ def testDense(self, units, use_bias):
     channels_dim = mtf.Dimension("channels", channels)
     depth_dim = mtf.Dimension("depth", units)
 
-    mtf_inputs = mtf.infeed(mesh, inputs,
-                            shape=mtf.TensorShape([batch_dim, channels_dim]))
+    mtf_inputs = mtf.import_tf_tensor(
+        mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
     mtf_outputs = mtf_layers.dense(mtf_inputs,
                                    output_dim=depth_dim,
                                    reduced_dims=[channels_dim],
                                    activation=mtf.relu,
                                    use_bias=use_bias)
     mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[1], layout={}, devices=[""])
+        shape=[], layout={}, devices=[""])
     lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-    actual_outputs = lowering.outfeed(mtf_outputs)
+    actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
 
     expected_outputs = tf.keras.layers.Dense(units=units,
                                              activation=tf.nn.relu,
@@ -79,14 +79,14 @@ def testLayerNorm(self):
     batch_dim = mtf.Dimension("batch", batch)
     channels_dim = mtf.Dimension("channels", channels)
 
-    mtf_inputs = mtf.infeed(mesh, inputs,
-                            shape=mtf.TensorShape([batch_dim, channels_dim]))
+    mtf_inputs = mtf.import_tf_tensor(
+        mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
     mtf_outputs = mtf_layers.layer_norm(mtf_inputs,
                                         dim=channels_dim)
     mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[1], layout={}, devices=[""])
+        shape=[], layout={}, devices=[""])
     lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-    actual_outputs = lowering.outfeed(mtf_outputs)
+    actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
 
     expected_outputs = common_layers.layer_norm(inputs)
     tf_group = lowering.copy_masters_to_slices()
@@ -106,13 +106,13 @@ def testWeightsNonzero(self):
     batch_dim = mtf.Dimension("batch", inputs.shape.as_list()[0])
     channels_dim = mtf.Dimension("channels", inputs.shape.as_list()[1])
 
-    mtf_inputs = mtf.infeed(mesh, inputs,
-                            shape=mtf.TensorShape([batch_dim, channels_dim]))
+    mtf_inputs = mtf.import_tf_tensor(
+        mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
     mtf_outputs = mtf_layers.weights_nonzero(mtf_inputs)
     mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[1], layout={}, devices=[""])
+        shape=[], layout={}, devices=[""])
     lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-    actual_outputs = lowering.outfeed(mtf_outputs)
+    actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
 
     expected_outputs = common_layers.weights_nonzero(inputs)
     tf_group = lowering.copy_masters_to_slices()
@@ -134,14 +134,14 @@ def testDenseReluDense(self):
     channels_dim = mtf.Dimension("channels", channels)
     hidden_dim = mtf.Dimension("hidden", hidden)
 
-    mtf_inputs = mtf.infeed(mesh, inputs,
-                            shape=mtf.TensorShape([batch_dim, channels_dim]))
+    mtf_inputs = mtf.import_tf_tensor(
+        mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
     mtf_outputs = mtf_layers.dense_relu_dense(mtf_inputs,
                                               hidden_channels=hidden_dim)
     mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[1], layout={}, devices=[""])
+        shape=[], layout={}, devices=[""])
     lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-    actual_outputs = lowering.outfeed(mtf_outputs)
+    actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
 
     tf_group = lowering.copy_masters_to_slices()
     init = tf.global_variables_initializer()
@@ -172,12 +172,12 @@ def testMaskedLocalAttention1D(self, kv_channels, heads):
     kv_channels_dim = mtf.Dimension("kv_channels", kv_channels)
     heads_dim = mtf.Dimension("heads", heads)
 
-    mtf_query = mtf.infeed(
+    mtf_query = mtf.import_tf_tensor(
         mesh, query,
-        shape=mtf.TensorShape([batch_dim, length_q_dim, channels_dim]))
-    mtf_memory = mtf.infeed(
+        shape=mtf.Shape([batch_dim, length_q_dim, channels_dim]))
+    mtf_memory = mtf.import_tf_tensor(
         mesh, memory,
-        shape=mtf.TensorShape([batch_dim, length_m_dim, channels_dim]))
+        shape=mtf.Shape([batch_dim, length_m_dim, channels_dim]))
     mtf_outputs = mtf_layers.masked_local_attention_1d(
         mtf_query,
         mtf_memory,
@@ -185,9 +185,9 @@ def testMaskedLocalAttention1D(self, kv_channels, heads):
         heads=heads_dim,
         block_length=2)
     mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[1], layout={}, devices=[""])
+        shape=[], layout={}, devices=[""])
     lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-    actual_outputs = lowering.outfeed(mtf_outputs)
+    actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
 
     tf_group = lowering.copy_masters_to_slices()
     init = tf.global_variables_initializer()
@@ -216,17 +216,17 @@ def testDotProductAttention(
     depth_k_dim = mtf.Dimension("depth_k", depth_k)
     depth_v_dim = mtf.Dimension("depth_v", depth_v)
 
-    mtf_query = mtf.infeed(
+    mtf_query = mtf.import_tf_tensor(
         mesh, query,
-        shape=mtf.TensorShape(
+        shape=mtf.Shape(
             [batch_dim, heads_dim, length_q_dim, depth_k_dim]))
-    mtf_key = mtf.infeed(
+    mtf_key = mtf.import_tf_tensor(
         mesh, key,
-        shape=mtf.TensorShape(
+        shape=mtf.Shape(
             [batch_dim, heads_dim, length_kv_dim, depth_k_dim]))
-    mtf_value = mtf.infeed(
+    mtf_value = mtf.import_tf_tensor(
         mesh, value,
-        shape=mtf.TensorShape(
+        shape=mtf.Shape(
             [batch_dim, heads_dim, length_kv_dim, depth_v_dim]))
     mtf_outputs = mtf_layers.dot_product_attention(
         mtf_query,
@@ -234,9 +234,9 @@ def testDotProductAttention(
         mtf_value,
         mask=None)
     mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[1], layout={}, devices=[""])
+        shape=[], layout={}, devices=[""])
     lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-    actual_outputs = lowering.outfeed(mtf_outputs)
+    actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
 
     tf_group = lowering.copy_masters_to_slices()
     init = tf.global_variables_initializer()
@@ -265,9 +265,9 @@ def testMultiheadAttention(self, kv_channels, heads):
     kv_channels_dim = mtf.Dimension("kv_channels", kv_channels)
     heads_dim = mtf.Dimension("heads", heads)
 
-    mtf_query = mtf.infeed(
+    mtf_query = mtf.import_tf_tensor(
         mesh, query,
-        shape=mtf.TensorShape([batch_dim, length_dim, channels_dim]))
+        shape=mtf.Shape([batch_dim, length_dim, channels_dim]))
     mtf_outputs = mtf_layers.multihead_attention(
         mtf_query,
         memory_antecedent=None,
@@ -275,9 +275,9 @@ def testMultiheadAttention(self, kv_channels, heads):
         kv_channels=kv_channels_dim,
         heads=heads_dim)
     mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[1], layout={}, devices=[""])
+        shape=[], layout={}, devices=[""])
     lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-    actual_outputs = lowering.outfeed(mtf_outputs)
+    actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
 
     tf_group = lowering.copy_masters_to_slices()
     init = tf.global_variables_initializer()
diff --git a/tensor2tensor/mesh_tensorflow/mtf_model.py b/tensor2tensor/mesh_tensorflow/mtf_model.py
index 59f975d83..ff752191d 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_model.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_model.py
@@ -76,21 +76,21 @@ def estimator_model_fn(cls,
     graph = mtf.Graph()
     mesh = mtf.Mesh(graph, "my_mesh")
 
-    mesh_shape = mtf.parse_mesh_shape(hparams.mesh_shape)
-    mesh_size = mtf.list_product(mesh_shape)
+    mesh_shape = mtf.convert_to_shape(hparams.mesh_shape)
+    computation_layout = mtf.ComputationLayout(hparams.layout)
     if use_tpu:
-      mesh_devices = [""] * mesh_size
+      mesh_devices = [""] * mesh_shape.size
       mesh_impl = simd_mesh_impl.SimdMeshImpl(
-          mesh_shape, mtf.parse_layout(hparams.layout), mesh_devices,
+          mesh_shape, computation_layout, mesh_devices,
           params["context"].device_assignment)
     else:
       if len(data_parallelism.ps_devices) == 1:
-        mesh_devices = [""] * mesh_size
+        mesh_devices = [""] * mesh_shape.size
       else:
-        assert len(data_parallelism.ps_devices) == mesh_size
+        assert len(data_parallelism.ps_devices) == mesh_shape.size
         mesh_devices = data_parallelism.ps_devices
       mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-          mesh_shape, mtf.parse_layout(hparams.layout), mesh_devices)
+          mesh_shape, computation_layout, mesh_devices)
 
     # PREDICT mode
     if mode == tf.estimator.ModeKeys.PREDICT:
@@ -105,8 +105,8 @@ def estimator_model_fn(cls,
       var_grads = mtf.gradients(
           [loss], [v.outputs[0] for v in graph.trainable_variables])
       lr = learning_rate.learning_rate_schedule(hparams)
-      mtf_lr = mtf.infeed(
-          mesh, tf.convert_to_tensor(lr, dtype=tf.float32), mtf.TensorShape([]))
+      mtf_lr = mtf.import_tf_tensor(
+          mesh, tf.convert_to_tensor(lr, dtype=tf.float32), mtf.Shape([]))
       optimizer = mtf_optimize.make_optimizer(hparams, mtf_lr)
       update_ops = []
       for grad, var in zip(var_grads, graph.trainable_variables):
@@ -114,10 +114,10 @@ def estimator_model_fn(cls,
 
     lowering = mtf.Lowering(graph, {mesh: mesh_impl})
 
-    tf_loss = lowering.outfeed(loss)
+    tf_loss = lowering.export_to_tf_tensor(loss)
     tf_loss = tf.to_float(tf_loss)
     if logits and mode != tf.estimator.ModeKeys.TRAIN:
-      tf_logits = lowering.outfeed(logits)
+      tf_logits = lowering.export_to_tf_tensor(logits)
 
     if mode == tf.estimator.ModeKeys.TRAIN:
       tf_update_ops = [lowering.lowered_operation(op) for op in update_ops]
@@ -145,7 +145,7 @@ def estimator_model_fn(cls,
 
     # EVAL mode
     if mode == tf.estimator.ModeKeys.EVAL:
-      tf_logits = lowering.outfeed(logits)
+      tf_logits = lowering.export_to_tf_tensor(logits)
       return model.estimator_spec_eval(features, tf_logits, labels, tf_loss,
                                        restore_hook, use_tpu)
 
@@ -201,7 +201,7 @@ def metric_fn(tf_logits, labels):
   def estimator_spec_predict(self, features, mesh, mesh_impl, use_tpu):
     mtf_samples = self.sample(features, mesh)
     lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
-    outputs = lowering.outfeed(mtf_samples)
+    outputs = lowering.export_to_tf_tensor(mtf_samples)
     if self.has_input:
       ndims = len(outputs.shape.as_list())
       actual_batch_size = tf.shape(features["inputs"])[0]
diff --git a/tensor2tensor/mesh_tensorflow/mtf_optimize.py b/tensor2tensor/mesh_tensorflow/mtf_optimize.py
index df822fb27..0dcd0a79e 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_optimize.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_optimize.py
@@ -106,7 +106,7 @@ def _factored_dims(self, shape):
     If there are not two dimensions of size >=128, then we do not factor.
 
     Args:
-      shape: a TensorShape
+      shape: a Shape
     Returns:
       either a list of 2 Dimensions or None
     """
diff --git a/tensor2tensor/mesh_tensorflow/mtf_toy.py b/tensor2tensor/mesh_tensorflow/mtf_toy.py
deleted file mode 100644
index 811d47086..000000000
--- a/tensor2tensor/mesh_tensorflow/mtf_toy.py
+++ /dev/null
@@ -1,178 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-"""Toy model."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.layers import common_hparams
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-from tensor2tensor.mesh_tensorflow import mtf_layers
-from tensor2tensor.mesh_tensorflow import mtf_model
-from tensor2tensor.utils import registry
-import tensorflow as tf
-
-from tensorflow.contrib.tpu.python.ops import tpu_ops
-
-
-@registry.register_model
-class MtfToy(mtf_model.MtfModel):
-  """Toy model to test mesh_tensorflow."""
-
-  def mtf_model_fn(self, features, mesh):
-    hparams = self._hparams
-    # tf_x = tf.random_uniform([hparams.batch_size, hparams.io_size])
-    tf_x = tf.matmul(
-        tf.reshape(
-            tf.lin_space(0., 1.0, hparams.batch_size), [hparams.batch_size, 1]),
-        tf.reshape(
-            tf.lin_space(0., 1.0, hparams.io_size), [1, hparams.io_size]))
-    batch_dim = mtf.Dimension("batch", hparams.batch_size)
-
-    hidden_dim = mtf.Dimension("hidden", hparams.hidden_size)
-    io_dim = mtf.Dimension("io", hparams.io_size)
-    x = mtf.infeed_fully_replicated(
-        mesh, tf_x, mtf.TensorShape([batch_dim, io_dim]))
-    h = mtf_layers.dense(x, hidden_dim, name="layer1", use_bias=False)
-    y = mtf_layers.dense(h, io_dim, name="layer2", use_bias=False)
-
-    loss = mtf.reduce_sum(mtf.square(y - x))
-    return None, loss
-
-
-@registry.register_model
-class MtfSimple(mtf_model.MtfModel):
-  """Toy model to test mesh_tensorflow."""
-
-  def mtf_model_fn(self, features, mesh):
-    hparams = self._hparams
-    # tf_x = tf.random_uniform([hparams.batch_size, hparams.io_size])
-    tf_x = tf.matmul(
-        tf.reshape(
-            tf.lin_space(0., 1.0, hparams.batch_size), [hparams.batch_size, 1]),
-        tf.reshape(
-            tf.lin_space(0., 1.0, hparams.io_size), [1, hparams.io_size]))
-    batch_dim = mtf.Dimension("batch", hparams.batch_size)
-    hidden_dim = mtf.Dimension("hidden", hparams.hidden_size)
-    io_dim = mtf.Dimension("io", hparams.io_size)
-
-    x = mtf.infeed_fully_replicated(
-        mesh, tf_x, mtf.TensorShape([batch_dim, io_dim]))
-    h = mtf_layers.dense(x, hidden_dim, name="layer1", use_bias=False)
-    y = mtf_layers.dense(h, io_dim, name="layer2", use_bias=False)
-    loss = mtf.reduce_sum(mtf.square(y - x))
-    return None, loss
-
-
-@registry.register_model
-class MtfToyNormal(mtf_model.MtfModel):
-  """Toy model to test mesh_tensorflow."""
-
-  def mtf_model_fn(self, features, mesh):
-    hparams = self._hparams
-    hparams.batch_size = 10
-    hparams.io_size = 4
-    hparams.hidden_size = 2
-    tf_x = tf.matmul(
-        tf.reshape(
-            tf.lin_space(0., 1.0, hparams.batch_size), [hparams.batch_size, 1]),
-        tf.reshape(
-            tf.lin_space(0., 1.0, hparams.io_size), [1, hparams.io_size]))
-    # tf_x = tf.random_uniform([hparams.batch_size, hparams.io_size])
-
-    hidden_1_variable = tf.get_variable(
-        "a",
-        shape=[hparams.io_size, hparams.hidden_size],
-        initializer=tf.random_normal_initializer())
-    hidden_2_variable = tf.get_variable(
-        "b",
-        shape=[hparams.hidden_size, hparams.io_size],
-        initializer=tf.random_normal_initializer())
-
-    hidden_layer_1 = tf.matmul(tf_x, hidden_1_variable)
-    hidden_layer_2 = tf.matmul(hidden_layer_1, hidden_2_variable)
-    hidden_layer_2 = tpu_ops.cross_replica_sum(hidden_layer_2)
-    loss = tf.reduce_mean(tf.square(hidden_layer_2 - tf_x))
-    return None, loss
-
-
-def set_sgd_optimizer(hparams):
-  hparams.optimizer = "SGD"
-  hparams.learning_rate_schedule = "constant"
-  hparams.learning_rate_constant = 0.01
-
-
-def set_adafactor_optimizer(hparams):
-  hparams.optimizer = "Adafactor"
-  hparams.learning_rate_schedule = "rsqrt_decay"
-  hparams.optimizer_adafactor_factored = True
-  hparams.learning_rate_warmup_steps = 1000
-
-
-@registry.register_hparams
-def mtf_toy_base():
-  """Set of hyperparameters."""
-  hparams = common_hparams.basic_params1()
-  hparams.no_data_parallelism = True
-  hparams.use_fixed_batch_size = True
-  hparams.add_hparam("mtf_mode", True)
-  hparams.batch_size = 64
-  set_adafactor_optimizer(hparams)
-  hparams.add_hparam("io_size", 32)
-  hparams.hidden_size = 32
-  hparams.add_hparam("mesh_shape", "4.2")
-  hparams.add_hparam("layout", "batch:0;hidden:1")
-  return hparams
-
-
-@registry.register_hparams
-def mtf_toy_data_parallel():
-  """Set of hyperparameters."""
-  hparams = mtf_toy_base()
-  hparams.add_hparam("layout", "batch:0")
-  return hparams
-
-
-@registry.register_hparams
-def mtf_toy_model_parallel():
-  """Set of hyperparameters."""
-  hparams = mtf_toy_base()
-  hparams.add_hparam("layout", "hidden:0")
-  return hparams
-
-
-@registry.register_hparams
-def mtf_toy_data_parallel_m2():
-  """Set of hyperparameters."""
-  hparams = mtf_toy_data_parallel()
-  hparams.mesh_shape = "2"
-  return hparams
-
-
-@registry.register_hparams
-def mtf_toy_model_parallel_m2():
-  """Set of hyperparameters."""
-  hparams = mtf_toy_model_parallel()
-  hparams.mesh_shape = "2"
-  return hparams
-
-
-@registry.register_hparams
-def mtf_toy_m32():
-  """Set of hyperparameters."""
-  hparams = mtf_toy_base()
-  hparams.mesh_shape = "8;4"
-  return hparams
diff --git a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
index 8dfea0f7c..7366e253d 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
@@ -40,13 +40,12 @@
 tf.flags.DEFINE_integer('batch_size', 64, 'Training batch size.')
 tf.flags.DEFINE_integer('io_size', 2, 'Number of channels per feature.')
 tf.flags.DEFINE_integer('hidden_size', 2, 'Size of each hidden layer.')
-tf.flags.DEFINE_string('mesh_shape', '2;1', 'mesh shape')
-tf.flags.DEFINE_string('layout', 'batch:0', 'computation layout')
+tf.flags.DEFINE_string('mesh_shape', 'all:8', 'mesh shape')
+tf.flags.DEFINE_string('layout', 'hidden:all', 'computation layout')
 tf.flags.DEFINE_integer('iterations', 100,
                         'Number of iterations per training loop.')
 tf.flags.DEFINE_integer('train_steps', 10000, 'max steps')
 tf.flags.DEFINE_integer('steps_per_checkpoint', 200, 'steps_per_checkpoint')
-tf.flags.DEFINE_integer('num_shards', 2, 'Number of shards.')
 tf.flags.DEFINE_string('master', 'local',
                        'BNS name of the TensorFlow master to use.')
 tf.flags.DEFINE_string('evaluation_master', 'local',
@@ -92,7 +91,7 @@ def toy_model(features, mesh):
   hidden_dim = mtf.Dimension('hidden', FLAGS.hidden_size)
   io_dim = mtf.Dimension('io', FLAGS.io_size)
 
-  x = mtf.infeed(mesh, features, mtf.TensorShape([batch_dim, io_dim]))
+  x = mtf.import_tf_tensor(mesh, features, mtf.Shape([batch_dim, io_dim]))
   h = mtf_layers.dense(x, hidden_dim, name='layer1', use_bias=False)
   y = mtf_layers.dense(h, io_dim, name='layer2', use_bias=False)
 
@@ -106,11 +105,11 @@ def model_fn(features, labels, mode, params):
   global_step = tf.train.get_global_step()
   graph = mtf.Graph()
   mesh = mtf.Mesh(graph, 'my_mesh')
-  mesh_shape = mtf.parse_mesh_shape(FLAGS.mesh_shape)
-  mesh_size = mtf.list_product(mesh_shape)
-  mesh_devices = [''] * mesh_size
-  mesh_impl = SimdMeshImpl(mesh_shape, mtf.parse_layout(FLAGS.layout),
-                           mesh_devices, params['context'].device_assignment)
+  mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
+  mesh_devices = [''] * mesh_shape.size
+  mesh_impl = SimdMeshImpl(
+      mesh_shape, mtf.convert_to_computation_layout(FLAGS.layout),
+      mesh_devices, params['context'].device_assignment)
   with mtf_utils.outside_all_rewrites():
     logits, loss = toy_model(features, mesh)
 
@@ -122,17 +121,21 @@ def model_fn(features, labels, mode, params):
     update_ops = []
     for grad, var in zip(var_grads, graph.trainable_variables):
       update_ops.extend(optimizer.apply_grad(grad, var))
+  else:
+    # for now, we can only export fully-replicated tensors.
+    fully_replicated_logits = mtf.anonymize(logits)
 
   lowering = mtf.Lowering(graph, {mesh: mesh_impl})
 
-  tf_loss = lowering.outfeed(loss)
-  tf_logits = lowering.outfeed(logits)
+  tf_loss = lowering.export_to_tf_tensor(loss)
 
   if mode == tf.estimator.ModeKeys.TRAIN:
     tf_update_ops = [lowering.lowered_operation(op) for op in update_ops]
     tf_update_ops.append(tf.assign_add(global_step, 1))
     tf.logging.info('tf_update_ops: {}'.format(tf_update_ops))
     train_op = tf.group(tf_update_ops)
+  else:
+    tf_logits = lowering.export_to_tf_tensor(fully_replicated_logits)
 
   with mtf_utils.outside_all_rewrites():
     # Copy master variables to slices. Must be called first.
@@ -176,6 +179,7 @@ def metric_fn(tf_logits):
 def run_toy_model_tpu():
   """Run a toy model on TPU."""
   iterations_per_loop = FLAGS.iterations
+  mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
   config = tpu_config.RunConfig(
       master=FLAGS.master,
       evaluation_master=FLAGS.evaluation_master,
@@ -184,7 +188,7 @@ def run_toy_model_tpu():
       save_checkpoints_secs=None,  # Disable the default saver
       log_step_count_steps=iterations_per_loop,
       tpu_config=tpu_config.TPUConfig(
-          num_shards=FLAGS.num_shards,
+          num_shards=mesh_shape.size,
           iterations_per_loop=iterations_per_loop,
           num_cores_per_replica=1,
           per_host_input_for_training=tpu_config.InputPipelineConfig.BROADCAST))
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
index 550af3391..69c855c7a 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
@@ -89,16 +89,16 @@ def activation_dtype(self):
           "unknown hparams.activation_dtype %s"
           % self._hparams.activation_dtype)
 
-  def _infeed_to_batch_by_length(self, x, name, mesh, hparams):
+  def _import_to_batch_by_length(self, x, name, mesh, hparams):
     x = tf.reshape(x, [self.batch_dim.size, self.length_dim.size])
-    return mtf.infeed_fully_replicated(
-        mesh, x, mtf.TensorShape([self.batch_dim, self.length_dim]), name=name)
+    return mtf.import_fully_replicated(
+        mesh, x, mtf.Shape([self.batch_dim, self.length_dim]), name=name)
 
   def _embedding_and_softmax_vars(self, mesh):
     hparams = self._hparams
     targets_embedding_var = mtf.get_variable(
         mesh, "targets_embedding",
-        mtf.TensorShape([self.targets_vocab_dim, self.model_dim]),
+        mtf.Shape([self.targets_vocab_dim, self.model_dim]),
         initializer=tf.random_normal_initializer(),
         activation_dtype=self.activation_dtype)
     if self.has_input:
@@ -107,7 +107,7 @@ def _embedding_and_softmax_vars(self, mesh):
       else:
         inputs_embedding_var = mtf.get_variable(
             mesh, "inputs_embedding",
-            mtf.TensorShape([self.inputs_vocab_dim, self.model_dim]),
+            mtf.Shape([self.inputs_vocab_dim, self.model_dim]),
             initializer=tf.random_normal_initializer(),
             activation_dtype=self.activation_dtype)
     else:
@@ -118,13 +118,13 @@ def _embedding_and_softmax_vars(self, mesh):
       softmax_var = mtf.get_variable(
           mesh,
           "softmax",
-          mtf.TensorShape([self.targets_vocab_dim, self.model_dim]),
+          mtf.Shape([self.targets_vocab_dim, self.model_dim]),
           initializer=tf.random_normal_initializer(
               stddev=self.model_dim.size**-0.5),
           activation_dtype=self.activation_dtype)
     positional_embedding_var = mtf.get_variable(
         mesh, "positional_embedding",
-        mtf.TensorShape([self.max_length_dim, self.model_dim]),
+        mtf.Shape([self.max_length_dim, self.model_dim]),
         initializer=tf.random_normal_initializer(),
         activation_dtype=self.activation_dtype)
     return (inputs_embedding_var, targets_embedding_var,
@@ -150,16 +150,16 @@ def pad_to_max_length(x):
         features[key] = pad_to_max_length(features[key])
     shifted_targets = common_layers.shift_right_2d(targets)
 
-    targets = self._infeed_to_batch_by_length(targets, "targets", mesh, hparams)
-    shifted_targets = self._infeed_to_batch_by_length(
+    targets = self._import_to_batch_by_length(targets, "targets", mesh, hparams)
+    shifted_targets = self._import_to_batch_by_length(
         shifted_targets, "shifted_targets", mesh, hparams)
 
     if "targets_segmentation" in features:
       # "Packed" dataset - keep the examples from seeing each other.
-      targets_segmentation = self._infeed_to_batch_by_length(
+      targets_segmentation = self._import_to_batch_by_length(
           features["targets_segmentation"], "targets_segmentation",
           mesh, hparams)
-      targets_position = self._infeed_to_batch_by_length(
+      targets_position = self._import_to_batch_by_length(
           features["targets_position"], "targets_position",
           mesh, hparams)
       decoder_self_attention_mask = (
@@ -175,7 +175,7 @@ def pad_to_max_length(x):
     def layer_prepostprocess_dropout(x):
       return mtf.dropout(
           x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
-          noise_shape=mtf.TensorShape([self.batch_dim, self.model_dim]))
+          noise_shape=mtf.Shape([self.batch_dim, self.model_dim]))
 
     extra_losses = []
     (inputs_embedding_var,
@@ -185,13 +185,13 @@ def layer_prepostprocess_dropout(x):
     if self.has_input:
       inputs = tf.squeeze(tf.to_int32(features["inputs"]), [2, 3])
       inputs = pad_to_max_length(inputs)
-      inputs = self._infeed_to_batch_by_length(inputs, "inputs", mesh, hparams)
+      inputs = self._import_to_batch_by_length(inputs, "inputs", mesh, hparams)
       if "inputs_segmentation" in features:
         # "Packed" dataset - keep the examples from seeing each other.
-        inputs_segmentation = self._infeed_to_batch_by_length(
+        inputs_segmentation = self._import_to_batch_by_length(
             features["inputs_segmentation"], "inputs_segmentation",
             mesh, hparams)
-        inputs_position = self._infeed_to_batch_by_length(
+        inputs_position = self._import_to_batch_by_length(
             features["inputs_position"], "inputs_position",
             mesh, hparams)
         encoder_self_attention_mask = (
@@ -338,13 +338,13 @@ def _layer_stack(self,
     def layer_prepostprocess_dropout(x):
       return mtf.dropout(
           x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
-          noise_shape=mtf.TensorShape([self.batch_dim, self.model_dim]))
+          noise_shape=mtf.Shape([self.batch_dim, self.model_dim]))
     num_layer_norms = num_layers * (2 if encoder_output is None else 3) + 1
     layer_norms_dim = mtf.Dimension("layer_norms", num_layer_norms)
     layer_norm_combined_var = mtf.get_variable(
         x.mesh,
         "layer_norm_scale",
-        mtf.TensorShape([layer_norms_dim, self.model_dim]),
+        mtf.Shape([layer_norms_dim, self.model_dim]),
         initializer=tf.ones_initializer(),
         activation_dtype=x.dtype)
     layer_norm_vars = mtf.unstack(layer_norm_combined_var, layer_norms_dim)
@@ -375,7 +375,7 @@ def normalize(x):
         # ffn layer
         x += layer_prepostprocess_dropout(
             self._feedforward_layer(normalize(x), losses=losses))
-    x = normalize(x)
+    x = layer_prepostprocess_dropout(normalize(x))
     assert not layer_norm_vars
     return x
 
@@ -398,11 +398,11 @@ def _sample(self, features, mesh):
       inputs = tf.pad(
           inputs, [[0, hparams.batch_size - actual_batch_size],
                    [0, hparams.max_length - actual_length]])
-      inputs = self._infeed_to_batch_by_length(
+      inputs = self._import_to_batch_by_length(
           inputs, "inputs", mesh, hparams)
       x = (mtf.gather(inputs_embedding_var, inputs, self.inputs_vocab_dim) +
            mtf.reshape(positional_embedding_var,
-                       mtf.TensorShape([self.length_dim, self.model_dim])))
+                       mtf.Shape([self.length_dim, self.model_dim])))
       encoder_attention_mask = (
           mtf_layers.attention_mask_ignore_padding(
               inputs, dtype=self.activation_dtype))
@@ -420,12 +420,12 @@ def _sample(self, features, mesh):
               self.kv_dim, self.activation_dtype)
           k = mtf.einsum(
               [encoder_output, k_var],
-              mtf.TensorShape(
+              mtf.Shape(
                   [self.batch_dim, self.heads_dim,
                    self.memory_length_dim, self.kv_dim]))
           v = mtf.einsum(
               [encoder_output, v_var],
-              mtf.TensorShape(
+              mtf.Shape(
                   [self.batch_dim, self.heads_dim,
                    self.memory_length_dim, self.kv_dim]))
         encdec_tensors.append((q_var, o_var, k, v))
@@ -448,18 +448,18 @@ def _sample(self, features, mesh):
         partial_targets = tf.pad(
             partial_targets, [[0, hparams.batch_size - partial_targets_batch],
                               [0, hparams.max_length - partial_targets_length]])
-        partial_targets = self._infeed_to_batch_by_length(
+        partial_targets = self._import_to_batch_by_length(
             partial_targets, "partial_targets", mesh, hparams)
 
     if hparams.beam_size == 1:
-      ids_shape = mtf.TensorShape([self.batch_dim, self.length_dim])
-      kv_shape = mtf.TensorShape([self.batch_dim, self.heads_dim,
-                                  self.memory_length_dim, self.kv_dim])
+      ids_shape = mtf.Shape([self.batch_dim, self.length_dim])
+      kv_shape = mtf.Shape([self.batch_dim, self.heads_dim,
+                            self.memory_length_dim, self.kv_dim])
     else:
       beam_dim = mtf.Dimension("beam", hparams.beam_size)
-      ids_shape = mtf.TensorShape([self.batch_dim, beam_dim, self.length_dim])
-      kv_shape = mtf.TensorShape([self.batch_dim, beam_dim, self.heads_dim,
-                                  self.memory_length_dim, self.kv_dim])
+      ids_shape = mtf.Shape([self.batch_dim, beam_dim, self.length_dim])
+      kv_shape = mtf.Shape([self.batch_dim, beam_dim, self.heads_dim,
+                            self.memory_length_dim, self.kv_dim])
 
     initial_ids = mtf.constant(mesh, 0, ids_shape, dtype=tf.int32)
     initial_kv_states = (
@@ -563,7 +563,7 @@ def _decoder_layer_stack_incremental(self,
     layer_norm_combined_var = mtf.get_variable(
         x.mesh,
         "layer_norm_scale",
-        mtf.TensorShape([layer_norms_dim, self.model_dim]),
+        mtf.Shape([layer_norms_dim, self.model_dim]),
         initializer=tf.ones_initializer(),
         activation_dtype=x.dtype)
     layer_norm_vars = mtf.unstack(layer_norm_combined_var, layer_norms_dim)
@@ -614,8 +614,8 @@ def mtf_transformer_base():
   hparams.add_hparam("d_kv", 128)
   hparams.label_smoothing = 0.1
   # 8-way model-parallelism
-  hparams.add_hparam("mesh_shape", "8")
-  hparams.add_hparam("layout", "vocab:0;d_ff:0;heads:0")
+  hparams.add_hparam("mesh_shape", "model:8")
+  hparams.add_hparam("layout", "batch:batch;vocab:model;d_ff:model;heads:model")
   hparams.add_hparam("num_heads", 8)
   hparams.add_hparam("d_ff", 2048)
   hparams.add_hparam("num_encoder_layers", 6)
@@ -639,7 +639,7 @@ def mtf_transformer_base():
   # Reuse targets_embedding_var as inputs_embedding_var
   hparams.shared_embedding = True
   hparams.optimizer = "Adafactor"
-  hparams.learning_rate_schedule = "rsqrt_decay*linear_decay"
+  hparams.learning_rate_schedule = "linear_warmup*rsqrt_decay*linear_decay"
   hparams.learning_rate_warmup_steps = 10000
   hparams.activation_dtype = "float32"
 
@@ -659,22 +659,6 @@ def mtf_transformer_base():
   return hparams
 
 
-@registry.register_hparams
-def mtf_transformer_shared_embedding():
-  hparams = mtf_transformer_base()
-  hparams.shared_embedding_and_softmax_weights = False
-  hparams.shared_embedding = True
-  return hparams
-
-
-@registry.register_hparams
-def mtf_transformer_no_share():
-  hparams = mtf_transformer_base()
-  hparams.shared_embedding_and_softmax_weights = False
-  hparams.shared_embedding = False
-  return hparams
-
-
 @registry.register_hparams
 def mtf_transformer_tiny():
   """Catch bugs locally..."""
@@ -686,8 +670,7 @@ def mtf_transformer_tiny():
   hparams.num_decoder_layers = 2
   hparams.num_heads = 4
   # data parallelism and model-parallelism
-  hparams.mesh_shape = "2.2"
-  hparams.layout = "batch:0;vocab:1;d_ff:1;heads:1"
+  hparams.mesh_shape = "batch:2;model:2"
   return hparams
 
 
@@ -695,15 +678,14 @@ def mtf_transformer_tiny():
 def mtf_transformer_single():
   hparams = mtf_transformer_tiny()
   hparams.mesh_shape = ""
-  hparams.layout = ""
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_tiny_moe():
   hparams = mtf_transformer_tiny()
-  hparams.mesh_shape = "4"
-  hparams.layout = "batch:0,experts:0"
+  hparams.mesh_shape = "all:4"
+  hparams.layout = "batch:all;experts:all"
   hparams.feedforward_layer = "moe"
   return hparams
 
@@ -711,24 +693,7 @@ def mtf_transformer_tiny_moe():
 @registry.register_hparams
 def mtf_transformer_tiny_8gpu():
   hparams = mtf_transformer_tiny()
-  hparams.mesh_shape = "8"
-  hparams.layout = "vocab:0;d_ff:0;heads:0"
-  return hparams
-
-
-@registry.register_hparams
-def mtf_transformer_length_sharded():
-  hparams = mtf_transformer_tiny()
-  hparams.mesh_shape = "2"
-  hparams.layout = "length:0"
-  return hparams
-
-
-@registry.register_hparams
-def mtf_transformer_lm():
-  """Set of hyperparameters."""
-  hparams = mtf_transformer_base()
-  hparams.label_smoothing = 0.0
+  hparams.mesh_shape = "model:8"
   return hparams
 
 
@@ -737,85 +702,68 @@ def mtf_transformer_paper_lm(sz):
   n = 2 ** sz
   hparams = mtf_transformer_base()
   hparams.label_smoothing = 0.0
-  hparams.batch_size = 128
+  hparams.batch_size = 256
   hparams.d_model = 1024
   hparams.d_ff = int(8192 * n)
+  hparams.d_kv = 256
   hparams.num_heads = int(8 * n)
   hparams.shared_embedding_and_softmax_weights = False
-  hparams.learning_rate_decay_steps = 27300  # one epoch for lm1b
+  # one epoch for languagemodel_lm1b32k_packed = 13600 steps
+  hparams.learning_rate_decay_steps = 13600
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_paper_lm_m1():
   hparams = mtf_transformer_paper_lm(-1)
-  hparams.mesh_shape = "32"
-  hparams.layout = "batch:0"
+  hparams.mesh_shape = "batch:32"
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_paper_lm_0():
   hparams = mtf_transformer_paper_lm(0)
-  hparams.mesh_shape = "32"
-  hparams.layout = "batch:0"
+  hparams.mesh_shape = "batch:32"
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_paper_lm_1():
   hparams = mtf_transformer_paper_lm(1)
-  hparams.mesh_shape = "4;8"
-  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  hparams.mesh_shape = "model:4;batch:8"
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_paper_lm_2():
   hparams = mtf_transformer_paper_lm(2)
-  hparams.mesh_shape = "4;8"
-  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  hparams.mesh_shape = "model:4;batch:8"
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_paper_lm_3():
   hparams = mtf_transformer_paper_lm(3)
-  hparams.mesh_shape = "8;16"
-  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  hparams.mesh_shape = "model:8;batch:16"
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_paper_lm_4():
   hparams = mtf_transformer_paper_lm(4)
-  hparams.mesh_shape = "8;16"
-  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  hparams.mesh_shape = "batch:16;model:32"
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_paper_lm_5():
   hparams = mtf_transformer_paper_lm(5)
-  hparams.mesh_shape = "8;16"
-  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  hparams.mesh_shape = "batch:16;model:32"
   return hparams
 
 
 def mtf_transformer_paper_tr(sz):
-  """Config for translation experiments.
-
-  translate_enfr_wmt32k_packed - tokens=2.385e9
-  batch:128 * length:256 = 32768 tokens/step
-  steps/epoch = 72800
-  Let's run for 3 epochs (218K steps)
-  Leaerning-rate-decay for last epoch
-
-  Args:
-    sz: an integer
-  Returns:
-    hyperparameters
-  """
+  """Config for translation experiments."""
   n = 2 ** sz
   hparams = mtf_transformer_base()
   hparams.label_smoothing = 0.1
@@ -824,58 +772,51 @@ def mtf_transformer_paper_tr(sz):
   hparams.d_ff = int(4096 * n)
   hparams.num_heads = int(8 * n)
   hparams.shared_embedding_and_softmax_weights = False
-  hparams.learning_rate_decay_steps = 51400  # one epoch for enfr
+  # one epoch for translate_enfr_wmt32k_packed = 51400 steps
+  hparams.learning_rate_decay_steps = 51400
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_paper_tr_m1():
   hparams = mtf_transformer_paper_tr(-1)
-  hparams.mesh_shape = "32"
-  hparams.layout = "batch:0"
+  hparams.mesh_shape = "batch:32"
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_paper_tr_0():
   hparams = mtf_transformer_paper_tr(0)
-  hparams.mesh_shape = "32"
-  hparams.layout = "batch:0"
+  hparams.mesh_shape = "batch:32"
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_paper_tr_1():
   hparams = mtf_transformer_paper_tr(1)
-  hparams.mesh_shape = "4;8"
-  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  hparams.mesh_shape = "model:4;batch:8"
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_paper_tr_2():
   hparams = mtf_transformer_paper_tr(2)
-  hparams.mesh_shape = "4;8"
-  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  hparams.mesh_shape = "model:4;batch:8"
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_paper_tr_3():
   hparams = mtf_transformer_paper_tr(3)
-  hparams.mesh_shape = "8;16"
-  hparams.layout = "batch:1;vocab:0;d_ff:0;heads:0"
+  hparams.mesh_shape = "model:8;batch:16"
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_paper_tr_4():
-  return mtf_transformer_paper_tr(4)
-
-
-@registry.register_hparams
-def mtf_transformer_paper_tr_5():
-  return mtf_transformer_paper_tr(5)
+  hparams = mtf_transformer_paper_tr(4)
+  hparams.mesh_shape = "model:8;batch:16"
+  return hparams
 
 
 @registry.register_hparams
@@ -890,6 +831,7 @@ def mtf_transformer_lm_moe():
   hparams.attention_value_channels = 1024
   hparams.shared_embedding_and_softmax_weights = False
   hparams.num_heads = 8
-  hparams.layout = "batch:0,experts:0"
+  hparams.mesh_shape = "all:8"
+  hparams.layout = "batch:all;experts:all"
   hparams.feedforward_layer = "moe"
   return hparams
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer_compat.py b/tensor2tensor/mesh_tensorflow/mtf_transformer_compat.py
deleted file mode 100644
index 2e48f9de7..000000000
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer_compat.py
+++ /dev/null
@@ -1,926 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-"""Temporary hack for decoding mtf_transformer models.
-
-This is a transformer implementation in regular TensorFlow which is
-checkpoint-compatible with MtfTransformer for eval/inference.
-
-The purpose of this model is to run inference on MtfTransformer models.
-We are working on native decoding in MtfTransformer which will be faster and
-cleaner.
-
-TODO(noam): Remove once we can decode in mtf.
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-from six.moves import range  # pylint: disable=redefined-builtin
-
-from tensor2tensor.layers import common_attention
-from tensor2tensor.layers import common_layers
-from tensor2tensor.utils import beam_search
-from tensor2tensor.utils import registry
-from tensor2tensor.utils import t2t_model
-
-import tensorflow as tf
-
-from tensorflow.python.util import nest
-
-
-@registry.register_model
-class MtfTransformerCompat(t2t_model.T2TModel):
-  """Attention net.  See file docstring."""
-
-  def __init__(self, *args, **kwargs):
-    with tf.variable_scope("transformer"):
-      self._top_scope = tf.get_variable_scope()
-    kwargs["_scope"] = "transformer"
-    super(MtfTransformerCompat, self).__init__(*args, **kwargs)
-    self._name = "transformer"
-    self._base_name = "transformer"
-
-  @property
-  def _targets_vocab_size(self):
-    targets_vocab_size = self._problem_hparams.vocabulary["targets"].vocab_size
-    targets_vocab_size += (-targets_vocab_size) % self._hparams.vocab_divisor
-    return targets_vocab_size
-
-  @property
-  def _inputs_vocab_size(self):
-    if not self.has_input:
-      return None
-    inputs_vocab_size = self._problem_hparams.vocabulary["inputs"].vocab_size
-    inputs_vocab_size += (-inputs_vocab_size) % self._hparams.vocab_divisor
-    return inputs_vocab_size
-
-  @property
-  def _embedding_and_softmax_var_names(self):
-    """Figure out the variable names for the embedding and softmax variables.
-
-    Equality between the returned names means that we should share the
-    variables.
-
-    Returns:
-      inputs_embedding_name: a string or None
-      targets_embedding_name: a string
-      softmax_var_name: a string
-    Raises:
-      ValueError: if we try to share embeddings with different vocab sizes.
-    """
-    hparams = self._hparams
-    inputs_embedding_name = "input_emb"
-    targets_embedding_name = "target_emb"
-    softmax_var_name = "softmax"
-    if (self.has_input and
-        (hparams.shared_embedding or
-         hparams.shared_embedding_and_softmax_weights) and
-        self._inputs_vocab_size != self._targets_vocab_size):
-      raise ValueError(
-          "hparams.shared_embedding_and_softmax_weights "
-          " or hparams.shared_embedding require "
-          "that input and target vocabulary sizes be equal %s vs %s"
-          % (self._inputs_vocab_size, self._targets_vocab_size))
-    if hparams.shared_embedding_and_softmax_weights:
-      inputs_embedding_name = "shared"
-      targets_embedding_name = "shared"
-      softmax_var_name = "shared"
-    elif hparams.shared_embedding:
-      inputs_embedding_name = "shared"
-      targets_embedding_name = "shared"
-    targets_embedding_name = (
-        "symbol_modality_%d_%d/%s/weights_0" %
-        (self._targets_vocab_size, hparams.d_model, targets_embedding_name))
-    softmax_var_name = (
-        "symbol_modality_%d_%d/%s/weights_0" %
-        (self._targets_vocab_size, hparams.d_model, softmax_var_name))
-    if self.has_input:
-      inputs_embedding_name = (
-          "symbol_modality_%d_%d/%s/weights_0" %
-          (self._inputs_vocab_size, hparams.d_model, inputs_embedding_name))
-    else:
-      inputs_embedding_name = None
-    return inputs_embedding_name, targets_embedding_name, softmax_var_name
-
-  @property
-  def _get_targets_emb_var(self):
-    with tf.variable_scope(self._top_scope, reuse=tf.AUTO_REUSE):
-      return tf.get_variable(
-          "targets_embedding",
-          [self._targets_vocab_size, self._hparams.d_model])
-
-  @property
-  def _get_inputs_emb_var(self):
-    if self._hparams.shared_embedding:
-      return self._get_targets_emb_var
-    with tf.variable_scope(self._top_scope, reuse=tf.AUTO_REUSE):
-      return tf.get_variable(
-          "inputs_embedding",
-          [self._inputs_vocab_size, self._hparams.d_model])
-
-  @property
-  def _get_softmax_var(self):
-    if self._hparams.shared_embedding_and_softmax_weights:
-      return self._get_targets_emb_var * (self._hparams.d_model ** -0.5)
-    with tf.variable_scope(self._top_scope, reuse=tf.AUTO_REUSE):
-      return tf.get_variable(
-          "softmax",
-          [self._targets_vocab_size, self._hparams.d_model])
-
-  def encode(self, inputs, hparams, features=None, losses=None):
-    """Encode transformer inputs.
-
-    Args:
-      inputs: Transformer inputs [batch_size, input_length, input_height,
-        hidden_dim] which will be flattened along the two spatial dimensions.
-      hparams: hyperparameters for model.
-      features: optionally pass the entire features dictionary as well.
-        This is needed now for "packed" datasets.
-      losses: optional list onto which to append extra training losses
-
-    Returns:
-      Tuple of:
-          encoder_output: Encoder representation.
-              [batch_size, input_length, hidden_dim]
-          encoder_decoder_attention_bias: Bias and mask weights for
-              encoder-decoder attention. [batch_size, input_length]
-    """
-    tf.logging.info("Encode inputs = %s" % inputs)
-    encoder_input, self_attention_bias, encoder_decoder_attention_bias = (
-        transformer_prepare_encoder(
-            self._get_inputs_emb_var, inputs, hparams, features=features))
-
-    encoder_input = tf.nn.dropout(encoder_input,
-                                  1.0 - hparams.layer_prepostprocess_dropout)
-
-    encoder_output = transformer_encoder(
-        encoder_input,
-        self_attention_bias,
-        hparams,
-        losses=losses)
-
-    return encoder_output, encoder_decoder_attention_bias
-
-  def decode(self,
-             decoder_input,
-             encoder_output,
-             encoder_decoder_attention_bias,
-             decoder_self_attention_bias,
-             hparams,
-             cache=None,
-             losses=None):
-    """Decode Transformer outputs from encoder representation.
-
-    Args:
-      decoder_input: inputs to bottom of the model.
-          [batch_size, decoder_length, hidden_dim]
-      encoder_output: Encoder representation.
-          [batch_size, input_length, hidden_dim]
-      encoder_decoder_attention_bias: Bias and mask weights for
-          encoder-decoder attention. [batch_size, input_length]
-      decoder_self_attention_bias: Bias and mask weights for decoder
-          self-attention. [batch_size, decoder_length]
-      hparams: hyperparameters for model.
-      cache: dict, containing tensors which are the results of previous
-          attentions, used for fast decoding.
-      losses: optional list onto which to append extra training losses
-
-    Returns:
-      Final decoder representation. [batch_size, decoder_length, hidden_dim]
-    """
-    decoder_input = tf.nn.dropout(decoder_input,
-                                  1.0 - hparams.layer_prepostprocess_dropout)
-
-    decoder_output = transformer_decoder(
-        decoder_input,
-        encoder_output,
-        decoder_self_attention_bias,
-        encoder_decoder_attention_bias,
-        hparams,
-        cache=cache,
-        losses=losses)
-
-    ret = tf.tensordot(decoder_output, self._get_softmax_var, axes=[[-1], [1]])
-    ret = tf.expand_dims(tf.expand_dims(ret, 2), 3)
-    return ret
-
-  def body(self, features):
-    """Transformer main model_fn.
-
-    Args:
-      features: Map of features to the model. Should contain the following:
-          "inputs": Transformer inputs [batch_size, input_length, hidden_dim]
-          "targets": Target decoder outputs.
-              [batch_size, decoder_length, hidden_dim]
-
-    Returns:
-      Final decoder representation. [batch_size, decoder_length, hidden_dim]
-    """
-    with tf.variable_scope(self._top_scope):
-      hparams = self._hparams
-      losses = []
-
-      if self.has_input:
-        inputs = tf.squeeze(features["inputs_raw"], (2, 3))
-        encoder_output, encoder_decoder_attention_bias = self.encode(
-            inputs, hparams, features=features, losses=losses)
-      else:
-        encoder_output, encoder_decoder_attention_bias = (None, None)
-
-      targets = tf.squeeze(features["targets_raw"], (2, 3))
-      decoder_input, decoder_self_attention_bias = transformer_prepare_decoder(
-          self._get_targets_emb_var, targets, hparams, features=features)
-
-      decoder_output = self.decode(
-          decoder_input,
-          encoder_output,
-          encoder_decoder_attention_bias,
-          decoder_self_attention_bias,
-          hparams,
-          losses=losses)
-
-      if losses:
-        return decoder_output, {"extra_loss": tf.add_n(losses)}
-      else:
-        return decoder_output
-
-  def _greedy_infer(self, features, decode_length, use_tpu=False):
-    """Fast version of greedy decoding.
-
-    Args:
-      features: an map of string to `Tensor`
-      decode_length: an integer.  How many additional timesteps to decode.
-      use_tpu: a boolean
-
-    Returns:
-      A dict of decoding results {
-          "outputs": integer `Tensor` of decoded ids of shape
-              [batch_size, <= decode_length] if beam_size == 1 or
-              [batch_size, top_beams, <= decode_length]
-          "scores": decoding log probs from the beam search,
-              None if using greedy decoding (beam_size=1)
-      }
-
-    Raises:
-      NotImplementedError: If there are multiple data shards.
-    """
-    with tf.variable_scope(self.name):
-      return  self._fast_decode(features, decode_length)
-
-  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
-    """Beam search decoding.
-
-    Args:
-      features: an map of string to `Tensor`
-      decode_length: an integer.  How many additional timesteps to decode.
-      beam_size: number of beams.
-      top_beams: an integer. How many of the beams to return.
-      alpha: Float that controls the length penalty. larger the alpha, stronger
-        the preference for longer translations.
-
-    Returns:
-      A dict of decoding results {
-          "outputs": integer `Tensor` of decoded ids of shape
-              [batch_size, <= decode_length] if beam_size == 1 or
-              [batch_size, top_beams, <= decode_length]
-          "scores": decoding log probs from the beam search,
-              None if using greedy decoding (beam_size=1)
-      }
-    """
-    with tf.variable_scope(self.name):
-      return self._fast_decode(features, decode_length, beam_size, top_beams,
-                               alpha)
-
-  def _fast_decode(self,
-                   features,
-                   decode_length,
-                   beam_size=1,
-                   top_beams=1,
-                   alpha=1.0):
-    """Fast decoding.
-
-    Implements both greedy and beam search decoding, uses beam search iff
-    beam_size > 1, otherwise beam search related arguments are ignored.
-
-    Args:
-      features: a map of string to model  features.
-      decode_length: an integer.  How many additional timesteps to decode.
-      beam_size: number of beams.
-      top_beams: an integer. How many of the beams to return.
-      alpha: Float that controls the length penalty. larger the alpha, stronger
-        the preference for longer translations.
-
-    Returns:
-      A dict of decoding results {
-          "outputs": integer `Tensor` of decoded ids of shape
-              [batch_size, <= decode_length] if beam_size == 1 or
-              [batch_size, top_beams, <= decode_length]
-          "scores": decoding log probs from the beam search,
-              None if using greedy decoding (beam_size=1)
-      }
-
-    Raises:
-      NotImplementedError: If there are multiple data shards.
-    """
-    if self._num_datashards != 1:
-      raise NotImplementedError("Fast decoding only supports a single shard.")
-    dp = self._data_parallelism
-    hparams = self._hparams
-    target_modality = self._problem_hparams.target_modality
-    if "targets_segmentation" in features:
-      raise NotImplementedError(
-          "Decoding not supported on packed datasets "
-          " If you want to decode from a dataset, use the non-packed version"
-          " of the dataset when decoding.")
-    if self.has_input:
-      inputs = features["inputs"]
-      if target_modality.is_class_modality:
-        decode_length = 1
-      else:
-        decode_length = (
-            common_layers.shape_list(inputs)[1] + features.get(
-                "decode_length", decode_length))
-
-      # TODO(llion): Clean up this reshaping logic.
-      inputs = tf.expand_dims(inputs, axis=1)
-      if len(inputs.shape) < 5:
-        inputs = tf.expand_dims(inputs, axis=4)
-      s = common_layers.shape_list(inputs)
-      batch_size = s[0]
-      inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
-      inputs = tf.squeeze(inputs, (2, 3))
-      # _shard_features called to ensure that the variable names match
-      inputs = self._shard_features({"inputs": inputs})["inputs"]
-
-      # input_modality = self._problem_hparams.input_modality["inputs"]
-      # with tf.variable_scope(input_modality.name):
-      #   inputs = input_modality.bottom_sharded(inputs, dp)
-      encoder_output, encoder_decoder_attention_bias = dp(
-          self.encode,
-          inputs,
-          hparams,
-          features=features)
-      encoder_output = encoder_output[0]
-      encoder_decoder_attention_bias = encoder_decoder_attention_bias[0]
-      partial_targets = None
-    else:
-      # The problem has no inputs.
-      encoder_output = None
-      encoder_decoder_attention_bias = None
-
-      # Prepare partial targets.
-      # In either features["inputs"] or features["targets"].
-      # We force the outputs to begin with these sequences.
-      partial_targets = features.get("inputs")
-      if partial_targets is None:
-        partial_targets = features["targets"]
-      assert partial_targets is not None
-      partial_targets = common_layers.expand_squeeze_to_nd(partial_targets, 2)
-      partial_targets = tf.to_int64(partial_targets)
-      partial_targets_shape = common_layers.shape_list(partial_targets)
-      partial_targets_length = partial_targets_shape[1]
-      decode_length = (
-          partial_targets_length + features.get("decode_length", decode_length))
-      batch_size = partial_targets_shape[0]
-
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      positional_encoding = common_attention.add_positional_embedding(
-          tf.zeros([1, decode_length, hparams.d_model]),
-          hparams.max_length, "positional_embedding", None)
-
-    def preprocess_targets(targets, i):
-      """Performs preprocessing steps on the targets to prepare for the decoder.
-
-      This includes:
-        - Embedding the ids.
-        - Flattening to 3D tensor.
-        - Optionally adding timing signals.
-
-      Args:
-        targets: inputs ids to the decoder. [batch_size, 1]
-        i: scalar, Step number of the decoding loop.
-
-      Returns:
-        Processed targets [batch_size, 1, hidden_dim]
-      """
-      targets_emb_var = self._get_targets_emb_var
-      targets = tf.gather(targets_emb_var, targets)
-      tf.logging.info("targets = %s" % targets)
-      targets = tf.squeeze(targets, (2, 3))
-      if positional_encoding is not None:
-        targets += positional_encoding[:, i:i + 1]
-      return targets
-
-    def symbols_to_logits_fn(ids, i, cache):
-      """Go from ids to logits for next symbol."""
-      ids = ids[:, -1:]
-      targets = tf.expand_dims(tf.expand_dims(ids, axis=2), axis=3)
-      targets = preprocess_targets(targets, i)
-
-      bias = None  # decoder_self_attention_bias[:, :, i:i + 1, :i + 1]
-
-      body_outputs = dp(
-          self.decode,
-          targets,
-          cache.get("encoder_output"),
-          cache.get("encoder_decoder_attention_bias"),
-          bias,
-          hparams,
-          cache)
-
-      logits = body_outputs[0]
-      # with tf.variable_scope(target_modality.name):
-      #   logits = target_modality.top_sharded(body_outputs, None, dp)[0]
-
-      ret = tf.squeeze(logits, axis=[1, 2, 3])
-      if partial_targets is not None:
-        # If the position is within the given partial targets, we alter the
-        # logits to always return those values.
-        # A faster approach would be to process the partial targets in one
-        # iteration in order to fill the corresponding parts of the cache.
-        # This would require broader changes, though.
-        vocab_size = tf.shape(ret)[1]
-
-        def forced_logits():
-          return tf.one_hot(
-              tf.tile(partial_targets[:, i], [beam_size]), vocab_size, 0.0,
-              -1e9)
-
-        ret = tf.cond(
-            tf.less(i, partial_targets_length), forced_logits, lambda: ret)
-      return ret, cache
-
-    ret = fast_decode(
-        encoder_output=encoder_output,
-        encoder_decoder_attention_bias=encoder_decoder_attention_bias,
-        symbols_to_logits_fn=symbols_to_logits_fn,
-        hparams=hparams,
-        decode_length=decode_length,
-        vocab_size=target_modality.top_dimensionality,
-        beam_size=beam_size,
-        top_beams=top_beams,
-        alpha=alpha,
-        batch_size=batch_size,
-        force_decode_length=self._decode_hparams.force_decode_length)
-    if partial_targets is not None:
-      if beam_size <= 1 or top_beams <= 1:
-        ret["outputs"] = ret["outputs"][:, partial_targets_length:]
-      else:
-        ret["outputs"] = ret["outputs"][:, :, partial_targets_length:]
-    return ret
-
-
-def fast_decode(encoder_output,
-                encoder_decoder_attention_bias,
-                symbols_to_logits_fn,
-                hparams,
-                decode_length,
-                vocab_size,
-                beam_size=1,
-                top_beams=1,
-                alpha=1.0,
-                eos_id=beam_search.EOS_ID,
-                batch_size=None,
-                force_decode_length=False):
-  """Given encoder output and a symbols to logits function, does fast decoding.
-
-  Implements both greedy and beam search decoding, uses beam search iff
-  beam_size > 1, otherwise beam search related arguments are ignored.
-
-  Args:
-    encoder_output: Output from encoder.
-    encoder_decoder_attention_bias: a bias tensor for use in encoder-decoder
-      attention
-    symbols_to_logits_fn: Incremental decoding; function mapping triple
-      `(ids, step, cache)` to symbol logits.
-    hparams: run hyperparameters
-    decode_length: an integer.  How many additional timesteps to decode.
-    vocab_size: Output vocabulary size.
-    beam_size: number of beams.
-    top_beams: an integer. How many of the beams to return.
-    alpha: Float that controls the length penalty. larger the alpha, stronger
-      the preference for longer translations.
-    eos_id: End-of-sequence symbol in beam search.
-    batch_size: an integer scalar - must be passed if there is no input
-    force_decode_length: bool, whether to force the full decode length, or if
-      False, stop when all beams hit eos_id.
-
-  Returns:
-      A dict of decoding results {
-          "outputs": integer `Tensor` of decoded ids of shape
-              [batch_size, <= decode_length] if top_beams == 1 or
-              [batch_size, top_beams, <= decode_length] otherwise
-          "scores": decoding log probs from the beam search,
-              None if using greedy decoding (beam_size=1)
-      }
-
-    Raises:
-      NotImplementedError: If beam size > 1 with partial targets.
-  """
-  if encoder_output is not None:
-    batch_size = common_layers.shape_list(encoder_output)[0]
-
-  num_layers = hparams.num_decoder_layers
-  cache = {
-      "layer_%d" % layer: {
-          "k": tf.zeros([batch_size, hparams.num_heads,
-                         0, hparams.d_kv]),
-          "v": tf.zeros([batch_size, hparams.num_heads,
-                         0, hparams.d_kv]),
-      } for layer in range(num_layers)
-  }
-
-  if encoder_output is not None:
-    for layer in range(num_layers):
-      layer_name = "layer_%d" % layer
-      with tf.variable_scope("decoder/%s" % layer_name):
-        k_encdec, v_encdec = multihead_attention_compat(
-            None,
-            encoder_output,
-            None,
-            hparams.d_kv,
-            hparams.num_heads,
-            name="encdec_attention")
-      cache[layer_name]["k_encdec"] = k_encdec
-      cache[layer_name]["v_encdec"] = v_encdec
-
-    cache["encoder_output"] = encoder_output
-    cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias
-
-  if beam_size > 1:  # Beam Search
-    initial_ids = tf.zeros([batch_size], dtype=tf.int32)
-    decoded_ids, scores = beam_search.beam_search(
-        symbols_to_logits_fn,
-        initial_ids,
-        beam_size,
-        decode_length,
-        vocab_size,
-        alpha,
-        states=cache,
-        eos_id=eos_id,
-        stop_early=(top_beams == 1))
-
-    if top_beams == 1:
-      decoded_ids = decoded_ids[:, 0, 1:]
-      scores = scores[:, 0]
-    else:
-      decoded_ids = decoded_ids[:, :top_beams, 1:]
-      scores = scores[:, :top_beams]
-  else:  # Greedy
-
-    def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
-      """One step of greedy decoding."""
-      logits, cache = symbols_to_logits_fn(next_id, i, cache)
-      log_probs = common_layers.log_prob_from_logits(logits)
-      temperature = (0.0 if hparams.sampling_method == "argmax" else
-                     hparams.sampling_temp)
-      next_id = common_layers.sample_with_temperature(logits, temperature)
-      hit_eos |= tf.equal(next_id, eos_id)
-
-      log_prob_indices = tf.stack(
-          [tf.range(tf.to_int64(batch_size)), next_id], axis=1)
-      log_prob += tf.gather_nd(log_probs, log_prob_indices)
-
-      next_id = tf.expand_dims(next_id, axis=1)
-      decoded_ids = tf.concat([decoded_ids, next_id], axis=1)
-      return i + 1, hit_eos, next_id, decoded_ids, cache, log_prob
-
-    def is_not_finished(i, hit_eos, *_):
-      finished = i >= decode_length
-      if not force_decode_length:
-        finished |= tf.reduce_all(hit_eos)
-      return tf.logical_not(finished)
-
-    decoded_ids = tf.zeros([batch_size, 0], dtype=tf.int64)
-    hit_eos = tf.fill([batch_size], False)
-    next_id = tf.zeros([batch_size, 1], dtype=tf.int64)
-    initial_log_prob = tf.zeros([batch_size], dtype=tf.float32)
-    _, _, _, decoded_ids, _, log_prob = tf.while_loop(
-        is_not_finished,
-        inner_loop, [
-            tf.constant(0), hit_eos, next_id, decoded_ids, cache,
-            initial_log_prob
-        ],
-        shape_invariants=[
-            tf.TensorShape([]),
-            tf.TensorShape([None]),
-            tf.TensorShape([None, None]),
-            tf.TensorShape([None, None]),
-            nest.map_structure(beam_search.get_state_shape_invariants, cache),
-            tf.TensorShape([None]),
-        ])
-    scores = log_prob
-
-  return {"outputs": decoded_ids, "scores": scores}
-
-
-def transformer_prepare_encoder(
-    inputs_emb_var, inputs, hparams, features=None):
-  """Prepare one shard of the model for the encoder.
-
-  Args:
-    inputs_emb_var: a Tensor
-    inputs: a Tensor.
-    hparams: run hyperparameters
-    features: optionally pass the entire features dictionary as well.
-      This is needed now for "packed" datasets.
-
-  Returns:
-    encoder_input: a Tensor, bottom of encoder stack
-    encoder_self_attention_bias: a bias tensor for use in encoder self-attention
-    encoder_decoder_attention_bias: a bias tensor for use in encoder-decoder
-      attention
-  """
-  encoder_input = tf.gather(inputs_emb_var, inputs)
-
-  if features and "inputs_segmentation" in features:
-    # Packed dataset.  Keep the examples from seeing each other.
-    inputs_segmentation = features["inputs_segmentation"]
-    inputs_position = features["inputs_position"]
-    targets_segmentation = features["targets_segmentation"]
-    encoder_self_attention_bias = common_attention.attention_bias_same_segment(
-        inputs_segmentation, inputs_segmentation)
-    encoder_decoder_attention_bias = (
-        common_attention.attention_bias_same_segment(targets_segmentation,
-                                                     inputs_segmentation))
-  else:
-    # Usual case - not a packed dataset.
-    encoder_padding = tf.to_float(tf.equal(inputs, 0))
-    ignore_padding = common_attention.attention_bias_ignore_padding(
-        encoder_padding)
-    encoder_self_attention_bias = ignore_padding
-    encoder_decoder_attention_bias = ignore_padding
-    inputs_position = None
-  with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-    encoder_input = common_attention.add_positional_embedding(
-        encoder_input, hparams.max_length, "positional_embedding",
-        inputs_position)
-  if hparams.activation_dtype == "bfloat16":
-    encoder_self_attention_bias = tf.cast(encoder_self_attention_bias,
-                                          tf.bfloat16)
-    encoder_decoder_attention_bias = tf.cast(encoder_decoder_attention_bias,
-                                             tf.bfloat16)
-  return (encoder_input, encoder_self_attention_bias,
-          encoder_decoder_attention_bias)
-
-
-def transformer_prepare_decoder(
-    targets_emb_var, targets, hparams, features=None):
-  """Prepare one shard of the model for the decoder.
-
-  Args:
-    targets_emb_var: a Tensor
-    targets: a Tensor.
-    hparams: run hyperparameters
-    features: optionally pass the entire features dictionary as well.
-      This is needed now for "packed" datasets.
-
-  Returns:
-    decoder_input: a Tensor, bottom of decoder stack
-    decoder_self_attention_bias: a bias tensor for use in decoder self-attention
-  """
-  decoder_self_attention_bias = (
-      common_attention.attention_bias_lower_triangle(
-          common_layers.shape_list(targets)[1]))
-
-  if features and "targets_segmentation" in features:
-    # "Packed" dataset - keep the examples from seeing each other.
-    targets_segmentation = features["targets_segmentation"]
-    targets_position = features["targets_position"]
-    decoder_self_attention_bias += common_attention.attention_bias_same_segment(
-        targets_segmentation, targets_segmentation)
-  else:
-    targets_position = None
-  decoder_input = tf.gather(
-      targets_emb_var, common_layers.shift_right_2d(targets))
-  with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-    decoder_input = common_attention.add_positional_embedding(
-        decoder_input, hparams.max_length, "positional_embedding",
-        targets_position)
-
-  if hparams.activation_dtype == "bfloat16":
-    decoder_self_attention_bias = tf.cast(decoder_self_attention_bias,
-                                          tf.bfloat16)
-  return (decoder_input, decoder_self_attention_bias)
-
-
-def transformer_encoder(encoder_input,
-                        encoder_self_attention_bias,
-                        hparams,
-                        name="encoder",
-                        losses=None):
-  """A stack of transformer layers.
-
-  Args:
-    encoder_input: a Tensor
-    encoder_self_attention_bias: bias Tensor for self-attention
-       (see common_attention.attention_bias())
-    hparams: hyperparameters for model
-    name: a string
-    losses: optional list onto which to append extra training losses
-
-  Returns:
-    y: a Tensors
-  """
-  x = encoder_input
-  with tf.variable_scope(name):
-    num_layer_norms = hparams.num_encoder_layers * 2 + 1
-    layer_norm_combined_var = tf.get_variable(
-        "layer_norm_scale", [num_layer_norms, hparams.d_model])
-    layer_norm_vars = tf.unstack(layer_norm_combined_var, num_layer_norms)
-    def normalize(x):
-      scale = layer_norm_vars.pop(0)
-      variance = tf.reduce_mean(tf.square(x), -1, keep_dims=True)
-      return x * tf.rsqrt(variance + hparams.norm_epsilon) * scale
-    for layer in range(hparams.num_encoder_layers):
-      with tf.variable_scope("layer_%d" % layer):
-        x += multihead_attention_compat(
-            normalize(x),
-            None,
-            encoder_self_attention_bias,
-            kv_channels=hparams.d_kv,
-            heads=hparams.num_heads,
-            name="self_attention")
-        x += transformer_feedforward_layer(normalize(x), hparams, losses=losses)
-    x = normalize(x)
-    return x
-
-
-def transformer_decoder(decoder_input,
-                        encoder_output,
-                        decoder_self_attention_bias,
-                        encoder_decoder_attention_bias,
-                        hparams,
-                        cache=None,
-                        name="decoder",
-                        losses=None):
-  """A stack of transformer layers.
-
-  Args:
-    decoder_input: a Tensor
-    encoder_output: a Tensor
-    decoder_self_attention_bias: bias Tensor for self-attention
-      (see common_attention.attention_bias())
-    encoder_decoder_attention_bias: bias Tensor for encoder-decoder attention
-      (see common_attention.attention_bias())
-    hparams: hyperparameters for model
-    cache: dict, containing tensors which are the results of previous
-        attentions, used for fast decoding.
-    name: a string
-    losses: optional list onto which to append extra training losses
-
-  Returns:
-    y: a Tensors
-  """
-  x = decoder_input
-  with tf.variable_scope(name):
-    num_layer_norms = (
-        hparams.num_decoder_layers * (2 if encoder_output is None else 3) + 1)
-    layer_norm_combined_var = tf.get_variable(
-        "layer_norm_scale", [num_layer_norms, hparams.d_model])
-    layer_norm_vars = tf.unstack(layer_norm_combined_var, num_layer_norms)
-    def normalize(x):
-      scale = layer_norm_vars.pop(0)
-      variance = tf.reduce_mean(tf.square(x), -1, keep_dims=True)
-      return x * tf.rsqrt(variance + hparams.norm_epsilon) * scale
-    for layer in range(hparams.num_decoder_layers):
-      layer_name = "layer_%d" % layer
-      layer_cache = cache[layer_name] if cache is not None else None
-      with tf.variable_scope(layer_name):
-        x += multihead_attention_compat(
-            normalize(x),
-            None,
-            decoder_self_attention_bias,
-            kv_channels=hparams.d_kv,
-            heads=hparams.num_heads,
-            cache=layer_cache,
-            name="self_attention")
-        if encoder_output is not None:
-          x += multihead_attention_compat(
-              normalize(x),
-              encoder_output,
-              encoder_decoder_attention_bias,
-              kv_channels=hparams.d_kv,
-              heads=hparams.num_heads,
-              cache=layer_cache,
-              name="encdec_attention")
-        x += transformer_feedforward_layer(normalize(x), hparams, losses=losses)
-    x = normalize(x)
-    return x
-
-
-def transformer_feedforward_layer(x, hparams, losses=None):
-  """Feed-forward layer in the transformer.
-
-  Args:
-    x: a Tensor of shape [batch_size, length, hparams.d_model]
-    hparams: hyperparameters for model
-    losses: an optional list
-
-  Returns:
-    a Tensor of shape [batch_size, length, hparams.d_model]
-
-  Raises:
-    ValueError: If losses arg is None, but layer generates extra losses.
-  """
-  del losses
-  feedforward_layer = hparams.feedforward_layer
-  if feedforward_layer == "dense_relu_dense":
-    return dense_relu_dense_compat(x, hparams.d_ff)
-  else:
-    raise ValueError("Unknown hparams.feedforward_layer = %s"
-                     % hparams.feedforward_layer)
-
-
-def dense_relu_dense_compat(x, filter_depth, name=None):
-  """Hidden layer with RELU activation followed by linear projection.
-
-  Args:
-    x: a Tensor
-    filter_depth: integer
-    name: an optional string
-
-  Returns:
-    a tf.Tensor
-  """
-  with tf.variable_scope(name, default_name="dense_relu_dense"):
-    io_channels = x.shape.as_list()[-1]
-    w = tf.get_variable("kernel", [2, io_channels, filter_depth])
-    wi, wo = tf.unstack(w, num=2, axis=0)
-    h = tf.nn.relu(tf.tensordot(x, wi, axes=[[-1], [0]]))
-    return tf.tensordot(h, wo, axes=[[-1], [1]])
-
-
-def multihead_attention_compat(query_antecedent,
-                               memory_antecedent,
-                               mask,
-                               kv_channels,
-                               heads,
-                               cache=None,
-                               name="multihead_attention"):
-  """Multihead scaled-dot-product attention with input/output transformations.
-
-  In order to use only one variable containing the four weight matrices
-  packed together, we insist that the query and memory antecedents have the
-  same dimensionality (io_channels) and that the keys and values have the
-  same dimensionality (kv_channels).
-
-  Args:
-    query_antecedent: a Tensor with shape [batch, query_length, io_channels]
-    memory_antecedent: a Tensor with shape
-      [batch, memory_length, io_channels] (optional)
-    mask: mask Tensor (see attention_mask())
-    kv_channels: integer
-    heads: integer
-    cache: an optional dict
-    name: an optional string.
-
-  Returns:
-    A Tensor with shape [batch, qlen, io_channels]
-
-  Raises:
-    ValueError: if the dimensions do not match.
-  """
-  memory_or_query_antecedent = (
-      memory_antecedent if memory_antecedent is not None
-      else query_antecedent)
-  io_channels = memory_or_query_antecedent.shape.as_list()[-1]
-  with tf.variable_scope(name,
-                         default_name="multihead_attention",
-                         values=[query_antecedent, memory_antecedent],
-                         reuse=tf.AUTO_REUSE):
-    var = tf.get_variable("qkvo", [4, heads, io_channels, kv_channels])
-    q_var, k_var, v_var, o_var = tf.unstack(var, num=4, axis=0)
-    if cache is None or memory_antecedent is None:
-      k = tf.einsum("bmi,hik->bhmk", memory_or_query_antecedent, k_var)
-      v = tf.einsum("bmi,hiv->bhmv", memory_or_query_antecedent, v_var)
-      if query_antecedent is None:
-        # we are computing the cache.
-        return k, v
-      q = tf.einsum("bqi,hik->bhqk", query_antecedent, q_var)
-    if cache is not None:
-      if memory_antecedent is not None:
-        q = tf.einsum("bqi,hik->bhqk", query_antecedent, q_var)
-        k = cache["k_encdec"]
-        v = cache["v_encdec"]
-      else:
-        k = cache["k"] = tf.concat([cache["k"], k], axis=2)
-        v = cache["v"] = tf.concat([cache["v"], v], axis=2)
-    logits = tf.einsum("bhqk,bhmk->bhqm", q, k)
-    if mask is not None:
-      logits += mask
-    weights = tf.nn.softmax(logits)
-    o = tf.einsum("bhqm,bhmv->bhqv", weights, v)
-    return tf.einsum("bhqv,hiv->bqi", o, o_var)
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py b/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py
index fc95edc63..68f3a6741 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py
@@ -63,12 +63,11 @@ def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
 def get_placement_mesh(hparams):
   graph = mtf.Graph()
   mesh = mtf.Mesh(graph, "my_mesh")
-  mesh_shape = mtf.parse_mesh_shape(hparams.mesh_shape)
-  mesh_size = mtf.list_product(mesh_shape)
+  mesh_shape = mtf.convert_to_shape(hparams.mesh_shape)
 
-  mesh_devices = [""] * mesh_size
+  mesh_devices = [""] * mesh_shape.size
   mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-      mesh_shape, mtf.parse_layout(hparams.layout), mesh_devices)
+      mesh_shape, hparams.layout, mesh_devices)
   return mesh, mesh_impl
 
 
@@ -85,7 +84,7 @@ def testMtfTransformer(self):
     logits, _ = model.mtf_model_fn(features, mesh)
     lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
     tf_group = lowering.copy_masters_to_slices()
-    tf_logits = lowering.outfeed(logits)
+    tf_logits = lowering.export_to_tf_tensor(logits)
 
     with self.test_session() as session:
       session.run(tf.global_variables_initializer())
@@ -97,14 +96,14 @@ def testMtfTransformerDataParallel(self):
     hparams = mtf_transformer.mtf_transformer_single()
 
     model, features, hparams = get_model(hparams)
-    hparams.mesh_shape = "2"
-    hparams.layout = "batch:0"
+    hparams.mesh_shape = "all:2"
+    hparams.layout = "batch:all"
     mesh, mesh_impl = get_placement_mesh(hparams)
 
     logits, _ = model.mtf_model_fn(features, mesh)
     lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
     tf_group = lowering.copy_masters_to_slices()
-    tf_logits = lowering.outfeed(logits)
+    tf_logits = lowering.export_to_tf_tensor(logits)
 
     with self.test_session() as session:
       session.run(tf.global_variables_initializer())
@@ -116,14 +115,14 @@ def testMtfTransformerModelParallel(self):
     hparams = mtf_transformer.mtf_transformer_single()
 
     model, features, hparams = get_model(hparams)
-    hparams.mesh_shape = "2"
-    hparams.layout = "length:0"
+    hparams.mesh_shape = "all:2"
+    hparams.layout = "length:all"
     mesh, mesh_impl = get_placement_mesh(hparams)
 
     logits, _ = model.mtf_model_fn(features, mesh)
     lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
     tf_group = lowering.copy_masters_to_slices()
-    tf_logits = lowering.outfeed(logits)
+    tf_logits = lowering.export_to_tf_tensor(logits)
 
     with self.test_session() as session:
       session.run(tf.global_variables_initializer())
@@ -135,14 +134,14 @@ def testMtfTransformerDataModelParallel(self):
     hparams = mtf_transformer.mtf_transformer_single()
 
     model, features, hparams = get_model(hparams)
-    hparams.mesh_shape = "2.2"
-    hparams.layout = "batch:0;vocab:1;d_ff:1;heads:1"
+    hparams.mesh_shape = "batch:2;model:2"
+    hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
     mesh, mesh_impl = get_placement_mesh(hparams)
 
     logits, _ = model.mtf_model_fn(features, mesh)
     lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
     tf_group = lowering.copy_masters_to_slices()
-    tf_logits = lowering.outfeed(logits)
+    tf_logits = lowering.export_to_tf_tensor(logits)
 
     with self.test_session() as session:
       session.run(tf.global_variables_initializer())
diff --git a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
index f6a4d4a31..0496b45ca 100644
--- a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
@@ -229,7 +229,7 @@ def random(self, shape, tf_fn, kwargs):
     """Call a random tf operation (e.g. random_uniform).
 
     Args:
-      shape: a TensorShape
+      shape: a Shape
       tf_fn: a function such as tf.random_uniform
       kwargs: kwargs to pass to tf_fn, except for seed
 
@@ -253,7 +253,7 @@ def laid_out_pnum(self):
   def devices(self):
     return self._devices
 
-  def outfeed(self, x, laid_out_x):
+  def export_to_tf_tensor(self, x, laid_out_x):
     """Turn a Tensor into a tf.Tensor.
 
     Args:
@@ -264,8 +264,8 @@ def outfeed(self, x, laid_out_x):
     """
     return self.combine_slices(laid_out_x.all_slices, x.shape)
 
-  def infeed(self, x, tf_x):
-    """Infeed a tf.Tensor, producing a LaidOutTensor.
+  def import_tf_tensor(self, x, tf_x):
+    """Import a tf.Tensor, producing a LaidOutTensor.
 
     Args:
       x: a Tensor
diff --git a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
index fbd9ae66c..4edb7073d 100644
--- a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
@@ -201,7 +201,7 @@ def allconcat(self, x, mesh_axis, concat_axis, stack=False):
     coord = self.laid_out_pcoord(mesh_axis)
     t = x.one_slice
     old_shape = t.shape.as_list()
-    num_parts = self.shape[mesh_axis]
+    num_parts = self.shape[mesh_axis].size
     t = tf.expand_dims(t, concat_axis)
     t *= tf.reshape(
         tf.one_hot(coord.one_slice, num_parts, dtype=t.dtype),
@@ -282,7 +282,7 @@ def random(self, shape, tf_fn, kwargs):
     """Call a random tf operation (e.g. random_uniform).
 
     Args:
-      shape: a TensorShape
+      shape: a Shape
       tf_fn: a function such as tf.random_uniform
       kwargs: kwargs to pass to tf_fn, except for seed
 
@@ -308,7 +308,7 @@ def random(self, shape, tf_fn, kwargs):
     x = self.allreduce(x, mesh_axes, "SUM")
     return x
 
-  def outfeed(self, x, laid_out_x):
+  def export_to_tf_tensor(self, x, laid_out_x):
     """Turn a Tensor into a tf.Tensor.
 
     Args:
@@ -320,12 +320,14 @@ def outfeed(self, x, laid_out_x):
     tensor_layout = self.tensor_layout(x.shape)
     if not tensor_layout.is_fully_replicated:
       raise NotImplementedError(
-          "SimdMeshImpl only supports outfeed of fully-replicated Tensors."
-          " Try reshaping to new dimension names.")
+          "SimdMeshImpl only supports export_to_tf_tensor of fully-replicated "
+          "Tensors.  Try reshaping to new dimension names. "
+          " x.shape = %s tensor_layout=%s"
+          % (x.shape, tensor_layout))
     return laid_out_x.one_slice
 
-  def infeed(self, x, tf_x):
-    """Infeed a tf.Tensor, producing a LaidOutTensor.
+  def import_tf_tensor(self, x, tf_x):
+    """Import a tf.Tensor, producing a LaidOutTensor.
 
     Args:
       x: a Tensor

From 6ba9bfe84620dbae325f6000f3c8b5ab33450bb4 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 8 Aug 2018 11:03:05 -0700
Subject: [PATCH 0498/2720] removing unnecessary const.

PiperOrigin-RevId: 207913779
---
 tensor2tensor/models/research/next_frame_sv2p.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index c3cefa93b..99f64d644 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -36,8 +36,6 @@
 tfl = tf.layers
 tfcl = tf.contrib.layers
 
-_LARGE_STEP_NUMBER = 100000
-
 
 @registry.register_model
 class NextFrameStochastic(next_frame.NextFrameBasic):

From c30080ef406407da1e0796379ab3b262778df2ce Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 8 Aug 2018 11:26:46 -0700
Subject: [PATCH 0499/2720] small tweak for mtf_toy_tpu_model

PiperOrigin-RevId: 207918076
---
 tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
index 7366e253d..d5c4cdd86 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
@@ -48,8 +48,6 @@
 tf.flags.DEFINE_integer('steps_per_checkpoint', 200, 'steps_per_checkpoint')
 tf.flags.DEFINE_string('master', 'local',
                        'BNS name of the TensorFlow master to use.')
-tf.flags.DEFINE_string('evaluation_master', 'local',
-                       'BNS name of the TensorFlow master to use.')
 tf.flags.DEFINE_string(
     'model_dir',
     default='',
@@ -182,7 +180,7 @@ def run_toy_model_tpu():
   mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
   config = tpu_config.RunConfig(
       master=FLAGS.master,
-      evaluation_master=FLAGS.evaluation_master,
+      evaluation_master=FLAGS.master,
       model_dir=FLAGS.model_dir,
       save_checkpoints_steps=None,  # Disable the default saver
       save_checkpoints_secs=None,  # Disable the default saver

From 0b6e09100c57b19e848584010faa10254509dc82 Mon Sep 17 00:00:00 2001
From: cbockman <c.bockman@gmail.com>
Date: Wed, 8 Aug 2018 13:31:39 -0700
Subject: [PATCH 0500/2720] Fix broken universal_transformer_teeny

universal_transformer_teeny() always fails with

```
raise ValueError('Hyperparameter name is reserved: %s' % name)
ValueError: Hyperparameter name is reserved: num_rec_steps
```

This is because universal_transformer_teeny does

transformer_teeny => hparams.num_rec_steps = 2 => update_hparams_for_universal_transformer => tries to add num_rec_steps to hparams.

num_rec_steps is already in hparams, so it fails.

Deleting num_rec_steps from transformer_teeny seemed like the most logical solution here.
---
 tensor2tensor/models/research/universal_transformer.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 9558122a3..81380afc8 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -477,7 +477,6 @@ def universal_transformer_tiny():
 @registry.register_hparams
 def transformer_teeny():
   hparams = transformer.transformer_base()
-  hparams.num_rec_steps = 2
   hparams.hidden_size = 128
   hparams.filter_size = 128
   hparams.num_heads = 2

From b3a29bebf514439f81bc5d7e575279eb98614a40 Mon Sep 17 00:00:00 2001
From: "K1ngP1r@t&" <agozzoli.ai@gmail.com>
Date: Wed, 8 Aug 2018 22:56:49 +0200
Subject: [PATCH 0501/2720] Run on FloydHub integration (#961)

* Run on FH integration.

* Update integration
---
 README.md              | 25 +++++++++++++++++++++++++
 floyd.yml              |  2 ++
 floyd_requirements.txt |  1 +
 3 files changed, 28 insertions(+)
 create mode 100644 floyd.yml
 create mode 100644 floyd_requirements.txt

diff --git a/README.md b/README.md
index 814e58404..37d96829b 100644
--- a/README.md
+++ b/README.md
@@ -9,6 +9,7 @@ welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CO
 [![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
 [![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
 [![Travis](https://img.shields.io/travis/tensorflow/tensor2tensor.svg)](https://travis-ci.org/tensorflow/tensor2tensor)
+[![Run on FH](https://static.floydhub.com/button/button-small.svg)](https://floydhub.com/run)
 
 [Tensor2Tensor](https://github.com/tensorflow/tensor2tensor), or
 [T2T](https://github.com/tensorflow/tensor2tensor) for short, is a library
@@ -66,6 +67,7 @@ pip install tensor2tensor && t2t-trainer \
 * [Adding your own components](#adding-your-own-components)
 * [Adding a dataset](#adding-a-dataset)
 * [Papers](#papers)
+* [Run on FloydHub](#run-on-floydhub)
 
 ## Suggested Datasets and Models
 
@@ -408,6 +410,29 @@ paper](https://arxiv.org/abs/1803.07416).
 }
 ```
 
+## Run on FloydHub
+
+[![Run on FloydHub](https://static.floydhub.com/button/button.svg)](https://floydhub.com/run)
+
+Click this button to open a [Workspace](https://blog.floydhub.com/workspaces/) on [FloydHub](https://www.floydhub.com/?utm_medium=readme&utm_source=tensor2tensor&utm_campaign=jul_2018). You can use the workspace to develop and test your code on a fully configured cloud GPU machine.
+
+Tensor2Tensor comes preinstalled in the environment, you can simply open a [Terminal](https://docs.floydhub.com/guides/workspace/#using-terminal) and run your code.
+
+```bash
+# Test the quick-start on a Workspace's Terminal with this command
+t2t-trainer \
+  --generate_data \
+  --data_dir=./t2t_data \
+  --output_dir=./t2t_train/mnist \
+  --problem=image_mnist \
+  --model=shake_shake \
+  --hparams_set=shake_shake_quick \
+  --train_steps=1000 \
+  --eval_steps=100
+```
+
+Note: Ensure compliance with the FloydHub [Terms of Service](https://www.floydhub.com/about/terms).
+
 Tensor2Tensor was used to develop a number of state-of-the-art models
 and deep learning methods. Here we list some papers that were based on T2T
 from the start and benefited from its features and architecture in ways
diff --git a/floyd.yml b/floyd.yml
new file mode 100644
index 000000000..16ca4fd92
--- /dev/null
+++ b/floyd.yml
@@ -0,0 +1,2 @@
+env: tensorflow-1.9
+machine: gpu
diff --git a/floyd_requirements.txt b/floyd_requirements.txt
new file mode 100644
index 000000000..fe920060a
--- /dev/null
+++ b/floyd_requirements.txt
@@ -0,0 +1 @@
+tensor2tensor

From a88e3998e257037c98c1f43d87dc0ef504f49725 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 8 Aug 2018 13:55:51 -0700
Subject: [PATCH 0502/2720] internal

PiperOrigin-RevId: 207941037
---
 README.md              | 25 -------------------------
 floyd.yml              |  2 --
 floyd_requirements.txt |  1 -
 3 files changed, 28 deletions(-)
 delete mode 100644 floyd.yml
 delete mode 100644 floyd_requirements.txt

diff --git a/README.md b/README.md
index 37d96829b..814e58404 100644
--- a/README.md
+++ b/README.md
@@ -9,7 +9,6 @@ welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CO
 [![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
 [![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
 [![Travis](https://img.shields.io/travis/tensorflow/tensor2tensor.svg)](https://travis-ci.org/tensorflow/tensor2tensor)
-[![Run on FH](https://static.floydhub.com/button/button-small.svg)](https://floydhub.com/run)
 
 [Tensor2Tensor](https://github.com/tensorflow/tensor2tensor), or
 [T2T](https://github.com/tensorflow/tensor2tensor) for short, is a library
@@ -67,7 +66,6 @@ pip install tensor2tensor && t2t-trainer \
 * [Adding your own components](#adding-your-own-components)
 * [Adding a dataset](#adding-a-dataset)
 * [Papers](#papers)
-* [Run on FloydHub](#run-on-floydhub)
 
 ## Suggested Datasets and Models
 
@@ -410,29 +408,6 @@ paper](https://arxiv.org/abs/1803.07416).
 }
 ```
 
-## Run on FloydHub
-
-[![Run on FloydHub](https://static.floydhub.com/button/button.svg)](https://floydhub.com/run)
-
-Click this button to open a [Workspace](https://blog.floydhub.com/workspaces/) on [FloydHub](https://www.floydhub.com/?utm_medium=readme&utm_source=tensor2tensor&utm_campaign=jul_2018). You can use the workspace to develop and test your code on a fully configured cloud GPU machine.
-
-Tensor2Tensor comes preinstalled in the environment, you can simply open a [Terminal](https://docs.floydhub.com/guides/workspace/#using-terminal) and run your code.
-
-```bash
-# Test the quick-start on a Workspace's Terminal with this command
-t2t-trainer \
-  --generate_data \
-  --data_dir=./t2t_data \
-  --output_dir=./t2t_train/mnist \
-  --problem=image_mnist \
-  --model=shake_shake \
-  --hparams_set=shake_shake_quick \
-  --train_steps=1000 \
-  --eval_steps=100
-```
-
-Note: Ensure compliance with the FloydHub [Terms of Service](https://www.floydhub.com/about/terms).
-
 Tensor2Tensor was used to develop a number of state-of-the-art models
 and deep learning methods. Here we list some papers that were based on T2T
 from the start and benefited from its features and architecture in ways
diff --git a/floyd.yml b/floyd.yml
deleted file mode 100644
index 16ca4fd92..000000000
--- a/floyd.yml
+++ /dev/null
@@ -1,2 +0,0 @@
-env: tensorflow-1.9
-machine: gpu
diff --git a/floyd_requirements.txt b/floyd_requirements.txt
deleted file mode 100644
index fe920060a..000000000
--- a/floyd_requirements.txt
+++ /dev/null
@@ -1 +0,0 @@
-tensor2tensor

From b20f0073f09ece90b822628a43ad5299b4ecc445 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 8 Aug 2018 13:57:32 -0700
Subject: [PATCH 0503/2720] internal merge of PR #961

PiperOrigin-RevId: 207941265
---
 README.md              | 25 +++++++++++++++++++++++++
 docs/walkthrough.md    | 25 +++++++++++++++++++++++++
 floyd.yml              |  2 ++
 floyd_requirements.txt |  1 +
 4 files changed, 53 insertions(+)
 create mode 100644 floyd.yml
 create mode 100644 floyd_requirements.txt

diff --git a/README.md b/README.md
index 814e58404..37d96829b 100644
--- a/README.md
+++ b/README.md
@@ -9,6 +9,7 @@ welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CO
 [![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
 [![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
 [![Travis](https://img.shields.io/travis/tensorflow/tensor2tensor.svg)](https://travis-ci.org/tensorflow/tensor2tensor)
+[![Run on FH](https://static.floydhub.com/button/button-small.svg)](https://floydhub.com/run)
 
 [Tensor2Tensor](https://github.com/tensorflow/tensor2tensor), or
 [T2T](https://github.com/tensorflow/tensor2tensor) for short, is a library
@@ -66,6 +67,7 @@ pip install tensor2tensor && t2t-trainer \
 * [Adding your own components](#adding-your-own-components)
 * [Adding a dataset](#adding-a-dataset)
 * [Papers](#papers)
+* [Run on FloydHub](#run-on-floydhub)
 
 ## Suggested Datasets and Models
 
@@ -408,6 +410,29 @@ paper](https://arxiv.org/abs/1803.07416).
 }
 ```
 
+## Run on FloydHub
+
+[![Run on FloydHub](https://static.floydhub.com/button/button.svg)](https://floydhub.com/run)
+
+Click this button to open a [Workspace](https://blog.floydhub.com/workspaces/) on [FloydHub](https://www.floydhub.com/?utm_medium=readme&utm_source=tensor2tensor&utm_campaign=jul_2018). You can use the workspace to develop and test your code on a fully configured cloud GPU machine.
+
+Tensor2Tensor comes preinstalled in the environment, you can simply open a [Terminal](https://docs.floydhub.com/guides/workspace/#using-terminal) and run your code.
+
+```bash
+# Test the quick-start on a Workspace's Terminal with this command
+t2t-trainer \
+  --generate_data \
+  --data_dir=./t2t_data \
+  --output_dir=./t2t_train/mnist \
+  --problem=image_mnist \
+  --model=shake_shake \
+  --hparams_set=shake_shake_quick \
+  --train_steps=1000 \
+  --eval_steps=100
+```
+
+Note: Ensure compliance with the FloydHub [Terms of Service](https://www.floydhub.com/about/terms).
+
 Tensor2Tensor was used to develop a number of state-of-the-art models
 and deep learning methods. Here we list some papers that were based on T2T
 from the start and benefited from its features and architecture in ways
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 814e58404..37d96829b 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -9,6 +9,7 @@ welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CO
 [![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
 [![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
 [![Travis](https://img.shields.io/travis/tensorflow/tensor2tensor.svg)](https://travis-ci.org/tensorflow/tensor2tensor)
+[![Run on FH](https://static.floydhub.com/button/button-small.svg)](https://floydhub.com/run)
 
 [Tensor2Tensor](https://github.com/tensorflow/tensor2tensor), or
 [T2T](https://github.com/tensorflow/tensor2tensor) for short, is a library
@@ -66,6 +67,7 @@ pip install tensor2tensor && t2t-trainer \
 * [Adding your own components](#adding-your-own-components)
 * [Adding a dataset](#adding-a-dataset)
 * [Papers](#papers)
+* [Run on FloydHub](#run-on-floydhub)
 
 ## Suggested Datasets and Models
 
@@ -408,6 +410,29 @@ paper](https://arxiv.org/abs/1803.07416).
 }
 ```
 
+## Run on FloydHub
+
+[![Run on FloydHub](https://static.floydhub.com/button/button.svg)](https://floydhub.com/run)
+
+Click this button to open a [Workspace](https://blog.floydhub.com/workspaces/) on [FloydHub](https://www.floydhub.com/?utm_medium=readme&utm_source=tensor2tensor&utm_campaign=jul_2018). You can use the workspace to develop and test your code on a fully configured cloud GPU machine.
+
+Tensor2Tensor comes preinstalled in the environment, you can simply open a [Terminal](https://docs.floydhub.com/guides/workspace/#using-terminal) and run your code.
+
+```bash
+# Test the quick-start on a Workspace's Terminal with this command
+t2t-trainer \
+  --generate_data \
+  --data_dir=./t2t_data \
+  --output_dir=./t2t_train/mnist \
+  --problem=image_mnist \
+  --model=shake_shake \
+  --hparams_set=shake_shake_quick \
+  --train_steps=1000 \
+  --eval_steps=100
+```
+
+Note: Ensure compliance with the FloydHub [Terms of Service](https://www.floydhub.com/about/terms).
+
 Tensor2Tensor was used to develop a number of state-of-the-art models
 and deep learning methods. Here we list some papers that were based on T2T
 from the start and benefited from its features and architecture in ways
diff --git a/floyd.yml b/floyd.yml
new file mode 100644
index 000000000..16ca4fd92
--- /dev/null
+++ b/floyd.yml
@@ -0,0 +1,2 @@
+env: tensorflow-1.9
+machine: gpu
diff --git a/floyd_requirements.txt b/floyd_requirements.txt
new file mode 100644
index 000000000..fe920060a
--- /dev/null
+++ b/floyd_requirements.txt
@@ -0,0 +1 @@
+tensor2tensor

From b706e3f4a79f3b4430ec5e7c1c55b79cf881ba8a Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 8 Aug 2018 14:26:19 -0700
Subject: [PATCH 0504/2720] Add README.md. rename ComputationLayout ->
 LayoutRules.

PiperOrigin-RevId: 207945957
---
 tensor2tensor/data_generators/problem.py      |   4 +-
 tensor2tensor/mesh_tensorflow/README.md       | 288 ++++++++++++++++++
 .../mesh_tensorflow/mesh_tensorflow.py        |  24 +-
 tensor2tensor/mesh_tensorflow/mnist.py        |  15 +-
 tensor2tensor/mesh_tensorflow/mtf_model.py    |   6 +-
 .../mesh_tensorflow/mtf_toy_model_tpu.py      |   4 +-
 6 files changed, 310 insertions(+), 31 deletions(-)
 create mode 100644 tensor2tensor/mesh_tensorflow/README.md

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 9124a6bee..3ffc6b01b 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -753,12 +753,10 @@ def _dataset_partition(self, mode, config):
       self._next_partition_id = 0
       return 0, 1
     phift = config.tpu_config.per_host_input_for_training
-    # BEGIN GOOGLE-INTERNAL
-    # This is the mesh-tensorflow case.  Still requires patch of cl/204685944
+    # This is the mesh-tensorflow case.
     if (hasattr(tpu_config.InputPipelineConfig, "BROADCAST") and
         phift == tpu_config.InputPipelineConfig.BROADCAST):
       return 0, 1
-    # END GOOOGLE-INTERNAL
     if phift:
       num_partitions = max(config.tpu_config.num_shards // 8, 1)
     else:
diff --git a/tensor2tensor/mesh_tensorflow/README.md b/tensor2tensor/mesh_tensorflow/README.md
new file mode 100644
index 000000000..e78c4f983
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/README.md
@@ -0,0 +1,288 @@
+# Mesh TensorFlow - Model Parallelism Made Easier
+
+# Introduction
+
+Mesh TensorFlow (mtf) is a language for distributed deep
+learning, capable of specifying a broad class of distributed tensor
+computations.  The purpose of mesh-tensorflow is to formalize and implement
+distribution strategies for your computation graph over your hardware/processors
+For example: "Split the batch over rows of processors and split
+the units in the hidden layer across columns of processors." Mesh-TensorFlow is
+implemented as a layer over TensorFlow.
+
+## Do I need Mesh-TensorFlow?
+If you just want data-parallel training (batch-splitting), then you do not need
+mesh-tensorflow, though Mesh-TensorFlow can do this.  The most common reasons
+for more sophisticated parallel computation are:
+
+* The parameters of the model do not fit on one device - e.g. a
+5-billion-parameter language model.
+
+* An example is so large that the activations do not fit on one device. - e.g.
+large images.  TODO(noam): we still need to implement spatially-partitioned
+convolutions
+
+* Lower-latency parallel inference (at batch size 1).
+
+## The Mesh-TensorFlow Approach to Distributed Computation
+
+* A "Mesh" is an n-dimensional array of processors, connected by a network.
+
+* Each tensor is distributed (split and/or replicated) across all processors
+  in a mesh.
+
+* The "layout" of a tensor on a mesh is an injective partial map from the
+  dimensions of the tensor to the dimensions of the mesh, specifying which
+  dimensions of the tensor are split across which dimensions of the mesh.  An
+  empty layout means that the tensor is fully replicated across all processors.
+
+* Tensor dimensions and mesh dimensions are named.  The layouts of all tensors
+  follow from a set of user-defined layout rules which specify which
+  tensor-dimensions are split across which mesh-dimensions.  This ensures that
+  the corresponding dimensions in different tensors are split in the same
+  manner.
+
+* Layouts do not affect results - only performance.
+
+* The implementation of an operation involves parallel computation on all
+  processors in the mesh, and sometimes also collective communication.  A
+  processor usually just manipulates the slices of the input tensors already
+  resident on that processor, and produces the slice of the output that goes on
+  that processor.
+  
+## Example Models
+
+This directory contains code for running several well-known models across
+different tasks. We outline an example below.
+
+### Example Network (MNIST)
+
+To illustrate, let us consider a simple model for the MNIST image-classification
+task.  Our network has one hidden layer with 1024 units, and an output layer
+with 10 units (corresponding to the 10 digit classes).  
+
+The code consists of two parts, the first describing the mathematical
+operations, and the second describing the devices and tensor/computation layout.
+For the full example, see [`mnist.py`](mnist.py).
+TODO(noam): verify that this code works.
+
+```Python
+# tf_images is a tf.Tensor with shape [100, 28, 28] and dtype tf.float32
+# tf_labels is a tf.Tensor with shape [100] and dtype tf.int32
+graph = mtf.Graph()
+mesh = mtf.Mesh(graph, "my_mesh")
+batch_dim = mtf.Dimension("batch", 100)
+rows_dim = mtf.Dimension("rows", 28)
+cols_dim = mtf.Dimension("cols", 28)
+hidden_dim = mtf.Dimension("hidden", 1024)
+classes_dim = mtf.Dimension("classes", 10)
+images = mtf.import_tf_tensor(
+    mesh, tf_images, shape=[batch_dim, rows_dim, cols_dim])
+labels = mtf.import_tf_tensor(mesh, tf_labels, [batch_dim])
+w1 = mtf.get_variable(mesh, "w1", [rows_dim, cols_dim, hidden_dim])
+w2 = mtf.get_variable(mesh, "w2", [hidden_dim, classes_dim])
+# einsum is a generalization of matrix multiplication (see numpy.einsum)
+hidden = mtf.relu(mtf.einsum(images, w1, output_shape=[batch_dim, hidden_dim]))
+logits = mtf.einsum(hidden, w2, output_shape=[batch_dim, classes_dim])
+loss = mtf.reduce_mean(mtf_layers.softmax_cross_entropy_with_logits(
+    logits, mtf.one_hot(labels, classes_dim), classes_dim))
+w1_grad, w2_grad = mtf.gradients([loss], [w1, w2])
+update_w1_op = mtf.assign(w1, w1 - w1_grad * 0.001)
+update_w2_op = mtf.assign(w1, w1 - w1_grad * 0.001)
+```
+
+In the code above, we have built a mesh-tensorflow graph, which is simply
+a Python structure.  We have completely defined the mathematical operations.
+In the code below, we specify the mesh of processors and the layout of the
+computation.
+
+```Python
+devices = ["gpu:0", "gpu:1", "gpu:2", "gpu:3"]
+mesh_shape = [("all_processors", 4)]
+layout_rules = [("batch", "all_processors")]
+mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+    mesh_shape, layout_rules, devices)
+lowering = mtf.Lowering(graph, {mesh:mesh_impl})
+tf_update_ops = [lowering.lowered_operation(update_w1_op),
+                 lowering.lowered_operation(update_w2_op)]
+```
+
+The particular layout above implements data-parallelism, splitting the batch of
+examples evenly across all four processors.  Any Tensor with a "batch" dimension
+(e.g. `images`, `h`, `logits`, and their gradients) is split in that dimension
+across all processors, while any tensor without a "batch" dimension (e.g. the
+model parameters) is replicated identically on every processor.
+
+Alternatively, for model-parallelism, we can set
+`layout_rules=[("hidden", "all_processors")]`.  In this case,
+any tensor with a "hidden" dimension (e.g. `hidden`, `w1`, `w2`)  is split,
+while any other tensor (e.g. `image`, `logits`) is fully replicated.
+
+We can even combine data-parallelism and model-parallelism on a 2-dimensional
+mesh of processors.  We split the batch along one dimension of the mesh, and the
+units in the hidden layer along the other dimension of the mesh, as below.  In
+this case, the hidden layer is actually tiled between the four processors, being
+split in both the "batch" and "hidden_units" dimensions.
+
+```Python
+mesh_shape = [("processor_rows", 2), ("processor_cols", 2)]
+layout_rules = [("batch", "processor_rows"), ("hidden", "processor_cols")]
+```
+
+## Where does the network communication happen?
+
+Some mesh-tensorflow operations cause network communication.  For example, an
+einsum (generalized matrix multiplication) is computed as follows:
+
+* On each processor, compute the einsum of the slices of the two operands that
+  are local to that processor.
+* If no reduced-out dimensions are split, then we are done.
+* If reduced-out dimensions are split, then perform an "allreduce" operation 
+  on the resulting slices - summing across any mesh dimensions over which the
+  reduced-out dimensions are split.
+
+Where the allreduces happen depends will depend on the computation layout.
+For example, in a data-parallel layout where the "batch" dimension is split,
+allreduces will happen when computing the parameter gradients, since this
+involves matrix multiplications which reduce out the "batch" dimension.
+
+## How do I pick a layout?
+
+While results do not depend on layout (except in the realm of roundoff errors
+and random seeds), performance and memory consumption depend heavily on layout.
+One day we hope to automate the process of choosing a layout.  For now, you
+really need to fully understand the performance implications and pick one
+yourself.  Mesh-tensorflow helps by accumulating and printing counters of
+computation/communication.  To start, here are some tricks/guidelines.
+
+* It is illegal for two dimensions of the same tensor to be split across the
+  same batch dimension.
+* For any compute-intense operation (e.g. einsum), make sure that all
+  mesh-dimensions are used to split dimensions of the inputs or outputs.
+  Otherwise, computation is duplicated.
+* To keep the ratio of compute/commuication high (i.e. not be bandwidth-bound),
+  split dimensions into large chunks.  This should be familiar in the
+  data-parallelism case, where we want a large batch size per processor to avoid
+  spending most of our time communicating.
+
+# The Mesh-TensorFlow Language
+
+Mesh-TensorFlow (v0.0) is implemented as a Python library which can generate
+part of a TensorFlow graph.  The user first builds a `mtf.Graph` (the analog of
+a TensorFlow graph) made up of `mtf.Tensor`s and `mtf.Operation`s.  As in
+TensorFlow, this graph consists of simple Python objects.  The user then creates
+a `mtf.Lowering` object, which lowers the `mtf.Graph` into TensorFlow, adding to
+the default TensorFlow graph.
+
+The Mesh-TensorFlow language is nearly identical to TensorFlow, with the
+familiar notion of a Graph, Tensors, Operations, and automatic gradient
+computation.  The principal differences are as follows:
+
+## Meshes replace devices
+
+A `Mesh` is a n-dimensional array of processors with named dimensions.  Each
+`Tensor` is assigned to a `Mesh`, instead of a device.
+
+## Tensor dimensions are named
+
+Each `Tensor` has a static `Shape`, which is a tuple of different "Dimensions".
+A `Dimension` is a `(name, size)` pair. For example, the shape of a `Tensor`
+representing a batch of images might be:
+
+`[("batch", 100), ("rows", 28"), ("cols", 28), ("channels", 3)]`.
+
+## Layouts
+
+A `Tensor` is laid out on its mesh with one slice on each processor.  A `Tensor`
+"layout", is an injective partial map specifying which dimensions of the tensor
+are (evenly) split across which dimensions of the mesh.  No dimension of a
+tensor may be split across two dimensions of its mesh and no two dimensions of a
+tensor may be split across the same dimension of its mesh.  The user defines a
+global set of layout rules in the form of (tensor-dimension-name,
+mesh-dimension-name) pairs.  A dimension of a tensor is split across a dimension
+of its mesh if there is a matching rule.
+
+### Example Layouts
+
+Take our example `Tensor` `image_batch` with shape: 
+`[("batch", 100), ("rows", 28"), ("cols", 28), ("channels", 3)]`
+
+Assume that this `Tensor` is assigned to a mesh of 8 processors with shape:
+`[("processor_rows", 2), ("processor_cols", 4)]`
+
+* If we use an empty set of layout rules `[]`, we get no splitting.  Each
+  processor contains the whole `Tensor`.
+
+* If we use the layout rules `"batch:processor_cols"`, then the `"batch"`
+  dimension of the `Tensor` is split across the `"processor_cols"` dimension of
+  the batch.  This means that each processor contains a Tensor slice with shape
+  `[25, 28, 28, 3]`.  For example, processors (0, 3) and (1, 3) contain
+  identical slices - `image_batch[75:100, :, :, :]`.
+
+* If we use the layout rules `"rows:processor_rows;cols:processor_cols"`, 
+  then the image is split in two dimensions, with each processor containing one
+  spatial tile with shape `[100, 14, 7, 3]`.   For example, processor (0, 1)
+  contains the slice `image_batch[:, 0:14, 7:14, :]`.
+
+Some layout rules would lead to illegal layouts:
+
+* `"batch:processor_rows;rows:processor_rows"` is illegal because two tensor
+  dimensions could be split across the same mesh dimension.
+
+* `"channels:processor_rows"` is illegal because the size of the tensor
+  dimension is not evenly divisible by the size of the mesh dimension.
+
+## Einsum
+
+Mesh-TensorFlow uses Einstein-summation notation, `mtf.einsum(inputs,
+output_shape)`, using the (named) `Dimensions` as the symbols.  Matrix-
+multiplication, broadcast, sum-reduction, and transposition can all be expressed
+as special cases of `mtf.einsum`, though the familiar interfaces are also
+supported.  The operation is lowered to slice-wise `tf.einsum`s, followed by
+allreduce across any mesh-dimensions corresponding to the summed-out Tensor
+dimensions.
+
+## Reshape can be expensive
+
+`mtf.reshape(x, new_shape)` is used to change a `Tensor`'s shape, potentially
+leading to a new tensor layout and hence network communication.
+
+# CPU/GPU/TPU implementations
+
+Mesh-TensorFlow works on CPU, GPU and TPU.  The TPU implementation is very
+different from the CPU/GPU implementation.
+
+Multi-CPU/GPU meshes are implemented with `PlacementMeshImpl`.  In this case
+mesh-tensorflow emits separate tensorflow operations placed on the different
+devices, all in one big tensorflow graph.
+
+TPU meshes are implemented in with `SimdMeshImpl`.  In this case,
+mesh-tensorflow emits tensorflow operations (and communication collectives) from
+the perspective of one core, and this same program runs on every core, relying
+on the fact that each core actually performs the same operations.  This
+piggy-backs on the TPU data-parallelism infrastructure, which operates the same
+way.  This "SIMD" approach keeps the tensorflow and xla graphs from growing with
+the number of cores.  The differences between cores are as follows:
+
+* different slices of the variables (this works now)
+* different positions in the collective communication (this works now)
+* different slices of the infed and outfed tensors.  We currently work around
+  this by requiring that all imported/exported tensors be fully-replicated.  In
+  the future, we should handle this correctly.
+
+
+# TODO LIST (please add items)
+
+We are actively working on improving Mesh-TensorFlow in a variety of ways.  Some
+of the top-priority items are:
+`Contact us if you'd like to help!`
+
+* Instructions for running on cloud-tpu.
+* Operations necessary for spatial-partitioning (spatially-partitioned
+  convolution, etc)
+* Examples of image-classification models.
+* Support for multiple meshes and efficient communication between them.  For
+  example, we may want to load training data on a mesh of 64 cpu-machines and
+  infeed them to a mesh of 512 tpu-cores.  We do not need this for language
+  tasks where the data is tiny, but it will be important for other tasks.
+
diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index 12f076534..1d2c18ad4 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -161,7 +161,7 @@ def convert_to_shape(x):
   return Shape([convert_to_dimension(d) for d in x])
 
 
-class ComputationLayout(object):
+class LayoutRules(object):
   """Represents layout of a computation.
 
   Consists of a set of pairs of strings (tensor_dim_name, mesh_dim_name)
@@ -173,7 +173,7 @@ def __init__(self, pairs):
     self._pairs = set(pairs)
 
   def __repr__(self):
-    return "ComputationLayout%s" % self._pairs
+    return "LayoutRules%s" % self._pairs
 
   def tensor_dimension_to_mesh_axis(self, tensor_dimension, mesh_shape):
     """Mesh axis associated with tensor dimension (or None).
@@ -217,11 +217,11 @@ def tensor_layout(self, tensor_shape, mesh_shape):
     return TensorLayout(ret)
 
 
-def convert_to_computation_layout(x):
-  if isinstance(x, ComputationLayout):
+def convert_to_layout_rules(x):
+  if isinstance(x, LayoutRules):
     return x
   else:
-    return ComputationLayout(x)
+    return LayoutRules(x)
 
 
 class TensorLayout(object):
@@ -428,15 +428,15 @@ class MeshImpl(object):
   Subclasses will include PlacementMeshImpl and SimdMeshImpl
   """
 
-  def __init__(self, shape, layout):
+  def __init__(self, shape, layout_rules):
     """Create a mesh.
 
     Args:
       shape: a list of ints
-      layout: a ComputationLayout
+      layout_rules: a LayoutRules
     """
     self._shape = convert_to_shape(shape)
-    self._layout = convert_to_computation_layout(layout)
+    self._layout_rules = convert_to_layout_rules(layout_rules)
 
   @property
   def shape(self):
@@ -447,8 +447,8 @@ def ndims(self):
     return len(self._shape)
 
   @property
-  def layout(self):
-    return self._layout
+  def layout_rules(self):
+    return self._layout_rules
 
   @property
   def size(self):
@@ -462,7 +462,7 @@ def tensor_dimension_to_mesh_axis(self, tensor_dimension):
     Returns:
       an integer or None
     """
-    return self.layout.tensor_dimension_to_mesh_axis(
+    return self.layout_rules.tensor_dimension_to_mesh_axis(
         tensor_dimension, self.shape)
 
   def tensor_layout(self, arg):
@@ -475,7 +475,7 @@ def tensor_layout(self, arg):
     """
     if isinstance(arg, Tensor):
       arg = arg.shape
-    return self.layout.tensor_layout(arg, self.shape)
+    return self.layout_rules.tensor_layout(arg, self.shape)
 
   def mesh_axis_to_cumprod(self, tensor_shape):
     """For each mesh axis, give the product of previous tensor axes.
diff --git a/tensor2tensor/mesh_tensorflow/mnist.py b/tensor2tensor/mesh_tensorflow/mnist.py
index 60f81fe1b..4a7a9d5e3 100644
--- a/tensor2tensor/mesh_tensorflow/mnist.py
+++ b/tensor2tensor/mesh_tensorflow/mnist.py
@@ -41,7 +41,7 @@
                         "after training is skipped.")
 tf.flags.DEFINE_string("mesh_shape", "rows:2;cols:2", "mesh shape")
 tf.flags.DEFINE_string("layout", "batch:rows;hidden1:cols",
-                       "computation layout")
+                       "layout rules")
 
 FLAGS = tf.flags.FLAGS
 
@@ -92,11 +92,11 @@ def model_fn(features, labels, mode, params):
   mesh = mtf.Mesh(graph, "my_mesh")
   logits, loss = mnist_model(features, labels, mesh)
   mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
-  computation_layout = mtf.ComputationLayout(FLAGS.layout)
+  layout_rules = mtf.LayoutRules(FLAGS.layout)
   mesh_size = mesh_shape.size
   mesh_devices = [""] * mesh_size
   mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-      mesh_shape, computation_layout, mesh_devices)
+      mesh_shape, layout_rules, mesh_devices)
 
   if mode == tf.estimator.ModeKeys.TRAIN:
     var_grads = mtf.gradients(
@@ -175,14 +175,7 @@ def run_mnist():
   """Run MNIST training and eval loop."""
   mnist_classifier = tf.estimator.Estimator(
       model_fn=model_fn,
-      model_dir=FLAGS.model_dir,
-      params={
-          "model_dir": FLAGS.model_dir,
-          "mesh_shape": "2.2",
-          "layout": "batch.0;hidden1.1",
-          "batch_size": FLAGS.batch_size,
-          "hidden_size": 512,
-      })
+      model_dir=FLAGS.model_dir)
 
   # Set up training and evaluation input functions.
   def train_input_fn():
diff --git a/tensor2tensor/mesh_tensorflow/mtf_model.py b/tensor2tensor/mesh_tensorflow/mtf_model.py
index ff752191d..9ff075242 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_model.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_model.py
@@ -77,11 +77,11 @@ def estimator_model_fn(cls,
     mesh = mtf.Mesh(graph, "my_mesh")
 
     mesh_shape = mtf.convert_to_shape(hparams.mesh_shape)
-    computation_layout = mtf.ComputationLayout(hparams.layout)
+    layout_rules = mtf.LayoutRules(hparams.layout)
     if use_tpu:
       mesh_devices = [""] * mesh_shape.size
       mesh_impl = simd_mesh_impl.SimdMeshImpl(
-          mesh_shape, computation_layout, mesh_devices,
+          mesh_shape, layout_rules, mesh_devices,
           params["context"].device_assignment)
     else:
       if len(data_parallelism.ps_devices) == 1:
@@ -90,7 +90,7 @@ def estimator_model_fn(cls,
         assert len(data_parallelism.ps_devices) == mesh_shape.size
         mesh_devices = data_parallelism.ps_devices
       mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-          mesh_shape, computation_layout, mesh_devices)
+          mesh_shape, layout_rules, mesh_devices)
 
     # PREDICT mode
     if mode == tf.estimator.ModeKeys.PREDICT:
diff --git a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
index d5c4cdd86..8c7d26f38 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
@@ -41,7 +41,7 @@
 tf.flags.DEFINE_integer('io_size', 2, 'Number of channels per feature.')
 tf.flags.DEFINE_integer('hidden_size', 2, 'Size of each hidden layer.')
 tf.flags.DEFINE_string('mesh_shape', 'all:8', 'mesh shape')
-tf.flags.DEFINE_string('layout', 'hidden:all', 'computation layout')
+tf.flags.DEFINE_string('layout', 'hidden:all', 'layout rules')
 tf.flags.DEFINE_integer('iterations', 100,
                         'Number of iterations per training loop.')
 tf.flags.DEFINE_integer('train_steps', 10000, 'max steps')
@@ -106,7 +106,7 @@ def model_fn(features, labels, mode, params):
   mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
   mesh_devices = [''] * mesh_shape.size
   mesh_impl = SimdMeshImpl(
-      mesh_shape, mtf.convert_to_computation_layout(FLAGS.layout),
+      mesh_shape, mtf.convert_to_layout_rules(FLAGS.layout),
       mesh_devices, params['context'].device_assignment)
   with mtf_utils.outside_all_rewrites():
     logits, loss = toy_model(features, mesh)

From 1adbb459d7af84116f539da97ccaa081959a0210 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 8 Aug 2018 16:02:09 -0700
Subject: [PATCH 0505/2720] The gan-only implementation should use frames
 sampled from the prior latents and not the conditional latents.

PiperOrigin-RevId: 207961576
---
 tensor2tensor/models/research/next_frame_savp.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_savp.py b/tensor2tensor/models/research/next_frame_savp.py
index 1b4678995..5f497fd57 100644
--- a/tensor2tensor/models/research/next_frame_savp.py
+++ b/tensor2tensor/models/research/next_frame_savp.py
@@ -392,7 +392,7 @@ def construct_model(self, images, actions, rewards):
         cond_pred_images, prior_pred_images = \
           all_pred_images[:batch_size], all_pred_images[batch_size:]
 
-        if train_mode:
+        if train_mode and self.hparams.use_vae:
           pred_image = cond_pred_images
         else:
           pred_image = prior_pred_images
@@ -406,7 +406,7 @@ def construct_model(self, images, actions, rewards):
     gen_prior_video = tf.stack(gen_prior_video, axis=0)
     fake_rewards = tf.stack(fake_rewards, axis=0)
 
-    if train_mode:
+    if train_mode and self.hparams.use_vae:
       return gen_cond_video, fake_rewards, latent_means, latent_stds
     else:
       return gen_prior_video, fake_rewards, latent_means, latent_stds

From c28c356be5946b3798c36688c23eb52af02c244a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 8 Aug 2018 16:36:14 -0700
Subject: [PATCH 0506/2720] Internal change

PiperOrigin-RevId: 207966766
---
 tensor2tensor/models/research/next_frame_sv2p.py |  2 ++
 tensor2tensor/models/research/next_frame_test.py | 10 +++++-----
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 99f64d644..b72a62f32 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -544,6 +544,8 @@ def process_single_frame(prev_outputs, inputs):
   def get_extra_loss(self, latent_means=None, latent_stds=None,
                      true_frames=None, gen_frames=None, beta=1.0):
     """Losses in addition to the default modality losses."""
+    del true_frames
+    del gen_frames
     kl_loss = 0.0
     if self.is_training:
       for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index b7bd2c13a..88aa77a98 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -231,11 +231,11 @@ def testStochasticSavp(self):
         next_frame_savp.NextFrameSAVP,
         1)
 
-  def testStochasticSavpGAN(self):
-    hparams = next_frame_params.next_frame_savp()
-    hparams.use_gan = True
-    hparams.use_vae = False
-    self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
+  # def testStochasticSavpGAN(self):
+  #   hparams = next_frame_params.next_frame_savp()
+  #   hparams.use_gan = True
+  #   hparams.use_vae = False
+  #   self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
 
   def testStochasticInvalidVAEGANCombinations(self):
     hparams = next_frame_params.next_frame_savp()

From 94461c714792b38b126692207020430c85f9b92e Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 8 Aug 2018 16:59:50 -0700
Subject: [PATCH 0507/2720] Write an accessor to set tensor lowerings and check
 for shape errors.

PiperOrigin-RevId: 207969927
---
 .../mesh_tensorflow/mesh_tensorflow.py        | 104 +++++++++++-------
 .../mesh_tensorflow/placement_mesh_impl.py    |   4 +
 .../mesh_tensorflow/simd_mesh_impl.py         |   4 +
 3 files changed, 74 insertions(+), 38 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index 1d2c18ad4..c5e2fa1e7 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -378,9 +378,6 @@ def copy_slices_to_masters(self):
     return tf.group(
         [v.copy_slices_to_master for _, v in six.iteritems(self.variables)])
 
-  # def tensor_layout(self, t):
-  #   return self.mesh_impl(t).tensor_layout(t)
-
   def add_counter(self, key, value):
     assert isinstance(value, int)
     self._counters.append((key, value))
@@ -400,6 +397,19 @@ def laid_out_size(self, tensor):
     """
     return self.mesh_impl(tensor).laid_out_size(tensor.shape)
 
+  def set_tensor_lowering(self, tensor, laid_out_tensor):
+    self.verify_slice_shapes(tensor, laid_out_tensor)
+    self.tensors[tensor] = laid_out_tensor
+
+  def verify_slice_shapes(self, tensor, laid_out_tensor):
+    mesh_impl = self.mesh_impl(tensor)
+    correct_shape = mesh_impl.slice_shape(tensor.shape)
+    actual_shape = laid_out_tensor.slice_shape
+    if actual_shape != correct_shape:
+      raise ValueError(
+          "Wrong slice shape: correct_shape = %s actual shape = %s"
+          % (correct_shape, actual_shape))
+
 
 class Mesh(object):
   """A placeholder with no functionality.
@@ -818,6 +828,10 @@ def __add__(self, other):
       return self.mesh_impl.slicewise(
           tf.add, self.to_laid_out_tensor(), other.to_laid_out_tensor())
 
+  @property
+  def slice_shape(self):
+    return self.laid_out_input.slice_shape
+
 
 def convert_args_to_laid_out_tensors(xs):
   """Convert list elements to laid-out-tensors when possible.
@@ -1043,8 +1057,10 @@ def lower(self, lowering):
       for d, mesh_axis in zip(t.shape.dims, layout.tensor_axis_to_mesh_axis):
         if (mesh_axis is not None and d not in self._splittable_dims):
           raise ValueError("dimension %s is not declared as splittable" % d)
-    lowering.tensors[self.outputs[0]] = mesh_impl.slicewise(
-        self._tf_fn, *[lowering.tensors[x] for x in self.inputs])
+    lowering.set_tensor_lowering(
+        self.outputs[0],
+        mesh_impl.slicewise(
+            self._tf_fn, *[lowering.tensors[x] for x in self.inputs]))
 
 
 def slicewise(tf_fn,
@@ -1205,7 +1221,8 @@ def lower(self, lowering):
                    ys, xs, grad_ys in zip(all_ys, all_xs, all_grad_ys)]
     grad_xs = transpose_list_of_lists(all_grad_xs)
     for out, grad_x in zip(self.outputs, grad_xs):
-      lowering.tensors[out] = (
+      lowering.set_tensor_lowering(
+          out,
           lowering.mesh_impl(self).LaidOutTensor.from_tensor_list(grad_x))
 
 
@@ -1223,8 +1240,10 @@ def gradient(self, grad_ys):
     return [dy * self._scalar]
 
   def lower(self, lowering):
-    lowering.tensors[self.outputs[0]] = lowering.mesh_impl(self).slicewise(
-        lambda x: x * self._scalar, lowering.tensors[self.inputs[0]])
+    lowering.set_tensor_lowering(
+        self.outputs[0],
+        lowering.mesh_impl(self).slicewise(
+            lambda x: x * self._scalar, lowering.tensors[self.inputs[0]]))
 
 
 class ScalarAddOperation(Operation):
@@ -1239,8 +1258,10 @@ def gradient(self, grad_ys):
     return grad_ys
 
   def lower(self, lowering):
-    lowering.tensors[self.outputs[0]] = lowering.mesh_impl(self).slicewise(
-        lambda x: x + self._scalar, lowering.tensors[self.inputs[0]])
+    lowering.set_tensor_lowering(
+        self.outputs[0],
+        lowering.mesh_impl(self).slicewise(
+            lambda x: x + self._scalar, lowering.tensors[self.inputs[0]]))
 
 
 class BinaryOpWithBroadcasting(Operation):
@@ -1269,8 +1290,10 @@ def lower(self, lowering):
     if x2.shape != output.shape:
       laid_out_x2 = mesh_impl.slicewise(
           _expand_dims, laid_out_x2, x2.shape, output.shape)
-    lowering.tensors[self.outputs[0]] = mesh_impl.slicewise(
-        self._tf_fn, laid_out_x1, laid_out_x2)
+    lowering.set_tensor_lowering(
+        self.outputs[0],
+        mesh_impl.slicewise(
+            self._tf_fn, laid_out_x1, laid_out_x2))
 
 
 def binary_arguments_to_tensors(x1, x2):
@@ -1391,7 +1414,7 @@ def lower(self, lowering):
     ret = lowering.mesh_impl(self).broadcast_impl(
         lowering.tensors[self.inputs[0]], self.inputs[0].shape,
         self.outputs[0].shape)
-    lowering.tensors[self.outputs[0]] = ret
+    lowering.set_tensor_lowering(self.outputs[0], ret)
 
 
 def broadcast(x, new_shape):
@@ -1465,7 +1488,7 @@ def add_counter_fn():
         y = mesh_impl.allreduce(
             y, reduced_mesh_axes, self._reduction_fn_string)
         add_counter_fn()
-    lowering.tensors[self.outputs[0]] = y
+    lowering.set_tensor_lowering(self.outputs[0], y)
 
 
 class ConcatOperation(Operation):
@@ -1509,7 +1532,7 @@ def slicewise_fn(*args):
       return tf.concat(args, axis=self._axis, name="concat")
     y = mesh_impl.slicewise(
         slicewise_fn, *[lowering.tensors[x] for x in self._inputs])
-    lowering.tensors[self.outputs[0]] = y
+    lowering.set_tensor_lowering(self.outputs[0], y)
 
 
 def concat(xs, concat_dim_name, name=None):
@@ -1572,7 +1595,7 @@ def slicewise_fn(x):
     values = mesh_impl.slicewise(
         slicewise_fn, lowering.tensors[self.inputs[0]])
     for t, v in zip(self._outputs, values):
-      lowering.tensors[t] = v
+      lowering.set_tensor_lowering(t, v)
 
 
 def split(x, split_dim, num_or_size_splits, name=None):
@@ -1617,7 +1640,7 @@ def lower(self, lowering):
     def slicewise_fn(*args):
       return tf.stack(args, axis=self._axis)
     ret = mesh_impl.slicewise(slicewise_fn, *inputs)
-    lowering.tensors[self.outputs[0]] = ret
+    lowering.set_tensor_lowering(self.outputs[0], ret)
 
 
 def stack(xs, dim_name, axis, name=None):
@@ -1659,7 +1682,7 @@ def slicewise_fn(x):
     output_values = mesh_impl.slicewise(
         slicewise_fn, lowering.tensors[self._inputs[0]])
     for t, v in zip(self.outputs, list(output_values)):
-      lowering.tensors[t] = v
+      lowering.set_tensor_lowering(t, v)
 
 
 def unstack(x, dim, name=None):
@@ -1766,7 +1789,7 @@ def add_counter_fn():
     # broadcast from intersection_shape to output_shape
     if intersection_shape != output_shape:
       y = mesh_impl.broadcast_impl(y, intersection_shape, output_shape)
-    lowering.tensors[self.outputs[0]] = y
+    lowering.set_tensor_lowering(self.outputs[0], y)
     computation_shape = Shape(list(input_shape_set))
     lowering.add_counter("einsum", mesh_impl.laid_out_size(computation_shape))
     lowering.add_counter("einsum_unique", computation_shape.size)
@@ -1811,7 +1834,7 @@ def slicewise_fn(x, begin, size):
       return tf.slice(x, begin, size, name="slice")
     y = mesh_impl.slicewise(
         slicewise_fn, lowering.tensors[inputs], begin, size)
-    lowering.tensors[self.outputs[0]] = y
+    lowering.set_tensor_lowering(self.outputs[0], y)
 
 
 class PadOperation(Operation):
@@ -1857,7 +1880,7 @@ def slicewise_fn(x, paddings):
       return tf.pad(x, paddings, name="pad")
     y = mesh_impl.slicewise(
         slicewise_fn, lowering.tensors[inputs], paddings)
-    lowering.tensors[self.outputs[0]] = y
+    lowering.set_tensor_lowering(self.outputs[0], y)
 
 
 class OneHotOperation(Operation):
@@ -1897,7 +1920,7 @@ def slicewise_fn(indices_slice, offset):
                         dtype=self._dtype)
     y = mesh_impl.slicewise(
         slicewise_fn, lowering.tensors[indices], offset)
-    lowering.tensors[self.outputs[0]] = y
+    lowering.set_tensor_lowering(self.outputs[0], y)
 
 
 class ImportOperation(Operation):
@@ -1910,8 +1933,9 @@ def __init__(self, mesh, tf_tensor, shape, name=None):
 
   def lower(self, lowering):
     mesh_impl = lowering.mesh_impl(self)
-    lowering.tensors[self.outputs[0]] = mesh_impl.import_tf_tensor(
-        self.outputs[0], self._tf_tensor)
+    lowering.set_tensor_lowering(
+        self.outputs[0],
+        mesh_impl.import_tf_tensor(self.outputs[0], self._tf_tensor))
 
 
 def anonymous_shape(shape):
@@ -1960,7 +1984,7 @@ def lower(self, lowering):
     with mtf_utils.outside_all_rewrites():
       sv = mesh_impl.LaidOutVariable(self, mesh_impl)
     lowering.variables[self] = sv
-    lowering.tensors[self.outputs[0]] = sv.laid_out_tensor
+    lowering.set_tensor_lowering(self.outputs[0], sv.laid_out_tensor)
     if self._trainable:
       lowering.add_counter("variables/trainable", self.outputs[0].size)
     else:
@@ -2038,9 +2062,10 @@ def lower(self, lowering):
       raise ValueError("Mesh does not suppport control dependencies.")
     with tf.control_dependencies(
         [lowering.operations[d] for d in self._dependencies]):
-      lowering.tensors[self.outputs[0]] = mesh_impl.slicewise(
-          tf.identity,
-          lowering.tensors[self.inputs[0]])
+      lowering.set_tensor_lowering(
+          self.outputs[0],
+          mesh_impl.slicewise(tf.identity,
+                              lowering.tensors[self.inputs[0]]))
 
   def gradient(self, grad_ys):
     return grad_ys
@@ -2073,7 +2098,7 @@ def tf_fn():
       return tf.constant(value=self._value,
                          dtype=self.outputs[0].dtype,
                          shape=slice_shape)
-    lowering.tensors[self.outputs[0]] = mesh_impl.slicewise(tf_fn)
+    lowering.set_tensor_lowering(self.outputs[0], mesh_impl.slicewise(tf_fn))
 
 
 def constant(mesh, value, shape=None, dtype=tf.float32):
@@ -2100,7 +2125,8 @@ def __init__(self, x, name=None):
     self._outputs = [Tensor(self, x.shape, x.dtype)]
 
   def lower(self, lowering):
-    lowering.tensors[self.outputs[0]] = lowering.tensors[self.inputs[0]]
+    lowering.set_tensor_lowering(self.outputs[0],
+                                 lowering.tensors[self.inputs[0]])
 
   @property
   def has_gradient(self):
@@ -2123,10 +2149,12 @@ def __init__(self, x, data, message, name=None, **kwargs):
     self._kwargs = kwargs
 
   def lower(self, lowering):
-    lowering.tensors[self.outputs[0]] = lowering.mesh_impl(self).Print(
-        lowering.tensors[self.inputs[0]],
-        [lowering.tensors[d] for d in self._data], self._message,
-        **self._kwargs)
+    lowering.set_tensor_lowering(
+        self.outputs[0],
+        lowering.mesh_impl(self).Print(
+            lowering.tensors[self.inputs[0]],
+            [lowering.tensors[d] for d in self._data], self._message,
+            **self._kwargs))
 
   def gradient(self, grad_ys):
     return grad_ys
@@ -2232,7 +2260,7 @@ def lower(self, lowering):
       def reshape_fn(x):
         return tf.reshape(x, new_slice_shape)
       slices = mesh_impl.slicewise(reshape_fn, slices)
-    lowering.tensors[self.outputs[0]] = slices
+    lowering.set_tensor_lowering(self.outputs[0], slices)
 
   def gradient(self, grad_ys):
     return [reshape(grad_ys[0], self.inputs[0].shape)]
@@ -3111,8 +3139,8 @@ def __init__(self, mesh, shape, tf_fn, **kwargs):
   def lower(self, lowering):
     mesh_impl = lowering.mesh_impl(self)
     output_shape = self.outputs[0].shape
-    lowering.tensors[self.outputs[0]] = (
-        mesh_impl.random(output_shape, self._tf_fn, self._kwargs))
+    lowering.set_tensor_lowering(self.outputs[0], (
+        mesh_impl.random(output_shape, self._tf_fn, self._kwargs)))
 
 
 def random_uniform(mesh, shape, **kwargs):
@@ -3254,7 +3282,7 @@ def tf_body_fn(*tf_inputs):
                             back_prop=False,
                             **self._tf_kwargs)
     for tf_out, mtf_out in zip(tf_outs, self._outputs):
-      lowering.tensors[mtf_out] = mesh_impl.LaidOutTensor(tf_out)
+      lowering.set_tensor_lowering(mtf_out, mesh_impl.LaidOutTensor(tf_out))
 
 
 def while_loop(cond_fn, body_fn, inputs, num_loop_vars=None, **kwargs):
diff --git a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
index 0496b45ca..9302f916f 100644
--- a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
@@ -52,6 +52,10 @@ def from_tensor_list(cls, tensor_list):
     def all_slices(self):
       return self._tensor_list
 
+    @property
+    def slice_shape(self):
+      return self.tensor_list[0].shape.as_list()
+
     def to_laid_out_tensor(self):
       return self
 
diff --git a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
index 4edb7073d..3033e53d3 100644
--- a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
@@ -74,6 +74,10 @@ def from_tensor_list(cls, tensor_list):
     def all_slices(self):
       return self._tensor_list
 
+    @property
+    def slice_shape(self):
+      return self.one_slice.shape.as_list()
+
     def to_laid_out_tensor(self):
       return self
 

From 04706714efc11cef678d7eab21100c4ea16912ce Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 8 Aug 2018 17:17:36 -0700
Subject: [PATCH 0508/2720] fix typo in asr_transformer.ipynb thanks to @Qiaoxl
 for reporting it in issues #980

PiperOrigin-RevId: 207972493
---
 tensor2tensor/notebooks/asr_transformer.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/notebooks/asr_transformer.ipynb b/tensor2tensor/notebooks/asr_transformer.ipynb
index 85593bf72..d1802d727 100644
--- a/tensor2tensor/notebooks/asr_transformer.ipynb
+++ b/tensor2tensor/notebooks/asr_transformer.ipynb
@@ -158,7 +158,7 @@
         "def decode(integers):\n",
         "  integers = list(np.squeeze(integers))\n",
         "  if 1 in integers:\n",
-        "    integets = integers[:integers.index(1)]\n",
+        "    integers = integers[:integers.index(1)]\n",
         "  return encoders[\"targets\"].decode(np.squeeze(integers))\n"
       ]
     },

From 05a1f5f88785b0cff8792ea06073efb1db53779f Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 8 Aug 2018 17:18:36 -0700
Subject: [PATCH 0509/2720] Formatting corrections in docs.

PiperOrigin-RevId: 207972614
---
 README.md           | 38 +++++++++++++++++++-------------------
 docs/index.md       |  2 +-
 docs/walkthrough.md | 38 +++++++++++++++++++-------------------
 3 files changed, 39 insertions(+), 39 deletions(-)

diff --git a/README.md b/README.md
index 37d96829b..e1702cd2c 100644
--- a/README.md
+++ b/README.md
@@ -391,25 +391,6 @@ for an example.
 Also see the [data generators
 README](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/README.md).
 
-## Papers
-
-When referencing Tensor2Tensor, please cite [this
-paper](https://arxiv.org/abs/1803.07416).
-
-```
-@article{tensor2tensor,
-  author    = {Ashish Vaswani and Samy Bengio and Eugene Brevdo and
-    Francois Chollet and Aidan N. Gomez and Stephan Gouws and Llion Jones and
-    \L{}ukasz Kaiser and Nal Kalchbrenner and Niki Parmar and Ryan Sepassi and
-    Noam Shazeer and Jakob Uszkoreit},
-  title     = {Tensor2Tensor for Neural Machine Translation},
-  journal   = {CoRR},
-  volume    = {abs/1803.07416},
-  year      = {2018},
-  url       = {http://arxiv.org/abs/1803.07416},
-}
-```
-
 ## Run on FloydHub
 
 [![Run on FloydHub](https://static.floydhub.com/button/button.svg)](https://floydhub.com/run)
@@ -433,6 +414,25 @@ t2t-trainer \
 
 Note: Ensure compliance with the FloydHub [Terms of Service](https://www.floydhub.com/about/terms).
 
+## Papers
+
+When referencing Tensor2Tensor, please cite [this
+paper](https://arxiv.org/abs/1803.07416).
+
+```
+@article{tensor2tensor,
+  author    = {Ashish Vaswani and Samy Bengio and Eugene Brevdo and
+    Francois Chollet and Aidan N. Gomez and Stephan Gouws and Llion Jones and
+    \L{}ukasz Kaiser and Nal Kalchbrenner and Niki Parmar and Ryan Sepassi and
+    Noam Shazeer and Jakob Uszkoreit},
+  title     = {Tensor2Tensor for Neural Machine Translation},
+  journal   = {CoRR},
+  volume    = {abs/1803.07416},
+  year      = {2018},
+  url       = {http://arxiv.org/abs/1803.07416},
+}
+```
+
 Tensor2Tensor was used to develop a number of state-of-the-art models
 and deep learning methods. Here we list some papers that were based on T2T
 from the start and benefited from its features and architecture in ways
diff --git a/docs/index.md b/docs/index.md
index 58dffb134..7a7287851 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -33,7 +33,7 @@ research](https://research.googleblog.com/2017/06/accelerating-deep-learning-res
 * [Training on Google Cloud ML](cloud_mlengine.md)
 * [Training on Google Cloud TPUs](cloud_tpu.md)
 * [Distributed Training](distributed_training.md)
-# [Automatic Speech Recognition (ASR) with Transformer](tutorials/asr_with_transformer.md)
+* [Automatic Speech Recognition (ASR) with Transformer](tutorials/asr_with_transformer.md)
 
 ## Solving your task
 
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 37d96829b..e1702cd2c 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -391,25 +391,6 @@ for an example.
 Also see the [data generators
 README](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/README.md).
 
-## Papers
-
-When referencing Tensor2Tensor, please cite [this
-paper](https://arxiv.org/abs/1803.07416).
-
-```
-@article{tensor2tensor,
-  author    = {Ashish Vaswani and Samy Bengio and Eugene Brevdo and
-    Francois Chollet and Aidan N. Gomez and Stephan Gouws and Llion Jones and
-    \L{}ukasz Kaiser and Nal Kalchbrenner and Niki Parmar and Ryan Sepassi and
-    Noam Shazeer and Jakob Uszkoreit},
-  title     = {Tensor2Tensor for Neural Machine Translation},
-  journal   = {CoRR},
-  volume    = {abs/1803.07416},
-  year      = {2018},
-  url       = {http://arxiv.org/abs/1803.07416},
-}
-```
-
 ## Run on FloydHub
 
 [![Run on FloydHub](https://static.floydhub.com/button/button.svg)](https://floydhub.com/run)
@@ -433,6 +414,25 @@ t2t-trainer \
 
 Note: Ensure compliance with the FloydHub [Terms of Service](https://www.floydhub.com/about/terms).
 
+## Papers
+
+When referencing Tensor2Tensor, please cite [this
+paper](https://arxiv.org/abs/1803.07416).
+
+```
+@article{tensor2tensor,
+  author    = {Ashish Vaswani and Samy Bengio and Eugene Brevdo and
+    Francois Chollet and Aidan N. Gomez and Stephan Gouws and Llion Jones and
+    \L{}ukasz Kaiser and Nal Kalchbrenner and Niki Parmar and Ryan Sepassi and
+    Noam Shazeer and Jakob Uszkoreit},
+  title     = {Tensor2Tensor for Neural Machine Translation},
+  journal   = {CoRR},
+  volume    = {abs/1803.07416},
+  year      = {2018},
+  url       = {http://arxiv.org/abs/1803.07416},
+}
+```
+
 Tensor2Tensor was used to develop a number of state-of-the-art models
 and deep learning methods. Here we list some papers that were based on T2T
 from the start and benefited from its features and architecture in ways

From 39b34345b9292c1f4cbebbb4f99014e741863f28 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 8 Aug 2018 17:43:56 -0700
Subject: [PATCH 0510/2720] Rely on convolutional arithmetic instead of
 tf-shape magic to figure out fully-connected dimensions.

PiperOrigin-RevId: 207975405
---
 .../models/research/next_frame_savp.py        | 23 +++++++++++++++----
 .../models/research/next_frame_test.py        | 10 ++++----
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_savp.py b/tensor2tensor/models/research/next_frame_savp.py
index 5f497fd57..ad4136ea5 100644
--- a/tensor2tensor/models/research/next_frame_savp.py
+++ b/tensor2tensor/models/research/next_frame_savp.py
@@ -21,6 +21,7 @@
 from __future__ import division
 from __future__ import print_function
 import numbers
+import numpy as np
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
@@ -102,6 +103,19 @@ def encoder(self, inputs, n_layers=3):
         z_log_var, (batch_size, -1, latent_dims))
     return z_mu, z_log_var
 
+  def expected_output_shape(self, input_shape, stride, padding, kernel_size):
+    return (input_shape + 2*padding - kernel_size) // stride + 1
+
+  def get_fc_dimensions(self, strides, kernel_sizes):
+    """Get expected fully connected shape after a series of convolutions."""
+    output_height, output_width, _ = self.hparams.problem.frame_shape
+    output_steps = self.hparams.video_num_target_frames
+    output_shape = np.array([output_steps, output_height, output_width])
+    for curr_stride, kernel_size in zip(strides, kernel_sizes):
+      output_shape = self.expected_output_shape(
+          output_shape, np.array(curr_stride), 1, kernel_size)
+    return np.prod(output_shape) * self.hparams.num_discriminator_filters * 8
+
   def discriminator(self, frames):
     """3-D SNGAN discriminator.
 
@@ -122,7 +136,9 @@ def discriminator(self, frames):
     # 3-D Conv-net mapping inputs to activations.
     num_outputs = [ndf, ndf*2, ndf*2, ndf*4, ndf*4, ndf*8, ndf*8]
     kernel_sizes = [3, 4, 3, 4, 3, 4, 3]
-    strides = [1, [1, 2, 2], 1, [1, 2, 2], 1, 2, 1]
+    strides = [[1, 1, 1], [1, 2, 2], [1, 1, 1], [1, 2, 2], [1, 1, 1],
+               [2, 2, 2], [1, 1, 1]]
+
     names = ["video_sn_conv0_0", "video_sn_conv0_1", "video_sn_conv1_0",
              "video_sn_conv1_1", "video_sn_conv2_0", "video_sn_conv2_1",
              "video_sn_conv3_0"]
@@ -131,10 +147,7 @@ def discriminator(self, frames):
     for num_filters, kernel_size, stride, name in iterable:
       activations = self.pad_conv3d_lrelu(activations, num_filters, kernel_size,
                                           stride, name)
-
-    # Flatten and apply fully-connected layer.
-    num_fc_dimensions = tf.reduce_prod(
-        common_layers.shape_list(activations)[1:])
+    num_fc_dimensions = self.get_fc_dimensions(strides, kernel_sizes)
     activations = tf.reshape(activations, (-1, num_fc_dimensions))
     return tf.squeeze(tf.layers.dense(activations, 1))
 
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index 88aa77a98..b7bd2c13a 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -231,11 +231,11 @@ def testStochasticSavp(self):
         next_frame_savp.NextFrameSAVP,
         1)
 
-  # def testStochasticSavpGAN(self):
-  #   hparams = next_frame_params.next_frame_savp()
-  #   hparams.use_gan = True
-  #   hparams.use_vae = False
-  #   self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
+  def testStochasticSavpGAN(self):
+    hparams = next_frame_params.next_frame_savp()
+    hparams.use_gan = True
+    hparams.use_vae = False
+    self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
 
   def testStochasticInvalidVAEGANCombinations(self):
     hparams = next_frame_params.next_frame_savp()

From 4b7a75766bde366664666bc20a3b60b0651b9c40 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Wed, 8 Aug 2018 17:52:59 -0700
Subject: [PATCH 0511/2720] Make splitting work for image transformer. Fix bug
 in SliceOperation , add test for image_transformer

PiperOrigin-RevId: 207976391
---
 .../mesh_tensorflow/mesh_tensorflow.py        |  2 +-
 .../mesh_tensorflow/mtf_image_transformer.py  |  2 +
 .../mtf_image_transformer_test.py             | 38 +++++++++++++++++++
 3 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index c5e2fa1e7..bf149997c 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -1828,7 +1828,7 @@ def lower(self, lowering):
     ndims = self._inputs[0].shape.ndims
     axis = self._axis
     begin = [0] * axis + [self._begin] + [0] * (ndims - axis - 1)
-    size = self._outputs[0].shape.to_integer_list
+    size = [-1] * axis + [self._slice_dim[1]] + [-1] * (ndims - axis - 1)
 
     def slicewise_fn(x, begin, size):
       return tf.slice(x, begin, size, name="slice")
diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
index 0b80c4c7a..ab37e7a4d 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
@@ -135,6 +135,7 @@ def layer_prepostprocess_dropout(x):
                 None,
                 kv_channels,
                 heads,
+                block_length=hparams.block_length,
                 name="self_att"))
         # ffn layer
         x += layer_prepostprocess_dropout(mtf_layers.dense_relu_dense(
@@ -231,6 +232,7 @@ def mtf_image_transformer_single():
   hparams.num_heads = 2
   hparams.attention_key_size = 32
   hparams.attention_value_size = 32
+  hparams.block_length = 16
   return hparams
 
 
diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py b/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
index 770b53af5..70ca839ab 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
@@ -89,6 +89,44 @@ def testMtfImageTransformer(self):
       res = session.run(tf_logits)
     self.assertEqual(res.shape, (BATCH_SIZE, IMG_LENGTH*IMG_LENGTH, VOCAB_SIZE))
 
+  def testMtfImageTransformerDataParallel(self):
+    hparams = mtf_image_transformer.mtf_image_transformer_single()
+
+    model, features, hparams = get_model(hparams)
+    hparams.mesh_shape = "all:2"
+    hparams.layout = "batch:all"
+    mesh, mesh_impl = get_placement_mesh(hparams)
+
+    logits, _ = model.mtf_model_fn(features, mesh)
+    lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
+    tf_group = lowering.copy_masters_to_slices()
+    tf_logits = lowering.export_to_tf_tensor(logits)
+
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      session.run(tf_group)
+      res = session.run(tf_logits)
+    self.assertEqual(res.shape, (BATCH_SIZE, IMG_LENGTH*IMG_LENGTH, VOCAB_SIZE))
+
+  def testMtfImageTransformerModelParallel(self):
+    hparams = mtf_image_transformer.mtf_image_transformer_single()
+
+    model, features, hparams = get_model(hparams)
+    hparams.mesh_shape = "all:2"
+    hparams.layout = "length:all"
+    mesh, mesh_impl = get_placement_mesh(hparams)
+
+    logits, _ = model.mtf_model_fn(features, mesh)
+    lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
+    tf_group = lowering.copy_masters_to_slices()
+    tf_logits = lowering.export_to_tf_tensor(logits)
+
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      session.run(tf_group)
+      res = session.run(tf_logits)
+    self.assertEqual(res.shape, (BATCH_SIZE, IMG_LENGTH*IMG_LENGTH, VOCAB_SIZE))
+
 
 if __name__ == "__main__":
   tf.test.main()

From 271da5289a10249f41fcebf3b7fbb50bbbd9041e Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 8 Aug 2018 21:57:35 -0700
Subject: [PATCH 0512/2720] Add Mesh TensorFlow plots with preliminary numbers.

PiperOrigin-RevId: 207994768
---
 tensor2tensor/mesh_tensorflow/README.md       |  11 ++++-
 .../mesh_tensorflow/mtf_transformer.py        |  45 ++++++++++++++++++
 .../mtf_transformer_data_splitting.png        | Bin 0 -> 19717 bytes
 .../mtf_transformer_model_splitting.png       | Bin 0 -> 20232 bytes
 4 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_transformer_data_splitting.png
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_transformer_model_splitting.png

diff --git a/tensor2tensor/mesh_tensorflow/README.md b/tensor2tensor/mesh_tensorflow/README.md
index e78c4f983..ec3f20d68 100644
--- a/tensor2tensor/mesh_tensorflow/README.md
+++ b/tensor2tensor/mesh_tensorflow/README.md
@@ -1,5 +1,9 @@
 # Mesh TensorFlow - Model Parallelism Made Easier
 
+Transformer for EN-FR WMT with model splitting |  Transformer for EN-FR WMT with data splitting
+:-------------------------:|:-------------------------:
+![model_splitting](./mtf_transformer_model_splitting.png) | ![data_splitting](./mtf_transformer_data_splitting.png)
+
 # Introduction
 
 Mesh TensorFlow (mtf) is a language for distributed deep
@@ -53,7 +57,12 @@ convolutions
 ## Example Models
 
 This directory contains code for running several well-known models across
-different tasks. We outline an example below.
+different tasks.
+
+We outline an example below. In the above figures, Mesh-TensorFlow scales
+linearly as the number of TPU shards increases. For model splitting, we varied
+the number of hidden units in the feedforward layer and the number of heads; for
+data splitting, we varied the batch size.
 
 ### Example Network (MNIST)
 
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
index 69c855c7a..f9ed35fe2 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
@@ -819,6 +819,51 @@ def mtf_transformer_paper_tr_4():
   return hparams
 
 
+@registry.register_hparams
+def mtf_transformer_paper_tr_0_mesh_8():
+  hparams = mtf_transformer_paper_tr(0)
+  hparams.mesh_shape = "batch:8"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_tr_4_mesh_16_8():
+  hparams = mtf_transformer_paper_tr(4)
+  hparams.mesh_shape = "model:16;batch:8"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_tr_6_mesh_64_8():
+  hparams = mtf_transformer_paper_tr(6)
+  hparams.mesh_shape = "model:64;batch:8"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_tr_0_mesh_8_v2():
+  hparams = mtf_transformer_paper_tr(0)
+  hparams.batch_size = int(hparams.batch_size / 4)
+  hparams.mesh_shape = "batch:8"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_tr_0_mesh_128():
+  hparams = mtf_transformer_paper_tr(0)
+  hparams.batch_size = int(hparams.batch_size * 4)
+  hparams.mesh_shape = "batch:128"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_tr_0_mesh_512():
+  hparams = mtf_transformer_paper_tr(0)
+  hparams.batch_size = int(hparams.batch_size * 16)
+  hparams.mesh_shape = "batch:512"
+  return hparams
+
+
 @registry.register_hparams
 def mtf_transformer_lm_moe():
   """Mixture of experts language model."""
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer_data_splitting.png b/tensor2tensor/mesh_tensorflow/mtf_transformer_data_splitting.png
new file mode 100644
index 0000000000000000000000000000000000000000..b83a7cdc2f1d3d773e2afa3063e643f3c423cdd5
GIT binary patch
literal 19717
zcmcG$bySsG*FKD(bVx`^C?FvzARUTGhjd7n(g@Ne4I&^QT@upWU0X^7l<p7#X%s1`
z-`t+(JmY=G@B8Z;-^Vy-;O5?Y-RoX!t~sxH&1>$kCrZ+I*p%2PC@6TcGLot&D5#7m
zC|BNMqQhU3Yl^Di*HtG8S#?bK@xn9<h1XaPGCEEuC<MmH|F7hW=UJkl+(MC+e5mf0
zx;5wSrRUso`Ew#k^+x%?XA;6$84J}Ktun=`S8V>(NuxoJrx|LrRt#zsb!UCtW*DOn
z3g6~pX_k6e#x&|l%8$sZ)N{Pe*}7GobXA4*-RF<j9zJK{VfpZ`R<KPxm&YtdmRr_>
zo2T#e#O?dn!IN~sHP1n|HJ*i`qf^mf%Vb76__1t{*}x8kpFkEaLPEmm>dbaY_!$mH
zC5D%>H?bAq6@xAF|6doQ1An@S)a!9LIyqHnvwdG(RaI4$u(q}?EaJ1fcJ12d3~}+X
zv0j;jgM+Q<8tacek4ki^(8IfkCJa?n2vNi`>_0khxd`96bLXYA^W*Kj>}<5~E(t}&
zf&7CyZ8pn;>4dmA;|hhgGON0oh%OU~gO+<!w(1%hkG*MOpkN~|fu-}lQ2~~*8*kS$
z3kzTW`Qw3$cwb*np3omOvH78;CGGIg<<qB6ldayKg+&zn2rgCqpf@DA=mmv^CL_6u
zYzED4lgrxJI5_K@o8h|S<Kt3X<T=&VyjBe(>u)@4s-~DZI3h*xucErO7zhaoF*7r_
zZ@X04OjTZ`*0&-LzPCLe)nW8@tVm0FqrQ7Ocx3(K+P*P-ij@_yc`!UW`Z1DPc<;pP
zbDFSg=V?3x&tt~E<)4Yc#)<oDL;R<wr`ezUJKpQ4s+#DskybS{2tK|~)#mSy%EQmk
z59=ndy}g~=Wpb~nsmW9<QtRPEKg$YjDYhezmRw_U3wL+-&%ahzEk?L-N9^jP#|F4?
zs~Q_UM!jzaM{AaO>*!47uh9|5#m6VZ{Yb;RwtoDOIp*8g*ib-leZIk}DalUI#wjN!
z$EfEj+qym>#;jJ_(n5{k$}cK19ZVHeF?U)U{P5w!wGm!RYioHGJoCbrE-t9$4ojVh
zQF)&VHHsVrsIH=3@$mG_d%3_ZR5zn&ZM|7OUM6X0dV07`9{kYB$!YZ0-f31#i#I&{
zcQqrrSeM22D<}#hZ*<b*;<|DbQyd!fqRshIB5JHBzU`kNU$tIYSy?&uYE7J(s<tE!
z3JQvLo+#GgS$A$R;50WkpWHs|y)a#Qk8pm|Wl~Y@vOO(JJtK3Mm6cxbS7)b0VNp@`
zr~dwaFvgdzt^}TT0+;h&zFgNK9da)yE;gHQc8|V%%0^miHz#BziZ~bUJP1D!m=!lG
z6rjSt^=>|?I5(FO8yh?8Px;H2FLPctnYQrXI?1T|iH#nhuCD$BjQjE9$HTRRN=vxs
znOE(bIZ4WME7bX?WLBM;=3rBD2P}9}l9RuAM8ZGxIT9?ZrWBgy?BWv=Qs7aLd%5k>
zw7!3mH9j;n#LrF|uA{B}yi})(v9c6)R7-0jaD01f>kaQUGE&lHScu#?JAPZf6c{Ee
zd0(GnB6MIKHq&~%IBBn<iwN_oLI{FTqi%Lq4~(Fyyqv3Y{`q>?w{IpRCKZV-EiKDJ
zoN;k+$U%w<3+Y%I!N65EXJ%(@`1tu#_)-Sg*w}dMU@f+Ox+QNhrF^<K@MY7eLi@83
z*Yo7nk7|NT)enwqaPt^VTaLQv=eO|{S3aK3V`c=Hc@HvX!7=%jY%^VLDN&H>R{M2q
z?EbekwzJl^zdy7~5@j_vi~PD~Q8o3svXZ;1prG%!N6Wp5x+xo2q}AaM^4ak{pRP=d
zjU|C$&em9KjW+7*>myp<^i<g$%e$=&rm1X_l91pw2}JO7g&j>wsj8?LBhN@?=4~r?
zbMp^y5I+C;^C$njXM@<7JXl9v-7I<acu*cp<h#G%eCzFCW89*dFQ%bAW~j00=`>aK
z^#U8mica7B`o7${-^>m^=%6u5t9gc+GV#q?OLXy&&aLaAoA%|@!Y;`3`LxdNzT!yg
z6C-YLIO^&jw>;_AANz$dC}?RVG*}E!!U>(I)@F--(6#;&-&z*y(hBP`nU>8<1^aTi
zX3EAwR<^I4v!%7QRi>S5-WfsdV(R@csGpHI@uQ~=Jw*xH>K7c{<IelXWZ%AU?)^+f
z>|NtHn{^8sdMfY*M?76-t=kPPHu{2%G+fYrzBzxwD9f);rmX9E+oxs1U6&Aa#$|VV
zhD|5F=>~#~%Z`k&2e*uLOUCn2ESn#8nbi_vUZqPjUmMC0UtLpHl}7~KG5+=ISN(SF
z49bt)-8^uM-Q7~tGc#}`a;mCg*`L-*j@Rjyq&ko5AZn{jy09A?8z-xt5rmU*GG530
zGCvR2e6!ZhE&u((?z5fQh8Ja<2Od|lf|G;Ub*nLp#)SNM_JZA_*yA0<#5Lz`@ix0S
zhzEZ9BsXivZ&#P%yg7cw{^bQLCnpB^gdpzIN;6rJ5X$w~D$gT3T*NIVrVt}8a(A1$
z87%YwGjsE<nwFpBLVD!wc6Dmn+P!3L$6bM2+11tW#>X`w%pm7nU!V9moHRNifq!L&
zde;*zW?wKoZ$E#|F6zv8tWe|k(b34{q`>OF6!%wK-uw3}V80+N<m3d#$CKjSP;$an
zU|?AtOzW>S@7J!ip_rSSgFui7t`yuy-{2tn{EtAs(M6^3fVNOLH&U}1b9S^7jRs|8
z2rj8&#3JX|*EhU6Y9c*+F%DBQ^!QZx7}V6%a91AJomm-_3h&cH8KdFpX*qXyL8sNv
z6cg?gY-d6dABv0F95cCb1J2IQWOR-zjLCoh{{69K!0iSZ8I!k9Dol=&lJdEYO-^GY
z<p^(uZppiZ1T`%!{-syebjatVZ??&CM$Ay{dYa|D7mAOJ#MgcL)Y8i8HTK)jghKsH
zRyUX29Uw}T&usJZ^E-9j(E8Y-EK7Zjk=v3qJP5*y%kDgNl?W$i)a%#S8=IRyJzAo}
z)#}Gg^7kBz#wbtds&z|DeJFz4a=z(IEIAk@Ga8ifuQvL$iF0$C^ZK;;j+&BNo;+vg
zUaCw_PgkkCcdTW2JmKo=D}FTRN*iqaL`MgUtJFz=Fon;ybN5J1Ma9pU9M<1Bd9}PJ
z3}24Xs`@<_)pLAr@9~nSJ&n$rS-Pm;EPrfmQT_C$jrwWlO1B#!={soj{gqdKadC04
z&)vIsgCnxAlzTDX2vN*S)>gRvG-+&VV&UXm8x-DsQ`8zmLz38=BS?iWLZdudR-M^t
zm=oaJ610{`?$Ob>+j>g3X<=?|T%#*(Bjh5Krp-orxIM#P>WnyeZf$J}rlPXx-eQ32
z_dw52)-63U^P6C8^!M-IQQo*wc^+<29G*D{J%d<QT~h;NG#q~mztZh&z(|+sv3Gi2
ztX)A@xsVy9R#IMmo&2O3HM+qCUD<nKax%1k#iTFhjtaZqQf5nb%dRx{*KMWmEzv)Q
zD6aF&@gy3PgGX_gv8`Kj5TJYb>q}#>z3I|RA?b64H?B7|G&CS`scC3*4e&3BSn1!v
z`%vl`u0W=!s2HtOS|PQ>%f}}rlok{uyn;uaNo3&afKj(sK9)(()tQZn9;HM+)VMV0
zYgZ?1HMjRXO`BQ)!?kjAbCV}cX{|icXBi$aP`Vp;n^$Bw6n6yISq3+XS+~i>R(;BC
zP`-YD)9os%Up}GWBibk}x?r51Z;lN}v;Yvw*xc6EHh0YGK0B$r*>f8=@d}}o;E*l+
z8`hGaWDcU9H9Zn;i~PrtN=x@7`kFf9(v&P)iL5zP+XjY6%UdM7_{oacTgW!ww;q>_
z=zIaO{)8A8LCz;Zch*Ii&Yy!k9x7@4_VeYOHG49&c_WsZnwnDM)FM<bn>Rf~(q+<c
zu(2VWRLZmfV9Vy9Tq2}i3fF~zY(^3Q4vG{U`J@TfRaA&~QKKc5mBTNE(&%ZbEQ2|W
zmIw(eR#QWdniUwhu!gXRTOZ(hix~7_E5P1=7YQ6H>0f!BI-{K(-MJqwkfu~sT>Kdf
z@QvPH2zuz!k*tW?bWhWq_I@ISKbh4vn@T?0X_9?WBC(K1Yi41tH<ylrA;!Y=q|Rc0
z$KCVm18DPEF@-M)!jT2+sQT1mxYV({@ehnr2!%Fwc6?y>CKs=77U^}5fBkA4|L$F)
zOOq(KpPI38=G2!jTsCe7$5}i3Cv?OqQ`;_1U9WG@8`KvS710xmt*ovl-@kvKe*Dh7
zv-Gqt+(8{D$}-skQ1|!n832N{mUT1Py3CA@HWgO`I`Wq$0C{B7IJx1!$vmLhMdBdd
zX9LMMTtyXmexFc?n7RL0{&|;~nY;$ypbQz)h|{C8T3hD<8@y@O0t`0*Sn(s3rgd`W
zKHqT-=iR%Me3ucHtv@M0dX}ZR)jrkK<lE9iF3jPY4Y!olHf9FG%|BCn22<R(x=p9b
z!m`(al1IUbN<Ui-JX><aO@1kwkdOzp)zv7Hb37eZHTeA>0FMykqGwC7ZOyj`R}NHs
zy7LqA0@&JF&)t@SQn^mWHPhINYXJ{mhlewbvGA$<Hh8D-?(pan#Ej~?x{-wiS{+R_
zwSbxJ$i@jj#L_KJPEL!VONe9nGkKhbt*N=Wx%K^zCBKc0kCU!HXQf+CIP}8`^e88w
z`9-?pF6uTVgD2i$bUkB;pUntN2##Qs4hJ1kS!pRpa!N`;QIWi}h9tM<Gy&P*r%$hE
zTlk9e@}5n8*JfvD|D3{4Fv6QX%taru1X0wjWu6c%+2O?Ond`erTtw(+M!`7+nEh9X
zmCPl>j2@hFjB=;#4Dox%mkz8UkTR^llp4q@hagpJGsXSws<ulLb=!jYx8jH;N7=~C
zj@6gkX;F&Y;q&1!Vcc|aJc&Ka#uKIbTay*^m2B{wh34515a9vyeOMqRjJe=0?aIl`
zMUU!1qzXDo=`o22XH9HU|4e7{X`^kL6SHt<J3ZQ6oA=zkBP^T}t1RZ8Dt!w+=Mh3c
zn_E|x@IgMoaoRHPL!*$r%(~Iai^Ae!D!xOp>zIjrDePA6ZsUDBJ9a$RF+VmBc!{i3
zFcxqtjx;Yv3M)MD{WFqr@jI{W&1Dn3CUT5WdiTci9)-s<-Hah=E^-D4^;ub2`6~C9
z1}kvLIIc=$y9Lh=i>CjI7HSPXa_*saEjQ}$H(t>)^uC|MZ=e2oiP71UCvGr1wD;!W
z&RqX9@zbg5h7JpDsjG|tEFRo?psY+FDk}Q>;Na{0{2&~<L}6DBjs#Y#bFi{X%aPpP
z<-)?dvv#g#s5P00lkU5;cT-z+%8ee5vd<&HlbCR1<V)Jy=MKA*5K?f=H8SLQ8xi56
zLwup-`wsaOYzS(g9kqd>Awi(7JmgAQJLFAPV|O7=Pc}NUdV6~xpZi!@TQ75vLGoP?
zeds5vsY!~4hGtq?=X3r?r5{mpo0uJW?R$UyQYY*Vw>JxytbZJ<R_?)}#oaV%!sy?5
z?F?8Xi8)OFdGY@)e@}XdPb1Q+)#ibgpjXK_rV+-ny&FdlHw=5RxA<^S)K$>)$S2Lu
zum4r2x|St*O(JjxpRZ5ojX2OTpO7=98Rs+#-TTWSB;DR=SW^*jXE}6CJhCq+RkV>?
zLB1KKCo+ZVsl&BrrF>j$4o04Z)+8xNmLb}5r+{p!5d)V^zU?Eca^lsUd*p6+#%f+G
z7EmFt7VdJ5zHDEaS^Dv!e(|K_YVATw?L(~mE_hC}X$mgmQ(EKYYzHqCyLwEqquT;r
zQ`qT$4jd16?o*Hph>nfDk!43R9V5NFmsnahAshJGF=$7ISQ$eBv20w=2B~O#o8vfr
zPp}3>_6yau?lWHd`|bye_v<al^jw&z1o3xNkCm{Rw?}^*tbR^Hl+?}lUY;-578Lv-
zfA<MUl0G9e#_Cj=8LFE=ys$h<0Ym|hYnGp9=)buWQc7k#b(*<~DsVnpk&!U?&$YyM
z>nE4wXIqwT3nI^1f}G})Hgxa&rCg$M?-_MQiU$qzmC=T5at5-JS=dc8lK!0(g^|Pb
z>T6=ll)7hb|NhOA?U<5R1m9Hzufb)ql}hZaRDI&E``^3dWL_bjWg6y_jzruiYj5KW
zoT^gQullZP1gjBaadFJ8=PB-H@Tz}-3|nsgVqAT~%IYu8BTL*t@6f-zH?j~gg6_%@
z;*|H|-@CE!8Fdb}JUPXr5<_38N4YRX&#^U8`8%M5uZJM-w8apP!<6UJbkyh6YDKoc
z<dNzvx2t@Y_!mbk4wFVR+yc9E1*IAr|5%DdE|%~mF26W#8BM^({t}Km505Qz_+O)7
zsXQ{T-?=h*PEh~5CG#uavq4$slHq@@F>KPltL>}p94dW<GatL5`+@22`pB9jJ52Co
zRu5!mmoG_9uPVyTypiz>3k{VaI&=4ZUOf<d`AU8@L~4fFmWS`}C$L`bVcPhh`V5y;
z4@my=@F*jJh9YSEqPS(v_?trkoE)mt7R`UxQ(>6r^ij5Lx*LsHG|q%R=I)be3yFWO
zeCeRyn`2t+nRC!2J^QWZu%rXeXaRf$oL8k>tTr!nI}Zku)M8=&ceU<SrG>G%YRqs6
ze~1S2gRRez>#-Y$-VveO+AM6spC`zV5rt?Oggni7MeB@cTI<}-@DG1iD_j2tO95c_
zmB|)Qe*BOCWFU+KqQ{gRJ=YSvj!vDOz68W5GYf9T`Zb;7Lqz7yU43u-;z6xxMjf_-
z(gz}tpXjvth;B}l<pD8MHX#ojP)>70Ly{3!FAy@?Egts;1qIc$w3v$ZBqc3!pAocw
zO52ND=*;4IzwETPQg!1UL#%OfZqkO_8%A>6O0zy9U{Jbd8!{OTw6w^8u>+o~vZ|_Y
zU;s&;ZEbD!?n4}gxM4KB^2l9;2DlM`>kQ)J;uRQdbaX67pEmavlsgHU7WJJL+r+SG
z2b7#kh*<T@%qhN_0Hd=9DMk-|2pS5^m5Ut6vB%eLp!hv_wX?G$n-Ct_`<YN_hV8`i
zN3iGRE>-<*)d(lE>@&*+a$qgG{LnD6;tziRHX=vI523$nMDf|G=G%b>MZrVn`CrY=
z&3)9b;G70}`_y0fxNYs0k5VKvY|l~J@~jMtafN-Ts31(}f9&GI2La;V1Xi92_<KeT
z4sazinQ#b8^!Ek@_m=yOUMLk>-*+<CGy3+^ZL-Fij3$a{OamMju0`_R=Gx*@0>4Wx
zuLGNWwpVp#vvDsx`szmv0ck)G9!&kBh`HtB`T5|?N8F{!dw2zm;#_1c<R;B4=?nBb
zkx!;qMJsN8F@J5YZd3{Jo$}^Q`iEG}UkPj5Ag8jzkVw-?Q+k%8wki1uF+<Zgre!F~
zcB9&HMNYF+PcU<t7{Z3f0khYx?DNC+B&S8`gpL@Bx3yEXwl`;LZNqg*anaS(Q@)xj
z7!A+PDmXhkGrR?^@oK<BNcMr7J3hb2Z|~&7CFRh=<XZYwA=GL&-z-=XtQ{6AW3yB}
zP)F;J7P7?3!)8moDedG|@6gy6y43mKxC$w4{hS@pg>>wc0AcIrc6$E|3=E8pjXkDO
zt_Abr)>5YR+97ioD_b(0N8io}C`;pU@{du7B_uAMvf%^hqiaWw-V;e3E@JomdAn83
z@We#O!fsw|EuYnSUpzs4Y3y()M0NvF>;+$<WD$q%r%T=8@7;dmp1+AHov5siBCD#b
zB!?fTmWmx)o+(Fg*0^vSxBcwJ$mM3d3NXh>3xgQcewY(MoIy-r@QYv)TwEi47NRRV
zSee_X$xVZFVGPUEV%Q4{)N(VaK6lUdR6p2q$02=GpLQ<R(P2A3t^avh=QL`4p^RmU
zLT-W*JlXJK^+4BDxoy%g--zGe_cDh>brKvtFl9>h!R9qY<-owDXl9|1L%Jukt0RGY
zoM#gOd&5PUSi`OAm`p=T&I1h5NarFmy@}TzLC|^L;_y|+q36b1Q+ekS28w?!d2ExP
z8S&e!2iB#bDkcQlHRe?hAlWk@@Ku}mM;i`4JbDsMV2Jc#jDOvV<q9W6geDK=99yvu
z)dLDyf8X0_-LqILn%Op@<kzj-shT&;#}k2^v}{8p7&@7u-)n!cM3N7-77R$Y8N=<1
z7aV+9d%8%!>o+5Q4IeJ>A}?^S7zm?3J#4?+x^mj*FvGOC+#oj-iX17)gu3;~+Bg}0
z=Z4qurzyFa-kZoL=zVu=&GoXXg^yEj&+fQyWl?=vD){e2We=t+zIfaElncMBJ;_jR
zXL?r~O^Ec+(K!umWOi;S>Ef70WDq9vT3N<wVqWGVAB1Sf{(P?_WO1HG;Bsyc(Qni$
z3fC0g|Gl<aWBTR0->x{;cwIYz7=A`-S#`umW#sxZd_-%3sIfrMNk-P8H56QT*Oq=c
z2l*k!mECiQJ+R3QyZ)>4OD+4V`mBBGNTV>5h@=XzAn>y1H*2A2y?jwGX?Zj0@A7cl
z5nSAXSw9u`x=qmP5;@ElL;lx{?1qjIbu_^XA{ewjWW4%W)e#Kw$i<-;i8d{-!D&5W
z`0UPecGR<2c?2^E%R|mU>flRQ^?;%OUeM#cnyN<~VKUr*S0VA@1sAXEbXi6KK^t+c
z7ty?{ZlU#^ku>CZVZZ*onB4HY^x+=OJQ)AraB>QZSMay6Lq~JMs&|R<8+}py*W!i~
z^6%!d<H0~X^#P~4zPpR@zNjc13E(f!kC?0zU{4=%QUw<I;AJ%6tft|%UGUdmNL;(a
zeGo&GSY2DIs;m3@A4(5JyF~MFcR~5Zi+iqf&RX0qtaNnCzhNSk7^IgZ>0UH^oI&xb
zZf$>JWL`7e;VLL8`A}IIGuPxQ9l!GZ`<JPyJ|NTTPgYZ6E1trVgh*dt)1Fe;c;IUK
zpt$H`_0N~JGQ{Z0VkwBw5CA~8n937Bw*=DW?~Pu#{k(Df%U;6r=4szvQJlm1V=gVr
zo3d!|cHq{4i3!uG1+tI|U;6Lcv)#@G;LwEIhxA;cr^p*gv#_~6EwpC*6`b?tQ$c?I
z``p}M_x;rhj{_@U`Q(A0f-G%xVgk6dgg0}qLwhl4&YwQ7j;JHJ@V>2gzVs#HNOD+F
z6*U&`CZ7WoD8|bShl25L()BBj<YZ;}1AlqPv<{Z?V|ZD+A^~D<NMb1o46ZFdNOD+v
zc=aA+hs!q?j?tl71K|^9d~$LEED{59{T(~MF7GYpEjG0U#47tg)hoBTqCHxnF|4P`
zqz$wT!lDYwMNszuXKl6vuIE+HCH}-5fvC5`O332o8W>U%y`^8@>l>^NO>>F4VO=Jf
zkn9Z)t11a^^ZK05EHGNb)Og0l4gFGR{0zm=u089#ED0oFwxy?FWXW~}#4uH1iV+bJ
zfs2dF#l?lv-W_q%{Z}7T!$x6I4D3}j#rlsI&kq9`Wj()_eOW|xP4Xw|Q-GBR@)q2F
zv2M*>D|fJpBCRqDpu+klOG10!>^HR`m})x)qS}95bGVFqa>%(MPnX2X;ui*0xBPr&
zi>aw+8H%^A%XHg4mb7@|OXFX5PO4emSx~i<u8iOp$JVizbsmspWMc7?_3;sdg}c{3
zJw4q!G^E^Z2PwVmRB5P8qW~X@_?CZH9paI1kYZPP6_uEq2W)c;x$OFxuD-q^(5u&z
zZSW*q$y}P|fM;BCY3gYi09q8{tfi$TMXlG~?<1wWSO_?PY3TbQk-vwFHo^-xK&+}z
z6zNGaP>Y;+0kY)5SgfM;cV6??U}xYS0Q+WPXE*PUX9pw&OfqKsua$nFV}W_AgwkJq
z-PqVz7Z8cS6F{NLc*XC*EyrDIU{G(oeZc&KoXQ4yvh?)_V5OJuk`c<*B;Y^XrZ3+g
zD*y2Y9I@d-`i5>8LkJRweWr%O=hD$>#mqLFZuRrd8B1GR^X2Xc6<0QVPS*xKC>-*y
z?gQ;@X>T8etqowqOewD7lPOePNk~ZG*d+lm9@<-RrT(1dcyM=-0Y^Q9V1@QckEZOK
zmo4#kK1B5yg_E*p<>zA}xPbdt{v;;#y2QWZ7L;DKS><kBf8W&P`Pq>$T7BwVM6VLB
z_2jzVg_zso!qNz>`V+blLkYCl`FR7)V(o{I9z8PCv9@M~Y8_C?HheehW*z0}T(KU=
z%i{>U?!3t{?rb)&|9nuQ$!bvMiX5}x@Yn6*Ki*F^g87I1o5b<>*)<J`k_3sydjJRm
z;rT*<3P@B`%o^=?adFWZ8Or+l)NoR5Cd;pcAUbYn=|#Ui7fMq~Rbm2|ivpSUc(D#5
zp8ctmLQkODeH|U0?D+6dg*a0*$aN6xQ6a;8Hxt^iKiHa~%V`^_YG?$FteY*gdPkct
z5%L18+1j#+XVdW?S~G>6#llhOydtX&h7C{?idZbSXJDsl=PofLYESlkj_U|I|LnKB
zID%J>p~RA3aTB3|B4c)ZO-;>`V?)*-hsSu1^#5uBpa6cjH5Hpbhfmz<SuwHMWmngu
z$23dfS`4dT=5ccUY|Z;gBbvq2#RS&H)jQa7W)||=Mu25u1yuT>5R<)1?)T&}LUQ1Y
z5*Hn8;rHR86gwIhF1>0-W||$E-vfqiUe%K?rV!KL!Q;l}NyO&VoW6ItcV=>pO@w~z
zKCx^|;0H%=km_x3fU<%O?&|7N*#z{=@H@*;U7ci%Ze(O6z?hs(ub!*WARvUeM@@|P
zF-boXa~5ZO`V_~i)8M#bhA%6bTOoS+YLn(w3n$_eK_@yaC$`N@ZG`Pi?PGTNiK$^?
z(nK)TP@vX&BTnI@!NEXq?7{l_`}@DJzijtIL&O6qiDw->RsX_N{W9ji3=9-l08W16
zpdDBpNExx(ste621L)N)ZQPTw_$V(#SLWBK1f3egc+~2i2;(}VF>x_6<-ozIn6iQT
zEJTi*d^YcEYe`>0IbHSX)9{|<%&ID$Qc^58(^V}{2n?-dgTIEKd||SQi3zmK@tDre
zPGi^m{}X>kjk3)D((Xx)1oSEMM^Fq!+>>&H76UC`h*s?UN@Cn^og7K21i+txgJ&B*
z-SSIHy4Lnp^z^9imo#;m09^@S`!hHM;+p<<G81HKlEd<AIDj&Iy#z_#sHYLQe6VFu
z%-O7A%RKtjH#GDN0Pwwj4-XGIdJ9R(+$uBxaO}R=a|cIctP)4{gwYjuRv3LMn3)Nq
z*LRLl&dvzR*DJm(+Z$3PCjcgiK33F$gpxbWxdn4#p1{|Qh1b-2%$;lZo%)0v!#)yB
zBwf8U&i5qQIivSlv;{ATJ*ePZBFUWqSRxurK<qlU)CSEp;^T4YC7#Q|>WE&mJ6rVX
zYHE*MU0t*InCN0^8w~@MZA0w0{E3)meedWsJb4%<L&vaXst)Bhh=YmZzGBuBr5bM`
zb&9;?Qg;pXk1@4iwjEjY9~n)MWH2U*|D@CE<!<H^10_1U))Nb`4bV;G0)gHYPR3-V
zMMtN$ag1VfNk1xzj)t~pNE<W+fh~qa_C%W<xLGKCfj5<gv-q^uCagLKq#~>oO(cY}
zuk#E7S#{rAkkvaHO~mc(0vcsd92VCN2zWO<JS+nfiW&Ta|40!-ua!EI2E&zt_g(vp
zk|IQw+&UKOxGG3*;ly*t!B`XtS9splP8~mI4st>~-fJM}khQu^m_$U$!3HsU<Pujg
zX8c@PyA{Q1PzxulR~?!B^B64N&-Z0}Nw#>Dfxe^4?C}h|0XGkDAJG)n?Sb(Cu1#S|
zIoAL@bJL5S6Y077{-ELSpC<sqW*vED_T5#+^Yin6Pyxs((q8Mo4*ke0V7ou>{zVmg
z4We`mAJMeY)XYpjg!t#rpNFzp+uC+nRH;09@@l|hIsU1g?01gGlHZ7EY@L2??}a~G
zY{Xg8V!iP$G$$qS8lzYSK?vgEvu8|@W4|)ulF`wj0C)eb^>hbW-2zYQ*R<2fs?u}=
zP%5o~zd3t(CN+nb^7EK#P7Z%0p~OQ-ao4yY70xKQS%~l9Y}S^RS3u+vo0c{Zb4MN$
zN(gyDtflpIDovz>VJ!bQUnM6g6<bb99M$s(j{RsQL|Z)$vZ<-5XmgH)a8b`g1}i`)
z2;9W-Wo1B3gW^a!t4!-IGhyFVq!&(mieYVU-?QTQn7vB31j<j6+nSOLpJ)w(hpEd;
z@i^QFw*F8Rf6|(kVo+mXi(7UOfZ|(p56A@|(}i+C_W0)dIt6KaXsCkb9_H*XfwSL7
zi`5-z{Y!c^Z$r&e18e^4ua=?gU98Cp(L`mlVb1Q?{yuTdQ^KSrbOnb?EYeH+9}CGi
zi~N_l?<E=3?i3?Q^kC4GfU^ogEVVEHEH_33%bX0R2?P2Cp8#;<@5ehXemlZ~igLD_
z?*>_mp+0NAKUPB&-6x|Rjs!OVUyXl$>q^Y+`*W}fpNx2y5FbB0GV(}XzQ0_PgfRLi
zCXL@s+CXKyfOf*kTHMy77uW94jXaci<}>g=+(&O;pX~P?NweU4iuIdA4*T>DKbRrb
z0}skyU|Q5$9sQ1F-P6+(f``h;5G#0-zR6Q=y(5&Z3k<zK^Ixtn@~q`FYSP6#$4ZOw
z*`BWXoO~ZAZzK$CY^!PrXn+3brt8rZxtEw{6Jga6mt1FC*vA(=U}p+iSl0In%F6O8
zDxv{*!rl&v9#Mg!3#!xm-(!^-t@!D<A0;x*9^lTNIWXH+EqNXL5A)HD7<8h>!h-8@
zXqi!ZpyZA8fdXuOdpjceTtm`~*LR+l(suv5_+re7!?rZ|eQq^|m#Y&%(&U5P;EeD>
zbvkT1B2>W?BMO<LE1yx}1510hNQ^BfX%Q3X;!`uToh=3?*~Pm=HiQJ`tYwHl`x35~
z8q-o{tFoW+sR*OpM61Mkr&q()mWUf1=7RR{!6ZIwW8o=)b|i}?WH$|}PQ(p236s$A
zY5aD@u88jhS5+?NW%%O^$$=AP7$c%gL<)aYRPNch1Ew4`MQ{nYY-!WEveQv?@-FQ>
zTbdyV@sD+03m)dP*E}r=ecRq{wug5nDQV(<_<-3KiIjVsmKvklh1ViacnOmp5|Jhr
z8@7q5q!2}0uEtF)jWaDu&+Y9rIGpA!(tWU%WC+gTOLuW~T|PPmc#zy9X@-vE10|;!
z-L963VQ8+g+S1d}%O6jG1Lp@+-o5-21xoe7#)~h7ENyJQe+%bS`pxl}0Z-HurPLRN
z=2Do=x$}Umv*}$eVe~zQ-v)wscz8(gI5E)!2o)YG6O3|tS~rqbTnT*eKd_F44mr6;
zu9xS>LrXycd$IkQ>qKCeL>X|Nj=jYAY$hKd*Q5G3ihMj991J|baF53YH`(qaKF%cu
z7Ir(pUKH1FV(~IfW~TYE<xNNm;iDkl1LPNS?M}e<WJN$~Dpg@&ArTxxS+t>(B_7v~
z1Rpe<1$-o#lv?s$QnH)V*zhL6$w{v}6(nS2BVWHp=g+;C!KU`0hJ4#*8+%~<&wn_)
zJ9J5G#_WtVQ5bn6f~y|DWi}RM$OPl$*~Xvo@?Je=eTMMP(0D>e7n=ay3cRMW@=ziY
zAuT!xAZL(d0y7B4Sd4D->n8-$OfQBAE14S(UKU(x=eCR5)2G549{&QJw(RVH?(XhH
zZ|(n)Run(2<H%g0LoK7h`(T?a>aa?6{Pjb*WZOMSh7O351qB6}d3hm4rHxh}odL`S
z0AK;Z+0M3$v$F~$4F8A)5uZ|0QXofx)ax}U3u9wr9|5=5`RTDmqs;WG%dxQ*#NI<4
z*o+tk3slHOlK!J71SOZ1%gdLE;2JHQoT`#iQmiS1jmIY@l%GRgcm@i(RJF6~FO$VG
zIIWTf0D`)+rKP1Y8o9W*bZfas<Bl!!4-D<_4jg$+@V4#AkDSTh8oN&vEh#%;ts{Z4
zZwt}n?92<a07#xW*f=nN7Mv-l0@!pa8H`VH&;yX!(8wGH>J`vv^!4>6sk<8(7<6vj
zjgduid<%GrX-JMQBtfD=S4Pqdo53sI5437<6fw~Qwx+5uUV*o*tRz|xw#)e{s#$cS
zY$Dt;IrgrW5gkv2Wv9L8VhTa0A7bwVo%YlG6BuEGVSAI2m6g?EQw9<fSc$EfI{V>M
zh!6J?X0~0<d@e&fo;0t}@r0ET*tuP&Im2Na48F<fUsYUC8W;MRm*s>y$6Y~kxM#^>
zx*IJo1JF@uj;0j<m+9@p#QyVtN!^ho%gz?jy@3#;%cAmHl+wGZ(^8VEtl4e;oy0W<
ze`P;6>p1YPq6(x9%H)lQy|X3~NE$@Xvla6oJ@cY4LmpVmTiIx_pl|?o0Q59>%eim^
zf^f)C73xLOBg>%Pc~I_imsP9OY~@p2tPn@EnV=VmPdNsq@X~L!_a`x5wn~`Lzl8Pj
z*4WPQ)!NRG*07P|=H~Nd_b0Z02gH3ZK?}~6OA~!n4OG&RJV*JjuQv0|Ngy@CqK@iS
z=T>~a-{~ug$=QM(cWv*3gn+vXWYTnMx5uL_qI!5TlO|&axTGn0gq0c^8lv6hrNH|@
zdP^nlBT|aAPS7;D03jfjfs_uxk9LifhFMy0FgjSk-zw;Z@yIvjTCe=QaESZ`?UKXu
zhNuwj|I(g&hGY*T*^x0(T^>vc!UA~tdeqo*?W5*ClJ!3(uW=8j?;U5%m46^s621n7
zKsltZA|5`_jW1Upgi}FO!$O#p<6@s5_x-<=sFWo4qfWbTbuQbwU~a@Uf`m!0+O96{
zw*_*}qv5m8o>FN0ru!2S_E7=KfIKCeQ3*_x<i-saZNenfGU|YVHna)qy|wf5m00k-
z?U}6%bTP&;LU;43s?^k)!WndWfVvAWgP@a&mj_&rB)57uE*WX`^Z!AgUY?T(j7iYZ
zQ7HL^$ppJ^Yd`m|rrkpSmp-M0ogWE6C62}&UJ8I1{||xs8loex47^+*I=)$JBjS_-
zgggvtS39v9gfovP-k=BWGAW&i?=HrBNN7MqaStS)y>ss2;o<cCoxV(IIl+|CZp|*Z
zcHXG3eX)JrurAh!fnA2sX!b`=JFxu4vQ$lED8m<6e+1gz*+cCXW~go+C?f4%#_<jt
z*HpRjqb$BV+}Y<Pw5<o%XJ^pXT5x>aC_dNa0ZO1#qX$rUvcFaZpjFaV$*fz@5rByq
zIG%I5v&H4TeRJfkEMx)VC@U)?2U-^Sa#L>R?EIG33qmNDxD_EE60Lx&PzL+F1sQ{Z
z)jiVx;>FVJY5*{rE=CBV@aTyKgP3&W!zzr{x|@~rxC@s8dfbMwIbor&IyzuudTB;@
zh(#os^u6>kcXxm<)(6{*aRS3$C^lrnX_e6if&!ulWWIF#EZnvt3j_-u<UadUUdQ@h
zUWUOO=pNnS<vQu+Z|wkrlBhv8`i<*%atIs|^-22{$GcG-=DUCwBpFDZeQwBRg(&6N
zUJwNit#w=qGm)eMe{TcaZ_FYs#-rGr^FM;ed&0*3OiSLNj(}BK=QQ?{wX`hwlhFI-
z=+6uA%s(#p?7fL+8%a2PGmtm`2}}f}Dlz$rVP-kNiB5>4v|Krk5~<gP9gcpXc30Xg
z#y|jb!H4|9o;LM7@(cB=1t7LVdjlb}OL{*Mg@|fg9A{b@oho-{(aS_$$w6Lu1rCes
z)ulG3mM{4#<Wh9pmN3umr7kTA%M{P_pe)zG9hBxpTjYNC{s!)t0bd$9)>4ylI}20C
zDpCEnb4XA^8Ir*`?};4NJ(<qUBe*yRlp&9tSq`xH-2T4lGBdO1Y$WMeZMs;OKzN7B
zs6zHb@$JWp<DY7Ka|<1POdaeT+O;Tbe+L8lI{G014#<O;ZJqulFzLACpjdESZf5p3
z=GiZbUUA|J`#W9_khft0lH3A!g!f)BZnmHYmCrAAAmdV=Jq*B!M>@?IZfpfJU17Jq
zTOd?u7d;aB^ffz=+w<@?N6T+7yGBcDU)(`Z9xMuqdmS1Y&o#U~h*zJ`2D^mpZDhP7
z07LL-c4ces)EOP!?kaBk_u3iSMJ+UCN1F>G(zGkCzLXaSze%(-B_&ZnBN?IZq=*Q6
z|J^k68foG_VL)+|;pb@6GPUDhoKq)b{c*KGzR<_3Ydv7lK?)?$(c@9~EHF|l@QV`=
znR9!c16A%D&la=*V8mFAn$V5JSSC}6dfbf{M6c~%9HkQ&#^LnAh`gX~U;Ob^&BP`2
z24VCl)ZRyiUyLl~WD(83hzxU|U2IQt_<_}qTV{-WObGFGmQsFU3c^hpA)7QP`&l*3
ze<l|&&o)GHn<-l&yM8_8Xxj3BH>wPkhFmd5Zrf%r%vK-Tca2MmRf~iD1Wv0NnjO|{
zOA)z5#NjLaY}(~X_P(d$YC8H*W_ZZX81pj{uU*-fpiTszNyliJNc-oNa`L=i$IB&p
zp4;rR-bbzi9=}N}hJ51sIfOIpeVgJy(3;b@0fa;#DskWQrpqf63k04=h&J!*;~5yW
zu?x%@JkRj!<~``^<|#c_wP4LSL2XOF?R$ey3|08BGejxGu=G{Xvo{x1$fsL}3XU02
zS`IxU4*jM^6%V$*L0LZDLq0520cV7R3t$jPf<H7;#$5Hh0900WWPbb`!x-UiB<}t>
z4+@#+Nu++V($W}EQ}{Kx{&E;PH~`L=={Py@1RDb^hVBgD?;qQm7CXwO{Yy4U991`1
zUt^=B2OFbHvee~$1(*fsCBTkl6aB5;r@`0mzy~#oG?6s~M@Pr}($dH#J=Nu7l-Uas
zJul(e2{e9vstx%pfT(7FxD`ftdkolHYHDjiF9{MPyU)Ksl8@x|{BJbPMK0_{6?rj3
ze*Ey^6#!h{*VdS}1is^?lPyl4E=6%7^2AXz{-XLAnRO5L6jFA^>y+4(|JKW?uUna~
z@9fZ2>E(|5X&FBf*87SYe8n~Hk$we4x9B7j-QV}FJh!wohUNxHGI!>hfOn%uK}vW3
z{KZhhCJ%@?hWh*Eb<<YxdN3B(d|&<Pk#MT^IsNvfnE@H1xT_Sx+d+o*?WGX?!$2qq
z4J6-(tl#|OO9cT9o?*=Nud0GE+WMT3;~)vB-*p=ziLP=;FHTp$#qJZey0Y@P#s-Pc
zM@J=neSNX5@0k9i4K8#!9(CbToYcWjc@2qV=x4K$tS+;BaM!oqbwi?hZ{p|aIW%<K
z>%a3elAt}!H?UFqk!6a#H;LGO8PcKj@?t6gl?jptH;{i~j2m%b2Ik%(z+eGRMJ3w>
z)el7}kLOpa_6mN8>8Lk=s4;}<!+%kfXt8C{>ErqxpW}0RC{R4{7E8?~g$2u*{%`%f
zRWmo`4S@p6>({Tdq{cvv4N?Q7F0FmZVe#lR8#T)9Q_se6Bh-?rDNW*#RMVWP!y2to
z$ISa0T<nEtpz%XL<5u|ahJ@hcrM$5h83Ju!^1%%wVDOQg@~`^BLhQO3YWG>F&jTu-
zUEzL##&~vSdHiF0ujtuQcT}fnh9-Cnx(dZ`tZ@_&hxk3QmNhoMU8Gs^`mX)|u9~~(
ztpt47%;8;MEg*YzRO`Nx@uc<^6cLe(SBa$L)!`#uK6X_>6i|oI)_4_rF8r7m2pe$~
z;HB*wx<tWILUW}cDF7HF`HjmK6$s-i#I+gZwoW$xDz8|UiG$%_Ymg+$NQj&dePOVi
zy9`{xc`pT?nlK6JB3b+<AiQLV;_aqj2+J*Pxk8N{++6`^umf=uzx<M(bfdF%IuCMK
zk$^;2(##Pq)Qg_00z2|w9hE+EG+7vpsBRHdR06X<J>7Qys;B_-5k_AHeUm9ASI0S1
zzBh@1I6eq<w0tlvAR8Zm0UN<#X8%FRXUYzXaJJ{oV=NqFHvcms5sav|cmKH?)s7pb
z<94~XAV2_Q_zP8V1rPM;6;_F7<{_cmHo4a_>Ty)iU;7$`Fh>Z)E$?*h{<ZouZ-;?v
zdxv#@pNcyc;;g^)qHq*qyA2}`Z$}RwS&akD%lmdSbRNaYlV|lO#sXe!u)J%?;1Ers
zYQ+8MeBHeM;;-y!tHohJdJz}75DZ-;iJ!U7VX(G<OK)r{KA?F9Jzf7EAfl{xvtX|O
zi_>92-pW$auMfBUXa%+cJZ(ZNKftKZf`1>X*8y80x8eeHOA8&!>HapXz}M6-ivDx0
zmn6{)1)?ghDb*%|GaV{}P~+!I#W2*)ffyc|BYut1a$~V0K+j1L?PzwNSvo5H)7Df-
z{qN!ZhV@+M&q+LvG^<JHfZvIS5KL=`8d^VBpeh08O8CiA3PC(xOK(bT@oXDxZ;yC&
zCGl(#<EYYl!V$E1&e{%+{}pAyz05e|3t4X={*A4H{zNDNetuCEuDfH%5EJdJNbK|T
zS^crg@yST}g~;*tGZOF=&_`($6a1er+s$^9{FvNz;hY&|h~|}YUFXCO;G6bRR%6@D
zKQZv;Rn@yp1K|1q9Q=}!?|XtS9pUb^c{`TY^=4|^HSu#2J-2c7aaGs{LC{%JAh;4m
zPp=yhu?h&@Suh)~+;!?1LBDk8DAv7dKVAgFu8}}hLCYV*@<Tl)4g6M9vi2b(W&R_j
zyaFR`Vfq;XoHtmQVY{InNS3k@mWxe$kD4uh!KiuaqX$K5Lx)-X!F8s~+~C@><`maK
zMQ8TMAZ<`?s~h<Oa$L8q3Qe)LpLJV;2Pms4Yepw?L)tE@`=dL@-TC`}g@@d*InOg_
zOWS<$J?9O&Rhd9q7VBDe3cYRMZ)09!D-8E$B59+gkiQb9*#4z&HstPsk>D?XXSFln
zB{H5_(Rytf{}9GC6vnvwYl+>%8O~1d*AJXs>H7#gf_8e}03r-MU&7|6)K5Wg77U9`
z7h({yD~k2BNrHds^rv~b02AbEa8TK=F{*v$U78r`sG~u>GzRg<6P@=+Bf@0nN>UZ|
z1{*>6UQ0kLZM!}n8;5j=p*aNbnZ>c}%xc9dj6pZnvmdJ3-(-+3N^&>~HI{)}@pu3-
zD@f#5=(N@Hk%RThybMzyvZtLrz9n)@1`1DpgTI#At4)%~Fj*wMblL=gp)LbkyqV-R
z@~{zxg`>t^{WH>EzCkm2c~?&Z)B5k9DaiEsh<}Lc>@B%=XoV|O+9pdGz_Wt_kf9gw
zV8d!MI?}3iBmbAdH{L7<GfIUwwf}F>$nbJRG?w;k(EyUn!R4sBy8_5t2izIEk%(W~
zkh>)L<k!9VdgAhq_V%91rX#D;mcw{;O+aG|rp#QjuRM3$-}(7{xYv)4F4jHgDlXJ@
zAJ()_L2d;Xib!va$rE#Bxya4t8l6qWdUzEhX+-1rb>)J&k)TUq<RjPUZNR2xyjJ%F
zGyfkneo-eP|1ssr^58x=zsQpl`thM(9~SUM@8aX7X#Q2cr+cn{STn7d6@K1V+j|Md
z>NsUirch`HGP0YVcPL^(VJ-7V%?K43buTem$lz7_(fm-LFUnNxGB;(s!l=WwT-)bS
z{)*X6Ba@1=;A=R1``h-HYdITvI94V<L|*^(DX%6rBN}F{Y9h7yCXuFo_EWR<Y=x<J
zkYkD6-aTe(iwOCWS6v$(ifZ)a#(#L4L&E;xWI&B6ZJuZo<IvKvvI59+F042U4C4)W
zTcRDGi?XwY6t=UKF|FgxOs(UYq7VTgu69(XoXB#dNR#p!W1y~-3BZYa{rI$ExDS>C
zj`nxf)||Tg>~F8J=Q&*T<vGB=SDtm;0)3zH0S=?v{}7hrG`~=kC&N}i2|-+W0CEvX
zdL!5fuA)X)c0hOQ6}O-9m#bZbwK`9py!|4?=<&}dWvZbtJ}7>9rmjWb$)&hLjj}qF
zq5Sme)6w2k$k<<;9se}m+*xSLfDVQPs68u*oqT60{s4w%LM7&v2p^8;(EntamlnUV
z0pe9b<IQno^`qCOYzb@PTN04u;5Hul_1yxx7Q{KglEHe7j*sUy%E2S^yhciepu;=a
zND+Ck{-Zu#-aq-Vu_5PtKSlx(M0|@J7g+{3CZ`f`KtXUZL+eJ>vJFsWK-!T9`&lMK
zVG4>MfESlTht@YXvOqK0wh%S3X%1Kt45-rMU_+8tbGo3kG#9$(pqLlcI4?Q)RnqVg
zC<dO}+vfpWS2QLEfJL(??!hfE0BDD-H0_DZU*|NDU3J~aN#te$$P2y74HVGd5f8nv
zRqx*$P5vqe-m|Kx$a>N-q2F96Z7^}j+RBO>h|lbJsQHdqR7D)8u0zcw5o$`vwpn2z
zAt`-Vy0zz7WtpKeg`oCGg0683JG)Zo1Vy&kl0w5QG;LDvH}qX7rLkROj8%W~gv$gH
zfhtf62--AEY~II1aDnC~Bo288=$CD7_L_7uDG@&wgrSy;y*-dqfTmp?eSKTL`}ZSG
zZYTc`OzyR2ri+n-_N&~XwVOvyx!;7HD@hU+{Ck$~AsdrT+;3ffZaLkCf`S``{9XVq
zWTVx}e;T)8KohZ{fBLi8p_P#GKDTp&9y|}l(arm6f#O%Qy+autBN_it!U-CdM2+o;
zIY>*!+&B}$GAK)^t*s?QvXIcq`m_~lnLyAFBAdd#IfW<%EzHj=Au}DJvrO2+D;B?7
zjQ$--!L^`q2i$WYkgLrGoKU-s(AjZmg2o@J3BH+`nJb_-Qt97z5e66y7N?T2dvw}U
z)3$C#x5i8X%iH_%ZQ*95)QW-xG&6t(2qd2sAir|intZ6FgnxabYwN(H0h&p>YX%w{
zg|=Ky%g5iH?ljp?DxzcTd$gdCkdi7b%H)QGrDmHo1Ef^z;8lc9#h)i$(K`J|byddO
z?Yxb~>~^yl=$%O7x2J)dK-U$B13<P4UG+OSx@u~~eh*MGhRCYCsxu8?IpI3UTDxn8
z)-7O~(a344L+@i2*wKb(S!-*0ObpRZ({5|v7X6)xw7^5{q>fp;I;3$RxD-+b_*SnM
zcMKIl?^Bvm9yIVKW8t-cHi`#Yup3~3UWoIneH>A-=KYzP9oCARvu!ng9$$*>$FID`
zYCqG2RIn9@-b0m?d$0+f+3^K(Gbdb{4a&eNaFwn>83EZf3`cc&S_@=7A<XGc@A4;}
zKn1DO>D?Q~sfB{rOn6>Bj)`)z+M*&>?)8B$Co<j8!fi5W%1s0%i*M!2>t@`vSrvmS
zS%1$LR7SX=+3vHf2qNYrO-P3_jt2|MZcw_xO&R2eE>7q(Mdqi#%Qm@ezxGT4JBLHt
zebD^tmm@U6OH;duj6^s`#{5KDSIo%zRimMFQLKeSbv;)wU<p%GdSoA2ygk%+fhRWJ
zH-n+yrlbge+lPdROVMfc<UTCX3!yY@IVr*+!Q>yvlW-UGGT8=?IzfsgaU`rU0qv!+
zi~A?t&9bVGA9t1CTZB9Z<SWAq3*sA2(v-sUe*+!(Eqqp>Urk3nEJBZ{6o21`@_yXl
zpc1NId+5a5y(0}b)N0AZPn(gbv9e$U^n;;f<4_}g`0(K<cS3)c(~v0tw`(tscIL=~
z37u@*z=w4hajm8x*7m=C{YrRwMM2FpEqjl_h3hAB(KVn-2|Z4<KPes8(p*idrh4I4
zy%)KEkVZ!&Cod25Cn6MtTdzPY151@14_`D9kw}RNKEGng@a!AerLX5v<-6LW_A6&8
z#|xUby-^xII+t|s6v<57l{|8moY=<vgHrO^vw!7^WNfLh<N7>Kz1^K-)X5x9hYOq#
z1$v{d>5LbD&OdXh#iYtR$hD?B5s98sMxg7Kkj}UHE>zikP+j+Uciy(jk)iVJ;zs@D
z(A&@-6~)EHw^9ZUXP@bi4h{{yaC1ApWh>h4EKMX4ic#{D7P{CmAFGtt)C5IG6S{h4
ze|^9h+W}v3F%&Xn#+{UoFzGWRGHmndnwp}Jl9ry^jyJze6GcxNey}CoG$&P}`KW2m
z=RMDUg4_pxH<3s$u_*Y8f%Kl)Bc8_HzqdE5(vgYzj)j}H+g?+1zuEI>XZ)A%zVLza
zy+n~$x+THN8Dj1ilTR6l)&;i?ilnla#GI?=><}-g6ciMGv*i_9EiiSuNM@^efs(+7
zY3Wg}C0#;)M<v@(gESqXBpqRnD$D1gp{AM3@BiLWy-0(KhUVMG26nA2Jvpw4lM^wg
zLGusyGvAZ6KTPBbfM|O8>$?Ntw^u;N|F^$>gbh1WL)#<@#Kv3w=a<5pjrcn8_TC_t
zyn+J0)ni9X<HnTDxSIXlxW<~f=e)0>ITiP%!HE~mPdBPcwnr)?#iga4FH-$B>SZ-F
zY;0`KRHy^2YINV4Q>?A6U0OCtD$iZgG@(hp6hH8oYng<uX%iQh+T99>bK1JPjtjoW
zdsj*{ak@+h`7hh+8D!rwV8P(fhlt|f>>PwZJjxld)Q%5oX%QW&IG?R8OGrzjHshxF
z(TI=WGSVT1(%uUfmtd$6`fdCQcsP_13JaFb)?djbD2TNvl=JdpuK)O0pu4AM?g+kY
z#DNL}+Bv=Zo4rrR_kG;mDH==ZGlrTvzGng+07EhL^<9`F#Sh8c6tb$J%+1TYlxy<s
zANV3I+167)mBp09BYLZ!n?be{w5?`3QXf>o^!S2<gD<R3(^Qh+TQ~gr8DnMh3#Xk<
zCnw3Ndi<6JWEf-9&nvWXB``-f4_cj{#s^JJ=`g=W)rrrvq@;|Hns^Z;aiy^EP?RxE
zsp<5_W53ePbV~*e4HY)h=%KYtW^6Ivi&+g#6%sZMj?UU`&x7b6hCN0SzXJmi#;zXp
zL}BffY^iB!S5OMa?j{ec8roid`-Y10p7~R|IjAhYSk*L+>w#+4j5`bt3~$OnX>_t<
z!&Y#)4<i%P3wQT08@_NaF@l*|!tgGUp6`x=-}l{*%ymd#>boB0sj_VC>@<$Z6sTk8
zTAIKY6Q~V!bX?)%<124$d>s~sg)-q+qE@TThKgcgVIiU>=(<i%_~LHJk3Os*iE`VU
z#Kgqp)YM%^y9>8vi459&&fjlW>$3B|($m+^)Z?H}8ITy^#U&sh&|B#2L=C*o94EI^
z9b-)XmZ7iZI7H&Ma*{_&WL(@0$?UJKYZF#AivsuxBQNtjcoiRUBp~3c9p+l)V!oHE
zC8ehaUq^-3U^`JbLCEcV?&Cuv?sN8_M)#oF70d&eKM~JET$C@*d7qSMsxXs?m^<r-
zSSF(#M28BPNc?6<1GBWU!frq4YT8dZeMBg6i!y?a92fZ?WrT^XEnaf&(&FOc`~f2o
z7J2YetGIwb`^S%xTDM+MVZh=f=RVZXNKH>aEa9QTr@xMQD^4!+33K}ksy9(lruOzj
zP92xdFR!A|MTf3=7pOZ11T712;iARKeJE<tW?Ky=LJ<%Y^kb0Krxz_yw^|e<A`HO#
z8avz&QEN~}pQiLE-W)7`bki9Z56{@u_2}VrGy_(yWtM6|eVIj7*z)pn>GgY|>UXd8
zSI*wQ8*J<*#g6h`wcr*N91&a+>}yoCw3mn9^vfn-i{JK{UB^ZDoS&w`yz28*T7)Jf
zi#b%1PC!ItOJ~*MeFa<5+6(@?rk!DFIuowj9;&#A2Mv1Y4&GPl*l?m{t^#&^&X=OU
zdHdqy>?{?Wk2$Tf2fv!a&&S5%))JGFC@^mWfv7ODMQNF@emAsd+3jHC1_vi+lho}$
zU=x!wGvI%N&2`{V$mEUectkaQ;l+EM_5SPg_j7jH>UZZX@nm=@w{?GDqinmJcAU!m
z<}`w7NCbz37|+x;<uZ_9H@JERN@QMo=^DX?Nnny-2bFpU8@Jcmeq4SrbSO#ZcYf~u
zPVNo^5!R5H{sduE&(D>9vZ0rKwhAMGS<E}k3d|&(JyEpq1tB}l<n>Oj3J^bj{K#)=
zsVYC_gN?-edV|ZcVK(3p&RIuSnWXX?bDf=JqW(%0Jwqy51X!v`1-Yh{mI_D0z5Dlr
z0P<&N(<kO)DUNuQwj`#dJ<k5>Anp5V)4^zlPfAM4L+J%fB@nSxx(!YM>@!-Xep4ti
zGICKsNsD<Sk4LnswKejk5FX0p{Je|3ikP_gkqE}B7UZojXIks)+eb%9<UYi#?PQ1i
zp#LiV==#^#7zM_z>FMO_5}|S*=*)ClwWv)^O+}C3TXRA)Gc)}*SV^$qU}ZADenRs2
zc`_piX-;WqanZZr3ztM^a1qxl*^az#U<b7xXcTE=S|)oIgDq(mHJoZ_+(=UBR$E<t
zrd6UDR9I+tVp74ZTN3|7L*se!cwKs=0^|Rga!9CX*P@Q2M_nIU^i4>OG?D7H2Iir4
vo=zVwwg5X}Pxxb!oq+Vdg9(Fx`dBZ*QgZFj!A-#Wje)__)z4*}Q$iB}Q2aIZ

literal 0
HcmV?d00001

diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer_model_splitting.png b/tensor2tensor/mesh_tensorflow/mtf_transformer_model_splitting.png
new file mode 100644
index 0000000000000000000000000000000000000000..77270b3af99cbf6c94332d07f9950048636fa463
GIT binary patch
literal 20232
zcmce;by!tz*DgwkC?H5khb~&W8x-kQRJuzfr5kAw1zB{7LHty@yOz=+ARrAQB`w{1
zEZ_J2_V-=;I_Ix*HkZL-G3T1|c}Cph9`~40S{jPi@Tl=HFfgtuE6F{;z`$a`z`%Tc
z<ud#vt>SYT{CCM!R{7x-`0~4A840hiIxFe9Vqg%Nq5oq3Ba>%~fpHT<S?=ybuZ*=R
zUq8b~4F}u(^lVpBtod>3RryKf<Q&qnul#%^9FM@0P0t=V84y&KlJ@!J5Py3~OnT{N
zjjX&}TK06@xZ49s=Z80yrEWU}3C6RThapp+Bi^oCvMfvspB!0;-(yM5u378sVVXHu
zO&sMY8-3F-I_)~vq~}|(C!jVYNDN=fn3a|~cjV;cm}G;ZA|oTS32;~t2n3R!ECpUE
zb!nNxD+K$!|KHx!ZCl;@=~36&+xykJEG|7=_1QDN_4W16!tsO~WMt0gM@y`A)V3c>
zO5&#KJY!G1%0GR&Z7)2LeV)W^g!g=AlfjcN0(BP-YG`QSJ6ZYtyVbg^#CvD)>Z()O
z*n{5{JX~B~t<!x?hUsH*(|d*Id>IjBc<R5;O6*1pQ<fg;=(PJW41fR5^yl>#L>@~W
z7QYDmZ0N_2NHa>Z8#m@W>b8&k+&nx6$HtUxZP_I})|5PiSbKS2TI$^5lIOmB?b<a?
zgG&3Jc`rPCd@5SnmwF=u19#cT-e+esS=V$gyztIQOq9QW|NibXUdk^^I|IYR@^|lI
z(#J+5a_H#0xpkURwq^Lrgv;I<DKwH;neXc9$$tCRbkZfOUl9&)tWr!~SXvU`<-HRS
zAf1zwbLb~cO-;T0`}h00I^VYTn^!Gg^t>Kj6}pDv?U%XmbA9rJEtiuFFZJu!uerU}
z0tN;K@effSIoQ~EU0q${jb4=xjg8rk+Y6T!74dA$cix6K9E2yT7o@ZL1qo1`930UG
z)8pe$sl(00X=74qslv@hzJFJ%c0%#~Yb`Uk46{EVkt9RB1Vh$ksYXUd0@H807x;DZ
zb&35AO-!8NNXO}VG3$okzkipXaHqvC8lSxFsqO32AV9rsY<%<b<;#@PFno%^k&$rQ
zG&eUlrjOzL20lJM7`pXd&nDZ;a_Z}8^M;gpDJ@M+O&=d{k?wd&mRD89+McBq{{H>+
zG(e`ew|BrQZZmtJ%JDV_$DUEbWwMlxmvqv8lo&&T^GP<=)`Q=^J)c;Q;qLD0GM}n-
zKRgSVbo4BDx}&E@70@KJ;Q>p?#>x4z(BRP{p?P<IX$DwXQQDa0^Ly{UrR*WxU$%o0
zIZoA5(M@>4M|?Qwz{kgL_Zvt`Of)MSTV7qI*_rH+I?c_^T`uj+Nbc%VA?MO3JDY){
z940ESu2R*sGki3uZ;v~_RrYno4Qp;fLST4sQ1<TKyJ`$cg@uLLjt7>T*OX3|>_5rA
zyME)w4KfM}?fdun=GTv4gK21NpFEkZ-|36yj;-Ghx^vU|Y@mvOS)RG(tHL5iY`X<4
zvbB{}Di;@5ypP7?$4M(5b-nrdoI`WYDOOfi)@Ek!Y|F;fB2Z?>XXhz<J3Bk->AlwE
z<m73O9z7ax2V3#=_06kr)YQ~07po~Q{s#LlKt~vnnwV$>=lv?UR5!KpEPV)m@$;uo
zx&HRsmqP9it+=&)=MoVaMcWB_@k9pJqu#$qpv?dq8(UkPCowTG7d6NG`Sa%wRm<ZO
z6Ri18c6Oh~z?{$*8x|I}Sm`j)CizAs{!GnjA;tR8y?D&tnvINK-SxQ21yO~AD6;xB
zVp)GM{ZP%L?t{~U+I|DE7;e;YhPtA%vhuDFtV+suLB{*{?^7n+iH(ayd}xh}Mg@Df
zmZV&!JvOJM*8Np}vhxzFCV<a1HJ6r__6!a_D9g$Dx_jjJY_z|ti~kxbvWCY#t!V!m
z8#i|h_vfOb*b}S4i3y3tZm-*`3XP$j>*F`d9Pi$3Nr@hMh#lWa4;FCywi>LTu#nJK
z&Z7K$K?mXVu^&Gu%Pze95d<jSWNxFUX^Y?cRAK3NgT;m$Z-tw^{5Y{}k@7P263f}y
znKih7&vyFaE_VFD-KUlY!s%R#?pdEY9X;Jl`4bNm-b>rGm^wcQWL%%D8G_ZFtaM1X
zO<R149p6C|-LbC_z3WeNWb-JCT$vMx#XNO-+PLh~r>>sfULglz13|54HHH@Y4-ie(
zj7^J*FANh9O>Sjy=Rbe`<TLRJ2?-f^ZuSZ8to(72v6S^iTq#^syj113+ILLXT^uf#
zIqIWT&s}-b+e^GIK{T8Afs%1M<X3kQv3ahyQtVCnr?ViAP=I;wP^*dXQf*c&udio5
zTphJM5SaE*#AVU;c4GSNDm+jVn{g4B6=ll%WPE);1;aKssMD)}IoCD~eF-CRH#i=r
z2nh)dkbjKQA7c`FQeCXe!ow3g;aGR=`MY<w#3dwb9UXJa%HqM9AR7Frb+>c#ubZNG
zH4!!6`#rq7zh4@75s<`XpukI+4ZoJ3pU=0d5KeO$3zMANP;>?D#dmtNh9a3&lfDp0
zWW;pR8%V~T9Bu!;JU3UtcyjCE<HtX@y|%n0?SB+wKBtrLFr~Z<JBFbx9zool%#+pg
zPT}O2o2RD*<z+E3u~6A}%NCGW*vJq90<I4x(q9aq!Z2wGalwrJ=I~95Eg6`ZnISI4
zL&~zUVwRVe=hQ0=t*@`wo7s3cL-RVNU(d#7NjbKE&ROn)hlj@o9hpOZt9dmT*VfiP
zvA2I|(CByMa{2P*<+Zh6f4X^Mxt*Mx^2*D}ZWiS7q&ShKD9LgfR=<WcQR=dwa(sNO
zTWL>KP*Cvw`}fzGnb+#a?>dynWs)lf5bNmZ#Ma3~p!jmoivY%F5D{nci8N~TzQ-@V
zmsNi62~*=rw-vZVGd19&DXFQ==UT%CFwEOy=mwgz8yYf8OZnj|*UebrMZE^Aw!S`f
zgBDfz-GG3Aa?U?LcnOYAPFNdyJHZcB4smh|3SQ;rvZSS_fBg8dMNoasLjsoZ`{<~q
z#B5rhN8a&tBqxp2REKchF+CN}KQVcqKS%HHd+6)yn`1G#X7v4uADEncX-1iEROg|V
zI<uKqRh103n)3FN>F}pVzm}Hfx4q2!;<CwcSokf3(_aM#2R~!w>gw+P@cw;pTbqKP
zw8`#IGTRDDuF^?Jk4;zBJx-LH2pCP=;`K&X+|1Khq_cQ7Y9VTh^IIW4_@2c|hi@lo
zen*h9uWoK;InyvUJ!EWZj0g{}-fTSKB?tk7d(NJTT6yy1$<6CmwC>*zGNUXnFE>qF
zOlg0N%RD4>UyJgppkk9Wjj+8bD*PHBR}ci=-QC@uC#8iS9X-W2uYH!W;SIPWt#55@
z?Ir`!+(O}}9E;_pTtO;hJ39`@#m}&#J66-<<4F%{z<|J*V{C6Fe1F$0XWg^n_H)mN
z^%RHX2tjwW%EjOQWf_0O|6q#3NJKz@v~Eh;aj2ww0Z`4et=Sg#0-BGCii)@k3k!U!
z^44(~YNjNqc^2iJv=d$klYi#v+Rx8U+9hRNBy3=JA-kGZO?+q~&vP-(Wpr7+b^S`Q
z&))Cl07V&%%GOp{)6*S1xuq&fu|0oRR{_VW^o8!^_VP;#;rkj($p>y#5}60#!&w41
zH`d?8LH36AN4sm@5Y;uL;;0a^t#!dg{V<iXbAH>DQeii#=Skf9ON!U>Sd9;_CNRe^
z?K{PZz0b*E9`6*IX3JPyTwE?5id4$0A7FflSdCd@qM)i2uN~62Bx~syxT>uuGI-QK
zzx^ejhN^(HC}gE01|7U|yBDa32ApoU3%}{;NEyfsZ>NbfOPeE%iX)a^9Vv)rPw(=M
zatp|OVaiJ>fUwz}|5$2X?JVh7H35MXVg#)=;4!oDbq`2W@)2zc`=mu>w%M)=Pqr1Z
zU+P?R+IzY!jee{bA=wjK{ry{fWu{a7vgY?A_xG-pVtrAS3@iB($7NE9{co9Ku}Dkg
zAb<9-(N;~+Wy-YhzT{(MLhRg@TfD~3#r4YgAp}XmMO9Y*!S)3gNTrDb?uvxCnZIk}
zY6uY&yMu#+PZ|IuT5wax5!1y*>qE$JBfIzb@kmOu<?VKIO<qbtIzq*AM|2v6yFN%!
z7x2T3ei<1lvh>BWzuOHfzT;bOhrBi1qsl{55`xPgmTtK?PC0t+PPWm!BBSb^bE)~8
z@4%gXS{UeVDonG+rw#^`PP`I{?X;Ysi8{^beR%KQy~zWmU$ilA;+f6spF1;U^;5XU
zC@>s!DtuLT$&p8&xHl<@9Gx}u3k!ugI5^tnhQQ#h(YbtZkmNvTT&;WI7eo?BG3Ed{
z<XD|=b8-r%r=_*|5zBG^281PEk_tdT2NKRkSRDil@i}cr$Bh(8${{_^`ZJx!kK4+(
zQgUr8tH6J&;ilcv2_qPD!_38D8Q92PPz4z2>Wa?sDK0is8f|`jg9cbT3ECQ|Zc^$3
zr$<yr93HR;fL&9H2?4m4kJ&$q@YPhXu=t>-t4rx$1g4huS*IvQB7Z~>wdC;2uE}j>
zkZ)%uW$|-Si>2$jnsprjK8wDu=tVit=>6iw#S-UA$wP^EIb`5gGpoQc?E;(v>`Pk<
zd6{4+<L>CP-h)Sv<YD`FjXfScx&a}^C$_wzB2kNtOaSq&t}eZ{R$R5~s}EsBNRDc<
z%Vxs?OWp0H^zL>RRtmf@Hi(>21eHEKD&i{9hGMWnx!Ork98i5WF)?9?r}u}*2O&rI
zF*kx=kB*LRAf^(*?{3R?+s3VYW@d(g&YYYoM5Z>v9daKamW+%Hs>FUb5@zC586RI?
zTL%ZLrM}G0!WT;8D$&iBu%(KS6UuR{>hqI5>gcliMK+kE<GFN9k(A+PayIc9c8h;X
zI=>`w>LDLgsqY#mxFtLeQiezc+bk#|qG)rS)+v|+vfs_?l%9{<I+Ik*X{G%{#~vxA
z)cDyWjsG4&?yrV`nuTUPLEJdz61tQ$u^l!ixo>3jYGOUNyxeUvgC*z=2Bf9=KVF1%
z;$IpLrd2&Wgl#gOwe1qhLd6+n?6)ls2~fAh9LFZh-=p;N&Zl27kELq>LnI|7wY`*z
zRV(~HK2Du!=Cb?44aM8j*x0ySeXx?yo}S$$sJ^p5;Ut22wxSt$Jd|}lc)kg*tCn9l
zyh?E5i*oazOLexj4DapfQ8VNoO$k*CH!JqqHouu6VJ5yX@9y{N<!sH_{zT1ID}H-I
zkxA1VE`!RR*?&mu8je(!?#umpgiA;}tDSZHuxsI(%2BgrCzAQ$!-roossTD^X=z<o
zCwvUI+CCqp@p@!@T*cqt-|Z|lDVDs(-*qmGqAlAiwxg%#Wu0Gb{`9n8&JG)b|M4ex
zWu>F$OK)>?!_(5J2~f>~>ay>&dM?>PI2RNWg4#j&GA{1bOIY7WM)biJ5Qs8HF0A;T
zo}PFOA_m)arH3_8lrtIYU!x!=CXTC74^qa;*DGVmsS+_+k$e%EEseu+^=&hg(0^Y1
zU#;Z-FR%w1>K-`2RaoxRin)ZUUj*PcdAWgEMyYb4NsFB?bF;*FBj+#c`lp(=CZsdg
zet!{K!iv{oOIW*&zPbc6(v4I6dZ!6K<J#(AigRb`h(;%-I{jxky#*NnEih{9>Yj;b
zeuo;p8{(UD-PAxkcXD)03}I<m*@td@0|Pz+>><IcVX}ZBZ&{zPa!`QTD>nIJJw0g7
zy0|L2eMPg%sI4TBsqXQ-YBYH~?6R=qWJlX88e2O%)Y8uV$B)Umh%Sc!2&^~ek<N+@
z!6k~hssbe-SO5T`q>K#Or%#`<7{T5n9q9werEBn|&+tWe9zNj_5(}_lY(I}n<6<U;
z%GtV<gD*~t?aQBq?y&}65qt{nOK$;51qWmKNi&R%jR7>Bhazg>14J?)FQ`cKMQI7K
zf||3t6nA_Y$jHdP4GrB32xt_-y^TPa2eID!R!y<tL*PQE<w`}i@Z9W<h6XV<Hr;bf
zub}UBNjdV-Q8SzEOFQxm5}6v$Q>^B{>b%ps(lo<B>UO6(F{h|#GkXEHU}bg-f2((<
zwKS_aMDy;HOZ7qgtVehV$@9;JwFK$w)4v|?GCgRFRkpTfi6CN#(HohWi8Z4fz`%J6
zs8T+f+$@b{YHDggnvF3eIS!Debe7ZPbCw5hl9O!);$RDPH^V*@7f_`6wyxh?<q7Iv
zYN5h!D?BN1axgvEctk5|$}j1+kw5Cr|3I2IDzB!73R1Phf^&^wwPE3t>%><s{qyOQ
zImGq(>6oa=Pn|tI8O4SWiEhx*gww?~C*j`Q<@Y<D+RDk96}32Bvb3};b~2F|HS%PX
z@H|h;%X-srtnqjphsC%k8Q7@#J)aNV3RO>mC%Uhvr+n`o%3enk&h}1?bYndtykV9m
z)NY#4q>gNwq$z!Yl8TCI(r-uPXw9Z5_wC!8C5f5g<@jpCPSY<(SD6_ZEhgPvTwI=S
zInYL+h&?m!+<cCj6SW10SNgOZfamxFdpd6}WT}OVAw0TI*`pvo5kfX(A|ONf?kYo>
z!pe$UG)mJ@L7Hj_^wg8r{DOjH0#8{f34XfNyzh_FkuD`tm0TSyVsPbn(wim(al?AN
zoehEDAGe5j7eM58JeRrD+FuxU(BS^3f2Pr2Vxr0^RxjxjDKb~xbPyt(y1KgH>#dWs
zfwz_`gE@R3o+2_WD0qwsLKU|!VFjf$TktQit{&D}btffIn=>O2uH8e5v$?p<t{Bq^
zm{o?D4E{7faz$?_*E96SedKwdu1`R*?<EP$0+bhM-fn(d%fuvhXoX2YK)`xL=Q(&_
z$~p7!@UWNUo8|eDk;lJRS6R`X!JjV66JekveK_yT(Rp!(v35>i=-=`^1nCw}PMQzJ
zh#6%`QK%Uut%L{Xu_(Wle_)^tpzeg^<o5iT@83hG9BrJPU*EZT$H8Gep-XY_k{x*4
z6&{9=tW)B~3A*zpNru}yx8~he4-!9(mL@bM#;2xsZf$jH3ph<yKWqIHztDK!krq5Z
z$&55$o1*#|gX9^l-znm-i+QWI3~5w0mi!0Jmv$IhrEN{HwpKZu`&}t?ZqY2n(Ro~I
zmn4a&20;ut8{03tepY<oRJ{AD_%i|C^y@W!<}55yi0pd1a@VO3TBc%KjtCFW9~Db3
z4f_f2Wd7?t=VhGgg<FSXsN!X{Cv3rwC|W<-@=1H_v+<&n@Y^L5Z=ud<b_K8F!+x%q
zEKAM5OfJ)xf%!Qm-<bejNuEv{3RDxh?8qhMeYwB9ZjKshz?rDv22SAhLa8{*&E9ym
zQjr-TQizF5IzlwV6vua^Ni|A6R};5l#F8v|2L7J1=i7&fw*fqEM{`C?t#ll7rBx92
zR+RtVxuBH)B}YGuaiUPCVZQ^^R@}}~fONXt12mh=FaIui*7syC!hdUgXjVr0N|qOE
zhKC9)%=zE*stu~P&$b6XH>jIi=6!ghdHfRdSbEXlw=%wLn;m}r05{SviEB4~)x?`9
zLxL!NjXCY_(mYFHIo%}H@{~DU;Pn%waBL#a*b(hCe*gDQqUMSE@2u1RI1usbPoM4M
z@6bUuglhl2pw>IQl49c*w4qg8!3>5t-3gXXBY#)(R)*!~BJNy_OpBO}?0giHZysCc
zbpGGF*csm`y7QSrzkcqf<$TUojilGB!q>|0hQK~|1_$Pg&ot*ABN>Y_luE0}){2xp
zdKX2|%sPJ&T;p#!WJ8eYp+~O{Pml<Er?JXkmVKuUI}o3BY}J=s+KqnFgW2#OCSP`W
z^m6ui2Z~+V^EExd;0&O%H@Uj`2-%R3-3&&wn4cGJ+KoLchOuPTdQce3;VQH%l8C<F
zI5E~;lF3?zfbY0ec`bnfY?7}|v6#nHz0t;J+2^%FdAbajVSy?x<b}Je$YZXv=GQ~N
z$+S+#5N;G0l)m30V=M~hU@oetygwL&USef~+^bDgS3K>{y3Gzdj^?fox+GZgIQ;#<
z^6v*~Udx-<nk~#u!57BYyo)qF3M_Nc8>ghdO(ct&-eR();0$IkzLsobm)&a!<*%IF
zV5>}4gp<q_nZQeZQlyS<s0t_8m>LE2eP@?1DKL#0n1-ltUR$cT4?kAs`3yhK{?Lxz
z6Hg7Mu@gi3Q@dN1w>zgZa?z)I@WQs0hbTOLO^P@ei9h!_xs^r$X(W((?eCpgt(|dv
zUk{Rz#P2SBcl2RvzbxK}BUSzHdn<&e0<^EeeqJAl|0u!Z_I53`WA1z{9HeYv!H5A_
z3Fs5xw9UJKQ~z%;pk>enykuJYp7T6*@hXPpM53kUO|Hv&&{HUWHm?BbU&+9L22wq+
zkx#q<W53gi{PV{Jkh7|W1|cwFNr};skrlx_EqhH7-GNvDCeBw~mJc6#Ek4F!-bpu0
z%S~NVQDY**D*?I)NY38YFnmnl;Bwx+#Q>lWh`Oex209eT%uG7}H>f$_Ffe+m{RE3f
zRiV)!B_#!vih*x4#N!{T?X6USPE!nhPTG1H3#zl6;8^Y-eYgOskT%W4_Af<Cdx&RZ
zK#u_s2^1+T|6p9X9J0H&*INx)(t6{Nw_FN~;h0ZwBgBsT`Ah(L;~#n)W)cpwpw2?0
zIy8)j>Js39#n%@%P8UU*QU^Ww1q8(W4tc%3y$7s-3Z)DW8d`aIRac>OV*Svskt#Vl
zvST-u#qDVJg@maED_K2etu$Su`-h$;5=~jGc;JmxR8+29!Wx+1@u<tZ7eRO*Ab6(d
zwAhDr5fKppIGufc^amcQzaICfb<4|H`qXQpetR6AP8BZ`RNDzn`dC%Rf9~)iAxQy9
z{Q#1*r>Ccz9xObpOTJ+ZlE$&xY<yzhr(tcnAw5&VQ~&-W*iGGwn-9-O3LQe;CC{}p
zgsgLzmGED5F)=amZu>M;Pi<6NF<czB!}mke`-{El@QtNqXdZ@}DDAUR{7YESN%(tX
zzfR}958E7_FUZahIqliQO$vjEtn9L1sTs>I@&V3Vn`z1ddM}1Oy5}B09n_ZS+A9Rv
z(TEWq9{#@S)y+xo;RP3Yi3*LpQ_DTWrU0qb)KtFik$gQ=nvh+tyC2}K<r<y0HLMv=
zHp{i&F8Xh!dUA9c-yZu<_9n**Lq1H8V*cd-&4m`YRK4;i#L)P%0s@fW+O>g^5h$7&
z_xAR*r=h&(K=ZSAS~m|!@xnT~yHP;+@7iDjL-sYA0$Z)l&oW}dtI~c9`QTnYSi=wp
zBXUjZ0MAdFXapU@UL1T7<2Mu?X8QHGFIPSNN3j_;jff)?bd<C;#KNFE0ra76oyU5D
z&7-~lBkeOd4E$4A*fzf{H|pfz@ak%SwLS#WB5L8ejrO<8d0PPWyIp4kWoR7P?`t6h
z62Dy?0Ta#4%miD_&|p0}I{NkNmzp<?o}M1%2S+XH@Rkyb_V=;CNfb}46A%)T0P+Kb
zsTY{Lz?kzpTM&VWs$444@|BT1C1UFeQr`^2hd9QS^SgrcohJq^E^cR67lMNWziL9N
zV6F<5oq&ia%dH{4qP3^oP|i~NWCxN#lRs8cJ^i@Ktdx@qL^^I2UDHyZy{D4&M7g%+
zWn)nELl-c#tY13accX^&e%%dXV*E3QD{1>JJm*+q9zTwWt6OvK3nN9Zu2?+K)<*SZ
zqB(GDFE7#YpZB3Ld5=iy3q{l6SGl#va&nyfhltbB%~s|8jX;vZCfbb?GC$qaK;{;(
zZ7)gsPoF=t1D&pC4E`}*VK-J{p@{I{Kp+Hm2l)<-Mf_$JAtM2<oy!o;@)6+_Wge@D
zKe~A<D5`(S0?^#EZtliI^WYZKzXuJu*;59bvmVaI&gJ5oE`)8>^|N=r>1hnT!~!b~
zv8Ih{b$MCo<8%`wyHs1*2%^G6s)wE}U!4e8$4&wng3J2VJk~9xRg{)Sp#<vc>UzFq
zX?MMkdu`k6iLGr^ssI}Qd{|o9`ZlxCgz}vs5V?V>Gf7_4@AON4-JA`pR$x#COL7lr
zNsP~r3tAnkh9)OzJwN#n6Q|Ohg?uZ<-B^k6%(x~~b6I-pvHR1+@2`vJT7h|m68+o6
z1dV!R6%_1p>l$)$8b8OD<`!{hmtvb~jPVT~Hw#fbSI3OXejS>pt6yq)6+#w7BEG0U
zkomS_>Ir&{e=LR#X@(Uw@kn+lDDN}f&>KH;H#L1e{8?w{`*$*<k?CpZxsl`G;8<E&
zJ)3_eqF<v*Oib0z(T;{RK@9g)iB46Gy|UH48-q!djgoL!0-6}0qY2$M!5U*|4VRZE
zK$F6G#@`9EYvMmzfP?$S#%RB(-szJl7%zy7e93R$8YqL;9e{OlB;hQ9!c^X~dhP8~
zYRAgY*KNKCtx!X53S1)km)v%CcG#_?g;q8;WsoAEL26LtxG-*4RPl5afpBR()0ADO
zKOB;2#q1<146ot>Ta{z!FC{013;|%drxfbjNk<W=83(L@13o)BFsLvfCdQ;cV)Z-L
zIQ*6AJkCDL9^|0I&K9?8p%)D%6FuD;g>GA{nq2tw@YhWhrkcphasC;-^<*4Maxz1`
zZ>Pk%@0fg%ZJhXO;GaOoA6G+ls>LRU?uGe4Ca}PwguP5_j^3Jzn9U=KIC=G*JESyv
z_Iqg05*NQOr!PJa^i!biu{^)HXbq(xq!7MUZ%KM}WgY~-IuCuAfhN<~#J0`hpRaQ<
z1zBFn!6d0PoOXO;_1*%2_Br*-ujCCyjjmb^<*E<-fo8*f1B2+9&C!Hi1VS4O+~b_8
zQu4YZ*U@erDm9RTJ?m8uAM6zvIK*WcYCrqQ>bUXo@h9HO%2y%xR905<6+pQ}V#3sE
zd}iz=iNR2ZA#8}})T`#vD@TGFWcT>9voluXYBO1SejWg&T>fNcWJgBj#U{YjldSNA
z2AswQtZ8>gfE7sHGU0=)U=rfEPH#z!Yrb>S0bcE5TSIoU*h|cHj}c|};z^<-vm<S?
zLb!-pV-#m!QO%JfW6@41fCei_`(>*i<JdEsVFDUqtJ6eTK02avZ#w)wR5e5@A?y7J
z90MsGy{{NcUN+IidW2J~b!2ZzU*=Z5^?UfU&`$-qUrgcZ)vNcPjN@w31r1ZQ93M2V
z8mF(NMKNVr1u2`E-AWgBfIOi6@S&WQRZcJ&aoqiY3pvrwX~v6CbJo)>3+G10u?iAp
zI#@7pPjjFzF@SJzMO_mC?^C^hpBVU6&_(ooGKz}K&WyOg-1D71(3EYVY1YP5ok>~K
zkIH@%s%BdMboe%`v8dP9O<=?Ar)m>-`em}*S4T>G4>si19>_5<OHl>$%wRK25eCX&
z6)=Q(;4mZRRlJr)<G9#Ce1JX#$YDZy$&_WI=t!pCIf~i3sjVTI;CX_73<z^JwWMj{
z;!}t3i$ROe=IGBt8Te|P-d^j#R`E=Mlqtr={+q2Q=>p_Tb=HWLyQnZQ<Lk<$o<rEY
zhX}t#U;?611<VNkzo^-|S#VC;q}syn;-;Z1rk^^LPr2%VQ#%CRvy@6lcLe=r3$`*j
zW;0bvS@{{!`;@y}f`Z+rldfK-JshjpYOG>`-E~FfUt;+icF?&Cm>abp0g%m7Q=iJ*
zi2LHTSmVrq%nS?{u?jU%D)dQ5u0`o~4mN`v;yE}y$X>d_uhF0bDmd?eJ4hK0$4485
z7aTm@-J!FlW?(Rpvfx-Har|S!WdO7&9v&WUdXU`1%}mo+-f6MvS9K;ucbp3E^4m`l
z&1{AUP&7jg?e`OX)yOMJ{`n9<xRdeohC0Dmf}HUcK=uMo0S&&Pk&!ltP|o6+hQ8|+
z@#Bt8PCOupxd{}*!AW}pZpd9AwJ<#Dq*!hSIe+8h|M3JGn!t5JA+MYWAG!CZWJ{8%
z$}Qr;B%LVn=rsfgM3jYoc6s@CXlB3C8yOrl2W10j7MA&|CLG8qC^WrK=wsW_YJ)mS
zGQ3wOlF8jS$^wT3LF|yPoZLP%G-QEgt*P1H$-#}VT{<kbtz_WsJqVZ)bafQ&P16|i
zTdc7hRr#V17`M!xQ1*WWJWn-_@QHQni}SO`>*Q$R2scrf?BUrtczKKIYhaq=(aU6x
zuHa`GJ)4hx=ia-{HEP~nck!Ci>nzOdI5izB$WU4uiG$wt(eY>#eOZrV6)C~xf^jYk
z2L}gq02&fp+1QAg4hId<v(c)*koOsOno><EV{{J#kqk}WaUGzL!*FNmblHV6;#7C=
zCyl)LJ2$WUPl6vouK4F4HtUA-^K(AT9<^#iZbNPz!AQWv4<9`$0XIoGe)-E;6kgBA
z%tj?dcT65wUZFl{HHy7kh+S!fBQIw9B>R=p16>ae8=$XY3l$jvHd|Cpw3mwmXn=qP
z7mE(8Cup#elM@D=`0o8%2-kk_qx7aQE14kNtl;AX=S|j@^@myrq*_u%iv4tbDl`v4
z@e(?;@?~|e)OINM$q^I)fUv%99zgpxe|()1dgyl!g=0{2k#!n3uQ$&*X8>?od_xh_
zT50@58mh|A$<0l_AfhRKD$f7PBqNgriKondj9+cYa=z^qz;%UgRTY)x&CR@mg5J_b
z8lW2^P%X{PcY63SP`vM}UMf_*xfIf}{B3c&tDt_03DN1fD|)(ScGw&{n-pQXC%<*l
z-$Idmn~hkRIZ1;kAuFpN0M>lc-ob&Qfk1vaw1iUb3;nH0T7G_w4$?Vke3|jggMT6Z
z>YI@w6PetsE<oA3W!AlFgvxT9mTAtKGJ<QEK#_R+LL@XG5xH8S)b+oi_R+1ZrzHGN
zbIr{d(9Bviml=BdYikZ=kz{3U-PUuAr=_E#2u{EI%-vmxK9*HHNFVaHnn_B0qNZhc
z5?9iGGBL58)Gpbiv&^nEHiINlqpy`{)vFFa(oMy((N!q-ZYm2=XPL!_SbI>96c{L~
zs;cfjgAxyH&0_It>Whw^BG9-s&1lNY12(|{x=V4?*53Y7YXLLlLq01EAn<;yd;F`z
zz*m*RqAjie>+9Oe%6(G=GMHY1V9)pM5YeVM#Zzu=q*UC%gFSs~J=4J-Ka}8#LF`aH
z<wzfR`V7=b?3|px-bK)tnp3bdJ~B~@JK9)bgfRC8duuVQc%=c@9w>sD-aKqb((~Z>
zHwznY*}vcdnksQ<_*j2`3QS1cw{B+xk<+(Ced1ieK*Vnb4`~2TB4_!DGxMkdx`x)N
z6YFh)>VDE_QmkZBl@+<>e*m0C7uG*oaDnC4wo9V}u7OXuuM;bW&ZD+wVr7S<y%xoE
zcQg+F_$Kr#SaK_M{6k9|Cl|>y_IUyum#`(mbuMo-n1f?Xh?MxVE_Di&25pitztWY5
zO9l~0>GMN-U@#h+0Lg|*bl<$V)>)VssoOi)oDIG27cXA$p-s}9p5||Pkv~}x*>hj>
z(z|C2lV_Cm4#cOkw{>$Nrnx3r1$A1oGLsUgzP!iN>`9NO$r!kw`c3N+Ges+V67%5o
zxQCW{rlmj+d=_R#?EKg%s!y}rB*PXqq}C`spx%olO{E|?zQ4J;%6R?Cted!yp{^W@
z8(8-#arIHlRy|Xe5rK(5-u@^$Yw@RZ5~-o(A$1;H#zhnm$$*_m?v@AWoNN+MGh9a&
zCup6vwc7sfnG@J;7n&3awQ;-=cho*uR{=Hwy#QWOQ6(cII<O>=3}_h`bZ^hMv){h`
zvyH~7+$9--Y@_$XgL2rtDY(kIwMZIM<OT;DacV!Ht?iRD`|<K9BodioyMmd-V}i!8
zY|subbe;BI(yOo|Ru+h*4?b@Awh-Z&g+HOw>N|sk^sjUqygs|OmO8mxh=q;H7W$V^
z!D$&8ISf0>Grc-JyJElmoMYp8JGS#`#9U!4d|sT6UK)qZt{?L@5bmI(v_k*pwNT^W
z&GiU|{T4gT5GP<aU)4jt?05ok3Ysj?_){JQ@niGSjvioIWiuW+e4-Qoi`ZG#z|c=v
zhf_<fjE-apkxG*A`@ru26RlZm40N$Z-q5G17fN{I(gdF8IF6IJ{n8AMbVpsjvxdxL
zq;c_Z@#uyiFw^mcOWR%@{ry!>s}M+8AW{V|bfz&)R-f08-doC>Gw6oXcl^>N{~I;|
z7&ab$<%y}&0OmfRDXTYK1dX+oi37bZOE=<6AI!`t4+NA@JDyw-O+3<(%WW<-5JM;5
z;o<vG4cXYR0v!(PH0vz-qbX1Z3isWEgQ0KV?id@tc~o~Jj)~@qEXw$B;cm4lSmRUO
zqyjKUeyOCs(^}gh76|tsRIpkfub^aYOjR2Kb?MN+Knrw`3MO+RnL2@pOJ@jP4-9tg
z3uO_>_p3v%WRI>ivp)bj&|-wj-rC=sA-V(j5Gvxwn<~Iz*XWq0|C<E$opX%_zn9C8
zD=tDp&$7CJ5rn)TS8@j7irl)s9LenRw)1L7FHvip&aA8V)>m0rOJY0b1(BFVK4N7C
zpcZ{-srQ>nZ?V-}A2GpldLT_htc)Sl_@bkF|LofDjfh5@xj*(#&e`hE)*!FTE%F1U
z%HL)}cz5c}Rt8hF39gI@PI99dK-0LRJAlmOq=YVsNMWqw&ep3PPA58B6^5U`=vae*
zhd|qn2P*ag<A1ZGkb&d$Amky_(`hoRED=F>nm-x6=pbltBM8{Vt*)SnxfY=~mfssu
zrlE?CQ;nsZyn3aUo!JnclpE(!#j8BdL(+%s_*+q)t;sy$t`LIckUgc%mSEkF9aQ8<
zz6hEfV0_;UQY1nt_D5VxmW}x{7OC{2*#6HlWb=X1QJdjKxi97o{mtzYb32}QN2r`f
zpH}aw0O&>PLn#HKD+?A4{4iNgquPfpFO-l+m(8Yl=F;e{`&|v#<vTo5(LFu3995IE
z%nyy}T7l4RqZ7w;w6wFej*A-VXxiYc{zzF^a|e(sQ-y7HaR=?rB*(Ww?1F~k=F7v$
zEjH-=8w_pyB)1r;gJxIzfabVxL3A4Df-NXfE2lSSgPaQrGC=KZ2^E!<=9HInZAtLL
zfbI#&SwMX35u8iLeP_laNfnIrY)13j*61LG(6(8bT5M@);p?^*uTlo!2wc<JWDU@c
zGP(7Cbrf>m-eN#aF|)Ep0B6wMCI4V-x~X+oCL-`m=$H8Hw5Dvl!Cj~(cu<1sLocq4
z*;n%2jkHaBhE0m%74=yAkx)4Ug*ikqplbQLLG8S<w#FnOAyL}{5*3N5zH&+W|4xZ&
zu87LzJI`im)!_!5T;lYW`fhkRnXU6pq}tt2#@(ECiP}0kBC7&Tpv*oE2vpO;yiRa=
z1p+cFtBt|}K&B`uj5zk=30R*x!aJ38j&&r``Dka0K)d>QdFLl}_3;jG?SB!S(rtDG
zY4$Z(q{=s1NR{^k{^06*i5mur<u@w9R#<M7<i~a%EoJosh{_+I%zo)v$Lzvkox6bO
z+#lqtc95yRh;VL4Z564`Rb_=~HlWv`M;$^Ic)vjfdz|;?A+Hx<sM7?k52o}{p);}p
zO~CZZ$&qk4gyXl*JB$A0Ci)LA*&{p@$}RJKGpwBu^~ebt4&9Zk=p8CUL5)!<(IlR!
zf#x4|tDVi<yg`EgZw6Ab)FE!{r$6=6&4`F|LU6JSOU)yDaDoy+USOWG9<u79c)iw2
z+O6xGKyvY~qMxPsf3t`DbS9%kCYm6T>nRno8!5K#O{1bq?%@G8r<wxVPljH?r-v3>
zEjxmP3Uc{n{wE%lPBZjnAZBO&C`?)=$lKZ3sjRDtnnOj*a~J)KEVQwMl$F!iI9>WD
z%6_brIj1c^)yI4PXlr+S;^I<KAQ!N%(1I1lVM3Ufc*&$O?xbW>tFr2Xq#W4a!56*q
zd!Q?|XeR{D4#M$^7kKc+20U6JFg^dhPw3DiY^^V?|4BoIab>Qf``gd3$2=JUxBYC{
zp<N7x6QYdcQIQq<#iMG=P7=`7$Ql_<Rj+1nR`micoxj!f4BB<h-S_4QnzQ;t+*NPw
z?4rGt$vxZ86k(pY{eX#Sy!h(H4yE7*e*dJtiy0(p8g%rqB<RM92f6(#%eVyBUme7a
z+pkzQ8hwcrW8ZuY#_4w*ew4fo&F`CxjN5T7K=e$}dmS=pC1B4q;Z0Yv%>m`8YydR{
ztP;eZFC>aFv~z>1I?raCb*SjiSxhlls-}H->O70gXBziIaaaS-ZA5QI0j;8kMbX55
ztB!;?Rhu&%Y0Nb2H&3I^8jR$)g3F9?8HLhuo^AFitlLoNJ~$e^by(#)*CjDqbPG*s
zmZ^btc386W#)++N)E)3j)*~Fpwzh6ex{uO#?+g89F-^xbTdSNJD+t7;3wnS=)XwQE
z;%80G?Hu@Ev;9=H^T}nnL9uzl5}<h)_QOcmuVG&-nX{Q-*>I=qv&hRW3V}G3l`6ai
z)&QGjJA^bI2mei$hn9;yrDR|@P$qK}tz!Wf{SS1r!6y*{?+wKGZH(N9t$x@av~iV<
zZOr}tJR`E;kjCP6IQt^&t~-U{slxKHpx;gio{F4Y{#$a=)Uw9Cf@{zgc!7=ELriC?
zZeQtJz17NP80deJP&w)fF%hzLX;o<6>%h@fk_Ya@ao(yyQ9XLl=3ko_?Q<WU>M-)0
z`T}=UYTANt=)Y%W7<l$80`GuG_Kd!|iedU0;pr~5*M-6G$47EZ#4y5;krWr#(kl`e
zpYHVLxTmOJnrJ2f!$XqfP0ro~$E*bfIyxyKU}3#2O7C{YC7an1`8A0EbY9&AH<mEU
zA}5Xu9nvZ2xXQyI5ocLMV=Rhi(h>4WRu?*n;NZ|x>+wag09X~=2+NDPGbUkr#=Lbr
z%7S+P9nd$Xjsmxgma?2<e$_Z^8Acm6QVmbd()bwYLq>;y119LuNHP~c1kaa4YQXt8
zG{a}x{I{YNTM}`&7_L<%JXtP^RC)>x%^lfU`NNHSTSE~o=hT0>nCIF;2PKh8Dv*9%
z$D<WdsDP^d(6U9pBlE*xNH3AuAaCbq1t|uQA6!JC2CC9$8_tU+?NptGsgaRt(5voY
zk7M%P^PSti<CQydExFY|Al^s-XhaxV;<s_TpaY0-rWxTnHd&7~?Tsedz_g&tI5^#c
zA-EQF;^#ad=6oWwb-sf%;()^laafR*Z~O@>B`E=Sx4m@6@BG(2kAV7moQfK%)R)?G
zGc!b4e#z(y?7=y{om{^Chp+m~3|jDB+ms1BL0!;Q1*ESEhjh^h&owcG<@#^!66^XC
zpHBGBx}(>Se~FDZE@8ip71;apQ%&8uV1__Y+MtVzNTwH%6_klC7$tWP&YRvSmFoPP
zAd6IjzI9x}pMO{MVexAf0c4YZPv?X_o%111bq<Z=g@xC+*WWA1f-88?z>Da7gPGrS
zNjeMH25m4M_0T9r{`>*0JCZjy9k%2Td+O#?|DI+O&aY>c=IjJ@ZP@wU>P4)>H1utD
zT!k$kag)IjS;O2M+6t;_^(=&fi<mg-Uur5j>)e2VJY+i|h%m(c=I;Em8Erk8P6JdQ
zsdPJyG)P%ThvNU4fQr0EBmBjF-g1H8U6bMJ!L6|)8!v5z>abs*s(=ju-rGsz5=)Ee
z^`cSy=TnZcSCwwky1gV2hItauo4+e8bbPSDFZ3?iZ1o(7yE2WCi<2Noc>gLL<(P)4
zVB(;&@QYf-U;UFFWU<LVlaq$P(O6)S|CKbRJUn^=0v=FG!JGjaxvz;1PIb4VOM*?0
zhzp+nh%Ubl^!>&q#Do4rY``3g*Vyj{UKnU!EcEXOx7-A+AG+P73$mo-x77V~D3;To
z_p6`r{7)GO7^bp_D+@13zbppc%cta`n~%G~FM&R<Yn*ZfG1V)YZ*pBXQeO58&UtEv
zyN;h8eFHo(ir%(JwV8}7AY7#2PUfJNyM%@L>C-2eiCC?!vOrWlZKG)s7$~3(ipUD1
z*@$QW*}&+;BXW&ORw%1UT@k!gQ=gzw3fwI-%%>@T8kpRZyPmcWCX~7PeA5J{+cg1j
zo?Jd3D^~o#e?&)&=l)^;5gmon`CtNn9I)@~HGX)!p!FUdKg>(YiPV6n1ZKypVfqQ~
zh!)$W=OEs?X$7MV*55V$g%z?$izZjxG^#iER!3uNZj@X1QbWU1U<jo`#X}2OQ}Cdf
z{j%R4$sq;WXq7oYWn^=*zs}c<80~r(zaMn9c`!%SLSY!74NPK0Q}A4FUU7r=ojF!)
zXF^u~{3?uqJa`(Jm>5GJ0g|mM27-RoKdGLm<Rsaozd<Auo##<u+UnyXUtkpM-_#4m
zP!qb^jBO|NpZu%VJj?E0A{`1fp$q*vF;_d&d@_SqCxd-)jtN*Bf%I}U4;Fe3lE9mj
z`VVw<l>qRmd24HvLL;E1<t|nbMny%1_H73CN31!iA0yu2#g?!5UzYZw2N1a_$_n)d
zkPb@2c4HaUiP#_mXuI02nRH?;z0~Ht#tQ-jZkc*j0Dd5WT~sqgXdmnnS8roJk-=~>
zzA|&h>-WnKy*6>b?wj%xh5fs}{p%a(p@^LhRS1GaEgbFdN-=K#-N*-#Sn=Mohg)Dg
zC$p(bro9q#9GgL|XfQAL;3ZbP7@+dNtmns4&a>A+7Zc#aRb2|A2oQ-JlBM+AhvE~0
z{(rF^CV#K&D?4wh#pPhYX6H*c{~Ni1K(fFw#1Y2N(-}e~&Mz%xTzw^9{99I(d3|aB
z265bxT##~_?|CBlBsrRIXj7eoQjfZ!6jn9d?MOviF&5R&aJ32<Y8>9rO~^V{*2IEl
zV~A5hylUGfIHwaEaS}&x$qd?#!>X+URWCsBApviJSe5`a8`N}n{p{NTr{A(96m;k?
zhH#C>iYIgb@8C#ariGp1^pYsg-21ZFGXDkk`U9@c!grt_+#^fLe#@vB<Ghmd*=ud3
zD8A%pdPc?f*=4i%cMJ#HUpUgeMu*q$D3pF?2ZI;8#-e_g=HjgGbgCk8)am@XO+Cug
znL^UYX3NisK2WUq&1zvC-`FEE@boJ<EbPB>EQ%{|1J1&Eo{d<E5xBR^6h$m;;KoFZ
zQqeaC*KLf)fA@jgVS3@Xko|lf<-BA93E_?ZwB=#*=Fv$Z2bDC5lS9Y{#r~Hbk<cto
zRsC*`l*_$Zcjml6w^c#d;Vb19Fh2WzX8Ow(v1`3eiD}CY8m?GsJTOHv`Ts2%O*Zf}
zU+hN@J5MM!r{d@sxl;*GG(5|)DVstK>ohqf*+Ea3U&4*UCg}*uG$$`)IX-y{QBinz
znMM4X|EdAaj|SfxxP&PLL?)P&&!c)hh$xJ!!&Zx&G&L6Fp`cI;>g0ra(s}1z@flEY
zXqv$Hl*#x14fOa+P*G9e=7f!k$BMt?>j;NwI}~p>Y}I^a)htc^FU+GrX9Nn4X)ROa
zO|TKVz~D{ED~iTfZBEZdWjitykWy&R<TqjCElK%5G(PSJaB;t~$x?Kew=w+oqz_SL
z#hf&0GI&lh5@>G(b^%X~^Oq<#zeWh%=UJo(Yn?D!uBsRm6-H~*-2a%RjvW13l{6o8
zN9eI<Bu1{-wA&xFUX2)sg<Q<#*A<iVKY_2~wZ9llyId0!llMF7zG#VKd?v7sK7{tS
zg?Z4wElbpS_A>!dV0;JJpZ~(w(!C=fwB@II%2(1X(m3FTV2k=yfyqk=hJlSc;yUFn
z!cyh~pJofA);MU0k3w<$Ou9y7(qiPfp1MGdmP4`Pun1VDv7o{KXD;RsRJ2k)x#C~I
zqSS)1;;HxvLI4Hhw&6>vk#l+(A2~vko#if^t+4htKcLt7q%r#g=e^<0^nai2RTtH}
zii#U3US4(3NPmH@63bq_lfNA{6a+t;F<s?wbxG6L>ACBj@kdQ&-hA@`p;FI{r+i{?
z6bwH6kAgi-SxSVFF~t95{Ae(lo=KuEMWkWH!qSFfCx=JspLeZrXWDO@e?ddn2NXT;
zV|5Dj<yBR?v$L}Pwg03|O(Ego!-r>dXpJn_RTU`yq~{$X`#wY}iEEi+0rv!bDd<u2
zb%*>-z}nN({ekDM(ND>9@jnT<(aXy{0Gb$xzOP0w{L{__5rue_ofi`;=+5K<VT8pb
zKU&TBMq#0+*<VxU>n4L#W>eEHLbK@*3;p-P?}&LCx%L9R%<MQUp#};+L5KgW8$F_c
z;W#<fB8DxN;d*Q*da^q@D|zNN4c{L+@J5cVR^$h~)9M&p@6{2H1<8uI^DM^JZy~Nn
zY`mCI1Y-c@4`Y11wUU$s3f&6Sl$02=-_)y|uZUi^Ne7lbX=LE{-$NCi{)T4=5F*jf
z1wbH5fF-Y9EA2)ay<$T>B%#3QgW+y*adC5n#SE7&c3PPLzn<C%Lb{*RzMISv66pec
zq88+K=$Q42!o)8wC<ya$zHZByO;+aGZ#i~eF|&zE2XwQ<z9{bMD2v1isOTblGO*%_
z)BdMIM%^P3tY-1XzaIi17z^w1pJxKtwrw5^O-!V(J!J|F$2Yi1hIbzZAV6jdgV!it
z-CDO5Wr;o*=>bj=Cg@UfK_|cnZ5`<TJ<EdQRDqC&d62CB_V@D9Z0zi2z(YP}c-)0X
zb=B1=ARAV%o_6aSa5K3_9A^pYg6z^#lDa7kU>hOpn({{WFQK4=;JVb8vcmII`IqPK
zmnr;=pyRUE(@Uko3o}phm836&+4-J7k?8rQsw&SO7ic5a8WvXf+u7Rkfo{(t$-BW=
zfP=$sa;XOgMu$1MxL}4heQbDGXs&ihOmWEnn!Fqmde>h$xw(abLqg9@i_;R8!HhIa
z<xVJPr7v%OPkms@e-HL<J_!i0G#J;iNILP8hKcyHl9Fd?ssEbBhtcykN!ICO#w=|^
zL&>k>d6jTcL+c*tKzvz)<txj`4M`~ud80q@BYHyl<V-uD&)zSuF9Cdc%csHm(WBw&
zg%;CqZZO}Q<mus&CY;_oS?}#Sxdc=B=;?ebAcIL2a>3kT0!lmJLvD0Ogz0;tf&W^>
zG!0Y!2k#rT3*Pn>3C!*sEiW%qR!LsJVoKST0;7IBONZO@lg_AX^U|KX4-Y|2Wocuh
zn1^EEl96FBHC_r%?GsDeu!G5!QRda!56zaBp#`#a<d<(y#qC(rKL4(?)J0`tZEfvR
z^SVbnr>7__Xv&BQF0=Xt+`OJLc)o~S$@m~w9}bcKZfT(eb)o<^@W$|X1#J@(O;1k|
z3C|5CP_17%FyBA&oAj7)&>pw;_m=`^RkQ1Y2}lu^&RHetETZSFt*y4Sos+zot_hK&
zT5^`cfTewXCD4Orc4EyiUBNXL03U#-X8liPWfF`LpbfKo;eBHxJxD;zD772kR+N_m
zfd}1+XWgvS$HQY|QQZp&9v&x^lMK6D*O+`i*az4aRs3`jr<kK@JA3<VFv^pSx_+cT
z(?YF<ZET`<&-``@>;ViG_H=cLuUyZM>;DYn;6~iZdhi^9f_FN=$D<#*;0crH3ro4t
zoo?ULIxPQf(PP5*BZi=D{Y-I4uBXlOw;MZe$sKohcee~r_1+7o&49|mHhR(Lfuw;G
zACcY|N2H_kcT&!DO-CGVdCbHt`z&7B$IN{+Y0!*&UENwL538(5)=-mq3nL_c7iYt1
z?v~AiJMzVC+b3cXt|~}hQe@Xzny)LD&nts>S}Fg~PgKey;!^%G%o`#`GDQJH8nwUo
zO>O6%TfD2_XFDZ1x0%bX35E-WP1<=aru;c>;!<vAT-|?I8jf6~k*lIH0xtCN&&*4H
z{AiE4P|ZEMRl7mkdrVD2Je4Q=?xv+4y;#&zvr#~vMzrRF+PjChZ^bDQs4vyk*T2-v
zlAc@23Jf=+d?K7KQs^2GKtE=`udv`|6j8XUBkghe!^ksc<LBQ~>2hAi5S5b3^qw;L
zFlEM@DqCaNHZn5uV{db&Nl}uXA&r$WPU+60;gSVpCsXKkoxY;`oaDsOUnx}zo#S2I
zCeZg`^5^I@(V^q{dp3co&C#cxlxn%();)^O7-784?q2_MgQPIyqKN)1852E^$+G$3
zr#FyHbL$gT_f`!@8y5dzh)93)VZ!(x;{D0s(~4Qqz8}JtTmC{R7!Kr()^#XT{;(&4
z;vNe1K^l1>a)>#DB#pdh@L-P0yOLsJK}^bwS=U6=NefqFPWuPeJqDf#I-X;me<+kt
z3{hTtK4F4BEm9+I`(P$8*5Mu@A>nnru<L{oHl<BPQyaV~YPZsR%^~Ql`LsVX)gyN^
zI^cqysYTso+V<@?*8=x`Cpu+X=`DNI9qx^qC?*w*uHMrglMxmXQw%xRncW|K9l6|?
zUr+$+|H{y2Lb$K5&)LmQFy4Pkv5xWTtuJbJewL4Qbl$PyO*i;P936Q()mZ7Vkw^cS
zo~9Q&QY+D=^)}#4?w*6^h{$~&ZvFZ5eHFw!OMa^C2Os0w2R07|?1kZ9o88Xf9OEL!
z_|CTvKF&gw$5>UfE>^oP#Du8j$$4-U$#jH;nw4<2u)uCQ{nMSeu=z*AU7vtO;G4bh
z_W7Cb41;g9!Uskt-?MZpz3Kc9X8{a_s|fnoSkAmpzkg7f;hL9Fo*%CkP54J*@T-$9
zxKLF(oJji$Qp>|1>8$phrT&J6p^X3R(i1_r$<XL%e9@Q?Ew)e*r|H?c%|@D<v(ZzH
z%*!EHA?xE@Zcn*ROnl#z_y6Pnq_k=aE#N*e;86mW-lwr;g{!Zvj=Dx+N*2S^#>VEC
z&!b28ZzuZgE$$p~9MD96-5S!AQQxh(oruSuN?@rM{+><O#3W+1mUMWV85A8A6+(~b
z;(vS8y|$u)!%f_q<cOv&Hcl`imp#VUD@T94dGiL5sS!K14~;{wqdz$Z=T$QS_o^no
zo}WuMo?_1<neHN)9^SqIXbGeF>sOqZm>3vPHJuRN-{1F}h;bdRy$^CAEl%=xtKtNx
z=MED;CMVH*4b-}vJ5CFru}Dg_NP5nR$~Yw^48!n*mL^3#W8=0ze@YL|XJ;`ASAUL`
z*;sVf`|NccoWtIKbEivG`HIg{w!bkY0j5GrOY7B*h`Zuuq*ylg1%#iOYX^mXRL7Ca
zK5JdKeC}MH!Y<{!hh$<Rx!O_s(x)M2ecb`HKT|8`xSm!%4bxw&Tx)0_KBBAfv$efi
zxZ3J&^vgwT=J?@loLiURu@XOO9VQgSGhy9eh0=0cO7q{Xe;U3G2h^#*)7l#0c!6a!
z>OZ9fLtcS_42dc@FW%3BvoYNM%!UW1KPTE(5EHjtV!;cSf{2S8Z6JQ7U^XGFUqFuW
zcHqg|p_R#%^Zx#ROpF)r)eQF`MS`>4v*1=6lH*QBot{dHXZC%|Oq-v#;0b*duf+6i
zbx-SdqSnWUv!a9#J|;M&yVSf?$@$q}KM`KII(_UX9XmB;R7As`)ek)@@XVv%Vv>?y
zMn=d-Mn{Qb+c6C<SKb{vJM+hA9?H}7Hxd&OQ$U#O5DtWp5=P9468_v41r5~p$`Gp(
zGL!+N7r7_!pci9|n(FH2?QK4l91=lwR6MIDOW`Icw6Dv)ef-&bN+sOv4w9*8n<xU+
z$iMDE17R=ntGR`YN!3KyYAvt&(CY62!E4}ESu%n9F^`Z;P8qOGc%j$95)Xbqqs2w<
zrve1XcA1OQ5QPtO6DFUBK{<sIneoHMDCWd}rGP#BbXe!{sC!DT*Z?01@n*_(DW$x;
zJSCQR3~rju7N?CeMCct(@_TG#vhatK9M#r#8Lo70ZmwXp*Sd_C0OcF#<rV2ADNUwi
zI0`-;k!!2<=mwm^V)dmQLRVM!j&zn{g+5%bE*RhR^qZoVi-Rd*j7WtKjRzI_z4+u{
z8oUJGSNGnX<vd{LNx*3C?7YgPOkFz^-@n8v4XX${taUkdpmM^23=aoP&AcSiqwZBr
zjo8}8MxUh&k_nz%@qN{Ob9lvV3&alZuJc8KkbR=Ef$w&>+2Mvh#+?RZMJ5sevj05(
zd<ZzH;+1FJl)~*q*%DnjHZtm?Cp&;jD%`>5*<J)c5q$CC!-qJlMgqY~huF!D?v289
zT1H|%SQ)l(v)RK#%4REU#|h#018l+~A_FUK;cfTS=E0S8_4KxlKVWsbzdHN+HOwhJ
zPcuvQ3MPsd9zO9cGaURwjU|L3;D7B0<KDRZ7kWo59u`OdPwsi~{r84k>76lrt-uA9
z?+b-zEprSMiI~<}x4%DXZP{W&g^dw884I)-fKAGo^2!&J4<tB!X|SpM#M1Y;BH8QE
z!wuiJxLe5ldv%xTWQvdzN66}{3rpJ)4;0D(J#+t4j#+neGxL=!)0EAT_pM^DPdu6O
z;kbPLk6qsDfM;unxPE+hsO50tp4#7Hz%!A|<#=o=td>RSBo!83jJA`As1fzOtk5FC
zCGoVVa@lc@FA)VFv@bFz-`-}v?8vrUhUvYJYb!oJVhvg;0?Y%2H8njcMxCv#tTjJA
zFfPk<y1Zp#vi_x)COW4-z5B;7{UOjzXU?DRfAi=>+U5=yCFj$@D@A0EJDSb@daM+r
zW!-DO)<uPEca3Myo~<Fm#n1pOPTmVX-MrcOzKy~EZ-2kvul{XkZ|`n2(`9YgXWLqn
z?qroJ6Ba>Rjh!)mIXO8TiY+Uyzjm7TH2UhsEn7^A?w&i>Hv?$aihWsIwM=}4yB2Y{
z2DSiK^8-h2-+Wsb;8FAY?e=}CZHXL_(~j?3b@kPaUsAno7fjCH4!djutRDA%UA3w!
zY2yp+&rf$|n@Bl@hOYnf@83RPDFS3J%)E7@immJDqJYp)jV>i%c?3-0PSa8oBd_)L
zP3+L~bSlVl0``-+*!m0=m6R48NH7=#)DQbUgYzbDLl?dU9%aYi>FVdQ&MBb@0G~HB
AoB#j-

literal 0
HcmV?d00001


From 008699a0e9390bee106d3d8f578c0de2cb354988 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 9 Aug 2018 10:59:42 -0700
Subject: [PATCH 0513/2720] Add __init__.py for fixing the python2.7 issue in
 open source version.

PiperOrigin-RevId: 208076695
---
 tensor2tensor/mesh_tensorflow/__init__.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)
 create mode 100644 tensor2tensor/mesh_tensorflow/__init__.py

diff --git a/tensor2tensor/mesh_tensorflow/__init__.py b/tensor2tensor/mesh_tensorflow/__init__.py
new file mode 100644
index 000000000..dba7ece95
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/__init__.py
@@ -0,0 +1,14 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.

From d4040ae13e90e0a4e43c0669d6ce33ecd23cc0eb Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Thu, 9 Aug 2018 11:38:57 -0700
Subject: [PATCH 0514/2720] Exponential schedule for mixing datasets.

PiperOrigin-RevId: 208084245
---
 .../data_generators/multi_problem.py          | 34 ++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index b59a686cd..199eabf18 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -27,6 +27,12 @@
 import tensorflow as tf
 
 
+class MixingSchedule(object):
+  """Available schedules for mixing datasets."""
+  EXPONENTIAL = "exponential"
+  CONSTANT = "constant"
+
+
 class MultiProblem(problem.Problem):
   """MultiProblem base class."""
 
@@ -82,6 +88,10 @@ def get_hparams(self, model_hparams=None):
 
     return self._hparams
 
+  @property
+  def mixing_schedule(self):
+    return MixingSchedule.EXPONENTIAL
+
   def flatten_zip(self, *args):
     """A list of examples to a dataset containing mixed examples.
 
@@ -158,15 +168,37 @@ def dataset(self,
     self.get_hparams()
 
     if is_training:
+      problem_step = tf.get_variable("problem_step",
+                                     shape=[],
+                                     dtype=tf.float32,
+                                     initializer=tf.zeros_initializer(),
+                                     trainable=False,
+                                     use_resource=True)
       dataset_iterators = [d.make_one_shot_iterator() for d in datasets]
 
       def get_next_from_dataset(dataset_iter):
         return dataset_iter.get_next()
 
+      def get_exp_sched_prob():
+        with tf.control_dependencies([problem_step.assign_add(1)]):
+          # TODO(urvashik): Make 5e-8 a parameter.
+          # In the current setup, with about 100 examples per batch on average,
+          # the model converges to 50-50 mixing by ~140k problem steps.
+          return tf.minimum(1. - tf.exp(-5e-8 * problem_step), 0.5)
+
+      def get_const_sched_prob():
+        return 0.5
+
       def mix_data(example):
         del example
+        if self.mixing_schedule == MixingSchedule.EXPONENTIAL:
+          prob = get_exp_sched_prob()
+        elif self.mixing_schedule == MixingSchedule.CONSTANT:
+          prob = get_const_sched_prob()
+        else:
+          raise ValueError("Unknown schedule %s" % str(self.mixing_schedule))
         return tf.data.Dataset.from_tensors(tf.cond(
-            tf.less(tf.random_uniform([]), 0.5),
+            tf.greater(tf.random_uniform([]), prob),
             lambda d=dataset_iterators[0]: get_next_from_dataset(d),
             lambda d=dataset_iterators[1]: get_next_from_dataset(d)
         ))

From 86bace390c19153dd42245835a3b96b41a2f1e53 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 9 Aug 2018 15:05:58 -0700
Subject: [PATCH 0515/2720] Make mtf_toy_model run on Cloud TPU.

PiperOrigin-RevId: 208117966
---
 .../mesh_tensorflow/mtf_toy_model_tpu.py      | 28 +++++++++++++++++--
 1 file changed, 25 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
index 8c7d26f38..9797c439e 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
@@ -53,6 +53,25 @@
     default='',
     help='The directory where the model will be stored.')
 
+# Cloud TPU Cluster Resolvers
+tf.flags.DEFINE_string(
+    'tpu',
+    default=None,
+    help='The Cloud TPU to use for training. This should be either the name '
+    'used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 url.')
+
+tf.flags.DEFINE_string(
+    'gcp_project',
+    default=None,
+    help='Project name for the Cloud TPU-enabled project. If not specified, we '
+    'will attempt to automatically detect the GCE project from metadata.')
+
+tf.flags.DEFINE_string(
+    'tpu_zone',
+    default=None,
+    help='GCE zone where the Cloud TPU is located in. If not specified, we '
+    'will attempt to automatically detect the GCE project from metadata.')
+
 
 class ToyModelInput(object):
   """Wrapper class that acts as the input_fn to TPUEstimator."""
@@ -176,11 +195,13 @@ def metric_fn(tf_logits):
 
 def run_toy_model_tpu():
   """Run a toy model on TPU."""
+  tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
+      FLAGS.tpu, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
+
   iterations_per_loop = FLAGS.iterations
   mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
   config = tpu_config.RunConfig(
-      master=FLAGS.master,
-      evaluation_master=FLAGS.master,
+      cluster=tpu_cluster_resolver,
       model_dir=FLAGS.model_dir,
       save_checkpoints_steps=None,  # Disable the default saver
       save_checkpoints_secs=None,  # Disable the default saver
@@ -204,7 +225,7 @@ def run_toy_model_tpu():
     classifier.train(input_fn=ToyModelInput(), max_steps=next_checkpoint)
     current_step = next_checkpoint
 
-    tf.logging.info('Starting to evaluate.')
+    logging.info('Starting to evaluate.')
     eval_results = classifier.evaluate(
         input_fn=ToyModelInput(),
         steps=156)  # since we have 10000 examples and batch_size = 64 per host
@@ -217,4 +238,5 @@ def main(_):
 
 
 if __name__ == '__main__':
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()

From f7d518a6d623a8e50dbcf15c05f595363836f60b Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Thu, 9 Aug 2018 15:17:56 -0700
Subject: [PATCH 0516/2720] Bug fix in growing the vocab.

PiperOrigin-RevId: 208119731
---
 tensor2tensor/data_generators/multi_problem.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 199eabf18..5623fb18b 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -19,7 +19,6 @@
 from __future__ import print_function
 
 from tensor2tensor.data_generators import problem
-from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.layers import discretization
 from tensor2tensor.utils import metrics
@@ -58,8 +57,7 @@ def add_task_id(self, task, example, encoder):
         example["targets"] = tf.squeeze(example["targets"], axis=[-1])
       elif self.vocab_type == text_problems.VocabType.SUBWORD:
         offset = encoder.vocab_size + len(self.task_list)
-        # An additional +1 because of 0-indexing
-        example["targets"] = offset + example["targets"] + 1
+        example["targets"] = offset + example["targets"]
 
     if task.has_inputs:
       inputs = example.pop("inputs")
@@ -83,6 +81,8 @@ def get_hparams(self, model_hparams=None):
     vocab_size_inc = len(self.task_list)
     vocab_size_inc += self.get_max_num_classes()
     vocab_size = self._hparams.vocabulary["targets"].vocab_size
+    tf.logging.info("Old vocabulary size: %d" % vocab_size)
+    tf.logging.info("New vocabulary size: %d" % (vocab_size + vocab_size_inc))
     self._hparams.target_modality = (registry.Modalities.SYMBOL,
                                      vocab_size + vocab_size_inc)
 
@@ -190,6 +190,7 @@ def get_const_sched_prob():
         return 0.5
 
       def mix_data(example):
+        """Function to mix the different datasets according to a schedule."""
         del example
         if self.mixing_schedule == MixingSchedule.EXPONENTIAL:
           prob = get_exp_sched_prob()
@@ -197,6 +198,9 @@ def mix_data(example):
           prob = get_const_sched_prob()
         else:
           raise ValueError("Unknown schedule %s" % str(self.mixing_schedule))
+        tf.logging.info("Using the %s schedule to "
+                        "train the MultiProblem." % str(self.mixing_schedule))
+
         return tf.data.Dataset.from_tensors(tf.cond(
             tf.greater(tf.random_uniform([]), prob),
             lambda d=dataset_iterators[0]: get_next_from_dataset(d),
@@ -226,14 +230,10 @@ def update_task_ids(self, encoder):
       encoder: this provides the size of the vocab which is used to compute
         the index offset.
     """
-    primary_task = self.task_list[0]
-    id_offset = encoder.vocab_size + text_encoder.NUM_RESERVED_TOKENS
-    if hasattr(primary_task, "additional_reserved_tokens"):
-      id_offset += len(primary_task.additional_reserved_tokens)
+    offset = encoder.vocab_size
 
     for idx, _ in enumerate(self.task_list):
-      # Subtract one to get actual indices in the context of 0-indexing
-      self.task_list[idx].set_task_id(idx + id_offset - 1)
+      self.task_list[idx].set_task_id(idx + offset)
       print(self.task_list[idx].task_id)
 
   def get_max_num_classes(self):

From a4b148fd2fb10473ac2454e902f8ac960c75160f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 9 Aug 2018 16:16:21 -0700
Subject: [PATCH 0517/2720] Make t2t_bleu / bleu_hook robust to unexpected
 unicode line breaks.

PiperOrigin-RevId: 208128811
---
 tensor2tensor/utils/bleu_hook.py      |  7 +++--
 tensor2tensor/utils/bleu_hook_test.py | 43 +++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/bleu_hook.py b/tensor2tensor/utils/bleu_hook.py
index 73ef1081e..1e75a1269 100644
--- a/tensor2tensor/utils/bleu_hook.py
+++ b/tensor2tensor/utils/bleu_hook.py
@@ -201,10 +201,11 @@ def bleu_tokenize(string):
 def bleu_wrapper(ref_filename, hyp_filename, case_sensitive=False):
   """Compute BLEU for two files (reference and hypothesis translation)."""
   ref_lines = text_encoder.native_to_unicode(
-      tf.gfile.Open(ref_filename, "r").read()).splitlines()
+      tf.gfile.Open(ref_filename, "r").read()).split("\n")
   hyp_lines = text_encoder.native_to_unicode(
-      tf.gfile.Open(hyp_filename, "r").read()).splitlines()
-  assert len(ref_lines) == len(hyp_lines)
+      tf.gfile.Open(hyp_filename, "r").read()).split("\n")
+  assert len(ref_lines) == len(hyp_lines), ("{} != {}".format(
+      len(ref_lines), len(hyp_lines)))
   if not case_sensitive:
     ref_lines = [x.lower() for x in ref_lines]
     hyp_lines = [x.lower() for x in hyp_lines]
diff --git a/tensor2tensor/utils/bleu_hook_test.py b/tensor2tensor/utils/bleu_hook_test.py
index 017236912..e291adbac 100644
--- a/tensor2tensor/utils/bleu_hook_test.py
+++ b/tensor2tensor/utils/bleu_hook_test.py
@@ -18,6 +18,12 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+import os
+import tempfile
+import six
+
+from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.utils import bleu_hook
 
 import tensorflow as tf
@@ -58,6 +64,43 @@ def testBleuTokenize(self):
     self.assertEqual(bleu_hook.bleu_tokenize(u"hi, “there”"),
                      [u"hi", u",", u"“", u"there", u"”"])
 
+  def _generate_test_data(self, name, hyps, refs):
+    """Writes test data to temporary files.
+
+    Args:
+      name: str, used for making temp files unique across tests
+      hyps: list of unicode strings serving as translation hypotheses
+      refs: list of unicode strings serving as references
+
+    Returns:
+      hyp_file: path to temporary file containing the hypotheses
+      refs_file: path to temporary file containing the references
+    """
+    assert len(hyps) == len(refs)
+    hyp_file = os.path.join(tempfile.gettempdir(), "{}.hyps".format(name))
+    refs_file = os.path.join(tempfile.gettempdir(), "{}.refs".format(name))
+    for filename, items in zip([hyp_file, refs_file], [hyps, refs]):
+      with (open(filename, "wb")
+            if six.PY2 else open(filename, "w", encoding="utf-8")) as out:
+        content = text_encoder.unicode_to_native(u"\n".join(items))
+        out.write(content)
+    return hyp_file, refs_file
+
+  def testBleuWrapper(self):
+    hyp_filename, ref_filename = self._generate_test_data(
+        "standard", [u"a b a c", u"e f g d"], [u"a b a z", u"y f g d k l m"])
+    bleu = bleu_hook.bleu_wrapper(ref_filename, hyp_filename)
+    actual_bleu = 0.3436
+    self.assertAllClose(bleu, actual_bleu, atol=1e-03)
+
+  def testBleuWrapperWithUnicodeLineSeparator(self):
+    hyp_filename, ref_filename = self._generate_test_data(
+        "unicode-linesep", [u"a b a c", u"e f \u2028 d"],
+        [u"a b a z", u"y f g d k l m"])
+    bleu = bleu_hook.bleu_wrapper(ref_filename, hyp_filename)
+    actual_bleu = 0.2638
+    self.assertAllClose(bleu, actual_bleu, atol=1e-03)
+
 
 if __name__ == "__main__":
   tf.test.main()

From 3d16cfa530bb21d176576765cfb5468705a03a2d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 9 Aug 2018 16:40:58 -0700
Subject: [PATCH 0518/2720] Fix a lint error.

PiperOrigin-RevId: 208132355
---
 tensor2tensor/mesh_tensorflow/mtf_transformer.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
index f9ed35fe2..7c4674a52 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
@@ -90,6 +90,7 @@ def activation_dtype(self):
           % self._hparams.activation_dtype)
 
   def _import_to_batch_by_length(self, x, name, mesh, hparams):
+    del hparams
     x = tf.reshape(x, [self.batch_dim.size, self.length_dim.size])
     return mtf.import_fully_replicated(
         mesh, x, mtf.Shape([self.batch_dim, self.length_dim]), name=name)

From 8db9ab2ea211d76b853fb7a84f5c9876ebc7ef80 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 9 Aug 2018 16:57:52 -0700
Subject: [PATCH 0519/2720] Improve docs & test
 convert_to_{tensor,shape,layout_rules}.

PiperOrigin-RevId: 208134574
---
 .../mesh_tensorflow/mesh_tensorflow.py        | 105 ++++++++++++------
 .../mesh_tensorflow/mesh_tensorflow_test.py   |  83 ++++++++++++++
 tensor2tensor/mesh_tensorflow/mnist.py        |   2 +-
 tensor2tensor/mesh_tensorflow/mtf_model.py    |   2 +-
 4 files changed, 154 insertions(+), 38 deletions(-)
 create mode 100644 tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index bf149997c..9d324623e 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -32,14 +32,16 @@
 
 
 def convert_to_dimension(d):
-  """Convert something to a Dimension.
+  """Converts input to a Dimension.
 
   Args:
-    d: either a Dimension or a pair of (string, int) or None
+    d: Dimension, tuple (string, int), or None.
+
   Returns:
-    a Dimension, or None
+    Dimension or None.
+
   Raises:
-    ValueError: if d cannot be converted to a Dimension
+    ValueError: If d cannot be converted to a Dimension.
   """
   if d is None:
     return None
@@ -53,11 +55,18 @@ def convert_to_dimension(d):
 
 
 class Shape(object):
-  """Shape of a Tensor."""
+  """Shape of a Tensor or Mesh."""
 
   def __init__(self, dims):
-    self._dims = tuple(dims)
-    # verify no repeated dims
+    """Constructs a shape for a Tensor or Mesh.
+
+    Args:
+      dims: List-like of Dimensions.
+
+    Raises:
+      ValueError: If Dimensions are repeated.
+    """
+    self._dims = [convert_to_dimension(d) for d in tuple(dims)]
     if len(set(dims)) != len(dims):
       raise ValueError("Shape must not have repeated dimensions %s" % dims)
 
@@ -118,7 +127,7 @@ def to_string(self):
 
   @property
   def cumprod(self):
-    """cumulative product (exclusive) of dimension sizes."""
+    """Cumulative product (exclusive) of Dimension sizes."""
     return _cumprod(self.to_integer_list)[::-1]
 
   def cumprod_to_tensor_axis(self, cumprod):
@@ -152,24 +161,35 @@ def resize_dimension(self, name, new_size):
 
 
 def convert_to_shape(x):
+  """Converts input to a Shape.
+
+  Args:
+    x: Shape, str, or None.
+
+  Returns:
+    Shape or None.
+
+  Raises:
+    ValueError: If x cannot be converted to a Shape.
+  """
   if x is None:
     return None
   if isinstance(x, Shape):
     return x
   if isinstance(x, str):
     x = _parse_string_to_list_of_pairs(x, seconds_to_int=True)
-  return Shape([convert_to_dimension(d) for d in x])
+  return Shape(x)
 
 
 class LayoutRules(object):
-  """Represents layout of a computation.
-
-  Consists of a set of pairs of strings (tensor_dim_name, mesh_dim_name)
-  """
+  """Represents layout of a computation."""
 
   def __init__(self, pairs):
-    if isinstance(pairs, str):
-      pairs = _parse_string_to_list_of_pairs(pairs)
+    """Constructs a layout.
+
+    Args:
+      pairs: Set-like of string pairs (tensor_dim_name, mesh_dim_name).
+    """
     self._pairs = set(pairs)
 
   def __repr__(self):
@@ -179,12 +199,14 @@ def tensor_dimension_to_mesh_axis(self, tensor_dimension, mesh_shape):
     """Mesh axis associated with tensor dimension (or None).
 
     Args:
-      tensor_dimension: a Dimension
-      mesh_shape: a Shape
+      tensor_dimension: Dimension.
+      mesh_shape: Shape.
+
     Returns:
-      an integer or None
+      Integer or None.
+
     Raises:
-      ValueError: if one Tensor dimension maps to two mesh dimensions.
+      ValueError: If one Tensor dimension maps to two mesh dimensions.
     """
     val = [i for i, mesh_dimension in enumerate(mesh_shape)
            if (tensor_dimension.name, mesh_dimension.name) in self._pairs]
@@ -196,15 +218,17 @@ def tensor_dimension_to_mesh_axis(self, tensor_dimension, mesh_shape):
     return val[0] if val else None
 
   def tensor_layout(self, tensor_shape, mesh_shape):
-    """Compute TensorLayout given a tensor shape and a mesh shape.
+    """Computes TensorLayout given a tensor shape and a mesh shape.
 
     Args:
-      tensor_shape: a Shape
-      mesh_shape: a Shape
+      tensor_shape: Shape.
+      mesh_shape: Shape.
+
     Returns:
-      a TensorLayout
+      TensorLayout.
+
     Raises:
-      ValueError: if two tensor dimensions map to the same mesh dimension.
+      ValueError: If two tensor dimensions map to the same mesh dimension.
     """
     ret = [self.tensor_dimension_to_mesh_axis(d, mesh_shape)
            for d in tensor_shape]
@@ -218,10 +242,19 @@ def tensor_layout(self, tensor_shape, mesh_shape):
 
 
 def convert_to_layout_rules(x):
+  """Converts input to a LayoutRules.
+
+  Args:
+    x: LayoutRules, str, or set-like of string pairs.
+
+  Returns:
+    LayoutRules.
+  """
   if isinstance(x, LayoutRules):
     return x
-  else:
-    return LayoutRules(x)
+  if isinstance(x, str):
+    x = _parse_string_to_list_of_pairs(x)
+  return LayoutRules(x)
 
 
 class TensorLayout(object):
@@ -2994,23 +3027,23 @@ def pretty_print_counters(counters):
 
 
 def _parse_string_to_list_of_pairs(s, seconds_to_int=False):
-  r"""Parase a string into a list of pairs.
-
-  If seconds_to_int, then the second elements are integers, otherwise, they
-  are strings.
+  r"""Parses a string into a list of pairs.
 
-  In the input string, each pair is separated by a colon, and the delimeters
-  between paris are any of " ,.;"
+  In the input string, each pair is separated by a colon, and the delimiters
+  between pairs are any of " ,.;".
 
   e.g. "rows:32,cols:32"
 
   Args:
-    s: a string
-    seconds_to_int: a boolean
+    s: str to parse.
+    seconds_to_int: Boolean. If True, then the second elements are returned
+      as integers;  otherwise they are strings.
+
   Returns:
-    a list of pairs
+    List of tuple pairs.
+
   Raises:
-    ValueError: on badly formatted string
+    ValueError: Badly formatted string.
   """
   ret = []
   for p in [s.split(":") for s in re.sub("[,.;]", " ", s).split()]:
diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
new file mode 100644
index 000000000..2564bb11a
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
@@ -0,0 +1,83 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 Mesh TensorFlow."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+
+import tensorflow as tf
+
+
+class MeshTensorFlowTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(
+      (mtf.Dimension(name="x", size=5),),
+      (("x", 5),),
+  )
+  def testConvertToDimension(self, inputs):
+    dimension = mtf.convert_to_dimension(inputs)
+    self.assertEqual(dimension.name, "x")
+    self.assertEqual(dimension.size, 5)
+
+  def testConvertToDimensionGenericInputs(self):
+    dimension = mtf.convert_to_dimension(None)
+    self.assertEqual(dimension, None)
+    with self.assertRaises(TypeError):
+      mtf.convert_to_dimension(5)
+
+  @parameterized.parameters(
+      (mtf.Shape([mtf.Dimension(name="x", size=4),
+                  mtf.Dimension(name="y", size=8)]),),
+      ("x:4;y:8",),
+      ("x:4.y:8",),
+      ("x:4 y:8",),
+      ("x:4,y:8",),
+  )
+  def testConvertToShape(self, inputs):
+    shape = mtf.convert_to_shape(inputs)
+    self.assertEqual(shape, mtf.Shape([mtf.Dimension(name="x", size=4),
+                                       mtf.Dimension(name="y", size=8)]))
+
+  def testConvertToShapeGenericInputs(self):
+    shape = mtf.convert_to_shape(None)
+    self.assertEqual(shape, None)
+    with self.assertRaises(ValueError):
+      mtf.convert_to_shape("x;4")
+
+  @parameterized.parameters(
+      (mtf.LayoutRules([("d_ff", "model"), ("heads", "model")]),),
+      ("d_ff:model;heads:model",),
+      ("d_ff:model.heads:model",),
+      ("d_ff:model heads:model",),
+      ("d_ff:model,heads:model",),
+      ([("d_ff", "model"), ("heads", "model")],),
+  )
+  def testConvertToLayoutRules(self, inputs):
+    layout_rules = mtf.convert_to_layout_rules(inputs)
+    self.assertEqual(
+        layout_rules._pairs,
+        mtf.LayoutRules([("d_ff", "model"), ("heads", "model")])._pairs)
+
+  def testConvertToLayoutRulesGenericInputs(self):
+    with self.assertRaises(ValueError):
+      mtf.convert_to_layout_rules("d_ff;heads")
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/mesh_tensorflow/mnist.py b/tensor2tensor/mesh_tensorflow/mnist.py
index 4a7a9d5e3..c0460506b 100644
--- a/tensor2tensor/mesh_tensorflow/mnist.py
+++ b/tensor2tensor/mesh_tensorflow/mnist.py
@@ -92,7 +92,7 @@ def model_fn(features, labels, mode, params):
   mesh = mtf.Mesh(graph, "my_mesh")
   logits, loss = mnist_model(features, labels, mesh)
   mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
-  layout_rules = mtf.LayoutRules(FLAGS.layout)
+  layout_rules = mtf.convert_to_layout_rules(FLAGS.layout)
   mesh_size = mesh_shape.size
   mesh_devices = [""] * mesh_size
   mesh_impl = placement_mesh_impl.PlacementMeshImpl(
diff --git a/tensor2tensor/mesh_tensorflow/mtf_model.py b/tensor2tensor/mesh_tensorflow/mtf_model.py
index 9ff075242..8bd15fa39 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_model.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_model.py
@@ -77,7 +77,7 @@ def estimator_model_fn(cls,
     mesh = mtf.Mesh(graph, "my_mesh")
 
     mesh_shape = mtf.convert_to_shape(hparams.mesh_shape)
-    layout_rules = mtf.LayoutRules(hparams.layout)
+    layout_rules = mtf.convert_to_layout_rules(hparams.layout)
     if use_tpu:
       mesh_devices = [""] * mesh_shape.size
       mesh_impl = simd_mesh_impl.SimdMeshImpl(

From d2e64dacff0cf40ed79acd593a1c6d7f0f4498cd Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 9 Aug 2018 17:20:29 -0700
Subject: [PATCH 0520/2720] Move TF version forward (add 1.10, drop 1.8)

PiperOrigin-RevId: 208137816
---
 .travis.yml                              |  6 +++---
 setup.py                                 |  4 ++--
 tensor2tensor/layers/common_attention.py |  5 +++--
 tensor2tensor/layers/common_layers.py    | 19 +++++--------------
 tensor2tensor/layers/vqa_layers.py       |  9 +++++----
 tensor2tensor/models/transformer.py      |  3 ++-
 tensor2tensor/models/transformer_test.py | 15 ---------------
 tensor2tensor/utils/cloud_mlengine.py    |  2 +-
 tensor2tensor/utils/cloud_tpu.py         |  4 ++--
 tensor2tensor/utils/t2t_model.py         | 13 ++++---------
 10 files changed, 27 insertions(+), 53 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 398c15703..e89d0181c 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -7,20 +7,20 @@ env:
     - T2T_PROBLEM=algorithmic_reverse_binary40_test
     - T2T_DATA_DIR=/tmp/t2t-data
     - T2T_TRAIN_DIR=/tmp/t2t-train
-    - TF_LATEST="1.9.*"
+    - TF_LATEST="1.10.*"
     # This is necessary to have gsutil work with Python 2.7
     - BOTO_CONFIG=/dev/null
   matrix:
     # We test against recent versions of TensorFlow and tf-nightly.
     # If updating, also update TF_LATEST above
-    - TF_VERSION="1.8.*"
     - TF_VERSION="1.9.*"
+    - TF_VERSION="1.10.*"
     - TF_VERSION="tf-nightly"
 matrix:
   exclude:
     # We test against all versions in Python 2 but only the latest in Python 3
     - python: "3.6"
-      env: TF_VERSION="1.8.*"
+      env: TF_VERSION="1.9.*"
     - python: "3.6"
       env: TF_VERSION="tf-nightly"
 before_install:
diff --git a/setup.py b/setup.py
index ff69570f3..f868adcfd 100644
--- a/setup.py
+++ b/setup.py
@@ -49,8 +49,8 @@
         'tqdm',
     ],
     extras_require={
-        'tensorflow': ['tensorflow>=1.8.0'],
-        'tensorflow_gpu': ['tensorflow-gpu>=1.8.0'],
+        'tensorflow': ['tensorflow>=1.9.0'],
+        'tensorflow_gpu': ['tensorflow-gpu>=1.9.0'],
         'tests': [
             'absl-py', 'pytest', 'mock', 'pylint', 'jupyter', 'gsutil'
             # Need atari extras for Travis tests, but because gym is already in
diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 6393de99b..672b5d79a 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -33,6 +33,7 @@
 import tensorflow as tf
 
 from tensorflow.python.framework import function
+from tensorflow.python.ops import inplace_ops
 
 # Struct containing the sequences ids and order on a batch (are send to the
 # expert to allow them to compute the bias mask)
@@ -3316,11 +3317,11 @@ def multihead_attention(query_antecedent,
           # the tensor by adding the result of matmul(one_hot,
           # update_in_current_step)
           tmp_k = tf.transpose(cache["k"], perm=[2, 0, 1, 3])
-          tmp_k = common_layers.tf_inplace_ops().alias_inplace_update(
+          tmp_k = inplace_ops.alias_inplace_update(
               tmp_k, decode_loop_step, tf.squeeze(k, axis=2))
           k = cache["k"] = tf.transpose(tmp_k, perm=[1, 2, 0, 3])
           tmp_v = tf.transpose(cache["v"], perm=[2, 0, 1, 3])
-          tmp_v = common_layers.tf_inplace_ops().alias_inplace_update(
+          tmp_v = inplace_ops.alias_inplace_update(
               tmp_v, decode_loop_step, tf.squeeze(v, axis=2))
           v = cache["v"] = tf.transpose(tmp_v, perm=[1, 2, 0, 3])
 
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 851a0222d..514ab9d8b 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -31,17 +31,13 @@
 
 from tensorflow.python.framework import function
 from tensorflow.python.framework import ops
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.ops import inplace_ops
 
 # This is a global setting. When turned off, no @function.Defun is used.
 allow_defun = False
 
 
-# Lazy load inplace_ops
-def tf_inplace_ops():
-  from tensorflow.python.ops import inplace_ops  # pylint: disable=g-import-not-at-top
-  return inplace_ops
-
-
 @function.Defun(
     python_grad_func=lambda x, dy: tf.convert_to_tensor(dy),
     shape_func=lambda op: [op.inputs[0].get_shape()])
@@ -66,13 +62,8 @@ def convert_gradient_to_tensor(x):
 
 
 def is_on_tpu():
-  # Support TF versions 1.5+
-  try:
-    from tensorflow.python.ops import control_flow_util  # pylint: disable=g-import-not-at-top
-    ctxt = tf.get_default_graph()._get_control_flow_context()  # pylint: disable=protected-access
-    return control_flow_util.GetContainingXLAContext(ctxt) is not None
-  except (ImportError, AttributeError):
-    return tf.contrib.framework.get_name_scope().startswith("TPUReplicate")
+  ctxt = tf.get_default_graph()._get_control_flow_context()  # pylint: disable=protected-access
+  return control_flow_util.GetContainingXLAContext(ctxt) is not None
 
 
 def dropout_with_broadcast_dims(x, keep_prob, broadcast_dims=None, **kwargs):
@@ -1644,7 +1635,7 @@ def conv_relu_conv(inputs,
         # the tensor by adding the result of matmul(one_hot,
         # update_in_current_step)
         tmp_f = tf.transpose(cache["f"], perm=[1, 0, 2])
-        tmp_f = tf_inplace_ops().alias_inplace_update(
+        tmp_f = inplace_ops.alias_inplace_update(
             tmp_f,
             decode_loop_step * tf.shape(inputs)[1],
             tf.transpose(inputs, perm=[1, 0, 2]))
diff --git a/tensor2tensor/layers/vqa_layers.py b/tensor2tensor/layers/vqa_layers.py
index 74044429d..587ef3a8e 100644
--- a/tensor2tensor/layers/vqa_layers.py
+++ b/tensor2tensor/layers/vqa_layers.py
@@ -23,10 +23,11 @@
 
 import tensorflow as tf
 
-# pylint: disable=unused-import
+
 from tensorflow.contrib import slim
 from tensorflow.contrib.slim.python.slim.nets.resnet_v1 import resnet_v1_152
-from tensorflow.contrib.slim.python.slim.nets.resnet_v2 import resnet_v2_152
+from tensorflow.contrib.slim.python.slim.nets.resnet_v2 import resnet_v2_152  # pylint: disable=unused-import
+from tensorflow.python.ops import inplace_ops
 
 
 def summarize_tensors(tensor_dict, tag=None):
@@ -249,11 +250,11 @@ def multihead_attention(query_antecedent,
           # the tensor by adding the result of matmul(one_hot,
           # update_in_current_step)
           tmp_k = tf.transpose(cache["k"], perm=[2, 0, 1, 3])
-          tmp_k = common_layers.tf_inplace_ops().alias_inplace_update(
+          tmp_k = inplace_ops.alias_inplace_update(
               tmp_k, decode_loop_step, tf.squeeze(k, axis=2))
           k = cache["k"] = tf.transpose(tmp_k, perm=[1, 2, 0, 3])
           tmp_v = tf.transpose(cache["v"], perm=[2, 0, 1, 3])
-          tmp_v = common_layers.tf_inplace_ops().alias_inplace_update(
+          tmp_v = inplace_ops.alias_inplace_update(
               tmp_v, decode_loop_step, tf.squeeze(v, axis=2))
           v = cache["v"] = tf.transpose(tmp_v, perm=[1, 2, 0, 3])
 
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 5d3733e88..2a7e1a01d 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -38,6 +38,7 @@
 
 import tensorflow as tf
 
+from tensorflow.python.ops import inplace_ops
 from tensorflow.python.util import nest
 
 
@@ -760,7 +761,7 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
 
     next_id = tf.expand_dims(next_id, axis=1)
     decoded_ids = tf.transpose(decoded_ids)
-    decoded_ids = common_layers.tf_inplace_ops().alias_inplace_update(
+    decoded_ids = inplace_ops.alias_inplace_update(
         decoded_ids, i, tf.squeeze(next_id, axis=1))
     decoded_ids = tf.transpose(decoded_ids)
     return i + 1, hit_eos, next_id, decoded_ids, cache, log_prob
diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index d4bb1253a..4450100c9 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -31,12 +31,6 @@
 VOCAB_SIZE = 10
 
 
-def tf_version_has_inplace_ops():
-  # Available in TF 1.8+
-  major, minor = [int(el) for el in tf.__version__.split(".")[:2]]
-  return major > 1 or (major == 1 and minor >= 8)
-
-
 def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
               has_input=True, model_cls=transformer.Transformer):
   if hparams is None:
@@ -273,9 +267,6 @@ def _create_greedy_infer_model(self):
     return model, features
 
   def testGreedySlowTPUVsNonTPU(self):
-    if not tf_version_has_inplace_ops():
-      return
-
     decode_length = 3
 
     model, features = self._create_greedy_infer_model()
@@ -298,9 +289,6 @@ def testGreedySlowTPUVsNonTPU(self):
     self.assertAllClose(slow_tpu_res, slow_non_tpu_res)
 
   def testGreedyFastTPUVsNonTPU(self):
-    if not tf_version_has_inplace_ops():
-      return
-
     decode_length = 3
 
     model, features = self._create_greedy_infer_model()
@@ -321,9 +309,6 @@ def testGreedyFastTPUVsNonTPU(self):
     self.assertAllClose(fast_tpu_res, fast_non_tpu_res)
 
   def testGreedyTPUSlowVsFast(self):
-    if not tf_version_has_inplace_ops():
-      return
-
     decode_length = 3
 
     model, features = self._create_greedy_infer_model()
diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 88e87778d..443aff5dc 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -33,7 +33,7 @@
 FLAGS = tf.flags.FLAGS
 
 CONSOLE_URL = "https://console.cloud.google.com/mlengine/jobs/"
-RUNTIME_VERSION = "1.8"
+RUNTIME_VERSION = "1.9"
 
 # TODO(rsepassi):
 # * Enable multi-machine sync/async training
diff --git a/tensor2tensor/utils/cloud_tpu.py b/tensor2tensor/utils/cloud_tpu.py
index 7b8b38037..c7edac244 100644
--- a/tensor2tensor/utils/cloud_tpu.py
+++ b/tensor2tensor/utils/cloud_tpu.py
@@ -162,8 +162,8 @@ def cloud_tpu(vm_name, tpu_name, delete_on_done=False, skip_confirmation=False):
 class Gcloud(object):
   """gcloud command strings."""
   # Note these can be modified by set_versions
-  VM_VERSION = "tf-1-8"
-  TPU_VERSION = "1.8"
+  VM_VERSION = "tf-1-9"
+  TPU_VERSION = "1.9"
 
   @classmethod
   def set_versions(cls, vm, tpu):
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 4313cb02a..a83ccface 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -41,6 +41,7 @@
 import tensorflow as tf
 
 from tensorflow.python.layers import base
+from tensorflow.python.ops import inplace_ops
 from tensorflow.python.ops import variable_scope
 
 _no_problem_err_str = (
@@ -52,12 +53,6 @@
     lambda method_name: _no_problem_err_str % (method_name, method_name))
 
 
-# Lazy load inplace_ops because moudle is only available in TF 1.8+
-def tf_inplace_ops():
-  from tensorflow.python.ops import inplace_ops  # pylint: disable=g-import-not-at-top
-  return inplace_ops
-
-
 class T2TModel(base.Layer):
   """Abstract base class for models.
 
@@ -792,15 +787,15 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
       else:
         cur_sample = samples[:, i, :, :]
       samples = tf.transpose(recent_output, perm=[1, 0, 2, 3])
-      samples = tf_inplace_ops().alias_inplace_update(samples, i,
-                                                      tf.to_int64(cur_sample))
+      samples = inplace_ops.alias_inplace_update(samples, i,
+                                                 tf.to_int64(cur_sample))
       samples = tf.transpose(samples, perm=[1, 0, 2, 3])
       if not tf.contrib.eager.in_eager_mode():
         samples.set_shape([None, None, None, 1])
 
       # Assuming we have one shard for logits.
       recent_logits = tf.transpose(recent_logits, perm=[1, 0, 2, 3, 4])
-      recent_logits = tf_inplace_ops().alias_inplace_update(
+      recent_logits = inplace_ops.alias_inplace_update(
           recent_logits, i, tf.squeeze(logits[:, -1:], axis=1))
       logits = tf.transpose(recent_logits, perm=[1, 0, 2, 3, 4])
       loss = sum([l for l in losses.values() if l is not None])

From c36b081e47e799c695814b4ae64a3f2449341f09 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 9 Aug 2018 18:01:35 -0700
Subject: [PATCH 0521/2720] Internal change

PiperOrigin-RevId: 208142594
---
 tensor2tensor/layers/common_layers.py         |  38 ++++++
 .../common_message_passing_attention.py       | 122 +++++++++++++++++-
 2 files changed, 156 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 514ab9d8b..1e2fe001a 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3742,3 +3742,41 @@ def kl_divergence(mu, log_sigma):
   return kl / tf.to_float(batch_size)
 
 
+def sparse_equals_constant(constant, tensor):
+  return tf.SparseTensor(
+      indices=tensor.indices,
+      dense_shape=tensor.dense_shape,
+      values=tf.equal(tensor.values, constant))
+
+
+def sparse_expand_dims(tensor, current_num_dims, axis=0):
+  if axis == -1:
+    axis = current_num_dims
+
+  new_col = tf.zeros([tf.shape(tensor.indices)[0]], dtype=tf.int64)
+  cols = tf.unstack(tensor.indices, axis=1, num=current_num_dims)
+  shape = tf.unstack(tensor.dense_shape, num=current_num_dims)
+  new_indices = tf.stack(cols[:axis] + [new_col] + cols[axis:], axis=1)
+  return tf.SparseTensor(
+      indices=new_indices,
+      values=tensor.values,
+      dense_shape=tf.stack(shape[:axis] + [1] + shape[axis:]))
+
+
+def sparse_add_constant(constant, tensor):
+  return tf.SparseTensor(
+      indices=tensor.indices,
+      values=constant + tensor.values,
+      dense_shape=tensor.dense_shape)
+
+
+def sparse_eye(size):
+  indices = tf.cast(tf.stack([tf.range(size), tf.range(size)]), tf.int64)
+  values = tf.ones(size)
+  dense_shape = [tf.cast(size, tf.int64), tf.cast(size, tf.int64)]
+
+  return tf.SparseTensor(
+      indices=indices,
+      values=values,
+      dense_shape=dense_shape
+  )
diff --git a/tensor2tensor/layers/common_message_passing_attention.py b/tensor2tensor/layers/common_message_passing_attention.py
index af005c35d..5de2adeef 100644
--- a/tensor2tensor/layers/common_message_passing_attention.py
+++ b/tensor2tensor/layers/common_message_passing_attention.py
@@ -153,7 +153,7 @@ def make_edge_vectors(adjacency_matrix,
   """Gets edge vectors for the edge types in the adjacency matrix.
 
   Args:
-    adjacency_matrix: A [batch, num_nodes, num_nodes] tensor of ints.
+    adjacency_matrix: A [batch, num_nodes, num_nodes, num_edge_types] tensor.
     num_edge_types: Number of different edge types
     depth: Number of channels
     name: A optional string name for scoping
@@ -363,6 +363,123 @@ def compute_mpnn_qkv(node_states,
   return q, k, v
 
 
+def sparse_message_pass_batched(node_states,
+                                adjacency_matrices,
+                                num_edge_types,
+                                hidden_size,
+                                name="sparse_ggnn_batched"):
+  """Identical to sparse_ggnn except that each input has a batch dimension.
+
+  B = The batch size.
+  N = The number of nodes in each batch.
+  H = The size of the hidden states.
+  T = The number of edge types.
+
+  Args:
+    node_states: Initial states of each node in the graph. Shape: [B, N, H]
+    adjacency_matrices: Adjacency matrices of directed edges for each edge
+      type and batch. Shape: [B, N, N, T] (sparse).
+    num_edge_types: The number of edge types. T.
+    hidden_size: The size of the hidden layer. H.
+    name: (optional) The scope within which tf variables should be created.
+
+  Returns:
+    The result of one round of message-passing of shape [B, N, H].
+  """
+
+  b, n = tf.shape(node_states)[0], tf.shape(node_states)[1]
+
+  # Flatten the batch dimension of the node states.
+  node_states = tf.reshape(node_states, [b*n, hidden_size])
+
+  # Flatten the batch dimension of the adjacency matrices.
+  indices = adjacency_matrices.indices
+  new_index2 = indices[:, 3]  # The edge type dimension.
+
+  # Offset N x N adjacency matrix by the batch number in which it appears.
+  new_index0 = indices[:, 1] + indices[:, 0] * tf.cast(n, tf.int64)
+  new_index1 = indices[:, 2] + indices[:, 0] * tf.cast(n, tf.int64)
+
+  # Combine these indices as triples.
+  new_indices = tf.stack([new_index0, new_index1, new_index2], axis=1)
+
+  # Build the new sparse matrix.
+  new_shape = [tf.cast(b*n, tf.int64), tf.cast(b*n, tf.int64), num_edge_types]
+  adjacency_matrices = tf.SparseTensor(indices=new_indices,
+                                       values=adjacency_matrices.values,
+                                       dense_shape=new_shape)
+
+  # Run a message-passing step and return the result with the batch dimension.
+  node_states = sparse_message_pass(node_states, adjacency_matrices,
+                                    num_edge_types, hidden_size, name)
+  return tf.reshape(node_states, [b, n, hidden_size])
+
+
+def sparse_message_pass(node_states,
+                        adjacency_matrices,
+                        num_edge_types,
+                        hidden_size,
+                        name="sparse_ggnn"):
+  """One message-passing step for a GNN with a sparse adjacency matrix.
+
+  Implements equation 2 (the message passing step) in
+  [Li et al. 2015](https://arxiv.org/abs/1511.05493).
+
+  N = The number of nodes in each batch.
+  H = The size of the hidden states.
+  T = The number of edge types.
+
+  Args:
+    node_states: Initial states of each node in the graph. Shape is [N, H].
+    adjacency_matrices: Adjacency matrix of directed edges for each edge
+      type. Shape is [N, N, T] (sparse tensor).
+    num_edge_types: The number of edge types. T.
+    hidden_size: The size of the hidden state. H.
+    name: (optional) The scope within which tf variables should be created.
+
+  Returns:
+    The result of one step of Gated Graph Neural Network (GGNN) message passing.
+    Shape: [N, H]
+  """
+  n = tf.shape(node_states)[0]
+  t = num_edge_types
+
+  # Convert the adjacency matrix into shape [T, N, N] - one [N, N] adjacency
+  # matrix for each edge type. Since sparse tensor multiplication only supports
+  # two-dimensional tensors, we actually convert the adjacency matrix into a
+  # [T * N, N] tensor.
+  adjacency_matrices = tf.sparse_transpose(adjacency_matrices, [2, 0, 1])
+  adjacency_matrices = tf.sparse_reshape(adjacency_matrices, [t * n, n])
+
+  # Multiply the adjacency matrix by the node states, producing a [T * N, H]
+  # tensor. For each (edge type, node) pair, this tensor stores the sum of
+  # the hidden states of the node's neighbors over incoming edges of that type.
+  messages = tf.sparse_tensor_dense_matmul(adjacency_matrices, node_states)
+
+  # Rearrange this tensor to have shape [N, T * H]. The incoming states of each
+  # nodes neighbors are summed by edge type and then concatenated together into
+  # a single T * H vector.
+  messages = tf.reshape(messages, [t, n, hidden_size])
+  messages = tf.transpose(messages, [1, 0, 2])
+  messages = tf.reshape(messages, [n, t * hidden_size])
+
+  # Run each of those T * H vectors through a linear layer that produces
+  # a vector of size H. This process is equivalent to running each H-sized
+  # vector through a separate linear layer for each edge type and then adding
+  # the results together.
+  #
+  # Note that, earlier on, we added together all of the states of neighbors
+  # that were connected by edges of the same edge type. Since addition and
+  # multiplying by a linear layer are commutative, this process was equivalent
+  # to running each incoming edge through a linear layer separately and then
+  # adding everything at the end.
+  with tf.variable_scope(name, default_name="sparse_ggnn"):
+    final_node_states = common_layers.dense(
+        messages, hidden_size, use_bias=True)
+
+  return final_node_states
+
+
 def multihead_mpnn_attention(node_states,
                              total_key_depth,
                              total_value_depth,
@@ -541,9 +658,6 @@ def dot_product_mpnn_attention(q,
     ValueError: if num_transforms doesn't equal num_edge_types and not using
       weighted sum.
   """
-  # TODO(jfrankle): Consider ways to handle graphs that have multiple edges
-  # between the same nodes (with only one edge of each type. adjacency_matrix
-  # will need to be converted to shape [B, T, N, N].
   with tf.variable_scope(
       name,
       default_name="dot_product_mpnn_attention",

From ab420b5ab080f3bc10ee126b9adbe66dcac555cb Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 9 Aug 2018 18:02:54 -0700
Subject: [PATCH 0522/2720] Second round to fix pylint errors in
 mesh_tensorflow.

PiperOrigin-RevId: 208142742
---
 tensor2tensor/mesh_tensorflow/mesh_tensorflow.py     | 10 +++++-----
 tensor2tensor/mesh_tensorflow/mnist.py               |  6 +++---
 tensor2tensor/mesh_tensorflow/mtf_model.py           |  1 +
 tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py   |  9 +++------
 tensor2tensor/mesh_tensorflow/placement_mesh_impl.py |  4 ++--
 5 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index 9d324623e..7acb83b09 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -21,8 +21,8 @@
 from functools import reduce  # pylint: disable=redefined-builtin; for py3
 from operator import mul
 import re
-from past.builtins import xrange
 import six
+from six.moves import xrange  # pylint: disable=redefined-builtin
 
 from tensor2tensor.mesh_tensorflow import mtf_utils
 import tensorflow as tf
@@ -1088,7 +1088,7 @@ def lower(self, lowering):
     for t in self.inputs + self.outputs:
       layout = mesh_impl.tensor_layout(t)
       for d, mesh_axis in zip(t.shape.dims, layout.tensor_axis_to_mesh_axis):
-        if (mesh_axis is not None and d not in self._splittable_dims):
+        if mesh_axis is not None and d not in self._splittable_dims:
           raise ValueError("dimension %s is not declared as splittable" % d)
     lowering.set_tensor_lowering(
         self.outputs[0],
@@ -2528,7 +2528,7 @@ def top_1(x, reduced_dim, dtype=tf.int32, name=None):
   with tf.name_scope(name, default_name="top_1"):
     max_val = reduce_max(x, reduced_dim=reduced_dim)
     is_max = to_float(equal(x, max_val))
-    pos = range(x.mesh, reduced_dim, tf.float32)
+    pos = xrange(x.mesh, reduced_dim, tf.float32)
     ret = reduce_max(is_max * pos, reduced_dim=reduced_dim)
     ret = cast(ret, dtype)
     return ret, max_val
@@ -2841,8 +2841,8 @@ def _einsum_equation(input_shapes, output_shape):
         dim_to_letter[d] = chr(next_letter)
         next_letter += 1
       ret.append(dim_to_letter[d])
-  ret = "".join(ret)
-  return ret
+
+  return "".join(ret)
 
 
 def is_subsequence(short_seq, long_seq):
diff --git a/tensor2tensor/mesh_tensorflow/mnist.py b/tensor2tensor/mesh_tensorflow/mnist.py
index c0460506b..94449c5e9 100644
--- a/tensor2tensor/mesh_tensorflow/mnist.py
+++ b/tensor2tensor/mesh_tensorflow/mnist.py
@@ -185,11 +185,11 @@ def train_input_fn():
     # randomness, while smaller sizes use less memory. MNIST is a small
     # enough dataset that we can easily shuffle the full epoch.
     ds = dataset.train(FLAGS.data_dir)
-    ds = ds.cache().shuffle(buffer_size=50000).batch(FLAGS.batch_size)
+    ds_batched = ds.cache().shuffle(buffer_size=50000).batch(FLAGS.batch_size)
 
     # Iterate through the dataset a set number (`epochs_between_evals`) of times
     # during each training session.
-    ds = ds.repeat(FLAGS.epochs_between_evals)
+    ds = ds_batched.repeat(FLAGS.epochs_between_evals)
     return ds
 
   def eval_input_fn():
@@ -197,7 +197,7 @@ def eval_input_fn():
         FLAGS.batch_size).make_one_shot_iterator().get_next()
 
   # Train and evaluate model.
-  for _ in range(FLAGS.train_epochs // FLAGS.epochs_between_evals):
+  for _ in xrange(FLAGS.train_epochs // FLAGS.epochs_between_evals):
     mnist_classifier.train(input_fn=train_input_fn, hooks=None)
     eval_results = mnist_classifier.evaluate(input_fn=eval_input_fn)
     print("\nEvaluation results:\n\t%s\n" % eval_results)
diff --git a/tensor2tensor/mesh_tensorflow/mtf_model.py b/tensor2tensor/mesh_tensorflow/mtf_model.py
index 8bd15fa39..e39aa6ee7 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_model.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_model.py
@@ -52,6 +52,7 @@ def estimator_model_fn(cls,
                          decode_hparams=None,
                          use_tpu=False,
                          xla_compile=False):
+    del xla_compile
     hparams = copy.deepcopy(hparams)
     hparams.use_tpu = use_tpu
     # merge decode_hparams into hparams if present
diff --git a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
index 9797c439e..cfcd893d2 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
@@ -91,14 +91,11 @@ def __call__(self, params):
     batch_size = params['batch_size']
     logging.info('call ToyModelInput() with batch size {}'.format(batch_size))
 
-    dataset = Dataset.from_tensor_slices((self._images, self._labels))
-    dataset = dataset.repeat()
+    ds = Dataset.from_tensor_slices((self._images, self._labels)).repeat()
 
-    dataset = dataset.prefetch(batch_size)
-    dataset = dataset.apply(
-        tf.contrib.data.batch_and_drop_remainder(batch_size))
+    dataset = ds.apply(
+        tf.contrib.data.batch_and_drop_remainder(batch_size)).prefetch(2)
 
-    dataset = dataset.prefetch(2)  # Prefetch overlaps in-feed with training
     return dataset
 
 
diff --git a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
index 9302f916f..2466a8011 100644
--- a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
@@ -18,7 +18,7 @@
 from __future__ import print_function
 
 import functools
-from past.builtins import xrange
+from six.moves import xrange  # pylint: disable=redefined-builtin
 
 from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
 import tensorflow as tf
@@ -251,7 +251,7 @@ def my_fn(pnum):
 
   def laid_out_pnum(self):
     """Returns a LaidOutTensor containing the processor number."""
-    return self.LaidOutTensor(list(range(self.size)))
+    return self.LaidOutTensor(list(xrange(self.size)))
 
   @property
   def devices(self):

From 7986075bfe273c96eaf6c6e2526b89e69fed7b99 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 9 Aug 2018 18:44:23 -0700
Subject: [PATCH 0523/2720] fixing visualization bug on resolutions other than
 64x64

PiperOrigin-RevId: 208146799
---
 tensor2tensor/models/research/next_frame_sv2p.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index b72a62f32..2a0905138 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -557,7 +557,8 @@ def get_extra_loss(self, latent_means=None, latent_stds=None,
 
   def body(self, features):
     hparams = self.hparams
-    batch_size = common_layers.shape_list(features["inputs"])[0]
+    input_shape = common_layers.shape_list(features["inputs"])
+    batch_size, _, _, frame_height, frame_channels = input_shape
 
     # Swap time and batch axes.
     input_frames = common_video.swap_time_and_batch_axes(features["inputs"])
@@ -606,7 +607,7 @@ def body(self, features):
     # TODO(mbz): clean this up!
     def fix_video_dims_and_concat_on_x_axis(x):
       x = tf.transpose(x, [1, 3, 4, 0, 2])
-      x = tf.reshape(x, [batch_size, 64, 3, -1])
+      x = tf.reshape(x, [batch_size, frame_height, frame_channels, -1])
       x = tf.transpose(x, [0, 3, 1, 2])
       return x
 

From 106ebb83101ece69b3cc74072729c4568c68b6ed Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 9 Aug 2018 21:54:25 -0700
Subject: [PATCH 0524/2720] Fix bugs.

PiperOrigin-RevId: 208161193
---
 tensor2tensor/mesh_tensorflow/mesh_tensorflow.py     | 2 +-
 tensor2tensor/mesh_tensorflow/mnist.py               | 2 +-
 tensor2tensor/mesh_tensorflow/placement_mesh_impl.py | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index 7acb83b09..36ecc8665 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -2528,7 +2528,7 @@ def top_1(x, reduced_dim, dtype=tf.int32, name=None):
   with tf.name_scope(name, default_name="top_1"):
     max_val = reduce_max(x, reduced_dim=reduced_dim)
     is_max = to_float(equal(x, max_val))
-    pos = xrange(x.mesh, reduced_dim, tf.float32)
+    pos = range(x.mesh, reduced_dim, tf.float32)
     ret = reduce_max(is_max * pos, reduced_dim=reduced_dim)
     ret = cast(ret, dtype)
     return ret, max_val
diff --git a/tensor2tensor/mesh_tensorflow/mnist.py b/tensor2tensor/mesh_tensorflow/mnist.py
index 94449c5e9..2bf870160 100644
--- a/tensor2tensor/mesh_tensorflow/mnist.py
+++ b/tensor2tensor/mesh_tensorflow/mnist.py
@@ -197,7 +197,7 @@ def eval_input_fn():
         FLAGS.batch_size).make_one_shot_iterator().get_next()
 
   # Train and evaluate model.
-  for _ in xrange(FLAGS.train_epochs // FLAGS.epochs_between_evals):
+  for _ in range(FLAGS.train_epochs // FLAGS.epochs_between_evals):
     mnist_classifier.train(input_fn=train_input_fn, hooks=None)
     eval_results = mnist_classifier.evaluate(input_fn=eval_input_fn)
     print("\nEvaluation results:\n\t%s\n" % eval_results)
diff --git a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
index 2466a8011..915cb9bec 100644
--- a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
@@ -251,7 +251,7 @@ def my_fn(pnum):
 
   def laid_out_pnum(self):
     """Returns a LaidOutTensor containing the processor number."""
-    return self.LaidOutTensor(list(xrange(self.size)))
+    return self.LaidOutTensor(list(range(self.size)))
 
   @property
   def devices(self):

From 85979c327997c16bdc90ed968d2a9da34c7044e0 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 9 Aug 2018 23:47:51 -0700
Subject: [PATCH 0525/2720] Improve docs & test
 TensorLayout,Graph,Lowering,Mesh{,Impl}.

PiperOrigin-RevId: 208168342
---
 .../mesh_tensorflow/mesh_tensorflow.py        | 293 +++++++++++-------
 .../mesh_tensorflow/mesh_tensorflow_test.py   |  96 +++++-
 2 files changed, 271 insertions(+), 118 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index 36ecc8665..7c31ddf7a 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -55,7 +55,15 @@ def convert_to_dimension(d):
 
 
 class Shape(object):
-  """Shape of a Tensor or Mesh."""
+  """Shape of a Tensor or Mesh.
+
+  #### Examples
+
+  ```python
+  # Create shape [4, 8] with names "x" and "y" respectively.
+  shape = mtf.Shape([mtf.Dimension("x", 4), mtf.Dimension("y", 8)])
+  ```
+  """
 
   def __init__(self, dims):
     """Constructs a shape for a Tensor or Mesh.
@@ -182,7 +190,15 @@ def convert_to_shape(x):
 
 
 class LayoutRules(object):
-  """Represents layout of a computation."""
+  """Represents layout of a computation.
+
+  #### Examples
+
+  ```python
+  # Map "d_ff" and "heads" Tensor Dimensions to the "model" Mesh Dimension.
+  layout_rules = mtf.LayoutRules([("d_ff", "model"), ("heads", "model")])
+  ```
+  """
 
   def __init__(self, pairs):
     """Constructs a layout.
@@ -218,7 +234,7 @@ def tensor_dimension_to_mesh_axis(self, tensor_dimension, mesh_shape):
     return val[0] if val else None
 
   def tensor_layout(self, tensor_shape, mesh_shape):
-    """Computes TensorLayout given a tensor shape and a mesh shape.
+    """Computes TensorLayout given a Tensor Shape and a Mesh Shape.
 
     Args:
       tensor_shape: Shape.
@@ -228,14 +244,14 @@ def tensor_layout(self, tensor_shape, mesh_shape):
       TensorLayout.
 
     Raises:
-      ValueError: If two tensor dimensions map to the same mesh dimension.
+      ValueError: If two Tensor Dimensions map to the same Mesh Dimensions.
     """
     ret = [self.tensor_dimension_to_mesh_axis(d, mesh_shape)
            for d in tensor_shape]
     not_nones = [a for a in ret if a is not None]
     if len(not_nones) != len(set(not_nones)):
       raise ValueError(
-          "Two tensor dimensions may not map to the same mesh dimesnsion:"
+          "Two Tensor Dimensions may not map to the same Mesh Dimension:"
           " layout=%s tensor_shape=%s mesh_shape=%s " %
           (self, tensor_shape, mesh_shape))
     return TensorLayout(ret)
@@ -258,19 +274,26 @@ def convert_to_layout_rules(x):
 
 
 class TensorLayout(object):
-  """Mapping from tensor dimension to mesh dimension.
+  """Injective partial map between Tensor axes and Mesh axes.
+
+  TensorLayout is a tuple of optional integers with length tensor.ndims. Each
+  item is either a unique integer indicating the mesh axis over which that
+  tensor dimension is split or None, indicating that this tensor dimension is
+  not split.
 
-  Represented as a tuple of optional integers with length tensor.ndims.
-  Each item is either a unique integer inicating the mesh axis over
-  which that tensor dimension is split, or None, indicating that this
-  tensor dimension is not split.
+  #### Examples
+
+  ```python
+  # Split first and last Tensor dimensions according to mesh axes 0 and 1.
+  tensor_layout = mtf.TensorLayout([0, None, 1])
+  ```
   """
 
   def __init__(self, tensor_axis_to_mesh_axis):
-    """Create a TensorLayout.
+    """Creates a TensorLayout.
 
     Args:
-      tensor_axis_to_mesh_axis: a sequence of optional integers.
+      tensor_axis_to_mesh_axis: List-like where each element is an int or None.
     """
     self._tensor_axis_to_mesh_axis = tuple(tensor_axis_to_mesh_axis)
 
@@ -294,22 +317,22 @@ def __iter__(self):
 
   @property
   def tensor_axis_to_mesh_axis(self):
-    """Convert to a tuple of optional integers."""
+    """Converts to a tuple of optional integers."""
     return self._tensor_axis_to_mesh_axis
 
   @property
   def is_fully_replicated(self):
-    """Do all tensor dimensions map to None."""
+    """Whether all tensor dimensions map to None."""
     return self.tensor_axis_to_mesh_axis == (None,) * len(self)
 
   def mesh_axis_to_tensor_axis(self, mesh_ndims):
     """For each mesh axis, which Tensor axis maps to it.
 
     Args:
-      mesh_ndims: an integer
+      mesh_ndims: int.
 
     Returns:
-      a tuple of optional integers, with length mesh_ndims
+      Tuple of optional integers, with length mesh_ndims.
     """
     return tuple(
         [self._tensor_axis_to_mesh_axis.index(mesh_axis)
@@ -318,7 +341,7 @@ def mesh_axis_to_tensor_axis(self, mesh_ndims):
 
 
 class Graph(object):
-  """Distributed-TF graph."""
+  """Mesh-TensorFlow graph."""
 
   def __init__(self):
     self._operations = []
@@ -351,18 +374,37 @@ def to_string(self):
 
 
 class Lowering(object):
-  """Lowering of a Graph from mesh-tensorflow to tensorflow."""
+  """Lowering of a Graph from Mesh-TensorFlow to TensorFlow.
 
-  def __init__(self, graph, mesh_to_impl):
-    """Create a Lowering of a graph.
+  #### Examples
+
+  Below we form a Graph with one Tensor and lower it to recover the original
+  tf.Tensor.
 
-    layout is a dictionary whose keys are the meshes in the graph
-    and whose values are themselves dictionaries mapping tensor-dimension
-    names to mesh dimensions (integers).
+  ```python
+  from tensor2tensor.mesh_tensorflow import placement_mesh_impl
+
+  graph = mtf.Graph()
+  mesh = mtf.Mesh(graph, "my_mesh")
+  inputs = tf.constant(0.)
+  mtf_inputs = mtf.import_tf_tensor(mesh,
+                                    inputs=inputs,
+                                    shape=mtf.Shape([]))
+  mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+      shape=[], layout={}, devices=[""])
+  lowering = mtf.Lowering(graph, {mesh: mesh_impl})
+  outputs = lowering.export_to_tf_tensor(mtf_inputs)  # tf.constant(0.)
+  ```
+  """
+
+  def __init__(self, graph, mesh_to_impl):
+    """Creates a Lowering of a Graph.
 
     Args:
-      graph: a Graph
-      mesh_to_impl: {Mesh: MeshImpl}
+      graph: Graph.
+      mesh_to_impl: {Mesh: MeshImpl}. Keys are the Mesh's in the graph and
+        their values are MeshImpl's, which map Tensor Dimension names to
+        Mesh Dimension names.
     """
     # tf.logging.info("LOWERING GRAPH:\n%s" % graph.to_string)
     self.mesh_to_impl = mesh_to_impl   # {Mesh: MeshImpl}
@@ -392,9 +434,10 @@ def export_to_tf_tensor(self, x):
     """Turn a Tensor into a tf.Tensor.
 
     Args:
-      x: a Tensor
+      x: Tensor.
+
     Returns:
-      a tf.Tensor
+      tf.Tensor.
     """
     mesh_impl = self.mesh_impl(x)
     return mesh_impl.export_to_tf_tensor(
@@ -405,11 +448,11 @@ def lowered_operation(self, op):
 
   def copy_masters_to_slices(self):
     return tf.group(
-        [v.copy_master_to_slices for _, v in six.iteritems(self.variables)])
+        [v.copy_master_to_slices for v in six.itervalues(self.variables)])
 
   def copy_slices_to_masters(self):
     return tf.group(
-        [v.copy_slices_to_master for _, v in six.iteritems(self.variables)])
+        [v.copy_slices_to_master for v in six.itervalues(self.variables)])
 
   def add_counter(self, key, value):
     assert isinstance(value, int)
@@ -423,10 +466,10 @@ def laid_out_size(self, tensor):
     """Total size of all slices.
 
     Args:
-      tensor: a Tensor
+      tensor: Tensor.
 
     Returns:
-      an integer
+      int.
     """
     return self.mesh_impl(tensor).laid_out_size(tensor.shape)
 
@@ -447,10 +490,10 @@ def verify_slice_shapes(self, tensor, laid_out_tensor):
 class Mesh(object):
   """A placeholder with no functionality.
 
-  A Graph is built with each tensor assigned to a mesh.  The mesh does not
+  A Graph is built with each Tensor assigned to a Mesh. The Mesh does not
   know its shape or its implementation.
 
-  A Lowering asssigns a MeshImpl to each mesh.
+  A Lowering assigns each Mesh to a MeshImpl.
   """
 
   def __init__(self, graph, name):
@@ -463,20 +506,29 @@ def graph(self):
 
 
 class MeshImpl(object):
-  """Implementation of a mesh.
+  """Implementation of a Mesh.
+
+  Unlike Mesh, MeshImpl carries Shape and LayoutRules. Subclasses of MeshImpl
+  also carry devices.
 
-  Knows its shape, its underlying devices, and its layout
-  (mapping from TensorDim to mesh-dimension).
+  #### Examples
 
-  Subclasses will include PlacementMeshImpl and SimdMeshImpl
+  ```python
+  shape = mtf.Shape([mtf.Dimension("batch", 4),
+                     mtf.Dimension("model", 8)])
+  layout_rules = mtf.LayoutRules([("batch", "batch"),
+                                  ("d_ff", "model"),
+                                  ("heads", "model")])
+  mesh_impl = mtf.MeshImpl(shape=shape, layout_rules=layout_rules)
+  ```
   """
 
   def __init__(self, shape, layout_rules):
-    """Create a mesh.
+    """Creates a mesh implementation.
 
     Args:
-      shape: a list of ints
-      layout_rules: a LayoutRules
+      shape: Shape.
+      layout_rules: LayoutRules.
     """
     self._shape = convert_to_shape(shape)
     self._layout_rules = convert_to_layout_rules(layout_rules)
@@ -497,13 +549,18 @@ def layout_rules(self):
   def size(self):
     return self.shape.size
 
+  @property
+  def supports_control_dependencies(self):
+    return True
+
   def tensor_dimension_to_mesh_axis(self, tensor_dimension):
     """Mesh axis associated with tensor dimension (or None).
 
     Args:
-      tensor_dimension: a Dimension
+      tensor_dimension: Dimension.
+
     Returns:
-      an integer or None
+      int or None.
     """
     return self.layout_rules.tensor_dimension_to_mesh_axis(
         tensor_dimension, self.shape)
@@ -512,9 +569,10 @@ def tensor_layout(self, arg):
     """Compute TensorLayout for a Tensor or a Shape.
 
     Args:
-      arg: a Tensor or Shape
+      arg: Tensor or Shape.
+
     Returns:
-      a TensorLayout
+      TensorLayout.
     """
     if isinstance(arg, Tensor):
       arg = arg.shape
@@ -524,9 +582,10 @@ def mesh_axis_to_cumprod(self, tensor_shape):
     """For each mesh axis, give the product of previous tensor axes.
 
     Args:
-      tensor_shape: a Shape
+      tensor_shape: Shape.
+
     Returns:
-      a list with length self.ndims where each element is an integer or None.
+      list with length self.ndims where each element is an integer or None.
     """
     tensor_layout = self.tensor_layout(tensor_shape)
     ma2ta = tensor_layout.mesh_axis_to_tensor_axis(self.ndims)
@@ -534,15 +593,17 @@ def mesh_axis_to_cumprod(self, tensor_shape):
     return [None if ta is None else ta2cumprod[ta] for ta in ma2ta]
 
   def slice_shape(self, tensor_shape):
-    """Shape of each slice of the tensor.
+    """Shape of each slice of the Tensor.
 
     Args:
-      tensor_shape: a Shape
+      tensor_shape: Shape.
+
     Returns:
-      a list of integers with length tensor_shape.ndims
+      list of integers with length tensor_shape.ndims.
+
     Raises:
-      ValueError: if a tensor dimension is not divisible by the corresponding
-        mesh dimension.
+      ValueError: If a Tensor dimension is not divisible by the corresponding
+        Mesh dimension.
     """
     tensor_layout = self.tensor_layout(tensor_shape)
     ret = []
@@ -564,10 +625,11 @@ def slice_begin(self, tensor_shape, pnum):
     """Begin position for the tensor slice for the given processor.
 
     Args:
-      tensor_shape: a Shape
-      pnum: an integer <= self.size
+      tensor_shape: Shape.
+      pnum: int <= self.size.
+
     Returns:
-      a list of integers with length tensor_shape.ndims
+      list of integers with length tensor_shape.ndims.
     """
     tensor_layout = self.tensor_layout(tensor_shape)
     coordinates = pnum_to_processor_coordinates(self.shape, pnum)
@@ -585,35 +647,37 @@ def laid_out_size(self, tensor_shape):
     """Total size of all slices.
 
     Args:
-      tensor_shape: a Shape
+      tensor_shape: Shape.
 
     Returns:
-      an integer
+      int.
     """
     return list_product(self.slice_shape(tensor_shape)) * self.size
 
   def slicewise(self, fn, *inputs):
-    """Execute a function in parallel on all slices.
+    """Executes a function in parallel on all slices.
 
     Args:
-      fn: a function from tf.Tensors to tf.Tensor or a tuple of tf.Tensors.
-      *inputs: a list of inputs.  Each input is either a LaidOutTensor or
+      fn: function from tf.Tensors to tf.Tensor or a tuple of tf.Tensors.
+      *inputs: list of inputs.  Each input is either a LaidOutTensor or
         is convertible to a tf.Tensor.
+
     Returns:
-      a LaidOutTensor, or a tuple of LaidOutTensors if fn returns a tuple.
+      LaidOutTensor, or a tuple of LaidOutTensors if fn returns a tuple.
     """
     raise NotImplementedError("Slicewise not implemented")
 
   def Print(self, x, data, message, **kwargs):  # pylint: disable=invalid-name
-    """call tf.Print.
+    """Calls tf.Print.
 
     Args:
-      x: a LaidOutTensor
-      data: a list of LaidOutTensor
-      message: a string
-      **kwargs: keyword arguments to tf.print
+      x: LaidOutTensor.
+      data: list of LaidOutTensor.
+      message: str.
+      **kwargs: keyword arguments to tf.print.
+
     Returns:
-      a LaidOutTensor
+      LaidOutTensor.
     """
     del data, message, kwargs
     tf.logging.warning("Warning - mtf.Print not implemented for this mesh type")
@@ -623,11 +687,12 @@ def allreduce(self, x, mesh_axes, reduction_fn_string):
     """Grouped allreduce, (summed across the given dimensions).
 
     Args:
-      x: a LaidOutTensor
-      mesh_axes: a list of integers - the mesh dimensions to be reduced
-      reduction_fn_string: "SUM" or "MAX"
+      x: LaidOutTensor.
+      mesh_axes: list of integers, the mesh dimensions to be reduced.
+      reduction_fn_string: "SUM" or "MAX".
+
     Returns:
-      a LaidOutTensor
+      LaidOutTensor.
     """
     raise NotImplementedError("Allreduce not implemented")
 
@@ -635,14 +700,15 @@ def allsplit(self, x, mesh_axis, split_axis):
     """Inverse of allconcat - split each slice and keep only one piece of it.
 
     The number of ways to split is the number of processors in the group.
-    The part that is kept corrseponds to the processor's index in the group.
+    The part that is kept corresponds to the processor's index in the group.
 
     Args:
-      x: a LaidOutTensor
-      mesh_axis: an integer the mesh axis along which to split
-      split_axis: an integer (the Tensor axis along which to split)
+      x: LaidOutTensor.
+      mesh_axis: int, the mesh axis along which to split.
+      split_axis: int, the Tensor axis along which to split.
+
     Returns:
-      a LaidOutTensor
+      LaidOutTensor.
     """
     num_splits = self.shape[mesh_axis].size
     def my_fn(x, coordinate):
@@ -659,11 +725,12 @@ def allconcat(self, x, mesh_axis, concat_axis):
     """Grouped allconcat (like MPI allgather followed by concat).
 
     Args:
-      x: a LaidOutTensor
-      mesh_axis: an integer - the mesh axis along which to group
-      concat_axis: an integer (the Tensor axis along which to concatenate)
+      x: LaidOutTensor.
+      mesh_axis: int, the mesh axis along which to group.
+      concat_axis: int, the Tensor axis along which to concatenate.
+
     Returns:
-      a LaidOutTensor
+      LaidOutTensor.
     """
     raise NotImplementedError("Allconcat not implemented")
 
@@ -671,12 +738,13 @@ def alltoall(self, x, mesh_axis, split_axis, concat_axis):
     """Grouped alltoall (like MPI alltoall with splitting and concatenation).
 
     Args:
-      x: a LaidOutTensor
-      mesh_axis: an integer the mesh axis along which to group
-      split_axis: an integer (the Tensor axis along which to split)
-      concat_axis: an integer (the Tensor axis along which to concatenate)
+      x: LaidOutTensor.
+      mesh_axis: int, the mesh axis along which to group.
+      split_axis: int, the Tensor axis along which to split.
+      concat_axis: int, the Tensor axis along which to concatenate.
+
     Returns:
-      a LaidOutTensor
+      LaidOutTensor.
     """
     raise NotImplementedError("Alltoall not implemented")
 
@@ -684,7 +752,7 @@ def laid_out_pnum(self):
     """Returns a LaidOutTensor containing the processor number.
 
     Returns:
-      a LaidOutTensor where each slice is an integer scalar
+      LaidOutTensor where each slice is an integer scalar.
     """
     raise NotImplementedError("laid_out_pnum not implemented")
 
@@ -692,9 +760,10 @@ def laid_out_pcoord(self, mesh_axis):
     """Returns a LaidOutTensor containing the processor coordinate.
 
     Args:
-      mesh_axis: an integer
+      mesh_axis: int.
+
     Returns:
-      a LaidOutTensor where each slice is an integer scalar
+      LaidOutTensor where each slice is an integer scalar.
     """
     divisor = list_product(self.shape.to_integer_list[mesh_axis + 1:])
     modulus = self.shape[mesh_axis].size
@@ -706,11 +775,12 @@ def broadcast_impl(self, old_slices, old_shape, new_shape):
     """Implementation of a broadcast operation.
 
     Args:
-      old_slices: a LaidOutTensor
-      old_shape: a Shape
-      new_shape: a Shape
+      old_slices: LaidOutTensor.
+      old_shape: Shape.
+      new_shape: Shape.
+
     Returns:
-      a LaidOutTensor
+      LaidOutTensor.
     """
     new_slice_shape = self.slice_shape(new_shape)
     def tf_fn(x):
@@ -719,14 +789,14 @@ def tf_fn(x):
     return self.slicewise(tf_fn, old_slices)
 
   def make_slices(self, tf_tensor, tensor_shape):
-    """Turn a single tf.Tensor into a list of slices, one for each processor.
+    """Turns a single tf.Tensor into a list of slices, one for each processor.
 
     Args:
-      tf_tensor: a tf.Tensor
-      tensor_shape: a Shape
+      tf_tensor: tf.Tensor.
+      tensor_shape: Shape.
 
     Returns:
-      a list of tf.tensor with length self.size
+      list of tf.tensor with length self.size.
     """
     tensor_layout = self.tensor_layout(tensor_shape)
     slice_shape = self.slice_shape(tensor_shape)
@@ -741,16 +811,15 @@ def my_fn(pnum):
                     list(xrange(self.size)))
 
   def combine_slices(self, slices, tensor_shape, device=None):
-    """Turn a set of slices into a single tensor.
+    """Turns a set of slices into a single tensor.
 
     Args:
-      slices: a list of tf.Tensor with length self.size
-      tensor_shape: a Shape
-      device: an optional device string.
-        if absent, we use the devices of the slices.
+      slices: list of tf.Tensor with length self.size.
+      tensor_shape: Shape.
+      device: optional str. If absent, we use the devices of the slices.
 
     Returns:
-      a tf.Tensor
+      tf.Tensor.
     """
     if tensor_shape.ndims == 0:
       return slices[0]
@@ -777,31 +846,29 @@ def combine_slices(self, slices, tensor_shape, device=None):
     return ret[0]
 
   def export_to_tf_tensor(self, x, laid_out_x):
-    """Turn a Tensor into a tf.Tensor.
+    """Turns a Tensor into a tf.Tensor.
 
     Args:
-      x: a Tensor
-      laid_out_x: a LaidOutTensor
+      x: Tensor.
+      laid_out_x: LaidOutTensor.
+
     Returns:
-      a tf.Tensor
+      tf.Tensor.
     """
     raise NotImplementedError("export_to_tf_tensor not implemented")
 
   def import_tf_tensor(self, x, tf_x):
-    """Import a tf.Tensor, producing a LaidOutTensor.
+    """Imports a tf.Tensor, producing a LaidOutTensor.
 
     Args:
-      x: a Tensor
-      tf_x: a tf.Tensor
+      x: Tensor.
+      tf_x: tf.Tensor.
+
     Returns:
-      a LaidOutTensor
+      LaidOutTensor.
     """
     raise NotImplementedError("Import not implemented")
 
-  @property
-  def supports_control_dependencies(self):
-    return True
-
 
 class LazyAllreduceSum(object):
   """Represents a LaidOutTensor with a lazy allreduce.
diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
index 2564bb11a..e73d7e4ec 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
@@ -21,6 +21,7 @@
 from absl.testing import parameterized
 
 from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+from tensor2tensor.mesh_tensorflow import placement_mesh_impl
 
 import tensorflow as tf
 
@@ -28,7 +29,7 @@
 class MeshTensorFlowTest(parameterized.TestCase, tf.test.TestCase):
 
   @parameterized.parameters(
-      (mtf.Dimension(name="x", size=5),),
+      (mtf.Dimension("x", 5),),
       (("x", 5),),
   )
   def testConvertToDimension(self, inputs):
@@ -43,8 +44,8 @@ def testConvertToDimensionGenericInputs(self):
       mtf.convert_to_dimension(5)
 
   @parameterized.parameters(
-      (mtf.Shape([mtf.Dimension(name="x", size=4),
-                  mtf.Dimension(name="y", size=8)]),),
+      (mtf.Shape([mtf.Dimension("x", 4),
+                  mtf.Dimension("y", 8)]),),
       ("x:4;y:8",),
       ("x:4.y:8",),
       ("x:4 y:8",),
@@ -52,10 +53,12 @@ def testConvertToDimensionGenericInputs(self):
   )
   def testConvertToShape(self, inputs):
     shape = mtf.convert_to_shape(inputs)
-    self.assertEqual(shape, mtf.Shape([mtf.Dimension(name="x", size=4),
-                                       mtf.Dimension(name="y", size=8)]))
+    self.assertEqual(shape, mtf.Shape([mtf.Dimension("x", 4),
+                                       mtf.Dimension("y", 8)]))
 
   def testConvertToShapeGenericInputs(self):
+    shape = mtf.convert_to_shape([])
+    self.assertEqual(shape.dims, [])
     shape = mtf.convert_to_shape(None)
     self.assertEqual(shape, None)
     with self.assertRaises(ValueError):
@@ -79,5 +82,88 @@ def testConvertToLayoutRulesGenericInputs(self):
     with self.assertRaises(ValueError):
       mtf.convert_to_layout_rules("d_ff;heads")
 
+  def testTensorLayout(self):
+    tensor_layout = mtf.TensorLayout([0, 2, 1])
+    self.assertEqual(tensor_layout.mesh_axis_to_tensor_axis(0), ())
+    self.assertEqual(tensor_layout.mesh_axis_to_tensor_axis(1), (0,))
+    self.assertEqual(tensor_layout.mesh_axis_to_tensor_axis(2), (0, 2))
+    tensor_layout = mtf.TensorLayout([None, 0])
+    self.assertFalse(tensor_layout.is_fully_replicated)
+    tensor_layout = mtf.TensorLayout([None, None, None])
+    self.assertTrue(tensor_layout.is_fully_replicated)
+
+  def testGraph(self):
+    graph = mtf.Graph()
+    self.assertLen(graph.operations, 0)
+    self.assertLen(graph.tensors, 0)
+    self.assertLen(graph.trainable_variables, 0)
+    self.assertLen(graph.all_variables, 0)
+    mesh = mtf.Mesh(graph, "mesh_test")
+    _ = mtf.import_tf_tensor(mesh,
+                             tf_tensor=tf.constant(0.),
+                             shape=mtf.Shape([]))
+    self.assertLen(graph.operations, 1)
+    self.assertLen(graph.tensors, 1)
+    self.assertLen(graph.trainable_variables, 0)
+    self.assertLen(graph.all_variables, 0)
+    _ = mtf.get_variable(mesh, "variable_0", mtf.Shape([]), trainable=True)
+    self.assertLen(graph.operations, 2)
+    self.assertLen(graph.tensors, 2)
+    self.assertLen(graph.trainable_variables, 1)
+    self.assertLen(graph.all_variables, 1)
+    _ = mtf.get_variable(mesh, "variable_1", mtf.Shape([]), trainable=False)
+    self.assertLen(graph.operations, 3)
+    self.assertLen(graph.tensors, 3)
+    self.assertLen(graph.trainable_variables, 1)
+    self.assertLen(graph.all_variables, 2)
+
+  def testLowering(self):
+    graph = mtf.Graph()
+    mesh = mtf.Mesh(graph, "my_mesh")
+    inputs = tf.constant(0.)
+    mtf_inputs = mtf.import_tf_tensor(mesh,
+                                      tf_tensor=inputs,
+                                      shape=mtf.Shape([]))
+    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+        shape=[], layout={}, devices=[""])
+    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
+
+    outputs = lowering.export_to_tf_tensor(mtf_inputs)
+    with self.test_session() as sess:
+      inputs_value, outputs_value = sess.run([inputs, outputs])
+    self.assertEqual(inputs_value, outputs_value)
+
+    # Check that methods run without error.
+    _ = lowering.copy_masters_to_slices()
+    _ = lowering.copy_slices_to_masters()
+
+  def testMesh(self):
+    graph = mtf.Graph()
+    mesh = mtf.Mesh(graph, "my_mesh")
+    self.assertEqual(mesh.graph, graph)
+
+  def testMeshImpl(self):
+    shape = mtf.Shape([mtf.Dimension("batch", 4),
+                       mtf.Dimension("model", 8)])
+    layout_rules = mtf.LayoutRules([("batch", "batch"),
+                                    ("d_ff", "model"),
+                                    ("heads", "model")])
+    mesh_impl = mtf.MeshImpl(shape=shape, layout_rules=layout_rules)
+    self.assertEqual(mesh_impl.shape, shape)
+    self.assertEqual(mesh_impl.ndims, len(shape))
+    self.assertEqual(mesh_impl.layout_rules, layout_rules)
+    self.assertEqual(mesh_impl.size, shape.size)
+    self.assertTrue(mesh_impl.supports_control_dependencies)
+
+    batch = mtf.Dimension("batch", 128)
+    length = mtf.Dimension("length", 500)
+    d_ff = mtf.Dimension("d_ff", 2048)
+    heads = mtf.Dimension("heads", 8)
+    self.assertEqual(mesh_impl.tensor_dimension_to_mesh_axis(batch), 0)
+    self.assertEqual(mesh_impl.tensor_dimension_to_mesh_axis(d_ff), 1)
+    self.assertEqual(mesh_impl.tensor_dimension_to_mesh_axis(heads), 1)
+    self.assertEqual(mesh_impl.tensor_layout(mtf.Shape([batch, length, d_ff])),
+                     mtf.TensorLayout([0, None, 1]))
+
 if __name__ == "__main__":
   tf.test.main()

From bbdc1cd88d0eb8b0d0fb21b83b1e7a075b058323 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Fri, 10 Aug 2018 00:29:09 -0700
Subject: [PATCH 0526/2720] Updated baseline and mixture-of-experts language
 model configs.

PiperOrigin-RevId: 208171247
---
 .../mesh_tensorflow/mtf_transformer.py        | 36 +++++++++++++------
 1 file changed, 26 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
index 7c4674a52..20bb91f06 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
@@ -866,17 +866,33 @@ def mtf_transformer_paper_tr_0_mesh_512():
 
 
 @registry.register_hparams
-def mtf_transformer_lm_moe():
-  """Mixture of experts language model."""
-  hparams = mtf_transformer_base()
-  hparams.label_smoothing = 0.0
+def mtf_transformer_lm_baseline():
+  """Small language model to run on 1 TPU.
+
+  Run this on 2x2 on languagemodel_lm1b32k_packed for 272000 steps (10 epochs)
+  140M params
+
+  Returns:
+    a hparams
+  """
+  hparams = mtf_transformer_paper_lm(-1)
   hparams.batch_size = 128
-  hparams.d_model = 1024
-  hparams.d_ff = 4096
-  hparams.attention_key_channels = 1024
-  hparams.attention_value_channels = 1024
-  hparams.shared_embedding_and_softmax_weights = False
-  hparams.num_heads = 8
+  hparams.learning_rate_decay_steps = 27200  # one epoch on lm1b
+  hparams.mesh_shape = "batch:8"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_lm_moe():
+  """Mixture of experts language model.
+
+  Run this on 2x2 on languagemodel_lm1b32k_packed for 272000 steps (10 epochs)
+  900M params.
+
+  Returns:
+    a hparams
+  """
+  hparams = mtf_transformer_lm_baseline()
   hparams.mesh_shape = "all:8"
   hparams.layout = "batch:all;experts:all"
   hparams.feedforward_layer = "moe"

From 9221fe8561ac8059144bbb82eab68976423ce4d4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 10 Aug 2018 09:39:36 -0700
Subject: [PATCH 0527/2720] Rename common_layers.is_on_tpu to is_xla_compiled

PiperOrigin-RevId: 208221937
---
 tensor2tensor/layers/common_layers.py         | 25 +++++++++++++------
 tensor2tensor/layers/modalities.py            |  2 +-
 tensor2tensor/models/image_transformer_2d.py  |  2 +-
 .../research/universal_transformer_util.py    |  2 +-
 tensor2tensor/models/transformer.py           |  4 +--
 tensor2tensor/utils/t2t_model.py              |  6 ++---
 6 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 1e2fe001a..9f332ab48 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -61,7 +61,18 @@ def convert_gradient_to_tensor(x):
   return x
 
 
-def is_on_tpu():
+def is_xla_compiled():
+  """Whether we are building graph that will be compiled by XLA.
+
+  This checks whether the code is executing within an XLA context.
+
+  If True, model authors should ensure the graph they build is compilable by
+  XLA. Specifically, they should ensure that all ops have XLA implementations
+  and that all shapes are statically known.
+
+  Returns:
+    bool, whether the current graph will be compiled for XLA.
+  """
   ctxt = tf.get_default_graph()._get_control_flow_context()  # pylint: disable=protected-access
   return control_flow_util.GetContainingXLAContext(ctxt) is not None
 
@@ -256,7 +267,7 @@ def flatten4d3d(x):
 # TODO(noam): remove this function after TPUs do gather faster.
 def gather(params, indices, dtype=tf.float32):
   """Version of tf.gather that works faster on tpu."""
-  if not is_on_tpu():
+  if not is_xla_compiled():
     return tf.gather(params, indices)
   vocab_size = params.get_shape().as_list()[0]
   indices_flat = tf.reshape(indices, [-1])
@@ -280,7 +291,7 @@ def cumsum(x, axis=0, exclusive=False):
   Returns:
     Tensor of the same shape as x.
   """
-  if not is_on_tpu():
+  if not is_xla_compiled():
     return tf.cumsum(x, axis=axis, exclusive=exclusive)
   x_shape = shape_list(x)
   rank = len(x_shape)
@@ -2497,7 +2508,7 @@ def sru(x,
   """
   if num_layers < 1:
     raise ValueError("Number of layers must be positive: %d" % num_layers)
-  if is_on_tpu():  # On TPU the XLA does a good job with while.
+  if is_xla_compiled():  # On TPU the XLA does a good job with while.
     return sru_with_scan(x, num_layers, activation, initial_state, name, reuse)
   try:
     from tensorflow.contrib.recurrent.python.ops import functional_rnn  # pylint: disable=g-import-not-at-top
@@ -3322,7 +3333,7 @@ def get_res():
       return get_res()
 
     # Prevent sampling after steps is passed to speed it up.
-    if is_on_tpu():
+    if is_xla_compiled():
       return get_res()
     else:
       return tf.cond(
@@ -3582,7 +3593,7 @@ def get_sorted_projections(x):
         proj = tf.tanh(x)
       proj = tf.transpose(proj, [1, 0])  # [num_vecs, batch] after this.
 
-      if is_on_tpu():
+      if is_xla_compiled():
         proj_dtype = proj.dtype
         proj = tf.cast(proj, tf.bfloat16)
 
@@ -3611,7 +3622,7 @@ def upscale(inputs, f, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR):
 
 
 def tpu_safe_image_summary(image):
-  if is_on_tpu():
+  if is_xla_compiled():
     # We only support float32 images at the moment due to casting complications.
     if image.dtype != tf.float32:
       image = tf.to_float(image)
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index d404495ab..f2a9d3b21 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -158,7 +158,7 @@ def top(self, body_output, _):
       else:
         body_output = tf.reshape(body_output, [-1, body_output_shape[-1]])
         logits = tf.matmul(body_output, var, transpose_b=True)
-        if (common_layers.is_on_tpu() and
+        if (common_layers.is_xla_compiled() and
             self._model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
           # TPU does not react kindly to extra dimensions.
           # TODO(noam): remove this once TPU is more forgiving of extra dims.
diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index a6cde618b..1ef6ede1a 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -104,7 +104,7 @@ class Img2imgTransformerBlockParallel(t2t_model.T2TModel):
 
   def body(self, features):
     assert self._hparams.block_size > 0
-    assert not common_layers.is_on_tpu()
+    assert not common_layers.is_xla_compiled()
 
     hparams = copy.copy(self._hparams)
     targets = features["targets"]
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index fb358f2e8..c76490a33 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -104,7 +104,7 @@ def universal_transformer_encoder(encoder_input,
           encoder_self_attention_bias)
       nonpadding = 1.0 - padding
     pad_remover = None
-    if hparams.use_pad_remover and not common_layers.is_on_tpu():
+    if hparams.use_pad_remover and not common_layers.is_xla_compiled():
       pad_remover = expert_utils.PadRemover(padding)
 
     ffn_unit = functools.partial(
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 2a7e1a01d..9af54dc90 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -135,7 +135,7 @@ def decode(self,
         save_weights_to=self.attention_weights,
         losses=losses)
 
-    if (common_layers.is_on_tpu() and
+    if (common_layers.is_xla_compiled() and
         hparams.mode == tf.estimator.ModeKeys.TRAIN):
       # TPU does not react kindly to extra dimensions.
       # TODO(noam): remove this once TPU is more forgiving of extra dims.
@@ -1208,7 +1208,7 @@ def transformer_encoder(encoder_input,
           encoder_self_attention_bias)
       nonpadding = 1.0 - padding
     pad_remover = None
-    if hparams.use_pad_remover and not common_layers.is_on_tpu():
+    if hparams.use_pad_remover and not common_layers.is_xla_compiled():
       pad_remover = expert_utils.PadRemover(padding)
     for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):
       with tf.variable_scope("layer_%d" % layer):
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index a83ccface..8f80111a4 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -456,7 +456,7 @@ def optimize(self, loss, num_async_replicas=1):
                num_async_replicas)
     lr /= math.sqrt(float(num_async_replicas))
     train_op = optimize.optimize(
-        loss, lr, self.hparams, use_tpu=common_layers.is_on_tpu())
+        loss, lr, self.hparams, use_tpu=common_layers.is_xla_compiled())
     return train_op
 
   def set_mode(self, mode):
@@ -1279,7 +1279,7 @@ def estimator_spec_train(self, loss, num_async_replicas=1):
               "Cannot find variable in checkpoint, skipping: %s", var_name)
       tf.train.init_from_checkpoint(pretrained_model_dir, variable_map)
 
-    if common_layers.is_on_tpu():
+    if common_layers.is_xla_compiled():
       host_call = _create_host_call(self.hparams.model_dir)
       _remove_summaries()
       return tf.contrib.tpu.TPUEstimatorSpec(
@@ -1300,7 +1300,7 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
       raise NotImplementedError(_no_problem_err("estimator_spec_eval"))
 
     problem = hparams.problem
-    if common_layers.is_on_tpu():
+    if common_layers.is_xla_compiled():
       _remove_summaries()
       if isinstance(logits, dict):
         eval_metrics_fn = _create_tpu_eval_metrics_fn(problem, hparams)

From d8ff7ae9091a903e03d67e1b6195fdf2e36915e4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 10 Aug 2018 10:06:55 -0700
Subject: [PATCH 0528/2720] Fix TPU hang issue when no summaries are generated

PiperOrigin-RevId: 208226269
---
 tensor2tensor/utils/t2t_model.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 8f80111a4..29b4d996a 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1591,6 +1591,11 @@ def _create_host_call(model_dir):
         continue
       # tensor = tf.to_float(tensor)
       summary_kwargs["ImageSummary" + name] = tensor
+  # When no supported summaries are found, don't create host_call. Otherwise,
+  # TPU outfeed queue would enqueue global_step while host_call doesn't dequeue
+  # it, eventually causing hang.
+  if not summary_kwargs:
+    return None
   summary_kwargs["global_step"] = gs_t
 
   def host_call_fn(**kwargs):

From 0b37ac4e89a40fd52dec296b49303b8487303c04 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 10 Aug 2018 10:33:59 -0700
Subject: [PATCH 0529/2720] Pylint fix.

PiperOrigin-RevId: 208230938
---
 tensor2tensor/bin/t2t_datagen.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index ec9fc2162..7375b3141 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -40,8 +40,6 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import usr_dir
 
-import tensorflow as tf
-
 try:
   # pylint: disable=g-import-not-at-top
   from tensor2tensor.data_generators import algorithmic_math
@@ -52,6 +50,10 @@
 except ImportError:
   pass
 
+# Improrting here to prevent pylint from ungrouped-imports warning.
+import tensorflow as tf  # pylint: disable=g-import-not-at-top
+
+
 
 flags = tf.flags
 FLAGS = flags.FLAGS

From 90a4e50c640e45845bf147295d26575855f4be6b Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 10 Aug 2018 10:41:51 -0700
Subject: [PATCH 0530/2720] A few readability changes for code simplification.

PiperOrigin-RevId: 208232446
---
 tensor2tensor/data_generators/text_encoder.py | 44 +++++++------------
 1 file changed, 15 insertions(+), 29 deletions(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 38ca98c84..64dbf39fa 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -89,6 +89,14 @@ def to_unicode_ignore_errors(s):
   return to_unicode(s, ignore_errors=True)
 
 
+def strip_ids(ids, ids_to_strip):
+  """Strip ids_to_strip from the end ids."""
+  ids = list(ids)
+  while ids and ids[-1] in ids_to_strip:
+    ids.pop()
+  return ids
+
+
 class TextEncoder(object):
   """Base class for converting from ints to/from human readable strings."""
 
@@ -209,13 +217,11 @@ class ClassLabelEncoder(TextEncoder):
   def __init__(self, class_labels=None, class_labels_fname=None):
     super(ClassLabelEncoder, self).__init__(num_reserved_ids=0)
 
-    assert class_labels or class_labels_fname
-    assert not (class_labels and class_labels_fname)
-
     if class_labels_fname:
       with tf.gfile.Open(class_labels_fname) as f:
         class_labels = [label.strip() for label in f.readlines()]
 
+    assert class_labels
     self._class_labels = class_labels
 
   def encode(self, s):
@@ -240,24 +246,11 @@ def vocab_size(self):
     return len(self._class_labels)
 
 
-class OneHotClassLabelEncoder(TextEncoder):
+class OneHotClassLabelEncoder(ClassLabelEncoder):
   """One-hot encoder for class labels."""
 
-  def __init__(self, class_labels=None, class_labels_fname=None):
-    super(OneHotClassLabelEncoder, self).__init__()
-    assert class_labels or class_labels_fname
-    assert not (class_labels and class_labels_fname)
-
-    if class_labels_fname:
-      with tf.gfile.Open(class_labels_fname) as f:
-        class_labels = [label.strip() for label in f.readlines()]
-
-    self._class_labels = class_labels
-
   def encode(self, label_str, on_value=1, off_value=0):  # pylint: disable=arguments-differ
-    e = np.zeros(self.vocab_size, dtype=np.int32)
-    if off_value != 0:
-      e.fill(off_value)
+    e = np.fill(self.vocab_size, off_value, dtype=np.int32)
     e[self._class_labels.index(label_str)] = on_value
     return e.tolist()
 
@@ -644,14 +637,14 @@ def _escaped_token_to_subtoken_ids(self, escaped_token):
   @classmethod
   def build_from_generator(cls,
                            generator,
-                           target_vocab_size,
+                           target_size,
                            max_subtoken_length=None,
                            reserved_tokens=None):
     """Builds a SubwordTextEncoder from the generated text.
 
     Args:
       generator: yields text.
-      target_vocab_size: int, approximate vocabulary size to create.
+      target_size: int, approximate vocabulary size to create.
       max_subtoken_length: Maximum length of a subtoken. If this is not set,
         then the runtime and memory use of creating the vocab is quadratic in
         the length of the longest token. If this is set, then it is instead
@@ -661,14 +654,14 @@ def build_from_generator(cls,
         argument is `None`, it will use `RESERVED_TOKENS`.
 
     Returns:
-      SubwordTextEncoder with `vocab_size` approximately `target_vocab_size`.
+      SubwordTextEncoder with `vocab_size` approximately `target_size`.
     """
     token_counts = collections.defaultdict(int)
     for item in generator:
       for tok in tokenizer.encode(native_to_unicode(item)):
         token_counts[tok] += 1
     encoder = cls.build_to_target_size(
-        target_vocab_size, token_counts, 1, 1e3,
+        target_size, token_counts, 1, 1e3,
         max_subtoken_length=max_subtoken_length,
         reserved_tokens=reserved_tokens)
     return encoder
@@ -1057,10 +1050,3 @@ def decode(self, ids, strip_extraneous=False):
     del strip_extraneous
     return " ".join([str(i) for i in ids])
 
-
-def strip_ids(ids, ids_to_strip):
-  """Strip ids_to_strip from the end ids."""
-  ids = list(ids)
-  while ids and ids[-1] in ids_to_strip:
-    ids.pop()
-  return ids

From 3a2e87f1d013d51398cba9db35fa23160d4b2aeb Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Fri, 10 Aug 2018 11:14:55 -0700
Subject: [PATCH 0531/2720] Update discrete autoencoders

PiperOrigin-RevId: 208238718
---
 tensor2tensor/layers/discretization.py        | 196 ++++++++++--------
 tensor2tensor/layers/discretization_test.py   |   5 +-
 tensor2tensor/models/research/autoencoders.py |  55 +++--
 3 files changed, 152 insertions(+), 104 deletions(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 4de85375a..4fe443cf2 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -595,9 +595,8 @@ def discrete_bottleneck(inputs,
   else:
     block_v_size = None
 
-  with tf.variable_scope(name,
-                         default_name="discrete_bottleneck",
-                         reuse=tf.AUTO_REUSE):
+  with tf.variable_scope(
+      name, default_name="discrete_bottleneck", reuse=tf.AUTO_REUSE):
     embed_fn = partial(
         embed,
         hidden_size=hidden_size,
@@ -614,24 +613,23 @@ def discrete_bottleneck(inputs,
     if bottleneck_kind == "dense":
       # Note discrete output is continuous here.
       outputs_discrete = tf.layers.dense(inputs, z_size, name="vcc")
-      outputs_dense = tf.layers.dense(outputs_discrete,
-                                      filter_size,
-                                      name="vch1")
+      outputs_dense = tf.layers.dense(
+          outputs_discrete, filter_size, name="vch1")
       extra_loss = tf.constant(0.0)
       neg_q_entropy = tf.constant(0.0)
     elif bottleneck_kind in ["dvq", "gumbel-softmax-dvq"]:
       if reshape_method == "slice":
-        x_reshaped = slice_hidden(inputs,
-                                  hidden_size=hidden_size,
-                                  num_blocks=num_blocks)
+        x_reshaped = slice_hidden(
+            inputs, hidden_size=hidden_size, num_blocks=num_blocks)
       elif reshape_method == "project":
         if projection_tensors is None:
           raise ValueError(
               "Projection tensors is None for reshape_method project")
-        x_reshaped = project_hidden(inputs,
-                                    projection_tensors=projection_tensors,
-                                    hidden_size=hidden_size,
-                                    num_blocks=num_blocks)
+        x_reshaped = project_hidden(
+            inputs,
+            projection_tensors=projection_tensors,
+            hidden_size=hidden_size,
+            num_blocks=num_blocks)
       else:
         raise ValueError("Unknown reshape_method")
 
@@ -642,19 +640,20 @@ def discrete_bottleneck(inputs,
       extra_loss = 0
       for i in range(num_residuals):
         x_means_hot_res, x_means_res, q_loss_res, e_loss_res, neg_q_entropy = (
-            embedding_lookup(x_reshaped,
-                             means=means[i],
-                             num_blocks=num_blocks,
-                             block_v_size=block_v_size,
-                             bottleneck_kind=bottleneck_kind,
-                             random_top_k=random_top_k,
-                             soft_em=soft_em,
-                             num_samples=num_samples,
-                             temperature_warmup_steps=temperature_warmup_steps,
-                             do_hard_gumbel_softmax=do_hard_gumbel_softmax,
-                             num_flows=num_flows,
-                             approximate_gs_entropy=approximate_gs_entropy,
-                             sum_over_latents=sum_over_latents))
+            embedding_lookup(
+                x_reshaped,
+                means=means[i],
+                num_blocks=num_blocks,
+                block_v_size=block_v_size,
+                bottleneck_kind=bottleneck_kind,
+                random_top_k=random_top_k,
+                soft_em=soft_em,
+                num_samples=num_samples,
+                temperature_warmup_steps=temperature_warmup_steps,
+                do_hard_gumbel_softmax=do_hard_gumbel_softmax,
+                num_flows=num_flows,
+                approximate_gs_entropy=approximate_gs_entropy,
+                sum_over_latents=sum_over_latents))
         # Update the EMA variables.
         if ema:
           tf.logging.info("Using EMA with beta = {}".format(beta))
@@ -667,8 +666,9 @@ def discrete_bottleneck(inputs,
               decay,
               zero_debias=False)
 
-          dw = tf.matmul(tf.transpose(x_means_hot_res, perm=[1, 2, 0]),
-                         tf.transpose(x_res, perm=[1, 0, 2]))
+          dw = tf.matmul(
+              tf.transpose(x_means_hot_res, perm=[1, 2, 0]),
+              tf.transpose(x_res, perm=[1, 0, 2]))
 
           updated_ema_means_res = moving_averages.assign_moving_average(
               ema_means[i], dw, decay, zero_debias=False)
@@ -705,9 +705,8 @@ def discrete_bottleneck(inputs,
       new_shape = shape[:-2]
       new_shape[-1] = z_size
       x_means_bits = tf.reshape(x_means_bits, shape=new_shape)
-      outputs_discrete = bit_to_int(tf.to_int32(x_means_bits),
-                                    num_bits=z_size,
-                                    base=2)
+      outputs_discrete = bit_to_int(
+          tf.to_int32(x_means_bits), num_bits=z_size, base=2)
 
       # Adjust shape of discrete outputs.
       inputs_shape = common_layers.shape_list(inputs)
@@ -732,9 +731,8 @@ def discrete_bottleneck(inputs,
           summary=summary,
           name=name)
       outputs_discrete = tf.argmax(outputs_hot, axis=-1)
-      outputs_dense = tf.layers.dense(outputs_hot,
-                                      hidden_size,
-                                      name="dae_dense")
+      outputs_dense = tf.layers.dense(
+          outputs_hot, hidden_size, name="dae_dense")
       neg_q_entropy = tf.constant(0.0)
     elif bottleneck_kind == "semhash":
       outputs_discrete = tf.layers.dense(inputs, z_size, name="vcc")
@@ -742,9 +740,10 @@ def discrete_bottleneck(inputs,
       if summary:
         tf.summary.histogram("y_clean", tf.reshape(y_clean, [-1]))
       if noise_dev > 0 and mode == tf.estimator.ModeKeys.TRAIN:
-        noise = tf.truncated_normal(common_layers.shape_list(outputs_discrete),
-                                    mean=0.0,
-                                    stddev=noise_dev)
+        noise = tf.truncated_normal(
+            common_layers.shape_list(outputs_discrete),
+            mean=0.0,
+            stddev=noise_dev)
         y = common_layers.saturating_sigmoid(outputs_discrete + noise)
       else:
         y = y_clean
@@ -756,24 +755,17 @@ def discrete_bottleneck(inputs,
       c = tf.where(
           tf.less(tf.random_uniform([common_layers.shape_list(y)[0]]), pd),
           y_discrete, y)
-      outputs_dense_a = tf.layers.dense(c,
-                                        filter_size,
-                                        name="vch1a")
-      outputs_dense_b = tf.layers.dense(1.0 - c,
-                                        filter_size,
-                                        name="vch1b")
+      outputs_dense_a = tf.layers.dense(c, filter_size, name="vch1a")
+      outputs_dense_b = tf.layers.dense(1.0 - c, filter_size, name="vch1b")
       outputs_dense = outputs_dense_a + outputs_dense_b
       dx = tf.to_int32(tf.stop_gradient(d))
       outputs_discrete = bit_to_int(dx, z_size)
       extra_loss = tf.constant(0.0)
       neg_q_entropy = tf.constant(0.0)
     elif bottleneck_kind == "vae":
-      outputs_discrete, extra_loss, _, _ = vae(inputs,
-                                               z_size,
-                                               name="vae")
-      outputs_dense = tf.layers.dense(outputs_discrete,
-                                      filter_size,
-                                      name="vch1")
+      outputs_discrete, extra_loss, _, _ = vae(inputs, z_size, name="vae")
+      outputs_dense = tf.layers.dense(
+          outputs_discrete, filter_size, name="vch1")
       neg_q_entropy = tf.constant(0.0)
     else:
       raise ValueError("Unknown discretization method.")
@@ -828,7 +820,7 @@ def vq_nearest_neighbor(x, means, soft_em=False, num_samples=10):
   x_means_hot_flat = tf.reshape(x_means_hot, [-1, bottleneck_size])
   x_means = tf.matmul(x_means_hot_flat, means)
   e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
-  return x_means_hot, e_loss
+  return x_means_hot, e_loss, dist
 
 
 def vq_discrete_bottleneck(x,
@@ -840,7 +832,7 @@ def vq_discrete_bottleneck(x,
                            num_samples=10):
   """Simple vector quantized discrete bottleneck."""
   bottleneck_size = 2**bottleneck_bits
-  return vq_body(
+  x_means_hot, e_loss, _ = vq_body(
       x,
       bottleneck_size,
       beta=beta,
@@ -848,6 +840,7 @@ def vq_discrete_bottleneck(x,
       epsilon=epsilon,
       soft_em=soft_em,
       num_samples=num_samples)
+  return x_means_hot, e_loss
 
 
 def vq_body(x,
@@ -862,7 +855,7 @@ def vq_body(x,
   hidden_size = x_shape[-1]
   means, ema_means, ema_count = get_vq_codebook(codebook_size, hidden_size)
   x = tf.reshape(x, [-1, hidden_size])
-  x_means_hot, e_loss = vq_nearest_neighbor(
+  x_means_hot, e_loss, distances = vq_nearest_neighbor(
       x, means, soft_em=soft_em, num_samples=num_samples)
 
   # Update the ema variables
@@ -886,30 +879,49 @@ def vq_body(x,
       loss = beta * e_loss
 
   d = tf.reshape(x_means_hot, x_shape[:-1] + [codebook_size])
-  return d, loss
+  return d, loss, distances
 
 
 def vq_loss(x,
-            one_hot_targets,
+            targets,
             codebook_size,
             beta=0.25,
             decay=0.999,
             epsilon=1e-5,
             soft_em=False,
             num_samples=10):
-  """Simple vector quantized discrete bottleneck."""
+  """Compute the loss of large vocab tensors using a VQAE codebook.
+
+  Args:
+    x: Tensor of inputs to be quantized to nearest code
+    targets: Tensor of target indices to target codes
+    codebook_size: Size of quantization codebook
+    beta: scalar float for moving averages
+    decay: scalar float for moving averages
+    epsilon: scalar float for moving averages
+    soft_em: boolean, whether to apply a soft sampling procedure
+    num_samples: if soft_em, number of samples to take
+
+  Returns:
+    discrete_x: one-hot Tensor indicating which codebook element is closest to x
+    x_means: Tensor, on the forward pass: closest codebook element to x, on the
+      backwards pass: soft convex-combination of codebook elements by proximity
+      to x
+    target_means: the codebook elements corresponding to the targets
+    code_loss: loss driving x closer to its nearest codebook element
+    targets_loss: cross-entropy loss driving x closer to code corresponding to
+      target
+  """
   x_shape = common_layers.shape_list(x)
-  target_shape = common_layers.shape_list(one_hot_targets)
+  target_shape = common_layers.shape_list(targets)
   hidden_size = x_shape[-1]
   means, _, _ = get_vq_codebook(codebook_size, hidden_size)
   x = tf.reshape(x, [-1, hidden_size])
-
-  one_hot_targets = tf.reshape(one_hot_targets, [-1, target_shape[-1]])
+  targets = tf.reshape(targets, [-1])
+  one_hot_targets = tf.one_hot(targets, codebook_size)
   target_means = tf.matmul(one_hot_targets, means)
-  targets_loss = tf.reduce_sum((x - target_means)**2, -1)
-  targets_loss = tf.reduce_mean(targets_loss)
 
-  discrete_x, code_loss = vq_body(
+  discrete_x, code_loss, distances = vq_body(
       x,
       codebook_size,
       beta=beta,
@@ -918,9 +930,17 @@ def vq_loss(x,
       soft_em=soft_em,
       num_samples=num_samples)
 
+  logits = -distances
+  targets_loss = tf.losses.sparse_softmax_cross_entropy(
+      logits=logits, labels=targets)
+  targets_loss = tf.reduce_mean(targets_loss)
+
+  x_means = tf.matmul(discrete_x, means)
+  x_means = x + tf.stop_gradient(x_means - x)
+
   discrete_x = tf.reshape(discrete_x, x_shape[:-1] + [codebook_size])
-  target_means = tf.reshape(target_means, target_shape[:-1] + [hidden_size])
-  return discrete_x, target_means, code_loss, targets_loss
+  target_means = tf.reshape(target_means, target_shape + [hidden_size])
+  return discrete_x, x_means, target_means, code_loss, targets_loss
 
 
 def vq_discrete_unbottleneck(x, hidden_size):
@@ -992,8 +1012,8 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
   # [batch_size * latent_dim, num_blocks, block_v_size] to
   # [batch_size * num_blocks, latent_dim, block_v_size].
   dist = tf.reshape(dist, [batch_size, latent_dim, num_blocks, -1])
-  dist = tf.reshape(tf.transpose(dist, perm=[0, 2, 1, 3]),
-                    [-1, latent_dim, block_v_size])
+  dist = tf.reshape(
+      tf.transpose(dist, perm=[0, 2, 1, 3]), [-1, latent_dim, block_v_size])
   log_class_probs = tf.nn.log_softmax(-dist)
 
   sample_shape = [num_samples] + common_layers.shape_list(dist)
@@ -1010,7 +1030,7 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
 
   gumbel_softmax_samples = tf.nn.softmax(
       (tf.expand_dims(log_class_probs, 0) + gumbel_samples) / temperature)
-  q_samples = tf.clip_by_value(gumbel_softmax_samples, 1e-6, 1-1e-6)
+  q_samples = tf.clip_by_value(gumbel_softmax_samples, 1e-6, 1 - 1e-6)
 
   if approximate_gs_entropy:
     q_dist = tf.contrib.distributions.Multinomial(total_count=1.0, logits=-dist)
@@ -1036,8 +1056,8 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
 
       # Project samples from  [batch_size, latent_size, block_v_size] to
       # [batch_size, latent_size, hidden_size].
-      shifted_samples = common_layers.dense(
-          shifted_samples, hparams.hidden_size)
+      shifted_samples = common_layers.dense(shifted_samples,
+                                            hparams.hidden_size)
       # TODO(vafa): Include masking as a flag.
       mask = True
       if mask:
@@ -1061,11 +1081,12 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
       # which we can do without recalculating probabilities because the last
       # dimension of log_pi and q_samples are deterministic given the others.
       # Flow 2: Centered-softmax.
-      chained_bijectors = tf.contrib.distributions.bijectors.Chain(
-          [tf.contrib.distributions.bijectors.SoftmaxCentered(),
-           tf.contrib.distributions.bijectors.Affine(
-               shift=log_pi[:, :, :-1],
-               scale_identity_multiplier=1./temperature)])
+      chained_bijectors = tf.contrib.distributions.bijectors.Chain([
+          tf.contrib.distributions.bijectors.SoftmaxCentered(),
+          tf.contrib.distributions.bijectors.Affine(
+              shift=log_pi[:, :, :-1],
+              scale_identity_multiplier=1. / temperature)
+      ])
       q_samples = chained_bijectors.forward(q_samples[:, :, :-1])
       log_det = chained_bijectors.inverse_log_det_jacobian(
           q_samples, event_ndims=1)
@@ -1084,9 +1105,9 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
 
     # Take average of one-hot vectors over samples.
     x_means_hot = tf.reduce_mean(tf.one_hot(x_means_idx, block_v_size), 0)
-    x_means_assignments = (tf.reduce_mean(q_samples, 0) +
-                           tf.stop_gradient(x_means_hot - tf.reduce_mean(
-                               q_samples, 0)))
+    x_means_assignments = (
+        tf.reduce_mean(q_samples, 0) +
+        tf.stop_gradient(x_means_hot - tf.reduce_mean(q_samples, 0)))
   else:
     x_means_assignments = tf.reduce_mean(gumbel_softmax_samples, 0)
 
@@ -1094,12 +1115,10 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
   # block_v_size]. We have to transpose between reshapes to make sure the
   # dimensions have the correct interpretation.
   x_means_assignments = tf.reshape(
-      x_means_assignments,
-      [batch_size, num_blocks, latent_dim, block_v_size])
+      x_means_assignments, [batch_size, num_blocks, latent_dim, block_v_size])
   x_means_assignments = tf.transpose(x_means_assignments, [0, 2, 1, 3])
   x_means_assignments = tf.reshape(
-      x_means_assignments,
-      [batch_size * latent_dim, num_blocks, block_v_size])
+      x_means_assignments, [batch_size * latent_dim, num_blocks, block_v_size])
 
   return x_means_assignments, neg_q_entropy
 
@@ -1309,14 +1328,15 @@ def parametrized_bottleneck(x, hparams):
         soft_em=True,
         num_samples=hparams.vq_num_samples)
   if hparams.bottleneck_kind == "gumbel_softmax":
-    return gumbel_softmax_discrete_bottleneck(x,
-                                              hparams.bottleneck_bits,
-                                              hparams.vq_beta,
-                                              hparams.vq_decay,
-                                              hparams.vq_epsilon,
-                                              hparams.temperature_warmup_steps,
-                                              hard=False,
-                                              summary=True)
+    return gumbel_softmax_discrete_bottleneck(
+        x,
+        hparams.bottleneck_bits,
+        hparams.vq_beta,
+        hparams.vq_decay,
+        hparams.vq_epsilon,
+        hparams.temperature_warmup_steps,
+        hard=False,
+        summary=True)
 
   raise ValueError(
       "Unsupported hparams.bottleneck_kind %s" % hparams.bottleneck_kind)
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index 3970d14d8..c75f61704 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -111,7 +111,8 @@ def testNearestNeighbors(self):
     means = tf.constant(
         [[1, 0, 0], [0, 1, 0], [0, 0, 1], [9, 9, 9]], dtype=tf.float32)
     means = tf.stack([means, means], axis=0)
-    x_means_hot, _ = discretization.nearest_neighbor(x, means, block_v_size=4)
+    x_means_hot, _ = discretization.nearest_neighbor(
+        x, means, block_v_size=4)
     x_means_hot_test = np.array([[0, 1, 0, 0], [1, 0, 0, 0]])
     x_means_hot_test = np.expand_dims(x_means_hot_test, axis=0)
     with self.test_session() as sess:
@@ -139,7 +140,7 @@ def testVQNearestNeighbors(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
     means = tf.constant(
         [[1, 0, 0], [0, 1, 0], [0, 0, 1], [9, 9, 9]], dtype=tf.float32)
-    x_means_hot, _ = discretization.vq_nearest_neighbor(x, means)
+    x_means_hot, _, _ = discretization.vq_nearest_neighbor(x, means)
     x_means_hot_test = np.array([[0, 1, 0, 0], [1, 0, 0, 0]])
     with self.test_session() as sess:
       tf.global_variables_initializer().run()
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 6b9da76a4..37fa81f23 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -200,17 +200,20 @@ def body(self, features):
     else:
       raise Exception("Unsupported problem type: %s" % self.hparams.problem)
 
-    one_hot_labels = tf.one_hot(labels, vocab_size)
-    code_loss_gan = 0.0
+    losses = {}
     if hparams.gan_loss_factor != 0.0:
       res_gan, res = tf.split(res, 2, axis=0)
       with tf.variable_scope("vq"):
-        reconstr_gan, _, code_loss_gan, _ = discretization.vq_loss(
-            res, one_hot_labels, vocab_size)
+        reconstr_gan, gan_codes, _, code_loss_gan, _ = discretization.vq_loss(
+            res_gan, labels, vocab_size)
+        losses["code_loss_gan"] = code_loss_gan
 
     with tf.variable_scope("vq", reuse=tf.AUTO_REUSE):
-      reconstr, target_codes, code_loss, targets_loss = discretization.vq_loss(
-          res, one_hot_labels, vocab_size)
+      (reconstr, _, target_codes, code_loss,
+       targets_loss) = discretization.vq_loss(res, labels, vocab_size)
+
+    losses["code_loss"] = code_loss
+    losses["training"] = targets_loss
 
     # Add GAN loss if requested.
     gan_loss = 0.0
@@ -224,8 +227,14 @@ def body(self, features):
       def discriminate(x):
         return self.discriminator(x, is_training=is_training)
 
+      tc_shape = common_layers.shape_list(target_codes)
+      if len(tc_shape) > 4:
+        target_codes = tf.reshape(target_codes,
+                                  tc_shape[:-2] + [tc_shape[-1] * tc_shape[-2]])
+        gan_codes = tf.reshape(gan_codes,
+                               tc_shape[:-2] + [tc_shape[-1] * tc_shape[-2]])
       gan_loss = common_layers.sliced_gan_loss(target_codes,
-                                               reverse_gradient(res_gan),
+                                               reverse_gradient(gan_codes),
                                                discriminate,
                                                self.hparams.num_sliced_vecs)
       gan_loss *= hparams.gan_loss_factor
@@ -236,13 +245,11 @@ def discriminate(x):
           common_layers.tpu_safe_image_summary(tf.argmax(reconstr, -1)),
           max_outputs=1)
 
-    return reconstr, {
-        "training": targets_loss,
-        "code_loss": code_loss,
-        "code_loss_gan": code_loss_gan,
-        "b_loss": b_loss,
-        "gan_loss": -gan_loss
-    }
+    losses["b_loss"] = b_loss
+    losses["gan_loss"] = -gan_loss
+
+    logits = reconstr
+    return logits, losses
 
   def sample(self, features=None, shape=None):
     del features, shape
@@ -787,6 +794,7 @@ def autoencoder_basic():
   hparams.add_hparam("discriminator_batchnorm", True)
   hparams.add_hparam("num_sliced_vecs", 4096)
   hparams.add_hparam("gan_loss_factor", 0.0)
+  hparams.add_hparam("use_vqloss", False)
   return hparams
 
 
@@ -875,7 +883,26 @@ def autoencoder_ordered_discrete():
   hparams.gan_loss_factor = 0.0
   hparams.dropout = 0.1
   hparams.residual_dropout = 0.3
+  hparams.use_vqloss = True
   hparams.add_hparam("unordered", False)
+
+  return hparams
+
+
+@registry.register_hparams
+def autoencoder_ordered_discrete_novq():
+  """Ordered discrete autoencoder model."""
+  hparams = autoencoder_ordered_discrete()
+  hparams.use_vqloss = False
+
+  return hparams
+
+
+@registry.register_hparams
+def autoencoder_ordered_discrete_hs256():
+  """Ordered discrete autoencoder model."""
+  hparams = autoencoder_ordered_discrete()
+  hparams.hidden_size = 256
   return hparams
 
 

From 61d9aeb53d9086167dda8a8dbfbb5e69a3c70cf3 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 10 Aug 2018 11:16:49 -0700
Subject: [PATCH 0532/2720] switching to Session.

PiperOrigin-RevId: 208239036
---
 tensor2tensor/models/research/next_frame_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index b7bd2c13a..309f51ee4 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -102,7 +102,7 @@ def get_tensor_shape(tensor):
 class NextFrameTest(tf.test.TestCase):
 
   def RunModel(self, model, hparams, features):
-    with self.test_session() as session:
+    with self.Session() as session:
       model = model(
           hparams, tf.estimator.ModeKeys.TRAIN)
       logits, _ = model(features)

From 2e512d607b7a1bceb88fde0695e2d8db8154a868 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 10 Aug 2018 11:26:04 -0700
Subject: [PATCH 0533/2720] adding the 1st stage of training from original SV2P

PiperOrigin-RevId: 208240720
---
 tensor2tensor/models/research/next_frame_params.py |  1 +
 tensor2tensor/models/research/next_frame_sv2p.py   | 14 +++++++++++---
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index 85bb29d7f..ebcf46fb0 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -78,6 +78,7 @@ def next_frame_stochastic():
   hparams.add_hparam("num_masks", 10)
   hparams.add_hparam("latent_channels", 1)
   hparams.add_hparam("latent_std_min", -5.0)
+  hparams.add_hparam("num_iterations_1st_stage", 10000)
   hparams.add_hparam("num_iterations_2nd_stage", 10000)
   hparams.add_hparam("latent_loss_multiplier", 1e-3)
   hparams.add_hparam("latent_loss_multiplier_schedule", "constant")
diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 2a0905138..5f08980fe 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -66,7 +66,8 @@ def get_beta(self):
     """Get KL multiplier (beta) based on the schedule."""
     step_num = self.get_iteration_num()
     schedule = self.hparams.latent_loss_multiplier_schedule
-    second_stage = self.hparams.num_iterations_2nd_stage
+    second_stage = (self.hparams.num_iterations_1st_stage +
+                    self.hparams.num_iterations_2nd_stage)
     # TODO(mechcoder): Add log_annealing schedule.
     if schedule == "constant":
       beta = tf.cond(tf.greater(step_num, second_stage),
@@ -307,10 +308,17 @@ def construct_latent_tower(self, images):
       std += self.hparams.latent_std_min
 
       # No latent tower at inference time, just standard gaussian.
-      if self.hparams.mode != tf.estimator.ModeKeys.TRAIN:
+      if not self.is_training:
         return tf.zeros_like(mean), tf.zeros_like(std)
 
-      return mean, std
+      # No latent in the first phase
+      iter_num = self.get_iteration_num()
+      ret_mean, ret_std = tf.cond(
+          iter_num < self.hparams.num_iterations_1st_stage,
+          lambda: (tf.zeros_like(mean), tf.zeros_like(std)),
+          lambda: (mean, std))
+
+      return ret_mean, ret_std
 
   def reward_prediction(self, input_image, input_reward, action, latent):
     """Builds a reward prediction network."""

From 553823d21ce89f34e7e312ba6b04d31561d06712 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 10 Aug 2018 11:54:05 -0700
Subject: [PATCH 0534/2720] Add use_tpu_estimator flag, which allows using
 TPUEstimator even when not using TPU

PiperOrigin-RevId: 208245305
---
 tensor2tensor/bin/t2t_attack.py            |  8 +++---
 tensor2tensor/bin/t2t_trainer.py           |  5 +++-
 tensor2tensor/mesh_tensorflow/mtf_model.py |  6 ++---
 tensor2tensor/utils/t2t_model.py           | 29 ++++++++++------------
 tensor2tensor/utils/trainer_lib.py         | 18 +++++++++-----
 5 files changed, 36 insertions(+), 30 deletions(-)

diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index b15caea65..00cb48ae2 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -71,7 +71,10 @@ def main(argv):
   attack_params.add_hparam("eps", 0.0)
 
   config = t2t_trainer.create_run_config(hparams)
-  params = {"batch_size": hparams.batch_size}
+  params = {
+      "batch_size": hparams.batch_size,
+      "use_tpu": FLAGS.use_tpu,
+  }
 
   # add "_rev" as a hack to avoid image standardization
   problem = registry.problem(FLAGS.problem + "_rev")
@@ -85,8 +88,7 @@ def main(argv):
 
   sess = tf.Session()
 
-  model_fn = t2t_model.T2TModel.make_estimator_model_fn(
-      FLAGS.model, hparams, use_tpu=FLAGS.use_tpu)
+  model_fn = t2t_model.T2TModel.make_estimator_model_fn(FLAGS.model, hparams)
   ch_model = adv_attack_utils.T2TAttackModel(model_fn, params, config)
 
   acc_mask = None
diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 9dc07dced..e4caf6cf1 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -48,6 +48,8 @@
 flags.DEFINE_integer("iterations_per_loop", 100,
                      "Number of iterations in a TPU training loop.")
 flags.DEFINE_bool("use_tpu", False, "Whether to use TPU.")
+flags.DEFINE_bool("use_tpu_estimator", False, "Whether to use TPUEstimator. "
+                  "This is always enabled when use_tpu is True.")
 flags.DEFINE_integer("tpu_infeed_sleep_secs", None,
                      "How long to sleep the infeed thread.")
 flags.DEFINE_bool("generate_data", False, "Generate data before training?")
@@ -174,6 +176,7 @@ def create_experiment_fn(**kwargs):
       eval_early_stopping_metric_minimize=FLAGS.
       eval_early_stopping_metric_minimize,
       use_tpu=FLAGS.use_tpu,
+      use_tpu_estimator=FLAGS.use_tpu_estimator,
       **kwargs)
 
 
@@ -216,6 +219,7 @@ def create_run_config(hp):
       gpu_mem_fraction=FLAGS.worker_gpu_memory_fraction,
       enable_graph_rewriter=FLAGS.enable_graph_rewriter,
       use_tpu=FLAGS.use_tpu,
+      use_tpu_estimator=FLAGS.use_tpu_estimator,
       schedule=FLAGS.schedule,
       no_data_parallelism=hp.no_data_parallelism,
       daisy_chain_variables=daisy_chain_variables,
@@ -340,7 +344,6 @@ def main(argv):
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
   maybe_log_registry_and_exit()
 
-
   if FLAGS.cloud_mlengine:
     cloud_mlengine.launch()
     return
diff --git a/tensor2tensor/mesh_tensorflow/mtf_model.py b/tensor2tensor/mesh_tensorflow/mtf_model.py
index e39aa6ee7..769033ff9 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_model.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_model.py
@@ -49,11 +49,9 @@ def estimator_model_fn(cls,
                          mode,
                          config=None,
                          params=None,
-                         decode_hparams=None,
-                         use_tpu=False,
-                         xla_compile=False):
-    del xla_compile
+                         decode_hparams=None):
     hparams = copy.deepcopy(hparams)
+    use_tpu = params and params.get("use_tpu", False)
     hparams.use_tpu = use_tpu
     # merge decode_hparams into hparams if present
     if mode == tf.estimator.ModeKeys.PREDICT and decode_hparams is not None:
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 29b4d996a..7123cba3f 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -448,15 +448,14 @@ def loss(self, logits, features):
         target_modality = target_modality["targets"]
       return self._loss_single(logits, target_modality, features["targets"])
 
-  def optimize(self, loss, num_async_replicas=1):
+  def optimize(self, loss, num_async_replicas=1, use_tpu=False):
     """Return a training op minimizing loss."""
     lr = learning_rate.learning_rate_schedule(self.hparams)
     if num_async_replicas > 1:
       log_info("Dividing learning rate by num_async_replicas: %d",
                num_async_replicas)
     lr /= math.sqrt(float(num_async_replicas))
-    train_op = optimize.optimize(
-        loss, lr, self.hparams, use_tpu=common_layers.is_xla_compiled())
+    train_op = optimize.optimize(loss, lr, self.hparams, use_tpu=use_tpu)
     return train_op
 
   def set_mode(self, mode):
@@ -1144,8 +1143,7 @@ def get_eval_hooks(model_name):
   @staticmethod
   def make_estimator_model_fn(model_name,
                               hparams,
-                              decode_hparams=None,
-                              use_tpu=False):
+                              decode_hparams=None):
     model_cls = registry.model(model_name)
 
     def wrapping_model_fn(features, labels, mode, params=None, config=None):
@@ -1156,8 +1154,7 @@ def wrapping_model_fn(features, labels, mode, params=None, config=None):
           mode,
           config=config,
           params=params,
-          decode_hparams=decode_hparams,
-          use_tpu=use_tpu)
+          decode_hparams=decode_hparams)
 
     return wrapping_model_fn
 
@@ -1169,8 +1166,7 @@ def estimator_model_fn(cls,
                          mode,
                          config=None,
                          params=None,
-                         decode_hparams=None,
-                         use_tpu=False):
+                         decode_hparams=None):
     """Model fn for Estimator.
 
     Args:
@@ -1179,9 +1175,8 @@ def estimator_model_fn(cls,
       labels: Tensor
       mode: tf.estimator.ModeKeys
       config: RunConfig, possibly with data_parallelism attribute
-      params: dict, may include batch_size
+      params: dict, may include batch_size, use_tpu
       decode_hparams: HParams, used when mode == PREDICT.
-      use_tpu: bool, whether using TPU
 
     Returns:
       TPUEstimatorSpec if use tpu else EstimatorSpec
@@ -1190,6 +1185,7 @@ def estimator_model_fn(cls,
       _create_dummy_vars()
     hparams = copy.deepcopy(hparams)
 
+    use_tpu = params and params.get("use_tpu", False)
     # Instantiate model
     data_parallelism = None
     if not use_tpu and config:
@@ -1213,7 +1209,7 @@ def estimator_model_fn(cls,
       logits, losses_dict = model(features)  # pylint: disable=not-callable
 
     # Set known shapes
-    if use_tpu:
+    if common_layers.is_xla_compiled():
       if isinstance(logits, dict):
         for k, v in sorted(six.iteritems(logits)):
           if "scalar/" in k:
@@ -1257,11 +1253,12 @@ def estimator_model_fn(cls,
     num_async_replicas = (1 if (use_tpu or not config) else
                           config.t2t_device_info["num_async_replicas"])
     return model.estimator_spec_train(
-        loss, num_async_replicas=num_async_replicas)
+        loss, num_async_replicas=num_async_replicas, use_tpu=use_tpu)
 
-  def estimator_spec_train(self, loss, num_async_replicas=1):
+  def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
     """Construct EstimatorSpec for TRAIN mode."""
-    train_op = self.optimize(loss, num_async_replicas=num_async_replicas)
+    train_op = self.optimize(loss, num_async_replicas=num_async_replicas,
+                             use_tpu=use_tpu)
 
     # TODO(mitchellstern): Add support for partitioned variables?
     if (tf.train.latest_checkpoint(self._hparams.model_dir) is None and
@@ -1279,7 +1276,7 @@ def estimator_spec_train(self, loss, num_async_replicas=1):
               "Cannot find variable in checkpoint, skipping: %s", var_name)
       tf.train.init_from_checkpoint(pretrained_model_dir, variable_map)
 
-    if common_layers.is_xla_compiled():
+    if use_tpu:
       host_call = _create_host_call(self.hparams.model_dir)
       _remove_summaries()
       return tf.contrib.tpu.TPUEstimatorSpec(
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index c5ea9df24..9b4762f34 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -125,6 +125,7 @@ def create_run_config(master="",
                       sync=False,
                       tpu_infeed_sleep_secs=None,
                       use_tpu=False,
+                      use_tpu_estimator=False,
                       inter_op_parallelism_threads=0,
                       log_step_count_steps=100,
                       intra_op_parallelism_threads=0,
@@ -154,8 +155,9 @@ def create_run_config(master="",
     del run_config_args["save_checkpoints_steps"]
   run_config_cls = tf.contrib.learn.RunConfig
 
-  # If using TPU, use TPU RunConfig, add TPUConfig, and add additional args
-  if use_tpu:
+  # If using TPUEstimator, use TPU RunConfig, add TPUConfig, and add additional
+  # args.
+  if use_tpu or use_tpu_estimator:
     tpu_config_kwargs = {
         "iterations_per_loop": iterations_per_loop,
         "num_shards": num_shards,
@@ -207,12 +209,13 @@ def create_estimator(model_name,
                      run_config,
                      schedule="train_and_evaluate",
                      decode_hparams=None,
-                     use_tpu=False):
+                     use_tpu=False,
+                     use_tpu_estimator=False):
   """Create a T2T Estimator."""
   model_fn = t2t_model.T2TModel.make_estimator_model_fn(
-      model_name, hparams, decode_hparams=decode_hparams, use_tpu=use_tpu)
+      model_name, hparams, decode_hparams=decode_hparams)
 
-  if use_tpu:
+  if use_tpu or use_tpu_estimator:
     problem = hparams.problem
     batch_size = (
         problem.tpu_batch_size_per_shard(hparams) *
@@ -226,6 +229,7 @@ def create_estimator(model_name,
         model_fn=model_fn,
         model_dir=run_config.model_dir,
         config=run_config,
+        use_tpu=use_tpu,
         train_batch_size=batch_size,
         eval_batch_size=batch_size if "eval" in schedule else None,
         predict_batch_size=predict_batch_size)
@@ -412,6 +416,7 @@ def create_experiment(
     eval_early_stopping_metric_minimize=True,
     autotune=False,
     use_tpu=False,
+    use_tpu_estimator=False,
     additional_train_hooks=None,
     additional_eval_hooks=None):
   """Create Experiment."""
@@ -430,7 +435,8 @@ def create_experiment(
       run_config,
       schedule=schedule,
       decode_hparams=decode_hparams,
-      use_tpu=use_tpu)
+      use_tpu=use_tpu,
+      use_tpu_estimator=use_tpu_estimator)
 
   # Input fns from Problem
   problem = hparams.problem

From c8ded0741d7aa46109765d81248ca0a59833388a Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 10 Aug 2018 12:09:57 -0700
Subject: [PATCH 0535/2720] Set the default value of
 Text2TextProblem.max_subtoken_length to 200 and added logging in
 SubwordTextEncoder.build_from_token_counts if processing the token takes a
 long time.

PiperOrigin-RevId: 208247864
---
 tensor2tensor/data_generators/text_encoder.py  | 7 +++++++
 tensor2tensor/data_generators/text_problems.py | 6 +++---
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 64dbf39fa..2063282bc 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -28,6 +28,7 @@
 import math
 import re
 import tempfile
+import time
 import numpy as np
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
@@ -798,6 +799,7 @@ def build_from_token_counts(self,
       # subtoken boundaries.
       subtoken_counts = collections.defaultdict(int)
       for token, count in six.iteritems(token_counts):
+        iter_start_time = time.time()
         escaped_token = _escape_token(token, self._alphabet)
         subtokens = self._escaped_token_to_subtoken_strings(escaped_token)
         start = 0
@@ -810,6 +812,11 @@ def build_from_token_counts(self,
             new_subtoken = escaped_token[start:end]
             subtoken_counts[new_subtoken] += count
           start += len(subtoken)
+        iter_time_secs = time.time() - iter_start_time
+        if iter_time_secs > 0.1:
+          tf.logging.info("Processing token [{0}] took {1} seconds, consider "
+                          "setting Text2TextProblem.max_subtoken_length to a "
+                          "smaller value.".format(token, iter_time_secs))
 
       # Array of sets of candidate subtoken strings, by length.
       len_to_subtoken_strings = []
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 4e2448f12..25227be0a 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -260,13 +260,13 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
   def max_subtoken_length(self):
     """Maximum subtoken length when generating vocab.
 
-    Override with a finite integer (e.g. 100) to avoid quadratic-time vocab
-    building.
+    SubwordTextEncoder vocabulary building is quadratic-time wrt this variable,
+    setting it to None uses the length of the longest token in the corpus.
 
     Returns:
       an integer or None
     """
-    return None
+    return 200
 
   @property
   def batch_size_means_tokens(self):

From bba9dba72f8d835c309f2bcf994b1356f89ba297 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Fri, 10 Aug 2018 12:10:39 -0700
Subject: [PATCH 0536/2720] Task sampling support for more than two tasks.

PiperOrigin-RevId: 208247967
---
 .../data_generators/multi_problem.py          | 42 ++++++++++++++++---
 1 file changed, 36 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 5623fb18b..c16077184 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -90,7 +90,7 @@ def get_hparams(self, model_hparams=None):
 
   @property
   def mixing_schedule(self):
-    return MixingSchedule.EXPONENTIAL
+    return MixingSchedule.CONSTANT
 
   def flatten_zip(self, *args):
     """A list of examples to a dataset containing mixed examples.
@@ -201,11 +201,41 @@ def mix_data(example):
         tf.logging.info("Using the %s schedule to "
                         "train the MultiProblem." % str(self.mixing_schedule))
 
-        return tf.data.Dataset.from_tensors(tf.cond(
-            tf.greater(tf.random_uniform([]), prob),
-            lambda d=dataset_iterators[0]: get_next_from_dataset(d),
-            lambda d=dataset_iterators[1]: get_next_from_dataset(d)
-        ))
+        def sample_task(curr_task, num_tasks_left):
+          """A recursive function to sample a task.
+
+          This function treats the probability as the threshold for the primary
+          task and divides the remaining probability mass across the other
+          tasks.
+
+          Args:
+            curr_task: The index of the task being considered for sampling.
+            num_tasks_left: Number of tasks remaining to possibly sample from.
+
+          Returns:
+            A Tensor representing an example from the task that was sampled
+            from.
+          """
+
+          if num_tasks_left == 0:
+            return get_next_from_dataset(dataset_iterators[curr_task])
+          elif curr_task == 0:
+            # primary task
+            return tf.cond(
+                tf.greater(tf.random_uniform([]), prob),
+                lambda d=dataset_iterators[0]: get_next_from_dataset(d),
+                lambda c=curr_task+1, n=num_tasks_left-1: sample_task(c, n)
+            )
+          # divide the probability mass across all the secondary tasks equally.
+          new_prob = prob - curr_task * prob / (len(self.task_list)-1)
+          return tf.cond(
+              tf.greater(tf.random_uniform([]), new_prob),
+              lambda d=dataset_iterators[curr_task]: get_next_from_dataset(d),
+              lambda c=curr_task+1, n=num_tasks_left-1: sample_task(c, n)
+          )
+
+        return tf.data.Dataset.from_tensors(
+            sample_task(0, len(self.task_list)-1))
 
       single_mtl_dataset = tf.data.Dataset.from_tensors(tf.zeros([1])).repeat()
       single_mtl_dataset = single_mtl_dataset.flat_map(mix_data)

From 821ae0bf31b793b40cdcc5007bb24c41a5d699eb Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Fri, 10 Aug 2018 12:19:57 -0700
Subject: [PATCH 0537/2720] Add t2t_prune: simple magnitude-based pruning
 script for T2T

PiperOrigin-RevId: 208249376
---
 tensor2tensor/bin/t2t_prune.py       | 116 +++++++++++++++++++++++++++
 tensor2tensor/models/resnet.py       |  17 ++++
 tensor2tensor/utils/pruning_utils.py |  91 +++++++++++++++++++++
 tensor2tensor/utils/registry.py      |  96 +++++++++++++++++++++-
 4 files changed, 316 insertions(+), 4 deletions(-)
 create mode 100644 tensor2tensor/bin/t2t_prune.py
 create mode 100644 tensor2tensor/utils/pruning_utils.py

diff --git a/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
new file mode 100644
index 000000000..eae69df7c
--- /dev/null
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -0,0 +1,116 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+r"""Prune T2TModels using some heuristic.
+
+This supports a very common form of pruning known as magnitude-based pruning.
+It ranks individual weights or units according to their magnitudes and zeros
+out the smallest k% of weights, effectively removing them from the graph.
+
+Example run:
+- train a resnet on cifar10:
+    bin/t2t_trainer.py --problem=image_cifar10 --hparams_set=resnet_cifar_32 \
+      --model=resnet
+
+- evaluate different pruning percentages using weight-level pruning:
+    bin/t2t_prune.py --pruning_params_set=resnet_weight --problem=image_cifar10\
+      --hparams_set=resnet_cifar_32 --model=resnet
+"""
+
+import os
+
+from tensor2tensor.bin import t2t_trainer
+from tensor2tensor.data_generators import problem as problem_lib  # pylint: disable=unused-import
+from tensor2tensor.utils import pruning_utils
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
+from tensor2tensor.utils import trainer_lib
+from tensor2tensor.utils import usr_dir
+
+import tensorflow as tf
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+# See flags.py for additional command-line flags.
+flags.DEFINE_string("pruning_params_set", None,
+                    "Which pruning parameters to use.")
+
+
+def create_pruning_params():
+  return registry.pruning_params(FLAGS.pruning_params_set)
+
+
+def create_pruning_strategy(name):
+  return registry.pruning_strategies(name)
+
+
+def main(argv):
+  tf.logging.set_verbosity(tf.logging.INFO)
+  trainer_lib.set_random_seed(FLAGS.random_seed)
+  usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
+  t2t_trainer.maybe_log_registry_and_exit()
+
+
+  if FLAGS.generate_data:
+    t2t_trainer.generate_data()
+
+  if argv:
+    t2t_trainer.set_hparams_from_args(argv[1:])
+  hparams = t2t_trainer.create_hparams()
+  trainer_lib.add_problem_hparams(hparams, FLAGS.problem)
+  pruning_params = create_pruning_params()
+  pruning_strategy = create_pruning_strategy(pruning_params.strategy)
+
+  config = t2t_trainer.create_run_config(hparams)
+  params = {"batch_size": hparams.batch_size}
+
+  # add "_rev" as a hack to avoid image standardization
+  problem = registry.problem(FLAGS.problem)
+  input_fn = problem.make_estimator_input_fn(tf.estimator.ModeKeys.EVAL,
+                                             hparams)
+  dataset = input_fn(params, config).repeat()
+  features, labels = dataset.make_one_shot_iterator().get_next()
+
+  sess = tf.Session()
+
+  model_fn = t2t_model.T2TModel.make_estimator_model_fn(
+      FLAGS.model, hparams, use_tpu=FLAGS.use_tpu)
+  spec = model_fn(
+      features,
+      labels,
+      tf.estimator.ModeKeys.EVAL,
+      params=hparams,
+      config=config)
+
+  # Restore weights
+  saver = tf.train.Saver()
+  checkpoint_path = os.path.expanduser(FLAGS.output_dir or
+                                       FLAGS.checkpoint_path)
+  saver.restore(sess, tf.train.latest_checkpoint(checkpoint_path))
+
+  def eval_model():
+    preds = spec.predictions["predictions"]
+    preds = tf.argmax(preds, -1, output_type=labels.dtype)
+    _, acc_update_op = tf.metrics.accuracy(labels=labels, predictions=preds)
+    sess.run(tf.initialize_local_variables())
+    for _ in range(FLAGS.eval_steps):
+      acc = sess.run(acc_update_op)
+    return acc
+
+  pruning_utils.sparsify(sess, eval_model, pruning_strategy, pruning_params)
+
+
+if __name__ == "__main__":
+  tf.app.run()
diff --git a/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index a6996f773..5a5b1ccbd 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -562,3 +562,20 @@ def resnet_200():
   hp = resnet_base()
   hp.layer_sizes = [3, 24, 36, 3]
   return hp
+
+
+@registry.register_pruning_params
+def resnet_weight():
+  hp = tf.contrib.training.HParams()
+  hp.add_hparam("strategy", "weight")
+  hp.add_hparam("black_list", ["logits", "bias"])
+  hp.add_hparam("white_list", None)
+  hp.add_hparam("sparsities", [0.1*i for i in range(10)])
+  return hp
+
+
+@registry.register_pruning_params
+def resnet_unit():
+  hp = resnet_weight()
+  hp.strategy = "unit"
+  return hp
diff --git a/tensor2tensor/utils/pruning_utils.py b/tensor2tensor/utils/pruning_utils.py
new file mode 100644
index 000000000..70cb3621a
--- /dev/null
+++ b/tensor2tensor/utils/pruning_utils.py
@@ -0,0 +1,91 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Utilities to assist in pruning models."""
+
+import numpy as np
+
+from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+@registry.register_pruning_strategy
+def weight(w, sparsity):
+  """Weight-level magnitude pruning."""
+  w_shape = common_layers.shape_list(w)
+  k = int(np.prod(w_shape[:-1]))
+  w = tf.reshape(w, [k, w_shape[-1]])
+
+  idx = int(k * sparsity)
+  thres = tf.contrib.framework.sort(tf.abs(w), axis=0)[idx]
+  mask = tf.to_float(thres[None, :] < tf.abs(w))
+
+  w = mask * w
+  return tf.reshape(w, w_shape)
+
+
+@registry.register_pruning_strategy
+def unit(w, sparsity):
+  """Unit-level magnitude pruning."""
+  w_shape = common_layers.shape_list(w)
+  k = int(np.prod(w_shape[:-1]))
+  w = tf.reshape(w, [k, w_shape[-1]])
+  idx = int(w_shape[-1] * sparsity)
+
+  norm = tf.norm(w, axis=0)
+  thres = tf.contrib.framework.sort(norm, axis=0)[idx]
+  mask = tf.to_float(thres < norm)[None, :]
+  mask = tf.tile(mask, [k, 1])
+
+  w = mask * w
+  return tf.reshape(w, w_shape)
+
+
+def sparsify(sess, eval_model, pruning_strategy, pruning_params):
+  """Prune the weights of a model and evaluate."""
+  weights = tf.trainable_variables()
+
+  def should_prune(name):
+    """Whether to prune a weight or not."""
+    in_whitelist = not pruning_params.white_list or any(
+        e in name for e in pruning_params.white_list)
+    in_blacklist = any(e in name for e in pruning_params.black_list)
+
+    if pruning_params.white_list and not in_whitelist:
+      return False
+    elif in_blacklist:
+      return False
+
+    return True
+
+  weights = [w for w in weights if should_prune(w.name)]
+  unpruned_weights = sess.run(weights)
+
+  reset_op = tf.no_op()
+  for w, ow in zip(weights, unpruned_weights):
+    op = tf.assign(w, ow)
+    reset_op = tf.group(reset_op, op)
+
+  for sparsity in pruning_params.sparsities:
+    set_weights_op = tf.no_op()
+    for w in weights:
+      op = tf.assign(w, pruning_strategy(w, sparsity))
+      set_weights_op = tf.group(set_weights_op, op)
+    sess.run(set_weights_op)
+
+    acc = eval_model()
+    tf.logging.info("\tPruning to sparsity = %f: acc = %f" % (sparsity, acc))
+    sess.run(reset_op)
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 59f6711ee..1b27ad118 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -48,12 +48,14 @@ class MyModel(T2TModel):
 import six
 import tensorflow as tf
 
-_MODELS = {}
-_HPARAMS = {}
-_RANGED_HPARAMS = {}
 _ATTACKS = {}
 _ATTACK_PARAMS = {}
+_HPARAMS = {}
+_MODELS = {}
 _PROBLEMS = {}
+_PRUNING_PARAMS = {}
+_PRUNING_STRATEGY = {}
+_RANGED_HPARAMS = {}
 
 
 class Modalities(object):
@@ -361,6 +363,84 @@ def list_attack_params(prefix=None):
   return list(_ATTACK_PARAMS)
 
 
+def register_pruning_params(name=None):
+  """Register an pruning HParams set. Same behaviour as register_hparams."""
+
+  def decorator(pp_fn, registration_name=None):
+    """Registers & returns pp_fn with registration_name or default name."""
+    pp_name = registration_name or default_name(pp_fn)
+    if pp_name in _PRUNING_PARAMS and not tf.contrib.eager.in_eager_mode():
+      raise LookupError("Pruning HParams set %s already registered." % pp_name)
+    _PRUNING_PARAMS[pp_name] = pp_fn
+    return pp_fn
+
+  # Handle if decorator was used without parens
+  if callable(name):
+    pp_fn = name
+    return decorator(pp_fn, registration_name=default_name(pp_fn))
+
+  return lambda pp_fn: decorator(pp_fn, name)
+
+
+def pruning_params(name):
+  """Retrieve registered pruning params by name."""
+  if name not in _PRUNING_PARAMS:
+    error_msg = "Pruning HParams set %s never registered. Sets registered:\n%s"
+    raise LookupError(error_msg % (
+        name, display_list_by_prefix(list_pruning_params(), starting_spaces=4)))
+  pp = _PRUNING_PARAMS[name]()
+  if pp is None:
+    raise TypeError("Pruning HParams %s is None. Make sure the registered "
+                    "function returns the HParams object." % name)
+  return pp
+
+
+def list_pruning_params(prefix=None):
+  if prefix:
+    return [name for name in _PRUNING_PARAMS if name.startswith(prefix)]
+  return list(_PRUNING_PARAMS)
+
+
+def register_pruning_strategy(name=None):
+  """Register an pruning strategy. Same behaviour as register_hparams."""
+
+  def decorator(ps_fn, registration_name=None):
+    """Registers & returns ps_fn with registration_name or default name."""
+    ps_name = registration_name or default_name(ps_fn)
+    if ps_name in _PRUNING_STRATEGY and not tf.contrib.eager.in_eager_mode():
+      raise LookupError("Pruning strategy %s already registered." % ps_name)
+    _PRUNING_STRATEGY[ps_name] = ps_fn
+    return ps_fn
+
+  # Handle if decorator was used without parens
+  if callable(name):
+    ps_fn = name
+    return decorator(ps_fn, registration_name=default_name(ps_fn))
+
+  return lambda ps_fn: decorator(ps_fn, name)
+
+
+def pruning_strategies(name):
+  """Retrieve registered pruning strategies by name."""
+  if name not in _PRUNING_STRATEGY:
+    error_msg = "Pruning strategy set %s never registered. Sets registered:\n%s"
+    raise LookupError(
+        error_msg % (name,
+                     display_list_by_prefix(
+                         list_pruning_strategies(), starting_spaces=4)))
+  ps = _PRUNING_STRATEGY[name]
+  if ps is None:
+    raise TypeError("Pruning strategy %s is None. Make sure to register the "
+                    "function." % name)
+  return ps
+
+
+def list_pruning_strategies(prefix=None):
+  if prefix:
+    return [name for name in _PRUNING_STRATEGY if name.startswith(prefix)]
+  return list(_PRUNING_STRATEGY)
+
+
 def _internal_get_modality(name, mod_collection, collection_str):
   if name is None:
     name = "default"
@@ -559,6 +639,12 @@ def help_string():
 
   Attack HParams:
 %s
+
+  Pruning HParams:
+%s
+
+  Pruning Strategies:
+%s
 """
   m, hp, rhp, mod, probs, atks, ap = [
       display_list_by_prefix(entries, starting_spaces=4) for entries in [
@@ -568,7 +654,9 @@ def help_string():
           list_modalities(),
           list_problems(),
           list_attacks(),
-          list_attack_params()
+          list_attack_params(),
+          list_pruning_params(),
+          list_pruning_strategies(),
       ]
   ]
   return help_str % (m, hp, rhp, mod, probs, atks, ap)

From 8c4893cad0cfad81a0dd087901cf8abe86b5cf73 Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Fri, 10 Aug 2018 12:31:35 -0700
Subject: [PATCH 0538/2720] Update adversarial attacks in T2T

PiperOrigin-RevId: 208250924
---
 tensor2tensor/bin/t2t_attack.py          | 193 ++++++++++++++++++++---
 tensor2tensor/data_generators/problem.py |   8 +-
 tensor2tensor/models/resnet.py           |  51 +++++-
 tensor2tensor/utils/adv_attack_utils.py  | 161 +++++++++++++++++--
 tensor2tensor/utils/t2t_model.py         |   4 +-
 5 files changed, 376 insertions(+), 41 deletions(-)

diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 00cb48ae2..7368242e9 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -12,7 +12,29 @@
 # 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.
-"""Adversarially attack a model."""
+r"""Adversarially attack a model.
+
+This script adversarially attacks a model and evaluates accuracy at various
+  epsilons.
+
+Params such as which epsilons to evaluate at and the attack algorithm are
+  specified by attack_params, see models/resnet.py for examples.
+
+--ignore_incorrect will only attack those examples that are already correctly
+  classified by the model.
+
+--surrogate_attack will attack a model (A) and evaluate adversarial examples for
+  A on a different model (B).
+
+Example run:
+- train a resnet on cifar10:
+    bin/t2t_trainer.py --problem=image_cifar10 --hparams_set=resnet_cifar_32 \
+      --model=resnet
+
+- evaluate robustness using the FGSM attack:
+    bin/t2t_attack.py --attack_params_set=resnet_fgsm --problem=image_cifar10\
+      --hparams_set=resnet_cifar_32 --model=resnet
+"""
 
 import os
 
@@ -33,6 +55,13 @@
 # See flags.py for additional command-line flags.
 flags.DEFINE_string("attack_params_set", None,
                     "Which attack parameters to use.")
+flags.DEFINE_boolean("surrogate_attack", False,
+                     "Perform an attack on a surrogate model.")
+flags.DEFINE_string("surrogate_model", None, "Surrogate model to attack.")
+flags.DEFINE_string("surrogate_hparams_set", None,
+                    "Surrogate model's hyperparameter set.")
+flags.DEFINE_string("surrogate_output_dir", None,
+                    "Directory storing surrogate model's weights.")
 flags.DEFINE_boolean(
     "ignore_incorrect", False, "Ignore examples that are "
     "incorrectly classified to begin with.")
@@ -46,6 +75,75 @@ def create_attack(attack):
   return registry.attacks(attack)
 
 
+def create_surrogate_hparams():
+  return trainer_lib.create_hparams(FLAGS.surrogate_hparams_set, None)
+
+
+def create_surrogate_run_config(hp):
+  """Create a run config.
+
+  Args:
+    hp: model hyperparameters
+  Returns:
+    a run config
+  """
+  save_ckpt_steps = max(FLAGS.iterations_per_loop, FLAGS.local_eval_frequency)
+  save_ckpt_secs = FLAGS.save_checkpoints_secs or None
+  if save_ckpt_secs:
+    save_ckpt_steps = None
+  assert FLAGS.surrogate_output_dir
+  # the various custom getters we have written do not play well together yet.
+  # TODO(noam): ask rsepassi for help here.
+  daisy_chain_variables = (
+      hp.daisy_chain_variables and hp.activation_dtype == "float32" and
+      hp.weight_dtype == "float32")
+  return trainer_lib.create_run_config(
+      model_dir=os.path.expanduser(FLAGS.surrogate_output_dir),
+      master=FLAGS.master,
+      iterations_per_loop=FLAGS.iterations_per_loop,
+      num_shards=FLAGS.tpu_num_shards,
+      log_device_placement=FLAGS.log_device_placement,
+      save_checkpoints_steps=save_ckpt_steps,
+      save_checkpoints_secs=save_ckpt_secs,
+      keep_checkpoint_max=FLAGS.keep_checkpoint_max,
+      keep_checkpoint_every_n_hours=FLAGS.keep_checkpoint_every_n_hours,
+      num_gpus=FLAGS.worker_gpu,
+      gpu_order=FLAGS.gpu_order,
+      shard_to_cpu=FLAGS.locally_shard_to_cpu,
+      num_async_replicas=FLAGS.worker_replicas,
+      gpu_mem_fraction=FLAGS.worker_gpu_memory_fraction,
+      enable_graph_rewriter=FLAGS.enable_graph_rewriter,
+      use_tpu=FLAGS.use_tpu,
+      schedule=FLAGS.schedule,
+      no_data_parallelism=hp.no_data_parallelism,
+      daisy_chain_variables=daisy_chain_variables,
+      ps_replicas=FLAGS.ps_replicas,
+      ps_job=FLAGS.ps_job,
+      ps_gpu=FLAGS.ps_gpu,
+      sync=FLAGS.sync,
+      worker_id=FLAGS.worker_id,
+      worker_job=FLAGS.worker_job,
+      random_seed=FLAGS.random_seed,
+      tpu_infeed_sleep_secs=FLAGS.tpu_infeed_sleep_secs,
+      inter_op_parallelism_threads=FLAGS.inter_op_parallelism_threads,
+      log_step_count_steps=FLAGS.log_step_count_steps,
+      intra_op_parallelism_threads=FLAGS.intra_op_parallelism_threads)
+
+
+def prepare_data(problem, hparams, params, config):
+  """Construct input pipeline."""
+  input_fn = problem.make_estimator_input_fn(
+      tf.estimator.ModeKeys.EVAL, hparams, force_repeat=True)
+  dataset = input_fn(params, config)
+  features, _ = dataset.make_one_shot_iterator().get_next()
+  inputs, labels = features["targets"], features["inputs"]
+  inputs = tf.to_float(inputs)
+  input_shape = inputs.shape.as_list()
+  inputs = tf.reshape(inputs, [hparams.batch_size] + input_shape[1:])
+  labels = tf.reshape(labels, [hparams.batch_size])
+  return inputs, labels, features
+
+
 def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
   trainer_lib.set_random_seed(FLAGS.random_seed)
@@ -65,11 +163,20 @@ def main(argv):
 
   if argv:
     t2t_trainer.set_hparams_from_args(argv[1:])
+
+  if FLAGS.surrogate_attack:
+    tf.logging.warn("Performing surrogate model attack.")
+    sur_hparams = create_surrogate_hparams()
+    trainer_lib.add_problem_hparams(sur_hparams, FLAGS.problem)
+
   hparams = t2t_trainer.create_hparams()
   trainer_lib.add_problem_hparams(hparams, FLAGS.problem)
+
   attack_params = create_attack_params()
-  attack_params.add_hparam("eps", 0.0)
+  attack_params.add_hparam(attack_params.epsilon_name, 0.0)
 
+  if FLAGS.surrogate_attack:
+    sur_config = create_surrogate_run_config(sur_hparams)
   config = t2t_trainer.create_run_config(hparams)
   params = {
       "batch_size": hparams.batch_size,
@@ -78,59 +185,93 @@ def main(argv):
 
   # add "_rev" as a hack to avoid image standardization
   problem = registry.problem(FLAGS.problem + "_rev")
-  input_fn = problem.make_estimator_input_fn(
-      tf.estimator.ModeKeys.EVAL, hparams)
-  dataset = input_fn(params, config).repeat()
-  features, _ = dataset.make_one_shot_iterator().get_next()
-  inputs, labels = features["targets"], features["inputs"]
-  inputs = tf.to_float(inputs)
-  labels = tf.squeeze(labels)
+
+  inputs, labels, features = prepare_data(problem, hparams, params, config)
 
   sess = tf.Session()
 
-  model_fn = t2t_model.T2TModel.make_estimator_model_fn(FLAGS.model, hparams)
-  ch_model = adv_attack_utils.T2TAttackModel(model_fn, params, config)
+  if FLAGS.surrogate_attack:
+    sur_model_fn = t2t_model.T2TModel.make_estimator_model_fn(
+        FLAGS.surrogate_model, sur_hparams, use_tpu=FLAGS.use_tpu)
+    sur_ch_model = adv_attack_utils.T2TAttackModel(
+        sur_model_fn, features, params, sur_config, scope="surrogate")
+    # Dummy call to construct graph
+    sur_ch_model.get_probs(inputs)
+
+    checkpoint_path = os.path.expanduser(FLAGS.surrogate_output_dir)
+    tf.contrib.framework.init_from_checkpoint(
+        tf.train.latest_checkpoint(checkpoint_path), {"/": "surrogate/"})
+    sess.run(tf.global_variables_initializer())
+
+  other_vars = set(tf.global_variables())
+
+  model_fn = t2t_model.T2TModel.make_estimator_model_fn(
+      FLAGS.model, hparams, use_tpu=FLAGS.use_tpu)
+  ch_model = adv_attack_utils.T2TAttackModel(model_fn, features, params, config)
 
   acc_mask = None
   probs = ch_model.get_probs(inputs)
   if FLAGS.ignore_incorrect:
-    preds = tf.argmax(probs, -1)
-    preds = tf.squeeze(preds)
+    preds = tf.argmax(probs, -1, output_type=labels.dtype)
+    preds = tf.reshape(preds, labels.shape)
     acc_mask = tf.to_float(tf.equal(labels, preds))
   one_hot_labels = tf.one_hot(labels, probs.shape[-1])
 
-  attack = create_attack(attack_params.attack)(ch_model, sess=sess)
+  if FLAGS.surrogate_attack:
+    attack = create_attack(attack_params.attack)(sur_ch_model, sess=sess)
+  else:
+    attack = create_attack(attack_params.attack)(ch_model, sess=sess)
+
+  new_vars = set(tf.global_variables()) - other_vars
 
   # Restore weights
-  saver = tf.train.Saver()
-  checkpoint_path = os.path.expanduser(FLAGS.output_dir or
-                                       FLAGS.checkpoint_path)
+  saver = tf.train.Saver(new_vars)
+  checkpoint_path = os.path.expanduser(FLAGS.output_dir)
   saver.restore(sess, tf.train.latest_checkpoint(checkpoint_path))
 
   # reuse variables
   tf.get_variable_scope().reuse_variables()
 
-  def compute_accuracy(x, labels, mask):
+  def compute_accuracy(x, l, mask):
+    """Compute model accuracy."""
     preds = ch_model.get_probs(x)
     preds = tf.squeeze(preds)
-    preds = tf.argmax(preds, -1, output_type=labels.dtype)
-    _, acc_update_op = tf.metrics.accuracy(
-        labels=labels, predictions=preds, weights=mask)
+    preds = tf.argmax(preds, -1, output_type=l.dtype)
+
+    _, acc_update_op = tf.metrics.accuracy(l, preds, weights=mask)
+
+    if FLAGS.surrogate_attack:
+      preds = sur_ch_model.get_probs(x)
+      preds = tf.squeeze(preds)
+      preds = tf.argmax(preds, -1, output_type=l.dtype)
+      acc_update_op = tf.tuple((acc_update_op,
+                                tf.metrics.accuracy(l, preds, weights=mask)[1]))
+
     sess.run(tf.initialize_local_variables())
-    for _ in range(FLAGS.eval_steps):
+    for i in range(FLAGS.eval_steps):
+      tf.logging.info(
+          "\tEvaluating batch [%d / %d]" % (i + 1, FLAGS.eval_steps))
       acc = sess.run(acc_update_op)
+    if FLAGS.surrogate_attack:
+      tf.logging.info("\tFinal acc: (%.4f, %.4f)" % (acc[0], acc[1]))
+    else:
+      tf.logging.info("\tFinal acc: %.4f" % acc)
     return acc
 
-  acc = compute_accuracy(inputs, labels, acc_mask)
-  epsilon_acc_pairs = [(0.0, acc)]
+  epsilon_acc_pairs = []
   for epsilon in attack_params.attack_epsilons:
-    attack_params.eps = epsilon
+    tf.logging.info("Attacking @ eps=%.4f" % epsilon)
+    attack_params.set_hparam(attack_params.epsilon_name, epsilon)
     adv_x = attack.generate(inputs, y=one_hot_labels, **attack_params.values())
     acc = compute_accuracy(adv_x, labels, acc_mask)
     epsilon_acc_pairs.append((epsilon, acc))
 
   for epsilon, acc in epsilon_acc_pairs:
-    tf.logging.info("Accuracy @ eps=%f: %f" % (epsilon, acc))
+    if FLAGS.surrogate_attack:
+      tf.logging.info(
+          "Accuracy @ eps=%.4f: (%.4f, %.4f)" % (epsilon, acc[0], acc[1]))
+    else:
+      tf.logging.info("Accuracy @ eps=%.4f: %.4f" % (epsilon, acc))
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 3ffc6b01b..061112687 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -720,6 +720,7 @@ def make_estimator_input_fn(self,
                               mode,
                               hparams,
                               data_dir=None,
+                              force_repeat=False,
                               dataset_kwargs=None):
     """Return input_fn wrapped for Estimator."""
 
@@ -730,6 +731,7 @@ def estimator_input_fn(params, config):
           data_dir=data_dir,
           params=params,
           config=config,
+          force_repeat=force_repeat,
           dataset_kwargs=dataset_kwargs)
 
     return estimator_input_fn
@@ -774,6 +776,7 @@ def input_fn(self,
                data_dir=None,
                params=None,
                config=None,
+               force_repeat=False,
                dataset_kwargs=None):
     """Builds input pipeline for problem.
 
@@ -784,6 +787,7 @@ def input_fn(self,
       params: dict, may include "batch_size"
       config: RunConfig; should have the data_parallelism attribute if not using
         TPU
+      force_repeat: bool, whether to repeat the data even if not training
       dataset_kwargs: dict, if passed, will pass as kwargs to self.dataset
         method when called
 
@@ -828,9 +832,11 @@ def define_shapes(example):
     })
 
     dataset = self.dataset(**dataset_kwargs)
-    if is_training:
+    if force_repeat or is_training:
       # Repeat and skip a random number of records
       dataset = dataset.repeat()
+
+    if is_training:
       data_files = tf.contrib.slim.parallel_reader.get_data_files(
           self.filepattern(data_dir, mode))
       #  In continuous_train_and_eval when switching between train and
diff --git a/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index 5a5b1ccbd..dd185a669 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -18,6 +18,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import registry
@@ -386,6 +387,8 @@ def body(self, features):
     }
     assert hp.block_fn in block_fns
     is_training = hp.mode == tf.estimator.ModeKeys.TRAIN
+    if is_training:
+      targets = features["targets_raw"]
 
     inputs = features["inputs"]
 
@@ -426,7 +429,21 @@ def body(self, features):
     if hp.use_nchw:
       out = tf.transpose(out, [0, 2, 3, 1])
 
-    return out
+    out = tf.reduce_mean(out, [1, 2])
+    num_classes = self._problem_hparams.target_modality.top_dimensionality
+    logits = tf.layers.dense(out, num_classes)
+
+    losses = {"training": 0.0}
+    if is_training:
+      loss = tf.losses.sparse_softmax_cross_entropy(
+          labels=tf.squeeze(targets), logits=logits)
+      loss = tf.reduce_mean(loss)
+
+      losses = {"training": loss}
+
+    logits = tf.reshape(logits, [-1, 1, 1, 1, logits.shape[1]])
+
+    return logits, losses
 
   def infer(self,
             features=None,
@@ -564,6 +581,7 @@ def resnet_200():
   return hp
 
 
+# Pruning parameters
 @registry.register_pruning_params
 def resnet_weight():
   hp = tf.contrib.training.HParams()
@@ -579,3 +597,34 @@ def resnet_unit():
   hp = resnet_weight()
   hp.strategy = "unit"
   return hp
+
+
+# Adversarial attack parameters
+@registry.register_attack_params
+def resnet_fgsm():
+  aparams = tf.contrib.training.HParams()
+  aparams.attack = "fgsm"
+  aparams.epsilon_name = "eps"
+  aparams.attack_epsilons = [i * 0.8 for i in range(20)]
+  aparams.add_hparam("clip_min", 0.0)
+  aparams.add_hparam("clip_max", 255.0)
+  return aparams
+
+
+@registry.register_attack_params
+def resnet_madry():
+  aparams = resnet_fgsm()
+  aparams.attack = "madry"
+  aparams.add_hparam("nb_iter", 40)
+  aparams.add_hparam("eps_iter", 1.0)
+  return aparams
+
+
+@registry.register_attack_params
+def resnet_random():
+  aparams = resnet_fgsm()
+  aparams.attack = "random"
+  aparams.epsilon_name = "eps"
+  aparams.add_hparam("num_samples", 10)
+  aparams.add_hparam("num_batches", 100)
+  return aparams
diff --git a/tensor2tensor/utils/adv_attack_utils.py b/tensor2tensor/utils/adv_attack_utils.py
index d543f7615..9418854e9 100644
--- a/tensor2tensor/utils/adv_attack_utils.py
+++ b/tensor2tensor/utils/adv_attack_utils.py
@@ -16,6 +16,9 @@
 
 from cleverhans import attacks
 from cleverhans import model
+from cleverhans import utils_tf
+
+import numpy as np
 
 from tensor2tensor.utils import registry
 
@@ -32,29 +35,165 @@ def madry():
   return attacks.MadryEtAl
 
 
+@registry.register_attack
+def random():
+  return RandomAttack
+
+
 class T2TAttackModel(model.Model):
   """Wrapper of Cleverhans Model object."""
 
-  def __init__(self, model_fn, params, config):
+  def __init__(self, model_fn, features, params, config, scope=None):
     self._model_fn = model_fn
     self._params = params
     self._config = config
     self._logits_dict = {}
+    self._additional_features = features
+    self._scope = scope
 
-  def get_logits(self, x):
+  def fprop(self, x):
     if x.name in self._logits_dict:
       return self._logits_dict[x.name]
 
     x = tf.map_fn(tf.image.per_image_standardization, x)
+    self._additional_features['inputs'] = x
 
-    logits = self._model_fn(
-        {
-            "inputs": x
-        },
-        None,
-        "attack",
-        params=self._params,
-        config=self._config)
+    if self._scope is None:
+      scope = tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE)
+    else:
+      scope = tf.variable_scope(self._scope, reuse=tf.AUTO_REUSE)
+
+    with scope:
+      logits = self._model_fn(
+          self._additional_features,
+          None,
+          'attack',
+          params=self._params,
+          config=self._config)
     self._logits_dict[x.name] = logits
 
-    return tf.squeeze(logits)
+    return {model.Model.O_LOGITS: tf.reshape(logits, [-1, logits.shape[-1]])}
+
+
+class RandomAttack(attacks.FastGradientMethod):
+  """Blackbox random sample attack."""
+
+  def __init__(self, m, back='tf', sess=None):
+    if not isinstance(m, model.Model):
+      m = model.CallableModelWrapper(m, 'probs')
+
+    super(RandomAttack, self).__init__(m, back, sess)
+    self.feedable_kwargs = {
+        'eps': np.float32,
+        'num_samples': np.float32,
+        'num_batches': np.float32,
+        'y': np.float32,
+        'y_target': np.float32,
+        'clip_min': np.float32,
+        'clip_max': np.float32
+    }
+    self.structural_kwargs = ['ord']
+
+  def generate(self, x, **kwargs):
+    # Parse and save attack-specific parameters
+    assert self.parse_params(**kwargs)
+
+    labels, _ = self.get_or_guess_labels(x, kwargs)
+
+    x_shape = x.shape.as_list()
+    deltas_shape = [x_shape[0], self.num_samples] + x_shape[1:]
+
+    def cond(i, old_adv_x, old_loss):
+      del old_adv_x, old_loss
+      return tf.less(i, self.num_batches)
+
+    def body(i, old_adv_x, old_loss, labels=labels):
+      """Find example with max loss value amongst batch of perturbations."""
+      deltas = tf.random_uniform(deltas_shape)
+
+      # generate uniform samples from the l^p unit ball interior
+      if self.ord == np.inf:
+        deltas *= 2. * self.eps
+        deltas -= self.eps
+      elif self.ord == 1:
+        # ref: https://mathoverflow.net/questions/9185/how-to-generate-random-points-in-ell-p-balls  pylint: disable=line-too-long
+        exp = -tf.log(deltas)
+        shift = -tf.log(tf.random_uniform(deltas_shape[:2]))
+        norm = tf.reduce_sum(tf.abs(exp), range(2, len(deltas_shape) - 2))
+        scale = tf.reshape(shift + norm,
+                           deltas_shape[:2] + [1] * (len(deltas_shape) - 2))
+        deltas = exp / scale
+      elif self.ord == 2:
+        # ref: https://blogs.sas.com/content/iml/2016/04/06/generate-points-uniformly-in-ball.html  pylint: disable=line-too-long
+        dims = tf.reduce_prod(deltas_shape[2:])
+        deltas = tf.pow(deltas, 1. / dims)
+        normal = tf.random_normal(deltas)
+        normal /= tf.sqrt(
+            tf.reduce_sum(normal**2, axis=range(2,
+                                                len(deltas_shape) - 2)),
+            keepdims=True)
+        deltas *= normal
+      else:
+        raise NotImplementedError('Only L-inf, L1 and L2 norms are '
+                                  'currently implemented.')
+
+      adv_x = tf.expand_dims(x, 1) + deltas
+      labels = tf.expand_dims(labels, 1)
+      labels = tf.tile(labels, [1, self.num_samples, 1])
+
+      if (self.clip_min is not None) and (self.clip_max is not None):
+        adv_x = tf.clip_by_value(adv_x, self.clip_min, self.clip_max)
+
+      adv_x_r = tf.reshape(adv_x, [-1] + deltas_shape[2:])
+      preds = self.model.get_probs(adv_x_r)
+      preds_shape = preds.shape.as_list()
+      preds = tf.reshape(preds, deltas_shape[:2] + preds_shape[1:])
+
+      if labels is None:
+        # Using model predictions as ground truth to avoid label leaking
+        preds_max = tf.reduce_max(preds, -1, keep_dims=True)
+        labels = tf.to_float(tf.equal(preds, preds_max))
+        labels = tf.stop_gradient(labels)
+      labels = labels / tf.reduce_sum(labels, -1, keep_dims=True)
+
+      # Compute loss
+      loss = utils_tf.model_loss(labels, preds, mean=False)
+      if self.y_target is not None:
+        loss = -loss
+
+      # find the maximum loss value
+      input_idx = tf.one_hot(tf.argmax(loss, axis=1), self.num_samples, axis=1)
+      loss = tf.reduce_sum(loss * input_idx, axis=1)
+      input_idx = tf.reshape(input_idx,
+                             deltas_shape[:2] + [1] * (len(deltas_shape) - 2))
+      adv_x = tf.reduce_sum(adv_x * input_idx, axis=1)
+
+      condition = tf.greater(old_loss, loss)
+      new_loss = tf.where(condition, old_loss, loss)
+      new_adv_x = tf.where(condition, old_adv_x, adv_x)
+      print(new_loss, new_adv_x)
+
+      return i + 1, new_adv_x, new_loss
+
+    _, adv_x, _ = tf.while_loop(
+        cond, body,
+        [tf.zeros([]),
+         tf.zeros_like(x), -1e10 * tf.ones(x_shape[0])], back_prop=False)
+
+    return adv_x
+
+  def parse_params(
+      self,
+      eps=0.3,
+      num_samples=100,
+      num_batches=100,
+      ord=np.inf,  # pylint: disable=redefined-builtin
+      y=None,
+      y_target=None,
+      clip_min=None,
+      clip_max=None,
+      **kwargs):
+    self.num_samples = num_samples
+    self.num_batches = num_batches
+    return super(RandomAttack, self).parse_params(eps, ord, y, y_target,
+                                                  clip_min, clip_max, **kwargs)
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 7123cba3f..d4c28dcf1 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1611,11 +1611,11 @@ def host_call_fn(**kwargs):
         # We need to use tf.contrib.summary in order to feed the `step`.
         for name, value in sorted(six.iteritems(kwargs)):
           if name.startswith("ScalarSummary"):
-            name = name.strip("ScalarSummary")
+            name = name[len("ScalarSummary"):]
             tf.contrib.summary.scalar(
                 name, tf.reduce_mean(tf.to_float(value)), step=gs)
           elif name.startswith("ImageSummary"):
-            name = name.strip("ImageSummary")
+            name = name[len("ImageSummary"):]
             tf.contrib.summary.image(name, value, step=gs)
 
         return tf.contrib.summary.all_summary_ops()

From 995b5b868bce363a3a53ca3cc6e721dcdab97e78 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 10 Aug 2018 14:11:37 -0700
Subject: [PATCH 0539/2720] Add xla_compile flag as placeholder for future xla
 compile support.

PiperOrigin-RevId: 208265440
---
 tensor2tensor/bin/t2t_trainer.py   | 2 ++
 tensor2tensor/utils/trainer_lib.py | 7 +++++--
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index e4caf6cf1..c3e699c17 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -50,6 +50,7 @@
 flags.DEFINE_bool("use_tpu", False, "Whether to use TPU.")
 flags.DEFINE_bool("use_tpu_estimator", False, "Whether to use TPUEstimator. "
                   "This is always enabled when use_tpu is True.")
+flags.DEFINE_bool("xla_compile", False, "Whether to use XLA to compile graph.")
 flags.DEFINE_integer("tpu_infeed_sleep_secs", None,
                      "How long to sleep the infeed thread.")
 flags.DEFINE_bool("generate_data", False, "Generate data before training?")
@@ -177,6 +178,7 @@ def create_experiment_fn(**kwargs):
       eval_early_stopping_metric_minimize,
       use_tpu=FLAGS.use_tpu,
       use_tpu_estimator=FLAGS.use_tpu_estimator,
+      use_xla=FLAGS.xla_compile,
       **kwargs)
 
 
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 9b4762f34..34a2fa165 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -210,7 +210,8 @@ def create_estimator(model_name,
                      schedule="train_and_evaluate",
                      decode_hparams=None,
                      use_tpu=False,
-                     use_tpu_estimator=False):
+                     use_tpu_estimator=False,
+                     use_xla=False):
   """Create a T2T Estimator."""
   model_fn = t2t_model.T2TModel.make_estimator_model_fn(
       model_name, hparams, decode_hparams=decode_hparams)
@@ -417,6 +418,7 @@ def create_experiment(
     autotune=False,
     use_tpu=False,
     use_tpu_estimator=False,
+    use_xla=False,
     additional_train_hooks=None,
     additional_eval_hooks=None):
   """Create Experiment."""
@@ -436,7 +438,8 @@ def create_experiment(
       schedule=schedule,
       decode_hparams=decode_hparams,
       use_tpu=use_tpu,
-      use_tpu_estimator=use_tpu_estimator)
+      use_tpu_estimator=use_tpu_estimator,
+      use_xla=use_xla)
 
   # Input fns from Problem
   problem = hparams.problem

From f09597a7ea413b0e47dde9b459c7ac88d9290556 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 10 Aug 2018 14:36:52 -0700
Subject: [PATCH 0540/2720] Add enable_summary flag to control whether summary
 ops are generated

PiperOrigin-RevId: 208269839
---
 tensor2tensor/layers/common_layers.py |  5 +++--
 tensor2tensor/utils/flags.py          |  1 +
 tensor2tensor/utils/optimize.py       | 18 +++++++++++-------
 tensor2tensor/utils/t2t_model.py      | 18 ++++++++++++------
 4 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 9f332ab48..e4b11a60b 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3448,13 +3448,14 @@ def should_generate_summaries():
   Returns:
     a boolean
   """
-  if "while/" in tf.contrib.framework.get_name_scope():
+  name_scope = tf.contrib.framework.get_name_scope()
+  if name_scope and "while/" in name_scope:
     # Summaries don't work well within tf.while_loop()
     return False
   if tf.get_variable_scope().reuse:
     # Avoid generating separate summaries for different data shards
     return False
-  return True
+  return getattr(tf.flags.FLAGS, "enable_summaries", True)
 
 
 def reshape_like(a, b):
diff --git a/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index 3ae12988a..51d9d1a86 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -73,6 +73,7 @@
                      "How many recent checkpoints to keep.")
 flags.DEFINE_bool("enable_graph_rewriter", False,
                   "Enable graph optimizations that are not on by default.")
+flags.DEFINE_bool("enable_summaries", True, "Enable creating summary ops.")
 flags.DEFINE_integer("keep_checkpoint_every_n_hours", 10000,
                      "Number of hours between each checkpoint to be saved. "
                      "The default value 10,000 hours effectively disables it.")
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 7249a76d8..b2ac2a325 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -44,11 +44,14 @@ def optimize(loss, learning_rate, hparams, use_tpu=False):
   if use_tpu:
     opt = tf.contrib.tpu.CrossShardOptimizer(opt)
 
-  tf.summary.scalar("learning_rate", learning_rate)
-  opt_summaries = ["loss"]
-  if hparams.summarize_grads:
-    tf.logging.info("Summarizing gradients")
-    opt_summaries.extend(["gradients", "gradient_norm", "global_gradient_norm"])
+  opt_summaries = []
+  if common_layers.should_generate_summaries():
+    tf.summary.scalar("learning_rate", learning_rate)
+    opt_summaries.append("loss")
+    if hparams.summarize_grads:
+      tf.logging.info("Summarizing gradients")
+      opt_summaries.extend(
+          ["gradients", "gradient_norm", "global_gradient_norm"])
 
   if hparams.clip_grad_norm:
     tf.logging.info("Clipping gradients, norm: %0.5f", hparams.clip_grad_norm)
@@ -136,7 +139,7 @@ def weight_decay_and_noise(loss, hparams, learning_rate, var_list=None):
   noise_vars = [v for v in var_list if "/body/" in v.name]
 
   weight_decay_loss = weight_decay(hparams.weight_decay, decay_vars)
-  if hparams.weight_decay:
+  if hparams.weight_decay and common_layers.should_generate_summaries():
     tf.summary.scalar("losses/weight_decay", weight_decay_loss)
   weight_noise_ops = weight_noise(hparams.weight_noise, learning_rate,
                                   noise_vars)
@@ -161,7 +164,8 @@ def weight_noise(noise_rate, learning_rate, var_list):
   for v in var_list:
     with tf.device(v._ref().device):  # pylint: disable=protected-access
       scale = noise_rate * learning_rate * 0.001
-      tf.summary.scalar("weight_noise_scale", scale)
+      if common_layers.should_generate_summaries():
+        tf.summary.scalar("weight_noise_scale", scale)
       noise = tf.truncated_normal(v.shape) * scale
       noise_op = v.assign_add(noise)
       noise_ops.append(noise_op)
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index d4c28dcf1..acbbe7aaf 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -431,9 +431,10 @@ def loss(self, logits, features):
         losses[k] = self._loss_single(v, target_modality[k], features[k])
 
         n, d = losses[k]
-        tf.summary.scalar(k + "_loss", n / d)
-        tf.summary.scalar(k + "_loss_num", n)
-        tf.summary.scalar(k + "_loss_den", d)
+        if common_layers.should_generate_summaries():
+          tf.summary.scalar(k + "_loss", n / d)
+          tf.summary.scalar(k + "_loss_num", n)
+          tf.summary.scalar(k + "_loss_den", d)
 
       return tf.add_n([n / d for n, d in losses.values()])
     else:
@@ -1236,9 +1237,10 @@ def estimator_model_fn(cls,
       return logits
 
     # Summarize losses
-    with tf.name_scope("losses"):
-      for loss_name, loss_val in sorted(losses_dict.items()):
-        tf.summary.scalar(loss_name, loss_val)
+    if common_layers.should_generate_summaries():
+      with tf.name_scope("losses"):
+        for loss_name, loss_val in sorted(losses_dict.items()):
+          tf.summary.scalar(loss_name, loss_val)
 
     # Accumulate losses
     loss = sum(losses_dict[key] for key in sorted(losses_dict.keys()))
@@ -1722,6 +1724,10 @@ def average_sharded_losses(sharded_losses):
 
 
 def summarize_features(features, num_shards=1):
+  """Generate summaries for features."""
+  if not common_layers.should_generate_summaries():
+    return
+
   with tf.name_scope("input_stats"):
     for (k, v) in sorted(six.iteritems(features)):
       if isinstance(v, tf.Tensor) and v.get_shape().ndims > 1:

From 468c0c95b95fafbe478308d220f9a2a328acfe9f Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 10 Aug 2018 15:04:34 -0700
Subject: [PATCH 0541/2720] Fix broken registry help_string() function. Added
 tests.

PiperOrigin-RevId: 208274407
---
 tensor2tensor/utils/registry.py      | 4 ++--
 tensor2tensor/utils/registry_test.py | 8 ++++++++
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 1b27ad118..14600d216 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -646,7 +646,7 @@ def help_string():
   Pruning Strategies:
 %s
 """
-  m, hp, rhp, mod, probs, atks, ap = [
+  m, hp, rhp, mod, probs, atks, ap, pp, ps = [
       display_list_by_prefix(entries, starting_spaces=4) for entries in [
           list_models(),
           list_hparams(),
@@ -659,4 +659,4 @@ def help_string():
           list_pruning_strategies(),
       ]
   ]
-  return help_str % (m, hp, rhp, mod, probs, atks, ap)
+  return help_str % (m, hp, rhp, mod, probs, atks, ap, pp, ps)
diff --git a/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index aa57e3482..7324f1712 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -273,5 +273,13 @@ class MyClassLabelModality(modality.Modality):
     self.assertSetEqual(set(registry.list_modalities()), set(expected))
 
 
+class RegistryTest(tf.test.TestCase):
+  """ Test class for common functions."""
+
+  def testRegistryHelp(self):
+    help_str = registry.help_string()
+    self.assertIsNotNone(help_str)
+    self.assertGreater(len(help_str), 0)
+
 if __name__ == "__main__":
   tf.test.main()

From 85c6c5a774c70d100f9fc79254ab641092e0f24e Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 10 Aug 2018 15:08:26 -0700
Subject: [PATCH 0542/2720] Allow sequential optimization of the discriminator
 and generator via hparams.gan_optimization. This is done using a tf.cond op
 based on the global step.

PiperOrigin-RevId: 208275102
---
 tensor2tensor/models/research/next_frame_params.py |  1 +
 tensor2tensor/models/research/next_frame_savp.py   | 12 +++++++++++-
 tensor2tensor/models/research/next_frame_sv2p.py   | 13 +++----------
 tensor2tensor/models/research/next_frame_test.py   |  5 ++++-
 4 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index ebcf46fb0..b8769671b 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -120,6 +120,7 @@ def next_frame_savp():
   hparams.add_hparam("use_spectral_norm", True)
   hparams.add_hparam("gan_loss", "cross_entropy")
   hparams.add_hparam("gan_loss_multiplier", 0.01)
+  hparams.add_hparam("gan_optimization", "joint")
   hparams.target_modality = "video:l1raw"
   hparams.input_modalities = "inputs:video:l1raw"
   hparams.latent_loss_multiplier_schedule = "linear_anneal"
diff --git a/tensor2tensor/models/research/next_frame_savp.py b/tensor2tensor/models/research/next_frame_savp.py
index ad4136ea5..7f18afcf4 100644
--- a/tensor2tensor/models/research/next_frame_savp.py
+++ b/tensor2tensor/models/research/next_frame_savp.py
@@ -248,8 +248,15 @@ def get_gan_loss(self, true_frames, gen_frames):
       gan_g_loss_pos_d, gan_g_loss_neg_d = self.g_step(
           gen_frames, fake_logits_stop)
     gan_g_loss = gan_g_loss_pos_d + gan_g_loss_neg_d
-    gan_loss = gan_g_loss + gan_d_loss
     tf.summary.scalar("gan_loss", gan_g_loss_pos_d + gan_d_loss)
+
+    if self.hparams.gan_optimization == "joint":
+      gan_loss = gan_g_loss + gan_d_loss
+    else:
+      curr_step = tf.train.get_or_create_global_step()
+      gan_loss = tf.cond(
+          tf.logical_not(curr_step % 2 == 0), lambda: gan_g_loss,
+          lambda: gan_d_loss)
     return self.hparams.gan_loss_multiplier * gan_loss
 
   def get_extra_loss(self, latent_means=None, latent_stds=None,
@@ -331,6 +338,9 @@ def construct_model(self, images, actions, rewards):
       raise ValueError("VAE + GAN variant not implemented")
     if not self.hparams.use_vae and not self.hparams.use_gan:
       raise ValueError("Set at least one of use_vae or use_gan to be True")
+    if self.hparams.gan_optimization not in ["joint", "sequential"]:
+      raise ValueError("self.hparams.gan_optimization should be either joint "
+                       "or sequential got %s" % self.hparams.gan_optimization)
 
     images = tf.unstack(images, axis=0)
     actions = tf.unstack(actions, axis=0)
diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 5f08980fe..4acc3b2ee 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -55,16 +55,9 @@ def get_gaussian_latent(self, latent_mean, latent_std):
     latent = latent_mean + tf.exp(latent_std / 2.0) * latent
     return latent
 
-  def get_iteration_num(self):
-    step_num = tf.train.get_global_step()
-    # TODO(lukaszkaiser): what should it be if it"s undefined?
-    if step_num is None:
-      step_num = 1000000
-    return step_num
-
   def get_beta(self):
     """Get KL multiplier (beta) based on the schedule."""
-    step_num = self.get_iteration_num()
+    step_num = tf.train.get_or_create_global_step()
     schedule = self.hparams.latent_loss_multiplier_schedule
     second_stage = (self.hparams.num_iterations_1st_stage +
                     self.hparams.num_iterations_2nd_stage)
@@ -103,7 +96,7 @@ def anneal_loss(step_num):
   def get_scheduled_sample_func(self, batch_size):
     """Creates a function for scheduled sampling based on given hparams."""
     with tf.variable_scope("scheduled_sampling_func", reuse=False):
-      iter_num = self.get_iteration_num()
+      iter_num = tf.train.get_or_create_global_step()
       if self.hparams.scheduled_sampling_mode == "prob":
         decay_steps = self.hparams.scheduled_sampling_decay_steps
         probability = tf.train.polynomial_decay(
@@ -312,7 +305,7 @@ def construct_latent_tower(self, images):
         return tf.zeros_like(mean), tf.zeros_like(std)
 
       # No latent in the first phase
-      iter_num = self.get_iteration_num()
+      iter_num = tf.train.get_or_create_global_step()
       ret_mean, ret_std = tf.cond(
           iter_num < self.hparams.num_iterations_1st_stage,
           lambda: (tf.zeros_like(mean), tf.zeros_like(std)),
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index 309f51ee4..f289758c8 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -102,7 +102,7 @@ def get_tensor_shape(tensor):
 class NextFrameTest(tf.test.TestCase):
 
   def RunModel(self, model, hparams, features):
-    with self.Session() as session:
+    with tf.Session() as session:
       model = model(
           hparams, tf.estimator.ModeKeys.TRAIN)
       logits, _ = model(features)
@@ -237,6 +237,9 @@ def testStochasticSavpGAN(self):
     hparams.use_vae = False
     self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
 
+    hparams.gan_optimization = "sequential"
+    self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
+
   def testStochasticInvalidVAEGANCombinations(self):
     hparams = next_frame_params.next_frame_savp()
     for use_vae, use_gan in [[True, True], [False, False]]:

From 2c060839cae831768575df6ffee591980bd53c68 Mon Sep 17 00:00:00 2001
From: Youlong Cheng <ylc@google.com>
Date: Fri, 10 Aug 2018 15:11:50 -0700
Subject: [PATCH 0543/2720] Pylint against tf.nightly version

PiperOrigin-RevId: 208275662
---
 .travis.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.travis.yml b/.travis.yml
index e89d0181c..1d25a3440 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -98,7 +98,7 @@ script:
 
   # Do some things only on Python 2 and the latest TF version
   # Each should be in a separate block to get proper errors on Travis.
-  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
+  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "tf-nightly"  ]]; then
         pylint -j 2 tensor2tensor;
     fi
   - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then

From 2621d15596cad8d62ace9bae1d10488cc03b39b4 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Fri, 10 Aug 2018 15:40:18 -0700
Subject: [PATCH 0544/2720] Fixing indices of rewards and observation in gym
 problems.

PiperOrigin-RevId: 208280025
---
 tensor2tensor/data_generators/gym_problems.py |  8 ++++++--
 tensor2tensor/rl/collect.py                   | 18 ++++++------------
 2 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 09f981661..214bb724a 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -149,6 +149,8 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
       memory_index = 0
       memory = None
       pieces_generated = 0
+      prev_reward = 0
+      prev_done = False
 
       # TODO(piotrmilos): self.settable_eval_phase possibly violates sematics
       # of VideoProblem
@@ -171,8 +173,8 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
             "image/height": [self.frame_height],
             "image/width": [self.frame_width],
             "action": [int(action)],
-            "done": [int(done)],
-            "reward": [int(reward - self.min_reward)]
+            "done": [int(prev_done)],
+            "reward": [int(prev_reward - self.min_reward)]
         }
 
         if debug_image is not None:
@@ -183,6 +185,8 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
         if done and self.settable_eval_phase:
           return
 
+        prev_done, prev_reward = done, reward
+
         pieces_generated += 1
         frame_counter += 1
         if done:
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 3c7441a21..f32fca858 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -48,7 +48,6 @@ class _MemoryWrapper(WrapperBase):
 
   def __init__(self, batch_env):
     super(_MemoryWrapper, self).__init__(batch_env)
-    self._is_simple = False  # TODO(lukaszkaiser): why do we need it? mbz?
     infinity = 10000000
     meta_data = list(zip(*_rollout_metadata(batch_env)))
     # In memory wrapper we do not collect pdfs neither value_function
@@ -70,17 +69,12 @@ def simulate(self, action):
 
     reward, done = self._batch_env.simulate(action)
 
-    if self._is_simple:
-      with tf.control_dependencies([reward, done]):
-        enqueue_op = self.speculum.enqueue(
-            [self._batch_env.observ, reward, done, action])
-    else:
-      with tf.control_dependencies([reward, done]):
-        enqueue_op = self.speculum.enqueue(
-            [self._observ.read_value(), reward, done, action])
-
-      with tf.control_dependencies([enqueue_op]):
-        assign = self._observ.assign(self._batch_env.observ)
+    with tf.control_dependencies([reward, done]):
+      enqueue_op = self.speculum.enqueue(
+          [self._observ.read_value(), reward, done, action])
+
+    with tf.control_dependencies([enqueue_op]):
+      assign = self._observ.assign(self._batch_env.observ)
 
     with tf.control_dependencies([assign]):
       return tf.identity(reward), tf.identity(done)

From 32409299527391e39c3a05c101a78eb298fdb95d Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 10 Aug 2018 16:44:35 -0700
Subject: [PATCH 0545/2720] Part 1 Glow in Tensor2Tensor. Add reversible ops
 necessary for one step of flow.

PiperOrigin-RevId: 208289539
---
 tensor2tensor/models/research/glow_ops.py     | 374 ++++++++++++++++++
 .../models/research/glow_ops_test.py          | 126 ++++++
 2 files changed, 500 insertions(+)
 create mode 100644 tensor2tensor/models/research/glow_ops.py
 create mode 100644 tensor2tensor/models/research/glow_ops_test.py

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
new file mode 100644
index 000000000..f240fee94
--- /dev/null
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -0,0 +1,374 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Glow generative model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import scipy
+from tensor2tensor.layers import common_layers
+import tensorflow as tf
+
+arg_scope = tf.contrib.framework.arg_scope
+add_arg_scope = tf.contrib.framework.add_arg_scope
+
+
+def default_initializer(std=0.05):
+  return tf.random_normal_initializer(0., std)
+
+
+@add_arg_scope
+def get_variable_ddi(name, shape, initial_value, dtype=tf.float32, init=False,
+                     trainable=True):
+  """Wrapper for data-dependent initialization."""
+  w = tf.get_variable(name, shape, dtype, None, trainable=trainable)
+  if init:
+    w = w.assign(initial_value)
+    with tf.control_dependencies([w]):
+      return w
+  return w
+
+
+@add_arg_scope
+def actnorm(name, x, logscale_factor=3., reverse=False, init=False,
+            trainable=True):
+  """x_{ij} = s x x_{ij} + b. Per-channel scaling and bias.
+
+  If init is set to True, the scaling and bias are initialized such
+  that the mean and variance of the output activations of the first minibatch
+  are zero and one respectively.
+
+  Args:
+    name: variable scope.
+    x: input
+    logscale_factor: Used in actnorm_scale. Optimizes f(ls*s') instead of f(s)
+                     where s' = s / ls. Helps in faster convergence.
+    reverse: forward or reverse operation.
+    init: Whether or not to do data-dependent initialization.
+    trainable:
+
+  Returns:
+    x: output after adding bias and scaling.
+    objective: log(sum(s))
+  """
+  var_arg_scope = arg_scope([get_variable_ddi], trainable=trainable)
+  var_scope = tf.variable_scope(name)
+  with var_scope, var_arg_scope:
+    if not reverse:
+      x = actnorm_center(name + "_center", x, reverse, init=init)
+      x, objective = actnorm_scale(
+          name + "_scale", x, logscale_factor=logscale_factor,
+          reverse=reverse, init=init)
+    else:
+      x, objective = actnorm_scale(
+          name + "_scale", x, logscale_factor=logscale_factor,
+          reverse=reverse, init=init)
+      x = actnorm_center(name + "_center", x, reverse)
+    return x, objective
+
+
+@add_arg_scope
+def actnorm_center(name, x, reverse=False, init=False):
+  """Add a bias to x.
+
+  Initialize such that the output of the first minibatch is zero centered
+  per channel.
+
+  Args:
+    name: scope
+    x: 2-D or 4-D Tensor.
+    reverse: Forward or backward operation.
+    init: data-dependent initialization.
+
+  Returns:
+    x_center: (x + b), if reverse is True and (x - b) otherwise.
+  """
+  shape = common_layers.shape_list(x)
+  with tf.variable_scope(name):
+    assert len(shape) == 2 or len(shape) == 4
+    if len(shape) == 2:
+      x_mean = tf.reduce_mean(x, [0], keepdims=True)
+      b = get_variable_ddi(
+          "b", (1, shape[1]), initial_value=-x_mean, init=init)
+    elif len(shape) == 4:
+      x_mean = tf.reduce_mean(x, [0, 1, 2], keepdims=True)
+      b = get_variable_ddi(
+          "b", (1, 1, 1, shape[3]), initial_value=-x_mean, init=init)
+
+    if not reverse:
+      x += b
+    else:
+      x -= b
+    return x
+
+
+@add_arg_scope
+def actnorm_scale(name, x, logscale_factor=3., reverse=False, init=False):
+  """Per-channel scaling of x."""
+  x_shape = common_layers.shape_list(x)
+  with tf.variable_scope(name):
+
+    # Variance initialization logic.
+    assert len(x_shape) == 2 or len(x_shape) == 4
+    if len(x_shape) == 2:
+      x_var = tf.reduce_mean(x**2, [0], keepdims=True)
+      logdet_factor = 1
+      var_shape = (1, x_shape[1])
+    elif len(x_shape) == 4:
+      x_var = tf.reduce_mean(x**2, [0, 1, 2], keepdims=True)
+      logdet_factor = x_shape[1]*x_shape[2]
+      var_shape = (1, 1, 1, x_shape[3])
+
+    init_value = tf.log(1.0 / (tf.sqrt(x_var) + 1e-6)) / logscale_factor
+    logs = get_variable_ddi(
+        "logs", var_shape, initial_value=init_value, init=init)
+    logs = logs * logscale_factor
+
+    # Function and reverse function.
+    if not reverse:
+      x = x * tf.exp(logs)
+    else:
+      x = x * tf.exp(-logs)
+
+    # Objective calculation, h * w * sum(log|s|)
+    dlogdet = tf.reduce_sum(logs) * logdet_factor
+    if reverse:
+      dlogdet *= -1
+    return x, dlogdet
+
+
+@add_arg_scope
+def invertible_1x1_conv(name, x, reverse=False):
+  """1X1 convolution on x.
+
+  The 1X1 convolution is parametrized as P*L*(U + sign(s)*exp(log(s))) where
+  1. P is a permutation matrix.
+  2. L is a lower triangular matrix with diagonal entries unity.
+  3. U is a upper triangular matrix where the diagonal entries zero.
+  4. s is a vector.
+
+  sign(s) and P are fixed and the remaining are optimized. P, L, U and s are
+  initialized by the PLU decomposition of a random rotation matrix.
+
+  Args:
+    name: scope
+    x: Input Tensor.
+    reverse: whether the pass is from z -> x or x -> z.
+
+  Returns:
+    x_conv: x after a 1X1 convolution is applied on x.
+    objective: sum(log(s))
+  """
+  _, height, width, channels = common_layers.shape_list(x)
+  w_shape = [channels, channels]
+
+  # Random rotation-matrix Q
+  random_matrix = np.random.rand(channels, channels)
+  np_w = scipy.linalg.qr(random_matrix)[0].astype("float32")
+
+  # Initialize P,L,U and s from the LU decomposition of a random rotation matrix
+  np_p, np_l, np_u = scipy.linalg.lu(np_w)
+  np_s = np.diag(np_u)
+  np_sign_s = np.sign(np_s)
+  np_log_s = np.log(np.abs(np_s))
+  np_u = np.triu(np_u, k=1)
+
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    p = tf.get_variable("P", initializer=np_p, trainable=False)
+    l = tf.get_variable("L", initializer=np_l)
+    sign_s = tf.get_variable(
+        "sign_S", initializer=np_sign_s, trainable=False)
+    log_s = tf.get_variable("log_S", initializer=np_log_s)
+    u = tf.get_variable("U", initializer=np_u)
+
+    # W = P * L * (U + sign_s * exp(log_s))
+    l_mask = np.tril(np.ones([channels, channels], dtype=np.float32), -1)
+    l = l * l_mask + tf.eye(channels, channels)
+    u = u * np.transpose(l_mask) + tf.diag(sign_s * tf.exp(log_s))
+    w = tf.matmul(p, tf.matmul(l, u))
+
+    objective = tf.reduce_sum(log_s) * height * width
+    if not reverse:
+      w = tf.reshape(w, [1, 1] + w_shape)
+      x = tf.nn.conv2d(x, w, [1, 1, 1, 1], "SAME", data_format="NHWC")
+    else:
+      u_inv = tf.matrix_inverse(u)
+      l_inv = tf.matrix_inverse(l)
+      p_inv = tf.matrix_inverse(p)
+      w_inv = tf.matmul(u_inv, tf.matmul(l_inv, p_inv))
+      w_inv = tf.reshape(w_inv, [1, 1]+w_shape)
+      x = tf.nn.conv2d(
+          x, w_inv, [1, 1, 1, 1], "SAME", data_format="NHWC")
+      objective *= -1
+  return x, objective
+
+
+def add_edge_bias(x, filter_size):
+  """Pad x and concatenates an edge bias across the depth of x.
+
+  The edge bias can be thought of as a binary feature which is unity when
+  the filter is being convolved over an edge and zero otherwise.
+
+  Args:
+    x: Input tensor, shape (NHWC)
+    filter_size: filter_size to determine padding.
+  Returns:
+    x_pad: Input tensor, shape (NHW(c+1))
+  """
+  x_shape = common_layers.shape_list(x)
+  if filter_size[0] == 1 and filter_size[1] == 1:
+    return x
+  a = (filter_size[0] - 1) // 2  # vertical padding size
+  b = (filter_size[1] - 1) // 2  # horizontal padding size
+  padding = [[0, 0], [a, a], [b, b], [0, 0]]
+  x_bias = tf.zeros(x_shape[:-1] + [1])
+
+  x = tf.pad(x, padding)
+  x_pad = tf.pad(x_bias, padding, constant_values=1)
+  return tf.concat([x, x_pad], axis=3)
+
+
+@add_arg_scope
+def conv2d(name, x, output_channels, filter_size=None, stride=None,
+           logscale_factor=3.0, init=True, apply_actnorm=True,
+           conv_init="default"):
+  """conv2d layer with edge bias padding and optional actnorm.
+
+  Args:
+    name: variable scope.
+    x: 4-D Tensor of shape (NHWC)
+    output_channels: Number of output channels.
+    filter_size:
+    stride:
+    logscale_factor: see actnorm for parameter meaning.
+    init: Whether to apply data-dependent initialization Valid only if
+          apply_actnorm is set to True.
+    apply_actnorm: if apply_actnorm the activations of the first minibatch
+                   have zero mean and unit variance. Else, there is no scaling
+                   applied.
+    conv_init: default or zeros. default is a normal distribution with 0.05 std.
+  Returns:
+    x: actnorm(conv2d(x))
+  Raises:
+    ValueError: if init is set to "zeros" and apply_actnorm is set to True.
+  """
+  if init == "zeros" and apply_actnorm:
+    raise ValueError("apply_actnorm is unstable when init is set to zeros.")
+
+  if filter_size is None:
+    filter_size = [3, 3]
+  if stride is None:
+    stride = [1, 1]
+
+  x = add_edge_bias(x, filter_size=filter_size)
+  _, _, _, in_channels = common_layers.shape_list(x)
+
+  filter_shape = filter_size + [in_channels, output_channels]
+  stride_shape = [1, 1] + stride
+
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+
+    if conv_init == "default":
+      initializer = default_initializer()
+    elif conv_init == "zeros":
+      initializer = tf.zeros_initializer()
+
+    w = tf.get_variable("W", filter_shape, tf.float32,
+                        initializer=initializer)
+    x = tf.nn.conv2d(x, w, stride_shape, padding="VALID", data_format="NHWC")
+
+    if apply_actnorm:
+      x, _ = actnorm("actnorm", x, logscale_factor=logscale_factor, init=init,
+                     trainable=True)
+    else:
+      x += tf.get_variable("b", [1, 1, 1, output_channels],
+                           initializer=tf.zeros_initializer())
+      logs = tf.get_variable("logs", [1, output_channels],
+                             initializer=tf.zeros_initializer())
+      x *= tf.exp(logs * logscale_factor)
+    return x
+
+
+@add_arg_scope
+def nn(name, x, mid_channels, output_channels):
+  """3-layer conv2d.
+
+  Args:
+    name:
+    x:
+    mid_channels: Number of output channels of the first layer.
+    output_channels: Number of output channels.
+
+  Returns:
+    output:
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+
+    # Edge Padding + conv2d + actnorm + relu:
+    # [output: 512 channels]
+    x = conv2d("1_1", x, output_channels=mid_channels, filter_size=[3, 3],
+               stride=[1, 1])
+    x = tf.nn.relu(x)
+
+    # Padding + conv2d + actnorm + relu
+    # [input, output: 512 channels]
+    x = conv2d("1_2", x, output_channels=mid_channels, filter_size=[1, 1],
+               stride=[1, 1])
+    x = tf.nn.relu(x)
+
+    # Final layer.
+    x = conv2d("zeros", x, filter_size=[3, 3], stride=[1, 1],
+               output_channels=output_channels, apply_actnorm=False,
+               conv_init="zeros")
+  return x
+
+
+@add_arg_scope
+def affine_coupling(name, x, mid_channels, reverse=False):
+  """Reversible affine coupling layer.
+
+  Args:
+    name:
+    x:
+    mid_channels: intermediate
+    reverse: Forward or reverse operation.
+  Returns:
+    output:
+    objective:
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    x_shape = common_layers.shape_list(x)
+    x1, x2 = tf.split(x, num_or_size_splits=2, axis=-1)
+
+    # scale, shift = NN(x1)
+    # If reverse:
+    # z2 = scale * (x2 + shift)
+    # Else:
+    # z2 = (x2 / scale) - shift
+    z1 = x1
+    log_scale_and_shift = nn("nn", x1, mid_channels, x_shape[-1])
+    shift = log_scale_and_shift[:, :, :, 0::2]
+    scale = tf.nn.sigmoid(log_scale_and_shift[:, :, :, 1::2] + 2.0)
+    if not reverse:
+      z2 = (x2 + shift) * scale
+    else:
+      z2 = x2 / scale - shift
+
+    objective = tf.reduce_sum(tf.log(scale), axis=[1, 2, 3])
+    if reverse:
+      objective *= -1
+    return tf.concat([z1, z2], axis=3), objective
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
new file mode 100644
index 000000000..960ec7a7d
--- /dev/null
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -0,0 +1,126 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 google3.third_party.py.tensor2tensor.models.research.glow_ops."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.models.research import glow_ops
+import tensorflow as tf
+
+
+class GlowOpsTest(tf.test.TestCase):
+
+  def test_get_variable_ddi(self):
+    with tf.Graph().as_default():
+      x_t = tf.random_normal((5, 5))
+      ddi = glow_ops.get_variable_ddi(
+          "x", (5, 5), x_t, init=True)
+      with tf.Session() as session:
+        diff = ddi - x_t
+        self.assertTrue(np.allclose(session.run(diff), 0.0))
+
+  def test_actnorm(self):
+    """Test that actnorm provides activations with zero channel-mean."""
+    with tf.Graph().as_default():
+      x_t = tf.random_normal((16, 32, 32, 3), mean=50.0, stddev=2.0)
+      x_act = glow_ops.actnorm("actnorm", x_t, init=True)
+      with tf.Session() as session:
+        x_act_np, _ = session.run(x_act)
+        channel_mean = np.mean(x_act_np, axis=(0, 1, 2))
+        channel_var = np.var(x_act_np, axis=(0, 1, 2))
+        self.assertTrue(np.allclose(channel_mean, 0.0, atol=1e-3))
+        self.assertTrue(np.allclose(channel_var, 1.0, atol=1e-3))
+
+  def test_invertible_conv(self):
+    with tf.Graph().as_default():
+      x_t = tf.random_uniform(shape=(16, 32, 32, 3))
+      activation, _ = glow_ops.invertible_1x1_conv("inv", x_t, reverse=False)
+      inv_activation, _ = glow_ops.invertible_1x1_conv(
+          "inv", activation, reverse=True)
+      with tf.Session() as session:
+        session.run(tf.global_variables_initializer())
+        diff = session.run(x_t - inv_activation)
+
+        # Test reversibility.
+        self.assertTrue(np.allclose(diff, 0.0, atol=1e-5))
+
+  def test_add_edge_bias(self):
+    with tf.Graph().as_default():
+      x = tf.random_uniform(shape=(16, 32, 32, 3))
+      x_pad = glow_ops.add_edge_bias(x, [3, 3])
+      with tf.Session() as session:
+        x_pad_np = session.run(x_pad)
+
+        # Test expected output shape.
+        self.assertEqual(x_pad_np.shape, (16, 34, 34, 4))
+
+  def test_conv2d(self):
+    with tf.Graph().as_default():
+      x = 10.0 * tf.random_uniform(shape=(16, 5, 5, 32))
+      actnorm_conv2d = glow_ops.conv2d(
+          "actnorm_conv2d", x, output_channels=64, init=True,
+          apply_actnorm=True)
+      actnorm_zeros2d = glow_ops.conv2d(
+          "actnorm_zeros2d", x, output_channels=64, init=True,
+          apply_actnorm=False)
+      with tf.Session() as session:
+        session.run(tf.global_variables_initializer())
+
+        # test if apply_actnorm is set to True, the first minibatch has
+        # zero mean and unit variance.
+        actnorm_np, zeros_np = session.run([actnorm_conv2d, actnorm_zeros2d])
+        self.assertEqual(actnorm_np.shape, (16, 5, 5, 64))
+        mean = np.mean(actnorm_np, axis=(0, 1, 2))
+        var = np.var(actnorm_np, axis=(0, 1, 2))
+        self.assertTrue(np.allclose(mean, 0.0, atol=1e-5))
+        self.assertTrue(np.allclose(var, 1.0, atol=1e-5))
+
+        # test shape in case apply_actnorm is set to False,
+        self.assertEqual(zeros_np.shape, (16, 5, 5, 64))
+
+  def test_nn(self):
+    """Test output shape."""
+    with tf.Graph().as_default():
+      x = 10.0 * tf.random_uniform(shape=(16, 5, 5, 32))
+      nn = glow_ops.nn("nn", x, 512, 64)
+
+      with tf.Session() as session:
+        session.run(tf.global_variables_initializer())
+        nn_np = session.run(nn)
+        self.assertEqual(nn_np.shape, (16, 5, 5, 64))
+
+        # Initialized with zeros.
+        self.assertTrue(np.allclose(nn_np, 0.0))
+
+  def test_affine_coupling(self):
+    """Test affine coupling reversibility."""
+    with tf.Graph().as_default():
+      rng = np.random.RandomState(0)
+      x = np.asarray(rng.rand(16, 3, 3, 32), dtype=np.float32)
+      x_t = tf.convert_to_tensor(x)
+      x_inv, _ = glow_ops.affine_coupling("affine", x_t, 512)
+      x_inv_inv, _ = glow_ops.affine_coupling(
+          "affine", x_inv, 512, reverse=True)
+
+      with tf.Session() as session:
+        session.run(tf.global_variables_initializer())
+        diff = tf.reduce_max(tf.abs(x - x_inv_inv))
+        self.assertTrue(np.allclose(session.run(diff), 0.0, atol=1e-5))
+
+if __name__ == "__main__":
+  tf.test.main()

From 30fe32c0a4f0a5333dd4e921f538dac93c03e099 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 10 Aug 2018 16:48:15 -0700
Subject: [PATCH 0546/2720] Use tensorflow's bucket_by_sequene_length instead
 of internal one.

PiperOrigin-RevId: 208289987
---
 tensor2tensor/data_generators/problem.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 061112687..107409c5c 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -894,9 +894,10 @@ def define_shapes(example):
           # Here  batch_size really means examples per datashard.
           batching_scheme["batch_sizes"] = [hparams.batch_size]
           batching_scheme["boundaries"] = []
-        dataset = data_reader.bucket_by_sequence_length(
-            dataset, data_reader.example_length, batching_scheme["boundaries"],
-            batching_scheme["batch_sizes"])
+        dataset = dataset.apply(
+            tf.contrib.data.bucket_by_sequence_length(
+                data_reader.example_length, batching_scheme["boundaries"],
+                batching_scheme["batch_sizes"]))
 
         if not is_training:
           batch_multiple = shard_multiplier

From e6caaf2afe0e2dda58aa43483ea2d214dcdc4be1 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 10 Aug 2018 17:30:24 -0700
Subject: [PATCH 0547/2720] Fix description

PiperOrigin-RevId: 208295334
---
 tensor2tensor/models/research/glow_ops_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 960ec7a7d..0e42f71ea 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -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.
-"""Tests for google3.third_party.py.tensor2tensor.models.research.glow_ops."""
+"""Tests for tensor2tensor.models.research.glow_ops."""
 
 from __future__ import absolute_import
 from __future__ import division

From ac332b3e9a404514be5b2b2b2a16b7498ceee9f9 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 10 Aug 2018 18:04:32 -0700
Subject: [PATCH 0548/2720] Remove flag access in should_generate_summaries

PiperOrigin-RevId: 208298728
---
 tensor2tensor/layers/common_layers.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index e4b11a60b..5845b15c9 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3455,7 +3455,10 @@ def should_generate_summaries():
   if tf.get_variable_scope().reuse:
     # Avoid generating separate summaries for different data shards
     return False
-  return getattr(tf.flags.FLAGS, "enable_summaries", True)
+  # TODO(rsepassi): Figure out a way to re-enable this line. Breaks all tests
+  # on Travis.
+  # return getattr(tf.flags.FLAGS, "enable_summaries", True)
+  return True
 
 
 def reshape_like(a, b):

From 185f3861e1cc68fa9c0342b544f0e66858615896 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 10 Aug 2018 19:57:53 -0700
Subject: [PATCH 0549/2720] Lint fixes

PiperOrigin-RevId: 208305606
---
 tensor2tensor/bin/t2t_attack.py    | 4 ++--
 tensor2tensor/bin/t2t_prune.py     | 2 +-
 tensor2tensor/utils/trainer_lib.py | 1 +
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 7368242e9..09d190b68 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -192,7 +192,7 @@ def main(argv):
 
   if FLAGS.surrogate_attack:
     sur_model_fn = t2t_model.T2TModel.make_estimator_model_fn(
-        FLAGS.surrogate_model, sur_hparams, use_tpu=FLAGS.use_tpu)
+        FLAGS.surrogate_model, sur_hparams)
     sur_ch_model = adv_attack_utils.T2TAttackModel(
         sur_model_fn, features, params, sur_config, scope="surrogate")
     # Dummy call to construct graph
@@ -206,7 +206,7 @@ def main(argv):
   other_vars = set(tf.global_variables())
 
   model_fn = t2t_model.T2TModel.make_estimator_model_fn(
-      FLAGS.model, hparams, use_tpu=FLAGS.use_tpu)
+      FLAGS.model, hparams)
   ch_model = adv_attack_utils.T2TAttackModel(model_fn, features, params, config)
 
   acc_mask = None
diff --git a/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
index eae69df7c..99950b933 100644
--- a/tensor2tensor/bin/t2t_prune.py
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -86,7 +86,7 @@ def main(argv):
   sess = tf.Session()
 
   model_fn = t2t_model.T2TModel.make_estimator_model_fn(
-      FLAGS.model, hparams, use_tpu=FLAGS.use_tpu)
+      FLAGS.model, hparams)
   spec = model_fn(
       features,
       labels,
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 34a2fa165..3c082a3d1 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -216,6 +216,7 @@ def create_estimator(model_name,
   model_fn = t2t_model.T2TModel.make_estimator_model_fn(
       model_name, hparams, decode_hparams=decode_hparams)
 
+  del use_xla
   if use_tpu or use_tpu_estimator:
     problem = hparams.problem
     batch_size = (

From 2c2cf7c2331166f1d13045ced7468eb8ccb1e247 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Sat, 11 Aug 2018 22:27:30 -0700
Subject: [PATCH 0550/2720] Make minor improvement to model splitting numbers.

PiperOrigin-RevId: 208370829
---
 .../mesh_tensorflow/mtf_transformer.py        |   3 ++-
 .../mtf_transformer_model_splitting.png       | Bin 20232 -> 20556 bytes
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
index 20bb91f06..119eecbb0 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
@@ -830,12 +830,13 @@ def mtf_transformer_paper_tr_0_mesh_8():
 @registry.register_hparams
 def mtf_transformer_paper_tr_4_mesh_16_8():
   hparams = mtf_transformer_paper_tr(4)
-  hparams.mesh_shape = "model:16;batch:8"
+  hparams.mesh_shape = "batch:8;model:16"
   return hparams
 
 
 @registry.register_hparams
 def mtf_transformer_paper_tr_6_mesh_64_8():
+  # Note: This mesh shape does align well with physical [16, 16, 2] topology.
   hparams = mtf_transformer_paper_tr(6)
   hparams.mesh_shape = "model:64;batch:8"
   return hparams
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer_model_splitting.png b/tensor2tensor/mesh_tensorflow/mtf_transformer_model_splitting.png
index 77270b3af99cbf6c94332d07f9950048636fa463..6104a68e379ee7e1565877c03586e3fc0904c471 100644
GIT binary patch
literal 20556
zcmd43bySt#_CAWB0!m0I4bq_02Bbrg5Tpbpq#LBWL8JviK)NIZ5fG%iHWF?^M5Vh_
zq`QA}^Ev1DJ?A&>z2p9K$G8p)Ht)OFyVhKDKJ$5=wZm1F<?ydiUctb?z?Xj@t%iYt
z$%27#;Wf@>_)bb?VLAMB(OF9VF%JCl!7&YkuW=n7XggzI5E`NXTzD^;Z-s$z6GLA5
z-eZrnzZ0Iix;u?$8$-$bt3gJyYKgb-EK@(=8#&gF?xx%|wynH#;hRhC1o2JVl4`5C
zyU+ahcHR9eaW7I0KBH=}(;Oq`9oAVfzWD7<xB~fmou}~>&-DJp8inx7Dz)wHb-Ku2
z4k^VK=3DoZp!NKfwl<U2FU;3-ur5BewxrcT9vBuD_JG)7nC3b9KS9JKL_|dLlnRQ{
z($Y*;YzX)&CO4}czLJg&`Ty)idGeB>hmMX2VPRn*Y62%GC#@O>`rh7Nw@?A|gM$NI
zkC|Y%(3n^OyYa5RbkR^>aV;}5W`a<Y$7ffrUcIn>bb|ai=TwWnC_OzLab#s}jht~P
zwVS9j=~#0Riv82b&B^(BCdtd-2W>QVDAKW~UskNRxR{iTOm1qqtc+JG>+9<ssjO>H
zo;*QJ$^H28Bm9a1%fQ%Qg1dPkA3ofI_YE`CCcS?Bk*;oJ|58q6W&F>d8VU~|j#SyE
zP8da{m}$F(zK&%e3mhGNy4as7IqxcV>C&a`-Q5ljYbz^y(kqYD)QHQ)L`6jxU1eo2
z%gD$K47D0w%j}oa(V>C|^bn)H!L6V6Hm5w)Oj~_*i~LHk)+gi7NzsK_Spi?ZDDd*~
zvT<@UljF^9cv!;2t?lty|1A1d=jrLWy}kW9F;ULSid`Zz8s2_%bkx$`US*qJo`E6;
zchjq}qm09OCYUUxsi`TcN|nv@2A8fSH9@GAjZFo6+b2|izaY#y%&yZ@9i3Om86G`7
zJx7CR4L`poM<Hrt)sc^ciFN|}?b{J!lf%Q794KC{#M>+as4_)Y*B!*3t+s|*<KDY>
z?>fcE3f|}DwreGrq3>>OP4%PyVqj~{6dyi5e#eYMdqc^cJ9l6(<*+PCwvV4Yc{lDp
z;KILih`mr_ZEO3vjN{p}XAd4ejIc_*V|_zONr~LV(#q-+yOglLh=>SAlBm1g_)sTm
za{=X8KXEbeo|l(b)N<3puox{NM9O`yX4x{^F#!%#c2kqYqPO|ym)Bvzi104wpctm)
z1MMdsKYskHp6wDFA_jM?*<TW)BdKQP_wV1|fBsa{)uq-;h>2<a5`TN+=;Zz?+0n^q
z%~(Nk@w>NgZ`RqFnJokj?Hu^LNl2J=s*Un&79lnfBM1$QzpbhKYdA#x(IcWpIek@C
zp81HXq9VLd<ULKzq1TzO#*rn3g%Mm-y`BZ{-y81!9T_-t7or~bTuSeZ<l}8(BsORm
zirJv=j2K(F99C=iO7s4GEC>dKi54uP&BW@1O#Gzdsq#-(f~DRhCYC$QNPR(}I=9Pr
zc6Nl)l|R-lrqJ7py8WH@o;uRscAak|yoiZd#&0rXlmD_vOO=c%MjQQ!17&JzniLWe
z^5{fSRyKs4i%Wolf}&%=DH+0ET~*Z(7R#|#w_|OuvnVq$(E`3?XJg}ukB=X4<wWtq
zcA9%yos+{jnyw2!{N>A+NYnx^+<}FWkuguSxu(Wrc+~}h0p_SCnlr<Y$8y}xUC{<Y
z3YHq~e@$;qtD5!Y%a@T5TH@m3CL-L|uU|*M|A!ABEQdUuowpEkwl+5WuU@@^dm>3m
zNs~m^8l1b92i!)e_%?KeJA$TQ%Nj{#y+)WHS35COUt(A@oBwIr&_Aga=~f?TzYw`w
zP9%&Sx8p59%goF?u$Nm^<)#QfnaOMtX=P=FfL)zeRrS-Oagu)7t8wy|qx7h#h)DL5
z?NVPliNuX$F$F~?dU&*K14DSSyo!qN{QUgnnf+w$Cy|kO5dGQ}mXCTiyM~_!UnP@G
z{PF#J`w8#wzkjJW&BPL184cR@CcW)@>R{>ab19oBDav}=AlbRO1qnit-+%o0%u(|G
zJrDffX1oSrezwse4O;}#)+>aGD0(git}beIhG-?+VXA8F<e?;ySFEeAS9<j5QO^by
z-)Xz4LoDS;0&*tIbE@-M#TZRmQshE8Lc7_b)zIPb+2zqD-+EYNbWnAW^DVYB5!*XD
zI&Krb>~dFb=)Y(BcdW(VWF&J~E<Bt`(ZV9Xps0u+76)NKIWjVG-%E|bzgPA&g^tbX
z(aYm^sWCA|kRF8QAC^>C3#RZyx=wT@{`~gtVd!7UZvrvxrWKt6X(sqB)IRqjbZai3
z?uhO)@FX~s*S5&`xm7SMamN0V(T`q|;u&o>t?`FP?#U1rT@)v`H|8Dp9vB%B{xDfG
zuxe4WLZGbJu{n}dG*xbre$;h9w=W7C_0{0-Ncq<tu}RXGn`EQJW5o_mPAydiTY`#U
zFr%VnmIDmxdU`YvP>nXR)ckr{#4Lh>f?k=`B_+|2j{ht#TiM#a%g%lg6-7{GGeWF4
z_1nd%^%5E0NFO``te5H9WMfaJWTqj{18eKsu+;Z|zETv#yNKD7a>x4n@}ZB!^Q3}2
zqOHYV<ZqX#IlKn5iq0D<^L~B}!#Bkm-KK?tUcS5n>4r8s^hcp4mOVZx$L-r#<b%SI
zTBU}qmIQZ1Sn?YiX;5WdU%oJqUXd}V*X4qu0!b!kbW#g1g9~rksg{k6ZQ42WNLOkX
zYJ!xrGasx5q$btJkJ}qOx4L?J#aG8@mMySYmiCG&D_z^?oNILl&J3D-gnfN|zYh&X
z?>l!IYu})aW)(>i;p2;=5p{#RYBl@$Zf`XkwqinBw6wK-MG^OHd|VN#!d-=U?j*Z6
zNnN=~kX9qexXm|aKfAcPBK%K%D5<ClK7M53<0J8~em-i5h0RHD%<%EU2TC`8u0(4^
z^o($zC<Q_S3_f%c+rRzWVnS|r@aeLt>}W`8<+NKpWceJJeC={`8ImyDu_~G7q}0?@
zs1=NyoVJQfW^cQW_+!6)`}QmNOBu)J=4Q5S%us^az{G_3AKNI6M~?<_4GdFKQr;I9
z<}@{B{Q2`N5&?+~cH;Yj0t)xGI~glcJLN>$0^@ZtwoyBLyt-;MDo`*X-NOcLXL?|s
z)VjGL-4{Hz5+#Z(eE&Xdanbzn@NgK?J0va+-5Q6kyCcjG6QDqbcFu(h%f;%#L<9*_
z6Krm6g>pW9lbZU_(^HsA(5`)HONLuHMn9sJ0hO+7u^sUhB4FQ){2-gEx6ZDb$^=g$
zlY6x;;O<Qkz6Gtz(T8+Cct}o3jcCQ(;`;idq%rdSca#Uq^)MUQbUxQmW$)j=$AHM7
zMwevF0E~r&1+4@b$nPCV^d%(eZz6|2#kM`j7T!P6ZuGiSs96&A8`s9t@@-CziPY!1
zj=G5r9|7VM@5@Nm`M`Ji`Eq;e`uY;4EXuhz8F`Em4FlW8moWo=9i_&{Uw2P`@d8ss
zR5Wz^nga2s4S`riBG}rDQK?~(9ns^W(c?!^?$HvwsixC)SXa&P=iTZ<$JOb2>tZ$t
zk@y+ZE97xqT2US*in7i&ND7n?rpY9w<iOXjh5&Cu(*i%GA-l-qK7HxxB6p(g5EbGl
zDtEdZ5veYuu~E@=RO__x<(D61az$(F<=3v*(e>mJpKJt)OcY`80tpBf09kouWlmkq
zU6GlPJX4@#MhPvmI?>LiuXVY<M%Xv@Bk_5o@w^h}-?dA#6lG=6i*((Y<Vb^M{8cd+
z)<w2B%^lrUhlSlOlUGd8gd<~4gqN1i18J8`GP;wb+-^SN_9Vt&0M?F|w|9;~*amUj
z0QHyy1E0xdLP}aY*$~7`gp8@|qxm*UM8XztMmv)eA`I^?^|4nG3nK#2A#LFTyARI_
zU>ZU>e^7`z>X!^}S*Az$T~mhz!wvIXJhC&(Dp4T>1<Jyt%&#9Rv+S{;+NN(&Ysp66
zjDx?Tbe<mJagSVVeQWcW$pH=N5g?49_eDjK$-(LXEQ&J?wST^*AdYG7V{p9>lbS>J
zW2D@I@IWo1XS=(vPXJv4T(On0X_o6^o?^#b^@5Vc|MA6(7fBJaxpE9nk8X2wNB0?r
z4XN`am_)FW<54D}uD>eI&TcccupspA=?LOBjA_=GJJ}H)J%2zo5*DWz3iJcsM&Mo!
z9N4-hT>bg;=K*IbdjU050Xa9fDt1oJJLje6bowG<Q<y?Y{%F8XP@g7lz)yf>YtJV8
znYHztpDe?YgM-oylSu_3h9wT-sbDgF^1-;-tpnq}f$2F~74Af8N%oi|z+!_CFO-y&
zicg-z1JogpRTAXqPwME9h3{Y6E2#*tLUqB!^|S?FeT7<hdBforKnzz`*ZWO43}mj&
zQ*@j*i%zvY=t9%zpsA@TSMTra<P;)?Z6dZzP<;mIA&Hh2YPp?Vl%He0?#EvOdR&@4
z2G(=JRUT4Vsy*AwFkL3J&bgO6rfWLA^l#m|)e}j<Crgu$uoo|3a`Fi@h<@VdcZx7s
zKkz|g(ezbuO4#-M0MZBgpD8OxY1g?}jayZA1oe-#>@}5^Iy^X`GMV%ulszg+AF!#e
zsBr13C0~IE@b70auAz<V5~k*&c-i32_*rnDdSQM(2?F~TyZnU&gb6vNZ5%-MTeolL
z4E}<@Hx7Kn_NN0e_yq)ruLK*LGgVYqmm48Nf+`-)*6MNr5*Ybp5;~E^H2zpZfI^#%
zH2eK~Huem|!v5^s>^B}7QB!Nt$%2{b=`?(r&3Sn&iyOZlCU9iIPBB_xV95FL?VHrS
zdjS(hVYx+blDNc##@1F>qm&U5J*{+N*Ae7+-+ug%(a`uAG|Ku^T3$LYy01?e^8N_m
zhKVfnH4hXoZyxwmR94D(tA=?0c{lhQk|DlMyG6x{f+7p_BLmN3)q7!V!d0GIi_AZ4
zNqlcGUw`-5_}lRC1ITEAMhC9EKwhtLs@1)S8DLb%8&##p(s_Dvv_fF?O>w~_?tAw6
zqmY{rLr?Sa*W2%}s02r8qM+Yo(o)z-l7?OY8Y?JPzf1f1aW7)l=zhel?%euQSW;%v
zNx;j;Cwwp;*Xe$Zfq_9l#O2pDfRP+eE5Zg<BqvrKCXLr@7Gf?1-cx`2R3z6U%1k@&
z(<j7JBd7IoG%nTDOoTn)h$4^9^Q|H7cAxUCzOb!w_Rnv%)2Rvfsr0wE-He%hqvcIY
zms_y8@O|3My}U%w@QAj{R9j=4{f%9<x9$oO-{)1&bnVe<eKMa#-|O8cg&Lo%2J>*K
zcjoQv>%MuC5_zJUe2@0qStVll35O#e9JZTIHh2PZ?@X(K;uw>eN#~pE`sB%VsNGL*
z7JvE_1GVqhA&D&=htrXHUz#u?+9Udbib{07PXlpd<8OtfN7AP6O{rpEZ-?XFb#=W1
zeW_MLFd{^ija1#x_F!`X(52&2y8E}Wv4=OW8J-;O;98SGzdpyvOdJP&ArT5G=OxNg
zs0YnEi-Cd?*$sHHKq<VNNvq06Th0ETQOo<mG6xXH&F5~hF~lCfNka1g-&BoZwi)E_
zhlit=37?q-zKZ4l|JjR!Jz2<U&u^=2v)-u_7hs@W8;<*IHvYtR-sJioUsf$=)uUnj
zU8F<y)2I1we{pG5lYA5+{3B`U{&=IH6Q;%VCGTawGmh_nU$E*2So6-*(*;HJW0VZ7
z{+d-s&j9ZwPlKS8iS@6}sty7++ySV*O3@E7>}XmfSt?_ckjfnyDB&?=v){VIBYBkr
zwb~^MNE`i@>(vr9EA*C=*h{@H#%{FYleI?P%aJFf7Ha7g9GclU9?iHb_qH8|z!cng
zNI5=9;4XQkFK+tm^fpKRMdc)f%LEoSfx|-l!{O>R)b`#K1}XKLgB<??ahxifw3k;M
zZ1}J+#Sd8acXP-nMU6r(ojD5o2#_g@8|`P-9xT@RvoH@~b_*ZgF(ZzH{f7B<ax!Xc
z#T*%eKTuUwCGavLBB!Lhyc~)RJs%%mUD2mcM5r>Jj%&a?G<v!A?8-fZP?En~LKjd^
z8!S3HQ;A&vt*ezxNWmI?T)sEe9C3PDqT&@v5<?s=52;*{K}|zLR$H3_8g~q+2n}ax
zd-E<LCd9-ztTb40uND^<lWe2Db1^U&S2Lc>Y%o@fU*x6<DvOwr<O(bwl-=9vASa-c
zxc*@EQE39AW0<D>Q;-z>4GygqAbgSmxncmKn_pU*+dAm#>B+6FO`e^dwLCHe^juL@
z)xpHUl^AjNCeMf6)&_lU=)_!HTmbbYCShGCA__;6d%P7D8_cUGlW?K;Sx|7Wi1gLb
z)|T1cuDGVG+;DeO52cp9Sva-U2{6#0e!}i~k{NPAGfNexc}g;u40snvqIZnzVaD}N
z_c?q#s2-mZ_bCM$_>ivF*!jrgw(JZU1S3B&e!9n^^=;~lyo!p&+35)a1)ba{j@v13
zG5GJ?L3r=lL?R&h%h@CG6!kAViKOrQG#QbkJn*)&vwIPRjz*W$ZTI&R`2NRK^BI>U
z_3$(sO=ia$7ZQH_9j^`YJ;FOVu~7C2jJx@pB+di+9q8txcAWqQ#6fjXVgQ^8eY?#_
znc46dB1Y#9vUgheXE8Nw`tEf_j6b_AC7A)#Wc)rjH|j_<sW~}0Ux|CC9E=XxnH}vK
zY|Tc5kJvH1YcEa<fGHGpUo}12T^TT|t*u3fQvqWIko)V<kcC_rYzvdkwhLNgS3bKV
zbGYy6A;k6_wce!fzbuAmvY9+~_1lP8V|GL0#yW|!EdMM<F7L!92psdw7Zq{b^r7lp
zaQdvrdJPBK7FinzvSf~c`!iO}?&6dmB{%&iuhKP>Dl66&I?v{Q&(hS;P%|)~i+cmX
zP*zs9?$tP;BMR$+fhx<ES1>Rz=qXBvcFxVsZQEkDl8_on-2dT8VaWKEr&Tr-s02=`
zJ;ckfuJ7*aN;Tyqlnpc5f|=aU7d#zC5ifI;+J6@7jMTYW_mp*2@{Z&xo>w%8W&Rjo
z6bEv1OG|moHIz&pVDg;<);Tin-n<s)5t-+k8&HxeV5{(fT?2CB_@oe~q%XCd|HZ?}
zAQk;925dni`HHK4b0-7T?7lHJYQ~Y_opbaMLK728(7ZtNn36&{ZYKzBA>xU_1p%Ak
zySBEr+y5#Ihjstdg-6v9u)W&{{A%3kdYZ4E9_;}VJ6@ccmYkA2keB};AC74a9jsQM
z?3q@==NLx0>=H)i<@)-1U~deeG$59Te*Bnm5&6Bc!h8*<)bHdFQAwDKX$>ijvg2-{
zq}a}0!O2W(QDI3?g!vea_NG26_EbX|3g5}e3FF6V^WGGbs_5;%RN_=}($d7gvzFeP
z3wjs?-R=)O)4Q@YY3Kg9oBKBju*ETtH}_F9w6Rks*SW+=R!?zsef?(hPn4gWbI8PR
zS#hdY1CBB<t4W%?X>=2Ff=mIYpmSHt!?WkTgvct^nQe7LBqp1$U;`&}ULq6Sd9j(h
z=0(#~P3E`HD)ig7bk|^7-5GhtJ(oYo3XgC(lmJOsKkZH9u+LW16}Tyu<eEJZqlAC{
zma;SJk7|~}*KJd$m)Ac%oQVi_SS)ap6R13AvxP(au=XUc(+(fn;Mn<)71d$(MV~o%
zCEq-MS4VQJ=HCIY>Sww&v)HauN?cqMnON=06F7e{u%a2GvVJ;-Y`V<3Mg2<f(lE`B
z?eiQypLu%xtR^(~2jQ9ojo-G>{?OIO5%V5u=n^tGSx&$I^Zt6Xe~5c7vbrRAnHw{d
zICf4MT_+TEN~nX@Dv&C4SWEKjXCDLJks)#U{zlmx{&TJJ6qV}K1<7A8s<qb_Zch)^
zI($*AXCy%nhu5~j@1ox^>;CGbLsmuqqLkUanri9TKr~68ycxBP+-04q3Po1iC+J}c
z>Rq--H+MJXO*((I>=6I>kmCIqe$&lyhe!KWN*W;(G2IU6W$Aijf_FTl4M`}v5pS&-
zqiO3x%{;tiTnu{lV%Xnpt*Mq@-4(<3v_MudJAC$ZlVi?-)=gSEFDQTa{sist06tl7
zh-e96VvMF^#rJIVOKPLez7AvezaCo%vEmDG1`#%sm<#=lQtUW~7rD77C9O0o4nrod
z-7t@rj;2i#Tf*zg-tjm3!v0b_-C!C29emk-U!b?y;BqD^SY4%5lvn;DdJbZV^(8SS
zJ@M@9Lj&#GntOx$$(}Ep!f2(>r%!&xQgZyrVtD&n!#*LKx8QN-!psH}@6P$7<eDma
zgeXnWdWc{QEwr5ZTy)ZPlVgdKM=uLg`h@@0A%@>xxF7wn$DiUkPeHspiT!ik#_+?L
zXT{C_Hx?mB*Ajh{m|A`RueGjPUJcgSYvIle`HMHad9Ox5=kq`wBYLF9sbqvLi*j#&
z;|eW5e$V-W*?HRbR+NA(8cM{>ecZd|Q-bC9{Zf8UP>DW%rUzkKZsE6z{Adtms0Fd7
zg6LfC`C?PwQB5Yiy?Jv*8lKSfT$Ev%NR=~Mo?B0sySaVYXI5=riayQn<wT5p*&}r9
z#VuM&9v`#s|IAoN?k4mEaO+)IOEPmejXmFj585cOL{DZ4n*ApUuTHF=EYGCpkI^JE
ziJnibYP1*~S5(!WOwFwBzBa{tx~==;*js8Wh>;pcMvRq}m6{1VNM6?lK^6dYgxmxt
zSgJ_BQ6vkx5+@s#CgN5410_OTAK4n8<D))JPGXAGH%aXJdV0OUwrQ1{;{#II_mM#J
zt}8(PLV{1=Ht2hK{Wh>?*>!ay=<bHMq51Mm9`-brWN^mK@zV?YpLspbE-5IO{^c-<
z{3IG?f*UNwp;h`kUDTb|TEnB_v4sUI6o0rabPe6z-C3a2^z`)n>W4-Ps$)yg81|M|
zBfyBkfq{Ga`ufI%TGC814KCN3PI$(*F7`-JG`R@BX`(XkgAk*5W!jkVPC7dLe{cbY
zj|AO0Dxiydj7B+ib#*M`fN4HS+238^@|cPSm}%JC+_EN#Pgnm~K3s7;mzf3FVeW_<
zoY^HMggARn=+0&8cu~+~6rF63_SRT|#H=f-uctzlIn+-ea8Y-a^GP!M1!vrdh;T(+
z$gBPa@164KGOajoO%rC~lF6%o){uWAFRz#aR2w@x{WTnE`(~ic^nG`pK&QMv((fmb
zwK6m`1nh&Nx%pygwGKnIG0P|y1MdWl#1+^rlY1t^HQ06@tjDW^kXdi)UDsi5Wnl(V
z*-so;0W?HOG&3_B7#L`s{Ty+yDA+8#YU|yo@3uPLy}L3-Li8al3@>M@#n(R|yMB5x
zeVu!$wNIr@@FT9IgNiAyqJr`4ju#r6jO`9Y`;nFmiLTm~=u{Jc#`p9o<+at3uz#vQ
zdg^#5AmgqC5ZVR%`(un#tQSVNE?UoE=1&ZCKNe(nzcgVd=)JdU2At*nQ3lrs9XVxX
zvA_*ksu7Qn-G%zUT}2dxU1DxwA%|TVmG7G?q*?q_1~|~3GBZ@dC%{6EG?`&x6C(O}
zgf_mvJlTY(#F$k*K2AL9mbRU!i!aux?%=o@EF~brhslQv0u)-N00_WBdWFIy5~2&5
z$pMfg%N^%B5ccI+S*`mgr@yQ-`z?TRdxhwjpPz59fS&|@{`5?G414+4FP`VWfBhmK
zy9ZgMK(g7mXK9sh@pvp!+anQCEManZv=XJjqF$ztVp-7W8?CfP^M9b_(K9jzQt+9-
zj)}PkjO1;NB68;^Pog{T^N{D(*QWwt2i_U<vjN&ypht?)s<@m0y#nLSE@4Hl?30sX
z!lvV!9EZ5a2Ty8jiH;BFN^d^=JkU-X_tcP+^vb~K=*#{k<ceDyXobKABvr@6#01>=
zCD>lZ5k)Ob?5n4#89Mkox4Js<?#-6Non@1snU+U)yPv+1l{T%=_r-?{TTgdsN#0@5
zY<=M2^u-_a$pW!EPX!A2`n|oqySut>!UF&#cT{RtQc=loZuW2Ak^K{#Y*`=H_4w=F
zWqz-jKIJ5%;KE8jH`+Ab!PPG5HWp3<@Gk5i$(rDSoR*!<prD{2bh16@IURt_VqML^
zK-{@3H~mp=Cd$oYg~x}5Oyfq7_e1H}?%YTO5ZC4aP=OI2oSLFr*%&T0!ipB##K}B6
zUSc(YYk9k3{A!jQ_lM|^hmvjBPHhWwLwd!z(y@cNOe}!HU>+gi(3Us6e;;=Jwnizm
z=-tD`VPVH{i1osJqnQcSAD0g-?k}?%onU$#S-Huv@;1P(gA5QH9PBv8RsWawMeDTb
zUkz#L=aLdv1;?>&`(J`)&wwLP^y5`B1rm9;xY(rmg$0;<%Ie~G@0bf@&t8TpUiUu~
z_4%v6zu#{df>p&O^Hi<o-(G-_vjZGmue^zv@W{zrJip?PAM-#@e7?Eqbo1E72%-Ug
zDfV82`krKgsalQgcY>$;>qQ7*t5o%$*AJehylKzP%?;wzsZ!O`Q}Xo{#{i<J^|uQe
zL03O+3=50wy^1|u6>4-8=df$#V2v2x|L9;J{@_DO#qTy?sDY{&92ueTy$Q6%7K^nz
zU>DhGRgcj%z34V8)|N9*Zl7Z{AxE{_>fzLMai@q3L3>&dC7_F;C;UP}6f+>YZ(hUm
z6e&sf^)y&&f4srL2_loJ0$GxuANuZLqMD{Anch#3B|rj)EDge>!*X#oi&G*4gZd<u
z+;phk`XpU2s*Agk-0zHiKN=PPrt5oLX?y$H-ke)K?5Z#cU2f3QpFaHp1N2yL#nuK{
z#CquCf`#OfcDCDX23|9omC!|Bo-_aHW`8=r(n<R>{L?L8o+Hxw9Fk;e-1>O+Cr@PD
z+yqo4Ph4EI)zwL^an&RwCwITNbQLf^q>ErG`v%0L`Wr;aYLNjq=xoNZR8~T*%V*^1
zaQkbBlK&9@G=by>BxjiRPhj4_vyhyeJYWX=>feca!AJGi;X+z30y~=**PEXYs)S^6
z2Oko$v&Z~qG!kb44Z-UwYc>?662F~71S(BA|5o$MFucuSRH>uj@oLG;3(J2a3hAik
zAZ8Z3>Qv}Dd?!Xm<o_WC%^O!&SLtyh!@67+)`E5DUL0lqZ)81d$?MDfSECOm2LzPx
zrjI2hoxfdTQ&XwRGdkxU$$Jqc@9H07`&8oJ7{9c_eYY3W@_@-_U5c5Rw~@K1({!TF
z735bQ=u@FR(Jx7?v>uWjTWj3Y|5;)XJ+-EZzCqQC66D@)Nd8%d7|Dl{5Pxq#Dzlah
zcQHxg8XClQS0|{HBz_iYBO!(4tbiH|g$Y=$knL>1x~?TO1gwlYyKO{R3kRwh*yYL<
zqP#Klv0tghf#O~ntM0U2UR`CmefxGyJqeM#mpyH;H*KiBB=IWe!T!7T&r(8jj{*G3
zD|Z%esmh0kDYBkiUOjr~VYoacH50mA3{YOB8hyvS`V<BRTvNQ`JXYv)0Ll;JzoL8L
zo^I!E%203*0N#Ta6_SK^>XGF(T%y=153{DJ>vM?ZInnI{i{;eqnavGbKaV|%^&CQE
zv;JGFW`?<#QW113n^RbL`%x--aQz|E4@J?y!~%0Y-FhYwHJzf!%*-$sUrHMXg3Qm4
zdJ$7=`4A_d87f`95ySp_ArN4i;qu(H7;cjzK1U<@Im(Fno><y$+8xWxp<Q<IIS5!y
z2fg;+V7<Xr0JaLyD<vP*6BB)sIqw9f*6>*A4Qh!c#07QPl%$DG)H;&zl+{4)l#-DN
zjwL0Mmj}j%?u1?U)6Nd#`iSoPp3v<$fZC>ilDAt+u;ShT`C}Qoxjm?&ving+nq~fh
z<iGebUwjEK#JG!a+B;a!xW$^jFWRMs7p~sm5Xz`0C88+z!)o@s{KqYj8%I!=w8=K^
zH+u03R0e}(q))ra_mSl@Pnn{dsWbGPk6ax#C;&_I`l12z1(7p?JN1VB&aeFi_~Tm^
zmPTpfyS8(nFUYI>cgL`nhCkIzd&9JVAUsSD+K{e_ZC$$Y-D~`36n*l6N)gjYpwEKh
z(@A4Cgt)@xA)iTP5>}7(7HOBW*h*%phHblu*QrZ~RG;0c$)Q>G!ZT0DINR~lazmuO
zmmHLim6Gzm?(OCE{=<jt($d#x@peXp1r2drLFhR6qS?f|Xj@%b`N+UPMO}UH^#WQh
z83G2lqN3t;ax#V9&;NnA1H-aoTnq(kEmRH`to&ZBMy>zRl@L*|8HEF0WEB;qg5iOn
z5#T^53o0xab-sW9{uP`9Q1ZdQGVU?=DreLV;ObbFEhG4P_Krd;uz@zy=f+UZ=VPG{
zaQKGc8L_mp%d4wP!KV_4{GcbTHU`y@HM(_W#bI4(|GVzK&jr&o8;EJ*METd>zNy8<
z#gUYQgw_fMgJVB`7i~w!PuevOGfCBJ8y<@H?%f+-$WpZ&&P4U~MU6BgIv@c@)`@3;
zT#l6#6%`eOMMk7UE%B_Us`FLttg;EOAXRzQ=cJ{hzflTIo=nS?97f0n7$87I1nTM2
zr)+@met(9)%|TZWALT?lYEbvHuK6|lKj2mX;8Iv<-R?#tFW%9D)!&kx*^;7><StO*
zNBrz%TcEBCQ!V5wWheo6K_v9qDKvPcz33tWWIg&;&>Mhv0AyF_GX}<Z(BkXiKd{(F
zio|MYFh)iV^|ATh3fC_$+(~Z<o2NKy!l>2v(xx252}#hHU<TaXc7@r6(8$itPTBz?
zE|iIfBqlvw3AiHuaGZ;nk2N&V7%|D?^1p-&i*@Ny8Kyg8yJ|I0m+3NsYb{z;?1BpW
zt3_`SQA7t#d`X6ShsIz-=LVcU`bGkP5e*Fu(klt$<Ks!b;5P!{6pc#TcIg?q3xIAs
zyS@Kae0^5(?P=ip);k3SX*KOOFGb*;<nG<OP;SxRvEnwj<E35iW?OMG>AhwK9RB`^
zAK<%eH&JN~jT_gRUV*V~_9F?-MTc5l#;S-1PF@>#_rsgVGVvsTyPFVkWMJ8=JMzc8
zNh~#(`zHgVjt{mF9i`^Il#sXSY`NLlvj9To<>dvvnt&dyMwbilv2nr$t%U21X8bjK
zOyqd7`2vJP8C_#a46)>mq0<SHEl2)%goo3Gnn-(zOJ-!Z*@=>G5d)I^k31c<YtE!*
zWMtG+l(M>N$AL0l9j|TCF)_HTkw&9*WEq<e0tr4vewob(Zvv-|K+JvOn0o|GxKnE{
zLI{2dmX={8^aoautVAg?-Ld&M5|22h&$MBAhvTfGLI^gqMPtc<TvK3D?<ar=-uc+~
z5*p=IR#pxSfv3kqEW`3bWk*Lx(kC9qN;*`YQYx{CzO1xyqbI4~p^(p#(JyZj2~T#P
ztZU|(#^}_~JW1CJb?x}rCzXLR!z?6~Pn6)c=Eq#n(hfx=B{N};=v}2c7}O4a%xz5X
z9f6656~4WIN|YNycfHqEIe}gav%&($0$4=;(Otr7Kl7*Twldh8%gKcjlf=k7RM7dI
zlKNN^iuq$CG+b?Rxs)d889hpz>;`~20h0Y($1Akt1@;es5B*|WcpydQzJ2>#j@FZj
zsQv+E_o4aAToIGEy@ZwZGu%f;j|7UEh`Ca1;StnTpmQL0yPXey2|(LPp?}RREXti0
z6kthNz5~0bKE>FN4AJ|#!Jz9%dt8LK*;fVOD?CXo5H+#11<;AIyr`yz+@z=Ux|WP|
zhQ}p|6KA%a7Am=!kI^mS5%F)r6g6>ajzwyAL!SkNbvbywa9F0TS!3^VYNMfhzRYzo
z6zBs6<9zw#ct!@^v_yioYPt~bjOQb)$dF5_yx8Y84RixsB}EYAP4wk51mlXF-^tAH
z<2DlLh)iSaIn3f|ZE>&vu62k#6(f$hKztals;{r?>3Q&a0RRJ{gcCK}6~h=lUI`7<
z{m>aX{}*ek!ALn|@W7WXdcDsB(o9i+jB+>?SQrM8GdZOzxWmUd7YSNBxdQx~@%qGm
zKySg6%*=rRUy4$;05b!7(Nqg}#uY-F6KfoaO|JdT>kfX0U6KG73*6o7>*^pGK+5=m
zAWEK%4468-qf5Qgt!a4ghRwmv9$_T<Zt+ZFh~IzyfO>h|qfs9t5&>ooK#`CKbe((e
zgipQ;8@5@yIukTMvy`<?z5eFzOfB$Y3<Ydigjc`=$Ou6+IjN7PdI53*@m^%moRK3J
ztp=u_;_AtiF1?pwuTD2_aWvB<*n~C-PKOG%xB-~py@0!n07qhxUVTDrY;025HU`vV
z!<ErWup-J1)lR>-N?d78B;0rU6FZgmh4)^o=r^M+hMJl5TTv_wB#8h;pvx&MFArs6
z#<HDiN(TZS%o@a%Ix&wPP6t-x7i$u<`JdeQcz`%QEm@yByJDM(N9gcGje#Njbtx^H
zg~-*H7X$Ji<YDG#)eN03(=8Y5nrfqt)}}tsAoF>h7w>&+KD}jEYjy{?f49&?-owe~
zm$`kKNhZx(Kir<?`xjvuQ@{fF7xv!Yzki8_z;rrv|75;K$YB~oQ@=JybhRcnBjZMf
z|LGD)G7g<Inob~m!K=?F=vf+vqJv0c_fByrMW&}FvHbV(e6r5bUsTWzI@LpJm(wi)
z{-k;={r$vCyyeBz<$p9Y15v{Q8UM4&zhf5rXq97Yv1IJJzMAMa!VEMuIs<@gAYvvD
zHM5dGLX!;m6lNQp<7uNU+IOJehN2xlTHIB-+U!quauP9HwX<h{^4(hz=_6v6G?`9R
zLAKRreZarchZ{|Qu!so@=H$Z!VOMQOGchsYbGnsD=avpufMY*NfQ}Ep+sVxGb~x;>
z7qmDq@^XUOC4*&*y+2p(fD~UpFlhBGp9U1UOL!N}#)fyf#Fpv^!w_tSO9Q|f1@7E}
z4<8tA-@clN`aId_T~JVfPU~PVa&mSiffcQR6uj6I_?+I)$v8=P>t@2c0w!oRzyA%<
z^AGg25A{xHiC-F_iAQhn;^R8Em8amy0#~sa1oZ&7uL+)Veq2Jr4UpVkLp<UwA1;^7
zMv1!NQJN4hI*7I43+me5f80Iy8{_**2zK3!6dAuC&dZ6nx=PZ2->k=OLy&;q85|D7
z_3qplV64c^U&hsU@|kc#J#$;tZENyf%<K-oL{1%roy2EBGPN&TK0!Y^d(gE0R-<1+
zXo<h3k{tN>o(GmRmP9cd$rYjd=$yYo1FjD=2FlO$^D1-W#JmeB%zKxAtFmO`i|?Lq
z1Ts{`Xt8IgW+tR?ze6x3UymlPg}$bdX4Q{WS@Z%le;@%&x6Av!$q(xq9^Q=46=88|
z$bI+jljqjZNW(!fwt}+YnFpUPc`^O5u2PdKPG%~ZvNE{cU!7{iTA9=_Guu-KWJ(ZB
zc`~J?`D}{M<hHw(!TN;Zsf34jkvc=N-(6b2zpKBG{RU025-jzViSiP~phe<j26PYk
zeU+fU=HM>D@h`?13PA>j0x9%`rAxx^M=quNhZas=aQ*!H$PGFWB8qy*Y#?do)P>UF
zTqLf17-O~eDbZ-oZhyjHM%<$9f8&vujfYysy!Mj~R)RIE?Odn&H9Pb*XY5^l&)4+G
znp|@X?h(nm;yzTYvY*O$uN2<Jb9zcRe9U9xjoBx5{n?GX-Ssm-KeI5XOSd)kX?)a`
z2iHxu04oF5MT5{C|De4S9N-p;>M8w>*!O=i18y;WCO@qC6uN!InF<h~jf}ElugBm5
zD>HtI)36u1E*UK)BXN8>Sjs@C{OzkNQSu9xox<Xrocn=#G95F$IS&1E>nn0?jm*Gp
z`($DC38AyJG)v;?BC$Nf$z{J&GPFigeGTwT;m7Ch?}bvr6iJ{Ti+zqLEr%}7eX6kZ
zE55rQ&ZK7@diJRe*yEBQEBza?R{<Tyfodba&qivz^XCVyEm^YSO8@;{J4^oyT#{s~
zoPo#Xz`RMf)8f+fPmPU<nvO_Yr@FYgsR1Dkq$b+;JxWDIg(j2@n$my}dL17x159As
zs0<6LNpqS{GL)E*S`1I}l*dLVI~TeaK@{52LLZ7mFyT_HBDNCD9Ot`Y5Vp-AP6J{G
z!JW2TjFmMMys5xO=b&`CG+n-~NRQ2J{09X+nG|(@e&pJC9jjQv{0o=;FY)VH5iMEJ
zf;mw$qak;-6O;{y@!QQIQuF}o71b_W!@CT%2r_}S)C#Dxao!lb#SK0@VkmRyEk!8o
zk_pY$LXts5Qj`YR4*3YsWwyYbXIRU5cp%Ax9|>=Wf`8PFnI#q{OYm(Imfu$jWvQRK
zxS1(eS^J@@BZ`@iD@Cx7>3h(loz2+-g13+pcH!eMfz9iZ2Gyy`B`6}tx0ysD^{W3d
z|Kl5M>_kKq$^dL37n4-7$Gma-w+myJ?nSgdo}eCqMFFg6SD2c4or4%4yVcFRH=m;x
zQXJ~Lw6>}Ce+&Yfb$oAsU3BKQX3HgI=<<o=X#hh+?M-V|kKqHRrYkpZd_u`bju#7B
zH_dmEjNPs(8Chez<I{7osu^#{w2fw{^|Qq=r1A%U<SmWMd=54r;lHTlh{(_Ide(OI
zP(`J4#SM_OG4WMY8Js8JYN;8M<6C$A@!?4>q#6#Le>CINua5G6)|<#dTmMhn$1S25
zEuCsRGf3qK6u8lq2B>wUNaBi$isYBAz;fz8zVI(lTKV%IZkG6`AIG^+Ts!4-ELJJ$
zgEE&BY;YDhp7u)0ywq7cs48vxcouZx0+U!9K-ax<8>Jo_v(h$PAIi()fRCKn@Bmf{
zXt}t|%zildkh2VqpZ^kdlWY|$d`nF^6&T~zY9&DE0=Y|wl@k!mdB^-y6^JaKZZFP<
zhwV%h%*(0laZ5jjrlZ_yjt~zdiB3;=wcK#i-a@p+W~0IBw{j<GmazcAw*yz*-9jFA
zOu1;+&~_KFqUW^h`<qfUzbMV8ZTiLI1ZVt&4O}Zg>1aMX`FA7Apsv$K90xF3+E$@m
zK<+Th6WzswzVd}D<g5m-Pq`c&UBB;DJ4xdVF?hbE5?Bpc)!+XbJQGm7iickHWPNbx
zR#Y$BtJYtw(|@s-PMAv8`pyDfx4O1-+oki0<qWQ9TAv`xUwhv;c0uUw5TFxh@+A7Q
z!HKJ`^)s6NccZTUk;I88dqS-w%|royOsl!iSTn4rcK*bZew5#@BWK;w;rM_TL^5^#
zHToJ8?&JQ<wjDV;zjKOHKGP3QJ>`^`Czt(C$1)?rl$nZc#X1Ue5rzjYA4TDRDnqp-
zrmuY)D|awAY1?0Q^tO|h1=%w%GUsG{tGT6^u=OgoVK?>uw*`ku4)im|41hYx-cN41
z6B?0oV&?A=(o1On94*7cog?FkN#v~vB?+nA0?^u<*#6Wzp0yP;oMe~-p;hf0Nl8#)
z<>_vTLwB;nklb|RQsE6?v_61eiT&MwiS0f&v8fdvK&H9~Tb&rq{2dY1Ea3{BkPeq(
zC`{p_8%L0IW#DK)<=5rXcaR7!vGzT$Jj2XPDTJ6{P>*k#WM-!k!IS*^V$HP$u(l8F
zo2C18FLQu=vw~9-1=JE(#xkDwMh89Y_W!H4Z%m)&`2n4tyEfjo1MSiraFhwS+gIQs
zohO;k!4)BJ>_OT;S4*pSS;Q8q-N?iaZgF(II3#RKJ5O2Cf8k1-^1Vy%Cuq20H9QDB
z6p&}vA!r_khf}-z``yu%X!=3utW1~~-OQp6-GH<-ydc*fVmNqzy<8G0`fX-xsBMP~
zy})4&@N9~80CHDW*5f)`Jc*B+tqp&Fc4#|~wHQ-ip;(lrsJ>1IuerC4acH?-0+B&E
z*$P@52DC={psmsG;_1?(njCRmww3Z)Rc&Yi!TTkA17?&*?hu&t(gP!pgPT6v#x=j4
z&Se4+d(jnMF$rV@i?+`Lb>1@}5%6yr?|Epby&@(_ehq@bo%-p4#F-PaAd695r@V<r
z=rx46nESS~=O9&fB-|xRRmB-_yRRJjHz?)`il&Mg1a~LwJl;}qrQo~k6Uu0#D>y?q
z2*I)o2eA2h_}2oRl>_ZLgQ6S$5!hkviJ4iP6v1FKYw0e21z3FlCPE3=^Y2x_f2sI_
z-Q5!^9%R2U8rd&LovAZzmiUwF%iK5GmKIyBScj<+;iqiX84zfze@%_hA;?~-fD+FO
z8rVobaa=XQ1Cj+OwN4w8+}+=^uNY^(ZB$o>#>J^LVFR4{7zZsrj3e+Lr*OmMdgku`
z4a4OtA~Z_pS2I~n`-1H3!n?H)G@8mqLvx^d|1pJ0-w_l9U9$&>zuBs#`T4J)t;CfA
zrGXjH@_zKc#oe|jBtb_@pu}*EsMy>hmhgJ~?HQ@W^XN2H_-Q@zCRrV!>v#Q1Fyz?#
zUXNb1-h`uXj8!UT-Wh>Zsc(<T9DWu!OsRp656GUuIFsG+HuhWK)xo`O5}DQ5DDFko
zv9$FWear+r(IHBf4Cm$eM!6oXxA6G=&HCmx@4leSaU9J|d-unn#;Y?S!0^oirvOu=
zFJfK*Z!lQ3P^Hx-4Asv<Qd^yibc74O3{m-I^m1o_fna~qbw)?<e2u$}6_{88z3+g&
z`~Lm=w{Vy#FDw={bw=bjyH>?I-J_mWg#+?_wNf&bCn#S}6vNw1OaJY*dMRnh@GrNO
zidpN>jJNkps2=!8wA@`i4%g(>H}%NC7J74m`w=W8T-2a%?}A~?MvB=|R%99bZD@#$
z&l+s!=(>B7?6G23eUx;hU-|6EYn^-Gurh&XAQ3OZ^<3dOeoOL4bUG~HrH@t&n%uv`
zK}7Va`@c<AKE{};T3Y0Mg?Z<SH@XUA{3CMwAJ(e>B>&sJ+pg=VNPb4=j&|jZQ_Cnt
zGGzgF$uBp-3Y_=nc1s98m2pQRkSb6yvr7QRL3c#q>Rq-h2yuUP^~KpSXTuuiesBKt
z@|q@ew+sbl3Zz<C!~dJpIk!x$0New;0(c8BAw!Gqr?`T}{O}Q_iJ^sx$Dg$(O`|L-
z4FlXYv5BhJ;M<g8LRH$7qNBnPEDfDDbiYRg3wCYNc9*iu1;Hwmn0a+C;oj-UdFtjx
zsIwrFxpo^dE9;Ts%w(At07y2RVg6tM{|9Ii@-ICENC9y`!MOC>;Rt&=knF{MT(F`j
z(km#&j-#+=M59|&=&*;n9Z)`CW$8&XLMJJP@cecr`QS;MA^L-R7DBk%&dff~xYESX
z`YOQdXyQ1gtwW;yfz0XpmGzEC`?_{Dcb^{0133?2a_}aJDI6(>mIKF-rDxavNn+qn
zK^G|^VZS#?!)WFtlYr1>2YUv&7F42pflTOpJJ0W{8yl0DOz-te^>TU4UQf!zooXbQ
z^hI~BVPWXgX7~-<ArU#^+tj3K_+t;jF@;9o((Pb2j>^g34MlZI{t9&O4WS#qw#xnp
z>JU*34>L<=>(n6#ly5!h(gd&C>5$Ss%tqHChEH#j6tguA0>{Q?R^?hDjZgUpn<cj3
zns7v!N4Ba}zl3?h0vmJdIdplI&Vzgbxg#?jJ3lf@Xfyx|yXIg~1l`7m5@vqsNye@{
zJuIvNTP5^?O8V>Y1_Elvv3LK|>zb^qh0gU95&z}F^`|MtK+?GXHC*ar{yQ>QG!K)U
z#}qnoL6I3`Lj3w^$TCMN$hyMLj5iqLKF$-8ak198tjriUPELl(?5A!u?N;iLkXM4k
zdQ%nNaZUKDoBxfU+|R2s1x{DS{m-}#kH|jTF0VBw9IcMsIei5tKqdq;3$DNb@#eWf
zHon?R;c~f<&pv4+f@CZPeRdm$Zu(E}-{KG+3e9j}$~#8LHnFhyE+;#hGY)^l#{Sj9
z`=&7x!_rE?%4g}=Y4@fHPqoeL-?86GgT~wdeC;ev8?Q=+DKfwbmix#!HV95{4VUL9
zbNg)=@$!*5t=7`9@Hx{P9M2+CZgIeyr2I)@DyXR77qCASNfj*IEQFPy`-70b06n9$
zY2E)j8+kX=S7#fT$_}}3`4`ke4$o8V4!O(9UV3zjN=RvF5D89f8^E_)R=>gzsfAoA
zVs1gN>E;*azhir5heyrI``s-eqN~HLed!M$7e9e_e)HWV7DI4KcdtKp{%M;$?(=_$
zpz4<;zAGA{NQiHF3I#>t>C+?kfvILTJ00MO#@(kr&QYW1^qHD5@4f$HzVXArN=$Ky
znT5nu`v#u{?ZuFXVtDZH_bodZl`L#wOkMwzCqmSsvsJ>>gbql2O}eEp^nl~c7x2FR
zkAegyDX+(p?kLW}ky>Zmqut~SScdonomf4|orT#k@Bslo(Vq)Y9&YcG<Eb_Wp8e!o
z-%U1E{x3XIa2qJb9U0=iJ{u8)amlr>o%>ICl9R6<%#zXld~5`p3ZZJesRE+Kn>TNy
z4Y;+-g3w2S&wPMHNfq>`X>JZ49euZd!U?Y3oggM==x@^31*GBN0n>Fh#05CK=+~SH
z2CAfmR>RjHKYm2muVr|z%Bz|h)xBj2>C1sIRyW2313(R&np-Nhg`E`Zp`iF2G;8--
z)B<Id5<EY#O1TrubSqir5~06v44J-Udqk-I+~@lx8(*kOwg~F6(2Ma%K@n#!vK_j!
z*(w?l7o6uh?o|%5DmN_BJc#o%J&0FHROUBHRA${<Z|au_8f^qiWd;4`v5V)`AFaBj
z!g1kL;)P$!U|%WH9)Lmfj0&iNmVi;>w=CLl<!?lJ@A2_*P)E|Q6(voPXh<x;o7epL
zQ23XC@e8k4*i{U4(+dPGg3sB!C#iKHM=rGJ=)Vu5M4P-p=UvBNU+LjutqMybpr-03
zW~A8>2l;SxInI!Wym*(Ay+&a4PYL`aucLzE<^L26&<8c#LYd%|VzfSr{(J=yNXP#^
zUXe@QJf#9BbtjAzuPBizKL@r2&Wt3Oeeyr^gVXJBHmN<J9F9POHfgARqk|j_DgUs9
zM#GtK$Tb&~0qCv!mrUTmIjnjqkk-<48sE@<nVGg=(38)SmR~JVzuFTgCHiz94#I+-
z2UYLNl`9S+>ELOD`+@RlGu0%XcPt|ocaSbZFUyjflY@K{AD;|IPs+#VHj-c#=~0a6
zBVQI4OXqqi8rH1#-#I_`)sbYTUor&pBHj_7dcfXUwY#8>EtM~M3thPT*z@?SEn3EU
zFJ<;Fjp?y|lG&C2q+R`)h3uMfPsUB>3qD<Xd(V~jF4+9UPpEx98b7?(Q{VMTIa0EP
z&X6c$N?3yRr={#B-p^yk-bj>L%mbz%n3Tu2cE;Ubi%?eJhlh21=2s$$_#BjHC6!>7
z#rE>6tKQ$$)yb5-rji1ru_14+gV6XL&dmM@Q{IuPi7H!q`Pi-nCtIrpC+!EZOn693
z)S%T_$yU4ZCyD;aDaAq30n5jxyc9+$i^T@n_m(YR3otQIvJUdg6l*Mt_TSf45B=O?
ziUU{Ysug_v9W{sZ<jGhSS%eHds!YqkU@b`{mMM!(0afK9GQr3DN;XF-VWH_vl`dj2
z@$FEe<-Z4Q^}1!=de2pf<`RJ)%?sNxAqk1qlu1keoAPR7yWkx<-kb9GI6L(nY7k6p
z<GvArb2)I*X<x9nzaOh5?vyj>-XF`aFtEd)uL`G(40(edFky~WeX6S)MNE`&FvWE@
zec`wff$`juf3!^R5vi_$0n#xo<ACawtU>2d0-l=O-av@7^wbN-9nrLSzM$@I`P@M(
z({k;!jEvd!^=oiK1fwa^xXPAlY;4RoKDpN%9%*diO-@cu6&=NKitQK%hTahNKGRgl
zy+6MsMqh{hqQ{|!U{Njn(eMOyVH!qHCx0;a$DKS0&&tFE3qB)5P9z;ZTcgXo95l(-
zOa*HMz5Gf-QGrj3e<B)(zLyw(f{8>33Jwh+n(qXDA7C4Un`R~Dg@fplu&mFOc%S12
z5-f@jq7Vsw7P*WX89b;mhy|%wCL)}RcIhv6T$@aJ*&jc-Uhw$Y+XBrI8;!0q4rOKK
z!@eV}akIA?-0lTA_s(y+L+L<?bFr>3&&n&z`p;YPq{zD>>AQ=b7Q7TC&JXhkedA5W
zs&GP)z6*sKg>e5*m6bn+KGCKcHA%xg9#zs44b<EjV*9;-jxa-rFqq|^rSS@%b+#85
zy_l*Usq^PVSSH?o7<C=ps5VXF9<X)g|AP*rd>0W;Q|;B7FpEK7eF!9N;rkGa@K&~{
zgTCJ0nXRS%{e?sW*yk|tRo%t;d8XNo{q#+WkPAp2FJ)1`KX!#VN0#>)vK`f)T*q(#
zO=`1@F8e$J0InC6l(;LM?QbwHE-j7y?1uY(<uA5w3J$)6L+|ymB&OT+?T1@|v9zmd
za)DLor+`4o@O==M&BiFcdVfb0PcZ}*(C#EA7tWPb2~Wo>TmqiBrm1N)xq57T!$*Qp
zKDHDLk_n4-tDF`m{Nt5Tmu=C55tft;ofFbKdcC;wp%!D5P0X|z37aHK_5Jco52YW^
zuFp8snCb9SeqLBG8zR`)-X@J@a!sSsXCs}iv+I<MdRiQ!R^YNnxc;iHI4UJ#_(EVc
zdU(<~kB+zcCE>o%eL)8WpDz+ay_R)EbD8CSSm&IUj?Qw*su01g_$#oDM=~BT)k#n~
zP<N(Moi6sKCR8ZAdjU6|@;fRvnKb@1T=nKe-8oU@nS}qP%l7C|Yin!Y<4RmkG9MG*
z;^K1f@DPeMd6hF26LifumstXVLZMViFVRW(w9L=vRyV2@+>GnKncC~7>%$V)?b5)R
z@PKK3$iRQ*_?UWdaB$LXtP0Z5wQy;X47o=A&hBp9KRbtE63NY^5i;$@2#Dlr`zg0^
zi?9DmTg-{5d$G9*&$B%7b@lW<mo%S_@|~SWMT_Z0SpM~mxy}SSZ!oLU*n9gHs}*<@
zM*UoNt@#)MUb?6Qex>C)1{=`(_K((@v5&4vJZ2*;EicC!tJ0D1W+Rn?-@M7~s6_6@
z73BAUx3&+Y`MKVy;f98WTJpmrVF%pB02>DqHtLk)S!hEZL0|wb6sLpBok?IgjIrn8
zp@$kZWoRe`CB3P&B?yji`2724ut?ywP5jF&@Y*~To+7-~B}o4NeBl728Tn-mP0b4!
z31;^a>iyI4k<U3$;s$Ui4dXe~x2WGKdGv{>$~UQsEZ(%5Po5mZ+WY!)4f4aZ-bC?|
z%G|pZ|6;k$H{LNW&k9m4-G@lOAJ6V|l|@EGL|}_^4VdgTJ7!`K;pgdcd8e;{voU+X
za>#e8(K~dg#K3!Ps@Xpp;$dP!Z~QpGQ9kyjJh9@S^v!D+82ps@)U>n=>B`6|TN``(
z>c3-3B^Am7f`YJP49PEZpddj*upti)gq2lPj&7j{4^NJz+gpNu4V2kZ%{Y*^B;I^?
zu%--$-gpEBml{ED2wbiMYFB1R46WICp6#yakcV6t<Zo;-{Vw6N|6GeB1VY+=tfvPL
z57`2-a)XR)ZJ}mw!ThV7boE>gb+%mYfydnhvj?$34nla^36D5##PpgY5D3STPW#FD
z+Pb<wqV^0i&-1Kks{8x;+FbP@Fr*m<Vrj#rZwB41)m?pBWjhuon*$+S#&Lhh_3^Fv
zs*}wFF?8IhK)U;OT_(VYoHIIxUkSOo%B46Ml+Ct37cDuudF$m1*rq~Jd|!o$#a&@>
zX=A_>wYDM4Mv5*}l9JdMZxomx7yj@TKV4X0#<=_FP19DB-w9>q(_-5mnJm@23^A}}
zm_Gb?HkO_;H)q0)v&=$GF!$%rm0DebJgYWGYH5a;1Ya=b7^dXy!3>5NdHn7+CDzf=
znKa2Y<qen0dJYLF!HIx?ptPwe>hZ1Mg9A?t1I#M(NmxCM7Vtiya?<M47>PC>5e(#z
z;9Z^(B7mHS5Ep+1mF9jd(=iy{7{!%~HMTHc%jN{YAm*Kff7L}<^k?koO~U6d+ze_t
zJ#E_4>MrAW5IgUB+G<$JtOB2l^X1Ex$4^!c51v(0>JGuPtkdqmWsqB;gnBRCl8+Ts
zHgj|&lETSVV=tlZbc81$BX9IMAU-=i;#PR4hT7=s?{{={4PI+HBEV2Bc&vIW{&oM7
z&L(9@0DS7ykT5aPTGBOf1q%xcEiEmK^><r#r~ST<gRZGQxei+ret-C2?CN^G?&I9_
zwBfYV+=dVbYIBl~n%dCMuZi=fqaGQ|RNcq>gZydf>8_m4=H{2x)zu%c#C0ugIp~o&
zJ39|gPTKumZ*CIB#DE06a01Usrd9J1G!Rve8)H#q;s0wcdls-T028tzu&pt79<W*T
z9q1}%;YIFCEqk}c=qb<jbDow8YywR=o$BM`laQUAt#W*#3UJ(z@o3V;>d@7pHNU>B
z)a`iTZZxwc_x3#r(bJov-bFDuY>YSpOwYh7*>3*v<^J>Y-tGaO9i%jo!)PWCaDCSC
zm1ltt-98PNn-^$s0lOM}{QTC>c1m8jvF&c*{@MMz-xc+N^Hd;kkn#BO;|%*s9st+M
ziEt(6=JH<6+L~)~#2{`@MdQ(<uJ5mED(qalIp0J|_qUF&t|JE{&~F+dx89jR%Gy7`
zguG+duB_eM7kv#Rj(q!O_Ufxu&TTiPiJ*%9SoNLOb+3W>0$9C$t&;uy`**n%^P*K*
zw;jS6R$PCrZ128kc|#-TMdeR#ZW>$kJ?_;6E>gRf)1P4Q<3T^Gn3xzaf7#8?_uj0-
z-S=2Qtb1!t|MAB+Zr?uq@1Gs;ye{uqi4q-0lYRgzu4uXGUxGA57HD*pi>pcSBmo;j
zz#u4><GC|OcJj#-;Q2^zbMyIh-nm{c<xsrE)C1h|I&q%z#o&Vn*Sqxh0E2dh(CuxE
zH7{k3KLifAcDN{cEDZuS6ptM`!~~pkDJUsvaass$hgtXS1!nXWAzJtD%0KWfcwoV&
zrmk*n_Xt=;&X_qf^ER+Z0E#{OxWjzi;}<VD*1g`w>U!5i{`{J?YZd4EF*9tsxyIb}
z=c7l9Qe|hX|NG%E|Gka+`uZC;Y-j-P4V-q<eX)i6B6Vkhi;E|Ds2q9qYL$RfN76=(
zlt{T1J;0trUncMjn53knMe9F(`nd<#`qt=D5^)vWwXPeu&Trjsb|%J;zy)g4a;K$o
zMo!zk4|w*#^-o?)CtWDLdprDkX<<>(qovui)Qo2Oc+Z<%8m=`J*pQF=2h8SytFMCE
z^w+<xduozuB)KqiR)C1Bj=nzk_0m>g7a(ZmotHA5FStE9^n5*l)zH5KhNRR+;Hr+C
fYK&NoKkWUrp5MM-x$Xw=I6wwZS3j3^P6<r_=qaHN

literal 20232
zcmce;by!tz*DgwkC?H5khb~&W8x-kQRJuzfr5kAw1zB{7LHty@yOz=+ARrAQB`w{1
zEZ_J2_V-=;I_Ix*HkZL-G3T1|c}Cph9`~40S{jPi@Tl=HFfgtuE6F{;z`$a`z`%Tc
z<ud#vt>SYT{CCM!R{7x-`0~4A840hiIxFe9Vqg%Nq5oq3Ba>%~fpHT<S?=ybuZ*=R
zUq8b~4F}u(^lVpBtod>3RryKf<Q&qnul#%^9FM@0P0t=V84y&KlJ@!J5Py3~OnT{N
zjjX&}TK06@xZ49s=Z80yrEWU}3C6RThapp+Bi^oCvMfvspB!0;-(yM5u378sVVXHu
zO&sMY8-3F-I_)~vq~}|(C!jVYNDN=fn3a|~cjV;cm}G;ZA|oTS32;~t2n3R!ECpUE
zb!nNxD+K$!|KHx!ZCl;@=~36&+xykJEG|7=_1QDN_4W16!tsO~WMt0gM@y`A)V3c>
zO5&#KJY!G1%0GR&Z7)2LeV)W^g!g=AlfjcN0(BP-YG`QSJ6ZYtyVbg^#CvD)>Z()O
z*n{5{JX~B~t<!x?hUsH*(|d*Id>IjBc<R5;O6*1pQ<fg;=(PJW41fR5^yl>#L>@~W
z7QYDmZ0N_2NHa>Z8#m@W>b8&k+&nx6$HtUxZP_I})|5PiSbKS2TI$^5lIOmB?b<a?
zgG&3Jc`rPCd@5SnmwF=u19#cT-e+esS=V$gyztIQOq9QW|NibXUdk^^I|IYR@^|lI
z(#J+5a_H#0xpkURwq^Lrgv;I<DKwH;neXc9$$tCRbkZfOUl9&)tWr!~SXvU`<-HRS
zAf1zwbLb~cO-;T0`}h00I^VYTn^!Gg^t>Kj6}pDv?U%XmbA9rJEtiuFFZJu!uerU}
z0tN;K@effSIoQ~EU0q${jb4=xjg8rk+Y6T!74dA$cix6K9E2yT7o@ZL1qo1`930UG
z)8pe$sl(00X=74qslv@hzJFJ%c0%#~Yb`Uk46{EVkt9RB1Vh$ksYXUd0@H807x;DZ
zb&35AO-!8NNXO}VG3$okzkipXaHqvC8lSxFsqO32AV9rsY<%<b<;#@PFno%^k&$rQ
zG&eUlrjOzL20lJM7`pXd&nDZ;a_Z}8^M;gpDJ@M+O&=d{k?wd&mRD89+McBq{{H>+
zG(e`ew|BrQZZmtJ%JDV_$DUEbWwMlxmvqv8lo&&T^GP<=)`Q=^J)c;Q;qLD0GM}n-
zKRgSVbo4BDx}&E@70@KJ;Q>p?#>x4z(BRP{p?P<IX$DwXQQDa0^Ly{UrR*WxU$%o0
zIZoA5(M@>4M|?Qwz{kgL_Zvt`Of)MSTV7qI*_rH+I?c_^T`uj+Nbc%VA?MO3JDY){
z940ESu2R*sGki3uZ;v~_RrYno4Qp;fLST4sQ1<TKyJ`$cg@uLLjt7>T*OX3|>_5rA
zyME)w4KfM}?fdun=GTv4gK21NpFEkZ-|36yj;-Ghx^vU|Y@mvOS)RG(tHL5iY`X<4
zvbB{}Di;@5ypP7?$4M(5b-nrdoI`WYDOOfi)@Ek!Y|F;fB2Z?>XXhz<J3Bk->AlwE
z<m73O9z7ax2V3#=_06kr)YQ~07po~Q{s#LlKt~vnnwV$>=lv?UR5!KpEPV)m@$;uo
zx&HRsmqP9it+=&)=MoVaMcWB_@k9pJqu#$qpv?dq8(UkPCowTG7d6NG`Sa%wRm<ZO
z6Ri18c6Oh~z?{$*8x|I}Sm`j)CizAs{!GnjA;tR8y?D&tnvINK-SxQ21yO~AD6;xB
zVp)GM{ZP%L?t{~U+I|DE7;e;YhPtA%vhuDFtV+suLB{*{?^7n+iH(ayd}xh}Mg@Df
zmZV&!JvOJM*8Np}vhxzFCV<a1HJ6r__6!a_D9g$Dx_jjJY_z|ti~kxbvWCY#t!V!m
z8#i|h_vfOb*b}S4i3y3tZm-*`3XP$j>*F`d9Pi$3Nr@hMh#lWa4;FCywi>LTu#nJK
z&Z7K$K?mXVu^&Gu%Pze95d<jSWNxFUX^Y?cRAK3NgT;m$Z-tw^{5Y{}k@7P263f}y
znKih7&vyFaE_VFD-KUlY!s%R#?pdEY9X;Jl`4bNm-b>rGm^wcQWL%%D8G_ZFtaM1X
zO<R149p6C|-LbC_z3WeNWb-JCT$vMx#XNO-+PLh~r>>sfULglz13|54HHH@Y4-ie(
zj7^J*FANh9O>Sjy=Rbe`<TLRJ2?-f^ZuSZ8to(72v6S^iTq#^syj113+ILLXT^uf#
zIqIWT&s}-b+e^GIK{T8Afs%1M<X3kQv3ahyQtVCnr?ViAP=I;wP^*dXQf*c&udio5
zTphJM5SaE*#AVU;c4GSNDm+jVn{g4B6=ll%WPE);1;aKssMD)}IoCD~eF-CRH#i=r
z2nh)dkbjKQA7c`FQeCXe!ow3g;aGR=`MY<w#3dwb9UXJa%HqM9AR7Frb+>c#ubZNG
zH4!!6`#rq7zh4@75s<`XpukI+4ZoJ3pU=0d5KeO$3zMANP;>?D#dmtNh9a3&lfDp0
zWW;pR8%V~T9Bu!;JU3UtcyjCE<HtX@y|%n0?SB+wKBtrLFr~Z<JBFbx9zool%#+pg
zPT}O2o2RD*<z+E3u~6A}%NCGW*vJq90<I4x(q9aq!Z2wGalwrJ=I~95Eg6`ZnISI4
zL&~zUVwRVe=hQ0=t*@`wo7s3cL-RVNU(d#7NjbKE&ROn)hlj@o9hpOZt9dmT*VfiP
zvA2I|(CByMa{2P*<+Zh6f4X^Mxt*Mx^2*D}ZWiS7q&ShKD9LgfR=<WcQR=dwa(sNO
zTWL>KP*Cvw`}fzGnb+#a?>dynWs)lf5bNmZ#Ma3~p!jmoivY%F5D{nci8N~TzQ-@V
zmsNi62~*=rw-vZVGd19&DXFQ==UT%CFwEOy=mwgz8yYf8OZnj|*UebrMZE^Aw!S`f
zgBDfz-GG3Aa?U?LcnOYAPFNdyJHZcB4smh|3SQ;rvZSS_fBg8dMNoasLjsoZ`{<~q
z#B5rhN8a&tBqxp2REKchF+CN}KQVcqKS%HHd+6)yn`1G#X7v4uADEncX-1iEROg|V
zI<uKqRh103n)3FN>F}pVzm}Hfx4q2!;<CwcSokf3(_aM#2R~!w>gw+P@cw;pTbqKP
zw8`#IGTRDDuF^?Jk4;zBJx-LH2pCP=;`K&X+|1Khq_cQ7Y9VTh^IIW4_@2c|hi@lo
zen*h9uWoK;InyvUJ!EWZj0g{}-fTSKB?tk7d(NJTT6yy1$<6CmwC>*zGNUXnFE>qF
zOlg0N%RD4>UyJgppkk9Wjj+8bD*PHBR}ci=-QC@uC#8iS9X-W2uYH!W;SIPWt#55@
z?Ir`!+(O}}9E;_pTtO;hJ39`@#m}&#J66-<<4F%{z<|J*V{C6Fe1F$0XWg^n_H)mN
z^%RHX2tjwW%EjOQWf_0O|6q#3NJKz@v~Eh;aj2ww0Z`4et=Sg#0-BGCii)@k3k!U!
z^44(~YNjNqc^2iJv=d$klYi#v+Rx8U+9hRNBy3=JA-kGZO?+q~&vP-(Wpr7+b^S`Q
z&))Cl07V&%%GOp{)6*S1xuq&fu|0oRR{_VW^o8!^_VP;#;rkj($p>y#5}60#!&w41
zH`d?8LH36AN4sm@5Y;uL;;0a^t#!dg{V<iXbAH>DQeii#=Skf9ON!U>Sd9;_CNRe^
z?K{PZz0b*E9`6*IX3JPyTwE?5id4$0A7FflSdCd@qM)i2uN~62Bx~syxT>uuGI-QK
zzx^ejhN^(HC}gE01|7U|yBDa32ApoU3%}{;NEyfsZ>NbfOPeE%iX)a^9Vv)rPw(=M
zatp|OVaiJ>fUwz}|5$2X?JVh7H35MXVg#)=;4!oDbq`2W@)2zc`=mu>w%M)=Pqr1Z
zU+P?R+IzY!jee{bA=wjK{ry{fWu{a7vgY?A_xG-pVtrAS3@iB($7NE9{co9Ku}Dkg
zAb<9-(N;~+Wy-YhzT{(MLhRg@TfD~3#r4YgAp}XmMO9Y*!S)3gNTrDb?uvxCnZIk}
zY6uY&yMu#+PZ|IuT5wax5!1y*>qE$JBfIzb@kmOu<?VKIO<qbtIzq*AM|2v6yFN%!
z7x2T3ei<1lvh>BWzuOHfzT;bOhrBi1qsl{55`xPgmTtK?PC0t+PPWm!BBSb^bE)~8
z@4%gXS{UeVDonG+rw#^`PP`I{?X;Ysi8{^beR%KQy~zWmU$ilA;+f6spF1;U^;5XU
zC@>s!DtuLT$&p8&xHl<@9Gx}u3k!ugI5^tnhQQ#h(YbtZkmNvTT&;WI7eo?BG3Ed{
z<XD|=b8-r%r=_*|5zBG^281PEk_tdT2NKRkSRDil@i}cr$Bh(8${{_^`ZJx!kK4+(
zQgUr8tH6J&;ilcv2_qPD!_38D8Q92PPz4z2>Wa?sDK0is8f|`jg9cbT3ECQ|Zc^$3
zr$<yr93HR;fL&9H2?4m4kJ&$q@YPhXu=t>-t4rx$1g4huS*IvQB7Z~>wdC;2uE}j>
zkZ)%uW$|-Si>2$jnsprjK8wDu=tVit=>6iw#S-UA$wP^EIb`5gGpoQc?E;(v>`Pk<
zd6{4+<L>CP-h)Sv<YD`FjXfScx&a}^C$_wzB2kNtOaSq&t}eZ{R$R5~s}EsBNRDc<
z%Vxs?OWp0H^zL>RRtmf@Hi(>21eHEKD&i{9hGMWnx!Ork98i5WF)?9?r}u}*2O&rI
zF*kx=kB*LRAf^(*?{3R?+s3VYW@d(g&YYYoM5Z>v9daKamW+%Hs>FUb5@zC586RI?
zTL%ZLrM}G0!WT;8D$&iBu%(KS6UuR{>hqI5>gcliMK+kE<GFN9k(A+PayIc9c8h;X
zI=>`w>LDLgsqY#mxFtLeQiezc+bk#|qG)rS)+v|+vfs_?l%9{<I+Ik*X{G%{#~vxA
z)cDyWjsG4&?yrV`nuTUPLEJdz61tQ$u^l!ixo>3jYGOUNyxeUvgC*z=2Bf9=KVF1%
z;$IpLrd2&Wgl#gOwe1qhLd6+n?6)ls2~fAh9LFZh-=p;N&Zl27kELq>LnI|7wY`*z
zRV(~HK2Du!=Cb?44aM8j*x0ySeXx?yo}S$$sJ^p5;Ut22wxSt$Jd|}lc)kg*tCn9l
zyh?E5i*oazOLexj4DapfQ8VNoO$k*CH!JqqHouu6VJ5yX@9y{N<!sH_{zT1ID}H-I
zkxA1VE`!RR*?&mu8je(!?#umpgiA;}tDSZHuxsI(%2BgrCzAQ$!-roossTD^X=z<o
zCwvUI+CCqp@p@!@T*cqt-|Z|lDVDs(-*qmGqAlAiwxg%#Wu0Gb{`9n8&JG)b|M4ex
zWu>F$OK)>?!_(5J2~f>~>ay>&dM?>PI2RNWg4#j&GA{1bOIY7WM)biJ5Qs8HF0A;T
zo}PFOA_m)arH3_8lrtIYU!x!=CXTC74^qa;*DGVmsS+_+k$e%EEseu+^=&hg(0^Y1
zU#;Z-FR%w1>K-`2RaoxRin)ZUUj*PcdAWgEMyYb4NsFB?bF;*FBj+#c`lp(=CZsdg
zet!{K!iv{oOIW*&zPbc6(v4I6dZ!6K<J#(AigRb`h(;%-I{jxky#*NnEih{9>Yj;b
zeuo;p8{(UD-PAxkcXD)03}I<m*@td@0|Pz+>><IcVX}ZBZ&{zPa!`QTD>nIJJw0g7
zy0|L2eMPg%sI4TBsqXQ-YBYH~?6R=qWJlX88e2O%)Y8uV$B)Umh%Sc!2&^~ek<N+@
z!6k~hssbe-SO5T`q>K#Or%#`<7{T5n9q9werEBn|&+tWe9zNj_5(}_lY(I}n<6<U;
z%GtV<gD*~t?aQBq?y&}65qt{nOK$;51qWmKNi&R%jR7>Bhazg>14J?)FQ`cKMQI7K
zf||3t6nA_Y$jHdP4GrB32xt_-y^TPa2eID!R!y<tL*PQE<w`}i@Z9W<h6XV<Hr;bf
zub}UBNjdV-Q8SzEOFQxm5}6v$Q>^B{>b%ps(lo<B>UO6(F{h|#GkXEHU}bg-f2((<
zwKS_aMDy;HOZ7qgtVehV$@9;JwFK$w)4v|?GCgRFRkpTfi6CN#(HohWi8Z4fz`%J6
zs8T+f+$@b{YHDggnvF3eIS!Debe7ZPbCw5hl9O!);$RDPH^V*@7f_`6wyxh?<q7Iv
zYN5h!D?BN1axgvEctk5|$}j1+kw5Cr|3I2IDzB!73R1Phf^&^wwPE3t>%><s{qyOQ
zImGq(>6oa=Pn|tI8O4SWiEhx*gww?~C*j`Q<@Y<D+RDk96}32Bvb3};b~2F|HS%PX
z@H|h;%X-srtnqjphsC%k8Q7@#J)aNV3RO>mC%Uhvr+n`o%3enk&h}1?bYndtykV9m
z)NY#4q>gNwq$z!Yl8TCI(r-uPXw9Z5_wC!8C5f5g<@jpCPSY<(SD6_ZEhgPvTwI=S
zInYL+h&?m!+<cCj6SW10SNgOZfamxFdpd6}WT}OVAw0TI*`pvo5kfX(A|ONf?kYo>
z!pe$UG)mJ@L7Hj_^wg8r{DOjH0#8{f34XfNyzh_FkuD`tm0TSyVsPbn(wim(al?AN
zoehEDAGe5j7eM58JeRrD+FuxU(BS^3f2Pr2Vxr0^RxjxjDKb~xbPyt(y1KgH>#dWs
zfwz_`gE@R3o+2_WD0qwsLKU|!VFjf$TktQit{&D}btffIn=>O2uH8e5v$?p<t{Bq^
zm{o?D4E{7faz$?_*E96SedKwdu1`R*?<EP$0+bhM-fn(d%fuvhXoX2YK)`xL=Q(&_
z$~p7!@UWNUo8|eDk;lJRS6R`X!JjV66JekveK_yT(Rp!(v35>i=-=`^1nCw}PMQzJ
zh#6%`QK%Uut%L{Xu_(Wle_)^tpzeg^<o5iT@83hG9BrJPU*EZT$H8Gep-XY_k{x*4
z6&{9=tW)B~3A*zpNru}yx8~he4-!9(mL@bM#;2xsZf$jH3ph<yKWqIHztDK!krq5Z
z$&55$o1*#|gX9^l-znm-i+QWI3~5w0mi!0Jmv$IhrEN{HwpKZu`&}t?ZqY2n(Ro~I
zmn4a&20;ut8{03tepY<oRJ{AD_%i|C^y@W!<}55yi0pd1a@VO3TBc%KjtCFW9~Db3
z4f_f2Wd7?t=VhGgg<FSXsN!X{Cv3rwC|W<-@=1H_v+<&n@Y^L5Z=ud<b_K8F!+x%q
zEKAM5OfJ)xf%!Qm-<bejNuEv{3RDxh?8qhMeYwB9ZjKshz?rDv22SAhLa8{*&E9ym
zQjr-TQizF5IzlwV6vua^Ni|A6R};5l#F8v|2L7J1=i7&fw*fqEM{`C?t#ll7rBx92
zR+RtVxuBH)B}YGuaiUPCVZQ^^R@}}~fONXt12mh=FaIui*7syC!hdUgXjVr0N|qOE
zhKC9)%=zE*stu~P&$b6XH>jIi=6!ghdHfRdSbEXlw=%wLn;m}r05{SviEB4~)x?`9
zLxL!NjXCY_(mYFHIo%}H@{~DU;Pn%waBL#a*b(hCe*gDQqUMSE@2u1RI1usbPoM4M
z@6bUuglhl2pw>IQl49c*w4qg8!3>5t-3gXXBY#)(R)*!~BJNy_OpBO}?0giHZysCc
zbpGGF*csm`y7QSrzkcqf<$TUojilGB!q>|0hQK~|1_$Pg&ot*ABN>Y_luE0}){2xp
zdKX2|%sPJ&T;p#!WJ8eYp+~O{Pml<Er?JXkmVKuUI}o3BY}J=s+KqnFgW2#OCSP`W
z^m6ui2Z~+V^EExd;0&O%H@Uj`2-%R3-3&&wn4cGJ+KoLchOuPTdQce3;VQH%l8C<F
zI5E~;lF3?zfbY0ec`bnfY?7}|v6#nHz0t;J+2^%FdAbajVSy?x<b}Je$YZXv=GQ~N
z$+S+#5N;G0l)m30V=M~hU@oetygwL&USef~+^bDgS3K>{y3Gzdj^?fox+GZgIQ;#<
z^6v*~Udx-<nk~#u!57BYyo)qF3M_Nc8>ghdO(ct&-eR();0$IkzLsobm)&a!<*%IF
zV5>}4gp<q_nZQeZQlyS<s0t_8m>LE2eP@?1DKL#0n1-ltUR$cT4?kAs`3yhK{?Lxz
z6Hg7Mu@gi3Q@dN1w>zgZa?z)I@WQs0hbTOLO^P@ei9h!_xs^r$X(W((?eCpgt(|dv
zUk{Rz#P2SBcl2RvzbxK}BUSzHdn<&e0<^EeeqJAl|0u!Z_I53`WA1z{9HeYv!H5A_
z3Fs5xw9UJKQ~z%;pk>enykuJYp7T6*@hXPpM53kUO|Hv&&{HUWHm?BbU&+9L22wq+
zkx#q<W53gi{PV{Jkh7|W1|cwFNr};skrlx_EqhH7-GNvDCeBw~mJc6#Ek4F!-bpu0
z%S~NVQDY**D*?I)NY38YFnmnl;Bwx+#Q>lWh`Oex209eT%uG7}H>f$_Ffe+m{RE3f
zRiV)!B_#!vih*x4#N!{T?X6USPE!nhPTG1H3#zl6;8^Y-eYgOskT%W4_Af<Cdx&RZ
zK#u_s2^1+T|6p9X9J0H&*INx)(t6{Nw_FN~;h0ZwBgBsT`Ah(L;~#n)W)cpwpw2?0
zIy8)j>Js39#n%@%P8UU*QU^Ww1q8(W4tc%3y$7s-3Z)DW8d`aIRac>OV*Svskt#Vl
zvST-u#qDVJg@maED_K2etu$Su`-h$;5=~jGc;JmxR8+29!Wx+1@u<tZ7eRO*Ab6(d
zwAhDr5fKppIGufc^amcQzaICfb<4|H`qXQpetR6AP8BZ`RNDzn`dC%Rf9~)iAxQy9
z{Q#1*r>Ccz9xObpOTJ+ZlE$&xY<yzhr(tcnAw5&VQ~&-W*iGGwn-9-O3LQe;CC{}p
zgsgLzmGED5F)=amZu>M;Pi<6NF<czB!}mke`-{El@QtNqXdZ@}DDAUR{7YESN%(tX
zzfR}958E7_FUZahIqliQO$vjEtn9L1sTs>I@&V3Vn`z1ddM}1Oy5}B09n_ZS+A9Rv
z(TEWq9{#@S)y+xo;RP3Yi3*LpQ_DTWrU0qb)KtFik$gQ=nvh+tyC2}K<r<y0HLMv=
zHp{i&F8Xh!dUA9c-yZu<_9n**Lq1H8V*cd-&4m`YRK4;i#L)P%0s@fW+O>g^5h$7&
z_xAR*r=h&(K=ZSAS~m|!@xnT~yHP;+@7iDjL-sYA0$Z)l&oW}dtI~c9`QTnYSi=wp
zBXUjZ0MAdFXapU@UL1T7<2Mu?X8QHGFIPSNN3j_;jff)?bd<C;#KNFE0ra76oyU5D
z&7-~lBkeOd4E$4A*fzf{H|pfz@ak%SwLS#WB5L8ejrO<8d0PPWyIp4kWoR7P?`t6h
z62Dy?0Ta#4%miD_&|p0}I{NkNmzp<?o}M1%2S+XH@Rkyb_V=;CNfb}46A%)T0P+Kb
zsTY{Lz?kzpTM&VWs$444@|BT1C1UFeQr`^2hd9QS^SgrcohJq^E^cR67lMNWziL9N
zV6F<5oq&ia%dH{4qP3^oP|i~NWCxN#lRs8cJ^i@Ktdx@qL^^I2UDHyZy{D4&M7g%+
zWn)nELl-c#tY13accX^&e%%dXV*E3QD{1>JJm*+q9zTwWt6OvK3nN9Zu2?+K)<*SZ
zqB(GDFE7#YpZB3Ld5=iy3q{l6SGl#va&nyfhltbB%~s|8jX;vZCfbb?GC$qaK;{;(
zZ7)gsPoF=t1D&pC4E`}*VK-J{p@{I{Kp+Hm2l)<-Mf_$JAtM2<oy!o;@)6+_Wge@D
zKe~A<D5`(S0?^#EZtliI^WYZKzXuJu*;59bvmVaI&gJ5oE`)8>^|N=r>1hnT!~!b~
zv8Ih{b$MCo<8%`wyHs1*2%^G6s)wE}U!4e8$4&wng3J2VJk~9xRg{)Sp#<vc>UzFq
zX?MMkdu`k6iLGr^ssI}Qd{|o9`ZlxCgz}vs5V?V>Gf7_4@AON4-JA`pR$x#COL7lr
zNsP~r3tAnkh9)OzJwN#n6Q|Ohg?uZ<-B^k6%(x~~b6I-pvHR1+@2`vJT7h|m68+o6
z1dV!R6%_1p>l$)$8b8OD<`!{hmtvb~jPVT~Hw#fbSI3OXejS>pt6yq)6+#w7BEG0U
zkomS_>Ir&{e=LR#X@(Uw@kn+lDDN}f&>KH;H#L1e{8?w{`*$*<k?CpZxsl`G;8<E&
zJ)3_eqF<v*Oib0z(T;{RK@9g)iB46Gy|UH48-q!djgoL!0-6}0qY2$M!5U*|4VRZE
zK$F6G#@`9EYvMmzfP?$S#%RB(-szJl7%zy7e93R$8YqL;9e{OlB;hQ9!c^X~dhP8~
zYRAgY*KNKCtx!X53S1)km)v%CcG#_?g;q8;WsoAEL26LtxG-*4RPl5afpBR()0ADO
zKOB;2#q1<146ot>Ta{z!FC{013;|%drxfbjNk<W=83(L@13o)BFsLvfCdQ;cV)Z-L
zIQ*6AJkCDL9^|0I&K9?8p%)D%6FuD;g>GA{nq2tw@YhWhrkcphasC;-^<*4Maxz1`
zZ>Pk%@0fg%ZJhXO;GaOoA6G+ls>LRU?uGe4Ca}PwguP5_j^3Jzn9U=KIC=G*JESyv
z_Iqg05*NQOr!PJa^i!biu{^)HXbq(xq!7MUZ%KM}WgY~-IuCuAfhN<~#J0`hpRaQ<
z1zBFn!6d0PoOXO;_1*%2_Br*-ujCCyjjmb^<*E<-fo8*f1B2+9&C!Hi1VS4O+~b_8
zQu4YZ*U@erDm9RTJ?m8uAM6zvIK*WcYCrqQ>bUXo@h9HO%2y%xR905<6+pQ}V#3sE
zd}iz=iNR2ZA#8}})T`#vD@TGFWcT>9voluXYBO1SejWg&T>fNcWJgBj#U{YjldSNA
z2AswQtZ8>gfE7sHGU0=)U=rfEPH#z!Yrb>S0bcE5TSIoU*h|cHj}c|};z^<-vm<S?
zLb!-pV-#m!QO%JfW6@41fCei_`(>*i<JdEsVFDUqtJ6eTK02avZ#w)wR5e5@A?y7J
z90MsGy{{NcUN+IidW2J~b!2ZzU*=Z5^?UfU&`$-qUrgcZ)vNcPjN@w31r1ZQ93M2V
z8mF(NMKNVr1u2`E-AWgBfIOi6@S&WQRZcJ&aoqiY3pvrwX~v6CbJo)>3+G10u?iAp
zI#@7pPjjFzF@SJzMO_mC?^C^hpBVU6&_(ooGKz}K&WyOg-1D71(3EYVY1YP5ok>~K
zkIH@%s%BdMboe%`v8dP9O<=?Ar)m>-`em}*S4T>G4>si19>_5<OHl>$%wRK25eCX&
z6)=Q(;4mZRRlJr)<G9#Ce1JX#$YDZy$&_WI=t!pCIf~i3sjVTI;CX_73<z^JwWMj{
z;!}t3i$ROe=IGBt8Te|P-d^j#R`E=Mlqtr={+q2Q=>p_Tb=HWLyQnZQ<Lk<$o<rEY
zhX}t#U;?611<VNkzo^-|S#VC;q}syn;-;Z1rk^^LPr2%VQ#%CRvy@6lcLe=r3$`*j
zW;0bvS@{{!`;@y}f`Z+rldfK-JshjpYOG>`-E~FfUt;+icF?&Cm>abp0g%m7Q=iJ*
zi2LHTSmVrq%nS?{u?jU%D)dQ5u0`o~4mN`v;yE}y$X>d_uhF0bDmd?eJ4hK0$4485
z7aTm@-J!FlW?(Rpvfx-Har|S!WdO7&9v&WUdXU`1%}mo+-f6MvS9K;ucbp3E^4m`l
z&1{AUP&7jg?e`OX)yOMJ{`n9<xRdeohC0Dmf}HUcK=uMo0S&&Pk&!ltP|o6+hQ8|+
z@#Bt8PCOupxd{}*!AW}pZpd9AwJ<#Dq*!hSIe+8h|M3JGn!t5JA+MYWAG!CZWJ{8%
z$}Qr;B%LVn=rsfgM3jYoc6s@CXlB3C8yOrl2W10j7MA&|CLG8qC^WrK=wsW_YJ)mS
zGQ3wOlF8jS$^wT3LF|yPoZLP%G-QEgt*P1H$-#}VT{<kbtz_WsJqVZ)bafQ&P16|i
zTdc7hRr#V17`M!xQ1*WWJWn-_@QHQni}SO`>*Q$R2scrf?BUrtczKKIYhaq=(aU6x
zuHa`GJ)4hx=ia-{HEP~nck!Ci>nzOdI5izB$WU4uiG$wt(eY>#eOZrV6)C~xf^jYk
z2L}gq02&fp+1QAg4hId<v(c)*koOsOno><EV{{J#kqk}WaUGzL!*FNmblHV6;#7C=
zCyl)LJ2$WUPl6vouK4F4HtUA-^K(AT9<^#iZbNPz!AQWv4<9`$0XIoGe)-E;6kgBA
z%tj?dcT65wUZFl{HHy7kh+S!fBQIw9B>R=p16>ae8=$XY3l$jvHd|Cpw3mwmXn=qP
z7mE(8Cup#elM@D=`0o8%2-kk_qx7aQE14kNtl;AX=S|j@^@myrq*_u%iv4tbDl`v4
z@e(?;@?~|e)OINM$q^I)fUv%99zgpxe|()1dgyl!g=0{2k#!n3uQ$&*X8>?od_xh_
zT50@58mh|A$<0l_AfhRKD$f7PBqNgriKondj9+cYa=z^qz;%UgRTY)x&CR@mg5J_b
z8lW2^P%X{PcY63SP`vM}UMf_*xfIf}{B3c&tDt_03DN1fD|)(ScGw&{n-pQXC%<*l
z-$Idmn~hkRIZ1;kAuFpN0M>lc-ob&Qfk1vaw1iUb3;nH0T7G_w4$?Vke3|jggMT6Z
z>YI@w6PetsE<oA3W!AlFgvxT9mTAtKGJ<QEK#_R+LL@XG5xH8S)b+oi_R+1ZrzHGN
zbIr{d(9Bviml=BdYikZ=kz{3U-PUuAr=_E#2u{EI%-vmxK9*HHNFVaHnn_B0qNZhc
z5?9iGGBL58)Gpbiv&^nEHiINlqpy`{)vFFa(oMy((N!q-ZYm2=XPL!_SbI>96c{L~
zs;cfjgAxyH&0_It>Whw^BG9-s&1lNY12(|{x=V4?*53Y7YXLLlLq01EAn<;yd;F`z
zz*m*RqAjie>+9Oe%6(G=GMHY1V9)pM5YeVM#Zzu=q*UC%gFSs~J=4J-Ka}8#LF`aH
z<wzfR`V7=b?3|px-bK)tnp3bdJ~B~@JK9)bgfRC8duuVQc%=c@9w>sD-aKqb((~Z>
zHwznY*}vcdnksQ<_*j2`3QS1cw{B+xk<+(Ced1ieK*Vnb4`~2TB4_!DGxMkdx`x)N
z6YFh)>VDE_QmkZBl@+<>e*m0C7uG*oaDnC4wo9V}u7OXuuM;bW&ZD+wVr7S<y%xoE
zcQg+F_$Kr#SaK_M{6k9|Cl|>y_IUyum#`(mbuMo-n1f?Xh?MxVE_Di&25pitztWY5
zO9l~0>GMN-U@#h+0Lg|*bl<$V)>)VssoOi)oDIG27cXA$p-s}9p5||Pkv~}x*>hj>
z(z|C2lV_Cm4#cOkw{>$Nrnx3r1$A1oGLsUgzP!iN>`9NO$r!kw`c3N+Ges+V67%5o
zxQCW{rlmj+d=_R#?EKg%s!y}rB*PXqq}C`spx%olO{E|?zQ4J;%6R?Cted!yp{^W@
z8(8-#arIHlRy|Xe5rK(5-u@^$Yw@RZ5~-o(A$1;H#zhnm$$*_m?v@AWoNN+MGh9a&
zCup6vwc7sfnG@J;7n&3awQ;-=cho*uR{=Hwy#QWOQ6(cII<O>=3}_h`bZ^hMv){h`
zvyH~7+$9--Y@_$XgL2rtDY(kIwMZIM<OT;DacV!Ht?iRD`|<K9BodioyMmd-V}i!8
zY|subbe;BI(yOo|Ru+h*4?b@Awh-Z&g+HOw>N|sk^sjUqygs|OmO8mxh=q;H7W$V^
z!D$&8ISf0>Grc-JyJElmoMYp8JGS#`#9U!4d|sT6UK)qZt{?L@5bmI(v_k*pwNT^W
z&GiU|{T4gT5GP<aU)4jt?05ok3Ysj?_){JQ@niGSjvioIWiuW+e4-Qoi`ZG#z|c=v
zhf_<fjE-apkxG*A`@ru26RlZm40N$Z-q5G17fN{I(gdF8IF6IJ{n8AMbVpsjvxdxL
zq;c_Z@#uyiFw^mcOWR%@{ry!>s}M+8AW{V|bfz&)R-f08-doC>Gw6oXcl^>N{~I;|
z7&ab$<%y}&0OmfRDXTYK1dX+oi37bZOE=<6AI!`t4+NA@JDyw-O+3<(%WW<-5JM;5
z;o<vG4cXYR0v!(PH0vz-qbX1Z3isWEgQ0KV?id@tc~o~Jj)~@qEXw$B;cm4lSmRUO
zqyjKUeyOCs(^}gh76|tsRIpkfub^aYOjR2Kb?MN+Knrw`3MO+RnL2@pOJ@jP4-9tg
z3uO_>_p3v%WRI>ivp)bj&|-wj-rC=sA-V(j5Gvxwn<~Iz*XWq0|C<E$opX%_zn9C8
zD=tDp&$7CJ5rn)TS8@j7irl)s9LenRw)1L7FHvip&aA8V)>m0rOJY0b1(BFVK4N7C
zpcZ{-srQ>nZ?V-}A2GpldLT_htc)Sl_@bkF|LofDjfh5@xj*(#&e`hE)*!FTE%F1U
z%HL)}cz5c}Rt8hF39gI@PI99dK-0LRJAlmOq=YVsNMWqw&ep3PPA58B6^5U`=vae*
zhd|qn2P*ag<A1ZGkb&d$Amky_(`hoRED=F>nm-x6=pbltBM8{Vt*)SnxfY=~mfssu
zrlE?CQ;nsZyn3aUo!JnclpE(!#j8BdL(+%s_*+q)t;sy$t`LIckUgc%mSEkF9aQ8<
zz6hEfV0_;UQY1nt_D5VxmW}x{7OC{2*#6HlWb=X1QJdjKxi97o{mtzYb32}QN2r`f
zpH}aw0O&>PLn#HKD+?A4{4iNgquPfpFO-l+m(8Yl=F;e{`&|v#<vTo5(LFu3995IE
z%nyy}T7l4RqZ7w;w6wFej*A-VXxiYc{zzF^a|e(sQ-y7HaR=?rB*(Ww?1F~k=F7v$
zEjH-=8w_pyB)1r;gJxIzfabVxL3A4Df-NXfE2lSSgPaQrGC=KZ2^E!<=9HInZAtLL
zfbI#&SwMX35u8iLeP_laNfnIrY)13j*61LG(6(8bT5M@);p?^*uTlo!2wc<JWDU@c
zGP(7Cbrf>m-eN#aF|)Ep0B6wMCI4V-x~X+oCL-`m=$H8Hw5Dvl!Cj~(cu<1sLocq4
z*;n%2jkHaBhE0m%74=yAkx)4Ug*ikqplbQLLG8S<w#FnOAyL}{5*3N5zH&+W|4xZ&
zu87LzJI`im)!_!5T;lYW`fhkRnXU6pq}tt2#@(ECiP}0kBC7&Tpv*oE2vpO;yiRa=
z1p+cFtBt|}K&B`uj5zk=30R*x!aJ38j&&r``Dka0K)d>QdFLl}_3;jG?SB!S(rtDG
zY4$Z(q{=s1NR{^k{^06*i5mur<u@w9R#<M7<i~a%EoJosh{_+I%zo)v$Lzvkox6bO
z+#lqtc95yRh;VL4Z564`Rb_=~HlWv`M;$^Ic)vjfdz|;?A+Hx<sM7?k52o}{p);}p
zO~CZZ$&qk4gyXl*JB$A0Ci)LA*&{p@$}RJKGpwBu^~ebt4&9Zk=p8CUL5)!<(IlR!
zf#x4|tDVi<yg`EgZw6Ab)FE!{r$6=6&4`F|LU6JSOU)yDaDoy+USOWG9<u79c)iw2
z+O6xGKyvY~qMxPsf3t`DbS9%kCYm6T>nRno8!5K#O{1bq?%@G8r<wxVPljH?r-v3>
zEjxmP3Uc{n{wE%lPBZjnAZBO&C`?)=$lKZ3sjRDtnnOj*a~J)KEVQwMl$F!iI9>WD
z%6_brIj1c^)yI4PXlr+S;^I<KAQ!N%(1I1lVM3Ufc*&$O?xbW>tFr2Xq#W4a!56*q
zd!Q?|XeR{D4#M$^7kKc+20U6JFg^dhPw3DiY^^V?|4BoIab>Qf``gd3$2=JUxBYC{
zp<N7x6QYdcQIQq<#iMG=P7=`7$Ql_<Rj+1nR`micoxj!f4BB<h-S_4QnzQ;t+*NPw
z?4rGt$vxZ86k(pY{eX#Sy!h(H4yE7*e*dJtiy0(p8g%rqB<RM92f6(#%eVyBUme7a
z+pkzQ8hwcrW8ZuY#_4w*ew4fo&F`CxjN5T7K=e$}dmS=pC1B4q;Z0Yv%>m`8YydR{
ztP;eZFC>aFv~z>1I?raCb*SjiSxhlls-}H->O70gXBziIaaaS-ZA5QI0j;8kMbX55
ztB!;?Rhu&%Y0Nb2H&3I^8jR$)g3F9?8HLhuo^AFitlLoNJ~$e^by(#)*CjDqbPG*s
zmZ^btc386W#)++N)E)3j)*~Fpwzh6ex{uO#?+g89F-^xbTdSNJD+t7;3wnS=)XwQE
z;%80G?Hu@Ev;9=H^T}nnL9uzl5}<h)_QOcmuVG&-nX{Q-*>I=qv&hRW3V}G3l`6ai
z)&QGjJA^bI2mei$hn9;yrDR|@P$qK}tz!Wf{SS1r!6y*{?+wKGZH(N9t$x@av~iV<
zZOr}tJR`E;kjCP6IQt^&t~-U{slxKHpx;gio{F4Y{#$a=)Uw9Cf@{zgc!7=ELriC?
zZeQtJz17NP80deJP&w)fF%hzLX;o<6>%h@fk_Ya@ao(yyQ9XLl=3ko_?Q<WU>M-)0
z`T}=UYTANt=)Y%W7<l$80`GuG_Kd!|iedU0;pr~5*M-6G$47EZ#4y5;krWr#(kl`e
zpYHVLxTmOJnrJ2f!$XqfP0ro~$E*bfIyxyKU}3#2O7C{YC7an1`8A0EbY9&AH<mEU
zA}5Xu9nvZ2xXQyI5ocLMV=Rhi(h>4WRu?*n;NZ|x>+wag09X~=2+NDPGbUkr#=Lbr
z%7S+P9nd$Xjsmxgma?2<e$_Z^8Acm6QVmbd()bwYLq>;y119LuNHP~c1kaa4YQXt8
zG{a}x{I{YNTM}`&7_L<%JXtP^RC)>x%^lfU`NNHSTSE~o=hT0>nCIF;2PKh8Dv*9%
z$D<WdsDP^d(6U9pBlE*xNH3AuAaCbq1t|uQA6!JC2CC9$8_tU+?NptGsgaRt(5voY
zk7M%P^PSti<CQydExFY|Al^s-XhaxV;<s_TpaY0-rWxTnHd&7~?Tsedz_g&tI5^#c
zA-EQF;^#ad=6oWwb-sf%;()^laafR*Z~O@>B`E=Sx4m@6@BG(2kAV7moQfK%)R)?G
zGc!b4e#z(y?7=y{om{^Chp+m~3|jDB+ms1BL0!;Q1*ESEhjh^h&owcG<@#^!66^XC
zpHBGBx}(>Se~FDZE@8ip71;apQ%&8uV1__Y+MtVzNTwH%6_klC7$tWP&YRvSmFoPP
zAd6IjzI9x}pMO{MVexAf0c4YZPv?X_o%111bq<Z=g@xC+*WWA1f-88?z>Da7gPGrS
zNjeMH25m4M_0T9r{`>*0JCZjy9k%2Td+O#?|DI+O&aY>c=IjJ@ZP@wU>P4)>H1utD
zT!k$kag)IjS;O2M+6t;_^(=&fi<mg-Uur5j>)e2VJY+i|h%m(c=I;Em8Erk8P6JdQ
zsdPJyG)P%ThvNU4fQr0EBmBjF-g1H8U6bMJ!L6|)8!v5z>abs*s(=ju-rGsz5=)Ee
z^`cSy=TnZcSCwwky1gV2hItauo4+e8bbPSDFZ3?iZ1o(7yE2WCi<2Noc>gLL<(P)4
zVB(;&@QYf-U;UFFWU<LVlaq$P(O6)S|CKbRJUn^=0v=FG!JGjaxvz;1PIb4VOM*?0
zhzp+nh%Ubl^!>&q#Do4rY``3g*Vyj{UKnU!EcEXOx7-A+AG+P73$mo-x77V~D3;To
z_p6`r{7)GO7^bp_D+@13zbppc%cta`n~%G~FM&R<Yn*ZfG1V)YZ*pBXQeO58&UtEv
zyN;h8eFHo(ir%(JwV8}7AY7#2PUfJNyM%@L>C-2eiCC?!vOrWlZKG)s7$~3(ipUD1
z*@$QW*}&+;BXW&ORw%1UT@k!gQ=gzw3fwI-%%>@T8kpRZyPmcWCX~7PeA5J{+cg1j
zo?Jd3D^~o#e?&)&=l)^;5gmon`CtNn9I)@~HGX)!p!FUdKg>(YiPV6n1ZKypVfqQ~
zh!)$W=OEs?X$7MV*55V$g%z?$izZjxG^#iER!3uNZj@X1QbWU1U<jo`#X}2OQ}Cdf
z{j%R4$sq;WXq7oYWn^=*zs}c<80~r(zaMn9c`!%SLSY!74NPK0Q}A4FUU7r=ojF!)
zXF^u~{3?uqJa`(Jm>5GJ0g|mM27-RoKdGLm<Rsaozd<Auo##<u+UnyXUtkpM-_#4m
zP!qb^jBO|NpZu%VJj?E0A{`1fp$q*vF;_d&d@_SqCxd-)jtN*Bf%I}U4;Fe3lE9mj
z`VVw<l>qRmd24HvLL;E1<t|nbMny%1_H73CN31!iA0yu2#g?!5UzYZw2N1a_$_n)d
zkPb@2c4HaUiP#_mXuI02nRH?;z0~Ht#tQ-jZkc*j0Dd5WT~sqgXdmnnS8roJk-=~>
zzA|&h>-WnKy*6>b?wj%xh5fs}{p%a(p@^LhRS1GaEgbFdN-=K#-N*-#Sn=Mohg)Dg
zC$p(bro9q#9GgL|XfQAL;3ZbP7@+dNtmns4&a>A+7Zc#aRb2|A2oQ-JlBM+AhvE~0
z{(rF^CV#K&D?4wh#pPhYX6H*c{~Ni1K(fFw#1Y2N(-}e~&Mz%xTzw^9{99I(d3|aB
z265bxT##~_?|CBlBsrRIXj7eoQjfZ!6jn9d?MOviF&5R&aJ32<Y8>9rO~^V{*2IEl
zV~A5hylUGfIHwaEaS}&x$qd?#!>X+URWCsBApviJSe5`a8`N}n{p{NTr{A(96m;k?
zhH#C>iYIgb@8C#ariGp1^pYsg-21ZFGXDkk`U9@c!grt_+#^fLe#@vB<Ghmd*=ud3
zD8A%pdPc?f*=4i%cMJ#HUpUgeMu*q$D3pF?2ZI;8#-e_g=HjgGbgCk8)am@XO+Cug
znL^UYX3NisK2WUq&1zvC-`FEE@boJ<EbPB>EQ%{|1J1&Eo{d<E5xBR^6h$m;;KoFZ
zQqeaC*KLf)fA@jgVS3@Xko|lf<-BA93E_?ZwB=#*=Fv$Z2bDC5lS9Y{#r~Hbk<cto
zRsC*`l*_$Zcjml6w^c#d;Vb19Fh2WzX8Ow(v1`3eiD}CY8m?GsJTOHv`Ts2%O*Zf}
zU+hN@J5MM!r{d@sxl;*GG(5|)DVstK>ohqf*+Ea3U&4*UCg}*uG$$`)IX-y{QBinz
znMM4X|EdAaj|SfxxP&PLL?)P&&!c)hh$xJ!!&Zx&G&L6Fp`cI;>g0ra(s}1z@flEY
zXqv$Hl*#x14fOa+P*G9e=7f!k$BMt?>j;NwI}~p>Y}I^a)htc^FU+GrX9Nn4X)ROa
zO|TKVz~D{ED~iTfZBEZdWjitykWy&R<TqjCElK%5G(PSJaB;t~$x?Kew=w+oqz_SL
z#hf&0GI&lh5@>G(b^%X~^Oq<#zeWh%=UJo(Yn?D!uBsRm6-H~*-2a%RjvW13l{6o8
zN9eI<Bu1{-wA&xFUX2)sg<Q<#*A<iVKY_2~wZ9llyId0!llMF7zG#VKd?v7sK7{tS
zg?Z4wElbpS_A>!dV0;JJpZ~(w(!C=fwB@II%2(1X(m3FTV2k=yfyqk=hJlSc;yUFn
z!cyh~pJofA);MU0k3w<$Ou9y7(qiPfp1MGdmP4`Pun1VDv7o{KXD;RsRJ2k)x#C~I
zqSS)1;;HxvLI4Hhw&6>vk#l+(A2~vko#if^t+4htKcLt7q%r#g=e^<0^nai2RTtH}
zii#U3US4(3NPmH@63bq_lfNA{6a+t;F<s?wbxG6L>ACBj@kdQ&-hA@`p;FI{r+i{?
z6bwH6kAgi-SxSVFF~t95{Ae(lo=KuEMWkWH!qSFfCx=JspLeZrXWDO@e?ddn2NXT;
zV|5Dj<yBR?v$L}Pwg03|O(Ego!-r>dXpJn_RTU`yq~{$X`#wY}iEEi+0rv!bDd<u2
zb%*>-z}nN({ekDM(ND>9@jnT<(aXy{0Gb$xzOP0w{L{__5rue_ofi`;=+5K<VT8pb
zKU&TBMq#0+*<VxU>n4L#W>eEHLbK@*3;p-P?}&LCx%L9R%<MQUp#};+L5KgW8$F_c
z;W#<fB8DxN;d*Q*da^q@D|zNN4c{L+@J5cVR^$h~)9M&p@6{2H1<8uI^DM^JZy~Nn
zY`mCI1Y-c@4`Y11wUU$s3f&6Sl$02=-_)y|uZUi^Ne7lbX=LE{-$NCi{)T4=5F*jf
z1wbH5fF-Y9EA2)ay<$T>B%#3QgW+y*adC5n#SE7&c3PPLzn<C%Lb{*RzMISv66pec
zq88+K=$Q42!o)8wC<ya$zHZByO;+aGZ#i~eF|&zE2XwQ<z9{bMD2v1isOTblGO*%_
z)BdMIM%^P3tY-1XzaIi17z^w1pJxKtwrw5^O-!V(J!J|F$2Yi1hIbzZAV6jdgV!it
z-CDO5Wr;o*=>bj=Cg@UfK_|cnZ5`<TJ<EdQRDqC&d62CB_V@D9Z0zi2z(YP}c-)0X
zb=B1=ARAV%o_6aSa5K3_9A^pYg6z^#lDa7kU>hOpn({{WFQK4=;JVb8vcmII`IqPK
zmnr;=pyRUE(@Uko3o}phm836&+4-J7k?8rQsw&SO7ic5a8WvXf+u7Rkfo{(t$-BW=
zfP=$sa;XOgMu$1MxL}4heQbDGXs&ihOmWEnn!Fqmde>h$xw(abLqg9@i_;R8!HhIa
z<xVJPr7v%OPkms@e-HL<J_!i0G#J;iNILP8hKcyHl9Fd?ssEbBhtcykN!ICO#w=|^
zL&>k>d6jTcL+c*tKzvz)<txj`4M`~ud80q@BYHyl<V-uD&)zSuF9Cdc%csHm(WBw&
zg%;CqZZO}Q<mus&CY;_oS?}#Sxdc=B=;?ebAcIL2a>3kT0!lmJLvD0Ogz0;tf&W^>
zG!0Y!2k#rT3*Pn>3C!*sEiW%qR!LsJVoKST0;7IBONZO@lg_AX^U|KX4-Y|2Wocuh
zn1^EEl96FBHC_r%?GsDeu!G5!QRda!56zaBp#`#a<d<(y#qC(rKL4(?)J0`tZEfvR
z^SVbnr>7__Xv&BQF0=Xt+`OJLc)o~S$@m~w9}bcKZfT(eb)o<^@W$|X1#J@(O;1k|
z3C|5CP_17%FyBA&oAj7)&>pw;_m=`^RkQ1Y2}lu^&RHetETZSFt*y4Sos+zot_hK&
zT5^`cfTewXCD4Orc4EyiUBNXL03U#-X8liPWfF`LpbfKo;eBHxJxD;zD772kR+N_m
zfd}1+XWgvS$HQY|QQZp&9v&x^lMK6D*O+`i*az4aRs3`jr<kK@JA3<VFv^pSx_+cT
z(?YF<ZET`<&-``@>;ViG_H=cLuUyZM>;DYn;6~iZdhi^9f_FN=$D<#*;0crH3ro4t
zoo?ULIxPQf(PP5*BZi=D{Y-I4uBXlOw;MZe$sKohcee~r_1+7o&49|mHhR(Lfuw;G
zACcY|N2H_kcT&!DO-CGVdCbHt`z&7B$IN{+Y0!*&UENwL538(5)=-mq3nL_c7iYt1
z?v~AiJMzVC+b3cXt|~}hQe@Xzny)LD&nts>S}Fg~PgKey;!^%G%o`#`GDQJH8nwUo
zO>O6%TfD2_XFDZ1x0%bX35E-WP1<=aru;c>;!<vAT-|?I8jf6~k*lIH0xtCN&&*4H
z{AiE4P|ZEMRl7mkdrVD2Je4Q=?xv+4y;#&zvr#~vMzrRF+PjChZ^bDQs4vyk*T2-v
zlAc@23Jf=+d?K7KQs^2GKtE=`udv`|6j8XUBkghe!^ksc<LBQ~>2hAi5S5b3^qw;L
zFlEM@DqCaNHZn5uV{db&Nl}uXA&r$WPU+60;gSVpCsXKkoxY;`oaDsOUnx}zo#S2I
zCeZg`^5^I@(V^q{dp3co&C#cxlxn%();)^O7-784?q2_MgQPIyqKN)1852E^$+G$3
zr#FyHbL$gT_f`!@8y5dzh)93)VZ!(x;{D0s(~4Qqz8}JtTmC{R7!Kr()^#XT{;(&4
z;vNe1K^l1>a)>#DB#pdh@L-P0yOLsJK}^bwS=U6=NefqFPWuPeJqDf#I-X;me<+kt
z3{hTtK4F4BEm9+I`(P$8*5Mu@A>nnru<L{oHl<BPQyaV~YPZsR%^~Ql`LsVX)gyN^
zI^cqysYTso+V<@?*8=x`Cpu+X=`DNI9qx^qC?*w*uHMrglMxmXQw%xRncW|K9l6|?
zUr+$+|H{y2Lb$K5&)LmQFy4Pkv5xWTtuJbJewL4Qbl$PyO*i;P936Q()mZ7Vkw^cS
zo~9Q&QY+D=^)}#4?w*6^h{$~&ZvFZ5eHFw!OMa^C2Os0w2R07|?1kZ9o88Xf9OEL!
z_|CTvKF&gw$5>UfE>^oP#Du8j$$4-U$#jH;nw4<2u)uCQ{nMSeu=z*AU7vtO;G4bh
z_W7Cb41;g9!Uskt-?MZpz3Kc9X8{a_s|fnoSkAmpzkg7f;hL9Fo*%CkP54J*@T-$9
zxKLF(oJji$Qp>|1>8$phrT&J6p^X3R(i1_r$<XL%e9@Q?Ew)e*r|H?c%|@D<v(ZzH
z%*!EHA?xE@Zcn*ROnl#z_y6Pnq_k=aE#N*e;86mW-lwr;g{!Zvj=Dx+N*2S^#>VEC
z&!b28ZzuZgE$$p~9MD96-5S!AQQxh(oruSuN?@rM{+><O#3W+1mUMWV85A8A6+(~b
z;(vS8y|$u)!%f_q<cOv&Hcl`imp#VUD@T94dGiL5sS!K14~;{wqdz$Z=T$QS_o^no
zo}WuMo?_1<neHN)9^SqIXbGeF>sOqZm>3vPHJuRN-{1F}h;bdRy$^CAEl%=xtKtNx
z=MED;CMVH*4b-}vJ5CFru}Dg_NP5nR$~Yw^48!n*mL^3#W8=0ze@YL|XJ;`ASAUL`
z*;sVf`|NccoWtIKbEivG`HIg{w!bkY0j5GrOY7B*h`Zuuq*ylg1%#iOYX^mXRL7Ca
zK5JdKeC}MH!Y<{!hh$<Rx!O_s(x)M2ecb`HKT|8`xSm!%4bxw&Tx)0_KBBAfv$efi
zxZ3J&^vgwT=J?@loLiURu@XOO9VQgSGhy9eh0=0cO7q{Xe;U3G2h^#*)7l#0c!6a!
z>OZ9fLtcS_42dc@FW%3BvoYNM%!UW1KPTE(5EHjtV!;cSf{2S8Z6JQ7U^XGFUqFuW
zcHqg|p_R#%^Zx#ROpF)r)eQF`MS`>4v*1=6lH*QBot{dHXZC%|Oq-v#;0b*duf+6i
zbx-SdqSnWUv!a9#J|;M&yVSf?$@$q}KM`KII(_UX9XmB;R7As`)ek)@@XVv%Vv>?y
zMn=d-Mn{Qb+c6C<SKb{vJM+hA9?H}7Hxd&OQ$U#O5DtWp5=P9468_v41r5~p$`Gp(
zGL!+N7r7_!pci9|n(FH2?QK4l91=lwR6MIDOW`Icw6Dv)ef-&bN+sOv4w9*8n<xU+
z$iMDE17R=ntGR`YN!3KyYAvt&(CY62!E4}ESu%n9F^`Z;P8qOGc%j$95)Xbqqs2w<
zrve1XcA1OQ5QPtO6DFUBK{<sIneoHMDCWd}rGP#BbXe!{sC!DT*Z?01@n*_(DW$x;
zJSCQR3~rju7N?CeMCct(@_TG#vhatK9M#r#8Lo70ZmwXp*Sd_C0OcF#<rV2ADNUwi
zI0`-;k!!2<=mwm^V)dmQLRVM!j&zn{g+5%bE*RhR^qZoVi-Rd*j7WtKjRzI_z4+u{
z8oUJGSNGnX<vd{LNx*3C?7YgPOkFz^-@n8v4XX${taUkdpmM^23=aoP&AcSiqwZBr
zjo8}8MxUh&k_nz%@qN{Ob9lvV3&alZuJc8KkbR=Ef$w&>+2Mvh#+?RZMJ5sevj05(
zd<ZzH;+1FJl)~*q*%DnjHZtm?Cp&;jD%`>5*<J)c5q$CC!-qJlMgqY~huF!D?v289
zT1H|%SQ)l(v)RK#%4REU#|h#018l+~A_FUK;cfTS=E0S8_4KxlKVWsbzdHN+HOwhJ
zPcuvQ3MPsd9zO9cGaURwjU|L3;D7B0<KDRZ7kWo59u`OdPwsi~{r84k>76lrt-uA9
z?+b-zEprSMiI~<}x4%DXZP{W&g^dw884I)-fKAGo^2!&J4<tB!X|SpM#M1Y;BH8QE
z!wuiJxLe5ldv%xTWQvdzN66}{3rpJ)4;0D(J#+t4j#+neGxL=!)0EAT_pM^DPdu6O
z;kbPLk6qsDfM;unxPE+hsO50tp4#7Hz%!A|<#=o=td>RSBo!83jJA`As1fzOtk5FC
zCGoVVa@lc@FA)VFv@bFz-`-}v?8vrUhUvYJYb!oJVhvg;0?Y%2H8njcMxCv#tTjJA
zFfPk<y1Zp#vi_x)COW4-z5B;7{UOjzXU?DRfAi=>+U5=yCFj$@D@A0EJDSb@daM+r
zW!-DO)<uPEca3Myo~<Fm#n1pOPTmVX-MrcOzKy~EZ-2kvul{XkZ|`n2(`9YgXWLqn
z?qroJ6Ba>Rjh!)mIXO8TiY+Uyzjm7TH2UhsEn7^A?w&i>Hv?$aihWsIwM=}4yB2Y{
z2DSiK^8-h2-+Wsb;8FAY?e=}CZHXL_(~j?3b@kPaUsAno7fjCH4!djutRDA%UA3w!
zY2yp+&rf$|n@Bl@hOYnf@83RPDFS3J%)E7@immJDqJYp)jV>i%c?3-0PSa8oBd_)L
zP3+L~bSlVl0``-+*!m0=m6R48NH7=#)DQbUgYzbDLl?dU9%aYi>FVdQ&MBb@0G~HB
AoB#j-


From 419e13c3e04af91d2ba158ec311ed23603185f37 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Sun, 12 Aug 2018 01:02:47 -0700
Subject: [PATCH 0551/2720] allow mtf.slicewise with no gradient.  Fix bug in
 mtf.Print.

PiperOrigin-RevId: 208376978
---
 .../mesh_tensorflow/mesh_tensorflow.py        | 18 ++++--
 tensor2tensor/mesh_tensorflow/mtf_layers.py   |  2 +
 .../mesh_tensorflow/mtf_transformer.py        | 58 ++++++++++++++++---
 .../mesh_tensorflow/simd_mesh_impl.py         |  3 +
 4 files changed, 70 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index 7c31ddf7a..011083762 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -1119,6 +1119,7 @@ def __init__(self,
                output_shape,
                output_dtype,
                splittable_dims,
+               has_gradient=True,
                grad_function=None,
                name=None):
     """Create a SlicewiseOperation.
@@ -1135,15 +1136,21 @@ def _square_grad(op, dy):
       output_shape: a Shape
       output_dtype: a dtype
       splittable_dims: a list of Dimensions which are ok to split
-      grad_function: an optional python function
+      has_gradient: a boolean
+      grad_function: an optional python function. Default to using tf.gradients
       name: an optional string
     """
     super(SlicewiseOperation, self).__init__(inputs, name=name or "slicewise")
     self._tf_fn = tf_fn
     self._outputs = [Tensor(self, output_shape, output_dtype)]
     self._splittable_dims = splittable_dims
+    self._has_gradient = has_gradient
     self._grad_function = grad_function
 
+  @property
+  def has_gradient(self):
+    return self._has_gradient
+
   def gradient(self, grad_ys):
     if self._grad_function is not None:
       return self._grad_function(self, grad_ys[0])
@@ -1168,6 +1175,7 @@ def slicewise(tf_fn,
               output_shape=None,
               output_dtype=None,
               splittable_dims=None,
+              has_gradient=True,
               grad_function=None,
               name=None):
   """Slice-wise call to any tensorflow function.
@@ -1182,7 +1190,8 @@ def slicewise(tf_fn,
     output_shape: a Shape
     output_dtype: a dtype
     splittable_dims: a list of Dimensions which are ok to split
-    grad_function: an optional gradients function
+    has_gradient: a boolean
+    grad_function: an optional gradients function.  If None, use tf gradient.
     name: an optional string
 
   Returns:
@@ -1194,6 +1203,7 @@ def slicewise(tf_fn,
       convert_to_shape(output_shape) or xs[0].shape,
       output_dtype or xs[0].dtype,
       splittable_dims,
+      has_gradient,
       grad_function,
       name=name).outputs[0]
 
@@ -2253,8 +2263,8 @@ def lower(self, lowering):
         self.outputs[0],
         lowering.mesh_impl(self).Print(
             lowering.tensors[self.inputs[0]],
-            [lowering.tensors[d] for d in self._data], self._message,
-            **self._kwargs))
+            [lowering.tensors[d].to_laid_out_tensor() for d in self._data],
+            self._message, **self._kwargs))
 
   def gradient(self, grad_ys):
     return grad_ys
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers.py b/tensor2tensor/mesh_tensorflow/mtf_layers.py
index bb23a800c..effbfc553 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers.py
@@ -601,6 +601,8 @@ def attention_bias_local_block(mesh, block_length, memory_length,
   return mask
 
 
+
+
 def moe_v0(inputs,
            hidden_dim,
            output_dim,
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
index 119eecbb0..dc25e1f92 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
@@ -698,9 +698,33 @@ def mtf_transformer_tiny_8gpu():
   return hparams
 
 
-def mtf_transformer_paper_lm(sz):
-  """Config for language-model experiments."""
-  n = 2 ** sz
+def mtf_transformer_paper_lm(size):
+  """Config for language-model experiments.
+
+  Train these on languagemodel_lm1b32k_packed for 136000 steps (10 epochs)
+
+  The size parameter is an integer that controls the number of heads and the
+  size of the size of the feedforward hidden layers.  Increasing size by 1
+  doubles each of these.
+
+  Results:
+  size   params/10^9  log-ppl(per-token)
+  -1     0.14         3.209
+  0      0.22         3.119
+  1      0.37         3.037
+  2      0.67         2.969
+  3      1.28         2.912
+  4      2.48         2.874
+  5      4.90         2.871
+
+  (to get word-level log-ppl, multiply by 1.1078)
+
+  Args:
+    size: an integer
+  Returns:
+    a hparams object
+  """
+  n = 2 ** size
   hparams = mtf_transformer_base()
   hparams.label_smoothing = 0.0
   hparams.batch_size = 256
@@ -763,9 +787,21 @@ def mtf_transformer_paper_lm_5():
   return hparams
 
 
-def mtf_transformer_paper_tr(sz):
-  """Config for translation experiments."""
-  n = 2 ** sz
+def mtf_transformer_paper_tr(size):
+  """Config for translation experiments.
+
+  Train these on translate_enfr_wmt32k_packed for 154000 steps (3 epochs)
+
+  The size parameter is an integer that controls the number of heads and the
+  size of the size of the feedforward hidden layers.  Increasing size by 1
+  doubles each of these.
+
+  Args:
+    size: an integer
+  Returns:
+    a hparams object
+  """
+  n = 2 ** size
   hparams = mtf_transformer_base()
   hparams.label_smoothing = 0.1
   hparams.batch_size = 128
@@ -871,7 +907,9 @@ def mtf_transformer_lm_baseline():
   """Small language model to run on 1 TPU.
 
   Run this on 2x2 on languagemodel_lm1b32k_packed for 272000 steps (10 epochs)
-  140M params
+  Results:
+         params/10^9  log-ppl(per-token)
+         0.14         3.202
 
   Returns:
     a hparams
@@ -890,6 +928,10 @@ def mtf_transformer_lm_moe():
   Run this on 2x2 on languagemodel_lm1b32k_packed for 272000 steps (10 epochs)
   900M params.
 
+  Results on LM1B:
+         params/10^9  log-ppl(per-token)
+         0.90         3.002
+
   Returns:
     a hparams
   """
@@ -898,3 +940,5 @@ def mtf_transformer_lm_moe():
   hparams.layout = "batch:all;experts:all"
   hparams.feedforward_layer = "moe"
   return hparams
+
+
diff --git a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
index 3033e53d3..b689e5c7f 100644
--- a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
@@ -220,6 +220,9 @@ def allconcat(self, x, mesh_axis, concat_axis, stack=False):
   def alltoall(self, x, mesh_axis, split_axis, concat_axis):
     """Grouped alltoall (like MPI alltoall with splitting and concatenation).
 
+    TODO(noam): this is a terribly inefficient implementation using allreduce.
+    Replace this with a native xla alltoall once it is ready.
+
     Args:
       x: a LaidOutTensor
       mesh_axis: an integer the mesh axis along which to group

From 0e2b974511848de4cdd029057057313e3266df48 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Sun, 12 Aug 2018 21:14:28 -0700
Subject: [PATCH 0552/2720] change the default for
 reward_prediction_stop_gradient to True.

PiperOrigin-RevId: 208426661
---
 tensor2tensor/models/research/next_frame_params.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index b8769671b..dd0559818 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -73,7 +73,7 @@ def next_frame_stochastic():
   hparams.video_modality_loss_cutoff = 0.0
   hparams.add_hparam("stochastic_model", True)
   hparams.add_hparam("reward_prediction", True)
-  hparams.add_hparam("reward_prediction_stop_gradient", False)
+  hparams.add_hparam("reward_prediction_stop_gradient", True)
   hparams.add_hparam("model_options", "CDNA")
   hparams.add_hparam("num_masks", 10)
   hparams.add_hparam("latent_channels", 1)

From 408c48c602bb36100ab50cd41735c25cbf2c9ee7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sun, 12 Aug 2018 21:26:47 -0700
Subject: [PATCH 0553/2720] Internal change

PiperOrigin-RevId: 208427326
---
 tensor2tensor/bin/t2t_trainer.py   |  3 ++-
 tensor2tensor/utils/trainer_lib.py | 15 ++++++++++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index c3e699c17..685979cb0 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -236,7 +236,8 @@ def create_run_config(hp):
       inter_op_parallelism_threads=FLAGS.inter_op_parallelism_threads,
       log_step_count_steps=FLAGS.log_step_count_steps,
       intra_op_parallelism_threads=FLAGS.intra_op_parallelism_threads,
-      tpu_config_extra_kwargs=tpu_config_extra_kwargs)
+      tpu_config_extra_kwargs=tpu_config_extra_kwargs,
+      cloud_tpu_name=FLAGS.cloud_tpu_name)
 
 
 def generate_data():
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 3c082a3d1..9b20ef57f 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -129,7 +129,8 @@ def create_run_config(master="",
                       inter_op_parallelism_threads=0,
                       log_step_count_steps=100,
                       intra_op_parallelism_threads=0,
-                      tpu_config_extra_kwargs=None):
+                      tpu_config_extra_kwargs=None,
+                      cloud_tpu_name=""):
   """Create RunConfig, TPUConfig, and Parallelism object."""
   session_config = create_session_config(
       log_device_placement=log_device_placement,
@@ -176,6 +177,14 @@ def create_run_config(master="",
       run_config_args["master"] = os.environ[
           "KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS"]
       run_config_args["evaluation_master"] = run_config_args["master"]
+    elif not master and cloud_tpu_name:
+      # Update run_config to use cluster instead of master/evaluation_master
+      # as we need the cluster spec to use Cloud Pods
+      tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
+          cloud_tpu_name)
+      run_config_args["cluster"] = tpu_cluster_resolver
+      del run_config_args["master"]
+      del run_config_args["evaluation_master"]
 
   config = run_config_cls(**run_config_args)
   config.warm_start_from = warm_start_from
@@ -471,6 +480,10 @@ def create_experiment(
       plateau_delta=eval_early_stopping_metric_delta,
       every_n_steps=min_eval_frequency)
 
+  # Eval on TPU Pods is not supported yet
+  if use_tpu and run_config.tpu_config.num_shards > 8 and "eval" in schedule:
+    raise ValueError("Eval is not currently supported on a TPU Pod")
+
   # In-process eval (and possible early stopping)
   if schedule == "continuous_train_and_eval" and min_eval_frequency:
     tf.logging.warn("ValidationMonitor only works with "

From 4daee8581c8a089062ef60742d9db6e18d90316c Mon Sep 17 00:00:00 2001
From: me-sh <me-sh@users.noreply.github.com>
Date: Mon, 13 Aug 2018 10:28:51 -0400
Subject: [PATCH 0554/2720] fix unexpected keyword 'drop_remainder' on TPU

---
 tensor2tensor/data_generators/problem.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 107409c5c..e1884ff98 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -866,7 +866,7 @@ def define_shapes(example):
         # on TPU, we use params["batch_size"], which specifies the number of
         # examples across all datashards
         batch_size = params["batch_size"]
-        dataset = dataset.batch(batch_size, drop_remainder=True)
+        dataset = dataset.batch(batch_size)
       else:
         num_shards = config.data_parallelism.n if config else 1
         batch_size = hparams.batch_size * num_shards

From 5b278c1f78b8938fe4e15e5b8e3db4c45a406e1a Mon Sep 17 00:00:00 2001
From: me-sh <me-sh@users.noreply.github.com>
Date: Mon, 13 Aug 2018 10:29:44 -0400
Subject: [PATCH 0555/2720] create problem transformer_timeseries_tpu

---
 tensor2tensor/models/transformer.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 9af54dc90..c03dd6a9a 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1953,6 +1953,13 @@ def transformer_tpu():
   update_hparams_for_tpu(hparams)
   return hparams
 
+@registry.register_hparams
+def transformer_timeseries_tpu():
+  """HParams for running Transformer model on timeseries on TPU."""
+  hparams = transformer_timeseries()
+  update_hparams_for_tpu(hparams)
+  hparams.batch_size = 256 # revert to value set in transformer_timeseries
+  return hparams
 
 @registry.register_hparams
 def transformer_tpu_bf16_activation():

From e518cd6bbd416022f5e041a242db3dfd238ff08e Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 13 Aug 2018 09:31:53 -0700
Subject: [PATCH 0556/2720] Rm data_reader.bucket_by_sequence_length. Use
 tf.contrib.data.bucket_by_sequence_length

PiperOrigin-RevId: 208492918
---
 tensor2tensor/utils/data_reader.py      | 53 -------------------------
 tensor2tensor/utils/data_reader_test.py | 41 -------------------
 2 files changed, 94 deletions(-)

diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index fa9facb64..14bd04624 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -16,7 +16,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-import numpy as np
 
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
@@ -52,58 +51,6 @@ def example_valid_size(example, min_length, max_length):
   )
 
 
-def bucket_by_sequence_length(dataset,
-                              example_length_fn,
-                              bucket_boundaries,
-                              bucket_batch_sizes,
-                              padded_shapes=None):
-  """Bucket entries in dataset by length.
-
-  Args:
-    dataset: Dataset of dict<feature name, Tensor>.
-    example_length_fn: function from example to int, determines the length of
-      the example, which will determine the bucket it goes into.
-    bucket_boundaries: list<int>, boundaries of the buckets.
-    bucket_batch_sizes: list<int>, batch size per bucket.
-    padded_shapes: dict<feature name, list<int>>, optional, shapes of the
-      features with None where feature should be padded to max in that dim.
-
-  Returns:
-    Dataset of padded and batched examples.
-  """
-  with tf.name_scope("bucket_by_seq_length"):
-
-    def example_to_bucket_id(example):
-      """Return int64 id of the length bucket for this example."""
-      seq_length = example_length_fn(example)
-
-      boundaries = list(bucket_boundaries)
-      buckets_min = [np.iinfo(np.int32).min] + boundaries
-      buckets_max = boundaries + [np.iinfo(np.int32).max]
-      conditions_c = tf.logical_and(
-          tf.less_equal(buckets_min, seq_length),
-          tf.less(seq_length, buckets_max))
-      bucket_id = tf.reduce_min(tf.where(conditions_c))
-
-      return bucket_id
-
-    def window_size_fn(bucket_id):
-      # window size = batch size
-      batch_sizes = tf.constant(bucket_batch_sizes, dtype=tf.int64)
-      window_size = batch_sizes[bucket_id]
-      return window_size
-
-    def batching_fn(bucket_id, grouped_dataset):
-      batch_sizes = tf.constant(bucket_batch_sizes, dtype=tf.int64)
-      batch_size = batch_sizes[bucket_id]
-      return padded_batch(grouped_dataset, batch_size, padded_shapes)
-
-    dataset = dataset.apply(
-        tf.contrib.data.group_by_window(example_to_bucket_id, batching_fn, None,
-                                        window_size_fn))
-    return dataset
-
-
 def padded_batch(dataset, batch_size, padded_shapes=None):
   padded_shapes = padded_shapes or dict(
       [(name, [None] * len(shape))
diff --git a/tensor2tensor/utils/data_reader_test.py b/tensor2tensor/utils/data_reader_test.py
index a677eec98..6b3a0f313 100644
--- a/tensor2tensor/utils/data_reader_test.py
+++ b/tensor2tensor/utils/data_reader_test.py
@@ -205,47 +205,6 @@ def testBatchingSchemeBuckets(self):
     self.assertEqual([max(1, bs // 2)
                       for bs in expected_batch_sizes], batch_sizes)
 
-  def testBucketBySeqLength(self):
-
-    def example_len(ex):
-      return tf.shape(ex["inputs"])[0]
-
-    boundaries = [10, 20, 30]
-    batch_sizes = [10, 8, 4, 2]
-
-    dataset = self.problem.dataset(
-        tf.estimator.ModeKeys.TRAIN,
-        data_dir=self.data_dir,
-        shuffle_files=False)
-    dataset = data_reader.bucket_by_sequence_length(
-        dataset, example_len, boundaries, batch_sizes)
-    batch = dataset.make_one_shot_iterator().get_next()
-
-    input_vals = []
-    obs_batch_sizes = []
-    with tf.train.MonitoredSession() as sess:
-      # Until OutOfRangeError
-      while True:
-        batch_val = sess.run(batch)
-        batch_inputs = batch_val["inputs"]
-        batch_size, max_len = batch_inputs.shape
-        obs_batch_sizes.append(batch_size)
-        for inputs in batch_inputs:
-          input_val = inputs[0]
-          input_vals.append(input_val)
-          # The inputs were constructed such that they were repeated value+1
-          # times (i.e. if the inputs value is 7, the example has 7 repeated 8
-          # times).
-          repeat = input_val + 1
-          # Check padding
-          self.assertAllEqual([input_val] * repeat + [0] * (max_len - repeat),
-                              inputs)
-
-    # Check that all inputs came through
-    self.assertEqual(list(range(30)), sorted(input_vals))
-    # Check that we saw variable batch size
-    self.assertTrue(len(set(obs_batch_sizes)) > 1)
-
 
 if __name__ == "__main__":
   tf.test.main()

From 348ca2412ea418f7c71e1aaac3053bcefcee205a Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Mon, 13 Aug 2018 09:33:59 -0700
Subject: [PATCH 0557/2720] Undo bad change to slicewise operations in
 mesh-tensorflow - mixture of experts seems to work again.

PiperOrigin-RevId: 208493241
---
 tensor2tensor/mesh_tensorflow/mesh_tensorflow.py | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index 011083762..f3e5bdde7 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -1119,7 +1119,6 @@ def __init__(self,
                output_shape,
                output_dtype,
                splittable_dims,
-               has_gradient=True,
                grad_function=None,
                name=None):
     """Create a SlicewiseOperation.
@@ -1136,7 +1135,6 @@ def _square_grad(op, dy):
       output_shape: a Shape
       output_dtype: a dtype
       splittable_dims: a list of Dimensions which are ok to split
-      has_gradient: a boolean
       grad_function: an optional python function. Default to using tf.gradients
       name: an optional string
     """
@@ -1144,13 +1142,8 @@ def _square_grad(op, dy):
     self._tf_fn = tf_fn
     self._outputs = [Tensor(self, output_shape, output_dtype)]
     self._splittable_dims = splittable_dims
-    self._has_gradient = has_gradient
     self._grad_function = grad_function
 
-  @property
-  def has_gradient(self):
-    return self._has_gradient
-
   def gradient(self, grad_ys):
     if self._grad_function is not None:
       return self._grad_function(self, grad_ys[0])
@@ -1175,7 +1168,6 @@ def slicewise(tf_fn,
               output_shape=None,
               output_dtype=None,
               splittable_dims=None,
-              has_gradient=True,
               grad_function=None,
               name=None):
   """Slice-wise call to any tensorflow function.
@@ -1190,7 +1182,6 @@ def slicewise(tf_fn,
     output_shape: a Shape
     output_dtype: a dtype
     splittable_dims: a list of Dimensions which are ok to split
-    has_gradient: a boolean
     grad_function: an optional gradients function.  If None, use tf gradient.
     name: an optional string
 
@@ -1203,7 +1194,6 @@ def slicewise(tf_fn,
       convert_to_shape(output_shape) or xs[0].shape,
       output_dtype or xs[0].dtype,
       splittable_dims,
-      has_gradient,
       grad_function,
       name=name).outputs[0]
 

From 7742afcc2b95a11e2d0d7364f06e689434d173b9 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 13 Aug 2018 11:16:46 -0700
Subject: [PATCH 0558/2720] weight normalization

PiperOrigin-RevId: 208511842
---
 tensor2tensor/layers/common_layers.py | 124 ++++++++++++++++++++++++++
 1 file changed, 124 insertions(+)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 5845b15c9..5ac71bcca 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3795,3 +3795,127 @@ def sparse_eye(size):
       values=values,
       dense_shape=dense_shape
   )
+
+
+# modification from https://github.com/tensorflow/tensorflow/pull/21276
+# without special initialization for g
+class WeightNorm(tf.keras.layers.Wrapper):
+  """ This wrapper reparameterizes a layer by decoupling the weight's
+  magnitude and direction. This speeds up convergence by improving the
+  conditioning of the optimization problem.
+
+  Weight Normalization: A Simple Reparameterization to Accelerate
+  Training of Deep Neural Networks: https://arxiv.org/abs/1602.07868
+  Tim Salimans, Diederik P. Kingma (2016)
+
+  WeightNorm wrapper works for keras and tf layers.
+
+  ```python
+    net = WeightNorm(tf.keras.layers.Conv2D(2, 2, activation='relu'),
+           input_shape=(32, 32, 3), data_init=True)(x)
+    net = WeightNorm(tf.keras.layers.Conv2D(16, 5, activation='relu'),
+                     data_init=True)
+    net = WeightNorm(tf.keras.layers.Dense(120, activation='relu'),
+                     data_init=True)(net)
+    net = WeightNorm(tf.keras.layers.Dense(n_classes),
+                     data_init=True)(net)
+  ```
+
+  Arguments:
+    layer: a layer instance.
+    data_init: If `True` use data dependent variable initialization
+
+  Raises:
+    ValueError: If not initialized with a `Layer` instance.
+    ValueError: If `Layer` does not contain a `kernel` of weights
+    NotImplementedError: If `data_init` is True and running graph execution
+  """
+
+  def __init__(self, layer, data_init=False, **kwargs):
+    if not isinstance(layer, tf.keras.layers.Layer):
+      raise ValueError(
+          "Please initialize `WeightNorm` layer with a "
+          "`Layer` instance. You passed: {input}".format(input=layer))
+
+    super(WeightNorm, self).__init__(layer, **kwargs)
+    self._track_checkpointable(layer, name="layer")
+
+  def _compute_weights(self):
+    """Generate weights with normalization."""
+    with tf.variable_scope("compute_weights"):
+      self.layer.kernel = tf.nn.l2_normalize(
+          self.layer.v, axis=self.norm_axes) * self.layer.g
+
+  def _init_norm(self, weights):
+    """Set the norm of the weight vector."""
+    with tf.variable_scope("init_norm"):
+      flat = tf.reshape(weights, [-1, self.layer_depth])
+      return tf.reshape(tf.norm(flat, axis=0), (self.layer_depth,))
+
+  def _data_dep_init(self, inputs):
+    """Data dependent initialization for eager execution."""
+
+    with tf.variable_scope("data_dep_init"):
+      # Generate data dependent init values
+      activation = self.layer.activation
+      self.layer.activation = None
+      x_init = self.layer.call(inputs)
+      m_init, v_init = tf.moments(x_init, self.norm_axes)
+      scale_init = 1. / tf.sqrt(v_init + 1e-10)
+
+    # Assign data dependent init values
+    self.layer.g = self.layer.g * scale_init
+    self.layer.bias = (-m_init * scale_init)
+    self.layer.activation = activation
+    self.initialized = True
+
+  def build(self, input_shape):
+    """Build `Layer`."""
+    input_shape = tf.TensorShape(input_shape).as_list()
+    self.input_spec = tf.layers.InputSpec(shape=input_shape)
+
+    if not self.layer.built:
+      self.layer.build(input_shape)
+      self.layer.built = False
+
+      if not hasattr(self.layer, "kernel"):
+        raise ValueError(
+            "`WeightNorm` must wrap a layer that"
+            " contains a `kernel` for weights"
+        )
+
+      # The kernel's filter or unit dimension is -1
+      self.layer_depth = int(self.layer.kernel.shape[-1])
+      self.norm_axes = list(range(self.layer.kernel.shape.ndims - 1))
+
+      self.layer.v = self.layer.kernel
+      self.layer.g = self.layer.add_variable(
+          name="g",
+          shape=(self.layer_depth,),
+          initializer=tf.ones_initializer,
+          dtype=self.layer.kernel.dtype,
+          trainable=True)
+
+      # with ops.control_dependencies([self.layer.g.assign(
+      #     self._init_norm(self.layer.v))]):
+      #   self._compute_weights()
+      self._compute_weights()
+
+      self.layer.built = True
+
+    super(WeightNorm, self).build()
+    self.built = True
+
+  def call(self, inputs):
+    """Call `Layer`."""
+    # if context.executing_eagerly():
+    #   if not self.initialized:
+    #     self._data_dep_init(inputs)
+    self._compute_weights()  # Recompute weights for each forward pass
+
+    output = self.layer.call(inputs)
+    return output
+
+  def compute_output_shape(self, input_shape):
+    return tf.TensorShape(
+        self.layer.compute_output_shape(input_shape).as_list())

From 539a32b656d40da8bd5850f4f737c55a64cffe0d Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 13 Aug 2018 11:56:19 -0700
Subject: [PATCH 0559/2720] Allow name as kwarg and test pos. embeddings,
 make_edge_vectors.

PiperOrigin-RevId: 208518880
---
 tensor2tensor/layers/common_attention.py      | 91 +++++++++----------
 tensor2tensor/layers/common_attention_test.py | 32 ++++++-
 2 files changed, 75 insertions(+), 48 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 672b5d79a..c1ac1973f 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -645,70 +645,67 @@ def add_timing_signal_nd(x, min_timescale=1.0, max_timescale=1.0e4):
   return x
 
 
-@expert_utils.add_name_scope()
 def add_positional_embedding(x, max_length, name, positions=None):
-  """Add positional embedding.
+  """Adds positional embedding.
 
   Args:
-    x: a Tensor with shape [batch, length, depth]
-    max_length: an integer.  static maximum size of any dimension.
-    name: a name for this layer.
-    positions: an optional tensor with shape [batch, length]
+    x: Tensor with shape [batch, length, depth].
+    max_length: int representing static maximum size of any dimension.
+    name: str representing name of the embedding tf.Variable.
+    positions: Tensor with shape [batch, length].
 
   Returns:
-    a Tensor the same shape as x.
+    Tensor of same shape as x.
   """
-  _, length, depth = common_layers.shape_list(x)
-  var = tf.cast(tf.get_variable(name, [max_length, depth]), x.dtype)
-  if positions is None:
-    sliced = tf.cond(
-        tf.less(length, max_length),
-        lambda: tf.slice(var, [0, 0], [length, -1]),
-        lambda: tf.pad(var, [[0, length - max_length], [0, 0]]))
-    return x + tf.expand_dims(sliced, 0)
-  else:
-    return x + tf.gather(var, tf.to_int32(positions))
+  with tf.name_scope("add_positional_embedding"):
+    _, length, depth = common_layers.shape_list(x)
+    var = tf.cast(tf.get_variable(name, [max_length, depth]), x.dtype)
+    if positions is None:
+      sliced = tf.cond(
+          tf.less(length, max_length),
+          lambda: tf.slice(var, [0, 0], [length, -1]),
+          lambda: tf.pad(var, [[0, max(0, length - max_length)], [0, 0]]))
+      return x + tf.expand_dims(sliced, 0)
+    else:
+      return x + tf.gather(var, tf.to_int32(positions))
 
 
-@expert_utils.add_name_scope()
 def add_positional_embedding_nd(x, max_length, name):
-  """Add n-dimensional positional embedding.
+  """Adds n-dimensional positional embedding.
 
-  Adds embeddings to represent the positional dimensions of the tensor.
-  The input tensor has n positional dimensions - i.e. 1 for text, 2 for images,
-  3 for video, etc.
+  The embeddings add to all positional dimensions of the tensor.
 
   Args:
-    x: a Tensor with shape [batch, p1 ... pn, depth]
-    max_length: an integer.  static maximum size of any dimension.
-    name: a name for this layer.
+    x: Tensor with shape [batch, p1 ... pn, depth]. It has n positional
+      dimensions, i.e., 1 for text, 2 for images, 3 for video, etc.
+    max_length: int representing static maximum size of any dimension.
+    name: str representing name of the embedding tf.Variable.
 
   Returns:
-    a Tensor the same shape as x.
+    Tensor of same shape as x.
   """
-  x_shape = common_layers.shape_list(x)
-  num_dims = len(x_shape) - 2
-  depth = x_shape[-1]
-  base_shape = [1] * (num_dims + 1) + [depth]
-  base_start = [0] * (num_dims + 2)
-  base_size = [-1] + [1] * num_dims + [depth]
-  for i in range(num_dims):
-    shape = base_shape[:]
-    start = base_start[:]
-    size = base_size[:]
-    shape[i + 1] = max_length
-    size[i + 1] = x_shape[i + 1]
-    var = (
-        tf.get_variable(
-            name + "_%d" % i,
-            shape,
-            initializer=tf.random_normal_initializer(0, depth**-0.5)) * (depth**
-                                                                         0.5))
-    x += tf.slice(var, start, size)
-  return x
+  with tf.name_scope("add_positional_embedding_nd"):
+    x_shape = common_layers.shape_list(x)
+    num_dims = len(x_shape) - 2
+    depth = x_shape[-1]
+    base_shape = [1] * (num_dims + 1) + [depth]
+    base_start = [0] * (num_dims + 2)
+    base_size = [-1] + [1] * num_dims + [depth]
+    for i in range(num_dims):
+      shape = base_shape[:]
+      start = base_start[:]
+      size = base_size[:]
+      shape[i + 1] = max_length
+      size[i + 1] = x_shape[i + 1]
+      var = tf.get_variable(
+          name + "_%d" % i,
+          shape,
+          initializer=tf.random_normal_initializer(0, depth**-0.5))
+      var *= depth**0.5
+      x += tf.slice(var, start, size)
+    return x
 
 
-@expert_utils.add_name_scope()
 def make_edge_vectors(adjacency_matrix, num_edge_types, depth, name=None):
   """Gets edge vectors for the edge types in the adjacency matrix.
 
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index cf47eff8b..7ae4bc660 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -17,14 +17,44 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+from absl.testing import parameterized
 import numpy as np
+
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 
 import tensorflow as tf
 
 
-class CommonAttentionTest(tf.test.TestCase):
+class CommonAttentionTest(parameterized.TestCase, tf.test.TestCase):
+
+  def testAddPositionalEmbedding(self):
+    x = np.random.rand(5, 3, 12)
+    with self.test_session() as session:
+      y = common_attention.add_positional_embedding(
+          tf.constant(x, dtype=tf.float32),
+          max_length=4,
+          name="pos_embedding")
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
+    self.assertEqual(res.shape, (5, 3, 12))
+
+  @parameterized.parameters(
+      ((5, 3, 12),),
+      ((5, 5, 5, 12),),
+      ((5, 3, 3, 3, 12),),
+  )
+  def testAddPositionalEmbeddingNd(self, input_shape):
+    x = np.random.rand(*input_shape)
+    with self.test_session() as session:
+      y = common_attention.add_positional_embedding_nd(
+          tf.constant(x, dtype=tf.float32),
+          max_length=5,
+          name="pos_embedding")
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
+    self.assertEqual(res.shape, input_shape)
 
   def testDotProductAttention(self):
     x = np.random.rand(5, 7, 12, 32)

From 06818cf4596d4cc5a5a78aa0c86203a338f59850 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 13 Aug 2018 12:01:49 -0700
Subject: [PATCH 0560/2720] partially reverting 208275102 since it broken
 simple_planning and Suraj's experiment.

PiperOrigin-RevId: 208519782
---
 tensor2tensor/models/research/next_frame_sv2p.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 4acc3b2ee..5f08980fe 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -55,9 +55,16 @@ def get_gaussian_latent(self, latent_mean, latent_std):
     latent = latent_mean + tf.exp(latent_std / 2.0) * latent
     return latent
 
+  def get_iteration_num(self):
+    step_num = tf.train.get_global_step()
+    # TODO(lukaszkaiser): what should it be if it"s undefined?
+    if step_num is None:
+      step_num = 1000000
+    return step_num
+
   def get_beta(self):
     """Get KL multiplier (beta) based on the schedule."""
-    step_num = tf.train.get_or_create_global_step()
+    step_num = self.get_iteration_num()
     schedule = self.hparams.latent_loss_multiplier_schedule
     second_stage = (self.hparams.num_iterations_1st_stage +
                     self.hparams.num_iterations_2nd_stage)
@@ -96,7 +103,7 @@ def anneal_loss(step_num):
   def get_scheduled_sample_func(self, batch_size):
     """Creates a function for scheduled sampling based on given hparams."""
     with tf.variable_scope("scheduled_sampling_func", reuse=False):
-      iter_num = tf.train.get_or_create_global_step()
+      iter_num = self.get_iteration_num()
       if self.hparams.scheduled_sampling_mode == "prob":
         decay_steps = self.hparams.scheduled_sampling_decay_steps
         probability = tf.train.polynomial_decay(
@@ -305,7 +312,7 @@ def construct_latent_tower(self, images):
         return tf.zeros_like(mean), tf.zeros_like(std)
 
       # No latent in the first phase
-      iter_num = tf.train.get_or_create_global_step()
+      iter_num = self.get_iteration_num()
       ret_mean, ret_std = tf.cond(
           iter_num < self.hparams.num_iterations_1st_stage,
           lambda: (tf.zeros_like(mean), tf.zeros_like(std)),

From 3acb9d74d8fd9a2ed73c69529ed2e6fbf02d7804 Mon Sep 17 00:00:00 2001
From: Keyon Vafa <vafa@google.com>
Date: Mon, 13 Aug 2018 13:01:05 -0700
Subject: [PATCH 0561/2720] Revert loss() in transformer_vae to original loss

PiperOrigin-RevId: 208529058
---
 .../models/research/transformer_vae.py        | 20 -------------------
 1 file changed, 20 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 98ecb4250..8e5e8c8b4 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -605,26 +605,6 @@ def __init__(self, *args, **kwargs):
   def has_input(self):
     return self._problem_hparams.input_modality
 
-  def loss(self, logits, features):
-    """Computes cross-entropy loss and scales by 1/batch_size."""
-    labels = features["targets"]
-    logits_shape = common_layers.shape_list(logits)
-    vocab_size = logits_shape[-1]
-    with tf.name_scope("padded_cross_entropy", values=[logits, labels]):
-      logits, labels = common_layers.pad_with_zeros(logits, labels)
-      logits = tf.reshape(
-          logits,
-          common_layers.shape_list(labels) + [vocab_size],
-          name="padded_cross_entropy_size_check")
-      logits = tf.cast(logits, tf.float32)
-      xent = common_layers.smoothing_cross_entropy(
-          logits, labels, vocab_size, confidence=1.0, gaussian=False)
-      if self._hparams.sum_over_latents:
-        recon_loss = tf.reduce_sum(xent) / tf.cast(logits_shape[0], tf.float32)
-      else:
-        recon_loss = tf.reduce_mean(xent)
-      return recon_loss
-
   def body(self, features):
     inputs = features["inputs"] if "inputs" in features else None
     if self._hparams.drop_inputs:

From f5938ccc10c4cf37965c42a3c8c5b5349310cb3d Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 13 Aug 2018 13:55:55 -0700
Subject: [PATCH 0562/2720] Tweak model-based RL parameters to run with more
 epochs, set default to 6 (measured slightly better performance than previous
 default 3 with same frame and computation budget).

PiperOrigin-RevId: 208537676
---
 .../models/research/next_frame_params.py      |  3 +--
 tensor2tensor/rl/model_rl_experiment.py       | 25 +++++++++++++------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index dd0559818..55786ec60 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -53,7 +53,7 @@ def next_frame():
 def next_frame_pixel_noise():
   """Basic 2-frame conv model with pixel noise."""
   hparams = next_frame()
-  hparams.add_hparam("video_modality_input_noise", 0.25)
+  hparams.add_hparam("video_modality_input_noise", 0.05)
   hparams.input_modalities = "inputs:video:pixel_noise"
   return hparams
 
@@ -252,4 +252,3 @@ def next_frame_ae_range(rhp):
   rhp.set_float("learning_rate_constant", 1., 2.)
   rhp.set_float("initializer_gain", 0.8, 1.5)
   rhp.set_int("filter_double_steps", 2, 3)
-
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/model_rl_experiment.py
index 5a563231a..2d53c79d4 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/model_rl_experiment.py
@@ -205,12 +205,22 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
 def train_world_model(problem_name, data_dir, output_dir, hparams, epoch):
   """Train the world model on problem_name."""
   train_steps = hparams.model_train_steps * (epoch + 2)
+  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
+  learning_rate = model_hparams.learning_rate_constant
+  # Bump learning rate after first epoch by 3x.
+  # We picked 3x because our default learning rate schedule decreases with
+  # 1/square root of the time step; 1/sqrt(10k) = 0.01 and 1/sqrt(100k) ~ 0.0032
+  # so by bumping it up 3x we about "go back" from 100k steps to 10k, which is
+  # approximately as much as "going back 1 epoch" would be in default schedule.
+  # In your experiments, you may want to optimize this rate to your schedule.
+  if epoch > 0: learning_rate *= 3
   with temporary_flags({
       "data_dir": data_dir,
       "output_dir": output_dir,
       "problem": problem_name,
       "model": hparams.generative_model,
       "hparams_set": hparams.generative_model_params,
+      "hparams": "learning_rate_constant=%.6f" % learning_rate,
       "eval_steps": 100,
       "train_steps": train_steps,
   }):
@@ -490,21 +500,22 @@ def combine_training_data(problem, final_data_dir, old_data_dirs,
 @registry.register_hparams
 def rl_modelrl_base():
   return tf.contrib.training.HParams(
-      epochs=3,
+      epochs=6,
       # Total frames used for training =
       # steps * (1 - 1/11) * epochs
-      # 1/11 steps are used for evaluation data
-      # 100k frames for training = 36666
-      true_env_generator_num_steps=36666,
+      # 1/11 steps are used for evaluation data.
+      # So to use N frames set steps = N / (epochs * (1 - 1/11)).
+      # We set it to use 100k frames for training.
+      true_env_generator_num_steps=int(100000 / (6 * (1.0 - 1.0/11.0))),
       generative_model="next_frame_basic",
-      generative_model_params="next_frame",
+      generative_model_params="next_frame_pixel_noise",
       ppo_params="ppo_pong_base",
       autoencoder_train_steps=0,
-      model_train_steps=100000,
+      model_train_steps=50000,
       simulated_env_generator_num_steps=2000,
       simulation_random_starts=True,
       intrinsic_reward_scale=0.,
-      ppo_epochs_num=400,  # This should be enough to see something
+      ppo_epochs_num=200,  # This should be enough to see something
       # Our simulated envs do not know how to reset.
       # You should set ppo_time_limit to the value you believe that
       # the simulated env produces a reasonable output.

From d6d4dae5861e695b02a1aa14c187447f4c4a3624 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 13 Aug 2018 13:59:57 -0700
Subject: [PATCH 0563/2720] fixing OSS version.

PiperOrigin-RevId: 208538283
---
 tensor2tensor/models/research/next_frame_sv2p.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 5f08980fe..f5da10330 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -314,7 +314,7 @@ def construct_latent_tower(self, images):
       # No latent in the first phase
       iter_num = self.get_iteration_num()
       ret_mean, ret_std = tf.cond(
-          iter_num < self.hparams.num_iterations_1st_stage,
+          tf.less(iter_num, self.hparams.num_iterations_1st_stage),
           lambda: (tf.zeros_like(mean), tf.zeros_like(std)),
           lambda: (mean, std))
 

From 1747148b6f7e02a22af514799cb0b43e74facabc Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 13 Aug 2018 14:16:46 -0700
Subject: [PATCH 0564/2720] Fill in other required features with dummy values
 in serving_utils

PiperOrigin-RevId: 208541089
---
 tensor2tensor/serving/README.md        |  7 ++++
 tensor2tensor/serving/serving_utils.py | 48 ++++++++++++++++++++++++--
 2 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/serving/README.md b/tensor2tensor/serving/README.md
index 156a2503d..8e3ca3350 100644
--- a/tensor2tensor/serving/README.md
+++ b/tensor2tensor/serving/README.md
@@ -34,6 +34,13 @@ tensorflow_model_server \
 
 ## 3. Query the Server
 
+**Note**: The `t2t-query-server` is meant only as an example. You may need to
+modify it to suit your needs. The exported model expects an input
+example that is structured identically to what would be found on disk during
+training (serialized `tf.train.Example`). For text problems, that means that
+it expects the inputs to already be encoded as integers. You can see how the
+`t2t-query-server` does this by reading the code.
+
 Install some dependencies:
 
 ```
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index c11689d93..2004aa7a6 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -19,6 +19,7 @@
 from __future__ import print_function
 
 import base64
+import functools
 from googleapiclient import discovery
 import grpc
 
@@ -31,11 +32,51 @@
 
 
 
-def _make_example(input_ids, feature_name="inputs"):
+def _make_example(input_ids, problem, input_feature_name="inputs"):
+  """Make a tf.train.Example for the problem.
+
+  features[input_feature_name] = input_ids
+
+  Also fills in any other required features with dummy values.
+
+  Args:
+    input_ids: list<int>.
+    problem: Problem.
+    input_feature_name: name of feature for input_ids.
+
+  Returns:
+    tf.train.Example
+  """
   features = {
-      feature_name:
+      input_feature_name:
           tf.train.Feature(int64_list=tf.train.Int64List(value=input_ids))
   }
+
+  # Fill in dummy values for any other required features that presumably
+  # will not actually be used for prediction.
+  data_fields, _ = problem.example_reading_spec()
+  for fname, ftype in data_fields.items():
+    if fname == input_feature_name:
+      continue
+    if not isinstance(ftype, tf.FixedLenFeature):
+      # Only FixedLenFeatures are required
+      continue
+    if ftype.default_value is not None:
+      # If there's a default value, no need to fill it in
+      continue
+    num_elements = functools.reduce(lambda acc, el: acc * el, ftype.shape, 1)
+    if ftype.dtype in [tf.int32, tf.int64]:
+      value = tf.train.Feature(
+          int64_list=tf.train.Int64List(value=[0] * num_elements))
+    if ftype.dtype in [tf.float32, tf.float64]:
+      value = tf.train.Feature(
+          float_list=tf.train.FloatList(value=[0.] * num_elements))
+    if ftype.dtype == tf.bytes:
+      value = tf.train.Feature(
+          bytes_list=tf.train.BytesList(value=[""] * num_elements))
+    tf.logging.info("Adding dummy value for feature %s as it is required by "
+                    "the Problem.", fname)
+    features[fname] = value
   return tf.train.Example(features=tf.train.Features(feature=features))
 
 
@@ -110,7 +151,8 @@ def predict(inputs_list, problem, request_fn):
       _encode(inputs, input_encoder, add_eos=problem.has_inputs)
       for inputs in inputs_list
   ]
-  examples = [_make_example(input_ids, fname) for input_ids in input_ids_list]
+  examples = [_make_example(input_ids, problem, fname)
+              for input_ids in input_ids_list]
   predictions = request_fn(examples)
   output_decoder = problem.feature_info["targets"].encoder
   outputs = [

From 3d194cb14c4ae563895fbd910e6ed82809baee3f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 13 Aug 2018 14:36:21 -0700
Subject: [PATCH 0565/2720] moving vqa to outside

PiperOrigin-RevId: 208544629
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/vqa.py          | 453 ++++++++++++++++++
 tensor2tensor/data_generators/vqa_utils.py    | 235 +++++++++
 3 files changed, 689 insertions(+)
 create mode 100644 tensor2tensor/data_generators/vqa.py
 create mode 100644 tensor2tensor/data_generators/vqa_utils.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 792cd1896..0a68b662e 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -77,6 +77,7 @@
     "tensor2tensor.data_generators.translate_enzh",
     "tensor2tensor.data_generators.twentybn",
     "tensor2tensor.data_generators.video_generated",
+    "tensor2tensor.data_generators.vqa",
     "tensor2tensor.data_generators.wiki",
     "tensor2tensor.data_generators.wikisum.wikisum",
     "tensor2tensor.data_generators.wikitext103",
diff --git a/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
new file mode 100644
index 000000000..10bfcc599
--- /dev/null
+++ b/tensor2tensor/data_generators/vqa.py
@@ -0,0 +1,453 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generators for VQA data sets."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import base64
+import csv
+import json
+import os
+import random
+import sys
+import tarfile
+import zipfile
+import numpy as np
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import image_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.google.data_generators import vqa_utils
+from tensor2tensor.utils import metrics
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+def _get_vqa_v2_annotations(directory,
+                            annotation_url,
+                            annotation_filename="vqa_v2.tar.gz"):
+  """Extract the VQA V2 annotation files to directory unless it's there."""
+  annotation_file = generator_utils.maybe_download_from_drive(
+      directory, annotation_filename, annotation_url)
+  with tarfile.open(annotation_file, "r:gz") as annotation_tar:
+    annotation_tar.extractall(directory)
+
+
+def _get_vqa_v2_image_raw_dataset(directory, image_root_url, image_urls):
+  """Extract the VQA V2 image data set to directory unless it's there."""
+  for url in image_urls:
+    filename = os.path.basename(url)
+    download_url = os.path.join(image_root_url, url)
+    path = generator_utils.maybe_download(directory, filename, download_url)
+    unzip_dir = os.path.join(directory, filename.strip(".zip"))
+    if not tf.gfile.Exists(unzip_dir):
+      zipfile.ZipFile(path, "r").extractall(directory)
+
+
+def _get_vqa_v2_image_feature_dataset(
+    directory, feature_url, feature_filename="mscoco_feat.tar.gz"):
+  """Extract the VQA V2 feature data set to directory unless it's there."""
+  feature_file = generator_utils.maybe_download_from_drive(
+      directory, feature_filename, feature_url)
+  with tarfile.open(feature_file, "r:gz") as feature_tar:
+    feature_tar.extractall(directory)
+
+
+class ImageQuestion2MultilabelProblem(image_utils.ImageProblem):
+  """Base class for image question answer problem."""
+
+  @property
+  def target_space_id(self):
+    raise NotImplementedError()
+
+  @property
+  def vocab_size(self):
+    raise NotImplementedError
+
+  @property
+  def num_classes(self):
+    raise NotImplementedError()
+
+  @property
+  def vocab_filename(self):
+    raise NotImplementedError()
+
+  @property
+  def label_filename(self):
+    raise NotImplementedError()
+
+  @property
+  def train_shards(self):
+    raise NotImplementedError()
+
+  @property
+  def dev_shards(self):
+    raise NotImplementedError()
+
+  def source_data_files(self, dataset_split):
+    raise NotImplementedError()
+
+  def generator(self, data_dir, tmp_dir, dataset_split):
+    raise NotImplementedError()
+
+  def eval_metrics(self):
+    return [
+        metrics.Metrics.ACC_MULTILABEL_MATCH3,
+    ]
+
+  def feature_encoders(self, data_dir):
+    input_encoder = text_encoder.ImageEncoder(channels=self.num_channels)
+    vocab_file = os.path.join(data_dir, self.vocab_filename)
+    question_encoder = text_encoder.TokenTextEncoder(
+        vocab_file, replace_oov="UNK")
+    label_file = os.path.join(data_dir, self.label_filename)
+    target_encoder = text_encoder.ClassLabelEncoder(
+        class_labels_fname=label_file)
+    return {"inputs": input_encoder,
+            "question": question_encoder,
+            "targets": target_encoder}
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    question_encoder = self._encoders["question"]
+    targets_encoder = self._encoders["targets"]
+
+    p.input_modality = {
+        "inputs": (registry.Modalities.IMAGE + ":identity", None),
+        "question": (registry.Modalities.SYMBOL, question_encoder.vocab_size)
+    }
+    p.target_modality = (registry.Modalities.CLASS_LABEL + ":multi_label",
+                         targets_encoder.vocab_size)
+    p.input_space_id = problem.SpaceID.IMAGE  # multiple input features?
+    p.target_space_id = self.target_space_id
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    generator_utils.generate_dataset_and_shuffle(
+        self.generator(data_dir, tmp_dir, problem.DatasetSplit.TRAIN),
+        self.training_filepaths(data_dir, self.train_shards, shuffled=False),
+        self.generator(data_dir, tmp_dir, problem.DatasetSplit.EVAL),
+        self.dev_filepaths(data_dir, self.dev_shards, shuffled=False))
+
+
+@registry.register_problem
+class ImageVqav2Tokens10kLabels3k(ImageQuestion2MultilabelProblem):
+  """VQA V2, raw images, 10k question vocab, 3k answer label."""
+  _MSCOCO_ROOT_URL = "http://msvocds.blob.core.windows.net/"
+  _MSCOCO_IMAGE_URLS = [
+      "coco2014/train2014.zip", "coco2014/val2014.zip", "coco2014/test2014.zip",
+  ]
+  _VQA_V2_ANNOTATION_URL = ("https://drive.google.com/uc?export=download&id="
+                            "1xfMU54ObCLvMRAekT3cfcIg-AgY39fWB")
+
+  _VQA_V2_TRAIN_DATASETS = [
+      ("trainval_resnet101_faster_rcnn_genome_36.tsv",
+       "v2_train2014_annotations.json"),
+  ]
+  _VQA_V2_DEV_DATASETS = [
+      ("trainval_resnet101_faster_rcnn_genome_36.tsv",
+       "v2_val2014_annotations.json"),
+  ]
+  _VQA_V2_TEST_DATASETS = [
+      ("test2015_resnet101_faster_rcnn_genome_36.tsv",
+       "v2_test2015_annotations.json"),
+  ]
+
+  def source_data_files(self, dataset_split):
+    train = dataset_split == problem.DatasetSplit.TRAIN
+    return self._VQA_V2_TRAIN_DATASETS if train else self._VQA_V2_DEV_DATASETS
+
+  @property
+  def target_space_id(self):
+    return problem.SpaceID.GENERIC
+
+  @property
+  def vocab_size(self):
+    return 10000
+
+  @property
+  def num_classes(self):
+    return 3000
+
+  @property
+  def vocab_filename(self):
+    return "question.vocab.%d" % self.vocab_size
+
+  @property
+  def label_filename(self):
+    return "answer.label.%d" % self.num_classes
+
+  @property
+  def train_shards(self):
+    return 128
+
+  @property
+  def dev_shards(self):
+    return 64
+
+  def example_reading_spec(self):
+    data_fields, data_items_to_decoders = (
+        super(ImageVqav2Tokens10kLabels3k, self).example_reading_spec())
+    data_fields["image/image_id"] = tf.FixedLenFeature((), tf.int64)
+    data_fields["image/question_id"] = tf.FixedLenFeature((), tf.int64)
+    data_fields["image/question"] = tf.FixedLenSequenceFeature(
+        (), tf.int64, allow_missing=True)
+    data_fields["image/answer"] = tf.FixedLenSequenceFeature(
+        (), tf.int64, allow_missing=True)
+
+    data_items_to_decoders[
+        "question"] = tf.contrib.slim.tfexample_decoder.Tensor(
+            "image/question")
+    data_items_to_decoders[
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(
+            "image/answer")
+    return data_fields, data_items_to_decoders
+
+  def preprocess_example(self, example, mode, hparams):
+    # hparams is model_hparams
+    image = example["inputs"]
+    example["inputs"] = vqa_utils.vqa_v2_preprocess_image(
+        image, hparams.height, hparams.width, mode,
+        resize_side=hparams.resize_side, distort=hparams.distort,
+        image_model_fn=hparams.image_model_fn)
+    return example
+
+  def generator(self, data_dir, tmp_dir, dataset_split):
+    datasets = self.source_data_files(dataset_split)
+    return self.vqa_v2_generator(data_dir, tmp_dir, datasets)
+
+  def vqa_v2_generator(self, data_dir, tmp_dir, datasets):
+    """VQA v2 generator using raw images."""
+    _get_vqa_v2_annotations(tmp_dir, self._VQA_V2_ANNOTATION_URL)
+    _get_vqa_v2_image_raw_dataset(tmp_dir, self._MSCOCO_ROOT_URL,
+                                  self._MSCOCO_IMAGE_URLS)
+    vocab_path = os.path.join(data_dir, self.vocab_filename)
+    if not tf.gfile.Exists(vocab_path):
+      vocab_tmp_path = os.path.join(tmp_dir, self.vocab_filename)
+      tf.gfile.Copy(vocab_tmp_path, vocab_path)
+      with tf.gfile.GFile(vocab_path, mode="r") as f:
+        vocab_data = "<pad>\n<EOS>\n" + f.read() + "UNK\n"
+      with tf.gfile.GFile(vocab_path, mode="w") as f:
+        f.write(vocab_data)
+    label_path = os.path.join(data_dir, self.label_filename)
+    if not tf.gfile.Exists(label_path):
+      label_tmp_path = os.path.join(tmp_dir, self.label_filename)
+      tf.gfile.Copy(label_tmp_path, label_path)
+
+    vocab_encoder = text_encoder.TokenTextEncoder(vocab_path, replace_oov="UNK")
+    label_encoder = text_encoder.ClassLabelEncoder(
+        class_labels_fname=label_path)
+
+    prefix_annotation = []
+    for prefix, annotation_file in datasets:
+      annotation_path = os.path.join(tmp_dir, annotation_file)
+      with tf.gfile.Open(annotation_path) as f:
+        annotation_json = json.loads(f.read())
+      prefix_annotation += [(prefix, anno) for anno in annotation_json]
+    random.shuffle(prefix_annotation)
+    annotation_count = len(prefix_annotation)
+    tf.logging.info("Processing %d annotations for vqa v2" %(annotation_count))
+
+    for prefix, anno in prefix_annotation:
+      image_id = anno["image_id"]
+      question = vocab_encoder.encode(anno["question"])
+      answer = [label_encoder.encode(ans) for ans in anno["answer"]]
+      answer = answer if answer else [0]  # 0 indicates padding
+      image_filename = "COCO_" + prefix + "_" + str(image_id).zfill(12) + ".jpg"
+      image_filepath = os.path.join(tmp_dir, prefix, image_filename)
+      with tf.gfile.Open(image_filepath, "r") as f:
+        encoded_image_data = f.read()
+        yield {
+            "image/encoded": [encoded_image_data],
+            "image/format": ["jpeg"],
+            "image/image_id": [image_id],
+            "image/question_id": [anno["question_id"]],
+            "image/question": question,
+            "image/answer": answer,
+        }
+
+
+@registry.register_problem
+class ImageVqav2RcnnFeatureTokens10kLabels3k(ImageVqav2Tokens10kLabels3k):
+  """VQA V2, image feature, 10k question vocab, 3k answer label."""
+  _VQA_V2_FEATURE_URL = ("https://drive.google.com/uc?export=download&id="
+                         "1yTTFUWqx1SScC-Whs2vRbF3tDsEEjrtt")
+
+  @property
+  def num_boxes(self):
+    return 36
+
+  @property
+  def feature_dimension(self):
+    return 2048
+
+  @property
+  def spatial_feature_dimension(self):
+    return 6
+
+  @property
+  def feature_file_field_names(self):
+    return ["image_id",
+            "image_w",
+            "image_h",
+            "num_boxes",
+            "boxes",
+            "features"]
+
+  def preprocess_example(self, example, mode, hparams):
+    # reshape some features
+    example["inputs"] = tf.reshape(
+        example["inputs"], [self.num_boxes, 1, self.feature_dimension])
+    example["spatial_feature"] = tf.reshape(
+        example["spatial_feature"],
+        [self.num_boxes, 1, self.spatial_feature_dimension])
+    return example
+
+  def example_reading_spec(self):
+    data_fields, data_items_to_decoders = {}, {}
+    data_fields["image/feature"] = tf.FixedLenSequenceFeature(
+        (), tf.float32, allow_missing=True)
+    data_fields["image/spatial_feature"] = tf.FixedLenSequenceFeature(
+        (), tf.float32, allow_missing=True)
+    data_fields["image/image_id"] = tf.FixedLenFeature((), tf.int64)
+    data_fields["image/question_id"] = tf.FixedLenFeature((), tf.int64)
+    data_fields["image/question"] = tf.FixedLenSequenceFeature(
+        (), tf.int64, allow_missing=True)
+    data_fields["image/answer"] = tf.FixedLenSequenceFeature(
+        (), tf.int64, allow_missing=True)
+
+    data_items_to_decoders[
+        "inputs"] = tf.contrib.slim.tfexample_decoder.Tensor(
+            "image/feature")
+    data_items_to_decoders[
+        "question_id"] = tf.contrib.slim.tfexample_decoder.Tensor(
+            "image/question_id")
+    data_items_to_decoders[
+        "image_id"] = tf.contrib.slim.tfexample_decoder.Tensor(
+            "image/image_id")
+
+    data_items_to_decoders[
+        "spatial_feature"] = tf.contrib.slim.tfexample_decoder.Tensor(
+            "image/spatial_feature")
+    data_items_to_decoders[
+        "question"] = tf.contrib.slim.tfexample_decoder.Tensor(
+            "image/question")
+    data_items_to_decoders[
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(
+            "image/answer")
+
+    return data_fields, data_items_to_decoders
+
+  def vqa_v2_generator(self, data_dir, tmp_dir, datasets):
+    """VQA v2 generator using image features."""
+    _get_vqa_v2_annotations(tmp_dir, self._VQA_V2_ANNOTATION_URL)
+    _get_vqa_v2_image_feature_dataset(tmp_dir, self._VQA_V2_FEATURE_URL)
+    vocab_path = os.path.join(data_dir, self.vocab_filename)
+    if not tf.gfile.Exists(vocab_path):
+      vocab_tmp_path = os.path.join(tmp_dir, self.vocab_filename)
+      tf.gfile.Copy(vocab_tmp_path, vocab_path)
+      with tf.gfile.GFile(vocab_path, mode="r") as f:
+        vocab_data = "<pad>\n<EOS>\n" + f.read() + "UNK\n"
+      with tf.gfile.GFile(vocab_path, mode="w") as f:
+        f.write(vocab_data)
+    label_path = os.path.join(data_dir, self.label_filename)
+    if not tf.gfile.Exists(label_path):
+      label_tmp_path = os.path.join(tmp_dir, self.label_filename)
+      tf.gfile.Copy(label_tmp_path, label_path)
+
+    vocab_encoder = text_encoder.TokenTextEncoder(vocab_path, replace_oov="UNK")
+    label_encoder = text_encoder.ClassLabelEncoder(
+        class_labels_fname=label_path)
+
+    # merge annotations
+    annotation_json = []
+    for _, annotation_file in datasets:
+      annotation_path = os.path.join(tmp_dir, annotation_file)
+      with tf.gfile.Open(annotation_path) as f:
+        annotation_json += json.loads(f.read())
+    annotation_count = len(annotation_json)
+    tf.logging.info("Processing %d annotations for vqa v2" %(annotation_count))
+
+    imageid2annotation = {}
+    for anno in annotation_json:
+      if anno["image_id"] not in imageid2annotation:
+        imageid2annotation[anno["image_id"]] = [anno]
+      else:
+        imageid2annotation[anno["image_id"]].append(anno)
+
+    csv.field_size_limit(sys.maxsize)
+    for feature_file, _ in datasets:
+      feature_file_path = os.path.join(tmp_dir, feature_file)
+      with open(feature_file_path, "r+b") as tsv_file:
+        csv_reader = csv.DictReader(
+            tsv_file, delimiter="\t", fieldnames=self.feature_file_field_names)
+        for item in csv_reader:
+          item["num_boxes"] = int(item["num_boxes"])
+          image_id = int(item["image_id"])
+          image_w = float(item["image_w"])
+          image_h = float(item["image_h"])
+          bboxes = np.frombuffer(base64.decodestring(item["boxes"]),
+                                 dtype=np.float32).reshape(
+                                     (item["num_boxes"], -1))
+
+          box_width = bboxes[:, 2] - bboxes[:, 0]
+          box_height = bboxes[:, 3] - bboxes[:, 1]
+          scaled_width = box_width / image_w
+          scaled_height = box_height / image_h
+          scaled_x = bboxes[:, 0] / image_w
+          scaled_y = bboxes[:, 1] / image_h
+
+          box_width = box_width[..., np.newaxis]
+          box_height = box_height[..., np.newaxis]
+          scaled_width = scaled_width[..., np.newaxis]
+          scaled_height = scaled_height[..., np.newaxis]
+          scaled_x = scaled_x[..., np.newaxis]
+          scaled_y = scaled_y[..., np.newaxis]
+
+          spatial_features = np.concatenate(
+              (scaled_x,
+               scaled_y,
+               scaled_x + scaled_width,
+               scaled_y + scaled_height,
+               scaled_width,
+               scaled_height),
+              axis=1)
+
+          if image_id in imageid2annotation:
+            for anno in imageid2annotation[image_id]:
+              question = vocab_encoder.encode(anno["question"])
+              answer = [label_encoder.encode(ans) for ans in anno["answer"]]
+              answer = answer if answer else [0]  # 0 indicates padding
+              yield {
+                  "image/feature":
+                  np.frombuffer(base64.decodestring(item["features"]),
+                                dtype=np.float32).tolist(),
+                  "image/spatial_feature": spatial_features.flatten().tolist(),
+                  "image/height": [image_h],
+                  "image/width": [image_w],
+                  "image/bboxes": bboxes.flatten().tolist(),
+                  "image/image_id": [image_id],
+                  "image/question_id": [anno["question_id"]],
+                  "image/question": question,
+                  "image/answer": answer,
+              }
+
+            del imageid2annotation[image_id]
+
+    # assert all annotations are included
+    assert not imageid2annotation
diff --git a/tensor2tensor/data_generators/vqa_utils.py b/tensor2tensor/data_generators/vqa_utils.py
new file mode 100644
index 000000000..1599879b2
--- /dev/null
+++ b/tensor2tensor/data_generators/vqa_utils.py
@@ -0,0 +1,235 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Utilities for VQA data sets."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+import tensorflow as tf
+
+from tensorflow.python.ops import control_flow_ops
+
+# some functions are copied and modified from
+# vgg_preprocessing and inception_preprocessing in
+# models/research/slim/preprocessing/
+
+_R_MEAN = 123.68
+_G_MEAN = 116.78
+_B_MEAN = 103.94
+
+
+def _smallest_size_at_least(height, width, smallest_side):
+  """Computes new shape with the smallest side equal to `smallest_side`.
+
+    Computes new shape with the smallest side equal to `smallest_side` while
+    preserving the original aspect ratio.
+
+  Args:
+    height: an int32 scalar tensor indicating the current height.
+    width: an int32 scalar tensor indicating the current width.
+    smallest_side: A python integer or scalar `Tensor` indicating the size of
+    the smallest side after resize.
+
+  Returns:
+    new_height: an int32 scalar tensor indicating the new height.
+    new_width: and int32 scalar tensor indicating the new width.
+  """
+  smallest_side = tf.convert_to_tensor(smallest_side, dtype=tf.int32)
+
+  height = tf.to_float(height)
+  width = tf.to_float(width)
+  smallest_side = tf.to_float(smallest_side)
+
+  scale = tf.cond(
+      tf.greater(height, width), lambda: smallest_side / width,
+      lambda: smallest_side / height)
+  new_height = tf.to_int32(height * scale)
+  new_width = tf.to_int32(width * scale)
+  return new_height, new_width
+
+
+def _aspect_preserving_resize(image, smallest_side):
+  """Resize images preserving the original aspect ratio.
+
+  Args:
+    image: A 3-D image `Tensor`.
+    smallest_side: A python integer or scalar `Tensor` indicating the size of
+    the smallest side after resize.
+
+  Returns:
+    resized_image: A 3-D tensor containing the resized image.
+  """
+  smallest_side = tf.convert_to_tensor(smallest_side, dtype=tf.int32)
+
+  shape = tf.shape(image)
+  height = shape[0]
+  width = shape[1]
+  new_height, new_width = _smallest_size_at_least(height, width, smallest_side)
+  image = tf.expand_dims(image, 0)
+  resized_image = tf.image.resize_images(
+      image, size=[new_height, new_width], method=tf.image.ResizeMethod.BICUBIC)
+
+  resized_image = tf.squeeze(resized_image)
+  resized_image.set_shape([None, None, 3])
+  return resized_image
+
+
+def _flip(image):
+  """Random horizontal image flip."""
+  image = tf.image.random_flip_left_right(image)
+  return image
+
+
+def _distort_color(image, color_ordering=0, scope=None):
+  """Distort the color of a Tensor image.
+
+  Each color distortion is non-commutative and thus ordering of the color ops
+  matters. Ideally we would randomly permute the ordering of the color ops.
+  Rather then adding that level of complication, we select a distinct ordering
+  of color ops for each preprocessing thread.
+
+  Args:
+    image: 3-D Tensor containing single image in [0, 1].
+    color_ordering: Python int, a type of distortion (valid values: 0-3).
+    scope: Optional scope for name_scope.
+  Returns:
+    3-D Tensor color-distorted image on range [0, 1]
+  Raises:
+    ValueError: if color_ordering not in [0, 3]
+  """
+  with tf.name_scope(scope, "distort_color", [image]):
+    if color_ordering == 0:
+      image = tf.image.random_brightness(image, max_delta=32. / 255.)
+      image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
+      image = tf.image.random_hue(image, max_delta=0.2)
+      image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
+    elif color_ordering == 1:
+      image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
+      image = tf.image.random_brightness(image, max_delta=32. / 255.)
+      image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
+      image = tf.image.random_hue(image, max_delta=0.2)
+    elif color_ordering == 2:
+      image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
+      image = tf.image.random_hue(image, max_delta=0.2)
+      image = tf.image.random_brightness(image, max_delta=32. / 255.)
+      image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
+    elif color_ordering == 3:
+      image = tf.image.random_hue(image, max_delta=0.2)
+      image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
+      image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
+      image = tf.image.random_brightness(image, max_delta=32. / 255.)
+    else:
+      raise ValueError("color_ordering must be in [0, 3]")
+
+    # The random_* ops do not necessarily clamp.
+    return tf.clip_by_value(image, 0.0, 1.0)
+
+
+def _apply_with_random_selector(x, func, num_cases):
+  """Computes func(x, sel), with sel sampled from [0...num_cases-1].
+
+  Args:
+    x: input Tensor.
+    func: Python function to apply.
+    num_cases: Python int32, number of cases to sample sel from.
+
+  Returns:
+    The result of func(x, sel), where func receives the value of the
+    selector as a python integer, but sel is sampled dynamically.
+  """
+  sel = tf.random_uniform([], maxval=num_cases, dtype=tf.int32)
+  # Pass the real x only to one of the func calls.
+  return control_flow_ops.merge([
+      func(control_flow_ops.switch(x, tf.equal(sel, case))[1], case)
+      for case in range(num_cases)
+  ])[0]
+
+
+def _mean_image_subtraction(image, means):
+  """Subtracts the given means from each image channel.
+
+  For example:
+    means = [123.68, 116.779, 103.939]
+    image = _mean_image_subtraction(image, means)
+
+  Note that the rank of `image` must be known.
+
+  Args:
+    image: a tensor of size [height, width, C].
+    means: a C-vector of values to subtract from each channel.
+
+  Returns:
+    the centered image.
+
+  Raises:
+    ValueError: If the rank of `image` is unknown, if `image` has a rank other
+      than three or if the number of channels in `image` doesn't match the
+      number of values in `means`.
+  """
+  if image.get_shape().ndims != 3:
+    raise ValueError("Input must be of size [height, width, C>0]")
+  num_channels = image.get_shape().as_list()[-1]
+  if len(means) != num_channels:
+    raise ValueError("len(means) must match the number of channels")
+
+  channels = tf.split(axis=2, num_or_size_splits=num_channels, value=image)
+  for i in range(num_channels):
+    channels[i] -= means[i]
+  return tf.concat(axis=2, values=channels)
+
+
+def vqa_v2_preprocess_image(
+    image,
+    height,
+    width,
+    mode,
+    resize_side=512,
+    distort=True,
+    image_model_fn="resnet_v1_152",
+):
+  """vqa v2 preprocess image."""
+
+  image = tf.image.convert_image_dtype(image, dtype=tf.float32)
+  assert resize_side > 0
+  if resize_side:
+    image = _aspect_preserving_resize(image, resize_side)
+  if mode == tf.estimator.ModeKeys.TRAIN:
+    image = tf.random_crop(image, [height, width, 3])
+  else:
+    # Central crop, assuming resize_height > height, resize_width > width.
+    image = tf.image.resize_image_with_crop_or_pad(image, height, width)
+
+  image = tf.clip_by_value(image, 0.0, 1.0)
+
+  if mode == tf.estimator.ModeKeys.TRAIN and distort:
+    image = _flip(image)
+    num_distort_cases = 4
+    # pylint: disable=unnecessary-lambda
+    image = _apply_with_random_selector(
+        image, lambda x, ordering: _distort_color(x, ordering),
+        num_cases=num_distort_cases)
+
+  if image_model_fn.startswith("resnet_v1"):
+    # resnet_v1 uses vgg preprocessing
+    image = image * 255.
+    image = _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
+  elif image_model_fn.startswith("resnet_v2"):
+    # resnet v2 uses inception preprocessing
+    image = tf.subtract(image, 0.5)
+    image = tf.multiply(image, 2.0)
+
+  return image

From 967a578beef2064a9824febb80cc54f248329e6d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 13 Aug 2018 16:06:57 -0700
Subject: [PATCH 0566/2720] fix vqa_utils import

PiperOrigin-RevId: 208559948
---
 tensor2tensor/data_generators/vqa.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
index 10bfcc599..27f025c1d 100644
--- a/tensor2tensor/data_generators/vqa.py
+++ b/tensor2tensor/data_generators/vqa.py
@@ -32,7 +32,7 @@
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
-from tensor2tensor.google.data_generators import vqa_utils
+from tensor2tensor.data_generators import vqa_utils
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 

From d8c3419b2ad808e6d2b42bfe63a50a29e50062ee Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 13 Aug 2018 16:29:35 -0700
Subject: [PATCH 0567/2720] moving shuffle_buffer_size to hparams.

PiperOrigin-RevId: 208563357
---
 tensor2tensor/data_generators/video_utils.py       | 2 +-
 tensor2tensor/models/research/next_frame_params.py | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index f8e7ef26e..38c8d9f16 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -300,7 +300,7 @@ def check_integrity_and_batch(*datasets):
       batch_dataset = preprocessed_dataset.apply(
           tf.contrib.data.batch_and_drop_remainder(num_frames))
     dataset = batch_dataset.map(features_from_batch)
-    dataset = dataset.shuffle(256)
+    dataset = dataset.shuffle(hparams.shuffle_buffer_size)
     return dataset
 
   def eval_metrics(self):
diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index 55786ec60..5a7db28eb 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -46,6 +46,7 @@ def next_frame():
   hparams.add_hparam("preprocess_resize_frames", None)
   hparams.add_hparam("concatenate_actions", True)
   hparams.add_hparam("tiny_mode", False)
+  hparams.add_hparam("shuffle_buffer_size", 128)
   return hparams
 
 

From a6f62a9c4fc777ea9710402f8444853000fd773e Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 13 Aug 2018 16:44:13 -0700
Subject: [PATCH 0568/2720] internal merge of PR #988

PiperOrigin-RevId: 208565631
---
 tensor2tensor/data_generators/problem.py | 2 +-
 tensor2tensor/models/transformer.py      | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index e1884ff98..107409c5c 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -866,7 +866,7 @@ def define_shapes(example):
         # on TPU, we use params["batch_size"], which specifies the number of
         # examples across all datashards
         batch_size = params["batch_size"]
-        dataset = dataset.batch(batch_size)
+        dataset = dataset.batch(batch_size, drop_remainder=True)
       else:
         num_shards = config.data_parallelism.n if config else 1
         batch_size = hparams.batch_size * num_shards
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index c03dd6a9a..c31b800c2 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1953,14 +1953,16 @@ def transformer_tpu():
   update_hparams_for_tpu(hparams)
   return hparams
 
+
 @registry.register_hparams
 def transformer_timeseries_tpu():
   """HParams for running Transformer model on timeseries on TPU."""
   hparams = transformer_timeseries()
   update_hparams_for_tpu(hparams)
-  hparams.batch_size = 256 # revert to value set in transformer_timeseries
+  hparams.batch_size = 256  # revert to value set in transformer_timeseries
   return hparams
 
+
 @registry.register_hparams
 def transformer_tpu_bf16_activation():
   """HParams for Transformer model with BF16 activation on TPU."""

From 964cde8bdc2326371b7f9d28fc1f5e8b49561ea6 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 13 Aug 2018 18:45:53 -0700
Subject: [PATCH 0569/2720] Adding BAIR with actions.

PiperOrigin-RevId: 208579946
---
 .../data_generators/bair_robot_pushing.py     | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index 784010542..eca4a2792 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -157,3 +157,23 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
           "state": state,
           "action": action,
       }
+
+
+@registry.register_problem
+class VideoBairRobotPushingWithActions(VideoBairRobotPushing):
+  """Berkeley (BAIR) robot pushing dataset with actions."""
+
+  @property
+  def extra_reading_spec(self):
+    """Additional data fields to store on disk and their decoders."""
+    data_fields = {
+        "frame_number": tf.FixedLenFeature([1], tf.int64),
+        "action": tf.FixedLenFeature([4], tf.float32),
+    }
+    decoders = {
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
+            tensor_key="frame_number"),
+        "action": tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="action"),
+    }
+    return data_fields, decoders
+

From 8218bafef5333f2f71076c583f3b437669353349 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 13 Aug 2018 22:43:13 -0700
Subject: [PATCH 0570/2720] adding support for full internal loss. the main
 difference is that the loss is being calculated also on the input frames.
 e.g. if the number input frames is 4, sv2p still predicts 3 of them and
 currently we do not apply the losses on those.

PiperOrigin-RevId: 208596982
---
 tensor2tensor/models/research/next_frame_params.py | 1 +
 tensor2tensor/models/research/next_frame_sv2p.py   | 4 ++++
 2 files changed, 5 insertions(+)

diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index 5a7db28eb..d9cb7efe4 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -93,6 +93,7 @@ def next_frame_stochastic():
   hparams.add_hparam("latent_num_frames", 0)  # 0 means use all frames.
   hparams.add_hparam("anneal_end", 100000)
   hparams.add_hparam("upsample_method", "conv2d_transpose")
+  hparams.add_hparam("internal_loss", False)
   return hparams
 
 
diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index f5da10330..c107f45b7 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -632,6 +632,10 @@ def fix_video_dims_and_concat_on_x_axis(x):
     if "target_reward" in features:
       return_targets = {"targets": predictions, "target_reward": reward_pred}
 
+    if hparams.internal_loss:
+      loss = tf.losses.mean_squared_error(all_frames[1:], gen_images)
+      extra_loss = {"training": loss + extra_loss}
+
     return return_targets, extra_loss
 
 

From 83e9df7013d56bd0e13cff39ad6a8b4ee10ddd8e Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 13 Aug 2018 23:08:04 -0700
Subject: [PATCH 0571/2720] Code cleaup in autoencoder, works both on image and
 text.

PiperOrigin-RevId: 208598788
---
 tensor2tensor/models/research/autoencoders.py | 142 ++++++++++--------
 1 file changed, 83 insertions(+), 59 deletions(-)

diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 37fa81f23..ffc4226f6 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -18,8 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.data_generators import image_utils
-from tensor2tensor.data_generators import text_problems
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
@@ -47,6 +45,36 @@ def __init__(self, *args, **kwargs):
     self._cur_bottleneck_tensor = None
     self.is1d = None
 
+  @property
+  def num_channels(self):
+    # TODO(lukaszkaiser): is this a universal enough way to get channels?
+    try:
+      num_channels = self.hparams.problem.num_channels
+    except AttributeError:
+      num_channels = 1
+    return num_channels
+
+  def image_summary(self, name, image_logits, max_outputs=1):
+    """Helper for image summaries that are safe on TPU."""
+    if len(image_logits.get_shape()) != 5:
+      tf.logging.info("Not generating image summary, maybe not an image.")
+      return
+    return tf.summary.image(
+        name,
+        common_layers.tpu_safe_image_summary(tf.argmax(image_logits, -1)),
+        max_outputs=max_outputs)
+
+  def embed(self, x):
+    """Input embedding with a non-zero bias for uniform inputs."""
+    with tf.variable_scope("embed", reuse=tf.AUTO_REUSE):
+      x = tf.layers.dense(
+          x,
+          self.hparams.hidden_size,
+          name="embed",
+          activation=common_layers.belu,
+          bias_initializer=tf.random_normal_initializer(stddev=0.01))
+      return common_attention.add_timing_signal_nd(x)
+
   def bottleneck(self, x):
     with tf.variable_scope("bottleneck"):
       hparams = self.hparams
@@ -83,7 +111,11 @@ def discriminator(self, x, is_training):
             net, training=is_training, momentum=0.999, name="d_bn2")
       net = lrelu(net)
       size = height * width
-      net = tf.reshape(net, [batch_size, size * 8])  # [bs, h * w * 8]
+      x_shape = x.get_shape().as_list()
+      if x_shape[1] is None or x_shape[2] is None:
+        net = tf.reduce_mean(net, axis=[1, 2])  # [bs, 128]
+      else:
+        net = tf.reshape(net, [batch_size, size * 8])  # [bs, h * w * 8]
       net = tf.layers.dense(net, 1024, name="d_fc3")  # [bs, 1024]
       if hparams.discriminator_batchnorm:
         net = tf.layers.batch_normalization(
@@ -146,22 +178,27 @@ def body(self, features):
     hparams = self.hparams
     is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
-      x = features["targets"]
       labels = features["targets_raw"]
-      shape = common_layers.shape_list(x)
+      vocab_size = self._problem_hparams.target_modality.top_dimensionality
+      shape = common_layers.shape_list(labels)
+      x = tf.one_hot(labels, vocab_size)
+      x = tf.reshape(x, shape[:-1] + [shape[-1] * vocab_size])
+      x = self.embed(x)
       is1d = shape[2] == 1
       self.is1d = is1d
       # Run encoder.
       x = self.encoder(x)
       # Bottleneck (mix during early training, not too important but stable).
       b, b_loss = self.bottleneck(x)
+      b_shape = common_layers.shape_list(b)
       self._cur_bottleneck_tensor = b
       b = self.unbottleneck(b, common_layers.shape_list(x)[-1])
       b = common_layers.mix(b, x, hparams.bottleneck_warmup_steps, is_training)
       if hparams.gan_loss_factor != 0.0:
         # Add a purely sampled batch on which we'll compute the GAN loss.
         g = self.unbottleneck(
-            self.sample(), common_layers.shape_list(x)[-1], reuse=True)
+            self.sample(shape=b_shape),
+            common_layers.shape_list(x)[-1], reuse=True)
         b = tf.concat([g, b], axis=0)
       # With probability bottleneck_max_prob use the bottleneck, otherwise x.
       if hparams.bottleneck_max_prob < -1.0:
@@ -184,45 +221,36 @@ def body(self, features):
     # Cut to the right size and mix before returning.
     res = x[:, :shape[1], :shape[2], :]
 
-    is_image = isinstance(self.hparams.problem, image_utils.ImageProblem)
-    if is_image:
-      vocab_size = self.hparams.problem.vocab_size
-
-      res = tf.layers.dense(
-          res, self.hparams.problem.num_channels * self.hparams.hidden_size)
-      output_shape = common_layers.shape_list(res)[:-1] + [
-          self.hparams.problem.num_channels, self.hparams.hidden_size
-      ]
-      res = tf.reshape(res, output_shape)
-    elif isinstance(self.hparams.problem, text_problems.Text2TextProblem):
-      vocab_size = self._problem_hparams.target_modality.top_dimensionality
-      res = tf.layers.dense(res, self.hparams.hidden_size)
-    else:
-      raise Exception("Unsupported problem type: %s" % self.hparams.problem)
+    # Final dense layer.
+    res = tf.layers.dense(
+        res, self.num_channels * hparams.hidden_size, name="res_dense")
+
+    output_shape = common_layers.shape_list(res)[:-1] + [
+        self.num_channels, self.hparams.hidden_size
+    ]
+    res = tf.reshape(res, output_shape)
 
+    # Losses.
     losses = {}
     if hparams.gan_loss_factor != 0.0:
       res_gan, res = tf.split(res, 2, axis=0)
       with tf.variable_scope("vq"):
         reconstr_gan, gan_codes, _, code_loss_gan, _ = discretization.vq_loss(
             res_gan, labels, vocab_size)
-        losses["code_loss_gan"] = code_loss_gan
+        losses["code_loss_gan"] = (code_loss_gan * hparams.code_loss_factor *
+                                   hparams.gan_loss_factor)
 
     with tf.variable_scope("vq", reuse=tf.AUTO_REUSE):
       (reconstr, _, target_codes, code_loss,
        targets_loss) = discretization.vq_loss(res, labels, vocab_size)
 
-    losses["code_loss"] = code_loss
+    losses["code_loss"] = code_loss * hparams.code_loss_factor
     losses["training"] = targets_loss
 
     # Add GAN loss if requested.
     gan_loss = 0.0
     if hparams.gan_loss_factor != 0.0:
-      if is_image:
-        tf.summary.image(
-            "gan",
-            common_layers.tpu_safe_image_summary(tf.argmax(reconstr_gan, -1)),
-            max_outputs=1)
+      self.image_summary("gan", reconstr_gan)
 
       def discriminate(x):
         return self.discriminator(x, is_training=is_training)
@@ -239,11 +267,7 @@ def discriminate(x):
                                                self.hparams.num_sliced_vecs)
       gan_loss *= hparams.gan_loss_factor
 
-    if is_image:
-      tf.summary.image(
-          "ae",
-          common_layers.tpu_safe_image_summary(tf.argmax(reconstr, -1)),
-          max_outputs=1)
+    self.image_summary("ae", reconstr)
 
     losses["b_loss"] = b_loss
     losses["gan_loss"] = -gan_loss
@@ -252,7 +276,7 @@ def discriminate(x):
     return logits, losses
 
   def sample(self, features=None, shape=None):
-    del features, shape
+    del features
     hp = self.hparams
     div_x = 2**hp.num_hidden_layers
     div_y = 1 if self.is1d else 2**hp.num_hidden_layers
@@ -260,6 +284,7 @@ def sample(self, features=None, shape=None):
         hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
         hp.bottleneck_bits
     ]
+    size = size if shape is None else shape
     # Sample in [-1, 1] as the bottleneck is under tanh.
     return 2.0 * tf.random_uniform(size) - 1.0
 
@@ -283,11 +308,7 @@ def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
       features["inputs"] = tf.expand_dims(features["inputs"], 2)
 
     # Sample and decode.
-    # TODO(lukaszkaiser): is this a universal enough way to get channels?
-    try:
-      num_channels = self.hparams.problem.num_channels
-    except AttributeError:
-      num_channels = 1
+    num_channels = self.num_channels
     if "targets" not in features:
       features["targets"] = tf.zeros(
           [self.hparams.batch_size, 1, 1, num_channels], dtype=tf.int32)
@@ -340,22 +361,30 @@ def body(self, features):
     if hparams.autoregressive_mode == "none":
       assert not hparams.autoregressive_forget_base
       return basic_result, losses
+    if "training" in losses:
+      plain_training_loss = losses.pop("training")
+      losses["plain"] = plain_training_loss
+    basic_result = self.embed(basic_result)
     shape = common_layers.shape_list(basic_result)
-    basic1d = tf.reshape(basic_result, [shape[0], -1, shape[3]])
+    basic1d = tf.reshape(basic_result, [shape[0], -1, shape[-1]])
+    vocab_size = self._problem_hparams.target_modality.top_dimensionality
+    targets = tf.one_hot(features["targets_raw"], vocab_size)
+    targets = tf.reshape(targets, shape[:-2] + [shape[-2] * vocab_size])
+    targets = self.embed(targets)
+    targets = tf.reshape(targets, common_layers.shape_list(basic_result))
     # During autoregressive inference, don't resample.
     if hparams.mode == tf.estimator.ModeKeys.PREDICT:
       if hasattr(hparams, "sampled_basic1d_tensor"):
         basic1d = hparams.sampled_basic1d_tensor
       else:
         hparams.sampled_basic1d_tensor = basic1d
-    print(common_layers.shape_list(features["targets"]))
     # Prepare inputs for autoregressive modes.
     if common_layers.shape_list(features["targets"])[1] == 1:
       # This happens on the first step of predicitions.
       assert hparams.mode == tf.estimator.ModeKeys.PREDICT
-      features["targets"] = tf.zeros_like(basic_result)
+      targets = tf.zeros_like(basic_result)
     targets_dropout = common_layers.mix(
-        features["targets"],
+        targets,
         tf.zeros_like(basic_result),
         hparams.bottleneck_warmup_steps,
         is_training,
@@ -366,18 +395,18 @@ def body(self, features):
         hparams.autoregressive_eval_pure_autoencoder):
       targets_dropout = tf.zeros_like(basic_result)
     # Now combine the basic reconstruction with shifted targets.
-    targets1d = tf.reshape(targets_dropout, [shape[0], -1, shape[3]])
+    targets1d = tf.reshape(targets_dropout, [shape[0], -1, shape[-1]])
     targets_shifted = common_layers.shift_right_3d(targets1d)
     concat1d = tf.concat([basic1d, targets_shifted], axis=-1)
     # The forget_base hparam sets purely-autoregressive mode, no autoencoder.
     if hparams.autoregressive_forget_base:
-      concat1d = tf.reshape(features["targets"], [shape[0], -1, shape[3]])
+      concat1d = tf.reshape(targets, [shape[0], -1, shape[-1]])
       concat1d = common_layers.shift_right_3d(concat1d)
     # The autoregressive part depends on the mode.
     if hparams.autoregressive_mode == "conv3":
       res = common_layers.conv1d(
           concat1d,
-          shape[3],
+          hparams.hidden_size,
           3,
           padding="LEFT",
           activation=common_layers.belu,
@@ -386,7 +415,7 @@ def body(self, features):
     if hparams.autoregressive_mode == "conv5":
       res = common_layers.conv1d(
           concat1d,
-          shape[3],
+          hparams.hidden_size,
           5,
           padding="LEFT",
           activation=common_layers.belu,
@@ -395,7 +424,7 @@ def body(self, features):
     if hparams.autoregressive_mode == "sru":
       res = common_layers.conv1d(
           concat1d,
-          shape[3],
+          hparams.hidden_size,
           3,
           padding="LEFT",
           activation=common_layers.belu,
@@ -486,14 +515,6 @@ def encoder(self, x):
       residual_conv = tf.layers.conv2d
       if hparams.residual_use_separable_conv:
         residual_conv = tf.layers.separable_conv2d
-      # Input embedding with a non-zero bias for uniform inputs.
-      x = tf.layers.dense(
-          x,
-          hparams.hidden_size,
-          name="embed",
-          activation=common_layers.belu,
-          bias_initializer=tf.random_normal_initializer(stddev=0.01))
-      x = common_attention.add_timing_signal_nd(x)
       # Down-convolutions.
       for i in range(hparams.num_hidden_layers):
         with tf.variable_scope("layer_%d" % i):
@@ -588,7 +609,7 @@ def bottleneck(self, x):
                           hparams.mode == tf.estimator.ModeKeys.TRAIN)
     return x, 0.0
 
-  def sample(self, features=None):
+  def sample(self, features=None, shape=None):
     del features
     hp = self.hparams
     div_x = 2**hp.num_hidden_layers
@@ -597,6 +618,7 @@ def sample(self, features=None):
         hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
         hp.bottleneck_bits
     ]
+    size = size if shape is None else shape
     rand = tf.random_uniform(size)
     return 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
 
@@ -625,7 +647,7 @@ def unbottleneck(self, x, res_size, reuse=None):
     with tf.variable_scope("unbottleneck", reuse=reuse):
       return discretization.parametrized_unbottleneck(x, res_size, self.hparams)
 
-  def sample(self, features=None):
+  def sample(self, features=None, shape=None):
     del features
     hp = self.hparams
     div_x = 2**hp.num_hidden_layers
@@ -634,6 +656,7 @@ def sample(self, features=None):
         hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
         hp.bottleneck_bits
     ]
+    size = size if shape is None else shape
     rand = tf.random_uniform(size)
     res = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
     # If you want to set some first bits to a fixed value, do this:
@@ -793,6 +816,7 @@ def autoencoder_basic():
   hparams.add_hparam("sample_width", 32)
   hparams.add_hparam("discriminator_batchnorm", True)
   hparams.add_hparam("num_sliced_vecs", 4096)
+  hparams.add_hparam("code_loss_factor", 1.0)
   hparams.add_hparam("gan_loss_factor", 0.0)
   hparams.add_hparam("use_vqloss", False)
   return hparams
@@ -913,7 +937,7 @@ def autoencoder_ordered_text():
   hparams.learning_rate_constant = 2.0
   hparams.learning_rate_warmup_steps = 2000
   hparams.bottleneck_bits = 1024
-  hparams.batch_size = 2048
+  hparams.batch_size = 1024
   hparams.autoregressive_mode = "sru"
   hparams.hidden_size = 256
   hparams.max_hidden_size = 4096

From 7e5f0c2def82949abb5aac0ab321b2af184a164f Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Tue, 14 Aug 2018 09:25:43 -0700
Subject: [PATCH 0572/2720] Making schedules parametrized.

PiperOrigin-RevId: 208658788
---
 .../data_generators/multi_problem.py          | 29 ++++++++++---------
 tensor2tensor/layers/common_hparams.py        | 14 +++++++++
 tensor2tensor/layers/common_layers.py         | 10 ++++---
 3 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index c16077184..e121c9199 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -20,6 +20,7 @@
 
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
@@ -88,10 +89,6 @@ def get_hparams(self, model_hparams=None):
 
     return self._hparams
 
-  @property
-  def mixing_schedule(self):
-    return MixingSchedule.CONSTANT
-
   def flatten_zip(self, *args):
     """A list of examples to a dataset containing mixed examples.
 
@@ -180,26 +177,32 @@ def get_next_from_dataset(dataset_iter):
         return dataset_iter.get_next()
 
       def get_exp_sched_prob():
+        """Inverse decay exponential to mix datasets."""
         with tf.control_dependencies([problem_step.assign_add(1)]):
-          # TODO(urvashik): Make 5e-8 a parameter.
-          # In the current setup, with about 100 examples per batch on average,
-          # the model converges to 50-50 mixing by ~140k problem steps.
-          return tf.minimum(1. - tf.exp(-5e-8 * problem_step), 0.5)
+          inv_exp_decay = common_layers.inverse_exp_decay(
+              max_step=hparams.multiproblem_schedule_max_examples,
+              min_value=1e-4,
+              step=problem_step
+          )
+          # inv_exp_decay is bounded above by 1.0
+          return inv_exp_decay * hparams.multiproblem_schedule_threshold
 
       def get_const_sched_prob():
-        return 0.5
+        return hparams.multiproblem_schedule_threshold
 
       def mix_data(example):
         """Function to mix the different datasets according to a schedule."""
         del example
-        if self.mixing_schedule == MixingSchedule.EXPONENTIAL:
+        if hparams.multiproblem_mixing_schedule == MixingSchedule.EXPONENTIAL:
           prob = get_exp_sched_prob()
-        elif self.mixing_schedule == MixingSchedule.CONSTANT:
+        elif hparams.multiproblem_mixing_schedule == MixingSchedule.CONSTANT:
           prob = get_const_sched_prob()
         else:
-          raise ValueError("Unknown schedule %s" % str(self.mixing_schedule))
+          raise ValueError("Unknown schedule %s" % str(
+              hparams.multiproblem_mixing_schedule))
         tf.logging.info("Using the %s schedule to "
-                        "train the MultiProblem." % str(self.mixing_schedule))
+                        "train the MultiProblem." % str(
+                            hparams.multiproblem_mixing_schedule))
 
         def sample_task(curr_task, num_tasks_left):
           """A recursive function to sample a task.
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index f202cec8e..bbe8c2d18 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -253,6 +253,20 @@ def basic_params1():
       # pretrained model will be randomly initialized. Superfluous parameters in
       # the pretrained model will be ignored.
       pretrained_model_dir="",
+      # Threshold used for two cases: the primary task probability for the
+      # constant mixing schedule, and the exponential schedule limit for when
+      # mixing should stop (eg: 0.5 means stop at 50-50 mixing, 0.8 means stop
+      # at 20-80 mixing for the primary-others mixing case.)
+      multiproblem_schedule_threshold=0.5,
+      # The number of examples at which the proportion of the mixed in datasets
+      # is multiproblem_schedule_threshold
+      multiproblem_schedule_max_examples=1e7,
+      # When training multiproblems, we can mix the data according to different
+      # schedules. Example: a constant schedule mixing 20-80 between the primary
+      # and other tasks.
+      # A list of supported schedules can be found in
+      # `data_generators.multi_problem.py`.
+      multiproblem_mixing_schedule="constant"
   )
 
 
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 5ac71bcca..b8973c950 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -129,16 +129,18 @@ def hard_tanh(x, saturation_limit=0.9):
   return tf.minimum(1.0, tf.maximum(x, -1.0)), saturation_cost
 
 
-def inverse_exp_decay(max_step, min_value=0.01):
+def inverse_exp_decay(max_step, min_value=0.01, step=None):
   """Inverse-decay exponentially from 0.01 to 1.0 reached at max_step."""
   inv_base = tf.exp(tf.log(min_value) / float(max_step))
-  step = tf.to_float(tf.train.get_global_step())
+  if step is None:
+    step = tf.to_float(tf.train.get_global_step())
   return inv_base**tf.maximum(float(max_step) - step, 0.0)
 
 
-def inverse_lin_decay(max_step, min_value=0.01):
+def inverse_lin_decay(max_step, min_value=0.01, step=None):
   """Inverse-decay linearly from 0.01 to 1.0 reached at max_step."""
-  step = tf.to_float(tf.train.get_global_step())
+  if step is None:
+    step = tf.to_float(tf.train.get_global_step())
   progress = tf.minimum(step / float(max_step), 1.0)
   return progress * (1.0 - min_value) + min_value
 

From 8648666759aca6b7d8ea30b8e99b6d9e1ed4c676 Mon Sep 17 00:00:00 2001
From: Tomasz Latkowski <tlatkowski@gmail.com>
Date: Tue, 14 Aug 2018 19:05:13 +0200
Subject: [PATCH 0573/2720] added paraphrase generation problem

---
 tensor2tensor/data_generators/all_problems.py |   1 +
 .../data_generators/paraphrase_ms_coco.py     | 192 ++++++++++++++++++
 .../paraphrase_ms_coco_test.py                |  69 +++++++
 3 files changed, 262 insertions(+)
 create mode 100644 tensor2tensor/data_generators/paraphrase_ms_coco.py
 create mode 100644 tensor2tensor/data_generators/paraphrase_ms_coco_test.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 0a68b662e..129489aa2 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -53,6 +53,7 @@
     "tensor2tensor.data_generators.mrpc",
     "tensor2tensor.data_generators.mscoco",
     "tensor2tensor.data_generators.multinli",
+    "tensor2tensor.data_generators.paraphrase_ms_coco",
     "tensor2tensor.data_generators.program_search",
     "tensor2tensor.data_generators.ocr",
     "tensor2tensor.data_generators.pointer_generator_word",
diff --git a/tensor2tensor/data_generators/paraphrase_ms_coco.py b/tensor2tensor/data_generators/paraphrase_ms_coco.py
new file mode 100644
index 000000000..f8c2b0449
--- /dev/null
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco.py
@@ -0,0 +1,192 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 classes for paraphrase generation problems."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import json
+import io
+import collections
+import os
+import zipfile
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators import problem
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+_MS_COCO_DOWNLOAD_URL = "http://msvocds.blob.core.windows.net/annotations-1-0-3"
+_MS_COCO_ZIPPED_FILE = "captions_train-val2014.zip"
+
+_MS_COCO_TRAIN_FILE = "captions_train2014.json"
+_MS_COCO_DEV_FILE = "captions_val2014.json"
+
+
+def create_combination(list_of_sentences):
+  """ Generates all possible pair combinations for
+  the input list of sentences, for example:
+
+  input = ["paraphrase1", "paraphrase2", "paraphrase3"]
+
+  output = [("paraphrase1", "paraphrase2"),
+            ("paraphrase1", "paraphrase3"),
+            ("paraphrase2", "paraphrase3")]
+
+  Args:
+    list_of_sentences: the list of input sentences.
+  Returns:
+    the list of all possible sentence pairs.
+  """
+  num_sentences = len(list_of_sentences) - 1
+  combinations = []
+  for i, sentence in enumerate(list_of_sentences):
+    if i == num_sentences:
+      break
+    num_pairs = num_sentences - i
+    populated = num_pairs * [list_of_sentences[i]]
+    zipped = list(zip(populated, list_of_sentences[i + 1:]))
+    combinations += zipped
+  return combinations
+
+
+class ParaphraseGenerationProblem(text_problems.Text2TextProblem):
+
+  @property
+  def bidirectional(self):
+    """If set to true, generates data in the following way:
+
+    sentence1 -> sentence2
+    sentence2 -> sentence1
+    """
+    raise NotImplementedError()
+
+  def prepare_data(self, data_dir, tmp_dir, dataset_split):
+    raise NotImplementedError()
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    paraphrase_pairs = self.prepare_data(data_dir, tmp_dir, dataset_split)
+    for (caption1, caption2) in paraphrase_pairs:
+      caption_pairs = [(caption1, caption2)]
+      if self.bidirectional:
+        caption_pairs += [(caption2, caption1)]
+      for caption_pair in caption_pairs:
+        yield {
+          'inputs': caption_pair[0],
+          'targets': caption_pair[1]
+        }
+
+
+class ParaphraseGenerationMsCocoProblem(ParaphraseGenerationProblem):
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+      "split": problem.DatasetSplit.TRAIN,
+      "shards": 10,
+    }, {
+      "split": problem.DatasetSplit.EVAL,
+      "shards": 1,
+    }]
+
+  @property
+  def approx_vocab_size(self):
+    return 2 ** 13
+
+  def prepare_data(self, data_dir, tmp_dir, dataset_split):
+    ms_coco_path = self._maybe_download(tmp_dir, dataset_split)
+    captions = self._get_captions(ms_coco_path)
+    tf.logging.info("Retrieved %d captions\n" % (len(captions)))
+    paraphrase_pairs = []
+
+    tf.logging.info("Generating input combinations...")
+    for captions_for_image in captions:
+      combinations_of_captions = create_combination(captions_for_image)
+      paraphrase_pairs += combinations_of_captions
+
+    tf.logging.info("Created %d combinations pairs." % (len(paraphrase_pairs)))
+    return paraphrase_pairs
+
+  def _maybe_download(self, tmp_dir, dataset_split):
+    filename = os.path.basename(_MS_COCO_ZIPPED_FILE)
+    download_url = os.path.join(_MS_COCO_DOWNLOAD_URL, filename)
+    path = generator_utils.maybe_download(tmp_dir, filename, download_url)
+    unzip_dir = os.path.join(tmp_dir, filename.strip(".zip"))
+    if not tf.gfile.Exists(unzip_dir):
+      tf.logging.info("Unzipping data to {}".format(unzip_dir))
+      zipfile.ZipFile(path, "r").extractall(unzip_dir)
+
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      ms_coco_file = _MS_COCO_TRAIN_FILE
+    else:
+      ms_coco_file = _MS_COCO_DEV_FILE
+    ms_coco_path = os.path.join(unzip_dir, "annotations", ms_coco_file)
+    return ms_coco_path
+
+  def _get_captions(self, ms_coco_path):
+    caption_file = io.open(ms_coco_path)
+    caption_json = json.load(caption_file)
+    annotations = caption_json["annotations"]
+    captions_for_image = collections.defaultdict(list)
+
+    for annotation in annotations:
+      image_id = annotation["image_id"]
+      captions_for_image[image_id].append(annotation["caption"])
+
+    captions = list(captions_for_image.values())
+    return captions
+
+
+@registry.register_problem
+class ParaphraseGenerationMsCocoProblem2d(
+  ParaphraseGenerationMsCocoProblem):
+
+  @property
+  def bidirectional(self):
+    return True
+
+
+@registry.register_problem
+class ParaphraseGenerationMsCocoProblem1d(
+  ParaphraseGenerationMsCocoProblem):
+
+  @property
+  def bidirectional(self):
+    return False
+
+
+@registry.register_problem
+class ParaphraseGenerationMsCocoProblem2dCharacters(
+  ParaphraseGenerationMsCocoProblem2d):
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+
+@registry.register_problem
+class ParaphraseGenerationMsCocoProblem1dCharacters(
+  ParaphraseGenerationMsCocoProblem1d):
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
diff --git a/tensor2tensor/data_generators/paraphrase_ms_coco_test.py b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
new file mode 100644
index 000000000..a20e01ecb
--- /dev/null
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
@@ -0,0 +1,69 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.data_generators.paraphrase_ms_coco."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import mock
+
+from tensor2tensor.data_generators import paraphrase_ms_coco
+
+import tensorflow as tf
+
+
+class ParaphraseGenerationProblemTest(tf.test.TestCase):
+
+  def testCombinationPairs(self):
+    inputs = ['A', 'B', 'C']
+    expected_combination = [('A', 'B'), ('A', 'C'), ('B', 'C')]
+    actual_combination = paraphrase_ms_coco.create_combination(inputs)
+    self.assertEqual(actual_combination, expected_combination)
+
+  @mock.patch("tensor2tensor.data_generators"
+              ".paraphrase_mscoco.ParaphraseGenerationProblem.prepare_data",
+              return_value=[("sentence1", "sentence2")])
+  @mock.patch("tensor2tensor.data_generators"
+              ".paraphrase_mscoco.ParaphraseGenerationProblem.bidirectional")
+  def testBidirectionalTrue(self, data, bidirectional):
+    paraphrase_problem = paraphrase_ms_coco.ParaphraseGenerationProblem()
+    paraphrase_problem.bidirectional = True
+
+    expected_generated_data = [{"inputs": "sentence1", "targets": "sentence2"},
+                               {"inputs": "sentence2", "targets": "sentence1"}]
+    actual_generated_data = list(paraphrase_problem.generate_samples("data_dir",
+                                                                     "tmp_dir",
+                                                                     "dataset_split"))
+    self.assertEqual(actual_generated_data, expected_generated_data)
+
+  @mock.patch("tensor2tensor.data_generators"
+              ".paraphrase_mscoco.ParaphraseGenerationProblem.prepare_data",
+              return_value=[("sentence1", "sentence2")])
+  @mock.patch("tensor2tensor.data_generators"
+              ".paraphrase_mscoco.ParaphraseGenerationProblem.bidirectional")
+  def testBidirectionalFalse(self, data, bidirectional):
+    paraphrase_problem = paraphrase_ms_coco.ParaphraseGenerationProblem()
+    paraphrase_problem.bidirectional = False
+
+    expected_generated_data = [{"inputs": "sentence1", "targets": "sentence2"}]
+    actual_generated_data = list(paraphrase_problem.generate_samples("data_dir",
+                                                                     "tmp_dir",
+                                                                     "dataset_split"))
+    self.assertEqual(actual_generated_data, expected_generated_data)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 2ce54dfc6a2ca8fc5383292333ef9ddc97c50efb Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 14 Aug 2018 10:20:49 -0700
Subject: [PATCH 0574/2720] include original images in TB.

PiperOrigin-RevId: 208668158
---
 tensor2tensor/models/research/next_frame_sv2p.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index c107f45b7..df7ff0c60 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -606,12 +606,6 @@ def body(self, features):
         latent_stds=latent_stds, beta=beta, true_frames=all_frames,
         gen_frames=gen_images)
 
-    # Ignore the predictions from the input frames.
-    # This is NOT the same as original paper/implementation.
-    predictions = gen_images[hparams.video_num_input_frames-1:]
-    reward_pred = gen_rewards[hparams.video_num_input_frames-1:]
-    reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove undeeded dimension.
-
     # TODO(mbz): clean this up!
     def fix_video_dims_and_concat_on_x_axis(x):
       x = tf.transpose(x, [1, 3, 4, 0, 2])
@@ -619,11 +613,17 @@ def fix_video_dims_and_concat_on_x_axis(x):
       x = tf.transpose(x, [0, 3, 1, 2])
       return x
 
-    frames_gd = fix_video_dims_and_concat_on_x_axis(target_frames)
-    frames_pd = fix_video_dims_and_concat_on_x_axis(predictions)
+    frames_gd = fix_video_dims_and_concat_on_x_axis(all_frames[1:])
+    frames_pd = fix_video_dims_and_concat_on_x_axis(gen_images)
     side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
     tf.summary.image("full_video", side_by_side_video)
 
+    # Ignore the predictions from the input frames.
+    # This is NOT the same as original paper/implementation.
+    predictions = gen_images[hparams.video_num_input_frames-1:]
+    reward_pred = gen_rewards[hparams.video_num_input_frames-1:]
+    reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove undeeded dimension.
+
     # Swap back time and batch axes.
     predictions = common_video.swap_time_and_batch_axes(predictions)
     reward_pred = common_video.swap_time_and_batch_axes(reward_pred)

From 5e8009e6d4863c3f95a673dc50b36b05da7441ad Mon Sep 17 00:00:00 2001
From: Youlong Cheng <ylc@google.com>
Date: Tue, 14 Aug 2018 10:21:25 -0700
Subject: [PATCH 0575/2720] Expose MeshTensorFlow in t2t_trainer.

PiperOrigin-RevId: 208668268
---
 tensor2tensor/bin/t2t_trainer.py | 9 +++++++++
 tensor2tensor/models/__init__.py | 2 ++
 2 files changed, 11 insertions(+)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 685979cb0..c69ded415 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -32,6 +32,7 @@
 from tensor2tensor.utils import usr_dir
 import tensorflow as tf
 
+from tensorflow.contrib.tpu.python.tpu import tpu_config
 
 flags = tf.flags
 FLAGS = flags.FLAGS
@@ -197,6 +198,14 @@ def create_run_config(hp):
   assert FLAGS.output_dir or FLAGS.checkpoint_path
   tpu_config_extra_kwargs = {}
 
+  if getattr(hp, "mtf_mode", False):
+    save_ckpt_steps = None  # Disable the default saver
+    save_ckpt_secs = None  # Disable the default saver
+    tpu_config_extra_kwargs = {
+        "num_cores_per_replica": 1,
+        "per_host_input_for_training": tpu_config.InputPipelineConfig.BROADCAST,
+    }
+
   # the various custom getters we have written do not play well together yet.
   # TODO(noam): ask rsepassi for help here.
   daisy_chain_variables = (
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 311956c50..e4989cb96 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -19,6 +19,8 @@
 # pylint: disable=unused-import
 
 from tensor2tensor.layers import modalities  # pylint: disable=g-import-not-at-top
+from tensor2tensor.mesh_tensorflow import mtf_image_transformer
+from tensor2tensor.mesh_tensorflow import mtf_transformer
 from tensor2tensor.models import basic
 from tensor2tensor.models import bytenet
 from tensor2tensor.models import distillation

From d84f5474d5c42ffc9bd9dc35472c40aba45970dc Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 14 Aug 2018 11:29:08 -0700
Subject: [PATCH 0576/2720] move vqa models to outside

PiperOrigin-RevId: 208681659
---
 tensor2tensor/models/__init__.py              |   4 +
 .../models/research/vqa_attention.py          | 586 +++++++++++++
 .../models/research/vqa_attention_test.py     |  69 ++
 .../research/vqa_recurrent_self_attention.py  | 316 +++++++
 .../models/research/vqa_self_attention.py     | 823 ++++++++++++++++++
 5 files changed, 1798 insertions(+)
 create mode 100644 tensor2tensor/models/research/vqa_attention.py
 create mode 100644 tensor2tensor/models/research/vqa_attention_test.py
 create mode 100644 tensor2tensor/models/research/vqa_recurrent_self_attention.py
 create mode 100644 tensor2tensor/models/research/vqa_self_attention.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index e4989cb96..84f2404b9 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -59,4 +59,8 @@
 from tensor2tensor.models.research import transformer_symshard
 from tensor2tensor.models.research import transformer_vae
 from tensor2tensor.models.research import universal_transformer
+from tensor2tensor.models.research import vqa_attention
+from tensor2tensor.models.research import vqa_recurrent_self_attention
+from tensor2tensor.models.research import vqa_self_attention
+
 # pylint: enable=unused-import
diff --git a/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
new file mode 100644
index 000000000..920ef64a6
--- /dev/null
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -0,0 +1,586 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Attention models for VQA."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_hparams
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import vqa_layers
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import restore_hook
+from tensor2tensor.utils import t2t_model
+
+import tensorflow as tf
+
+# pylint: disable=unused-import
+from tensorflow.contrib.layers.python.layers import utils
+from tensorflow.contrib.slim.python.slim.nets.resnet_v1 import resnet_v1_152
+from tensorflow.contrib.slim.python.slim.nets.resnet_v2 import resnet_v2_152
+
+
+@registry.register_model
+class VqaAttentionBaseline(t2t_model.T2TModel):
+  """Attention baseline model for VQA."""
+
+  @staticmethod
+  def train_hooks():
+    restore_resnet_hook = restore_hook.RestoreHook(
+        # TODO(zichaoy): hard code the path given static function.
+        checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
+        new_model_scope="vqa_attention_baseline/body/",
+        old_model_scope="resnet_v1_152/",
+    )
+    return [restore_resnet_hook]
+
+  def body(self, features):
+    hp = self.hparams
+    # pylint: disable=eval-used
+    if hp.image_input_type == "image":
+      image_feat = vqa_layers.image_embedding(
+          features["inputs"],
+          model_fn=eval(hp.image_model_fn),
+          trainable=hp.train_resnet,
+          is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
+    else:
+      image_feat = features["inputs"]
+
+    if hp.image_feat_size:
+      image_feat = common_layers.dense(image_feat, hp.image_feat_size)
+
+    # apply layer normalization and dropout on image_feature
+    utils.collect_named_outputs("norms", "image_feat_before_l2",
+                                tf.norm(image_feat, axis=-1))
+    image_feat = common_layers.l2_norm(image_feat)
+    utils.collect_named_outputs("norms", "image_feat_after_l2",
+                                tf.norm(image_feat, axis=-1))
+
+    image_feat = tf.nn.dropout(image_feat, keep_prob=1.-hp.dropout)
+
+    query = question_encoder(features["question"], hp)
+    utils.collect_named_outputs("norms", "query",
+                                tf.norm(query, axis=-1))
+
+    image_ave = attn(image_feat, query, hp)
+    utils.collect_named_outputs("norms", "image_ave",
+                                tf.norm(image_ave, axis=-1))
+
+    image_question = tf.concat([image_ave, query], axis=1)
+    utils.collect_named_outputs("norms", "image_question",
+                                tf.norm(image_question, axis=-1))
+
+    image_question = tf.nn.dropout(image_question, 1. - hp.dropout)
+
+    output = mlp(image_question, hp)
+    utils.collect_named_outputs("norms", "output",
+                                tf.norm(output, axis=-1))
+
+    norm_tensors = utils.convert_collection_to_dict("norms")
+    vqa_layers.summarize_tensors(norm_tensors, tag="norms/")
+
+    # Expand dimension 1 and 2
+    return tf.expand_dims(tf.expand_dims(output, axis=1), axis=2)
+
+  def infer(self,
+            features,
+            decode_length=1,
+            beam_size=1,
+            top_beams=1,
+            alpha=0.0,
+            use_tpu=False):
+    """Predict."""
+    del decode_length, beam_size, top_beams, alpha, use_tpu
+    assert features is not None
+    logits, _ = self(features)
+    assert len(logits.get_shape()) == 5
+    logits = tf.squeeze(logits, [1, 2, 3])
+    log_probs = common_layers.log_prob_from_logits(logits)
+    predictions, scores = common_layers.argmax_with_score(log_probs)
+    return {
+        "outputs": predictions,
+        "scores": scores,
+    }
+
+
+@registry.register_model
+class VqaSimpleImageSelfAttention(VqaAttentionBaseline):
+  """Attention baseline model for VQA."""
+
+  def body(self, features):
+    hp = self.hparams
+    # pylint: disable=eval-used
+    if hp.image_input_type == "image":
+      image_feat = vqa_layers.image_embedding(
+          features["inputs"],
+          model_fn=eval(hp.image_model_fn),
+          trainable=hp.train_resnet,
+          is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
+    else:
+      image_feat = features["inputs"]
+
+    image_feat = common_layers.flatten4d3d(image_feat)
+    # image feature self attention
+    # image_feat = tf.nn.dropout(
+    #     image_feat, keep_prob=1.-hp.layer_prepostprocess_dropout)
+
+    # image_feat = image_feat - tf.reduce_mean(
+    #     image_feat, axis=-1, keepdims=True)
+    # image_feat = tf.nn.l2_normalize(image_feat, -1)
+    # utils.collect_named_outputs("norms", "image_feat_after_l2",
+    #                             tf.norm(image_feat, axis=-1))
+
+    image_feat = tf.nn.dropout(image_feat, keep_prob=1.-hp.dropout)
+
+    image_feat = image_encoder(image_feat, hp)
+    utils.collect_named_outputs("norms", "image_feat_encoded",
+                                tf.norm(image_feat, axis=-1))
+    image_feat = common_layers.l2_norm(image_feat)
+    utils.collect_named_outputs("norms", "image_feat_encoded_l2",
+                                tf.norm(image_feat, axis=-1))
+
+    query = question_encoder(features["question"], hp)
+    utils.collect_named_outputs("norms", "query",
+                                tf.norm(query, axis=-1))
+
+    image_ave = attn(image_feat, query, hp)
+    utils.collect_named_outputs("norms", "image_ave",
+                                tf.norm(image_ave, axis=-1))
+
+    image_question = tf.concat([image_ave, query], axis=1)
+    utils.collect_named_outputs("norms", "image_question",
+                                tf.norm(image_question, axis=-1))
+
+    image_question = tf.nn.dropout(image_question, 1. - hp.dropout)
+
+    output = mlp(image_question, hp)
+    utils.collect_named_outputs("norms", "output",
+                                tf.norm(output, axis=-1))
+
+    norm_tensors = utils.convert_collection_to_dict("norms")
+    vqa_layers.summarize_tensors(norm_tensors, tag="norms/")
+
+    # Expand dimension 1 and 2
+    return tf.expand_dims(tf.expand_dims(output, axis=1), axis=2)
+
+
+def image_encoder(image_feat,
+                  hparams,
+                  name="image_encoder",
+                  save_weights_to=None,
+                  make_image_summary=True):
+  """A stack of self attention layers."""
+
+  x = image_feat
+  with tf.variable_scope(name):
+    for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):
+      with tf.variable_scope("layer_%d" % layer):
+        with tf.variable_scope("self_attention"):
+          y = vqa_layers.multihead_attention(
+              common_layers.layer_preprocess(x, hparams),
+              None,
+              None,
+              hparams.attention_key_channels or hparams.image_hidden_size,
+              hparams.attention_value_channels or hparams.image_hidden_size,
+              hparams.image_hidden_size,
+              hparams.num_heads,
+              hparams.attention_dropout,
+              attention_type=hparams.self_attention_type,
+              save_weights_to=save_weights_to,
+              max_relative_position=None,
+              make_image_summary=make_image_summary,
+              dropout_broadcast_dims=None,
+              max_length=None,
+              vars_3d=False,
+              scale_otproduct=hparams.scale_dotproduct)
+          utils.collect_named_outputs("norms", "image_feat_self_attention",
+                                      tf.norm(y, axis=-1))
+          x = common_layers.layer_postprocess(x, y, hparams)
+          utils.collect_named_outputs(
+              "norms", "image_feat_self_attention_zero_add",
+              tf.norm(x, axis=-1))
+        with tf.variable_scope("ffn"):
+          y = common_layers.dense_relu_dense(
+              common_layers.layer_preprocess(x, hparams),
+              hparams.image_filter_size,
+              hparams.image_hidden_size,
+              dropout=hparams.relu_dropout,
+              dropout_broadcast_dims=None)
+          utils.collect_named_outputs("norms", "image_feat_ffn",
+                                      tf.norm(y, axis=-1))
+          x = common_layers.layer_postprocess(x, y, hparams)
+          utils.collect_named_outputs("norms", "image_feat_ffn_zero_add",
+                                      tf.norm(x, axis=-1))
+    # if normalization is done in layer_preprocess, then it should also be done
+    # on the output, since the output can grow very large, being the sum of
+    # a whole stack of unnormalized layer outputs.
+    return common_layers.layer_preprocess(x, hparams)
+
+
+def _get_rnn_cell(hparams):
+  if hparams.rnn_type == "lstm":
+    rnn_cell = tf.contrib.rnn.BasicLSTMCell
+  elif hparams.rnn_type == "lstm_layernorm":
+    rnn_cell = tf.contrib.rnn.LayerNormBasicLSTMCell
+  return tf.contrib.rnn.DropoutWrapper(
+      rnn_cell(hparams.hidden_size),
+      output_keep_prob=1.0-hparams.dropout)
+
+
+def question_encoder(question, hparams, name="encoder"):
+  """Question encoder, run LSTM encoder and get the last output as encoding."""
+  with tf.variable_scope(name, "encoder", values=[question]):
+    question = common_layers.flatten4d3d(question)
+    padding = common_attention.embedding_to_padding(question)
+    length = common_attention.padding_to_length(padding)
+
+    max_question_length = hparams.max_question_length
+    question = question[:, :max_question_length, :]
+    actual_question_length = common_layers.shape_list(question)[1]
+    length = tf.minimum(length, max_question_length)
+    padding = [[0, 0],
+               [0, max_question_length-actual_question_length],
+               [0, 0]]
+    question = tf.pad(question, padding)
+    question_shape = question.get_shape().as_list()
+    question_shape[1] = max_question_length
+    question.set_shape(question_shape)
+
+    # apply tanh dropout on question embedding
+    question = tf.tanh(question)
+    question = tf.nn.dropout(question, keep_prob=1.-hparams.dropout)
+
+    question = [question[:, i, :] for i in range(max_question_length)]
+
+    # rnn_layers = [_get_rnn_cell(hparams)
+    #               for _ in range(hparams.num_rnn_layers)]
+    # rnn_multi_cell = tf.contrib.rnn.MultiRNNCell(rnn_layers)
+    rnn_cell = _get_rnn_cell(hparams)
+    # outputs, _ = tf.nn.dynamic_rnn(
+    #     rnn_cell, question, length, dtype=tf.float32)
+    _, state = tf.nn.static_rnn(rnn_cell, question, sequence_length=length,
+                                dtype=tf.float32)
+    # outputs = [tf.expand_dims(output, axis=1) for output in outputs]
+    # outputs = tf.concat(outputs, axis=1)
+
+    # utils.collect_named_outputs("vqa_attention_debug", "question_output",
+    #                             outputs)
+    # utils.collect_named_outputs("vqa_attention_debug", "question_state",
+    #                             state.h)
+
+    # batch_size = common_layers.shape_list(outputs)[0]
+    # row_indices = tf.range(batch_size)
+    # # length - 1 as index
+    # indices = tf.transpose([row_indices, tf.maximum(length-1, 0)])
+    # last_output = tf.gather_nd(outputs, indices)
+
+    # utils.collect_named_outputs("vqa_attention_debug",
+    #                             "question_final_output", last_output)
+
+  return state.h
+
+
+def attn(image_feat, query, hparams, name="attn"):
+  """Attention on image feature with question as query."""
+  with tf.variable_scope(name, "attn", values=[image_feat, query]):
+    attn_dim = hparams.attn_dim
+    num_glimps = hparams.num_glimps
+    num_channels = common_layers.shape_list(image_feat)[-1]
+    if len(common_layers.shape_list(image_feat)) == 4:
+      image_feat = common_layers.flatten4d3d(image_feat)
+    query = tf.expand_dims(query, 1)
+    image_proj = common_attention.compute_attention_component(
+        image_feat, attn_dim, name="image_proj")
+    query_proj = common_attention.compute_attention_component(
+        query, attn_dim, name="query_proj")
+    h = tf.nn.relu(image_proj + query_proj)
+    h_proj = common_attention.compute_attention_component(
+        h, num_glimps, name="h_proj")
+    p = tf.nn.softmax(h_proj, axis=1)
+    image_ave = tf.matmul(image_feat, p, transpose_a=True)
+    image_ave = tf.reshape(image_ave, [-1, num_channels*num_glimps])
+
+    return image_ave
+
+
+def mlp(feature, hparams, name="mlp"):
+  """Multi layer perceptron with dropout and relu activation."""
+  with tf.variable_scope(name, "mlp", values=[feature]):
+    num_mlp_layers = hparams.num_mlp_layers
+    mlp_dim = hparams.mlp_dim
+    for _ in range(num_mlp_layers):
+      feature = common_layers.dense(feature, mlp_dim, activation=tf.nn.relu)
+      feature = tf.nn.dropout(feature, keep_prob=1.-hparams.dropout)
+    return feature
+
+
+@registry.register_hparams
+def vqa_attention_base():
+  """VQA attention baseline hparams."""
+  hparams = common_hparams.basic_params1()
+  hparams.batch_size = 128
+  hparams.use_fixed_batch_size = True,
+  hparams.optimizer = "Adam"
+  hparams.optimizer_adam_beta1 = 0.9
+  hparams.optimizer_adam_beta2 = 0.999
+  hparams.optimizer_adam_epsilon = 1e-8
+  hparams.weight_decay = 0.
+  hparams.clip_grad_norm = 0.
+  hparams.initializer = "xavier"
+  hparams.learning_rate = 0.5
+  hparams.learning_rate_schedule = "legacy"
+  hparams.learning_rate_warmup_steps = 0
+  hparams.learning_rate_decay_scheme = "exp"
+  hparams.learning_rate_decay_rate = 0.5
+  hparams.learning_rate_decay_steps = 50000
+  hparams.dropout = 0.5
+  hparams.summarize_grads = True
+  hparams.summarize_vars = True
+
+  # not used hparams
+  hparams.label_smoothing = 0.
+  hparams.multiply_embedding_mode = ""
+
+  # add new hparams
+  # preprocess
+  hparams.add_hparam("resize_side", 512)
+  hparams.add_hparam("height", 448)
+  hparams.add_hparam("width", 448)
+  hparams.add_hparam("distort", True)
+
+  hparams.add_hparam("train_resnet", False)
+  hparams.add_hparam("rnn_type", "lstm")
+  hparams.add_hparam("num_rnn_layers", 1)
+  hparams.add_hparam("max_question_length", 15)
+  # lstm hidden size
+  hparams.hidden_size = 512
+
+  hparams.add_hparam("attn_dim", 512)
+  hparams.add_hparam("num_glimps", 2)
+
+  hparams.add_hparam("num_mlp_layers", 1)
+  hparams.add_hparam("mlp_dim", 1024)
+
+  hparams.add_hparam("image_input_type", "image")
+  hparams.add_hparam("image_model_fn", "resnet_v1_152")
+  hparams.add_hparam("image_feat_size", 0)
+
+  # self attention parts
+  hparams.norm_type = "layer"
+  hparams.layer_preprocess_sequence = "n"
+  hparams.layer_postprocess_sequence = "da"
+  hparams.layer_prepostprocess_dropout = 0.3
+  hparams.attention_dropout = 0.1
+  hparams.relu_dropout = 0.1
+  hparams.image_hidden_size = 2048
+  hparams.add_hparam("num_encoder_layers", 1)
+  # Attention-related flags.
+  hparams.add_hparam("num_heads", 8)
+  hparams.add_hparam("attention_key_channels", 0)
+  hparams.add_hparam("attention_value_channels", 0)
+  hparams.add_hparam("image_filter_size", 1024)
+  hparams.add_hparam("self_attention_type", "dot_product")
+  hparams.add_hparam("scale_dotproduct", True)
+
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_base():
+  hparams = vqa_attention_base()
+  hparams.image_input_type = "feature"
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_lstmlayernorm():
+  hparams = vqa_attention_feature_base()
+  hparams.rnn_type = "lstm_layernorm"
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_initializer():
+  hparams = vqa_attention_feature_base()
+  hparams.initializer = "uniform_unit_scaling"
+  hparams.initializer_gain = 1.0
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_batch512():
+  hparams = vqa_attention_feature_base()
+  hparams.batch_size = 512
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_hidden1024():
+  hparams = vqa_attention_feature_base()
+  hparams.hidden_size = 1024
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_imagefeat512():
+  hparams = vqa_attention_feature_base()
+  hparams.image_feat_size = 512
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_imagefeat1024():
+  hparams = vqa_attention_feature_base()
+  hparams.image_feat_size = 1024
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_batch1024_lstmlayernorm():
+  hparams = vqa_attention_feature_lstmlayernorm()
+  hparams.batch_size = 1024
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_numglimps1():
+  hparams = vqa_attention_base()
+  hparams.num_glimps = 1
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_numglimps1():
+  hparams = vqa_attention_feature_base()
+  hparams.num_glimps = 1
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_batch1024_numglimps1():
+  hparams = vqa_attention_feature_numglimps1()
+  hparams.batch_size = 1024
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_batch1024():
+  hparams = vqa_attention_feature_base()
+  hparams.batch_size = 1024
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_batch1024_dnz():
+  hparams = vqa_attention_feature_batch1024()
+  hparams.layer_preprocess_sequence = ""
+  hparams.layer_postprocess_sequence = "dnz"
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_batch1024_dnz_l2():
+  hparams = vqa_attention_feature_batch1024_dnz()
+  hparams.norm_type = "l2"
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_dnz():
+  hparams = vqa_attention_feature_base()
+  hparams.layer_preprocess_sequence = ""
+  hparams.layer_postprocess_sequence = "dnz"
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_dna():
+  hparams = vqa_attention_feature_base()
+  hparams.layer_preprocess_sequence = ""
+  hparams.layer_postprocess_sequence = "dna"
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_dnz_noscaledp():
+  hparams = vqa_attention_feature_dnz()
+  hparams.scale_dotproduct = False
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_dnz_l2():
+  hparams = vqa_attention_feature_dnz()
+  hparams.norm_type = "l2"
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_batch1024_dnz_noscaledp():
+  hparams = vqa_attention_feature_batch1024_dnz()
+  hparams.scale_dotproduct = False
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_batch1024_drop01():
+  hparams = vqa_attention_feature_batch1024()
+  hparams.layer_prepostprocess_dropout = 0.1
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_batch1024_drop01_dna():
+  hparams = vqa_attention_feature_batch1024_drop01()
+  hparams.layer_preprocess_sequence = ""
+  hparams.layer_postprocess_sequence = "dna"
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_drop01_dna():
+  hparams = vqa_attention_feature_batch1024_drop01_dna()
+  hparams.batch_size = 128
+  hparams.image_input_type = "image"
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_batch1024_drop01_dna_concat():
+  hparams = vqa_attention_feature_batch1024_drop01()
+  hparams.layer_preprocess_sequence = ""
+  hparams.layer_postprocess_sequence = "dna"
+  hparams.num_glimps = 1
+  return hparams
+
+
+@registry.register_hparams
+def vqa_attention_feature_nonormalization():
+  hparams = vqa_attention_feature_base()
+  hparams.layer_preprocess_sequence = ""
+  return hparams
+
+
+@registry.register_ranged_hparams
+def vqa_attention_base_range(rhp):
+  """Small range of hyperparameters."""
+  # After starting from base, set intervals for some parameters.
+  rhp.set_float("learning_rate", 0.1, 1.0, scale=rhp.LOG_SCALE)
+  rhp.set_float("clip_grad_norm", 0.1, 10, scale=rhp.LOG_SCALE)
+  rhp.set_discrete("batch_size", [128, 256, 512, 1024])
+  rhp.set_float("weight_decay", 0.0, 1e-4)
+  rhp.set_categorical("rnn_type", ["lstm", "lstm_layernorm"])
diff --git a/tensor2tensor/models/research/vqa_attention_test.py b/tensor2tensor/models/research/vqa_attention_test.py
new file mode 100644
index 000000000..d3a9e5ecb
--- /dev/null
+++ b/tensor2tensor/models/research/vqa_attention_test.py
@@ -0,0 +1,69 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Vqa_attention_baseline tests."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.models.research import vqa_attention
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+class VqaAttentionBaselineTest(tf.test.TestCase):
+
+  def testVqaAttentionBaseline(self):
+
+    batch_size = 3
+    image_size = 448
+    vocab_size = 100
+    num_classes = 10
+    question_length = 5
+    answer_length = 10
+    x = 2 * np.random.rand(batch_size, image_size, image_size, 3) - 1
+    q = np.random.random_integers(
+        1, high=vocab_size - 1, size=(batch_size, question_length, 1, 1))
+    a = np.random.random_integers(
+        0, high=num_classes, size=(batch_size, answer_length, 1, 1))
+    hparams = vqa_attention.vqa_attention_base()
+    p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
+    p_hparams.input_modality["inputs"] = (registry.Modalities.IMAGE, None)
+    p_hparams.input_modality["question"] = (registry.Modalities.SYMBOL,
+                                            vocab_size)
+    p_hparams.target_modality = (registry.Modalities.CLASS_LABEL
+                                 + ":multi_label", num_classes + 1)
+    with self.test_session() as session:
+      features = {
+          "inputs": tf.constant(x, dtype=tf.float32),
+          "question": tf.constant(q, dtype=tf.int32),
+          "targets": tf.constant(a, dtype=tf.int32),
+      }
+      model = vqa_attention.VqaAttentionBaseline(
+          hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
+      logits, losses = model(features)
+      session.run(tf.global_variables_initializer())
+      logits_, losses_ = session.run([logits, losses])
+
+    self.assertEqual(logits_.shape, (batch_size, 1, 1, 1, num_classes + 1))
+    self.assertEqual(losses_["training"].shape, ())
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/research/vqa_recurrent_self_attention.py b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
new file mode 100644
index 000000000..e54de3630
--- /dev/null
+++ b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
@@ -0,0 +1,316 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Recurrent self attention models for VQA."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import vqa_layers
+from tensor2tensor.models.research import universal_transformer
+from tensor2tensor.models.research import universal_transformer_util
+from tensor2tensor.models.research import vqa_attention
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import restore_hook
+
+import tensorflow as tf
+
+from tensorflow.contrib.layers.python.layers import utils
+
+
+@registry.register_model
+class VqaRecurrentSelfAttention(vqa_attention.VqaAttentionBaseline):
+  """Recurrent Self attention both on image and question."""
+
+  @staticmethod
+  def train_hooks():
+    restore_resnet_hook = restore_hook.RestoreHook(
+        # TODO(zichaoy): hard code the path given static function.
+        checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
+        new_model_scope="vqa_recurrent_self_attention/body/",
+        old_model_scope="resnet_v1_152/",
+    )
+    return [restore_resnet_hook]
+
+  def body(self, features):
+    hp = self.hparams
+    # pylint: disable=eval-used
+    if hp.image_input_type == "image":
+      image_feat = vqa_layers.image_embedding(
+          features["inputs"],
+          model_fn=eval(hp.image_model_fn),
+          trainable=hp.train_resnet,
+          is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
+    else:
+      image_feat = features["inputs"]
+
+    image_feat = common_layers.flatten4d3d(image_feat)
+    image_feat = common_layers.dense(image_feat, hp.hidden_size)
+    utils.collect_named_outputs("norms", "image_feat_after_proj",
+                                tf.norm(image_feat, axis=-1))
+
+    question = common_layers.flatten4d3d(features["question"])
+    utils.collect_named_outputs("norms", "question_embedding",
+                                tf.norm(question, axis=-1))
+    (encoder_input, encoder_self_attention_bias,
+     encoder_decoder_attention_bias) = prepare_image_question_encoder(
+         image_feat, question, hp)
+
+    encoder_input = tf.nn.dropout(
+        encoder_input, keep_prob=1.-hp.layer_prepostprocess_dropout)
+
+    encoder_output, _ = recurrent_transformer_decoder(
+        encoder_input, None, encoder_self_attention_bias, None,
+        hp, name="encoder")
+    utils.collect_named_outputs(
+        "norms", "encoder_output", tf.norm(encoder_output, axis=-1))
+
+    # scale query by sqrt(hidden_size)
+    query = tf.get_variable("query", [hp.hidden_size]) * hp.hidden_size **0.5
+    query = tf.expand_dims(tf.expand_dims(query, axis=0), axis=0)
+    batch_size = common_layers.shape_list(encoder_input)[0]
+    query = tf.tile(query, [batch_size, 1, 1])
+    query = tf.nn.dropout(
+        query, keep_prob=1.-hp.layer_prepostprocess_dropout)
+
+    decoder_output, _ = recurrent_transformer_decoder(
+        query, encoder_output, None, encoder_decoder_attention_bias,
+        hp, name="decoder")
+    utils.collect_named_outputs("norms", "decoder_output",
+                                tf.norm(decoder_output, axis=-1))
+
+    norm_tensors = utils.convert_collection_to_dict("norms")
+    vqa_layers.summarize_tensors(norm_tensors, tag="norms/")
+
+    # Expand dimension 1 and 2
+    return tf.expand_dims(decoder_output, axis=1)
+
+
+def prepare_image_question_encoder(image_feat, question, hparams):
+  """Prepare encoder.
+
+  Args:
+    image_feat: a Tensor.
+    question: a Tensor.
+    hparams: run hyperparameters
+
+  Returns:
+    encoder_input: a Tensor, bottom of encoder stack
+    encoder_self_attention_bias: a bias tensor for use in encoder self-attention
+  """
+
+  encoder_input = tf.concat([image_feat, question], axis=1)
+  encoder_padding = common_attention.embedding_to_padding(encoder_input)
+  ignore_padding = common_attention.attention_bias_ignore_padding(
+      encoder_padding)
+  encoder_self_attention_bias = ignore_padding
+  encoder_decoder_attention_bias = ignore_padding
+  # Usual case - not a packed dataset.
+  if hparams.pos == "timing":
+    question = common_attention.add_timing_signal_1d(question)
+  elif hparams.pos == "emb":
+    question = common_attention.add_positional_embedding(
+        question, hparams.max_length, "inputs_positional_embedding",
+        None)
+  encoder_input = tf.concat([image_feat, question], axis=1)
+
+  return (encoder_input, encoder_self_attention_bias,
+          encoder_decoder_attention_bias)
+
+
+def recurrent_transformer_decoder(
+    decoder_input,
+    encoder_output,
+    decoder_self_attention_bias,
+    encoder_decoder_attention_bias,
+    hparams,
+    name="decoder",
+    nonpadding=None,
+    save_weights_to=None,
+    make_image_summary=True):
+  """Recurrent decoder function."""
+  x = decoder_input
+  attention_dropout_broadcast_dims = (
+      common_layers.comma_separated_string_to_integer_list(
+          getattr(hparams, "attention_dropout_broadcast_dims", "")))
+  with tf.variable_scope(name):
+    ffn_unit = functools.partial(
+        # use encoder ffn, since decoder ffn use left padding
+        universal_transformer_util.transformer_encoder_ffn_unit,
+        hparams=hparams,
+        nonpadding_mask=nonpadding)
+
+    attention_unit = functools.partial(
+        universal_transformer_util.transformer_decoder_attention_unit,
+        hparams=hparams,
+        encoder_output=encoder_output,
+        decoder_self_attention_bias=decoder_self_attention_bias,
+        encoder_decoder_attention_bias=encoder_decoder_attention_bias,
+        attention_dropout_broadcast_dims=attention_dropout_broadcast_dims,
+        save_weights_to=save_weights_to,
+        make_image_summary=make_image_summary)
+
+    x, extra_output = universal_transformer_util.universal_transformer_layer(
+        x, hparams, ffn_unit, attention_unit)
+
+    return common_layers.layer_preprocess(x, hparams), extra_output
+
+
+@registry.register_hparams
+def vqa_recurrent_self_attention_base():
+  """VQA attention baseline hparams."""
+  hparams = universal_transformer.universal_transformer_base()
+  hparams.batch_size = 1024
+  hparams.use_fixed_batch_size = True
+  hparams.weight_decay = 0.
+  hparams.clip_grad_norm = 0.
+  # use default initializer
+  # hparams.initializer = "xavier"
+  hparams.learning_rate_schedule = (
+      "constant*linear_warmup*rsqrt_normalized_decay")
+  hparams.learning_rate_warmup_steps = 8000
+  hparams.learning_rate_constant = 7e-4
+  hparams.learning_rate_decay_rate = 0.5
+  hparams.learning_rate_decay_steps = 50000
+  # hparams.dropout = 0.5
+  hparams.summarize_grads = True
+  hparams.summarize_vars = True
+
+  # not used hparams
+  hparams.label_smoothing = 0.1
+  hparams.multiply_embedding_mode = "sqrt_depth"
+
+  # add new hparams
+  # use raw image as input
+  hparams.add_hparam("image_input_type", "feature")
+  hparams.add_hparam("image_model_fn", "resnet_v1_152")
+  hparams.add_hparam("resize_side", 512)
+  hparams.add_hparam("height", 448)
+  hparams.add_hparam("width", 448)
+  hparams.add_hparam("distort", True)
+  hparams.add_hparam("train_resnet", False)
+
+  # question hidden size
+  # hparams.hidden_size = 512
+  # hparams.filter_size = 1024
+  # hparams.num_hidden_layers = 4
+
+  # self attention parts
+  # hparams.norm_type = "layer"
+  # hparams.layer_preprocess_sequence = "n"
+  # hparams.layer_postprocess_sequence = "da"
+  # hparams.layer_prepostprocess_dropout = 0.1
+  # hparams.attention_dropout = 0.1
+  # hparams.relu_dropout = 0.1
+  # hparams.add_hparam("pos", "timing")
+  # hparams.add_hparam("num_encoder_layers", 0)
+  # hparams.add_hparam("num_decoder_layers", 0)
+  # hparams.add_hparam("num_heads", 8)
+  # hparams.add_hparam("attention_key_channels", 0)
+  # hparams.add_hparam("attention_value_channels", 0)
+  # hparams.add_hparam("self_attention_type", "dot_product")
+
+  # iterative part
+  hparams.transformer_ffn_type = "fc"
+
+  return hparams
+
+
+@registry.register_hparams
+def vqa_recurrent_self_attention_small():
+  hparams = vqa_recurrent_self_attention_base()
+  hparams.learning_rate_constant = 1e-3
+  hparams.hidden_size = 512
+  hparams.filter_size = 2048
+  hparams.num_heads = 8
+  hparams.layer_prepostprocess_dropout = 0.1
+  return hparams
+
+
+@registry.register_hparams
+def vqa_recurrent_self_attention_big():
+  hparams = vqa_recurrent_self_attention_base()
+  hparams.learning_rate_constant = 5e-4
+  hparams.hidden_size = 2048
+  hparams.filter_size = 8192
+  return hparams
+
+
+@registry.register_hparams
+def vqa_recurrent_self_attention_big_l4():
+  hparams = vqa_recurrent_self_attention_big()
+  hparams.num_rec_steps = 4
+  return hparams
+
+
+@registry.register_hparams
+def vqa_recurrent_self_attention_highway():
+  hparams = vqa_recurrent_self_attention_base()
+  hparams.recurrence_type = "highway"
+  return hparams
+
+
+@registry.register_hparams
+def vqa_recurrent_self_attention_gru():
+  hparams = vqa_recurrent_self_attention_base()
+  hparams.recurrence_type = "gru"
+  return hparams
+
+
+@registry.register_hparams
+def vqa_recurrent_self_attention_l8():
+  hparams = vqa_recurrent_self_attention_base()
+  hparams.num_rec_steps = 8
+  return hparams
+
+
+@registry.register_hparams
+def vqa_recurrent_self_attention_mix_before_ut():
+  hparams = vqa_recurrent_self_attention_base()
+  hparams.mix_with_transformer = "before_ut"
+  return hparams
+
+
+@registry.register_hparams
+def vqa_recurrent_self_attention_l4():
+  hparams = vqa_recurrent_self_attention_base()
+  hparams.num_rec_steps = 4
+  return hparams
+
+
+@registry.register_hparams
+def vqa_recurrent_self_attention_ls2():
+  hparams = vqa_recurrent_self_attention_base()
+  hparams.label_smoothing = 0.2
+  return hparams
+
+
+@registry.register_hparams
+def vqa_recurrent_self_attention_drop1():
+  hparams = vqa_recurrent_self_attention_base()
+  hparams.layer_prepostprocess_dropout = 0.1
+  return hparams
+
+
+@registry.register_hparams
+def vqa_recurrent_self_attention_drop3():
+  hparams = vqa_recurrent_self_attention_base()
+  hparams.relu_dropout = 0.3
+  hparams.attention_dropout = 0.3
+  return hparams
diff --git a/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
new file mode 100644
index 000000000..f556b1c78
--- /dev/null
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -0,0 +1,823 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Self attention models for VQA."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_hparams
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import vqa_layers
+from tensor2tensor.models.research import vqa_attention
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import restore_hook
+
+import tensorflow as tf
+
+from tensorflow.contrib.layers.python.layers import utils
+
+
+@registry.register_model
+class VqaSelfAttention(vqa_attention.VqaAttentionBaseline):
+  """Self attention both on image and question."""
+
+  @staticmethod
+  def train_hooks():
+    restore_resnet_hook = restore_hook.RestoreHook(
+        # TODO(zichaoy): hard code the path given static function.
+        checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
+        new_model_scope="vqa_self_attention/body/",
+        old_model_scope="resnet_v1_152/",
+    )
+    return [restore_resnet_hook]
+
+  def body(self, features):
+    hp = self.hparams
+    # pylint: disable=eval-used
+    if hp.image_input_type == "image":
+      image_feat = vqa_layers.image_embedding(
+          features["inputs"],
+          model_fn=eval(hp.image_model_fn),
+          trainable=hp.train_resnet,
+          is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
+    else:
+      image_feat = features["inputs"]
+
+    image_feat = common_layers.flatten4d3d(image_feat)
+    image_hidden_size = hp.image_hidden_size or hp.hidden_size
+    if hp.image_feat_preprocess_proj:
+      image_feat = common_layers.dense(image_feat, image_hidden_size)
+      utils.collect_named_outputs("norms", "image_feat_after_proj",
+                                  tf.norm(image_feat, axis=-1))
+    else:
+      assert image_hidden_size == 2048
+
+    image_feat = tf.nn.dropout(
+        image_feat, keep_prob=1.-hp.layer_prepostprocess_dropout)
+
+    if hp.image_feat_encode:
+      image_feat = image_encoder(image_feat, hp)
+      utils.collect_named_outputs("norms", "image_feat_encoded",
+                                  tf.norm(image_feat, axis=-1))
+    else:
+      image_feat = common_layers.layer_norm(image_feat)
+      utils.collect_named_outputs("norms", "image_feat_after_layer",
+                                  tf.norm(image_feat, axis=-1))
+
+    question = common_layers.flatten4d3d(features["question"])
+    utils.collect_named_outputs("norms", "question_embedding",
+                                tf.norm(question, axis=-1))
+    question, question_self_attention_bias = prepare_question_encoder(
+        question, hp)
+    question = tf.nn.dropout(
+        question, keep_prob=1.-hp.layer_prepostprocess_dropout)
+    query = question_encoder(question, question_self_attention_bias, hp)
+    utils.collect_named_outputs(
+        "norms", "query_encode", tf.norm(query, axis=-1))
+    query = (query + tf.expand_dims(
+        tf.squeeze(question_self_attention_bias, [1, 2]), axis=2))
+    query = tf.reduce_max(query, axis=1)
+    utils.collect_named_outputs(
+        "norms", "query_maxpool", tf.norm(query, axis=-1))
+
+    # query = common_layers.l2_norm(query)
+    # utils.collect_named_outputs("norms", "query_after_l2",
+    #                             tf.norm(query, axis=-1))
+
+    image_ave = attn(image_feat, query, hp)
+    utils.collect_named_outputs("norms", "image_ave",
+                                tf.norm(image_ave, axis=-1))
+
+    if hp.multimodal_combine == "concat":
+      image_question = tf.concat([image_ave, query], axis=1)
+    elif hp.multimodal_combine == "sum":
+      image_question = image_ave + query
+    elif hp.multimodal_combine == "product":
+      image_question = image_ave * query
+
+    utils.collect_named_outputs("norms", "image_question",
+                                tf.norm(image_question, axis=-1))
+
+    image_question = tf.nn.dropout(image_question, 1. - hp.dropout)
+
+    output = mlp(image_question, hp)
+    utils.collect_named_outputs("norms", "output",
+                                tf.norm(output, axis=-1))
+
+    norm_tensors = utils.convert_collection_to_dict("norms")
+    vqa_layers.summarize_tensors(norm_tensors, tag="norms/")
+
+    # Expand dimension 1 and 2
+    return tf.expand_dims(tf.expand_dims(output, axis=1), axis=2)
+
+
+@registry.register_model
+class VqaCombinedSelfAttention(VqaSelfAttention):
+  """Combined Self attention both on image and question."""
+
+  @staticmethod
+  def train_hooks():
+    restore_resnet_hook = restore_hook.RestoreHook(
+        # TODO(zichaoy): hard code the path given static function.
+        checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
+        new_model_scope="vqa_combined_self_attention/body/",
+        old_model_scope="resnet_v1_152/",
+    )
+    return [restore_resnet_hook]
+
+  def body(self, features):
+    hp = self.hparams
+    # pylint: disable=eval-used
+    if hp.image_input_type == "image":
+      image_feat = vqa_layers.image_embedding(
+          features["inputs"],
+          model_fn=eval(hp.image_model_fn),
+          trainable=hp.train_resnet,
+          is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
+    else:
+      image_feat = features["inputs"]
+
+    image_feat = common_layers.flatten4d3d(image_feat)
+    image_hidden_size = hp.hidden_size
+    image_feat = common_layers.dense(image_feat, image_hidden_size)
+    utils.collect_named_outputs("norms", "image_feat_after_proj",
+                                tf.norm(image_feat, axis=-1))
+
+    question = common_layers.flatten4d3d(features["question"])
+    utils.collect_named_outputs("norms", "question_embedding",
+                                tf.norm(question, axis=-1))
+    (encoder_input, encoder_self_attention_bias,
+     encoder_decoder_attention_bias) = prepare_image_question_encoder(
+         image_feat, question, hp)
+    encoder_input = tf.nn.dropout(
+        encoder_input, keep_prob=1.-hp.layer_prepostprocess_dropout)
+    encoder_output = image_question_encoder(
+        encoder_input, encoder_self_attention_bias, hp)
+    utils.collect_named_outputs(
+        "norms", "encoder_output", tf.norm(encoder_output, axis=-1))
+
+    # scale query by sqrt(hidden_size)
+    query = tf.get_variable("query", [hp.hidden_size]) * hp.hidden_size **0.5
+    query = tf.expand_dims(tf.expand_dims(query, axis=0), axis=0)
+    batch_size = common_layers.shape_list(encoder_input)[0]
+    query = tf.tile(query, [batch_size, 1, 1])
+    query = tf.nn.dropout(
+        query, keep_prob=1.-hp.layer_prepostprocess_dropout)
+
+    decoder_output = decoder(
+        query, encoder_output, None, encoder_decoder_attention_bias, hp)
+    utils.collect_named_outputs("norms", "decoder_output",
+                                tf.norm(decoder_output, axis=-1))
+
+    norm_tensors = utils.convert_collection_to_dict("norms")
+    vqa_layers.summarize_tensors(norm_tensors, tag="norms/")
+
+    # Expand dimension 1 and 2
+    return tf.expand_dims(decoder_output, axis=1)
+
+
+@registry.register_model
+class VqaIterativeCombinedSelfAttention(VqaSelfAttention):
+  """Combined Self attention both on image and question."""
+
+  @staticmethod
+  def train_hooks():
+    restore_resnet_hook = restore_hook.RestoreHook(
+        # TODO(zichaoy): hard code the path given static function.
+        checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
+        new_model_scope="vqa_combined_self_attention/body/",
+        old_model_scope="resnet_v1_152/",
+    )
+    return [restore_resnet_hook]
+
+  def body(self, features):
+    hp = self.hparams
+    # pylint: disable=eval-used
+    if hp.image_input_type == "image":
+      image_feat = vqa_layers.image_embedding(
+          features["inputs"],
+          model_fn=eval(hp.image_model_fn),
+          trainable=hp.train_resnet,
+          is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
+    else:
+      image_feat = features["inputs"]
+
+    image_feat = common_layers.flatten4d3d(image_feat)
+    image_hidden_size = hp.hidden_size
+    image_feat = common_layers.dense(image_feat, image_hidden_size)
+    utils.collect_named_outputs("norms", "image_feat_after_proj",
+                                tf.norm(image_feat, axis=-1))
+
+    question = common_layers.flatten4d3d(features["question"])
+    utils.collect_named_outputs("norms", "question_embedding",
+                                tf.norm(question, axis=-1))
+    (encoder_input, encoder_self_attention_bias,
+     encoder_decoder_attention_bias) = prepare_image_question_encoder(
+         image_feat, question, hp)
+    encoder_input = tf.nn.dropout(
+        encoder_input, keep_prob=1.-hp.layer_prepostprocess_dropout)
+
+    # scale query by sqrt(hidden_size)
+    query = tf.get_variable("query", [hp.hidden_size]) * hp.hidden_size **0.5
+    query = tf.expand_dims(tf.expand_dims(query, axis=0), axis=0)
+    batch_size = common_layers.shape_list(encoder_input)[0]
+    query = tf.tile(query, [batch_size, 1, 1])
+    query = tf.nn.dropout(
+        query, keep_prob=1.-hp.layer_prepostprocess_dropout)
+
+    decoder_output = iterative_encoder_decoder(
+        encoder_input,
+        encoder_self_attention_bias,
+        encoder_decoder_attention_bias,
+        query,
+        hp)
+
+    utils.collect_named_outputs("norms", "decoder_output",
+                                tf.norm(decoder_output, axis=-1))
+
+    norm_tensors = utils.convert_collection_to_dict("norms")
+    vqa_layers.summarize_tensors(norm_tensors, tag="norms/")
+
+    # Expand dimension 1 and 2
+    return tf.expand_dims(decoder_output, axis=1)
+
+
+def image_encoder(image_feat,
+                  hparams,
+                  name="image_encoder",
+                  save_weights_to=None,
+                  make_image_summary=True):
+  """A stack of self attention layers."""
+
+  x = image_feat
+  image_hidden_size = hparams.image_hidden_size or hparams.hidden_size
+  image_filter_size = hparams.image_filter_size or hparams.filter_size
+  with tf.variable_scope(name):
+    for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):
+      with tf.variable_scope("layer_%d" % layer):
+        with tf.variable_scope("self_attention"):
+          y = vqa_layers.multihead_attention(
+              common_layers.layer_preprocess(x, hparams),
+              None,
+              None,
+              hparams.attention_key_channels or image_hidden_size,
+              hparams.attention_value_channels or image_hidden_size,
+              image_hidden_size,
+              hparams.num_heads,
+              hparams.attention_dropout,
+              attention_type=hparams.image_self_attention_type,
+              save_weights_to=save_weights_to,
+              make_image_summary=make_image_summary,
+              scale_dotproduct=hparams.scale_dotproduct,
+          )
+          utils.collect_named_outputs(
+              "norms", "image_feat_self_attention_%d"%(layer),
+              tf.norm(y, axis=-1))
+          x = common_layers.layer_postprocess(x, y, hparams)
+          utils.collect_named_outputs(
+              "norms", "image_feat_self_attention_postprocess_%d"%(layer),
+              tf.norm(x, axis=-1))
+        with tf.variable_scope("ffn"):
+          y = common_layers.dense_relu_dense(
+              common_layers.layer_preprocess(x, hparams),
+              image_filter_size,
+              image_hidden_size,
+              dropout=hparams.relu_dropout,
+          )
+          utils.collect_named_outputs(
+              "norms", "image_feat_ffn_%d"%(layer), tf.norm(y, axis=-1))
+          x = common_layers.layer_postprocess(x, y, hparams)
+          utils.collect_named_outputs(
+              "norms", "image_feat_ffn_postprocess_%d"%(layer),
+              tf.norm(x, axis=-1))
+    # if normalization is done in layer_preprocess, then it should also be done
+    # on the output, since the output can grow very large, being the sum of
+    # a whole stack of unnormalized layer outputs.
+    return common_layers.layer_preprocess(x, hparams)
+
+
+def prepare_question_encoder(inputs, hparams):
+  """Prepare question encoder.
+
+  Args:
+    inputs: a Tensor.
+    hparams: run hyperparameters
+
+  Returns:
+    encoder_input: a Tensor, bottom of encoder stack
+    encoder_self_attention_bias: a bias tensor for use in encoder self-attention
+  """
+  encoder_input = inputs
+  # Usual case - not a packed dataset.
+  encoder_padding = common_attention.embedding_to_padding(encoder_input)
+  ignore_padding = common_attention.attention_bias_ignore_padding(
+      encoder_padding)
+  encoder_self_attention_bias = ignore_padding
+  if hparams.pos == "timing":
+    encoder_input = common_attention.add_timing_signal_1d(encoder_input)
+  elif hparams.pos == "emb":
+    encoder_input = common_attention.add_positional_embedding(
+        encoder_input, hparams.max_length, "inputs_positional_embedding",
+        None)
+  return (encoder_input, encoder_self_attention_bias)
+
+
+def question_encoder(question,
+                     question_self_attention_bias,
+                     hparams,
+                     name="question_encoder",
+                     save_weights_to=None,
+                     make_image_summary=True):
+  """A stack of self attention layers."""
+  x = question
+  with tf.variable_scope(name):
+    for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):
+      with tf.variable_scope("layer_%d" % layer):
+        with tf.variable_scope("self_attention"):
+          y = vqa_layers.multihead_attention(
+              common_layers.layer_preprocess(x, hparams),
+              None,
+              question_self_attention_bias,
+              hparams.attention_key_channels or hparams.hidden_size,
+              hparams.attention_value_channels or hparams.hidden_size,
+              hparams.hidden_size,
+              hparams.num_heads,
+              hparams.attention_dropout,
+              attention_type=hparams.question_self_attention_type,
+              block_length=hparams.block_length,
+              save_weights_to=save_weights_to,
+              make_image_summary=make_image_summary,
+              scale_dotproduct=hparams.scale_dotproduct,
+          )
+          utils.collect_named_outputs(
+              "norms", "query_self_attention_%d"%(layer),
+              tf.norm(y, axis=-1))
+          x = common_layers.layer_postprocess(x, y, hparams)
+          utils.collect_named_outputs(
+              "norms", "query_self_attention_postprocess_%d"%(layer),
+              tf.norm(x, axis=-1))
+        with tf.variable_scope("ffn"):
+          y = common_layers.dense_relu_dense(
+              common_layers.layer_preprocess(x, hparams),
+              hparams.filter_size,
+              hparams.hidden_size,
+              dropout=hparams.relu_dropout,
+              )
+          utils.collect_named_outputs(
+              "norms", "query_ffn_%d"%(layer), tf.norm(y, axis=-1))
+          x = common_layers.layer_postprocess(x, y, hparams)
+          utils.collect_named_outputs(
+              "norms", "query_ffn_postprocess_%d"%(layer),
+              tf.norm(x, axis=-1))
+    # if normalization is done in layer_preprocess, then it should also be done
+    # on the output, since the output can grow very large, being the sum of
+    # a whole stack of unnormalized layer outputs.
+    return common_layers.layer_preprocess(x, hparams)
+
+
+def attn(image_feat,
+         query,
+         hparams,
+         name="attn",
+         save_weights_to=None,
+         make_image_summary=True):
+  """Attention on image feature with question as query."""
+  with tf.variable_scope(name, "attn", values=[image_feat, query]):
+    total_key_depth = hparams.attention_key_channels or hparams.hidden_size
+    total_value_depth = hparams.attention_value_channels or hparams.hidden_size
+    num_heads = hparams.num_heads
+    query = tf.expand_dims(query, 1)
+    q, k, v = common_attention.compute_qkv(
+        query,
+        image_feat,
+        total_key_depth,
+        total_value_depth,
+    )
+    q = common_attention.split_heads(q, num_heads)
+    k = common_attention.split_heads(k, num_heads)
+    v = common_attention.split_heads(v, num_heads)
+
+    if hparams.scale_dotproduct:
+      key_depth_per_head = total_key_depth // num_heads
+      q *= key_depth_per_head**-0.5
+
+    # image_feat is input as v
+    x = common_attention.dot_product_attention(
+        q, k, v, None,
+        dropout_rate=hparams.attention_dropout,
+        image_shapes=None,
+        save_weights_to=save_weights_to,
+        make_image_summary=make_image_summary)
+    x = common_attention.combine_heads(x)
+
+    return tf.squeeze(x, axis=1)
+
+
+def mlp(feature, hparams, name="mlp"):
+  """Multi layer perceptron with dropout and relu activation."""
+  with tf.variable_scope(name, "mlp", values=[feature]):
+    num_mlp_layers = hparams.num_mlp_layers
+    mlp_size = hparams.mlp_size
+    for _ in range(num_mlp_layers):
+      feature = common_layers.dense(feature, mlp_size, activation=None)
+      utils.collect_named_outputs("norms", "mlp_feature",
+                                  tf.norm(feature, axis=-1))
+      feature = common_layers.layer_norm(feature)
+      feature = tf.nn.relu(feature)
+      feature = tf.nn.dropout(feature, keep_prob=1.-hparams.dropout)
+    return feature
+
+
+def prepare_image_question_encoder(image_feat, question, hparams):
+  """Prepare encoder.
+
+  Args:
+    image_feat: a Tensor.
+    question: a Tensor.
+    hparams: run hyperparameters
+
+  Returns:
+    encoder_input: a Tensor, bottom of encoder stack
+    encoder_self_attention_bias: a bias tensor for use in encoder self-attention
+  """
+
+  encoder_input = tf.concat([image_feat, question], axis=1)
+  encoder_padding = common_attention.embedding_to_padding(encoder_input)
+  ignore_padding = common_attention.attention_bias_ignore_padding(
+      encoder_padding)
+  encoder_self_attention_bias = ignore_padding
+  encoder_decoder_attention_bias = ignore_padding
+  # Usual case - not a packed dataset.
+  if hparams.pos == "timing":
+    question = common_attention.add_timing_signal_1d(question)
+  elif hparams.pos == "emb":
+    question = common_attention.add_positional_embedding(
+        question, hparams.max_length, "inputs_positional_embedding",
+        None)
+  encoder_input = tf.concat([image_feat, question], axis=1)
+
+  return (encoder_input, encoder_self_attention_bias,
+          encoder_decoder_attention_bias)
+
+
+def image_question_encoder(encoder_inputs,
+                           encoder_self_attention_bias,
+                           hparams,
+                           query=None,
+                           name="image_question_encoder",
+                           save_weights_to=None,
+                           make_image_summary=True):
+  """A stack of self attention layers."""
+  x = encoder_inputs
+  with tf.variable_scope(name):
+    for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):
+      with tf.variable_scope("layer_%d" % layer):
+        with tf.variable_scope("self_attention"):
+          y = vqa_layers.multihead_attention(
+              common_layers.layer_preprocess(x, hparams),
+              None,
+              encoder_self_attention_bias,
+              hparams.attention_key_channels or hparams.hidden_size,
+              hparams.attention_value_channels or hparams.hidden_size,
+              hparams.hidden_size,
+              hparams.num_heads,
+              hparams.attention_dropout,
+              attention_type=hparams.self_attention_type,
+              block_length=hparams.block_length,
+              save_weights_to=save_weights_to,
+              make_image_summary=make_image_summary,
+              scale_dotproduct=hparams.scale_dotproduct,
+          )
+          utils.collect_named_outputs(
+              "norms", "encoder_self_attention_%d"%(layer),
+              tf.norm(y, axis=-1))
+          x = common_layers.layer_postprocess(x, y, hparams)
+          utils.collect_named_outputs(
+              "norms", "encoder_self_attention_postprocess_%d"%(layer),
+              tf.norm(x, axis=-1))
+        if query is not None:
+          with tf.variable_scope("encdec_attention"):
+            y = common_attention.multihead_attention(
+                common_layers.layer_preprocess(x, hparams),
+                query,
+                None,
+                hparams.attention_key_channels or hparams.hidden_size,
+                hparams.attention_value_channels or hparams.hidden_size,
+                hparams.hidden_size,
+                hparams.num_heads,
+                hparams.attention_dropout,
+                attention_type=hparams.self_attention_type,
+                block_length=hparams.block_length,
+                save_weights_to=save_weights_to,
+                make_image_summary=make_image_summary,
+                scale_dotproduct=hparams.scale_dotproduct,
+            )
+            utils.collect_named_outputs(
+                "norms",
+                "encoder_decoder_attention_%d"%(layer),
+                tf.norm(y, axis=-1))
+            x = common_layers.layer_postprocess(x, y, hparams)
+            utils.collect_named_outputs(
+                "norms",
+                "encoder_decoder_attention_post_%d"%(layer),
+                tf.norm(x, axis=-1))
+        with tf.variable_scope("ffn"):
+          y = common_layers.dense_relu_dense(
+              common_layers.layer_preprocess(x, hparams),
+              hparams.filter_size,
+              hparams.hidden_size,
+              dropout=hparams.relu_dropout,
+              )
+          utils.collect_named_outputs(
+              "norms", "encoder_ffn_%d"%(layer), tf.norm(y, axis=-1))
+          x = common_layers.layer_postprocess(x, y, hparams)
+          utils.collect_named_outputs(
+              "norms", "encoder_ffn_postprocess_%d"%(layer),
+              tf.norm(x, axis=-1))
+    # if normalization is done in layer_preprocess, then it should also be done
+    # on the output, since the output can grow very large, being the sum of
+    # a whole stack of unnormalized layer outputs.
+    return common_layers.layer_preprocess(x, hparams)
+
+
+def decoder(decoder_input,
+            encoder_output,
+            decoder_self_attention_bias,
+            encoder_decoder_attention_bias,
+            hparams,
+            name="decoder",
+            save_weights_to=None,
+            make_image_summary=True,):
+  """A stack of transformer layers.
+
+  Args:
+    decoder_input: a Tensor
+    encoder_output: a Tensor
+    decoder_self_attention_bias: bias Tensor for self-attention
+      (see common_attention.attention_bias())
+    encoder_decoder_attention_bias: bias Tensor for encoder-decoder attention
+      (see common_attention.attention_bias())
+    hparams: hyperparameters for model
+    name: a string
+    save_weights_to: an optional dictionary to capture attention weights
+      for visualization; the weights tensor will be appended there under
+      a string key created from the variable scope (including name).
+    make_image_summary: Whether to make an attention image summary.
+
+  Returns:
+    y: a Tensors
+  """
+  x = decoder_input
+  with tf.variable_scope(name):
+    for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
+      layer_name = "layer_%d" % layer
+      with tf.variable_scope(layer_name):
+        with tf.variable_scope("self_attention"):
+          y = common_attention.multihead_attention(
+              common_layers.layer_preprocess(x, hparams),
+              None,
+              decoder_self_attention_bias,
+              hparams.attention_key_channels or hparams.hidden_size,
+              hparams.attention_value_channels or hparams.hidden_size,
+              hparams.hidden_size,
+              hparams.num_heads,
+              hparams.attention_dropout,
+              attention_type=hparams.self_attention_type,
+              save_weights_to=save_weights_to,
+              make_image_summary=make_image_summary,
+              )
+          utils.collect_named_outputs("norms",
+                                      "decoder_self_attention_%d"%(layer),
+                                      tf.norm(y, axis=-1))
+          x = common_layers.layer_postprocess(x, y, hparams)
+          utils.collect_named_outputs("norms",
+                                      "decoder_self_attention_post_%d"%(layer),
+                                      tf.norm(x, axis=-1))
+        if encoder_output is not None:
+          with tf.variable_scope("encdec_attention"):
+            y = common_attention.multihead_attention(
+                common_layers.layer_preprocess(x, hparams),
+                encoder_output,
+                encoder_decoder_attention_bias,
+                hparams.attention_key_channels or hparams.hidden_size,
+                hparams.attention_value_channels or hparams.hidden_size,
+                hparams.hidden_size,
+                hparams.num_heads,
+                hparams.attention_dropout,
+                save_weights_to=save_weights_to,
+                make_image_summary=make_image_summary,
+                )
+            utils.collect_named_outputs(
+                "norms",
+                "decoder_encoder_attention_%d"%(layer),
+                tf.norm(y, axis=-1))
+            x = common_layers.layer_postprocess(x, y, hparams)
+            utils.collect_named_outputs(
+                "norms",
+                "decoder_encoder_attention_post_%d"%(layer),
+                tf.norm(x, axis=-1))
+        with tf.variable_scope("ffn"):
+          y = common_layers.dense_relu_dense(
+              common_layers.layer_preprocess(x, hparams),
+              hparams.filter_size,
+              hparams.hidden_size,
+              dropout=hparams.relu_dropout,
+          )
+          utils.collect_named_outputs("norms", "decoder_ffn_%d"%(layer),
+                                      tf.norm(y, axis=-1))
+          x = common_layers.layer_postprocess(x, y, hparams)
+          utils.collect_named_outputs("norms", "decoder_ffn_post_%d"%(layer),
+                                      tf.norm(x, axis=-1))
+    # if normalization is done in layer_preprocess, then it should also be done
+    # on the output, since the output can grow very large, being the sum of
+    # a whole stack of unnormalized layer outputs.
+    return common_layers.layer_preprocess(x, hparams)
+
+
+def iterative_encoder_decoder(encoder_input,
+                              encoder_self_attention_bias,
+                              encoder_decoder_attention_bias,
+                              query,
+                              hparams):
+  """Iterative encoder decoder."""
+  for _ in xrange(hparams.num_rec_steps):
+    with tf.variable_scope("step", reuse=tf.AUTO_REUSE):
+      encoder_output = image_question_encoder(
+          encoder_input,
+          encoder_self_attention_bias,
+          hparams,
+          query)
+
+      decoder_output = decoder(
+          query,
+          encoder_output,
+          None,
+          encoder_decoder_attention_bias,
+          hparams)
+
+      encoder_input = encoder_output
+      query = decoder_output
+
+      return decoder_output
+
+
+@registry.register_hparams
+def vqa_self_attention_base():
+  """VQA attention baseline hparams."""
+  hparams = common_hparams.basic_params1()
+  hparams.batch_size = 128
+  hparams.use_fixed_batch_size = True,
+  hparams.optimizer = "Adam"
+  hparams.optimizer_adam_beta1 = 0.9
+  hparams.optimizer_adam_beta2 = 0.997
+  hparams.optimizer_adam_epsilon = 1e-9
+  hparams.weight_decay = 0.
+  hparams.clip_grad_norm = 0.
+  hparams.initializer = "xavier"
+  hparams.learning_rate_schedule = (
+      "constant*linear_warmup*rsqrt_normalized_decay")
+  hparams.learning_rate_warmup_steps = 8000
+  hparams.learning_rate_constant = 1e-3
+  hparams.learning_rate_decay_rate = 0.5
+  hparams.learning_rate_decay_steps = 50000
+  hparams.dropout = 0.5
+  hparams.summarize_grads = True
+  hparams.summarize_vars = True
+
+  # not used hparams
+  hparams.label_smoothing = 0.
+  hparams.multiply_embedding_mode = "sqrt_depth"
+
+  # add new hparams
+  # use raw image as input
+  hparams.add_hparam("image_input_type", "image")
+  hparams.add_hparam("image_model_fn", "resnet_v1_152")
+  hparams.add_hparam("resize_side", 512)
+  hparams.add_hparam("height", 448)
+  hparams.add_hparam("width", 448)
+  hparams.add_hparam("distort", True)
+  hparams.add_hparam("train_resnet", False)
+
+  # image parts
+  hparams.add_hparam("image_feat_preprocess_proj", True)
+  hparams.add_hparam("image_feat_preprocess_layernorm", True)
+  hparams.add_hparam("image_feat_encode", True)
+  hparams.add_hparam("image_hidden_size", 0)  # default to hidden_size
+  hparams.add_hparam("image_filter_size", 0)  # defaults to filter_size
+
+  # question hidden size
+  hparams.hidden_size = 512
+  hparams.filter_size = 1024
+  hparams.num_hidden_layers = 4
+
+  hparams.add_hparam("multimodal_combine", "concat")
+  hparams.add_hparam("num_mlp_layers", 1)
+  hparams.add_hparam("mlp_size", 1024)
+
+  # self attention parts
+  hparams.norm_type = "layer"
+  hparams.layer_preprocess_sequence = "n"
+  hparams.layer_postprocess_sequence = "da"
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.attention_dropout = 0.1
+  hparams.relu_dropout = 0.1
+  hparams.add_hparam("pos", "timing")
+  hparams.add_hparam("num_encoder_layers", 0)
+  hparams.add_hparam("num_decoder_layers", 0)
+  hparams.add_hparam("num_heads", 8)
+  hparams.add_hparam("attention_key_channels", 0)
+  hparams.add_hparam("attention_value_channels", 0)
+  hparams.add_hparam("self_attention_type", "dot_product")
+  hparams.add_hparam("image_self_attention_type", "dot_product")
+  hparams.add_hparam("question_self_attention_type", "dot_product")
+  hparams.add_hparam("block_length", 1)
+  hparams.add_hparam("scale_dotproduct", True)
+
+  # iterative part
+  hparams.add_hparam("num_rec_steps", 3)
+
+  return hparams
+
+
+@registry.register_hparams
+def vqa_self_attention_feature():
+  hparams = vqa_self_attention_base()
+  hparams.image_input_type = "feature"
+  return hparams
+
+
+@registry.register_hparams
+def vqa_self_attention_feature_batch1024():
+  hparams = vqa_self_attention_feature()
+  hparams.batch_size = 1024
+  return hparams
+
+
+@registry.register_hparams
+def vqa_self_attention_feature_batch1024_big():
+  """Big model."""
+  hparams = vqa_self_attention_feature_batch1024()
+  hparams.learning_rate_constant = 7e-4
+  hparams.batch_size = 256
+  hparams.hidden_size = 1024
+  hparams.filter_size = 4096
+  hparams.num_heads = 16
+  hparams.layer_prepostprocess_dropout = 0.3
+  hparams.attention_dropout = 0.3
+  hparams.relu_dropout = 0.3
+  return hparams
+
+
+@registry.register_hparams
+def vqa_self_attention_feature_batch1024_exp():
+  hparams = vqa_self_attention_feature_batch1024()
+  hparams.learning_rate_schedule = (
+      "constant*linear_warmup*exp_decay")
+  hparams.learning_rate_decay_steps = 4000
+  return hparams
+
+
+@registry.register_hparams
+def vqa_self_attention_feature_batch1024_hidden6():
+  hparams = vqa_self_attention_feature_batch1024()
+  hparams.num_hidden_layers = 6
+  return hparams
+
+
+@registry.register_hparams
+def vqa_self_attention_feature_batch1024_hidden6_big():
+  hparams = vqa_self_attention_feature_batch1024_hidden6()
+  hparams.batch_size = 256
+  hparams.hidden_size = 1024
+  hparams.filter_size = 4096
+  hparams.num_heads = 16
+  hparams.layer_prepostprocess_dropout = 0.3
+  return hparams
+
+
+@registry.register_hparams
+def vqa_self_attention_feature_batch1024_drop03():
+  hparams = vqa_self_attention_feature_batch1024()
+  hparams.layer_prepostprocess_dropout = 0.3
+  return hparams
+
+
+@registry.register_hparams
+def vqa_self_attention_feature_lr5():
+  hparams = vqa_self_attention_feature()
+  hparams.learning_rate_constant = 5e-4
+  return hparams

From a4c893b74c44f9cb3d672ab30f674b03ed1c1b21 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 14 Aug 2018 11:42:07 -0700
Subject: [PATCH 0577/2720] Rm unused flag

PiperOrigin-RevId: 208684087
---
 tensor2tensor/utils/flags.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index 51d9d1a86..3ae12988a 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -73,7 +73,6 @@
                      "How many recent checkpoints to keep.")
 flags.DEFINE_bool("enable_graph_rewriter", False,
                   "Enable graph optimizations that are not on by default.")
-flags.DEFINE_bool("enable_summaries", True, "Enable creating summary ops.")
 flags.DEFINE_integer("keep_checkpoint_every_n_hours", 10000,
                      "Number of hours between each checkpoint to be saved. "
                      "The default value 10,000 hours effectively disables it.")

From 951c66e303971f0f34572d7df575ee40fc4b5797 Mon Sep 17 00:00:00 2001
From: Tomasz Latkowski <tlatkowski@gmail.com>
Date: Tue, 14 Aug 2018 20:47:09 +0200
Subject: [PATCH 0578/2720] fixed test bug and added mock requirement

---
 setup.py                                      |  1 +
 .../paraphrase_ms_coco_test.py                | 22 ++++++++++---------
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/setup.py b/setup.py
index f868adcfd..2096d4c21 100644
--- a/setup.py
+++ b/setup.py
@@ -40,6 +40,7 @@
         'gunicorn',
         'gym',
         'h5py',
+        'mock',
         'numpy',
         'oauth2client',
         'requests',
diff --git a/tensor2tensor/data_generators/paraphrase_ms_coco_test.py b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
index a20e01ecb..046e42dea 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
@@ -34,34 +34,36 @@ def testCombinationPairs(self):
     self.assertEqual(actual_combination, expected_combination)
 
   @mock.patch("tensor2tensor.data_generators"
-              ".paraphrase_mscoco.ParaphraseGenerationProblem.prepare_data",
+              ".paraphrase_ms_coco.ParaphraseGenerationProblem.prepare_data",
               return_value=[("sentence1", "sentence2")])
   @mock.patch("tensor2tensor.data_generators"
-              ".paraphrase_mscoco.ParaphraseGenerationProblem.bidirectional")
+              ".paraphrase_ms_coco.ParaphraseGenerationProblem.bidirectional")
   def testBidirectionalTrue(self, data, bidirectional):
     paraphrase_problem = paraphrase_ms_coco.ParaphraseGenerationProblem()
     paraphrase_problem.bidirectional = True
 
     expected_generated_data = [{"inputs": "sentence1", "targets": "sentence2"},
                                {"inputs": "sentence2", "targets": "sentence1"}]
-    actual_generated_data = list(paraphrase_problem.generate_samples("data_dir",
-                                                                     "tmp_dir",
-                                                                     "dataset_split"))
+    actual_generated_data = list(paraphrase_problem
+                                 .generate_samples("data_dir",
+                                                   "tmp_dir",
+                                                   "dataset_split"))
     self.assertEqual(actual_generated_data, expected_generated_data)
 
   @mock.patch("tensor2tensor.data_generators"
-              ".paraphrase_mscoco.ParaphraseGenerationProblem.prepare_data",
+              ".paraphrase_ms_coco.ParaphraseGenerationProblem.prepare_data",
               return_value=[("sentence1", "sentence2")])
   @mock.patch("tensor2tensor.data_generators"
-              ".paraphrase_mscoco.ParaphraseGenerationProblem.bidirectional")
+              ".paraphrase_ms_coco.ParaphraseGenerationProblem.bidirectional")
   def testBidirectionalFalse(self, data, bidirectional):
     paraphrase_problem = paraphrase_ms_coco.ParaphraseGenerationProblem()
     paraphrase_problem.bidirectional = False
 
     expected_generated_data = [{"inputs": "sentence1", "targets": "sentence2"}]
-    actual_generated_data = list(paraphrase_problem.generate_samples("data_dir",
-                                                                     "tmp_dir",
-                                                                     "dataset_split"))
+    actual_generated_data = list(paraphrase_problem
+                                 .generate_samples("data_dir",
+                                                   "tmp_dir",
+                                                   "dataset_split"))
     self.assertEqual(actual_generated_data, expected_generated_data)
 
 

From 7e16303c859db7d68db51d319125e4f6e264213f Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Tue, 14 Aug 2018 12:30:42 -0700
Subject: [PATCH 0579/2720] Created file
 mesh_tensorflow/research/expereiments_moe.py for a sequence of

PiperOrigin-RevId: 208692274
---
 .../research/experiments_moe.py               | 103 ++++++++++++++++++
 tensor2tensor/models/__init__.py              |   1 +
 2 files changed, 104 insertions(+)
 create mode 100644 tensor2tensor/mesh_tensorflow/research/experiments_moe.py

diff --git a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
new file mode 100644
index 000000000..3ea2b4db6
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
@@ -0,0 +1,103 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Experiments with mixture-of-experts architectures.
+
+For all of these architectures, we run on languagemodel_lm1b8k_packed
+for 32k-96 steps (1-3 epochs) on one TPU (8 cores).
+
+All log-perplexities are per-token - multiply by 1.298 for per-word
+
+Results:
+model      params(M)  einsum  alltoall  mxu-util  log-ppl(1ep) (3ep)
+dense_4k   30         3.0e12  0         45%        3.31
+dense_8k   46         4.7e12  0         49%        3.24
+dense_64k                     0                    3.06
+v0         282        4.9e12  5.4e8     35%        3.06
+v0_o75     282        4.0e12  4.0e8     34%
+k_means    282        4.0e12  4.0e8                3.12
+k_means_o2 282        4.9e12  5.4e8     33%
+
+Note: configurations and code are likely to change without notice.
+"""
+
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.mesh_tensorflow import mtf_transformer
+from tensor2tensor.utils import registry
+
+
+@registry.register_hparams
+def xmoe_dense_4k():
+  """Small transformer language model."""
+  hparams = mtf_transformer.mtf_transformer_base()
+
+  # The following hparams are constant across all these experiments.
+  hparams.label_smoothing = 0.0
+  hparams.batch_size = 128
+  hparams.d_model = 512
+  hparams.d_kv = 128
+  hparams.num_heads = 4
+  hparams.num_decoder_layers = 4
+  hparams.shared_embedding_and_softmax_weights = False
+  hparams.learning_rate_schedule = "rsqrt_decay"
+
+  # We will vary the following parameters related to the ffn/moe layers.
+  hparams.feedforward_layer = "dense_relu_dense"
+  hparams.d_ff = 4096
+  hparams.moe_num_experts = 16
+  hparams.moe_overhead_train = 1.0
+  hparams.moe_overhead_eval = 2.0
+  hparams.moe_loss_coef = 1e-3
+  hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
+  hparams.mesh_shape = "batch:8"
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_dense_8k():
+  hparams = xmoe_dense_4k()
+  hparams.d_ff = 8192
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_dense_64k():
+  hparams = xmoe_dense_4k()
+  hparams.d_ff = 65536
+  hparams.mesh_shape = "model:4,batch:8"
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_v0():
+  """Mixture of experts."""
+  hparams = xmoe_dense_4k()
+  hparams.feedforward_layer = "moe"
+  hparams.mesh_shape = "all:8"
+  hparams.layout = "batch:all;experts:all"
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_v0_o75():
+  """Mixture of experts."""
+  hparams = xmoe_v0()
+  hparams.moe_overhead_train = 0.75
+  return hparams
+
+
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 84f2404b9..b6af0ed00 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -21,6 +21,7 @@
 from tensor2tensor.layers import modalities  # pylint: disable=g-import-not-at-top
 from tensor2tensor.mesh_tensorflow import mtf_image_transformer
 from tensor2tensor.mesh_tensorflow import mtf_transformer
+from tensor2tensor.mesh_tensorflow.research import experiments_moe
 from tensor2tensor.models import basic
 from tensor2tensor.models import bytenet
 from tensor2tensor.models import distillation

From a8e5975dace0cdc16b39214a5201635fcb0b3f2c Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Tue, 14 Aug 2018 12:59:03 -0700
Subject: [PATCH 0580/2720] Add targeted dropout to ResNet and update pruning

PiperOrigin-RevId: 208696866
---
 tensor2tensor/layers/common_layers.py | 117 +++++++++--
 tensor2tensor/models/resnet.py        | 266 ++++++++++++++++++++++----
 tensor2tensor/utils/pruning_utils.py  |  26 +--
 3 files changed, 347 insertions(+), 62 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index b8973c950..52dae54f5 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -738,7 +738,8 @@ def apply_spectral_norm(x):
   # v = Wu / ||W u||
   with tf.variable_scope("u", reuse=tf.AUTO_REUSE):
     u = tf.get_variable(
-        "u", [num_filters, 1], initializer=tf.truncated_normal_initializer(),
+        "u", [num_filters, 1],
+        initializer=tf.truncated_normal_initializer(),
         trainable=False)
   v = tf.nn.l2_normalize(tf.matmul(weights_2d, u))
 
@@ -746,9 +747,8 @@ def apply_spectral_norm(x):
   u_new = tf.nn.l2_normalize(tf.matmul(tf.transpose(v), weights_2d))
 
   # s = v*W*u
-  spectral_norm = tf.squeeze(tf.matmul(
-      tf.transpose(v),
-      tf.matmul(weights_2d, tf.transpose(u_new))))
+  spectral_norm = tf.squeeze(
+      tf.matmul(tf.transpose(v), tf.matmul(weights_2d, tf.transpose(u_new))))
 
   # set u equal to u_new in the next iteration.
   assign_op = tf.assign(u, tf.transpose(u_new))
@@ -3701,6 +3701,103 @@ def cyclegan_upsample(net, num_outputs, stride, method="conv2d_transpose"):
     return net
 
 
+def weight_targeting(w, k):
+  """Weight-level magnitude pruning."""
+  k = tf.to_int32(k)
+  w_shape = shape_list(w)
+  size = tf.to_int32(tf.reduce_prod(w_shape[:-1]))
+  w = tf.reshape(w, [size, w_shape[-1]])
+
+  transpose_w = tf.transpose(w)
+  thres = tf.contrib.framework.sort(tf.abs(transpose_w), axis=1)[:, k]
+  mask = tf.to_float(thres[None, :] >= tf.abs(w))
+
+  return tf.reshape(mask, w_shape)
+
+
+def unit_targeting(w, k):
+  """Unit-level magnitude pruning."""
+  k = tf.to_int32(k)
+  w_shape = shape_list(w)
+  size = tf.to_int32(tf.reduce_prod(w_shape[:-1]))
+  w = tf.reshape(w, [size, w_shape[-1]])
+
+  norm = tf.norm(w, axis=0)
+  thres = tf.contrib.framework.sort(norm, axis=0)[k]
+  mask = tf.to_float(thres >= norm)[None, :]
+  mask = tf.tile(mask, [size, 1])
+
+  return tf.reshape(mask, w_shape)
+
+
+def td_conv(inputs,
+            filters,
+            kernel_size,
+            targeting_count,
+            targeting_fn,
+            keep_prob,
+            is_training,
+            do_prune=True,
+            strides=(1, 1),
+            padding="valid",
+            data_format="channels_last",
+            dilation_rate=(1, 1),
+            activation=None,
+            use_bias=True,
+            kernel_initializer=None,
+            bias_initializer=tf.zeros_initializer(),
+            name=None,
+            reuse=None):
+  """Apply targeted dropout to the weights of a convolution."""
+  with tf.variable_scope(name, default_name="td_conv", reuse=reuse):
+    nhwc = data_format == "channels_last"
+    in_dim = shape_list(inputs)[-1] if nhwc else shape_list(inputs)[1]
+
+    kernel_shape = [kernel_size, kernel_size, in_dim, filters]
+    w = tf.get_variable(
+        "DW", shape=kernel_shape, initializer=kernel_initializer)
+    if use_bias:
+      b = tf.get_variable("b", shape=[filters], initializer=bias_initializer)
+
+    if keep_prob < 1.0:
+      w = targeted_dropout(
+          w,
+          targeting_count,
+          keep_prob,
+          targeting_fn,
+          is_training,
+          do_prune=do_prune)
+
+    if isinstance(strides, int):
+      strides = [strides, strides]
+    if isinstance(dilation_rate, int):
+      dilation_rate = [dilation_rate, dilation_rate]
+
+    if nhwc:
+      strides = [1, strides[0], strides[1], 1]
+      dilation_rate = [1, dilation_rate[0], dilation_rate[1], 1]
+    else:
+      strides = [1, 1, strides[0], strides[1]]
+      dilation_rate = [1, 1, dilation_rate[0], dilation_rate[1]]
+
+    y = tf.nn.conv2d(
+        inputs,
+        w,
+        strides,
+        padding,
+        data_format="NHWC" if nhwc else "NCHW",
+        dilations=dilation_rate,
+        name=None)
+
+    if use_bias:
+      y += b
+
+    if activation:
+      y = activation(y)
+
+    return y
+
+
 def targeted_dropout(inputs,
                      k,
                      keep_prob,
@@ -3709,8 +3806,8 @@ def targeted_dropout(inputs,
                      do_prune=False):
   """Applies targeted dropout.
 
-  Applies dropout at a rate of `1 - keep_prob` to only those elements of `x`
-  marked by `targeting_fn`. See below and paper for more detail:
+  Applies dropout at a rate of `1 - keep_prob` to only those elements of
+  `inputs` marked by `targeting_fn`. See below and paper for more detail:
 
   "Targeted Dropout for Posthoc Pruning" Aidan N. Gomez, Ivan Zhang,
     Kevin Swersky, Yarin Gal, and Geoffrey E. Hinton.
@@ -3727,11 +3824,12 @@ def targeted_dropout(inputs,
     is_training: bool, indicates whether currently training.
     do_prune: bool, indicates whether to prune the `k * (1 - keep_prob)`
       elements of `inputs` expected to be dropped each forwards pass.
+
   Returns:
     Tensor, same shape and dtype as `inputs`.
   """
   if not is_training and do_prune:
-    k = tf.round(k * (1 - keep_prob))
+    k = tf.round(tf.to_float(k) * tf.to_float(1. - keep_prob))
 
   mask = targeting_fn(inputs, k)
   mask = tf.cast(mask, inputs.dtype)
@@ -3793,10 +3891,7 @@ def sparse_eye(size):
   dense_shape = [tf.cast(size, tf.int64), tf.cast(size, tf.int64)]
 
   return tf.SparseTensor(
-      indices=indices,
-      values=values,
-      dense_shape=dense_shape
-  )
+      indices=indices, values=values, dense_shape=dense_shape)
 
 
 # modification from https://github.com/tensorflow/tensorflow/pull/21276
diff --git a/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index dd185a669..4b2ae4cb7 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -107,7 +107,11 @@ def conv2d_fixed_padding(inputs,
                          filters,
                          kernel_size,
                          strides,
-                         data_format="channels_first"):
+                         data_format="channels_first",
+                         use_td=False,
+                         targeting_rate=None,
+                         keep_prob=None,
+                         is_training=None):
   """Strided 2-D convolution with explicit padding.
 
   The padding is consistent and is based only on `kernel_size`, not on the
@@ -120,22 +124,63 @@ def conv2d_fixed_padding(inputs,
     strides: `int` strides of the convolution.
     data_format: `str` either "channels_first" for `[batch, channels, height,
         width]` or "channels_last for `[batch, height, width, channels]`.
+    use_td: `str` one of "weight" or "unit". Set to False or "" to disable
+      targeted dropout.
+    targeting_rate: `float` proportion of weights to target with targeted
+      dropout.
+    keep_prob: `float` keep probability for targeted dropout.
+    is_training: `bool` for whether the model is in training.
 
   Returns:
     A `Tensor` of shape `[batch, filters, height_out, width_out]`.
+
+  Raises:
+    Exception: if use_td is not valid.
   """
   if strides > 1:
     inputs = fixed_padding(inputs, kernel_size, data_format=data_format)
 
-  return tf.layers.conv2d(
-      inputs=inputs,
-      filters=filters,
-      kernel_size=kernel_size,
-      strides=strides,
-      padding=("SAME" if strides == 1 else "VALID"),
-      use_bias=False,
-      kernel_initializer=tf.variance_scaling_initializer(),
-      data_format=data_format)
+  if use_td:
+    inputs_shape = common_layers.shape_list(inputs)
+    if use_td == "weight":
+      if data_format == "channels_last":
+        size = kernel_size * kernel_size * inputs_shape[-1]
+      else:
+        size = kernel_size * kernel_size * inputs_shape[1]
+      targeting_count = targeting_rate * tf.to_float(size)
+      targeting_fn = common_layers.weight_targeting
+    elif use_td == "unit":
+      targeting_count = targeting_rate * filters
+      targeting_fn = common_layers.unit_targeting
+    else:
+      raise Exception("Unrecognized targeted dropout type: %s" % use_td)
+
+    y = common_layers.td_conv(
+        inputs,
+        filters,
+        kernel_size,
+        targeting_count,
+        targeting_fn,
+        keep_prob,
+        is_training,
+        do_prune=True,
+        strides=strides,
+        padding=("SAME" if strides == 1 else "VALID"),
+        data_format=data_format,
+        use_bias=False,
+        kernel_initializer=tf.variance_scaling_initializer())
+  else:
+    y = tf.layers.conv2d(
+        inputs=inputs,
+        filters=filters,
+        kernel_size=kernel_size,
+        strides=strides,
+        padding=("SAME" if strides == 1 else "VALID"),
+        use_bias=False,
+        kernel_initializer=tf.variance_scaling_initializer(),
+        data_format=data_format)
+
+  return y
 
 
 def residual_block(inputs,
@@ -144,7 +189,10 @@ def residual_block(inputs,
                    projection_shortcut,
                    strides,
                    final_block,
-                   data_format="channels_first"):
+                   data_format="channels_first",
+                   use_td=False,
+                   targeting_rate=None,
+                   keep_prob=None):
   """Standard building block for residual networks with BN before convolutions.
 
   Args:
@@ -162,6 +210,11 @@ def residual_block(inputs,
         `bottleneck_block`.
     data_format: `str` either "channels_first" for `[batch, channels, height,
         width]` or "channels_last for `[batch, height, width, channels]`.
+    use_td: `str` one of "weight" or "unit". Set to False or "" to disable
+      targeted dropout.
+    targeting_rate: `float` proportion of weights to target with targeted
+      dropout.
+    keep_prob: `float` keep probability for targeted dropout.
 
   Returns:
     The output `Tensor` of the block.
@@ -178,7 +231,11 @@ def residual_block(inputs,
       filters=filters,
       kernel_size=3,
       strides=strides,
-      data_format=data_format)
+      data_format=data_format,
+      use_td=use_td,
+      targeting_rate=targeting_rate,
+      keep_prob=keep_prob,
+      is_training=is_training)
 
   inputs = batch_norm_relu(inputs, is_training, data_format=data_format)
   inputs = conv2d_fixed_padding(
@@ -186,7 +243,11 @@ def residual_block(inputs,
       filters=filters,
       kernel_size=3,
       strides=1,
-      data_format=data_format)
+      data_format=data_format,
+      use_td=use_td,
+      targeting_rate=targeting_rate,
+      keep_prob=keep_prob,
+      is_training=is_training)
 
   return inputs + shortcut
 
@@ -197,7 +258,10 @@ def bottleneck_block(inputs,
                      projection_shortcut,
                      strides,
                      final_block,
-                     data_format="channels_first"):
+                     data_format="channels_first",
+                     use_td=False,
+                     targeting_rate=None,
+                     keep_prob=None):
   """Bottleneck block variant for residual networks with BN after convolutions.
 
   Args:
@@ -216,6 +280,11 @@ def bottleneck_block(inputs,
         the final batch norm in a block.
     data_format: `str` either "channels_first" for `[batch, channels, height,
         width]` or "channels_last for `[batch, height, width, channels]`.
+    use_td: `str` one of "weight" or "unit". Set to False or "" to disable
+      targeted dropout.
+    targeting_rate: `float` proportion of weights to target with targeted
+      dropout.
+    keep_prob: `float` keep probability for targeted dropout.
 
   Returns:
     The output `Tensor` of the block.
@@ -232,7 +301,11 @@ def bottleneck_block(inputs,
       filters=filters,
       kernel_size=1,
       strides=1,
-      data_format=data_format)
+      data_format=data_format,
+      use_td=use_td,
+      targeting_rate=targeting_rate,
+      keep_prob=keep_prob,
+      is_training=is_training)
 
   inputs = batch_norm_relu(inputs, is_training, data_format=data_format)
   inputs = conv2d_fixed_padding(
@@ -240,7 +313,11 @@ def bottleneck_block(inputs,
       filters=filters,
       kernel_size=3,
       strides=strides,
-      data_format=data_format)
+      data_format=data_format,
+      use_td=use_td,
+      targeting_rate=targeting_rate,
+      keep_prob=keep_prob,
+      is_training=is_training)
 
   inputs = batch_norm_relu(inputs, is_training, data_format=data_format)
   inputs = conv2d_fixed_padding(
@@ -248,7 +325,11 @@ def bottleneck_block(inputs,
       filters=4 * filters,
       kernel_size=1,
       strides=1,
-      data_format=data_format)
+      data_format=data_format,
+      use_td=use_td,
+      targeting_rate=targeting_rate,
+      keep_prob=keep_prob,
+      is_training=is_training)
   inputs = batch_norm_relu(
       inputs,
       is_training,
@@ -266,7 +347,10 @@ def block_layer(inputs,
                 strides,
                 is_training,
                 name,
-                data_format="channels_first"):
+                data_format="channels_first",
+                use_td=False,
+                targeting_rate=None,
+                keep_prob=None):
   """Creates one layer of blocks for the ResNet model.
 
   Args:
@@ -280,6 +364,11 @@ def block_layer(inputs,
     name: `str`name for the Tensor output of the block layer.
     data_format: `str` either "channels_first" for `[batch, channels, height,
         width]` or "channels_last for `[batch, height, width, channels]`.
+    use_td: `str` one of "weight" or "unit". Set to False or "" to disable
+      targeted dropout.
+    targeting_rate: `float` proportion of weights to target with targeted
+      dropout.
+    keep_prob: `float` keep probability for targeted dropout.
 
   Returns:
     The output `Tensor` of the block layer.
@@ -288,22 +377,44 @@ def block_layer(inputs,
   filters_out = 4 * filters if block_fn is bottleneck_block else filters
 
   def projection_shortcut(inputs):
+    """Project identity branch."""
     inputs = conv2d_fixed_padding(
         inputs=inputs,
         filters=filters_out,
         kernel_size=1,
         strides=strides,
-        data_format=data_format)
+        data_format=data_format,
+        use_td=use_td,
+        targeting_rate=targeting_rate,
+        keep_prob=keep_prob,
+        is_training=is_training)
     return batch_norm_relu(
         inputs, is_training, relu=False, data_format=data_format)
 
   # Only the first block per block_layer uses projection_shortcut and strides
-  inputs = block_fn(inputs, filters, is_training, projection_shortcut, strides,
-                    False, data_format)
+  inputs = block_fn(
+      inputs,
+      filters,
+      is_training,
+      projection_shortcut,
+      strides,
+      False,
+      data_format,
+      use_td=use_td,
+      targeting_rate=targeting_rate,
+      keep_prob=keep_prob)
 
   for i in range(1, blocks):
-    inputs = block_fn(inputs, filters, is_training, None, 1, (i + 1 == blocks),
-                      data_format)
+    inputs = block_fn(
+        inputs,
+        filters,
+        is_training,
+        None,
+        1, (i + 1 == blocks),
+        data_format,
+        use_td=use_td,
+        targeting_rate=targeting_rate,
+        keep_prob=keep_prob)
 
   return tf.identity(inputs, name)
 
@@ -314,7 +425,10 @@ def resnet_v2(inputs,
               filters,
               data_format="channels_first",
               is_training=False,
-              is_cifar=False):
+              is_cifar=False,
+              use_td=False,
+              targeting_rate=None,
+              keep_prob=None):
   """Resnet model.
 
   Args:
@@ -330,6 +444,11 @@ def resnet_v2(inputs,
         width]` or "channels_last" `[batch, height, width, channels]`.
     is_training: bool, build in training mode or not.
     is_cifar: bool, whether the data is CIFAR or not.
+    use_td: `str` one of "weight" or "unit". Set to False or "" to disable
+      targeted dropout.
+    targeting_rate: `float` proportion of weights to target with targeted
+      dropout.
+    keep_prob: `float` keep probability for targeted dropout.
 
   Returns:
     Pre-logit activations.
@@ -342,7 +461,10 @@ def resnet_v2(inputs,
       strides=1,
       is_training=is_training,
       name="block_layer1",
-      data_format=data_format)
+      data_format=data_format,
+      use_td=use_td,
+      targeting_rate=targeting_rate,
+      keep_prob=keep_prob)
   inputs = block_layer(
       inputs=inputs,
       filters=filters[2],
@@ -351,7 +473,10 @@ def resnet_v2(inputs,
       strides=2,
       is_training=is_training,
       name="block_layer2",
-      data_format=data_format)
+      data_format=data_format,
+      use_td=use_td,
+      targeting_rate=targeting_rate,
+      keep_prob=keep_prob)
   inputs = block_layer(
       inputs=inputs,
       filters=filters[3],
@@ -360,7 +485,10 @@ def resnet_v2(inputs,
       strides=2,
       is_training=is_training,
       name="block_layer3",
-      data_format=data_format)
+      data_format=data_format,
+      use_td=use_td,
+      targeting_rate=targeting_rate,
+      keep_prob=keep_prob)
   if not is_cifar:
     inputs = block_layer(
         inputs=inputs,
@@ -370,7 +498,10 @@ def resnet_v2(inputs,
         strides=2,
         is_training=is_training,
         name="block_layer4",
-        data_format=data_format)
+        data_format=data_format,
+        use_td=use_td,
+        targeting_rate=targeting_rate,
+        keep_prob=keep_prob)
 
   return inputs
 
@@ -424,14 +555,20 @@ def body(self, features):
         hp.filter_sizes,
         data_format,
         is_training=is_training,
-        is_cifar=hp.is_cifar)
+        is_cifar=hp.is_cifar,
+        use_td=hp.use_td,
+        targeting_rate=hp.targeting_rate,
+        keep_prob=hp.keep_prob)
 
     if hp.use_nchw:
       out = tf.transpose(out, [0, 2, 3, 1])
 
+    if not hp.is_cifar:
+      return out
+
     out = tf.reduce_mean(out, [1, 2])
     num_classes = self._problem_hparams.target_modality.top_dimensionality
-    logits = tf.layers.dense(out, num_classes)
+    logits = tf.layers.dense(out, num_classes, name="logits")
 
     losses = {"training": 0.0}
     if is_training:
@@ -482,6 +619,11 @@ def resnet_base():
   hparams.add_hparam("use_nchw", True)
   hparams.add_hparam("is_cifar", False)
 
+  # Targeted dropout
+  hparams.add_hparam("use_td", False)
+  hparams.add_hparam("targeting_rate", None)
+  hparams.add_hparam("keep_prob", None)
+
   # Variable init
   hparams.initializer = "normal_unit_scaling"
   hparams.initializer_gain = 2.
@@ -527,6 +669,39 @@ def resnet_imagenet_34():
   return hp
 
 
+@registry.register_hparams
+def resnet_imagenet_34_td_weight_05_05():
+  """Set of hyperparameters."""
+  hp = resnet_imagenet_34()
+  hp.use_td = "weight"
+  hp.targeting_rate = 0.5
+  hp.keep_prob = 0.5
+
+  return hp
+
+
+@registry.register_hparams
+def resnet_imagenet_34_td_unit_05_05():
+  """Set of hyperparameters."""
+  hp = resnet_imagenet_34()
+  hp.use_td = "unit"
+  hp.targeting_rate = 0.5
+  hp.keep_prob = 0.5
+
+  return hp
+
+
+@registry.register_hparams
+def resnet_imagenet_34_td_unit_no_drop():
+  """Set of hyperparameters."""
+  hp = resnet_imagenet_34()
+  hp.use_td = "unit"
+  hp.targeting_rate = 0.0
+  hp.keep_prob = 1.0
+
+  return hp
+
+
 @registry.register_hparams
 def resnet_imagenet_102():
   hp = resnet_imagenet_34()
@@ -553,6 +728,33 @@ def resnet_cifar_32():
   return hp
 
 
+@registry.register_hparams
+def resnet_cifar_32_td_weight_05_05():
+  hp = resnet_cifar_32()
+  hp.use_td = "weight"
+  hp.targeting_rate = 0.5
+  hp.keep_prob = 0.5
+  return hp
+
+
+@registry.register_hparams
+def resnet_cifar_32_td_unit_05_05():
+  hp = resnet_cifar_32()
+  hp.use_td = "unit"
+  hp.targeting_rate = 0.5
+  hp.keep_prob = 0.5
+  return hp
+
+
+@registry.register_hparams
+def resnet_cifar_32_td_unit_no_drop():
+  hp = resnet_cifar_32()
+  hp.use_td = "unit"
+  hp.targeting_rate = 0.0
+  hp.keep_prob = 1.0
+  return hp
+
+
 @registry.register_hparams
 def resnet_34():
   hp = resnet_base()
@@ -587,8 +789,8 @@ def resnet_weight():
   hp = tf.contrib.training.HParams()
   hp.add_hparam("strategy", "weight")
   hp.add_hparam("black_list", ["logits", "bias"])
-  hp.add_hparam("white_list", None)
-  hp.add_hparam("sparsities", [0.1*i for i in range(10)])
+  hp.add_hparam("white_list", ["td_conv"])
+  hp.add_hparam("sparsities", [0.1 * i for i in range(10)])
   return hp
 
 
diff --git a/tensor2tensor/utils/pruning_utils.py b/tensor2tensor/utils/pruning_utils.py
index 70cb3621a..446e1c59c 100644
--- a/tensor2tensor/utils/pruning_utils.py
+++ b/tensor2tensor/utils/pruning_utils.py
@@ -27,31 +27,18 @@ def weight(w, sparsity):
   """Weight-level magnitude pruning."""
   w_shape = common_layers.shape_list(w)
   k = int(np.prod(w_shape[:-1]))
-  w = tf.reshape(w, [k, w_shape[-1]])
-
-  idx = int(k * sparsity)
-  thres = tf.contrib.framework.sort(tf.abs(w), axis=0)[idx]
-  mask = tf.to_float(thres[None, :] < tf.abs(w))
-
-  w = mask * w
-  return tf.reshape(w, w_shape)
+  count = tf.to_int32(k * sparsity)
+  mask = common_layers.weight_targeting(w, count)
+  return (1 - mask) * w
 
 
 @registry.register_pruning_strategy
 def unit(w, sparsity):
   """Unit-level magnitude pruning."""
   w_shape = common_layers.shape_list(w)
-  k = int(np.prod(w_shape[:-1]))
-  w = tf.reshape(w, [k, w_shape[-1]])
-  idx = int(w_shape[-1] * sparsity)
-
-  norm = tf.norm(w, axis=0)
-  thres = tf.contrib.framework.sort(norm, axis=0)[idx]
-  mask = tf.to_float(thres < norm)[None, :]
-  mask = tf.tile(mask, [k, 1])
-
-  w = mask * w
-  return tf.reshape(w, w_shape)
+  count = tf.to_int32(w_shape[-1] * sparsity)
+  mask = common_layers.unit_targeting(w, count)
+  return (1 - mask) * w
 
 
 def sparsify(sess, eval_model, pruning_strategy, pruning_params):
@@ -72,6 +59,7 @@ def should_prune(name):
     return True
 
   weights = [w for w in weights if should_prune(w.name)]
+  tf.logging.info("Pruning weights: %s" % weights)
   unpruned_weights = sess.run(weights)
 
   reset_op = tf.no_op()

From 53e51bfdd7f0bf0212831ec78d15131f27f0ffae Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 14 Aug 2018 14:06:41 -0700
Subject: [PATCH 0581/2720] Stochastic Adversarial Video Prediction (Part 3 of
 3). Adds an adversarial loss component using a discriminator to the videos
 generated by the conditional latents, via hparams
 "use_vae=True,use_gan=True".

PiperOrigin-RevId: 208709463
---
 .../models/research/next_frame_params.py      |  1 +
 .../models/research/next_frame_savp.py        | 43 +++++++++++++------
 .../models/research/next_frame_test.py        | 29 +++++++------
 3 files changed, 46 insertions(+), 27 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index d9cb7efe4..a0d4294ed 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -122,6 +122,7 @@ def next_frame_savp():
   hparams.add_hparam("use_spectral_norm", True)
   hparams.add_hparam("gan_loss", "cross_entropy")
   hparams.add_hparam("gan_loss_multiplier", 0.01)
+  hparams.add_hparam("gan_vae_loss_multiplier", 0.01)
   hparams.add_hparam("gan_optimization", "joint")
   hparams.target_modality = "video:l1raw"
   hparams.input_modalities = "inputs:video:l1raw"
diff --git a/tensor2tensor/models/research/next_frame_savp.py b/tensor2tensor/models/research/next_frame_savp.py
index 7f18afcf4..49775f646 100644
--- a/tensor2tensor/models/research/next_frame_savp.py
+++ b/tensor2tensor/models/research/next_frame_savp.py
@@ -224,7 +224,7 @@ def g_step(self, gen_frames, fake_logits_stop):
         discriminator_gen_outputs=fake_logits_stop, add_summaries=True)
     return gan_g_loss_pos_d, gan_g_loss_neg_d
 
-  def get_gan_loss(self, true_frames, gen_frames):
+  def get_gan_loss(self, true_frames, gen_frames, name):
     """Get the discriminator + generator loss at every step.
 
     This performs an 1:1 update of the discriminator and generator at every
@@ -235,20 +235,21 @@ def get_gan_loss(self, true_frames, gen_frames):
                    Assumed to be ground truth.
       gen_frames: 5-D Tensor of shape (num_steps, batch_size, H, W, C)
                   Assumed to be fake.
+      name: discriminator scope.
     Returns:
       loss: 0-D Tensor, with d_loss + g_loss
     """
     # D - STEP
-    with tf.variable_scope("gan_discriminator", reuse=tf.AUTO_REUSE):
+    with tf.variable_scope("%s_discriminator" % name, reuse=tf.AUTO_REUSE):
       gan_d_loss, _, fake_logits_stop = self.d_step(
           true_frames, gen_frames)
 
     # G - STEP
-    with tf.variable_scope("gan_discriminator", reuse=True):
+    with tf.variable_scope("%s_discriminator" % name, reuse=True):
       gan_g_loss_pos_d, gan_g_loss_neg_d = self.g_step(
           gen_frames, fake_logits_stop)
     gan_g_loss = gan_g_loss_pos_d + gan_g_loss_neg_d
-    tf.summary.scalar("gan_loss", gan_g_loss_pos_d + gan_d_loss)
+    tf.summary.scalar("gan_loss_%s" % name, gan_g_loss_pos_d + gan_d_loss)
 
     if self.hparams.gan_optimization == "joint":
       gan_loss = gan_g_loss + gan_d_loss
@@ -257,24 +258,40 @@ def get_gan_loss(self, true_frames, gen_frames):
       gan_loss = tf.cond(
           tf.logical_not(curr_step % 2 == 0), lambda: gan_g_loss,
           lambda: gan_d_loss)
-    return self.hparams.gan_loss_multiplier * gan_loss
+    return gan_loss
 
   def get_extra_loss(self, latent_means=None, latent_stds=None,
                      true_frames=None, gen_frames=None, beta=1.0):
     if not self.is_training:
       return 0.0
+
+    vae_loss, d_vae_loss, d_gan_loss = 0.0, 0.0, 0.0
+    # Use next_frame_sv2p's KL divergence computation.
     if self.hparams.use_vae:
-      return super(NextFrameSAVP, self).get_extra_loss(
+      vae_loss = super(NextFrameSAVP, self).get_extra_loss(
           latent_means=latent_means, latent_stds=latent_stds, beta=beta)
-    elif self.hparams.use_gan:
+
+    if self.hparams.use_gan:
       # Strip out the first context_frames for the true_frames
       # Strip out the first context_frames - 1 for the gen_frames
       context_frames = self.hparams.video_num_input_frames
       true_frames = tf.stack(
           tf.unstack(true_frames, axis=0)[context_frames:])
-      gen_frames = tf.stack(
-          tf.unstack(gen_frames, axis=0)[context_frames-1:])
-      return self.get_gan_loss(true_frames, gen_frames)
+
+      # discriminator for VAE.
+      if self.hparams.use_vae:
+        gen_enc_frames = tf.stack(
+            tf.unstack(gen_frames, axis=0)[context_frames-1:])
+        d_vae_loss = self.get_gan_loss(true_frames, gen_enc_frames, name="vae")
+
+      # discriminator for GAN.
+      gen_prior_frames = tf.stack(
+          tf.unstack(self.gen_prior_video, axis=0)[context_frames-1:])
+      d_gan_loss = self.get_gan_loss(true_frames, gen_prior_frames, name="gan")
+
+    return (
+        vae_loss + self.hparams.gan_loss_multiplier * d_gan_loss +
+        self.hparams.gan_vae_loss_multiplier * d_vae_loss)
 
   def pad_conv3d_lrelu(self, activations, n_filters, kernel_size, strides,
                        scope):
@@ -334,8 +351,6 @@ def construct_model(self, images, actions, rewards):
       ValueError: If not exactly one of self.hparams.vae or self.hparams.gan
                   is set to True.
     """
-    if self.hparams.use_vae and self.hparams.use_gan:
-      raise ValueError("VAE + GAN variant not implemented")
     if not self.hparams.use_vae and not self.hparams.use_gan:
       raise ValueError("Set at least one of use_vae or use_gan to be True")
     if self.hparams.gan_optimization not in ["joint", "sequential"]:
@@ -426,10 +441,10 @@ def construct_model(self, images, actions, rewards):
         latent_stds.append(log_sigma_sq)
 
     gen_cond_video = tf.stack(gen_cond_video, axis=0)
-    gen_prior_video = tf.stack(gen_prior_video, axis=0)
+    self.gen_prior_video = tf.stack(gen_prior_video, axis=0)
     fake_rewards = tf.stack(fake_rewards, axis=0)
 
     if train_mode and self.hparams.use_vae:
       return gen_cond_video, fake_rewards, latent_means, latent_stds
     else:
-      return gen_prior_video, fake_rewards, latent_means, latent_stds
+      return self.gen_prior_video, fake_rewards, latent_means, latent_stds
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index f289758c8..b797d2b49 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -221,15 +221,14 @@ def testStochasticEmily(self):
         next_frame_emily.NextFrameStochasticEmily,
         1)
 
-  def testStochasticSavp(self):
+  def testStochasticSavpVAE(self):
+    savp_hparams = next_frame_params.next_frame_savp()
+    savp_hparams.use_vae = True
+    savp_hparams.use_gan = False
     self.TestOnVariousInputOutputSizes(
-        next_frame_params.next_frame_savp(),
-        next_frame_savp.NextFrameSAVP,
-        1)
+        savp_hparams, next_frame_savp.NextFrameSAVP, 1)
     self.TestOnVariousUpSampleLayers(
-        next_frame_params.next_frame_savp(),
-        next_frame_savp.NextFrameSAVP,
-        1)
+        savp_hparams, next_frame_savp.NextFrameSAVP, 1)
 
   def testStochasticSavpGAN(self):
     hparams = next_frame_params.next_frame_savp()
@@ -240,14 +239,18 @@ def testStochasticSavpGAN(self):
     hparams.gan_optimization = "sequential"
     self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
 
-  def testStochasticInvalidVAEGANCombinations(self):
+  def testStochasticSavpGANVAE(self):
     hparams = next_frame_params.next_frame_savp()
-    for use_vae, use_gan in [[True, True], [False, False]]:
-      hparams.use_gan = use_gan
-      hparams.use_vae = use_vae
-      self.assertRaises(ValueError, self.TestVideoModel,
-                        7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
+    hparams.use_vae = True
+    hparams.use_gan = True
+    self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
 
+  def testStochasticInvalidVAEGANCombinations(self):
+    hparams = next_frame_params.next_frame_savp()
+    hparams.use_gan = False
+    hparams.use_vae = False
+    self.assertRaises(ValueError, self.TestVideoModel,
+                      7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
 
 if __name__ == "__main__":
   tf.test.main()

From 1e49db976a1275c1918813d86951fb2fa7f13717 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 14 Aug 2018 14:35:28 -0700
Subject: [PATCH 0582/2720] Part 2 of Glow, adds and test all the other
 operations to construct the entire network.

PiperOrigin-RevId: 208715014
---
 tensor2tensor/models/research/glow_ops.py     | 207 +++++++++++++++++-
 .../models/research/glow_ops_test.py          |  85 +++++--
 2 files changed, 271 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index f240fee94..619f7d080 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -18,19 +18,44 @@
 from __future__ import division
 from __future__ import print_function
 
+from functools import partial
 import numpy as np
 import scipy
+from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import registry
 import tensorflow as tf
 
 arg_scope = tf.contrib.framework.arg_scope
 add_arg_scope = tf.contrib.framework.add_arg_scope
 
 
+@registry.register_hparams
+def glow_hparams():
+  """Glow Hparams."""
+  hparams = common_hparams.basic_params1()
+  hparams.add_hparam("n_levels", 3)
+  hparams.add_hparam("n_bits_x", 8)
+  hparams.add_hparam("depth", 32)
+  hparams.add_hparam("affine_coupling_width", 512)
+  hparams.add_hparam("learn_prior", True)
+  return hparams
+
+
 def default_initializer(std=0.05):
   return tf.random_normal_initializer(0., std)
 
 
+def get_eps(dist, x):
+  """Z = (X - mu) / sigma."""
+  return (x - dist.loc) / dist.scale
+
+
+def set_eps(dist, eps):
+  """Z = eps * sigma + mu."""
+  return eps * dist.scale + dist.loc
+
+
 @add_arg_scope
 def get_variable_ddi(name, shape, initial_value, dtype=tf.float32, init=False,
                      trainable=True):
@@ -66,7 +91,7 @@ def actnorm(name, x, logscale_factor=3., reverse=False, init=False,
     objective: log(sum(s))
   """
   var_arg_scope = arg_scope([get_variable_ddi], trainable=trainable)
-  var_scope = tf.variable_scope(name)
+  var_scope = tf.variable_scope(name, reuse=tf.AUTO_REUSE)
   with var_scope, var_arg_scope:
     if not reverse:
       x = actnorm_center(name + "_center", x, reverse, init=init)
@@ -98,7 +123,7 @@ def actnorm_center(name, x, reverse=False, init=False):
     x_center: (x + b), if reverse is True and (x - b) otherwise.
   """
   shape = common_layers.shape_list(x)
-  with tf.variable_scope(name):
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     assert len(shape) == 2 or len(shape) == 4
     if len(shape) == 2:
       x_mean = tf.reduce_mean(x, [0], keepdims=True)
@@ -338,7 +363,7 @@ def nn(name, x, mid_channels, output_channels):
 
 
 @add_arg_scope
-def affine_coupling(name, x, mid_channels, reverse=False):
+def affine_coupling(name, x, mid_channels=512, reverse=False):
   """Reversible affine coupling layer.
 
   Args:
@@ -372,3 +397,179 @@ def affine_coupling(name, x, mid_channels, reverse=False):
     if reverse:
       objective *= -1
     return tf.concat([z1, z2], axis=3), objective
+
+
+@add_arg_scope
+def squeeze(name, x, factor=2, reverse=True):
+  """Block-wise spatial squeezing of x to increase the number of channels.
+
+  Args:
+    name: Used for variable scoping.
+    x: 4-D Tensor of shape (batch_size X H X W X C)
+    factor: Factor by which the spatial dimensions should be squeezed.
+    reverse: Squueze or unsqueeze operation.
+
+  Returns:
+    x: 4-D Tensor of shape (batch_size X (H//factor) X (W//factor) X
+       (cXfactor^2). If reverse is True, then it is factor = (1 / factor)
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    shape = common_layers.shape_list(x)
+    if factor == 1:
+      return x
+    height = int(shape[1])
+    width = int(shape[2])
+    n_channels = int(shape[3])
+    assert height % factor == 0 and width % factor == 0
+    if not reverse:
+      x = tf.reshape(x, [-1, height//factor, factor,
+                         width//factor, factor, n_channels])
+      x = tf.transpose(x, [0, 1, 3, 5, 2, 4])
+      x = tf.reshape(x, [-1, height//factor, width //
+                         factor, n_channels*factor*factor])
+    else:
+      x = tf.reshape(
+          x, (-1, height, width, int(n_channels/factor**2), factor, factor))
+      x = tf.transpose(x, [0, 1, 4, 2, 5, 3])
+      x = tf.reshape(x, (-1, int(height*factor),
+                         int(width*factor), int(n_channels/factor**2)))
+    return x
+
+
+@add_arg_scope
+def split_prior(name, x):
+  """Map x to the mean and log-scale of a Gaussian distribution."""
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    x_shape = common_layers.shape_list(x)
+    mean_log_scale = conv2d("conv2d", x, output_channels=2*x_shape[-1],
+                            apply_actnorm=False, conv_init="zeros")
+    mean = mean_log_scale[:, :, :, 0::2]
+    log_scale = mean_log_scale[:, :, :, 1::2]
+    return tf.distributions.Normal(mean, tf.exp(log_scale))
+
+
+@add_arg_scope
+def split(name, x, reverse=False, eps=None, eps_std=None):
+  """Splits / concatenates x into x1 and x2 across number of channels.
+
+  For the forward pass, x2 is assumed be gaussian,
+  i.e P(x2 | x1) ~ N(mu(x1), sigma(x1)) where mu and sigma are the outputs of
+  a network. For the reverse pass, x2 is determined from mu(x1) and sigma(x1).
+  This is deterministic/stochastic depending on whether eps is provided.
+
+  Args:
+    name:
+    x:
+    reverse: Forward or reverse pass.
+    eps: If eps is provided, x2
+    eps_std: Sample x2
+
+  Returns:
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    if not reverse:
+      x1, x2 = tf.split(x, num_or_size_splits=2, axis=-1)
+
+      # objective: P(x2|x1) ~N(x2 ; NN(x1))
+      x1_dist = split_prior("split_prior", x1)
+      logpb = tf.reduce_sum(x1_dist.log_prob(x2), axis=[1, 2, 3])
+
+      eps = get_eps(x1_dist, x2)
+      return x1, logpb, eps
+    else:
+      x1_dist = split_prior("split_prior", x)
+      if eps is not None:
+        x2 = set_eps(x1_dist, eps)
+      elif eps_std is not None:
+        x2 = eps_std * tf.random_normal(common_layers.shape_list(x))
+      else:
+        x2 = x1_dist.sample()
+      return tf.concat([x, x2], 3)
+
+
+@add_arg_scope
+def revnet_step(name, x, hparams, reverse=True):
+  """One step of glow generative flow.
+
+  Actnorm + invertible 1X1 conv + affine_coupling.
+
+  Args:
+    name: used for variable scope.
+    x: input
+    hparams: affine_coupling_width is the only hparam that is being used in
+             this function.
+    reverse: forward or reverse pass.
+  Returns:
+    z: Output of one step of reversible flow.
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    ops = [
+        partial(actnorm, name="actnorm", reverse=reverse),
+        partial(invertible_1x1_conv, name="invertible", reverse=reverse),
+        partial(affine_coupling, name="affine", reverse=reverse,
+                mid_channels=hparams.affine_coupling_width)]
+
+    if reverse:
+      ops = ops[::-1]
+
+    objective = 0.0
+    for op in ops:
+      x, curr_obj = op(x=x)
+      objective += curr_obj
+    return x, objective
+
+
+def revnet(name, x, hparams, reverse=True):
+  """'hparams.depth' steps of generative flow."""
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    steps = np.arange(hparams.depth)
+    if reverse:
+      steps = steps[::-1]
+
+    objective = 0.0
+    for step in steps:
+      x, curr_obj = revnet_step(
+          "revnet_%d" % step, x, hparams, reverse=reverse)
+      objective += curr_obj
+    return x, objective
+
+
+def encoder_decoder(name, x, hparams, eps=None, reverse=False):
+  """Glow encoder-decoder. n_levels of (Squeeze + Flow + Split.) operations."""
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+
+    objective = 0.0
+    all_eps = []
+
+    if not reverse:
+      # Squeeze + Flow + Split
+      for level in range(hparams.n_levels):
+        x = squeeze("squeeze_%d" % level, x, factor=2, reverse=False)
+
+        x, obj = revnet("revnet_%d" % level, x, hparams, reverse=False)
+        objective += obj
+
+        if level < hparams.n_levels - 1:
+          x, obj, eps = split("split_%d" % level, x, reverse=False)
+          objective += obj
+          all_eps.append(eps)
+      return x, objective, all_eps
+
+    else:
+      if eps and len(eps) != hparams.n_levels - 1:
+        raise ValueError("Expected length of eps to be %d, got %d" %
+                         (hparams.n_levels - 1, len(eps)))
+
+      for level in reversed(range(hparams.n_levels)):
+        if level < hparams.n_levels - 1:
+
+          curr_eps = None
+          if eps:
+            curr_eps = eps[level]
+          x = split("split_%d" % level, x, eps=curr_eps, reverse=True)
+
+        x, obj = revnet(
+            "revnet_%d" % level, x, hparams=hparams, reverse=True)
+        objective += obj
+        x = squeeze("squeeze_%d" % level, x, reverse=True)
+      return x, objective
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 0e42f71ea..271a58657 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -46,19 +46,24 @@ def test_actnorm(self):
         self.assertTrue(np.allclose(channel_mean, 0.0, atol=1e-3))
         self.assertTrue(np.allclose(channel_var, 1.0, atol=1e-3))
 
-  def test_invertible_conv(self):
+  def check_invertibility(self, op, name):
     with tf.Graph().as_default():
-      x_t = tf.random_uniform(shape=(16, 32, 32, 3))
-      activation, _ = glow_ops.invertible_1x1_conv("inv", x_t, reverse=False)
-      inv_activation, _ = glow_ops.invertible_1x1_conv(
-          "inv", activation, reverse=True)
+      x = tf.random_uniform(shape=(16, 32, 32, 4))
+
+      x_inv, _ = op(name, x, reverse=False)
+      x_inv_inv, _ = op(name, x_inv, reverse=True)
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())
-        diff = session.run(x_t - inv_activation)
-
-        # Test reversibility.
+        diff = session.run(x - x_inv_inv)
         self.assertTrue(np.allclose(diff, 0.0, atol=1e-5))
 
+  def test_invertibility(self):
+    rev_ops = [glow_ops.invertible_1x1_conv, glow_ops.affine_coupling,
+               glow_ops.actnorm]
+    names = ["inv_1X1_conv", "affine_coupling", "actnorm"]
+    for rev_op, name in zip(rev_ops, names):
+      self.check_invertibility(rev_op, name)
+
   def test_add_edge_bias(self):
     with tf.Graph().as_default():
       x = tf.random_uniform(shape=(16, 32, 32, 3))
@@ -107,20 +112,64 @@ def test_nn(self):
         # Initialized with zeros.
         self.assertTrue(np.allclose(nn_np, 0.0))
 
-  def test_affine_coupling(self):
-    """Test affine coupling reversibility."""
+  def test_split_prior(self):
     with tf.Graph().as_default():
-      rng = np.random.RandomState(0)
-      x = np.asarray(rng.rand(16, 3, 3, 32), dtype=np.float32)
-      x_t = tf.convert_to_tensor(x)
-      x_inv, _ = glow_ops.affine_coupling("affine", x_t, 512)
-      x_inv_inv, _ = glow_ops.affine_coupling(
-          "affine", x_inv, 512, reverse=True)
+      x = tf.random_uniform(shape=(16, 5, 5, 32))
+      x_prior = glow_ops.split_prior("split_prior", x)
+      mean_t, scale_t = x_prior.loc, x_prior.scale
+      with tf.Session() as session:
+        session.run(tf.global_variables_initializer())
+        mean, scale = session.run([mean_t, scale_t])
+        self.assertTrue(np.allclose(mean, 0.0))
+        self.assertTrue(np.allclose(scale, 1.0))
 
+  def test_split(self):
+    with tf.Graph().as_default():
+      x = tf.random_uniform(shape=(16, 5, 5, 32))
+      x_inv, _, eps = glow_ops.split("split", x)
+      x_inv_inv = glow_ops.split("split", x_inv, reverse=True, eps=eps)
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())
-        diff = tf.reduce_max(tf.abs(x - x_inv_inv))
-        self.assertTrue(np.allclose(session.run(diff), 0.0, atol=1e-5))
+        x_inv_np, diff = session.run([x_inv, x - x_inv_inv])
+        self.assertEqual(x_inv_np.shape, (16, 5, 5, 16))
+        self.assertTrue(np.allclose(diff, 0.0, atol=1e-5))
+
+  def check_revnet_reversibility(self, op, name):
+    with tf.Graph().as_default():
+      hparams = glow_ops.glow_hparams()
+      hparams.depth = 2
+      x = tf.random_uniform(shape=(16, 32, 32, 4), seed=0)
+      x_inv, _ = op(name, x, hparams, reverse=False)
+      x_inv_inv, _ = op(name, x_inv, hparams, reverse=True)
+      with tf.Session() as session:
+        session.run(tf.global_variables_initializer())
+        diff = session.run(x - x_inv_inv)
+        self.assertTrue(np.allclose(diff, 0.0, atol=1e-3))
+
+  def test_revnet_reversibility(self):
+    ops = [glow_ops.revnet_step, glow_ops.revnet]
+    names = ["revnet_step", "revnet"]
+    for op, name in zip(ops, names):
+      self.check_revnet_reversibility(op, name)
+
+  def test_encoder_decoder(self):
+    with tf.Graph().as_default():
+      hparams = glow_ops.glow_hparams()
+      hparams.n_levels = 2
+      hparams.depth = 2
+
+      x = tf.random_uniform(shape=(16, 64, 64, 4), seed=0)
+      x_inv, _, eps = glow_ops.encoder_decoder(
+          "encoder_decoder", x, hparams, reverse=False)
+      x_inv_inv, _ = glow_ops.encoder_decoder(
+          "encoder_decoder", x_inv, hparams, eps=eps, reverse=True)
+
+      with tf.Session() as session:
+        session.run(tf.global_variables_initializer())
+        diff, x_inv_np = session.run([x - x_inv_inv, x_inv])
+        self.assertTrue(x_inv_np.shape, (16, 8, 8, 64))
+        self.assertTrue(np.allclose(diff, 0.0, atol=1e-2))
+
 
 if __name__ == "__main__":
   tf.test.main()

From 71ef8c710a1fecc3a795e70e54fd378ae7990928 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 14 Aug 2018 15:00:49 -0700
Subject: [PATCH 0583/2720] Clean up old flag

PiperOrigin-RevId: 208719839
---
 tensor2tensor/layers/common_layers.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 52dae54f5..c3171b47a 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3457,9 +3457,6 @@ def should_generate_summaries():
   if tf.get_variable_scope().reuse:
     # Avoid generating separate summaries for different data shards
     return False
-  # TODO(rsepassi): Figure out a way to re-enable this line. Breaks all tests
-  # on Travis.
-  # return getattr(tf.flags.FLAGS, "enable_summaries", True)
   return True
 
 

From d3e2c4a9374cce370d04458197b6bbf30abff37d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 14 Aug 2018 15:51:09 -0700
Subject: [PATCH 0584/2720] Adding the edge network mpnn (handles vector valued
 edge features, no attention).

PiperOrigin-RevId: 208729162
---
 .../common_message_passing_attention.py       | 77 ++++++++++++++++++-
 1 file changed, 76 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_message_passing_attention.py b/tensor2tensor/layers/common_message_passing_attention.py
index 5de2adeef..2c5c61b3e 100644
--- a/tensor2tensor/layers/common_message_passing_attention.py
+++ b/tensor2tensor/layers/common_message_passing_attention.py
@@ -75,7 +75,7 @@ def multihead_graph_attention(query_antecedent,
 
   Returns:
     The result of the attention transformation. The output shape is
-        [batch_size, length_q, hidden_dim]
+        [batch_size, length_q, output_depth]
 
   Raises:
     ValueError: if the key depth or value depth are not divisible by the
@@ -827,3 +827,78 @@ def compute_values(edge_compatibility, v):
   # single N x V matrix for each batch.
   output = tf.reduce_sum(all_edge_values, axis=1)  # Shape [B, N, V].
   return output
+
+
+def precompute_edge_matrices(adjacency, hparams):
+  """Precompute the a_in and a_out tensors.
+
+  (we don't want to add to the graph everytime _fprop is called)
+  Args:
+    adjacency: placeholder of real valued vectors of shape [B, L, L, E]
+    hparams: tf.HParams object
+  Returns:
+    edge_matrices: [batch, L * D, L * D] the dense matrix for message passing
+    viewed as a block matrix (L,L) blocks of size (D,D). Each plot is a function
+    of the edge vector of the adjacency matrix at that spot.
+  """
+  batch_size, num_nodes, _, edge_dim = common_layers.shape_list(adjacency)
+
+  # build the edge_network for incoming edges
+  with tf.variable_scope("edge_network"):
+    x = tf.reshape(
+        adjacency, [batch_size * num_nodes * num_nodes, edge_dim],
+        name="adj_reshape_in")
+
+    for ip_layer in range(hparams.edge_network_layers):
+      name = "edge_network_layer_%d"%ip_layer
+      x = tf.layers.dense(common_layers.layer_preprocess(x, hparams),
+                          hparams.edge_network_hidden_size,
+                          activation=tf.nn.relu,
+                          name=name)
+    x = tf.layers.dense(common_layers.layer_preprocess(x, hparams),
+                        hparams.hidden_size**2,
+                        activation=None,
+                        name="edge_network_output")
+
+  # x = [batch * l * l, d *d]
+  edge_matrices_flat = tf.reshape(x, [batch_size, num_nodes,
+                                      num_nodes, hparams.hidden_size,
+                                      hparams.hidden_size])
+
+  # reshape to [batch, l * d, l *d]
+  edge_matrices = tf.reshape(
+      tf.transpose(edge_matrices_flat, [0, 1, 3, 2, 4]), [
+          -1, num_nodes * hparams.hidden_size,
+          num_nodes * hparams.hidden_size
+      ],
+      name="edge_matrices")
+
+  return edge_matrices
+
+
+def dense_message_pass(node_states, edge_matrices):
+  """Computes a_t from h_{t-1}, see bottom of page 3 in the paper.
+
+  Args:
+    node_states: [B, L, D] tensor (h_{t-1})
+    edge_matrices (tf.float32): [B, L*D, L*D]
+
+  Returns:
+    messages (tf.float32): [B, L, D] For each pair
+      of nodes in the graph a message is sent along both the incoming and
+      outgoing edge.
+  """
+  batch_size, num_nodes, node_dim = common_layers.shape_list(node_states)
+
+  # Stack the nodes as a big column vector.
+  h_flat = tf.reshape(
+      node_states, [batch_size, num_nodes * node_dim, 1], name="h_flat")
+
+  messages = tf.reshape(
+      tf.matmul(edge_matrices, h_flat), [batch_size * num_nodes, node_dim],
+      name="messages_matmul")
+
+  message_bias = tf.get_variable("message_bias", shape=node_dim)
+  messages = messages + message_bias
+  messages = tf.reshape(messages, [batch_size, num_nodes, node_dim])
+  return messages

From de31dea48d7115554578ef1d7232ab3f5260d3c7 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 14 Aug 2018 16:09:00 -0700
Subject: [PATCH 0585/2720] Internal change

PiperOrigin-RevId: 208732434
---
 tensor2tensor/utils/t2t_model.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index acbbe7aaf..ced52766d 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -118,6 +118,19 @@ def __init__(self,
     self._eager_var_store = create_eager_var_store()
     if self._problem_hparams:
       self._create_modalities(self._problem_hparams, self._hparams)
+    if not common_layers.is_xla_compiled():
+      self.summarize_hparams()
+
+  def summarize_hparams(self):
+    def create_hparams_summary(hparams, name):
+      hparams_strs = [tf.convert_to_tensor([k, str(v)])
+                      for k, v in hparams.values().items()]
+      tf.summary.text(name, tf.stack(hparams_strs))
+
+    create_hparams_summary(self._hparams, "%s_hparams" % self.name)
+    if self._problem_hparams:
+      create_hparams_summary(self._problem_hparams,
+                             "%s_problem_hparams" % self.name)
 
   # Replace the two methods below in order to add custom SessionRunHooks to
   # the training procedure.

From 91e8c5cb69051f58c5d95672282cdc1f0e68a0b2 Mon Sep 17 00:00:00 2001
From: Youlong Cheng <ylc@google.com>
Date: Tue, 14 Aug 2018 17:17:23 -0700
Subject: [PATCH 0586/2720] Add an empty __init__.py in the new research
 folder.

PiperOrigin-RevId: 208743414
---
 tensor2tensor/mesh_tensorflow/research/__init__.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)
 create mode 100644 tensor2tensor/mesh_tensorflow/research/__init__.py

diff --git a/tensor2tensor/mesh_tensorflow/research/__init__.py b/tensor2tensor/mesh_tensorflow/research/__init__.py
new file mode 100644
index 000000000..dba7ece95
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/research/__init__.py
@@ -0,0 +1,14 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.

From 3ef7c41554b6446b533a5542d42b6f8fc7b0f8cf Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Tue, 14 Aug 2018 18:32:58 -0700
Subject: [PATCH 0587/2720] Rewrite mixture-of-experts gating function in mtf.

PiperOrigin-RevId: 208752587
---
 .../mesh_tensorflow/mesh_tensorflow.py        |  36 ++++
 tensor2tensor/mesh_tensorflow/mtf_layers.py   | 155 ++++++++----------
 2 files changed, 100 insertions(+), 91 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index f3e5bdde7..3adc48804 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -1453,6 +1453,16 @@ def greater(x1, x2, output_shape=None):
       tf.greater, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
 
 
+def less_equal(x1, x2, output_shape=None):
+  return binary_op_with_broadcasting(
+      tf.less_equal, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
+
+
+def greater_equal(x1, x2, output_shape=None):
+  return binary_op_with_broadcasting(
+      tf.greater_equal, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
+
+
 def equal(x1, x2, output_shape=None):
   return binary_op_with_broadcasting(
       tf.equal, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
@@ -1567,6 +1577,10 @@ def __init__(self, x, output_shape, reduction_fn_string, name=None):
   def gradient(self, grad_ys):
     if self._reduction_fn_string == "SUM":
       return [broadcast(grad_ys[0], self.inputs[0].shape)]
+    elif (self._reduction_fn_string == "MAX" or
+          self._reduction_fn_string == "MIN"):
+      return [cast(equal(self.inputs[0], self.outputs[0]), self.inputs[0].dtype)
+              * grad_ys[0]]
     else:
       raise ValueError("Gradients to other reductions not implemented")
 
@@ -1799,6 +1813,28 @@ def unstack(x, dim, name=None):
   return UnstackOperation(x, dim, name).outputs
 
 
+def cumsum(x, dim, exclusive=False):
+  """Cumulative sum.
+
+  Args:
+    x: a Tensor
+    dim: a Dimension
+    exclusive: a boolean
+
+  Returns:
+    a Tensor with the same shape as x.
+  """
+  new_name = "tmp_dim_cumsum"
+  new_dim = Dimension(new_name, dim.size)
+  new_shape = x.shape.rename_dimension(dim.name, new_name)
+  comparator = less if exclusive else less_equal
+  m = cast(
+      comparator(range(x.mesh, dim, dtype=tf.float32),
+                 range(x.mesh, new_dim, dtype=tf.float32)), x.dtype)
+  ret = einsum([x, m], output_shape=new_shape)
+  return reshape(ret, x.shape)
+
+
 def _einsum_helper(input_shapes, output_shape, mesh_impl):
   """Returns slicewise function and reduced mesh dimensions.
 
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers.py b/tensor2tensor/mesh_tensorflow/mtf_layers.py
index effbfc553..17f789a1b 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers.py
@@ -17,8 +17,8 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-import functools
-from tensor2tensor.layers import common_layers
+
+
 from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
 import tensorflow as tf
 
@@ -676,16 +676,8 @@ def moe_v0(inputs,
   # shape = [batch_dim, length_dim, experts_dim_unsplit]
   gates = mtf.softmax(dense(inputs, experts_dim_unsplit), experts_dim_unsplit)
 
-  assignment_shape = mtf.Shape(
-      [batch_dim, length_dim, experts_dim_unsplit, expert_capacity_dim])
-
-  backward_assignment = mtf.slicewise(
-      functools.partial(
-          _truncated_top_2_gating, expert_capacity=expert_capacity),
-      [gates],
-      output_shape=assignment_shape,
-      splittable_dims=[batch_dim],
-      name="backward_assignment")
+  backward_assignment = _truncated_top_2_gating_mtf(
+      gates, length_dim, experts_dim_unsplit, expert_capacity_dim)
 
   forward_assignment = mtf.cast(
       mtf.cast(backward_assignment, tf.bool), inputs.dtype)
@@ -734,7 +726,8 @@ def cv_squared(x):
   return variance / (mtf.square(mean) + epsilon)
 
 
-def _truncated_top_2_gating(gates, expert_capacity, show_summaries=False):
+def _truncated_top_2_gating_mtf(
+    gates, group_dim, experts_dim, expert_capacity_dim):
   """Compute gating for mixture-of-experts in TensorFlow.
 
   gates is usually the output of a softmax function.
@@ -746,108 +739,88 @@ def _truncated_top_2_gating(gates, expert_capacity, show_summaries=False):
   call it from both places.
 
   Args:
-    gates: a Tensor with shape [batch, length, num_experts]
-    expert_capacity: an integer
-    show_summaries: a boolean
+    gates: a Tensor
+    group_dim: one dimension of gates
+    experts_dim: one dimension of gates
+    expert_capacity_dim: a Dimension not in gates
 
   Returns:
-    a Tensor with shape [batch, length, num_experts, expert_capacity]
+    a Tensor with shape gates.shape + expert_capacity_dim
+
+  Raises:
+    ValueError: if group_dim has size >256
   """
-  def _to_float(x):
-    return tf.cast(x, gates.dtype)
-  batch = tf.shape(gates)[0]
-  length = tf.shape(gates)[1]
-  num_experts = tf.shape(gates)[2]
-  expert_capacity_f = _to_float(expert_capacity)
-  # Find the top expert for each position.
-  gate_1, index_1 = common_layers.top_1_tpu(gates)
-  # [batch, length, num_experts]
-  mask_1 = tf.one_hot(index_1, num_experts, dtype=gates.dtype)
-  # [batch, length, num_experts]
+  gates = mtf.to_float(gates)
+  expert_capacity_f = float(expert_capacity_dim.size)
+  # Find the top expert for each position. shape=[batch, group]
+  index_1, gate_1 = mtf.top_1(gates, experts_dim)
+  # [batch, group, experts]
+  mask_1 = mtf.one_hot(index_1, experts_dim, dtype=gates.dtype)
+
+  if expert_capacity_dim.size > 256:
+    # using mtf.cumsum (implemented on TPU as bfloat16 matmul) to compute
+    # position in the mini-batch sent to the expert.  This will cause
+    # very bad things to happen if expert_capacity_dim > 256.
+    raise ValueError(
+        "expert_capacity_dim.size must be <=256 to avoid roundoff errors in"
+        " indices - got %s" % (expert_capacity_dim,))
+  # [batch, group, experts]
   # This is the position within the expert's mini-batch for this sequence
-  position_in_expert_1 = common_layers.cumsum(
-      mask_1, axis=1, exclusive=True) * mask_1
-  # Remove the elements that don't fit.
-  mask_1 *= _to_float(tf.less(position_in_expert_1, expert_capacity_f))
-  # [batch, 1, num_experts]
+  position_in_expert_1 = mtf.cumsum(mask_1, group_dim, exclusive=True) * mask_1
+  # Remove the elements that don't fit. [batch, group, experts]
+  mask_1 *= mtf.to_float(mtf.less(position_in_expert_1, expert_capacity_f))
+  # [batch, experts]
   # How many examples in this sequence go to this expert
-  mask_1_count = tf.reduce_sum(mask_1, axis=1, keepdims=True)
-  # [batch, length] - mostly ones, but zeros where something didn't fit
-  mask_1_flat = tf.reduce_sum(mask_1, axis=2)
-  position_in_expert_1 = tf.reduce_sum(position_in_expert_1, axis=2)
-  # Weight assigned to first expert.
+  mask_1_count = mtf.reduce_sum(mask_1, reduced_dim=group_dim)
+  # [batch, group] - mostly ones, but zeros where something didn't fit
+  mask_1_flat = mtf.reduce_sum(mask_1, reduced_dim=experts_dim)
+  # [batch, group]
+  position_in_expert_1 = mtf.reduce_sum(
+      position_in_expert_1, reduced_dim=experts_dim)
+  # Weight assigned to first expert.  [batch, group]
   gate_1 *= mask_1_flat
 
   # Pick a second-place expert for each position.
   # We first mask out the experts that we expect to be over-capacity
+  # [batch, experts]
   space_remaining = expert_capacity_f - mask_1_count
-  use_rate = (mask_1_count + 1.0) / _to_float(length)
+  use_rate = (mask_1_count + 1.0) / float(group_dim.size)
   # At what point in the sequence do we expect the expert to be full.
+  # [batch, experts]
   expected_exhaustion_pos = space_remaining / use_rate
-  # A Tensor with shape [batch, length, num_experts] representing a boolean
+  # A Tensor with shape [batch, group, experts] representing a boolean
   #   - whether we expect that the expert will already be full.
-  expected_exhausted = _to_float(tf.greater(
-      tf.reshape(_to_float(tf.range(length)), [1, length, 1]),
-      expected_exhaustion_pos))
+  expected_exhausted = mtf.to_float(mtf.greater(
+      mtf.range(gates.mesh, group_dim, tf.float32), expected_exhaustion_pos))
   masked_gates = gates - mask_1 - expected_exhausted
   # This section is similar to the section above.
-  gate_2, index_2 = common_layers.top_1_tpu(masked_gates)
-  # [batch, length, num_experts]
-  mask_2 = tf.one_hot(index_2, num_experts, dtype=gates.dtype)
+  # [batch, group]
+  index_2, gate_2 = mtf.top_1(masked_gates, experts_dim)
+  # [batch, group, experts]
+  mask_2 = mtf.one_hot(index_2, experts_dim, dtype=gates.dtype)
+  # [batch, group, experts]
   position_in_expert_2 = (
-      common_layers.cumsum(mask_2, axis=1, exclusive=True) + mask_1_count)
+      mtf.cumsum(mask_2, group_dim, exclusive=True) + mask_1_count)
   position_in_expert_2 *= mask_2
-  mask_2 *= _to_float(tf.less(position_in_expert_2, expert_capacity_f))
-  mask_2_count = tf.reduce_sum(mask_2, axis=1, keepdims=True)
-  mask_2_flat = tf.reduce_sum(mask_2, axis=2)
-  position_in_expert_2 = tf.reduce_sum(position_in_expert_2, axis=2)
+  mask_2 *= mtf.to_float(mtf.less(position_in_expert_2, expert_capacity_f))
+  # mask_2_count = mtf.reduce_sum(mask_2, reduced_dim=experts_dim)
+  mask_2_flat = mtf.reduce_sum(mask_2, reduced_dim=experts_dim)
+  position_in_expert_2 = mtf.reduce_sum(
+      position_in_expert_2, reduced_dim=experts_dim)
   gate_2 *= mask_2_flat
 
-  # What fraction didn't fit - show summaries
-  if show_summaries:
-    miss_rate_1 = (
-        1.0 - tf.reduce_sum(mask_1_count) / _to_float(batch * length))
-    miss_rate_2 = (
-        1.0 - tf.reduce_sum(mask_2_count) / _to_float(batch * length))
-    tf.summary.scalar("miss_rate_1", miss_rate_1)
-    tf.summary.scalar("miss_rate_2", miss_rate_2)
-
   # renormalize the two gate values to add up to 1
   denom = gate_1 + gate_2 + 1e-9
   gate_1 /= denom
   gate_2 /= denom
 
-  # inputs: [batch, length, input_size]
-  # forward_assignment: [batch, length, num_experts * expert_capacity]
-  # expert_inputs: [batch, num_experts * expert_capacity, input_size]
-
-  segment_ids_forward_1 = (
-      (index_1 * expert_capacity) +
-      tf.to_int32(position_in_expert_1) +
-      tf.to_int32(1.0 - mask_1_flat) * (num_experts * expert_capacity))
-
-  segment_ids_forward_2 = (
-      (index_2 * expert_capacity) +
-      tf.to_int32(position_in_expert_2) +
-      tf.to_int32(1.0 - mask_2_flat) * (num_experts * expert_capacity))
-
-  # Gather and scatter are painfully slow on TPU.
-  # We will use one_hot and matmul instead.
-
-  # [batch, length, num_experts * expert_capacity]
-  one_hot_1 = tf.one_hot(
-      segment_ids_forward_1, num_experts * expert_capacity, dtype=gates.dtype)
-  one_hot_2 = tf.one_hot(
-      segment_ids_forward_2, num_experts * expert_capacity, dtype=gates.dtype)
-
-  # expert_output: [batch, num_experts * expert_capacity, output_size]
-  # backward_assignment: [batch, length, num_experts * expert_capacity]
-  # output: [batch, length, output_size]
+  # [batch, group, experts, expert_capacity]
   assignment = (
-      one_hot_1 * tf.cast(tf.expand_dims(gate_1, 2), gates.dtype) +
-      one_hot_2 * tf.cast(tf.expand_dims(gate_2, 2), gates.dtype))
-
-  assignment = tf.reshape(
-      assignment, [batch, length, num_experts, expert_capacity])
+      gate_1 * mask_1_flat
+      * mtf.one_hot(index_1, experts_dim)
+      * mtf.one_hot(mtf.to_int32(position_in_expert_1), expert_capacity_dim) +
+      gate_2 * mask_2_flat
+      * mtf.one_hot(index_2, experts_dim)
+      * mtf.one_hot(mtf.to_int32(position_in_expert_2), expert_capacity_dim))
 
   return assignment

From 609fffe18cc022463824185b312cb6402c371337 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 14 Aug 2018 18:53:37 -0700
Subject: [PATCH 0588/2720] remove restore hook from vqa

PiperOrigin-RevId: 208754647
---
 .../models/research/vqa_attention.py          | 20 +++----
 .../research/vqa_recurrent_self_attention.py  | 20 +++----
 .../models/research/vqa_self_attention.py     | 56 +++++++++----------
 3 files changed, 48 insertions(+), 48 deletions(-)

diff --git a/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index 920ef64a6..37787fe87 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -23,7 +23,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import vqa_layers
 from tensor2tensor.utils import registry
-from tensor2tensor.utils import restore_hook
+# from tensor2tensor.utils import restore_hook
 from tensor2tensor.utils import t2t_model
 
 import tensorflow as tf
@@ -38,15 +38,15 @@
 class VqaAttentionBaseline(t2t_model.T2TModel):
   """Attention baseline model for VQA."""
 
-  @staticmethod
-  def train_hooks():
-    restore_resnet_hook = restore_hook.RestoreHook(
-        # TODO(zichaoy): hard code the path given static function.
-        checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
-        new_model_scope="vqa_attention_baseline/body/",
-        old_model_scope="resnet_v1_152/",
-    )
-    return [restore_resnet_hook]
+  # @staticmethod
+  # def train_hooks():
+  #   restore_resnet_hook = restore_hook.RestoreHook(
+  #       # TODO(zichaoy): hard code the path given static function.
+  #       checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
+  #       new_model_scope="vqa_attention_baseline/body/",
+  #       old_model_scope="resnet_v1_152/",
+  #   )
+  #   return [restore_resnet_hook]
 
   def body(self, features):
     hp = self.hparams
diff --git a/tensor2tensor/models/research/vqa_recurrent_self_attention.py b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
index e54de3630..387bc0484 100644
--- a/tensor2tensor/models/research/vqa_recurrent_self_attention.py
+++ b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
@@ -27,7 +27,7 @@
 from tensor2tensor.models.research import universal_transformer_util
 from tensor2tensor.models.research import vqa_attention
 from tensor2tensor.utils import registry
-from tensor2tensor.utils import restore_hook
+# from tensor2tensor.utils import restore_hook
 
 import tensorflow as tf
 
@@ -38,15 +38,15 @@
 class VqaRecurrentSelfAttention(vqa_attention.VqaAttentionBaseline):
   """Recurrent Self attention both on image and question."""
 
-  @staticmethod
-  def train_hooks():
-    restore_resnet_hook = restore_hook.RestoreHook(
-        # TODO(zichaoy): hard code the path given static function.
-        checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
-        new_model_scope="vqa_recurrent_self_attention/body/",
-        old_model_scope="resnet_v1_152/",
-    )
-    return [restore_resnet_hook]
+  # @staticmethod
+  # def train_hooks():
+  #   restore_resnet_hook = restore_hook.RestoreHook(
+  #       # TODO(zichaoy): hard code the path given static function.
+  #       checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
+  #       new_model_scope="vqa_recurrent_self_attention/body/",
+  #       old_model_scope="resnet_v1_152/",
+  #   )
+  #   return [restore_resnet_hook]
 
   def body(self, features):
     hp = self.hparams
diff --git a/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
index f556b1c78..e6794dde8 100644
--- a/tensor2tensor/models/research/vqa_self_attention.py
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -24,7 +24,7 @@
 from tensor2tensor.layers import vqa_layers
 from tensor2tensor.models.research import vqa_attention
 from tensor2tensor.utils import registry
-from tensor2tensor.utils import restore_hook
+# from tensor2tensor.utils import restore_hook
 
 import tensorflow as tf
 
@@ -35,15 +35,15 @@
 class VqaSelfAttention(vqa_attention.VqaAttentionBaseline):
   """Self attention both on image and question."""
 
-  @staticmethod
-  def train_hooks():
-    restore_resnet_hook = restore_hook.RestoreHook(
-        # TODO(zichaoy): hard code the path given static function.
-        checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
-        new_model_scope="vqa_self_attention/body/",
-        old_model_scope="resnet_v1_152/",
-    )
-    return [restore_resnet_hook]
+  # @staticmethod
+  # def train_hooks():
+  #   restore_resnet_hook = restore_hook.RestoreHook(
+  #       # TODO(zichaoy): hard code the path given static function.
+  #       checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
+  #       new_model_scope="vqa_self_attention/body/",
+  #       old_model_scope="resnet_v1_152/",
+  #   )
+  #   return [restore_resnet_hook]
 
   def body(self, features):
     hp = self.hparams
@@ -129,15 +129,15 @@ def body(self, features):
 class VqaCombinedSelfAttention(VqaSelfAttention):
   """Combined Self attention both on image and question."""
 
-  @staticmethod
-  def train_hooks():
-    restore_resnet_hook = restore_hook.RestoreHook(
-        # TODO(zichaoy): hard code the path given static function.
-        checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
-        new_model_scope="vqa_combined_self_attention/body/",
-        old_model_scope="resnet_v1_152/",
-    )
-    return [restore_resnet_hook]
+  # @staticmethod
+  # def train_hooks():
+  #   restore_resnet_hook = restore_hook.RestoreHook(
+  #       # TODO(zichaoy): hard code the path given static function.
+  #       checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
+  #       new_model_scope="vqa_combined_self_attention/body/",
+  #       old_model_scope="resnet_v1_152/",
+  #   )
+  #   return [restore_resnet_hook]
 
   def body(self, features):
     hp = self.hparams
@@ -194,15 +194,15 @@ def body(self, features):
 class VqaIterativeCombinedSelfAttention(VqaSelfAttention):
   """Combined Self attention both on image and question."""
 
-  @staticmethod
-  def train_hooks():
-    restore_resnet_hook = restore_hook.RestoreHook(
-        # TODO(zichaoy): hard code the path given static function.
-        checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
-        new_model_scope="vqa_combined_self_attention/body/",
-        old_model_scope="resnet_v1_152/",
-    )
-    return [restore_resnet_hook]
+  # @staticmethod
+  # def train_hooks():
+  #   restore_resnet_hook = restore_hook.RestoreHook(
+  #       # TODO(zichaoy): hard code the path given static function.
+  #       checkpoint_path="/home/zichaoy/resnet_v1_152.ckpt",
+  #       new_model_scope="vqa_combined_self_attention/body/",
+  #       old_model_scope="resnet_v1_152/",
+  #   )
+  #   return [restore_resnet_hook]
 
   def body(self, features):
     hp = self.hparams

From 3d418439076c1732b0703b7dbc1b6269aef7c469 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 15 Aug 2018 10:03:38 -0700
Subject: [PATCH 0589/2720] Refactor RL code

PiperOrigin-RevId: 208837067
---
 .../data_generators/gym_problems_specs.py     | 235 +++-----------
 .../data_generators/gym_problems_test.py      |   2 +-
 tensor2tensor/data_generators/video_utils.py  |   2 +-
 tensor2tensor/models/research/next_frame.py   |   5 -
 tensor2tensor/rl/README.md                    | 100 ++++--
 tensor2tensor/rl/datagen_with_agent.py        |  68 +++++
 tensor2tensor/rl/rl_trainer_lib.py            |   5 +-
 tensor2tensor/rl/rl_trainer_lib_test.py       |  54 ++--
 ...l_experiment.py => trainer_model_based.py} | 288 ++++++++++++------
 ...=> trainer_model_based_stochastic_test.py} |   7 +-
 ...nt_test.py => trainer_model_based_test.py} |   7 +-
 ...2t_rl_trainer.py => trainer_model_free.py} |   0
 12 files changed, 406 insertions(+), 367 deletions(-)
 create mode 100644 tensor2tensor/rl/datagen_with_agent.py
 rename tensor2tensor/rl/{model_rl_experiment.py => trainer_model_based.py} (78%)
 rename tensor2tensor/rl/{model_rl_experiment_stochastic_test.py => trainer_model_based_stochastic_test.py} (83%)
 rename tensor2tensor/rl/{model_rl_experiment_test.py => trainer_model_based_test.py} (83%)
 rename tensor2tensor/rl/{t2t_rl_trainer.py => trainer_model_free.py} (100%)

diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index e957b67ab..f42c6bc82 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -30,6 +30,11 @@
 # pylint: enable=g-multiple-import
 from tensor2tensor.utils import registry
 
+# Game list from our list of ROMs
+# Removed because XDeterministic-v4 did not exist:
+# * adventure
+# * defender
+# * kaboom
 ATARI_GAMES = [
     "air_raid", "alien", "amidar", "assault", "asterix", "asteroids",
     "atlantis", "bank_heist", "battle_zone", "beam_rider", "berzerk", "bowling",
@@ -44,10 +49,6 @@
     "up_n_down", "venture", "video_pinball", "wizard_of_wor", "yars_revenge",
     "zaxxon"
 ]
-# Removed because XDeterministic-v4 did not exist:
-# * adventure
-# * defender
-# * kaboom
 
 # Subset of games with promissing results on model based training.
 ATARI_WHITELIST_GAMES = [
@@ -65,56 +66,9 @@
     "pong",
     "road_runner",
     "seaquest",
-    "wrapped_full_pong",  # TODO(blazej): check if we get equally good results
-]                         # on vanilla pong.
-
-
-@registry.register_problem
-class GymPongRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  # Hard-coding num_actions, frame_height, frame_width to avoid loading
-  # libale.so file.
-  @property
-  def num_actions(self):
-    return 6
-
-  @property
-  def frame_height(self):
-    return 210
-
-  @property
-  def frame_width(self):
-    return 160
-
-  @property
-  def env_name(self):
-    return "PongDeterministic-v4"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-
-@registry.register_problem
-class GymWrappedPongRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  @property
-  def env_name(self):
-    return "T2TPongWarmUp20RewSkip200Steps-v1"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
+    # TODO(blazej): check if we get equally good results on vanilla pong.
+    "wrapped_full_pong",
+]
 
 
 @registry.register_problem
@@ -138,79 +92,6 @@ def num_testing_steps(self):
     return 100
 
 
-@registry.register_problem
-class GymWrappedBreakoutRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  @property
-  def env_name(self):
-    return "T2TBreakoutWarmUp20RewSkip500Steps-v1"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnPong(GymSimulatedDiscreteProblem,
-                                                 GymPongRandom):
-  """Simulated pong."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_pong"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymFreewayRandom(GymDiscreteProblem):
-  """Freeway game, random actions."""
-
-  @property
-  def env_name(self):
-    return "FreewayDeterministic-v4"
-
-  @property
-  def min_reward(self):
-    return 0
-
-  @property
-  def num_rewards(self):
-    return 2
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnPong(GymRealDiscreteProblem, GymPongRandom):
-  pass
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnFreeway(GymRealDiscreteProblem,
-                                           GymFreewayRandom):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedPong(
-    GymSimulatedDiscreteProblem, GymWrappedPongRandom):
-  """Simulated pong."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_pong"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
 @registry.register_problem
 class GymDiscreteProblemWithAgentOnWrappedFullPong(GymRealDiscreteProblem,
                                                    GymWrappedFullPongRandom):
@@ -257,74 +138,8 @@ def num_testing_steps(self):
     return 100
 
 
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedBreakout(GymRealDiscreteProblem,
-                                                   GymWrappedBreakoutRandom):
-  pass
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedBreakoutAe(
-    GymDiscreteProblemWithAgentOnWrappedBreakout):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedBreakout(
-    GymSimulatedDiscreteProblem, GymWrappedBreakoutRandom):
-  """Simulated breakout."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_breakout"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedPong(GymRealDiscreteProblem,
-                                               GymWrappedPongRandom):
-  """GymDiscreteProblemWithAgentOnWrappedPong."""
-
-  # Hard-coding num_actions, frame_height, frame_width to avoid loading
-  # libale.so file.
-  @property
-  def num_actions(self):
-    return 6
-
-  @property
-  def frame_height(self):
-    return 210
-
-  @property
-  def frame_width(self):
-    return 160
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedPongAe(  # With autoencoder.
-    GymDiscreteProblemWithAgentOnWrappedPong):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnFreeway(GymSimulatedDiscreteProblem,
-                                                    GymFreewayRandom):
-  """Simulated freeway."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_freeway"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
 class GymClippedRewardRandom(GymDiscreteProblem):
-  """Base class for clipped reward games."""
+  """Abstract base class for clipped reward games."""
 
   @property
   def env_name(self):
@@ -339,10 +154,25 @@ def num_rewards(self):
     return 3
 
 
-def dynamically_create_gym_clipped_reward_problem(game_name):
-  """Dynamically create env wrapper and Problems for game."""
-  # e.g. game_name == bank_heist
-  assert game_name in ATARI_GAMES
+def create_problems_for_game(game_name, clipped_reward=True):
+  """Create and register problems for game_name.
+
+  Args:
+    game_name: str, one of the games in ATARI_GAMES, e.g. "bank_heist".
+    clipped_reward: bool, whether the rewards should be clipped. False is not
+      yet supported.
+
+  Returns:
+    dict of problems with keys ("base", "agent", "simulated").
+
+  Raises:
+    ValueError: if clipped_reward=False or game_name not in ATARI_GAMES.
+  """
+  if not clipped_reward:
+    raise ValueError("Creating problems without clipped reward is not "
+                     "yet supported.")
+  if game_name not in ATARI_GAMES:
+    raise ValueError("Game %s not in ATARI_GAMES" % game_name)
   camel_game_name = "".join(
       [w[0].upper() + w[1:] for w in game_name.split("_")])
   env_name = "%sDeterministic-v4" % camel_game_name
@@ -355,7 +185,8 @@ def dynamically_create_gym_clipped_reward_problem(game_name):
           env=env_name, reward_clipping=True))
 
   # Create and register the Random and WithAgent Problem classes
-  problem_cls = type(camel_game_name + "Random", (GymClippedRewardRandom,),
+  problem_cls = type("Gym%sRandom" % camel_game_name,
+                     (GymClippedRewardRandom,),
                      {"env_name": wrapped_env_name})
   with_agent_cls = type("GymDiscreteProblemWithAgentOn%s" % camel_game_name,
                         (GymRealDiscreteProblem, problem_cls), {})
@@ -369,3 +200,9 @@ def dynamically_create_gym_clipped_reward_problem(game_name):
           "num_testing_steps": 100
       })
   registry.register_problem(simulated_cls)
+
+  return {
+      "base": problem_cls,
+      "agent": with_agent_cls,
+      "simulated": simulated_cls,
+  }
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index 061a53dc9..ebdbcf021 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -35,7 +35,7 @@ def setUpClass(cls):
     os.mkdir(cls.tmp_dir)
 
   def testGymAtariBoots(self):
-    problem = gym_problems_specs.GymPongRandom()
+    problem = gym_problems_specs.create_problems_for_game("pong")["base"]()
     self.assertEqual(210, problem.frame_height)
 
 
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 38c8d9f16..2877b0655 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -300,7 +300,7 @@ def check_integrity_and_batch(*datasets):
       batch_dataset = preprocessed_dataset.apply(
           tf.contrib.data.batch_and_drop_remainder(num_frames))
     dataset = batch_dataset.map(features_from_batch)
-    dataset = dataset.shuffle(hparams.shuffle_buffer_size)
+    dataset = dataset.shuffle(hparams.get("shuffle_buffer_size", 128))
     return dataset
 
   def eval_metrics(self):
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame.py
index a49f86fa4..6f35f6a78 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame.py
@@ -28,9 +28,6 @@
 
 import tensorflow as tf
 
-tfl = tf.layers
-tfcl = tf.contrib.layers
-
 
 @registry.register_model
 class NextFrameBasic(t2t_model.T2TModel):
@@ -161,5 +158,3 @@ def logits_to_samples(logits):
 
     # Return results.
     return results
-
-
diff --git a/tensor2tensor/rl/README.md b/tensor2tensor/rl/README.md
index 50fc73cde..e38c9bd5e 100644
--- a/tensor2tensor/rl/README.md
+++ b/tensor2tensor/rl/README.md
@@ -1,53 +1,97 @@
 # Tensor2Tensor experimental Model-Based Reinforcement Learning.
 
-The rl package intention is to provide the ability to run reinforcement
-algorithms within TensorFlow's computation graph, in order to do model-based
-RL using environment models from Tensor2Tensor. It's very experimental
-for now and under heavy development.
+**Note**: Experimental and under development.
 
-Currently the only supported algorithm is Proximal Policy Optimization - PPO.
+The `rl` package provides the ability to run model-based reinforcement learning
+algorithms using models trained with Tensor2Tensor.
 
-# Sample usages
+Currently this entails alternating model training and agent training using
+Proximal Policy Optimization (PPO). See `trainer_model_based.py`.
 
-## Training agent in the Pendulum-v0 environment.
+As a baseline, you can also run PPO without the model using
+`trainer_model_free.py`.
+
+## Model-based training
+
+Alternate training a world model and a PPO agent within that model using the
+base hyperparameters on Freeway:
+
+```
+python -m tensor2tensor.rl.trainer_model_based \
+  --output_dir=$OUT_DIR \
+  --loop_hparams_set=rl_modelrl_base \
+  --loop_hparams='game=freeway'
+```
+
+All hyperparameter sets are defined in `trainer_model_based.py` and are derived
+from `rl_modelrl_base`.
+
+The hyperparameters for the environment model and agent are nested within the
+`loop_hparams` by name. For example:
+
+```
+  generative_model="next_frame_basic",
+  generative_model_params="next_frame_pixel_noise",
+  ppo_params="ppo_pong_base",
+```
+
+## Model-free training
+
+**TODO(piotrmilos): Update**
+
+Training an agent in `Pendulum-v0`:
 
 ```
-python rl/t2t_rl_trainer.py \
+python -m tensor2tensor.rl.trainer_model_free \
   --problem=Pendulum-v0 \
   --hparams_set ppo_continuous_action_base \
-  [--output_dir dir_location]
+  --output_dir $OUT_DIR
 ```
 
-## Training agent in the PongNoFrameskip-v0 environment.
+Training an agent in `PongNoFrameskip-v0`:
 
 ```
-python tensor2tensor/rl/t2t_rl_trainer.py \
+python -m tensor2tensor.rl.trainer_model_free \
   --problem stacked_pong \
   --hparams_set ppo_atari_base \
   --hparams num_agents=5 \
-  [--output_dir dir_location]
+  --output_dir dir_location
+```
+
+## Model training on random trajectories
+
+Generate trajectories with a random policy:
+
+```
+python -m tensor2tensor.rl.datagen_with_agent \
+  --data_dir=$HOME/t2t/data \
+  --tmp_dir=$HOME/t2t/tmp \
+  --game=pong \
+  --num_env_steps=30000
 ```
 
-## Generation of trajectories data
+Train model on trajectories:
 
 ```
-python tensor2tensor/bin/t2t-datagen \
-  --data_dir=~/t2t_data \
-  --tmp_dir=~/t2t_data/tmp \
-  --problem=gym_pong_trajectories_from_policy \
-  --agent_policy_path [model]
+python -m tensor2tensor.bin.t2t_trainer \
+  --data_dir=$HOME/t2t/data \
+  --output_dir=$HOME/t2t/train/pong_model \
+  --problem=gym_pong_random \
+  --model=next_frame_basic \
+  --hparams_set=next_frame
 ```
 
-## Training model for frames generation based on randomly played games
+
+## Collect trajectories using a trained agent
 
 ```
-python tensor2tensor/bin/t2t-trainer \
-  --generate_data \
-  --data_dir=~/t2t_data \
-  --output_dir=~/t2t_data/output \
-  --problem=gym_pong_random5k \
-  --model=basic_conv_gen \
-  --hparams_set=basic_conv_small \
-  --train_steps=1000 \
-  --eval_steps=10
+python -m tensor2tensor.rl.datagen_with_agent \
+  --data_dir=$HOME/t2t/data \
+  --tmp_dir=$HOME/t2t/tmp \
+  --game=pong \
+  --num_env_steps=30000 \
+  --agent_policy_path=$AGENT_CKPT_PATH
 ```
+
+Add `--eval` if you want to evaluate the agent against the environment instead
+of generating trajectories for training the world model.
diff --git a/tensor2tensor/rl/datagen_with_agent.py b/tensor2tensor/rl/datagen_with_agent.py
new file mode 100644
index 000000000..1cdbc9d5d
--- /dev/null
+++ b/tensor2tensor/rl/datagen_with_agent.py
@@ -0,0 +1,68 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Generate trajectories to disk with random or ckpt agent.
+
+TODO: Usage
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators import gym_problems_specs
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("data_dir", "", "Data directory.")
+flags.DEFINE_string("tmp_dir", "/tmp/t2t_datagen",
+                    "Temporary storage directory.")
+flags.DEFINE_string("game", None, "Atari game to generate data for.")
+flags.DEFINE_integer("num_env_steps", 5000, "Number of steps to roll out.")
+flags.DEFINE_boolean("eval", False, "Whether to run in eval mode.")
+
+
+def main(_):
+
+  tf.gfile.MakeDirs(FLAGS.data_dir)
+  tf.gfile.MakeDirs(FLAGS.tmp_dir)
+
+  # Create problem if not already defined
+  problem_name = "gym_discrete_problem_with_agent_on_%s" % FLAGS.game
+  if problem_name not in registry.list_problems():
+    gym_problems_specs.create_problems_for_game(FLAGS.game)
+
+  # Generate
+  tf.logging.info("Running %s environment for %d steps for trajectories.",
+                  FLAGS.game, FLAGS.num_env_steps)
+  problem = registry.problem(problem_name)
+  problem.settable_num_steps = FLAGS.num_env_steps
+  problem.settable_eval_phase = FLAGS.eval
+  problem.generate_data(FLAGS.data_dir, FLAGS.tmp_dir)
+
+  # Log stats
+  if problem.statistics.number_of_dones:
+    mean_reward = (problem.statistics.sum_of_rewards /
+                   problem.statistics.number_of_dones)
+    tf.logging.info("Mean reward: %.2f, Num dones: %d",
+                    mean_reward,
+                    problem.statistics.number_of_dones)
+
+
+if __name__ == "__main__":
+  tf.app.run(main)
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 3d908495a..4646ed6de 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -43,7 +43,7 @@ def define_train(hparams):
 def train(hparams, event_dir=None, model_dir=None,
           restore_agent=True, epoch=0):
   """Train."""
-  with tf.name_scope("rl_train"):
+  with tf.Graph().as_default():
     train_summary_op, _, initialization = define_train(hparams)
     if event_dir:
       summary_writer = tf.summary.FileWriter(
@@ -68,7 +68,8 @@ def train(hparams, event_dir=None, model_dir=None,
       initialization(sess)
       if env_model_loader:
         trainer_lib.restore_checkpoint(
-            hparams.world_model_dir, env_model_loader, sess, must_restore=True)
+            hparams.world_model_dir, env_model_loader, sess,
+            must_restore=True)
       start_step = 0
       if model_saver and restore_agent:
         start_step = trainer_lib.restore_checkpoint(
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index 963b86dec..fbcfe7b14 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -17,11 +17,9 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.data_generators.gym_problems import standard_atari_env_spec
-from tensor2tensor.models.research.rl import feed_forward_cnn_small_categorical_fun
-from tensor2tensor.models.research.rl import simple_gym_spec
+from tensor2tensor.data_generators import gym_problems
+from tensor2tensor.models.research import rl as rl_models
 from tensor2tensor.rl import rl_trainer_lib
-from tensor2tensor.utils import registry  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -37,7 +35,8 @@ def test_no_crash_pendulum(self):
         "ppo_continuous_action_base",
         TrainTest.test_config)
 
-    hparams.add_hparam("environment_spec", simple_gym_spec("Pendulum-v0"))
+    hparams.add_hparam(
+        "environment_spec", rl_models.simple_gym_spec("Pendulum-v0"))
     rl_trainer_lib.train(hparams)
 
   def test_no_crash_cartpole(self):
@@ -45,35 +44,36 @@ def test_no_crash_cartpole(self):
         "ppo_discrete_action_base",
         TrainTest.test_config)
 
-    hparams.add_hparam("environment_spec",
-                       standard_atari_env_spec("CartPole-v0"))
+    hparams.add_hparam(
+        "environment_spec", gym_problems.standard_atari_env_spec("CartPole-v0"))
     rl_trainer_lib.train(hparams)
 
   # This test should successfully train pong.
   # It should get train mean_score around 0 after 200 epoch
   # By default the test is disabled to avoid travis timeouts
   def test_train_pong(self):
-    hparams = tf.contrib.training.\
-      HParams(epochs_num=300,
-              eval_every_epochs=10,
-              num_agents=10,
-              optimization_epochs=3,
-              epoch_length=200,
-              entropy_loss_coef=0.003,
-              learning_rate=8e-05,
-              optimizer="Adam",
-              policy_network=feed_forward_cnn_small_categorical_fun,
-              gae_lambda=0.985,
-              num_eval_agents=1,
-              max_gradients_norm=0.5,
-              gae_gamma=0.985,
-              optimization_batch_size=4,
-              clipping_coef=0.2,
-              value_loss_coef=1,
-              save_models_every_epochs=False)
+    hparams = tf.contrib.training.HParams(
+        epochs_num=300,
+        eval_every_epochs=10,
+        num_agents=10,
+        optimization_epochs=3,
+        epoch_length=200,
+        entropy_loss_coef=0.003,
+        learning_rate=8e-05,
+        optimizer="Adam",
+        policy_network=rl_models.feed_forward_cnn_small_categorical_fun,
+        gae_lambda=0.985,
+        num_eval_agents=1,
+        max_gradients_norm=0.5,
+        gae_gamma=0.985,
+        optimization_batch_size=4,
+        clipping_coef=0.2,
+        value_loss_coef=1,
+        save_models_every_epochs=False)
 
-    hparams.add_hparam("environment_spec",
-                       standard_atari_env_spec("PongNoFrameskip-v4"))
+    hparams.add_hparam(
+        "environment_spec",
+        gym_problems.standard_atari_env_spec("PongNoFrameskip-v4"))
     # TODO(lukaszkaiser): enable tests with Atari.
     # rl_trainer_lib.train(hparams)
 
diff --git a/tensor2tensor/rl/model_rl_experiment.py b/tensor2tensor/rl/trainer_model_based.py
similarity index 78%
rename from tensor2tensor/rl/model_rl_experiment.py
rename to tensor2tensor/rl/trainer_model_based.py
index 2d53c79d4..a5c9c6516 100644
--- a/tensor2tensor/rl/model_rl_experiment.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -16,7 +16,7 @@
 
 Example invocation:
 
-python -m tensor2tensor.rl.model_rl_experiment \
+python -m tensor2tensor.rl.trainer_model_based \
     --output_dir=$HOME/t2t/rl_v1 \
     --loop_hparams_set=rl_modelrl_base \
     --loop_hparams='true_env_generator_num_steps=10000,epochs=3'
@@ -32,6 +32,8 @@
 import os
 import time
 
+import six
+
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import gym_problems_specs
@@ -51,6 +53,9 @@
 flags.DEFINE_string("loop_hparams", "", "Overrides for overall loop HParams.")
 
 
+HP_SCOPES = ["loop", "model", "ppo"]
+
+
 def setup_directories(base_dir, subdirs):
   base_dir = os.path.expanduser(base_dir)
   tf.gfile.MakeDirs(base_dir)
@@ -359,8 +364,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     if problem_name not in registry.list_problems():
       tf.logging.info("Game Problem %s not found; dynamically registering",
                       problem_name)
-      gym_problems_specs.dynamically_create_gym_clipped_reward_problem(
-          hparams.game)
+      gym_problems_specs.create_problems_for_game(hparams.game)
 
   # Autoencoder model dir
   autoencoder_model_dir = directories.get("autoencoder")
@@ -570,19 +574,28 @@ def rl_modelrl_short():
   return hparams
 
 
+@registry.register_hparams
+def rl_modelrl_model_only():
+  hp = rl_modelrl_base()
+  hp.epochs = 1
+  hp.ppo_epochs_num = 0
+  return hp
+
+
 @registry.register_hparams
 def rl_modelrl_tiny():
   """Tiny set for testing."""
   return rl_modelrl_base().override_from_dict(
       tf.contrib.training.HParams(
           epochs=2,
-          true_env_generator_num_steps=100,
-          simulated_env_generator_num_steps=100,
+          true_env_generator_num_steps=64,
+          simulated_env_generator_num_steps=64,
           model_train_steps=2,
           ppo_epochs_num=2,
           ppo_time_limit=5,
           ppo_epoch_length=5,
           ppo_num_agents=2,
+          generative_model_params="next_frame_tiny",
       ).values())
 
 
@@ -723,125 +736,208 @@ def rl_modelrl_ae_tiny():
 
 
 @registry.register_hparams
-def rl_modelrl_breakout_tiny():
-  """Tiny set for testing Breakout."""
-  hparams = rl_modelrl_tiny()
-  hparams.game = "wrapped_breakout"
-  return hparams
+def rl_modelrl_tiny_simulation_deterministic_starts():
+  hp = rl_modelrl_tiny()
+  hp.simulation_random_starts = False
+  return hp
 
 
-@registry.register_hparams
-def rl_modelrl_breakout_base():
-  """Base set for testing Breakout."""
-  hparams = rl_modelrl_base()
-  hparams.game = "wrapped_breakout"
-  return hparams
+# RangedHParams for tuning
+# ==============================================================================
+# Note that the items here must be scoped with one of
+# HP_SCOPES={loop, model, ppo}, which set hyperparameters for the top-level
+# hparams, hp.generative_model_params, and hp.ppo_params, respectively.
+@registry.register_ranged_hparams
+def rl_modelrl_grid(rhp):
+  """Grid over games and frames, and 5 runs each for variance."""
+  rhp.set_categorical("loop.game",
+                      ["breakout", "wrapped_long_pong", "freeway"])
 
+  # 100k, 50k, 25k frames
+  base = 36666
+  medium = base // 2
+  small = medium // 2
+  rhp.set_discrete("loop.true_env_generator_num_steps", [base, medium, small])
 
-@registry.register_hparams
-def rl_modelrl_breakout_ae_base():
-  """Base set for testing Breakout with an autoencoder."""
-  hparams = rl_modelrl_ae_base()
-  hparams.game = "wrapped_breakout"
-  return hparams
+  # Dummy parameter to get 5 runs for each configuration
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
 
 
-@registry.register_hparams
-def rl_modelrl_breakout_medium():
-  """Medium set for testing Breakout."""
-  hparams = rl_modelrl_medium()
-  hparams.game = "wrapped_breakout"
-  return hparams
+@registry.register_ranged_hparams
+def rl_modelrl_variance(rhp):
+  # Dummy parameter to get 5 runs for each configuration
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game",
+                      ["breakout", "wrapped_long_pong", "freeway"])
 
 
-@registry.register_hparams
-def rl_modelrl_breakout_ae_medium():
-  """Medium set for testing Breakout with an autoencoder."""
-  hparams = rl_modelrl_ae_medium()
-  hparams.game = "wrapped_breakout"
-  return hparams
+@registry.register_ranged_hparams
+def rl_modelrl_variance_nogame(rhp):
+  # Dummy parameter to get 5 runs for each configuration
+  rhp.set_discrete("model.moe_loss_coef", list(range(500)))
 
 
-@registry.register_hparams
-def rl_modelrl_breakout_short():
-  """Short set for testing Breakout."""
-  hparams = rl_modelrl_short()
-  hparams.game = "wrapped_breakout"
-  return hparams
+@registry.register_ranged_hparams
+def rl_modelrl_all_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_GAMES)
 
 
-@registry.register_hparams
-def rl_modelrl_breakout_ae_short():
-  """Short set for testing Breakout with an autoencoder."""
-  hparams = rl_modelrl_ae_short()
-  hparams.game = "wrapped_breakout"
-  return hparams
+@registry.register_ranged_hparams
+def rl_modelrl_whitelisted_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
 
 
-@registry.register_hparams
-def rl_modelrl_freeway_tiny():
-  """Tiny set for testing Freeway."""
-  hparams = rl_modelrl_tiny()
-  hparams.game = "freeway"
-  return hparams
+@registry.register_ranged_hparams
+def rl_modelrl_ae_variance(rhp):
+  # Dummy parameter to get 5 runs for each configuration
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game",
+                      ["breakout", "wrapped_long_pong", "freeway"])
+  # 100k, 25k frames
+  base = 36666
+  small = base // 4
+  rhp.set_discrete("loop.true_env_generator_num_steps", [base, small])
 
 
-@registry.register_hparams
-def rl_modelrl_freeway_base():
-  """Base set for testing Freeway."""
-  hparams = rl_modelrl_base()
-  hparams.game = "freeway"
-  return hparams
+@registry.register_ranged_hparams
+def rl_modelrl_ppolr_game(rhp):
+  rhp.set_categorical("loop.game",
+                      ["breakout", "wrapped_long_pong", "freeway"])
+  base_lr = 2e-4
+  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
 
 
-@registry.register_hparams
-def rl_modelrl_freeway_ae_base():
-  """Base set for testing Freeway with an autoencoder."""
-  hparams = rl_modelrl_ae_base()
-  hparams.game = "freeway"
-  return hparams
+@registry.register_ranged_hparams
+def rl_modelrl_ppolr(rhp):
+  base_lr = 2e-4
+  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
 
 
-@registry.register_hparams
-def rl_modelrl_freeway_medium():
-  """Medium set for testing Freeway."""
-  hparams = rl_modelrl_medium()
-  hparams.game = "freeway"
-  return hparams
+@registry.register_ranged_hparams
+def rl_modelrl_ae_ppo_lr(rhp):
+  rhp.set_categorical("loop.game",
+                      ["breakout", "wrapped_long_pong", "freeway"])
+  base_lr = 2e-4
+  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
 
 
-@registry.register_hparams
-def rl_modelrl_freeway_ae_medium():
-  """Medium set for testing Freeway with an autoencoder."""
-  hparams = rl_modelrl_ae_medium()
-  hparams.game = "freeway"
-  return hparams
+@registry.register_ranged_hparams
+def rl_modelrl_dropout_range(rhp):
+  rhp.set_float("model.dropout", 0.2, 0.4)
 
 
-@registry.register_hparams
-def rl_modelrl_freeway_short():
-  """Short set for testing Freeway."""
-  hparams = rl_modelrl_freeway_medium()
-  hparams.true_env_generator_num_steps //= 5
-  hparams.model_train_steps //= 2
-  hparams.ppo_epochs_num //= 2
-  hparams.intrinsic_reward_scale = 0.1
-  return hparams
+@registry.register_ranged_hparams
+def rl_modelrl_intrinsic_reward_scale(rhp):
+  rhp.set_float("loop.intrinsic_reward_scale", 0.01, 10.)
 
 
-@registry.register_hparams
-def rl_modelrl_freeway_ae_short():
-  """Short set for testing Freeway with an autoencoder."""
-  hparams = rl_modelrl_ae_short()
-  hparams.game = "freeway"
-  return hparams
+@registry.register_ranged_hparams
+def rl_modelrl_l1l2cutoff_range(rhp):
+  """Loss and loss-cutoff tuning grid."""
+  rhp.set_float("model.video_modality_loss_cutoff", 1.4, 3.4)
 
 
-@registry.register_hparams
-def rl_modelrl_tiny_simulation_deterministic_starts():
-  hp = rl_modelrl_tiny()
-  hp.simulation_random_starts = False
-  return hp
+@registry.register_ranged_hparams
+def rl_modelrl_xentcutoff_range(rhp):
+  """Cross entropy cutoff tuning grid."""
+  rhp.set_float("model.video_modality_loss_cutoff", 0.01, 0.05)
+
+
+@registry.register_ranged_hparams
+def rl_modelrl_pixel_noise(rhp):
+  """Input pixel noise tuning grid."""
+  rhp.set_categorical("loop.generative_model_params",
+                      ["next_frame_pixel_noise"])
+  rhp.set_discrete("model.video_modality_input_noise",
+                   [0.0025 * i for i in range(200)])
+
+
+@registry.register_ranged_hparams
+def rl_modelrl_dummy_range(rhp):
+  """Dummy tuning grid just to get the variance."""
+  rhp.set_float("model.moe_loss_coef", 0.01, 0.02)
+
+
+def merge_unscoped_hparams(scopes_and_hparams):
+  """Merge multiple HParams into one with scopes."""
+  merged_values = {}
+  for (scope, hparams) in scopes_and_hparams:
+    for key, value in six.iteritems(hparams.values()):
+      scoped_key = "%s.%s" % (scope, key)
+      merged_values[scoped_key] = value
+
+  return tf.contrib.training.HParams(**merged_values)
+
+
+def split_scoped_hparams(scopes, merged_hparams):
+  """Split single HParams with scoped keys into multiple."""
+  split_values = dict([(scope, dict()) for scope in scopes])
+  merged_values = merged_hparams.values()
+  for scoped_key, value in six.iteritems(merged_values):
+    scope = scoped_key.split(".")[0]
+    key = scoped_key[len(scope) + 1:]
+    split_values[scope][key] = value
+
+  return [
+      tf.contrib.training.HParams(**split_values[scope]) for scope in scopes
+  ]
+
+
+def training_loop_hparams_from_scoped_overrides(scoped_overrides, trial_id):
+  """Create HParams suitable for training loop from scoped HParams.
+
+  Args:
+    scoped_overrides: HParams, with keys all scoped by one of HP_SCOPES. These
+      parameters are overrides for the base HParams created by
+      create_loop_hparams.
+    trial_id: str, trial identifier. This is used to register unique HParams
+      names for the underlying model and ppo HParams.
+
+  Returns:
+    HParams suitable for passing to training_loop.
+  """
+  trial_hp_overrides = scoped_overrides.values()
+
+  # Create loop, model, and ppo base HParams
+  loop_hp = create_loop_hparams()
+  model_hp_name = trial_hp_overrides.get(
+      "loop.generative_model_params", loop_hp.generative_model_params)
+  model_hp = registry.hparams(model_hp_name).parse(FLAGS.hparams)
+  ppo_params_name = trial_hp_overrides.get(
+      "loop.ppo_params", loop_hp.ppo_params)
+  ppo_hp = registry.hparams(ppo_params_name)
+
+  # Merge them and then override with the scoped overrides
+  combined_hp = merge_unscoped_hparams(
+      zip(HP_SCOPES, [loop_hp, model_hp, ppo_hp]))
+  combined_hp.override_from_dict(trial_hp_overrides)
+
+  # Split out the component hparams
+  loop_hp, model_hp, ppo_hp = (
+      split_scoped_hparams(HP_SCOPES, combined_hp))
+
+  # Dynamic register the model hp and set the new name in loop_hp
+  model_hp_name = "model_hp_%s" % str(trial_id)
+  dynamic_register_hparams(model_hp_name, model_hp)
+  loop_hp.generative_model_params = model_hp_name
+
+  # Dynamic register the PPO hp and set the new name in loop_hp
+  ppo_hp_name = "ppo_hp_%s" % str(trial_id)
+  dynamic_register_hparams(ppo_hp_name, ppo_hp)
+  loop_hp.ppo_params = ppo_hp_name
+
+  return loop_hp
+
+
+def dynamic_register_hparams(name, hparams):
+
+  @registry.register_hparams(name)
+  def new_hparams_set():
+    return tf.contrib.training.HParams(**hparams.values())
+
+  return new_hparams_set
 
 
 def create_loop_hparams():
diff --git a/tensor2tensor/rl/model_rl_experiment_stochastic_test.py b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
similarity index 83%
rename from tensor2tensor/rl/model_rl_experiment_stochastic_test.py
rename to tensor2tensor/rl/trainer_model_based_stochastic_test.py
index 51f207364..d55951256 100644
--- a/tensor2tensor/rl/model_rl_experiment_stochastic_test.py
+++ b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
@@ -12,12 +12,12 @@
 # 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.
-"""Tiny run of model_rl_experiment. Smoke test."""
+"""Tiny run of trainer_model_based with stochastic model. Smoke test."""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.rl import model_rl_experiment
+from tensor2tensor.rl import trainer_model_based
 
 import tensorflow as tf
 
@@ -29,9 +29,8 @@ class ModelRLExperimentStochasticTest(tf.test.TestCase):
   def test_stochastic(self):
     FLAGS.output_dir = tf.test.get_temp_dir()
     FLAGS.loop_hparams_set = "rl_modelrl_tiny_stochastic"
-    FLAGS.loop_hparams = "generative_model_params=next_frame_stochastic_tiny"
     FLAGS.schedule = "train"  # skip evaluation for world model training
-    model_rl_experiment.main(None)
+    trainer_model_based.main(None)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/rl/model_rl_experiment_test.py b/tensor2tensor/rl/trainer_model_based_test.py
similarity index 83%
rename from tensor2tensor/rl/model_rl_experiment_test.py
rename to tensor2tensor/rl/trainer_model_based_test.py
index 90388d161..48e2f40ea 100644
--- a/tensor2tensor/rl/model_rl_experiment_test.py
+++ b/tensor2tensor/rl/trainer_model_based_test.py
@@ -12,12 +12,12 @@
 # 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.
-"""Tiny run of model_rl_experiment. Smoke test."""
+"""Tiny run of trainer_model_based. Smoke test."""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.rl import model_rl_experiment
+from tensor2tensor.rl import trainer_model_based
 
 import tensorflow as tf
 
@@ -29,9 +29,8 @@ class ModelRLExperimentTest(tf.test.TestCase):
   def test_basic(self):
     FLAGS.output_dir = tf.test.get_temp_dir()
     FLAGS.loop_hparams_set = "rl_modelrl_tiny"
-    FLAGS.loop_hparams = "generative_model_params=next_frame_tiny"
     FLAGS.schedule = "train"  # skip evaluation for world model training
-    model_rl_experiment.main(None)
+    trainer_model_based.main(None)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/rl/t2t_rl_trainer.py b/tensor2tensor/rl/trainer_model_free.py
similarity index 100%
rename from tensor2tensor/rl/t2t_rl_trainer.py
rename to tensor2tensor/rl/trainer_model_free.py

From 8ff6ec409e0206edcb26ae9092ea174aa8fd8d0c Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 15 Aug 2018 14:11:41 -0700
Subject: [PATCH 0590/2720] internal

PiperOrigin-RevId: 208880139
---
 tensor2tensor/layers/common_video.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 8b41d221b..f04938108 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -267,3 +267,4 @@ def tile_and_concat(image, latent, concat_latent=True):
   latent = tf.pad(latent, [[0, 0], [pad // 2, pad // 2], [0, 0], [0, 0]])
   return tf.concat([image, latent], axis=-1)
 
+

From 64dc8cfca1f649d4995fe12292002b6c79da65fc Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Wed, 15 Aug 2018 16:57:45 -0700
Subject: [PATCH 0591/2720] Fix for dataset mixing.

PiperOrigin-RevId: 208908276
---
 .../data_generators/multi_problem.py          | 32 +++++++++++--------
 1 file changed, 18 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index e121c9199..4c4e5d70d 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -193,6 +193,9 @@ def get_const_sched_prob():
       def mix_data(example):
         """Function to mix the different datasets according to a schedule."""
         del example
+        # This block computes the probability of mixing the primary task with
+        # the secondary tasks. 0 = only the primary task, 1 = only the secondary
+        # tasks.
         if hparams.multiproblem_mixing_schedule == MixingSchedule.EXPONENTIAL:
           prob = get_exp_sched_prob()
         elif hparams.multiproblem_mixing_schedule == MixingSchedule.CONSTANT:
@@ -203,8 +206,10 @@ def mix_data(example):
         tf.logging.info("Using the %s schedule to "
                         "train the MultiProblem." % str(
                             hparams.multiproblem_mixing_schedule))
+        tf.logging.info("Schedule mixing threshold "
+                        "%.2f" % hparams.multiproblem_schedule_threshold)
 
-        def sample_task(curr_task, num_tasks_left):
+        def sample_task(curr_task, num_tasks_left, randnum):
           """A recursive function to sample a task.
 
           This function treats the probability as the threshold for the primary
@@ -214,6 +219,7 @@ def sample_task(curr_task, num_tasks_left):
           Args:
             curr_task: The index of the task being considered for sampling.
             num_tasks_left: Number of tasks remaining to possibly sample from.
+            randnum: The random number used to select the dataset.
 
           Returns:
             A Tensor representing an example from the task that was sampled
@@ -222,23 +228,21 @@ def sample_task(curr_task, num_tasks_left):
 
           if num_tasks_left == 0:
             return get_next_from_dataset(dataset_iterators[curr_task])
-          elif curr_task == 0:
-            # primary task
-            return tf.cond(
-                tf.greater(tf.random_uniform([]), prob),
-                lambda d=dataset_iterators[0]: get_next_from_dataset(d),
-                lambda c=curr_task+1, n=num_tasks_left-1: sample_task(c, n)
-            )
-          # divide the probability mass across all the secondary tasks equally.
-          new_prob = prob - curr_task * prob / (len(self.task_list)-1)
+
+          # When curr_task is 0, the primary task, the new prob is the same as
+          # the original probability. `tf.greater` indicates that the primary
+          # task receives (1-prob) of the probability mass.
+          # Otherwise, `prob` is divided equally amongst all the secondary
+          # tasks.
+          new_prob = prob - (curr_task * prob / (len(self.task_list)-1))
           return tf.cond(
-              tf.greater(tf.random_uniform([]), new_prob),
-              lambda d=dataset_iterators[curr_task]: get_next_from_dataset(d),
-              lambda c=curr_task+1, n=num_tasks_left-1: sample_task(c, n)
+              tf.greater(randnum, new_prob),
+              lambda: get_next_from_dataset(dataset_iterators[curr_task]),
+              lambda: sample_task(curr_task+1, num_tasks_left-1, randnum)
           )
 
         return tf.data.Dataset.from_tensors(
-            sample_task(0, len(self.task_list)-1))
+            sample_task(0, len(self.task_list)-1, tf.random_uniform([])))
 
       single_mtl_dataset = tf.data.Dataset.from_tensors(tf.zeros([1])).repeat()
       single_mtl_dataset = single_mtl_dataset.flat_map(mix_data)

From 77ce1c55e02b21665ed33c484df9124236483c5b Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 15 Aug 2018 16:58:57 -0700
Subject: [PATCH 0592/2720] Disable unused-argument pylint check.

PiperOrigin-RevId: 208908419
---
 pylintrc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pylintrc b/pylintrc
index 47e9c7ab3..ab45e0220 100644
--- a/pylintrc
+++ b/pylintrc
@@ -37,7 +37,7 @@ msg-template={msg_id}:{line:3} {obj}: {msg} [{symbol}]
 enable=indexing-exception,old-raise-syntax
 
 # List of checkers and warnings to disable.
-disable=design,similarities,no-self-use,attribute-defined-outside-init,locally-disabled,star-args,pointless-except,bad-option-value,global-statement,fixme,suppressed-message,useless-suppression,locally-enabled,file-ignored,multiple-imports,c-extension-no-member,trailing-newlines,unsubscriptable-object,misplaced-comparison-constant,no-member,abstract-method,no-else-return,missing-docstring,wrong-import-order,protected-access,inconsistent-return-statements,invalid-unary-operand-type,import-error,no-name-in-module,arguments-differ,not-context-manager
+disable=design,similarities,no-self-use,attribute-defined-outside-init,locally-disabled,star-args,pointless-except,bad-option-value,global-statement,fixme,suppressed-message,useless-suppression,locally-enabled,file-ignored,multiple-imports,c-extension-no-member,trailing-newlines,unsubscriptable-object,misplaced-comparison-constant,no-member,abstract-method,no-else-return,missing-docstring,wrong-import-order,protected-access,inconsistent-return-statements,invalid-unary-operand-type,import-error,no-name-in-module,arguments-differ,not-context-manager,unused-argument
 
 [BASIC]
 

From 6c8bec5d5687174d617bf8212c1f7647e5b8728f Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 15 Aug 2018 17:14:03 -0700
Subject: [PATCH 0593/2720] adding multiplicative actions to sv2p.

PiperOrigin-RevId: 208910844
---
 tensor2tensor/models/research/next_frame_sv2p.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index df7ff0c60..4cae407de 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -231,9 +231,19 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
 
     # Pass in action if exists.
     if action is not None:
-      emb_action = common_video.encode_to_shape(
-          action, enc2.get_shape(), "action_enc")
-      enc2 = tf.concat(values=[enc2, emb_action], axis=3)
+      if self.hparams.concatenate_actions:
+        emb_action = common_video.encode_to_shape(
+            action, enc2.get_shape(), "action_enc")
+        enc2 = tf.concat(values=[enc2, emb_action], axis=3)
+      else:
+        action_shape = common_layers.shape_list(action)
+        enc2_shape = common_layers.shape_list(enc2)
+        filters = enc2_shape[-1]
+        action_reshaped = tf.reshape(action, [-1, 1, 1, action_shape[-1]])
+        action_mask = tf.layers.dense(action_reshaped, filters)
+        zeros_mask = tf.zeros(enc2_shape, dtype=tf.float32)
+        action_broad = action_mask + zeros_mask
+        enc2 *= action_broad
 
     # Pass in reward if exists.
     if input_reward is not None:

From a705063fc48dee34740a5fa1808c4fe5b495777f Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 15 Aug 2018 17:32:31 -0700
Subject: [PATCH 0594/2720] Deprecate AttentionLM in favor of Transformer

PiperOrigin-RevId: 208913126
---
 tensor2tensor/models/research/attention_lm.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tensor2tensor/models/research/attention_lm.py b/tensor2tensor/models/research/attention_lm.py
index ee3d9b9f2..d3de21044 100644
--- a/tensor2tensor/models/research/attention_lm.py
+++ b/tensor2tensor/models/research/attention_lm.py
@@ -14,6 +14,8 @@
 # limitations under the License.
 """Self-attention based language model.
 
+DEPRECATED. Use Transformer which supports running the decoder only.
+
 Like transformer.py, but no encoder
 
 decoder: [Self-Attention, Feed-forward] x n
@@ -34,6 +36,10 @@
 import tensorflow as tf
 
 
+@tf.contrib.framework.deprecated(
+    "2018-09-15",
+    "Use Transformer, which supports decoder-only mode when "
+    "Transformer.has_input=False.")
 @registry.register_model
 class AttentionLM(t2t_model.T2TModel):
   """Attention net.  See file docstring."""

From 921e5196b0246b0bbfe8a763ceb8b7daaba7be86 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 15 Aug 2018 19:52:26 -0700
Subject: [PATCH 0595/2720] internal merge of PR #996

PiperOrigin-RevId: 208925627
---
 setup.py                                      |  1 -
 .../data_generators/paraphrase_ms_coco.py     | 37 ++++++++++---------
 .../paraphrase_ms_coco_test.py                |  4 +-
 3 files changed, 22 insertions(+), 20 deletions(-)

diff --git a/setup.py b/setup.py
index 2096d4c21..f868adcfd 100644
--- a/setup.py
+++ b/setup.py
@@ -40,7 +40,6 @@
         'gunicorn',
         'gym',
         'h5py',
-        'mock',
         'numpy',
         'oauth2client',
         'requests',
diff --git a/tensor2tensor/data_generators/paraphrase_ms_coco.py b/tensor2tensor/data_generators/paraphrase_ms_coco.py
index f8c2b0449..ed5c4dd02 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco.py
@@ -18,15 +18,15 @@
 from __future__ import division
 from __future__ import print_function
 
-import json
-import io
 import collections
+import io
+import json
 import os
 import zipfile
 
 from tensor2tensor.data_generators import generator_utils
-from tensor2tensor.data_generators import text_problems
 from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -39,8 +39,9 @@
 
 
 def create_combination(list_of_sentences):
-  """ Generates all possible pair combinations for
-  the input list of sentences, for example:
+  """Generates all possible pair combinations for the input list of sentences.
+
+  For example:
 
   input = ["paraphrase1", "paraphrase2", "paraphrase3"]
 
@@ -55,7 +56,7 @@ def create_combination(list_of_sentences):
   """
   num_sentences = len(list_of_sentences) - 1
   combinations = []
-  for i, sentence in enumerate(list_of_sentences):
+  for i, _ in enumerate(list_of_sentences):
     if i == num_sentences:
       break
     num_pairs = num_sentences - i
@@ -66,10 +67,11 @@ def create_combination(list_of_sentences):
 
 
 class ParaphraseGenerationProblem(text_problems.Text2TextProblem):
+  """Paraphrase problem."""
 
   @property
   def bidirectional(self):
-    """If set to true, generates data in the following way:
+    """If set to true, generates data in the following way.
 
     sentence1 -> sentence2
     sentence2 -> sentence1
@@ -87,12 +89,13 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
         caption_pairs += [(caption2, caption1)]
       for caption_pair in caption_pairs:
         yield {
-          'inputs': caption_pair[0],
-          'targets': caption_pair[1]
+            "inputs": caption_pair[0],
+            "targets": caption_pair[1]
         }
 
 
 class ParaphraseGenerationMsCocoProblem(ParaphraseGenerationProblem):
+  """Paraphrase problem."""
 
   @property
   def is_generate_per_split(self):
@@ -101,11 +104,11 @@ def is_generate_per_split(self):
   @property
   def dataset_splits(self):
     return [{
-      "split": problem.DatasetSplit.TRAIN,
-      "shards": 10,
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 10,
     }, {
-      "split": problem.DatasetSplit.EVAL,
-      "shards": 1,
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
     }]
 
   @property
@@ -158,7 +161,7 @@ def _get_captions(self, ms_coco_path):
 
 @registry.register_problem
 class ParaphraseGenerationMsCocoProblem2d(
-  ParaphraseGenerationMsCocoProblem):
+    ParaphraseGenerationMsCocoProblem):
 
   @property
   def bidirectional(self):
@@ -167,7 +170,7 @@ def bidirectional(self):
 
 @registry.register_problem
 class ParaphraseGenerationMsCocoProblem1d(
-  ParaphraseGenerationMsCocoProblem):
+    ParaphraseGenerationMsCocoProblem):
 
   @property
   def bidirectional(self):
@@ -176,7 +179,7 @@ def bidirectional(self):
 
 @registry.register_problem
 class ParaphraseGenerationMsCocoProblem2dCharacters(
-  ParaphraseGenerationMsCocoProblem2d):
+    ParaphraseGenerationMsCocoProblem2d):
 
   @property
   def vocab_type(self):
@@ -185,7 +188,7 @@ def vocab_type(self):
 
 @registry.register_problem
 class ParaphraseGenerationMsCocoProblem1dCharacters(
-  ParaphraseGenerationMsCocoProblem1d):
+    ParaphraseGenerationMsCocoProblem1d):
 
   @property
   def vocab_type(self):
diff --git a/tensor2tensor/data_generators/paraphrase_ms_coco_test.py b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
index 046e42dea..439e2ec6c 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
@@ -28,8 +28,8 @@
 class ParaphraseGenerationProblemTest(tf.test.TestCase):
 
   def testCombinationPairs(self):
-    inputs = ['A', 'B', 'C']
-    expected_combination = [('A', 'B'), ('A', 'C'), ('B', 'C')]
+    inputs = ["A", "B", "C"]
+    expected_combination = [("A", "B"), ("A", "C"), ("B", "C")]
     actual_combination = paraphrase_ms_coco.create_combination(inputs)
     self.assertEqual(actual_combination, expected_combination)
 

From adf79aea80b38b9224f35d9f5dec4a7fe2d04537 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 15 Aug 2018 20:15:33 -0700
Subject: [PATCH 0596/2720] cleanup the visualization and adding internal gifs.

PiperOrigin-RevId: 208927281
---
 .../models/research/next_frame_sv2p.py        | 29 ++++++++++---------
 1 file changed, 16 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 4cae407de..02d6e7b26 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -50,6 +50,19 @@ def tinyify(self, array):
       return [1 for _ in array]
     return array
 
+  def visualize_predictions(self, real_frames, gen_frames):
+    def concat_on_y_axis(x):
+      x = tf.unstack(x, axis=1)
+      x = tf.concat(x, axis=1)
+      return x
+
+    frames_gd = common_video.swap_time_and_batch_axes(real_frames)
+    frames_pd = common_video.swap_time_and_batch_axes(gen_frames)
+    frames_gd = concat_on_y_axis(frames_gd)
+    frames_pd = concat_on_y_axis(frames_pd)
+    side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
+    tf.summary.image("full_video", side_by_side_video)
+
   def get_gaussian_latent(self, latent_mean, latent_std):
     latent = tf.random_normal(tf.shape(latent_mean), 0, 1, dtype=tf.float32)
     latent = latent_mean + tf.exp(latent_std / 2.0) * latent
@@ -575,8 +588,7 @@ def get_extra_loss(self, latent_means=None, latent_stds=None,
 
   def body(self, features):
     hparams = self.hparams
-    input_shape = common_layers.shape_list(features["inputs"])
-    batch_size, _, _, frame_height, frame_channels = input_shape
+    batch_size = common_layers.shape_list(features["inputs"])[0]
 
     # Swap time and batch axes.
     input_frames = common_video.swap_time_and_batch_axes(features["inputs"])
@@ -616,17 +628,8 @@ def body(self, features):
         latent_stds=latent_stds, beta=beta, true_frames=all_frames,
         gen_frames=gen_images)
 
-    # TODO(mbz): clean this up!
-    def fix_video_dims_and_concat_on_x_axis(x):
-      x = tf.transpose(x, [1, 3, 4, 0, 2])
-      x = tf.reshape(x, [batch_size, frame_height, frame_channels, -1])
-      x = tf.transpose(x, [0, 3, 1, 2])
-      return x
-
-    frames_gd = fix_video_dims_and_concat_on_x_axis(all_frames[1:])
-    frames_pd = fix_video_dims_and_concat_on_x_axis(gen_images)
-    side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
-    tf.summary.image("full_video", side_by_side_video)
+    # Visualize predictions in Tensorboard
+    self.visualize_predictions(all_frames[1:], gen_images)
 
     # Ignore the predictions from the input frames.
     # This is NOT the same as original paper/implementation.

From 730201a7d28e2ded76145ae7dd7a663311c86d4a Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 15 Aug 2018 23:10:08 -0700
Subject: [PATCH 0597/2720] Add examples to distributed_training.md, update
 support for async training, and simplify run_std_server codepath

PiperOrigin-RevId: 208939618
---
 docs/distributed_training.md         | 247 ++++++++++++++++++++-------
 tensor2tensor/bin/make_tf_configs.py |  87 ++++++----
 tensor2tensor/bin/t2t_trainer.py     |   7 +
 tensor2tensor/utils/devices.py       |   5 +-
 tensor2tensor/utils/trainer_lib.py   |  26 +--
 5 files changed, 270 insertions(+), 102 deletions(-)

diff --git a/docs/distributed_training.md b/docs/distributed_training.md
index 48ef14a34..1fa00cb7a 100644
--- a/docs/distributed_training.md
+++ b/docs/distributed_training.md
@@ -3,87 +3,218 @@
 The `t2t-trainer` supports both synchronous and asynchronous distributed
 training.
 
+Note that it is almost always more efficient to train on a single machine with
+multiple GPUs/TPUs. Async training is less stable than sync training, and sync
+training is much faster on 1 machine than on multiple. For these reasons, we
+almost always train on single machines with multiple GPUs/TPUs.
+
 T2T uses TensorFlow Estimators and so distributed training is configured with
 the `TF_CONFIG` environment variable that is read by the
 [RunConfig](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/estimator/run_config.py)
-along with a set of flags.
+along with a set of flags that T2T uses to distribute the computation.
+
+## Shared output directory
+
+When using multiple machines, it is necessary that all nodes use the same
+`--output_dir`, which means that it should be set to a Google Cloud Storage
+bucket (`gs://...`) or a directory on a shared network filesystem.
+
+## Utility to produce `TF_CONFIG` and flags
+
+[`t2t-make-tf-configs`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/bin/t2t-make-tf-configs))
+generates the `TF_CONFIG` json strings and the necessary command-line flags for
+the jobs.
+
+Given a set of master and parameter server addresses, the script outputs, for
+each job, a line with the `TF_CONFIG` environment variable and the command-line
+flags necessary for distributed training. For each job, you should invoke the
+`t2t-trainer` with the `TF_CONFIG` value and flags that are output.
+
+## Eval jobs
+
+Eval jobs should set the following flags and do not need the `TF_CONFIG`
+environment variable to be set as the eval jobs run locally and do not
+communicate to the other jobs (the eval jobs read the model checkpoints that the
+trainer writes out):
+
+- `--schedule=continuous_eval_on_train_data` or
+  `--schedule=continuous_eval` (for dev data)
+- `--worker_job='/job:localhost'`
+- `--output_dir=$TRAIN_DIR`
+
+**Note that evaluation does not work distributed.** That is, distributed jobs
+should always use `--schedule=train`.
+
+## Examples
 
-## `TF_CONFIG`
+### Sync training across multiple workers
 
-Both masters and parameter servers must have the `TF_CONFIG` environment
-variable set.
+In this scenario, you wish to do synchronous training across multiple workers.
+Note that it is easier to simply use 1 worker with multiple GPUs and set
+`--worker_gpu=8`, but there may be cases where you may want to have multiple
+machines.
 
-The `TF_CONFIG` environment variable is a json-encoded string with the addresses
-of the masters and parameter servers (in the `'cluster'` key) and the
-identification of the current task (in the `'task'` key).
+You will need 1 `ip:port` for the master and then 1 `ip:port` for each worker.
 
-For example:
+For this example we'll use 2 workers and these addresses:
 
 ```
-cluster = {
-    'ps': ['host1:2222', 'host2:2222'],
-    'master': ['host3:2222', 'host4:2222', 'host5:2222']
-}
-os.environ['TF_CONFIG'] = json.dumps({
-    'cluster': cluster,
-    'task': {'type': 'master', 'index': 1},
-    'environment': 'cloud',
-})
+# Master
+10.0.0.1:5555
+
+# Worker 1
+10.0.0.2:5555
+
+# Worker 2
+10.0.0.3:5555
 ```
 
-## Command-line flags
+Next we generate the `TF_CONFIG` and command-line-flags for each job.
 
-The following T2T command-line flags must also be set on the masters for
-distributed training:
+```
+$ t2t-make-tf-configs --masters='10.0.0.1:5555' --ps='10.0.0.2:5555,10.0.0.3:5555'
+Assuming SYNC distributed training with a single master and 2 workers
+'{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 0, "type": "master"}}'      --master=grpc://10.0.0.1:5555 --ps_replicas=2 --worker_replicas=1 --worker_gpu=0 --worker_id=0 --ps_gpu=1 --sync --schedule=train --worker_job='/job:master'
+'{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 0, "type": "ps"}}'  --schedule=run_std_server
+'{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 1, "type": "ps"}}'  --schedule=run_std_server
+```
 
-- `--master=grpc://$ADDRESS`
-- `--worker_replicas=$NUM_MASTERS`
-- `--worker_gpu=$NUM_GPUS_PER_MASTER`
-- `--worker_id=$MASTER_ID`
-- `--worker_job='/job:master'`
-- `--ps_replicas=$NUM_PS`
-- `--ps_gpu=$NUM_GPUS_PER_PS`
-- `--schedule=train`
-- `--sync`, if you want synchronous training, i.e. for there to be a single
-  master coordinating the work across "ps" jobs. If not set, then each master
-  operates independently while variables are shared on the parameter servers.
+The output here is 1 line per job. Each line contains the `TF_CONFIG` to set
+for that job as well as the command-line flags to set for that job.
 
-Parameter servers only need `--master=grpc://$ADDRESS` and
-`--schedule=run_std_server`.
+It is a bit confusing that the workers are being passed to the `--ps` flag, but
+this is correct. When running in `--sync` mode, the `ps` are actually the
+workers. You can see in the next example below that when `--sync=False`, i.e.
+async mode, that the `ps` are in fact being used as parameter servers.
 
->> Note about `--output_dir`: All the nodes should use the same `--output_dir`.
->> When using multiple machines, `output_dir` should point to a shared
->> filesystem like NFS or an object store like Google Cloud Storage
->> (`gs://...`).
+Here's how we would start each job on their respective machines (the
+commands below assume that you're ssh'd into that job's machine):
 
-## Utility to produce `TF_CONFIG` and flags
+**Master**:
 
-[`t2t-make-tf-configs`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/bin/t2t-make-tf-configs))
-generates the `TF_CONFIG` json strings and the above-mentioned command-line
-flags for the masters and parameter servers.
+```
+$ export TF_CONFIG='{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 0, "type": "master"}}'
+$ t2t-trainer \
+    --master=grpc://10.0.0.1:5555 \
+    --ps_replicas=2 \
+    --worker_replicas=1 \
+    --worker_gpu=0 \
+    --worker_id=0 \
+    --ps_gpu=1 \
+    --sync \
+    --schedule=train \
+    --worker_job='/job:master' \
+    --model=transformer \
+    --hparams_set=transformer_base \
+    --problem=translate_ende_wmt32k
+```
 
-Given a set of master and parameter server addresses, the script outputs, for
-each job, a line with the `TF_CONFIG` environment variable and the command-line
-flags necessary for distributed training. For each job, you should invoke the
-`t2t-trainer` with the `TF_CONFIG` value and flags that are output.
+**Worker 1**:
 
-For example:
+```
+$ export TF_CONFIG='{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 0, "type": "ps"}}'
+$ t2t-trainer --schedule=run_std_server
+```
+
+**Worker 2**:
 
 ```
-TF_CONFIG=$JOB_TF_CONFIG t2t-trainer $JOB_FLAGS --model=transformer ...
+$ export TF_CONFIG='{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 1, "type": "ps"}}'
+$ t2t-trainer --schedule=run_std_server
 ```
 
-Modify the `--worker_gpu` and `--ps_gpu` flags, which specify how many GPUs are
-on each master and ps, respectively, as needed for your machine/cluster setup.
+Note that if you have more than 1 GPU on each worker machine, make sure to
+modify the `--ps_gpu` passed to the master.
 
-## Command-line flags for eval jobs
+### Async training across multiple workers
 
-Eval jobs should set the following flags and do not need the `TF_CONFIG`
-environment variable to be set as the eval jobs run locally and do not
-communicate to the other jobs (the eval jobs read the model checkpoints that the
-trainer writes out):
+In this scenario, you wish to do asynchronous training across multiple workers
+with 1+ shared parameter servers.
 
-- `--schedule=continuous_eval_on_train_data` or
-  `--schedule=continuous_eval` (for test data)
-- `--worker_job='/job:localhost'`
-- `--output_dir=$TRAIN_DIR`
+Note that async training is usually less stable than sync training and for that
+reason we almost always prefer sync training, but there may be cases where you
+want to do async distributed training.
+
+For this example we'll use 2 workers and 2 parameter servers:
+
+```
+# Worker 1
+10.0.0.1:5555
+
+# Worker 2
+10.0.0.2:5555
+
+# PS 1
+10.0.0.3:5555
+
+# PS 2
+10.0.0.4:5555
+```
+
+Next we generate the `TF_CONFIG` and command-line-flags for each job.
+
+```
+$ t2t-make-tf-configs --masters='10.0.0.1:5555,10.0.0.2:5555' --ps='10.0.0.3:5555,10.0.0.4:5555'
+Assuming ASYNC distributed training with 2 workers and 2 parameter servers
+'{"task": {"index": 0, "type": "chief"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}' --master=grpc://10.0.0.1:5555 --ps_replicas=2 --worker_replicas=2 --worker_gpu=1 --worker_id=0 --ps_gpu=0  --schedule=train --worker_job='/job:chief'
+'{"task": {"index": 0, "type": "worker"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'        --master=grpc://10.0.0.2:5555 --ps_replicas=2 --worker_replicas=2 --worker_gpu=1 --worker_id=1 --ps_gpu=0 --schedule=train --worker_job='/job:worker'
+'{"task": {"index": 0, "type": "ps"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'    --schedule=run_std_server
+'{"task": {"index": 1, "type": "ps"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'    --schedule=run_std_server
+```
+
+Here's how we would start each job on their respective machines (the
+commands below assume that you're ssh'd into that job's machine):
+
+**Worker 1**:
+
+```
+$ export TF_CONFIG='{"task": {"index": 0, "type": "chief"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'
+$ t2t-trainer \
+    --master=grpc://10.0.0.1:5555 \
+    --ps_replicas=2 \
+    --worker_replicas=2 \
+    --worker_gpu=1 \
+    --worker_id=0 \
+    --ps_gpu=0 \
+    --schedule=train \
+    --worker_job='/job:chief' \
+    --model=transformer \
+    --hparams_set=transformer_base \
+    --problem=translate_ende_wmt32k
+```
+
+**Worker 2**:
+
+```
+$ export TF_CONFIG='{"task": {"index": 0, "type": "worker"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'
+$ t2t-trainer \
+    --master=grpc://10.0.0.2:5555 \
+    --ps_replicas=2 \
+    --worker_replicas=2 \
+    --worker_gpu=1 \
+    --worker_id=1 \
+    --ps_gpu=0 \
+    --schedule=train \
+    --worker_job='/job:worker' \
+    --model=transformer \
+    --hparams_set=transformer_base \
+    --problem=translate_ende_wmt32k
+```
+
+**PS 1**:
+
+```
+$ export TF_CONFIG='{"task": {"index": 0, "type": "ps"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'
+$ t2t-trainer --schedule=run_std_server
+```
+
+**PS 2**:
+
+```
+$ export TF_CONFIG='{"task": {"index": 1, "type": "ps"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'
+$ t2t-trainer --schedule=run_std_server
+```
+
+Increase `--worker_gpu` on each of the workers if you have multiple GPUs. If the
+parameter servers are also using GPUs, set `--ps_gpu` to the number of GPUs on
+the parameter servers.
diff --git a/tensor2tensor/bin/make_tf_configs.py b/tensor2tensor/bin/make_tf_configs.py
index d896e6b91..b38f16530 100644
--- a/tensor2tensor/bin/make_tf_configs.py
+++ b/tensor2tensor/bin/make_tf_configs.py
@@ -45,37 +45,64 @@ def main(_):
   masters = FLAGS.masters.split(",")
   ps = FLAGS.ps.split(",")
 
-  cluster = {"ps": ps, "master": masters}
-
-  for task_type, jobs in (("master", masters), ("ps", ps)):
-    for idx, job in enumerate(jobs):
-      if task_type == "master":
-        cmd_line_flags = " ".join([
-            "--master=grpc://%s" % job,
-            "--ps_replicas=%d" % len(ps),
-            "--worker_replicas=%d" % len(masters),
-            "--worker_gpu=1",
-            "--worker_id=%d" % idx,
-            "--worker_job='/job:master'",
-            "--ps_gpu=1",
-            "--schedule=train",
-            "--sync" if len(masters) == 1 else "",
-        ])
+  is_sync = len(masters) == 1
+  if is_sync:
+    print("Assuming SYNC distributed training with a single master and %d "
+          "workers" % len(ps))
+    cluster = {"ps": ps, "master": masters}
+  else:
+    print("Assuming ASYNC distributed training with %d workers and %d "
+          "parameter servers" % (len(masters), len(ps)))
+    cluster = {"ps": ps, "chief": [masters[0]], "worker": masters[1:]}
+
+  # Trainer configs
+  for idx, addr in enumerate(masters):
+    cmd_line_flags = [
+        "--master=grpc://%s" % addr,
+        "--ps_replicas=%d" % len(ps),
+        "--worker_replicas=%d" % len(masters),
+        "--worker_gpu=%d" % (0 if is_sync else 1),
+        "--worker_id=%d" % idx,
+        "--ps_gpu=%d" % (1 if is_sync else 0),
+        "--sync" if is_sync else "",
+        "--schedule=train",
+    ]
+    if is_sync:
+      task_type = "master"
+      cmd_line_flags.append("--worker_job='/job:master'")
+    else:
+      if idx == 0:
+        task_type = "chief"
+        idx = 0
+        cmd_line_flags.append("--worker_job='/job:chief'")
       else:
-        cmd_line_flags = " ".join([
-            "--master=grpc://%s" % job,
-            "--schedule=run_std_server",
-        ])
-
-      tf_config = json.dumps({
-          "cluster": cluster,
-          "task": {
-              "type": task_type,
-              "index": idx
-          },
-          "environment": "cloud",
-      })
-      print("'%s'\t%s" % (tf_config, cmd_line_flags))
+        task_type = "worker"
+        idx -= 1
+        cmd_line_flags.append("--worker_job='/job:worker'")
+
+    tf_config = json.dumps({
+        "cluster": cluster,
+        "task": {
+            "type": task_type,
+            "index": idx
+        },
+        "environment": "cloud",
+    })
+    cmd_line_flags = " ".join(cmd_line_flags)
+    print("'%s'\t%s" % (tf_config, cmd_line_flags))
+
+  # Std server configs
+  for idx, addr in enumerate(ps):
+    tf_config = json.dumps({
+        "cluster": cluster,
+        "task": {
+            "type": "ps",
+            "index": idx
+        },
+        "environment": "cloud",
+    })
+    cmd_line_flags = "--schedule=run_std_server"
+    print("'%s'\t%s" % (tf_config, cmd_line_flags))
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index c69ded415..8ebff6094 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -350,8 +350,15 @@ def maybe_cloud_tpu():
     yield
 
 
+def run_std_server():
+  exp = trainer_lib.T2TExperiment(*([None] * 5))
+  exp.run_std_server()
+
+
 def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
+  if FLAGS.schedule == "run_std_server":
+    run_std_server()
   trainer_lib.set_random_seed(FLAGS.random_seed)
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
   maybe_log_registry_and_exit()
diff --git a/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index 061be8f81..7aae0a3e0 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -130,10 +130,12 @@ def _replica_device_setter(worker_device):
         ps_tasks=ps_replicas,
         ps_device=ps_job + "/GPU:0" if ps_gpu > 0 else ps_job)
 
+  is_single_machine = ps_replicas == 0 and worker_replicas == 1
+
   if no_data_parallelism:
     datashard_devices = [""]
     caching_devices = None
-  elif schedule in ["train_and_evaluate", "continuous_train_and_eval"]:
+  elif is_single_machine:
     assert not sync
     tf.logging.warn(
         "Schedule=%s. Assuming that training is running on a single machine.",
@@ -162,7 +164,6 @@ def _replica_device_setter(worker_device):
           _replica_device_setter(worker_job + "/GPU:%d" % d)
           for d in _gpu_order(worker_gpu)
       ]
-      # caching_devices = [worker_job + "/GPU:0"] * worker_gpu
       caching_devices = None
     else:
       datashard_devices = [_replica_device_setter(worker_job)]
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 9b20ef57f..6a56c1f05 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import json
 import os
 import random
 import numpy as np
@@ -97,6 +98,11 @@ def create_hparams(hparams_set,
   return hparams
 
 
+def is_cloud_async_distributed():
+  return ("chief" in
+          json.loads(os.environ.get("TF_CONFIG", "{}")).get("cluster", {}))
+
+
 def create_run_config(master="",
                       model_dir=None,
                       warm_start_from=None,
@@ -156,9 +162,9 @@ def create_run_config(master="",
     del run_config_args["save_checkpoints_steps"]
   run_config_cls = tf.contrib.learn.RunConfig
 
-  # If using TPUEstimator, use TPU RunConfig, add TPUConfig, and add additional
-  # args.
   if use_tpu or use_tpu_estimator:
+    # If using TPUEstimator, use TPU RunConfig, add TPUConfig, and add
+    # additional args.
     tpu_config_kwargs = {
         "iterations_per_loop": iterations_per_loop,
         "num_shards": num_shards,
@@ -185,6 +191,10 @@ def create_run_config(master="",
       run_config_args["cluster"] = tpu_cluster_resolver
       del run_config_args["master"]
       del run_config_args["evaluation_master"]
+  elif is_cloud_async_distributed():
+    run_config_cls = tf.estimator.RunConfig
+    del run_config_args["master"]
+    del run_config_args["evaluation_master"]
 
   config = run_config_cls(**run_config_args)
   config.warm_start_from = warm_start_from
@@ -380,19 +390,11 @@ def run_std_server(self):
       ValueError: if not enough information is available in the estimator's
         config to create a server.
     """
-    config = self._estimator.config
-    if (not config.cluster_spec or not config.task_type or not config.master or
-        config.task_id is None):
-      raise ValueError("Could not start server; be sure to specify "
-                       "cluster_spec, task_type, master, and task in "
-                       "RunConfig or set the TF_CONFIG environment variable.")
+    config = tf.estimator.RunConfig()
     server = tf.train.Server(
         config.cluster_spec,
         job_name=config.task_type,
-        task_index=config.task_id,
-        config=config.tf_config,
-        start=False)
-    server.start()
+        task_index=config.task_id)
     server.join()
 
   def decode(self):

From 2250c2c2fa286b525760455f72f743fadf4fb45e Mon Sep 17 00:00:00 2001
From: Stefan Schweter <stefan@schweter.it>
Date: Thu, 16 Aug 2018 17:45:27 +0200
Subject: [PATCH 0598/2720] Minor doc fixes (#1000)

* docs: markdown tweak for gcloud-compute reference

* docs: fix reference to asr tutorial

* docs: minor markdown tweak for t2t-make-tf-configs reference
---
 docs/cloud_mlengine.md       | 4 ++--
 docs/cloud_tpu.md            | 2 +-
 docs/distributed_training.md | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/docs/cloud_mlengine.md b/docs/cloud_mlengine.md
index 8b40947db..83ebe7e57 100644
--- a/docs/cloud_mlengine.md
+++ b/docs/cloud_mlengine.md
@@ -8,8 +8,8 @@ you can easily launch Tensor2Tensor on it, including for hyperparameter tuning.
 
 It's the same `t2t-trainer` you know and love with the addition of the
 `--cloud_mlengine` flag, which by default will launch on a 1-GPU machine
-in the default compute region. See the [docs for `gcloud compute`]
-(https://cloud.google.com/compute/docs/gcloud-compute/#set_default_zone_and_region_in_your_local_client)
+in the default compute region. See the
+[docs for `gcloud compute`](https://cloud.google.com/compute/docs/gcloud-compute/#set_default_zone_and_region_in_your_local_client)
 to learn how to set the default compute region.
 
 ```
diff --git a/docs/cloud_tpu.md b/docs/cloud_tpu.md
index 2b81c19f5..b5f4e13a6 100644
--- a/docs/cloud_tpu.md
+++ b/docs/cloud_tpu.md
@@ -20,7 +20,7 @@ for some examples and try out your own problems.
 
 You can train an Automatic Speech Recognition (ASR) model with Transformer
 on TPU by using `transformer` as `model` with `transformer_librispeech_tpu` as
-`hparams_set` and `librispeech` as `problem`. See this [tutorial](tutorials/ast_with_transformer.md) for more details on training it and this
+`hparams_set` and `librispeech` as `problem`. See this [tutorial](tutorials/asr_with_transformer.md) for more details on training it and this
 [notebook](https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/asr_transformer.ipynb) to see how the resulting model transcribes your speech to text.
 
 Image Transformer:
diff --git a/docs/distributed_training.md b/docs/distributed_training.md
index 1fa00cb7a..f59974623 100644
--- a/docs/distributed_training.md
+++ b/docs/distributed_training.md
@@ -21,7 +21,7 @@ bucket (`gs://...`) or a directory on a shared network filesystem.
 
 ## Utility to produce `TF_CONFIG` and flags
 
-[`t2t-make-tf-configs`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/bin/t2t-make-tf-configs))
+[`t2t-make-tf-configs`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/bin/t2t-make-tf-configs)
 generates the `TF_CONFIG` json strings and the necessary command-line flags for
 the jobs.
 

From 9e3fadb9e294302b6c3215e8275eb51d60b2e61f Mon Sep 17 00:00:00 2001
From: Giovanni Campagna <scampa.giovanni@gmail.com>
Date: Thu, 16 Aug 2018 09:20:24 -0700
Subject: [PATCH 0599/2720] Undeprecate exporting the model from the trainer
 (#974)

But make use of tf.estimator.Exporter
---
 tensor2tensor/data_generators/problem.py |  9 +++++++++
 tensor2tensor/utils/trainer_lib.py       | 15 ++++++++++++---
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 107409c5c..d1d43c4f7 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -938,6 +938,15 @@ def prepare_for_output(example):
 
     return dataset
 
+  @property
+  def export_assets(self):
+    """Assets to export with the model.
+    
+    This property contains a dictionary of assets, such as vocabulary files,
+    that should be exported together with the model, or None if no assets
+    are needed."""
+    return None
+
   def serving_input_fn(self, hparams):
     """Input fn for serving export, starting from serialized example."""
     mode = tf.estimator.ModeKeys.PREDICT
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 6a56c1f05..f9e87d1b2 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -462,8 +462,16 @@ def create_experiment(
 
   # Export
   if export:
-    tf.logging.warn("Exporting from the trainer is deprecated. "
-                    "See serving/export.py.")
+    def compare_fn(best_eval_result, current_eval_result):
+        metric = eval_early_stopping_metric or "loss"
+        return current_eval_result[metric] < best_eval_result[metric]
+      
+    exporter = tf.estimator.BestExporter(name="best",
+                                         serving_input_receiver_fn=lambda: problem.serving_input_fn(hparams),
+                                         compare_fn=compare_fn,
+                                         assets_extra=problem.export_assets)
+  else:
+    exporter = None
 
   # Hooks
   validation_monitor_kwargs = dict(
@@ -523,7 +531,8 @@ def create_experiment(
       steps=eval_steps,
       hooks=eval_hooks,
       start_delay_secs=0 if hparams.schedule == "evaluate" else 120,
-      throttle_secs=eval_throttle_seconds)
+      throttle_secs=eval_throttle_seconds,
+      exporters=exporter)
 
   if autotune:
     hooks_kwargs = {"train_monitors": train_hooks, "eval_hooks": eval_hooks}

From 4bbd59569fd5efef4604605bb839f11fd09987ef Mon Sep 17 00:00:00 2001
From: Xiaoqing Zhou <zxqchat@hotmail.com>
Date: Fri, 17 Aug 2018 00:22:13 +0800
Subject: [PATCH 0600/2720] Update universal_transformer_util.py to fix
 TypeError (#987)

---
 tensor2tensor/models/research/universal_transformer_util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index c76490a33..62c71590c 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -578,7 +578,7 @@ def universal_transformer_basic(layer_inputs,
     layer_output:
          new_state: new state
   """
-  state, inputs, memory = layer_inputs
+  state, inputs, memory = tf.unstack(layer_inputs,num=None,axis=0,name="unstack")
   state = step_preprocess(state, step, hparams)
 
   new_state = ffn_unit(attention_unit(state))

From 0405ed11dd0a32973d59d89d5ac569b233d2f368 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 16 Aug 2018 08:49:25 -0700
Subject: [PATCH 0601/2720] internal merge of PR #1000

PiperOrigin-RevId: 208990232
---
 tensor2tensor/data_generators/problem.py          |  9 ---------
 .../models/research/universal_transformer_util.py |  2 +-
 tensor2tensor/utils/trainer_lib.py                | 15 +++------------
 3 files changed, 4 insertions(+), 22 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index d1d43c4f7..107409c5c 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -938,15 +938,6 @@ def prepare_for_output(example):
 
     return dataset
 
-  @property
-  def export_assets(self):
-    """Assets to export with the model.
-    
-    This property contains a dictionary of assets, such as vocabulary files,
-    that should be exported together with the model, or None if no assets
-    are needed."""
-    return None
-
   def serving_input_fn(self, hparams):
     """Input fn for serving export, starting from serialized example."""
     mode = tf.estimator.ModeKeys.PREDICT
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 62c71590c..c76490a33 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -578,7 +578,7 @@ def universal_transformer_basic(layer_inputs,
     layer_output:
          new_state: new state
   """
-  state, inputs, memory = tf.unstack(layer_inputs,num=None,axis=0,name="unstack")
+  state, inputs, memory = layer_inputs
   state = step_preprocess(state, step, hparams)
 
   new_state = ffn_unit(attention_unit(state))
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index f9e87d1b2..6a56c1f05 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -462,16 +462,8 @@ def create_experiment(
 
   # Export
   if export:
-    def compare_fn(best_eval_result, current_eval_result):
-        metric = eval_early_stopping_metric or "loss"
-        return current_eval_result[metric] < best_eval_result[metric]
-      
-    exporter = tf.estimator.BestExporter(name="best",
-                                         serving_input_receiver_fn=lambda: problem.serving_input_fn(hparams),
-                                         compare_fn=compare_fn,
-                                         assets_extra=problem.export_assets)
-  else:
-    exporter = None
+    tf.logging.warn("Exporting from the trainer is deprecated. "
+                    "See serving/export.py.")
 
   # Hooks
   validation_monitor_kwargs = dict(
@@ -531,8 +523,7 @@ def compare_fn(best_eval_result, current_eval_result):
       steps=eval_steps,
       hooks=eval_hooks,
       start_delay_secs=0 if hparams.schedule == "evaluate" else 120,
-      throttle_secs=eval_throttle_seconds,
-      exporters=exporter)
+      throttle_secs=eval_throttle_seconds)
 
   if autotune:
     hooks_kwargs = {"train_monitors": train_hooks, "eval_hooks": eval_hooks}

From 013f3ddda4dbccb8acbc8b1b59334ce6a0d8dcc6 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 16 Aug 2018 09:21:48 -0700
Subject: [PATCH 0602/2720] internal merge of PR #974

PiperOrigin-RevId: 208994727
---
 tensor2tensor/data_generators/problem.py | 11 +++++++++++
 tensor2tensor/utils/trainer_lib.py       | 15 ++++++++++++---
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 107409c5c..70e2daa97 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -938,6 +938,17 @@ def prepare_for_output(example):
 
     return dataset
 
+  @property
+  def export_assets(self):
+    """Assets to export with the model.
+
+    This property contains a dictionary of assets, such as vocabulary files,
+    that should be exported together with the model, or None if no assets
+    are needed.
+    """
+
+    return None
+
   def serving_input_fn(self, hparams):
     """Input fn for serving export, starting from serialized example."""
     mode = tf.estimator.ModeKeys.PREDICT
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 6a56c1f05..7e69211a3 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -461,9 +461,17 @@ def create_experiment(
                                                   hparams)
 
   # Export
+  exporter = None
   if export:
-    tf.logging.warn("Exporting from the trainer is deprecated. "
-                    "See serving/export.py.")
+    def compare_fn(best_eval_result, current_eval_result):
+      metric = eval_early_stopping_metric or "loss"
+      return current_eval_result[metric] < best_eval_result[metric]
+
+    exporter = tf.estimator.BestExporter(
+        name="best",
+        serving_input_receiver_fn=lambda: problem.serving_input_fn(hparams),
+        compare_fn=compare_fn,
+        assets_extra=problem.export_assets)
 
   # Hooks
   validation_monitor_kwargs = dict(
@@ -523,7 +531,8 @@ def create_experiment(
       steps=eval_steps,
       hooks=eval_hooks,
       start_delay_secs=0 if hparams.schedule == "evaluate" else 120,
-      throttle_secs=eval_throttle_seconds)
+      throttle_secs=eval_throttle_seconds,
+      exporters=exporter)
 
   if autotune:
     hooks_kwargs = {"train_monitors": train_hooks, "eval_hooks": eval_hooks}

From fd212f570ace6d5f402fce053d642d6b69e10b1b Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 16 Aug 2018 09:22:34 -0700
Subject: [PATCH 0603/2720] internal merge of PR #987

PiperOrigin-RevId: 208994831
---
 tensor2tensor/models/research/universal_transformer_util.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index c76490a33..82747e0b3 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -578,7 +578,8 @@ def universal_transformer_basic(layer_inputs,
     layer_output:
          new_state: new state
   """
-  state, inputs, memory = layer_inputs
+  state, inputs, memory = tf.unstack(layer_inputs, num=None, axis=0,
+                                     name="unstack")
   state = step_preprocess(state, step, hparams)
 
   new_state = ffn_unit(attention_unit(state))

From 5e27952be06e182565e7c47d0ad9026a419331bd Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 16 Aug 2018 11:11:39 -0700
Subject: [PATCH 0604/2720] Fix trainer_model_based name in travis

PiperOrigin-RevId: 209014290
---
 .travis.yml | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 1d25a3440..7f2b4a5e1 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -58,9 +58,9 @@ script:
   #   * registry_test
   #   * trainer_lib_test
   #   * visualization_test
-  #   * model_rl_experiment_test
+  #   * trainer_model_based_test
   #   * allen_brain_test
-  #   * model_rl_experiment_stochastic_test
+  #   * trainer_model_based_stochastic_test
   #   * models/research
   # algorithmic_math_test: flaky
   # universal_transformer_test: requires new feature in tf.foldl (rm with TF 1.9)
@@ -71,9 +71,9 @@ script:
     --ignore=tensor2tensor/bin/t2t_trainer_test.py
     --ignore=tensor2tensor/data_generators/algorithmic_math_test.py
     --ignore=tensor2tensor/models/research/universal_transformer_test.py
-    --ignore=tensor2tensor/rl/model_rl_experiment_test.py
+    --ignore=tensor2tensor/rl/trainer_model_based_test.py
     --ignore=tensor2tensor/data_generators/allen_brain_test.py
-    --ignore=tensor2tensor/rl/model_rl_experiment_stochastic_test.py
+    --ignore=tensor2tensor/rl/trainer_model_based_stochastic_test.py
     --ignore=tensor2tensor/models/research
   - pytest tensor2tensor/utils/registry_test.py
   - pytest tensor2tensor/utils/trainer_lib_test.py
@@ -102,7 +102,7 @@ script:
         pylint -j 2 tensor2tensor;
     fi
   - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
-        pytest tensor2tensor/rl/model_rl_experiment_test.py;
+        pytest tensor2tensor/rl/trainer_model_based_test.py;
     fi
   - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
         jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/hello_t2t.ipynb;

From c6e7ffe50a0e124c4bb635e651b6c9449f961eaa Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 16 Aug 2018 13:45:08 -0700
Subject: [PATCH 0605/2720] Calculated lengths of sequences using _raw.

PiperOrigin-RevId: 209039877
---
 tensor2tensor/models/lstm.py | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index beed64b37..c8634ced2 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -165,13 +165,10 @@ def lstm_seq2seq_internal(inputs, targets, hparams, train):
     return tf.expand_dims(decoder_outputs, axis=2)
 
 
-def lstm_seq2seq_internal_attention(inputs, targets, hparams, train):
+def lstm_seq2seq_internal_attention(inputs, targets, hparams, train,
+                                    inputs_length, targets_length):
   """LSTM seq2seq model with attention, main step used for training."""
   with tf.variable_scope("lstm_seq2seq_attention"):
-    # This is a temporary fix for varying-length sequences within in a batch.
-    # A more complete fix should pass a length tensor from outside so that
-    # all the lstm variants can use it.
-    inputs_length = common_layers.length_from_embedding(inputs)
     # Flatten inputs.
     inputs = common_layers.flatten4d3d(inputs)
 
@@ -183,7 +180,7 @@ def lstm_seq2seq_internal_attention(inputs, targets, hparams, train):
     # LSTM decoder with attention.
     shifted_targets = common_layers.shift_right(targets)
     # Add 1 to account for the padding added to the left from shift_right
-    targets_length = common_layers.length_from_embedding(shifted_targets) + 1
+    targets_length = targets_length + 1
     decoder_outputs = lstm_attention_decoder(
         common_layers.flatten4d3d(shifted_targets), hparams, train, "decoder",
         final_encoder_state, encoder_outputs, inputs_length, targets_length)
@@ -323,14 +320,27 @@ def body(self, features):
 
 @registry.register_model
 class LSTMSeq2seqAttention(t2t_model.T2TModel):
+  """Seq to seq LSTM with attention."""
 
   def body(self, features):
     # TODO(lukaszkaiser): investigate this issue and repair.
     if self._hparams.initializer == "orthogonal":
       raise ValueError("LSTM models fail with orthogonal initializer.")
     train = self._hparams.mode == tf.estimator.ModeKeys.TRAIN
+    # This is a temporary fix for varying-length sequences within in a batch.
+    # A more complete fix should pass a length tensor from outside so that
+    # all the lstm variants can use it.
+    input_shape = common_layers.shape_list(features["inputs_raw"])
+    flat_input = tf.reshape(features["inputs_raw"],
+                            [input_shape[0], input_shape[1]])
+    inputs_length = tf.reduce_sum(tf.minimum(flat_input, 1), -1)
+    target_shape = common_layers.shape_list(features["targets_raw"])
+    flat_target = tf.reshape(features["targets_raw"],
+                             [target_shape[0], target_shape[1]])
+    targets_length = tf.reduce_sum(tf.minimum(flat_target, 1), -1)
     return lstm_seq2seq_internal_attention(
-        features.get("inputs"), features["targets"], self._hparams, train)
+        features["inputs"], features["targets"], self._hparams, train,
+        inputs_length, targets_length)
 
 
 @registry.register_model

From 1c8226c62b91b22d33250db8393ab3c8441fda40 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 16 Aug 2018 15:11:36 -0700
Subject: [PATCH 0606/2720] Add NAC and NALU to common layers.

PiperOrigin-RevId: 209056430
---
 tensor2tensor/layers/common_layers.py      | 28 ++++++++++++++++++++++
 tensor2tensor/layers/common_layers_test.py | 25 +++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index c3171b47a..a4bc3c09e 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3360,6 +3360,34 @@ def belu(x):
   return tf.reshape(tf.concat([y1, y2], axis=-1), x_shape)
 
 
+def nac(x, depth, name=None, reuse=None):
+  """NAC as in https://arxiv.org/abs/1808.00508."""
+  with tf.variable_scope(
+      name, default_name="nac", values=[x], reuse=reuse):
+    x_shape = shape_list(x)
+    w = tf.get_variable("w", [x_shape[-1], depth])
+    m = tf.get_variable("m", [x_shape[-1], depth])
+    w = tf.tanh(w) * tf.nn.sigmoid(m)
+    x_flat = tf.reshape(x, [-1, x_shape[-1]])
+    res_flat = tf.matmul(x_flat, w)
+    return tf.reshape(res_flat, x_shape[:-1] + [depth])
+
+
+def nalu(x, depth, epsilon=1e-30, name=None, reuse=None):
+  """NALU as in https://arxiv.org/abs/1808.00508."""
+  with tf.variable_scope(
+      name, default_name="nalu", values=[x], reuse=reuse):
+    x_shape = shape_list(x)
+    x_flat = tf.reshape(x, [-1, x_shape[-1]])
+    gw = tf.get_variable("w", [x_shape[-1], depth])
+    g = tf.nn.sigmoid(tf.matmul(x_flat, gw))
+    g = tf.reshape(g, x_shape[:-1] + [depth])
+    a = nac(x, depth, name="nac_lin")
+    log_x = tf.log(tf.abs(x) + epsilon)
+    m = nac(log_x, depth, name="nac_log")
+    return g * a + (1 - g) * tf.exp(m)
+
+
 def argmax_with_score(logits, axis=None):
   """Argmax along with the value."""
   axis = axis or len(logits.get_shape()) - 1
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index b510e84e5..d4a6df588 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -417,6 +417,31 @@ def testBELU(self):
       actual = session.run(y)
     self.assertEqual(actual.shape, (5, 2, 1, 12))
 
+  def testNAC(self):
+    with self.test_session() as session:
+      x = np.random.rand(5, 2, 1, 12)
+      y = common_layers.nac(tf.constant(x, dtype=tf.float32), 14)
+      session.run(tf.global_variables_initializer())
+      actual = session.run(y)
+    self.assertEqual(actual.shape, (5, 2, 1, 14))
+
+  def testNALU(self):
+    with self.test_session() as session:
+      x = np.random.rand(5, 2, 1, 12)
+      y = common_layers.nalu(tf.constant(x, dtype=tf.float32), 14)
+      session.run(tf.global_variables_initializer())
+      actual = session.run(y)
+    self.assertEqual(actual.shape, (5, 2, 1, 14))
+
+  def testNALUzeros(self):
+    with self.test_session() as session:
+      x = np.random.rand(5, 2, 1, 12)
+      y = common_layers.nalu(tf.zeros_like(x, dtype=tf.float32), 14)
+      session.run(tf.global_variables_initializer())
+      actual = session.run(y)
+    self.assertTrue(np.all(np.isfinite(actual)))
+    self.assertEqual(actual.shape, (5, 2, 1, 14))
+
   def testPaddingCrossEntropyFactored(self):
     vocab_size = 19
     rows = 5

From fb16924d117424534cd9fd6b1afd4bcae0273ce2 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 16 Aug 2018 17:34:37 -0700
Subject: [PATCH 0607/2720] Internal change

PiperOrigin-RevId: 209079337
---
 tensor2tensor/layers/discretization.py        | 78 ++++++++++++-------
 tensor2tensor/models/research/autoencoders.py | 64 ++++++++++-----
 2 files changed, 95 insertions(+), 47 deletions(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 4fe443cf2..b40050767 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -802,7 +802,8 @@ def get_vq_codebook(codebook_size, hidden_size):
   return means, ema_means, ema_count
 
 
-def vq_nearest_neighbor(x, means, soft_em=False, num_samples=10):
+def vq_nearest_neighbor(x, means,
+                        soft_em=False, num_samples=10, temperature=None):
   """Find the nearest element in means to elements in x."""
   bottleneck_size = common_layers.shape_list(means)[0]
   x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keepdims=True)
@@ -815,7 +816,15 @@ def vq_nearest_neighbor(x, means, soft_em=False, num_samples=10):
         x_means_idx, depth=common_layers.shape_list(means)[0])
     x_means_hot = tf.reduce_mean(x_means_hot, axis=1)
   else:
-    x_means_idx = tf.argmax(-dist, axis=-1)
+    if temperature is None:
+      x_means_idx = tf.argmax(-dist, axis=-1)
+    else:
+      sm_dist = tf.nn.softmax(-dist)
+      x_means_idx = tf.multinomial(sm_dist / temperature, 1)
+      x_means_idx = tf.squeeze(x_means_idx, axis=-1)
+    if (common_layers.should_generate_summaries() and
+        not common_layers.is_xla_compiled()):
+      tf.summary.histogram("means_idx", tf.reshape(x_means_idx, [-1]))
     x_means_hot = tf.one_hot(x_means_idx, bottleneck_size)
   x_means_hot_flat = tf.reshape(x_means_hot, [-1, bottleneck_size])
   x_means = tf.matmul(x_means_hot_flat, means)
@@ -849,34 +858,45 @@ def vq_body(x,
             decay=0.999,
             epsilon=1e-5,
             soft_em=False,
-            num_samples=10):
+            num_samples=10,
+            temperature=None,
+            do_update=True):
   """Discretize each x into one of codebook_size codes."""
   x_shape = common_layers.shape_list(x)
   hidden_size = x_shape[-1]
   means, ema_means, ema_count = get_vq_codebook(codebook_size, hidden_size)
   x = tf.reshape(x, [-1, hidden_size])
   x_means_hot, e_loss, distances = vq_nearest_neighbor(
-      x, means, soft_em=soft_em, num_samples=num_samples)
-
-  # Update the ema variables
-  updated_ema_count = moving_averages.assign_moving_average(
-      ema_count,
-      tf.reduce_sum(tf.reshape(x_means_hot, shape=[-1, codebook_size]), axis=0),
-      decay,
-      zero_debias=False)
-
-  dw = tf.matmul(x_means_hot, x, transpose_a=True)
-  updated_ema_means = tf.identity(
-      moving_averages.assign_moving_average(
-          ema_means, dw, decay, zero_debias=False))
-  n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
-  updated_ema_count = (
-      (updated_ema_count + epsilon) / (n + codebook_size * epsilon) * n)
-  updated_ema_means /= tf.expand_dims(updated_ema_count, axis=-1)
-  with tf.control_dependencies([e_loss]):
-    update_means = means.assign(updated_ema_means)
-    with tf.control_dependencies([update_means]):
-      loss = beta * e_loss
+      x, means, soft_em=soft_em, num_samples=num_samples,
+      temperature=temperature)
+
+  def loss_with_update():
+    """Update the ema variables and return loss triggering the update."""
+    updated_ema_count = moving_averages.assign_moving_average(
+        ema_count,
+        tf.reduce_sum(tf.reshape(x_means_hot, shape=[-1, codebook_size]),
+                      axis=0),
+        decay,
+        zero_debias=False)
+
+    dw = tf.matmul(x_means_hot, x, transpose_a=True)
+    updated_ema_means = tf.identity(
+        moving_averages.assign_moving_average(
+            ema_means, dw, decay, zero_debias=False))
+    n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
+    updated_ema_count = (
+        (updated_ema_count + epsilon) / (n + codebook_size * epsilon) * n)
+    updated_ema_means /= tf.expand_dims(updated_ema_count, axis=-1)
+    with tf.control_dependencies([e_loss]):
+      update_means = means.assign(updated_ema_means)
+      with tf.control_dependencies([update_means]):
+        return beta * e_loss
+
+  # Loss, also do update if requested.
+  if do_update is True:
+    loss = loss_with_update()
+  else:
+    loss = tf.cond(do_update, loss_with_update, lambda: beta * e_loss)
 
   d = tf.reshape(x_means_hot, x_shape[:-1] + [codebook_size])
   return d, loss, distances
@@ -889,7 +909,9 @@ def vq_loss(x,
             decay=0.999,
             epsilon=1e-5,
             soft_em=False,
-            num_samples=10):
+            num_samples=10,
+            temperature=None,
+            do_update=True):
   """Compute the loss of large vocab tensors using a VQAE codebook.
 
   Args:
@@ -901,6 +923,8 @@ def vq_loss(x,
     epsilon: scalar float for moving averages
     soft_em: boolean, whether to apply a soft sampling procedure
     num_samples: if soft_em, number of samples to take
+    temperature: temperature if we want to sample nearest neighbors or None
+    do_update: whether to update the means; True by default, can be a Tensor
 
   Returns:
     discrete_x: one-hot Tensor indicating which codebook element is closest to x
@@ -928,7 +952,9 @@ def vq_loss(x,
       decay=decay,
       epsilon=epsilon,
       soft_em=soft_em,
-      num_samples=num_samples)
+      num_samples=num_samples,
+      temperature=temperature,
+      do_update=do_update)
 
   logits = -distances
   targets_loss = tf.losses.sparse_softmax_cross_entropy(
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index ffc4226f6..118b3bf35 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -32,8 +32,8 @@ def lrelu(input_, leak=0.2, name="lrelu"):
   return tf.maximum(input_, leak * input_, name=name)
 
 
-def reverse_gradient(x):
-  return -x + tf.stop_gradient(2 * x)
+def reverse_gradient(x, lr=1.0):
+  return -lr * x + tf.stop_gradient((1.0 + lr) * x)
 
 
 @registry.register_model
@@ -67,6 +67,9 @@ def image_summary(self, name, image_logits, max_outputs=1):
   def embed(self, x):
     """Input embedding with a non-zero bias for uniform inputs."""
     with tf.variable_scope("embed", reuse=tf.AUTO_REUSE):
+      x_shape = common_layers.shape_list(x)
+      # Merge channels and depth before embedding.
+      x = tf.reshape(x, x_shape[:-2] + [x_shape[-2] * x_shape[-1]])
       x = tf.layers.dense(
           x,
           self.hparams.hidden_size,
@@ -182,7 +185,6 @@ def body(self, features):
       vocab_size = self._problem_hparams.target_modality.top_dimensionality
       shape = common_layers.shape_list(labels)
       x = tf.one_hot(labels, vocab_size)
-      x = tf.reshape(x, shape[:-1] + [shape[-1] * vocab_size])
       x = self.embed(x)
       is1d = shape[2] == 1
       self.is1d = is1d
@@ -230,23 +232,36 @@ def body(self, features):
     ]
     res = tf.reshape(res, output_shape)
 
-    # Losses.
-    losses = {}
     if hparams.gan_loss_factor != 0.0:
       res_gan, res = tf.split(res, 2, axis=0)
-      with tf.variable_scope("vq"):
-        reconstr_gan, gan_codes, _, code_loss_gan, _ = discretization.vq_loss(
-            res_gan, labels, vocab_size)
-        losses["code_loss_gan"] = (code_loss_gan * hparams.code_loss_factor *
-                                   hparams.gan_loss_factor)
 
-    with tf.variable_scope("vq", reuse=tf.AUTO_REUSE):
+    # Losses.
+    losses = {}
+    vq_temperature = 0.001 / common_layers.inverse_exp_decay(
+        hparams.gan_codes_warmup_steps * 1.2, min_value=0.002)
+    if hparams.mode != tf.estimator.ModeKeys.TRAIN:
+      vq_temperature = None
+
+    with tf.variable_scope("vq_loss"):
       (reconstr, _, target_codes, code_loss,
-       targets_loss) = discretization.vq_loss(res, labels, vocab_size)
+       targets_loss) = discretization.vq_loss(
+           res, labels, vocab_size, temperature=vq_temperature)
 
     losses["code_loss"] = code_loss * hparams.code_loss_factor
     losses["training"] = targets_loss
 
+    # GAN losses.
+    if hparams.gan_loss_factor != 0.0:
+      with tf.variable_scope("vq_loss", reuse=True):
+        update_means_factor = common_layers.inverse_exp_decay(
+            hparams.gan_codes_warmup_steps, min_value=0.0001)
+        update_means = tf.less(tf.random_uniform([]), update_means_factor)
+        reconstr_gan, gan_codes, _, code_loss_gan, _ = discretization.vq_loss(
+            res_gan, labels, vocab_size, do_update=update_means,
+            temperature=vq_temperature)
+        code_loss_gan *= hparams.code_loss_factor * update_means_factor
+        losses["code_loss_gan"] = code_loss_gan
+
     # Add GAN loss if requested.
     gan_loss = 0.0
     if hparams.gan_loss_factor != 0.0:
@@ -261,11 +276,13 @@ def discriminate(x):
                                   tc_shape[:-2] + [tc_shape[-1] * tc_shape[-2]])
         gan_codes = tf.reshape(gan_codes,
                                tc_shape[:-2] + [tc_shape[-1] * tc_shape[-2]])
-      gan_loss = common_layers.sliced_gan_loss(target_codes,
-                                               reverse_gradient(gan_codes),
-                                               discriminate,
-                                               self.hparams.num_sliced_vecs)
-      gan_loss *= hparams.gan_loss_factor
+      gan_lr = common_layers.inverse_exp_decay(
+          hparams.gan_codes_warmup_steps * 1.5)
+      rev_grad_gan_codes = reverse_gradient(gan_codes, lr=gan_lr)
+      gan_loss = common_layers.sliced_gan_loss(
+          target_codes, rev_grad_gan_codes, discriminate,
+          self.hparams.num_sliced_vecs)
+      gan_loss *= hparams.gan_loss_factor  * update_means_factor
 
     self.image_summary("ae", reconstr)
 
@@ -364,12 +381,13 @@ def body(self, features):
     if "training" in losses:
       plain_training_loss = losses.pop("training")
       losses["plain"] = plain_training_loss
+    res_shape = common_layers.shape_list(basic_result)
+    vocab_size = self._problem_hparams.target_modality.top_dimensionality
     basic_result = self.embed(basic_result)
     shape = common_layers.shape_list(basic_result)
     basic1d = tf.reshape(basic_result, [shape[0], -1, shape[-1]])
     vocab_size = self._problem_hparams.target_modality.top_dimensionality
     targets = tf.one_hot(features["targets_raw"], vocab_size)
-    targets = tf.reshape(targets, shape[:-2] + [shape[-2] * vocab_size])
     targets = self.embed(targets)
     targets = tf.reshape(targets, common_layers.shape_list(basic_result))
     # During autoregressive inference, don't resample.
@@ -411,7 +429,8 @@ def body(self, features):
           padding="LEFT",
           activation=common_layers.belu,
           name="autoregressive_conv3")
-      return tf.reshape(res, shape), losses
+      res = tf.layers.dense(res, vocab_size, "autoregressive_final")
+      return tf.reshape(res, res_shape), losses
     if hparams.autoregressive_mode == "conv5":
       res = common_layers.conv1d(
           concat1d,
@@ -420,7 +439,8 @@ def body(self, features):
           padding="LEFT",
           activation=common_layers.belu,
           name="autoregressive_conv5")
-      return tf.reshape(res, shape), losses
+      res = tf.layers.dense(res, vocab_size, "autoregressive_final")
+      return tf.reshape(res, res_shape), losses
     if hparams.autoregressive_mode == "sru":
       res = common_layers.conv1d(
           concat1d,
@@ -430,7 +450,8 @@ def body(self, features):
           activation=common_layers.belu,
           name="autoregressive_sru_conv3")
       res = common_layers.sru(res)
-      return tf.reshape(res, shape), losses
+      res = tf.layers.dense(res, vocab_size, "autoregressive_final")
+      return tf.reshape(res, res_shape), losses
 
     raise ValueError(
         "Unsupported autoregressive mode: %s" % hparams.autoregressive_mode)
@@ -817,6 +838,7 @@ def autoencoder_basic():
   hparams.add_hparam("discriminator_batchnorm", True)
   hparams.add_hparam("num_sliced_vecs", 4096)
   hparams.add_hparam("code_loss_factor", 1.0)
+  hparams.add_hparam("gan_codes_warmup_steps", 5000)
   hparams.add_hparam("gan_loss_factor", 0.0)
   hparams.add_hparam("use_vqloss", False)
   return hparams

From 3b12e6c39326e1d118700692e87ad6f85c533a7c Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 16 Aug 2018 18:50:02 -0700
Subject: [PATCH 0608/2720] Re-enable serving tests in Travis using Docker

PiperOrigin-RevId: 209087220
---
 .travis.yml | 29 +++++++++++++----------------
 1 file changed, 13 insertions(+), 16 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 7f2b4a5e1..bc14f26e0 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -1,4 +1,9 @@
+sudo: required
 language: python
+git:
+  depth: 3
+services:
+  - docker
 python:
   - "2.7"
   - "3.6"
@@ -24,13 +29,8 @@ matrix:
     - python: "3.6"
       env: TF_VERSION="tf-nightly"
 before_install:
-  # Disabled TensorFlow Serving install until bug fixed. See "Export and query"
-  # section below.
-  # - echo "deb [arch=amd64] http://storage.googleapis.com/tensorflow-serving-apt stable tensorflow-model-server tensorflow-model-server-universal" | sudo tee /etc/apt/sources.list.d/tensorflow-serving.list
-  # - curl https://storage.googleapis.com/tensorflow-serving-apt/tensorflow-serving.release.pub.gpg | sudo apt-key add -
   - sudo apt-get update -qq
   - sudo apt-get install -qq libhdf5-dev
-  # - sudo apt-get install -qq tensorflow-model-server
 install:
   - if [[ "$TF_VERSION" == "tf-nightly"  ]];
     then
@@ -111,14 +111,11 @@ script:
         jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/t2t_problem.ipynb;
     fi
 
-  # Export and query (on Python 2 only)
-  # Bug: https://github.com/tensorflow/serving/issues/819
-  #- if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "1.6.*"  ]]; then
-  #      t2t-exporter --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR --model=transformer --hparams_set=transformer_tiny --output_dir=$T2T_TRAIN_DIR;
-  #      pip install tensorflow-serving-api;
-  #      tensorflow_model_server --port=9000 --model_name=my_model --model_base_path=$T2T_TRAIN_DIR/export/Servo &
-  #      sleep 10;
-  #      t2t-query-server --problem=$T2T_PROBLEM --server=localhost:9000 --servable_name=my_model --data_dir=$T2T_DATA_DIR --inputs_once='1 0 1 0 1 0';
-  #  fi
-git:
-  depth: 3
+  # Export and query
+  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
+        t2t-exporter --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR --model=transformer --hparams_set=transformer_tiny --output_dir=$T2T_TRAIN_DIR;
+        pip install tensorflow-serving-api;
+        docker run -d -p 8500:8500 --mount type=bind,source=$T2T_TRAIN_DIR/export/Servo,target=/models/my_model -e MODEL_NAME=my_model -t tensorflow/serving;
+        sleep 10;
+        t2t-query-server --problem=$T2T_PROBLEM --server=localhost:8500 --servable_name=my_model --data_dir=$T2T_DATA_DIR --inputs_once='1 0 1 0 1 0';
+    fi

From 36e144600f319500cdf7259ca4ad91682b982f1b Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Thu, 16 Aug 2018 21:16:20 -0700
Subject: [PATCH 0609/2720] Introducing StackWrapper.

PiperOrigin-RevId: 209098352
---
 tensor2tensor/data_generators/gym_problems.py |  2 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py    | 41 +++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 214bb724a..d907bdda1 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -48,7 +48,7 @@
 
 def standard_atari_env_spec(env):
   """Parameters of environment specification."""
-  standard_wrappers = [[tf_atari_wrappers.StackAndSkipWrapper, {"skip": 4}]]
+  standard_wrappers = [[tf_atari_wrappers.StackWrapper, {"history": 4}]]
   env_lambda = None
   if isinstance(env, str):
     env_lambda = lambda: gym.make(env)
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index f01020abc..4d5b34479 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -151,6 +151,47 @@ def _reset_non_empty(self, indices):
       return tf.gather(self.observ, indices)
 
 
+class StackWrapper(WrapperBase):
+  """ A wrapper which stacks previously seen frames. """
+
+  def __init__(self, batch_env, history=4):
+    super(StackWrapper, self).__init__(batch_env)
+    self.history = history
+    self.old_shape = batch_env.observ.shape.as_list()
+    observs_shape = self.old_shape[:-1] + [self.old_shape[-1] * self.history]
+    observ_dtype = tf.float32
+    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
+                               trainable=False)
+
+  def simulate(self, action):
+    reward, done = self._batch_env.simulate(action)
+    with tf.control_dependencies([reward, done]):
+      new_observ = self._batch_env.observ + 0
+      old_observ = tf.gather(
+          self._observ.read_value(),
+          range(self.old_shape[-1], self.old_shape[-1] * self.history),
+          axis=-1)
+      with tf.control_dependencies([new_observ, old_observ]):
+        with tf.control_dependencies([self._observ.assign(
+            tf.concat([old_observ, new_observ], axis=-1))]):
+          return tf.identity(reward), tf.identity(done)
+
+  def _reset_non_empty(self, indices):
+    # pylint: disable=protected-access
+    new_values = self._batch_env._reset_non_empty(indices)
+    # pylint: enable=protected-access
+    inx = tf.concat(
+        [
+            tf.ones(tf.size(tf.shape(new_values)), dtype=tf.int32)[:-1],
+            [self.history]
+        ],
+        axis=0)
+    assign_op = tf.scatter_update(self._observ, indices, tf.tile(
+        new_values, inx))
+    with tf.control_dependencies([assign_op]):
+      return tf.gather(self.observ, indices)
+
+
 class AutoencoderWrapper(WrapperBase):
   """ Transforms the observations taking the bottleneck
       state of an autoencoder"""

From 2f8423a7daf39c549fa4f87d369d3ff95e719e6c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 17 Aug 2018 02:06:38 -0700
Subject: [PATCH 0610/2720] Override serving_input_fn for video problems.

PiperOrigin-RevId: 209120097
---
 tensor2tensor/data_generators/video_utils.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 2877b0655..267ef0fb5 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -190,6 +190,20 @@ def example_reading_spec(self):
 
     return data_fields, data_items_to_decoders
 
+  def serving_input_fn(self, hparams):
+    """For serving/predict, assume that only video frames are provided."""
+    video_input_frames = tf.placeholder(
+        dtype=tf.float32,
+        shape=[
+            None, hparams.video_num_input_frames, self.frame_width,
+            self.frame_height, self.num_channels
+        ])
+
+    # TODO(michalski): add support for passing input_action and input_reward.
+    return tf.estimator.export.ServingInputReceiver(
+        features={"inputs": video_input_frames},
+        receiver_tensors=video_input_frames)
+
   def preprocess(self, dataset, mode, hparams, interleave=True):
     del interleave
     def split_on_batch(x):

From 837990f7a3b1466d04a4fff944ea583c9d4885be Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 17 Aug 2018 12:56:45 -0700
Subject: [PATCH 0611/2720] Fixing pylint signature-differs on various methods.

PiperOrigin-RevId: 209193066
---
 tensor2tensor/layers/common_layers.py          | 2 +-
 tensor2tensor/models/research/vqa_attention.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index a4bc3c09e..ce06e341a 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3991,7 +3991,7 @@ def _data_dep_init(self, inputs):
     self.layer.activation = activation
     self.initialized = True
 
-  def build(self, input_shape):
+  def build(self, input_shape=None):
     """Build `Layer`."""
     input_shape = tf.TensorShape(input_shape).as_list()
     self.input_spec = tf.layers.InputSpec(shape=input_shape)
diff --git a/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index 37787fe87..c7af49e30 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -97,7 +97,7 @@ def body(self, features):
     return tf.expand_dims(tf.expand_dims(output, axis=1), axis=2)
 
   def infer(self,
-            features,
+            features=None,
             decode_length=1,
             beam_size=1,
             top_beams=1,

From 8277f506964c5e5eae13933be52c42d49f4fb37d Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 17 Aug 2018 15:43:29 -0700
Subject: [PATCH 0612/2720] Store variable scopes in T2TModel; add
 T2TModel.initialize_from_ckpt

PiperOrigin-RevId: 209218783
---
 tensor2tensor/bin/t2t_trainer.py   |  2 +-
 tensor2tensor/utils/t2t_model.py   | 66 +++++++++++++++++++-----------
 tensor2tensor/utils/trainer_lib.py |  7 ++--
 3 files changed, 47 insertions(+), 28 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 8ebff6094..09cba7013 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -180,6 +180,7 @@ def create_experiment_fn(**kwargs):
       use_tpu=FLAGS.use_tpu,
       use_tpu_estimator=FLAGS.use_tpu_estimator,
       use_xla=FLAGS.xla_compile,
+      warm_start_from=FLAGS.warm_start_from,
       **kwargs)
 
 
@@ -214,7 +215,6 @@ def create_run_config(hp):
       hp.weight_dtype == "float32")
   return trainer_lib.create_run_config(
       model_dir=os.path.expanduser(FLAGS.output_dir),
-      warm_start_from=FLAGS.warm_start_from,
       master=FLAGS.master,
       iterations_per_loop=FLAGS.iterations_per_loop,
       num_shards=FLAGS.tpu_num_shards,
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index ced52766d..b5cdcde81 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -120,6 +120,11 @@ def __init__(self,
       self._create_modalities(self._problem_hparams, self._hparams)
     if not common_layers.is_xla_compiled():
       self.summarize_hparams()
+    self._variable_scopes = {}
+
+  def _add_variable_scope(self, key, vs):
+    if key not in self._variable_scopes:
+      self._variable_scopes[key] = vs
 
   def summarize_hparams(self):
     def create_hparams_summary(hparams, name):
@@ -261,7 +266,8 @@ def model_fn_sharded(self, sharded_features):
     return sharded_logits, losses
 
   def model_fn(self, features):
-    with tf.variable_scope(tf.get_variable_scope(), use_resource=True):
+    with tf.variable_scope(tf.get_variable_scope(), use_resource=True) as vs:
+      self._add_variable_scope("model_fn", vs)
       transformed_features = self.bottom(features)
 
       if self.hparams.activation_dtype == "bfloat16":
@@ -269,7 +275,8 @@ def model_fn(self, features):
           if v.dtype == tf.float32:
             transformed_features[k] = tf.cast(v, tf.bfloat16)
 
-      with tf.variable_scope("body"):
+      with tf.variable_scope("body") as body_vs:
+        self._add_variable_scope("body", body_vs)
         log_info("Building model body")
         body_out = self.body(transformed_features)
       output, losses = self._normalize_body_output(body_out)
@@ -302,7 +309,8 @@ def bottom(self, features):
         tf.logging.warning("Missing feature %s - ignoring." % key)
         continue
       do_reuse = input_modality.name in all_previous_modalities
-      with tf.variable_scope(input_modality.name, reuse=do_reuse):
+      with tf.variable_scope(input_modality.name, reuse=do_reuse) as im_vs:
+        self._add_variable_scope(input_modality.name, im_vs)
         log_info("Transforming feature '%s' with %s.bottom", key,
                  input_modality.name)
         transformed_features[key] = input_modality.bottom(features[key])
@@ -313,14 +321,16 @@ def bottom(self, features):
     if isinstance(target_modality, dict):
       for k, v in six.iteritems(target_modality):
         if k in features:
-          with tf.variable_scope(
-              "%s/%s" % (v.name, k)):  # TODO(aidangomez): share variables?
+          # TODO(aidangomez): share variables?
+          with tf.variable_scope("%s/%s" % (v.name, k)) as tm_vs:
+            self._add_variable_scope("%s/%s" % (v.name, k), tm_vs)
             log_info("Transforming '%s' with %s.targets_bottom", k, v.name)
             transformed_features[k] = v.targets_bottom(features[k])
         else:
           tf.logging.warn("Modality not found in features: %s", k)
     else:
-      with tf.variable_scope(target_modality.name):
+      with tf.variable_scope(target_modality.name) as tm_vs:
+        self._add_variable_scope(target_modality.name, tm_vs)
         if "targets" in features:
           log_info("Transforming 'targets' with %s.targets_bottom",
                    target_modality.name)
@@ -359,7 +369,8 @@ def _top_single(self, body_output, target_modality, features):
       log_warn("Without a Problem, T2TModel.top is a passthrough.")
       return body_output
 
-    with tf.variable_scope(target_modality.name):
+    with tf.variable_scope(target_modality.name) as tm_vs:
+      self._add_variable_scope(tm_vs.name, tm_vs)
       log_info("Transforming body output with %s.top", target_modality.name)
       last_only = (
           target_modality.top_is_pointwise and
@@ -401,7 +412,9 @@ def top(self, body_output, features):
             "problem_hparams.target_modality's dict." % k)
       logits = {}
       for k, v in six.iteritems(body_output):
-        with tf.variable_scope(k):  # TODO(aidangomez): share variables here?
+        # TODO(aidangomez): share variables here?
+        with tf.variable_scope(k) as top_vs:
+          self._add_variable_scope("top_%s" % k, top_vs)
           logits[k] = self._top_single(v, target_modality[k], features)
       return logits
     else:
@@ -1270,26 +1283,33 @@ def estimator_model_fn(cls,
     return model.estimator_spec_train(
         loss, num_async_replicas=num_async_replicas, use_tpu=use_tpu)
 
+  def initialize_from_ckpt(self, ckpt_dir):
+    model_dir = self._hparams.get("model_dir", None)
+    already_has_ckpt = (
+        model_dir and tf.train.latest_checkpoint(model_dir) is not None)
+    if already_has_ckpt:
+      return
+
+    # TODO(mitchellstern): Add support for partitioned variables?
+    reader = tf.contrib.framework.load_checkpoint(ckpt_dir)
+    variable_map = {}
+    for var in tf.contrib.framework.get_trainable_variables():
+      var_name = var.name.split(":")[0]
+      if reader.has_tensor(var_name):
+        tf.logging.info("Loading variable from checkpoint: %s", var_name)
+        variable_map[var_name] = var
+      else:
+        tf.logging.info(
+            "Cannot find variable in checkpoint, skipping: %s", var_name)
+    tf.train.init_from_checkpoint(ckpt_dir, variable_map)
+
   def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
     """Construct EstimatorSpec for TRAIN mode."""
     train_op = self.optimize(loss, num_async_replicas=num_async_replicas,
                              use_tpu=use_tpu)
 
-    # TODO(mitchellstern): Add support for partitioned variables?
-    if (tf.train.latest_checkpoint(self._hparams.model_dir) is None and
-        self._hparams.pretrained_model_dir):
-      pretrained_model_dir = self._hparams.pretrained_model_dir
-      reader = tf.contrib.framework.load_checkpoint(pretrained_model_dir)
-      variable_map = {}
-      for var in tf.contrib.framework.get_trainable_variables():
-        var_name = var.name.split(":")[0]
-        if reader.has_tensor(var_name):
-          tf.logging.info("Loading variable from checkpoint: %s", var_name)
-          variable_map[var_name] = var
-        else:
-          tf.logging.info(
-              "Cannot find variable in checkpoint, skipping: %s", var_name)
-      tf.train.init_from_checkpoint(pretrained_model_dir, variable_map)
+    if self._hparams.warm_start_from:
+      self.initialize_from_ckpt(self._hparams.warm_start_from)
 
     if use_tpu:
       host_call = _create_host_call(self.hparams.model_dir)
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 7e69211a3..03f9b3894 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -105,7 +105,6 @@ def is_cloud_async_distributed():
 
 def create_run_config(master="",
                       model_dir=None,
-                      warm_start_from=None,
                       iterations_per_loop=1000,
                       num_shards=8,
                       log_device_placement=False,
@@ -197,7 +196,6 @@ def create_run_config(master="",
     del run_config_args["evaluation_master"]
 
   config = run_config_cls(**run_config_args)
-  config.warm_start_from = warm_start_from
 
   # If not using TPU, add device info for data_parallelism
   config.use_tpu = use_tpu
@@ -259,7 +257,6 @@ def create_estimator(model_name,
         model_fn=model_fn,
         model_dir=run_config.model_dir,
         config=run_config,
-        warm_start_from=run_config.warm_start_from
     )
   return estimator
 
@@ -432,7 +429,8 @@ def create_experiment(
     use_tpu_estimator=False,
     use_xla=False,
     additional_train_hooks=None,
-    additional_eval_hooks=None):
+    additional_eval_hooks=None,
+    warm_start_from=None):
   """Create Experiment."""
   # HParams
   hparams.add_hparam("model_dir", run_config.model_dir)
@@ -440,6 +438,7 @@ def create_experiment(
   hparams.add_hparam("train_steps", train_steps)
   hparams.add_hparam("eval_steps", eval_steps)
   hparams.add_hparam("schedule", schedule)
+  hparams.add_hparam("warm_start_from", warm_start_from)
   add_problem_hparams(hparams, problem_name)
 
   # Estimator

From 178738ded5d013d7e7d3cfed75746f68219a4e7a Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Fri, 17 Aug 2018 20:19:38 -0700
Subject: [PATCH 0613/2720] fixing envs specs.

PiperOrigin-RevId: 209242641
---
 tensor2tensor/data_generators/gym_problems.py | 2 +-
 tensor2tensor/models/research/rl.py           | 2 +-
 tensor2tensor/rl/rl_trainer_lib_test.py       | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index d907bdda1..8034a2b96 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -54,7 +54,7 @@ def standard_atari_env_spec(env):
     env_lambda = lambda: gym.make(env)
   if callable(env):
     env_lambda = env
-  assert env is not None, "Unknown specification of environment"
+  assert env_lambda is not None, "Unknown specification of environment"
 
   return tf.contrib.training.HParams(
       env_lambda=env_lambda, wrappers=standard_wrappers, simulated_env=False)
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index de71cde6f..0c300544e 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -130,7 +130,7 @@ def simple_gym_spec(env):
     env_lambda = lambda: gym.make(env)
   if callable(env):
     env_lambda = env
-  assert env is not None, "Unknown specification of environment"
+  assert env_lambda is not None, "Unknown specification of environment"
 
   return tf.contrib.training.HParams(env_lambda=env_lambda,
                                      wrappers=standard_wrappers,
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index fbcfe7b14..267b0c94f 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -45,7 +45,7 @@ def test_no_crash_cartpole(self):
         TrainTest.test_config)
 
     hparams.add_hparam(
-        "environment_spec", gym_problems.standard_atari_env_spec("CartPole-v0"))
+        "environment_spec", rl_models.simple_gym_spec("CartPole-v0"))
     rl_trainer_lib.train(hparams)
 
   # This test should successfully train pong.

From 43979260b7cbe75f53caf98b37260d5da629c2fc Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 20 Aug 2018 10:32:45 -0700
Subject: [PATCH 0614/2720] v 1.8.0

PiperOrigin-RevId: 209442376
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index f868adcfd..aa007f92b 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.7.0',
+    version='1.8.0',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',

From a913de429c3e05766ef9c7ecbb4e66cf3009846b Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 20 Aug 2018 11:24:36 -0700
Subject: [PATCH 0615/2720] add support for non-deterministic ATARI modes and
 sticky keys.

PiperOrigin-RevId: 209453006
---
 .../data_generators/gym_problems_specs.py     | 45 ++++++++++++++++++-
 .../data_generators/gym_problems_test.py      |  7 +--
 2 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index f42c6bc82..e54208916 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -70,6 +70,25 @@
     "wrapped_full_pong",
 ]
 
+ATARI_ALL_MODES_SHORT_LIST = [
+    "pong",
+    "boxing",
+]
+
+# Different ATARI game modes in OpenAI Gym. Full list here:
+# https://github.com/openai/gym/blob/master/gym/envs/__init__.py
+ATARI_GAME_MODES = [
+    "Deterministic-v0",  # 0.25 repeat action probability, 4 frame skip.
+    "Deterministic-v4",  # 0.00 repeat action probability, 4 frame skip.
+    "NoFrameskip-v0",    # 0.25 repeat action probability, 1 frame skip.
+    "NoFrameskip-v4",    # 0.00 repeat action probability, 1 frame skip.
+    "-v0",               # 0.25 repeat action probability, (2 to 5) frame skip.
+    "-v4"                # 0.00 repeat action probability, (2 to 5) frame skip.
+]
+
+# List of all ATARI envs in all modes.
+ATARI_PROBLEMS = {}
+
 
 @registry.register_problem
 class GymWrappedFullPongRandom(GymDiscreteProblem):
@@ -154,13 +173,17 @@ def num_rewards(self):
     return 3
 
 
-def create_problems_for_game(game_name, clipped_reward=True):
+def create_problems_for_game(
+    game_name,
+    clipped_reward=True,
+    game_mode="Deterministic-v4"):
   """Create and register problems for game_name.
 
   Args:
     game_name: str, one of the games in ATARI_GAMES, e.g. "bank_heist".
     clipped_reward: bool, whether the rewards should be clipped. False is not
       yet supported.
+    game_mode: the frame skip and sticky keys config.
 
   Returns:
     dict of problems with keys ("base", "agent", "simulated").
@@ -173,9 +196,12 @@ def create_problems_for_game(game_name, clipped_reward=True):
                      "yet supported.")
   if game_name not in ATARI_GAMES:
     raise ValueError("Game %s not in ATARI_GAMES" % game_name)
+  if game_mode not in ATARI_GAME_MODES:
+    raise ValueError("Unknown ATARI game mode: %s." % game_mode)
   camel_game_name = "".join(
       [w[0].upper() + w[1:] for w in game_name.split("_")])
-  env_name = "%sDeterministic-v4" % camel_game_name
+  camel_game_name += game_mode
+  env_name = camel_game_name
   wrapped_env_name = "T2T%s" % env_name
 
   # Register an environment that does the reward clipping
@@ -188,8 +214,11 @@ def create_problems_for_game(game_name, clipped_reward=True):
   problem_cls = type("Gym%sRandom" % camel_game_name,
                      (GymClippedRewardRandom,),
                      {"env_name": wrapped_env_name})
+  registry.register_problem(problem_cls)
+
   with_agent_cls = type("GymDiscreteProblemWithAgentOn%s" % camel_game_name,
                         (GymRealDiscreteProblem, problem_cls), {})
+
   registry.register_problem(with_agent_cls)
 
   # Create and register the simulated Problem
@@ -206,3 +235,15 @@ def create_problems_for_game(game_name, clipped_reward=True):
       "agent": with_agent_cls,
       "simulated": simulated_cls,
   }
+
+# Register the atari games with all of the possible modes.
+for game in ATARI_ALL_MODES_SHORT_LIST:
+  ATARI_PROBLEMS[game] = {}
+  for mode in ATARI_GAME_MODES:
+    classes = create_problems_for_game(
+        game,
+        clipped_reward=True,
+        game_mode=mode)
+    ATARI_PROBLEMS[game][mode] = classes
+
+
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index ebdbcf021..6b6193bef 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -34,9 +34,10 @@ def setUpClass(cls):
     shutil.rmtree(cls.tmp_dir)
     os.mkdir(cls.tmp_dir)
 
-  def testGymAtariBoots(self):
-    problem = gym_problems_specs.create_problems_for_game("pong")["base"]()
-    self.assertEqual(210, problem.frame_height)
+  def testGymAtariGameModes(self):
+    for mode in gym_problems_specs.ATARI_GAME_MODES:
+      problem = gym_problems_specs.ATARI_PROBLEMS["pong"][mode]["base"]()
+      self.assertEqual(210, problem.frame_height)
 
 
 if __name__ == "__main__":

From 8dd9c45f0246be4b7f94d60158c1e3f61fea0f44 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 20 Aug 2018 12:04:55 -0700
Subject: [PATCH 0616/2720] Multi-frame reward prediction.

PiperOrigin-RevId: 209460402
---
 .../models/research/next_frame_params.py      |  1 +
 .../models/research/next_frame_sv2p.py        | 45 ++++++++++++++++---
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index a0d4294ed..8253cbed3 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -75,6 +75,7 @@ def next_frame_stochastic():
   hparams.add_hparam("stochastic_model", True)
   hparams.add_hparam("reward_prediction", True)
   hparams.add_hparam("reward_prediction_stop_gradient", True)
+  hparams.add_hparam("reward_prediction_buffer_size", 4)
   hparams.add_hparam("model_options", "CDNA")
   hparams.add_hparam("num_masks", 10)
   hparams.add_hparam("latent_channels", 1)
diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 02d6e7b26..46986d107 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -344,10 +344,28 @@ def construct_latent_tower(self, images):
       return ret_mean, ret_std
 
   def reward_prediction(self, input_image, input_reward, action, latent):
-    """Builds a reward prediction network."""
+    """Builds a reward prediction network.
+
+    Args:
+      input_image: image sequences (buffer_size x batch x W x H x C)
+      input_reward: previous reward (batch + reward_size)
+      action: current action (batch + action_size)
+      latent: current latent (batch + latent_size)
+    Returns:
+      latent_mean: predicted latent mean
+      latent_std: predicted latent standard deviation
+      latent_loss: loss of the latent twoer
+      samples: random samples sampled from standard guassian
+    """
+
     del action
     del latent
 
+    # Unstack buffer images on buffer axis
+    input_image = tf.unstack(input_image, axis=0)
+    # Concat buffer images on channels.
+    input_image = tf.concat(input_image, axis=3)
+
     conv_size = self.tinyify([32, 32, 16, 4])
 
     with tf.variable_scope("reward_pred", reuse=tf.AUTO_REUSE):
@@ -513,12 +531,14 @@ def construct_model(self,
     context_frames = self.hparams.video_num_input_frames
 
     batch_size = common_layers.shape_list(images)[1]
+    reward_buf_size = self.hparams.reward_prediction_buffer_size
     ss_func = self.get_scheduled_sample_func(batch_size)
 
     def process_single_frame(prev_outputs, inputs):
       """Process a single frame of the video."""
       cur_image, cur_reward, action = inputs
-      time_step, prev_image, prev_reward, lstm_states = prev_outputs
+      (time_step, prev_image, prev_reward,
+       lstm_states, frames_buffer) = prev_outputs
 
       generated_items = [prev_image, prev_reward]
       groundtruth_items = [cur_image, cur_reward]
@@ -534,16 +554,28 @@ def process_single_frame(prev_outputs, inputs):
         reward_input_image = pred_image
         if self.hparams.reward_prediction_stop_gradient:
           reward_input_image = tf.stop_gradient(reward_input_image)
-        pred_reward = self.reward_prediction(
-            reward_input_image, input_reward, action, latent)
+
+        reward_input_image = tf.expand_dims(reward_input_image, axis=0)
+        frames_buffer = tf.concat(
+            [frames_buffer[1:], reward_input_image], axis=0)
+
+        pred_reward = tf.cond(
+            tf.less(time_step, reward_buf_size),
+            lambda: input_reward,  # HACK. just return something.
+            lambda: self.reward_prediction(  # pylint: disable=g-long-lambda
+                frames_buffer, input_reward, action, latent)
+        )
       else:
         pred_reward = input_reward
 
       time_step += 1
-      outputs = (time_step, pred_image, pred_reward, lstm_states)
+      outputs = (time_step, pred_image, pred_reward, lstm_states, frames_buffer)
 
       return outputs
 
+    # Create frames buffer for reward prediction
+    frames_buffer = tf.identity(images[:reward_buf_size])
+
     # Latent tower
     latent = None
     if self.hparams.stochastic_model:
@@ -556,7 +588,8 @@ def process_single_frame(prev_outputs, inputs):
     prev_outputs = (tf.constant(0),
                     tf.zeros_like(images[0]),
                     tf.zeros_like(rewards[0]),
-                    lstm_states)
+                    lstm_states,
+                    frames_buffer)
 
     initializers = process_single_frame(prev_outputs, inputs)
     first_gen_images = tf.expand_dims(initializers[1], axis=0)

From ee52b1527d20e9efd22ab0d49d6974fc5b33bcd4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 20 Aug 2018 13:44:55 -0700
Subject: [PATCH 0617/2720] Change cloud tpu zone.

PiperOrigin-RevId: 209476297
---
 docs/cloud_tpu.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/cloud_tpu.md b/docs/cloud_tpu.md
index b5f4e13a6..cf34523e7 100644
--- a/docs/cloud_tpu.md
+++ b/docs/cloud_tpu.md
@@ -49,7 +49,7 @@ Configure the `gcloud` CLI:
 gcloud components update
 gcloud auth application-default login
 # Set your default zone to a TPU-enabled zone.
-gcloud config set compute/zone us-central1-f
+gcloud config set compute/zone us-central1-b
 ```
 
 Generate data to GCS.

From b634c5e8d378f6c92cfd712f098e0032f0c65e03 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 20 Aug 2018 14:29:09 -0700
Subject: [PATCH 0618/2720] Don't daisy-chain variables in Universal
 Transformer.

PiperOrigin-RevId: 209484453
---
 tensor2tensor/models/research/universal_transformer.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 81380afc8..a38e4c187 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -353,6 +353,8 @@ def update_hparams_for_universal_transformer(hparams):
     hparams with default values for Universal Transformers hyper-parameters
 
   """
+  hparams.daisy_chain_variables = False  # Breaks multi-gpu in while loops.
+
   # If not None, mixes vanilla transformer with Universal Transformer.
   # Options: None, "before_ut", and "after_ut".
   hparams.add_hparam("mix_with_transformer", None)

From fda65a132a967637193f5d9b8a6579d264aef79b Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Mon, 20 Aug 2018 14:34:09 -0700
Subject: [PATCH 0619/2720] Adding a pretrain schedule.

PiperOrigin-RevId: 209485527
---
 tensor2tensor/data_generators/multi_problem.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 4c4e5d70d..d5a744dcc 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -31,6 +31,7 @@ class MixingSchedule(object):
   """Available schedules for mixing datasets."""
   EXPONENTIAL = "exponential"
   CONSTANT = "constant"
+  PRETRAIN = "pretrain"
 
 
 class MultiProblem(problem.Problem):
@@ -190,6 +191,14 @@ def get_exp_sched_prob():
       def get_const_sched_prob():
         return hparams.multiproblem_schedule_threshold
 
+      def get_pretrain_sched_prob():
+        """Pretrain the primary tasks for max examples."""
+        with tf.control_dependencies([problem_step.assign_add(1)]):
+          return tf.cond(
+              tf.greater(problem_step,
+                         hparams.multiproblem_schedule_max_examples),
+              lambda: 1.0, lambda: 0.0)
+
       def mix_data(example):
         """Function to mix the different datasets according to a schedule."""
         del example
@@ -200,6 +209,8 @@ def mix_data(example):
           prob = get_exp_sched_prob()
         elif hparams.multiproblem_mixing_schedule == MixingSchedule.CONSTANT:
           prob = get_const_sched_prob()
+        elif hparams.multiproblem_mixing_schedule == MixingSchedule.PRETRAIN:
+          prob = get_pretrain_sched_prob()
         else:
           raise ValueError("Unknown schedule %s" % str(
               hparams.multiproblem_mixing_schedule))
@@ -208,6 +219,10 @@ def mix_data(example):
                             hparams.multiproblem_mixing_schedule))
         tf.logging.info("Schedule mixing threshold "
                         "%.2f" % hparams.multiproblem_schedule_threshold)
+        prob = tf.cond(
+            tf.equal(tf.floormod(problem_step, 5e6), 0),
+            lambda: tf.Print(prob, [prob], message="Probability"),
+            lambda: prob)
 
         def sample_task(curr_task, num_tasks_left, randnum):
           """A recursive function to sample a task.

From 828030da9986bcf13c6ba9ad09f32ff1437158e9 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 20 Aug 2018 15:10:29 -0700
Subject: [PATCH 0620/2720] Multi-frame reward prediction.

PiperOrigin-RevId: 209492494
---
 .../models/research/next_frame_params.py      |  1 -
 .../models/research/next_frame_sv2p.py        | 45 +++----------------
 2 files changed, 6 insertions(+), 40 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_params.py
index 8253cbed3..a0d4294ed 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_params.py
@@ -75,7 +75,6 @@ def next_frame_stochastic():
   hparams.add_hparam("stochastic_model", True)
   hparams.add_hparam("reward_prediction", True)
   hparams.add_hparam("reward_prediction_stop_gradient", True)
-  hparams.add_hparam("reward_prediction_buffer_size", 4)
   hparams.add_hparam("model_options", "CDNA")
   hparams.add_hparam("num_masks", 10)
   hparams.add_hparam("latent_channels", 1)
diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 46986d107..02d6e7b26 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -344,28 +344,10 @@ def construct_latent_tower(self, images):
       return ret_mean, ret_std
 
   def reward_prediction(self, input_image, input_reward, action, latent):
-    """Builds a reward prediction network.
-
-    Args:
-      input_image: image sequences (buffer_size x batch x W x H x C)
-      input_reward: previous reward (batch + reward_size)
-      action: current action (batch + action_size)
-      latent: current latent (batch + latent_size)
-    Returns:
-      latent_mean: predicted latent mean
-      latent_std: predicted latent standard deviation
-      latent_loss: loss of the latent twoer
-      samples: random samples sampled from standard guassian
-    """
-
+    """Builds a reward prediction network."""
     del action
     del latent
 
-    # Unstack buffer images on buffer axis
-    input_image = tf.unstack(input_image, axis=0)
-    # Concat buffer images on channels.
-    input_image = tf.concat(input_image, axis=3)
-
     conv_size = self.tinyify([32, 32, 16, 4])
 
     with tf.variable_scope("reward_pred", reuse=tf.AUTO_REUSE):
@@ -531,14 +513,12 @@ def construct_model(self,
     context_frames = self.hparams.video_num_input_frames
 
     batch_size = common_layers.shape_list(images)[1]
-    reward_buf_size = self.hparams.reward_prediction_buffer_size
     ss_func = self.get_scheduled_sample_func(batch_size)
 
     def process_single_frame(prev_outputs, inputs):
       """Process a single frame of the video."""
       cur_image, cur_reward, action = inputs
-      (time_step, prev_image, prev_reward,
-       lstm_states, frames_buffer) = prev_outputs
+      time_step, prev_image, prev_reward, lstm_states = prev_outputs
 
       generated_items = [prev_image, prev_reward]
       groundtruth_items = [cur_image, cur_reward]
@@ -554,28 +534,16 @@ def process_single_frame(prev_outputs, inputs):
         reward_input_image = pred_image
         if self.hparams.reward_prediction_stop_gradient:
           reward_input_image = tf.stop_gradient(reward_input_image)
-
-        reward_input_image = tf.expand_dims(reward_input_image, axis=0)
-        frames_buffer = tf.concat(
-            [frames_buffer[1:], reward_input_image], axis=0)
-
-        pred_reward = tf.cond(
-            tf.less(time_step, reward_buf_size),
-            lambda: input_reward,  # HACK. just return something.
-            lambda: self.reward_prediction(  # pylint: disable=g-long-lambda
-                frames_buffer, input_reward, action, latent)
-        )
+        pred_reward = self.reward_prediction(
+            reward_input_image, input_reward, action, latent)
       else:
         pred_reward = input_reward
 
       time_step += 1
-      outputs = (time_step, pred_image, pred_reward, lstm_states, frames_buffer)
+      outputs = (time_step, pred_image, pred_reward, lstm_states)
 
       return outputs
 
-    # Create frames buffer for reward prediction
-    frames_buffer = tf.identity(images[:reward_buf_size])
-
     # Latent tower
     latent = None
     if self.hparams.stochastic_model:
@@ -588,8 +556,7 @@ def process_single_frame(prev_outputs, inputs):
     prev_outputs = (tf.constant(0),
                     tf.zeros_like(images[0]),
                     tf.zeros_like(rewards[0]),
-                    lstm_states,
-                    frames_buffer)
+                    lstm_states)
 
     initializers = process_single_frame(prev_outputs, inputs)
     first_gen_images = tf.expand_dims(initializers[1], axis=0)

From c14819b29044f4fba04aea4877a2a07de5999a7a Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Mon, 20 Aug 2018 16:06:00 -0700
Subject: [PATCH 0621/2720] adjusting trainer_model_based for different modes
 in atari envs.

PiperOrigin-RevId: 209502286
---
 tensor2tensor/rl/trainer_model_based.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index a5c9c6516..73f6ca1b2 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -347,24 +347,27 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     subdirectories.append("autoencoder")
   directories = setup_directories(output_dir, subdirectories)
 
+  game_with_mode = hparams.game + "deterministic-v4"
   # Problems
   if using_autoencoder:
     problem_name = (
-        "gym_discrete_problem_with_agent_on_%s_with_autoencoder" % hparams.game)
+        "gym_discrete_problem_with_agent_on_%s_with_autoencoder"
+        % game_with_mode)
     world_model_problem = (
-        "gym_discrete_problem_with_agent_on_%s_autoencoded" % hparams.game)
+        "gym_discrete_problem_with_agent_on_%s_autoencoded" % game_with_mode)
     simulated_problem_name = (
         "gym_simulated_discrete_problem_with_agent_on_%s_autoencoded"
-        % hparams.game)
+        % game_with_mode)
   else:
-    problem_name = ("gym_discrete_problem_with_agent_on_%s" % hparams.game)
+    problem_name = ("gym_discrete_problem_with_agent_on_%s" % game_with_mode)
     world_model_problem = problem_name
     simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
-                              % hparams.game)
+                              % game_with_mode)
     if problem_name not in registry.list_problems():
       tf.logging.info("Game Problem %s not found; dynamically registering",
                       problem_name)
-      gym_problems_specs.create_problems_for_game(hparams.game)
+      gym_problems_specs.create_problems_for_game(hparams.game,
+                                                  game_mode="Deterministic-v4")
 
   # Autoencoder model dir
   autoencoder_model_dir = directories.get("autoencoder")

From 9a5a39aee3f53a832395e03b746cb94a20b2cee7 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Mon, 20 Aug 2018 17:40:19 -0700
Subject: [PATCH 0622/2720] six.moves.range

PiperOrigin-RevId: 209516631
---
 tensor2tensor/rl/envs/tf_atari_wrappers.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 4d5b34479..5f4489fe4 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -20,6 +20,8 @@
 
 import math
 
+from six.moves import range  # pylint: disable=redefined-builtin
+
 from tensor2tensor.layers import discretization
 from tensor2tensor.models.research import autoencoders
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
@@ -169,7 +171,7 @@ def simulate(self, action):
       new_observ = self._batch_env.observ + 0
       old_observ = tf.gather(
           self._observ.read_value(),
-          range(self.old_shape[-1], self.old_shape[-1] * self.history),
+          list(range(self.old_shape[-1], self.old_shape[-1] * self.history)),
           axis=-1)
       with tf.control_dependencies([new_observ, old_observ]):
         with tf.control_dependencies([self._observ.assign(

From c25fc2d37ce1cd7b79e4e4d66b3e6993a76dc99c Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 20 Aug 2018 18:49:04 -0700
Subject: [PATCH 0623/2720] Rm tensor2tensor rev_block in favor of
 tf.contrib.layers.rev_block

PiperOrigin-RevId: 209524010
---
 tensor2tensor/layers/rev_block.py             | 345 ------------------
 tensor2tensor/layers/rev_block_test.py        | 167 ---------
 .../models/research/transformer_revnet.py     |   5 +-
 tensor2tensor/models/revnet.py                |  26 +-
 4 files changed, 18 insertions(+), 525 deletions(-)
 delete mode 100644 tensor2tensor/layers/rev_block.py
 delete mode 100644 tensor2tensor/layers/rev_block_test.py

diff --git a/tensor2tensor/layers/rev_block.py b/tensor2tensor/layers/rev_block.py
deleted file mode 100644
index 3d8ad5b2e..000000000
--- a/tensor2tensor/layers/rev_block.py
+++ /dev/null
@@ -1,345 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-"""Reversible Residual Block.
-
-From
-[The Reversible Residual Network: Backpropagation Without Storing
-Activations](https://arxiv.org/abs/1707.04585).
-"""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import re
-from six.moves import range  # pylint: disable=redefined-builtin
-
-from tensor2tensor.layers import common_layers
-import tensorflow as tf
-
-LAYER_RE = re.compile(".*revlayer_([0-9]*)/([fg])/.*")
-
-
-def _acc_grads(*lists_of_grads):
-  """Accumulates lists of gradients."""
-  acc_grads = []
-  for grads in zip(*lists_of_grads):
-    grads = [g for g in grads if g is not None]
-    if grads:
-      acc_grads.append(tf.add_n(grads))
-    else:
-      acc_grads.append(None)
-  return acc_grads
-
-
-def _rev_layer_forward(xs, f, g, f_side_input, g_side_input,
-                       gate_outputs=False):
-  """Forward for 1 reversible layer."""
-  x1, x2 = xs
-  y1 = x1 + (f(x2, f_side_input) if f_side_input else f(x2))
-  y2 = x2 + (g(y1, g_side_input) if g_side_input else g(y1))
-  out = (y1, y2)
-  if gate_outputs:
-    out = tf.tuple(out)
-  return out
-
-
-def _rev_layer_backward(ys, grad_ys, f, g, f_vars, f_side_input, g_vars,
-                        g_side_input):
-  """Backprop for 1 layer."""
-  y1, y2 = ys
-  grad_y1, grad_y2 = grad_ys
-
-  # Reconstruct intermediates and inputs (x1, x2)
-  # stop_gradients required on fn inputs to prevent infinite recursion into this
-  # grad function on the calls to tf.gradients.
-  y1_stop = tf.stop_gradient(y1)
-  g_side_input = [tf.stop_gradient(t) for t in g_side_input]
-  gy1 = g(y1_stop, g_side_input) if g_side_input else g(y1_stop)
-
-  x2 = y2 - gy1
-  x2_stop = tf.stop_gradient(x2)
-  f_side_input = [tf.stop_gradient(t) for t in f_side_input]
-  fx2 = f(x2_stop, f_side_input) if f_side_input else f(x2_stop)
-
-  x1 = y1 - fx2
-
-  # Compute gradients wrt to inputs
-  # dL/dy2 * dG(y1)/y1
-  grad_gy1_y2 = tf.gradients(gy1, y1_stop, grad_y2)[0]
-  grad_x1 = grad_y1 + grad_gy1_y2
-  grad_x2 = (
-      tf.gradients(fx2, x2_stop, grad_y1)[0] + grad_y2 +
-      tf.gradients(fx2, x2_stop, grad_gy1_y2)[0])
-
-  # Compute gradients wrt to vars and side inputs in f and g
-  grads1 = tf.gradients(gy1, g_vars + g_side_input, grad_y2)
-  grad_g_vars, grad_g_side = grads1[:len(g_vars)], grads1[len(g_vars):]
-  grads2 = tf.gradients(fx2, f_vars + f_side_input, grad_y1)
-  grad_f_y1, grad_f_side1 = grads2[:len(f_vars)], grads2[len(f_vars):]
-  grads3 = tf.gradients(fx2, f_vars + f_side_input, grad_gy1_y2)
-  grad_f_y2, grad_f_side2 = grads3[:len(f_vars)], grads3[len(f_vars):]
-  grad_f_vars = _acc_grads(grad_f_y1, grad_f_y2)
-
-  grad_f_side = _acc_grads(grad_f_side1, grad_f_side2)
-
-  # Put returns in a tuple to ensure a constant memory budget (i.e. don't want
-  # the subsequent layer to start computing and consuming memory based on a
-  # subset of these values).
-  outputs = ((x1, x2), (grad_x1, grad_x2), (grad_f_vars, grad_f_side),
-             (grad_g_vars, grad_g_side))
-  tupled = tf.tuple(tf.contrib.framework.nest.flatten(outputs))
-  return tf.contrib.framework.nest.pack_sequence_as(outputs, tupled)
-
-
-def _rev_block_forward(x1,
-                       x2,
-                       f,
-                       g,
-                       num_layers=1,
-                       f_side_input=None,
-                       g_side_input=None,
-                       gate_outputs=False):
-  """Forward for a series of reversible layers."""
-  out = (x1, x2)
-  for i in range(num_layers):
-    out = _rev_layer_forward(
-        out, f[i], g[i], f_side_input, g_side_input, gate_outputs=gate_outputs)
-
-  y1, y2 = out
-  return y1, y2
-
-
-class RevBlock(object):
-  """Block of reversible layers. See rev_block."""
-
-  def __init__(self,
-               f,
-               g,
-               num_layers=1,
-               f_side_input=None,
-               g_side_input=None,
-               use_efficient_backprop=True):
-
-    if isinstance(f, list):
-      assert len(f) == num_layers
-    else:
-      f = [f] * num_layers
-
-    if isinstance(g, list):
-      assert len(g) == num_layers
-    else:
-      g = [g] * num_layers
-
-    scope_prefix = "revblock/revlayer_%d/"
-    f_scope = scope_prefix + "f"
-    g_scope = scope_prefix + "g"
-
-    f = [
-        tf.make_template(f_scope % i, fn, create_scope_now_=True)
-        for i, fn in enumerate(f)
-    ]
-    g = [
-        tf.make_template(g_scope % i, fn, create_scope_now_=True)
-        for i, fn in enumerate(g)
-    ]
-
-    self.f = f
-    self.g = g
-
-    self.num_layers = num_layers
-    self.f_side_input = f_side_input or []
-    self.g_side_input = g_side_input or []
-
-    self._use_efficient_backprop = use_efficient_backprop
-
-  def _efficient_grad_fn(self, inputs, variables, ys, grad_ys):
-    """Custom gradient fn for a block of reversible residual layers."""
-    side_inputs = inputs[2:]
-    f_side_idxs = [None] * len(self.f_side_input)
-    g_side_idxs = [None] * len(self.g_side_input)
-    assert len(side_inputs) == len(self.f_side_input) + len(self.g_side_input)
-
-    for i, t in enumerate(side_inputs):
-      if t in self.f_side_input:
-        f_side_idxs[self.f_side_input.index(t)] = i
-      elif t in self.g_side_input:
-        g_side_idxs[self.g_side_input.index(t)] = i
-      else:
-        assert False
-
-    f_vars = [[] for _ in range(self.num_layers)]
-    g_vars = [[] for _ in range(self.num_layers)]
-    f_vars_idxs = [[] for _ in range(self.num_layers)]
-    g_vars_idxs = [[] for _ in range(self.num_layers)]
-
-    for i, t in enumerate(variables):
-      ref = common_layers.underlying_variable_ref(t)
-
-      # Use the name to identify the layer number and function (f or g)
-      regex = LAYER_RE.match(ref.name)
-      layer_no = int(regex.group(1))
-      fn_name = regex.group(2)
-      if fn_name == "f":
-        f_vars[layer_no].append(ref)
-        f_vars_idxs[layer_no].append(i)
-      else:
-        assert fn_name == "g"
-        g_vars[layer_no].append(ref)
-        g_vars_idxs[layer_no].append(i)
-
-    f_var_grads = []
-    g_var_grads = []
-    f_side_grads = []
-    g_side_grads = []
-
-    # Reverse variable containers to go backward
-    f_vars.reverse()
-    g_vars.reverse()
-    f = list(self.f)
-    g = list(self.g)
-    f.reverse()
-    g.reverse()
-
-    for i in range(self.num_layers):
-      ys, grad_ys, f_ret, g_ret = _rev_layer_backward(
-          ys, grad_ys, f[i], g[i], f_vars[i], self.f_side_input, g_vars[i],
-          self.g_side_input)
-
-      grad_f_vars, grad_f_side = f_ret
-      grad_g_vars, grad_g_side = g_ret
-      f_var_grads.append(grad_f_vars)
-      g_var_grads.append(grad_g_vars)
-      f_side_grads.append(grad_f_side)
-      g_side_grads.append(grad_g_side)
-
-    # Accumulate layer gradients for f_side_input and g_side_input
-    acc_f_side_grads = _acc_grads(*f_side_grads)
-    acc_g_side_grads = _acc_grads(*g_side_grads)
-
-    # Use the stored idxs to put gradients in the passed-in order.
-    side_input_grads = [None] * len(side_inputs)
-    variable_grads = [None] * len(variables)
-
-    # Variable gradients were collected in reverse layer order. Reverse to match
-    # idxs.
-    f_var_grads.reverse()
-    g_var_grads.reverse()
-    for idxs, grads in list(zip(f_vars_idxs, f_var_grads)) + list(
-        zip(g_vars_idxs, g_var_grads)):
-      for i, grad in zip(idxs, grads):
-        variable_grads[i] = grad
-
-    for i, grad in zip(f_side_idxs, acc_f_side_grads):
-      side_input_grads[i] = grad
-    for i, grad in zip(g_side_idxs, acc_g_side_grads):
-      side_input_grads[i] = grad
-
-    grad_x1, grad_x2 = grad_ys
-    return [grad_x1, grad_x2] + side_input_grads, variable_grads
-
-  def forward(self, x1, x2):
-    """Run forward through the reversible layers."""
-
-    side_inputs = [self.f_side_input, self.g_side_input]
-    flat_side_inputs = tf.contrib.framework.nest.flatten(side_inputs)
-
-    custom_grad_fn = (
-        self._efficient_grad_fn if self._use_efficient_backprop else None)
-
-    @common_layers.fn_with_custom_grad(custom_grad_fn)
-    def _forward(x1_, x2_, *flat_side_inputs):
-      f_side, g_side = tf.contrib.framework.nest.pack_sequence_as(
-          side_inputs, flat_side_inputs)
-      return _rev_block_forward(
-          x1_,
-          x2_,
-          self.f,
-          self.g,
-          num_layers=self.num_layers,
-          f_side_input=f_side,
-          g_side_input=g_side,
-          gate_outputs=self._use_efficient_backprop)
-
-    return _forward(x1, x2, *flat_side_inputs)
-
-  def backward(self, y1, y2):
-    """Run backward through the reversible layers."""
-
-    f = list(self.f)
-    g = list(self.g)
-    f.reverse()
-    g.reverse()
-
-    for i in range(self.num_layers):
-      gy1 = g[i](y1, self.g_side_input) if self.g_side_input else g[i](y1)
-      x2 = y2 - gy1
-      fx2 = f[i](x2, self.f_side_input) if self.f_side_input else f[i](x2)
-      x1 = y1 - fx2
-
-      y1, y2 = x1, x2
-
-    return x1, x2
-
-
-def rev_block(x1,
-              x2,
-              f,
-              g,
-              num_layers=1,
-              f_side_input=None,
-              g_side_input=None,
-              is_training=True):
-  """A block of reversible residual layers.
-
-  A reversible residual layer is defined as:
-
-  ```
-  y1 = x1 + f(x2, f_side_input)
-  y2 = x2 + g(y1, g_side_input)
-  ```
-
-  A reversible residual block, defined here, is a series of reversible residual
-  layers.
-
-  Limitations:
-  * f and g must not close over any Tensors; all side inputs to f and g should
-    be passed in with f_side_input and g_side_input which will be forwarded to
-    f and g.
-  * f and g must not change the dimensionality of their inputs in order for the
-    addition in the equations above to work.
-
-  Args:
-    x1: a float Tensor.
-    x2: a float Tensor.
-    f: a function, (Tensor) -> (Tensor) (or list of such of length num_layers).
-      Should not change the shape of the Tensor. Expected to create variables.
-      See f_side_input if there are side inputs.
-    g: a function, (Tensor) -> (Tensor) (or list of such of length num_layers).
-      Should not change the shape of the Tensor. Expected to create variables.
-      See g_side_input if there are side inputs.
-    num_layers: int, number of reversible residual layers. Each layer will
-      apply f and g according to the equations above, with new variables in each
-      layer.
-    f_side_input: list of Tensors, side input to f. If not None, signature of f
-      should be (Tensor, list<Tensor>) -> (Tensor).
-    g_side_input: list of Tensors, side input to g. If not None, signature of g
-      should be (Tensor, list<Tensor>) -> (Tensor).
-    is_training: bool, whether to actually use the efficient backprop codepath.
-
-  Returns:
-    y1, y2: tuple of float Tensors.
-  """
-  block = RevBlock(f, g, num_layers, f_side_input, g_side_input, is_training)
-  return block.forward(x1, x2)
diff --git a/tensor2tensor/layers/rev_block_test.py b/tensor2tensor/layers/rev_block_test.py
deleted file mode 100644
index 93944373f..000000000
--- a/tensor2tensor/layers/rev_block_test.py
+++ /dev/null
@@ -1,167 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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 RevBlock."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-from tensor2tensor.layers import rev_block
-
-import tensorflow as tf
-
-
-class RevBlockTest(tf.test.TestCase):
-  CHANNELS = 8
-  NUM_LAYERS = 4
-  BATCH_SIZE = 16
-
-  def testForwardBackward(self):
-
-    def f(x):
-      return tf.layers.dense(x, self.CHANNELS // 2, use_bias=True)
-
-    def g(x):
-      return tf.layers.dense(x, self.CHANNELS // 2, use_bias=True)
-
-    x = tf.random_uniform([self.BATCH_SIZE, self.CHANNELS], dtype=tf.float32)
-    x1, x2 = tf.split(x, 2, axis=-1)
-
-    block = rev_block.RevBlock(f, g, num_layers=3)
-    y1, y2 = block.forward(x1, x2)
-    x1_inv, x2_inv = block.backward(y1, y2)
-
-    with self.test_session() as sess:
-      sess.run(tf.global_variables_initializer())
-      x1, x2, x1_inv, x2_inv = sess.run([x1, x2, x1_inv, x2_inv])
-
-      self.assertAllClose(x1, x1_inv)
-      self.assertAllClose(x2, x2_inv)
-
-  def testBackwardForward(self):
-
-    def f(x):
-      return tf.layers.dense(x, self.CHANNELS // 2, use_bias=True)
-
-    def g(x):
-      return tf.layers.dense(x, self.CHANNELS // 2, use_bias=True)
-
-    y = tf.random_uniform([self.BATCH_SIZE, self.CHANNELS], dtype=tf.float32)
-    y1, y2 = tf.split(y, 2, axis=-1)
-
-    block = rev_block.RevBlock(f, g, num_layers=3)
-    x1, x2 = block.backward(y1, y2)
-    y1_inv, y2_inv = block.forward(x1, x2)
-
-    with self.test_session() as sess:
-      sess.run(tf.global_variables_initializer())
-      y1, y2, y1_inv, y2_inv = sess.run([y1, y2, y1_inv, y2_inv])
-
-      self.assertAllClose(y1, y1_inv)
-      self.assertAllClose(y2, y2_inv)
-
-  def _test_rev_block(self,
-                      x=None,
-                      f=None,
-                      g=None,
-                      f_side_input=None,
-                      g_side_input=None):
-    tf.set_random_seed(1234)
-
-    if f is None:
-
-      def f(x):  # pylint: disable=function-redefined
-        return tf.layers.dense(x, self.CHANNELS // 2, use_bias=True)
-
-    if g is None:
-
-      def g(x):  # pylint: disable=function-redefined
-        return tf.layers.dense(x, self.CHANNELS // 2, use_bias=True)
-
-    if f_side_input is None:
-      f_side_input = []
-
-    if g_side_input is None:
-      g_side_input = []
-
-    if x is None:
-      x = tf.random_uniform([self.BATCH_SIZE, self.CHANNELS], dtype=tf.float32)
-    x1, x2 = tf.split(x, 2, axis=-1)
-
-    with tf.variable_scope("rev_test") as vs:
-      y1_rev, y2_rev = rev_block.rev_block(
-          x1,
-          x2,
-          f,
-          g,
-          f_side_input=f_side_input,
-          g_side_input=g_side_input,
-          num_layers=self.NUM_LAYERS)
-      y_rev = tf.concat([y1_rev, y2_rev], axis=1)
-      fg_vars = vs.trainable_variables()
-
-    num_vars = len(tf.global_variables())
-    with tf.variable_scope(vs, reuse=True):
-      y1, y2 = rev_block.rev_block(
-          x1,
-          x2,
-          f,
-          g,
-          f_side_input=f_side_input,
-          g_side_input=g_side_input,
-          num_layers=self.NUM_LAYERS,
-          is_training=False)
-      y = tf.concat([y1, y2], axis=1)
-    # Ensure no new vars were created - full reuse
-    assert len(tf.global_variables()) == num_vars
-
-    loss_rev = tf.reduce_mean(y_rev + 10.)
-    loss = tf.reduce_mean(y + 10.)
-
-    wrt = [x] + f_side_input + g_side_input + fg_vars
-    grads_rev = tf.gradients(loss_rev, wrt)
-    grads = tf.gradients(loss, wrt)
-
-    with self.test_session() as sess:
-      sess.run(tf.global_variables_initializer())
-      y_val, yd_val, gd_val, g_val = sess.run([y, y_rev, grads_rev, grads])
-      self.assertAllClose(y_val, yd_val)
-      for g1, g2 in zip(gd_val, g_val):
-        self.assertAllClose(g1, g2)
-
-  def testRevBlock(self):
-    self._test_rev_block()
-
-  def testSideInput(self):
-    f_side_input = tf.random_uniform([self.BATCH_SIZE, self.CHANNELS // 2])
-
-    def f(x, side_input):
-      return tf.layers.dense(
-          x, self.CHANNELS // 2, use_bias=True) + side_input[0]
-
-    self._test_rev_block(f=f, f_side_input=[f_side_input])
-
-  def testMultipleFns(self):
-
-    def f1(x):
-      return tf.layers.dense(x, self.CHANNELS // 2)
-
-    def f2(x):
-      return tf.layers.dense(x, self.CHANNELS // 2, activation=tf.nn.relu)
-
-    self._test_rev_block(f=[f1, f2, f1, f2])
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/models/research/transformer_revnet.py b/tensor2tensor/models/research/transformer_revnet.py
index 5935e1c0d..f31b262bc 100644
--- a/tensor2tensor/models/research/transformer_revnet.py
+++ b/tensor2tensor/models/research/transformer_revnet.py
@@ -19,7 +19,6 @@
 from __future__ import print_function
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
-from tensor2tensor.layers import rev_block
 from tensor2tensor.models import transformer
 from tensor2tensor.utils import registry
 
@@ -120,7 +119,7 @@ def g(x):
   x1, x2 = tf.split(encoder_input, 2, axis=-1)
 
   with tf.variable_scope(name):
-    y1, y2 = rev_block.rev_block(
+    y1, y2 = tf.contrib.layers.rev_block(
         x1,
         x2,
         f,
@@ -198,7 +197,7 @@ def g(x):
   x1, x2 = tf.split(decoder_input, 2, axis=-1)
 
   with tf.variable_scope(name):
-    y1, y2 = rev_block.rev_block(
+    y1, y2 = tf.contrib.layers.rev_block(
         x1,
         x2,
         f,
diff --git a/tensor2tensor/models/revnet.py b/tensor2tensor/models/revnet.py
index 6db6f8940..49304105b 100644
--- a/tensor2tensor/models/revnet.py
+++ b/tensor2tensor/models/revnet.py
@@ -36,23 +36,29 @@
 
 import functools
 from tensor2tensor.layers import common_hparams
-from tensor2tensor.layers import rev_block
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
 import tensorflow as tf
 
+
+def wrapped_partial(fn, *args, **kwargs):
+  partial = functools.partial(fn, *args, **kwargs)
+  wrapped = functools.update_wrapper(partial, fn)
+  return wrapped
+
+
 conv_initializer = tf.contrib.layers.variance_scaling_initializer(
     factor=2.0, mode='FAN_OUT')
 
-CONFIG = {'2d': {'conv': functools.partial(
+CONFIG = {'2d': {'conv': wrapped_partial(
     tf.layers.conv2d, kernel_initializer=conv_initializer),
                  'max_pool': tf.layers.max_pooling2d,
                  'avg_pool': tf.layers.average_pooling2d,
                  'split_axis': 3,
                  'reduction_dimensions': [1, 2]
                 },
-          '3d': {'conv': functools.partial(
+          '3d': {'conv': wrapped_partial(
               tf.layers.conv3d, kernel_initializer=conv_initializer),
                  'max_pool': tf.layers.max_pooling3d,
                  'avg_pool': tf.layers.average_pooling2d,
@@ -225,9 +231,9 @@ def unit(x1, x2, block_num, depth, num_layers, dim='2d',
   else:
     depth1 = depth2 = depth
 
-  residual = functools.partial(f,
-                               depth1=depth1, depth2=depth2, dim=dim,
-                               training=training, bottleneck=bottleneck)
+  residual = wrapped_partial(f,
+                             depth1=depth1, depth2=depth2, dim=dim,
+                             training=training, bottleneck=bottleneck)
 
   with tf.variable_scope(scope_name):
     downsample = downsample_bottleneck if bottleneck else downsample_residual
@@ -244,10 +250,10 @@ def unit(x1, x2, block_num, depth, num_layers, dim='2d',
 
     # Full block using memory-efficient rev_block implementation.
     with tf.variable_scope('full_block'):
-      x1, x2 = rev_block.rev_block(x1, x2,
-                                   residual,
-                                   residual,
-                                   num_layers=num_layers)
+      x1, x2 = tf.contrib.layers.rev_block(x1, x2,
+                                           residual,
+                                           residual,
+                                           num_layers=num_layers)
       return x1, x2
 
 

From 0f28f4df829f2fad0c4626b485611820930b2c5d Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Mon, 20 Aug 2018 20:26:35 -0700
Subject: [PATCH 0624/2720] fix atari env registration.

PiperOrigin-RevId: 209531611
---
 tensor2tensor/rl/trainer_model_based.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 73f6ca1b2..93e807b61 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -347,7 +347,10 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     subdirectories.append("autoencoder")
   directories = setup_directories(output_dir, subdirectories)
 
-  game_with_mode = hparams.game + "deterministic-v4"
+  if hparams.game in gym_problems_specs.ATARI_GAMES:
+    game_with_mode = hparams.game + "_deterministic-v4"
+  else:
+    game_with_mode = hparams.game
   # Problems
   if using_autoencoder:
     problem_name = (

From d94b8eedb9a5bb5def1630532e4045a37f1badfe Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Tue, 21 Aug 2018 10:44:18 -0700
Subject: [PATCH 0625/2720] move mixture-of-experts code and experiments into
 research directory.  Some changes to mixture-of-experts.

PiperOrigin-RevId: 209620750
---
 tensor2tensor/mesh_tensorflow/mtf_layers.py   | 225 -----------
 .../mesh_tensorflow/mtf_transformer.py        |  52 +--
 .../research/experiments_moe.py               |  80 ++--
 tensor2tensor/mesh_tensorflow/research/moe.py | 365 ++++++++++++++++++
 4 files changed, 420 insertions(+), 302 deletions(-)
 create mode 100644 tensor2tensor/mesh_tensorflow/research/moe.py

diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers.py b/tensor2tensor/mesh_tensorflow/mtf_layers.py
index 17f789a1b..c5d121c09 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers.py
@@ -599,228 +599,3 @@ def attention_bias_local_block(mesh, block_length, memory_length,
                   dtype=dtype)
   mask = mtf.cast(mask, dtype=tf.float32)  * -1e9
   return mask
-
-
-
-
-def moe_v0(inputs,
-           hidden_dim,
-           output_dim,
-           experts_dim,
-           loss_coef=1e-3,
-           overhead=1.0):
-  """Local mixture of experts that works well on TPU.
-
-  See https://arxiv.org/abs/1701.06538
-
-  There are num_experts expert networks, each containing a relu-activated
-  hidden layer of size hidden_size, followed by an output projection.
-
-  The number of parameters is thus:
-    num_experts * (input_size * hidden_size + hidden_size * output_size)
-
-  The input is 3d: [batch, length, depth], consisting of the representations
-  of all positions in a batch of sequences.
-
-  Each position of each sequence is sent to 0-2 experts.  The expert
-  choices and the combination weights are determined by a learned gating
-  function.
-
-  This function returns a small auxiliary loss that should be added to the
-  training loss of the model.  This loss helps to balance expert usage.
-  Without the loss, it is very likely that a few experts will be trained and
-  the rest will starve.
-
-  Several hacks are necessary to get around current TPU limitations:
-
-  - To ensure static shapes, we enforce (by truncation/padding)
-    that each sequence send the same number of elements to each expert.
-
-    It would make more sense to enforce this equality over the entire batch,
-    as opposed to on individual sequences.  This would allow more freedom
-    for individual sequences to be unbalanced.  Unfortunately, that would
-    slow down our hacked-up gather-by-matmul implementation.
-
-    TODO(noam): There is no real reason for a single sequence to be the unit
-      of equal allocation.  Reshaping the inputs would allow us to pick a
-      different unit of equal allocation.
-
-  TODO(noam): Factor this code better.  We want to be able to substitute
-  different code for the experts themselves.  We also want to integrate this
-  gating/dispatching logic into multi-device mixtures-of-experts.
-
-  Args:
-    inputs: a mtf.Tensor with shape [batch_dim, length_dim, input_dim]
-    hidden_dim: a mtf.Dimension
-    output_dim: a mtf.Dimension
-    experts_dim: a mtf.Dimension
-    loss_coef: a float scalar
-    overhead: multiplicative factor of how much spare capacity to assign
-
-  Returns:
-    outputs: a Tensor with shape [batch_dim, length_dim, output_dim]
-    loss: a mtf scalar
-  """
-  batch_dim, length_dim, input_dim = inputs.shape.dims
-
-  # Each sequence sends expert_capacity positions to each expert.
-  expert_capacity = min(
-      length_dim.size,
-      int((length_dim.size * 2 * overhead) / experts_dim.size))
-  expert_capacity_dim = mtf.Dimension("expert_capacity", expert_capacity)
-
-  experts_dim_unsplit = mtf.Dimension("expert_unsplit", experts_dim.size)
-  batch_dim_unsplit = mtf.Dimension("batch_unsplit", batch_dim.size)
-
-  # This is the learned gating function.
-  # shape = [batch_dim, length_dim, experts_dim_unsplit]
-  gates = mtf.softmax(dense(inputs, experts_dim_unsplit), experts_dim_unsplit)
-
-  backward_assignment = _truncated_top_2_gating_mtf(
-      gates, length_dim, experts_dim_unsplit, expert_capacity_dim)
-
-  forward_assignment = mtf.cast(
-      mtf.cast(backward_assignment, tf.bool), inputs.dtype)
-
-  # put num_experts dimension first to make split easier in alltoall
-  expert_inputs = mtf.einsum([inputs, forward_assignment], mtf.Shape(
-      [experts_dim_unsplit, batch_dim, expert_capacity_dim, input_dim]))
-
-  expert_inputs = mtf.reshape(expert_inputs, mtf.Shape(
-      [experts_dim, batch_dim_unsplit, expert_capacity_dim, input_dim]))
-
-  # Now feed the expert inputs through the experts.
-  h = dense(expert_inputs, hidden_dim, expert_dims=[experts_dim],
-            activation=mtf.relu, name="x0")
-  expert_output = dense(h, output_dim, expert_dims=[experts_dim], name="x1")
-
-  expert_output = mtf.reshape(expert_output, mtf.Shape(
-      [experts_dim_unsplit, batch_dim, expert_capacity_dim, input_dim]))
-
-  output = mtf.einsum([expert_output, backward_assignment], mtf.Shape(
-      [batch_dim, length_dim, output_dim]))
-
-  importance = mtf.reduce_sum(backward_assignment, output_shape=mtf.Shape(
-      [batch_dim, experts_dim_unsplit]))
-
-  loss = cv_squared(importance) * loss_coef
-  return output, loss
-
-
-def cv_squared(x):
-  """The squared coefficient of variation of a sample.
-
-  Useful as a loss to encourage a positive distribution to be more uniform.
-  Epsilons added for numerical stability.
-  Returns 0 for an empty Tensor.
-
-  Args:
-    x: a mtf.Tensor
-
-  Returns:
-    a mtf Scalar
-  """
-  epsilon = 1e-10
-  mean = mtf.reduce_mean(x)
-  variance = mtf.reduce_mean(mtf.square(x - mean))
-  return variance / (mtf.square(mean) + epsilon)
-
-
-def _truncated_top_2_gating_mtf(
-    gates, group_dim, experts_dim, expert_capacity_dim):
-  """Compute gating for mixture-of-experts in TensorFlow.
-
-  gates is usually the output of a softmax function.
-  The return value is a dense representation of the mapping between
-  the input positions in the positions in the batches sent to the experts.
-
-  TODO(noam): this function contains code factored out of
-  expert_utils.local_moe_tpu.  Move this function to that file and
-  call it from both places.
-
-  Args:
-    gates: a Tensor
-    group_dim: one dimension of gates
-    experts_dim: one dimension of gates
-    expert_capacity_dim: a Dimension not in gates
-
-  Returns:
-    a Tensor with shape gates.shape + expert_capacity_dim
-
-  Raises:
-    ValueError: if group_dim has size >256
-  """
-  gates = mtf.to_float(gates)
-  expert_capacity_f = float(expert_capacity_dim.size)
-  # Find the top expert for each position. shape=[batch, group]
-  index_1, gate_1 = mtf.top_1(gates, experts_dim)
-  # [batch, group, experts]
-  mask_1 = mtf.one_hot(index_1, experts_dim, dtype=gates.dtype)
-
-  if expert_capacity_dim.size > 256:
-    # using mtf.cumsum (implemented on TPU as bfloat16 matmul) to compute
-    # position in the mini-batch sent to the expert.  This will cause
-    # very bad things to happen if expert_capacity_dim > 256.
-    raise ValueError(
-        "expert_capacity_dim.size must be <=256 to avoid roundoff errors in"
-        " indices - got %s" % (expert_capacity_dim,))
-  # [batch, group, experts]
-  # This is the position within the expert's mini-batch for this sequence
-  position_in_expert_1 = mtf.cumsum(mask_1, group_dim, exclusive=True) * mask_1
-  # Remove the elements that don't fit. [batch, group, experts]
-  mask_1 *= mtf.to_float(mtf.less(position_in_expert_1, expert_capacity_f))
-  # [batch, experts]
-  # How many examples in this sequence go to this expert
-  mask_1_count = mtf.reduce_sum(mask_1, reduced_dim=group_dim)
-  # [batch, group] - mostly ones, but zeros where something didn't fit
-  mask_1_flat = mtf.reduce_sum(mask_1, reduced_dim=experts_dim)
-  # [batch, group]
-  position_in_expert_1 = mtf.reduce_sum(
-      position_in_expert_1, reduced_dim=experts_dim)
-  # Weight assigned to first expert.  [batch, group]
-  gate_1 *= mask_1_flat
-
-  # Pick a second-place expert for each position.
-  # We first mask out the experts that we expect to be over-capacity
-  # [batch, experts]
-  space_remaining = expert_capacity_f - mask_1_count
-  use_rate = (mask_1_count + 1.0) / float(group_dim.size)
-  # At what point in the sequence do we expect the expert to be full.
-  # [batch, experts]
-  expected_exhaustion_pos = space_remaining / use_rate
-  # A Tensor with shape [batch, group, experts] representing a boolean
-  #   - whether we expect that the expert will already be full.
-  expected_exhausted = mtf.to_float(mtf.greater(
-      mtf.range(gates.mesh, group_dim, tf.float32), expected_exhaustion_pos))
-  masked_gates = gates - mask_1 - expected_exhausted
-  # This section is similar to the section above.
-  # [batch, group]
-  index_2, gate_2 = mtf.top_1(masked_gates, experts_dim)
-  # [batch, group, experts]
-  mask_2 = mtf.one_hot(index_2, experts_dim, dtype=gates.dtype)
-  # [batch, group, experts]
-  position_in_expert_2 = (
-      mtf.cumsum(mask_2, group_dim, exclusive=True) + mask_1_count)
-  position_in_expert_2 *= mask_2
-  mask_2 *= mtf.to_float(mtf.less(position_in_expert_2, expert_capacity_f))
-  # mask_2_count = mtf.reduce_sum(mask_2, reduced_dim=experts_dim)
-  mask_2_flat = mtf.reduce_sum(mask_2, reduced_dim=experts_dim)
-  position_in_expert_2 = mtf.reduce_sum(
-      position_in_expert_2, reduced_dim=experts_dim)
-  gate_2 *= mask_2_flat
-
-  # renormalize the two gate values to add up to 1
-  denom = gate_1 + gate_2 + 1e-9
-  gate_1 /= denom
-  gate_2 /= denom
-
-  # [batch, group, experts, expert_capacity]
-  assignment = (
-      gate_1 * mask_1_flat
-      * mtf.one_hot(index_1, experts_dim)
-      * mtf.one_hot(mtf.to_int32(position_in_expert_1), expert_capacity_dim) +
-      gate_2 * mask_2_flat
-      * mtf.one_hot(index_2, experts_dim)
-      * mtf.one_hot(mtf.to_int32(position_in_expert_2), expert_capacity_dim))
-
-  return assignment
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
index dc25e1f92..9f2e8a176 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
@@ -25,6 +25,7 @@
 from tensor2tensor.mesh_tensorflow import mtf_beam_search
 from tensor2tensor.mesh_tensorflow import mtf_layers
 from tensor2tensor.mesh_tensorflow import mtf_model
+from tensor2tensor.mesh_tensorflow.research import moe
 from tensor2tensor.utils import registry
 import tensorflow as tf
 
@@ -74,10 +75,6 @@ def kv_dim(self):
   def feedforward_dim(self):
     return mtf.Dimension("d_ff", self._hparams.d_ff)
 
-  @property
-  def experts_dim(self):
-    return mtf.Dimension("experts", self._hparams.moe_num_experts)
-
   @property
   def activation_dtype(self):
     if self._hparams.activation_dtype == "float32":
@@ -292,17 +289,11 @@ def _feedforward_layer(self, x, losses=None):
           x, self.feedforward_dim, dropout=hparams.relu_dropout,
           dropout_broadcast_dims=[self.length_dim])
     elif feedforward_layer == "moe":
-      overhead = (
-          hparams.moe_overhead_train
-          if hparams.mode == tf.estimator.ModeKeys.TRAIN else
-          hparams.moe_overhead_eval)
-      output, loss = mtf_layers.moe_v0(
+      output, loss = moe.transformer_moe_layer_v1(
           x,
-          self.feedforward_dim,
           self.model_dim,
-          self.experts_dim,
-          loss_coef=hparams.moe_loss_coef,
-          overhead=overhead)
+          hparams,
+          hparams.mode == tf.estimator.ModeKeys.TRAIN)
       if losses is not None:
         losses.append(loss)
         return output
@@ -628,12 +619,7 @@ def mtf_transformer_base():
   # round up vocab sizes to be a multiple of this value
   hparams.vocab_divisor = 128
 
-  # mixture of experts hparams
   hparams.add_hparam("feedforward_layer", "dense_relu_dense")
-  hparams.add_hparam("moe_overhead_train", 1.0)
-  hparams.add_hparam("moe_overhead_eval", 2.0)
-  hparams.moe_num_experts = 16
-  hparams.moe_loss_coef = 1e-3
 
   # Use targets_embedding_var * rsqrt(d_model) as softmax_var
   hparams.shared_embedding_and_softmax_weights = True
@@ -682,15 +668,6 @@ def mtf_transformer_single():
   return hparams
 
 
-@registry.register_hparams
-def mtf_transformer_tiny_moe():
-  hparams = mtf_transformer_tiny()
-  hparams.mesh_shape = "all:4"
-  hparams.layout = "batch:all;experts:all"
-  hparams.feedforward_layer = "moe"
-  return hparams
-
-
 @registry.register_hparams
 def mtf_transformer_tiny_8gpu():
   hparams = mtf_transformer_tiny()
@@ -921,24 +898,3 @@ def mtf_transformer_lm_baseline():
   return hparams
 
 
-@registry.register_hparams
-def mtf_transformer_lm_moe():
-  """Mixture of experts language model.
-
-  Run this on 2x2 on languagemodel_lm1b32k_packed for 272000 steps (10 epochs)
-  900M params.
-
-  Results on LM1B:
-         params/10^9  log-ppl(per-token)
-         0.90         3.002
-
-  Returns:
-    a hparams
-  """
-  hparams = mtf_transformer_lm_baseline()
-  hparams.mesh_shape = "all:8"
-  hparams.layout = "batch:all;experts:all"
-  hparams.feedforward_layer = "moe"
-  return hparams
-
-
diff --git a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
index 3ea2b4db6..d976c5d78 100644
--- a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
@@ -12,25 +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.
-"""Experiments with mixture-of-experts architectures.
-
-For all of these architectures, we run on languagemodel_lm1b8k_packed
-for 32k-96 steps (1-3 epochs) on one TPU (8 cores).
-
-All log-perplexities are per-token - multiply by 1.298 for per-word
-
-Results:
-model      params(M)  einsum  alltoall  mxu-util  log-ppl(1ep) (3ep)
-dense_4k   30         3.0e12  0         45%        3.31
-dense_8k   46         4.7e12  0         49%        3.24
-dense_64k                     0                    3.06
-v0         282        4.9e12  5.4e8     35%        3.06
-v0_o75     282        4.0e12  4.0e8     34%
-k_means    282        4.0e12  4.0e8                3.12
-k_means_o2 282        4.9e12  5.4e8     33%
-
-Note: configurations and code are likely to change without notice.
-"""
+"""Experiments with mixture-of-experts architectures."""
 
 
 from __future__ import absolute_import
@@ -38,12 +20,32 @@
 from __future__ import print_function
 
 from tensor2tensor.mesh_tensorflow import mtf_transformer
+from tensor2tensor.mesh_tensorflow.research import moe
 from tensor2tensor.utils import registry
 
 
 @registry.register_hparams
 def xmoe_dense_4k():
-  """Small transformer language model."""
+  """Series of architectural experiments on cheap language models.
+
+  For all of these architectures, we run on languagemodel_lm1b8k_packed
+  for 32k-96 steps (1-3 epochs) on one TPU (8 cores).
+
+  All log-perplexities are per-token - multiply by 1.298 for per-word
+
+  Results:
+  model             params(M)  einsum  alltoall  mxu-util  log-ppl(1ep) (3ep)
+  xmoe_dense_4k     30         3.0e12  0         45%        3.31
+  xmoe_dense_8k     46         4.7e12  0         49%        3.24
+  xmoe_dense_64k                       0                    3.06
+  xmoe_top_2        282        4.0e12  3.4e8     36%
+  xmoe_top_2_c15    282        4.5e12  4.0e8     38%
+
+  Note: configurations and code are likely to change without notice.
+
+  Returns:
+    a hparams
+  """
   hparams = mtf_transformer.mtf_transformer_base()
 
   # The following hparams are constant across all these experiments.
@@ -59,10 +61,6 @@ def xmoe_dense_4k():
   # We will vary the following parameters related to the ffn/moe layers.
   hparams.feedforward_layer = "dense_relu_dense"
   hparams.d_ff = 4096
-  hparams.moe_num_experts = 16
-  hparams.moe_overhead_train = 1.0
-  hparams.moe_overhead_eval = 2.0
-  hparams.moe_loss_coef = 1e-3
   hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
   hparams.mesh_shape = "batch:8"
   return hparams
@@ -84,20 +82,44 @@ def xmoe_dense_64k():
 
 
 @registry.register_hparams
-def xmoe_v0():
+def xmoe_top_2():
   """Mixture of experts."""
   hparams = xmoe_dense_4k()
-  hparams.feedforward_layer = "moe"
+  moe.set_default_moe_hparams(hparams)
   hparams.mesh_shape = "all:8"
   hparams.layout = "batch:all;experts:all"
   return hparams
 
 
 @registry.register_hparams
-def xmoe_v0_o75():
+def xmoe_top_2_c15():
   """Mixture of experts."""
-  hparams = xmoe_v0()
-  hparams.moe_overhead_train = 0.75
+  hparams = xmoe_top_2()
+  hparams.moe_capacity_factor_train = 1.5
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_lm_moe():
+  """Mixture of experts language model.
+
+  Compare to mtf_transformer.mtf_transformer_lm_baseline()
+
+  Run this on 2x2 on languagemodel_lm1b32k_packed for 272000 steps (10 epochs)
+  900M params.
+
+  Results on LM1B:
+         params/10^9  log-ppl(per-token)
+         0.90         TODO(noam): rerun experiment
+
+  Returns:
+    a hparams
+  """
+  hparams = mtf_transformer.mtf_transformer_lm_baseline()
+  moe.set_default_moe_hparams(hparams)
+  hparams.mesh_shape = "all:8"
+  hparams.layout = "batch:all;experts:all"
+  hparams.feedforward_layer = "moe"
   return hparams
 
 
diff --git a/tensor2tensor/mesh_tensorflow/research/moe.py b/tensor2tensor/mesh_tensorflow/research/moe.py
new file mode 100644
index 000000000..c1896ca7b
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/research/moe.py
@@ -0,0 +1,365 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Mixture-of-experts code.
+
+Interfaces and algorithms are under development and subject to rapid change
+without notice.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+from tensor2tensor.mesh_tensorflow import mtf_layers
+import tensorflow as tf
+
+
+
+
+def transformer_moe_layer_v1(inputs, output_dim, hparams, train):
+  """Local mixture of experts that works well on TPU.
+
+  Adapted from the paper https://arxiv.org/abs/1701.06538
+
+  Note: until the algorithm and inferface solidify, we pass in a hyperparameters
+  dictionary in order not to complicate the interface in mtf_transformer.py .
+  Once this code moves out of "research", we should pass the hyperparameters
+  separately.
+
+  Hyperparameters used:
+    hparams.moe_num_experts: number of experts
+    hparams.moe_hidden_size: size of hidden layer in each expert
+    hparams.moe_group_size: size of each "group" for gating purposes
+    hparams.moe_capacity_factor_train: a float
+    hparams.moe_capacity_factor_eval: a float
+    hparams.moe_gating: a string
+    + all hyperparmeters used by _top_2_gating()
+
+  The number of parameters in the gating network is:
+    (input_dim.size * hparams.num_experts) +
+
+  The number of parameters in the experts themselves is:
+    (hparams.num_experts
+     * (input_dim.size + output_dim.size)
+     * hparams.moe_hidden_size)
+
+  The input is n-dimensional: [<batch_and_length_dims>, input_dim], consisting
+  of the representations of all positions in a batch of sequences.
+
+  Each position of each sequence is sent to 0-2 experts.  The expert
+  choices and the combination weights are determined by a learned gating
+  function.
+
+  This function returns a small auxiliary loss that should be added to the
+  training loss of the model.  This loss helps to balance expert usage.
+  Without the loss, it is very likely that a few experts will be trained and
+  the rest will starve.
+
+  Several hacks are necessary to get around current TPU limitations:
+
+  - To ensure static shapes, we enforce (by truncation/padding)
+    that each sequence send the same number of elements to each expert.
+
+    It would make more sense to enforce this equality over the entire batch,
+    but due to our hacked-up gather-by-matmul implementation, we need to divide
+    the batch into "groups".  For each group, the same number of elements
+    are sent to each expert.
+
+  TODO(noam): Factor this code better.  We want to be able to substitute
+  different code for the experts themselves.
+
+  Args:
+    inputs: a mtf.Tensor with shape [<batch_dims...>, length_dim, input_dim]
+    output_dim: a mtf.Dimension (for Transformer, this is input_dim)
+    hparams: model hyperparameters
+    train: a boolean
+
+  Returns:
+    outputs: a Tensor with shape [<batch_dims...>, length_dim, output_dim]
+    loss: a mtf scalar
+
+  Raises:
+    ValueError: on unrecognized hparams.moe_gating
+  """
+  orig_inputs = inputs
+  input_dim = inputs.shape.dims[-1]
+  hidden_dim = mtf.Dimension("expert_hidden", hparams.moe_hidden_size)
+  experts_dim = mtf.Dimension("experts", hparams.moe_num_experts)
+  group_dim = mtf.Dimension("group", hparams.moe_group_size)
+  batch_dim = mtf.Dimension(
+      orig_inputs.shape[0].name,
+      orig_inputs.shape.size // (group_dim.size * input_dim.size))
+  inputs = mtf.reshape(inputs, [batch_dim, group_dim, input_dim])
+
+  # Each sequence sends expert_capacity positions to each expert.
+  capacity_factor = (
+      hparams.moe_capacity_factor_train if train else
+      hparams.moe_capacity_factor_eval)
+  expert_capacity = min(
+      group_dim.size,
+      int((group_dim.size * capacity_factor) / experts_dim.size))
+  expert_capacity_dim = mtf.Dimension("expert_capacity", expert_capacity)
+
+  experts_dim_unsplit = mtf.Dimension("expert_unsplit", experts_dim.size)
+  batch_dim_unsplit = mtf.Dimension("batch_unsplit", batch_dim.size)
+
+  if hparams.moe_gating == "top_2":
+    forward_assignment, backward_assignment, loss = _top_2_gating(
+        inputs=inputs,
+        experts_dim=experts_dim_unsplit,
+        expert_capacity_dim=expert_capacity_dim,
+        max_experts=None,
+        hparams=hparams,
+        train=train)
+  else:
+    raise ValueError("unknown hparams.moe_gating=%s" % hparams.moe_gating)
+
+  # put num_experts dimension first to make split easier in alltoall
+  expert_inputs = mtf.einsum([inputs, forward_assignment], mtf.Shape(
+      [experts_dim_unsplit, batch_dim, expert_capacity_dim, input_dim]))
+
+  expert_inputs = mtf.reshape(expert_inputs, mtf.Shape(
+      [experts_dim, batch_dim_unsplit, expert_capacity_dim, input_dim]))
+
+  # Now feed the expert inputs through the experts.
+  h = mtf_layers.dense(
+      expert_inputs, hidden_dim, expert_dims=[experts_dim],
+      activation=mtf.relu, use_bias=False, name="x0")
+  expert_output = mtf_layers.dense(
+      h, output_dim, expert_dims=[experts_dim], use_bias=False, name="x1")
+
+  expert_output = mtf.reshape(expert_output, mtf.Shape(
+      [experts_dim_unsplit, batch_dim, expert_capacity_dim, input_dim]))
+
+  output = mtf.einsum([expert_output, backward_assignment], mtf.Shape(
+      [batch_dim, group_dim, output_dim]))
+
+  output = mtf.reshape(output, orig_inputs.shape.dims[:-1] + [output_dim])
+
+  return output, loss * hparams.moe_loss_coef
+
+
+def _top_2_gating(
+    inputs, experts_dim, expert_capacity_dim, max_experts, hparams, train):
+  """Compute gating for mixture-of-experts in TensorFlow.
+
+  Note: until the algorithm and inferface solidify, we pass in a hyperparameters
+  dictionary in order not to complicate the interface in mtf_transformer.py .
+  Once this code moves out of "research", we should pass the hyperparameters
+  separately.
+
+  Hyperparameters used:
+    hparams.moe_use_second_place_loss: a boolean
+    hparams.moe_second_policy_train: a string
+    hparams.moe_second_policy_eval: a string
+    hparams.moe_second_threshold: a float
+
+  max_experts is an float tensor with shape [batch_dim, group_dim]
+  indicating at most how many experts to use per example.  This can be
+  used to prevent padding from going to experts.
+
+  The returned forward assignment is a tensor used to map (via einsum) from the
+  inputs to the expert_inputs.  Likewise, the returned backward_assignment is
+  used to map (via einsum) from the expert outputs to the outputs.  Both the
+  forward and backward assignments are mostly zeros.  The shapes of all of these
+  are as follows.
+
+  inputs: [batch_dim, group_dim, input_dim]
+  forward_assignment: [batch_dim, group_dim, experts_dim, expert_capacity_dim]
+  expert_inputs: [batch_dim, experts_dim, expert_capacity_dim, input_dim]
+
+  expert_outputs: [batch_dim, experts_dim, expert_capacity_dim, output_dim]
+  backward_assignment: [batch_dim, group_dim, experts_dim, expert_capacity_dim]
+  outputs: [batch_dim, group_dim, output_dim]
+
+  Args:
+    inputs: a mtf.Tensor with shape [batch_dim, group_dim, input_dim]
+    experts_dim: a Dimension (the number of experts)
+    expert_capacity_dim: a Dimension (number of examples per group per expert)
+    max_experts: optional mtf.Tensor with shape [batch_dim, group_dim]
+    hparams: model hyperparameters.
+    train: a boolean
+
+  Returns:
+    forward_assignment: a Tensor with shape
+      [batch_dim, group_dim, experts_dim, expert_capacity_dim]
+    backward_assignment: a Tensor with shape
+      [batch_dim, group_dim, experts_dim, expert_capacity_dim]
+    loss: a mtf scalar
+
+  Raises:
+    ValueError: on illegal hyperparameters
+  """
+  unused_batch_dim, group_dim, unused_input_dim = inputs.shape.dims
+
+  raw_gates = mtf.softmax(mtf_layers.dense(
+      inputs, experts_dim, use_bias=False), experts_dim)
+
+  expert_capacity_f = float(expert_capacity_dim.size)
+
+  # FIND TOP 2 EXPERTS PER POSITON
+  # Find the top expert for each position. shape=[batch, group]
+  index_1, gate_1 = mtf.top_1(raw_gates, experts_dim)
+  # [batch, group, experts]
+  mask_1 = mtf.one_hot(index_1, experts_dim, dtype=raw_gates.dtype)
+  gates_without_top_1 = raw_gates * (1.0 - mask_1)
+  # [batch, group]
+  index_2, gate_2 = mtf.top_1(gates_without_top_1, experts_dim)
+  # [batch, group, experts]
+  mask_2 = mtf.one_hot(index_2, experts_dim, dtype=raw_gates.dtype)
+
+  if max_experts is not None:
+    geq1 = mtf.to_float(mtf.greater_equal(max_experts, 1.0))
+    geq2 = mtf.to_float(mtf.greater_equal(max_experts, 2.0))
+    mask_1 *= geq1
+    mask_2 *= geq2
+    raw_gates *= geq1
+    gates_without_top_1 *= geq2
+
+  # BALANCING LOSSES
+  # shape = [batch, experts]
+  # We want to equalize the fraction of the batch assigned to each expert
+  density_1 = mtf.reduce_mean(mask_1, reduced_dim=group_dim)
+  # Something continuous that is correlated with what we want to equalize.
+  density_1_proxy = mtf.reduce_mean(raw_gates, reduced_dim=group_dim)
+  density_1 = mtf.Print(
+      density_1, [mtf.reduce_mean(density_1, output_shape=[experts_dim])],
+      "density_1", summarize=1000)
+  loss = (mtf.reduce_mean(density_1_proxy * density_1)
+          * float(experts_dim.size * experts_dim.size))
+
+  if hparams.moe_use_second_place_loss:
+    # Also add a loss to encourage all experts to be used equally also as the
+    # second-place expert.  Experimentally, this seems to be a wash.
+    # We want to equalize the fraction of the batch assigned to each expert:
+    density_2 = mtf.reduce_mean(mask_2, reduced_dim=group_dim)
+    # As a proxy for density_2, we renormalize the raw gates after the top one
+    # has been removed.
+    normalized = gates_without_top_1 / (
+        mtf.reduce_sum(gates_without_top_1, reduced_dim=experts_dim) + 1e-9)
+    density_2_proxy = mtf.reduce_mean(normalized, reduced_dim=group_dim)
+    loss_2 = (mtf.reduce_mean(density_2_proxy * density_2)
+              * float(experts_dim.size * experts_dim.size))
+    loss += loss_2 * 0.5
+
+  # Depending on the policy in the hparams, we may drop out some of the
+  # second-place experts.
+  policy = (
+      hparams.moe_second_policy_train if train else
+      hparams.moe_second_policy_eval)
+  threshold = (
+      hparams.moe_second_threshold_train if train else
+      hparams.moe_second_threshold_eval)
+  if policy == "all":
+    # Use second-place experts for all examples.
+    pass
+  elif policy == "none":
+    # Never use second-place experts for all examples.
+    mask_2 = mtf.zeros_like(mask_2)
+  elif policy == "threshold":
+    # Use second-place experts if gate_2 > threshold.
+    mask_2 *= mtf.to_float(mtf.greater(gate_2, threshold))
+  elif policy == "random":
+    # Use second-place experts with probablity min(1.0, gate_2 / threshold).
+    mask_2 *= mtf.to_float(
+        mtf.less(mtf.random_uniform(gate_2.mesh, gate_2.shape),
+                 gate_2 / max(threshold, 1e-9)))
+  else:
+    raise ValueError("Unknown policy %s" % policy)
+  mask_2 = mtf.Print(
+      mask_2, [mtf.reduce_mean(mask_2, output_shape=[experts_dim])],
+      "density_2", summarize=1000)
+
+  # COMPUTE ASSIGNMENT TO EXPERTS
+  # [batch, group, experts]
+  # This is the position within the expert's mini-batch for this sequence
+  position_in_expert_1 = mtf.cumsum(mask_1, group_dim, exclusive=True) * mask_1
+  # Remove the elements that don't fit. [batch, group, experts]
+  mask_1 *= mtf.to_float(mtf.less(position_in_expert_1, expert_capacity_f))
+  # [batch, experts]
+  # How many examples in this sequence go to this expert
+  mask_1_count = mtf.reduce_sum(mask_1, reduced_dim=group_dim)
+  # [batch, group] - mostly ones, but zeros where something didn't fit
+  mask_1_flat = mtf.reduce_sum(mask_1, reduced_dim=experts_dim)
+  # [batch, group]
+  position_in_expert_1 = mtf.reduce_sum(
+      position_in_expert_1, reduced_dim=experts_dim)
+  # Weight assigned to first expert.  [batch, group]
+  gate_1 *= mask_1_flat
+
+  # [batch, group, experts]
+  position_in_expert_2 = (
+      mtf.cumsum(mask_2, group_dim, exclusive=True) + mask_1_count)
+  position_in_expert_2 *= mask_2
+  mask_2 *= mtf.to_float(mtf.less(position_in_expert_2, expert_capacity_f))
+  # mask_2_count = mtf.reduce_sum(mask_2, reduced_dim=experts_dim)
+  mask_2_flat = mtf.reduce_sum(mask_2, reduced_dim=experts_dim)
+  gate_2 *= mask_2_flat
+  position_in_expert_2 = mtf.reduce_sum(
+      position_in_expert_2, reduced_dim=experts_dim)
+
+  # renormalize the two gate values to add up to 1
+  denom = gate_1 + gate_2 + 1e-9
+  gate_1 /= denom
+  gate_2 /= denom
+
+  # [batch, group, experts, expert_capacity]
+  backward_assignment = (
+      gate_1 * mask_1_flat
+      * mtf.one_hot(index_1, experts_dim)
+      * mtf.one_hot(mtf.to_int32(position_in_expert_1), expert_capacity_dim) +
+      gate_2 * mask_2_flat
+      * mtf.one_hot(index_2, experts_dim)
+      * mtf.one_hot(mtf.to_int32(position_in_expert_2), expert_capacity_dim))
+
+  forward_assignment = mtf.cast(
+      mtf.cast(backward_assignment, tf.bool), backward_assignment.dtype)
+
+  return forward_assignment, backward_assignment, loss
+
+
+def set_default_moe_hparams(hparams):
+  """Add necessary hyperparameters for mixture-of-experts."""
+  hparams.feedforward_layer = "moe"
+  hparams.moe_num_experts = 16
+  hparams.moe_loss_coef = 1e-2
+  hparams.add_hparam("moe_gating", "top_2")
+  # Experts have fixed capacity per batch.  We need some extra capacity
+  # in case gating is not perfectly balanced.
+  # moe_capacity_factor_* should be set to a value >=1.
+  hparams.add_hparam("moe_capacity_factor_train", 1.25)
+  hparams.add_hparam("moe_capacity_factor_eval", 2.0)
+  # Each expert has a hidden layer with this size.
+  hparams.add_hparam("moe_hidden_size", 4096)
+  # For gating, divide inputs into groups of this size before gating.
+  # Each group sends the same number of inputs to each expert.
+  # Ideally, the group size would be the whole batch, but this is expensive
+  # due to our use of matrix multiplication for reordering.
+  hparams.add_hparam("moe_group_size", 1024)
+  # For top_2 gating, whether to impose an additional loss in order to make
+  # the experts equally used as the second-place expert.
+  hparams.add_hparam("moe_use_second_place_loss", 0)
+  # In top_2 gating, policy for whether to use a second-place expert.
+  # Legal values are:
+  #    "all": always
+  #    "none": never
+  #    "threshold": if gate value > the given threshold
+  #    "random": if gate value > threshold*random_uniform(0,1)
+  hparams.add_hparam("moe_second_policy_train", "random")
+  hparams.add_hparam("moe_second_policy_eval", "random")
+  hparams.add_hparam("moe_second_threshold_train", 0.2)
+  hparams.add_hparam("moe_second_threshold_eval", 0.2)

From 5f26b79ffd5612bc5c9534def7c1d585866b0770 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 21 Aug 2018 11:14:01 -0700
Subject: [PATCH 0626/2720] Internal change

PiperOrigin-RevId: 209626961
---
 tensor2tensor/layers/discretization.py        |   3 +-
 tensor2tensor/layers/latent_layers.py         |   6 +-
 tensor2tensor/models/research/autoencoders.py | 256 +++++++++++-------
 tensor2tensor/utils/expert_utils.py           |   2 -
 4 files changed, 168 insertions(+), 99 deletions(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index b40050767..3188855cc 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -819,8 +819,7 @@ def vq_nearest_neighbor(x, means,
     if temperature is None:
       x_means_idx = tf.argmax(-dist, axis=-1)
     else:
-      sm_dist = tf.nn.softmax(-dist)
-      x_means_idx = tf.multinomial(sm_dist / temperature, 1)
+      x_means_idx = tf.multinomial(- dist / temperature, 1)
       x_means_idx = tf.squeeze(x_means_idx, axis=-1)
     if (common_layers.should_generate_summaries() and
         not common_layers.is_xla_compiled()):
diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 00c16a358..2e66ed3c7 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -46,7 +46,8 @@ def compress_self_attention_layer(x, hparams, name):
     return tf.reshape(res, xshape)
 
 
-def multinomial_sample(x, vocab_size, sampling_method, temperature):
+def multinomial_sample(x, vocab_size=None, sampling_method="random",
+                       temperature=1.0):
   """Multinomial sampling from a n-dimensional tensor.
 
   Args:
@@ -58,7 +59,8 @@ def multinomial_sample(x, vocab_size, sampling_method, temperature):
   Returns:
     Tensor of shape [...].
   """
-  if sampling_method == "random":
+  vocab_size = vocab_size or common_layers.shape_list(x)[-1]
+  if sampling_method == "random" and temperature > 0.0:
     samples = tf.multinomial(tf.reshape(x, [-1, vocab_size]) / temperature, 1)
   else:
     samples = tf.argmax(x, axis=-1)
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 118b3bf35..cde25ba12 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -22,6 +22,7 @@
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
+from tensor2tensor.layers import latent_layers
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -76,6 +77,7 @@ def embed(self, x):
           name="embed",
           activation=common_layers.belu,
           bias_initializer=tf.random_normal_initializer(stddev=0.01))
+      x = common_layers.layer_norm(x, name="ln_embed")
       return common_attention.add_timing_signal_nd(x)
 
   def bottleneck(self, x):
@@ -144,10 +146,12 @@ def make_even_size(self, x):
   def encoder(self, x):
     with tf.variable_scope("encoder"):
       hparams = self.hparams
+      layers = []
       kernel, strides = self._get_kernel_and_strides()
       # Down-convolutions.
       for i in range(hparams.num_hidden_layers):
         x = self.make_even_size(x)
+        layers.append(x)
         x = tf.layers.conv2d(
             x,
             hparams.hidden_size * 2**(i + 1),
@@ -156,10 +160,11 @@ def encoder(self, x):
             padding="SAME",
             activation=common_layers.belu,
             name="conv_%d" % i)
-        x = common_layers.layer_norm(x)
-      return x
+        x = common_layers.layer_norm(x, name="ln_%d" % i)
+      return x, layers
 
-  def decoder(self, x):
+  def decoder(self, x, encoder_layers):
+    del encoder_layers
     with tf.variable_scope("decoder"):
       hparams = self.hparams
       kernel, strides = self._get_kernel_and_strides()
@@ -174,7 +179,7 @@ def decoder(self, x):
             padding="SAME",
             activation=common_layers.belu,
             name="deconv_%d" % j)
-        x = common_layers.layer_norm(x)
+        x = common_layers.layer_norm(x, name="ln_%d" % i)
       return x
 
   def body(self, features):
@@ -186,28 +191,45 @@ def body(self, features):
       shape = common_layers.shape_list(labels)
       x = tf.one_hot(labels, vocab_size)
       x = self.embed(x)
+      target_codes = x
       is1d = shape[2] == 1
       self.is1d = is1d
       # Run encoder.
-      x = self.encoder(x)
-      # Bottleneck (mix during early training, not too important but stable).
+      x, encoder_layers = self.encoder(x)
+      # Bottleneck.
       b, b_loss = self.bottleneck(x)
+      xb_loss = 0.0
       b_shape = common_layers.shape_list(b)
       self._cur_bottleneck_tensor = b
       b = self.unbottleneck(b, common_layers.shape_list(x)[-1])
-      b = common_layers.mix(b, x, hparams.bottleneck_warmup_steps, is_training)
+      if not is_training:
+        x = b
+      else:
+        l = 2**hparams.num_hidden_layers
+        warm_step = int(hparams.bottleneck_warmup_steps * 0.25 * l)
+        nomix_p = common_layers.inverse_lin_decay(warm_step) + 0.01
+        if common_layers.should_generate_summaries():
+          tf.summary.scalar("nomix_p_bottleneck", nomix_p)
+        rand = tf.random_uniform(common_layers.shape_list(x))
+        # This is the distance between b and x. Having this as loss helps learn
+        # the bottleneck function, but if we back-propagated to x it would be
+        # minimized by just setting x=0 and b=0 -- so we don't want too much
+        # of the influence of this, and we stop-gradient to not zero-out x.
+        x_stop = tf.stop_gradient(x)
+        xb_loss = tf.reduce_mean(tf.reduce_sum(tf.square(x_stop - b), axis=-1))
+        # To prevent this loss from exploding we clip at 1, but anneal clipping.
+        clip_max = 1.0 / common_layers.inverse_exp_decay(
+            warm_step, min_value=0.001)
+        xb_clip = tf.maximum(tf.stop_gradient(xb_loss), clip_max)
+        xb_loss *= clip_max / xb_clip
+        x = tf.where(tf.less(rand, nomix_p), b, x)
       if hparams.gan_loss_factor != 0.0:
         # Add a purely sampled batch on which we'll compute the GAN loss.
         g = self.unbottleneck(
             self.sample(shape=b_shape),
             common_layers.shape_list(x)[-1], reuse=True)
-        b = tf.concat([g, b], axis=0)
-      # With probability bottleneck_max_prob use the bottleneck, otherwise x.
-      if hparams.bottleneck_max_prob < -1.0:
-        x = tf.where(
-            tf.less(tf.random_uniform([]), hparams.bottleneck_max_prob), b, x)
-      else:
-        x = b
+        x = tf.concat([g, x], axis=0)
+        encoder_layers = [tf.concat([l, l], axis=0) for l in encoder_layers]
     else:
       if self._cur_bottleneck_tensor is None:
         b = self.sample()
@@ -217,7 +239,7 @@ def body(self, features):
       res_size = min(res_size, hparams.max_hidden_size)
       x = self.unbottleneck(b, res_size)
     # Run decoder.
-    x = self.decoder(x)
+    x = self.decoder(x, encoder_layers)
     if hparams.mode == tf.estimator.ModeKeys.PREDICT:
       return x, {"bottleneck_loss": 0.0}
     # Cut to the right size and mix before returning.
@@ -236,31 +258,60 @@ def body(self, features):
       res_gan, res = tf.split(res, 2, axis=0)
 
     # Losses.
-    losses = {}
-    vq_temperature = 0.001 / common_layers.inverse_exp_decay(
-        hparams.gan_codes_warmup_steps * 1.2, min_value=0.002)
-    if hparams.mode != tf.estimator.ModeKeys.TRAIN:
-      vq_temperature = None
-
-    with tf.variable_scope("vq_loss"):
-      (reconstr, _, target_codes, code_loss,
-       targets_loss) = discretization.vq_loss(
-           res, labels, vocab_size, temperature=vq_temperature)
-
-    losses["code_loss"] = code_loss * hparams.code_loss_factor
-    losses["training"] = targets_loss
+    losses = {"bottleneck_extra": b_loss,
+              "bottleneck_l2": hparams.bottleneck_l2_factor * xb_loss}
+
+    if hparams.use_vq_loss:
+      vq_temperature = hparams.vq_temperature / common_layers.inverse_exp_decay(
+          hparams.gan_codes_warmup_steps * 1.2,
+          min_value=hparams.vq_temperature * 2)
+      if hparams.mode != tf.estimator.ModeKeys.TRAIN:
+        vq_temperature = None
+      with tf.variable_scope("vq_loss"):
+        (reconstr, _, target_codes, code_loss,
+         targets_loss) = discretization.vq_loss(
+             res, labels, vocab_size, temperature=vq_temperature)
+      losses["code_loss"] = code_loss * hparams.code_loss_factor
+      losses["training"] = targets_loss
+    else:
+      reconstr = tf.layers.dense(res, vocab_size, name="autoencoder_final")
+      targets_loss = tf.losses.sparse_softmax_cross_entropy(
+          logits=reconstr, labels=labels)
+      losses["training"] = targets_loss
 
     # GAN losses.
     if hparams.gan_loss_factor != 0.0:
-      with tf.variable_scope("vq_loss", reuse=True):
-        update_means_factor = common_layers.inverse_exp_decay(
-            hparams.gan_codes_warmup_steps, min_value=0.0001)
-        update_means = tf.less(tf.random_uniform([]), update_means_factor)
-        reconstr_gan, gan_codes, _, code_loss_gan, _ = discretization.vq_loss(
-            res_gan, labels, vocab_size, do_update=update_means,
-            temperature=vq_temperature)
-        code_loss_gan *= hparams.code_loss_factor * update_means_factor
-        losses["code_loss_gan"] = code_loss_gan
+      update_means_factor = common_layers.inverse_exp_decay(
+          hparams.gan_codes_warmup_steps, min_value=0.0001)
+      if hparams.use_vq_loss:
+        with tf.variable_scope("vq_loss", reuse=True):
+          update_means = tf.less(tf.random_uniform([]), update_means_factor)
+          reconstr_gan, gan_codes, _, code_loss_gan, _ = discretization.vq_loss(
+              res_gan, labels, vocab_size, do_update=update_means,
+              temperature=vq_temperature)
+          code_loss_gan *= hparams.code_loss_factor * update_means_factor
+          losses["code_loss_gan"] = code_loss_gan
+      else:
+        reconstr_gan = tf.layers.dense(
+            res_gan, vocab_size, name="autoencoder_final", reuse=True)
+        reconstr_gan = tf.nn.log_softmax(reconstr_gan)
+        if is_training and hparams.gumbel_temperature > 0.0:
+          gumbel_samples = discretization.gumbel_sample(
+              common_layers.shape_list(reconstr_gan))
+          gumbel_samples *= hparams.gumbel_noise_factor
+          reconstr_gan += gumbel_samples
+          reconstr_sample = latent_layers.multinomial_sample(
+              reconstr_gan, temperature=hparams.gumbel_temperature)
+          reconstr_gan = tf.nn.softmax(
+              reconstr_gan / hparams.gumbel_temperature)
+        else:
+          reconstr_sample = tf.argmax(reconstr_gan, axis=-1)
+          reconstr_gan = tf.nn.softmax(reconstr_gan / 0.1)  # Sharpen a bit.
+        # Use 1-hot forward, softmax backward.
+        reconstr_hot = tf.one_hot(reconstr_sample, vocab_size)
+        reconstr_gan += reconstr_hot - tf.stop_gradient(reconstr_gan)
+        # Embed to codes.
+        gan_codes = self.embed(reconstr_gan)
 
     # Add GAN loss if requested.
     gan_loss = 0.0
@@ -283,12 +334,9 @@ def discriminate(x):
           target_codes, rev_grad_gan_codes, discriminate,
           self.hparams.num_sliced_vecs)
       gan_loss *= hparams.gan_loss_factor  * update_means_factor
+      losses["gan_loss"] = -gan_loss
 
     self.image_summary("ae", reconstr)
-
-    losses["b_loss"] = b_loss
-    losses["gan_loss"] = -gan_loss
-
     logits = reconstr
     return logits, losses
 
@@ -386,7 +434,6 @@ def body(self, features):
     basic_result = self.embed(basic_result)
     shape = common_layers.shape_list(basic_result)
     basic1d = tf.reshape(basic_result, [shape[0], -1, shape[-1]])
-    vocab_size = self._problem_hparams.target_modality.top_dimensionality
     targets = tf.one_hot(features["targets_raw"], vocab_size)
     targets = self.embed(targets)
     targets = tf.reshape(targets, common_layers.shape_list(basic_result))
@@ -512,22 +559,19 @@ class AutoencoderResidual(AutoencoderAutoregressive):
   """Residual autoencoder."""
 
   def dropout(self, x):
-    if self.hparams.dropout <= 0.0:
-      return x
-    # For simple dropout just do this:
-    # return tf.nn.dropout(x, 1.0 - self.hparams.dropout)
     is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
-    return common_layers.mix(
-        tf.zeros_like(x),
-        x,
-        self.hparams.bottleneck_warmup_steps,
-        is_training,
-        max_prob=self.hparams.dropout,
-        broadcast_last=True)
+    hparams = self.hparams
+    if hparams.dropout <= 0.0 or not is_training:
+      return x
+    warm_step = hparams.bottleneck_warmup_steps * 2**hparams.num_hidden_layers
+    dropout = common_layers.inverse_lin_decay(warm_step // 2) * hparams.dropout
+    return common_layers.dropout_with_broadcast_dims(
+        x, 1.0 - dropout, broadcast_dims=[-1])
 
   def encoder(self, x):
     with tf.variable_scope("encoder"):
       hparams = self.hparams
+      layers = []
       kernel, strides = self._get_kernel_and_strides()
       residual_kernel = (hparams.residual_kernel_height,
                          hparams.residual_kernel_width)
@@ -540,6 +584,7 @@ def encoder(self, x):
       for i in range(hparams.num_hidden_layers):
         with tf.variable_scope("layer_%d" % i):
           x = self.make_even_size(x)
+          layers.append(x)
           x = self.dropout(x)
           filters = hparams.hidden_size * 2**(i + 1)
           filters = min(filters, hparams.max_hidden_size)
@@ -565,12 +610,13 @@ def encoder(self, x):
                 activation=common_layers.belu,
                 name="residual_%d" % r)
           x += tf.nn.dropout(y, 1.0 - hparams.residual_dropout)
-          x = common_layers.layer_norm(x)
-      return x
+          x = common_layers.layer_norm(x, name="ln")
+      return x, layers
 
-  def decoder(self, x):
+  def decoder(self, x, encoder_layers=None):
     with tf.variable_scope("decoder"):
       hparams = self.hparams
+      is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
       kernel, strides = self._get_kernel_and_strides()
       residual_kernel = (hparams.residual_kernel_height,
                          hparams.residual_kernel_width)
@@ -582,11 +628,14 @@ def decoder(self, x):
       # Up-convolutions.
       for i in range(hparams.num_hidden_layers):
         j = hparams.num_hidden_layers - i - 1
+        nomix_p = common_layers.inverse_lin_decay(
+            int(hparams.bottleneck_warmup_steps * 0.25 * 2**j)) + 0.01
+        if common_layers.should_generate_summaries():
+          tf.summary.scalar("nomix_p_%d" % j, nomix_p)
         filters = hparams.hidden_size * 2**j
         filters = min(filters, hparams.max_hidden_size)
         with tf.variable_scope("layer_%d" % i):
           j = hparams.num_hidden_layers - i - 1
-          filters = hparams.hidden_size * 2**j
           x = tf.layers.conv2d_transpose(
               x,
               filters,
@@ -609,8 +658,15 @@ def decoder(self, x):
                 activation=common_layers.belu,
                 name="residual_%d" % r)
           x += tf.nn.dropout(y, 1.0 - hparams.residual_dropout)
-          x = common_layers.layer_norm(x)
+          x = common_layers.layer_norm(x, name="ln")
           x = common_attention.add_timing_signal_nd(x)
+          if encoder_layers is not None:
+            enc_x = encoder_layers[j]
+            enc_shape = common_layers.shape_list(enc_x)
+            x = x[:, :enc_shape[1], :enc_shape[2], :]
+            if is_training:  # Mix at the beginning of training.
+              rand = tf.random_uniform(common_layers.shape_list(x))
+              x = tf.where(tf.less(rand, nomix_p), x, enc_x)
       return x
 
 
@@ -787,12 +843,7 @@ def body(self, features):
                           is_training, num_stacks - 1)
       b = self.unbottleneck(b, x_size)
       b = common_layers.mix(b, x, hparams.bottleneck_warmup_steps, is_training)
-      # With probability bottleneck_max_prob use the bottleneck, otherwise x.
-      if hparams.bottleneck_max_prob < 1.0:
-        x = tf.where(
-            tf.less(tf.random_uniform([]), hparams.bottleneck_max_prob), b, x)
-      else:
-        x = b
+      x = b
     else:
       b = self.sample()
       res_size = self.hparams.hidden_size * 2**self.hparams.num_hidden_layers
@@ -831,16 +882,19 @@ def autoencoder_basic():
   hparams.add_hparam("max_hidden_size", 1024)
   hparams.add_hparam("bottleneck_bits", 128)
   hparams.add_hparam("bottleneck_noise", 0.1)
-  hparams.add_hparam("bottleneck_warmup_steps", 3000)
-  hparams.add_hparam("bottleneck_max_prob", 1.0)
+  hparams.add_hparam("bottleneck_warmup_steps", 500)
   hparams.add_hparam("sample_height", 32)
   hparams.add_hparam("sample_width", 32)
   hparams.add_hparam("discriminator_batchnorm", True)
   hparams.add_hparam("num_sliced_vecs", 4096)
   hparams.add_hparam("code_loss_factor", 1.0)
-  hparams.add_hparam("gan_codes_warmup_steps", 5000)
+  hparams.add_hparam("gan_codes_warmup_steps", 6000)
   hparams.add_hparam("gan_loss_factor", 0.0)
-  hparams.add_hparam("use_vqloss", False)
+  hparams.add_hparam("bottleneck_l2_factor", 0.05)
+  hparams.add_hparam("gumbel_temperature", 0.05)
+  hparams.add_hparam("gumbel_noise_factor", 0.2)
+  hparams.add_hparam("vq_temperature", 0.001)
+  hparams.add_hparam("use_vq_loss", 0)
   return hparams
 
 
@@ -865,7 +919,6 @@ def autoencoder_residual():
   hparams.learning_rate_constant = 0.5
   hparams.learning_rate_warmup_steps = 500
   hparams.learning_rate_schedule = "constant * linear_warmup * rsqrt_decay"
-  hparams.dropout = 0.05
   hparams.num_hidden_layers = 5
   hparams.hidden_size = 64
   hparams.max_hidden_size = 1024
@@ -878,6 +931,21 @@ def autoencoder_residual():
   return hparams
 
 
+@registry.register_hparams
+def autoencoder_residual_text():
+  """Residual autoencoder model for text."""
+  hparams = autoencoder_residual()
+  hparams.bottleneck_bits = 32
+  hparams.batch_size = 1024
+  hparams.hidden_size = 64
+  hparams.max_hidden_size = 512
+  hparams.bottleneck_noise = 0.0
+  hparams.target_modality = "symbol:identity"
+  hparams.input_modalities = "symbol:identity"
+  hparams.autoregressive_mode = "none"
+  return hparams
+
+
 @registry.register_hparams
 def autoencoder_basic_discrete():
   """Basic autoencoder model."""
@@ -886,7 +954,6 @@ def autoencoder_basic_discrete():
   hparams.hidden_size = 64
   hparams.bottleneck_bits = 4096
   hparams.bottleneck_noise = 0.1
-  hparams.bottleneck_warmup_steps = 3000
   hparams.add_hparam("discretize_warmup_steps", 5000)
   return hparams
 
@@ -897,7 +964,6 @@ def autoencoder_residual_discrete():
   hparams = autoencoder_residual()
   hparams.bottleneck_bits = 4096
   hparams.bottleneck_noise = 0.1
-  hparams.bottleneck_warmup_steps = 3000
   hparams.add_hparam("discretize_warmup_steps", 5000)
   hparams.add_hparam("bottleneck_kind", "tanh_discrete")
   hparams.add_hparam("isemhash_noise_dev", 0.5)
@@ -916,7 +982,6 @@ def autoencoder_residual_discrete_big():
   hparams.hidden_size = 128
   hparams.max_hidden_size = 4096
   hparams.bottleneck_noise = 0.1
-  hparams.dropout = 0.1
   hparams.residual_dropout = 0.4
   return hparams
 
@@ -925,22 +990,10 @@ def autoencoder_residual_discrete_big():
 def autoencoder_ordered_discrete():
   """Ordered discrete autoencoder model."""
   hparams = autoencoder_residual_discrete()
-  hparams.bottleneck_noise = 1.0
-  hparams.gan_loss_factor = 0.0
-  hparams.dropout = 0.1
-  hparams.residual_dropout = 0.3
-  hparams.use_vqloss = True
+  hparams.bottleneck_noise = 0.8
+  hparams.gan_loss_factor = 0.02
+  hparams.use_vq_loss = 0
   hparams.add_hparam("unordered", False)
-
-  return hparams
-
-
-@registry.register_hparams
-def autoencoder_ordered_discrete_novq():
-  """Ordered discrete autoencoder model."""
-  hparams = autoencoder_ordered_discrete()
-  hparams.use_vqloss = False
-
   return hparams
 
 
@@ -956,20 +1009,28 @@ def autoencoder_ordered_discrete_hs256():
 def autoencoder_ordered_text():
   """Ordered discrete autoencoder model for text."""
   hparams = autoencoder_ordered_discrete()
-  hparams.learning_rate_constant = 2.0
-  hparams.learning_rate_warmup_steps = 2000
   hparams.bottleneck_bits = 1024
   hparams.batch_size = 1024
   hparams.autoregressive_mode = "sru"
   hparams.hidden_size = 256
   hparams.max_hidden_size = 4096
-  hparams.bottleneck_warmup_steps = 10000
-  hparams.discretize_warmup_steps = 15000
   hparams.target_modality = "symbol:identity"
   hparams.input_modalities = "symbol:identity"
   return hparams
 
 
+@registry.register_hparams
+def autoencoder_ordered_text_small():
+  """Ordered discrete autoencoder model for text, small version."""
+  hparams = autoencoder_ordered_text()
+  hparams.bottleneck_bits = 64
+  hparams.hidden_size = 64
+  hparams.max_hidden_size = 512
+  hparams.bottleneck_noise = 0.0
+  hparams.autoregressive_mode = "none"
+  return hparams
+
+
 @registry.register_hparams
 def autoencoder_ordered_discrete_vq():
   """Ordered discrete autoencoder model with VQ bottleneck."""
@@ -985,7 +1046,6 @@ def autoencoder_discrete_pong():
   hparams = autoencoder_ordered_discrete()
   hparams.num_hidden_layers = 2
   hparams.bottleneck_bits = 24
-  hparams.dropout = 0.1
   hparams.batch_size = 2
   hparams.bottleneck_noise = 0.2
   hparams.max_hidden_size = 1024
@@ -1005,10 +1065,20 @@ def autoencoder_discrete_cifar():
   hparams.num_residual_layers = 4
   hparams.batch_size = 32
   hparams.learning_rate_constant = 1.0
-  hparams.dropout = 0.1
   return hparams
 
 
+@registry.register_ranged_hparams
+def autoencoder_range(rhp):
+  """Tuning grid of the main autoencoder params."""
+  rhp.set_float("dropout", 0.01, 0.3)
+  rhp.set_float("gan_loss_factor", 0.01, 0.1)
+  rhp.set_float("bottleneck_l2_factor", 0.001, 0.1, scale=rhp.LOG_SCALE)
+  rhp.set_discrete("bottleneck_warmup_steps", [200, 500, 1000, 2000])
+  rhp.set_float("gumbel_temperature", 0, 1)
+  rhp.set_float("gumbel_noise_factor", 0, 0.5)
+
+
 @registry.register_ranged_hparams
 def autoencoder_discrete_pong_range(rhp):
   """Narrow tuning grid."""
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index a60f3fa72..0be91d7d5 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -181,7 +181,6 @@ def daisy_chain_getter(getter, name, *args, **kwargs):
           v = tf.identity(last_device_v)
         else:
           var = getter(name, *args, **kwargs)
-          # v = tf.identity(var._ref())  # pylint: disable=protected-access
           v = var.read_value()
 
         # keep track of the original variable
@@ -203,7 +202,6 @@ def caching_getter(getter, name, *args, **kwargs):
 
         v = getter(name, *args, **kwargs)
         with tf.device(self._caching_devices[i]):
-          # ret = tf.identity(v._ref())  # pylint: disable=protected-access
           ret = v.read_value()
         _add_variable_proxy_methods(v, ret)
         cache[key] = ret

From 377febf14e788231c609eb0dc21ad0ef7d8fb831 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 21 Aug 2018 15:04:20 -0700
Subject: [PATCH 0627/2720] set max examples to use for subword vocab for
 wsj_parsing problem.

PiperOrigin-RevId: 209666801
---
 tensor2tensor/data_generators/wsj_parsing.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/wsj_parsing.py b/tensor2tensor/data_generators/wsj_parsing.py
index 3db8b7060..82a810a2f 100644
--- a/tensor2tensor/data_generators/wsj_parsing.py
+++ b/tensor2tensor/data_generators/wsj_parsing.py
@@ -19,12 +19,12 @@
 from __future__ import print_function
 
 import os
+
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 import tensorflow as tf
 
-
 tf.flags.DEFINE_string("parsing_path", "", "Path to parsing files in tmp_dir.")
 
 
@@ -62,10 +62,17 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
   def generate_text_for_vocab(self, data_dir, tmp_dir):
     files = [os.path.join(tmp_dir, f) for f in self.TRAIN_FILES]
     inputs_file, targets_file = files
-    for sample in text_problems.text2text_txt_iterator(inputs_file,
-                                                       targets_file):
+    for i, sample in enumerate(text_problems.text2text_txt_iterator(inputs_file,
+                                                                    targets_file
+                                                                   )):
       yield sample["inputs"]
       yield sample["targets"]
+      if self.max_samples_for_vocab and (i + 1) >= self.max_samples_for_vocab:
+        break
+
+  @property
+  def max_samples_for_vocab(self):
+    return 1000
 
 
 def words_and_tags_from_wsj_tree(tree_string):

From ca776803384cc0a4e27d990f4071f93e5d342b58 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Tue, 21 Aug 2018 15:11:24 -0700
Subject: [PATCH 0628/2720] Clean up shape info in local_attention_1d.

PiperOrigin-RevId: 209668171
---
 tensor2tensor/layers/common_attention.py      | 14 ++-----
 tensor2tensor/layers/common_attention_test.py | 39 +++++++------------
 2 files changed, 18 insertions(+), 35 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index c1ac1973f..d82b6060e 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2311,11 +2311,8 @@ def local_attention_1d(q, k, v, block_length=128, filter_width=100, name=None):
   """
   with tf.variable_scope(
       name, default_name="local_self_attention_1d", values=[q, k, v]):
-    v_shape = v.get_shape()
-    depth_v = common_layers.shape_list(v)[3]
-    batch_size = common_layers.shape_list(q)[0]
-    num_heads = common_layers.shape_list(q)[1]
-    original_length = common_layers.shape_list(q)[2]
+    batch_size, num_heads, original_length, _ = common_layers.shape_list(q)
+    depth_v = common_layers.shape_list(v)[-1]
 
     # Pad query, key, value to ensure multiple of corresponding lengths.
     def pad_to_multiple(x, pad_length):
@@ -2331,10 +2328,7 @@ def pad_l_and_r(x, pad_length):
 
     # Set up query blocks.
     new_q_shape = common_layers.shape_list(q)
-    q = tf.reshape(q, [
-        new_q_shape[0], new_q_shape[1], new_q_shape[2] // block_length,
-        block_length, new_q_shape[3]
-    ])
+    q = reshape_by_blocks(q, new_q_shape, block_length)
 
     # Set up key and value blocks.
     # Get gather indices.
@@ -2378,7 +2372,7 @@ def pad_l_and_r(x, pad_length):
 
     # Remove the padding if introduced.
     output = tf.slice(output, [0, 0, 0, 0], [-1, -1, original_length, -1])
-    output.set_shape(v_shape)
+    output.set_shape([batch_size, num_heads, original_length, depth_v])
     return output
 
 
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 7ae4bc660..07d1fde16 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -88,33 +88,22 @@ def testMaskedLocalAttention1D(self):
 
     self.assertEqual(res.shape, (1, 1, 8, 1))
 
-  def testLocalUnmaskedAttention1D(self):
-    x = np.random.rand(5, 4, 25, 16)
-    y = np.random.rand(5, 4, 25, 16)
+  @parameterized.named_parameters(
+      ("matching_block_length", 3, 4, 25, 16, 16, 5),
+      ("unmatching_block_length", 3, 4, 25, 16, 16, 4),
+      ("different_depth_v", 3, 4, 25, 16, 17, 5),
+  )
+  def testLocalUnmaskedAttention1D(self, batch, heads, length,
+                                   depth_k, depth_v, block_length):
+    q = tf.random_normal([batch, heads, length, depth_k])
+    k = tf.random_normal([batch, heads, length, depth_k])
+    v = tf.random_normal([batch, heads, length, depth_v])
+    output = common_attention.local_attention_1d(
+        q, k, v, block_length=block_length, filter_width=3)
     with self.test_session() as session:
-      a = common_attention.local_attention_1d(
-          tf.constant(x, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          block_length=4,
-          filter_width=3)
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
-    self.assertEqual(res.shape, (5, 4, 25, 16))
+      res = session.run(output)
 
-  def testLocalUnmaskedAttention1DMatchingBlockLength(self):
-    x = np.random.rand(5, 4, 25, 16)
-    y = np.random.rand(5, 4, 25, 16)
-    with self.test_session() as session:
-      a = common_attention.local_attention_1d(
-          tf.constant(x, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          block_length=5,
-          filter_width=3)
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
-    self.assertEqual(res.shape, (5, 4, 25, 16))
+    self.assertEqual(res.shape, (batch, heads, length, depth_v))
 
   def testLocalUnmaskedAttention2D(self):
     x = np.random.rand(5, 4, 25, 25, 16)

From b31428f458aa4517b1c8712905971a2ecde5c048 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Tue, 21 Aug 2018 15:29:14 -0700
Subject: [PATCH 0629/2720] Replace manual ops w/ dot_product_attention in
 masked_local_attention_1d.

PiperOrigin-RevId: 209671048
---
 tensor2tensor/layers/common_attention.py      | 35 +++++++++---------
 tensor2tensor/layers/common_attention_test.py | 31 ++++++++--------
 tensor2tensor/mesh_tensorflow/mtf_layers.py   | 36 ++++++-------------
 .../mesh_tensorflow/mtf_layers_test.py        | 26 +++++++-------
 4 files changed, 56 insertions(+), 72 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index d82b6060e..a328556ba 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2073,27 +2073,28 @@ def masked_local_attention_1d(q,
                         [batch, heads, num_blocks - 1, block_length, depth_k])
     local_length = common_layers.shape_list(local_k)[3]
 
-    # [batch, heads, num_blocks - 1, block_length, local_length]
-    attention = tf.matmul(tail_q, local_k, transpose_b=True)
-
     # make sure source_pos <= target_pos
-    good_part = common_layers.ones_matrix_band_part(block_length,
-                                                    local_length,
-                                                    -1, block_length)
-    mask = (1.0 - good_part) * -1e9
-    mask = common_layers.cast_like(mask, attention)
-    attention += tf.reshape(mask, [1, 1, 1, block_length, local_length])
-    attention = tf.nn.softmax(attention)
-    attention = common_layers.dropout_with_broadcast_dims(
-        attention, 1.0 - dropout_rate,
-        broadcast_dims=None)
+    good_part = common_layers.ones_matrix_band_part(
+        block_length,
+        local_length,
+        -1,
+        block_length,
+        out_shape=[1, 1, 1, block_length, local_length])
+    bias = (1.0 - good_part) * -1e9
     # TODO(noam): figure out how to show a summary for the remaining blocks.
     # The naive way currently causes errors due to empty tensors.
     # output: [batch, heads, num_blocks-1, block_length, depth_v]
-    output = tf.matmul(attention, local_v)
-    output = tf.reshape(
-        output, [batch, heads, (num_blocks - 1) * block_length, depth_v])
-    output = tf.concat([first_output, output], axis=2)
+    tail_output = dot_product_attention(
+        tail_q,
+        local_k,
+        local_v,
+        bias,
+        dropout_rate=dropout_rate,
+        make_image_summary=False,
+        name="tail_block")
+    tail_output = tf.reshape(
+        tail_output, [batch, heads, (num_blocks - 1) * block_length, depth_v])
+    output = tf.concat([first_output, tail_output], axis=2)
 
     # Remove the padding if introduced.
     output = tf.slice(output, [0, 0, 0, 0], [-1, -1, original_length, -1])
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 07d1fde16..52f4ad27c 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -68,25 +68,22 @@ def testDotProductAttention(self):
       res = session.run(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
-  def testMaskedLocalAttention1D(self):
-    q = np.array([[[[1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0],
-                    [1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0],
-                    [1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0],
-                    [1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0]]]])
-    k = np.array([[[[1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0],
-                    [1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0],
-                    [1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0],
-                    [1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0]]]])
-    v = np.ones((1, 1, 8, 1))
+  @parameterized.parameters(
+      (1, 1, 8, 4, 1, 2),
+      (4, 1, 8, 4, 1, 4),
+      (3, 2, 8, 4, 3, 4),
+  )
+  def testMaskedLocalAttention1D(self, batch, heads, length, depth_k, depth_v,
+                                 block_length):
+    q = tf.random_normal([batch, heads, length, depth_k])
+    k = tf.random_normal([batch, heads, length, depth_k])
+    v = tf.random_normal([batch, heads, length, depth_v])
+    output = common_attention.masked_local_attention_1d(
+        q, k, v, block_length=block_length)
     with self.test_session() as session:
-      q_ = tf.constant(q, dtype=tf.float32)
-      k_ = tf.constant(k, dtype=tf.float32)
-      v_ = tf.constant(v, dtype=tf.float32)
-      y = common_attention.masked_local_attention_1d(
-          q_, k_, v_, block_length=tf.constant(2))
-      res = session.run(y)
+      res = session.run(output)
 
-    self.assertEqual(res.shape, (1, 1, 8, 1))
+    self.assertEqual(res.shape, (batch, heads, length, depth_v))
 
   @parameterized.named_parameters(
       ("matching_block_length", 3, 4, 25, 16, 16, 5),
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers.py b/tensor2tensor/mesh_tensorflow/mtf_layers.py
index c5d121c09..ca931561f 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers.py
@@ -249,15 +249,10 @@ def first_block_attention():
       first_q = mtf.slice(q, 0, 1, num_blocks.name)
       first_k = mtf.slice(k, 0, 1, num_blocks.name)
       first_v = mtf.slice(v, 0, 1, num_blocks.name)
-      block = first_q.shape.dims[2]
-
-      first_logits = mtf.einsum(
-          [first_q, first_k],
-          mtf.Shape([batch, heads, block, blength, mlength]))
-      weights = mtf.softmax(first_logits, mlength)
-      first_output = mtf.einsum(
-          [weights, first_v],
-          mtf.Shape([batch, heads, block, blength, kv_channels]))
+      first_output = dot_product_attention(first_q,
+                                           first_k,
+                                           first_v,
+                                           mask=None)
       return first_output
 
     # Attention for first block, since query_length = key_length.
@@ -273,31 +268,22 @@ def local(x):
 
     local_k = local(k)
     local_v = local(v)
-    mblocks = local_k.shape.dims[2]
-    mlength = local_k.shape.dims[3]
     # Calculate the causal mask to avoid peeking into the future. We compute
     # this once and reuse it for all blocks since the block_size is known.
+    mlength = local_k.shape.dims[3]
     mask = attention_bias_local_block(query_antecedent.mesh,
                                       blength, mlength)
 
     # Remove the first block from q since we already computed that.
     tail_q = mtf.slice(q, 1, num_blocks.size-1, num_blocks.name)
 
-    # Compatibility between q and k for rest of the blocks.
-    # Shape [batch, heads, num_blocks - 1, block_length, local_length]
-    attention = mtf.einsum(
-        [tail_q, local_k],
-        mtf.Shape([batch, heads, mblocks, blength, mlength]))
-    attention += mask
-    attention = mtf.softmax(attention, mlength)
-
-    # Run attention for rest of the blocks.
-    # Shape [batch, heads, num_blocks-1, block_length, kv_channels]
-    output = mtf.einsum(
-        [attention, local_v],
-        mtf.Shape([batch, heads, mblocks, blength, kv_channels]))
+    tail_output = dot_product_attention(tail_q,
+                                        local_k,
+                                        local_v,
+                                        mask=mask)
+
     # Now concatenate the first and rest of the blocks.
-    final_output = mtf.concat([first_output, output], num_blocks.name)
+    final_output = mtf.concat([first_output, tail_output], num_blocks.name)
     final_output = mtf.reshape(final_output, mtf.Shape(
         [batch, heads, query_length, kv_channels]))
     return mtf.einsum([final_output, o_var],
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers_test.py b/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
index 1c37d4c8d..26590b2c3 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
@@ -153,37 +153,37 @@ def testDenseReluDense(self):
     self.assertEqual(actual.shape, inputs.shape)
 
   @parameterized.parameters(
-      (4, 2),
+      (2, 16, 3, 4, 2, 2),
+      (1, 8, 5, 3, 1, 4),
   )
-  def testMaskedLocalAttention1D(self, kv_channels, heads):
-    batch = 2
-    length_q = 16
-    length_m = 16
-    channels = 3
-    query = tf.random_normal([batch, length_q, channels])
-    memory = tf.random_normal([batch, length_m, channels])
+  def testMaskedLocalAttention1D(self, batch, length, io_channels, kv_channels,
+                                 heads, block_length):
+    length_q = length
+    length_m = length
+    query = tf.random_normal([batch, length_q, io_channels])
+    memory = tf.random_normal([batch, length_m, io_channels])
 
     graph = mtf.Graph()
     mesh = mtf.Mesh(graph, "my_mesh")
     batch_dim = mtf.Dimension("batch", batch)
     length_q_dim = mtf.Dimension("length_q", length_q)
     length_m_dim = mtf.Dimension("length_m", length_m)
-    channels_dim = mtf.Dimension("channels", channels)
+    io_channels_dim = mtf.Dimension("io_channels", io_channels)
     kv_channels_dim = mtf.Dimension("kv_channels", kv_channels)
     heads_dim = mtf.Dimension("heads", heads)
 
     mtf_query = mtf.import_tf_tensor(
         mesh, query,
-        shape=mtf.Shape([batch_dim, length_q_dim, channels_dim]))
+        shape=mtf.Shape([batch_dim, length_q_dim, io_channels_dim]))
     mtf_memory = mtf.import_tf_tensor(
         mesh, memory,
-        shape=mtf.Shape([batch_dim, length_m_dim, channels_dim]))
+        shape=mtf.Shape([batch_dim, length_m_dim, io_channels_dim]))
     mtf_outputs = mtf_layers.masked_local_attention_1d(
         mtf_query,
         mtf_memory,
         kv_channels=kv_channels_dim,
         heads=heads_dim,
-        block_length=2)
+        block_length=block_length)
     mesh_impl = placement_mesh_impl.PlacementMeshImpl(
         shape=[], layout={}, devices=[""])
     lowering = mtf.Lowering(graph, {mesh: mesh_impl})
@@ -196,7 +196,7 @@ def testMaskedLocalAttention1D(self, kv_channels, heads):
       sess.run(tf_group)
       actual = sess.run(actual_outputs)
 
-    self.assertEqual(actual.shape, (batch, length_q, channels))
+    self.assertEqual(actual.shape, (batch, length_q, io_channels))
 
   @parameterized.parameters(
       (2, 4, 5, 7, 3, 1),

From 61ea8b0a66f249dee38180f893cf796c60513c08 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Tue, 21 Aug 2018 18:56:21 -0700
Subject: [PATCH 0630/2720] Reweighting losses in multiproblems.

PiperOrigin-RevId: 209699337
---
 .../data_generators/multi_problem.py          | 72 +++++++++++++++++--
 tensor2tensor/layers/common_hparams.py        | 10 ++-
 tensor2tensor/layers/common_layers.py         | 25 +++++++
 tensor2tensor/utils/modality.py               |  6 +-
 tensor2tensor/utils/t2t_model.py              | 15 ++++
 5 files changed, 121 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index d5a744dcc..44485dc5c 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -168,7 +168,7 @@ def dataset(self,
     if is_training:
       problem_step = tf.get_variable("problem_step",
                                      shape=[],
-                                     dtype=tf.float32,
+                                     dtype=tf.int64,
                                      initializer=tf.zeros_initializer(),
                                      trainable=False,
                                      use_resource=True)
@@ -183,7 +183,7 @@ def get_exp_sched_prob():
           inv_exp_decay = common_layers.inverse_exp_decay(
               max_step=hparams.multiproblem_schedule_max_examples,
               min_value=1e-4,
-              step=problem_step
+              step=tf.to_float(problem_step)
           )
           # inv_exp_decay is bounded above by 1.0
           return inv_exp_decay * hparams.multiproblem_schedule_threshold
@@ -196,7 +196,8 @@ def get_pretrain_sched_prob():
         with tf.control_dependencies([problem_step.assign_add(1)]):
           return tf.cond(
               tf.greater(problem_step,
-                         hparams.multiproblem_schedule_max_examples),
+                         tf.cast(hparams.multiproblem_schedule_max_examples,
+                                 dtype=tf.int64)),
               lambda: 1.0, lambda: 0.0)
 
       def mix_data(example):
@@ -220,7 +221,8 @@ def mix_data(example):
         tf.logging.info("Schedule mixing threshold "
                         "%.2f" % hparams.multiproblem_schedule_threshold)
         prob = tf.cond(
-            tf.equal(tf.floormod(problem_step, 5e6), 0),
+            tf.equal(tf.floormod(
+                problem_step, tf.cast(5e6, dtype=tf.int64)), 0),
             lambda: tf.Print(prob, [prob], message="Probability"),
             lambda: prob)
 
@@ -306,3 +308,65 @@ def get_max_num_classes(self):
           num = task.num_classes
 
     return num
+
+
+def aggregate_task_losses(hparams,
+                          problem_hparams,
+                          logits,
+                          target_modality,
+                          feature):
+  """Multiproblem loss function."""
+  summaries = []
+  main_task_id = hparams.problem.task_list[0].task_id
+  # Primary task loss
+  loss_num, loss_den = target_modality.loss(
+      logits, feature,
+      weights_fn=
+      lambda x: common_layers.weights_multi_problem_all(x, main_task_id))
+
+  loss_val = loss_num / tf.maximum(1.0, loss_den)
+  summaries.append([hparams.problem.task_list[0].name+"_loss", loss_val])
+
+  for task in hparams.problem.task_list[1:]:
+    if hasattr(task, "num_classes"):
+      task_loss_num_seq, task_loss_den_seq = target_modality.loss(
+          logits, feature,
+          weights_fn=
+          lambda x: common_layers.weights_multi_problem_input(x, task.task_id))  # pylint: disable=cell-var-from-loop
+      task_loss_num_seq *= problem_hparams.loss_multiplier
+
+      task_loss_num_label, task_loss_den_label = target_modality.loss(
+          logits, feature,
+          weights_fn=
+          lambda x: common_layers.weights_multi_problem(x, task.task_id))  # pylint: disable=cell-var-from-loop
+      task_loss_num_label *= problem_hparams.loss_multiplier
+
+      if hparams.multiproblem_reweight_label_loss:
+        task_loss_num = (1 - hparams.multiproblem_label_weight) * \
+                        task_loss_num_seq
+        task_loss_num += hparams.multiproblem_label_weight * task_loss_num_label
+      elif hparams.multiproblem_class_loss_multiplier > 0:
+        task_loss_num = task_loss_num_seq
+        task_loss_num += hparams.multiproblem_class_loss_multiplier * \
+                         task_loss_num_label
+      else:
+        task_loss_num = task_loss_num_seq + task_loss_num_label
+
+      task_loss_den = task_loss_den_seq + task_loss_den_label
+
+      # Log the unscaled versions of the losses to tensorboard.
+      task_loss_val = (task_loss_num_seq + task_loss_num_label) / tf.maximum(
+          1.0, task_loss_den)
+      summaries.append([task.name+"_loss", task_loss_val])
+
+      task_loss_val_label = task_loss_num_label / tf.maximum(
+          1.0, task_loss_den_label)
+      summaries.append([task.name+"_only_label_loss", task_loss_val_label])
+
+      loss_num += task_loss_num
+      loss_den += task_loss_den
+
+    else:
+      raise ValueError("Non-classification secondary tasks are not supported.")
+
+  return loss_num, loss_den, summaries
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index bbe8c2d18..4d408880d 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -266,7 +266,15 @@ def basic_params1():
       # and other tasks.
       # A list of supported schedules can be found in
       # `data_generators.multi_problem.py`.
-      multiproblem_mixing_schedule="constant"
+      multiproblem_mixing_schedule="constant",
+      # A scalar to upweight the classifier loss in a multiproblem setting.
+      multiproblem_class_loss_multiplier=0.0,
+      # A boolean that decides whether input sequence losses and target label
+      # losses in classification problems should be reweighted.
+      multiproblem_reweight_label_loss=False,
+      # How much weight the targets in classification problems receive. Inputs
+      # receive 1 minus this weight.
+      multiproblem_label_weight=0.5
   )
 
 
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index ce06e341a..5bae4ea80 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1956,6 +1956,31 @@ def weights_multi_problem(labels, taskid=-1):
   return tf.to_float(tf.not_equal(past_taskid * non_taskid, 0))
 
 
+def weights_multi_problem_all(labels, taskid=-1):
+  """Assign weight 1.0 to only examples from the given task."""
+  weights = tf.to_float(tf.not_equal(labels, 0))
+  if taskid < 0:
+    raise ValueError("Task ID must be non-negative.")
+
+  past_taskid = tf.cumsum(tf.to_float(tf.equal(labels, taskid)), axis=1)
+  # Additionally zero out the task id location
+  past_taskid *= tf.to_float(tf.not_equal(labels, taskid))
+  non_taskid = tf.to_float(labels)
+  example_mask = tf.to_float(tf.not_equal(past_taskid * non_taskid, 0))
+  example_mask = tf.reduce_sum(example_mask, axis=1)
+  example_mask = tf.to_float(
+      tf.greater(example_mask, tf.zeros_like(example_mask)))
+
+  return weights * tf.expand_dims(example_mask, axis=-1)
+
+
+def weights_multi_problem_input(labels, taskid=-1):
+  """Assign weight 1.0 to only the inputs for the given task."""
+  weights_all_tokens = weights_multi_problem_all(labels, taskid)
+  weights_target = weights_multi_problem(labels, taskid)
+  return weights_all_tokens - weights_target
+
+
 def weights_all(labels):
   """Assign weight 1.0 to all labels."""
   return tf.ones_like(labels, dtype=tf.float32)
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index 815d6d5ed..ded5577ef 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -177,14 +177,16 @@ def top_sharded(self, sharded_body_output, sharded_targets, data_parallelism):
     """
     return data_parallelism(self.top, sharded_body_output, sharded_targets)
 
-  def loss(self, top_out, targets):
+  def loss(self, top_out, targets, weights_fn=None):
     """Compute loss numerator and denominator for one shard of output."""
     logits = top_out
+    if weights_fn is None:
+      weights_fn = self.targets_weights_fn
     return common_layers.padded_cross_entropy(
         logits,
         targets,
         self._model_hparams.label_smoothing,
-        weights_fn=self.targets_weights_fn)
+        weights_fn=weights_fn)
 
   def loss_sharded(self, sharded_top_out, sharded_targets, data_parallelism):
     """Compute loss for all shards."""
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index b5cdcde81..28c088ca1 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -25,6 +25,7 @@
 import time
 import six
 
+from tensor2tensor.data_generators import multi_problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators.problem import problem_hparams_to_features
 from tensor2tensor.layers import common_layers
@@ -440,6 +441,20 @@ def _loss_single(self, logits, target_modality, feature):
 
     loss_num, loss_den = target_modality.loss(logits, feature)
     loss_num *= self._problem_hparams.loss_multiplier
+
+    if hasattr(self.hparams, "problem") and hasattr(
+        self.hparams.problem, "task_list"):
+      loss_num, loss_den, summaries = multi_problem.aggregate_task_losses(
+          self.hparams,
+          self._problem_hparams,
+          logits,
+          target_modality,
+          feature
+      )
+
+      for key, val in summaries:
+        tf.summary.scalar(key, val)
+
     return loss_num, loss_den
 
   def loss(self, logits, features):

From e1152bea99b3e8367bca6e9064cf77d1241d29e0 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 21 Aug 2018 19:40:40 -0700
Subject: [PATCH 0631/2720] replace standard python files with tf.gFile for
 input/output of

PiperOrigin-RevId: 209702762
---
 tensor2tensor/bin/t2t_decoder.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 8df00b163..399d2a4da 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -135,7 +135,11 @@ def score_file(filename):
     saver.restore(sess, ckpt)
     # Run on each line.
     results = []
-    for line in open(filename):
+    lines = []
+    with tf.gfile.Open(filename) as f:
+      text = f.read()
+      lines = [l.strip() for l in text.split("\n")]
+    for line in lines:
       tab_split = line.split("\t")
       if len(tab_split) > 2:
         raise ValueError("Each line must have at most one tab separator.")
@@ -173,7 +177,7 @@ def main(_):
     results = score_file(filename)
     if not FLAGS.decode_to_file:
       raise ValueError("To score a file, specify --decode_to_file for results.")
-    write_file = open(os.path.expanduser(FLAGS.decode_to_file), "w")
+    write_file = tf.gfile.Open(os.path.expanduser(FLAGS.decode_to_file), "w")
     for score in results:
       write_file.write("%.6f\n" % score)
     write_file.close()

From 16c142b92c330d6ed55b74ca06417379036723e9 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Tue, 21 Aug 2018 19:52:15 -0700
Subject: [PATCH 0632/2720] Clean up shape info in
 masked_within_block_local_attention_1d.

PiperOrigin-RevId: 209703515
---
 tensor2tensor/layers/common_attention.py      | 17 +++----
 tensor2tensor/layers/common_attention_test.py | 49 ++++++++++++++++---
 2 files changed, 51 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index a328556ba..3c472a3d3 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1918,16 +1918,14 @@ def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
   """
   with tf.variable_scope(
       name, default_name="within_local_attention_1d", values=[q, k, v]):
-    v_shape = v.get_shape()
-    batch, heads, length, _ = common_layers.shape_list(q)
+    batch, heads, length, depth_k = common_layers.shape_list(q)
+    depth_v = common_layers.shape_list(v)[-1]
     if isinstance(block_length, tf.Tensor):
       const = tf.contrib.util.constant_value(block_length)
       if const is not None:
         block_length = int(const)
 
     # Pad query, key, value to ensure multiple of block length.
-    depth_k = common_layers.shape_list(k)[3]
-    depth_v = common_layers.shape_list(v)[3]
     original_length = length
     padding_size = tf.mod(-length, block_length)
     length += padding_size
@@ -1943,9 +1941,8 @@ def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
     v = tf.reshape(v, [batch, heads, num_blocks, block_length, depth_v])
     # [batch, heads, num_blocks, block_length, block_length]
     attention = tf.matmul(q, k, transpose_b=True)
-    attention += tf.reshape(
-        attention_bias_lower_triangle(block_length),
-        [1, 1, 1, block_length, block_length])
+    attention += tf.reshape(attention_bias_lower_triangle(block_length),
+                            [1, 1, 1, block_length, block_length])
     attention = tf.nn.softmax(attention)
     # [batch, heads, num_blocks, block_length, depth_v]
     output = tf.matmul(attention, v)
@@ -1953,7 +1950,8 @@ def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
 
     # Remove the padding if introduced.
     output = tf.slice(output, [0, 0, 0, 0], [-1, -1, original_length, -1])
-    output.set_shape(v_shape)
+    output.set_shape([None if isinstance(dim, tf.Tensor) else dim for dim in
+                      (batch, heads, length, depth_v)])
     return output
 
 
@@ -2373,7 +2371,8 @@ def pad_l_and_r(x, pad_length):
 
     # Remove the padding if introduced.
     output = tf.slice(output, [0, 0, 0, 0], [-1, -1, original_length, -1])
-    output.set_shape([batch_size, num_heads, original_length, depth_v])
+    output.set_shape([None if isinstance(dim, tf.Tensor) else dim for dim in
+                      (batch_size, num_heads, original_length, depth_v)])
     return output
 
 
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 52f4ad27c..a14f43252 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -68,37 +68,74 @@ def testDotProductAttention(self):
       res = session.run(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
-  @parameterized.parameters(
-      (1, 1, 8, 4, 1, 2),
-      (4, 1, 8, 4, 1, 4),
-      (3, 2, 8, 4, 3, 4),
+  @parameterized.named_parameters(
+      ("", 1, 1, 8, 4, 1, 2),
+      ("dynamic_batch", None, 1, 8, 4, 1, 2),
+      ("batches", 4, 3, 8, 4, 1, 2),
+      ("depth_v", 1, 1, 8, 4, 3, 2),
+      ("block_length", 1, 1, 8, 4, 1, 4),
+  )
+  def testMaskedWithinBlockLocalAttention1D(self, batch, heads, length,
+                                            depth_k, depth_v, block_length):
+    if batch is None:
+      batch = tf.random_uniform([], minval=0, maxval=5, dtype=tf.int32)
+    q = tf.random_normal([batch, heads, length, depth_k])
+    k = tf.random_normal([batch, heads, length, depth_k])
+    v = tf.random_normal([batch, heads, length, depth_v])
+    output = common_attention.masked_within_block_local_attention_1d(
+        q, k, v, block_length=block_length)
+    with self.test_session() as session:
+      if isinstance(batch, tf.Tensor):
+        batch, res = session.run([batch, output])
+      else:
+        res = session.run(output)
+
+    self.assertEqual(res.shape, (batch, heads, length, depth_v))
+
+  @parameterized.named_parameters(
+      ("", 1, 1, 8, 4, 1, 2),
+      ("dynamic_batch", None, 1, 8, 4, 1, 2),
+      ("batches", 4, 3, 8, 4, 1, 2),
+      ("depth_v", 1, 1, 8, 4, 3, 2),
+      ("block_length", 1, 1, 8, 4, 1, 4),
   )
   def testMaskedLocalAttention1D(self, batch, heads, length, depth_k, depth_v,
                                  block_length):
+    if batch is None:
+      batch = tf.random_uniform([], minval=0, maxval=5, dtype=tf.int32)
     q = tf.random_normal([batch, heads, length, depth_k])
     k = tf.random_normal([batch, heads, length, depth_k])
     v = tf.random_normal([batch, heads, length, depth_v])
     output = common_attention.masked_local_attention_1d(
         q, k, v, block_length=block_length)
     with self.test_session() as session:
-      res = session.run(output)
+      if isinstance(batch, tf.Tensor):
+        batch, res = session.run([batch, output])
+      else:
+        res = session.run(output)
 
     self.assertEqual(res.shape, (batch, heads, length, depth_v))
 
   @parameterized.named_parameters(
       ("matching_block_length", 3, 4, 25, 16, 16, 5),
       ("unmatching_block_length", 3, 4, 25, 16, 16, 4),
+      ("dynamic_batch", None, 4, 25, 16, 16, 5),
       ("different_depth_v", 3, 4, 25, 16, 17, 5),
   )
   def testLocalUnmaskedAttention1D(self, batch, heads, length,
                                    depth_k, depth_v, block_length):
+    if batch is None:
+      batch = tf.random_uniform([], minval=0, maxval=5, dtype=tf.int32)
     q = tf.random_normal([batch, heads, length, depth_k])
     k = tf.random_normal([batch, heads, length, depth_k])
     v = tf.random_normal([batch, heads, length, depth_v])
     output = common_attention.local_attention_1d(
         q, k, v, block_length=block_length, filter_width=3)
     with self.test_session() as session:
-      res = session.run(output)
+      if isinstance(batch, tf.Tensor):
+        batch, res = session.run([batch, output])
+      else:
+        res = session.run(output)
 
     self.assertEqual(res.shape, (batch, heads, length, depth_v))
 

From 561e11bb80bde2cbf522bf060316573ad5eb790c Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Tue, 21 Aug 2018 20:13:32 -0700
Subject: [PATCH 0633/2720] Add vq_discrete, add vq to mixture of experts and
 new hparams

PiperOrigin-RevId: 209705147
---
 tensor2tensor/layers/common_attention.py |   2 +-
 tensor2tensor/layers/vq_discrete.py      | 358 +++++++++++++++++++++++
 tensor2tensor/models/transformer.py      | 232 ++++++++++++++-
 tensor2tensor/utils/expert_utils.py      | 142 +++++++--
 4 files changed, 712 insertions(+), 22 deletions(-)
 create mode 100644 tensor2tensor/layers/vq_discrete.py

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 3c472a3d3..9c2a3398f 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -231,7 +231,7 @@ def memeff_attention_fn(*args, **kwargs):
           num_experts=hparams.moe_num_experts,
           k=hparams.moe_k,
           loss_coef=hparams.moe_loss_coef,
-      ),
+          hparams=hparams),
       use_dp=False,
   )
 
diff --git a/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
new file mode 100644
index 000000000..2078660e1
--- /dev/null
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -0,0 +1,358 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Clean discrete bottleneck as in https://arxiv.org/abs/1805.11063."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from functools import partial
+
+from tensor2tensor.layers import common_layers
+
+import tensorflow as tf
+
+from tensorflow.python.training import moving_averages
+
+
+class DiscreteBottleneck(object):
+  """Discrete bottleneck class."""
+
+  def __init__(self, hparams):
+    self.hparams = hparams
+
+    # Set the discretization bottleneck specific things here
+    self.hparams.z_size_per_residual = self.hparams.z_size // \
+                                       self.hparams.num_residuals
+    self.hparams.block_dim = int(
+        self.hparams.hidden_size // self.hparams.num_blocks)
+    self.hparams.block_v_size = 2**(
+        self.hparams.z_size_per_residual / self.hparams.num_blocks)
+    self.hparams.block_v_size = int(self.hparams.block_v_size)
+    self.hparams.means = tf.get_variable(
+        name="means",
+        shape=[
+            self.hparams.num_residuals, self.hparams.num_blocks,
+            self.hparams.block_v_size, self.hparams.block_dim
+        ],
+        initializer=tf.uniform_unit_scaling_initializer())
+    tf.logging.info("Done creating means")
+
+    # Create the shadow variables if we are using EMA
+    self.hparams.ema_count = None
+    self.hparams.ema_means = None
+    if self.hparams.ema:
+      self.hparams.ema_count = []
+      self.hparams.ema_means = []
+      for i in range(hparams.num_residuals):
+        ema_count_i = tf.get_variable(
+            "ema_count_{}".format(i),
+            [self.hparams.num_blocks, self.hparams.block_v_size],
+            initializer=tf.constant_initializer(0),
+            trainable=False)
+        self.hparams.ema_count.append(ema_count_i)
+
+      with tf.colocate_with(self.hparams.means):
+        self.ema_means = []
+        for i in range(hparams.num_residuals):
+          ema_means_i = tf.get_variable(
+              "ema_means_{}".format(i),
+              initializer=self.hparams.means.initialized_value()[i],
+              trainable=False)
+          self.hparams.ema_means.append(ema_means_i)
+
+  def slice_hidden(self, x):
+    """Slice encoder hidden state into block_dim.
+
+    Args:
+        x: Encoder hidden state of shape [-1, hidden_size].
+
+    Returns:
+        Sliced states of shape [-1, num_blocks, block_dim].
+    """
+    x_sliced = tf.reshape(
+        x, shape=[-1, self.hparams.num_blocks, self.hparams.block_dim])
+    return x_sliced
+
+  def nearest_neighbor(self, x, means):
+    """Find the nearest element in means to elements in x.
+
+    Args:
+        x: Batch of encoder continuous latent states sliced/projected into
+           shape [-1, num_blocks, block_dim].
+        means: Embedding means of shape.
+
+    Returns:
+      Tensor with nearest element in mean encoded in one-hot notation.
+    """
+    x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keep_dims=True)
+    means_norm_sq = tf.reduce_sum(tf.square(means), axis=-1, keep_dims=True)
+    scalar_prod = tf.matmul(
+        tf.transpose(x, perm=[1, 0, 2]), tf.transpose(means, perm=[0, 2, 1]))
+    scalar_prod = tf.transpose(scalar_prod, perm=[1, 0, 2])
+    dist = x_norm_sq + tf.transpose(
+        means_norm_sq, perm=[2, 0, 1]) - 2 * scalar_prod
+
+    if self.hparams.soft_em:
+      nearest_idx = tf.stack(
+          [
+              tf.multinomial(
+                  -dist[:, i, :], num_samples=self.hparams.num_samples)
+              for i in range(self.hparams.num_blocks)
+          ],
+          axis=1)
+      nearest_hot = tf.one_hot(nearest_idx, depth=self.hparams.block_v_size)
+      nearest_hot = tf.reduce_mean(nearest_hot, axis=-2)
+    else:
+      if self.hparams.random_top_k > 1:
+        _, top_k_idx = tf.nn.top_k(-dist, k=self.hparams.random_top_k)
+        nearest_idx = tf.gather(
+            top_k_idx,
+            tf.random_uniform(
+                [1],
+                minval=0,
+                maxval=self.hparams.random_top_k - 1,
+                dtype=tf.int32),
+            axis=-1)
+      else:
+        if self.hparams.use_scales:
+          dist /= tf.reshape(self.hparams.scales,
+                             [1, 1, self.hparams.moe_num_experts])
+        nearest_idx = tf.argmax(-dist, axis=-1)
+      nearest_hot = tf.one_hot(nearest_idx, self.hparams.block_v_size)
+    return nearest_hot
+
+  def embedding_lookup(self, x, means):
+    """Compute nearest neighbors and loss for training the embeddings.
+
+    Args:
+        x: Batch of encoder continuous latent states sliced/projected into
+        shape
+        [-1, num_blocks, block_dim].
+        means: Embedding means.
+
+    Returns:
+        The nearest neighbor in one hot form, the nearest neighbor
+        itself, the
+        commitment loss, embedding training loss.
+    """
+    x_means_hot = self.nearest_neighbor(x, means)
+    x_means_hot_flat = tf.reshape(
+        x_means_hot, [-1, self.hparams.num_blocks, self.hparams.block_v_size])
+    x_means = tf.matmul(tf.transpose(x_means_hot_flat, perm=[1, 0, 2]), means)
+    x_means = tf.transpose(x_means, [1, 0, 2])
+    q_loss = tf.reduce_mean(tf.square((tf.stop_gradient(x) - x_means)))
+    e_loss = tf.reduce_mean((x - tf.stop_gradient(x_means))**2)
+    return x_means_hot, x_means, q_loss, e_loss
+
+  def bit_to_int(self, x_bit, num_bits, base=2):
+    """Turn x_bit representing numbers bitwise (lower-endian) to int tensor.
+
+    Args:
+        x_bit: Tensor containing numbers in a particular base to be
+        converted to
+        int.
+        num_bits: Number of bits in the representation.
+        base: Base of the representation.
+
+    Returns:
+        Integer representation of this number.
+    """
+    x_l = tf.stop_gradient(tf.to_int32(tf.reshape(x_bit, [-1, num_bits])))
+    x_labels = []
+    for i in range(num_bits):
+      x_labels.append(x_l[:, i] * tf.to_int32(base)**tf.to_int32(i))
+    res = sum(x_labels)
+    return tf.to_int32(tf.reshape(res, common_layers.shape_list(x_bit)[:-1]))
+
+  def int_to_bit(self, x_int, num_bits, base=2):
+    """Turn x_int representing numbers into a bitwise (lower-endian) tensor.
+
+    Args:
+        x_int: Tensor containing integer to be converted into base
+        notation.
+        num_bits: Number of bits in the representation.
+        base: Base of the representation.
+
+    Returns:
+        Corresponding number expressed in base.
+    """
+    x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
+    x_labels = []
+    for i in range(num_bits):
+      x_labels.append(
+          tf.floormod(
+              tf.floordiv(tf.to_int32(x_l),
+                          tf.to_int32(base)**i), tf.to_int32(base)))
+    res = tf.concat(x_labels, axis=-1)
+    return tf.to_float(res)
+
+  def embed(self, x, scope="bottleneck"):
+    """Embedding function that takes discrete latent and returns embedding.
+
+    Args:
+        x: Input to the discretization bottleneck.
+        scope: Scope name of the function.
+
+    Returns:
+        Continuous embedding to be passed on to the decoder.
+
+    Raises:
+        ValueError: For unknown or missing arguments.
+    """
+    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
+      shape_x = common_layers.shape_list(x)
+      x_flat = tf.reshape(x, [-1, 1])
+      c = self.int_to_bit(x_flat, num_bits=self.hparams.z_size, base=2)
+      shape = common_layers.shape_list(c)
+      new_shape = shape
+      new_shape[-1] = self.hparams.num_residuals
+      new_shape.append(self.hparams.num_blocks)
+      new_shape.append(
+          int(self.hparams.z_size /
+              (self.hparams.num_residuals * self.hparams.num_blocks)))
+      c = tf.to_int32(tf.reshape(c, shape=new_shape))
+      h1_shape = shape_x
+      h1_shape.append(self.hparams.hidden_size)
+      h1 = tf.zeros(dtype=tf.float32, shape=h1_shape)
+      for i in range(self.hparams.num_residuals):
+        c_residual = self.bit_to_int(
+            c[:, :, i, :, :],
+            num_bits=int(
+                self.hparams.z_size /
+                (self.hparams.num_residuals * self.hparams.num_blocks)),
+            base=2)
+        c_hot = tf.one_hot(c_residual, depth=self.hparams.block_v_size, axis=-1)
+        c_hot_flat = tf.reshape(
+            c_hot,
+            shape=[-1, self.hparams.num_blocks, self.hparams.block_v_size])
+        h1_residual = tf.matmul(
+            tf.transpose(c_hot_flat, perm=[1, 0, 2]), self.hparams.means[i])
+        h1_residual = tf.transpose(h1_residual, perm=[1, 0, 2])
+        h1_residual = tf.reshape(h1_residual, shape=h1_shape)
+        h1 += h1_residual
+
+      # Add Gaussian noise
+      h1_shape[0] = self.hparams.batch_size
+      h2 = tf.layers.dense(
+          tf.nn.relu(h1), self.hparams.filter_size, name="vch2")
+      res = tf.layers.dense(
+          tf.nn.relu(h2), self.hparams.hidden_size, name="vcfin")
+      return res
+
+  def discrete_bottleneck(self, x, scope="bottleneck"):
+    """Discretization bottleneck for latent variables.
+
+    Args:
+        x: Input to the discretization bottleneck.
+        scope: Scope of the function.
+
+    Returns:
+        Embedding to pass to the decoder, discrete latent, loss, and the
+        embedding
+        function.
+
+    Raises:
+        ValueError: If projection_tensors is None for reshape_method
+        project, or
+        ema_count or ema_means is None if we are using ema, or unknown
+        args.
+    """
+    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
+      x_reshaped = self.slice_hidden(x)
+      x_res = x_reshaped
+      x_means_hot = []
+      x_means = 0
+      loss = 0
+      for i in range(self.hparams.num_residuals):
+        x_means_hot_res, x_means_res, q_loss_res, e_loss_res = \
+            self.embedding_lookup(x_reshaped, self.hparams.means[i])
+
+        # Update the ema variables
+        if self.hparams.ema:
+          tf.logging.info("Using EMA with beta = {}".format(self.hparams.beta))
+          updated_ema_count_res = \
+              moving_averages.assign_moving_average(
+                  self.hparams.ema_count[i],
+                  tf.reduce_sum(
+                      tf.reshape(
+                          x_means_hot_res,
+                          shape=[-1, self.hparams.num_blocks,
+                                 self.hparams.block_v_size]),
+                      axis=0),
+                  self.hparams.decay,
+                  zero_debias=False)
+
+          dw = tf.matmul(
+              tf.transpose(x_means_hot_res, perm=[1, 2, 0]),
+              tf.transpose(x_res, perm=[1, 0, 2]))
+
+          updated_ema_means_res = \
+              moving_averages.assign_moving_average(
+                  self.hparams.ema_means[i], dw, self.hparams.decay,
+                  zero_debias=False)
+          n = tf.reduce_sum(updated_ema_count_res, axis=-1, keep_dims=True)
+          updated_ema_count_res = (
+              (updated_ema_count_res + self.hparams.epsilon) /
+              (n + 2**self.hparams.z_size * self.hparams.epsilon) * n)
+          updated_ema_means_res = updated_ema_means_res/tf.expand_dims(
+              updated_ema_count_res, axis=-1)
+
+          with tf.control_dependencies([e_loss_res]):
+            update_means_res = tf.assign(self.hparams.means[i],
+                                         updated_ema_means_res)
+            with tf.control_dependencies([update_means_res]):
+              loss += self.hparams.beta * e_loss_res
+        else:
+          loss += q_loss_res + self.hparams.beta * e_loss_res
+
+        # Update the residuals
+        x_res -= x_means_res
+        x_means += x_means_res
+        x_means_hot.append(x_means_hot_res)
+
+      # Get the discrete latent representation
+      x_means_hot = tf.stack(x_means_hot, axis=1)
+      x_means_idx = tf.argmax(x_means_hot, axis=-1)
+
+      # Get the binary representation
+      num_bits = int(self.hparams.z_size //
+                     (self.hparams.num_blocks * self.hparams.num_residuals))
+      x_means_bits = self.int_to_bit(x_means_idx, num_bits=num_bits, base=2)
+      shape = common_layers.shape_list(x_means_bits)
+      new_shape = shape[:-2]
+      new_shape[0] = -1
+      new_shape[-1] = self.hparams.z_size
+      x_means_bits = tf.reshape(x_means_bits, new_shape)
+      x_discrete = self.bit_to_int(
+          tf.to_int32(x_means_bits), num_bits=self.hparams.z_size, base=2)
+
+      # Reshape x_discrete
+      shape_x = common_layers.shape_list(x)
+      shape_discrete = shape_x[:-1]
+      x_discrete = tf.reshape(x_discrete, shape_discrete)
+      x_means = tf.reshape(x_means, shape=shape_x)
+      h1 = x + tf.stop_gradient(x_means - x)
+
+      h2 = tf.layers.dense(
+          tf.nn.relu(h1), self.hparams.filter_size, name="vch2")
+      res = tf.layers.dense(
+          tf.nn.relu(h2), self.hparams.hidden_size, name="vcfin")
+      embed_fn = partial(self.embed, scope=scope)
+      return {
+          "dense": res,
+          "discrete": x_discrete,
+          "loss": loss,
+          "embed": embed_fn
+      }
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index c31b800c2..64c47c55b 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1453,10 +1453,19 @@ def transformer_ffn_layer(x,
         hparams.moe_num_experts,
         overhead=overhead,
         loss_coef=hparams.moe_loss_coef)
-    if losses is None:
-      raise ValueError(
-          "transformer_ffn_layer with type local_moe_tpu must pass in "
-          "a losses list")
+  elif ffn_layer == "local_moe":
+    overhead = (
+        hparams.moe_overhead_train
+        if hparams.mode == tf.estimator.ModeKeys.TRAIN else
+        hparams.moe_overhead_eval)
+    ret, loss = expert_utils.local_moe(
+        x,
+        True,
+        expert_utils.ffn_expert_fn(hparams.hidden_size, [hparams.filter_size],
+                                   hparams.hidden_size),
+        hparams.moe_num_experts,
+        k=hparams.moe_k,
+        hparams=hparams)
     losses.append(loss)
     return ret
   else:
@@ -1539,6 +1548,221 @@ def transformer_base_v2():
   return hparams
 
 
+@registry.register_hparams
+def transformer_base_topk_32():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq_16_nb1()
+  hparams.batch_size = 1024
+  hparams.moe_num_experts = 32
+  hparams.filter_size = 4096
+  hparams.hidden_size = 1024
+  hparams.ffn_layers = "local_moe"
+  hparams.moe_k = 2
+  hparams.gating_type = "topk"
+  hparams.num_decoder_layers = 6
+  hparams.label_smoothing = 0.
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_topk_16():
+  """Set of hyperparameters."""
+  hparams = transformer_base_topk_32()
+  hparams.moe_num_experts = 16
+  hparams.moe_k = 2
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_v2_32_nb1():
+  """Set of hyperparameters."""
+  hparams = transformer_base_v2()
+  expert_utils.update_hparams_for_vq_gating(hparams)
+  hparams.batch_size = 1024
+  hparams.moe_num_experts = 32
+  hparams.filter_size = 4096
+  hparams.hidden_size = 1024
+  hparams.shared_embedding_and_softmax_weights = False
+  # one epoch for languagemodel_lm1b32k_packed w/ 256 batch = 13600 steps
+  hparams.learning_rate_decay_steps = 10000
+  hparams.num_heads = 4
+  hparams.num_blocks = 1
+  hparams.moe_k = 1
+  hparams.ffn_layer = "local_moe"
+  hparams.num_decoder_layers = 6
+  hparams.label_smoothing = 0.
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_vq_16_nb1():
+  """Set of hyperparameters."""
+  hparams = transformer_base_v2_32_nb1()
+  hparams.gating_type = "vq"
+  hparams.batch_size = 1024
+  hparams.moe_num_experts = 16
+  hparams.num_blocks = 1
+  hparams.moe_k = 1
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_vq_ada_32ex_packed():
+  """Set of hyperparameters."""
+  hparams = transformer_base_v2_32_nb1()
+  hparams.moe_num_experts = 32
+  hparams.gating_type = "vq"
+  # this gives us a batch size of 16 because each seq is len 256
+  hparams.batch_size = 5072
+  hparams.shared_embedding_and_softmax_weights = False
+  hparams.learning_rate_warmup_steps = 10000
+  # one epoch for languagemodel_lm1b32k_packed = 27200 steps
+  hparams.learning_rate_decay_steps = 27200
+  hparams.num_heads = 4
+  hparams.num_blocks = 1
+  hparams.moe_k = 1
+  hparams.num_decoder_layers = 6
+  hparams.label_smoothing = 0.
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.layer_postprocess_sequence = "dan"
+  hparams.layer_preprocess_sequence = "none"
+  hparams.weight_decay = 1e-06
+  hparams.attention_dropout = 0.1
+  hparams.optimizer = "Adafactor"
+  hparams.learning_rate_schedule = "linear_warmup*rsqrt_decay*linear_decay"
+  hparams.activation_dtype = "float32"
+  hparams.learning_rate = 0.1
+  # doesn't seem to feature in the learning rate schedule but making sure
+  hparams.learning_rate_constant = 1.0
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_vq_16_nb1_packed_nda():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq_16_nb1()
+  hparams.layer_preprocess_sequence = "n"
+  hparams.layer_postprocess_sequence = "da"
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_vq_16_nb1_packed_nda_b01():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq_16_nb1()
+  hparams.layer_preprocess_sequence = "n"
+  hparams.layer_postprocess_sequence = "da"
+  hparams.beta = 0.1
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_vqtopk2_16_nb1_packed_nda():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq_16_nb1()
+  hparams.layer_preprocess_sequence = "n"
+  hparams.layer_postprocess_sequence = "da"
+  hparams.moe_k = 2
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_vq2_16_nb1_packed_nda_b01_scales():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vqtopk2_16_nb1_packed_nda()
+  hparams.use_scales = int(True)
+  hparams.moe_k = 2
+  hparams.beta = 0.1
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_vq1_16_nb1_packed_nda_b01_scales():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq_16_nb1()
+  hparams.use_scales = int(True)
+  hparams.moe_k = 1
+  hparams.beta = 0.1
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_vq1k_16_nb1_packed_nda_small_scale():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq1_16_nb1_packed_nda_b01_scales()
+  hparams.hidden_size = 128
+  hparams.filter_size = 128
+  hparams.expert_filter_size = 128
+  hparams.ema = int(False)
+  hparams.moe_k = 1
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_vq1k_16_nb1_packed_nda_small_scale_cent():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq1k_16_nb1_packed_nda_small_scale()
+  hparams.use_scales = int(True)
+  hparams.residual_centroids = int(True)
+  hparams.moe_k = 1
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_vq_32small_nb1_packed_nda():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq_16_nb1_packed_nda()
+  hparams.expert_filter_size = 2048
+  hparams.moe_num_experts = 32
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_vq_32_nb1_packed_nda():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq_16_nb1_packed_nda()
+  hparams.moe_num_experts = 32
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_topk_16_packed_small():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq_16_nb1()
+  hparams.moe_num_experts = 16
+  hparams.moe_k = 2
+  hparams.gating_type = "topk"
+  hparams.expert_filter_size = 128
+  hparams.filter_size = 128
+  hparams.hidden_size = 128
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_v2_16_nb1():
+  """Set of hyperparameters."""
+  hparams = transformer_base_v2_32_nb1()
+  hparams.moe_num_experts = 16
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_v2_lmbase():
+  """Set of hyperparameters."""
+  hparams = transformer_base_v2_32_nb1()
+  hparams.expert_layer = 500  # basically expert layer will never be called
+  hparams.batch_size = 32768
+  return hparams
+
+
+@registry.register_hparams
+def transformer_ada_lmpackedbase():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq_ada_32ex_packed()
+  hparams.expert_layer = 500  # basically expert layer will never be called
+  return hparams
+
+
 @registry.register_hparams
 def transformer_base():
   """Base parameters for Transformer model."""
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 0be91d7d5..82fadb5eb 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -29,6 +29,7 @@
 from six.moves import zip  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers.vq_discrete import DiscreteBottleneck
 
 import tensorflow as tf
 
@@ -379,6 +380,27 @@ def _gates_to_load(gates):
   return tf.reduce_sum(tf.to_float(gates > 0), 0)
 
 
+def update_hparams_for_vq_gating(hparams):
+  """VQ Gating hparams."""
+  hparams.add_hparam("z_size", 14)
+  hparams.add_hparam("noise_dev", 0.5)
+  # Bottleneck kinds supported: dense, vae, dvq.
+  hparams.add_hparam("bottleneck_kind", "dvq")
+  hparams.add_hparam("num_blocks", 1)
+  hparams.add_hparam("num_residuals", 1)
+  # Reshape method for DVQ: slice, project
+  hparams.add_hparam("beta", 0.25)
+  hparams.add_hparam("epsilon", 1e-5)
+  hparams.add_hparam("decay", 0.999)
+  hparams.add_hparam("ema", True)
+  hparams.add_hparam("random_top_k", 1)
+  hparams.add_hparam("soft_em", False)
+  hparams.add_hparam("num_samples", 10)
+  hparams.add_hparam("gating_type", "vq")
+  hparams.add_hparam("use_scales", int(True))
+  hparams.add_hparam("residual_centroids", int(False))
+
+
 def _my_top_k(x, k):
   """GPU-compatible version of top-k that works for very small constant k.
 
@@ -411,6 +433,81 @@ def _my_top_k(x, k):
   return tf.stack(values, axis=1), tf.to_int32(tf.stack(indices, axis=1))
 
 
+def vq_gating(x,
+              num_experts,
+              k=2,
+              hparams=None,
+              name="vq_gating"):
+  """VQ gating.
+
+  Args:
+    x: input Tensor with shape [batch_size, input_size]
+    num_experts: an integer
+    k: an integer - number of experts per example
+    hparams: optional hparams
+    name: an optional string
+
+  Returns:
+    gates: a Tensor with shape [batch_size, num_experts]
+    load: a Tensor with shape [num_experts]
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    if hparams.use_scales:
+      scales = tf.get_variable(
+          "scales", [num_experts],
+          tf.float32,
+          initializer=tf.ones_initializer())
+      scales = tf.nn.softmax(scales)
+      hparams.scales = scales
+    input_size = x.get_shape().as_list()[-1]
+    batch_size = common_layers.shape_list(x)[0]
+
+    if k > 1:
+      # first project into two dense layers, chop and discretize, and gate
+      # TODO(avaswani): Maybe scale the embeddings flowing out of the experts.
+      # We might want to do this to match the computation being done by topk
+      x = tf.layers.dense(x, input_size * k)
+      # x goes from [batch_size, input_size*k] to [batch_size*k, input_size]
+      x = tf.reshape(x, [batch_size * k, input_size])
+    inputs = tf.expand_dims(x, axis=1)
+    inputs = tf.expand_dims(inputs, axis=1)
+    # VQ hparams
+    hparams.z_size = int(math.log(num_experts, 2))
+    hparams.hidden_size = input_size
+    hparams.top_k = k
+    bneck = DiscreteBottleneck(hparams)
+    d = bneck.discrete_bottleneck(inputs)
+    centroids = None
+    exp_discrete = d["discrete"]
+    embed_lookup = d["embed"]
+    extra_loss = d["loss"]
+    if hparams.residual_centroids:
+      centroids = embed_lookup(exp_discrete)  # gives the centroids
+    top_k_indices = tf.squeeze(exp_discrete, axis=1)
+    tf.summary.histogram("discrete_counts", top_k_indices)
+    # if k > 1, then we need to reshape top_k_indices from [batch_size*k, 1]
+    # to [batch_size, k]
+    if k > 1:
+      top_k_indices = tf.reshape(top_k_indices, [batch_size, k])
+    # get the top k gates
+    top_k_gates = tf.ones([batch_size, k])
+    # This will be a `Tensor` of shape `[batch_size, n]`, with zeros in the
+    # positions corresponding to all but the top k experts per example.
+    gates = _rowwise_unsorted_segment_sum(top_k_gates, top_k_indices,
+                                          num_experts)
+    # Compute count per expert from the gates.
+    # gates has shape [batch_size, num_experts]
+    # count per expert has shape [num_experts, 1]
+    count_per_expert = tf.reduce_sum(gates, axis=0)
+    if hparams.use_scales:
+      scale_loss = tf.reduce_mean(tf.to_float(count_per_expert) * scales)
+      extra_loss += scale_loss
+    if common_layers.should_generate_summaries():
+      tf.summary.histogram("vq_loss", extra_loss)
+      tf.summary.historgram("scale_loss", scale_loss)
+    return gates, extra_loss, centroids
+
+
 def noisy_top_k_gating(x,
                        num_experts,
                        train,
@@ -459,6 +556,7 @@ def noisy_top_k_gating(x,
     else:
       logits = clean_logits
     top_logits, top_indices = _my_top_k(logits, min(k + 1, num_experts))
+    # top k logits has shape [batch, k]
     top_k_logits = tf.slice(top_logits, [0, 0], [-1, k])
     top_k_indices = tf.slice(top_indices, [0, 0], [-1, k])
     top_k_gates = tf.nn.softmax(top_k_logits)
@@ -963,8 +1061,9 @@ def local_moe(x,
               train,
               expert_fn,
               num_experts,
-              k=2,
+              k=1,
               loss_coef=1e-2,
+              hparams=None,
               pass_x=True,
               pass_gates=False,
               additional_dispatch_params=None,
@@ -978,6 +1077,7 @@ def local_moe(x,
     num_experts: an integer - number of experts
     k: an integer - how many experts to use for each batch element
     loss_coef: a scalar - multiplier on load-balancing losses
+    hparams: optional hparams for vq gating
     pass_x: a boolean. If true, x will also be dispatched to the experts.
     pass_gates: a boolean. If true, gates will be passed to experts. Might be
       necessary when dealing with sparse encoder-encoder decoder attention
@@ -996,20 +1096,28 @@ def local_moe(x,
 
   with tf.variable_scope(name, default_name="local_moe"):
     x_flat = flatten_all_but_last(x)
-
-    # The gates indicate which batch elements go to which tensors.
-    # load is a measure of approximately how many examples go to each expert
-    gates, load = noisy_top_k_gating(
-        x_flat,
-        num_experts,
-        train,
-        k,
-        initializer=tf.zeros_initializer(),
-        noisy_gating=True,
-        noise_epsilon=1e-2)
-    # This magic object helps us shuffle data between datashards and experts.
+    if hparams.gating_type == "topk":
+      tf.logging.info("Using noisy top_k with k = {}".format(k))
+      # The gates indicate which batch elements go to which tensors.
+      # load is a measure of approximately how many examples go to each expert
+      gates, load = noisy_top_k_gating(
+          x_flat,
+          num_experts,
+          train,
+          k,
+          initializer=tf.zeros_initializer(),
+          noisy_gating=True,
+          noise_epsilon=1e-2)
+      importance = tf.reduce_sum(gates, 0)
+      loss = loss_coef * (cv_squared(importance) + cv_squared(load))
+    else:
+      assert hparams.gating_type == "vq"
+      tf.logging.info("Using VQ gating")
+      gates, loss, centroids = vq_gating(
+          x_flat, num_experts, k, hparams=hparams)
+    loss *= loss_coef
+    # Shuffle data between datashards and experts.
     dispatcher = SparseDispatcher(num_experts, gates)
-
     # Set up expert_fn arguments
     expert_kwargs = {}
     if pass_x:
@@ -1024,10 +1132,10 @@ def local_moe(x,
     expert_outputs = ep(expert_fn, **expert_kwargs)
 
     y_flat = dispatcher.combine(expert_outputs)
+    if centroids is not None:
+      centroids = tf.squeeze(centroids, axis=[1, 2])
+      y_flat += centroids
     y = common_layers.reshape_like(y_flat, x)
-
-    importance = tf.reduce_sum(gates, 0)
-    loss = loss_coef * (cv_squared(importance) + cv_squared(load))
     return y, loss
 
 

From 57444300243f068bad88eb5ed51a9793c4bde172 Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Tue, 21 Aug 2018 21:06:16 -0700
Subject: [PATCH 0634/2720] Add vq_discrete, add vq to mixture of experts and
 new hparams

PiperOrigin-RevId: 209709134
---
 tensor2tensor/layers/common_attention.py |   2 +-
 tensor2tensor/layers/vq_discrete.py      | 358 -----------------------
 tensor2tensor/models/transformer.py      | 232 +--------------
 tensor2tensor/utils/expert_utils.py      | 142 ++-------
 4 files changed, 22 insertions(+), 712 deletions(-)
 delete mode 100644 tensor2tensor/layers/vq_discrete.py

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 9c2a3398f..3c472a3d3 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -231,7 +231,7 @@ def memeff_attention_fn(*args, **kwargs):
           num_experts=hparams.moe_num_experts,
           k=hparams.moe_k,
           loss_coef=hparams.moe_loss_coef,
-          hparams=hparams),
+      ),
       use_dp=False,
   )
 
diff --git a/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
deleted file mode 100644
index 2078660e1..000000000
--- a/tensor2tensor/layers/vq_discrete.py
+++ /dev/null
@@ -1,358 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-"""Clean discrete bottleneck as in https://arxiv.org/abs/1805.11063."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-from functools import partial
-
-from tensor2tensor.layers import common_layers
-
-import tensorflow as tf
-
-from tensorflow.python.training import moving_averages
-
-
-class DiscreteBottleneck(object):
-  """Discrete bottleneck class."""
-
-  def __init__(self, hparams):
-    self.hparams = hparams
-
-    # Set the discretization bottleneck specific things here
-    self.hparams.z_size_per_residual = self.hparams.z_size // \
-                                       self.hparams.num_residuals
-    self.hparams.block_dim = int(
-        self.hparams.hidden_size // self.hparams.num_blocks)
-    self.hparams.block_v_size = 2**(
-        self.hparams.z_size_per_residual / self.hparams.num_blocks)
-    self.hparams.block_v_size = int(self.hparams.block_v_size)
-    self.hparams.means = tf.get_variable(
-        name="means",
-        shape=[
-            self.hparams.num_residuals, self.hparams.num_blocks,
-            self.hparams.block_v_size, self.hparams.block_dim
-        ],
-        initializer=tf.uniform_unit_scaling_initializer())
-    tf.logging.info("Done creating means")
-
-    # Create the shadow variables if we are using EMA
-    self.hparams.ema_count = None
-    self.hparams.ema_means = None
-    if self.hparams.ema:
-      self.hparams.ema_count = []
-      self.hparams.ema_means = []
-      for i in range(hparams.num_residuals):
-        ema_count_i = tf.get_variable(
-            "ema_count_{}".format(i),
-            [self.hparams.num_blocks, self.hparams.block_v_size],
-            initializer=tf.constant_initializer(0),
-            trainable=False)
-        self.hparams.ema_count.append(ema_count_i)
-
-      with tf.colocate_with(self.hparams.means):
-        self.ema_means = []
-        for i in range(hparams.num_residuals):
-          ema_means_i = tf.get_variable(
-              "ema_means_{}".format(i),
-              initializer=self.hparams.means.initialized_value()[i],
-              trainable=False)
-          self.hparams.ema_means.append(ema_means_i)
-
-  def slice_hidden(self, x):
-    """Slice encoder hidden state into block_dim.
-
-    Args:
-        x: Encoder hidden state of shape [-1, hidden_size].
-
-    Returns:
-        Sliced states of shape [-1, num_blocks, block_dim].
-    """
-    x_sliced = tf.reshape(
-        x, shape=[-1, self.hparams.num_blocks, self.hparams.block_dim])
-    return x_sliced
-
-  def nearest_neighbor(self, x, means):
-    """Find the nearest element in means to elements in x.
-
-    Args:
-        x: Batch of encoder continuous latent states sliced/projected into
-           shape [-1, num_blocks, block_dim].
-        means: Embedding means of shape.
-
-    Returns:
-      Tensor with nearest element in mean encoded in one-hot notation.
-    """
-    x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keep_dims=True)
-    means_norm_sq = tf.reduce_sum(tf.square(means), axis=-1, keep_dims=True)
-    scalar_prod = tf.matmul(
-        tf.transpose(x, perm=[1, 0, 2]), tf.transpose(means, perm=[0, 2, 1]))
-    scalar_prod = tf.transpose(scalar_prod, perm=[1, 0, 2])
-    dist = x_norm_sq + tf.transpose(
-        means_norm_sq, perm=[2, 0, 1]) - 2 * scalar_prod
-
-    if self.hparams.soft_em:
-      nearest_idx = tf.stack(
-          [
-              tf.multinomial(
-                  -dist[:, i, :], num_samples=self.hparams.num_samples)
-              for i in range(self.hparams.num_blocks)
-          ],
-          axis=1)
-      nearest_hot = tf.one_hot(nearest_idx, depth=self.hparams.block_v_size)
-      nearest_hot = tf.reduce_mean(nearest_hot, axis=-2)
-    else:
-      if self.hparams.random_top_k > 1:
-        _, top_k_idx = tf.nn.top_k(-dist, k=self.hparams.random_top_k)
-        nearest_idx = tf.gather(
-            top_k_idx,
-            tf.random_uniform(
-                [1],
-                minval=0,
-                maxval=self.hparams.random_top_k - 1,
-                dtype=tf.int32),
-            axis=-1)
-      else:
-        if self.hparams.use_scales:
-          dist /= tf.reshape(self.hparams.scales,
-                             [1, 1, self.hparams.moe_num_experts])
-        nearest_idx = tf.argmax(-dist, axis=-1)
-      nearest_hot = tf.one_hot(nearest_idx, self.hparams.block_v_size)
-    return nearest_hot
-
-  def embedding_lookup(self, x, means):
-    """Compute nearest neighbors and loss for training the embeddings.
-
-    Args:
-        x: Batch of encoder continuous latent states sliced/projected into
-        shape
-        [-1, num_blocks, block_dim].
-        means: Embedding means.
-
-    Returns:
-        The nearest neighbor in one hot form, the nearest neighbor
-        itself, the
-        commitment loss, embedding training loss.
-    """
-    x_means_hot = self.nearest_neighbor(x, means)
-    x_means_hot_flat = tf.reshape(
-        x_means_hot, [-1, self.hparams.num_blocks, self.hparams.block_v_size])
-    x_means = tf.matmul(tf.transpose(x_means_hot_flat, perm=[1, 0, 2]), means)
-    x_means = tf.transpose(x_means, [1, 0, 2])
-    q_loss = tf.reduce_mean(tf.square((tf.stop_gradient(x) - x_means)))
-    e_loss = tf.reduce_mean((x - tf.stop_gradient(x_means))**2)
-    return x_means_hot, x_means, q_loss, e_loss
-
-  def bit_to_int(self, x_bit, num_bits, base=2):
-    """Turn x_bit representing numbers bitwise (lower-endian) to int tensor.
-
-    Args:
-        x_bit: Tensor containing numbers in a particular base to be
-        converted to
-        int.
-        num_bits: Number of bits in the representation.
-        base: Base of the representation.
-
-    Returns:
-        Integer representation of this number.
-    """
-    x_l = tf.stop_gradient(tf.to_int32(tf.reshape(x_bit, [-1, num_bits])))
-    x_labels = []
-    for i in range(num_bits):
-      x_labels.append(x_l[:, i] * tf.to_int32(base)**tf.to_int32(i))
-    res = sum(x_labels)
-    return tf.to_int32(tf.reshape(res, common_layers.shape_list(x_bit)[:-1]))
-
-  def int_to_bit(self, x_int, num_bits, base=2):
-    """Turn x_int representing numbers into a bitwise (lower-endian) tensor.
-
-    Args:
-        x_int: Tensor containing integer to be converted into base
-        notation.
-        num_bits: Number of bits in the representation.
-        base: Base of the representation.
-
-    Returns:
-        Corresponding number expressed in base.
-    """
-    x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
-    x_labels = []
-    for i in range(num_bits):
-      x_labels.append(
-          tf.floormod(
-              tf.floordiv(tf.to_int32(x_l),
-                          tf.to_int32(base)**i), tf.to_int32(base)))
-    res = tf.concat(x_labels, axis=-1)
-    return tf.to_float(res)
-
-  def embed(self, x, scope="bottleneck"):
-    """Embedding function that takes discrete latent and returns embedding.
-
-    Args:
-        x: Input to the discretization bottleneck.
-        scope: Scope name of the function.
-
-    Returns:
-        Continuous embedding to be passed on to the decoder.
-
-    Raises:
-        ValueError: For unknown or missing arguments.
-    """
-    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
-      shape_x = common_layers.shape_list(x)
-      x_flat = tf.reshape(x, [-1, 1])
-      c = self.int_to_bit(x_flat, num_bits=self.hparams.z_size, base=2)
-      shape = common_layers.shape_list(c)
-      new_shape = shape
-      new_shape[-1] = self.hparams.num_residuals
-      new_shape.append(self.hparams.num_blocks)
-      new_shape.append(
-          int(self.hparams.z_size /
-              (self.hparams.num_residuals * self.hparams.num_blocks)))
-      c = tf.to_int32(tf.reshape(c, shape=new_shape))
-      h1_shape = shape_x
-      h1_shape.append(self.hparams.hidden_size)
-      h1 = tf.zeros(dtype=tf.float32, shape=h1_shape)
-      for i in range(self.hparams.num_residuals):
-        c_residual = self.bit_to_int(
-            c[:, :, i, :, :],
-            num_bits=int(
-                self.hparams.z_size /
-                (self.hparams.num_residuals * self.hparams.num_blocks)),
-            base=2)
-        c_hot = tf.one_hot(c_residual, depth=self.hparams.block_v_size, axis=-1)
-        c_hot_flat = tf.reshape(
-            c_hot,
-            shape=[-1, self.hparams.num_blocks, self.hparams.block_v_size])
-        h1_residual = tf.matmul(
-            tf.transpose(c_hot_flat, perm=[1, 0, 2]), self.hparams.means[i])
-        h1_residual = tf.transpose(h1_residual, perm=[1, 0, 2])
-        h1_residual = tf.reshape(h1_residual, shape=h1_shape)
-        h1 += h1_residual
-
-      # Add Gaussian noise
-      h1_shape[0] = self.hparams.batch_size
-      h2 = tf.layers.dense(
-          tf.nn.relu(h1), self.hparams.filter_size, name="vch2")
-      res = tf.layers.dense(
-          tf.nn.relu(h2), self.hparams.hidden_size, name="vcfin")
-      return res
-
-  def discrete_bottleneck(self, x, scope="bottleneck"):
-    """Discretization bottleneck for latent variables.
-
-    Args:
-        x: Input to the discretization bottleneck.
-        scope: Scope of the function.
-
-    Returns:
-        Embedding to pass to the decoder, discrete latent, loss, and the
-        embedding
-        function.
-
-    Raises:
-        ValueError: If projection_tensors is None for reshape_method
-        project, or
-        ema_count or ema_means is None if we are using ema, or unknown
-        args.
-    """
-    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
-      x_reshaped = self.slice_hidden(x)
-      x_res = x_reshaped
-      x_means_hot = []
-      x_means = 0
-      loss = 0
-      for i in range(self.hparams.num_residuals):
-        x_means_hot_res, x_means_res, q_loss_res, e_loss_res = \
-            self.embedding_lookup(x_reshaped, self.hparams.means[i])
-
-        # Update the ema variables
-        if self.hparams.ema:
-          tf.logging.info("Using EMA with beta = {}".format(self.hparams.beta))
-          updated_ema_count_res = \
-              moving_averages.assign_moving_average(
-                  self.hparams.ema_count[i],
-                  tf.reduce_sum(
-                      tf.reshape(
-                          x_means_hot_res,
-                          shape=[-1, self.hparams.num_blocks,
-                                 self.hparams.block_v_size]),
-                      axis=0),
-                  self.hparams.decay,
-                  zero_debias=False)
-
-          dw = tf.matmul(
-              tf.transpose(x_means_hot_res, perm=[1, 2, 0]),
-              tf.transpose(x_res, perm=[1, 0, 2]))
-
-          updated_ema_means_res = \
-              moving_averages.assign_moving_average(
-                  self.hparams.ema_means[i], dw, self.hparams.decay,
-                  zero_debias=False)
-          n = tf.reduce_sum(updated_ema_count_res, axis=-1, keep_dims=True)
-          updated_ema_count_res = (
-              (updated_ema_count_res + self.hparams.epsilon) /
-              (n + 2**self.hparams.z_size * self.hparams.epsilon) * n)
-          updated_ema_means_res = updated_ema_means_res/tf.expand_dims(
-              updated_ema_count_res, axis=-1)
-
-          with tf.control_dependencies([e_loss_res]):
-            update_means_res = tf.assign(self.hparams.means[i],
-                                         updated_ema_means_res)
-            with tf.control_dependencies([update_means_res]):
-              loss += self.hparams.beta * e_loss_res
-        else:
-          loss += q_loss_res + self.hparams.beta * e_loss_res
-
-        # Update the residuals
-        x_res -= x_means_res
-        x_means += x_means_res
-        x_means_hot.append(x_means_hot_res)
-
-      # Get the discrete latent representation
-      x_means_hot = tf.stack(x_means_hot, axis=1)
-      x_means_idx = tf.argmax(x_means_hot, axis=-1)
-
-      # Get the binary representation
-      num_bits = int(self.hparams.z_size //
-                     (self.hparams.num_blocks * self.hparams.num_residuals))
-      x_means_bits = self.int_to_bit(x_means_idx, num_bits=num_bits, base=2)
-      shape = common_layers.shape_list(x_means_bits)
-      new_shape = shape[:-2]
-      new_shape[0] = -1
-      new_shape[-1] = self.hparams.z_size
-      x_means_bits = tf.reshape(x_means_bits, new_shape)
-      x_discrete = self.bit_to_int(
-          tf.to_int32(x_means_bits), num_bits=self.hparams.z_size, base=2)
-
-      # Reshape x_discrete
-      shape_x = common_layers.shape_list(x)
-      shape_discrete = shape_x[:-1]
-      x_discrete = tf.reshape(x_discrete, shape_discrete)
-      x_means = tf.reshape(x_means, shape=shape_x)
-      h1 = x + tf.stop_gradient(x_means - x)
-
-      h2 = tf.layers.dense(
-          tf.nn.relu(h1), self.hparams.filter_size, name="vch2")
-      res = tf.layers.dense(
-          tf.nn.relu(h2), self.hparams.hidden_size, name="vcfin")
-      embed_fn = partial(self.embed, scope=scope)
-      return {
-          "dense": res,
-          "discrete": x_discrete,
-          "loss": loss,
-          "embed": embed_fn
-      }
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 64c47c55b..c31b800c2 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1453,19 +1453,10 @@ def transformer_ffn_layer(x,
         hparams.moe_num_experts,
         overhead=overhead,
         loss_coef=hparams.moe_loss_coef)
-  elif ffn_layer == "local_moe":
-    overhead = (
-        hparams.moe_overhead_train
-        if hparams.mode == tf.estimator.ModeKeys.TRAIN else
-        hparams.moe_overhead_eval)
-    ret, loss = expert_utils.local_moe(
-        x,
-        True,
-        expert_utils.ffn_expert_fn(hparams.hidden_size, [hparams.filter_size],
-                                   hparams.hidden_size),
-        hparams.moe_num_experts,
-        k=hparams.moe_k,
-        hparams=hparams)
+    if losses is None:
+      raise ValueError(
+          "transformer_ffn_layer with type local_moe_tpu must pass in "
+          "a losses list")
     losses.append(loss)
     return ret
   else:
@@ -1548,221 +1539,6 @@ def transformer_base_v2():
   return hparams
 
 
-@registry.register_hparams
-def transformer_base_topk_32():
-  """Set of hyperparameters."""
-  hparams = transformer_base_vq_16_nb1()
-  hparams.batch_size = 1024
-  hparams.moe_num_experts = 32
-  hparams.filter_size = 4096
-  hparams.hidden_size = 1024
-  hparams.ffn_layers = "local_moe"
-  hparams.moe_k = 2
-  hparams.gating_type = "topk"
-  hparams.num_decoder_layers = 6
-  hparams.label_smoothing = 0.
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_topk_16():
-  """Set of hyperparameters."""
-  hparams = transformer_base_topk_32()
-  hparams.moe_num_experts = 16
-  hparams.moe_k = 2
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_v2_32_nb1():
-  """Set of hyperparameters."""
-  hparams = transformer_base_v2()
-  expert_utils.update_hparams_for_vq_gating(hparams)
-  hparams.batch_size = 1024
-  hparams.moe_num_experts = 32
-  hparams.filter_size = 4096
-  hparams.hidden_size = 1024
-  hparams.shared_embedding_and_softmax_weights = False
-  # one epoch for languagemodel_lm1b32k_packed w/ 256 batch = 13600 steps
-  hparams.learning_rate_decay_steps = 10000
-  hparams.num_heads = 4
-  hparams.num_blocks = 1
-  hparams.moe_k = 1
-  hparams.ffn_layer = "local_moe"
-  hparams.num_decoder_layers = 6
-  hparams.label_smoothing = 0.
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_vq_16_nb1():
-  """Set of hyperparameters."""
-  hparams = transformer_base_v2_32_nb1()
-  hparams.gating_type = "vq"
-  hparams.batch_size = 1024
-  hparams.moe_num_experts = 16
-  hparams.num_blocks = 1
-  hparams.moe_k = 1
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_vq_ada_32ex_packed():
-  """Set of hyperparameters."""
-  hparams = transformer_base_v2_32_nb1()
-  hparams.moe_num_experts = 32
-  hparams.gating_type = "vq"
-  # this gives us a batch size of 16 because each seq is len 256
-  hparams.batch_size = 5072
-  hparams.shared_embedding_and_softmax_weights = False
-  hparams.learning_rate_warmup_steps = 10000
-  # one epoch for languagemodel_lm1b32k_packed = 27200 steps
-  hparams.learning_rate_decay_steps = 27200
-  hparams.num_heads = 4
-  hparams.num_blocks = 1
-  hparams.moe_k = 1
-  hparams.num_decoder_layers = 6
-  hparams.label_smoothing = 0.
-  hparams.layer_prepostprocess_dropout = 0.1
-  hparams.layer_postprocess_sequence = "dan"
-  hparams.layer_preprocess_sequence = "none"
-  hparams.weight_decay = 1e-06
-  hparams.attention_dropout = 0.1
-  hparams.optimizer = "Adafactor"
-  hparams.learning_rate_schedule = "linear_warmup*rsqrt_decay*linear_decay"
-  hparams.activation_dtype = "float32"
-  hparams.learning_rate = 0.1
-  # doesn't seem to feature in the learning rate schedule but making sure
-  hparams.learning_rate_constant = 1.0
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_vq_16_nb1_packed_nda():
-  """Set of hyperparameters."""
-  hparams = transformer_base_vq_16_nb1()
-  hparams.layer_preprocess_sequence = "n"
-  hparams.layer_postprocess_sequence = "da"
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_vq_16_nb1_packed_nda_b01():
-  """Set of hyperparameters."""
-  hparams = transformer_base_vq_16_nb1()
-  hparams.layer_preprocess_sequence = "n"
-  hparams.layer_postprocess_sequence = "da"
-  hparams.beta = 0.1
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_vqtopk2_16_nb1_packed_nda():
-  """Set of hyperparameters."""
-  hparams = transformer_base_vq_16_nb1()
-  hparams.layer_preprocess_sequence = "n"
-  hparams.layer_postprocess_sequence = "da"
-  hparams.moe_k = 2
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_vq2_16_nb1_packed_nda_b01_scales():
-  """Set of hyperparameters."""
-  hparams = transformer_base_vqtopk2_16_nb1_packed_nda()
-  hparams.use_scales = int(True)
-  hparams.moe_k = 2
-  hparams.beta = 0.1
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_vq1_16_nb1_packed_nda_b01_scales():
-  """Set of hyperparameters."""
-  hparams = transformer_base_vq_16_nb1()
-  hparams.use_scales = int(True)
-  hparams.moe_k = 1
-  hparams.beta = 0.1
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_vq1k_16_nb1_packed_nda_small_scale():
-  """Set of hyperparameters."""
-  hparams = transformer_base_vq1_16_nb1_packed_nda_b01_scales()
-  hparams.hidden_size = 128
-  hparams.filter_size = 128
-  hparams.expert_filter_size = 128
-  hparams.ema = int(False)
-  hparams.moe_k = 1
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_vq1k_16_nb1_packed_nda_small_scale_cent():
-  """Set of hyperparameters."""
-  hparams = transformer_base_vq1k_16_nb1_packed_nda_small_scale()
-  hparams.use_scales = int(True)
-  hparams.residual_centroids = int(True)
-  hparams.moe_k = 1
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_vq_32small_nb1_packed_nda():
-  """Set of hyperparameters."""
-  hparams = transformer_base_vq_16_nb1_packed_nda()
-  hparams.expert_filter_size = 2048
-  hparams.moe_num_experts = 32
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_vq_32_nb1_packed_nda():
-  """Set of hyperparameters."""
-  hparams = transformer_base_vq_16_nb1_packed_nda()
-  hparams.moe_num_experts = 32
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_topk_16_packed_small():
-  """Set of hyperparameters."""
-  hparams = transformer_base_vq_16_nb1()
-  hparams.moe_num_experts = 16
-  hparams.moe_k = 2
-  hparams.gating_type = "topk"
-  hparams.expert_filter_size = 128
-  hparams.filter_size = 128
-  hparams.hidden_size = 128
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_v2_16_nb1():
-  """Set of hyperparameters."""
-  hparams = transformer_base_v2_32_nb1()
-  hparams.moe_num_experts = 16
-  return hparams
-
-
-@registry.register_hparams
-def transformer_base_v2_lmbase():
-  """Set of hyperparameters."""
-  hparams = transformer_base_v2_32_nb1()
-  hparams.expert_layer = 500  # basically expert layer will never be called
-  hparams.batch_size = 32768
-  return hparams
-
-
-@registry.register_hparams
-def transformer_ada_lmpackedbase():
-  """Set of hyperparameters."""
-  hparams = transformer_base_vq_ada_32ex_packed()
-  hparams.expert_layer = 500  # basically expert layer will never be called
-  return hparams
-
-
 @registry.register_hparams
 def transformer_base():
   """Base parameters for Transformer model."""
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 82fadb5eb..0be91d7d5 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -29,7 +29,6 @@
 from six.moves import zip  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_layers
-from tensor2tensor.layers.vq_discrete import DiscreteBottleneck
 
 import tensorflow as tf
 
@@ -380,27 +379,6 @@ def _gates_to_load(gates):
   return tf.reduce_sum(tf.to_float(gates > 0), 0)
 
 
-def update_hparams_for_vq_gating(hparams):
-  """VQ Gating hparams."""
-  hparams.add_hparam("z_size", 14)
-  hparams.add_hparam("noise_dev", 0.5)
-  # Bottleneck kinds supported: dense, vae, dvq.
-  hparams.add_hparam("bottleneck_kind", "dvq")
-  hparams.add_hparam("num_blocks", 1)
-  hparams.add_hparam("num_residuals", 1)
-  # Reshape method for DVQ: slice, project
-  hparams.add_hparam("beta", 0.25)
-  hparams.add_hparam("epsilon", 1e-5)
-  hparams.add_hparam("decay", 0.999)
-  hparams.add_hparam("ema", True)
-  hparams.add_hparam("random_top_k", 1)
-  hparams.add_hparam("soft_em", False)
-  hparams.add_hparam("num_samples", 10)
-  hparams.add_hparam("gating_type", "vq")
-  hparams.add_hparam("use_scales", int(True))
-  hparams.add_hparam("residual_centroids", int(False))
-
-
 def _my_top_k(x, k):
   """GPU-compatible version of top-k that works for very small constant k.
 
@@ -433,81 +411,6 @@ def _my_top_k(x, k):
   return tf.stack(values, axis=1), tf.to_int32(tf.stack(indices, axis=1))
 
 
-def vq_gating(x,
-              num_experts,
-              k=2,
-              hparams=None,
-              name="vq_gating"):
-  """VQ gating.
-
-  Args:
-    x: input Tensor with shape [batch_size, input_size]
-    num_experts: an integer
-    k: an integer - number of experts per example
-    hparams: optional hparams
-    name: an optional string
-
-  Returns:
-    gates: a Tensor with shape [batch_size, num_experts]
-    load: a Tensor with shape [num_experts]
-  """
-  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
-    if hparams.use_scales:
-      scales = tf.get_variable(
-          "scales", [num_experts],
-          tf.float32,
-          initializer=tf.ones_initializer())
-      scales = tf.nn.softmax(scales)
-      hparams.scales = scales
-    input_size = x.get_shape().as_list()[-1]
-    batch_size = common_layers.shape_list(x)[0]
-
-    if k > 1:
-      # first project into two dense layers, chop and discretize, and gate
-      # TODO(avaswani): Maybe scale the embeddings flowing out of the experts.
-      # We might want to do this to match the computation being done by topk
-      x = tf.layers.dense(x, input_size * k)
-      # x goes from [batch_size, input_size*k] to [batch_size*k, input_size]
-      x = tf.reshape(x, [batch_size * k, input_size])
-    inputs = tf.expand_dims(x, axis=1)
-    inputs = tf.expand_dims(inputs, axis=1)
-    # VQ hparams
-    hparams.z_size = int(math.log(num_experts, 2))
-    hparams.hidden_size = input_size
-    hparams.top_k = k
-    bneck = DiscreteBottleneck(hparams)
-    d = bneck.discrete_bottleneck(inputs)
-    centroids = None
-    exp_discrete = d["discrete"]
-    embed_lookup = d["embed"]
-    extra_loss = d["loss"]
-    if hparams.residual_centroids:
-      centroids = embed_lookup(exp_discrete)  # gives the centroids
-    top_k_indices = tf.squeeze(exp_discrete, axis=1)
-    tf.summary.histogram("discrete_counts", top_k_indices)
-    # if k > 1, then we need to reshape top_k_indices from [batch_size*k, 1]
-    # to [batch_size, k]
-    if k > 1:
-      top_k_indices = tf.reshape(top_k_indices, [batch_size, k])
-    # get the top k gates
-    top_k_gates = tf.ones([batch_size, k])
-    # This will be a `Tensor` of shape `[batch_size, n]`, with zeros in the
-    # positions corresponding to all but the top k experts per example.
-    gates = _rowwise_unsorted_segment_sum(top_k_gates, top_k_indices,
-                                          num_experts)
-    # Compute count per expert from the gates.
-    # gates has shape [batch_size, num_experts]
-    # count per expert has shape [num_experts, 1]
-    count_per_expert = tf.reduce_sum(gates, axis=0)
-    if hparams.use_scales:
-      scale_loss = tf.reduce_mean(tf.to_float(count_per_expert) * scales)
-      extra_loss += scale_loss
-    if common_layers.should_generate_summaries():
-      tf.summary.histogram("vq_loss", extra_loss)
-      tf.summary.historgram("scale_loss", scale_loss)
-    return gates, extra_loss, centroids
-
-
 def noisy_top_k_gating(x,
                        num_experts,
                        train,
@@ -556,7 +459,6 @@ def noisy_top_k_gating(x,
     else:
       logits = clean_logits
     top_logits, top_indices = _my_top_k(logits, min(k + 1, num_experts))
-    # top k logits has shape [batch, k]
     top_k_logits = tf.slice(top_logits, [0, 0], [-1, k])
     top_k_indices = tf.slice(top_indices, [0, 0], [-1, k])
     top_k_gates = tf.nn.softmax(top_k_logits)
@@ -1061,9 +963,8 @@ def local_moe(x,
               train,
               expert_fn,
               num_experts,
-              k=1,
+              k=2,
               loss_coef=1e-2,
-              hparams=None,
               pass_x=True,
               pass_gates=False,
               additional_dispatch_params=None,
@@ -1077,7 +978,6 @@ def local_moe(x,
     num_experts: an integer - number of experts
     k: an integer - how many experts to use for each batch element
     loss_coef: a scalar - multiplier on load-balancing losses
-    hparams: optional hparams for vq gating
     pass_x: a boolean. If true, x will also be dispatched to the experts.
     pass_gates: a boolean. If true, gates will be passed to experts. Might be
       necessary when dealing with sparse encoder-encoder decoder attention
@@ -1096,28 +996,20 @@ def local_moe(x,
 
   with tf.variable_scope(name, default_name="local_moe"):
     x_flat = flatten_all_but_last(x)
-    if hparams.gating_type == "topk":
-      tf.logging.info("Using noisy top_k with k = {}".format(k))
-      # The gates indicate which batch elements go to which tensors.
-      # load is a measure of approximately how many examples go to each expert
-      gates, load = noisy_top_k_gating(
-          x_flat,
-          num_experts,
-          train,
-          k,
-          initializer=tf.zeros_initializer(),
-          noisy_gating=True,
-          noise_epsilon=1e-2)
-      importance = tf.reduce_sum(gates, 0)
-      loss = loss_coef * (cv_squared(importance) + cv_squared(load))
-    else:
-      assert hparams.gating_type == "vq"
-      tf.logging.info("Using VQ gating")
-      gates, loss, centroids = vq_gating(
-          x_flat, num_experts, k, hparams=hparams)
-    loss *= loss_coef
-    # Shuffle data between datashards and experts.
+
+    # The gates indicate which batch elements go to which tensors.
+    # load is a measure of approximately how many examples go to each expert
+    gates, load = noisy_top_k_gating(
+        x_flat,
+        num_experts,
+        train,
+        k,
+        initializer=tf.zeros_initializer(),
+        noisy_gating=True,
+        noise_epsilon=1e-2)
+    # This magic object helps us shuffle data between datashards and experts.
     dispatcher = SparseDispatcher(num_experts, gates)
+
     # Set up expert_fn arguments
     expert_kwargs = {}
     if pass_x:
@@ -1132,10 +1024,10 @@ def local_moe(x,
     expert_outputs = ep(expert_fn, **expert_kwargs)
 
     y_flat = dispatcher.combine(expert_outputs)
-    if centroids is not None:
-      centroids = tf.squeeze(centroids, axis=[1, 2])
-      y_flat += centroids
     y = common_layers.reshape_like(y_flat, x)
+
+    importance = tf.reduce_sum(gates, 0)
+    loss = loss_coef * (cv_squared(importance) + cv_squared(load))
     return y, loss
 
 

From 57d4f06285ddfd75a14b466f3b8a6727becacae8 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Tue, 21 Aug 2018 22:32:43 -0700
Subject: [PATCH 0635/2720] Add Eager tests in mesh_tensorflow/.

PiperOrigin-RevId: 209715363
---
 .../mesh_tensorflow/mesh_tensorflow_test.py   |  4 +-
 .../mesh_tensorflow/mtf_layers_test.py        | 48 ++++++++-----------
 2 files changed, 23 insertions(+), 29 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
index e73d7e4ec..efec84659 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
@@ -117,6 +117,7 @@ def testGraph(self):
     self.assertLen(graph.trainable_variables, 1)
     self.assertLen(graph.all_variables, 2)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes
   def testLowering(self):
     graph = mtf.Graph()
     mesh = mtf.Mesh(graph, "my_mesh")
@@ -129,8 +130,7 @@ def testLowering(self):
     lowering = mtf.Lowering(graph, {mesh: mesh_impl})
 
     outputs = lowering.export_to_tf_tensor(mtf_inputs)
-    with self.test_session() as sess:
-      inputs_value, outputs_value = sess.run([inputs, outputs])
+    inputs_value, outputs_value = self.evaluate([inputs, outputs])
     self.assertEqual(inputs_value, outputs_value)
 
     # Check that methods run without error.
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers_test.py b/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
index 26590b2c3..4d6cba965 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
@@ -28,6 +28,7 @@
 import tensorflow as tf
 
 
+@tf.contrib.eager.run_all_tests_in_graph_and_eager_modes
 class MtfLayersTest(parameterized.TestCase, tf.test.TestCase):
 
   @parameterized.parameters(
@@ -62,10 +63,9 @@ def testDense(self, units, use_bias):
                                              use_bias=use_bias)(inputs)
     tf_group = lowering.copy_masters_to_slices()
     init = tf.global_variables_initializer()
-    with self.test_session() as sess:
-      sess.run(init)
-      sess.run(tf_group)
-      actual, expected = sess.run([actual_outputs, expected_outputs])
+    self.evaluate(init)
+    self.evaluate(tf_group)
+    actual, expected = self.evaluate([actual_outputs, expected_outputs])
 
     self.assertEqual(actual.shape, expected.shape)
 
@@ -91,10 +91,9 @@ def testLayerNorm(self):
     expected_outputs = common_layers.layer_norm(inputs)
     tf_group = lowering.copy_masters_to_slices()
     init = tf.global_variables_initializer()
-    with self.test_session() as sess:
-      sess.run(init)
-      sess.run(tf_group)
-      actual, expected = sess.run([actual_outputs, expected_outputs])
+    self.evaluate(init)
+    self.evaluate(tf_group)
+    actual, expected = self.evaluate([actual_outputs, expected_outputs])
 
     self.assertEqual(actual.shape, expected.shape)
 
@@ -116,9 +115,8 @@ def testWeightsNonzero(self):
 
     expected_outputs = common_layers.weights_nonzero(inputs)
     tf_group = lowering.copy_masters_to_slices()
-    with self.test_session() as sess:
-      sess.run(tf_group)
-      actual, expected = sess.run([actual_outputs, expected_outputs])
+    self.evaluate(tf_group)
+    actual, expected = self.evaluate([actual_outputs, expected_outputs])
 
     self.assertAllEqual(actual, expected)
 
@@ -145,10 +143,9 @@ def testDenseReluDense(self):
 
     tf_group = lowering.copy_masters_to_slices()
     init = tf.global_variables_initializer()
-    with self.test_session() as sess:
-      sess.run(init)
-      sess.run(tf_group)
-      actual = sess.run(actual_outputs)
+    self.evaluate(init)
+    self.evaluate(tf_group)
+    actual = self.evaluate(actual_outputs)
 
     self.assertEqual(actual.shape, inputs.shape)
 
@@ -191,10 +188,9 @@ def testMaskedLocalAttention1D(self, batch, length, io_channels, kv_channels,
 
     tf_group = lowering.copy_masters_to_slices()
     init = tf.global_variables_initializer()
-    with self.test_session() as sess:
-      sess.run(init)
-      sess.run(tf_group)
-      actual = sess.run(actual_outputs)
+    self.evaluate(init)
+    self.evaluate(tf_group)
+    actual = self.evaluate(actual_outputs)
 
     self.assertEqual(actual.shape, (batch, length_q, io_channels))
 
@@ -240,10 +236,9 @@ def testDotProductAttention(
 
     tf_group = lowering.copy_masters_to_slices()
     init = tf.global_variables_initializer()
-    with self.test_session() as sess:
-      sess.run(init)
-      sess.run(tf_group)
-      actual = sess.run(actual_outputs)
+    self.evaluate(init)
+    self.evaluate(tf_group)
+    actual = self.evaluate(actual_outputs)
 
     self.assertEqual(actual.shape, (batch, heads, length_q, depth_v))
 
@@ -281,10 +276,9 @@ def testMultiheadAttention(self, kv_channels, heads):
 
     tf_group = lowering.copy_masters_to_slices()
     init = tf.global_variables_initializer()
-    with self.test_session() as sess:
-      sess.run(init)
-      sess.run(tf_group)
-      actual = sess.run(actual_outputs)
+    self.evaluate(init)
+    self.evaluate(tf_group)
+    actual = self.evaluate(actual_outputs)
 
     self.assertEqual(actual.shape, query.shape)
 

From 66b5ce69587ab6c271c5797df60e946d74b5114a Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Tue, 21 Aug 2018 22:45:32 -0700
Subject: [PATCH 0636/2720] Clean up shape info in masked_local_attention_1d.

PiperOrigin-RevId: 209716252
---
 tensor2tensor/layers/common_attention.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 3c472a3d3..ebf70535d 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2012,10 +2012,8 @@ def masked_local_attention_1d(q,
   """
   with tf.variable_scope(
       name, default_name="local_attention_1d", values=[q, k, v]):
-
-    batch = common_layers.shape_list(q)[0]
-    heads = common_layers.shape_list(q)[1]
-    length = common_layers.shape_list(q)[2]
+    batch, heads, length, depth_k = common_layers.shape_list(q)
+    depth_v = common_layers.shape_list(v)[-1]
     if isinstance(block_length, tf.Tensor):
       const = tf.contrib.util.constant_value(block_length)
       if const is not None:
@@ -2028,8 +2026,6 @@ def masked_local_attention_1d(q,
           tf.less(length, block_length * 2), length, block_length)
 
     # Pad query, key, value to ensure multiple of block length.
-    depth_k = common_layers.shape_list(k)[3]
-    depth_v = common_layers.shape_list(v)[3]
     original_length = length
     padding_size = tf.mod(-length, block_length)
     length += padding_size

From 8ac94737827d134d6291838182d72a2865d6a301 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Tue, 21 Aug 2018 23:22:33 -0700
Subject: [PATCH 0637/2720] add prerequisites for spatially-partitioned
 convolutions.

PiperOrigin-RevId: 209718944
---
 .../mesh_tensorflow/mesh_tensorflow.py        | 276 +++++++++++++++++-
 tensor2tensor/mesh_tensorflow/mnist.py        |  31 +-
 .../mesh_tensorflow/placement_mesh_impl.py    |  37 +++
 .../mesh_tensorflow/simd_mesh_impl.py         |  43 +++
 4 files changed, 379 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index 3adc48804..e21a8fa83 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -696,7 +696,7 @@ def allreduce(self, x, mesh_axes, reduction_fn_string):
     """
     raise NotImplementedError("Allreduce not implemented")
 
-  def allsplit(self, x, mesh_axis, split_axis):
+  def allsplit(self, x, mesh_axis, split_axis, which=None):
     """Inverse of allconcat - split each slice and keep only one piece of it.
 
     The number of ways to split is the number of processors in the group.
@@ -706,20 +706,24 @@ def allsplit(self, x, mesh_axis, split_axis):
       x: LaidOutTensor.
       mesh_axis: int, the mesh axis along which to split.
       split_axis: int, the Tensor axis along which to split.
+      which: an optional LaidOutTensor of integer scalars. Selects the slice to
+        to keep, instead of the coordinate.
 
     Returns:
       LaidOutTensor.
     """
+    if which is None:
+      which = self.laid_out_pcoord(mesh_axis)
     num_splits = self.shape[mesh_axis].size
-    def my_fn(x, coordinate):
+    def my_fn(x, which):
       slice_begin = [
-          dimsize // num_splits * coordinate if i == split_axis
+          dimsize // num_splits * which if i == split_axis
           else 0 for i, dimsize in enumerate(x.shape.as_list())]
       slice_size = [
           dimsize // num_splits if i == split_axis
           else dimsize for i, dimsize in enumerate(x.shape.as_list())]
       return tf.slice(x, slice_begin, slice_size)
-    return self.slicewise(my_fn, x, self.laid_out_pcoord(mesh_axis))
+    return self.slicewise(my_fn, x, which)
 
   def allconcat(self, x, mesh_axis, concat_axis):
     """Grouped allconcat (like MPI allgather followed by concat).
@@ -748,6 +752,51 @@ def alltoall(self, x, mesh_axis, split_axis, concat_axis):
     """
     raise NotImplementedError("Alltoall not implemented")
 
+  def receive(self, x, mesh_axis, source_pcoord):
+    """Collective receive in groups.
+
+    Each group contains the processors that differ only in mesh_axis.
+
+    ```python
+    group_size = self.shape[mesh_axis].size
+    ```
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axis: an integer
+      source_pcoord: a list of optional integers. Each element is either None
+        or an integer in [0, group_size). If source_pcoord[k] is None, then the
+        output for the k-th processor in each group is a zero tensor. If
+        source_pcoord[k] is not None, then the output for the k-th processor in
+        each group is equal to the input for the source_pcoord[k]-th processor
+        in that group.
+
+    Returns:
+      a LaidOutTensor
+    """
+    raise NotImplementedError("Alltoall not implemented")
+
+  def shift_by_n_processors(self, x, mesh_axis, offset, wrap):
+    """Receive the slice from processor pcoord - offset.
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axis: an integer
+      offset: an integer
+      wrap: a boolean. If True, then wrap around. Otherwise, pad with zeros.
+    """
+    n = self.shape[mesh_axis].size
+    source_pcoord = []
+    for i in xrange(n):
+      c = i - offset
+      if c != c % n:
+        if wrap:
+          c = c % n
+        else:
+          c = None
+      source_pcoord.append(c)
+    return self.receive(x, mesh_axis, source_pcoord)
+
   def laid_out_pnum(self):
     """Returns a LaidOutTensor containing the processor number.
 
@@ -1931,6 +1980,225 @@ def add_counter_fn():
     lowering.add_counter("einsum_unique", computation_shape.size)
 
 
+class Conv2dOperation(Operation):
+  """like tf.nn.conv2d.
+
+  Always "NHWC".
+  Always padding="SAME"
+  Always stride 1
+  Always dilation 1
+
+  TODO(noam): implement more options.
+  """
+
+  def __init__(self, conv_input, conv_filter, is_backprop=False, name=None):
+    super(Conv2dOperation, self).__init__(
+        [conv_input, conv_filter], name=name or "conv2d")
+    self._n_dim, self._h_dim, self._w_dim, self._in_dim = conv_input.shape.dims
+    self._fh_dim, self._fw_dim = conv_filter.shape.dims[:2]
+    if is_backprop:
+      self._out_dim, f_in_dim = conv_filter.shape.dims[2:]
+    else:
+      f_in_dim, self._out_dim = conv_filter.shape.dims[2:]
+    self._is_backprop = is_backprop
+    if f_in_dim != self._in_dim:
+      raise ValueError("Dimensions do not match input=%s filter=%s"
+                       % (conv_input, conv_filter))
+    output_shape = Shape([self._n_dim, self._h_dim, self._w_dim, self._out_dim])
+    self._outputs = [Tensor(self, output_shape, conv_input.dtype)]
+
+  def gradient(self, grad_ys):
+    if self._is_backprop:
+      raise ValueError("Gradient not implemented for conv backprop")
+    dy = grad_ys[0]
+    conv_input, conv_filter = self.inputs
+    return [
+        conv2d(dy, conv_filter, is_backprop=True),
+        conv2d_backprop_filter(conv_input, self.inputs[1].shape, dy)]
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    conv_input, conv_filter = self.inputs
+    # TODO(noam): support splitting h_dim, w_dim
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._h_dim) is not None:
+      raise ValueError("can't slice along dimension h")
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._w_dim) is not None:
+      raise ValueError("can't slice along dimension w")
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._fh_dim) is not None:
+      raise ValueError("can't slice along dimension fh")
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._fw_dim) is not None:
+      raise ValueError("can't slice along dimension fw")
+    def tf_fn(tf_input, tf_filter):
+      if self._is_backprop:
+        input_sizes = mesh_impl.slice_shape(self.outputs[0].shape)
+        return tf.nn.conv2d_backprop_input(
+            input_sizes, tf_filter, tf_input,
+            strides=[1, 1, 1, 1], padding="SAME")
+      else:
+        return tf.nn.conv2d(
+            tf_input, tf_filter, strides=[1, 1, 1, 1], padding="SAME")
+    y = mesh_impl.slicewise(
+        tf_fn, lowering.tensors[conv_input], lowering.tensors[conv_filter])
+    out_mesh_axis = mesh_impl.tensor_dimension_to_mesh_axis(self._out_dim)
+    if out_mesh_axis is not None:
+      def add_counter_fn():
+        lowering.add_counter(
+            "allreduce/%s/conv2d_op" % [out_mesh_axis],
+            mesh_impl.laid_out_size(self.outputs[0].shape))
+      y = LazyAllreduceSum(mesh_impl, y, [out_mesh_axis], add_counter_fn)
+    lowering.set_tensor_lowering(self.outputs[0], y)
+    input_shape_set = set(sum([x.shape.dims for x in self.inputs], []))
+    computation_shape = Shape(list(input_shape_set))
+    lowering.add_counter("conv2d", mesh_impl.laid_out_size(computation_shape))
+    lowering.add_counter("conv2d_unique", computation_shape.size)
+
+
+def conv2d(conv_input, conv_filter, is_backprop=False, name=None):
+  """conv2d."""
+  return Conv2dOperation(
+      conv_input, conv_filter, is_backprop, name=name).outputs[0]
+
+
+class Conv2dBackpropFilterOperation(Operation):
+  """like tf.nn.conv2d_backprop_filter."""
+
+  def __init__(self, conv_input, filter_shape, dy, name=None):
+    super(Conv2dBackpropFilterOperation, self).__init__(
+        [conv_input, dy], name=name or "conv2d_backprop_filter")
+    self._n_dim, self._h_dim, self._w_dim, self._in_dim = conv_input.shape.dims
+    dy_n_dim, dy_h_dim, dy_w_dim, self._out_dim = dy.shape.dims
+    self._fh_dim, self._fw_dim, f_in_dim, f_out_dim = filter_shape.dims
+    if (dy_n_dim != self._n_dim or
+        dy_h_dim != self._h_dim or
+        dy_w_dim != self._w_dim or
+        f_in_dim != self._in_dim or
+        f_out_dim != self._out_dim):
+      raise ValueError("Dimensions do not match input=%s dy=%s filter=%s"
+                       % (conv_input, dy, filter_shape))
+    self._outputs = [Tensor(self, filter_shape, conv_input.dtype)]
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    conv_input, dy = self.inputs
+    # TODO(noam): support splitting h_dim, w_dim
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._h_dim) is not None:
+      raise ValueError("can't slice along dimension h")
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._w_dim) is not None:
+      raise ValueError("can't slice along dimension w")
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._fh_dim) is not None:
+      raise ValueError("can't slice along dimension fh")
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._fw_dim) is not None:
+      raise ValueError("can't slice along dimension fw")
+    filter_sizes = mesh_impl.slice_shape(self.outputs[0].shape)
+    def tf_fn(tf_input, tf_dy):
+      return tf.nn.conv2d_backprop_filter(
+          tf_input, filter_sizes, tf_dy, strides=[1, 1, 1, 1], padding="SAME")
+    y = mesh_impl.slicewise(
+        tf_fn, lowering.tensors[conv_input], lowering.tensors[dy])
+    reduced_mesh_axes = [
+        mesh_impl.tensor_dimension_to_mesh_axis(d)
+        for d in [self._n_dim, self._h_dim, self._w_dim]]
+    reduced_mesh_axes = [a for a in reduced_mesh_axes if a is not None]
+    if reduced_mesh_axes:
+      def add_counter_fn():
+        lowering.add_counter(
+            "allreduce/%s/conv2d_op" % (reduced_mesh_axes,),
+            mesh_impl.laid_out_size(self.outputs[0].shape))
+      y = LazyAllreduceSum(mesh_impl, y, reduced_mesh_axes, add_counter_fn)
+    lowering.set_tensor_lowering(self.outputs[0], y)
+    input_shape_set = set(sum([x.shape.dims for x in self.inputs], []))
+    computation_shape = Shape(list(input_shape_set))
+    lowering.add_counter("conv2d", mesh_impl.laid_out_size(computation_shape))
+    lowering.add_counter("conv2d_unique", computation_shape.size)
+
+
+def conv2d_backprop_filter(
+    conv_input, filter_shape, dy, name=None):
+  """conv2d."""
+  return Conv2dBackpropFilterOperation(
+      conv_input, filter_shape, dy, name=name).outputs[0]
+
+
+class ShiftOperation(Operation):
+  """Shift by a static offset in one dimension."""
+
+  def __init__(self, x, offset, dim, wrap, name=None):
+    """Create a shift operation.
+
+    Shift x right by +offset in dimension dim.
+    If offset is negative, shift left.
+    If wrap is true then wrap-around.  Else, pad with zeros.
+
+    Args:
+      x: a Tensor
+      offset: an integer
+      dim: a Dimension of x
+      wrap: a boolean - whether to wrap or pad.
+      name: an optional string
+    """
+    super(ShiftOperation, self).__init__([x], name=name or "shift")
+    self._dim = dim
+    self._axis = x.shape.dims.index(dim)
+    self._offset = offset
+    self._wrap = wrap
+    self._outputs = [Tensor(self, x.shape, x.dtype)]
+
+  def gradient(self, grad_ys):
+    return [shift(grad_ys[0], -self._offset, self._dim, self._wrap)]
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    mesh_axis = mesh_impl.tensor_dimension_to_mesh_axis(self._dim)
+    inputs = self._inputs[0]
+    ndims = self._inputs[0].shape.ndims
+    axis = self._axis
+    dim = self._dim
+    if mesh_axis is None:
+      def slicewise_fn(x):
+        """Slicewise function."""
+        def my_slice(start, size):
+          begin = [0] * axis + [start] + [0] * (ndims - axis - 1)
+          size = [-1] * axis + [size] + [-1] * (ndims - axis - 1)
+          return tf.slice(x, begin, size)
+        def my_pad(s, begin_pad, end_pad):
+          paddings = ([[0, 0]] * axis + [begin_pad, end_pad]
+                      + [[0, 0]] * (ndims - axis - 1))
+          return tf.pad(s, paddings)
+        if self._wrap:
+          offset = self._offset % dim.size
+          return tf.concat([my_slice(dim.size - offset, offset),
+                            my_slice(0, dim.size - offset)], axis=axis)
+        elif self._offset > 0:
+          return my_pad(my_slice(0, dim.size - self._offset), self._offset, 0)
+        else:
+          neg_offset = -self._offset
+          return my_pad(
+              my_slice(neg_offset, dim.size - neg_offset), 0, neg_offset)
+      y = mesh_impl.slicewise(slicewise_fn, lowering.tensors[inputs])
+    else:
+      raise NotImplementedError(
+          "TODO(noam): implement this using mesh_impl.shift_by_n_processors")
+    lowering.set_tensor_lowering(self.outputs[0], y)
+
+
+def shift(x, offset, dim, wrap, name=None):
+  """Shift operation.
+
+  Shift x right by +offset in dimension dim.
+
+  Args:
+    x: a Tensor
+    offset: an integer. If negative, shift left instead of right.
+    dim: a Dimension of x
+    wrap: a boolean - whether to wrap (True) or pad with zeros (False).
+    name: an optional string
+
+  Returns:
+    a Tensor with the same shape and dtype as x
+  """
+  return ShiftOperation(x, offset, dim, wrap, name=name).outputs[0]
+
+
 class SliceOperation(Operation):
   """tf.slice.
 
diff --git a/tensor2tensor/mesh_tensorflow/mnist.py b/tensor2tensor/mesh_tensorflow/mnist.py
index 2bf870160..74df1aec5 100644
--- a/tensor2tensor/mesh_tensorflow/mnist.py
+++ b/tensor2tensor/mesh_tensorflow/mnist.py
@@ -12,7 +12,10 @@
 # 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.
-"""Mnist using mesh-tensrflow and tf.Estimator."""
+"""Mnist using mesh-tensorflow and tf.Estimator.
+
+This is an illustration of mesh-tensorflow, not a good model.
+"""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -40,7 +43,7 @@
                         "Total number of evaluation steps. If `0`, evaluation "
                         "after training is skipped.")
 tf.flags.DEFINE_string("mesh_shape", "rows:2;cols:2", "mesh shape")
-tf.flags.DEFINE_string("layout", "batch:rows;hidden1:cols",
+tf.flags.DEFINE_string("layout", "batch:rows;hidden1:cols,filters1:cols",
                        "layout rules")
 
 FLAGS = tf.flags.FLAGS
@@ -62,11 +65,31 @@ def mnist_model(image, labels, mesh):
   rows_dim = mtf.Dimension("rows", 28)
   cols_dim = mtf.Dimension("cols", 28)
   classes_dim = mtf.Dimension("classes", 10)
-  hidden_dim1 = mtf.Dimension("hidden1", FLAGS.hidden_size)
-  hidden_dim2 = mtf.Dimension("hidden2", FLAGS.hidden_size)
+  one_channel_dim = mtf.Dimension("one_channel", 1)
 
   x = mtf.import_tf_tensor(mesh, tf.reshape(image, [-1, 28, 28]),
                            mtf.Shape([batch_dim, rows_dim, cols_dim]))
+  x = mtf.reshape(x, [batch_dim, rows_dim, cols_dim, one_channel_dim])
+
+  # add some convolutional layers to demonstrate that convolution works.
+  # TODO(noam): get spatially-partitioned convolution working.
+  fh_dim = mtf.Dimension("fh", 3)
+  fw_dim = mtf.Dimension("fw", 3)
+  filters1_dim = mtf.Dimension("filters1", 32)
+  filters2_dim = mtf.Dimension("filters2", 32)
+  kernel1 = mtf.get_variable(
+      mesh, "kernel1", [fh_dim, fw_dim, one_channel_dim, filters1_dim])
+  kernel2 = mtf.get_variable(
+      mesh, "kernel2", [fh_dim, fw_dim, filters1_dim, filters2_dim])
+
+  f1 = mtf.relu(mtf.conv2d(x, kernel1))
+  f2 = mtf.relu(mtf.conv2d(f1, kernel2))
+  x = mtf.reduce_mean(f2, reduced_dim=filters2_dim)
+
+  # add some fully-connected dense layers.
+  hidden_dim1 = mtf.Dimension("hidden1", FLAGS.hidden_size)
+  hidden_dim2 = mtf.Dimension("hidden2", FLAGS.hidden_size)
+
   h1 = mtf_layers.dense(
       x, hidden_dim1, reduced_dims=[rows_dim, cols_dim],
       activation=mtf.relu, name="hidden1")
diff --git a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
index 915cb9bec..0d0c87c0c 100644
--- a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
@@ -203,6 +203,43 @@ def alltoall(self, x, mesh_axis, split_axis, concat_axis):
         functools.partial(
             alltoall_ring, split_axis=split_axis, concat_axis=concat_axis))
 
+  def receive(self, x, mesh_axis, source_pcoord):
+    """Collective receive in groups.
+
+    Each group contains the processors that differ only in mesh_axis.
+
+    ```python
+    group_size = self.shape[mesh_axis].size
+    ```
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axis: an integer
+      source_pcoord: a list of optional integers. Each element is either None
+        or an integer in [0, group_size). If source_pcoord[k] is None, then the
+        output for the k-th processor in each group is a zero tensor. If
+        source_pcoord[k] is not None, then the output for the k-th processor in
+        each group is equal to the input for the source_pcoord[k]-th processor
+        in that group.
+
+    Returns:
+      a LaidOutTensor
+    """
+    x = x.to_laid_out_tensor()
+    shape = x.tensor_list[0].shape
+    dtype = x.tensor_list[0].dtype
+    def _collective_receive(tensor_list, device_list):
+      ret = []
+      for pcoord, device in enumerate(device_list):
+        with tf.device(device):
+          if source_pcoord[pcoord] is None:
+            ret.append(tf.zeros(shape, dtype))
+          else:
+            ret.append(tf.identity(tensor_list[source_pcoord[pcoord]]))
+      return ret
+    return self._collective_with_groups(
+        x, [mesh_axis], _collective_receive)
+
   def _collective_with_groups(self, x, mesh_axes, collective):
     """Grouped collective, (across the given dimensions).
 
diff --git a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
index b689e5c7f..8b0b08565 100644
--- a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
@@ -193,6 +193,8 @@ def allreduce(self, x, mesh_axes, reduction_fn_string):
   def allconcat(self, x, mesh_axis, concat_axis, stack=False):
     """Grouped allconcat (like MPI allgather followed by concat).
 
+    TODO(noam): inefficient - replace with a XLA allconcat when available
+
     Args:
       x: a LaidOutTensor
       mesh_axis: an integer - the mesh axis along which to group
@@ -236,6 +238,47 @@ def alltoall(self, x, mesh_axis, split_axis, concat_axis):
     x = self.allsplit(x, mesh_axis, split_axis)
     return x
 
+  def receive(self, x, mesh_axis, source_pcoord):
+    """Collective receive in groups.
+
+    TODO(noam): inefficient - replace with XLA collective-receive when available
+
+    Each group contains the processors that differ only in mesh_axis.
+
+    ```python
+    group_size = self.shape[mesh_axis].size
+    ```
+
+    Args:
+      x: a LaidOutTensor
+      mesh_axis: an integer
+      source_pcoord: a list of optional integers. Each element is either None
+        or an integer in [0, group_size). If source_pcoord[k] is None, then the
+        output for the k-th processor in each group is a zero tensor. If
+        source_pcoord[k] is not None, then the output for the k-th processor in
+        each group is equal to the input for the source_pcoord[k]-th processor
+        in that group.
+
+    Returns:
+      a LaidOutTensor
+    """
+    x = x.to_laid_out_tensor()
+    x = self.allconcat(x, mesh_axis, concat_axis=0)
+    pcoord = self.laid_out_pcoord(mesh_axis).one_slice
+    # allsplit will barf on Nones, so replace them with something legal.
+    # we will zero out below.
+    source_pcoord_no_nones = [
+        i if c is None else c for i, c in enumerate(source_pcoord)]
+    which = tf.gather(source_pcoord_no_nones, pcoord)
+    x = self.allsplit(
+        x, mesh_axis, split_axis=0, which=self.LaidOutTensor([which]))
+    if None in source_pcoord:
+      # zero out the outputs for which source_pcoord[pcoord]==None
+      source_pcoord_mask = [0.0 if c is None else 1.0 for c in source_pcoord]
+      gathered_mask = tf.gather(source_pcoord_mask, pcoord)
+      x = self.LaidOutTensor([x.one_slice * gathered_mask])
+    return x
+
   def slice(self, tf_tensor, tensor_shape):
     """"Slice out the correspoding part of tensor given the pnum variable."""
     tensor_layout = self.tensor_layout(tensor_shape)

From 01fc64788d05cb785f8acc38dcfad47baad5c359 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 22 Aug 2018 10:48:18 -0700
Subject: [PATCH 0638/2720] stochastic version of the basic model.

PiperOrigin-RevId: 209791306
---
 tensor2tensor/layers/common_video.py          | 107 +++++++++++++
 tensor2tensor/models/__init__.py              |   3 +-
 .../models/research/next_frame_base_vae.py    |  84 ++++++++++
 ...e.py => next_frame_basic_deterministic.py} |  18 ++-
 ... next_frame_basic_deterministic_params.py} | 117 ++------------
 .../research/next_frame_basic_stochastic.py   |  72 +++++++++
 .../models/research/next_frame_emily.py       |  18 ++-
 .../models/research/next_frame_savp.py        |   6 +-
 .../models/research/next_frame_savp_params.py |  41 +++++
 .../models/research/next_frame_sv2p.py        | 143 ++----------------
 .../models/research/next_frame_sv2p_params.py |  74 +++++++++
 .../models/research/next_frame_test.py        |  59 +++++---
 tensor2tensor/rl/trainer_model_based.py       |   2 +-
 tensor2tensor/utils/t2t_model.py              |   4 +
 14 files changed, 475 insertions(+), 273 deletions(-)
 create mode 100644 tensor2tensor/models/research/next_frame_base_vae.py
 rename tensor2tensor/models/research/{next_frame.py => next_frame_basic_deterministic.py} (92%)
 rename tensor2tensor/models/research/{next_frame_params.py => next_frame_basic_deterministic_params.py} (54%)
 create mode 100644 tensor2tensor/models/research/next_frame_basic_stochastic.py
 create mode 100644 tensor2tensor/models/research/next_frame_savp_params.py
 create mode 100644 tensor2tensor/models/research/next_frame_sv2p_params.py

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index f04938108..1f7a27459 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -21,6 +21,7 @@
 import tensorflow as tf
 
 tfl = tf.layers
+tfcl = tf.contrib.layers
 
 
 def swap_time_and_batch_axes(inputs):
@@ -268,3 +269,109 @@ def tile_and_concat(image, latent, concat_latent=True):
   return tf.concat([image, latent], axis=-1)
 
 
+
+
+def tinyify(array, tiny_mode):
+  if tiny_mode:
+    return [1 for _ in array]
+  return array
+
+
+def get_gaussian_tensor(mean, log_var):
+  z = tf.random_normal(tf.shape(mean), 0, 1, dtype=tf.float32)
+  z = mean + tf.exp(log_var / 2.0) * z
+  return z
+
+
+def conv_latent_tower(images, time_axis, latent_channels=1, min_logvar=-5,
+                      is_training=False, random_latent=False, tiny_mode=False):
+  """Builds convolutional latent tower for stochastic model.
+
+  At training time this tower generates a latent distribution (mean and std)
+  conditioned on the entire video. This latent variable will be fed to the
+  main tower as an extra variable to be used for future frames prediction.
+  At inference time, the tower is disabled and only returns latents sampled
+  from N(0,1).
+  If the multi_latent flag is on, a different latent for every timestep would
+  be generated.
+
+  Args:
+    images: tensor of ground truth image sequences
+    time_axis: the time axis  in images tensor
+    latent_channels: number of latent channels
+    min_logvar: minimum value for log_var
+    is_training: whether or not it is training mode
+    random_latent: whether or not generate random latents
+    tiny_mode: whether or not it is tiny_mode
+  Returns:
+    latent_mean: predicted latent mean
+    latent_logvar: predicted latent log variance
+  """
+  conv_size = tinyify([32, 64, 64], tiny_mode)
+  with tf.variable_scope("latent", reuse=tf.AUTO_REUSE):
+    images = tf.to_float(images)
+    images = tf.unstack(images, axis=time_axis)
+    images = tf.concat(images, axis=3)
+
+    x = images
+    x = common_layers.make_even_size(x)
+    x = tfl.conv2d(x, conv_size[0], [3, 3], strides=(2, 2),
+                   padding="SAME", activation=tf.nn.relu, name="latent_conv1")
+    x = tfcl.batch_norm(x, updates_collections=None,
+                        is_training=is_training, scope="latent_bn1")
+    x = common_layers.make_even_size(x)
+    x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
+                   padding="SAME", activation=tf.nn.relu, name="latent_conv2")
+    x = tfcl.batch_norm(x, updates_collections=None,
+                        is_training=is_training, scope="latent_bn2")
+    x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(1, 1),
+                   padding="SAME", activation=tf.nn.relu, name="latent_conv3")
+    x = tfcl.batch_norm(x, updates_collections=None,
+                        is_training=is_training, scope="latent_bn3")
+
+    nc = latent_channels
+    mean = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
+                      padding="SAME", activation=None, name="latent_mean")
+    logv = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
+                      padding="SAME", activation=tf.nn.relu, name="latent_std")
+    logvar = logv + min_logvar
+
+    # No latent tower at inference time, just standard gaussian.
+    if not is_training:
+      return tf.zeros_like(mean), tf.zeros_like(logvar)
+
+    # No latent in the first phase
+    ret_mean, ret_logvar = tf.cond(
+        random_latent,
+        lambda: (tf.zeros_like(mean), tf.zeros_like(logvar)),
+        lambda: (mean, logvar))
+
+    return ret_mean, ret_logvar
+
+
+def beta_schedule(schedule, global_step, final_beta, decay_start, decay_end):
+  """Get KL multiplier (beta) based on the schedule."""
+  # TODO(mechcoder): Add log_annealing schedule.
+  if schedule == "constant":
+    beta = tf.cond(
+        tf.less(global_step, decay_start), lambda: 0.0, lambda: final_beta)
+  elif schedule == "linear_anneal":
+    # Linearly anneal beta from 0.0 to self.hparams.latent_loss_multiplier.
+    # between self.hparams.num_iterations_2nd_stage to anneal_end.
+    # beta = latent_loss * (1 - (global_step - 2nd_stage) / (anneal_end - 2nd_stage))  # pylint:disable=line-too-long
+    if decay_start > decay_end:
+      raise ValueError("decay_end is smaller than decay_end.")
+
+    def anneal_loss(step_num):
+      step_num = tf.cast(step_num, dtype=tf.float32)
+      fraction = (float(decay_end) - step_num) / (decay_end - decay_start)
+      return final_beta * (1 - fraction)
+
+    beta = tf.case(
+        pred_fn_pairs={
+            tf.less(global_step, decay_start): lambda: 0.0,
+            tf.greater(global_step, decay_end): lambda: final_beta},
+        default=lambda: anneal_loss(global_step))
+  else:
+    raise ValueError("Unknown beta schedule.")
+  return beta
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index b6af0ed00..b41892bd2 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -46,7 +46,8 @@
 from tensor2tensor.models.research import gene_expression
 from tensor2tensor.models.research import lm_experiments
 from tensor2tensor.models.research import multimodel
-from tensor2tensor.models.research import next_frame
+from tensor2tensor.models.research import next_frame_basic_deterministic
+from tensor2tensor.models.research import next_frame_basic_stochastic
 from tensor2tensor.models.research import next_frame_emily
 from tensor2tensor.models.research import next_frame_savp
 from tensor2tensor.models.research import next_frame_sv2p
diff --git a/tensor2tensor/models/research/next_frame_base_vae.py b/tensor2tensor/models/research/next_frame_base_vae.py
new file mode 100644
index 000000000..8d5da6303
--- /dev/null
+++ b/tensor2tensor/models/research/next_frame_base_vae.py
@@ -0,0 +1,84 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Basic models for testing simple tasks."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
+
+import tensorflow as tf
+
+
+class NextFrameBaseVae(object):
+  """Basic function for stochastic variational video prediction."""
+
+  def __init__(self, hparams):
+    self.hparams = hparams
+
+  def get_iteration_num(self):
+    step_num = tf.train.get_global_step()
+    # TODO(lukaszkaiser): what should it be if it's undefined?
+    if step_num is None:
+      step_num = 1000000
+    return step_num
+
+  def get_beta(self):
+    beta = common_video.beta_schedule(
+        schedule=self.hparams.latent_loss_multiplier_schedule,
+        global_step=self.get_iteration_num(),
+        final_beta=self.hparams.latent_loss_multiplier,
+        decay_start=(self.hparams.num_iterations_1st_stage +
+                     self.hparams.num_iterations_2nd_stage),
+        decay_end=self.hparams.anneal_end)
+    tf.summary.scalar("beta", beta)
+    return beta
+
+  def get_extra_loss(self, mean, std):
+    """Losses in addition to the default modality losses."""
+    if self.is_training:
+      beta = self.get_beta()
+      kl_loss = common_layers.kl_divergence(mean, std)
+      tf.summary.histogram("posterior_mean", mean)
+      tf.summary.histogram("posterior_std", std)
+      tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
+    return beta * kl_loss
+
+  def construct_latent_tower(self, images, time_axis):
+    """Create the latent tower."""
+    # No latent in the first phase
+    first_phase = tf.less(
+        self.get_iteration_num(), self.hparams.num_iterations_1st_stage)
+
+    # use all frames by default but this allows more
+    # predicted frames at inference time
+    latent_num_frames = self.hparams.latent_num_frames
+    tf.logging.info("Creating latent tower with %d frames." % latent_num_frames)
+    if latent_num_frames > 0:
+      images = images[:latent_num_frames]
+
+    return common_video.conv_latent_tower(
+        images=images,
+        time_axis=time_axis,
+        latent_channels=self.hparams.latent_channels,
+        min_logvar=self.hparams.latent_std_min,
+        is_training=self.is_training,
+        random_latent=first_phase,
+        tiny_mode=self.hparams.tiny_mode)
+
+
+
diff --git a/tensor2tensor/models/research/next_frame.py b/tensor2tensor/models/research/next_frame_basic_deterministic.py
similarity index 92%
rename from tensor2tensor/models/research/next_frame.py
rename to tensor2tensor/models/research/next_frame_basic_deterministic.py
index 6f35f6a78..5cc6b2bd4 100644
--- a/tensor2tensor/models/research/next_frame.py
+++ b/tensor2tensor/models/research/next_frame_basic_deterministic.py
@@ -22,17 +22,26 @@
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
-from tensor2tensor.models.research import next_frame_params  # pylint: disable=unused-import
+from tensor2tensor.models.research import next_frame_basic_deterministic_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
 import tensorflow as tf
 
 
+tfl = tf.layers
+tfcl = tf.contrib.layers
+
+
 @registry.register_model
-class NextFrameBasic(t2t_model.T2TModel):
+class NextFrameBasicDeterministic(t2t_model.T2TModel):
   """Basic next-frame model, may take actions and predict rewards too."""
 
+  def inject_latent(self, layer, features, filters):
+    """Do nothing for deterministic model."""
+    del features, filters
+    return layer, 0.0
+
   def body(self, features):
     hparams = self.hparams
     filters = hparams.hidden_size
@@ -70,6 +79,8 @@ def body(self, features):
       else:
         x *= action_mask + zeros_mask
 
+    x, extra_loss = self.inject_latent(x, features, filters)
+
     # Run a stack of convolutions.
     for i in range(hparams.num_hidden_layers):
       with tf.variable_scope("layer%d" % i):
@@ -103,7 +114,7 @@ def body(self, features):
     if "target_reward" not in features:
       return x
     reward_pred = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
-    return {"targets": x, "target_reward": reward_pred}
+    return {"targets": x, "target_reward": reward_pred}, extra_loss
 
   def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
     """Produce predictions from the model by running it."""
@@ -158,3 +169,4 @@ def logits_to_samples(logits):
 
     # Return results.
     return results
+
diff --git a/tensor2tensor/models/research/next_frame_params.py b/tensor2tensor/models/research/next_frame_basic_deterministic_params.py
similarity index 54%
rename from tensor2tensor/models/research/next_frame_params.py
rename to tensor2tensor/models/research/next_frame_basic_deterministic_params.py
index a0d4294ed..2fd5a6958 100644
--- a/tensor2tensor/models/research/next_frame_params.py
+++ b/tensor2tensor/models/research/next_frame_basic_deterministic_params.py
@@ -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.
-"""Param sets for next frame prediction models."""
+"""Param sets for deterministic basic next frame prediction model."""
 
 from __future__ import division
 from __future__ import print_function
@@ -22,7 +22,7 @@
 
 
 @registry.register_hparams
-def next_frame():
+def next_frame_basic_deterministic():
   """Basic 2-frame conv model."""
   hparams = common_hparams.basic_params1()
   hparams.video_num_input_frames = 4
@@ -45,119 +45,24 @@ def next_frame():
   hparams.add_hparam("video_modality_loss_cutoff", 0.02)
   hparams.add_hparam("preprocess_resize_frames", None)
   hparams.add_hparam("concatenate_actions", True)
-  hparams.add_hparam("tiny_mode", False)
   hparams.add_hparam("shuffle_buffer_size", 128)
+  hparams.add_hparam("tiny_mode", False)
+  hparams.add_hparam("stochastic_model", False)
   return hparams
 
 
 @registry.register_hparams
 def next_frame_pixel_noise():
   """Basic 2-frame conv model with pixel noise."""
-  hparams = next_frame()
+  hparams = next_frame_basic_deterministic()
   hparams.add_hparam("video_modality_input_noise", 0.05)
   hparams.input_modalities = "inputs:video:pixel_noise"
   return hparams
 
 
-@registry.register_hparams
-def next_frame_stochastic():
-  """SV2P model."""
-  hparams = next_frame()
-  hparams.optimizer = "TrueAdam"
-  hparams.learning_rate_schedule = "constant"
-  hparams.learning_rate_constant = 1e-3
-  hparams.video_num_input_frames = 1
-  hparams.video_num_target_frames = 3
-  hparams.batch_size = 16
-  hparams.target_modality = "video:l2raw"
-  hparams.input_modalities = "inputs:video:l2raw"
-  hparams.video_modality_loss_cutoff = 0.0
-  hparams.add_hparam("stochastic_model", True)
-  hparams.add_hparam("reward_prediction", True)
-  hparams.add_hparam("reward_prediction_stop_gradient", True)
-  hparams.add_hparam("model_options", "CDNA")
-  hparams.add_hparam("num_masks", 10)
-  hparams.add_hparam("latent_channels", 1)
-  hparams.add_hparam("latent_std_min", -5.0)
-  hparams.add_hparam("num_iterations_1st_stage", 10000)
-  hparams.add_hparam("num_iterations_2nd_stage", 10000)
-  hparams.add_hparam("latent_loss_multiplier", 1e-3)
-  hparams.add_hparam("latent_loss_multiplier_schedule", "constant")
-  hparams.add_hparam("multi_latent", False)
-  hparams.add_hparam("relu_shift", 1e-12)
-  hparams.add_hparam("dna_kernel_size", 5)
-  # Scheduled sampling method. Choose between prob or count.
-  hparams.add_hparam("scheduled_sampling_mode", "count")
-  hparams.add_hparam("scheduled_sampling_decay_steps", 10000)
-  hparams.add_hparam("scheduled_sampling_k", 900.0)
-  hparams.add_hparam("latent_num_frames", 0)  # 0 means use all frames.
-  hparams.add_hparam("anneal_end", 100000)
-  hparams.add_hparam("upsample_method", "conv2d_transpose")
-  hparams.add_hparam("internal_loss", False)
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_stochastic_emily():
-  """Emily's model."""
-  hparams = next_frame_stochastic()
-  hparams.latent_loss_multiplier = 1e-4
-  hparams.learning_rate_constant = 0.002
-  hparams.add_hparam("z_dim", 10)
-  hparams.add_hparam("g_dim", 128)
-  hparams.add_hparam("rnn_size", 256)
-  hparams.add_hparam("posterior_rnn_layers", 1)
-  hparams.add_hparam("predictor_rnn_layers", 2)
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_savp():
-  """SAVP model."""
-  hparams = next_frame_stochastic()
-  hparams.add_hparam("z_dim", 8)
-  hparams.add_hparam("num_discriminator_filters", 32)
-  hparams.add_hparam("use_vae", True)
-  hparams.add_hparam("use_gan", False)
-  hparams.add_hparam("use_spectral_norm", True)
-  hparams.add_hparam("gan_loss", "cross_entropy")
-  hparams.add_hparam("gan_loss_multiplier", 0.01)
-  hparams.add_hparam("gan_vae_loss_multiplier", 0.01)
-  hparams.add_hparam("gan_optimization", "joint")
-  hparams.target_modality = "video:l1raw"
-  hparams.input_modalities = "inputs:video:l1raw"
-  hparams.latent_loss_multiplier_schedule = "linear_anneal"
-  hparams.anneal_end = 100000
-  hparams.upsample_method = "bilinear_upsample_conv"
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_stochastic_cutoff():
-  """SV2P model with additional cutoff in L2 loss for environments like pong."""
-  hparams = next_frame_stochastic()
-  hparams.video_modality_loss_cutoff = 0.4
-  hparams.video_num_input_frames = 4
-  hparams.video_num_target_frames = 1
-  return hparams
-
-
-@registry.register_hparams
-def next_frame_stochastic_tiny():
-  """SV2P model with additional cutoff in L2 loss for environments like pong."""
-  hparams = next_frame_stochastic()
-  hparams.batch_size = 2
-  hparams.tiny_mode = True
-  hparams.num_masks = 1
-  hparams.video_modality_loss_cutoff = 0.4
-  hparams.video_num_input_frames = 4
-  hparams.video_num_target_frames = 1
-  return hparams
-
-
 @registry.register_hparams
 def next_frame_tpu():
-  hparams = next_frame()
+  hparams = next_frame_basic_deterministic()
   hparams.batch_size = 1
   return hparams
 
@@ -165,7 +70,7 @@ def next_frame_tpu():
 @registry.register_hparams
 def next_frame_ae():
   """Conv autoencoder."""
-  hparams = next_frame()
+  hparams = next_frame_basic_deterministic()
   hparams.input_modalities = "inputs:video:bitwise"
   hparams.hidden_size = 256
   hparams.batch_size = 8
@@ -178,7 +83,7 @@ def next_frame_ae():
 @registry.register_hparams
 def next_frame_small():
   """Small conv model."""
-  hparams = next_frame()
+  hparams = next_frame_basic_deterministic()
   hparams.hidden_size = 32
   return hparams
 
@@ -186,7 +91,7 @@ def next_frame_small():
 @registry.register_hparams
 def next_frame_tiny():
   """Tiny for testing."""
-  hparams = next_frame()
+  hparams = next_frame_basic_deterministic()
   hparams.hidden_size = 32
   hparams.num_hidden_layers = 1
   hparams.num_compress_steps = 2
@@ -197,7 +102,7 @@ def next_frame_tiny():
 @registry.register_hparams
 def next_frame_l1():
   """Basic conv model with L1 modality."""
-  hparams = next_frame()
+  hparams = next_frame_basic_deterministic()
   hparams.target_modality = "video:l1"
   hparams.video_modality_loss_cutoff = 2.4
   return hparams
@@ -206,7 +111,7 @@ def next_frame_l1():
 @registry.register_hparams
 def next_frame_l2():
   """Basic conv model with L2 modality."""
-  hparams = next_frame()
+  hparams = next_frame_basic_deterministic()
   hparams.target_modality = "video:l2"
   hparams.video_modality_loss_cutoff = 2.4
   return hparams
diff --git a/tensor2tensor/models/research/next_frame_basic_stochastic.py b/tensor2tensor/models/research/next_frame_basic_stochastic.py
new file mode 100644
index 000000000..4356c43c1
--- /dev/null
+++ b/tensor2tensor/models/research/next_frame_basic_stochastic.py
@@ -0,0 +1,72 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Basic models for testing simple tasks."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
+from tensor2tensor.models.research import next_frame_base_vae
+from tensor2tensor.models.research import next_frame_basic_deterministic
+from tensor2tensor.models.research import next_frame_basic_deterministic_params
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+@registry.register_model
+class NextFrameBasicStochastic(
+    next_frame_basic_deterministic.NextFrameBasicDeterministic,
+    next_frame_base_vae.NextFrameBaseVae):
+  """Stochastic version of basic next-frame model."""
+
+  def inject_latent(self, layer, features, filters):
+    """Do nothing for deterministic model."""
+    # Latent for stochastic model
+    full_video = tf.concat(
+        [features["inputs_raw"], features["targets_raw"]], axis=1)
+    latent_mean, latent_std = self.construct_latent_tower(
+        full_video, time_axis=1)
+    latent = common_video.get_gaussian_tensor(latent_mean, latent_std)
+    latent = tf.layers.flatten(latent)
+    latent = tf.expand_dims(latent, axis=1)
+    latent = tf.expand_dims(latent, axis=1)
+    latent_mask = tf.layers.dense(latent, filters, name="latent_mask")
+    zeros_mask = tf.zeros(
+        common_layers.shape_list(layer)[:-1] + [filters], dtype=tf.float32)
+    layer = tf.concat([layer, latent_mask + zeros_mask], axis=-1)
+    extra_loss = self.get_extra_loss(latent_mean, latent_std)
+    return layer, extra_loss
+
+
+@registry.register_hparams
+def next_frame_basic_stochastic():
+  """Basic 2-frame conv model with stochastic tower."""
+  base = next_frame_basic_deterministic_params
+  hparams = base.next_frame_basic_deterministic()
+  hparams.stochastic_model = True
+  hparams.add_hparam("latent_channels", 1)
+  hparams.add_hparam("latent_std_min", -5.0)
+  hparams.add_hparam("num_iterations_1st_stage", 25000)
+  hparams.add_hparam("num_iterations_2nd_stage", 25000)
+  hparams.add_hparam("latent_loss_multiplier", 1e-3)
+  hparams.add_hparam("latent_loss_multiplier_schedule", "constant")
+  hparams.add_hparam("latent_num_frames", 0)  # 0 means use all frames.
+  hparams.add_hparam("anneal_end", 100000)
+  return hparams
+
+
diff --git a/tensor2tensor/models/research/next_frame_emily.py b/tensor2tensor/models/research/next_frame_emily.py
index ff2275ede..9790c2148 100644
--- a/tensor2tensor/models/research/next_frame_emily.py
+++ b/tensor2tensor/models/research/next_frame_emily.py
@@ -29,8 +29,8 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
-from tensor2tensor.models.research import next_frame_params  # pylint: disable=unused-import
 from tensor2tensor.models.research import next_frame_sv2p
+from tensor2tensor.models.research import next_frame_sv2p_params
 from tensor2tensor.utils import registry
 import tensorflow as tf
 
@@ -39,7 +39,7 @@
 
 
 @registry.register_model
-class NextFrameStochasticEmily(next_frame_sv2p.NextFrameStochastic):
+class NextFrameEmily(next_frame_sv2p.NextFrameSv2p):
   """Stochastic Variational Video Prediction Without Learned Prior."""
 
   def encoder(self, inputs, nout):
@@ -261,3 +261,17 @@ def construct_model(self, images, actions, rewards):
     tf.logging.info(">>>> Done")
     gen_images = tf.stack(gen_images, axis=0)
     return gen_images, fake_reward_prediction, pred_mu, pred_logvar
+
+
+@registry.register_hparams
+def next_frame_emily():
+  """Emily's model hparams."""
+  hparams = next_frame_sv2p_params.next_frame_sv2p()
+  hparams.latent_loss_multiplier = 1e-4
+  hparams.learning_rate_constant = 0.002
+  hparams.add_hparam("z_dim", 10)
+  hparams.add_hparam("g_dim", 128)
+  hparams.add_hparam("rnn_size", 256)
+  hparams.add_hparam("posterior_rnn_layers", 1)
+  hparams.add_hparam("predictor_rnn_layers", 2)
+  return hparams
diff --git a/tensor2tensor/models/research/next_frame_savp.py b/tensor2tensor/models/research/next_frame_savp.py
index 49775f646..3d2307784 100644
--- a/tensor2tensor/models/research/next_frame_savp.py
+++ b/tensor2tensor/models/research/next_frame_savp.py
@@ -25,7 +25,7 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
-from tensor2tensor.models.research import next_frame_params  # pylint: disable=unused-import
+from tensor2tensor.models.research import next_frame_savp_params  # pylint: disable=unused-import
 from tensor2tensor.models.research import next_frame_sv2p
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import update_ops_hook
@@ -35,7 +35,7 @@
 
 
 @registry.register_model
-class NextFrameSAVP(next_frame_sv2p.NextFrameStochastic):
+class NextFrameSAVP(next_frame_sv2p.NextFrameSv2p):
   """Stochastic Adversarial Video Prediction."""
 
   def encoder(self, inputs, n_layers=3):
@@ -397,7 +397,7 @@ def construct_model(self, images, actions, rewards):
 
       for step, (image, action, reward, mu, log_sigma_sq) in enumerate(iterable):  # pylint:disable=line-too-long
         # Sample latents using a gaussian centered at conditional mu and std.
-        latent = self.get_gaussian_latent(mu, log_sigma_sq)
+        latent = common_video.get_gaussian_tensor(mu, log_sigma_sq)
 
         # Sample prior latents from isotropic normal distribution.
         prior_latent = tf.random_normal(tf.shape(latent), dtype=tf.float32)
diff --git a/tensor2tensor/models/research/next_frame_savp_params.py b/tensor2tensor/models/research/next_frame_savp_params.py
new file mode 100644
index 000000000..253c0a707
--- /dev/null
+++ b/tensor2tensor/models/research/next_frame_savp_params.py
@@ -0,0 +1,41 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Param sets for SAVP model."""
+
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.models.research import next_frame_sv2p_params
+from tensor2tensor.utils import registry
+
+
+@registry.register_hparams
+def next_frame_savp():
+  """SAVP model hparams."""
+  hparams = next_frame_sv2p_params.next_frame_sv2p()
+  hparams.add_hparam("z_dim", 8)
+  hparams.add_hparam("num_discriminator_filters", 32)
+  hparams.add_hparam("use_vae", True)
+  hparams.add_hparam("use_gan", False)
+  hparams.add_hparam("use_spectral_norm", True)
+  hparams.add_hparam("gan_loss", "cross_entropy")
+  hparams.add_hparam("gan_loss_multiplier", 0.01)
+  hparams.add_hparam("gan_vae_loss_multiplier", 0.01)
+  hparams.add_hparam("gan_optimization", "joint")
+  hparams.target_modality = "video:l1raw"
+  hparams.input_modalities = "inputs:video:l1raw"
+  hparams.latent_loss_multiplier_schedule = "linear_anneal"
+  hparams.upsample_method = "bilinear_upsample_conv"
+  return hparams
diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 02d6e7b26..324adf70d 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -28,8 +28,8 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
-from tensor2tensor.models.research import next_frame
-from tensor2tensor.models.research import next_frame_params  # pylint: disable=unused-import
+from tensor2tensor.models.research import next_frame_basic_stochastic
+from tensor2tensor.models.research import next_frame_sv2p_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
 import tensorflow as tf
 
@@ -38,17 +38,11 @@
 
 
 @registry.register_model
-class NextFrameStochastic(next_frame.NextFrameBasic):
+class NextFrameSv2p(next_frame_basic_stochastic.NextFrameBasicStochastic):
   """Stochastic Variational Video Prediction."""
 
-  @property
-  def is_training(self):
-    return self.hparams.mode == tf.estimator.ModeKeys.TRAIN
-
   def tinyify(self, array):
-    if self.hparams.tiny_mode:
-      return [1 for _ in array]
-    return array
+    return common_video.tinyify(array, self.hparams.tiny_mode)
 
   def visualize_predictions(self, real_frames, gen_frames):
     def concat_on_y_axis(x):
@@ -63,56 +57,6 @@ def concat_on_y_axis(x):
     side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
     tf.summary.image("full_video", side_by_side_video)
 
-  def get_gaussian_latent(self, latent_mean, latent_std):
-    latent = tf.random_normal(tf.shape(latent_mean), 0, 1, dtype=tf.float32)
-    latent = latent_mean + tf.exp(latent_std / 2.0) * latent
-    return latent
-
-  def get_iteration_num(self):
-    step_num = tf.train.get_global_step()
-    # TODO(lukaszkaiser): what should it be if it"s undefined?
-    if step_num is None:
-      step_num = 1000000
-    return step_num
-
-  def get_beta(self):
-    """Get KL multiplier (beta) based on the schedule."""
-    step_num = self.get_iteration_num()
-    schedule = self.hparams.latent_loss_multiplier_schedule
-    second_stage = (self.hparams.num_iterations_1st_stage +
-                    self.hparams.num_iterations_2nd_stage)
-    # TODO(mechcoder): Add log_annealing schedule.
-    if schedule == "constant":
-      beta = tf.cond(tf.greater(step_num, second_stage),
-                     lambda: self.hparams.latent_loss_multiplier,
-                     lambda: 0.0)
-    elif schedule == "linear_anneal":
-      # Linearly anneal beta from 0.0 to self.hparams.latent_loss_multiplier.
-      # between self.hparams.num_iterations_2nd_stage to anneal_end.
-      # beta = latent_loss * (1 - (global_step - 2nd_stage) / (anneal_end - 2nd_stage))  # pylint:disable=line-too-long
-      anneal_end = self.hparams.anneal_end
-      latent_multiplier = self.hparams.latent_loss_multiplier
-      if anneal_end < second_stage:
-        raise ValueError("Expected hparams.num_iterations_2nd_stage < "
-                         "hparams.anneal_end %d, got %d." %
-                         (second_stage, anneal_end))
-
-      def anneal_loss(step_num):
-        step_num = tf.cast(step_num, dtype=tf.float32)
-        fraction = (float(anneal_end) - step_num) / (anneal_end - second_stage)
-        return self.hparams.latent_loss_multiplier * (1 - fraction)
-
-      beta = tf.case(
-          pred_fn_pairs={
-              tf.less(step_num, second_stage): lambda: 0.0,
-              tf.greater(step_num, anneal_end): lambda: latent_multiplier},
-          default=lambda: anneal_loss(step_num))
-    else:
-      raise ValueError("Unknown beta schedule.")
-
-    tf.summary.scalar("beta", beta)
-    return beta
-
   def get_scheduled_sample_func(self, batch_size):
     """Creates a function for scheduled sampling based on given hparams."""
     with tf.variable_scope("scheduled_sampling_func", reuse=False):
@@ -277,72 +221,6 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
     hidden5 = tile_and_concat(hidden5, latent, concat_latent=concat_latent)
     return hidden5, (enc0, enc1)
 
-  def construct_latent_tower(self, images):
-    """Builds convolutional latent tower for stochastic model.
-
-    At training time this tower generates a latent distribution (mean and std)
-    conditioned on the entire video. This latent variable will be fed to the
-    main tower as an extra variable to be used for future frames prediction.
-    At inference time, the tower is disabled and only returns latents sampled
-    from N(0,1).
-    If the multi_latent flag is on, a different latent for every timestep would
-    be generated.
-
-    Args:
-      images: tensor of ground truth image sequences
-    Returns:
-      latent_mean: predicted latent mean
-      latent_std: predicted latent standard deviation
-      latent_loss: loss of the latent twoer
-      samples: random samples sampled from standard guassian
-    """
-    conv_size = self.tinyify([32, 64, 64])
-    with tf.variable_scope("latent", reuse=tf.AUTO_REUSE):
-      # this allows more predicted frames at inference time
-      latent_num_frames = self.hparams.latent_num_frames
-      if latent_num_frames == 0:  # use all frames by default.
-        latent_num_frames = (self.hparams.video_num_input_frames +
-                             self.hparams.video_num_target_frames)
-      tf.logging.info("Creating latent tower with %d frames."%latent_num_frames)
-      latent_images = tf.unstack(images[:latent_num_frames], axis=0)
-      images = tf.concat(latent_images, 3)
-
-      x = images
-      x = common_layers.make_even_size(x)
-      x = tfl.conv2d(x, conv_size[0], [3, 3], strides=(2, 2),
-                     padding="SAME", activation=tf.nn.relu, name="latent_conv1")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="latent_bn1")
-      x = common_layers.make_even_size(x)
-      x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
-                     padding="SAME", activation=tf.nn.relu, name="latent_conv2")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="latent_bn2")
-      x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(1, 1),
-                     padding="SAME", activation=tf.nn.relu, name="latent_conv3")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="latent_bn3")
-
-      nc = self.hparams.latent_channels
-      mean = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
-                        padding="SAME", activation=None, name="latent_mean")
-      std = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
-                       padding="SAME", activation=tf.nn.relu, name="latent_std")
-      std += self.hparams.latent_std_min
-
-      # No latent tower at inference time, just standard gaussian.
-      if not self.is_training:
-        return tf.zeros_like(mean), tf.zeros_like(std)
-
-      # No latent in the first phase
-      iter_num = self.get_iteration_num()
-      ret_mean, ret_std = tf.cond(
-          tf.less(iter_num, self.hparams.num_iterations_1st_stage),
-          lambda: (tf.zeros_like(mean), tf.zeros_like(std)),
-          lambda: (mean, std))
-
-      return ret_mean, ret_std
-
   def reward_prediction(self, input_image, input_reward, action, latent):
     """Builds a reward prediction network."""
     del action
@@ -547,8 +425,8 @@ def process_single_frame(prev_outputs, inputs):
     # Latent tower
     latent = None
     if self.hparams.stochastic_model:
-      latent_mean, latent_std = self.construct_latent_tower(images)
-      latent = self.get_gaussian_latent(latent_mean, latent_std)
+      latent_mean, latent_std = self.construct_latent_tower(images, time_axis=0)
+      latent = common_video.get_gaussian_tensor(latent_mean, latent_std)
 
     # HACK: Do first step outside to initialize all the variables
     lstm_states = [None] * 7
@@ -653,7 +531,7 @@ def body(self, features):
 
 
 @registry.register_model
-class NextFrameStochasticTwoFrames(NextFrameStochastic):
+class NextFrameSv2pTwoFrames(NextFrameSv2p):
   """Stochastic next-frame model with 2 frames posterior."""
 
   def construct_model(self, images, actions, rewards):
@@ -687,9 +565,10 @@ def construct_model(self, images, actions, rewards):
 
       # Latent
       # TODO(mbz): should we use input_image iunstead of image?
-      latent_images = [image, images[timestep+1]]
-      latent_mean, latent_std = self.construct_latent_tower(latent_images)
-      latent = self.get_gaussian_latent(latent_mean, latent_std)
+      latent_images = tf.stack([image, images[timestep+1]], axis=0)
+      latent_mean, latent_std = self.construct_latent_tower(
+          latent_images, time_axis=0)
+      latent = common_video.get_gaussian_tensor(latent_mean, latent_std)
       latent_means.append(latent_mean)
       latent_stds.append(latent_std)
 
diff --git a/tensor2tensor/models/research/next_frame_sv2p_params.py b/tensor2tensor/models/research/next_frame_sv2p_params.py
new file mode 100644
index 000000000..8ecb4d08b
--- /dev/null
+++ b/tensor2tensor/models/research/next_frame_sv2p_params.py
@@ -0,0 +1,74 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Param sets for SV2P model."""
+
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.models.research import next_frame_basic_stochastic
+from tensor2tensor.utils import registry
+
+
+@registry.register_hparams
+def next_frame_sv2p():
+  """SV2P model hparams."""
+  hparams = next_frame_basic_stochastic.next_frame_basic_stochastic()
+  hparams.optimizer = "TrueAdam"
+  hparams.learning_rate_schedule = "constant"
+  hparams.learning_rate_constant = 1e-3
+  hparams.video_num_input_frames = 1
+  hparams.video_num_target_frames = 3
+  hparams.batch_size = 16
+  hparams.target_modality = "video:l2raw"
+  hparams.input_modalities = "inputs:video:l2raw"
+  hparams.video_modality_loss_cutoff = 0.0
+  hparams.add_hparam("reward_prediction", True)
+  hparams.add_hparam("reward_prediction_stop_gradient", True)
+  hparams.add_hparam("model_options", "CDNA")
+  hparams.add_hparam("num_masks", 10)
+  hparams.add_hparam("multi_latent", False)
+  hparams.add_hparam("relu_shift", 1e-12)
+  hparams.add_hparam("dna_kernel_size", 5)
+  # Scheduled sampling method. Choose between prob or count.
+  hparams.add_hparam("scheduled_sampling_mode", "count")
+  hparams.add_hparam("scheduled_sampling_decay_steps", 10000)
+  hparams.add_hparam("scheduled_sampling_k", 900.0)
+  hparams.add_hparam("upsample_method", "conv2d_transpose")
+  hparams.add_hparam("internal_loss", False)
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_sv2p_tiny():
+  """Tiny SV2P model."""
+  hparams = next_frame_sv2p()
+  hparams.batch_size = 2
+  hparams.tiny_mode = True
+  hparams.num_masks = 1
+  hparams.video_modality_loss_cutoff = 0.4
+  hparams.video_num_input_frames = 4
+  hparams.video_num_target_frames = 1
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_sv2p_cutoff():
+  """SV2P model with additional cutoff in L2 loss for environments like pong."""
+  hparams = next_frame_sv2p()
+  hparams.video_modality_loss_cutoff = 0.4
+  hparams.video_num_input_frames = 4
+  hparams.video_num_target_frames = 1
+  return hparams
+
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index b797d2b49..73ad1ea31 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -20,11 +20,14 @@
 import numpy as np
 
 from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
-from tensor2tensor.models.research import next_frame
+from tensor2tensor.models.research import next_frame_basic_deterministic
+from tensor2tensor.models.research import next_frame_basic_deterministic_params
+from tensor2tensor.models.research import next_frame_basic_stochastic
 from tensor2tensor.models.research import next_frame_emily
-from tensor2tensor.models.research import next_frame_params
 from tensor2tensor.models.research import next_frame_savp
+from tensor2tensor.models.research import next_frame_savp_params
 from tensor2tensor.models.research import next_frame_sv2p
+from tensor2tensor.models.research import next_frame_sv2p_params
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -191,38 +194,44 @@ def TestOnVariousUpSampleLayers(self, hparams, model, expected_last_dim):
     self.TestVideoModel(4, 1, hparams, model, expected_last_dim,
                         upsample_method="nn_upsample_conv")
 
-  def testBasic(self):
+  def testBasicDeterministic(self):
     self.TestOnVariousInputOutputSizes(
-        next_frame_params.next_frame(),
-        next_frame.NextFrameBasic,
+        next_frame_basic_deterministic_params.next_frame_basic_deterministic(),
+        next_frame_basic_deterministic.NextFrameBasicDeterministic,
         256)
 
-  def testStochastic(self):
+  def testBasicStochastic(self):
     self.TestOnVariousInputOutputSizes(
-        next_frame_params.next_frame_stochastic(),
-        next_frame_sv2p.NextFrameStochastic,
+        next_frame_basic_stochastic.next_frame_basic_stochastic(),
+        next_frame_basic_stochastic.NextFrameBasicStochastic,
+        256)
+
+  def testSv2p(self):
+    self.TestOnVariousInputOutputSizes(
+        next_frame_sv2p_params.next_frame_sv2p(),
+        next_frame_sv2p.NextFrameSv2p,
         1)
 
-  def testStochasticWithActionsAndRewards(self):
+  def testSv2pWithActionsAndRewards(self):
     self.TestWithActionAndRewards(
-        next_frame_params.next_frame_stochastic(),
-        next_frame_sv2p.NextFrameStochastic,
+        next_frame_sv2p_params.next_frame_sv2p(),
+        next_frame_sv2p.NextFrameSv2p,
         1)
 
-  def testStochasticTwoFrames(self):
+  def testSv2pTwoFrames(self):
     self.TestOnVariousInputOutputSizes(
-        next_frame_params.next_frame_stochastic(),
-        next_frame_sv2p.NextFrameStochasticTwoFrames,
+        next_frame_sv2p_params.next_frame_sv2p(),
+        next_frame_sv2p.NextFrameSv2pTwoFrames,
         1)
 
-  def testStochasticEmily(self):
+  def testEmily(self):
     self.TestOnVariousInputOutputSizes(
-        next_frame_params.next_frame_stochastic_emily(),
-        next_frame_emily.NextFrameStochasticEmily,
+        next_frame_emily.next_frame_emily(),
+        next_frame_emily.NextFrameEmily,
         1)
 
-  def testStochasticSavpVAE(self):
-    savp_hparams = next_frame_params.next_frame_savp()
+  def testSavpVAE(self):
+    savp_hparams = next_frame_savp_params.next_frame_savp()
     savp_hparams.use_vae = True
     savp_hparams.use_gan = False
     self.TestOnVariousInputOutputSizes(
@@ -230,8 +239,8 @@ def testStochasticSavpVAE(self):
     self.TestOnVariousUpSampleLayers(
         savp_hparams, next_frame_savp.NextFrameSAVP, 1)
 
-  def testStochasticSavpGAN(self):
-    hparams = next_frame_params.next_frame_savp()
+  def testSavpGAN(self):
+    hparams = next_frame_savp_params.next_frame_savp()
     hparams.use_gan = True
     hparams.use_vae = False
     self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
@@ -239,14 +248,14 @@ def testStochasticSavpGAN(self):
     hparams.gan_optimization = "sequential"
     self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
 
-  def testStochasticSavpGANVAE(self):
-    hparams = next_frame_params.next_frame_savp()
+  def testSavpGANVAE(self):
+    hparams = next_frame_savp_params.next_frame_savp()
     hparams.use_vae = True
     hparams.use_gan = True
     self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
 
-  def testStochasticInvalidVAEGANCombinations(self):
-    hparams = next_frame_params.next_frame_savp()
+  def testInvalidVAEGANCombinations(self):
+    hparams = next_frame_savp_params.next_frame_savp()
     hparams.use_gan = False
     hparams.use_vae = False
     self.assertRaises(ValueError, self.TestVideoModel,
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 93e807b61..e9a3d4771 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -517,7 +517,7 @@ def rl_modelrl_base():
       # So to use N frames set steps = N / (epochs * (1 - 1/11)).
       # We set it to use 100k frames for training.
       true_env_generator_num_steps=int(100000 / (6 * (1.0 - 1.0/11.0))),
-      generative_model="next_frame_basic",
+      generative_model="next_frame_basic_deterministic",
       generative_model_params="next_frame_pixel_noise",
       ppo_params="ppo_pong_base",
       autoencoder_train_steps=0,
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 28c088ca1..de192b6a2 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -152,6 +152,10 @@ def eval_hooks():
   def hparams(self):
     return self._hparams
 
+  @property
+  def is_training(self):
+    return self._hparams.mode == tf.estimator.ModeKeys.TRAIN
+
   @property
   def has_input(self):
     if self._problem_hparams:

From bfb629a9c00455a23c0b1d038d3349e0ba1a7db5 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 22 Aug 2018 12:24:11 -0700
Subject: [PATCH 0639/2720] Add mtf.sigmoid, mtf.tanh, and gradients to
 mtf.maximum, mtf.minimum.

PiperOrigin-RevId: 209809199
---
 .../mesh_tensorflow/mesh_tensorflow.py        | 77 ++++++++++++++++---
 1 file changed, 67 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index e21a8fa83..f0623064d 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -1298,6 +1298,20 @@ def exp(x, name="exp"):
                grad_function=lambda op, dy: [dy * op.outputs[0]])
 
 
+def sigmoid(x, name="sigmoid"):
+  def grad_function(op, dy):
+    y = op.outputs[0]
+    return [y * (1.0 - y) * dy]
+  return cwise(tf.sigmoid, [x], name=name, grad_function=grad_function)
+
+
+def tanh(x, name="tanh"):
+  def grad_function(op, dy):
+    y = op.outputs[0]
+    return [(1.0 - square(y)) * dy]
+  return cwise(tf.tanh, [x], name=name, grad_function=grad_function)
+
+
 def pow(x, y):  # pylint: disable=redefined-builtin
   return exp(log(x) * y)
 
@@ -1482,16 +1496,6 @@ def binary_op_with_broadcasting(
       output_dtype).outputs[0]
 
 
-def maximum(x1, x2, output_shape=None):
-  return binary_op_with_broadcasting(
-      tf.maximum, x1, x2, output_shape=output_shape)
-
-
-def minimum(x1, x2, output_shape=None):
-  return binary_op_with_broadcasting(
-      tf.minimum, x1, x2, output_shape=output_shape)
-
-
 def less(x1, x2, output_shape=None):
   return binary_op_with_broadcasting(
       tf.less, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
@@ -1559,6 +1563,59 @@ def gradient(self, grad_ys):
             reduce_sum(dy, output_shape=self.inputs[1].shape)]
 
 
+class MinMaxOperation(BinaryOpWithBroadcasting):
+  """Binary minimum/maximum with broadcasting."""
+
+  def __init__(self, tf_fn, x1, x2, output_shape, name=None):
+    super(MinMaxOperation, self).__init__(
+        tf_fn, x1, x2, output_shape, x1.dtype, name=name or "add")
+    if x1.dtype != x2.dtype:
+      raise ValueError("Dtypes must be equal.")
+
+  def gradient(self, grad_ys):
+    dy = grad_ys[0]
+    return [dy * cast(equal(self.inputs[0], self.outputs[0]), dy.dtype),
+            dy * cast(equal(self.inputs[1], self.outputs[0]), dy.dtype)]
+
+
+def minimum(x1, x2, output_shape=None, name=None):
+  """Binary minimum with broadcsting.
+
+  Args:
+    x1: a Tensor
+    x2: a Tensor
+    output_shape: an optional Shape
+    name: an optional string
+  Returns:
+    a Tensor
+  """
+  output_shape = convert_to_shape(output_shape)
+  with tf.name_scope(name, default_name="minimum"):
+    x1, x2 = binary_arguments_to_tensors(x1, x2)
+    return MinMaxOperation(
+        tf.minimum, x1, x2, output_shape=_infer_binary_broadcast_shape(
+            x1.shape, x2.shape, output_shape)).outputs[0]
+
+
+def maximum(x1, x2, output_shape=None, name=None):
+  """Binary maximum with broadcsting.
+
+  Args:
+    x1: a Tensor
+    x2: a Tensor
+    output_shape: an optional Shape
+    name: an optional string
+  Returns:
+    a Tensor
+  """
+  output_shape = convert_to_shape(output_shape)
+  with tf.name_scope(name, default_name="maximum"):
+    x1, x2 = binary_arguments_to_tensors(x1, x2)
+    return MinMaxOperation(
+        tf.maximum, x1, x2, output_shape=_infer_binary_broadcast_shape(
+            x1.shape, x2.shape, output_shape)).outputs[0]
+
+
 class BroadcastOperation(Operation):
   """Broadcast - output dims are a superset of input dims, in any order."""
 

From b54d5cd02296d28c5f77555d8c63ea70ac3af05d Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 22 Aug 2018 13:06:58 -0700
Subject: [PATCH 0640/2720] Internal change

PiperOrigin-RevId: 209816150
---
 tensor2tensor/models/research/autoencoders.py | 144 +++++++++++-------
 1 file changed, 93 insertions(+), 51 deletions(-)

diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index cde25ba12..7ccc318bd 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -182,18 +182,42 @@ def decoder(self, x, encoder_layers):
         x = common_layers.layer_norm(x, name="ln_%d" % i)
       return x
 
+  def gumbel_sample(self, reconstr_gan):
+    hparams = self.hparams
+    is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+    vocab_size = self._problem_hparams.target_modality.top_dimensionality
+    reconstr_gan = tf.nn.log_softmax(reconstr_gan)
+    if is_training and hparams.gumbel_temperature > 0.0:
+      gumbel_samples = discretization.gumbel_sample(
+          common_layers.shape_list(reconstr_gan))
+      gumbel_samples *= hparams.gumbel_noise_factor
+      reconstr_gan += gumbel_samples
+      reconstr_sample = latent_layers.multinomial_sample(
+          reconstr_gan, temperature=hparams.gumbel_temperature)
+      reconstr_gan = tf.nn.softmax(
+          reconstr_gan / hparams.gumbel_temperature)
+    else:
+      reconstr_sample = tf.argmax(reconstr_gan, axis=-1)
+      reconstr_gan = tf.nn.softmax(reconstr_gan / 0.1)  # Sharpen a bit.
+    # Use 1-hot forward, softmax backward.
+    reconstr_hot = tf.one_hot(reconstr_sample, vocab_size)
+    reconstr_gan += reconstr_hot - tf.stop_gradient(reconstr_gan)
+    return reconstr_gan
+
   def body(self, features):
     hparams = self.hparams
     is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+    vocab_size = self._problem_hparams.target_modality.top_dimensionality
+    encoder_layers = None
+    self.is1d = hparams.sample_width == 1
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
       labels = features["targets_raw"]
-      vocab_size = self._problem_hparams.target_modality.top_dimensionality
       shape = common_layers.shape_list(labels)
       x = tf.one_hot(labels, vocab_size)
       x = self.embed(x)
       target_codes = x
-      is1d = shape[2] == 1
-      self.is1d = is1d
+      if shape[2] == 1:
+        self.is1d = True
       # Run encoder.
       x, encoder_layers = self.encoder(x)
       # Bottleneck.
@@ -240,10 +264,11 @@ def body(self, features):
       x = self.unbottleneck(b, res_size)
     # Run decoder.
     x = self.decoder(x, encoder_layers)
-    if hparams.mode == tf.estimator.ModeKeys.PREDICT:
-      return x, {"bottleneck_loss": 0.0}
+
     # Cut to the right size and mix before returning.
-    res = x[:, :shape[1], :shape[2], :]
+    res = x
+    if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+      res = x[:, :shape[1], :shape[2], :]
 
     # Final dense layer.
     res = tf.layers.dense(
@@ -254,6 +279,14 @@ def body(self, features):
     ]
     res = tf.reshape(res, output_shape)
 
+    if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+      if hparams.use_vq_loss:
+        (reconstr, _, _, _, _) = discretization.vq_loss(
+            res, labels, vocab_size)
+      else:
+        reconstr = tf.layers.dense(res, vocab_size, name="autoencoder_final")
+      return reconstr, {"bottleneck_loss": 0.0}
+
     if hparams.gan_loss_factor != 0.0:
       res_gan, res = tf.split(res, 2, axis=0)
 
@@ -294,22 +327,7 @@ def body(self, features):
       else:
         reconstr_gan = tf.layers.dense(
             res_gan, vocab_size, name="autoencoder_final", reuse=True)
-        reconstr_gan = tf.nn.log_softmax(reconstr_gan)
-        if is_training and hparams.gumbel_temperature > 0.0:
-          gumbel_samples = discretization.gumbel_sample(
-              common_layers.shape_list(reconstr_gan))
-          gumbel_samples *= hparams.gumbel_noise_factor
-          reconstr_gan += gumbel_samples
-          reconstr_sample = latent_layers.multinomial_sample(
-              reconstr_gan, temperature=hparams.gumbel_temperature)
-          reconstr_gan = tf.nn.softmax(
-              reconstr_gan / hparams.gumbel_temperature)
-        else:
-          reconstr_sample = tf.argmax(reconstr_gan, axis=-1)
-          reconstr_gan = tf.nn.softmax(reconstr_gan / 0.1)  # Sharpen a bit.
-        # Use 1-hot forward, softmax backward.
-        reconstr_hot = tf.one_hot(reconstr_sample, vocab_size)
-        reconstr_gan += reconstr_hot - tf.stop_gradient(reconstr_gan)
+        reconstr_gan = self.gumbel_sample(reconstr_gan)
         # Embed to codes.
         gan_codes = self.embed(reconstr_gan)
 
@@ -420,7 +438,6 @@ class AutoencoderAutoregressive(AutoencoderBasic):
 
   def body(self, features):
     hparams = self.hparams
-    is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
     # Run the basic autoencoder part first.
     basic_result, losses = super(AutoencoderAutoregressive, self).body(features)
     if hparams.autoregressive_mode == "none":
@@ -431,11 +448,17 @@ def body(self, features):
       losses["plain"] = plain_training_loss
     res_shape = common_layers.shape_list(basic_result)
     vocab_size = self._problem_hparams.target_modality.top_dimensionality
-    basic_result = self.embed(basic_result)
-    shape = common_layers.shape_list(basic_result)
-    basic1d = tf.reshape(basic_result, [shape[0], -1, shape[-1]])
     targets = tf.one_hot(features["targets_raw"], vocab_size)
+    # Prepare inputs for autoregressive modes.
+    if common_layers.shape_list(features["targets"])[1] == 1:
+      # This happens on the first step of predicitions.
+      assert hparams.mode == tf.estimator.ModeKeys.PREDICT
+      targets = tf.zeros_like(basic_result)
     targets = self.embed(targets)
+    basic_hot = self.gumbel_sample(basic_result)
+    basic_result = self.embed(basic_hot)
+    shape = common_layers.shape_list(basic_result)
+    basic1d = tf.reshape(basic_result, [shape[0], -1, shape[-1]])
     targets = tf.reshape(targets, common_layers.shape_list(basic_result))
     # During autoregressive inference, don't resample.
     if hparams.mode == tf.estimator.ModeKeys.PREDICT:
@@ -443,19 +466,8 @@ def body(self, features):
         basic1d = hparams.sampled_basic1d_tensor
       else:
         hparams.sampled_basic1d_tensor = basic1d
-    # Prepare inputs for autoregressive modes.
-    if common_layers.shape_list(features["targets"])[1] == 1:
-      # This happens on the first step of predicitions.
-      assert hparams.mode == tf.estimator.ModeKeys.PREDICT
-      targets = tf.zeros_like(basic_result)
-    targets_dropout = common_layers.mix(
-        targets,
-        tf.zeros_like(basic_result),
-        hparams.bottleneck_warmup_steps,
-        is_training,
-        max_prob=1.0 - hparams.autoregressive_dropout,
-        broadcast_last=True)
     # Sometimes it's useful to look at non-autoregressive evals.
+    targets_dropout = targets
     if (hparams.mode == tf.estimator.ModeKeys.EVAL and
         hparams.autoregressive_eval_pure_autoencoder):
       targets_dropout = tf.zeros_like(basic_result)
@@ -476,7 +488,7 @@ def body(self, features):
           padding="LEFT",
           activation=common_layers.belu,
           name="autoregressive_conv3")
-      res = tf.layers.dense(res, vocab_size, "autoregressive_final")
+      res = tf.layers.dense(res, vocab_size, name="autoregressive_final")
       return tf.reshape(res, res_shape), losses
     if hparams.autoregressive_mode == "conv5":
       res = common_layers.conv1d(
@@ -486,7 +498,7 @@ def body(self, features):
           padding="LEFT",
           activation=common_layers.belu,
           name="autoregressive_conv5")
-      res = tf.layers.dense(res, vocab_size, "autoregressive_final")
+      res = tf.layers.dense(res, vocab_size, name="autoregressive_final")
       return tf.reshape(res, res_shape), losses
     if hparams.autoregressive_mode == "sru":
       res = common_layers.conv1d(
@@ -497,7 +509,7 @@ def body(self, features):
           activation=common_layers.belu,
           name="autoregressive_sru_conv3")
       res = common_layers.sru(res)
-      res = tf.layers.dense(res, vocab_size, "autoregressive_final")
+      res = tf.layers.dense(res, vocab_size, name="autoregressive_final")
       return tf.reshape(res, res_shape), losses
 
     raise ValueError(
@@ -527,13 +539,10 @@ def infer(self, features, *args, **kwargs):
 
     # Sample again if requested for the autoregressive part.
     extra_samples = self.hparams.autoregressive_decode_steps
-    self.hparams.autoregressive_dropout = 0.2
     for i in range(extra_samples):
       if i == extra_samples - 2:
-        self.hparams.autoregressive_dropout -= 0.1
         self.hparams.sampling_temp /= 2
       if i == extra_samples - 1:
-        self.hparams.autoregressive_dropout -= 0.1
         self.hparams.sampling_temp = 0.0
       features["targets"] = samples
       old_samples1d = tf.reshape(samples, [shape[0], -1, shape[3]])
@@ -670,6 +679,38 @@ def decoder(self, x, encoder_layers=None):
       return x
 
 
+@registry.register_model
+class AutoencoderResidualVAE(AutoencoderResidual):
+  """Residual VAE autoencoder."""
+
+  def bottleneck(self, x):
+    hparams = self.hparams
+    z_size = hparams.bottleneck_bits
+    x_shape = common_layers.shape_list(x)
+    with tf.variable_scope("vae"):
+      mu = tf.layers.dense(x, z_size, name="mu")
+      if hparams.mode != tf.estimator.ModeKeys.TRAIN:
+        return mu, 0.0  # No sampling or kl loss on eval.
+      log_sigma = tf.layers.dense(x, z_size, name="log_sigma")
+      epsilon = tf.random_normal(x_shape[:-1] + [z_size])
+      z = mu + tf.exp(log_sigma / 2) * epsilon
+      kl = 0.5 * tf.reduce_mean(
+          tf.exp(log_sigma) + tf.square(mu) - 1. - log_sigma, axis=-1)
+      free_bits = z_size // 4
+      kl_loss = tf.reduce_mean(tf.maximum(kl - free_bits, 0.0))
+    return z, kl_loss * hparams.kl_beta
+
+  def sample(self, features=None, shape=None):
+    del features
+    hparams = self.hparams
+    div_x = 2**hparams.num_hidden_layers
+    div_y = 1 if self.is1d else 2**hparams.num_hidden_layers
+    size = [hparams.batch_size, hparams.sample_height // div_x,
+            hparams.sample_width // div_y, hparams.bottleneck_bits]
+    size = size if shape is None else shape
+    return tf.random_normal(size)
+
+
 @registry.register_model
 class AutoencoderBasicDiscrete(AutoencoderAutoregressive):
   """Discrete autoencoder."""
@@ -891,10 +932,10 @@ def autoencoder_basic():
   hparams.add_hparam("gan_codes_warmup_steps", 6000)
   hparams.add_hparam("gan_loss_factor", 0.0)
   hparams.add_hparam("bottleneck_l2_factor", 0.05)
-  hparams.add_hparam("gumbel_temperature", 0.05)
-  hparams.add_hparam("gumbel_noise_factor", 0.2)
+  hparams.add_hparam("gumbel_temperature", 0.2)
+  hparams.add_hparam("gumbel_noise_factor", 0.4)
   hparams.add_hparam("vq_temperature", 0.001)
-  hparams.add_hparam("use_vq_loss", 0)
+  hparams.add_hparam("use_vq_loss", int(False))
   return hparams
 
 
@@ -904,7 +945,6 @@ def autoencoder_autoregressive():
   hparams = autoencoder_basic()
   hparams.add_hparam("autoregressive_forget_base", False)
   hparams.add_hparam("autoregressive_mode", "none")
-  hparams.add_hparam("autoregressive_dropout", 0.4)
   hparams.add_hparam("autoregressive_decode_steps", 0)
   hparams.add_hparam("autoregressive_eval_pure_autoencoder", False)
   return hparams
@@ -928,6 +968,7 @@ def autoencoder_residual():
   hparams.add_hparam("residual_filter_multiplier", 2.0)
   hparams.add_hparam("residual_dropout", 0.2)
   hparams.add_hparam("residual_use_separable_conv", int(True))
+  hparams.add_hparam("kl_beta", 1.0)
   return hparams
 
 
@@ -943,6 +984,7 @@ def autoencoder_residual_text():
   hparams.target_modality = "symbol:identity"
   hparams.input_modalities = "symbol:identity"
   hparams.autoregressive_mode = "none"
+  hparams.sample_width = 1
   return hparams
 
 
@@ -992,7 +1034,6 @@ def autoencoder_ordered_discrete():
   hparams = autoencoder_residual_discrete()
   hparams.bottleneck_noise = 0.8
   hparams.gan_loss_factor = 0.02
-  hparams.use_vq_loss = 0
   hparams.add_hparam("unordered", False)
   return hparams
 
@@ -1016,6 +1057,7 @@ def autoencoder_ordered_text():
   hparams.max_hidden_size = 4096
   hparams.target_modality = "symbol:identity"
   hparams.input_modalities = "symbol:identity"
+  hparams.sample_width = 1
   return hparams
 
 
@@ -1026,8 +1068,8 @@ def autoencoder_ordered_text_small():
   hparams.bottleneck_bits = 64
   hparams.hidden_size = 64
   hparams.max_hidden_size = 512
-  hparams.bottleneck_noise = 0.0
-  hparams.autoregressive_mode = "none"
+  hparams.bottleneck_noise = 0.3
+  hparams.autoregressive_mode = "conv5"
   return hparams
 
 

From 0b6b83c7784daa84ac580f67896cf9b1d14b40d8 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Wed, 22 Aug 2018 13:40:16 -0700
Subject: [PATCH 0641/2720] Fix some bugs in mtf image transformer and add
 hparams

PiperOrigin-RevId: 209821452
---
 .../mesh_tensorflow/mtf_image_transformer.py  | 78 ++++++++++++++++---
 1 file changed, 68 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
index ab37e7a4d..fe01c528e 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
@@ -35,7 +35,7 @@
 
 @registry.register_model
 class MtfImageTransformer(mtf_model.MtfModel):
-  """Transformer in mesh_tensorflow."""
+  """Image Transformer in mesh_tensorflow."""
 
   def set_activation_type(self):
     hparams = self._hparams
@@ -66,10 +66,11 @@ def mtf_model_fn(self, features, mesh):
     shifted_targets = common_layers.shift_right_2d(targets)
 
     # Declare all the dimensions
-    model_dim = mtf.Dimension("model", hparams.hidden_size)
+    model_dim = mtf.Dimension("d_model", hparams.hidden_size)
     batch_dim = mtf.Dimension("batch", hparams.batch_size)
     length_dim = mtf.Dimension("length", length)
-    filter_dim = mtf.Dimension("filter_size", hparams.filter_size)
+    max_length_dim = mtf.Dimension("max_length", hparams.max_length)
+    filter_dim = mtf.Dimension("d_ff", hparams.d_ff)
     kv_channels = mtf.Dimension("kv_channels", hparams.d_kv)
     heads = mtf.Dimension("heads", hparams.num_heads)
 
@@ -116,12 +117,12 @@ def layer_prepostprocess_dropout(x):
 
     positional_embedding_var = mtf.get_variable(
         mesh, "positional_embedding",
-        mtf.Shape([targets_vocab_dim, model_dim]),
+        mtf.Shape([max_length_dim, model_dim]),
         initializer=tf.random_normal_initializer(),
         activation_dtype=activation_dtype)
     x = (mtf.gather(targets_embedding_var, shifted_targets, targets_vocab_dim) +
          mtf.gather(
-             positional_embedding_var, targets_position, targets_vocab_dim))
+             positional_embedding_var, targets_position, max_length_dim))
 
     # Image Transformer Decoder
     # [ self attention - ffn - residual + dropout] x n
@@ -164,14 +165,15 @@ def mtf_image_transformer_base():
   hparams.no_data_parallelism = True
   hparams.use_fixed_batch_size = True
   hparams.batch_size = 1
-  hparams.max_length = 256
+  hparams.max_length = 3072
   hparams.hidden_size = 256
   hparams.label_smoothing = 0.0
   # 8-way model-parallelism
-  hparams.add_hparam("mesh_shape", "8")
-  hparams.add_hparam("layout", "vocab:0;filter_size:0;heads:0")
+  hparams.add_hparam("mesh_shape", "batch:8")
+  hparams.add_hparam("layout", "batch:batch")
+  hparams.add_hparam("mtf_mode", True)
   hparams.add_hparam("num_heads", 8)
-  hparams.add_hparam("filter_size", 512)
+  hparams.add_hparam("filter_size", 1024)
   hparams.add_hparam("num_encoder_layers", 0)
   hparams.add_hparam("num_decoder_layers", 6)
   hparams.add_hparam("attention_key_size", 256)
@@ -191,6 +193,7 @@ def mtf_image_transformer_base():
   hparams.learning_rate_schedule = "rsqrt_decay"
   hparams.learning_rate_warmup_steps = 10000
   hparams.add_hparam("d_kv", 32)
+  hparams.add_hparam("d_ff", 2048)
 
   # Image related hparams
   hparams.add_hparam("img_len", 32)
@@ -205,7 +208,7 @@ def mtf_image_transformer_tiny():
   """Catch bugs locally..."""
   hparams = mtf_image_transformer_base()
   hparams.hidden_size = 128
-  hparams.filter_size = 256
+  hparams.d_ff = 256
   hparams.batch_size = 4
   hparams.num_encoder_layers = 1
   hparams.num_decoder_layers = 1
@@ -248,6 +251,61 @@ def mtf_image_transformer_base_single():
   return hparams
 
 
+@registry.register_hparams
+def mtf_image_transformer_base_cifar():
+  """Data parallel CIFAR parameters."""
+  hparams = mtf_image_transformer_base()
+  hparams.mesh_shape = "batch:32"
+  hparams.layout = "batch:batch"
+  hparams.batch_size = 128
+  hparams.num_heads = 4
+  hparams.num_decoder_layers = 12
+  hparams.block_length = 256
+  hparams.hidden_size = 512
+  hparams.d_ff = 2048
+  hparams.learning_rate = 0.5
+  hparams.learning_rate_warmup_steps = 6000
+  hparams.layer_preprocess_sequence = "none"
+  hparams.layer_postprocess_sequence = "dan"
+  hparams.layer_prepostprocess_dropout = 0.3
+  hparams.unconditional = True
+  return hparams
+
+
+@registry.register_hparams
+def mtf_image_transformer_base_imagenet():
+  """Data parallel CIFAR parameters."""
+  hparams = mtf_image_transformer_base_cifar()
+  hparams.mesh_shape = "batch:32"
+  hparams.layout = "batch:batch"
+  hparams.batch_size = 64
+  hparams.d_ff = 2048
+  hparams.hidden_size = 512
+  hparams.num_decoder_layers = 12
+  hparams.learning_rate = 0.5
+  hparams.learning_rate_warmup_steps = 6000
+  hparams.layer_preprocess_sequence = "none"
+  hparams.layer_postprocess_sequence = "dan"
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.unconditional = True
+  return hparams
+
+
+@registry.register_hparams
+def mtf_image_transformer_base_imagenet_mp():
+  """Model parallel ImageNet parameters."""
+  hparams = mtf_image_transformer_base_imagenet()
+  hparams.mesh_shape = "model:4;batch:8"
+  hparams.layout = "batch:batch;d_ff:model;heads:model"
+  hparams.batch_size = 32
+  hparams.num_heads = 4
+  hparams.d_ff = 8192
+  hparams.learning_rate_warmup_steps = 6000
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.unconditional = True
+  return hparams
+
+
 @registry.register_hparams
 def mtf_image_transformer_tiny_moe():
   hparams = mtf_image_transformer_tiny()

From 3e383c30cb17131c073988a956bc201d1c5a93b0 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Wed, 22 Aug 2018 14:28:14 -0700
Subject: [PATCH 0642/2720] Add the Stanford Natural Language Inference data.

PiperOrigin-RevId: 209830291
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/stanford_nli.py | 134 ++++++++++++++++++
 2 files changed, 135 insertions(+)
 create mode 100644 tensor2tensor/data_generators/stanford_nli.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 129489aa2..dd3ee3dd8 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -63,6 +63,7 @@
     "tensor2tensor.data_generators.quora_qpairs",
     "tensor2tensor.data_generators.rte",
     "tensor2tensor.data_generators.snli",
+    "tensor2tensor.data_generators.stanford_nli",
     "tensor2tensor.data_generators.style_transfer",
     "tensor2tensor.data_generators.squad",
     "tensor2tensor.data_generators.sst_binary",
diff --git a/tensor2tensor/data_generators/stanford_nli.py b/tensor2tensor/data_generators/stanford_nli.py
new file mode 100644
index 000000000..16efea71e
--- /dev/null
+++ b/tensor2tensor/data_generators/stanford_nli.py
@@ -0,0 +1,134 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generators for StanfordNLI."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import zipfile
+import six
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+EOS = text_encoder.EOS
+
+
+@registry.register_problem
+class StanfordNLI(text_problems.TextConcat2ClassProblem):
+  """StanfordNLI classification problems."""
+
+  # Link to data from GLUE: https://gluebenchmark.com/tasks
+  _SNLI_URL = ("https://nlp.stanford.edu/projects/snli/snli_1.0.zip")
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 100,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  @property
+  def approx_vocab_size(self):
+    return 2**15
+
+  @property
+  def vocab_filename(self):
+    return "vocab.snli.%d" % self.approx_vocab_size
+
+  @property
+  def num_classes(self):
+    return 3
+
+  def class_labels(self, data_dir):
+    del data_dir
+    # Note this binary classification is different from usual SNLI.
+    return ["contradiction", "entailment", "neutral"]
+
+  def _maybe_download_corpora(self, tmp_dir):
+    snli_filename = "SNLI.zip"
+    snli_finalpath = os.path.join(tmp_dir, "snli_1.0")
+    if not tf.gfile.Exists(snli_finalpath):
+      zip_filepath = generator_utils.maybe_download(
+          tmp_dir, snli_filename, self._SNLI_URL)
+      zip_ref = zipfile.ZipFile(zip_filepath, "r")
+      zip_ref.extractall(tmp_dir)
+      zip_ref.close()
+
+    return snli_finalpath
+
+  def example_generator(self, filename):
+    label_list = self.class_labels(data_dir=None)
+    for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
+      if idx == 0: continue  # skip header
+      if six.PY2:
+        line = unicode(line.strip(), "utf-8")
+      else:
+        line = line.strip().decode("utf-8")
+      split_line = line.split("\t")
+      # Works for both splits even though dev has some extra human labels.
+      s1, s2 = split_line[5:7]
+      if split_line[0] == "-":
+        continue
+      l = label_list.index(split_line[0])
+      inputs = [s1, s2]
+      yield {
+          "inputs": inputs,
+          "label": l
+      }
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    snli_dir = self._maybe_download_corpora(tmp_dir)
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      filesplit = "snli_1.0_train.txt"
+    else:
+      filesplit = "snli_1.0_dev.txt"
+
+    filename = os.path.join(snli_dir, filesplit)
+    for example in self.example_generator(filename):
+      yield example
+
+
+@registry.register_problem
+class StanfordNLICharacters(StanfordNLI):
+  """StanfordNLI classification problems, character level"""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  def global_task_id(self):
+    return problem.TaskID.THREE_CL_NLI
+
+
+@registry.register_problem
+class StanfordNLISharedVocab(StanfordNLI):
+  """StanfordNLI classification problems with the LM1b vocabulary"""
+
+  @property
+  def vocab_filename(self):
+    return "vocab.lm1b.en.%d" % 2**15

From c3ab3b9ebc3e0c6f27603eb79790e8d392cb3209 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 22 Aug 2018 14:39:30 -0700
Subject: [PATCH 0643/2720] Add Eager tests to existing T2T layers; move to
 self.session().

PiperOrigin-RevId: 209832467
---
 tensor2tensor/layers/common_attention.py      |   2 +-
 tensor2tensor/layers/common_attention_test.py | 217 ++++----
 .../layers/common_image_attention_test.py     |   1 +
 tensor2tensor/layers/common_layers_test.py    | 468 ++++++++----------
 tensor2tensor/layers/common_video_test.py     |  73 ++-
 tensor2tensor/layers/discretization_test.py   |  91 ++--
 tensor2tensor/layers/modalities_test.py       |  34 +-
 7 files changed, 404 insertions(+), 482 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index ebf70535d..ad9c612a4 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -701,7 +701,7 @@ def add_positional_embedding_nd(x, max_length, name):
           name + "_%d" % i,
           shape,
           initializer=tf.random_normal_initializer(0, depth**-0.5))
-      var *= depth**0.5
+      var = var * depth**0.5
       x += tf.slice(var, start, size)
     return x
 
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index a14f43252..6a833713b 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -27,17 +27,17 @@
 import tensorflow as tf
 
 
+@tf.contrib.eager.run_all_tests_in_graph_and_eager_modes
 class CommonAttentionTest(parameterized.TestCase, tf.test.TestCase):
 
   def testAddPositionalEmbedding(self):
     x = np.random.rand(5, 3, 12)
-    with self.test_session() as session:
-      y = common_attention.add_positional_embedding(
-          tf.constant(x, dtype=tf.float32),
-          max_length=4,
-          name="pos_embedding")
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_attention.add_positional_embedding(
+        tf.constant(x, dtype=tf.float32),
+        max_length=4,
+        name="pos_embedding")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 3, 12))
 
   @parameterized.parameters(
@@ -47,25 +47,22 @@ def testAddPositionalEmbedding(self):
   )
   def testAddPositionalEmbeddingNd(self, input_shape):
     x = np.random.rand(*input_shape)
-    with self.test_session() as session:
-      y = common_attention.add_positional_embedding_nd(
-          tf.constant(x, dtype=tf.float32),
-          max_length=5,
-          name="pos_embedding")
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_attention.add_positional_embedding_nd(
+        tf.constant(x, dtype=tf.float32),
+        max_length=5,
+        name="pos_embedding")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, input_shape)
 
   def testDotProductAttention(self):
     x = np.random.rand(5, 7, 12, 32)
     y = np.random.rand(5, 7, 12, 32)
-    with self.test_session() as session:
-      a = common_attention.dot_product_attention(
-          tf.constant(x, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32), None)
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_attention.dot_product_attention(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32), None)
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
   @parameterized.named_parameters(
@@ -84,11 +81,10 @@ def testMaskedWithinBlockLocalAttention1D(self, batch, heads, length,
     v = tf.random_normal([batch, heads, length, depth_v])
     output = common_attention.masked_within_block_local_attention_1d(
         q, k, v, block_length=block_length)
-    with self.test_session() as session:
-      if isinstance(batch, tf.Tensor):
-        batch, res = session.run([batch, output])
-      else:
-        res = session.run(output)
+    if isinstance(batch, tf.Tensor):
+      batch, res = self.evaluate([batch, output])
+    else:
+      res = self.evaluate(output)
 
     self.assertEqual(res.shape, (batch, heads, length, depth_v))
 
@@ -108,11 +104,10 @@ def testMaskedLocalAttention1D(self, batch, heads, length, depth_k, depth_v,
     v = tf.random_normal([batch, heads, length, depth_v])
     output = common_attention.masked_local_attention_1d(
         q, k, v, block_length=block_length)
-    with self.test_session() as session:
-      if isinstance(batch, tf.Tensor):
-        batch, res = session.run([batch, output])
-      else:
-        res = session.run(output)
+    if isinstance(batch, tf.Tensor):
+      batch, res = self.evaluate([batch, output])
+    else:
+      res = self.evaluate(output)
 
     self.assertEqual(res.shape, (batch, heads, length, depth_v))
 
@@ -131,43 +126,41 @@ def testLocalUnmaskedAttention1D(self, batch, heads, length,
     v = tf.random_normal([batch, heads, length, depth_v])
     output = common_attention.local_attention_1d(
         q, k, v, block_length=block_length, filter_width=3)
-    with self.test_session() as session:
-      if isinstance(batch, tf.Tensor):
-        batch, res = session.run([batch, output])
-      else:
-        res = session.run(output)
+    if isinstance(batch, tf.Tensor):
+      batch, res = self.evaluate([batch, output])
+    else:
+      res = self.evaluate(output)
 
     self.assertEqual(res.shape, (batch, heads, length, depth_v))
 
   def testLocalUnmaskedAttention2D(self):
     x = np.random.rand(5, 4, 25, 25, 16)
     y = np.random.rand(5, 4, 25, 25, 16)
-    with self.test_session() as session:
-      a = common_attention.local_attention_2d(
-          tf.constant(x, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          query_shape=(4, 4),
-          memory_flange=(3, 3))
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_attention.local_attention_2d(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        query_shape=(4, 4),
+        memory_flange=(3, 3))
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 4, 25, 25, 16))
 
   def testLocalUnmaskedAttention2DMatchingBlockLength(self):
     x = np.random.rand(5, 4, 25, 25, 16)
     y = np.random.rand(5, 4, 25, 25, 16)
-    with self.test_session() as session:
-      a = common_attention.local_attention_2d(
-          tf.constant(x, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          query_shape=(5, 5),
-          memory_flange=(3, 3))
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_attention.local_attention_2d(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        query_shape=(5, 5),
+        memory_flange=(3, 3))
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 4, 25, 25, 16))
 
   def testMultiheadSelfAttentionMemoryEfficient(self):
+    if tf.executing_eagerly():
+      return  # don't run test in Eager mode
+
     num_heads = 4
     io_size = 16
     batch = 2
@@ -175,7 +168,7 @@ def testMultiheadSelfAttentionMemoryEfficient(self):
     head_size = 5
     x = np.random.rand(batch, length, io_size)
     dy = np.random.rand(batch, length, io_size)
-    with self.test_session() as session:
+    with self.session() as session:
       x = tf.to_float(x)
       dy = tf.to_float(dy)
       bias = common_attention.attention_bias_lower_triangle(length)
@@ -220,25 +213,22 @@ def test2dGatherAndScatterInvertibility(self):
     depth = 8
     query_shape = (2, 3)
     x = np.random.rand(batch_size, num_heads, height, width, depth)
-    with self.test_session() as session:
-      x_indices = common_attention.gather_indices_2d(
-          x, query_shape, query_shape)
-      gathered_x = common_attention.gather_blocks_2d(x, x_indices)
-      x_shape = tf.constant([batch_size, num_heads, height, width, depth])
-      scattered_x = common_attention.scatter_blocks_2d(
-          gathered_x, x_indices, x_shape)
-      session.run(tf.global_variables_initializer())
-      res = session.run(scattered_x)
+    x_indices = common_attention.gather_indices_2d(
+        x, query_shape, query_shape)
+    gathered_x = common_attention.gather_blocks_2d(x, x_indices)
+    x_shape = tf.constant([batch_size, num_heads, height, width, depth])
+    scattered_x = common_attention.scatter_blocks_2d(
+        gathered_x, x_indices, x_shape)
+    res = self.evaluate(scattered_x)
     self.assertAllClose(x, res)
 
   def test2dBlockRasterScanMask(self):
     """Testing the 2d block raster scan mask."""
     query_shape = (2, 3)
     memory_flange = (2, 1)
-    with self.test_session() as session:
-      mask = common_attention.make_2d_block_raster_mask(
-          query_shape, memory_flange)
-      res = session.run(mask)
+    mask = common_attention.make_2d_block_raster_mask(
+        query_shape, memory_flange)
+    res = self.evaluate(mask)
     correct_mask = np.array(
         [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0,
           1.0, 0.0, 1.0, 1.0, 1.0, 1.0],
@@ -285,11 +275,10 @@ def test2dGather(self):
                             y[1, 1, correct_indices[2]],
                             y[1, 1, correct_indices[3]]]]]
 
-    with self.test_session() as session:
-      x_indices = common_attention.gather_indices_2d(
-          x, query_shape, query_shape)
-      gathered_x = common_attention.gather_blocks_2d(x, x_indices)
-      x_indices, gathered_x = session.run([x_indices, gathered_x])
+    x_indices = common_attention.gather_indices_2d(
+        x, query_shape, query_shape)
+    gathered_x = common_attention.gather_blocks_2d(x, x_indices)
+    x_indices, gathered_x = self.evaluate([x_indices, gathered_x])
     self.assertAllEqual(correct_indices, x_indices)
     self.assertAllClose(correct_gathered_x, gathered_x)
 
@@ -358,16 +347,14 @@ def testGetMemoryRegion(self):
                           y[1, 1, [12, 13, 14, 18, 19, 20]],
                           y[1, 1, [15, 16, 17, 21, 22, 23]]]]]
     correct_x_center = np.array(correct_x_center)
-    with self.test_session() as session:
-      x_indices = common_attention.gather_indices_2d(
-          x, query_shape, query_shape)
-      x_flange, x_center = common_attention.get_memory_region(
-          tf.constant(x, dtype=tf.float32),
-          query_shape,
-          memory_flange,
-          x_indices)
-      session.run(tf.global_variables_initializer())
-      [x_flange, x_center] = session.run([x_flange, x_center])
+    x_indices = common_attention.gather_indices_2d(
+        x, query_shape, query_shape)
+    x_flange, x_center = common_attention.get_memory_region(
+        tf.constant(x, dtype=tf.float32),
+        query_shape,
+        memory_flange,
+        x_indices)
+    [x_flange, x_center] = self.evaluate([x_flange, x_center])
     self.assertAllClose(correct_x_flange, x_flange)
     self.assertAllClose(correct_x_center, x_center)
 
@@ -427,42 +414,38 @@ def testGetShiftedCenterBlocks(self):
                                             y[1, 1, [15, 16, 17, 21, 22]]),
                                            axis=0)]]]
     correct_gathered_x = np.array(correct_gathered_x)
-    with self.test_session() as session:
-      x_indices = common_attention.gather_indices_2d(
-          x, query_shape, query_shape)
-      gathered_x = common_attention.get_shifted_center_blocks(
-          tf.constant(x, dtype=tf.float32),
-          x_indices)
-      session.run(tf.global_variables_initializer())
-      x_indices, gathered_x = session.run([x_indices, gathered_x])
+    x_indices = common_attention.gather_indices_2d(
+        x, query_shape, query_shape)
+    gathered_x = common_attention.get_shifted_center_blocks(
+        tf.constant(x, dtype=tf.float32),
+        x_indices)
+    x_indices, gathered_x = self.evaluate([x_indices, gathered_x])
     self.assertAllClose(correct_gathered_x, gathered_x)
 
   def testDotProductAttentionRelative(self):
     x = np.random.rand(5, 7, 12, 32)
     y = np.random.rand(5, 7, 12, 32)
-    with self.test_session() as session:
-      a = common_attention.dot_product_attention_relative(
-          tf.constant(x, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          None,
-          max_relative_position=3)
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_attention.dot_product_attention_relative(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        None,
+        max_relative_position=3)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
   def testDotProductUnMaskedAttentionRelativeV2(self):
     x = np.random.rand(5, 7, 12, 32)
     y = np.random.rand(5, 7, 12, 32)
-    with self.test_session() as session:
-      a = common_attention.dot_product_unmasked_self_attention_relative_v2(
-          tf.constant(x, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          None,
-          35)
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_attention.dot_product_unmasked_self_attention_relative_v2(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        None,
+        35)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
   def testBiasBatchCoordinates(self):
@@ -486,13 +469,7 @@ def testBiasBatchCoordinates(self):
     ], np.float32) * -1e9
 
     bias = common_attention.attention_bias_coordinates(q, k)
-
-    with self.test_session() as session:
-      session.run(tf.global_variables_initializer())
-      self.assertAllClose(
-          bias.eval(),
-          ground_truth,
-      )
+    self.assertAllClose(self.evaluate(bias), ground_truth)
 
   def testBiasFuture(self):
     """Testing the sequence order mask."""
@@ -515,13 +492,7 @@ def testBiasFuture(self):
     ], np.float32) * -1e9
 
     bias = common_attention.attention_bias_future(q, k)
-
-    with self.test_session() as session:
-      session.run(tf.global_variables_initializer())
-      self.assertAllClose(
-          bias.eval(),
-          ground_truth,
-      )
+    self.assertAllClose(self.evaluate(bias), ground_truth)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index 6d57e7413..c0845ef10 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -24,6 +24,7 @@
 import tensorflow as tf
 
 
+@tf.contrib.eager.run_all_tests_in_graph_and_eager_modes
 class CommonImageAttentionTest(parameterized.TestCase, tf.test.TestCase):
 
   @parameterized.parameters(
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index d4a6df588..da5647c09 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -26,6 +26,7 @@
 import tensorflow as tf
 
 
+@tf.contrib.eager.run_all_tests_in_graph_and_eager_modes
 class CommonLayersTest(parameterized.TestCase, tf.test.TestCase):
 
   def testIndexLastDimWithIndices(self):
@@ -35,149 +36,134 @@ def testIndexLastDimWithIndices(self):
     x_idx = common_layers.index_last_dim_with_indices(x, indices)
 
     expected = np.array([4., 6.])
-    with self.test_session() as sess:
-      self.assertAllEqual(expected, sess.run(x_idx))
+    self.assertAllEqual(expected, self.evaluate(x_idx))
 
   def testSaturatingSigmoid(self):
     x = np.array([-120.0, -100.0, 0.0, 100.0, 120.0], dtype=np.float32)
-    with self.test_session() as session:
-      y = common_layers.saturating_sigmoid(tf.constant(x))
-      res = session.run(y)
+    y = common_layers.saturating_sigmoid(tf.constant(x))
+    res = self.evaluate(y)
     self.assertAllClose(res, [0.0, 0.0, 0.5, 1.0, 1.0])
 
   def testFlatten4D3D(self):
     x = np.random.random_integers(1, high=8, size=(3, 5, 2))
-    with self.test_session() as session:
-      y = common_layers.flatten4d3d(common_layers.embedding(x, 10, 7))
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.flatten4d3d(common_layers.embedding(x, 10, 7))
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (3, 5 * 2, 7))
 
   def testEmbedding(self):
     x = np.random.random_integers(1, high=8, size=(3, 5))
-    with self.test_session() as session:
-      y = common_layers.embedding(x, 10, 16)
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.embedding(x, 10, 16)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (3, 5, 16))
 
   def testShakeShake(self):
+    if tf.executing_eagerly():
+      return  # don't run test in Eager mode
+
     x = np.random.rand(5, 7)
-    with self.test_session() as session:
+    with self.session() as session:
       x = tf.constant(x, dtype=tf.float32)
       y = common_layers.shakeshake([x, x, x, x, x])
-      session.run(tf.global_variables_initializer())
       inp, res = session.run([x, y])
     self.assertAllClose(res, inp)
 
   def testConv(self):
     x = np.random.rand(5, 7, 1, 11)
-    with self.test_session() as session:
-      y = common_layers.conv(tf.constant(x, dtype=tf.float32), 13, (3, 1))
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.conv(tf.constant(x, dtype=tf.float32), 13, (3, 1))
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 5, 1, 13))
 
   def testConv1d(self):
     x = np.random.rand(5, 7, 11)
-    with self.test_session() as session:
-      y = common_layers.conv1d(tf.constant(x, dtype=tf.float32), 13, 1)
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.conv1d(tf.constant(x, dtype=tf.float32), 13, 1)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 13))
 
   def testSeparableConv(self):
     x = np.random.rand(5, 7, 1, 11)
-    with self.test_session() as session:
-      y = common_layers.separable_conv(
-          tf.constant(x, dtype=tf.float32), 13, (3, 1))
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.separable_conv(
+        tf.constant(x, dtype=tf.float32), 13, (3, 1))
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 5, 1, 13))
 
   def testSubSeparableConv(self):
     for sep in [0, 1, 2, 4]:
       x = np.random.rand(5, 7, 1, 12)
-      with self.test_session() as session:
-        with tf.variable_scope("sep_%d" % sep):
-          y = common_layers.subseparable_conv(
-              tf.constant(x, dtype=tf.float32), 16, (3, 1), separability=sep)
-        session.run(tf.global_variables_initializer())
-        res = session.run(y)
+      with tf.variable_scope("sep_%d" % sep):
+        y = common_layers.subseparable_conv(
+            tf.constant(x, dtype=tf.float32), 16, (3, 1), separability=sep)
+      self.evaluate(tf.global_variables_initializer())
+      res = self.evaluate(y)
       self.assertEqual(res.shape, (5, 5, 1, 16))
 
   def testConvBlock(self):
     x = np.random.rand(5, 7, 1, 11)
-    with self.test_session() as session:
-      y = common_layers.conv_block(
-          tf.constant(x, dtype=tf.float32),
-          13, [(1, (3, 3)), (1, (3, 3))],
-          padding="SAME",
-          normalizer_fn=common_layers.noam_norm)
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.conv_block(
+        tf.constant(x, dtype=tf.float32),
+        13, [(1, (3, 3)), (1, (3, 3))],
+        padding="SAME",
+        normalizer_fn=common_layers.noam_norm)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 1, 13))
 
   def testSeparableConvBlock(self):
     x = np.random.rand(5, 7, 1, 11)
-    with self.test_session() as session:
-      y = common_layers.separable_conv_block(
-          tf.constant(x, dtype=tf.float32),
-          13, [(1, (3, 3)), (1, (3, 3))],
-          padding="SAME")
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.separable_conv_block(
+        tf.constant(x, dtype=tf.float32),
+        13, [(1, (3, 3)), (1, (3, 3))],
+        padding="SAME")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 1, 13))
 
   def testSubSeparableConvBlock(self):
     for sep in [0, 1, 2, 4]:
       x = np.random.rand(5, 7, 1, 12)
-      with self.test_session() as session:
-        with tf.variable_scope("sep_%d" % sep):
-          y = common_layers.subseparable_conv_block(
-              tf.constant(x, dtype=tf.float32),
-              16, [(1, (3, 3)), (1, (3, 3))],
-              padding="SAME",
-              separability=sep)
-        session.run(tf.global_variables_initializer())
-        res = session.run(y)
+      with tf.variable_scope("sep_%d" % sep):
+        y = common_layers.subseparable_conv_block(
+            tf.constant(x, dtype=tf.float32),
+            16, [(1, (3, 3)), (1, (3, 3))],
+            padding="SAME",
+            separability=sep)
+      self.evaluate(tf.global_variables_initializer())
+      res = self.evaluate(y)
       self.assertEqual(res.shape, (5, 7, 1, 16))
 
   def testPool(self):
     x = np.random.rand(5, 8, 1, 11)
-    with self.test_session() as session:
-      y = common_layers.pool(
-          tf.constant(x, dtype=tf.float32), (2, 2), "AVG", "SAME")
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.pool(
+        tf.constant(x, dtype=tf.float32), (2, 2), "AVG", "SAME")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 8, 1, 11))
 
   def testConvBlockDownsample(self):
     x = np.random.rand(5, 7, 1, 11)
-    with self.test_session() as session:
-      y = common_layers.conv_block_downsample(
-          tf.constant(x, dtype=tf.float32), (3, 1), (2, 1), "SAME")
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.conv_block_downsample(
+        tf.constant(x, dtype=tf.float32), (3, 1), (2, 1), "SAME")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 4, 1, 27))
 
   def testSimpleAttention(self):
     x = np.random.rand(5, 7, 1, 11)
     y = np.random.rand(5, 9, 1, 11)
-    with self.test_session() as session:
-      a = common_layers.simple_attention(
-          tf.constant(x, dtype=tf.float32), tf.constant(y, dtype=tf.float32))
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_layers.simple_attention(
+        tf.constant(x, dtype=tf.float32), tf.constant(y, dtype=tf.float32))
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 1, 11))
 
   def testGetTimingSignal(self):
     length = 7
     num_timescales = 10
-    with self.test_session() as session:
-      a = common_layers.get_timing_signal(length, num_timescales=num_timescales)
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_layers.get_timing_signal(length, num_timescales=num_timescales)
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (length, 2 * num_timescales))
 
   def testAddTimingSignal(self):
@@ -186,10 +172,8 @@ def testAddTimingSignal(self):
     height = 3
     depth = 35
     x = np.random.rand(batch, length, height, depth)
-    with self.test_session() as session:
-      a = common_layers.add_timing_signal(tf.constant(x, dtype=tf.float32))
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_layers.add_timing_signal(tf.constant(x, dtype=tf.float32))
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (batch, length, height, depth))
 
   def testAttention1D(self):
@@ -204,42 +188,42 @@ def testAttention1D(self):
     source = np.random.rand(batch, source_length, source_depth)
     target = np.random.rand(batch, target_length, target_depth)
     mask = np.random.rand(batch, target_length, source_length)
-    with self.test_session() as session:
-      a = common_layers.attention_1d_v0(
-          tf.constant(source, dtype=tf.float32),
-          tf.constant(target, dtype=tf.float32), attention_size, output_size,
-          num_heads, tf.constant(mask, dtype=tf.float32))
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_layers.attention_1d_v0(
+        tf.constant(source, dtype=tf.float32),
+        tf.constant(target, dtype=tf.float32), attention_size, output_size,
+        num_heads, tf.constant(mask, dtype=tf.float32))
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (batch, target_length, output_size))
 
   def testMultiscaleConvSum(self):
     x = np.random.rand(5, 9, 1, 11)
-    with self.test_session() as session:
-      y = common_layers.multiscale_conv_sum(
-          tf.constant(x, dtype=tf.float32),
-          13, [((1, 1), (5, 5)), ((2, 2), (3, 3))],
-          "AVG",
-          padding="SAME")
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.multiscale_conv_sum(
+        tf.constant(x, dtype=tf.float32),
+        13, [((1, 1), (5, 5)), ((2, 2), (3, 3))],
+        "AVG",
+        padding="SAME")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 9, 1, 13))
 
   def testConvGRU(self):
     x = np.random.rand(5, 7, 3, 11)
-    with self.test_session() as session:
-      y = common_layers.conv_gru(tf.constant(x, dtype=tf.float32), (1, 3), 11)
-      z = common_layers.conv_gru(
-          tf.constant(x, dtype=tf.float32), (1, 3), 11, padding="LEFT")
-      session.run(tf.global_variables_initializer())
-      res1 = session.run(y)
-      res2 = session.run(z)
+    y = common_layers.conv_gru(tf.constant(x, dtype=tf.float32), (1, 3), 11)
+    z = common_layers.conv_gru(
+        tf.constant(x, dtype=tf.float32), (1, 3), 11, padding="LEFT")
+    self.evaluate(tf.global_variables_initializer())
+    res1 = self.evaluate(y)
+    res2 = self.evaluate(z)
     self.assertEqual(res1.shape, (5, 7, 3, 11))
     self.assertEqual(res2.shape, (5, 7, 3, 11))
 
   def testSRU(self):
+    if tf.executing_eagerly():
+      return  # don't run test in Eager mode
+
     x = np.random.rand(5, 7, 3, 11)
-    with self.test_session() as session:
+    with self.session() as session:
       y = common_layers.sru(tf.constant(x, dtype=tf.float32))
       session.run(tf.global_variables_initializer())
       res = session.run(y)
@@ -247,40 +231,36 @@ def testSRU(self):
 
   def testLayerNorm(self):
     x = np.random.rand(5, 7, 11)
-    with self.test_session() as session:
-      y = common_layers.layer_norm(tf.constant(x, dtype=tf.float32), 11)
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.layer_norm(tf.constant(x, dtype=tf.float32), 11)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 11))
 
   def testGroupNorm(self):
     x = np.random.rand(5, 7, 3, 16)
-    with self.test_session() as session:
-      y = common_layers.group_norm(tf.constant(x, dtype=tf.float32))
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.group_norm(tf.constant(x, dtype=tf.float32))
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 3, 16))
 
   def testConvLSTM(self):
     x = np.random.rand(5, 7, 11, 13)
-    with self.test_session() as session:
-      y = common_layers.conv_lstm(tf.constant(x, dtype=tf.float32), (1, 3), 13)
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.conv_lstm(tf.constant(x, dtype=tf.float32), (1, 3), 13)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 11, 13))
 
   def testPadToSameLength(self):
     x1 = np.random.rand(5, 7, 11)
     x2 = np.random.rand(5, 9, 11)
-    with self.test_session() as session:
-      a, b = common_layers.pad_to_same_length(
-          tf.constant(x1, dtype=tf.float32), tf.constant(x2, dtype=tf.float32))
-      c, d = common_layers.pad_to_same_length(
-          tf.constant(x1, dtype=tf.float32),
-          tf.constant(x2, dtype=tf.float32),
-          final_length_divisible_by=4)
-      res1, res2 = session.run([a, b])
-      res1a, res2a = session.run([c, d])
+    a, b = common_layers.pad_to_same_length(
+        tf.constant(x1, dtype=tf.float32), tf.constant(x2, dtype=tf.float32))
+    c, d = common_layers.pad_to_same_length(
+        tf.constant(x1, dtype=tf.float32),
+        tf.constant(x2, dtype=tf.float32),
+        final_length_divisible_by=4)
+    res1, res2 = self.evaluate([a, b])
+    res1a, res2a = self.evaluate([c, d])
     self.assertEqual(res1.shape, (5, 9, 11))
     self.assertEqual(res2.shape, (5, 9, 11))
     self.assertEqual(res1a.shape, (5, 12, 11))
@@ -291,63 +271,56 @@ def testShiftLeft(self):
     x1[:, 0, :] = np.ones_like(x1[:, 0, :])
     expected = np.zeros((5, 7, 1, 11))
     expected[:, 1, :] = np.ones_like(expected[:, 1, :])
-    with self.test_session() as session:
-      a = common_layers.shift_right(tf.constant(x1, dtype=tf.float32))
-      actual = session.run(a)
+    a = common_layers.shift_right(tf.constant(x1, dtype=tf.float32))
+    actual = self.evaluate(a)
     self.assertAllEqual(actual, expected)
 
   def testConvStride2MultiStep(self):
     x1 = np.random.rand(5, 32, 16, 11)
-    with self.test_session() as session:
-      a = common_layers.conv_stride2_multistep(
-          tf.constant(x1, dtype=tf.float32), 4, 16)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(a[0])
+    a = common_layers.conv_stride2_multistep(
+        tf.constant(x1, dtype=tf.float32), 4, 16)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(a[0])
     self.assertEqual(actual.shape, (5, 2, 1, 16))
 
   def testDeconvStride2MultiStep(self):
     x1 = np.random.rand(5, 2, 1, 11)
-    with self.test_session() as session:
-      a = common_layers.deconv_stride2_multistep(
-          tf.constant(x1, dtype=tf.float32), 4, 16)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(a)
+    a = common_layers.deconv_stride2_multistep(
+        tf.constant(x1, dtype=tf.float32), 4, 16)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(a)
     self.assertEqual(actual.shape, (5, 32, 1, 16))
 
   def testApplyNormLayer(self):
-    with self.test_session() as session:
-      x1 = np.random.rand(5, 2, 1, 11)
-      x2 = common_layers.apply_norm(
-          tf.constant(x1, dtype=tf.float32), "layer", depth=11, epsilon=1e-6)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(x2)
+    x1 = np.random.rand(5, 2, 1, 11)
+    x2 = common_layers.apply_norm(
+        tf.constant(x1, dtype=tf.float32), "layer", depth=11, epsilon=1e-6)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
   def testApplyNormNoam(self):
-    with self.test_session() as session:
-      x1 = np.random.rand(5, 2, 1, 11)
-      x2 = common_layers.apply_norm(
-          tf.constant(x1, dtype=tf.float32), "noam", depth=11, epsilon=1e-6)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(x2)
+    x1 = np.random.rand(5, 2, 1, 11)
+    x2 = common_layers.apply_norm(
+        tf.constant(x1, dtype=tf.float32), "noam", depth=11, epsilon=1e-6)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
   def testApplyNormBatch(self):
-    with self.test_session() as session:
-      x1 = np.random.rand(5, 2, 1, 11)
-      x2 = common_layers.apply_norm(
-          tf.constant(x1, dtype=tf.float32), "batch", depth=11, epsilon=1e-6)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(x2)
+    x1 = np.random.rand(5, 2, 1, 11)
+    x2 = common_layers.apply_norm(
+        tf.constant(x1, dtype=tf.float32), "batch", depth=11, epsilon=1e-6)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
   def testApplyNormNone(self):
-    with self.test_session() as session:
-      x1 = np.random.rand(5, 2, 1, 11)
-      x2 = common_layers.apply_norm(
-          tf.constant(x1, dtype=tf.float32), "none", depth=11, epsilon=1e-6)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(x2)
+    x1 = np.random.rand(5, 2, 1, 11)
+    x2 = common_layers.apply_norm(
+        tf.constant(x1, dtype=tf.float32), "none", depth=11, epsilon=1e-6)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
     self.assertAllClose(actual, x1, atol=1e-03)
 
@@ -356,93 +329,88 @@ def testGlobalPool1d(self):
     no_mask = np.ones((5, 4))
     full_mask = np.zeros((5, 4))
 
-    with self.test_session() as session:
-      x1_ = tf.Variable(x1, dtype=tf.float32)
-      no_mask_ = tf.Variable(no_mask, dtype=tf.float32)
-      full_mask_ = tf.Variable(full_mask, dtype=tf.float32)
+    x1_ = tf.Variable(x1, dtype=tf.float32)
+    no_mask_ = tf.Variable(no_mask, dtype=tf.float32)
+    full_mask_ = tf.Variable(full_mask, dtype=tf.float32)
 
-      none_mask_max = common_layers.global_pool_1d(x1_)
-      no_mask_max = common_layers.global_pool_1d(x1_, mask=no_mask_)
-      result1 = tf.reduce_sum(none_mask_max - no_mask_max)
+    none_mask_max = common_layers.global_pool_1d(x1_)
+    no_mask_max = common_layers.global_pool_1d(x1_, mask=no_mask_)
+    result1 = tf.reduce_sum(none_mask_max - no_mask_max)
 
-      full_mask_max = common_layers.global_pool_1d(x1_, mask=full_mask_)
-      result2 = tf.reduce_sum(full_mask_max)
+    full_mask_max = common_layers.global_pool_1d(x1_, mask=full_mask_)
+    result2 = tf.reduce_sum(full_mask_max)
 
-      none_mask_avr = common_layers.global_pool_1d(x1_, "AVR")
-      no_mask_avr = common_layers.global_pool_1d(x1_, "AVR", no_mask_)
-      result3 = tf.reduce_sum(none_mask_avr - no_mask_avr)
+    none_mask_avr = common_layers.global_pool_1d(x1_, "AVR")
+    no_mask_avr = common_layers.global_pool_1d(x1_, "AVR", no_mask_)
+    result3 = tf.reduce_sum(none_mask_avr - no_mask_avr)
 
-      full_mask_avr = common_layers.global_pool_1d(x1_, "AVR", full_mask_)
-      result4 = tf.reduce_sum(full_mask_avr)
+    full_mask_avr = common_layers.global_pool_1d(x1_, "AVR", full_mask_)
+    result4 = tf.reduce_sum(full_mask_avr)
 
-      session.run(tf.global_variables_initializer())
-      actual = session.run([result1, result2, result3, result4])
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate([result1, result2, result3, result4])
     self.assertAllEqual(actual[:3], [0.0, 0.0, 0.0])
 
   def testLinearSetLayer(self):
     x1 = np.random.rand(5, 4, 11)
     cont = np.random.rand(5, 13)
-    with self.test_session() as session:
-      x1_ = tf.Variable(x1, dtype=tf.float32)
-      cont_ = tf.Variable(cont, dtype=tf.float32)
+    x1_ = tf.Variable(x1, dtype=tf.float32)
+    cont_ = tf.Variable(cont, dtype=tf.float32)
 
-      simple_ff = common_layers.linear_set_layer(32, x1_)
-      cont_ff = common_layers.linear_set_layer(32, x1_, context=cont_)
+    simple_ff = common_layers.linear_set_layer(32, x1_)
+    cont_ff = common_layers.linear_set_layer(32, x1_, context=cont_)
 
-      session.run(tf.global_variables_initializer())
-      actual = session.run([simple_ff, cont_ff])
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate([simple_ff, cont_ff])
     self.assertEqual(actual[0].shape, (5, 4, 32))
     self.assertEqual(actual[1].shape, (5, 4, 32))
 
   def testRavanbakhshSetLayer(self):
     x1 = np.random.rand(5, 4, 11)
-    with self.test_session() as session:
-      x1_ = tf.Variable(x1, dtype=tf.float32)
-      layer = common_layers.ravanbakhsh_set_layer(32, x1_)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(layer)
+    x1_ = tf.Variable(x1, dtype=tf.float32)
+    layer = common_layers.ravanbakhsh_set_layer(32, x1_)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(layer)
     self.assertEqual(actual.shape, (5, 4, 32))
 
   def testBReLU(self):
-    with self.test_session() as session:
-      x = np.random.rand(5, 2, 1, 12)
-      y = common_layers.brelu(tf.constant(x, dtype=tf.float32))
-      actual = session.run(y)
+    x = np.random.rand(5, 2, 1, 12)
+    y = common_layers.brelu(tf.constant(x, dtype=tf.float32))
+    actual = self.evaluate(y)
     self.assertEqual(actual.shape, (5, 2, 1, 12))
 
   def testBELU(self):
-    with self.test_session() as session:
-      x = np.random.rand(5, 2, 1, 12)
-      y = common_layers.belu(tf.constant(x, dtype=tf.float32))
-      actual = session.run(y)
+    x = np.random.rand(5, 2, 1, 12)
+    y = common_layers.belu(tf.constant(x, dtype=tf.float32))
+    actual = self.evaluate(y)
     self.assertEqual(actual.shape, (5, 2, 1, 12))
 
   def testNAC(self):
-    with self.test_session() as session:
-      x = np.random.rand(5, 2, 1, 12)
-      y = common_layers.nac(tf.constant(x, dtype=tf.float32), 14)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(y)
+    x = np.random.rand(5, 2, 1, 12)
+    y = common_layers.nac(tf.constant(x, dtype=tf.float32), 14)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(y)
     self.assertEqual(actual.shape, (5, 2, 1, 14))
 
   def testNALU(self):
-    with self.test_session() as session:
-      x = np.random.rand(5, 2, 1, 12)
-      y = common_layers.nalu(tf.constant(x, dtype=tf.float32), 14)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(y)
+    x = np.random.rand(5, 2, 1, 12)
+    y = common_layers.nalu(tf.constant(x, dtype=tf.float32), 14)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(y)
     self.assertEqual(actual.shape, (5, 2, 1, 14))
 
   def testNALUzeros(self):
-    with self.test_session() as session:
-      x = np.random.rand(5, 2, 1, 12)
-      y = common_layers.nalu(tf.zeros_like(x, dtype=tf.float32), 14)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(y)
+    x = np.random.rand(5, 2, 1, 12)
+    y = common_layers.nalu(tf.zeros_like(x, dtype=tf.float32), 14)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(y)
     self.assertTrue(np.all(np.isfinite(actual)))
     self.assertEqual(actual.shape, (5, 2, 1, 14))
 
   def testPaddingCrossEntropyFactored(self):
+    if tf.executing_eagerly():
+      return  # don't run test in Eager mode
+
     vocab_size = 19
     rows = 5
     cols = 4
@@ -451,7 +419,7 @@ def testPaddingCrossEntropyFactored(self):
     features = np.random.rand(rows, cols, depth)
     weights = np.random.rand(vocab_size, depth)
     labels = np.random.randint(0, vocab_size - 1, size=(rows, cols))
-    with self.test_session() as session:
+    with self.session() as session:
       features = tf.to_float(features)
       weights = tf.to_float(weights)
       labels = tf.to_int32(labels)
@@ -476,6 +444,9 @@ def testPaddingCrossEntropyFactored(self):
     self.assertAllClose(den, den_f)
 
   def testPaddingCrossEntropyFactoredGrad(self):
+    if tf.executing_eagerly():
+      return  # don't run test in Eager mode
+
     vocab_size = 19
     rows = 5
     cols = 4
@@ -484,7 +455,7 @@ def testPaddingCrossEntropyFactoredGrad(self):
     features = np.random.rand(rows, cols, depth)
     weights = np.random.rand(vocab_size, depth)
     labels = np.random.randint(0, vocab_size - 1, size=(rows, cols))
-    with self.test_session() as session:
+    with self.session() as session:
       features = tf.to_float(features)
       weights = tf.to_float(weights)
       labels = tf.to_int32(labels)
@@ -533,9 +504,8 @@ def testDmlLoss(self, batch, height, width, num_mixtures, reduce_sum):
     if reduce_sum:
       expected_loss = tf.reduce_mean(expected_loss)
 
-    with self.test_session() as sess:
-      actual_loss_val, expected_loss_val = sess.run(
-          [actual_loss, expected_loss])
+    actual_loss_val, expected_loss_val = self.evaluate(
+        [actual_loss, expected_loss])
     self.assertAllClose(actual_loss_val, expected_loss_val)
 
   def testDiscretizedMixLogisticLoss(self):
@@ -570,9 +540,8 @@ def testDiscretizedMixLogisticLoss(self):
 
     actual_loss = common_layers.discretized_mix_logistic_loss(
         pred=pred, labels=labels)
-    with self.test_session() as session:
-      actual_loss_val, expected_loss_val = session.run(
-          [actual_loss, expected_loss])
+    actual_loss_val, expected_loss_val = self.evaluate(
+        [actual_loss, expected_loss])
     self.assertAllClose(actual_loss_val, expected_loss_val, rtol=1e-5)
 
   def testSampleFromDiscretizedMixLogistic(self):
@@ -596,9 +565,8 @@ def testSampleFromDiscretizedMixLogistic(self):
 
     actual_sample = common_layers.sample_from_discretized_mix_logistic(
         pred, seed=seed)
-    with self.test_session() as session:
-      actual_sample_val, expected_sample_val = session.run(
-          [actual_sample, expected_sample])
+    actual_sample_val, expected_sample_val = self.evaluate(
+        [actual_sample, expected_sample])
     # Use a low tolerance: samples numerically differ, as the actual
     # implementation clips log-scales so they always contribute to sampling.
     self.assertAllClose(actual_sample_val, expected_sample_val, atol=1e-2)
@@ -607,22 +575,24 @@ def testFactoredTensorImplicitConversion(self):
     a = np.random.rand(3, 4, 5)
     b = np.random.rand(6, 5)
     c = np.random.rand(3, 4, 6)
-    with self.test_session() as session:
-      # a factored representation of a Tensor of shape (3, 4, 6)
-      factored = common_layers.FactoredTensor(tf.to_float(a), tf.to_float(b))
-      # implicitly converts factored to a Tensor (performing the matmul)
-      d = factored + tf.to_float(c)
-      out = session.run(d)
+    # a factored representation of a Tensor of shape (3, 4, 6)
+    factored = common_layers.FactoredTensor(tf.to_float(a), tf.to_float(b))
+    # implicitly converts factored to a Tensor (performing the matmul)
+    d = factored + tf.to_float(c)
+    out = self.evaluate(d)
     self.assertEqual(out.shape, (3, 4, 6))
 
   def testConvHiddenReluMemoryEfficient(self):
+    if tf.executing_eagerly():
+      return  # don't run test in Eager mode
+
     batch = 3
     length = 23
     io_size = 16
     filter_size = 7
     x = np.random.rand(batch, length, io_size)
     dy = np.random.rand(batch, length, io_size)
-    with self.test_session() as session:
+    with self.session() as session:
       x = tf.to_float(x)
       dy = tf.to_float(dy)
       f1 = tf.get_variable("f1", [1, io_size, filter_size])
@@ -666,11 +636,10 @@ def testCycleGANUpsampleNnUpsampleConv(self):
     upsampled_output = common_layers.cyclegan_upsample(
         random_input, output_filters, stride, "nn_upsample_conv")
     upsampled_output_shape = tf.shape(upsampled_output)
-    with self.test_session() as session:
-      session.run(tf.global_variables_initializer())
-      self.assertAllEqual(
-          [batch, height * stride[0], width * stride[1], output_filters],
-          session.run(upsampled_output_shape))
+    self.evaluate(tf.global_variables_initializer())
+    self.assertAllEqual(
+        [batch, height * stride[0], width * stride[1], output_filters],
+        self.evaluate(upsampled_output_shape))
 
   def testCycleGANUpsampleBilinearUpsampleConv(self):
     batch = 8
@@ -686,11 +655,10 @@ def testCycleGANUpsampleBilinearUpsampleConv(self):
     upsampled_output = common_layers.cyclegan_upsample(
         random_input, output_filters, stride, "bilinear_upsample_conv")
     upsampled_output_shape = tf.shape(upsampled_output)
-    with self.test_session() as session:
-      session.run(tf.global_variables_initializer())
-      self.assertAllEqual(
-          [batch, height * stride[0], width * stride[1], output_filters],
-          session.run(upsampled_output_shape))
+    self.evaluate(tf.global_variables_initializer())
+    self.assertAllEqual(
+        [batch, height * stride[0], width * stride[1], output_filters],
+        self.evaluate(upsampled_output_shape))
 
   def testCycleGANUpsampleConv2dTranspose(self):
     batch = 8
@@ -711,13 +679,15 @@ def testCycleGANUpsampleConv2dTranspose(self):
                                                        output_filters, stride,
                                                        "conv2d_transpose")
     upsampled_output_shape = tf.shape(upsampled_output)
-    with self.test_session() as session:
-      session.run(tf.global_variables_initializer())
-      self.assertAllEqual(
-          [batch, upsampled_height, upsampled_width, output_filters],
-          session.run(upsampled_output_shape))
+    self.evaluate(tf.global_variables_initializer())
+    self.assertAllEqual(
+        [batch, upsampled_height, upsampled_width, output_filters],
+        self.evaluate(upsampled_output_shape))
 
   def testSpectralNorm(self):
+    if tf.executing_eagerly():
+      return  # don't run test in Eager mode
+
     # Test that after 20 calls to apply_spectral_norm, the spectral
     # norm of the normalized matrix is close to 1.0
     with tf.Graph().as_default():
@@ -776,7 +746,7 @@ def grad_fn(inputs, variables, outputs, grad_outputs):
     custom_grads = tf.gradients(custom_loss,
                                 [a, b, c] + [tf.trainable_variables()[1]])
 
-    with self.test_session() as sess:
+    with self.session() as sess:
       sess.run(tf.global_variables_initializer())
       out_val, custom_out_val, grads_val, custom_grads_val = sess.run(
           [out, custom_out, grads, custom_grads])
@@ -807,7 +777,7 @@ def grad_fn(inputs, variables, unused_outputs, unused_grad_outputs):
     expected_grads = [
         tf.ones_like(t) * (i + 1.) for i, t in enumerate([a, b, c, w])
     ]
-    with self.test_session() as sess:
+    with self.session() as sess:
       sess.run(tf.global_variables_initializer())
       g_val, eg_val = sess.run([grads, expected_grads])
       for g1, g2 in zip(g_val, eg_val):
@@ -853,7 +823,7 @@ def fn_recompute(x):
     grad1 = tf.gradients(out1, recompute_vars)
     grad2 = tf.gradients(out2, reg_vars)
 
-    with self.test_session() as sess:
+    with self.session() as sess:
       sess.run(tf.global_variables_initializer())
       outs = sess.run([out1, out2, grad1, grad2])
       self.assertAllClose(outs[0], outs[1])
diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index efa031957..fe4672410 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -22,79 +22,76 @@
 import tensorflow as tf
 
 
+@tf.contrib.eager.run_all_tests_in_graph_and_eager_modes
 class CommonVideoTest(tf.test.TestCase):
 
-  @staticmethod
-  def RunScheduledSampleFunc(func, var, batch_size):
+  def runScheduledSampleFunc(self, func, var, batch_size):
     ground_truth_x = list(range(1, batch_size+1))
     generated_x = [-x for x in ground_truth_x]
     ground_truth_x = tf.convert_to_tensor(ground_truth_x)
     generated_x = tf.convert_to_tensor(generated_x)
     ss_out = func(ground_truth_x, generated_x, batch_size, var)
-    with tf.Session() as session:
-      output = session.run([ground_truth_x, generated_x, ss_out])
+    output = self.evaluate([ground_truth_x, generated_x, ss_out])
     return output
 
   def testScheduledSampleProbStart(self):
-    ground_truth_x, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
+    ground_truth_x, _, ss_out = self.runScheduledSampleFunc(
         common_video.scheduled_sample_prob, 1.0, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
   def testScheduledSampleProbMid(self):
-    _, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
+    _, _, ss_out = self.runScheduledSampleFunc(
         common_video.scheduled_sample_prob, 0.5, 1000)
     positive_count = np.sum(ss_out > 0)
-    self.assertAlmostEqual(positive_count / 1000.0, 0.5, places=2)
+    self.assertAlmostEqual(positive_count / 1000.0, 0.5, places=1)
 
   def testScheduledSampleProbEnd(self):
-    _, generated_x, ss_out = CommonVideoTest.RunScheduledSampleFunc(
+    _, generated_x, ss_out = self.runScheduledSampleFunc(
         common_video.scheduled_sample_prob, 0.0, 10)
     self.assertAllEqual(generated_x, ss_out)
 
   def testScheduledSampleCountStart(self):
-    ground_truth_x, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
+    ground_truth_x, _, ss_out = self.runScheduledSampleFunc(
         common_video.scheduled_sample_count, 10, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
   def testScheduledSampleCountMid(self):
-    _, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
+    _, _, ss_out = self.runScheduledSampleFunc(
         common_video.scheduled_sample_count, 5, 10)
     positive_count = np.sum(ss_out > 0)
     self.assertEqual(positive_count, 5)
 
   def testScheduledSampleCountEnd(self):
-    _, generated_x, ss_out = CommonVideoTest.RunScheduledSampleFunc(
+    _, generated_x, ss_out = self.runScheduledSampleFunc(
         common_video.scheduled_sample_count, 0, 10)
     self.assertAllEqual(generated_x, ss_out)
 
   def testDynamicTileAndConcat(self):
-    with tf.Graph().as_default():
-      # image = (1 X 4 X 4 X 1)
-      image = [[1, 2, 3, 4],
-               [2, 4, 5, 6],
-               [7, 8, 9, 10],
-               [7, 9, 10, 1]]
-      image_t = tf.expand_dims(tf.expand_dims(image, axis=0), axis=-1)
-      image_t = tf.cast(image_t, dtype=tf.float32)
-
-      # latent = (1 X 2)
-      latent = np.array([[90, 100]])
-      latent_t = tf.cast(tf.convert_to_tensor(latent), dtype=tf.float32)
-
-      with tf.Session() as session:
-        tiled = common_video.tile_and_concat(
-            image_t, latent_t)
-        tiled_np, image_np = session.run([tiled, image_t])
-        tiled_latent = tiled_np[0, :, :, -1]
-        self.assertAllEqual(tiled_np.shape, (1, 4, 4, 2))
-
-        self.assertAllEqual(tiled_np[:, :, :, :1], image_np)
-        self.assertAllEqual(
-            tiled_latent,
-            [[90, 90, 90, 90],
-             [100, 100, 100, 100],
-             [90, 90, 90, 90],
-             [100, 100, 100, 100]])
+    # image = (1 X 4 X 4 X 1)
+    image = [[1, 2, 3, 4],
+             [2, 4, 5, 6],
+             [7, 8, 9, 10],
+             [7, 9, 10, 1]]
+    image_t = tf.expand_dims(tf.expand_dims(image, axis=0), axis=-1)
+    image_t = tf.cast(image_t, dtype=tf.float32)
+
+    # latent = (1 X 2)
+    latent = np.array([[90, 100]])
+    latent_t = tf.cast(tf.convert_to_tensor(latent), dtype=tf.float32)
+
+    tiled = common_video.tile_and_concat(
+        image_t, latent_t)
+    tiled_np, image_np = self.evaluate([tiled, image_t])
+    tiled_latent = tiled_np[0, :, :, -1]
+    self.assertAllEqual(tiled_np.shape, (1, 4, 4, 2))
+
+    self.assertAllEqual(tiled_np[:, :, :, :1], image_np)
+    self.assertAllEqual(
+        tiled_latent,
+        [[90, 90, 90, 90],
+         [100, 100, 100, 100],
+         [90, 90, 90, 90],
+         [100, 100, 100, 100]])
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index c75f61704..20dfc22b2 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -23,6 +23,7 @@
 import tensorflow as tf
 
 
+@tf.contrib.eager.run_all_tests_in_graph_and_eager_modes
 class DiscretizationTest(tf.test.TestCase):
   """Tests for discretization layers."""
 
@@ -34,37 +35,29 @@ def testBitToIntZeros(self):
     x_bit = tf.zeros(shape=[1, 10], dtype=tf.float32)
     x_int = tf.zeros(shape=[1], dtype=tf.int32)
     diff = discretization.bit_to_int(x_bit, num_bits=10) - x_int
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      d = sess.run(diff)
-      self.assertEqual(d, 0)
+    d = self.evaluate(diff)
+    self.assertEqual(d, 0)
 
   def testBitToIntOnes(self):
     x_bit = tf.ones(shape=[1, 3], dtype=tf.float32)
     x_int = 7 * tf.ones(shape=[1], dtype=tf.int32)
     diff = discretization.bit_to_int(x_bit, num_bits=3) - x_int
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      d = sess.run(diff)
-      self.assertEqual(d, 0)
+    d = self.evaluate(diff)
+    self.assertEqual(d, 0)
 
   def testIntToBitZeros(self):
     x_bit = tf.zeros(shape=[1, 10], dtype=tf.float32)
     x_int = tf.zeros(shape=[1], dtype=tf.int32)
     diff = discretization.int_to_bit(x_int, num_bits=10) - x_bit
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      d = sess.run(diff)
-      self.assertTrue(np.all(d == 0))
+    d = self.evaluate(diff)
+    self.assertTrue(np.all(d == 0))
 
   def testIntToBitOnes(self):
     x_bit = tf.ones(shape=[1, 3], dtype=tf.float32)
     x_int = 7 * tf.ones(shape=[1], dtype=tf.int32)
     diff = discretization.int_to_bit(x_int, num_bits=3) - x_bit
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      d = sess.run(diff)
-      self.assertTrue(np.all(d == 0))
+    d = self.evaluate(diff)
+    self.assertTrue(np.all(d == 0))
 
   def testProjectHidden(self):
     hidden_size = 60
@@ -75,11 +68,9 @@ def testProjectHidden(self):
         shape=[num_blocks, hidden_size, block_dim], dtype=tf.float32)
     x_projected = discretization.project_hidden(x, projection_tensors,
                                                 hidden_size, num_blocks)
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_projected_eval = sess.run(x_projected)
-      self.assertEqual(np.shape(x_projected_eval), (1, num_blocks, block_dim))
-      self.assertTrue(np.all(x_projected_eval == 0))
+    x_projected_eval = self.evaluate(x_projected)
+    self.assertEqual(np.shape(x_projected_eval), (1, num_blocks, block_dim))
+    self.assertTrue(np.all(x_projected_eval == 0))
 
   def testSliceHiddenZeros(self):
     hidden_size = 60
@@ -87,11 +78,9 @@ def testSliceHiddenZeros(self):
     num_blocks = 3
     x = tf.zeros(shape=[1, hidden_size], dtype=tf.float32)
     x_sliced = discretization.slice_hidden(x, hidden_size, num_blocks)
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_sliced_eval = sess.run(x_sliced)
-      self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
-      self.assertTrue(np.all(x_sliced_eval == 0))
+    x_sliced_eval = self.evaluate(x_sliced)
+    self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
+    self.assertTrue(np.all(x_sliced_eval == 0))
 
   def testSliceHiddenOnes(self):
     hidden_size = 60
@@ -99,11 +88,9 @@ def testSliceHiddenOnes(self):
     num_blocks = 3
     x = tf.ones(shape=[1, hidden_size], dtype=tf.float32)
     x_sliced = discretization.slice_hidden(x, hidden_size, num_blocks)
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_sliced_eval = sess.run(x_sliced)
-      self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
-      self.assertTrue(np.all(x_sliced_eval == 1))
+    x_sliced_eval = self.evaluate(x_sliced)
+    self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
+    self.assertTrue(np.all(x_sliced_eval == 1))
 
   def testNearestNeighbors(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
@@ -115,13 +102,14 @@ def testNearestNeighbors(self):
         x, means, block_v_size=4)
     x_means_hot_test = np.array([[0, 1, 0, 0], [1, 0, 0, 0]])
     x_means_hot_test = np.expand_dims(x_means_hot_test, axis=0)
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_means_hot_eval = sess.run(x_means_hot)
-      self.assertEqual(np.shape(x_means_hot_eval), (1, 2, 4))
-      self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
+    x_means_hot_eval = self.evaluate(x_means_hot)
+    self.assertEqual(np.shape(x_means_hot_eval), (1, 2, 4))
+    self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
 
   def testGetVQBottleneck(self):
+    if tf.executing_eagerly():
+      return  # don't run test in Eager mode
+
     bottleneck_bits = 2
     bottleneck_size = 2**bottleneck_bits
     hidden_size = 3
@@ -130,7 +118,7 @@ def testGetVQBottleneck(self):
     assign_op = means.assign(tf.zeros(shape=[bottleneck_size, hidden_size]))
     means_new, _, _ = discretization.get_vq_codebook(bottleneck_size,
                                                      hidden_size)
-    with self.test_session() as sess:
+    with self.session() as sess:
       tf.global_variables_initializer().run()
       sess.run(assign_op)
       self.assertTrue(np.all(sess.run(means_new) == 0))
@@ -142,37 +130,32 @@ def testVQNearestNeighbors(self):
         [[1, 0, 0], [0, 1, 0], [0, 0, 1], [9, 9, 9]], dtype=tf.float32)
     x_means_hot, _, _ = discretization.vq_nearest_neighbor(x, means)
     x_means_hot_test = np.array([[0, 1, 0, 0], [1, 0, 0, 0]])
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_means_hot_eval = sess.run(x_means_hot)
-      self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
-      self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
+    x_means_hot_eval = self.evaluate(x_means_hot)
+    self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
+    self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
 
   def testVQDiscreteBottleneck(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
     x_means_hot, _ = discretization.vq_discrete_bottleneck(x, bottleneck_bits=2)
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_means_hot_eval = sess.run(x_means_hot)
-      self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
+    self.evaluate(tf.global_variables_initializer())
+    x_means_hot_eval = self.evaluate(x_means_hot)
+    self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
 
   def testVQDiscreteUnbottlenck(self):
     x = tf.constant([[1, 0, 0, 0], [0, 0, 1, 0]], dtype=tf.int32)
     x_means = discretization.vq_discrete_unbottleneck(x, hidden_size=3)
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_means_eval = sess.run(x_means)
-      self.assertEqual(np.shape(x_means_eval), (2, 3))
+    self.evaluate(tf.global_variables_initializer())
+    x_means_eval = self.evaluate(x_means)
+    self.assertEqual(np.shape(x_means_eval), (2, 3))
 
   def testGumbelSoftmaxDiscreteBottleneck(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
     tf.add_to_collection(tf.GraphKeys.GLOBAL_STEP, tf.constant(1))
     x_means_hot, _ = discretization.gumbel_softmax_discrete_bottleneck(
         x, bottleneck_bits=2)
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_means_hot_eval = sess.run(x_means_hot)
-      self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
+    self.evaluate(tf.global_variables_initializer())
+    x_means_hot_eval = self.evaluate(x_means_hot)
+    self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
 
 
 if __name__ == '__main__':
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 04fba3bff..0a8ef0311 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -27,6 +27,7 @@
 
 class ModalityTest(tf.test.TestCase):
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes
   def testSymbolModalityInputs(self):
     batch_size = 10
     num_datashards = 5
@@ -41,14 +42,14 @@ def testSymbolModalityInputs(self):
     m = modalities.SymbolModality(model_hparams, vocab_size)
     data_parallelism = expert_utils.Parallelism(
         ["/device:CPU:0"] * num_datashards)
-    with self.test_session() as session:
-      xs = tf.split(x, num_datashards)
-      sharded_output = m.bottom_sharded(xs, data_parallelism)
-      output = tf.concat(sharded_output, 0)
-      session.run(tf.global_variables_initializer())
-      res = session.run(output)
+    xs = tf.split(x, num_datashards)
+    sharded_output = m.bottom_sharded(xs, data_parallelism)
+    output = tf.concat(sharded_output, 0)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(output)
     self.assertEqual(res.shape, (batch_size, length, 1, hidden_size))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes
   def testSymbolModalityTargets(self):
     batch_size = 10
     num_datashards = 5
@@ -66,16 +67,15 @@ def testSymbolModalityTargets(self):
     m = modalities.SymbolModality(model_hparams, vocab_size)
     data_parallelism = expert_utils.Parallelism(
         ["/device:CPU:0"] * num_datashards)
-    with self.test_session() as session:
-      sharded_body_output = tf.split(tf.to_float(body_output), num_datashards)
-      sharded_targets = tf.split(targets, num_datashards)
-      sharded_logits = m.top_sharded(sharded_body_output, sharded_targets,
-                                     data_parallelism)
-      train_loss = m.loss_sharded(sharded_logits, sharded_targets,
-                                  data_parallelism)
-      logits = tf.concat(sharded_logits, 0)
-      session.run(tf.global_variables_initializer())
-      res1, res2 = session.run((logits, train_loss))
+    sharded_body_output = tf.split(tf.to_float(body_output), num_datashards)
+    sharded_targets = tf.split(targets, num_datashards)
+    sharded_logits = m.top_sharded(sharded_body_output, sharded_targets,
+                                   data_parallelism)
+    train_loss = m.loss_sharded(sharded_logits, sharded_targets,
+                                data_parallelism)
+    logits = tf.concat(sharded_logits, 0)
+    self.evaluate(tf.global_variables_initializer())
+    res1, res2 = self.evaluate((logits, train_loss))
     self.assertEqual(res1.shape, (batch_size, length, height, 1, vocab_size))
     self.assertEqual(res2.shape, ())
 
@@ -97,7 +97,7 @@ def testSymbolModalityTargetsFactored(self):
     m = modalities.SymbolModality(model_hparams, vocab_size)
     data_parallelism = expert_utils.Parallelism(
         ["/device:CPU:0"] * num_datashards)
-    with self.test_session() as session:
+    with self.session() as session:
       sharded_body_output = tf.split(tf.to_float(body_output), num_datashards)
       sharded_targets = tf.split(targets, num_datashards)
       sharded_logits = m.top_sharded(sharded_body_output, sharded_targets,

From f7e1f959749899e154e37fe4b0d65c158748e6c9 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Wed, 22 Aug 2018 15:40:20 -0700
Subject: [PATCH 0644/2720] use agg as the matplotlib backend.

PiperOrigin-RevId: 209842898
---
 tensor2tensor/data_generators/video_generated.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 60088be97..14fe3d25d 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -29,6 +29,8 @@
 import tensorflow as tf
 
 try:
+  import matplotlib  # pylint: disable=g-import-not-at-top
+  matplotlib.use("agg")
   import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top
 except ImportError:
   pass

From f3ca97cfc4a9002003a41b5489cc50f5cbbc06b0 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 22 Aug 2018 17:38:14 -0700
Subject: [PATCH 0645/2720] Internal change

PiperOrigin-RevId: 209859953
---
 tensor2tensor/rl/trainer_model_based.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index e9a3d4771..d1717da7d 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -549,8 +549,8 @@ def rl_modelrl_base():
 def rl_modelrl_base_stochastic():
   """Base setting with a stochastic next-frame model."""
   hparams = rl_modelrl_base()
-  hparams.generative_model = "next_frame_stochastic"
-  hparams.generative_model_params = "next_frame_stochastic_cutoff"
+  hparams.generative_model = "next_frame_basic_stochastic"
+  hparams.generative_model_params = "next_frame_basic_stochastic"
   return hparams
 
 

From acb7002e91460a327924ed9d9b4df30f3c0749fd Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 22 Aug 2018 17:47:10 -0700
Subject: [PATCH 0646/2720] fixing beta bug

PiperOrigin-RevId: 209861048
---
 tensor2tensor/models/research/next_frame_base_vae.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_base_vae.py b/tensor2tensor/models/research/next_frame_base_vae.py
index 8d5da6303..e9c4bad6e 100644
--- a/tensor2tensor/models/research/next_frame_base_vae.py
+++ b/tensor2tensor/models/research/next_frame_base_vae.py
@@ -50,12 +50,11 @@ def get_beta(self):
 
   def get_extra_loss(self, mean, std):
     """Losses in addition to the default modality losses."""
-    if self.is_training:
-      beta = self.get_beta()
-      kl_loss = common_layers.kl_divergence(mean, std)
-      tf.summary.histogram("posterior_mean", mean)
-      tf.summary.histogram("posterior_std", std)
-      tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
+    beta = self.get_beta()
+    kl_loss = common_layers.kl_divergence(mean, std)
+    tf.summary.histogram("posterior_mean", mean)
+    tf.summary.histogram("posterior_std", std)
+    tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
     return beta * kl_loss
 
   def construct_latent_tower(self, images, time_axis):

From bcd815fb2d11c837c182d1c2912cca83a60200ea Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 22 Aug 2018 17:49:54 -0700
Subject: [PATCH 0647/2720] Accept sync=true for single machine training

PiperOrigin-RevId: 209861368
---
 tensor2tensor/utils/devices.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index 7aae0a3e0..2d0cbe48e 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -136,7 +136,6 @@ def _replica_device_setter(worker_device):
     datashard_devices = [""]
     caching_devices = None
   elif is_single_machine:
-    assert not sync
     tf.logging.warn(
         "Schedule=%s. Assuming that training is running on a single machine.",
         schedule)

From b73c520c98300f75b9a75241c48f081df6226e2a Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 22 Aug 2018 18:43:11 -0700
Subject: [PATCH 0648/2720] Multi-frame reward prediction. 2nd try!

PiperOrigin-RevId: 209866728
---
 tensor2tensor/layers/common_video.py          |   4 +-
 .../models/research/next_frame_sv2p.py        | 105 ++++++++++--------
 .../models/research/next_frame_sv2p_params.py |   1 +
 3 files changed, 64 insertions(+), 46 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 1f7a27459..0698387a1 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -33,7 +33,7 @@ def swap_time_and_batch_axes(inputs):
 def encode_to_shape(inputs, shape, scope):
   """Encode the given tensor to given image shape."""
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
-    w, h = shape[1].value, shape[2].value
+    w, h = shape[1], shape[2]
     x = inputs
     x = tf.contrib.layers.flatten(x)
     x = tfl.dense(x, w * h, activation=None, name="enc_dense")
@@ -46,7 +46,7 @@ def decode_to_shape(inputs, shape, scope):
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
     x = inputs
     x = tf.contrib.layers.flatten(x)
-    x = tfl.dense(x, shape[2].value, activation=None, name="dec_dense")
+    x = tfl.dense(x, shape[2], activation=None, name="dec_dense")
     x = tf.expand_dims(x, axis=1)
     return x
 
diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 324adf70d..4762663a4 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -123,6 +123,21 @@ def get_input_if_exists(self, features, key, batch_size, num_frames):
       x = tf.zeros((batch_size, num_frames, 1, self.hparams.hidden_size))
     return common_video.swap_time_and_batch_axes(x)
 
+  def inject_additional_input(self, layer, inputs, scope, concatenate=True):
+    layer_shape = common_layers.shape_list(layer)
+    input_shape = common_layers.shape_list(inputs)
+    if concatenate:
+      emb = common_video.encode_to_shape(inputs, layer_shape, scope)
+      layer = tf.concat(values=[layer, emb], axis=-1)
+    else:
+      filters = layer_shape[-1]
+      input_reshaped = tf.reshape(inputs, [-1, 1, 1, input_shape[-1]])
+      input_mask = tf.layers.dense(input_reshaped, filters, name=scope)
+      zeros_mask = tf.zeros(layer_shape, dtype=tf.float32)
+      input_broad = input_mask + zeros_mask
+      layer *= input_broad
+    return layer
+
   def bottom_part_tower(self, input_image, input_reward, action, latent,
                         lstm_state, lstm_size, conv_size, concat_latent=False):
     """The bottom part of predictive towers.
@@ -186,28 +201,11 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
     enc2 = tfl.conv2d(hidden4, hidden4.get_shape()[3], [3, 3], strides=(2, 2),
                       padding="SAME", activation=tf.nn.relu, name="conv3")
 
-    # Pass in action if exists.
     if action is not None:
-      if self.hparams.concatenate_actions:
-        emb_action = common_video.encode_to_shape(
-            action, enc2.get_shape(), "action_enc")
-        enc2 = tf.concat(values=[enc2, emb_action], axis=3)
-      else:
-        action_shape = common_layers.shape_list(action)
-        enc2_shape = common_layers.shape_list(enc2)
-        filters = enc2_shape[-1]
-        action_reshaped = tf.reshape(action, [-1, 1, 1, action_shape[-1]])
-        action_mask = tf.layers.dense(action_reshaped, filters)
-        zeros_mask = tf.zeros(enc2_shape, dtype=tf.float32)
-        action_broad = action_mask + zeros_mask
-        enc2 *= action_broad
-
-    # Pass in reward if exists.
+      enc2 = self.inject_additional_input(
+          enc2, action, "action_enc", self.hparams.concatenate_actions)
     if input_reward is not None:
-      emb_reward = common_video.encode_to_shape(
-          input_reward, enc2.get_shape(), "reward_enc")
-      enc2 = tf.concat(values=[enc2, emb_reward], axis=3)
-
+      enc2 = self.inject_additional_input(enc2, input_reward, "reward_enc")
     if latent is not None and not concat_latent:
       with tf.control_dependencies([latent]):
         enc2 = tf.concat([enc2, latent], axis=3)
@@ -221,32 +219,38 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
     hidden5 = tile_and_concat(hidden5, latent, concat_latent=concat_latent)
     return hidden5, (enc0, enc1)
 
-  def reward_prediction(self, input_image, input_reward, action, latent):
+  def reward_prediction(self, input_images, input_reward, action, latent):
     """Builds a reward prediction network."""
-    del action
-    del latent
-
-    conv_size = self.tinyify([32, 32, 16, 4])
+    conv_size = self.tinyify([32, 32, 16, 8])
 
     with tf.variable_scope("reward_pred", reuse=tf.AUTO_REUSE):
-      x = input_image
+      x = tf.concat(input_images, axis=3)
       x = tfcl.batch_norm(x, updates_collections=None,
                           is_training=self.is_training, scope="reward_bn0")
       x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
-                     padding="SAME", activation=tf.nn.relu, name="reward_conv1")
+                     activation=tf.nn.relu, name="reward_conv1")
       x = tfcl.batch_norm(x, updates_collections=None,
                           is_training=self.is_training, scope="reward_bn1")
+
+      # Inject additional inputs
+      if action is not None:
+        x = self.inject_additional_input(
+            x, action, "action_enc", self.hparams.concatenate_actions)
+      if input_reward is not None:
+        x = self.inject_additional_input(x, input_reward, "reward_enc")
+      if latent is not None:
+        latent = tfl.flatten(latent)
+        latent = tf.expand_dims(latent, axis=1)
+        latent = tf.expand_dims(latent, axis=1)
+        x = self.inject_additional_input(x, latent, "latent_enc")
+
       x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(2, 2),
-                     padding="SAME", activation=tf.nn.relu, name="reward_conv2")
+                     activation=tf.nn.relu, name="reward_conv2")
       x = tfcl.batch_norm(x, updates_collections=None,
                           is_training=self.is_training, scope="reward_bn2")
       x = tfl.conv2d(x, conv_size[3], [3, 3], strides=(2, 2),
-                     padding="SAME", activation=tf.nn.relu, name="reward_conv3")
-
-      pred_reward = common_video.decode_to_shape(
-          x, input_reward.shape, "reward_dec")
-
-      return pred_reward
+                     activation=tf.nn.relu, name="reward_conv3")
+      return x
 
   def construct_predictive_tower(
       self, input_image, input_reward, action, lstm_state, latent,
@@ -389,36 +393,45 @@ def construct_model(self,
       ValueError: if more than 1 mask specified for DNA model.
     """
     context_frames = self.hparams.video_num_input_frames
+    buffer_size = self.hparams.reward_prediction_buffer_size
+    if buffer_size == 0:
+      buffer_size = context_frames
+    if buffer_size > context_frames:
+      raise ValueError("Buffer size is bigger than context frames %d %d." %
+                       (buffer_size, context_frames))
 
     batch_size = common_layers.shape_list(images)[1]
     ss_func = self.get_scheduled_sample_func(batch_size)
 
     def process_single_frame(prev_outputs, inputs):
       """Process a single frame of the video."""
-      cur_image, cur_reward, action = inputs
-      time_step, prev_image, prev_reward, lstm_states = prev_outputs
+      cur_image, input_reward, action = inputs
+      time_step, prev_image, prev_reward, frame_buf, lstm_states = prev_outputs
 
-      generated_items = [prev_image, prev_reward]
-      groundtruth_items = [cur_image, cur_reward]
+      generated_items = [prev_image]
+      groundtruth_items = [cur_image]
       done_warm_start = tf.greater(time_step, context_frames - 1)
-      input_image, input_reward = self.get_scheduled_sample_inputs(
+      input_image, = self.get_scheduled_sample_inputs(
           done_warm_start, groundtruth_items, generated_items, ss_func)
 
       # Prediction
       pred_image, lstm_states = self.construct_predictive_tower(
-          input_image, input_reward, action, lstm_states, latent)
+          input_image, None, action, lstm_states, latent)
 
       if self.hparams.reward_prediction:
         reward_input_image = pred_image
         if self.hparams.reward_prediction_stop_gradient:
           reward_input_image = tf.stop_gradient(reward_input_image)
-        pred_reward = self.reward_prediction(
-            reward_input_image, input_reward, action, latent)
+        with tf.control_dependencies([time_step]):
+          frame_buf = [reward_input_image] + frame_buf[:-1]
+        pred_reward = self.reward_prediction(frame_buf, None, action, latent)
+        pred_reward = common_video.decode_to_shape(
+            pred_reward, common_layers.shape_list(input_reward), "reward_dec")
       else:
-        pred_reward = input_reward
+        pred_reward = prev_reward
 
       time_step += 1
-      outputs = (time_step, pred_image, pred_reward, lstm_states)
+      outputs = (time_step, pred_image, pred_reward, frame_buf, lstm_states)
 
       return outputs
 
@@ -430,10 +443,12 @@ def process_single_frame(prev_outputs, inputs):
 
     # HACK: Do first step outside to initialize all the variables
     lstm_states = [None] * 7
+    frame_buffer = [tf.zeros_like(images[0])] * buffer_size
     inputs = images[0], rewards[0], actions[0]
     prev_outputs = (tf.constant(0),
                     tf.zeros_like(images[0]),
                     tf.zeros_like(rewards[0]),
+                    frame_buffer,
                     lstm_states)
 
     initializers = process_single_frame(prev_outputs, inputs)
@@ -579,6 +594,8 @@ def construct_model(self, images, actions, rewards):
       if self.hparams.reward_prediction:
         pred_reward = self.reward_prediction(
             pred_image, input_reward, action, latent)
+        pred_reward = common_video.decode_to_shape(
+            pred_reward, common_layers.shape_list(input_reward), "reward_dec")
       else:
         pred_reward = input_reward
 
diff --git a/tensor2tensor/models/research/next_frame_sv2p_params.py b/tensor2tensor/models/research/next_frame_sv2p_params.py
index 8ecb4d08b..5d6b5e9eb 100644
--- a/tensor2tensor/models/research/next_frame_sv2p_params.py
+++ b/tensor2tensor/models/research/next_frame_sv2p_params.py
@@ -36,6 +36,7 @@ def next_frame_sv2p():
   hparams.video_modality_loss_cutoff = 0.0
   hparams.add_hparam("reward_prediction", True)
   hparams.add_hparam("reward_prediction_stop_gradient", True)
+  hparams.add_hparam("reward_prediction_buffer_size", 0)
   hparams.add_hparam("model_options", "CDNA")
   hparams.add_hparam("num_masks", 10)
   hparams.add_hparam("multi_latent", False)

From 1de75bda4bd4c98ca50bcdbcf5e94b388bf9a044 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 22 Aug 2018 20:07:12 -0700
Subject: [PATCH 0649/2720] fix scoring crash on empty targets.

PiperOrigin-RevId: 209872868
---
 tensor2tensor/bin/t2t_decoder.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 399d2a4da..f8f844779 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -134,11 +134,9 @@ def score_file(filename):
     ckpt = ckpts.model_checkpoint_path
     saver.restore(sess, ckpt)
     # Run on each line.
-    results = []
-    lines = []
     with tf.gfile.Open(filename) as f:
-      text = f.read()
-      lines = [l.strip() for l in text.split("\n")]
+      lines = f.readlines()
+    results = []
     for line in lines:
       tab_split = line.split("\t")
       if len(tab_split) > 2:

From 99ea0d29b7fb312ebe66cd26b535bc1e9a257ac7 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Thu, 23 Aug 2018 16:41:50 +0200
Subject: [PATCH 0650/2720] ppo summary writer bugfix

---
 tensor2tensor/rl/rl_trainer_lib.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 4646ed6de..c73a3b599 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -48,11 +48,13 @@ def train(hparams, event_dir=None, model_dir=None,
     if event_dir:
       summary_writer = tf.summary.FileWriter(
           event_dir, graph=tf.get_default_graph(), flush_secs=60)
+    else:
+      summary_writer = None
+
     if model_dir:
       model_saver = tf.train.Saver(
           tf.global_variables(".*network_parameters.*"))
     else:
-      summary_writer = None
       model_saver = None
 
     # TODO(piotrmilos): This should be refactored, possibly with

From 6d67daef94435b1bf4f9b064e6e62768b8a175d9 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Thu, 23 Aug 2018 16:53:33 +0200
Subject: [PATCH 0651/2720] reward clipping wrapper

---
 tensor2tensor/data_generators/gym_problems.py |  3 ++-
 tensor2tensor/rl/envs/tf_atari_wrappers.py    | 15 +++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 214bb724a..995af8e13 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -48,7 +48,8 @@
 
 def standard_atari_env_spec(env):
   """Parameters of environment specification."""
-  standard_wrappers = [[tf_atari_wrappers.StackAndSkipWrapper, {"skip": 4}]]
+  standard_wrappers = [[tf_atari_wrappers.RewardClippingWrapper, {}],
+                       [tf_atari_wrappers.StackAndSkipWrapper, {"skip": 4}]]
   env_lambda = None
   if isinstance(env, str):
     env_lambda = lambda: gym.make(env)
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index f01020abc..bbcd54100 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -57,6 +57,21 @@ def _reset_non_empty(self, indices):
     with tf.control_dependencies([assign_op]):
       return tf.identity(new_values)
 
+class RewardClippingWrapper(WrapperBase):
+  """ Reward clipping wrapper.
+      The rewards are clipped to -1, 0, 1
+      This is a common strategy to ensure learning stability
+      of rl algorithms
+  """
+
+  def __init__(self, batch_env):
+    super(RewardClippingWrapper, self).__init__(batch_env)
+
+  def simulate(self, action):
+    reward, done = self._batch_env.simulate(action)
+    with tf.control_dependencies([reward, done]):
+      return tf.sign(reward), tf.identity(done)
+
 
 class MaxAndSkipWrapper(WrapperBase):
   """ Max and skip wrapper.

From 7e520596668941126df21db13f5f8ec5383fd3b8 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 23 Aug 2018 09:47:37 -0700
Subject: [PATCH 0652/2720] fixing the float issue.

PiperOrigin-RevId: 209951402
---
 tensor2tensor/models/research/next_frame_basic_stochastic.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/next_frame_basic_stochastic.py b/tensor2tensor/models/research/next_frame_basic_stochastic.py
index 4356c43c1..311cca896 100644
--- a/tensor2tensor/models/research/next_frame_basic_stochastic.py
+++ b/tensor2tensor/models/research/next_frame_basic_stochastic.py
@@ -37,8 +37,9 @@ class NextFrameBasicStochastic(
   def inject_latent(self, layer, features, filters):
     """Do nothing for deterministic model."""
     # Latent for stochastic model
-    full_video = tf.concat(
-        [features["inputs_raw"], features["targets_raw"]], axis=1)
+    input_frames = tf.to_float(features["inputs_raw"])
+    target_frames = tf.to_float(features["targets_raw"])
+    full_video = tf.concat([input_frames, target_frames], axis=1)
     latent_mean, latent_std = self.construct_latent_tower(
         full_video, time_axis=1)
     latent = common_video.get_gaussian_tensor(latent_mean, latent_std)

From 18e739868b7c867045a80626be33872ae019a76b Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 23 Aug 2018 10:50:19 -0700
Subject: [PATCH 0653/2720] Add a base class Text2TextRemotedir for easy adding
 of problems with data in a persistent remote directory (GCS bucket).

PiperOrigin-RevId: 209962813
---
 .../data_generators/text_problems.py          | 27 ++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 25227be0a..0258892ee 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -646,10 +646,35 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
     del data_dir
     is_training = dataset_split == problem.DatasetSplit.TRAIN
     files = self.TRAIN_FILES if is_training else self.EVAL_FILES
-    files = [os.path.join(tmp_dir, f) for f in files]
+    files = [os.path.join(self._tmp_dir_override or tmp_dir, f) for f in files]
     inputs_file, targets_file = files
     return text2text_txt_iterator(inputs_file, targets_file)
 
+  @property
+  def _tmp_dir_override(self):
+    return None
+
+
+class Text2TextRemotedir(Text2textTmpdir):
+  """Text2TextProblem from files in a remote directory.
+
+  SRC_REMOTE_DIR should be a remote directory, e.g. a GCS bucket (gs://...),
+  that contains the following files, 1 record per line:
+
+    * inputs.train.txt
+    * targets.train.txt
+    * inputs.eval.txt
+    * targets.eval.txt
+
+  """
+  # Override in subclass.
+  SRC_REMOTE_DIR = None
+
+  @property
+  def _tmp_dir_override(self):
+    assert self.SRC_REMOTE_DIR
+    return self.SRC_REMOTE_DIR
+
 
 @registry.register_problem
 class Text2textTmpdirTokens(Text2textTmpdir):

From ed7c781065d407b6e2fccbe72d0080b26e5d75a1 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 23 Aug 2018 11:07:22 -0700
Subject: [PATCH 0654/2720] Internal change

PiperOrigin-RevId: 209966032
---
 tensor2tensor/bin/build_vocab.py | 70 ++++++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)
 create mode 100644 tensor2tensor/bin/build_vocab.py

diff --git a/tensor2tensor/bin/build_vocab.py b/tensor2tensor/bin/build_vocab.py
new file mode 100644
index 000000000..f1235e28e
--- /dev/null
+++ b/tensor2tensor/bin/build_vocab.py
@@ -0,0 +1,70 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+r"""Build vocab for a subclass of Text2TextProblem.
+
+build_vocab \
+    --problem=program_search_algolisp \
+    --data_dir=~/t2t_data \
+    --tmp_dir=~/t2t_data/tmp
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("data_dir", "/tmp/t2t/data_dir",
+                    "Directory to place the generated vocabulary file in.")
+
+flags.DEFINE_string("tmp_dir", "/tmp/t2t/tmp_dir",
+                    "Temporary storage directory.")
+
+flags.DEFINE_string("problem", "",
+                    "Problem to generate the vocabulary file for.")
+
+flags.mark_flag_as_required("problem")
+
+
+def main(_):
+  problem = registry.problem(FLAGS.problem)
+
+  # We make the assumption that the problem is a subclass of Text2TextProblem.
+  assert isinstance(problem, text_problems.Text2TextProblem)
+
+  data_dir = os.path.expanduser(FLAGS.data_dir)
+  tmp_dir = os.path.expanduser(FLAGS.tmp_dir)
+
+  tf.gfile.MakeDirs(data_dir)
+  tf.gfile.MakeDirs(tmp_dir)
+
+  tf.logging.info("Saving vocabulary to data_dir: %s" % data_dir)
+
+  problem.get_or_create_vocab(data_dir, tmp_dir)
+
+  tf.logging.info("Saved vocabulary file: " +
+                  os.path.join(data_dir, problem.vocab_filename))
+
+
+if __name__ == "__main__":
+  tf.app.run()

From 3f3a5f9c65967894653e9b63b385aa85b4030579 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 23 Aug 2018 12:00:21 -0700
Subject: [PATCH 0655/2720] Revert moving to eager tests

PiperOrigin-RevId: 209975398
---
 tensor2tensor/layers/common_attention.py      |   2 +-
 tensor2tensor/layers/common_attention_test.py | 217 ++++----
 .../layers/common_image_attention_test.py     |   1 -
 tensor2tensor/layers/common_layers_test.py    | 468 ++++++++++--------
 tensor2tensor/layers/common_video_test.py     |  73 +--
 tensor2tensor/layers/discretization_test.py   |  91 ++--
 tensor2tensor/layers/modalities_test.py       |  34 +-
 7 files changed, 482 insertions(+), 404 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index ad9c612a4..ebf70535d 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -701,7 +701,7 @@ def add_positional_embedding_nd(x, max_length, name):
           name + "_%d" % i,
           shape,
           initializer=tf.random_normal_initializer(0, depth**-0.5))
-      var = var * depth**0.5
+      var *= depth**0.5
       x += tf.slice(var, start, size)
     return x
 
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 6a833713b..a14f43252 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -27,17 +27,17 @@
 import tensorflow as tf
 
 
-@tf.contrib.eager.run_all_tests_in_graph_and_eager_modes
 class CommonAttentionTest(parameterized.TestCase, tf.test.TestCase):
 
   def testAddPositionalEmbedding(self):
     x = np.random.rand(5, 3, 12)
-    y = common_attention.add_positional_embedding(
-        tf.constant(x, dtype=tf.float32),
-        max_length=4,
-        name="pos_embedding")
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_attention.add_positional_embedding(
+          tf.constant(x, dtype=tf.float32),
+          max_length=4,
+          name="pos_embedding")
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, (5, 3, 12))
 
   @parameterized.parameters(
@@ -47,22 +47,25 @@ def testAddPositionalEmbedding(self):
   )
   def testAddPositionalEmbeddingNd(self, input_shape):
     x = np.random.rand(*input_shape)
-    y = common_attention.add_positional_embedding_nd(
-        tf.constant(x, dtype=tf.float32),
-        max_length=5,
-        name="pos_embedding")
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_attention.add_positional_embedding_nd(
+          tf.constant(x, dtype=tf.float32),
+          max_length=5,
+          name="pos_embedding")
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, input_shape)
 
   def testDotProductAttention(self):
     x = np.random.rand(5, 7, 12, 32)
     y = np.random.rand(5, 7, 12, 32)
-    a = common_attention.dot_product_attention(
-        tf.constant(x, dtype=tf.float32),
-        tf.constant(y, dtype=tf.float32),
-        tf.constant(y, dtype=tf.float32), None)
-    res = self.evaluate(a)
+    with self.test_session() as session:
+      a = common_attention.dot_product_attention(
+          tf.constant(x, dtype=tf.float32),
+          tf.constant(y, dtype=tf.float32),
+          tf.constant(y, dtype=tf.float32), None)
+      session.run(tf.global_variables_initializer())
+      res = session.run(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
   @parameterized.named_parameters(
@@ -81,10 +84,11 @@ def testMaskedWithinBlockLocalAttention1D(self, batch, heads, length,
     v = tf.random_normal([batch, heads, length, depth_v])
     output = common_attention.masked_within_block_local_attention_1d(
         q, k, v, block_length=block_length)
-    if isinstance(batch, tf.Tensor):
-      batch, res = self.evaluate([batch, output])
-    else:
-      res = self.evaluate(output)
+    with self.test_session() as session:
+      if isinstance(batch, tf.Tensor):
+        batch, res = session.run([batch, output])
+      else:
+        res = session.run(output)
 
     self.assertEqual(res.shape, (batch, heads, length, depth_v))
 
@@ -104,10 +108,11 @@ def testMaskedLocalAttention1D(self, batch, heads, length, depth_k, depth_v,
     v = tf.random_normal([batch, heads, length, depth_v])
     output = common_attention.masked_local_attention_1d(
         q, k, v, block_length=block_length)
-    if isinstance(batch, tf.Tensor):
-      batch, res = self.evaluate([batch, output])
-    else:
-      res = self.evaluate(output)
+    with self.test_session() as session:
+      if isinstance(batch, tf.Tensor):
+        batch, res = session.run([batch, output])
+      else:
+        res = session.run(output)
 
     self.assertEqual(res.shape, (batch, heads, length, depth_v))
 
@@ -126,41 +131,43 @@ def testLocalUnmaskedAttention1D(self, batch, heads, length,
     v = tf.random_normal([batch, heads, length, depth_v])
     output = common_attention.local_attention_1d(
         q, k, v, block_length=block_length, filter_width=3)
-    if isinstance(batch, tf.Tensor):
-      batch, res = self.evaluate([batch, output])
-    else:
-      res = self.evaluate(output)
+    with self.test_session() as session:
+      if isinstance(batch, tf.Tensor):
+        batch, res = session.run([batch, output])
+      else:
+        res = session.run(output)
 
     self.assertEqual(res.shape, (batch, heads, length, depth_v))
 
   def testLocalUnmaskedAttention2D(self):
     x = np.random.rand(5, 4, 25, 25, 16)
     y = np.random.rand(5, 4, 25, 25, 16)
-    a = common_attention.local_attention_2d(
-        tf.constant(x, dtype=tf.float32),
-        tf.constant(y, dtype=tf.float32),
-        tf.constant(y, dtype=tf.float32),
-        query_shape=(4, 4),
-        memory_flange=(3, 3))
-    res = self.evaluate(a)
+    with self.test_session() as session:
+      a = common_attention.local_attention_2d(
+          tf.constant(x, dtype=tf.float32),
+          tf.constant(y, dtype=tf.float32),
+          tf.constant(y, dtype=tf.float32),
+          query_shape=(4, 4),
+          memory_flange=(3, 3))
+      session.run(tf.global_variables_initializer())
+      res = session.run(a)
     self.assertEqual(res.shape, (5, 4, 25, 25, 16))
 
   def testLocalUnmaskedAttention2DMatchingBlockLength(self):
     x = np.random.rand(5, 4, 25, 25, 16)
     y = np.random.rand(5, 4, 25, 25, 16)
-    a = common_attention.local_attention_2d(
-        tf.constant(x, dtype=tf.float32),
-        tf.constant(y, dtype=tf.float32),
-        tf.constant(y, dtype=tf.float32),
-        query_shape=(5, 5),
-        memory_flange=(3, 3))
-    res = self.evaluate(a)
+    with self.test_session() as session:
+      a = common_attention.local_attention_2d(
+          tf.constant(x, dtype=tf.float32),
+          tf.constant(y, dtype=tf.float32),
+          tf.constant(y, dtype=tf.float32),
+          query_shape=(5, 5),
+          memory_flange=(3, 3))
+      session.run(tf.global_variables_initializer())
+      res = session.run(a)
     self.assertEqual(res.shape, (5, 4, 25, 25, 16))
 
   def testMultiheadSelfAttentionMemoryEfficient(self):
-    if tf.executing_eagerly():
-      return  # don't run test in Eager mode
-
     num_heads = 4
     io_size = 16
     batch = 2
@@ -168,7 +175,7 @@ def testMultiheadSelfAttentionMemoryEfficient(self):
     head_size = 5
     x = np.random.rand(batch, length, io_size)
     dy = np.random.rand(batch, length, io_size)
-    with self.session() as session:
+    with self.test_session() as session:
       x = tf.to_float(x)
       dy = tf.to_float(dy)
       bias = common_attention.attention_bias_lower_triangle(length)
@@ -213,22 +220,25 @@ def test2dGatherAndScatterInvertibility(self):
     depth = 8
     query_shape = (2, 3)
     x = np.random.rand(batch_size, num_heads, height, width, depth)
-    x_indices = common_attention.gather_indices_2d(
-        x, query_shape, query_shape)
-    gathered_x = common_attention.gather_blocks_2d(x, x_indices)
-    x_shape = tf.constant([batch_size, num_heads, height, width, depth])
-    scattered_x = common_attention.scatter_blocks_2d(
-        gathered_x, x_indices, x_shape)
-    res = self.evaluate(scattered_x)
+    with self.test_session() as session:
+      x_indices = common_attention.gather_indices_2d(
+          x, query_shape, query_shape)
+      gathered_x = common_attention.gather_blocks_2d(x, x_indices)
+      x_shape = tf.constant([batch_size, num_heads, height, width, depth])
+      scattered_x = common_attention.scatter_blocks_2d(
+          gathered_x, x_indices, x_shape)
+      session.run(tf.global_variables_initializer())
+      res = session.run(scattered_x)
     self.assertAllClose(x, res)
 
   def test2dBlockRasterScanMask(self):
     """Testing the 2d block raster scan mask."""
     query_shape = (2, 3)
     memory_flange = (2, 1)
-    mask = common_attention.make_2d_block_raster_mask(
-        query_shape, memory_flange)
-    res = self.evaluate(mask)
+    with self.test_session() as session:
+      mask = common_attention.make_2d_block_raster_mask(
+          query_shape, memory_flange)
+      res = session.run(mask)
     correct_mask = np.array(
         [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0,
           1.0, 0.0, 1.0, 1.0, 1.0, 1.0],
@@ -275,10 +285,11 @@ def test2dGather(self):
                             y[1, 1, correct_indices[2]],
                             y[1, 1, correct_indices[3]]]]]
 
-    x_indices = common_attention.gather_indices_2d(
-        x, query_shape, query_shape)
-    gathered_x = common_attention.gather_blocks_2d(x, x_indices)
-    x_indices, gathered_x = self.evaluate([x_indices, gathered_x])
+    with self.test_session() as session:
+      x_indices = common_attention.gather_indices_2d(
+          x, query_shape, query_shape)
+      gathered_x = common_attention.gather_blocks_2d(x, x_indices)
+      x_indices, gathered_x = session.run([x_indices, gathered_x])
     self.assertAllEqual(correct_indices, x_indices)
     self.assertAllClose(correct_gathered_x, gathered_x)
 
@@ -347,14 +358,16 @@ def testGetMemoryRegion(self):
                           y[1, 1, [12, 13, 14, 18, 19, 20]],
                           y[1, 1, [15, 16, 17, 21, 22, 23]]]]]
     correct_x_center = np.array(correct_x_center)
-    x_indices = common_attention.gather_indices_2d(
-        x, query_shape, query_shape)
-    x_flange, x_center = common_attention.get_memory_region(
-        tf.constant(x, dtype=tf.float32),
-        query_shape,
-        memory_flange,
-        x_indices)
-    [x_flange, x_center] = self.evaluate([x_flange, x_center])
+    with self.test_session() as session:
+      x_indices = common_attention.gather_indices_2d(
+          x, query_shape, query_shape)
+      x_flange, x_center = common_attention.get_memory_region(
+          tf.constant(x, dtype=tf.float32),
+          query_shape,
+          memory_flange,
+          x_indices)
+      session.run(tf.global_variables_initializer())
+      [x_flange, x_center] = session.run([x_flange, x_center])
     self.assertAllClose(correct_x_flange, x_flange)
     self.assertAllClose(correct_x_center, x_center)
 
@@ -414,38 +427,42 @@ def testGetShiftedCenterBlocks(self):
                                             y[1, 1, [15, 16, 17, 21, 22]]),
                                            axis=0)]]]
     correct_gathered_x = np.array(correct_gathered_x)
-    x_indices = common_attention.gather_indices_2d(
-        x, query_shape, query_shape)
-    gathered_x = common_attention.get_shifted_center_blocks(
-        tf.constant(x, dtype=tf.float32),
-        x_indices)
-    x_indices, gathered_x = self.evaluate([x_indices, gathered_x])
+    with self.test_session() as session:
+      x_indices = common_attention.gather_indices_2d(
+          x, query_shape, query_shape)
+      gathered_x = common_attention.get_shifted_center_blocks(
+          tf.constant(x, dtype=tf.float32),
+          x_indices)
+      session.run(tf.global_variables_initializer())
+      x_indices, gathered_x = session.run([x_indices, gathered_x])
     self.assertAllClose(correct_gathered_x, gathered_x)
 
   def testDotProductAttentionRelative(self):
     x = np.random.rand(5, 7, 12, 32)
     y = np.random.rand(5, 7, 12, 32)
-    a = common_attention.dot_product_attention_relative(
-        tf.constant(x, dtype=tf.float32),
-        tf.constant(y, dtype=tf.float32),
-        tf.constant(y, dtype=tf.float32),
-        None,
-        max_relative_position=3)
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(a)
+    with self.test_session() as session:
+      a = common_attention.dot_product_attention_relative(
+          tf.constant(x, dtype=tf.float32),
+          tf.constant(y, dtype=tf.float32),
+          tf.constant(y, dtype=tf.float32),
+          None,
+          max_relative_position=3)
+      session.run(tf.global_variables_initializer())
+      res = session.run(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
   def testDotProductUnMaskedAttentionRelativeV2(self):
     x = np.random.rand(5, 7, 12, 32)
     y = np.random.rand(5, 7, 12, 32)
-    a = common_attention.dot_product_unmasked_self_attention_relative_v2(
-        tf.constant(x, dtype=tf.float32),
-        tf.constant(y, dtype=tf.float32),
-        tf.constant(y, dtype=tf.float32),
-        None,
-        35)
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(a)
+    with self.test_session() as session:
+      a = common_attention.dot_product_unmasked_self_attention_relative_v2(
+          tf.constant(x, dtype=tf.float32),
+          tf.constant(y, dtype=tf.float32),
+          tf.constant(y, dtype=tf.float32),
+          None,
+          35)
+      session.run(tf.global_variables_initializer())
+      res = session.run(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
   def testBiasBatchCoordinates(self):
@@ -469,7 +486,13 @@ def testBiasBatchCoordinates(self):
     ], np.float32) * -1e9
 
     bias = common_attention.attention_bias_coordinates(q, k)
-    self.assertAllClose(self.evaluate(bias), ground_truth)
+
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      self.assertAllClose(
+          bias.eval(),
+          ground_truth,
+      )
 
   def testBiasFuture(self):
     """Testing the sequence order mask."""
@@ -492,7 +515,13 @@ def testBiasFuture(self):
     ], np.float32) * -1e9
 
     bias = common_attention.attention_bias_future(q, k)
-    self.assertAllClose(self.evaluate(bias), ground_truth)
+
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      self.assertAllClose(
+          bias.eval(),
+          ground_truth,
+      )
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index c0845ef10..6d57e7413 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -24,7 +24,6 @@
 import tensorflow as tf
 
 
-@tf.contrib.eager.run_all_tests_in_graph_and_eager_modes
 class CommonImageAttentionTest(parameterized.TestCase, tf.test.TestCase):
 
   @parameterized.parameters(
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index da5647c09..d4a6df588 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -26,7 +26,6 @@
 import tensorflow as tf
 
 
-@tf.contrib.eager.run_all_tests_in_graph_and_eager_modes
 class CommonLayersTest(parameterized.TestCase, tf.test.TestCase):
 
   def testIndexLastDimWithIndices(self):
@@ -36,134 +35,149 @@ def testIndexLastDimWithIndices(self):
     x_idx = common_layers.index_last_dim_with_indices(x, indices)
 
     expected = np.array([4., 6.])
-    self.assertAllEqual(expected, self.evaluate(x_idx))
+    with self.test_session() as sess:
+      self.assertAllEqual(expected, sess.run(x_idx))
 
   def testSaturatingSigmoid(self):
     x = np.array([-120.0, -100.0, 0.0, 100.0, 120.0], dtype=np.float32)
-    y = common_layers.saturating_sigmoid(tf.constant(x))
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_layers.saturating_sigmoid(tf.constant(x))
+      res = session.run(y)
     self.assertAllClose(res, [0.0, 0.0, 0.5, 1.0, 1.0])
 
   def testFlatten4D3D(self):
     x = np.random.random_integers(1, high=8, size=(3, 5, 2))
-    y = common_layers.flatten4d3d(common_layers.embedding(x, 10, 7))
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_layers.flatten4d3d(common_layers.embedding(x, 10, 7))
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, (3, 5 * 2, 7))
 
   def testEmbedding(self):
     x = np.random.random_integers(1, high=8, size=(3, 5))
-    y = common_layers.embedding(x, 10, 16)
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_layers.embedding(x, 10, 16)
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, (3, 5, 16))
 
   def testShakeShake(self):
-    if tf.executing_eagerly():
-      return  # don't run test in Eager mode
-
     x = np.random.rand(5, 7)
-    with self.session() as session:
+    with self.test_session() as session:
       x = tf.constant(x, dtype=tf.float32)
       y = common_layers.shakeshake([x, x, x, x, x])
+      session.run(tf.global_variables_initializer())
       inp, res = session.run([x, y])
     self.assertAllClose(res, inp)
 
   def testConv(self):
     x = np.random.rand(5, 7, 1, 11)
-    y = common_layers.conv(tf.constant(x, dtype=tf.float32), 13, (3, 1))
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_layers.conv(tf.constant(x, dtype=tf.float32), 13, (3, 1))
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, (5, 5, 1, 13))
 
   def testConv1d(self):
     x = np.random.rand(5, 7, 11)
-    y = common_layers.conv1d(tf.constant(x, dtype=tf.float32), 13, 1)
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_layers.conv1d(tf.constant(x, dtype=tf.float32), 13, 1)
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, (5, 7, 13))
 
   def testSeparableConv(self):
     x = np.random.rand(5, 7, 1, 11)
-    y = common_layers.separable_conv(
-        tf.constant(x, dtype=tf.float32), 13, (3, 1))
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_layers.separable_conv(
+          tf.constant(x, dtype=tf.float32), 13, (3, 1))
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, (5, 5, 1, 13))
 
   def testSubSeparableConv(self):
     for sep in [0, 1, 2, 4]:
       x = np.random.rand(5, 7, 1, 12)
-      with tf.variable_scope("sep_%d" % sep):
-        y = common_layers.subseparable_conv(
-            tf.constant(x, dtype=tf.float32), 16, (3, 1), separability=sep)
-      self.evaluate(tf.global_variables_initializer())
-      res = self.evaluate(y)
+      with self.test_session() as session:
+        with tf.variable_scope("sep_%d" % sep):
+          y = common_layers.subseparable_conv(
+              tf.constant(x, dtype=tf.float32), 16, (3, 1), separability=sep)
+        session.run(tf.global_variables_initializer())
+        res = session.run(y)
       self.assertEqual(res.shape, (5, 5, 1, 16))
 
   def testConvBlock(self):
     x = np.random.rand(5, 7, 1, 11)
-    y = common_layers.conv_block(
-        tf.constant(x, dtype=tf.float32),
-        13, [(1, (3, 3)), (1, (3, 3))],
-        padding="SAME",
-        normalizer_fn=common_layers.noam_norm)
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_layers.conv_block(
+          tf.constant(x, dtype=tf.float32),
+          13, [(1, (3, 3)), (1, (3, 3))],
+          padding="SAME",
+          normalizer_fn=common_layers.noam_norm)
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, (5, 7, 1, 13))
 
   def testSeparableConvBlock(self):
     x = np.random.rand(5, 7, 1, 11)
-    y = common_layers.separable_conv_block(
-        tf.constant(x, dtype=tf.float32),
-        13, [(1, (3, 3)), (1, (3, 3))],
-        padding="SAME")
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_layers.separable_conv_block(
+          tf.constant(x, dtype=tf.float32),
+          13, [(1, (3, 3)), (1, (3, 3))],
+          padding="SAME")
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, (5, 7, 1, 13))
 
   def testSubSeparableConvBlock(self):
     for sep in [0, 1, 2, 4]:
       x = np.random.rand(5, 7, 1, 12)
-      with tf.variable_scope("sep_%d" % sep):
-        y = common_layers.subseparable_conv_block(
-            tf.constant(x, dtype=tf.float32),
-            16, [(1, (3, 3)), (1, (3, 3))],
-            padding="SAME",
-            separability=sep)
-      self.evaluate(tf.global_variables_initializer())
-      res = self.evaluate(y)
+      with self.test_session() as session:
+        with tf.variable_scope("sep_%d" % sep):
+          y = common_layers.subseparable_conv_block(
+              tf.constant(x, dtype=tf.float32),
+              16, [(1, (3, 3)), (1, (3, 3))],
+              padding="SAME",
+              separability=sep)
+        session.run(tf.global_variables_initializer())
+        res = session.run(y)
       self.assertEqual(res.shape, (5, 7, 1, 16))
 
   def testPool(self):
     x = np.random.rand(5, 8, 1, 11)
-    y = common_layers.pool(
-        tf.constant(x, dtype=tf.float32), (2, 2), "AVG", "SAME")
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_layers.pool(
+          tf.constant(x, dtype=tf.float32), (2, 2), "AVG", "SAME")
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, (5, 8, 1, 11))
 
   def testConvBlockDownsample(self):
     x = np.random.rand(5, 7, 1, 11)
-    y = common_layers.conv_block_downsample(
-        tf.constant(x, dtype=tf.float32), (3, 1), (2, 1), "SAME")
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_layers.conv_block_downsample(
+          tf.constant(x, dtype=tf.float32), (3, 1), (2, 1), "SAME")
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, (5, 4, 1, 27))
 
   def testSimpleAttention(self):
     x = np.random.rand(5, 7, 1, 11)
     y = np.random.rand(5, 9, 1, 11)
-    a = common_layers.simple_attention(
-        tf.constant(x, dtype=tf.float32), tf.constant(y, dtype=tf.float32))
-    res = self.evaluate(a)
+    with self.test_session() as session:
+      a = common_layers.simple_attention(
+          tf.constant(x, dtype=tf.float32), tf.constant(y, dtype=tf.float32))
+      session.run(tf.global_variables_initializer())
+      res = session.run(a)
     self.assertEqual(res.shape, (5, 7, 1, 11))
 
   def testGetTimingSignal(self):
     length = 7
     num_timescales = 10
-    a = common_layers.get_timing_signal(length, num_timescales=num_timescales)
-    res = self.evaluate(a)
+    with self.test_session() as session:
+      a = common_layers.get_timing_signal(length, num_timescales=num_timescales)
+      session.run(tf.global_variables_initializer())
+      res = session.run(a)
     self.assertEqual(res.shape, (length, 2 * num_timescales))
 
   def testAddTimingSignal(self):
@@ -172,8 +186,10 @@ def testAddTimingSignal(self):
     height = 3
     depth = 35
     x = np.random.rand(batch, length, height, depth)
-    a = common_layers.add_timing_signal(tf.constant(x, dtype=tf.float32))
-    res = self.evaluate(a)
+    with self.test_session() as session:
+      a = common_layers.add_timing_signal(tf.constant(x, dtype=tf.float32))
+      session.run(tf.global_variables_initializer())
+      res = session.run(a)
     self.assertEqual(res.shape, (batch, length, height, depth))
 
   def testAttention1D(self):
@@ -188,42 +204,42 @@ def testAttention1D(self):
     source = np.random.rand(batch, source_length, source_depth)
     target = np.random.rand(batch, target_length, target_depth)
     mask = np.random.rand(batch, target_length, source_length)
-    a = common_layers.attention_1d_v0(
-        tf.constant(source, dtype=tf.float32),
-        tf.constant(target, dtype=tf.float32), attention_size, output_size,
-        num_heads, tf.constant(mask, dtype=tf.float32))
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(a)
+    with self.test_session() as session:
+      a = common_layers.attention_1d_v0(
+          tf.constant(source, dtype=tf.float32),
+          tf.constant(target, dtype=tf.float32), attention_size, output_size,
+          num_heads, tf.constant(mask, dtype=tf.float32))
+      session.run(tf.global_variables_initializer())
+      res = session.run(a)
     self.assertEqual(res.shape, (batch, target_length, output_size))
 
   def testMultiscaleConvSum(self):
     x = np.random.rand(5, 9, 1, 11)
-    y = common_layers.multiscale_conv_sum(
-        tf.constant(x, dtype=tf.float32),
-        13, [((1, 1), (5, 5)), ((2, 2), (3, 3))],
-        "AVG",
-        padding="SAME")
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_layers.multiscale_conv_sum(
+          tf.constant(x, dtype=tf.float32),
+          13, [((1, 1), (5, 5)), ((2, 2), (3, 3))],
+          "AVG",
+          padding="SAME")
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, (5, 9, 1, 13))
 
   def testConvGRU(self):
     x = np.random.rand(5, 7, 3, 11)
-    y = common_layers.conv_gru(tf.constant(x, dtype=tf.float32), (1, 3), 11)
-    z = common_layers.conv_gru(
-        tf.constant(x, dtype=tf.float32), (1, 3), 11, padding="LEFT")
-    self.evaluate(tf.global_variables_initializer())
-    res1 = self.evaluate(y)
-    res2 = self.evaluate(z)
+    with self.test_session() as session:
+      y = common_layers.conv_gru(tf.constant(x, dtype=tf.float32), (1, 3), 11)
+      z = common_layers.conv_gru(
+          tf.constant(x, dtype=tf.float32), (1, 3), 11, padding="LEFT")
+      session.run(tf.global_variables_initializer())
+      res1 = session.run(y)
+      res2 = session.run(z)
     self.assertEqual(res1.shape, (5, 7, 3, 11))
     self.assertEqual(res2.shape, (5, 7, 3, 11))
 
   def testSRU(self):
-    if tf.executing_eagerly():
-      return  # don't run test in Eager mode
-
     x = np.random.rand(5, 7, 3, 11)
-    with self.session() as session:
+    with self.test_session() as session:
       y = common_layers.sru(tf.constant(x, dtype=tf.float32))
       session.run(tf.global_variables_initializer())
       res = session.run(y)
@@ -231,36 +247,40 @@ def testSRU(self):
 
   def testLayerNorm(self):
     x = np.random.rand(5, 7, 11)
-    y = common_layers.layer_norm(tf.constant(x, dtype=tf.float32), 11)
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_layers.layer_norm(tf.constant(x, dtype=tf.float32), 11)
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, (5, 7, 11))
 
   def testGroupNorm(self):
     x = np.random.rand(5, 7, 3, 16)
-    y = common_layers.group_norm(tf.constant(x, dtype=tf.float32))
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_layers.group_norm(tf.constant(x, dtype=tf.float32))
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, (5, 7, 3, 16))
 
   def testConvLSTM(self):
     x = np.random.rand(5, 7, 11, 13)
-    y = common_layers.conv_lstm(tf.constant(x, dtype=tf.float32), (1, 3), 13)
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
+    with self.test_session() as session:
+      y = common_layers.conv_lstm(tf.constant(x, dtype=tf.float32), (1, 3), 13)
+      session.run(tf.global_variables_initializer())
+      res = session.run(y)
     self.assertEqual(res.shape, (5, 7, 11, 13))
 
   def testPadToSameLength(self):
     x1 = np.random.rand(5, 7, 11)
     x2 = np.random.rand(5, 9, 11)
-    a, b = common_layers.pad_to_same_length(
-        tf.constant(x1, dtype=tf.float32), tf.constant(x2, dtype=tf.float32))
-    c, d = common_layers.pad_to_same_length(
-        tf.constant(x1, dtype=tf.float32),
-        tf.constant(x2, dtype=tf.float32),
-        final_length_divisible_by=4)
-    res1, res2 = self.evaluate([a, b])
-    res1a, res2a = self.evaluate([c, d])
+    with self.test_session() as session:
+      a, b = common_layers.pad_to_same_length(
+          tf.constant(x1, dtype=tf.float32), tf.constant(x2, dtype=tf.float32))
+      c, d = common_layers.pad_to_same_length(
+          tf.constant(x1, dtype=tf.float32),
+          tf.constant(x2, dtype=tf.float32),
+          final_length_divisible_by=4)
+      res1, res2 = session.run([a, b])
+      res1a, res2a = session.run([c, d])
     self.assertEqual(res1.shape, (5, 9, 11))
     self.assertEqual(res2.shape, (5, 9, 11))
     self.assertEqual(res1a.shape, (5, 12, 11))
@@ -271,56 +291,63 @@ def testShiftLeft(self):
     x1[:, 0, :] = np.ones_like(x1[:, 0, :])
     expected = np.zeros((5, 7, 1, 11))
     expected[:, 1, :] = np.ones_like(expected[:, 1, :])
-    a = common_layers.shift_right(tf.constant(x1, dtype=tf.float32))
-    actual = self.evaluate(a)
+    with self.test_session() as session:
+      a = common_layers.shift_right(tf.constant(x1, dtype=tf.float32))
+      actual = session.run(a)
     self.assertAllEqual(actual, expected)
 
   def testConvStride2MultiStep(self):
     x1 = np.random.rand(5, 32, 16, 11)
-    a = common_layers.conv_stride2_multistep(
-        tf.constant(x1, dtype=tf.float32), 4, 16)
-    self.evaluate(tf.global_variables_initializer())
-    actual = self.evaluate(a[0])
+    with self.test_session() as session:
+      a = common_layers.conv_stride2_multistep(
+          tf.constant(x1, dtype=tf.float32), 4, 16)
+      session.run(tf.global_variables_initializer())
+      actual = session.run(a[0])
     self.assertEqual(actual.shape, (5, 2, 1, 16))
 
   def testDeconvStride2MultiStep(self):
     x1 = np.random.rand(5, 2, 1, 11)
-    a = common_layers.deconv_stride2_multistep(
-        tf.constant(x1, dtype=tf.float32), 4, 16)
-    self.evaluate(tf.global_variables_initializer())
-    actual = self.evaluate(a)
+    with self.test_session() as session:
+      a = common_layers.deconv_stride2_multistep(
+          tf.constant(x1, dtype=tf.float32), 4, 16)
+      session.run(tf.global_variables_initializer())
+      actual = session.run(a)
     self.assertEqual(actual.shape, (5, 32, 1, 16))
 
   def testApplyNormLayer(self):
-    x1 = np.random.rand(5, 2, 1, 11)
-    x2 = common_layers.apply_norm(
-        tf.constant(x1, dtype=tf.float32), "layer", depth=11, epsilon=1e-6)
-    self.evaluate(tf.global_variables_initializer())
-    actual = self.evaluate(x2)
+    with self.test_session() as session:
+      x1 = np.random.rand(5, 2, 1, 11)
+      x2 = common_layers.apply_norm(
+          tf.constant(x1, dtype=tf.float32), "layer", depth=11, epsilon=1e-6)
+      session.run(tf.global_variables_initializer())
+      actual = session.run(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
   def testApplyNormNoam(self):
-    x1 = np.random.rand(5, 2, 1, 11)
-    x2 = common_layers.apply_norm(
-        tf.constant(x1, dtype=tf.float32), "noam", depth=11, epsilon=1e-6)
-    self.evaluate(tf.global_variables_initializer())
-    actual = self.evaluate(x2)
+    with self.test_session() as session:
+      x1 = np.random.rand(5, 2, 1, 11)
+      x2 = common_layers.apply_norm(
+          tf.constant(x1, dtype=tf.float32), "noam", depth=11, epsilon=1e-6)
+      session.run(tf.global_variables_initializer())
+      actual = session.run(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
   def testApplyNormBatch(self):
-    x1 = np.random.rand(5, 2, 1, 11)
-    x2 = common_layers.apply_norm(
-        tf.constant(x1, dtype=tf.float32), "batch", depth=11, epsilon=1e-6)
-    self.evaluate(tf.global_variables_initializer())
-    actual = self.evaluate(x2)
+    with self.test_session() as session:
+      x1 = np.random.rand(5, 2, 1, 11)
+      x2 = common_layers.apply_norm(
+          tf.constant(x1, dtype=tf.float32), "batch", depth=11, epsilon=1e-6)
+      session.run(tf.global_variables_initializer())
+      actual = session.run(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
   def testApplyNormNone(self):
-    x1 = np.random.rand(5, 2, 1, 11)
-    x2 = common_layers.apply_norm(
-        tf.constant(x1, dtype=tf.float32), "none", depth=11, epsilon=1e-6)
-    self.evaluate(tf.global_variables_initializer())
-    actual = self.evaluate(x2)
+    with self.test_session() as session:
+      x1 = np.random.rand(5, 2, 1, 11)
+      x2 = common_layers.apply_norm(
+          tf.constant(x1, dtype=tf.float32), "none", depth=11, epsilon=1e-6)
+      session.run(tf.global_variables_initializer())
+      actual = session.run(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
     self.assertAllClose(actual, x1, atol=1e-03)
 
@@ -329,88 +356,93 @@ def testGlobalPool1d(self):
     no_mask = np.ones((5, 4))
     full_mask = np.zeros((5, 4))
 
-    x1_ = tf.Variable(x1, dtype=tf.float32)
-    no_mask_ = tf.Variable(no_mask, dtype=tf.float32)
-    full_mask_ = tf.Variable(full_mask, dtype=tf.float32)
+    with self.test_session() as session:
+      x1_ = tf.Variable(x1, dtype=tf.float32)
+      no_mask_ = tf.Variable(no_mask, dtype=tf.float32)
+      full_mask_ = tf.Variable(full_mask, dtype=tf.float32)
 
-    none_mask_max = common_layers.global_pool_1d(x1_)
-    no_mask_max = common_layers.global_pool_1d(x1_, mask=no_mask_)
-    result1 = tf.reduce_sum(none_mask_max - no_mask_max)
+      none_mask_max = common_layers.global_pool_1d(x1_)
+      no_mask_max = common_layers.global_pool_1d(x1_, mask=no_mask_)
+      result1 = tf.reduce_sum(none_mask_max - no_mask_max)
 
-    full_mask_max = common_layers.global_pool_1d(x1_, mask=full_mask_)
-    result2 = tf.reduce_sum(full_mask_max)
+      full_mask_max = common_layers.global_pool_1d(x1_, mask=full_mask_)
+      result2 = tf.reduce_sum(full_mask_max)
 
-    none_mask_avr = common_layers.global_pool_1d(x1_, "AVR")
-    no_mask_avr = common_layers.global_pool_1d(x1_, "AVR", no_mask_)
-    result3 = tf.reduce_sum(none_mask_avr - no_mask_avr)
+      none_mask_avr = common_layers.global_pool_1d(x1_, "AVR")
+      no_mask_avr = common_layers.global_pool_1d(x1_, "AVR", no_mask_)
+      result3 = tf.reduce_sum(none_mask_avr - no_mask_avr)
 
-    full_mask_avr = common_layers.global_pool_1d(x1_, "AVR", full_mask_)
-    result4 = tf.reduce_sum(full_mask_avr)
+      full_mask_avr = common_layers.global_pool_1d(x1_, "AVR", full_mask_)
+      result4 = tf.reduce_sum(full_mask_avr)
 
-    self.evaluate(tf.global_variables_initializer())
-    actual = self.evaluate([result1, result2, result3, result4])
+      session.run(tf.global_variables_initializer())
+      actual = session.run([result1, result2, result3, result4])
     self.assertAllEqual(actual[:3], [0.0, 0.0, 0.0])
 
   def testLinearSetLayer(self):
     x1 = np.random.rand(5, 4, 11)
     cont = np.random.rand(5, 13)
-    x1_ = tf.Variable(x1, dtype=tf.float32)
-    cont_ = tf.Variable(cont, dtype=tf.float32)
+    with self.test_session() as session:
+      x1_ = tf.Variable(x1, dtype=tf.float32)
+      cont_ = tf.Variable(cont, dtype=tf.float32)
 
-    simple_ff = common_layers.linear_set_layer(32, x1_)
-    cont_ff = common_layers.linear_set_layer(32, x1_, context=cont_)
+      simple_ff = common_layers.linear_set_layer(32, x1_)
+      cont_ff = common_layers.linear_set_layer(32, x1_, context=cont_)
 
-    self.evaluate(tf.global_variables_initializer())
-    actual = self.evaluate([simple_ff, cont_ff])
+      session.run(tf.global_variables_initializer())
+      actual = session.run([simple_ff, cont_ff])
     self.assertEqual(actual[0].shape, (5, 4, 32))
     self.assertEqual(actual[1].shape, (5, 4, 32))
 
   def testRavanbakhshSetLayer(self):
     x1 = np.random.rand(5, 4, 11)
-    x1_ = tf.Variable(x1, dtype=tf.float32)
-    layer = common_layers.ravanbakhsh_set_layer(32, x1_)
-    self.evaluate(tf.global_variables_initializer())
-    actual = self.evaluate(layer)
+    with self.test_session() as session:
+      x1_ = tf.Variable(x1, dtype=tf.float32)
+      layer = common_layers.ravanbakhsh_set_layer(32, x1_)
+      session.run(tf.global_variables_initializer())
+      actual = session.run(layer)
     self.assertEqual(actual.shape, (5, 4, 32))
 
   def testBReLU(self):
-    x = np.random.rand(5, 2, 1, 12)
-    y = common_layers.brelu(tf.constant(x, dtype=tf.float32))
-    actual = self.evaluate(y)
+    with self.test_session() as session:
+      x = np.random.rand(5, 2, 1, 12)
+      y = common_layers.brelu(tf.constant(x, dtype=tf.float32))
+      actual = session.run(y)
     self.assertEqual(actual.shape, (5, 2, 1, 12))
 
   def testBELU(self):
-    x = np.random.rand(5, 2, 1, 12)
-    y = common_layers.belu(tf.constant(x, dtype=tf.float32))
-    actual = self.evaluate(y)
+    with self.test_session() as session:
+      x = np.random.rand(5, 2, 1, 12)
+      y = common_layers.belu(tf.constant(x, dtype=tf.float32))
+      actual = session.run(y)
     self.assertEqual(actual.shape, (5, 2, 1, 12))
 
   def testNAC(self):
-    x = np.random.rand(5, 2, 1, 12)
-    y = common_layers.nac(tf.constant(x, dtype=tf.float32), 14)
-    self.evaluate(tf.global_variables_initializer())
-    actual = self.evaluate(y)
+    with self.test_session() as session:
+      x = np.random.rand(5, 2, 1, 12)
+      y = common_layers.nac(tf.constant(x, dtype=tf.float32), 14)
+      session.run(tf.global_variables_initializer())
+      actual = session.run(y)
     self.assertEqual(actual.shape, (5, 2, 1, 14))
 
   def testNALU(self):
-    x = np.random.rand(5, 2, 1, 12)
-    y = common_layers.nalu(tf.constant(x, dtype=tf.float32), 14)
-    self.evaluate(tf.global_variables_initializer())
-    actual = self.evaluate(y)
+    with self.test_session() as session:
+      x = np.random.rand(5, 2, 1, 12)
+      y = common_layers.nalu(tf.constant(x, dtype=tf.float32), 14)
+      session.run(tf.global_variables_initializer())
+      actual = session.run(y)
     self.assertEqual(actual.shape, (5, 2, 1, 14))
 
   def testNALUzeros(self):
-    x = np.random.rand(5, 2, 1, 12)
-    y = common_layers.nalu(tf.zeros_like(x, dtype=tf.float32), 14)
-    self.evaluate(tf.global_variables_initializer())
-    actual = self.evaluate(y)
+    with self.test_session() as session:
+      x = np.random.rand(5, 2, 1, 12)
+      y = common_layers.nalu(tf.zeros_like(x, dtype=tf.float32), 14)
+      session.run(tf.global_variables_initializer())
+      actual = session.run(y)
     self.assertTrue(np.all(np.isfinite(actual)))
     self.assertEqual(actual.shape, (5, 2, 1, 14))
 
   def testPaddingCrossEntropyFactored(self):
-    if tf.executing_eagerly():
-      return  # don't run test in Eager mode
-
     vocab_size = 19
     rows = 5
     cols = 4
@@ -419,7 +451,7 @@ def testPaddingCrossEntropyFactored(self):
     features = np.random.rand(rows, cols, depth)
     weights = np.random.rand(vocab_size, depth)
     labels = np.random.randint(0, vocab_size - 1, size=(rows, cols))
-    with self.session() as session:
+    with self.test_session() as session:
       features = tf.to_float(features)
       weights = tf.to_float(weights)
       labels = tf.to_int32(labels)
@@ -444,9 +476,6 @@ def testPaddingCrossEntropyFactored(self):
     self.assertAllClose(den, den_f)
 
   def testPaddingCrossEntropyFactoredGrad(self):
-    if tf.executing_eagerly():
-      return  # don't run test in Eager mode
-
     vocab_size = 19
     rows = 5
     cols = 4
@@ -455,7 +484,7 @@ def testPaddingCrossEntropyFactoredGrad(self):
     features = np.random.rand(rows, cols, depth)
     weights = np.random.rand(vocab_size, depth)
     labels = np.random.randint(0, vocab_size - 1, size=(rows, cols))
-    with self.session() as session:
+    with self.test_session() as session:
       features = tf.to_float(features)
       weights = tf.to_float(weights)
       labels = tf.to_int32(labels)
@@ -504,8 +533,9 @@ def testDmlLoss(self, batch, height, width, num_mixtures, reduce_sum):
     if reduce_sum:
       expected_loss = tf.reduce_mean(expected_loss)
 
-    actual_loss_val, expected_loss_val = self.evaluate(
-        [actual_loss, expected_loss])
+    with self.test_session() as sess:
+      actual_loss_val, expected_loss_val = sess.run(
+          [actual_loss, expected_loss])
     self.assertAllClose(actual_loss_val, expected_loss_val)
 
   def testDiscretizedMixLogisticLoss(self):
@@ -540,8 +570,9 @@ def testDiscretizedMixLogisticLoss(self):
 
     actual_loss = common_layers.discretized_mix_logistic_loss(
         pred=pred, labels=labels)
-    actual_loss_val, expected_loss_val = self.evaluate(
-        [actual_loss, expected_loss])
+    with self.test_session() as session:
+      actual_loss_val, expected_loss_val = session.run(
+          [actual_loss, expected_loss])
     self.assertAllClose(actual_loss_val, expected_loss_val, rtol=1e-5)
 
   def testSampleFromDiscretizedMixLogistic(self):
@@ -565,8 +596,9 @@ def testSampleFromDiscretizedMixLogistic(self):
 
     actual_sample = common_layers.sample_from_discretized_mix_logistic(
         pred, seed=seed)
-    actual_sample_val, expected_sample_val = self.evaluate(
-        [actual_sample, expected_sample])
+    with self.test_session() as session:
+      actual_sample_val, expected_sample_val = session.run(
+          [actual_sample, expected_sample])
     # Use a low tolerance: samples numerically differ, as the actual
     # implementation clips log-scales so they always contribute to sampling.
     self.assertAllClose(actual_sample_val, expected_sample_val, atol=1e-2)
@@ -575,24 +607,22 @@ def testFactoredTensorImplicitConversion(self):
     a = np.random.rand(3, 4, 5)
     b = np.random.rand(6, 5)
     c = np.random.rand(3, 4, 6)
-    # a factored representation of a Tensor of shape (3, 4, 6)
-    factored = common_layers.FactoredTensor(tf.to_float(a), tf.to_float(b))
-    # implicitly converts factored to a Tensor (performing the matmul)
-    d = factored + tf.to_float(c)
-    out = self.evaluate(d)
+    with self.test_session() as session:
+      # a factored representation of a Tensor of shape (3, 4, 6)
+      factored = common_layers.FactoredTensor(tf.to_float(a), tf.to_float(b))
+      # implicitly converts factored to a Tensor (performing the matmul)
+      d = factored + tf.to_float(c)
+      out = session.run(d)
     self.assertEqual(out.shape, (3, 4, 6))
 
   def testConvHiddenReluMemoryEfficient(self):
-    if tf.executing_eagerly():
-      return  # don't run test in Eager mode
-
     batch = 3
     length = 23
     io_size = 16
     filter_size = 7
     x = np.random.rand(batch, length, io_size)
     dy = np.random.rand(batch, length, io_size)
-    with self.session() as session:
+    with self.test_session() as session:
       x = tf.to_float(x)
       dy = tf.to_float(dy)
       f1 = tf.get_variable("f1", [1, io_size, filter_size])
@@ -636,10 +666,11 @@ def testCycleGANUpsampleNnUpsampleConv(self):
     upsampled_output = common_layers.cyclegan_upsample(
         random_input, output_filters, stride, "nn_upsample_conv")
     upsampled_output_shape = tf.shape(upsampled_output)
-    self.evaluate(tf.global_variables_initializer())
-    self.assertAllEqual(
-        [batch, height * stride[0], width * stride[1], output_filters],
-        self.evaluate(upsampled_output_shape))
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      self.assertAllEqual(
+          [batch, height * stride[0], width * stride[1], output_filters],
+          session.run(upsampled_output_shape))
 
   def testCycleGANUpsampleBilinearUpsampleConv(self):
     batch = 8
@@ -655,10 +686,11 @@ def testCycleGANUpsampleBilinearUpsampleConv(self):
     upsampled_output = common_layers.cyclegan_upsample(
         random_input, output_filters, stride, "bilinear_upsample_conv")
     upsampled_output_shape = tf.shape(upsampled_output)
-    self.evaluate(tf.global_variables_initializer())
-    self.assertAllEqual(
-        [batch, height * stride[0], width * stride[1], output_filters],
-        self.evaluate(upsampled_output_shape))
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      self.assertAllEqual(
+          [batch, height * stride[0], width * stride[1], output_filters],
+          session.run(upsampled_output_shape))
 
   def testCycleGANUpsampleConv2dTranspose(self):
     batch = 8
@@ -679,15 +711,13 @@ def testCycleGANUpsampleConv2dTranspose(self):
                                                        output_filters, stride,
                                                        "conv2d_transpose")
     upsampled_output_shape = tf.shape(upsampled_output)
-    self.evaluate(tf.global_variables_initializer())
-    self.assertAllEqual(
-        [batch, upsampled_height, upsampled_width, output_filters],
-        self.evaluate(upsampled_output_shape))
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      self.assertAllEqual(
+          [batch, upsampled_height, upsampled_width, output_filters],
+          session.run(upsampled_output_shape))
 
   def testSpectralNorm(self):
-    if tf.executing_eagerly():
-      return  # don't run test in Eager mode
-
     # Test that after 20 calls to apply_spectral_norm, the spectral
     # norm of the normalized matrix is close to 1.0
     with tf.Graph().as_default():
@@ -746,7 +776,7 @@ def grad_fn(inputs, variables, outputs, grad_outputs):
     custom_grads = tf.gradients(custom_loss,
                                 [a, b, c] + [tf.trainable_variables()[1]])
 
-    with self.session() as sess:
+    with self.test_session() as sess:
       sess.run(tf.global_variables_initializer())
       out_val, custom_out_val, grads_val, custom_grads_val = sess.run(
           [out, custom_out, grads, custom_grads])
@@ -777,7 +807,7 @@ def grad_fn(inputs, variables, unused_outputs, unused_grad_outputs):
     expected_grads = [
         tf.ones_like(t) * (i + 1.) for i, t in enumerate([a, b, c, w])
     ]
-    with self.session() as sess:
+    with self.test_session() as sess:
       sess.run(tf.global_variables_initializer())
       g_val, eg_val = sess.run([grads, expected_grads])
       for g1, g2 in zip(g_val, eg_val):
@@ -823,7 +853,7 @@ def fn_recompute(x):
     grad1 = tf.gradients(out1, recompute_vars)
     grad2 = tf.gradients(out2, reg_vars)
 
-    with self.session() as sess:
+    with self.test_session() as sess:
       sess.run(tf.global_variables_initializer())
       outs = sess.run([out1, out2, grad1, grad2])
       self.assertAllClose(outs[0], outs[1])
diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index fe4672410..efa031957 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -22,76 +22,79 @@
 import tensorflow as tf
 
 
-@tf.contrib.eager.run_all_tests_in_graph_and_eager_modes
 class CommonVideoTest(tf.test.TestCase):
 
-  def runScheduledSampleFunc(self, func, var, batch_size):
+  @staticmethod
+  def RunScheduledSampleFunc(func, var, batch_size):
     ground_truth_x = list(range(1, batch_size+1))
     generated_x = [-x for x in ground_truth_x]
     ground_truth_x = tf.convert_to_tensor(ground_truth_x)
     generated_x = tf.convert_to_tensor(generated_x)
     ss_out = func(ground_truth_x, generated_x, batch_size, var)
-    output = self.evaluate([ground_truth_x, generated_x, ss_out])
+    with tf.Session() as session:
+      output = session.run([ground_truth_x, generated_x, ss_out])
     return output
 
   def testScheduledSampleProbStart(self):
-    ground_truth_x, _, ss_out = self.runScheduledSampleFunc(
+    ground_truth_x, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
         common_video.scheduled_sample_prob, 1.0, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
   def testScheduledSampleProbMid(self):
-    _, _, ss_out = self.runScheduledSampleFunc(
+    _, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
         common_video.scheduled_sample_prob, 0.5, 1000)
     positive_count = np.sum(ss_out > 0)
-    self.assertAlmostEqual(positive_count / 1000.0, 0.5, places=1)
+    self.assertAlmostEqual(positive_count / 1000.0, 0.5, places=2)
 
   def testScheduledSampleProbEnd(self):
-    _, generated_x, ss_out = self.runScheduledSampleFunc(
+    _, generated_x, ss_out = CommonVideoTest.RunScheduledSampleFunc(
         common_video.scheduled_sample_prob, 0.0, 10)
     self.assertAllEqual(generated_x, ss_out)
 
   def testScheduledSampleCountStart(self):
-    ground_truth_x, _, ss_out = self.runScheduledSampleFunc(
+    ground_truth_x, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
         common_video.scheduled_sample_count, 10, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
   def testScheduledSampleCountMid(self):
-    _, _, ss_out = self.runScheduledSampleFunc(
+    _, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
         common_video.scheduled_sample_count, 5, 10)
     positive_count = np.sum(ss_out > 0)
     self.assertEqual(positive_count, 5)
 
   def testScheduledSampleCountEnd(self):
-    _, generated_x, ss_out = self.runScheduledSampleFunc(
+    _, generated_x, ss_out = CommonVideoTest.RunScheduledSampleFunc(
         common_video.scheduled_sample_count, 0, 10)
     self.assertAllEqual(generated_x, ss_out)
 
   def testDynamicTileAndConcat(self):
-    # image = (1 X 4 X 4 X 1)
-    image = [[1, 2, 3, 4],
-             [2, 4, 5, 6],
-             [7, 8, 9, 10],
-             [7, 9, 10, 1]]
-    image_t = tf.expand_dims(tf.expand_dims(image, axis=0), axis=-1)
-    image_t = tf.cast(image_t, dtype=tf.float32)
-
-    # latent = (1 X 2)
-    latent = np.array([[90, 100]])
-    latent_t = tf.cast(tf.convert_to_tensor(latent), dtype=tf.float32)
-
-    tiled = common_video.tile_and_concat(
-        image_t, latent_t)
-    tiled_np, image_np = self.evaluate([tiled, image_t])
-    tiled_latent = tiled_np[0, :, :, -1]
-    self.assertAllEqual(tiled_np.shape, (1, 4, 4, 2))
-
-    self.assertAllEqual(tiled_np[:, :, :, :1], image_np)
-    self.assertAllEqual(
-        tiled_latent,
-        [[90, 90, 90, 90],
-         [100, 100, 100, 100],
-         [90, 90, 90, 90],
-         [100, 100, 100, 100]])
+    with tf.Graph().as_default():
+      # image = (1 X 4 X 4 X 1)
+      image = [[1, 2, 3, 4],
+               [2, 4, 5, 6],
+               [7, 8, 9, 10],
+               [7, 9, 10, 1]]
+      image_t = tf.expand_dims(tf.expand_dims(image, axis=0), axis=-1)
+      image_t = tf.cast(image_t, dtype=tf.float32)
+
+      # latent = (1 X 2)
+      latent = np.array([[90, 100]])
+      latent_t = tf.cast(tf.convert_to_tensor(latent), dtype=tf.float32)
+
+      with tf.Session() as session:
+        tiled = common_video.tile_and_concat(
+            image_t, latent_t)
+        tiled_np, image_np = session.run([tiled, image_t])
+        tiled_latent = tiled_np[0, :, :, -1]
+        self.assertAllEqual(tiled_np.shape, (1, 4, 4, 2))
+
+        self.assertAllEqual(tiled_np[:, :, :, :1], image_np)
+        self.assertAllEqual(
+            tiled_latent,
+            [[90, 90, 90, 90],
+             [100, 100, 100, 100],
+             [90, 90, 90, 90],
+             [100, 100, 100, 100]])
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index 20dfc22b2..c75f61704 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -23,7 +23,6 @@
 import tensorflow as tf
 
 
-@tf.contrib.eager.run_all_tests_in_graph_and_eager_modes
 class DiscretizationTest(tf.test.TestCase):
   """Tests for discretization layers."""
 
@@ -35,29 +34,37 @@ def testBitToIntZeros(self):
     x_bit = tf.zeros(shape=[1, 10], dtype=tf.float32)
     x_int = tf.zeros(shape=[1], dtype=tf.int32)
     diff = discretization.bit_to_int(x_bit, num_bits=10) - x_int
-    d = self.evaluate(diff)
-    self.assertEqual(d, 0)
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      d = sess.run(diff)
+      self.assertEqual(d, 0)
 
   def testBitToIntOnes(self):
     x_bit = tf.ones(shape=[1, 3], dtype=tf.float32)
     x_int = 7 * tf.ones(shape=[1], dtype=tf.int32)
     diff = discretization.bit_to_int(x_bit, num_bits=3) - x_int
-    d = self.evaluate(diff)
-    self.assertEqual(d, 0)
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      d = sess.run(diff)
+      self.assertEqual(d, 0)
 
   def testIntToBitZeros(self):
     x_bit = tf.zeros(shape=[1, 10], dtype=tf.float32)
     x_int = tf.zeros(shape=[1], dtype=tf.int32)
     diff = discretization.int_to_bit(x_int, num_bits=10) - x_bit
-    d = self.evaluate(diff)
-    self.assertTrue(np.all(d == 0))
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      d = sess.run(diff)
+      self.assertTrue(np.all(d == 0))
 
   def testIntToBitOnes(self):
     x_bit = tf.ones(shape=[1, 3], dtype=tf.float32)
     x_int = 7 * tf.ones(shape=[1], dtype=tf.int32)
     diff = discretization.int_to_bit(x_int, num_bits=3) - x_bit
-    d = self.evaluate(diff)
-    self.assertTrue(np.all(d == 0))
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      d = sess.run(diff)
+      self.assertTrue(np.all(d == 0))
 
   def testProjectHidden(self):
     hidden_size = 60
@@ -68,9 +75,11 @@ def testProjectHidden(self):
         shape=[num_blocks, hidden_size, block_dim], dtype=tf.float32)
     x_projected = discretization.project_hidden(x, projection_tensors,
                                                 hidden_size, num_blocks)
-    x_projected_eval = self.evaluate(x_projected)
-    self.assertEqual(np.shape(x_projected_eval), (1, num_blocks, block_dim))
-    self.assertTrue(np.all(x_projected_eval == 0))
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      x_projected_eval = sess.run(x_projected)
+      self.assertEqual(np.shape(x_projected_eval), (1, num_blocks, block_dim))
+      self.assertTrue(np.all(x_projected_eval == 0))
 
   def testSliceHiddenZeros(self):
     hidden_size = 60
@@ -78,9 +87,11 @@ def testSliceHiddenZeros(self):
     num_blocks = 3
     x = tf.zeros(shape=[1, hidden_size], dtype=tf.float32)
     x_sliced = discretization.slice_hidden(x, hidden_size, num_blocks)
-    x_sliced_eval = self.evaluate(x_sliced)
-    self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
-    self.assertTrue(np.all(x_sliced_eval == 0))
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      x_sliced_eval = sess.run(x_sliced)
+      self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
+      self.assertTrue(np.all(x_sliced_eval == 0))
 
   def testSliceHiddenOnes(self):
     hidden_size = 60
@@ -88,9 +99,11 @@ def testSliceHiddenOnes(self):
     num_blocks = 3
     x = tf.ones(shape=[1, hidden_size], dtype=tf.float32)
     x_sliced = discretization.slice_hidden(x, hidden_size, num_blocks)
-    x_sliced_eval = self.evaluate(x_sliced)
-    self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
-    self.assertTrue(np.all(x_sliced_eval == 1))
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      x_sliced_eval = sess.run(x_sliced)
+      self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
+      self.assertTrue(np.all(x_sliced_eval == 1))
 
   def testNearestNeighbors(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
@@ -102,14 +115,13 @@ def testNearestNeighbors(self):
         x, means, block_v_size=4)
     x_means_hot_test = np.array([[0, 1, 0, 0], [1, 0, 0, 0]])
     x_means_hot_test = np.expand_dims(x_means_hot_test, axis=0)
-    x_means_hot_eval = self.evaluate(x_means_hot)
-    self.assertEqual(np.shape(x_means_hot_eval), (1, 2, 4))
-    self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      x_means_hot_eval = sess.run(x_means_hot)
+      self.assertEqual(np.shape(x_means_hot_eval), (1, 2, 4))
+      self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
 
   def testGetVQBottleneck(self):
-    if tf.executing_eagerly():
-      return  # don't run test in Eager mode
-
     bottleneck_bits = 2
     bottleneck_size = 2**bottleneck_bits
     hidden_size = 3
@@ -118,7 +130,7 @@ def testGetVQBottleneck(self):
     assign_op = means.assign(tf.zeros(shape=[bottleneck_size, hidden_size]))
     means_new, _, _ = discretization.get_vq_codebook(bottleneck_size,
                                                      hidden_size)
-    with self.session() as sess:
+    with self.test_session() as sess:
       tf.global_variables_initializer().run()
       sess.run(assign_op)
       self.assertTrue(np.all(sess.run(means_new) == 0))
@@ -130,32 +142,37 @@ def testVQNearestNeighbors(self):
         [[1, 0, 0], [0, 1, 0], [0, 0, 1], [9, 9, 9]], dtype=tf.float32)
     x_means_hot, _, _ = discretization.vq_nearest_neighbor(x, means)
     x_means_hot_test = np.array([[0, 1, 0, 0], [1, 0, 0, 0]])
-    x_means_hot_eval = self.evaluate(x_means_hot)
-    self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
-    self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      x_means_hot_eval = sess.run(x_means_hot)
+      self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
+      self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
 
   def testVQDiscreteBottleneck(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
     x_means_hot, _ = discretization.vq_discrete_bottleneck(x, bottleneck_bits=2)
-    self.evaluate(tf.global_variables_initializer())
-    x_means_hot_eval = self.evaluate(x_means_hot)
-    self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      x_means_hot_eval = sess.run(x_means_hot)
+      self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
 
   def testVQDiscreteUnbottlenck(self):
     x = tf.constant([[1, 0, 0, 0], [0, 0, 1, 0]], dtype=tf.int32)
     x_means = discretization.vq_discrete_unbottleneck(x, hidden_size=3)
-    self.evaluate(tf.global_variables_initializer())
-    x_means_eval = self.evaluate(x_means)
-    self.assertEqual(np.shape(x_means_eval), (2, 3))
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      x_means_eval = sess.run(x_means)
+      self.assertEqual(np.shape(x_means_eval), (2, 3))
 
   def testGumbelSoftmaxDiscreteBottleneck(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
     tf.add_to_collection(tf.GraphKeys.GLOBAL_STEP, tf.constant(1))
     x_means_hot, _ = discretization.gumbel_softmax_discrete_bottleneck(
         x, bottleneck_bits=2)
-    self.evaluate(tf.global_variables_initializer())
-    x_means_hot_eval = self.evaluate(x_means_hot)
-    self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      x_means_hot_eval = sess.run(x_means_hot)
+      self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
 
 
 if __name__ == '__main__':
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 0a8ef0311..04fba3bff 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -27,7 +27,6 @@
 
 class ModalityTest(tf.test.TestCase):
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes
   def testSymbolModalityInputs(self):
     batch_size = 10
     num_datashards = 5
@@ -42,14 +41,14 @@ def testSymbolModalityInputs(self):
     m = modalities.SymbolModality(model_hparams, vocab_size)
     data_parallelism = expert_utils.Parallelism(
         ["/device:CPU:0"] * num_datashards)
-    xs = tf.split(x, num_datashards)
-    sharded_output = m.bottom_sharded(xs, data_parallelism)
-    output = tf.concat(sharded_output, 0)
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(output)
+    with self.test_session() as session:
+      xs = tf.split(x, num_datashards)
+      sharded_output = m.bottom_sharded(xs, data_parallelism)
+      output = tf.concat(sharded_output, 0)
+      session.run(tf.global_variables_initializer())
+      res = session.run(output)
     self.assertEqual(res.shape, (batch_size, length, 1, hidden_size))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes
   def testSymbolModalityTargets(self):
     batch_size = 10
     num_datashards = 5
@@ -67,15 +66,16 @@ def testSymbolModalityTargets(self):
     m = modalities.SymbolModality(model_hparams, vocab_size)
     data_parallelism = expert_utils.Parallelism(
         ["/device:CPU:0"] * num_datashards)
-    sharded_body_output = tf.split(tf.to_float(body_output), num_datashards)
-    sharded_targets = tf.split(targets, num_datashards)
-    sharded_logits = m.top_sharded(sharded_body_output, sharded_targets,
-                                   data_parallelism)
-    train_loss = m.loss_sharded(sharded_logits, sharded_targets,
-                                data_parallelism)
-    logits = tf.concat(sharded_logits, 0)
-    self.evaluate(tf.global_variables_initializer())
-    res1, res2 = self.evaluate((logits, train_loss))
+    with self.test_session() as session:
+      sharded_body_output = tf.split(tf.to_float(body_output), num_datashards)
+      sharded_targets = tf.split(targets, num_datashards)
+      sharded_logits = m.top_sharded(sharded_body_output, sharded_targets,
+                                     data_parallelism)
+      train_loss = m.loss_sharded(sharded_logits, sharded_targets,
+                                  data_parallelism)
+      logits = tf.concat(sharded_logits, 0)
+      session.run(tf.global_variables_initializer())
+      res1, res2 = session.run((logits, train_loss))
     self.assertEqual(res1.shape, (batch_size, length, height, 1, vocab_size))
     self.assertEqual(res2.shape, ())
 
@@ -97,7 +97,7 @@ def testSymbolModalityTargetsFactored(self):
     m = modalities.SymbolModality(model_hparams, vocab_size)
     data_parallelism = expert_utils.Parallelism(
         ["/device:CPU:0"] * num_datashards)
-    with self.session() as session:
+    with self.test_session() as session:
       sharded_body_output = tf.split(tf.to_float(body_output), num_datashards)
       sharded_targets = tf.split(targets, num_datashards)
       sharded_logits = m.top_sharded(sharded_body_output, sharded_targets,

From 1ad804ff172dd2728af9afdf260a34842db2f84b Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Thu, 23 Aug 2018 13:35:15 -0700
Subject: [PATCH 0656/2720] Hierarchical mixture of experts on mtf.

PiperOrigin-RevId: 209990439
---
 tensor2tensor/mesh_tensorflow/mtf_layers.py   |  19 +-
 .../mesh_tensorflow/mtf_transformer.py        |  73 +++--
 .../research/experiments_moe.py               |  13 +
 tensor2tensor/mesh_tensorflow/research/moe.py | 305 +++++++++++++++---
 4 files changed, 341 insertions(+), 69 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers.py b/tensor2tensor/mesh_tensorflow/mtf_layers.py
index ca931561f..611cc4ac7 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers.py
@@ -384,7 +384,8 @@ def multihead_attention(query_antecedent,
   same dimensionality (kv_channels).
 
   Args:
-    query_antecedent: a mtf.Tensor with shape [batch, query_length, io_channels]
+    query_antecedent: a mtf.Tensor with shape
+      [<batch_dims>, query_length, io_channels]
     memory_antecedent: a mtf.Tensor with shape
       [batch, memory_length, io_channels] (optional)
     mask: mask Tensor (see attention_mask())
@@ -400,7 +401,8 @@ def multihead_attention(query_antecedent,
   Raises:
     ValueError: if the dimensions do not match.
   """
-  batch, query_length, io_channels = query_antecedent.shape.dims
+  batch_dims = query_antecedent.shape.dims[:-2]
+  query_length, io_channels = query_antecedent.shape.dims[-2:]
   with tf.variable_scope(name,
                          default_name="multihead_attention",
                          values=[query_antecedent, memory_antecedent]):
@@ -410,24 +412,25 @@ def multihead_attention(query_antecedent,
     if memory_antecedent is None:
       memory_antecedent = rename_length_to_memory_length(
           query_antecedent, query_length.name)
-    memory_batch, memory_length, memory_channels = memory_antecedent.shape.dims
-    if memory_batch != batch:
+    memory_batch_dims = memory_antecedent.shape.dims[:-2]
+    memory_length, memory_channels = memory_antecedent.shape.dims[-2:]
+    if memory_batch_dims != batch_dims:
       raise ValueError("memory batch must equal query batch")
     if memory_channels != io_channels:
       raise ValueError("memory channels must equal query channels")
     q = mtf.einsum(
         [query_antecedent, q_var],
-        mtf.Shape([batch, heads, query_length, kv_channels]))
+        mtf.Shape(batch_dims + [heads, query_length, kv_channels]))
     k = mtf.einsum(
         [memory_antecedent, k_var],
-        mtf.Shape([batch, heads, memory_length, kv_channels]))
+        mtf.Shape(batch_dims + [heads, memory_length, kv_channels]))
     v = mtf.einsum(
         [memory_antecedent, v_var],
-        mtf.Shape([batch, heads, memory_length, kv_channels]))
+        mtf.Shape(batch_dims + [heads, memory_length, kv_channels]))
     o = dot_product_attention(
         q, k, v, mask, dropout, dropout_broadcast_dims)
     return mtf.einsum(
-        [o, o_var], mtf.Shape([batch, query_length, io_channels]))
+        [o, o_var], mtf.Shape(batch_dims + [query_length, io_channels]))
 
 
 def multihead_self_attention_incremental(query_antecedent,
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
index 9f2e8a176..26f994f2a 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
@@ -35,8 +35,18 @@ class MtfTransformer(mtf_model.MtfModel):
   """Transformer in mesh_tensorflow."""
 
   @property
-  def batch_dim(self):
-    return mtf.Dimension("batch", self._hparams.batch_size)
+  def batch_dims(self):
+    hparams = self._hparams
+    if hparams.outer_batch_size == 0:
+      return [mtf.Dimension("batch", hparams.batch_size)]
+    else:
+      if hparams.batch_size % hparams.outer_batch_size != 0:
+        raise ValueError(
+            "hparams.outer_batch_size must divide hparams.batch_size")
+      return [
+          mtf.Dimension("outer_batch", hparams.outer_batch_size),
+          mtf.Dimension("inner_batch",
+                        hparams.batch_size // hparams.outer_batch_size)]
 
   @property
   def inputs_vocab_dim(self):
@@ -88,9 +98,9 @@ def activation_dtype(self):
 
   def _import_to_batch_by_length(self, x, name, mesh, hparams):
     del hparams
-    x = tf.reshape(x, [self.batch_dim.size, self.length_dim.size])
-    return mtf.import_fully_replicated(
-        mesh, x, mtf.Shape([self.batch_dim, self.length_dim]), name=name)
+    mtf_shape = mtf.Shape(self.batch_dims + [self.length_dim])
+    x = tf.reshape(x, mtf_shape.to_integer_list)
+    return mtf.import_fully_replicated(mesh, x, mtf_shape, name=name)
 
   def _embedding_and_softmax_vars(self, mesh):
     hparams = self._hparams
@@ -173,7 +183,7 @@ def pad_to_max_length(x):
     def layer_prepostprocess_dropout(x):
       return mtf.dropout(
           x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
-          noise_shape=mtf.Shape([self.batch_dim, self.model_dim]))
+          noise_shape=mtf.Shape(self.batch_dims + [self.model_dim]))
 
     extra_losses = []
     (inputs_embedding_var,
@@ -275,10 +285,10 @@ def _feedforward_layer(self, x, losses=None):
     """Feed-forward layer.
 
     Args:
-      x: a mtf.Tensor with shape [batch_dim, length_dim, model_dim]
+      x: a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim]
       losses: a list to be appended-to
     Returns:
-      a mtf.Tensor with shape [batch_dim, length_dim, model_dim]
+      a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim]
     Raises:
       ValueError: if hparams make no sense
     """
@@ -294,9 +304,15 @@ def _feedforward_layer(self, x, losses=None):
           self.model_dim,
           hparams,
           hparams.mode == tf.estimator.ModeKeys.TRAIN)
+    elif feedforward_layer == "hmoe":
+      output, loss = moe.transformer_moe_layer_v2(
+          x,
+          self.model_dim,
+          hparams,
+          hparams.mode == tf.estimator.ModeKeys.TRAIN)
       if losses is not None:
         losses.append(loss)
-        return output
+      return output
     else:
       raise ValueError(
           "hparams.feedforward_layer not recognized %s" % feedforward_layer)
@@ -311,26 +327,25 @@ def _layer_stack(self,
     """Encoder or decoder stack.
 
     Args:
-      x: a mtf.Tensor with shape [batch_dim, length_dim, model_dim]
+      x: a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim]
       num_layers: an integer
       encoder_output: an optional mtf.Tensor with shape
-        [batch_dim, encoder_length_dim, model_dim]
+        [<batch_dims>, encoder_length_dim, model_dim]
       self_attention_mask: an optional mtf.Tensor with shape
         [batch, length_dim, memory_length_dim] containing values 0 or -inf.
       encdec_attention_mask: an optional mtf.Tensor with shape
         [batch, length_dim, encoder_length_dim] containing values 0 or -inf.
       losses: a list to be appended-to
     Returns:
-      a mtf.Tensor with shape [batch_dim, length_dim, model_dim]
+      a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim]
     Raises:
       ValueError: if hparams make no sense
     """
     hparams = self._hparams
-
     def layer_prepostprocess_dropout(x):
       return mtf.dropout(
           x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
-          noise_shape=mtf.Shape([self.batch_dim, self.model_dim]))
+          noise_shape=mtf.Shape(self.batch_dims + [self.model_dim]))
     num_layer_norms = num_layers * (2 if encoder_output is None else 3) + 1
     layer_norms_dim = mtf.Dimension("layer_norms", num_layer_norms)
     layer_norm_combined_var = mtf.get_variable(
@@ -413,13 +428,13 @@ def _sample(self, features, mesh):
           k = mtf.einsum(
               [encoder_output, k_var],
               mtf.Shape(
-                  [self.batch_dim, self.heads_dim,
-                   self.memory_length_dim, self.kv_dim]))
+                  self.batch_dims + [self.heads_dim,
+                                     self.memory_length_dim, self.kv_dim]))
           v = mtf.einsum(
               [encoder_output, v_var],
               mtf.Shape(
-                  [self.batch_dim, self.heads_dim,
-                   self.memory_length_dim, self.kv_dim]))
+                  self.batch_dims + [self.heads_dim,
+                                     self.memory_length_dim, self.kv_dim]))
         encdec_tensors.append((q_var, o_var, k, v))
       partial_targets = None
     else:
@@ -444,13 +459,15 @@ def _sample(self, features, mesh):
             partial_targets, "partial_targets", mesh, hparams)
 
     if hparams.beam_size == 1:
-      ids_shape = mtf.Shape([self.batch_dim, self.length_dim])
-      kv_shape = mtf.Shape([self.batch_dim, self.heads_dim,
+      ids_shape = mtf.Shape(self.batch_dims + [self.length_dim])
+      kv_shape = mtf.Shape(self.batch_dims +
+                           [self.heads_dim,
                             self.memory_length_dim, self.kv_dim])
     else:
       beam_dim = mtf.Dimension("beam", hparams.beam_size)
-      ids_shape = mtf.Shape([self.batch_dim, beam_dim, self.length_dim])
-      kv_shape = mtf.Shape([self.batch_dim, beam_dim, self.heads_dim,
+      ids_shape = mtf.Shape(self.batch_dims + [beam_dim, self.length_dim])
+      kv_shape = mtf.Shape(self.batch_dims +
+                           [beam_dim, self.heads_dim,
                             self.memory_length_dim, self.kv_dim])
 
     initial_ids = mtf.constant(mesh, 0, ids_shape, dtype=tf.int32)
@@ -527,7 +544,7 @@ def _decoder_layer_stack_incremental(self,
     attention as well as the weight matrices q_var and o_var.
 
     Args:
-      x: a mtf.Tensor with shape [batch_dim, model_dim]
+      x: a mtf.Tensor with shape [<batch_dims>, model_dim]
       step_num: an mtf integer Scalar
       encdec_tensors: an optional list of num_layers tuples, each of the form
         (q_var, o_var, k, v)
@@ -539,7 +556,7 @@ def _decoder_layer_stack_incremental(self,
         [batch, length_dim, encoder_length_dim] containing values 0 or -inf.
 
     Returns:
-      y: a mtf.Tensor with shape [batch_dim, model_dim]
+      y: a mtf.Tensor with shape [<batch_dims>, model_dim]
       new_self_attention_k: a list of num_layers mtf.Tensors, with the same
         shapes as the elements of self_attention_k
       new_self_attention_v: a list of num_layers mtf.Tensors, with the same
@@ -619,6 +636,7 @@ def mtf_transformer_base():
   # round up vocab sizes to be a multiple of this value
   hparams.vocab_divisor = 128
 
+  # options are dense_relu_dense, moe, hmoe
   hparams.add_hparam("feedforward_layer", "dense_relu_dense")
 
   # Use targets_embedding_var * rsqrt(d_model) as softmax_var
@@ -643,6 +661,13 @@ def mtf_transformer_base():
   hparams.add_hparam("decode_length_multiplier", 1.5)
   hparams.add_hparam("decode_length_constant", 10.0)
 
+  # If nonzero, we split the batch across two tensor-dimensions named
+  # "outer_batch" and "inner_batch", allowing for splitting across two mesh
+  # dimensions.  This is necessary for hierarchical mixture of experts.
+  # The two tensor dimensions have sizes hparams.outer_batch_size and
+  # hparams.batch_size // hparams.outer_batch_size.
+  hparams.add_hparam("outer_batch_size", 0)
+
   return hparams
 
 
diff --git a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
index d976c5d78..f95c2df02 100644
--- a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
@@ -123,3 +123,16 @@ def mtf_transformer_lm_moe():
   return hparams
 
 
+@registry.register_hparams
+def xmoe_2d():
+  """Two-dimensional hierarchical mixture of experts."""
+  hparams = xmoe_top_2()
+  hparams.mesh_shape = "b0:2;b1:4"
+  hparams.outer_batch_size = 4
+  hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
+  hparams.moe_num_experts = [4, 4]
+  hparams.moe_group_size = [256, 256]
+  hparams.feedforward_layer = "hmoe"
+  return hparams
+
+
diff --git a/tensor2tensor/mesh_tensorflow/research/moe.py b/tensor2tensor/mesh_tensorflow/research/moe.py
index c1896ca7b..917d2ad6c 100644
--- a/tensor2tensor/mesh_tensorflow/research/moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/moe.py
@@ -99,27 +99,28 @@ def transformer_moe_layer_v1(inputs, output_dim, hparams, train):
   input_dim = inputs.shape.dims[-1]
   hidden_dim = mtf.Dimension("expert_hidden", hparams.moe_hidden_size)
   experts_dim = mtf.Dimension("experts", hparams.moe_num_experts)
-  group_dim = mtf.Dimension("group", hparams.moe_group_size)
+  group_size_dim = mtf.Dimension("group", hparams.moe_group_size)
   batch_dim = mtf.Dimension(
       orig_inputs.shape[0].name,
-      orig_inputs.shape.size // (group_dim.size * input_dim.size))
-  inputs = mtf.reshape(inputs, [batch_dim, group_dim, input_dim])
+      orig_inputs.shape.size // (group_size_dim.size * input_dim.size))
+  inputs = mtf.reshape(inputs, [batch_dim, group_size_dim, input_dim])
 
   # Each sequence sends expert_capacity positions to each expert.
   capacity_factor = (
       hparams.moe_capacity_factor_train if train else
       hparams.moe_capacity_factor_eval)
   expert_capacity = min(
-      group_dim.size,
-      int((group_dim.size * capacity_factor) / experts_dim.size))
+      group_size_dim.size,
+      int((group_size_dim.size * capacity_factor) / experts_dim.size))
   expert_capacity_dim = mtf.Dimension("expert_capacity", expert_capacity)
 
   experts_dim_unsplit = mtf.Dimension("expert_unsplit", experts_dim.size)
   batch_dim_unsplit = mtf.Dimension("batch_unsplit", batch_dim.size)
 
   if hparams.moe_gating == "top_2":
-    forward_assignment, backward_assignment, loss = _top_2_gating(
+    dispatch_tensor, combine_tensor, loss = _top_2_gating(
         inputs=inputs,
+        outer_expert_dims=None,
         experts_dim=experts_dim_unsplit,
         expert_capacity_dim=expert_capacity_dim,
         max_experts=None,
@@ -129,7 +130,7 @@ def transformer_moe_layer_v1(inputs, output_dim, hparams, train):
     raise ValueError("unknown hparams.moe_gating=%s" % hparams.moe_gating)
 
   # put num_experts dimension first to make split easier in alltoall
-  expert_inputs = mtf.einsum([inputs, forward_assignment], mtf.Shape(
+  expert_inputs = mtf.einsum([inputs, dispatch_tensor], mtf.Shape(
       [experts_dim_unsplit, batch_dim, expert_capacity_dim, input_dim]))
 
   expert_inputs = mtf.reshape(expert_inputs, mtf.Shape(
@@ -145,16 +146,239 @@ def transformer_moe_layer_v1(inputs, output_dim, hparams, train):
   expert_output = mtf.reshape(expert_output, mtf.Shape(
       [experts_dim_unsplit, batch_dim, expert_capacity_dim, input_dim]))
 
-  output = mtf.einsum([expert_output, backward_assignment], mtf.Shape(
-      [batch_dim, group_dim, output_dim]))
+  output = mtf.einsum([expert_output, combine_tensor], mtf.Shape(
+      [batch_dim, group_size_dim, output_dim]))
 
   output = mtf.reshape(output, orig_inputs.shape.dims[:-1] + [output_dim])
 
   return output, loss * hparams.moe_loss_coef
 
 
+def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
+  """2-level mixture of experts.
+
+  Adapted from the paper https://arxiv.org/abs/1701.06538
+
+  Note: until the algorithm and inferface solidify, we pass in a hyperparameters
+  dictionary in order not to complicate the interface in mtf_transformer.py .
+  Once this code moves out of "research", we should pass the hyperparameters
+  separately.
+
+  Hyperparameters used:
+    hparams.moe_num_experts: number of experts
+    hparams.moe_hidden_size: size of hidden layer in each expert
+    hparams.moe_group_size: size of each "group" for gating purposes
+    hparams.moe_capacity_factor_train: a float
+    hparams.moe_capacity_factor_eval: a float
+    hparams.moe_gating: a string
+    + all hyperparmeters used by _top_2_gating()
+
+  One set of params for experts in first level and different of hparams
+  per expert in the second level.
+  The number of parameters in the gating network is:
+    (input_dim.size * (hparams.num_experts) +
+      (moe_hidden_size * hparams.num_experts) * hparams.num_experts
+
+
+  The number of parameters in the experts themselves is:
+    (hparams.num_experts
+     * (input_dim.size + output_dim.size)
+     * hparams.moe_hidden_size)
+
+  The input is n-dimensional: [<batch_and_length_dims>, input_dim], consisting
+  of the representations of all positions in a batch of sequences.
+
+  Each position of each sequence is sent to 0-2 experts.  The expert
+  choices and the combination weights are determined by a learned gating
+  function.
+
+  This function returns a small auxiliary loss that should be added to the
+  training loss of the model.  This loss helps to balance expert usage.
+  Without the loss, it is very likely that a few experts will be trained and
+  the rest will starve.
+
+  Several hacks are necessary to get around current TPU limitations:
+
+  - To ensure static shapes, we enforce (by truncation/padding)
+    that each sequence send the same number of elements to each expert.
+
+    It would make more sense to enforce this equality over the entire batch,
+    but due to our hacked-up gather-by-matmul implementation, we need to divide
+    the batch into "groups".  For each group, the same number of elements
+    are sent to each expert.
+
+  TODO(noam): Factor this code better.  We want to be able to substitute
+  different code for the experts themselves.
+
+  Dimensions cheat sheet:
+  a, b: batch size
+  l: original sequence length
+  m: input depth
+  n: output depth
+  g, h: number of groups
+  s, t: group size
+  x, y: number of experts
+  c, d: expert capacity
+
+  input: [a0, b1, l, m]
+  input: [a0, g1, s, m]
+  dispatch_tensor_x: [a0, g1, s, x, c]
+  expert_input: [a0, g1, x, c, m]
+  alltoall: [a0, g, x1, c, m]
+  alltoall: [a0, g, x1, c, m]
+  transpose: [x1, a0, g, c, m]
+  reshape: [x1, h0, s, m]
+  assignment2: [x1, h0, t, y, d]
+  expert_input2: [x1, h0, y, d, m]
+  alltoall: [x1, h, y0, d, m]
+  ...
+  reverse of that
+
+  gating params 0: [m, x]
+  gating params 1: [x1, m, y]
+
+  expert params:
+     [x1, y0, m, hidden]
+     [x1, y0, hidden, n]
+
+  Args:
+    inputs: a mtf.Tensor with shape [a, b, l, m]
+    output_dim: a mtf.Dimension (for Transformer, this is input_dim)
+    hparams: model hyperparameters
+    train: a boolean
+
+  Returns:
+    outputs: a Tensor with shape [a, b, l, n]
+    loss: a mtf scalar
+
+  Raises:
+    ValueError: on unrecognized hparams.moe_gating
+  """
+  assert len(hparams.moe_num_experts) == 2
+  a0, b1, l, m = inputs.shape.dims
+  # a = mtf.Dimension("a_unsplit", a0.size)
+  # b = mtf.Dimension("b_unsplit", b1.size)
+  hidden_dim = mtf.Dimension("expert_hidden", hparams.moe_hidden_size)
+  x1 = mtf.Dimension("expert_x", hparams.moe_num_experts[0])
+  y0 = mtf.Dimension("expert_y", hparams.moe_num_experts[1])
+  x = mtf.Dimension("expert_x_unsplit", hparams.moe_num_experts[0])
+  y = mtf.Dimension("expert_y_unsplit", hparams.moe_num_experts[1])
+  s = mtf.Dimension("group_size_x", hparams.moe_group_size[0])
+  t = mtf.Dimension("group_size_y", hparams.moe_group_size[1])
+  g1 = mtf.Dimension(b1.name, b1.size * l.size // s.size)
+  g = mtf.Dimension(b1.name + "_unsplit", b1.size * l.size // s.size)
+  n = output_dim
+
+  # First level of expert routing
+  inputs = mtf.reshape(inputs, [a0, g1, s, m])
+
+  # Each sequence sends (at most?) expert_capacity positions to each expert.
+  # Static expert_capacity dimension is needed for expert batch sizes
+  capacity_factor = (
+      hparams.moe_capacity_factor_train if train else
+      hparams.moe_capacity_factor_eval)
+  expert_capacity = min(
+      s.size,
+      int((s.size * capacity_factor) / x.size))
+  c = mtf.Dimension("expert_capacity_x", expert_capacity)
+
+  # Get the assignments for the first level.
+  # dispatch_tensor_x has shape [a0, g1, s, x, c]
+  if hparams.moe_gating == "top_2":
+    dispatch_tensor_x, combine_tensor_x, loss_outer = _top_2_gating(
+        inputs=inputs,
+        outer_expert_dims=None,
+        experts_dim=x,
+        expert_capacity_dim=c,
+        max_experts=None,
+        hparams=hparams,
+        train=train)
+  else:
+    raise ValueError("unknown hparams.moe_gating=%s" % hparams.moe_gating)
+
+  # Now create expert_inputs based on the assignments.
+  # put num_experts dimension first to make split easier in alltoall
+  expert_inputs_x = mtf.einsum([inputs, dispatch_tensor_x], mtf.Shape(
+      [x, a0, g1, c, m]))
+
+  # First level, all to all. Here we change the split dimension from g1 to x1.
+  expert_inputs_x = mtf.reshape(expert_inputs_x, mtf.Shape(
+      [x1, a0, g, c, m]))
+
+  # Second level of expert routing
+
+  numerator = a0.size * g.size * c.size
+  if numerator % t.size != 0:
+    raise ValueError("cannont divide evenly %s / %s" % (numerator, t.size))
+  h0 = mtf.Dimension(a0.name, numerator // t.size)
+  h = mtf.Dimension(a0.name + "_unsplit", h0.size)
+  inputs_y = mtf.reshape(expert_inputs_x, [x1, h0, t, m])
+
+  expert_capacity = min(
+      t.size,
+      int((t.size * capacity_factor) / y.size))
+  d = mtf.Dimension("expert_capacity_y", expert_capacity)
+
+  # Get the assignments for the second level.
+  # dispatch_tensor_x has shape [x1, h0, t, y, d]
+  if hparams.moe_gating == "top_2":
+    dispatch_tensor_y, combine_tensor_y, loss_inner = _top_2_gating(
+        inputs=inputs_y,
+        outer_expert_dims=[x1],
+        experts_dim=y,
+        expert_capacity_dim=d,
+        max_experts=None,
+        hparams=hparams,
+        train=train)
+  else:
+    raise ValueError("unknown hparams.moe_gating=%s" % hparams.moe_gating)
+
+  # Now create expert_inputs based on the assignments.
+  # put num_experts dimension first to make split easier in alltoall
+  expert_inputs_y = mtf.einsum([inputs_y, dispatch_tensor_y], mtf.Shape(
+      [y, x1, h0, d, m]))
+
+  # Second level, all to all. Here we change the split dimension from h0 to y0.
+  expert_inputs_y = mtf.reshape(expert_inputs_y, mtf.Shape(
+      [y0, x1, h, d, m]))
+
+  # Now feed the expert inputs through the experts.
+  hidden_output = mtf_layers.dense(
+      expert_inputs_y, hidden_dim, expert_dims=[y0, x1],
+      activation=mtf.relu, use_bias=False, name="expert0")
+  expert_output = mtf_layers.dense(
+      hidden_output, output_dim, expert_dims=[y0, x1],
+      use_bias=False, name="expert1")
+
+  # NOW COMBINE EXPERT OUTPUTS (reversing everything we have done)
+  # expert_output has shape [y0, x1, h, d, n]
+
+  # alltoall
+  expert_output = mtf.reshape(expert_output, mtf.Shape(
+      [y, x1, h0, d, n]))
+
+  # combine results from inner level
+  output_y = mtf.einsum([expert_output, combine_tensor_y], mtf.Shape(
+      [x1, h0, t, n]))
+
+  # simple reshape
+  output = mtf.reshape(output_y, [x1, a0, g, c, n])
+
+  # alltoall
+  expert_output_x = mtf.reshape(output, mtf.Shape([x, a0, g1, c, n]))
+
+  # combine results from outer level
+  output_x = mtf.einsum([expert_output_x, combine_tensor_x], mtf.Shape(
+      [a0, g1, s, m]))
+
+  # simple reshape
+  output = mtf.reshape(output_x, [a0, b1, l, n])
+  return output, (loss_outer + loss_inner) * hparams.moe_loss_coef
+
+
 def _top_2_gating(
-    inputs, experts_dim, expert_capacity_dim, max_experts, hparams, train):
+    inputs, outer_expert_dims, experts_dim, expert_capacity_dim, max_experts,
+    hparams, train):
   """Compute gating for mixture-of-experts in TensorFlow.
 
   Note: until the algorithm and inferface solidify, we pass in a hyperparameters
@@ -168,46 +392,52 @@ def _top_2_gating(
     hparams.moe_second_policy_eval: a string
     hparams.moe_second_threshold: a float
 
-  max_experts is an float tensor with shape [batch_dim, group_dim]
+  max_experts is an float tensor with shape [<batch_dims>, group_size_dim]
   indicating at most how many experts to use per example.  This can be
   used to prevent padding from going to experts.
 
   The returned forward assignment is a tensor used to map (via einsum) from the
-  inputs to the expert_inputs.  Likewise, the returned backward_assignment is
+  inputs to the expert_inputs.  Likewise, the returned combine_tensor is
   used to map (via einsum) from the expert outputs to the outputs.  Both the
   forward and backward assignments are mostly zeros.  The shapes of all of these
   are as follows.
 
-  inputs: [batch_dim, group_dim, input_dim]
-  forward_assignment: [batch_dim, group_dim, experts_dim, expert_capacity_dim]
-  expert_inputs: [batch_dim, experts_dim, expert_capacity_dim, input_dim]
+  inputs: [<batch_dims>, group_size_dim, input_dim]
+  dispatch_tensor:
+    [<batch_dims>, group_size_dim, experts_dim, expert_capacity_dim]
+  expert_inputs:
+    [<batch_dims>, experts_dim, expert_capacity_dim, input_dim]
 
-  expert_outputs: [batch_dim, experts_dim, expert_capacity_dim, output_dim]
-  backward_assignment: [batch_dim, group_dim, experts_dim, expert_capacity_dim]
-  outputs: [batch_dim, group_dim, output_dim]
+  expert_outputs: [<batch_dims>, experts_dim, expert_capacity_dim, output_dim]
+  combine_tensor:
+    [<batch_dims>, group_size_dim, experts_dim, expert_capacity_dim]
+  outputs: [<batch_dims>, group_size_dim, output_dim]
 
   Args:
-    inputs: a mtf.Tensor with shape [batch_dim, group_dim, input_dim]
+    inputs: a mtf.Tensor with shape [<batch_dims>, group_size_dim, input_dim]
+    outer_expert_dims: an optional list of dimensions.  This is for the case
+      where we are at an inner level of a hierarchical MoE.
     experts_dim: a Dimension (the number of experts)
     expert_capacity_dim: a Dimension (number of examples per group per expert)
-    max_experts: optional mtf.Tensor with shape [batch_dim, group_dim]
+    max_experts: optional mtf.Tensor with shape [<batch_dims>, group_size_dim]
     hparams: model hyperparameters.
     train: a boolean
 
   Returns:
-    forward_assignment: a Tensor with shape
-      [batch_dim, group_dim, experts_dim, expert_capacity_dim]
-    backward_assignment: a Tensor with shape
-      [batch_dim, group_dim, experts_dim, expert_capacity_dim]
+    dispatch_tensor: a Tensor with shape
+      [<batch_dims>, group_size_dim, experts_dim, expert_capacity_dim]
+    combine_tensor: a Tensor with shape
+      [<batch_dims>, group_size_dim, experts_dim, expert_capacity_dim]
     loss: a mtf scalar
 
   Raises:
     ValueError: on illegal hyperparameters
   """
-  unused_batch_dim, group_dim, unused_input_dim = inputs.shape.dims
+  group_size_dim, unused_input_dim = inputs.shape.dims[-2:]
 
   raw_gates = mtf.softmax(mtf_layers.dense(
-      inputs, experts_dim, use_bias=False), experts_dim)
+      inputs, experts_dim, use_bias=False,
+      expert_dims=outer_expert_dims), experts_dim)
 
   expert_capacity_f = float(expert_capacity_dim.size)
 
@@ -233,9 +463,9 @@ def _top_2_gating(
   # BALANCING LOSSES
   # shape = [batch, experts]
   # We want to equalize the fraction of the batch assigned to each expert
-  density_1 = mtf.reduce_mean(mask_1, reduced_dim=group_dim)
+  density_1 = mtf.reduce_mean(mask_1, reduced_dim=group_size_dim)
   # Something continuous that is correlated with what we want to equalize.
-  density_1_proxy = mtf.reduce_mean(raw_gates, reduced_dim=group_dim)
+  density_1_proxy = mtf.reduce_mean(raw_gates, reduced_dim=group_size_dim)
   density_1 = mtf.Print(
       density_1, [mtf.reduce_mean(density_1, output_shape=[experts_dim])],
       "density_1", summarize=1000)
@@ -246,12 +476,12 @@ def _top_2_gating(
     # Also add a loss to encourage all experts to be used equally also as the
     # second-place expert.  Experimentally, this seems to be a wash.
     # We want to equalize the fraction of the batch assigned to each expert:
-    density_2 = mtf.reduce_mean(mask_2, reduced_dim=group_dim)
+    density_2 = mtf.reduce_mean(mask_2, reduced_dim=group_size_dim)
     # As a proxy for density_2, we renormalize the raw gates after the top one
     # has been removed.
     normalized = gates_without_top_1 / (
         mtf.reduce_sum(gates_without_top_1, reduced_dim=experts_dim) + 1e-9)
-    density_2_proxy = mtf.reduce_mean(normalized, reduced_dim=group_dim)
+    density_2_proxy = mtf.reduce_mean(normalized, reduced_dim=group_size_dim)
     loss_2 = (mtf.reduce_mean(density_2_proxy * density_2)
               * float(experts_dim.size * experts_dim.size))
     loss += loss_2 * 0.5
@@ -287,12 +517,13 @@ def _top_2_gating(
   # COMPUTE ASSIGNMENT TO EXPERTS
   # [batch, group, experts]
   # This is the position within the expert's mini-batch for this sequence
-  position_in_expert_1 = mtf.cumsum(mask_1, group_dim, exclusive=True) * mask_1
+  position_in_expert_1 = mtf.cumsum(
+      mask_1, group_size_dim, exclusive=True) * mask_1
   # Remove the elements that don't fit. [batch, group, experts]
   mask_1 *= mtf.to_float(mtf.less(position_in_expert_1, expert_capacity_f))
   # [batch, experts]
   # How many examples in this sequence go to this expert
-  mask_1_count = mtf.reduce_sum(mask_1, reduced_dim=group_dim)
+  mask_1_count = mtf.reduce_sum(mask_1, reduced_dim=group_size_dim)
   # [batch, group] - mostly ones, but zeros where something didn't fit
   mask_1_flat = mtf.reduce_sum(mask_1, reduced_dim=experts_dim)
   # [batch, group]
@@ -303,7 +534,7 @@ def _top_2_gating(
 
   # [batch, group, experts]
   position_in_expert_2 = (
-      mtf.cumsum(mask_2, group_dim, exclusive=True) + mask_1_count)
+      mtf.cumsum(mask_2, group_size_dim, exclusive=True) + mask_1_count)
   position_in_expert_2 *= mask_2
   mask_2 *= mtf.to_float(mtf.less(position_in_expert_2, expert_capacity_f))
   # mask_2_count = mtf.reduce_sum(mask_2, reduced_dim=experts_dim)
@@ -318,7 +549,7 @@ def _top_2_gating(
   gate_2 /= denom
 
   # [batch, group, experts, expert_capacity]
-  backward_assignment = (
+  combine_tensor = (
       gate_1 * mask_1_flat
       * mtf.one_hot(index_1, experts_dim)
       * mtf.one_hot(mtf.to_int32(position_in_expert_1), expert_capacity_dim) +
@@ -326,10 +557,10 @@ def _top_2_gating(
       * mtf.one_hot(index_2, experts_dim)
       * mtf.one_hot(mtf.to_int32(position_in_expert_2), expert_capacity_dim))
 
-  forward_assignment = mtf.cast(
-      mtf.cast(backward_assignment, tf.bool), backward_assignment.dtype)
+  dispatch_tensor = mtf.cast(
+      mtf.cast(combine_tensor, tf.bool), combine_tensor.dtype)
 
-  return forward_assignment, backward_assignment, loss
+  return dispatch_tensor, combine_tensor, loss
 
 
 def set_default_moe_hparams(hparams):

From 7b7cac378a8939f8ffc64c1970a5eab7d7bec304 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Thu, 23 Aug 2018 14:42:29 -0700
Subject: [PATCH 0657/2720] Restructure all dimensions to be declared in one
 place and add some comments

PiperOrigin-RevId: 210003442
---
 tensor2tensor/mesh_tensorflow/research/moe.py | 46 ++++++++++---------
 1 file changed, 25 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/research/moe.py b/tensor2tensor/mesh_tensorflow/research/moe.py
index 917d2ad6c..67439593a 100644
--- a/tensor2tensor/mesh_tensorflow/research/moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/moe.py
@@ -256,8 +256,6 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
   """
   assert len(hparams.moe_num_experts) == 2
   a0, b1, l, m = inputs.shape.dims
-  # a = mtf.Dimension("a_unsplit", a0.size)
-  # b = mtf.Dimension("b_unsplit", b1.size)
   hidden_dim = mtf.Dimension("expert_hidden", hparams.moe_hidden_size)
   x1 = mtf.Dimension("expert_x", hparams.moe_num_experts[0])
   y0 = mtf.Dimension("expert_y", hparams.moe_num_experts[1])
@@ -265,13 +263,8 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
   y = mtf.Dimension("expert_y_unsplit", hparams.moe_num_experts[1])
   s = mtf.Dimension("group_size_x", hparams.moe_group_size[0])
   t = mtf.Dimension("group_size_y", hparams.moe_group_size[1])
-  g1 = mtf.Dimension(b1.name, b1.size * l.size // s.size)
-  g = mtf.Dimension(b1.name + "_unsplit", b1.size * l.size // s.size)
   n = output_dim
 
-  # First level of expert routing
-  inputs = mtf.reshape(inputs, [a0, g1, s, m])
-
   # Each sequence sends (at most?) expert_capacity positions to each expert.
   # Static expert_capacity dimension is needed for expert batch sizes
   capacity_factor = (
@@ -281,6 +274,24 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
       s.size,
       int((s.size * capacity_factor) / x.size))
   c = mtf.Dimension("expert_capacity_x", expert_capacity)
+  expert_capacity = min(
+      t.size,
+      int((t.size * capacity_factor) / y.size))
+  d = mtf.Dimension("expert_capacity_y", expert_capacity)
+
+  g1 = mtf.Dimension(b1.name, b1.size * l.size // s.size)
+  g = mtf.Dimension(b1.name + "_unsplit", b1.size * l.size // s.size)
+
+  numerator = a0.size * g.size * c.size
+  if numerator % t.size != 0:
+    raise ValueError("cannot divide evenly %s / %s" % (numerator, t.size))
+  h0 = mtf.Dimension(a0.name, numerator // t.size)
+  h = mtf.Dimension(a0.name + "_unsplit", h0.size)
+
+  # First level of expert routing
+  # Reshape the inner batch size to a multiple of group_dim g1 and
+  # group_size_dim s.
+  inputs = mtf.reshape(inputs, [a0, g1, s, m])
 
   # Get the assignments for the first level.
   # dispatch_tensor_x has shape [a0, g1, s, x, c]
@@ -306,19 +317,10 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
       [x1, a0, g, c, m]))
 
   # Second level of expert routing
-
-  numerator = a0.size * g.size * c.size
-  if numerator % t.size != 0:
-    raise ValueError("cannont divide evenly %s / %s" % (numerator, t.size))
-  h0 = mtf.Dimension(a0.name, numerator // t.size)
-  h = mtf.Dimension(a0.name + "_unsplit", h0.size)
+  # Reshape the expert_inputs outer batch dim to be a multiple of group_dim h0
+  # and group_size_dim t.
   inputs_y = mtf.reshape(expert_inputs_x, [x1, h0, t, m])
 
-  expert_capacity = min(
-      t.size,
-      int((t.size * capacity_factor) / y.size))
-  d = mtf.Dimension("expert_capacity_y", expert_capacity)
-
   # Get the assignments for the second level.
   # dispatch_tensor_x has shape [x1, h0, t, y, d]
   if hparams.moe_gating == "top_2":
@@ -361,17 +363,19 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
   output_y = mtf.einsum([expert_output, combine_tensor_y], mtf.Shape(
       [x1, h0, t, n]))
 
-  # simple reshape
+  # Reshape the combined tensor from inner level to now contain outer_batch_dim
+  # a0 and group_dim g
   output = mtf.reshape(output_y, [x1, a0, g, c, n])
 
-  # alltoall
+  # alltoall from expert_dim x to group_dim g1
   expert_output_x = mtf.reshape(output, mtf.Shape([x, a0, g1, c, n]))
 
   # combine results from outer level
   output_x = mtf.einsum([expert_output_x, combine_tensor_x], mtf.Shape(
       [a0, g1, s, m]))
 
-  # simple reshape
+  # Reshape the combined tensor to now contain inner_batch_dim
+  # b1 and the original sequence length
   output = mtf.reshape(output_x, [a0, b1, l, n])
   return output, (loss_outer + loss_inner) * hparams.moe_loss_coef
 

From 7032c14d87a5c4a54bfcfa81adaa67401ee72138 Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Thu, 23 Aug 2018 14:44:17 -0700
Subject: [PATCH 0658/2720] Trying a new discriminator

PiperOrigin-RevId: 210003709
---
 tensor2tensor/data_generators/celeba.py       |  34 +++++
 tensor2tensor/layers/common_layers.py         | 139 +++++++++++++++++-
 tensor2tensor/models/research/autoencoders.py | 106 ++++++-------
 3 files changed, 212 insertions(+), 67 deletions(-)

diff --git a/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index cf068f040..7c75f8d69 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -236,3 +236,37 @@ def preprocess_example(self, example, unused_mode, unused_hparams):
     example["inputs"] = image_8
     example["targets"] = image_64
     return example
+
+
+@registry.register_problem
+class ImageCeleba32(Img2imgCeleba):
+  """CelebA resized to spatial dims [32, 32]."""
+
+  def preprocess_example(self, example, unused_mode, unused_hparams):
+    image = example["inputs"]
+    # Remove boundaries in CelebA images. Remove 40 pixels each side
+    # vertically and 20 pixels each side horizontally.
+    image = tf.image.crop_to_bounding_box(image, 40, 20, 218 - 80, 178 - 40)
+    image = image_utils.resize_by_area(image, 32)
+
+    example["inputs"] = image
+    example["targets"] = image
+    return example
+
+
+@registry.register_problem
+class ImageCeleba64(Img2imgCeleba):
+  """CelebA resized to spatial dims [64, 64]."""
+
+  def preprocess_example(self, example, unused_mode, unused_hparams):
+    image = example["inputs"]
+    # Remove boundaries in CelebA images. Remove 40 pixels each side
+    # vertically and 20 pixels each side horizontally.
+    image = tf.image.crop_to_bounding_box(image, 40, 20, 218 - 80, 178 - 40)
+    image = image_utils.resize_by_area(image, 64)
+
+    example["inputs"] = image
+    example["targets"] = image
+    return example
+
+
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 5bae4ea80..01f477b98 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3387,8 +3387,7 @@ def belu(x):
 
 def nac(x, depth, name=None, reuse=None):
   """NAC as in https://arxiv.org/abs/1808.00508."""
-  with tf.variable_scope(
-      name, default_name="nac", values=[x], reuse=reuse):
+  with tf.variable_scope(name, default_name="nac", values=[x], reuse=reuse):
     x_shape = shape_list(x)
     w = tf.get_variable("w", [x_shape[-1], depth])
     m = tf.get_variable("m", [x_shape[-1], depth])
@@ -3400,8 +3399,7 @@ def nac(x, depth, name=None, reuse=None):
 
 def nalu(x, depth, epsilon=1e-30, name=None, reuse=None):
   """NALU as in https://arxiv.org/abs/1808.00508."""
-  with tf.variable_scope(
-      name, default_name="nalu", values=[x], reuse=reuse):
+  with tf.variable_scope(name, default_name="nalu", values=[x], reuse=reuse):
     x_shape = shape_list(x)
     x_flat = tf.reshape(x, [-1, x_shape[-1]])
     gw = tf.get_variable("w", [x_shape[-1], depth])
@@ -3671,6 +3669,133 @@ def get_sorted_projections(x):
     return dist
 
 
+def lrelu(input_, leak=0.2, name="lrelu"):
+  return tf.maximum(input_, leak * input_, name=name)
+
+
+def deep_discriminator(x,
+                       batch_norm,
+                       is_training,
+                       filters=64,
+                       filter_size=4,
+                       stride=2,
+                       output_size=1024):
+  """Discriminator architecture based on InfoGAN."""
+  with tf.variable_scope(
+      "discriminator", initializer=tf.random_normal_initializer(stddev=0.02)):
+    batch_size, height, width = shape_list(x)[:3]
+    net = tf.layers.conv2d(
+        x, filters, filter_size, strides=stride, padding="SAME", name="d_conv1")
+    net = lrelu(net)
+    net = tf.layers.conv2d(
+        net,
+        2 * filters,
+        filter_size,
+        strides=stride,
+        padding="SAME",
+        name="d_conv2")
+    # [bs, h/4, w/4, 128]
+    if batch_norm:
+      net = tf.layers.batch_normalization(
+          net, training=is_training, momentum=0.999, name="d_bn2")
+    net = lrelu(net)
+    size = height * width
+    x_shape = x.get_shape().as_list()
+    if x_shape[1] is None or x_shape[2] is None:
+      net = tf.reduce_mean(net, axis=[1, 2])
+    else:
+      net = tf.reshape(net, [batch_size, size * 8])
+    net = tf.layers.dense(net, output_size, name="d_fc3")
+    if batch_norm:
+      net = tf.layers.batch_normalization(
+          net, training=is_training, momentum=0.999, name="d_bn3")
+    net = lrelu(net)
+    return net
+
+
+def instance_norm(x):
+  """Instance normalization layer."""
+  with tf.variable_scope("instance_norm"):
+    epsilon = 1e-5
+    mean, var = tf.nn.moments(x, [1, 2], keep_dims=True)
+    scale = tf.get_variable(
+        "scale", [x.get_shape()[-1]],
+        initializer=tf.truncated_normal_initializer(mean=1.0, stddev=0.02))
+    offset = tf.get_variable(
+        "offset", [x.get_shape()[-1]], initializer=tf.constant_initializer(0.0))
+    out = scale * tf.div(x - mean, tf.sqrt(var + epsilon)) + offset
+
+    return out
+
+
+def general_conv(x,
+                 num_filters=64,
+                 filter_size=7,
+                 stride=1,
+                 stddev=0.02,
+                 padding="VALID",
+                 name="conv",
+                 do_norm="instance",
+                 do_relu=True,
+                 relufactor=0):
+  """Generalized convolution layer."""
+  with tf.variable_scope(name):
+    x = tf.layers.conv2d(
+        x,
+        num_filters,
+        filter_size,
+        stride,
+        padding,
+        activation=None,
+        kernel_initializer=tf.truncated_normal_initializer(stddev=stddev),
+        bias_initializer=tf.constant_initializer(0.0))
+    if do_norm == "layer":
+      x = tf.contrib.layers.layer_norm(x)
+    elif do_norm == "instance":
+      x = instance_norm(x)
+
+    if do_relu:
+      if relufactor == 0:
+        x = tf.nn.relu(x, "relu")
+      else:
+        x = lrelu(x, leak=relufactor)
+
+    return x
+
+
+def patch_discriminator(x, filters=64, filter_size=5, n=4,
+                        name="patch_discrim"):
+  """Patch descriminator."""
+  with tf.variable_scope(name):
+    x_shape = shape_list(x)
+    spatial_dims = [x_shape[1] // 4, x_shape[2] // 4]
+    x = tf.random_crop(x, [x_shape[0]] + spatial_dims + [x_shape[3]])
+    for i in range(n):
+      x = general_conv(
+          x=x,
+          num_filters=filters * 2**i,
+          filter_size=filter_size,
+          stride=2 if i != n - 1 else 1,
+          stddev=0.02,
+          padding="SAME",
+          name="c%d" % i,
+          do_norm="instance" if i != 0 else False,
+          do_relu=i != n - 1,
+          relufactor=0.2)
+    x = tf.reduce_mean(x, [1, 2])
+    return x
+
+
+def simple_discriminator(x, filters=128, filter_size=7, stride=4):
+  """A very simple convolutional discriminator."""
+  with tf.variable_scope("discriminator"):
+    net = tf.layers.conv2d(
+        x, filters, filter_size, strides=stride, padding="SAME", name="d_conv1")
+    net = tf.nn.relu(net)
+    net = tf.reduce_mean(net, [1, 2])
+    return net
+
+
 def upscale(inputs, f, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR):
   """Upscaling the image by a factor of f."""
   height, width = shape_list(inputs)[1:3]
@@ -4026,10 +4151,8 @@ def build(self, input_shape=None):
       self.layer.built = False
 
       if not hasattr(self.layer, "kernel"):
-        raise ValueError(
-            "`WeightNorm` must wrap a layer that"
-            " contains a `kernel` for weights"
-        )
+        raise ValueError("`WeightNorm` must wrap a layer that"
+                         " contains a `kernel` for weights")
 
       # The kernel's filter or unit dimension is -1
       self.layer_depth = int(self.layer.kernel.shape[-1])
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 7ccc318bd..6e69ca72a 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -29,10 +29,6 @@
 import tensorflow as tf
 
 
-def lrelu(input_, leak=0.2, name="lrelu"):
-  return tf.maximum(input_, leak * input_, name=name)
-
-
 def reverse_gradient(x, lr=1.0):
   return -lr * x + tf.stop_gradient((1.0 + lr) * x)
 
@@ -89,45 +85,6 @@ def bottleneck(self, x):
         return tf.tanh(x) + noise * hparams.bottleneck_noise, 0.0
       return tf.tanh(x), 0.0
 
-  def discriminator(self, x, is_training):
-    """Discriminator architecture based on InfoGAN.
-
-    Args:
-      x: input images, shape [bs, h, w, channels]
-      is_training: boolean, are we in train or eval model.
-
-    Returns:
-      out_logit: the output logits (before sigmoid).
-    """
-    hparams = self.hparams
-    with tf.variable_scope(
-        "discriminator", initializer=tf.random_normal_initializer(stddev=0.02)):
-      batch_size, height, width = common_layers.shape_list(x)[:3]
-      # Mapping x from [bs, h, w, c] to [bs, 1]
-      net = tf.layers.conv2d(
-          x, 64, (4, 4), strides=(2, 2), padding="SAME", name="d_conv1")
-      # [bs, h/2, w/2, 64]
-      net = lrelu(net)
-      net = tf.layers.conv2d(
-          net, 128, (4, 4), strides=(2, 2), padding="SAME", name="d_conv2")
-      # [bs, h/4, w/4, 128]
-      if hparams.discriminator_batchnorm:
-        net = tf.layers.batch_normalization(
-            net, training=is_training, momentum=0.999, name="d_bn2")
-      net = lrelu(net)
-      size = height * width
-      x_shape = x.get_shape().as_list()
-      if x_shape[1] is None or x_shape[2] is None:
-        net = tf.reduce_mean(net, axis=[1, 2])  # [bs, 128]
-      else:
-        net = tf.reshape(net, [batch_size, size * 8])  # [bs, h * w * 8]
-      net = tf.layers.dense(net, 1024, name="d_fc3")  # [bs, 1024]
-      if hparams.discriminator_batchnorm:
-        net = tf.layers.batch_normalization(
-            net, training=is_training, momentum=0.999, name="d_bn3")
-      net = lrelu(net)
-      return net
-
   def unbottleneck(self, x, res_size, reuse=None):
     with tf.variable_scope("unbottleneck", reuse=reuse):
       x = tf.layers.dense(x, res_size, name="dense")
@@ -194,8 +151,7 @@ def gumbel_sample(self, reconstr_gan):
       reconstr_gan += gumbel_samples
       reconstr_sample = latent_layers.multinomial_sample(
           reconstr_gan, temperature=hparams.gumbel_temperature)
-      reconstr_gan = tf.nn.softmax(
-          reconstr_gan / hparams.gumbel_temperature)
+      reconstr_gan = tf.nn.softmax(reconstr_gan / hparams.gumbel_temperature)
     else:
       reconstr_sample = tf.argmax(reconstr_gan, axis=-1)
       reconstr_gan = tf.nn.softmax(reconstr_gan / 0.1)  # Sharpen a bit.
@@ -251,7 +207,8 @@ def body(self, features):
         # Add a purely sampled batch on which we'll compute the GAN loss.
         g = self.unbottleneck(
             self.sample(shape=b_shape),
-            common_layers.shape_list(x)[-1], reuse=True)
+            common_layers.shape_list(x)[-1],
+            reuse=True)
         x = tf.concat([g, x], axis=0)
         encoder_layers = [tf.concat([l, l], axis=0) for l in encoder_layers]
     else:
@@ -281,8 +238,7 @@ def body(self, features):
 
     if hparams.mode == tf.estimator.ModeKeys.PREDICT:
       if hparams.use_vq_loss:
-        (reconstr, _, _, _, _) = discretization.vq_loss(
-            res, labels, vocab_size)
+        (reconstr, _, _, _, _) = discretization.vq_loss(res, labels, vocab_size)
       else:
         reconstr = tf.layers.dense(res, vocab_size, name="autoencoder_final")
       return reconstr, {"bottleneck_loss": 0.0}
@@ -291,8 +247,10 @@ def body(self, features):
       res_gan, res = tf.split(res, 2, axis=0)
 
     # Losses.
-    losses = {"bottleneck_extra": b_loss,
-              "bottleneck_l2": hparams.bottleneck_l2_factor * xb_loss}
+    losses = {
+        "bottleneck_extra": b_loss,
+        "bottleneck_l2": hparams.bottleneck_l2_factor * xb_loss
+    }
 
     if hparams.use_vq_loss:
       vq_temperature = hparams.vq_temperature / common_layers.inverse_exp_decay(
@@ -320,7 +278,10 @@ def body(self, features):
         with tf.variable_scope("vq_loss", reuse=True):
           update_means = tf.less(tf.random_uniform([]), update_means_factor)
           reconstr_gan, gan_codes, _, code_loss_gan, _ = discretization.vq_loss(
-              res_gan, labels, vocab_size, do_update=update_means,
+              res_gan,
+              labels,
+              vocab_size,
+              do_update=update_means,
               temperature=vq_temperature)
           code_loss_gan *= hparams.code_loss_factor * update_means_factor
           losses["code_loss_gan"] = code_loss_gan
@@ -337,7 +298,15 @@ def body(self, features):
       self.image_summary("gan", reconstr_gan)
 
       def discriminate(x):
-        return self.discriminator(x, is_training=is_training)
+        if hparams.discriminator == "default":
+          return common_layers.deep_discriminator(
+              x, hparams.discriminator_batchnorm, is_training)
+        elif hparams.discriminator == "patched":
+          return common_layers.patch_discriminator(x)
+        elif hparams.discriminator == "simple":
+          return common_layers.simple_discriminator(x)
+        else:
+          raise Exception("Unknown discriminator %s" % hparams.discriminator)
 
       tc_shape = common_layers.shape_list(target_codes)
       if len(tc_shape) > 4:
@@ -348,10 +317,10 @@ def discriminate(x):
       gan_lr = common_layers.inverse_exp_decay(
           hparams.gan_codes_warmup_steps * 1.5)
       rev_grad_gan_codes = reverse_gradient(gan_codes, lr=gan_lr)
-      gan_loss = common_layers.sliced_gan_loss(
-          target_codes, rev_grad_gan_codes, discriminate,
-          self.hparams.num_sliced_vecs)
-      gan_loss *= hparams.gan_loss_factor  * update_means_factor
+      gan_loss = common_layers.sliced_gan_loss(target_codes, rev_grad_gan_codes,
+                                               discriminate,
+                                               self.hparams.num_sliced_vecs)
+      gan_loss *= hparams.gan_loss_factor * update_means_factor
       losses["gan_loss"] = -gan_loss
 
     self.image_summary("ae", reconstr)
@@ -705,8 +674,10 @@ def sample(self, features=None, shape=None):
     hparams = self.hparams
     div_x = 2**hparams.num_hidden_layers
     div_y = 1 if self.is1d else 2**hparams.num_hidden_layers
-    size = [hparams.batch_size, hparams.sample_height // div_x,
-            hparams.sample_width // div_y, hparams.bottleneck_bits]
+    size = [
+        hparams.batch_size, hparams.sample_height // div_x,
+        hparams.sample_width // div_y, hparams.bottleneck_bits
+    ]
     size = size if shape is None else shape
     return tf.random_normal(size)
 
@@ -927,7 +898,7 @@ def autoencoder_basic():
   hparams.add_hparam("sample_height", 32)
   hparams.add_hparam("sample_width", 32)
   hparams.add_hparam("discriminator_batchnorm", True)
-  hparams.add_hparam("num_sliced_vecs", 4096)
+  hparams.add_hparam("num_sliced_vecs", 20000)
   hparams.add_hparam("code_loss_factor", 1.0)
   hparams.add_hparam("gan_codes_warmup_steps", 6000)
   hparams.add_hparam("gan_loss_factor", 0.0)
@@ -936,6 +907,7 @@ def autoencoder_basic():
   hparams.add_hparam("gumbel_noise_factor", 0.4)
   hparams.add_hparam("vq_temperature", 0.001)
   hparams.add_hparam("use_vq_loss", int(False))
+  hparams.add_hparam("discriminator", "default")
   return hparams
 
 
@@ -1038,6 +1010,22 @@ def autoencoder_ordered_discrete():
   return hparams
 
 
+@registry.register_hparams
+def autoencoder_ordered_discrete_patched():
+  """Ordered discrete autoencoder model."""
+  hparams = autoencoder_ordered_discrete()
+  hparams.discriminator = "patched"
+  return hparams
+
+
+@registry.register_hparams
+def autoencoder_ordered_discrete_simple():
+  """Ordered discrete autoencoder model."""
+  hparams = autoencoder_ordered_discrete()
+  hparams.discriminator = "simple"
+  return hparams
+
+
 @registry.register_hparams
 def autoencoder_ordered_discrete_hs256():
   """Ordered discrete autoencoder model."""

From 73e3732301c227a5440492bba7262bffc6a35ac3 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Thu, 23 Aug 2018 15:14:33 -0700
Subject: [PATCH 0659/2720] figure out group sizes automatically to be divisors
 of what they need

PiperOrigin-RevId: 210009228
---
 .../research/experiments_moe.py               | 18 +++++++--
 tensor2tensor/mesh_tensorflow/research/moe.py | 38 +++++++++++++------
 2 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
index f95c2df02..5cee29a7b 100644
--- a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
@@ -37,9 +37,9 @@ def xmoe_dense_4k():
   model             params(M)  einsum  alltoall  mxu-util  log-ppl(1ep) (3ep)
   xmoe_dense_4k     30         3.0e12  0         45%        3.31
   xmoe_dense_8k     46         4.7e12  0         49%        3.24
-  xmoe_dense_64k                       0                    3.06
-  xmoe_top_2        282        4.0e12  3.4e8     36%
-  xmoe_top_2_c15    282        4.5e12  4.0e8     38%
+  xmoe_dense_64k    282        2.8e13  0                    3.06
+  xmoe_top_2        282        4.0e12  3.4e8     36%        3.07
+  xmoe_top_2_c15    282        4.5e12  4.0e8     38%        3.07
 
   Note: configurations and code are likely to change without notice.
 
@@ -131,8 +131,18 @@ def xmoe_2d():
   hparams.outer_batch_size = 4
   hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
   hparams.moe_num_experts = [4, 4]
-  hparams.moe_group_size = [256, 256]
   hparams.feedforward_layer = "hmoe"
   return hparams
 
 
+@registry.register_hparams
+def xmoe_2d_88():
+  """Two-dimensional hierarchical mixture of experts."""
+  hparams = xmoe_2d()
+  hparams.mesh_shape = "b0:4;b1:8"
+  hparams.batch_size = 512
+  hparams.outer_batch_size = 4
+  hparams.moe_num_experts = [8, 8]
+  return hparams
+
+
diff --git a/tensor2tensor/mesh_tensorflow/research/moe.py b/tensor2tensor/mesh_tensorflow/research/moe.py
index 67439593a..e09f262ac 100644
--- a/tensor2tensor/mesh_tensorflow/research/moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/moe.py
@@ -254,6 +254,11 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
   Raises:
     ValueError: on unrecognized hparams.moe_gating
   """
+  insert_outer_batch_dim = (len(inputs.shape.dims) == 3)
+  if insert_outer_batch_dim:
+    inputs = mtf.reshape(
+        inputs, [mtf.Dimension("outer_batch", 1)] + inputs.shape.dims)
+
   assert len(hparams.moe_num_experts) == 2
   a0, b1, l, m = inputs.shape.dims
   hidden_dim = mtf.Dimension("expert_hidden", hparams.moe_hidden_size)
@@ -261,10 +266,14 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
   y0 = mtf.Dimension("expert_y", hparams.moe_num_experts[1])
   x = mtf.Dimension("expert_x_unsplit", hparams.moe_num_experts[0])
   y = mtf.Dimension("expert_y_unsplit", hparams.moe_num_experts[1])
-  s = mtf.Dimension("group_size_x", hparams.moe_group_size[0])
-  t = mtf.Dimension("group_size_y", hparams.moe_group_size[1])
   n = output_dim
 
+  numerator = b1.size * l.size
+  s = mtf.Dimension("group_size_x", _largest_divisor_leq(
+      numerator, hparams.moe_group_size))
+  g1 = mtf.Dimension(b1.name, numerator // s.size)
+  g = mtf.Dimension(b1.name + "_unsplit", g1.size)
+
   # Each sequence sends (at most?) expert_capacity positions to each expert.
   # Static expert_capacity dimension is needed for expert batch sizes
   capacity_factor = (
@@ -274,20 +283,18 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
       s.size,
       int((s.size * capacity_factor) / x.size))
   c = mtf.Dimension("expert_capacity_x", expert_capacity)
-  expert_capacity = min(
-      t.size,
-      int((t.size * capacity_factor) / y.size))
-  d = mtf.Dimension("expert_capacity_y", expert_capacity)
-
-  g1 = mtf.Dimension(b1.name, b1.size * l.size // s.size)
-  g = mtf.Dimension(b1.name + "_unsplit", b1.size * l.size // s.size)
 
   numerator = a0.size * g.size * c.size
-  if numerator % t.size != 0:
-    raise ValueError("cannot divide evenly %s / %s" % (numerator, t.size))
+  t = mtf.Dimension("group_size_y", _largest_divisor_leq(
+      numerator, hparams.moe_group_size))
   h0 = mtf.Dimension(a0.name, numerator // t.size)
   h = mtf.Dimension(a0.name + "_unsplit", h0.size)
 
+  expert_capacity = min(
+      t.size,
+      int((t.size * capacity_factor) / y.size))
+  d = mtf.Dimension("expert_capacity_y", expert_capacity)
+
   # First level of expert routing
   # Reshape the inner batch size to a multiple of group_dim g1 and
   # group_size_dim s.
@@ -377,6 +384,8 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
   # Reshape the combined tensor to now contain inner_batch_dim
   # b1 and the original sequence length
   output = mtf.reshape(output_x, [a0, b1, l, n])
+  if insert_outer_batch_dim:
+    output = mtf.reshape(output, [b1, l, n])
   return output, (loss_outer + loss_inner) * hparams.moe_loss_coef
 
 
@@ -598,3 +607,10 @@ def set_default_moe_hparams(hparams):
   hparams.add_hparam("moe_second_policy_eval", "random")
   hparams.add_hparam("moe_second_threshold_train", 0.2)
   hparams.add_hparam("moe_second_threshold_eval", 0.2)
+
+
+def _largest_divisor_leq(numerator, maximum):
+  x = maximum
+  while numerator % x != 0:
+    x -= 1
+  return x

From 7c895d935e87958cf9ce0dfb3550db55ef7aec9d Mon Sep 17 00:00:00 2001
From: jurasofish <33386122+jurasofish@users.noreply.github.com>
Date: Fri, 24 Aug 2018 10:03:32 +1000
Subject: [PATCH 0660/2720] Update universal_transformer_util.py (#1017)

---
 tensor2tensor/models/research/universal_transformer_util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 82747e0b3..4c2e76fb0 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -197,7 +197,7 @@ def universal_transformer_layer(x,
                                 ffn_unit,
                                 attention_unit,
                                 pad_remover=None):
-  """Core function applying the universal transforemr layer.
+  """Core function applying the universal transformer layer.
 
   Args:
     x: input

From 3c1af812547fa0d520232c5240ce483a4ad2933b Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 23 Aug 2018 17:39:21 -0700
Subject: [PATCH 0661/2720] infer func for SV2P so it runs in PPO.

PiperOrigin-RevId: 210031063
---
 tensor2tensor/layers/common_video.py          |  9 ++++-
 .../models/research/next_frame_sv2p.py        | 28 +++++++++++++-
 .../models/research/next_frame_test.py        |  2 +
 tensor2tensor/rl/trainer_model_based.py       | 13 ++++++-
 .../rl/trainer_model_based_stochastic_test.py |  2 +-
 .../rl/trainer_model_based_sv2p_test.py       | 37 +++++++++++++++++++
 6 files changed, 84 insertions(+), 7 deletions(-)
 create mode 100644 tensor2tensor/rl/trainer_model_based_sv2p_test.py

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 0698387a1..a89567f16 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -102,8 +102,13 @@ def scheduled_sample_count(ground_truth_x,
 
   ground_truth_examps = tf.gather(ground_truth_x, ground_truth_idx)
   generated_examps = tf.gather(generated_x, generated_idx)
-  return tf.dynamic_stitch([ground_truth_idx, generated_idx],
-                           [ground_truth_examps, generated_examps])
+
+  output = tf.dynamic_stitch([ground_truth_idx, generated_idx],
+                             [ground_truth_examps, generated_examps])
+  # if batch size is known set it.
+  if isinstance(batch_size, int):
+    output.set_shape([batch_size] + common_layers.shape_list(output)[1:])
+  return output
 
 
 def scheduled_sample_prob(ground_truth_x,
diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 4762663a4..800730abc 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -479,6 +479,29 @@ def get_extra_loss(self, latent_means=None, latent_stds=None,
       tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
     return beta * kl_loss
 
+  def infer(self, features, *args, **kwargs):
+    """Produce predictions from the model by running it."""
+    del args, kwargs
+    if "targets" not in features:
+      if "infer_targets" in features:
+        targets_shape = common_layers.shape_list(features["infer_targets"])
+      elif "inputs" in features:
+        targets_shape = common_layers.shape_list(features["inputs"])
+        targets_shape[1] = self.hparams.video_num_target_frames
+      else:
+        raise ValueError("no inputs are given.")
+      features["targets"] = tf.zeros(targets_shape, dtype=tf.float32)
+
+    output, _ = self(features)  # pylint: disable=not-callable
+
+    output["targets"] = tf.squeeze(output["targets"], axis=-1)
+    output["target_reward"] = tf.argmax(output["target_reward"], axis=-1)
+
+    # only required for decoding.
+    output["outputs"] = output["targets"]
+    output["scores"] = output["targets"]
+    return output
+
   def body(self, features):
     hparams = self.hparams
     batch_size = common_layers.shape_list(features["inputs"])[0]
@@ -528,14 +551,15 @@ def body(self, features):
     # This is NOT the same as original paper/implementation.
     predictions = gen_images[hparams.video_num_input_frames-1:]
     reward_pred = gen_rewards[hparams.video_num_input_frames-1:]
-    reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove undeeded dimension.
+    if self.is_training:
+      reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove extra dimension.
 
     # Swap back time and batch axes.
     predictions = common_video.swap_time_and_batch_axes(predictions)
     reward_pred = common_video.swap_time_and_batch_axes(reward_pred)
 
     return_targets = predictions
-    if "target_reward" in features:
+    if hparams.reward_prediction:
       return_targets = {"targets": predictions, "target_reward": reward_pred}
 
     if hparams.internal_loss:
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/research/next_frame_test.py
index 73ad1ea31..13f69e766 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/research/next_frame_test.py
@@ -41,6 +41,7 @@ def fill_hparams(hparams, in_frames, out_frames):
   hparams.problem = problem
   hparams.problem_hparams = p_hparams
   hparams.tiny_mode = True
+  hparams.reward_prediction = False
   return hparams
 
 
@@ -155,6 +156,7 @@ def TestVideoModelWithActionAndRewards(self,
                                          expected_last_dim):
     hparams = fill_hparams(hparams, in_frames, out_frames)
     hparams = full_modalities(hparams)
+    hparams.reward_prediction = True
 
     features = create_full_features(in_frames, out_frames)
 
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index d1717da7d..b8ad0bf8d 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -609,8 +609,17 @@ def rl_modelrl_tiny():
 def rl_modelrl_tiny_stochastic():
   """Tiny setting with a stochastic next-frame model."""
   hparams = rl_modelrl_tiny()
-  hparams.generative_model = "next_frame_stochastic"
-  hparams.generative_model_params = "next_frame_stochastic_tiny"
+  hparams.generative_model = "next_frame_basic_stochastic"
+  hparams.generative_model_params = "next_frame_basic_stochastic"
+  return hparams
+
+
+@registry.register_hparams
+def rl_modelrl_tiny_sv2p():
+  """Tiny setting with a sv2p model."""
+  hparams = rl_modelrl_tiny()
+  hparams.generative_model = "next_frame_sv2p"
+  hparams.generative_model_params = "next_frame_sv2p_tiny"
   return hparams
 
 
diff --git a/tensor2tensor/rl/trainer_model_based_stochastic_test.py b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
index d55951256..24f69df2a 100644
--- a/tensor2tensor/rl/trainer_model_based_stochastic_test.py
+++ b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
@@ -26,7 +26,7 @@
 
 class ModelRLExperimentStochasticTest(tf.test.TestCase):
 
-  def test_stochastic(self):
+  def test_basic_stochastic(self):
     FLAGS.output_dir = tf.test.get_temp_dir()
     FLAGS.loop_hparams_set = "rl_modelrl_tiny_stochastic"
     FLAGS.schedule = "train"  # skip evaluation for world model training
diff --git a/tensor2tensor/rl/trainer_model_based_sv2p_test.py b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
new file mode 100644
index 000000000..1b5f76a61
--- /dev/null
+++ b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
@@ -0,0 +1,37 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Tiny run of trainer_model_based with stochastic model. Smoke test."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.rl import trainer_model_based
+
+import tensorflow as tf
+
+FLAGS = tf.flags.FLAGS
+
+
+class ModelRLExperimentSv2pTest(tf.test.TestCase):
+
+  def test_sv2p(self):
+    FLAGS.output_dir = tf.test.get_temp_dir()
+    FLAGS.loop_hparams_set = "rl_modelrl_tiny_sv2p"
+    FLAGS.schedule = "train"
+    trainer_model_based.main(None)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From ed0f9c18472ee6d839cd62f7df20a35ffbcf360b Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Thu, 23 Aug 2018 17:57:35 -0700
Subject: [PATCH 0662/2720] Fix optimizer issue with TPUs and create base
 hparams for TPU

PiperOrigin-RevId: 210033389
---
 tensor2tensor/models/image_transformer.py | 62 ++++++++++++++++-------
 1 file changed, 44 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 0c894bc27..af34960e9 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -298,6 +298,39 @@ def imagetransformer_cifar10_base_dmol():
   return hparams
 
 
+@registry.register_hparams
+def imagetransformer_base_tpu():
+  """Transformer base params for cifar-10."""
+  hparams = imagetransformer_bas8l_8h_big_uncond_dr03_imgnet()
+  update_hparams_for_tpu(hparams)
+  hparams.batch_size = 4
+  hparams.num_heads = 4   # heads are expensive on tpu
+  hparams.num_decoder_layers = 12
+  hparams.block_length = 128
+  hparams.hidden_size = 512
+  hparams.filter_size = 2048
+  hparams.learning_rate = 0.2
+  hparams.learning_rate_warmup_steps = 6000
+  hparams.layer_preprocess_sequence = "none"
+  hparams.layer_postprocess_sequence = "dan"
+  hparams.layer_prepostprocess_dropout = 0.3
+  return hparams
+
+
+@registry.register_hparams
+def imagetransformer_base_imagenet_tpu():
+  """Transformer base params for cifar-10."""
+  hparams = imagetransformer_base_tpu()
+  hparams.batch_size = 4
+  hparams.num_heads = 4   # heads are expensive on tpu
+  hparams.num_decoder_layers = 12
+  hparams.block_length = 128
+  hparams.layer_preprocess_sequence = "none"
+  hparams.layer_postprocess_sequence = "dan"
+  hparams.layer_prepostprocess_dropout = 0.1
+  return hparams
+
+
 @registry.register_hparams
 def imagetransformer_imagenet32_base():
   """Best config for ImageNet-32 with 3.77 bits/dim using cross entropy."""
@@ -897,23 +930,10 @@ def imagetransformer_moe_tiny():
 
 
 def update_hparams_for_tpu(hparams):
-  hparams.optimizer = "TrueAdam"
-  hparams.batch_size = 4
-
-
-@registry.register_hparams
-def imagetransformer_base_tpu():
-  """Transformer base params for cifar-10."""
-  hparams = imagetransformer_bas8l_8h_big_uncond_dr03_imgnet()
-  update_hparams_for_tpu(hparams)
+  hparams.optimizer = "Adafactor"
+  hparams.learning_rate_schedule = "rsqrt_decay"
+  hparams.learning_rate_warmup_steps = 6000
   hparams.batch_size = 4
-  hparams.num_heads = 4   # heads are expensive on tpu
-  hparams.num_decoder_layers = 12
-  hparams.block_length = 128
-  hparams.layer_preprocess_sequence = "none"
-  hparams.layer_postprocess_sequence = "dan"
-  hparams.layer_prepostprocess_dropout = 0.3
-  return hparams
 
 
 @registry.register_hparams
@@ -1061,11 +1081,14 @@ def imagetransformer_b12l_4h_b128_h512_uncond_dr03_tpu():
 
 
 @registry.register_hparams
-def imagetransformer_b12l_4h_b128_h512_uncond_dr03_im():
+def imagetransformer_b12l_4h_b128_h512_uncond_dr01_im():
   """TPU related imagenet model."""
   hparams = imagetransformer_b12l_4h_b256_uncond_dr03_tpu()
   update_hparams_for_tpu(hparams)
   hparams.batch_size = 4
+  hparams.optimizer = "Adafactor"
+  hparams.learning_rate_schedule = "rsqrt_decay"
+  hparams.learning_rate_warmup_steps = 6000
   hparams.layer_prepostprocess_dropout = 0.1
   return hparams
 
@@ -1095,7 +1118,10 @@ def imagetransformer_b12l_4h_b128_uncond_dr03_tpu():
   hparams.filter_size = 2048
   hparams.layer_preprocess_sequence = "none"
   hparams.layer_postprocess_sequence = "dan"
-  hparams.layer_prepostprocess_dropout = 0.3
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.optimizer = "Adafactor"
+  hparams.learning_rate_schedule = "rsqrt_decay"
+  hparams.learning_rate_warmup_steps = 10000
   return hparams
 
 

From f37a12e78bb51bb25b7e4d72c7075a33522aeada Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 23 Aug 2018 19:48:19 -0700
Subject: [PATCH 0663/2720] Clean up shape info in
 {masked_,}local_attention_2d.

PiperOrigin-RevId: 210042530
---
 tensor2tensor/layers/common_attention.py      | 38 ++++-----
 tensor2tensor/layers/common_attention_test.py | 81 ++++++++++++-------
 2 files changed, 71 insertions(+), 48 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index ebf70535d..8a4a6c2f4 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2689,7 +2689,8 @@ def local_attention_2d(q,
   Args:
     q: a Tensor with shape [batch, heads, h, w, depth_k]
     k: a Tensor with shape [batch, heads, h, w, depth_k]
-    v: a Tensor with shape [batch, heads, h, w, depth_v]
+    v: a Tensor with shape [batch, heads, h, w, depth_v]. In the current
+      implementation, depth_v must be equal to depth_k.
     query_shape: an tuple indicating the height and width of each query block.
     memory_flange: an integer indicating how much to look in height and width
       from each query block.
@@ -2700,14 +2701,12 @@ def local_attention_2d(q,
   """
   with tf.variable_scope(
       name, default_name="local_self_attention_2d", values=[q, k, v]):
-    q_shape = q.get_shape().as_list()
     v_shape = common_layers.shape_list(v)
 
     # Pad query, key, value to ensure multiple of corresponding lengths.
     q = pad_to_multiple_2d(q, query_shape)
     k = pad_to_multiple_2d(k, query_shape)
     v = pad_to_multiple_2d(v, query_shape)
-    padded_q_shape = common_layers.shape_list(q)
     paddings = [[0, 0], [0, 0], [memory_flange[0], memory_flange[1]],
                 [memory_flange[0], memory_flange[1]], [0, 0]]
     k = tf.pad(k, paddings)
@@ -2726,7 +2725,6 @@ def local_attention_2d(q,
 
     attention_bias = tf.expand_dims(
         tf.to_float(embedding_to_padding(k_new)) * -1e9, axis=-2)
-
     output = dot_product_attention(
         q_new,
         k_new,
@@ -2736,12 +2734,12 @@ def local_attention_2d(q,
         name="local_2d",
         make_image_summary=False)
     # Put representations back into original shapes.
+    padded_q_shape = common_layers.shape_list(q)
     output = scatter_blocks_2d(output, q_indices, padded_q_shape)
 
     # Remove the padding if introduced.
     output = tf.slice(output, [0, 0, 0, 0, 0],
                       [-1, -1, v_shape[2], v_shape[3], -1])
-    output.set_shape(q_shape)
     return output
 
 
@@ -2949,19 +2947,19 @@ def right_shift_blockwise(x, query_shape, name=None):
       name, default_name="right_shift_blockwise", values=[x]):
     x_list_shape = x.get_shape().as_list()
     x_shape = common_layers.shape_list(x)
-    # Add a dummy dimension for heads
+    # Add a dummy dimension for heads.
     x = tf.expand_dims(x, axis=1)
     x = pad_to_multiple_2d(x, query_shape)
     padded_x_shape = common_layers.shape_list(x)
-    # Setting up q blocks
+    # Set up q blocks.
     x_indices = gather_indices_2d(x, query_shape, query_shape)
     x_new = get_shifted_center_blocks(x, x_indices)
 
-    # putting the representations back in the right place
+    # Put representations back into original shapes.
     output = scatter_blocks_2d(x_new, x_indices, padded_x_shape)
-    # Removing the dummy head dimension
+    # Remove the dummy head dimension.
     output = tf.squeeze(output, axis=1)
-    # Remove the padding if introduced
+    # Remove the padding if introduced.
     output = tf.slice(output, [0, 0, 0, 0], [-1, x_shape[1], x_shape[2], -1])
     output.set_shape(x_list_shape)
     return output
@@ -2975,17 +2973,18 @@ def masked_local_attention_2d(q,
                               name=None):
   """Strided block local self-attention.
 
-    Each position in a query block can attend to all the generated queries in
-    the query block, which are generated in raster scan, and positions that are
-    generated to the left and top. The shapes are specified by query shape and
-    memory flange. Note that if you're using this function, you do not need to
-    right shift. Right shifting happens inside this function separately for each
-    block.
+  Each position in a query block can attend to all the generated queries in
+  the query block, which are generated in raster scan, and positions that are
+  generated to the left and top. The shapes are specified by query shape and
+  memory flange. Note that if you're using this function, you do not need to
+  right shift. Right shifting happens inside this function separately for each
+  block.
 
   Args:
     q: a Tensor with shape [batch, heads, h, w, depth_k]
     k: a Tensor with shape [batch, heads, h, w, depth_k]
-    v: a Tensor with shape [batch, heads, h, w, depth_v]
+    v: a Tensor with shape [batch, heads, h, w, depth_v]. In the current
+      implementation, depth_v must be equal to depth_k.
     query_shape: an tuple indicating the height and width of each query block.
       query_shape = block_shape
     memory_flange: an integer indicating how much to look in height and width
@@ -2998,12 +2997,10 @@ def masked_local_attention_2d(q,
   """
   with tf.variable_scope(
       name, default_name="local_masked_self_attention_2d", values=[q, k, v]):
-    q_shape = q.get_shape().as_list()
     v_shape = common_layers.shape_list(v)
 
     # Pad query to ensure multiple of corresponding lengths.
     q = pad_to_multiple_2d(q, query_shape)
-    padded_q_shape = common_layers.shape_list(q)
 
     # Set up query blocks.
     q_indices = gather_indices_2d(q, query_shape, query_shape)
@@ -3051,14 +3048,13 @@ def masked_local_attention_2d(q,
         dropout_rate=0.,
         name="masked_local_2d",
         make_image_summary=False)
-
     # Put representations back into original shapes.
+    padded_q_shape = common_layers.shape_list(q)
     output = scatter_blocks_2d(output, q_indices, padded_q_shape)
 
     # Remove the padding if introduced.
     output = tf.slice(output, [0, 0, 0, 0, 0],
                       [-1, -1, v_shape[2], v_shape[3], -1])
-    output.set_shape(q_shape)
     return output
 
 
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index a14f43252..0585e7ef8 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -116,6 +116,34 @@ def testMaskedLocalAttention1D(self, batch, heads, length, depth_k, depth_v,
 
     self.assertEqual(res.shape, (batch, heads, length, depth_v))
 
+  @parameterized.named_parameters(
+      ("", 1, 1, 8, 4, 4, (2, 2)),
+      ("dynamic_batch", None, 1, 8, 4, 4, (2, 2)),
+      ("batches", 3, 2, 8, 4, 4, (2, 2)),
+      # TODO(trandustin): Extend function to enable depth_k != depth_v.
+      # ("depth_v", 1, 1, 8, 4, 1, (2, 2)),
+      ("query_shape", 1, 1, 8, 4, 4, (4, 4)),
+  )
+  def testMaskedLocalAttention2D(self, batch, heads, length, depth_k, depth_v,
+                                 query_shape):
+    if batch is None:
+      batch = tf.random_uniform([], minval=0, maxval=5, dtype=tf.int32)
+    q = tf.random_normal([batch, heads, length, length, depth_k])
+    k = tf.random_normal([batch, heads, length, length, depth_k])
+    v = tf.random_normal([batch, heads, length, length, depth_v])
+    output = common_attention.masked_local_attention_2d(
+        q,
+        k,
+        v,
+        query_shape=query_shape,
+        memory_flange=(2, 2))
+    if isinstance(batch, tf.Tensor):
+      batch, res = self.evaluate([batch, output])
+    else:
+      res = self.evaluate(output)
+
+    self.assertEqual(res.shape, (batch, heads, length, length, depth_v))
+
   @parameterized.named_parameters(
       ("matching_block_length", 3, 4, 25, 16, 16, 5),
       ("unmatching_block_length", 3, 4, 25, 16, 16, 4),
@@ -139,33 +167,32 @@ def testLocalUnmaskedAttention1D(self, batch, heads, length,
 
     self.assertEqual(res.shape, (batch, heads, length, depth_v))
 
-  def testLocalUnmaskedAttention2D(self):
-    x = np.random.rand(5, 4, 25, 25, 16)
-    y = np.random.rand(5, 4, 25, 25, 16)
-    with self.test_session() as session:
-      a = common_attention.local_attention_2d(
-          tf.constant(x, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          query_shape=(4, 4),
-          memory_flange=(3, 3))
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
-    self.assertEqual(res.shape, (5, 4, 25, 25, 16))
-
-  def testLocalUnmaskedAttention2DMatchingBlockLength(self):
-    x = np.random.rand(5, 4, 25, 25, 16)
-    y = np.random.rand(5, 4, 25, 25, 16)
-    with self.test_session() as session:
-      a = common_attention.local_attention_2d(
-          tf.constant(x, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          query_shape=(5, 5),
-          memory_flange=(3, 3))
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
-    self.assertEqual(res.shape, (5, 4, 25, 25, 16))
+  @parameterized.named_parameters(
+      ("matching_block_length", 3, 4, 25, 16, 16, (4, 4)),
+      ("unmatching_block_length", 3, 4, 25, 16, 16, (5, 5)),
+      ("dynamic_batch", None, 4, 25, 16, 16, (4, 4)),
+      # TODO(trandustin): Extend function to enable depth_k != depth_v.
+      # ("different_depth_v", 3, 4, 25, 16, 17, (4, 4)),
+  )
+  def testLocalUnmaskedAttention2D(self, batch, heads, length,
+                                   depth_k, depth_v, query_shape):
+    if batch is None:
+      batch = tf.random_uniform([], minval=0, maxval=5, dtype=tf.int32)
+    q = tf.random_normal([batch, heads, length, length, depth_k])
+    k = tf.random_normal([batch, heads, length, length, depth_k])
+    v = tf.random_normal([batch, heads, length, length, depth_v])
+    output = common_attention.local_attention_2d(
+        q,
+        k,
+        v,
+        query_shape=query_shape,
+        memory_flange=(3, 3))
+    if isinstance(batch, tf.Tensor):
+      batch, res = self.evaluate([batch, output])
+    else:
+      res = self.evaluate(output)
+
+    self.assertEqual(res.shape, (batch, heads, length, length, depth_v))
 
   def testMultiheadSelfAttentionMemoryEfficient(self):
     num_heads = 4

From 023de5122ab5eb1a11f43ddeb4e4927955f879e7 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 23 Aug 2018 22:51:55 -0700
Subject: [PATCH 0664/2720] Add a 2-layer discriminator, remove relu and add
 options to not do mean to the simple one, adjust hparams for text.

PiperOrigin-RevId: 210055227
---
 tensor2tensor/layers/common_layers.py         | 41 +++++++++++++++----
 tensor2tensor/models/research/autoencoders.py | 37 +++++++++++++----
 2 files changed, 63 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 01f477b98..c6c158eee 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3639,7 +3639,7 @@ def get_sorted_projections(x):
       batch_size = shape_list(x)[0]
       if do_random_vecs and do_tanh:
         n = tf.nn.l2_normalize(x, axis=1)
-        proj = tf.concat([tf.matmul(n, random_vecs), tf.tanh(x)], axis=1)
+        proj = tf.concat([tf.matmul(n, random_vecs), tf.tanh(n)], axis=1)
       elif do_random_vecs:
         n = tf.nn.l2_normalize(x, axis=1)
         proj = tf.matmul(n, random_vecs)
@@ -3685,7 +3685,7 @@ def deep_discriminator(x,
       "discriminator", initializer=tf.random_normal_initializer(stddev=0.02)):
     batch_size, height, width = shape_list(x)[:3]
     net = tf.layers.conv2d(
-        x, filters, filter_size, strides=stride, padding="SAME", name="d_conv1")
+        x, filters, filter_size, strides=stride, padding="SAME", name="conv1")
     net = lrelu(net)
     net = tf.layers.conv2d(
         net,
@@ -3693,7 +3693,7 @@ def deep_discriminator(x,
         filter_size,
         strides=stride,
         padding="SAME",
-        name="d_conv2")
+        name="conv2")
     # [bs, h/4, w/4, 128]
     if batch_norm:
       net = tf.layers.batch_normalization(
@@ -3786,16 +3786,43 @@ def patch_discriminator(x, filters=64, filter_size=5, n=4,
     return x
 
 
-def simple_discriminator(x, filters=128, filter_size=7, stride=4):
+def simple_discriminator(x, filters=128, kernel_size=7,
+                         strides=4, do_mean=True):
   """A very simple convolutional discriminator."""
   with tf.variable_scope("discriminator"):
     net = tf.layers.conv2d(
-        x, filters, filter_size, strides=stride, padding="SAME", name="d_conv1")
-    net = tf.nn.relu(net)
-    net = tf.reduce_mean(net, [1, 2])
+        x, filters, kernel_size, strides=strides, padding="SAME", name="conv1")
+    if do_mean:
+      net = tf.reduce_mean(net, [1, 2])
+    else:
+      batch_size = shape_list(x)[0]
+      net = tf.reshape(net, [batch_size, -1])
     return net
 
 
+def double_discriminator(x, filters1=128, filters2=None,
+                         kernel_size=7, strides=4, do_mean=True):
+  """A convolutional discriminator with 2 layers and concatenated output."""
+  if filters2 is None:
+    filters2 = 4 * filters1
+  with tf.variable_scope("discriminator"):
+    batch_size = shape_list(x)[0]
+    net = tf.layers.conv2d(
+        x, filters1, kernel_size, strides=strides, padding="SAME", name="conv1")
+    if do_mean:
+      net1 = tf.reduce_mean(net, [1, 2])
+    else:
+      net1 = tf.reshape(net, [batch_size, -1])
+    net = tf.nn.relu(net)
+    net = tf.layers.conv2d(
+        x, filters2, kernel_size, strides=strides, padding="SAME", name="conv2")
+    if do_mean:
+      net2 = tf.reduce_mean(net, [1, 2])
+    else:
+      net2 = tf.reshape(net, [batch_size, -1])
+    return tf.concat([net1, net2], axis=-1)
+
+
 def upscale(inputs, f, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR):
   """Upscaling the image by a factor of f."""
   height, width = shape_list(inputs)[1:3]
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 6e69ca72a..7cc089d07 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -283,11 +283,13 @@ def body(self, features):
               vocab_size,
               do_update=update_means,
               temperature=vq_temperature)
+          reconstr_gan_nonoise = reconstr_gan
           code_loss_gan *= hparams.code_loss_factor * update_means_factor
           losses["code_loss_gan"] = code_loss_gan
       else:
         reconstr_gan = tf.layers.dense(
             res_gan, vocab_size, name="autoencoder_final", reuse=True)
+        reconstr_gan_nonoise = reconstr_gan
         reconstr_gan = self.gumbel_sample(reconstr_gan)
         # Embed to codes.
         gan_codes = self.embed(reconstr_gan)
@@ -295,16 +297,29 @@ def body(self, features):
     # Add GAN loss if requested.
     gan_loss = 0.0
     if hparams.gan_loss_factor != 0.0:
-      self.image_summary("gan", reconstr_gan)
+      self.image_summary("gan", reconstr_gan_nonoise)
 
       def discriminate(x):
+        """Run a dioscriminator depending on the hparams."""
         if hparams.discriminator == "default":
           return common_layers.deep_discriminator(
               x, hparams.discriminator_batchnorm, is_training)
         elif hparams.discriminator == "patched":
           return common_layers.patch_discriminator(x)
         elif hparams.discriminator == "simple":
-          return common_layers.simple_discriminator(x)
+          return common_layers.simple_discriminator(
+              x,
+              hparams.discriminator_size,
+              hparams.discriminator_kernel_size,
+              hparams.discriminator_strides,
+              do_mean=hparams.discriminator_do_mean)
+        elif hparams.discriminator == "double":
+          return common_layers.double_discriminator(
+              x,
+              hparams.discriminator_size,
+              hparams.discriminator_kernel_size,
+              hparams.discriminator_strides,
+              do_mean=hparams.discriminator_do_mean)
         else:
           raise Exception("Unknown discriminator %s" % hparams.discriminator)
 
@@ -317,9 +332,9 @@ def discriminate(x):
       gan_lr = common_layers.inverse_exp_decay(
           hparams.gan_codes_warmup_steps * 1.5)
       rev_grad_gan_codes = reverse_gradient(gan_codes, lr=gan_lr)
-      gan_loss = common_layers.sliced_gan_loss(target_codes, rev_grad_gan_codes,
-                                               discriminate,
-                                               self.hparams.num_sliced_vecs)
+      gan_loss = common_layers.sliced_gan_loss(
+          target_codes, rev_grad_gan_codes, discriminate,
+          self.hparams.num_sliced_vecs, do_tanh=hparams.sliced_do_tanh)
       gan_loss *= hparams.gan_loss_factor * update_means_factor
       losses["gan_loss"] = -gan_loss
 
@@ -899,6 +914,11 @@ def autoencoder_basic():
   hparams.add_hparam("sample_width", 32)
   hparams.add_hparam("discriminator_batchnorm", True)
   hparams.add_hparam("num_sliced_vecs", 20000)
+  hparams.add_hparam("sliced_do_tanh", int(True))
+  hparams.add_hparam("discriminator_size", 256)
+  hparams.add_hparam("discriminator_kernel_size", 6)
+  hparams.add_hparam("discriminator_strides", 4)
+  hparams.add_hparam("discriminator_do_mean", int(True))
   hparams.add_hparam("code_loss_factor", 1.0)
   hparams.add_hparam("gan_codes_warmup_steps", 6000)
   hparams.add_hparam("gan_loss_factor", 0.0)
@@ -907,7 +927,7 @@ def autoencoder_basic():
   hparams.add_hparam("gumbel_noise_factor", 0.4)
   hparams.add_hparam("vq_temperature", 0.001)
   hparams.add_hparam("use_vq_loss", int(False))
-  hparams.add_hparam("discriminator", "default")
+  hparams.add_hparam("discriminator", "simple")
   return hparams
 
 
@@ -1053,10 +1073,11 @@ def autoencoder_ordered_text():
 def autoencoder_ordered_text_small():
   """Ordered discrete autoencoder model for text, small version."""
   hparams = autoencoder_ordered_text()
-  hparams.bottleneck_bits = 64
+  hparams.bottleneck_bits = 14
+  hparams.num_hidden_layers = 2
   hparams.hidden_size = 64
   hparams.max_hidden_size = 512
-  hparams.bottleneck_noise = 0.3
+  hparams.bottleneck_noise = 0.0
   hparams.autoregressive_mode = "conv5"
   return hparams
 

From b3ceeb827cbafe04fa2b3034a8ca07596a44a0ec Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Fri, 24 Aug 2018 00:11:16 -0700
Subject: [PATCH 0665/2720] More correct logic for selecting group sizes for
 hierarchical mixture of experts - We need the number of groups to be a
 multiple of the mesh dimension over which the groups are split.

PiperOrigin-RevId: 210060581
---
 .../research/experiments_moe.py               |  4 +
 tensor2tensor/mesh_tensorflow/research/moe.py | 75 +++++++++++++++----
 2 files changed, 66 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
index 5cee29a7b..43d2b19db 100644
--- a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
@@ -40,6 +40,10 @@ def xmoe_dense_4k():
   xmoe_dense_64k    282        2.8e13  0                    3.06
   xmoe_top_2        282        4.0e12  3.4e8     36%        3.07
   xmoe_top_2_c15    282        4.5e12  4.0e8     38%        3.07
+  xmoe_2d           282        5.3e12  7.6e8     34%        3.06
+
+  Trained at 4x the batch size:
+  xmoe_2d_88        1090       2.1e13  3.0e9     24%
 
   Note: configurations and code are likely to change without notice.
 
diff --git a/tensor2tensor/mesh_tensorflow/research/moe.py b/tensor2tensor/mesh_tensorflow/research/moe.py
index e09f262ac..88e9c27e1 100644
--- a/tensor2tensor/mesh_tensorflow/research/moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/moe.py
@@ -268,11 +268,15 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
   y = mtf.Dimension("expert_y_unsplit", hparams.moe_num_experts[1])
   n = output_dim
 
-  numerator = b1.size * l.size
-  s = mtf.Dimension("group_size_x", _largest_divisor_leq(
-      numerator, hparams.moe_group_size))
-  g1 = mtf.Dimension(b1.name, numerator // s.size)
+  # We "cheat" here and look at the mesh shape and layout. This is to ensure
+  # that the number of groups (g.size) is a multiple of the mesh dimension
+  # over which those groups are split.
+  num_groups, group_size = _split_into_groups(
+      b1.size * l.size, hparams.moe_group_size,
+      _tensor_dim_to_mesh_dim_size(hparams, b1))
+  g1 = mtf.Dimension(b1.name, num_groups)
   g = mtf.Dimension(b1.name + "_unsplit", g1.size)
+  s = mtf.Dimension("group_size_x", group_size)
 
   # Each sequence sends (at most?) expert_capacity positions to each expert.
   # Static expert_capacity dimension is needed for expert batch sizes
@@ -284,10 +288,15 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
       int((s.size * capacity_factor) / x.size))
   c = mtf.Dimension("expert_capacity_x", expert_capacity)
 
-  numerator = a0.size * g.size * c.size
-  t = mtf.Dimension("group_size_y", _largest_divisor_leq(
-      numerator, hparams.moe_group_size))
-  h0 = mtf.Dimension(a0.name, numerator // t.size)
+  # We "cheat" here and look at the mesh shape and layout. This is to ensure
+  # that the number of groups (h.size) is a multiple of the mesh dimension
+  # over which those groups are split.
+  num_groups, group_size = _split_into_groups(
+      a0.size * g.size * c.size,
+      hparams.moe_group_size,
+      _tensor_dim_to_mesh_dim_size(hparams, a0))
+  t = mtf.Dimension("group_size_y", group_size)
+  h0 = mtf.Dimension(a0.name, num_groups)
   h = mtf.Dimension(a0.name + "_unsplit", h0.size)
 
   expert_capacity = min(
@@ -609,8 +618,48 @@ def set_default_moe_hparams(hparams):
   hparams.add_hparam("moe_second_threshold_eval", 0.2)
 
 
-def _largest_divisor_leq(numerator, maximum):
-  x = maximum
-  while numerator % x != 0:
-    x -= 1
-  return x
+def _split_into_groups(n, max_group_size, mesh_dim_size):
+  """Helper function for figuring out how to split a dimensino into groups.
+
+  We have a dimension with size n and we want to split it into
+  two dimensions: n = num_groups * group_size
+
+  group_size should be the largest possible value meeting the constraints:
+    group_size <= max_group_size
+    (num_groups = n/group_size) is a multiple of mesh_dim_size
+
+  Args:
+    n: an integer
+    max_group_size: an integer
+    mesh_dim_size: an integer
+
+  Returns:
+    num_groups: an integer
+    group_size: an integer
+
+  Raises:
+    ValueError: if n is not a multiple of mesh_dim_size
+  """
+  if n % mesh_dim_size != 0:
+    raise ValueError(
+        "n=%d is not a multiple of mesh_dim_size=%d" % (n, mesh_dim_size))
+  num_groups = max(1, n // max_group_size)
+  while (num_groups % mesh_dim_size != 0 or n % num_groups != 0):
+    num_groups += 1
+  group_size = n // num_groups
+  tf.logging.info(
+      "_split_into_groups(n=%d, max_group_size=%d, mesh_dim_size=%d)"
+      " = (num_groups=%d group_size=%d)" %
+      (n, max_group_size, mesh_dim_size, num_groups, group_size))
+  return num_groups, group_size
+
+
+def _tensor_dim_to_mesh_dim_size(hparams, tensor_dim):
+  """Inspect hparams to figure out how many ways tensor_dim gets split."""
+  layout_rules = mtf.convert_to_layout_rules(hparams.layout)
+  mesh_shape = mtf.convert_to_shape(hparams.mesh_shape)
+  mesh_axis = layout_rules.tensor_dimension_to_mesh_axis(tensor_dim, mesh_shape)
+  if mesh_axis is None:
+    return 1
+  else:
+    return mesh_shape.dims[mesh_axis].size

From 9966af9d46a762d0d96f04341b58eaba1880c43e Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 24 Aug 2018 11:40:15 -0700
Subject: [PATCH 0666/2720] Internal change

PiperOrigin-RevId: 210131357
---
 tensor2tensor/data_generators/cifar.py        |  16 +++
 tensor2tensor/data_generators/image_utils.py  |  20 ++++
 .../data_generators/image_utils_test.py       |   4 +
 tensor2tensor/models/__init__.py              |   1 +
 tensor2tensor/models/research/glow.py         | 109 ++++++++++++++++++
 tensor2tensor/models/research/glow_ops.py     |  76 +++++++-----
 .../models/research/glow_ops_test.py          |  28 +++--
 tensor2tensor/models/research/glow_test.py    |  64 ++++++++++
 8 files changed, 276 insertions(+), 42 deletions(-)
 create mode 100644 tensor2tensor/models/research/glow.py
 create mode 100644 tensor2tensor/models/research/glow_test.py

diff --git a/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index f19deb84c..c070d237d 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -177,6 +177,22 @@ def preprocess_example(self, example, mode, unused_hparams):
     return example
 
 
+@registry.register_problem
+class ImageCifar10PlainRandomShift(ImageCifar10Plain):
+  """CIFAR-10 32x32 for image generation with random shift data-augmentation."""
+
+  def dataset_filename(self):
+    return "image_cifar10_plain"  # Reuse CIFAR-10 plain data.
+
+  def preprocess_example(self, example, mode, unused_hparams):
+    example["inputs"].set_shape([_CIFAR10_IMAGE_SIZE, _CIFAR10_IMAGE_SIZE, 3])
+    example["inputs"] = tf.to_int64(example["inputs"])
+    if mode == tf.estimator.ModeKeys.TRAIN:
+      example["inputs"] = image_utils.random_shift(
+          example["inputs"], wsr=0.1, hsr=0.1)
+    return example
+
+
 @registry.register_problem
 class ImageCifar10PlainGenDmol(ImageCifar10PlainGen):
   """Discretized mixture of logistics problem."""
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 98ec9da81..c7d1f5f7f 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -337,3 +337,23 @@ def cifar_image_augmentation(images):
   images = tf.random_crop(images, [32, 32, 3])
   images = tf.image.random_flip_left_right(images)
   return images
+
+
+def random_shift(image, wsr=0.1, hsr=0.1):
+  """Apply random horizontal and vertical shift to images.
+
+  This is the default data-augmentation strategy used on CIFAR in Glow.
+
+  Args:
+    image: a 3-D Tensor
+    wsr: Width shift range, as a float fraction of the width.
+    hsr: Height shift range, as a float fraction of the width.
+  Returns:
+    images: images translated by the provided wsr and hsr.
+  """
+  height, width, _ = common_layers.shape_list(image)
+  width_range, height_range = wsr*width, hsr*height
+  height_translations = tf.random_uniform((1,), -height_range, height_range)
+  width_translations = tf.random_uniform((1,), -width_range, width_range)
+  translations = tf.concat((height_translations, width_translations), axis=0)
+  return tf.contrib.image.translate(image, translations=translations)
diff --git a/tensor2tensor/data_generators/image_utils_test.py b/tensor2tensor/data_generators/image_utils_test.py
index c8f8ca47e..867c5638f 100644
--- a/tensor2tensor/data_generators/image_utils_test.py
+++ b/tensor2tensor/data_generators/image_utils_test.py
@@ -112,6 +112,10 @@ def testMakeMultiscaleDilatedLarger(self):
     with self.assertRaisesRegexp(ValueError, "strides.* must be non-zero"):
       _ = image_utils.make_multiscale_dilated(image, resolutions)
 
+  def testRandomShift(self):
+    image = tf.random_normal([256, 256, 3])
+    image_shift = image_utils.random_shift(image, wsr=0.1, hsr=0.1)
+    self.assertEqual(image_shift.shape, [256, 256, 3])
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index b41892bd2..46fda37e2 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -44,6 +44,7 @@
 from tensor2tensor.models.research import autoencoders
 from tensor2tensor.models.research import cycle_gan
 from tensor2tensor.models.research import gene_expression
+from tensor2tensor.models.research import glow
 from tensor2tensor.models.research import lm_experiments
 from tensor2tensor.models.research import multimodel
 from tensor2tensor.models.research import next_frame_basic_deterministic
diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
new file mode 100644
index 000000000..f4cc8aa27
--- /dev/null
+++ b/tensor2tensor/models/research/glow.py
@@ -0,0 +1,109 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Glow generative model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.layers import common_hparams
+from tensor2tensor.layers import common_layers
+from tensor2tensor.models.research import glow_ops
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
+import tensorflow as tf
+
+arg_scope = tf.contrib.framework.arg_scope
+add_arg_scope = tf.contrib.framework.add_arg_scope
+
+
+@registry.register_hparams
+def glow_hparams():
+  """Glow Hparams."""
+  hparams = common_hparams.basic_params1()
+  hparams.clip_grad_norm = None
+  hparams.weight_decay = 0.0
+  hparams.learning_rate_constant = 3e-4
+  hparams.batch_size = 32
+  hparams.add_hparam("n_levels", 3)
+  hparams.add_hparam("n_bits_x", 8)
+  hparams.add_hparam("depth", 32)
+  hparams.add_hparam("affine_coupling_width", 512)
+  hparams.add_hparam("learn_prior", True)
+  return hparams
+
+
+@registry.register_model
+class Glow(t2t_model.T2TModel):
+  """Glow generative model.
+
+  Reference: https://arxiv.org/abs/1807.03039"""
+
+  def preprocess(self, x):
+    """Normalize x.
+
+    Args:
+      x: 4-D Tensor.
+
+    Returns:
+      x: Scaled such that x lies in-between -0.5 and 0.5
+    """
+    n_bits_x = self.hparams.n_bits_x
+    n_bins = 2**n_bits_x
+    x = tf.cast(x, dtype=tf.float32)
+    if n_bits_x < 8:
+      x = tf.floor(x / 2 ** (8 - n_bits_x))
+    x = x / n_bins - 0.5
+    return x
+
+  @property
+  def is_training(self):
+    return self.hparams.mode == tf.estimator.ModeKeys.TRAIN
+
+  def body(self, features):
+    x = features["inputs"]
+
+    # Scale x such that the pixels lie in-between -0.5 and.0.5
+    x = self.preprocess(x)
+
+    n_bins = 2**self.hparams.n_bits_x
+    batch_size, height, width, n_channels = common_layers.shape_list(x)
+    hwc = float(height * width * n_channels)
+
+    x = x + tf.random_uniform(
+        shape=(batch_size, height, width, n_channels),
+        minval=0.0, maxval=1.0/n_bins)
+    objective = -np.log(n_bins) * hwc * tf.ones(batch_size)
+
+    # The arg_scope call ensures that the actnorm parameters are set such that
+    # the per-channel output activations have zero mean and unit variance
+    # ONLY during the first step. After that the parameters are learned
+    # through optimisation.
+    global_step = tf.train.get_or_create_global_step()
+    init_op = tf.logical_and(tf.equal(global_step, 0), self.is_training)
+    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
+    with arg_scope(ops, init=init_op):
+      z, encoder_objective, _ = glow_ops.encoder_decoder(
+          "codec", x, self.hparams, eps=None, reverse=False)
+      objective += encoder_objective
+
+      prior_objective = glow_ops.top_prior(
+          "top_prior", z, learn_prior=self.hparams.learn_prior)
+      objective += prior_objective
+
+    # bits per pixel
+    objective = -objective / (np.log(2) * hwc)
+    return tf.zeros_like(features["targets"]), {"training": objective}
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 619f7d080..2d86e2921 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -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.
-"""Glow generative model."""
+"""Various reversible ops for the glow generative model."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -21,27 +21,13 @@
 from functools import partial
 import numpy as np
 import scipy
-from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
-from tensor2tensor.utils import registry
 import tensorflow as tf
 
 arg_scope = tf.contrib.framework.arg_scope
 add_arg_scope = tf.contrib.framework.add_arg_scope
 
 
-@registry.register_hparams
-def glow_hparams():
-  """Glow Hparams."""
-  hparams = common_hparams.basic_params1()
-  hparams.add_hparam("n_levels", 3)
-  hparams.add_hparam("n_bits_x", 8)
-  hparams.add_hparam("depth", 32)
-  hparams.add_hparam("affine_coupling_width", 512)
-  hparams.add_hparam("learn_prior", True)
-  return hparams
-
-
 def default_initializer(std=0.05):
   return tf.random_normal_initializer(0., std)
 
@@ -56,16 +42,22 @@ def set_eps(dist, eps):
   return eps * dist.scale + dist.loc
 
 
+@add_arg_scope
+def assign(w, initial_value):
+  w = w.assign(initial_value)
+  with tf.control_dependencies([w]):
+    return w
+
+
 @add_arg_scope
 def get_variable_ddi(name, shape, initial_value, dtype=tf.float32, init=False,
                      trainable=True):
   """Wrapper for data-dependent initialization."""
+  # Cast from python bool to TF bool for usage in tf.cond
+  if isinstance(init, bool):
+    init = tf.constant(init, dtype=tf.bool)
   w = tf.get_variable(name, shape, dtype, None, trainable=trainable)
-  if init:
-    w = w.assign(initial_value)
-    with tf.control_dependencies([w]):
-      return w
-  return w
+  return tf.cond(init, lambda: assign(w, initial_value), lambda: w)
 
 
 @add_arg_scope
@@ -92,6 +84,7 @@ def actnorm(name, x, logscale_factor=3., reverse=False, init=False,
   """
   var_arg_scope = arg_scope([get_variable_ddi], trainable=trainable)
   var_scope = tf.variable_scope(name, reuse=tf.AUTO_REUSE)
+
   with var_scope, var_arg_scope:
     if not reverse:
       x = actnorm_center(name + "_center", x, reverse, init=init)
@@ -127,8 +120,8 @@ def actnorm_center(name, x, reverse=False, init=False):
     assert len(shape) == 2 or len(shape) == 4
     if len(shape) == 2:
       x_mean = tf.reduce_mean(x, [0], keepdims=True)
-      b = get_variable_ddi(
-          "b", (1, shape[1]), initial_value=-x_mean, init=init)
+      b = get_variable_ddi("b", (1, shape[1]), initial_value=-x_mean,
+                           init=init)
     elif len(shape) == 4:
       x_mean = tf.reduce_mean(x, [0, 1, 2], keepdims=True)
       b = get_variable_ddi(
@@ -159,8 +152,8 @@ def actnorm_scale(name, x, logscale_factor=3., reverse=False, init=False):
       var_shape = (1, 1, 1, x_shape[3])
 
     init_value = tf.log(1.0 / (tf.sqrt(x_var) + 1e-6)) / logscale_factor
-    logs = get_variable_ddi(
-        "logs", var_shape, initial_value=init_value, init=init)
+    logs = get_variable_ddi("logs", var_shape, initial_value=init_value,
+                            init=init)
     logs = logs * logscale_factor
 
     # Function and reverse function.
@@ -269,8 +262,7 @@ def add_edge_bias(x, filter_size):
 
 @add_arg_scope
 def conv2d(name, x, output_channels, filter_size=None, stride=None,
-           logscale_factor=3.0, init=True, apply_actnorm=True,
-           conv_init="default"):
+           logscale_factor=3.0, apply_actnorm=True, conv_init="default"):
   """conv2d layer with edge bias padding and optional actnorm.
 
   Args:
@@ -280,8 +272,6 @@ def conv2d(name, x, output_channels, filter_size=None, stride=None,
     filter_size:
     stride:
     logscale_factor: see actnorm for parameter meaning.
-    init: Whether to apply data-dependent initialization Valid only if
-          apply_actnorm is set to True.
     apply_actnorm: if apply_actnorm the activations of the first minibatch
                    have zero mean and unit variance. Else, there is no scaling
                    applied.
@@ -291,7 +281,7 @@ def conv2d(name, x, output_channels, filter_size=None, stride=None,
   Raises:
     ValueError: if init is set to "zeros" and apply_actnorm is set to True.
   """
-  if init == "zeros" and apply_actnorm:
+  if conv_init == "zeros" and apply_actnorm:
     raise ValueError("apply_actnorm is unstable when init is set to zeros.")
 
   if filter_size is None:
@@ -317,7 +307,7 @@ def conv2d(name, x, output_channels, filter_size=None, stride=None,
     x = tf.nn.conv2d(x, w, stride_shape, padding="VALID", data_format="NHWC")
 
     if apply_actnorm:
-      x, _ = actnorm("actnorm", x, logscale_factor=logscale_factor, init=init,
+      x, _ = actnorm("actnorm", x, logscale_factor=logscale_factor,
                      trainable=True)
     else:
       x += tf.get_variable("b", [1, 1, 1, output_channels],
@@ -529,11 +519,35 @@ def revnet(name, x, hparams, reverse=True):
     objective = 0.0
     for step in steps:
       x, curr_obj = revnet_step(
-          "revnet_%d" % step, x, hparams, reverse=reverse)
+          "revnet_step_%d" % step, x, hparams, reverse=reverse)
       objective += curr_obj
     return x, objective
 
 
+@add_arg_scope
+def top_prior(name, x, learn_prior=False):
+  """Log probability of x being gaussian.
+
+  Args:
+    name: variable scope
+    x: input, 4-D Tensor shape=(batch_size, width, height, channels)
+    learn_prior: If set to true, then the mean and the standard deviation
+                 are the output of a single conv layer initialized with
+                 zeros. Otherwise the mean and logstd are zeros and ones
+                 respectively.
+  Returns:
+    objective: 1-D Tensor shape=(batch_size,) summed across spatial components.
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    h = tf.zeros_like(x)
+    if not learn_prior:
+      prior_dist = tf.distributions.Normal(h, tf.exp(h))
+    else:
+      prior_dist = split_prior("top_learn_prior", h)
+    return tf.reduce_sum(prior_dist.log_prob(x), axis=[1, 2, 3])
+
+
+@add_arg_scope
 def encoder_decoder(name, x, hparams, eps=None, reverse=False):
   """Glow encoder-decoder. n_levels of (Squeeze + Flow + Split.) operations."""
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 271a58657..f37d789dc 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -19,17 +19,22 @@
 from __future__ import print_function
 
 import numpy as np
+from tensor2tensor.models.research import glow
 from tensor2tensor.models.research import glow_ops
 import tensorflow as tf
 
 
+arg_scope = tf.contrib.framework.arg_scope
+add_arg_scope = tf.contrib.framework.add_arg_scope
+
+
 class GlowOpsTest(tf.test.TestCase):
 
   def test_get_variable_ddi(self):
     with tf.Graph().as_default():
       x_t = tf.random_normal((5, 5))
       ddi = glow_ops.get_variable_ddi(
-          "x", (5, 5), x_t, init=True)
+          "x", (5, 5), initial_value=x_t, init=True)
       with tf.Session() as session:
         diff = ddi - x_t
         self.assertTrue(np.allclose(session.run(diff), 0.0))
@@ -77,12 +82,13 @@ def test_add_edge_bias(self):
   def test_conv2d(self):
     with tf.Graph().as_default():
       x = 10.0 * tf.random_uniform(shape=(16, 5, 5, 32))
-      actnorm_conv2d = glow_ops.conv2d(
-          "actnorm_conv2d", x, output_channels=64, init=True,
-          apply_actnorm=True)
-      actnorm_zeros2d = glow_ops.conv2d(
-          "actnorm_zeros2d", x, output_channels=64, init=True,
-          apply_actnorm=False)
+
+      with arg_scope([glow_ops.actnorm], init=True):
+        actnorm_conv2d = glow_ops.conv2d(
+            "actnorm_conv2d", x, output_channels=64, apply_actnorm=True)
+        actnorm_zeros2d = glow_ops.conv2d(
+            "actnorm_zeros2d", x, output_channels=64, apply_actnorm=False)
+
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())
 
@@ -136,7 +142,7 @@ def test_split(self):
 
   def check_revnet_reversibility(self, op, name):
     with tf.Graph().as_default():
-      hparams = glow_ops.glow_hparams()
+      hparams = glow.glow_hparams()
       hparams.depth = 2
       x = tf.random_uniform(shape=(16, 32, 32, 4), seed=0)
       x_inv, _ = op(name, x, hparams, reverse=False)
@@ -144,7 +150,7 @@ def check_revnet_reversibility(self, op, name):
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())
         diff = session.run(x - x_inv_inv)
-        self.assertTrue(np.allclose(diff, 0.0, atol=1e-3))
+        self.assertTrue(np.allclose(diff, 0.0, atol=1e-2))
 
   def test_revnet_reversibility(self):
     ops = [glow_ops.revnet_step, glow_ops.revnet]
@@ -154,7 +160,7 @@ def test_revnet_reversibility(self):
 
   def test_encoder_decoder(self):
     with tf.Graph().as_default():
-      hparams = glow_ops.glow_hparams()
+      hparams = glow.glow_hparams()
       hparams.n_levels = 2
       hparams.depth = 2
 
@@ -168,7 +174,7 @@ def test_encoder_decoder(self):
         session.run(tf.global_variables_initializer())
         diff, x_inv_np = session.run([x - x_inv_inv, x_inv])
         self.assertTrue(x_inv_np.shape, (16, 8, 8, 64))
-        self.assertTrue(np.allclose(diff, 0.0, atol=1e-2))
+        self.assertTrue(np.allclose(diff, 0.0, atol=1e-3))
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
new file mode 100644
index 000000000..19c0e3bea
--- /dev/null
+++ b/tensor2tensor/models/research/glow_test.py
@@ -0,0 +1,64 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.models.research.glow_model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor import problems
+from tensor2tensor.data_generators import cifar  # pylint: disable=unused-import
+from tensor2tensor.models.research import glow
+from tensor2tensor.utils import registry  # pylint: disable=unused-import
+import tensorflow as tf
+
+MODES = tf.estimator.ModeKeys
+
+
+class GlowModelTest(tf.test.TestCase):
+
+  def batch(self, one_shot_iterator, batch_size=16):
+    x_batch, y_batch = [], []
+    for _ in range(batch_size):
+      curr = one_shot_iterator.get_next()
+      x_batch.append(curr['inputs'])
+      y_batch.append(curr['targets'])
+    return tf.stack(x_batch), tf.stack(y_batch)
+
+  def test_glow(self):
+    with tf.Graph().as_default():
+      hparams = glow.glow_hparams()
+      model = glow.Glow(hparams, tf.estimator.ModeKeys.TRAIN)
+      cifar_problem = problems.problem('image_cifar10_plain_random_shift')
+      train_dataset = cifar_problem.dataset(MODES.TRAIN)
+      one_shot = train_dataset.make_one_shot_iterator()
+      x_batch, y_batch = self.batch(one_shot)
+      features = {'inputs': x_batch, 'targets': y_batch}
+      _, obj_dict = model.body(features)
+      objective = obj_dict['training']
+      with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+        obj_np = sess.run(objective)
+        mean_obj = np.mean(obj_np)
+
+        # Check that one forward-propagation does not NaN, i.e
+        # initialization etc works as expected.
+        is_undefined = np.isnan(mean_obj) or np.isinf(mean_obj)
+        self.assertTrue(not is_undefined)
+
+
+if __name__ == '__main__':
+  tf.test.main()

From a40b9ea43eed58d77b0a4b087b9fbabbc5905507 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 24 Aug 2018 11:56:20 -0700
Subject: [PATCH 0667/2720] Test encoder_decoder in a practical setting. More
 specifically, test the following sequence of operations.

PiperOrigin-RevId: 210133924
---
 tensor2tensor/models/research/glow_ops.py     |  7 +--
 .../models/research/glow_ops_test.py          | 53 +++++++++++++++++++
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 2d86e2921..025b18705 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -95,7 +95,7 @@ def actnorm(name, x, logscale_factor=3., reverse=False, init=False,
       x, objective = actnorm_scale(
           name + "_scale", x, logscale_factor=logscale_factor,
           reverse=reverse, init=init)
-      x = actnorm_center(name + "_center", x, reverse)
+      x = actnorm_center(name + "_center", x, reverse, init=init)
     return x, objective
 
 
@@ -138,7 +138,7 @@ def actnorm_center(name, x, reverse=False, init=False):
 def actnorm_scale(name, x, logscale_factor=3., reverse=False, init=False):
   """Per-channel scaling of x."""
   x_shape = common_layers.shape_list(x)
-  with tf.variable_scope(name):
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
 
     # Variance initialization logic.
     assert len(x_shape) == 2 or len(x_shape) == 4
@@ -410,8 +410,9 @@ def squeeze(name, x, factor=2, reverse=True):
     height = int(shape[1])
     width = int(shape[2])
     n_channels = int(shape[3])
-    assert height % factor == 0 and width % factor == 0
+
     if not reverse:
+      assert height % factor == 0 and width % factor == 0
       x = tf.reshape(x, [-1, height//factor, factor,
                          width//factor, factor, n_channels])
       x = tf.transpose(x, [0, 1, 3, 5, 2, 4])
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index f37d789dc..8e90117f4 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -18,11 +18,16 @@
 from __future__ import division
 from __future__ import print_function
 
+import os
+import tempfile
 import numpy as np
 from tensor2tensor.models.research import glow
 from tensor2tensor.models.research import glow_ops
 import tensorflow as tf
 
+arg_scope = tf.contrib.framework.arg_scope
+add_arg_scope = tf.contrib.framework.add_arg_scope
+
 
 arg_scope = tf.contrib.framework.arg_scope
 add_arg_scope = tf.contrib.framework.add_arg_scope
@@ -176,6 +181,54 @@ def test_encoder_decoder(self):
         self.assertTrue(x_inv_np.shape, (16, 8, 8, 64))
         self.assertTrue(np.allclose(diff, 0.0, atol=1e-3))
 
+  def test_encoder_decoder_practical_usage(self):
+    """Tests the following sequence of operations.
+
+    1. Define forward network with arg_scope(init=True).
+    2. Run one-forward pass to do data-dependent initialization and save.
+    3. Define forward and reverse network with arg_scope(init=False)
+    4. Check that reverse(forward(x)) == x
+    """
+    hparams = glow.glow_hparams()
+    hparams.n_levels = 2
+    hparams.depth = 12
+
+    with tf.Graph().as_default():
+      rng = np.random.RandomState(0)
+      x_rand = np.asarray(rng.rand(1, 4, 4, 4), dtype=np.float32)
+      x_t = tf.convert_to_tensor(x_rand)
+
+      ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
+      with arg_scope(ops, init=True):
+        x_inv, _, _ = glow_ops.encoder_decoder(
+            "revnet", x_t, hparams, reverse=False)
+      curr_dir = tempfile.mkdtemp()
+      model_path = os.path.join(curr_dir, "model")
+
+      with tf.Session() as session:
+        saver = tf.train.Saver()
+        session.run(tf.global_variables_initializer())
+        session.run(x_inv)
+        saver.save(session, model_path)
+
+    with tf.Graph().as_default():
+      rng = np.random.RandomState(0)
+      x_rand = np.asarray(rng.rand(1, 4, 4, 4), dtype=np.float32)
+      x_t = tf.convert_to_tensor(x_rand)
+      ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
+      with arg_scope(ops, init=False):
+        x_inv2, _, all_eps = glow_ops.encoder_decoder(
+            "revnet", x_t, hparams, reverse=False)
+        x_inv_inv_, _ = glow_ops.encoder_decoder(
+            "revnet", x_inv2, hparams, eps=all_eps, reverse=True)
+
+      with tf.Session() as session:
+        saver = tf.train.Saver()
+        saver.restore(session, model_path)
+        x_inv_inv_np = session.run(x_inv_inv_)
+        diff = np.abs(x_inv_inv_np - x_rand)
+        self.assertTrue(np.allclose(diff, 0.0, atol=1e-3))
+
 
 if __name__ == "__main__":
   tf.test.main()

From b8b5f52a0b701f3fe8537dceb02bc11e69b9c8e6 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 24 Aug 2018 17:10:14 -0700
Subject: [PATCH 0668/2720] Internal change

PiperOrigin-RevId: 210182059
---
 tensor2tensor/data_generators/image_utils.py  | 53 +++++++++++++++++++
 .../data_generators/image_utils_test.py       | 29 ++++++++++
 tensor2tensor/models/research/glow.py         | 31 +++++++++--
 tensor2tensor/models/research/glow_ops.py     |  3 +-
 tensor2tensor/models/research/glow_test.py    | 45 ++++++++++++++++
 tensor2tensor/utils/decoding.py               |  3 +-
 6 files changed, 159 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index c7d1f5f7f..f1e898f7c 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -18,7 +18,11 @@
 from __future__ import division
 from __future__ import print_function
 
+
+import io
 import os
+import matplotlib.pyplot as plt
+import numpy as np
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -29,6 +33,51 @@
 import tensorflow as tf
 
 
+def image_to_tf_summary_value(image, tag):
+  """Converts a NumPy image to a tf.Summary.Value object.
+
+  Args:
+    image: 3-D NumPy array.
+    tag: name for tf.Summary.Value for display in tensorboard.
+  Returns:
+    image_summary: A tf.Summary.Value object.
+  """
+  curr_image = np.asarray(image, dtype=np.uint8)
+  height, width, n_channels = curr_image.shape
+  s = io.BytesIO()
+  plt.imsave(s, curr_image, format="png")
+  img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(),
+                             height=height, width=width,
+                             colorspace=n_channels)
+  return tf.Summary.Value(tag=tag, image=img_sum)
+
+
+def convert_predictions_to_image_summaries(hook_args):
+  """Optionally converts images from hooks_args to image summaries.
+
+  Args:
+    hook_args: DecodeHookArgs namedtuple
+  Returns:
+    summaries: list of tf.Summary values if hook_args.decode_hpara
+  """
+  decode_hparams = hook_args.decode_hparams
+  if not decode_hparams.display_decoded_images:
+    return []
+  predictions = hook_args.predictions[0]
+
+  # Display ten random inputs and outputs so that tensorboard does not hang.
+  all_summaries = []
+  rand_predictions = np.random.choice(predictions, size=10)
+  for ind, prediction in enumerate(rand_predictions):
+    output_summary = image_to_tf_summary_value(
+        prediction["outputs"], tag="%d_output" % ind)
+    input_summary = image_to_tf_summary_value(
+        prediction["inputs"], tag="%d_input" % ind)
+    all_summaries.append(input_summary)
+    all_summaries.append(output_summary)
+  return all_summaries
+
+
 def resize_by_area(img, size):
   """image resize function used by quite a few image problems."""
   return tf.to_int64(
@@ -135,6 +184,10 @@ def eval_metrics(self):
       eval_metrics += [metrics.Metrics.IMAGE_SUMMARY]
     return eval_metrics
 
+  @property
+  def decode_hooks(self):
+    return [convert_predictions_to_image_summaries]
+
 
 class Image2ClassProblem(ImageProblem):
   """Base class for image classification problems."""
diff --git a/tensor2tensor/data_generators/image_utils_test.py b/tensor2tensor/data_generators/image_utils_test.py
index 867c5638f..1dca736e5 100644
--- a/tensor2tensor/data_generators/image_utils_test.py
+++ b/tensor2tensor/data_generators/image_utils_test.py
@@ -19,6 +19,7 @@
 from __future__ import print_function
 import numpy as np
 from tensor2tensor.data_generators import image_utils
+from tensor2tensor.utils import decoding
 
 import tensorflow as tf
 
@@ -117,5 +118,33 @@ def testRandomShift(self):
     image_shift = image_utils.random_shift(image, wsr=0.1, hsr=0.1)
     self.assertEqual(image_shift.shape, [256, 256, 3])
 
+  def testImageToSummaryValue(self):
+    rng = np.random.RandomState(0)
+    x = rng.randint(0, 255, (32, 32, 3))
+    x_summary = image_utils.image_to_tf_summary_value(x, "X_image")
+    self.assertEqual(x_summary.tag, "X_image")
+
+  def testConvertPredictionsToImageSummaries(self):
+    # Initialize predictions.
+    rng = np.random.RandomState(0)
+    x = rng.randint(0, 255, (32, 32, 3))
+    predictions = [[{"outputs": x, "inputs": x}] * 50]
+
+    decode_hparams = decoding.decode_hparams()
+    # should return 20 summaries of images, 10 outputs and 10 inputs if
+    # display_decoded_images is set to True.
+    for display, summaries_length in zip([True, False], [20, 0]):
+      decode_hparams.display_decoded_images = display
+      decode_hooks = decoding.DecodeHookArgs(
+          estimator=None, problem=None, output_dirs=None,
+          hparams=decode_hparams, decode_hparams=decode_hparams,
+          predictions=predictions)
+      summaries = image_utils.convert_predictions_to_image_summaries(
+          decode_hooks)
+      self.assertEqual(len(summaries), summaries_length)
+      if summaries:
+        self.assertTrue(isinstance(summaries[0], tf.Summary.Value))
+
+
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index f4cc8aa27..629a882a8 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -29,6 +29,9 @@
 arg_scope = tf.contrib.framework.arg_scope
 add_arg_scope = tf.contrib.framework.add_arg_scope
 
+GLOW_DECODE_HPARAMS = ("identity_output=True,log_results=False,"
+                       "decode_in_memory=True,display_decoded_images=True")
+
 
 @registry.register_hparams
 def glow_hparams():
@@ -69,10 +72,31 @@ def preprocess(self, x):
     x = x / n_bins - 0.5
     return x
 
+  def scale(self, x):
+    """Scale x from -0.5 - 0.5 to 0 - 255."""
+    x = (x + 0.5) * 2**self.hparams.n_bits_x
+    return tf.cast(tf.clip_by_value(x, 0, 255), dtype=tf.uint8)
+
   @property
   def is_training(self):
     return self.hparams.mode == tf.estimator.ModeKeys.TRAIN
 
+  def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
+    del args, kwargs
+    x = features["inputs"]
+    batch_size = common_layers.shape_list(x)[0]
+    features["targets"] = tf.zeros(shape=(batch_size, 1, 1, 1))
+    _, _ = self(features)  # pylint: disable=not-callable
+
+    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
+    var_scope = tf.variable_scope("glow/body", reuse=True)
+    # If eps=None, images are sampled from the prior.
+    with arg_scope(ops, init=False), var_scope:
+      predictions, _ = glow_ops.encoder_decoder(
+          "codec", self.z_sample, self.hparams, eps=None, reverse=True)
+
+    return self.scale(predictions)
+
   def body(self, features):
     x = features["inputs"]
 
@@ -96,12 +120,13 @@ def body(self, features):
     init_op = tf.logical_and(tf.equal(global_step, 0), self.is_training)
     ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
     with arg_scope(ops, init=init_op):
-      z, encoder_objective, _ = glow_ops.encoder_decoder(
+      self.z, encoder_objective, self.eps = glow_ops.encoder_decoder(
           "codec", x, self.hparams, eps=None, reverse=False)
       objective += encoder_objective
 
-      prior_objective = glow_ops.top_prior(
-          "top_prior", z, learn_prior=self.hparams.learn_prior)
+      prior_objective, prior_dist = glow_ops.top_prior(
+          "top_prior", self.z, learn_prior=self.hparams.learn_prior)
+      self.z_sample = prior_dist.sample()
       objective += prior_objective
 
     # bits per pixel
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 025b18705..d1ca6acfe 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -545,7 +545,8 @@ def top_prior(name, x, learn_prior=False):
       prior_dist = tf.distributions.Normal(h, tf.exp(h))
     else:
       prior_dist = split_prior("top_learn_prior", h)
-    return tf.reduce_sum(prior_dist.log_prob(x), axis=[1, 2, 3])
+    objective = tf.reduce_sum(prior_dist.log_prob(x), axis=[1, 2, 3])
+    return objective, prior_dist
 
 
 @add_arg_scope
diff --git a/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
index 19c0e3bea..43557cbf7 100644
--- a/tensor2tensor/models/research/glow_test.py
+++ b/tensor2tensor/models/research/glow_test.py
@@ -18,6 +18,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import os
+import tempfile
 import numpy as np
 from tensor2tensor import problems
 from tensor2tensor.data_generators import cifar  # pylint: disable=unused-import
@@ -59,6 +61,49 @@ def test_glow(self):
         is_undefined = np.isnan(mean_obj) or np.isinf(mean_obj)
         self.assertTrue(not is_undefined)
 
+  def test_glow_inference(self):
+    hparams = glow.glow_hparams()
+    hparams.depth = 15
+    hparams.n_levels = 2
+    curr_dir = tempfile.mkdtemp()
+
+    # Training pipeline
+    with tf.Graph().as_default():
+      model = glow.Glow(hparams, tf.estimator.ModeKeys.TRAIN)
+      cifar_problem = problems.problem('image_cifar10_plain_random_shift')
+      train_dataset = cifar_problem.dataset(MODES.TRAIN)
+      one_shot = train_dataset.make_one_shot_iterator()
+      x_batch, y_batch = self.batch(one_shot)
+      features = {'inputs': x_batch, 'targets': y_batch}
+      model_path = os.path.join(curr_dir, 'model')
+
+      model(features)
+      with tf.Session() as session:
+        saver = tf.train.Saver()
+        session.run(tf.global_variables_initializer())
+        z = session.run([model.z])
+        mean_z = np.mean(z)
+        is_undefined = np.isnan(mean_z) or np.isinf(mean_z)
+        self.assertTrue(not is_undefined)
+        saver.save(session, model_path)
+
+    # Inference pipeline
+    with tf.Graph().as_default():
+      model = glow.Glow(hparams, tf.estimator.ModeKeys.PREDICT)
+      cifar_problem = problems.problem('image_cifar10_plain_random_shift')
+      test_dataset = cifar_problem.dataset(MODES.EVAL)
+      one_shot = test_dataset.make_one_shot_iterator()
+      x_batch, y_batch = self.batch(one_shot)
+      features = {'inputs': x_batch, 'targets': y_batch}
+      model_path = os.path.join(curr_dir, 'model')
+
+      predictions = model.infer(features)
+      with tf.Session() as session:
+        saver = tf.train.Saver()
+        saver.restore(session, model_path)
+        predictions_np = session.run(predictions)
+        self.assertTrue(np.all(predictions_np <= 255))
+        self.assertTrue(np.all(predictions_np >= 0))
 
 if __name__ == '__main__':
   tf.test.main()
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 12eb1aa96..170c045e3 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -63,7 +63,8 @@ def decode_hparams(overrides=""):
       shards=1,
       shard_id=0,
       num_decodes=1,
-      force_decode_length=False)
+      force_decode_length=False,
+      display_decoded_images=False)
   hp.parse(overrides)
   return hp
 

From 408625068db7f930a4434a1ce079f1503c531483 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 24 Aug 2018 18:02:12 -0700
Subject: [PATCH 0669/2720] Use unicode when logging tokens during vocab gen

PiperOrigin-RevId: 210187068
---
 tensor2tensor/data_generators/text_encoder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 2063282bc..e25ac33b3 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -814,7 +814,7 @@ def build_from_token_counts(self,
           start += len(subtoken)
         iter_time_secs = time.time() - iter_start_time
         if iter_time_secs > 0.1:
-          tf.logging.info("Processing token [{0}] took {1} seconds, consider "
+          tf.logging.info(u"Processing token [{0}] took {1} seconds, consider "
                           "setting Text2TextProblem.max_subtoken_length to a "
                           "smaller value.".format(token, iter_time_secs))
 

From a23dbb9a02b5c7e663bd036d48ca1cd9b9d66fc1 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 27 Aug 2018 10:22:14 -0700
Subject: [PATCH 0670/2720] adding specific hparams for Atari.

PiperOrigin-RevId: 210387277
---
 .../models/research/next_frame_sv2p_params.py   | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/tensor2tensor/models/research/next_frame_sv2p_params.py b/tensor2tensor/models/research/next_frame_sv2p_params.py
index 5d6b5e9eb..d4d14ffbf 100644
--- a/tensor2tensor/models/research/next_frame_sv2p_params.py
+++ b/tensor2tensor/models/research/next_frame_sv2p_params.py
@@ -51,6 +51,23 @@ def next_frame_sv2p():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_sv2p_atari():
+  """SV2P model for atari."""
+  hparams = next_frame_sv2p()
+  hparams.video_num_input_frames = 4
+  hparams.video_num_target_frames = 4
+  hparams.concatenate_actions = False
+  hparams.num_iterations_1st_stage = 15000
+  hparams.num_iterations_2nd_stage = 15000
+  hparams.latent_loss_multiplier_schedule = "noisy_linear_cosine_decay"
+  hparams.latent_loss_multiplier = 1e-3
+  hparams.anneal_end = 50000
+  hparams.preprocess_resize_frames = [96, 96]
+  hparams.information_capacity = 0.0
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_sv2p_tiny():
   """Tiny SV2P model."""

From ed75e92a9acd1379982ac1acf7594d0397e02719 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 27 Aug 2018 10:44:15 -0700
Subject: [PATCH 0671/2720] adding noisy cosine schedule + other cleanups.

PiperOrigin-RevId: 210391700
---
 tensor2tensor/layers/common_video.py          | 48 +++++++++++--------
 .../models/research/next_frame_savp_params.py |  2 +-
 2 files changed, 29 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index a89567f16..ced53ead5 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -356,27 +356,35 @@ def conv_latent_tower(images, time_axis, latent_channels=1, min_logvar=-5,
 
 def beta_schedule(schedule, global_step, final_beta, decay_start, decay_end):
   """Get KL multiplier (beta) based on the schedule."""
-  # TODO(mechcoder): Add log_annealing schedule.
+  if decay_start > decay_end:
+    raise ValueError("decay_end is smaller than decay_end.")
+
+  # Since some of the TF schedules do not support incrementing a value,
+  # in all of the schedules, we anneal the beta from final_beta to zero
+  # and then reverse it at the bottom.
   if schedule == "constant":
-    beta = tf.cond(
-        tf.less(global_step, decay_start), lambda: 0.0, lambda: final_beta)
-  elif schedule == "linear_anneal":
-    # Linearly anneal beta from 0.0 to self.hparams.latent_loss_multiplier.
-    # between self.hparams.num_iterations_2nd_stage to anneal_end.
-    # beta = latent_loss * (1 - (global_step - 2nd_stage) / (anneal_end - 2nd_stage))  # pylint:disable=line-too-long
-    if decay_start > decay_end:
-      raise ValueError("decay_end is smaller than decay_end.")
-
-    def anneal_loss(step_num):
-      step_num = tf.cast(step_num, dtype=tf.float32)
-      fraction = (float(decay_end) - step_num) / (decay_end - decay_start)
-      return final_beta * (1 - fraction)
-
-    beta = tf.case(
-        pred_fn_pairs={
-            tf.less(global_step, decay_start): lambda: 0.0,
-            tf.greater(global_step, decay_end): lambda: final_beta},
-        default=lambda: anneal_loss(global_step))
+    decayed_value = 0.0
+  elif schedule == "linear":
+    decayed_value = tf.train.polynomial_decay(
+        learning_rate=final_beta,
+        global_step=global_step - decay_start,
+        decay_steps=decay_end - decay_start,
+        end_learning_rate=0.0)
+  elif schedule == "noisy_linear_cosine_decay":
+    decayed_value = tf.train.noisy_linear_cosine_decay(
+        learning_rate=final_beta,
+        global_step=global_step - decay_start,
+        decay_steps=decay_end - decay_start)
+  # TODO(mechcoder): Add log_annealing schedule.
   else:
     raise ValueError("Unknown beta schedule.")
+
+  increased_value = final_beta - decayed_value
+  increased_value = tf.maximum(0.0, increased_value)
+
+  beta = tf.case(
+      pred_fn_pairs={
+          tf.less(global_step, decay_start): lambda: 0.0,
+          tf.greater(global_step, decay_end): lambda: final_beta},
+      default=lambda: increased_value)
   return beta
diff --git a/tensor2tensor/models/research/next_frame_savp_params.py b/tensor2tensor/models/research/next_frame_savp_params.py
index 253c0a707..e76d33717 100644
--- a/tensor2tensor/models/research/next_frame_savp_params.py
+++ b/tensor2tensor/models/research/next_frame_savp_params.py
@@ -36,6 +36,6 @@ def next_frame_savp():
   hparams.add_hparam("gan_optimization", "joint")
   hparams.target_modality = "video:l1raw"
   hparams.input_modalities = "inputs:video:l1raw"
-  hparams.latent_loss_multiplier_schedule = "linear_anneal"
+  hparams.latent_loss_multiplier_schedule = "linear"
   hparams.upsample_method = "bilinear_upsample_conv"
   return hparams

From aff492e7bfe893b7c5d8cea76ab2af93b7034d2e Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Mon, 27 Aug 2018 12:00:56 -0700
Subject: [PATCH 0672/2720] RL ablation studies.

PiperOrigin-RevId: 210406286
---
 tensor2tensor/models/research/rl.py          |  2 +-
 tensor2tensor/rl/envs/simulated_batch_env.py |  2 +-
 tensor2tensor/rl/trainer_model_based.py      | 72 +++++++++++++-------
 3 files changed, 51 insertions(+), 25 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 0c300544e..d22b600b0 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -112,7 +112,7 @@ def ppo_pong_base():
   hparams.gae_lambda = 0.985
   hparams.entropy_loss_coef = 0.003
   hparams.value_loss_coef = 1
-  hparams.optimization_epochs = 2
+  hparams.optimization_epochs = 3
   hparams.epochs_num = 1000
   hparams.num_eval_agents = 1
   hparams.policy_network = feed_forward_cnn_small_categorical_fun
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index f6515becc..51bcaceed 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -125,7 +125,7 @@ def __init__(self, environment_spec, length):
                                                FLAGS.data_dir,
                                                shuffle_files=True,
                                                hparams=hparams)
-      dataset = dataset.shuffle(buffer_size=100)
+      dataset = dataset.shuffle(buffer_size=1000)
     else:
       dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
                                                FLAGS.data_dir,
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index b8ad0bf8d..ed8c30736 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -19,7 +19,7 @@
 python -m tensor2tensor.rl.trainer_model_based \
     --output_dir=$HOME/t2t/rl_v1 \
     --loop_hparams_set=rl_modelrl_base \
-    --loop_hparams='true_env_generator_num_steps=10000,epochs=3'
+    --loop_hparams='num_real_env_frames=10000,epochs=3'
 """
 from __future__ import absolute_import
 from __future__ import division
@@ -103,7 +103,9 @@ def generate_real_env_data(problem_name, agent_policy_path, hparams, data_dir,
       "autoencoder_path": autoencoder_path,
   }):
     gym_problem = registry.problem(problem_name)
-    gym_problem.settable_num_steps = hparams.true_env_generator_num_steps
+    env_steps_per_epoch = (
+        hparams.num_real_env_frames / (hparams.epochs * (1. - 1./11.)))
+    gym_problem.settable_num_steps = env_steps_per_epoch
     gym_problem.settable_eval_phase = eval_phase
     gym_problem.generate_data(data_dir, tmp_dir)
     mean_reward = None
@@ -511,12 +513,9 @@ def combine_training_data(problem, final_data_dir, old_data_dirs,
 def rl_modelrl_base():
   return tf.contrib.training.HParams(
       epochs=6,
-      # Total frames used for training =
-      # steps * (1 - 1/11) * epochs
-      # 1/11 steps are used for evaluation data.
-      # So to use N frames set steps = N / (epochs * (1 - 1/11)).
-      # We set it to use 100k frames for training.
-      true_env_generator_num_steps=int(100000 / (6 * (1.0 - 1.0/11.0))),
+      # Total frames used for training. This will be distributed evenly across
+      # hparams.epochs.
+      num_real_env_frames=100000,
       generative_model="next_frame_basic_deterministic",
       generative_model_params="next_frame_pixel_noise",
       ppo_params="ppo_pong_base",
@@ -525,14 +524,14 @@ def rl_modelrl_base():
       simulated_env_generator_num_steps=2000,
       simulation_random_starts=True,
       intrinsic_reward_scale=0.,
-      ppo_epochs_num=200,  # This should be enough to see something
+      ppo_epochs_num=2000,  # This should be enough to see something
       # Our simulated envs do not know how to reset.
       # You should set ppo_time_limit to the value you believe that
       # the simulated env produces a reasonable output.
       ppo_time_limit=200,  # TODO(blazej): this param is unused
       # It makes sense to have ppo_time_limit=ppo_epoch_length,
       # though it is not necessary.
-      ppo_epoch_length=30,
+      ppo_epoch_length=50,
       ppo_num_agents=16,
       ppo_learning_rate=2e-4,  # Will be changed, just so it exists.
       # Whether the PPO agent should be restored from the previous iteration, or
@@ -558,7 +557,7 @@ def rl_modelrl_base_stochastic():
 def rl_modelrl_medium():
   """Small set for larger testing."""
   hparams = rl_modelrl_base()
-  hparams.true_env_generator_num_steps //= 2
+  hparams.num_real_env_frames //= 2
   return hparams
 
 
@@ -566,7 +565,7 @@ def rl_modelrl_medium():
 def rl_modelrl_25k():
   """Small set for larger testing."""
   hparams = rl_modelrl_medium()
-  hparams.true_env_generator_num_steps //= 2
+  hparams.num_real_env_frames //= 2
   return hparams
 
 
@@ -574,7 +573,7 @@ def rl_modelrl_25k():
 def rl_modelrl_short():
   """Small set for larger testing."""
   hparams = rl_modelrl_base()
-  hparams.true_env_generator_num_steps //= 5
+  hparams.num_real_env_frames //= 5
   hparams.model_train_steps //= 10
   hparams.ppo_epochs_num //= 10
   return hparams
@@ -594,7 +593,7 @@ def rl_modelrl_tiny():
   return rl_modelrl_base().override_from_dict(
       tf.contrib.training.HParams(
           epochs=2,
-          true_env_generator_num_steps=64,
+          num_real_env_frames=128,
           simulated_env_generator_num_steps=64,
           model_train_steps=2,
           ppo_epochs_num=2,
@@ -700,7 +699,7 @@ def rl_modelrl_ae_base():
 @registry.register_hparams
 def rl_modelrl_ae_25k():
   hparams = rl_modelrl_ae_base()
-  hparams.true_env_generator_num_steps //= 4
+  hparams.num_real_env_frames //= 4
   return hparams
 
 
@@ -724,7 +723,7 @@ def rl_modelrl_ae_l2_base():
 def rl_modelrl_ae_medium():
   """Medium parameter set for autoencoders."""
   hparams = rl_modelrl_ae_base()
-  hparams.true_env_generator_num_steps //= 2
+  hparams.num_real_env_frames //= 2
   return hparams
 
 
@@ -733,7 +732,7 @@ def rl_modelrl_ae_short():
   """Small parameter set for autoencoders."""
   hparams = rl_modelrl_ae_base()
   hparams.autoencoder_train_steps //= 10
-  hparams.true_env_generator_num_steps //= 5
+  hparams.num_real_env_frames //= 5
   hparams.model_train_steps //= 10
   hparams.ppo_epochs_num //= 10
   return hparams
@@ -768,11 +767,10 @@ def rl_modelrl_grid(rhp):
   rhp.set_categorical("loop.game",
                       ["breakout", "wrapped_long_pong", "freeway"])
 
-  # 100k, 50k, 25k frames
-  base = 36666
+  base = 100000
   medium = base // 2
   small = medium // 2
-  rhp.set_discrete("loop.true_env_generator_num_steps", [base, medium, small])
+  rhp.set_discrete("loop.num_real_env_frames", [base, medium, small])
 
   # Dummy parameter to get 5 runs for each configuration
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
@@ -810,10 +808,9 @@ def rl_modelrl_ae_variance(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
   rhp.set_categorical("loop.game",
                       ["breakout", "wrapped_long_pong", "freeway"])
-  # 100k, 25k frames
-  base = 36666
+  base = 100000
   small = base // 4
-  rhp.set_discrete("loop.true_env_generator_num_steps", [base, small])
+  rhp.set_discrete("loop.num_real_env_frames", [base, small])
 
 
 @registry.register_ranged_hparams
@@ -875,6 +872,35 @@ def rl_modelrl_dummy_range(rhp):
   rhp.set_float("model.moe_loss_coef", 0.01, 0.02)
 
 
+@registry.register_ranged_hparams
+def rl_modelrl_epochs_num(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.epochs", [3, 6, 12])
+
+
+@registry.register_ranged_hparams
+def rl_modelrl_ppo_epochs_num(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.ppo_epochs_num", [200, 1000, 2000, 4000])
+
+
+@registry.register_ranged_hparams
+def rl_modelrl_ppo_epoch_len(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.ppo_epoch_length", [25, 50, 100])
+
+
+@registry.register_ranged_hparams
+def rl_modelrl_num_frames(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.num_real_env_frames",
+                   [1000*el for el in [30, 100, 500, 1000]])
+
+
 def merge_unscoped_hparams(scopes_and_hparams):
   """Merge multiple HParams into one with scopes."""
   merged_values = {}

From a9e4cfb0e4efaa37a58eed5409d391586a60f73c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 27 Aug 2018 14:03:49 -0700
Subject: [PATCH 0673/2720] Fix OneHotClassLabelEncoder, add some tests

PiperOrigin-RevId: 210427289
---
 tensor2tensor/data_generators/text_encoder.py   |  2 +-
 .../data_generators/text_encoder_test.py        | 17 +++++++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index e25ac33b3..42d136a04 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -251,7 +251,7 @@ class OneHotClassLabelEncoder(ClassLabelEncoder):
   """One-hot encoder for class labels."""
 
   def encode(self, label_str, on_value=1, off_value=0):  # pylint: disable=arguments-differ
-    e = np.fill(self.vocab_size, off_value, dtype=np.int32)
+    e = np.full(self.vocab_size, off_value, dtype=np.int32)
     e[self._class_labels.index(label_str)] = on_value
     return e.tolist()
 
diff --git a/tensor2tensor/data_generators/text_encoder_test.py b/tensor2tensor/data_generators/text_encoder_test.py
index ead1de54f..26c5ce65a 100644
--- a/tensor2tensor/data_generators/text_encoder_test.py
+++ b/tensor2tensor/data_generators/text_encoder_test.py
@@ -363,5 +363,22 @@ def gen():
     self.assertEqual(corpus, reconstructed_corpus)
 
 
+class OneHotClassLabelEncoderTest(tf.test.TestCase):
+
+  def test_one_hot_encode(self):
+    encoder = text_encoder.OneHotClassLabelEncoder(
+        class_labels=["zero", "one", "two"])
+    self.assertEqual(encoder.encode("zero"), [1, 0, 0])
+    self.assertEqual(encoder.encode("one"), [0, 1, 0])
+    self.assertEqual(encoder.encode("two"), [0, 0, 1])
+
+  def test_one_hot_decode(self):
+    encoder = text_encoder.OneHotClassLabelEncoder(
+        class_labels=["zero", "one", "two"])
+    self.assertEqual(encoder.decode([1, 0, 0]), "zero")
+    self.assertEqual(encoder.decode([0, 1, 0]), "one")
+    self.assertEqual(encoder.decode([0, 0, 1]), "two")
+
+
 if __name__ == "__main__":
   tf.test.main()

From 40a7dc0f47b33770c2177c75ee78e16e8d493f09 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 27 Aug 2018 14:15:49 -0700
Subject: [PATCH 0674/2720] BugFix: UserWarning: Flag --problem has a non-None
 default value; therefore, mark_flag_as_required will pass even if flag is not
 specified in the command line!

PiperOrigin-RevId: 210429581
---
 tensor2tensor/bin/build_vocab.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/build_vocab.py b/tensor2tensor/bin/build_vocab.py
index f1235e28e..0638ae47d 100644
--- a/tensor2tensor/bin/build_vocab.py
+++ b/tensor2tensor/bin/build_vocab.py
@@ -40,7 +40,7 @@
 flags.DEFINE_string("tmp_dir", "/tmp/t2t/tmp_dir",
                     "Temporary storage directory.")
 
-flags.DEFINE_string("problem", "",
+flags.DEFINE_string("problem", None,
                     "Problem to generate the vocabulary file for.")
 
 flags.mark_flag_as_required("problem")

From 36b6f4920f3e70c8fb33581a4d6e1b78116e752d Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 27 Aug 2018 14:59:24 -0700
Subject: [PATCH 0675/2720] Speed up RL test.

PiperOrigin-RevId: 210437369
---
 tensor2tensor/rl/trainer_model_based.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index ed8c30736..27a0959d7 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -608,6 +608,7 @@ def rl_modelrl_tiny():
 def rl_modelrl_tiny_stochastic():
   """Tiny setting with a stochastic next-frame model."""
   hparams = rl_modelrl_tiny()
+  hparams.epochs = 1  # Too slow with 2 for regular runs.
   hparams.generative_model = "next_frame_basic_stochastic"
   hparams.generative_model_params = "next_frame_basic_stochastic"
   return hparams

From d1946d4c258eeea9f8f89ec9e5814e6429432259 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 27 Aug 2018 15:06:16 -0700
Subject: [PATCH 0676/2720] Guard matplotlib import. Also use agg backend so we
 don't need to link TK with matplotlib.

PiperOrigin-RevId: 210438548
---
 tensor2tensor/data_generators/image_utils.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index f1e898f7c..18dd3f0b4 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -21,7 +21,6 @@
 
 import io
 import os
-import matplotlib.pyplot as plt
 import numpy as np
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -33,6 +32,13 @@
 import tensorflow as tf
 
 
+def matplotlib_pyplot():
+  import matplotlib  # pylint: disable=g-import-not-at-top
+  matplotlib.use("agg")
+  import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top
+  return plt
+
+
 def image_to_tf_summary_value(image, tag):
   """Converts a NumPy image to a tf.Summary.Value object.
 
@@ -45,7 +51,7 @@ def image_to_tf_summary_value(image, tag):
   curr_image = np.asarray(image, dtype=np.uint8)
   height, width, n_channels = curr_image.shape
   s = io.BytesIO()
-  plt.imsave(s, curr_image, format="png")
+  matplotlib_pyplot().imsave(s, curr_image, format="png")
   img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(),
                              height=height, width=width,
                              colorspace=n_channels)

From 815673a6a5273af09d6b1f5ca9b431ce40402841 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 27 Aug 2018 15:19:30 -0700
Subject: [PATCH 0677/2720] internal merge of PR #1014

PiperOrigin-RevId: 210440979
---
 tensor2tensor/data_generators/gym_problems.py | 2 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py    | 4 +---
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index a594d9ae9..0fda7aefc 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -49,7 +49,7 @@
 def standard_atari_env_spec(env):
   """Parameters of environment specification."""
   standard_wrappers = [[tf_atari_wrappers.RewardClippingWrapper, {}],
-                       [tf_atari_wrappers.StackAndSkipWrapper, {"skip": 4}]]
+                       [tf_atari_wrappers.StackWrapper, {"history": 4}]]
   env_lambda = None
   if isinstance(env, str):
     env_lambda = lambda: gym.make(env)
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 6ca76de3f..fcf744f0c 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -59,6 +59,7 @@ def _reset_non_empty(self, indices):
     with tf.control_dependencies([assign_op]):
       return tf.identity(new_values)
 
+
 class RewardClippingWrapper(WrapperBase):
   """ Reward clipping wrapper.
       The rewards are clipped to -1, 0, 1
@@ -66,9 +67,6 @@ class RewardClippingWrapper(WrapperBase):
       of rl algorithms
   """
 
-  def __init__(self, batch_env):
-    super(RewardClippingWrapper, self).__init__(batch_env)
-
   def simulate(self, action):
     reward, done = self._batch_env.simulate(action)
     with tf.control_dependencies([reward, done]):

From 06d8728cb84c41f028ccd3299335bcaacbde313c Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 27 Aug 2018 15:25:15 -0700
Subject: [PATCH 0678/2720] adding support for information cap.

PiperOrigin-RevId: 210442013
---
 tensor2tensor/models/research/next_frame_base_vae.py         | 3 +++
 tensor2tensor/models/research/next_frame_basic_stochastic.py | 1 +
 2 files changed, 4 insertions(+)

diff --git a/tensor2tensor/models/research/next_frame_base_vae.py b/tensor2tensor/models/research/next_frame_base_vae.py
index e9c4bad6e..ff27d76f6 100644
--- a/tensor2tensor/models/research/next_frame_base_vae.py
+++ b/tensor2tensor/models/research/next_frame_base_vae.py
@@ -55,6 +55,9 @@ def get_extra_loss(self, mean, std):
     tf.summary.histogram("posterior_mean", mean)
     tf.summary.histogram("posterior_std", std)
     tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
+    # information capacity from "Understanding disentangling in beta-VAE"
+    if self.hparams.information_capacity > 0.0:
+      kl_loss = tf.abs(kl_loss - self.hparams.information_capacity)
     return beta * kl_loss
 
   def construct_latent_tower(self, images, time_axis):
diff --git a/tensor2tensor/models/research/next_frame_basic_stochastic.py b/tensor2tensor/models/research/next_frame_basic_stochastic.py
index 311cca896..9585c1745 100644
--- a/tensor2tensor/models/research/next_frame_basic_stochastic.py
+++ b/tensor2tensor/models/research/next_frame_basic_stochastic.py
@@ -68,6 +68,7 @@ def next_frame_basic_stochastic():
   hparams.add_hparam("latent_loss_multiplier_schedule", "constant")
   hparams.add_hparam("latent_num_frames", 0)  # 0 means use all frames.
   hparams.add_hparam("anneal_end", 100000)
+  hparams.add_hparam("information_capacity", 0.0)
   return hparams
 
 

From 6ba6f953e8460812d1bba19813f871e36f71fa83 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 27 Aug 2018 15:31:30 -0700
Subject: [PATCH 0679/2720] replacing all batch_norms  with layer_norms.

PiperOrigin-RevId: 210443127
---
 tensor2tensor/layers/common_video.py             | 10 +++-------
 tensor2tensor/models/research/next_frame_sv2p.py |  9 +++------
 2 files changed, 6 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index ced53ead5..1876e9ce3 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -322,17 +322,13 @@ def conv_latent_tower(images, time_axis, latent_channels=1, min_logvar=-5,
     x = common_layers.make_even_size(x)
     x = tfl.conv2d(x, conv_size[0], [3, 3], strides=(2, 2),
                    padding="SAME", activation=tf.nn.relu, name="latent_conv1")
-    x = tfcl.batch_norm(x, updates_collections=None,
-                        is_training=is_training, scope="latent_bn1")
-    x = common_layers.make_even_size(x)
+    x = tfcl.layer_norm(x)
     x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
                    padding="SAME", activation=tf.nn.relu, name="latent_conv2")
-    x = tfcl.batch_norm(x, updates_collections=None,
-                        is_training=is_training, scope="latent_bn2")
+    x = tfcl.layer_norm(x)
     x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(1, 1),
                    padding="SAME", activation=tf.nn.relu, name="latent_conv3")
-    x = tfcl.batch_norm(x, updates_collections=None,
-                        is_training=is_training, scope="latent_bn3")
+    x = tfcl.layer_norm(x)
 
     nc = latent_channels
     mean = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/research/next_frame_sv2p.py
index 800730abc..fbfae182d 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/research/next_frame_sv2p.py
@@ -225,12 +225,10 @@ def reward_prediction(self, input_images, input_reward, action, latent):
 
     with tf.variable_scope("reward_pred", reuse=tf.AUTO_REUSE):
       x = tf.concat(input_images, axis=3)
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="reward_bn0")
+      x = tfcl.layer_norm(x)
       x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
                      activation=tf.nn.relu, name="reward_conv1")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="reward_bn1")
+      x = tfcl.layer_norm(x)
 
       # Inject additional inputs
       if action is not None:
@@ -246,8 +244,7 @@ def reward_prediction(self, input_images, input_reward, action, latent):
 
       x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(2, 2),
                      activation=tf.nn.relu, name="reward_conv2")
-      x = tfcl.batch_norm(x, updates_collections=None,
-                          is_training=self.is_training, scope="reward_bn2")
+      x = tfcl.layer_norm(x)
       x = tfl.conv2d(x, conv_size[3], [3, 3], strides=(2, 2),
                      activation=tf.nn.relu, name="reward_conv3")
       return x

From 68b10abe5e38f0810eecc72183bd2047b262d759 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 27 Aug 2018 16:19:47 -0700
Subject: [PATCH 0680/2720] Add Eager tests to T2T layers.

PiperOrigin-RevId: 210451153
---
 tensor2tensor/layers/common_attention.py      |   2 +-
 tensor2tensor/layers/common_attention_test.py | 197 ++++----
 tensor2tensor/layers/common_layers_test.py    | 474 +++++++++---------
 tensor2tensor/layers/common_video_test.py     |  79 +--
 tensor2tensor/layers/discretization_test.py   |  94 ++--
 tensor2tensor/layers/modalities_test.py       |  32 +-
 .../mesh_tensorflow/mesh_tensorflow_test.py   |   2 +-
 .../mesh_tensorflow/mtf_layers_test.py        |   4 +-
 8 files changed, 426 insertions(+), 458 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 8a4a6c2f4..a35081581 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -701,7 +701,7 @@ def add_positional_embedding_nd(x, max_length, name):
           name + "_%d" % i,
           shape,
           initializer=tf.random_normal_initializer(0, depth**-0.5))
-      var *= depth**0.5
+      var = var * depth**0.5
       x += tf.slice(var, start, size)
     return x
 
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 0585e7ef8..926337105 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -29,43 +29,42 @@
 
 class CommonAttentionTest(parameterized.TestCase, tf.test.TestCase):
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testAddPositionalEmbedding(self):
     x = np.random.rand(5, 3, 12)
-    with self.test_session() as session:
-      y = common_attention.add_positional_embedding(
-          tf.constant(x, dtype=tf.float32),
-          max_length=4,
-          name="pos_embedding")
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_attention.add_positional_embedding(
+        tf.constant(x, dtype=tf.float32),
+        max_length=4,
+        name="pos_embedding")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 3, 12))
 
   @parameterized.parameters(
-      ((5, 3, 12),),
-      ((5, 5, 5, 12),),
-      ((5, 3, 3, 3, 12),),
+      {"input_shape": (5, 3, 12)},
+      {"input_shape": (5, 5, 5, 12)},
+      {"input_shape": (5, 3, 3, 3, 12)},
   )
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testAddPositionalEmbeddingNd(self, input_shape):
     x = np.random.rand(*input_shape)
-    with self.test_session() as session:
-      y = common_attention.add_positional_embedding_nd(
-          tf.constant(x, dtype=tf.float32),
-          max_length=5,
-          name="pos_embedding")
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_attention.add_positional_embedding_nd(
+        tf.constant(x, dtype=tf.float32),
+        max_length=5,
+        name="pos_embedding")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, input_shape)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDotProductAttention(self):
     x = np.random.rand(5, 7, 12, 32)
     y = np.random.rand(5, 7, 12, 32)
-    with self.test_session() as session:
-      a = common_attention.dot_product_attention(
-          tf.constant(x, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32), None)
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_attention.dot_product_attention(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32), None)
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
   @parameterized.named_parameters(
@@ -84,11 +83,10 @@ def testMaskedWithinBlockLocalAttention1D(self, batch, heads, length,
     v = tf.random_normal([batch, heads, length, depth_v])
     output = common_attention.masked_within_block_local_attention_1d(
         q, k, v, block_length=block_length)
-    with self.test_session() as session:
-      if isinstance(batch, tf.Tensor):
-        batch, res = session.run([batch, output])
-      else:
-        res = session.run(output)
+    if isinstance(batch, tf.Tensor):
+      batch, res = self.evaluate([batch, output])
+    else:
+      res = self.evaluate(output)
 
     self.assertEqual(res.shape, (batch, heads, length, depth_v))
 
@@ -108,11 +106,10 @@ def testMaskedLocalAttention1D(self, batch, heads, length, depth_k, depth_v,
     v = tf.random_normal([batch, heads, length, depth_v])
     output = common_attention.masked_local_attention_1d(
         q, k, v, block_length=block_length)
-    with self.test_session() as session:
-      if isinstance(batch, tf.Tensor):
-        batch, res = session.run([batch, output])
-      else:
-        res = session.run(output)
+    if isinstance(batch, tf.Tensor):
+      batch, res = self.evaluate([batch, output])
+    else:
+      res = self.evaluate(output)
 
     self.assertEqual(res.shape, (batch, heads, length, depth_v))
 
@@ -159,11 +156,10 @@ def testLocalUnmaskedAttention1D(self, batch, heads, length,
     v = tf.random_normal([batch, heads, length, depth_v])
     output = common_attention.local_attention_1d(
         q, k, v, block_length=block_length, filter_width=3)
-    with self.test_session() as session:
-      if isinstance(batch, tf.Tensor):
-        batch, res = session.run([batch, output])
-      else:
-        res = session.run(output)
+    if isinstance(batch, tf.Tensor):
+      batch, res = self.evaluate([batch, output])
+    else:
+      res = self.evaluate(output)
 
     self.assertEqual(res.shape, (batch, heads, length, depth_v))
 
@@ -238,6 +234,7 @@ def testMultiheadSelfAttentionMemoryEfficient(self):
     self.assertAllClose(dnorm_bias, dnorm_bias_f)
     self.assertAllClose(dx, dx_f)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def test2dGatherAndScatterInvertibility(self):
     """2d gather and scatter invertibility test."""
     batch_size = 2
@@ -247,25 +244,23 @@ def test2dGatherAndScatterInvertibility(self):
     depth = 8
     query_shape = (2, 3)
     x = np.random.rand(batch_size, num_heads, height, width, depth)
-    with self.test_session() as session:
-      x_indices = common_attention.gather_indices_2d(
-          x, query_shape, query_shape)
-      gathered_x = common_attention.gather_blocks_2d(x, x_indices)
-      x_shape = tf.constant([batch_size, num_heads, height, width, depth])
-      scattered_x = common_attention.scatter_blocks_2d(
-          gathered_x, x_indices, x_shape)
-      session.run(tf.global_variables_initializer())
-      res = session.run(scattered_x)
+    x_indices = common_attention.gather_indices_2d(
+        x, query_shape, query_shape)
+    gathered_x = common_attention.gather_blocks_2d(x, x_indices)
+    x_shape = tf.constant([batch_size, num_heads, height, width, depth])
+    scattered_x = common_attention.scatter_blocks_2d(
+        gathered_x, x_indices, x_shape)
+    res = self.evaluate(scattered_x)
     self.assertAllClose(x, res)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def test2dBlockRasterScanMask(self):
     """Testing the 2d block raster scan mask."""
     query_shape = (2, 3)
     memory_flange = (2, 1)
-    with self.test_session() as session:
-      mask = common_attention.make_2d_block_raster_mask(
-          query_shape, memory_flange)
-      res = session.run(mask)
+    mask = common_attention.make_2d_block_raster_mask(
+        query_shape, memory_flange)
+    res = self.evaluate(mask)
     correct_mask = np.array(
         [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0,
           1.0, 0.0, 1.0, 1.0, 1.0, 1.0],
@@ -281,6 +276,7 @@ def test2dBlockRasterScanMask(self):
           1.0, 0.0, 0.0, 0.0, 0.0, 1.0]])
     self.assertAllClose(correct_mask, res)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def test2dGather(self):
     """Testing 2d index gather and block gather functions."""
     batch_size = 2
@@ -312,14 +308,14 @@ def test2dGather(self):
                             y[1, 1, correct_indices[2]],
                             y[1, 1, correct_indices[3]]]]]
 
-    with self.test_session() as session:
-      x_indices = common_attention.gather_indices_2d(
-          x, query_shape, query_shape)
-      gathered_x = common_attention.gather_blocks_2d(x, x_indices)
-      x_indices, gathered_x = session.run([x_indices, gathered_x])
+    x_indices = common_attention.gather_indices_2d(
+        x, query_shape, query_shape)
+    gathered_x = common_attention.gather_blocks_2d(x, x_indices)
+    x_indices, gathered_x = self.evaluate([x_indices, gathered_x])
     self.assertAllEqual(correct_indices, x_indices)
     self.assertAllClose(correct_gathered_x, gathered_x)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testGetMemoryRegion(self):
     """Testing the function that gathers the flanged memory region."""
     np.set_printoptions(threshold=np.inf)
@@ -385,19 +381,18 @@ def testGetMemoryRegion(self):
                           y[1, 1, [12, 13, 14, 18, 19, 20]],
                           y[1, 1, [15, 16, 17, 21, 22, 23]]]]]
     correct_x_center = np.array(correct_x_center)
-    with self.test_session() as session:
-      x_indices = common_attention.gather_indices_2d(
-          x, query_shape, query_shape)
-      x_flange, x_center = common_attention.get_memory_region(
-          tf.constant(x, dtype=tf.float32),
-          query_shape,
-          memory_flange,
-          x_indices)
-      session.run(tf.global_variables_initializer())
-      [x_flange, x_center] = session.run([x_flange, x_center])
+    x_indices = common_attention.gather_indices_2d(
+        x, query_shape, query_shape)
+    x_flange, x_center = common_attention.get_memory_region(
+        tf.constant(x, dtype=tf.float32),
+        query_shape,
+        memory_flange,
+        x_indices)
+    [x_flange, x_center] = self.evaluate([x_flange, x_center])
     self.assertAllClose(correct_x_flange, x_flange)
     self.assertAllClose(correct_x_center, x_center)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testGetShiftedCenterBlocks(self):
     """Testing the function that gathers the flanged memory region."""
     np.set_printoptions(threshold=np.inf)
@@ -454,44 +449,43 @@ def testGetShiftedCenterBlocks(self):
                                             y[1, 1, [15, 16, 17, 21, 22]]),
                                            axis=0)]]]
     correct_gathered_x = np.array(correct_gathered_x)
-    with self.test_session() as session:
-      x_indices = common_attention.gather_indices_2d(
-          x, query_shape, query_shape)
-      gathered_x = common_attention.get_shifted_center_blocks(
-          tf.constant(x, dtype=tf.float32),
-          x_indices)
-      session.run(tf.global_variables_initializer())
-      x_indices, gathered_x = session.run([x_indices, gathered_x])
+    x_indices = common_attention.gather_indices_2d(
+        x, query_shape, query_shape)
+    gathered_x = common_attention.get_shifted_center_blocks(
+        tf.constant(x, dtype=tf.float32),
+        x_indices)
+    x_indices, gathered_x = self.evaluate([x_indices, gathered_x])
     self.assertAllClose(correct_gathered_x, gathered_x)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDotProductAttentionRelative(self):
     x = np.random.rand(5, 7, 12, 32)
     y = np.random.rand(5, 7, 12, 32)
-    with self.test_session() as session:
-      a = common_attention.dot_product_attention_relative(
-          tf.constant(x, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          None,
-          max_relative_position=3)
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_attention.dot_product_attention_relative(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        None,
+        max_relative_position=3)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDotProductUnMaskedAttentionRelativeV2(self):
     x = np.random.rand(5, 7, 12, 32)
     y = np.random.rand(5, 7, 12, 32)
-    with self.test_session() as session:
-      a = common_attention.dot_product_unmasked_self_attention_relative_v2(
-          tf.constant(x, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          tf.constant(y, dtype=tf.float32),
-          None,
-          35)
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_attention.dot_product_unmasked_self_attention_relative_v2(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        None,
+        35)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testBiasBatchCoordinates(self):
     """Testing the batch coordinates mask."""
     q = tf.constant([0, 0, 1, 1, 1, 1, 2, 2, 2], dtype=tf.int32)
@@ -513,14 +507,9 @@ def testBiasBatchCoordinates(self):
     ], np.float32) * -1e9
 
     bias = common_attention.attention_bias_coordinates(q, k)
+    self.assertAllClose(self.evaluate(bias), ground_truth)
 
-    with self.test_session() as session:
-      session.run(tf.global_variables_initializer())
-      self.assertAllClose(
-          bias.eval(),
-          ground_truth,
-      )
-
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testBiasFuture(self):
     """Testing the sequence order mask."""
     q = tf.constant([0, 1, 2, 3, 0, 1, 2, 0, 1], dtype=tf.int32)
@@ -542,13 +531,7 @@ def testBiasFuture(self):
     ], np.float32) * -1e9
 
     bias = common_attention.attention_bias_future(q, k)
-
-    with self.test_session() as session:
-      session.run(tf.global_variables_initializer())
-      self.assertAllClose(
-          bias.eval(),
-          ground_truth,
-      )
+    self.assertAllClose(self.evaluate(bias), ground_truth)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index d4a6df588..c7031b2e5 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -28,6 +28,7 @@
 
 class CommonLayersTest(parameterized.TestCase, tf.test.TestCase):
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testIndexLastDimWithIndices(self):
     x = np.array([[2., 3., 4., 5.],
                   [6., 7., 8., 9.]])
@@ -35,30 +36,29 @@ def testIndexLastDimWithIndices(self):
     x_idx = common_layers.index_last_dim_with_indices(x, indices)
 
     expected = np.array([4., 6.])
-    with self.test_session() as sess:
-      self.assertAllEqual(expected, sess.run(x_idx))
+    self.assertAllEqual(expected, self.evaluate(x_idx))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testSaturatingSigmoid(self):
     x = np.array([-120.0, -100.0, 0.0, 100.0, 120.0], dtype=np.float32)
-    with self.test_session() as session:
-      y = common_layers.saturating_sigmoid(tf.constant(x))
-      res = session.run(y)
+    y = common_layers.saturating_sigmoid(tf.constant(x))
+    res = self.evaluate(y)
     self.assertAllClose(res, [0.0, 0.0, 0.5, 1.0, 1.0])
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testFlatten4D3D(self):
     x = np.random.random_integers(1, high=8, size=(3, 5, 2))
-    with self.test_session() as session:
-      y = common_layers.flatten4d3d(common_layers.embedding(x, 10, 7))
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.flatten4d3d(common_layers.embedding(x, 10, 7))
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (3, 5 * 2, 7))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testEmbedding(self):
     x = np.random.random_integers(1, high=8, size=(3, 5))
-    with self.test_session() as session:
-      y = common_layers.embedding(x, 10, 16)
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.embedding(x, 10, 16)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (3, 5, 16))
 
   def testShakeShake(self):
@@ -66,132 +66,129 @@ def testShakeShake(self):
     with self.test_session() as session:
       x = tf.constant(x, dtype=tf.float32)
       y = common_layers.shakeshake([x, x, x, x, x])
-      session.run(tf.global_variables_initializer())
       inp, res = session.run([x, y])
     self.assertAllClose(res, inp)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testConv(self):
     x = np.random.rand(5, 7, 1, 11)
-    with self.test_session() as session:
-      y = common_layers.conv(tf.constant(x, dtype=tf.float32), 13, (3, 1))
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.conv(tf.constant(x, dtype=tf.float32), 13, (3, 1))
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 5, 1, 13))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testConv1d(self):
     x = np.random.rand(5, 7, 11)
-    with self.test_session() as session:
-      y = common_layers.conv1d(tf.constant(x, dtype=tf.float32), 13, 1)
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.conv1d(tf.constant(x, dtype=tf.float32), 13, 1)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 13))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testSeparableConv(self):
     x = np.random.rand(5, 7, 1, 11)
-    with self.test_session() as session:
-      y = common_layers.separable_conv(
-          tf.constant(x, dtype=tf.float32), 13, (3, 1))
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.separable_conv(
+        tf.constant(x, dtype=tf.float32), 13, (3, 1))
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 5, 1, 13))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testSubSeparableConv(self):
     for sep in [0, 1, 2, 4]:
       x = np.random.rand(5, 7, 1, 12)
-      with self.test_session() as session:
-        with tf.variable_scope("sep_%d" % sep):
-          y = common_layers.subseparable_conv(
-              tf.constant(x, dtype=tf.float32), 16, (3, 1), separability=sep)
-        session.run(tf.global_variables_initializer())
-        res = session.run(y)
+      with tf.variable_scope("sep_%d" % sep):
+        y = common_layers.subseparable_conv(
+            tf.constant(x, dtype=tf.float32), 16, (3, 1), separability=sep)
+      self.evaluate(tf.global_variables_initializer())
+      res = self.evaluate(y)
       self.assertEqual(res.shape, (5, 5, 1, 16))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testConvBlock(self):
     x = np.random.rand(5, 7, 1, 11)
-    with self.test_session() as session:
-      y = common_layers.conv_block(
-          tf.constant(x, dtype=tf.float32),
-          13, [(1, (3, 3)), (1, (3, 3))],
-          padding="SAME",
-          normalizer_fn=common_layers.noam_norm)
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.conv_block(
+        tf.constant(x, dtype=tf.float32),
+        13, [(1, (3, 3)), (1, (3, 3))],
+        padding="SAME",
+        normalizer_fn=common_layers.noam_norm)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 1, 13))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testSeparableConvBlock(self):
     x = np.random.rand(5, 7, 1, 11)
-    with self.test_session() as session:
-      y = common_layers.separable_conv_block(
-          tf.constant(x, dtype=tf.float32),
-          13, [(1, (3, 3)), (1, (3, 3))],
-          padding="SAME")
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.separable_conv_block(
+        tf.constant(x, dtype=tf.float32),
+        13, [(1, (3, 3)), (1, (3, 3))],
+        padding="SAME")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 1, 13))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testSubSeparableConvBlock(self):
     for sep in [0, 1, 2, 4]:
       x = np.random.rand(5, 7, 1, 12)
-      with self.test_session() as session:
-        with tf.variable_scope("sep_%d" % sep):
-          y = common_layers.subseparable_conv_block(
-              tf.constant(x, dtype=tf.float32),
-              16, [(1, (3, 3)), (1, (3, 3))],
-              padding="SAME",
-              separability=sep)
-        session.run(tf.global_variables_initializer())
-        res = session.run(y)
+      with tf.variable_scope("sep_%d" % sep):
+        y = common_layers.subseparable_conv_block(
+            tf.constant(x, dtype=tf.float32),
+            16, [(1, (3, 3)), (1, (3, 3))],
+            padding="SAME",
+            separability=sep)
+      self.evaluate(tf.global_variables_initializer())
+      res = self.evaluate(y)
       self.assertEqual(res.shape, (5, 7, 1, 16))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testPool(self):
     x = np.random.rand(5, 8, 1, 11)
-    with self.test_session() as session:
-      y = common_layers.pool(
-          tf.constant(x, dtype=tf.float32), (2, 2), "AVG", "SAME")
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.pool(
+        tf.constant(x, dtype=tf.float32), (2, 2), "AVG", "SAME")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 8, 1, 11))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testConvBlockDownsample(self):
     x = np.random.rand(5, 7, 1, 11)
-    with self.test_session() as session:
-      y = common_layers.conv_block_downsample(
-          tf.constant(x, dtype=tf.float32), (3, 1), (2, 1), "SAME")
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.conv_block_downsample(
+        tf.constant(x, dtype=tf.float32), (3, 1), (2, 1), "SAME")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 4, 1, 27))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testSimpleAttention(self):
     x = np.random.rand(5, 7, 1, 11)
     y = np.random.rand(5, 9, 1, 11)
-    with self.test_session() as session:
-      a = common_layers.simple_attention(
-          tf.constant(x, dtype=tf.float32), tf.constant(y, dtype=tf.float32))
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_layers.simple_attention(
+        tf.constant(x, dtype=tf.float32), tf.constant(y, dtype=tf.float32))
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 1, 11))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testGetTimingSignal(self):
     length = 7
     num_timescales = 10
-    with self.test_session() as session:
-      a = common_layers.get_timing_signal(length, num_timescales=num_timescales)
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_layers.get_timing_signal(length, num_timescales=num_timescales)
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (length, 2 * num_timescales))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testAddTimingSignal(self):
     batch = 5
     length = 7
     height = 3
     depth = 35
     x = np.random.rand(batch, length, height, depth)
-    with self.test_session() as session:
-      a = common_layers.add_timing_signal(tf.constant(x, dtype=tf.float32))
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_layers.add_timing_signal(tf.constant(x, dtype=tf.float32))
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (batch, length, height, depth))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testAttention1D(self):
     batch = 5
     target_length = 7
@@ -204,36 +201,35 @@ def testAttention1D(self):
     source = np.random.rand(batch, source_length, source_depth)
     target = np.random.rand(batch, target_length, target_depth)
     mask = np.random.rand(batch, target_length, source_length)
-    with self.test_session() as session:
-      a = common_layers.attention_1d_v0(
-          tf.constant(source, dtype=tf.float32),
-          tf.constant(target, dtype=tf.float32), attention_size, output_size,
-          num_heads, tf.constant(mask, dtype=tf.float32))
-      session.run(tf.global_variables_initializer())
-      res = session.run(a)
+    a = common_layers.attention_1d_v0(
+        tf.constant(source, dtype=tf.float32),
+        tf.constant(target, dtype=tf.float32), attention_size, output_size,
+        num_heads, tf.constant(mask, dtype=tf.float32))
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
     self.assertEqual(res.shape, (batch, target_length, output_size))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testMultiscaleConvSum(self):
     x = np.random.rand(5, 9, 1, 11)
-    with self.test_session() as session:
-      y = common_layers.multiscale_conv_sum(
-          tf.constant(x, dtype=tf.float32),
-          13, [((1, 1), (5, 5)), ((2, 2), (3, 3))],
-          "AVG",
-          padding="SAME")
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.multiscale_conv_sum(
+        tf.constant(x, dtype=tf.float32),
+        13, [((1, 1), (5, 5)), ((2, 2), (3, 3))],
+        "AVG",
+        padding="SAME")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 9, 1, 13))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testConvGRU(self):
     x = np.random.rand(5, 7, 3, 11)
-    with self.test_session() as session:
-      y = common_layers.conv_gru(tf.constant(x, dtype=tf.float32), (1, 3), 11)
-      z = common_layers.conv_gru(
-          tf.constant(x, dtype=tf.float32), (1, 3), 11, padding="LEFT")
-      session.run(tf.global_variables_initializer())
-      res1 = session.run(y)
-      res2 = session.run(z)
+    y = common_layers.conv_gru(tf.constant(x, dtype=tf.float32), (1, 3), 11)
+    z = common_layers.conv_gru(
+        tf.constant(x, dtype=tf.float32), (1, 3), 11, padding="LEFT")
+    self.evaluate(tf.global_variables_initializer())
+    res1 = self.evaluate(y)
+    res2 = self.evaluate(z)
     self.assertEqual(res1.shape, (5, 7, 3, 11))
     self.assertEqual(res2.shape, (5, 7, 3, 11))
 
@@ -245,109 +241,109 @@ def testSRU(self):
       res = session.run(y)
     self.assertEqual(res.shape, (5, 7, 3, 11))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testLayerNorm(self):
     x = np.random.rand(5, 7, 11)
-    with self.test_session() as session:
-      y = common_layers.layer_norm(tf.constant(x, dtype=tf.float32), 11)
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.layer_norm(tf.constant(x, dtype=tf.float32), 11)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 11))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testGroupNorm(self):
     x = np.random.rand(5, 7, 3, 16)
-    with self.test_session() as session:
-      y = common_layers.group_norm(tf.constant(x, dtype=tf.float32))
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.group_norm(tf.constant(x, dtype=tf.float32))
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 3, 16))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testConvLSTM(self):
     x = np.random.rand(5, 7, 11, 13)
-    with self.test_session() as session:
-      y = common_layers.conv_lstm(tf.constant(x, dtype=tf.float32), (1, 3), 13)
-      session.run(tf.global_variables_initializer())
-      res = session.run(y)
+    y = common_layers.conv_lstm(tf.constant(x, dtype=tf.float32), (1, 3), 13)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 11, 13))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testPadToSameLength(self):
     x1 = np.random.rand(5, 7, 11)
     x2 = np.random.rand(5, 9, 11)
-    with self.test_session() as session:
-      a, b = common_layers.pad_to_same_length(
-          tf.constant(x1, dtype=tf.float32), tf.constant(x2, dtype=tf.float32))
-      c, d = common_layers.pad_to_same_length(
-          tf.constant(x1, dtype=tf.float32),
-          tf.constant(x2, dtype=tf.float32),
-          final_length_divisible_by=4)
-      res1, res2 = session.run([a, b])
-      res1a, res2a = session.run([c, d])
+    a, b = common_layers.pad_to_same_length(
+        tf.constant(x1, dtype=tf.float32), tf.constant(x2, dtype=tf.float32))
+    c, d = common_layers.pad_to_same_length(
+        tf.constant(x1, dtype=tf.float32),
+        tf.constant(x2, dtype=tf.float32),
+        final_length_divisible_by=4)
+    res1, res2 = self.evaluate([a, b])
+    res1a, res2a = self.evaluate([c, d])
     self.assertEqual(res1.shape, (5, 9, 11))
     self.assertEqual(res2.shape, (5, 9, 11))
     self.assertEqual(res1a.shape, (5, 12, 11))
     self.assertEqual(res2a.shape, (5, 12, 11))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testShiftLeft(self):
     x1 = np.zeros((5, 7, 1, 11))
     x1[:, 0, :] = np.ones_like(x1[:, 0, :])
     expected = np.zeros((5, 7, 1, 11))
     expected[:, 1, :] = np.ones_like(expected[:, 1, :])
-    with self.test_session() as session:
-      a = common_layers.shift_right(tf.constant(x1, dtype=tf.float32))
-      actual = session.run(a)
+    a = common_layers.shift_right(tf.constant(x1, dtype=tf.float32))
+    actual = self.evaluate(a)
     self.assertAllEqual(actual, expected)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testConvStride2MultiStep(self):
     x1 = np.random.rand(5, 32, 16, 11)
-    with self.test_session() as session:
-      a = common_layers.conv_stride2_multistep(
-          tf.constant(x1, dtype=tf.float32), 4, 16)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(a[0])
+    a = common_layers.conv_stride2_multistep(
+        tf.constant(x1, dtype=tf.float32), 4, 16)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(a[0])
     self.assertEqual(actual.shape, (5, 2, 1, 16))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDeconvStride2MultiStep(self):
     x1 = np.random.rand(5, 2, 1, 11)
-    with self.test_session() as session:
-      a = common_layers.deconv_stride2_multistep(
-          tf.constant(x1, dtype=tf.float32), 4, 16)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(a)
+    a = common_layers.deconv_stride2_multistep(
+        tf.constant(x1, dtype=tf.float32), 4, 16)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(a)
     self.assertEqual(actual.shape, (5, 32, 1, 16))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testApplyNormLayer(self):
-    with self.test_session() as session:
-      x1 = np.random.rand(5, 2, 1, 11)
-      x2 = common_layers.apply_norm(
-          tf.constant(x1, dtype=tf.float32), "layer", depth=11, epsilon=1e-6)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(x2)
+    x1 = np.random.rand(5, 2, 1, 11)
+    x2 = common_layers.apply_norm(
+        tf.constant(x1, dtype=tf.float32), "layer", depth=11, epsilon=1e-6)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testApplyNormNoam(self):
-    with self.test_session() as session:
-      x1 = np.random.rand(5, 2, 1, 11)
-      x2 = common_layers.apply_norm(
-          tf.constant(x1, dtype=tf.float32), "noam", depth=11, epsilon=1e-6)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(x2)
+    x1 = np.random.rand(5, 2, 1, 11)
+    x2 = common_layers.apply_norm(
+        tf.constant(x1, dtype=tf.float32), "noam", depth=11, epsilon=1e-6)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testApplyNormBatch(self):
-    with self.test_session() as session:
-      x1 = np.random.rand(5, 2, 1, 11)
-      x2 = common_layers.apply_norm(
-          tf.constant(x1, dtype=tf.float32), "batch", depth=11, epsilon=1e-6)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(x2)
+    x1 = np.random.rand(5, 2, 1, 11)
+    x2 = common_layers.apply_norm(
+        tf.constant(x1, dtype=tf.float32), "batch", depth=11, epsilon=1e-6)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testApplyNormNone(self):
-    with self.test_session() as session:
-      x1 = np.random.rand(5, 2, 1, 11)
-      x2 = common_layers.apply_norm(
-          tf.constant(x1, dtype=tf.float32), "none", depth=11, epsilon=1e-6)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(x2)
+    x1 = np.random.rand(5, 2, 1, 11)
+    x2 = common_layers.apply_norm(
+        tf.constant(x1, dtype=tf.float32), "none", depth=11, epsilon=1e-6)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
     self.assertAllClose(actual, x1, atol=1e-03)
 
@@ -356,89 +352,86 @@ def testGlobalPool1d(self):
     no_mask = np.ones((5, 4))
     full_mask = np.zeros((5, 4))
 
-    with self.test_session() as session:
-      x1_ = tf.Variable(x1, dtype=tf.float32)
-      no_mask_ = tf.Variable(no_mask, dtype=tf.float32)
-      full_mask_ = tf.Variable(full_mask, dtype=tf.float32)
+    x1_ = tf.Variable(x1, dtype=tf.float32)
+    no_mask_ = tf.Variable(no_mask, dtype=tf.float32)
+    full_mask_ = tf.Variable(full_mask, dtype=tf.float32)
 
-      none_mask_max = common_layers.global_pool_1d(x1_)
-      no_mask_max = common_layers.global_pool_1d(x1_, mask=no_mask_)
-      result1 = tf.reduce_sum(none_mask_max - no_mask_max)
+    none_mask_max = common_layers.global_pool_1d(x1_)
+    no_mask_max = common_layers.global_pool_1d(x1_, mask=no_mask_)
+    result1 = tf.reduce_sum(none_mask_max - no_mask_max)
 
-      full_mask_max = common_layers.global_pool_1d(x1_, mask=full_mask_)
-      result2 = tf.reduce_sum(full_mask_max)
+    full_mask_max = common_layers.global_pool_1d(x1_, mask=full_mask_)
+    result2 = tf.reduce_sum(full_mask_max)
 
-      none_mask_avr = common_layers.global_pool_1d(x1_, "AVR")
-      no_mask_avr = common_layers.global_pool_1d(x1_, "AVR", no_mask_)
-      result3 = tf.reduce_sum(none_mask_avr - no_mask_avr)
+    none_mask_avr = common_layers.global_pool_1d(x1_, "AVR")
+    no_mask_avr = common_layers.global_pool_1d(x1_, "AVR", no_mask_)
+    result3 = tf.reduce_sum(none_mask_avr - no_mask_avr)
 
-      full_mask_avr = common_layers.global_pool_1d(x1_, "AVR", full_mask_)
-      result4 = tf.reduce_sum(full_mask_avr)
+    full_mask_avr = common_layers.global_pool_1d(x1_, "AVR", full_mask_)
+    result4 = tf.reduce_sum(full_mask_avr)
 
-      session.run(tf.global_variables_initializer())
-      actual = session.run([result1, result2, result3, result4])
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate([result1, result2, result3, result4])
     self.assertAllEqual(actual[:3], [0.0, 0.0, 0.0])
 
   def testLinearSetLayer(self):
     x1 = np.random.rand(5, 4, 11)
     cont = np.random.rand(5, 13)
-    with self.test_session() as session:
-      x1_ = tf.Variable(x1, dtype=tf.float32)
-      cont_ = tf.Variable(cont, dtype=tf.float32)
+    x1_ = tf.Variable(x1, dtype=tf.float32)
+    cont_ = tf.Variable(cont, dtype=tf.float32)
 
-      simple_ff = common_layers.linear_set_layer(32, x1_)
-      cont_ff = common_layers.linear_set_layer(32, x1_, context=cont_)
+    simple_ff = common_layers.linear_set_layer(32, x1_)
+    cont_ff = common_layers.linear_set_layer(32, x1_, context=cont_)
 
-      session.run(tf.global_variables_initializer())
-      actual = session.run([simple_ff, cont_ff])
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate([simple_ff, cont_ff])
     self.assertEqual(actual[0].shape, (5, 4, 32))
     self.assertEqual(actual[1].shape, (5, 4, 32))
 
   def testRavanbakhshSetLayer(self):
     x1 = np.random.rand(5, 4, 11)
-    with self.test_session() as session:
-      x1_ = tf.Variable(x1, dtype=tf.float32)
-      layer = common_layers.ravanbakhsh_set_layer(32, x1_)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(layer)
+    x1_ = tf.Variable(x1, dtype=tf.float32)
+    layer = common_layers.ravanbakhsh_set_layer(32, x1_)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(layer)
     self.assertEqual(actual.shape, (5, 4, 32))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testBReLU(self):
-    with self.test_session() as session:
-      x = np.random.rand(5, 2, 1, 12)
-      y = common_layers.brelu(tf.constant(x, dtype=tf.float32))
-      actual = session.run(y)
+    x = np.random.rand(5, 2, 1, 12)
+    y = common_layers.brelu(tf.constant(x, dtype=tf.float32))
+    actual = self.evaluate(y)
     self.assertEqual(actual.shape, (5, 2, 1, 12))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testBELU(self):
-    with self.test_session() as session:
-      x = np.random.rand(5, 2, 1, 12)
-      y = common_layers.belu(tf.constant(x, dtype=tf.float32))
-      actual = session.run(y)
+    x = np.random.rand(5, 2, 1, 12)
+    y = common_layers.belu(tf.constant(x, dtype=tf.float32))
+    actual = self.evaluate(y)
     self.assertEqual(actual.shape, (5, 2, 1, 12))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testNAC(self):
-    with self.test_session() as session:
-      x = np.random.rand(5, 2, 1, 12)
-      y = common_layers.nac(tf.constant(x, dtype=tf.float32), 14)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(y)
+    x = np.random.rand(5, 2, 1, 12)
+    y = common_layers.nac(tf.constant(x, dtype=tf.float32), 14)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(y)
     self.assertEqual(actual.shape, (5, 2, 1, 14))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testNALU(self):
-    with self.test_session() as session:
-      x = np.random.rand(5, 2, 1, 12)
-      y = common_layers.nalu(tf.constant(x, dtype=tf.float32), 14)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(y)
+    x = np.random.rand(5, 2, 1, 12)
+    y = common_layers.nalu(tf.constant(x, dtype=tf.float32), 14)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(y)
     self.assertEqual(actual.shape, (5, 2, 1, 14))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testNALUzeros(self):
-    with self.test_session() as session:
-      x = np.random.rand(5, 2, 1, 12)
-      y = common_layers.nalu(tf.zeros_like(x, dtype=tf.float32), 14)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(y)
+    x = np.random.rand(5, 2, 1, 12)
+    y = common_layers.nalu(tf.zeros_like(x, dtype=tf.float32), 14)
+    self.evaluate(tf.global_variables_initializer())
+    actual = self.evaluate(y)
     self.assertTrue(np.all(np.isfinite(actual)))
     self.assertEqual(actual.shape, (5, 2, 1, 14))
 
@@ -533,11 +526,11 @@ def testDmlLoss(self, batch, height, width, num_mixtures, reduce_sum):
     if reduce_sum:
       expected_loss = tf.reduce_mean(expected_loss)
 
-    with self.test_session() as sess:
-      actual_loss_val, expected_loss_val = sess.run(
-          [actual_loss, expected_loss])
+    actual_loss_val, expected_loss_val = self.evaluate(
+        [actual_loss, expected_loss])
     self.assertAllClose(actual_loss_val, expected_loss_val)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDiscretizedMixLogisticLoss(self):
     batch = 2
     height = 4
@@ -570,11 +563,11 @@ def testDiscretizedMixLogisticLoss(self):
 
     actual_loss = common_layers.discretized_mix_logistic_loss(
         pred=pred, labels=labels)
-    with self.test_session() as session:
-      actual_loss_val, expected_loss_val = session.run(
-          [actual_loss, expected_loss])
+    actual_loss_val, expected_loss_val = self.evaluate(
+        [actual_loss, expected_loss])
     self.assertAllClose(actual_loss_val, expected_loss_val, rtol=1e-5)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testSampleFromDiscretizedMixLogistic(self):
     batch = 2
     height = 4
@@ -596,23 +589,22 @@ def testSampleFromDiscretizedMixLogistic(self):
 
     actual_sample = common_layers.sample_from_discretized_mix_logistic(
         pred, seed=seed)
-    with self.test_session() as session:
-      actual_sample_val, expected_sample_val = session.run(
-          [actual_sample, expected_sample])
+    actual_sample_val, expected_sample_val = self.evaluate(
+        [actual_sample, expected_sample])
     # Use a low tolerance: samples numerically differ, as the actual
     # implementation clips log-scales so they always contribute to sampling.
     self.assertAllClose(actual_sample_val, expected_sample_val, atol=1e-2)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testFactoredTensorImplicitConversion(self):
     a = np.random.rand(3, 4, 5)
     b = np.random.rand(6, 5)
     c = np.random.rand(3, 4, 6)
-    with self.test_session() as session:
-      # a factored representation of a Tensor of shape (3, 4, 6)
-      factored = common_layers.FactoredTensor(tf.to_float(a), tf.to_float(b))
-      # implicitly converts factored to a Tensor (performing the matmul)
-      d = factored + tf.to_float(c)
-      out = session.run(d)
+    # a factored representation of a Tensor of shape (3, 4, 6)
+    factored = common_layers.FactoredTensor(tf.to_float(a), tf.to_float(b))
+    # implicitly converts factored to a Tensor (performing the matmul)
+    d = factored + tf.to_float(c)
+    out = self.evaluate(d)
     self.assertEqual(out.shape, (3, 4, 6))
 
   def testConvHiddenReluMemoryEfficient(self):
@@ -652,6 +644,7 @@ def testConvHiddenReluMemoryEfficient(self):
     self.assertAllClose(dnorm_bias, dnorm_bias_f)
     self.assertAllClose(dx, dx_f)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testCycleGANUpsampleNnUpsampleConv(self):
     batch = 8
     height = 32
@@ -666,12 +659,12 @@ def testCycleGANUpsampleNnUpsampleConv(self):
     upsampled_output = common_layers.cyclegan_upsample(
         random_input, output_filters, stride, "nn_upsample_conv")
     upsampled_output_shape = tf.shape(upsampled_output)
-    with self.test_session() as session:
-      session.run(tf.global_variables_initializer())
-      self.assertAllEqual(
-          [batch, height * stride[0], width * stride[1], output_filters],
-          session.run(upsampled_output_shape))
+    self.evaluate(tf.global_variables_initializer())
+    self.assertAllEqual(
+        [batch, height * stride[0], width * stride[1], output_filters],
+        self.evaluate(upsampled_output_shape))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testCycleGANUpsampleBilinearUpsampleConv(self):
     batch = 8
     height = 32
@@ -686,12 +679,12 @@ def testCycleGANUpsampleBilinearUpsampleConv(self):
     upsampled_output = common_layers.cyclegan_upsample(
         random_input, output_filters, stride, "bilinear_upsample_conv")
     upsampled_output_shape = tf.shape(upsampled_output)
-    with self.test_session() as session:
-      session.run(tf.global_variables_initializer())
-      self.assertAllEqual(
-          [batch, height * stride[0], width * stride[1], output_filters],
-          session.run(upsampled_output_shape))
+    self.evaluate(tf.global_variables_initializer())
+    self.assertAllEqual(
+        [batch, height * stride[0], width * stride[1], output_filters],
+        self.evaluate(upsampled_output_shape))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testCycleGANUpsampleConv2dTranspose(self):
     batch = 8
     height = 32
@@ -711,11 +704,10 @@ def testCycleGANUpsampleConv2dTranspose(self):
                                                        output_filters, stride,
                                                        "conv2d_transpose")
     upsampled_output_shape = tf.shape(upsampled_output)
-    with self.test_session() as session:
-      session.run(tf.global_variables_initializer())
-      self.assertAllEqual(
-          [batch, upsampled_height, upsampled_width, output_filters],
-          session.run(upsampled_output_shape))
+    self.evaluate(tf.global_variables_initializer())
+    self.assertAllEqual(
+        [batch, upsampled_height, upsampled_width, output_filters],
+        self.evaluate(upsampled_output_shape))
 
   def testSpectralNorm(self):
     # Test that after 20 calls to apply_spectral_norm, the spectral
diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index efa031957..6f88c5fb6 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -24,77 +24,80 @@
 
 class CommonVideoTest(tf.test.TestCase):
 
-  @staticmethod
-  def RunScheduledSampleFunc(func, var, batch_size):
+  def runScheduledSampleFunc(self, func, var, batch_size):
     ground_truth_x = list(range(1, batch_size+1))
     generated_x = [-x for x in ground_truth_x]
     ground_truth_x = tf.convert_to_tensor(ground_truth_x)
     generated_x = tf.convert_to_tensor(generated_x)
     ss_out = func(ground_truth_x, generated_x, batch_size, var)
-    with tf.Session() as session:
-      output = session.run([ground_truth_x, generated_x, ss_out])
+    output = self.evaluate([ground_truth_x, generated_x, ss_out])
     return output
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testScheduledSampleProbStart(self):
-    ground_truth_x, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
+    ground_truth_x, _, ss_out = self.runScheduledSampleFunc(
         common_video.scheduled_sample_prob, 1.0, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testScheduledSampleProbMid(self):
-    _, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
+    _, _, ss_out = self.runScheduledSampleFunc(
         common_video.scheduled_sample_prob, 0.5, 1000)
     positive_count = np.sum(ss_out > 0)
-    self.assertAlmostEqual(positive_count / 1000.0, 0.5, places=2)
+    self.assertAlmostEqual(positive_count / 1000.0, 0.5, places=1)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testScheduledSampleProbEnd(self):
-    _, generated_x, ss_out = CommonVideoTest.RunScheduledSampleFunc(
+    _, generated_x, ss_out = self.runScheduledSampleFunc(
         common_video.scheduled_sample_prob, 0.0, 10)
     self.assertAllEqual(generated_x, ss_out)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testScheduledSampleCountStart(self):
-    ground_truth_x, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
+    ground_truth_x, _, ss_out = self.runScheduledSampleFunc(
         common_video.scheduled_sample_count, 10, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testScheduledSampleCountMid(self):
-    _, _, ss_out = CommonVideoTest.RunScheduledSampleFunc(
+    _, _, ss_out = self.runScheduledSampleFunc(
         common_video.scheduled_sample_count, 5, 10)
     positive_count = np.sum(ss_out > 0)
     self.assertEqual(positive_count, 5)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testScheduledSampleCountEnd(self):
-    _, generated_x, ss_out = CommonVideoTest.RunScheduledSampleFunc(
+    _, generated_x, ss_out = self.runScheduledSampleFunc(
         common_video.scheduled_sample_count, 0, 10)
     self.assertAllEqual(generated_x, ss_out)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDynamicTileAndConcat(self):
-    with tf.Graph().as_default():
-      # image = (1 X 4 X 4 X 1)
-      image = [[1, 2, 3, 4],
-               [2, 4, 5, 6],
-               [7, 8, 9, 10],
-               [7, 9, 10, 1]]
-      image_t = tf.expand_dims(tf.expand_dims(image, axis=0), axis=-1)
-      image_t = tf.cast(image_t, dtype=tf.float32)
-
-      # latent = (1 X 2)
-      latent = np.array([[90, 100]])
-      latent_t = tf.cast(tf.convert_to_tensor(latent), dtype=tf.float32)
-
-      with tf.Session() as session:
-        tiled = common_video.tile_and_concat(
-            image_t, latent_t)
-        tiled_np, image_np = session.run([tiled, image_t])
-        tiled_latent = tiled_np[0, :, :, -1]
-        self.assertAllEqual(tiled_np.shape, (1, 4, 4, 2))
-
-        self.assertAllEqual(tiled_np[:, :, :, :1], image_np)
-        self.assertAllEqual(
-            tiled_latent,
-            [[90, 90, 90, 90],
-             [100, 100, 100, 100],
-             [90, 90, 90, 90],
-             [100, 100, 100, 100]])
+    # image = (1 X 4 X 4 X 1)
+    image = [[1, 2, 3, 4],
+             [2, 4, 5, 6],
+             [7, 8, 9, 10],
+             [7, 9, 10, 1]]
+    image_t = tf.expand_dims(tf.expand_dims(image, axis=0), axis=-1)
+    image_t = tf.cast(image_t, dtype=tf.float32)
+
+    # latent = (1 X 2)
+    latent = np.array([[90, 100]])
+    latent_t = tf.cast(tf.convert_to_tensor(latent), dtype=tf.float32)
+
+    tiled = common_video.tile_and_concat(
+        image_t, latent_t)
+    tiled_np, image_np = self.evaluate([tiled, image_t])
+    tiled_latent = tiled_np[0, :, :, -1]
+    self.assertAllEqual(tiled_np.shape, (1, 4, 4, 2))
+
+    self.assertAllEqual(tiled_np[:, :, :, :1], image_np)
+    self.assertAllEqual(
+        tiled_latent,
+        [[90, 90, 90, 90],
+         [100, 100, 100, 100],
+         [90, 90, 90, 90],
+         [100, 100, 100, 100]])
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index c75f61704..d26a92042 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -30,42 +30,39 @@ def setUp(self):
     tf.set_random_seed(1234)
     np.random.seed(123)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testBitToIntZeros(self):
     x_bit = tf.zeros(shape=[1, 10], dtype=tf.float32)
     x_int = tf.zeros(shape=[1], dtype=tf.int32)
     diff = discretization.bit_to_int(x_bit, num_bits=10) - x_int
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      d = sess.run(diff)
-      self.assertEqual(d, 0)
+    d = self.evaluate(diff)
+    self.assertEqual(d, 0)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testBitToIntOnes(self):
     x_bit = tf.ones(shape=[1, 3], dtype=tf.float32)
     x_int = 7 * tf.ones(shape=[1], dtype=tf.int32)
     diff = discretization.bit_to_int(x_bit, num_bits=3) - x_int
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      d = sess.run(diff)
-      self.assertEqual(d, 0)
+    d = self.evaluate(diff)
+    self.assertEqual(d, 0)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testIntToBitZeros(self):
     x_bit = tf.zeros(shape=[1, 10], dtype=tf.float32)
     x_int = tf.zeros(shape=[1], dtype=tf.int32)
     diff = discretization.int_to_bit(x_int, num_bits=10) - x_bit
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      d = sess.run(diff)
-      self.assertTrue(np.all(d == 0))
+    d = self.evaluate(diff)
+    self.assertTrue(np.all(d == 0))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testIntToBitOnes(self):
     x_bit = tf.ones(shape=[1, 3], dtype=tf.float32)
     x_int = 7 * tf.ones(shape=[1], dtype=tf.int32)
     diff = discretization.int_to_bit(x_int, num_bits=3) - x_bit
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      d = sess.run(diff)
-      self.assertTrue(np.all(d == 0))
+    d = self.evaluate(diff)
+    self.assertTrue(np.all(d == 0))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testProjectHidden(self):
     hidden_size = 60
     block_dim = 20
@@ -75,36 +72,33 @@ def testProjectHidden(self):
         shape=[num_blocks, hidden_size, block_dim], dtype=tf.float32)
     x_projected = discretization.project_hidden(x, projection_tensors,
                                                 hidden_size, num_blocks)
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_projected_eval = sess.run(x_projected)
-      self.assertEqual(np.shape(x_projected_eval), (1, num_blocks, block_dim))
-      self.assertTrue(np.all(x_projected_eval == 0))
+    x_projected_eval = self.evaluate(x_projected)
+    self.assertEqual(np.shape(x_projected_eval), (1, num_blocks, block_dim))
+    self.assertTrue(np.all(x_projected_eval == 0))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testSliceHiddenZeros(self):
     hidden_size = 60
     block_dim = 20
     num_blocks = 3
     x = tf.zeros(shape=[1, hidden_size], dtype=tf.float32)
     x_sliced = discretization.slice_hidden(x, hidden_size, num_blocks)
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_sliced_eval = sess.run(x_sliced)
-      self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
-      self.assertTrue(np.all(x_sliced_eval == 0))
+    x_sliced_eval = self.evaluate(x_sliced)
+    self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
+    self.assertTrue(np.all(x_sliced_eval == 0))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testSliceHiddenOnes(self):
     hidden_size = 60
     block_dim = 20
     num_blocks = 3
     x = tf.ones(shape=[1, hidden_size], dtype=tf.float32)
     x_sliced = discretization.slice_hidden(x, hidden_size, num_blocks)
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_sliced_eval = sess.run(x_sliced)
-      self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
-      self.assertTrue(np.all(x_sliced_eval == 1))
+    x_sliced_eval = self.evaluate(x_sliced)
+    self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
+    self.assertTrue(np.all(x_sliced_eval == 1))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testNearestNeighbors(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
     x = tf.reshape(x, [1, 1, 2, 3])
@@ -115,11 +109,9 @@ def testNearestNeighbors(self):
         x, means, block_v_size=4)
     x_means_hot_test = np.array([[0, 1, 0, 0], [1, 0, 0, 0]])
     x_means_hot_test = np.expand_dims(x_means_hot_test, axis=0)
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_means_hot_eval = sess.run(x_means_hot)
-      self.assertEqual(np.shape(x_means_hot_eval), (1, 2, 4))
-      self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
+    x_means_hot_eval = self.evaluate(x_means_hot)
+    self.assertEqual(np.shape(x_means_hot_eval), (1, 2, 4))
+    self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
 
   def testGetVQBottleneck(self):
     bottleneck_bits = 2
@@ -136,43 +128,39 @@ def testGetVQBottleneck(self):
       self.assertTrue(np.all(sess.run(means_new) == 0))
       self.assertTrue(np.all(sess.run(ema_count) == 0))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testVQNearestNeighbors(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
     means = tf.constant(
         [[1, 0, 0], [0, 1, 0], [0, 0, 1], [9, 9, 9]], dtype=tf.float32)
     x_means_hot, _, _ = discretization.vq_nearest_neighbor(x, means)
     x_means_hot_test = np.array([[0, 1, 0, 0], [1, 0, 0, 0]])
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_means_hot_eval = sess.run(x_means_hot)
-      self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
-      self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
+    x_means_hot_eval = self.evaluate(x_means_hot)
+    self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
+    self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
 
   def testVQDiscreteBottleneck(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
     x_means_hot, _ = discretization.vq_discrete_bottleneck(x, bottleneck_bits=2)
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_means_hot_eval = sess.run(x_means_hot)
-      self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
+    self.evaluate(tf.global_variables_initializer())
+    x_means_hot_eval = self.evaluate(x_means_hot)
+    self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
 
   def testVQDiscreteUnbottlenck(self):
     x = tf.constant([[1, 0, 0, 0], [0, 0, 1, 0]], dtype=tf.int32)
     x_means = discretization.vq_discrete_unbottleneck(x, hidden_size=3)
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_means_eval = sess.run(x_means)
-      self.assertEqual(np.shape(x_means_eval), (2, 3))
+    self.evaluate(tf.global_variables_initializer())
+    x_means_eval = self.evaluate(x_means)
+    self.assertEqual(np.shape(x_means_eval), (2, 3))
 
   def testGumbelSoftmaxDiscreteBottleneck(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
     tf.add_to_collection(tf.GraphKeys.GLOBAL_STEP, tf.constant(1))
     x_means_hot, _ = discretization.gumbel_softmax_discrete_bottleneck(
         x, bottleneck_bits=2)
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      x_means_hot_eval = sess.run(x_means_hot)
-      self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
+    self.evaluate(tf.global_variables_initializer())
+    x_means_hot_eval = self.evaluate(x_means_hot)
+    self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
 
 
 if __name__ == '__main__':
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 04fba3bff..e0f631e30 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -27,6 +27,7 @@
 
 class ModalityTest(tf.test.TestCase):
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testSymbolModalityInputs(self):
     batch_size = 10
     num_datashards = 5
@@ -41,14 +42,14 @@ def testSymbolModalityInputs(self):
     m = modalities.SymbolModality(model_hparams, vocab_size)
     data_parallelism = expert_utils.Parallelism(
         ["/device:CPU:0"] * num_datashards)
-    with self.test_session() as session:
-      xs = tf.split(x, num_datashards)
-      sharded_output = m.bottom_sharded(xs, data_parallelism)
-      output = tf.concat(sharded_output, 0)
-      session.run(tf.global_variables_initializer())
-      res = session.run(output)
+    xs = tf.split(x, num_datashards)
+    sharded_output = m.bottom_sharded(xs, data_parallelism)
+    output = tf.concat(sharded_output, 0)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(output)
     self.assertEqual(res.shape, (batch_size, length, 1, hidden_size))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testSymbolModalityTargets(self):
     batch_size = 10
     num_datashards = 5
@@ -66,16 +67,15 @@ def testSymbolModalityTargets(self):
     m = modalities.SymbolModality(model_hparams, vocab_size)
     data_parallelism = expert_utils.Parallelism(
         ["/device:CPU:0"] * num_datashards)
-    with self.test_session() as session:
-      sharded_body_output = tf.split(tf.to_float(body_output), num_datashards)
-      sharded_targets = tf.split(targets, num_datashards)
-      sharded_logits = m.top_sharded(sharded_body_output, sharded_targets,
-                                     data_parallelism)
-      train_loss = m.loss_sharded(sharded_logits, sharded_targets,
-                                  data_parallelism)
-      logits = tf.concat(sharded_logits, 0)
-      session.run(tf.global_variables_initializer())
-      res1, res2 = session.run((logits, train_loss))
+    sharded_body_output = tf.split(tf.to_float(body_output), num_datashards)
+    sharded_targets = tf.split(targets, num_datashards)
+    sharded_logits = m.top_sharded(sharded_body_output, sharded_targets,
+                                   data_parallelism)
+    train_loss = m.loss_sharded(sharded_logits, sharded_targets,
+                                data_parallelism)
+    logits = tf.concat(sharded_logits, 0)
+    self.evaluate(tf.global_variables_initializer())
+    res1, res2 = self.evaluate((logits, train_loss))
     self.assertEqual(res1.shape, (batch_size, length, height, 1, vocab_size))
     self.assertEqual(res2.shape, ())
 
diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
index efec84659..10e36fe2a 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
@@ -117,7 +117,7 @@ def testGraph(self):
     self.assertLen(graph.trainable_variables, 1)
     self.assertLen(graph.all_variables, 2)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testLowering(self):
     graph = mtf.Graph()
     mesh = mtf.Mesh(graph, "my_mesh")
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers_test.py b/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
index 4d6cba965..f3c3501e0 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
@@ -28,7 +28,6 @@
 import tensorflow as tf
 
 
-@tf.contrib.eager.run_all_tests_in_graph_and_eager_modes
 class MtfLayersTest(parameterized.TestCase, tf.test.TestCase):
 
   @parameterized.parameters(
@@ -69,6 +68,7 @@ def testDense(self, units, use_bias):
 
     self.assertEqual(actual.shape, expected.shape)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testLayerNorm(self):
     batch = 2
     channels = 3
@@ -97,6 +97,7 @@ def testLayerNorm(self):
 
     self.assertEqual(actual.shape, expected.shape)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testWeightsNonzero(self):
     inputs = tf.constant([[3, 1, 0], [1, 0, 0]])
 
@@ -120,6 +121,7 @@ def testWeightsNonzero(self):
 
     self.assertAllEqual(actual, expected)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDenseReluDense(self):
     batch = 2
     channels = 3

From 321bacaa3abcca5dbf341ed6fb3d4a1531e513ff Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 27 Aug 2018 16:26:01 -0700
Subject: [PATCH 0681/2720] Fix translate data generation; use standard vocab
 naming

PiperOrigin-RevId: 210452200
---
 tensor2tensor/bin/t2t_datagen.py              |  2 +
 .../data_generators/cnn_dailymail.py          |  4 --
 tensor2tensor/data_generators/cola.py         |  4 --
 tensor2tensor/data_generators/image_utils.py  |  4 +-
 tensor2tensor/data_generators/imdb.py         |  4 --
 tensor2tensor/data_generators/lm1b.py         |  8 ++--
 tensor2tensor/data_generators/mrpc.py         |  4 --
 tensor2tensor/data_generators/mscoco.py       |  7 +--
 tensor2tensor/data_generators/multinli.py     |  7 +--
 tensor2tensor/data_generators/qnli.py         |  4 --
 tensor2tensor/data_generators/quora_qpairs.py |  4 --
 tensor2tensor/data_generators/rte.py          |  4 --
 tensor2tensor/data_generators/sst_binary.py   |  4 --
 tensor2tensor/data_generators/stanford_nli.py |  7 +--
 .../data_generators/translate_encs.py         |  8 +---
 .../data_generators/translate_ende.py         | 48 +++++++++----------
 .../data_generators/translate_enet.py         |  4 --
 .../data_generators/translate_enfr.py         | 13 ++---
 .../data_generators/translate_enid.py         |  4 --
 .../data_generators/translate_enmk.py         |  4 --
 .../data_generators/translate_envi.py         |  4 --
 .../data_generators/translate_enzh.py         |  8 ++--
 tensor2tensor/data_generators/wiki.py         |  8 ----
 tensor2tensor/notebooks/hello_t2t.ipynb       |  4 +-
 ...ocab.translate_ende_wmt32k.32768.subwords} |  0
 ... vocab.translate_ende_wmt8k.8192.subwords} |  0
 26 files changed, 51 insertions(+), 121 deletions(-)
 rename tensor2tensor/test_data/{vocab.ende.32768 => vocab.translate_ende_wmt32k.32768.subwords} (100%)
 rename tensor2tensor/test_data/{vocab.ende.8192 => vocab.translate_ende_wmt8k.8192.subwords} (100%)

diff --git a/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index 7375b3141..d8fa6ce1a 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -147,6 +147,8 @@ def main(_):
       problems = [p for p in problems if exclude not in p]
   if FLAGS.problem and FLAGS.problem[-1] == "*":
     problems = [p for p in problems if p.startswith(FLAGS.problem[:-1])]
+  elif FLAGS.problem and "," in FLAGS.problem:
+    problems = [p for p in problems if p in FLAGS.problem.split(",")]
   elif FLAGS.problem:
     problems = [p for p in problems if p == FLAGS.problem]
   else:
diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index ce147ca75..00464db0f 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -203,10 +203,6 @@ def write_to_file(all_files, urls_path, tmp_dir, filename):
 class SummarizeCnnDailymail32k(text_problems.Text2TextProblem):
   """Summarize CNN and Daily Mail articles to their summary highlights."""
 
-  @property
-  def vocab_filename(self):
-    return "vocab.cnndailymail.%d" % self.approx_vocab_size
-
   def generate_text_for_vocab(self, data_dir, tmp_dir):
     del data_dir
     all_files, urls_path = _maybe_download_corpora(tmp_dir, True)
diff --git a/tensor2tensor/data_generators/cola.py b/tensor2tensor/data_generators/cola.py
index 0028c4f26..4e5f4592f 100644
--- a/tensor2tensor/data_generators/cola.py
+++ b/tensor2tensor/data_generators/cola.py
@@ -59,10 +59,6 @@ def dataset_splits(self):
   def approx_vocab_size(self):
     return 2**13  # 8k vocab suffices for this small dataset.
 
-  @property
-  def vocab_filename(self):
-    return "vocab.cola.%d" % self.approx_vocab_size
-
   @property
   def num_classes(self):
     return 2
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 18dd3f0b4..3b59060e5 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -313,7 +313,7 @@ def is_character_level(self):
     raise NotImplementedError()
 
   @property
-  def targeted_vocab_size(self):
+  def vocab_problem(self):
     raise NotImplementedError()  # Not needed if self.is_character_level.
 
   @property
@@ -345,7 +345,7 @@ def feature_encoders(self, data_dir):
       encoder = text_encoder.ByteTextEncoder()
     else:
       vocab_filename = os.path.join(
-          data_dir, "vocab.ende.%d" % self.targeted_vocab_size)
+          data_dir, self.vocab_problem.vocab_filename)
       encoder = text_encoder.SubwordTextEncoder(vocab_filename)
     input_encoder = text_encoder.ImageEncoder(channels=self.num_channels)
     return {"inputs": input_encoder, "targets": encoder}
diff --git a/tensor2tensor/data_generators/imdb.py b/tensor2tensor/data_generators/imdb.py
index cbb37d3df..2f02a471b 100644
--- a/tensor2tensor/data_generators/imdb.py
+++ b/tensor2tensor/data_generators/imdb.py
@@ -47,10 +47,6 @@ def dataset_splits(self):
         "shards": 1,
     }]
 
-  @property
-  def vocab_filename(self):
-    return "sentiment_imdb.vocab.%d" % self.approx_vocab_size
-
   @property
   def approx_vocab_size(self):
     return 2**13  # 8k vocab suffices for this small dataset.
diff --git a/tensor2tensor/data_generators/lm1b.py b/tensor2tensor/data_generators/lm1b.py
index dca331c6c..4d9e5e5a6 100644
--- a/tensor2tensor/data_generators/lm1b.py
+++ b/tensor2tensor/data_generators/lm1b.py
@@ -110,10 +110,6 @@ class LanguagemodelLm1b32k(text_problems.Text2SelfProblem):
   176884 / 159658 = 1.107893; multiply log_ppl by this to compare results.
   """
 
-  @property
-  def vocab_filename(self):
-    return "vocab.lm1b.en.%d" % self.approx_vocab_size
-
   @property
   def approx_vocab_size(self):
     return 2**15  # 32768
@@ -149,6 +145,10 @@ class LanguagemodelLm1b32kPacked(LanguagemodelLm1b32k):
   def packed_length(self):
     return 256
 
+  @property
+  def vocab_filename(self):
+    return LanguagemodelLm1b32k().vocab_filename
+
 
 @registry.register_problem
 class LanguagemodelLm1b8kPacked(LanguagemodelLm1b32kPacked):
diff --git a/tensor2tensor/data_generators/mrpc.py b/tensor2tensor/data_generators/mrpc.py
index f6e2b504b..c56bac32c 100644
--- a/tensor2tensor/data_generators/mrpc.py
+++ b/tensor2tensor/data_generators/mrpc.py
@@ -63,10 +63,6 @@ def dataset_splits(self):
   def approx_vocab_size(self):
     return 2**13  # 8k vocab suffices for this small dataset.
 
-  @property
-  def vocab_filename(self):
-    return "vocab.mrpc.%d" % self.approx_vocab_size
-
   @property
   def num_classes(self):
     return 2
diff --git a/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index 28a1ea935..b103aa135 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -28,6 +28,7 @@
 from tensor2tensor.data_generators import imagenet
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import translate_ende
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -180,8 +181,8 @@ def is_character_level(self):
     return False
 
   @property
-  def targeted_vocab_size(self):
-    return 2**15  # 32768
+  def vocab_problem(self):
+    return translate_ende.TranslateEndeWmt32k()
 
   @property
   def target_space_id(self):
@@ -199,7 +200,7 @@ def generator(self, data_dir, tmp_dir, is_training):
     # We use the translate vocab file as the vocabulary for captions.
     # This requires having the vocab file present in the data_dir for the
     # generation pipeline to succeed.
-    vocab_filename = "vocab.ende.%d" % self.targeted_vocab_size
+    vocab_filename = self.vocab_problem.vocab_filename
     if is_training:
       return mscoco_generator(
           data_dir,
diff --git a/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index be879be2f..2b4dddfcb 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -22,6 +22,7 @@
 import zipfile
 import six
 from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import lm1b
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
@@ -59,10 +60,6 @@ def dataset_splits(self):
   def approx_vocab_size(self):
     return 2**15
 
-  @property
-  def vocab_filename(self):
-    return "vocab.mnli.%d" % self.approx_vocab_size
-
   @property
   def num_classes(self):
     return 3
@@ -133,4 +130,4 @@ class MultiNLISharedVocab(MultiNLI):
 
   @property
   def vocab_filename(self):
-    return "vocab.lm1b.en.%d" % 2**15
+    return lm1b.LanguagemodelLm1b32k().vocab_filename
diff --git a/tensor2tensor/data_generators/qnli.py b/tensor2tensor/data_generators/qnli.py
index a075e8fb3..310aadcd5 100644
--- a/tensor2tensor/data_generators/qnli.py
+++ b/tensor2tensor/data_generators/qnli.py
@@ -59,10 +59,6 @@ def dataset_splits(self):
   def approx_vocab_size(self):
     return 2**15
 
-  @property
-  def vocab_filename(self):
-    return "vocab.qnli.%d" % self.approx_vocab_size
-
   @property
   def num_classes(self):
     return 2
diff --git a/tensor2tensor/data_generators/quora_qpairs.py b/tensor2tensor/data_generators/quora_qpairs.py
index 89736d7f3..7e189bfc0 100644
--- a/tensor2tensor/data_generators/quora_qpairs.py
+++ b/tensor2tensor/data_generators/quora_qpairs.py
@@ -59,10 +59,6 @@ def dataset_splits(self):
   def approx_vocab_size(self):
     return 2**15
 
-  @property
-  def vocab_filename(self):
-    return "vocab.qqp.%d" % self.approx_vocab_size
-
   @property
   def num_classes(self):
     return 2
diff --git a/tensor2tensor/data_generators/rte.py b/tensor2tensor/data_generators/rte.py
index 0ea78e144..04e322ca1 100644
--- a/tensor2tensor/data_generators/rte.py
+++ b/tensor2tensor/data_generators/rte.py
@@ -59,10 +59,6 @@ def dataset_splits(self):
   def approx_vocab_size(self):
     return 2**13  # 8k vocab suffices for this small dataset.
 
-  @property
-  def vocab_filename(self):
-    return "vocab.rte.%d" % self.approx_vocab_size
-
   @property
   def num_classes(self):
     return 2
diff --git a/tensor2tensor/data_generators/sst_binary.py b/tensor2tensor/data_generators/sst_binary.py
index 17b1a3ef5..1fceacdee 100644
--- a/tensor2tensor/data_generators/sst_binary.py
+++ b/tensor2tensor/data_generators/sst_binary.py
@@ -59,10 +59,6 @@ def dataset_splits(self):
   def approx_vocab_size(self):
     return 2**14
 
-  @property
-  def vocab_filename(self):
-    return "vocab.sst_binary.%d" % self.approx_vocab_size
-
   @property
   def num_classes(self):
     return 2
diff --git a/tensor2tensor/data_generators/stanford_nli.py b/tensor2tensor/data_generators/stanford_nli.py
index 16efea71e..ad244303b 100644
--- a/tensor2tensor/data_generators/stanford_nli.py
+++ b/tensor2tensor/data_generators/stanford_nli.py
@@ -22,6 +22,7 @@
 import zipfile
 import six
 from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import lm1b
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
@@ -56,10 +57,6 @@ def dataset_splits(self):
   def approx_vocab_size(self):
     return 2**15
 
-  @property
-  def vocab_filename(self):
-    return "vocab.snli.%d" % self.approx_vocab_size
-
   @property
   def num_classes(self):
     return 3
@@ -131,4 +128,4 @@ class StanfordNLISharedVocab(StanfordNLI):
 
   @property
   def vocab_filename(self):
-    return "vocab.lm1b.en.%d" % 2**15
+    return lm1b.LanguagemodelLm1b32k().vocab_filename
diff --git a/tensor2tensor/data_generators/translate_encs.py b/tensor2tensor/data_generators/translate_encs.py
index be7341d53..9224e0d3d 100644
--- a/tensor2tensor/data_generators/translate_encs.py
+++ b/tensor2tensor/data_generators/translate_encs.py
@@ -36,8 +36,8 @@
      ("tsv", 3, 2, "data.plaintext-format/*train.gz")],
     [
         "http://data.statmt.org/wmt18/translation-task/training-parallel-nc-v13.tgz",  # pylint: disable=line-too-long
-        ("training/news-commentary-v13.cs-en.en",
-         "training/news-commentary-v13.cs-en.cs")
+        ("training-parallel-nc-v13/news-commentary-v13.cs-en.en",
+         "training-parallel-nc-v13/news-commentary-v13.cs-en.cs")
     ],
     [
         "http://www.statmt.org/wmt13/training-parallel-commoncrawl.tgz",
@@ -64,10 +64,6 @@ class TranslateEncsWmt32k(translate.TranslateProblem):
   def approx_vocab_size(self):
     return 2**15  # 32768
 
-  @property
-  def vocab_filename(self):
-    return "vocab.encs.%d" % self.approx_vocab_size
-
   def source_data_files(self, dataset_split):
     train = dataset_split == problem.DatasetSplit.TRAIN
     return _ENCS_TRAIN_DATASETS if train else _ENCS_TEST_DATASETS
diff --git a/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index 095d5fcff..73b1c9ddc 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -32,8 +32,8 @@
 _ENDE_TRAIN_DATASETS = [
     [
         "http://data.statmt.org/wmt18/translation-task/training-parallel-nc-v13.tgz",  # pylint: disable=line-too-long
-        ("training/news-commentary-v13.de-en.en",
-         "training/news-commentary-v13.de-en.de")
+        ("training-parallel-nc-v13/news-commentary-v13.de-en.en",
+         "training-parallel-nc-v13/news-commentary-v13.de-en.de")
     ],
     [
         "http://www.statmt.org/wmt13/training-parallel-commoncrawl.tgz",
@@ -71,18 +71,12 @@ class TranslateEndeWmtBpe32k(translate.TranslateProblem):
   """Problem spec for WMT En-De translation, BPE version."""
 
   @property
-  def approx_vocab_size(self):
-    return 32000
+  def vocab_type(self):
+    return text_problems.VocabType.TOKEN
 
   @property
-  def vocab_filename(self):
-    return "vocab.bpe.%d" % self.approx_vocab_size
-
-  def get_or_create_vocab(self, data_dir, tmp_dir, force_get=False):
-    vocab_filename = os.path.join(data_dir, self.vocab_filename)
-    if not tf.gfile.Exists(vocab_filename) and force_get:
-      raise ValueError("Vocab %s not found" % vocab_filename)
-    return text_encoder.TokenTextEncoder(vocab_filename, replace_oov="UNK")
+  def oov_token(self):
+    return "UNK"
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
     """Instance of token generator for the WMT en->de task, training set."""
@@ -92,14 +86,14 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
     train_path = _get_wmt_ende_bpe_dataset(tmp_dir, dataset_path)
 
     # Vocab
-    token_path = os.path.join(data_dir, self.vocab_filename)
-    if not tf.gfile.Exists(token_path):
-      token_tmp_path = os.path.join(tmp_dir, self.vocab_filename)
-      tf.gfile.Copy(token_tmp_path, token_path)
-      with tf.gfile.GFile(token_path, mode="r") as f:
-        vocab_data = "<pad>\n<EOS>\n" + f.read() + "UNK\n"
-      with tf.gfile.GFile(token_path, mode="w") as f:
-        f.write(vocab_data)
+    vocab_path = os.path.join(data_dir, self.vocab_filename)
+    if not tf.gfile.Exists(vocab_path):
+      bpe_vocab = os.path.join(tmp_dir, "vocab.bpe.32000")
+      with tf.gfile.Open(bpe_vocab) as f:
+        vocab_list = f.read().split("\n")
+      vocab_list.append(self.oov_token)
+      text_encoder.TokenTextEncoder(
+          None, vocab_list=vocab_list).store_to_file(vocab_path)
 
     return text_problems.text2text_txt_iterator(train_path + ".en",
                                                 train_path + ".de")
@@ -113,10 +107,6 @@ class TranslateEndeWmt8k(translate.TranslateProblem):
   def approx_vocab_size(self):
     return 2**13  # 8192
 
-  @property
-  def vocab_filename(self):
-    return "vocab.ende.%d" % self.approx_vocab_size
-
   def source_data_files(self, dataset_split):
     train = dataset_split == problem.DatasetSplit.TRAIN
     return _ENDE_TRAIN_DATASETS if train else _ENDE_TEST_DATASETS
@@ -137,6 +127,10 @@ class TranslateEndeWmt32kPacked(TranslateEndeWmt32k):
   def packed_length(self):
     return 256
 
+  @property
+  def vocab_filename(self):
+    return TranslateEndeWmt32k().vocab_filename
+
 
 @registry.register_problem
 class TranslateEndeWmt8kPacked(TranslateEndeWmt8k):
@@ -145,9 +139,13 @@ class TranslateEndeWmt8kPacked(TranslateEndeWmt8k):
   def packed_length(self):
     return 256
 
+  @property
+  def vocab_filename(self):
+    return TranslateEndeWmt8k().vocab_filename
+
 
 @registry.register_problem
-class TranslateEndeWmtCharacters(translate.TranslateProblem):
+class TranslateEndeWmtCharacters(TranslateEndeWmt8k):
   """Problem spec for WMT En-De translation."""
 
   @property
diff --git a/tensor2tensor/data_generators/translate_enet.py b/tensor2tensor/data_generators/translate_enet.py
index ec98db06e..33150a8e4 100644
--- a/tensor2tensor/data_generators/translate_enet.py
+++ b/tensor2tensor/data_generators/translate_enet.py
@@ -59,10 +59,6 @@ class TranslateEnetWmt32k(translate.TranslateProblem):
   def approx_vocab_size(self):
     return 2**15  # 32768
 
-  @property
-  def vocab_filename(self):
-    return "vocab.enet.%d" % self.approx_vocab_size
-
   def source_data_files(self, dataset_split):
     train = dataset_split == problem.DatasetSplit.TRAIN
     return _ENET_TRAIN_DATASETS if train else _ENET_TEST_DATASETS
diff --git a/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index e190f9982..028ea0ac5 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -84,11 +84,6 @@ class TranslateEnfrWmtSmall8k(translate.TranslateProblem):
   def approx_vocab_size(self):
     return 2**13  # 8192
 
-  @property
-  def vocab_filename(self):
-    return "vocab.enfr.%s.%d" % (
-        "small" if self.use_small_dataset else "large", self.approx_vocab_size)
-
   @property
   def use_small_dataset(self):
     return True
@@ -137,6 +132,10 @@ class TranslateEnfrWmt32kPacked(TranslateEnfrWmt32k):
   def packed_length(self):
     return 256
 
+  @property
+  def vocab_filename(self):
+    return TranslateEnfrWmt32k().vocab_filename
+
 
 @registry.register_problem
 class TranslateEnfrWmtSmallCharacters(translate.TranslateProblem):
@@ -150,10 +149,6 @@ def vocab_type(self):
   def use_small_dataset(self):
     return True
 
-  @property
-  def vocab_filename(self):
-    return "vocab.enfr.%d" % self.approx_vocab_size
-
   def source_data_files(self, dataset_split):
     train = dataset_split == problem.DatasetSplit.TRAIN
     if self.use_small_dataset:
diff --git a/tensor2tensor/data_generators/translate_enid.py b/tensor2tensor/data_generators/translate_enid.py
index df49d31ec..0331b664b 100644
--- a/tensor2tensor/data_generators/translate_enid.py
+++ b/tensor2tensor/data_generators/translate_enid.py
@@ -79,10 +79,6 @@ class TranslateEnidIwslt32k(translate.TranslateProblem):
   def approx_vocab_size(self):
     return 2**15  # 32768
 
-  @property
-  def vocab_filename(self):
-    return "vocab.enid.%d" % self.approx_vocab_size
-
   def source_data_files(self, dataset_split):
     train = dataset_split == problem.DatasetSplit.TRAIN
     return _ENID_TRAIN_DATASETS if train else _ENID_TEST_DATASETS
diff --git a/tensor2tensor/data_generators/translate_enmk.py b/tensor2tensor/data_generators/translate_enmk.py
index 67dfefd99..fc832c3e4 100644
--- a/tensor2tensor/data_generators/translate_enmk.py
+++ b/tensor2tensor/data_generators/translate_enmk.py
@@ -54,10 +54,6 @@ class TranslateEnmkSetimes32k(translate.TranslateProblem):
   def approx_vocab_size(self):
     return 2**15  # 32768
 
-  @property
-  def vocab_filename(self):
-    return "vocab.enmk.%d" % self.approx_vocab_size
-
   def source_data_files(self, dataset_split):
     train = dataset_split == problem.DatasetSplit.TRAIN
     return _ENMK_TRAIN_DATASETS if train else _ENMK_DEV_DATASETS
diff --git a/tensor2tensor/data_generators/translate_envi.py b/tensor2tensor/data_generators/translate_envi.py
index 3d1709c79..97023cbd2 100644
--- a/tensor2tensor/data_generators/translate_envi.py
+++ b/tensor2tensor/data_generators/translate_envi.py
@@ -52,10 +52,6 @@ class TranslateEnviIwslt32k(translate.TranslateProblem):
   def approx_vocab_size(self):
     return 2**15  # 32768
 
-  @property
-  def vocab_filename(self):
-    return "vocab.envi.%d" % self.approx_vocab_size
-
   def source_data_files(self, dataset_split):
     train = dataset_split == problem.DatasetSplit.TRAIN
     return _ENVI_TRAIN_DATASETS if train else _ENVI_TEST_DATASETS
diff --git a/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index 500614b9f..db030120d 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -42,8 +42,8 @@
 _STAT_MT_URL = "http://data.statmt.org/wmt18/translation-task/"
 _NC_TRAIN_DATASETS = [[
     _STAT_MT_URL + "training-parallel-nc-v13.tgz", [
-        "training/news-commentary-v13.zh-en.en",
-        "training/news-commentary-v13.zh-en.zh"
+        "training-parallel-nc-v13/news-commentary-v13.zh-en.en",
+        "training-parallel-nc-v13/news-commentary-v13.zh-en.zh"
     ]
 ]]
 
@@ -182,11 +182,11 @@ def approx_vocab_size(self):
 
   @property
   def source_vocab_name(self):
-    return "vocab.enzh-en.%d" % self.approx_vocab_size
+    return "%s.en" % self.vocab_filename
 
   @property
   def target_vocab_name(self):
-    return "vocab.enzh-zh.%d" % self.approx_vocab_size
+    return "%s.zh" % self.vocab_filename
 
   def get_training_dataset(self, tmp_dir):
     """UN Parallel Corpus and CWMT Corpus need to be downloaded manually.
diff --git a/tensor2tensor/data_generators/wiki.py b/tensor2tensor/data_generators/wiki.py
index 45e3387e8..3b76457f2 100644
--- a/tensor2tensor/data_generators/wiki.py
+++ b/tensor2tensor/data_generators/wiki.py
@@ -85,10 +85,6 @@ def corpus_url(self):
     return ("https://archive.org/download/enwiki-20171201/"
             "enwiki-20171201-pages-articles.xml.bz2")
 
-  @property
-  def vocab_filename(self):
-    return "vocab.wiki_xml.%d" % self.approx_vocab_size
-
   @property
   def approx_vocab_size(self):
     return 2**13  # 8192
@@ -219,10 +215,6 @@ class LanguagemodelWikiNorefV8kL1k(LanguagemodelWikiXmlV8kL1k):
   without regard to article boundaries.
   """
 
-  @property
-  def vocab_filename(self):
-    return "vocab.wiki_noref.%d" % self.approx_vocab_size
-
   def filepath_to_unicode_strings(self, filepath):
     """Overrides the base class to clean up the xml dump before tokenizing."""
     dump = text_encoder.to_unicode_ignore_errors(tf.gfile.Open(filepath).read())
diff --git a/tensor2tensor/notebooks/hello_t2t.ipynb b/tensor2tensor/notebooks/hello_t2t.ipynb
index e28344d20..5e163d876 100644
--- a/tensor2tensor/notebooks/hello_t2t.ipynb
+++ b/tensor2tensor/notebooks/hello_t2t.ipynb
@@ -423,7 +423,7 @@
         "\n",
         "# Copy the vocab file locally so we can encode inputs and decode model outputs\n",
         "# All vocabs are stored on GCS\n",
-        "vocab_name = \"vocab.ende.32768\"\n",
+        "vocab_name = \"vocab.translate_ende_wmt32k.32768.subwords\"\n",
         "vocab_file = os.path.join(gs_data_dir, vocab_name)\n",
         "!gsutil cp {vocab_file} {data_dir}\n",
         "\n",
@@ -455,7 +455,7 @@
             "please run:\r\n",
             "  $ gcloud components update\r\n",
             "\n",
-            "Copying gs://tensor2tensor-data/vocab.ende.32768...\n",
+            "Copying gs://tensor2tensor-data/vocab.translate_ende_wmt32k.32768.subwords...\n",
             "/ [1 files][316.4 KiB/316.4 KiB]                                                \n",
             "Operation completed over 1 objects/316.4 KiB.                                    \n"
           ],
diff --git a/tensor2tensor/test_data/vocab.ende.32768 b/tensor2tensor/test_data/vocab.translate_ende_wmt32k.32768.subwords
similarity index 100%
rename from tensor2tensor/test_data/vocab.ende.32768
rename to tensor2tensor/test_data/vocab.translate_ende_wmt32k.32768.subwords
diff --git a/tensor2tensor/test_data/vocab.ende.8192 b/tensor2tensor/test_data/vocab.translate_ende_wmt8k.8192.subwords
similarity index 100%
rename from tensor2tensor/test_data/vocab.ende.8192
rename to tensor2tensor/test_data/vocab.translate_ende_wmt8k.8192.subwords

From 03e889baf68664dda80414fcefd0d8bbaa5355c3 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 27 Aug 2018 16:30:22 -0700
Subject: [PATCH 0682/2720] Move next frame models into video sub-directory.

PiperOrigin-RevId: 210452811
---
 tensor2tensor/models/__init__.py              | 11 ++--
 .../next_frame_test.py => video/base_test.py} | 62 +++++++++----------
 .../base_vae.py}                              |  0
 .../basic_deterministic.py}                   |  3 +-
 .../basic_deterministic_params.py}            |  0
 .../basic_stochastic.py}                      | 17 +++--
 .../next_frame_emily.py => video/emily.py}    |  9 +--
 .../next_frame_savp.py => video/savp.py}      |  9 +--
 .../savp_params.py}                           |  4 +-
 .../next_frame_sv2p.py => video/sv2p.py}      |  8 ++-
 .../sv2p_params.py}                           |  5 +-
 11 files changed, 65 insertions(+), 63 deletions(-)
 rename tensor2tensor/models/{research/next_frame_test.py => video/base_test.py} (81%)
 rename tensor2tensor/models/{research/next_frame_base_vae.py => video/base_vae.py} (100%)
 rename tensor2tensor/models/{research/next_frame_basic_deterministic.py => video/basic_deterministic.py} (98%)
 rename tensor2tensor/models/{research/next_frame_basic_deterministic_params.py => video/basic_deterministic_params.py} (100%)
 rename tensor2tensor/models/{research/next_frame_basic_stochastic.py => video/basic_stochastic.py} (86%)
 rename tensor2tensor/models/{research/next_frame_emily.py => video/emily.py} (97%)
 rename tensor2tensor/models/{research/next_frame_savp.py => video/savp.py} (98%)
 rename tensor2tensor/models/{research/next_frame_savp_params.py => video/savp_params.py} (92%)
 rename tensor2tensor/models/{research/next_frame_sv2p.py => video/sv2p.py} (99%)
 rename tensor2tensor/models/{research/next_frame_sv2p_params.py => video/sv2p_params.py} (95%)

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 46fda37e2..1177bc756 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -47,11 +47,6 @@
 from tensor2tensor.models.research import glow
 from tensor2tensor.models.research import lm_experiments
 from tensor2tensor.models.research import multimodel
-from tensor2tensor.models.research import next_frame_basic_deterministic
-from tensor2tensor.models.research import next_frame_basic_stochastic
-from tensor2tensor.models.research import next_frame_emily
-from tensor2tensor.models.research import next_frame_savp
-from tensor2tensor.models.research import next_frame_sv2p
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research import similarity_transformer
 from tensor2tensor.models.research import super_lm
@@ -66,4 +61,10 @@
 from tensor2tensor.models.research import vqa_recurrent_self_attention
 from tensor2tensor.models.research import vqa_self_attention
 
+from tensor2tensor.models.video import basic_deterministic
+from tensor2tensor.models.video import basic_stochastic
+from tensor2tensor.models.video import emily
+from tensor2tensor.models.video import savp
+from tensor2tensor.models.video import sv2p
+
 # pylint: enable=unused-import
diff --git a/tensor2tensor/models/research/next_frame_test.py b/tensor2tensor/models/video/base_test.py
similarity index 81%
rename from tensor2tensor/models/research/next_frame_test.py
rename to tensor2tensor/models/video/base_test.py
index 13f69e766..0a981de3b 100644
--- a/tensor2tensor/models/research/next_frame_test.py
+++ b/tensor2tensor/models/video/base_test.py
@@ -20,14 +20,15 @@
 import numpy as np
 
 from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
-from tensor2tensor.models.research import next_frame_basic_deterministic
-from tensor2tensor.models.research import next_frame_basic_deterministic_params
-from tensor2tensor.models.research import next_frame_basic_stochastic
-from tensor2tensor.models.research import next_frame_emily
-from tensor2tensor.models.research import next_frame_savp
-from tensor2tensor.models.research import next_frame_savp_params
-from tensor2tensor.models.research import next_frame_sv2p
-from tensor2tensor.models.research import next_frame_sv2p_params
+from tensor2tensor.models.video import basic_deterministic
+from tensor2tensor.models.video import basic_deterministic_params
+from tensor2tensor.models.video import basic_stochastic
+from tensor2tensor.models.video import emily
+from tensor2tensor.models.video import savp
+from tensor2tensor.models.video import savp_params
+from tensor2tensor.models.video import sv2p
+from tensor2tensor.models.video import sv2p_params
+
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -198,70 +199,69 @@ def TestOnVariousUpSampleLayers(self, hparams, model, expected_last_dim):
 
   def testBasicDeterministic(self):
     self.TestOnVariousInputOutputSizes(
-        next_frame_basic_deterministic_params.next_frame_basic_deterministic(),
-        next_frame_basic_deterministic.NextFrameBasicDeterministic,
-        256)
+        basic_deterministic_params.next_frame_basic_deterministic(),
+        basic_deterministic.NextFrameBasicDeterministic, 256)
 
   def testBasicStochastic(self):
     self.TestOnVariousInputOutputSizes(
-        next_frame_basic_stochastic.next_frame_basic_stochastic(),
-        next_frame_basic_stochastic.NextFrameBasicStochastic,
+        basic_stochastic.next_frame_basic_stochastic(),
+        basic_stochastic.NextFrameBasicStochastic,
         256)
 
   def testSv2p(self):
     self.TestOnVariousInputOutputSizes(
-        next_frame_sv2p_params.next_frame_sv2p(),
-        next_frame_sv2p.NextFrameSv2p,
+        sv2p_params.next_frame_sv2p(),
+        sv2p.NextFrameSv2p,
         1)
 
   def testSv2pWithActionsAndRewards(self):
     self.TestWithActionAndRewards(
-        next_frame_sv2p_params.next_frame_sv2p(),
-        next_frame_sv2p.NextFrameSv2p,
+        sv2p_params.next_frame_sv2p(),
+        sv2p.NextFrameSv2p,
         1)
 
   def testSv2pTwoFrames(self):
     self.TestOnVariousInputOutputSizes(
-        next_frame_sv2p_params.next_frame_sv2p(),
-        next_frame_sv2p.NextFrameSv2pTwoFrames,
+        sv2p_params.next_frame_sv2p(),
+        sv2p.NextFrameSv2pTwoFrames,
         1)
 
   def testEmily(self):
     self.TestOnVariousInputOutputSizes(
-        next_frame_emily.next_frame_emily(),
-        next_frame_emily.NextFrameEmily,
+        emily.next_frame_emily(),
+        emily.NextFrameEmily,
         1)
 
   def testSavpVAE(self):
-    savp_hparams = next_frame_savp_params.next_frame_savp()
+    savp_hparams = savp_params.next_frame_savp()
     savp_hparams.use_vae = True
     savp_hparams.use_gan = False
     self.TestOnVariousInputOutputSizes(
-        savp_hparams, next_frame_savp.NextFrameSAVP, 1)
+        savp_hparams, savp.NextFrameSAVP, 1)
     self.TestOnVariousUpSampleLayers(
-        savp_hparams, next_frame_savp.NextFrameSAVP, 1)
+        savp_hparams, savp.NextFrameSAVP, 1)
 
   def testSavpGAN(self):
-    hparams = next_frame_savp_params.next_frame_savp()
+    hparams = savp_params.next_frame_savp()
     hparams.use_gan = True
     hparams.use_vae = False
-    self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
+    self.TestVideoModel(7, 5, hparams, savp.NextFrameSAVP, 1)
 
     hparams.gan_optimization = "sequential"
-    self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
+    self.TestVideoModel(7, 5, hparams, savp.NextFrameSAVP, 1)
 
   def testSavpGANVAE(self):
-    hparams = next_frame_savp_params.next_frame_savp()
+    hparams = savp_params.next_frame_savp()
     hparams.use_vae = True
     hparams.use_gan = True
-    self.TestVideoModel(7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
+    self.TestVideoModel(7, 5, hparams, savp.NextFrameSAVP, 1)
 
   def testInvalidVAEGANCombinations(self):
-    hparams = next_frame_savp_params.next_frame_savp()
+    hparams = savp_params.next_frame_savp()
     hparams.use_gan = False
     hparams.use_vae = False
     self.assertRaises(ValueError, self.TestVideoModel,
-                      7, 5, hparams, next_frame_savp.NextFrameSAVP, 1)
+                      7, 5, hparams, savp.NextFrameSAVP, 1)
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/models/research/next_frame_base_vae.py b/tensor2tensor/models/video/base_vae.py
similarity index 100%
rename from tensor2tensor/models/research/next_frame_base_vae.py
rename to tensor2tensor/models/video/base_vae.py
diff --git a/tensor2tensor/models/research/next_frame_basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
similarity index 98%
rename from tensor2tensor/models/research/next_frame_basic_deterministic.py
rename to tensor2tensor/models/video/basic_deterministic.py
index 5cc6b2bd4..91c00daf7 100644
--- a/tensor2tensor/models/research/next_frame_basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -22,7 +22,7 @@
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
-from tensor2tensor.models.research import next_frame_basic_deterministic_params  # pylint: disable=unused-import
+from tensor2tensor.models.video import basic_deterministic_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -169,4 +169,3 @@ def logits_to_samples(logits):
 
     # Return results.
     return results
-
diff --git a/tensor2tensor/models/research/next_frame_basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
similarity index 100%
rename from tensor2tensor/models/research/next_frame_basic_deterministic_params.py
rename to tensor2tensor/models/video/basic_deterministic_params.py
diff --git a/tensor2tensor/models/research/next_frame_basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
similarity index 86%
rename from tensor2tensor/models/research/next_frame_basic_stochastic.py
rename to tensor2tensor/models/video/basic_stochastic.py
index 9585c1745..26ba47a4f 100644
--- a/tensor2tensor/models/research/next_frame_basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -20,9 +20,11 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
-from tensor2tensor.models.research import next_frame_base_vae
-from tensor2tensor.models.research import next_frame_basic_deterministic
-from tensor2tensor.models.research import next_frame_basic_deterministic_params
+
+from tensor2tensor.models.video import base_vae
+from tensor2tensor.models.video import basic_deterministic
+from tensor2tensor.models.video import basic_deterministic_params
+
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -30,8 +32,8 @@
 
 @registry.register_model
 class NextFrameBasicStochastic(
-    next_frame_basic_deterministic.NextFrameBasicDeterministic,
-    next_frame_base_vae.NextFrameBaseVae):
+    basic_deterministic.NextFrameBasicDeterministic,
+    base_vae.NextFrameBaseVae):
   """Stochastic version of basic next-frame model."""
 
   def inject_latent(self, layer, features, filters):
@@ -57,8 +59,7 @@ def inject_latent(self, layer, features, filters):
 @registry.register_hparams
 def next_frame_basic_stochastic():
   """Basic 2-frame conv model with stochastic tower."""
-  base = next_frame_basic_deterministic_params
-  hparams = base.next_frame_basic_deterministic()
+  hparams = basic_deterministic_params.next_frame_basic_deterministic()
   hparams.stochastic_model = True
   hparams.add_hparam("latent_channels", 1)
   hparams.add_hparam("latent_std_min", -5.0)
@@ -70,5 +71,3 @@ def next_frame_basic_stochastic():
   hparams.add_hparam("anneal_end", 100000)
   hparams.add_hparam("information_capacity", 0.0)
   return hparams
-
-
diff --git a/tensor2tensor/models/research/next_frame_emily.py b/tensor2tensor/models/video/emily.py
similarity index 97%
rename from tensor2tensor/models/research/next_frame_emily.py
rename to tensor2tensor/models/video/emily.py
index 9790c2148..cbdc0f479 100644
--- a/tensor2tensor/models/research/next_frame_emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -29,9 +29,10 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
-from tensor2tensor.models.research import next_frame_sv2p
-from tensor2tensor.models.research import next_frame_sv2p_params
+from tensor2tensor.models.video import sv2p
+from tensor2tensor.models.video import sv2p_params
 from tensor2tensor.utils import registry
+
 import tensorflow as tf
 
 tfl = tf.layers
@@ -39,7 +40,7 @@
 
 
 @registry.register_model
-class NextFrameEmily(next_frame_sv2p.NextFrameSv2p):
+class NextFrameEmily(sv2p.NextFrameSv2p):
   """Stochastic Variational Video Prediction Without Learned Prior."""
 
   def encoder(self, inputs, nout):
@@ -266,7 +267,7 @@ def construct_model(self, images, actions, rewards):
 @registry.register_hparams
 def next_frame_emily():
   """Emily's model hparams."""
-  hparams = next_frame_sv2p_params.next_frame_sv2p()
+  hparams = sv2p_params.next_frame_sv2p()
   hparams.latent_loss_multiplier = 1e-4
   hparams.learning_rate_constant = 0.002
   hparams.add_hparam("z_dim", 10)
diff --git a/tensor2tensor/models/research/next_frame_savp.py b/tensor2tensor/models/video/savp.py
similarity index 98%
rename from tensor2tensor/models/research/next_frame_savp.py
rename to tensor2tensor/models/video/savp.py
index 3d2307784..5cd2afb50 100644
--- a/tensor2tensor/models/research/next_frame_savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -25,17 +25,18 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
-from tensor2tensor.models.research import next_frame_savp_params  # pylint: disable=unused-import
-from tensor2tensor.models.research import next_frame_sv2p
+from tensor2tensor.models.video import savp_params  # pylint: disable=unused-import
+from tensor2tensor.models.video import sv2p
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import update_ops_hook
+
 import tensorflow as tf
 
 gan_losses = tf.contrib.gan.losses.wargs
 
 
 @registry.register_model
-class NextFrameSAVP(next_frame_sv2p.NextFrameSv2p):
+class NextFrameSAVP(sv2p.NextFrameSv2p):
   """Stochastic Adversarial Video Prediction."""
 
   def encoder(self, inputs, n_layers=3):
@@ -266,7 +267,7 @@ def get_extra_loss(self, latent_means=None, latent_stds=None,
       return 0.0
 
     vae_loss, d_vae_loss, d_gan_loss = 0.0, 0.0, 0.0
-    # Use next_frame_sv2p's KL divergence computation.
+    # Use sv2p's KL divergence computation.
     if self.hparams.use_vae:
       vae_loss = super(NextFrameSAVP, self).get_extra_loss(
           latent_means=latent_means, latent_stds=latent_stds, beta=beta)
diff --git a/tensor2tensor/models/research/next_frame_savp_params.py b/tensor2tensor/models/video/savp_params.py
similarity index 92%
rename from tensor2tensor/models/research/next_frame_savp_params.py
rename to tensor2tensor/models/video/savp_params.py
index e76d33717..23c6590d8 100644
--- a/tensor2tensor/models/research/next_frame_savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -17,14 +17,14 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.models.research import next_frame_sv2p_params
+from tensor2tensor.models.video import sv2p_params
 from tensor2tensor.utils import registry
 
 
 @registry.register_hparams
 def next_frame_savp():
   """SAVP model hparams."""
-  hparams = next_frame_sv2p_params.next_frame_sv2p()
+  hparams = sv2p_params.next_frame_sv2p()
   hparams.add_hparam("z_dim", 8)
   hparams.add_hparam("num_discriminator_filters", 32)
   hparams.add_hparam("use_vae", True)
diff --git a/tensor2tensor/models/research/next_frame_sv2p.py b/tensor2tensor/models/video/sv2p.py
similarity index 99%
rename from tensor2tensor/models/research/next_frame_sv2p.py
rename to tensor2tensor/models/video/sv2p.py
index fbfae182d..9d7ba1f1f 100644
--- a/tensor2tensor/models/research/next_frame_sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -28,9 +28,11 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
-from tensor2tensor.models.research import next_frame_basic_stochastic
-from tensor2tensor.models.research import next_frame_sv2p_params  # pylint: disable=unused-import
+
+from tensor2tensor.models.video import basic_stochastic
+from tensor2tensor.models.video import sv2p_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
+
 import tensorflow as tf
 
 tfl = tf.layers
@@ -38,7 +40,7 @@
 
 
 @registry.register_model
-class NextFrameSv2p(next_frame_basic_stochastic.NextFrameBasicStochastic):
+class NextFrameSv2p(basic_stochastic.NextFrameBasicStochastic):
   """Stochastic Variational Video Prediction."""
 
   def tinyify(self, array):
diff --git a/tensor2tensor/models/research/next_frame_sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
similarity index 95%
rename from tensor2tensor/models/research/next_frame_sv2p_params.py
rename to tensor2tensor/models/video/sv2p_params.py
index d4d14ffbf..74f98c365 100644
--- a/tensor2tensor/models/research/next_frame_sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -17,14 +17,14 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.models.research import next_frame_basic_stochastic
+from tensor2tensor.models.video import basic_stochastic
 from tensor2tensor.utils import registry
 
 
 @registry.register_hparams
 def next_frame_sv2p():
   """SV2P model hparams."""
-  hparams = next_frame_basic_stochastic.next_frame_basic_stochastic()
+  hparams = basic_stochastic.next_frame_basic_stochastic()
   hparams.optimizer = "TrueAdam"
   hparams.learning_rate_schedule = "constant"
   hparams.learning_rate_constant = 1e-3
@@ -89,4 +89,3 @@ def next_frame_sv2p_cutoff():
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 1
   return hparams
-

From 80c7ea0635ffd24ea98c3ab44c05a5d5ed291898 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 27 Aug 2018 17:14:11 -0700
Subject: [PATCH 0683/2720] fix bug in transformer_vae

PiperOrigin-RevId: 210459174
---
 tensor2tensor/models/research/transformer_vae.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 8e5e8c8b4..e0c34dca4 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -415,7 +415,7 @@ def bn_inputs():
                 mode=hparams.mode,
                 name="vc")
           return bn
-        inputs_c = bn_inputs
+        inputs_c = bn_inputs()
         ptc = 1.0 - common_layers.inverse_lin_decay(200000) * 0.5
         ptc = ptc if hparams.mode == tf.estimator.ModeKeys.TRAIN else 1.0
         latents_dense = tf.where(tf.less(tf.random_uniform([batch_size]), ptc),

From ba047aafffe0a8924b7c861fb1266c097f1c68d2 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 27 Aug 2018 17:14:27 -0700
Subject: [PATCH 0684/2720] refactor ae_transformer_internal a bit

PiperOrigin-RevId: 210459212
---
 tensor2tensor/models/research/transformer_vae.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index e0c34dca4..1026b814d 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -354,7 +354,10 @@ def ae_transformer_internal(inputs,
     if hparams.task == "image":
       cia.maybe_reshape_4d_to_3d(targets)
     if hparams.task == "translate":
-      max_targets_len_from_inputs = tf.concat([inputs, inputs], axis=1)
+      if inputs is not None:
+        max_targets_len_from_inputs = tf.concat([inputs, inputs], axis=1)
+      else:
+        max_targets_len_from_inputs = targets
     else:
       assert hparams.task == "image"
       max_targets_len_from_inputs = targets
@@ -454,7 +457,7 @@ def bn_inputs():
     for i in range(hparams.num_compress_steps):
       j = hparams.num_compress_steps - i - 1
       d = residual_conv(d, 1, (3, 1), hparams, "decompress_rc_%d" % j)
-      if hparams.do_attend_decompress:
+      if inputs is not None and hparams.do_attend_decompress:
         d = attend(d, inputs, hparams, "decompress_attend_%d" % j)
       d = decompress_step(d, hparams, i > 0, False, "decompress_%d" % j)
 
@@ -497,6 +500,11 @@ def refine_res():
     latent_time = tf.less(nonlatent_steps,
                           tf.to_int32(tf.train.get_global_step()))
     losses["latent_pred"] *= tf.to_float(latent_time)
+
+  # res was generated from padded targets, which means it has some extra
+  # elements. These can cause shape problems when computing loss with respect to
+  # the original (unpadded) targets. So we remove their extra elements here.
+  res = res[:, :original_targets_shape[1], :, :]
   return res, losses, cache
 
 

From 42c3f377f441e5a0f431127d63e71414ead291c4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 27 Aug 2018 17:14:44 -0700
Subject: [PATCH 0685/2720] allow real input modality in TransformerAE.infer

PiperOrigin-RevId: 210459243
---
 .../models/research/transformer_vae.py        | 26 ++++++++++++++-----
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 1026b814d..c1835ec38 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -661,15 +661,24 @@ def infer(self, features=None, decode_length=50, beam_size=1, top_beams=1,
     if "partial_targets" in features:
       initial_output = tf.convert_to_tensor(features["partial_targets"])
     else:
-      batch_size = common_layers.shape_list(features["inputs"])[0]
-      length = common_layers.shape_list(features["inputs"])[1]
+      # inputs might not be present in features (e.g.: language modeling),
+      # in which case we fallback to 'infer_targets' for calculating initial
+      # input shape, type, etc.
+      inputs_or_targets = features.get("inputs", features["infer_targets"])
+      batch_size = common_layers.shape_list(inputs_or_targets)[0]
+      length = common_layers.shape_list(inputs_or_targets)[1]
+      hidden_dim = common_layers.shape_list(inputs_or_targets)[-1]
       target_length = tf.to_int32(2.0 * tf.to_float(length))
-      initial_output = tf.zeros((batch_size, target_length, 1, 1),
-                                dtype=tf.int64)
+      initial_output = tf.zeros((batch_size, target_length, 1, hidden_dim),
+                                dtype=inputs_or_targets.dtype)
 
     features["targets"] = initial_output
     logits, _ = self(features)  # pylint: disable=not-callable
-    samples = tf.argmax(logits, axis=-1)
+    # this should only happen if we're doing target_modality not real
+    if inputs_or_targets.dtype == tf.float32:
+      samples = logits
+    else:
+      samples = tf.argmax(logits, axis=-1)
 
     # More steps.
     self.predict_mask = 0.0  # Use the provided targets this time.
@@ -678,7 +687,12 @@ def infer(self, features=None, decode_length=50, beam_size=1, top_beams=1,
       with tf.variable_scope(tf.get_variable_scope(), reuse=True):
         features["targets"] = samples
         logits, _ = self(features)  # pylint: disable=not-callable
-        samples = tf.argmax(logits, axis=-1)
+        if inputs_or_targets.dtype == tf.float32:
+          # When target_modality is real, the last axis does not represent
+          # classes, so it should not be argmax'ed
+          samples = logits
+        else:
+          samples = tf.argmax(logits, axis=-1)
 
     self.predict_mask = 1.0
     if inputs_old is not None:  # Restore to not confuse Estimator.

From bdae1e420598bc76c363b272667108715b32998c Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 27 Aug 2018 21:18:11 -0700
Subject: [PATCH 0686/2720] Fix LM1B8kPacked vocab name

PiperOrigin-RevId: 210486770
---
 tensor2tensor/data_generators/lm1b.py | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/lm1b.py b/tensor2tensor/data_generators/lm1b.py
index 4d9e5e5a6..67f0cc695 100644
--- a/tensor2tensor/data_generators/lm1b.py
+++ b/tensor2tensor/data_generators/lm1b.py
@@ -137,6 +137,14 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
         yield {"targets": txt}
 
 
+@registry.register_problem
+class LanguagemodelLm1b8k(LanguagemodelLm1b32k):
+
+  @property
+  def approx_vocab_size(self):
+    return 2**13  # 8192
+
+
 @registry.register_problem
 class LanguagemodelLm1b32kPacked(LanguagemodelLm1b32k):
   """Packed version for TPU training."""
@@ -151,7 +159,7 @@ def vocab_filename(self):
 
 
 @registry.register_problem
-class LanguagemodelLm1b8kPacked(LanguagemodelLm1b32kPacked):
+class LanguagemodelLm1b8kPacked(LanguagemodelLm1b8k):
   """Packed version, 8k vocabulary.
 
   Ratio of dev tokens (including eos) to dev words (including eos)
@@ -159,8 +167,12 @@ class LanguagemodelLm1b8kPacked(LanguagemodelLm1b32kPacked):
   """
 
   @property
-  def approx_vocab_size(self):
-    return 2**13  # 8192
+  def packed_length(self):
+    return 256
+
+  @property
+  def vocab_filename(self):
+    return LanguagemodelLm1b8k().vocab_filename
 
 
 @registry.register_problem

From 5874f6623f12455ade2680ef02d77a8658af6ea8 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 28 Aug 2018 10:45:35 -0700
Subject: [PATCH 0687/2720] Add video/__init__.py

PiperOrigin-RevId: 210566678
---
 tensor2tensor/models/video/__init__.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)
 create mode 100644 tensor2tensor/models/video/__init__.py

diff --git a/tensor2tensor/models/video/__init__.py b/tensor2tensor/models/video/__init__.py
new file mode 100644
index 000000000..dba7ece95
--- /dev/null
+++ b/tensor2tensor/models/video/__init__.py
@@ -0,0 +1,14 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.

From d2fce2c17646c86ddd42c228183d254bc2260627 Mon Sep 17 00:00:00 2001
From: Youlong Cheng <ylc@google.com>
Date: Tue, 28 Aug 2018 12:12:14 -0700
Subject: [PATCH 0688/2720] Instruction for mesh-tensorflow on Cloud TPU.

PiperOrigin-RevId: 210582762
---
 tensor2tensor/mesh_tensorflow/README.md | 54 ++++++++++++++++++++++++-
 1 file changed, 53 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/mesh_tensorflow/README.md b/tensor2tensor/mesh_tensorflow/README.md
index ec3f20d68..d3d9acf71 100644
--- a/tensor2tensor/mesh_tensorflow/README.md
+++ b/tensor2tensor/mesh_tensorflow/README.md
@@ -279,6 +279,59 @@ the number of cores.  The differences between cores are as follows:
   this by requiring that all imported/exported tensors be fully-replicated.  In
   the future, we should handle this correctly.
 
+# Instructions for running on cloud-tpu
+
+Note: It will be available in `tensorflow>=1.11.0`. For early adoption, use
+`tf-nightly`. Please contact the Google Cloud TPU team if you need to obtain
+`tf-nightly`.
+
+## Prerequisite
+
+Please go through the
+[Transformer tutorial](https://cloud.google.com/tpu/docs/tutorials/transformer).
+
+## Create VM and TPU instance in Cloud console
+
+```sh
+ctpu up -name=ylc-mtf-donut -tf-version=nightly -tpu-size=v2-8 -zone=us-central1-b
+```
+
+## SSH into VM
+
+```sh
+git clone https://github.com/tensorflow/tensor2tensor.git
+cd tensor2tensor/
+pip install --user .
+```
+
+## Run the model
+
+Before run the model, you need to prepare the training data and bucket for
+storing checkpoints. Refer to the
+[Transformer tutorial](https://cloud.google.com/tpu/docs/tutorials/transformer)
+to learn how to generate the training data and create buckets.
+
+```sh
+CONF=mtf_transformer_paper_tr_0_mesh_8
+NAME=ende_$CONF\_0828
+MODEL=mtf_transformer
+PROBLEM=translate_ende_wmt32k_packed
+
+DATA_DIR=gs://xxxx
+OUT_DIR=gs://xxxx
+TPU_NAME=ylc-mtf-donut
+
+tensor2tensor/bin/t2t-trainer \
+  --model=$MODEL \
+  --hparams_set=$CONF \
+  --problem=$PROBLEM \
+  --train_steps=10000 \
+  --eval_steps=200 \
+  --data_dir=$DATA_DIR \
+  --output_dir=$OUT_DIR \
+  --use_tpu=True \
+  --cloud_tpu_name=$TPU_NAME
+```
 
 # TODO LIST (please add items)
 
@@ -286,7 +339,6 @@ We are actively working on improving Mesh-TensorFlow in a variety of ways.  Some
 of the top-priority items are:
 `Contact us if you'd like to help!`
 
-* Instructions for running on cloud-tpu.
 * Operations necessary for spatial-partitioning (spatially-partitioned
   convolution, etc)
 * Examples of image-classification models.

From 5cf0b5c4f30a9e3d8f0ec242b08222461229a515 Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Tue, 28 Aug 2018 12:30:04 -0700
Subject: [PATCH 0689/2720] Add vq_discrete, add vq to mixture of experts and
 new hparams

PiperOrigin-RevId: 210585720
---
 tensor2tensor/layers/vq_discrete.py | 362 ++++++++++++++++++++++++++++
 tensor2tensor/models/transformer.py | 110 ++++++++-
 tensor2tensor/utils/expert_utils.py | 146 +++++++++--
 3 files changed, 595 insertions(+), 23 deletions(-)
 create mode 100644 tensor2tensor/layers/vq_discrete.py

diff --git a/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
new file mode 100644
index 000000000..dc4104204
--- /dev/null
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -0,0 +1,362 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Clean discrete bottleneck as in https://arxiv.org/abs/1805.11063."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from functools import partial
+
+from tensor2tensor.layers import common_layers
+
+import tensorflow as tf
+
+from tensorflow.python.training import moving_averages
+
+
+class DiscreteBottleneck(object):
+  """Discrete bottleneck class."""
+
+  def __init__(self, hparams):
+    self.hparams = hparams
+    print ("self.hparams.z_size", self.hparams.z_size)
+    # Set the discretization bottleneck specific things here
+    self.hparams.z_size_per_residual = self.hparams.z_size // \
+                                       self.hparams.num_residuals
+    print ("self.hparams.num_residuals", self.hparams.num_residuals)
+    self.hparams.block_dim = int(
+        self.hparams.hidden_size // self.hparams.num_blocks)
+    self.hparams.block_v_size = 2**(
+        self.hparams.z_size_per_residual / self.hparams.num_blocks)
+    self.hparams.block_v_size = int(self.hparams.block_v_size)
+    # TODO(avaswani): Figure out why tf.get_variable doesn't work with assign
+    self.hparams.means = tf.Variable(
+        tf.random_normal([
+            self.hparams.num_residuals, self.hparams.num_blocks,
+            self.hparams.block_v_size, self.hparams.block_dim
+        ], stddev=0.1),
+        name="means")
+    tf.logging.info("means = {}".format(self.hparams.means))
+    tf.logging.info("Done creating means")
+
+    # Create the shadow variables if we are using EMA
+    self.hparams.ema_count = None
+    self.hparams.ema_means = None
+    if self.hparams.ema:
+      self.hparams.ema_count = []
+      self.hparams.ema_means = []
+      for i in range(hparams.num_residuals):
+        ema_count_i = tf.get_variable(
+            "ema_count_{}".format(i),
+            [self.hparams.num_blocks, self.hparams.block_v_size],
+            initializer=tf.constant_initializer(0),
+            trainable=False)
+        self.hparams.ema_count.append(ema_count_i)
+
+      with tf.colocate_with(self.hparams.means):
+        self.ema_means = []
+        for i in range(hparams.num_residuals):
+          ema_means_i = tf.get_variable(
+              "ema_means_{}".format(i),
+              initializer=self.hparams.means.initialized_value()[i],
+              trainable=False)
+          self.hparams.ema_means.append(ema_means_i)
+
+  def slice_hidden(self, x):
+    """Slice encoder hidden state into block_dim.
+
+    Args:
+        x: Encoder hidden state of shape [-1, hidden_size].
+
+    Returns:
+        Sliced states of shape [-1, num_blocks, block_dim].
+    """
+    x_sliced = tf.reshape(
+        x, shape=[-1, self.hparams.num_blocks, self.hparams.block_dim])
+    return x_sliced
+
+  def nearest_neighbor(self, x, means):
+    """Find the nearest element in means to elements in x.
+
+    Args:
+        x: Batch of encoder continuous latent states sliced/projected into
+           shape [-1, num_blocks, block_dim].
+        means: Embedding means of shape.
+
+    Returns:
+      Tensor with nearest element in mean encoded in one-hot notation.
+    """
+    x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keep_dims=True)
+    means_norm_sq = tf.reduce_sum(tf.square(means), axis=-1, keep_dims=True)
+    scalar_prod = tf.matmul(
+        tf.transpose(x, perm=[1, 0, 2]), tf.transpose(means, perm=[0, 2, 1]))
+    scalar_prod = tf.transpose(scalar_prod, perm=[1, 0, 2])
+    dist = x_norm_sq + tf.transpose(
+        means_norm_sq, perm=[2, 0, 1]) - 2 * scalar_prod
+
+    if self.hparams.soft_em:
+      nearest_idx = tf.stack(
+          [
+              tf.multinomial(
+                  -dist[:, i, :], num_samples=self.hparams.num_samples)
+              for i in range(self.hparams.num_blocks)
+          ],
+          axis=1)
+      nearest_hot = tf.one_hot(nearest_idx, depth=self.hparams.block_v_size)
+      nearest_hot = tf.reduce_mean(nearest_hot, axis=-2)
+    else:
+      if self.hparams.random_top_k > 1:
+        _, top_k_idx = tf.nn.top_k(-dist, k=self.hparams.random_top_k)
+        nearest_idx = tf.gather(
+            top_k_idx,
+            tf.random_uniform(
+                [1],
+                minval=0,
+                maxval=self.hparams.random_top_k - 1,
+                dtype=tf.int32),
+            axis=-1)
+      else:
+        if self.hparams.use_scales:
+          dist /= tf.reshape(self.hparams.scales,
+                             [1, 1, self.hparams.moe_num_experts])
+        nearest_idx = tf.argmax(-dist, axis=-1)
+      nearest_hot = tf.one_hot(nearest_idx, self.hparams.block_v_size)
+    return nearest_hot
+
+  def embedding_lookup(self, x, means):
+    """Compute nearest neighbors and loss for training the embeddings.
+
+    Args:
+        x: Batch of encoder continuous latent states sliced/projected into
+        shape
+        [-1, num_blocks, block_dim].
+        means: Embedding means.
+
+    Returns:
+        The nearest neighbor in one hot form, the nearest neighbor
+        itself, the
+        commitment loss, embedding training loss.
+    """
+    x_means_hot = self.nearest_neighbor(x, means)
+    x_means_hot_flat = tf.reshape(
+        x_means_hot, [-1, self.hparams.num_blocks, self.hparams.block_v_size])
+    x_means = tf.matmul(tf.transpose(x_means_hot_flat, perm=[1, 0, 2]), means)
+    x_means = tf.transpose(x_means, [1, 0, 2])
+    q_loss = tf.reduce_mean(tf.square((tf.stop_gradient(x) - x_means)))
+    e_loss = tf.reduce_mean((x - tf.stop_gradient(x_means))**2)
+    return x_means_hot, x_means, q_loss, e_loss
+
+  def bit_to_int(self, x_bit, num_bits, base=2):
+    """Turn x_bit representing numbers bitwise (lower-endian) to int tensor.
+
+    Args:
+        x_bit: Tensor containing numbers in a particular base to be
+        converted to
+        int.
+        num_bits: Number of bits in the representation.
+        base: Base of the representation.
+
+    Returns:
+        Integer representation of this number.
+    """
+    x_l = tf.stop_gradient(tf.to_int32(tf.reshape(x_bit, [-1, num_bits])))
+    x_labels = []
+    for i in range(num_bits):
+      x_labels.append(x_l[:, i] * tf.to_int32(base)**tf.to_int32(i))
+    res = sum(x_labels)
+    return tf.to_int32(tf.reshape(res, common_layers.shape_list(x_bit)[:-1]))
+
+  def int_to_bit(self, x_int, num_bits, base=2):
+    """Turn x_int representing numbers into a bitwise (lower-endian) tensor.
+
+    Args:
+        x_int: Tensor containing integer to be converted into base
+        notation.
+        num_bits: Number of bits in the representation.
+        base: Base of the representation.
+
+    Returns:
+        Corresponding number expressed in base.
+    """
+    x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
+    x_labels = []
+    for i in range(num_bits):
+      x_labels.append(
+          tf.floormod(
+              tf.floordiv(tf.to_int32(x_l),
+                          tf.to_int32(base)**i), tf.to_int32(base)))
+    res = tf.concat(x_labels, axis=-1)
+    return tf.to_float(res)
+
+  def embed(self, x, scope="bottleneck"):
+    """Embedding function that takes discrete latent and returns embedding.
+
+    Args:
+        x: Input to the discretization bottleneck.
+        scope: Scope name of the function.
+
+    Returns:
+        Continuous embedding to be passed on to the decoder.
+
+    Raises:
+        ValueError: For unknown or missing arguments.
+    """
+    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
+      shape_x = common_layers.shape_list(x)
+      x_flat = tf.reshape(x, [-1, 1])
+      c = self.int_to_bit(x_flat, num_bits=self.hparams.z_size, base=2)
+      shape = common_layers.shape_list(c)
+      new_shape = shape
+      new_shape[-1] = self.hparams.num_residuals
+      new_shape.append(self.hparams.num_blocks)
+      new_shape.append(
+          int(self.hparams.z_size /
+              (self.hparams.num_residuals * self.hparams.num_blocks)))
+      c = tf.to_int32(tf.reshape(c, shape=new_shape))
+      h1_shape = shape_x
+      h1_shape.append(self.hparams.hidden_size)
+      h1 = tf.zeros(dtype=tf.float32, shape=h1_shape)
+      for i in range(self.hparams.num_residuals):
+        c_residual = self.bit_to_int(
+            c[:, :, i, :, :],
+            num_bits=int(
+                self.hparams.z_size /
+                (self.hparams.num_residuals * self.hparams.num_blocks)),
+            base=2)
+        c_hot = tf.one_hot(c_residual, depth=self.hparams.block_v_size, axis=-1)
+        c_hot_flat = tf.reshape(
+            c_hot,
+            shape=[-1, self.hparams.num_blocks, self.hparams.block_v_size])
+        h1_residual = tf.matmul(
+            tf.transpose(c_hot_flat, perm=[1, 0, 2]), self.hparams.means[i])
+        h1_residual = tf.transpose(h1_residual, perm=[1, 0, 2])
+        h1_residual = tf.reshape(h1_residual, shape=h1_shape)
+        h1 += h1_residual
+
+      # Add Gaussian noise
+      h1_shape[0] = self.hparams.batch_size
+      h2 = tf.layers.dense(
+          tf.nn.relu(h1), self.hparams.filter_size, name="vch2")
+      res = tf.layers.dense(
+          tf.nn.relu(h2), self.hparams.hidden_size, name="vcfin")
+      return res
+
+  def discrete_bottleneck(self, x, scope="bottleneck"):
+    """Discretization bottleneck for latent variables.
+
+    Args:
+        x: Input to the discretization bottleneck.
+        scope: Scope of the function.
+
+    Returns:
+        Embedding to pass to the decoder, discrete latent, loss, and the
+        embedding
+        function.
+
+    Raises:
+        ValueError: If projection_tensors is None for reshape_method
+        project, or
+        ema_count or ema_means is None if we are using ema, or unknown
+        args.
+    """
+    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
+      x_reshaped = self.slice_hidden(x)
+      x_res = x_reshaped
+      x_means_hot = []
+      x_means = 0
+      loss = 0
+      for i in range(self.hparams.num_residuals):
+        x_means_hot_res, x_means_res, q_loss_res, e_loss_res = \
+            self.embedding_lookup(x_reshaped, self.hparams.means[i])
+
+        # Update the ema variables
+        if self.hparams.ema:
+          tf.logging.info("Using EMA with beta = {}".format(self.hparams.beta))
+          updated_ema_count_res = \
+              moving_averages.assign_moving_average(
+                  self.hparams.ema_count[i],
+                  tf.reduce_sum(
+                      tf.reshape(
+                          x_means_hot_res,
+                          shape=[-1, self.hparams.num_blocks,
+                                 self.hparams.block_v_size]),
+                      axis=0),
+                  self.hparams.decay,
+                  zero_debias=False)
+
+          dw = tf.matmul(
+              tf.transpose(x_means_hot_res, perm=[1, 2, 0]),
+              tf.transpose(x_res, perm=[1, 0, 2]))
+
+          updated_ema_means_res = \
+              moving_averages.assign_moving_average(
+                  self.hparams.ema_means[i], dw, self.hparams.decay,
+                  zero_debias=False)
+          n = tf.reduce_sum(updated_ema_count_res, axis=-1, keep_dims=True)
+          updated_ema_count_res = (
+              (updated_ema_count_res + self.hparams.epsilon) /
+              (n + 2**self.hparams.z_size * self.hparams.epsilon) * n)
+          updated_ema_means_res = updated_ema_means_res/tf.expand_dims(
+              updated_ema_count_res, axis=-1)
+          with tf.control_dependencies([e_loss_res]):
+            print ("self.hparams.means[i]", self.hparams.means[i])
+            # raw_input()
+            update_means_res = tf.assign(self.hparams.means[i],
+                                         updated_ema_means_res)
+            # update_means_res = self.hparams.means[i]
+            with tf.control_dependencies([update_means_res]):
+              loss += self.hparams.beta * e_loss_res
+        else:
+          loss += q_loss_res + self.hparams.beta * e_loss_res
+
+        # Update the residuals
+        x_res -= x_means_res
+        x_means += x_means_res
+        x_means_hot.append(x_means_hot_res)
+
+      # Get the discrete latent representation
+      x_means_hot = tf.stack(x_means_hot, axis=1)
+      x_means_idx = tf.argmax(x_means_hot, axis=-1)
+
+      # Get the binary representation
+      num_bits = int(self.hparams.z_size //
+                     (self.hparams.num_blocks * self.hparams.num_residuals))
+      x_means_bits = self.int_to_bit(x_means_idx, num_bits=num_bits, base=2)
+      shape = common_layers.shape_list(x_means_bits)
+      new_shape = shape[:-2]
+      new_shape[0] = -1
+      new_shape[-1] = self.hparams.z_size
+      x_means_bits = tf.reshape(x_means_bits, new_shape)
+      x_discrete = self.bit_to_int(
+          tf.to_int32(x_means_bits), num_bits=self.hparams.z_size, base=2)
+
+      # Reshape x_discrete
+      shape_x = common_layers.shape_list(x)
+      shape_discrete = shape_x[:-1]
+      x_discrete = tf.reshape(x_discrete, shape_discrete)
+      x_means = tf.reshape(x_means, shape=shape_x)
+      h1 = x + tf.stop_gradient(x_means - x)
+
+      h2 = tf.layers.dense(
+          tf.nn.relu(h1), self.hparams.filter_size, name="vch2")
+      res = tf.layers.dense(
+          tf.nn.relu(h2), self.hparams.hidden_size, name="vcfin")
+      embed_fn = partial(self.embed, scope=scope)
+      return {
+          "dense": res,
+          "discrete": x_discrete,
+          "loss": loss,
+          "embed": embed_fn
+      }
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index c31b800c2..35b04c9d4 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -173,10 +173,8 @@ def body(self, features):
     targets = features["targets"]
     targets_shape = common_layers.shape_list(targets)
     targets = common_layers.flatten4d3d(targets)
-
     decoder_input, decoder_self_attention_bias = transformer_prepare_decoder(
         targets, hparams, features=features)
-
     decoder_output = self.decode(
         decoder_input,
         encoder_output,
@@ -1453,10 +1451,19 @@ def transformer_ffn_layer(x,
         hparams.moe_num_experts,
         overhead=overhead,
         loss_coef=hparams.moe_loss_coef)
-    if losses is None:
-      raise ValueError(
-          "transformer_ffn_layer with type local_moe_tpu must pass in "
-          "a losses list")
+  elif ffn_layer == "local_moe":
+    overhead = (
+        hparams.moe_overhead_train
+        if hparams.mode == tf.estimator.ModeKeys.TRAIN else
+        hparams.moe_overhead_eval)
+    ret, loss = expert_utils.local_moe(
+        x,
+        True,
+        expert_utils.ffn_expert_fn(hparams.hidden_size, [hparams.filter_size],
+                                   hparams.hidden_size),
+        hparams.moe_num_experts,
+        k=hparams.moe_k,
+        hparams=hparams)
     losses.append(loss)
     return ret
   else:
@@ -1539,6 +1546,97 @@ def transformer_base_v2():
   return hparams
 
 
+@registry.register_hparams
+def transformer_base_vq_ada_32ex_packed():
+  """Set of hyperparameters for lm1b packed following tpu params."""
+  hparams = transformer_base_v2()
+  expert_utils.update_hparams_for_vq_gating(hparams)
+  hparams.moe_num_experts = 32
+  hparams.gating_type = "vq"
+  # this gives us a batch size of 16 because each seq is len 256
+  hparams.batch_size = 5072
+  hparams.ffn_layer = "local_moe"
+  hparams.shared_embedding_and_softmax_weights = False
+  hparams.learning_rate_warmup_steps = 10000
+  # one epoch for languagemodel_lm1b32k_packed = 27200 steps w/ bsize 128
+  hparams.learning_rate_decay_steps = 27200
+  hparams.num_heads = 4
+  hparams.num_blocks = 1
+  hparams.moe_k = 1
+  hparams.num_decoder_layers = 6
+  hparams.label_smoothing = 0.
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.layer_postprocess_sequence = "dan"
+  hparams.layer_preprocess_sequence = "none"
+  hparams.weight_decay = 1e-06
+  hparams.attention_dropout = 0.1
+  hparams.optimizer = "Adafactor"
+  hparams.learning_rate_schedule = "linear_warmup*rsqrt_decay*linear_decay"
+  hparams.activation_dtype = "float32"
+  hparams.learning_rate = 0.1
+  hparams.learning_rate_constant = 1.0
+  return hparams
+
+
+@registry.register_hparams
+def transformer_topk_16_packed():
+  hparams = transformer_base_vq_ada_32ex_packed()
+  hparams.gating_type = "topk"
+  hparams.moe_num_experts = 16
+  hparams.moe_k = 2
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_vq1_16_nb1_packed_nda_b01_scales():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq_ada_32ex_packed()
+  hparams.use_scales = int(True)
+  hparams.moe_num_experts = 16
+  hparams.moe_k = 1
+  hparams.beta = 0.1
+  hparams.layer_preprocess_sequence = "n"
+  hparams.layer_postprocess_sequence = "da"
+  hparams.ema = False
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_vq1_16_nb1_packed_nda_b01_scales_dialog():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq1_16_nb1_packed_nda_b01_scales()
+  hparams.batch_size = 2048
+  hparams.max_length = 1024
+  hparams.filter_size = 3072
+  return hparams
+
+
+@registry.register_hparams
+def transformer_ada_lmpackedbase():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq_ada_32ex_packed()
+  hparams.ffn_layer = "dense_relu_dense"
+  return hparams
+
+
+@registry.register_hparams
+def transformer_ada_lmpackedbase_dialog():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq_ada_32ex_packed()
+  hparams.max_length = 1024
+  hparams.ffn_layer = "dense_relu_dense"
+  hparams.batch_size = 4096
+  return hparams
+
+
+@registry.register_hparams
+def transformer_ada_lmpackedbase_relative():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq_ada_32ex_packed()
+  hparams.ffn_layer = "dense_relu_dense"
+  return hparams
+
+
 @registry.register_hparams
 def transformer_base():
   """Base parameters for Transformer model."""
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 0be91d7d5..dbdd30290 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -29,6 +29,7 @@
 from six.moves import zip  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers.vq_discrete import DiscreteBottleneck
 
 import tensorflow as tf
 
@@ -379,6 +380,27 @@ def _gates_to_load(gates):
   return tf.reduce_sum(tf.to_float(gates > 0), 0)
 
 
+def update_hparams_for_vq_gating(hparams):
+  """VQ Gating hparams."""
+  hparams.add_hparam("z_size", 4)
+  hparams.add_hparam("noise_dev", 0.5)
+  # Bottleneck kinds supported: dense, vae, dvq.
+  hparams.add_hparam("bottleneck_kind", "dvq")
+  hparams.add_hparam("num_blocks", 1)
+  hparams.add_hparam("num_residuals", 1)
+  # Reshape method for DVQ: slice, project
+  hparams.add_hparam("beta", 0.25)
+  hparams.add_hparam("epsilon", 1e-5)
+  hparams.add_hparam("decay", 0.999)
+  hparams.add_hparam("ema", True)
+  hparams.add_hparam("random_top_k", 1)
+  hparams.add_hparam("soft_em", False)
+  hparams.add_hparam("num_samples", 10)
+  hparams.add_hparam("gating_type", "vq")
+  hparams.add_hparam("use_scales", int(True))
+  hparams.add_hparam("residual_centroids", int(False))
+
+
 def _my_top_k(x, k):
   """GPU-compatible version of top-k that works for very small constant k.
 
@@ -411,6 +433,83 @@ def _my_top_k(x, k):
   return tf.stack(values, axis=1), tf.to_int32(tf.stack(indices, axis=1))
 
 
+def vq_gating(x,
+              num_experts,
+              k,
+              bneck,
+              hparams=None,
+              name="vq_gating"):
+  """VQ gating.
+
+  Args:
+    x: input Tensor with shape [batch_size, input_size]
+    num_experts: an integer
+    k: an integer - number of experts per example
+    bneck: a bottleneck object
+    hparams: optional hparams
+    name: an optional string
+
+  Returns:
+    gates: a Tensor with shape [batch_size, num_experts]
+    load: a Tensor with shape [num_experts]
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+
+    if hparams.use_scales:
+      scales = tf.get_variable(
+          "scales", [num_experts],
+          tf.float32,
+          initializer=tf.ones_initializer())
+      scales = tf.nn.softmax(scales)
+      hparams.scales = scales
+    input_size = x.get_shape().as_list()[-1]
+    batch_size = common_layers.shape_list(x)[0]
+
+    if k > 1:
+      # first project into two dense layers, chop and discretize, and gate
+      # TODO(avaswani): Maybe scale the embeddings flowing out of the experts.
+      # We might want to do this to match the computation being done by topk
+      x = tf.layers.dense(x, input_size * k)
+      # x goes from [batch_size, input_size*k] to [batch_size*k, input_size]
+      x = tf.reshape(x, [batch_size * k, input_size])
+    inputs = tf.expand_dims(x, axis=1)
+    inputs = tf.expand_dims(inputs, axis=1)
+    # VQ hparams
+    hparams.z_size = int(math.log(num_experts, 2))
+    hparams.hidden_size = input_size
+    hparams.top_k = k
+    d = bneck.discrete_bottleneck(inputs, scope=name)
+    centroids = None
+    exp_discrete = d["discrete"]
+    embed_lookup = d["embed"]
+    extra_loss = d["loss"]
+    if hparams.residual_centroids:
+      centroids = embed_lookup(exp_discrete)  # gives the centroids
+    top_k_indices = tf.squeeze(exp_discrete, axis=1)
+    tf.summary.histogram("discrete_counts", top_k_indices)
+    # if k > 1, then we need to reshape top_k_indices from [batch_size*k, 1]
+    # to [batch_size, k]
+    if k > 1:
+      top_k_indices = tf.reshape(top_k_indices, [batch_size, k])
+    # get the top k gates
+    top_k_gates = tf.ones([batch_size, k])
+    # This will be a `Tensor` of shape `[batch_size, n]`, with zeros in the
+    # positions corresponding to all but the top k experts per example.
+    gates = _rowwise_unsorted_segment_sum(top_k_gates, top_k_indices,
+                                          num_experts)
+    # Compute count per expert from the gates.
+    # gates has shape [batch_size, num_experts]
+    # count per expert has shape [num_experts, 1]
+    count_per_expert = tf.reduce_sum(gates, axis=0)
+    if hparams.use_scales:
+      scale_loss = tf.reduce_mean(tf.to_float(count_per_expert) * scales)
+      extra_loss += scale_loss
+    if common_layers.should_generate_summaries():
+      tf.summary.histogram("vq_loss", extra_loss)
+      tf.summary.historgram("scale_loss", scale_loss)
+    return gates, extra_loss, centroids
+
+
 def noisy_top_k_gating(x,
                        num_experts,
                        train,
@@ -459,6 +558,7 @@ def noisy_top_k_gating(x,
     else:
       logits = clean_logits
     top_logits, top_indices = _my_top_k(logits, min(k + 1, num_experts))
+    # top k logits has shape [batch, k]
     top_k_logits = tf.slice(top_logits, [0, 0], [-1, k])
     top_k_indices = tf.slice(top_indices, [0, 0], [-1, k])
     top_k_gates = tf.nn.softmax(top_k_logits)
@@ -963,8 +1063,9 @@ def local_moe(x,
               train,
               expert_fn,
               num_experts,
-              k=2,
+              k=1,
               loss_coef=1e-2,
+              hparams=None,
               pass_x=True,
               pass_gates=False,
               additional_dispatch_params=None,
@@ -978,6 +1079,7 @@ def local_moe(x,
     num_experts: an integer - number of experts
     k: an integer - how many experts to use for each batch element
     loss_coef: a scalar - multiplier on load-balancing losses
+    hparams: optional hparams for vq gating
     pass_x: a boolean. If true, x will also be dispatched to the experts.
     pass_gates: a boolean. If true, gates will be passed to experts. Might be
       necessary when dealing with sparse encoder-encoder decoder attention
@@ -995,21 +1097,31 @@ def local_moe(x,
   """
 
   with tf.variable_scope(name, default_name="local_moe"):
+    centroids = None
     x_flat = flatten_all_but_last(x)
-
-    # The gates indicate which batch elements go to which tensors.
-    # load is a measure of approximately how many examples go to each expert
-    gates, load = noisy_top_k_gating(
-        x_flat,
-        num_experts,
-        train,
-        k,
-        initializer=tf.zeros_initializer(),
-        noisy_gating=True,
-        noise_epsilon=1e-2)
-    # This magic object helps us shuffle data between datashards and experts.
+    if hparams.gating_type == "topk":
+      tf.logging.info("Using noisy top_k with k = {}".format(k))
+      # The gates indicate which batch elements go to which tensors.
+      # load is a measure of approximately how many examples go to each expert
+      gates, load = noisy_top_k_gating(
+          x_flat,
+          num_experts,
+          train,
+          k,
+          initializer=tf.zeros_initializer(),
+          noisy_gating=True,
+          noise_epsilon=1e-2)
+      importance = tf.reduce_sum(gates, 0)
+      loss = loss_coef * (cv_squared(importance) + cv_squared(load))
+    else:
+      assert hparams.gating_type == "vq"
+      tf.logging.info("Using VQ gating")
+      bneck = DiscreteBottleneck(hparams)
+      gates, loss, centroids = vq_gating(
+          x_flat, num_experts, k, bneck, hparams=hparams)
+    loss *= loss_coef
+    # Shuffle data between datashards and experts.
     dispatcher = SparseDispatcher(num_experts, gates)
-
     # Set up expert_fn arguments
     expert_kwargs = {}
     if pass_x:
@@ -1024,10 +1136,10 @@ def local_moe(x,
     expert_outputs = ep(expert_fn, **expert_kwargs)
 
     y_flat = dispatcher.combine(expert_outputs)
+    if centroids is not None:
+      centroids = tf.squeeze(centroids, axis=[1, 2])
+      y_flat += centroids
     y = common_layers.reshape_like(y_flat, x)
-
-    importance = tf.reduce_sum(gates, 0)
-    loss = loss_coef * (cv_squared(importance) + cv_squared(load))
     return y, loss
 
 

From 6bc242127ed4bb60a0d2c42a70010fdc4cd1f7ff Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 28 Aug 2018 12:38:13 -0700
Subject: [PATCH 0690/2720] continuously decode from train set

PiperOrigin-RevId: 210587093
---
 tensor2tensor/utils/decoding.py    |  9 ++++++---
 tensor2tensor/utils/trainer_lib.py | 10 ++++++++--
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 170c045e3..84419f982 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -201,7 +201,7 @@ def decode_from_dataset(estimator,
       hparams=hparams,
       decode_hparams=decode_hp,
       predictions=predictions
-  ))
+  ), dataset_split)
   return predictions
 
 
@@ -785,7 +785,7 @@ class DecodeHookArgs(collections.namedtuple(
   pass
 
 
-def run_postdecode_hooks(decode_hook_args):
+def run_postdecode_hooks(decode_hook_args, dataset_split):
   """Run hooks after decodes have run."""
   hooks = decode_hook_args.problem.decode_hooks
   if not hooks:
@@ -797,7 +797,10 @@ def run_postdecode_hooks(decode_hook_args):
     return
   tf.logging.info("Running decode hooks.")
   parent_dir = os.path.join(decode_hook_args.output_dirs[0], os.pardir)
-  final_dir = os.path.join(parent_dir, "decode")
+  child_dir = "decode"
+  if dataset_split is not None:
+    child_dir += "_{}".format(dataset_split)
+  final_dir = os.path.join(parent_dir, child_dir)
   summary_writer = tf.summary.FileWriter(final_dir)
 
   for hook in hooks:
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 03f9b3894..74c4597b9 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -394,16 +394,22 @@ def run_std_server(self):
         task_index=config.task_id)
     server.join()
 
-  def decode(self):
+  def decode(self, dataset_split=None):
     """Decodes from dataset."""
     decoding.decode_from_dataset(self._estimator, self._hparams.problem.name,
-                                 self._hparams, self._decode_hparams)
+                                 self._hparams, self._decode_hparams,
+                                 dataset_split=dataset_split)
 
   def continuous_decode(self):
     """Decode from dataset on new checkpoint."""
     for _ in next_checkpoint(self._hparams.model_dir):
       self.decode()
 
+  def continuous_decode_on_train_data(self):
+    """Decode from dataset on new checkpoint."""
+    for _ in next_checkpoint(self._hparams.model_dir):
+      self.decode(dataset_split=tf.estimator.ModeKeys.TRAIN)
+
 
 def create_experiment(
     run_config,

From 19ef2ce70603ad715b4c24b7cfa023976e450dc4 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Tue, 28 Aug 2018 14:04:29 -0700
Subject: [PATCH 0691/2720] Change to old discretization

PiperOrigin-RevId: 210602220
---
 tensor2tensor/layers/latent_layers.py | 30 +++++++++++++--------------
 1 file changed, 14 insertions(+), 16 deletions(-)

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 2e66ed3c7..81a156ba9 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -14,14 +14,11 @@
 # limitations under the License.
 """Utils for latent variable models."""
 
-import functools
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
-from tensor2tensor.layers import discretization
 from tensor2tensor.models import transformer
 from tensor2tensor.utils import beam_search
 
@@ -463,18 +460,19 @@ def transformer_latent_decoder(x,
     return decoder_output
 
 
-def bottleneck_layer(targets_c, hparams):
+def bottleneck_layer(targets_c,
+                     hparams,
+                     name="bottleneck_d"):
   """Compute latents from compressed targets."""
-  latents_discrete_hot, extra_loss = discretization.parametrized_bottleneck(
-      targets_c, hparams)
-  latents_dense = discretization.parametrized_unbottleneck(
-      latents_discrete_hot, hparams.hidden_size, hparams)
-  latents_dense = targets_c + tf.stop_gradient(latents_dense - targets_c)
-  latents_discrete = tf.argmax(latents_discrete_hot, axis=-1)
-
+  latents_dense, latents_discrete, extra_loss, embed_func = (
+      hparams.bottleneck(
+          inputs=targets_c,
+          filter_size=hparams.compress_filter_size,
+          name=name,
+          mode=hparams.mode))
   if DO_SUMMARIES:
-    tf.summary.histogram("discrete_latents", tf.reshape(latents_discrete, [-1]))
-  return latents_dense, latents_discrete_hot, extra_loss
+    tf.summary.histogram("b0", tf.reshape(latents_discrete, [-1]))
+  return latents_dense, latents_discrete, extra_loss, embed_func
 
 
 def latent_prediction_model(inputs,
@@ -557,7 +555,7 @@ def transformer_autoencoder(inputs,
   # Call bottleneck layer, that takes encoder output and outputs the latents.
   # Returns embedded latents, discrete latent codes, loss.
   if hparams.mode != tf.estimator.ModeKeys.PREDICT:
-    latents_dense, latents_discrete, extra_loss = (
+    latents_dense, latents_discrete, extra_loss, _ = (
         bottleneck_layer(targets_c, hparams))
     extra_loss = tf.reduce_mean(extra_loss) * tf.to_float(cond)
 
@@ -586,8 +584,8 @@ def transformer_autoencoder(inputs,
     latent_len = (
         hparams.img_len * hparams.img_len * hparams.num_latents) / 2**(
             hparams.num_compress_steps)
-    embed = functools.partial(
-        discretization.parametrized_unbottleneck, hparams=hparams)
+    _, _, _, embed = (
+        bottleneck_layer(targets_c, hparams))
     latents_dense = tf.zeros([batch_size, latent_len, 1, hparams.hidden_size])
     if cache is None:
       cache = ae_latent_sample_beam(latents_dense, inputs, ed_attention_bias,

From 14ea28183f08f1983b75b59d0e0f7160b649cbaf Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 28 Aug 2018 15:36:11 -0700
Subject: [PATCH 0692/2720] Add INFO verbosity to binaries

PiperOrigin-RevId: 210619888
---
 tensor2tensor/bin/build_vocab.py       | 1 +
 tensor2tensor/bin/make_tf_configs.py   | 1 +
 tensor2tensor/bin/t2t-avg-all          | 1 +
 tensor2tensor/bin/t2t-bleu             | 1 +
 tensor2tensor/bin/t2t-datagen          | 1 +
 tensor2tensor/bin/t2t-decoder          | 1 +
 tensor2tensor/bin/t2t-exporter         | 1 +
 tensor2tensor/bin/t2t-insights-server  | 1 +
 tensor2tensor/bin/t2t-make-tf-configs  | 1 +
 tensor2tensor/bin/t2t-query-server     | 1 +
 tensor2tensor/bin/t2t-trainer          | 1 +
 tensor2tensor/bin/t2t-translate-all    | 1 +
 tensor2tensor/bin/t2t_attack.py        | 1 +
 tensor2tensor/bin/t2t_avg_all.py       | 1 +
 tensor2tensor/bin/t2t_bleu.py          | 1 +
 tensor2tensor/bin/t2t_decoder.py       | 1 +
 tensor2tensor/bin/t2t_distill.py       | 1 +
 tensor2tensor/bin/t2t_prune.py         | 1 +
 tensor2tensor/bin/t2t_trainer.py       | 1 +
 tensor2tensor/bin/t2t_translate_all.py | 1 +
 20 files changed, 20 insertions(+)

diff --git a/tensor2tensor/bin/build_vocab.py b/tensor2tensor/bin/build_vocab.py
index 0638ae47d..0f363fa3c 100644
--- a/tensor2tensor/bin/build_vocab.py
+++ b/tensor2tensor/bin/build_vocab.py
@@ -67,4 +67,5 @@ def main(_):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/make_tf_configs.py b/tensor2tensor/bin/make_tf_configs.py
index b38f16530..321b51d55 100644
--- a/tensor2tensor/bin/make_tf_configs.py
+++ b/tensor2tensor/bin/make_tf_configs.py
@@ -106,4 +106,5 @@ def main(_):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t-avg-all b/tensor2tensor/bin/t2t-avg-all
index abef8b755..696a20b5b 100755
--- a/tensor2tensor/bin/t2t-avg-all
+++ b/tensor2tensor/bin/t2t-avg-all
@@ -13,4 +13,5 @@ def main(argv):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t-bleu b/tensor2tensor/bin/t2t-bleu
index 966f50a81..6c51e480b 100755
--- a/tensor2tensor/bin/t2t-bleu
+++ b/tensor2tensor/bin/t2t-bleu
@@ -14,4 +14,5 @@ def main(argv):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t-datagen b/tensor2tensor/bin/t2t-datagen
index 332f063ad..a49ed5615 100755
--- a/tensor2tensor/bin/t2t-datagen
+++ b/tensor2tensor/bin/t2t-datagen
@@ -24,4 +24,5 @@ def main(argv):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t-decoder b/tensor2tensor/bin/t2t-decoder
index 612117c22..9bcca1e9b 100755
--- a/tensor2tensor/bin/t2t-decoder
+++ b/tensor2tensor/bin/t2t-decoder
@@ -13,4 +13,5 @@ def main(argv):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t-exporter b/tensor2tensor/bin/t2t-exporter
index cfd4f5ff8..3166b0ee3 100755
--- a/tensor2tensor/bin/t2t-exporter
+++ b/tensor2tensor/bin/t2t-exporter
@@ -13,4 +13,5 @@ def main(argv):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t-insights-server b/tensor2tensor/bin/t2t-insights-server
index 102202c9b..e757b783c 100755
--- a/tensor2tensor/bin/t2t-insights-server
+++ b/tensor2tensor/bin/t2t-insights-server
@@ -13,4 +13,5 @@ def main(argv):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t-make-tf-configs b/tensor2tensor/bin/t2t-make-tf-configs
index b481ea910..7142e9673 100755
--- a/tensor2tensor/bin/t2t-make-tf-configs
+++ b/tensor2tensor/bin/t2t-make-tf-configs
@@ -13,4 +13,5 @@ def main(argv):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t-query-server b/tensor2tensor/bin/t2t-query-server
index 91ede7ce7..a354819db 100755
--- a/tensor2tensor/bin/t2t-query-server
+++ b/tensor2tensor/bin/t2t-query-server
@@ -13,4 +13,5 @@ def main(argv):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t-trainer b/tensor2tensor/bin/t2t-trainer
index 1d848d04d..c2f129409 100755
--- a/tensor2tensor/bin/t2t-trainer
+++ b/tensor2tensor/bin/t2t-trainer
@@ -29,4 +29,5 @@ def main(argv):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t-translate-all b/tensor2tensor/bin/t2t-translate-all
index fed5d3045..9e8ee219b 100755
--- a/tensor2tensor/bin/t2t-translate-all
+++ b/tensor2tensor/bin/t2t-translate-all
@@ -14,4 +14,5 @@ def main(argv):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 09d190b68..faabdcd8e 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -275,4 +275,5 @@ def compute_accuracy(x, l, mask):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t_avg_all.py b/tensor2tensor/bin/t2t_avg_all.py
index a3f96d3a4..2843ff3ce 100644
--- a/tensor2tensor/bin/t2t_avg_all.py
+++ b/tensor2tensor/bin/t2t_avg_all.py
@@ -107,4 +107,5 @@ def main(_):
       avg_values[name] -= reader.get_tensor(name) / FLAGS.n
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t_bleu.py b/tensor2tensor/bin/t2t_bleu.py
index 338b2adaa..db18ec91f 100644
--- a/tensor2tensor/bin/t2t_bleu.py
+++ b/tensor2tensor/bin/t2t_bleu.py
@@ -170,4 +170,5 @@ def main(_):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index f8f844779..c7d2814d9 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -195,4 +195,5 @@ def main(_):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t_distill.py b/tensor2tensor/bin/t2t_distill.py
index 91f9caece..adbeac667 100644
--- a/tensor2tensor/bin/t2t_distill.py
+++ b/tensor2tensor/bin/t2t_distill.py
@@ -155,4 +155,5 @@ def student_experiment_fn(run_config, hparams):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
index 99950b933..6e32597b5 100644
--- a/tensor2tensor/bin/t2t_prune.py
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -113,4 +113,5 @@ def eval_model():
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 09cba7013..01bfe7c12 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -386,4 +386,5 @@ def main(argv):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/bin/t2t_translate_all.py b/tensor2tensor/bin/t2t_translate_all.py
index 9194fa43d..361720a0b 100644
--- a/tensor2tensor/bin/t2t_translate_all.py
+++ b/tensor2tensor/bin/t2t_translate_all.py
@@ -102,4 +102,5 @@ def main(_):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()

From db395a47af1522310cb7c888fac9295b177c5f14 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 28 Aug 2018 16:10:19 -0700
Subject: [PATCH 0693/2720] Internal

PiperOrigin-RevId: 210625665
---
 tensor2tensor/bin/t2t_datagen.py              | 1 -
 tensor2tensor/data_generators/gym_problems.py | 2 ++
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index d8fa6ce1a..30e21e1b4 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -54,7 +54,6 @@
 import tensorflow as tf  # pylint: disable=g-import-not-at-top
 
 
-
 flags = tf.flags
 FLAGS = flags.FLAGS
 
diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 0fda7aefc..71893f11d 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -40,6 +40,8 @@
 flags = tf.flags
 FLAGS = flags.FLAGS
 
+
+
 flags.DEFINE_string("agent_policy_path", None, "File with model for agent.")
 
 flags.DEFINE_string("autoencoder_path", None,

From 2244fdd3652743f285ad18c928cdf3054d8f9a0c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 28 Aug 2018 16:20:00 -0700
Subject: [PATCH 0694/2720] Internal change

PiperOrigin-RevId: 210627141
---
 .../data_generators/cnn_dailymail.py          | 50 ++++++++++++-------
 1 file changed, 33 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index 00464db0f..79e7d011c 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -57,12 +57,12 @@
 ]
 
 
-def _maybe_download_corpora(tmp_dir, is_training):
+def _maybe_download_corpora(tmp_dir, dataset_split):
   """Download corpora if necessary and unzip them.
 
   Args:
     tmp_dir: directory containing dataset.
-    is_training: whether we're in training mode or not.
+    dataset_split: whether we're in train/dev/test mode.
 
   Returns:
     List of all files generated and path to file containing
@@ -87,12 +87,15 @@ def _maybe_download_corpora(tmp_dir, is_training):
   dailymail_files = tf.gfile.Glob(dailymail_finalpath + "*")
   all_files = cnn_files + dailymail_files
 
-  if is_training:
+  if dataset_split == problem.DatasetSplit.TRAIN:
     urls_path = generator_utils.maybe_download(tmp_dir, "all_train.txt",
                                                _TRAIN_URLS)
-  else:
+  elif dataset_split == problem.DatasetSplit.EVAL:
     urls_path = generator_utils.maybe_download(tmp_dir, "all_val.txt",
                                                _DEV_URLS)
+  else:
+    urls_path = generator_utils.maybe_download(tmp_dir, "all_test.txt",
+                                               _TEST_URLS)
 
   return all_files, urls_path
 
@@ -175,7 +178,7 @@ def _story_summary_split(story):
   return story[:split_pos], story[split_pos + split_str_len:]  # story, summary
 
 
-def write_raw_text_to_files(all_files, urls_path, tmp_dir, is_training):
+def write_raw_text_to_files(all_files, urls_path, dataset_split, tmp_dir):
   """Write text to files."""
 
   def write_to_file(all_files, urls_path, tmp_dir, filename):
@@ -187,17 +190,16 @@ def write_to_file(all_files, urls_path, tmp_dir, filename):
           fstory.write(story + "\n")
           fsummary.write(summary + "\n")
 
-  filename = "cnndm.train" if is_training else "cnndm.dev"
+  if dataset_split == problem.DatasetSplit.TRAIN:
+    filename = "cnndm.train"
+  elif dataset_split == problem.DatasetSplit.EVAL:
+    filename = "cnndm.dev"
+  else:
+    filename = "cnndm.test"
+
   tf.logging.info("Writing %s" % filename)
   write_to_file(all_files, urls_path, tmp_dir, filename)
 
-  if not is_training:
-    test_urls_path = generator_utils.maybe_download(tmp_dir, "all_test.txt",
-                                                    _TEST_URLS)
-    filename = "cnndm.test"
-    tf.logging.info("Writing %s" % filename)
-    write_to_file(all_files, test_urls_path, tmp_dir, filename)
-
 
 @registry.register_problem
 class SummarizeCnnDailymail32k(text_problems.Text2TextProblem):
@@ -205,17 +207,31 @@ class SummarizeCnnDailymail32k(text_problems.Text2TextProblem):
 
   def generate_text_for_vocab(self, data_dir, tmp_dir):
     del data_dir
-    all_files, urls_path = _maybe_download_corpora(tmp_dir, True)
+    all_files, urls_path = _maybe_download_corpora(tmp_dir,
+                                                   problem.DatasetSplit.TRAIN)
     return example_generator(all_files, urls_path, sum_token=False)
 
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each."""
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 100,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 10,
+    }, {
+        "split": problem.DatasetSplit.TEST,
+        "shards": 10,
+    }]
+
   def is_generate_per_split(self):
     return True
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
     del data_dir
-    is_training = dataset_split == problem.DatasetSplit.TRAIN
-    all_files, urls_path = _maybe_download_corpora(tmp_dir, is_training)
-    write_raw_text_to_files(all_files, urls_path, tmp_dir, is_training)
+    all_files, urls_path = _maybe_download_corpora(tmp_dir, dataset_split)
+    write_raw_text_to_files(all_files, urls_path, dataset_split, tmp_dir)
     for example in example_generator(all_files, urls_path, sum_token=True):
       story, summary = _story_summary_split(example)
       yield {"inputs": story, "targets": summary}

From b5df6389ae2692841f6fe376f9bc5e2bff77ce16 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 28 Aug 2018 16:56:17 -0700
Subject: [PATCH 0695/2720] Discriminator work, try a simple form of attention
 in addition to mean.

PiperOrigin-RevId: 210632857
---
 tensor2tensor/layers/common_layers.py         | 37 +++++++++-----
 tensor2tensor/models/research/autoencoders.py | 48 ++++++++++---------
 2 files changed, 52 insertions(+), 33 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index c6c158eee..9720a7a07 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3786,22 +3786,36 @@ def patch_discriminator(x, filters=64, filter_size=5, n=4,
     return x
 
 
-def simple_discriminator(x, filters=128, kernel_size=7,
-                         strides=4, do_mean=True):
-  """A very simple convolutional discriminator."""
+def mean_with_attention(x, name, num_heads=4):
+  """Mean and attention to reduce spatial dimensions."""
+  with tf.variable_scope(name):
+    shape = shape_list(x)
+    m = tf.reduce_mean(x, [1, 2])
+    a = tf.layers.dense(x, num_heads, name="mean_attn")
+    s = tf.reshape(a, [shape[0], -1, num_heads])
+    s = tf.nn.softmax(s, axis=1)
+    s = tf.reshape(s, shape[:-1] + [1, num_heads])
+    am = tf.reduce_mean(tf.expand_dims(x, axis=-1) * s, [1, 2])
+    l = tf.concat([am, tf.expand_dims(m, axis=-1)], axis=-1)
+    return tf.layers.dense(tf.reshape(l, [shape[0], (num_heads+1) * shape[-1]]),
+                           2 * shape[-1], name="mean_attn_final")
+
+
+def single_discriminator(x, filters=128, kernel_size=7,
+                         strides=4, pure_mean=True):
+  """A simple single-layer convolutional discriminator."""
   with tf.variable_scope("discriminator"):
     net = tf.layers.conv2d(
         x, filters, kernel_size, strides=strides, padding="SAME", name="conv1")
-    if do_mean:
+    if pure_mean:
       net = tf.reduce_mean(net, [1, 2])
     else:
-      batch_size = shape_list(x)[0]
-      net = tf.reshape(net, [batch_size, -1])
+      net = mean_with_attention(net, "mean_with_attention")
     return net
 
 
 def double_discriminator(x, filters1=128, filters2=None,
-                         kernel_size=7, strides=4, do_mean=True):
+                         kernel_size=7, strides=4, pure_mean=True):
   """A convolutional discriminator with 2 layers and concatenated output."""
   if filters2 is None:
     filters2 = 4 * filters1
@@ -3809,17 +3823,18 @@ def double_discriminator(x, filters1=128, filters2=None,
     batch_size = shape_list(x)[0]
     net = tf.layers.conv2d(
         x, filters1, kernel_size, strides=strides, padding="SAME", name="conv1")
-    if do_mean:
+    if pure_mean:
       net1 = tf.reduce_mean(net, [1, 2])
     else:
-      net1 = tf.reshape(net, [batch_size, -1])
+      net1 = mean_with_attention(net, "mean_with_attention1")
+      tf.reshape(net, [batch_size, -1])
     net = tf.nn.relu(net)
     net = tf.layers.conv2d(
         x, filters2, kernel_size, strides=strides, padding="SAME", name="conv2")
-    if do_mean:
+    if pure_mean:
       net2 = tf.reduce_mean(net, [1, 2])
     else:
-      net2 = tf.reshape(net, [batch_size, -1])
+      net2 = mean_with_attention(net, "mean_with_attention2")
     return tf.concat([net1, net2], axis=-1)
 
 
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 7cc089d07..dedaeac7b 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -306,20 +306,20 @@ def discriminate(x):
               x, hparams.discriminator_batchnorm, is_training)
         elif hparams.discriminator == "patched":
           return common_layers.patch_discriminator(x)
-        elif hparams.discriminator == "simple":
-          return common_layers.simple_discriminator(
+        elif hparams.discriminator == "single":
+          return common_layers.single_discriminator(
               x,
               hparams.discriminator_size,
               hparams.discriminator_kernel_size,
               hparams.discriminator_strides,
-              do_mean=hparams.discriminator_do_mean)
+              pure_mean=hparams.discriminator_pure_mean)
         elif hparams.discriminator == "double":
           return common_layers.double_discriminator(
               x,
               hparams.discriminator_size,
               hparams.discriminator_kernel_size,
               hparams.discriminator_strides,
-              do_mean=hparams.discriminator_do_mean)
+              pure_mean=hparams.discriminator_pure_mean)
         else:
           raise Exception("Unknown discriminator %s" % hparams.discriminator)
 
@@ -439,7 +439,10 @@ def body(self, features):
       assert hparams.mode == tf.estimator.ModeKeys.PREDICT
       targets = tf.zeros_like(basic_result)
     targets = self.embed(targets)
-    basic_hot = self.gumbel_sample(basic_result)
+    if hparams.autoregressive_gumbel_sample:
+      basic_hot = self.gumbel_sample(basic_result)
+    else:
+      basic_hot = basic_result
     basic_result = self.embed(basic_hot)
     shape = common_layers.shape_list(basic_result)
     basic1d = tf.reshape(basic_result, [shape[0], -1, shape[-1]])
@@ -905,11 +908,11 @@ def autoencoder_basic():
   hparams.weight_decay = 0.0
   hparams.kernel_height = 4
   hparams.kernel_width = 4
-  hparams.dropout = 0.1
+  hparams.dropout = 0.05
   hparams.add_hparam("max_hidden_size", 1024)
   hparams.add_hparam("bottleneck_bits", 128)
   hparams.add_hparam("bottleneck_noise", 0.1)
-  hparams.add_hparam("bottleneck_warmup_steps", 500)
+  hparams.add_hparam("bottleneck_warmup_steps", 2000)
   hparams.add_hparam("sample_height", 32)
   hparams.add_hparam("sample_width", 32)
   hparams.add_hparam("discriminator_batchnorm", True)
@@ -918,16 +921,16 @@ def autoencoder_basic():
   hparams.add_hparam("discriminator_size", 256)
   hparams.add_hparam("discriminator_kernel_size", 6)
   hparams.add_hparam("discriminator_strides", 4)
-  hparams.add_hparam("discriminator_do_mean", int(True))
+  hparams.add_hparam("discriminator_pure_mean", int(False))
   hparams.add_hparam("code_loss_factor", 1.0)
-  hparams.add_hparam("gan_codes_warmup_steps", 6000)
+  hparams.add_hparam("gan_codes_warmup_steps", 16000)
   hparams.add_hparam("gan_loss_factor", 0.0)
   hparams.add_hparam("bottleneck_l2_factor", 0.05)
-  hparams.add_hparam("gumbel_temperature", 0.2)
-  hparams.add_hparam("gumbel_noise_factor", 0.4)
+  hparams.add_hparam("gumbel_temperature", 0.5)
+  hparams.add_hparam("gumbel_noise_factor", 0.5)
   hparams.add_hparam("vq_temperature", 0.001)
   hparams.add_hparam("use_vq_loss", int(False))
-  hparams.add_hparam("discriminator", "simple")
+  hparams.add_hparam("discriminator", "double")
   return hparams
 
 
@@ -939,6 +942,7 @@ def autoencoder_autoregressive():
   hparams.add_hparam("autoregressive_mode", "none")
   hparams.add_hparam("autoregressive_decode_steps", 0)
   hparams.add_hparam("autoregressive_eval_pure_autoencoder", False)
+  hparams.add_hparam("autoregressive_gumbel_sample", False)
   return hparams
 
 
@@ -988,7 +992,7 @@ def autoencoder_basic_discrete():
   hparams.hidden_size = 64
   hparams.bottleneck_bits = 4096
   hparams.bottleneck_noise = 0.1
-  hparams.add_hparam("discretize_warmup_steps", 5000)
+  hparams.add_hparam("discretize_warmup_steps", 16000)
   return hparams
 
 
@@ -998,7 +1002,7 @@ def autoencoder_residual_discrete():
   hparams = autoencoder_residual()
   hparams.bottleneck_bits = 4096
   hparams.bottleneck_noise = 0.1
-  hparams.add_hparam("discretize_warmup_steps", 5000)
+  hparams.add_hparam("discretize_warmup_steps", 16000)
   hparams.add_hparam("bottleneck_kind", "tanh_discrete")
   hparams.add_hparam("isemhash_noise_dev", 0.5)
   hparams.add_hparam("isemhash_mix_prob", 0.5)
@@ -1025,7 +1029,7 @@ def autoencoder_ordered_discrete():
   """Ordered discrete autoencoder model."""
   hparams = autoencoder_residual_discrete()
   hparams.bottleneck_noise = 0.8
-  hparams.gan_loss_factor = 0.02
+  hparams.gan_loss_factor = 0.05
   hparams.add_hparam("unordered", False)
   return hparams
 
@@ -1039,10 +1043,10 @@ def autoencoder_ordered_discrete_patched():
 
 
 @registry.register_hparams
-def autoencoder_ordered_discrete_simple():
+def autoencoder_ordered_discrete_single():
   """Ordered discrete autoencoder model."""
   hparams = autoencoder_ordered_discrete()
-  hparams.discriminator = "simple"
+  hparams.discriminator = "single"
   return hparams
 
 
@@ -1058,11 +1062,11 @@ def autoencoder_ordered_discrete_hs256():
 def autoencoder_ordered_text():
   """Ordered discrete autoencoder model for text."""
   hparams = autoencoder_ordered_discrete()
-  hparams.bottleneck_bits = 1024
+  hparams.bottleneck_bits = 2048
+  hparams.num_hidden_layers = 7
   hparams.batch_size = 1024
-  hparams.autoregressive_mode = "sru"
-  hparams.hidden_size = 256
-  hparams.max_hidden_size = 4096
+  hparams.autoregressive_mode = "conv5"
+  hparams.max_hidden_size = 1024
   hparams.target_modality = "symbol:identity"
   hparams.input_modalities = "symbol:identity"
   hparams.sample_width = 1
@@ -1125,7 +1129,7 @@ def autoencoder_range(rhp):
   rhp.set_float("dropout", 0.01, 0.3)
   rhp.set_float("gan_loss_factor", 0.01, 0.1)
   rhp.set_float("bottleneck_l2_factor", 0.001, 0.1, scale=rhp.LOG_SCALE)
-  rhp.set_discrete("bottleneck_warmup_steps", [200, 500, 1000, 2000])
+  rhp.set_discrete("bottleneck_warmup_steps", [200, 2000])
   rhp.set_float("gumbel_temperature", 0, 1)
   rhp.set_float("gumbel_noise_factor", 0, 0.5)
 

From 7fcc3748cac87f27d99e100a68458a70858b52ae Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Tue, 28 Aug 2018 17:28:44 -0700
Subject: [PATCH 0696/2720] Bring back uniform_unit_scaling. Results without it
 are bad.

PiperOrigin-RevId: 210637952
---
 tensor2tensor/layers/vq_discrete.py | 68 ++++-------------------------
 tensor2tensor/models/transformer.py | 12 +++++
 tensor2tensor/utils/expert_utils.py |  2 +-
 3 files changed, 21 insertions(+), 61 deletions(-)

diff --git a/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index dc4104204..4b0180269 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -23,8 +23,6 @@
 
 import tensorflow as tf
 
-from tensorflow.python.training import moving_averages
-
 
 class DiscreteBottleneck(object):
   """Discrete bottleneck class."""
@@ -42,12 +40,13 @@ def __init__(self, hparams):
         self.hparams.z_size_per_residual / self.hparams.num_blocks)
     self.hparams.block_v_size = int(self.hparams.block_v_size)
     # TODO(avaswani): Figure out why tf.get_variable doesn't work with assign
-    self.hparams.means = tf.Variable(
-        tf.random_normal([
+    self.hparams.means = tf.get_variable(
+        name="means",
+        shape=[
             self.hparams.num_residuals, self.hparams.num_blocks,
             self.hparams.block_v_size, self.hparams.block_dim
-        ], stddev=0.1),
-        name="means")
+            ],
+        initializer=tf.uniform_unit_scaling_initializer())
     tf.logging.info("means = {}".format(self.hparams.means))
     tf.logging.info("Done creating means")
 
@@ -55,24 +54,7 @@ def __init__(self, hparams):
     self.hparams.ema_count = None
     self.hparams.ema_means = None
     if self.hparams.ema:
-      self.hparams.ema_count = []
-      self.hparams.ema_means = []
-      for i in range(hparams.num_residuals):
-        ema_count_i = tf.get_variable(
-            "ema_count_{}".format(i),
-            [self.hparams.num_blocks, self.hparams.block_v_size],
-            initializer=tf.constant_initializer(0),
-            trainable=False)
-        self.hparams.ema_count.append(ema_count_i)
-
-      with tf.colocate_with(self.hparams.means):
-        self.ema_means = []
-        for i in range(hparams.num_residuals):
-          ema_means_i = tf.get_variable(
-              "ema_means_{}".format(i),
-              initializer=self.hparams.means.initialized_value()[i],
-              trainable=False)
-          self.hparams.ema_means.append(ema_means_i)
+      raise NotImplementedError("ema updates not implemented")
 
   def slice_hidden(self, x):
     """Slice encoder hidden state into block_dim.
@@ -281,43 +263,9 @@ def discrete_bottleneck(self, x, scope="bottleneck"):
         x_means_hot_res, x_means_res, q_loss_res, e_loss_res = \
             self.embedding_lookup(x_reshaped, self.hparams.means[i])
 
-        # Update the ema variables
+        # TODO(avaswani,nikip,aurkor): Implement ema
         if self.hparams.ema:
-          tf.logging.info("Using EMA with beta = {}".format(self.hparams.beta))
-          updated_ema_count_res = \
-              moving_averages.assign_moving_average(
-                  self.hparams.ema_count[i],
-                  tf.reduce_sum(
-                      tf.reshape(
-                          x_means_hot_res,
-                          shape=[-1, self.hparams.num_blocks,
-                                 self.hparams.block_v_size]),
-                      axis=0),
-                  self.hparams.decay,
-                  zero_debias=False)
-
-          dw = tf.matmul(
-              tf.transpose(x_means_hot_res, perm=[1, 2, 0]),
-              tf.transpose(x_res, perm=[1, 0, 2]))
-
-          updated_ema_means_res = \
-              moving_averages.assign_moving_average(
-                  self.hparams.ema_means[i], dw, self.hparams.decay,
-                  zero_debias=False)
-          n = tf.reduce_sum(updated_ema_count_res, axis=-1, keep_dims=True)
-          updated_ema_count_res = (
-              (updated_ema_count_res + self.hparams.epsilon) /
-              (n + 2**self.hparams.z_size * self.hparams.epsilon) * n)
-          updated_ema_means_res = updated_ema_means_res/tf.expand_dims(
-              updated_ema_count_res, axis=-1)
-          with tf.control_dependencies([e_loss_res]):
-            print ("self.hparams.means[i]", self.hparams.means[i])
-            # raw_input()
-            update_means_res = tf.assign(self.hparams.means[i],
-                                         updated_ema_means_res)
-            # update_means_res = self.hparams.means[i]
-            with tf.control_dependencies([update_means_res]):
-              loss += self.hparams.beta * e_loss_res
+          raise NotImplementedError("ema updates not implemented")
         else:
           loss += q_loss_res + self.hparams.beta * e_loss_res
 
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 35b04c9d4..8ac1a16a7 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1601,6 +1601,18 @@ def transformer_base_vq1_16_nb1_packed_nda_b01_scales():
   return hparams
 
 
+@registry.register_hparams
+def transformer_base_vq1_16_nb1_packed_dan_b01_scales():
+  """Set of hyperparameters."""
+  hparams = transformer_base_vq_ada_32ex_packed()
+  hparams.use_scales = int(True)
+  hparams.moe_num_experts = 16
+  hparams.moe_k = 1
+  hparams.beta = 0.1
+  hparams.ema = False
+  return hparams
+
+
 @registry.register_hparams
 def transformer_base_vq1_16_nb1_packed_nda_b01_scales_dialog():
   """Set of hyperparameters."""
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index dbdd30290..489ddaf3b 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -392,7 +392,7 @@ def update_hparams_for_vq_gating(hparams):
   hparams.add_hparam("beta", 0.25)
   hparams.add_hparam("epsilon", 1e-5)
   hparams.add_hparam("decay", 0.999)
-  hparams.add_hparam("ema", True)
+  hparams.add_hparam("ema", False)  # default is false until ema is implemented
   hparams.add_hparam("random_top_k", 1)
   hparams.add_hparam("soft_em", False)
   hparams.add_hparam("num_samples", 10)

From 722eb6ec09a765f82d04c88a32dd91f74d599e27 Mon Sep 17 00:00:00 2001
From: Youlong Cheng <ylc@google.com>
Date: Tue, 28 Aug 2018 20:40:03 -0700
Subject: [PATCH 0697/2720] [TF:XLA] Change group_assignment from 1d array
 attribute to 2d array input tensor with shape [num_groups,
 num_replica_per_group].

PiperOrigin-RevId: 210656091
---
 tensor2tensor/mesh_tensorflow/simd_mesh_impl.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
index 8b0b08565..db9317b9d 100644
--- a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
@@ -178,11 +178,18 @@ def allreduce(self, x, mesh_axes, reduction_fn_string):
       return x
     x = x.to_laid_out_tensor()
     if reduction_fn_string == "SUM":
-      partitioning = [
-          mtf.pnum_to_group(self.shape, mesh_axes, pnum)
-          for pnum in xrange(self.size)]
+      partitioning = {}
+      for pnum in xrange(self.size):
+        group = mtf.pnum_to_group(self.shape, mesh_axes, pnum)
+        if group not in partitioning:
+          partitioning[group] = []
+        partitioning[group].append(pnum)
+      group_assignment = []
+      for group, pnums in partitioning.items():
+        group_assignment.append(pnums)
+
       return self.LaidOutTensor(
-          [tpu_ops.cross_replica_sum(x.one_slice, partitioning)])
+          [tpu_ops.cross_replica_sum(x.one_slice, group_assignment)])
     else:
       for axis in mesh_axes:
         x = self.allconcat(x, axis, 0, stack=True)

From 04d5805973e2101d059d0f0141ca940a7ae957e2 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 29 Aug 2018 09:23:54 -0700
Subject: [PATCH 0698/2720] fixing video models infer/decode

PiperOrigin-RevId: 210733500
---
 tensor2tensor/models/video/base_test.py | 129 ++++++++++++++++++++----
 tensor2tensor/models/video/sv2p.py      |   6 +-
 2 files changed, 115 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/models/video/base_test.py b/tensor2tensor/models/video/base_test.py
index 0a981de3b..e7e467974 100644
--- a/tensor2tensor/models/video/base_test.py
+++ b/tensor2tensor/models/video/base_test.py
@@ -69,6 +69,7 @@ def full_modalities(hparams):
       "target_reward": ("symbol:one_hot", 3),
       "target_action": ("symbol:one_hot", 5),
   }
+  hparams.force_full_predict = True
   return hparams
 
 
@@ -108,13 +109,20 @@ class NextFrameTest(tf.test.TestCase):
 
   def RunModel(self, model, hparams, features):
     with tf.Session() as session:
-      model = model(
-          hparams, tf.estimator.ModeKeys.TRAIN)
+      model = model(hparams, tf.estimator.ModeKeys.TRAIN)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
     return res
 
+  def InferModel(self, model, hparams, features):
+    with tf.Session() as session:
+      model = model(hparams, tf.estimator.ModeKeys.PREDICT)
+      output = model.infer(features)
+      session.run(tf.global_variables_initializer())
+      res = session.run(output)
+    return res
+
   def TestVideoModel(self,
                      in_frames,
                      out_frames,
@@ -132,6 +140,27 @@ def TestVideoModel(self,
     expected_shape = get_tensor_shape(targets) + (expected_last_dim,)
     self.assertEqual(output.shape, expected_shape)
 
+  def TestVideoModelInfer(self,
+                          in_frames,
+                          out_frames,
+                          hparams,
+                          model,
+                          expected_last_dim,
+                          upsample_method="conv2d_transpose"):
+    del expected_last_dim
+    hparams = fill_hparams(hparams, in_frames, out_frames)
+    hparams.upsample_method = upsample_method
+
+    features = create_basic_features(in_frames, out_frames)
+    output = self.InferModel(model, hparams, features)
+
+    self.assertTrue(isinstance(output, dict))
+    self.assertTrue("outputs" in output.keys())
+    self.assertTrue("scores" in output.keys())
+    self.assertTrue("targets" in output.keys())
+    expected_shape = get_tensor_shape(features["targets"])
+    self.assertEqual(output["targets"].shape, expected_shape)
+
   def TestVideoModelWithActions(self,
                                 in_frames,
                                 out_frames,
@@ -149,6 +178,27 @@ def TestVideoModelWithActions(self,
     expected_shape = get_tensor_shape(targets) + (expected_last_dim,)
     self.assertEqual(output.shape, expected_shape)
 
+  def TestVideoModelWithActionsInfer(self,
+                                     in_frames,
+                                     out_frames,
+                                     hparams,
+                                     model,
+                                     expected_last_dim):
+    del expected_last_dim
+    hparams = fill_hparams(hparams, in_frames, out_frames)
+    hparams = action_modalities(hparams)
+    hparams.reward_prediction = False
+
+    features = create_action_features(in_frames, out_frames)
+    output = self.InferModel(model, hparams, features)
+
+    self.assertTrue(isinstance(output, dict))
+    self.assertTrue("outputs" in output.keys())
+    self.assertTrue("scores" in output.keys())
+    self.assertTrue("targets" in output.keys())
+    expected_shape = get_tensor_shape(features["targets"])
+    self.assertEqual(output["targets"].shape, expected_shape)
+
   def TestVideoModelWithActionAndRewards(self,
                                          in_frames,
                                          out_frames,
@@ -171,25 +221,57 @@ def TestVideoModelWithActionAndRewards(self,
     expected_shape = get_tensor_shape(targets)[:2] + (3,)
     self.assertEqual(output.shape, expected_shape)
 
-  def TestOnVariousInputOutputSizes(self, hparams, model, expected_last_dim):
-    self.TestVideoModel(1, 1, hparams, model, expected_last_dim)
-    self.TestVideoModel(1, 6, hparams, model, expected_last_dim)
-    self.TestVideoModel(4, 1, hparams, model, expected_last_dim)
-    self.TestVideoModel(7, 5, hparams, model, expected_last_dim)
+  def TestVideoModelWithActionAndRewardsInfer(self,
+                                              in_frames,
+                                              out_frames,
+                                              hparams,
+                                              model,
+                                              expected_last_dim):
+    del expected_last_dim
+    hparams = fill_hparams(hparams, in_frames, out_frames)
+    hparams = full_modalities(hparams)
+    hparams.reward_prediction = True
+
+    features = create_full_features(in_frames, out_frames)
+
+    output = self.InferModel(model, hparams, features)
+
+    self.assertTrue(isinstance(output, dict))
+    self.assertTrue("outputs" in output.keys())
+    self.assertTrue("scores" in output.keys())
+    self.assertTrue("targets" in output.keys())
+    self.assertTrue("target_reward" in output.keys())
+    expected_shape = get_tensor_shape(features["targets"])
+    self.assertEqual(output["targets"].shape, expected_shape)
+    expected_shape = get_tensor_shape(features["target_reward"])[:2] + (1,)
+    self.assertEqual(output["target_reward"].shape, expected_shape)
+
+  def TestOnVariousInputOutputSizes(
+      self, hparams, model, expected_last_dim, test_infer=True):
+    test_funcs = [self.TestVideoModel]
+    if test_infer:
+      test_funcs += [self.TestVideoModelInfer]
+    for test_func in test_funcs:
+      test_func(1, 1, hparams, model, expected_last_dim)
+      test_func(1, 6, hparams, model, expected_last_dim)
+      test_func(4, 1, hparams, model, expected_last_dim)
+      test_func(7, 5, hparams, model, expected_last_dim)
 
   def TestWithActions(self, hparams, model, expected_last_dim):
-    test_func = self.TestVideoModelWithActionAndRewards
-    test_func(1, 1, hparams, model, expected_last_dim)
-    test_func(1, 6, hparams, model, expected_last_dim)
-    test_func(4, 1, hparams, model, expected_last_dim)
-    test_func(7, 5, hparams, model, expected_last_dim)
+    for test_func in [self.TestVideoModelWithActions,
+                      self.TestVideoModelWithActionsInfer]:
+      test_func(1, 1, hparams, model, expected_last_dim)
+      test_func(1, 6, hparams, model, expected_last_dim)
+      test_func(4, 1, hparams, model, expected_last_dim)
+      test_func(7, 5, hparams, model, expected_last_dim)
 
   def TestWithActionAndRewards(self, hparams, model, expected_last_dim):
-    test_func = self.TestVideoModelWithActionAndRewards
-    test_func(1, 1, hparams, model, expected_last_dim)
-    test_func(1, 6, hparams, model, expected_last_dim)
-    test_func(4, 1, hparams, model, expected_last_dim)
-    test_func(7, 5, hparams, model, expected_last_dim)
+    for test_func in [self.TestVideoModelWithActionAndRewards,
+                      self.TestVideoModelWithActionAndRewardsInfer]:
+      test_func(1, 1, hparams, model, expected_last_dim)
+      test_func(1, 6, hparams, model, expected_last_dim)
+      test_func(4, 1, hparams, model, expected_last_dim)
+      test_func(7, 5, hparams, model, expected_last_dim)
 
   def TestOnVariousUpSampleLayers(self, hparams, model, expected_last_dim):
     self.TestVideoModel(4, 1, hparams, model, expected_last_dim,
@@ -200,13 +282,16 @@ def TestOnVariousUpSampleLayers(self, hparams, model, expected_last_dim):
   def testBasicDeterministic(self):
     self.TestOnVariousInputOutputSizes(
         basic_deterministic_params.next_frame_basic_deterministic(),
-        basic_deterministic.NextFrameBasicDeterministic, 256)
+        basic_deterministic.NextFrameBasicDeterministic,
+        256,
+        False)
 
   def testBasicStochastic(self):
     self.TestOnVariousInputOutputSizes(
         basic_stochastic.next_frame_basic_stochastic(),
         basic_stochastic.NextFrameBasicStochastic,
-        256)
+        256,
+        False)
 
   def testSv2p(self):
     self.TestOnVariousInputOutputSizes(
@@ -214,6 +299,12 @@ def testSv2p(self):
         sv2p.NextFrameSv2p,
         1)
 
+  def testSv2pWithActions(self):
+    self.TestWithActions(
+        sv2p_params.next_frame_sv2p(),
+        sv2p.NextFrameSv2p,
+        1)
+
   def testSv2pWithActionsAndRewards(self):
     self.TestWithActionAndRewards(
         sv2p_params.next_frame_sv2p(),
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 9d7ba1f1f..4f36b296b 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -493,8 +493,12 @@ def infer(self, features, *args, **kwargs):
 
     output, _ = self(features)  # pylint: disable=not-callable
 
+    if not isinstance(output, dict):
+      output = {"targets": output}
+
     output["targets"] = tf.squeeze(output["targets"], axis=-1)
-    output["target_reward"] = tf.argmax(output["target_reward"], axis=-1)
+    if self.hparams.reward_prediction:
+      output["target_reward"] = tf.argmax(output["target_reward"], axis=-1)
 
     # only required for decoding.
     output["outputs"] = output["targets"]

From 1daf42b1bca2a59c4ada5435670e25c3913ee9f1 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 29 Aug 2018 13:24:25 -0700
Subject: [PATCH 0699/2720] Improve docstrings for T2TModel.

PiperOrigin-RevId: 210778457
---
 tensor2tensor/utils/t2t_model.py | 48 +++++++++++++++++++++++---------
 1 file changed, 35 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index de192b6a2..ab4bcab5b 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -57,7 +57,28 @@
 class T2TModel(base.Layer):
   """Abstract base class for models.
 
-  Subclassess generally only need to override `body`.
+  `T2TModel` has three typical usages:
+
+  1. Estimator: The method `make_estimator_model_fn` builds a `model_fn` for
+     the tf.Estimator workflow of training, evaluation, and prediction.
+     Its core computation comes from the method `call`, which proceeds to call
+     the following methods:
+
+     * `bottom`, which transforms features according to `problem_hparams`' input
+       and target `Modality`s;
+     * `body`, which takes features and performs the core model computation to
+        return output and any auxiliary loss terms;
+     * `top`, which takes features and the body output, and transforms them
+       according to `problem_hparams`' input and target `Modality`s to return
+       the final logits;
+     * `loss`, which takes the logits, forms any missing training loss, and sums
+       all loss terms.
+  2. Layer: The method `call` enables `T2TModel` to be used a callable by
+     itself. For example, it can be composed with any other Keras layer.
+  3. Inference: The method `infer` enables `T2TModel` to make sequence
+     predictions by itself.
+
+  Subclasses generally only need to override `body`.
   """
   REGISTERED_NAME = None  # Updated on registration.
 
@@ -68,7 +89,7 @@ def __init__(self,
                data_parallelism=None,
                decode_hparams=None,
                **kwargs):
-    """Create a T2TModel.
+    """Creates a T2TModel.
 
     Args:
       hparams: tf.contrib.training.HParams, model hyperparameters.
@@ -82,9 +103,6 @@ def __init__(self,
       decode_hparams: a hyperparameter object with decoding parameters.
         See decoding.decode_hparams.
       **kwargs: arguments to pass to base.Layer constructor.
-
-    Returns:
-      a T2TModel
     """
     # Determine name first: use registered name if possible, class name else.
     default_name = registry.default_name(type(self))
@@ -353,19 +371,22 @@ def bottom(self, features):
     return transformed_features
 
   def body(self, features):
-    """Most models will override this function.
+    """Computes the targets' logits for one shard given transformed inputs.
 
-    Compute label logits for one shard as a function of the transformed
-    features.
+    Most `T2TModel` subclasses will override this method.
 
     Args:
-      features: A dictionary of key to Tensor.  Each Tensor has shape
-         [batch_size, ?, ?, hidden_size].
+      features: dict of str to Tensor, where each Tensor has shape [batch_size,
+        ..., hidden_size]. It typically contains keys `inputs` and `targets`.
 
     Returns:
-      output: tensor of logits with shape [batch_size, O, P, body_output_size.
-      losses: either single loss as a scalar, a list, a tensor (to be averaged)
-              or a dictionary of losses.
+      output: Tensor of pre-logit activations with shape [batch_size, ...,
+              hidden_size].
+      losses: Either single loss as a scalar, a list, a Tensor (to be averaged),
+              or a dictionary of losses. If losses is a dictionary with the key
+              "training", losses["training"] is considered the final training
+              loss and output is considered logits; self.top and self.loss will
+              be skipped.
     """
     raise NotImplementedError("Abstract Method")
 
@@ -406,6 +427,7 @@ def _top_single(self, body_output, target_modality, features):
     return logits
 
   def top(self, body_output, features):
+    """Returns `logits` given body output and features."""
     if isinstance(body_output, dict):
       if self._problem_hparams:
         target_modality = self._problem_hparams.target_modality

From 1e7af9bd1a9bd864fa3fbe7fc2b39302803db7f3 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 29 Aug 2018 14:57:19 -0700
Subject: [PATCH 0700/2720] Remove unused hparams.sep_rgb_embed.

PiperOrigin-RevId: 210796475
---
 tensor2tensor/models/image_transformer.py        | 2 --
 tensor2tensor/models/research/transformer_vae.py | 1 -
 2 files changed, 3 deletions(-)

diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index af34960e9..d4145d04f 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -221,7 +221,6 @@ def image_transformer_base():
   hparams.add_hparam("block_width", 128)
   hparams.add_hparam("num_encoder_layers", 4)
   hparams.add_hparam("num_decoder_layers", 12)
-  hparams.sep_rgb_embed = False
   hparams.add_hparam("dec_attention_type", cia.AttentionType.LOCAL_1D)
   hparams.add_hparam("block_raster_scan", False)
 
@@ -789,7 +788,6 @@ def imagetransformer_sep_channels_16l_16h_imgnet_lrg_loc_128():
 def imagetransformer_sep_output_channels_8l_local_and_global_att():
   """separate rgb embeddings."""
   hparams = imagetransformer_sep_channels_8l()
-  hparams.sep_rgb_embed = True
   hparams.sampling_method = "random"
   hparams.local_and_global_att = True
   return hparams
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index c1835ec38..f2a4a7dfa 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -829,7 +829,6 @@ def imagetransformer_ae_cifar():
   hparams.add_hparam("block_width", 128)
   hparams.num_encoder_layers = 4
   hparams.num_decoder_layers = 12
-  hparams.sep_rgb_embed = False
   hparams.add_hparam("dec_attention_type", cia.AttentionType.LOCAL_1D)
   hparams.add_hparam("block_raster_scan", False)
   hparams.add_hparam("shared_rel", False)

From 5df7a2c4c14944d44d4a58eefc783d7e372c3817 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 29 Aug 2018 18:17:07 -0700
Subject: [PATCH 0701/2720] adding sv2p config for RL.

PiperOrigin-RevId: 210829422
---
 tensor2tensor/rl/trainer_model_based.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 27a0959d7..36731645f 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -553,6 +553,15 @@ def rl_modelrl_base_stochastic():
   return hparams
 
 
+@registry.register_hparams
+def rl_modelrl_base_sv2p():
+  """Base setting with sv2p as world model."""
+  hparams = rl_modelrl_base()
+  hparams.generative_model = "next_frame_sv2p"
+  hparams.generative_model_params = "next_frame_sv2p_atari"
+  return hparams
+
+
 @registry.register_hparams
 def rl_modelrl_medium():
   """Small set for larger testing."""

From ea9934874031a66dc5bd74def5d87fa377131ac8 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 29 Aug 2018 18:45:58 -0700
Subject: [PATCH 0702/2720] Allow scheduled sampling in basic video model,
 simplify default video modality.

PiperOrigin-RevId: 210832375
---
 tensor2tensor/layers/common_layers.py         | 10 ++-
 tensor2tensor/layers/modalities.py            | 41 +++++-----
 tensor2tensor/models/research/autoencoders.py |  6 +-
 .../models/video/basic_deterministic.py       | 76 ++++++++++++++++++-
 .../video/basic_deterministic_params.py       | 10 +++
 tensor2tensor/rl/envs/simulated_batch_env.py  |  4 +
 tensor2tensor/rl/trainer_model_based.py       | 19 ++++-
 tensor2tensor/utils/metrics.py                |  8 +-
 8 files changed, 144 insertions(+), 30 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 9720a7a07..dc2e4d3f4 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -133,14 +133,20 @@ def inverse_exp_decay(max_step, min_value=0.01, step=None):
   """Inverse-decay exponentially from 0.01 to 1.0 reached at max_step."""
   inv_base = tf.exp(tf.log(min_value) / float(max_step))
   if step is None:
-    step = tf.to_float(tf.train.get_global_step())
+    step = tf.train.get_global_step()
+  if step is None:
+    return 1.0
+  step = tf.to_float(step)
   return inv_base**tf.maximum(float(max_step) - step, 0.0)
 
 
 def inverse_lin_decay(max_step, min_value=0.01, step=None):
   """Inverse-decay linearly from 0.01 to 1.0 reached at max_step."""
   if step is None:
-    step = tf.to_float(tf.train.get_global_step())
+    step = tf.train.get_global_step()
+  if step is None:
+    return 1.0
+  step = tf.to_float(step)
   progress = tf.minimum(step / float(max_step), 1.0)
   return progress * (1.0 - min_value) + min_value
 
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index f2a9d3b21..68e6d8d3d 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -511,14 +511,13 @@ def xnet_resblock(x, filters, res_relu, name):
                            "compress_block_final")
 
 
-@registry.register_video_modality("default")
 class VideoModality(modality.Modality):
   """Modality for videos, i.e., time-sequences of frames."""
   PIXEL_EMBEDDING_SIZE = 64
 
   def bottom(self, x):
     inputs = x
-    with tf.variable_scope(self.name):
+    with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
       common_layers.summarize_video(inputs, "inputs")
       inputs_shape = common_layers.shape_list(inputs)
       # Standardize frames.
@@ -553,24 +552,20 @@ def targets_bottom(self, x, summary_prefix="targets_bottom"):  # pylint: disable
           self._body_input_depth,
           name="merge_pixel_embedded_frames")
 
-  def top(self, body_output, _):
+  def top(self, body_output, targets):
     num_channels = self._model_hparams.problem.num_channels
-    num_frames = self._model_hparams.video_num_target_frames
-    with tf.variable_scope("rgb_softmax"):
-      body_output_shape = common_layers.shape_list(body_output)
-      reshape_shape = body_output_shape[:3]
-      reshape_shape.extend([num_channels, num_frames, self.top_dimensionality])
-      res = tf.layers.dense(body_output,
-                            self.top_dimensionality * num_channels * num_frames)
-      res = tf.reshape(res, reshape_shape)
-      res = tf.transpose(res, [0, 4, 1, 2, 3, 5])
-      if not tf.get_variable_scope().reuse:
-        res_argmax = tf.argmax(res[:, -1, :, :, :, :], axis=-1)
-        tf.summary.image(
-            "result",
-            common_layers.tpu_safe_image_summary(res_argmax),
-            max_outputs=1)
-      return res
+    num_frames = common_layers.shape_list(targets)[1]
+    body_output_shape = common_layers.shape_list(body_output)
+    # We assume the body output is of this shape and layout.
+    reshape_shape = body_output_shape[:-1] + [
+        num_channels, self.top_dimensionality, num_frames]
+    res = tf.reshape(body_output, reshape_shape)
+    res = tf.transpose(res, [0, 5, 1, 2, 3, 4])
+    res_shape = common_layers.shape_list(res)
+    res_argmax = tf.argmax(tf.reshape(res, [-1, res_shape[-1]]), axis=-1)
+    res_argmax = tf.reshape(res_argmax, res_shape[:-1])
+    common_layers.summarize_video(res_argmax, "result")
+    return res
 
   def loss(self, top_out, targets):
     """Compute loss numerator and denominator for one shard of output."""
@@ -586,6 +581,14 @@ def loss(self, top_out, targets):
         weights_fn=self.targets_weights_fn)
 
 
+@registry.register_video_modality("default")
+class VideoModalityNoEmbed(VideoModality):
+  """Video Modality where target_bottom does not embeds pixels."""
+
+  def targets_bottom(self, x):
+    return super(VideoModalityNoEmbed, self).bottom(x)
+
+
 @registry.register_video_modality("embed")
 class VideoModalityEmbed(VideoModality):
   """Video Modality where bottom embeds pixels."""
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index dedaeac7b..6ff3fa17e 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1062,13 +1062,16 @@ def autoencoder_ordered_discrete_hs256():
 def autoencoder_ordered_text():
   """Ordered discrete autoencoder model for text."""
   hparams = autoencoder_ordered_discrete()
-  hparams.bottleneck_bits = 2048
+  hparams.bottleneck_bits = 1024
+  hparams.unordered = True
+  hparams.bottleneck_noise = 0.05
   hparams.num_hidden_layers = 7
   hparams.batch_size = 1024
   hparams.autoregressive_mode = "conv5"
   hparams.max_hidden_size = 1024
   hparams.target_modality = "symbol:identity"
   hparams.input_modalities = "symbol:identity"
+  hparams.sample_height = 128
   hparams.sample_width = 1
   return hparams
 
@@ -1083,6 +1086,7 @@ def autoencoder_ordered_text_small():
   hparams.max_hidden_size = 512
   hparams.bottleneck_noise = 0.0
   hparams.autoregressive_mode = "conv5"
+  hparams.sample_height = 4
   return hparams
 
 
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 91c00daf7..c0691f73c 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -42,7 +42,7 @@ def inject_latent(self, layer, features, filters):
     del features, filters
     return layer, 0.0
 
-  def body(self, features):
+  def body_single(self, features):
     hparams = self.hparams
     filters = hparams.hidden_size
     kernel1, kernel2 = (3, 3), (4, 4)
@@ -109,13 +109,84 @@ def body(self, features):
 
     # Cut down to original size.
     x = x[:, :inputs_shape[1], :inputs_shape[2], :]
+    x = tf.layers.dense(x, hparams.problem.num_channels * 256, name="logits")
 
     # Reward prediction if needed.
     if "target_reward" not in features:
       return x
-    reward_pred = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
+    reward_pred = tf.expand_dims(  # Add a fake channels dim.
+        tf.reduce_mean(x, axis=[1, 2], keepdims=True), axis=3)
     return {"targets": x, "target_reward": reward_pred}, extra_loss
 
+  def body(self, features):
+    hparams = self.hparams
+    is_predicting = hparams.mode == tf.estimator.ModeKeys.PREDICT
+    if hparams.video_num_target_frames < 2:
+      res = self.body_single(features)
+      return res
+
+    # TODO(lukaszkaiser): the split axes and the argmax below heavily depend on
+    # using the default (a bit strange) video modality - we should change that.
+
+    # Split inputs and targets into lists.
+    input_frames = list(tf.split(
+        features["inputs"], hparams.video_num_input_frames, axis=-1))
+    target_frames = list(tf.split(
+        features["targets"], hparams.video_num_target_frames, axis=-1))
+    all_frames = input_frames + target_frames
+    if "input_action" in features:
+      input_actions = list(tf.split(
+          features["input_action"], hparams.video_num_input_frames, axis=1))
+      target_actions = list(tf.split(
+          features["target_action"], hparams.video_num_target_frames, axis=1))
+      all_actions = input_actions + target_actions
+
+    # Run a number of steps.
+    res_frames = []
+    if "target_reward" in features:
+      res_rewards, extra_loss = [], 0.0
+    sample_prob = common_layers.inverse_exp_decay(
+        hparams.scheduled_sampling_warmup_steps)
+    sample_prob *= hparams.scheduled_sampling_prob
+    for i in range(hparams.video_num_target_frames):
+      cur_frames = all_frames[i:i + hparams.video_num_input_frames]
+      features["inputs"] = tf.concat(cur_frames, axis=-1)
+      if "input_action" in features:
+        cur_actions = all_actions[i:i + hparams.video_num_input_frames]
+        features["input_action"] = tf.concat(cur_actions, axis=1)
+
+      # Run model.
+      with tf.variable_scope(tf.get_variable_scope(), reuse=i > 0):
+        if "target_reward" not in features:
+          res_frames.append(self.body_single(features))
+        else:
+          res_dict, res_extra_loss = self.body_single(features)
+          extra_loss += res_extra_loss
+          res_frames.append(res_dict["targets"])
+          res_rewards.append(res_dict["target_reward"])
+
+      # When predicting, use the generated frame.
+      orig_frame = all_frames[i + hparams.video_num_input_frames]
+      shape = common_layers.shape_list(orig_frame)
+      sampled_frame = tf.reshape(
+          res_frames[-1], shape[:-1] + [hparams.problem.num_channels, 256])
+      sampled_frame = tf.to_float(tf.argmax(sampled_frame, axis=-1))
+      if is_predicting:
+        all_frames[i + hparams.video_num_input_frames] = sampled_frame
+
+      # Scheduled sampling during training.
+      if (hparams.scheduled_sampling_prob > 0.0 and self.is_training):
+        do_sample = tf.less(tf.random_uniform([shape[0]]), sample_prob)
+        sampled_frame = tf.where(do_sample, sampled_frame, orig_frame)
+        all_frames[i + hparams.video_num_input_frames] = sampled_frame
+
+    # Concatenate results and return them.
+    frames = tf.concat(res_frames, axis=-1)
+    if "target_reward" not in features:
+      return frames
+    rewards = tf.concat(res_rewards, axis=1)
+    return {"targets": frames, "target_reward": rewards}, extra_loss
+
   def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
     """Produce predictions from the model by running it."""
     del args, kwargs
@@ -150,6 +221,7 @@ def logits_to_samples(logits):
       tf.logging.warn("Guessing targets shape as no inputs are given.")
       targets_shape = [self.hparams.batch_size,
                        self.hparams.video_num_target_frames, 1, 1, num_channels]
+
     features["targets"] = tf.zeros(targets_shape, dtype=tf.int32)
     if "target_reward" in self.hparams.problem_hparams.target_modality:
       features["target_reward"] = tf.zeros(
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 2fd5a6958..6d234adab 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -60,6 +60,16 @@ def next_frame_pixel_noise():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_sampling():
+  """Basic conv model with scheduled sampling."""
+  hparams = next_frame_basic_deterministic()
+  hparams.video_num_target_frames = 2
+  hparams.scheduled_sampling_warmup_steps = 30000
+  hparams.scheduled_sampling_prob = 0.1
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_tpu():
   hparams = next_frame_basic_deterministic()
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 51bcaceed..99f2d0f3a 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -153,8 +153,12 @@ def simulate(self, action):
                           axis=1)
       history = self.history_buffer.get_all_elements()
       with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+        # We only need 1 target frame here, set it.
+        hparams_target_frames = self._model.hparams.video_num_target_frames
+        self._model.hparams.video_num_target_frames = 1
         model_output = self._model.infer(
             {"inputs": history, "input_action": actions})
+        self._model.hparams.video_num_target_frames = hparams_target_frames
 
       observ = tf.to_float(tf.squeeze(model_output["targets"], axis=1))
 
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 36731645f..242779bb4 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -512,7 +512,7 @@ def combine_training_data(problem, final_data_dir, old_data_dirs,
 @registry.register_hparams
 def rl_modelrl_base():
   return tf.contrib.training.HParams(
-      epochs=6,
+      epochs=2,
       # Total frames used for training. This will be distributed evenly across
       # hparams.epochs.
       num_real_env_frames=100000,
@@ -524,7 +524,7 @@ def rl_modelrl_base():
       simulated_env_generator_num_steps=2000,
       simulation_random_starts=True,
       intrinsic_reward_scale=0.,
-      ppo_epochs_num=2000,  # This should be enough to see something
+      ppo_epochs_num=1000,  # This should be enough to see something
       # Our simulated envs do not know how to reset.
       # You should set ppo_time_limit to the value you believe that
       # the simulated env produces a reasonable output.
@@ -562,6 +562,14 @@ def rl_modelrl_base_sv2p():
   return hparams
 
 
+@registry.register_hparams
+def rl_modelrl_base_sampling():
+  """Base setting with a stochastic next-frame model."""
+  hparams = rl_modelrl_base()
+  hparams.generative_model_params = "next_frame_sampling"
+  return hparams
+
+
 @registry.register_hparams
 def rl_modelrl_medium():
   """Small set for larger testing."""
@@ -599,7 +607,7 @@ def rl_modelrl_model_only():
 @registry.register_hparams
 def rl_modelrl_tiny():
   """Tiny set for testing."""
-  return rl_modelrl_base().override_from_dict(
+  return rl_modelrl_base_sampling().override_from_dict(
       tf.contrib.training.HParams(
           epochs=2,
           num_real_env_frames=128,
@@ -800,6 +808,11 @@ def rl_modelrl_variance_nogame(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(500)))
 
 
+@registry.register_ranged_hparams
+def rl_modelrl_scheduled_sampling(rhp):
+  rhp.set_float("model.scheduled_sampling_prob", 0.0, 1.0)
+
+
 @registry.register_ranged_hparams
 def rl_modelrl_all_games(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index f5ab0033e..505af9611 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -525,11 +525,13 @@ def reduce_dimensions(predictions, labels):
     """Reduce dimensions for high-dimensional predictions and labels."""
     # We will treat first dimensions as batch. One example are video frames.
     if len(predictions.get_shape()) > 5:
+      predictions_shape = common_layers.shape_list(predictions)
       predictions = tf.reshape(
-          predictions, [-1] + common_layers.shape_list(predictions)[-4:])
-    if len(labels.get_shape()) > 4:
+          predictions, [predictions_shape[0], predictions_shape[1], -1,
+                        predictions_shape[-1]])
+      labels_shape = common_layers.shape_list(labels)
       labels = tf.reshape(
-          labels, [-1] + common_layers.shape_list(labels)[-3:])
+          labels, [labels_shape[0], labels_shape[1], -1])
     return predictions, labels
 
   def make_problem_specific_metric_fn(metric_fn, weights_fn):

From 17345a809e3d6219d349f8a2ef3d169d5e324723 Mon Sep 17 00:00:00 2001
From: cbockman <c.bockman@gmail.com>
Date: Thu, 30 Aug 2018 12:52:51 -0700
Subject: [PATCH 0703/2720] Fix to universal transformer update weights (#1011)

* Fix to universal transformer update weights

This was identified (not by me) in https://github.com/tensorflow/tensor2tensor/issues/1004.

Putting in a PR to ensure this is fixed.

If this is not a correct fix, then the underlying code needs to be rejiggered for readability...

* right paren

* one more case
---
 tensor2tensor/models/research/universal_transformer_util.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 4c2e76fb0..87c20e8cd 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -1204,7 +1204,7 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
 
     # update running part in the weighted state and keep the rest
     new_state = ((transformed_state * update_weights) +
-                 (previous_state * 1 - update_weights))
+                 (previous_state * (1 - update_weights)))
 
     # remind TensorFlow of everything's shape
     transformed_state.set_shape(state_shape)
@@ -1495,7 +1495,7 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
 
     # Add in the weighted state
     new_state = ((transformed_state * update_weights) +
-                 (previous_state * 1 - update_weights))
+                 (previous_state * (1 - update_weights)))
 
     # Remind TensorFlow of everything's shape
     state.set_shape(state_shape)
@@ -1642,7 +1642,7 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
 
     # update running part in the weighted state and keep the rest
     new_state = ((transformed_state * update_weights) +
-                 (previous_state * 1 - update_weights))
+                 (previous_state * (1 - update_weights)))
 
     # remind TensorFlow of everything's shape
     transformed_state.set_shape(state_shape)

From 1d87d013e1d1204995eec65af188e678d4cacf41 Mon Sep 17 00:00:00 2001
From: Tomasz Latkowski <13836101+tlatkowski@users.noreply.github.com>
Date: Thu, 30 Aug 2018 23:58:07 +0200
Subject: [PATCH 0704/2720] bug fix for python 3 in Common Voice problem
 (#1028)

---
 tensor2tensor/data_generators/common_voice.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/common_voice.py b/tensor2tensor/data_generators/common_voice.py
index 4bad08fdb..79934a4c6 100644
--- a/tensor2tensor/data_generators/common_voice.py
+++ b/tensor2tensor/data_generators/common_voice.py
@@ -20,6 +20,7 @@
 package installed. The original samples will be downsampled by the encoder.
 """
 
+import six
 import csv
 import os
 import tarfile
@@ -56,7 +57,11 @@ def _collect_data(directory):
     transcript_path = os.path.join(directory, transcript)
     with open(transcript_path, "r") as transcript_file:
       transcript_reader = csv.reader(transcript_file)
-      _ = transcript_reader.next()  # Skip headers.
+      # skip header
+      if six.PY3:
+        _ = next(transcript_reader)
+      else:
+        _ = transcript_reader.next()
       for transcript_line in transcript_reader:
         media_name, label = transcript_line[0:2]
         filename = os.path.join(directory, media_name)

From 88861622ef67af35f9f496b59d6f9f65fc63c861 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 30 Aug 2018 15:11:31 -0700
Subject: [PATCH 0705/2720] internal merge of PR #1011

PiperOrigin-RevId: 210989025
---
 tensor2tensor/data_generators/common_voice.py | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/common_voice.py b/tensor2tensor/data_generators/common_voice.py
index 79934a4c6..4bad08fdb 100644
--- a/tensor2tensor/data_generators/common_voice.py
+++ b/tensor2tensor/data_generators/common_voice.py
@@ -20,7 +20,6 @@
 package installed. The original samples will be downsampled by the encoder.
 """
 
-import six
 import csv
 import os
 import tarfile
@@ -57,11 +56,7 @@ def _collect_data(directory):
     transcript_path = os.path.join(directory, transcript)
     with open(transcript_path, "r") as transcript_file:
       transcript_reader = csv.reader(transcript_file)
-      # skip header
-      if six.PY3:
-        _ = next(transcript_reader)
-      else:
-        _ = transcript_reader.next()
+      _ = transcript_reader.next()  # Skip headers.
       for transcript_line in transcript_reader:
         media_name, label = transcript_line[0:2]
         filename = os.path.join(directory, media_name)

From c15c287878f23b9ffa4dd2aed96892c21a886d4f Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 30 Aug 2018 15:12:01 -0700
Subject: [PATCH 0706/2720] internal merge of PR #1028

PiperOrigin-RevId: 210989131
---
 tensor2tensor/data_generators/common_voice.py |  3 +-
 .../data_generators/common_voice_test.py      | 42 +++++++++++++++++++
 tensor2tensor/data_generators/test_data/1.csv |  2 +
 3 files changed, 46 insertions(+), 1 deletion(-)
 create mode 100644 tensor2tensor/data_generators/common_voice_test.py
 create mode 100644 tensor2tensor/data_generators/test_data/1.csv

diff --git a/tensor2tensor/data_generators/common_voice.py b/tensor2tensor/data_generators/common_voice.py
index 4bad08fdb..fa9609aed 100644
--- a/tensor2tensor/data_generators/common_voice.py
+++ b/tensor2tensor/data_generators/common_voice.py
@@ -56,7 +56,8 @@ def _collect_data(directory):
     transcript_path = os.path.join(directory, transcript)
     with open(transcript_path, "r") as transcript_file:
       transcript_reader = csv.reader(transcript_file)
-      _ = transcript_reader.next()  # Skip headers.
+      # skip header
+      _ = next(transcript_reader)
       for transcript_line in transcript_reader:
         media_name, label = transcript_line[0:2]
         filename = os.path.join(directory, media_name)
diff --git a/tensor2tensor/data_generators/common_voice_test.py b/tensor2tensor/data_generators/common_voice_test.py
new file mode 100644
index 000000000..8da86a7e5
--- /dev/null
+++ b/tensor2tensor/data_generators/common_voice_test.py
@@ -0,0 +1,42 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.data_generators.common_voice."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+from tensor2tensor.data_generators import common_voice
+
+import tensorflow as tf
+
+pkg_dir, _ = os.path.split(__file__)
+_TESTDATA = os.path.join(pkg_dir, "test_data")
+
+
+class CommonVoiceTest(tf.test.TestCase):
+
+  def testCollectData(self):
+    output = common_voice._collect_data(_TESTDATA)
+    self.assertEqual(1, len(output))
+
+    # NOTE: No header.
+    self.assertTrue("my_media" == output[0][0])
+    self.assertTrue("my_label" == output[0][2])
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/data_generators/test_data/1.csv b/tensor2tensor/data_generators/test_data/1.csv
new file mode 100644
index 000000000..fcb33cc3c
--- /dev/null
+++ b/tensor2tensor/data_generators/test_data/1.csv
@@ -0,0 +1,2 @@
+media_name,label
+my_media,my_label

From 37625c7e24b965cf27b15f1c4f8dcd7b20d07405 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 30 Aug 2018 15:13:36 -0700
Subject: [PATCH 0707/2720] Clean up ImageChannelCompressModality.

PiperOrigin-RevId: 210989363
---
 tensor2tensor/layers/modalities.py | 60 +++++++++++++++++++-----------
 1 file changed, 38 insertions(+), 22 deletions(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 68e6d8d3d..c97df5ef5 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -303,18 +303,19 @@ def num_channels(self):
     return 3
 
   def bottom_compress(self, inputs, name="bottom"):
-    """Transform input from data space to model space.
+    """Compresses channel-wise input pixels into whole pixel representions.
 
-    Perform conversion of RGB pixel values to a real number in the range -1 to 1
-    and combine channel values for each pixel to form a representation of
-    size image_length x image_length dims.
+    Perform conversion of RGB pixel values to a real number in the range -1 to
+    1. This combines pixel channels to form a representation of shape
+    [img_len, img_len].
 
     Args:
-      inputs: A Tensor representing RGB pixel intensities as integers.
-        [batch, ...]
+      inputs: Tensor representing RGB pixel intensities as integers, of shape
+        [batch, img_len, img_len, channels].
       name: string, scope.
+
     Returns:
-      body_input: A Tensor with shape [batch, ?, ?, body_input_depth].
+      body_input: Tensor of shape [batch, img_len, img_len, body_input_depth].
     """
     with tf.variable_scope(name):
       inputs = tf.to_float(inputs)
@@ -325,19 +326,22 @@ def bottom_compress(self, inputs, name="bottom"):
             common_layers.tpu_safe_image_summary(inputs),
             max_outputs=2)
       inputs = common_layers.convert_rgb_to_symmetric_real(inputs)
-      ishape = common_layers.shape_list(inputs)
-      inputs = tf.reshape(inputs, [-1, ishape[1], ishape[2] * ishape[3], 1])
-      inputs.set_shape([None, None, None, 1])
-      # We compress RGB intensities for each pixel using a conv.
-      x = tf.layers.conv2d(
+
+      # Reshape inputs to apply convolutions across [img_len, img_len*channels].
+      inputs_shape = common_layers.shape_list(inputs)
+      inputs = tf.reshape(
+          inputs, [-1, inputs_shape[1], inputs_shape[2] * inputs_shape[3], 1])
+
+      # Compress RGB intensities for each pixel using a convolution.
+      outputs = tf.layers.conv2d(
           inputs,
-          self._body_input_depth, (1, self.num_channels),
+          self._body_input_depth,
+          kernel_size=(1, self.num_channels),
           padding="VALID",
           strides=(1, self.num_channels),
           activation=tf.nn.relu,
           name="conv_input")
-      x.set_shape([None, None, None, self._body_input_depth])
-      return x
+      return outputs
 
   def bottom(self, x):
     return self.bottom_compress(x, "input_bottom")
@@ -346,24 +350,36 @@ def targets_bottom(self, x):
     return self.bottom_compress(x, "output_bottom")
 
   def top(self, body_output, _):
+    """Transforms body output to return logits.
+
+    Args:
+      body_output: Tensor of shape [batch, img_len, img_len, depth].
+
+    Returns:
+      Tensor of shape [batch, img_len, img_len, channels, top_dimensionality].
+    """
     with tf.variable_scope(self.name):
-      hidden_dim = self._model_hparams.hidden_size
+      hidden_size = self._model_hparams.hidden_size
       img_len = self._model_hparams.img_len
       channels = self.num_channels  # RGB
       batch = common_layers.shape_list(body_output)[0]
       x = tf.layers.conv2d(
           body_output,
-          hidden_dim * channels, (1, 1),
+          hidden_size * channels,
+          kernel_size=(1, 1),
           strides=(1, 1),
           padding="VALID",
           activation=tf.nn.relu,
           name="decompress_conv")
-      x = tf.reshape(x, [batch, img_len, img_len * channels, hidden_dim])
+      x = tf.reshape(x, [batch, img_len, img_len * channels, hidden_size])
       x = common_layers.layer_preprocess(x, self._model_hparams)
-      x = tf.layers.dense(
-          x, 256, use_bias=True, activation=None, name="output_conv")
-      x = tf.reshape(x,
-                     [-1, img_len, img_len, channels, self.top_dimensionality])
+      x = tf.layers.dense(x,
+                          self.top_dimensionality,
+                          use_bias=True,
+                          activation=None,
+                          name="output_conv")
+      x = tf.reshape(
+          x, [batch, img_len, img_len, channels, self.top_dimensionality])
       return x
 
 

From d8ff992335c8e1966e00098ec58e28c6cc50368d Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 31 Aug 2018 00:21:02 +0200
Subject: [PATCH 0708/2720] Infer observation datatype and shape from the
 environment (#1025)

* Infer observation datatype and shape from the environment

* fix dtypes in Stack Wrappers

* Fix pylint errors
---
 tensor2tensor/rl/collect.py                  |  8 +--
 tensor2tensor/rl/envs/batch_env.py           | 16 ++---
 tensor2tensor/rl/envs/in_graph_batch_env.py  | 25 ++++++-
 tensor2tensor/rl/envs/py_func_batch_env.py   | 19 ++---
 tensor2tensor/rl/envs/simulated_batch_env.py | 34 +++++----
 tensor2tensor/rl/envs/tf_atari_wrappers.py   | 75 ++++++++++++--------
 tensor2tensor/rl/envs/utils.py               | 24 ++++---
 7 files changed, 125 insertions(+), 76 deletions(-)

diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index f32fca858..bde467969 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -33,10 +33,11 @@ def _rollout_metadata(batch_env):
   batch_size = [batch_env_shape[0]]
   shapes_types_names = [
       # TODO(piotrmilos): possibly retrieve the observation type for batch_env
-      (batch_size + batch_env_shape[1:], tf.float32, "observation"),
+      (batch_size + batch_env_shape[1:], batch_env.observ_dtype, "observation"),
       (batch_size, tf.float32, "reward"),
       (batch_size, tf.bool, "done"),
-      (batch_size + batch_env.action_shape, batch_env.action_dtype, "action"),
+      (batch_size + list(batch_env.action_shape), batch_env.action_dtype,
+       "action"),
       (batch_size, tf.float32, "pdf"),
       (batch_size, tf.float32, "value_function"),
   ]
@@ -57,8 +58,7 @@ def __init__(self, batch_env):
     self.speculum = tf.FIFOQueue(infinity, shapes=shapes, dtypes=dtypes)
     observs_shape = batch_env.observ.shape
     # TODO(piotrmilos): possibly retrieve the observation type for batch_env
-    observ_dtype = tf.float32
-    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
+    self._observ = tf.Variable(tf.zeros(observs_shape, self.observ_dtype),
                                trainable=False)
 
   def simulate(self, action):
diff --git a/tensor2tensor/rl/envs/batch_env.py b/tensor2tensor/rl/envs/batch_env.py
index bf433066f..230bd3a43 100644
--- a/tensor2tensor/rl/envs/batch_env.py
+++ b/tensor2tensor/rl/envs/batch_env.py
@@ -43,11 +43,12 @@ def __init__(self, envs, blocking):
     """
     self._envs = envs
     self._blocking = blocking
-    observ_space = self._envs[0].observation_space
-    if not all(env.observation_space == observ_space for env in self._envs):
+    self.observ_space = self._envs[0].observation_space
+    if not all(env.observation_space == self.observ_space
+               for env in self._envs):
       raise ValueError('All environments must use the same observation space.')
-    action_space = self._envs[0].action_space
-    if not all(env.action_space == action_space for env in self._envs):
+    self.action_space = self._envs[0].action_space
+    if not all(env.action_space == self.action_space for env in self._envs):
       raise ValueError('All environments must use the same observation space.')
 
   def __len__(self):
@@ -96,8 +97,8 @@ def step(self, actions):
       transitions = [transition() for transition in transitions]
     observs, rewards, dones, infos = zip(*transitions)
 
+    observ = np.stack(observs).astype(self.observ_space.dtype)
     # TODO(piotrmilos): Do we really want cast to float32
-    observ = np.stack(observs).astype(np.float32)
     reward = np.stack(rewards).astype(np.float32)
     done = np.stack(dones)
     info = tuple(infos)
@@ -119,9 +120,8 @@ def reset(self, indices=None):
     else:
       observs = [self._envs[index].reset(blocking=False) for index in indices]
       observs = [observ() for observ in observs]
-    observ = np.stack(observs)
-    # TODO(piotrmilos): Do we really want this?
-    observ = observ.astype(np.float32)
+    observ = np.stack(observs).astype(self.observ_space.dtype)
+
     return observ
 
   def close(self):
diff --git a/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index 56e13878d..8db9b56ef 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -23,11 +23,17 @@
 
 import tensorflow as tf
 
+from tensor2tensor.rl.envs import utils
+
 
 class InGraphBatchEnv(object):
   """Abstract class for batch of environments inside the TensorFlow graph.
   """
 
+  def __init__(self, observ_space, action_space):
+    self.observ_space = observ_space
+    self.action_space = action_space
+
   def __getattr__(self, name):
     """Forward unimplemented attributes to one of the original environments.
 
@@ -71,7 +77,24 @@ def reset(self, indices=None):
     """
     return tf.cond(
         tf.cast(tf.reduce_sum(indices + 1), tf.bool),
-        lambda: self._reset_non_empty(indices), lambda: 0.0)
+        lambda: self._reset_non_empty(indices),
+        lambda: tf.cast(0, self.observ_dtype))
+
+  @property
+  def observ_dtype(self):
+    return utils.parse_dtype(self.observ_space)
+
+  @property
+  def observ_shape(self):
+    return utils.parse_shape(self.observ_space)
+
+  @property
+  def action_dtype(self):
+    return utils.parse_dtype(self.action_space)
+
+  @property
+  def action_shape(self):
+    return utils.parse_shape(self.action_space)
 
   @property
   def observ(self):
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index 76916d991..ac1ad1dd5 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -21,7 +21,6 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.rl.envs import utils
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
 import tensorflow as tf
 
@@ -40,14 +39,13 @@ def __init__(self, batch_env):
     Args:
       batch_env: Batch environment.
     """
+    super(PyFuncBatchEnv, self).__init__(batch_env.observation_space,
+                                         batch_env.action_space)
     self._batch_env = batch_env
-    observ_shape = utils.parse_shape(self._batch_env.observation_space)
-    observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
-    self.action_shape = list(utils.parse_shape(self._batch_env.action_space))
-    self.action_dtype = utils.parse_dtype(self._batch_env.action_space)
     with tf.variable_scope('env_temporary'):
       self._observ = tf.Variable(
-          tf.zeros((len(self._batch_env),) + observ_shape, observ_dtype),
+          tf.zeros((len(self._batch_env),) + self.observ_shape,
+                   self.observ_dtype),
           name='observ', trainable=False)
 
   def __getattr__(self, name):
@@ -86,11 +84,9 @@ def simulate(self, action):
     with tf.name_scope('environment/simulate'):
       if action.dtype in (tf.float16, tf.float32, tf.float64):
         action = tf.check_numerics(action, 'action')
-      observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
       observ, reward, done = tf.py_func(
           lambda a: self._batch_env.step(a)[:3], [action],
-          [observ_dtype, tf.float32, tf.bool], name='step')
-      observ = tf.check_numerics(observ, 'observ')
+          [self.observ_dtype, tf.float32, tf.bool], name='step')
       reward = tf.check_numerics(reward, 'reward')
       reward.set_shape((len(self),))
       done.set_shape((len(self),))
@@ -106,10 +102,9 @@ def _reset_non_empty(self, indices):
     Returns:
       Batch tensor of the new observations.
     """
-    observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
     observ = tf.py_func(
-        self._batch_env.reset, [indices], observ_dtype, name='reset')
-    observ = tf.check_numerics(observ, 'observ')
+        self._batch_env.reset, [indices], self.observ_dtype, name='reset')
+    observ.set_shape(indices.get_shape().concatenate(self.observ_shape))
     with tf.control_dependencies([
         tf.scatter_update(self._observ, indices, observ)]):
       return tf.identity(observ)
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 99f2d0f3a..184728d58 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -22,8 +22,7 @@
 from __future__ import print_function
 
 from tensor2tensor.layers import common_layers
-from tensor2tensor.rl.envs import in_graph_batch_env
-from tensor2tensor.rl.envs.utils import get_action_space
+from tensor2tensor.rl.envs import in_graph_batch_env, utils
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
@@ -39,20 +38,21 @@
 class HistoryBuffer(object):
   """History Buffer."""
 
-  def __init__(self, input_dataset, length):
+  def __init__(self, input_dataset, length, observ_dtype):
     self.input_data_iterator = (
         input_dataset.batch(length).make_initializable_iterator())
     self.length = length
+    self._observ_dtype = observ_dtype
     initial_frames = self.get_initial_observations()
     initial_shape = [length] + common_layers.shape_list(initial_frames)[1:]
-    self._history_buff = tf.Variable(tf.zeros(initial_shape, tf.float32),
+    self._history_buff = tf.Variable(tf.zeros(initial_shape, observ_dtype),
                                      trainable=False)
 
   def initialize(self, sess):
     sess.run(self.input_data_iterator.initializer)
 
   def get_initial_observations(self):
-    return tf.cast(self.input_data_iterator.get_next(), tf.float32)
+    return tf.cast(self.input_data_iterator.get_next(), self._observ_dtype)
 
   def get_all_elements(self):
     return self._history_buff.read_value()
@@ -100,8 +100,16 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   def __init__(self, environment_spec, length):
     """Batch of environments inside the TensorFlow graph."""
 
-    self.length = length
+    observ_space = utils.get_observation_space(environment_spec)
     initial_frames_problem = environment_spec.initial_frames_problem
+    observ_shape = (initial_frames_problem.frame_height,
+                    initial_frames_problem.frame_width,
+                    initial_frames_problem.num_channels)
+    observ_space.shape = observ_shape
+    action_space = utils.get_action_space(environment_spec)
+    super(SimulatedBatchEnv, self).__init__(observ_space, action_space)
+
+    self.length = length
     self._min_reward = initial_frames_problem.min_reward
     self._num_frames = environment_spec.video_num_input_frames
     self._intrinsic_reward_scale = environment_spec.intrinsic_reward_scale
@@ -112,8 +120,6 @@ def __init__(self, environment_spec, length):
     self._model = registry.model(FLAGS.model)(
         model_hparams, tf.estimator.ModeKeys.PREDICT)
 
-    _, self.action_shape, self.action_dtype = get_action_space(environment_spec)
-
     hparams = HParams(video_num_input_frames=
                       environment_spec.video_num_input_frames,
                       video_num_target_frames=
@@ -133,12 +139,11 @@ def __init__(self, environment_spec, length):
                                                hparams=hparams).take(1)
 
     dataset = dataset.map(lambda x: x["inputs"]).repeat()
-    self.history_buffer = HistoryBuffer(dataset, self.length)
+    self.history_buffer = HistoryBuffer(dataset, self.length, self.observ_dtype)
 
-    shape = (self.length, initial_frames_problem.frame_height,
-             initial_frames_problem.frame_width,
-             initial_frames_problem.num_channels)
-    self._observ = tf.Variable(tf.zeros(shape, tf.float32), trainable=False)
+    self._observ = tf.Variable(
+        tf.zeros((len(self),) + observ_shape, self.observ_dtype),
+        trainable=False)
 
   def initialize(self, sess):
     self.history_buffer.initialize(sess)
@@ -160,7 +165,8 @@ def simulate(self, action):
             {"inputs": history, "input_action": actions})
         self._model.hparams.video_num_target_frames = hparams_target_frames
 
-      observ = tf.to_float(tf.squeeze(model_output["targets"], axis=1))
+      observ = tf.cast(tf.squeeze(model_output["targets"], axis=1),
+                       self.observ_dtype)
 
       reward = tf.to_float(model_output["target_reward"])
       reward = tf.reshape(reward, shape=(self.length,)) + self._min_reward
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index fcf744f0c..4df4df6bd 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -33,10 +33,10 @@ class WrapperBase(InGraphBatchEnv):
   """Base wrapper class."""
 
   def __init__(self, batch_env):
+    super(WrapperBase, self).__init__(
+        batch_env.observ_space, batch_env.action_space)
     self._length = len(batch_env)
     self._batch_env = batch_env
-    self.action_shape = batch_env.action_shape
-    self.action_dtype = batch_env.action_dtype
 
   def initialize(self, sess):
     """Initializations to be run once the tf.Session is available."""
@@ -83,8 +83,7 @@ def __init__(self, batch_env, skip=4):
     super(MaxAndSkipWrapper, self).__init__(batch_env)
     self.skip = skip
     observs_shape = batch_env.observ.shape
-    observ_dtype = tf.float32
-    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
+    self._observ = tf.Variable(tf.zeros(observs_shape, self.observ_dtype),
                                trainable=False)
 
   def simulate(self, action):
@@ -120,16 +119,19 @@ class StackAndSkipWrapper(WrapperBase):
   def __init__(self, batch_env, skip=4):
     super(StackAndSkipWrapper, self).__init__(batch_env)
     self.skip = skip
-    self._observ = None
-    self.old_shape = batch_env.observ.shape.as_list()
-    observs_shape = self.old_shape[:-1] + [self.old_shape[-1] * self.skip]
-    observ_dtype = tf.float32
-    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
-                               trainable=False)
+    self.old_shape = self._batch_env.observ_shape
+    self._observ = tf.Variable(
+        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
+        trainable=False)
+
+  @property
+  def observ_shape(self):
+    return self.old_shape[:-1] + (self.old_shape[-1] * self.skip,)
 
   def simulate(self, action):
     with tf.name_scope("environment/simulate"):  # Do we need this?
-      initializer = (tf.zeros(self.old_shape, dtype=tf.float32),
+      initializer = (tf.zeros((len(self),) + self.old_shape,
+                              dtype=self.observ_dtype),
                      tf.fill((len(self),), 0.0), tf.fill((len(self),), False))
 
       def not_done_step(a, _):
@@ -156,7 +158,8 @@ def _reset_non_empty(self, indices):
     # pylint: enable=protected-access
     inx = tf.concat(
         [
-            tf.ones(tf.size(tf.shape(new_values)), dtype=tf.int32)[:-1],
+            tf.ones(tf.size(tf.shape(new_values)),
+                    dtype=tf.int32)[:-1],
             [self.skip]
         ],
         axis=0)
@@ -172,11 +175,14 @@ class StackWrapper(WrapperBase):
   def __init__(self, batch_env, history=4):
     super(StackWrapper, self).__init__(batch_env)
     self.history = history
-    self.old_shape = batch_env.observ.shape.as_list()
-    observs_shape = self.old_shape[:-1] + [self.old_shape[-1] * self.history]
-    observ_dtype = tf.float32
-    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
-                               trainable=False)
+    self.old_shape = batch_env.observ_shape
+    self._observ = tf.Variable(
+        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
+        trainable=False)
+
+  @property
+  def observ_shape(self):
+    return self.old_shape[:-1] + (self.old_shape[-1] * self.history,)
 
   def simulate(self, action):
     reward, done = self._batch_env.simulate(action)
@@ -197,7 +203,8 @@ def _reset_non_empty(self, indices):
     # pylint: enable=protected-access
     inx = tf.concat(
         [
-            tf.ones(tf.size(tf.shape(new_values)), dtype=tf.int32)[:-1],
+            tf.ones(tf.size(tf.shape(new_values)),
+                    dtype=tf.int32)[:-1],
             [self.history]
         ],
         axis=0)
@@ -213,19 +220,23 @@ class AutoencoderWrapper(WrapperBase):
 
   def __init__(self, batch_env):
     super(AutoencoderWrapper, self).__init__(batch_env)
-    batch_size, height, width, _ = self._batch_env.observ.get_shape().as_list()
-    ae_height = int(math.ceil(height / self.autoencoder_factor))
-    ae_width = int(math.ceil(width / self.autoencoder_factor))
-    ae_channels = 24  # TODO(piotrmilos): make it better
-    observ_shape = (batch_size, ae_height, ae_width, ae_channels)
     self._observ = self._observ = tf.Variable(
-        tf.zeros(observ_shape, tf.float32), trainable=False)
+        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
+        trainable=False)
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       autoencoder_hparams = autoencoders.autoencoder_discrete_pong()
       self.autoencoder_model = autoencoders.AutoencoderOrderedDiscrete(
           autoencoder_hparams, tf.estimator.ModeKeys.EVAL)
     self.autoencoder_model.set_mode(tf.estimator.ModeKeys.EVAL)
 
+  @property
+  def observ_shape(self):
+    height, width, _ = self._batch_env.observ_shape
+    ae_height = int(math.ceil(height / self.autoencoder_factor))
+    ae_width = int(math.ceil(width / self.autoencoder_factor))
+    ae_channels = 24  # TODO(piotrmilos): make it better
+    return (ae_height, ae_width, ae_channels)
+
   @property
   def autoencoder_factor(self):
     """By how much to divide sizes when using autoencoders."""
@@ -237,6 +248,7 @@ def simulate(self, action):
     with tf.control_dependencies([reward, done]):
       with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
         ret = self.autoencoder_model.encode(self._batch_env.observ)
+        ret = tf.cast(ret, self.observ_dtype)
         assign_op = self._observ.assign(ret)
         with tf.control_dependencies([assign_op]):
           return tf.identity(reward), tf.identity(done)
@@ -245,6 +257,7 @@ def _reset_non_empty(self, indices):
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       new_values = self._batch_env._reset_non_empty(indices)  # pylint: disable=protected-access
       ret = self.autoencoder_model.encode(new_values)
+      ret = tf.cast(ret, self.observ_dtype)
       assign_op = tf.scatter_update(self._observ, indices, ret)
       with tf.control_dependencies([assign_op]):
         return tf.gather(self.observ, indices)
@@ -255,14 +268,16 @@ class IntToBitWrapper(WrapperBase):
 
   def __init__(self, batch_env):
     super(IntToBitWrapper, self).__init__(batch_env)
-    batch_size, height, width, channels = \
-      self._batch_env.observ.get_shape().as_list()
-    # We treat each channel as 8-bit integer to be expanded to 8 channels
-    self.observ_shape = (height, width, channels*8)
     self._observ = self._observ = tf.Variable(
-        tf.zeros((batch_size,) + self.observ_shape, tf.float32),
+        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
         trainable=False)
 
+  @property
+  def observ_shape(self):
+    height, width, channels = self._batch_env.observ_shape
+    # We treat each channel as 8-bit integer to be expanded to 8 channels
+    return (height, width, channels*8)
+
   def simulate(self, action):
     action = tf.Print(action, [action], message="action=", summarize=200)
 
@@ -272,6 +287,7 @@ def simulate(self, action):
       with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
         unpacked = discretization.int_to_bit(self._batch_env.observ, 8)
         unpacked = tf.reshape(unpacked, (-1,)+self.observ_shape)
+        unpacked = tf.cast(unpacked, self.observ_dtype)
         assign_op = self._observ.assign(unpacked)
         with tf.control_dependencies([assign_op]):
           return tf.identity(reward), tf.identity(done)
@@ -282,6 +298,7 @@ def _reset_non_empty(self, indices):
     new_values_unpacked = discretization.int_to_bit(new_values, 8)
     new_values_unpacked = tf.reshape(new_values_unpacked, (-1,)
                                      +self.observ_shape)
+    new_values_unpacked = tf.cast(new_values_unpacked, self.observ_dtype)
     # pylint: enable=protected-access
     assign_op = tf.scatter_update(self._observ, indices, new_values_unpacked)
     with tf.control_dependencies([assign_op]):
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index b23dc358e..badd48485 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -65,8 +65,20 @@ def _reset(self, **kwargs):
     return self._last_returned[0]
 
 
+def get_observation_space(environment_spec):
+  """Get observation space associated with environment spec.
+
+  Args:
+     environment_spec:  EnvironmentSpec object
+
+  Returns:
+    OpenAi Gym observation space
+  """
+  return environment_spec.env_lambda().observation_space
+
+
 def get_action_space(environment_spec):
-  """Get action spece associated with environment spec.
+  """Get action space associated with environment spec.
 
   Args:
      environment_spec:  EnvironmentSpec object
@@ -74,11 +86,7 @@ def get_action_space(environment_spec):
   Returns:
     OpenAi Gym action space
   """
-  action_space = environment_spec.env_lambda().action_space
-  action_shape = list(parse_shape(action_space))
-  action_dtype = parse_dtype(action_space)
-
-  return action_space, action_shape, action_dtype
+  return environment_spec.env_lambda().action_space
 
 
 def get_policy(observations, hparams):
@@ -92,7 +100,7 @@ def get_policy(observations, hparams):
     Tensor with policy and value function output
   """
   policy_network_lambda = hparams.policy_network
-  action_space, _, _ = get_action_space(hparams.environment_spec)
+  action_space = get_action_space(hparams.environment_spec)
   return policy_network_lambda(action_space, hparams, observations)
 
 
@@ -124,5 +132,5 @@ def parse_dtype(space):
   if isinstance(space, gym.spaces.Discrete):
     return tf.int32
   if isinstance(space, gym.spaces.Box):
-    return tf.float32
+    return tf.as_dtype(space.dtype)
   raise NotImplementedError()

From ee6d538944fb05c2ac853828ef4bd5275f01eac5 Mon Sep 17 00:00:00 2001
From: Keyon Vafa <vafa@google.com>
Date: Thu, 30 Aug 2018 17:43:26 -0700
Subject: [PATCH 0709/2720] Changed project_hidden so its output is 4D

PiperOrigin-RevId: 211013213
---
 tensor2tensor/layers/discretization.py       | 35 +++++----
 tensor2tensor/layers/discretization_test.py  | 12 ++--
 tensor2tensor/rl/collect.py                  |  8 +--
 tensor2tensor/rl/envs/batch_env.py           | 16 ++---
 tensor2tensor/rl/envs/in_graph_batch_env.py  | 25 +------
 tensor2tensor/rl/envs/py_func_batch_env.py   | 19 +++--
 tensor2tensor/rl/envs/simulated_batch_env.py | 34 ++++-----
 tensor2tensor/rl/envs/tf_atari_wrappers.py   | 75 ++++++++------------
 tensor2tensor/rl/envs/utils.py               | 24 +++----
 9 files changed, 104 insertions(+), 144 deletions(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 3188855cc..8a63b8531 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -29,40 +29,44 @@
 
 
 def project_hidden(x, projection_tensors, hidden_size, num_blocks):
-  """Project encoder hidden state into block_dim using projection tensors.
+  """Project encoder hidden state under num_blocks using projection tensors.
 
   Args:
-    x: Encoder hidden state of shape [-1, hidden_size].
+    x: Encoder hidden state of shape [batch_size, latent_dim,  hidden_size].
     projection_tensors: Projection tensors used to project the hidden state.
     hidden_size: Dimension of the latent space.
     num_blocks: Number of blocks in DVQ.
 
   Returns:
-    Projected states of shape [-1, num_blocks, block_dim].
+    x_projected: Projected states of shape [batch_size, latent_dim, num_blocks,
+      hidden_size / num_blocks].
   """
+  batch_size, latent_dim, _ = common_layers.shape_list(x)
   x = tf.reshape(x, shape=[1, -1, hidden_size])
   x_tiled = tf.reshape(
       tf.tile(x, multiples=[num_blocks, 1, 1]),
       shape=[num_blocks, -1, hidden_size])
   x_projected = tf.matmul(x_tiled, projection_tensors)
   x_projected = tf.transpose(x_projected, perm=[1, 0, 2])
-  return x_projected
+  x_4d = tf.reshape(x_projected, [batch_size, latent_dim, num_blocks, -1])
+  return x_4d
 
 
 def slice_hidden(x, hidden_size, num_blocks):
-  """Slice encoder hidden state into block_dim.
+  """Slice encoder hidden state under num_blocks.
 
   Args:
-    x: Encoder hidden state of shape [..., 1, hidden_size].
+    x: Encoder hidden state of shape [batch_size, latent_dim, hidden_size].
     hidden_size: Dimension of the latent space.
     num_blocks: Number of blocks in DVQ.
 
   Returns:
-    Sliced states of shape [..., num_blocks, block_dim].
+    Sliced states of shape [batch_size, latent_dim, num_blocks, block_dim].
   """
+  batch_size, latent_dim, _ = common_layers.shape_list(x)
   block_dim = hidden_size // num_blocks
-  x_shape = common_layers.shape_list(x)
-  x_sliced = tf.reshape(x, shape=(x_shape[:-2] + [num_blocks, block_dim]))
+  x_sliced = tf.reshape(x,
+                        shape=[batch_size, latent_dim, num_blocks, block_dim])
   return x_sliced
 
 
@@ -92,7 +96,8 @@ def nearest_neighbor(x,
     Tensor with nearest element in mean encoded in one-hot notation
     and distances.
   """
-  x = tf.reshape(x, [-1] + common_layers.shape_list(x)[2:])
+  batch_size, latent_dim, num_blocks, block_dim = common_layers.shape_list(x)
+  x = tf.reshape(x, [batch_size * latent_dim, num_blocks, block_dim])
   x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keep_dims=True)
   means_norm_sq = tf.reduce_sum(tf.square(means), axis=-1, keep_dims=True)
   scalar_prod = tf.matmul(
@@ -205,7 +210,8 @@ def embedding_lookup(x,
   x_means_hot_flat = tf.reshape(x_means_hot, [-1, num_blocks, block_v_size])
   x_means = tf.matmul(tf.transpose(x_means_hot_flat, perm=[1, 0, 2]), means)
   x_means = tf.transpose(x_means, [1, 0, 2])
-  x = tf.reshape(x, [-1] + common_layers.shape_list(x)[2:])
+  batch_size, latent_dim, num_blocks, block_dim = common_layers.shape_list(x)
+  x = tf.reshape(x, [batch_size * latent_dim, num_blocks, block_dim])
 
   # Currently, we use the mean scaling for the commitment loss, as opposed to
   # summing across all non-batch dimensions.
@@ -618,15 +624,18 @@ def discrete_bottleneck(inputs,
       extra_loss = tf.constant(0.0)
       neg_q_entropy = tf.constant(0.0)
     elif bottleneck_kind in ["dvq", "gumbel-softmax-dvq"]:
+      inputs_3d = inputs
+      if len(inputs.shape) == 4:
+        inputs_3d = tf.squeeze(inputs, axis=2)
       if reshape_method == "slice":
         x_reshaped = slice_hidden(
-            inputs, hidden_size=hidden_size, num_blocks=num_blocks)
+            inputs_3d, hidden_size=hidden_size, num_blocks=num_blocks)
       elif reshape_method == "project":
         if projection_tensors is None:
           raise ValueError(
               "Projection tensors is None for reshape_method project")
         x_reshaped = project_hidden(
-            inputs,
+            inputs_3d,
             projection_tensors=projection_tensors,
             hidden_size=hidden_size,
             num_blocks=num_blocks)
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index d26a92042..bf82b1ae3 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -67,13 +67,13 @@ def testProjectHidden(self):
     hidden_size = 60
     block_dim = 20
     num_blocks = 3
-    x = tf.zeros(shape=[1, hidden_size], dtype=tf.float32)
+    x = tf.zeros(shape=[1, 1, hidden_size], dtype=tf.float32)
     projection_tensors = tf.random_normal(
         shape=[num_blocks, hidden_size, block_dim], dtype=tf.float32)
     x_projected = discretization.project_hidden(x, projection_tensors,
                                                 hidden_size, num_blocks)
     x_projected_eval = self.evaluate(x_projected)
-    self.assertEqual(np.shape(x_projected_eval), (1, num_blocks, block_dim))
+    self.assertEqual(np.shape(x_projected_eval), (1, 1, num_blocks, block_dim))
     self.assertTrue(np.all(x_projected_eval == 0))
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
@@ -81,10 +81,10 @@ def testSliceHiddenZeros(self):
     hidden_size = 60
     block_dim = 20
     num_blocks = 3
-    x = tf.zeros(shape=[1, hidden_size], dtype=tf.float32)
+    x = tf.zeros(shape=[1, 1, hidden_size], dtype=tf.float32)
     x_sliced = discretization.slice_hidden(x, hidden_size, num_blocks)
     x_sliced_eval = self.evaluate(x_sliced)
-    self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
+    self.assertEqual(np.shape(x_sliced_eval), (1, 1, num_blocks, block_dim))
     self.assertTrue(np.all(x_sliced_eval == 0))
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
@@ -92,10 +92,10 @@ def testSliceHiddenOnes(self):
     hidden_size = 60
     block_dim = 20
     num_blocks = 3
-    x = tf.ones(shape=[1, hidden_size], dtype=tf.float32)
+    x = tf.ones(shape=[1, 1, hidden_size], dtype=tf.float32)
     x_sliced = discretization.slice_hidden(x, hidden_size, num_blocks)
     x_sliced_eval = self.evaluate(x_sliced)
-    self.assertEqual(np.shape(x_sliced_eval), (num_blocks, block_dim))
+    self.assertEqual(np.shape(x_sliced_eval), (1, 1, num_blocks, block_dim))
     self.assertTrue(np.all(x_sliced_eval == 1))
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index bde467969..f32fca858 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -33,11 +33,10 @@ def _rollout_metadata(batch_env):
   batch_size = [batch_env_shape[0]]
   shapes_types_names = [
       # TODO(piotrmilos): possibly retrieve the observation type for batch_env
-      (batch_size + batch_env_shape[1:], batch_env.observ_dtype, "observation"),
+      (batch_size + batch_env_shape[1:], tf.float32, "observation"),
       (batch_size, tf.float32, "reward"),
       (batch_size, tf.bool, "done"),
-      (batch_size + list(batch_env.action_shape), batch_env.action_dtype,
-       "action"),
+      (batch_size + batch_env.action_shape, batch_env.action_dtype, "action"),
       (batch_size, tf.float32, "pdf"),
       (batch_size, tf.float32, "value_function"),
   ]
@@ -58,7 +57,8 @@ def __init__(self, batch_env):
     self.speculum = tf.FIFOQueue(infinity, shapes=shapes, dtypes=dtypes)
     observs_shape = batch_env.observ.shape
     # TODO(piotrmilos): possibly retrieve the observation type for batch_env
-    self._observ = tf.Variable(tf.zeros(observs_shape, self.observ_dtype),
+    observ_dtype = tf.float32
+    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
                                trainable=False)
 
   def simulate(self, action):
diff --git a/tensor2tensor/rl/envs/batch_env.py b/tensor2tensor/rl/envs/batch_env.py
index 230bd3a43..bf433066f 100644
--- a/tensor2tensor/rl/envs/batch_env.py
+++ b/tensor2tensor/rl/envs/batch_env.py
@@ -43,12 +43,11 @@ def __init__(self, envs, blocking):
     """
     self._envs = envs
     self._blocking = blocking
-    self.observ_space = self._envs[0].observation_space
-    if not all(env.observation_space == self.observ_space
-               for env in self._envs):
+    observ_space = self._envs[0].observation_space
+    if not all(env.observation_space == observ_space for env in self._envs):
       raise ValueError('All environments must use the same observation space.')
-    self.action_space = self._envs[0].action_space
-    if not all(env.action_space == self.action_space for env in self._envs):
+    action_space = self._envs[0].action_space
+    if not all(env.action_space == action_space for env in self._envs):
       raise ValueError('All environments must use the same observation space.')
 
   def __len__(self):
@@ -97,8 +96,8 @@ def step(self, actions):
       transitions = [transition() for transition in transitions]
     observs, rewards, dones, infos = zip(*transitions)
 
-    observ = np.stack(observs).astype(self.observ_space.dtype)
     # TODO(piotrmilos): Do we really want cast to float32
+    observ = np.stack(observs).astype(np.float32)
     reward = np.stack(rewards).astype(np.float32)
     done = np.stack(dones)
     info = tuple(infos)
@@ -120,8 +119,9 @@ def reset(self, indices=None):
     else:
       observs = [self._envs[index].reset(blocking=False) for index in indices]
       observs = [observ() for observ in observs]
-    observ = np.stack(observs).astype(self.observ_space.dtype)
-
+    observ = np.stack(observs)
+    # TODO(piotrmilos): Do we really want this?
+    observ = observ.astype(np.float32)
     return observ
 
   def close(self):
diff --git a/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index 8db9b56ef..56e13878d 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -23,17 +23,11 @@
 
 import tensorflow as tf
 
-from tensor2tensor.rl.envs import utils
-
 
 class InGraphBatchEnv(object):
   """Abstract class for batch of environments inside the TensorFlow graph.
   """
 
-  def __init__(self, observ_space, action_space):
-    self.observ_space = observ_space
-    self.action_space = action_space
-
   def __getattr__(self, name):
     """Forward unimplemented attributes to one of the original environments.
 
@@ -77,24 +71,7 @@ def reset(self, indices=None):
     """
     return tf.cond(
         tf.cast(tf.reduce_sum(indices + 1), tf.bool),
-        lambda: self._reset_non_empty(indices),
-        lambda: tf.cast(0, self.observ_dtype))
-
-  @property
-  def observ_dtype(self):
-    return utils.parse_dtype(self.observ_space)
-
-  @property
-  def observ_shape(self):
-    return utils.parse_shape(self.observ_space)
-
-  @property
-  def action_dtype(self):
-    return utils.parse_dtype(self.action_space)
-
-  @property
-  def action_shape(self):
-    return utils.parse_shape(self.action_space)
+        lambda: self._reset_non_empty(indices), lambda: 0.0)
 
   @property
   def observ(self):
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index ac1ad1dd5..76916d991 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -21,6 +21,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.rl.envs import utils
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
 import tensorflow as tf
 
@@ -39,13 +40,14 @@ def __init__(self, batch_env):
     Args:
       batch_env: Batch environment.
     """
-    super(PyFuncBatchEnv, self).__init__(batch_env.observation_space,
-                                         batch_env.action_space)
     self._batch_env = batch_env
+    observ_shape = utils.parse_shape(self._batch_env.observation_space)
+    observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
+    self.action_shape = list(utils.parse_shape(self._batch_env.action_space))
+    self.action_dtype = utils.parse_dtype(self._batch_env.action_space)
     with tf.variable_scope('env_temporary'):
       self._observ = tf.Variable(
-          tf.zeros((len(self._batch_env),) + self.observ_shape,
-                   self.observ_dtype),
+          tf.zeros((len(self._batch_env),) + observ_shape, observ_dtype),
           name='observ', trainable=False)
 
   def __getattr__(self, name):
@@ -84,9 +86,11 @@ def simulate(self, action):
     with tf.name_scope('environment/simulate'):
       if action.dtype in (tf.float16, tf.float32, tf.float64):
         action = tf.check_numerics(action, 'action')
+      observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
       observ, reward, done = tf.py_func(
           lambda a: self._batch_env.step(a)[:3], [action],
-          [self.observ_dtype, tf.float32, tf.bool], name='step')
+          [observ_dtype, tf.float32, tf.bool], name='step')
+      observ = tf.check_numerics(observ, 'observ')
       reward = tf.check_numerics(reward, 'reward')
       reward.set_shape((len(self),))
       done.set_shape((len(self),))
@@ -102,9 +106,10 @@ def _reset_non_empty(self, indices):
     Returns:
       Batch tensor of the new observations.
     """
+    observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
     observ = tf.py_func(
-        self._batch_env.reset, [indices], self.observ_dtype, name='reset')
-    observ.set_shape(indices.get_shape().concatenate(self.observ_shape))
+        self._batch_env.reset, [indices], observ_dtype, name='reset')
+    observ = tf.check_numerics(observ, 'observ')
     with tf.control_dependencies([
         tf.scatter_update(self._observ, indices, observ)]):
       return tf.identity(observ)
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 184728d58..99f2d0f3a 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -22,7 +22,8 @@
 from __future__ import print_function
 
 from tensor2tensor.layers import common_layers
-from tensor2tensor.rl.envs import in_graph_batch_env, utils
+from tensor2tensor.rl.envs import in_graph_batch_env
+from tensor2tensor.rl.envs.utils import get_action_space
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
@@ -38,21 +39,20 @@
 class HistoryBuffer(object):
   """History Buffer."""
 
-  def __init__(self, input_dataset, length, observ_dtype):
+  def __init__(self, input_dataset, length):
     self.input_data_iterator = (
         input_dataset.batch(length).make_initializable_iterator())
     self.length = length
-    self._observ_dtype = observ_dtype
     initial_frames = self.get_initial_observations()
     initial_shape = [length] + common_layers.shape_list(initial_frames)[1:]
-    self._history_buff = tf.Variable(tf.zeros(initial_shape, observ_dtype),
+    self._history_buff = tf.Variable(tf.zeros(initial_shape, tf.float32),
                                      trainable=False)
 
   def initialize(self, sess):
     sess.run(self.input_data_iterator.initializer)
 
   def get_initial_observations(self):
-    return tf.cast(self.input_data_iterator.get_next(), self._observ_dtype)
+    return tf.cast(self.input_data_iterator.get_next(), tf.float32)
 
   def get_all_elements(self):
     return self._history_buff.read_value()
@@ -100,16 +100,8 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   def __init__(self, environment_spec, length):
     """Batch of environments inside the TensorFlow graph."""
 
-    observ_space = utils.get_observation_space(environment_spec)
-    initial_frames_problem = environment_spec.initial_frames_problem
-    observ_shape = (initial_frames_problem.frame_height,
-                    initial_frames_problem.frame_width,
-                    initial_frames_problem.num_channels)
-    observ_space.shape = observ_shape
-    action_space = utils.get_action_space(environment_spec)
-    super(SimulatedBatchEnv, self).__init__(observ_space, action_space)
-
     self.length = length
+    initial_frames_problem = environment_spec.initial_frames_problem
     self._min_reward = initial_frames_problem.min_reward
     self._num_frames = environment_spec.video_num_input_frames
     self._intrinsic_reward_scale = environment_spec.intrinsic_reward_scale
@@ -120,6 +112,8 @@ def __init__(self, environment_spec, length):
     self._model = registry.model(FLAGS.model)(
         model_hparams, tf.estimator.ModeKeys.PREDICT)
 
+    _, self.action_shape, self.action_dtype = get_action_space(environment_spec)
+
     hparams = HParams(video_num_input_frames=
                       environment_spec.video_num_input_frames,
                       video_num_target_frames=
@@ -139,11 +133,12 @@ def __init__(self, environment_spec, length):
                                                hparams=hparams).take(1)
 
     dataset = dataset.map(lambda x: x["inputs"]).repeat()
-    self.history_buffer = HistoryBuffer(dataset, self.length, self.observ_dtype)
+    self.history_buffer = HistoryBuffer(dataset, self.length)
 
-    self._observ = tf.Variable(
-        tf.zeros((len(self),) + observ_shape, self.observ_dtype),
-        trainable=False)
+    shape = (self.length, initial_frames_problem.frame_height,
+             initial_frames_problem.frame_width,
+             initial_frames_problem.num_channels)
+    self._observ = tf.Variable(tf.zeros(shape, tf.float32), trainable=False)
 
   def initialize(self, sess):
     self.history_buffer.initialize(sess)
@@ -165,8 +160,7 @@ def simulate(self, action):
             {"inputs": history, "input_action": actions})
         self._model.hparams.video_num_target_frames = hparams_target_frames
 
-      observ = tf.cast(tf.squeeze(model_output["targets"], axis=1),
-                       self.observ_dtype)
+      observ = tf.to_float(tf.squeeze(model_output["targets"], axis=1))
 
       reward = tf.to_float(model_output["target_reward"])
       reward = tf.reshape(reward, shape=(self.length,)) + self._min_reward
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 4df4df6bd..fcf744f0c 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -33,10 +33,10 @@ class WrapperBase(InGraphBatchEnv):
   """Base wrapper class."""
 
   def __init__(self, batch_env):
-    super(WrapperBase, self).__init__(
-        batch_env.observ_space, batch_env.action_space)
     self._length = len(batch_env)
     self._batch_env = batch_env
+    self.action_shape = batch_env.action_shape
+    self.action_dtype = batch_env.action_dtype
 
   def initialize(self, sess):
     """Initializations to be run once the tf.Session is available."""
@@ -83,7 +83,8 @@ def __init__(self, batch_env, skip=4):
     super(MaxAndSkipWrapper, self).__init__(batch_env)
     self.skip = skip
     observs_shape = batch_env.observ.shape
-    self._observ = tf.Variable(tf.zeros(observs_shape, self.observ_dtype),
+    observ_dtype = tf.float32
+    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
                                trainable=False)
 
   def simulate(self, action):
@@ -119,19 +120,16 @@ class StackAndSkipWrapper(WrapperBase):
   def __init__(self, batch_env, skip=4):
     super(StackAndSkipWrapper, self).__init__(batch_env)
     self.skip = skip
-    self.old_shape = self._batch_env.observ_shape
-    self._observ = tf.Variable(
-        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
-        trainable=False)
-
-  @property
-  def observ_shape(self):
-    return self.old_shape[:-1] + (self.old_shape[-1] * self.skip,)
+    self._observ = None
+    self.old_shape = batch_env.observ.shape.as_list()
+    observs_shape = self.old_shape[:-1] + [self.old_shape[-1] * self.skip]
+    observ_dtype = tf.float32
+    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
+                               trainable=False)
 
   def simulate(self, action):
     with tf.name_scope("environment/simulate"):  # Do we need this?
-      initializer = (tf.zeros((len(self),) + self.old_shape,
-                              dtype=self.observ_dtype),
+      initializer = (tf.zeros(self.old_shape, dtype=tf.float32),
                      tf.fill((len(self),), 0.0), tf.fill((len(self),), False))
 
       def not_done_step(a, _):
@@ -158,8 +156,7 @@ def _reset_non_empty(self, indices):
     # pylint: enable=protected-access
     inx = tf.concat(
         [
-            tf.ones(tf.size(tf.shape(new_values)),
-                    dtype=tf.int32)[:-1],
+            tf.ones(tf.size(tf.shape(new_values)), dtype=tf.int32)[:-1],
             [self.skip]
         ],
         axis=0)
@@ -175,14 +172,11 @@ class StackWrapper(WrapperBase):
   def __init__(self, batch_env, history=4):
     super(StackWrapper, self).__init__(batch_env)
     self.history = history
-    self.old_shape = batch_env.observ_shape
-    self._observ = tf.Variable(
-        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
-        trainable=False)
-
-  @property
-  def observ_shape(self):
-    return self.old_shape[:-1] + (self.old_shape[-1] * self.history,)
+    self.old_shape = batch_env.observ.shape.as_list()
+    observs_shape = self.old_shape[:-1] + [self.old_shape[-1] * self.history]
+    observ_dtype = tf.float32
+    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
+                               trainable=False)
 
   def simulate(self, action):
     reward, done = self._batch_env.simulate(action)
@@ -203,8 +197,7 @@ def _reset_non_empty(self, indices):
     # pylint: enable=protected-access
     inx = tf.concat(
         [
-            tf.ones(tf.size(tf.shape(new_values)),
-                    dtype=tf.int32)[:-1],
+            tf.ones(tf.size(tf.shape(new_values)), dtype=tf.int32)[:-1],
             [self.history]
         ],
         axis=0)
@@ -220,23 +213,19 @@ class AutoencoderWrapper(WrapperBase):
 
   def __init__(self, batch_env):
     super(AutoencoderWrapper, self).__init__(batch_env)
+    batch_size, height, width, _ = self._batch_env.observ.get_shape().as_list()
+    ae_height = int(math.ceil(height / self.autoencoder_factor))
+    ae_width = int(math.ceil(width / self.autoencoder_factor))
+    ae_channels = 24  # TODO(piotrmilos): make it better
+    observ_shape = (batch_size, ae_height, ae_width, ae_channels)
     self._observ = self._observ = tf.Variable(
-        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
-        trainable=False)
+        tf.zeros(observ_shape, tf.float32), trainable=False)
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       autoencoder_hparams = autoencoders.autoencoder_discrete_pong()
       self.autoencoder_model = autoencoders.AutoencoderOrderedDiscrete(
           autoencoder_hparams, tf.estimator.ModeKeys.EVAL)
     self.autoencoder_model.set_mode(tf.estimator.ModeKeys.EVAL)
 
-  @property
-  def observ_shape(self):
-    height, width, _ = self._batch_env.observ_shape
-    ae_height = int(math.ceil(height / self.autoencoder_factor))
-    ae_width = int(math.ceil(width / self.autoencoder_factor))
-    ae_channels = 24  # TODO(piotrmilos): make it better
-    return (ae_height, ae_width, ae_channels)
-
   @property
   def autoencoder_factor(self):
     """By how much to divide sizes when using autoencoders."""
@@ -248,7 +237,6 @@ def simulate(self, action):
     with tf.control_dependencies([reward, done]):
       with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
         ret = self.autoencoder_model.encode(self._batch_env.observ)
-        ret = tf.cast(ret, self.observ_dtype)
         assign_op = self._observ.assign(ret)
         with tf.control_dependencies([assign_op]):
           return tf.identity(reward), tf.identity(done)
@@ -257,7 +245,6 @@ def _reset_non_empty(self, indices):
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       new_values = self._batch_env._reset_non_empty(indices)  # pylint: disable=protected-access
       ret = self.autoencoder_model.encode(new_values)
-      ret = tf.cast(ret, self.observ_dtype)
       assign_op = tf.scatter_update(self._observ, indices, ret)
       with tf.control_dependencies([assign_op]):
         return tf.gather(self.observ, indices)
@@ -268,16 +255,14 @@ class IntToBitWrapper(WrapperBase):
 
   def __init__(self, batch_env):
     super(IntToBitWrapper, self).__init__(batch_env)
+    batch_size, height, width, channels = \
+      self._batch_env.observ.get_shape().as_list()
+    # We treat each channel as 8-bit integer to be expanded to 8 channels
+    self.observ_shape = (height, width, channels*8)
     self._observ = self._observ = tf.Variable(
-        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
+        tf.zeros((batch_size,) + self.observ_shape, tf.float32),
         trainable=False)
 
-  @property
-  def observ_shape(self):
-    height, width, channels = self._batch_env.observ_shape
-    # We treat each channel as 8-bit integer to be expanded to 8 channels
-    return (height, width, channels*8)
-
   def simulate(self, action):
     action = tf.Print(action, [action], message="action=", summarize=200)
 
@@ -287,7 +272,6 @@ def simulate(self, action):
       with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
         unpacked = discretization.int_to_bit(self._batch_env.observ, 8)
         unpacked = tf.reshape(unpacked, (-1,)+self.observ_shape)
-        unpacked = tf.cast(unpacked, self.observ_dtype)
         assign_op = self._observ.assign(unpacked)
         with tf.control_dependencies([assign_op]):
           return tf.identity(reward), tf.identity(done)
@@ -298,7 +282,6 @@ def _reset_non_empty(self, indices):
     new_values_unpacked = discretization.int_to_bit(new_values, 8)
     new_values_unpacked = tf.reshape(new_values_unpacked, (-1,)
                                      +self.observ_shape)
-    new_values_unpacked = tf.cast(new_values_unpacked, self.observ_dtype)
     # pylint: enable=protected-access
     assign_op = tf.scatter_update(self._observ, indices, new_values_unpacked)
     with tf.control_dependencies([assign_op]):
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index badd48485..b23dc358e 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -65,20 +65,8 @@ def _reset(self, **kwargs):
     return self._last_returned[0]
 
 
-def get_observation_space(environment_spec):
-  """Get observation space associated with environment spec.
-
-  Args:
-     environment_spec:  EnvironmentSpec object
-
-  Returns:
-    OpenAi Gym observation space
-  """
-  return environment_spec.env_lambda().observation_space
-
-
 def get_action_space(environment_spec):
-  """Get action space associated with environment spec.
+  """Get action spece associated with environment spec.
 
   Args:
      environment_spec:  EnvironmentSpec object
@@ -86,7 +74,11 @@ def get_action_space(environment_spec):
   Returns:
     OpenAi Gym action space
   """
-  return environment_spec.env_lambda().action_space
+  action_space = environment_spec.env_lambda().action_space
+  action_shape = list(parse_shape(action_space))
+  action_dtype = parse_dtype(action_space)
+
+  return action_space, action_shape, action_dtype
 
 
 def get_policy(observations, hparams):
@@ -100,7 +92,7 @@ def get_policy(observations, hparams):
     Tensor with policy and value function output
   """
   policy_network_lambda = hparams.policy_network
-  action_space = get_action_space(hparams.environment_spec)
+  action_space, _, _ = get_action_space(hparams.environment_spec)
   return policy_network_lambda(action_space, hparams, observations)
 
 
@@ -132,5 +124,5 @@ def parse_dtype(space):
   if isinstance(space, gym.spaces.Discrete):
     return tf.int32
   if isinstance(space, gym.spaces.Box):
-    return tf.as_dtype(space.dtype)
+    return tf.float32
   raise NotImplementedError()

From 1638c02f12bbdadd53f83c576a1a9895c2237fa1 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 30 Aug 2018 19:10:09 -0700
Subject: [PATCH 0710/2720] Revert recent unintended parameter overwrite.

PiperOrigin-RevId: 211022104
---
 tensor2tensor/models/research/autoencoders.py   | 17 +++++++----------
 .../models/video/basic_deterministic.py         |  1 +
 tensor2tensor/rl/trainer_model_based.py         |  4 ++--
 3 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 6ff3fa17e..f77b0a9c4 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -584,6 +584,7 @@ def encoder(self, x):
           x = self.dropout(x)
           filters = hparams.hidden_size * 2**(i + 1)
           filters = min(filters, hparams.max_hidden_size)
+          x = common_attention.add_timing_signal_nd(x)
           x = tf.layers.conv2d(
               x,
               filters,
@@ -990,7 +991,7 @@ def autoencoder_basic_discrete():
   hparams = autoencoder_autoregressive()
   hparams.num_hidden_layers = 5
   hparams.hidden_size = 64
-  hparams.bottleneck_bits = 4096
+  hparams.bottleneck_bits = 1024
   hparams.bottleneck_noise = 0.1
   hparams.add_hparam("discretize_warmup_steps", 16000)
   return hparams
@@ -1000,8 +1001,8 @@ def autoencoder_basic_discrete():
 def autoencoder_residual_discrete():
   """Residual discrete autoencoder model."""
   hparams = autoencoder_residual()
-  hparams.bottleneck_bits = 4096
-  hparams.bottleneck_noise = 0.1
+  hparams.bottleneck_bits = 1024
+  hparams.bottleneck_noise = 0.05
   hparams.add_hparam("discretize_warmup_steps", 16000)
   hparams.add_hparam("bottleneck_kind", "tanh_discrete")
   hparams.add_hparam("isemhash_noise_dev", 0.5)
@@ -1028,9 +1029,9 @@ def autoencoder_residual_discrete_big():
 def autoencoder_ordered_discrete():
   """Ordered discrete autoencoder model."""
   hparams = autoencoder_residual_discrete()
-  hparams.bottleneck_noise = 0.8
+  hparams.bottleneck_noise = 0.05  # Use 0.8 for ordered.
   hparams.gan_loss_factor = 0.05
-  hparams.add_hparam("unordered", False)
+  hparams.add_hparam("unordered", True)
   return hparams
 
 
@@ -1062,9 +1063,7 @@ def autoencoder_ordered_discrete_hs256():
 def autoencoder_ordered_text():
   """Ordered discrete autoencoder model for text."""
   hparams = autoencoder_ordered_discrete()
-  hparams.bottleneck_bits = 1024
-  hparams.unordered = True
-  hparams.bottleneck_noise = 0.05
+  hparams.bottleneck_bits = 512
   hparams.num_hidden_layers = 7
   hparams.batch_size = 1024
   hparams.autoregressive_mode = "conv5"
@@ -1108,7 +1107,6 @@ def autoencoder_discrete_pong():
   hparams.batch_size = 2
   hparams.bottleneck_noise = 0.2
   hparams.max_hidden_size = 1024
-  hparams.unordered = True
   return hparams
 
 
@@ -1118,7 +1116,6 @@ def autoencoder_discrete_cifar():
   hparams = autoencoder_ordered_discrete()
   hparams.bottleneck_noise = 0.0
   hparams.bottleneck_bits = 90
-  hparams.unordered = True
   hparams.num_hidden_layers = 2
   hparams.hidden_size = 256
   hparams.num_residual_layers = 4
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index c0691f73c..364625dfd 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -63,6 +63,7 @@ def body_single(self, features):
         x = common_layers.make_even_size(x)
         if i < hparams.filter_double_steps:
           filters *= 2
+        x = common_attention.add_timing_signal_nd(x)
         x = tf.layers.conv2d(x, filters, kernel2, activation=common_layers.belu,
                              strides=(2, 2), padding="SAME")
         x = common_layers.layer_norm(x)
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 242779bb4..c89f91d5b 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -512,7 +512,7 @@ def combine_training_data(problem, final_data_dir, old_data_dirs,
 @registry.register_hparams
 def rl_modelrl_base():
   return tf.contrib.training.HParams(
-      epochs=2,
+      epochs=6,
       # Total frames used for training. This will be distributed evenly across
       # hparams.epochs.
       num_real_env_frames=100000,
@@ -524,7 +524,7 @@ def rl_modelrl_base():
       simulated_env_generator_num_steps=2000,
       simulation_random_starts=True,
       intrinsic_reward_scale=0.,
-      ppo_epochs_num=1000,  # This should be enough to see something
+      ppo_epochs_num=2000,  # This should be enough to see something
       # Our simulated envs do not know how to reset.
       # You should set ppo_time_limit to the value you believe that
       # the simulated env produces a reasonable output.

From 8d63cc381cd47df68f663fc823dc5c2bb42a4c4e Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Thu, 30 Aug 2018 21:37:15 -0700
Subject: [PATCH 0711/2720] Make EMA work with VQ gating, by removing residual
 quantization.

PiperOrigin-RevId: 211032799
---
 tensor2tensor/layers/vq_discrete.py | 225 ++++++++++++++--------------
 tensor2tensor/utils/expert_utils.py |   5 +-
 2 files changed, 113 insertions(+), 117 deletions(-)

diff --git a/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index 4b0180269..9477708f3 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -13,15 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Clean discrete bottleneck as in https://arxiv.org/abs/1805.11063."""
-
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 from functools import partial
-
 from tensor2tensor.layers import common_layers
-
 import tensorflow as tf
+from tensorflow.python.training import moving_averages
 
 
 class DiscreteBottleneck(object):
@@ -39,22 +37,25 @@ def __init__(self, hparams):
     self.hparams.block_v_size = 2**(
         self.hparams.z_size_per_residual / self.hparams.num_blocks)
     self.hparams.block_v_size = int(self.hparams.block_v_size)
-    # TODO(avaswani): Figure out why tf.get_variable doesn't work with assign
-    self.hparams.means = tf.get_variable(
+    self.means = tf.get_variable(
         name="means",
         shape=[
-            self.hparams.num_residuals, self.hparams.num_blocks,
-            self.hparams.block_v_size, self.hparams.block_dim
-            ],
-        initializer=tf.uniform_unit_scaling_initializer())
-    tf.logging.info("means = {}".format(self.hparams.means))
-    tf.logging.info("Done creating means")
+            self.hparams.num_blocks, self.hparams.block_v_size,
+            self.hparams.block_dim
+        ],
+        initializer=tf.initializers.variance_scaling(distribution="uniform"))
 
     # Create the shadow variables if we are using EMA
-    self.hparams.ema_count = None
-    self.hparams.ema_means = None
     if self.hparams.ema:
-      raise NotImplementedError("ema updates not implemented")
+      self.ema_count = tf.get_variable(
+          "ema_count", [self.hparams.num_blocks, self.hparams.block_v_size],
+          initializer=tf.constant_initializer(0),
+          trainable=False)
+      with tf.colocate_with(self.means):
+        self.ema_means = tf.get_variable(
+            "ema_means",
+            initializer=self.means.initialized_value(),
+            trainable=False)
 
   def slice_hidden(self, x):
     """Slice encoder hidden state into block_dim.
@@ -182,65 +183,47 @@ def int_to_bit(self, x_int, num_bits, base=2):
     res = tf.concat(x_labels, axis=-1)
     return tf.to_float(res)
 
-  def embed(self, x, scope="bottleneck"):
+  def embed(self, x):
     """Embedding function that takes discrete latent and returns embedding.
 
     Args:
         x: Input to the discretization bottleneck.
-        scope: Scope name of the function.
-
     Returns:
         Continuous embedding to be passed on to the decoder.
 
     Raises:
         ValueError: For unknown or missing arguments.
     """
-    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
-      shape_x = common_layers.shape_list(x)
-      x_flat = tf.reshape(x, [-1, 1])
-      c = self.int_to_bit(x_flat, num_bits=self.hparams.z_size, base=2)
-      shape = common_layers.shape_list(c)
-      new_shape = shape
-      new_shape[-1] = self.hparams.num_residuals
-      new_shape.append(self.hparams.num_blocks)
-      new_shape.append(
-          int(self.hparams.z_size /
-              (self.hparams.num_residuals * self.hparams.num_blocks)))
-      c = tf.to_int32(tf.reshape(c, shape=new_shape))
-      h1_shape = shape_x
-      h1_shape.append(self.hparams.hidden_size)
-      h1 = tf.zeros(dtype=tf.float32, shape=h1_shape)
-      for i in range(self.hparams.num_residuals):
-        c_residual = self.bit_to_int(
-            c[:, :, i, :, :],
-            num_bits=int(
-                self.hparams.z_size /
-                (self.hparams.num_residuals * self.hparams.num_blocks)),
-            base=2)
-        c_hot = tf.one_hot(c_residual, depth=self.hparams.block_v_size, axis=-1)
-        c_hot_flat = tf.reshape(
-            c_hot,
-            shape=[-1, self.hparams.num_blocks, self.hparams.block_v_size])
-        h1_residual = tf.matmul(
-            tf.transpose(c_hot_flat, perm=[1, 0, 2]), self.hparams.means[i])
-        h1_residual = tf.transpose(h1_residual, perm=[1, 0, 2])
-        h1_residual = tf.reshape(h1_residual, shape=h1_shape)
-        h1 += h1_residual
-
-      # Add Gaussian noise
-      h1_shape[0] = self.hparams.batch_size
-      h2 = tf.layers.dense(
-          tf.nn.relu(h1), self.hparams.filter_size, name="vch2")
-      res = tf.layers.dense(
-          tf.nn.relu(h2), self.hparams.hidden_size, name="vcfin")
-      return res
-
-  def discrete_bottleneck(self, x, scope="bottleneck"):
+    shape_x = common_layers.shape_list(x)
+    x_flat = tf.reshape(x, [-1, 1])
+    c = self.int_to_bit(x_flat, num_bits=self.hparams.z_size, base=2)
+    shape = common_layers.shape_list(c)
+    new_shape = shape
+    new_shape.append(self.hparams.num_blocks)
+    new_shape.append(int(self.hparams.z_size / self.hparams.num_blocks))
+    c = tf.to_int32(tf.reshape(c, shape=new_shape))
+    h1_shape = shape_x
+    h1_shape.append(self.hparams.hidden_size)
+    h1 = tf.zeros(dtype=tf.float32, shape=h1_shape)
+    c_int = self.bit_to_int(
+        c, num_bits=int(self.hparams.z_size / self.hparams.num_blocks), base=2)
+    c_hot = tf.one_hot(c_int, depth=self.hparams.block_v_size, axis=-1)
+    c_hot_flat = tf.reshape(
+        c_hot, shape=[-1, self.hparams.num_blocks, self.hparams.block_v_size])
+    h1 = tf.matmul(tf.transpose(c_hot_flat, perm=[1, 0, 2]), self.means)
+    h1 = tf.transpose(h1, perm=[1, 0, 2])
+    h1 = tf.reshape(h1, shape=h1_shape)
+    h1_shape[0] = self.hparams.batch_size
+    h2 = tf.layers.dense(tf.nn.relu(h1), self.hparams.filter_size, name="vch2")
+    res = tf.layers.dense(
+        tf.nn.relu(h2), self.hparams.hidden_size, name="vcfin")
+    return res
+
+  def discrete_bottleneck(self, x):
     """Discretization bottleneck for latent variables.
 
     Args:
         x: Input to the discretization bottleneck.
-        scope: Scope of the function.
 
     Returns:
         Embedding to pass to the decoder, discrete latent, loss, and the
@@ -253,58 +236,72 @@ def discrete_bottleneck(self, x, scope="bottleneck"):
         ema_count or ema_means is None if we are using ema, or unknown
         args.
     """
-    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
-      x_reshaped = self.slice_hidden(x)
-      x_res = x_reshaped
-      x_means_hot = []
-      x_means = 0
-      loss = 0
-      for i in range(self.hparams.num_residuals):
-        x_means_hot_res, x_means_res, q_loss_res, e_loss_res = \
-            self.embedding_lookup(x_reshaped, self.hparams.means[i])
-
-        # TODO(avaswani,nikip,aurkor): Implement ema
-        if self.hparams.ema:
-          raise NotImplementedError("ema updates not implemented")
-        else:
-          loss += q_loss_res + self.hparams.beta * e_loss_res
-
-        # Update the residuals
-        x_res -= x_means_res
-        x_means += x_means_res
-        x_means_hot.append(x_means_hot_res)
-
-      # Get the discrete latent representation
-      x_means_hot = tf.stack(x_means_hot, axis=1)
-      x_means_idx = tf.argmax(x_means_hot, axis=-1)
-
-      # Get the binary representation
-      num_bits = int(self.hparams.z_size //
-                     (self.hparams.num_blocks * self.hparams.num_residuals))
-      x_means_bits = self.int_to_bit(x_means_idx, num_bits=num_bits, base=2)
-      shape = common_layers.shape_list(x_means_bits)
-      new_shape = shape[:-2]
-      new_shape[0] = -1
-      new_shape[-1] = self.hparams.z_size
-      x_means_bits = tf.reshape(x_means_bits, new_shape)
-      x_discrete = self.bit_to_int(
-          tf.to_int32(x_means_bits), num_bits=self.hparams.z_size, base=2)
-
-      # Reshape x_discrete
-      shape_x = common_layers.shape_list(x)
-      shape_discrete = shape_x[:-1]
-      x_discrete = tf.reshape(x_discrete, shape_discrete)
-      x_means = tf.reshape(x_means, shape=shape_x)
-      h1 = x + tf.stop_gradient(x_means - x)
-
-      h2 = tf.layers.dense(
-          tf.nn.relu(h1), self.hparams.filter_size, name="vch2")
-      res = tf.layers.dense(
-          tf.nn.relu(h2), self.hparams.hidden_size, name="vcfin")
-      embed_fn = partial(self.embed, scope=scope)
-      return {
-          "dense": res,
-          "discrete": x_discrete,
-          "loss": loss,
-          "embed": embed_fn
-      }
+    x_reshaped = self.slice_hidden(x)
+    x_means_hot = []
+    x_means = 0
+    loss = 0
+    x_means_hot, x_means, q_loss, e_loss = self.embedding_lookup(
+        x_reshaped, self.means)
+
+    if self.hparams.ema:
+      tf.logging.info("Using EMA with beta = {}".format(self.hparams.beta))
+      updated_ema_count = \
+          moving_averages.assign_moving_average(
+              self.ema_count,
+              tf.reduce_sum(
+                  tf.reshape(
+                      x_means_hot,
+                      shape=[-1, self.hparams.num_blocks,
+                             self.hparams.block_v_size]),
+                  axis=0),
+              self.hparams.decay,
+              zero_debias=False)
+
+      dw = tf.matmul(
+          tf.transpose(x_means_hot, perm=[1, 2, 0]),
+          tf.transpose(x_reshaped, perm=[1, 0, 2]))
+
+      updated_ema_means = \
+          moving_averages.assign_moving_average(
+              self.ema_means, dw, self.hparams.decay,
+              zero_debias=False)
+      n = tf.reduce_sum(updated_ema_count, axis=-1, keep_dims=True)
+      updated_ema_count = ((updated_ema_count + self.hparams.epsilon) / (
+          n + 2**self.hparams.z_size * self.hparams.epsilon) * n)
+      updated_ema_means = updated_ema_means / tf.expand_dims(
+          updated_ema_count, axis=-1)
+
+      with tf.control_dependencies([e_loss]):
+        update_means = tf.assign(self.means, updated_ema_means)
+        with tf.control_dependencies([update_means]):
+          loss += self.hparams.beta * e_loss
+    else:
+      # Use a gradient based loss for learning the cluster centers
+      loss += q_loss + self.hparams.beta * e_loss
+
+    # Get the discrete latent representation
+    x_means_idx = tf.argmax(x_means_hot, axis=-1)
+
+    # Get the binary representation
+    num_bits = int(self.hparams.z_size // self.hparams.num_blocks)
+    x_means_bits = self.int_to_bit(x_means_idx, num_bits=num_bits, base=2)
+    x_discrete = self.bit_to_int(
+        tf.to_int32(x_means_bits), num_bits=self.hparams.z_size, base=2)
+
+    # Reshape x_discrete
+    shape_x = common_layers.shape_list(x)
+    shape_discrete = shape_x[:-1]
+    x_discrete = tf.reshape(x_discrete, shape_discrete)
+    x_means = tf.reshape(x_means, shape=shape_x)
+    h1 = x + tf.stop_gradient(x_means - x)
+
+    h2 = tf.layers.dense(tf.nn.relu(h1), self.hparams.filter_size, name="vch2")
+    res = tf.layers.dense(
+        tf.nn.relu(h2), self.hparams.hidden_size, name="vcfin")
+    embed_fn = partial(self.embed)
+    return {
+        "dense": res,
+        "discrete": x_discrete,
+        "loss": loss,
+        "embed": embed_fn
+    }
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 489ddaf3b..0c8e26ef6 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -478,7 +478,7 @@ def vq_gating(x,
     hparams.z_size = int(math.log(num_experts, 2))
     hparams.hidden_size = input_size
     hparams.top_k = k
-    d = bneck.discrete_bottleneck(inputs, scope=name)
+    d = bneck.discrete_bottleneck(inputs)
     centroids = None
     exp_discrete = d["discrete"]
     embed_lookup = d["embed"]
@@ -1095,7 +1095,7 @@ def local_moe(x,
       training loss of the model.  The backpropagation of this loss
       encourages all experts to be approximately equally used across a batch.
   """
-
+  bneck = DiscreteBottleneck(hparams)
   with tf.variable_scope(name, default_name="local_moe"):
     centroids = None
     x_flat = flatten_all_but_last(x)
@@ -1116,7 +1116,6 @@ def local_moe(x,
     else:
       assert hparams.gating_type == "vq"
       tf.logging.info("Using VQ gating")
-      bneck = DiscreteBottleneck(hparams)
       gates, loss, centroids = vq_gating(
           x_flat, num_experts, k, bneck, hparams=hparams)
     loss *= loss_coef

From ab2f0a6ca945e16919bc734dee214f76638cffd5 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 31 Aug 2018 09:23:14 -0700
Subject: [PATCH 0712/2720] internal merge of PR #1025

PiperOrigin-RevId: 211095470
---
 tensor2tensor/rl/collect.py                  |  8 +--
 tensor2tensor/rl/envs/batch_env.py           | 16 ++---
 tensor2tensor/rl/envs/in_graph_batch_env.py  | 25 ++++++-
 tensor2tensor/rl/envs/py_func_batch_env.py   | 19 ++---
 tensor2tensor/rl/envs/simulated_batch_env.py | 33 +++++----
 tensor2tensor/rl/envs/tf_atari_wrappers.py   | 75 ++++++++++++--------
 tensor2tensor/rl/envs/utils.py               | 24 ++++---
 7 files changed, 125 insertions(+), 75 deletions(-)

diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index f32fca858..bde467969 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -33,10 +33,11 @@ def _rollout_metadata(batch_env):
   batch_size = [batch_env_shape[0]]
   shapes_types_names = [
       # TODO(piotrmilos): possibly retrieve the observation type for batch_env
-      (batch_size + batch_env_shape[1:], tf.float32, "observation"),
+      (batch_size + batch_env_shape[1:], batch_env.observ_dtype, "observation"),
       (batch_size, tf.float32, "reward"),
       (batch_size, tf.bool, "done"),
-      (batch_size + batch_env.action_shape, batch_env.action_dtype, "action"),
+      (batch_size + list(batch_env.action_shape), batch_env.action_dtype,
+       "action"),
       (batch_size, tf.float32, "pdf"),
       (batch_size, tf.float32, "value_function"),
   ]
@@ -57,8 +58,7 @@ def __init__(self, batch_env):
     self.speculum = tf.FIFOQueue(infinity, shapes=shapes, dtypes=dtypes)
     observs_shape = batch_env.observ.shape
     # TODO(piotrmilos): possibly retrieve the observation type for batch_env
-    observ_dtype = tf.float32
-    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
+    self._observ = tf.Variable(tf.zeros(observs_shape, self.observ_dtype),
                                trainable=False)
 
   def simulate(self, action):
diff --git a/tensor2tensor/rl/envs/batch_env.py b/tensor2tensor/rl/envs/batch_env.py
index bf433066f..230bd3a43 100644
--- a/tensor2tensor/rl/envs/batch_env.py
+++ b/tensor2tensor/rl/envs/batch_env.py
@@ -43,11 +43,12 @@ def __init__(self, envs, blocking):
     """
     self._envs = envs
     self._blocking = blocking
-    observ_space = self._envs[0].observation_space
-    if not all(env.observation_space == observ_space for env in self._envs):
+    self.observ_space = self._envs[0].observation_space
+    if not all(env.observation_space == self.observ_space
+               for env in self._envs):
       raise ValueError('All environments must use the same observation space.')
-    action_space = self._envs[0].action_space
-    if not all(env.action_space == action_space for env in self._envs):
+    self.action_space = self._envs[0].action_space
+    if not all(env.action_space == self.action_space for env in self._envs):
       raise ValueError('All environments must use the same observation space.')
 
   def __len__(self):
@@ -96,8 +97,8 @@ def step(self, actions):
       transitions = [transition() for transition in transitions]
     observs, rewards, dones, infos = zip(*transitions)
 
+    observ = np.stack(observs).astype(self.observ_space.dtype)
     # TODO(piotrmilos): Do we really want cast to float32
-    observ = np.stack(observs).astype(np.float32)
     reward = np.stack(rewards).astype(np.float32)
     done = np.stack(dones)
     info = tuple(infos)
@@ -119,9 +120,8 @@ def reset(self, indices=None):
     else:
       observs = [self._envs[index].reset(blocking=False) for index in indices]
       observs = [observ() for observ in observs]
-    observ = np.stack(observs)
-    # TODO(piotrmilos): Do we really want this?
-    observ = observ.astype(np.float32)
+    observ = np.stack(observs).astype(self.observ_space.dtype)
+
     return observ
 
   def close(self):
diff --git a/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index 56e13878d..2385369e7 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -21,6 +21,8 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.rl.envs import utils
+
 import tensorflow as tf
 
 
@@ -28,6 +30,10 @@ class InGraphBatchEnv(object):
   """Abstract class for batch of environments inside the TensorFlow graph.
   """
 
+  def __init__(self, observ_space, action_space):
+    self.observ_space = observ_space
+    self.action_space = action_space
+
   def __getattr__(self, name):
     """Forward unimplemented attributes to one of the original environments.
 
@@ -71,7 +77,24 @@ def reset(self, indices=None):
     """
     return tf.cond(
         tf.cast(tf.reduce_sum(indices + 1), tf.bool),
-        lambda: self._reset_non_empty(indices), lambda: 0.0)
+        lambda: self._reset_non_empty(indices),
+        lambda: tf.cast(0, self.observ_dtype))
+
+  @property
+  def observ_dtype(self):
+    return utils.parse_dtype(self.observ_space)
+
+  @property
+  def observ_shape(self):
+    return utils.parse_shape(self.observ_space)
+
+  @property
+  def action_dtype(self):
+    return utils.parse_dtype(self.action_space)
+
+  @property
+  def action_shape(self):
+    return utils.parse_shape(self.action_space)
 
   @property
   def observ(self):
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index 76916d991..ac1ad1dd5 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -21,7 +21,6 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.rl.envs import utils
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
 import tensorflow as tf
 
@@ -40,14 +39,13 @@ def __init__(self, batch_env):
     Args:
       batch_env: Batch environment.
     """
+    super(PyFuncBatchEnv, self).__init__(batch_env.observation_space,
+                                         batch_env.action_space)
     self._batch_env = batch_env
-    observ_shape = utils.parse_shape(self._batch_env.observation_space)
-    observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
-    self.action_shape = list(utils.parse_shape(self._batch_env.action_space))
-    self.action_dtype = utils.parse_dtype(self._batch_env.action_space)
     with tf.variable_scope('env_temporary'):
       self._observ = tf.Variable(
-          tf.zeros((len(self._batch_env),) + observ_shape, observ_dtype),
+          tf.zeros((len(self._batch_env),) + self.observ_shape,
+                   self.observ_dtype),
           name='observ', trainable=False)
 
   def __getattr__(self, name):
@@ -86,11 +84,9 @@ def simulate(self, action):
     with tf.name_scope('environment/simulate'):
       if action.dtype in (tf.float16, tf.float32, tf.float64):
         action = tf.check_numerics(action, 'action')
-      observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
       observ, reward, done = tf.py_func(
           lambda a: self._batch_env.step(a)[:3], [action],
-          [observ_dtype, tf.float32, tf.bool], name='step')
-      observ = tf.check_numerics(observ, 'observ')
+          [self.observ_dtype, tf.float32, tf.bool], name='step')
       reward = tf.check_numerics(reward, 'reward')
       reward.set_shape((len(self),))
       done.set_shape((len(self),))
@@ -106,10 +102,9 @@ def _reset_non_empty(self, indices):
     Returns:
       Batch tensor of the new observations.
     """
-    observ_dtype = utils.parse_dtype(self._batch_env.observation_space)
     observ = tf.py_func(
-        self._batch_env.reset, [indices], observ_dtype, name='reset')
-    observ = tf.check_numerics(observ, 'observ')
+        self._batch_env.reset, [indices], self.observ_dtype, name='reset')
+    observ.set_shape(indices.get_shape().concatenate(self.observ_shape))
     with tf.control_dependencies([
         tf.scatter_update(self._observ, indices, observ)]):
       return tf.identity(observ)
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 99f2d0f3a..395f8f34a 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -23,7 +23,7 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.rl.envs import in_graph_batch_env
-from tensor2tensor.rl.envs.utils import get_action_space
+from tensor2tensor.rl.envs import utils
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
@@ -39,20 +39,21 @@
 class HistoryBuffer(object):
   """History Buffer."""
 
-  def __init__(self, input_dataset, length):
+  def __init__(self, input_dataset, length, observ_dtype):
     self.input_data_iterator = (
         input_dataset.batch(length).make_initializable_iterator())
     self.length = length
+    self._observ_dtype = observ_dtype
     initial_frames = self.get_initial_observations()
     initial_shape = [length] + common_layers.shape_list(initial_frames)[1:]
-    self._history_buff = tf.Variable(tf.zeros(initial_shape, tf.float32),
+    self._history_buff = tf.Variable(tf.zeros(initial_shape, observ_dtype),
                                      trainable=False)
 
   def initialize(self, sess):
     sess.run(self.input_data_iterator.initializer)
 
   def get_initial_observations(self):
-    return tf.cast(self.input_data_iterator.get_next(), tf.float32)
+    return tf.cast(self.input_data_iterator.get_next(), self._observ_dtype)
 
   def get_all_elements(self):
     return self._history_buff.read_value()
@@ -100,8 +101,16 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   def __init__(self, environment_spec, length):
     """Batch of environments inside the TensorFlow graph."""
 
-    self.length = length
+    observ_space = utils.get_observation_space(environment_spec)
     initial_frames_problem = environment_spec.initial_frames_problem
+    observ_shape = (initial_frames_problem.frame_height,
+                    initial_frames_problem.frame_width,
+                    initial_frames_problem.num_channels)
+    observ_space.shape = observ_shape
+    action_space = utils.get_action_space(environment_spec)
+    super(SimulatedBatchEnv, self).__init__(observ_space, action_space)
+
+    self.length = length
     self._min_reward = initial_frames_problem.min_reward
     self._num_frames = environment_spec.video_num_input_frames
     self._intrinsic_reward_scale = environment_spec.intrinsic_reward_scale
@@ -112,8 +121,6 @@ def __init__(self, environment_spec, length):
     self._model = registry.model(FLAGS.model)(
         model_hparams, tf.estimator.ModeKeys.PREDICT)
 
-    _, self.action_shape, self.action_dtype = get_action_space(environment_spec)
-
     hparams = HParams(video_num_input_frames=
                       environment_spec.video_num_input_frames,
                       video_num_target_frames=
@@ -133,12 +140,11 @@ def __init__(self, environment_spec, length):
                                                hparams=hparams).take(1)
 
     dataset = dataset.map(lambda x: x["inputs"]).repeat()
-    self.history_buffer = HistoryBuffer(dataset, self.length)
+    self.history_buffer = HistoryBuffer(dataset, self.length, self.observ_dtype)
 
-    shape = (self.length, initial_frames_problem.frame_height,
-             initial_frames_problem.frame_width,
-             initial_frames_problem.num_channels)
-    self._observ = tf.Variable(tf.zeros(shape, tf.float32), trainable=False)
+    self._observ = tf.Variable(
+        tf.zeros((len(self),) + observ_shape, self.observ_dtype),
+        trainable=False)
 
   def initialize(self, sess):
     self.history_buffer.initialize(sess)
@@ -160,7 +166,8 @@ def simulate(self, action):
             {"inputs": history, "input_action": actions})
         self._model.hparams.video_num_target_frames = hparams_target_frames
 
-      observ = tf.to_float(tf.squeeze(model_output["targets"], axis=1))
+      observ = tf.cast(tf.squeeze(model_output["targets"], axis=1),
+                       self.observ_dtype)
 
       reward = tf.to_float(model_output["target_reward"])
       reward = tf.reshape(reward, shape=(self.length,)) + self._min_reward
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index fcf744f0c..4df4df6bd 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -33,10 +33,10 @@ class WrapperBase(InGraphBatchEnv):
   """Base wrapper class."""
 
   def __init__(self, batch_env):
+    super(WrapperBase, self).__init__(
+        batch_env.observ_space, batch_env.action_space)
     self._length = len(batch_env)
     self._batch_env = batch_env
-    self.action_shape = batch_env.action_shape
-    self.action_dtype = batch_env.action_dtype
 
   def initialize(self, sess):
     """Initializations to be run once the tf.Session is available."""
@@ -83,8 +83,7 @@ def __init__(self, batch_env, skip=4):
     super(MaxAndSkipWrapper, self).__init__(batch_env)
     self.skip = skip
     observs_shape = batch_env.observ.shape
-    observ_dtype = tf.float32
-    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
+    self._observ = tf.Variable(tf.zeros(observs_shape, self.observ_dtype),
                                trainable=False)
 
   def simulate(self, action):
@@ -120,16 +119,19 @@ class StackAndSkipWrapper(WrapperBase):
   def __init__(self, batch_env, skip=4):
     super(StackAndSkipWrapper, self).__init__(batch_env)
     self.skip = skip
-    self._observ = None
-    self.old_shape = batch_env.observ.shape.as_list()
-    observs_shape = self.old_shape[:-1] + [self.old_shape[-1] * self.skip]
-    observ_dtype = tf.float32
-    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
-                               trainable=False)
+    self.old_shape = self._batch_env.observ_shape
+    self._observ = tf.Variable(
+        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
+        trainable=False)
+
+  @property
+  def observ_shape(self):
+    return self.old_shape[:-1] + (self.old_shape[-1] * self.skip,)
 
   def simulate(self, action):
     with tf.name_scope("environment/simulate"):  # Do we need this?
-      initializer = (tf.zeros(self.old_shape, dtype=tf.float32),
+      initializer = (tf.zeros((len(self),) + self.old_shape,
+                              dtype=self.observ_dtype),
                      tf.fill((len(self),), 0.0), tf.fill((len(self),), False))
 
       def not_done_step(a, _):
@@ -156,7 +158,8 @@ def _reset_non_empty(self, indices):
     # pylint: enable=protected-access
     inx = tf.concat(
         [
-            tf.ones(tf.size(tf.shape(new_values)), dtype=tf.int32)[:-1],
+            tf.ones(tf.size(tf.shape(new_values)),
+                    dtype=tf.int32)[:-1],
             [self.skip]
         ],
         axis=0)
@@ -172,11 +175,14 @@ class StackWrapper(WrapperBase):
   def __init__(self, batch_env, history=4):
     super(StackWrapper, self).__init__(batch_env)
     self.history = history
-    self.old_shape = batch_env.observ.shape.as_list()
-    observs_shape = self.old_shape[:-1] + [self.old_shape[-1] * self.history]
-    observ_dtype = tf.float32
-    self._observ = tf.Variable(tf.zeros(observs_shape, observ_dtype),
-                               trainable=False)
+    self.old_shape = batch_env.observ_shape
+    self._observ = tf.Variable(
+        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
+        trainable=False)
+
+  @property
+  def observ_shape(self):
+    return self.old_shape[:-1] + (self.old_shape[-1] * self.history,)
 
   def simulate(self, action):
     reward, done = self._batch_env.simulate(action)
@@ -197,7 +203,8 @@ def _reset_non_empty(self, indices):
     # pylint: enable=protected-access
     inx = tf.concat(
         [
-            tf.ones(tf.size(tf.shape(new_values)), dtype=tf.int32)[:-1],
+            tf.ones(tf.size(tf.shape(new_values)),
+                    dtype=tf.int32)[:-1],
             [self.history]
         ],
         axis=0)
@@ -213,19 +220,23 @@ class AutoencoderWrapper(WrapperBase):
 
   def __init__(self, batch_env):
     super(AutoencoderWrapper, self).__init__(batch_env)
-    batch_size, height, width, _ = self._batch_env.observ.get_shape().as_list()
-    ae_height = int(math.ceil(height / self.autoencoder_factor))
-    ae_width = int(math.ceil(width / self.autoencoder_factor))
-    ae_channels = 24  # TODO(piotrmilos): make it better
-    observ_shape = (batch_size, ae_height, ae_width, ae_channels)
     self._observ = self._observ = tf.Variable(
-        tf.zeros(observ_shape, tf.float32), trainable=False)
+        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
+        trainable=False)
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       autoencoder_hparams = autoencoders.autoencoder_discrete_pong()
       self.autoencoder_model = autoencoders.AutoencoderOrderedDiscrete(
           autoencoder_hparams, tf.estimator.ModeKeys.EVAL)
     self.autoencoder_model.set_mode(tf.estimator.ModeKeys.EVAL)
 
+  @property
+  def observ_shape(self):
+    height, width, _ = self._batch_env.observ_shape
+    ae_height = int(math.ceil(height / self.autoencoder_factor))
+    ae_width = int(math.ceil(width / self.autoencoder_factor))
+    ae_channels = 24  # TODO(piotrmilos): make it better
+    return (ae_height, ae_width, ae_channels)
+
   @property
   def autoencoder_factor(self):
     """By how much to divide sizes when using autoencoders."""
@@ -237,6 +248,7 @@ def simulate(self, action):
     with tf.control_dependencies([reward, done]):
       with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
         ret = self.autoencoder_model.encode(self._batch_env.observ)
+        ret = tf.cast(ret, self.observ_dtype)
         assign_op = self._observ.assign(ret)
         with tf.control_dependencies([assign_op]):
           return tf.identity(reward), tf.identity(done)
@@ -245,6 +257,7 @@ def _reset_non_empty(self, indices):
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       new_values = self._batch_env._reset_non_empty(indices)  # pylint: disable=protected-access
       ret = self.autoencoder_model.encode(new_values)
+      ret = tf.cast(ret, self.observ_dtype)
       assign_op = tf.scatter_update(self._observ, indices, ret)
       with tf.control_dependencies([assign_op]):
         return tf.gather(self.observ, indices)
@@ -255,14 +268,16 @@ class IntToBitWrapper(WrapperBase):
 
   def __init__(self, batch_env):
     super(IntToBitWrapper, self).__init__(batch_env)
-    batch_size, height, width, channels = \
-      self._batch_env.observ.get_shape().as_list()
-    # We treat each channel as 8-bit integer to be expanded to 8 channels
-    self.observ_shape = (height, width, channels*8)
     self._observ = self._observ = tf.Variable(
-        tf.zeros((batch_size,) + self.observ_shape, tf.float32),
+        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
         trainable=False)
 
+  @property
+  def observ_shape(self):
+    height, width, channels = self._batch_env.observ_shape
+    # We treat each channel as 8-bit integer to be expanded to 8 channels
+    return (height, width, channels*8)
+
   def simulate(self, action):
     action = tf.Print(action, [action], message="action=", summarize=200)
 
@@ -272,6 +287,7 @@ def simulate(self, action):
       with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
         unpacked = discretization.int_to_bit(self._batch_env.observ, 8)
         unpacked = tf.reshape(unpacked, (-1,)+self.observ_shape)
+        unpacked = tf.cast(unpacked, self.observ_dtype)
         assign_op = self._observ.assign(unpacked)
         with tf.control_dependencies([assign_op]):
           return tf.identity(reward), tf.identity(done)
@@ -282,6 +298,7 @@ def _reset_non_empty(self, indices):
     new_values_unpacked = discretization.int_to_bit(new_values, 8)
     new_values_unpacked = tf.reshape(new_values_unpacked, (-1,)
                                      +self.observ_shape)
+    new_values_unpacked = tf.cast(new_values_unpacked, self.observ_dtype)
     # pylint: enable=protected-access
     assign_op = tf.scatter_update(self._observ, indices, new_values_unpacked)
     with tf.control_dependencies([assign_op]):
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index b23dc358e..badd48485 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -65,8 +65,20 @@ def _reset(self, **kwargs):
     return self._last_returned[0]
 
 
+def get_observation_space(environment_spec):
+  """Get observation space associated with environment spec.
+
+  Args:
+     environment_spec:  EnvironmentSpec object
+
+  Returns:
+    OpenAi Gym observation space
+  """
+  return environment_spec.env_lambda().observation_space
+
+
 def get_action_space(environment_spec):
-  """Get action spece associated with environment spec.
+  """Get action space associated with environment spec.
 
   Args:
      environment_spec:  EnvironmentSpec object
@@ -74,11 +86,7 @@ def get_action_space(environment_spec):
   Returns:
     OpenAi Gym action space
   """
-  action_space = environment_spec.env_lambda().action_space
-  action_shape = list(parse_shape(action_space))
-  action_dtype = parse_dtype(action_space)
-
-  return action_space, action_shape, action_dtype
+  return environment_spec.env_lambda().action_space
 
 
 def get_policy(observations, hparams):
@@ -92,7 +100,7 @@ def get_policy(observations, hparams):
     Tensor with policy and value function output
   """
   policy_network_lambda = hparams.policy_network
-  action_space, _, _ = get_action_space(hparams.environment_spec)
+  action_space = get_action_space(hparams.environment_spec)
   return policy_network_lambda(action_space, hparams, observations)
 
 
@@ -124,5 +132,5 @@ def parse_dtype(space):
   if isinstance(space, gym.spaces.Discrete):
     return tf.int32
   if isinstance(space, gym.spaces.Box):
-    return tf.float32
+    return tf.as_dtype(space.dtype)
   raise NotImplementedError()

From 62331ebc97a1a3438157109d6c39f659b3a08540 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 31 Aug 2018 10:13:56 -0700
Subject: [PATCH 0713/2720] internal Conditional Glow, first pass and adds
 boilerplate code for extending glow to video models.

PiperOrigin-RevId: 211102670
---
 tensor2tensor/models/research/glow.py     | 24 +++++++++++++++++++----
 tensor2tensor/models/research/glow_ops.py | 23 ++++++++++++++--------
 2 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 629a882a8..90e787577 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -45,7 +45,7 @@ def glow_hparams():
   hparams.add_hparam("n_bits_x", 8)
   hparams.add_hparam("depth", 32)
   hparams.add_hparam("affine_coupling_width", 512)
-  hparams.add_hparam("learn_prior", True)
+  hparams.add_hparam("top_prior", "single_conv")
   return hparams
 
 
@@ -74,7 +74,11 @@ def preprocess(self, x):
 
   def scale(self, x):
     """Scale x from -0.5 - 0.5 to 0 - 255."""
-    x = (x + 0.5) * 2**self.hparams.n_bits_x
+    x = tf.where(tf.is_nan(x), tf.ones_like(x), x)
+    x = tf.where(tf.is_inf(x), tf.ones_like(x), x)
+    x = tf.clip_by_value(x, -0.5, 0.5)
+    x += 0.5
+    x = x * 2**self.hparams.n_bits_x
     return tf.cast(tf.clip_by_value(x, 0, 255), dtype=tf.uint8)
 
   @property
@@ -97,6 +101,18 @@ def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
 
     return self.scale(predictions)
 
+  def top_prior(self, z):
+    """Objective based on the prior over latent z.
+
+    Args:
+      z: 4-D Tensor, (batch_size, height, width, num_channels)
+    Returns:
+      objective: float, log-likelihood of z under the prior.
+      dist: instance of tf.distributions.Normal, prior distribution.
+    """
+    return glow_ops.top_prior(
+        "top_prior", z, learn_prior=self.hparams.top_prior)
+
   def body(self, features):
     x = features["inputs"]
 
@@ -124,8 +140,8 @@ def body(self, features):
           "codec", x, self.hparams, eps=None, reverse=False)
       objective += encoder_objective
 
-      prior_objective, prior_dist = glow_ops.top_prior(
-          "top_prior", self.z, learn_prior=self.hparams.learn_prior)
+      prior_objective, prior_dist = self.top_prior(self.z)
+      tf.summary.scalar("top_prior", tf.reduce_mean(prior_objective))
       self.z_sample = prior_dist.sample()
       objective += prior_objective
 
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index d1ca6acfe..e23e1433f 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -293,7 +293,7 @@ def conv2d(name, x, output_channels, filter_size=None, stride=None,
   _, _, _, in_channels = common_layers.shape_list(x)
 
   filter_shape = filter_size + [in_channels, output_channels]
-  stride_shape = [1, 1] + stride
+  stride_shape = [1] + stride + [1]
 
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
 
@@ -526,25 +526,32 @@ def revnet(name, x, hparams, reverse=True):
 
 
 @add_arg_scope
-def top_prior(name, x, learn_prior=False):
+def top_prior(name, x, learn_prior="normal"):
   """Log probability of x being gaussian.
 
   Args:
     name: variable scope
     x: input, 4-D Tensor shape=(batch_size, width, height, channels)
-    learn_prior: If set to true, then the mean and the standard deviation
-                 are the output of a single conv layer initialized with
-                 zeros. Otherwise the mean and logstd are zeros and ones
-                 respectively.
+    learn_prior: Possible options are "normal" and "single_conv".
+                 If set to "single_conv", the gaussian is parametrized by a
+                 single convolutional layer whose input are an array of zeros
+                 and initialized such that the mean and std are zero and one.
+                 If set to "normal", the prior is just a Gaussian with zero
+                 mean and unit variance.
   Returns:
     objective: 1-D Tensor shape=(batch_size,) summed across spatial components.
+  Raises:
+    ValueError: If learn_prior not in "normal" or "single_conv"
   """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     h = tf.zeros_like(x)
-    if not learn_prior:
+    if learn_prior == "normal":
       prior_dist = tf.distributions.Normal(h, tf.exp(h))
-    else:
+    elif learn_prior == "single_conv":
       prior_dist = split_prior("top_learn_prior", h)
+    else:
+      raise ValueError("Expected learn_prior to be normal or single_conv "
+                       "got %s" % learn_prior)
     objective = tf.reduce_sum(prior_dist.log_prob(x), axis=[1, 2, 3])
     return objective, prior_dist
 

From c0123b38bcf29b0dd7e0a886a5ee1022aa5fd04f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 31 Aug 2018 12:11:25 -0700
Subject: [PATCH 0714/2720] Add support for test sets in old problems

PiperOrigin-RevId: 211123440
---
 tensor2tensor/bin/t2t_datagen.py | 42 +++++++++++++++++++++++---------
 1 file changed, 30 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index 30e21e1b4..4dfdc4d6f 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -89,40 +89,46 @@
 _SUPPORTED_PROBLEM_GENERATORS = {
     "algorithmic_algebra_inverse": (
         lambda: algorithmic_math.algebra_inverse(26, 0, 2, 100000),
-        lambda: algorithmic_math.algebra_inverse(26, 3, 3, 10000)),
+        lambda: algorithmic_math.algebra_inverse(26, 3, 3, 10000),
+        lambda: None),  # test set
     "parsing_english_ptb8k": (
         lambda: wsj_parsing.parsing_token_generator(
             FLAGS.data_dir, FLAGS.tmp_dir, True, 2**13, 2**9),
         lambda: wsj_parsing.parsing_token_generator(
-            FLAGS.data_dir, FLAGS.tmp_dir, False, 2**13, 2**9)),
+            FLAGS.data_dir, FLAGS.tmp_dir, False, 2**13, 2**9),
+        lambda: None),  # test set
     "parsing_english_ptb16k": (
         lambda: wsj_parsing.parsing_token_generator(
             FLAGS.data_dir, FLAGS.tmp_dir, True, 2**14, 2**9),
         lambda: wsj_parsing.parsing_token_generator(
-            FLAGS.data_dir, FLAGS.tmp_dir, False, 2**14, 2**9)),
+            FLAGS.data_dir, FLAGS.tmp_dir, False, 2**14, 2**9),
+        lambda: None),  # test set
     "inference_snli32k": (
         lambda: snli.snli_token_generator(FLAGS.tmp_dir, True, 2**15),
         lambda: snli.snli_token_generator(FLAGS.tmp_dir, False, 2**15),
-    ),
+        lambda: None),  # test set
     "audio_timit_characters_test": (
         lambda: audio.timit_generator(
             FLAGS.data_dir, FLAGS.tmp_dir, True, 1718),
         lambda: audio.timit_generator(
-            FLAGS.data_dir, FLAGS.tmp_dir, False, 626)),
+            FLAGS.data_dir, FLAGS.tmp_dir, False, 626),
+        lambda: None),  # test set
     "audio_timit_tokens_8k_test": (
         lambda: audio.timit_generator(
             FLAGS.data_dir, FLAGS.tmp_dir, True, 1718,
             vocab_filename="vocab.endefr.%d" % 2**13, vocab_size=2**13),
         lambda: audio.timit_generator(
             FLAGS.data_dir, FLAGS.tmp_dir, False, 626,
-            vocab_filename="vocab.endefr.%d" % 2**13, vocab_size=2**13)),
+            vocab_filename="vocab.endefr.%d" % 2**13, vocab_size=2**13),
+        lambda: None),  # test set
     "audio_timit_tokens_32k_test": (
         lambda: audio.timit_generator(
             FLAGS.data_dir, FLAGS.tmp_dir, True, 1718,
             vocab_filename="vocab.endefr.%d" % 2**15, vocab_size=2**15),
         lambda: audio.timit_generator(
             FLAGS.data_dir, FLAGS.tmp_dir, False, 626,
-            vocab_filename="vocab.endefr.%d" % 2**15, vocab_size=2**15)),
+            vocab_filename="vocab.endefr.%d" % 2**15, vocab_size=2**15),
+        lambda: None),  # test set
 }
 
 # pylint: enable=g-long-lambda
@@ -193,19 +199,31 @@ def main(_):
 
 def generate_data_for_problem(problem):
   """Generate data for a problem in _SUPPORTED_PROBLEM_GENERATORS."""
-  training_gen, dev_gen = _SUPPORTED_PROBLEM_GENERATORS[problem]
+  training_gen, dev_gen, test_gen = _SUPPORTED_PROBLEM_GENERATORS[problem]
 
-  num_shards = FLAGS.num_shards or 10
+  num_train_shards = FLAGS.num_shards or 10
   tf.logging.info("Generating training data for %s.", problem)
   train_output_files = generator_utils.train_data_filenames(
-      problem + generator_utils.UNSHUFFLED_SUFFIX, FLAGS.data_dir, num_shards)
+      problem + generator_utils.UNSHUFFLED_SUFFIX, FLAGS.data_dir,
+      num_train_shards)
   generator_utils.generate_files(training_gen(), train_output_files,
                                  FLAGS.max_cases)
+  num_dev_shards = int(num_train_shards * 0.1)
   tf.logging.info("Generating development data for %s.", problem)
   dev_output_files = generator_utils.dev_data_filenames(
-      problem + generator_utils.UNSHUFFLED_SUFFIX, FLAGS.data_dir, 1)
+      problem + generator_utils.UNSHUFFLED_SUFFIX, FLAGS.data_dir,
+      num_dev_shards)
   generator_utils.generate_files(dev_gen(), dev_output_files)
-  all_output_files = train_output_files + dev_output_files
+  num_test_shards = int(num_train_shards * 0.1)
+  test_output_files = []
+  test_gen_data = test_gen()
+  if test_gen_data is not None:
+    tf.logging.info("Generating test data for %s.", problem)
+    test_output_files = generator_utils.test_data_filenames(
+        problem + generator_utils.UNSHUFFLED_SUFFIX, FLAGS.data_dir,
+        num_test_shards)
+    generator_utils.generate_files(test_gen_data, test_output_files)
+  all_output_files = train_output_files + dev_output_files + test_output_files
   generator_utils.shuffle_dataset(all_output_files)
 
 

From faed43975ff2da667583a5ddc8a85957f78bcbc2 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Fri, 31 Aug 2018 13:20:26 -0700
Subject: [PATCH 0715/2720] StackWrapper should initialize with consecutive
 frames, whenever possible.

PiperOrigin-RevId: 211133859
---
 tensor2tensor/data_generators/gym_problems.py |  2 +-
 tensor2tensor/rl/envs/batch_env_factory.py    | 10 ++-
 tensor2tensor/rl/envs/tf_atari_wrappers.py    | 89 +++++--------------
 3 files changed, 30 insertions(+), 71 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 71893f11d..adedf0325 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -65,7 +65,7 @@ def standard_atari_env_spec(env):
 
 def standard_atari_ae_env_spec(env):
   """Parameters of environment specification."""
-  standard_wrappers = [[tf_atari_wrappers.StackAndSkipWrapper, {"skip": 4}],
+  standard_wrappers = [[tf_atari_wrappers.StackWrapper, {"history": 4}],
                        [tf_atari_wrappers.AutoencoderWrapper, {}]]
   env_lambda = None
   if isinstance(env, str):
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index 8f610f42a..52ad1af81 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -81,6 +81,7 @@ class ExternalProcessEnv(object):
   _RESULT = 3
   _EXCEPTION = 4
   _CLOSE = 5
+  _ATTRIBUTE_EXCEPTION = 6
 
   def __init__(self, constructor, xvfb):
     """Step environment in a separate process for lock free parallelism.
@@ -226,6 +227,8 @@ def _receive(self):
     if message == self._EXCEPTION:
       stacktrace = payload
       raise Exception(stacktrace)
+    if message == self._ATTRIBUTE_EXCEPTION:
+      raise AttributeError(payload)
     if message == self._RESULT:
       return payload
     raise KeyError("Received message of unexpected type {}".format(message))
@@ -249,8 +252,11 @@ def _worker(self, constructor, conn):
           break
         if message == self._ACCESS:
           name = payload
-          result = getattr(env, name)
-          conn.send((self._RESULT, result))
+          try:
+            result = getattr(env, name)
+            conn.send((self._RESULT, result))
+          except AttributeError as err:
+            conn.send((self._ATTRIBUTE_EXCEPTION, err.args))
           continue
         if message == self._CALL:
           name, args, kwargs = payload
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 4df4df6bd..39006878d 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -110,65 +110,6 @@ def not_done_step(a, _):
         return tf.identity(simulate_ret[1]), tf.identity(simulate_ret[2])
 
 
-class StackAndSkipWrapper(WrapperBase):
-  """ Stack and skip wrapper.
-      The wrapper works under assumptions that issuing an action
-      to an environment with done=True has not effect.
-  """
-
-  def __init__(self, batch_env, skip=4):
-    super(StackAndSkipWrapper, self).__init__(batch_env)
-    self.skip = skip
-    self.old_shape = self._batch_env.observ_shape
-    self._observ = tf.Variable(
-        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
-        trainable=False)
-
-  @property
-  def observ_shape(self):
-    return self.old_shape[:-1] + (self.old_shape[-1] * self.skip,)
-
-  def simulate(self, action):
-    with tf.name_scope("environment/simulate"):  # Do we need this?
-      initializer = (tf.zeros((len(self),) + self.old_shape,
-                              dtype=self.observ_dtype),
-                     tf.fill((len(self),), 0.0), tf.fill((len(self),), False))
-
-      def not_done_step(a, _):
-        reward, done = self._batch_env.simulate(action)
-        with tf.control_dependencies([reward, done]):
-          r0 = self._batch_env.observ + 0
-          r1 = tf.add(a[1], reward)
-          r2 = tf.logical_or(a[2], done)
-          return (r0, r1, r2)
-
-      simulate_ret = tf.scan(not_done_step, tf.range(self.skip),
-                             initializer=initializer, parallel_iterations=1,
-                             infer_shape=False)
-      observations, rewards, dones = simulate_ret
-      split_observations = tf.split(observations, self.skip, axis=0)
-      split_observations = [tf.squeeze(o, axis=0) for o in split_observations]
-      observation = tf.concat(split_observations, axis=-1)
-      with tf.control_dependencies([self._observ.assign(observation)]):
-        return tf.identity(rewards[-1, ...]), tf.identity(dones[-1, ...])
-
-  def _reset_non_empty(self, indices):
-    # pylint: disable=protected-access
-    new_values = self._batch_env._reset_non_empty(indices)
-    # pylint: enable=protected-access
-    inx = tf.concat(
-        [
-            tf.ones(tf.size(tf.shape(new_values)),
-                    dtype=tf.int32)[:-1],
-            [self.skip]
-        ],
-        axis=0)
-    assign_op = tf.scatter_update(self._observ, indices, tf.tile(
-        new_values, inx))
-    with tf.control_dependencies([assign_op]):
-      return tf.gather(self.observ, indices)
-
-
 class StackWrapper(WrapperBase):
   """ A wrapper which stacks previously seen frames. """
 
@@ -201,15 +142,27 @@ def _reset_non_empty(self, indices):
     # pylint: disable=protected-access
     new_values = self._batch_env._reset_non_empty(indices)
     # pylint: enable=protected-access
-    inx = tf.concat(
-        [
-            tf.ones(tf.size(tf.shape(new_values)),
-                    dtype=tf.int32)[:-1],
-            [self.history]
-        ],
-        axis=0)
-    assign_op = tf.scatter_update(self._observ, indices, tf.tile(
-        new_values, inx))
+    history_buffer = getattr(self._batch_env, "history_buffer", None)
+    if history_buffer:
+      # Using history buffer frames for initialization, if they are available.
+      # This assumes that wrappers don't alter the observations.
+      with tf.control_dependencies([new_values]):
+        initial_frames = history_buffer.get_all_elements()
+        # Transpose to [batch, height, width, history, channels] and merge
+        # history and channels into one dimension.
+        initial_frames = tf.transpose(initial_frames, [0, 2, 3, 1, 4])
+        initial_frames = tf.reshape(initial_frames,
+                                    (len(self),) + self.observ_shape)
+    else:
+      inx = tf.concat(
+          [
+              tf.ones(tf.size(tf.shape(new_values)),
+                      dtype=tf.int32)[:-1],
+              [self.history]
+          ],
+          axis=0)
+      initial_frames = tf.tile(new_values, inx)
+    assign_op = tf.scatter_update(self._observ, indices, initial_frames)
     with tf.control_dependencies([assign_op]):
       return tf.gather(self.observ, indices)
 

From 41c84a57c8a43109330fcf203092170f40b6969f Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 31 Aug 2018 16:10:59 -0700
Subject: [PATCH 0716/2720] Minor clean up to estimator method docstrings.

PiperOrigin-RevId: 211160106
---
 tensor2tensor/utils/t2t_model.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index ab4bcab5b..f91600bfa 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1345,7 +1345,7 @@ def initialize_from_ckpt(self, ckpt_dir):
     tf.train.init_from_checkpoint(ckpt_dir, variable_map)
 
   def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
-    """Construct EstimatorSpec for TRAIN mode."""
+    """Constructs `tf.estimator.EstimatorSpec` for TRAIN (training) mode."""
     train_op = self.optimize(loss, num_async_replicas=num_async_replicas,
                              use_tpu=use_tpu)
 
@@ -1365,7 +1365,7 @@ def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
           tf.estimator.ModeKeys.TRAIN, loss=loss, train_op=train_op)
 
   def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
-    """Construct EstimatorSpec for EVAL mode."""
+    """Constructs `tf.estimator.EstimatorSpec` for EVAL (evaluation) mode."""
     del losses_dict
     hparams = self.hparams
 
@@ -1425,7 +1425,7 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
           loss=loss)
 
   def estimator_spec_predict(self, features, use_tpu=False):
-    """Construct EstimatorSpec for PREDICT mode."""
+    """Constructs `tf.estimator.EstimatorSpec` for PREDICT (inference) mode."""
     decode_hparams = self._decode_hparams
     infer_out = self.infer(
         features,

From 9e590afc6c06cff65610ce7357be14ba4e3ddba9 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Fri, 31 Aug 2018 17:23:36 -0700
Subject: [PATCH 0717/2720] Add 2d positional embeddings to mesh for image
 transformer, add label conditioning and update hparams

PiperOrigin-RevId: 211169110
---
 .../mesh_tensorflow/mtf_image_transformer.py  | 98 ++++++++++++-------
 1 file changed, 65 insertions(+), 33 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
index fe01c528e..a854745a6 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
@@ -50,6 +50,39 @@ def set_activation_type(self):
           "unknown hparams.activation_dtype %s" % hparams.activation_dtype)
     return activation_dtype
 
+  def create_positional_emb_2d(self, targets, max_length_dim, model_dim):
+    """Learned 2d positional embedding for images."""
+    mesh = targets.mesh
+    hparams = self._hparams
+    activation_dtype = self.set_activation_type()
+
+    rows_dim = mtf.Dimension("rows", hparams.img_len)
+    cols_dim = mtf.Dimension("cols", hparams.img_len*hparams.num_channels)
+
+    positional_emb_rows_var = mtf.get_variable(
+        mesh, "positional_emb_rows",
+        mtf.Shape([max_length_dim, model_dim]),
+        initializer=tf.random_normal_initializer(),
+        activation_dtype=activation_dtype)
+    positional_emb_cols_var = mtf.get_variable(
+        mesh, "positional_emb_cols",
+        mtf.Shape([max_length_dim, model_dim]),
+        initializer=tf.random_normal_initializer(),
+        activation_dtype=activation_dtype)
+
+    targets_position_x = mtf.range(mesh, rows_dim, dtype=tf.int32)
+    targets_position_y = mtf.range(mesh, cols_dim, dtype=tf.int32)
+    position_x = mtf.broadcast(
+        mtf.gather(positional_emb_rows_var, targets_position_x,
+                   max_length_dim),
+        mtf.Shape([rows_dim, cols_dim, model_dim]))
+
+    position_y = mtf.broadcast(
+        mtf.gather(positional_emb_cols_var, targets_position_y,
+                   max_length_dim),
+        mtf.Shape([rows_dim, cols_dim, model_dim]))
+    return position_x + position_y
+
   def mtf_model_fn(self, features, mesh):
     features = copy.copy(features)
     tf.logging.info("features = %s" % features)
@@ -89,21 +122,7 @@ def layer_prepostprocess_dropout(x):
 
     extra_losses = []
 
-    # TODO(nikip): Verify conditional.
-    if self.has_input and not hparams.unconditional:
-      vocab_size = hparams.num_classes
-      inputs_vocab_dim = mtf.Dimension("vocab", vocab_size)
-      inputs = tf.squeeze(tf.to_int32(features["inputs"]), [2, 3])
-      inputs = import_to_batch_by_length(inputs, "inputs")
-
-      # Input embeddings
-      inputs, _ = mtf_layers.embedding(
-          inputs, inputs_vocab_dim, model_dim,
-          activation_dtype=activation_dtype,
-          name="inputs_embedding")
-
     # Create targets content and position embeddings.
-    targets_position = mtf.range(mesh, length_dim, dtype=tf.int32)
     targets_vocab_size = 256 * hparams.num_channels
     targets_vocab_dim = mtf.Dimension("vocab", targets_vocab_size)
     outputs_vocab_dim = mtf.Dimension("output_vocab", 256)
@@ -115,14 +134,27 @@ def layer_prepostprocess_dropout(x):
         initializer=tf.random_normal_initializer(),
         activation_dtype=activation_dtype)
 
-    positional_embedding_var = mtf.get_variable(
-        mesh, "positional_embedding",
-        mtf.Shape([max_length_dim, model_dim]),
-        initializer=tf.random_normal_initializer(),
-        activation_dtype=activation_dtype)
-    x = (mtf.gather(targets_embedding_var, shifted_targets, targets_vocab_dim) +
-         mtf.gather(
-             positional_embedding_var, targets_position, max_length_dim))
+    x = mtf.gather(targets_embedding_var, shifted_targets, targets_vocab_dim)
+    # Add positional embeddings
+    x += mtf.reshape(
+        self.create_positional_emb_2d(targets, max_length_dim, model_dim),
+        [length_dim, model_dim])
+
+    # If conditional and input is given, add the input embedding to the target.
+    # TODO(nikip): Verify conditional.
+    if self.has_input and not hparams.unconditional:
+      vocab_size = hparams.num_classes
+      inputs_vocab_dim = mtf.Dimension("vocab", vocab_size)
+      inputs = tf.squeeze(tf.to_int32(features["inputs"]), [2, 3])
+      inputs = import_to_batch_by_length(inputs, "inputs")
+
+      # Input embeddings
+      inputs_embedding_var = mtf_layers.embedding(
+          mesh, "input_embedding",
+          mtf.Shape([inputs_vocab_dim, model_dim]),
+          activation_dtype=activation_dtype)
+      inputs_emb = mtf.gather(inputs_embedding_var, inputs, inputs_vocab_dim)
+      x += inputs_emb
 
     # Image Transformer Decoder
     # [ self attention - ffn - residual + dropout] x n
@@ -211,13 +243,13 @@ def mtf_image_transformer_tiny():
   hparams.d_ff = 256
   hparams.batch_size = 4
   hparams.num_encoder_layers = 1
-  hparams.num_decoder_layers = 1
+  hparams.num_decoder_layers = 2
   hparams.num_heads = 4
   hparams.attention_key_size = 128
   hparams.attention_value_size = 128
   # data parallelism and model-parallelism
-  hparams.mesh_shape = "2.2"
-  hparams.layout = "batch:0;filter_size:1"
+  hparams.mesh_shape = "batch:2"
+  hparams.layout = "batch:batch"
   return hparams
 
 
@@ -255,9 +287,9 @@ def mtf_image_transformer_base_single():
 def mtf_image_transformer_base_cifar():
   """Data parallel CIFAR parameters."""
   hparams = mtf_image_transformer_base()
-  hparams.mesh_shape = "batch:32"
+  hparams.mesh_shape = "batch:8"
   hparams.layout = "batch:batch"
-  hparams.batch_size = 128
+  hparams.batch_size = 32
   hparams.num_heads = 4
   hparams.num_decoder_layers = 12
   hparams.block_length = 256
@@ -309,8 +341,8 @@ def mtf_image_transformer_base_imagenet_mp():
 @registry.register_hparams
 def mtf_image_transformer_tiny_moe():
   hparams = mtf_image_transformer_tiny()
-  hparams.mesh_shape = "4"
-  hparams.layout = "batch:0,experts:0"
+  hparams.mesh_shape = "all:4"
+  hparams.layout = "batch:all,experts:all"
   hparams.ffn_layer = "moe"
   return hparams
 
@@ -318,14 +350,14 @@ def mtf_image_transformer_tiny_moe():
 @registry.register_hparams
 def mtf_image_transformer_tiny_8gpu():
   hparams = mtf_image_transformer_tiny()
-  hparams.mesh_shape = "8"
-  hparams.layout = "vocab:0;filter_size:0;heads:0"
+  hparams.mesh_shape = "all:8"
+  hparams.layout = "vocab:all;filter_size:all;heads:all"
   return hparams
 
 
 @registry.register_hparams
 def mtf_image_transformer_length_sharded():
   hparams = mtf_image_transformer_tiny()
-  hparams.mesh_shape = "2"
-  hparams.layout = "length:0"
+  hparams.mesh_shape = "all:2"
+  hparams.layout = "length:all"
   return hparams

From 5deeac83c895d71828d54c6f37b0ed624703330b Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Fri, 31 Aug 2018 17:30:10 -0700
Subject: [PATCH 0718/2720] fix moe return bug

PiperOrigin-RevId: 211169707
---
 tensor2tensor/mesh_tensorflow/mtf_transformer.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
index 26f994f2a..d3d165bf6 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
@@ -304,6 +304,9 @@ def _feedforward_layer(self, x, losses=None):
           self.model_dim,
           hparams,
           hparams.mode == tf.estimator.ModeKeys.TRAIN)
+      if losses is not None:
+        losses.append(loss)
+      return output
     elif feedforward_layer == "hmoe":
       output, loss = moe.transformer_moe_layer_v2(
           x,

From 0d1e3cc38e1b7b43237cb5ef0cbf6a0fecf60ed7 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 31 Aug 2018 17:35:09 -0700
Subject: [PATCH 0719/2720] DistributedText2TextProblem, a base class for
 Text2TextProblem for large-datasets.

PiperOrigin-RevId: 211170129
---
 .../data_generators/text_problems.py          | 220 +++++++++++++++++-
 .../data_generators/text_problems_test.py     | 206 ++++++++++++++++
 2 files changed, 425 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 0258892ee..e5e6a3159 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -533,7 +533,7 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
       for idx, inp in enumerate(sample["inputs"]):
         inputs += encoder.encode(inp)
         inputs.append(text_encoder.EOS_ID)
-        if idx < len(sample["inputs"])-1:
+        if idx < len(sample["inputs"]) - 1:
           inputs.append(encoder.encode(self.CONCAT_TOKEN)[0])
       label = sample["label"]
       yield {"inputs": inputs, "targets": [label]}
@@ -947,3 +947,221 @@ def num_generate_tasks(self):
 
   def eval_metrics(self):
     return [metrics.Metrics.ACC, metrics.Metrics.NEG_LOG_PERPLEXITY]
+
+
+class DistributedText2TextProblem(Text2TextProblem):
+  """Base class for text-to-text problems for large-datasets.
+
+  Text2TextProblem doesn't support data generation in a distributed manner.
+
+  Use DistributedText2TextProblem if you have a sharded dataset(s) and want to
+  create tf.Examples from them in a distributed manner.
+
+  Every task will write to one output shard and will read from specific input
+  shards.
+
+  Subclasses should override `generate_samples`, `input_dataset_files`
+  and `is_generate_per_split` as described below.
+
+  Users need to generate the vocabulary before generating data.
+  See tensor2tensor/bin/build_vocab.py.
+  """
+
+  # START: Subclass interface
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split, input_files):
+    """Generate samples of input text and target text pairs.
+
+    Subclasses should generate the samples using only files from `input_files`.
+
+    Please see Text2TextProblem.generate_samples for a fuller explanation.
+
+    Args:
+      data_dir: final data directory.
+      tmp_dir: temporary directory that you can use for downloading and scratch.
+      dataset_split: problem.DatasetSplit, which data split to generate samples
+        for (for example, training and evaluation).
+      input_files: Generate samples using only these input dataset files.
+
+    Yields:
+      {"inputs": text, "targets": text}
+    """
+    raise NotImplementedError()
+
+  def input_files(self, dataset_split=problem.DatasetSplit.TRAIN):
+    """The input files of the input dataset.
+
+    If you don't have a separate dev/test split then returning []
+    suffices for dataset_split != problem.DatasetSplit.TRAIN
+
+    Args:
+      dataset_split: The split for which to return the input files for.
+
+    Returns:
+      list of strings: The files for the supplied datasplit
+    """
+
+    raise NotImplementedError()
+
+  # END: Subclass interface
+
+  @property
+  def num_output_shards(self):
+    # Returns the total number of output shards.
+    num_output_shards = 0
+    for split in self.dataset_splits:
+      num_output_shards += split["shards"]
+    return num_output_shards
+
+  @property
+  def split_to_input_filenames(self):
+    # Dictionary of dataset split to input dataset filenames.
+    split_to_input_filenames = {}
+    num_input_files = 0
+    if not self.is_generate_per_split:
+      # We just have a single input dataset file.
+      split_to_input_filenames[problem.DatasetSplit.TRAIN] = (
+          self.input_files(problem.DatasetSplit.TRAIN))
+      num_input_files += len(
+          split_to_input_filenames[problem.DatasetSplit.TRAIN])
+    else:
+      # We have separate input dataset files.
+      for dataset_split in self.dataset_splits:
+        split = dataset_split["split"]
+        split_to_input_filenames[split] = self.input_files(split)
+        num_input_files += len(split_to_input_filenames[split])
+
+    # Number of input files >= number of output files. So that every task should
+    # have some work to do!
+    assert num_input_files >= self.num_output_shards
+
+    return split_to_input_filenames
+
+  def _task_id_to_output_split(self, task_id):
+    # Takes a task_id and returns a tuple of
+    # (split of the dataset to operate on, number of shards in that split,
+    # offset of this task from the first task to operate on that split)
+    num_output_shards = 0
+    for dataset_split in self.dataset_splits:
+      num_output_shards += dataset_split["shards"]
+      if task_id < num_output_shards:
+        return (dataset_split["split"], dataset_split["shards"],
+                (task_id - num_output_shards + dataset_split["shards"]))
+
+  def _task_id_to_output_file(self, data_dir, task_id):
+    # Returns the output filename that this task will write.
+
+    dataset_split, shards, offset = self._task_id_to_output_split(task_id)
+
+    filepath_fns = {
+        problem.DatasetSplit.TRAIN: self.training_filepaths,
+        problem.DatasetSplit.EVAL: self.dev_filepaths,
+        problem.DatasetSplit.TEST: self.test_filepaths,
+    }
+
+    return filepath_fns[dataset_split](data_dir, shards, False)[offset]
+
+  @staticmethod
+  def _divide_equally(input_files, num_tasks, task_id):
+    # There are num_tasks total tasks, we need to divide these
+    # input files among them equally and return the slice that task_id should
+    # read from.
+    task_load, remainder = divmod(len(input_files), num_tasks)
+
+    # This is the slice of almost equal sized chunks of files for a task_id to
+    # handle -- this distributes the excess remainder tasks among the first
+    # "remainder" task_ids.
+
+    # The extra min(task_id, remainder) in the end comes from assigning the
+    # remainder of the tasks to task_ids [0, remainder), so we need to advance
+    # the start by how many ever remainder tasks already assigned.
+    start_idx = task_id * task_load + min(task_id, remainder)
+
+    # This will handle atleast `task_load` files, plus an extra one if `task_id`
+    # is still less than remainder.
+    num_elements = task_load + int(task_id < remainder)
+
+    return input_files[start_idx : start_idx + num_elements]
+
+  def _task_id_to_input_files(self, task_id):
+    # Returns a list of input files that this task should read and process.
+
+    if not self.is_generate_per_split:
+      # We just have one unified input dataset to handle, so all tasks will read
+      # from the TRAIN dataset.
+      input_files = self.split_to_input_filenames[problem.DatasetSplit.TRAIN]
+
+      return self._divide_equally(input_files, self.num_output_shards, task_id)
+
+    # self.is_generate_per_split is True.
+    dataset_split, num_shards, offset = self._task_id_to_output_split(task_id)
+    input_files = self.split_to_input_filenames[dataset_split]
+    return self._divide_equally(input_files, num_shards, offset)
+
+  def generate_text_for_vocab(self, data_dir, tmp_dir):
+    # We need to override this because we'll be reading from specific files
+    # instead
+
+    # What files should we read for creating the vocabulary?
+    input_files_for_vocab = []
+    if self.is_generate_per_split:
+      input_files_for_vocab = (
+          self.split_to_input_filenames[problem.DatasetSplit.TRAIN])
+    else:
+      # We need to compute the 'train' shards from the whole input.
+      # Go over all task_ids that output training data, collect their input
+      # files.
+      for task_id in range(self.num_output_shards):
+        split, _, _ = self._task_id_to_output_split(task_id)
+        if split == problem.DatasetSplit.TRAIN:
+          input_files_for_vocab.extend(self._task_id_to_input_files(task_id))
+
+    # Generate samples only from the above generated files.
+    for i, sample in enumerate(
+        self.generate_samples(data_dir, tmp_dir, problem.DatasetSplit.TRAIN,
+                              input_files_for_vocab)):
+      if self.has_inputs:
+        yield sample["inputs"]
+      yield sample["targets"]
+      if self.max_samples_for_vocab and (i + 1) >= self.max_samples_for_vocab:
+        break
+
+  def generate_encoded_samples(self,
+                               data_dir,
+                               tmp_dir,
+                               dataset_split,
+                               input_files):
+    # Since this is a distributed problem, we don't want every task to create
+    # its own vocabulary, so we assume that the dictionary is already created
+    # for example by using build_vocab.py
+    vocab_filepath = os.path.join(data_dir, self.vocab_filename)
+    if not tf.gfile.Exists(vocab_filepath):
+      raise ValueError("Vocab file: %s doesn't exist, please use "
+                       "build_vocab.py to create one." % vocab_filepath)
+    encoder = self.get_or_create_vocab(data_dir, tmp_dir, force_get=True)
+    generator = self.generate_samples(data_dir, tmp_dir, dataset_split,
+                                      input_files)
+    return text2text_generate_encoded(
+        generator, encoder, has_inputs=self.has_inputs)
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    # task_id should be in [0, self.num_output_shards)
+    assert (0 <= task_id) and (task_id < self.num_output_shards)
+
+    # A task_id is only supposed to write only one output shard, it can operate
+    # over multiple *input* shards.
+    input_files = self._task_id_to_input_files(task_id)
+    output_file = self._task_id_to_output_file(data_dir, task_id)
+
+    # Which output split is this task writing to?
+    split, _, _ = self._task_id_to_output_split(task_id)
+
+    # Actually generate examples.
+    generator_utils.generate_files(
+        self._maybe_pack_examples(
+            self.generate_encoded_samples(
+                data_dir, tmp_dir, split, input_files)),
+        [output_file])
+
+    # Shuffle the output.
+    generator_utils.shuffle_dataset([output_file])
diff --git a/tensor2tensor/data_generators/text_problems_test.py b/tensor2tensor/data_generators/text_problems_test.py
index ab50cdefc..a42342404 100644
--- a/tensor2tensor/data_generators/text_problems_test.py
+++ b/tensor2tensor/data_generators/text_problems_test.py
@@ -183,5 +183,211 @@ def testText2TextTmpDir(self):
     self.assertTrue(encoder.decode(targets_encoded) in self.targets)
 
 
+class FakeDistributedProblem(text_problems.DistributedText2TextProblem):
+
+  def __init__(self):
+    self.name = "fake_distributed_problem"
+    # Call the base class ctor.
+    super(FakeDistributedProblem, self).__init__()
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split, input_files):
+    # Read all lines from all the input_files and return the same word as input
+    # and target.
+    for input_file in input_files:
+      with tf.gfile.Open(input_file, "r") as f:
+        for line in f.read().strip().split("\n"):
+          yield {"inputs": line.strip(), "targets": line.strip()}
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem_lib.DatasetSplit.TRAIN,
+        "shards": 2,
+    }, {
+        "split": problem_lib.DatasetSplit.EVAL,
+        "shards": 3,
+    }, {
+        "split": problem_lib.DatasetSplit.TEST,
+        "shards": 4,
+    }]
+
+  def input_files(self, dataset_split=problem_lib.DatasetSplit.TRAIN):
+    if dataset_split == problem_lib.DatasetSplit.TRAIN:
+      return self.train_files
+    elif dataset_split == problem_lib.DatasetSplit.EVAL:
+      return self.dev_files
+    return self.test_files
+
+  @classmethod
+  def setup_for_test(cls):
+    # First setup the temp train, dev, test files and then call the ctor.
+    cls.tmp_dir = tf.test.get_temp_dir()
+    shutil.rmtree(cls.tmp_dir)
+    os.mkdir(cls.tmp_dir)
+
+    # Write 25 train files, 5 dev files, 11 test files.
+    train_pattern = os.path.join(cls.tmp_dir, "train-%05d-of-00025")
+    dev_pattern = os.path.join(cls.tmp_dir, "dev-%05d-of-00005")
+    test_pattern = os.path.join(cls.tmp_dir, "test-%05d-of-00011")
+    cls.train_files, cls.dev_files, cls.test_files = [], [], []
+    for i in range(25):
+      cls.train_files.append(train_pattern % i)
+      with tf.gfile.Open(cls.train_files[-1], "w") as f:
+        f.write("train_%d\n" % i)
+    for i in range(5):
+      cls.dev_files.append(dev_pattern % i)
+      with tf.gfile.Open(cls.dev_files[-1], "w") as f:
+        f.write("dev_%d\n" % i)
+    for i in range(11):
+      cls.test_files.append(test_pattern % i)
+      with tf.gfile.Open(cls.test_files[-1], "w") as f:
+        f.write("test_%d\n" % i)
+
+
+class FakeDistributedProblemNotPerSplit(FakeDistributedProblem):
+
+  @property
+  def is_generate_per_split(self):
+    return False
+
+
+class DistributedText2TextProblemsTest(tf.test.TestCase):
+
+  def setUp(self):
+    FakeDistributedProblem.setup_for_test()
+
+  def testOutputSharding(self):
+    problem = FakeDistributedProblemNotPerSplit()
+
+    # self.dataset_split is 2, 3, 4
+    # So:
+    # num output shards = 2 + 3 + 4 = 9
+    # task_ids will be in range = [0, 9)
+
+    expected_split_shard_and_offset = [
+        (problem_lib.DatasetSplit.TRAIN, 2, 0),
+        (problem_lib.DatasetSplit.TRAIN, 2, 1),
+        (problem_lib.DatasetSplit.EVAL, 3, 0),
+        (problem_lib.DatasetSplit.EVAL, 3, 1),
+        (problem_lib.DatasetSplit.EVAL, 3, 2),
+        (problem_lib.DatasetSplit.TEST, 4, 0),
+        (problem_lib.DatasetSplit.TEST, 4, 1),
+        (problem_lib.DatasetSplit.TEST, 4, 2),
+        (problem_lib.DatasetSplit.TEST, 4, 3),
+    ]
+
+    expected_output_filenames = [
+        "/tmp/fake_distributed_problem-unshuffled-train-00000-of-00002",
+        "/tmp/fake_distributed_problem-unshuffled-train-00001-of-00002",
+        "/tmp/fake_distributed_problem-unshuffled-dev-00000-of-00003",
+        "/tmp/fake_distributed_problem-unshuffled-dev-00001-of-00003",
+        "/tmp/fake_distributed_problem-unshuffled-dev-00002-of-00003",
+        "/tmp/fake_distributed_problem-unshuffled-test-00000-of-00004",
+        "/tmp/fake_distributed_problem-unshuffled-test-00001-of-00004",
+        "/tmp/fake_distributed_problem-unshuffled-test-00002-of-00004",
+        "/tmp/fake_distributed_problem-unshuffled-test-00003-of-00004"
+    ]
+
+    actual_split_shard_and_offset = []
+    actual_output_filenames = []
+    for task_id in range(9):
+      actual_split_shard_and_offset.append(
+          problem._task_id_to_output_split(task_id))
+      actual_output_filenames.append(
+          problem._task_id_to_output_file("/tmp", task_id))
+
+    self.assertSequenceEqual(expected_split_shard_and_offset,
+                             actual_split_shard_and_offset)
+
+    self.assertSequenceEqual(expected_output_filenames, actual_output_filenames)
+
+  def testInputShardingNoGeneratePerSplit(self):
+    # 25 input shards (train only, is_generate_per_split = False).
+    # 9 output tasks in all (2 + 3 + 4), so
+    #
+    # Division should be like:
+    # task_id 0 -> 0, 1, 2
+    # task_id 1 -> 3, 4, 5
+    # ...
+    # task_id 6 -> 18, 19, 20
+    # task_id 7 -> 21, 22
+    # task_id 8 -> 23, 24
+
+    # tasks 0 to 6
+    expected_input_file_sharding = [[
+        "train-%05d-of-00025" % j for j in [i, i + 1, i + 2]
+    ] for i in range(0, 20, 3)]
+    # tasks 7 and 8
+    expected_input_file_sharding.extend(
+        [["train-%05d-of-00025" % i for i in [21, 22]],
+         ["train-%05d-of-00025" % i for i in [23, 24]]])
+
+    problem = FakeDistributedProblemNotPerSplit()
+
+    list_input_files = []
+    for task_id in range(9):
+      input_files = problem._task_id_to_input_files(task_id)
+      list_input_files.append(
+          [os.path.basename(input_file) for input_file in input_files])
+
+    self.assertSequenceEqual(expected_input_file_sharding, list_input_files)
+
+  def testInputShardingWithGeneratePerSplit(self):
+    # 25, 5, 11 train, dev, test input shards
+    # 9 output tasks in all (2 + 3 + 4), so
+    #
+    # Division should be like:
+    #
+    # Train
+    # task_id 0 -> 0, .. 12
+    # task_id 1 -> 13 .. 24
+    #
+    # Dev
+    # task_id 2 -> 0, 1
+    # task_id 3 -> 2, 3,
+    # task_id 4 -> 4
+    #
+    # Test
+    # task_id 5 -> 0, 1, 2
+    # task_id 6 -> 3, 4, 5
+    # task_id 7 -> 6, 7, 8
+    # task_id 8 -> 9, 10
+
+    expected_input_file_sharding = [
+        ["train-%05d-of-00025" % i for i in range(13)],      # task_id 0
+        ["train-%05d-of-00025" % i for i in range(13, 25)],  # task_id 1
+        ["dev-%05d-of-00005" % i for i in [0, 1]],           # task_id 2
+        ["dev-%05d-of-00005" % i for i in [2, 3]],           # task_id 3
+        ["dev-%05d-of-00005" % i for i in [4]],              # task_id 4
+        ["test-%05d-of-00011" % i for i in [0, 1, 2]],       # task_id 5
+        ["test-%05d-of-00011" % i for i in [3, 4, 5]],       # task_id 6
+        ["test-%05d-of-00011" % i for i in [6, 7, 8]],       # task_id 7
+        ["test-%05d-of-00011" % i for i in [9, 10]],         # task_id 8
+    ]
+
+    problem = FakeDistributedProblem()
+
+    list_input_files = []
+    for task_id in range(9):
+      input_files = problem._task_id_to_input_files(task_id)
+      list_input_files.append(
+          [os.path.basename(input_file) for input_file in input_files])
+
+    self.assertSequenceEqual(expected_input_file_sharding, list_input_files)
+
+  def testVocabularyIsAllTrain(self):
+    problem = FakeDistributedProblem()
+
+    tmp_dir = problem.tmp_dir
+
+    for text in problem.generate_text_for_vocab(tmp_dir, tmp_dir):
+      # All the vocabulary is coming from training input shards.
+      self.assertTrue("train_" in text, "train is not in %s" % text)
+
+
 if __name__ == "__main__":
   tf.test.main()

From 19953703430ecfccf68a4ddf8cb3c8ad2c618e69 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 31 Aug 2018 19:25:27 -0700
Subject: [PATCH 0720/2720] Add INFO verbosity to model exporter.

PiperOrigin-RevId: 211177510
---
 tensor2tensor/serving/export.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index f3dc0ac76..5be45bc45 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -68,4 +68,5 @@ def main(_):
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()

From e1cf9f4c76033ff9cc004964af079e95f338c6d3 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 31 Aug 2018 22:27:12 -0700
Subject: [PATCH 0721/2720] Remove internal cloud_tpu support in favor of
 official tutorial (using ctpu)

PiperOrigin-RevId: 211186051
---
 docs/cloud_tpu.md                             | 121 +-----
 docs/index.md                                 |   2 -
 docs/tutorials/asr_with_transformer.md        |  94 ----
 tensor2tensor/bin/t2t_distill.py              |  81 ++--
 tensor2tensor/bin/t2t_trainer.py              |  45 +-
 .../wikisum/parallel_launch.py                |   2 +-
 tensor2tensor/serving/serving_utils.py        |   2 +-
 tensor2tensor/utils/cloud_mlengine.py         |  88 +++-
 tensor2tensor/utils/cloud_tpu.py              | 409 ------------------
 9 files changed, 149 insertions(+), 695 deletions(-)
 delete mode 100644 docs/tutorials/asr_with_transformer.md
 delete mode 100644 tensor2tensor/utils/cloud_tpu.py

diff --git a/docs/cloud_tpu.md b/docs/cloud_tpu.md
index cf34523e7..c0625e132 100644
--- a/docs/cloud_tpu.md
+++ b/docs/cloud_tpu.md
@@ -1,114 +1,41 @@
 # Running on Cloud TPUs
 
 Tensor2Tensor supports running on Google Cloud Platforms TPUs, chips
-specialized for ML training. See the official tutorial for [running Transformer
-on Cloud TPUs](https://cloud.google.com/tpu/docs/tutorials/transformer) or
-read on for more T2T models on TPUs.
+specialized for ML training. See the official tutorials for [running the
+T2T Transformer for text on Cloud TPUs](https://cloud.google.com/tpu/docs/tutorials/transformer) and
+[Transformer for Speech Recognition](https://cloud.google.com/tpu/docs/tutorials/automated-speech-recognition).
 
-## Models and hparams for TPU:
+## Other models on TPU
 
-Transformer:
-* `transformer` with `transformer_tpu` (or `transformer_packed_tpu`,
-    `transformer_tiny_tpu`, `transformer_big_tpu`)
-* `transformer_encoder` with `transformer_tpu` (and the above ones)
+Many of Tensor2Tensor's models work on TPU.
 
-You can run the Transformer model on a number of problems,
-from translation through language modeling to sentiment analysis.
-See the official tutorial for [running Transformer
-on Cloud TPUs](https://cloud.google.com/tpu/docs/tutorials/transformer)
-for some examples and try out your own problems.
+You can provision a VM and TPU with `ctpu up`. Use the `t2t-trainer` command
+on the VM as usual with the additional flags `--use_tpu` and
+`--cloud_tpu_name=$TPU_NAME`.
 
-You can train an Automatic Speech Recognition (ASR) model with Transformer
-on TPU by using `transformer` as `model` with `transformer_librispeech_tpu` as
-`hparams_set` and `librispeech` as `problem`. See this [tutorial](tutorials/asr_with_transformer.md) for more details on training it and this
-[notebook](https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/asr_transformer.ipynb) to see how the resulting model transcribes your speech to text.
+Note that because the `TPUEstimator` does not catch the `OutOfRangeError`
+during evaluation, you should ensure that `--eval_steps` is small enough to
+not exhaust the evaluation data.
 
-Image Transformer:
-* `imagetransformer` with `imagetransformer_base_tpu` (or
-    `imagetransformer_tiny_tpu`)
-* `img2img_transformer` with `img2img_transformer_base_tpu` (or
-    `img2img_transformer_tiny_tpu`)
+A non-exhaustive list of T2T models that work on TPU:
 
-You can run the `ImageTransformer` model on problems like unconditional or
-conditional Image generation and `Img2ImgTransformer` model on Super Resolution.
-We run on datasets like CelebA, CIFAR and ImageNet but they should work with any
-other image dataset.
-
-Residual networks:
+* Image generation: `imagetransformer` with `imagetransformer_base_tpu` (or
+  `imagetransformer_tiny_tpu`)
+* Super-resolution: `img2img_transformer` with `img2img_transformer_base_tpu`
+  (or `img2img_transformer_tiny_tpu`)
 * `resnet` with `resnet_50` (or `resnet_18` or `resnet_34`)
 * `revnet` with `revnet_104` (or `revnet_38_cifar`)
 * `shake_shake` with `shakeshake_tpu` (or `shakeshake_small`)
 
-We run residual networks on MNIST, CIFAR and ImageNet, but they should
-work on any image classification data-set.
-
-## Tutorial: Transformer En-De translation on TPU
-
-Configure the `gcloud` CLI:
-```
-gcloud components update
-gcloud auth application-default login
-# Set your default zone to a TPU-enabled zone.
-gcloud config set compute/zone us-central1-b
-```
-
-Generate data to GCS.
-If you already have the data, use `gsutil cp` to copy to GCS.
-```
-GCS_BUCKET=gs://my-bucket
-DATA_DIR=$GCS_BUCKET/t2t/data/
-t2t-datagen --problem=translate_ende_wmt8k --data_dir=$DATA_DIR
-```
-
-Specify an output directory and launch TensorBoard to monitor training:
-```
-OUT_DIR=$GCS_BUCKET/t2t/training/transformer_v1
-tensorboard --logdir=$OUT_DIR
-```
+## Example invocation
 
-Note that both the data and output directories must be Google Cloud Storage
-buckets (i.e. start with `gs://`).
+Use `ctpu up` to bring up the VM and TPU machines; once the machines are ready
+it will SSH you into the VM and you can run the following:
 
-Launch! It's as simple as adding the `--cloud_tpu` flag.
 ```
-t2t-trainer \
-  --model=transformer \
-  --hparams_set=transformer_tpu \
-  --problem=translate_ende_wmt8k \
-  --train_steps=10 \
-  --eval_steps=10 \
-  --local_eval_frequency=10 \
-  --data_dir=$DATA_DIR \
-  --output_dir=$OUT_DIR \
-  --cloud_tpu \
-  --cloud_delete_on_done
-```
-
-The above command will train for 10 steps, then evaluate for 10 steps. You can
-(and should) increase the number of total training steps with the
-`--train_steps` flag. Evaluation will happen every `--local_eval_frequency`
-steps, each time for `--eval_steps`. The `--cloud_delete_on_done` flag has the
-trainer delete the VMs on completion.
-
-Voila. Enjoy your new supercomputer.
+# DATA_DIR and OUT_DIR should be GCS buckets
+# TPU_NAME should have been set automatically by the ctpu tool
 
-Note that checkpoints are compatible between CPU, GPU, and TPU models so you can
-switch between hardware at will.
-
-## Additional flags
-
-* `--cloud_vm_name`: The name of the VM to use or create. This can be reused
-  across multiple concurrent runs.
-* `--cloud_tpu_name`: The name of the TPU instance to use or create. If you want
-  to launch multiple jobs on TPU, provide different names here for each one.
-  Each TPU instance can only be training one model at a time.
-
-## Other T2T models on TPU
-
-To run other models on TPU, proceed exactly as in the tutorial above,
-just with different model, problem and hparams_set (and directories).
-For example, to train a shake-shake model on CIFAR you can run this command.
-```
 t2t-trainer \
   --model=shake_shake \
   --hparams_set=shakeshake_tpu \
@@ -118,8 +45,6 @@ t2t-trainer \
   --local_eval_frequency=100 \
   --data_dir=$DATA_DIR \
   --output_dir=$OUT_DIR \
-  --cloud_tpu \
-  --cloud_delete_on_done
+  --use_tpu \
+  --cloud_tpu_name=$TPU_NAME
 ```
-Note that `eval_steps` should not be too high so as not to run out
-of evaluation data.
diff --git a/docs/index.md b/docs/index.md
index 7a7287851..38c6120eb 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -20,7 +20,6 @@ research](https://research.googleblog.com/2017/06/accelerating-deep-learning-res
 
 * [Walkthrough](walkthrough.md): Install and run.
 * [IPython notebook](https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/hello_t2t.ipynb): Get a hands-on experience.
-* [Automatic Speech Recognition notebook](https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/asr_transformer.ipynb): Transcribe speech to text with a T2T model.
 
 ## Basics
 
@@ -33,7 +32,6 @@ research](https://research.googleblog.com/2017/06/accelerating-deep-learning-res
 * [Training on Google Cloud ML](cloud_mlengine.md)
 * [Training on Google Cloud TPUs](cloud_tpu.md)
 * [Distributed Training](distributed_training.md)
-* [Automatic Speech Recognition (ASR) with Transformer](tutorials/asr_with_transformer.md)
 
 ## Solving your task
 
diff --git a/docs/tutorials/asr_with_transformer.md b/docs/tutorials/asr_with_transformer.md
deleted file mode 100644
index 9feb381b3..000000000
--- a/docs/tutorials/asr_with_transformer.md
+++ /dev/null
@@ -1,94 +0,0 @@
-# Automatic Speech Recognition (ASR) with Transformer
-
-Check out the [Automatic Speech Recognition notebook](https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/asr_transformer.ipynb) to see how the resulting model transcribes your speech to text.
-
-## Data set
-
-This tutorial uses the publicly available
-[Librispeech](http://www.openslr.org/12/) ASR corpus.
-
-
-## Generate the dataset
-
-To generate the dataset use `t2t-datagen`. You need to create environment
-variables for a data directory `DATA_DIR` where the data is stored and for a
-temporary directory `TMP_DIR` where necessary data is downloaded.
-
-As the audio import in `t2t-datagen` uses `sox` to generate normalized
-waveforms, please install it as appropriate (e.g. `apt-get install sox`).
-
-```
-# Generate both the full dataset and the small clean version, which we use for
-# evaluation.
-t2t-datagen --problem=librispeech --data_dir=$DATA_DIR --tmp_dir=$TMP_DIR
-t2t-datagen --problem=librispeech_clean --data_dir=$DATA_DIR --tmp_dir=$TMP_DIR
-```
-
-The problem `librispeech_train_full_test_clean` will train on the full dataset
-but evaluate on the clean dataset.
-
-You can also use `librispeech_clean_small` which is a small version of the
-clean dataset.
-
-## Training on Cloud TPUs
-
-To train a model on TPU set up `OUT_DIR` and run the trainer with big batches
-and truncated sequences:
-
-```
-t2t-trainer \
-  --model=transformer \
-  --hparams_set=transformer_librispeech_tpu \
-  --hparams=max_length=125550,max_input_seq_length=1550,max_target_seq_length=350,batch_size=16 \
-  --problem=librispeech_train_full_test_clean \
-  --train_steps=210000 \
-  --eval_steps=3 \
-  --local_eval_frequency=100 \
-  --data_dir=$DATA_DIR \
-  --output_dir=$OUT_DIR \
-  --cloud_tpu \
-  --cloud_delete_on_done
-```
-
-After this step is compleated run the training again for more steps with smaller
-batch size and full sequences:
-
-```
-t2t-trainer \
-  --model=transformer \
-  --hparams_set=transformer_librispeech_tpu \
-  --hparams=max_length=295650,max_input_seq_length=3650,max_target_seq_length=650,batch_size=6 \
-  --problem=librispeech_train_full_test_clean \
-  --train_steps=230000 \
-  --eval_steps=3 \
-  --local_eval_frequency=100 \
-  --data_dir=$DATA_DIR \
-  --output_dir=$OUT_DIR \
-  --cloud_tpu \
-  --cloud_delete_on_done
-```
-
-For more information, see [Tensor2Tensor's
-documentation](https://github.com/tensorflow/tensor2tensor/tree/master/docs/cloud_tpu.md)
-for Tensor2Tensor on Cloud TPUs, or the [official Google Cloud Platform
-documentation](https://cloud.google.com/tpu/docs/tutorials/transformer) for
-Cloud TPUs.
-
-## Training on GPUs
-
-To train a model on GPU set up`OUT_DIR` and run the trainer:
-
-```
-t2t-trainer \
-  --model=transformer \
-  --hparams_set=transformer_librispeech_tpu \
-  --problem=librispeech \
-  --train_steps=120000 \
-  --eval_steps=3 \
-  --local_eval_frequency=100 \
-  --data_dir=$DATA_DIR \
-  --output_dir=$OUT_DIR
-```
-
-This model should achieve approximately 22% accuracy per sequence after
-approximately 80,000 steps.
diff --git a/tensor2tensor/bin/t2t_distill.py b/tensor2tensor/bin/t2t_distill.py
index adbeac667..30822b1ae 100644
--- a/tensor2tensor/bin/t2t_distill.py
+++ b/tensor2tensor/bin/t2t_distill.py
@@ -55,37 +55,36 @@ def main(argv):
   if argv:
     t2t_trainer.set_hparams_from_args(argv[1:])
 
-  with t2t_trainer.maybe_cloud_tpu():
-    root_output_dir = FLAGS.output_dir
-
-    # Train Teacher ============
-    hparams = t2t_trainer.create_hparams()
-    hparams.distill_phase = "train"
-    teacher_dir = os.path.join(root_output_dir, "teacher")
-    FLAGS.output_dir = teacher_dir
-
-    exp_fn = t2t_trainer.create_experiment_fn()
-    run_config = t2t_trainer.create_run_config(hparams)
-    exp = exp_fn(run_config, hparams)
-    if t2t_trainer.is_chief():
-      t2t_trainer.save_metadata(hparams)
-    t2t_trainer.execute_schedule(exp)
-    # ==========================
-    # Train Student ============
-    hparams = t2t_trainer.create_hparams()
-    hparams.add_hparam("teacher_dir", teacher_dir)
-    hparams.distill_phase = "distill"
-    student_dir = os.path.join(root_output_dir, "student")
-    FLAGS.output_dir = student_dir
-
-    exp_fn = t2t_trainer.create_experiment_fn()
-    run_config = t2t_trainer.create_run_config(hparams)
-    exp = exp_fn(run_config, hparams)
-
-    if t2t_trainer.is_chief():
-      t2t_trainer.save_metadata(hparams)
-    t2t_trainer.execute_schedule(exp)
-    # ==========================
+  root_output_dir = FLAGS.output_dir
+
+  # Train Teacher ============
+  hparams = t2t_trainer.create_hparams()
+  hparams.distill_phase = "train"
+  teacher_dir = os.path.join(root_output_dir, "teacher")
+  FLAGS.output_dir = teacher_dir
+
+  exp_fn = t2t_trainer.create_experiment_fn()
+  run_config = t2t_trainer.create_run_config(hparams)
+  exp = exp_fn(run_config, hparams)
+  if t2t_trainer.is_chief():
+    t2t_trainer.save_metadata(hparams)
+  t2t_trainer.execute_schedule(exp)
+  # ==========================
+  # Train Student ============
+  hparams = t2t_trainer.create_hparams()
+  hparams.add_hparam("teacher_dir", teacher_dir)
+  hparams.distill_phase = "distill"
+  student_dir = os.path.join(root_output_dir, "student")
+  FLAGS.output_dir = student_dir
+
+  exp_fn = t2t_trainer.create_experiment_fn()
+  run_config = t2t_trainer.create_run_config(hparams)
+  exp = exp_fn(run_config, hparams)
+
+  if t2t_trainer.is_chief():
+    t2t_trainer.save_metadata(hparams)
+  t2t_trainer.execute_schedule(exp)
+  # ==========================
 
 
 def create_teacher_experiment(run_config, hparams, argv):
@@ -108,11 +107,10 @@ def create_teacher_experiment(run_config, hparams, argv):
   if argv:
     t2t_trainer.set_hparams_from_args(argv[1:])
 
-  with t2t_trainer.maybe_cloud_tpu():
-    hparams.distill_phase = "train"
-    exp_fn = t2t_trainer.create_experiment_fn()
-    exp = exp_fn(run_config, hparams)
-    return exp
+  hparams.distill_phase = "train"
+  exp_fn = t2t_trainer.create_experiment_fn()
+  exp = exp_fn(run_config, hparams)
+  return exp
 
 
 def create_student_experiment(run_config, hparams, argv):
@@ -135,12 +133,11 @@ def create_student_experiment(run_config, hparams, argv):
   if argv:
     t2t_trainer.set_hparams_from_args(argv[1:])
 
-  with t2t_trainer.maybe_cloud_tpu():
-    hparams.add_hparam("teacher_dir", FLAGS.teacher_dir)
-    hparams.distill_phase = "distill"
-    exp_fn = t2t_trainer.create_experiment_fn()
-    exp = exp_fn(run_config, hparams)
-    return exp
+  hparams.add_hparam("teacher_dir", FLAGS.teacher_dir)
+  hparams.distill_phase = "distill"
+  exp_fn = t2t_trainer.create_experiment_fn()
+  exp = exp_fn(run_config, hparams)
+  return exp
 
 
 def create_experiment_fn(argv, train_teacher):
diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 01bfe7c12..b4bc01a4a 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -24,7 +24,6 @@
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
 from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
 from tensor2tensor.utils import cloud_mlengine
-from tensor2tensor.utils import cloud_tpu
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
@@ -81,15 +80,8 @@
   pass
 
 # Google Cloud TPUs
-flags.DEFINE_bool("cloud_tpu", False, "Whether to launch on Cloud TPUs.")
-flags.DEFINE_string("cloud_vm_name", "%s-vm" % os.getenv("USER"),
-                    "Name of Cloud VM to use or create.")
 flags.DEFINE_string("cloud_tpu_name", "%s-tpu" % os.getenv("USER"),
                     "Name of Cloud TPU instance to use or create.")
-flags.DEFINE_bool("cloud_delete_on_done", False,
-                  "Whether to delete the VM and TPU instance when done.")
-flags.DEFINE_bool("cloud_skip_confirmation", False,
-                  "Whether to skip launch confirmations.")
 
 # Google Cloud ML Engine
 flags.DEFINE_bool("cloud_mlengine", False,
@@ -151,7 +143,7 @@ def set_hparams_from_args(args):
 
 
 def create_hparams():
-  if (FLAGS.cloud_tpu or FLAGS.use_tpu) and "tpu" not in FLAGS.hparams_set:
+  if FLAGS.use_tpu and "tpu" not in FLAGS.hparams_set:
     tf.logging.warn("Not all hyperparameter sets work on TPU. "
                     "Prefer hparams_sets with a '_tpu' suffix, "
                     "e.g. transformer_tpu, if available for your model.")
@@ -326,30 +318,6 @@ def execute_schedule(exp):
     getattr(exp, FLAGS.schedule)()
 
 
-@contextlib.contextmanager
-def maybe_cloud_tpu():
-  """If FLAGS.cloud_tpu is set, setup Cloud instances."""
-  if not FLAGS.cloud_tpu:
-    yield
-    return
-
-  tf.logging.info("Running on Cloud TPU")
-
-  if (not FLAGS.data_dir.startswith("gs://") or
-      not FLAGS.output_dir.startswith("gs://")):
-    raise ValueError("To run on Cloud TPUs, data_dir and output_dir need to "
-                     "be gs:// paths, i.e. on Google Cloud Storage.")
-
-  FLAGS.use_tpu = True
-  with cloud_tpu.cloud_tpu(
-      FLAGS.cloud_vm_name,
-      FLAGS.cloud_tpu_name,
-      delete_on_done=FLAGS.cloud_delete_on_done,
-      skip_confirmation=FLAGS.cloud_skip_confirmation) as tpu_master:
-    FLAGS.master = tpu_master
-    yield
-
-
 def run_std_server():
   exp = trainer_lib.T2TExperiment(*([None] * 5))
   exp.run_std_server()
@@ -377,12 +345,11 @@ def main(argv):
     set_hparams_from_args(argv[1:])
   hparams = create_hparams()
 
-  with maybe_cloud_tpu():
-    exp_fn = create_experiment_fn()
-    exp = exp_fn(create_run_config(hparams), hparams)
-    if is_chief():
-      save_metadata(hparams)
-    execute_schedule(exp)
+  exp_fn = create_experiment_fn()
+  exp = exp_fn(create_run_config(hparams), hparams)
+  if is_chief():
+    save_metadata(hparams)
+  execute_schedule(exp)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/data_generators/wikisum/parallel_launch.py b/tensor2tensor/data_generators/wikisum/parallel_launch.py
index 3d7f5b69a..811466745 100644
--- a/tensor2tensor/data_generators/wikisum/parallel_launch.py
+++ b/tensor2tensor/data_generators/wikisum/parallel_launch.py
@@ -51,7 +51,7 @@
 import subprocess as sp
 import time
 
-from tensor2tensor.utils import cloud_tpu as cloud
+from tensor2tensor.utils import cloud_mlengine as cloud
 import tensorflow as tf
 
 flags = tf.flags
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 2004aa7a6..58aac9b0e 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -25,7 +25,7 @@
 
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
 from tensor2tensor.data_generators import text_encoder
-from tensor2tensor.utils import cloud_tpu as cloud
+from tensor2tensor.utils import cloud_mlengine as cloud
 import tensorflow as tf
 from tensorflow_serving.apis import predict_pb2
 from tensorflow_serving.apis import prediction_service_pb2_grpc
diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 443aff5dc..8a9509467 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -17,6 +17,7 @@
 import datetime
 import os
 import shutil
+import subprocess as sp
 import sys
 import tempfile
 
@@ -25,7 +26,6 @@
 
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.layers import common_hparams
-from tensor2tensor.utils import cloud_tpu as cloud
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import usr_dir as usr_dir_lib
 import tensorflow as tf
@@ -35,8 +35,73 @@
 CONSOLE_URL = "https://console.cloud.google.com/mlengine/jobs/"
 RUNTIME_VERSION = "1.9"
 
-# TODO(rsepassi):
-# * Enable multi-machine sync/async training
+
+class Gcloud(object):
+  """gcloud command strings."""
+  # Note these can be modified by set_versions
+  VM_VERSION = "tf-1-9"
+  TPU_VERSION = "1.9"
+
+  @classmethod
+  def set_versions(cls, vm, tpu):
+    cls.VM_VERSION = vm
+    cls.TPU_VERSION = tpu
+
+  @classmethod
+  def create_vm(cls):
+    create_vm_str = """
+    gcloud compute instances create {name} \
+      --machine-type=n1-standard-8 \
+      --image-family=%s \
+      --image-project=ml-images \
+      --scopes=https://www.googleapis.com/auth/cloud-platform
+    """ % cls.VM_VERSION
+    return create_vm_str
+
+  DELETE_VM = "gcloud compute instances delete {name} --quiet"
+
+  @classmethod
+  def create_tpu(cls):
+    create_tpu_str = """
+    gcloud beta compute tpus create \
+      {name} \
+      --range={tpu_ip}/29 \
+      --version=%s
+    """ % cls.TPU_VERSION
+    return create_tpu_str
+
+  DELETE_TPU = "gcloud beta compute tpus delete {name} --quiet"
+
+  LIST_TPU = "gcloud beta compute tpus list"
+  LIST_VM = "gcloud compute instances list"
+
+  SSH_LOCAL_PORT_FORWARD = "-L {local_port}:{host}:{remote_port}"
+  SSH_TUNNEL = """
+  gcloud compute ssh {name} -- -N
+  """
+
+  DEFAULT_PROJECT = "gcloud config get-value project"
+  DEFAULT_REGION = "gcloud config get-value compute/region"
+
+
+def shell_output(cmd_, **kwargs):
+  return text_encoder.to_unicode(sp.check_output(format_cmd(cmd_, **kwargs)))
+
+
+def shell_run(cmd_, **kwargs):
+  return sp.check_call(format_cmd(cmd_, **kwargs))
+
+
+def format_cmd(cmd_, **kwargs):
+  return cmd_.format(**kwargs).strip().split()
+
+
+def default_region():
+  return shell_output(Gcloud.DEFAULT_REGION).strip()
+
+
+def default_project():
+  return shell_output(Gcloud.DEFAULT_PROJECT).strip()
 
 
 def get_setup_file(name, packages=None):
@@ -112,7 +177,7 @@ def configure_job():
   training_input = {
       "pythonModule": "tensor2tensor.bin.t2t_trainer",
       "args": flags_as_args(),
-      "region": text_encoder.native_to_unicode(cloud.default_region()),
+      "region": text_encoder.native_to_unicode(default_region()),
       "runtimeVersion": RUNTIME_VERSION,
       "pythonVersion": "3.5" if sys.version_info.major == 3 else "2.7",
       "jobDir": FLAGS.output_dir,
@@ -144,7 +209,7 @@ def configure_job():
 def launch_job(job_spec):
   """Launch job on ML Engine."""
   project_id = "projects/{}".format(
-      text_encoder.native_to_unicode(cloud.default_project()))
+      text_encoder.native_to_unicode(default_project()))
   credentials = GoogleCredentials.get_application_default()
   cloudml = discovery.build("ml", "v1", credentials=credentials,
                             cache_discovery=False)
@@ -158,13 +223,13 @@ def _tar_and_copy(src_dir, target_dir):
   target_dir = target_dir.rstrip("/")
   tmp_dir = tempfile.gettempdir().rstrip("/")
   src_base = os.path.basename(src_dir)
-  cloud.shell_run(
+  shell_run(
       "tar -zcf {tmp_dir}/{src_base}.tar.gz -C {src_dir} .",
       src_dir=src_dir,
       src_base=src_base,
       tmp_dir=tmp_dir)
   final_destination = "%s/%s.tar.gz" % (target_dir, src_base)
-  cloud.shell_run(
+  shell_run(
       ("gsutil cp {tmp_dir}/{src_base}.tar.gz "
        "{final_destination}"),
       tmp_dir=tmp_dir,
@@ -177,7 +242,7 @@ def tar_and_copy_t2t(train_dir):
   """Tar Tensor2Tensor and cp to train_dir."""
   tf.logging.info("Tarring and pushing local Tensor2Tensor package.")
 
-  output = text_encoder.native_to_unicode(cloud.shell_output(
+  output = text_encoder.native_to_unicode(shell_output(
       "pip show tensor2tensor")).split("\n")
   assert output[1].startswith("Version")
   assert output[7].startswith("Location")
@@ -295,6 +360,11 @@ def validate_flags():
                                                   "complex_model_l"]
 
 
+def confirm():
+  out = input("Confirm (Y/n)? > ")
+  return out == "Y"
+
+
 def launch():
   """Launch t2t_trainer on Cloud ML Engine."""
   validate_flags()
@@ -302,7 +372,7 @@ def launch():
   job_name = job_spec["jobId"]
   tf.logging.info("Launching job %s with ML Engine spec:\n%s", job_name,
                   job_spec)
-  assert cloud.confirm()
+  assert confirm()
   train_dir = FLAGS.output_dir
   t2t_tar = tar_and_copy_t2t(train_dir)
   configure_trainer_package(job_spec, t2t_tar)
diff --git a/tensor2tensor/utils/cloud_tpu.py b/tensor2tensor/utils/cloud_tpu.py
deleted file mode 100644
index c7edac244..000000000
--- a/tensor2tensor/utils/cloud_tpu.py
+++ /dev/null
@@ -1,409 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-"""Launch on TPU on GCP."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import contextlib
-import json
-import multiprocessing.pool as mp
-import os
-import random
-import signal
-import socket
-import subprocess as sp
-import time
-
-from six.moves import input  # pylint: disable=redefined-builtin
-from tensor2tensor.data_generators import text_encoder
-import tensorflow as tf
-
-TPU_IP = "10.240.%d.2"
-TPU_PORT = 8470
-TPU_PROFILE_PORT = 8466
-TB_PORT = 6006
-
-# TODO(rsepassi):
-# --cloud_zone
-# --cloud_project
-
-
-class CloudState(object):
-  """Manage state across multiple trainer runs."""
-
-  def __init__(self):
-    self._tmp_dir = os.path.expanduser("~/.t2t/cloud_state")
-    tf.gfile.MakeDirs(self._tmp_dir)
-
-  def cleanup(self, current_vm_name=None, current_tpu_name=None,
-              skip_confirmation=False):
-    """Delete old instances and cleanup old trainer and tunnel processes."""
-    process_pids = os.listdir(self._tmp_dir)
-    for pid in process_pids:
-      try:
-        # Check if trainer pid is still running
-        os.kill(int(pid), 0)
-      except OSError:
-        # Trainer died ungracefully
-        pid_file = os.path.join(self._tmp_dir, pid)
-        with tf.gfile.Open(pid_file) as f:
-          info = json.loads(f.read())
-
-        # Kill possibly zombie tunnel process
-        try:
-          os.kill(info["tunnel_pid"], signal.SIGTERM)
-        except OSError:
-          pass
-
-        # Delete VM and TPU if requested
-        del_vm = False
-        del_tpu = False
-        if info["delete_on_done"]:
-          if (info["vm_name"] != current_vm_name and
-              info["vm_name"] in list(zip(*list_vm_names_and_ips()))[0]):
-            print("Old VM %s found. Delete?" % info["vm_name"])
-            if skip_confirmation:
-              del_vm = True
-            else:
-              if confirm():
-                del_vm = True
-          if (info["tpu_name"] != current_tpu_name and
-              info["tpu_name"] in list(zip(*list_tpu_names_and_ips()))[0]):
-            print("Old TPU %s found. Delete?" % info["tpu_name"])
-            if skip_confirmation:
-              del_tpu = True
-            else:
-              if confirm():
-                del_tpu = True
-
-        results = []
-        pool = mp.Pool(2)
-        if del_vm:
-          results.append(pool.apply_async(delete_vm, (info["vm_name"],)))
-        if del_tpu:
-          results.append(pool.apply_async(delete_tpu, (info["tpu_name"],)))
-        _ = [res.get() for res in results]
-
-        # Remove the now cleaned up state file
-        tf.gfile.Remove(pid_file)
-
-  def delete_current(self):
-    pid_file = os.path.join(self._tmp_dir, str(os.getpid()))
-    if tf.gfile.Exists(pid_file):
-      tf.gfile.Remove(pid_file)
-
-  def add_current(self, tunnel_pid, vm_name, tpu_name, delete_on_done):
-    state = {
-        "tunnel_pid": tunnel_pid,
-        "vm_name": vm_name,
-        "tpu_name": tpu_name,
-        "delete_on_done": delete_on_done,
-    }
-
-    with tf.gfile.Open(os.path.join(self._tmp_dir, str(os.getpid())), "w") as f:
-      f.write(json.dumps(state))
-
-
-@contextlib.contextmanager
-def cloud_tpu(vm_name, tpu_name, delete_on_done=False, skip_confirmation=False):
-  """Gets or creates a VM and TPU instance, and forwards ports.
-
-  Args:
-    vm_name: str, name of VM.
-    tpu_name: str, name of TPU instance.
-    delete_on_done: bool, whether to delete the instances when done.
-    skip_confirmation: bool, whether to skip launch confirmations.
-
-  Yields:
-    master: str, grpc master pointing to the TPU instance.
-  """
-  state = CloudState()
-  # Read state from previous processes and possibly cleanup
-  state.cleanup(current_vm_name=vm_name, current_tpu_name=tpu_name,
-                skip_confirmation=skip_confirmation)
-
-  done_str = "" if delete_on_done else "NOT "
-  print("Will %sdelete VM and TPU instance on done." % done_str)
-  if not skip_confirmation:
-    assert confirm()
-  _, tpu_ip = create_vm_tpu_pair(vm_name, tpu_name,
-                                 skip_confirmation=skip_confirmation)
-  with tpu_tunnel(vm_name, tpu_ip) as (local_ports, tunnel_pid):
-    master = "grpc://localhost:%d" % local_ports["tpu"]
-
-    state.add_current(tunnel_pid, vm_name, tpu_name, delete_on_done)
-
-    yield master
-
-  if delete_on_done:
-    pool = mp.Pool(2)
-    vm_res = pool.apply_async(delete_vm, (vm_name,))
-    tpu_res = pool.apply_async(delete_tpu, (tpu_name,))
-    vm_res.get()
-    tpu_res.get()
-
-  # Cleanup state from this process
-  state.delete_current()
-
-
-class Gcloud(object):
-  """gcloud command strings."""
-  # Note these can be modified by set_versions
-  VM_VERSION = "tf-1-9"
-  TPU_VERSION = "1.9"
-
-  @classmethod
-  def set_versions(cls, vm, tpu):
-    cls.VM_VERSION = vm
-    cls.TPU_VERSION = tpu
-
-  @classmethod
-  def create_vm(cls):
-    create_vm_str = """
-    gcloud compute instances create {name} \
-      --machine-type=n1-standard-8 \
-      --image-family=%s \
-      --image-project=ml-images \
-      --scopes=https://www.googleapis.com/auth/cloud-platform
-    """ % cls.VM_VERSION
-    return create_vm_str
-
-  DELETE_VM = "gcloud compute instances delete {name} --quiet"
-
-  @classmethod
-  def create_tpu(cls):
-    create_tpu_str = """
-    gcloud beta compute tpus create \
-      {name} \
-      --range={tpu_ip}/29 \
-      --version=%s
-    """ % cls.TPU_VERSION
-    return create_tpu_str
-
-  DELETE_TPU = "gcloud beta compute tpus delete {name} --quiet"
-
-  LIST_TPU = "gcloud beta compute tpus list"
-  LIST_VM = "gcloud compute instances list"
-
-  SSH_LOCAL_PORT_FORWARD = "-L {local_port}:{host}:{remote_port}"
-  SSH_TUNNEL = """
-  gcloud compute ssh {name} -- -N
-  """
-
-  DEFAULT_PROJECT = "gcloud config get-value project"
-  DEFAULT_REGION = "gcloud config get-value compute/region"
-
-
-@contextlib.contextmanager
-def shell_background(cmd_, **kwargs):
-  """Run process in background, join on exit."""
-  args = format_cmd(cmd_, **kwargs)
-  process = sp.Popen(args)
-  try:
-    yield process
-  finally:
-    if process.poll() is None:
-      process.terminate()
-      time.sleep(1)
-    if process.poll() is None:
-      process.kill()
-      time.sleep(1)
-    if process.poll() is None:
-      raise ValueError(
-          "Cannot kill process %d - please kill manually" % process.pid)
-    time.sleep(1)
-
-
-def shell_output(cmd_, **kwargs):
-  return text_encoder.to_unicode(sp.check_output(format_cmd(cmd_, **kwargs)))
-
-
-def shell_run(cmd_, **kwargs):
-  return sp.check_call(format_cmd(cmd_, **kwargs))
-
-
-def format_cmd(cmd_, **kwargs):
-  return cmd_.format(**kwargs).strip().split()
-
-
-def default_region():
-  return shell_output(Gcloud.DEFAULT_REGION).strip()
-
-
-def default_project():
-  return shell_output(Gcloud.DEFAULT_PROJECT).strip()
-
-
-def create_vm(vm_name):
-  out = shell_output(Gcloud.create_vm(), name=vm_name)
-  return out.split("\n")[1:-1][0].split()[4]
-
-
-def list_tpu_names_and_ips():
-  list_out = shell_output(Gcloud.LIST_TPU)
-  lines = [l.split() for l in list_out.split("\n")[1:-1]]
-  names_and_ips = [(l[0].strip(), l[3].strip().split(":")[0]) for l in lines]
-  return names_and_ips
-
-
-def list_vm_names_and_ips():
-  list_out = shell_output(Gcloud.LIST_VM)
-  lines = [l.split() for l in list_out.split("\n")[1:-1]]
-  names_and_ips = [(l[0].strip(), l[4].strip()) for l in lines]
-  return names_and_ips
-
-
-def unique_tpu_ip(tpu_names_and_ips):
-  inuse = [el[1].split(".")[2] for el in tpu_names_and_ips]
-  selection = random.choice(list(set(range(256)) - set(inuse)))
-  return TPU_IP % selection
-
-
-def delete_tpu(tpu_name):
-  shell_run(Gcloud.DELETE_TPU, name=tpu_name)
-
-
-def delete_vm(vm_name):
-  shell_run(Gcloud.DELETE_VM, name=vm_name)
-
-
-def create_tpu(tpu_name, tpu_names_and_ips=None):
-  tpu_names_and_ips = tpu_names_and_ips or list_tpu_names_and_ips()
-  tpu_ip = unique_tpu_ip(tpu_names_and_ips)
-
-  rounded_tpu_ip = tpu_ip
-  if rounded_tpu_ip.endswith("2"):
-    rounded_tpu_ip = rounded_tpu_ip[:-1] + "0"
-
-  shell_run(Gcloud.create_tpu(), name=tpu_name, tpu_ip=rounded_tpu_ip)
-  return tpu_ip
-
-
-@contextlib.contextmanager
-def tpu_tunnel(vm_name, tpu_ip):
-  """Forward TPU and TPU profiling ports."""
-  local_ports = {
-      "tpu": get_open_port(),
-      "tpu_profile": get_open_port(),
-  }
-
-  tpu = format_cmd(
-      Gcloud.SSH_LOCAL_PORT_FORWARD,
-      local_port=local_ports["tpu"],
-      host=tpu_ip,
-      remote_port=TPU_PORT)
-  tpu_profile = format_cmd(
-      Gcloud.SSH_LOCAL_PORT_FORWARD,
-      local_port=local_ports["tpu_profile"],
-      host=tpu_ip,
-      remote_port=TPU_PROFILE_PORT)
-
-  args = format_cmd(Gcloud.SSH_TUNNEL, name=vm_name) + tpu + tpu_profile
-  # Launch process running in background
-  with shell_background(" ".join(args)) as tunnel_process:
-    time.sleep(1)
-    if tunnel_process.poll() is not None:
-      raise ValueError("SSH failed")
-    tf.logging.info("Set up port forwarding. Local ports: %s", local_ports)
-    yield local_ports, tunnel_process.pid
-
-
-def create_vm_tpu_pair(vm_name, tpu_name, reuse_if_exists=True,
-                       skip_confirmation=False):
-  """Create a VM and paired TPU instance.
-
-  Args:
-    vm_name: str, name for VM.
-    tpu_name: str, name for TPU instance.
-    reuse_if_exists: bool, if True, this will act as a get or create. If False
-      and vm_name or tpu_name already exists, will error.
-    skip_confirmation: bool, whether to skip launch confirmations.
-
-  Returns:
-    tuple: (vm_ip, tpu_ip)
-
-  Raises:
-    ValueError: if instance exists but reuse_if_exists=False.
-  """
-  vm_info = list_vm_names_and_ips()
-  tpu_info = list_tpu_names_and_ips()
-
-  vm_names = list(zip(*vm_info))[0] if vm_info else []
-  tpu_names = list(zip(*tpu_info))[0] if tpu_info else []
-
-  make_vm = False
-  vm_ip = None
-  if vm_name in vm_names:
-    if not reuse_if_exists:
-      raise ValueError(
-          "VM %s already exists and reuse_if_exists=False" % vm_name)
-    tf.logging.info("VM %s already exists, reusing.", vm_name)
-    vm_ip = vm_info[vm_names.index(vm_name)][1]
-  else:
-    print("Creating VM %s" % vm_name)
-    if not skip_confirmation:
-      assert confirm()
-    make_vm = True
-
-  make_tpu = False
-  tpu_ip = None
-  if tpu_name in tpu_names:
-    if not reuse_if_exists:
-      raise ValueError(
-          "TPU instance %s already exists and reuse_if_exists=False" % tpu_name)
-    tf.logging.info("TPU %s already exists, reusing.", tpu_name)
-    tpu_ip = tpu_info[tpu_names.index(tpu_name)][1]
-  else:
-    print("Creating TPU instance %s" % tpu_name)
-    if not skip_confirmation:
-      assert confirm()
-    make_tpu = True
-
-  # Create VM and TPU in parallel
-  pool = mp.Pool(2)
-  vm_res = None
-  tpu_res = None
-  if make_vm:
-    vm_res = pool.apply_async(create_vm, (vm_name,))
-  if make_tpu:
-    tpu_res = pool.apply_async(create_tpu, (tpu_name, tpu_info))
-  if vm_res is not None:
-    vm_ip = vm_res.get()
-  if tpu_res is not None:
-    tpu_ip = tpu_res.get()
-
-  tf.logging.info("VM (Name, IP): %s, %s", vm_name, vm_ip)
-  tf.logging.info("TPU (Name, IP): %s, %s", tpu_name, tpu_ip)
-  tf.logging.info(
-      "To delete the VM, run: %s", Gcloud.DELETE_VM.format(name=vm_name))
-  tf.logging.info(
-      "To delete the TPU instance, run: %s",
-      Gcloud.DELETE_TPU.format(name=tpu_name))
-  return vm_ip, tpu_ip
-
-
-def get_open_port():
-  s = socket.socket()
-  s.bind(("", 0))
-  s.listen(1)
-  port = s.getsockname()[1]
-  s.close()
-  return port
-
-
-def confirm():
-  out = input("Confirm (Y/n)? > ")
-  return out == "Y"

From 6dcf54a1ce7a2064494c71527e0cb71a235cf979 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Fri, 31 Aug 2018 23:21:54 -0700
Subject: [PATCH 0722/2720] Spatially partitioned convolutions and its use in
 mnist example.

PiperOrigin-RevId: 211188260
---
 .../mesh_tensorflow/mesh_tensorflow.py        | 426 ++++++++++++++----
 tensor2tensor/mesh_tensorflow/mnist.py        |  49 +-
 tensor2tensor/mesh_tensorflow/mtf_layers.py   |   3 +-
 3 files changed, 367 insertions(+), 111 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index f0623064d..47d746160 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -1954,13 +1954,12 @@ def _einsum_helper(input_shapes, output_shape, mesh_impl):
     einsum_slice_fn: a function from tf.Tensors to tf.Tensor
     reduced_mesh_axes: a list of integers
   """
-  input_shape_set = set(sum([s.dims for s in input_shapes], []))
-  total_num_dims = len(input_shape_set)
+  input_shape_union = _shape_union(input_shapes)
+  total_num_dims = input_shape_union.ndims
   # list of input shapes that contain all dimensions.
   full_shapes = [
       s for s in input_shapes + [output_shape] if s.ndims == total_num_dims]
-  full_shape = (
-      full_shapes[0] if full_shapes else Shape(list(input_shape_set)))
+  full_shape = full_shapes[0] if full_shapes else input_shape_union
   reduce_slice_fn, reduced_mesh_axes = _reduce_helper(
       full_shape, output_shape, mesh_impl.tensor_layout(full_shape))
   def einsum_slice_fn_naive(*slices):
@@ -2040,140 +2039,200 @@ def add_counter_fn():
 class Conv2dOperation(Operation):
   """like tf.nn.conv2d.
 
-  Always "NHWC".
-  Always padding="SAME"
-  Always stride 1
-  Always dilation 1
+  Always data format "NHWC".
+  # TODO(nikip): support dilations
+  Always dilation rate of 1
+  padding: "SAME" or "VALID"
 
   TODO(noam): implement more options.
   """
 
-  def __init__(self, conv_input, conv_filter, is_backprop=False, name=None):
+  def __init__(self, conv_input, conv_filter, strides, padding, name=None):
     super(Conv2dOperation, self).__init__(
         [conv_input, conv_filter], name=name or "conv2d")
-    self._n_dim, self._h_dim, self._w_dim, self._in_dim = conv_input.shape.dims
+    self._padding = padding
+    self._batch_dims = conv_input.shape.dims[:-3]
+    self._in_h_dim, self._in_w_dim, self._in_dim = conv_input.shape.dims[-3:]
     self._fh_dim, self._fw_dim = conv_filter.shape.dims[:2]
-    if is_backprop:
-      self._out_dim, f_in_dim = conv_filter.shape.dims[2:]
-    else:
-      f_in_dim, self._out_dim = conv_filter.shape.dims[2:]
-    self._is_backprop = is_backprop
+    f_in_dim, self._out_dim = conv_filter.shape.dims[2:]
     if f_in_dim != self._in_dim:
       raise ValueError("Dimensions do not match input=%s filter=%s"
                        % (conv_input, conv_filter))
-    output_shape = Shape([self._n_dim, self._h_dim, self._w_dim, self._out_dim])
+    out_h = self._in_h_dim.size
+    out_w = self._in_w_dim.size
+    if padding == "VALID":
+      out_h -= (self._fh_dim.size - 1)
+      out_w -= (self._fw_dim.size - 1)
+
+    self._strides = strides
+    if strides is not None:
+      out_h //= strides[1]
+      out_w //= strides[2]
+    self._out_h_dim = Dimension(self._in_h_dim.name, out_h)
+    self._out_w_dim = Dimension(self._in_w_dim.name, out_w)
+    output_shape = Shape(
+        self._batch_dims + [self._out_h_dim, self._out_w_dim, self._out_dim])
     self._outputs = [Tensor(self, output_shape, conv_input.dtype)]
 
   def gradient(self, grad_ys):
-    if self._is_backprop:
-      raise ValueError("Gradient not implemented for conv backprop")
     dy = grad_ys[0]
     conv_input, conv_filter = self.inputs
     return [
-        conv2d(dy, conv_filter, is_backprop=True),
-        conv2d_backprop_filter(conv_input, self.inputs[1].shape, dy)]
+        conv2d_backprop_input(self._inputs[0].shape,
+                              conv_filter,
+                              dy,
+                              self._strides,
+                              self._padding),
+        conv2d_backprop_filter(conv_input,
+                               self._inputs[1].shape,
+                               dy,
+                               self._strides,
+                               self._padding)]
 
   def lower(self, lowering):
     mesh_impl = lowering.mesh_impl(self)
     conv_input, conv_filter = self.inputs
-    # TODO(noam): support splitting h_dim, w_dim
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._h_dim) is not None:
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._in_h_dim) is not None:
       raise ValueError("can't slice along dimension h")
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._w_dim) is not None:
+    if mesh_impl.tensor_dimension_to_mesh_axis(self._in_w_dim) is not None:
       raise ValueError("can't slice along dimension w")
     if mesh_impl.tensor_dimension_to_mesh_axis(self._fh_dim) is not None:
       raise ValueError("can't slice along dimension fh")
     if mesh_impl.tensor_dimension_to_mesh_axis(self._fw_dim) is not None:
       raise ValueError("can't slice along dimension fw")
     def tf_fn(tf_input, tf_filter):
-      if self._is_backprop:
-        input_sizes = mesh_impl.slice_shape(self.outputs[0].shape)
-        return tf.nn.conv2d_backprop_input(
-            input_sizes, tf_filter, tf_input,
-            strides=[1, 1, 1, 1], padding="SAME")
-      else:
-        return tf.nn.conv2d(
-            tf_input, tf_filter, strides=[1, 1, 1, 1], padding="SAME")
+      output = tf.nn.conv2d(
+          _tf_flatten_batch_dims(tf_input, 3),
+          tf_filter, self._strides, self._padding)
+      return _tf_restore_batch_dims(output, 3, tf_input)
     y = mesh_impl.slicewise(
         tf_fn, lowering.tensors[conv_input], lowering.tensors[conv_filter])
-    out_mesh_axis = mesh_impl.tensor_dimension_to_mesh_axis(self._out_dim)
-    if out_mesh_axis is not None:
+    # reducing out input channels - may need to allreduce
+    in_mesh_axis = mesh_impl.tensor_dimension_to_mesh_axis(self._in_dim)
+    if in_mesh_axis is not None:
       def add_counter_fn():
         lowering.add_counter(
-            "allreduce/%s/conv2d_op" % [out_mesh_axis],
+            "allreduce/%s/conv2d_op" % [in_mesh_axis],
             mesh_impl.laid_out_size(self.outputs[0].shape))
-      y = LazyAllreduceSum(mesh_impl, y, [out_mesh_axis], add_counter_fn)
+      y = LazyAllreduceSum(mesh_impl, y, [in_mesh_axis], add_counter_fn)
     lowering.set_tensor_lowering(self.outputs[0], y)
-    input_shape_set = set(sum([x.shape.dims for x in self.inputs], []))
-    computation_shape = Shape(list(input_shape_set))
-    lowering.add_counter("conv2d", mesh_impl.laid_out_size(computation_shape))
-    lowering.add_counter("conv2d_unique", computation_shape.size)
+    computation_shape = _shape_union([conv_filter.shape, self.outputs[0].shape])
+    lowering.add_counter("conv2d/forward",
+                         mesh_impl.laid_out_size(computation_shape))
+    lowering.add_counter("conv2d_unique/forward", computation_shape.size)
 
 
-def conv2d(conv_input, conv_filter, is_backprop=False, name=None):
+def conv2d(conv_input, conv_filter, strides, padding, name=None):
   """conv2d."""
   return Conv2dOperation(
-      conv_input, conv_filter, is_backprop, name=name).outputs[0]
+      conv_input, conv_filter, strides, padding, name=name).outputs[0]
+
+
+class Conv2dBackpropInputOperation(Operation):
+  """like tf.nn.conv2d_backprop_input"""
+
+  def __init__(self, input_shape, conv_filter, dy, strides, padding, name=None):
+    super(Conv2dBackpropInputOperation, self).__init__(
+        [dy, conv_filter], name=name or "conv2d_backprop")
+    self._padding = padding
+    self._strides = strides
+    self._input_shape = input_shape
+    self._outputs = [Tensor(self, input_shape, dy.dtype)]
+
+  def lower(self, lowering):
+    mesh_impl = lowering.mesh_impl(self)
+    dy, conv_filter = self.inputs
+    input_sizes = mesh_impl.slice_shape(self.outputs[0].shape)
+    input_sizes = [list_product(input_sizes[:-3])] + input_sizes[-3:]
+    def tf_fn(tf_dy, tf_filter):
+      return _tf_restore_batch_dims(
+          tf.nn.conv2d_backprop_input(
+              input_sizes, tf_filter, _tf_flatten_batch_dims(tf_dy, 3),
+              self._strides, self._padding), 3, tf_dy)
+    dx = mesh_impl.slicewise(
+        tf_fn, lowering.tensors[dy], lowering.tensors[conv_filter])
+    # reducing out output channels - may need to allreduce
+    out_mesh_axis = mesh_impl.tensor_dimension_to_mesh_axis(dy.shape.dims[-1])
+    if out_mesh_axis is not None:
+      def add_counter_fn():
+        lowering.add_counter(
+            "allreduce/%s/conv2d_op" % [out_mesh_axis],
+            mesh_impl.laid_out_size(self.outputs[0].shape))
+      dx = LazyAllreduceSum(mesh_impl, dx, [out_mesh_axis], add_counter_fn)
+    lowering.set_tensor_lowering(self.outputs[0], dx)
+    computation_shape = _shape_union([conv_filter.shape, dy.shape])
+    lowering.add_counter("conv2d/backprop_input",
+                         mesh_impl.laid_out_size(computation_shape))
+    lowering.add_counter("conv2d_unique/backprop_input", computation_shape.size)
+
+
+def conv2d_backprop_input(input_shape,
+                          conv_filter,
+                          dy,
+                          strides,
+                          padding, name=None):
+  return Conv2dBackpropInputOperation(input_shape,
+                                      conv_filter,
+                                      dy,
+                                      strides,
+                                      padding,
+                                      name=name).outputs[0]
 
 
 class Conv2dBackpropFilterOperation(Operation):
-  """like tf.nn.conv2d_backprop_filter."""
+  """like tf.nn.conv2d_backprop_input"""
 
-  def __init__(self, conv_input, filter_shape, dy, name=None):
+  def __init__(self, conv_input, filter_shape, dy, strides, padding, name=None):
     super(Conv2dBackpropFilterOperation, self).__init__(
         [conv_input, dy], name=name or "conv2d_backprop_filter")
-    self._n_dim, self._h_dim, self._w_dim, self._in_dim = conv_input.shape.dims
-    dy_n_dim, dy_h_dim, dy_w_dim, self._out_dim = dy.shape.dims
-    self._fh_dim, self._fw_dim, f_in_dim, f_out_dim = filter_shape.dims
-    if (dy_n_dim != self._n_dim or
-        dy_h_dim != self._h_dim or
-        dy_w_dim != self._w_dim or
-        f_in_dim != self._in_dim or
-        f_out_dim != self._out_dim):
-      raise ValueError("Dimensions do not match input=%s dy=%s filter=%s"
-                       % (conv_input, dy, filter_shape))
-    self._outputs = [Tensor(self, filter_shape, conv_input.dtype)]
+    self._padding = padding
+    self._strides = strides
+    self._filter_shape = filter_shape
+    self._outputs = [Tensor(self, filter_shape, dy.dtype)]
 
   def lower(self, lowering):
     mesh_impl = lowering.mesh_impl(self)
     conv_input, dy = self.inputs
-    # TODO(noam): support splitting h_dim, w_dim
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._h_dim) is not None:
-      raise ValueError("can't slice along dimension h")
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._w_dim) is not None:
-      raise ValueError("can't slice along dimension w")
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._fh_dim) is not None:
-      raise ValueError("can't slice along dimension fh")
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._fw_dim) is not None:
-      raise ValueError("can't slice along dimension fw")
     filter_sizes = mesh_impl.slice_shape(self.outputs[0].shape)
     def tf_fn(tf_input, tf_dy):
       return tf.nn.conv2d_backprop_filter(
-          tf_input, filter_sizes, tf_dy, strides=[1, 1, 1, 1], padding="SAME")
-    y = mesh_impl.slicewise(
+          _tf_flatten_batch_dims(tf_input, 3), filter_sizes,
+          _tf_flatten_batch_dims(tf_dy, 3), self._strides, self._padding)
+    df = mesh_impl.slicewise(
         tf_fn, lowering.tensors[conv_input], lowering.tensors[dy])
+
+    # reducing out batch dimensions - may need to allreduce
     reduced_mesh_axes = [
         mesh_impl.tensor_dimension_to_mesh_axis(d)
-        for d in [self._n_dim, self._h_dim, self._w_dim]]
+        for d in dy.shape.dims[:-3]]
     reduced_mesh_axes = [a for a in reduced_mesh_axes if a is not None]
+
     if reduced_mesh_axes:
       def add_counter_fn():
         lowering.add_counter(
-            "allreduce/%s/conv2d_op" % (reduced_mesh_axes,),
+            "allreduce/%s/conv2d_backprop_filter" % (reduced_mesh_axes,),
             mesh_impl.laid_out_size(self.outputs[0].shape))
-      y = LazyAllreduceSum(mesh_impl, y, reduced_mesh_axes, add_counter_fn)
-    lowering.set_tensor_lowering(self.outputs[0], y)
-    input_shape_set = set(sum([x.shape.dims for x in self.inputs], []))
-    computation_shape = Shape(list(input_shape_set))
-    lowering.add_counter("conv2d", mesh_impl.laid_out_size(computation_shape))
-    lowering.add_counter("conv2d_unique", computation_shape.size)
+      df = LazyAllreduceSum(mesh_impl, df, reduced_mesh_axes, add_counter_fn)
 
+    lowering.set_tensor_lowering(self.outputs[0], df)
+    computation_shape = _shape_union([self.outputs[0].shape, dy.shape])
+    lowering.add_counter("conv2d/backprop_filter",
+                         mesh_impl.laid_out_size(computation_shape))
+    lowering.add_counter(
+        "conv2d_unique/backprop_filter", computation_shape.size)
 
-def conv2d_backprop_filter(
-    conv_input, filter_shape, dy, name=None):
-  """conv2d."""
-  return Conv2dBackpropFilterOperation(
-      conv_input, filter_shape, dy, name=name).outputs[0]
+
+def conv2d_backprop_filter(conv_input,
+                           filter_shape,
+                           dy,
+                           strides,
+                           padding, name=None):
+  return Conv2dBackpropFilterOperation(conv_input,
+                                       filter_shape,
+                                       dy,
+                                       strides,
+                                       padding,
+                                       name=name).outputs[0]
 
 
 class ShiftOperation(Operation):
@@ -2210,32 +2269,60 @@ def lower(self, lowering):
     ndims = self._inputs[0].shape.ndims
     axis = self._axis
     dim = self._dim
+    lowered_x = lowering.tensors[inputs]
+    def my_slice(x, start, size):
+      begin = [0] * axis + [start] + [0] * (ndims - axis - 1)
+      size = [-1] * axis + [size] + [-1] * (ndims - axis - 1)
+      return tf.slice(x, begin, size)
     if mesh_axis is None:
       def slicewise_fn(x):
         """Slicewise function."""
-        def my_slice(start, size):
-          begin = [0] * axis + [start] + [0] * (ndims - axis - 1)
-          size = [-1] * axis + [size] + [-1] * (ndims - axis - 1)
-          return tf.slice(x, begin, size)
         def my_pad(s, begin_pad, end_pad):
-          paddings = ([[0, 0]] * axis + [begin_pad, end_pad]
+          paddings = ([[0, 0]] * axis + [[begin_pad, end_pad]]
                       + [[0, 0]] * (ndims - axis - 1))
           return tf.pad(s, paddings)
         if self._wrap:
           offset = self._offset % dim.size
-          return tf.concat([my_slice(dim.size - offset, offset),
-                            my_slice(0, dim.size - offset)], axis=axis)
+          return tf.concat([my_slice(x, dim.size - offset, offset),
+                            my_slice(x, 0, dim.size - offset)], axis=axis)
         elif self._offset > 0:
-          return my_pad(my_slice(0, dim.size - self._offset), self._offset, 0)
+          return my_pad(
+              my_slice(x, 0, dim.size - self._offset), self._offset, 0)
         else:
           neg_offset = -self._offset
           return my_pad(
-              my_slice(neg_offset, dim.size - neg_offset), 0, neg_offset)
-      y = mesh_impl.slicewise(slicewise_fn, lowering.tensors[inputs])
+              my_slice(x, neg_offset, dim.size - neg_offset), 0, neg_offset)
+      lowered_y = mesh_impl.slicewise(slicewise_fn, lowered_x)
     else:
-      raise NotImplementedError(
-          "TODO(noam): implement this using mesh_impl.shift_by_n_processors")
-    lowering.set_tensor_lowering(self.outputs[0], y)
+      mesh_dim_size = mesh_impl.shape.dims[mesh_axis].size
+      tensor_dim_size = self._dim.size
+      block_size = tensor_dim_size // mesh_dim_size
+      odiv = self._offset // block_size
+      omod = self._offset % block_size
+      laid_out_size = mesh_impl.laid_out_size(inputs.shape)
+      if omod == 0:
+        # shift by an integral number of processors.
+        lowered_y = mesh_impl.shift_by_n_processors(
+            lowered_x, mesh_axis, odiv, self._wrap)
+        lowering.add_counter("shift[%d]" % odiv, laid_out_size)
+      else:
+        # shift by odiv processors + omod positions
+        sliced = mesh_impl.slicewise(
+            lambda x: my_slice(x, 0, block_size - omod), lowered_x)
+        second_part = mesh_impl.shift_by_n_processors(
+            sliced, mesh_axis, odiv, self._wrap)
+        lowering.add_counter(
+            "shift[%d]" % odiv,
+            laid_out_size * (block_size - omod) // block_size)
+        sliced = mesh_impl.slicewise(
+            lambda x: my_slice(x, block_size - omod, omod), lowered_x)
+        first_part = mesh_impl.shift_by_n_processors(
+            sliced, mesh_axis, odiv + 1, self._wrap)
+        lowered_y = mesh_impl.slicewise(
+            lambda a, b: tf.concat([a, b], axis), first_part, second_part)
+        lowering.add_counter(
+            "shift[%d]" % (odiv + 1), laid_out_size * omod // block_size)
+    lowering.set_tensor_lowering(self.outputs[0], lowered_y)
 
 
 def shift(x, offset, dim, wrap, name=None):
@@ -2289,7 +2376,7 @@ def lower(self, lowering):
     ndims = self._inputs[0].shape.ndims
     axis = self._axis
     begin = [0] * axis + [self._begin] + [0] * (ndims - axis - 1)
-    size = [-1] * axis + [self._slice_dim[1]] + [-1] * (ndims - axis - 1)
+    size = [-1] * axis + [self._slice_dim.size] + [-1] * (ndims - axis - 1)
 
     def slicewise_fn(x, begin, size):
       return tf.slice(x, begin, size, name="slice")
@@ -2323,7 +2410,6 @@ def __init__(self, x, paddings, pad_dim_name, name=None):
     self._outputs = [Tensor(self, output_shape, x.dtype)]
 
   def gradient(self, grad_ys):
-    # slice_dim = self._inputs[0].shape.dims[self._axis]
     slice_dim_name = self._output_dim.name
     slice_size = self._inputs[0].shape.dims[self._axis].size
     return [slice(grad_ys[0], self._paddings[0], slice_size, slice_dim_name)]
@@ -3801,3 +3887,157 @@ def where(condition, if_true, if_false):
   return (
       if_true * cast(condition, dtype) +
       if_false * cast(logical_not(condition), dtype))
+
+
+def _shape_union(shapes):
+  """A shape containing the union of all dimensions in the input shapes.
+
+  Args:
+    shapes: a list of Shapes
+
+  Returns:
+    a Shape
+  """
+  return Shape(list(set(sum([s.dims for s in shapes], []))))
+
+
+def _tf_flatten_batch_dims(x, num_nonbatch_dims):
+  """Flatten all but last num_nonbatch_dims into one dimension.
+
+  Args:
+    x: a tf.Tensor:
+    num_nonbatch_dims: an integer
+
+  Returns:
+    a tf.Tensor with 1 + num_nonbatch_dims dimensions.
+  """
+  shape = x.shape.as_list()
+  assert None not in shape
+  new_shape = ([list_product(shape[:-num_nonbatch_dims])]
+               + shape[-num_nonbatch_dims:])
+  if new_shape != shape:
+    x = tf.reshape(x, new_shape)
+  return x
+
+
+def _tf_restore_batch_dims(x, num_nonbatch_dims, prototype):
+  """Reverse op of _tf_flatten_batch_dims.
+
+  Un-flatten the first dimension of x to match all but the last
+  num_nonbatch_dims dimensions of prototype.
+
+  Args:
+    x: a tf.Tensor with 1 + num_nonbatch_dims dimensions
+    num_nonbatch_dims: an integer
+    prototype: a tf.Tensor
+
+  Returns:
+    a tf.Tensor
+  """
+  assert x.shape.ndims == 1 + num_nonbatch_dims
+  new_shape = (
+      prototype.shape.as_list()[:-num_nonbatch_dims] + x.shape.as_list()[1:])
+  assert None not in new_shape
+  if new_shape != x.shape.as_list():
+    x = tf.reshape(x, new_shape)
+  return x
+
+
+def halo_exchange(x, blocks_dim, block_size_dim, halo_size, wrap=False):
+  """Concat each block with the margins of adjacent blocks.
+
+  Get left and right blocks_dim and concatenate along block_size_dim.
+
+  Args:
+    x: a Tensor.
+    blocks_dim: a Dimension in x.shape
+    block_size_dim: a Dimension in x.shape
+    halo_size: an integer
+    wrap: a boolean
+
+  Returns:
+    a Tensor with the same shape as x, other than in block_size_dim, whose
+    size is increased by 2*halo_size.
+  """
+  if halo_size == 0:
+    return x
+
+  block_size = block_size_dim.size
+  partial_size = halo_size % block_size
+  num_complete_blocks = halo_size // block_size
+  parts = [x]
+
+  for i in xrange(1, num_complete_blocks + 1):
+    parts = ([shift(x, i, blocks_dim, wrap)] + parts +
+             [shift(x, -i, blocks_dim, wrap)])
+  if partial_size > 0:
+    left_margin = slice(x, 0, partial_size, block_size_dim.name)
+    right_margin = slice(x, block_size_dim.size - partial_size, partial_size,
+                         block_size_dim.name)
+    parts = (
+        [shift(right_margin, num_complete_blocks + 1, blocks_dim, wrap)]
+        + parts +
+        [shift(left_margin, -(num_complete_blocks + 1), blocks_dim, wrap)])
+  return concat(parts, block_size_dim.name)
+
+
+def conv2d_with_blocks(
+    conv_input,
+    conv_filter,
+    strides,
+    padding,
+    h_blocks_dim=None,
+    w_blocks_dim=None,
+    name=None):
+  """conv2d operation with spatial partitioning.
+
+  Spatial partitioning is implemented by decomposing the image into blocks.
+  Block dimensions represented as h_blocks_dim and w_blocks_dim can be split
+  along the mesh axis. If split, then we do a halo exchange where each block
+  receives the part of the image from its left and right neighbors necessary to
+  do the convolution. Exchange can involve complete or partial blocks depending
+  on the filter height and width.
+
+  Currently, only "SAME" padding with dilation rate of 1 is supported.
+
+  Args:
+    conv_input: a Tensor of shape
+      [batch, h_blocks_dim, w_blocks_dim, h_dim, w_dim, in_channels_dim]
+    conv_filter: a Tensor of shape
+      [filter_height, filter_width, in_channels_dim, out_channels_dim]
+    strides: A list of ints. 1-D tensor of length 4.
+    padding: string, "SAME". The type of padding algorithm to use.
+      Valid is not currently supported.
+    h_blocks_dim: Dimension representing number of height blocks.
+    w_blocks_dim: Dimension representing number of height blocks.
+    name: A name for the operation (optional).
+
+  Returns:
+    A Tensor of shape
+      [batch, h_blocks_dim, w_blocks_dim, h_dim, w_dim, out_channels_dim]
+  """
+  filter_h_dim, filter_w_dim = conv_filter.shape.dims[:2]
+  assert filter_h_dim.size % 2 == 1
+  assert filter_w_dim.size % 2 == 1
+  h_dim, w_dim = conv_input.shape.dims[-3:-1]
+
+  # If h_blocks_dim and w_blocks_dim is not split, directly call conv2d.
+  if h_blocks_dim is None and w_blocks_dim is None:
+    return conv2d(conv_input, conv_filter, strides, padding, name)
+
+  # Padding 'VALID' is not supported yet.
+  if padding != "SAME":
+    raise NotImplementedError("conv2d_with_blocks requires padding=SAME")
+
+  # Halo exchange for h_blocks and w_blocks.
+  for blocks_dim, block_size_dim, halo_size in [
+      (h_blocks_dim, h_dim, filter_h_dim.size // 2),
+      (w_blocks_dim, w_dim, filter_w_dim.size // 2)]:
+    if halo_size > 0:
+      if blocks_dim is not None:
+        conv_input = halo_exchange(
+            conv_input, blocks_dim, block_size_dim, halo_size)
+      else:
+        conv_input = pad(
+            conv_input, [halo_size, halo_size], block_size_dim.name)
+  return conv2d(conv_input, conv_filter, strides, "VALID", name)
diff --git a/tensor2tensor/mesh_tensorflow/mnist.py b/tensor2tensor/mesh_tensorflow/mnist.py
index 74df1aec5..3f226b218 100644
--- a/tensor2tensor/mesh_tensorflow/mnist.py
+++ b/tensor2tensor/mesh_tensorflow/mnist.py
@@ -42,8 +42,8 @@
 tf.flags.DEFINE_integer("eval_steps", 0,
                         "Total number of evaluation steps. If `0`, evaluation "
                         "after training is skipped.")
-tf.flags.DEFINE_string("mesh_shape", "rows:2;cols:2", "mesh shape")
-tf.flags.DEFINE_string("layout", "batch:rows;hidden1:cols,filters1:cols",
+tf.flags.DEFINE_string("mesh_shape", "b1:2;b2:2", "mesh shape")
+tf.flags.DEFINE_string("layout", "col_blocks:b1;hidden1:b2;filters2:b2",
                        "layout rules")
 
 FLAGS = tf.flags.FLAGS
@@ -62,28 +62,40 @@ def mnist_model(image, labels, mesh):
     loss: a mtf.Tensor with shape []
   """
   batch_dim = mtf.Dimension("batch", FLAGS.batch_size)
-  rows_dim = mtf.Dimension("rows", 28)
-  cols_dim = mtf.Dimension("cols", 28)
+  row_blocks_dim = mtf.Dimension("row_blocks", 4)
+  col_blocks_dim = mtf.Dimension("col_blocks", 4)
+  rows_dim = mtf.Dimension("rows_size", 7)
+  cols_dim = mtf.Dimension("cols_size", 7)
+
   classes_dim = mtf.Dimension("classes", 10)
   one_channel_dim = mtf.Dimension("one_channel", 1)
 
-  x = mtf.import_tf_tensor(mesh, tf.reshape(image, [-1, 28, 28]),
-                           mtf.Shape([batch_dim, rows_dim, cols_dim]))
-  x = mtf.reshape(x, [batch_dim, rows_dim, cols_dim, one_channel_dim])
+  x = mtf.import_tf_tensor(
+      mesh, tf.reshape(image, [FLAGS.batch_size, 4, 7, 4, 7]),
+      mtf.Shape(
+          [batch_dim, row_blocks_dim, rows_dim, col_blocks_dim, cols_dim]))
+  x = mtf.reshape(x, [
+      batch_dim, row_blocks_dim, col_blocks_dim,
+      rows_dim, cols_dim, one_channel_dim])
 
   # add some convolutional layers to demonstrate that convolution works.
-  # TODO(noam): get spatially-partitioned convolution working.
-  fh_dim = mtf.Dimension("fh", 3)
-  fw_dim = mtf.Dimension("fw", 3)
-  filters1_dim = mtf.Dimension("filters1", 32)
-  filters2_dim = mtf.Dimension("filters2", 32)
+  # TODO(nikip): Currently spatial conv works only when splitting column blocks.
+  # Make it work for both height and width dimension of the image.
+  fh_dim = mtf.Dimension("fh", 9)
+  fw_dim = mtf.Dimension("fw", 9)
+  filters1_dim = mtf.Dimension("filters1", 16)
+  filters2_dim = mtf.Dimension("filters2", 16)
   kernel1 = mtf.get_variable(
       mesh, "kernel1", [fh_dim, fw_dim, one_channel_dim, filters1_dim])
   kernel2 = mtf.get_variable(
       mesh, "kernel2", [fh_dim, fw_dim, filters1_dim, filters2_dim])
 
-  f1 = mtf.relu(mtf.conv2d(x, kernel1))
-  f2 = mtf.relu(mtf.conv2d(f1, kernel2))
+  f1 = mtf.relu(mtf.conv2d_with_blocks(
+      x, kernel1, strides=[1, 1, 1, 1], padding="SAME",
+      h_blocks_dim=None, w_blocks_dim=col_blocks_dim))
+  f2 = mtf.relu(mtf.conv2d_with_blocks(
+      f1, kernel2, strides=[1, 1, 1, 1], padding="SAME",
+      h_blocks_dim=None, w_blocks_dim=None))
   x = mtf.reduce_mean(f2, reduced_dim=filters2_dim)
 
   # add some fully-connected dense layers.
@@ -91,15 +103,18 @@ def mnist_model(image, labels, mesh):
   hidden_dim2 = mtf.Dimension("hidden2", FLAGS.hidden_size)
 
   h1 = mtf_layers.dense(
-      x, hidden_dim1, reduced_dims=[rows_dim, cols_dim],
+      x, hidden_dim1,
+      reduced_dims=x.shape.dims[-4:],
       activation=mtf.relu, name="hidden1")
   h2 = mtf_layers.dense(
-      h1, hidden_dim2, activation=mtf.relu, name="hidden2")
+      h1, hidden_dim2,
+      activation=mtf.relu, name="hidden2")
   logits = mtf_layers.dense(h2, classes_dim, name="logits")
   if labels is None:
     loss = None
   else:
-    labels = mtf.import_tf_tensor(mesh, labels, mtf.Shape([batch_dim]))
+    labels = mtf.import_tf_tensor(
+        mesh, tf.reshape(labels, [FLAGS.batch_size]), mtf.Shape([batch_dim]))
     loss = mtf_layers.softmax_cross_entropy_with_logits(
         logits, mtf.one_hot(labels, classes_dim), classes_dim)
     loss = mtf.reduce_mean(loss)
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers.py b/tensor2tensor/mesh_tensorflow/mtf_layers.py
index 611cc4ac7..b7dd41084 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers.py
@@ -48,6 +48,7 @@ def dense(x, output_dim, reduced_dims=None, expert_dims=None,
   w_shape = mtf.Shape(expert_dims + reduced_dims + [output_dim])
   output_shape = mtf.Shape(
       [d for d in x.shape.dims if d not in reduced_dims] + [output_dim])
+
   with tf.variable_scope(name, default_name="dense"):
     stddev = mtf.list_product(d.size for d in reduced_dims) ** -0.5
     w = mtf.get_variable(
@@ -56,7 +57,7 @@ def dense(x, output_dim, reduced_dims=None, expert_dims=None,
         w_shape,
         initializer=tf.random_normal_initializer(stddev=stddev),
         activation_dtype=x.dtype)
-    y = mtf.matmul(x, w, output_shape=output_shape)
+    y = mtf.einsum([x, w], output_shape)
     if use_bias:
       b = mtf.get_variable(
           x.mesh,

From 6401045cf5ae3f86238b407b0516708d41834b85 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sat, 1 Sep 2018 20:48:27 -0700
Subject: [PATCH 0723/2720] Internal change

PiperOrigin-RevId: 211243289
---
 tensor2tensor/bin/t2t_decoder.py           |  4 ---
 tensor2tensor/bin/t2t_trainer.py           |  3 ++
 tensor2tensor/data_generators/translate.py | 28 ++++++++++++++++++
 tensor2tensor/utils/decoding.py            | 14 +++++++++
 tensor2tensor/utils/flags.py               | 31 ++++++++++++--------
 tensor2tensor/utils/trainer_lib.py         | 33 ++++++++++++++++++----
 6 files changed, 92 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index c7d2814d9..61e50da1a 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -48,10 +48,6 @@
 # Additional flags in bin/t2t_trainer.py and utils/flags.py
 flags.DEFINE_string("checkpoint_path", None,
                     "Path to the model checkpoint. Overrides output_dir.")
-flags.DEFINE_string("decode_from_file", None,
-                    "Path to the source file for decoding")
-flags.DEFINE_string("decode_to_file", None,
-                    "Path to the decoded (output) file")
 flags.DEFINE_bool("keep_timestamp", False,
                   "Set the mtime of the decoded file to the "
                   "checkpoint_path+'.index' mtime.")
diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index b4bc01a4a..ca78354be 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -173,6 +173,9 @@ def create_experiment_fn(**kwargs):
       use_tpu_estimator=FLAGS.use_tpu_estimator,
       use_xla=FLAGS.xla_compile,
       warm_start_from=FLAGS.warm_start_from,
+      decode_from_file=FLAGS.decode_from_file,
+      decode_to_file=FLAGS.decode_to_file,
+      decode_reference=FLAGS.decode_reference,
       **kwargs)
 
 
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index b8bc020cc..478ed80c2 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -24,6 +24,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import bleu_hook
 
 import tensorflow as tf
 
@@ -60,6 +61,33 @@ def generate_text_for_vocab(self, data_dir, tmp_dir):
     return generator_utils.generate_lines_for_vocab(tmp_dir,
                                                     self.vocab_data_files())
 
+  @property
+  def decode_hooks(self):
+    return [compute_bleu_summaries]
+
+
+def compute_bleu_summaries(hook_args):
+  """Compute BLEU core summaries using the decoder output.
+
+  Args:
+    hook_args: DecodeHookArgs namedtuple
+  Returns:
+    A list of tf.Summary values if hook_args.hparams contains the
+    reference file and the translated file.
+  """
+  decode_hparams = hook_args.decode_hparams
+
+  if (decode_hparams.decode_reference is None or
+      decode_hparams.decode_to_file is None):
+    return None
+
+  values = []
+  bleu = 100 * bleu_hook.bleu_wrapper(
+      decode_hparams.decode_reference, decode_hparams.decode_to_file)
+  values.append(tf.Summary.Value(tag="BLEU", simple_value=bleu))
+  tf.logging.info("%s: BLEU = %6.2f" % (decode_hparams.decode_to_file, bleu))
+  return values
+
 
 def _preprocess_sgm(line, is_sgm):
   """Preprocessing to strip tags in SGM files."""
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 84419f982..4a2d888f0 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -427,6 +427,20 @@ def timer(gen):
   outfile = tf.gfile.Open(decode_filename, "w")
   for index in range(len(sorted_inputs)):
     outfile.write("%s%s" % (decodes[sorted_keys[index]], decode_hp.delimiter))
+  outfile.flush()
+  outfile.close()
+
+  output_dir = os.path.join(estimator.model_dir, "decode")
+  tf.gfile.MakeDirs(output_dir)
+
+  run_postdecode_hooks(DecodeHookArgs(
+      estimator=estimator,
+      problem=hparams.problem,
+      output_dirs=[output_dir],
+      hparams=hparams,
+      decode_hparams=decode_hp,
+      predictions=list(result_iter)
+  ), None)
 
 
 def _decode_filename(base_filename, problem_name, decode_hp):
diff --git a/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index 3ae12988a..3cd0c0906 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -34,13 +34,13 @@
 flags.DEFINE_string("model", None, "Which model to use.")
 flags.DEFINE_string("hparams_set", None, "Which parameters to use.")
 flags.DEFINE_string("hparams_range", None, "Parameters range.")
-flags.DEFINE_string(
-    "hparams", "",
-    """A comma-separated list of `name=value` hyperparameter values. This flag
-    is used to override hyperparameter settings either when manually selecting
-    hyperparameters or when using Vizier. If a hyperparameter setting is
-    specified by this flag then it must be a valid hyperparameter name for the
-    model.""")
+flags.DEFINE_string("hparams", "",
+                    "A comma-separated list of `name=value` hyperparameter "
+                    "values. This flag is used to override hyperparameter "
+                    "settings either when manually selecting hyperparameters "
+                    "or when using Vizier. If a hyperparameter setting is "
+                    "specified by this flag then it must be a valid "
+                    "hyperparameter name for the model.")
 flags.DEFINE_string("problem", None, "Problem name.")
 
 # data_dir is a common flag name - catch conflicts and define it once.
@@ -108,7 +108,16 @@
 flags.DEFINE_integer("ps_replicas", 0, "How many ps replicas.")
 
 # Decoding flags
-flags.DEFINE_string(
-    "decode_hparams", "",
-    "Comma-separated list of name=value pairs to control decode behavior. "
-    "See decoding.decode_hparams for defaults.")
+flags.DEFINE_string("decode_hparams", "",
+                    "Comma-separated list of name=value pairs to control "
+                    "decode behavior. See decoding.decode_hparams for "
+                    "defaults.")
+flags.DEFINE_string("decode_from_file", None,
+                    "Path to the source file for decoding, used by "
+                    "continuous_decode_from_file.")
+flags.DEFINE_string("decode_to_file", None,
+                    "Path to the decoded file generated by decoding, used by "
+                    "continuous_decode_from_file.")
+flags.DEFINE_string("decode_reference", None,
+                    "Path to the reference file for decoding, used by "
+                    "continuous_decode_from_file to compute BLEU score.")
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 74c4597b9..b2dd55114 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -394,11 +394,20 @@ def run_std_server(self):
         task_index=config.task_id)
     server.join()
 
-  def decode(self, dataset_split=None):
-    """Decodes from dataset."""
-    decoding.decode_from_dataset(self._estimator, self._hparams.problem.name,
-                                 self._hparams, self._decode_hparams,
-                                 dataset_split=dataset_split)
+  def decode(self, dataset_split=None, decode_from_file=False):
+    """Decodes from dataset or file."""
+    if decode_from_file:
+      decoding.decode_from_file(self._estimator,
+                                self._decode_hparams.decode_from_file,
+                                self._hparams,
+                                self._decode_hparams,
+                                self._decode_hparams.decode_to_file)
+    else:
+      decoding.decode_from_dataset(self._estimator,
+                                   self._hparams.problem.name,
+                                   self._hparams,
+                                   self._decode_hparams,
+                                   dataset_split=dataset_split)
 
   def continuous_decode(self):
     """Decode from dataset on new checkpoint."""
@@ -410,6 +419,11 @@ def continuous_decode_on_train_data(self):
     for _ in next_checkpoint(self._hparams.model_dir):
       self.decode(dataset_split=tf.estimator.ModeKeys.TRAIN)
 
+  def continuous_decode_from_file(self):
+    """Decode from file on new checkpoint."""
+    for _ in next_checkpoint(self._hparams.model_dir):
+      self.decode(decode_from_file=True)
+
 
 def create_experiment(
     run_config,
@@ -436,7 +450,10 @@ def create_experiment(
     use_xla=False,
     additional_train_hooks=None,
     additional_eval_hooks=None,
-    warm_start_from=None):
+    warm_start_from=None,
+    decode_from_file=None,
+    decode_to_file=None,
+    decode_reference=None):
   """Create Experiment."""
   # HParams
   hparams.add_hparam("model_dir", run_config.model_dir)
@@ -445,6 +462,10 @@ def create_experiment(
   hparams.add_hparam("eval_steps", eval_steps)
   hparams.add_hparam("schedule", schedule)
   hparams.add_hparam("warm_start_from", warm_start_from)
+  if decode_hparams is not None:
+    decode_hparams.add_hparam("decode_from_file", decode_from_file)
+    decode_hparams.add_hparam("decode_to_file", decode_to_file)
+    decode_hparams.add_hparam("decode_reference", decode_reference)
   add_problem_hparams(hparams, problem_name)
 
   # Estimator

From a564e9b863837e6173b872118c74051495228c75 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 23 Aug 2018 20:42:24 +0200
Subject: [PATCH 0724/2720] Flip the order of wrappers in the AE experiment so
 that the AE encodes single frames

---
 tensor2tensor/data_generators/gym_problems.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index adedf0325..2fc7922b2 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -65,8 +65,8 @@ def standard_atari_env_spec(env):
 
 def standard_atari_ae_env_spec(env):
   """Parameters of environment specification."""
-  standard_wrappers = [[tf_atari_wrappers.StackWrapper, {"history": 4}],
-                       [tf_atari_wrappers.AutoencoderWrapper, {}]]
+  standard_wrappers = [[tf_atari_wrappers.AutoencoderWrapper, {}],
+                       [tf_atari_wrappers.StackWrapper, {"history": 4}]]
   env_lambda = None
   if isinstance(env, str):
     env_lambda = lambda: gym.make(env)
@@ -594,7 +594,10 @@ class GymSimulatedDiscreteProblemAutoencoded(GymSimulatedDiscreteProblem):
 
   def get_environment_spec(self):
     env_spec = standard_atari_env_spec(self.env_name)
-    env_spec.wrappers = [[tf_atari_wrappers.IntToBitWrapper, {}]]
+    env_spec.wrappers = [
+        [tf_atari_wrappers.IntToBitWrapper, {}],
+        [tf_atari_wrappers.StackWrapper, {"history": 4}]
+    ]
     env_spec.simulated_env = True
     env_spec.add_hparam("simulation_random_starts", False)
 

From 484b1285bae8a966e512c8a2e089ef0a635ecfd2 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 23 Aug 2018 20:43:01 +0200
Subject: [PATCH 0725/2720] Change GymDiscreteProblemWithAutoencoder to feed
 images instead of videos to the AE

---
 tensor2tensor/data_generators/gym_problems.py | 17 +++++++++++++++++
 tensor2tensor/rl/envs/tf_atari_wrappers.py    | 12 ++++++++++--
 tensor2tensor/rl/trainer_model_based.py       |  2 +-
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 2fc7922b2..ddf3b20c1 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -401,6 +401,23 @@ def restore_networks(self, sess):
       ckpt = ckpts.model_checkpoint_path
       autoencoder_saver.restore(sess, ckpt)
 
+  def hparams(self, defaults, unused_model_hparams):
+    """Overrides VideoProblem.hparams to work on images instead of videos."""
+    p = defaults
+    p.input_modality = {
+        "inputs": ("image", 256),
+    }
+    p.target_modality = ("image", 256)
+    p.input_space_id = problem.SpaceID.IMAGE
+    p.target_space_id = problem.SpaceID.IMAGE
+
+  def preprocess(self, dataset, mode, hparams, interleave=True):
+    """Overrides VideoProblem.preprocess to work on images instead of videos."""
+    def set_targets(example):
+      example["targets"] = example["frame"]
+      return example
+    return dataset.map(set_targets)
+
 
 class GymDiscreteProblemAutoencoded(GymRealDiscreteProblem):
   """Gym discrete problem with frames already autoencoded."""
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 39006878d..eaf0cfb74 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -25,6 +25,7 @@
 from tensor2tensor.layers import discretization
 from tensor2tensor.models.research import autoencoders
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
+from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -178,9 +179,14 @@ def __init__(self, batch_env):
         trainable=False)
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       autoencoder_hparams = autoencoders.autoencoder_discrete_pong()
+      problem = registry.problem(
+          "gym_discrete_problem_with_agent_on_wrapped_full_pong_with"
+          "_autoencoder")
+      autoencoder_hparams.problem_hparams = problem.get_hparams(
+          autoencoder_hparams)
+      autoencoder_hparams.problem = problem
       self.autoencoder_model = autoencoders.AutoencoderOrderedDiscrete(
           autoencoder_hparams, tf.estimator.ModeKeys.EVAL)
-    self.autoencoder_model.set_mode(tf.estimator.ModeKeys.EVAL)
 
   @property
   def observ_shape(self):
@@ -200,7 +206,8 @@ def simulate(self, action):
     reward, done = self._batch_env.simulate(action)
     with tf.control_dependencies([reward, done]):
       with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-        ret = self.autoencoder_model.encode(self._batch_env.observ)
+        observ = tf.cast(self._batch_env.observ, tf.int32)
+        ret = self.autoencoder_model.encode(observ)
         ret = tf.cast(ret, self.observ_dtype)
         assign_op = self._observ.assign(ret)
         with tf.control_dependencies([assign_op]):
@@ -209,6 +216,7 @@ def simulate(self, action):
   def _reset_non_empty(self, indices):
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       new_values = self._batch_env._reset_non_empty(indices)  # pylint: disable=protected-access
+      new_values = tf.cast(new_values, tf.int32)
       ret = self.autoencoder_model.encode(new_values)
       ret = tf.cast(ret, self.observ_dtype)
       assign_op = tf.scatter_update(self._observ, indices, ret)
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index c89f91d5b..7aabd6882 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -241,7 +241,7 @@ def encode_dataset(model, dataset, problem, ae_hparams, autoencoder_path,
   dataset = dataset.batch(batch_size)
   examples = dataset.make_one_shot_iterator().get_next()
   images = examples.pop("frame")
-  images = tf.expand_dims(images, 1)
+  images = tf.cast(images, tf.int32)
 
   encoded = model.encode(images)
   encoded_frame_height = int(

From b2d4ddbdc570563447e9ccf52b26fde4893ac92a Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 24 Aug 2018 17:18:10 +0200
Subject: [PATCH 0726/2720] Don't compute nomix_p when not training

---
 tensor2tensor/models/research/autoencoders.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index f77b0a9c4..99ff653b4 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -625,10 +625,11 @@ def decoder(self, x, encoder_layers=None):
       # Up-convolutions.
       for i in range(hparams.num_hidden_layers):
         j = hparams.num_hidden_layers - i - 1
-        nomix_p = common_layers.inverse_lin_decay(
-            int(hparams.bottleneck_warmup_steps * 0.25 * 2**j)) + 0.01
-        if common_layers.should_generate_summaries():
-          tf.summary.scalar("nomix_p_%d" % j, nomix_p)
+        if is_training:
+          nomix_p = common_layers.inverse_lin_decay(
+              int(hparams.bottleneck_warmup_steps * 0.25 * 2**j)) + 0.01
+          if common_layers.should_generate_summaries():
+            tf.summary.scalar("nomix_p_%d" % j, nomix_p)
         filters = hparams.hidden_size * 2**j
         filters = min(filters, hparams.max_hidden_size)
         with tf.variable_scope("layer_%d" % i):

From a4203915c34b6166b31929c42190cdc614594ae2 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 24 Aug 2018 17:18:51 +0200
Subject: [PATCH 0727/2720] Disable GAN loss for AE in the pong experiment

---
 tensor2tensor/models/research/autoencoders.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 99ff653b4..36de1865e 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1108,6 +1108,7 @@ def autoencoder_discrete_pong():
   hparams.batch_size = 2
   hparams.bottleneck_noise = 0.2
   hparams.max_hidden_size = 1024
+  hparams.gan_loss_factor = 0.0
   return hparams
 
 

From b9671e96e40b38d0662dbe0e32dca0ca0c5fe62e Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 23 Aug 2018 23:17:23 +0200
Subject: [PATCH 0728/2720] Add a test for the AE experiment

---
 tensor2tensor/rl/trainer_model_based_test.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based_test.py b/tensor2tensor/rl/trainer_model_based_test.py
index 48e2f40ea..67e06427a 100644
--- a/tensor2tensor/rl/trainer_model_based_test.py
+++ b/tensor2tensor/rl/trainer_model_based_test.py
@@ -17,6 +17,9 @@
 from __future__ import division
 from __future__ import print_function
 
+import os
+import shutil
+
 from tensor2tensor.rl import trainer_model_based
 
 import tensorflow as tf
@@ -26,10 +29,19 @@
 
 class ModelRLExperimentTest(tf.test.TestCase):
 
-  def test_basic(self):
+  def setUp(self):
+    super(ModelRLExperimentTest, self).setUp()
     FLAGS.output_dir = tf.test.get_temp_dir()
-    FLAGS.loop_hparams_set = "rl_modelrl_tiny"
+    shutil.rmtree(FLAGS.output_dir)
+    os.mkdir(FLAGS.output_dir)
     FLAGS.schedule = "train"  # skip evaluation for world model training
+
+  def test_basic(self):
+    FLAGS.loop_hparams_set = "rl_modelrl_tiny"
+    trainer_model_based.main(None)
+
+  def test_ae(self):
+    FLAGS.loop_hparams_set = "rl_modelrl_ae_tiny"
     trainer_model_based.main(None)
 
 

From 7dd15907445de23d619bf01b67c2cb8566bacb15 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 31 Aug 2018 18:52:56 +0200
Subject: [PATCH 0729/2720] Add a dummy problem for running the AE inside
 AutoencoderWrapper

---
 tensor2tensor/data_generators/gym_problems.py | 9 +++++++++
 tensor2tensor/rl/envs/tf_atari_wrappers.py    | 4 +---
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index ddf3b20c1..317ed7a06 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -642,3 +642,12 @@ def frame_height(self):
   def frame_width(self):
     width = self.env.observation_space.shape[1]
     return int(math.ceil(width / self.autoencoder_factor))
+
+
+@registry.register_problem
+class DummyAutoencoderProblem(GymDiscreteProblemWithAutoencoder):
+  """Dummy problem for running the autoencoder inside AutoencoderWrapper."""
+
+  @property
+  def env_name(self):
+    return "DummyAutoencoder"
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index eaf0cfb74..55c823ebb 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -179,9 +179,7 @@ def __init__(self, batch_env):
         trainable=False)
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       autoencoder_hparams = autoencoders.autoencoder_discrete_pong()
-      problem = registry.problem(
-          "gym_discrete_problem_with_agent_on_wrapped_full_pong_with"
-          "_autoencoder")
+      problem = registry.problem("dummy_autoencoder_problem")
       autoencoder_hparams.problem_hparams = problem.get_hparams(
           autoencoder_hparams)
       autoencoder_hparams.problem = problem

From 1b98c4bb644534c8068968c49c2ef7dff3541858 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 4 Sep 2018 13:18:20 +0200
Subject: [PATCH 0730/2720] Add property history_frames to WrapperBase to
 enable wrappers to modify history

---
 tensor2tensor/rl/envs/simulated_batch_env.py |  4 ++
 tensor2tensor/rl/envs/tf_atari_wrappers.py   | 55 ++++++++++++++++++--
 2 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 395f8f34a..e8076e5bb 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -215,3 +215,7 @@ def _reset_non_empty(self, indices):
   def observ(self):
     """Access the variable holding the current observation."""
     return self._observ.read_value()
+
+  @property
+  def history_frames(self):
+    return self.history_buffer.get_all_elements()
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 39006878d..94ff50b5c 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -59,6 +59,30 @@ def _reset_non_empty(self, indices):
     with tf.control_dependencies([assign_op]):
       return tf.identity(new_values)
 
+  def _transform_history_frames(self, frames):
+    """Applies a wrapper-specific transformation to the history frames.
+
+    Overridden in wrappers that alter observations.
+
+    Args:
+      frames: A tensor of history frames to transform.
+
+    Returns a tensor of transformed frames.
+    """
+    return frames
+
+  @property
+  def history_frames(self):
+    """Returns frames from the root simulated env's history_buffer.
+
+    Transforms them with a wrapper-specific function if necessary.
+
+    Raises:
+      AttributeError: if root env doesn't have a history_buffer (i.e. is not
+        simulated).
+    """
+    return self._transform_history_frames(self._batch_env.history_frames)
+
 
 class RewardClippingWrapper(WrapperBase):
   """ Reward clipping wrapper.
@@ -109,6 +133,11 @@ def not_done_step(a, _):
       with tf.control_dependencies([self._observ.assign(simulate_ret[0])]):
         return tf.identity(simulate_ret[1]), tf.identity(simulate_ret[2])
 
+  def _transform_history_frames(self, frames):
+    # Should be implemented if ever MaxAndSkipWrapper and StackWrapper are to
+    # be used together.
+    raise NotImplementedError
+
 
 class StackWrapper(WrapperBase):
   """ A wrapper which stacks previously seen frames. """
@@ -142,12 +171,10 @@ def _reset_non_empty(self, indices):
     # pylint: disable=protected-access
     new_values = self._batch_env._reset_non_empty(indices)
     # pylint: enable=protected-access
-    history_buffer = getattr(self._batch_env, "history_buffer", None)
-    if history_buffer:
+    initial_frames = getattr(self._batch_env, "history_frames", None)
+    if initial_frames is not None:
       # Using history buffer frames for initialization, if they are available.
-      # This assumes that wrappers don't alter the observations.
       with tf.control_dependencies([new_values]):
-        initial_frames = history_buffer.get_all_elements()
         # Transpose to [batch, height, width, history, channels] and merge
         # history and channels into one dimension.
         initial_frames = tf.transpose(initial_frames, [0, 2, 3, 1, 4])
@@ -166,6 +193,10 @@ def _reset_non_empty(self, indices):
     with tf.control_dependencies([assign_op]):
       return tf.gather(self.observ, indices)
 
+  def _transform_history_frames(self, frames):
+    # Should be implemented if ever two StackWrappers are to be used together.
+    raise NotImplementedError
+
 
 class AutoencoderWrapper(WrapperBase):
   """ Transforms the observations taking the bottleneck
@@ -215,6 +246,14 @@ def _reset_non_empty(self, indices):
       with tf.control_dependencies([assign_op]):
         return tf.gather(self.observ, indices)
 
+  def _transform_history_frames(self, frames):
+    batch_size, history_size = frames.get_shape().as_list()[:2]
+    new_frames = tf.reshape(frames, (-1,) + self._batch_env.observ_shape)
+    new_frames = tf.cast(new_frames, tf.int32)
+    new_frames = self.autoencoder_model.encode(new_frames)
+    new_frames = tf.cast(new_frames, self.observ_dtype)
+    return new_frames.reshape((batch_size, history_size) + self.observ_shape)
+
 
 class IntToBitWrapper(WrapperBase):
   """Unpacks the observations from integer values to bit values"""
@@ -256,3 +295,11 @@ def _reset_non_empty(self, indices):
     assign_op = tf.scatter_update(self._observ, indices, new_values_unpacked)
     with tf.control_dependencies([assign_op]):
       return tf.identity(new_values_unpacked)
+
+  def _transform_history_frames(self, frames):
+    batch_size, history_size = frames.get_shape().as_list()[:2]
+    new_frames = discretization.int_to_bit(frames, 8)
+    new_frames = tf.reshape(
+        new_frames, (batch_size, history_size) + self.observ_shape
+    )
+    return tf.cast(new_frames, self.observ_dtype)

From 089bbbf3afc721eb0efbc6d173f16e04108457cd Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 4 Sep 2018 18:43:18 +0200
Subject: [PATCH 0731/2720] history_frames -> history_observations for
 consistency

---
 tensor2tensor/rl/envs/simulated_batch_env.py |  2 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py   | 22 +++++++++++---------
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index e8076e5bb..7065e28b8 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -217,5 +217,5 @@ def observ(self):
     return self._observ.read_value()
 
   @property
-  def history_frames(self):
+  def history_observations(self):
     return self.history_buffer.get_all_elements()
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 94ff50b5c..2a591a625 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -59,8 +59,8 @@ def _reset_non_empty(self, indices):
     with tf.control_dependencies([assign_op]):
       return tf.identity(new_values)
 
-  def _transform_history_frames(self, frames):
-    """Applies a wrapper-specific transformation to the history frames.
+  def _transform_history_observations(self, frames):
+    """Applies a wrapper-specific transformation to the history observations.
 
     Overridden in wrappers that alter observations.
 
@@ -72,8 +72,8 @@ def _transform_history_frames(self, frames):
     return frames
 
   @property
-  def history_frames(self):
-    """Returns frames from the root simulated env's history_buffer.
+  def history_observations(self):
+    """Returns observations from the root simulated env's history_buffer.
 
     Transforms them with a wrapper-specific function if necessary.
 
@@ -81,7 +81,9 @@ def history_frames(self):
       AttributeError: if root env doesn't have a history_buffer (i.e. is not
         simulated).
     """
-    return self._transform_history_frames(self._batch_env.history_frames)
+    return self._transform_history_observations(
+        self._batch_env.history_observations
+    )
 
 
 class RewardClippingWrapper(WrapperBase):
@@ -133,7 +135,7 @@ def not_done_step(a, _):
       with tf.control_dependencies([self._observ.assign(simulate_ret[0])]):
         return tf.identity(simulate_ret[1]), tf.identity(simulate_ret[2])
 
-  def _transform_history_frames(self, frames):
+  def _transform_history_observations(self, frames):
     # Should be implemented if ever MaxAndSkipWrapper and StackWrapper are to
     # be used together.
     raise NotImplementedError
@@ -171,7 +173,7 @@ def _reset_non_empty(self, indices):
     # pylint: disable=protected-access
     new_values = self._batch_env._reset_non_empty(indices)
     # pylint: enable=protected-access
-    initial_frames = getattr(self._batch_env, "history_frames", None)
+    initial_frames = getattr(self._batch_env, "history_observations", None)
     if initial_frames is not None:
       # Using history buffer frames for initialization, if they are available.
       with tf.control_dependencies([new_values]):
@@ -193,7 +195,7 @@ def _reset_non_empty(self, indices):
     with tf.control_dependencies([assign_op]):
       return tf.gather(self.observ, indices)
 
-  def _transform_history_frames(self, frames):
+  def _transform_history_observations(self, frames):
     # Should be implemented if ever two StackWrappers are to be used together.
     raise NotImplementedError
 
@@ -246,7 +248,7 @@ def _reset_non_empty(self, indices):
       with tf.control_dependencies([assign_op]):
         return tf.gather(self.observ, indices)
 
-  def _transform_history_frames(self, frames):
+  def _transform_history_observations(self, frames):
     batch_size, history_size = frames.get_shape().as_list()[:2]
     new_frames = tf.reshape(frames, (-1,) + self._batch_env.observ_shape)
     new_frames = tf.cast(new_frames, tf.int32)
@@ -296,7 +298,7 @@ def _reset_non_empty(self, indices):
     with tf.control_dependencies([assign_op]):
       return tf.identity(new_values_unpacked)
 
-  def _transform_history_frames(self, frames):
+  def _transform_history_observations(self, frames):
     batch_size, history_size = frames.get_shape().as_list()[:2]
     new_frames = discretization.int_to_bit(frames, 8)
     new_frames = tf.reshape(

From 00e1f6b30abe16d54a16050b5b303dd55c828a84 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Tue, 4 Sep 2018 15:29:15 -0700
Subject: [PATCH 0732/2720] Experiments of large mixture-of-experts (6B, 26B
 params) models on wikipedia dataset.

PiperOrigin-RevId: 211536056
---
 .../research/experiments_moe.py               | 88 +++++++++++++++++++
 1 file changed, 88 insertions(+)

diff --git a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
index 43d2b19db..fb8bcbd79 100644
--- a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
@@ -150,3 +150,91 @@ def xmoe_2d_88():
   return hparams
 
 
+@registry.register_hparams
+def xmoe_wiki_base():
+  """Series of architectural experiments on wikipedia text.
+
+  For all of these architectures, we run on languagemodel_wiki_noref_v8k_l1k
+  for 3 epochs.  (training set has ~7390100 sequences each of length 1024)
+  1 epoch = 115000 steps at batch_size=64
+
+  Results:
+  model             params(M)  einsum  alltoall  mxu-util  log-ppl(1ep) (3ep)
+
+  Note: configurations and code are likely to change without notice.
+
+  Returns:
+    a hparams
+  """
+  hparams = mtf_transformer.mtf_transformer_base()
+
+  # The following hparams are constant across all these experiments.
+  hparams.label_smoothing = 0.0
+  hparams.max_length = 1024
+  hparams.batch_size = 64
+  hparams.d_model = 1024
+  hparams.d_kv = 128
+  hparams.num_heads = 8
+  hparams.shared_embedding_and_softmax_weights = False
+  hparams.learning_rate_decay_steps = 115000
+
+  # We will vary the following parameters related to the ffn/moe layers.
+  hparams.feedforward_layer = "dense_relu_dense"
+  hparams.d_ff = 8192
+  hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
+  hparams.mesh_shape = "batch:32"
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_wiki_f64k():
+  """d_ff = 64k.
+
+  Returns:
+    a hparams object.
+  """
+  hparams = xmoe_wiki_base()
+  hparams.moe_hidden_size = 8192
+  hparams.d_ff = 65536
+  hparams.mesh_shape = "model:8;batch:16"
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_wiki_x64():
+  """Two-dimensional hierarchical mixture of experts.
+
+  (8x8 experts) * (16M params/expert) * 6 layers = 6B params
+
+  Returns:
+    a hparams object.
+  """
+  hparams = xmoe_wiki_base()
+  moe.set_default_moe_hparams(hparams)
+  hparams.feedforward_layer = "hmoe"
+  hparams.moe_hidden_size = 8192
+  hparams.mesh_shape = "b0:4;b1:8"
+  hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
+  hparams.outer_batch_size = 4
+  hparams.moe_num_experts = [8, 8]
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_wiki_x256():
+  """Two-dimensional hierarchical mixture of experts.
+
+  (16x16 experts) * (16M params/expert) * 6 layers = 24B params
+
+  Returns:
+    a hparams object.
+  """
+  hparams = xmoe_wiki_x64()
+  hparams.mesh_shape = "b0:8;b1:16"
+  hparams.outer_batch_size = 8
+  hparams.moe_num_experts = [16, 16]
+  hparams.batch_size = 256
+  hparams.learning_rate_decay_steps = 28750
+  return hparams
+
+

From 583b27f58334b840e46627631d2a9e91208ade86 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 4 Sep 2018 16:51:56 -0700
Subject: [PATCH 0733/2720] Alters the way sparse_message_pass computes bias
 and adds the option of average aggregation.

PiperOrigin-RevId: 211550217
---
 .../common_message_passing_attention.py       | 36 +++++++++++++++++--
 1 file changed, 33 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/common_message_passing_attention.py b/tensor2tensor/layers/common_message_passing_attention.py
index 2c5c61b3e..f9464bb82 100644
--- a/tensor2tensor/layers/common_message_passing_attention.py
+++ b/tensor2tensor/layers/common_message_passing_attention.py
@@ -367,6 +367,8 @@ def sparse_message_pass_batched(node_states,
                                 adjacency_matrices,
                                 num_edge_types,
                                 hidden_size,
+                                use_bias=True,
+                                average_aggregation=False,
                                 name="sparse_ggnn_batched"):
   """Identical to sparse_ggnn except that each input has a batch dimension.
 
@@ -381,6 +383,10 @@ def sparse_message_pass_batched(node_states,
       type and batch. Shape: [B, N, N, T] (sparse).
     num_edge_types: The number of edge types. T.
     hidden_size: The size of the hidden layer. H.
+    use_bias: Whether to use bias in the hidden layer.
+    average_aggregation: How to aggregate the incoming node messages. If
+      average_aggregation is true, the messages are averaged. If it is false,
+      they are summed.
     name: (optional) The scope within which tf variables should be created.
 
   Returns:
@@ -410,8 +416,14 @@ def sparse_message_pass_batched(node_states,
                                        dense_shape=new_shape)
 
   # Run a message-passing step and return the result with the batch dimension.
-  node_states = sparse_message_pass(node_states, adjacency_matrices,
-                                    num_edge_types, hidden_size, name)
+  node_states = sparse_message_pass(
+      node_states,
+      adjacency_matrices,
+      num_edge_types,
+      hidden_size,
+      use_bias=use_bias,
+      average_aggregation=average_aggregation,
+      name=name)
   return tf.reshape(node_states, [b, n, hidden_size])
 
 
@@ -419,6 +431,8 @@ def sparse_message_pass(node_states,
                         adjacency_matrices,
                         num_edge_types,
                         hidden_size,
+                        use_bias=True,
+                        average_aggregation=False,
                         name="sparse_ggnn"):
   """One message-passing step for a GNN with a sparse adjacency matrix.
 
@@ -435,6 +449,10 @@ def sparse_message_pass(node_states,
       type. Shape is [N, N, T] (sparse tensor).
     num_edge_types: The number of edge types. T.
     hidden_size: The size of the hidden state. H.
+    use_bias: Whether to use bias in the hidden layer.
+    average_aggregation: How to aggregate the incoming node messages. If
+      average_aggregation is true, the messages are averaged. If it is false,
+      they are summed.
     name: (optional) The scope within which tf variables should be created.
 
   Returns:
@@ -443,6 +461,7 @@ def sparse_message_pass(node_states,
   """
   n = tf.shape(node_states)[0]
   t = num_edge_types
+  incoming_edges_per_type = tf.sparse_reduce_sum(adjacency_matrices, axis=1)
 
   # Convert the adjacency matrix into shape [T, N, N] - one [N, N] adjacency
   # matrix for each edge type. Since sparse tensor multiplication only supports
@@ -475,7 +494,18 @@ def sparse_message_pass(node_states,
   # adding everything at the end.
   with tf.variable_scope(name, default_name="sparse_ggnn"):
     final_node_states = common_layers.dense(
-        messages, hidden_size, use_bias=True)
+        messages, hidden_size, use_bias=False)
+
+    # Multiply the bias by for each edge type by the number of incoming nodes
+    # of that edge type.
+    if use_bias:
+      bias = tf.get_variable("bias", initializer=tf.zeros([t, hidden_size]))
+      final_node_states += tf.matmul(incoming_edges_per_type, bias)
+
+    if average_aggregation:
+      incoming_edges = tf.reduce_sum(incoming_edges_per_type, -1, keepdims=True)
+      incoming_edges = tf.tile(incoming_edges, [1, hidden_size])
+      final_node_states /= incoming_edges + 1e-7
 
   return final_node_states
 

From 6647090a125ee8d1d1d791f4080b71aa0cef5298 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Tue, 4 Sep 2018 18:44:33 -0700
Subject: [PATCH 0734/2720] always use read_value when accessing Variables.

PiperOrigin-RevId: 211564108
---
 tensor2tensor/data_generators/gym_problems.py |  2 --
 tensor2tensor/rl/collect.py                   | 10 ++++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index adedf0325..7a052c971 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -164,8 +164,6 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
         data = [memory[i][memory_index][0] for i in range(4)]
         memory_index += 1
         observation, reward, done, action = data
-        # TODO(piotrmilos): cleanup types management
-        observation = observation.astype(np.uint8)
 
         debug_image = self.collect_statistics_and_generate_debug_image(
             pieces_generated, *data)
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index bde467969..7e864e5df 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -145,11 +145,13 @@ def initialization_lambda(sess):
     force_beginning_resets = False
   force_beginning_resets = tf.convert_to_tensor(force_beginning_resets)
 
-  def group():
+  def reset_ops_group():
     return tf.group(batch_env.reset(tf.range(len(batch_env))),
                     tf.assign(cumulative_rewards, zeros_tensor))
+
   reset_op = tf.cond(
-      tf.logical_or(should_reset_var, force_beginning_resets), group, tf.no_op)
+      tf.logical_or(should_reset_var.read_value(), force_beginning_resets),
+      reset_ops_group, tf.no_op)
 
   with tf.control_dependencies([reset_op]):
     reset_once_op = tf.assign(should_reset_var, False)
@@ -215,7 +217,7 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
       with tf.control_dependencies([cumulate_rewards_op]):
         # TODO(piotrmilos): possibly we need cumulative_rewards.read_value()
         scores_sum_delta = tf.reduce_sum(
-            tf.gather(cumulative_rewards, agent_indices_to_reset))
+            tf.gather(cumulative_rewards.read_value(), agent_indices_to_reset))
         scores_num_delta = tf.count_nonzero(done, dtype=tf.int32)
       with tf.control_dependencies(save_ops + [scores_sum_delta,
                                                scores_num_delta]):
@@ -258,7 +260,7 @@ def stop_condition(i, _, resets):
                        lambda: 0.)
   printing = tf.Print(0, [mean_score, scores_sum, scores_num], "mean_score: ")
   with tf.control_dependencies([index, printing]):
-    memory = [tf.identity(mem) for mem in memory]
+    memory = [mem.read_value() for mem in memory]
     mean_score_summary = tf.cond(
         tf.greater(scores_num, 0),
         lambda: tf.summary.scalar("mean_score_this_iter", mean_score),

From 714399b29521eb5978e80c59c8f969f87b75e7dd Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 4 Sep 2018 20:02:58 -0700
Subject: [PATCH 0735/2720] Bump to v1.9

PiperOrigin-RevId: 211570044
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index aa007f92b..8dcaf1e28 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.8.0',
+    version='1.9.0',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',

From 10de664389d652ce4340fa6a04f4f7dd7ece1d5b Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 4 Sep 2018 21:31:53 -0700
Subject: [PATCH 0736/2720] Allow having the beginning always in RL with random
 starts.

PiperOrigin-RevId: 211576360
---
 tensor2tensor/data_generators/gym_problems.py |  1 +
 tensor2tensor/models/research/rl.py           |  1 +
 tensor2tensor/rl/envs/simulated_batch_env.py  | 26 +++++++++++++------
 tensor2tensor/rl/trainer_model_based.py       | 15 ++++++++++-
 4 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 7a052c971..49089509c 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -528,6 +528,7 @@ def get_environment_spec(self):
     env_spec = standard_atari_env_spec(self.env_name)
     env_spec.simulated_env = True
     env_spec.add_hparam("simulation_random_starts", False)
+    env_spec.add_hparam("simulation_flip_first_random_for_beginning", False)
     env_spec.add_hparam("intrinsic_reward_scale", 0.0)
     initial_frames_problem = registry.problem(self.initial_frames_problem)
     env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index d22b600b0..8c0ad9466 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -51,6 +51,7 @@ def ppo_base_v1():
   hparams.add_hparam("optimization_batch_size", 50)
   hparams.add_hparam("max_gradients_norm", 0.5)
   hparams.add_hparam("simulation_random_starts", False)
+  hparams.add_hparam("simulation_flip_first_random_for_beginning", False)
   hparams.add_hparam("intrinsic_reward_scale", 0.)
   return hparams
 
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 395f8f34a..e8dc9f333 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -39,9 +39,13 @@
 class HistoryBuffer(object):
   """History Buffer."""
 
-  def __init__(self, input_dataset, length, observ_dtype):
-    self.input_data_iterator = (
-        input_dataset.batch(length).make_initializable_iterator())
+  def __init__(self, input_dataset, length, observ_dtype, start_frame=None):
+    if start_frame is None:
+      dataset = input_dataset.batch(length)
+    else:
+      dataset = input_dataset.batch(length - 1)
+      dataset = dataset.map(lambda x: tf.concat([start_frame, x], axis=0))
+    self.input_data_iterator = dataset.make_initializable_iterator()
     self.length = length
     self._observ_dtype = observ_dtype
     initial_frames = self.get_initial_observations()
@@ -127,20 +131,26 @@ def __init__(self, environment_spec, length):
                       environment_spec.video_num_target_frames,
                       environment_spec=environment_spec)
 
+    initial_frames_dataset = initial_frames_problem.dataset(
+        tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir, shuffle_files=False,
+        hparams=hparams).take(1)
+    start_frame = None
     if environment_spec.simulation_random_starts:
       dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
                                                FLAGS.data_dir,
                                                shuffle_files=True,
                                                hparams=hparams)
       dataset = dataset.shuffle(buffer_size=1000)
+      if environment_spec.simulation_flip_first_random_for_beginning:
+        # Later flip the first random frame in PPO batch for the true beginning.
+        start = initial_frames_dataset.make_one_shot_iterator().get_next()
+        start_frame = tf.expand_dims(start["inputs"], axis=0)
     else:
-      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
-                                               FLAGS.data_dir,
-                                               shuffle_files=False,
-                                               hparams=hparams).take(1)
+      dataset = initial_frames_dataset
 
     dataset = dataset.map(lambda x: x["inputs"]).repeat()
-    self.history_buffer = HistoryBuffer(dataset, self.length, self.observ_dtype)
+    self.history_buffer = HistoryBuffer(
+        dataset, self.length, self.observ_dtype, start_frame=start_frame)
 
     self._observ = tf.Variable(
         tf.zeros((len(self),) + observ_shape, self.observ_dtype),
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index c89f91d5b..4c853653c 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -170,6 +170,8 @@ def train_agent(problem_name, agent_model_dir,
 
   environment_spec = copy.copy(gym_problem.environment_spec)
   environment_spec.simulation_random_starts = hparams.simulation_random_starts
+  do_flip = hparams.simulation_flip_first_random_for_beginning
+  environment_spec.simulation_flip_first_random_for_beginning = do_flip
   environment_spec.intrinsic_reward_scale = hparams.intrinsic_reward_scale
 
   ppo_hparams.add_hparam("environment_spec", environment_spec)
@@ -522,7 +524,9 @@ def rl_modelrl_base():
       autoencoder_train_steps=0,
       model_train_steps=50000,
       simulated_env_generator_num_steps=2000,
-      simulation_random_starts=True,
+      simulation_random_starts=True,  # Use random starts in PPO.
+      # Flip the first random frame in PPO batch for the true beginning.
+      simulation_flip_first_random_for_beginning=True,
       intrinsic_reward_scale=0.,
       ppo_epochs_num=2000,  # This should be enough to see something
       # Our simulated envs do not know how to reset.
@@ -544,6 +548,15 @@ def rl_modelrl_base():
   )
 
 
+@registry.register_hparams
+def rl_modelrl_base_quick():
+  """Base setting with only 2 epochs and 500 PPO steps per epoch."""
+  hparams = rl_modelrl_base()
+  hparams.epochs = 2
+  hparams.ppo_epochs_num = 500
+  return hparams
+
+
 @registry.register_hparams
 def rl_modelrl_base_stochastic():
   """Base setting with a stochastic next-frame model."""

From 54510673d1ded46575c68091d2aaf791f9007a34 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Tue, 4 Sep 2018 22:54:06 -0700
Subject: [PATCH 0737/2720] ppo_optimization_batch hparam sweep.

PiperOrigin-RevId: 211582184
---
 tensor2tensor/rl/trainer_model_based.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 4c853653c..cda71681e 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -937,6 +937,13 @@ def rl_modelrl_num_frames(rhp):
                    [1000*el for el in [30, 100, 500, 1000]])
 
 
+@registry.register_ranged_hparams
+def rl_modelrl_ppo_optimization_batch_size(rhp):
+  rhp.set_categorical("loop.game", ["pong", "wrapped_full_pong", "seaquest"])
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_discrete("ppo.optimization_batch_size", [4, 10, 20])
+
+
 def merge_unscoped_hparams(scopes_and_hparams):
   """Merge multiple HParams into one with scopes."""
   merged_values = {}

From 77545ce9dbb6decdbe32aba3562651607bb9cfab Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Wed, 5 Sep 2018 11:09:08 -0700
Subject: [PATCH 0738/2720] clipping logits for policy entropy control.

PiperOrigin-RevId: 211663759
---
 tensor2tensor/models/research/rl.py     | 13 ++++++++++++-
 tensor2tensor/rl/trainer_model_based.py |  7 +++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 8c0ad9466..a7d7cb270 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -53,6 +53,7 @@ def ppo_base_v1():
   hparams.add_hparam("simulation_random_starts", False)
   hparams.add_hparam("simulation_flip_first_random_for_beginning", False)
   hparams.add_hparam("intrinsic_reward_scale", 0.)
+  hparams.add_hparam("logits_clip", 0.)
   return hparams
 
 
@@ -192,6 +193,15 @@ def feed_forward_gaussian_fun(action_space, config, observations):
   return NetworkOutput(policy, value, lambda a: tf.clip_by_value(a, -2., 2))
 
 
+def clip_logits(logits, config):
+  logits_clip = getattr(config, "logits_clip", 0.)
+  if logits_clip > 0:
+    min_logit = tf.reduce_min(logits)
+    return tf.minimum(logits - min_logit, logits_clip)
+  else:
+    return logits
+
+
 def feed_forward_categorical_fun(action_space, config, observations):
   """Feed-forward categorical."""
   if not isinstance(action_space, gym.spaces.Discrete):
@@ -211,13 +221,13 @@ def feed_forward_categorical_fun(action_space, config, observations):
       for size in config.value_layers:
         x = tf.contrib.layers.fully_connected(x, size, tf.nn.relu)
       value = tf.contrib.layers.fully_connected(x, 1, None)[..., 0]
+  logits = clip_logits(logits, config)
   policy = tf.contrib.distributions.Categorical(logits=logits)
   return NetworkOutput(policy, value, lambda a: a)
 
 
 def feed_forward_cnn_small_categorical_fun(action_space, config, observations):
   """Small cnn network with categorical output."""
-  del config
   obs_shape = common_layers.shape_list(observations)
   x = tf.reshape(observations, [-1] + obs_shape[2:])
 
@@ -237,6 +247,7 @@ def feed_forward_cnn_small_categorical_fun(action_space, config, observations):
 
       logits = tf.contrib.layers.fully_connected(x, action_space.n,
                                                  activation_fn=None)
+      logits = clip_logits(logits, config)
 
       value = tf.contrib.layers.fully_connected(
           x, 1, activation_fn=None)[..., 0]
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index cda71681e..b2c5cf939 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -944,6 +944,13 @@ def rl_modelrl_ppo_optimization_batch_size(rhp):
   rhp.set_discrete("ppo.optimization_batch_size", [4, 10, 20])
 
 
+@registry.register_ranged_hparams
+def rl_modelrl_logits_clip(rhp):
+  rhp.set_categorical("loop.game", ["pong", "wrapped_full_pong", "seaquest"])
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_discrete("ppo.logits_clip", [0., 5.])
+
+
 def merge_unscoped_hparams(scopes_and_hparams):
   """Merge multiple HParams into one with scopes."""
   merged_values = {}

From 8c02799bbc6270ea97ca4c64444bb93ea97f80a1 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 5 Sep 2018 12:03:12 -0700
Subject: [PATCH 0739/2720] Add matplotlib to [tests] and ignore some tests in
 travis

PiperOrigin-RevId: 211673360
---
 .travis.yml | 6 +++---
 setup.py    | 3 ++-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index bc14f26e0..941e6b25e 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -40,7 +40,7 @@ install:
     fi
   # First ensure that the base dependencies are sufficient for a full import
   - pip install -q .
-  - t2t-trainer --registry_help
+  - t2t-trainer --registry_help 2>&1 >/dev/null
   # Then install the test dependencies
   - pip install -q .[tests,allen]
   # Make sure to install the atari extras for gym
@@ -60,7 +60,6 @@ script:
   #   * visualization_test
   #   * trainer_model_based_test
   #   * allen_brain_test
-  #   * trainer_model_based_stochastic_test
   #   * models/research
   # algorithmic_math_test: flaky
   # universal_transformer_test: requires new feature in tf.foldl (rm with TF 1.9)
@@ -74,6 +73,7 @@ script:
     --ignore=tensor2tensor/rl/trainer_model_based_test.py
     --ignore=tensor2tensor/data_generators/allen_brain_test.py
     --ignore=tensor2tensor/rl/trainer_model_based_stochastic_test.py
+    --ignore=tensor2tensor/rl/trainer_model_based_sv2p_test.py
     --ignore=tensor2tensor/models/research
   - pytest tensor2tensor/utils/registry_test.py
   - pytest tensor2tensor/utils/trainer_lib_test.py
@@ -81,7 +81,7 @@ script:
   - pytest tensor2tensor/data_generators/allen_brain_test.py
   - if [[ "$TF_VERSION" == "$TF_LATEST"  ]] || [[ "$TF_VERSION" == "tf-nightly"  ]];
     then
-      pytest tensor2tensor/models/research;
+      pytest tensor2tensor/models/research --ignore=tensor2tensor/models/research/glow_test.py;
     fi
 
   # Run installed scripts
diff --git a/setup.py b/setup.py
index 8dcaf1e28..a61df9889 100644
--- a/setup.py
+++ b/setup.py
@@ -52,7 +52,8 @@
         'tensorflow': ['tensorflow>=1.9.0'],
         'tensorflow_gpu': ['tensorflow-gpu>=1.9.0'],
         'tests': [
-            'absl-py', 'pytest', 'mock', 'pylint', 'jupyter', 'gsutil'
+            'absl-py', 'pytest', 'mock', 'pylint', 'jupyter', 'gsutil',
+            'matplotlib',
             # Need atari extras for Travis tests, but because gym is already in
             # install_requires, pip skips the atari extras, so we instead do an
             # explicit pip install gym[atari] for the tests.

From 8d33f87848ecec407c84ba1e9caecfab696bb7f8 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Wed, 5 Sep 2018 13:31:29 -0700
Subject: [PATCH 0740/2720] Internal change

PiperOrigin-RevId: 211688554
---
 tensor2tensor/models/research/rl.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index a7d7cb270..ab373e06d 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -119,7 +119,7 @@ def ppo_pong_base():
   hparams.num_eval_agents = 1
   hparams.policy_network = feed_forward_cnn_small_categorical_fun
   hparams.clipping_coef = 0.2
-  hparams.optimization_batch_size = 4
+  hparams.optimization_batch_size = 20
   hparams.max_gradients_norm = 0.5
   return hparams
 

From 4e0db2131caebea54231a1b2d27202b7ba6a64f2 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 5 Sep 2018 13:45:39 -0700
Subject: [PATCH 0741/2720] internal merge of PR #1040

PiperOrigin-RevId: 211690975
---
 tensor2tensor/rl/envs/simulated_batch_env.py |  4 ++
 tensor2tensor/rl/envs/tf_atari_wrappers.py   | 58 ++++++++++++++++++--
 2 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index e8dc9f333..57d149bd5 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -225,3 +225,7 @@ def _reset_non_empty(self, indices):
   def observ(self):
     """Access the variable holding the current observation."""
     return self._observ.read_value()
+
+  @property
+  def history_observations(self):
+    return self.history_buffer.get_all_elements()
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 39006878d..fe6e7f311 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -59,6 +59,33 @@ def _reset_non_empty(self, indices):
     with tf.control_dependencies([assign_op]):
       return tf.identity(new_values)
 
+  def _transform_history_observations(self, frames):
+    """Applies a wrapper-specific transformation to the history observations.
+
+    Overridden in wrappers that alter observations.
+
+    Args:
+      frames: A tensor of history frames to transform.
+
+    Returns:
+      a tensor of transformed frames.
+    """
+    return frames
+
+  @property
+  def history_observations(self):
+    """Returns observations from the root simulated env's history_buffer.
+
+    Transforms them with a wrapper-specific function if necessary.
+
+    Raises:
+      AttributeError: if root env doesn't have a history_buffer (i.e. is not
+        simulated).
+    """
+    return self._transform_history_observations(
+        self._batch_env.history_observations
+    )
+
 
 class RewardClippingWrapper(WrapperBase):
   """ Reward clipping wrapper.
@@ -109,6 +136,11 @@ def not_done_step(a, _):
       with tf.control_dependencies([self._observ.assign(simulate_ret[0])]):
         return tf.identity(simulate_ret[1]), tf.identity(simulate_ret[2])
 
+  def _transform_history_observations(self, frames):
+    # Should be implemented if ever MaxAndSkipWrapper and StackWrapper are to
+    # be used together.
+    raise NotImplementedError
+
 
 class StackWrapper(WrapperBase):
   """ A wrapper which stacks previously seen frames. """
@@ -142,12 +174,10 @@ def _reset_non_empty(self, indices):
     # pylint: disable=protected-access
     new_values = self._batch_env._reset_non_empty(indices)
     # pylint: enable=protected-access
-    history_buffer = getattr(self._batch_env, "history_buffer", None)
-    if history_buffer:
+    initial_frames = getattr(self._batch_env, "history_observations", None)
+    if initial_frames is not None:
       # Using history buffer frames for initialization, if they are available.
-      # This assumes that wrappers don't alter the observations.
       with tf.control_dependencies([new_values]):
-        initial_frames = history_buffer.get_all_elements()
         # Transpose to [batch, height, width, history, channels] and merge
         # history and channels into one dimension.
         initial_frames = tf.transpose(initial_frames, [0, 2, 3, 1, 4])
@@ -166,6 +196,10 @@ def _reset_non_empty(self, indices):
     with tf.control_dependencies([assign_op]):
       return tf.gather(self.observ, indices)
 
+  def _transform_history_observations(self, frames):
+    # Should be implemented if ever two StackWrappers are to be used together.
+    raise NotImplementedError
+
 
 class AutoencoderWrapper(WrapperBase):
   """ Transforms the observations taking the bottleneck
@@ -215,6 +249,14 @@ def _reset_non_empty(self, indices):
       with tf.control_dependencies([assign_op]):
         return tf.gather(self.observ, indices)
 
+  def _transform_history_observations(self, frames):
+    batch_size, history_size = frames.get_shape().as_list()[:2]
+    new_frames = tf.reshape(frames, (-1,) + self._batch_env.observ_shape)
+    new_frames = tf.cast(new_frames, tf.int32)
+    new_frames = self.autoencoder_model.encode(new_frames)
+    new_frames = tf.cast(new_frames, self.observ_dtype)
+    return new_frames.reshape((batch_size, history_size) + self.observ_shape)
+
 
 class IntToBitWrapper(WrapperBase):
   """Unpacks the observations from integer values to bit values"""
@@ -256,3 +298,11 @@ def _reset_non_empty(self, indices):
     assign_op = tf.scatter_update(self._observ, indices, new_values_unpacked)
     with tf.control_dependencies([assign_op]):
       return tf.identity(new_values_unpacked)
+
+  def _transform_history_observations(self, frames):
+    batch_size, history_size = frames.get_shape().as_list()[:2]
+    new_frames = discretization.int_to_bit(frames, 8)
+    new_frames = tf.reshape(
+        new_frames, (batch_size, history_size) + self.observ_shape
+    )
+    return tf.cast(new_frames, self.observ_dtype)

From a8931568fd7609a7a37e6a05040e99be36916976 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 5 Sep 2018 15:01:40 -0700
Subject: [PATCH 0742/2720] Add a simple discrete autoencoder video model.

PiperOrigin-RevId: 211705537
---
 .../models/video/basic_stochastic.py          | 61 ++++++++++++++++++-
 tensor2tensor/rl/trainer_model_based.py       |  9 +++
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 26ba47a4f..450713acb 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 
@@ -37,7 +38,7 @@ class NextFrameBasicStochastic(
   """Stochastic version of basic next-frame model."""
 
   def inject_latent(self, layer, features, filters):
-    """Do nothing for deterministic model."""
+    """Inject a VAE-style latent."""
     # Latent for stochastic model
     input_frames = tf.to_float(features["inputs_raw"])
     target_frames = tf.to_float(features["targets_raw"])
@@ -56,6 +57,53 @@ def inject_latent(self, layer, features, filters):
     return layer, extra_loss
 
 
+@registry.register_model
+class NextFrameBasicStochasticDiscrete(
+    basic_deterministic.NextFrameBasicDeterministic):
+  """Basic next-frame model with a tiny discrete latent."""
+
+  def inject_latent(self, layer, features, filters):
+    """Inject a deterministic latent based on the target frame."""
+    del filters
+    hparams = self.hparams
+    final_filters = common_layers.shape_list(layer)[-1]
+    filters = hparams.hidden_size
+    kernel = (4, 4)
+
+    if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+      layer_shape = common_layers.shape_list(layer)
+      rand = tf.random_uniform(layer_shape[:-1] + [hparams.bottleneck_bits])
+      d = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
+      z = tf.layers.dense(d, final_filters, name="unbottleneck")
+      return layer + z, 0.0
+
+    # Embed.
+    x = tf.layers.dense(
+        features["targets"], filters, name="latent_embed",
+        bias_initializer=tf.random_normal_initializer(stddev=0.01))
+    x = common_attention.add_timing_signal_nd(x)
+
+    for i in range(hparams.num_compress_steps):
+      with tf.variable_scope("latent_downstride%d" % i):
+        x = common_layers.make_even_size(x)
+        if i < hparams.filter_double_steps:
+          filters *= 2
+        x = common_attention.add_timing_signal_nd(x)
+        x = tf.layers.conv2d(x, filters, kernel, activation=common_layers.belu,
+                             strides=(2, 2), padding="SAME")
+        x = common_layers.layer_norm(x)
+
+    x = tf.tanh(tf.layers.dense(x, hparams.bottleneck_bits, name="bottleneck"))
+    d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
+    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+      noise = tf.random_uniform(common_layers.shape_list(x))
+      noise = 2.0 * tf.to_float(tf.less(hparams.bottleneck_noise, noise)) - 1.0
+      d *= noise
+
+    z = tf.layers.dense(d, final_filters, name="unbottleneck")
+    return layer + z, 0.0
+
+
 @registry.register_hparams
 def next_frame_basic_stochastic():
   """Basic 2-frame conv model with stochastic tower."""
@@ -71,3 +119,14 @@ def next_frame_basic_stochastic():
   hparams.add_hparam("anneal_end", 100000)
   hparams.add_hparam("information_capacity", 0.0)
   return hparams
+
+
+@registry.register_hparams
+def next_frame_basic_stochastic_discrete():
+  """Basic 2-frame conv model with stochastic discrete latent."""
+  hparams = basic_deterministic_params.next_frame_basic_deterministic()
+  hparams.num_compress_steps = 8
+  hparams.filter_double_steps = 3
+  hparams.add_hparam("bottleneck_bits", 32)
+  hparams.add_hparam("bottleneck_noise", 0.05)
+  return hparams
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index b2c5cf939..034643da2 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -557,6 +557,15 @@ def rl_modelrl_base_quick():
   return hparams
 
 
+@registry.register_hparams
+def rl_modelrl_base_quick_sd():
+  """Quick setting with stochastic discrete model."""
+  hparams = rl_modelrl_base_quick()
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  return hparams
+
+
 @registry.register_hparams
 def rl_modelrl_base_stochastic():
   """Base setting with a stochastic next-frame model."""

From f7695e8d5c9d403fed763c0cc351c6835516e6f3 Mon Sep 17 00:00:00 2001
From: Mostafa Dehghani <dehghani.mostafa@gmail.com>
Date: Wed, 5 Sep 2018 15:52:16 -0700
Subject: [PATCH 0743/2720] Cleaning up the codes for gru/lstm as transition
 function for universal transformer (#1036)

* fixing the code for gru/lstm as transition function in universal transformer

* fixing ut_lstm and ut_gru hparams_sets

* cleaning up hparams_set for different types of universal transformer

* change the default of transformer_ffn_type to fc instead of sepconv as it gets the mt results reported in the paper.
---
 .../models/research/universal_transformer.py  | 284 +++----------
 .../research/universal_transformer_util.py    | 391 ++++++++----------
 2 files changed, 239 insertions(+), 436 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index a38e4c187..212294c44 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -86,6 +86,7 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
 
     return encoder_output, encoder_decoder_attention_bias, encoder_extra_output
 
+
   def decode(self,
              decoder_input,
              encoder_output,
@@ -93,6 +94,7 @@ def decode(self,
              decoder_self_attention_bias,
              hparams,
              cache=None,
+             decode_loop_step=None,
              nonpadding=None,
              losses=None):
     """Decode Universal Transformer outputs from encoder representation.
@@ -112,6 +114,7 @@ def decode(self,
         self-attention. [batch_size, decoder_length]
       hparams: hyperparmeters for model.
       cache: Unimplemented.
+      decode_loop_step: Unused.
       nonpadding: optional Tensor with shape [batch_size, decoder_length]
       losses: Unused.
 
@@ -123,6 +126,7 @@ def decode(self,
             variants of the model (e.g. in ACT, to pass the ponder-time to body)
 
     """
+    del decode_loop_step
     del losses
     # TODO(dehghani): enable caching.
     del cache
@@ -392,7 +396,7 @@ def update_hparams_for_universal_transformer(hparams):
 
   # Default ffn layer is separable convolution.
   # Options: "fc" and "sepconv".
-  hparams.add_hparam("transformer_ffn_type", "sepconv")
+  hparams.add_hparam("transformer_ffn_type", "fc")
 
   # Transform bias (in models with highway or skip connection).
   hparams.add_hparam("transform_bias_init", -1.0)
@@ -409,30 +413,13 @@ def update_hparams_for_universal_transformer(hparams):
   # With dense_relu_dense, the bias/kernel initializations will not be applied.
   hparams.add_hparam("gate_ffn_layer", "dense")
 
-  # Config for all rnn style recurrencies (rnn, lstm, gru):
-  # Input of the gate functions: i:input/s:state/t:transformed state.
-  # or any combination: e.g. is, ts, ist, etc.
-  hparams.add_hparam("gates_inputs", "i")
-
-  # LSTEM forget bias.
+  # LSTM forget bias for lstm style recurrence.
   hparams.add_hparam("lstm_forget_bias", 1.0)
+  # Uses the memory at the last step as the final ouput, if true.
+  hparams.add_hparam("use_memory_as_final_state", True)
+  # if also add a ffn unit to the transition function when using gru/lstm
+  hparams.add_hparam("add_ffn_unit_to_the_transition_function", False)
 
-  # How to combine state and input in each step:
-  # "mh_attention_ffn_add" or "add_mh_attention_ffn" or "dense_mh_attention"
-  # or "mh_attention_dense".
-  # Interpretation for e.g. "mh_attention_ffn_add":
-  # Apply transformer attention then transformer ffn, then add.
-  hparams.add_hparam("inputs_states_combination", "mh_attention_ffn_add")
-
-  # Config for gru_style recurrency:
-  # What to transform in gru: state/output/candidate/combination of them.
-  hparams.add_hparam("gru_transformation", ["state_transformation"])
-
-  # Config for lstm_style Recurrency:
-  # What to transform in lstm: state/modulated_input/memory.
-  hparams.add_hparam("lstm_transformation", ["state_transformation"])
-  # Uses the mememory at the last step as the final touput, if true.
-  hparams.add_hparam("use_memory_as_final_state", False)
 
   # Type of act: basic/accumulated/global (instead of position-wise!)/random.
   hparams.add_hparam("act_type", "basic")
@@ -494,24 +481,17 @@ def universal_transformer_teeny():
 
 
 @registry.register_hparams
-def universal_transformer_small_dropconnect():
+def universal_transformer_tall():
   hparams = universal_transformer_small()
-  hparams.gate_ffn_layer = "dense_dropconnect"
-  hparams.add_hparam("dropconnect_dropout", 0.5)
+  hparams.num_rec_steps = 16
   return hparams
 
 
 @registry.register_hparams
-def adaptive_universal_transformer_small():
+def universal_transformer_small_dropconnect():
   hparams = universal_transformer_small()
-  hparams.recurrence_type = "act"
-  return hparams
-
-
-@registry.register_hparams
-def adaptive_universal_transformer_tiny():
-  hparams = universal_transformer_tiny()
-  hparams.recurrence_type = "act"
+  hparams.gate_ffn_layer = "dense_dropconnect"
+  hparams.add_hparam("dropconnect_dropout", 0.5)
   return hparams
 
 
@@ -521,73 +501,28 @@ def adaptive_universal_transformer_base():
   hparams.recurrence_type = "act"
   return hparams
 
-
-@registry.register_hparams
-def adaptive_universal_transformer_random_small():
-  hparams = universal_transformer_small()
-  hparams.recurrence_type = "act"
-  hparams.act_type = "random"
-  return hparams
-
-
-@registry.register_hparams
-def adaptive_universal_transformer_accumulated_small():
-  hparams = universal_transformer_small()
-  hparams.recurrence_type = "act"
-  hparams.act_type = "accumulated"
-  return hparams
-
-
 @registry.register_hparams
-def adaptive_universal_transformer_global_small():
+def adaptive_universal_transformer_small():
   hparams = universal_transformer_small()
   hparams.recurrence_type = "act"
-  hparams.act_type = "global"
   return hparams
 
 
 @registry.register_hparams
-def adaptive_universal_transformer_accumulated_tiny():
+def adaptive_universal_transformer_tiny():
   hparams = universal_transformer_tiny()
   hparams.recurrence_type = "act"
-  hparams.act_type = "accumulated"
   return hparams
 
 
 @registry.register_hparams
-def adaptive_universal_transformer_global_tiny():
-  hparams = universal_transformer_tiny()
+def adaptive_universal_transformer_global_base():
+  hparams = universal_transformer_base()
   hparams.recurrence_type = "act"
   hparams.act_type = "global"
   return hparams
 
 
-@registry.register_hparams
-def adaptive_universal_transformer_random_tiny():
-  hparams = universal_transformer_tiny()
-  hparams.recurrence_type = "act"
-  hparams.act_type = "random"
-  return hparams
-
-
-@registry.register_hparams
-def adaptive_universal_transformer_small_sb():
-  hparams = universal_transformer_small()
-  hparams.recurrence_type = "act"
-  hparams.batch_size = 2048
-  return hparams
-
-
-@registry.register_hparams
-def adaptive_universal_transformer_large():
-  hparams = universal_transformer_small()
-  hparams.recurrence_type = "act"
-  hparams.hidden_size = 1024
-  hparams.batch_size = 2048
-  hparams.filter_size = 2048
-  return hparams
-
-
 @registry.register_hparams
 def adaptive_universal_transformer_tall():
   hparams = universal_transformer_small()
@@ -605,6 +540,7 @@ def adaptive_universal_transformer_tall_actlossw0():
   hparams.num_hidden_layers = 16
   hparams.batch_size = 1024
   hparams.act_max_steps = 24
+  hparams.act_loss_weight = 0.0
   return hparams
 
 
@@ -615,21 +551,13 @@ def adaptive_universal_transformer_tall_actlossw001():
   hparams.num_hidden_layers = 16
   hparams.batch_size = 1024
   hparams.act_max_steps = 24
+  hparams.act_loss_weight = 0.001
   return hparams
 
 
-@registry.register_hparams
-def adaptive_universal_transformer_small_d03():
-  hparams = universal_transformer_small()
-  hparams.recurrence_type = "act"
-  hparams.layer_prepostprocess_dropout = 0.3
-  hparams.attention_dropout = 0.3
-  hparams.relu_dropout = 0.3
-  return hparams
-
 
 @registry.register_hparams
-def adaptive_universal_transformer_base_d03():
+def adaptive_universal_transformer_base_dropout03():
   hparams = universal_transformer_base()
   hparams.recurrence_type = "act"
   hparams.layer_prepostprocess_dropout = 0.3
@@ -639,37 +567,8 @@ def adaptive_universal_transformer_base_d03():
 
 
 @registry.register_hparams
-def adaptive_universal_transformer_tiny_d02():
-  hparams = universal_transformer_tiny()
-  hparams.recurrence_type = "act"
-  hparams.layer_prepostprocess_dropout = 0.2
-  hparams.attention_dropout = 0.2
-  hparams.relu_dropout = 0.2
-  return hparams
-
-
-@registry.register_hparams
-def adaptive_universal_transformer_tiny_d02_sb():
-  hparams = universal_transformer_tiny()
-  hparams.recurrence_type = "act"
-  hparams.layer_prepostprocess_dropout = 0.2
-  hparams.attention_dropout = 0.2
-  hparams.relu_dropout = 0.2
-  hparams.batch_size = 2048
-  return hparams
-
-
-@registry.register_hparams
-def adaptive_universal_transformer_tiny_sb():
-  hparams = universal_transformer_tiny()
-  hparams.recurrence_type = "act"
-  hparams.batch_size = 2048
-  return hparams
-
-
-@registry.register_hparams
-def adaptive_universal_transformer_tiny_d05():
-  hparams = universal_transformer_tiny()
+def adaptive_universal_transformer_base_dropout05():
+  hparams = universal_transformer_base()
   hparams.recurrence_type = "act"
   hparams.layer_prepostprocess_dropout = 0.5
   hparams.attention_dropout = 0.5
@@ -677,87 +576,47 @@ def adaptive_universal_transformer_tiny_d05():
   return hparams
 
 
-@registry.register_hparams
-def universal_transformer_small_sb():
-  hparams = universal_transformer_small()
-  hparams.batch_size = 2048
-  return hparams
-
-
-@registry.register_hparams
-def universal_transformer_skip_small():
-  hparams = universal_transformer_small()
-  hparams.recurrence_type = "skip"
-  return hparams
-
 
 @registry.register_hparams
-def universal_transformer_skip_tiny():
-  hparams = universal_transformer_tiny()
+def universal_transformer_skip_base():
+  hparams = universal_transformer_base()
   hparams.recurrence_type = "skip"
   return hparams
 
 
 @registry.register_hparams
-def universal_transformer_highway_small():
-  hparams = universal_transformer_small()
-  hparams.recurrence_type = "highway"
-  return hparams
-
-
-@registry.register_hparams
-def universal_transformer_highway_tiny():
-  hparams = universal_transformer_tiny()
+def universal_transformer_highway_base():
+  hparams = universal_transformer_base()
   hparams.recurrence_type = "highway"
   return hparams
 
-
-@registry.register_hparams
-def universal_transformer_dwa_small():
-  hparams = universal_transformer_small()
-  hparams.recurrence_type = "dwa"
-  return hparams
-
-
-@registry.register_hparams
-def universal_transformer_dwa_tiny():
-  hparams = universal_transformer_tiny()
-  hparams.recurrence_type = "dwa"
-  return hparams
-
-
 @registry.register_hparams
-def universal_transformer_dwa_tiny_test():
-  hparams = universal_transformer_tiny()
+def universal_transformer_dwa_base():
+  hparams = universal_transformer_base()
   hparams.recurrence_type = "dwa"
   return hparams
 
-
 @registry.register_hparams
-def universal_transformer_rnn_small():
-  hparams = universal_transformer_small()
-  hparams.recurrence_type = "rnn"
+def universal_transformer_lstm_base():
+  hparams = universal_transformer_base()
+  hparams.recurrence_type = "lstm"
+  hparams.add_step_timing_signal = False  # Let lstm count in depth for us!
   return hparams
 
 
 @registry.register_hparams
-def universal_transformer_gru_small():
-  hparams = universal_transformer_small()
+def universal_transformer_gru_base():
+  hparams = universal_transformer_base()
   hparams.recurrence_type = "gru"
+  hparams.add_step_timing_signal = False  # Let gru count in depth for us!
   return hparams
 
 
 @registry.register_hparams
-def universal_transformer_lstm_small():
-  hparams = universal_transformer_small()
+def universal_transformer_lstm_tall():
+  hparams = universal_transformer_tall()
   hparams.recurrence_type = "lstm"
-  return hparams
-
-
-@registry.register_hparams
-def universal_transformer_position_random_timing_small():
-  hparams = universal_transformer_small()
-  hparams.position_start_index = "random"
+  hparams.add_step_timing_signal = False  # Let lstm count in depth for us!
   return hparams
 
 
@@ -791,39 +650,30 @@ def adaptive_universal_transformer_position_random_timing_tiny():
 
 
 @registry.register_hparams
-def adaptive_universal_transformer_position_step_timing_tiny():
-  hparams = universal_transformer_tiny()
-  hparams.recurrence_type = "act"
-  hparams.position_start_index = "step"
-  return hparams
-
-
-@registry.register_hparams
-def adaptive_universal_transformer_step_sinusoid_timing_tiny():
-  hparams = universal_transformer_tiny()
-  hparams.recurrence_type = "act"
-  hparams.step_timing_signal_type = "sinusoid"
+def universal_transformer_mix_before_ut_base():
+  hparams = universal_transformer_base()
+  hparams.mix_with_transformer = "before_ut"
   return hparams
 
 
 @registry.register_hparams
-def universal_transformer_mix_after_ut_small():
-  hparams = universal_transformer_small()
-  hparams.mix_with_transformer = "before_ut"
+def universal_transformer_mix_after_ut_base():
+  hparams = universal_transformer_base()
+  hparams.mix_with_transformer = "after_ut"
   return hparams
 
 
 @registry.register_hparams
-def adaptive_universal_transformer_mix_before_ut_small():
-  hparams = universal_transformer_small()
+def adaptive_universal_transformer_mix_before_ut_base():
+  hparams = universal_transformer_base()
   hparams.mix_with_transformer = "before_ut"
   hparams.recurrence_type = "act"
   return hparams
 
 
 @registry.register_hparams
-def adaptive_universal_transformer_mix_after_ut_small():
-  hparams = universal_transformer_small()
+def adaptive_universal_transformer_mix_after_ut_base():
+  hparams = universal_transformer_base()
   hparams.mix_with_transformer = "after_ut"
   hparams.recurrence_type = "act"
   return hparams
@@ -838,45 +688,29 @@ def adaptive_universal_transformer_concat_tiny():
 
 
 @registry.register_hparams
-def adaptive_universal_transformer_concat_small():
-  hparams = universal_transformer_small()
-  hparams.recurrence_type = "act"
-  hparams.add_or_concat_timing_signal = "concat"
-  return hparams
-
-
-@registry.register_hparams
-def adaptive_universal_transformer_with_sru_small():
-  hparams = universal_transformer_small()
+def adaptive_universal_transformer_with_sru_base():
+  hparams = universal_transformer_base()
   hparams.recurrence_type = "act"
   hparams.add_sru = True
   return hparams
 
 
 @registry.register_hparams
-def universal_transformer_fc_small():
-  hparams = universal_transformer_small()
-  hparams.transformer_ffn_type = "fc"
+def universal_transformer_sepconv_big():
+  hparams = universal_transformer_big()
+  hparams.transformer_ffn_type = "sepconv"
   return hparams
 
 
 @registry.register_hparams
-def universal_transformer_fc_base():
+def universal_transformer_sepconv_base():
   hparams = universal_transformer_base()
-  hparams.transformer_ffn_type = "fc"
-  return hparams
-
-
-@registry.register_hparams
-def universal_transformer_fc_big():
-  hparams = universal_transformer_big()
-  hparams.transformer_ffn_type = "fc"
+  hparams.transformer_ffn_type = "sepconv"
   return hparams
 
-
 @registry.register_ranged_hparams
 def universal_transformer_base_range(rhp):
-  """Small range of hyperparameters."""
+  """Range of hyperparameters."""
   # After starting from base, set intervals for some parameters.
   rhp.set_discrete("num_rec_steps", [6, 8, 10])
   rhp.set_discrete("hidden_size", [1024, 2048, 4096])
@@ -889,7 +723,7 @@ def universal_transformer_base_range(rhp):
 
 @registry.register_ranged_hparams
 def adaptive_universal_transformer_base_range(rhp):
-  """Small range of hyperparameters."""
+  """Range of hyperparameters."""
   # After starting from base, set intervals for some parameters.
   rhp.set_discrete("act_max_steps", [8, 16, 32])
   rhp.set_float("act_loss_weight", 0.0, 0.5)
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 87c20e8cd..82f0d53f8 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -37,7 +37,7 @@
 each other in depth. For instance, the recurrent transition, can be a simple
 identity function which passes the output of a step as the input to next step.
 Or it can be an LSTM (filliped vertically) next to the transformer which
-controls how state  of the model changes in depth, Or even another transformer.
+controls how state of the model changes in depth.
 
 """
 
@@ -250,7 +250,7 @@ def add_vanilla_transformer_layer(x, num_layers):
       output, _, extra_output = tf.foldl(
           ut_function, tf.range(hparams.num_rec_steps), initializer=initializer)
 
-      # This can be the if we use universal_transformer_lstm layer.
+      # This is possible only when we are using lstm as transition function.
       if hparams.get("use_memory_as_final_state", False):
         output = extra_output
 
@@ -324,20 +324,12 @@ def get_ut_layer(x,
         ffn_unit=ffn_unit,
         attention_unit=attention_unit)
 
-  elif hparams.recurrence_type == "rnn":
-    ut_initializer = (x, x, x)  # (state, input, memory)
-    ut_function = functools.partial(
-        universal_transformer_rnn,
-        hparams=hparams,
-        ffn_unit=ffn_unit,
-        attention_unit=attention_unit,
-        pad_remover=pad_remover)
-
   elif hparams.recurrence_type == "gru":
     ut_initializer = (x, x, x)  # (state, input, memory)
     ut_function = functools.partial(
-        universal_transformer_gru,
+        universal_transformer_with_gru_as_transition_function,
         hparams=hparams,
+        ffn_unit=ffn_unit,
         attention_unit=attention_unit,
         pad_remover=pad_remover)
 
@@ -345,8 +337,9 @@ def get_ut_layer(x,
     memory = tf.zeros(common_layers.shape_list(x))
     ut_initializer = (x, x, memory)  # (state, input, memory)
     ut_function = functools.partial(
-        universal_transformer_lstm,
+        universal_transformer_with_lstm_as_transition_function,
         hparams=hparams,
+        ffn_unit=ffn_unit,
         attention_unit=attention_unit,
         pad_remover=pad_remover)
 
@@ -559,17 +552,17 @@ def universal_transformer_basic(layer_inputs,
                                 step, hparams,
                                 ffn_unit,
                                 attention_unit):
-  """Basic universal_transformer.
+  """Basic Universal Transformer.
 
-  This is in fact vanilla transformer in which weights are shared between
-  layers. For some tasks, this simple idea brings a generalization that is not
-  achievable by playing with the size of the model or drop_out parameters in
-  the vanilla transformer.
+  This model is pretty similar to the vanilla transformer in which weights are
+  shared between layers. For some tasks, this simple idea brings a
+  generalization that is not achievable by playing with the size of the model
+  or drop_out parameters in the vanilla transformer.
 
   Args:
     layer_inputs:
         - state: state
-    step: indicating number of steps take so far
+    step: indicates number of steps taken so far
     hparams: model hyper-parameters
     ffn_unit: feed-forward unit
     attention_unit: multi-head attention unit
@@ -592,12 +585,13 @@ def universal_transformer_highway(layer_inputs,
                                   ffn_unit,
                                   attention_unit,
                                   pad_remover=None):
-  """universal_transformer with highway connection.
+  """Universal Transformer with highway connection.
 
 
-  It transforms the state using attention and ffn and wrap this transformation
-  with a highway connection. (the new state is a combination of the state and
-  the transformed-state based on cary/transform gates.)
+  It transforms the state using a block contaaining sel-attention and transition
+  function  and wrap the whole block with a highway connection.
+  (the new state is a combination of the state and the transformed-state
+  based on cary/transform gates.)
 
   Interesting observation:
     Controlling the cary/transform gate with the original inputs works usually
@@ -607,7 +601,7 @@ def universal_transformer_highway(layer_inputs,
     layer_inputs:
       - state: state
       - inputs: the original embedded inputs (= inputs to the first step)
-    step: indicating number of steps take so far
+    step: indicates number of steps taken so far
     hparams: model hyper-parameters.
     ffn_unit: feed-forward unit
     attention_unit: multi-head attention unit
@@ -681,7 +675,7 @@ def universal_transformer_skip(layer_inputs,
                                ffn_unit,
                                attention_unit,
                                pad_remover=None):
-  """universal_transformer with highway connection.
+  """Universal Transformer with highway connection.
 
 
   It transforms the state using attention and ffn and wrap this transformation
@@ -696,7 +690,7 @@ def universal_transformer_skip(layer_inputs,
     layer_inputs:
       - state: state
       - inputs: the original embedded inputs (= inputs to the first step)
-    step: indicating number of steps take so far
+    step: indicates number of steps taken so far
     hparams: model hyper-parameters.
     ffn_unit: feed-forward unit
     attention_unit: multi-head attention unit
@@ -820,234 +814,209 @@ def universal_transformer_depthwise_attention(layer_inputs,
   return new_state, inputs, memory
 
 
-def universal_transformer_rnn(layer_inputs,
-                              step,
-                              hparams,
-                              ffn_unit,
-                              attention_unit,
-                              pad_remover=None):
-  """The UT layer which models recurencey similar to basic RNN cell.
+def universal_transformer_with_gru_as_transition_function(
+    layer_inputs, step, hparams, ffn_unit, attention_unit, pad_remover=None):
+  """Universal Transformer which uses a gru as transition function.
 
-    It's an U-Transformer with an RNN applied over the stats on depth.
+  It's kind of like having a gru, filliped vertically next to the Universal
+  Transformer that controls the flow of the information in depth,
+  over different steps of the Universal Transformer.
 
   Args:
     layer_inputs:
       - state: state
-      - inputs: the original embedded inputs (= inputs to the first step)
-    step: indicating number of steps take so far
+      - inputs: not used here
+      - memory: not used here
+    step: indicates number of steps taken so far
     hparams: model hyper-parameters.
     ffn_unit: feed-forward unit
     attention_unit: multi-head attention unit
     pad_remover: to mask out padding in convolutional layers (efficiency).
-
   Returns:
     layer_output:
-      new_state: new state
-      inputs: the original embedded inputs (= inputs to the first step)
-
-  Raises:
-    ValueError: Unknown inputs_states_combination type
-
+        new_state: new state
+        inputs: not uesed
+        memory: not used
   """
 
-  state, inputs, memory = layer_inputs
-  state = step_preprocess(state, step, hparams)
-
-  # TODO(dehghani) keep only the meaningful cases:
-  if hparams.inputs_states_combination == "mh_attention_ffn_add":
-    state.get_shape().assert_is_compatible_with(inputs.get_shape())
-    state = ffn_unit(attention_unit(state))
-    new_state = state + inputs
-
-  elif hparams.inputs_states_combination == "add_mh_attention_ffn":
-    state.get_shape().assert_is_compatible_with(inputs.get_shape())
-    state += inputs
-    new_state = ffn_unit(attention_unit(state))
-
-  elif hparams.inputs_states_combination == "dense_mh_attention":
-    state = _ffn_layer_multi_inputs(
-        [state, inputs],
-        hparams=hparams,
-        ffn_layer_type="dense_relu_dense",
-        name="rnn",
-        activation=tf.tanh,
-        pad_remover=pad_remover)
-
-    new_state = attention_unit(state)
-
-  elif hparams.inputs_states_combination == "mh_attention_dense":
-    state = attention_unit(state)
-    new_state = _ffn_layer_multi_inputs(
-        [state, inputs],
-        hparams=hparams,
-        ffn_layer_type="dense_relu_dense",
-        name="rnn",
-        activation=tf.tanh,
-        pad_remover=pad_remover)
-
-  else:
-    raise ValueError("Unknown inputs_states_combination type: %s" %
-                     hparams.inputs_states_combination)
-
-  return new_state, inputs, memory
-
-
-def universal_transformer_gru(layer_inputs,
-                              step,
-                              hparams,
-                              attention_unit,
-                              pad_remover=None):
-  """The RT layer which models recurencey similar to GRU cell.
-
-    It's an R-transformer with a gru applied over the stats on depth.
-    Based on GRU paper: http://arxiv.org/abs/1406.1078
-
-  Args:
-    layer_inputs:
-      - state: state
-      - inputs: the original embedded inputs (= inputs to the first step)
-    step: indicating number of steps take so far
-    hparams: model hyper-parameters.
-    attention_unit: multi-head attention unit
-    pad_remover: to mask out padding in convolutional layers (efficiency).
-
-
-  Returns:
-    layer_output:
-      new_state: new state
-      inputs: the original embedded inputs (= inputs to the first step)
-  """
+  state, unused_inputs, unused_memory = tf.unstack(
+      layer_inputs, num=None, axis=0, name="unstack")
 
-  state, inputs, memory = layer_inputs
-  state = step_preprocess(state, step, hparams)
+  # state (ut_state): output of the gru in the previous step
 
-  # TODO(dehghani): do we need preprocess here?
-  state = common_layers.layer_preprocess(state, hparams)
-  inputs = common_layers.layer_preprocess(inputs, hparams)
+  # Multi_head_attention:
+  assert not hparams.add_step_timing_signal   # Let gru count for us!
+  mh_attention_input = step_preprocess(state, step, hparams)
+  transition_function_input = attention_unit(mh_attention_input)
 
-  update_gate = _ffn_layer_multi_inputs(
-      [inputs, state],
-      hparams,
-      name="update",
-      bias_initializer=tf.constant_initializer(1.0),
-      activation=tf.sigmoid,
-      pad_remover=pad_remover)
+  # Transition Function:
+  if hparams.add_ffn_unit_to_the_transition_function:
+    transition_function_input = ffn_unit(transition_function_input)
 
-  reset_gate = _ffn_layer_multi_inputs(
-      [inputs, state],
-      hparams,
-      name="reset",
-      bias_initializer=tf.constant_initializer(1.0),
-      activation=tf.sigmoid,
-      pad_remover=pad_remover)
+  transition_function_input = common_layers.layer_preprocess(
+      transition_function_input, hparams)
+  with tf.variable_scope("gru"):
+    # gru update gate: z_t = sigmoid(W_z.x_t + U_z.h_{t-1})
+    transition_function_update_gate = _ffn_layer_multi_inputs(
+        [transition_function_input, state],
+        hparams,
+        name="update",
+        bias_initializer=tf.constant_initializer(1.0),
+        activation=tf.sigmoid,
+        pad_remover=pad_remover,
+        preprocess=False,
+        postprocess=False)
 
-  reset_state = reset_gate * state
+    tf.contrib.summary.scalar("gru_update_gate",
+                              tf.reduce_mean(transition_function_update_gate))
 
-  candidate = _ffn_layer_multi_inputs(
-      [inputs, reset_state],
-      hparams,
-      name="candidate",
-      bias_initializer=tf.zeros_initializer(),
-      activation=tf.tanh,
-      pad_remover=pad_remover)
+    # gru reset gate: r_t = sigmoid(W_r.x_t + U_r.h_{t-1})
+    transition_function_reset_gate = _ffn_layer_multi_inputs(
+        [transition_function_input, state],
+        hparams,
+        name="reset",
+        bias_initializer=tf.constant_initializer(1.0),
+        activation=tf.sigmoid,
+        pad_remover=pad_remover,
+        preprocess=False,
+        postprocess=False)
 
-  if "candidate_transformation" in hparams.gru_transformation:
-    candidate = attention_unit(candidate)
+    tf.contrib.summary.scalar("gru_reset_gate",
+                              tf.reduce_mean(transition_function_reset_gate))
+    reset_state = transition_function_reset_gate * state
 
-  if "state_transformation" in hparams.gru_transformation:
-    state = attention_unit(state)
+    # gru_candidate_activation: h' = tanh(W_{x_t} + U (r_t h_{t-1})
+    transition_function_candidate = _ffn_layer_multi_inputs(
+        [transition_function_input, reset_state],
+        hparams,
+        name="candidate",
+        bias_initializer=tf.zeros_initializer(),
+        activation=tf.tanh,
+        pad_remover=pad_remover,
+        preprocess=False,
+        postprocess=False)
 
-  state = update_gate * state + (1 - update_gate) * candidate
+    transition_function_output = (
+        (1 - transition_function_update_gate) * transition_function_input +
+        transition_function_update_gate * transition_function_candidate)
 
-  if "state_transformation" in hparams.gru_transformation:
-    state = attention_unit(state)
-  # normalization on the output
-  new_state = common_layers.layer_preprocess(state, hparams)
+  transition_function_output = common_layers.layer_preprocess(
+      transition_function_output, hparams)
 
-  return new_state, inputs, memory
+  return transition_function_output, unused_inputs, unused_memory
 
 
-def universal_transformer_lstm(layer_inputs,
-                               step,
-                               hparams,
-                               attention_unit,
-                               pad_remover=None):
-  """The UT layer which models recurencey similar to GRU cell.
+def universal_transformer_with_lstm_as_transition_function(
+    layer_inputs, step, hparams, ffn_unit, attention_unit, pad_remover=None):
+  """Universal Transformer which uses a lstm as transition function.
 
-  It's an R-transformer with a gru applied over the stats on depth.
-  based on LSTM paper: https://arxiv.org/pdf/1409.2329.pdf
+  It's kind of like having a lstm, filliped vertically next to the Universal
+  Transformer that controls the flow of the  information in depth,
+  over different steps of the Universal Transformer.
 
   Args:
     layer_inputs:
       - state: state
       - inputs: the original embedded inputs (= inputs to the first step)
       - memory: memory used in lstm.
-    step: indicating number of steps take so far
+    step: indicates number of steps taken so far
     hparams: model hyper-parameters.
+    ffn_unit: feed-forward unit
     attention_unit: multi-head attention unit
     pad_remover: to mask out padding in convolutional layers (efficiency).
-
   Returns:
     layer_output:
         new_state: new state
         inputs: the original embedded inputs (= inputs to the first step)
-        memory: contains states from all the previous steps.
+        memory: contains information of state from all the previous steps.
   """
-  state, inputs, memory = layer_inputs
-  state = step_preprocess(state, step, hparams)
-
-  state = common_layers.layer_preprocess(state, hparams)
-  inputs = common_layers.layer_preprocess(inputs, hparams)
 
-  input_gate = _ffn_layer_multi_inputs(
-      [inputs, state],
-      hparams,
-      name="input_g",
-      bias_initializer=tf.zeros_initializer(),
-      activation=tf.sigmoid,
-      pad_remover=pad_remover)
-
-  forget_gate = _ffn_layer_multi_inputs(
-      [inputs, state],
-      hparams,
-      name="forget_g",
-      bias_initializer=tf.zeros_initializer(),
-      activation=None,
-      pad_remover=pad_remover)
-
-  output_gate = _ffn_layer_multi_inputs(
-      [inputs, state],
-      hparams,
-      name="output_g",
-      bias_initializer=tf.zeros_initializer(),
-      activation=tf.sigmoid,
-      pad_remover=pad_remover)
+  state, unused_inputs, memory = tf.unstack(
+      layer_inputs, num=None, axis=0, name="unstack")
+  # NOTE:
+  # state (ut_state): output of the lstm in the previous step
+  # inputs (ut_input): original input --> we don't use it here
+  # memory: lstm memory
+
+  # Multi_head_attention:
+  assert not hparams.add_step_timing_signal  # Let lstm count for us!
+  mh_attention_input = step_preprocess(state, step, hparams)
+  transition_function_input = attention_unit(mh_attention_input)
+
+  # Transition Function:
+  if hparams.add_ffn_unit_to_the_transition_function:
+    transition_function_input = ffn_unit(transition_function_input)
+
+  transition_function_input = common_layers.layer_preprocess(
+      transition_function_input, hparams)
+  with tf.variable_scope("lstm"):
+    # lstm input gate: i_t = sigmoid(W_i.x_t + U_i.h_{t-1})
+    transition_function_input_gate = _ffn_layer_multi_inputs(
+        [transition_function_input, state],
+        hparams,
+        name="input",
+        bias_initializer=tf.zeros_initializer(),
+        activation=tf.sigmoid,
+        pad_remover=pad_remover,
+        preprocess=False,
+        postprocess=False)
 
-  input_modulation = _ffn_layer_multi_inputs(
-      [inputs, state],
-      hparams,
-      name="input_modulation",
-      bias_initializer=tf.zeros_initializer(),
-      activation=tf.tanh,
-      pad_remover=pad_remover)
+    tf.contrib.summary.scalar("lstm_input_gate",
+                              tf.reduce_mean(transition_function_input_gate))
 
-  forget_bias_tensor = tf.constant(hparams.lstm_forget_bias)
-  forget_gate = tf.sigmoid(forget_gate + forget_bias_tensor)
+    # lstm forget gate: f_t = sigmoid(W_f.x_t + U_f.h_{t-1})
+    transition_function_forget_gate = _ffn_layer_multi_inputs(
+        [transition_function_input, state],
+        hparams,
+        name="forget",
+        bias_initializer=tf.zeros_initializer(),
+        activation=None,
+        pad_remover=pad_remover,
+        preprocess=False,
+        postprocess=False)
+    forget_bias_tensor = tf.constant(hparams.lstm_forget_bias)
+    transition_function_forget_gate = tf.sigmoid(
+        transition_function_forget_gate + forget_bias_tensor)
+
+    tf.contrib.summary.scalar("lstm_forget_gate",
+                              tf.reduce_mean(transition_function_forget_gate))
+
+    # lstm output gate: o_t = sigmoid(W_o.x_t + U_o.h_{t-1})
+    transition_function_output_gate = _ffn_layer_multi_inputs(
+        [transition_function_input, state],
+        hparams,
+        name="output",
+        bias_initializer=tf.zeros_initializer(),
+        activation=tf.sigmoid,
+        pad_remover=pad_remover,
+        preprocess=False,
+        postprocess=False)
 
-  if "modulated_input_transformation" in hparams.lstm_transformation:
-    input_modulation = attention_unit(input_modulation)
+    tf.contrib.summary.scalar("lstm_output_gate",
+                              tf.reduce_mean(transition_function_output_gate))
 
-  memory = memory * forget_gate + input_gate * input_modulation
+    # lstm input modulation
+    transition_function_input_modulation = _ffn_layer_multi_inputs(
+        [transition_function_input, state],
+        hparams,
+        name="input_modulation",
+        bias_initializer=tf.zeros_initializer(),
+        activation=tf.tanh,
+        pad_remover=pad_remover,
+        preprocess=False,
+        postprocess=False)
 
-  if "memory_transformation" in hparams.lstm_transformation:
-    memory = attention_unit(memory)
+    transition_function_memory = (
+        memory * transition_function_forget_gate +
+        transition_function_input_gate * transition_function_input_modulation)
 
-  new_state = tf.tanh(memory) * output_gate
+    transition_function_output = (
+        tf.tanh(transition_function_memory) * transition_function_output_gate)
 
-  if "state_transformation" in hparams.lstm_transformation:
-    new_state = attention_unit(new_state)
+  transition_function_output = common_layers.layer_preprocess(
+      transition_function_output, hparams)
 
-  return new_state, inputs, memory
+  return transition_function_output, unused_inputs, transition_function_memory
 
 
 def universal_transformer_act(x, hparams, ffn_unit, attention_unit):
@@ -1068,7 +1037,7 @@ def universal_transformer_act(x, hparams, ffn_unit, attention_unit):
     ValueError: Unknown act type
 
   """
-  # TODO(dehghani): Use pad_remover for the act computations.
+  # TODO(dehghani): Enable pad_remover for the act computations.
   if hparams.act_type == "basic":
     return universal_transformer_act_basic(
         x, hparams, ffn_unit, attention_unit)
@@ -1139,7 +1108,7 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
 
     Args:
       state: 3-D Tensor: [batch_size, length, channel]
-      step: indicating number of steps take so far
+      step: indicates number of steps taken so far
       halting_probability: halting probability
       remainders: act remainders
       n_updates: act n_updates
@@ -1291,7 +1260,7 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
 
     Args:
       state: 3-D Tensor: [batch_size, length, channel]
-      step: indicating number of steps take so far
+      step: indicates number of steps taken so far
       halting_probability: halting probability
       remainders: act remainders
       n_updates: act n_updates
@@ -1427,7 +1396,7 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
 
     Args:
       state: 3-D Tensor: [batch_size, length, channel]
-      step: indicating number of steps take so far
+      step: indicates number of steps taken so far
       halting_probability: halting probability
       remainders: act remainders
       n_updates: act n_updates
@@ -1583,7 +1552,7 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
 
     Args:
       state: 3-D Tensor: [batch_size, length, channel]
-      step: indicating number of steps take so far
+      step: indicates number of steps taken so far
       halting_probability: halting probability
       remainders: act remainders
       n_updates: act n_updates
@@ -1688,8 +1657,8 @@ def _ffn_layer_multi_inputs(inputs_list,
                             bias_initializer=None,
                             activation=None,
                             pad_remover=None,
-                            preprocess=True,
-                            postprocess=True):
+                            preprocess=False,
+                            postprocess=False):
   """Implements a Feed-forward layer with multiple inputs, pad-removing, etc.
 
   Args:

From 8b39868d9678da5524ebb9af51be3ba17be6dc20 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 5 Sep 2018 15:52:55 -0700
Subject: [PATCH 0744/2720] internal merge of PR #1036

PiperOrigin-RevId: 211714231
---
 tensor2tensor/models/research/universal_transformer.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 212294c44..48100f2c6 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -86,7 +86,6 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
 
     return encoder_output, encoder_decoder_attention_bias, encoder_extra_output
 
-
   def decode(self,
              decoder_input,
              encoder_output,
@@ -415,12 +414,11 @@ def update_hparams_for_universal_transformer(hparams):
 
   # LSTM forget bias for lstm style recurrence.
   hparams.add_hparam("lstm_forget_bias", 1.0)
-  # Uses the memory at the last step as the final ouput, if true.
+  # Uses the memory at the last step as the final output, if true.
   hparams.add_hparam("use_memory_as_final_state", True)
   # if also add a ffn unit to the transition function when using gru/lstm
   hparams.add_hparam("add_ffn_unit_to_the_transition_function", False)
 
-
   # Type of act: basic/accumulated/global (instead of position-wise!)/random.
   hparams.add_hparam("act_type", "basic")
   # Max number of steps (forces halting at this step).
@@ -501,6 +499,7 @@ def adaptive_universal_transformer_base():
   hparams.recurrence_type = "act"
   return hparams
 
+
 @registry.register_hparams
 def adaptive_universal_transformer_small():
   hparams = universal_transformer_small()
@@ -555,7 +554,6 @@ def adaptive_universal_transformer_tall_actlossw001():
   return hparams
 
 
-
 @registry.register_hparams
 def adaptive_universal_transformer_base_dropout03():
   hparams = universal_transformer_base()
@@ -576,7 +574,6 @@ def adaptive_universal_transformer_base_dropout05():
   return hparams
 
 
-
 @registry.register_hparams
 def universal_transformer_skip_base():
   hparams = universal_transformer_base()
@@ -590,12 +587,14 @@ def universal_transformer_highway_base():
   hparams.recurrence_type = "highway"
   return hparams
 
+
 @registry.register_hparams
 def universal_transformer_dwa_base():
   hparams = universal_transformer_base()
   hparams.recurrence_type = "dwa"
   return hparams
 
+
 @registry.register_hparams
 def universal_transformer_lstm_base():
   hparams = universal_transformer_base()
@@ -708,6 +707,7 @@ def universal_transformer_sepconv_base():
   hparams.transformer_ffn_type = "sepconv"
   return hparams
 
+
 @registry.register_ranged_hparams
 def universal_transformer_base_range(rhp):
   """Range of hyperparameters."""

From 84d04d278afc5a46861e85ead06416ad4d0234a3 Mon Sep 17 00:00:00 2001
From: Tomasz Latkowski <13836101+tlatkowski@users.noreply.github.com>
Date: Thu, 6 Sep 2018 01:37:52 +0200
Subject: [PATCH 0745/2720] small refactor in speech problems (#1035)

---
 .../data_generators/audio_encoder.py          | 96 +++++++++++++++++++
 tensor2tensor/data_generators/librispeech.py  |  4 +-
 .../data_generators/speech_recognition.py     | 84 +---------------
 3 files changed, 101 insertions(+), 83 deletions(-)
 create mode 100644 tensor2tensor/data_generators/audio_encoder.py

diff --git a/tensor2tensor/data_generators/audio_encoder.py b/tensor2tensor/data_generators/audio_encoder.py
new file mode 100644
index 000000000..5ff6f9f83
--- /dev/null
+++ b/tensor2tensor/data_generators/audio_encoder.py
@@ -0,0 +1,96 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Encoder for audio data."""
+
+import os
+import tempfile
+import numpy as np
+
+from scipy.io import wavfile
+from subprocess import call
+
+
+class AudioEncoder(object):
+  """Encoder class for saving and loading waveforms."""
+
+  def __init__(self, num_reserved_ids=0, sample_rate=16000):
+    assert num_reserved_ids == 0
+    self._sample_rate = sample_rate
+
+  @property
+  def num_reserved_ids(self):
+    return 0
+
+  def encode(self, s):
+    """Transform a string with a filename into a list of float32.
+
+    Args:
+      s: path to the file with a waveform.
+
+    Returns:
+      samples: list of int16s
+    """
+    # Make sure that the data is a single channel, 16bit, 16kHz wave.
+    # TODO(chorowski): the directory may not be writable, this should fallback
+    # to a temp path, and provide instructions for installing sox.
+    if s.endswith(".mp3"):
+      # TODO(dliebling) On Linux, check if libsox-fmt-mp3 is installed.
+      out_filepath = s[:-4] + ".wav"
+      call([
+          "sox", "--guard", s, "-r", "16k", "-b", "16", "-c", "1", out_filepath
+      ])
+      s = out_filepath
+    elif not s.endswith(".wav"):
+      out_filepath = s + ".wav"
+      if not os.path.exists(out_filepath):
+        call(["sox", "-r", "16k", "-b", "16", "-c", "1", s, out_filepath])
+      s = out_filepath
+    rate, data = wavfile.read(s)
+    assert rate == self._sample_rate
+    assert len(data.shape) == 1
+    if data.dtype not in [np.float32, np.float64]:
+      data = data.astype(np.float32) / np.iinfo(data.dtype).max
+    return data.tolist()
+
+  def decode(self, ids):
+    """Transform a sequence of float32 into a waveform.
+
+    Args:
+      ids: list of integers to be converted.
+
+    Returns:
+      Path to the temporary file where the waveform was saved.
+
+    Raises:
+      ValueError: if the ids are not of the appropriate size.
+    """
+    _, tmp_file_path = tempfile.mkstemp()
+    wavfile.write(tmp_file_path, self._sample_rate, np.asarray(ids))
+    return tmp_file_path
+
+  def decode_list(self, ids):
+    """Transform a sequence of int ids into an image file.
+
+    Args:
+      ids: list of integers to be converted.
+
+    Returns:
+      Singleton list: path to the temporary file where the wavfile was saved.
+    """
+    return [self.decode(ids)]
+
+  @property
+  def vocab_size(self):
+    return 256
\ No newline at end of file
diff --git a/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index 37adc4aa6..84498b380 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -16,6 +16,7 @@
 
 import os
 import tarfile
+import tqdm  # pylint: disable=g-bad-import-order
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import speech_recognition
@@ -141,7 +142,8 @@ def generator(self, data_dir, tmp_dir, datasets,
       audio_encoder = encoders["waveforms"]
       text_encoder = encoders["targets"]
 
-      for utt_id, media_file, text_data in sorted(data_pairs)[start_from:]:
+      for utt_id, media_file, text_data in tqdm.tqdm(
+          sorted(data_pairs)[start_from:]):
         if how_many > 0 and i == how_many:
           return
         i += 1
diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 9c0726a86..dc318a80f 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -19,14 +19,11 @@
 """
 
 import functools
-import os
-from subprocess import call
-import tempfile
 import numpy as np
-from scipy.io import wavfile
 import scipy.signal
 
 from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import audio_encoder
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
@@ -153,83 +150,6 @@ def compute_mel_filterbank_features(
   return tf.expand_dims(log_mel_sgram, -1, name="mel_sgrams")
 
 
-#
-# Audio problem definition
-#
-class AudioEncoder(object):
-  """Encoder class for saving and loading waveforms."""
-
-  def __init__(self, num_reserved_ids=0, sample_rate=16000):
-    assert num_reserved_ids == 0
-    self._sample_rate = sample_rate
-
-  @property
-  def num_reserved_ids(self):
-    return 0
-
-  def encode(self, s):
-    """Transform a string with a filename into a list of float32.
-
-    Args:
-      s: path to the file with a waveform.
-
-    Returns:
-      samples: list of int16s
-    """
-    # Make sure that the data is a single channel, 16bit, 16kHz wave.
-    # TODO(chorowski): the directory may not be writable, this should fallback
-    # to a temp path, and provide instructions for installing sox.
-    if s.endswith(".mp3"):
-      # TODO(dliebling) On Linux, check if libsox-fmt-mp3 is installed.
-      out_filepath = s[:-4] + ".wav"
-      call([
-          "sox", "--guard", s, "-r", "16k", "-b", "16", "-c", "1", out_filepath
-      ])
-      s = out_filepath
-    elif not s.endswith(".wav"):
-      out_filepath = s + ".wav"
-      if not os.path.exists(out_filepath):
-        call(["sox", "-r", "16k", "-b", "16", "-c", "1", s, out_filepath])
-      s = out_filepath
-    rate, data = wavfile.read(s)
-    assert rate == self._sample_rate
-    assert len(data.shape) == 1
-    if data.dtype not in [np.float32, np.float64]:
-      data = data.astype(np.float32) / np.iinfo(data.dtype).max
-    return data.tolist()
-
-  def decode(self, ids):
-    """Transform a sequence of float32 into a waveform.
-
-    Args:
-      ids: list of integers to be converted.
-
-    Returns:
-      Path to the temporary file where the waveform was saved.
-
-    Raises:
-      ValueError: if the ids are not of the appropriate size.
-    """
-    _, tmp_file_path = tempfile.mkstemp()
-    wavfile.write(tmp_file_path, self._sample_rate, np.asarray(ids))
-    return tmp_file_path
-
-  def decode_list(self, ids):
-    """Transform a sequence of int ids into an image file.
-
-    Args:
-      ids: list of integers to be converted.
-
-    Returns:
-      Singleton list: path to the temporary file where the wavfile was saved.
-    """
-    return [self.decode(ids)]
-
-  @property
-  def vocab_size(self):
-    return 256
-
-
 class ByteTextEncoderWithEos(text_encoder.ByteTextEncoder):
   """Encodes each byte to an id and appends the EOS token."""
 
@@ -279,7 +199,7 @@ def feature_encoders(self, _):
         "inputs": None,  # Put None to make sure that the logic in
                          # decoding.py doesn't try to convert the floats
                          # into text...
-        "waveforms": AudioEncoder(),
+        "waveforms": audio_encoder.AudioEncoder(),
         "targets": ByteTextEncoderWithEos(),
     }
 

From 6217df3df411129daa92b4cb569517b3f93a9c1e Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 5 Sep 2018 16:48:53 -0700
Subject: [PATCH 0746/2720] internal merge of PR #1035

PiperOrigin-RevId: 211722821
---
 tensor2tensor/data_generators/audio_encoder.py      | 5 ++---
 tensor2tensor/data_generators/librispeech.py        | 4 +---
 tensor2tensor/data_generators/speech_recognition.py | 2 +-
 3 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/audio_encoder.py b/tensor2tensor/data_generators/audio_encoder.py
index 5ff6f9f83..74cdb73e9 100644
--- a/tensor2tensor/data_generators/audio_encoder.py
+++ b/tensor2tensor/data_generators/audio_encoder.py
@@ -15,11 +15,10 @@
 """Encoder for audio data."""
 
 import os
+from subprocess import call
 import tempfile
 import numpy as np
-
 from scipy.io import wavfile
-from subprocess import call
 
 
 class AudioEncoder(object):
@@ -93,4 +92,4 @@ def decode_list(self, ids):
 
   @property
   def vocab_size(self):
-    return 256
\ No newline at end of file
+    return 256
diff --git a/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index 84498b380..37adc4aa6 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -16,7 +16,6 @@
 
 import os
 import tarfile
-import tqdm  # pylint: disable=g-bad-import-order
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import speech_recognition
@@ -142,8 +141,7 @@ def generator(self, data_dir, tmp_dir, datasets,
       audio_encoder = encoders["waveforms"]
       text_encoder = encoders["targets"]
 
-      for utt_id, media_file, text_data in tqdm.tqdm(
-          sorted(data_pairs)[start_from:]):
+      for utt_id, media_file, text_data in sorted(data_pairs)[start_from:]:
         if how_many > 0 and i == how_many:
           return
         i += 1
diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index dc318a80f..aaee536ea 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -22,8 +22,8 @@
 import numpy as np
 import scipy.signal
 
-from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import audio_encoder
+from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers

From d0e440a6b1977f25b801ee125a6f27f8392ada9d Mon Sep 17 00:00:00 2001
From: Youlong Cheng <ylc@google.com>
Date: Wed, 5 Sep 2018 16:58:52 -0700
Subject: [PATCH 0747/2720] Use tpu_ops.all_to_all in mesh_tensorflow.

PiperOrigin-RevId: 211724278
---
 .../mesh_tensorflow/simd_mesh_impl.py         | 41 ++++++++++++-------
 1 file changed, 26 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
index db9317b9d..369b4806b 100644
--- a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
@@ -162,6 +162,20 @@ def laid_out_pnum(self):
     """
     return self.LaidOutTensor([self.pnum_tensor])
 
+  def _create_group_assignment(self, mesh_axes):
+    """Create group assignment for XLA cross replica ops."""
+
+    partitioning = {}
+    for pnum in xrange(self.size):
+      group = mtf.pnum_to_group(self.shape, mesh_axes, pnum)
+      if group not in partitioning:
+        partitioning[group] = []
+      partitioning[group].append(pnum)
+    group_assignment = []
+    for group, pnums in partitioning.items():
+      group_assignment.append(pnums)
+    return group_assignment
+
   def allreduce(self, x, mesh_axes, reduction_fn_string):
     """Grouped allreduce, (summed across the given dimensions).
 
@@ -178,16 +192,7 @@ def allreduce(self, x, mesh_axes, reduction_fn_string):
       return x
     x = x.to_laid_out_tensor()
     if reduction_fn_string == "SUM":
-      partitioning = {}
-      for pnum in xrange(self.size):
-        group = mtf.pnum_to_group(self.shape, mesh_axes, pnum)
-        if group not in partitioning:
-          partitioning[group] = []
-        partitioning[group].append(pnum)
-      group_assignment = []
-      for group, pnums in partitioning.items():
-        group_assignment.append(pnums)
-
+      group_assignment = self._create_group_assignment(mesh_axes)
       return self.LaidOutTensor(
           [tpu_ops.cross_replica_sum(x.one_slice, group_assignment)])
     else:
@@ -229,9 +234,6 @@ def allconcat(self, x, mesh_axis, concat_axis, stack=False):
   def alltoall(self, x, mesh_axis, split_axis, concat_axis):
     """Grouped alltoall (like MPI alltoall with splitting and concatenation).
 
-    TODO(noam): this is a terribly inefficient implementation using allreduce.
-    Replace this with a native xla alltoall once it is ready.
-
     Args:
       x: a LaidOutTensor
       mesh_axis: an integer the mesh axis along which to group
@@ -241,8 +243,17 @@ def alltoall(self, x, mesh_axis, split_axis, concat_axis):
       a LaidOutTensor
     """
     x = x.to_laid_out_tensor()
-    x = self.allconcat(x, mesh_axis, concat_axis)
-    x = self.allsplit(x, mesh_axis, split_axis)
+    t = x.one_slice
+    group_assignment = self._create_group_assignment([mesh_axis])
+
+    t = tpu_ops.all_to_all(
+        t,
+        concat_dimension=concat_axis,
+        split_dimension=split_axis,
+        split_count=len(group_assignment[0]),
+        group_assignment=group_assignment)
+    x = self.LaidOutTensor([t])
+
     return x
 
   def receive(self, x, mesh_axis, source_pcoord):

From 91534ac0f3f218f2751b9ed38635d5b76f683845 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Wed, 5 Sep 2018 18:09:38 -0700
Subject: [PATCH 0748/2720] Adding short ppo training on real env between
 epochs

---
 tensor2tensor/rl/trainer_model_based.py | 46 +++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index c89f91d5b..09969edb8 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -183,6 +183,38 @@ def train_agent(problem_name, agent_model_dir,
   }):
     rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=epoch)
 
+def train_agent_real_env(problem_name, agent_model_dir,
+                event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
+                is_final_epoch=False):
+  """Train the PPO agent in the simulated environment."""
+  gym_problem = registry.problem(problem_name)
+  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
+  ppo_params_names = ["epochs_num", "epoch_length",
+                      "learning_rate", "num_agents",
+                      "optimization_epochs"]
+
+  for param_name in ppo_params_names:
+    ppo_param_name = "real_ppo_"+ param_name
+    if ppo_param_name in hparams:
+      ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
+
+  ppo_epochs_num = hparams.real_ppo_epochs_num
+  if ppo_epochs_num == 0:
+    return
+
+  ppo_hparams.save_models_every_epochs = ppo_epochs_num #check this
+
+  environment_spec = copy.copy(gym_problem.environment_spec)
+
+  ppo_hparams.add_hparam("environment_spec", environment_spec)
+
+  with temporary_flags({
+      "problem": problem_name,
+      "output_dir": world_model_dir,
+      "data_dir": epoch_data_dir,
+  }):
+    # epoch = 0 is a hackish way to avoid skiping training
+    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=0)
 
 def evaluate_world_model(simulated_problem_name, problem_name, hparams,
                          world_model_dir, epoch_data_dir, tmp_dir):
@@ -452,6 +484,12 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
                 ppo_event_dir, directories["world_model"], epoch_data_dir,
                 hparams, epoch=epoch, is_final_epoch=is_final_epoch)
 
+    # Train PPO on real env (short)
+    train_agent_real_env(problem_name, ppo_model_dir,
+                ppo_event_dir, directories["world_model"], epoch_data_dir,
+                hparams, epoch=epoch, is_final_epoch=is_final_epoch)
+
+
     # Collect data from the real environment.
     log("Generating real environment data")
     eval_data_dir = os.path.join(epoch_data_dir, "eval")
@@ -537,6 +575,13 @@ def rl_modelrl_base():
       # Whether the PPO agent should be restored from the previous iteration, or
       # should start fresh each time.
       ppo_continue_training=True,
+
+      real_ppo_epochs_num=30,
+      real_ppo_epoch_length=200,
+      real_ppo_num_agents=16,
+      real_ppo_learning_rate=2e-4,
+      real_ppo_continue_training=True,
+
       game="wrapped_full_pong",
       # Whether to evaluate the world model in each iteration of the loop to get
       # the model_reward_accuracy metric.
@@ -617,6 +662,7 @@ def rl_modelrl_tiny():
           ppo_time_limit=5,
           ppo_epoch_length=5,
           ppo_num_agents=2,
+          real_ppo_epochs_num=4,
           generative_model_params="next_frame_tiny",
       ).values())
 

From 24d1f140d0a57fbe5fd21e0c9e3f73adcc9d391a Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 5 Sep 2018 18:49:08 -0700
Subject: [PATCH 0749/2720] Abstract T2TModel's summarize losses functionality.

PiperOrigin-RevId: 211736697
---
 tensor2tensor/utils/t2t_model.py      | 12 +++++---
 tensor2tensor/utils/t2t_model_test.py | 40 +++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 4 deletions(-)
 create mode 100644 tensor2tensor/utils/t2t_model_test.py

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index f91600bfa..f8b2aae60 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1304,10 +1304,7 @@ def estimator_model_fn(cls,
       return logits
 
     # Summarize losses
-    if common_layers.should_generate_summaries():
-      with tf.name_scope("losses"):
-        for loss_name, loss_val in sorted(losses_dict.items()):
-          tf.summary.scalar(loss_name, loss_val)
+    model._summarize_losses(losses_dict)  # pylint: disable=protected-access
 
     # Accumulate losses
     loss = sum(losses_dict[key] for key in sorted(losses_dict.keys()))
@@ -1501,6 +1498,13 @@ def _normalize_body_output(self, body_out):
 
     return output, losses
 
+  def _summarize_losses(self, losses_dict):
+    """Adds `tf.summary`s to all terms in the losses dictionary."""
+    if common_layers.should_generate_summaries():
+      with tf.name_scope("losses"):
+        for loss_name, loss_val in sorted(losses_dict.items()):
+          tf.summary.scalar(loss_name, loss_val)
+
 
 def _warn_changed_modality_type(new_name, old_name, feature_name):
   new_type, new_name = registry.parse_modality_name(new_name)
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
new file mode 100644
index 000000000..66bf5eb8e
--- /dev/null
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -0,0 +1,40 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 T2TModel."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.utils import t2t_model
+
+import tensorflow as tf
+
+
+class T2TModelTest(tf.test.TestCase):
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testSummarizeLosses(self):
+    with tf.Graph().as_default():
+      model = t2t_model.T2TModel(tf.contrib.training.HParams())
+      losses = {"training": tf.random_normal([]),
+                "extra": tf.random_normal([])}
+      outputs = model._summarize_losses(losses)
+      self.assertIsNone(outputs, None)
+      self.assertLen(tf.get_collection(tf.GraphKeys.SUMMARIES, scope="losses"),
+                     len(losses))
+
+if __name__ == "__main__":
+  tf.test.main()

From 20c2007b7d96b6f0895713558afa96470c0548e1 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Wed, 5 Sep 2018 18:50:18 -0700
Subject: [PATCH 0750/2720] Implement word shuffle as in
 https://arxiv.org/abs/1804.07755

PiperOrigin-RevId: 211736784
---
 .../models/research/transformer_vae.py        | 31 ++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index f2a4a7dfa..10be631f1 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -361,6 +361,20 @@ def ae_transformer_internal(inputs,
     else:
       assert hparams.task == "image"
       max_targets_len_from_inputs = targets
+    if hparams.word_shuffle:
+      tf.logging.info("Using word shuffle with rate = {}".format(
+          hparams.word_shuffle))
+      targets_idx = tf.range(start=0,
+                             limit=common_layers.shape_list(targets)[1],
+                             delta=1)
+      targets_idx = tf.to_float(targets_idx)
+      noise = tf.random_uniform(shape=common_layers.shape_list(targets_idx),
+                                minval=0,
+                                maxval=1 + hparams.word_shuffle)
+      targets_idx += noise
+      permutation = tf.contrib.framework.argsort(targets_idx)
+      targets_permuted = tf.gather(targets, indices=permutation, axis=1)
+      targets = targets_permuted
     targets, _ = common_layers.pad_to_same_length(
         targets, max_targets_len_from_inputs,
         final_length_divisible_by=2**hparams.num_compress_steps)
@@ -720,7 +734,7 @@ def transformer_ae_small():
   hparams.add_hparam("noise_dev", 0.5)
   hparams.add_hparam("d_mix", 0.5)
   hparams.add_hparam("logit_normalization", True)
-  hparams.add_hparam("word_dropout", 0.0)
+  hparams.add_hparam("word_dropout", True)
   # Bottleneck kinds supported: dense, vae, semhash, gumbel-softmax, dvq.
   hparams.add_hparam("bottleneck_kind", "semhash")
   hparams.add_hparam("num_blocks", 1)
@@ -728,6 +742,7 @@ def transformer_ae_small():
   # Add an hparam for number of reiduals
   hparams.add_hparam("num_residuals", 1)
   # Reshape method for DVQ: slice, project
+  hparams.add_hparam("word_shuffle", 0.5)
   hparams.add_hparam("causal", True)
   hparams.add_hparam("reshape_method", "slice")
   hparams.add_hparam("trainable_projections", False)
@@ -922,6 +937,20 @@ def transformer_ae_base_noatt():
   return hparams
 
 
+@registry.register_hparams
+def transformer_ae_small_noatt():
+  """Set of hyperparameters."""
+  hparams = transformer_ae_small()
+  hparams.reshape_method = "slice"
+  hparams.bottleneck_kind = "dvq"
+  hparams.hidden_size = 512
+  hparams.num_blocks = 1
+  hparams.num_decode_blocks = 1
+  hparams.z_size = 12
+  hparams.do_attend_decompress = False
+  return hparams
+
+
 @registry.register_hparams
 def transformer_ae_base_ablation_1():
   hparams = transformer_ae_base_noatt()

From f23805b7652e2b4447ec6f372711b5db447b1021 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 5 Sep 2018 20:23:14 -0700
Subject: [PATCH 0751/2720] Use max_relative_position instead of max_length for
 setting the maximum relative distance to consider in relative attention. 
 Setting max_relative_position to be less than max_length allows for better
 generalization and for inference beyond max_length.  Refactored the three
 different situations of using relative attention (masked, unmasked, local) to
 share common helper functions. Also adding support to not share embeddings
 across heads and to compute relative values whenever missing.

PiperOrigin-RevId: 211744110
---
 tensor2tensor/layers/common_attention.py      | 465 ++++++++++++------
 tensor2tensor/layers/common_attention_test.py | 172 +++++++
 tensor2tensor/layers/common_hparams.py        |   9 +-
 tensor2tensor/models/transformer.py           |  18 +-
 4 files changed, 504 insertions(+), 160 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index a35081581..b84b1b3e8 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1678,16 +1678,68 @@ def _absolute_position_to_relative_position_masked(x):
   return x
 
 
+def get_relative_embeddings_left(max_relative_position, length, depth,
+                                 num_heads, heads_share_relative_embedding,
+                                 name):
+  """Instantiate or retrieve relative embeddings, sliced according to length.
+
+  Use for masked case where the relative attention is only looking left.
+
+  Args:
+    max_relative_position: an Integer for the number of entries in the relative
+      embedding, which corresponds to the max relative distance that is
+      considered.
+    length: an Integer, specifies the length of the input sequence for which
+      this relative embedding is retrieved for.
+    depth: an Integer, specifies the depth for relative embeddings.
+    num_heads: an Integer, specifies the number of heads.
+    heads_share_relative_embedding: a Boolean specifying if the relative
+      embedding is shared across heads.
+    name: a string giving the name of the embedding variables.
+
+  Returns:
+    a Tensor with shape [length, depth]
+  """
+  initializer_stddev = depth**-0.5
+  if heads_share_relative_embedding:
+    embedding_shape = (max_relative_position, depth)
+  else:
+    embedding_shape = (num_heads, max_relative_position, depth)
+  relative_embeddings = tf.get_variable(
+      name=name, shape=embedding_shape,
+      initializer=tf.random_normal_initializer(stddev=initializer_stddev))
+  # Pad first before slice to avoid using tf.cond.
+  pad_length = tf.maximum(length - max_relative_position, 0)
+  start_slice_position = tf.maximum(max_relative_position - length, 0)
+  if heads_share_relative_embedding:
+    padded_relative_embeddings = tf.pad(
+        relative_embeddings,
+        [[pad_length, 0], [0, 0]])
+    used_relative_embeddings = tf.slice(
+        padded_relative_embeddings,
+        [start_slice_position, 0], [length, -1])
+  else:
+    padded_relative_embeddings = tf.pad(
+        relative_embeddings,
+        [[0, 0], [pad_length, 0], [0, 0]])
+    used_relative_embeddings = tf.slice(
+        padded_relative_embeddings,
+        [0, start_slice_position, 0], [-1, length, -1])
+  return used_relative_embeddings
+
+
 def dot_product_self_attention_relative_v2(q,
                                            k,
                                            v,
                                            bias,
-                                           max_length=None,
+                                           max_relative_position=None,
                                            dropout_rate=0.0,
                                            image_shapes=None,
                                            name=None,
                                            make_image_summary=True,
-                                           dropout_broadcast_dims=None):
+                                           dropout_broadcast_dims=None,
+                                           heads_share_relative_embedding=False,
+                                           add_relative_to_values=False):
   """Calculate relative position-aware dot-product self-attention.
 
   Only works for masked self-attention (no looking forward).
@@ -1700,7 +1752,8 @@ def dot_product_self_attention_relative_v2(q,
     k: a Tensor with shape [batch, heads, length, depth].
     v: a Tensor with shape [batch, heads, length, depth].
     bias: bias Tensor.
-    max_length: an integer - changing this invalidates checkpoints
+    max_relative_position: an integer indicating the maximum relative distance
+      to look back - changing this invalidates checkpoints
     dropout_rate: a floating point number.
     image_shapes: optional tuple of integer scalars.
     name: an optional string.
@@ -1708,10 +1761,20 @@ def dot_product_self_attention_relative_v2(q,
     dropout_broadcast_dims:  an optional list of integers less than 4
       specifying in which dimensions to broadcast the dropout decisions.
       saves memory.
+    heads_share_relative_embedding: a boolean indicating wheather to share
+      relative embeddings between attention heads.
+    add_relative_to_values: a boolean for whether to add relative component to
+      values.
 
   Returns:
     A Tensor.
+
+  Raises:
+    ValueError: if max_relative_position is not > 0.
   """
+  if not max_relative_position:
+    raise ValueError("Max relative position (%s) should be > 0 when using "
+                     "relative self attention." % (max_relative_position))
   with tf.variable_scope(
       name,
       default_name="dot_product_self_attention_relative_v2",
@@ -1723,39 +1786,39 @@ def dot_product_self_attention_relative_v2(q,
     q.get_shape().assert_is_compatible_with(v.get_shape())
 
     # Use separate embeddings suitable for keys and values.
-    length = common_layers.shape_list(q)[2]
-    assert max_length is not None
+    _, num_heads, length, depth_k = common_layers.shape_list(k)
 
     # [batch, num_heads, query_length, memory_length]
     logits = tf.matmul(q, k, transpose_b=True)
+    key_relative_embeddings = get_relative_embeddings_left(
+        max_relative_position, length, depth_k, num_heads,
+        heads_share_relative_embedding, "key_relative_embededings")
 
-    # now add relative logits
-    # [batch, num_heads, query_length, max_length]
-    rel_logits = common_layers.dense(q, max_length, name="rel0")
-    # [batch, num_heads, query_length, max_length]
-    rel_logits = tf.slice(rel_logits, [0, 0, 0, max_length - length],
-                          [-1, -1, -1, -1])
+    rel_logits = matmul_with_relative_keys(q, key_relative_embeddings,
+                                           heads_share_relative_embedding)
     rel_logits = _relative_position_to_absolute_position_masked(rel_logits)
     logits += rel_logits
-
     if bias is not None:
       logits += bias
+
     weights = tf.nn.softmax(logits, name="attention_weights")
-    # dropping out the attention links for each of the heads
+    # Dropping out the attention links for each of the heads.
     weights = common_layers.dropout_with_broadcast_dims(
         weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
     if common_layers.should_generate_summaries() and make_image_summary:
       attention_image_summary(weights, image_shapes)
-    ret = tf.matmul(weights, v)
-    # [batch, num_heads, query_length, memory_length]
-    relative_weights = _absolute_position_to_relative_position_masked(weights)
-    # [batch, num_heads, query_length, memory_length]
-    relative_weights = tf.pad(
-        relative_weights, [[0, 0], [0, 0], [0, 0], [max_length - length, 0]])
-    relative_weights.set_shape([None, None, None, max_length])
-    depth_v = common_layers.shape_list(v)[3]
-    ret += common_layers.dense(relative_weights, depth_v, name="rel1")
-    return ret
+    output = tf.matmul(weights, v)
+    if add_relative_to_values:
+      # [batch, num_heads, query_length, memory_length]
+      relative_weights = _absolute_position_to_relative_position_masked(weights)
+      depth_v = common_layers.shape_list(v)[3]
+      value_relative_embeddings = get_relative_embeddings_left(
+          max_relative_position, length, depth_v, num_heads,
+          heads_share_relative_embedding, "value_relative_embeddings")
+      output += matmul_with_relative_values(
+          relative_weights, value_relative_embeddings,
+          heads_share_relative_embedding)
+    return output
 
 
 def _absolute_position_to_relative_position_unmasked(x):
@@ -1789,9 +1852,63 @@ def _absolute_position_to_relative_position_unmasked(x):
   return x
 
 
+def get_relative_embeddings_left_right(max_relative_position, length, depth,
+                                       num_heads,
+                                       heads_share_relative_embedding,
+                                       name):
+  """Instantiate or retrieve relative embeddings, sliced according to length.
+
+  Use for unmasked case where the relative attention looks both left and right.
+
+  Args:
+    max_relative_position: an Integer for the number of entries in the relative
+      embedding, which corresponds to the max relative distance that is
+      considered.
+    length: an Integer, specifies the length of the input sequence for which
+      this relative embedding is retrieved for.
+    depth: an Integer, specifies the depth for relative embeddings.
+    num_heads: an Integer, specifies the number of heads.
+    heads_share_relative_embedding: a Boolean specifying if the relative
+      embedding is shared across heads.
+    name: a string giving the name of the embedding variables.
+
+  Returns:
+    a Tensor with shape [length, depth]
+  """
+  initializer_stddev = depth**-0.5
+  max_relative_position_unmasked = 2 * max_relative_position - 1
+  if heads_share_relative_embedding:
+    embedding_shape = (max_relative_position_unmasked, depth)
+  else:
+    embedding_shape = (num_heads, max_relative_position_unmasked, depth)
+  relative_embeddings = tf.get_variable(
+      name=name, shape=embedding_shape,
+      initializer=tf.random_normal_initializer(stddev=initializer_stddev))
+  # Pad first before slice to avoid using tf.cond.
+  pad_length = tf.maximum(length - max_relative_position, 0)
+  slice_start_position = tf.maximum(max_relative_position-length, 0)
+  if heads_share_relative_embedding:
+    padded_relative_embeddings = tf.pad(
+        relative_embeddings,
+        [[pad_length, pad_length], [0, 0]])
+    used_relative_embeddings = tf.slice(
+        padded_relative_embeddings,
+        [slice_start_position, 0], [2 * length - 1, -1])
+  else:
+    padded_relative_embeddings = tf.pad(
+        relative_embeddings,
+        [[0, 0], [pad_length, pad_length], [0, 0]])
+    used_relative_embeddings = tf.slice(
+        padded_relative_embeddings,
+        [0, slice_start_position, 0], [-1, 2 * length - 1, -1])
+  return used_relative_embeddings
+
+
 def dot_product_unmasked_self_attention_relative_v2(
-    q, k, v, bias, max_length=None, dropout_rate=0.0, image_shapes=None,
-    name=None, make_image_summary=True, dropout_broadcast_dims=None):
+    q, k, v, bias, max_relative_position=None, dropout_rate=0.0,
+    image_shapes=None, name=None, make_image_summary=True,
+    dropout_broadcast_dims=None, heads_share_relative_embedding=False,
+    add_relative_to_values=False):
   """Calculate relative position-aware dot-product self-attention.
 
   The attention calculation is augmented with learned representations for the
@@ -1802,7 +1919,8 @@ def dot_product_unmasked_self_attention_relative_v2(
     k: a Tensor with shape [batch, heads, length, depth].
     v: a Tensor with shape [batch, heads, length, depth].
     bias: bias Tensor.
-    max_length: an integer - changing this invalidates checkpoints
+    max_relative_position: an integer the max relative embedding considered.
+      Changing this invalidates checkpoints.
     dropout_rate: a floating point number.
     image_shapes: optional tuple of integer scalars.
     name: an optional string.
@@ -1810,10 +1928,21 @@ def dot_product_unmasked_self_attention_relative_v2(
     dropout_broadcast_dims:  an optional list of integers less than 4
       specifying in which dimensions to broadcast the dropout decisions.
       saves memory.
+    heads_share_relative_embedding: a boolean indicating wheather to share
+      relative embeddings between attention heads.
+    add_relative_to_values: a boolean for whether to add relative component to
+      values.
 
   Returns:
     A Tensor.
+
+  Raises:
+    ValueError: if max_relative_position is not > 0.
   """
+  if not max_relative_position:
+    raise ValueError("Max relative position (%s) should be > 0 when using "
+                     "relative self attention." % (max_relative_position))
+
   with tf.variable_scope(
       name,
       default_name="dot_product_unmasked_self_attention_relative_v2",
@@ -1824,37 +1953,20 @@ def dot_product_unmasked_self_attention_relative_v2(
     q.get_shape().assert_is_compatible_with(k.get_shape())
     q.get_shape().assert_is_compatible_with(v.get_shape())
 
-    # Use separate embeddings suitable for keys and values.
+    # [batch, num_heads, query_length, memory_length]
+    logits = tf.matmul(q, k, transpose_b=True)
+
     length = common_layers.shape_list(q)[2]
-    assert max_length is not None
     k_shape = common_layers.shape_list(k)
+    num_heads = k_shape[1]
     depth_k = k_shape[-1]
-    initializer_stddev = depth_k**-0.5
-    # TODO(avaswani): Add option for unshared relative embeddings
-    key_relative_embeddings = (
-        tf.get_variable(name="key_relative_embeddings",
-                        shape=(2*max_length-1, depth_k),
-                        initializer=tf.random_normal_initializer(
-                            stddev=initializer_stddev)))
-    # [batch, num_heads, query_length, memory_length]
-    logits = tf.matmul(q, k, transpose_b=True)
-    # slice out the right band of rel embeddings to save computation
-    # First pad the relative embeddings with zeros if the sequence length
-    # is longer than max length
-    padded_key_relative_embeddings = tf.pad(key_relative_embeddings,
-                                            [[tf.maximum(
-                                                length-max_length, 0),
-                                              tf.maximum(
-                                                  length-max_length, 0)],
-                                             [0, 0]])
-    used_key_relative_embeddings = tf.slice(padded_key_relative_embeddings,
-                                            [tf.maximum(
-                                                max_length-length,
-                                                length-max_length),
-                                             0],
-                                            [2*length -1, -1])
-    unmasked_rel_logits = tf.einsum("bhld,md->bhlm", q,
-                                    used_key_relative_embeddings)
+
+    key_relative_embeddings = get_relative_embeddings_left_right(
+        max_relative_position, length, depth_k, num_heads,
+        heads_share_relative_embedding,
+        "key_relative_embeddings")
+    unmasked_rel_logits = matmul_with_relative_keys(
+        q, key_relative_embeddings, heads_share_relative_embedding)
     unmasked_rel_logits = _relative_position_to_absolute_position_unmasked(
         unmasked_rel_logits)
     logits += unmasked_rel_logits
@@ -1869,33 +1981,18 @@ def dot_product_unmasked_self_attention_relative_v2(
     if common_layers.should_generate_summaries() and make_image_summary:
       attention_image_summary(weights, image_shapes)
     ret = tf.matmul(weights, v)
-    # getting the contribution of the relative embeddings to the values
-    # [batch, num_heads, query_length, 2*memory_length-1]
-    relative_weights = _absolute_position_to_relative_position_unmasked(
-        weights)
-    depth_v = common_layers.shape_list(v)[3]
-    initializer_stddev = depth_v**-0.5
-    value_relative_embeddings = (
-        tf.get_variable(name="value_relative_embeddings",
-                        shape=(2*max_length-1, depth_v),
-                        initializer=tf.random_normal_initializer(
-                            stddev=initializer_stddev)))
-    # slice out the right band of rel embeddings to save computation
-    padded_value_relative_embeddings = tf.pad(value_relative_embeddings,
-                                              [[tf.maximum(
-                                                  length-max_length, 0),
-                                                tf.maximum(
-                                                    length-max_length, 0)],
-                                               [0, 0]])
-    used_value_relative_embeddings = tf.slice(padded_value_relative_embeddings,
-                                              [tf.maximum(
-                                                  max_length-length,
-                                                  length-max_length),
-                                               0],
-                                              [2*length -1, -1])
-
-    ret += tf.einsum("bhlm,md->bhld", relative_weights,
-                     used_value_relative_embeddings)
+    if add_relative_to_values:
+      # Adds the contribution of the weighted relative embeddings to the values.
+      # [batch, num_heads, query_length, 2*memory_length-1]
+      relative_weights = _absolute_position_to_relative_position_unmasked(
+          weights)
+      depth_v = common_layers.shape_list(v)[3]
+      value_relative_embeddings = get_relative_embeddings_left_right(
+          max_relative_position, length, depth_v, num_heads,
+          heads_share_relative_embedding, "value_relative_embeddings")
+      ret += matmul_with_relative_values(
+          relative_weights, value_relative_embeddings,
+          heads_share_relative_embedding)
     return ret
 
 
@@ -2106,14 +2203,15 @@ def _make_local_block(x, depth, batch, heads, num_blocks, block_length):
                     [batch, heads, num_blocks - 1, block_length * 2, depth])
 
 
-def masked_rel_local_attention_1d(q,
-                                  k,
-                                  v,
-                                  block_length=128,
-                                  make_image_summary=False,
-                                  dropout_rate=0.,
-                                  share_rel_embed=False,
-                                  name=None):
+def masked_relative_local_attention_1d(q,
+                                       k,
+                                       v,
+                                       block_length=128,
+                                       make_image_summary=False,
+                                       dropout_rate=0.,
+                                       heads_share_relative_embedding=False,
+                                       add_relative_to_values=False,
+                                       name=None):
   """Masked local 1d attention with relative positions.
 
   The sequence is divided into blocks of length block_size.
@@ -2131,21 +2229,38 @@ def masked_rel_local_attention_1d(q,
     block_length: an integer
     make_image_summary: a boolean, whether to make an attention image summary.
     dropout_rate: Dropout rate for attention dropout
-    share_rel_embed: Boolean for sharing relative embeddings
+    heads_share_relative_embedding: a boolean for sharing relative embeddings.
+    add_relative_to_values: a boolean for whether to add relative component to
+        values.
     name: an optional string
 
   Returns:
     a Tensor of shape [batch, heads, length, depth_v]
+
+  Raises:
+    ValueError: wwhen the name for the variable scope is not passed.
   """
+  if not name:
+    raise ValueError("Name must be assigned since reuse for variable scope is "
+                     "set to tf.AUTO_REUSE, in order to reuse relative "
+                     "embeddings of keys and values.")
+
+  # Reuse flag is set to auto_reuse to reuse relative embeddings of keys and
+  # values across blocks (first and tail blocks).
   with tf.variable_scope(
-      name, default_name="local_attention_1d", values=[q, k, v]):
+      name, default_name="masked_relative_local_attention_1d",
+      values=[q, k, v], reuse=tf.AUTO_REUSE):
 
     default_block_length = block_length
     batch = common_layers.shape_list(q)[0]
     heads = common_layers.shape_list(q)[1]
     length = common_layers.shape_list(q)[2]
     # If (length < 2 * block_length), then we use only one block.
-    block_length = length if length < block_length * 2 else block_length
+    if isinstance(length, int) and isinstance(block_length, int):
+      block_length = length if length < block_length * 2 else block_length
+    else:
+      block_length = tf.where(
+          tf.less(length, block_length * 2), length, block_length)
     depth_k = common_layers.shape_list(k)[3]
     depth_v = common_layers.shape_list(v)[3]
     original_length = length
@@ -2163,35 +2278,18 @@ def masked_rel_local_attention_1d(q,
     first_v = tf.slice(v, [0, 0, 0, 0], [-1, -1, block_length, -1])
     # Relative embeddings will be used later as well.
     # TODO(avaswani,annahuang): check why 2*bl was breaking for music
-    # We only multiply with the needed embeddings as we slice them out.
+    # Needs to be known at static shape inference time, hence cannot be
+    # 2 * block_length.
     rel_embed_length = 4 * default_block_length
-    # Relative embeddings can be shared or unshared
+    # We only multiply with the needed embeddings as we slice them out.
+    first_rel_embeddings = get_relative_embeddings_left(
+        rel_embed_length, block_length, depth_k, heads,
+        heads_share_relative_embedding, "relative_embeddings")
+    first_rel_logits = matmul_with_relative_keys(
+        first_q, first_rel_embeddings, heads_share_relative_embedding)
     first_logits = tf.matmul(first_q, first_k, transpose_b=True)
-    initializer_stddev = depth_k**-0.5
-    if share_rel_embed:
-      relative_embeddings = (
-          tf.get_variable(name="relative_embeddings",
-                          shape=(rel_embed_length, depth_k),
-                          initializer=tf.random_normal_initializer(
-                              stddev=initializer_stddev)))
-      masked_relative_embeddings = tf.slice(
-          relative_embeddings,
-          [rel_embed_length - block_length, 0], [-1, -1])
-      first_relative_logits = tf.einsum(
-          "bhld,md->bhlm", first_q, masked_relative_embeddings)
-    else:
-      relative_embeddings = (
-          tf.get_variable(name="relative_embeddings",
-                          shape=(heads, rel_embed_length, depth_k),
-                          initializer=tf.random_normal_initializer(
-                              stddev=initializer_stddev)))
-      masked_relative_embeddings = tf.slice(
-          relative_embeddings,
-          [0, rel_embed_length - block_length, 0], [-1, -1, -1])
-      first_relative_logits = tf.einsum(
-          "bhld,hmd->bhlm", first_q, masked_relative_embeddings)
     first_logits += (
-        _relative_position_to_absolute_position_masked(first_relative_logits))
+        _relative_position_to_absolute_position_masked(first_rel_logits))
     # adding a mask
     first_logits += (
         common_layers.cast_like(attention_bias_lower_triangle(block_length),
@@ -2232,21 +2330,13 @@ def _reshape_for_relative(x):
     rel_tail_q = _reshape_for_relative(tail_q)
     rel_k = _reshape_for_relative(local_k)
     rel_v = _reshape_for_relative(local_v)
+    rel_embeddings = get_relative_embeddings_left(
+        rel_embed_length, 2 * block_length, depth_k, heads,
+        heads_share_relative_embedding, "relative_embeddings")
+    rel_logits = matmul_with_relative_keys(
+        rel_tail_q, rel_embeddings, heads_share_relative_embedding)
     # Computing relative logits separately for the masked and unmasked parts
     # because the reshaping logic is different for both
-    if share_rel_embed:
-      used_relative_embeddings = tf.slice(
-          relative_embeddings,
-          [rel_embed_length - 2*block_length, 0], [-1, -1])
-      rel_logits = tf.einsum(
-          "bhld,md->bhlm", rel_tail_q, used_relative_embeddings)
-    else:
-      used_relative_embeddings = tf.slice(
-          relative_embeddings,
-          [0, rel_embed_length - 2*block_length, 0], [-1, -1, -1])
-      rel_logits = tf.einsum(
-          "bhld,hmd->bhlm", rel_tail_q, used_relative_embeddings)
-
     masked_rel_logits = tf.slice(rel_logits, [0, 0, 0, block_length],
                                  [-1, -1, -1, -1])
     masked_rel_logits = _relative_position_to_absolute_position_masked(
@@ -2266,11 +2356,50 @@ def _reshape_for_relative(x):
     mask = (1.0 - good_part) * -1e9
     mask = common_layers.cast_like(mask, all_logits)
     all_logits += tf.reshape(mask, [1, 1, block_length, local_length])
-    attention = tf.nn.softmax(all_logits)
-    attention = common_layers.dropout_with_broadcast_dims(
-        attention, 1.0 - dropout_rate,
+    weights = tf.nn.softmax(all_logits, name="attention_weights")
+    # [batch (* num_blocks), heads, query_length (=block_length),
+    # key_length (=2*block_length)]
+    weights = common_layers.dropout_with_broadcast_dims(
+        weights, 1.0 - dropout_rate,
         broadcast_dims=None)
-    output = tf.matmul(attention, rel_v)
+
+    output = tf.matmul(weights, rel_v)
+    if add_relative_to_values:
+      # Adds the contribution of the weighted relative embeddings to the values.
+      weights_for_unmasked, weights_for_masked = (
+          tf.split(weights, 2, axis=3))
+      rel_weights_unmasked = _absolute_position_to_relative_position_unmasked(
+          weights_for_unmasked)
+      rel_weights_masked = _absolute_position_to_relative_position_masked(
+          weights_for_masked)
+
+      value_rel_embeddings_unmasked = get_relative_embeddings_left(
+          rel_embed_length, 2 * block_length, depth_v,
+          heads, heads_share_relative_embedding,
+          "value_relative_embeddings")
+      # The unmasked part starts with index -1 as opposed 0 has take uptil last.
+      if heads_share_relative_embedding:
+        value_rel_embeddings_unmasked = value_rel_embeddings_unmasked[:-1, :]
+      else:
+        value_rel_embeddings_unmasked = value_rel_embeddings_unmasked[:, :-1, :]
+      value_rel_embeddings_masked = get_relative_embeddings_left(
+          rel_embed_length, block_length, depth_v,
+          heads, heads_share_relative_embedding,
+          "value_relative_embeddings")
+
+      # [batch (*num_blocks), heads, query length, key length]
+      rel_weights = tf.concat(
+          [rel_weights_unmasked, rel_weights_masked], axis=3)
+      if heads_share_relative_embedding:
+        value_rel_embeddings_concat_axis = 0
+      else:
+        value_rel_embeddings_concat_axis = 1
+      value_rel_embeddings = tf.concat(
+          [value_rel_embeddings_unmasked, value_rel_embeddings_masked],
+          axis=value_rel_embeddings_concat_axis)
+      output_rel = matmul_with_relative_values(
+          rel_weights, value_rel_embeddings, heads_share_relative_embedding)
+      output += output_rel
 
     # bring to [batch, heads, num_blocks-1, block_length, depth]
     output = tf.reshape(output,
@@ -2282,10 +2411,25 @@ def _reshape_for_relative(x):
     output = tf.concat([first_output, output], axis=2)
     output = tf.slice(output, [0, 0, 0, 0], [-1, -1, original_length, -1])
     output = tf.reshape(output, [batch, heads, original_length, depth_v])
-
     return output
 
 
+def matmul_with_relative_values(x, y, heads_share_relative_embedding):
+  if heads_share_relative_embedding:
+    ret = tf.einsum("bhlm,md->bhld", x, y)
+  else:
+    ret = tf.einsum("bhlm,hmd->bhld", x, y)
+  return ret
+
+
+def matmul_with_relative_keys(x, y, heads_share_relative_embedding):
+  if heads_share_relative_embedding:
+    ret = tf.einsum("bhld,md->bhlm", x, y)
+  else:
+    ret = tf.einsum("bhld,hmd->bhlm", x, y)
+  return ret
+
+
 def local_attention_1d(q, k, v, block_length=128, filter_width=100, name=None):
   """Strided block local self-attention.
 
@@ -3161,10 +3305,11 @@ def multihead_attention(query_antecedent,
                         output_depth,
                         num_heads,
                         dropout_rate,
-                        shared_rel=False,
+                        attention_type="dot_product",
                         max_relative_position=None,
+                        heads_share_relative_embedding=False,
+                        add_relative_to_values=False,
                         image_shapes=None,
-                        attention_type="dot_product",
                         block_length=128,
                         block_width=128,
                         q_filter_width=1,
@@ -3178,7 +3323,6 @@ def multihead_attention(query_antecedent,
                         save_weights_to=None,
                         make_image_summary=True,
                         dropout_broadcast_dims=None,
-                        max_length=None,
                         vars_3d=False,
                         **kwargs):
   """Multihead scaled-dot-product attention with input/output transformations.
@@ -3189,19 +3333,21 @@ def multihead_attention(query_antecedent,
     bias: bias Tensor (see attention_bias())
     total_key_depth: an integer
     total_value_depth: an integer
-    output_depth: an integer
+    output_depth: an integerg
     num_heads: an integer dividing total_key_depth and total_value_depth
     dropout_rate: a floating point number
-    shared_rel: boolean to share relative embeddings
+    attention_type: a string, either "dot_product", "dot_product_relative",
+                    "local_mask_right", "local_unmasked", "masked_dilated_1d",
+                    "unmasked_dilated_1d", graph, or any attention function
+                    with the signature (query, key, value, **kwargs)
     max_relative_position: Maximum distance between inputs to generate
                            unique relation embeddings for. Only relevant
                            when using "dot_product_relative" attention.
+    heads_share_relative_embedding: boolean to share relative embeddings
+    add_relative_to_values: a boolean for whether to add relative component to
+                            values.
     image_shapes: optional tuple of integer scalars.
                   see comments for attention_image_summary()
-    attention_type: a string, either "dot_product", "dot_product_relative",
-                    "local_mask_right", "local_unmasked", "masked_dilated_1d",
-                    "unmasked_dilated_1d", graph, or any attention function
-                    with the signature (query, key, value, **kwargs)
     block_length: an integer - relevant for "local_mask_right"
     block_width: an integer - relevant for "local_unmasked"
     q_filter_width: An integer specifying how wide you want the query to be.
@@ -3228,7 +3374,6 @@ def multihead_attention(query_antecedent,
     dropout_broadcast_dims:  an optional list of integers less than 4
       specifying in which dimensions to broadcast the dropout decisions.
       saves memory.
-    max_length: an integer - needed by relative attention
     vars_3d: use 3-dimensional variables for input/output transformations
     **kwargs (dict): Parameters for the attention function
 
@@ -3343,30 +3488,40 @@ def multihead_attention(query_antecedent,
           k,
           v,
           bias,
-          max_length,
+          max_relative_position,
           dropout_rate,
           image_shapes,
           make_image_summary=make_image_summary,
-          dropout_broadcast_dims=dropout_broadcast_dims)
+          dropout_broadcast_dims=dropout_broadcast_dims,
+          heads_share_relative_embedding=heads_share_relative_embedding,
+          add_relative_to_values=add_relative_to_values)
     elif attention_type == "dot_product_relative_v2":
       x = dot_product_self_attention_relative_v2(
           q,
           k,
           v,
           bias,
-          max_length,
+          max_relative_position,
           dropout_rate,
           image_shapes,
           make_image_summary=make_image_summary,
-          dropout_broadcast_dims=dropout_broadcast_dims)
+          dropout_broadcast_dims=dropout_broadcast_dims,
+          heads_share_relative_embedding=heads_share_relative_embedding,
+          add_relative_to_values=add_relative_to_values)
     elif attention_type == "local_within_block_mask_right":
       x = masked_within_block_local_attention_1d(
           q, k, v, block_length=block_length)
-    elif attention_type == "rel_local_mask_right":
-      x = masked_rel_local_attention_1d(q, k, v, block_length=block_length,
-                                        make_image_summary=make_image_summary,
-                                        dropout_rate=dropout_rate,
-                                        share_rel_embed=shared_rel)
+    elif attention_type == "local_relative_mask_right":
+      x = masked_relative_local_attention_1d(
+          q,
+          k,
+          v,
+          block_length=block_length,
+          make_image_summary=make_image_summary,
+          dropout_rate=dropout_rate,
+          heads_share_relative_embedding=heads_share_relative_embedding,
+          add_relative_to_values=add_relative_to_values,
+          name="masked_relative_local_attention_1d")
     elif attention_type == "local_mask_right":
       x = masked_local_attention_1d(
           q,
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 926337105..3015535e7 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -471,6 +471,57 @@ def testDotProductAttentionRelative(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testRelativeAttentionV2(self):
+    # (batch, heads, length, depth)
+    x = np.random.rand(5, 4, 16, 7)
+    y = np.random.rand(5, 4, 16, 7)
+    max_relative_position = 3
+    a = common_attention.dot_product_self_attention_relative_v2(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        None,
+        max_relative_position=max_relative_position,
+        heads_share_relative_embedding=False)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
+    self.assertEqual(res.shape, (5, 4, 16, 7))
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testRelativeAttentionV2SharedRel(self):
+    # (batch, heads, length, depth)
+    x = np.random.rand(5, 4, 16, 7)
+    y = np.random.rand(5, 4, 16, 7)
+    max_relative_position = 3
+    a = common_attention.dot_product_self_attention_relative_v2(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        None,
+        max_relative_position=max_relative_position,
+        heads_share_relative_embedding=True)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
+    self.assertEqual(res.shape, (5, 4, 16, 7))
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testRelativeAttentionV2MaxRelativeLargerThanLength(self):
+    # (batch, heads, length, depth)
+    x = np.random.rand(5, 4, 3, 7)
+    y = np.random.rand(5, 4, 3, 7)
+    max_relative_position = 16
+    a = common_attention.dot_product_self_attention_relative_v2(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        None,
+        max_relative_position=max_relative_position,
+        heads_share_relative_embedding=False)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
+    self.assertEqual(res.shape, (5, 4, 3, 7))
+
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDotProductUnMaskedAttentionRelativeV2(self):
     x = np.random.rand(5, 7, 12, 32)
@@ -486,6 +537,126 @@ def testDotProductUnMaskedAttentionRelativeV2(self):
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testRelativeAttentionV2Unmasked(self):
+    # (batch, heads, length, depth)
+    x = np.random.rand(5, 4, 16, 7)
+    y = np.random.rand(5, 4, 16, 7)
+    max_relative_position = 3
+    a = common_attention.dot_product_unmasked_self_attention_relative_v2(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        None,
+        max_relative_position=max_relative_position,
+        heads_share_relative_embedding=False)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
+    self.assertEqual(res.shape, (5, 4, 16, 7))
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testRelativeAttentionV2UnmaskedSharedRel(self):
+    # (batch, heads, length, depth)
+    x = np.random.rand(5, 4, 16, 7)
+    y = np.random.rand(5, 4, 16, 7)
+    max_relative_position = 3
+    a = common_attention.dot_product_unmasked_self_attention_relative_v2(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        None,
+        max_relative_position=max_relative_position,
+        heads_share_relative_embedding=True)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
+    self.assertEqual(res.shape, (5, 4, 16, 7))
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testRelativeAttentionV2UnmaskedRelativeLargerThanLength(self):
+    # (batch, heads, length, depth)
+    x = np.random.rand(5, 4, 3, 7)
+    y = np.random.rand(5, 4, 3, 7)
+    max_relative_position = 16
+    a = common_attention.dot_product_unmasked_self_attention_relative_v2(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        None,
+        max_relative_position=max_relative_position,
+        heads_share_relative_embedding=False)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
+    self.assertEqual(res.shape, (5, 4, 3, 7))
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testMaskedRelativeLocalAttentionV2(self):
+    # (batch, heads, length, depth)
+    x = np.random.rand(5, 4, 16, 7)
+    y = np.random.rand(5, 4, 16, 7)
+    block_length = 3
+    a = common_attention.masked_relative_local_attention_1d(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        block_length=block_length,
+        heads_share_relative_embedding=True,
+        add_relative_to_values=False,
+        name="masked_relative_local_attention_1d")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
+    self.assertEqual(res.shape, (5, 4, 16, 7))
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testMaskedRelativeLocalAttentionV2AddRelativeValues(self):
+    # (batch, heads, length, depth)
+    x = np.random.rand(5, 4, 16, 7)
+    y = np.random.rand(5, 4, 16, 7)
+    block_length = 3
+    a = common_attention.masked_relative_local_attention_1d(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        block_length=block_length,
+        heads_share_relative_embedding=True,
+        add_relative_to_values=False,
+        name="masked_relative_local_attention_1d")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
+    self.assertEqual(res.shape, (5, 4, 16, 7))
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testMaskedRelativeLocalAttentionV2SeqShorterThanBlockLength(self):
+    # (batch, heads, length, depth)
+    x = np.random.rand(5, 7, 2, 7)
+    y = np.random.rand(5, 7, 2, 7)
+    block_length = 3
+    a = common_attention.masked_relative_local_attention_1d(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        block_length=block_length,
+        heads_share_relative_embedding=True,
+        name="masked_relative_local_attention_1d")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
+    self.assertEqual(res.shape, (5, 7, 2, 7))
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testMaskedRelativeLocalAttentionV2SeqShorterThanTwiceBlockLength(self):
+    # (batch, heads, length, depth)
+    x = np.random.rand(5, 7, 5, 7)
+    y = np.random.rand(5, 7, 5, 7)
+    block_length = 3
+    a = common_attention.masked_relative_local_attention_1d(
+        tf.constant(x, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        tf.constant(y, dtype=tf.float32),
+        block_length=block_length,
+        heads_share_relative_embedding=True,
+        name="masked_relative_local_attention_1d")
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
+    self.assertEqual(res.shape, (5, 7, 5, 7))
+
   def testBiasBatchCoordinates(self):
     """Testing the batch coordinates mask."""
     q = tf.constant([0, 0, 1, 1, 1, 1, 2, 2, 2], dtype=tf.int32)
@@ -536,3 +707,4 @@ def testBiasFuture(self):
 
 if __name__ == "__main__":
   tf.test.main()
+
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 4d408880d..0d5c5895f 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -274,7 +274,14 @@ def basic_params1():
       multiproblem_reweight_label_loss=False,
       # How much weight the targets in classification problems receive. Inputs
       # receive 1 minus this weight.
-      multiproblem_label_weight=0.5
+      multiproblem_label_weight=0.5,
+      # Hyperparameters for relative attention.
+      # The maximum relative positional distance to learn an embedding for.
+      max_relative_position=0,
+      # If heads share the same relative embedding.
+      heads_share_relative_embedding=False,
+      # If relative embedding terms are added to values too.
+      add_relative_to_values=False,
   )
 
 
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 8ac1a16a7..3175eb1ac 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -219,7 +219,8 @@ def _greedy_infer(self, features, decode_length, use_tpu=False):
       NotImplementedError: If there are multiple data shards.
     """
     # For real-valued modalities use the slow decode path for now.
-    if self._target_modality_is_real:
+    if (self._target_modality_is_real or
+        self._hparams.self_attention_type != "dot_product"):
       return  super(Transformer, self)._greedy_infer(features, decode_length)
     with tf.variable_scope(self.name):
       return (self._fast_decode_tpu(features, decode_length) if use_tpu else
@@ -1221,8 +1222,11 @@ def transformer_encoder(encoder_input,
               hparams.num_heads,
               hparams.attention_dropout,
               attention_type=hparams.self_attention_type,
-              save_weights_to=save_weights_to,
               max_relative_position=hparams.max_relative_position,
+              heads_share_relative_embedding=(
+                  hparams.heads_share_relative_embedding),
+              add_relative_to_values=hparams.add_relative_to_values,
+              save_weights_to=save_weights_to,
               make_image_summary=make_image_summary,
               dropout_broadcast_dims=attention_dropout_broadcast_dims,
               max_length=hparams.get("max_length"),
@@ -1304,8 +1308,11 @@ def transformer_decoder(decoder_input,
               hparams.num_heads,
               hparams.attention_dropout,
               attention_type=hparams.self_attention_type,
-              save_weights_to=save_weights_to,
               max_relative_position=hparams.max_relative_position,
+              heads_share_relative_embedding=(
+                  hparams.heads_share_relative_embedding),
+              add_relative_to_values=hparams.add_relative_to_values,
+              save_weights_to=save_weights_to,
               cache=layer_cache,
               make_image_summary=make_image_summary,
               dropout_broadcast_dims=attention_dropout_broadcast_dims,
@@ -1324,6 +1331,10 @@ def transformer_decoder(decoder_input,
                 hparams.hidden_size,
                 hparams.num_heads,
                 hparams.attention_dropout,
+                max_relative_position=hparams.max_relative_position,
+                heads_share_relative_embedding=(
+                    hparams.heads_share_relative_embedding),
+                add_relative_to_values=hparams.add_relative_to_values,
                 save_weights_to=save_weights_to,
                 cache=layer_cache,
                 make_image_summary=make_image_summary,
@@ -1520,7 +1531,6 @@ def transformer_base_v1():
   hparams.add_hparam("causal_decoder_self_attention", True)
   hparams.add_hparam("use_pad_remover", True)
   hparams.add_hparam("self_attention_type", "dot_product")
-  hparams.add_hparam("max_relative_position", 0)
   hparams.add_hparam("conv_first_kernel", 3)
   hparams.add_hparam("attention_variables_3d", False)
   hparams.add_hparam("use_target_space_embedding", True)

From 2917c2db5f72847c1668ccb76ca6d5d8c33209d5 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Wed, 5 Sep 2018 22:51:50 -0700
Subject: [PATCH 0752/2720] Fix attention mask for local attention.

PiperOrigin-RevId: 211754966
---
 tensor2tensor/mesh_tensorflow/mtf_layers.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers.py b/tensor2tensor/mesh_tensorflow/mtf_layers.py
index b7dd41084..26f2cc31c 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers.py
@@ -575,6 +575,10 @@ def attention_bias_local_block(mesh, block_length, memory_length,
                                dtype=tf.int32):
   """Bias for attention for local blocks where attention to right is disallowed.
 
+  Create the bias matrix by using two separate masks, one for the memory part
+  which doesn't overlap with the query and second which interacts with the query
+  and should be disallowed to look to the right of the current query position.
+
   Args:
     mesh: a MeshTensorflow object
     block_length: a mtf.Dimension
@@ -582,10 +586,14 @@ def attention_bias_local_block(mesh, block_length, memory_length,
     dtype: a tf.dtype
 
   Returns:
-    a mtf.Tensor with shape [rows, cols]
+    a mtf.Tensor with shape [block_length, memory_length]
   """
+  memory_length = mtf.Dimension(memory_length.name, block_length.size)
+  memory_mask = mtf.zeros(mesh, [block_length, memory_length], dtype=dtype)
+
   mask = mtf.cast(mtf.less(mtf.range(mesh, block_length, dtype=dtype),
                            mtf.range(mesh, memory_length, dtype=dtype)),
                   dtype=dtype)
-  mask = mtf.cast(mask, dtype=tf.float32)  * -1e9
+  mask = mtf.cast(mtf.concat([memory_mask, mask], memory_length.name),
+                  dtype=tf.float32)  * -1e9
   return mask

From a3c556d1d42b932e105c0d156c2a8585650c2115 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 6 Sep 2018 10:50:21 -0700
Subject: [PATCH 0753/2720] Add gif summary.

PiperOrigin-RevId: 211833011
---
 .travis.yml                               |   1 +
 tensor2tensor/layers/common_video.py      | 137 ++++++++++++++++++++++
 tensor2tensor/layers/common_video_test.py |  19 +++
 3 files changed, 157 insertions(+)

diff --git a/.travis.yml b/.travis.yml
index 941e6b25e..405d60602 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -75,6 +75,7 @@ script:
     --ignore=tensor2tensor/rl/trainer_model_based_stochastic_test.py
     --ignore=tensor2tensor/rl/trainer_model_based_sv2p_test.py
     --ignore=tensor2tensor/models/research
+    --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary
   - pytest tensor2tensor/utils/registry_test.py
   - pytest tensor2tensor/utils/trainer_lib_test.py
   - pytest tensor2tensor/visualization/visualization_test.py
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 1876e9ce3..156e332d5 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -17,9 +17,13 @@
 from __future__ import division
 from __future__ import print_function
 
+import numpy as np
+
 from tensor2tensor.layers import common_layers
 import tensorflow as tf
 
+from tensorflow.python.ops import summary_op_util
+
 tfl = tf.layers
 tfcl = tf.contrib.layers
 
@@ -274,6 +278,139 @@ def tile_and_concat(image, latent, concat_latent=True):
   return tf.concat([image, latent], axis=-1)
 
 
+def _encode_gif(images, fps):
+  """Encodes numpy images into gif string.
+
+  Args:
+    images: A 5-D `uint8` `np.array` (or a list of 4-D images) of shape
+      `[batch_size, time, height, width, channels]` where `channels` is 1 or 3.
+    fps: frames per second of the animation
+
+  Returns:
+    The encoded gif string.
+
+  Raises:
+    IOError: If the ffmpeg command returns an error.
+  """
+  from subprocess import Popen, PIPE  # pylint: disable=g-import-not-at-top,g-multiple-import
+  ffmpeg = "ffmpeg"
+  height, width, channels = images[0].shape
+  cmd = [
+      ffmpeg, "-y",
+      "-f", "rawvideo",
+      "-vcodec", "rawvideo",
+      "-r", "%.02f" % fps,
+      "-s", "%dx%d" % (width, height),
+      "-pix_fmt", {1: "gray", 3: "rgb24"}[channels],
+      "-i", "-",
+      "-filter_complex", "[0:v]split[x][z];[z]palettegen[y];[x][y]paletteuse",
+      "-r", "%.02f" % fps,
+      "-f", "gif",
+      "-"
+  ]
+  proc = Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
+  for image in images:
+    proc.stdin.write(image.tostring())
+  out, err = proc.communicate()
+  if proc.returncode:
+    err = "\n".join([" ".join(cmd), err.decode("utf8")])
+    raise IOError(err)
+  del proc
+  return out
+
+
+def py_gif_summary(tag, images, max_outputs, fps):
+  """Outputs a `Summary` protocol buffer with gif animations.
+
+  Args:
+    tag: Name of the summary.
+    images: A 5-D `uint8` `np.array` of shape `[batch_size, time, height, width,
+      channels]` where `channels` is 1 or 3.
+    max_outputs: Max number of batch elements to generate gifs for.
+    fps: frames per second of the animation
+
+  Returns:
+    The serialized `Summary` protocol buffer.
+
+  Raises:
+    ValueError: If `images` is not a 5-D `uint8` array with 1 or 3 channels.
+  """
+  images = np.asarray(images)
+  if images.dtype != np.uint8:
+    raise ValueError("Tensor must have dtype uint8 for gif summary.")
+  if images.ndim != 5:
+    raise ValueError("Tensor must be 5-D for gif summary.")
+  batch_size, _, height, width, channels = images.shape
+  if channels not in (1, 3):
+    raise ValueError("Tensors must have 1 or 3 channels for gif summary.")
+
+  summ = tf.Summary()
+  num_outputs = min(batch_size, max_outputs)
+  for i in range(num_outputs):
+    image_summ = tf.Summary.Image()
+    image_summ.height = height
+    image_summ.width = width
+    image_summ.colorspace = channels  # 1: grayscale, 3: RGB
+    try:
+      image_summ.encoded_image_string = _encode_gif(images[i], fps)
+    except (IOError, OSError) as e:
+      tf.logging.warning(
+          "Unable to encode images to a gif string because either ffmpeg is "
+          "not installed or ffmpeg returned an error: %s. Falling back to an "
+          "image summary of the first frame in the sequence.", e)
+      try:
+        from PIL import Image  # pylint: disable=g-import-not-at-top
+        import io  # pylint: disable=g-import-not-at-top
+        with io.BytesIO() as output:
+          Image.fromarray(images[i][0]).save(output, "PNG")
+          image_summ.encoded_image_string = output.getvalue()
+      except ImportError as e:
+        tf.logging.warning(
+            "Gif summaries requires ffmpeg or PIL to be installed: %s", e)
+        image_summ.encoded_image_string = ""
+    if num_outputs == 1:
+      summ_tag = "{}/gif".format(tag)
+    else:
+      summ_tag = "{}/gif/{}".format(tag, i)
+    summ.value.add(tag=summ_tag, image=image_summ)
+  summ_str = summ.SerializeToString()
+  return summ_str
+
+
+def gif_summary(name, tensor, max_outputs=3, fps=10, collections=None,
+                family=None):
+  """Outputs a `Summary` protocol buffer with gif animations.
+
+  Args:
+    name: Name of the summary.
+    tensor: A 5-D `uint8` `Tensor` of shape `[batch_size, time, height, width,
+      channels]` where `channels` is 1 or 3.
+    max_outputs: Max number of batch elements to generate gifs for.
+    fps: frames per second of the animation
+    collections: Optional list of tf.GraphKeys.  The collections to add the
+      summary to.  Defaults to [tf.GraphKeys.SUMMARIES]
+    family: Optional; if provided, used as the prefix of the summary tag name,
+      which controls the tab name used for display on Tensorboard.
+
+  Returns:
+    A scalar `Tensor` of type `string`. The serialized `Summary` protocol
+    buffer.
+  """
+  tensor = tf.convert_to_tensor(tensor)
+  if summary_op_util.skip_summary():
+    return tf.constant("")
+  with summary_op_util.summary_scope(
+      name, family, values=[tensor]) as (tag, scope):
+    val = tf.py_func(
+        py_gif_summary,
+        [tag, tensor, max_outputs, fps],
+        tf.string,
+        stateful=False,
+        name=scope)
+    summary_op_util.collect(val, collections, [tf.GraphKeys.SUMMARIES])
+  return val
+
+
 
 
 def tinyify(array, tiny_mode):
diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index 6f88c5fb6..06e650797 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -99,6 +99,25 @@ def testDynamicTileAndConcat(self):
          [90, 90, 90, 90],
          [100, 100, 100, 100]])
 
+  def testGifSummary(self):
+    for c in (1, 3):
+      images_shape = (1, 12, 48, 64, c)  # batch, time, height, width, channels
+      images = np.random.randint(256, size=images_shape).astype(np.uint8)
+
+      with self.test_session():
+        summary = common_video.gif_summary(
+            "gif", tf.convert_to_tensor(images), fps=10)
+        summary_string = summary.eval()
+
+      summary = tf.Summary()
+      summary.ParseFromString(summary_string)
+
+      self.assertEqual(1, len(summary.value))
+      self.assertTrue(summary.value[0].HasField("image"))
+      encoded = summary.value[0].image.encoded_image_string
+
+      self.assertEqual(encoded, common_video._encode_gif(images[0], fps=10))  # pylint: disable=protected-access
+
 
 if __name__ == "__main__":
   tf.test.main()

From f0797f66af131458db7c0849bd33e0f09d20b2d3 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Thu, 6 Sep 2018 11:55:11 -0700
Subject: [PATCH 0754/2720] Correct a small typo introduced in the last CL

PiperOrigin-RevId: 211844289
---
 tensor2tensor/models/research/transformer_vae.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 10be631f1..9b6d20737 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -734,7 +734,7 @@ def transformer_ae_small():
   hparams.add_hparam("noise_dev", 0.5)
   hparams.add_hparam("d_mix", 0.5)
   hparams.add_hparam("logit_normalization", True)
-  hparams.add_hparam("word_dropout", True)
+  hparams.add_hparam("word_dropout", 0.)
   # Bottleneck kinds supported: dense, vae, semhash, gumbel-softmax, dvq.
   hparams.add_hparam("bottleneck_kind", "semhash")
   hparams.add_hparam("num_blocks", 1)

From fcb0fa2e6493b4da6d8a3cbb81864d4ed0e74b32 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 6 Sep 2018 14:06:52 -0700
Subject: [PATCH 0755/2720] internal merge of PR #1043

PiperOrigin-RevId: 211866054
---
 tensor2tensor/rl/trainer_model_based.py | 34 +++++++++++++++----------
 1 file changed, 20 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 05195fcb4..365b82f41 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -185,10 +185,12 @@ def train_agent(problem_name, agent_model_dir,
   }):
     rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=epoch)
 
-def train_agent_real_env(problem_name, agent_model_dir,
-                event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
-                is_final_epoch=False):
-  """Train the PPO agent in the simulated environment."""
+
+def train_agent_real_env(
+    problem_name, agent_model_dir, event_dir, world_model_dir, epoch_data_dir,
+    hparams, epoch=0, is_final_epoch=False):
+  """Train the PPO agent in the real environment."""
+  del epoch, is_final_epoch
   gym_problem = registry.problem(problem_name)
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_params_names = ["epochs_num", "epoch_length",
@@ -204,7 +206,7 @@ def train_agent_real_env(problem_name, agent_model_dir,
   if ppo_epochs_num == 0:
     return
 
-  ppo_hparams.save_models_every_epochs = ppo_epochs_num #check this
+  ppo_hparams.save_models_every_epochs = ppo_epochs_num
 
   environment_spec = copy.copy(gym_problem.environment_spec)
 
@@ -215,8 +217,9 @@ def train_agent_real_env(problem_name, agent_model_dir,
       "output_dir": world_model_dir,
       "data_dir": epoch_data_dir,
   }):
-    # epoch = 0 is a hackish way to avoid skiping training
-    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=0)
+    # epoch = 10**20 is a hackish way to avoid skiping training
+    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=10**20)
+
 
 def evaluate_world_model(simulated_problem_name, problem_name, hparams,
                          world_model_dir, epoch_data_dir, tmp_dir):
@@ -476,7 +479,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       log("World model reward accuracy: %.4f", model_reward_accuracy)
 
     # Train PPO
-    log("Training PPO")
+    log("Training PPO in simulated environment.")
     ppo_event_dir = os.path.join(directories["world_model"],
                                  "ppo_summaries", str(epoch))
     ppo_model_dir = directories["ppo"]
@@ -487,10 +490,11 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
                 hparams, epoch=epoch, is_final_epoch=is_final_epoch)
 
     # Train PPO on real env (short)
-    train_agent_real_env(problem_name, ppo_model_dir,
-                ppo_event_dir, directories["world_model"], epoch_data_dir,
-                hparams, epoch=epoch, is_final_epoch=is_final_epoch)
-
+    log("Training PPO in real environment.")
+    train_agent_real_env(
+        problem_name, ppo_model_dir,
+        ppo_event_dir, directories["world_model"], epoch_data_dir,
+        hparams, epoch=epoch, is_final_epoch=is_final_epoch)
 
     # Collect data from the real environment.
     log("Generating real environment data")
@@ -676,7 +680,7 @@ def rl_modelrl_tiny():
   """Tiny set for testing."""
   return rl_modelrl_base_sampling().override_from_dict(
       tf.contrib.training.HParams(
-          epochs=2,
+          epochs=1,
           num_real_env_frames=128,
           simulated_env_generator_num_steps=64,
           model_train_steps=2,
@@ -684,7 +688,9 @@ def rl_modelrl_tiny():
           ppo_time_limit=5,
           ppo_epoch_length=5,
           ppo_num_agents=2,
-          real_ppo_epochs_num=4,
+          real_ppo_epochs_num=1,
+          real_ppo_epoch_length=5,
+          real_ppo_num_agents=2,
           generative_model_params="next_frame_tiny",
       ).values())
 

From 32ae3ee32abbfae4a7f1ffd2ef3098beb4d639db Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 6 Sep 2018 14:55:53 -0700
Subject: [PATCH 0756/2720] Move travis commands into scripts

PiperOrigin-RevId: 211874719
---
 .travis.yml                                   | 91 ++-----------------
 .../oss_scripts/oss_integration_test.sh       | 47 ++++++++++
 tensor2tensor/oss_scripts/oss_pip_install.sh  | 27 ++++++
 tensor2tensor/oss_scripts/oss_release.sh      | 31 +++++++
 tensor2tensor/oss_scripts/oss_tests.sh        | 91 +++++++++++++++++++
 5 files changed, 203 insertions(+), 84 deletions(-)
 create mode 100755 tensor2tensor/oss_scripts/oss_integration_test.sh
 create mode 100755 tensor2tensor/oss_scripts/oss_pip_install.sh
 create mode 100755 tensor2tensor/oss_scripts/oss_release.sh
 create mode 100755 tensor2tensor/oss_scripts/oss_tests.sh

diff --git a/.travis.yml b/.travis.yml
index 405d60602..628f13548 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -1,7 +1,8 @@
 sudo: required
 language: python
 git:
-  depth: 3
+  depth: 10
+  quiet: true
 services:
   - docker
 python:
@@ -32,91 +33,13 @@ before_install:
   - sudo apt-get update -qq
   - sudo apt-get install -qq libhdf5-dev
 install:
-  - if [[ "$TF_VERSION" == "tf-nightly"  ]];
-    then
-      pip install tf-nightly;
-    else
-      pip install -q "tensorflow==$TF_VERSION";
-    fi
-  # First ensure that the base dependencies are sufficient for a full import
-  - pip install -q .
-  - t2t-trainer --registry_help 2>&1 >/dev/null
-  # Then install the test dependencies
-  - pip install -q .[tests,allen]
-  # Make sure to install the atari extras for gym
-  - pip install "gym[atari]"
-  # Make sure we have the latest version of numpy - avoid problems we were
-  # seeing with Python 3
-  - pip install -q -U numpy
+  - ./oss_scripts/oss_pip_install.sh
 script:
-  # Check import
-  - python -c "from tensor2tensor.models import transformer; print(transformer.Transformer.__name__)"
-
-  # Run tests
-  # Ignores:
-  # Tested separately:
-  #   * registry_test
-  #   * trainer_lib_test
-  #   * visualization_test
-  #   * trainer_model_based_test
-  #   * allen_brain_test
-  #   * models/research
-  # algorithmic_math_test: flaky
-  # universal_transformer_test: requires new feature in tf.foldl (rm with TF 1.9)
-  - pytest
-    --ignore=tensor2tensor/utils/registry_test.py
-    --ignore=tensor2tensor/utils/trainer_lib_test.py
-    --ignore=tensor2tensor/visualization/visualization_test.py
-    --ignore=tensor2tensor/bin/t2t_trainer_test.py
-    --ignore=tensor2tensor/data_generators/algorithmic_math_test.py
-    --ignore=tensor2tensor/models/research/universal_transformer_test.py
-    --ignore=tensor2tensor/rl/trainer_model_based_test.py
-    --ignore=tensor2tensor/data_generators/allen_brain_test.py
-    --ignore=tensor2tensor/rl/trainer_model_based_stochastic_test.py
-    --ignore=tensor2tensor/rl/trainer_model_based_sv2p_test.py
-    --ignore=tensor2tensor/models/research
-    --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary
-  - pytest tensor2tensor/utils/registry_test.py
-  - pytest tensor2tensor/utils/trainer_lib_test.py
-  - pytest tensor2tensor/visualization/visualization_test.py
-  - pytest tensor2tensor/data_generators/allen_brain_test.py
-  - if [[ "$TF_VERSION" == "$TF_LATEST"  ]] || [[ "$TF_VERSION" == "tf-nightly"  ]];
-    then
-      pytest tensor2tensor/models/research --ignore=tensor2tensor/models/research/glow_test.py;
-    fi
-
-  # Run installed scripts
-  - t2t-datagen 2>&1 | grep translate && echo passed
-  - t2t-trainer --registry_help
-
-  # Test --t2t_usr_dir
-  - t2t-trainer --registry_help --t2t_usr_dir=./tensor2tensor/test_data/example_usr_dir 2>&1 | grep my_very_own_hparams && echo passed
-
-  # Run data generation, training, and decoding on a dummy problem
-  - t2t-datagen --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR
-  - t2t-trainer --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR --model=transformer --hparams_set=transformer_tiny --train_steps=5 --eval_steps=5 --output_dir=$T2T_TRAIN_DIR
-  - t2t-decoder --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR --model=transformer --hparams_set=transformer_tiny --output_dir=$T2T_TRAIN_DIR --decode_hparams='num_samples=10'
+  - ./oss_scripts/oss_tests.sh
+  - ./oss_scripts/oss_integration_test.sh
 
-  # Do some things only on Python 2 and the latest TF version
-  # Each should be in a separate block to get proper errors on Travis.
+  # Conditional commands should each be in a separate block to get proper
+  # errors on Travis.
   - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "tf-nightly"  ]]; then
         pylint -j 2 tensor2tensor;
     fi
-  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
-        pytest tensor2tensor/rl/trainer_model_based_test.py;
-    fi
-  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
-        jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/hello_t2t.ipynb;
-    fi
-  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
-        jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/t2t_problem.ipynb;
-    fi
-
-  # Export and query
-  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]; then
-        t2t-exporter --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR --model=transformer --hparams_set=transformer_tiny --output_dir=$T2T_TRAIN_DIR;
-        pip install tensorflow-serving-api;
-        docker run -d -p 8500:8500 --mount type=bind,source=$T2T_TRAIN_DIR/export/Servo,target=/models/my_model -e MODEL_NAME=my_model -t tensorflow/serving;
-        sleep 10;
-        t2t-query-server --problem=$T2T_PROBLEM --server=localhost:8500 --servable_name=my_model --data_dir=$T2T_DATA_DIR --inputs_once='1 0 1 0 1 0';
-    fi
diff --git a/tensor2tensor/oss_scripts/oss_integration_test.sh b/tensor2tensor/oss_scripts/oss_integration_test.sh
new file mode 100755
index 000000000..87038d9d3
--- /dev/null
+++ b/tensor2tensor/oss_scripts/oss_integration_test.sh
@@ -0,0 +1,47 @@
+#!/bin/bash
+
+set -v  # print commands as they're executed
+set -e  # fail and exit on any command erroring
+
+: "${TF_VERSION:?}"
+: "${TF_LATEST:?}"
+: "${T2T_DATA_DIR:?}"
+: "${T2T_TRAIN_DIR:?}"
+: "${T2T_PROBLEM:?}"
+
+# Test --t2t_usr_dir
+t2t-trainer --registry_help --t2t_usr_dir=./tensor2tensor/test_data/example_usr_dir 2>&1 | grep my_very_own_hparams && echo passed
+
+# Run data generation, training, and decoding on a dummy problem
+t2t-datagen --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR
+t2t-trainer --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR --model=transformer --hparams_set=transformer_tiny --train_steps=5 --eval_steps=5 --output_dir=$T2T_TRAIN_DIR
+t2t-decoder --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR --model=transformer --hparams_set=transformer_tiny --output_dir=$T2T_TRAIN_DIR --decode_hparams='num_samples=10'
+
+# Test serving
+if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]
+then
+  # Export for serving
+  t2t-exporter \
+      --problem=$T2T_PROBLEM \
+      --data_dir=$T2T_DATA_DIR \
+      --model=transformer \
+      --hparams_set=transformer_tiny \
+      --output_dir=$T2T_TRAIN_DIR
+
+  # Run model server
+  server_port=8500
+  model_name=my_model
+  docker run -d -p $server_port:$server_port \
+      --mount type=bind,source=$T2T_TRAIN_DIR/export/Servo,target=/models/$model_name \
+      -e MODEL_NAME=$model_name -t tensorflow/serving
+  sleep 10
+
+  # Query
+  pip install tensorflow-serving-api
+  t2t-query-server \
+      --server=localhost:$server_port \
+      --servable_name=$model_name \
+      --problem=$T2T_PROBLEM \
+      --data_dir=$T2T_DATA_DIR \
+      --inputs_once='1 0 1 0 1 0'
+fi
diff --git a/tensor2tensor/oss_scripts/oss_pip_install.sh b/tensor2tensor/oss_scripts/oss_pip_install.sh
new file mode 100755
index 000000000..12ffeb2c1
--- /dev/null
+++ b/tensor2tensor/oss_scripts/oss_pip_install.sh
@@ -0,0 +1,27 @@
+#!/bin/bash
+
+set -v  # print commands as they're executed
+set -e  # fail and exit on any command erroring
+
+: "${TF_VERSION:?}"
+
+if [[ "$TF_VERSION" == "tf-nightly"  ]]
+then
+  pip install tf-nightly;
+else
+  pip install -q "tensorflow==$TF_VERSION"
+fi
+
+# First ensure that the base dependencies are sufficient for a full import
+pip install -q -e .
+t2t-trainer --registry_help 2>&1 >/dev/null
+t2t-datagen 2>&1 | grep translate_ende 2>&1 >/dev/null && echo passed
+
+# Then install the test dependencies
+pip install -q -e .[tests,allen]
+# Make sure to install the atari extras for gym
+pip install "gym[atari]"
+
+# Make sure we have the latest version of numpy - avoid problems we were
+# seeing with Python 3
+pip install -q -U numpy
diff --git a/tensor2tensor/oss_scripts/oss_release.sh b/tensor2tensor/oss_scripts/oss_release.sh
new file mode 100755
index 000000000..6038bae21
--- /dev/null
+++ b/tensor2tensor/oss_scripts/oss_release.sh
@@ -0,0 +1,31 @@
+#!/bin/bash
+
+set -v  # print commands as they're executed
+set -e  # fail and exit on any command erroring
+
+GIT_COMMIT_ID=${1:-""}
+[[ -z $GIT_COMMIT_ID ]] && echo "Must provide a commit" && exit 1
+
+TMP_DIR=$(mktemp -d)
+pushd $TMP_DIR
+
+echo "Cloning tensor2tensor and checking out commit $GIT_COMMIT_ID"
+git clone https://github.com/tensorflow/tensor2tensor.git
+cd tensor2tensor
+git checkout $GIT_COMMIT_ID
+
+pip install wheel twine pyopenssl
+
+# Build the distribution
+echo "Building distribution"
+python setup.py sdist
+python setup.py bdist_wheel --universal
+
+# Publish to PyPI
+echo "Publishing to PyPI"
+twine upload dist/*
+
+# Cleanup
+rm -rf build/ dist/ tensor2tensor.egg-info/
+popd
+rm -rf $TMP_DIR
diff --git a/tensor2tensor/oss_scripts/oss_tests.sh b/tensor2tensor/oss_scripts/oss_tests.sh
new file mode 100755
index 000000000..f5a33a276
--- /dev/null
+++ b/tensor2tensor/oss_scripts/oss_tests.sh
@@ -0,0 +1,91 @@
+#!/bin/bash
+
+set -v  # print commands as they're executed
+
+# Instead of exiting on any failure with "set -e", we'll call set_status after
+# each command and exit $STATUS at the end.
+STATUS=0
+function set_status() {
+    local last_status=$?
+    if [[ $last_status -ne 0 ]]
+    then
+      echo "<<<<<<FAILED>>>>>> Exit code: $last_status"
+    fi
+    STATUS=$(($last_status || $STATUS))
+}
+
+# Check env vars set
+echo "${TF_VERSION:?}" && \
+echo "${TF_LATEST:?}" && \
+echo "${TRAVIS_PYTHON_VERSION:?}"
+set_status
+if [[ $STATUS -ne 0 ]]
+then
+  exit $STATUS
+fi
+
+# Check import
+python -c "from tensor2tensor.models import transformer; print(transformer.Transformer.__name__)"
+set_status
+
+# Run tests
+# Ignores:
+# Tested separately:
+#   * registry_test
+#   * trainer_lib_test
+#   * visualization_test
+#   * trainer_model_based_test
+#   * allen_brain_test
+#   * models/research
+# algorithmic_math_test: flaky
+pytest
+  --ignore=tensor2tensor/utils/registry_test.py \
+  --ignore=tensor2tensor/utils/trainer_lib_test.py \
+  --ignore=tensor2tensor/visualization/visualization_test.py \
+  --ignore=tensor2tensor/bin/t2t_trainer_test.py \
+  --ignore=tensor2tensor/data_generators/algorithmic_math_test.py \
+  --ignore=tensor2tensor/rl/trainer_model_based_test.py \
+  --ignore=tensor2tensor/data_generators/allen_brain_test.py \
+  --ignore=tensor2tensor/rl/trainer_model_based_stochastic_test.py \
+  --ignore=tensor2tensor/rl/trainer_model_based_sv2p_test.py \
+  --ignore=tensor2tensor/models/research \
+  --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary
+set_status
+
+pytest tensor2tensor/utils/registry_test.py
+set_status
+
+pytest tensor2tensor/utils/trainer_lib_test.py
+set_status
+
+pytest tensor2tensor/visualization/visualization_test.py
+set_status
+
+pytest tensor2tensor/data_generators/allen_brain_test.py
+set_status
+
+
+# Test models/research only against tf-nightly
+if [[ "$TF_VERSION" == "$TF_LATEST"  ]] || [[ "$TF_VERSION" == "tf-nightly"  ]]
+then
+  # Ignores:
+  # * Glow requires the CIFAR-10 dataset to be generated
+  pytest tensor2tensor/models/research --ignore=tensor2tensor/models/research/glow_test.py
+  set_status
+fi
+
+if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]
+then
+    pytest tensor2tensor/rl/trainer_model_based_test.py
+    set_status
+    jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/hello_t2t.ipynb
+    set_status
+    jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/t2t_problem.ipynb;
+    set_status
+fi
+
+# Test --t2t_usr_dir
+t2t-trainer --registry_help --t2t_usr_dir=./tensor2tensor/test_data/example_usr_dir 2>&1 | grep my_very_own_hparams && echo passed
+set_status
+
+exit $STATUS

From d1e34cad08b12e8f74b2ba92418de881bdaa26b0 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 6 Sep 2018 15:22:16 -0700
Subject: [PATCH 0757/2720] Travis fixes. Require pytest 3.8.0 for --deselect
 and s/assertLen/assertEquals

PiperOrigin-RevId: 211879271
---
 setup.py                              | 2 +-
 tensor2tensor/utils/t2t_model_test.py | 5 +++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/setup.py b/setup.py
index a61df9889..78d6de769 100644
--- a/setup.py
+++ b/setup.py
@@ -52,7 +52,7 @@
         'tensorflow': ['tensorflow>=1.9.0'],
         'tensorflow_gpu': ['tensorflow-gpu>=1.9.0'],
         'tests': [
-            'absl-py', 'pytest', 'mock', 'pylint', 'jupyter', 'gsutil',
+            'absl-py', 'pytest>=3.8.0', 'mock', 'pylint', 'jupyter', 'gsutil',
             'matplotlib',
             # Need atari extras for Travis tests, but because gym is already in
             # install_requires, pip skips the atari extras, so we instead do an
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index 66bf5eb8e..1a218eb0c 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -33,8 +33,9 @@ def testSummarizeLosses(self):
                 "extra": tf.random_normal([])}
       outputs = model._summarize_losses(losses)
       self.assertIsNone(outputs, None)
-      self.assertLen(tf.get_collection(tf.GraphKeys.SUMMARIES, scope="losses"),
-                     len(losses))
+      self.assertEquals(
+          len(tf.get_collection(tf.GraphKeys.SUMMARIES, scope="losses")),
+          len(losses))
 
 if __name__ == "__main__":
   tf.test.main()

From c77ef147ebc8c462a23a0ceb85399d104339cbf0 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 6 Sep 2018 16:13:05 -0700
Subject: [PATCH 0758/2720] Fixes to script locations in .travis.yml

PiperOrigin-RevId: 211887398
---
 .travis.yml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 628f13548..68cda574b 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -33,10 +33,10 @@ before_install:
   - sudo apt-get update -qq
   - sudo apt-get install -qq libhdf5-dev
 install:
-  - ./oss_scripts/oss_pip_install.sh
+  - ./tensor2tensor/oss_scripts/oss_pip_install.sh
 script:
-  - ./oss_scripts/oss_tests.sh
-  - ./oss_scripts/oss_integration_test.sh
+  - ./tensor2tensor/oss_scripts/oss_tests.sh
+  - ./tensor2tensor/oss_scripts/oss_integration_test.sh
 
   # Conditional commands should each be in a separate block to get proper
   # errors on Travis.

From 863741420a92d6cd1c910415a89d68c954586436 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 6 Sep 2018 16:54:40 -0700
Subject: [PATCH 0759/2720] More oss script fixes

PiperOrigin-RevId: 211893226
---
 tensor2tensor/oss_scripts/oss_tests.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/oss_scripts/oss_tests.sh b/tensor2tensor/oss_scripts/oss_tests.sh
index f5a33a276..de53f604a 100755
--- a/tensor2tensor/oss_scripts/oss_tests.sh
+++ b/tensor2tensor/oss_scripts/oss_tests.sh
@@ -38,7 +38,7 @@ set_status
 #   * allen_brain_test
 #   * models/research
 # algorithmic_math_test: flaky
-pytest
+pytest \
   --ignore=tensor2tensor/utils/registry_test.py \
   --ignore=tensor2tensor/utils/trainer_lib_test.py \
   --ignore=tensor2tensor/visualization/visualization_test.py \

From 4e353498d70168b8e159ce525affe03e6638f6cf Mon Sep 17 00:00:00 2001
From: Marcin Michalski <michalski@google.com>
Date: Fri, 7 Sep 2018 01:44:24 -0700
Subject: [PATCH 0760/2720] Update the tensor2tensor export to use the
 estimator Exporter library.

PiperOrigin-RevId: 211941142
---
 tensor2tensor/oss_scripts/oss_integration_test.sh |  2 +-
 tensor2tensor/serving/export.py                   | 12 ++++++++----
 tensor2tensor/utils/trainer_lib.py                |  5 -----
 3 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/oss_scripts/oss_integration_test.sh b/tensor2tensor/oss_scripts/oss_integration_test.sh
index 87038d9d3..9e82908a2 100755
--- a/tensor2tensor/oss_scripts/oss_integration_test.sh
+++ b/tensor2tensor/oss_scripts/oss_integration_test.sh
@@ -32,7 +32,7 @@ then
   server_port=8500
   model_name=my_model
   docker run -d -p $server_port:$server_port \
-      --mount type=bind,source=$T2T_TRAIN_DIR/export/Servo,target=/models/$model_name \
+      --mount type=bind,source=$T2T_TRAIN_DIR/export,target=/models/$model_name \
       -e MODEL_NAME=$model_name -t tensorflow/serving
   sleep 10
 
diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 5be45bc45..b5856d3ec 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -58,13 +58,17 @@ def main(_):
   estimator = create_estimator(run_config, hparams)
 
   problem = hparams.problem
-  strategy = trainer_lib.create_export_strategy(problem, hparams)
 
-  export_dir = os.path.join(ckpt_dir, "export", strategy.name)
-  strategy.export(
+  exporter = tf.estimator.FinalExporter(
+      "exporter", lambda: problem.serving_input_fn(hparams), as_text=True)
+
+  export_dir = os.path.join(ckpt_dir, "export")
+  exporter.export(
       estimator,
       export_dir,
-      checkpoint_path=tf.train.latest_checkpoint(ckpt_dir))
+      checkpoint_path=tf.train.latest_checkpoint(ckpt_dir),
+      eval_result=None,
+      is_the_final_export=True)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index b2dd55114..5f9c439a5 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -585,11 +585,6 @@ def experiment_fn(run_config, hparams):
   return experiment_fn
 
 
-def create_export_strategy(problem, hparams):
-  return tf.contrib.learn.make_export_strategy(
-      lambda: problem.serving_input_fn(hparams), as_text=True)
-
-
 def add_problem_hparams(hparams, problem_name):
   """Add problem hparams for the problems."""
   problem = registry.problem(problem_name)

From 250667158911fa3a1aab14764f47c8624bb75149 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 7 Sep 2018 10:39:22 -0700
Subject: [PATCH 0761/2720] refactoring the gif summary into writers so they
 can be used more for other purposes.

PiperOrigin-RevId: 211998264
---
 tensor2tensor/layers/common_video.py | 114 +++++++++++++++++++++------
 1 file changed, 89 insertions(+), 25 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 156e332d5..16a24d4e8 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -292,31 +292,9 @@ def _encode_gif(images, fps):
   Raises:
     IOError: If the ffmpeg command returns an error.
   """
-  from subprocess import Popen, PIPE  # pylint: disable=g-import-not-at-top,g-multiple-import
-  ffmpeg = "ffmpeg"
-  height, width, channels = images[0].shape
-  cmd = [
-      ffmpeg, "-y",
-      "-f", "rawvideo",
-      "-vcodec", "rawvideo",
-      "-r", "%.02f" % fps,
-      "-s", "%dx%d" % (width, height),
-      "-pix_fmt", {1: "gray", 3: "rgb24"}[channels],
-      "-i", "-",
-      "-filter_complex", "[0:v]split[x][z];[z]palettegen[y];[x][y]paletteuse",
-      "-r", "%.02f" % fps,
-      "-f", "gif",
-      "-"
-  ]
-  proc = Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
-  for image in images:
-    proc.stdin.write(image.tostring())
-  out, err = proc.communicate()
-  if proc.returncode:
-    err = "\n".join([" ".join(cmd), err.decode("utf8")])
-    raise IOError(err)
-  del proc
-  return out
+  writer = VideoWriter(fps)
+  writer.write_multi(images)
+  return writer.finish()
 
 
 def py_gif_summary(tag, images, max_outputs, fps):
@@ -521,3 +499,89 @@ def beta_schedule(schedule, global_step, final_beta, decay_start, decay_end):
           tf.greater(global_step, decay_end): lambda: final_beta},
       default=lambda: increased_value)
   return beta
+
+
+class VideoWriter(object):
+  """Helper class for writing videos."""
+
+  def __init__(self, fps, file_format="gif"):
+    self.fps = fps
+    self.file_format = file_format
+    self.proc = None
+
+  def __init_ffmpeg(self, image_shape):
+    """Initializes ffmpeg to write frames."""
+    from subprocess import Popen, PIPE  # pylint: disable=g-import-not-at-top,g-multiple-import
+    ffmpeg = "ffmpeg"
+    height, width, channels = image_shape
+    self.cmd = [
+        ffmpeg, "-y",
+        "-f", "rawvideo",
+        "-vcodec", "rawvideo",
+        "-r", "%.02f" % self.fps,
+        "-s", "%dx%d" % (width, height),
+        "-pix_fmt", {1: "gray", 3: "rgb24"}[channels],
+        "-i", "-",
+        "-filter_complex", "[0:v]split[x][z];[z]palettegen[y];[x][y]paletteuse",
+        "-r", "%.02f" % self.fps,
+        "-f", self.file_format,
+        "-"
+    ]
+    self.proc = Popen(self.cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
+
+  def write(self, frame):
+    if self.proc is None:
+      self.__init_ffmpeg(frame.shape)
+    self.proc.stdin.write(frame.tostring())
+
+  def write_multi(self, frames):
+    for frame in frames:
+      self.write(frame)
+
+  def finish(self):
+    if self.proc is None:
+      return None
+    out, err = self.proc.communicate()
+    if self.proc.returncode:
+      err = "\n".join([" ".join(self.cmd), err.decode("utf8")])
+      raise IOError(err)
+    del self.proc
+    self.proc = None
+    return out
+
+  def finish_to_file(self, path):
+    with tf.gfile.open(path) as f:
+      f.write(self.finish())
+
+  def __del__(self):
+    self.finish()
+
+
+class BatchVideoWriter(object):
+  """Helper class for writing videos in batch."""
+
+  def __init__(self, fps, file_format="gif"):
+    self.fps = fps
+    self.file_format = file_format
+    self.writers = None
+
+  def write(self, batch_frame):
+    if self.writers is None:
+      self.writers = [
+          VideoWriter(self.fps, self.file_format) for _ in batch_frame]
+    for i, frame in enumerate(batch_frame):
+      self.writers[i].write(frame)
+
+  def write_multi(self, batch_frames):
+    for batch_frame in batch_frames:
+      self.write(batch_frame)
+
+  def finish(self):
+    outs = [w.finish() for w in self.writers]
+    return outs
+
+  def finish_to_files(self, path_template):
+    for i, writer in enumerate(self.writers):
+      path = path_template.format(i)
+      writer.finish_to_file(path)
+

From 8bdaeb281c1deafce95c14fd80276d15724f2703 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 7 Sep 2018 13:24:14 -0700
Subject: [PATCH 0762/2720] RL bugfixes.

PiperOrigin-RevId: 212026070
---
 tensor2tensor/layers/common_layers.py         |  10 +-
 tensor2tensor/layers/modalities.py            |  18 ++--
 tensor2tensor/models/research/autoencoders.py |   3 +-
 .../models/video/basic_deterministic.py       |   5 +-
 .../video/basic_deterministic_params.py       |   4 +-
 tensor2tensor/rl/rl_trainer_lib.py            | 101 +++++++++---------
 tensor2tensor/rl/trainer_model_based.py       |  22 +++-
 7 files changed, 87 insertions(+), 76 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index dc2e4d3f4..b7ab9dacc 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -253,16 +253,16 @@ def expand_squeeze_to_nd(x, n, squeeze_dim=2, expand_dim=-1):
 
 
 def standardize_images(x):
-  """Image standardization on batches."""
+  """Image standardization on batches and videos."""
   with tf.name_scope("standardize_images", [x]):
-    x = tf.to_float(x)
+    x_shape = shape_list(x)
+    x = tf.to_float(tf.reshape(x, [-1] + x_shape[-3:]))
     x_mean = tf.reduce_mean(x, axis=[1, 2, 3], keepdims=True)
     x_variance = tf.reduce_mean(
         tf.square(x - x_mean), axis=[1, 2, 3], keepdims=True)
-    x_shape = shape_list(x)
-    num_pixels = tf.to_float(x_shape[1] * x_shape[2] * x_shape[3])
+    num_pixels = tf.to_float(x_shape[-1] * x_shape[-2] * x_shape[-3])
     x = (x - x_mean) / tf.maximum(tf.sqrt(x_variance), tf.rsqrt(num_pixels))
-    return x
+    return tf.reshape(x, x_shape)
 
 
 def flatten4d3d(x):
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index c97df5ef5..a896ccefe 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -535,17 +535,8 @@ def bottom(self, x):
     inputs = x
     with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
       common_layers.summarize_video(inputs, "inputs")
-      inputs_shape = common_layers.shape_list(inputs)
-      # Standardize frames.
-      inputs = tf.reshape(inputs, [-1] + inputs_shape[2:])
       inputs = common_layers.standardize_images(inputs)
-      inputs = tf.reshape(inputs, inputs_shape)
-      # Concatenate the time dimension on channels for image models to work.
-      transposed = tf.transpose(inputs, [0, 2, 3, 1, 4])
-      return tf.reshape(transposed, [
-          inputs_shape[0], inputs_shape[2], inputs_shape[3],
-          inputs_shape[1] * inputs_shape[4]
-      ])
+      return common_layers.time_to_channels(inputs)
 
   def targets_bottom(self, x, summary_prefix="targets_bottom"):  # pylint: disable=arguments-differ
     inputs = x
@@ -573,10 +564,13 @@ def top(self, body_output, targets):
     num_frames = common_layers.shape_list(targets)[1]
     body_output_shape = common_layers.shape_list(body_output)
     # We assume the body output is of this shape and layout.
+    # Note: if you tf.concat([frames], axis=-1) at the end of your model,
+    # then you need to reshape to [..., num_frames, depth] like below, not
+    # into [..., depth, num_frames] due to memory layout of concat/reshape.
     reshape_shape = body_output_shape[:-1] + [
-        num_channels, self.top_dimensionality, num_frames]
+        num_channels, num_frames, self.top_dimensionality]
     res = tf.reshape(body_output, reshape_shape)
-    res = tf.transpose(res, [0, 5, 1, 2, 3, 4])
+    res = tf.transpose(res, [0, 4, 1, 2, 3, 5])
     res_shape = common_layers.shape_list(res)
     res_argmax = tf.argmax(tf.reshape(res, [-1, res_shape[-1]]), axis=-1)
     res_argmax = tf.reshape(res_argmax, res_shape[:-1])
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index f77b0a9c4..8685a3fb9 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -594,6 +594,7 @@ def encoder(self, x):
               activation=common_layers.belu,
               name="strided")
           y = x
+          y = tf.nn.dropout(y, 1.0 - hparams.residual_dropout)
           for r in range(hparams.num_residual_layers):
             residual_filters = filters
             if r < hparams.num_residual_layers - 1:
@@ -606,7 +607,7 @@ def encoder(self, x):
                 padding="SAME",
                 activation=common_layers.belu,
                 name="residual_%d" % r)
-          x += tf.nn.dropout(y, 1.0 - hparams.residual_dropout)
+          x += y
           x = common_layers.layer_norm(x, name="ln")
       return x, layers
 
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 364625dfd..5501e523c 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -85,9 +85,9 @@ def body_single(self, features):
     # Run a stack of convolutions.
     for i in range(hparams.num_hidden_layers):
       with tf.variable_scope("layer%d" % i):
-        y = tf.layers.conv2d(x, filters, kernel1, activation=common_layers.belu,
+        y = tf.nn.dropout(x, 1.0 - hparams.dropout)
+        y = tf.layers.conv2d(y, filters, kernel1, activation=common_layers.belu,
                              strides=(1, 1), padding="SAME")
-        y = tf.nn.dropout(y, 1.0 - hparams.dropout)
         if i == 0:
           x = y
         else:
@@ -172,6 +172,7 @@ def body(self, features):
       sampled_frame = tf.reshape(
           res_frames[-1], shape[:-1] + [hparams.problem.num_channels, 256])
       sampled_frame = tf.to_float(tf.argmax(sampled_frame, axis=-1))
+      sampled_frame = common_layers.standardize_images(sampled_frame)
       if is_predicting:
         all_frames[i + hparams.video_num_input_frames] = sampled_frame
 
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 6d234adab..61bd2fd03 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -65,8 +65,8 @@ def next_frame_sampling():
   """Basic conv model with scheduled sampling."""
   hparams = next_frame_basic_deterministic()
   hparams.video_num_target_frames = 2
-  hparams.scheduled_sampling_warmup_steps = 30000
-  hparams.scheduled_sampling_prob = 0.1
+  hparams.scheduled_sampling_warmup_steps = 50000
+  hparams.scheduled_sampling_prob = 0.5
   return hparams
 
 
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index c73a3b599..be7544d2c 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -41,61 +41,62 @@ def define_train(hparams):
 
 
 def train(hparams, event_dir=None, model_dir=None,
-          restore_agent=True, epoch=0):
+          restore_agent=True, epoch=0, name_scope="rl_train"):
   """Train."""
   with tf.Graph().as_default():
-    train_summary_op, _, initialization = define_train(hparams)
-    if event_dir:
-      summary_writer = tf.summary.FileWriter(
-          event_dir, graph=tf.get_default_graph(), flush_secs=60)
-    else:
-      summary_writer = None
+    with tf.name_scope(name_scope):
+      train_summary_op, _, initialization = define_train(hparams)
+      if event_dir:
+        summary_writer = tf.summary.FileWriter(
+            event_dir, graph=tf.get_default_graph(), flush_secs=60)
+      else:
+        summary_writer = None
 
-    if model_dir:
-      model_saver = tf.train.Saver(
-          tf.global_variables(".*network_parameters.*"))
-    else:
-      model_saver = None
+      if model_dir:
+        model_saver = tf.train.Saver(
+            tf.global_variables(".*network_parameters.*"))
+      else:
+        model_saver = None
 
-    # TODO(piotrmilos): This should be refactored, possibly with
-    # handlers for each type of env
-    if hparams.environment_spec.simulated_env:
-      env_model_loader = tf.train.Saver(
-          tf.global_variables("next_frame*"))
-    else:
-      env_model_loader = None
+      # TODO(piotrmilos): This should be refactored, possibly with
+      # handlers for each type of env
+      if hparams.environment_spec.simulated_env:
+        env_model_loader = tf.train.Saver(
+            tf.global_variables("next_frame*"))
+      else:
+        env_model_loader = None
 
-    with tf.Session() as sess:
-      sess.run(tf.global_variables_initializer())
-      initialization(sess)
-      if env_model_loader:
-        trainer_lib.restore_checkpoint(
-            hparams.world_model_dir, env_model_loader, sess,
-            must_restore=True)
-      start_step = 0
-      if model_saver and restore_agent:
-        start_step = trainer_lib.restore_checkpoint(
-            model_dir, model_saver, sess)
+      with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+        initialization(sess)
+        if env_model_loader:
+          trainer_lib.restore_checkpoint(
+              hparams.world_model_dir, env_model_loader, sess,
+              must_restore=True)
+        start_step = 0
+        if model_saver and restore_agent:
+          start_step = trainer_lib.restore_checkpoint(
+              model_dir, model_saver, sess)
 
-      # Fail-friendly, don't train if already trained for this epoch
-      if start_step >= ((hparams.epochs_num * (epoch + 1))):
-        tf.logging.info("Skipping PPO training for epoch %d as train steps "
-                        "(%d) already reached", epoch, start_step)
-        return
+        # Fail-friendly, don't train if already trained for this epoch
+        if start_step >= ((hparams.epochs_num * (epoch + 1))):
+          tf.logging.info("Skipping PPO training for epoch %d as train steps "
+                          "(%d) already reached", epoch, start_step)
+          return
 
-      for epoch_index in range(hparams.epochs_num):
-        summary = sess.run(train_summary_op)
-        if summary_writer:
-          summary_writer.add_summary(summary, epoch_index)
-        if (hparams.eval_every_epochs and
-            epoch_index % hparams.eval_every_epochs == 0):
-          if summary_writer and summary:
+        for epoch_index in range(hparams.epochs_num):
+          summary = sess.run(train_summary_op)
+          if summary_writer:
             summary_writer.add_summary(summary, epoch_index)
-          else:
-            tf.logging.info("Eval summary not saved")
-        if (model_saver and hparams.save_models_every_epochs and
-            (epoch_index % hparams.save_models_every_epochs == 0 or
-             (epoch_index + 1) == hparams.epochs_num)):
-          ckpt_path = os.path.join(
-              model_dir, "model.ckpt-{}".format(epoch_index + 1 + start_step))
-          model_saver.save(sess, ckpt_path)
+          if (hparams.eval_every_epochs and
+              epoch_index % hparams.eval_every_epochs == 0):
+            if summary_writer and summary:
+              summary_writer.add_summary(summary, epoch_index)
+            else:
+              tf.logging.info("Eval summary not saved")
+          if (model_saver and hparams.save_models_every_epochs and
+              (epoch_index % hparams.save_models_every_epochs == 0 or
+               (epoch_index + 1) == hparams.epochs_num)):
+            ckpt_path = os.path.join(
+                model_dir, "model.ckpt-{}".format(epoch_index + 1 + start_step))
+            model_saver.save(sess, ckpt_path)
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 365b82f41..ede7d175c 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -183,7 +183,8 @@ def train_agent(problem_name, agent_model_dir,
       "output_dir": world_model_dir,
       "data_dir": epoch_data_dir,
   }):
-    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=epoch)
+    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=epoch,
+                         name_scope="ppo_sim")
 
 
 def train_agent_real_env(
@@ -218,7 +219,8 @@ def train_agent_real_env(
       "data_dir": epoch_data_dir,
   }):
     # epoch = 10**20 is a hackish way to avoid skiping training
-    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=10**20)
+    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=10**20,
+                         name_scope="ppo_real")
 
 
 def evaluate_world_model(simulated_problem_name, problem_name, hparams,
@@ -266,6 +268,7 @@ def train_world_model(problem_name, data_dir, output_dir, hparams, epoch):
       "hparams_set": hparams.generative_model_params,
       "hparams": "learning_rate_constant=%.6f" % learning_rate,
       "eval_steps": 100,
+      "local_eval_frequency": 2000,
       "train_steps": train_steps,
   }):
     t2t_trainer.main([])
@@ -519,6 +522,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     mean_reward_summary.value[0].simple_value = mean_reward
     eval_metrics_writer.add_summary(model_reward_accuracy_summary, epoch)
     eval_metrics_writer.add_summary(mean_reward_summary, epoch)
+    eval_metrics_writer.flush()
 
     # Report metrics
     eval_metrics = {"model_reward_accuracy": model_reward_accuracy,
@@ -599,10 +603,12 @@ def rl_modelrl_base():
 
 @registry.register_hparams
 def rl_modelrl_base_quick():
-  """Base setting with only 2 epochs and 500 PPO steps per epoch."""
+  """Base setting but quicker with only 2 epochs."""
   hparams = rl_modelrl_base()
   hparams.epochs = 2
-  hparams.ppo_epochs_num = 500
+  hparams.ppo_epochs_num = 1000
+  hparams.ppo_epoch_length = 50
+  hparams.real_ppo_epochs_num = 10
   return hparams
 
 
@@ -615,6 +621,14 @@ def rl_modelrl_base_quick_sd():
   return hparams
 
 
+@registry.register_hparams
+def rl_modelrl_base_quick_sm():
+  """Quick setting with sampling."""
+  hparams = rl_modelrl_base_quick()
+  hparams.generative_model_params = "next_frame_sampling"
+  return hparams
+
+
 @registry.register_hparams
 def rl_modelrl_base_stochastic():
   """Base setting with a stochastic next-frame model."""

From b661f940cad3aba8740e4aa2191fd1f496ace6aa Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 7 Sep 2018 13:59:11 -0700
Subject: [PATCH 0763/2720] Move oss_scripts to top-level

PiperOrigin-RevId: 212031390
---
 .travis.yml                                                 | 6 +++---
 .../oss_scripts => oss_scripts}/oss_integration_test.sh     | 2 ++
 .../oss_scripts => oss_scripts}/oss_pip_install.sh          | 0
 {tensor2tensor/oss_scripts => oss_scripts}/oss_release.sh   | 0
 {tensor2tensor/oss_scripts => oss_scripts}/oss_tests.sh     | 0
 5 files changed, 5 insertions(+), 3 deletions(-)
 rename {tensor2tensor/oss_scripts => oss_scripts}/oss_integration_test.sh (95%)
 rename {tensor2tensor/oss_scripts => oss_scripts}/oss_pip_install.sh (100%)
 rename {tensor2tensor/oss_scripts => oss_scripts}/oss_release.sh (100%)
 rename {tensor2tensor/oss_scripts => oss_scripts}/oss_tests.sh (100%)

diff --git a/.travis.yml b/.travis.yml
index 68cda574b..628f13548 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -33,10 +33,10 @@ before_install:
   - sudo apt-get update -qq
   - sudo apt-get install -qq libhdf5-dev
 install:
-  - ./tensor2tensor/oss_scripts/oss_pip_install.sh
+  - ./oss_scripts/oss_pip_install.sh
 script:
-  - ./tensor2tensor/oss_scripts/oss_tests.sh
-  - ./tensor2tensor/oss_scripts/oss_integration_test.sh
+  - ./oss_scripts/oss_tests.sh
+  - ./oss_scripts/oss_integration_test.sh
 
   # Conditional commands should each be in a separate block to get proper
   # errors on Travis.
diff --git a/tensor2tensor/oss_scripts/oss_integration_test.sh b/oss_scripts/oss_integration_test.sh
similarity index 95%
rename from tensor2tensor/oss_scripts/oss_integration_test.sh
rename to oss_scripts/oss_integration_test.sh
index 9e82908a2..580a76866 100755
--- a/tensor2tensor/oss_scripts/oss_integration_test.sh
+++ b/oss_scripts/oss_integration_test.sh
@@ -1,5 +1,7 @@
 #!/bin/bash
 
+# Note that this test script requires docker to be installed and running.
+
 set -v  # print commands as they're executed
 set -e  # fail and exit on any command erroring
 
diff --git a/tensor2tensor/oss_scripts/oss_pip_install.sh b/oss_scripts/oss_pip_install.sh
similarity index 100%
rename from tensor2tensor/oss_scripts/oss_pip_install.sh
rename to oss_scripts/oss_pip_install.sh
diff --git a/tensor2tensor/oss_scripts/oss_release.sh b/oss_scripts/oss_release.sh
similarity index 100%
rename from tensor2tensor/oss_scripts/oss_release.sh
rename to oss_scripts/oss_release.sh
diff --git a/tensor2tensor/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
similarity index 100%
rename from tensor2tensor/oss_scripts/oss_tests.sh
rename to oss_scripts/oss_tests.sh

From 97e4beaa97a75848ba7499290c3e265aa3bc7b1f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 7 Sep 2018 16:41:12 -0700
Subject: [PATCH 0764/2720] Fix typo in metrics.py comment.

PiperOrigin-RevId: 212057129
---
 tensor2tensor/utils/metrics.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 505af9611..db9ddb183 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -33,7 +33,7 @@
 
 class Metrics(object):
   """Available evaluation metrics."""
-  # Entries here should match the keys in METRICS_FN below
+  # Entries here should match the keys in METRICS_FNS below
   ACC = "accuracy"
   ACC_TOP5 = "accuracy_top5"
   ACC_PER_SEQ = "accuracy_per_sequence"

From d11504132ebfafddf500f2f63b7cf3a144523833 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 7 Sep 2018 16:45:02 -0700
Subject: [PATCH 0765/2720] video/base_test.py only on latest TF (very large
 test)

PiperOrigin-RevId: 212057679
---
 oss_scripts/oss_tests.sh | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index de53f604a..d4c0225ea 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -49,6 +49,7 @@ pytest \
   --ignore=tensor2tensor/rl/trainer_model_based_stochastic_test.py \
   --ignore=tensor2tensor/rl/trainer_model_based_sv2p_test.py \
   --ignore=tensor2tensor/models/research \
+  --ignore=tensor2tensor/models/video/base_test.py \
   --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary
 set_status
 
@@ -78,6 +79,8 @@ if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"
 then
     pytest tensor2tensor/rl/trainer_model_based_test.py
     set_status
+    pytest tensor2tensor/models/video/base_test.py
+    set_status
     jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/hello_t2t.ipynb
     set_status
     jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/t2t_problem.ipynb;

From 06862886125f5b5d262e2b0e6ffeba059dc0d57d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 7 Sep 2018 16:45:37 -0700
Subject: [PATCH 0766/2720] s/runScheduledSampleFunc/_run_scheduled_sample_func
 to keep linter happy.

PiperOrigin-RevId: 212057756
---
 tensor2tensor/layers/common_video_test.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index 06e650797..f87c50146 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -24,7 +24,7 @@
 
 class CommonVideoTest(tf.test.TestCase):
 
-  def runScheduledSampleFunc(self, func, var, batch_size):
+  def _run_scheduled_sample_func(self, func, var, batch_size):
     ground_truth_x = list(range(1, batch_size+1))
     generated_x = [-x for x in ground_truth_x]
     ground_truth_x = tf.convert_to_tensor(ground_truth_x)
@@ -35,39 +35,39 @@ def runScheduledSampleFunc(self, func, var, batch_size):
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testScheduledSampleProbStart(self):
-    ground_truth_x, _, ss_out = self.runScheduledSampleFunc(
+    ground_truth_x, _, ss_out = self._run_scheduled_sample_func(
         common_video.scheduled_sample_prob, 1.0, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testScheduledSampleProbMid(self):
-    _, _, ss_out = self.runScheduledSampleFunc(
+    _, _, ss_out = self._run_scheduled_sample_func(
         common_video.scheduled_sample_prob, 0.5, 1000)
     positive_count = np.sum(ss_out > 0)
     self.assertAlmostEqual(positive_count / 1000.0, 0.5, places=1)
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testScheduledSampleProbEnd(self):
-    _, generated_x, ss_out = self.runScheduledSampleFunc(
+    _, generated_x, ss_out = self._run_scheduled_sample_func(
         common_video.scheduled_sample_prob, 0.0, 10)
     self.assertAllEqual(generated_x, ss_out)
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testScheduledSampleCountStart(self):
-    ground_truth_x, _, ss_out = self.runScheduledSampleFunc(
+    ground_truth_x, _, ss_out = self._run_scheduled_sample_func(
         common_video.scheduled_sample_count, 10, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testScheduledSampleCountMid(self):
-    _, _, ss_out = self.runScheduledSampleFunc(
+    _, _, ss_out = self._run_scheduled_sample_func(
         common_video.scheduled_sample_count, 5, 10)
     positive_count = np.sum(ss_out > 0)
     self.assertEqual(positive_count, 5)
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testScheduledSampleCountEnd(self):
-    _, generated_x, ss_out = self.runScheduledSampleFunc(
+    _, generated_x, ss_out = self._run_scheduled_sample_func(
         common_video.scheduled_sample_count, 0, 10)
     self.assertAllEqual(generated_x, ss_out)
 

From 2bf2b1d4aa0262f338e594211a742a458a19c267 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 7 Sep 2018 17:17:52 -0700
Subject: [PATCH 0767/2720] Internal change

PiperOrigin-RevId: 212062268
---
 tensor2tensor/models/research/glow.py         |  18 +-
 tensor2tensor/models/research/glow_ops.py     | 168 +++++++++++++++---
 .../models/research/glow_ops_test.py          |  66 ++++++-
 3 files changed, 212 insertions(+), 40 deletions(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 90e787577..183843729 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -41,6 +41,9 @@ def glow_hparams():
   hparams.weight_decay = 0.0
   hparams.learning_rate_constant = 3e-4
   hparams.batch_size = 32
+  # can be prev_level, prev_step or normal.
+  # see: glow_ops.merge_level_and_latent_dist
+  hparams.add_hparam("level_prior_scale", "prev_level")
   hparams.add_hparam("n_levels", 3)
   hparams.add_hparam("n_bits_x", 8)
   hparams.add_hparam("depth", 32)
@@ -118,15 +121,7 @@ def body(self, features):
 
     # Scale x such that the pixels lie in-between -0.5 and.0.5
     x = self.preprocess(x)
-
-    n_bins = 2**self.hparams.n_bits_x
-    batch_size, height, width, n_channels = common_layers.shape_list(x)
-    hwc = float(height * width * n_channels)
-
-    x = x + tf.random_uniform(
-        shape=(batch_size, height, width, n_channels),
-        minval=0.0, maxval=1.0/n_bins)
-    objective = -np.log(n_bins) * hwc * tf.ones(batch_size)
+    x, objective = glow_ops.uniform_binning_correction(x)
 
     # The arg_scope call ensures that the actnorm parameters are set such that
     # the per-channel output activations have zero mean and unit variance
@@ -136,7 +131,7 @@ def body(self, features):
     init_op = tf.logical_and(tf.equal(global_step, 0), self.is_training)
     ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
     with arg_scope(ops, init=init_op):
-      self.z, encoder_objective, self.eps = glow_ops.encoder_decoder(
+      self.z, encoder_objective, self.eps, _ = glow_ops.encoder_decoder(
           "codec", x, self.hparams, eps=None, reverse=False)
       objective += encoder_objective
 
@@ -146,5 +141,6 @@ def body(self, features):
       objective += prior_objective
 
     # bits per pixel
-    objective = -objective / (np.log(2) * hwc)
+    _, h, w, c = common_layers.shape_list(x)
+    objective = -objective / (np.log(2) * h * w * c)
     return tf.zeros_like(features["targets"]), {"training": objective}
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index e23e1433f..6dc6f64be 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -440,41 +440,98 @@ def split_prior(name, x):
 
 
 @add_arg_scope
-def split(name, x, reverse=False, eps=None, eps_std=None):
+def merge_level_and_latent_dist(level_dist, latent_dist,
+                                merge_std="prev_level"):
+  """Merge level_dist and latent_dist.
+
+  new_dist ~ N(level_dist.mean + latent_dis.mean, std) where std is determined
+  according to merge_std.
+
+  Args:
+    level_dist: instance of tf.distributions.Normal
+    latent_dist: instance of tf.distributions.Normal
+    merge_std: can be "prev_level", "prev_step" or "normal".
+  Returns:
+    merged_dist: instance of tf.distributions.Normal
+  """
+  level_mean, level_std = level_dist.loc, level_dist.scale
+  latent_mean, latent_std = latent_dist.loc, latent_dist.scale
+  new_mean = level_mean + latent_mean
+  if merge_std == "normal":
+    z_shape = common_layers.shape_list(latent_mean)
+    log_scale = tf.get_variable(
+        "merge_std", shape=z_shape, dtype=tf.float32,
+        initializer=tf.zeros_initializer(), trainable=False)
+    scale = tf.exp(log_scale * 3.0)
+  elif merge_std == "prev_level":
+    scale = level_std
+  elif merge_std == "prev_step":
+    scale = latent_std
+  tf.summary.scalar("latent_scale", tf.reduce_mean(scale))
+  return tf.distributions.Normal(loc=new_mean, scale=scale)
+
+
+@add_arg_scope
+def compute_prior(name, z, latent, merge_std):
+  """Distribution condtioned on both z and latent."""
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    prior_dist = split_prior("level_prior", z)
+    if latent is not None:
+      latent_shape = common_layers.shape_list(latent)
+      z_shape = common_layers.shape_list(z)
+      if latent_shape != z_shape:
+        raise ValueError("Expected latent_shape to be %s, got %s" %
+                         (latent_shape, z_shape))
+      latent_dist = scale_gaussian_prior(
+          "latent_prior", latent, logscale_factor=3.0)
+      prior_dist = merge_level_and_latent_dist(prior_dist, latent_dist,
+                                               merge_std=merge_std)
+  return prior_dist
+
+
+@add_arg_scope
+def split(name, x, reverse=False, eps=None, eps_std=None, cond_latent=None,
+          merge_std="normal"):
   """Splits / concatenates x into x1 and x2 across number of channels.
 
   For the forward pass, x2 is assumed be gaussian,
   i.e P(x2 | x1) ~ N(mu(x1), sigma(x1)) where mu and sigma are the outputs of
-  a network. For the reverse pass, x2 is determined from mu(x1) and sigma(x1).
-  This is deterministic/stochastic depending on whether eps is provided.
+  a one-layer network. For the reverse pass, x2 is determined
+  from mu(x1) and sigma(x1). This is deterministic/stochastic depending on
+  whether eps is provided.
 
   Args:
     name:
     x:
     reverse: Forward or reverse pass.
-    eps: If eps is provided, x2
-    eps_std: Sample x2
+    eps: If eps is provided, x2 is set to be
+    eps_std: Sample x2.
+    cond_latent: optionally condition x2 on cond_latent.
+    merge_std: used to determine the std of the gaussian prior on x2 if
+               cond_latent is provided.
 
   Returns:
+  Raises:
+    ValueError: If latent is provided and shape is not equal to NHW(C/2)
+                where (NHWC) is the size of x.
   """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     if not reverse:
       x1, x2 = tf.split(x, num_or_size_splits=2, axis=-1)
 
       # objective: P(x2|x1) ~N(x2 ; NN(x1))
-      x1_dist = split_prior("split_prior", x1)
-      logpb = tf.reduce_sum(x1_dist.log_prob(x2), axis=[1, 2, 3])
-
-      eps = get_eps(x1_dist, x2)
-      return x1, logpb, eps
+      prior_dist = compute_prior("prior_on_z2", x1, cond_latent, merge_std)
+      logpb = tf.reduce_sum(prior_dist.log_prob(x2), axis=[1, 2, 3])
+      eps = get_eps(prior_dist, x2)
+      return x1, logpb, eps, x2
     else:
-      x1_dist = split_prior("split_prior", x)
+      prior_dist = compute_prior("prior_on_z2", x, cond_latent, merge_std)
       if eps is not None:
-        x2 = set_eps(x1_dist, eps)
+        x2 = set_eps(prior_dist, eps)
       elif eps_std is not None:
         x2 = eps_std * tf.random_normal(common_layers.shape_list(x))
       else:
-        x2 = x1_dist.sample()
+        x2 = prior_dist.sample()
       return tf.concat([x, x2], 3)
 
 
@@ -525,6 +582,34 @@ def revnet(name, x, hparams, reverse=True):
     return x, objective
 
 
+@add_arg_scope
+def scale_gaussian_prior(name, z, logscale_factor=3.0, trainable=True):
+  """Returns N(s^i * z^i, std^i) where s^i and std^i are pre-component.
+
+  s^i is a learnable parameter with identity initialization.
+  std^i is optionally learnable with identity initialization.
+
+  Args:
+    name: variable scope.
+    z: input_tensor
+    logscale_factor: equivalent to scaling up the learning_rate by a factor
+                     of logscale_factor.
+    trainable: Whether or not std^i is learnt.
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    z_shape = common_layers.shape_list(z)
+    latent_multiplier = tf.get_variable(
+        "latent_multiplier", shape=z_shape, dtype=tf.float32,
+        initializer=tf.ones_initializer())
+    log_scale = tf.get_variable(
+        "log_scale_latent", shape=z_shape, dtype=tf.float32,
+        initializer=tf.zeros_initializer(), trainable=trainable)
+    log_scale = log_scale * logscale_factor
+    tf.summary.scalar("gaussian_log_scale", tf.reduce_mean(log_scale))
+    return tf.distributions.Normal(
+        loc=latent_multiplier * z, scale=tf.exp(log_scale))
+
+
 @add_arg_scope
 def top_prior(name, x, learn_prior="normal"):
   """Log probability of x being gaussian.
@@ -556,13 +641,43 @@ def top_prior(name, x, learn_prior="normal"):
     return objective, prior_dist
 
 
+def uniform_binning_correction(x, n_bits=8):
+  """Replaces x^i with q^i(x) = U(x, x + 1.0 / 256.0).
+
+  Args:
+    x: 4-D Tensor of shape (NHWC)
+    n_bits: optional.
+  Returns:
+    x: x ~ U(x, x + 1.0 / 256)
+    objective: Equivalent to -q(x)*log(q(x)).
+  """
+  n_bins = 2**n_bits
+  batch_size, height, width, n_channels = common_layers.shape_list(x)
+  hwc = float(height * width * n_channels)
+
+  x = x + tf.random_uniform(
+      shape=(batch_size, height, width, n_channels),
+      minval=0.0, maxval=1.0/n_bins)
+  objective = -np.log(n_bins) * hwc * tf.ones(batch_size)
+  return x, objective
+
+
 @add_arg_scope
-def encoder_decoder(name, x, hparams, eps=None, reverse=False):
+def encoder_decoder(name, x, hparams, eps=None, reverse=False,
+                    cond_latents=None):
   """Glow encoder-decoder. n_levels of (Squeeze + Flow + Split.) operations."""
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
 
+    if eps and len(eps) != hparams.n_levels - 1:
+      raise ValueError("Expected length of eps to be %d, got %d" %
+                       (hparams.n_levels - 1, len(eps)))
+    if cond_latents and len(cond_latents) != hparams.n_levels - 1:
+      raise ValueError("Expected level_latets to be %d, got %d" %
+                       (hparams.n_levels - 1, len(cond_latents)))
+
     objective = 0.0
     all_eps = []
+    all_latents = []
 
     if not reverse:
       # Squeeze + Flow + Split
@@ -573,23 +688,34 @@ def encoder_decoder(name, x, hparams, eps=None, reverse=False):
         objective += obj
 
         if level < hparams.n_levels - 1:
-          x, obj, eps = split("split_%d" % level, x, reverse=False)
+
+          curr_latent = None
+          if cond_latents is not None:
+            curr_latent = cond_latents[level]
+
+          x, obj, eps, z = split(
+              "split_%d" % level, x, reverse=False, cond_latent=curr_latent,
+              merge_std=hparams.level_prior_scale)
           objective += obj
           all_eps.append(eps)
-      return x, objective, all_eps
+          all_latents.append(z)
+      return x, objective, all_eps, all_latents
 
     else:
-      if eps and len(eps) != hparams.n_levels - 1:
-        raise ValueError("Expected length of eps to be %d, got %d" %
-                         (hparams.n_levels - 1, len(eps)))
-
       for level in reversed(range(hparams.n_levels)):
         if level < hparams.n_levels - 1:
 
           curr_eps = None
           if eps:
             curr_eps = eps[level]
-          x = split("split_%d" % level, x, eps=curr_eps, reverse=True)
+
+          curr_latent = None
+          if cond_latents is not None:
+            curr_latent = cond_latents[level]
+
+          x = split("split_%d" % level, x, eps=curr_eps, reverse=True,
+                    cond_latent=curr_latent,
+                    merge_std=hparams.level_prior_scale)
 
         x, obj = revnet(
             "revnet_%d" % level, x, hparams=hparams, reverse=True)
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 8e90117f4..d48dd244c 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -137,11 +137,12 @@ def test_split_prior(self):
   def test_split(self):
     with tf.Graph().as_default():
       x = tf.random_uniform(shape=(16, 5, 5, 32))
-      x_inv, _, eps = glow_ops.split("split", x)
+      x_inv, _, eps, z = glow_ops.split("split", x)
       x_inv_inv = glow_ops.split("split", x_inv, reverse=True, eps=eps)
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())
-        x_inv_np, diff = session.run([x_inv, x - x_inv_inv])
+        x_inv_np, diff, z_np = session.run([x_inv, x - x_inv_inv, z])
+        self.assertEqual(z_np.shape, (16, 5, 5, 16))
         self.assertEqual(x_inv_np.shape, (16, 5, 5, 16))
         self.assertTrue(np.allclose(diff, 0.0, atol=1e-5))
 
@@ -166,20 +167,25 @@ def test_revnet_reversibility(self):
   def test_encoder_decoder(self):
     with tf.Graph().as_default():
       hparams = glow.glow_hparams()
-      hparams.n_levels = 2
+      hparams.n_levels = 3
       hparams.depth = 2
 
       x = tf.random_uniform(shape=(16, 64, 64, 4), seed=0)
-      x_inv, _, eps = glow_ops.encoder_decoder(
+      x_inv, _, eps, z_levels = glow_ops.encoder_decoder(
           "encoder_decoder", x, hparams, reverse=False)
       x_inv_inv, _ = glow_ops.encoder_decoder(
           "encoder_decoder", x_inv, hparams, eps=eps, reverse=True)
 
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())
-        diff, x_inv_np = session.run([x - x_inv_inv, x_inv])
+        diff, x_inv_np, z_levels_np = session.run(
+            [x - x_inv_inv, x_inv, z_levels])
+        self.assertEqual(len(z_levels_np), 2)
+        # (h_i, w_i, c_i) = (h_{i-1}/f, w_{i-1}/f, c_{i-1}*(2f)/2) where (f=2)
+        self.assertEqual(z_levels_np[0].shape, (16, 32, 32, 8))
+        self.assertEqual(z_levels_np[1].shape, (16, 16, 16, 16))
         self.assertTrue(x_inv_np.shape, (16, 8, 8, 64))
-        self.assertTrue(np.allclose(diff, 0.0, atol=1e-3))
+        self.assertTrue(np.allclose(diff, 0.0, atol=1e-2))
 
   def test_encoder_decoder_practical_usage(self):
     """Tests the following sequence of operations.
@@ -200,7 +206,7 @@ def test_encoder_decoder_practical_usage(self):
 
       ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
       with arg_scope(ops, init=True):
-        x_inv, _, _ = glow_ops.encoder_decoder(
+        x_inv, _, _, _ = glow_ops.encoder_decoder(
             "revnet", x_t, hparams, reverse=False)
       curr_dir = tempfile.mkdtemp()
       model_path = os.path.join(curr_dir, "model")
@@ -217,7 +223,7 @@ def test_encoder_decoder_practical_usage(self):
       x_t = tf.convert_to_tensor(x_rand)
       ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
       with arg_scope(ops, init=False):
-        x_inv2, _, all_eps = glow_ops.encoder_decoder(
+        x_inv2, _, all_eps, _ = glow_ops.encoder_decoder(
             "revnet", x_t, hparams, reverse=False)
         x_inv_inv_, _ = glow_ops.encoder_decoder(
             "revnet", x_inv2, hparams, eps=all_eps, reverse=True)
@@ -229,6 +235,50 @@ def test_encoder_decoder_practical_usage(self):
         diff = np.abs(x_inv_inv_np - x_rand)
         self.assertTrue(np.allclose(diff, 0.0, atol=1e-3))
 
+  def test_scale_gaussian_prior(self):
+    with tf.Graph().as_default():
+      rng = np.random.RandomState(0)
+      img_shape = (16, 2, 2, 2)
+      x_rand = np.asarray(rng.randint(0, 10, img_shape), dtype=np.float32)
+      z_rand = np.asarray(rng.randint(0, 10, img_shape), dtype=np.float32)
+      x_t = tf.convert_to_tensor(x_rand)
+      z_t = tf.convert_to_tensor(z_rand)
+      dist = glow_ops.scale_gaussian_prior(
+          "scale_gaussian_prior", z_t, x_t, trainable=True)
+      with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+        mean, scale = sess.run([dist.loc, dist.scale])
+        self.assertTrue(np.allclose(mean, z_rand))
+        self.assertTrue(np.allclose(scale, 1.0))
+
+  def check_split_latent_conditioning(self, merge_std):
+    with tf.Graph().as_default():
+      rng = np.random.RandomState(0)
+      x_rand = rng.randn(12, 32, 32, 32).astype(np.float32)
+      latent_rand = rng.randn(12, 32, 32, 16).astype(np.float32)
+      x_t = tf.convert_to_tensor(x_rand)
+      latent_t = tf.convert_to_tensor(latent_rand)
+
+      # Test initalization.
+      # x2 ~ N(scale * latent, 1.0) where initial scale is 1.0
+      exp_x2 = x_rand[:, :, :, 16:]
+      exp_eps = x_rand[:, :, :, 16:] - latent_rand
+      x_inv, _, eps, x2_t = glow_ops.split(merge_std, x_t, cond_latent=latent_t,
+                                           merge_std=merge_std)
+      # Test reversibility.
+      x_inv_inv = glow_ops.split(merge_std, x_inv, cond_latent=latent_t,
+                                 merge_std=merge_std, eps=eps, reverse=True)
+      with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+        actual_eps, actual_x2, diff_np = sess.run([eps, x2_t, x_inv_inv - x_t])
+        self.assertTrue(np.allclose(diff_np, 0.0, atol=1e-5))
+        self.assertTrue(np.allclose(actual_eps, exp_eps))
+        self.assertTrue(np.allclose(exp_x2, actual_x2))
+
+  def test_split_latent_conditioning(self):
+    for merge_std in ["normal", "prev_level", "prev_step"]:
+      self.check_split_latent_conditioning(merge_std)
+
 
 if __name__ == "__main__":
   tf.test.main()

From ee6b794e7ca4a98c0156668d885b18b904ec9d0c Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 7 Sep 2018 17:33:09 -0700
Subject: [PATCH 0768/2720] breaking down video tests.

PiperOrigin-RevId: 212064133
---
 oss_scripts/oss_tests.sh                      |  3 -
 .../models/video/basic_deterministic_test.py  | 38 ++++++++
 .../models/video/basic_stochastic_test.py     | 37 ++++++++
 tensor2tensor/models/video/emily_test.py      | 38 ++++++++
 tensor2tensor/models/video/savp_test.py       | 63 +++++++++++++
 tensor2tensor/models/video/sv2p_test.py       | 55 +++++++++++
 .../video/{base_test.py => tests_utils.py}    | 92 +------------------
 7 files changed, 235 insertions(+), 91 deletions(-)
 create mode 100644 tensor2tensor/models/video/basic_deterministic_test.py
 create mode 100644 tensor2tensor/models/video/basic_stochastic_test.py
 create mode 100644 tensor2tensor/models/video/emily_test.py
 create mode 100644 tensor2tensor/models/video/savp_test.py
 create mode 100644 tensor2tensor/models/video/sv2p_test.py
 rename tensor2tensor/models/video/{base_test.py => tests_utils.py} (78%)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index d4c0225ea..de53f604a 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -49,7 +49,6 @@ pytest \
   --ignore=tensor2tensor/rl/trainer_model_based_stochastic_test.py \
   --ignore=tensor2tensor/rl/trainer_model_based_sv2p_test.py \
   --ignore=tensor2tensor/models/research \
-  --ignore=tensor2tensor/models/video/base_test.py \
   --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary
 set_status
 
@@ -79,8 +78,6 @@ if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"
 then
     pytest tensor2tensor/rl/trainer_model_based_test.py
     set_status
-    pytest tensor2tensor/models/video/base_test.py
-    set_status
     jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/hello_t2t.ipynb
     set_status
     jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/t2t_problem.ipynb;
diff --git a/tensor2tensor/models/video/basic_deterministic_test.py b/tensor2tensor/models/video/basic_deterministic_test.py
new file mode 100644
index 000000000..bd95d76c8
--- /dev/null
+++ b/tensor2tensor/models/video/basic_deterministic_test.py
@@ -0,0 +1,38 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Basic tests for basic deterministic model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.models.video import basic_deterministic
+from tensor2tensor.models.video import basic_deterministic_params
+from tensor2tensor.models.video import tests_utils
+
+import tensorflow as tf
+
+
+class NextFrameTest(tests_utils.BaseNextFrameTest):
+
+  def testBasicDeterministic(self):
+    self.TestOnVariousInputOutputSizes(
+        basic_deterministic_params.next_frame_basic_deterministic(),
+        basic_deterministic.NextFrameBasicDeterministic,
+        256,
+        False)
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/video/basic_stochastic_test.py b/tensor2tensor/models/video/basic_stochastic_test.py
new file mode 100644
index 000000000..fdc4951f6
--- /dev/null
+++ b/tensor2tensor/models/video/basic_stochastic_test.py
@@ -0,0 +1,37 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Basic tests for basic stochastic model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.models.video import basic_stochastic
+from tensor2tensor.models.video import tests_utils
+
+import tensorflow as tf
+
+
+class NextFrameTest(tests_utils.BaseNextFrameTest):
+
+  def testBasicStochastic(self):
+    self.TestOnVariousInputOutputSizes(
+        basic_stochastic.next_frame_basic_stochastic(),
+        basic_stochastic.NextFrameBasicStochastic,
+        256,
+        False)
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/video/emily_test.py b/tensor2tensor/models/video/emily_test.py
new file mode 100644
index 000000000..51a8e2e28
--- /dev/null
+++ b/tensor2tensor/models/video/emily_test.py
@@ -0,0 +1,38 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Basic tests for emily's model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.models.video import emily
+from tensor2tensor.models.video import tests_utils
+
+
+import tensorflow as tf
+
+
+class NextFrameTest(tests_utils.BaseNextFrameTest):
+
+  def testEmily(self):
+    self.TestOnVariousInputOutputSizes(
+        emily.next_frame_emily(),
+        emily.NextFrameEmily,
+        1)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/video/savp_test.py b/tensor2tensor/models/video/savp_test.py
new file mode 100644
index 000000000..80ad6f47b
--- /dev/null
+++ b/tensor2tensor/models/video/savp_test.py
@@ -0,0 +1,63 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Basic tests for SAVP model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.models.video import savp
+from tensor2tensor.models.video import savp_params
+from tensor2tensor.models.video import tests_utils
+
+
+import tensorflow as tf
+
+
+class NextFrameTest(tests_utils.BaseNextFrameTest):
+
+  def testSavpVAE(self):
+    savp_hparams = savp_params.next_frame_savp()
+    savp_hparams.use_vae = True
+    savp_hparams.use_gan = False
+    self.TestOnVariousInputOutputSizes(
+        savp_hparams, savp.NextFrameSAVP, 1)
+    self.TestOnVariousUpSampleLayers(
+        savp_hparams, savp.NextFrameSAVP, 1)
+
+  def testSavpGAN(self):
+    hparams = savp_params.next_frame_savp()
+    hparams.use_gan = True
+    hparams.use_vae = False
+    self.TestVideoModel(7, 5, hparams, savp.NextFrameSAVP, 1)
+
+    hparams.gan_optimization = "sequential"
+    self.TestVideoModel(7, 5, hparams, savp.NextFrameSAVP, 1)
+
+  def testSavpGANVAE(self):
+    hparams = savp_params.next_frame_savp()
+    hparams.use_vae = True
+    hparams.use_gan = True
+    self.TestVideoModel(7, 5, hparams, savp.NextFrameSAVP, 1)
+
+  def testInvalidVAEGANCombinations(self):
+    hparams = savp_params.next_frame_savp()
+    hparams.use_gan = False
+    hparams.use_vae = False
+    self.assertRaises(ValueError, self.TestVideoModel,
+                      7, 5, hparams, savp.NextFrameSAVP, 1)
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/video/sv2p_test.py b/tensor2tensor/models/video/sv2p_test.py
new file mode 100644
index 000000000..42c253792
--- /dev/null
+++ b/tensor2tensor/models/video/sv2p_test.py
@@ -0,0 +1,55 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Basic tests for SV2P model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.models.video import sv2p
+from tensor2tensor.models.video import sv2p_params
+from tensor2tensor.models.video import tests_utils
+
+import tensorflow as tf
+
+
+class NextFrameTest(tests_utils.BaseNextFrameTest):
+
+  def testSv2p(self):
+    self.TestOnVariousInputOutputSizes(
+        sv2p_params.next_frame_sv2p(),
+        sv2p.NextFrameSv2p,
+        1)
+
+  def testSv2pWithActions(self):
+    self.TestWithActions(
+        sv2p_params.next_frame_sv2p(),
+        sv2p.NextFrameSv2p,
+        1)
+
+  def testSv2pWithActionsAndRewards(self):
+    self.TestWithActionAndRewards(
+        sv2p_params.next_frame_sv2p(),
+        sv2p.NextFrameSv2p,
+        1)
+
+  def testSv2pTwoFrames(self):
+    self.TestOnVariousInputOutputSizes(
+        sv2p_params.next_frame_sv2p(),
+        sv2p.NextFrameSv2pTwoFrames,
+        1)
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/video/base_test.py b/tensor2tensor/models/video/tests_utils.py
similarity index 78%
rename from tensor2tensor/models/video/base_test.py
rename to tensor2tensor/models/video/tests_utils.py
index e7e467974..67aa3dd81 100644
--- a/tensor2tensor/models/video/base_test.py
+++ b/tensor2tensor/models/video/tests_utils.py
@@ -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.
-"""Basic tests for video prediction models."""
+"""Utilties for testing video models."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -20,14 +20,6 @@
 import numpy as np
 
 from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
-from tensor2tensor.models.video import basic_deterministic
-from tensor2tensor.models.video import basic_deterministic_params
-from tensor2tensor.models.video import basic_stochastic
-from tensor2tensor.models.video import emily
-from tensor2tensor.models.video import savp
-from tensor2tensor.models.video import savp_params
-from tensor2tensor.models.video import sv2p
-from tensor2tensor.models.video import sv2p_params
 
 from tensor2tensor.utils import registry
 
@@ -59,6 +51,7 @@ def action_modalities(hparams):
 
 
 def full_modalities(hparams):
+  """Full modalities with actions and rewards."""
   hparams.problem_hparams.input_modality = {
       "inputs": ("video:l2raw", 256),
       "input_reward": ("symbol:one_hot", 3),
@@ -105,7 +98,8 @@ def get_tensor_shape(tensor):
   return tuple([d.value for d in tensor.shape])
 
 
-class NextFrameTest(tf.test.TestCase):
+class BaseNextFrameTest(tf.test.TestCase):
+  """Base helper class for next frame tests."""
 
   def RunModel(self, model, hparams, features):
     with tf.Session() as session:
@@ -278,81 +272,3 @@ def TestOnVariousUpSampleLayers(self, hparams, model, expected_last_dim):
                         upsample_method="bilinear_upsample_conv")
     self.TestVideoModel(4, 1, hparams, model, expected_last_dim,
                         upsample_method="nn_upsample_conv")
-
-  def testBasicDeterministic(self):
-    self.TestOnVariousInputOutputSizes(
-        basic_deterministic_params.next_frame_basic_deterministic(),
-        basic_deterministic.NextFrameBasicDeterministic,
-        256,
-        False)
-
-  def testBasicStochastic(self):
-    self.TestOnVariousInputOutputSizes(
-        basic_stochastic.next_frame_basic_stochastic(),
-        basic_stochastic.NextFrameBasicStochastic,
-        256,
-        False)
-
-  def testSv2p(self):
-    self.TestOnVariousInputOutputSizes(
-        sv2p_params.next_frame_sv2p(),
-        sv2p.NextFrameSv2p,
-        1)
-
-  def testSv2pWithActions(self):
-    self.TestWithActions(
-        sv2p_params.next_frame_sv2p(),
-        sv2p.NextFrameSv2p,
-        1)
-
-  def testSv2pWithActionsAndRewards(self):
-    self.TestWithActionAndRewards(
-        sv2p_params.next_frame_sv2p(),
-        sv2p.NextFrameSv2p,
-        1)
-
-  def testSv2pTwoFrames(self):
-    self.TestOnVariousInputOutputSizes(
-        sv2p_params.next_frame_sv2p(),
-        sv2p.NextFrameSv2pTwoFrames,
-        1)
-
-  def testEmily(self):
-    self.TestOnVariousInputOutputSizes(
-        emily.next_frame_emily(),
-        emily.NextFrameEmily,
-        1)
-
-  def testSavpVAE(self):
-    savp_hparams = savp_params.next_frame_savp()
-    savp_hparams.use_vae = True
-    savp_hparams.use_gan = False
-    self.TestOnVariousInputOutputSizes(
-        savp_hparams, savp.NextFrameSAVP, 1)
-    self.TestOnVariousUpSampleLayers(
-        savp_hparams, savp.NextFrameSAVP, 1)
-
-  def testSavpGAN(self):
-    hparams = savp_params.next_frame_savp()
-    hparams.use_gan = True
-    hparams.use_vae = False
-    self.TestVideoModel(7, 5, hparams, savp.NextFrameSAVP, 1)
-
-    hparams.gan_optimization = "sequential"
-    self.TestVideoModel(7, 5, hparams, savp.NextFrameSAVP, 1)
-
-  def testSavpGANVAE(self):
-    hparams = savp_params.next_frame_savp()
-    hparams.use_vae = True
-    hparams.use_gan = True
-    self.TestVideoModel(7, 5, hparams, savp.NextFrameSAVP, 1)
-
-  def testInvalidVAEGANCombinations(self):
-    hparams = savp_params.next_frame_savp()
-    hparams.use_gan = False
-    hparams.use_vae = False
-    self.assertRaises(ValueError, self.TestVideoModel,
-                      7, 5, hparams, savp.NextFrameSAVP, 1)
-
-if __name__ == "__main__":
-  tf.test.main()

From fdf993b7149cc2d5ab2941d050973817a62f1918 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 7 Sep 2018 17:50:19 -0700
Subject: [PATCH 0769/2720] Internal change

PiperOrigin-RevId: 212065960
---
 tensor2tensor/data_generators/video_utils.py  | 31 ++++++++++-
 .../data_generators/video_utils_test.py       | 53 +++++++++++++++++++
 tensor2tensor/layers/common_video.py          | 11 +++-
 3 files changed, 92 insertions(+), 3 deletions(-)
 create mode 100644 tensor2tensor/data_generators/video_utils_test.py

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 267ef0fb5..79029b4de 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -19,11 +19,13 @@
 from __future__ import print_function
 
 import os
+import numpy as np
 import six
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.layers import common_video
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import video_metrics
@@ -40,6 +42,33 @@ def resize_video_frames(images, size):
   return resized_images
 
 
+def display_video_hooks(hook_args):
+  """Hooks to display videos at decode time."""
+  predictions = hook_args.predictions
+
+  all_summaries = []
+  for decode_ind, decode in enumerate(predictions):
+
+    target_videos = video_metrics.stack_data_given_key(decode, "targets")
+    output_videos = video_metrics.stack_data_given_key(decode, "outputs")
+    input_videos = video_metrics.stack_data_given_key(decode, "inputs")
+    target_videos = np.asarray(target_videos, dtype=np.uint8)
+    output_videos = np.asarray(output_videos, dtype=np.uint8)
+    input_videos = np.asarray(input_videos, dtype=np.uint8)
+
+    input_videos = np.concatenate((input_videos, target_videos), axis=1)
+    output_videos = np.concatenate((input_videos, output_videos), axis=1)
+    input_summ_vals, _ = common_video.py_gif_summary(
+        "decode_%d/input" % decode_ind, input_videos, max_outputs=10, fps=10,
+        return_summary_value=True)
+    output_summ_vals, _ = common_video.py_gif_summary(
+        "decode_%d/output" % decode_ind, output_videos, max_outputs=10, fps=10,
+        return_summary_value=True)
+    all_summaries.extend(input_summ_vals)
+    all_summaries.extend(output_summ_vals)
+  return all_summaries
+
+
 def summarize_video_metrics(hook_args):
   """Computes video metrics summaries using the decoder output."""
   problem_name = hook_args.problem.name
@@ -149,7 +178,7 @@ def preprocess_example(self, example, mode, hparams):
 
   @property
   def decode_hooks(self):
-    return [summarize_video_metrics]
+    return [summarize_video_metrics, display_video_hooks]
 
   @property
   def is_generate_per_split(self):
diff --git a/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
new file mode 100644
index 000000000..02f467704
--- /dev/null
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -0,0 +1,53 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""video_utils test."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import numpy as np
+from tensor2tensor.data_generators import video_utils
+from tensor2tensor.utils import decoding
+
+import tensorflow as tf
+
+
+class VideoUtilsTest(tf.test.TestCase):
+
+  def testConvertPredictionsToVideoSummaries(self):
+    # Initialize predictions.
+    rng = np.random.RandomState(0)
+    inputs = rng.randint(0, 255, (2, 32, 32, 3))
+    outputs = rng.randint(0, 255, (5, 32, 32, 3))
+    targets = rng.randint(0, 255, (5, 32, 32, 3))
+
+    # batch it up.
+    prediction = [{"outputs": outputs, "inputs": inputs, "targets": targets}]*50
+    predictions = [prediction]
+    decode_hparams = decoding.decode_hparams()
+
+    decode_hooks = decoding.DecodeHookArgs(
+        estimator=None, problem=None, output_dirs=None,
+        hparams=decode_hparams, decode_hparams=decode_hparams,
+        predictions=predictions)
+    summaries = video_utils.display_video_hooks(decode_hooks)
+    # ground_truth + output.
+    self.assertEqual(len(summaries), 20)
+    for summary in summaries:
+      self.assertTrue(isinstance(summary, tf.Summary.Value))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 16a24d4e8..2b5158af5 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -297,7 +297,7 @@ def _encode_gif(images, fps):
   return writer.finish()
 
 
-def py_gif_summary(tag, images, max_outputs, fps):
+def py_gif_summary(tag, images, max_outputs, fps, return_summary_value=False):
   """Outputs a `Summary` protocol buffer with gif animations.
 
   Args:
@@ -305,7 +305,9 @@ def py_gif_summary(tag, images, max_outputs, fps):
     images: A 5-D `uint8` `np.array` of shape `[batch_size, time, height, width,
       channels]` where `channels` is 1 or 3.
     max_outputs: Max number of batch elements to generate gifs for.
-    fps: frames per second of the animation
+    fps: frames per second of the animation.
+    return_summary_value: If set to True, return a list of tf.Summary.Value
+                          objects in addition to the protocol buffer.
 
   Returns:
     The serialized `Summary` protocol buffer.
@@ -323,6 +325,7 @@ def py_gif_summary(tag, images, max_outputs, fps):
     raise ValueError("Tensors must have 1 or 3 channels for gif summary.")
 
   summ = tf.Summary()
+  all_summ_values = []
   num_outputs = min(batch_size, max_outputs)
   for i in range(num_outputs):
     image_summ = tf.Summary.Image()
@@ -350,8 +353,12 @@ def py_gif_summary(tag, images, max_outputs, fps):
       summ_tag = "{}/gif".format(tag)
     else:
       summ_tag = "{}/gif/{}".format(tag, i)
+    curr_summ_value = tf.Summary.Value(tag=summ_tag, image=image_summ)
+    all_summ_values.append(curr_summ_value)
     summ.value.add(tag=summ_tag, image=image_summ)
   summ_str = summ.SerializeToString()
+  if return_summary_value:
+    return all_summ_values, summ_str
   return summ_str
 
 

From f797a588c26a405dbeeebc066b313e88cfeaaea7 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 7 Sep 2018 18:28:19 -0700
Subject: [PATCH 0770/2720] saving a video instead of all the eval pngs.

PiperOrigin-RevId: 212069544
---
 tensor2tensor/data_generators/video_utils.py | 65 ++++++++++----------
 tensor2tensor/layers/common_video.py         |  7 ++-
 2 files changed, 38 insertions(+), 34 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 79029b4de..f665d0ca3 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -352,6 +352,19 @@ def eval_metrics(self):
         metrics.Metrics.NEG_LOG_PERPLEXITY, metrics.Metrics.IMAGE_SUMMARY]
     return eval_metrics
 
+  def validate_frame(self, frame):
+    height, width, channels = frame.shape
+    if channels != self.num_channels:
+      raise ValueError("Generated frame has %d channels while the class "
+                       "assumes %d channels." % (channels,
+                                                 self.num_channels))
+    if height != self.frame_height:
+      raise ValueError("Generated frame has height %d while the class "
+                       "assumes height %d." % (height, self.frame_height))
+    if width != self.frame_width:
+      raise ValueError("Generated frame has width %d while the class "
+                       "assumes width %d." % (width, self.frame_width))
+
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
     """Generate samples of the frames with possible extra data.
 
@@ -391,6 +404,9 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
     Raises:
       ValueError: if the frame has a different number of channels than required.
     """
+    if self.debug_dump_frames_path:
+      writer = common_video.VideoWriter(fps=10, file_format="avi")
+
     with tf.Graph().as_default():
       image_t = tf.placeholder(
           dtype=tf.uint8, shape=(None, None, None))
@@ -398,48 +414,33 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
       with tf.Session() as sess:
         for features in self.generate_samples(data_dir, tmp_dir, dataset_split):
           unencoded_frame = features.pop("frame")
-          height, width, channels = unencoded_frame.shape
-          if channels != self.num_channels:
-            raise ValueError("Generated frame has %d channels while the class "
-                             "assumes %d channels." % (channels,
-                                                       self.num_channels))
-          if height != self.frame_height:
-            raise ValueError("Generated frame has height %d while the class "
-                             "assumes height %d." % (height, self.frame_height))
-          if width != self.frame_width:
-            raise ValueError("Generated frame has width %d while the class "
-                             "assumes width %d." % (width, self.frame_width))
+          self.validate_frame(unencoded_frame)
+          height, width, _ = unencoded_frame.shape
           encoded_frame = sess.run(encoded_image_t, feed_dict={
               image_t: unencoded_frame})
           features["image/encoded"] = [encoded_frame]
           features["image/format"] = ["png"]
           features["image/height"] = [height]
           features["image/width"] = [width]
-          if "image/debug" in features:
+
+          has_debug_image = "image/debug" in features
+          if has_debug_image:
             unencoded_debug = features.pop("image/debug")
             encoded_debug = sess.run(encoded_image_t, feed_dict={
                 image_t: unencoded_debug})
             features["image/encoded_debug"] = [encoded_debug]
+
+          if self.debug_dump_frames_path:
+            img = unencoded_debug if has_debug_image else unencoded_frame
+            writer.write(img)
+
           yield features
 
-  def generate_encoded_samples_debug(self, data_dir, tmp_dir, dataset_split):
-    """Generate samples of the encoded frames and dump for debug if needed."""
-    counter = 0
-    for sample in self.generate_encoded_samples(
-        data_dir, tmp_dir, dataset_split):
-      if self.debug_dump_frames_path:
-        if not tf.gfile.Exists(self.debug_dump_frames_path):
-          tf.gfile.MkDir(self.debug_dump_frames_path)
-        path = os.path.join(self.debug_dump_frames_path,
-                            "frame_%05d.png" % counter)
-        with tf.gfile.Open(path, "wb") as f:
-          if "image/encoded_debug" in sample:
-            img_to_save = sample["image/encoded_debug"][0]
-          else:
-            img_to_save = sample["image/encoded"][0]
-          f.write(img_to_save)
-        counter += 1
-      yield sample
+    if self.debug_dump_frames_path:
+      if not tf.gfile.Exists(self.debug_dump_frames_path):
+        tf.gfile.MkDir(self.debug_dump_frames_path)
+      path = os.path.join(self.debug_dump_frames_path, "video.avi")
+      writer.finish_to_file(path)
 
   def generate_data(self, data_dir, tmp_dir, task_id=-1):
     """The function generating the data."""
@@ -460,12 +461,12 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
     if self.is_generate_per_split:
       for split, paths in split_paths:
         generator_utils.generate_files(
-            self.generate_encoded_samples_debug(
+            self.generate_encoded_samples(
                 data_dir, tmp_dir, split), paths,
             cycle_every_n=self.total_number_of_frames // len(paths))
     else:
       generator_utils.generate_files(
-          self.generate_encoded_samples_debug(
+          self.generate_encoded_samples(
               data_dir, tmp_dir, problem.DatasetSplit.TRAIN),
           all_paths,
           cycle_every_n=self.total_number_of_frames // len(all_paths))
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 2b5158af5..469156f23 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -532,6 +532,7 @@ def __init_ffmpeg(self, image_shape):
         "-filter_complex", "[0:v]split[x][z];[z]palettegen[y];[x][y]paletteuse",
         "-r", "%.02f" % self.fps,
         "-f", self.file_format,
+        "-qscale", "0",
         "-"
     ]
     self.proc = Popen(self.cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
@@ -557,8 +558,10 @@ def finish(self):
     return out
 
   def finish_to_file(self, path):
-    with tf.gfile.open(path) as f:
-      f.write(self.finish())
+    out = self.finish()
+    if out is not None:
+      with tf.gfile.Open(path, "w") as f:
+        f.write(out)
 
   def __del__(self):
     self.finish()

From d8f9b5564b621f9b0e60d9b6175b85278f1e266c Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Fri, 7 Sep 2018 19:07:42 -0700
Subject: [PATCH 0771/2720] fixing epoch_num after restarting

---
 tensor2tensor/rl/rl_trainer_lib.py      | 18 ++++++-----
 tensor2tensor/rl/trainer_model_based.py | 42 +++++++++++++++----------
 2 files changed, 35 insertions(+), 25 deletions(-)

diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index be7544d2c..330b9d5a3 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -41,7 +41,7 @@ def define_train(hparams):
 
 
 def train(hparams, event_dir=None, model_dir=None,
-          restore_agent=True, epoch=0, name_scope="rl_train"):
+          restore_agent=True, name_scope="rl_train"):
   """Train."""
   with tf.Graph().as_default():
     with tf.name_scope(name_scope):
@@ -78,13 +78,15 @@ def train(hparams, event_dir=None, model_dir=None,
           start_step = trainer_lib.restore_checkpoint(
               model_dir, model_saver, sess)
 
-        # Fail-friendly, don't train if already trained for this epoch
-        if start_step >= ((hparams.epochs_num * (epoch + 1))):
-          tf.logging.info("Skipping PPO training for epoch %d as train steps "
-                          "(%d) already reached", epoch, start_step)
+        # Fail-friendly, complete only unfinished epoch
+        steps_to_go = hparams.epochs_num - start_step
+
+        if steps_to_go <= 0:
+          tf.logging.info("Skipping PPO training. Requested %d steps while %d train steps "
+                          "already reached", hparams.epochs_num, start_step)
           return
 
-        for epoch_index in range(hparams.epochs_num):
+        for epoch_index in range(steps_to_go):
           summary = sess.run(train_summary_op)
           if summary_writer:
             summary_writer.add_summary(summary, epoch_index)
@@ -95,8 +97,8 @@ def train(hparams, event_dir=None, model_dir=None,
             else:
               tf.logging.info("Eval summary not saved")
           if (model_saver and hparams.save_models_every_epochs and
-              (epoch_index % hparams.save_models_every_epochs == 0 or
-               (epoch_index + 1) == hparams.epochs_num)):
+              ((epoch_index + start_step) % hparams.save_models_every_epochs == 0 or
+               (epoch_index + 1) == steps_to_go)):
             ckpt_path = os.path.join(
                 model_dir, "model.ckpt-{}".format(epoch_index + 1 + start_step))
             model_saver.save(sess, ckpt_path)
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index ede7d175c..5c4ee0713 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -141,6 +141,20 @@ def train_autoencoder(problem_name, data_dir, output_dir, hparams, epoch):
     t2t_trainer.main([])
 
 
+def _ppo_training_epochs(hparams, epoch, is_final_epoch, real_env_training):
+  real_training_ppo_epochs_num = hparams.real_ppo_epochs_num
+  simulated_training_ppo_epochs_num = hparams.ppo_epochs_num
+
+  ppo_training_epochs = (epoch + 1)*simulated_training_ppo_epochs_num + \
+                        epoch * real_training_ppo_epochs_num
+  if real_env_training:
+    ppo_training_epochs += real_training_ppo_epochs_num
+  if is_final_epoch:
+    ppo_training_epochs += simulated_training_ppo_epochs_num
+
+  return ppo_training_epochs
+
+
 def train_agent(problem_name, agent_model_dir,
                 event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
                 is_final_epoch=False):
@@ -152,15 +166,13 @@ def train_agent(problem_name, agent_model_dir,
                       "optimization_epochs"]
 
   for param_name in ppo_params_names:
-    ppo_param_name = "ppo_"+ param_name
+    ppo_param_name = "ppo_" + param_name
     if ppo_param_name in hparams:
       ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
 
-  ppo_epochs_num = hparams.ppo_epochs_num
-  if is_final_epoch:
-    ppo_epochs_num *= 2
-    ppo_hparams.epoch_length *= 2
-  ppo_hparams.save_models_every_epochs = ppo_epochs_num
+  ppo_hparams.epochs_num = _ppo_training_epochs(hparams, epoch,
+                                                  is_final_epoch, False)
+  ppo_hparams.save_models_every_epochs = 10
   ppo_hparams.world_model_dir = world_model_dir
   ppo_hparams.add_hparam("force_beginning_resets", True)
 
@@ -183,15 +195,14 @@ def train_agent(problem_name, agent_model_dir,
       "output_dir": world_model_dir,
       "data_dir": epoch_data_dir,
   }):
-    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=epoch,
-                         name_scope="ppo_sim")
+    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, name_scope="ppo_sim")
 
 
 def train_agent_real_env(
     problem_name, agent_model_dir, event_dir, world_model_dir, epoch_data_dir,
     hparams, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the real environment."""
-  del epoch, is_final_epoch
+
   gym_problem = registry.problem(problem_name)
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_params_names = ["epochs_num", "epoch_length",
@@ -203,11 +214,9 @@ def train_agent_real_env(
     if ppo_param_name in hparams:
       ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
 
-  ppo_epochs_num = hparams.real_ppo_epochs_num
-  if ppo_epochs_num == 0:
-    return
-
-  ppo_hparams.save_models_every_epochs = ppo_epochs_num
+  ppo_hparams.epochs_num = _ppo_training_epochs(hparams, epoch,
+                                                is_final_epoch, True)
+  ppo_hparams.save_models_every_epochs = 10
 
   environment_spec = copy.copy(gym_problem.environment_spec)
 
@@ -219,8 +228,7 @@ def train_agent_real_env(
       "data_dir": epoch_data_dir,
   }):
     # epoch = 10**20 is a hackish way to avoid skiping training
-    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, epoch=10**20,
-                         name_scope="ppo_real")
+    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, name_scope="ppo_real")
 
 
 def evaluate_world_model(simulated_problem_name, problem_name, hparams,
@@ -694,7 +702,7 @@ def rl_modelrl_tiny():
   """Tiny set for testing."""
   return rl_modelrl_base_sampling().override_from_dict(
       tf.contrib.training.HParams(
-          epochs=1,
+          epochs=2,
           num_real_env_frames=128,
           simulated_env_generator_num_steps=64,
           model_train_steps=2,

From e484aeed793ce00b709066f07e10497074e8a109 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 7 Sep 2018 20:45:00 -0700
Subject: [PATCH 0772/2720] internal merge of PR #1034

PiperOrigin-RevId: 212079013
---
 tensor2tensor/data_generators/gym_problems.py |  1 +
 .../rl/trainer_model_based_ae_test.py         | 37 +++++++++++++++++++
 tensor2tensor/rl/trainer_model_based_test.py  | 16 +-------
 3 files changed, 40 insertions(+), 14 deletions(-)
 create mode 100644 tensor2tensor/rl/trainer_model_based_ae_test.py

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 8a6993f31..e22285e79 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -42,6 +42,7 @@
 
 
 
+
 flags.DEFINE_string("agent_policy_path", None, "File with model for agent.")
 
 flags.DEFINE_string("autoencoder_path", None,
diff --git a/tensor2tensor/rl/trainer_model_based_ae_test.py b/tensor2tensor/rl/trainer_model_based_ae_test.py
new file mode 100644
index 000000000..7888e365c
--- /dev/null
+++ b/tensor2tensor/rl/trainer_model_based_ae_test.py
@@ -0,0 +1,37 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Tiny run of trainer_model_based. Smoke test."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.rl import trainer_model_based
+
+import tensorflow as tf
+
+FLAGS = tf.flags.FLAGS
+
+
+class ModelRLExperimentTestAe(tf.test.TestCase):
+
+  def test_ae(self):
+    FLAGS.output_dir = tf.test.get_temp_dir()
+    FLAGS.loop_hparams_set = "rl_modelrl_ae_tiny"
+    FLAGS.schedule = "train"  # skip evaluation for world model training
+    trainer_model_based.main(None)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/rl/trainer_model_based_test.py b/tensor2tensor/rl/trainer_model_based_test.py
index 67e06427a..48e2f40ea 100644
--- a/tensor2tensor/rl/trainer_model_based_test.py
+++ b/tensor2tensor/rl/trainer_model_based_test.py
@@ -17,9 +17,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import os
-import shutil
-
 from tensor2tensor.rl import trainer_model_based
 
 import tensorflow as tf
@@ -29,19 +26,10 @@
 
 class ModelRLExperimentTest(tf.test.TestCase):
 
-  def setUp(self):
-    super(ModelRLExperimentTest, self).setUp()
-    FLAGS.output_dir = tf.test.get_temp_dir()
-    shutil.rmtree(FLAGS.output_dir)
-    os.mkdir(FLAGS.output_dir)
-    FLAGS.schedule = "train"  # skip evaluation for world model training
-
   def test_basic(self):
+    FLAGS.output_dir = tf.test.get_temp_dir()
     FLAGS.loop_hparams_set = "rl_modelrl_tiny"
-    trainer_model_based.main(None)
-
-  def test_ae(self):
-    FLAGS.loop_hparams_set = "rl_modelrl_ae_tiny"
+    FLAGS.schedule = "train"  # skip evaluation for world model training
     trainer_model_based.main(None)
 
 

From b78598271154be3b0bbff76b28b22f440f858372 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 7 Sep 2018 21:18:12 -0700
Subject: [PATCH 0773/2720] internal merge of PR #1048

PiperOrigin-RevId: 212080800
---
 tensor2tensor/rl/rl_trainer_lib.py      |  8 +++++---
 tensor2tensor/rl/trainer_model_based.py | 12 +++++++-----
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 330b9d5a3..6333febf9 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -82,8 +82,9 @@ def train(hparams, event_dir=None, model_dir=None,
         steps_to_go = hparams.epochs_num - start_step
 
         if steps_to_go <= 0:
-          tf.logging.info("Skipping PPO training. Requested %d steps while %d train steps "
-                          "already reached", hparams.epochs_num, start_step)
+          tf.logging.info("Skipping PPO training. Requested %d steps while "
+                          "%d train steps already reached",
+                          hparams.epochs_num, start_step)
           return
 
         for epoch_index in range(steps_to_go):
@@ -96,8 +97,9 @@ def train(hparams, event_dir=None, model_dir=None,
               summary_writer.add_summary(summary, epoch_index)
             else:
               tf.logging.info("Eval summary not saved")
+          epoch_index_and_start = epoch_index + start_step
           if (model_saver and hparams.save_models_every_epochs and
-              ((epoch_index + start_step) % hparams.save_models_every_epochs == 0 or
+              (epoch_index_and_start % hparams.save_models_every_epochs == 0 or
                (epoch_index + 1) == steps_to_go)):
             ckpt_path = os.path.join(
                 model_dir, "model.ckpt-{}".format(epoch_index + 1 + start_step))
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 534b9fd36..f40c2340f 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -142,6 +142,7 @@ def train_autoencoder(problem_name, data_dir, output_dir, hparams, epoch):
 
 
 def _ppo_training_epochs(hparams, epoch, is_final_epoch, real_env_training):
+  """Helper for PPO restarts."""
   real_training_ppo_epochs_num = hparams.real_ppo_epochs_num
   simulated_training_ppo_epochs_num = hparams.ppo_epochs_num
 
@@ -151,7 +152,6 @@ def _ppo_training_epochs(hparams, epoch, is_final_epoch, real_env_training):
     ppo_training_epochs += real_training_ppo_epochs_num
   if is_final_epoch:
     ppo_training_epochs += simulated_training_ppo_epochs_num
-
   return ppo_training_epochs
 
 
@@ -171,7 +171,7 @@ def train_agent(problem_name, agent_model_dir,
       ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
 
   ppo_hparams.epochs_num = _ppo_training_epochs(hparams, epoch,
-                                                  is_final_epoch, False)
+                                                is_final_epoch, False)
   ppo_hparams.save_models_every_epochs = 10
   ppo_hparams.world_model_dir = world_model_dir
   ppo_hparams.add_hparam("force_beginning_resets", True)
@@ -195,7 +195,8 @@ def train_agent(problem_name, agent_model_dir,
       "output_dir": world_model_dir,
       "data_dir": epoch_data_dir,
   }):
-    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, name_scope="ppo_sim")
+    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir,
+                         name_scope="ppo_sim")
 
 
 def train_agent_real_env(
@@ -228,7 +229,8 @@ def train_agent_real_env(
       "data_dir": epoch_data_dir,
   }):
     # epoch = 10**20 is a hackish way to avoid skiping training
-    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir, name_scope="ppo_real")
+    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir,
+                         name_scope="ppo_real")
 
 
 def evaluate_world_model(simulated_problem_name, problem_name, hparams,
@@ -702,7 +704,7 @@ def rl_modelrl_tiny():
   """Tiny set for testing."""
   return rl_modelrl_base_sampling().override_from_dict(
       tf.contrib.training.HParams(
-          epochs=2,
+          epochs=1,
           num_real_env_frames=128,
           simulated_env_generator_num_steps=64,
           model_train_steps=2,

From 61deed59bdd0b0fd73df311f3c1a5b7c5bba0c31 Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Fri, 7 Sep 2018 21:38:17 -0700
Subject: [PATCH 0774/2720] Optionally prevent repeats in dataset

PiperOrigin-RevId: 212081806
---
 tensor2tensor/data_generators/problem.py      |  7 +++++-
 tensor2tensor/models/research/autoencoders.py | 22 +++++++++++++++++--
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 70e2daa97..7ffdac942 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -721,6 +721,7 @@ def make_estimator_input_fn(self,
                               hparams,
                               data_dir=None,
                               force_repeat=False,
+                              prevent_repeat=False,
                               dataset_kwargs=None):
     """Return input_fn wrapped for Estimator."""
 
@@ -732,6 +733,7 @@ def estimator_input_fn(params, config):
           params=params,
           config=config,
           force_repeat=force_repeat,
+          prevent_repeat=prevent_repeat,
           dataset_kwargs=dataset_kwargs)
 
     return estimator_input_fn
@@ -777,6 +779,7 @@ def input_fn(self,
                params=None,
                config=None,
                force_repeat=False,
+               prevent_repeat=False,
                dataset_kwargs=None):
     """Builds input pipeline for problem.
 
@@ -788,6 +791,8 @@ def input_fn(self,
       config: RunConfig; should have the data_parallelism attribute if not using
         TPU
       force_repeat: bool, whether to repeat the data even if not training
+      prevent_repeat: bool, whether to not repeat when in training mode.
+        Overrides force_repeat.
       dataset_kwargs: dict, if passed, will pass as kwargs to self.dataset
         method when called
 
@@ -832,7 +837,7 @@ def define_shapes(example):
     })
 
     dataset = self.dataset(**dataset_kwargs)
-    if force_repeat or is_training:
+    if (force_repeat or is_training) and not prevent_repeat:
       # Repeat and skip a random number of records
       dataset = dataset.repeat()
 
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 7d3df327a..aba618855 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -168,6 +168,10 @@ def body(self, features):
     self.is1d = hparams.sample_width == 1
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
       labels = features["targets_raw"]
+      labels_shape = common_layers.shape_list(labels)
+      # handle videos
+      if len(labels.shape) == 5:
+        labels = common_layers.time_to_channels(labels)
       shape = common_layers.shape_list(labels)
       x = tf.one_hot(labels, vocab_size)
       x = self.embed(x)
@@ -216,6 +220,7 @@ def body(self, features):
         b = self.sample()
       else:
         b = self._cur_bottleneck_tensor
+      self._cur_bottleneck_tensor = b
       res_size = self.hparams.hidden_size * 2**self.hparams.num_hidden_layers
       res_size = min(res_size, hparams.max_hidden_size)
       x = self.unbottleneck(b, res_size)
@@ -267,7 +272,7 @@ def body(self, features):
     else:
       reconstr = tf.layers.dense(res, vocab_size, name="autoencoder_final")
       targets_loss = tf.losses.sparse_softmax_cross_entropy(
-          logits=reconstr, labels=labels)
+          logits=reconstr, labels=tf.reshape(labels, labels_shape))
       losses["training"] = targets_loss
 
     # GAN losses.
@@ -339,7 +344,8 @@ def discriminate(x):
       losses["gan_loss"] = -gan_loss
 
     self.image_summary("ae", reconstr)
-    logits = reconstr
+
+    logits = tf.reshape(reconstr, labels_shape + [vocab_size])
     return logits, losses
 
   def sample(self, features=None, shape=None):
@@ -1037,6 +1043,18 @@ def autoencoder_ordered_discrete():
   return hparams
 
 
+@registry.register_hparams
+def autoencoder_ordered_discrete_image64():
+  """Ordered discrete autoencoder model."""
+  hparams = autoencoder_ordered_discrete()
+  hparams.batch_size = 32
+  hparams.num_hidden_layers = 6
+  hparams.target_modality = "video:default"
+  hparams.input_modalities = "video:default"
+
+  return hparams
+
+
 @registry.register_hparams
 def autoencoder_ordered_discrete_patched():
   """Ordered discrete autoencoder model."""

From bfcb536ecbc6a9bad1ad914f781bf56fd9fcf061 Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Sat, 8 Sep 2018 00:09:14 -0700
Subject: [PATCH 0775/2720] Correct shapes for video

PiperOrigin-RevId: 212089692
---
 tensor2tensor/models/research/autoencoders.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index aba618855..5580893cb 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -272,7 +272,8 @@ def body(self, features):
     else:
       reconstr = tf.layers.dense(res, vocab_size, name="autoencoder_final")
       targets_loss = tf.losses.sparse_softmax_cross_entropy(
-          logits=reconstr, labels=tf.reshape(labels, labels_shape))
+          logits=tf.reshape(reconstr, labels_shape + [vocab_size]),
+          labels=tf.reshape(labels, labels_shape))
       losses["training"] = targets_loss
 
     # GAN losses.
@@ -338,8 +339,11 @@ def discriminate(x):
           hparams.gan_codes_warmup_steps * 1.5)
       rev_grad_gan_codes = reverse_gradient(gan_codes, lr=gan_lr)
       gan_loss = common_layers.sliced_gan_loss(
-          target_codes, rev_grad_gan_codes, discriminate,
-          self.hparams.num_sliced_vecs, do_tanh=hparams.sliced_do_tanh)
+          target_codes,
+          rev_grad_gan_codes,
+          discriminate,
+          self.hparams.num_sliced_vecs,
+          do_tanh=hparams.sliced_do_tanh)
       gan_loss *= hparams.gan_loss_factor * update_means_factor
       losses["gan_loss"] = -gan_loss
 
@@ -544,8 +548,8 @@ def infer(self, features, *args, **kwargs):
         samples = common_layers.sample_with_temperature(
             logits, self.hparams.sampling_temp)
         samples1d = tf.reshape(samples, [shape[0], -1, shape[3]])
-        samples1d = tf.concat(
-            [old_samples1d[:, :i, :], samples1d[:, i:, :]], axis=1)
+        samples1d = tf.concat([old_samples1d[:, :i, :], samples1d[:, i:, :]],
+                              axis=1)
         samples = tf.reshape(samples1d, shape)
 
     # Restore inputs to not confuse Estimator in edge cases.

From 0a34a1e99799b7af3b91dcaa8e50aa5077bf5bcf Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 8 Sep 2018 14:00:01 -0700
Subject: [PATCH 0776/2720] debug env

---
 .../rl/model_rl_experiment_player.py          | 90 +++++++++++++++++++
 1 file changed, 90 insertions(+)
 create mode 100644 tensor2tensor/rl/model_rl_experiment_player.py

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
new file mode 100644
index 000000000..f41505fe9
--- /dev/null
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -0,0 +1,90 @@
+# This is a starter for piotrmilos experiments. Should not be committed into repo
+# The intention is to keep my dirty configs outside of the repo.
+
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.rl.model_rl_experiment import create_loop_hparams, train, rl_modelrl_base
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+@registry.register_hparams
+def pm_rl_modelrl_tiny():
+  """Tiny set for testing."""
+  tiny_hp = tf.contrib.training.HParams(
+      epochs=2,
+      true_env_generator_num_steps=20,
+      model_train_steps=10,
+      simulated_env_generator_num_steps=20,
+      ppo_epochs_num=2,
+      ppo_time_limit=20,
+      ppo_epoch_length=20,
+
+  )
+  return rl_modelrl_base().override_from_dict(tiny_hp.values())
+
+@registry.register_hparams
+def pm_rl_modelrl_tiny_2agents():
+  """Tiny set for testing."""
+  tiny_hp = tf.contrib.training.HParams(
+      epochs=2,
+      true_env_generator_num_steps=200,
+      model_train_steps=2,
+      simulated_env_generator_num_steps=20,
+      ppo_epochs_num=2,
+      ppo_time_limit=20,
+      ppo_epoch_length=20,
+      ppo_num_agents=2
+
+  )
+  return rl_modelrl_base().override_from_dict(tiny_hp.values())
+
+
+@registry.register_hparams
+def pm_rl_modelrl_longpong_tiny():
+  """Tiny set for testing."""
+  tiny_hp = tf.contrib.training.HParams(
+      epochs=2,
+      true_env_generator_num_steps=20,
+      model_train_steps=10,
+      simulated_env_generator_num_steps=20,
+      ppo_epochs_num=2,
+      ppo_time_limit=20,
+      #The same as GymWrappedLongPongRandom.num_testing_steps
+      #both should be roughly similar
+      ppo_epoch_length=100,
+      game="wrapped_long_pong",
+
+  )
+  return rl_modelrl_base().override_from_dict(tiny_hp.values())
+
+
+@registry.register_hparams
+def pm_rl_modelrl_medium():
+  """Tiny set for testing."""
+  tiny_hp = tf.contrib.training.HParams(
+      epochs=2,
+      true_env_generator_num_steps=50000,
+      model_train_steps=15000,
+      simulated_env_generator_num_steps=10000,
+      ppo_epochs_num=2,
+      ppo_time_limit=20,
+      ppo_epoch_length=20,
+  )
+  return rl_modelrl_base().override_from_dict(tiny_hp.values())
+
+def main(_):
+  hp = create_loop_hparams()
+  output_dir = FLAGS.output_dir
+  train(hp, output_dir)
+
+
+if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
+  tf.app.run()

From 5bb487eae577258ac41cc3598f67e10dbe10a2f4 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 8 Sep 2018 14:10:10 -0700
Subject: [PATCH 0777/2720] added probs and vf

---
 .../rl/model_rl_experiment_player.py          | 433 +++++++++++++++---
 1 file changed, 361 insertions(+), 72 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index f41505fe9..d31b010f7 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -1,88 +1,377 @@
-# This is a starter for piotrmilos experiments. Should not be committed into repo
-# The intention is to keep my dirty configs outside of the repo.
-
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.rl.model_rl_experiment import create_loop_hparams, train, rl_modelrl_base
+import contextlib
+import copy
+import datetime
+import math
+import os
+import time
+
+import six
+
+from tensor2tensor.bin import t2t_trainer
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import gym_problems_specs
+from tensor2tensor.layers import discretization
+from tensor2tensor.rl import rl_trainer_lib
+from tensor2tensor.rl.envs.batch_env_factory import batch_env_factory
+from tensor2tensor.rl.envs.utils import get_policy
 from tensor2tensor.utils import registry
+from tensor2tensor.utils import trainer_lib
+from tensor2tensor.rl.trainer_model_based import FLAGS
 
 import tensorflow as tf
 
-flags = tf.flags
-FLAGS = flags.FLAGS
-
-@registry.register_hparams
-def pm_rl_modelrl_tiny():
-  """Tiny set for testing."""
-  tiny_hp = tf.contrib.training.HParams(
-      epochs=2,
-      true_env_generator_num_steps=20,
-      model_train_steps=10,
-      simulated_env_generator_num_steps=20,
-      ppo_epochs_num=2,
-      ppo_time_limit=20,
-      ppo_epoch_length=20,
-
-  )
-  return rl_modelrl_base().override_from_dict(tiny_hp.values())
-
-@registry.register_hparams
-def pm_rl_modelrl_tiny_2agents():
-  """Tiny set for testing."""
-  tiny_hp = tf.contrib.training.HParams(
-      epochs=2,
-      true_env_generator_num_steps=200,
-      model_train_steps=2,
-      simulated_env_generator_num_steps=20,
-      ppo_epochs_num=2,
-      ppo_time_limit=20,
-      ppo_epoch_length=20,
-      ppo_num_agents=2
-
-  )
-  return rl_modelrl_base().override_from_dict(tiny_hp.values())
-
-
-@registry.register_hparams
-def pm_rl_modelrl_longpong_tiny():
-  """Tiny set for testing."""
-  tiny_hp = tf.contrib.training.HParams(
-      epochs=2,
-      true_env_generator_num_steps=20,
-      model_train_steps=10,
-      simulated_env_generator_num_steps=20,
-      ppo_epochs_num=2,
-      ppo_time_limit=20,
-      #The same as GymWrappedLongPongRandom.num_testing_steps
-      #both should be roughly similar
-      ppo_epoch_length=100,
-      game="wrapped_long_pong",
-
-  )
-  return rl_modelrl_base().override_from_dict(tiny_hp.values())
-
-
-@registry.register_hparams
-def pm_rl_modelrl_medium():
-  """Tiny set for testing."""
-  tiny_hp = tf.contrib.training.HParams(
-      epochs=2,
-      true_env_generator_num_steps=50000,
-      model_train_steps=15000,
-      simulated_env_generator_num_steps=10000,
-      ppo_epochs_num=2,
-      ppo_time_limit=20,
-      ppo_epoch_length=20,
-  )
-  return rl_modelrl_base().override_from_dict(tiny_hp.values())
+
+HP_SCOPES = ["loop", "model", "ppo"]
+
+
+def setup_directories(base_dir, subdirs):
+  base_dir = os.path.expanduser(base_dir)
+  tf.gfile.MakeDirs(base_dir)
+
+  all_dirs = {}
+  for subdir in subdirs:
+    dir_name = os.path.join(base_dir, subdir)
+    tf.gfile.MakeDirs(dir_name)
+    all_dirs[subdir] = dir_name
+  return all_dirs
+
+
+def make_relative_timing_fn():
+  """Make a function that logs the duration since it was made."""
+  start_time = time.time()
+
+  def format_relative_time():
+    time_delta = time.time() - start_time
+    return str(datetime.timedelta(seconds=time_delta))
+
+  def log_relative_time():
+    tf.logging.info("Timing: %s", format_relative_time())
+
+  return log_relative_time
+
+
+@contextlib.contextmanager
+def temporary_flags(flag_settings):
+  old_values = {}
+  for flag_name, flag_value in flag_settings.items():
+    old_values[flag_name] = getattr(FLAGS, flag_name)
+    setattr(FLAGS, flag_name, flag_value)
+  yield
+  for flag_name, flag_value in old_values.items():
+    setattr(FLAGS, flag_name, flag_value)
+
+
+def make_log_fn(epoch, log_relative_time_fn):
+
+  def log(msg, *args):
+    msg %= args
+    tf.logging.info("%s Epoch %d: %s", ">>>>>>>", epoch, msg)
+    log_relative_time_fn()
+
+  return log
+
+
+
+def train_agent(problem_name, agent_model_dir,
+                event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
+                is_final_epoch=False):
+  """Train the PPO agent in the simulated environment."""
+  gym_problem = registry.problem(problem_name)
+  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
+  ppo_params_names = ["epochs_num", "epoch_length",
+                      "learning_rate", "num_agents",
+                      "optimization_epochs"]
+
+  for param_name in ppo_params_names:
+    ppo_param_name = "ppo_" + param_name
+    if ppo_param_name in hparams:
+      ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
+
+  # ppo_hparams.epochs_num = _ppo_training_epochs(hparams, epoch,
+  #                                                 is_final_epoch, False)
+  ppo_hparams.save_models_every_epochs = 10
+  ppo_hparams.world_model_dir = world_model_dir
+  ppo_hparams.add_hparam("force_beginning_resets", True)
+
+  # Adding model hparams for model specific adjustments
+  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
+  ppo_hparams.add_hparam("model_hparams", model_hparams)
+
+  environment_spec = copy.copy(gym_problem.environment_spec)
+  environment_spec.simulation_random_starts = hparams.simulation_random_starts
+  environment_spec.simulation_flip_first_random_for_beginning = False
+  environment_spec.intrinsic_reward_scale = hparams.intrinsic_reward_scale
+
+  ppo_hparams.add_hparam("environment_spec", environment_spec)
+  ppo_hparams.num_agents = 1
+
+  with temporary_flags({
+      "problem": problem_name,
+      "model": hparams.generative_model,
+      "hparams_set": hparams.generative_model_params,
+      "output_dir": world_model_dir,
+      "data_dir": epoch_data_dir,
+  }):
+
+
+    sess = tf.Session()
+    env = DebugBatchEnv(ppo_hparams, sess)
+    sess.run(tf.global_variables_initializer())
+    env.initialize()
+
+    r = env.step(0)
+    r = env.reset()
+    print("R:{}".format(r))
+
+from gym.core import Env
+
+
+class DebugBatchEnv(Env):
+
+  def __init__(self, hparams, sess = None):
+    if sess == None:
+      self.sess = tf.Session()
+    else:
+      self.sess = sess
+
+    batch_env = batch_env_factory(hparams)
+
+    self.action = tf.placeholder(shape=(1,), dtype=tf.int32)
+
+    self.reward, self.done = batch_env.simulate(self.action)
+    self.observation = batch_env.observ
+    self.reset_op = batch_env.reset(tf.constant([0], dtype=tf.int32))
+
+    environment_wrappers = hparams.environment_spec.wrappers
+    wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
+
+    to_initialize = [batch_env]
+    for w in wrappers:
+      batch_env = w[0](batch_env, **w[1])
+      to_initialize.append(batch_env)
+
+    def initialization_lambda():
+      for batch_env in to_initialize:
+        batch_env.initialize(sess)
+
+    self.initialize = initialization_lambda
+
+    obs_copy = batch_env.observ + 0
+
+    actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
+    self.policy_probs = actor_critic.policy.probs[0, 0, :]
+    self.value = actor_critic.value[0, :]
+    x = 1
+
+  def render(self, mode='human'):
+    raise NotImplemented()
+
+  def reset(self):
+    observ = self.sess.run(self.reset_op)
+    return observ
+
+  def step(self, action):
+    observ, rew, done, probs, vf = self.sess.\
+      run([self.observation, self.reward, self.done, self.policy_probs, self.value],
+          feed_dict={self.action: [action]})
+
+    return observ[0, ...], rew[0, ...], done[0, ...], probs, vf
+
+
+
+def encode_dataset(model, dataset, problem, ae_hparams, autoencoder_path,
+                   out_files):
+  """Encode all frames in dataset with model and write them out to out_files."""
+  batch_size = 8
+  dataset = dataset.batch(batch_size)
+  examples = dataset.make_one_shot_iterator().get_next()
+  images = examples.pop("frame")
+  images = tf.expand_dims(images, 1)
+
+  encoded = model.encode(images)
+  encoded_frame_height = int(
+      math.ceil(problem.frame_height / 2**ae_hparams.num_hidden_layers))
+  encoded_frame_width = int(
+      math.ceil(problem.frame_width / 2**ae_hparams.num_hidden_layers))
+  num_bits = 8
+  encoded = tf.reshape(
+      encoded, [-1, encoded_frame_height, encoded_frame_width, 3, num_bits])
+  encoded = tf.cast(discretization.bit_to_int(encoded, num_bits), tf.uint8)
+
+  pngs = tf.map_fn(tf.image.encode_png, encoded, dtype=tf.string,
+                   back_prop=False)
+
+  with tf.Session() as sess:
+    autoencoder_saver = tf.train.Saver(tf.global_variables("autoencoder.*"))
+    trainer_lib.restore_checkpoint(autoencoder_path, autoencoder_saver, sess,
+                                   must_restore=True)
+
+    def generator():
+      """Generate examples."""
+      while True:
+        try:
+          pngs_np, examples_np = sess.run([pngs, examples])
+          rewards = examples_np["reward"].tolist()
+          actions = examples_np["action"].tolist()
+          frame_numbers = examples_np["frame_number"].tolist()
+          for action, reward, frame_number, png in \
+                  zip(actions, rewards, frame_numbers, pngs_np):
+            yield {
+                "action": action,
+                "reward": reward,
+                "frame_number": frame_number,
+                "image/encoded": [png],
+                "image/format": ["png"],
+                "image/height": [encoded_frame_height],
+                "image/width": [encoded_frame_width],
+            }
+        except tf.errors.OutOfRangeError:
+          break
+
+    generator_utils.generate_files(
+        generator(), out_files,
+        cycle_every_n=problem.total_number_of_frames // 10)
+
+
+def encode_env_frames(problem_name, ae_problem_name, autoencoder_path,
+                      epoch_data_dir):
+  """Encode all frames from problem_name and write out as ae_problem_name."""
+  with tf.Graph().as_default():
+    ae_hparams = trainer_lib.create_hparams("autoencoder_discrete_pong",
+                                            problem_name=problem_name)
+    problem = ae_hparams.problem
+    model = registry.model("autoencoder_ordered_discrete")(
+        ae_hparams, tf.estimator.ModeKeys.EVAL)
+
+    ae_problem = registry.problem(ae_problem_name)
+    ae_training_paths = ae_problem.training_filepaths(epoch_data_dir, 10, True)
+    ae_eval_paths = ae_problem.dev_filepaths(epoch_data_dir, 1, True)
+
+    skip_train = False
+    skip_eval = False
+    for path in ae_training_paths:
+      if tf.gfile.Exists(path):
+        skip_train = True
+        break
+    for path in ae_eval_paths:
+      if tf.gfile.Exists(path):
+        skip_eval = True
+        break
+
+    # Encode train data
+    if not skip_train:
+      dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, epoch_data_dir,
+                                shuffle_files=False, output_buffer_size=100,
+                                preprocess=False)
+      encode_dataset(model, dataset, problem, ae_hparams, autoencoder_path,
+                     ae_training_paths)
+
+    # Encode eval data
+    if not skip_eval:
+      dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, epoch_data_dir,
+                                shuffle_files=False, output_buffer_size=100,
+                                preprocess=False)
+      encode_dataset(model, dataset, problem, ae_hparams, autoencoder_path,
+                     ae_eval_paths)
+
+
+def check_problems(problem_names):
+  for problem_name in problem_names:
+    registry.problem(problem_name)
+
+
+def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
+  """Run the main training loop."""
+  if report_fn:
+    assert report_metric is not None
+
+  # Global state
+
+  # Directories
+  subdirectories = ["data", "tmp", "world_model", "ppo"]
+  using_autoencoder = hparams.autoencoder_train_steps > 0
+  if using_autoencoder:
+    subdirectories.append("autoencoder")
+  directories = setup_directories(output_dir, subdirectories)
+
+  if hparams.game in gym_problems_specs.ATARI_GAMES:
+    game_with_mode = hparams.game + "_deterministic-v4"
+  else:
+    game_with_mode = hparams.game
+  # Problems
+  if using_autoencoder:
+    problem_name = (
+        "gym_discrete_problem_with_agent_on_%s_with_autoencoder"
+        % game_with_mode)
+    world_model_problem = (
+        "gym_discrete_problem_with_agent_on_%s_autoencoded" % game_with_mode)
+    simulated_problem_name = (
+        "gym_simulated_discrete_problem_with_agent_on_%s_autoencoded"
+        % game_with_mode)
+  else:
+    problem_name = ("gym_discrete_problem_with_agent_on_%s" % game_with_mode)
+    world_model_problem = problem_name
+    simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
+                              % game_with_mode)
+    if problem_name not in registry.list_problems():
+      tf.logging.info("Game Problem %s not found; dynamically registering",
+                      problem_name)
+      gym_problems_specs.create_problems_for_game(hparams.game,
+                                                  game_mode="Deterministic-v4")
+
+  # Autoencoder model dir
+  autoencoder_model_dir = directories.get("autoencoder")
+
+  # Timing log function
+  log_relative_time = make_relative_timing_fn()
+
+  # Per-epoch state
+  epoch_metrics = []
+  epoch_data_dirs = []
+
+
+  # Collect data from the real environment with random policy
+  data_dir = os.path.join(directories["data"], "random")
+  epoch_data_dirs.append(data_dir)
+
+  for epoch in range(hparams.epochs):
+    is_final_epoch = (epoch + 1) == hparams.epochs
+    # log = make_log_fn(epoch, log_relative_time)
+
+    # Combine all previously collected environment data
+    epoch_data_dir = os.path.join(directories["data"], str(epoch))
+
+
+    ppo_event_dir = os.path.join(directories["world_model"],
+                                 "ppo_summaries", str(epoch))
+    ppo_model_dir = directories["ppo"]
+    if not hparams.ppo_continue_training:
+      ppo_model_dir = ppo_event_dir
+    train_agent(simulated_problem_name, ppo_model_dir,
+                ppo_event_dir, directories["world_model"], epoch_data_dir,
+                hparams, epoch=epoch, is_final_epoch=is_final_epoch)
+
+
+  raise NotImplementedError()
+  return 1
+
+
+
+def create_loop_hparams():
+  hparams = registry.hparams(FLAGS.loop_hparams_set)
+  hparams.parse(FLAGS.loop_hparams)
+  return hparams
+
 
 def main(_):
   hp = create_loop_hparams()
   output_dir = FLAGS.output_dir
-  train(hp, output_dir)
+  training_loop(hp, output_dir)
 
 
 if __name__ == "__main__":

From 7162e34c40b10048367cc337191dcf92bea72b33 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 8 Sep 2018 14:24:51 -0700
Subject: [PATCH 0778/2720] some cleanups

---
 .../rl/model_rl_experiment_player.py          | 101 ------------------
 1 file changed, 101 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index d31b010f7..b92bbb8d6 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -184,107 +184,6 @@ def step(self, action):
 
 
 
-def encode_dataset(model, dataset, problem, ae_hparams, autoencoder_path,
-                   out_files):
-  """Encode all frames in dataset with model and write them out to out_files."""
-  batch_size = 8
-  dataset = dataset.batch(batch_size)
-  examples = dataset.make_one_shot_iterator().get_next()
-  images = examples.pop("frame")
-  images = tf.expand_dims(images, 1)
-
-  encoded = model.encode(images)
-  encoded_frame_height = int(
-      math.ceil(problem.frame_height / 2**ae_hparams.num_hidden_layers))
-  encoded_frame_width = int(
-      math.ceil(problem.frame_width / 2**ae_hparams.num_hidden_layers))
-  num_bits = 8
-  encoded = tf.reshape(
-      encoded, [-1, encoded_frame_height, encoded_frame_width, 3, num_bits])
-  encoded = tf.cast(discretization.bit_to_int(encoded, num_bits), tf.uint8)
-
-  pngs = tf.map_fn(tf.image.encode_png, encoded, dtype=tf.string,
-                   back_prop=False)
-
-  with tf.Session() as sess:
-    autoencoder_saver = tf.train.Saver(tf.global_variables("autoencoder.*"))
-    trainer_lib.restore_checkpoint(autoencoder_path, autoencoder_saver, sess,
-                                   must_restore=True)
-
-    def generator():
-      """Generate examples."""
-      while True:
-        try:
-          pngs_np, examples_np = sess.run([pngs, examples])
-          rewards = examples_np["reward"].tolist()
-          actions = examples_np["action"].tolist()
-          frame_numbers = examples_np["frame_number"].tolist()
-          for action, reward, frame_number, png in \
-                  zip(actions, rewards, frame_numbers, pngs_np):
-            yield {
-                "action": action,
-                "reward": reward,
-                "frame_number": frame_number,
-                "image/encoded": [png],
-                "image/format": ["png"],
-                "image/height": [encoded_frame_height],
-                "image/width": [encoded_frame_width],
-            }
-        except tf.errors.OutOfRangeError:
-          break
-
-    generator_utils.generate_files(
-        generator(), out_files,
-        cycle_every_n=problem.total_number_of_frames // 10)
-
-
-def encode_env_frames(problem_name, ae_problem_name, autoencoder_path,
-                      epoch_data_dir):
-  """Encode all frames from problem_name and write out as ae_problem_name."""
-  with tf.Graph().as_default():
-    ae_hparams = trainer_lib.create_hparams("autoencoder_discrete_pong",
-                                            problem_name=problem_name)
-    problem = ae_hparams.problem
-    model = registry.model("autoencoder_ordered_discrete")(
-        ae_hparams, tf.estimator.ModeKeys.EVAL)
-
-    ae_problem = registry.problem(ae_problem_name)
-    ae_training_paths = ae_problem.training_filepaths(epoch_data_dir, 10, True)
-    ae_eval_paths = ae_problem.dev_filepaths(epoch_data_dir, 1, True)
-
-    skip_train = False
-    skip_eval = False
-    for path in ae_training_paths:
-      if tf.gfile.Exists(path):
-        skip_train = True
-        break
-    for path in ae_eval_paths:
-      if tf.gfile.Exists(path):
-        skip_eval = True
-        break
-
-    # Encode train data
-    if not skip_train:
-      dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, epoch_data_dir,
-                                shuffle_files=False, output_buffer_size=100,
-                                preprocess=False)
-      encode_dataset(model, dataset, problem, ae_hparams, autoencoder_path,
-                     ae_training_paths)
-
-    # Encode eval data
-    if not skip_eval:
-      dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, epoch_data_dir,
-                                shuffle_files=False, output_buffer_size=100,
-                                preprocess=False)
-      encode_dataset(model, dataset, problem, ae_hparams, autoencoder_path,
-                     ae_eval_paths)
-
-
-def check_problems(problem_names):
-  for problem_name in problem_names:
-    registry.problem(problem_name)
-
-
 def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   """Run the main training loop."""
   if report_fn:

From f0adb3dae749cde738094fe2539e3e8da67fb83c Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 8 Sep 2018 14:52:27 -0700
Subject: [PATCH 0779/2720] cleanups and fake data

---
 .../rl/model_rl_experiment_player.py          | 30 +++++++++++--------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index b92bbb8d6..9dfcf9d2a 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -66,17 +66,6 @@ def temporary_flags(flag_settings):
     setattr(FLAGS, flag_name, flag_value)
 
 
-def make_log_fn(epoch, log_relative_time_fn):
-
-  def log(msg, *args):
-    msg %= args
-    tf.logging.info("%s Epoch %d: %s", ">>>>>>>", epoch, msg)
-    log_relative_time_fn()
-
-  return log
-
-
-
 def train_agent(problem_name, agent_model_dir,
                 event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
                 is_final_epoch=False):
@@ -175,6 +164,24 @@ def reset(self):
     observ = self.sess.run(self.reset_op)
     return observ
 
+
+  def _step_fake(self, action):
+    import numpy as np
+    observ = np.zeros(shape=(210, 160, 3), dtype=np.uint8)
+    rew = 1
+    done = False
+    probs = np.ones(shape=(6,), dtype=np.float32)/6
+    vf = 0.0
+
+    return observ, rew, done, probs, vf
+
+  def _env_step_fake(self, action):
+    observ, rew, done, probs, vf = self.sess.\
+      run([self.observation, self.reward, self.done, self.policy_probs, self.value],
+          feed_dict={self.action: [action]})
+
+    return observ[0, ...], rew[0, ...], done[0, ...], probs, vf
+
   def step(self, action):
     observ, rew, done, probs, vf = self.sess.\
       run([self.observation, self.reward, self.done, self.policy_probs, self.value],
@@ -240,7 +247,6 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   for epoch in range(hparams.epochs):
     is_final_epoch = (epoch + 1) == hparams.epochs
-    # log = make_log_fn(epoch, log_relative_time)
 
     # Combine all previously collected environment data
     epoch_data_dir = os.path.join(directories["data"], str(epoch))

From 75149abed873bdeacf9832dffd77e78d729711a7 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 8 Sep 2018 16:56:41 -0700
Subject: [PATCH 0780/2720] more cleaning

---
 .../rl/model_rl_experiment_player.py          | 75 ++++++-------------
 1 file changed, 21 insertions(+), 54 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index 9dfcf9d2a..b5140c2d0 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -9,8 +9,12 @@
 import math
 import os
 import time
+from gym.spaces import Box
+import numpy as np
 
 import six
+from gym.spaces import Discrete
+from gym.utils.play import PlayPlot
 
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.data_generators import generator_utils
@@ -21,7 +25,7 @@
 from tensor2tensor.rl.envs.utils import get_policy
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
-from tensor2tensor.rl.trainer_model_based import FLAGS
+from tensor2tensor.rl.trainer_model_based import FLAGS, setup_directories, temporary_flags
 
 import tensorflow as tf
 
@@ -29,43 +33,6 @@
 HP_SCOPES = ["loop", "model", "ppo"]
 
 
-def setup_directories(base_dir, subdirs):
-  base_dir = os.path.expanduser(base_dir)
-  tf.gfile.MakeDirs(base_dir)
-
-  all_dirs = {}
-  for subdir in subdirs:
-    dir_name = os.path.join(base_dir, subdir)
-    tf.gfile.MakeDirs(dir_name)
-    all_dirs[subdir] = dir_name
-  return all_dirs
-
-
-def make_relative_timing_fn():
-  """Make a function that logs the duration since it was made."""
-  start_time = time.time()
-
-  def format_relative_time():
-    time_delta = time.time() - start_time
-    return str(datetime.timedelta(seconds=time_delta))
-
-  def log_relative_time():
-    tf.logging.info("Timing: %s", format_relative_time())
-
-  return log_relative_time
-
-
-@contextlib.contextmanager
-def temporary_flags(flag_settings):
-  old_values = {}
-  for flag_name, flag_value in flag_settings.items():
-    old_values[flag_name] = getattr(FLAGS, flag_name)
-    setattr(FLAGS, flag_name, flag_value)
-  yield
-  for flag_name, flag_value in old_values.items():
-    setattr(FLAGS, flag_name, flag_value)
-
-
 def train_agent(problem_name, agent_model_dir,
                 event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
                 is_final_epoch=False):
@@ -81,8 +48,6 @@ def train_agent(problem_name, agent_model_dir,
     if ppo_param_name in hparams:
       ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
 
-  # ppo_hparams.epochs_num = _ppo_training_epochs(hparams, epoch,
-  #                                                 is_final_epoch, False)
   ppo_hparams.save_models_every_epochs = 10
   ppo_hparams.world_model_dir = world_model_dir
   ppo_hparams.add_hparam("force_beginning_resets", True)
@@ -113,9 +78,9 @@ def train_agent(problem_name, agent_model_dir,
     sess.run(tf.global_variables_initializer())
     env.initialize()
 
-    r = env.step(0)
-    r = env.reset()
-    print("R:{}".format(r))
+    key_mapping = {(ord('q'),):1, (ord('a'),):2}
+    from gym.utils import play
+    play.play(env, zoom=4, fps=40, keys_to_action=key_mapping)
 
 from gym.core import Env
 
@@ -128,6 +93,9 @@ def __init__(self, hparams, sess = None):
     else:
       self.sess = sess
 
+    self.action_space = Discrete(6)
+    self.observation_space = Box(low=0, high=255, shape=(210, 160, 3), dtype=np.uint8)
+
     batch_env = batch_env_factory(hparams)
 
     self.action = tf.placeholder(shape=(1,), dtype=tf.int32)
@@ -155,7 +123,7 @@ def initialization_lambda():
     actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
     self.policy_probs = actor_critic.policy.probs[0, 0, :]
     self.value = actor_critic.value[0, :]
-    x = 1
+    self._tmp = 1
 
   def render(self, mode='human'):
     raise NotImplemented()
@@ -166,8 +134,13 @@ def reset(self):
 
 
   def _step_fake(self, action):
-    import numpy as np
-    observ = np.zeros(shape=(210, 160, 3), dtype=np.uint8)
+
+    print("Action:{}".format(action))
+    observ = np.ones(shape=(210, 160, 3), dtype=np.uint8)*10*self._tmp
+    self._tmp += 1
+    if self._tmp>20:
+      self._tmp = 0
+
     rew = 1
     done = False
     probs = np.ones(shape=(6,), dtype=np.float32)/6
@@ -183,11 +156,8 @@ def _env_step_fake(self, action):
     return observ[0, ...], rew[0, ...], done[0, ...], probs, vf
 
   def step(self, action):
-    observ, rew, done, probs, vf = self.sess.\
-      run([self.observation, self.reward, self.done, self.policy_probs, self.value],
-          feed_dict={self.action: [action]})
-
-    return observ[0, ...], rew[0, ...], done[0, ...], probs, vf
+    observ, rew, done, probs, vf = self._step_fake(action)
+    return observ, rew, done, {"probs": probs, "vf": vf}
 
 
 
@@ -233,9 +203,6 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   # Autoencoder model dir
   autoencoder_model_dir = directories.get("autoencoder")
 
-  # Timing log function
-  log_relative_time = make_relative_timing_fn()
-
   # Per-epoch state
   epoch_metrics = []
   epoch_data_dirs = []

From 32d227325a984a0045c30ebd4abd909db6695566 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 8 Sep 2018 17:46:51 -0700
Subject: [PATCH 0781/2720] some debug information

---
 .../rl/model_rl_experiment_player.py          | 158 ++++++++++++------
 1 file changed, 109 insertions(+), 49 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index b5140c2d0..5592d392c 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -3,12 +3,11 @@
 from __future__ import division
 from __future__ import print_function
 
-import contextlib
+from PIL import Image
+from PIL import ImageFont
+from PIL import ImageDraw
 import copy
-import datetime
-import math
 import os
-import time
 from gym.spaces import Box
 import numpy as np
 
@@ -26,12 +25,54 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.rl.trainer_model_based import FLAGS, setup_directories, temporary_flags
-
+from gym.utils import play
 import tensorflow as tf
 
 
 HP_SCOPES = ["loop", "model", "ppo"]
 
+_font = None
+FONT_SIZE = 20
+
+def _get_font():
+  global _font
+  if _font is None:
+    #weirdness due to various working dirs
+    FONT_PATHS = ["tensor-2-tensor-with-mrunner/tensor-2-tensor-with-mrunner/deepsense_experiments/Xerox Serif Narrow.ttf",
+                  "tensor-2-tensor-with-mrunner/deepsense_experiments/Xerox Serif Narrow.ttf",
+                  "deepsense_experiments/Xerox Serif Narrow.ttf"]
+
+    for path in FONT_PATHS:
+      try:
+        _font = ImageFont.truetype(path, FONT_SIZE)
+        return _font
+      except:
+        pass
+
+
+def _assert_image(img):
+  if isinstance(img, np.ndarray):
+    img = Image.fromarray(np.ndarray.astype(img, np.uint8))
+  return img
+
+
+def write_on_image(img, text="", positon=(0,0), color=(255,255,255)):
+  img = _assert_image(img)
+  if text=="":
+    return img
+  draw = ImageDraw.Draw(img)
+  font = _get_font()
+  draw.text(positon, text, color, font=font)
+
+  return img
+
+def concatenate_images(*imgs, axis=1):
+  imgs = [_assert_image(img) for img in imgs]
+  imgs_np = [np.array(img) for img in imgs]
+  concatenated_im_np = np.concatenate(imgs_np, axis=axis)
+
+  return _assert_image(concatenated_im_np)
+
 
 def train_agent(problem_name, agent_model_dir,
                 event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
@@ -76,11 +117,13 @@ def train_agent(problem_name, agent_model_dir,
     sess = tf.Session()
     env = DebugBatchEnv(ppo_hparams, sess)
     sess.run(tf.global_variables_initializer())
-    env.initialize()
+    # env.initialize()
+
+    key_mapping = {(): 100, (ord('q'),):1, (ord('a'),):2,
+                   (ord('r'),):101,
+                   (ord('p'),):102}
 
-    key_mapping = {(ord('q'),):1, (ord('a'),):2}
-    from gym.utils import play
-    play.play(env, zoom=4, fps=40, keys_to_action=key_mapping)
+    play.play(env, zoom=2, fps=10, keys_to_action=key_mapping)
 
 from gym.core import Env
 
@@ -96,47 +139,56 @@ def __init__(self, hparams, sess = None):
     self.action_space = Discrete(6)
     self.observation_space = Box(low=0, high=255, shape=(210, 160, 3), dtype=np.uint8)
 
-    batch_env = batch_env_factory(hparams)
+    # batch_env = batch_env_factory(hparams)
 
     self.action = tf.placeholder(shape=(1,), dtype=tf.int32)
 
-    self.reward, self.done = batch_env.simulate(self.action)
-    self.observation = batch_env.observ
-    self.reset_op = batch_env.reset(tf.constant([0], dtype=tf.int32))
+    # self.reward, self.done = batch_env.simulate(self.action)
+    # self.observation = batch_env.observ
+    # self.reset_op = batch_env.reset(tf.constant([0], dtype=tf.int32))
 
     environment_wrappers = hparams.environment_spec.wrappers
     wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
 
-    to_initialize = [batch_env]
-    for w in wrappers:
-      batch_env = w[0](batch_env, **w[1])
-      to_initialize.append(batch_env)
-
-    def initialization_lambda():
-      for batch_env in to_initialize:
-        batch_env.initialize(sess)
+    # to_initialize = [batch_env]
+    # for w in wrappers:
+    #   batch_env = w[0](batch_env, **w[1])
+    #   to_initialize.append(batch_env)
+    #
+    # def initialization_lambda():
+    #   for batch_env in to_initialize:
+    #     batch_env.initialize(sess)
 
-    self.initialize = initialization_lambda
+    # self.initialize = initialization_lambda
 
-    obs_copy = batch_env.observ + 0
+    # obs_copy = batch_env.observ + 0
 
-    actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
-    self.policy_probs = actor_critic.policy.probs[0, 0, :]
-    self.value = actor_critic.value[0, :]
+    # actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
+    # self.policy_probs = actor_critic.policy.probs[0, 0, :]
+    # self.value = actor_critic.value[0, :]
     self._tmp = 1
+    self.res = None
 
   def render(self, mode='human'):
     raise NotImplemented()
 
   def reset(self):
-    observ = self.sess.run(self.reset_op)
+    # observ = self.sess.run(self.reset_op)
+    self._tmp = 0
+    _observ = np.ones(shape=(210, 160, 3), dtype=np.uint8) * 10 * self._tmp
+    _observ[0, 0, 0] = 0
+    _observ[0, 0, 1] = 255
+    self.res = (_observ, 0, False, "a", "b")
+    observ = self._augment_observation()
     return observ
 
 
   def _step_fake(self, action):
 
-    print("Action:{}".format(action))
     observ = np.ones(shape=(210, 160, 3), dtype=np.uint8)*10*self._tmp
+    observ[0, 0, 0] = 0
+    observ[0, 0, 1] = 255
+
     self._tmp += 1
     if self._tmp>20:
       self._tmp = 0
@@ -155,8 +207,36 @@ def _env_step_fake(self, action):
 
     return observ[0, ...], rew[0, ...], done[0, ...], probs, vf
 
+  def _augment_observation(self):
+    _observ, rew, probs, probs, vf = self.res
+    info_pane = np.zeros_like(_observ)
+    write_on_image(info_pane, "Policy:{}".format(probs))
+    write_on_image(info_pane, "Value function:{}".format(vf), positon=(0, 10))
+    write_on_image(info_pane, "Rew:{}".format(rew), positon=(0, 20))
+
   def step(self, action):
-    observ, rew, done, probs, vf = self._step_fake(action)
+    #Special codes
+    if action==100:
+      #Skip action
+      observ, rew, done, _, _ = self.res
+      return observ, rew, done, {}
+
+    if action == 101:
+      #reset
+      observ, rew, _ = self.res
+      return observ, rew, True, {}
+
+    if action == 102:
+      #play
+      raise NotImplemented()
+
+    #standard codes
+    _observ, rew, done, probs, vf = self._step_fake(action)
+    self.res = (_observ, rew, done, probs, vf)
+
+    observ = self._augment_observation()
+
+
     return observ, rew, done, {"probs": probs, "vf": vf}
 
 
@@ -181,34 +261,14 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     game_with_mode = hparams.game
   # Problems
   if using_autoencoder:
-    problem_name = (
-        "gym_discrete_problem_with_agent_on_%s_with_autoencoder"
-        % game_with_mode)
-    world_model_problem = (
-        "gym_discrete_problem_with_agent_on_%s_autoencoded" % game_with_mode)
     simulated_problem_name = (
         "gym_simulated_discrete_problem_with_agent_on_%s_autoencoded"
         % game_with_mode)
   else:
-    problem_name = ("gym_discrete_problem_with_agent_on_%s" % game_with_mode)
-    world_model_problem = problem_name
     simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
                               % game_with_mode)
-    if problem_name not in registry.list_problems():
-      tf.logging.info("Game Problem %s not found; dynamically registering",
-                      problem_name)
-      gym_problems_specs.create_problems_for_game(hparams.game,
-                                                  game_mode="Deterministic-v4")
 
-  # Autoencoder model dir
-  autoencoder_model_dir = directories.get("autoencoder")
-
-  # Per-epoch state
-  epoch_metrics = []
   epoch_data_dirs = []
-
-
-  # Collect data from the real environment with random policy
   data_dir = os.path.join(directories["data"], "random")
   epoch_data_dirs.append(data_dir)
 

From b3a9719dff71d7e292e97e7f120ea610a0f79ca8 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 8 Sep 2018 18:00:02 -0700
Subject: [PATCH 0782/2720] debug_info_pane

---
 .../rl/model_rl_experiment_player.py          | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index 5592d392c..495313b61 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -137,7 +137,7 @@ def __init__(self, hparams, sess = None):
       self.sess = sess
 
     self.action_space = Discrete(6)
-    self.observation_space = Box(low=0, high=255, shape=(210, 160, 3), dtype=np.uint8)
+    self.observation_space = Box(low=0, high=255, shape=(210, 320, 3), dtype=np.uint8)
 
     # batch_env = batch_env_factory(hparams)
 
@@ -185,7 +185,7 @@ def reset(self):
 
   def _step_fake(self, action):
 
-    observ = np.ones(shape=(210, 160, 3), dtype=np.uint8)*10*self._tmp
+    observ = np.ones(shape=(210, 320, 3), dtype=np.uint8)*10*self._tmp
     observ[0, 0, 0] = 0
     observ[0, 0, 1] = 255
 
@@ -210,15 +210,20 @@ def _env_step_fake(self, action):
   def _augment_observation(self):
     _observ, rew, probs, probs, vf = self.res
     info_pane = np.zeros_like(_observ)
-    write_on_image(info_pane, "Policy:{}".format(probs))
-    write_on_image(info_pane, "Value function:{}".format(vf), positon=(0, 10))
-    write_on_image(info_pane, "Rew:{}".format(rew), positon=(0, 20))
+    info_str = "Policy:{}\nValue function:{}\nReward:{}".format(probs, vf, rew)
+    info_pane = write_on_image(info_pane, info_str)
+
+    augmented_observ = concatenate_images(_observ, info_pane)
+    augmented_observ = np.array(augmented_observ)
+    return augmented_observ
+
 
   def step(self, action):
     #Special codes
     if action==100:
       #Skip action
-      observ, rew, done, _, _ = self.res
+      _, rew, done, _, _ = self.res
+      observ = self._augment_observation()
       return observ, rew, done, {}
 
     if action == 101:
@@ -235,8 +240,6 @@ def step(self, action):
     self.res = (_observ, rew, done, probs, vf)
 
     observ = self._augment_observation()
-
-
     return observ, rew, done, {"probs": probs, "vf": vf}
 
 

From d5f551b98ba89186b33186db1838b76a8b1d8e27 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 8 Sep 2018 18:16:19 -0700
Subject: [PATCH 0783/2720] more info

---
 .../rl/model_rl_experiment_player.py          | 23 +++++++++++--------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index 495313b61..df0bc672e 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -123,7 +123,7 @@ def train_agent(problem_name, agent_model_dir,
                    (ord('r'),):101,
                    (ord('p'),):102}
 
-    play.play(env, zoom=2, fps=10, keys_to_action=key_mapping)
+    play.play(env, zoom=1, fps=10, keys_to_action=key_mapping)
 
 from gym.core import Env
 
@@ -178,14 +178,14 @@ def reset(self):
     _observ = np.ones(shape=(210, 160, 3), dtype=np.uint8) * 10 * self._tmp
     _observ[0, 0, 0] = 0
     _observ[0, 0, 1] = 255
-    self.res = (_observ, 0, False, "a", "b")
+    self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)
     observ = self._augment_observation()
     return observ
 
 
   def _step_fake(self, action):
 
-    observ = np.ones(shape=(210, 320, 3), dtype=np.uint8)*10*self._tmp
+    observ = np.ones(shape=(210, 160+250, 3), dtype=np.uint8)*10*self._tmp
     observ[0, 0, 0] = 0
     observ[0, 0, 1] = 255
 
@@ -208,9 +208,16 @@ def _env_step_fake(self, action):
     return observ[0, ...], rew[0, ...], done[0, ...], probs, vf
 
   def _augment_observation(self):
-    _observ, rew, probs, probs, vf = self.res
-    info_pane = np.zeros_like(_observ)
-    info_str = "Policy:{}\nValue function:{}\nReward:{}".format(probs, vf, rew)
+    _observ, rew, done, probs, vf = self.res
+    info_pane = np.zeros(shape=(210, 250, 3), dtype=np.uint8)
+    probs_str = ""
+    for p in probs:
+      probs_str += "%.2f" % p +", "
+
+    action = np.argmax(probs)
+
+    info_str = "Policy:{}\nAction:{}\nValue function:{}\nReward:{}".format(probs_str, action,
+                                                                           vf, rew)
     info_pane = write_on_image(info_pane, info_str)
 
     augmented_observ = concatenate_images(_observ, info_pane)
@@ -246,10 +253,6 @@ def step(self, action):
 
 def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   """Run the main training loop."""
-  if report_fn:
-    assert report_metric is not None
-
-  # Global state
 
   # Directories
   subdirectories = ["data", "tmp", "world_model", "ppo"]

From 4d56c063488f22ec15bfd04dbb014f98ece9e178 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 8 Sep 2018 18:18:49 -0700
Subject: [PATCH 0784/2720] bug fixes and resets

---
 tensor2tensor/rl/model_rl_experiment_player.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index df0bc672e..fd3e181ff 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -137,7 +137,7 @@ def __init__(self, hparams, sess = None):
       self.sess = sess
 
     self.action_space = Discrete(6)
-    self.observation_space = Box(low=0, high=255, shape=(210, 320, 3), dtype=np.uint8)
+    self.observation_space = Box(low=0, high=255, shape=(210, 160+250, 3), dtype=np.uint8)
 
     # batch_env = batch_env_factory(hparams)
 
@@ -185,7 +185,7 @@ def reset(self):
 
   def _step_fake(self, action):
 
-    observ = np.ones(shape=(210, 160+250, 3), dtype=np.uint8)*10*self._tmp
+    observ = np.ones(shape=(210, 160, 3), dtype=np.uint8)*10*self._tmp
     observ[0, 0, 0] = 0
     observ[0, 0, 1] = 255
 
@@ -235,8 +235,10 @@ def step(self, action):
 
     if action == 101:
       #reset
-      observ, rew, _ = self.res
-      return observ, rew, True, {}
+      self.reset()
+      _, rew, done, _, _ = self.res
+      observ = self._augment_observation()
+      return observ, rew, done, {}
 
     if action == 102:
       #play

From d36d8100cba0736bb5087ff448750bdd2af7b697 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 8 Sep 2018 18:41:05 -0700
Subject: [PATCH 0785/2720] restoring world model

---
 .../rl/model_rl_experiment_player.py          | 101 ++++++++++--------
 1 file changed, 55 insertions(+), 46 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index fd3e181ff..247da2762 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -56,13 +56,13 @@ def _assert_image(img):
   return img
 
 
-def write_on_image(img, text="", positon=(0,0), color=(255,255,255)):
+def write_on_image(img, text="", position=(0, 0), color=(255, 255, 255)):
   img = _assert_image(img)
   if text=="":
     return img
   draw = ImageDraw.Draw(img)
   font = _get_font()
-  draw.text(positon, text, color, font=font)
+  draw.text(position, text, color, font=font)
 
   return img
 
@@ -117,13 +117,18 @@ def train_agent(problem_name, agent_model_dir,
     sess = tf.Session()
     env = DebugBatchEnv(ppo_hparams, sess)
     sess.run(tf.global_variables_initializer())
-    # env.initialize()
+    env.initialize()
+    env_model_loader = tf.train.Saver(
+      tf.global_variables("next_frame*"))
+
+    trainer_lib.restore_checkpoint(world_model_dir, env_model_loader, sess,
+      must_restore=True)
 
     key_mapping = {(): 100, (ord('q'),):1, (ord('a'),):2,
                    (ord('r'),):101,
                    (ord('p'),):102}
 
-    play.play(env, zoom=1, fps=10, keys_to_action=key_mapping)
+    play.play(env, zoom=2, fps=10, keys_to_action=key_mapping)
 
 from gym.core import Env
 
@@ -138,40 +143,51 @@ def __init__(self, hparams, sess = None):
 
     self.action_space = Discrete(6)
     self.observation_space = Box(low=0, high=255, shape=(210, 160+250, 3), dtype=np.uint8)
+    self._tmp = 1
+    self.res = None
+    self._prepare_networks(hparams, self.sess)
 
-    # batch_env = batch_env_factory(hparams)
-
+  def _prepare_networks(self, hparams, sess):
     self.action = tf.placeholder(shape=(1,), dtype=tf.int32)
-
-    # self.reward, self.done = batch_env.simulate(self.action)
-    # self.observation = batch_env.observ
-    # self.reset_op = batch_env.reset(tf.constant([0], dtype=tf.int32))
+    batch_env = batch_env_factory(hparams)
+    self.reward, self.done = batch_env.simulate(self.action)
+    self.observation = batch_env.observ
+    self.reset_op = batch_env.reset(tf.constant([0], dtype=tf.int32))
 
     environment_wrappers = hparams.environment_spec.wrappers
     wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
 
-    # to_initialize = [batch_env]
-    # for w in wrappers:
-    #   batch_env = w[0](batch_env, **w[1])
-    #   to_initialize.append(batch_env)
-    #
-    # def initialization_lambda():
-    #   for batch_env in to_initialize:
-    #     batch_env.initialize(sess)
+    to_initialize = [batch_env]
+    for w in wrappers:
+      batch_env = w[0](batch_env, **w[1])
+      to_initialize.append(batch_env)
 
-    # self.initialize = initialization_lambda
+    def initialization_lambda():
+      for batch_env in to_initialize:
+        batch_env.initialize(sess)
 
-    # obs_copy = batch_env.observ + 0
+    self.initialize = initialization_lambda
+
+    obs_copy = batch_env.observ + 0
+
+    actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
+    self.policy_probs = actor_critic.policy.probs[0, 0, :]
+    self.value = actor_critic.value[0, :]
 
-    # actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
-    # self.policy_probs = actor_critic.policy.probs[0, 0, :]
-    # self.value = actor_critic.value[0, :]
-    self._tmp = 1
-    self.res = None
 
   def render(self, mode='human'):
     raise NotImplemented()
 
+  def _fake_reset(self):
+    self._tmp = 0
+    _observ = np.ones(shape=(210, 160, 3), dtype=np.uint8) * 10 * self._tmp
+    _observ[0, 0, 0] = 0
+    _observ[0, 0, 1] = 255
+    self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)
+    observ = self._augment_observation()
+    return observ
+
+
   def reset(self):
     # observ = self.sess.run(self.reset_op)
     self._tmp = 0
@@ -200,7 +216,7 @@ def _step_fake(self, action):
 
     return observ, rew, done, probs, vf
 
-  def _env_step_fake(self, action):
+  def _step_env(self, action):
     observ, rew, done, probs, vf = self.sess.\
       run([self.observation, self.reward, self.done, self.policy_probs, self.value],
           feed_dict={self.action: [action]})
@@ -216,7 +232,7 @@ def _augment_observation(self):
 
     action = np.argmax(probs)
 
-    info_str = "Policy:{}\nAction:{}\nValue function:{}\nReward:{}".format(probs_str, action,
+    info_str = " Policy:{}\n Action:{}\n Value function:{}\n Reward:{}".format(probs_str, action,
                                                                            vf, rew)
     info_pane = write_on_image(info_pane, info_str)
 
@@ -245,7 +261,7 @@ def step(self, action):
       raise NotImplemented()
 
     #standard codes
-    _observ, rew, done, probs, vf = self._step_fake(action)
+    _observ, rew, done, probs, vf = self._step_env(action)
     self.res = (_observ, rew, done, probs, vf)
 
     observ = self._augment_observation()
@@ -280,26 +296,19 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   data_dir = os.path.join(directories["data"], "random")
   epoch_data_dirs.append(data_dir)
 
-  for epoch in range(hparams.epochs):
-    is_final_epoch = (epoch + 1) == hparams.epochs
-
-    # Combine all previously collected environment data
-    epoch_data_dir = os.path.join(directories["data"], str(epoch))
-
-
-    ppo_event_dir = os.path.join(directories["world_model"],
-                                 "ppo_summaries", str(epoch))
-    ppo_model_dir = directories["ppo"]
-    if not hparams.ppo_continue_training:
-      ppo_model_dir = ppo_event_dir
-    train_agent(simulated_problem_name, ppo_model_dir,
-                ppo_event_dir, directories["world_model"], epoch_data_dir,
-                hparams, epoch=epoch, is_final_epoch=is_final_epoch)
-
+  epoch = hparams.epochs-1
 
-  raise NotImplementedError()
-  return 1
+  # Combine all previously collected environment data
+  epoch_data_dir = os.path.join(directories["data"], str(epoch))
 
+  ppo_event_dir = os.path.join(directories["world_model"],
+                               "ppo_summaries", str(epoch))
+  ppo_model_dir = directories["ppo"]
+  if not hparams.ppo_continue_training:
+    ppo_model_dir = ppo_event_dir
+  train_agent(simulated_problem_name, ppo_model_dir,
+              ppo_event_dir, directories["world_model"], epoch_data_dir,
+              hparams, epoch=epoch, is_final_epoch=False)
 
 
 def create_loop_hparams():

From b17d6e219aac538432c79425e92ce5da1ce88107 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 8 Sep 2018 18:43:58 -0700
Subject: [PATCH 0786/2720] restring agents net

---
 tensor2tensor/rl/model_rl_experiment_player.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index 247da2762..a07a7522d 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -118,12 +118,16 @@ def train_agent(problem_name, agent_model_dir,
     env = DebugBatchEnv(ppo_hparams, sess)
     sess.run(tf.global_variables_initializer())
     env.initialize()
+
     env_model_loader = tf.train.Saver(
       tf.global_variables("next_frame*"))
-
     trainer_lib.restore_checkpoint(world_model_dir, env_model_loader, sess,
       must_restore=True)
 
+    model_saver = tf.train.Saver(
+      tf.global_variables(".*network_parameters.*"))
+    trainer_lib.restore_checkpoint(agent_model_dir, model_saver, sess)
+
     key_mapping = {(): 100, (ord('q'),):1, (ord('a'),):2,
                    (ord('r'),):101,
                    (ord('p'),):102}

From 42d9d4667a976e2d9df27ee597e2f74721eee8ba Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 8 Sep 2018 18:47:15 -0700
Subject: [PATCH 0787/2720] fixing resets

---
 tensor2tensor/rl/model_rl_experiment_player.py | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index a07a7522d..e8d6161ea 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -188,23 +188,22 @@ def _fake_reset(self):
     _observ[0, 0, 0] = 0
     _observ[0, 0, 1] = 255
     self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)
-    observ = self._augment_observation()
-    return observ
-
 
-  def reset(self):
-    # observ = self.sess.run(self.reset_op)
-    self._tmp = 0
-    _observ = np.ones(shape=(210, 160, 3), dtype=np.uint8) * 10 * self._tmp
+  def _reset_env(self):
+    _observ = self.sess.run(self.reset_op)[0, ...]
     _observ[0, 0, 0] = 0
     _observ[0, 0, 1] = 255
+    #TODO:(put correct numbers)
     self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)
+
+
+  def reset(self):
+    self._reset_env()
     observ = self._augment_observation()
     return observ
 
 
   def _step_fake(self, action):
-
     observ = np.ones(shape=(210, 160, 3), dtype=np.uint8)*10*self._tmp
     observ[0, 0, 0] = 0
     observ[0, 0, 1] = 255

From 3fca12df39e90916b53f7800055b42f934d4e544 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Sat, 8 Sep 2018 18:56:38 -0700
Subject: [PATCH 0788/2720] clean-up

---
 tensor2tensor/rl/model_rl_experiment_player.py | 16 ++--------------
 1 file changed, 2 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index e8d6161ea..f7b040bc2 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -74,24 +74,12 @@ def concatenate_images(*imgs, axis=1):
   return _assert_image(concatenated_im_np)
 
 
-def train_agent(problem_name, agent_model_dir,
+def show_agent(problem_name, agent_model_dir,
                 event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
                 is_final_epoch=False):
   """Train the PPO agent in the simulated environment."""
   gym_problem = registry.problem(problem_name)
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
-  ppo_params_names = ["epochs_num", "epoch_length",
-                      "learning_rate", "num_agents",
-                      "optimization_epochs"]
-
-  for param_name in ppo_params_names:
-    ppo_param_name = "ppo_" + param_name
-    if ppo_param_name in hparams:
-      ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
-
-  ppo_hparams.save_models_every_epochs = 10
-  ppo_hparams.world_model_dir = world_model_dir
-  ppo_hparams.add_hparam("force_beginning_resets", True)
 
   # Adding model hparams for model specific adjustments
   model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
@@ -309,7 +297,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   ppo_model_dir = directories["ppo"]
   if not hparams.ppo_continue_training:
     ppo_model_dir = ppo_event_dir
-  train_agent(simulated_problem_name, ppo_model_dir,
+  show_agent(simulated_problem_name, ppo_model_dir,
               ppo_event_dir, directories["world_model"], epoch_data_dir,
               hparams, epoch=epoch, is_final_epoch=False)
 

From 994905fc807aeb7ed5c558075e4eadff1359d3a0 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sat, 8 Sep 2018 19:52:21 -0700
Subject: [PATCH 0789/2720] Internal.

PiperOrigin-RevId: 212149739
---
 tensor2tensor/models/video/sv2p.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 4f36b296b..5dbaa5cd8 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -462,7 +462,10 @@ def process_single_frame(prev_outputs, inputs):
     gen_images = tf.concat((first_gen_images, gen_images), axis=0)
     gen_rewards = tf.concat((first_gen_rewards, gen_rewards), axis=0)
 
-    return gen_images, gen_rewards, [latent_mean], [latent_std]
+    if self.hparams.stochastic_model:
+      return gen_images, gen_rewards, [latent_mean], [latent_std]
+    else:
+      return gen_images, gen_rewards, None, None
 
   def get_extra_loss(self, latent_means=None, latent_stds=None,
                      true_frames=None, gen_frames=None, beta=1.0):
@@ -470,7 +473,7 @@ def get_extra_loss(self, latent_means=None, latent_stds=None,
     del true_frames
     del gen_frames
     kl_loss = 0.0
-    if self.is_training:
+    if self.is_training and self.hparams.stochastic_model:
       for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
         kl_loss += common_layers.kl_divergence(mean, std)
         tf.summary.histogram("posterior_mean_%d" % i, mean)

From 1f8dbb257210cea00f86b1545cf08af293d7f256 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Sun, 9 Sep 2018 16:32:18 -0700
Subject: [PATCH 0790/2720] allow real input modality in TransformerAE.infer

PiperOrigin-RevId: 212198451
---
 .../models/research/transformer_vae.py        | 26 +++++--------------
 1 file changed, 6 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 9b6d20737..3f5dc5012 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -675,24 +675,15 @@ def infer(self, features=None, decode_length=50, beam_size=1, top_beams=1,
     if "partial_targets" in features:
       initial_output = tf.convert_to_tensor(features["partial_targets"])
     else:
-      # inputs might not be present in features (e.g.: language modeling),
-      # in which case we fallback to 'infer_targets' for calculating initial
-      # input shape, type, etc.
-      inputs_or_targets = features.get("inputs", features["infer_targets"])
-      batch_size = common_layers.shape_list(inputs_or_targets)[0]
-      length = common_layers.shape_list(inputs_or_targets)[1]
-      hidden_dim = common_layers.shape_list(inputs_or_targets)[-1]
+      batch_size = common_layers.shape_list(features["inputs"])[0]
+      length = common_layers.shape_list(features["inputs"])[1]
       target_length = tf.to_int32(2.0 * tf.to_float(length))
-      initial_output = tf.zeros((batch_size, target_length, 1, hidden_dim),
-                                dtype=inputs_or_targets.dtype)
+      initial_output = tf.zeros((batch_size, target_length, 1, 1),
+                                dtype=tf.int64)
 
     features["targets"] = initial_output
     logits, _ = self(features)  # pylint: disable=not-callable
-    # this should only happen if we're doing target_modality not real
-    if inputs_or_targets.dtype == tf.float32:
-      samples = logits
-    else:
-      samples = tf.argmax(logits, axis=-1)
+    samples = tf.argmax(logits, axis=-1)
 
     # More steps.
     self.predict_mask = 0.0  # Use the provided targets this time.
@@ -701,12 +692,7 @@ def infer(self, features=None, decode_length=50, beam_size=1, top_beams=1,
       with tf.variable_scope(tf.get_variable_scope(), reuse=True):
         features["targets"] = samples
         logits, _ = self(features)  # pylint: disable=not-callable
-        if inputs_or_targets.dtype == tf.float32:
-          # When target_modality is real, the last axis does not represent
-          # classes, so it should not be argmax'ed
-          samples = logits
-        else:
-          samples = tf.argmax(logits, axis=-1)
+        samples = tf.argmax(logits, axis=-1)
 
     self.predict_mask = 1.0
     if inputs_old is not None:  # Restore to not confuse Estimator.

From d0cd2e65716afa4adec9dbd74b4e025114073623 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Sun, 9 Sep 2018 18:11:05 -0700
Subject: [PATCH 0791/2720] Allow configuring frames_per_second as a hparam at
 decode time.

PiperOrigin-RevId: 212202714
---
 tensor2tensor/data_generators/video_utils.py | 5 +++--
 tensor2tensor/utils/decoding.py              | 4 +++-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index f665d0ca3..7bd343faa 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -45,6 +45,7 @@ def resize_video_frames(images, size):
 def display_video_hooks(hook_args):
   """Hooks to display videos at decode time."""
   predictions = hook_args.predictions
+  fps = hook_args.decode_hparams.frames_per_second
 
   all_summaries = []
   for decode_ind, decode in enumerate(predictions):
@@ -59,10 +60,10 @@ def display_video_hooks(hook_args):
     input_videos = np.concatenate((input_videos, target_videos), axis=1)
     output_videos = np.concatenate((input_videos, output_videos), axis=1)
     input_summ_vals, _ = common_video.py_gif_summary(
-        "decode_%d/input" % decode_ind, input_videos, max_outputs=10, fps=10,
+        "decode_%d/input" % decode_ind, input_videos, max_outputs=10, fps=fps,
         return_summary_value=True)
     output_summ_vals, _ = common_video.py_gif_summary(
-        "decode_%d/output" % decode_ind, output_videos, max_outputs=10, fps=10,
+        "decode_%d/output" % decode_ind, output_videos, max_outputs=10, fps=fps,
         return_summary_value=True)
     all_summaries.extend(input_summ_vals)
     all_summaries.extend(output_summ_vals)
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 4a2d888f0..4d9024b18 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -64,7 +64,9 @@ def decode_hparams(overrides=""):
       shard_id=0,
       num_decodes=1,
       force_decode_length=False,
-      display_decoded_images=False)
+      display_decoded_images=False,
+      # Used for video decoding.
+      frames_per_second=10)
   hp.parse(overrides)
   return hp
 

From a7e2473ea1e95c8c99f0ac67ea63d0607900fa3e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 10 Sep 2018 10:58:08 -0700
Subject: [PATCH 0792/2720] Call summarize_features before sharding because it
 expects unsharded tensors.

PiperOrigin-RevId: 212295951
---
 tensor2tensor/utils/t2t_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index f8b2aae60..af434371d 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -209,6 +209,7 @@ def call(self, inputs, **kwargs):
         optimize.get_variable_initializer(self.hparams))
     with self._eager_var_store.as_default():
       self._fill_problem_hparams_features(features)
+      summarize_features(features, num_shards=self._num_datashards)
       sharded_features = self._shard_features(features)
       sharded_logits, losses = self.model_fn_sharded(sharded_features)
       if isinstance(sharded_logits, dict):
@@ -230,7 +231,6 @@ def body_sharded(self, sharded_features):
 
   def model_fn_sharded(self, sharded_features):
     dp = self._data_parallelism
-    summarize_features(sharded_features, num_shards=dp.n)
     datashard_to_features = self._to_features_per_datashard(sharded_features)
     if self.use_body_sharded:
       # MoE models override body_sharded

From 5554aefb866eb6d0ccb880353bc0b2f13ceeabb6 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 10 Sep 2018 11:06:45 -0700
Subject: [PATCH 0793/2720] adding reward loss to the internal loss.

PiperOrigin-RevId: 212297757
---
 tensor2tensor/models/video/sv2p.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 5dbaa5cd8..d13120c63 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -569,8 +569,15 @@ def body(self, features):
       return_targets = {"targets": predictions, "target_reward": reward_pred}
 
     if hparams.internal_loss:
-      loss = tf.losses.mean_squared_error(all_frames[1:], gen_images)
-      extra_loss = {"training": loss + extra_loss}
+      recon_loss = tf.losses.mean_squared_error(all_frames[1:], gen_images)
+      rew_loss = 0.0
+      if hparams.reward_prediction:
+        rew_loss = tf.losses.softmax_cross_entropy(all_rewards[1:], gen_rewards)
+        tf.summary.scalar("loss/reward", rew_loss)
+      tf.summary.scalar("loss/recon", recon_loss)
+      tf.summary.scalar("loss/kl", extra_loss)
+
+      extra_loss = {"training": recon_loss + rew_loss + extra_loss}
 
     return return_targets, extra_loss
 

From 926b55eff0f0bb863092c4d9e6e7f3ef8242ed8a Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 10 Sep 2018 11:09:19 -0700
Subject: [PATCH 0794/2720] Internal change

PiperOrigin-RevId: 212298331
---
 tensor2tensor/models/research/glow.py          |  2 +-
 tensor2tensor/models/research/glow_ops.py      | 11 ++++++-----
 tensor2tensor/models/research/glow_ops_test.py | 18 +++++++++++-------
 3 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 183843729..95e4b3a2b 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -99,7 +99,7 @@ def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
     var_scope = tf.variable_scope("glow/body", reuse=True)
     # If eps=None, images are sampled from the prior.
     with arg_scope(ops, init=False), var_scope:
-      predictions, _ = glow_ops.encoder_decoder(
+      predictions, _, _ = glow_ops.encoder_decoder(
           "codec", self.z_sample, self.hparams, eps=None, reverse=True)
 
     return self.scale(predictions)
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 6dc6f64be..19948be61 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -532,7 +532,7 @@ def split(name, x, reverse=False, eps=None, eps_std=None, cond_latent=None,
         x2 = eps_std * tf.random_normal(common_layers.shape_list(x))
       else:
         x2 = prior_dist.sample()
-      return tf.concat([x, x2], 3)
+      return tf.concat([x, x2], 3), x2
 
 
 @add_arg_scope
@@ -713,12 +713,13 @@ def encoder_decoder(name, x, hparams, eps=None, reverse=False,
           if cond_latents is not None:
             curr_latent = cond_latents[level]
 
-          x = split("split_%d" % level, x, eps=curr_eps, reverse=True,
-                    cond_latent=curr_latent,
-                    merge_std=hparams.level_prior_scale)
+          x, latent = split("split_%d" % level, x, eps=curr_eps, reverse=True,
+                            cond_latent=curr_latent,
+                            merge_std=hparams.level_prior_scale)
+          all_latents.append(latent)
 
         x, obj = revnet(
             "revnet_%d" % level, x, hparams=hparams, reverse=True)
         objective += obj
         x = squeeze("squeeze_%d" % level, x, reverse=True)
-      return x, objective
+      return x, objective, all_latents[::-1]
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index d48dd244c..0666e552d 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -138,7 +138,7 @@ def test_split(self):
     with tf.Graph().as_default():
       x = tf.random_uniform(shape=(16, 5, 5, 32))
       x_inv, _, eps, z = glow_ops.split("split", x)
-      x_inv_inv = glow_ops.split("split", x_inv, reverse=True, eps=eps)
+      x_inv_inv, _ = glow_ops.split("split", x_inv, reverse=True, eps=eps)
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())
         x_inv_np, diff, z_np = session.run([x_inv, x - x_inv_inv, z])
@@ -173,17 +173,21 @@ def test_encoder_decoder(self):
       x = tf.random_uniform(shape=(16, 64, 64, 4), seed=0)
       x_inv, _, eps, z_levels = glow_ops.encoder_decoder(
           "encoder_decoder", x, hparams, reverse=False)
-      x_inv_inv, _ = glow_ops.encoder_decoder(
+      x_inv_inv, _, z_inv_levels = glow_ops.encoder_decoder(
           "encoder_decoder", x_inv, hparams, eps=eps, reverse=True)
 
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())
-        diff, x_inv_np, z_levels_np = session.run(
-            [x - x_inv_inv, x_inv, z_levels])
+        diff, x_inv_np, z_levels_np, z_inv_levels_np = session.run(
+            [x - x_inv_inv, x_inv, z_levels, z_inv_levels])
+
         self.assertEqual(len(z_levels_np), 2)
+        self.assertEqual(len(z_inv_levels_np), 2)
         # (h_i, w_i, c_i) = (h_{i-1}/f, w_{i-1}/f, c_{i-1}*(2f)/2) where (f=2)
         self.assertEqual(z_levels_np[0].shape, (16, 32, 32, 8))
         self.assertEqual(z_levels_np[1].shape, (16, 16, 16, 16))
+        self.assertEqual(z_inv_levels_np[0].shape, (16, 32, 32, 8))
+        self.assertEqual(z_inv_levels_np[1].shape, (16, 16, 16, 16))
         self.assertTrue(x_inv_np.shape, (16, 8, 8, 64))
         self.assertTrue(np.allclose(diff, 0.0, atol=1e-2))
 
@@ -225,7 +229,7 @@ def test_encoder_decoder_practical_usage(self):
       with arg_scope(ops, init=False):
         x_inv2, _, all_eps, _ = glow_ops.encoder_decoder(
             "revnet", x_t, hparams, reverse=False)
-        x_inv_inv_, _ = glow_ops.encoder_decoder(
+        x_inv_inv_, _, _ = glow_ops.encoder_decoder(
             "revnet", x_inv2, hparams, eps=all_eps, reverse=True)
 
       with tf.Session() as session:
@@ -266,8 +270,8 @@ def check_split_latent_conditioning(self, merge_std):
       x_inv, _, eps, x2_t = glow_ops.split(merge_std, x_t, cond_latent=latent_t,
                                            merge_std=merge_std)
       # Test reversibility.
-      x_inv_inv = glow_ops.split(merge_std, x_inv, cond_latent=latent_t,
-                                 merge_std=merge_std, eps=eps, reverse=True)
+      x_inv_inv, _ = glow_ops.split(merge_std, x_inv, cond_latent=latent_t,
+                                    merge_std=merge_std, eps=eps, reverse=True)
       with tf.Session() as sess:
         sess.run(tf.global_variables_initializer())
         actual_eps, actual_x2, diff_np = sess.run([eps, x2_t, x_inv_inv - x_t])

From 99bbbead46a9e6ea281ee181105083f49fc459ee Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 10 Sep 2018 15:08:17 -0700
Subject: [PATCH 0795/2720] remove ppo_event_dir

---
 tensor2tensor/rl/model_rl_experiment_player.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index f7b040bc2..8cec7a6e0 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -1,3 +1,5 @@
+#--output_dir=/Users/piotr.milos/Downloads/18 --alsologtostderr --loop_hparams_set=rl_modelrl_base_quick
+#--output_dir=/Users/piotr.milos/t2t/rl_v1 --alsologtostderr --loop_hparams_set=rl_modelrl_tiny
 
 from __future__ import absolute_import
 from __future__ import division
@@ -74,9 +76,7 @@ def concatenate_images(*imgs, axis=1):
   return _assert_image(concatenated_im_np)
 
 
-def show_agent(problem_name, agent_model_dir,
-                event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
-                is_final_epoch=False):
+def show_agent(problem_name, agent_model_dir, world_model_dir, epoch_data_dir, hparams, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the simulated environment."""
   gym_problem = registry.problem(problem_name)
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
@@ -289,7 +289,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   epoch = hparams.epochs-1
 
-  # Combine all previously collected environment data
+
   epoch_data_dir = os.path.join(directories["data"], str(epoch))
 
   ppo_event_dir = os.path.join(directories["world_model"],
@@ -297,9 +297,8 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   ppo_model_dir = directories["ppo"]
   if not hparams.ppo_continue_training:
     ppo_model_dir = ppo_event_dir
-  show_agent(simulated_problem_name, ppo_model_dir,
-              ppo_event_dir, directories["world_model"], epoch_data_dir,
-              hparams, epoch=epoch, is_final_epoch=False)
+  show_agent(simulated_problem_name, ppo_model_dir, directories["world_model"], epoch_data_dir, hparams, epoch=epoch,
+             is_final_epoch=False)
 
 
 def create_loop_hparams():

From 1f23123bd6df78221bb0940b035f3be43b71a829 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 10 Sep 2018 15:15:30 -0700
Subject: [PATCH 0796/2720] more cleanup

---
 tensor2tensor/rl/model_rl_experiment_player.py | 11 ++---------
 1 file changed, 2 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index 8cec7a6e0..1501f689e 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -76,7 +76,7 @@ def concatenate_images(*imgs, axis=1):
   return _assert_image(concatenated_im_np)
 
 
-def show_agent(problem_name, agent_model_dir, world_model_dir, epoch_data_dir, hparams, epoch=0, is_final_epoch=False):
+def show_agent(problem_name, agent_model_dir, world_model_dir, epoch_data_dir, hparams):
   """Train the PPO agent in the simulated environment."""
   gym_problem = registry.problem(problem_name)
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
@@ -283,13 +283,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
                               % game_with_mode)
 
-  epoch_data_dirs = []
-  data_dir = os.path.join(directories["data"], "random")
-  epoch_data_dirs.append(data_dir)
-
   epoch = hparams.epochs-1
-
-
   epoch_data_dir = os.path.join(directories["data"], str(epoch))
 
   ppo_event_dir = os.path.join(directories["world_model"],
@@ -297,8 +291,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   ppo_model_dir = directories["ppo"]
   if not hparams.ppo_continue_training:
     ppo_model_dir = ppo_event_dir
-  show_agent(simulated_problem_name, ppo_model_dir, directories["world_model"], epoch_data_dir, hparams, epoch=epoch,
-             is_final_epoch=False)
+  show_agent(simulated_problem_name, ppo_model_dir, directories["world_model"], epoch_data_dir, hparams)
 
 
 def create_loop_hparams():

From 2fad2eded39b5ad1274eafa2d744fa4be229ab38 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 10 Sep 2018 15:20:13 -0700
Subject: [PATCH 0797/2720] more cleanups

---
 .../rl/model_rl_experiment_player.py          | 24 ++++++-------------
 1 file changed, 7 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index 1501f689e..7c55a7165 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -259,9 +259,8 @@ def step(self, action):
     return observ, rew, done, {"probs": probs, "vf": vf}
 
 
-
-def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
-  """Run the main training loop."""
+def setup_show_metadata(hparams, output_dir):
+  """Setup metadata"""
 
   # Directories
   subdirectories = ["data", "tmp", "world_model", "ppo"]
@@ -282,28 +281,19 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   else:
     simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
                               % game_with_mode)
-
   epoch = hparams.epochs-1
   epoch_data_dir = os.path.join(directories["data"], str(epoch))
-
-  ppo_event_dir = os.path.join(directories["world_model"],
-                               "ppo_summaries", str(epoch))
   ppo_model_dir = directories["ppo"]
-  if not hparams.ppo_continue_training:
-    ppo_model_dir = ppo_event_dir
-  show_agent(simulated_problem_name, ppo_model_dir, directories["world_model"], epoch_data_dir, hparams)
-
 
-def create_loop_hparams():
-  hparams = registry.hparams(FLAGS.loop_hparams_set)
-  hparams.parse(FLAGS.loop_hparams)
-  return hparams
+  show_agent(simulated_problem_name, ppo_model_dir,
+             directories["world_model"], epoch_data_dir, hparams)
 
 
 def main(_):
-  hp = create_loop_hparams()
+  hparams = registry.hparams(FLAGS.loop_hparams_set)
+  hparams.parse(FLAGS.loop_hparams)
   output_dir = FLAGS.output_dir
-  training_loop(hp, output_dir)
+  setup_show_metadata(hparams, output_dir)
 
 
 if __name__ == "__main__":

From 21188b8cc888d884866d2cdf392d238507e0be06 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 10 Sep 2018 15:24:40 -0700
Subject: [PATCH 0798/2720] even more cleanups

---
 tensor2tensor/rl/model_rl_experiment_player.py | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index 7c55a7165..f1a7ba026 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -13,15 +13,8 @@
 from gym.spaces import Box
 import numpy as np
 
-import six
 from gym.spaces import Discrete
-from gym.utils.play import PlayPlot
-
-from tensor2tensor.bin import t2t_trainer
-from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import gym_problems_specs
-from tensor2tensor.layers import discretization
-from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.rl.envs.batch_env_factory import batch_env_factory
 from tensor2tensor.rl.envs.utils import get_policy
 from tensor2tensor.utils import registry
@@ -100,8 +93,6 @@ def show_agent(problem_name, agent_model_dir, world_model_dir, epoch_data_dir, h
       "output_dir": world_model_dir,
       "data_dir": epoch_data_dir,
   }):
-
-
     sess = tf.Session()
     env = DebugBatchEnv(ppo_hparams, sess)
     sess.run(tf.global_variables_initializer())

From 4f6fd205f5fa8ddb76cd7306684602508994eb6e Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 10 Sep 2018 16:32:19 -0700
Subject: [PATCH 0799/2720] new binary for training a ppo agent with a
 pre-trained world-model.

PiperOrigin-RevId: 212356874
---
 .../rl/trainer_model_based_agent_only.py      | 71 +++++++++++++++++++
 1 file changed, 71 insertions(+)
 create mode 100644 tensor2tensor/rl/trainer_model_based_agent_only.py

diff --git a/tensor2tensor/rl/trainer_model_based_agent_only.py b/tensor2tensor/rl/trainer_model_based_agent_only.py
new file mode 100644
index 000000000..e15b01aaa
--- /dev/null
+++ b/tensor2tensor/rl/trainer_model_based_agent_only.py
@@ -0,0 +1,71 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+r"""Training of model-based RL agent assuming a fully trained world model.
+
+Example invocation:
+
+python -m tensor2tensor.rl.trainer_model_based_agent_only \
+    --loop_hparams_set=rl_modelrl_base \
+    --world_model_dir=$HOME/world_model/ \
+    --data_dir=$HOME/data/ \
+    --output_dir=$HOME/ppo_agent_only/ \
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
+from tensor2tensor.data_generators import gym_problems_specs
+from tensor2tensor.rl import trainer_model_based
+
+
+import tensorflow as tf
+
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("world_model_dir", "",
+                    "Directory containing checkpoints of the world model.")
+
+
+def get_simulated_problem_name(game):
+  game_with_mode = game
+  if game in gym_problems_specs.ATARI_GAMES:
+    game_with_mode += "_deterministic-v4"
+  return "gym_simulated_discrete_problem_with_agent_on_%s" % game_with_mode
+
+
+def main(_):
+  hparams = trainer_model_based.create_loop_hparams()
+  problem_name = get_simulated_problem_name(hparams.game)
+  world_model_dir = FLAGS.world_model_dir
+  agent_model_dir = FLAGS.output_dir
+  event_dir = FLAGS.output_dir
+  epoch_data_dir = FLAGS.data_dir  # only required for initial frames
+
+  trainer_model_based.train_agent(
+      problem_name,
+      agent_model_dir,
+      event_dir,
+      world_model_dir,
+      epoch_data_dir,
+      hparams,
+      epoch=0,
+      is_final_epoch=True)
+
+if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
+  tf.app.run()

From 9d597a1ad83ed939cb560e21b338ae9c3dbb07c0 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 10 Sep 2018 18:02:04 -0700
Subject: [PATCH 0800/2720] Fix slicing during top_is_pointwise inference
 without decode_loop_step.

PiperOrigin-RevId: 212370343
---
 tensor2tensor/utils/t2t_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index af434371d..4792a6f14 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -410,7 +410,7 @@ def _top_single(self, body_output, target_modality, features):
           last_position_body_output = tf.expand_dims(
               body_output[:, -1, :, :], axis=[1])
           last_position_targets = tf.expand_dims(
-              features["targets"][:, -1:, :, :], axis=[1])
+              features["targets"][:, -1, :, :], axis=[1])
         else:
           body_output_shape = body_output.shape.as_list()
           last_position_body_output = tf.slice(

From a05b1b423aecf915210333e6371d55ae1e1ec7e3 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Mon, 10 Sep 2018 18:02:10 -0700
Subject: [PATCH 0801/2720] MTF Image transformer: fix bug with eval, add
 property def for dims.

PiperOrigin-RevId: 212370353
---
 .../mesh_tensorflow/mtf_image_transformer.py  | 185 ++++++++++++------
 .../mtf_image_transformer_test.py             |  18 +-
 2 files changed, 138 insertions(+), 65 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
index a854745a6..8934a6b75 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
@@ -37,7 +37,63 @@
 class MtfImageTransformer(mtf_model.MtfModel):
   """Image Transformer in mesh_tensorflow."""
 
-  def set_activation_type(self):
+  @property
+  def inputs_vocab_dim(self):
+    assert self.has_input
+    return mtf.Dimension("inputs_vocab", self._hparams.num_classes)
+
+  @property
+  def targets_vocab_dim(self):
+    return mtf.Dimension(
+        "vocab", self._problem_hparams.target_modality._vocab_size)  # pylint: disable=protected-access
+
+  @property
+  def outputs_vocab_dim(self):
+    return mtf.Dimension("output_vocab", 256)
+
+  @property
+  def rows_dim(self):
+    return mtf.Dimension("rows", self._hparams.img_len)
+
+  @property
+  def cols_dim(self):
+    return mtf.Dimension(
+        "cols", self._hparams.img_len*self._hparams.num_channels)
+
+  @property
+  def orig_cols_dim(self):
+    return mtf.Dimension("orig_cols", self._hparams.img_len)
+
+  @property
+  def channels_dim(self):
+    return mtf.Dimension("channels", self._hparams.num_channels)
+
+  @property
+  def model_dim(self):
+    return mtf.Dimension("d_model", self._hparams.hidden_size)
+
+  @property
+  def max_length_dim(self):
+    return mtf.Dimension("max_length", self._hparams.max_length)
+
+  @property
+  def length_dim(self):
+    return mtf.Dimension("length", self._hparams.max_length)
+
+  @property
+  def heads_dim(self):
+    return mtf.Dimension("heads", self._hparams.num_heads)
+
+  @property
+  def kv_dim(self):
+    return mtf.Dimension("d_kv", self._hparams.d_kv)
+
+  @property
+  def feedforward_dim(self):
+    return mtf.Dimension("d_ff", self._hparams.d_ff)
+
+  @property
+  def activation_type(self):
     hparams = self._hparams
     if hparams.activation_dtype == "float32":
       activation_dtype = tf.float32
@@ -50,47 +106,41 @@ def set_activation_type(self):
           "unknown hparams.activation_dtype %s" % hparams.activation_dtype)
     return activation_dtype
 
-  def create_positional_emb_2d(self, targets, max_length_dim, model_dim):
+  def create_positional_emb_2d(self, targets):
     """Learned 2d positional embedding for images."""
     mesh = targets.mesh
-    hparams = self._hparams
-    activation_dtype = self.set_activation_type()
-
-    rows_dim = mtf.Dimension("rows", hparams.img_len)
-    cols_dim = mtf.Dimension("cols", hparams.img_len*hparams.num_channels)
 
     positional_emb_rows_var = mtf.get_variable(
         mesh, "positional_emb_rows",
-        mtf.Shape([max_length_dim, model_dim]),
+        mtf.Shape([self.max_length_dim, self.model_dim]),
         initializer=tf.random_normal_initializer(),
-        activation_dtype=activation_dtype)
+        activation_dtype=self.activation_type)
     positional_emb_cols_var = mtf.get_variable(
         mesh, "positional_emb_cols",
-        mtf.Shape([max_length_dim, model_dim]),
+        mtf.Shape([self.max_length_dim, self.model_dim]),
         initializer=tf.random_normal_initializer(),
-        activation_dtype=activation_dtype)
+        activation_dtype=self.activation_type)
 
-    targets_position_x = mtf.range(mesh, rows_dim, dtype=tf.int32)
-    targets_position_y = mtf.range(mesh, cols_dim, dtype=tf.int32)
+    targets_position_x = mtf.range(mesh, self.rows_dim, dtype=tf.int32)
+    targets_position_y = mtf.range(mesh, self.cols_dim, dtype=tf.int32)
     position_x = mtf.broadcast(
         mtf.gather(positional_emb_rows_var, targets_position_x,
-                   max_length_dim),
-        mtf.Shape([rows_dim, cols_dim, model_dim]))
+                   self.max_length_dim),
+        mtf.Shape([self.rows_dim, self.cols_dim, self.model_dim]))
 
     position_y = mtf.broadcast(
         mtf.gather(positional_emb_cols_var, targets_position_y,
-                   max_length_dim),
-        mtf.Shape([rows_dim, cols_dim, model_dim]))
+                   self.max_length_dim),
+        mtf.Shape([self.rows_dim, self.cols_dim, self.model_dim]))
     return position_x + position_y
 
   def mtf_model_fn(self, features, mesh):
     features = copy.copy(features)
     tf.logging.info("features = %s" % features)
     hparams = self._hparams
-    activation_dtype = self.set_activation_type()
+    activation_dtype = self.activation_type
 
     # We assume fixed vocab size for targets
-    targets_vocab_size = self._problem_hparams.target_modality._vocab_size  # pylint: disable=protected-access
     targets = tf.to_int32(features["targets"])
 
     # Image preprocessing, reshape into a 1D sequence and shift right.
@@ -99,22 +149,16 @@ def mtf_model_fn(self, features, mesh):
     shifted_targets = common_layers.shift_right_2d(targets)
 
     # Declare all the dimensions
-    model_dim = mtf.Dimension("d_model", hparams.hidden_size)
     batch_dim = mtf.Dimension("batch", hparams.batch_size)
-    length_dim = mtf.Dimension("length", length)
-    max_length_dim = mtf.Dimension("max_length", hparams.max_length)
-    filter_dim = mtf.Dimension("d_ff", hparams.d_ff)
-    kv_channels = mtf.Dimension("kv_channels", hparams.d_kv)
-    heads = mtf.Dimension("heads", hparams.num_heads)
 
     def import_to_batch_by_length(x, name):
       return mtf.import_tf_tensor(
-          mesh, x, mtf.Shape([batch_dim, length_dim]), name=name)
+          mesh, x, mtf.Shape([batch_dim, self.length_dim]), name=name)
 
     def layer_prepostprocess_dropout(x):
       return mtf.dropout(
           x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
-          noise_shape=mtf.Shape([batch_dim, model_dim]))
+          noise_shape=mtf.Shape([batch_dim, self.model_dim]))
 
     targets = import_to_batch_by_length(targets, "targets")
     shifted_targets = import_to_batch_by_length(
@@ -123,37 +167,32 @@ def layer_prepostprocess_dropout(x):
     extra_losses = []
 
     # Create targets content and position embeddings.
-    targets_vocab_size = 256 * hparams.num_channels
-    targets_vocab_dim = mtf.Dimension("vocab", targets_vocab_size)
-    outputs_vocab_dim = mtf.Dimension("output_vocab", 256)
-
     # Create embedding var for targets and positions and do a gather.
     targets_embedding_var = mtf.get_variable(
         mesh, "targets_embedding",
-        mtf.Shape([targets_vocab_dim, model_dim]),
+        mtf.Shape([self.targets_vocab_dim, self.model_dim]),
         initializer=tf.random_normal_initializer(),
         activation_dtype=activation_dtype)
 
-    x = mtf.gather(targets_embedding_var, shifted_targets, targets_vocab_dim)
+    x = mtf.gather(targets_embedding_var,
+                   shifted_targets, self.targets_vocab_dim)
     # Add positional embeddings
-    x += mtf.reshape(
-        self.create_positional_emb_2d(targets, max_length_dim, model_dim),
-        [length_dim, model_dim])
+    x += mtf.reshape(self.create_positional_emb_2d(targets),
+                     [self.length_dim, self.model_dim])
 
     # If conditional and input is given, add the input embedding to the target.
     # TODO(nikip): Verify conditional.
     if self.has_input and not hparams.unconditional:
-      vocab_size = hparams.num_classes
-      inputs_vocab_dim = mtf.Dimension("vocab", vocab_size)
       inputs = tf.squeeze(tf.to_int32(features["inputs"]), [2, 3])
       inputs = import_to_batch_by_length(inputs, "inputs")
 
       # Input embeddings
       inputs_embedding_var = mtf_layers.embedding(
           mesh, "input_embedding",
-          mtf.Shape([inputs_vocab_dim, model_dim]),
+          mtf.Shape([self.inputs_vocab_dim, self.model_dim]),
           activation_dtype=activation_dtype)
-      inputs_emb = mtf.gather(inputs_embedding_var, inputs, inputs_vocab_dim)
+      inputs_emb = mtf.gather(
+          inputs_embedding_var, inputs, self.inputs_vocab_dim)
       x += inputs_emb
 
     # Image Transformer Decoder
@@ -164,29 +203,37 @@ def layer_prepostprocess_dropout(x):
         # Self attention layer
         x += layer_prepostprocess_dropout(
             mtf_layers.masked_local_attention_1d(
-                mtf_layers.layer_norm(x, model_dim, name="layer_norm_self_att"),
+                mtf_layers.layer_norm(x, self.model_dim, name="layer_norm_att"),
                 None,
-                kv_channels,
-                heads,
+                self.kv_dim,
+                self.heads_dim,
                 block_length=hparams.block_length,
                 name="self_att"))
         # ffn layer
         x += layer_prepostprocess_dropout(mtf_layers.dense_relu_dense(
-            mtf_layers.layer_norm(x, model_dim, name="layer_norm_ffn"),
-            filter_dim, hparams.dropout, dropout_broadcast_dims=[length_dim]))
+            mtf_layers.layer_norm(x, self.model_dim, name="layer_norm_ffn"),
+            self.feedforward_dim,
+            hparams.dropout,
+            dropout_broadcast_dims=[self.length_dim]))
 
-    x = mtf_layers.layer_norm(x, model_dim, name="decoder_final_layer_norm")
+    x = mtf_layers.layer_norm(x, self.model_dim, name="final_layer_norm")
 
     # Calculate the logits and loss.
-    logits = mtf_layers.dense(x, outputs_vocab_dim, name="logits")
+    logits = mtf_layers.dense(x, self.outputs_vocab_dim, name="logits")
     soft_targets = mtf.one_hot(
-        targets, outputs_vocab_dim, dtype=activation_dtype)
+        targets, self.outputs_vocab_dim, dtype=activation_dtype)
     loss = mtf_layers.softmax_cross_entropy_with_logits(
-        logits, soft_targets, outputs_vocab_dim)
-
+        logits, soft_targets, self.outputs_vocab_dim)
     loss = mtf.reduce_mean(loss)
     for l in extra_losses:
       loss += l
+
+    # Reshape logits to original target shape.
+    logits = mtf.reshape(
+        logits,
+        mtf.Shape([batch_dim, self.rows_dim, self.orig_cols_dim,
+                   self.channels_dim, self.outputs_vocab_dim]))
+
     return logits, loss
 
 
@@ -224,7 +271,7 @@ def mtf_image_transformer_base():
   hparams.optimizer = "Adafactor"
   hparams.learning_rate_schedule = "rsqrt_decay"
   hparams.learning_rate_warmup_steps = 10000
-  hparams.add_hparam("d_kv", 32)
+  hparams.add_hparam("d_kv", 64)
   hparams.add_hparam("d_ff", 2048)
 
   # Image related hparams
@@ -243,10 +290,11 @@ def mtf_image_transformer_tiny():
   hparams.d_ff = 256
   hparams.batch_size = 4
   hparams.num_encoder_layers = 1
-  hparams.num_decoder_layers = 2
+  hparams.num_decoder_layers = 4
   hparams.num_heads = 4
   hparams.attention_key_size = 128
   hparams.attention_value_size = 128
+  hparams.block_length = 32
   # data parallelism and model-parallelism
   hparams.mesh_shape = "batch:2"
   hparams.layout = "batch:batch"
@@ -289,6 +337,7 @@ def mtf_image_transformer_base_cifar():
   hparams = mtf_image_transformer_base()
   hparams.mesh_shape = "batch:8"
   hparams.layout = "batch:batch"
+  hparams.learning_rate_decay_steps = 13600  # one epoch
   hparams.batch_size = 32
   hparams.num_heads = 4
   hparams.num_decoder_layers = 12
@@ -296,7 +345,6 @@ def mtf_image_transformer_base_cifar():
   hparams.hidden_size = 512
   hparams.d_ff = 2048
   hparams.learning_rate = 0.5
-  hparams.learning_rate_warmup_steps = 6000
   hparams.layer_preprocess_sequence = "none"
   hparams.layer_postprocess_sequence = "dan"
   hparams.layer_prepostprocess_dropout = 0.3
@@ -304,18 +352,40 @@ def mtf_image_transformer_base_cifar():
   return hparams
 
 
+@registry.register_hparams
+def mtf_image_transformer_cifar_4x():
+  """Data parallel CIFAR parameters."""
+  hparams = mtf_image_transformer_base_cifar()
+  hparams.mesh_shape = "batch:32"
+  hparams.layout = "batch:batch"
+  hparams.batch_size = 128
+  return hparams
+
+
+@registry.register_hparams
+def mtf_image_transformer_cifar_mp_4x():
+  """Data parallel CIFAR parameters."""
+  hparams = mtf_image_transformer_base_cifar()
+  hparams.mesh_shape = "model:4;batch:8"
+  hparams.layout = "batch:batch;d_ff:model;heads:model"
+  hparams.batch_size = 32
+  hparams.num_heads = 8
+  hparams.d_ff = 8192
+  return hparams
+
+
 @registry.register_hparams
 def mtf_image_transformer_base_imagenet():
   """Data parallel CIFAR parameters."""
   hparams = mtf_image_transformer_base_cifar()
   hparams.mesh_shape = "batch:32"
   hparams.layout = "batch:batch"
-  hparams.batch_size = 64
+  hparams.batch_size = 128
   hparams.d_ff = 2048
   hparams.hidden_size = 512
   hparams.num_decoder_layers = 12
   hparams.learning_rate = 0.5
-  hparams.learning_rate_warmup_steps = 6000
+  hparams.learning_rate_warmup_steps = 31250
   hparams.layer_preprocess_sequence = "none"
   hparams.layer_postprocess_sequence = "dan"
   hparams.layer_prepostprocess_dropout = 0.1
@@ -330,10 +400,9 @@ def mtf_image_transformer_base_imagenet_mp():
   hparams.mesh_shape = "model:4;batch:8"
   hparams.layout = "batch:batch;d_ff:model;heads:model"
   hparams.batch_size = 32
-  hparams.num_heads = 4
+  hparams.num_heads = 8
   hparams.d_ff = 8192
-  hparams.learning_rate_warmup_steps = 6000
-  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.learning_rate_warmup_steps = 31250
   hparams.unconditional = True
   return hparams
 
diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py b/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
index 70ca839ab..53197293c 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
@@ -29,7 +29,6 @@
 
 # Constants shared between all functions.
 BATCH_SIZE = 8
-INPUT_LENGTH = 8
 IMG_LENGTH = 8
 VOCAB_SIZE = 256
 
@@ -39,7 +38,7 @@ def get_model(hparams=None,
               model_cls=mtf_image_transformer.MtfImageTransformer):
   if hparams is None:
     hparams = mtf_image_transformer.mtf_image_transformer_single()
-  hparams.max_length = INPUT_LENGTH
+  hparams.max_length = IMG_LENGTH*IMG_LENGTH
   hparams.batch_size = BATCH_SIZE
   hparams.img_len = IMG_LENGTH
   hparams.num_channels = 1
@@ -49,7 +48,7 @@ def get_model(hparams=None,
   hparams.problem_hparams = p_hparams
 
   targets = -1 + np.random.random_integers(
-      VOCAB_SIZE, size=(BATCH_SIZE, IMG_LENGTH*IMG_LENGTH, 1, 1))
+      VOCAB_SIZE, size=(BATCH_SIZE, IMG_LENGTH, IMG_LENGTH, 1, 1))
   features = {
       "targets": tf.constant(targets, dtype=tf.int32, name="targets"),
   }
@@ -87,7 +86,9 @@ def testMtfImageTransformer(self):
       session.run(tf.global_variables_initializer())
       session.run(tf_group)
       res = session.run(tf_logits)
-    self.assertEqual(res.shape, (BATCH_SIZE, IMG_LENGTH*IMG_LENGTH, VOCAB_SIZE))
+    self.assertEqual(res.shape,
+                     (BATCH_SIZE, IMG_LENGTH, IMG_LENGTH,
+                      hparams.num_channels, VOCAB_SIZE))
 
   def testMtfImageTransformerDataParallel(self):
     hparams = mtf_image_transformer.mtf_image_transformer_single()
@@ -106,7 +107,9 @@ def testMtfImageTransformerDataParallel(self):
       session.run(tf.global_variables_initializer())
       session.run(tf_group)
       res = session.run(tf_logits)
-    self.assertEqual(res.shape, (BATCH_SIZE, IMG_LENGTH*IMG_LENGTH, VOCAB_SIZE))
+    self.assertEqual(res.shape,
+                     (BATCH_SIZE, IMG_LENGTH, IMG_LENGTH,
+                      hparams.num_channels, VOCAB_SIZE))
 
   def testMtfImageTransformerModelParallel(self):
     hparams = mtf_image_transformer.mtf_image_transformer_single()
@@ -125,8 +128,9 @@ def testMtfImageTransformerModelParallel(self):
       session.run(tf.global_variables_initializer())
       session.run(tf_group)
       res = session.run(tf_logits)
-    self.assertEqual(res.shape, (BATCH_SIZE, IMG_LENGTH*IMG_LENGTH, VOCAB_SIZE))
-
+    self.assertEqual(
+        res.shape,
+        (BATCH_SIZE, IMG_LENGTH, IMG_LENGTH, hparams.num_channels, VOCAB_SIZE))
 
 if __name__ == "__main__":
   tf.test.main()

From 59e14f56406cf27b434d0a553d474bbe1f237afa Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 10 Sep 2018 18:06:17 -0700
Subject: [PATCH 0802/2720] allow real input modality in Transformer.VAE (and
 fix bug)

PiperOrigin-RevId: 212371181
---
 .../models/research/transformer_vae.py        | 26 ++++++++++++++-----
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 3f5dc5012..24ad179d3 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -675,15 +675,24 @@ def infer(self, features=None, decode_length=50, beam_size=1, top_beams=1,
     if "partial_targets" in features:
       initial_output = tf.convert_to_tensor(features["partial_targets"])
     else:
-      batch_size = common_layers.shape_list(features["inputs"])[0]
-      length = common_layers.shape_list(features["inputs"])[1]
+      # inputs might not be present in features (e.g.: language modeling),
+      # in which case we fallback to 'infer_targets' for calculating initial
+      # input shape, type, etc.
+      inputs_or_targets = features.get("inputs", features.get("infer_targets"))
+      batch_size = common_layers.shape_list(inputs_or_targets)[0]
+      length = common_layers.shape_list(inputs_or_targets)[1]
+      hidden_dim = common_layers.shape_list(inputs_or_targets)[-1]
       target_length = tf.to_int32(2.0 * tf.to_float(length))
-      initial_output = tf.zeros((batch_size, target_length, 1, 1),
-                                dtype=tf.int64)
+      initial_output = tf.zeros((batch_size, target_length, 1, hidden_dim),
+                                dtype=inputs_or_targets.dtype)
 
     features["targets"] = initial_output
     logits, _ = self(features)  # pylint: disable=not-callable
-    samples = tf.argmax(logits, axis=-1)
+    # this should only happen if we're doing target_modality not real
+    if inputs_or_targets.dtype == tf.float32:
+      samples = logits
+    else:
+      samples = tf.argmax(logits, axis=-1)
 
     # More steps.
     self.predict_mask = 0.0  # Use the provided targets this time.
@@ -692,7 +701,12 @@ def infer(self, features=None, decode_length=50, beam_size=1, top_beams=1,
       with tf.variable_scope(tf.get_variable_scope(), reuse=True):
         features["targets"] = samples
         logits, _ = self(features)  # pylint: disable=not-callable
-        samples = tf.argmax(logits, axis=-1)
+        if inputs_or_targets.dtype == tf.float32:
+          # When target_modality is real, the last axis does not represent
+          # classes, so it should not be argmax'ed
+          samples = logits
+        else:
+          samples = tf.argmax(logits, axis=-1)
 
     self.predict_mask = 1.0
     if inputs_old is not None:  # Restore to not confuse Estimator.

From 4f495d64aebd7036a81a8c9cd28474eee38d1cde Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 10 Sep 2018 21:06:17 -0700
Subject: [PATCH 0803/2720] more cleanups

---
 .../rl/model_rl_experiment_player.py          | 94 +++++++++----------
 1 file changed, 43 insertions(+), 51 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index f1a7ba026..81d74e3e0 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -22,6 +22,7 @@
 from tensor2tensor.rl.trainer_model_based import FLAGS, setup_directories, temporary_flags
 from gym.utils import play
 import tensorflow as tf
+from gym.core import Env
 
 
 HP_SCOPES = ["loop", "model", "ppo"]
@@ -69,53 +70,6 @@ def concatenate_images(*imgs, axis=1):
   return _assert_image(concatenated_im_np)
 
 
-def show_agent(problem_name, agent_model_dir, world_model_dir, epoch_data_dir, hparams):
-  """Train the PPO agent in the simulated environment."""
-  gym_problem = registry.problem(problem_name)
-  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
-
-  # Adding model hparams for model specific adjustments
-  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
-  ppo_hparams.add_hparam("model_hparams", model_hparams)
-
-  environment_spec = copy.copy(gym_problem.environment_spec)
-  environment_spec.simulation_random_starts = hparams.simulation_random_starts
-  environment_spec.simulation_flip_first_random_for_beginning = False
-  environment_spec.intrinsic_reward_scale = hparams.intrinsic_reward_scale
-
-  ppo_hparams.add_hparam("environment_spec", environment_spec)
-  ppo_hparams.num_agents = 1
-
-  with temporary_flags({
-      "problem": problem_name,
-      "model": hparams.generative_model,
-      "hparams_set": hparams.generative_model_params,
-      "output_dir": world_model_dir,
-      "data_dir": epoch_data_dir,
-  }):
-    sess = tf.Session()
-    env = DebugBatchEnv(ppo_hparams, sess)
-    sess.run(tf.global_variables_initializer())
-    env.initialize()
-
-    env_model_loader = tf.train.Saver(
-      tf.global_variables("next_frame*"))
-    trainer_lib.restore_checkpoint(world_model_dir, env_model_loader, sess,
-      must_restore=True)
-
-    model_saver = tf.train.Saver(
-      tf.global_variables(".*network_parameters.*"))
-    trainer_lib.restore_checkpoint(agent_model_dir, model_saver, sess)
-
-    key_mapping = {(): 100, (ord('q'),):1, (ord('a'),):2,
-                   (ord('r'),):101,
-                   (ord('p'),):102}
-
-    play.play(env, zoom=2, fps=10, keys_to_action=key_mapping)
-
-from gym.core import Env
-
-
 class DebugBatchEnv(Env):
 
   def __init__(self, hparams, sess = None):
@@ -212,8 +166,9 @@ def _augment_observation(self):
     for p in probs:
       probs_str += "%.2f" % p +", "
 
-    action = np.argmax(probs)
+    probs_str = probs_str[:-1]
 
+    action = np.argmax(probs)
     info_str = " Policy:{}\n Action:{}\n Value function:{}\n Reward:{}".format(probs_str, action,
                                                                            vf, rew)
     info_pane = write_on_image(info_pane, info_str)
@@ -250,7 +205,7 @@ def step(self, action):
     return observ, rew, done, {"probs": probs, "vf": vf}
 
 
-def setup_show_metadata(hparams, output_dir):
+def play_agent(hparams, output_dir):
   """Setup metadata"""
 
   # Directories
@@ -276,9 +231,46 @@ def setup_show_metadata(hparams, output_dir):
   epoch_data_dir = os.path.join(directories["data"], str(epoch))
   ppo_model_dir = directories["ppo"]
 
-  show_agent(simulated_problem_name, ppo_model_dir,
-             directories["world_model"], epoch_data_dir, hparams)
+  world_model_dir = directories["world_model"]
+
+  gym_problem = registry.problem(simulated_problem_name)
+  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
+
+  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
+  ppo_hparams.add_hparam("model_hparams", model_hparams)
+
+  environment_spec = copy.copy(gym_problem.environment_spec)
+  environment_spec.simulation_random_starts = hparams.simulation_random_starts
+
+  ppo_hparams.add_hparam("environment_spec", environment_spec)
+  ppo_hparams.num_agents = 1
+
+  with temporary_flags({
+      "problem": simulated_problem_name,
+      "model": hparams.generative_model,
+      "hparams_set": hparams.generative_model_params,
+      "output_dir": world_model_dir,
+      "data_dir": epoch_data_dir,
+  }):
+    sess = tf.Session()
+    env = DebugBatchEnv(ppo_hparams, sess)
+    sess.run(tf.global_variables_initializer())
+    env.initialize()
+
+    env_model_loader = tf.train.Saver(
+      tf.global_variables("next_frame*"))
+    trainer_lib.restore_checkpoint(world_model_dir, env_model_loader, sess,
+      must_restore=True)
 
+    model_saver = tf.train.Saver(
+      tf.global_variables(".*network_parameters.*"))
+    trainer_lib.restore_checkpoint(ppo_model_dir, model_saver, sess)
+
+    key_mapping = {(): 100, (ord('q'),):1, (ord('a'),):2,
+                   (ord('r'),):101,
+                   (ord('p'),):102}
+
+    play.play(env, zoom=2, fps=10, keys_to_action=key_mapping)
 
 def main(_):
   hparams = registry.hparams(FLAGS.loop_hparams_set)

From 48a437571719df68ba7fb58923a405b6be6dbce6 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 10 Sep 2018 21:11:10 -0700
Subject: [PATCH 0804/2720] removing unnecessary call stack

---
 tensor2tensor/rl/model_rl_experiment_player.py | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index 81d74e3e0..c5dd1ba7c 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -205,10 +205,11 @@ def step(self, action):
     return observ, rew, done, {"probs": probs, "vf": vf}
 
 
-def play_agent(hparams, output_dir):
-  """Setup metadata"""
+def main(_):
+  hparams = registry.hparams(FLAGS.loop_hparams_set)
+  hparams.parse(FLAGS.loop_hparams)
+  output_dir = FLAGS.output_dir
 
-  # Directories
   subdirectories = ["data", "tmp", "world_model", "ppo"]
   using_autoencoder = hparams.autoencoder_train_steps > 0
   if using_autoencoder:
@@ -272,11 +273,6 @@ def play_agent(hparams, output_dir):
 
     play.play(env, zoom=2, fps=10, keys_to_action=key_mapping)
 
-def main(_):
-  hparams = registry.hparams(FLAGS.loop_hparams_set)
-  hparams.parse(FLAGS.loop_hparams)
-  output_dir = FLAGS.output_dir
-  setup_show_metadata(hparams, output_dir)
 
 
 if __name__ == "__main__":

From b582b071be6d44fcd13a09dd03787dc24d7cf5ff Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 10 Sep 2018 21:25:58 -0700
Subject: [PATCH 0805/2720] bugfix?

---
 tensor2tensor/rl/envs/simulated_batch_env.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 57d149bd5..0bc946542 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -131,6 +131,7 @@ def __init__(self, environment_spec, length):
                       environment_spec.video_num_target_frames,
                       environment_spec=environment_spec)
 
+    #TODO(piotrmilos): check if this should not be tf.estimator.ModeKeys.Predict
     initial_frames_dataset = initial_frames_problem.dataset(
         tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir, shuffle_files=False,
         hparams=hparams).take(1)

From 2034ca8e6925edde0f5bf0fb0059a04c802a9a34 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Mon, 10 Sep 2018 21:38:08 -0700
Subject: [PATCH 0806/2720] upload key mapping

---
 .../rl/model_rl_experiment_player.py          | 40 ++++++++-----------
 1 file changed, 17 insertions(+), 23 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index c5dd1ba7c..f73ed065b 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -25,8 +25,6 @@
 from gym.core import Env
 
 
-HP_SCOPES = ["loop", "model", "ppo"]
-
 _font = None
 FONT_SIZE = 20
 
@@ -72,16 +70,14 @@ def concatenate_images(*imgs, axis=1):
 
 class DebugBatchEnv(Env):
 
-  def __init__(self, hparams, sess = None):
-    if sess == None:
-      self.sess = tf.Session()
-    else:
-      self.sess = sess
+  INFO_PANE_WIDTH = 250
 
+  def __init__(self, hparams, sess=None):
     self.action_space = Discrete(6)
-    self.observation_space = Box(low=0, high=255, shape=(210, 160+250, 3), dtype=np.uint8)
+    self.observation_space = Box(low=0, high=255, shape=(210, 160+DebugBatchEnv.INFO_PANE_WIDTH, 3), dtype=np.uint8)
     self._tmp = 1
     self.res = None
+    self.sess = sess if sess is not None else tf.Session()
     self._prepare_networks(hparams, self.sess)
 
   def _prepare_networks(self, hparams, sess):
@@ -129,13 +125,11 @@ def _reset_env(self):
     #TODO:(put correct numbers)
     self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)
 
-
   def reset(self):
     self._reset_env()
     observ = self._augment_observation()
     return observ
 
-
   def _step_fake(self, action):
     observ = np.ones(shape=(210, 160, 3), dtype=np.uint8)*10*self._tmp
     observ[0, 0, 0] = 0
@@ -161,12 +155,12 @@ def _step_env(self, action):
 
   def _augment_observation(self):
     _observ, rew, done, probs, vf = self.res
-    info_pane = np.zeros(shape=(210, 250, 3), dtype=np.uint8)
+    info_pane = np.zeros(shape=(210, DebugBatchEnv.INFO_PANE_WIDTH, 3), dtype=np.uint8)
     probs_str = ""
     for p in probs:
       probs_str += "%.2f" % p +", "
 
-    probs_str = probs_str[:-1]
+    probs_str = probs_str[:-2]
 
     action = np.argmax(probs)
     info_str = " Policy:{}\n Action:{}\n Value function:{}\n Reward:{}".format(probs_str, action,
@@ -220,7 +214,7 @@ def main(_):
     game_with_mode = hparams.game + "_deterministic-v4"
   else:
     game_with_mode = hparams.game
-  # Problems
+
   if using_autoencoder:
     simulated_problem_name = (
         "gym_simulated_discrete_problem_with_agent_on_%s_autoencoded"
@@ -235,16 +229,15 @@ def main(_):
   world_model_dir = directories["world_model"]
 
   gym_problem = registry.problem(simulated_problem_name)
-  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
 
   model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
-  ppo_hparams.add_hparam("model_hparams", model_hparams)
-
   environment_spec = copy.copy(gym_problem.environment_spec)
   environment_spec.simulation_random_starts = hparams.simulation_random_starts
 
-  ppo_hparams.add_hparam("environment_spec", environment_spec)
-  ppo_hparams.num_agents = 1
+  batch_env_hparams = trainer_lib.create_hparams(hparams.ppo_params)
+  batch_env_hparams.add_hparam("model_hparams", model_hparams)
+  batch_env_hparams.add_hparam("environment_spec", environment_spec)
+  batch_env_hparams.num_agents = 1
 
   with temporary_flags({
       "problem": simulated_problem_name,
@@ -254,7 +247,7 @@ def main(_):
       "data_dir": epoch_data_dir,
   }):
     sess = tf.Session()
-    env = DebugBatchEnv(ppo_hparams, sess)
+    env = DebugBatchEnv(batch_env_hparams, sess)
     sess.run(tf.global_variables_initializer())
     env.initialize()
 
@@ -267,14 +260,15 @@ def main(_):
       tf.global_variables(".*network_parameters.*"))
     trainer_lib.restore_checkpoint(ppo_model_dir, model_saver, sess)
 
-    key_mapping = {(): 100, (ord('q'),):1, (ord('a'),):2,
-                   (ord('r'),):101,
-                   (ord('p'),):102}
+    key_mapping = gym_problem.env.env.get_keys_to_action()
+    #map special codes
+    key_mapping[()] = 100
+    key_mapping[(ord('r'),)] = 101
+    key_mapping[(ord('p'),)] = 102
 
     play.play(env, zoom=2, fps=10, keys_to_action=key_mapping)
 
 
-
 if __name__ == "__main__":
   tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()

From ae01963559991b410e2b33ad6d280a267ade8281 Mon Sep 17 00:00:00 2001
From: Marcin Michalski <michalski@google.com>
Date: Mon, 10 Sep 2018 22:43:19 -0700
Subject: [PATCH 0807/2720] Allow exporting T2T models as tfhub modules.

PiperOrigin-RevId: 212395830
---
 setup.py                        |  8 +++-
 tensor2tensor/serving/export.py | 75 ++++++++++++++++++++++++++++++++-
 2 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/setup.py b/setup.py
index 78d6de769..325b69673 100644
--- a/setup.py
+++ b/setup.py
@@ -51,8 +51,14 @@
     extras_require={
         'tensorflow': ['tensorflow>=1.9.0'],
         'tensorflow_gpu': ['tensorflow-gpu>=1.9.0'],
+        'tensorflow-hub': ['tensorflow-hub>=0.1.1'],
         'tests': [
-            'absl-py', 'pytest>=3.8.0', 'mock', 'pylint', 'jupyter', 'gsutil',
+            'absl-py',
+            'pytest>=3.8.0',
+            'mock',
+            'pylint',
+            'jupyter',
+            'gsutil',
             'matplotlib',
             # Need atari extras for Travis tests, but because gym is already in
             # install_requires, pip skips the atari extras, so we instead do an
diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index b5856d3ec..52b1dbe4c 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -20,13 +20,18 @@
 import os
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.utils import decoding
+from tensor2tensor.utils import t2t_model
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
 
 import tensorflow as tf
+import tensorflow_hub as hub
 
 FLAGS = tf.flags.FLAGS
 
+tf.flags.DEFINE_bool("export_as_tfhub", False,
+                     "If True, the model will be exported as tfHub module.")
+
 
 def create_estimator(run_config, hparams):
   return trainer_lib.create_estimator(
@@ -44,6 +49,68 @@ def create_hparams():
       problem_name=FLAGS.problem)
 
 
+# TODO(michalski): Move this method into tfhub utils.
+def export_module_spec_with_checkpoint(module_spec,
+                                       checkpoint_path,
+                                       export_path,
+                                       scope_prefix=""):
+  """Exports given checkpoint as tfhub module with given spec."""
+
+  # The main requirement is that it is possible to know how to map from
+  # module variable name to checkpoint variable name.
+  # This is trivial if the original code used variable scopes,
+  # but can be messy if the variables to export are interwined
+  # with variables not export.
+  with tf.Graph().as_default():
+    m = hub.Module(module_spec)
+    assign_map = {
+        scope_prefix + name: value for name, value in m.variable_map.items()
+    }
+    tf.train.init_from_checkpoint(checkpoint_path, assign_map)
+    init_op = tf.initializers.global_variables()
+    with tf.Session() as session:
+      session.run(init_op)
+      m.export(export_path, session)
+
+
+def export_as_tfhub_module(hparams, problem, ckpt_dir):
+  """Exports the last checkpoint from the directory as tfhub module.
+
+  It creates the Module spec and signature (based on T2T problem information),
+  which is later used to create and export the hub module.
+  Module will be saved inside the ckpt_dir.
+
+  Args:
+    hparams: T2T parameters, model graph will be based on them.
+    problem: the name of the problem
+    ckpt_dir: directory with the checkpoints. The final model will be exported
+      there too.
+  """
+
+  def hub_module_fn():
+    """Creates the TF graph for the hub module."""
+    model_fn = t2t_model.T2TModel.make_estimator_model_fn(
+        FLAGS.model,
+        hparams,
+        decode_hparams=decoding.decode_hparams(FLAGS.decode_hparams))
+    features = problem.serving_input_fn(hparams).features
+    spec = model_fn(features, labels=None, mode=tf.estimator.ModeKeys.PREDICT)
+
+    # Currently only supports a single input and single output.
+    hub.add_signature(
+        inputs=features, outputs=spec.export_outputs["serving_default"].outputs)
+
+  module_spec = hub.create_module_spec(hub_module_fn)
+  export_dir = os.path.join(ckpt_dir, "export_tfhub")
+  # Loads the weights from the checkpoint using the model above
+  # and saves it in the export_path.
+  export_module_spec_with_checkpoint(
+      module_spec,
+      checkpoint_path=tf.train.latest_checkpoint(ckpt_dir),
+      export_path=export_dir,
+      scope_prefix="")
+
+
 def main(_):
   tf.logging.set_verbosity(tf.logging.INFO)
   trainer_lib.set_random_seed(FLAGS.random_seed)
@@ -53,12 +120,16 @@ def main(_):
 
   hparams = create_hparams()
   hparams.no_data_parallelism = True  # To clear the devices
+  problem = hparams.problem
+
+  if FLAGS.export_as_tfhub:
+    export_as_tfhub_module(hparams, problem, ckpt_dir)
+    return
+
   run_config = t2t_trainer.create_run_config(hparams)
 
   estimator = create_estimator(run_config, hparams)
 
-  problem = hparams.problem
-
   exporter = tf.estimator.FinalExporter(
       "exporter", lambda: problem.serving_input_fn(hparams), as_text=True)
 

From 14018d5b29983d4ea96f89ea28aa776c50c52cdd Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 11 Sep 2018 08:00:57 -0700
Subject: [PATCH 0808/2720] style fixes

---
 .../rl/model_rl_experiment_player.py          | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index f73ed065b..e98f7695f 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -28,6 +28,7 @@
 _font = None
 FONT_SIZE = 20
 
+
 def _get_font():
   global _font
   if _font is None:
@@ -107,7 +108,6 @@ def initialization_lambda():
     self.policy_probs = actor_critic.policy.probs[0, 0, :]
     self.value = actor_critic.value[0, :]
 
-
   def render(self, mode='human'):
     raise NotImplemented()
 
@@ -158,7 +158,7 @@ def _augment_observation(self):
     info_pane = np.zeros(shape=(210, DebugBatchEnv.INFO_PANE_WIDTH, 3), dtype=np.uint8)
     probs_str = ""
     for p in probs:
-      probs_str += "%.2f" % p +", "
+      probs_str += "%.2f" % p + ", "
 
     probs_str = probs_str[:-2]
 
@@ -171,7 +171,6 @@ def _augment_observation(self):
     augmented_observ = np.array(augmented_observ)
     return augmented_observ
 
-
   def step(self, action):
     #Special codes
     if action==100:
@@ -251,10 +250,15 @@ def main(_):
     sess.run(tf.global_variables_initializer())
     env.initialize()
 
-    env_model_loader = tf.train.Saver(
-      tf.global_variables("next_frame*"))
-    trainer_lib.restore_checkpoint(world_model_dir, env_model_loader, sess,
-      must_restore=True)
+    # env_model_loader = tf.train.Saver(
+    #   tf.global_variables("next_frame*"))
+    # trainer_lib.restore_checkpoint(world_model_dir, env_model_loader, sess,
+    #   must_restore=True)
+
+    ckpt = tf.train.get_checkpoint_state(world_model_dir)
+    new_saver = tf.train.import_meta_graph("/Users/piotr.milos/Downloads/18/world_model/model.ckpt-130000.meta", clear_devices=True)
+    new_saver.restore(sess, "/Users/piotr.milos/Downloads/18/world_model/model.ckpt-130000")
+
 
     model_saver = tf.train.Saver(
       tf.global_variables(".*network_parameters.*"))
@@ -269,6 +273,7 @@ def main(_):
     play.play(env, zoom=2, fps=10, keys_to_action=key_mapping)
 
 
+
 if __name__ == "__main__":
   tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()

From 71634aa4c6cafdd9784736a740dd61bdb42ab900 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 11 Sep 2018 09:12:07 -0700
Subject: [PATCH 0809/2720] Internal change

PiperOrigin-RevId: 212462664
---
 tensor2tensor/models/transformer.py | 22 +++++++++++++++++++++-
 tensor2tensor/utils/optimize.py     |  4 ----
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 3175eb1ac..d7d6314fc 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1660,7 +1660,7 @@ def transformer_ada_lmpackedbase_relative():
 
 
 @registry.register_hparams
-def transformer_base():
+def transformer_base_v3():
   """Base parameters for Transformer model."""
   # Update parameters here, then occasionally cut a versioned set, e.g.
   # transformer_base_v2.
@@ -1674,6 +1674,13 @@ def transformer_base():
   return hparams
 
 
+@registry.register_hparams
+def transformer_base():
+  """Base parameters for Transformer model."""
+  hparams = transformer_base_v3()
+  return hparams
+
+
 @registry.register_hparams
 def transformer_big():
   """HParams for transformer big model on WMT."""
@@ -2029,6 +2036,19 @@ def transformer_timeseries():
   return hparams
 
 
+@registry.register_hparams
+def transformer_mlperf_tpu():
+  """HParams for Transformer model on TPU for MLPerf on TPU 2x2."""
+  hparams = transformer_base_v3()
+  hparams.symbol_modality_num_shards = 1
+  hparams.max_length = 64  # ignored when using "_packed" problems
+  hparams.batch_size = 512  # gloabl batch size matches the reference model
+  hparams.attention_dropout_broadcast_dims = "0,1"  # batch, heads
+  hparams.relu_dropout_broadcast_dims = "1"  # length
+  hparams.layer_prepostprocess_dropout_broadcast_dims = "1"  # length
+  return hparams
+
+
 def update_hparams_for_tpu(hparams):
   """Change hparams to be compatible with TPU training."""
 
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index b2ac2a325..c3164c3f9 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -76,10 +76,6 @@ class ConditionalOptimizer(tf.train.Optimizer):
   """Conditional optimizer."""
 
   def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disable=super-init-not-called
-    if optimizer_name == "Adam" and use_tpu:
-      # LazyAdamOptimizer does not work on TPU
-      optimizer_name = "TrueAdam"
-
     tf.logging.info("Using optimizer %s", optimizer_name)
 
     if optimizer_name == "Adam":

From c5bbc3c4cc57a25889011428456b37a025e415b6 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Tue, 11 Sep 2018 09:57:05 -0700
Subject: [PATCH 0810/2720] Add SciTail dataset.

PiperOrigin-RevId: 212470002
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/scitail.py      | 128 ++++++++++++++++++
 2 files changed, 129 insertions(+)
 create mode 100644 tensor2tensor/data_generators/scitail.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index dd3ee3dd8..e9fab0b2a 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -62,6 +62,7 @@
     "tensor2tensor.data_generators.qnli",
     "tensor2tensor.data_generators.quora_qpairs",
     "tensor2tensor.data_generators.rte",
+    "tensor2tensor.data_generators.scitail",
     "tensor2tensor.data_generators.snli",
     "tensor2tensor.data_generators.stanford_nli",
     "tensor2tensor.data_generators.style_transfer",
diff --git a/tensor2tensor/data_generators/scitail.py b/tensor2tensor/data_generators/scitail.py
new file mode 100644
index 000000000..789246481
--- /dev/null
+++ b/tensor2tensor/data_generators/scitail.py
@@ -0,0 +1,128 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Data generators for SciTail."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import zipfile
+import six
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import lm1b
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+EOS = text_encoder.EOS
+
+
+@registry.register_problem
+class SciTail(text_problems.TextConcat2ClassProblem):
+  """SciTail classification problems."""
+
+  # Data from allen institute for AI.
+  _SCITAIL_URL = ("http://data.allenai.org.s3.amazonaws.com/"
+                  "downloads/SciTailV1.1.zip")
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 10,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  @property
+  def approx_vocab_size(self):
+    return 2**13
+
+  @property
+  def num_classes(self):
+    return 2
+
+  def class_labels(self, data_dir):
+    del data_dir
+    # Note this binary classification is different from usual SNLI.
+    return ["neutral", "entails"]
+
+  def _maybe_download_corpora(self, tmp_dir):
+    scitail_filename = "SciTailV1.1.zip"
+    scitail_finalpath = os.path.join(tmp_dir, "SciTailV1.1")
+    if not tf.gfile.Exists(scitail_finalpath):
+      zip_filepath = generator_utils.maybe_download(
+          tmp_dir, scitail_filename, self._SCITAIL_URL)
+      zip_ref = zipfile.ZipFile(zip_filepath, "r")
+      zip_ref.extractall(tmp_dir)
+      zip_ref.close()
+
+    return scitail_finalpath
+
+  def example_generator(self, filename):
+    label_list = self.class_labels(data_dir=None)
+    for line in tf.gfile.Open(filename, "rb"):
+      if six.PY2:
+        line = unicode(line.strip(), "utf-8")
+      else:
+        line = line.strip().decode("utf-8")
+      split_line = line.split("\t")
+      s1, s2 = split_line[:2]
+      l = label_list.index(split_line[2])
+      inputs = [s1, s2]
+      yield {
+          "inputs": inputs,
+          "label": l
+      }
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    scitail_dir = self._maybe_download_corpora(tmp_dir)
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      filesplit = "tsv_format/scitail_1.0_train.tsv"
+    else:
+      filesplit = "tsv_format/scitail_1.0_dev.tsv"
+
+    filename = os.path.join(scitail_dir, filesplit)
+    for example in self.example_generator(filename):
+      yield example
+
+
+@registry.register_problem
+class SciTailCharacters(SciTail):
+  """SciTail classification problems, character level"""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  def global_task_id(self):
+    return problem.TaskID.EN_NLI
+
+
+@registry.register_problem
+class SciTailSharedVocab(SciTail):
+  """SciTail classification problems with the LM1b vocabulary"""
+
+  @property
+  def vocab_filename(self):
+    return lm1b.LanguagemodelLm1b32k().vocab_filename

From 92cefda13249ba9d81e8c3dfb93bb3ef5c7b5aeb Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 11 Sep 2018 11:00:53 -0700
Subject: [PATCH 0811/2720] player prepared for PR

---
 .../rl/model_rl_experiment_player.py          | 31 +++++++++++++------
 1 file changed, 21 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index e98f7695f..ad7cad49e 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -1,3 +1,18 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
 #--output_dir=/Users/piotr.milos/Downloads/18 --alsologtostderr --loop_hparams_set=rl_modelrl_base_quick
 #--output_dir=/Users/piotr.milos/t2t/rl_v1 --alsologtostderr --loop_hparams_set=rl_modelrl_tiny
 
@@ -165,7 +180,8 @@ def _augment_observation(self):
     action = np.argmax(probs)
     info_str = " Policy:{}\n Action:{}\n Value function:{}\n Reward:{}".format(probs_str, action,
                                                                            vf, rew)
-    info_pane = write_on_image(info_pane, info_str)
+    print("Info str:{}".format(info_str))
+    # info_pane = write_on_image(info_pane, info_str)
 
     augmented_observ = concatenate_images(_observ, info_pane)
     augmented_observ = np.array(augmented_observ)
@@ -250,15 +266,10 @@ def main(_):
     sess.run(tf.global_variables_initializer())
     env.initialize()
 
-    # env_model_loader = tf.train.Saver(
-    #   tf.global_variables("next_frame*"))
-    # trainer_lib.restore_checkpoint(world_model_dir, env_model_loader, sess,
-    #   must_restore=True)
-
-    ckpt = tf.train.get_checkpoint_state(world_model_dir)
-    new_saver = tf.train.import_meta_graph("/Users/piotr.milos/Downloads/18/world_model/model.ckpt-130000.meta", clear_devices=True)
-    new_saver.restore(sess, "/Users/piotr.milos/Downloads/18/world_model/model.ckpt-130000")
-
+    env_model_loader = tf.train.Saver(
+      tf.global_variables("next_frame*"))
+    trainer_lib.restore_checkpoint(world_model_dir, env_model_loader, sess,
+      must_restore=True)
 
     model_saver = tf.train.Saver(
       tf.global_variables(".*network_parameters.*"))

From 959fc004719036089808632013086d01ae118ab7 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 11 Sep 2018 11:38:18 -0700
Subject: [PATCH 0812/2720] dynamic beta adjustment.

PiperOrigin-RevId: 212490167
---
 tensor2tensor/models/video/base_vae.py        | 51 +++++++++++++++----
 .../models/video/basic_stochastic.py          |  3 ++
 tensor2tensor/models/video/savp.py            |  4 +-
 tensor2tensor/models/video/sv2p.py            | 14 +++--
 4 files changed, 54 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/models/video/base_vae.py b/tensor2tensor/models/video/base_vae.py
index ff27d76f6..8af32e374 100644
--- a/tensor2tensor/models/video/base_vae.py
+++ b/tensor2tensor/models/video/base_vae.py
@@ -37,21 +37,50 @@ def get_iteration_num(self):
       step_num = 1000000
     return step_num
 
-  def get_beta(self):
-    beta = common_video.beta_schedule(
-        schedule=self.hparams.latent_loss_multiplier_schedule,
-        global_step=self.get_iteration_num(),
-        final_beta=self.hparams.latent_loss_multiplier,
-        decay_start=(self.hparams.num_iterations_1st_stage +
-                     self.hparams.num_iterations_2nd_stage),
-        decay_end=self.hparams.anneal_end)
-    tf.summary.scalar("beta", beta)
-    return beta
+  def get_beta(self, kl_loss=0.0):
+    """Get the KL multiplier, either dynamically or schedule based.
+
+    if hparams.latent_loss_multiplier_dynamic is set to true, then beta
+    is being adjusted to keep KL under hparams.latent_loss_multiplier_epsilon.
+    In order to do so, the beta is being updated at each iteration
+    by taking steps of size hparams.latent_loss_multiplier_alpha.
+    The same formulation can be retrieved by solving the Lagrangian
+    with KL < epsilon as a constraint.
+
+    Args:
+      kl_loss: KL loss. Only used for dynamic adjustment.
+
+    Returns:
+      beta: the final value of beta.
+
+    """
+    if self.hparams.latent_loss_multiplier_dynamic:
+      beta = tf.Variable(self.hparams.latent_loss_multiplier,
+                         trainable=False, dtype=tf.float32)
+      alpha = self.hparams.latent_loss_multiplier_alpha
+      epsilon = self.hparams.latent_loss_multiplier_epsilon
+      shadow_beta = beta + alpha * (kl_loss - epsilon)
+      # Caping the beta between 0 and 1. May need to change this later on.
+      shadow_beta = tf.maximum(shadow_beta, 0.0)
+      shadow_beta = tf.minimum(shadow_beta, 1.0)
+      update_op = tf.assign(beta, shadow_beta)
+    else:
+      beta = common_video.beta_schedule(
+          schedule=self.hparams.latent_loss_multiplier_schedule,
+          global_step=self.get_iteration_num(),
+          final_beta=self.hparams.latent_loss_multiplier,
+          decay_start=(self.hparams.num_iterations_1st_stage +
+                       self.hparams.num_iterations_2nd_stage),
+          decay_end=self.hparams.anneal_end)
+      update_op = tf.identity(beta)  # fake update for regular beta.
+    with tf.control_dependencies([update_op]):
+      tf.summary.scalar("beta", beta)
+      return beta
 
   def get_extra_loss(self, mean, std):
     """Losses in addition to the default modality losses."""
-    beta = self.get_beta()
     kl_loss = common_layers.kl_divergence(mean, std)
+    beta = self.get_beta(kl_loss)
     tf.summary.histogram("posterior_mean", mean)
     tf.summary.histogram("posterior_std", std)
     tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 450713acb..a1c5e2847 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -114,6 +114,9 @@ def next_frame_basic_stochastic():
   hparams.add_hparam("num_iterations_1st_stage", 25000)
   hparams.add_hparam("num_iterations_2nd_stage", 25000)
   hparams.add_hparam("latent_loss_multiplier", 1e-3)
+  hparams.add_hparam("latent_loss_multiplier_dynamic", False)
+  hparams.add_hparam("latent_loss_multiplier_alpha", 1e-5)
+  hparams.add_hparam("latent_loss_multiplier_epsilon", 1.0)
   hparams.add_hparam("latent_loss_multiplier_schedule", "constant")
   hparams.add_hparam("latent_num_frames", 0)  # 0 means use all frames.
   hparams.add_hparam("anneal_end", 100000)
diff --git a/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index 5cd2afb50..f1a8d4986 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -262,7 +262,7 @@ def get_gan_loss(self, true_frames, gen_frames, name):
     return gan_loss
 
   def get_extra_loss(self, latent_means=None, latent_stds=None,
-                     true_frames=None, gen_frames=None, beta=1.0):
+                     true_frames=None, gen_frames=None):
     if not self.is_training:
       return 0.0
 
@@ -270,7 +270,7 @@ def get_extra_loss(self, latent_means=None, latent_stds=None,
     # Use sv2p's KL divergence computation.
     if self.hparams.use_vae:
       vae_loss = super(NextFrameSAVP, self).get_extra_loss(
-          latent_means=latent_means, latent_stds=latent_stds, beta=beta)
+          latent_means=latent_means, latent_stds=latent_stds)
 
     if self.hparams.use_gan:
       # Strip out the first context_frames for the true_frames
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index d13120c63..9298599a2 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -467,8 +467,9 @@ def process_single_frame(prev_outputs, inputs):
     else:
       return gen_images, gen_rewards, None, None
 
-  def get_extra_loss(self, latent_means=None, latent_stds=None,
-                     true_frames=None, gen_frames=None, beta=1.0):
+  def get_extra_loss(self,
+                     latent_means=None, latent_stds=None,
+                     true_frames=None, gen_frames=None):
     """Losses in addition to the default modality losses."""
     del true_frames
     del gen_frames
@@ -479,7 +480,10 @@ def get_extra_loss(self, latent_means=None, latent_stds=None,
         tf.summary.histogram("posterior_mean_%d" % i, mean)
         tf.summary.histogram("posterior_std_%d" % i, std)
       tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
-    return beta * kl_loss
+
+    beta = self.get_beta(kl_loss)
+    extra_loss = beta * kl_loss
+    return extra_loss
 
   def infer(self, features, *args, **kwargs):
     """Produce predictions from the model by running it."""
@@ -544,10 +548,10 @@ def body(self, features):
         rewards=all_rewards,
     )
 
-    beta = self.get_beta()
     extra_loss = self.get_extra_loss(
         latent_means=latent_means,
-        latent_stds=latent_stds, beta=beta, true_frames=all_frames,
+        latent_stds=latent_stds,
+        true_frames=all_frames,
         gen_frames=gen_images)
 
     # Visualize predictions in Tensorboard

From 373c50b1ab9f6f517f3ba5f1c4939c2f70095d9b Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 11 Sep 2018 12:59:53 -0700
Subject: [PATCH 0813/2720] Internal change

PiperOrigin-RevId: 212504087
---
 tensor2tensor/utils/t2t_model.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 4792a6f14..0ea74540f 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -174,6 +174,10 @@ def hparams(self):
   def is_training(self):
     return self._hparams.mode == tf.estimator.ModeKeys.TRAIN
 
+  @property
+  def is_predicting(self):
+    return self._hparams.mode == tf.estimator.ModeKeys.PREDICT
+
   @property
   def has_input(self):
     if self._problem_hparams:

From c371d228d1509ba57537a6cd4047fbadc5fcd2ff Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 11 Sep 2018 14:36:10 -0700
Subject: [PATCH 0814/2720] Corrections to RL and besic next frame model.

PiperOrigin-RevId: 212522977
---
 tensor2tensor/layers/common_layers.py         |  6 +--
 tensor2tensor/layers/modalities.py            | 46 +++----------------
 tensor2tensor/models/research/rl.py           |  2 +-
 .../models/video/basic_deterministic.py       | 21 ++++++---
 .../video/basic_deterministic_params.py       |  3 +-
 tensor2tensor/models/video/sv2p_params.py     |  1 +
 tensor2tensor/rl/trainer_model_based.py       |  4 +-
 7 files changed, 28 insertions(+), 55 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index b7ab9dacc..71546e6de 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -257,10 +257,10 @@ def standardize_images(x):
   with tf.name_scope("standardize_images", [x]):
     x_shape = shape_list(x)
     x = tf.to_float(tf.reshape(x, [-1] + x_shape[-3:]))
-    x_mean = tf.reduce_mean(x, axis=[1, 2, 3], keepdims=True)
+    x_mean = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
     x_variance = tf.reduce_mean(
-        tf.square(x - x_mean), axis=[1, 2, 3], keepdims=True)
-    num_pixels = tf.to_float(x_shape[-1] * x_shape[-2] * x_shape[-3])
+        tf.square(x - x_mean), axis=[1, 2], keepdims=True)
+    num_pixels = tf.to_float(x_shape[-2] * x_shape[-3])
     x = (x - x_mean) / tf.maximum(tf.sqrt(x_variance), tf.rsqrt(num_pixels))
     return tf.reshape(x, x_shape)
 
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index a896ccefe..ca19e8c29 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -527,9 +527,9 @@ def xnet_resblock(x, filters, res_relu, name):
                            "compress_block_final")
 
 
+@registry.register_video_modality("default")
 class VideoModality(modality.Modality):
   """Modality for videos, i.e., time-sequences of frames."""
-  PIXEL_EMBEDDING_SIZE = 64
 
   def bottom(self, x):
     inputs = x
@@ -538,26 +538,8 @@ def bottom(self, x):
       inputs = common_layers.standardize_images(inputs)
       return common_layers.time_to_channels(inputs)
 
-  def targets_bottom(self, x, summary_prefix="targets_bottom"):  # pylint: disable=arguments-differ
-    inputs = x
-    with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
-      common_layers.summarize_video(inputs, summary_prefix)
-      inputs_shape = common_layers.shape_list(inputs)
-      # We embed each of 256=self.top_dimensionality possible pixel values.
-      embedding_var = tf.get_variable(
-          "pixel_embedding",
-          [self.top_dimensionality, self.PIXEL_EMBEDDING_SIZE])
-      hot_inputs = tf.one_hot(tf.to_int32(inputs), self.top_dimensionality)
-      hot_inputs = tf.reshape(hot_inputs, [-1, self.top_dimensionality])
-      embedded = tf.matmul(hot_inputs, embedding_var)
-      # Let's now merge all channels that were embedded into a single vector.
-      merged_size = self.PIXEL_EMBEDDING_SIZE * inputs_shape[4]
-      embedded = tf.reshape(embedded, inputs_shape[:4] + [merged_size])
-      transposed = common_layers.time_to_channels(embedded)
-      return tf.layers.dense(
-          transposed,
-          self._body_input_depth,
-          name="merge_pixel_embedded_frames")
+  def targets_bottom(self, x):
+    return self.bottom(x)
 
   def top(self, body_output, targets):
     num_channels = self._model_hparams.problem.num_channels
@@ -568,9 +550,9 @@ def top(self, body_output, targets):
     # then you need to reshape to [..., num_frames, depth] like below, not
     # into [..., depth, num_frames] due to memory layout of concat/reshape.
     reshape_shape = body_output_shape[:-1] + [
-        num_channels, num_frames, self.top_dimensionality]
+        num_frames, num_channels, self.top_dimensionality]
     res = tf.reshape(body_output, reshape_shape)
-    res = tf.transpose(res, [0, 4, 1, 2, 3, 5])
+    res = tf.transpose(res, [0, 3, 1, 2, 4, 5])
     res_shape = common_layers.shape_list(res)
     res_argmax = tf.argmax(tf.reshape(res, [-1, res_shape[-1]]), axis=-1)
     res_argmax = tf.reshape(res_argmax, res_shape[:-1])
@@ -591,26 +573,10 @@ def loss(self, top_out, targets):
         weights_fn=self.targets_weights_fn)
 
 
-@registry.register_video_modality("default")
-class VideoModalityNoEmbed(VideoModality):
-  """Video Modality where target_bottom does not embeds pixels."""
-
-  def targets_bottom(self, x):
-    return super(VideoModalityNoEmbed, self).bottom(x)
-
-
-@registry.register_video_modality("embed")
-class VideoModalityEmbed(VideoModality):
-  """Video Modality where bottom embeds pixels."""
-
-  def bottom(self, x):
-    return super(VideoModalityEmbed, self).targets_bottom(
-        x, summary_prefix="bottom")
-
-
 @registry.register_video_modality("bitwise")
 class VideoModalityBitwise(VideoModality):
   """Video Modality where bottom embeds pixels bitwise."""
+  PIXEL_EMBEDDING_SIZE = 64
 
   def bottom(self, x):
     inputs = x
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index ab373e06d..494717582 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -53,7 +53,7 @@ def ppo_base_v1():
   hparams.add_hparam("simulation_random_starts", False)
   hparams.add_hparam("simulation_flip_first_random_for_beginning", False)
   hparams.add_hparam("intrinsic_reward_scale", 0.)
-  hparams.add_hparam("logits_clip", 0.)
+  hparams.add_hparam("logits_clip", 3.0)
   return hparams
 
 
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 5501e523c..144708ef5 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -33,6 +33,17 @@
 tfcl = tf.contrib.layers
 
 
+def inject_action(action, x):
+  """Inject the action into x."""
+  x_shape = common_layers.shape_list(x)
+  filters = x_shape[-1]
+  action_mask = tf.layers.dense(action, filters, name="action_mask")
+  action_add = tf.layers.dense(action, filters, name="action_add")
+  x *= tf.nn.sigmoid(action_mask)
+  x += action_add
+  return x
+
+
 @registry.register_model
 class NextFrameBasicDeterministic(t2t_model.T2TModel):
   """Basic next-frame model, may take actions and predict rewards too."""
@@ -72,13 +83,7 @@ def body_single(self, features):
     if "input_action" in features:
       action = tf.reshape(features["input_action"][:, -1, :],
                           [-1, 1, 1, hparams.hidden_size])
-      action_mask = tf.layers.dense(action, filters, name="action_mask")
-      zeros_mask = tf.zeros(common_layers.shape_list(x)[:-1] + [filters],
-                            dtype=tf.float32)
-      if hparams.concatenate_actions:
-        x = tf.concat([x, action_mask + zeros_mask], axis=-1)
-      else:
-        x *= action_mask + zeros_mask
+      x = inject_action(action, x)
 
     x, extra_loss = self.inject_latent(x, features, filters)
 
@@ -97,6 +102,8 @@ def body_single(self, features):
     layer_inputs = list(reversed(layer_inputs))
     for i in range(hparams.num_compress_steps):
       with tf.variable_scope("upstride%d" % i):
+        if "input_action" in features:
+          x = inject_action(action, x)
         if i >= hparams.num_compress_steps - hparams.filter_double_steps:
           filters //= 2
         x = tf.layers.conv2d_transpose(
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 61bd2fd03..c19c25985 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -44,7 +44,6 @@ def next_frame_basic_deterministic():
   hparams.add_hparam("filter_double_steps", 2)
   hparams.add_hparam("video_modality_loss_cutoff", 0.02)
   hparams.add_hparam("preprocess_resize_frames", None)
-  hparams.add_hparam("concatenate_actions", True)
   hparams.add_hparam("shuffle_buffer_size", 128)
   hparams.add_hparam("tiny_mode", False)
   hparams.add_hparam("stochastic_model", False)
@@ -64,7 +63,7 @@ def next_frame_pixel_noise():
 def next_frame_sampling():
   """Basic conv model with scheduled sampling."""
   hparams = next_frame_basic_deterministic()
-  hparams.video_num_target_frames = 2
+  hparams.video_num_target_frames = 4
   hparams.scheduled_sampling_warmup_steps = 50000
   hparams.scheduled_sampling_prob = 0.5
   return hparams
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 74f98c365..306eca995 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -48,6 +48,7 @@ def next_frame_sv2p():
   hparams.add_hparam("scheduled_sampling_k", 900.0)
   hparams.add_hparam("upsample_method", "conv2d_transpose")
   hparams.add_hparam("internal_loss", False)
+  hparams.add_hparam("concatenate_actions", True)
   return hparams
 
 
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index f40c2340f..dd5c7ecc0 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -570,7 +570,7 @@ def combine_training_data(problem, final_data_dir, old_data_dirs,
 @registry.register_hparams
 def rl_modelrl_base():
   return tf.contrib.training.HParams(
-      epochs=6,
+      epochs=3,
       # Total frames used for training. This will be distributed evenly across
       # hparams.epochs.
       num_real_env_frames=100000,
@@ -598,7 +598,7 @@ def rl_modelrl_base():
       # should start fresh each time.
       ppo_continue_training=True,
 
-      real_ppo_epochs_num=30,
+      real_ppo_epochs_num=10,
       real_ppo_epoch_length=200,
       real_ppo_num_agents=16,
       real_ppo_learning_rate=2e-4,

From 126fb6bb691048294c8a3e14bb313ff9f1b0d74f Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 11 Sep 2018 15:35:54 -0700
Subject: [PATCH 0815/2720] Add task id to multiproblem features.

PiperOrigin-RevId: 212535366
---
 tensor2tensor/data_generators/multi_problem.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 44485dc5c..99b1b5da1 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -68,6 +68,9 @@ def add_task_id(self, task, example, encoder):
       concat_list = [[task.task_id], example["targets"]]
 
     example["targets"] = tf.concat(concat_list, 0)
+    min_task_id = min([t.task_id for t in self.task_list])
+    example["task_id"] = tf.constant([task.task_id - min_task_id],
+                                     dtype=tf.int64)
     return example
 
   def filepattern(self, data_dir, mode, shard=None):
@@ -153,10 +156,11 @@ def dataset(self,
       task_dataset = task_dataset.map(lambda x: self.add_task_id(task, x, enc))
 
       if not is_training:
+        zeros = tf.zeros([self._ADDED_EVAL_COUNT, 1], dtype=tf.int64)
         pad_data = tf.data.Dataset.from_tensor_slices({
-            "targets": tf.zeros([self._ADDED_EVAL_COUNT, 1], dtype=tf.int64),
-            "batch_prediction_key": tf.zeros(
-                [self._ADDED_EVAL_COUNT, 1], dtype=tf.int64),
+            "targets": zeros,
+            "batch_prediction_key": zeros,
+            "task_id": zeros,
         })
         task_dataset = task_dataset.concatenate(pad_data)
 

From 01af43d2b3e806035de461048a1d9fbe20f77bee Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 11 Sep 2018 17:32:09 -0700
Subject: [PATCH 0816/2720] internal merge of PR #1057

PiperOrigin-RevId: 212554982
---
 tensor2tensor/rl/envs/simulated_batch_env.py  |   2 +-
 .../rl/model_rl_experiment_player.py          | 116 +++++++++---------
 2 files changed, 60 insertions(+), 58 deletions(-)

diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 0bc946542..c2838ac49 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -131,7 +131,7 @@ def __init__(self, environment_spec, length):
                       environment_spec.video_num_target_frames,
                       environment_spec=environment_spec)
 
-    #TODO(piotrmilos): check if this should not be tf.estimator.ModeKeys.Predict
+    # TODO(piotrmilos): check if this shouldn't be tf.estimator.ModeKeys.Predict
     initial_frames_dataset = initial_frames_problem.dataset(
         tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir, shuffle_files=False,
         hparams=hparams).take(1)
diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index ad7cad49e..3f53f6b07 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -12,32 +12,36 @@
 # 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.
-
-#--output_dir=/Users/piotr.milos/Downloads/18 --alsologtostderr --loop_hparams_set=rl_modelrl_base_quick
-#--output_dir=/Users/piotr.milos/t2t/rl_v1 --alsologtostderr --loop_hparams_set=rl_modelrl_tiny
+"""Play with a world model."""
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from PIL import Image
-from PIL import ImageFont
-from PIL import ImageDraw
 import copy
 import os
+
+from gym.core import Env
 from gym.spaces import Box
+from gym.spaces import Discrete
+from gym.utils import play
+
 import numpy as np
 
-from gym.spaces import Discrete
+from PIL import Image
+from PIL import ImageDraw
+from PIL import ImageFont
+
 from tensor2tensor.data_generators import gym_problems_specs
 from tensor2tensor.rl.envs.batch_env_factory import batch_env_factory
 from tensor2tensor.rl.envs.utils import get_policy
+from tensor2tensor.rl.trainer_model_based import FLAGS
+from tensor2tensor.rl.trainer_model_based import setup_directories
+from tensor2tensor.rl.trainer_model_based import temporary_flags
+
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
-from tensor2tensor.rl.trainer_model_based import FLAGS, setup_directories, temporary_flags
-from gym.utils import play
 import tensorflow as tf
-from gym.core import Env
 
 
 _font = None
@@ -47,16 +51,12 @@
 def _get_font():
   global _font
   if _font is None:
-    #weirdness due to various working dirs
-    FONT_PATHS = ["tensor-2-tensor-with-mrunner/tensor-2-tensor-with-mrunner/deepsense_experiments/Xerox Serif Narrow.ttf",
-                  "tensor-2-tensor-with-mrunner/deepsense_experiments/Xerox Serif Narrow.ttf",
-                  "deepsense_experiments/Xerox Serif Narrow.ttf"]
-
-    for path in FONT_PATHS:
+    font_paths = []
+    for path in font_paths:
       try:
         _font = ImageFont.truetype(path, FONT_SIZE)
         return _font
-      except:
+      except:  # pylint: disable=bare-except
         pass
 
 
@@ -68,29 +68,30 @@ def _assert_image(img):
 
 def write_on_image(img, text="", position=(0, 0), color=(255, 255, 255)):
   img = _assert_image(img)
-  if text=="":
+  if not text:
     return img
   draw = ImageDraw.Draw(img)
   font = _get_font()
   draw.text(position, text, color, font=font)
-
   return img
 
-def concatenate_images(*imgs, axis=1):
+
+def concatenate_images(imgs, axis=1):
   imgs = [_assert_image(img) for img in imgs]
   imgs_np = [np.array(img) for img in imgs]
   concatenated_im_np = np.concatenate(imgs_np, axis=axis)
-
   return _assert_image(concatenated_im_np)
 
 
 class DebugBatchEnv(Env):
-
+  """Debugging Environment."""
   INFO_PANE_WIDTH = 250
 
   def __init__(self, hparams, sess=None):
     self.action_space = Discrete(6)
-    self.observation_space = Box(low=0, high=255, shape=(210, 160+DebugBatchEnv.INFO_PANE_WIDTH, 3), dtype=np.uint8)
+    self.observation_space = Box(
+        low=0, high=255, shape=(210, 160+DebugBatchEnv.INFO_PANE_WIDTH, 3),
+        dtype=np.uint8)
     self._tmp = 1
     self.res = None
     self.sess = sess if sess is not None else tf.Session()
@@ -123,22 +124,22 @@ def initialization_lambda():
     self.policy_probs = actor_critic.policy.probs[0, 0, :]
     self.value = actor_critic.value[0, :]
 
-  def render(self, mode='human'):
-    raise NotImplemented()
+  def render(self, mode="human"):
+    raise NotImplementedError()
 
   def _fake_reset(self):
     self._tmp = 0
-    _observ = np.ones(shape=(210, 160, 3), dtype=np.uint8) * 10 * self._tmp
-    _observ[0, 0, 0] = 0
-    _observ[0, 0, 1] = 255
-    self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)
+    observ = np.ones(shape=(210, 160, 3), dtype=np.uint8) * 10 * self._tmp
+    observ[0, 0, 0] = 0
+    observ[0, 0, 1] = 255
+    self.res = (observ, 0, False, [0.1, 0.5, 0.5], 1.1)
 
   def _reset_env(self):
-    _observ = self.sess.run(self.reset_op)[0, ...]
-    _observ[0, 0, 0] = 0
-    _observ[0, 0, 1] = 255
-    #TODO:(put correct numbers)
-    self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)
+    observ = self.sess.run(self.reset_op)[0, ...]
+    observ[0, 0, 0] = 0
+    observ[0, 0, 1] = 255
+    # TODO(pmilos): put correct numbers
+    self.res = (observ, 0, False, [0.1, 0.5, 0.5], 1.1)
 
   def reset(self):
     self._reset_env()
@@ -151,7 +152,7 @@ def _step_fake(self, action):
     observ[0, 0, 1] = 255
 
     self._tmp += 1
-    if self._tmp>20:
+    if self._tmp > 20:
       self._tmp = 0
 
     rew = 1
@@ -163,14 +164,16 @@ def _step_fake(self, action):
 
   def _step_env(self, action):
     observ, rew, done, probs, vf = self.sess.\
-      run([self.observation, self.reward, self.done, self.policy_probs, self.value],
+      run([self.observation, self.reward, self.done, self.policy_probs,
+           self.value],
           feed_dict={self.action: [action]})
 
     return observ[0, ...], rew[0, ...], done[0, ...], probs, vf
 
   def _augment_observation(self):
-    _observ, rew, done, probs, vf = self.res
-    info_pane = np.zeros(shape=(210, DebugBatchEnv.INFO_PANE_WIDTH, 3), dtype=np.uint8)
+    observ, rew, _, probs, vf = self.res
+    info_pane = np.zeros(shape=(210, DebugBatchEnv.INFO_PANE_WIDTH, 3),
+                         dtype=np.uint8)
     probs_str = ""
     for p in probs:
       probs_str += "%.2f" % p + ", "
@@ -178,37 +181,37 @@ def _augment_observation(self):
     probs_str = probs_str[:-2]
 
     action = np.argmax(probs)
-    info_str = " Policy:{}\n Action:{}\n Value function:{}\n Reward:{}".format(probs_str, action,
-                                                                           vf, rew)
+    info_str = " Policy:{}\n Action:{}\n Value function:{}\n Reward:{}".format(
+        probs_str, action, vf, rew)
     print("Info str:{}".format(info_str))
     # info_pane = write_on_image(info_pane, info_str)
 
-    augmented_observ = concatenate_images(_observ, info_pane)
+    augmented_observ = concatenate_images(observ, info_pane)
     augmented_observ = np.array(augmented_observ)
     return augmented_observ
 
   def step(self, action):
-    #Special codes
-    if action==100:
-      #Skip action
+    # Special codes
+    if action == 100:
+      # skip action
       _, rew, done, _, _ = self.res
       observ = self._augment_observation()
       return observ, rew, done, {}
 
     if action == 101:
-      #reset
+      # reset
       self.reset()
       _, rew, done, _, _ = self.res
       observ = self._augment_observation()
       return observ, rew, done, {}
 
     if action == 102:
-      #play
-      raise NotImplemented()
+      # play
+      raise NotImplementedError()
 
-    #standard codes
-    _observ, rew, done, probs, vf = self._step_env(action)
-    self.res = (_observ, rew, done, probs, vf)
+    # standard codes
+    observ, rew, done, probs, vf = self._step_env(action)
+    self.res = (observ, rew, done, probs, vf)
 
     observ = self._augment_observation()
     return observ, rew, done, {"probs": probs, "vf": vf}
@@ -267,24 +270,23 @@ def main(_):
     env.initialize()
 
     env_model_loader = tf.train.Saver(
-      tf.global_variables("next_frame*"))
+        tf.global_variables("next_frame*"))
     trainer_lib.restore_checkpoint(world_model_dir, env_model_loader, sess,
-      must_restore=True)
+                                   must_restore=True)
 
     model_saver = tf.train.Saver(
-      tf.global_variables(".*network_parameters.*"))
+        tf.global_variables(".*network_parameters.*"))
     trainer_lib.restore_checkpoint(ppo_model_dir, model_saver, sess)
 
     key_mapping = gym_problem.env.env.get_keys_to_action()
-    #map special codes
+    # map special codes
     key_mapping[()] = 100
-    key_mapping[(ord('r'),)] = 101
-    key_mapping[(ord('p'),)] = 102
+    key_mapping[(ord("r"),)] = 101
+    key_mapping[(ord("p"),)] = 102
 
     play.play(env, zoom=2, fps=10, keys_to_action=key_mapping)
 
 
-
 if __name__ == "__main__":
   tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()

From 39f23702969f12d7d7544410e5d3959be48e2870 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 11 Sep 2018 18:59:53 -0700
Subject: [PATCH 0817/2720] introducing effective_num_agents

---
 tensor2tensor/layers/common_video.py    |  3 ++-
 tensor2tensor/rl/collect.py             | 18 ++++++++++++++++++
 tensor2tensor/rl/ppo.py                 |  4 ++++
 tensor2tensor/rl/trainer_model_based.py | 14 +++++++++-----
 4 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 469156f23..56cfb5d97 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -535,9 +535,10 @@ def __init_ffmpeg(self, image_shape):
         "-qscale", "0",
         "-"
     ]
-    self.proc = Popen(self.cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
+    # self.proc = Popen(self.cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
 
   def write(self, frame):
+    return
     if self.proc is None:
       self.__init_ffmpeg(frame.shape)
     self.proc.stdin.write(frame.tostring())
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 7e864e5df..b712953c6 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -261,6 +261,24 @@ def stop_condition(i, _, resets):
   printing = tf.Print(0, [mean_score, scores_sum, scores_num], "mean_score: ")
   with tf.control_dependencies([index, printing]):
     memory = [mem.read_value() for mem in memory]
+    new_memory = []
+    if hasattr(hparams, "effective_num_agents"):
+      effective_num_agents = hparams.effective_num_agents
+      new_epoch_length = int(hparams.epoch_length / effective_num_agents)
+      for mem, info in zip(memory, rollout_metadata):
+        shape, _, name = info
+        new_shape = [effective_num_agents, new_epoch_length] + shape[1:]
+        perm = list(range(len(shape)+1))
+        perm[0] = 1
+        perm[1] = 0
+        mem = tf.transpose(mem, perm=perm)
+        mem = tf.reshape(mem, shape=new_shape)
+        mem = tf.transpose(mem, perm=perm,
+                           name="collect_memory_%d_%s"
+                                % (new_epoch_length, name))
+        new_memory.append(mem)
+      memory = new_memory
+
     mean_score_summary = tf.cond(
         tf.greater(scores_num, 0),
         lambda: tf.summary.scalar("mean_score_this_iter", mean_score),
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index 750c02537..ca8cc1a58 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -100,6 +100,10 @@ def define_ppo_epoch(memory, hparams):
   number_of_batches = (hparams.epoch_length * hparams.optimization_epochs
                        / hparams.optimization_batch_size)
 
+  if hasattr(hparams, "effective_num_agents"):
+    number_of_batches *= hparams.num_agents
+    number_of_batches /= hparams.effective_num_agents
+    
   dataset = tf.data.Dataset.from_tensor_slices(
       (observation, action, discounted_reward, advantage_normalized, old_pdf))
   dataset = dataset.shuffle(buffer_size=hparams.epoch_length,
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index dd5c7ecc0..692d37f46 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -208,8 +208,10 @@ def train_agent_real_env(
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_params_names = ["epochs_num", "epoch_length",
                       "learning_rate", "num_agents",
-                      "optimization_epochs"]
+                      "optimization_epochs", "effective_num_agents"]
 
+  # This should be overriden
+  ppo_hparams.add_hparam("effective_num_agents", None)
   for param_name in ppo_params_names:
     ppo_param_name = "real_ppo_"+ param_name
     if ppo_param_name in hparams:
@@ -599,10 +601,11 @@ def rl_modelrl_base():
       ppo_continue_training=True,
 
       real_ppo_epochs_num=10,
-      real_ppo_epoch_length=200,
-      real_ppo_num_agents=16,
+      real_ppo_epoch_length=16*200,
+      real_ppo_num_agents=1,
       real_ppo_learning_rate=2e-4,
       real_ppo_continue_training=True,
+      real_ppo_effective_num_agents=16,
 
       game="wrapped_full_pong",
       # Whether to evaluate the world model in each iteration of the loop to get
@@ -713,8 +716,9 @@ def rl_modelrl_tiny():
           ppo_epoch_length=5,
           ppo_num_agents=2,
           real_ppo_epochs_num=1,
-          real_ppo_epoch_length=5,
-          real_ppo_num_agents=2,
+          real_ppo_epoch_length=10,
+          real_ppo_num_agents=1,
+          real_ppo_effective_num_agents=2,
           generative_model_params="next_frame_tiny",
       ).values())
 

From 910cdfa3fca5c3f78fa7fd58a188a73f40b43f21 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 11 Sep 2018 19:29:06 -0700
Subject: [PATCH 0818/2720] PyFuncWrapper

---
 tensor2tensor/rl/envs/tf_atari_wrappers.py | 29 ++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 002fbff13..d2c5ff813 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -312,3 +312,32 @@ def _transform_history_observations(self, frames):
         new_frames, (batch_size, history_size) + self.observ_shape
     )
     return tf.cast(new_frames, self.observ_dtype)
+
+
+class PyFuncWrapper(WrapperBase):
+  """Calls arbitrary python function on passing data"""
+
+  def __init__(self, batch_env, process_fun):
+    super(PyFuncWrapper, self).__init__(batch_env)
+    self.process_fun = process_fun
+
+  def simulate(self, action):
+    reward, done = self._batch_env.simulate(action)
+    with tf.control_dependencies([reward, done]):
+      inputs = [self._batch_env.observ, reward, done, action]
+      ret = tf.py_func(self.process_fun, inputs, tf.double)
+      with tf.control_dependencies([ret]):
+        return tf.identity(reward), tf.identity(done)
+
+  @property
+  def observ(self):
+    """Access the variable holding the current observation."""
+    return self._batch_env.observ
+
+  def __len__(self):
+    """Number of combined environments."""
+    return len(self._batch_env)
+
+  def _reset_non_empty(self, indices):
+    # pylint: disable=protected-access
+    return self._batch_env._reset_non_empty(indices)
\ No newline at end of file

From 1db22234bacfc67eeb1c7c4ad5fe161b36ac25a5 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 11 Sep 2018 20:16:40 -0700
Subject: [PATCH 0819/2720] ppo data dumper

---
 tensor2tensor/rl/trainer_model_based.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 692d37f46..fcdc8f431 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -39,8 +39,10 @@
 from tensor2tensor.data_generators import gym_problems_specs
 from tensor2tensor.layers import discretization
 from tensor2tensor.rl import rl_trainer_lib
+from tensor2tensor.rl.envs.tf_atari_wrappers import PyFuncWrapper
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
+import numpy as np
 
 import tensorflow as tf
 
@@ -198,11 +200,22 @@ def train_agent(problem_name, agent_model_dir,
     rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir,
                          name_scope="ppo_sim")
 
+ppo_data_dumper_counter = 0
+dumper_path = None
+
+def ppo_data_dumper(observ, reward, done, action):
+  global ppo_data_dumper_counter, dumper_path
+  np.savez_compressed("{}/frame_{}".format(dumper_path,
+                                           ppo_data_dumper_counter),
+                      observ=observ, reward=reward, done=done, action=action)
+  ppo_data_dumper_counter += 1
+  return 0.0
 
 def train_agent_real_env(
     problem_name, agent_model_dir, event_dir, world_model_dir, epoch_data_dir,
     hparams, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the real environment."""
+  global dumper_path, ppo_data_dumper_counter
 
   gym_problem = registry.problem(problem_name)
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
@@ -223,6 +236,14 @@ def train_agent_real_env(
 
   environment_spec = copy.copy(gym_problem.environment_spec)
 
+  #TODO(piotrmilos):This should be refactored
+  ppo_data_dumper_counter = 0
+  dumper_path = os.path.join(epoch_data_dir, "dumper")
+  tf.gfile.MakeDirs(dumper_path)
+  dumper_spec = [PyFuncWrapper, {"process_fun": ppo_data_dumper}]
+  environment_spec.wrappers.insert(1, dumper_spec)
+
+
   ppo_hparams.add_hparam("environment_spec", environment_spec)
 
   with temporary_flags({

From 340c489e1fef10d353bcab3fb55244831f0cb467 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 11 Sep 2018 22:24:40 -0700
Subject: [PATCH 0820/2720] reader of ppo dumped data

---
 tensor2tensor/data_generators/gym_problems.py | 173 ++++++++++--------
 tensor2tensor/rl/trainer_model_based.py       |   6 +-
 2 files changed, 100 insertions(+), 79 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index e22285e79..773ccd724 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -97,44 +97,72 @@ def __init__(self, *args, **kwargs):
     self._internal_memory_force_beginning_resets = False
     self._session = None
     self.statistics = BasicStatistics()
+    self._use_dumper_data = False
+    self._dumper_data_index = 0
 
-  def _setup(self):
+  def _setup(self, data_dir):
     # TODO(piotrmilos):this should be consistent with
     # ppo_params in model_rl_experiment
-    collect_hparams = rl.ppo_pong_base()
-    collect_hparams.add_hparam("environment_spec", self.environment_spec)
-    collect_hparams.add_hparam("force_beginning_resets",
-                               self._internal_memory_force_beginning_resets)
-    collect_hparams.epoch_length = self._internal_memory_size
-    collect_hparams.num_agents = 1
-
-    if not FLAGS.agent_policy_path:
-      collect_hparams.policy_network = rl.random_policy_fun
-
-    policy_to_actions_lambda = None
-    if self.settable_eval_phase:
-      policy_to_actions_lambda = lambda policy: policy.mode()
-
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      self.collect_memory, self.collect_trigger_op, collect_init = (
-          collect.define_collect(
-              collect_hparams,
-              scope="gym_problems",
-              eval_phase=False,
-              collect_level=0,
-              policy_to_actions_lambda=policy_to_actions_lambda))
-
-    self._session = tf.Session()
-    collect_init(self._session)
-    self._session.run(tf.global_variables_initializer())
-    self.restore_networks(self._session)
+    dumper_path = os.path.join(data_dir, "dumper")
+    if os.path.isdir(dumper_path):
+      self._use_dumper_data = True
+      self._dumper_data_index = 0
+      self._dumper_path = dumper_path
+    else:
+
+      collect_hparams = rl.ppo_pong_base()
+      collect_hparams.add_hparam("environment_spec", self.environment_spec)
+      collect_hparams.add_hparam("force_beginning_resets",
+                                 self._internal_memory_force_beginning_resets)
+      collect_hparams.epoch_length = self._internal_memory_size
+      collect_hparams.num_agents = 1
+
+      if not FLAGS.agent_policy_path:
+        collect_hparams.policy_network = rl.random_policy_fun
+
+      policy_to_actions_lambda = None
+      if self.settable_eval_phase:
+        policy_to_actions_lambda = lambda policy: policy.mode()
+
+      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+        self.collect_memory, self.collect_trigger_op, collect_init = (
+            collect.define_collect(
+                collect_hparams,
+                scope="gym_problems",
+                eval_phase=False,
+                collect_level=0,
+                policy_to_actions_lambda=policy_to_actions_lambda))
+
+      self._session = tf.Session()
+      collect_init(self._session)
+      self._session.run(tf.global_variables_initializer())
+      self.restore_networks(self._session)
+      self.memory_index = 0
+      self.memory = None
+
 
   @property
   def random_skip(self):
     return False
 
+  def _get_data(self):
+    if self._use_dumper_data:
+      file_path = os.path.join(self._dumper_path, "frame_{}.npz".format(self._dumper_data_index))
+      print(file_path)
+      data = np.load(file_path)
+      self._dumper_data_index += 1
+      return data["observ"][0, ...], data["reward"][0], data["done"][0], data["action"][0]
+    else:
+      if self.memory is None or self.memory_index >= self._internal_memory_size:
+        self.memory = self._session.run(self.collect_memory)
+        self.memory_index = 0
+      data = [self.memory[i][self.memory_index][0] for i in range(4)]
+      self.memory_index += 1
+
+      return data
+
   def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
-    self._setup()
+    self._setup(data_dir)
 
     # We only want to save frames for eval and simulated experience, not the
     # frames used for world model training.
@@ -148,51 +176,44 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
       # Disable frame saving
       self.debug_dump_frames_path = ""
 
-    with self._session as sess:
-      frame_counter = 0
-      memory_index = 0
-      memory = None
-      pieces_generated = 0
-      prev_reward = 0
-      prev_done = False
-
-      # TODO(piotrmilos): self.settable_eval_phase possibly violates sematics
-      # of VideoProblem
-      while pieces_generated < self.num_steps or self.settable_eval_phase:
-        if memory is None or memory_index >= self._internal_memory_size:
-          memory = sess.run(self.collect_memory)
-          memory_index = 0
-        data = [memory[i][memory_index][0] for i in range(4)]
-        memory_index += 1
-        observation, reward, done, action = data
-
-        debug_image = self.collect_statistics_and_generate_debug_image(
-            pieces_generated, *data)
-        ret_dict = {
-            "frame": observation,
-            "frame_number": [int(frame_counter)],
-            "image/format": ["png"],
-            "image/height": [self.frame_height],
-            "image/width": [self.frame_width],
-            "action": [int(action)],
-            "done": [int(prev_done)],
-            "reward": [int(prev_reward - self.min_reward)]
-        }
-
-        if debug_image is not None:
-          ret_dict["image/debug"] = debug_image
-
-        yield ret_dict
-
-        if done and self.settable_eval_phase:
-          return
-
-        prev_done, prev_reward = done, reward
-
-        pieces_generated += 1
-        frame_counter += 1
-        if done:
-          frame_counter = 0
+    frame_counter = 0
+    pieces_generated = 0
+    prev_reward = 0
+    prev_done = False
+
+    # TODO(piotrmilos): self.settable_eval_phase possibly violates sematics
+    # of VideoProblem
+    while pieces_generated < self.num_steps or self.settable_eval_phase:
+      data = self._get_data()
+      observation, reward, done, action = data
+
+      debug_image = self.collect_statistics_and_generate_debug_image(
+          pieces_generated, *data)
+      ret_dict = {
+          "frame": observation,
+          "frame_number": [int(frame_counter)],
+          "image/format": ["png"],
+          "image/height": [self.frame_height],
+          "image/width": [self.frame_width],
+          "action": [int(action)],
+          "done": [int(prev_done)],
+          "reward": [int(prev_reward - self.min_reward)]
+      }
+
+      if debug_image is not None:
+        ret_dict["image/debug"] = debug_image
+
+      yield ret_dict
+
+      if done and self.settable_eval_phase:
+        return
+
+      prev_done, prev_reward = done, reward
+
+      pieces_generated += 1
+      frame_counter += 1
+      if done:
+        frame_counter = 0
 
   def restore_networks(self, sess):
     if FLAGS.agent_policy_path:
@@ -497,8 +518,8 @@ def __init__(self, *args, **kwargs):
     self.statistics = RewardPerSequenceStatistics()
     self.statistics.real_env = real_env
 
-  def _setup(self):
-    super(GymSimulatedDiscreteProblem, self)._setup()
+  def _setup(self, data_dir):
+    super(GymSimulatedDiscreteProblem, self)._setup(data_dir)
 
     environment_spec = self.environment_spec
     hparams = HParams(
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index fcdc8f431..1e85c6e90 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -728,7 +728,7 @@ def rl_modelrl_tiny():
   """Tiny set for testing."""
   return rl_modelrl_base_sampling().override_from_dict(
       tf.contrib.training.HParams(
-          epochs=1,
+          epochs=2,
           num_real_env_frames=128,
           simulated_env_generator_num_steps=64,
           model_train_steps=2,
@@ -736,8 +736,8 @@ def rl_modelrl_tiny():
           ppo_time_limit=5,
           ppo_epoch_length=5,
           ppo_num_agents=2,
-          real_ppo_epochs_num=1,
-          real_ppo_epoch_length=10,
+          real_ppo_epochs_num=2,
+          real_ppo_epoch_length=36,
           real_ppo_num_agents=1,
           real_ppo_effective_num_agents=2,
           generative_model_params="next_frame_tiny",

From 7e3b06a81fc35243b2c9e7eb471132fdbb0fac0c Mon Sep 17 00:00:00 2001
From: Marcin Michalski <michalski@google.com>
Date: Wed, 12 Sep 2018 02:53:22 -0700
Subject: [PATCH 0821/2720] adding the export_dir flag that allows storing the
 exported model in separate directory.

PiperOrigin-RevId: 212607096
---
 tensor2tensor/serving/export.py | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 52b1dbe4c..6adba8e82 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -32,6 +32,11 @@
 tf.flags.DEFINE_bool("export_as_tfhub", False,
                      "If True, the model will be exported as tfHub module.")
 
+tf.flags.DEFINE_string(
+    "export_dir", None, "Directory, where export model should be stored."
+    "If None, the model will be stored in subdirectory "
+    "where checkpoints are: --output_dir")
+
 
 def create_estimator(run_config, hparams):
   return trainer_lib.create_estimator(
@@ -73,7 +78,7 @@ def export_module_spec_with_checkpoint(module_spec,
       m.export(export_path, session)
 
 
-def export_as_tfhub_module(hparams, problem, ckpt_dir):
+def export_as_tfhub_module(hparams, problem, ckpt_dir, export_dir):
   """Exports the last checkpoint from the directory as tfhub module.
 
   It creates the Module spec and signature (based on T2T problem information),
@@ -83,8 +88,8 @@ def export_as_tfhub_module(hparams, problem, ckpt_dir):
   Args:
     hparams: T2T parameters, model graph will be based on them.
     problem: the name of the problem
-    ckpt_dir: directory with the checkpoints. The final model will be exported
-      there too.
+    ckpt_dir: directory with the checkpoints.
+    export_dir: Directory to write the exported model to.
   """
 
   def hub_module_fn():
@@ -101,7 +106,6 @@ def hub_module_fn():
         inputs=features, outputs=spec.export_outputs["serving_default"].outputs)
 
   module_spec = hub.create_module_spec(hub_module_fn)
-  export_dir = os.path.join(ckpt_dir, "export_tfhub")
   # Loads the weights from the checkpoint using the model above
   # and saves it in the export_path.
   export_module_spec_with_checkpoint(
@@ -122,8 +126,10 @@ def main(_):
   hparams.no_data_parallelism = True  # To clear the devices
   problem = hparams.problem
 
+  export_dir = FLAGS.export_dir or os.path.join(ckpt_dir, "export")
+
   if FLAGS.export_as_tfhub:
-    export_as_tfhub_module(hparams, problem, ckpt_dir)
+    export_as_tfhub_module(hparams, problem, ckpt_dir, export_dir)
     return
 
   run_config = t2t_trainer.create_run_config(hparams)
@@ -133,7 +139,6 @@ def main(_):
   exporter = tf.estimator.FinalExporter(
       "exporter", lambda: problem.serving_input_fn(hparams), as_text=True)
 
-  export_dir = os.path.join(ckpt_dir, "export")
   exporter.export(
       estimator,
       export_dir,

From 60f20ac02767ce5d13c54eee6c48ef0bb75a9d0d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 12 Sep 2018 10:09:20 -0700
Subject: [PATCH 0822/2720] Add hparam for controlling eos processing on
 decode.

PiperOrigin-RevId: 212659081
---
 tensor2tensor/utils/decoding.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 4d9024b18..ffc07334a 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -66,7 +66,8 @@ def decode_hparams(overrides=""):
       force_decode_length=False,
       display_decoded_images=False,
       # Used for video decoding.
-      frames_per_second=10)
+      frames_per_second=10,
+      skip_eos_postprocess=False)
   hp.parse(overrides)
   return hp
 
@@ -485,7 +486,8 @@ def decode_interactively(estimator, hparams, decode_hp, checkpoint_path=None):
   is_image = "image" in hparams.problem.name
   is_text2class = isinstance(hparams.problem,
                              text_problems.Text2ClassProblem)
-  skip_eos_postprocess = is_image or is_text2class
+  skip_eos_postprocess = (
+      is_image or is_text2class or decode_hp.skip_eos_postprocess)
 
   def input_fn():
     gen_fn = make_input_fn_from_generator(

From 0a8831fce9b5cc4abc2766a9072b2c4c49339cfd Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 12 Sep 2018 10:29:19 -0700
Subject: [PATCH 0823/2720] Move SpeechRecognitionModality to modalities.py;
 create common_audio.py.

PiperOrigin-RevId: 212662556
---
 .../data_generators/speech_recognition.py     | 208 +-----------------
 tensor2tensor/layers/common_audio.py          | 137 ++++++++++++
 tensor2tensor/layers/modalities.py            |  85 +++++++
 3 files changed, 225 insertions(+), 205 deletions(-)
 create mode 100644 tensor2tensor/layers/common_audio.py

diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index aaee536ea..1e422c8b2 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -18,138 +18,19 @@
 it as appropriate (e.g. using apt-get or yum).
 """
 
-import functools
 import numpy as np
-import scipy.signal
 
 from tensor2tensor.data_generators import audio_encoder
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
-from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_audio
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import metrics
-from tensor2tensor.utils import modality
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
 
-#
-# ASR Feature pipeline in TF.
-#
-def add_delta_deltas(filterbanks, name=None):
-  """Compute time first and second-order derivative channels.
-
-  Args:
-    filterbanks: float32 tensor with shape [batch_size, len, num_bins, 1]
-    name: scope name
-
-  Returns:
-    float32 tensor with shape [batch_size, len, num_bins, 3]
-  """
-  delta_filter = np.array([2, 1, 0, -1, -2])
-  delta_delta_filter = scipy.signal.convolve(delta_filter, delta_filter, "full")
-
-  delta_filter_stack = np.array(
-      [[0] * 4 + [1] + [0] * 4, [0] * 2 + list(delta_filter) + [0] * 2,
-       list(delta_delta_filter)],
-      dtype=np.float32).T[:, None, None, :]
-
-  delta_filter_stack /= np.sqrt(
-      np.sum(delta_filter_stack**2, axis=0, keepdims=True))
-
-  filterbanks = tf.nn.conv2d(
-      filterbanks, delta_filter_stack, [1, 1, 1, 1], "SAME", data_format="NHWC",
-      name=name)
-  return filterbanks
-
-
-def compute_mel_filterbank_features(
-    waveforms,
-    sample_rate=16000, dither=1.0 / np.iinfo(np.int16).max, preemphasis=0.97,
-    frame_length=25, frame_step=10, fft_length=None,
-    window_fn=functools.partial(tf.contrib.signal.hann_window, periodic=True),
-    lower_edge_hertz=80.0, upper_edge_hertz=7600.0, num_mel_bins=80,
-    log_noise_floor=1e-3, apply_mask=True):
-  """Implement mel-filterbank extraction using tf ops.
-
-  Args:
-    waveforms: float32 tensor with shape [batch_size, max_len]
-    sample_rate: sampling rate of the waveform
-    dither: stddev of Gaussian noise added to waveform to prevent quantization
-      artefacts
-    preemphasis: waveform high-pass filtering constant
-    frame_length: frame length in ms
-    frame_step: frame_Step in ms
-    fft_length: number of fft bins
-    window_fn: windowing function
-    lower_edge_hertz: lowest frequency of the filterbank
-    upper_edge_hertz: highest frequency of the filterbank
-    num_mel_bins: filterbank size
-    log_noise_floor: clip small values to prevent numeric overflow in log
-    apply_mask: When working on a batch of samples, set padding frames to zero
-  Returns:
-    filterbanks: a float32 tensor with shape [batch_size, len, num_bins, 1]
-  """
-  # `stfts` is a complex64 Tensor representing the short-time Fourier
-  # Transform of each signal in `signals`. Its shape is
-  # [batch_size, ?, fft_unique_bins]
-  # where fft_unique_bins = fft_length // 2 + 1
-
-  # Find the wave length: the largest index for which the value is !=0
-  # note that waveforms samples that are exactly 0.0 are quite common, so
-  # simply doing sum(waveforms != 0, axis=-1) will not work correctly.
-  wav_lens = tf.reduce_max(
-      tf.expand_dims(tf.range(tf.shape(waveforms)[1]), 0) *
-      tf.to_int32(tf.not_equal(waveforms, 0.0)),
-      axis=-1) + 1
-  if dither > 0:
-    waveforms += tf.random_normal(tf.shape(waveforms), stddev=dither)
-  if preemphasis > 0:
-    waveforms = waveforms[:, 1:] - preemphasis * waveforms[:, :-1]
-    wav_lens -= 1
-  frame_length = int(frame_length * sample_rate / 1e3)
-  frame_step = int(frame_step * sample_rate / 1e3)
-  if fft_length is None:
-    fft_length = int(2**(np.ceil(np.log2(frame_length))))
-
-  stfts = tf.contrib.signal.stft(
-      waveforms,
-      frame_length=frame_length,
-      frame_step=frame_step,
-      fft_length=fft_length,
-      window_fn=window_fn,
-      pad_end=True)
-
-  stft_lens = (wav_lens + (frame_step - 1)) // frame_step
-  masks = tf.to_float(tf.less_equal(
-      tf.expand_dims(tf.range(tf.shape(stfts)[1]), 0),
-      tf.expand_dims(stft_lens, 1)))
-
-  # An energy spectrogram is the magnitude of the complex-valued STFT.
-  # A float32 Tensor of shape [batch_size, ?, 257].
-  magnitude_spectrograms = tf.abs(stfts)
-
-  # Warp the linear-scale, magnitude spectrograms into the mel-scale.
-  num_spectrogram_bins = magnitude_spectrograms.shape[-1].value
-  linear_to_mel_weight_matrix = (
-      tf.contrib.signal.linear_to_mel_weight_matrix(
-          num_mel_bins, num_spectrogram_bins, sample_rate, lower_edge_hertz,
-          upper_edge_hertz))
-  mel_spectrograms = tf.tensordot(
-      magnitude_spectrograms, linear_to_mel_weight_matrix, 1)
-  # Note: Shape inference for tensordot does not currently handle this case.
-  mel_spectrograms.set_shape(magnitude_spectrograms.shape[:-1].concatenate(
-      linear_to_mel_weight_matrix.shape[-1:]))
-
-  log_mel_sgram = tf.log(tf.maximum(log_noise_floor, mel_spectrograms))
-
-  if apply_mask:
-    log_mel_sgram *= tf.expand_dims(tf.to_float(masks), -1)
-
-  return tf.expand_dims(log_mel_sgram, -1, name="mel_sgrams")
-
-
 class ByteTextEncoderWithEos(text_encoder.ByteTextEncoder):
   """Encodes each byte to an id and appends the EOS token."""
 
@@ -220,7 +101,7 @@ def preprocess_example(self, example, mode, hparams):
           tf.expand_dims(example["waveforms"], -1), -1)
     else:
       waveforms = tf.expand_dims(example["waveforms"], 0)
-      mel_fbanks = compute_mel_filterbank_features(
+      mel_fbanks = common_audio.compute_mel_filterbank_features(
           waveforms,
           sample_rate=p.audio_sample_rate,
           dither=p.audio_dither,
@@ -232,7 +113,7 @@ def preprocess_example(self, example, mode, hparams):
           num_mel_bins=p.audio_num_mel_bins,
           apply_mask=False)
       if p.audio_add_delta_deltas:
-        mel_fbanks = add_delta_deltas(mel_fbanks)
+        mel_fbanks = common_audio.add_delta_deltas(mel_fbanks)
       fbank_size = common_layers.shape_list(mel_fbanks)
       assert fbank_size[0] == 1
 
@@ -257,86 +138,3 @@ def preprocess_example(self, example, mode, hparams):
   def eval_metrics(self):
     defaults = super(SpeechRecognitionProblem, self).eval_metrics()
     return defaults + [metrics.Metrics.EDIT_DISTANCE]
-
-
-@registry.register_audio_modality
-class SpeechRecognitionModality(modality.Modality):
-  """Common ASR filterbank processing."""
-
-  def bottom(self, x):
-    """Use batchnorm instead of CMVN and shorten the stft with strided convs.
-
-    Args:
-      x: float32 tensor with shape [batch_size, len, 1, freqs * channels]
-
-    Returns:
-      float32 tensor with shape [batch_size, shorter_len, 1, hidden_size]
-    """
-    inputs = x
-    p = self._model_hparams
-
-    num_mel_bins = p.audio_num_mel_bins
-    num_channels = 3 if p.audio_add_delta_deltas else 1
-
-    with tf.variable_scope(self.name):
-      if p.audio_preproc_in_bottom:
-        # Compute filterbanks
-        with tf.variable_scope("fbanks"):
-          waveforms = tf.squeeze(inputs, [2, 3])
-          mel_fbanks = compute_mel_filterbank_features(
-              waveforms,
-              sample_rate=p.audio_sample_rate,
-              dither=p.audio_dither,
-              preemphasis=p.audio_preemphasis,
-              frame_length=p.audio_frame_length,
-              frame_step=p.audio_frame_step,
-              lower_edge_hertz=p.audio_lower_edge_hertz,
-              upper_edge_hertz=p.audio_upper_edge_hertz,
-              num_mel_bins=p.audio_num_mel_bins,
-              apply_mask=True)
-          if p.audio_add_delta_deltas:
-            mel_fbanks = add_delta_deltas(mel_fbanks)
-          x = tf.reshape(mel_fbanks,
-                         common_layers.shape_list(mel_fbanks)[:2] +
-                         [num_mel_bins, num_channels])
-
-          nonpadding_mask = 1. - common_attention.embedding_to_padding(x)
-          num_of_nonpadding_elements = tf.reduce_sum(
-              nonpadding_mask) * num_mel_bins * num_channels
-
-          # This replaces CMVN estimation on data
-          var_epsilon = 1e-09
-          mean = tf.reduce_sum(
-              x, axis=[1], keepdims=True) / num_of_nonpadding_elements
-          variance = (num_of_nonpadding_elements * mean**2. -
-                      2. * mean * tf.reduce_sum(x, axis=[1], keepdims=True) +
-                      tf.reduce_sum(x**2, axis=[1], keepdims=True)
-                     ) / num_of_nonpadding_elements
-          x = (x - mean) * tf.rsqrt(variance + var_epsilon) * tf.expand_dims(
-              nonpadding_mask, -1)
-      else:
-        x = inputs
-
-      # The convention is that the models are flattened along the spatial,
-      # dimensions, thus the speech preprocessor treats frequencies and
-      # channels as image colors (last axis)
-      x.set_shape([None, None, num_mel_bins, num_channels])
-
-      # TODO(chorowski): how to specify bottom's hparams and avoid hardcoding?
-      x = tf.pad(x, [[0, 0], [0, 8], [0, 0], [0, 0]])
-      for _ in range(2):
-        x = tf.layers.conv2d(
-            x, 128, (3, 3), (2, 2), use_bias=False)
-        x = common_layers.layer_norm(x)
-        x = tf.nn.relu(x)
-
-      xshape = common_layers.shape_list(x)
-      # apply a conv that will remove all frequencies and at the same time
-      # project the output into desired hidden_size
-      x = tf.pad(x, [[0, 0], [0, 2], [0, 0], [0, 0]])
-      x = tf.layers.conv2d(x, p.hidden_size, (3, xshape[2]), use_bias=False)
-
-      assert common_layers.shape_list(x)[2] == 1
-      x = common_layers.layer_norm(x)
-      x = tf.nn.relu(x)
-    return x
diff --git a/tensor2tensor/layers/common_audio.py b/tensor2tensor/layers/common_audio.py
new file mode 100644
index 000000000..30cc9ae45
--- /dev/null
+++ b/tensor2tensor/layers/common_audio.py
@@ -0,0 +1,137 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Utils for audio."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+import numpy as np
+import scipy.signal
+import tensorflow as tf
+
+
+def add_delta_deltas(filterbanks, name=None):
+  """Compute time first and second-order derivative channels.
+
+  Args:
+    filterbanks: float32 tensor with shape [batch_size, len, num_bins, 1]
+    name: scope name
+
+  Returns:
+    float32 tensor with shape [batch_size, len, num_bins, 3]
+  """
+  delta_filter = np.array([2, 1, 0, -1, -2])
+  delta_delta_filter = scipy.signal.convolve(delta_filter, delta_filter, "full")
+
+  delta_filter_stack = np.array(
+      [[0] * 4 + [1] + [0] * 4, [0] * 2 + list(delta_filter) + [0] * 2,
+       list(delta_delta_filter)],
+      dtype=np.float32).T[:, None, None, :]
+
+  delta_filter_stack /= np.sqrt(
+      np.sum(delta_filter_stack**2, axis=0, keepdims=True))
+
+  filterbanks = tf.nn.conv2d(
+      filterbanks, delta_filter_stack, [1, 1, 1, 1], "SAME", data_format="NHWC",
+      name=name)
+  return filterbanks
+
+
+def compute_mel_filterbank_features(
+    waveforms,
+    sample_rate=16000, dither=1.0 / np.iinfo(np.int16).max, preemphasis=0.97,
+    frame_length=25, frame_step=10, fft_length=None,
+    window_fn=functools.partial(tf.contrib.signal.hann_window, periodic=True),
+    lower_edge_hertz=80.0, upper_edge_hertz=7600.0, num_mel_bins=80,
+    log_noise_floor=1e-3, apply_mask=True):
+  """Implement mel-filterbank extraction using tf ops.
+
+  Args:
+    waveforms: float32 tensor with shape [batch_size, max_len]
+    sample_rate: sampling rate of the waveform
+    dither: stddev of Gaussian noise added to waveform to prevent quantization
+      artefacts
+    preemphasis: waveform high-pass filtering constant
+    frame_length: frame length in ms
+    frame_step: frame_Step in ms
+    fft_length: number of fft bins
+    window_fn: windowing function
+    lower_edge_hertz: lowest frequency of the filterbank
+    upper_edge_hertz: highest frequency of the filterbank
+    num_mel_bins: filterbank size
+    log_noise_floor: clip small values to prevent numeric overflow in log
+    apply_mask: When working on a batch of samples, set padding frames to zero
+  Returns:
+    filterbanks: a float32 tensor with shape [batch_size, len, num_bins, 1]
+  """
+  # `stfts` is a complex64 Tensor representing the short-time Fourier
+  # Transform of each signal in `signals`. Its shape is
+  # [batch_size, ?, fft_unique_bins]
+  # where fft_unique_bins = fft_length // 2 + 1
+
+  # Find the wave length: the largest index for which the value is !=0
+  # note that waveforms samples that are exactly 0.0 are quite common, so
+  # simply doing sum(waveforms != 0, axis=-1) will not work correctly.
+  wav_lens = tf.reduce_max(
+      tf.expand_dims(tf.range(tf.shape(waveforms)[1]), 0) *
+      tf.to_int32(tf.not_equal(waveforms, 0.0)),
+      axis=-1) + 1
+  if dither > 0:
+    waveforms += tf.random_normal(tf.shape(waveforms), stddev=dither)
+  if preemphasis > 0:
+    waveforms = waveforms[:, 1:] - preemphasis * waveforms[:, :-1]
+    wav_lens -= 1
+  frame_length = int(frame_length * sample_rate / 1e3)
+  frame_step = int(frame_step * sample_rate / 1e3)
+  if fft_length is None:
+    fft_length = int(2**(np.ceil(np.log2(frame_length))))
+
+  stfts = tf.contrib.signal.stft(
+      waveforms,
+      frame_length=frame_length,
+      frame_step=frame_step,
+      fft_length=fft_length,
+      window_fn=window_fn,
+      pad_end=True)
+
+  stft_lens = (wav_lens + (frame_step - 1)) // frame_step
+  masks = tf.to_float(tf.less_equal(
+      tf.expand_dims(tf.range(tf.shape(stfts)[1]), 0),
+      tf.expand_dims(stft_lens, 1)))
+
+  # An energy spectrogram is the magnitude of the complex-valued STFT.
+  # A float32 Tensor of shape [batch_size, ?, 257].
+  magnitude_spectrograms = tf.abs(stfts)
+
+  # Warp the linear-scale, magnitude spectrograms into the mel-scale.
+  num_spectrogram_bins = magnitude_spectrograms.shape[-1].value
+  linear_to_mel_weight_matrix = (
+      tf.contrib.signal.linear_to_mel_weight_matrix(
+          num_mel_bins, num_spectrogram_bins, sample_rate, lower_edge_hertz,
+          upper_edge_hertz))
+  mel_spectrograms = tf.tensordot(
+      magnitude_spectrograms, linear_to_mel_weight_matrix, 1)
+  # Note: Shape inference for tensordot does not currently handle this case.
+  mel_spectrograms.set_shape(magnitude_spectrograms.shape[:-1].concatenate(
+      linear_to_mel_weight_matrix.shape[-1:]))
+
+  log_mel_sgram = tf.log(tf.maximum(log_noise_floor, mel_spectrograms))
+
+  if apply_mask:
+    log_mel_sgram *= tf.expand_dims(tf.to_float(masks), -1)
+
+  return tf.expand_dims(log_mel_sgram, -1, name="mel_sgrams")
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index ca19e8c29..b71287dbb 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -18,6 +18,8 @@
 from __future__ import print_function
 from six.moves import range  # pylint: disable=redefined-builtin
 
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_audio
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
 from tensor2tensor.utils import modality
@@ -527,6 +529,89 @@ def xnet_resblock(x, filters, res_relu, name):
                            "compress_block_final")
 
 
+@registry.register_audio_modality
+class SpeechRecognitionModality(modality.Modality):
+  """Common ASR filterbank processing."""
+
+  def bottom(self, x):
+    """Use batchnorm instead of CMVN and shorten the stft with strided convs.
+
+    Args:
+      x: float32 tensor with shape [batch_size, len, 1, freqs * channels]
+
+    Returns:
+      float32 tensor with shape [batch_size, shorter_len, 1, hidden_size]
+    """
+    inputs = x
+    p = self._model_hparams
+
+    num_mel_bins = p.audio_num_mel_bins
+    num_channels = 3 if p.audio_add_delta_deltas else 1
+
+    with tf.variable_scope(self.name):
+      if p.audio_preproc_in_bottom:
+        # Compute filterbanks
+        with tf.variable_scope("fbanks"):
+          waveforms = tf.squeeze(inputs, [2, 3])
+          mel_fbanks = common_audio.compute_mel_filterbank_features(
+              waveforms,
+              sample_rate=p.audio_sample_rate,
+              dither=p.audio_dither,
+              preemphasis=p.audio_preemphasis,
+              frame_length=p.audio_frame_length,
+              frame_step=p.audio_frame_step,
+              lower_edge_hertz=p.audio_lower_edge_hertz,
+              upper_edge_hertz=p.audio_upper_edge_hertz,
+              num_mel_bins=p.audio_num_mel_bins,
+              apply_mask=True)
+          if p.audio_add_delta_deltas:
+            mel_fbanks = common_audio.add_delta_deltas(mel_fbanks)
+          x = tf.reshape(mel_fbanks,
+                         common_layers.shape_list(mel_fbanks)[:2] +
+                         [num_mel_bins, num_channels])
+
+          nonpadding_mask = 1. - common_attention.embedding_to_padding(x)
+          num_of_nonpadding_elements = tf.reduce_sum(
+              nonpadding_mask) * num_mel_bins * num_channels
+
+          # This replaces CMVN estimation on data
+          var_epsilon = 1e-09
+          mean = tf.reduce_sum(
+              x, axis=[1], keepdims=True) / num_of_nonpadding_elements
+          variance = (num_of_nonpadding_elements * mean**2. -
+                      2. * mean * tf.reduce_sum(x, axis=[1], keepdims=True) +
+                      tf.reduce_sum(x**2, axis=[1], keepdims=True)
+                     ) / num_of_nonpadding_elements
+          x = (x - mean) * tf.rsqrt(variance + var_epsilon) * tf.expand_dims(
+              nonpadding_mask, -1)
+      else:
+        x = inputs
+
+      # The convention is that the models are flattened along the spatial,
+      # dimensions, thus the speech preprocessor treats frequencies and
+      # channels as image colors (last axis)
+      x.set_shape([None, None, num_mel_bins, num_channels])
+
+      # TODO(chorowski): how to specify bottom's hparams and avoid hardcoding?
+      x = tf.pad(x, [[0, 0], [0, 8], [0, 0], [0, 0]])
+      for _ in range(2):
+        x = tf.layers.conv2d(
+            x, 128, (3, 3), (2, 2), use_bias=False)
+        x = common_layers.layer_norm(x)
+        x = tf.nn.relu(x)
+
+      xshape = common_layers.shape_list(x)
+      # apply a conv that will remove all frequencies and at the same time
+      # project the output into desired hidden_size
+      x = tf.pad(x, [[0, 0], [0, 2], [0, 0], [0, 0]])
+      x = tf.layers.conv2d(x, p.hidden_size, (3, xshape[2]), use_bias=False)
+
+      assert common_layers.shape_list(x)[2] == 1
+      x = common_layers.layer_norm(x)
+      x = tf.nn.relu(x)
+    return x
+
+
 @registry.register_video_modality("default")
 class VideoModality(modality.Modality):
   """Modality for videos, i.e., time-sequences of frames."""

From f0cf931a2c5c90f751274ee10212749c62e060ab Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 12 Sep 2018 11:02:51 -0700
Subject: [PATCH 0824/2720] simplifying the VideoModality and Basic model.
 cleaning up a lot of reshaping, etc. The code still works with both L2 and
 Softmax loss. Will try to clean up VideoModalities i.e. merging them into one
 next.

PiperOrigin-RevId: 212668869
---
 tensor2tensor/layers/common_layers.py         |  14 --
 tensor2tensor/layers/modalities.py            |  19 +--
 .../models/video/basic_deterministic.py       | 131 ++++++++++++------
 .../video/basic_deterministic_params.py       |   3 +
 tensor2tensor/models/video/sv2p.py            |  19 +--
 tensor2tensor/models/video/sv2p_params.py     |   5 +-
 6 files changed, 100 insertions(+), 91 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 71546e6de..2eccea693 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3547,20 +3547,6 @@ def summarize_video(video, prefix, max_outputs=1):
           max_outputs=max_outputs)
 
 
-def time_to_channels(embedded_video):
-  """Put time dimension on channels in an embedded video."""
-  video_shape = shape_list(embedded_video)
-  if len(video_shape) != 5:
-    raise ValueError("Assuming videos given as tensors in the format "
-                     "[batch, time, height, width, channels] but got one "
-                     "of shape: %s" % str(video_shape))
-  transposed = tf.transpose(embedded_video, [0, 2, 3, 1, 4])
-  return tf.reshape(transposed, [
-      video_shape[0], video_shape[2], video_shape[3],
-      video_shape[1] * video_shape[4]
-  ])
-
-
 def cast_like(x, y):
   """Cast x to y's dtype, if necessary."""
   x = tf.convert_to_tensor(x)
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index b71287dbb..e78cd0072 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -621,27 +621,17 @@ def bottom(self, x):
     with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
       common_layers.summarize_video(inputs, "inputs")
       inputs = common_layers.standardize_images(inputs)
-      return common_layers.time_to_channels(inputs)
+      return inputs
 
   def targets_bottom(self, x):
     return self.bottom(x)
 
   def top(self, body_output, targets):
     num_channels = self._model_hparams.problem.num_channels
-    num_frames = common_layers.shape_list(targets)[1]
     body_output_shape = common_layers.shape_list(body_output)
-    # We assume the body output is of this shape and layout.
-    # Note: if you tf.concat([frames], axis=-1) at the end of your model,
-    # then you need to reshape to [..., num_frames, depth] like below, not
-    # into [..., depth, num_frames] due to memory layout of concat/reshape.
     reshape_shape = body_output_shape[:-1] + [
-        num_frames, num_channels, self.top_dimensionality]
+        num_channels, self.top_dimensionality]
     res = tf.reshape(body_output, reshape_shape)
-    res = tf.transpose(res, [0, 3, 1, 2, 4, 5])
-    res_shape = common_layers.shape_list(res)
-    res_argmax = tf.argmax(tf.reshape(res, [-1, res_shape[-1]]), axis=-1)
-    res_argmax = tf.reshape(res_argmax, res_shape[:-1])
-    common_layers.summarize_video(res_argmax, "result")
     return res
 
   def loss(self, top_out, targets):
@@ -671,10 +661,9 @@ def bottom(self, x):
       assert self.top_dimensionality == 256
       embedded = discretization.int_to_bit_embed(inputs, 8,
                                                  self.PIXEL_EMBEDDING_SIZE)
-      # Transpose and project.
-      transposed = common_layers.time_to_channels(embedded)
+      # Project.
       return tf.layers.dense(
-          transposed,
+          embedded,
           self._body_input_depth,
           name="merge_pixel_embedded_frames")
 
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 144708ef5..3f4e4da08 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -22,6 +22,7 @@
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
 from tensor2tensor.models.video import basic_deterministic_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
@@ -33,26 +34,54 @@
 tfcl = tf.contrib.layers
 
 
-def inject_action(action, x):
-  """Inject the action into x."""
-  x_shape = common_layers.shape_list(x)
-  filters = x_shape[-1]
-  action_mask = tf.layers.dense(action, filters, name="action_mask")
-  action_add = tf.layers.dense(action, filters, name="action_add")
-  x *= tf.nn.sigmoid(action_mask)
-  x += action_add
-  return x
-
-
 @registry.register_model
 class NextFrameBasicDeterministic(t2t_model.T2TModel):
   """Basic next-frame model, may take actions and predict rewards too."""
 
+  @property
+  def is_per_pixel_softmax(self):
+    # TODO(mbz): this should not be a hyper parameter.
+    return self.hparams.per_pixel_softmax
+
   def inject_latent(self, layer, features, filters):
     """Do nothing for deterministic model."""
     del features, filters
     return layer, 0.0
 
+  def inject_additional_input(self, layer, inputs, scope, mode="concat"):
+    layer_shape = common_layers.shape_list(layer)
+    input_shape = common_layers.shape_list(inputs)
+    zeros_mask = tf.zeros(layer_shape, dtype=tf.float32)
+    if mode == "concat":
+      emb = common_video.encode_to_shape(inputs, layer_shape, scope)
+      layer = tf.concat(values=[layer, emb], axis=-1)
+    elif mode == "multiplicative":
+      filters = layer_shape[-1]
+      input_reshaped = tf.reshape(inputs, [-1, 1, 1, input_shape[-1]])
+      input_mask = tf.layers.dense(input_reshaped, filters, name=scope)
+      input_broad = input_mask + zeros_mask
+      layer *= input_broad
+    elif mode == "multi_additive":
+      filters = layer_shape[-1]
+      input_reshaped = tf.reshape(inputs, [-1, 1, 1, input_shape[-1]])
+      input_mul_mask = tf.layers.dense(input_reshaped, filters, name=scope+"_m")
+      input_mul = input_mul_mask + zeros_mask
+      layer *= input_mul
+      input_add_mask = tf.layers.dense(input_reshaped, filters, name=scope+"_a")
+      input_add = input_add_mask + zeros_mask
+      layer += input_add
+    else:
+      raise ValueError("Unknown injection mode: %s" % mode)
+
+    return layer
+
+  def get_sampled_frame(self, res_frame, orig_frame_shape):
+    target_shape = orig_frame_shape[:-1] + [self.hparams.problem.num_channels]
+    if self.is_per_pixel_softmax:
+      sampled_frame = tf.reshape(res_frame, target_shape + [256])
+      sampled_frame = tf.to_float(tf.argmax(sampled_frame, axis=-1))
+    return sampled_frame
+
   def body_single(self, features):
     hparams = self.hparams
     filters = hparams.hidden_size
@@ -81,9 +110,9 @@ def body_single(self, features):
 
     # Add embedded action if present.
     if "input_action" in features:
-      action = tf.reshape(features["input_action"][:, -1, :],
-                          [-1, 1, 1, hparams.hidden_size])
-      x = inject_action(action, x)
+      action = features["input_action"][:, -1, :]
+      x = self.inject_additional_input(
+          x, action, "action_enc", hparams.action_injection)
 
     x, extra_loss = self.inject_latent(x, features, filters)
 
@@ -103,7 +132,8 @@ def body_single(self, features):
     for i in range(hparams.num_compress_steps):
       with tf.variable_scope("upstride%d" % i):
         if "input_action" in features:
-          x = inject_action(action, x)
+          x = self.inject_additional_input(
+              x, action, "action_enc", hparams.action_injection)
         if i >= hparams.num_compress_steps - hparams.filter_double_steps:
           filters //= 2
         x = tf.layers.conv2d_transpose(
@@ -117,7 +147,10 @@ def body_single(self, features):
 
     # Cut down to original size.
     x = x[:, :inputs_shape[1], :inputs_shape[2], :]
-    x = tf.layers.dense(x, hparams.problem.num_channels * 256, name="logits")
+    if self.is_per_pixel_softmax:
+      x = tf.layers.dense(x, hparams.problem.num_channels * 256, name="logits")
+    else:
+      x = tf.layers.dense(x, hparams.problem.num_channels, name="logits")
 
     # Reward prediction if needed.
     if "target_reward" not in features:
@@ -129,18 +162,13 @@ def body_single(self, features):
   def body(self, features):
     hparams = self.hparams
     is_predicting = hparams.mode == tf.estimator.ModeKeys.PREDICT
-    if hparams.video_num_target_frames < 2:
-      res = self.body_single(features)
-      return res
 
     # TODO(lukaszkaiser): the split axes and the argmax below heavily depend on
     # using the default (a bit strange) video modality - we should change that.
 
     # Split inputs and targets into lists.
-    input_frames = list(tf.split(
-        features["inputs"], hparams.video_num_input_frames, axis=-1))
-    target_frames = list(tf.split(
-        features["targets"], hparams.video_num_target_frames, axis=-1))
+    input_frames = tf.unstack(features["inputs"], axis=1)
+    target_frames = tf.unstack(features["targets"], axis=1)
     all_frames = input_frames + target_frames
     if "input_action" in features:
       input_actions = list(tf.split(
@@ -149,8 +177,10 @@ def body(self, features):
           features["target_action"], hparams.video_num_target_frames, axis=1))
       all_actions = input_actions + target_actions
 
+    orig_frame_shape = common_layers.shape_list(all_frames[0])
+
     # Run a number of steps.
-    res_frames = []
+    res_frames, sampled_frames, sampled_frames_raw = [], [], []
     if "target_reward" in features:
       res_rewards, extra_loss = [], 0.0
     sample_prob = common_layers.inverse_exp_decay(
@@ -166,31 +196,42 @@ def body(self, features):
       # Run model.
       with tf.variable_scope(tf.get_variable_scope(), reuse=i > 0):
         if "target_reward" not in features:
-          res_frames.append(self.body_single(features))
+          res_frame = self.body_single(features)
         else:
           res_dict, res_extra_loss = self.body_single(features)
           extra_loss += res_extra_loss
-          res_frames.append(res_dict["targets"])
-          res_rewards.append(res_dict["target_reward"])
-
-      # When predicting, use the generated frame.
-      orig_frame = all_frames[i + hparams.video_num_input_frames]
-      shape = common_layers.shape_list(orig_frame)
-      sampled_frame = tf.reshape(
-          res_frames[-1], shape[:-1] + [hparams.problem.num_channels, 256])
-      sampled_frame = tf.to_float(tf.argmax(sampled_frame, axis=-1))
-      sampled_frame = common_layers.standardize_images(sampled_frame)
+          res_frame = res_dict["targets"]
+          res_reward = res_dict["target_reward"]
+          res_rewards.append(res_reward)
+      res_frames.append(res_frame)
+
+      # Only for Softmax loss: sample frame so we can keep iterating.
+      sampled_frame_raw = self.get_sampled_frame(res_frame, orig_frame_shape)
+      sampled_frames_raw.append(sampled_frame_raw)
+      # TODO(lukaszkaiser): this should be consistent with modality.bottom()
+      sampled_frame = common_layers.standardize_images(sampled_frame_raw)
+      sampled_frames.append(sampled_frame)
+
       if is_predicting:
         all_frames[i + hparams.video_num_input_frames] = sampled_frame
 
       # Scheduled sampling during training.
       if (hparams.scheduled_sampling_prob > 0.0 and self.is_training):
-        do_sample = tf.less(tf.random_uniform([shape[0]]), sample_prob)
+        do_sample = tf.less(
+            tf.random_uniform([orig_frame_shape[0]]), sample_prob)
+        orig_frame = all_frames[i + hparams.video_num_input_frames]
         sampled_frame = tf.where(do_sample, sampled_frame, orig_frame)
         all_frames[i + hparams.video_num_input_frames] = sampled_frame
 
     # Concatenate results and return them.
-    frames = tf.concat(res_frames, axis=-1)
+    frames = tf.stack(res_frames, axis=1)
+
+    if self.is_per_pixel_softmax:
+      def make_gif_ready(tensor_list):
+        return tf.cast(tf.stack(tensor_list, axis=1), tf.uint8)
+      summary = common_video.gif_summary
+      summary("pred", make_gif_ready(sampled_frames_raw))
+
     if "target_reward" not in features:
       return frames
     rewards = tf.concat(res_rewards, axis=1)
@@ -202,6 +243,7 @@ def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
     # Inputs and features preparation needed to handle edge cases.
     if not features:
       features = {}
+    hparams = self.hparams
     inputs_old = None
     if "inputs" in features and len(features["inputs"].shape) < 4:
       inputs_old = features["inputs"]
@@ -219,22 +261,27 @@ def logits_to_samples(logits):
 
     # Get predictions.
     try:
-      num_channels = self.hparams.problem.num_channels
+      num_channels = hparams.problem.num_channels
     except AttributeError:
       num_channels = 1
     if "inputs" in features:
       inputs_shape = common_layers.shape_list(features["inputs"])
-      targets_shape = [inputs_shape[0], self.hparams.video_num_target_frames,
+      targets_shape = [inputs_shape[0], hparams.video_num_target_frames,
                        inputs_shape[2], inputs_shape[3], num_channels]
     else:
       tf.logging.warn("Guessing targets shape as no inputs are given.")
-      targets_shape = [self.hparams.batch_size,
-                       self.hparams.video_num_target_frames, 1, 1, num_channels]
+      targets_shape = [hparams.batch_size,
+                       hparams.video_num_target_frames, 1, 1, num_channels]
 
     features["targets"] = tf.zeros(targets_shape, dtype=tf.int32)
-    if "target_reward" in self.hparams.problem_hparams.target_modality:
+    reward_in_mod = "target_reward" in hparams.problem_hparams.target_modality
+    action_in_mod = "target_action" in hparams.problem_hparams.target_modality
+    if reward_in_mod:
       features["target_reward"] = tf.zeros(
           [targets_shape[0], 1, 1], dtype=tf.int32)
+    if action_in_mod and "target_action" not in features:
+      features["target_action"] = tf.zeros(
+          [targets_shape[0], 1, 1], dtype=tf.int32)
     logits, _ = self(features)  # pylint: disable=not-callable
     if isinstance(logits, dict):
       results = {}
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index c19c25985..5168823f5 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -40,6 +40,9 @@ def next_frame_basic_deterministic():
   hparams.weight_decay = 0.0
   hparams.clip_grad_norm = 1.0
   hparams.dropout = 0.5
+  hparams.add_hparam("per_pixel_softmax", True)
+  # choose from: concat, multiplicative, multi_additive
+  hparams.add_hparam("action_injection", "multi_additive")
   hparams.add_hparam("num_compress_steps", 6)
   hparams.add_hparam("filter_double_steps", 2)
   hparams.add_hparam("video_modality_loss_cutoff", 0.02)
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 9298599a2..9275fb30d 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -125,21 +125,6 @@ def get_input_if_exists(self, features, key, batch_size, num_frames):
       x = tf.zeros((batch_size, num_frames, 1, self.hparams.hidden_size))
     return common_video.swap_time_and_batch_axes(x)
 
-  def inject_additional_input(self, layer, inputs, scope, concatenate=True):
-    layer_shape = common_layers.shape_list(layer)
-    input_shape = common_layers.shape_list(inputs)
-    if concatenate:
-      emb = common_video.encode_to_shape(inputs, layer_shape, scope)
-      layer = tf.concat(values=[layer, emb], axis=-1)
-    else:
-      filters = layer_shape[-1]
-      input_reshaped = tf.reshape(inputs, [-1, 1, 1, input_shape[-1]])
-      input_mask = tf.layers.dense(input_reshaped, filters, name=scope)
-      zeros_mask = tf.zeros(layer_shape, dtype=tf.float32)
-      input_broad = input_mask + zeros_mask
-      layer *= input_broad
-    return layer
-
   def bottom_part_tower(self, input_image, input_reward, action, latent,
                         lstm_state, lstm_size, conv_size, concat_latent=False):
     """The bottom part of predictive towers.
@@ -205,7 +190,7 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
 
     if action is not None:
       enc2 = self.inject_additional_input(
-          enc2, action, "action_enc", self.hparams.concatenate_actions)
+          enc2, action, "action_enc", self.hparams.action_injection)
     if input_reward is not None:
       enc2 = self.inject_additional_input(enc2, input_reward, "reward_enc")
     if latent is not None and not concat_latent:
@@ -235,7 +220,7 @@ def reward_prediction(self, input_images, input_reward, action, latent):
       # Inject additional inputs
       if action is not None:
         x = self.inject_additional_input(
-            x, action, "action_enc", self.hparams.concatenate_actions)
+            x, action, "action_enc", self.hparams.action_injection)
       if input_reward is not None:
         x = self.inject_additional_input(x, input_reward, "reward_enc")
       if latent is not None:
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 306eca995..6b4ada8b6 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -48,7 +48,6 @@ def next_frame_sv2p():
   hparams.add_hparam("scheduled_sampling_k", 900.0)
   hparams.add_hparam("upsample_method", "conv2d_transpose")
   hparams.add_hparam("internal_loss", False)
-  hparams.add_hparam("concatenate_actions", True)
   return hparams
 
 
@@ -58,12 +57,12 @@ def next_frame_sv2p_atari():
   hparams = next_frame_sv2p()
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 4
-  hparams.concatenate_actions = False
+  hparams.action_injection = "multiplicative"
   hparams.num_iterations_1st_stage = 15000
   hparams.num_iterations_2nd_stage = 15000
+  hparams.anneal_end = 50000
   hparams.latent_loss_multiplier_schedule = "noisy_linear_cosine_decay"
   hparams.latent_loss_multiplier = 1e-3
-  hparams.anneal_end = 50000
   hparams.preprocess_resize_frames = [96, 96]
   hparams.information_capacity = 0.0
   return hparams

From f29917e2b9aa1b19a4155b839e5623d868f1de11 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Wed, 12 Sep 2018 11:16:48 -0700
Subject: [PATCH 0825/2720] plint fixes

---
 tensor2tensor/rl/envs/tf_atari_wrappers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index d2c5ff813..3aa696d07 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -340,4 +340,4 @@ def __len__(self):
 
   def _reset_non_empty(self, indices):
     # pylint: disable=protected-access
-    return self._batch_env._reset_non_empty(indices)
\ No newline at end of file
+    return self._batch_env._reset_non_empty(indices)

From 6e66063ccbf0e98215b925491a62aa82cb57a55f Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 12 Sep 2018 11:27:06 -0700
Subject: [PATCH 0826/2720] Bring back sigmoid in action injection.

PiperOrigin-RevId: 212673660
---
 tensor2tensor/models/video/basic_deterministic.py | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 3f4e4da08..f1861f899 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -48,27 +48,25 @@ def inject_latent(self, layer, features, filters):
     del features, filters
     return layer, 0.0
 
-  def inject_additional_input(self, layer, inputs, scope, mode="concat"):
+  def inject_additional_input(self, layer, inputs, name, mode="concat"):
     layer_shape = common_layers.shape_list(layer)
     input_shape = common_layers.shape_list(inputs)
     zeros_mask = tf.zeros(layer_shape, dtype=tf.float32)
     if mode == "concat":
-      emb = common_video.encode_to_shape(inputs, layer_shape, scope)
+      emb = common_video.encode_to_shape(inputs, layer_shape, name)
       layer = tf.concat(values=[layer, emb], axis=-1)
     elif mode == "multiplicative":
       filters = layer_shape[-1]
       input_reshaped = tf.reshape(inputs, [-1, 1, 1, input_shape[-1]])
-      input_mask = tf.layers.dense(input_reshaped, filters, name=scope)
+      input_mask = tf.layers.dense(input_reshaped, filters, name=name)
       input_broad = input_mask + zeros_mask
       layer *= input_broad
     elif mode == "multi_additive":
       filters = layer_shape[-1]
       input_reshaped = tf.reshape(inputs, [-1, 1, 1, input_shape[-1]])
-      input_mul_mask = tf.layers.dense(input_reshaped, filters, name=scope+"_m")
-      input_mul = input_mul_mask + zeros_mask
-      layer *= input_mul
-      input_add_mask = tf.layers.dense(input_reshaped, filters, name=scope+"_a")
-      input_add = input_add_mask + zeros_mask
+      input_mul = tf.layers.dense(input_reshaped, filters, name=name + "_mul")
+      layer *= tf.nn.sigmoid(input_mul)
+      input_add = tf.layers.dense(input_reshaped, filters, name=name + "_add")
       layer += input_add
     else:
       raise ValueError("Unknown injection mode: %s" % mode)

From 9907e4c06f73de050107eaf98efcee4371e96b08 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Wed, 12 Sep 2018 13:45:19 -0700
Subject: [PATCH 0827/2720] First version of LTVP-EPVA model.

PiperOrigin-RevId: 212696436
---
 tensor2tensor/layers/common_video.py      |  29 +
 tensor2tensor/models/__init__.py          |   1 +
 tensor2tensor/models/video/epva.py        | 719 ++++++++++++++++++++++
 tensor2tensor/models/video/epva_params.py |  49 ++
 4 files changed, 798 insertions(+)
 create mode 100644 tensor2tensor/models/video/epva.py
 create mode 100644 tensor2tensor/models/video/epva_params.py

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 469156f23..9aceab42d 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -65,6 +65,35 @@ def basic_lstm(inputs, state, num_units, name=None):
   return outputs, new_state
 
 
+def lstm_cell(inputs,
+              state,
+              num_units,
+              use_peepholes=False,
+              cell_clip=0.0,
+              initializer=None,
+              num_proj=None,
+              num_unit_shards=None,
+              num_proj_shards=None,
+              reuse=None,
+              name=None):
+  """Full LSTM cell."""
+  input_shape = common_layers.shape_list(inputs)
+  cell = tf.contrib.rnn.LSTMCell(num_units,
+                                 use_peepholes=use_peepholes,
+                                 cell_clip=cell_clip,
+                                 initializer=initializer,
+                                 num_proj=num_proj,
+                                 num_unit_shards=num_unit_shards,
+                                 num_proj_shards=num_proj_shards,
+                                 reuse=reuse,
+                                 name=name,
+                                 state_is_tuple=False)
+  if state is None:
+    state = cell.zero_state(input_shape[0], tf.float32)
+  outputs, new_state = cell(inputs, state)
+  return outputs, new_state
+
+
 def conv_lstm_2d(inputs, state, output_channels,
                  kernel_size=5, name=None, spatial_dims=None):
   """2D Convolutional LSTM."""
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 1177bc756..e5e6ed265 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -64,6 +64,7 @@
 from tensor2tensor.models.video import basic_deterministic
 from tensor2tensor.models.video import basic_stochastic
 from tensor2tensor.models.video import emily
+from tensor2tensor.models.video import epva
 from tensor2tensor.models.video import savp
 from tensor2tensor.models.video import sv2p
 
diff --git a/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
new file mode 100644
index 000000000..7aa7c3f76
--- /dev/null
+++ b/tensor2tensor/models/video/epva.py
@@ -0,0 +1,719 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Model architecture for video prediction model.
+
+based on following paper:
+"Hierarchical Long-term Video Prediction without Supervision"
+http://web.eecs.umich.edu/~honglak/icml2018-unsupHierarchicalVideoPred.pdf
+by Nevan Wichers, Ruben Villegas, Dumitru Erhan and Honglak Lee.
+
+This code is based on the original code:
+https://github.com/brain-research/long-term-video-prediction-without-supervision
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
+from tensor2tensor.models.video import epva_params  # pylint: disable=unused-import
+from tensor2tensor.models.video import sv2p
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+from tensorflow.contrib.framework.python.ops import arg_scope
+from tensorflow.contrib.slim.python.slim.nets import vgg
+
+tfl = tf.layers
+tfcl = tf.contrib.layers
+
+IMG_WIDTH = 64
+IMG_HEIGHT = 64
+VGG_IMAGE_SIZE = 224
+COLOR_NORMALIZATION_VECTOR = [123.68, 116.78, 103.94]
+
+
+def van_image_enc_2d(x, first_depth, reuse=False, hparams=None):
+  """The image encoder for the VAN.
+
+  Similar architecture as Ruben's paper
+  (http://proceedings.mlr.press/v70/villegas17a/villegas17a.pdf).
+
+  Args:
+    x: The image to encode.
+    first_depth: The depth of the first layer. Depth is increased in subsequent
+      layers.
+    reuse: To reuse in variable scope or not.
+    hparams: The python hparams.
+
+  Returns:
+    The encoded image.
+  """
+  with tf.variable_scope('van_image_enc', reuse=reuse):
+    enc_history = [x]
+
+    enc = tf.layers.conv2d(
+        x, first_depth, 3, padding='same', activation=tf.nn.relu, strides=1)
+    enc = tf.contrib.layers.layer_norm(enc)
+    enc = tf.layers.conv2d(
+        enc, first_depth, 3, padding='same', activation=tf.nn.relu, strides=1)
+    enc = tf.nn.max_pool(enc, [1, 2, 2, 1], [1, 2, 2, 1], 'SAME')
+    enc = tf.nn.dropout(enc, hparams.van_keep_prob)
+    enc = tf.contrib.layers.layer_norm(enc)
+    enc_history.append(enc)
+
+    enc = tf.layers.conv2d(
+        enc,
+        first_depth * 2,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=1)
+    enc = tf.layers.conv2d(
+        enc,
+        first_depth * 2,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=1)
+    enc = tf.nn.max_pool(enc, [1, 2, 2, 1], [1, 2, 2, 1], 'SAME')
+    enc = tf.nn.dropout(enc, hparams.van_keep_prob)
+    enc = tf.contrib.layers.layer_norm(enc)
+    enc_history.append(enc)
+
+    enc = tf.layers.conv2d(
+        enc,
+        first_depth * 4,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=1)
+    enc = tf.layers.conv2d(
+        enc,
+        first_depth * 4,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=1)
+    enc = tf.layers.conv2d(
+        enc,
+        first_depth * 4,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=1)
+    enc = tf.nn.max_pool(enc, [1, 2, 2, 1], [1, 2, 2, 1], 'SAME')
+
+    return enc, enc_history
+
+
+def van_enc_2d(x, first_depth, reuse=False):
+  """The higher level structure encoder for the VAN.
+
+  The high level structure is a vector instead of an image.
+
+  Args:
+    x: The higher level structure to encode.
+    first_depth: The depth of the first layer. Depth is increased in subsequent
+      layers.
+    reuse: To reuse in variable scope or not.
+
+  Returns:
+    The encoded image.
+  """
+  with tf.variable_scope('van_enc', reuse=reuse):
+    a = 4  # depends on the inputs size
+    b = 4
+    # a, b = 4,4
+    enc = tf.nn.relu(x)
+    enc = tf.layers.dense(enc, first_depth * a * b, tf.nn.relu)
+    enc = tf.contrib.layers.layer_norm(enc)
+
+    enc = tf.reshape(enc, [-1, a, b, first_depth])
+
+    enc = tf.layers.conv2d_transpose(
+        enc, first_depth, 3, padding='same', activation=tf.nn.relu, strides=1)
+    enc = tf.contrib.layers.layer_norm(enc)
+    enc = tf.layers.conv2d_transpose(
+        enc,
+        first_depth * 2,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=2)
+    van_higher_level_2 = tf.reshape(enc, [-1, a * 2 * b * 2 * first_depth * 2])
+
+    enc = tf.layers.conv2d_transpose(
+        enc,
+        first_depth * 2,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=1)
+    enc = tf.contrib.layers.layer_norm(enc)
+    enc = tf.layers.conv2d_transpose(
+        enc,
+        first_depth * 4,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=1)
+    van_higher_level_4 = tf.reshape(enc, [-1, a * 2 * b * 2 * first_depth * 4])
+
+    van_higher_level = tf.concat([x, van_higher_level_2, van_higher_level_4], 1)
+
+    return enc, van_higher_level
+
+
+def van_dec_2d(x, skip_connections, output_shape, first_depth, hparams=None):
+  """The VAN decoder.
+
+  Args:
+    x: The analogy information to decode.
+    skip_connections: The encoder layers which can be used as skip connections.
+    output_shape: The shape of the desired output image.
+    first_depth: The depth of the first layer of the van image encoder.
+    hparams: The python hparams.
+
+  Returns:
+    The decoded image prediction.
+  """
+  with tf.variable_scope('van_dec'):
+    dec = tf.layers.conv2d_transpose(
+        x, first_depth * 4, 3, padding='same', activation=tf.nn.relu, strides=2)
+    dec = tf.nn.dropout(dec, hparams.van_keep_prob)
+    dec = tf.contrib.layers.layer_norm(dec)
+    dec = tf.layers.conv2d_transpose(
+        dec,
+        first_depth * 4,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=1)
+    dec = tf.nn.dropout(dec, hparams.van_keep_prob)
+    dec = tf.layers.conv2d_transpose(
+        dec,
+        first_depth * 2,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=1)
+    dec = tf.nn.dropout(dec, hparams.van_keep_prob)
+    dec = tf.contrib.layers.layer_norm(dec)
+
+    dec = tf.layers.conv2d_transpose(
+        dec,
+        first_depth * 2,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=2)
+    dec = tf.nn.dropout(dec, hparams.van_keep_prob)
+    dec = tf.layers.conv2d_transpose(
+        dec, first_depth, 3, padding='same', activation=tf.nn.relu, strides=1)
+    dec = tf.nn.dropout(dec, hparams.van_keep_prob)
+    dec = tf.contrib.layers.layer_norm(dec)
+
+    dec = tf.layers.conv2d_transpose(
+        dec,
+        output_shape[3] + 1,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=2)
+    dec = tf.nn.dropout(dec, hparams.van_keep_prob)
+
+    out_mask = tf.layers.conv2d_transpose(
+        dec, output_shape[3] + 1, 3, strides=1, padding='same', activation=None)
+
+    mask = tf.nn.sigmoid(out_mask[:, :, :, 3:4])
+    out = out_mask[:, :, :, :3]
+
+    return out * mask + skip_connections[0] * (1 - mask)
+
+
+def analogy_computation_2d(f_first_enc,
+                           f_first_frame,
+                           f_current_enc,
+                           first_depth):
+  """Implements the deep analogy computation."""
+  with tf.variable_scope('analogy_computation'):
+
+    frame_enc_diff = f_first_frame - f_first_enc
+
+    frame_enc_diff_enc = tf.layers.conv2d(
+        frame_enc_diff,
+        first_depth * 4,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=1)
+    f_current_enc_enc = tf.layers.conv2d(
+        f_current_enc,
+        first_depth * 4,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=1)
+
+    analogy = tf.concat([frame_enc_diff_enc, f_current_enc_enc], 3)
+    analogy = tf.layers.conv2d(
+        analogy,
+        first_depth * 4,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=1)
+    analogy = tf.contrib.layers.layer_norm(analogy)
+    analogy = tf.layers.conv2d(
+        analogy,
+        first_depth * 4,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=1)
+    return tf.layers.conv2d(
+        analogy,
+        first_depth * 4,
+        3,
+        padding='same',
+        activation=tf.nn.relu,
+        strides=1)
+
+
+def van(first_enc,
+        first_frame,
+        current_enc,
+        gt_image,
+        reuse=False,
+        scope_prefix='',
+        hparams=None):
+  """Implements a VAN.
+
+  Args:
+    first_enc: The first encoding.
+    first_frame: The first ground truth frame.
+    current_enc: The encoding of the frame to generate.
+    gt_image: The ground truth image, only used for regularization.
+    reuse: To reuse in variable scope or not.
+    scope_prefix: The prefix before the scope name.
+    hparams: The python hparams.
+
+  Returns:
+    The generated image.
+  """
+  with tf.variable_scope(scope_prefix + 'van', reuse=reuse):
+    output_shape = first_frame.get_shape().as_list()
+    output_shape[0] = -1
+
+    first_depth = 64
+
+    f_first_enc, _ = van_enc_2d(first_enc, first_depth)
+    f_first_frame, image_enc_history = van_image_enc_2d(
+        first_frame, first_depth, hparams=hparams)
+    f_current_enc, van_higher_level = van_enc_2d(
+        current_enc, first_depth, reuse=True)
+    f_gt_image, _ = van_image_enc_2d(gt_image, first_depth, True,
+                                     hparams=hparams)
+
+    analogy_t = analogy_computation_2d(
+        f_first_enc, f_first_frame, f_current_enc, first_depth)
+    enc_img = f_current_enc + analogy_t
+
+    img = van_dec_2d(
+        enc_img, image_enc_history, output_shape, first_depth, hparams=hparams)
+
+    batch_size = tf.to_float(tf.shape(first_enc)[0])
+    r_loss = tf.nn.l2_loss(f_gt_image - f_current_enc - analogy_t) / batch_size
+
+    return img, r_loss, van_higher_level
+
+
+def encoder_vgg(x, enc_final_size, reuse=False, scope_prefix='', hparams=None,
+                is_training=True):
+  """VGG network to use as encoder without the top few layers.
+
+  Can be pretrained.
+
+  Args:
+    x: The image to encode. In the range 0 to 1.
+    enc_final_size: The desired size of the encoding.
+    reuse: To reuse in variable scope or not.
+    scope_prefix: The prefix before the scope name.
+    hparams: The python hparams.
+    is_training: boolean value indicating if training is happening.
+
+  Returns:
+    The generated image.
+  """
+  with tf.variable_scope(scope_prefix + 'encoder', reuse=reuse):
+
+    # Preprocess input
+    x *= 256
+    x = x - COLOR_NORMALIZATION_VECTOR
+
+    with arg_scope(vgg.vgg_arg_scope()):
+      # Padding because vgg_16 accepts images of size at least VGG_IMAGE_SIZE.
+      x = tf.pad(x, [[0, 0], [0, VGG_IMAGE_SIZE - IMG_WIDTH],
+                     [0, VGG_IMAGE_SIZE - IMG_HEIGHT], [0, 0]])
+      _, end_points = vgg.vgg_16(
+          x,
+          num_classes=enc_final_size,
+          is_training=is_training)
+      pool5_key = [key for key in end_points.keys() if 'pool5' in key]
+      assert len(pool5_key) == 1
+      enc = end_points[pool5_key[0]]
+      # Undoing padding.
+      enc = tf.slice(enc, [0, 0, 0, 0], [-1, 2, 2, -1])
+
+    enc_shape = enc.get_shape().as_list()
+    enc_shape[0] = -1
+    enc_size = enc_shape[1] * enc_shape[2] * enc_shape[3]
+
+    enc_flat = tf.reshape(enc, (-1, enc_size))
+    enc_flat = tf.nn.dropout(enc_flat, hparams.enc_keep_prob)
+
+    enc_flat = tf.layers.dense(
+        enc_flat,
+        enc_final_size,
+        kernel_initializer=tf.truncated_normal_initializer(stddev=1e-4,))
+
+    if hparams.enc_pred_use_l2norm:
+      enc_flat = tf.nn.l2_normalize(enc_flat, 1)
+
+  return enc_flat
+
+
+def predictor(enc_flat,
+              action,
+              lstm_states,
+              pred_depth,
+              reuse=False,
+              scope_prefix='',
+              hparams=None):
+  """LSTM predictor network."""
+  with tf.variable_scope(scope_prefix + 'predict', reuse=reuse):
+
+    enc_final_size = enc_flat.get_shape().as_list()[1]
+    action_size = action.get_shape().as_list()[1]
+    initial_size = (enc_final_size + action_size)
+
+    batch_size = tf.shape(enc_flat)[0]
+
+    init_stddev = 1e-2
+
+    pre_pred = tf.concat([enc_flat, action], 1)
+    pre_pred = tf.layers.dense(
+        pre_pred,
+        initial_size,
+        kernel_initializer=tf.truncated_normal_initializer(stddev=init_stddev))
+
+    # This is only needed or the GAN version.
+    if hparams.pred_noise_std > 0:
+      # Add the noise like this so a pretrained model can be used.
+      pred_noise = tf.random_normal(
+          shape=[batch_size, 100], stddev=hparams.pred_noise_std)
+      pre_pred += tf.layers.dense(
+          pred_noise,
+          initial_size,
+          kernel_initializer=tf.truncated_normal_initializer(
+              stddev=init_stddev),
+          name='noise_dense')
+
+    pre_pred = tf.nn.relu(pre_pred)
+
+    if lstm_states[pred_depth - 2] is None:
+      back_connect = tf.tile(
+          tf.get_variable(
+              'back_connect_init',
+              shape=[1, initial_size * 2],
+              initializer=tf.truncated_normal_initializer(stddev=init_stddev))
+          , (batch_size, 1))
+    else:
+      back_connect = lstm_states[pred_depth - 2]
+
+    lstm_init_stddev = 1e-4
+
+    part_pred, lstm_states[0] = common_video.lstm_cell(
+        tf.concat([pre_pred, back_connect], 1),
+        lstm_states[0],
+        initial_size,
+        use_peepholes=True,
+        initializer=tf.truncated_normal_initializer(stddev=lstm_init_stddev),
+        num_proj=initial_size)
+    part_pred = tf.contrib.layers.layer_norm(part_pred)
+    pred = part_pred
+
+    for pred_layer_num in range(1, pred_depth, 2):
+      part_pred, lstm_states[pred_layer_num] = common_video.lstm_cell(
+          pred,
+          lstm_states[pred_layer_num],
+          initial_size,
+          use_peepholes=True,
+          initializer=tf.truncated_normal_initializer(stddev=lstm_init_stddev),
+          num_proj=initial_size)
+      pred += part_pred
+
+      part_pred, lstm_states[pred_layer_num + 1] = common_video.lstm_cell(
+          tf.concat([pred, pre_pred], 1),
+          lstm_states[pred_layer_num + 1],
+          initial_size,
+          use_peepholes=True,
+          initializer=tf.truncated_normal_initializer(stddev=lstm_init_stddev),
+          num_proj=initial_size)
+      part_pred = tf.contrib.layers.layer_norm(part_pred)
+      pred += part_pred
+
+    pred = tf.layers.dense(
+        pred,
+        enc_final_size,
+        kernel_initializer=tf.truncated_normal_initializer(stddev=init_stddev))
+
+    if hparams.enc_pred_use_l2norm:
+      pred = tf.nn.l2_normalize(pred, 1)
+
+    return pred
+
+
+def construct_model(images,
+                    actions=None,
+                    context_frames=2,
+                    hparams=None,
+                    is_training=True):
+  """Constructs the tensorflow graph of the hierarchical model."""
+
+  pred_depth = 20
+
+  enc_out_all, pred_out_all, van_out_all, van_on_enc_all = [], [], [], []
+
+  lstm_states = [None] * (pred_depth + 2)
+
+  enc_out = encoder_vgg(
+      images[0], hparams.enc_size, False, scope_prefix='timestep/',
+      hparams=hparams, is_training=is_training)
+  enc_out = tf.identity(enc_out, 'enc_out')
+  enc_out_all.append(enc_out)
+
+  num_timesteps = len(actions) - 1
+  sum_freq = int(num_timesteps / 4 + 1)
+
+  reuse = False
+  for timestep, action in zip(range(len(actions) - 1), actions[:-1]):
+    done_warm_start = timestep > context_frames - 1
+
+    with tf.variable_scope('timestep', reuse=reuse):
+      if done_warm_start:
+        pred_input = pred_out_all[-1]
+      else:
+        pred_input = enc_out_all[-1]
+      pred_out = predictor(
+          pred_input, action, lstm_states, pred_depth, False, hparams=hparams)
+      pred_out = tf.identity(pred_out, 'pred_out')
+      if timestep % sum_freq == 0:  # and not hparams.use_tpu:
+        tf.summary.histogram('pred_out', pred_out)
+      pred_out_all.append(pred_out)
+
+      if timestep % sum_freq == 0:  # and not hparams.use_tpu:
+        tf.summary.histogram('lstm_state', lstm_states[0])
+      van_out, _, _ = van(
+          enc_out_all[0],
+          images[0],
+          pred_out,
+          images[timestep + 1],
+          tf.AUTO_REUSE,
+          hparams=hparams)
+      van_out = tf.identity(van_out, 'van_out')
+      van_out_all.append(van_out)
+
+      enc_out = encoder_vgg(
+          images[timestep + 1], hparams.enc_size, True, hparams=hparams,
+          is_training=is_training)
+      enc_out = tf.identity(enc_out, 'enc_out')
+      if timestep % sum_freq == 0:  # and not hparams.use_tpu:
+        tf.summary.histogram('enc_out', enc_out)
+      enc_out_all.append(enc_out)
+
+      van_input = images[0]
+      enc_noise = tf.zeros_like(enc_out)
+      if timestep % sum_freq == 0:  # and not hparams.use_tpu:
+        tf.summary.histogram('enc_noise', enc_noise)
+      van_on_enc, _, _ = van(
+          enc_out_all[0],
+          van_input,
+          enc_out + enc_noise,
+          images[timestep + 1],
+          tf.AUTO_REUSE,
+          hparams=hparams)
+      van_on_enc = tf.identity(van_on_enc, 'van_on_enc')
+      van_on_enc_all.append(van_on_enc)
+
+      reuse = True
+
+  return enc_out_all, pred_out_all, van_out_all, van_on_enc_all
+
+
+def peak_signal_to_noise_ratio(true, pred):
+  """Image quality metric based on maximal signal power vs. power of the noise.
+
+  Args:
+    true: the ground truth image.
+    pred: the predicted image.
+  Returns:
+    peak signal to noise ratio (PSNR)
+  """
+  return 10.0 * tf.log(1.0 / mean_squared_error(true, pred)) / tf.log(10.0)
+
+
+def mean_squared_error(true, pred):
+  """L2 distance between tensors true and pred.
+
+  Args:
+    true: the ground truth image.
+    pred: the predicted image.
+  Returns:
+    mean squared error between ground truth and predicted image.
+  """
+  result = tf.reduce_sum(tf.square(true - pred)) / tf.to_float(tf.size(pred))
+  return result
+
+
+def l1_error(true, pred):
+  """L1 distance between tensors true and pred."""
+  return tf.reduce_sum(tf.abs(true - pred)) / tf.to_float(tf.size(pred))
+
+
+def calc_loss_psnr(gen_images, images, name, hparams=None, use_l1_loss=False):
+  """Calculates loss and psnr for predictions over multiple timesteps."""
+  del hparams
+  with tf.name_scope(name):
+    loss, error, psnr_all = 0.0, 0.0, 0.0
+    for _, x, gx in zip(range(len(gen_images)), images, gen_images):
+      recon_cost = mean_squared_error(x, gx)
+      if use_l1_loss:
+        recon_cost = l1_error(x, gx)
+
+      error_i = l1_error(x, gx)
+      psnr_i = peak_signal_to_noise_ratio(x, gx)
+      psnr_all += psnr_i
+      error += error_i
+      loss += recon_cost
+
+    psnr_all /= tf.to_float(len(gen_images))
+    loss /= tf.to_float(len(gen_images))
+    error /= tf.to_float(len(gen_images))
+
+    # if not hparams.use_tpu:
+    tf.summary.scalar('psnr_all', psnr_all)
+    tf.summary.scalar('loss', loss)
+
+    return loss, psnr_all
+
+
+@registry.register_model
+class NextFrameEpva(sv2p.NextFrameSv2p):
+  """Hierarchical Long-term Video Prediction without Supervision"""
+
+  def body(self, features):
+    hparams = self.hparams
+    input_shape = common_layers.shape_list(features['inputs'])
+    batch_size, _, frame_width, frame_height, frame_channels = input_shape  # pylint: disable=unused-variable
+
+    # Swap time and batch axes.
+    input_frames = common_video.swap_time_and_batch_axes(
+        tf.to_float(features['inputs']))
+    target_frames = common_video.swap_time_and_batch_axes(features['targets'])
+
+    # Get actions if exist otherwise use zeros
+    input_actions = self.get_input_if_exists(
+        features, 'input_action', batch_size, hparams.video_num_input_frames)
+    target_actions = self.get_input_if_exists(
+        features, 'target_action', batch_size, hparams.video_num_target_frames)
+
+    # Get rewards if exist otherwise use zeros
+    # TODO(blazej) enable rewards.
+    # input_rewards = self.get_input_if_exists(
+    #     features, 'input_reward', batch_size, hparams.video_num_input_frames)
+    # target_rewards = self.get_input_if_exists(
+    #     features, 'target_reward', batch_size,hparams.video_num_target_frames)
+    # all_rewards = tf.concat([input_rewards, target_rewards], axis=0)
+
+    all_actions = tf.concat([input_actions, target_actions], axis=0)
+    all_frames = tf.concat([input_frames, target_frames], axis=0)
+
+    all_frames = tf.unstack(all_frames, axis=0)
+    all_actions = tf.unstack(all_actions, axis=0)
+    all_actions = [tf.squeeze(a, 1) for a in all_actions]
+
+    # TODO(blazej) - most likely this downsize is too strong.
+    all_frames = [
+        tf.image.resize_images(
+            image, (IMG_HEIGHT, IMG_WIDTH),
+            method=tf.image.ResizeMethod.BICUBIC)
+        for image in all_frames
+    ]
+
+    enc_out_all, pred_out_all, _, van_on_enc_all = construct_model(
+        all_frames,
+        all_actions,
+        context_frames=hparams.context_frames,
+        hparams=hparams,
+        is_training=self.is_training)
+
+    enc_pred_loss, _ = calc_loss_psnr(
+        enc_out_all[1:],
+        pred_out_all,
+        'enc_pred_loss',
+        hparams=hparams,
+        use_l1_loss=hparams.enc_pred_use_l1_loss)
+
+    van_on_enc_loss, _ = calc_loss_psnr(
+        van_on_enc_all,
+        all_frames[1:],
+        'van_on_enc_loss',
+        hparams=hparams)
+
+    enc_pred_loss_scale_delay = max(hparams.enc_pred_loss_scale_delay, 1)
+    enc_pred_loss_scale = tf.nn.sigmoid(
+        (tf.to_float(tf.train.get_or_create_global_step()
+                    ) - enc_pred_loss_scale_delay) /
+        (enc_pred_loss_scale_delay * .1)) * hparams.enc_pred_loss_scale
+    tf.summary.scalar('enc_pred_loss_scale', enc_pred_loss_scale)
+    epva_loss = enc_pred_loss * enc_pred_loss_scale + van_on_enc_loss
+    tf.summary.scalar('epva_loss', epva_loss)
+
+    predictions = tf.stack(van_on_enc_all)
+
+    # TODO(mbz): clean this up!
+    def fix_video_dims_and_concat_on_x_axis(x):
+      x = tf.transpose(x, [1, 3, 4, 0, 2])
+      x = tf.reshape(x, [batch_size, frame_height, frame_channels, -1])
+      x = tf.transpose(x, [0, 3, 1, 2])
+      return x
+
+    frames_gd = fix_video_dims_and_concat_on_x_axis(target_frames)
+    frames_pd = fix_video_dims_and_concat_on_x_axis(predictions)
+    side_by_side_video = tf.concat([frames_gd, frames_pd], axis=1)
+    tf.summary.image('full_video', side_by_side_video)
+
+    predictions = common_video.swap_time_and_batch_axes(predictions)
+    predictions = tf.slice(predictions,
+                           [0, hparams.video_num_target_frames-1, 0, 0, 0],
+                           [-1]*5)
+
+    return predictions, {'extra': epva_loss}
+
diff --git a/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
new file mode 100644
index 000000000..2260b4c52
--- /dev/null
+++ b/tensor2tensor/models/video/epva_params.py
@@ -0,0 +1,49 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Param sets for EPVA model."""
+
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.models.video import basic_deterministic_params
+from tensor2tensor.utils import registry
+
+
+@registry.register_hparams
+def next_frame_epva():
+  """EPVA hparams."""
+  hparams = basic_deterministic_params.next_frame_basic_deterministic()
+  hparams.video_num_input_frames = 4
+  hparams.video_num_target_frames = 4
+  hparams.target_modality = "video:l2raw"
+  hparams.input_modalities = "inputs:video:l2raw"
+  hparams.learning_rate_constant = 1e-05
+  hparams.batch_size = 2
+  hparams.clip_grad_norm = 0.01
+  hparams.add_hparam("context_frames", 5)
+  hparams.add_hparam("enc_learning_rate", 1e-5)
+  hparams.add_hparam("enc_pred_loss_scale", 0.1)
+  hparams.add_hparam("enc_pred_loss_scale_delay", 6e5)
+  hparams.add_hparam("enc_size", 64)
+  hparams.add_hparam("enc_keep_prob", .65)
+  hparams.add_hparam("enc_pred_use_l1_loss", False)
+  hparams.add_hparam("enc_pred_use_l2norm", False)
+  hparams.add_hparam("van_learning_rate", 3e-5)
+  hparams.add_hparam("van_keep_prob", .9)
+  hparams.add_hparam("sequence_length ", 64)
+  hparams.add_hparam("skip_num", 2)
+  hparams.add_hparam("pred_noise_std", 0)
+  hparams.add_hparam("lstm_state_noise_stddev", 0)
+  return hparams

From ab42aab4762d2d47029a1794e6c3eb5f70ce94dc Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 12 Sep 2018 13:51:19 -0700
Subject: [PATCH 0828/2720] Flag no longer exists

PiperOrigin-RevId: 212697323
---
 tensor2tensor/utils/cloud_mlengine.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 8a9509467..8af62c112 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -333,7 +333,6 @@ def configure_usr_dir(job_spec, usr_tar):
 
 def validate_flags():
   """Validates flags are set to acceptable values for CloudML Engine runs."""
-  assert not FLAGS.cloud_tpu
   assert not job_dir()
   assert FLAGS.output_dir.startswith("gs://")
   assert FLAGS.data_dir.startswith("gs://")

From 360532576eb4c2cd94caa326be592b9c1b4468f9 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 12 Sep 2018 14:06:08 -0700
Subject: [PATCH 0829/2720] Update T2TModel docstring to reflect layer usage.

PiperOrigin-RevId: 212699928
---
 tensor2tensor/utils/t2t_model.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 0ea74540f..42641ed3e 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -61,8 +61,11 @@ class T2TModel(base.Layer):
 
   1. Estimator: The method `make_estimator_model_fn` builds a `model_fn` for
      the tf.Estimator workflow of training, evaluation, and prediction.
-     Its core computation comes from the method `call`, which proceeds to call
-     the following methods:
+     It performs the method `call`, which performs the core computation,
+     followed by `estimator_spec_train`, `estimator_spec_eval`, or
+     `estimator_spec_predict` depending on the tf.Estimator mode.
+  2. Layer: The method `call` enables `T2TModel` to be used a callable by
+     itself. It calls the following methods:
 
      * `bottom`, which transforms features according to `problem_hparams`' input
        and target `Modality`s;
@@ -73,8 +76,6 @@ class T2TModel(base.Layer):
        the final logits;
      * `loss`, which takes the logits, forms any missing training loss, and sums
        all loss terms.
-  2. Layer: The method `call` enables `T2TModel` to be used a callable by
-     itself. For example, it can be composed with any other Keras layer.
   3. Inference: The method `infer` enables `T2TModel` to make sequence
      predictions by itself.
 

From 10efdaead8c151ece6f70bebe4797fc0c2379c03 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 12 Sep 2018 15:36:20 -0700
Subject: [PATCH 0830/2720] switching to internal loss by default.

PiperOrigin-RevId: 212716273
---
 tensor2tensor/models/video/sv2p.py                | 12 +++++++-----
 tensor2tensor/models/video/sv2p_params.py         | 10 +++++++++-
 tensor2tensor/models/video/sv2p_test.py           |  8 ++++++++
 tensor2tensor/rl/trainer_model_based.py           | 11 ++++++++++-
 tensor2tensor/rl/trainer_model_based_sv2p_test.py |  6 ++++++
 5 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 9275fb30d..57626250f 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -553,10 +553,6 @@ def body(self, features):
     predictions = common_video.swap_time_and_batch_axes(predictions)
     reward_pred = common_video.swap_time_and_batch_axes(reward_pred)
 
-    return_targets = predictions
-    if hparams.reward_prediction:
-      return_targets = {"targets": predictions, "target_reward": reward_pred}
-
     if hparams.internal_loss:
       recon_loss = tf.losses.mean_squared_error(all_frames[1:], gen_images)
       rew_loss = 0.0
@@ -565,8 +561,14 @@ def body(self, features):
         tf.summary.scalar("loss/reward", rew_loss)
       tf.summary.scalar("loss/recon", recon_loss)
       tf.summary.scalar("loss/kl", extra_loss)
-
       extra_loss = {"training": recon_loss + rew_loss + extra_loss}
+      # also expand the last dimension of prediction since
+      # we all the modalities will be bypassed.
+      predictions = tf.expand_dims(predictions, axis=-1)
+
+    return_targets = predictions
+    if hparams.reward_prediction:
+      return_targets = {"targets": predictions, "target_reward": reward_pred}
 
     return return_targets, extra_loss
 
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 6b4ada8b6..deb96e662 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -47,7 +47,7 @@ def next_frame_sv2p():
   hparams.add_hparam("scheduled_sampling_decay_steps", 10000)
   hparams.add_hparam("scheduled_sampling_k", 900.0)
   hparams.add_hparam("upsample_method", "conv2d_transpose")
-  hparams.add_hparam("internal_loss", False)
+  hparams.add_hparam("internal_loss", True)
   return hparams
 
 
@@ -81,6 +81,14 @@ def next_frame_sv2p_tiny():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_sv2p_tiny_external():
+  """Tiny SV2P model with external loss."""
+  hparams = next_frame_sv2p_tiny()
+  hparams.internal_loss = False
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_sv2p_cutoff():
   """SV2P model with additional cutoff in L2 loss for environments like pong."""
diff --git a/tensor2tensor/models/video/sv2p_test.py b/tensor2tensor/models/video/sv2p_test.py
index 42c253792..7a59dc62e 100644
--- a/tensor2tensor/models/video/sv2p_test.py
+++ b/tensor2tensor/models/video/sv2p_test.py
@@ -45,6 +45,14 @@ def testSv2pWithActionsAndRewards(self):
         sv2p.NextFrameSv2p,
         1)
 
+  def testSv2pWithActionsAndRewardsExternalLoss(self):
+    hp = sv2p_params.next_frame_sv2p()
+    hp.internal_loss = False
+    self.TestWithActionAndRewards(
+        hp,
+        sv2p.NextFrameSv2p,
+        1)
+
   def testSv2pTwoFrames(self):
     self.TestOnVariousInputOutputSizes(
         sv2p_params.next_frame_sv2p(),
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index dd5c7ecc0..8e5f9b145 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -731,13 +731,22 @@ def rl_modelrl_tiny_stochastic():
 
 @registry.register_hparams
 def rl_modelrl_tiny_sv2p():
-  """Tiny setting with a sv2p model."""
+  """Tiny setting with a tiny sv2p model."""
   hparams = rl_modelrl_tiny()
   hparams.generative_model = "next_frame_sv2p"
   hparams.generative_model_params = "next_frame_sv2p_tiny"
   return hparams
 
 
+@registry.register_hparams
+def rl_modelrl_tiny_sv2p_external():
+  """Tiny setting with a tiny sv2p model and external loss."""
+  hparams = rl_modelrl_tiny()
+  hparams.generative_model = "next_frame_sv2p"
+  hparams.generative_model_params = "next_frame_sv2p_tiny_external"
+  return hparams
+
+
 @registry.register_hparams
 def rl_modelrl_l1_base():
   """Parameter set with L1 loss."""
diff --git a/tensor2tensor/rl/trainer_model_based_sv2p_test.py b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
index 1b5f76a61..361756f64 100644
--- a/tensor2tensor/rl/trainer_model_based_sv2p_test.py
+++ b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
@@ -32,6 +32,12 @@ def test_sv2p(self):
     FLAGS.schedule = "train"
     trainer_model_based.main(None)
 
+  def test_sv2p_external_loss(self):
+    FLAGS.output_dir = tf.test.get_temp_dir() + "_external"
+    FLAGS.loop_hparams_set = "rl_modelrl_tiny_sv2p_external"
+    FLAGS.schedule = "train"
+    trainer_model_based.main(None)
+
 
 if __name__ == "__main__":
   tf.test.main()

From e7baeaebb34d390b39d8700bf5f3b3f4359744ee Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 12 Sep 2018 19:10:18 -0700
Subject: [PATCH 0831/2720] INTERNAL

PiperOrigin-RevId: 212744025
---
 tensor2tensor/rl/datagen_with_agent.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/rl/datagen_with_agent.py b/tensor2tensor/rl/datagen_with_agent.py
index 1cdbc9d5d..a3912dcc2 100644
--- a/tensor2tensor/rl/datagen_with_agent.py
+++ b/tensor2tensor/rl/datagen_with_agent.py
@@ -38,7 +38,6 @@
 
 
 def main(_):
-
   tf.gfile.MakeDirs(FLAGS.data_dir)
   tf.gfile.MakeDirs(FLAGS.tmp_dir)
 

From ef90dbc0f6304acc15aa81b8f108c51597d9df67 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 13 Sep 2018 08:42:44 -0700
Subject: [PATCH 0832/2720] Fix typo: integerg --> integer.

PiperOrigin-RevId: 212819030
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index b84b1b3e8..4266d9787 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -3333,7 +3333,7 @@ def multihead_attention(query_antecedent,
     bias: bias Tensor (see attention_bias())
     total_key_depth: an integer
     total_value_depth: an integer
-    output_depth: an integerg
+    output_depth: an integer
     num_heads: an integer dividing total_key_depth and total_value_depth
     dropout_rate: a floating point number
     attention_type: a string, either "dot_product", "dot_product_relative",

From 3a26ad7b5c24e025a92e5ebaa88c19e4b63a33a3 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 13 Sep 2018 11:21:58 -0700
Subject: [PATCH 0833/2720] Correct trainer_lib_test setup and lower its BUILD
 deps.

PiperOrigin-RevId: 212845915
---
 tensor2tensor/utils/trainer_lib_test.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index aa9d57976..49e1a53f9 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -18,9 +18,9 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor import models  # pylint: disable=unused-import
 from tensor2tensor.data_generators import algorithmic
 from tensor2tensor.data_generators import problem as problem_lib
+from tensor2tensor.models import transformer  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 import tensorflow as tf
@@ -78,16 +78,16 @@ def testModel(self):
       self.assertEqual(loss_val.shape, tuple())
 
   def testMultipleTargetModalities(self):
-    # HParams
+    # Use existing hparams and override target modality.
     hparams = trainer_lib.create_hparams(
         "transformer_tiny", data_dir=algorithmic.TinyAlgo.data_dir,
         problem_name="tiny_algo")
-    tm = hparams.problem.get_hparams().target_modality
-    hparams.problem.get_hparams().target_modality = {
-        "targets": tm,
-        "A": tm,
-        "B": tm
+    hparams.problem_hparams.target_modality = {
+        "targets": hparams.problem_hparams.target_modality,
+        "A": hparams.problem_hparams.target_modality,
+        "B": hparams.problem_hparams.target_modality
     }
+    hparams.problem._hparams = hparams.problem_hparams
 
     # Dataset
     problem = hparams.problem

From d9a1e5bdc7b90159bb0422abc5ba795ac78d2eca Mon Sep 17 00:00:00 2001
From: Marcin Michalski <michalski@google.com>
Date: Thu, 13 Sep 2018 11:50:37 -0700
Subject: [PATCH 0834/2720] Model exporting: Fix the list of input features and
 drop unsupported collections.

PiperOrigin-RevId: 212851333
---
 tensor2tensor/serving/export.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 6adba8e82..b5600d3ef 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -99,13 +99,19 @@ def hub_module_fn():
         hparams,
         decode_hparams=decoding.decode_hparams(FLAGS.decode_hparams))
     features = problem.serving_input_fn(hparams).features
+
+    # we must do a copy of the features, as the model_fn can add additional
+    # entries there (like hyperparameter settings etc).
+    original_features = features.copy()
     spec = model_fn(features, labels=None, mode=tf.estimator.ModeKeys.PREDICT)
 
-    # Currently only supports a single input and single output.
     hub.add_signature(
-        inputs=features, outputs=spec.export_outputs["serving_default"].outputs)
+        inputs=original_features,
+        outputs=spec.export_outputs["serving_default"].outputs)
 
-  module_spec = hub.create_module_spec(hub_module_fn)
+  # TFHub doesn't support LOSSES collections.
+  module_spec = hub.create_module_spec(
+      hub_module_fn, drop_collections=[tf.GraphKeys.LOSSES])
   # Loads the weights from the checkpoint using the model above
   # and saves it in the export_path.
   export_module_spec_with_checkpoint(

From a57740a1753df2640874baec5e8b26ac9d3bb400 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Thu, 13 Sep 2018 12:13:15 -0700
Subject: [PATCH 0835/2720] more plint

---
 tensor2tensor/data_generators/gym_problems.py | 6 ++++--
 tensor2tensor/layers/common_video.py          | 3 +--
 tensor2tensor/rl/collect.py                   | 2 +-
 tensor2tensor/rl/ppo.py                       | 2 +-
 4 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 773ccd724..4f84c3191 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -147,11 +147,13 @@ def random_skip(self):
 
   def _get_data(self):
     if self._use_dumper_data:
-      file_path = os.path.join(self._dumper_path, "frame_{}.npz".format(self._dumper_data_index))
+      file_path = os.path.join(self._dumper_path,
+                               "frame_{}.npz".format(self._dumper_data_index))
       print(file_path)
       data = np.load(file_path)
       self._dumper_data_index += 1
-      return data["observ"][0, ...], data["reward"][0], data["done"][0], data["action"][0]
+      return data["observ"][0, ...], data["reward"][0], \
+             data["done"][0], data["action"][0]
     else:
       if self.memory is None or self.memory_index >= self._internal_memory_size:
         self.memory = self._session.run(self.collect_memory)
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 56cfb5d97..469156f23 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -535,10 +535,9 @@ def __init_ffmpeg(self, image_shape):
         "-qscale", "0",
         "-"
     ]
-    # self.proc = Popen(self.cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
+    self.proc = Popen(self.cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
 
   def write(self, frame):
-    return
     if self.proc is None:
       self.__init_ffmpeg(frame.shape)
     self.proc.stdin.write(frame.tostring())
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index b712953c6..bde510ea4 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -275,7 +275,7 @@ def stop_condition(i, _, resets):
         mem = tf.reshape(mem, shape=new_shape)
         mem = tf.transpose(mem, perm=perm,
                            name="collect_memory_%d_%s"
-                                % (new_epoch_length, name))
+                           % (new_epoch_length, name))
         new_memory.append(mem)
       memory = new_memory
 
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index ca8cc1a58..4f2f1c655 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -103,7 +103,7 @@ def define_ppo_epoch(memory, hparams):
   if hasattr(hparams, "effective_num_agents"):
     number_of_batches *= hparams.num_agents
     number_of_batches /= hparams.effective_num_agents
-    
+
   dataset = tf.data.Dataset.from_tensor_slices(
       (observation, action, discounted_reward, advantage_normalized, old_pdf))
   dataset = dataset.shuffle(buffer_size=hparams.epoch_length,

From cee7869d15b58436825ef37f351728d5316f07f4 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 13 Sep 2018 12:26:55 -0700
Subject: [PATCH 0836/2720] Expand docstring of T2TModel.{top,bottom}.

PiperOrigin-RevId: 212857847
---
 tensor2tensor/utils/t2t_model.py | 28 +++++++++++++++++++++++++---
 1 file changed, 25 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 42641ed3e..063f24f73 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -323,7 +323,16 @@ def model_fn(self, features):
       return logits, losses
 
   def bottom(self, features):
-    """Transform features to feed into body."""
+    """Transforms features to feed into body.
+
+    Args:
+      features: dict of str to Tensor. Typically it is the preprocessed data
+        batch after Problem's preprocess_example().
+
+    Returns:
+      transformed_features: dict of same key-value pairs as features. The value
+        Tensors are newly transformed.
+    """
     if not self._problem_hparams:
       log_warn("Without a Problem, T2TModel.bottom is a passthrough.")
       return features
@@ -376,7 +385,7 @@ def bottom(self, features):
     return transformed_features
 
   def body(self, features):
-    """Computes the targets' logits for one shard given transformed inputs.
+    """Computes the targets' pre-logit activations given transformed inputs.
 
     Most `T2TModel` subclasses will override this method.
 
@@ -432,7 +441,20 @@ def _top_single(self, body_output, target_modality, features):
     return logits
 
   def top(self, body_output, features):
-    """Returns `logits` given body output and features."""
+    """Computes logits given body output and features.
+
+    Args:
+      body_output: dict of str to Tensor, comprising one key-value pair for each
+        target. Each value denotes the target's pre-logit activations.
+        Alternatively, it may be a single Tensor denoting the pre-logits for
+        that target.
+      features: dict of str to Tensor. Typically it is the preprocessed data
+        batch after Problem's preprocess_example().
+
+    Returns:
+      logits: dict of str to Tensor, denoting each logits for each target; or
+        a single Tensor denoting the logits for that target.
+    """
     if isinstance(body_output, dict):
       if self._problem_hparams:
         target_modality = self._problem_hparams.target_modality

From 0a7c6f4942ba0dd2f6731b0f0b7389762165c2ed Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 13 Sep 2018 12:33:21 -0700
Subject: [PATCH 0837/2720] Pass model hparams into test_problem_hparams.

PiperOrigin-RevId: 212858905
---
 tensor2tensor/data_generators/problem_hparams.py            | 6 ++++--
 tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py | 4 +++-
 tensor2tensor/mesh_tensorflow/mtf_transformer_test.py       | 4 +++-
 tensor2tensor/models/bytenet_test.py                        | 4 +++-
 tensor2tensor/models/image_transformer_2d_test.py           | 4 +++-
 tensor2tensor/models/image_transformer_test.py              | 4 +++-
 tensor2tensor/models/lstm_test.py                           | 4 +++-
 tensor2tensor/models/neural_gpu_test.py                     | 3 ++-
 tensor2tensor/models/research/transformer_aux_test.py       | 4 +++-
 tensor2tensor/models/research/transformer_revnet_test.py    | 4 +++-
 tensor2tensor/models/research/transformer_vae_test.py       | 4 +++-
 tensor2tensor/models/research/universal_transformer_test.py | 4 +++-
 tensor2tensor/models/research/vqa_attention_test.py         | 4 +++-
 tensor2tensor/models/resnet_test.py                         | 4 +++-
 tensor2tensor/models/transformer_test.py                    | 4 +++-
 tensor2tensor/models/xception_test.py                       | 4 +++-
 16 files changed, 48 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/data_generators/problem_hparams.py b/tensor2tensor/data_generators/problem_hparams.py
index a6ffd7b0c..51459f1f5 100644
--- a/tensor2tensor/data_generators/problem_hparams.py
+++ b/tensor2tensor/data_generators/problem_hparams.py
@@ -198,7 +198,9 @@ def hparams(self, defaults, model_hparams):
     hp.target_modality = (registry.Modalities.SYMBOL, self.target_vocab_size)
 
 
-def test_problem_hparams(input_vocab_size=None, target_vocab_size=None):
+def test_problem_hparams(input_vocab_size=None,
+                         target_vocab_size=None,
+                         model_hparams=None):
   """Problem hparams for testing model bodies."""
   p = TestProblem(input_vocab_size, target_vocab_size)
-  return p.get_hparams()
+  return p.get_hparams(model_hparams)
diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py b/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
index 53197293c..53e8cff7a 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
@@ -43,7 +43,9 @@ def get_model(hparams=None,
   hparams.img_len = IMG_LENGTH
   hparams.num_channels = 1
 
-  p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE, VOCAB_SIZE)
+  p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE,
+                                                   VOCAB_SIZE,
+                                                   hparams)
   p_hparams.input_modality = {}
   hparams.problem_hparams = p_hparams
 
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py b/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py
index 68f3a6741..2d0695a5e 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py
@@ -41,7 +41,9 @@ def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
   hparams.max_length = INPUT_LENGTH
   hparams.batch_size = BATCH_SIZE
 
-  p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE, VOCAB_SIZE)
+  p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE,
+                                                   VOCAB_SIZE,
+                                                   hparams)
   if not has_input:
     p_hparams.input_modality = {}
   hparams.problem_hparams = p_hparams
diff --git a/tensor2tensor/models/bytenet_test.py b/tensor2tensor/models/bytenet_test.py
index 67fff4a50..1021ca432 100644
--- a/tensor2tensor/models/bytenet_test.py
+++ b/tensor2tensor/models/bytenet_test.py
@@ -32,7 +32,9 @@ def testByteNet(self):
     x = np.random.random_integers(1, high=vocab_size - 1, size=(3, 5, 1, 1))
     y = np.random.random_integers(1, high=vocab_size - 1, size=(3, 6, 1, 1))
     hparams = bytenet.bytenet_base()
-    p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
+    p_hparams = problem_hparams.test_problem_hparams(vocab_size,
+                                                     vocab_size,
+                                                     hparams)
     with self.test_session() as session:
       features = {
           "inputs": tf.constant(x, dtype=tf.int32),
diff --git a/tensor2tensor/models/image_transformer_2d_test.py b/tensor2tensor/models/image_transformer_2d_test.py
index b5dafdd99..3f52df180 100644
--- a/tensor2tensor/models/image_transformer_2d_test.py
+++ b/tensor2tensor/models/image_transformer_2d_test.py
@@ -59,7 +59,9 @@ def _test_imagetransformer_2d(self, net):
     size = 7
     vocab_size = 256
     hparams = image_transformer_2d.imagetransformer2d_tiny()
-    p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
+    p_hparams = problem_hparams.test_problem_hparams(vocab_size,
+                                                     vocab_size,
+                                                     hparams)
     inputs = -1 + np.random.random_integers(
         vocab_size, size=(batch_size, 1, 1, 1))
     targets = -1 + np.random.random_integers(
diff --git a/tensor2tensor/models/image_transformer_test.py b/tensor2tensor/models/image_transformer_test.py
index 8f4911f0d..8257a0fe9 100644
--- a/tensor2tensor/models/image_transformer_test.py
+++ b/tensor2tensor/models/image_transformer_test.py
@@ -42,7 +42,9 @@ def testImagetransformer(self, net, hparams):
     batch_size = 3
     size = 7
     vocab_size = 256
-    p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
+    p_hparams = problem_hparams.test_problem_hparams(vocab_size,
+                                                     vocab_size,
+                                                     hparams)
     inputs = -1 + np.random.random_integers(
         vocab_size, size=(batch_size, 1, 1, 1))
     targets = -1 + np.random.random_integers(
diff --git a/tensor2tensor/models/lstm_test.py b/tensor2tensor/models/lstm_test.py
index 4b2fcd640..f715c3672 100644
--- a/tensor2tensor/models/lstm_test.py
+++ b/tensor2tensor/models/lstm_test.py
@@ -32,7 +32,9 @@ def testLSTMSeq2Seq(self):
     x = np.random.random_integers(1, high=vocab_size - 1, size=(3, 5, 1, 1))
     y = np.random.random_integers(1, high=vocab_size - 1, size=(3, 6, 1, 1))
     hparams = lstm.lstm_seq2seq()
-    p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
+    p_hparams = problem_hparams.test_problem_hparams(vocab_size,
+                                                     vocab_size,
+                                                     hparams)
     with self.test_session() as session:
       features = {
           "inputs": tf.constant(x, dtype=tf.int32),
diff --git a/tensor2tensor/models/neural_gpu_test.py b/tensor2tensor/models/neural_gpu_test.py
index 9a2dab9f1..857e0e5ea 100644
--- a/tensor2tensor/models/neural_gpu_test.py
+++ b/tensor2tensor/models/neural_gpu_test.py
@@ -36,7 +36,8 @@ def testNeuralGPU(self):
     input_vocab_size = 9
     target_vocab_size = 11
     p_hparams = problem_hparams.test_problem_hparams(input_vocab_size,
-                                                     target_vocab_size)
+                                                     target_vocab_size,
+                                                     hparams)
     inputs = -1 + np.random.random_integers(
         input_vocab_size, size=(batch_size, input_length, 1, 1))
     targets = -1 + np.random.random_integers(
diff --git a/tensor2tensor/models/research/transformer_aux_test.py b/tensor2tensor/models/research/transformer_aux_test.py
index 3b3531b38..32e71d838 100644
--- a/tensor2tensor/models/research/transformer_aux_test.py
+++ b/tensor2tensor/models/research/transformer_aux_test.py
@@ -85,7 +85,9 @@ def test_transformer_aux_body(self):
     vocab_size = 9
     hparams = transformer_aux.transformer_aux_tiny()
     hparams.shift_values = "-5,1,2,3"
-    p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
+    p_hparams = problem_hparams.test_problem_hparams(vocab_size,
+                                                     vocab_size,
+                                                     hparams)
     hparams.problem_hparams = p_hparams
     inputs = -1 + np.random.random_integers(
         vocab_size, size=(batch_size, input_length, 1, 1))
diff --git a/tensor2tensor/models/research/transformer_revnet_test.py b/tensor2tensor/models/research/transformer_revnet_test.py
index 38571aefa..76e41c0d3 100644
--- a/tensor2tensor/models/research/transformer_revnet_test.py
+++ b/tensor2tensor/models/research/transformer_revnet_test.py
@@ -42,7 +42,9 @@ def testTransformer(self):
     target_length = 7
     vocab_size = 9
     hparams = transformer_revnet_test()
-    p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
+    p_hparams = problem_hparams.test_problem_hparams(vocab_size,
+                                                     vocab_size,
+                                                     hparams)
     hparams.problem_hparams = p_hparams
     inputs = -1 + np.random.random_integers(
         vocab_size, size=(batch_size, input_length, 1, 1))
diff --git a/tensor2tensor/models/research/transformer_vae_test.py b/tensor2tensor/models/research/transformer_vae_test.py
index 09efd32ec..e67734504 100644
--- a/tensor2tensor/models/research/transformer_vae_test.py
+++ b/tensor2tensor/models/research/transformer_vae_test.py
@@ -32,7 +32,9 @@ def testTransformerAEOnDVQ(self):
     hparams = transformer_vae.transformer_ae_small()
     hparams.bottleneck_kind = "dvq"
     hparams.dp_strength = 0
-    p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
+    p_hparams = problem_hparams.test_problem_hparams(vocab_size,
+                                                     vocab_size,
+                                                     hparams)
     hparams.problem_hparams = p_hparams
     inputs = -1 + np.random.random_integers(
         vocab_size, size=(batch_size, input_length, 1, 1))
diff --git a/tensor2tensor/models/research/universal_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
index da228c2d0..d8c389aac 100644
--- a/tensor2tensor/models/research/universal_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -40,7 +40,9 @@ def get_model(self,
     hparams.num_heads = 1
     hparams.layer_prepostprocess_dropout = 0.0
 
-    p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE, VOCAB_SIZE)
+    p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE,
+                                                     VOCAB_SIZE,
+                                                     hparams)
     if not has_input:
       p_hparams.input_modality = {}
     hparams.problems = [p_hparams]
diff --git a/tensor2tensor/models/research/vqa_attention_test.py b/tensor2tensor/models/research/vqa_attention_test.py
index d3a9e5ecb..4108572c2 100644
--- a/tensor2tensor/models/research/vqa_attention_test.py
+++ b/tensor2tensor/models/research/vqa_attention_test.py
@@ -43,7 +43,9 @@ def testVqaAttentionBaseline(self):
     a = np.random.random_integers(
         0, high=num_classes, size=(batch_size, answer_length, 1, 1))
     hparams = vqa_attention.vqa_attention_base()
-    p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
+    p_hparams = problem_hparams.test_problem_hparams(vocab_size,
+                                                     vocab_size,
+                                                     hparams)
     p_hparams.input_modality["inputs"] = (registry.Modalities.IMAGE, None)
     p_hparams.input_modality["question"] = (registry.Modalities.SYMBOL,
                                             vocab_size)
diff --git a/tensor2tensor/models/resnet_test.py b/tensor2tensor/models/resnet_test.py
index 18baff6e2..76458d58e 100644
--- a/tensor2tensor/models/resnet_test.py
+++ b/tensor2tensor/models/resnet_test.py
@@ -43,7 +43,9 @@ def _test_resnet(self, img_size, output_size):
     y = np.random.random_integers(
         1, high=vocab_size - 1, size=(batch_size, 1, 1, 1))
     hparams = resnet_tiny_cpu()
-    p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
+    p_hparams = problem_hparams.test_problem_hparams(vocab_size,
+                                                     vocab_size,
+                                                     hparams)
     p_hparams.input_modality["inputs"] = (registry.Modalities.IMAGE, None)
     p_hparams.target_modality = (registry.Modalities.CLASS_LABEL, vocab_size)
     with self.test_session() as session:
diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 4450100c9..d569e3b96 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -40,7 +40,9 @@ def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
   hparams.num_heads = 1
   hparams.layer_prepostprocess_dropout = 0.0
 
-  p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE, VOCAB_SIZE)
+  p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE,
+                                                   VOCAB_SIZE,
+                                                   hparams)
   if not has_input:
     p_hparams.input_modality = {}
   hparams.problem_hparams = p_hparams
diff --git a/tensor2tensor/models/xception_test.py b/tensor2tensor/models/xception_test.py
index b57a757b9..ab29a5e0e 100644
--- a/tensor2tensor/models/xception_test.py
+++ b/tensor2tensor/models/xception_test.py
@@ -36,7 +36,9 @@ def _test_xception(self, img_size):
     y = np.random.random_integers(
         1, high=vocab_size - 1, size=(batch_size, 1, 1, 1))
     hparams = xception.xception_tiny()
-    p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
+    p_hparams = problem_hparams.test_problem_hparams(vocab_size,
+                                                     vocab_size,
+                                                     hparams)
     p_hparams.input_modality["inputs"] = (registry.Modalities.IMAGE, None)
     p_hparams.target_modality = (registry.Modalities.CLASS_LABEL, vocab_size)
     with self.test_session() as session:

From 0845791b839b7f02fa5bb122096d7ba95577bc4b Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 13 Sep 2018 13:24:51 -0700
Subject: [PATCH 0838/2720] Internal change

PiperOrigin-RevId: 212867218
---
 tensor2tensor/rl/datagen_with_agent.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/rl/datagen_with_agent.py b/tensor2tensor/rl/datagen_with_agent.py
index a3912dcc2..1cdbc9d5d 100644
--- a/tensor2tensor/rl/datagen_with_agent.py
+++ b/tensor2tensor/rl/datagen_with_agent.py
@@ -38,6 +38,7 @@
 
 
 def main(_):
+
   tf.gfile.MakeDirs(FLAGS.data_dir)
   tf.gfile.MakeDirs(FLAGS.tmp_dir)
 

From fccbc751ca3f2a4d83ab06cd219238e9f8aa33db Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 13 Sep 2018 14:52:50 -0700
Subject: [PATCH 0839/2720] Add model hparams to test_problem_hparams in
 lstm_test.py.

PiperOrigin-RevId: 212883570
---
 tensor2tensor/models/lstm_test.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/lstm_test.py b/tensor2tensor/models/lstm_test.py
index f715c3672..a75e95cf5 100644
--- a/tensor2tensor/models/lstm_test.py
+++ b/tensor2tensor/models/lstm_test.py
@@ -53,7 +53,9 @@ def testLSTMSeq2SeqAttention(self):
     y = np.random.random_integers(1, high=vocab_size - 1, size=(3, 6, 1, 1))
     hparams = lstm.lstm_attention()
 
-    p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
+    p_hparams = problem_hparams.test_problem_hparams(vocab_size,
+                                                     vocab_size,
+                                                     hparams)
     x = tf.constant(x, dtype=tf.int32)
     x = tf.placeholder_with_default(x, shape=[None, None, 1, 1])
 
@@ -74,7 +76,9 @@ def testLSTMSeq2seqBidirectionalEncoder(self):
     x = np.random.random_integers(1, high=vocab_size - 1, size=(3, 5, 1, 1))
     y = np.random.random_integers(1, high=vocab_size - 1, size=(3, 6, 1, 1))
     hparams = lstm.lstm_seq2seq()
-    p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
+    p_hparams = problem_hparams.test_problem_hparams(vocab_size,
+                                                     vocab_size,
+                                                     hparams)
     with self.test_session() as session:
       features = {
           "inputs": tf.constant(x, dtype=tf.int32),

From 04a714db1b7fabb1ccb274da8415e2c1989eb023 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Thu, 13 Sep 2018 18:26:23 -0700
Subject: [PATCH 0840/2720] assert explaining effective_num_agents

---
 tensor2tensor/rl/collect.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index bde510ea4..10184d1cc 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -264,6 +264,9 @@ def stop_condition(i, _, resets):
     new_memory = []
     if hasattr(hparams, "effective_num_agents"):
       effective_num_agents = hparams.effective_num_agents
+      assert hparams.epoch_length % effective_num_agents == 0, \
+        "The rollout of hparams.epoch_length will be distributed amongst" \
+        "effective_num_agents of agents"
       new_epoch_length = int(hparams.epoch_length / effective_num_agents)
       for mem, info in zip(memory, rollout_metadata):
         shape, _, name = info

From 67c3d272e7f0afeb5cf50a521de8912b91684e94 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Thu, 13 Sep 2018 18:59:24 -0700
Subject: [PATCH 0841/2720] calculating correct number of
 real_training_ppo_epochs_num

---
 tensor2tensor/rl/trainer_model_based.py | 38 +++++++++++++++++--------
 1 file changed, 26 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 1e85c6e90..6668ebb4a 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -105,8 +105,11 @@ def generate_real_env_data(problem_name, agent_policy_path, hparams, data_dir,
       "autoencoder_path": autoencoder_path,
   }):
     gym_problem = registry.problem(problem_name)
-    env_steps_per_epoch = (
-        hparams.num_real_env_frames / (hparams.epochs * (1. - 1./11.)))
+    if hparams.gather_ppo_real_env_data:
+      env_steps_per_epoch = int(hparams.num_real_env_frames / hparams.epochs)
+    else:
+      env_steps_per_epoch = (
+          hparams.num_real_env_frames / (hparams.epochs * (1. - 1./11.)))
     gym_problem.settable_num_steps = env_steps_per_epoch
     gym_problem.settable_eval_phase = eval_phase
     gym_problem.generate_data(data_dir, tmp_dir)
@@ -145,7 +148,15 @@ def train_autoencoder(problem_name, data_dir, output_dir, hparams, epoch):
 
 def _ppo_training_epochs(hparams, epoch, is_final_epoch, real_env_training):
   """Helper for PPO restarts."""
-  real_training_ppo_epochs_num = hparams.real_ppo_epochs_num
+  if hparams.gather_ppo_real_env_data:
+    assert hparams.real_ppo_epochs_num is 0, \
+      "Should be put to 0 to enforce better readability"
+    real_training_ppo_epochs_num = \
+      math.ceil(hparams.num_real_env_frames/
+                (hparams.epochs*hparams.real_ppo_epoch_length))
+  else:
+    real_training_ppo_epochs_num = hparams.real_ppo_epochs_num
+
   simulated_training_ppo_epochs_num = hparams.ppo_epochs_num
 
   ppo_training_epochs = (epoch + 1)*simulated_training_ppo_epochs_num + \
@@ -236,12 +247,16 @@ def train_agent_real_env(
 
   environment_spec = copy.copy(gym_problem.environment_spec)
 
-  #TODO(piotrmilos):This should be refactored
-  ppo_data_dumper_counter = 0
-  dumper_path = os.path.join(epoch_data_dir, "dumper")
-  tf.gfile.MakeDirs(dumper_path)
-  dumper_spec = [PyFuncWrapper, {"process_fun": ppo_data_dumper}]
-  environment_spec.wrappers.insert(1, dumper_spec)
+  if hparams.gather_ppo_real_env_data:
+    #TODO(piotrmilos):This should be refactored
+    assert hparams.real_ppo_num_agents == 1, \
+      "It is required to use collect with pyfunc_wrapper"
+
+    ppo_data_dumper_counter = 0
+    dumper_path = os.path.join(epoch_data_dir, "dumper")
+    tf.gfile.MakeDirs(dumper_path)
+    dumper_spec = [PyFuncWrapper, {"process_fun": ppo_data_dumper}]
+    environment_spec.wrappers.insert(1, dumper_spec)
 
 
   ppo_hparams.add_hparam("environment_spec", environment_spec)
@@ -621,7 +636,8 @@ def rl_modelrl_base():
       # should start fresh each time.
       ppo_continue_training=True,
 
-      real_ppo_epochs_num=10,
+      gather_ppo_real_env_data=True,
+      real_ppo_epochs_num=0,
       real_ppo_epoch_length=16*200,
       real_ppo_num_agents=1,
       real_ppo_learning_rate=2e-4,
@@ -642,7 +658,6 @@ def rl_modelrl_base_quick():
   hparams.epochs = 2
   hparams.ppo_epochs_num = 1000
   hparams.ppo_epoch_length = 50
-  hparams.real_ppo_epochs_num = 10
   return hparams
 
 
@@ -736,7 +751,6 @@ def rl_modelrl_tiny():
           ppo_time_limit=5,
           ppo_epoch_length=5,
           ppo_num_agents=2,
-          real_ppo_epochs_num=2,
           real_ppo_epoch_length=36,
           real_ppo_num_agents=1,
           real_ppo_effective_num_agents=2,

From 3e1b32d28abedc07ac1e14fdf40852415ef55c09 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Thu, 13 Sep 2018 19:18:46 -0700
Subject: [PATCH 0842/2720] remove unnecessary print

---
 tensor2tensor/data_generators/gym_problems.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 4f84c3191..b0d5ecc4b 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -149,7 +149,6 @@ def _get_data(self):
     if self._use_dumper_data:
       file_path = os.path.join(self._dumper_path,
                                "frame_{}.npz".format(self._dumper_data_index))
-      print(file_path)
       data = np.load(file_path)
       self._dumper_data_index += 1
       return data["observ"][0, ...], data["reward"][0], \

From 6033e112745443414fb6ab34ce8d31ee738e9b19 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 14 Sep 2018 00:26:10 -0700
Subject: [PATCH 0843/2720] Corrections to autoencoders after recent modality
 change.

PiperOrigin-RevId: 212940799
---
 tensor2tensor/data_generators/video_utils.py  |  6 +-
 tensor2tensor/layers/common_video.py          |  9 ++-
 tensor2tensor/layers/discretization.py        |  7 +-
 tensor2tensor/layers/modalities.py            | 30 +++++---
 tensor2tensor/models/research/autoencoders.py | 36 +++++++--
 .../models/video/basic_deterministic.py       |  6 --
 tensor2tensor/rl/trainer_model_based.py       | 76 +------------------
 7 files changed, 64 insertions(+), 106 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 7bd343faa..d67d51866 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -235,7 +235,6 @@ def serving_input_fn(self, hparams):
         receiver_tensors=video_input_frames)
 
   def preprocess(self, dataset, mode, hparams, interleave=True):
-    del interleave
     def split_on_batch(x):
       """Split x on batch dimension into x[:size, ...] and x[size:, ...]."""
       length = len(x.get_shape())
@@ -335,7 +334,7 @@ def check_integrity_and_batch(*datasets):
     num_frames = (hparams.video_num_input_frames +
                   hparams.video_num_target_frames)
     # We jump by a random position at the beginning to add variety.
-    if self.random_skip:
+    if self.random_skip and interleave:
       random_skip = tf.random_uniform([], maxval=num_frames, dtype=tf.int64)
       preprocessed_dataset = preprocessed_dataset.skip(random_skip)
     if self.use_not_breaking_batching:
@@ -344,7 +343,8 @@ def check_integrity_and_batch(*datasets):
       batch_dataset = preprocessed_dataset.apply(
           tf.contrib.data.batch_and_drop_remainder(num_frames))
     dataset = batch_dataset.map(features_from_batch)
-    dataset = dataset.shuffle(hparams.get("shuffle_buffer_size", 128))
+    if interleave:
+      dataset = dataset.shuffle(hparams.get("shuffle_buffer_size", 128))
     return dataset
 
   def eval_metrics(self):
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 9aceab42d..80fccc30c 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -409,8 +409,16 @@ def gif_summary(name, tensor, max_outputs=3, fps=10, collections=None,
   Returns:
     A scalar `Tensor` of type `string`. The serialized `Summary` protocol
     buffer.
+
+  Raises:
+    ValueError: if the given tensor has the wrong shape.
   """
   tensor = tf.convert_to_tensor(tensor)
+  if len(tensor.get_shape()) != 5:
+    raise ValueError("Assuming videos given as tensors in the format "
+                     "[batch, time, height, width, channels] but got one "
+                     "of shape: %s" % str(tensor.get_shape()))
+  tensor = tf.cast(tensor, tf.uint8)
   if summary_op_util.skip_summary():
     return tf.constant("")
   with summary_op_util.summary_scope(
@@ -623,4 +631,3 @@ def finish_to_files(self, path_template):
     for i, writer in enumerate(self.writers):
       path = path_template.format(i)
       writer.finish_to_file(path)
-
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 8a63b8531..f1aace305 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -1279,8 +1279,11 @@ def gumbel_softmax_discrete_bottleneck(x,
 def tanh_discrete_bottleneck(x, bottleneck_bits, bottleneck_noise,
                              discretize_warmup_steps, mode):
   """Simple discretization through tanh, flip bottleneck_noise many bits."""
-  x = tf.tanh(
-      tf.layers.dense(x, bottleneck_bits, name="tanh_discrete_bottleneck"))
+  x = tf.layers.dense(x, bottleneck_bits, name="tanh_discrete_bottleneck")
+  if mode == tf.estimator.ModeKeys.TRAIN:
+    x += tf.truncated_normal(
+        common_layers.shape_list(x), mean=0.0, stddev=0.2)
+  x = tf.tanh(x)
   d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
   if mode == tf.estimator.ModeKeys.TRAIN:
     noise = tf.random_uniform(common_layers.shape_list(x))
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index e78cd0072..0dfad2e0b 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -21,6 +21,7 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_audio
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
 from tensor2tensor.layers import discretization
 from tensor2tensor.utils import modality
 from tensor2tensor.utils import registry
@@ -289,10 +290,12 @@ def top(self, body_output, _):
   def loss(self, top_out, targets):
     """Compute loss numerator and denominator for one shard of output."""
     logits = top_out
+    cutoff = getattr(self._model_hparams, "video_modality_loss_cutoff", 0.0)
     return common_layers.padded_cross_entropy(
         logits,
         targets,
         self._model_hparams.label_smoothing,
+        cutoff=cutoff,
         weights_fn=self.targets_weights_fn)
 
 
@@ -617,21 +620,24 @@ class VideoModality(modality.Modality):
   """Modality for videos, i.e., time-sequences of frames."""
 
   def bottom(self, x):
-    inputs = x
-    with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
-      common_layers.summarize_video(inputs, "inputs")
-      inputs = common_layers.standardize_images(inputs)
-      return inputs
+    common_video.gif_summary("inputs", x, max_outputs=1)
+    x = common_layers.standardize_images(x)
+    return x
 
   def targets_bottom(self, x):
-    return self.bottom(x)
+    common_video.gif_summary("targets", x, max_outputs=1)
+    x = common_layers.standardize_images(x)
+    return x
 
   def top(self, body_output, targets):
     num_channels = self._model_hparams.problem.num_channels
-    body_output_shape = common_layers.shape_list(body_output)
-    reshape_shape = body_output_shape[:-1] + [
-        num_channels, self.top_dimensionality]
+    shape = common_layers.shape_list(body_output)
+    reshape_shape = shape[:-1] + [num_channels, self.top_dimensionality]
     res = tf.reshape(body_output, reshape_shape)
+    # Calculate argmax so as to have a summary with the produced images.
+    x = tf.argmax(tf.reshape(res, [-1, self.top_dimensionality]), axis=-1)
+    x = tf.reshape(x, shape[:-1] + [num_channels])
+    common_video.gif_summary("results", x, max_outputs=1)
     return res
 
   def loss(self, top_out, targets):
@@ -768,11 +774,11 @@ def convert_rgb_to_real(self, prediction, targets):
     return prediction, targets
 
   def bottom(self, x):
-    common_layers.summarize_video(x, "inputs")
+    common_video.gif_summary("inputs", x)
     return common_layers.convert_rgb_to_real(x)
 
   def targets_bottom(self, x):  # pylint: disable=arguments-differ
-    common_layers.summarize_video(x, "targets_bottom")
+    common_video.gif_summary("targets_bottom", x)
     return common_layers.convert_rgb_to_real(x)
 
   def top(self, body_output, _):
@@ -780,7 +786,7 @@ def top(self, body_output, _):
     if isinstance(body_output, list):
       frames = tf.stack(body_output, axis=1)
     rgb_frames = common_layers.convert_real_to_rgb(frames)
-    common_layers.summarize_video(rgb_frames, "body_output")
+    common_video.gif_summary("body_output", rgb_frames)
     return tf.expand_dims(rgb_frames, axis=-1)
 
   def loss(self, top_out, targets):
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 5580893cb..0c7dea48d 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -33,6 +33,20 @@ def reverse_gradient(x, lr=1.0):
   return -lr * x + tf.stop_gradient((1.0 + lr) * x)
 
 
+def time_to_channels(embedded_video):
+  """Put time dimension on channels in an embedded video."""
+  video_shape = common_layers.shape_list(embedded_video)
+  if len(video_shape) != 5:
+    raise ValueError("Assuming videos given as tensors in the format "
+                     "[batch, time, height, width, channels] but got one "
+                     "of shape: %s" % str(video_shape))
+  transposed = tf.transpose(embedded_video, [0, 2, 3, 1, 4])
+  return tf.reshape(transposed, [
+      video_shape[0], video_shape[2], video_shape[3],
+      video_shape[1] * video_shape[4]
+  ])
+
+
 @registry.register_model
 class AutoencoderBasic(t2t_model.T2TModel):
   """A basic autoencoder, try with image_mnist_rev or image_cifar10_rev."""
@@ -142,7 +156,11 @@ def decoder(self, x, encoder_layers):
   def gumbel_sample(self, reconstr_gan):
     hparams = self.hparams
     is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
-    vocab_size = self._problem_hparams.target_modality.top_dimensionality
+    if isinstance(self._problem_hparams.target_modality, dict):
+      vocab_size = self._problem_hparams.target_modality[
+          "targets"].top_dimensionality
+    else:
+      vocab_size = self._problem_hparams.target_modality.top_dimensionality
     reconstr_gan = tf.nn.log_softmax(reconstr_gan)
     if is_training and hparams.gumbel_temperature > 0.0:
       gumbel_samples = discretization.gumbel_sample(
@@ -163,7 +181,11 @@ def gumbel_sample(self, reconstr_gan):
   def body(self, features):
     hparams = self.hparams
     is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
-    vocab_size = self._problem_hparams.target_modality.top_dimensionality
+    if isinstance(self._problem_hparams.target_modality, dict):
+      vocab_size = self._problem_hparams.target_modality[
+          "targets"].top_dimensionality
+    else:
+      vocab_size = self._problem_hparams.target_modality.top_dimensionality
     encoder_layers = None
     self.is1d = hparams.sample_width == 1
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
@@ -171,7 +193,7 @@ def body(self, features):
       labels_shape = common_layers.shape_list(labels)
       # handle videos
       if len(labels.shape) == 5:
-        labels = common_layers.time_to_channels(labels)
+        labels = time_to_channels(labels)
       shape = common_layers.shape_list(labels)
       x = tf.one_hot(labels, vocab_size)
       x = self.embed(x)
@@ -1126,12 +1148,12 @@ def autoencoder_ordered_discrete_vq():
 def autoencoder_discrete_pong():
   """Discrete autoencoder model for compressing pong frames."""
   hparams = autoencoder_ordered_discrete()
-  hparams.num_hidden_layers = 2
+  hparams.num_hidden_layers = 3
   hparams.bottleneck_bits = 24
   hparams.batch_size = 2
-  hparams.bottleneck_noise = 0.2
-  hparams.max_hidden_size = 1024
-  hparams.gan_loss_factor = 0.0
+  hparams.gan_loss_factor = 0.01
+  hparams.bottleneck_l2_factor = 0.001
+  hparams.add_hparam("video_modality_loss_cutoff", 0.02)
   return hparams
 
 
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index f1861f899..11df93cf9 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -224,12 +224,6 @@ def body(self, features):
     # Concatenate results and return them.
     frames = tf.stack(res_frames, axis=1)
 
-    if self.is_per_pixel_softmax:
-      def make_gif_ready(tensor_list):
-        return tf.cast(tf.stack(tensor_list, axis=1), tf.uint8)
-      summary = common_video.gif_summary
-      summary("pred", make_gif_ready(sampled_frames_raw))
-
     if "target_reward" not in features:
       return frames
     rewards = tf.concat(res_rewards, axis=1)
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 8e5f9b145..6b2e81890 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -755,22 +755,6 @@ def rl_modelrl_l1_base():
   return hparams
 
 
-@registry.register_hparams
-def rl_modelrl_l1_medium():
-  """Medium parameter set with L1 loss."""
-  hparams = rl_modelrl_medium()
-  hparams.generative_model_params = "next_frame_l1"
-  return hparams
-
-
-@registry.register_hparams
-def rl_modelrl_l1_short():
-  """Short parameter set with L1 loss."""
-  hparams = rl_modelrl_short()
-  hparams.generative_model_params = "next_frame_l1"
-  return hparams
-
-
 @registry.register_hparams
 def rl_modelrl_l1_tiny():
   """Tiny parameter set with L1 loss."""
@@ -787,22 +771,6 @@ def rl_modelrl_l2_base():
   return hparams
 
 
-@registry.register_hparams
-def rl_modelrl_l2_medium():
-  """Medium parameter set with L2 loss."""
-  hparams = rl_modelrl_medium()
-  hparams.generative_model_params = "next_frame_l2"
-  return hparams
-
-
-@registry.register_hparams
-def rl_modelrl_l2_short():
-  """Short parameter set with L2 loss."""
-  hparams = rl_modelrl_short()
-  hparams.generative_model_params = "next_frame_l2"
-  return hparams
-
-
 @registry.register_hparams
 def rl_modelrl_l2_tiny():
   """Tiny parameter set with L2 loss."""
@@ -817,49 +785,7 @@ def rl_modelrl_ae_base():
   hparams = rl_modelrl_base()
   hparams.ppo_params = "ppo_pong_ae_base"
   hparams.generative_model_params = "next_frame_ae"
-  hparams.autoencoder_train_steps = 50000
-  return hparams
-
-
-@registry.register_hparams
-def rl_modelrl_ae_25k():
-  hparams = rl_modelrl_ae_base()
-  hparams.num_real_env_frames //= 4
-  return hparams
-
-
-@registry.register_hparams
-def rl_modelrl_ae_l1_base():
-  """Parameter set for autoencoders and L1 loss."""
-  hparams = rl_modelrl_ae_base()
-  hparams.generative_model_params = "next_frame_l1"
-  return hparams
-
-
-@registry.register_hparams
-def rl_modelrl_ae_l2_base():
-  """Parameter set for autoencoders and L2 loss."""
-  hparams = rl_modelrl_ae_base()
-  hparams.generative_model_params = "next_frame_l2"
-  return hparams
-
-
-@registry.register_hparams
-def rl_modelrl_ae_medium():
-  """Medium parameter set for autoencoders."""
-  hparams = rl_modelrl_ae_base()
-  hparams.num_real_env_frames //= 2
-  return hparams
-
-
-@registry.register_hparams
-def rl_modelrl_ae_short():
-  """Small parameter set for autoencoders."""
-  hparams = rl_modelrl_ae_base()
-  hparams.autoencoder_train_steps //= 10
-  hparams.num_real_env_frames //= 5
-  hparams.model_train_steps //= 10
-  hparams.ppo_epochs_num //= 10
+  hparams.autoencoder_train_steps = 30000
   return hparams
 
 

From 59abe0841e94088f7e7b57ab9c55ed610c83ffbd Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Fri, 14 Sep 2018 09:31:37 -0700
Subject: [PATCH 0844/2720] Updates to discrete autoencoders

PiperOrigin-RevId: 212994702
---
 tensor2tensor/data_generators/video_utils.py  | 77 +++++++++++--------
 tensor2tensor/models/research/autoencoders.py | 20 +++--
 2 files changed, 57 insertions(+), 40 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index d67d51866..7a1b5728e 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -37,8 +37,9 @@ def resize_video_frames(images, size):
   resized_images = []
   for image in images:
     resized_images.append(
-        tf.to_int64(tf.image.resize_images(
-            image, [size, size], tf.image.ResizeMethod.BILINEAR)))
+        tf.to_int64(
+            tf.image.resize_images(image, [size, size],
+                                   tf.image.ResizeMethod.BILINEAR)))
   return resized_images
 
 
@@ -60,10 +61,16 @@ def display_video_hooks(hook_args):
     input_videos = np.concatenate((input_videos, target_videos), axis=1)
     output_videos = np.concatenate((input_videos, output_videos), axis=1)
     input_summ_vals, _ = common_video.py_gif_summary(
-        "decode_%d/input" % decode_ind, input_videos, max_outputs=10, fps=fps,
+        "decode_%d/input" % decode_ind,
+        input_videos,
+        max_outputs=10,
+        fps=fps,
         return_summary_value=True)
     output_summ_vals, _ = common_video.py_gif_summary(
-        "decode_%d/output" % decode_ind, output_videos, max_outputs=10, fps=fps,
+        "decode_%d/output" % decode_ind,
+        output_videos,
+        max_outputs=10,
+        fps=fps,
         return_summary_value=True)
     all_summaries.extend(input_summ_vals)
     all_summaries.extend(output_summ_vals)
@@ -88,8 +95,8 @@ def summarize_video_metrics(hook_args):
           predictions)
     else:
       metrics_results, _ = video_metrics.compute_video_metrics_from_png_files(
-          output_dirs, problem_name,
-          hparams.video_num_target_frames, frame_shape)
+          output_dirs, problem_name, hparams.video_num_target_frames,
+          frame_shape)
 
   summary_values = []
   for name, array in six.iteritems(metrics_results):
@@ -107,6 +114,10 @@ def __init__(self, *args, **kwargs):
     # Path to a directory to dump generated frames as png for debugging.
     # If empty, no debug frames will be generated.
     self.debug_dump_frames_path = ""
+    # Whether to skip random inputs at the beginning or not.
+    self.settable_random_skip = True
+    self.settable_use_not_breaking_batching = True
+    self.shuffle = True
 
   @property
   def num_channels(self):
@@ -172,8 +183,7 @@ def preprocess_example(self, example, mode, hparams):
     """Runtime preprocessing, e.g., resize example["frame"]."""
     if getattr(hparams, "preprocess_resize_frames", None) is not None:
       example["frame"] = tf.image.resize_images(
-          example["frame"],
-          hparams.preprocess_resize_frames,
+          example["frame"], hparams.preprocess_resize_frames,
           tf.image.ResizeMethod.BILINEAR)
     return example
 
@@ -235,6 +245,7 @@ def serving_input_fn(self, hparams):
         receiver_tensors=video_input_frames)
 
   def preprocess(self, dataset, mode, hparams, interleave=True):
+
     def split_on_batch(x):
       """Split x on batch dimension into x[:size, ...] and x[size:, ...]."""
       length = len(x.get_shape())
@@ -293,6 +304,7 @@ def avoid_break_batching(dataset):
         batched not-broken videos.
 
       """
+
       def check_integrity_and_batch(*datasets):
         """Checks whether a sequence of frames are from the same video.
 
@@ -306,11 +318,11 @@ def check_integrity_and_batch(*datasets):
         if "frame_number" in datasets[0]:
           frame_numbers = [dataset["frame_number"][0] for dataset in datasets]
 
-          not_broken = tf.equal(
-              frame_numbers[-1] - frame_numbers[0], num_frames-1)
+          not_broken = tf.equal(frame_numbers[-1] - frame_numbers[0],
+                                num_frames - 1)
           if self.only_keep_videos_from_0th_frame:
-            not_broken = tf.logical_and(not_broken,
-                                        tf.equal(frame_numbers[0], 0))
+            not_broken = tf.logical_and(not_broken, tf.equal(
+                frame_numbers[0], 0))
         else:
           tf.logging.warning("use_not_breaking_batching is True but "
                              "no frame_number is in the dataset.")
@@ -331,34 +343,36 @@ def check_integrity_and_batch(*datasets):
       return dataset
 
     preprocessed_dataset = dataset.map(_preprocess)
-    num_frames = (hparams.video_num_input_frames +
-                  hparams.video_num_target_frames)
+    num_frames = (
+        hparams.video_num_input_frames + hparams.video_num_target_frames)
     # We jump by a random position at the beginning to add variety.
-    if self.random_skip and interleave:
+    if (self.random_skip and self.settable_random_skip and interleave and
+        mode == tf.estimator.ModeKeys.TRAIN):
       random_skip = tf.random_uniform([], maxval=num_frames, dtype=tf.int64)
       preprocessed_dataset = preprocessed_dataset.skip(random_skip)
-    if self.use_not_breaking_batching:
+    if (self.use_not_breaking_batching and
+        self.settable_use_not_breaking_batching):
       batch_dataset = avoid_break_batching(preprocessed_dataset)
     else:
       batch_dataset = preprocessed_dataset.apply(
           tf.contrib.data.batch_and_drop_remainder(num_frames))
     dataset = batch_dataset.map(features_from_batch)
-    if interleave:
+    if self.shuffle and interleave and mode == tf.estimator.ModeKeys.TRAIN:
       dataset = dataset.shuffle(hparams.get("shuffle_buffer_size", 128))
     return dataset
 
   def eval_metrics(self):
     eval_metrics = [
         metrics.Metrics.ACC, metrics.Metrics.ACC_PER_SEQ,
-        metrics.Metrics.NEG_LOG_PERPLEXITY, metrics.Metrics.IMAGE_SUMMARY]
+        metrics.Metrics.NEG_LOG_PERPLEXITY, metrics.Metrics.IMAGE_SUMMARY
+    ]
     return eval_metrics
 
   def validate_frame(self, frame):
     height, width, channels = frame.shape
     if channels != self.num_channels:
       raise ValueError("Generated frame has %d channels while the class "
-                       "assumes %d channels." % (channels,
-                                                 self.num_channels))
+                       "assumes %d channels." % (channels, self.num_channels))
     if height != self.frame_height:
       raise ValueError("Generated frame has height %d while the class "
                        "assumes height %d." % (height, self.frame_height))
@@ -374,8 +388,8 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
         over user-supplied vocab files if there are extra fields needing them.
       tmp_dir: temporary directory that you can use for downloading and scratch.
       dataset_split: problem.DatasetSplit, which data split to generate samples
-        for (for example, training and evaluation). You can assume it's TRAIN
-        if self.
+        for (for example, training and evaluation). You can assume it's TRAIN if
+        self.
 
     Yields:
       Sample: dict<str feature_name, feature value>; we assume that there is
@@ -409,16 +423,15 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
       writer = common_video.VideoWriter(fps=10, file_format="avi")
 
     with tf.Graph().as_default():
-      image_t = tf.placeholder(
-          dtype=tf.uint8, shape=(None, None, None))
+      image_t = tf.placeholder(dtype=tf.uint8, shape=(None, None, None))
       encoded_image_t = tf.image.encode_png(image_t)
       with tf.Session() as sess:
         for features in self.generate_samples(data_dir, tmp_dir, dataset_split):
           unencoded_frame = features.pop("frame")
           self.validate_frame(unencoded_frame)
           height, width, _ = unencoded_frame.shape
-          encoded_frame = sess.run(encoded_image_t, feed_dict={
-              image_t: unencoded_frame})
+          encoded_frame = sess.run(
+              encoded_image_t, feed_dict={image_t: unencoded_frame})
           features["image/encoded"] = [encoded_frame]
           features["image/format"] = ["png"]
           features["image/height"] = [height]
@@ -427,8 +440,8 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
           has_debug_image = "image/debug" in features
           if has_debug_image:
             unencoded_debug = features.pop("image/debug")
-            encoded_debug = sess.run(encoded_image_t, feed_dict={
-                image_t: unencoded_debug})
+            encoded_debug = sess.run(
+                encoded_image_t, feed_dict={image_t: unencoded_debug})
             features["image/encoded_debug"] = [encoded_debug]
 
           if self.debug_dump_frames_path:
@@ -462,13 +475,13 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
     if self.is_generate_per_split:
       for split, paths in split_paths:
         generator_utils.generate_files(
-            self.generate_encoded_samples(
-                data_dir, tmp_dir, split), paths,
+            self.generate_encoded_samples(data_dir, tmp_dir, split),
+            paths,
             cycle_every_n=self.total_number_of_frames // len(paths))
     else:
       generator_utils.generate_files(
-          self.generate_encoded_samples(
-              data_dir, tmp_dir, problem.DatasetSplit.TRAIN),
+          self.generate_encoded_samples(data_dir, tmp_dir,
+                                        problem.DatasetSplit.TRAIN),
           all_paths,
           cycle_every_n=self.total_number_of_frames // len(all_paths))
 
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 0c7dea48d..d920c794f 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -236,7 +236,6 @@ def body(self, features):
             common_layers.shape_list(x)[-1],
             reuse=True)
         x = tf.concat([g, x], axis=0)
-        encoder_layers = [tf.concat([l, l], axis=0) for l in encoder_layers]
     else:
       if self._cur_bottleneck_tensor is None:
         b = self.sample()
@@ -694,10 +693,15 @@ def decoder(self, x, encoder_layers=None):
           if encoder_layers is not None:
             enc_x = encoder_layers[j]
             enc_shape = common_layers.shape_list(enc_x)
-            x = x[:, :enc_shape[1], :enc_shape[2], :]
+            x_mix = x[:enc_shape[0], :enc_shape[1], :enc_shape[2], :]
             if is_training:  # Mix at the beginning of training.
-              rand = tf.random_uniform(common_layers.shape_list(x))
-              x = tf.where(tf.less(rand, nomix_p), x, enc_x)
+              rand = tf.random_uniform(common_layers.shape_list(x_mix))
+              x_mix = tf.where(tf.less(rand, nomix_p), x_mix, enc_x)
+            if hparams.gan_loss_factor != 0:
+              x_gan = x[enc_shape[0]:, :enc_shape[1], :enc_shape[2], :]
+              x = tf.concat([x_mix, x_gan], axis=0)
+            else:
+              x = x_mix
       return x
 
 
@@ -1075,8 +1079,8 @@ def autoencoder_ordered_discrete_image64():
   hparams = autoencoder_ordered_discrete()
   hparams.batch_size = 32
   hparams.num_hidden_layers = 6
-  hparams.target_modality = "video:default"
-  hparams.input_modalities = "video:default"
+  hparams.bottleneck_warmup_steps *= 2
+  hparams.gan_codes_warmup_steps *= 2
 
   return hparams
 
@@ -1125,8 +1129,8 @@ def autoencoder_ordered_text():
 def autoencoder_ordered_text_small():
   """Ordered discrete autoencoder model for text, small version."""
   hparams = autoencoder_ordered_text()
-  hparams.bottleneck_bits = 14
-  hparams.num_hidden_layers = 2
+  hparams.bottleneck_bits = 32
+  hparams.num_hidden_layers = 3
   hparams.hidden_size = 64
   hparams.max_hidden_size = 512
   hparams.bottleneck_noise = 0.0

From 165ac62bdcedf670d06edd0ada72b491c5a53330 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 14 Sep 2018 09:56:44 -0700
Subject: [PATCH 0845/2720] Make t2t_model helper function public.

PiperOrigin-RevId: 212998328
---
 tensor2tensor/utils/t2t_model.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 063f24f73..3e600f613 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1378,7 +1378,7 @@ def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
 
     if use_tpu:
       host_call = _create_host_call(self.hparams.model_dir)
-      _remove_summaries()
+      remove_summaries()
       return tf.contrib.tpu.TPUEstimatorSpec(
           tf.estimator.ModeKeys.TRAIN,
           loss=loss,
@@ -1398,9 +1398,9 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
 
     problem = hparams.problem
     if common_layers.is_xla_compiled():
-      _remove_summaries()
+      remove_summaries()
       if isinstance(logits, dict):
-        eval_metrics_fn = _create_tpu_eval_metrics_fn(problem, hparams)
+        eval_metrics_fn = create_tpu_eval_metrics_fn(problem, hparams)
         # For TPU, logits dict will be passed as keyword arguments to
         # eval_metrics_fn. Here we add the labels to those arguments.
         logits.update({"labels": labels})
@@ -1409,7 +1409,7 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
             eval_metrics=(eval_metrics_fn, logits),
             loss=loss)
       else:
-        eval_metrics_fn = _create_tpu_eval_metrics_fn(problem, hparams)
+        eval_metrics_fn = create_tpu_eval_metrics_fn(problem, hparams)
         return tf.contrib.tpu.TPUEstimatorSpec(
             tf.estimator.ModeKeys.EVAL,
             eval_metrics=(eval_metrics_fn, [logits, labels]),
@@ -1497,7 +1497,7 @@ def estimator_spec_predict(self, features, use_tpu=False):
     if "batch_prediction_key" in predictions:
       export_out["batch_prediction_key"] = predictions["batch_prediction_key"]
 
-    _remove_summaries()
+    remove_summaries()
 
     export_outputs = {
         tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
@@ -1578,7 +1578,7 @@ def _create_dummy_vars():
 ])
 
 
-def _create_tpu_eval_metrics_fn(problem, hparams):
+def create_tpu_eval_metrics_fn(problem, hparams):
   """Create the metrics_fn that TPUEstimatorSpec expects."""
 
   metric_fns = []
@@ -1649,7 +1649,7 @@ def all_metrics_fn(logits=None, labels=None, **kwargs):
   return all_metrics_fn
 
 
-def _remove_summaries():
+def remove_summaries():
   g = tf.get_default_graph()
   key = tf.GraphKeys.SUMMARIES
   del g.get_collection_ref(key)[:]

From bc5b94f38ba4e7615e7f0107f2e0c056c50bd52f Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Fri, 14 Sep 2018 10:14:30 -0700
Subject: [PATCH 0846/2720] fixes to experiment player.

PiperOrigin-RevId: 213001695
---
 tensor2tensor/rl/model_rl_experiment_player.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index 3f53f6b07..b53d264d3 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -186,7 +186,7 @@ def _augment_observation(self):
     print("Info str:{}".format(info_str))
     # info_pane = write_on_image(info_pane, info_str)
 
-    augmented_observ = concatenate_images(observ, info_pane)
+    augmented_observ = concatenate_images([observ, info_pane])
     augmented_observ = np.array(augmented_observ)
     return augmented_observ
 
@@ -240,6 +240,12 @@ def main(_):
   else:
     simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
                               % game_with_mode)
+    if simulated_problem_name not in registry.list_problems():
+      tf.logging.info("Game Problem %s not found; dynamically registering",
+                      simulated_problem_name)
+      gym_problems_specs.create_problems_for_game(hparams.game,
+                                                  game_mode="Deterministic-v4")
+
   epoch = hparams.epochs-1
   epoch_data_dir = os.path.join(directories["data"], str(epoch))
   ppo_model_dir = directories["ppo"]

From c008a650428a3e054a098dcca5131ecdedb20e1a Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Fri, 14 Sep 2018 10:35:57 -0700
Subject: [PATCH 0847/2720] use only trajectories from last epoch for PPO
 random starts.

PiperOrigin-RevId: 213005590
---
 tensor2tensor/data_generators/problem.py     | 16 +++++++++++++---
 tensor2tensor/rl/envs/simulated_batch_env.py |  3 ++-
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 7ffdac942..4b9fd5ca0 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -544,7 +544,8 @@ def dataset(self,
               shard=None,
               partition_id=0,
               num_partitions=1,
-              max_records=-1):
+              max_records=-1,
+              only_last=False):
     """Build a Dataset for this problem.
 
     Args:
@@ -566,6 +567,7 @@ def dataset(self,
       partition_id: integer - which partition of the dataset to read from
       num_partitions: how many partitions in the dataset
       max_records: int, number of records to truncate to.
+      only_last: bool, whether we should include only files from last epoch.
 
     Returns:
       Dataset containing dict<feature name, Tensor>.
@@ -590,9 +592,17 @@ def dataset(self,
     _ = self.get_hparams(hparams)
 
     data_filepattern = self.filepattern(data_dir, dataset_split, shard=shard)
+    if only_last:
+      imprv_data_filepattern = data_filepattern + r"10.[\d+]"
+    else:
+      imprv_data_filepattern = data_filepattern
     tf.logging.info("Reading data files from %s", data_filepattern)
-    data_files = sorted(tf.contrib.slim.parallel_reader.get_data_files(
-        data_filepattern))
+    try:
+      data_files = sorted(tf.contrib.slim.parallel_reader.get_data_files(
+          imprv_data_filepattern))
+    except ValueError:
+      data_files = sorted(tf.contrib.slim.parallel_reader.get_data_files(
+          data_filepattern))
 
     # Functions used in dataset transforms below. `filenames` can be either a
     # `tf.string` tensor or `tf.data.Dataset` containing one or more filenames.
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index c2838ac49..0c12bbc10 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -140,7 +140,8 @@ def __init__(self, environment_spec, length):
       dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
                                                FLAGS.data_dir,
                                                shuffle_files=True,
-                                               hparams=hparams)
+                                               hparams=hparams,
+                                               only_last=True)
       dataset = dataset.shuffle(buffer_size=1000)
       if environment_spec.simulation_flip_first_random_for_beginning:
         # Later flip the first random frame in PPO batch for the true beginning.

From 5a7f04cbc85f79f137e87cd7f03f6569b8f50104 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 14 Sep 2018 14:16:38 -0700
Subject: [PATCH 0848/2720] Add support for frame-by-frame summaries at decode
 time.

PiperOrigin-RevId: 213042021
---
 tensor2tensor/data_generators/video_utils.py  | 25 +++++++++++++++----
 .../data_generators/video_utils_test.py       |  4 +--
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 7a1b5728e..80852441f 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -23,6 +23,7 @@
 import six
 
 from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.layers import common_video
@@ -58,22 +59,36 @@ def display_video_hooks(hook_args):
     output_videos = np.asarray(output_videos, dtype=np.uint8)
     input_videos = np.asarray(input_videos, dtype=np.uint8)
 
-    input_videos = np.concatenate((input_videos, target_videos), axis=1)
-    output_videos = np.concatenate((input_videos, output_videos), axis=1)
+    # Video gif.
+    all_input = np.concatenate((input_videos, target_videos), axis=1)
+    all_output = np.concatenate((input_videos, output_videos), axis=1)
     input_summ_vals, _ = common_video.py_gif_summary(
         "decode_%d/input" % decode_ind,
-        input_videos,
-        max_outputs=10,
+        all_input, max_outputs=10,
         fps=fps,
         return_summary_value=True)
     output_summ_vals, _ = common_video.py_gif_summary(
         "decode_%d/output" % decode_ind,
-        output_videos,
+        all_output,
         max_outputs=10,
         fps=fps,
         return_summary_value=True)
     all_summaries.extend(input_summ_vals)
     all_summaries.extend(output_summ_vals)
+
+    # Frame-by-frame summaries
+    iterable = zip(all_input[:10], all_output[:10])
+    for ind, (input_video, output_video) in enumerate(iterable):
+      t, h, w, c = input_video.shape
+      # Tile vertically
+      input_frames = np.reshape(input_video, (t*h, w, c))
+      output_frames = np.reshape(output_video, (t*h, w, c))
+
+      # Concat across width.
+      all_frames = np.concatenate((input_frames, output_frames), axis=1)
+      frame_by_frame_summ = image_utils.image_to_tf_summary_value(
+          all_frames, tag="input/output/decode_%d" % ind)
+      all_summaries.append(frame_by_frame_summ)
   return all_summaries
 
 
diff --git a/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index 02f467704..7c807bd88 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -43,8 +43,8 @@ def testConvertPredictionsToVideoSummaries(self):
         hparams=decode_hparams, decode_hparams=decode_hparams,
         predictions=predictions)
     summaries = video_utils.display_video_hooks(decode_hooks)
-    # ground_truth + output.
-    self.assertEqual(len(summaries), 20)
+    # 10 input vids + 10 output vids + 10 frame-by-frame.
+    self.assertEqual(len(summaries), 30)
     for summary in summaries:
       self.assertTrue(isinstance(summary, tf.Summary.Value))
 

From 31d1b88190e8ee9ff5063003785e01549dcb518c Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 14 Sep 2018 17:52:11 -0700
Subject: [PATCH 0849/2720] internal

PiperOrigin-RevId: 213072400
---
 tensor2tensor/mesh_tensorflow/OWNERS | 2 ++
 tensor2tensor/models/OWNERS          | 5 +++++
 2 files changed, 7 insertions(+)
 create mode 100644 tensor2tensor/mesh_tensorflow/OWNERS
 create mode 100644 tensor2tensor/models/OWNERS

diff --git a/tensor2tensor/mesh_tensorflow/OWNERS b/tensor2tensor/mesh_tensorflow/OWNERS
new file mode 100644
index 000000000..681d09a90
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/OWNERS
@@ -0,0 +1,2 @@
+nikip
+noam
diff --git a/tensor2tensor/models/OWNERS b/tensor2tensor/models/OWNERS
new file mode 100644
index 000000000..d3af3590b
--- /dev/null
+++ b/tensor2tensor/models/OWNERS
@@ -0,0 +1,5 @@
+avaswani
+mbz
+nikip
+noam
+uszkoreit

From af64ddb5ba598861c4e6b355c4d16a3e6eb56845 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 14 Sep 2018 19:55:26 -0700
Subject: [PATCH 0850/2720] more updates on the hparams of
 transformer_mlperf_tpu.

PiperOrigin-RevId: 213080819
---
 tensor2tensor/models/transformer.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index d7d6314fc..4ab85d284 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -2041,8 +2041,11 @@ def transformer_mlperf_tpu():
   """HParams for Transformer model on TPU for MLPerf on TPU 2x2."""
   hparams = transformer_base_v3()
   hparams.symbol_modality_num_shards = 1
-  hparams.max_length = 64  # ignored when using "_packed" problems
-  hparams.batch_size = 512  # gloabl batch size matches the reference model
+  hparams.max_length = 256  # ignored when using "_packed" problems
+  hparams.batch_size = 2048  # per-chip batch size matches the reference model
+  hparams.hidden_size = 1024
+  hparams.filter_size = 4096
+  hparams.num_heads = 16
   hparams.attention_dropout_broadcast_dims = "0,1"  # batch, heads
   hparams.relu_dropout_broadcast_dims = "1"  # length
   hparams.layer_prepostprocess_dropout_broadcast_dims = "1"  # length

From 1fca882b628a75858f0d126da3b447a683f47aa1 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 14 Sep 2018 21:21:19 -0700
Subject: [PATCH 0851/2720] Create Modalities in Problem, not T2TModel.

PiperOrigin-RevId: 213085818
---
 .../data_generators/multi_problem.py          |  6 +-
 .../data_generators/multi_problem_test.py     | 48 ++++++++++
 tensor2tensor/data_generators/problem.py      | 82 +++++++++++++++++
 tensor2tensor/data_generators/problem_test.py | 43 +++++++++
 .../models/research/vqa_attention_test.py     | 12 +--
 tensor2tensor/models/resnet_test.py           |  8 +-
 tensor2tensor/models/video/tests_utils.py     | 21 ++---
 tensor2tensor/models/xception_test.py         |  8 +-
 tensor2tensor/utils/data_reader_test.py       |  6 +-
 tensor2tensor/utils/metrics.py                | 12 +--
 tensor2tensor/utils/t2t_model.py              | 90 +++++--------------
 tensor2tensor/utils/trainer_lib_test.py       |  4 +-
 12 files changed, 235 insertions(+), 105 deletions(-)
 create mode 100644 tensor2tensor/data_generators/multi_problem_test.py

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 99b1b5da1..f8d32cc71 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -22,8 +22,8 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
-from tensor2tensor.utils import registry
 import tensorflow as tf
 
 
@@ -88,8 +88,8 @@ def get_hparams(self, model_hparams=None):
     vocab_size = self._hparams.vocabulary["targets"].vocab_size
     tf.logging.info("Old vocabulary size: %d" % vocab_size)
     tf.logging.info("New vocabulary size: %d" % (vocab_size + vocab_size_inc))
-    self._hparams.target_modality = (registry.Modalities.SYMBOL,
-                                     vocab_size + vocab_size_inc)
+    self._hparams.target_modality = modalities.SymbolModality(
+        vocab_size + vocab_size_inc)
 
     return self._hparams
 
diff --git a/tensor2tensor/data_generators/multi_problem_test.py b/tensor2tensor/data_generators/multi_problem_test.py
new file mode 100644
index 000000000..efd76fc0c
--- /dev/null
+++ b/tensor2tensor/data_generators/multi_problem_test.py
@@ -0,0 +1,48 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 MultiProblem."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators import multi_problem
+from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.layers import modalities
+
+import tensorflow as tf
+
+
+class TestMultiProblem(multi_problem.MultiProblem):
+  """Test multi-problem."""
+
+  def __init__(self):
+    super(TestMultiProblem, self).__init__()
+    self.task_list.append(problem_hparams.TestProblem(2, 3))
+    self.task_list.append(problem_hparams.TestProblem(4, 5))
+
+
+class MultiProblemTest(tf.test.TestCase):
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testProblemHparamsModality(self):
+    problem = TestMultiProblem()
+    p_hparams = problem.get_hparams()
+    self.assertIsInstance(p_hparams.input_modality["inputs"],
+                          modalities.SymbolModality)
+    self.assertIsInstance(p_hparams.target_modality, modalities.SymbolModality)
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 4b9fd5ca0..4c64ab27d 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -17,6 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 import collections
+import copy
 import functools
 import multiprocessing
 import os
@@ -25,8 +26,11 @@
 import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import text_encoder
+# Import modalities: they must be registered before we look them up here.
+from tensor2tensor.layers import modalities  # pylint: disable=unused-import
 from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import metrics
+from tensor2tensor.utils import registry
 import tensorflow as tf
 from tensorflow.contrib.tpu.python.tpu import tpu_config
 
@@ -467,6 +471,8 @@ def get_feature_encoders(self, data_dir=None):
 
   def get_hparams(self, model_hparams=None):
     """Returns problem_hparams."""
+    if model_hparams is None:
+      model_hparams = default_model_hparams()
     if self._hparams is not None:
       return self._hparams
 
@@ -490,6 +496,17 @@ def get_hparams(self, model_hparams=None):
     if self._was_copy:
       _copy_problem_hparams(hp)
 
+    model_hparams = copy.copy(model_hparams)
+    if (hasattr(model_hparams, "shared_embedding_and_softmax_weights") and
+        model_hparams.shared_embedding_and_softmax_weights):
+      # If vocabularies differ, unset shared_embedding_and_softmax_weights.
+      input_vocab_size = hp.input_modality.get("inputs")[1]
+      target_vocab_size = hp.target_modality[1]
+      if input_vocab_size != target_vocab_size:
+        tf.logging.warn("Unsetting shared_embedding_and_softmax_weights.")
+        model_hparams.shared_embedding_and_softmax_weights = 0
+
+    _create_modalities(hp, model_hparams)
     self._hparams = hp
     return self._hparams
 
@@ -1078,6 +1095,71 @@ def _reverse_problem_hparams(p_hparams):
   p.was_reversed = True
 
 
+def _create_modalities(problem_hparams, hparams):
+  """Converts string-type modalities to their modality object.
+
+  Args:
+    problem_hparams: tf.contrib.training.HParams for the Problem. It must have
+      input_modality and target_modality.
+    hparams: tf.contrib.training.HParams for the model. It may have
+      input_modalities and target_modality, which will override
+      problem_hparams's modalities.
+  """
+  input_modality_overrides = {}
+  if hasattr(hparams, "input_modalities"):
+    for override_str in hparams.input_modalities.split(";"):
+      if override_str != "default":
+        parts = override_str.split(":")
+        feature_name = parts[0]
+        modality_name = ":".join(parts[1:])
+        input_modality_overrides[feature_name] = modality_name
+
+  input_modality = {}
+  for f, modality_spec in six.iteritems(problem_hparams.input_modality):
+    if f in input_modality_overrides:
+      _warn_changed_modality_type(input_modality_overrides[f],
+                                  modality_spec[0], f)
+      modality_spec = (input_modality_overrides[f], modality_spec[1])
+    input_modality[f] = registry.create_modality(modality_spec, hparams)
+  problem_hparams.input_modality = input_modality
+
+  target_modality_name = None
+  if (hasattr(hparams, "target_modality") and
+      hparams.target_modality != "default"):
+    target_modality_name = hparams.target_modality
+
+  if problem_hparams.target_modality is None:
+    target_modality = None
+  elif isinstance(problem_hparams.target_modality, dict):
+    target_modality = {}
+    for f, modality_spec in six.iteritems(problem_hparams.target_modality):
+      # TODO(lukaszkaiser): allow overriding other target modalities.
+      if target_modality_name and f == "targets":
+        _warn_changed_modality_type(target_modality_name, modality_spec[0],
+                                    "target_modality/%s" % f)
+        modality_spec = (target_modality_name, modality_spec[1])
+      target_modality[f] = registry.create_modality(modality_spec, hparams)
+  else:
+    target_modality_spec = problem_hparams.target_modality
+    if target_modality_name:
+      _warn_changed_modality_type(target_modality_name,
+                                  target_modality_spec[0], "target")
+      target_modality_spec = (target_modality_name, target_modality_spec[1])
+    target_modality = registry.create_modality(target_modality_spec,
+                                               hparams)
+  problem_hparams.target_modality = target_modality
+
+
+def _warn_changed_modality_type(new_name, old_name, feature_name):
+  new_type, new_name = registry.parse_modality_name(new_name)
+  old_type, old_name = registry.parse_modality_name(old_name)
+  if new_type != old_type:
+    tf.logging.warn(
+        "%s has a designated modality type %s (%s) but has been "
+        "overridden with a modality of type %s (%s).", feature_name, old_type,
+        old_name, new_type, new_name)
+
+
 def _default_hparams():
   """A set of basic model hyperparameters."""
   return tf.contrib.training.HParams(
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 26e370bd8..36b1ecd59 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -19,7 +19,13 @@
 from __future__ import print_function
 
 import numpy as np
+
 from tensor2tensor.data_generators import algorithmic
+from tensor2tensor.data_generators import problem as problem_module
+from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.layers import modalities
+from tensor2tensor.utils import registry
+
 import tensorflow as tf
 
 
@@ -73,6 +79,43 @@ def testNoShufflePreprocess(self):
     with tf.Session() as sess:
       self.assertTrue(assert_tensors_equal(sess, tensor1, tensor2, 20))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testProblemHparamsModality(self):
+    problem = problem_hparams.TestProblem(input_vocab_size=2,
+                                          target_vocab_size=3)
+    p_hparams = problem.get_hparams()
+    self.assertIsInstance(p_hparams.input_modality["inputs"],
+                          modalities.SymbolModality)
+    self.assertIsInstance(p_hparams.target_modality, modalities.SymbolModality)
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testProblemHparamsInputOnlyModality(self):
+    class InputOnlyProblem(problem_module.Problem):
+
+      def hparams(self, defaults, model_hparams):
+        hp = defaults
+        hp.input_modality = {"inputs": (registry.Modalities.SYMBOL, 2)}
+        hp.target_modality = None
+
+    problem = InputOnlyProblem(False, False)
+    p_hparams = problem.get_hparams()
+    self.assertIsInstance(p_hparams.input_modality["inputs"],
+                          modalities.SymbolModality)
+    self.assertIsNone(p_hparams.target_modality)
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testProblemHparamsTargetOnlyModality(self):
+    class TargetOnlyProblem(problem_module.Problem):
+
+      def hparams(self, defaults, model_hparams):
+        hp = defaults
+        hp.input_modality = {}
+        hp.target_modality = (registry.Modalities.SYMBOL, 3)
+
+    problem = TargetOnlyProblem(False, False)
+    p_hparams = problem.get_hparams()
+    self.assertEqual(p_hparams.input_modality, {})
+    self.assertIsInstance(p_hparams.target_modality, modalities.SymbolModality)
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/models/research/vqa_attention_test.py b/tensor2tensor/models/research/vqa_attention_test.py
index 4108572c2..e3e88ba0d 100644
--- a/tensor2tensor/models/research/vqa_attention_test.py
+++ b/tensor2tensor/models/research/vqa_attention_test.py
@@ -21,8 +21,8 @@
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.layers import modalities
 from tensor2tensor.models.research import vqa_attention
-from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -46,11 +46,11 @@ def testVqaAttentionBaseline(self):
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
                                                      hparams)
-    p_hparams.input_modality["inputs"] = (registry.Modalities.IMAGE, None)
-    p_hparams.input_modality["question"] = (registry.Modalities.SYMBOL,
-                                            vocab_size)
-    p_hparams.target_modality = (registry.Modalities.CLASS_LABEL
-                                 + ":multi_label", num_classes + 1)
+    p_hparams.input_modality["inputs"] = modalities.ImageModality(hparams)
+    p_hparams.input_modality["question"] = modalities.SymbolModality(
+        hparams, vocab_size)
+    p_hparams.target_modality = modalities.MultiLabelModality(
+        hparams, num_classes + 1)
     with self.test_session() as session:
       features = {
           "inputs": tf.constant(x, dtype=tf.float32),
diff --git a/tensor2tensor/models/resnet_test.py b/tensor2tensor/models/resnet_test.py
index 76458d58e..50e84b581 100644
--- a/tensor2tensor/models/resnet_test.py
+++ b/tensor2tensor/models/resnet_test.py
@@ -17,11 +17,12 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.layers import modalities
 from tensor2tensor.models import resnet
-from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -46,8 +47,9 @@ def _test_resnet(self, img_size, output_size):
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
                                                      hparams)
-    p_hparams.input_modality["inputs"] = (registry.Modalities.IMAGE, None)
-    p_hparams.target_modality = (registry.Modalities.CLASS_LABEL, vocab_size)
+    p_hparams.input_modality["inputs"] = modalities.ImageModality(hparams)
+    p_hparams.target_modality = modalities.ClassLabelModality(
+        hparams, vocab_size)
     with self.test_session() as session:
       features = {
           "inputs": tf.constant(x, dtype=tf.int32),
diff --git a/tensor2tensor/models/video/tests_utils.py b/tensor2tensor/models/video/tests_utils.py
index 67aa3dd81..96fec0a2d 100644
--- a/tensor2tensor/models/video/tests_utils.py
+++ b/tensor2tensor/models/video/tests_utils.py
@@ -21,6 +21,7 @@
 
 from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
 
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -40,12 +41,12 @@ def fill_hparams(hparams, in_frames, out_frames):
 
 def action_modalities(hparams):
   hparams.problem_hparams.input_modality = {
-      "inputs": ("video:l2raw", 256),
-      "input_action": ("symbol:one_hot", 5)
+      "inputs": modalities.VideoModalityL2Raw(hparams, 256),
+      "input_action": modalities.SymbolModalityOneHot(hparams, 5),
   }
   hparams.problem_hparams.target_modality = {
-      "targets": ("video:l2raw", 256),
-      "target_action": ("symbol:one_hot", 5),
+      "targets": modalities.VideoModalityL2Raw(hparams, 256),
+      "target_action": modalities.SymbolModalityOneHot(hparams, 5),
   }
   return hparams
 
@@ -53,14 +54,14 @@ def action_modalities(hparams):
 def full_modalities(hparams):
   """Full modalities with actions and rewards."""
   hparams.problem_hparams.input_modality = {
-      "inputs": ("video:l2raw", 256),
-      "input_reward": ("symbol:one_hot", 3),
-      "input_action": ("symbol:one_hot", 5)
+      "inputs": modalities.VideoModalityL2Raw(hparams, 256),
+      "input_reward": modalities.SymbolModalityOneHot(hparams, 3),
+      "input_action": modalities.SymbolModalityOneHot(hparams, 5),
   }
   hparams.problem_hparams.target_modality = {
-      "targets": ("video:l2raw", 256),
-      "target_reward": ("symbol:one_hot", 3),
-      "target_action": ("symbol:one_hot", 5),
+      "targets": modalities.VideoModalityL2Raw(hparams, 256),
+      "target_reward": modalities.SymbolModalityOneHot(hparams, 3),
+      "target_action": modalities.SymbolModalityOneHot(hparams, 5),
   }
   hparams.force_full_predict = True
   return hparams
diff --git a/tensor2tensor/models/xception_test.py b/tensor2tensor/models/xception_test.py
index ab29a5e0e..c3c8e9bae 100644
--- a/tensor2tensor/models/xception_test.py
+++ b/tensor2tensor/models/xception_test.py
@@ -17,11 +17,12 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.layers import modalities
 from tensor2tensor.models import xception
-from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -39,8 +40,9 @@ def _test_xception(self, img_size):
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
                                                      hparams)
-    p_hparams.input_modality["inputs"] = (registry.Modalities.IMAGE, None)
-    p_hparams.target_modality = (registry.Modalities.CLASS_LABEL, vocab_size)
+    p_hparams.input_modality["inputs"] = modalities.ImageModality(hparams)
+    p_hparams.target_modality = modalities.ClassLabelModality(
+        hparams, vocab_size)
     with self.test_session() as session:
       features = {
           "inputs": tf.constant(x, dtype=tf.int32),
diff --git a/tensor2tensor/utils/data_reader_test.py b/tensor2tensor/utils/data_reader_test.py
index 6b3a0f313..ccee17c5c 100644
--- a/tensor2tensor/utils/data_reader_test.py
+++ b/tensor2tensor/utils/data_reader_test.py
@@ -48,7 +48,11 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
         self.generator(data_dir, tmp_dir, False), dev_paths)
 
   def hparams(self, defaults, model_hparams):
-    pass
+    hp = defaults
+    hp.input_modality = {
+        "inputs": (registry.Modalities.SYMBOL, 30)
+    }
+    hp.target_modality = (registry.Modalities.SYMBOL, 30)
 
   def example_reading_spec(self):
     data_fields = {
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index db9ddb183..207afe4c2 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -23,7 +23,6 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import bleu_hook
-from tensor2tensor.utils import registry
 from tensor2tensor.utils import rouge
 
 import tensorflow as tf
@@ -584,13 +583,11 @@ def weights_fn_for_mp(problem_task_id):
                                     metrics,
                                     list(METRICS_FNS.keys())))
 
-    tm = problem_instance.get_hparams().target_modality
+    tm = problem_instance.get_hparams(model_hparams).target_modality
     if not isinstance(tm, dict):
       tm = {"targets": tm}
 
     for target_name, modality in six.iteritems(tm):
-      if isinstance(modality, tuple):
-        modality = registry.create_modality(modality, model_hparams)
       weights_fn = modality.targets_weights_fn
       if hasattr(model_hparams.problem, "task_list"):
         ptid = problem_instance.task_id  # pylint: disable=cell-var-from-loop
@@ -608,13 +605,10 @@ def weights_fn_for_mp(problem_task_id):
   return eval_metrics
 
 
-def create_eager_metrics_for_problem(problem, model_hparams=None):
+def create_eager_metrics_for_problem(problem, model_hparams):
   """See create_eager_metrics."""
   metric_names = problem.eval_metrics()
-  tm = problem.get_hparams().target_modality
-  if isinstance(tm, tuple):
-    assert model_hparams is not None
-    tm = registry.create_modality(tm, model_hparams)
+  tm = problem.get_hparams(model_hparams).target_modality
   return create_eager_metrics(metric_names, weights_fn=tm.targets_weights_fn)
 
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 3e600f613..5dc61e9dc 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -29,12 +29,12 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators.problem import problem_hparams_to_features
 from tensor2tensor.layers import common_layers
-from tensor2tensor.layers import modalities  # pylint: disable=unused-import
 from tensor2tensor.utils import beam_search
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import expert_utils as eu
 from tensor2tensor.utils import learning_rate
 from tensor2tensor.utils import metrics
+from tensor2tensor.utils import modality
 from tensor2tensor.utils import optimize
 from tensor2tensor.utils import quantization
 from tensor2tensor.utils import registry
@@ -116,15 +116,15 @@ def __init__(self,
     self._problem_hparams = problem_hparams
 
     # Setup hparams
-    # If vocabularies differ, unset shared_embedding_and_softmax_weights.
     hparams = copy.copy(hparams)
     if self._problem_hparams and hparams.shared_embedding_and_softmax_weights:
-      same_vocab_sizes = True
-      if "inputs" in self._problem_hparams.input_modality:
-        if (self._problem_hparams.input_modality["inputs"] !=
-            self._problem_hparams.target_modality):
-          same_vocab_sizes = False
-      if not same_vocab_sizes:
+      # If vocabularies differ, unset shared_embedding_and_softmax_weights.
+      input_modality = self._problem_hparams.input_modality.get("inputs")
+      target_modality = self._problem_hparams.target_modality
+      if (isinstance(input_modality, modality.Modality) and
+          isinstance(target_modality, modality.Modality) and
+          input_modality.top_dimensionality !=
+          target_modality.top_dimensionality):
         log_info("Unsetting shared_embedding_and_softmax_weights.")
         hparams.shared_embedding_and_softmax_weights = 0
     self._original_hparams = hparams
@@ -136,8 +136,6 @@ def __init__(self,
     self._num_datashards = self._data_parallelism.n
     self._ps_devices = self._data_parallelism.ps_devices
     self._eager_var_store = create_eager_var_store()
-    if self._problem_hparams:
-      self._create_modalities(self._problem_hparams, self._hparams)
     if not common_layers.is_xla_compiled():
       self.summarize_hparams()
     self._variable_scopes = {}
@@ -566,47 +564,16 @@ def set_mode(self, mode):
           setattr(hparams, key, 0.0)
     self._hparams = hparams
 
-  def _create_modalities(self, problem_hparams, hparams):
-    """Construct modalities in problem_hparams."""
-
-    input_modality_overrides = {}
-    for override_str in hparams.input_modalities.split(";"):
-      if override_str != "default":
-        parts = override_str.split(":")
-        feature_name = parts[0]
-        modality_name = ":".join(parts[1:])
-        input_modality_overrides[feature_name] = modality_name
-
-    target_modality_name = None
-    if hparams.target_modality and hparams.target_modality != "default":
-      target_modality_name = hparams.target_modality
-
-    input_modality = {}
-    for f, modality_spec in six.iteritems(problem_hparams.input_modality):
-      if f in input_modality_overrides:
-        _warn_changed_modality_type(input_modality_overrides[f],
-                                    modality_spec[0], f)
-        modality_spec = (input_modality_overrides[f], modality_spec[1])
-      input_modality[f] = registry.create_modality(modality_spec, hparams)
-    problem_hparams.input_modality = input_modality
-
-    if isinstance(problem_hparams.target_modality, dict):
-      target_modality = {}
-      for f, modality_spec in six.iteritems(problem_hparams.target_modality):
-        # TODO(lukaszkaiser): allow overriding other target modalities.
-        if target_modality_name and f == "targets":
-          _warn_changed_modality_type(target_modality_name, modality_spec[0],
-                                      "target_modality/%s" % f)
-          modality_spec = (target_modality_name, modality_spec[1])
-        target_modality[f] = registry.create_modality(modality_spec, hparams)
-    else:
-      target_modality_spec = problem_hparams.target_modality
-      if target_modality_name:
-        _warn_changed_modality_type(target_modality_name,
-                                    target_modality_spec[0], "target")
-        target_modality_spec = (target_modality_name, target_modality_spec[1])
-      target_modality = registry.create_modality(target_modality_spec, hparams)
-    problem_hparams.target_modality = target_modality
+    if self._problem_hparams:
+      # Set model hparams in problem_hparams' modalities, which also store them.
+      for im in six.itervalues(self._problem_hparams.input_modality):
+        im._model_hparams = self._hparams  # pylint: disable=protected-access
+
+      if isinstance(self._problem_hparams.target_modality, dict):
+        for tm in six.itervalues(self._problem_hparams.target_modality):
+          tm._model_hparams = self._hparams  # pylint: disable=protected-access
+      elif self._problem_hparams.target_modality is not None:
+        self._problem_hparams.target_modality._model_hparams = self._hparams  # pylint: disable=protected-access
 
   def prepare_features_for_infer(self, features):
     """Called before inference to allow adding infer-specific features."""
@@ -750,8 +717,7 @@ def symbols_to_logits_fn(ids):
       # it has shape [batch_size] and contains floats between 0 and
       # source_length.
       if self._problem_hparams:
-        modality = self._problem_hparams.target_modality
-        if modality.top_is_pointwise:
+        if self._problem_hparams.target_modality.top_is_pointwise:
           return tf.squeeze(logits, axis=[1, 2, 3])
       # -1 due to the pad above.
       current_output_position = common_layers.shape_list(ids)[1] - 1
@@ -1533,16 +1499,6 @@ def _summarize_losses(self, losses_dict):
           tf.summary.scalar(loss_name, loss_val)
 
 
-def _warn_changed_modality_type(new_name, old_name, feature_name):
-  new_type, new_name = registry.parse_modality_name(new_name)
-  old_type, old_name = registry.parse_modality_name(old_name)
-  if new_type != old_type:
-    log_warn(
-        "%s has a designated modality type %s (%s) but has been "
-        "overridden with a modality of type %s (%s).", feature_name, old_type,
-        old_name, new_type, new_name)
-
-
 def _with_timing(fn, msg, silent=False):
 
   def fn_with_timing(*args, **kwargs):
@@ -1578,17 +1534,15 @@ def _create_dummy_vars():
 ])
 
 
-def create_tpu_eval_metrics_fn(problem, hparams):
+def create_tpu_eval_metrics_fn(problem, model_hparams):
   """Create the metrics_fn that TPUEstimatorSpec expects."""
 
   metric_fns = []
   eval_metrics = problem.eval_metrics()
 
-  tm = problem.get_hparams().target_modality
+  tm = problem.get_hparams(model_hparams).target_modality
   if isinstance(tm, dict):
     for k, v in six.iteritems(tm):
-      if isinstance(v, tuple):
-        v = registry.create_modality(v, hparams)
       weights_fn = v.targets_weights_fn
 
       def make_metric_fn(metric_fn):
@@ -1606,8 +1560,6 @@ def wrapped_metric_fn(logits, labels, weights_fn=weights_fn):
         name = "%s/metrics-%s/%s" % (k, problem.name, metric)
         metric_fns.append((name, make_metric_fn(metrics.METRICS_FNS[metric])))
   else:
-    if isinstance(tm, tuple):
-      tm = registry.create_modality(tm, hparams)
     weights_fn = tm.targets_weights_fn
 
     def make_metric_fn(metric_fn):
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index 49e1a53f9..bc18bad20 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -82,10 +82,12 @@ def testMultipleTargetModalities(self):
     hparams = trainer_lib.create_hparams(
         "transformer_tiny", data_dir=algorithmic.TinyAlgo.data_dir,
         problem_name="tiny_algo")
+    # Manually turn off sharing. It is not currently supported for multitargets.
+    hparams.shared_embedding_and_softmax_weights = 0  # pylint: disable=line-too-long
     hparams.problem_hparams.target_modality = {
         "targets": hparams.problem_hparams.target_modality,
         "A": hparams.problem_hparams.target_modality,
-        "B": hparams.problem_hparams.target_modality
+        "B": hparams.problem_hparams.target_modality,
     }
     hparams.problem._hparams = hparams.problem_hparams
 

From 8bc484f139e495721a31e25ccc91ad7f3b761b73 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Sat, 15 Sep 2018 23:37:58 -0700
Subject: [PATCH 0852/2720] Fix bug in Mesh TensorFlow.

PiperOrigin-RevId: 213153974
---
 tensor2tensor/mesh_tensorflow/mesh_tensorflow.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index 47d746160..3104487af 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -136,7 +136,7 @@ def to_string(self):
   @property
   def cumprod(self):
     """Cumulative product (exclusive) of Dimension sizes."""
-    return _cumprod(self.to_integer_list)[::-1]
+    return _cumprod(self.to_integer_list)[:-1]
 
   def cumprod_to_tensor_axis(self, cumprod):
     """Tensor axis i such that self.cumprod[i] == cumprod, or None."""

From 41647195df4c49f275b3ce4be469efbeb461bd11 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Mon, 17 Sep 2018 00:31:28 -0700
Subject: [PATCH 0853/2720] internal merge of PR #1062

PiperOrigin-RevId: 213226657
---
 tensor2tensor/data_generators/gym_problems.py |  8 ++--
 tensor2tensor/rl/collect.py                   | 13 +++--
 tensor2tensor/rl/envs/tf_atari_wrappers.py    | 24 ++++------
 tensor2tensor/rl/trainer_model_based.py       | 48 +++++++++++--------
 4 files changed, 49 insertions(+), 44 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index b0d5ecc4b..2097d19e8 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -140,7 +140,6 @@ def _setup(self, data_dir):
       self.memory_index = 0
       self.memory = None
 
-
   @property
   def random_skip(self):
     return False
@@ -149,10 +148,11 @@ def _get_data(self):
     if self._use_dumper_data:
       file_path = os.path.join(self._dumper_path,
                                "frame_{}.npz".format(self._dumper_data_index))
-      data = np.load(file_path)
+      with tf.gfile.Open(file_path) as gfile:
+        data = np.load(gfile)
       self._dumper_data_index += 1
-      return data["observ"][0, ...], data["reward"][0], \
-             data["done"][0], data["action"][0]
+      return (data["observ"][0, ...], data["reward"][0], data["done"][0],
+              data["action"][0])
     else:
       if self.memory is None or self.memory_index >= self._internal_memory_size:
         self.memory = self._session.run(self.collect_memory)
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 10184d1cc..95d422f1c 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -261,12 +261,15 @@ def stop_condition(i, _, resets):
   printing = tf.Print(0, [mean_score, scores_sum, scores_num], "mean_score: ")
   with tf.control_dependencies([index, printing]):
     memory = [mem.read_value() for mem in memory]
-    new_memory = []
-    if hasattr(hparams, "effective_num_agents"):
+    # When generating real data together with PPO training we must use single
+    # agent. For PPO to work we reshape the history, as if it was generated
+    # by real_ppo_effective_num_agents.
+    if getattr(hparams, "effective_num_agents", None):
+      new_memory = []
       effective_num_agents = hparams.effective_num_agents
-      assert hparams.epoch_length % effective_num_agents == 0, \
-        "The rollout of hparams.epoch_length will be distributed amongst" \
-        "effective_num_agents of agents"
+      assert hparams.epoch_length % effective_num_agents == 0, (
+          "The rollout of hparams.epoch_length will be distributed amongst"
+          "effective_num_agents of agents")
       new_epoch_length = int(hparams.epoch_length / effective_num_agents)
       for mem, info in zip(memory, rollout_metadata):
         shape, _, name = info
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 3aa696d07..2b293a91b 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -320,24 +320,16 @@ class PyFuncWrapper(WrapperBase):
   def __init__(self, batch_env, process_fun):
     super(PyFuncWrapper, self).__init__(batch_env)
     self.process_fun = process_fun
+    observs_shape = batch_env.observ.shape
+    self._observ = tf.Variable(tf.zeros(observs_shape, self.observ_dtype),
+                               trainable=False)
 
   def simulate(self, action):
     reward, done = self._batch_env.simulate(action)
     with tf.control_dependencies([reward, done]):
-      inputs = [self._batch_env.observ, reward, done, action]
+      inputs = [self._observ.read_value(), reward, done, action]
       ret = tf.py_func(self.process_fun, inputs, tf.double)
-      with tf.control_dependencies([ret]):
-        return tf.identity(reward), tf.identity(done)
-
-  @property
-  def observ(self):
-    """Access the variable holding the current observation."""
-    return self._batch_env.observ
-
-  def __len__(self):
-    """Number of combined environments."""
-    return len(self._batch_env)
-
-  def _reset_non_empty(self, indices):
-    # pylint: disable=protected-access
-    return self._batch_env._reset_non_empty(indices)
+    with tf.control_dependencies([ret]):
+      assign = self._observ.assign(self._batch_env.observ)
+    with tf.control_dependencies([assign]):
+      return tf.identity(reward), tf.identity(done)
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index d8bed1250..bee68cb67 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -31,7 +31,7 @@
 import math
 import os
 import time
-
+import numpy as np
 import six
 
 from tensor2tensor.bin import t2t_trainer
@@ -42,7 +42,6 @@
 from tensor2tensor.rl.envs.tf_atari_wrappers import PyFuncWrapper
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
-import numpy as np
 
 import tensorflow as tf
 
@@ -149,18 +148,18 @@ def train_autoencoder(problem_name, data_dir, output_dir, hparams, epoch):
 def _ppo_training_epochs(hparams, epoch, is_final_epoch, real_env_training):
   """Helper for PPO restarts."""
   if hparams.gather_ppo_real_env_data:
-    assert hparams.real_ppo_epochs_num is 0, \
-      "Should be put to 0 to enforce better readability"
-    real_training_ppo_epochs_num = \
-      math.ceil(hparams.num_real_env_frames/
-                (hparams.epochs*hparams.real_ppo_epoch_length))
+    assert hparams.real_ppo_epochs_num is 0, (
+        "Should be put to 0 to enforce better readability")
+    real_training_ppo_epochs_num = int(math.ceil(
+        hparams.num_real_env_frames /
+        (hparams.epochs*hparams.real_ppo_epoch_length)))
   else:
     real_training_ppo_epochs_num = hparams.real_ppo_epochs_num
 
   simulated_training_ppo_epochs_num = hparams.ppo_epochs_num
 
-  ppo_training_epochs = (epoch + 1)*simulated_training_ppo_epochs_num + \
-                        epoch * real_training_ppo_epochs_num
+  ppo_training_epochs = ((epoch + 1)*simulated_training_ppo_epochs_num
+                         + epoch * real_training_ppo_epochs_num)
   if real_env_training:
     ppo_training_epochs += real_training_ppo_epochs_num
   if is_final_epoch:
@@ -214,14 +213,22 @@ def train_agent(problem_name, agent_model_dir,
 ppo_data_dumper_counter = 0
 dumper_path = None
 
+
 def ppo_data_dumper(observ, reward, done, action):
+  """Save frames from PPO to a numpy file."""
   global ppo_data_dumper_counter, dumper_path
-  np.savez_compressed("{}/frame_{}".format(dumper_path,
-                                           ppo_data_dumper_counter),
-                      observ=observ, reward=reward, done=done, action=action)
+  file_path = "{}/frame_{}.npz".format(dumper_path, ppo_data_dumper_counter)
+  # np.savez_compressed can't create a tf.gfile, so we need to creat it
+  # beforehand.
+  with tf.gfile.Open(file_path, mode="wb+") as gfile:
+    gfile.write("1")
+  with tf.gfile.Open(file_path, mode="wb+") as gfile:
+    np.savez_compressed(gfile,
+                        observ=observ, reward=reward, done=done, action=action)
   ppo_data_dumper_counter += 1
   return 0.0
 
+
 def train_agent_real_env(
     problem_name, agent_model_dir, event_dir, world_model_dir, epoch_data_dir,
     hparams, epoch=0, is_final_epoch=False):
@@ -234,7 +241,7 @@ def train_agent_real_env(
                       "learning_rate", "num_agents",
                       "optimization_epochs", "effective_num_agents"]
 
-  # This should be overriden
+  # This should be overridden.
   ppo_hparams.add_hparam("effective_num_agents", None)
   for param_name in ppo_params_names:
     ppo_param_name = "real_ppo_"+ param_name
@@ -248,9 +255,9 @@ def train_agent_real_env(
   environment_spec = copy.copy(gym_problem.environment_spec)
 
   if hparams.gather_ppo_real_env_data:
-    #TODO(piotrmilos):This should be refactored
-    assert hparams.real_ppo_num_agents == 1, \
-      "It is required to use collect with pyfunc_wrapper"
+    # TODO(piotrmilos):This should be refactored
+    assert hparams.real_ppo_num_agents == 1, (
+        "It is required to use collect with pyfunc_wrapper")
 
     ppo_data_dumper_counter = 0
     dumper_path = os.path.join(epoch_data_dir, "dumper")
@@ -258,7 +265,6 @@ def train_agent_real_env(
     dumper_spec = [PyFuncWrapper, {"process_fun": ppo_data_dumper}]
     environment_spec.wrappers.insert(1, dumper_spec)
 
-
   ppo_hparams.add_hparam("environment_spec", environment_spec)
 
   with temporary_flags({
@@ -611,7 +617,9 @@ def rl_modelrl_base():
       epochs=3,
       # Total frames used for training. This will be distributed evenly across
       # hparams.epochs.
-      num_real_env_frames=100000,
+      # This number should be divisible by real_ppo_epoch_length*epochs
+      # for our frame accounting to be preceise.
+      num_real_env_frames=96000,
       generative_model="next_frame_basic_deterministic",
       generative_model_params="next_frame_pixel_noise",
       ppo_params="ppo_pong_base",
@@ -638,6 +646,7 @@ def rl_modelrl_base():
 
       gather_ppo_real_env_data=True,
       real_ppo_epochs_num=0,
+      # This needs to be divisible by real_ppo_effective_num_agents.
       real_ppo_epoch_length=16*200,
       real_ppo_num_agents=1,
       real_ppo_learning_rate=2e-4,
@@ -743,7 +752,7 @@ def rl_modelrl_tiny():
   """Tiny set for testing."""
   return rl_modelrl_base_sampling().override_from_dict(
       tf.contrib.training.HParams(
-          epochs=2,
+          epochs=1,
           num_real_env_frames=128,
           simulated_env_generator_num_steps=64,
           model_train_steps=2,
@@ -753,6 +762,7 @@ def rl_modelrl_tiny():
           ppo_num_agents=2,
           real_ppo_epoch_length=36,
           real_ppo_num_agents=1,
+          real_ppo_epochs_num=0,
           real_ppo_effective_num_agents=2,
           generative_model_params="next_frame_tiny",
       ).values())

From 10758c8763d2003bb89e5545b8c81724387be3cb Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Mon, 17 Sep 2018 17:17:25 +0200
Subject: [PATCH 0854/2720] Inherit wrapped env's observ_shape in WrapperBase

---
 tensor2tensor/rl/envs/tf_atari_wrappers.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 2b293a91b..f67d30d21 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -48,6 +48,10 @@ def observ(self):
     """Access the variable holding the current observation."""
     return self._observ.read_value()
 
+  @property
+  def observ_shape(self):
+    return self._batch_env.observ_shape
+
   def __len__(self):
     """Number of combined environments."""
     return self._length

From 8f05172861247b5e7692fa89b6a583f0de7f743c Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 17 Sep 2018 13:02:11 -0700
Subject: [PATCH 0855/2720] internal merge of PR #1071

PiperOrigin-RevId: 213322906
---
 tensor2tensor/rl/trainer_model_based.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index bee68cb67..40493203b 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -553,6 +553,8 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
         ppo_event_dir, directories["world_model"], epoch_data_dir,
         hparams, epoch=epoch, is_final_epoch=is_final_epoch)
 
+    if hparams.stop_loop_early:
+      return 0.0
     # Collect data from the real environment.
     log("Generating real environment data")
     eval_data_dir = os.path.join(epoch_data_dir, "eval")
@@ -657,6 +659,7 @@ def rl_modelrl_base():
       # Whether to evaluate the world model in each iteration of the loop to get
       # the model_reward_accuracy metric.
       eval_world_model=True,
+      stop_loop_early=False,  # To speed-up tests.
   )
 
 
@@ -765,6 +768,7 @@ def rl_modelrl_tiny():
           real_ppo_epochs_num=0,
           real_ppo_effective_num_agents=2,
           generative_model_params="next_frame_tiny",
+          stop_loop_early=True,
       ).values())
 
 

From 8e1b192324889e8eef7e5c17c6440da20b216a4f Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 17 Sep 2018 13:26:15 -0700
Subject: [PATCH 0856/2720] Internal change

PiperOrigin-RevId: 213326869
---
 tensor2tensor/layers/latent_layers.py         | 29 +++++++++++++
 .../models/research/transformer_vae.py        | 42 +++++++++++++++++--
 2 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 81a156ba9..5f4184168 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -43,6 +43,35 @@ def compress_self_attention_layer(x, hparams, name):
     return tf.reshape(res, xshape)
 
 
+def compute_nats_and_bits_per_dim(data_dim,
+                                  latent_dim,
+                                  average_reconstruction,
+                                  average_prior):
+  """Computes negative ELBO, which is an upper bound on the negative likelihood.
+
+  Args:
+    data_dim: int-like indicating data dimensionality.
+    latent_dim: int-like indicating latent dimensionality.
+    average_reconstruction: Scalar Tensor indicating the reconstruction cost
+      averaged over all data dimensions and any data batches.
+    average_prior: Scalar Tensor indicating the negative log-prior probability
+      averaged over all latent dimensions and any data batches.
+
+  Returns:
+    Tuple of scalar Tensors, representing the nats and bits per data dimension
+    (e.g., subpixels) respectively.
+  """
+  with tf.name_scope(None, default_name="compute_nats_per_dim"):
+    data_dim = tf.cast(data_dim, average_reconstruction.dtype)
+    latent_dim = tf.cast(latent_dim, average_prior.dtype)
+    negative_log_likelihood = data_dim * average_reconstruction
+    negative_log_prior = latent_dim * average_prior
+    negative_elbo = negative_log_likelihood + negative_log_prior
+    nats_per_dim = tf.divide(negative_elbo, data_dim, name="nats_per_dim")
+    bits_per_dim = tf.divide(nats_per_dim, tf.log(2.), name="bits_per_dim")
+    return nats_per_dim, bits_per_dim
+
+
 def multinomial_sample(x, vocab_size=None, sampling_method="random",
                        temperature=1.0):
   """Multinomial sampling from a n-dimensional tensor.
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 24ad179d3..4a2cd1014 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -20,12 +20,14 @@
 
 import functools
 import math
+import os
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
+from tensor2tensor.layers import latent_layers
 from tensor2tensor.models import transformer
 from tensor2tensor.utils import beam_search
 from tensor2tensor.utils import expert_utils
@@ -350,7 +352,8 @@ def ae_transformer_internal(inputs,
             "neg_q_entropy": tf.constant(0.0)}
   if hparams.do_ae:
     # flatten here
-    original_targets_shape = tf.shape(targets)
+    original_targets = targets
+    original_targets_shape = tf.shape(original_targets)
     if hparams.task == "image":
       cia.maybe_reshape_4d_to_3d(targets)
     if hparams.task == "translate":
@@ -519,7 +522,10 @@ def refine_res():
   # elements. These can cause shape problems when computing loss with respect to
   # the original (unpadded) targets. So we remove their extra elements here.
   res = res[:, :original_targets_shape[1], :, :]
-  return res, losses, cache
+
+  data_dim = common_layers.shape_list(res)[1]
+  latent_dim = common_layers.shape_list(targets_c)[1]
+  return res, losses, cache, data_dim, latent_dim
 
 
 @registry.register_model
@@ -633,7 +639,7 @@ def body(self, features):
       inputs = None
     reuse = "cache_raw" in features
     with tf.variable_scope(tf.get_variable_scope(), reuse=reuse):
-      res, loss, _ = ae_transformer_internal(
+      res, loss, _, self._data_dim, self._latent_dim = ae_transformer_internal(
           inputs,
           features["targets"],
           features["target_space_id"],
@@ -653,7 +659,7 @@ def prepare_features_for_infer(self, features):
       inputs = None
     targets = tf.zeros([beam_batch_size, 1, 1, self._hparams.hidden_size])
     with tf.variable_scope("body"):
-      _, _, cache = ae_transformer_internal(
+      _, _, cache, _, _ = ae_transformer_internal(
           inputs, targets, features["target_space_id"], self._hparams)
     features["cache_raw"] = cache
 
@@ -713,6 +719,34 @@ def infer(self, features=None, decode_length=50, beam_size=1, top_beams=1,
       features["inputs"] = inputs_old
     return samples
 
+  def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
+    """Constructs `tf.estimator.EstimatorSpec` for EVAL (evaluation) mode."""
+    estimator_spec = super(TransformerAE, self).estimator_spec_eval(
+        features, logits, labels, loss, losses_dict)
+
+    summary_op = tf.get_collection(tf.GraphKeys.SUMMARIES, scope="losses")
+    summary_op.extend(tf.get_collection(tf.GraphKeys.SUMMARIES, scope="loss"))
+    summary_op.append(tf.summary.scalar("loss", loss))
+    summary_saver_hook = tf.train.SummarySaverHook(
+        save_steps=100,
+        summary_op=summary_op,
+        output_dir=os.path.join(self.hparams.model_dir, "eval"))
+
+    hooks = list(estimator_spec.evaluation_hooks)
+    hooks.append(summary_saver_hook)
+    return estimator_spec._replace(evaluation_hooks=hooks)
+
+  def _summarize_losses(self, losses_dict):
+    """Adds `tf.summary`s to all terms in the losses dictionary."""
+    super(TransformerAE, self)._summarize_losses(losses_dict)
+    nats_per_dim, bits_per_dim = latent_layers.compute_nats_and_bits_per_dim(
+        data_dim=self._data_dim,
+        latent_dim=self._latent_dim,
+        average_reconstruction=losses_dict["training"],
+        average_prior=losses_dict["latent_pred"])
+    tf.summary.scalar("loss/nats_per_dim", nats_per_dim)
+    tf.summary.scalar("loss/bits_per_dim", bits_per_dim)
+
 
 @registry.register_hparams
 def transformer_ae_small():

From c36277e3287c696114e8152a898c84ca0c65cfd7 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 17 Sep 2018 13:44:22 -0700
Subject: [PATCH 0857/2720] Allow use of custom Modalities without registering
 them.

PiperOrigin-RevId: 213330078
---
 tensor2tensor/data_generators/problem.py      | 64 +++++++++++--------
 tensor2tensor/data_generators/problem_test.py | 16 +++++
 2 files changed, 54 insertions(+), 26 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 4c64ab27d..f11389475 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -1096,14 +1096,20 @@ def _reverse_problem_hparams(p_hparams):
 
 
 def _create_modalities(problem_hparams, hparams):
-  """Converts string-type modalities to their modality object.
+  """Converts string-type modalities to their corresponding Modality.
 
   Args:
     problem_hparams: tf.contrib.training.HParams for the Problem. It must have
-      input_modality and target_modality.
+      input_modality and target_modality as attributes. Modalities are either
+      tuples of type ("modality_type:modality_name", vocab_size), and they will
+      be converted to Modality objects; or they are already Modality objects,
+      and they remain the same.
     hparams: tf.contrib.training.HParams for the model. It may have
       input_modalities and target_modality, which will override
-      problem_hparams's modalities.
+      problem_hparams' modalities.
+
+  Returns:
+    None
   """
   input_modality_overrides = {}
   if hasattr(hparams, "input_modalities"):
@@ -1115,12 +1121,15 @@ def _create_modalities(problem_hparams, hparams):
         input_modality_overrides[feature_name] = modality_name
 
   input_modality = {}
-  for f, modality_spec in six.iteritems(problem_hparams.input_modality):
-    if f in input_modality_overrides:
-      _warn_changed_modality_type(input_modality_overrides[f],
-                                  modality_spec[0], f)
-      modality_spec = (input_modality_overrides[f], modality_spec[1])
-    input_modality[f] = registry.create_modality(modality_spec, hparams)
+  for feature_name, modality in six.iteritems(problem_hparams.input_modality):
+    if isinstance(modality, (list, tuple)):
+      if feature_name in input_modality_overrides:
+        _warn_changed_modality_type(input_modality_overrides[feature_name],
+                                    modality[0],
+                                    feature_name)
+        modality = (input_modality_overrides[feature_name], modality[1])
+      modality = registry.create_modality(modality, hparams)
+    input_modality[feature_name] = modality
   problem_hparams.input_modality = input_modality
 
   target_modality_name = None
@@ -1128,26 +1137,29 @@ def _create_modalities(problem_hparams, hparams):
       hparams.target_modality != "default"):
     target_modality_name = hparams.target_modality
 
-  if problem_hparams.target_modality is None:
-    target_modality = None
-  elif isinstance(problem_hparams.target_modality, dict):
+  if isinstance(problem_hparams.target_modality, dict):
     target_modality = {}
-    for f, modality_spec in six.iteritems(problem_hparams.target_modality):
-      # TODO(lukaszkaiser): allow overriding other target modalities.
-      if target_modality_name and f == "targets":
-        _warn_changed_modality_type(target_modality_name, modality_spec[0],
-                                    "target_modality/%s" % f)
-        modality_spec = (target_modality_name, modality_spec[1])
-      target_modality[f] = registry.create_modality(modality_spec, hparams)
-  else:
-    target_modality_spec = problem_hparams.target_modality
+    for feature_name, modality in six.iteritems(
+        problem_hparams.target_modality):
+      if isinstance(modality, (list, tuple)):
+        # TODO(lukaszkaiser): allow overriding other target modalities.
+        if target_modality_name and feature_name == "targets":
+          _warn_changed_modality_type(target_modality_name,
+                                      modality[0],
+                                      "target_modality/%s" % feature_name)
+          modality = (target_modality_name, modality[1])
+        modality = registry.create_modality(modality, hparams)
+      target_modality[feature_name] = modality
+    problem_hparams.target_modality = target_modality
+  elif isinstance(problem_hparams.target_modality, (list, tuple)):
+    modality = problem_hparams.target_modality
     if target_modality_name:
       _warn_changed_modality_type(target_modality_name,
-                                  target_modality_spec[0], "target")
-      target_modality_spec = (target_modality_name, target_modality_spec[1])
-    target_modality = registry.create_modality(target_modality_spec,
-                                               hparams)
-  problem_hparams.target_modality = target_modality
+                                  modality[0],
+                                  "target")
+      modality = (target_modality_name, modality[1])
+    modality = registry.create_modality(modality, hparams)
+    problem_hparams.target_modality = modality
 
 
 def _warn_changed_modality_type(new_name, old_name, feature_name):
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 36b1ecd59..4b6ccd4c1 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -88,6 +88,22 @@ def testProblemHparamsModality(self):
                           modalities.SymbolModality)
     self.assertIsInstance(p_hparams.target_modality, modalities.SymbolModality)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testProblemHparamsModalityObj(self):
+    class ModalityObjProblem(problem_module.Problem):
+
+      def hparams(self, defaults, model_hparams):
+        hp = defaults
+        hp.input_modality = {
+            "inputs": modalities.SymbolModality(model_hparams, 2)}
+        hp.target_modality = modalities.SymbolModality(model_hparams, 3)
+
+    problem = ModalityObjProblem(False, False)
+    p_hparams = problem.get_hparams()
+    self.assertIsInstance(p_hparams.input_modality["inputs"],
+                          modalities.SymbolModality)
+    self.assertIsInstance(p_hparams.target_modality, modalities.SymbolModality)
+
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testProblemHparamsInputOnlyModality(self):
     class InputOnlyProblem(problem_module.Problem):

From 7d17d2875f0229f97b6bfc35918567aa22850632 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 17 Sep 2018 13:44:42 -0700
Subject: [PATCH 0858/2720] Fix MultiProblem's target modality.

PiperOrigin-RevId: 213330135
---
 tensor2tensor/data_generators/multi_problem.py      | 2 +-
 tensor2tensor/data_generators/multi_problem_test.py | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index f8d32cc71..97d046d64 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -89,7 +89,7 @@ def get_hparams(self, model_hparams=None):
     tf.logging.info("Old vocabulary size: %d" % vocab_size)
     tf.logging.info("New vocabulary size: %d" % (vocab_size + vocab_size_inc))
     self._hparams.target_modality = modalities.SymbolModality(
-        vocab_size + vocab_size_inc)
+        model_hparams, vocab_size + vocab_size_inc)
 
     return self._hparams
 
diff --git a/tensor2tensor/data_generators/multi_problem_test.py b/tensor2tensor/data_generators/multi_problem_test.py
index efd76fc0c..e1f38b8d4 100644
--- a/tensor2tensor/data_generators/multi_problem_test.py
+++ b/tensor2tensor/data_generators/multi_problem_test.py
@@ -31,7 +31,7 @@ class TestMultiProblem(multi_problem.MultiProblem):
   def __init__(self):
     super(TestMultiProblem, self).__init__()
     self.task_list.append(problem_hparams.TestProblem(2, 3))
-    self.task_list.append(problem_hparams.TestProblem(4, 5))
+    self.task_list.append(problem_hparams.TestProblem(4, 6))
 
 
 class MultiProblemTest(tf.test.TestCase):
@@ -42,7 +42,9 @@ def testProblemHparamsModality(self):
     p_hparams = problem.get_hparams()
     self.assertIsInstance(p_hparams.input_modality["inputs"],
                           modalities.SymbolModality)
+    self.assertEqual(p_hparams.input_modality["inputs"].top_dimensionality, 3)
     self.assertIsInstance(p_hparams.target_modality, modalities.SymbolModality)
+    self.assertEqual(p_hparams.target_modality.top_dimensionality, 5)
 
 if __name__ == "__main__":
   tf.test.main()

From 4b6b7758850a5b2c202ced78851d0c39e15f8c61 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 17 Sep 2018 14:56:14 -0700
Subject: [PATCH 0859/2720] Use memory instead of files to keep PPO frames.

PiperOrigin-RevId: 213343257
---
 tensor2tensor/data_generators/gym_problems.py | 13 ++++-
 tensor2tensor/data_generators/problem.py      | 10 ++--
 tensor2tensor/layers/common_layers.py         |  6 +--
 .../models/video/basic_deterministic.py       |  2 +
 .../models/video/basic_stochastic.py          | 33 +++++++-----
 tensor2tensor/rl/trainer_model_based.py       | 51 ++++++++++++++-----
 6 files changed, 77 insertions(+), 38 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 2097d19e8..d56d389b1 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -80,6 +80,12 @@ def standard_atari_ae_env_spec(env):
                                      simulated_env=False)
 
 
+frame_dumper_use_disk = False  # Whether to use memory or disk to dump frames.
+
+
+frame_dumper = {}
+
+
 class GymDiscreteProblem(video_utils.VideoProblem):
   """Gym environment with discrete actions and rewards."""
 
@@ -148,8 +154,11 @@ def _get_data(self):
     if self._use_dumper_data:
       file_path = os.path.join(self._dumper_path,
                                "frame_{}.npz".format(self._dumper_data_index))
-      with tf.gfile.Open(file_path) as gfile:
-        data = np.load(gfile)
+      if frame_dumper_use_disk:
+        with tf.gfile.Open(file_path) as gfile:
+          data = np.load(gfile)
+      else:
+        data = frame_dumper.pop(file_path)
       self._dumper_data_index += 1
       return (data["observ"][0, ...], data["reward"][0], data["done"][0],
               data["action"][0])
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index f11389475..657cbae2c 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -723,15 +723,13 @@ def feature_info(self):
 
     features = collections.defaultdict(FeatureInfo)
 
-    for name, mod_spec in six.iteritems(input_mods):
-      mod, vocab_size = mod_spec
+    for name, mod in six.iteritems(input_mods):
       finfo = features[name]
       finfo.modality = mod
-      finfo.vocab_size = vocab_size
+      finfo.vocab_size = mod.top_dimensionality
 
-    mod, vocab_size = target_mod
-    features["targets"].modality = mod
-    features["targets"].vocab_size = vocab_size
+    features["targets"].modality = target_mod
+    features["targets"].vocab_size = target_mod.top_dimensionality
 
     for name, encoder in six.iteritems(vocabs):
       features[name].encoder = encoder
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 2eccea693..a875d1f22 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3793,8 +3793,8 @@ def mean_with_attention(x, name, num_heads=4):
                            2 * shape[-1], name="mean_attn_final")
 
 
-def single_discriminator(x, filters=128, kernel_size=7,
-                         strides=4, pure_mean=True):
+def single_discriminator(x, filters=128, kernel_size=8,
+                         strides=4, pure_mean=False):
   """A simple single-layer convolutional discriminator."""
   with tf.variable_scope("discriminator"):
     net = tf.layers.conv2d(
@@ -3807,7 +3807,7 @@ def single_discriminator(x, filters=128, kernel_size=7,
 
 
 def double_discriminator(x, filters1=128, filters2=None,
-                         kernel_size=7, strides=4, pure_mean=True):
+                         kernel_size=8, strides=4, pure_mean=False):
   """A convolutional discriminator with 2 layers and concatenated output."""
   if filters2 is None:
     filters2 = 4 * filters1
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 11df93cf9..00783b3ab 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -187,6 +187,8 @@ def body(self, features):
     for i in range(hparams.video_num_target_frames):
       cur_frames = all_frames[i:i + hparams.video_num_input_frames]
       features["inputs"] = tf.concat(cur_frames, axis=-1)
+      features["cur_target_frame"] = all_frames[
+          i + hparams.video_num_input_frames]
       if "input_action" in features:
         cur_actions = all_actions[i:i + hparams.video_num_input_frames]
         features["input_action"] = tf.concat(cur_actions, axis=1)
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index a1c5e2847..85b361665 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -79,20 +79,24 @@ def inject_latent(self, layer, features, filters):
 
     # Embed.
     x = tf.layers.dense(
-        features["targets"], filters, name="latent_embed",
+        features["cur_target_frame"], filters, name="latent_embed",
         bias_initializer=tf.random_normal_initializer(stddev=0.01))
     x = common_attention.add_timing_signal_nd(x)
 
-    for i in range(hparams.num_compress_steps):
-      with tf.variable_scope("latent_downstride%d" % i):
-        x = common_layers.make_even_size(x)
-        if i < hparams.filter_double_steps:
-          filters *= 2
-        x = common_attention.add_timing_signal_nd(x)
-        x = tf.layers.conv2d(x, filters, kernel, activation=common_layers.belu,
-                             strides=(2, 2), padding="SAME")
-        x = common_layers.layer_norm(x)
-
+    if hparams.full_latent_tower:
+      for i in range(hparams.num_compress_steps):
+        with tf.variable_scope("latent_downstride%d" % i):
+          x = common_layers.make_even_size(x)
+          if i < hparams.filter_double_steps:
+            filters *= 2
+          x = common_attention.add_timing_signal_nd(x)
+          x = tf.layers.conv2d(x, filters, kernel,
+                               activation=common_layers.belu,
+                               strides=(2, 2), padding="SAME")
+          x = common_layers.layer_norm(x)
+    else:
+      x = common_layers.double_discriminator(x)
+      x = tf.expand_dims(tf.expand_dims(x, axis=1), axis=1)
     x = tf.tanh(tf.layers.dense(x, hparams.bottleneck_bits, name="bottleneck"))
     d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
     if hparams.mode == tf.estimator.ModeKeys.TRAIN:
@@ -127,9 +131,10 @@ def next_frame_basic_stochastic():
 @registry.register_hparams
 def next_frame_basic_stochastic_discrete():
   """Basic 2-frame conv model with stochastic discrete latent."""
-  hparams = basic_deterministic_params.next_frame_basic_deterministic()
+  hparams = basic_deterministic_params.next_frame_sampling()
   hparams.num_compress_steps = 8
   hparams.filter_double_steps = 3
-  hparams.add_hparam("bottleneck_bits", 32)
-  hparams.add_hparam("bottleneck_noise", 0.05)
+  hparams.add_hparam("bottleneck_bits", 16)
+  hparams.add_hparam("bottleneck_noise", 0.02)
+  hparams.add_hparam("full_latent_tower", False)
   return hparams
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 40493203b..2a6eb284c 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -31,11 +31,14 @@
 import math
 import os
 import time
+
 import numpy as np
 import six
 
+
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import gym_problems
 from tensor2tensor.data_generators import gym_problems_specs
 from tensor2tensor.layers import discretization
 from tensor2tensor.rl import rl_trainer_lib
@@ -218,13 +221,17 @@ def ppo_data_dumper(observ, reward, done, action):
   """Save frames from PPO to a numpy file."""
   global ppo_data_dumper_counter, dumper_path
   file_path = "{}/frame_{}.npz".format(dumper_path, ppo_data_dumper_counter)
-  # np.savez_compressed can't create a tf.gfile, so we need to creat it
-  # beforehand.
-  with tf.gfile.Open(file_path, mode="wb+") as gfile:
-    gfile.write("1")
-  with tf.gfile.Open(file_path, mode="wb+") as gfile:
-    np.savez_compressed(gfile,
-                        observ=observ, reward=reward, done=done, action=action)
+  if gym_problems.frame_dumper_use_disk:
+    # np.savez_compressed can't create a tf.gfile, so we need to create it
+    # beforehand.
+    with tf.gfile.Open(file_path, mode="wb+") as gfile:
+      gfile.write("1")
+    with tf.gfile.Open(file_path, mode="wb+") as gfile:
+      np.savez_compressed(
+          gfile, observ=observ, reward=reward, done=done, action=action)
+  else:
+    data = {"observ": observ, "reward": reward, "done": done, "action": action}
+    gym_problems.frame_dumper[file_path] = data
   ppo_data_dumper_counter += 1
   return 0.0
 
@@ -621,7 +628,7 @@ def rl_modelrl_base():
       # hparams.epochs.
       # This number should be divisible by real_ppo_epoch_length*epochs
       # for our frame accounting to be preceise.
-      num_real_env_frames=96000,
+      num_real_env_frames=96000 * 2,
       generative_model="next_frame_basic_deterministic",
       generative_model_params="next_frame_pixel_noise",
       ppo_params="ppo_pong_base",
@@ -699,6 +706,15 @@ def rl_modelrl_base_stochastic():
   return hparams
 
 
+@registry.register_hparams
+def rl_modelrl_base_stochastic_discrete():
+  """Base setting with stochastic discrete model."""
+  hparams = rl_modelrl_base()
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  return hparams
+
+
 @registry.register_hparams
 def rl_modelrl_base_sv2p():
   """Base setting with sv2p as world model."""
@@ -869,7 +885,7 @@ def rl_modelrl_tiny_simulation_deterministic_starts():
 def rl_modelrl_grid(rhp):
   """Grid over games and frames, and 5 runs each for variance."""
   rhp.set_categorical("loop.game",
-                      ["breakout", "wrapped_long_pong", "freeway"])
+                      ["breakout", "wrapped_full_pong", "freeway"])
 
   base = 100000
   medium = base // 2
@@ -885,7 +901,7 @@ def rl_modelrl_variance(rhp):
   # Dummy parameter to get 5 runs for each configuration
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
   rhp.set_categorical("loop.game",
-                      ["breakout", "wrapped_long_pong", "freeway"])
+                      ["breakout", "wrapped_full_pong", "freeway"])
 
 
 @registry.register_ranged_hparams
@@ -894,6 +910,15 @@ def rl_modelrl_variance_nogame(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(500)))
 
 
+@registry.register_ranged_hparams
+def rl_modelrl_test1(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", ["breakout", "wrapped_full_pong", "boxing"])
+  rhp.set_discrete("loop.ppo_learning_rate", [1e-4, 2e-4, 4e-4])
+  rhp.set_discrete("ppo.optimization_batch_size", [20, 40])
+  rhp.set_discrete("loop.epochs", [3, 6])
+
+
 @registry.register_ranged_hparams
 def rl_modelrl_scheduled_sampling(rhp):
   rhp.set_float("model.scheduled_sampling_prob", 0.0, 1.0)
@@ -916,7 +941,7 @@ def rl_modelrl_ae_variance(rhp):
   # Dummy parameter to get 5 runs for each configuration
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
   rhp.set_categorical("loop.game",
-                      ["breakout", "wrapped_long_pong", "freeway"])
+                      ["breakout", "wrapped_full_pong", "freeway"])
   base = 100000
   small = base // 4
   rhp.set_discrete("loop.num_real_env_frames", [base, small])
@@ -925,7 +950,7 @@ def rl_modelrl_ae_variance(rhp):
 @registry.register_ranged_hparams
 def rl_modelrl_ppolr_game(rhp):
   rhp.set_categorical("loop.game",
-                      ["breakout", "wrapped_long_pong", "freeway"])
+                      ["breakout", "wrapped_full_pong", "freeway"])
   base_lr = 2e-4
   rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
 
@@ -939,7 +964,7 @@ def rl_modelrl_ppolr(rhp):
 @registry.register_ranged_hparams
 def rl_modelrl_ae_ppo_lr(rhp):
   rhp.set_categorical("loop.game",
-                      ["breakout", "wrapped_long_pong", "freeway"])
+                      ["breakout", "wrapped_full_pong", "freeway"])
   base_lr = 2e-4
   rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
 

From 38c062e727e875df57b96c87e64aa5ed8b91ea69 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 17 Sep 2018 15:15:16 -0700
Subject: [PATCH 0860/2720] Common layers cleanup (first small step): remove
 some unused functions.

PiperOrigin-RevId: 213347128
---
 tensor2tensor/layers/common_attention.py   |   6 +-
 tensor2tensor/layers/common_layers.py      | 361 +--------------------
 tensor2tensor/layers/common_layers_test.py |  42 ---
 3 files changed, 7 insertions(+), 402 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 4266d9787..cfe7e20d3 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -4675,8 +4675,7 @@ def multihead_self_attention_memory_efficient(x,
 
   def forward_internal(x, wqkv, wo, attention_bias, norm_scale, norm_bias):
     """Forward function."""
-    n = common_layers.layer_norm_compute_python(x, epsilon, norm_scale,
-                                                norm_bias)
+    n = common_layers.layer_norm_compute(x, epsilon, norm_scale, norm_bias)
     wqkv_split = tf.unstack(wqkv, num=num_heads)
     wo_split = tf.unstack(wo, num=num_heads)
     y = 0
@@ -4700,8 +4699,7 @@ def forward_internal(x, wqkv, wo, attention_bias, norm_scale, norm_bias):
     def grad_fn(x, wqkv, wo, attention_bias, norm_scale, norm_bias, dy):
       """Custom gradient function."""
       with tf.control_dependencies([dy]):
-        n = common_layers.layer_norm_compute_python(x, epsilon, norm_scale,
-                                                    norm_bias)
+        n = common_layers.layer_norm_compute(x, epsilon, norm_scale, norm_bias)
         wqkv_split = tf.unstack(wqkv, num=num_heads)
         wo_split = tf.unstack(wo, num=num_heads)
         deps = []
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index a875d1f22..a7f3c99a8 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -34,9 +34,6 @@
 from tensorflow.python.ops import control_flow_util
 from tensorflow.python.ops import inplace_ops
 
-# This is a global setting. When turned off, no @function.Defun is used.
-allow_defun = False
-
 
 @function.Defun(
     python_grad_func=lambda x, dy: tf.convert_to_tensor(dy),
@@ -633,7 +630,7 @@ def layer_norm_vars(filters):
   return scale, bias
 
 
-def layer_norm_compute_python(x, epsilon, scale, bias):
+def layer_norm_compute(x, epsilon, scale, bias):
   """Layer norm raw computation."""
   epsilon, scale, bias = [cast_like(t, x) for t in [epsilon, scale, bias]]
   mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
@@ -642,37 +639,14 @@ def layer_norm_compute_python(x, epsilon, scale, bias):
   return norm_x * scale + bias
 
 
-@function.Defun(compiled=True)
-def layer_norm_compute_grad(x, epsilon, scale, bias, dy):
-  y = layer_norm_compute_python(x, epsilon, scale, bias)
-  dx = tf.gradients(ys=[y], xs=[x, epsilon, scale, bias], grad_ys=[dy])
-  return dx
-
-
-@function.Defun(
-    compiled=True,
-    separate_compiled_gradients=True,
-    grad_func=layer_norm_compute_grad)
-def layer_norm_compute(x, epsilon, scale, bias):
-  return layer_norm_compute_python(x, epsilon, scale, bias)
-
-
 def layer_norm(x, filters=None, epsilon=1e-6, name=None, reuse=None):
   """Layer normalize the tensor x, averaging over the last dimension."""
   if filters is None:
     filters = shape_list(x)[-1]
   with tf.variable_scope(
       name, default_name="layer_norm", values=[x], reuse=reuse):
-    scale = tf.get_variable(
-        "layer_norm_scale", [filters], initializer=tf.ones_initializer())
-    bias = tf.get_variable(
-        "layer_norm_bias", [filters], initializer=tf.zeros_initializer())
-    if allow_defun:
-      result = layer_norm_compute(x, tf.constant(epsilon), scale, bias)
-      result.set_shape(x.get_shape())
-    else:
-      result = layer_norm_compute_python(x, epsilon, scale, bias)
-    return result
+    scale, bias = layer_norm_vars(filters)
+    return layer_norm_compute(x, epsilon, scale, bias)
 
 
 def group_norm(x, filters=None, num_groups=8, epsilon=1e-5):
@@ -1103,240 +1077,6 @@ def conv_block_downsample(x,
     return x
 
 
-def decompress_seqcnn(x,
-                      targets,
-                      targets_vocab_size,
-                      dilations_and_kernels,
-                      block_size,
-                      is_2d=False,
-                      embedding_var=None,
-                      name=None,
-                      reuse=None):
-  """Decompress x into targets size using a Sequence CNN at every element."""
-  with tf.variable_scope(
-      name,
-      default_name="decompress_batch_seqcnn",
-      values=[x, targets],
-      reuse=reuse):
-    # We assume targets are [batch x block_size * N x block_size * N x C] if
-    # is_2d=True or [batch, block_size * N, 1, C] otherwise, and C is static.
-    # Let's shift targets to depth and embed.
-    targets_shape = shape_list(targets)
-    channels = targets_shape[-1]
-    hidden_size = x.get_shape()[-1]
-    if is_2d:
-      depth_targets = tf.space_to_depth(targets, block_size)
-      factor = channels * block_size * block_size
-    else:
-      depth_targets = tf.reshape(targets, [
-          targets_shape[0], targets_shape[1] // block_size, 1,
-          channels * block_size
-      ])
-      factor = channels * block_size
-    if embedding_var is None:
-      embedding_var = tf.get_variable("targets_embedding",
-                                      [targets_vocab_size, hidden_size])
-    targets_emb = tf.gather(embedding_var, depth_targets)
-    # Flatten x and embedded targets. Flat targets are factor* larger on axis=1.
-    flat_x = tf.reshape(x, [-1, 1, 1, hidden_size])
-    flat_targets = tf.reshape(targets_emb, [-1, factor, 1, hidden_size])
-    shifted_targets = shift_right(flat_targets)
-    # Run a SeqCNN large-batch to produce factor outputs out of every target.
-    flat_x += tf.zeros_like(shifted_targets)  # Broadcast on axis=1.
-    flat_outputs = conv_block(
-        tf.concat([flat_x, shifted_targets], axis=3),
-        hidden_size,
-        dilations_and_kernels,
-        padding="LEFT")
-    # Reshape back to embedded targets shape.
-    targets_emb_shape = shape_list(targets_emb)
-    outputs = tf.reshape(flat_outputs, [
-        targets_emb_shape[0], targets_emb_shape[1], targets_emb_shape[2],
-        factor * hidden_size
-    ])
-    # Move depth back to target space.
-    if is_2d:
-      outputs = tf.depth_to_space(outputs, 2)
-    else:
-      outputs = tf.reshape(outputs, [
-          shape_list(outputs)[0], block_size * shape_list(outputs)[1], 1,
-          hidden_size
-      ])
-    # Final reshape before prediction to ensure target size.
-    outputs = tf.reshape(outputs, [
-        targets_shape[0], targets_shape[1], targets_shape[2], channels,
-        hidden_size
-    ])
-    return dense(outputs, targets_vocab_size)
-
-
-def simple_attention(target, source, bias=None):
-  """A simple attention function.
-
-  Args:
-    target: a `Tensor` with shape `[batch, target_timesteps, depth]` or
-     `[batch, target_timesteps_1, target_timesteps_2, depth]`
-    source: a `Tensor` with shape `[batch, source_timesteps, depth]` or
-     `[batch, source_timesteps_1, source_timesteps_2, depth]`
-    bias: an optional `Tensor` with shape `[batch, timesteps, 1, 1]` used
-     to mask the attention to not attend to padding of input.
-
-  Returns:
-    a `Tensor` with same shape as `target`
-  """
-  with tf.name_scope("simple_attention", values=[target, source]):
-    target_shape = shape_list(target)
-    source_shape = shape_list(source)
-    target = tf.reshape(
-        target,
-        [target_shape[0], target_shape[1] * target_shape[2], target_shape[3]])
-    source = tf.reshape(
-        source,
-        [source_shape[0], source_shape[1] * source_shape[2], source_shape[3]])
-    attention = tf.matmul(target, source, transpose_b=True)
-    attention *= tf.rsqrt(tf.to_float(shape_list(target)[2]))
-    if bias is not None:
-      attention += tf.expand_dims(tf.squeeze(bias, axis=[2, 3]), axis=1)
-    attention = tf.nn.softmax(attention)
-    if should_generate_summaries():
-      tf.summary.image("attention", tf.expand_dims(attention, 3), max_outputs=5)
-    attended = tf.matmul(attention, source)
-    return tf.reshape(attended, target_shape)
-
-
-def multiscale_conv_sum(inputs, output_size, dilation_rates_and_kernel_sizes,
-                        pooling_type, **kwargs):
-  """Sum of several dilated convolutions.
-
-  For all convolutions with dilation_rate > 1, we first pool the input with
-  width dilation_rate.
-
-  Args:
-    inputs: a Tensor
-    output_size: an Integer
-    dilation_rates_and_kernel_sizes: a list of pairs (dilation, kernel_size)
-    pooling_type: "AVG" or "MAX"
-    **kwargs: additional
-
-  Returns:
-     a Tensor.
-  """
-  name = kwargs.pop("name") if "name" in kwargs else None
-  with tf.variable_scope(name, "multiscale_conv_sum", [inputs]):
-    padding = kwargs["padding"]
-    results, counter = [], -1
-    for dilation_rate, kernel_size in dilation_rates_and_kernel_sizes:
-      counter += 1
-      if dilation_rate[0] > 1:
-        pooled = pool(inputs, kernel_size, pooling_type, padding)
-      else:
-        pooled = inputs
-      results.append(
-          conv(
-              pooled,
-              output_size,
-              kernel_size,
-              dilation_rate=dilation_rate,
-              name="conv_layer%d" % counter,
-              **kwargs))
-    return tf.add_n(results) * (len(results)**-0.5)
-
-
-def multiscale_conv_and_attention(x, padding, hparams, source=None):
-  """A common part of t2t layers.
-
-  First, do a linear multiscale convolution
-  Second, do attention (if source is not None)
-
-  Applies residuals and normalization on both steps.
-
-  Args:
-    x: a Tensor.
-    padding: a padding type
-    hparams: hyperparameters for model
-    source: optional source tensor for attention. (encoder output)
-
-  Returns:
-    a Tensor.
-  """
-  # TODO(noam): The number of different scales should be a hyperparameter.
-  conv_sum = multiscale_conv_sum(
-      x,
-      hparams.hidden_size,
-      [((hparams.kernel_height**i, hparams.kernel_width**i),
-        (hparams.kernel_height, hparams.kernel_width)) for i in range(3)],
-      "AVG",
-      padding=padding)
-  # For residuals a rescale if necessary if channels differ.
-  if x.get_shape().as_list()[-1] != conv_sum.get_shape().as_list()[-1]:
-    x = conv(x, hparams.hidden_size, (1, 1))
-  x = noam_norm(x + conv_sum)
-  if source is not None:
-    x = noam_norm(x + simple_attention(x, source))
-  return x
-
-
-def conv_with_pools(inputs, output_size, kernel_size, pool_sizes, pooling_type,
-                    **kwargs):
-  """Convolution plus 1x1 convolution applied to specified pools.
-
-  For example we might do a regular convolution with kernel size (3, 1),
-  and pools of sizes [(9, 1), (27, 1)].
-
-  Args:
-    inputs: a Tensor
-    output_size: an Integer
-    kernel_size: a tuple of integers
-    pool_sizes: a list of tuples of integers.
-    pooling_type: "AVG" or "MAX"
-    **kwargs: additional keyword args for conv
-
-  Returns:
-     a Tensor.
-  """
-  name = kwargs.pop("name") if "name" in kwargs else None
-  with tf.variable_scope(name, "conv_with_pools", [inputs]):
-    padding = kwargs["padding"]
-    results = []
-    results.append(conv(inputs, output_size, kernel_size, **kwargs))
-    for i, pool_size in enumerate(pool_sizes):
-      pooled = pool(inputs, pool_size, pooling_type, padding)
-      results.append(
-          conv(pooled, output_size, (1, 1), name="pool_%d" % i, **kwargs))
-    return tf.add_n(results) * (len(results)**-0.5)
-
-
-def conv_with_pools_and_attention(x, padding, hparams, source=None):
-  """A common part of t2t layers.
-
-  First, do conv_with_pools
-  Second, do attention (if source is not None)
-
-  Applies residuals and normalization on both steps.
-
-  Args:
-    x: a Tensor.
-    padding: a padding type
-    hparams: hyperparameters for model
-    source: optional source tensor for attention. (encoder output)
-
-  Returns:
-    a Tensor.
-  """
-  conv_sum = conv_with_pools(
-      x,
-      hparams.hidden_size, (hparams.kernel_height, hparams.kernel_width),
-      hparams.pool_sizes,
-      "AVG",
-      padding=padding)
-  if x.get_shape().as_list()[-1] == conv_sum.get_shape().as_list()[-1]:
-    conv_sum += x
-  x = noam_norm(conv_sum)
-  if source is not None:
-    x = noam_norm(x + simple_attention(x, source))
-  return x
-
-
 def get_timing_signal(length,
                       min_timescale=1,
                       max_timescale=1e4,
@@ -1437,97 +1177,6 @@ def mask_leq(target_length, source_length):
       out_shape=[1, target_length, source_length])
 
 
-def attention_1d_v0(source,
-                    target,
-                    attention_size,
-                    output_size,
-                    num_heads,
-                    mask=None,
-                    transform_source=True,
-                    transform_target=True,
-                    transform_output=True,
-                    name=None):
-  """multi-headed attention.
-
-  TODO(noam): this could probably be extended to 2d.
-
-  Args:
-    source: a Tensor of shape [batch, source_length, source_depth]
-    target: a Tensor of shape [batch, target_length, target_depth]
-    attention_size: an integer
-    output_size: an integer
-    num_heads: an integer divisor of attention_size
-    mask: a float32 Tensor of shape [batch, target_length, source_length]
-          1.0 means can-see; 0.0 means can't-see.
-          Any dimension can be 1 (supports broadcasting).
-    transform_source: a boolean
-    transform_target: a boolean
-    transform_output: a boolean
-    name: an optional string
-
-  Returns:
-    a Tensor of shape [batch, length, output_size]
-  """
-  with tf.variable_scope(name, default_name="attention", values=[target]):
-    source_shape = shape_list(source)
-    source_length = source_shape[1]
-    target_length = shape_list(target)[1]
-    batch = source_shape[0]
-
-    def _maybe_transform(t, size, should_transform, name):
-      if should_transform:
-        return conv1d(t, size, 1, name=name)
-      else:
-        assert t.get_shape()[-1] == size
-        return t
-
-    source_attention = _maybe_transform(source, attention_size,
-                                        transform_source, "source_attention")
-    target_attention = _maybe_transform(target, attention_size,
-                                        transform_target, "target_attention")
-    assert attention_size % num_heads == 0
-    size_per_head = attention_size // num_heads
-    source_attention = tf.reshape(
-        source_attention, [batch, source_length, num_heads, size_per_head])
-    target_attention = tf.reshape(
-        target_attention, [batch, target_length, num_heads, size_per_head])
-    # [batch, num_heads, length, size_per_head]
-    source_attention = tf.transpose(source_attention, [0, 2, 1, 3])
-    target_attention = tf.transpose(target_attention, [0, 2, 1, 3])
-
-    # [batch, num_heads, target_length, source_length]
-    attention = tf.matmul(target_attention, source_attention, transpose_b=True)
-    attention *= size_per_head**-0.5
-
-    if mask is not None:
-      mask = tf.expand_dims(mask, 1)
-      mask = (1.0 - mask) * -1e9
-      attention += mask
-    attention = tf.nn.softmax(attention)
-    if should_generate_summaries():
-      # Compute a color image summary.
-      image = tf.reshape(attention,
-                         [batch, num_heads, target_length, source_length])
-      image = tf.transpose(image, [0, 2, 3, 1])
-      image = tf.pow(image, 0.2)  # for high-dynamic-range
-      # Each head will correspond to one of RGB.
-      # pad the heads to be a multiple of 3
-      extra_heads = -num_heads % 3
-      image = tf.pad(image, [[0, 0], [0, 0], [0, 0], [0, -num_heads % 3]])
-      image = tf.reshape(image, [
-          batch, target_length, source_length, 3, (num_heads + extra_heads) // 3
-      ])
-      image = tf.reduce_max(image, 4)
-      tf.summary.image("local_attention", image, max_outputs=1)
-    # output: [batch, num_heads, target_length, size_per_head]
-    output = tf.matmul(attention, source_attention)
-    output = tf.transpose(output, [0, 2, 1, 3])
-    output = tf.reshape(output, [batch, target_length, attention_size])
-    output = _maybe_transform(output, output_size, transform_output,
-                              "attention_output")
-    return output
-
-
 def relu_density_logit(x, reduce_dims):
   """logit(density(x)).
 
@@ -3043,7 +2692,7 @@ def forward_internal(x, f1, f2, scale, bias):
     ys = []
     for i in range(num_splits):
       with tf.control_dependencies(ys[-1:]):
-        n = layer_norm_compute_python(xs[i], epsilon, scale, bias)
+        n = layer_norm_compute(xs[i], epsilon, scale, bias)
         y = tf.nn.conv1d(n, f1, 1, "SAME")
         y = tf.nn.relu(y)
         y = tf.nn.conv1d(y, f2, 1, "SAME")
@@ -3078,7 +2727,7 @@ def grad_fn(x, f1, f2, scale, bias, dy):
         deps = []
         for i in range(num_splits):
           with tf.control_dependencies(deps):
-            n = layer_norm_compute_python(xs[i], epsilon, scale, bias)
+            n = layer_norm_compute(xs[i], epsilon, scale, bias)
             y = tf.nn.conv1d(n, f1, 1, "SAME")
             y = tf.nn.relu(y)
             y = tf.nn.conv1d(y, f2, 1, "SAME")
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index c7031b2e5..01247db34 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -160,15 +160,6 @@ def testConvBlockDownsample(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 4, 1, 27))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
-  def testSimpleAttention(self):
-    x = np.random.rand(5, 7, 1, 11)
-    y = np.random.rand(5, 9, 1, 11)
-    a = common_layers.simple_attention(
-        tf.constant(x, dtype=tf.float32), tf.constant(y, dtype=tf.float32))
-    res = self.evaluate(a)
-    self.assertEqual(res.shape, (5, 7, 1, 11))
-
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testGetTimingSignal(self):
     length = 7
@@ -188,39 +179,6 @@ def testAddTimingSignal(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (batch, length, height, depth))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
-  def testAttention1D(self):
-    batch = 5
-    target_length = 7
-    source_length = 13
-    source_depth = 9
-    target_depth = 11
-    attention_size = 21
-    output_size = 15
-    num_heads = 7
-    source = np.random.rand(batch, source_length, source_depth)
-    target = np.random.rand(batch, target_length, target_depth)
-    mask = np.random.rand(batch, target_length, source_length)
-    a = common_layers.attention_1d_v0(
-        tf.constant(source, dtype=tf.float32),
-        tf.constant(target, dtype=tf.float32), attention_size, output_size,
-        num_heads, tf.constant(mask, dtype=tf.float32))
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(a)
-    self.assertEqual(res.shape, (batch, target_length, output_size))
-
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
-  def testMultiscaleConvSum(self):
-    x = np.random.rand(5, 9, 1, 11)
-    y = common_layers.multiscale_conv_sum(
-        tf.constant(x, dtype=tf.float32),
-        13, [((1, 1), (5, 5)), ((2, 2), (3, 3))],
-        "AVG",
-        padding="SAME")
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(y)
-    self.assertEqual(res.shape, (5, 9, 1, 13))
-
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testConvGRU(self):
     x = np.random.rand(5, 7, 3, 11)

From 450d471915b602786b050be00732eb06700664f7 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Mon, 17 Sep 2018 15:24:57 -0700
Subject: [PATCH 0861/2720] Fix another bug with spatial partitioning.

PiperOrigin-RevId: 213348919
---
 tensor2tensor/mesh_tensorflow/mesh_tensorflow.py |  4 ++++
 tensor2tensor/mesh_tensorflow/mnist.py           | 15 +++++++--------
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index 3104487af..c2ed10c09 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -2817,6 +2817,10 @@ def reshape(x, new_shape):
   return ReshapeOperation(x, convert_to_shape(new_shape)).outputs[0]
 
 
+def transpose(x, new_shape):
+  return einsum([x], output_shape=convert_to_shape(new_shape))
+
+
 def rename_dimension(x, old_name, new_name):
   """Reshape a Tensor, renaming one dimension.
 
diff --git a/tensor2tensor/mesh_tensorflow/mnist.py b/tensor2tensor/mesh_tensorflow/mnist.py
index 3f226b218..d02fa164b 100644
--- a/tensor2tensor/mesh_tensorflow/mnist.py
+++ b/tensor2tensor/mesh_tensorflow/mnist.py
@@ -43,7 +43,7 @@
                         "Total number of evaluation steps. If `0`, evaluation "
                         "after training is skipped.")
 tf.flags.DEFINE_string("mesh_shape", "b1:2;b2:2", "mesh shape")
-tf.flags.DEFINE_string("layout", "col_blocks:b1;hidden1:b2;filters2:b2",
+tf.flags.DEFINE_string("layout", "row_blocks:b1;col_blocks:b2",
                        "layout rules")
 
 FLAGS = tf.flags.FLAGS
@@ -71,16 +71,15 @@ def mnist_model(image, labels, mesh):
   one_channel_dim = mtf.Dimension("one_channel", 1)
 
   x = mtf.import_tf_tensor(
-      mesh, tf.reshape(image, [FLAGS.batch_size, 4, 7, 4, 7]),
+      mesh, tf.reshape(image, [FLAGS.batch_size, 4, 7, 4, 7, 1]),
       mtf.Shape(
-          [batch_dim, row_blocks_dim, rows_dim, col_blocks_dim, cols_dim]))
-  x = mtf.reshape(x, [
+          [batch_dim, row_blocks_dim, rows_dim,
+           col_blocks_dim, cols_dim, one_channel_dim]))
+  x = mtf.transpose(x, [
       batch_dim, row_blocks_dim, col_blocks_dim,
       rows_dim, cols_dim, one_channel_dim])
 
   # add some convolutional layers to demonstrate that convolution works.
-  # TODO(nikip): Currently spatial conv works only when splitting column blocks.
-  # Make it work for both height and width dimension of the image.
   fh_dim = mtf.Dimension("fh", 9)
   fw_dim = mtf.Dimension("fw", 9)
   filters1_dim = mtf.Dimension("filters1", 16)
@@ -92,10 +91,10 @@ def mnist_model(image, labels, mesh):
 
   f1 = mtf.relu(mtf.conv2d_with_blocks(
       x, kernel1, strides=[1, 1, 1, 1], padding="SAME",
-      h_blocks_dim=None, w_blocks_dim=col_blocks_dim))
+      h_blocks_dim=row_blocks_dim, w_blocks_dim=col_blocks_dim))
   f2 = mtf.relu(mtf.conv2d_with_blocks(
       f1, kernel2, strides=[1, 1, 1, 1], padding="SAME",
-      h_blocks_dim=None, w_blocks_dim=None))
+      h_blocks_dim=row_blocks_dim, w_blocks_dim=col_blocks_dim))
   x = mtf.reduce_mean(f2, reduced_dim=filters2_dim)
 
   # add some fully-connected dense layers.

From 21d508a254cfc2cc051d227b76b2801daf5c20e9 Mon Sep 17 00:00:00 2001
From: Keyon Vafa <vafa@google.com>
Date: Mon, 17 Sep 2018 15:31:10 -0700
Subject: [PATCH 0862/2720] Added IAF Flow layer to latent_layers.py

PiperOrigin-RevId: 213349934
---
 tensor2tensor/layers/latent_layers.py | 53 +++++++++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 5f4184168..9af39f50e 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -657,3 +657,56 @@ def transformer_autoencoder(inputs,
   else:
     decoder_output = output
   return decoder_output, losses, cache
+
+
+def iaf_flow(one_hot_assignments,
+             scale_weights,
+             num_codes,
+             summary=True,
+             name=None):
+  """Performs a single IAF flow using scale and normalization transformations.
+
+  Args:
+    one_hot_assignments: Assignments Tensor with shape [num_samples, batch_size,
+      latent_size, num_codes].
+    scale_weights: Tensor corresponding to lower triangular matrix used to
+      autoregressively generate scale matrix from assignments. To ensure the
+      lower-triangular matrix has length of latent_size, scale_weights should
+      be a rank-one tensor with size latent_size * (latent_size + 1) / 2.
+    num_codes: Number of codes in codebook.
+    summary: Whether to save summaries.
+    name: String used for name scope.
+
+  Returns:
+    flow_output: Transformed one-hot assignments.
+    inverse_log_det_jacobian: Inverse log deteriminant of Jacobian corresponding
+      to transformation.
+  """
+  with tf.name_scope(name, default_name="iaf"):
+    # Pad the one_hot_assignments by zeroing out the first latent dimension and
+    # shifting the rest down by one (and removing the last dimension).
+    padded_assignments = tf.pad(
+        one_hot_assignments, [[0, 0], [0, 0], [1, 0], [0, 0]])[:, :, :-1, :]
+    scale_bijector = tf.contrib.distributions.bijectors.Affine(
+        scale_tril=tf.contrib.distributions.fill_triangular(scale_weights))
+    scale = scale_bijector.forward(
+        tf.transpose(padded_assignments, [0, 1, 3, 2]))
+    # Transpose the bijector output since it performs a batch matmul.
+    scale = tf.transpose(scale, [0, 1, 3, 2])
+    scale = tf.nn.softplus(scale)
+    # Don't need last dimension since the transformation keeps it constant.
+    scale = scale[..., :-1]
+
+    z = one_hot_assignments[..., :-1]
+    unnormalized_probs = tf.concat([z * scale,
+                                    one_hot_assignments[..., -1, tf.newaxis]],
+                                   axis=-1)
+    normalizer = tf.reduce_sum(unnormalized_probs, axis=-1)
+    flow_output = unnormalized_probs / (normalizer[..., tf.newaxis])
+    inverse_log_det_jacobian = (-tf.reduce_sum(tf.log(scale), axis=-1)
+                                + num_codes * tf.log(normalizer))
+    if summary:
+      tf.summary.histogram("iaf/scale", tf.reshape(scale, [-1]))
+      tf.summary.histogram("iaf/inverse_log_det_jacobian",
+                           tf.reshape(inverse_log_det_jacobian, [-1]))
+    return flow_output, inverse_log_det_jacobian

From 9160c4fc9c84f4f32c1c636fd0eaae172cbff1ba Mon Sep 17 00:00:00 2001
From: Keyon Vafa <vafa@google.com>
Date: Mon, 17 Sep 2018 16:08:58 -0700
Subject: [PATCH 0863/2720] Added scale bias so IAF can initialize to identity

PiperOrigin-RevId: 213356296
---
 tensor2tensor/layers/latent_layers.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 9af39f50e..288a2f4d3 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -661,6 +661,7 @@ def transformer_autoencoder(inputs,
 
 def iaf_flow(one_hot_assignments,
              scale_weights,
+             scale_bias,
              num_codes,
              summary=True,
              name=None):
@@ -673,6 +674,9 @@ def iaf_flow(one_hot_assignments,
       autoregressively generate scale matrix from assignments. To ensure the
       lower-triangular matrix has length of latent_size, scale_weights should
       be a rank-one tensor with size latent_size * (latent_size + 1) / 2.
+    scale_bias: Bias tensor to be added to scale tensor, with shape
+      [latent_size, num_codes]. If scale weights are zero, initialize scale_bias
+      to be log(exp(1.) / 2. - 1) so initial transformation is identity.
     num_codes: Number of codes in codebook.
     summary: Whether to save summaries.
     name: String used for name scope.
@@ -694,6 +698,7 @@ def iaf_flow(one_hot_assignments,
     # Transpose the bijector output since it performs a batch matmul.
     scale = tf.transpose(scale, [0, 1, 3, 2])
     scale = tf.nn.softplus(scale)
+    scale = scale + tf.nn.softplus(scale_bias[tf.newaxis, tf.newaxis, ...])
     # Don't need last dimension since the transformation keeps it constant.
     scale = scale[..., :-1]
 

From 44ba523e5a0a6c4d456295676b63d2efbbfdb17a Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 17 Sep 2018 16:16:56 -0700
Subject: [PATCH 0864/2720] Remove the modality registry. Users should directly
 construct the modality objects in Problem.hparams

PiperOrigin-RevId: 213357662
---
 tensor2tensor/data_generators/problem.py |  18 +--
 tensor2tensor/layers/modalities.py       | 130 +++++++++++-----
 tensor2tensor/layers/modalities_test.py  |  18 +++
 tensor2tensor/utils/registry.py          | 183 +----------------------
 tensor2tensor/utils/registry_test.py     |  69 +--------
 5 files changed, 127 insertions(+), 291 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 657cbae2c..28fc88d7b 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -16,21 +16,21 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
 import collections
 import copy
 import functools
 import multiprocessing
 import os
 import random
-
 import six
+
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import text_encoder
-# Import modalities: they must be registered before we look them up here.
-from tensor2tensor.layers import modalities  # pylint: disable=unused-import
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import metrics
-from tensor2tensor.utils import registry
+
 import tensorflow as tf
 from tensorflow.contrib.tpu.python.tpu import tpu_config
 
@@ -1126,7 +1126,7 @@ def _create_modalities(problem_hparams, hparams):
                                     modality[0],
                                     feature_name)
         modality = (input_modality_overrides[feature_name], modality[1])
-      modality = registry.create_modality(modality, hparams)
+      modality = modalities.create_modality(modality, hparams)
     input_modality[feature_name] = modality
   problem_hparams.input_modality = input_modality
 
@@ -1146,7 +1146,7 @@ def _create_modalities(problem_hparams, hparams):
                                       modality[0],
                                       "target_modality/%s" % feature_name)
           modality = (target_modality_name, modality[1])
-        modality = registry.create_modality(modality, hparams)
+        modality = modalities.create_modality(modality, hparams)
       target_modality[feature_name] = modality
     problem_hparams.target_modality = target_modality
   elif isinstance(problem_hparams.target_modality, (list, tuple)):
@@ -1156,13 +1156,13 @@ def _create_modalities(problem_hparams, hparams):
                                   modality[0],
                                   "target")
       modality = (target_modality_name, modality[1])
-    modality = registry.create_modality(modality, hparams)
+    modality = modalities.create_modality(modality, hparams)
     problem_hparams.target_modality = modality
 
 
 def _warn_changed_modality_type(new_name, old_name, feature_name):
-  new_type, new_name = registry.parse_modality_name(new_name)
-  old_type, old_name = registry.parse_modality_name(old_name)
+  new_type, new_name = modalities.parse_modality_name(new_name)
+  old_type, old_name = modalities.parse_modality_name(old_name)
   if new_type != old_type:
     tf.logging.warn(
         "%s has a designated modality type %s (%s) but has been "
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 0dfad2e0b..316aa96c2 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -29,7 +29,6 @@
 import tensorflow as tf
 
 
-@registry.register_symbol_modality("default")
 class SymbolModality(modality.Modality):
   """Modality for sets of discrete symbols.
 
@@ -171,7 +170,6 @@ def top(self, body_output, _):
                             body_output_shape[:-1] + [1, self._vocab_size])
 
 
-@registry.register_symbol_modality("weights_all")
 class SymbolModalityWeightsAll(SymbolModality):
   """SymbolModality for features that do not have 0-padding."""
 
@@ -180,7 +178,6 @@ def targets_weights_fn(self):
     return common_layers.weights_all
 
 
-@registry.register_symbol_modality("one_hot")
 class SymbolModalityOneHot(SymbolModality):
   """Simple SymbolModality with one hot as embeddings."""
 
@@ -200,7 +197,6 @@ def loss(self, top_out, targets):
     return tf.reduce_mean(loss), tf.constant(1.0)
 
 
-@registry.register_symbol_modality("ctc")
 class CTCSymbolModality(SymbolModality):
   """SymbolModality that uses CTC loss."""
 
@@ -230,7 +226,6 @@ def loss(self, top_out, targets):
       return tf.reduce_sum(xent), tf.reduce_sum(weights)
 
 
-@registry.register_image_modality("default")
 class ImageModality(modality.Modality):
   """Modality for images."""
   PIXEL_EMBEDDING_SIZE = 64
@@ -299,7 +294,6 @@ def loss(self, top_out, targets):
         weights_fn=self.targets_weights_fn)
 
 
-@registry.register_image_modality("image_channel_compress")
 class ImageChannelCompressModality(modality.Modality):
   """Modality for images using channel compression for generation."""
 
@@ -388,14 +382,12 @@ def top(self, body_output, _):
       return x
 
 
-@registry.register_image_modality("image_channel_bottom_identity")
 class ImageChannelBottomIdentityModality(ImageChannelCompressModality):
 
   def top(self, body_output, _):
     return body_output
 
 
-@registry.register_image_modality("channel_embeddings_bottom")
 class ImageChannelEmbeddingsBottom(modality.Modality):
   """Modality for images using channel compression for generation."""
 
@@ -441,7 +433,6 @@ def top(self, body_output, _):
       return x
 
 
-@registry.register_audio_modality("default")
 class AudioModality(modality.Modality):
   """Performs strided conv compressions for audio data."""
 
@@ -486,7 +477,6 @@ def xnet_resblock(x, filters, res_relu, name):
                            "compress_block_final")
 
 
-@registry.register_audio_modality
 class AudioSpectralModality(modality.Modality):
   """Performs strided conv compressions for audio spectral data."""
 
@@ -532,7 +522,6 @@ def xnet_resblock(x, filters, res_relu, name):
                            "compress_block_final")
 
 
-@registry.register_audio_modality
 class SpeechRecognitionModality(modality.Modality):
   """Common ASR filterbank processing."""
 
@@ -615,7 +604,6 @@ def bottom(self, x):
     return x
 
 
-@registry.register_video_modality("default")
 class VideoModality(modality.Modality):
   """Modality for videos, i.e., time-sequences of frames."""
 
@@ -654,7 +642,6 @@ def loss(self, top_out, targets):
         weights_fn=self.targets_weights_fn)
 
 
-@registry.register_video_modality("bitwise")
 class VideoModalityBitwise(VideoModality):
   """Video Modality where bottom embeds pixels bitwise."""
   PIXEL_EMBEDDING_SIZE = 64
@@ -689,7 +676,6 @@ def targets_bottom(self, x):  # pylint: disable=arguments-differ
           name="merge_pixel_embedded_frames")
 
 
-@registry.register_video_modality("pixel_noise")
 class VideoModalityPixelNoise(VideoModality):
   """Video modality that introduces pixel noise on input during training."""
 
@@ -712,7 +698,6 @@ def input_noise(self):
     return getattr(self._model_hparams, "video_modality_input_noise", 0.25)
 
 
-@registry.register_video_modality("l1")
 class VideoModalityL1(VideoModality):
   """Video modality that predicts a scalar per channel with an L1 loss."""
 
@@ -754,7 +739,6 @@ def loss(self, top_out, targets):
     return tf.reduce_sum(loss * weights), tf.reduce_sum(weights)
 
 
-@registry.register_video_modality("l2")
 class VideoModalityL2(VideoModalityL1):
   """Modality for videos with L2 loss."""
 
@@ -762,7 +746,6 @@ def internal_loss(self, logits, targets):
     return tf.nn.relu((logits - targets)**2 - self.cutoff * self.cutoff)
 
 
-@registry.register_video_modality("l2raw")
 class VideoModalityL2Raw(VideoModalityL2):
   """Modality with L2 loss and raw input (sequences of frames)."""
 
@@ -795,7 +778,6 @@ def loss(self, top_out, targets):
     return loss, tf.constant(1.0)
 
 
-@registry.register_video_modality("l1raw")
 class VideoModalityL1Raw(VideoModalityL2Raw):
   """Modality with L1 loss and raw input (sequences of frames)."""
 
@@ -805,7 +787,6 @@ def loss(self, top_out, targets):
     return loss, tf.constant(1.0)
 
 
-@registry.register_class_label_modality("default")
 class ClassLabelModality(modality.Modality):
   """Used for label data."""
 
@@ -846,7 +827,6 @@ def top(self, body_output, _):
       return tf.expand_dims(res, 3)
 
 
-@registry.register_class_label_modality("multi_label")
 class MultiLabelModality(ClassLabelModality):
   """Used for multi label task."""
 
@@ -879,7 +859,6 @@ def loss(self, top_out, targets):
     return tf.reduce_sum(loss*weights), tf.reduce_sum(weights)
 
 
-@registry.register_class_label_modality("onehot")
 class OneHotClassLabelModality(ClassLabelModality):
   """Used for one-hot encoded class labels."""
 
@@ -899,12 +878,6 @@ def loss(self, top_out, targets):
     return loss_scale, loss_denom
 
 
-@registry.register_generic_modality("default")
-@registry.register_audio_modality("identity")
-@registry.register_image_modality("identity")
-@registry.register_video_modality("identity")
-@registry.register_class_label_modality("identity")
-@registry.register_real_modality("identity")
 class IdentityModality(modality.Modality):
   """Does nothing."""
 
@@ -915,7 +888,6 @@ def top(self, body_output, _):
     return body_output
 
 
-@registry.register_generic_modality("l2_loss")
 class GenericL2LossModality(IdentityModality):
   """Generic modality with L2 as Loss."""
 
@@ -951,8 +923,6 @@ def loss(self, top_out, targets):
     raise NotImplementedError()
 
 
-@registry.register_real_modality("default")
-@registry.register_real_modality("l2_loss")
 class RealL2LossModality(RealModality):
   """Modality for real (i.e. float) vectors with L2 (Gaussian) loss."""
 
@@ -967,7 +937,6 @@ def loss(self, top_out, targets):
       return tf.reduce_sum(l2 * weights), tf.reduce_sum(weights)
 
 
-@registry.register_real_modality("log_poisson_loss")
 class RealLogPoissonLossModality(RealModality):
   """Modality for real (i.e. float) vectors with log Poisson regression loss."""
 
@@ -982,7 +951,6 @@ def loss(self, top_out, targets):
       return tf.reduce_sum(lp_loss * weights), tf.reduce_sum(weights)
 
 
-@registry.register_symbol_modality("identity")
 class IdentitySymbolModality(SymbolModality):
   """Symbol modality with identity top and bottom transformations.
 
@@ -1005,7 +973,6 @@ def top_is_pointwise(self):
     return False
 
 
-@registry.register_class_label_modality("sigmoid")
 class SigmoidClassLabelModality(ClassLabelModality):
   """Sigmoid cross-entropy for independent class labels."""
 
@@ -1025,7 +992,6 @@ def loss(self, top_out, targets):
     return loss_scale, loss_denom
 
 
-@registry.register_class_label_modality("sigmoid_max_pooling")
 class SigmoidMaxPoolingClassLabelModality(ClassLabelModality):
   """Sigmoid cross-entropy applied on max-pooling over timesteps."""
 
@@ -1061,7 +1027,6 @@ def loss(self, top_out, targets):
     return loss_scale, loss_denom
 
 
-@registry.register_class_label_modality("onehot_softmax_max_pooling")
 class SoftmaxMaxPoolingClassLabelModality(OneHotClassLabelModality):
   """Softmax cross-entropy applied on max-pooling over timesteps."""
 
@@ -1077,7 +1042,6 @@ def top(self, body_output, _):
       return tf.layers.dense(x, self._vocab_size)
 
 
-@registry.register_class_label_modality("onehot_softmax_average_pooling")
 class SoftmaxAveragePoolingClassLabelModality(OneHotClassLabelModality):
   """Softmax cross-entropy applied on average-pooling over timesteps."""
 
@@ -1093,7 +1057,6 @@ def top(self, body_output, _):
       return tf.layers.dense(x, self._vocab_size)
 
 
-@registry.register_class_label_modality("onehot_softmax_last_timestep")
 class SoftmaxLastTimestepClassLabelModality(OneHotClassLabelModality):
   """Softmax cross-entropy applied on last-timestep encoding."""
 
@@ -1107,3 +1070,96 @@ def top(self, body_output, _):
       x = body_output
       x = tf.expand_dims(x[:, -1], 1)  # Pick the last timestep
       return tf.layers.dense(x, self._vocab_size)
+
+
+def create_modality(modality_spec, model_hparams):
+  """Creates modality.
+
+  Args:
+    modality_spec: tuple ("modality_type:modality_name", vocab_size).
+    model_hparams: tf.contrib.training.HParams.
+
+  Returns:
+    Modality.
+
+  Raises:
+    LookupError: if modality_type is not recognized. See registry.Modalities for
+      accepted types.
+  """
+  modality_full_name, vocab_size = modality_spec
+  modality_type, modality_name = parse_modality_name(modality_full_name)
+
+  if modality_type == registry.Modalities.SYMBOL:
+    modality_collection = {
+        "default": SymbolModality,
+        "identity": IdentitySymbolModality,
+        "weights_all": SymbolModalityWeightsAll,
+        "one_hot": SymbolModalityOneHot,
+        "ctc": CTCSymbolModality,
+    }
+  elif modality_type == registry.Modalities.IMAGE:
+    modality_collection = {
+        "default": ImageModality,
+        "identity": IdentityModality,
+        "image_channel_compress": ImageChannelCompressModality,
+        "image_channel_bottom_identity": ImageChannelBottomIdentityModality,
+        "channel_embeddings_bottom": ImageChannelEmbeddingsBottom,
+    }
+  elif modality_type == registry.Modalities.AUDIO:
+    modality_collection = {
+        "default": SpeechRecognitionModality,
+        "identity": IdentityModality,
+        "spectral": AudioSpectralModality,
+        "speech": SpeechRecognitionModality,
+    }
+  elif modality_type == registry.Modalities.VIDEO:
+    modality_collection = {
+        "default": VideoModality,
+        "identity": IdentityModality,
+        "bitwise": VideoModalityBitwise,
+        "pixel_noise": VideoModalityPixelNoise,
+        "l1": VideoModalityL1,
+        "l2": VideoModalityL2,
+        "l2raw": VideoModalityL2Raw,
+        "l1raw": VideoModalityL1Raw,
+    }
+  elif modality_type == registry.Modalities.CLASS_LABEL:
+    modality_collection = {
+        "default": ClassLabelModality,
+        "identity": IdentityModality,
+        "multi_label": MultiLabelModality,
+        "onehot": OneHotClassLabelModality,
+        "sigmoid": SigmoidClassLabelModality,
+        "sigmoid_max_pooling": SigmoidMaxPoolingClassLabelModality,
+        "onehot_softmax_max_pooling": SoftmaxMaxPoolingClassLabelModality,
+        "onehot_softmax_average_pooling":
+            SoftmaxAveragePoolingClassLabelModality,
+        "onehot_softmax_last_timestep": SoftmaxLastTimestepClassLabelModality,
+    }
+  elif modality_type == registry.Modalities.GENERIC:
+    modality_collection = {
+        "default": IdentityModality,
+        "l2_loss": GenericL2LossModality,
+    }
+  elif modality_type == registry.Modalities.REAL:
+    modality_collection = {
+        "default": RealL2LossModality,
+        "identity": IdentityModality,
+        "l2_loss": RealL2LossModality,
+        "log_poisson_loss": RealLogPoissonLossModality,
+    }
+  else:
+    modality_types = ("symbol", "image", "audio", "video", "class_label",
+                      "generic", "real")
+    raise LookupError("Modality type %s not recognized. Options are: %s" %
+                      (modality_type, list(modality_types)))
+
+  return modality_collection[modality_name](model_hparams, vocab_size)
+
+
+def parse_modality_name(name):
+  name_parts = name.split(":")
+  if len(name_parts) < 2:
+    name_parts.append("default")
+  modality_type, modality_name = name_parts
+  return modality_type, modality_name
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index e0f631e30..875b66918 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -16,11 +16,13 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
 import numpy as np
 
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import expert_utils
+from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -110,6 +112,22 @@ def testSymbolModalityTargetsFactored(self):
     self.assertEqual(res1.shape, (batch_size, length, height, 1, vocab_size))
     self.assertEqual(res2.shape, ())
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testCreateModality(self):
+    model_hparams = tf.contrib.training.HParams()
+
+    modality_spec = (registry.Modalities.SYMBOL, 2)
+    modality = modalities.create_modality(modality_spec, model_hparams)
+    self.assertIsInstance(modality, modalities.SymbolModality)
+
+    modality_spec = (registry.Modalities.CLASS_LABEL + ":onehot", None)
+    modality = modalities.create_modality(modality_spec, model_hparams)
+    self.assertIsInstance(modality, modalities.OneHotClassLabelModality)
+
+    modality_spec = (registry.Modalities.VIDEO + ":identity", None)
+    modality = modalities.create_modality(modality_spec, model_hparams)
+    self.assertIsInstance(modality, modalities.IdentityModality)
+
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 14600d216..af4d2b0e8 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -45,7 +45,6 @@ class MyModel(T2TModel):
 
 import inspect
 import re
-import six
 import tensorflow as tf
 
 _ATTACKS = {}
@@ -58,7 +57,10 @@ class MyModel(T2TModel):
 _RANGED_HPARAMS = {}
 
 
+# TODO(trandustin): Many files depend on this to specify modality strings; let's
+# remove it in the future.
 class Modalities(object):
+  """An enum-like object carrying the set of available modality types."""
   SYMBOL = "symbol"
   IMAGE = "image"
   AUDIO = "audio"
@@ -68,16 +70,6 @@ class Modalities(object):
   REAL = "real"
 
 
-_MODALITIES = {
-    Modalities.SYMBOL: {},
-    Modalities.IMAGE: {},
-    Modalities.AUDIO: {},
-    Modalities.VIDEO: {},
-    Modalities.CLASS_LABEL: {},
-    Modalities.GENERIC: {},
-    Modalities.REAL: {},
-}
-
 # Camel case to snake case utils
 _first_cap_re = re.compile("(.)([A-Z][a-z0-9]+)")
 _all_cap_re = re.compile("([a-z0-9])([A-Z])")
@@ -89,8 +81,7 @@ def _convert_camel_to_snake(name):
 
 
 def _reset():
-  for ctr in [_MODELS, _HPARAMS, _RANGED_HPARAMS, _ATTACK_PARAMS] + list(
-      _MODALITIES.values()):
+  for ctr in [_MODELS, _HPARAMS, _RANGED_HPARAMS, _ATTACK_PARAMS]:
     ctr.clear()
 
 
@@ -441,164 +432,6 @@ def list_pruning_strategies(prefix=None):
   return list(_PRUNING_STRATEGY)
 
 
-def _internal_get_modality(name, mod_collection, collection_str):
-  if name is None:
-    name = "default"
-  if name not in mod_collection:
-    raise LookupError(
-        "%s modality %s never registered." % (collection_str, name))
-  return mod_collection[name]
-
-
-def symbol_modality(name=None):
-  return _internal_get_modality(name, _MODALITIES[Modalities.SYMBOL],
-                                Modalities.SYMBOL.capitalize())
-
-
-def generic_modality(name=None):
-  return _internal_get_modality(name, _MODALITIES[Modalities.GENERIC],
-                                Modalities.GENERIC.capitalize())
-
-
-def video_modality(name=None):
-  return _internal_get_modality(name, _MODALITIES[Modalities.VIDEO],
-                                Modalities.VIDEO.capitalize())
-
-
-def audio_modality(name=None):
-  return _internal_get_modality(name, _MODALITIES[Modalities.AUDIO],
-                                Modalities.AUDIO.capitalize())
-
-
-def image_modality(name=None):
-  return _internal_get_modality(name, _MODALITIES[Modalities.IMAGE],
-                                Modalities.IMAGE.capitalize())
-
-
-def class_label_modality(name=None):
-  return _internal_get_modality(name, _MODALITIES[Modalities.CLASS_LABEL],
-                                Modalities.CLASS_LABEL.capitalize())
-
-
-def real_modality(name=None):
-  return _internal_get_modality(name, _MODALITIES[Modalities.REAL],
-                                Modalities.REAL.capitalize())
-
-
-def _internal_register_modality(name, mod_collection, collection_str):
-  """Register a modality into mod_collection."""
-
-  def decorator(mod_cls, registration_name=None):
-    """Registers & returns mod_cls with registration_name or default name."""
-    mod_name = registration_name or default_name(mod_cls)
-    if mod_name in mod_collection and not tf.contrib.eager.in_eager_mode():
-      raise LookupError(
-          "%s modality %s already registered." % (collection_str, mod_name))
-    mod_collection[mod_name] = mod_cls
-    return mod_cls
-
-  # Handle if decorator was used without parens
-  if callable(name):
-    mod_cls = name
-    return decorator(mod_cls, registration_name=default_name(mod_cls))
-
-  return lambda mod_cls: decorator(mod_cls, name)
-
-
-def register_symbol_modality(name=None):
-  """Register a symbol modality. name defaults to class name snake-cased."""
-  return _internal_register_modality(name, _MODALITIES[Modalities.SYMBOL],
-                                     Modalities.SYMBOL.capitalize())
-
-
-def register_generic_modality(name=None):
-  """Register a generic modality. name defaults to class name snake-cased."""
-  return _internal_register_modality(name, _MODALITIES[Modalities.GENERIC],
-                                     Modalities.GENERIC.capitalize())
-
-
-def register_real_modality(name=None):
-  """Register a real modality. name defaults to class name snake-cased."""
-  return _internal_register_modality(name, _MODALITIES[Modalities.REAL],
-                                     Modalities.REAL.capitalize())
-
-
-def register_audio_modality(name=None):
-  """Register an audio modality. name defaults to class name snake-cased."""
-  return _internal_register_modality(name, _MODALITIES[Modalities.AUDIO],
-                                     Modalities.AUDIO.capitalize())
-
-
-def register_image_modality(name=None):
-  """Register an image modality. name defaults to class name snake-cased."""
-  return _internal_register_modality(name, _MODALITIES[Modalities.IMAGE],
-                                     Modalities.IMAGE.capitalize())
-
-
-def register_video_modality(name=None):
-  """Register a video modality. name defaults to class name snake-cased."""
-  return _internal_register_modality(name, _MODALITIES[Modalities.VIDEO],
-                                     Modalities.VIDEO.capitalize())
-
-
-def register_class_label_modality(name=None):
-  """Register an image modality. name defaults to class name snake-cased."""
-  return _internal_register_modality(name, _MODALITIES[Modalities.CLASS_LABEL],
-                                     Modalities.CLASS_LABEL.capitalize())
-
-
-def list_modalities():
-  all_modalities = []
-  for modality_type, modalities in six.iteritems(_MODALITIES):
-    all_modalities.extend([
-        "%s:%s" % (mtype, modality)
-        for mtype, modality in zip([modality_type] *
-                                   len(modalities), modalities)
-    ])
-  return all_modalities
-
-
-def parse_modality_name(name):
-  name_parts = name.split(":")
-  if len(name_parts) < 2:
-    name_parts.append("default")
-  modality_type, modality_name = name_parts
-  return modality_type, modality_name
-
-
-def create_modality(modality_spec, model_hparams):
-  """Create modality.
-
-  Args:
-    modality_spec: tuple, ("modality_type:modality_name", vocab_size).
-    model_hparams: HParams object.
-
-  Returns:
-    Modality instance.
-
-  Raises:
-    LookupError: if modality_type is not recognized. See Modalities class for
-    accepted types.
-  """
-  retrieval_fns = {
-      Modalities.SYMBOL: symbol_modality,
-      Modalities.IMAGE: image_modality,
-      Modalities.AUDIO: audio_modality,
-      Modalities.VIDEO: video_modality,
-      Modalities.CLASS_LABEL: class_label_modality,
-      Modalities.GENERIC: generic_modality,
-      Modalities.REAL: real_modality,
-  }
-
-  modality_full_name, vocab_size = modality_spec
-  modality_type, modality_name = parse_modality_name(modality_full_name)
-  if modality_type not in retrieval_fns:
-    raise LookupError("Modality type %s not recognized. Options are: %s" %
-                      (modality_type, list(_MODALITIES)))
-
-  return retrieval_fns[modality_type](modality_name)(model_hparams, vocab_size)
-
-
 def display_list_by_prefix(names_list, starting_spaces=0):
   """Creates a help string for names_list grouped by prefix."""
   cur_prefix, result_lines = None, []
@@ -628,9 +461,6 @@ def help_string():
   RangedHParams:
 %s
 
-  Modalities:
-%s
-
   Problems:
 %s
 
@@ -646,12 +476,11 @@ def help_string():
   Pruning Strategies:
 %s
 """
-  m, hp, rhp, mod, probs, atks, ap, pp, ps = [
+  m, hp, rhp, probs, atks, ap, pp, ps = [
       display_list_by_prefix(entries, starting_spaces=4) for entries in [
           list_models(),
           list_hparams(),
           list_ranged_hparams(),
-          list_modalities(),
           list_problems(),
           list_attacks(),
           list_attack_params(),
@@ -659,4 +488,4 @@ def help_string():
           list_pruning_strategies(),
       ]
   ]
-  return help_str % (m, hp, rhp, mod, probs, atks, ap, pp, ps)
+  return help_str % (m, hp, rhp, probs, atks, ap, pp, ps)
diff --git a/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index 7324f1712..b7b328907 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -17,7 +17,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-from tensor2tensor.utils import modality
+
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -206,73 +206,6 @@ def rhp_bad2(a, b):  # pylint: disable=unused-argument
         pass
 
 
-class ModalityRegistryTest(tf.test.TestCase):
-
-  def setUp(self):
-    registry._reset()
-
-  def testModalityRegistration(self):
-
-    @registry.register_symbol_modality
-    class MySymbolModality(modality.Modality):
-      pass
-
-    @registry.register_audio_modality
-    class MyAudioModality(modality.Modality):
-      pass
-
-    @registry.register_image_modality
-    class MyImageModality(modality.Modality):
-      pass
-
-    @registry.register_class_label_modality
-    class MyClassLabelModality(modality.Modality):
-      pass
-
-    self.assertTrue(
-        registry.symbol_modality("my_symbol_modality") is MySymbolModality)
-    self.assertTrue(
-        registry.audio_modality("my_audio_modality") is MyAudioModality)
-    self.assertTrue(
-        registry.image_modality("my_image_modality") is MyImageModality)
-    self.assertTrue(
-        registry.class_label_modality("my_class_label_modality") is
-        MyClassLabelModality)
-
-  def testDefaultNameLookup(self):
-
-    @registry.register_symbol_modality("default")
-    class MyDefaultModality(modality.Modality):
-      pass
-
-    self.assertTrue(registry.symbol_modality() is MyDefaultModality)
-
-  def testList(self):
-
-    @registry.register_symbol_modality
-    class MySymbolModality(modality.Modality):
-      pass
-
-    @registry.register_audio_modality
-    class MyAudioModality(modality.Modality):
-      pass
-
-    @registry.register_image_modality
-    class MyImageModality(modality.Modality):
-      pass
-
-    @registry.register_class_label_modality
-    class MyClassLabelModality(modality.Modality):
-      pass
-
-    expected = [
-        "symbol:my_symbol_modality", "audio:my_audio_modality",
-        "image:my_image_modality", "class_label:my_class_label_modality"
-    ]
-
-    self.assertSetEqual(set(registry.list_modalities()), set(expected))
-
-
 class RegistryTest(tf.test.TestCase):
   """ Test class for common functions."""
 

From e451d0d412b629210a5fce155a96b0ad033cd4df Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Mon, 17 Sep 2018 19:53:11 -0700
Subject: [PATCH 0865/2720] Add training configs for a deeper transformer.

PiperOrigin-RevId: 213383428
---
 tensor2tensor/models/transformer.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 4ab85d284..708bb7d86 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1692,6 +1692,32 @@ def transformer_big():
   return hparams
 
 
+@registry.register_hparams
+def transformer_tall():
+  """Hparams for transformer on LM+MNLI."""
+  hparams = transformer_base()
+  hparams.batch_size = 2048
+  hparams.hidden_size = 768
+  hparams.filter_size = 3072
+  hparams.num_hidden_layers = 12
+  hparams.num_heads = 12
+  hparams.learning_rate_schedule = (
+      "constant*linear_warmup*rsqrt_hidden_size")
+  hparams.learning_rate_constant = 2e-3
+  hparams.label_smoothing = 0.0
+  hparams.max_length = 512
+  hparams.eval_drop_long_sequences = True
+  return hparams
+
+
+@registry.register_hparams
+def transformer_tall_big():
+  """Hparams for transformer on LM+MNLI."""
+  hparams = transformer_tall()
+  hparams.num_hidden_layers = 18
+  return hparams
+
+
 @registry.register_hparams
 def transformer_big_single_gpu():
   """HParams for transformer big model for single GPU."""

From 0d971f4b0df34a02ba0d6b9a3d21a75e7dde54b7 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Mon, 17 Sep 2018 23:37:54 -0700
Subject: [PATCH 0866/2720] Remove distributed_moe, it doesn't work anymore and
 using data parallelism is cumbersome. MeshTF already supports this for models
 that use distributed_moe

PiperOrigin-RevId: 213400840
---
 tensor2tensor/layers/common_attention.py      |  30 +--
 .../layers/common_image_attention.py          |  62 -----
 tensor2tensor/models/__init__.py              |   1 -
 tensor2tensor/models/research/aligned.py      |  24 --
 .../models/research/attention_lm_moe.py       |  26 +-
 tensor2tensor/models/research/multimodel.py   | 251 ------------------
 .../models/research/multimodel_test.py        |  54 ----
 tensor2tensor/utils/expert_utils.py           |  69 -----
 8 files changed, 2 insertions(+), 515 deletions(-)
 delete mode 100644 tensor2tensor/models/research/multimodel.py
 delete mode 100644 tensor2tensor/models/research/multimodel_test.py

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index cfe7e20d3..74f45d465 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -42,7 +42,7 @@
 _expert_count = 0
 
 
-def get_standardized_layers(hparams, dp=None, ps_devices=None):
+def get_standardized_layers(hparams, dp=None):
   """Get the common attention and feed-forward layers.
 
   The returned layer functions will have the following signature:
@@ -57,7 +57,6 @@ def get_standardized_layers(hparams, dp=None, ps_devices=None):
     hparams (tf.HParams): the model hparameters
     dp (expert_utils.Parallelism): A data parallelism object. If not given,
       the dp calls are simply ignored.
-    ps_devices: a reference to model._ps_devices (only used by the MOE layer)
 
   Returns:
     dict[str:fct]: A dictionary containing the standardized functions
@@ -119,14 +118,6 @@ def decorator(x, *args, **kwargs):
 
   total_key_depth = hparams.attention_key_channels or hparams.hidden_size
   total_value_depth = hparams.attention_value_channels or hparams.hidden_size
-  is_train = hparams.mode == tf.estimator.ModeKeys.TRAIN
-
-  moe_hidden_sizes = [int(s) for s in hparams.moe_hidden_sizes.split(",")]
-  # Use filter size if moe_hidden_sizes was not given
-  if not moe_hidden_sizes:
-    moe_hidden_sizes = [hparams.filter_size]
-  expert_fn = expert_utils.ffn_expert_fn(hparams.hidden_size, moe_hidden_sizes,
-                                         hparams.hidden_size)
 
   # Attention layers:
 
@@ -217,24 +208,6 @@ def memeff_attention_fn(*args, **kwargs):
 
   # Feed-forwards layers:
 
-  # === Mixture of expert layer ===
-  distributed_moe = register_layer(
-      expert_utils.distributed_moe,
-      default_args=[
-          dp,
-          ps_devices,
-      ],
-      default_kwargs=dict(
-          train=is_train,
-          input_size=hparams.hidden_size,
-          expert_fn=expert_fn,
-          num_experts=hparams.moe_num_experts,
-          k=hparams.moe_k,
-          loss_coef=hparams.moe_loss_coef,
-      ),
-      use_dp=False,
-  )
-
   # === FC layer ===
   conv_hidden_relu = register_layer(
       common_layers.conv_hidden_relu,
@@ -276,7 +249,6 @@ def memeff_attention_fn(*args, **kwargs):
       fc=conv_hidden_relu,  # Fully connected
       sep=sep_conv_relu,  # Separable convolution (unmasked)
       sepm=sep_conv_relu_masked,  # Separable convolution (masked)
-      moe=distributed_moe,  # Mixture of expert layer
   )
   return layers
 
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 9454e2c13..720d58c9b 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -463,68 +463,6 @@ def get_self_attention_bias(x):
   return self_attention_bias
 
 
-def transformer_layers_sharded(dp,
-                               ps_devices,
-                               inputs,
-                               num_layers,
-                               hparams,
-                               self_attention_bias=None,
-                               enc_output=None,
-                               attention_type=AttentionType.GLOBAL,
-                               name="transformer"):
-  """Multi layer transformer, sharded by the data parallelism dp."""
-  x = inputs
-  extra_loss = tf.constant(0.0)
-  moe_hidden_sizes = [int(s) for s in hparams.moe_hidden_sizes.split(",")]
-  expert_fn = expert_utils.ffn_expert_fn(
-      hparams.hidden_size, moe_hidden_sizes, hparams.hidden_size)
-  x = dp(tf.nn.dropout, x, 1.0 - hparams.layer_prepostprocess_dropout)
-  for layer in range(num_layers):
-    with tf.variable_scope("%s_layer_%d" % (name, layer)):
-      # self-attention
-      if attention_type == AttentionType.LOCAL_2D:
-        y = dp(local_attention_2d(common_layers.layer_preprocess(x, hparams),
-                                  hparams,
-                                  attention_type="masked_local_attention_2d"))
-      elif attention_type == AttentionType.LOCAL_1D:
-        y = dp(local_attention_1d(common_layers.layer_preprocess(x, hparams),
-                                  hparams,
-                                  attention_type="local_mask_right",
-                                  q_padding="LEFT", kv_padding="LEFT"))
-      elif attention_type == AttentionType.GLOCAL:
-        y = dp(local_global_attention(
-            common_layers.layer_preprocess(x, hparams), self_attention_bias,
-            hparams, q_padding="LEFT", kv_padding="LEFT"))
-      elif attention_type == AttentionType.GLOBAL:
-        self_attention_bias = dp(get_self_attention_bias(x))
-        y = dp(full_self_attention(common_layers.layer_preprocess(x, hparams),
-                                   self_attention_bias, hparams,
-                                   q_padding="LEFT", kv_padding="LEFT"))
-      x = common_layers.layer_postprocess(x, y, hparams)
-      if enc_output is not None:
-        y = dp(encdec_attention_1d(common_layers.layer_preprocess(x, hparams),
-                                   enc_output, None, hparams))
-        x = dp(common_layers.layer_postprocess, x, y, hparams)
-      with tf.variable_scope("ffn"):
-        if str(layer) in hparams.moe_layers_decoder.split(","):
-          y, loss = expert_utils.distributed_moe(
-              dp,
-              ps_devices,
-              common_layers.layer_preprocess(x, hparams),
-              hparams.mode == tf.estimator.ModeKeys.TRAIN,
-              input_size=hparams.hidden_size,
-              expert_fn=expert_fn,
-              num_experts=hparams.moe_num_experts,
-              k=hparams.moe_k,
-              loss_coef=hparams.moe_loss_coef)
-          extra_loss += loss
-          x = dp(common_layers.layer_postprocess, x, y, hparams)
-        else:
-          y = dp(ffn_layer, common_layers.layer_preprocess(x, hparams), hparams)
-          x = dp(common_layers.layer_postprocess, x, y, hparams)
-  return dp(common_layers.layer_preprocess, x, hparams), extra_loss
-
-
 def postprocess_image(x, rows, cols, hparams):
   """Postprocessing after decoding.
 
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index e5e6ed265..be08363f0 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -46,7 +46,6 @@
 from tensor2tensor.models.research import gene_expression
 from tensor2tensor.models.research import glow
 from tensor2tensor.models.research import lm_experiments
-from tensor2tensor.models.research import multimodel
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research import similarity_transformer
 from tensor2tensor.models.research import super_lm
diff --git a/tensor2tensor/models/research/aligned.py b/tensor2tensor/models/research/aligned.py
index 10a4f1c04..18a576b5d 100644
--- a/tensor2tensor/models/research/aligned.py
+++ b/tensor2tensor/models/research/aligned.py
@@ -28,7 +28,6 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
-from tensor2tensor.utils import diet
 from tensor2tensor.utils import expert_utils
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
@@ -69,7 +68,6 @@ def postprocess(x, y):
     x = dp(tf.nn.dropout, x, 1.0 - hparams.layer_prepostprocess_dropout)
     extra_loss = 0.0
     ffn_hidden_sizes = [int(s) for s in hparams.ffn_hidden_sizes.split(",")]
-    moe_hidden_sizes = [int(s) for s in hparams.moe_hidden_sizes.split(",")]
     if hparams.mask_right:
 
       def _bias(x):
@@ -79,16 +77,6 @@ def _bias(x):
       bias = dp(_bias, x)
     else:
       bias = tf.zeros([1, 1, 1, 1])
-    if hparams.diet_experts:
-      hsize, = moe_hidden_sizes
-
-      def _diet_expert(x):
-        return diet.diet_expert(x, hsize, diet.diet_adam_optimizer_params())
-
-      expert_fn = _diet_expert
-    else:
-      expert_fn = expert_utils.ffn_expert_fn(
-          hparams.hidden_size, moe_hidden_sizes, hparams.hidden_size)
 
     batch_coordinate = dp(get_batch_coordinate, x)
 
@@ -209,18 +197,6 @@ def _pseudolocal_bias(x):
               use_map_fn=False,
               experts_params=dict(nb_hyperplanes=4,))
           extra_loss += tf.add_n(loss) / dp.n
-        elif layer_type == "moe":
-          y, loss = expert_utils.distributed_moe(
-              dp,
-              self._ps_devices,
-              x,
-              hparams.mode == ModeKeys.TRAIN,
-              input_size=hparams.hidden_size,
-              expert_fn=expert_fn,
-              num_experts=hparams.moe_num_experts,
-              k=hparams.moe_k,
-              loss_coef=hparams.moe_loss_coef)
-          extra_loss += loss
         elif layer_type == "ffn":
           y = dp(
               expert_utils.ffn_expert_fn(hparams.hidden_size, ffn_hidden_sizes,
diff --git a/tensor2tensor/models/research/attention_lm_moe.py b/tensor2tensor/models/research/attention_lm_moe.py
index 29ff7a66f..a060e9ac1 100644
--- a/tensor2tensor/models/research/attention_lm_moe.py
+++ b/tensor2tensor/models/research/attention_lm_moe.py
@@ -30,7 +30,6 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
-from tensor2tensor.utils import diet
 from tensor2tensor.utils import expert_utils
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
@@ -106,17 +105,6 @@ def postprocess(x, y):
     x = dp(tf.nn.dropout, decoder_input,
            1.0 - hparams.layer_prepostprocess_dropout)
     extra_loss = 0.0
-    moe_hidden_sizes = [int(s) for s in hparams.moe_hidden_sizes.split(",")]
-    if hparams.diet_experts:
-      hsize, = moe_hidden_sizes
-
-      def _diet_expert(x):
-        return diet.diet_expert(x, hsize, diet.diet_adam_optimizer_params())
-
-      expert_fn = _diet_expert
-    else:
-      expert_fn = expert_utils.ffn_expert_fn(
-          hparams.hidden_size, moe_hidden_sizes, hparams.hidden_size)
 
     if not hparams.use_inputs:
       # As preprocess and postprocess are called with batch of size one (all
@@ -312,19 +300,7 @@ def print_shape(x, suffix, debug=False):
                 AttentionType.get_choices()))
           x = postprocess(x, y)
         with tf.variable_scope("ffn"):
-          if str(layer) in hparams.moe_layers.split(","):
-            y, loss = expert_utils.distributed_moe(
-                dp,
-                self._ps_devices,
-                preprocess(x),
-                hparams.mode == ModeKeys.TRAIN,
-                input_size=hparams.hidden_size,
-                expert_fn=expert_fn,
-                num_experts=hparams.moe_num_experts,
-                k=hparams.moe_k,
-                loss_coef=hparams.moe_loss_coef)
-            extra_loss += loss
-          elif hparams.memory_efficient_ffn:
+          if hparams.memory_efficient_ffn:
             assert hparams.layer_preprocess_sequence == "n"
             y = dp(
                 common_layers.conv_hidden_relu_memory_efficient,
diff --git a/tensor2tensor/models/research/multimodel.py b/tensor2tensor/models/research/multimodel.py
deleted file mode 100644
index 4152e20c3..000000000
--- a/tensor2tensor/models/research/multimodel.py
+++ /dev/null
@@ -1,251 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-"""MultiModel."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-from six.moves import range  # pylint: disable=redefined-builtin
-
-from tensor2tensor.layers import common_attention
-from tensor2tensor.layers import common_hparams
-from tensor2tensor.layers import common_layers
-from tensor2tensor.layers import modalities
-from tensor2tensor.models import slicenet
-from tensor2tensor.utils import expert_utils
-from tensor2tensor.utils import registry
-from tensor2tensor.utils import t2t_model
-
-import tensorflow as tf
-
-
-def conv_res_step(x, hparams, padding, mask):
-  """One step of convolutions and mid-residual."""
-  k = (hparams.kernel_height, hparams.kernel_width)
-  k2 = (hparams.large_kernel_size, 1)
-  dilations_and_kernels1 = [((1, 1), k), ((1, 1), k)]
-  dilations_and_kernels2 = [((1, 1), k2), ((4, 4), k2)]
-  with tf.variable_scope("conv_res_step"):
-    y = common_layers.subseparable_conv_block(
-        x,
-        hparams.filter_size,
-        dilations_and_kernels1,
-        padding=padding,
-        mask=mask,
-        separabilities=0,
-        name="residual1")
-    y = tf.nn.dropout(y, 1.0 - hparams.dropout)
-    return common_layers.subseparable_conv_block(
-        y,
-        hparams.hidden_size,
-        dilations_and_kernels2,
-        padding=padding,
-        mask=mask,
-        separabilities=0,
-        name="residual2")
-
-
-def residual_fn2(x, y, hparams):
-  y = tf.nn.dropout(y, 1.0 - hparams.dropout)
-  return common_layers.layer_norm(x + y)
-
-
-def residual_fn3(x, y, z, hparams):
-  y = tf.nn.dropout(y, 1.0 - hparams.dropout)
-  z = tf.nn.dropout(z, 1.0 - hparams.dropout)
-  return common_layers.layer_norm(x + y + z)
-
-
-def conv_experts(xs, hparams, dp, ps, padding, mask, layer_id):
-  """Convolutions + Mixture-of-Experts layer."""
-  del layer_id  # Unused.
-  train = hparams.mode == tf.estimator.ModeKeys.TRAIN,
-  conv_out = dp(conv_res_step, xs, hparams, padding, mask)
-  loss = 0.0
-  moe_hidden_sizes = [hparams.filter_size]
-  expert_fn = expert_utils.ffn_expert_fn(hparams.hidden_size, moe_hidden_sizes,
-                                         hparams.hidden_size)
-  moe_out, loss = expert_utils.distributed_moe(
-      dp,
-      ps,
-      xs,
-      train,
-      input_size=hparams.hidden_size,
-      expert_fn=expert_fn,
-      num_experts=hparams.moe_num_experts,
-      k=hparams.moe_k,
-      loss_coef=1.0)
-  return dp(residual_fn3, xs, moe_out, conv_out, hparams), loss
-
-
-def prepare_decoder(targets, target_space_emb):
-  """Prepare decoder."""
-  decoder_self_attention_bias = (
-      common_attention.attention_bias_lower_triangle(tf.shape(targets)[1]))
-  target_space_emb = tf.reshape(target_space_emb, [1, 1, -1])
-  target_space_emb = tf.tile(target_space_emb, [tf.shape(targets)[0], 1, 1])
-  decoder_input = common_layers.shift_right_3d(
-      targets, pad_value=target_space_emb)
-  decoder_input = common_attention.add_timing_signal_1d(decoder_input)
-  return (decoder_input, decoder_self_attention_bias)
-
-
-@registry.register_model
-class MultiModel(t2t_model.T2TModel):
-  """Model to train on multiple tasks simultaneously."""
-
-  @property
-  def use_body_sharded(self):
-    return True
-
-  def body_sharded(self, sharded_features):
-    train = self._hparams.mode == tf.estimator.ModeKeys.TRAIN
-    dp = self._data_parallelism
-    hparams = self._hparams
-
-    def project_to_hidden(inputs):
-      return common_layers.conv_block(
-          inputs,
-          hparams.hidden_size, [((1, 1), (3, 3))],
-          first_relu=False,
-          padding="SAME",
-          force2d=True)
-
-    def flatten(inputs):
-      return tf.expand_dims(common_layers.flatten4d3d(inputs), axis=2)
-
-    # Project to hidden size if necessary
-    if (sharded_features["inputs"][0].get_shape().as_list()[-1] !=
-        hparams.hidden_size):
-      inputs = dp(project_to_hidden, sharded_features["inputs"])
-
-    inputs = dp(flatten, inputs)
-    inputs_pad = dp(slicenet.embedding_to_padding, inputs)
-    inputs_mask = dp(lambda x: 1.0 - x, inputs_pad)
-    inputs_encoded = dp(common_layers.add_timing_signal, inputs)
-    expert_loss = 0.0
-    for i in range(hparams.num_hidden_layers):
-      with tf.variable_scope("enc_layer_%d" % i):
-        inputs_encoded, moe_loss = conv_experts(inputs_encoded, hparams, dp,
-                                                self._ps_devices, "SAME",
-                                                inputs_mask, i)
-        expert_loss += tf.reduce_mean(moe_loss) * hparams.moe_loss_coef
-
-    # If we're just predicing a class, there is no use for a decoder, return.
-    if isinstance(self._problem_hparams.target_modality,
-                  modalities.ClassLabelModality):
-      return inputs_encoded, tf.reduce_mean(expert_loss)
-
-    # Decoder.
-    inputs3d = dp(tf.squeeze, inputs, 2)
-    inputs_encoded3d = dp(tf.squeeze, inputs_encoded, 2)
-    encoder_padding = dp(common_attention.embedding_to_padding, inputs3d)
-    encoder_attention_bias = dp(common_attention.attention_bias_ignore_padding,
-                                encoder_padding)
-    targets = dp(common_layers.flatten4d3d, sharded_features["targets"])
-    target_space_emb = dp(slicenet.embed_target_space,
-                          sharded_features["target_space_id"],
-                          hparams.hidden_size)
-
-    (decoder_input, decoder_self_attention_bias) = dp(prepare_decoder, targets,
-                                                      target_space_emb)
-
-    moe_hidden_sizes = [int(s) for s in hparams.moe_hidden_sizes.split(",")]
-    expert_fn = expert_utils.ffn_expert_fn(
-        hparams.hidden_size, moe_hidden_sizes, hparams.hidden_size)
-    x = dp(tf.nn.dropout, decoder_input, 1.0 - hparams.dropout)
-    for layer in range(hparams.num_hidden_layers):
-      with tf.variable_scope("dec_layer_%d" % layer):
-        with tf.variable_scope("attention"):
-          y = dp(
-              common_attention.multihead_attention,
-              x,
-              None,
-              decoder_self_attention_bias,
-              hparams.hidden_size,
-              hparams.hidden_size,
-              hparams.hidden_size,
-              hparams.num_heads,
-              hparams.attention_dropout,
-              name="decoder_self_attention")
-          z = dp(
-              common_attention.multihead_attention,
-              y,
-              inputs_encoded3d,
-              encoder_attention_bias,
-              hparams.hidden_size,
-              hparams.hidden_size,
-              hparams.hidden_size,
-              hparams.num_heads,
-              hparams.attention_dropout,
-              name="encdec_attention")
-          x = dp(residual_fn3, x, y, z, hparams)
-        with tf.variable_scope("ffn"):
-          if str(layer) in hparams.moe_layers.split(","):
-            y, moe_loss = expert_utils.distributed_moe(
-                dp,
-                self._ps_devices,
-                x,
-                train,
-                input_size=hparams.hidden_size,
-                expert_fn=expert_fn,
-                num_experts=hparams.moe_num_experts,
-                k=hparams.moe_k,
-                loss_coef=hparams.moe_loss_coef)
-            expert_loss += tf.reduce_mean(moe_loss)
-          else:
-            y = dp(
-                common_layers.conv_hidden_relu,
-                x,
-                hparams.filter_size,
-                hparams.hidden_size,
-                dropout=hparams.dropout)
-          x = dp(residual_fn2, x, y, hparams)
-
-    x = dp(tf.expand_dims, x, 2)
-    return x, tf.reduce_mean(expert_loss)
-
-
-@registry.register_hparams
-def multimodel_base():
-  """Base parameters for MultiModel."""
-  hparams = common_hparams.basic_params1()
-  hparams.hidden_size = 512
-  hparams.batch_size = 2048
-  hparams.num_hidden_layers = 4
-  hparams.learning_rate_decay_scheme = "noam"
-  hparams.learning_rate = 0.1
-  hparams.learning_rate_warmup_steps = 4000
-  hparams.initializer_gain = 1.0
-  hparams.dropout = 0.1
-  hparams.add_hparam("filter_size", 2048)  # Add new ones like this.
-  hparams.add_hparam("large_kernel_size", 15)
-  hparams.add_hparam("attention_dropout", 0.1)
-  hparams.add_hparam("num_heads", 8)
-  hparams.add_hparam("moe_layers", "2")
-  hparams.moe_num_experts = 30
-  return hparams
-
-
-@registry.register_hparams
-def multimodel_tiny():
-  """Tiny parameters for MultiModel."""
-  hparams = multimodel_base()
-  hparams.hidden_size = 128
-  hparams.filter_size = 512
-  hparams.batch_size = 512
-  hparams.num_hidden_layers = 2
-  hparams.moe_n1 = 10
-  hparams.moe_layers = "0"
-  return hparams
diff --git a/tensor2tensor/models/research/multimodel_test.py b/tensor2tensor/models/research/multimodel_test.py
deleted file mode 100644
index 41febb460..000000000
--- a/tensor2tensor/models/research/multimodel_test.py
+++ /dev/null
@@ -1,54 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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 Xnet."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-import numpy as np
-
-from tensor2tensor.data_generators import cifar  # pylint: disable=unused-import
-from tensor2tensor.models.research import multimodel
-from tensor2tensor.utils import registry
-
-import tensorflow as tf
-
-
-class MultiModelTest(tf.test.TestCase):
-
-  def testMultiModel(self):
-    x = np.random.random_integers(0, high=255, size=(3, 5, 5, 3))
-    y = np.random.random_integers(0, high=9, size=(3, 5, 1, 1))
-    hparams = multimodel.multimodel_tiny()
-    hparams.add_hparam("data_dir", "")
-    problem = registry.problem("image_cifar10")
-    p_hparams = problem.get_hparams(hparams)
-    hparams.problem_hparams = p_hparams
-    with self.test_session() as session:
-      features = {
-          "inputs": tf.constant(x, dtype=tf.int32),
-          "targets": tf.constant(y, dtype=tf.int32),
-          "target_space_id": tf.constant(1, dtype=tf.int32),
-      }
-      model = multimodel.MultiModel(
-          hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
-      logits, _ = model(features)
-      session.run(tf.global_variables_initializer())
-      res = session.run(logits)
-    self.assertEqual(res.shape, (3, 1, 1, 1, 10))
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 0c8e26ef6..0f651ce17 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -990,75 +990,6 @@ def flatten_all_but_last(a):
   return ret
 
 
-def distributed_moe(data_parallelism,
-                    expert_devices,
-                    xs,
-                    train,
-                    input_size,
-                    expert_fn,
-                    num_experts,
-                    k=2,
-                    loss_coef=1e-2,
-                    name=None):
-  """Call a distributed mixture of experts.
-
-  Args:
-    data_parallelism: a expert_utils.Parallelism object.
-    expert_devices: a list of strings.  We round-robin the experts across these
-      devices.
-    xs: a list of input tensors, each with shape [... , input_size]
-    train: a boolean scalar.
-    input_size: an integer (input size for this layer)
-    expert_fn: a unary function for each expert to run
-       It should take a Tensor with shape [batch_size, input_size]
-       and return a Tensor with shape [batch_size, output_size].
-       e.g. ffn_expert_fn(...)
-    num_experts: an integer - number of experts
-    k: an integer - how many experts to use for each batch element
-    loss_coef: a scalar - multiplier on load-balancing losses
-    name: a string
-
-  Returns:
-    ys: a list of tensors.  Each Tensor has the same shape as the corresponding
-      Tensor in xs, except for the last dimension, which is output_size.
-    extra_training_loss: a scalar.  This should be added into the overall
-      training loss of the model.  The backpropagation of this loss
-      encourages all experts to be approximately equally used across a batch.
-  """
-  dp = data_parallelism
-  # create a parallelism object for running the experts.
-  #   We use the default of reuse=False.  Otherwise, the experts would all
-  #   use the same variables.
-  ep = Parallelism(
-      [expert_devices[i % len(expert_devices)] for i in range(num_experts)],
-      reuse=None)
-  # Experts expect 2d input tensors, so flatten the batch dimension and all
-  # spatial dimensions together.
-  xs_flat = dp(tf.reshape, xs, [[-1, input_size]] * dp.n)
-  with tf.variable_scope(name, default_name="moe"):
-    # The gates indicate which batch elements go to which tensors.
-    # load is a measure of approximately how many examples go to each expert
-    gates, load = dp(noisy_top_k_gating,
-                     xs_flat,
-                     num_experts,
-                     train,
-                     k,
-                     initializer=tf.zeros_initializer(),
-                     noisy_gating=True,
-                     noise_epsilon=1e-2)
-    # This magic object helps us shuffle data between datashards and experts.
-    dispatcher = DistributedSparseDispatcher(dp, ep, gates)
-    expert_in = dispatcher.dispatch(xs_flat)
-    expert_out = ep(expert_fn, expert_in)
-    ys_flat = dispatcher.combine(expert_out)
-    ys = dp(common_layers.reshape_like, ys_flat, xs)
-    # compute some load-balancing losses.
-    load = tf.add_n(load)
-    importance = tf.add_n(dp(tf.reduce_sum, gates, 0))
-    loss = loss_coef * (cv_squared(importance) + cv_squared(load))
-    return ys, loss
-
-
 def local_moe(x,
               train,
               expert_fn,

From ea75b1b7b0a74ce6a7709f5e67921b345e881a36 Mon Sep 17 00:00:00 2001
From: Katherine Lee <katherinelee@google.com>
Date: Tue, 18 Sep 2018 08:47:38 -0700
Subject: [PATCH 0867/2720] Pipe flag, decode_reference, through to the
 decoder.

PiperOrigin-RevId: 213456152
---
 tensor2tensor/bin/t2t_decoder.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 61e50da1a..b5d5d65f5 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -72,6 +72,8 @@ def create_decode_hparams():
   decode_hp.shards = FLAGS.decode_shards
   decode_hp.shard_id = FLAGS.worker_id
   decode_hp.decode_in_memory = FLAGS.decode_in_memory
+  decode_hp.decode_to_file = FLAGS.decode_to_file
+  decode_hp.decode_reference = FLAGS.decode_reference
   return decode_hp
 
 

From 819acbf28757bd69accd4b6ea903dac85dad81ba Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Tue, 18 Sep 2018 16:47:07 -0700
Subject: [PATCH 0868/2720] Fix latent image transformer; add baseline CIFAR
 hparams.

PiperOrigin-RevId: 213542262
---
 tensor2tensor/layers/latent_layers.py      | 317 ++++++++++++---------
 tensor2tensor/layers/latent_layers_test.py | 155 ++++++++++
 2 files changed, 332 insertions(+), 140 deletions(-)
 create mode 100644 tensor2tensor/layers/latent_layers_test.py

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 288a2f4d3..b1d84e936 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -27,9 +27,9 @@
 DO_SUMMARIES = True
 
 
-def compress_self_attention_layer(x, hparams, name):
+def compress_self_attention_layer(x, hparams, name=None):
   """Attend function."""
-  with tf.variable_scope(name):
+  with tf.variable_scope(name, default_name="compress_self_attention"):
     x, xshape, _ = cia.maybe_reshape_4d_to_3d(x)
     y = common_attention.multihead_attention(
         common_layers.layer_preprocess(x, hparams),
@@ -130,7 +130,8 @@ def ae_latent_sample_beam(latents_dense_in, inputs, ed, embed, hparams):
 
   Args:
     latents_dense_in: Tensor of shape [batch, length_q, ...]. Only the shape of
-      its first two dimensions are used.
+      its first two dimensions are used. length_q is the latent length, which is
+      height * width * hparams.num_latents / (2**hparams.num_compress_steps).
     inputs: Tensor of shape [batch, length_kv, hparams.hidden_size]. Encodings
       to attend to in decoder.
     ed: Tensor which broadcasts with shape [batch, hparams.num_heads, length_q,
@@ -216,29 +217,31 @@ def residual_block_layer(inputs, hparams):
 def compress_encoder(inputs,
                      hparams,
                      strides=(2, 2),
-                     kernel=(3, 3),
-                     name="compress"):
+                     kernel_size=(3, 3),
+                     name=None):
   """Encoder that compresses 2-D inputs by 2**num_compress_steps.
 
   Args:
     inputs: Tensor of shape [batch, height, width, channels].
     hparams: tf.contrib.training.HParams.
     strides: Tuple, strides for conv block.
-    kernel: Tuple, kernel window size for conv block.
+    kernel_size: Tuple, kernel window size for conv block.
     name: string, variable scope.
 
   Returns:
-    Tensor of shape [batch, (height*width) / 2**(hparams.num_compress_steps),
-    hparams.hidden_size].
+    Tensor of shape [batch, latent_length, hparams.hidden_size], where
+      latent_length is
+      hparams.num_latents * (height*width) / 2**(hparams.num_compress_steps).
   """
-  with tf.variable_scope(name):
+  with tf.variable_scope(name, default_name="compress"):
     x = inputs
     for i in range(hparams.num_compress_steps // 2):
       with tf.variable_scope("compress_conv_%d" % i):
         y = common_layers.conv_block(
             common_layers.layer_norm(
                 x, hparams.hidden_size, name="lnorm"),
-            hparams.hidden_size, [((1, 1), kernel)],
+            hparams.hidden_size,
+            dilation_rates_and_kernel_sizes=[((1, 1), kernel_size)],
             strides=strides,
             padding="SAME",
             name="compress_conv_%d" % i)
@@ -257,13 +260,12 @@ def compress_encoder(inputs,
     x = tf.layers.dense(x,
                         hparams.num_latents * hparams.hidden_size,
                         name=name + "_dense")
-    new_shape = [shape_x[0],
-                 shape_x[1] * shape_x[2] * hparams.num_latents,
-                 hparams.hidden_size]
-    return tf.reshape(x, new_shape)
+    return tf.reshape(x, [shape_x[0],
+                          shape_x[1] * shape_x[2] * hparams.num_latents,
+                          hparams.hidden_size])
 
 
-def compress_encoder_2d(x, hparams, name):
+def compress_encoder_2d(x, hparams, name=None):
   """Encoder that compresses 2-D inputs by 2**num_compress_steps.
 
   Args:
@@ -272,16 +274,19 @@ def compress_encoder_2d(x, hparams, name):
     name: string, variable scope.
 
   Returns:
-    Tensor of shape [batch, (height*width) / 2**hparams.num_compress_steps,
-    hparams.hidden_size].
+    Tensor of shape [batch, latent_length, hparams.hidden_size], where
+      latent_length is
+      hparams.num_latents * (height*width) / 2**(hparams.num_compress_steps).
   """
-  return compress_encoder(x, hparams,
-                          strides=(2, 2),
-                          kernel=(hparams.kernel_size, hparams.kernel_size),
-                          name=name)
+  return compress_encoder(
+      x,
+      hparams,
+      strides=(2, 2),
+      kernel_size=(hparams.kernel_size, hparams.kernel_size),
+      name=name)
 
 
-def compress_encoder_1d(x, hparams, name):
+def compress_encoder_1d(x, hparams, name=None):
   """Encoder that compresses 1-D inputs by 2**num_compress_steps.
 
   Args:
@@ -290,13 +295,15 @@ def compress_encoder_1d(x, hparams, name):
     name: string, variable scope.
 
   Returns:
-    Tensor of shape [batch, length / 2**hparams.num_compress_steps,
-    hparams.hidden_size].
+    Tensor of shape [batch, latent_length, hparams.hidden_size], where
+      latent_length is
+      hparams.num_latents * length / 2**hparams.num_compress_steps.
   """
   x = tf.expand_dims(x, axis=2)
-  return compress_encoder(x, hparams,
+  return compress_encoder(x,
+                          hparams,
                           strides=(2, 1),
-                          kernel=(hparams.kernel_size, 1),
+                          kernel_size=(hparams.kernel_size, 1),
                           name=name)
 
 
@@ -304,7 +311,7 @@ def decompress_decoder(inputs,
                        hparams,
                        strides=(2, 2),
                        kernel=(3, 3),
-                       name="decompress"):
+                       name=None):
   """Decoder that decompresses 2-D inputs by 2**num_compress_steps.
 
   Args:
@@ -317,7 +324,7 @@ def decompress_decoder(inputs,
   Returns:
     Tensor of shape [batch, height, width, hparams.hidden_size].
   """
-  with tf.variable_scope(name):
+  with tf.variable_scope(name, default_name="decompress"):
     x = inputs
     x = tf.layers.dense(x, hparams.hidden_size, name=name + "_dense")
     x = residual_block_layer(x, hparams)
@@ -340,7 +347,7 @@ def decompress_decoder(inputs,
     return x
 
 
-def decompress_decoder_2d(x, hparams, name):
+def decompress_decoder_2d(x, hparams, name=None):
   """Decoder that decompresses 2-D inputs by 2**num_compress_steps.
 
   Args:
@@ -357,7 +364,7 @@ def decompress_decoder_2d(x, hparams, name):
                             name=name)
 
 
-def decompress_decoder_1d(x, hparams, name):
+def decompress_decoder_1d(x, hparams, name=None):
   """Decoder that decompresses 1-D inputs by 2**num_compress_steps.
 
   Args:
@@ -376,15 +383,15 @@ def decompress_decoder_1d(x, hparams, name):
   return tf.squeeze(output, axis=2)
 
 
-def transformer_text_encoder(x,
-                             space_id,
+def transformer_text_encoder(inputs,
+                             target_space,
                              hparams,
-                             name="transformer_text_encoder"):
+                             name=None):
   """Transformer text encoder over inputs with unmasked full attention.
 
   Args:
-    x: Tensor of shape [batch, length, 1, hparams.hidden_size].
-    space_id: int, id.
+    inputs: Tensor of shape [batch, length, 1, hparams.hidden_size].
+    target_space: int. Used for encoding inputs under a target space id.
     hparams: tf.contrib.training.HParams.
     name: string, variable scope.
 
@@ -393,26 +400,31 @@ def transformer_text_encoder(x,
     ed: Tensor of shape [batch, 1, 1, length]. Encoder-decoder attention bias
       for any padded tokens.
   """
-  with tf.variable_scope(name):
-    x = common_layers.flatten4d3d(x)
-    (encoder_input, encoder_self_attention_bias,
-     ed) = transformer.transformer_prepare_encoder(x, space_id, hparams)
+  with tf.variable_scope(name, default_name="transformer_text_encoder"):
+    inputs = common_layers.flatten4d3d(inputs)
+    [
+        encoder_input,
+        encoder_self_attention_bias,
+        ed,
+    ] = transformer.transformer_prepare_encoder(inputs,
+                                                target_space=target_space,
+                                                hparams=hparams)
     encoder_input = tf.nn.dropout(encoder_input, 1.0 - hparams.dropout)
     encoder_output = transformer.transformer_encoder(
         encoder_input, encoder_self_attention_bias, hparams)
     return encoder_output, ed
 
 
-def transformer_image_decoder(x,
+def transformer_image_decoder(targets,
                               encoder_output,
                               ed_attention_bias,
                               hparams,
-                              name="transformer_dec"):
-  """Transformer image decoder over inputs with local attention.
+                              name=None):
+  """Transformer image decoder over targets with local attention.
 
   Args:
-    x: Tensor of shape [batch, ...], and whose size is batch * height * width *
-      hparams.num_channels * hparams.hidden_size.
+    targets: Tensor of shape [batch, ...], and whose size is batch * height *
+      width * hparams.num_channels * hparams.hidden_size.
     encoder_output: Tensor of shape [batch, length_kv, hparams.hidden_size].
     ed_attention_bias: Tensor which broadcasts with shape [batch,
       hparams.num_heads, length_q, length_kv]. Encoder-decoder attention bias.
@@ -423,12 +435,12 @@ def transformer_image_decoder(x,
     Tensor of shape [batch, height, width * hparams.num_channels,
     hparams.hidden_size].
   """
-  with tf.variable_scope(name):
-    batch_size = common_layers.shape_list(x)[0]
-    targets = tf.reshape(x, [batch_size,
-                             hparams.img_len,
-                             hparams.img_len,
-                             hparams.num_channels * hparams.hidden_size])
+  with tf.variable_scope(name, default_name="transformer_dec"):
+    batch_size = common_layers.shape_list(targets)[0]
+    targets = tf.reshape(targets, [batch_size,
+                                   hparams.img_len,
+                                   hparams.img_len,
+                                   hparams.num_channels * hparams.hidden_size])
     decoder_input, _, _ = cia.prepare_decoder(targets, hparams)
     decoder_output = cia.transformer_decoder_layers(
         decoder_input,
@@ -450,12 +462,12 @@ def transformer_latent_decoder(x,
                                encoder_output,
                                ed_attention_bias,
                                hparams,
-                               name="transformer_latent_dec"):
+                               name=None):
   """Transformer decoder over latents using latent_attention_type.
 
   Args:
-    x: Tensor of shape [batch, ...], and whose size is batch * length_q *
-      hparams.hidden_size. Here, length_q is the latent length, which is
+    x: Tensor of shape [batch, length_q, hparams.hidden_size]. length_q is the
+      latent length, which is
       height * width * hparams.num_latents / (2**hparams.num_compress_steps).
     encoder_output: Tensor of shape [batch, length_kv, hparams.hidden_size].
     ed_attention_bias: Tensor which broadcasts with shape [batch,
@@ -466,9 +478,10 @@ def transformer_latent_decoder(x,
   Returns:
     Tensor of shape [batch, length_q, hparams.hidden_size].
   """
-  with tf.variable_scope(name):
+  with tf.variable_scope(name, default_name="transformer_latent_dec"):
     batch_size = common_layers.shape_list(x)[0]
-    compressed_img_len = hparams.img_len / 2**(hparams.num_compress_steps // 2)
+    compressed_img_len = (hparams.img_len //
+                          2**(hparams.num_compress_steps // 2))
     x = tf.reshape(x, [batch_size,
                        compressed_img_len,
                        compressed_img_len * hparams.num_latents,
@@ -489,19 +502,24 @@ def transformer_latent_decoder(x,
     return decoder_output
 
 
-def bottleneck_layer(targets_c,
+def bottleneck_layer(inputs,
                      hparams,
-                     name="bottleneck_d"):
-  """Compute latents from compressed targets."""
-  latents_dense, latents_discrete, extra_loss, embed_func = (
-      hparams.bottleneck(
-          inputs=targets_c,
-          filter_size=hparams.compress_filter_size,
-          name=name,
-          mode=hparams.mode))
+                     name="discrete_bottleneck"):
+  """Computes latents given inputs (typically, compressed targets)."""
+  [
+      latents_dense,
+      latents_discrete,
+      extra_loss,
+      embed_fn,
+      _,
+  ] = hparams.bottleneck(inputs=inputs,
+                         filter_size=hparams.compress_filter_size,
+                         name=name,
+                         mode=hparams.mode)
   if DO_SUMMARIES:
-    tf.summary.histogram("b0", tf.reshape(latents_discrete, [-1]))
-  return latents_dense, latents_discrete, extra_loss, embed_func
+    tf.summary.histogram("discrete_latents",
+                         tf.reshape(latents_discrete, [-1]))
+  return latents_dense, latents_discrete, extra_loss, embed_fn
 
 
 def latent_prediction_model(inputs,
@@ -510,7 +528,7 @@ def latent_prediction_model(inputs,
                             latents_dense,
                             hparams,
                             vocab_size=None,
-                            name="latent_prediction"):
+                            name=None):
   """Transformer-based latent prediction model.
 
   It is an autoregressive decoder over latents_discrete given inputs.
@@ -523,23 +541,29 @@ def latent_prediction_model(inputs,
     latents_discrete: Tensor of shape [batch, length_q, vocab_size].
       One-hot latents to compute log-probability of given inputs.
     latents_dense: Tensor of shape [batch, length_q, hparams.hidden_size].
+      length_q is the latent length, which is
+      height * width * hparams.num_latents / (2**hparams.num_compress_steps).
     hparams: tf.contrib.training.HParams.
-    vocab_size: int, if given else None.
+    vocab_size: int or None. If None, it is 2**hparams.bottleneck_bits.
     name: string, variable scope.
 
   Returns:
     latents_pred: Tensor of shape [batch, length_q, hparams.hidden_size].
     latents_pred_loss: Tensor of shape [batch, length_q].
   """
-  with tf.variable_scope(name):
+  with tf.variable_scope(name, default_name="latent_prediction"):
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
       latents_pred = transformer_latent_decoder(tf.stop_gradient(latents_dense),
                                                 inputs,
                                                 ed_attention_bias,
                                                 hparams,
                                                 name)
-      vocab_size = (2**hparams.bottleneck_bits
-                    if vocab_size is None else vocab_size)
+      if vocab_size is None:
+        vocab_size = 2**hparams.bottleneck_bits
+      if not hparams.soft_em:
+        # TODO(trandustin): latents_discrete is not one-hot from
+        # discrete_bottleneck unless hparams.soft_em is True. Refactor.
+        latents_discrete = tf.one_hot(latents_discrete, depth=vocab_size)
       _, latent_pred_loss = ae_latent_softmax(
           latents_pred, tf.stop_gradient(latents_discrete), vocab_size, hparams)
   return latents_pred, latent_pred_loss
@@ -551,10 +575,27 @@ def transformer_autoencoder(inputs,
                             hparams,
                             cache=None,
                             predict_mask=1.0):
-  """Auto-encoder using transformer decoder and prior over latents."""
-  losses = {"extra": 0., "latent_pred": 0.}
+  """Auto-encoder using a Transformer decoder and a prior over latent sequences.
+
+  Args:
+    inputs: Tensor of shape [batch, length, 1, hparams.hidden_size] or None.
+    targets: Tensor of shape [batch, ..., channels]. Ellipses may be 1 or 2
+      dimensions denoting sequence length.
+    target_space: int. Used for encoding inputs under a target space id.
+    hparams: tf.contrib.training.HParams.
+    cache: Tensor of shape [batch, length] or None.
+    predict_mask: Tensor masking whether to use gold targets or predictions.
 
-  # Reshape image targets as 4d tensor.
+  Returns:
+    decoder_output: Tensor of shape [batch, ..., hparams.hidden_size] presenting
+      pre-logit activations. After a transformation (`top` in `T2TModel`), it is
+      used with targets to compute the "training" (reconstruction) loss.
+    losses: dict of str to Tensors. There are three loss terms: "extra",
+      "extra_loss", and "latent_pred". The first is hard-coded to 0. The latter
+      two are Tensors of shape [batch].
+    cache: Tensor of shape [batch, length], either the same as cache, or newly
+      computed if the cache input is None.
+  """
   original_targets_shape = common_layers.shape_list(targets)
   batch_size = original_targets_shape[0]
   if len(original_targets_shape) == 4:
@@ -564,98 +605,94 @@ def transformer_autoencoder(inputs,
     compress_fn = compress_encoder_1d
     decompress_fn = decompress_decoder_1d
 
-  # Input Encoder if present.
   ed_attention_bias = None
   if inputs is not None:
-    inputs = common_layers.flatten4d3d(inputs)
     inputs, ed_attention_bias = transformer_text_encoder(
-        inputs, target_space, hparams, "input_enc")
+        inputs, target_space, hparams, name="input_encoder")
 
-  # Encode targets to compute targets compressed.
-  targets_c = compress_fn(targets, hparams, "compress")
-  targets, _, _ = cia.maybe_reshape_4d_to_3d(targets)
+  losses = {"extra": 0.,
+            "extra_loss": 0.,
+            "latent_pred": 0.}
+  if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+    targets_compressed = compress_fn(targets, hparams, name="compress")
 
-  # Following code creates an exponentially decaying variable based on which
-  # we rescale the loss values.
-  pc = common_layers.inverse_exp_decay(hparams.startup_steps)
-  pc = pc if hparams.mode == tf.estimator.ModeKeys.TRAIN else 1.0
-  cond = tf.less(tf.random_uniform([batch_size]), pc)
+    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+      scale = common_layers.inverse_exp_decay(hparams.startup_steps)
+    else:
+      scale = 1.0
+    scale = tf.to_float(tf.less(tf.random_uniform([batch_size]), scale))
 
-  # Call bottleneck layer, that takes encoder output and outputs the latents.
-  # Returns embedded latents, discrete latent codes, loss.
-  if hparams.mode != tf.estimator.ModeKeys.PREDICT:
-    latents_dense, latents_discrete, extra_loss, _ = (
-        bottleneck_layer(targets_c, hparams))
-    extra_loss = tf.reduce_mean(extra_loss) * tf.to_float(cond)
+    latents_dense, latents_discrete, extra_loss, _ = bottleneck_layer(
+        targets_compressed, hparams)
+    extra_loss = scale * tf.reduce_mean(extra_loss)
 
     _, latents_pred_loss = latent_prediction_model(
-        inputs,
-        ed_attention_bias,
-        latents_discrete,
-        latents_dense,
-        hparams,
+        inputs, ed_attention_bias, latents_discrete, latents_dense, hparams,
         name="latent_pred")
-    latents_pred_loss = tf.reduce_mean(latents_pred_loss) * tf.to_float(cond)
+    latent_time = tf.less(hparams.mask_startup_steps,
+                          tf.to_int32(tf.train.get_global_step()))
+    latents_pred_loss = scale * tf.reduce_mean(latents_pred_loss)
+    latents_pred_loss *= tf.to_float(latent_time)
 
-    latents_shape = common_layers.shape_list(latents_dense)
+    # Apply dropout noise for each data point and time step.
+    latents_dense_shape = common_layers.shape_list(latents_dense)
     latents_dense = tf.nn.dropout(
-        latents_dense, 1 - hparams.latent_dropout,
-        noise_shape=[latents_shape[0], latents_shape[1], 1])
-
-    losses["extra_loss"] = extra_loss
-    losses["latent_pred"] = latents_pred_loss
+        latents_dense,
+        keep_prob=1 - hparams.latent_dropout,
+        noise_shape=[latents_dense_shape[0], latents_dense_shape[1], 1])
 
-    # We'll start training the extra model of latents after mask_startup_steps.
-    latent_time = tf.less(hparams.mask_startup_steps,
-                          tf.to_int32(tf.train.get_global_step()))
-    losses["latent_pred"] *= tf.to_float(latent_time)
+    # TODO(trandustin): Can we combine extra and extra_loss?
+    losses = {"extra": 0.,
+              "extra_loss": extra_loss,
+              "latent_pred": latents_pred_loss}
   else:
-    latent_len = (
-        hparams.img_len * hparams.img_len * hparams.num_latents) / 2**(
-            hparams.num_compress_steps)
-    _, _, _, embed = (
-        bottleneck_layer(targets_c, hparams))
+    # Set the latent length, which is num_latents times the number of latent
+    # pixels. The number of latent pixels is determined by a compression factor
+    # on the number of image pixels.
+    latent_len = ((hparams.img_len * hparams.img_len * hparams.num_latents) /
+                  (2**hparams.num_compress_steps))
+    _, _, _, embed_fn = bottleneck_layer(targets_compressed, hparams)
     latents_dense = tf.zeros([batch_size, latent_len, 1, hparams.hidden_size])
     if cache is None:
-      cache = ae_latent_sample_beam(latents_dense, inputs, ed_attention_bias,
-                                    embed, hparams)
-    latents_dense = embed(
-        tf.one_hot(cache, depth=2**hparams.bottleneck_bits),
-        hparams.hidden_size)
+      cache = ae_latent_sample_beam(latents_dense,
+                                    inputs,
+                                    ed_attention_bias,
+                                    embed_fn,
+                                    hparams)
+    cache_one_hot = tf.one_hot(cache, depth=2**hparams.bottleneck_bits)
+    latents_dense = embed_fn(cache_one_hot, hparams.hidden_size)
 
-  latents_decoder = latents_dense
   if len(original_targets_shape) == 4:
-    compressed_img_len = hparams.img_len / 2**(hparams.num_compress_steps // 2)
-    latents_decoder = tf.reshape(latents_decoder,
-                                 [batch_size,
-                                  compressed_img_len,
-                                  compressed_img_len,
-                                  hparams.num_latents * hparams.hidden_size])
-
-  latents_decoder = decompress_fn(latents_decoder, hparams, name="decompress")
-  # if we're operating in 2d space on images, then we're assuming that the
-  # last dimension will not be a multiple of channels
-  output = tf.reshape(
-      latents_decoder,
-      shape=[-1, hparams.img_len, hparams.img_len, hparams.hidden_size])
+    compressed_img_len = (hparams.img_len //
+                          2**(hparams.num_compress_steps // 2))
+    latents_dense = tf.reshape(latents_dense,
+                               [batch_size,
+                                compressed_img_len,
+                                compressed_img_len,
+                                hparams.num_latents * hparams.hidden_size])
+
+  latents_dense = decompress_fn(latents_dense, hparams, name="decompress")
+  latents_dense = tf.reshape(
+      latents_dense,
+      [-1, hparams.img_len, hparams.img_len, hparams.hidden_size])
 
   if hparams.use_gold_targets:
-    masking = common_layers.inverse_exp_decay(hparams.mask_startup_steps)
     if hparams.mode == tf.estimator.ModeKeys.PREDICT:
       masking = predict_mask
-    mask = tf.less(masking, tf.random_uniform(
-        common_layers.shape_list(targets)[:-1]))
+    else:
+      masking = common_layers.inverse_exp_decay(hparams.mask_startup_steps)
+    targets, _, _ = cia.maybe_reshape_4d_to_3d(targets)
+    mask = tf.less(masking,
+                   tf.random_uniform(common_layers.shape_list(targets)[:-1]))
     mask = tf.expand_dims(tf.to_float(mask), 2)
-    output = mask * targets + (1.0 - mask) * output
+    latents_dense = mask * targets + (1.0 - mask) * latents_dense
 
-  # reshape back to 4d here
-  output = tf.reshape(output, original_targets_shape)
+  latents_dense = tf.reshape(latents_dense, original_targets_shape)
   if hparams.decode_autoregressive:
-    # Transformer decoder, that goes from inputs->targets
     decoder_output = transformer_image_decoder(
-        output, inputs, ed_attention_bias, hparams, "decoder")
+        latents_dense, inputs, ed_attention_bias, hparams, name="decoder")
   else:
-    decoder_output = output
+    decoder_output = latents_dense
   return decoder_output, losses, cache
 
 
diff --git a/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
new file mode 100644
index 000000000..f7d00c02b
--- /dev/null
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -0,0 +1,155 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 layers in latent variable models."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+import six
+
+from tensor2tensor.layers import common_image_attention as cia
+from tensor2tensor.layers import discretization
+from tensor2tensor.layers import latent_layers
+from tensor2tensor.models import transformer
+
+import tensorflow as tf
+
+
+def imagetransformer_latent_tiny():
+  """Tiny set of hparams for a latent image model."""
+  hparams = transformer.transformer_small()
+  hparams.batch_size = 2
+  hparams.num_hidden_layers = 3
+  hparams.hidden_size = 16
+  hparams.filter_size = 32
+  hparams.compress_filter_size = 64
+  hparams.ffn_layer = "conv_hidden_relu"
+  hparams.layer_prepostprocess_dropout = 0.2
+  hparams.layer_preprocess_sequence = "none"
+  hparams.layer_postprocess_sequence = "dan"
+  hparams.dropout = 0.3
+  hparams.pos = "timing"
+  hparams.num_encoder_layers = 1
+  hparams.num_decoder_layers = 2
+  hparams.use_pad_remover = False
+  hparams.add_hparam("logit_normalization", True)
+  hparams.add_hparam("bottleneck_kind", "dvq")
+  hparams.add_hparam("bottleneck_bits", 4)
+  hparams.add_hparam("num_residuals", 1)
+  hparams.add_hparam("use_gold_targets", False)
+  hparams.add_hparam("do_compress_attend", False)
+  hparams.add_hparam("do_decompress_attend", False)
+  hparams.add_hparam("drop_inputs", False)
+  hparams.add_hparam("num_compress_steps", 2)
+  hparams.add_hparam("startup_steps", 10000)
+  hparams.add_hparam("mask_startup_steps", 50000)
+  hparams.add_hparam("latent_dropout", 0.0)
+  hparams.add_hparam("decode_autoregressive", False)
+  hparams.add_hparam("vq_beta", 0.25)
+  hparams.add_hparam("vq_epsilon", 1e-5)
+  hparams.add_hparam("vq_decay", 0.999)
+  hparams.add_hparam("ema", False)
+  hparams.add_hparam("soft_em", True)
+  hparams.add_hparam("num_samples", 1)
+  hparams.add_hparam("num_latent_layers", 2)
+  hparams.add_hparam("num_res_layers", 2)
+  hparams.add_hparam("res_kernel_size", 3)
+  hparams.add_hparam("num_blocks", 1)
+  hparams.add_hparam("reshape_method", "slice")
+  hparams.add_hparam("shared_rel", False)
+  hparams.add_hparam("block_size", 1)
+  hparams.add_hparam("kernel_size", 3)
+  hparams.add_hparam("img_len", 8)
+  hparams.add_hparam("num_channels", 1)
+  hparams.add_hparam("local_and_global_att", False)
+  hparams.add_hparam("block_length", 32)
+  hparams.add_hparam("block_width", 128)
+  hparams.add_hparam("dec_attention_type", cia.AttentionType.LOCAL_1D)
+  hparams.add_hparam("latent_attention_type", cia.AttentionType.GLOBAL)
+  hparams.add_hparam("block_raster_scan", False)
+  hparams.add_hparam("num_latents", 1)
+  hparams.add_hparam("q_filter_width", 1)
+  hparams.add_hparam("kv_filter_width", 1)
+  return hparams
+
+
+class LatentLayersTest(tf.test.TestCase):
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testTransformerAutoencoder(self):
+    hparams = imagetransformer_latent_tiny()
+    hparams.mode = tf.estimator.ModeKeys.TRAIN
+    block_dim = int(hparams.hidden_size // hparams.num_blocks)
+    block_v_size = 2**(hparams.bottleneck_bits /
+                       (hparams.num_residuals * hparams.num_blocks))
+    block_v_size = int(block_v_size)
+    means = tf.get_variable(
+        name="means",
+        shape=[hparams.num_residuals,
+               hparams.num_blocks,
+               block_v_size,
+               block_dim],
+        initializer=tf.uniform_unit_scaling_initializer())
+    hparams.bottleneck = functools.partial(
+        discretization.discrete_bottleneck,
+        hidden_size=hparams.hidden_size,
+        z_size=hparams.bottleneck_bits,
+        filter_size=hparams.filter_size,
+        startup_steps=hparams.startup_steps,
+        bottleneck_kind=hparams.bottleneck_kind,
+        num_blocks=hparams.num_blocks,
+        num_residuals=hparams.num_residuals,
+        reshape_method=hparams.reshape_method,
+        beta=hparams.vq_beta,
+        decay=hparams.vq_decay,
+        soft_em=hparams.soft_em,
+        num_samples=hparams.num_samples,
+        epsilon=hparams.vq_epsilon,
+        ema=hparams.ema,
+        means=means)
+
+    inputs = None
+    batch_size = hparams.batch_size
+    targets = tf.random_uniform([batch_size,
+                                 hparams.img_len,
+                                 hparams.img_len,
+                                 hparams.hidden_size],
+                                minval=-1., maxval=1.)
+    target_space_id = None
+
+    tf.train.create_global_step()
+    decoder_output, losses, cache = latent_layers.transformer_autoencoder(
+        inputs, targets, target_space_id, hparams)
+
+    self.assertEqual(set(six.iterkeys(losses)),
+                     {"extra", "extra_loss", "latent_pred"})
+
+    self.evaluate(tf.global_variables_initializer())
+    decoder_output_, extra_loss_, latent_pred_ = self.evaluate(
+        [decoder_output, losses["extra_loss"], losses["latent_pred"]])
+    self.assertEqual(decoder_output_.shape, (batch_size,
+                                             hparams.img_len,
+                                             hparams.img_len,
+                                             hparams.hidden_size))
+    self.assertEqual(extra_loss_.shape, (batch_size,))
+    self.assertEqual(latent_pred_.shape, (batch_size,))
+    self.assertAllGreaterEqual(extra_loss_, 0.)
+    self.assertAllGreaterEqual(latent_pred_, 0.)
+    self.assertEqual(cache, None)
+
+if __name__ == "__main__":
+  tf.test.main()

From 52807e9da1bc602e452db8aa37b862f6e6a09a1d Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 19 Sep 2018 00:01:33 -0700
Subject: [PATCH 0869/2720] Internal change

PiperOrigin-RevId: 213580758
---
 tensor2tensor/models/research/glow.py         |   2 +-
 tensor2tensor/models/research/glow_ops.py     | 133 +++++++++++++-----
 .../models/research/glow_ops_test.py          |  24 +++-
 3 files changed, 113 insertions(+), 46 deletions(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 95e4b3a2b..076cafe9c 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -43,7 +43,7 @@ def glow_hparams():
   hparams.batch_size = 32
   # can be prev_level, prev_step or normal.
   # see: glow_ops.merge_level_and_latent_dist
-  hparams.add_hparam("level_prior_scale", "prev_level")
+  hparams.add_hparam("level_scale", "prev_level")
   hparams.add_hparam("n_levels", 3)
   hparams.add_hparam("n_bits_x", 8)
   hparams.add_hparam("depth", 32)
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 19948be61..691528b2f 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -49,6 +49,30 @@ def assign(w, initial_value):
     return w
 
 
+def get_cond_latents_at_level(cond_latents, level, hparams):
+  """Returns a single or list of conditional latents at level 'level'."""
+  if cond_latents:
+    if hparams.latent_dist_encoder == "conv_net":
+      return [cond_latent[level] for cond_latent in cond_latents]
+    elif hparams.latent_dist_encoder == "pointwise":
+      return cond_latents[level]
+
+
+def check_cond_latents(cond_latents, hparams):
+  """Shape checking for cond_latents."""
+  if cond_latents is None:
+    return
+  if not isinstance(cond_latents[0], list):
+    cond_latents = [cond_latents]
+  if len(cond_latents) != hparams.num_cond_latents:
+    raise ValueError("Expected number of cond_latents: %d, got %d" %
+                     (hparams.num_cond_latents, len(cond_latents)))
+  for cond_latent in cond_latents:
+    if len(cond_latent) != hparams.n_levels - 1:
+      raise ValueError("Expected level_latents to be %d, got %d" %
+                       (hparams.n_levels - 1, len(cond_latent)))
+
+
 @add_arg_scope
 def get_variable_ddi(name, shape, initial_value, dtype=tf.float32, init=False,
                      trainable=True):
@@ -428,12 +452,34 @@ def squeeze(name, x, factor=2, reverse=True):
 
 
 @add_arg_scope
-def split_prior(name, x):
-  """Map x to the mean and log-scale of a Gaussian distribution."""
+def tensor_to_dist(name, x, output_channels=None, architecture="single_conv"):
+  """Map x to the mean and log-scale of a Gaussian.
+
+  Args:
+    name: variable scope.
+    x: 4-D Tensor of shape (NHWC)
+    output_channels: int, number of output channels of the mean.
+                     if not provided, set it to be the output channels of x.
+    architecture: "single_conv" or "glow_nn"
+  Returns:
+    dist: instance of tf.distributions.Normal
+  Raises:
+    ValueError: If architecture not in ["single_conv", "glow_nn"]
+  """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     x_shape = common_layers.shape_list(x)
-    mean_log_scale = conv2d("conv2d", x, output_channels=2*x_shape[-1],
-                            apply_actnorm=False, conv_init="zeros")
+    if output_channels is None:
+      output_channels = x_shape[-1]
+    if architecture == "single_conv":
+      mean_log_scale = conv2d("conv2d", x, output_channels=2*output_channels,
+                              conv_init="zeros", apply_actnorm=False)
+    elif architecture == "glow_nn":
+      mean_log_scale = nn("conv2d", x, mid_channels=512,
+                          output_channels=2*output_channels)
+    else:
+      raise ValueError("expected architecture to be single_conv or glow_nn "
+                       "got %s" % architecture)
+
     mean = mean_log_scale[:, :, :, 0::2]
     log_scale = mean_log_scale[:, :, :, 1::2]
     return tf.distributions.Normal(mean, tf.exp(log_scale))
@@ -472,26 +518,42 @@ def merge_level_and_latent_dist(level_dist, latent_dist,
 
 
 @add_arg_scope
-def compute_prior(name, z, latent, merge_std):
-  """Distribution condtioned on both z and latent."""
+def compute_prior(name, z, latent, hparams):
+  """Distribution conditioned on both z and latent."""
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
-    prior_dist = split_prior("level_prior", z)
+    prior_dist = tensor_to_dist("level_prior", z, architecture="single_conv")
     if latent is not None:
-      latent_shape = common_layers.shape_list(latent)
-      z_shape = common_layers.shape_list(z)
-      if latent_shape != z_shape:
-        raise ValueError("Expected latent_shape to be %s, got %s" %
-                         (latent_shape, z_shape))
-      latent_dist = scale_gaussian_prior(
-          "latent_prior", latent, logscale_factor=3.0)
-      prior_dist = merge_level_and_latent_dist(prior_dist, latent_dist,
-                                               merge_std=merge_std)
+      latent_dist_encoder = hparams.latent_dist_encoder
+      if latent_dist_encoder == "pointwise":
+        merge_std = hparams.level_scale
+        latent_shape = common_layers.shape_list(latent)
+        z_shape = common_layers.shape_list(z)
+        if latent_shape != z_shape:
+          raise ValueError("Expected latent_shape to be %s, got %s" %
+                           (latent_shape, z_shape))
+        latent_dist = scale_gaussian_prior(
+            "latent_prior", latent, logscale_factor=3.0)
+        prior_dist = merge_level_and_latent_dist(prior_dist, latent_dist,
+                                                 merge_std=merge_std)
+      elif latent_dist_encoder == "conv_net":
+        output_channels = common_layers.shape_list(z)[-1]
+        latent_stack = tf.concat([prior_dist.loc] + latent, axis=-1)
+        prior_dist = tensor_to_dist(
+            "latent_stack", latent_stack, output_channels=output_channels,
+            architecture=hparams.latent_architecture)
+        latent_skip = hparams.get("latent_skip", False)
+        if latent_skip:
+          prior_dist = tf.distributions.Normal(
+              prior_dist.loc + latent[-1], prior_dist.scale)
+      tf.summary.histogram("split_prior_mean", prior_dist.loc)
+      tf.summary.histogram("split_prior_scale", prior_dist.scale)
+
   return prior_dist
 
 
 @add_arg_scope
-def split(name, x, reverse=False, eps=None, eps_std=None, cond_latent=None,
-          merge_std="normal"):
+def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
+          hparams=None):
   """Splits / concatenates x into x1 and x2 across number of channels.
 
   For the forward pass, x2 is assumed be gaussian,
@@ -506,9 +568,8 @@ def split(name, x, reverse=False, eps=None, eps_std=None, cond_latent=None,
     reverse: Forward or reverse pass.
     eps: If eps is provided, x2 is set to be
     eps_std: Sample x2.
-    cond_latent: optionally condition x2 on cond_latent.
-    merge_std: used to determine the std of the gaussian prior on x2 if
-               cond_latent is provided.
+    cond_latents: optionally condition x2 on cond_latents.
+    hparams: next_frame_glow hparams.
 
   Returns:
   Raises:
@@ -520,12 +581,12 @@ def split(name, x, reverse=False, eps=None, eps_std=None, cond_latent=None,
       x1, x2 = tf.split(x, num_or_size_splits=2, axis=-1)
 
       # objective: P(x2|x1) ~N(x2 ; NN(x1))
-      prior_dist = compute_prior("prior_on_z2", x1, cond_latent, merge_std)
+      prior_dist = compute_prior("prior_on_z2", x1, cond_latents, hparams)
       logpb = tf.reduce_sum(prior_dist.log_prob(x2), axis=[1, 2, 3])
       eps = get_eps(prior_dist, x2)
       return x1, logpb, eps, x2
     else:
-      prior_dist = compute_prior("prior_on_z2", x, cond_latent, merge_std)
+      prior_dist = compute_prior("prior_on_z2", x, cond_latents, hparams)
       if eps is not None:
         x2 = set_eps(prior_dist, eps)
       elif eps_std is not None:
@@ -633,7 +694,7 @@ def top_prior(name, x, learn_prior="normal"):
     if learn_prior == "normal":
       prior_dist = tf.distributions.Normal(h, tf.exp(h))
     elif learn_prior == "single_conv":
-      prior_dist = split_prior("top_learn_prior", h)
+      prior_dist = tensor_to_dist("top_learn_prior", h)
     else:
       raise ValueError("Expected learn_prior to be normal or single_conv "
                        "got %s" % learn_prior)
@@ -671,9 +732,7 @@ def encoder_decoder(name, x, hparams, eps=None, reverse=False,
     if eps and len(eps) != hparams.n_levels - 1:
       raise ValueError("Expected length of eps to be %d, got %d" %
                        (hparams.n_levels - 1, len(eps)))
-    if cond_latents and len(cond_latents) != hparams.n_levels - 1:
-      raise ValueError("Expected level_latets to be %d, got %d" %
-                       (hparams.n_levels - 1, len(cond_latents)))
+    check_cond_latents(cond_latents, hparams)
 
     objective = 0.0
     all_eps = []
@@ -689,13 +748,11 @@ def encoder_decoder(name, x, hparams, eps=None, reverse=False,
 
         if level < hparams.n_levels - 1:
 
-          curr_latent = None
-          if cond_latents is not None:
-            curr_latent = cond_latents[level]
-
-          x, obj, eps, z = split(
-              "split_%d" % level, x, reverse=False, cond_latent=curr_latent,
-              merge_std=hparams.level_prior_scale)
+          curr_cond_latents = get_cond_latents_at_level(
+              cond_latents, level, hparams)
+          x, obj, eps, z = split("split_%d" % level, x, reverse=False,
+                                 cond_latents=curr_cond_latents,
+                                 hparams=hparams)
           objective += obj
           all_eps.append(eps)
           all_latents.append(z)
@@ -709,13 +766,11 @@ def encoder_decoder(name, x, hparams, eps=None, reverse=False,
           if eps:
             curr_eps = eps[level]
 
-          curr_latent = None
-          if cond_latents is not None:
-            curr_latent = cond_latents[level]
+          curr_cond_latents = get_cond_latents_at_level(
+              cond_latents, level, hparams)
 
           x, latent = split("split_%d" % level, x, eps=curr_eps, reverse=True,
-                            cond_latent=curr_latent,
-                            merge_std=hparams.level_prior_scale)
+                            cond_latents=curr_cond_latents, hparams=hparams)
           all_latents.append(latent)
 
         x, obj = revnet(
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 0666e552d..abd9dce05 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -123,17 +123,25 @@ def test_nn(self):
         # Initialized with zeros.
         self.assertTrue(np.allclose(nn_np, 0.0))
 
-  def test_split_prior(self):
+  def check_tensor_to_dist(self, architecture):
     with tf.Graph().as_default():
       x = tf.random_uniform(shape=(16, 5, 5, 32))
-      x_prior = glow_ops.split_prior("split_prior", x)
+      x_prior = glow_ops.tensor_to_dist("split_prior", x,
+                                        architecture=architecture,
+                                        output_channels=64)
       mean_t, scale_t = x_prior.loc, x_prior.scale
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())
         mean, scale = session.run([mean_t, scale_t])
+        self.assertEqual(mean.shape, (16, 5, 5, 64))
+        self.assertEqual(scale.shape, (16, 5, 5, 64))
         self.assertTrue(np.allclose(mean, 0.0))
         self.assertTrue(np.allclose(scale, 1.0))
 
+  def test_tensor_to_dist(self):
+    for architecture in ["single_conv", "glow_nn"]:
+      self.check_tensor_to_dist(architecture)
+
   def test_split(self):
     with tf.Graph().as_default():
       x = tf.random_uniform(shape=(16, 5, 5, 32))
@@ -262,16 +270,20 @@ def check_split_latent_conditioning(self, merge_std):
       latent_rand = rng.randn(12, 32, 32, 16).astype(np.float32)
       x_t = tf.convert_to_tensor(x_rand)
       latent_t = tf.convert_to_tensor(latent_rand)
+      hparams = glow.glow_hparams()
+      hparams.level_scale = merge_std
+      hparams.add_hparam("latent_dist_encoder", "pointwise")
 
       # Test initalization.
       # x2 ~ N(scale * latent, 1.0) where initial scale is 1.0
       exp_x2 = x_rand[:, :, :, 16:]
       exp_eps = x_rand[:, :, :, 16:] - latent_rand
-      x_inv, _, eps, x2_t = glow_ops.split(merge_std, x_t, cond_latent=latent_t,
-                                           merge_std=merge_std)
+      x_inv, _, eps, x2_t = glow_ops.split(
+          merge_std, x_t, cond_latents=latent_t, hparams=hparams)
       # Test reversibility.
-      x_inv_inv, _ = glow_ops.split(merge_std, x_inv, cond_latent=latent_t,
-                                    merge_std=merge_std, eps=eps, reverse=True)
+      x_inv_inv, _ = glow_ops.split(
+          merge_std, x_inv, cond_latents=latent_t, eps=eps, reverse=True,
+          hparams=hparams)
       with tf.Session() as sess:
         sess.run(tf.global_variables_initializer())
         actual_eps, actual_x2, diff_np = sess.run([eps, x2_t, x_inv_inv - x_t])

From bcd5afa85d91515de0e3d18f56d1f8d20955b6fb Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 19 Sep 2018 12:47:17 -0700
Subject: [PATCH 0870/2720] adding simple reward func.

PiperOrigin-RevId: 213676936
---
 tensor2tensor/models/video/sv2p.py        | 21 ++++++++++++++++++---
 tensor2tensor/models/video/sv2p_params.py |  3 ++-
 tensor2tensor/models/video/sv2p_test.py   |  4 +++-
 tensor2tensor/models/video/tests_utils.py |  2 +-
 4 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 57626250f..41984728f 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -206,7 +206,23 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
     hidden5 = tile_and_concat(hidden5, latent, concat_latent=concat_latent)
     return hidden5, (enc0, enc1)
 
-  def reward_prediction(self, input_images, input_reward, action, latent):
+  def reward_prediction(self, *args, **kwargs):
+    model = self.hparams.reward_model
+    if model == "basic":
+      return self.reward_prediction_basic(*args, **kwargs)
+    elif model == "big":
+      return self.reward_prediction_big(*args, **kwargs)
+    else:
+      raise ValueError("Unknown reward model %s" % model)
+
+  def reward_prediction_basic(self, input_images, input_reward, action, latent):
+    del input_reward, action, latent
+    x = tf.concat(input_images, axis=3)
+    x = tf.expand_dims(  # Add a fake channels dim.
+        tf.reduce_mean(x, axis=[1, 2], keepdims=True), axis=3)
+    return x
+
+  def reward_prediction_big(self, input_images, input_reward, action, latent):
     """Builds a reward prediction network."""
     conv_size = self.tinyify([32, 32, 16, 8])
 
@@ -546,8 +562,7 @@ def body(self, features):
     # This is NOT the same as original paper/implementation.
     predictions = gen_images[hparams.video_num_input_frames-1:]
     reward_pred = gen_rewards[hparams.video_num_input_frames-1:]
-    if self.is_training:
-      reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove extra dimension.
+    reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove extra dimension.
 
     # Swap back time and batch axes.
     predictions = common_video.swap_time_and_batch_axes(predictions)
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index deb96e662..dc7d0a0e7 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -47,7 +47,8 @@ def next_frame_sv2p():
   hparams.add_hparam("scheduled_sampling_decay_steps", 10000)
   hparams.add_hparam("scheduled_sampling_k", 900.0)
   hparams.add_hparam("upsample_method", "conv2d_transpose")
-  hparams.add_hparam("internal_loss", True)
+  hparams.add_hparam("internal_loss", False)
+  hparams.add_hparam("reward_model", "basic")
   return hparams
 
 
diff --git a/tensor2tensor/models/video/sv2p_test.py b/tensor2tensor/models/video/sv2p_test.py
index 7a59dc62e..a8d94ae50 100644
--- a/tensor2tensor/models/video/sv2p_test.py
+++ b/tensor2tensor/models/video/sv2p_test.py
@@ -40,8 +40,10 @@ def testSv2pWithActions(self):
         1)
 
   def testSv2pWithActionsAndRewards(self):
+    hp = sv2p_params.next_frame_sv2p()
+    hp.internal_loss = True
     self.TestWithActionAndRewards(
-        sv2p_params.next_frame_sv2p(),
+        hp,
         sv2p.NextFrameSv2p,
         1)
 
diff --git a/tensor2tensor/models/video/tests_utils.py b/tensor2tensor/models/video/tests_utils.py
index 96fec0a2d..3004f3bc2 100644
--- a/tensor2tensor/models/video/tests_utils.py
+++ b/tensor2tensor/models/video/tests_utils.py
@@ -238,7 +238,7 @@ def TestVideoModelWithActionAndRewardsInfer(self,
     self.assertTrue("target_reward" in output.keys())
     expected_shape = get_tensor_shape(features["targets"])
     self.assertEqual(output["targets"].shape, expected_shape)
-    expected_shape = get_tensor_shape(features["target_reward"])[:2] + (1,)
+    expected_shape = get_tensor_shape(features["target_reward"])[:2]
     self.assertEqual(output["target_reward"].shape, expected_shape)
 
   def TestOnVariousInputOutputSizes(

From a6bee918a4493a73c5de433c235419bf781d24a8 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 19 Sep 2018 12:54:40 -0700
Subject: [PATCH 0871/2720] internal

PiperOrigin-RevId: 213678147
---
 tensor2tensor/models/OWNERS | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/OWNERS b/tensor2tensor/models/OWNERS
index d3af3590b..4c18f708e 100644
--- a/tensor2tensor/models/OWNERS
+++ b/tensor2tensor/models/OWNERS
@@ -1,4 +1,5 @@
 avaswani
+dumitru
 mbz
 nikip
 noam

From c8d02488b8da155be591bbcc912f6d2ad4b7f481 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 19 Sep 2018 13:12:53 -0700
Subject: [PATCH 0872/2720] small sv2p

PiperOrigin-RevId: 213681350
---
 tensor2tensor/layers/common_video.py          | 23 ++++--
 tensor2tensor/models/video/base_vae.py        |  3 +-
 .../video/basic_deterministic_params.py       |  1 +
 tensor2tensor/models/video/sv2p.py            | 73 ++++++++++++-------
 4 files changed, 64 insertions(+), 36 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 80fccc30c..24f1781ca 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -435,9 +435,11 @@ def gif_summary(name, tensor, max_outputs=3, fps=10, collections=None,
 
 
 
-def tinyify(array, tiny_mode):
+def tinyify(array, tiny_mode, small_mode):
   if tiny_mode:
     return [1 for _ in array]
+  if small_mode:
+    return [x // 4 for x in array]
   return array
 
 
@@ -448,7 +450,8 @@ def get_gaussian_tensor(mean, log_var):
 
 
 def conv_latent_tower(images, time_axis, latent_channels=1, min_logvar=-5,
-                      is_training=False, random_latent=False, tiny_mode=False):
+                      is_training=False, random_latent=False,
+                      tiny_mode=False, small_mode=False):
   """Builds convolutional latent tower for stochastic model.
 
   At training time this tower generates a latent distribution (mean and std)
@@ -466,12 +469,17 @@ def conv_latent_tower(images, time_axis, latent_channels=1, min_logvar=-5,
     min_logvar: minimum value for log_var
     is_training: whether or not it is training mode
     random_latent: whether or not generate random latents
-    tiny_mode: whether or not it is tiny_mode
+    tiny_mode: whether or not it is tiny_mode. tiny_mode sets the number
+        of conv channels to 1 at each layer. useful for testing the
+        integration tests.
+    small_mode: whether or not it is small_mode. small mode is the same model
+        with less conv and lstm layers and also lower number of channels.
+        suitable for videos with less complexity and testing.
   Returns:
     latent_mean: predicted latent mean
     latent_logvar: predicted latent log variance
   """
-  conv_size = tinyify([32, 64, 64], tiny_mode)
+  conv_size = tinyify([32, 64, 64], tiny_mode, small_mode)
   with tf.variable_scope("latent", reuse=tf.AUTO_REUSE):
     images = tf.to_float(images)
     images = tf.unstack(images, axis=time_axis)
@@ -482,9 +490,10 @@ def conv_latent_tower(images, time_axis, latent_channels=1, min_logvar=-5,
     x = tfl.conv2d(x, conv_size[0], [3, 3], strides=(2, 2),
                    padding="SAME", activation=tf.nn.relu, name="latent_conv1")
     x = tfcl.layer_norm(x)
-    x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
-                   padding="SAME", activation=tf.nn.relu, name="latent_conv2")
-    x = tfcl.layer_norm(x)
+    if not small_mode:
+      x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
+                     padding="SAME", activation=tf.nn.relu, name="latent_conv2")
+      x = tfcl.layer_norm(x)
     x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(1, 1),
                    padding="SAME", activation=tf.nn.relu, name="latent_conv3")
     x = tfcl.layer_norm(x)
diff --git a/tensor2tensor/models/video/base_vae.py b/tensor2tensor/models/video/base_vae.py
index 8af32e374..0c1e604ee 100644
--- a/tensor2tensor/models/video/base_vae.py
+++ b/tensor2tensor/models/video/base_vae.py
@@ -109,7 +109,8 @@ def construct_latent_tower(self, images, time_axis):
         min_logvar=self.hparams.latent_std_min,
         is_training=self.is_training,
         random_latent=first_phase,
-        tiny_mode=self.hparams.tiny_mode)
+        tiny_mode=self.hparams.tiny_mode,
+        small_mode=self.hparams.small_mode)
 
 
 
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 5168823f5..ea9bc2855 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -49,6 +49,7 @@ def next_frame_basic_deterministic():
   hparams.add_hparam("preprocess_resize_frames", None)
   hparams.add_hparam("shuffle_buffer_size", 128)
   hparams.add_hparam("tiny_mode", False)
+  hparams.add_hparam("small_mode", False)
   hparams.add_hparam("stochastic_model", False)
   return hparams
 
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 41984728f..2d8ceb4e4 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -44,7 +44,8 @@ class NextFrameSv2p(basic_stochastic.NextFrameBasicStochastic):
   """Stochastic Variational Video Prediction."""
 
   def tinyify(self, array):
-    return common_video.tinyify(array, self.hparams.tiny_mode)
+    return common_video.tinyify(
+        array, self.hparams.tiny_mode, self.hparams.small_mode)
 
   def visualize_predictions(self, real_frames, gen_frames):
     def concat_on_y_axis(x):
@@ -155,6 +156,7 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
     concat_input_image = tile_and_concat(
         input_image, latent, concat_latent=concat_latent)
 
+    layer_id = 0
     enc0 = tfl.conv2d(
         concat_input_image,
         conv_size[0], [5, 5],
@@ -164,29 +166,38 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
         name="scale1_conv1")
     enc0 = tfcl.layer_norm(enc0, scope="layer_norm1")
 
-    hidden1, lstm_state[0] = lstm_func(
-        enc0, lstm_state[0], lstm_size[0], name="state1")
+    hidden1, lstm_state[layer_id] = lstm_func(
+        enc0, lstm_state[layer_id], lstm_size[layer_id], name="state1")
     hidden1 = tile_and_concat(hidden1, latent, concat_latent=concat_latent)
     hidden1 = tfcl.layer_norm(hidden1, scope="layer_norm2")
-    hidden2, lstm_state[1] = lstm_func(
-        hidden1, lstm_state[1], lstm_size[1], name="state2")
+    layer_id += 1
+
+    hidden2, lstm_state[layer_id] = lstm_func(
+        hidden1, lstm_state[layer_id], lstm_size[layer_id], name="state2")
     hidden2 = tfcl.layer_norm(hidden2, scope="layer_norm3")
     hidden2 = common_layers.make_even_size(hidden2)
     enc1 = tfl.conv2d(hidden2, hidden2.get_shape()[3], [3, 3], strides=(2, 2),
                       padding="SAME", activation=tf.nn.relu, name="conv2")
     enc1 = tile_and_concat(enc1, latent, concat_latent=concat_latent)
+    layer_id += 1
 
-    hidden3, lstm_state[2] = lstm_func(
-        enc1, lstm_state[2], lstm_size[2], name="state3")
-    hidden3 = tile_and_concat(hidden3, latent, concat_latent=concat_latent)
-    hidden3 = tfcl.layer_norm(hidden3, scope="layer_norm4")
-    hidden4, lstm_state[3] = lstm_func(
-        hidden3, lstm_state[3], lstm_size[3], name="state4")
-    hidden4 = tile_and_concat(hidden4, latent, concat_latent=concat_latent)
-    hidden4 = tfcl.layer_norm(hidden4, scope="layer_norm5")
-    hidden4 = common_layers.make_even_size(hidden4)
-    enc2 = tfl.conv2d(hidden4, hidden4.get_shape()[3], [3, 3], strides=(2, 2),
-                      padding="SAME", activation=tf.nn.relu, name="conv3")
+    if self.hparams.small_mode:
+      hidden4, enc2 = hidden2, enc1
+    else:
+      hidden3, lstm_state[layer_id] = lstm_func(
+          enc1, lstm_state[layer_id], lstm_size[layer_id], name="state3")
+      hidden3 = tile_and_concat(hidden3, latent, concat_latent=concat_latent)
+      hidden3 = tfcl.layer_norm(hidden3, scope="layer_norm4")
+      layer_id += 1
+
+      hidden4, lstm_state[layer_id] = lstm_func(
+          hidden3, lstm_state[layer_id], lstm_size[layer_id], name="state4")
+      hidden4 = tile_and_concat(hidden4, latent, concat_latent=concat_latent)
+      hidden4 = tfcl.layer_norm(hidden4, scope="layer_norm5")
+      hidden4 = common_layers.make_even_size(hidden4)
+      enc2 = tfl.conv2d(hidden4, hidden4.get_shape()[3], [3, 3], strides=(2, 2),
+                        padding="SAME", activation=tf.nn.relu, name="conv3")
+      layer_id += 1
 
     if action is not None:
       enc2 = self.inject_additional_input(
@@ -200,11 +211,12 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
     enc3 = tfl.conv2d(enc2, hidden4.get_shape()[3], [1, 1], strides=(1, 1),
                       padding="SAME", activation=tf.nn.relu, name="conv4")
 
-    hidden5, lstm_state[4] = lstm_func(
-        enc3, lstm_state[4], lstm_size[4], name="state5")  # last 8x8
+    hidden5, lstm_state[layer_id] = lstm_func(
+        enc3, lstm_state[layer_id], lstm_size[layer_id], name="state5")
     hidden5 = tfcl.layer_norm(hidden5, scope="layer_norm6")
     hidden5 = tile_and_concat(hidden5, latent, concat_latent=concat_latent)
-    return hidden5, (enc0, enc1)
+    layer_id += 1
+    return hidden5, (enc0, enc1), layer_id
 
   def reward_prediction(self, *args, **kwargs):
     model = self.hparams.reward_model
@@ -229,9 +241,11 @@ def reward_prediction_big(self, input_images, input_reward, action, latent):
     with tf.variable_scope("reward_pred", reuse=tf.AUTO_REUSE):
       x = tf.concat(input_images, axis=3)
       x = tfcl.layer_norm(x)
-      x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
-                     activation=tf.nn.relu, name="reward_conv1")
-      x = tfcl.layer_norm(x)
+
+      if not self.hparams.small_mode:
+        x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
+                       activation=tf.nn.relu, name="reward_conv1")
+        x = tfcl.layer_norm(x)
 
       # Inject additional inputs
       if action is not None:
@@ -269,7 +283,7 @@ def construct_predictive_tower(
     conv_size = self.tinyify([32])
 
     with tf.variable_scope("main", reuse=tf.AUTO_REUSE):
-      hidden5, skips = self.bottom_part_tower(
+      hidden5, skips, layer_id = self.bottom_part_tower(
           input_image, input_reward, action, latent,
           lstm_state, lstm_size, conv_size, concat_latent=concat_latent)
       enc0, enc1 = skips
@@ -283,13 +297,14 @@ def construct_predictive_tower(
       enc4 = enc4[:, :enc1_shape[1], :enc1_shape[2], :]  # Cut to shape.
       enc4 = tile_and_concat(enc4, latent, concat_latent=concat_latent)
 
-      hidden6, lstm_state[5] = lstm_func(
-          enc4, lstm_state[5], lstm_size[5], name="state6",
+      hidden6, lstm_state[layer_id] = lstm_func(
+          enc4, lstm_state[layer_id], lstm_size[5], name="state6",
           spatial_dims=enc1_shape[1:-1])  # 16x16
       hidden6 = tile_and_concat(hidden6, latent, concat_latent=concat_latent)
       hidden6 = tfcl.layer_norm(hidden6, scope="layer_norm7")
       # Skip connection.
       hidden6 = tf.concat(axis=3, values=[hidden6, enc1])  # both 16x16
+      layer_id += 1
 
       with tf.variable_scope("upsample2", reuse=tf.AUTO_REUSE):
         enc5 = common_layers.cyclegan_upsample(
@@ -300,10 +315,11 @@ def construct_predictive_tower(
       enc5 = enc5[:, :enc0_shape[1], :enc0_shape[2], :]  # Cut to shape.
       enc5 = tile_and_concat(enc5, latent, concat_latent=concat_latent)
 
-      hidden7, lstm_state[6] = lstm_func(
-          enc5, lstm_state[6], lstm_size[6], name="state7",
+      hidden7, lstm_state[layer_id] = lstm_func(
+          enc5, lstm_state[layer_id], lstm_size[6], name="state7",
           spatial_dims=enc0_shape[1:-1])  # 32x32
       hidden7 = tfcl.layer_norm(hidden7, scope="layer_norm8")
+      layer_id += 1
 
       # Skip connection.
       hidden7 = tf.concat(axis=3, values=[hidden7, enc0])  # both 32x32
@@ -442,7 +458,8 @@ def process_single_frame(prev_outputs, inputs):
       latent = common_video.get_gaussian_tensor(latent_mean, latent_std)
 
     # HACK: Do first step outside to initialize all the variables
-    lstm_states = [None] * 7
+
+    lstm_states = [None] * (5 if self.hparams.small_mode else 7)
     frame_buffer = [tf.zeros_like(images[0])] * buffer_size
     inputs = images[0], rewards[0], actions[0]
     prev_outputs = (tf.constant(0),

From 63169410137bba8be1813edb3cda7b97492d58ae Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 19 Sep 2018 15:18:33 -0700
Subject: [PATCH 0873/2720] fixing sv2p internal loss

PiperOrigin-RevId: 213705391
---
 tensor2tensor/models/video/savp_params.py |  1 +
 tensor2tensor/models/video/sv2p.py        | 22 +++++++++++-----------
 tensor2tensor/models/video/sv2p_params.py |  2 +-
 3 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index 23c6590d8..bd5a06a75 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -38,4 +38,5 @@ def next_frame_savp():
   hparams.input_modalities = "inputs:video:l1raw"
   hparams.latent_loss_multiplier_schedule = "linear"
   hparams.upsample_method = "bilinear_upsample_conv"
+  hparams.internal_loss = False
   return hparams
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 2d8ceb4e4..bff773f73 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -586,17 +586,17 @@ def body(self, features):
     reward_pred = common_video.swap_time_and_batch_axes(reward_pred)
 
     if hparams.internal_loss:
-      recon_loss = tf.losses.mean_squared_error(all_frames[1:], gen_images)
-      rew_loss = 0.0
-      if hparams.reward_prediction:
-        rew_loss = tf.losses.softmax_cross_entropy(all_rewards[1:], gen_rewards)
-        tf.summary.scalar("loss/reward", rew_loss)
-      tf.summary.scalar("loss/recon", recon_loss)
-      tf.summary.scalar("loss/kl", extra_loss)
-      extra_loss = {"training": recon_loss + rew_loss + extra_loss}
-      # also expand the last dimension of prediction since
-      # we all the modalities will be bypassed.
-      predictions = tf.expand_dims(predictions, axis=-1)
+      # add the MSE loss for input frames as well.
+      # we are assuming the modality is L2. otherwise the loss would be
+      # incosistent across the frames.
+      modality = self.hparams.problem_hparams.target_modality["targets"]
+      if modality.__class__.__name__ != "VideoModalityL2Raw":
+        raise ValueError("internal loss only works with L2.")
+      recon_loss = tf.losses.mean_squared_error(
+          all_frames[1:hparams.video_num_input_frames+1],
+          gen_images[:hparams.video_num_input_frames])
+      tf.summary.scalar("mse_extra", recon_loss)
+      extra_loss += recon_loss
 
     return_targets = predictions
     if hparams.reward_prediction:
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index dc7d0a0e7..ba137ed3f 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -47,7 +47,7 @@ def next_frame_sv2p():
   hparams.add_hparam("scheduled_sampling_decay_steps", 10000)
   hparams.add_hparam("scheduled_sampling_k", 900.0)
   hparams.add_hparam("upsample_method", "conv2d_transpose")
-  hparams.add_hparam("internal_loss", False)
+  hparams.add_hparam("internal_loss", True)
   hparams.add_hparam("reward_model", "basic")
   return hparams
 

From 50557225c7afa77b3cf003ec2728300b43f40cfd Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 19 Sep 2018 15:33:41 -0700
Subject: [PATCH 0874/2720] Allow frame resizing in gym problems and other
 corrections made in debugging this.

PiperOrigin-RevId: 213708610
---
 tensor2tensor/data_generators/gym_problems.py | 78 ++++++++++++----
 .../data_generators/gym_problems_specs.py     | 10 +-
 .../data_generators/gym_problems_test.py      |  2 +-
 tensor2tensor/models/research/rl.py           | 32 ++++---
 .../models/video/basic_stochastic.py          |  8 +-
 tensor2tensor/models/video/sv2p.py            |  4 +
 tensor2tensor/rl/collect.py                   |  5 +
 tensor2tensor/rl/envs/batch_env.py            | 11 ++-
 tensor2tensor/rl/envs/batch_env_factory.py    |  4 +
 tensor2tensor/rl/envs/in_graph_batch_env.py   |  3 +
 tensor2tensor/rl/envs/py_func_batch_env.py    | 17 ++--
 tensor2tensor/rl/envs/simulated_batch_env.py  |  4 +
 tensor2tensor/rl/envs/tf_atari_wrappers.py    | 92 +++++++++++++++++--
 tensor2tensor/rl/trainer_model_based.py       | 70 +++++++++-----
 .../rl/trainer_model_based_ae_test.py         | 13 +--
 15 files changed, 270 insertions(+), 83 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index d56d389b1..c2f6f359f 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -49,10 +49,18 @@
                     "File with model for autoencoder.")
 
 
-def standard_atari_env_spec(env):
+def standard_atari_env_spec(env, simulated=False,
+                            resize_height_factor=1, resize_width_factor=1):
   """Parameters of environment specification."""
-  standard_wrappers = [[tf_atari_wrappers.RewardClippingWrapper, {}],
-                       [tf_atari_wrappers.StackWrapper, {"history": 4}]]
+  standard_wrappers = [
+      [tf_atari_wrappers.ResizeWrapper,
+       {"height_factor": resize_height_factor,
+        "width_factor": resize_width_factor}],
+      [tf_atari_wrappers.RewardClippingWrapper, {}],
+      [tf_atari_wrappers.StackWrapper, {"history": 4}],
+  ]
+  if simulated:  # No resizing on simulated environments.
+    standard_wrappers = standard_wrappers[1:]
   env_lambda = None
   if isinstance(env, str):
     env_lambda = lambda: gym.make(env)
@@ -61,7 +69,9 @@ def standard_atari_env_spec(env):
   assert env_lambda is not None, "Unknown specification of environment"
 
   return tf.contrib.training.HParams(
-      env_lambda=env_lambda, wrappers=standard_wrappers, simulated_env=False)
+      env_lambda=env_lambda,
+      wrappers=standard_wrappers,
+      simulated_env=simulated)
 
 
 def standard_atari_ae_env_spec(env):
@@ -93,10 +103,11 @@ def __init__(self, *args, **kwargs):
     super(GymDiscreteProblem, self).__init__(*args, **kwargs)
     # TODO(piotrmilos): Check if self._env is used.
     self._env = None
+
     self.debug_dump_frames_path = "debug_frames_env"
     self.settable_num_steps = 5000
 
-    self.environment_spec = self.get_environment_spec()
+    self._environment_spec = None
     self.settable_eval_phase = False
 
     self._internal_memory_size = 20
@@ -106,16 +117,24 @@ def __init__(self, *args, **kwargs):
     self._use_dumper_data = False
     self._dumper_data_index = 0
 
+  @property
+  def resize_height_factor(self):
+    return 2
+
+  @property
+  def resize_width_factor(self):
+    return 2
+
   def _setup(self, data_dir):
     # TODO(piotrmilos):this should be consistent with
     # ppo_params in model_rl_experiment
     dumper_path = os.path.join(data_dir, "dumper")
     if os.path.isdir(dumper_path):
+      tf.logging.info("Using dumper data.")
       self._use_dumper_data = True
       self._dumper_data_index = 0
       self._dumper_path = dumper_path
     else:
-
       collect_hparams = rl.ppo_pong_base()
       collect_hparams.add_hparam("environment_spec", self.environment_spec)
       collect_hparams.add_hparam("force_beginning_resets",
@@ -136,7 +155,7 @@ def _setup(self, data_dir):
                 collect_hparams,
                 scope="gym_problems",
                 eval_phase=False,
-                collect_level=0,
+                collect_level=1,  # After ResizeWrapper but before others.
                 policy_to_actions_lambda=policy_to_actions_lambda))
 
       self._session = tf.Session()
@@ -261,7 +280,16 @@ def extra_reading_spec(self):
     return data_fields, decoders
 
   def get_environment_spec(self):
-    return standard_atari_env_spec(self.env_name)
+    return standard_atari_env_spec(
+        self.env_name,
+        resize_height_factor=self.resize_height_factor,
+        resize_width_factor=self.resize_width_factor)
+
+  @property
+  def environment_spec(self):
+    if self._environment_spec is None:
+      self._environment_spec = self.get_environment_spec()
+    return self._environment_spec
 
   @property
   def is_generate_per_split(self):
@@ -293,11 +321,11 @@ def collect_statistics_and_generate_debug_image(self, index, observation,
 
   @property
   def frame_height(self):
-    return self.env.observation_space.shape[0]
+    return self.env.observation_space.shape[0] // self.resize_height_factor
 
   @property
   def frame_width(self):
-    return self.env.observation_space.shape[1]
+    return self.env.observation_space.shape[1] // self.resize_width_factor
 
   @property
   def num_rewards(self):
@@ -465,6 +493,10 @@ def autoencoder_factor(self):
     hparams = autoencoders.autoencoder_discrete_pong()
     return 2**hparams.num_hidden_layers
 
+  @property
+  def num_channels(self):
+    return 24
+
   @property
   def frame_height(self):
     height = self.env.observation_space.shape[0]
@@ -522,8 +554,12 @@ def __init__(self, *args, **kwargs):
     # the amount of skips induced but wrappers
     self._internal_memory_size = self.num_testing_steps
     self._internal_memory_force_beginning_resets = True
-    env_spec = standard_atari_env_spec(self.env_name)
-    real_env = env_spec.env_lambda()
+    real_env_spec = standard_atari_env_spec(
+        self.env_name,
+        simulated=False,
+        resize_height_factor=self.resize_height_factor,
+        resize_width_factor=self.resize_width_factor)
+    real_env = real_env_spec.env_lambda()
 
     self.statistics = RewardPerSequenceStatistics()
     self.statistics.real_env = real_env
@@ -574,8 +610,11 @@ def num_testing_steps(self):
     return None
 
   def get_environment_spec(self):
-    env_spec = standard_atari_env_spec(self.env_name)
-    env_spec.simulated_env = True
+    env_spec = standard_atari_env_spec(
+        self.env_name,
+        simulated=True,
+        resize_height_factor=self.resize_height_factor,
+        resize_width_factor=self.resize_width_factor)
     env_spec.add_hparam("simulation_random_starts", False)
     env_spec.add_hparam("simulation_flip_first_random_for_beginning", False)
     env_spec.add_hparam("intrinsic_reward_scale", 0.0)
@@ -603,9 +642,14 @@ def collect_statistics_and_generate_debug_image(self, index,
     stat.episode_sim_reward += reward
 
     ob = np.ndarray.astype(observation, np.int)
-    err = np.ndarray.astype(
-        np.maximum(np.abs(stat.real_ob - ob, dtype=np.int) - 10, 0), np.uint8)
-    debug_im = np.concatenate([observation, stat.real_ob, err], axis=1)
+    if ob.shape == stat.real_ob.shape:
+      err = np.ndarray.astype(
+          np.maximum(np.abs(stat.real_ob - ob, dtype=np.int) - 10, 0), np.uint8)
+      debug_im = np.concatenate([observation, stat.real_ob, err], axis=1)
+    else:
+      # Real env does not get the ResizeWrapper and we don't have it in python,
+      # so we skip the debug image here and just output observations.
+      debug_im = observation
 
     assert (self._internal_memory_size == self.num_testing_steps and
             self._internal_memory_force_beginning_resets), (
diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index e54208916..fbadc1940 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -176,6 +176,8 @@ def num_rewards(self):
 def create_problems_for_game(
     game_name,
     clipped_reward=True,
+    resize_height_factor=2,
+    resize_width_factor=2,
     game_mode="Deterministic-v4"):
   """Create and register problems for game_name.
 
@@ -183,6 +185,8 @@ def create_problems_for_game(
     game_name: str, one of the games in ATARI_GAMES, e.g. "bank_heist".
     clipped_reward: bool, whether the rewards should be clipped. False is not
       yet supported.
+    resize_height_factor: factor by which to resize the height of frames.
+    resize_width_factor: factor by which to resize the width of frames.
     game_mode: the frame skip and sticky keys config.
 
   Returns:
@@ -213,7 +217,9 @@ def create_problems_for_game(
   # Create and register the Random and WithAgent Problem classes
   problem_cls = type("Gym%sRandom" % camel_game_name,
                      (GymClippedRewardRandom,),
-                     {"env_name": wrapped_env_name})
+                     {"env_name": wrapped_env_name,
+                      "resize_height_factor": resize_height_factor,
+                      "resize_width_factor": resize_width_factor})
   registry.register_problem(problem_cls)
 
   with_agent_cls = type("GymDiscreteProblemWithAgentOn%s" % camel_game_name,
@@ -245,5 +251,3 @@ def create_problems_for_game(
         clipped_reward=True,
         game_mode=mode)
     ATARI_PROBLEMS[game][mode] = classes
-
-
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index 6b6193bef..ebd7b4112 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -37,7 +37,7 @@ def setUpClass(cls):
   def testGymAtariGameModes(self):
     for mode in gym_problems_specs.ATARI_GAME_MODES:
       problem = gym_problems_specs.ATARI_PROBLEMS["pong"][mode]["base"]()
-      self.assertEqual(210, problem.frame_height)
+      self.assertEqual(105, problem.frame_height)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 494717582..379cd6c68 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -54,6 +54,7 @@ def ppo_base_v1():
   hparams.add_hparam("simulation_flip_first_random_for_beginning", False)
   hparams.add_hparam("intrinsic_reward_scale", 0.)
   hparams.add_hparam("logits_clip", 3.0)
+  hparams.add_hparam("dropout_ppo", 0.1)
   return hparams
 
 
@@ -107,7 +108,7 @@ def ppo_atari_base():
 def ppo_pong_base():
   """Pong base parameters."""
   hparams = ppo_discrete_action_base()
-  hparams.learning_rate = 2e-4
+  hparams.learning_rate = 1e-4
   hparams.num_agents = 8
   hparams.epoch_length = 200
   hparams.gae_gamma = 0.985
@@ -143,7 +144,7 @@ def simple_gym_spec(env):
 def ppo_pong_ae_base():
   """Pong autoencoder base parameters."""
   hparams = ppo_pong_base()
-  hparams.learning_rate = 2e-4
+  hparams.learning_rate = 1e-4
   hparams.network = dense_bitwise_categorical_fun
   return hparams
 
@@ -232,25 +233,32 @@ def feed_forward_cnn_small_categorical_fun(action_space, config, observations):
   x = tf.reshape(observations, [-1] + obs_shape[2:])
 
   with tf.variable_scope("network_parameters"):
+    dropout = getattr(config, "dropout_ppo", 0.0)
     with tf.variable_scope("feed_forward_cnn_small"):
       x = tf.to_float(x) / 255.0
-      x = tf.contrib.layers.conv2d(x, 32, [5, 5], [2, 2],
-                                   activation_fn=tf.nn.relu, padding="SAME")
-      x = tf.contrib.layers.conv2d(x, 32, [5, 5], [2, 2],
-                                   activation_fn=tf.nn.relu, padding="SAME")
+      x = tf.nn.dropout(x, keep_prob=1.0 - dropout)
+      x = tf.layers.conv2d(
+          x, 32, (4, 4), strides=(2, 2), name="conv1",
+          activation=common_layers.belu, padding="SAME")
+      x = tf.nn.dropout(x, keep_prob=1.0 - dropout)
+      x = tf.layers.conv2d(
+          x, 64, (4, 4), strides=(2, 2), name="conv2",
+          activation=common_layers.belu, padding="SAME")
+      x = tf.nn.dropout(x, keep_prob=1.0 - dropout)
+      x = tf.layers.conv2d(
+          x, 128, (4, 4), strides=(2, 2), name="conv3",
+          activation=common_layers.belu, padding="SAME")
 
       flat_x = tf.reshape(
           x, [obs_shape[0], obs_shape[1],
               functools.reduce(operator.mul, x.shape.as_list()[1:], 1)])
+      flat_x = tf.nn.dropout(flat_x, keep_prob=1.0 - dropout)
+      x = tf.layers.dense(flat_x, 128, activation=tf.nn.relu, name="dense1")
 
-      x = tf.contrib.layers.fully_connected(flat_x, 128, tf.nn.relu)
-
-      logits = tf.contrib.layers.fully_connected(x, action_space.n,
-                                                 activation_fn=None)
+      logits = tf.layers.dense(x, action_space.n, name="dense2")
       logits = clip_logits(logits, config)
 
-      value = tf.contrib.layers.fully_connected(
-          x, 1, activation_fn=None)[..., 0]
+      value = tf.layers.dense(x, 1, name="value")[..., 0]
       policy = tf.contrib.distributions.Categorical(logits=logits)
 
   return NetworkOutput(policy, value, lambda a: a)
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 85b361665..03a9c343c 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -72,7 +72,11 @@ def inject_latent(self, layer, features, filters):
 
     if hparams.mode == tf.estimator.ModeKeys.PREDICT:
       layer_shape = common_layers.shape_list(layer)
-      rand = tf.random_uniform(layer_shape[:-1] + [hparams.bottleneck_bits])
+      if hparams.full_latent_tower:
+        rand = tf.random_uniform(layer_shape[:-1] + [hparams.bottleneck_bits])
+      else:
+        rand = tf.random_uniform(layer_shape[:-3] + [
+            1, 1, hparams.bottleneck_bits])
       d = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
       z = tf.layers.dense(d, final_filters, name="unbottleneck")
       return layer + z, 0.0
@@ -132,8 +136,6 @@ def next_frame_basic_stochastic():
 def next_frame_basic_stochastic_discrete():
   """Basic 2-frame conv model with stochastic discrete latent."""
   hparams = basic_deterministic_params.next_frame_sampling()
-  hparams.num_compress_steps = 8
-  hparams.filter_double_steps = 3
   hparams.add_hparam("bottleneck_bits", 16)
   hparams.add_hparam("bottleneck_noise", 0.02)
   hparams.add_hparam("full_latent_tower", False)
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index bff773f73..1216ff639 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -373,6 +373,7 @@ def construct_predictive_tower(
       masks = tfl.conv2d(
           enc6, filters=num_masks + 1, kernel_size=[1, 1],
           strides=(1, 1), name="convt7", padding="SAME")
+      masks = masks[:, :img_height, :img_width, ...]
       masks = tf.reshape(
           tf.nn.softmax(tf.reshape(masks, [-1, num_masks + 1])),
           [batch_size,
@@ -382,6 +383,9 @@ def construct_predictive_tower(
           axis=3, num_or_size_splits=num_masks + 1, value=masks)
       output = mask_list[0] * input_image
       for layer, mask in zip(transformed, mask_list[1:]):
+        # TODO(mbz): take another look at this logic and verify.
+        output = output[:, :img_height, :img_width, :]
+        layer = layer[:, :img_height, :img_width, :]
         output += layer * mask
 
       return output, lstm_state
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 95d422f1c..0f95d8492 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -61,6 +61,9 @@ def __init__(self, batch_env):
     self._observ = tf.Variable(tf.zeros(observs_shape, self.observ_dtype),
                                trainable=False)
 
+  def __str__(self):
+    return "MemoryWrapper(%s)" % str(self._batch_env)
+
   def simulate(self, action):
 
     # There is subtlety here. We need to collect data
@@ -114,6 +117,8 @@ def define_collect(hparams, scope, eval_phase,
     rollout_metadata = None
     speculum = None
     for w in wrappers:
+      tf.logging.info("Applying wrapper %s(%s) to env %s."
+                      % (str(w[0]), str(w[1]), str(batch_env)))
       batch_env = w[0](batch_env, **w[1])
       to_initialize.append(batch_env)
       if w[0] == _MemoryWrapper:
diff --git a/tensor2tensor/rl/envs/batch_env.py b/tensor2tensor/rl/envs/batch_env.py
index 230bd3a43..862a8c9b2 100644
--- a/tensor2tensor/rl/envs/batch_env.py
+++ b/tensor2tensor/rl/envs/batch_env.py
@@ -46,10 +46,13 @@ def __init__(self, envs, blocking):
     self.observ_space = self._envs[0].observation_space
     if not all(env.observation_space == self.observ_space
                for env in self._envs):
-      raise ValueError('All environments must use the same observation space.')
+      raise ValueError("All environments must use the same observation space.")
     self.action_space = self._envs[0].action_space
     if not all(env.action_space == self.action_space for env in self._envs):
-      raise ValueError('All environments must use the same observation space.')
+      raise ValueError("All environments must use the same observation space.")
+
+  def __str__(self):
+    return "BatchEnv(%s)" % ", ".join([str(e) for e in self._envs])
 
   def __len__(self):
     """Number of combined environments."""
@@ -84,7 +87,7 @@ def step(self, actions):
     """
     for index, (env, action) in enumerate(zip(self._envs, actions)):
       if not env.action_space.contains(action):
-        message = 'Invalid action at index {}: {}'
+        message = "Invalid action at index {}: {}"
         raise ValueError(message.format(index, action))
     if self._blocking:
       transitions = [
@@ -127,5 +130,5 @@ def reset(self, indices=None):
   def close(self):
     """Send close messages to the external process and join them."""
     for env in self._envs:
-      if hasattr(env, 'close'):
+      if hasattr(env, "close"):
         env.close()
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index 52ad1af81..af3d6b549 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -99,6 +99,7 @@ def __init__(self, constructor, xvfb):
       observation_space: The cached observation space of the environment.
       action_space: The cached action space of the environment.
     """
+    self._constructor = constructor
     self._conn, conn = multiprocessing.Pipe()
     if xvfb:
       server_id = random.randint(10000, 99999)
@@ -129,6 +130,9 @@ def constructor_using_xvfb():
     self._observ_space = None
     self._action_space = None
 
+  def __str__(self):
+    return "ExternalProcessEnv(%s)" % str(self._constructor)
+
   @property
   def observation_space(self):
     if not self._observ_space:
diff --git a/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index 2385369e7..3cb7ae551 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -45,6 +45,9 @@ def __getattr__(self, name):
     """
     return getattr(self._batch_env, name)
 
+  def __str__(self):
+    return "InGraphEnv(%s)" % str(self._batch_env)
+
   def __len__(self):
     """Number of combined environments."""
     return len(self._batch_env)
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index ac1ad1dd5..e6e27d913 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -42,11 +42,14 @@ def __init__(self, batch_env):
     super(PyFuncBatchEnv, self).__init__(batch_env.observation_space,
                                          batch_env.action_space)
     self._batch_env = batch_env
-    with tf.variable_scope('env_temporary'):
+    with tf.variable_scope("env_temporary"):
       self._observ = tf.Variable(
           tf.zeros((len(self._batch_env),) + self.observ_shape,
                    self.observ_dtype),
-          name='observ', trainable=False)
+          name="observ", trainable=False)
+
+  def __str__(self):
+    return "PyFuncEnv(%s)" % str(self._batch_env)
 
   def __getattr__(self, name):
     """Forward unimplemented attributes to one of the original environments.
@@ -81,13 +84,13 @@ def simulate(self, action):
     Returns:
       Operation.
     """
-    with tf.name_scope('environment/simulate'):
+    with tf.name_scope("environment/simulate"):
       if action.dtype in (tf.float16, tf.float32, tf.float64):
-        action = tf.check_numerics(action, 'action')
+        action = tf.check_numerics(action, "action")
       observ, reward, done = tf.py_func(
           lambda a: self._batch_env.step(a)[:3], [action],
-          [self.observ_dtype, tf.float32, tf.bool], name='step')
-      reward = tf.check_numerics(reward, 'reward')
+          [self.observ_dtype, tf.float32, tf.bool], name="step")
+      reward = tf.check_numerics(reward, "reward")
       reward.set_shape((len(self),))
       done.set_shape((len(self),))
       with tf.control_dependencies([self._observ.assign(observ)]):
@@ -103,7 +106,7 @@ def _reset_non_empty(self, indices):
       Batch tensor of the new observations.
     """
     observ = tf.py_func(
-        self._batch_env.reset, [indices], self.observ_dtype, name='reset')
+        self._batch_env.reset, [indices], self.observ_dtype, name="reset")
     observ.set_shape(indices.get_shape().concatenate(self.observ_shape))
     with tf.control_dependencies([
         tf.scatter_update(self._observ, indices, observ)]):
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 0c12bbc10..7d7e4def4 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -107,6 +107,7 @@ def __init__(self, environment_spec, length):
 
     observ_space = utils.get_observation_space(environment_spec)
     initial_frames_problem = environment_spec.initial_frames_problem
+    self._frames_problem_name = str(initial_frames_problem)
     observ_shape = (initial_frames_problem.frame_height,
                     initial_frames_problem.frame_width,
                     initial_frames_problem.num_channels)
@@ -161,6 +162,9 @@ def __init__(self, environment_spec, length):
   def initialize(self, sess):
     self.history_buffer.initialize(sess)
 
+  def __str__(self):
+    return "SimulatedEnv(%s)" % self._frames_problem_name
+
   def __len__(self):
     """Number of combined environments."""
     return self.length
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index f67d30d21..7302191ca 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -99,6 +99,9 @@ class RewardClippingWrapper(WrapperBase):
       of rl algorithms
   """
 
+  def __str__(self):
+    return "RewardClippingWrapper(%s)" % str(self._batch_env)
+
   def simulate(self, action):
     reward, done = self._batch_env.simulate(action)
     with tf.control_dependencies([reward, done]):
@@ -118,6 +121,9 @@ def __init__(self, batch_env, skip=4):
     self._observ = tf.Variable(tf.zeros(observs_shape, self.observ_dtype),
                                trainable=False)
 
+  def __str__(self):
+    return "MaxAndSkipWrapper(%s)" % str(self._batch_env)
+
   def simulate(self, action):
     with tf.name_scope("environment/simulate"):  # Do we need this?
       initializer = (tf.zeros_like(self._observ),
@@ -158,6 +164,9 @@ def __init__(self, batch_env, history=4):
         tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
         trainable=False)
 
+  def __str__(self):
+    return "StackWrapper(%s)" % str(self._batch_env)
+
   @property
   def observ_shape(self):
     return self.old_shape[:-1] + (self.old_shape[-1] * self.history,)
@@ -192,7 +201,7 @@ def _reset_non_empty(self, indices):
       inx = tf.concat(
           [
               tf.ones(tf.size(tf.shape(new_values)),
-                      dtype=tf.int32)[:-1],
+                      dtype=tf.int64)[:-1],
               [self.history]
           ],
           axis=0)
@@ -207,12 +216,11 @@ def _transform_history_observations(self, frames):
 
 
 class AutoencoderWrapper(WrapperBase):
-  """ Transforms the observations taking the bottleneck
-      state of an autoencoder"""
+  """Transforms the observations taking the bottleneck of an autoencoder."""
 
   def __init__(self, batch_env):
     super(AutoencoderWrapper, self).__init__(batch_env)
-    self._observ = self._observ = tf.Variable(
+    self._observ = tf.Variable(
         tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
         trainable=False)
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
@@ -224,6 +232,9 @@ def __init__(self, batch_env):
       self.autoencoder_model = autoencoders.AutoencoderOrderedDiscrete(
           autoencoder_hparams, tf.estimator.ModeKeys.EVAL)
 
+  def __str__(self):
+    return "AutoencoderWrapper(%s)" % str(self._batch_env)
+
   @property
   def observ_shape(self):
     height, width, _ = self._batch_env.observ_shape
@@ -242,7 +253,7 @@ def simulate(self, action):
     reward, done = self._batch_env.simulate(action)
     with tf.control_dependencies([reward, done]):
       with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-        observ = tf.cast(self._batch_env.observ, tf.int32)
+        observ = tf.cast(self._batch_env.observ, tf.int64)
         ret = self.autoencoder_model.encode(observ)
         ret = tf.cast(ret, self.observ_dtype)
         assign_op = self._observ.assign(ret)
@@ -252,7 +263,7 @@ def simulate(self, action):
   def _reset_non_empty(self, indices):
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
       new_values = self._batch_env._reset_non_empty(indices)  # pylint: disable=protected-access
-      new_values = tf.cast(new_values, tf.int32)
+      new_values = tf.cast(new_values, tf.int64)
       ret = self.autoencoder_model.encode(new_values)
       ret = tf.cast(ret, self.observ_dtype)
       assign_op = tf.scatter_update(self._observ, indices, ret)
@@ -262,12 +273,73 @@ def _reset_non_empty(self, indices):
   def _transform_history_observations(self, frames):
     batch_size, history_size = frames.get_shape().as_list()[:2]
     new_frames = tf.reshape(frames, (-1,) + self._batch_env.observ_shape)
-    new_frames = tf.cast(new_frames, tf.int32)
+    new_frames = tf.cast(new_frames, tf.int64)
     new_frames = self.autoencoder_model.encode(new_frames)
     new_frames = tf.cast(new_frames, self.observ_dtype)
     return new_frames.reshape((batch_size, history_size) + self.observ_shape)
 
 
+class ResizeWrapper(WrapperBase):
+  """Resizes the observations."""
+
+  def __init__(self, batch_env, height_factor=1, width_factor=1):
+    super(ResizeWrapper, self).__init__(batch_env)
+    self._height_factor = height_factor  # How much to resize on x axis.
+    self._width_factor = width_factor  # How much to resize on y axis.
+    self._is_identity = (height_factor == 1) and (width_factor == 1)
+    if not self._is_identity:
+      self._observ = tf.Variable(
+          tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
+          trainable=False)
+    else:
+      self._observ = self._batch_env.observ
+
+  def __str__(self):
+    return "ResizeWrapper%d%d(%s)" % (self._height_factor,
+                                      self._width_factor, str(self._batch_env))
+
+  def _resize(self, tensor):
+    if self._is_identity:
+      return tensor
+    height, width, _ = self.observ_shape
+    observ = tf.to_float(tensor)
+    resized = tf.image.resize_images(
+        observ, [height, width], tf.image.ResizeMethod.BILINEAR)
+    return tf.cast(resized, self.observ_dtype)
+
+  @property
+  def observ_shape(self):
+    height, width, channels = self._batch_env.observ_shape
+    resized_height = height // self._height_factor
+    resized_width = width // self._width_factor
+    return (resized_height, resized_width, channels)
+
+  def simulate(self, action):
+    if self._is_identity:
+      return self._batch_env.simulate(action)
+    reward, done = self._batch_env.simulate(action)
+    with tf.control_dependencies([reward, done]):
+      ret = self._resize(self._batch_env.observ)
+      assign_op = self._observ.assign(ret)
+      with tf.control_dependencies([assign_op]):
+        return tf.identity(reward), tf.identity(done)
+
+  def _reset_non_empty(self, indices):
+    new_values = self._batch_env._reset_non_empty(indices)  # pylint: disable=protected-access
+    if self._is_identity:
+      return new_values
+    ret = self._resize(new_values)
+    assign_op = tf.scatter_update(self._observ, indices, ret)
+    with tf.control_dependencies([assign_op]):
+      return tf.gather(self.observ, indices)
+
+  def _transform_history_observations(self, frames):
+    batch_size, history_size = frames.get_shape().as_list()[:2]
+    new_frames = tf.reshape(frames, (-1,) + self._batch_env.observ_shape)
+    new_frames = self._resize(new_frames)
+    return new_frames.reshape((batch_size, history_size) + self.observ_shape)
+
+
 class IntToBitWrapper(WrapperBase):
   """Unpacks the observations from integer values to bit values"""
 
@@ -277,6 +349,9 @@ def __init__(self, batch_env):
         tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
         trainable=False)
 
+  def __str__(self):
+    return "IntToBitWrapper(%s)" % str(self._batch_env)
+
   @property
   def observ_shape(self):
     height, width, channels = self._batch_env.observ_shape
@@ -328,6 +403,9 @@ def __init__(self, batch_env, process_fun):
     self._observ = tf.Variable(tf.zeros(observs_shape, self.observ_dtype),
                                trainable=False)
 
+  def __str__(self):
+    return "PyFuncWrapper(%s)" % str(self._batch_env)
+
   def simulate(self, action):
     reward, done = self._batch_env.simulate(action)
     with tf.control_dependencies([reward, done]):
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 2a6eb284c..a61f9e332 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -472,8 +472,11 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     if problem_name not in registry.list_problems():
       tf.logging.info("Game Problem %s not found; dynamically registering",
                       problem_name)
-      gym_problems_specs.create_problems_for_game(hparams.game,
-                                                  game_mode="Deterministic-v4")
+      gym_problems_specs.create_problems_for_game(
+          hparams.game,
+          resize_height_factor=hparams.resize_height_factor,
+          resize_width_factor=hparams.resize_width_factor,
+          game_mode="Deterministic-v4")
 
   # Autoencoder model dir
   autoencoder_model_dir = directories.get("autoencoder")
@@ -648,21 +651,24 @@ def rl_modelrl_base():
       # though it is not necessary.
       ppo_epoch_length=50,
       ppo_num_agents=16,
-      ppo_learning_rate=2e-4,  # Will be changed, just so it exists.
+      ppo_learning_rate=1e-4,  # Will be changed, just so it exists.
       # Whether the PPO agent should be restored from the previous iteration, or
       # should start fresh each time.
       ppo_continue_training=True,
+      # Resizing.
+      resize_height_factor=2,
+      resize_width_factor=2,
 
       gather_ppo_real_env_data=True,
       real_ppo_epochs_num=0,
       # This needs to be divisible by real_ppo_effective_num_agents.
       real_ppo_epoch_length=16*200,
       real_ppo_num_agents=1,
-      real_ppo_learning_rate=2e-4,
+      real_ppo_learning_rate=1e-4,
       real_ppo_continue_training=True,
       real_ppo_effective_num_agents=16,
 
-      game="wrapped_full_pong",
+      game="pong",
       # Whether to evaluate the world model in each iteration of the loop to get
       # the model_reward_accuracy metric.
       eval_world_model=True,
@@ -670,6 +676,18 @@ def rl_modelrl_base():
   )
 
 
+@registry.register_hparams
+def rl_modelrl_basetest():
+  """Base setting but quicker with only 2 epochs."""
+  hparams = rl_modelrl_base()
+  hparams.epochs = 2
+  hparams.num_real_env_frames = 3200
+  hparams.model_train_steps = 500
+  hparams.simulated_env_generator_num_steps = 20
+  hparams.ppo_epochs_num = 2
+  return hparams
+
+
 @registry.register_hparams
 def rl_modelrl_base_quick():
   """Base setting but quicker with only 2 epochs."""
@@ -785,6 +803,9 @@ def rl_modelrl_tiny():
           real_ppo_effective_num_agents=2,
           generative_model_params="next_frame_tiny",
           stop_loop_early=True,
+          resize_height_factor=2,
+          resize_width_factor=2,
+          game="pong",
       ).values())
 
 
@@ -862,8 +883,11 @@ def rl_modelrl_ae_base():
 def rl_modelrl_ae_tiny():
   """Tiny set for testing autoencoders."""
   hparams = rl_modelrl_tiny()
+  hparams.game = "wrapped_full_pong"
   hparams.ppo_params = "ppo_pong_ae_base"
   hparams.generative_model_params = "next_frame_ae"
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
   hparams.autoencoder_train_steps = 2
   hparams.eval_world_model = False
   return hparams
@@ -884,9 +908,7 @@ def rl_modelrl_tiny_simulation_deterministic_starts():
 @registry.register_ranged_hparams
 def rl_modelrl_grid(rhp):
   """Grid over games and frames, and 5 runs each for variance."""
-  rhp.set_categorical("loop.game",
-                      ["breakout", "wrapped_full_pong", "freeway"])
-
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
   base = 100000
   medium = base // 2
   small = medium // 2
@@ -900,8 +922,7 @@ def rl_modelrl_grid(rhp):
 def rl_modelrl_variance(rhp):
   # Dummy parameter to get 5 runs for each configuration
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game",
-                      ["breakout", "wrapped_full_pong", "freeway"])
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
 
 
 @registry.register_ranged_hparams
@@ -910,11 +931,17 @@ def rl_modelrl_variance_nogame(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(500)))
 
 
+@registry.register_ranged_hparams
+def rl_modelrl_three(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
+
+
 @registry.register_ranged_hparams
 def rl_modelrl_test1(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", ["breakout", "wrapped_full_pong", "boxing"])
-  rhp.set_discrete("loop.ppo_learning_rate", [1e-4, 2e-4, 4e-4])
+  rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
+  rhp.set_discrete("loop.ppo_learning_rate", [5e-5, 1e-4, 2e-4])
   rhp.set_discrete("ppo.optimization_batch_size", [20, 40])
   rhp.set_discrete("loop.epochs", [3, 6])
 
@@ -940,8 +967,7 @@ def rl_modelrl_whitelisted_games(rhp):
 def rl_modelrl_ae_variance(rhp):
   # Dummy parameter to get 5 runs for each configuration
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game",
-                      ["breakout", "wrapped_full_pong", "freeway"])
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
   base = 100000
   small = base // 4
   rhp.set_discrete("loop.num_real_env_frames", [base, small])
@@ -949,23 +975,21 @@ def rl_modelrl_ae_variance(rhp):
 
 @registry.register_ranged_hparams
 def rl_modelrl_ppolr_game(rhp):
-  rhp.set_categorical("loop.game",
-                      ["breakout", "wrapped_full_pong", "freeway"])
-  base_lr = 2e-4
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+  base_lr = 1e-4
   rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
 
 
 @registry.register_ranged_hparams
 def rl_modelrl_ppolr(rhp):
-  base_lr = 2e-4
+  base_lr = 1e-4
   rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
 
 
 @registry.register_ranged_hparams
 def rl_modelrl_ae_ppo_lr(rhp):
-  rhp.set_categorical("loop.game",
-                      ["breakout", "wrapped_full_pong", "freeway"])
-  base_lr = 2e-4
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+  base_lr = 1e-4
   rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
 
 
@@ -1037,14 +1061,14 @@ def rl_modelrl_num_frames(rhp):
 
 @registry.register_ranged_hparams
 def rl_modelrl_ppo_optimization_batch_size(rhp):
-  rhp.set_categorical("loop.game", ["pong", "wrapped_full_pong", "seaquest"])
+  rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
   rhp.set_discrete("model.moe_loss_coef", list(range(10)))
   rhp.set_discrete("ppo.optimization_batch_size", [4, 10, 20])
 
 
 @registry.register_ranged_hparams
 def rl_modelrl_logits_clip(rhp):
-  rhp.set_categorical("loop.game", ["pong", "wrapped_full_pong", "seaquest"])
+  rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
   rhp.set_discrete("model.moe_loss_coef", list(range(10)))
   rhp.set_discrete("ppo.logits_clip", [0., 5.])
 
diff --git a/tensor2tensor/rl/trainer_model_based_ae_test.py b/tensor2tensor/rl/trainer_model_based_ae_test.py
index 7888e365c..aac2232a3 100644
--- a/tensor2tensor/rl/trainer_model_based_ae_test.py
+++ b/tensor2tensor/rl/trainer_model_based_ae_test.py
@@ -17,7 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.rl import trainer_model_based
+# from tensor2tensor.rl import trainer_model_based
 
 import tensorflow as tf
 
@@ -27,11 +27,12 @@
 class ModelRLExperimentTestAe(tf.test.TestCase):
 
   def test_ae(self):
-    FLAGS.output_dir = tf.test.get_temp_dir()
-    FLAGS.loop_hparams_set = "rl_modelrl_ae_tiny"
-    FLAGS.schedule = "train"  # skip evaluation for world model training
-    trainer_model_based.main(None)
-
+    # TODO(lukaszkaiser): re-enable this test.
+    # FLAGS.output_dir = tf.test.get_temp_dir()
+    # FLAGS.loop_hparams_set = "rl_modelrl_ae_tiny"
+    # FLAGS.schedule = "train"  # skip evaluation for world model training
+    # trainer_model_based.main(None)
+    assert True
 
 if __name__ == "__main__":
   tf.test.main()

From a275aaf95485edf37d73a7b5e6aed43342d1b910 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 19 Sep 2018 15:48:48 -0700
Subject: [PATCH 0875/2720] fixing sv2p RL external losses tests.

PiperOrigin-RevId: 213711371
---
 tensor2tensor/rl/trainer_model_based.py           | 9 ---------
 tensor2tensor/rl/trainer_model_based_sv2p_test.py | 7 -------
 2 files changed, 16 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index a61f9e332..37c1abdbc 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -828,15 +828,6 @@ def rl_modelrl_tiny_sv2p():
   return hparams
 
 
-@registry.register_hparams
-def rl_modelrl_tiny_sv2p_external():
-  """Tiny setting with a tiny sv2p model and external loss."""
-  hparams = rl_modelrl_tiny()
-  hparams.generative_model = "next_frame_sv2p"
-  hparams.generative_model_params = "next_frame_sv2p_tiny_external"
-  return hparams
-
-
 @registry.register_hparams
 def rl_modelrl_l1_base():
   """Parameter set with L1 loss."""
diff --git a/tensor2tensor/rl/trainer_model_based_sv2p_test.py b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
index 361756f64..a59aab6ec 100644
--- a/tensor2tensor/rl/trainer_model_based_sv2p_test.py
+++ b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
@@ -29,13 +29,6 @@ class ModelRLExperimentSv2pTest(tf.test.TestCase):
   def test_sv2p(self):
     FLAGS.output_dir = tf.test.get_temp_dir()
     FLAGS.loop_hparams_set = "rl_modelrl_tiny_sv2p"
-    FLAGS.schedule = "train"
-    trainer_model_based.main(None)
-
-  def test_sv2p_external_loss(self):
-    FLAGS.output_dir = tf.test.get_temp_dir() + "_external"
-    FLAGS.loop_hparams_set = "rl_modelrl_tiny_sv2p_external"
-    FLAGS.schedule = "train"
     trainer_model_based.main(None)
 
 

From f97127f6faee4b8ed82d7f096d81ea80902e9f46 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Wed, 19 Sep 2018 16:32:45 -0700
Subject: [PATCH 0876/2720] erase environment wrappers clipping rewards.

PiperOrigin-RevId: 213719554
---
 .../data_generators/gym_problems_specs.py     | 18 +--------------
 tensor2tensor/data_generators/gym_utils.py    | 23 -------------------
 2 files changed, 1 insertion(+), 40 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index fbadc1940..e5cb0bda3 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -18,8 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import gym
-
 # We need gym_utils for the game environments defined there.
 from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
 # pylint: disable=g-multiple-import
@@ -175,7 +173,6 @@ def num_rewards(self):
 
 def create_problems_for_game(
     game_name,
-    clipped_reward=True,
     resize_height_factor=2,
     resize_width_factor=2,
     game_mode="Deterministic-v4"):
@@ -183,8 +180,6 @@ def create_problems_for_game(
 
   Args:
     game_name: str, one of the games in ATARI_GAMES, e.g. "bank_heist".
-    clipped_reward: bool, whether the rewards should be clipped. False is not
-      yet supported.
     resize_height_factor: factor by which to resize the height of frames.
     resize_width_factor: factor by which to resize the width of frames.
     game_mode: the frame skip and sticky keys config.
@@ -195,9 +190,6 @@ def create_problems_for_game(
   Raises:
     ValueError: if clipped_reward=False or game_name not in ATARI_GAMES.
   """
-  if not clipped_reward:
-    raise ValueError("Creating problems without clipped reward is not "
-                     "yet supported.")
   if game_name not in ATARI_GAMES:
     raise ValueError("Game %s not in ATARI_GAMES" % game_name)
   if game_mode not in ATARI_GAME_MODES:
@@ -206,18 +198,11 @@ def create_problems_for_game(
       [w[0].upper() + w[1:] for w in game_name.split("_")])
   camel_game_name += game_mode
   env_name = camel_game_name
-  wrapped_env_name = "T2T%s" % env_name
-
-  # Register an environment that does the reward clipping
-  gym.envs.register(
-      id=wrapped_env_name,
-      entry_point=lambda: gym_utils.wrapped_factory(  # pylint: disable=g-long-lambda
-          env=env_name, reward_clipping=True))
 
   # Create and register the Random and WithAgent Problem classes
   problem_cls = type("Gym%sRandom" % camel_game_name,
                      (GymClippedRewardRandom,),
-                     {"env_name": wrapped_env_name,
+                     {"env_name": env_name,
                       "resize_height_factor": resize_height_factor,
                       "resize_width_factor": resize_width_factor})
   registry.register_problem(problem_cls)
@@ -248,6 +233,5 @@ def create_problems_for_game(
   for mode in ATARI_GAME_MODES:
     classes = create_problems_for_game(
         game,
-        clipped_reward=True,
         game_mode=mode)
     ATARI_PROBLEMS[game][mode] = classes
diff --git a/tensor2tensor/data_generators/gym_utils.py b/tensor2tensor/data_generators/gym_utils.py
index af14451ae..e821bc38b 100644
--- a/tensor2tensor/data_generators/gym_utils.py
+++ b/tensor2tensor/data_generators/gym_utils.py
@@ -299,26 +299,3 @@ def wrapped_freeway_factory(warm_up_examples=0,
                       easy_freeway=False
                   ),
                   max_episode_steps=500)
-
-
-class DefaultGymWrapper(gym.Wrapper):
-  """Warmup wrapper."""
-
-  def __init__(self, env, reward_clipping=True):
-    super(DefaultGymWrapper, self).__init__(env)
-    self.reward_clipping = reward_clipping
-
-  def step(self, action):
-    ob, rew, done, info = self.env.step(action)
-
-    if self.reward_clipping:
-      rew = np.sign(rew)
-
-    return ob, rew, done, info
-
-
-def wrapped_factory(env, reward_clipping):
-  """Wrapped games."""
-  env = gym.make(env)
-  env = DefaultGymWrapper(env, reward_clipping)
-  return env

From d08dc51d305beff2a0c707805f2ad4670c194729 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 19 Sep 2018 16:38:43 -0700
Subject: [PATCH 0877/2720] removing the resize frame for Atari params.

PiperOrigin-RevId: 213720585
---
 tensor2tensor/models/video/sv2p_params.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index ba137ed3f..78828363b 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -64,8 +64,8 @@ def next_frame_sv2p_atari():
   hparams.anneal_end = 50000
   hparams.latent_loss_multiplier_schedule = "noisy_linear_cosine_decay"
   hparams.latent_loss_multiplier = 1e-3
-  hparams.preprocess_resize_frames = [96, 96]
   hparams.information_capacity = 0.0
+  hparams.small_mode = True
   return hparams
 
 

From 83306ef45c110d28a34ccd74a950cdc563f90ddb Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Wed, 19 Sep 2018 16:53:51 -0700
Subject: [PATCH 0878/2720] Fixing multiproblem loss.

PiperOrigin-RevId: 213722937
---
 .../data_generators/multi_problem.py          | 60 +++++++++++--------
 1 file changed, 35 insertions(+), 25 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 97d046d64..5f28d23fd 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -331,44 +331,54 @@ def aggregate_task_losses(hparams,
   loss_val = loss_num / tf.maximum(1.0, loss_den)
   summaries.append([hparams.problem.task_list[0].name+"_loss", loss_val])
 
+  # Since the losses may undergo rescaling, they cannot exist as separate
+  # numerators and denominators. Set the denominators to 1 in order to faciliate
+  # loss averaging.
+  loss_num = loss_val
+  loss_den = tf.minimum(tf.convert_to_tensor(1, dtype=tf.float32), loss_den)
+
   for task in hparams.problem.task_list[1:]:
     if hasattr(task, "num_classes"):
-      task_loss_num_seq, task_loss_den_seq = target_modality.loss(
+      # Loss only from the input sequence -- the auxiliary LM loss.
+      seq_loss_num, seq_loss_den = target_modality.loss(
           logits, feature,
           weights_fn=
           lambda x: common_layers.weights_multi_problem_input(x, task.task_id))  # pylint: disable=cell-var-from-loop
-      task_loss_num_seq *= problem_hparams.loss_multiplier
+      seq_loss_num *= problem_hparams.loss_multiplier
 
-      task_loss_num_label, task_loss_den_label = target_modality.loss(
+      # Loss only from the classification label.
+      label_loss_num, label_loss_den = target_modality.loss(
           logits, feature,
           weights_fn=
           lambda x: common_layers.weights_multi_problem(x, task.task_id))  # pylint: disable=cell-var-from-loop
-      task_loss_num_label *= problem_hparams.loss_multiplier
+      label_loss_num *= problem_hparams.loss_multiplier
+
+      # Unscaled losses.
+      seq_loss = seq_loss_num / tf.maximum(1.0, seq_loss_den)
+      summaries.append([task.name+"_seq_loss", seq_loss])
+      label_loss = label_loss_num / tf.maximum(1.0, label_loss_den)
+      summaries.append([task.name+"_label_loss", label_loss])
 
+      # Scaling.
       if hparams.multiproblem_reweight_label_loss:
-        task_loss_num = (1 - hparams.multiproblem_label_weight) * \
-                        task_loss_num_seq
-        task_loss_num += hparams.multiproblem_label_weight * task_loss_num_label
-      elif hparams.multiproblem_class_loss_multiplier > 0:
-        task_loss_num = task_loss_num_seq
-        task_loss_num += hparams.multiproblem_class_loss_multiplier * \
-                         task_loss_num_label
-      else:
-        task_loss_num = task_loss_num_seq + task_loss_num_label
-
-      task_loss_den = task_loss_den_seq + task_loss_den_label
-
-      # Log the unscaled versions of the losses to tensorboard.
-      task_loss_val = (task_loss_num_seq + task_loss_num_label) / tf.maximum(
-          1.0, task_loss_den)
-      summaries.append([task.name+"_loss", task_loss_val])
+        label_loss *= hparams.multiproblem_label_weight
+        seq_loss *= (1 - hparams.multiproblem_label_weight)
 
-      task_loss_val_label = task_loss_num_label / tf.maximum(
-          1.0, task_loss_den_label)
-      summaries.append([task.name+"_only_label_loss", task_loss_val_label])
+      if hparams.multiproblem_class_loss_multiplier:
+        label_loss *= hparams.multiproblem_class_loss_multiplier
+        summaries.append([task.name+"_scaled_label_loss", label_loss])
+
+      # This is the training loss for the optimizer after all the scaling.
+      task_loss_val = seq_loss + label_loss
+      summaries.append([task.name+"_loss", task_loss_val])
 
-      loss_num += task_loss_num
-      loss_den += task_loss_den
+      # Adding 1 to the loss den for each task leads to averaging task losses,
+      # task with bigger loss will dominate.
+      # TODO(urvashik): Fix combination with other task losses - weighted
+      # average based on the number of examples from that task.
+      loss_num += task_loss_val
+      loss_den += tf.minimum(tf.convert_to_tensor(1, dtype=tf.float32),
+                             label_loss_den)
 
     else:
       raise ValueError("Non-classification secondary tasks are not supported.")

From 2661d78eaa8d12c01f0fe2aa9ef7891592e06fdd Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Wed, 19 Sep 2018 17:28:35 -0700
Subject: [PATCH 0879/2720] Use PPO in the initial round.

PiperOrigin-RevId: 213728405
---
 tensor2tensor/rl/trainer_model_based.py | 37 +++++++++++++++++--------
 1 file changed, 26 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 37c1abdbc..4326a9992 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -161,12 +161,16 @@ def _ppo_training_epochs(hparams, epoch, is_final_epoch, real_env_training):
 
   simulated_training_ppo_epochs_num = hparams.ppo_epochs_num
 
-  ppo_training_epochs = ((epoch + 1)*simulated_training_ppo_epochs_num
-                         + epoch * real_training_ppo_epochs_num)
+  if epoch == -1:
+    assert real_env_training, (
+        "Epoch -1 should only be used for PPO collection in real environment.")
+    return real_training_ppo_epochs_num
+  ppo_training_epochs = (epoch + 1) * (simulated_training_ppo_epochs_num
+                                       + real_training_ppo_epochs_num)
+  if is_final_epoch:  # Length of training in the final epoch is doubled.
+    ppo_training_epochs += simulated_training_ppo_epochs_num
   if real_env_training:
     ppo_training_epochs += real_training_ppo_epochs_num
-  if is_final_epoch:
-    ppo_training_epochs += simulated_training_ppo_epochs_num
   return ppo_training_epochs
 
 
@@ -210,7 +214,7 @@ def train_agent(problem_name, agent_model_dir,
       "output_dir": world_model_dir,
       "data_dir": epoch_data_dir,
   }):
-    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir,
+    rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
                          name_scope="ppo_sim")
 
 ppo_data_dumper_counter = 0
@@ -279,8 +283,7 @@ def train_agent_real_env(
       "output_dir": world_model_dir,
       "data_dir": epoch_data_dir,
   }):
-    # epoch = 10**20 is a hackish way to avoid skiping training
-    rl_trainer_lib.train(ppo_hparams, event_dir, agent_model_dir,
+    rl_trainer_lib.train(ppo_hparams, event_dir + "real", agent_model_dir,
                          name_scope="ppo_real")
 
 
@@ -488,12 +491,24 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   epoch_metrics = []
   epoch_data_dirs = []
 
-  # Collect data from the real environment with random policy
-  data_dir = os.path.join(directories["data"], "random")
+  ppo_model_dir = None
+  data_dir = os.path.join(directories["data"], "initial")
   epoch_data_dirs.append(data_dir)
-  tf.logging.info("Generating real environment data with random policy")
+  # Collect data from the real environment with PPO or random policy.
+  if hparams.gather_ppo_real_env_data:
+    ppo_model_dir = directories["ppo"]
+    tf.logging.info("Initial training of PPO in real environment.")
+    ppo_event_dir = os.path.join(directories["world_model"],
+                                 "ppo_summaries/initial")
+    train_agent_real_env(
+        problem_name, ppo_model_dir,
+        ppo_event_dir, directories["world_model"], data_dir,
+        hparams, epoch=-1, is_final_epoch=False)
+
+  tf.logging.info("Generating real environment data with %s policy",
+                  "PPO" if hparams.gather_ppo_real_env_data else "random")
   mean_reward = generate_real_env_data(
-      problem_name, None, hparams, data_dir, directories["tmp"])
+      problem_name, ppo_model_dir, hparams, data_dir, directories["tmp"])
   tf.logging.info("Mean reward (random): {}".format(mean_reward))
 
   eval_metrics_event_dir = os.path.join(directories["world_model"],

From c63127d7a1e24e9f0804df4dad13c2bba3147c6e Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 19 Sep 2018 17:54:31 -0700
Subject: [PATCH 0880/2720] fixing simple reward in sv2p.

PiperOrigin-RevId: 213731557
---
 tensor2tensor/models/video/sv2p.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 1216ff639..538083a32 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -229,7 +229,7 @@ def reward_prediction(self, *args, **kwargs):
 
   def reward_prediction_basic(self, input_images, input_reward, action, latent):
     del input_reward, action, latent
-    x = tf.concat(input_images, axis=3)
+    x = input_images[0]
     x = tf.expand_dims(  # Add a fake channels dim.
         tf.reduce_mean(x, axis=[1, 2], keepdims=True), axis=3)
     return x

From debc8935e22796bbf1de12865f52388d7093329a Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Wed, 19 Sep 2018 19:12:00 -0700
Subject: [PATCH 0881/2720] internal merge of PR #1067

PiperOrigin-RevId: 213739549
---
 tensor2tensor/data_generators/gym_problems.py | 206 ++++++++++++------
 .../data_generators/gym_problems_specs.py     |  44 +++-
 tensor2tensor/rl/envs/batch_env_factory.py    |  11 +-
 tensor2tensor/rl/envs/simulated_batch_env.py  |  48 +---
 tensor2tensor/rl/envs/utils.py                | 128 +++++++++++
 tensor2tensor/rl/trainer_model_based.py       |  18 +-
 6 files changed, 340 insertions(+), 115 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index c2f6f359f..9db1eaa6f 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -23,6 +23,7 @@
 import os
 import gym
 import numpy as np
+import six
 
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
@@ -30,13 +31,12 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import collect
 from tensor2tensor.rl.envs import tf_atari_wrappers
+from tensor2tensor.rl.envs.utils import InitialFrameChooser
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
-from tensorflow.contrib.training import HParams
-
 flags = tf.flags
 FLAGS = flags.FLAGS
 
@@ -125,9 +125,8 @@ def resize_height_factor(self):
   def resize_width_factor(self):
     return 2
 
-  def _setup(self, data_dir):
-    # TODO(piotrmilos):this should be consistent with
-    # ppo_params in model_rl_experiment
+  def _setup(self, data_dir, extra_collect_hparams=None,
+             override_collect_hparams=None):
     dumper_path = os.path.join(data_dir, "dumper")
     if os.path.isdir(dumper_path):
       tf.logging.info("Using dumper data.")
@@ -135,6 +134,8 @@ def _setup(self, data_dir):
       self._dumper_data_index = 0
       self._dumper_path = dumper_path
     else:
+      # TODO(piotrmilos):this should be consistent with
+      # ppo_params in model_rl_experiment
       collect_hparams = rl.ppo_pong_base()
       collect_hparams.add_hparam("environment_spec", self.environment_spec)
       collect_hparams.add_hparam("force_beginning_resets",
@@ -145,6 +146,16 @@ def _setup(self, data_dir):
       if not FLAGS.agent_policy_path:
         collect_hparams.policy_network = rl.random_policy_fun
 
+      if extra_collect_hparams is not None:
+        for (key, value) in six.iteritems(extra_collect_hparams):
+          collect_hparams.add_hparam(key, value)
+
+      if override_collect_hparams is not None:
+        # Override hparams manually - HParams.override_from_dict does not work
+        # with functions.
+        for (key, value) in six.iteritems(override_collect_hparams):
+          setattr(collect_hparams, key, value)
+
       policy_to_actions_lambda = None
       if self.settable_eval_phase:
         policy_to_actions_lambda = lambda policy: policy.mode()
@@ -525,8 +536,8 @@ def __init__(self):
     self.report_reward_statistics_every = 10
 
     # auxiliary objects
-    self.real_env = None
-    self.real_ob = None
+    self.real_obs = None
+    self.real_rewards = None
 
   def to_dict(self):
     stats_dict = super(RewardPerSequenceStatistics, self).to_dict()
@@ -554,38 +565,24 @@ def __init__(self, *args, **kwargs):
     # the amount of skips induced but wrappers
     self._internal_memory_size = self.num_testing_steps
     self._internal_memory_force_beginning_resets = True
-    real_env_spec = standard_atari_env_spec(
-        self.env_name,
-        simulated=False,
-        resize_height_factor=self.resize_height_factor,
-        resize_width_factor=self.resize_width_factor)
-    real_env = real_env_spec.env_lambda()
 
-    self.statistics = RewardPerSequenceStatistics()
-    self.statistics.real_env = real_env
+    self.statistics = BasicStatistics()
+    self._initial_frame_chooser = None
 
-  def _setup(self, data_dir):
-    super(GymSimulatedDiscreteProblem, self)._setup(data_dir)
-
-    environment_spec = self.environment_spec
-    hparams = HParams(
-        video_num_input_frames=environment_spec.video_num_input_frames,
-        video_num_target_frames=environment_spec.video_num_target_frames,
-        environment_spec=environment_spec)
-
-    initial_frames_problem = environment_spec.initial_frames_problem
-    dataset = initial_frames_problem.dataset(
-        tf.estimator.ModeKeys.TRAIN,
-        FLAGS.data_dir,
-        shuffle_files=False,
-        hparams=hparams)
-    dataset = dataset.map(lambda x: x["input_action"]).take(1)
-    input_data_iterator = (dataset.batch(1).make_initializable_iterator())
-    self._session.run(input_data_iterator.initializer)
-
-    res = self._session.run(input_data_iterator.get_next())
-    self._initial_actions = res[0, :, 0][:-1]
-    self._reset_real_env()
+  def _setup(self, data_dir, extra_collect_hparams=None,
+             override_collect_hparams=None):
+    if extra_collect_hparams is None:
+      extra_collect_hparams = {}
+
+    if self._initial_frame_chooser is None:
+      self._initial_frame_chooser = InitialFrameChooser(
+          self.environment_spec, mode=tf.estimator.ModeKeys.EVAL
+      )
+    extra_collect_hparams["initial_frame_chooser"] = self._initial_frame_chooser
+
+    super(GymSimulatedDiscreteProblem, self)._setup(
+        data_dir, extra_collect_hparams, override_collect_hparams
+    )
 
   @property
   def initial_frames_problem(self):
@@ -615,8 +612,8 @@ def get_environment_spec(self):
         simulated=True,
         resize_height_factor=self.resize_height_factor,
         resize_width_factor=self.resize_width_factor)
-    env_spec.add_hparam("simulation_random_starts", False)
-    env_spec.add_hparam("simulation_flip_first_random_for_beginning", False)
+    env_spec.add_hparam("simulation_random_starts", True)
+    env_spec.add_hparam("simulation_flip_first_random_for_beginning", True)
     env_spec.add_hparam("intrinsic_reward_scale", 0.0)
     initial_frames_problem = registry.problem(self.initial_frames_problem)
     env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
@@ -625,11 +622,70 @@ def get_environment_spec(self):
 
     return env_spec
 
-  def _reset_real_env(self):
-    stat = self.statistics
-    stat.real_env.reset()
-    for a in self._initial_actions:
-      stat.real_ob, _, _, _ = stat.real_env.step(a)
+  def restore_networks(self, sess):
+    super(GymSimulatedDiscreteProblem, self).restore_networks(sess)
+    # TODO(blazej): adjust regexp for different models.
+    # TODO(piotrmilos): move restoring networks to SimulatedBatchEnv.initialize
+    env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
+    ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
+    ckpt = ckpts.model_checkpoint_path
+    env_model_loader.restore(sess, ckpt)
+
+
+class GymSimulatedDiscreteProblemForWorldModelEval(GymSimulatedDiscreteProblem):
+  """Simulated gym environment for evaluating world model."""
+
+  def __init__(self, *args, **kwargs):
+    super(GymSimulatedDiscreteProblemForWorldModelEval, self).__init__(
+        *args, **kwargs
+    )
+    self.statistics = RewardPerSequenceStatistics()
+
+  def get_environment_spec(self):
+    env_spec = super(
+        GymSimulatedDiscreteProblemForWorldModelEval, self
+    ).get_environment_spec()
+    env_spec.simulation_flip_first_random_for_beginning = False
+    return env_spec
+
+  def _setup(self, data_dir):
+    trajectory_length = self.num_testing_steps
+    if self.num_steps < 1200:
+      # Decrease the trajectory length for tiny experiments, otherwise we don't
+      # have enough data to run the evaluation.
+      trajectory_length = 2
+    self._initial_frame_chooser = InitialFrameChooser(
+        self.environment_spec, mode=tf.estimator.ModeKeys.EVAL,
+        trajectory_length=trajectory_length
+    )
+
+    frame_index = tf.Variable(0, trainable=False)
+
+    def fixed_action_policy_fun(action_space, unused_config, observations):
+      """Policy which replays actions from a trajectory."""
+      action = self._initial_frame_chooser.trajectory["action"].read_value()[
+          :, frame_index.read_value(), :
+      ]
+      inc_frame_index = frame_index.assign(
+          (frame_index.read_value() + 1) % trajectory_length
+      )
+      with tf.control_dependencies([inc_frame_index]):
+        action = tf.identity(action)
+
+      obs_shape = observations.shape.as_list()
+      with tf.variable_scope("network_parameters"):
+        probs = tf.one_hot(
+            tf.transpose(action), depth=action_space.n
+        )
+        policy = tf.distributions.Categorical(probs=probs)
+        value = tf.zeros(obs_shape[:2])
+      return rl.NetworkOutput(policy, value, lambda a: a)
+
+    super(GymSimulatedDiscreteProblemForWorldModelEval, self)._setup(
+        data_dir, override_collect_hparams={
+            "policy_network": fixed_action_policy_fun
+        }
+    )
 
   def collect_statistics_and_generate_debug_image(self, index,
                                                   observation,
@@ -641,15 +697,20 @@ def collect_statistics_and_generate_debug_image(self, index,
     stat.sum_of_rewards += reward
     stat.episode_sim_reward += reward
 
-    ob = np.ndarray.astype(observation, np.int)
-    if ob.shape == stat.real_ob.shape:
-      err = np.ndarray.astype(
-          np.maximum(np.abs(stat.real_ob - ob, dtype=np.int) - 10, 0), np.uint8)
-      debug_im = np.concatenate([observation, stat.real_ob, err], axis=1)
-    else:
-      # Real env does not get the ResizeWrapper and we don't have it in python,
-      # so we skip the debug image here and just output observations.
-      debug_im = observation
+    if index % self._internal_memory_size == 0:
+      real_frame_tensor = {
+          key: var.read_value()[0, ...]
+          for (key, var) in six.iteritems(
+              self._initial_frame_chooser.trajectory
+          )
+      }
+      (stat.real_obs, stat.real_rewards) = self._session.run((
+          real_frame_tensor["inputs"], real_frame_tensor["reward"]
+      ))
+      stat.real_rewards += self.min_reward
+
+    real_ob = stat.real_obs[index % stat.real_obs.shape[0], ...]
+    debug_im = self._generate_debug_image(real_ob, observation)
 
     assert (self._internal_memory_size == self.num_testing_steps and
             self._internal_memory_force_beginning_resets), (
@@ -660,25 +721,27 @@ def collect_statistics_and_generate_debug_image(self, index,
 
       if stat.episode_sim_reward == stat.episode_real_reward:
         stat.successful_episode_reward_predictions += 1
-        stat.episode_sim_reward = 0.0
-        stat.episode_real_reward = 0.0
 
+      stat.episode_sim_reward = 0.0
+      stat.episode_real_reward = 0.0
       stat.number_of_dones += 1
-      self._reset_real_env()
     else:
-      stat.real_ob, real_reward, _, _ = stat.real_env.step(action)
+      real_reward = stat.real_rewards[index % stat.real_rewards.shape[0], 0]
       stat.episode_real_reward += real_reward
 
     return debug_im
 
-  def restore_networks(self, sess):
-    super(GymSimulatedDiscreteProblem, self).restore_networks(sess)
-    # TODO(blazej): adjust regexp for different models.
-    # TODO(piotrmilos): move restoring networks to SimulatedBatchEnv.initialize
-    env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
-    ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
-    ckpt = ckpts.model_checkpoint_path
-    env_model_loader.restore(sess, ckpt)
+  def _generate_debug_image(self, real_ob, sim_ob):
+    ob = np.ndarray.astype(sim_ob, np.int)
+    if ob.shape == real_ob.shape:
+      err = np.ndarray.astype(
+          np.maximum(np.abs(real_ob - ob, dtype=np.int) - 10, 0), np.uint8)
+      debug_im = np.concatenate([sim_ob, real_ob, err], axis=1)
+    else:
+      # Real env does not get the ResizeWrapper and we don't have it in python,
+      # so we skip the debug image here and just output observations.
+      debug_im = sim_ob
+    return debug_im
 
 
 class GymSimulatedDiscreteProblemAutoencoded(GymSimulatedDiscreteProblem):
@@ -691,8 +754,8 @@ def get_environment_spec(self):
         [tf_atari_wrappers.StackWrapper, {"history": 4}]
     ]
     env_spec.simulated_env = True
-    env_spec.add_hparam("simulation_random_starts", False)
-
+    env_spec.add_hparam("simulation_random_starts", True)
+    env_spec.add_hparam("simulation_flip_first_random_for_beginning", True)
     env_spec.add_hparam("intrinsic_reward_scale", 0.0)
     initial_frames_problem = registry.problem(self.initial_frames_problem)
     env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
@@ -719,6 +782,15 @@ def frame_width(self):
     return int(math.ceil(width / self.autoencoder_factor))
 
 
+class GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded(
+    GymSimulatedDiscreteProblemForWorldModelEval,
+    GymSimulatedDiscreteProblemAutoencoded):
+
+  def _generate_debug_image(self, real_ob, sim_ob):
+    # TODO(koz4k): Implement.
+    pass
+
+
 @registry.register_problem
 class DummyAutoencoderProblem(GymDiscreteProblemWithAutoencoder):
   """Dummy problem for running the autoencoder inside AutoencoderWrapper."""
diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index e5cb0bda3..136e8ec3f 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -24,7 +24,9 @@
 from tensor2tensor.data_generators.gym_problems import GymDiscreteProblem,\
   GymSimulatedDiscreteProblem, GymRealDiscreteProblem, \
   GymDiscreteProblemWithAutoencoder, GymDiscreteProblemAutoencoded, \
-  GymSimulatedDiscreteProblemAutoencoded
+  GymSimulatedDiscreteProblemAutoencoded, \
+  GymSimulatedDiscreteProblemForWorldModelEval, \
+  GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded
 # pylint: enable=g-multiple-import
 from tensor2tensor.utils import registry
 
@@ -141,6 +143,20 @@ def num_testing_steps(self):
     return 100
 
 
+@registry.register_problem
+class GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOnWrappedFullPong(
+    GymSimulatedDiscreteProblemForWorldModelEval, GymWrappedFullPongRandom):
+  """Simulated pong for world model evaluation."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_full_pong"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
 @registry.register_problem
 class GymSimulatedDiscreteProblemWithAgentOnWrappedFullPongAutoencoded(
     GymSimulatedDiscreteProblemAutoencoded, GymWrappedFullPongRandom):
@@ -155,6 +171,21 @@ def num_testing_steps(self):
     return 100
 
 
+@registry.register_problem
+class GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOnWrappedFullPongAutoencoded(  # pylint: disable=line-too-long
+    GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded,
+    GymWrappedFullPongRandom):
+  """Simulated pong for world model evaluation with encoded frames."""
+
+  @property
+  def initial_frames_problem(self):
+    return "gym_discrete_problem_with_agent_on_wrapped_full_pong_autoencoded"
+
+  @property
+  def num_testing_steps(self):
+    return 100
+
+
 class GymClippedRewardRandom(GymDiscreteProblem):
   """Abstract base class for clipped reward games."""
 
@@ -221,10 +252,21 @@ def create_problems_for_game(
       })
   registry.register_problem(simulated_cls)
 
+  # Create and register the simulated Problem
+  world_model_eval_cls = type(
+      "GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%s" %
+      camel_game_name,
+      (GymSimulatedDiscreteProblemForWorldModelEval, problem_cls), {
+          "initial_frames_problem": with_agent_cls.name,
+          "num_testing_steps": 100
+      })
+  registry.register_problem(world_model_eval_cls)
+
   return {
       "base": problem_cls,
       "agent": with_agent_cls,
       "simulated": simulated_cls,
+      "world_model_eval": world_model_eval_cls,
   }
 
 # Register the atari games with all of the possible modes.
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index af3d6b549..61b6cec8c 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -46,7 +46,8 @@ def batch_env_factory(hparams, xvfb=False):
   if environment_spec.simulated_env:
     # TODO(piotrmilos): Consider passing only relevant parameters
     cur_batch_env = _define_simulated_batch_env(
-        environment_spec, hparams.num_agents)
+        environment_spec, hparams.num_agents,
+        hparams.initial_frame_chooser)
   else:
     cur_batch_env = _define_batch_env(hparams.environment_spec,
                                       hparams.num_agents,
@@ -66,9 +67,11 @@ def _define_batch_env(environment_spec, num_agents, xvfb=False):
     return env
 
 
-def _define_simulated_batch_env(environment_spec, num_agents):
-  cur_batch_env = simulated_batch_env.SimulatedBatchEnv(environment_spec,
-                                                        num_agents)
+def _define_simulated_batch_env(environment_spec, num_agents,
+                                initial_frame_chooser):
+  cur_batch_env = simulated_batch_env.SimulatedBatchEnv(
+      environment_spec, num_agents, initial_frame_chooser
+  )
   return cur_batch_env
 
 
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 7d7e4def4..b478bd2bb 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -29,8 +29,6 @@
 
 import tensorflow as tf
 
-from tensorflow.contrib.training import HParams
-
 
 flags = tf.flags
 FLAGS = flags.FLAGS
@@ -39,13 +37,9 @@
 class HistoryBuffer(object):
   """History Buffer."""
 
-  def __init__(self, input_dataset, length, observ_dtype, start_frame=None):
-    if start_frame is None:
-      dataset = input_dataset.batch(length)
-    else:
-      dataset = input_dataset.batch(length - 1)
-      dataset = dataset.map(lambda x: tf.concat([start_frame, x], axis=0))
-    self.input_data_iterator = dataset.make_initializable_iterator()
+  def __init__(self, initial_frame_chooser, length, observ_dtype):
+    initial_frame_chooser.batch_size = length
+    self._initial_frame_chooser = initial_frame_chooser
     self.length = length
     self._observ_dtype = observ_dtype
     initial_frames = self.get_initial_observations()
@@ -54,10 +48,12 @@ def __init__(self, input_dataset, length, observ_dtype, start_frame=None):
                                      trainable=False)
 
   def initialize(self, sess):
-    sess.run(self.input_data_iterator.initializer)
+    self._initial_frame_chooser.initialize(sess)
 
   def get_initial_observations(self):
-    return tf.cast(self.input_data_iterator.get_next(), self._observ_dtype)
+    return tf.cast(
+        self._initial_frame_chooser.choose()["inputs"], self._observ_dtype
+    )
 
   def get_all_elements(self):
     return self._history_buff.read_value()
@@ -102,7 +98,7 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   flags are held in according variables.
   """
 
-  def __init__(self, environment_spec, length):
+  def __init__(self, environment_spec, length, initial_frame_chooser):
     """Batch of environments inside the TensorFlow graph."""
 
     observ_space = utils.get_observation_space(environment_spec)
@@ -126,34 +122,8 @@ def __init__(self, environment_spec, length):
     self._model = registry.model(FLAGS.model)(
         model_hparams, tf.estimator.ModeKeys.PREDICT)
 
-    hparams = HParams(video_num_input_frames=
-                      environment_spec.video_num_input_frames,
-                      video_num_target_frames=
-                      environment_spec.video_num_target_frames,
-                      environment_spec=environment_spec)
-
-    # TODO(piotrmilos): check if this shouldn't be tf.estimator.ModeKeys.Predict
-    initial_frames_dataset = initial_frames_problem.dataset(
-        tf.estimator.ModeKeys.TRAIN, FLAGS.data_dir, shuffle_files=False,
-        hparams=hparams).take(1)
-    start_frame = None
-    if environment_spec.simulation_random_starts:
-      dataset = initial_frames_problem.dataset(tf.estimator.ModeKeys.TRAIN,
-                                               FLAGS.data_dir,
-                                               shuffle_files=True,
-                                               hparams=hparams,
-                                               only_last=True)
-      dataset = dataset.shuffle(buffer_size=1000)
-      if environment_spec.simulation_flip_first_random_for_beginning:
-        # Later flip the first random frame in PPO batch for the true beginning.
-        start = initial_frames_dataset.make_one_shot_iterator().get_next()
-        start_frame = tf.expand_dims(start["inputs"], axis=0)
-    else:
-      dataset = initial_frames_dataset
-
-    dataset = dataset.map(lambda x: x["inputs"]).repeat()
     self.history_buffer = HistoryBuffer(
-        dataset, self.length, self.observ_dtype, start_frame=start_frame)
+        initial_frame_chooser, self.length, self.observ_dtype)
 
     self._observ = tf.Variable(
         tf.zeros((len(self),) + observ_shape, self.observ_dtype),
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index badd48485..9d8850a2c 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -19,6 +19,7 @@
 from __future__ import print_function
 
 import gym
+import six
 import tensorflow as tf
 
 
@@ -134,3 +135,130 @@ def parse_dtype(space):
   if isinstance(space, gym.spaces.Box):
     return tf.as_dtype(space.dtype)
   raise NotImplementedError()
+
+
+class InitialFrameChooser(object):
+  """Class for choosing the initial frame for simulation from the dataset.
+
+  Can also store a sequence of later frames, which is used for comparison in
+  world model evaluation.
+
+  Attributes:
+    batch_size (int): Batch size, should be set before calling choose().
+    trajectory (dict): Dict of Variables storing a sequence of frames after the
+        chosen one.
+  """
+
+  def __init__(self, environment_spec, mode, trajectory_length=1):
+    self._initial_frames_problem = environment_spec.initial_frames_problem
+    self._simulation_random_starts = environment_spec.simulation_random_starts
+    self._flip_first_random_for_beginning = \
+        environment_spec.simulation_flip_first_random_for_beginning
+    self._num_initial_frames = environment_spec.video_num_input_frames
+
+    def dataset_kwargs_lambda():
+      video_num_input_frames = environment_spec.video_num_input_frames
+      video_num_input_frames += trajectory_length - 1
+      dataset_hparams = tf.contrib.training.HParams(
+          video_num_input_frames=video_num_input_frames,
+          video_num_target_frames=environment_spec.video_num_target_frames,
+          environment_spec=environment_spec
+      )
+      return {
+          "mode": mode,
+          "data_dir": tf.flags.FLAGS.data_dir,
+          "hparams": dataset_hparams,
+          "only_last": True
+      }
+
+    self._dataset_kwargs_lambda = dataset_kwargs_lambda
+    self._start_frames = None
+
+  @property
+  def batch_size(self):
+    return self._batch_size
+
+  @batch_size.setter
+  def batch_size(self, batch_size):
+    self._batch_size = batch_size
+    self._iterator = \
+        self._create_initial_frame_dataset().make_initializable_iterator()
+
+    def fix_and_shorten(shape):
+      shape = shape.as_list()
+      shape[0] = batch_size
+      shape[1] -= self._num_initial_frames - 1
+      return shape
+
+    shapes = self._extract_input(self._iterator.output_shapes)
+    types = self._extract_input(self._iterator.output_types)
+    self.trajectory = {
+        key: tf.Variable(
+            tf.zeros(fix_and_shorten(shape), types[key]),
+            trainable=False
+        )
+        for (key, shape) in six.iteritems(shapes)
+    }
+
+  def initialize(self, sess):
+    sess.run(self._iterator.initializer)
+
+  def choose(self):
+    """Returns a dict of tensors of the chosen initial frame.
+
+    Also assigns the first trajectory_length frames after the initial frames to
+    self.trajectory.
+    """
+    if self._flip_first_random_for_beginning and self._start_frames is None:
+      ordered_dataset = self._create_dataset(shuffle_files=False)
+      # Later flip the first random frame in PPO batch for the true beginning.
+      self._start_frames = self._extract_input(
+          ordered_dataset.make_one_shot_iterator().get_next()
+      )
+
+    all_frames = self._extract_input(self._iterator.get_next())
+    if self._start_frames is not None:
+      all_frames = {
+          key: tf.concat([
+              tf.expand_dims(self._start_frames[key], axis=0),
+              value[1:, ...]
+          ], axis=0)
+          for (key, value) in six.iteritems(all_frames)
+      }
+    scatter_ops = [
+        tf.scatter_update(
+            self.trajectory[key], tf.range(tf.shape(value)[0]),
+            value[:, (self._num_initial_frames - 1):, ...]
+        )
+        for (key, value) in six.iteritems(all_frames)
+    ]
+
+    with tf.control_dependencies(scatter_ops):
+      return {
+          key: value[:, :self._num_initial_frames, ...]
+          for (key, value) in six.iteritems(all_frames)
+      }
+
+  def _create_dataset(self, **extra_dataset_kwargs):
+    dataset_kwargs = self._dataset_kwargs_lambda()
+    dataset_kwargs.update(extra_dataset_kwargs)
+    return self._initial_frames_problem.dataset(**dataset_kwargs)
+
+  def _create_initial_frame_dataset(self):
+    """Returns the dataset that consecutive initial frames will be taken from.
+    """
+    dataset = self._create_dataset(
+        shuffle_files=self._simulation_random_starts
+    )
+    if self._simulation_random_starts:
+      dataset = dataset.shuffle(buffer_size=1000)
+    return dataset.repeat().batch(self._batch_size)
+
+  def _extract_input(self, frame):
+    input_frame = {"inputs": frame["inputs"]}
+    input_frame.update({
+        key[len("input_"):]: value
+        for (key, value) in six.iteritems(frame)
+        if key.startswith("input_")
+    })
+    return input_frame
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 4326a9992..0a33c5127 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -43,6 +43,7 @@
 from tensor2tensor.layers import discretization
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.rl.envs.tf_atari_wrappers import PyFuncWrapper
+from tensor2tensor.rl.envs.utils import InitialFrameChooser
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
@@ -206,6 +207,9 @@ def train_agent(problem_name, agent_model_dir,
   environment_spec.intrinsic_reward_scale = hparams.intrinsic_reward_scale
 
   ppo_hparams.add_hparam("environment_spec", environment_spec)
+  ppo_hparams.add_hparam("initial_frame_chooser", InitialFrameChooser(
+      environment_spec, mode=tf.estimator.ModeKeys.TRAIN
+  ))
 
   with temporary_flags({
       "problem": problem_name,
@@ -467,11 +471,18 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     simulated_problem_name = (
         "gym_simulated_discrete_problem_with_agent_on_%s_autoencoded"
         % game_with_mode)
+    world_model_eval_problem_name = (
+        "gym_simulated_discrete_problem_for_world_model_eval_with_agent_on_%s"
+        "_autoencoded"
+        % game_with_mode)
   else:
     problem_name = ("gym_discrete_problem_with_agent_on_%s" % game_with_mode)
     world_model_problem = problem_name
     simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
                               % game_with_mode)
+    world_model_eval_problem_name = (
+        "gym_simulated_discrete_problem_for_world_model_eval_with_agent_on_%s"
+        % game_with_mode)
     if problem_name not in registry.list_problems():
       tf.logging.info("Game Problem %s not found; dynamically registering",
                       problem_name)
@@ -555,7 +566,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     if hparams.eval_world_model:
       log("Evaluating world model")
       model_reward_accuracy = evaluate_world_model(
-          simulated_problem_name, world_model_problem, hparams,
+          world_model_eval_problem_name, world_model_problem, hparams,
           directories["world_model"],
           epoch_data_dir, directories["tmp"])
       log("World model reward accuracy: %.4f", model_reward_accuracy)
@@ -895,7 +906,6 @@ def rl_modelrl_ae_tiny():
   hparams.resize_height_factor = 1
   hparams.resize_width_factor = 1
   hparams.autoencoder_train_steps = 2
-  hparams.eval_world_model = False
   return hparams
 
 
@@ -933,8 +943,8 @@ def rl_modelrl_variance(rhp):
 
 @registry.register_ranged_hparams
 def rl_modelrl_variance_nogame(rhp):
-  # Dummy parameter to get 5 runs for each configuration
-  rhp.set_discrete("model.moe_loss_coef", list(range(500)))
+  # Dummy parameter to get 10 runs for current configuration.
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
 
 
 @registry.register_ranged_hparams

From a27fdffb417c3d4e1979385f581f84319427242f Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Wed, 19 Sep 2018 20:27:26 -0700
Subject: [PATCH 0882/2720] updating whitelisted Atari games.

PiperOrigin-RevId: 213744699
---
 tensor2tensor/data_generators/gym_problems_specs.py | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index 136e8ec3f..02ce0626c 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -50,24 +50,19 @@
     "zaxxon"
 ]
 
-# Subset of games with promissing results on model based training.
 ATARI_WHITELIST_GAMES = [
     "amidar",
     "bank_heist",
     "berzerk",
     "boxing",
-    "breakout",
     "crazy_climber",
     "freeway",
     "frostbite",
     "gopher",
-    "hero",
     "kung_fu_master",
+    "ms_pacman",
     "pong",
-    "road_runner",
     "seaquest",
-    # TODO(blazej): check if we get equally good results on vanilla pong.
-    "wrapped_full_pong",
 ]
 
 ATARI_ALL_MODES_SHORT_LIST = [

From 9d03273967e19684bf2406256542737ec8935064 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 19 Sep 2018 20:27:31 -0700
Subject: [PATCH 0883/2720] Set default RL epochs to 6.

PiperOrigin-RevId: 213744704
---
 tensor2tensor/rl/trainer_model_based.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 0a33c5127..2a1bf5065 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -652,7 +652,7 @@ def combine_training_data(problem, final_data_dir, old_data_dirs,
 @registry.register_hparams
 def rl_modelrl_base():
   return tf.contrib.training.HParams(
-      epochs=3,
+      epochs=6,
       # Total frames used for training. This will be distributed evenly across
       # hparams.epochs.
       # This number should be divisible by real_ppo_epoch_length*epochs

From 9458bd22eba2f1b045a5cb04a5ec51c0385d694f Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 19 Sep 2018 20:41:45 -0700
Subject: [PATCH 0884/2720] adding stochastic hparams for Atari

PiperOrigin-RevId: 213745889
---
 tensor2tensor/models/video/sv2p_params.py | 14 +++++++++++---
 tensor2tensor/rl/trainer_model_based.py   |  9 +++++++++
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 78828363b..fdb442524 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -59,9 +59,9 @@ def next_frame_sv2p_atari():
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 4
   hparams.action_injection = "multiplicative"
-  hparams.num_iterations_1st_stage = 15000
-  hparams.num_iterations_2nd_stage = 15000
-  hparams.anneal_end = 50000
+  hparams.num_iterations_1st_stage = 20000
+  hparams.num_iterations_2nd_stage = 20000
+  hparams.anneal_end = 80000
   hparams.latent_loss_multiplier_schedule = "noisy_linear_cosine_decay"
   hparams.latent_loss_multiplier = 1e-3
   hparams.information_capacity = 0.0
@@ -69,6 +69,14 @@ def next_frame_sv2p_atari():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_sv2p_atari_deterministic():
+  """Deterministic for atari."""
+  hparams = next_frame_sv2p_atari()
+  hparams.stochastic_model = False
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_sv2p_tiny():
   """Tiny SV2P model."""
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 2a1bf5065..daa1c4a14 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -768,6 +768,15 @@ def rl_modelrl_base_sv2p():
   return hparams
 
 
+@registry.register_hparams
+def rl_modelrl_base_sv2p_deterministic():
+  """Base setting with deterministic sv2p as world model."""
+  hparams = rl_modelrl_base_sv2p()
+  hparams.generative_model = "next_frame_sv2p"
+  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
+  return hparams
+
+
 @registry.register_hparams
 def rl_modelrl_base_sampling():
   """Base setting with a stochastic next-frame model."""

From a94ce6a766ad7f80500b9f2378ba5f4055a1c8c8 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Wed, 19 Sep 2018 20:43:58 -0700
Subject: [PATCH 0885/2720] adding qbert to whitelisted games.

PiperOrigin-RevId: 213745996
---
 tensor2tensor/data_generators/gym_problems_specs.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index 02ce0626c..239d074ef 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -62,6 +62,7 @@
     "kung_fu_master",
     "ms_pacman",
     "pong",
+    "qbert",
     "seaquest",
 ]
 

From ae44ce492c9929a549abc26ae124b17b30ddbfbb Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 20 Sep 2018 00:32:01 -0700
Subject: [PATCH 0886/2720] Back to old style PPO network, lower learning rate
 and increase warmup for world models, better scopes for TB.

PiperOrigin-RevId: 213762596
---
 tensor2tensor/models/research/rl.py           | 29 +++++++++++++++++++
 .../video/basic_deterministic_params.py       |  4 +--
 tensor2tensor/rl/trainer_model_based.py       |  4 +--
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 379cd6c68..5c9d3fafb 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -231,7 +231,36 @@ def feed_forward_cnn_small_categorical_fun(action_space, config, observations):
   """Small cnn network with categorical output."""
   obs_shape = common_layers.shape_list(observations)
   x = tf.reshape(observations, [-1] + obs_shape[2:])
+  with tf.variable_scope("network_parameters"):
+    dropout = getattr(config, "dropout_ppo", 0.0)
+    with tf.variable_scope("feed_forward_cnn_small"):
+      x = tf.to_float(x) / 255.0
+      x = tf.contrib.layers.conv2d(x, 32, [5, 5], [2, 2],
+                                   activation_fn=tf.nn.relu, padding="SAME")
+      x = tf.contrib.layers.conv2d(x, 32, [5, 5], [2, 2],
+                                   activation_fn=tf.nn.relu, padding="SAME")
+
+      flat_x = tf.reshape(
+          x, [obs_shape[0], obs_shape[1],
+              functools.reduce(operator.mul, x.shape.as_list()[1:], 1)])
+      flat_x = tf.nn.dropout(flat_x, keep_prob=1.0 - dropout)
+      x = tf.contrib.layers.fully_connected(flat_x, 128, tf.nn.relu)
 
+      logits = tf.contrib.layers.fully_connected(x, action_space.n,
+                                                 activation_fn=None)
+      logits = clip_logits(logits, config)
+
+      value = tf.contrib.layers.fully_connected(
+          x, 1, activation_fn=None)[..., 0]
+      policy = tf.contrib.distributions.Categorical(logits=logits)
+  return NetworkOutput(policy, value, lambda a: a)
+
+
+def feed_forward_cnn_small_categorical_fun_new(
+    action_space, config, observations):
+  """Small cnn network with categorical output."""
+  obs_shape = common_layers.shape_list(observations)
+  x = tf.reshape(observations, [-1] + obs_shape[2:])
   with tf.variable_scope("network_parameters"):
     dropout = getattr(config, "dropout_ppo", 0.0)
     with tf.variable_scope("feed_forward_cnn_small"):
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index ea9bc2855..c4724f239 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -31,8 +31,8 @@ def next_frame_basic_deterministic():
   hparams.batch_size = 4
   hparams.num_hidden_layers = 2
   hparams.optimizer = "Adafactor"
-  hparams.learning_rate_constant = 1.5
-  hparams.learning_rate_warmup_steps = 1500
+  hparams.learning_rate_constant = 0.7
+  hparams.learning_rate_warmup_steps = 8000
   hparams.learning_rate_schedule = "linear_warmup * constant * rsqrt_decay"
   hparams.label_smoothing = 0.0
   hparams.initializer = "uniform_unit_scaling"
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index daa1c4a14..1da607be0 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -219,7 +219,7 @@ def train_agent(problem_name, agent_model_dir,
       "data_dir": epoch_data_dir,
   }):
     rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
-                         name_scope="ppo_sim")
+                         name_scope="ppo_sim%d" % (epoch + 1))
 
 ppo_data_dumper_counter = 0
 dumper_path = None
@@ -288,7 +288,7 @@ def train_agent_real_env(
       "data_dir": epoch_data_dir,
   }):
     rl_trainer_lib.train(ppo_hparams, event_dir + "real", agent_model_dir,
-                         name_scope="ppo_real")
+                         name_scope="ppo_real%d" % (epoch + 1))
 
 
 def evaluate_world_model(simulated_problem_name, problem_name, hparams,

From fcb7f2f66d9c49ad93969e57165b94d4a0a21341 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Thu, 20 Sep 2018 08:55:26 -0700
Subject: [PATCH 0887/2720] Fix for modalities in the problem class.

PiperOrigin-RevId: 213822117
---
 tensor2tensor/data_generators/problem.py | 49 +++++++++++++-----------
 1 file changed, 26 insertions(+), 23 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 28fc88d7b..fb74c2675 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -497,7 +497,8 @@ def get_hparams(self, model_hparams=None):
       _copy_problem_hparams(hp)
 
     model_hparams = copy.copy(model_hparams)
-    if (hasattr(model_hparams, "shared_embedding_and_softmax_weights") and
+    if (self.has_inputs and
+        hasattr(model_hparams, "shared_embedding_and_softmax_weights") and
         model_hparams.shared_embedding_and_softmax_weights):
       # If vocabularies differ, unset shared_embedding_and_softmax_weights.
       input_vocab_size = hp.input_modality.get("inputs")[1]
@@ -506,7 +507,7 @@ def get_hparams(self, model_hparams=None):
         tf.logging.warn("Unsetting shared_embedding_and_softmax_weights.")
         model_hparams.shared_embedding_and_softmax_weights = 0
 
-    _create_modalities(hp, model_hparams)
+    _create_modalities(hp, model_hparams, self.has_inputs)
     self._hparams = hp
     return self._hparams
 
@@ -1093,7 +1094,7 @@ def _reverse_problem_hparams(p_hparams):
   p.was_reversed = True
 
 
-def _create_modalities(problem_hparams, hparams):
+def _create_modalities(problem_hparams, hparams, has_inputs):
   """Converts string-type modalities to their corresponding Modality.
 
   Args:
@@ -1105,30 +1106,32 @@ def _create_modalities(problem_hparams, hparams):
     hparams: tf.contrib.training.HParams for the model. It may have
       input_modalities and target_modality, which will override
       problem_hparams' modalities.
+    has_inputs: A boolean that indicates whether to update the input modality.
 
   Returns:
     None
   """
-  input_modality_overrides = {}
-  if hasattr(hparams, "input_modalities"):
-    for override_str in hparams.input_modalities.split(";"):
-      if override_str != "default":
-        parts = override_str.split(":")
-        feature_name = parts[0]
-        modality_name = ":".join(parts[1:])
-        input_modality_overrides[feature_name] = modality_name
-
-  input_modality = {}
-  for feature_name, modality in six.iteritems(problem_hparams.input_modality):
-    if isinstance(modality, (list, tuple)):
-      if feature_name in input_modality_overrides:
-        _warn_changed_modality_type(input_modality_overrides[feature_name],
-                                    modality[0],
-                                    feature_name)
-        modality = (input_modality_overrides[feature_name], modality[1])
-      modality = modalities.create_modality(modality, hparams)
-    input_modality[feature_name] = modality
-  problem_hparams.input_modality = input_modality
+  if has_inputs:
+    input_modality_overrides = {}
+    if hasattr(hparams, "input_modalities"):
+      for override_str in hparams.input_modalities.split(";"):
+        if override_str != "default":
+          parts = override_str.split(":")
+          feature_name = parts[0]
+          modality_name = ":".join(parts[1:])
+          input_modality_overrides[feature_name] = modality_name
+
+    input_modality = {}
+    for feature_name, modality in six.iteritems(problem_hparams.input_modality):
+      if isinstance(modality, (list, tuple)):
+        if feature_name in input_modality_overrides:
+          _warn_changed_modality_type(input_modality_overrides[feature_name],
+                                      modality[0],
+                                      feature_name)
+          modality = (input_modality_overrides[feature_name], modality[1])
+        modality = modalities.create_modality(modality, hparams)
+      input_modality[feature_name] = modality
+    problem_hparams.input_modality = input_modality
 
   target_modality_name = None
   if (hasattr(hparams, "target_modality") and

From 4e671330b4b6eabfcd596c32fdacda88dd1fa45b Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 20 Sep 2018 11:43:56 -0700
Subject: [PATCH 0888/2720] Remove tf.contrib.learn.Experiment

PiperOrigin-RevId: 213853603
---
 tensor2tensor/bin/t2t_trainer.py   | 15 ++++++++-------
 tensor2tensor/utils/trainer_lib.py | 13 -------------
 2 files changed, 8 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index ca78354be..8ff348f2a 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -150,7 +150,7 @@ def create_hparams():
   return trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
 
 
-def create_experiment_fn(**kwargs):
+def create_experiment_fn():
   return trainer_lib.create_experiment_fn(
       model_name=FLAGS.model,
       problem_name=FLAGS.problem,
@@ -167,23 +167,24 @@ def create_experiment_fn(**kwargs):
       eval_early_stopping_steps=FLAGS.eval_early_stopping_steps,
       eval_early_stopping_metric=FLAGS.eval_early_stopping_metric,
       eval_early_stopping_metric_delta=FLAGS.eval_early_stopping_metric_delta,
-      eval_early_stopping_metric_minimize=FLAGS.
-      eval_early_stopping_metric_minimize,
+      eval_early_stopping_metric_minimize=FLAGS
+      .eval_early_stopping_metric_minimize,
       use_tpu=FLAGS.use_tpu,
       use_tpu_estimator=FLAGS.use_tpu_estimator,
       use_xla=FLAGS.xla_compile,
       warm_start_from=FLAGS.warm_start_from,
       decode_from_file=FLAGS.decode_from_file,
       decode_to_file=FLAGS.decode_to_file,
-      decode_reference=FLAGS.decode_reference,
-      **kwargs)
+      decode_reference=FLAGS.decode_reference)
 
 
-def create_run_config(hp):
+def create_run_config(hp, output_dir=None):
   """Create a run config.
 
   Args:
     hp: model hyperparameters
+    output_dir: model's output directory, defaults to output_dir flag.
+
   Returns:
     a run config
   """
@@ -209,7 +210,7 @@ def create_run_config(hp):
       hp.activation_dtype == "float32" and
       hp.weight_dtype == "float32")
   return trainer_lib.create_run_config(
-      model_dir=os.path.expanduser(FLAGS.output_dir),
+      model_dir=output_dir or os.path.expanduser(FLAGS.output_dir),
       master=FLAGS.master,
       iterations_per_loop=FLAGS.iterations_per_loop,
       num_shards=FLAGS.tpu_num_shards,
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 5f9c439a5..cf20300df 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -444,7 +444,6 @@ def create_experiment(
     eval_early_stopping_metric=None,
     eval_early_stopping_metric_delta=None,
     eval_early_stopping_metric_minimize=True,
-    autotune=False,
     use_tpu=False,
     use_tpu_estimator=False,
     use_xla=False,
@@ -560,18 +559,6 @@ def compare_fn(best_eval_result, current_eval_result):
       throttle_secs=eval_throttle_seconds,
       exporters=exporter)
 
-  if autotune:
-    hooks_kwargs = {"train_monitors": train_hooks, "eval_hooks": eval_hooks}
-    return tf.contrib.learn.Experiment(
-        estimator=estimator,
-        train_input_fn=train_input_fn,
-        eval_input_fn=eval_input_fn,
-        train_steps=train_steps,
-        eval_steps=eval_steps,
-        min_eval_frequency=min_eval_frequency,
-        train_steps_per_iteration=min(min_eval_frequency, train_steps),
-        eval_delay_secs=0 if schedule == "evaluate" else 120,
-        **hooks_kwargs if not use_tpu else {})
   return T2TExperiment(estimator, hparams, train_spec, eval_spec,
                        use_validation_monitor, decode_hparams)
 

From f9bde25b05d893a02bfb5e7d49e467cc3da46c61 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 20 Sep 2018 15:15:12 -0700
Subject: [PATCH 0889/2720] rounding to int.

PiperOrigin-RevId: 213891611
---
 tensor2tensor/models/video/sv2p.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 538083a32..58222eaf0 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -526,6 +526,7 @@ def infer(self, features, *args, **kwargs):
       output = {"targets": output}
 
     output["targets"] = tf.squeeze(output["targets"], axis=-1)
+    output["targets"] = tf.to_int64(tf.round(output["targets"]))
     if self.hparams.reward_prediction:
       output["target_reward"] = tf.argmax(output["target_reward"], axis=-1)
 

From 9af731f06a39f7f51fdd7be51538a10b83c59ca5 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 20 Sep 2018 15:33:43 -0700
Subject: [PATCH 0890/2720] back-prop through reward.

PiperOrigin-RevId: 213894761
---
 tensor2tensor/models/video/sv2p_params.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index fdb442524..a1278589c 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -35,7 +35,7 @@ def next_frame_sv2p():
   hparams.input_modalities = "inputs:video:l2raw"
   hparams.video_modality_loss_cutoff = 0.0
   hparams.add_hparam("reward_prediction", True)
-  hparams.add_hparam("reward_prediction_stop_gradient", True)
+  hparams.add_hparam("reward_prediction_stop_gradient", False)
   hparams.add_hparam("reward_prediction_buffer_size", 0)
   hparams.add_hparam("model_options", "CDNA")
   hparams.add_hparam("num_masks", 10)

From ac59e12bdbc8424ccd1e2b702a6f128e3055027e Mon Sep 17 00:00:00 2001
From: Kay Zhu <kayzhu@google.com>
Date: Thu, 20 Sep 2018 18:05:46 -0700
Subject: [PATCH 0891/2720] Fix for modalities in the problem class.

PiperOrigin-RevId: 213917724
---
 tensor2tensor/data_generators/problem.py | 49 +++++++++++-------------
 1 file changed, 23 insertions(+), 26 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index fb74c2675..28fc88d7b 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -497,8 +497,7 @@ def get_hparams(self, model_hparams=None):
       _copy_problem_hparams(hp)
 
     model_hparams = copy.copy(model_hparams)
-    if (self.has_inputs and
-        hasattr(model_hparams, "shared_embedding_and_softmax_weights") and
+    if (hasattr(model_hparams, "shared_embedding_and_softmax_weights") and
         model_hparams.shared_embedding_and_softmax_weights):
       # If vocabularies differ, unset shared_embedding_and_softmax_weights.
       input_vocab_size = hp.input_modality.get("inputs")[1]
@@ -507,7 +506,7 @@ def get_hparams(self, model_hparams=None):
         tf.logging.warn("Unsetting shared_embedding_and_softmax_weights.")
         model_hparams.shared_embedding_and_softmax_weights = 0
 
-    _create_modalities(hp, model_hparams, self.has_inputs)
+    _create_modalities(hp, model_hparams)
     self._hparams = hp
     return self._hparams
 
@@ -1094,7 +1093,7 @@ def _reverse_problem_hparams(p_hparams):
   p.was_reversed = True
 
 
-def _create_modalities(problem_hparams, hparams, has_inputs):
+def _create_modalities(problem_hparams, hparams):
   """Converts string-type modalities to their corresponding Modality.
 
   Args:
@@ -1106,32 +1105,30 @@ def _create_modalities(problem_hparams, hparams, has_inputs):
     hparams: tf.contrib.training.HParams for the model. It may have
       input_modalities and target_modality, which will override
       problem_hparams' modalities.
-    has_inputs: A boolean that indicates whether to update the input modality.
 
   Returns:
     None
   """
-  if has_inputs:
-    input_modality_overrides = {}
-    if hasattr(hparams, "input_modalities"):
-      for override_str in hparams.input_modalities.split(";"):
-        if override_str != "default":
-          parts = override_str.split(":")
-          feature_name = parts[0]
-          modality_name = ":".join(parts[1:])
-          input_modality_overrides[feature_name] = modality_name
-
-    input_modality = {}
-    for feature_name, modality in six.iteritems(problem_hparams.input_modality):
-      if isinstance(modality, (list, tuple)):
-        if feature_name in input_modality_overrides:
-          _warn_changed_modality_type(input_modality_overrides[feature_name],
-                                      modality[0],
-                                      feature_name)
-          modality = (input_modality_overrides[feature_name], modality[1])
-        modality = modalities.create_modality(modality, hparams)
-      input_modality[feature_name] = modality
-    problem_hparams.input_modality = input_modality
+  input_modality_overrides = {}
+  if hasattr(hparams, "input_modalities"):
+    for override_str in hparams.input_modalities.split(";"):
+      if override_str != "default":
+        parts = override_str.split(":")
+        feature_name = parts[0]
+        modality_name = ":".join(parts[1:])
+        input_modality_overrides[feature_name] = modality_name
+
+  input_modality = {}
+  for feature_name, modality in six.iteritems(problem_hparams.input_modality):
+    if isinstance(modality, (list, tuple)):
+      if feature_name in input_modality_overrides:
+        _warn_changed_modality_type(input_modality_overrides[feature_name],
+                                    modality[0],
+                                    feature_name)
+        modality = (input_modality_overrides[feature_name], modality[1])
+      modality = modalities.create_modality(modality, hparams)
+    input_modality[feature_name] = modality
+  problem_hparams.input_modality = input_modality
 
   target_modality_name = None
   if (hasattr(hparams, "target_modality") and

From 7469553e5933a9355d0eee4e06fd5cd916847c80 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 20 Sep 2018 18:26:02 -0700
Subject: [PATCH 0892/2720] Correct placement of resizing and reward clipping
 wrapper, save train data videos, use dumper TF paths. Also renaming
 rl_modelrl_ to rlmb_ to shorten the docs.

PiperOrigin-RevId: 213919766
---
 tensor2tensor/data_generators/gym_problems.py |  46 ++--
 .../data_generators/gym_problems_specs.py     |  20 +-
 .../data_generators/gym_problems_test.py      |   5 +-
 .../video/basic_deterministic_params.py       |   2 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py    |   5 +-
 tensor2tensor/rl/trainer_model_based.py       | 202 ++++++++++--------
 .../rl/trainer_model_based_ae_test.py         |   2 +-
 .../rl/trainer_model_based_stochastic_test.py |   2 +-
 .../rl/trainer_model_based_sv2p_test.py       |   2 +-
 tensor2tensor/rl/trainer_model_based_test.py  |   2 +-
 10 files changed, 158 insertions(+), 130 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 9db1eaa6f..dc07826b2 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -91,8 +91,6 @@ def standard_atari_ae_env_spec(env):
 
 
 frame_dumper_use_disk = False  # Whether to use memory or disk to dump frames.
-
-
 frame_dumper = {}
 
 
@@ -116,19 +114,22 @@ def __init__(self, *args, **kwargs):
     self.statistics = BasicStatistics()
     self._use_dumper_data = False
     self._dumper_data_index = 0
+    self._forced_collect_level = None
 
   @property
   def resize_height_factor(self):
-    return 2
+    return 1
 
   @property
   def resize_width_factor(self):
-    return 2
+    return 1
 
   def _setup(self, data_dir, extra_collect_hparams=None,
              override_collect_hparams=None):
     dumper_path = os.path.join(data_dir, "dumper")
-    if os.path.isdir(dumper_path):
+    dumper_exists = tf.gfile.Exists(dumper_path)
+    tf.logging.info("Dumper path %s." % dumper_path)
+    if dumper_exists and not self.settable_eval_phase:
       tf.logging.info("Using dumper data.")
       self._use_dumper_data = True
       self._dumper_data_index = 0
@@ -160,13 +161,19 @@ def _setup(self, data_dir, extra_collect_hparams=None,
       if self.settable_eval_phase:
         policy_to_actions_lambda = lambda policy: policy.mode()
 
+      collect_level = 2  # After Resize and RewardClipping.
+      if collect_hparams.environment_spec.simulated_env:
+        collect_level = 1  # We still have reward clipping.
+      if self._forced_collect_level is not None:  # For autoencoders.
+        collect_level = self._forced_collect_level
+
       with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
         self.collect_memory, self.collect_trigger_op, collect_init = (
             collect.define_collect(
                 collect_hparams,
                 scope="gym_problems",
                 eval_phase=False,
-                collect_level=1,  # After ResizeWrapper but before others.
+                collect_level=collect_level,
                 policy_to_actions_lambda=policy_to_actions_lambda))
 
       self._session = tf.Session()
@@ -204,17 +211,8 @@ def _get_data(self):
   def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
     self._setup(data_dir)
 
-    # We only want to save frames for eval and simulated experience, not the
-    # frames used for world model training.
-    base_dir = os.path.basename(os.path.dirname(data_dir + "/"))
-    if (base_dir == "eval" or self.debug_dump_frames_path in [
-        "debug_frames_sim_eval", "debug_frames_sim"
-    ]):
-      self.debug_dump_frames_path = os.path.join(data_dir,
-                                                 self.debug_dump_frames_path)
-    else:
-      # Disable frame saving
-      self.debug_dump_frames_path = ""
+    self.debug_dump_frames_path = os.path.join(
+        data_dir, self.debug_dump_frames_path)
 
     frame_counter = 0
     pieces_generated = 0
@@ -458,6 +456,10 @@ def collect_statistics_and_generate_debug_image(self, index, observation,
 class GymDiscreteProblemWithAutoencoder(GymRealDiscreteProblem):
   """Gym discrete problem with autoencoder."""
 
+  def __init__(self, *args, **kwargs):
+    super(GymDiscreteProblemWithAutoencoder, self).__init__(*args, **kwargs)
+    self._forced_collect_level = 0
+
   def get_environment_spec(self):
     return standard_atari_ae_env_spec(self.env_name)
 
@@ -491,6 +493,10 @@ def set_targets(example):
 class GymDiscreteProblemAutoencoded(GymRealDiscreteProblem):
   """Gym discrete problem with frames already autoencoded."""
 
+  def __init__(self, *args, **kwargs):
+    super(GymDiscreteProblemAutoencoded, self).__init__(*args, **kwargs)
+    self._forced_collect_level = 0
+
   def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
     raise RuntimeError("GymDiscreteProblemAutoencoded can be used only"
                        " for reading encoded frames")
@@ -534,7 +540,6 @@ def __init__(self):
     self.episode_real_reward = 0.0
     self.successful_episode_reward_predictions = 0
     self.report_reward_statistics_every = 10
-
     # auxiliary objects
     self.real_obs = None
     self.real_rewards = None
@@ -747,6 +752,11 @@ def _generate_debug_image(self, real_ob, sim_ob):
 class GymSimulatedDiscreteProblemAutoencoded(GymSimulatedDiscreteProblem):
   """Gym simulated discrete problem with frames already autoencoded."""
 
+  def __init__(self, *args, **kwargs):
+    super(GymSimulatedDiscreteProblemAutoencoded, self).__init__(
+        *args, **kwargs)
+    self._forced_collect_level = 0
+
   def get_environment_spec(self):
     env_spec = standard_atari_env_spec(self.env_name)
     env_spec.wrappers = [
diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index 239d074ef..75d165860 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -50,6 +50,8 @@
     "zaxxon"
 ]
 
+ATARI_ALL_MODES_SHORT_LIST = []
+
 ATARI_WHITELIST_GAMES = [
     "amidar",
     "bank_heist",
@@ -66,11 +68,6 @@
     "seaquest",
 ]
 
-ATARI_ALL_MODES_SHORT_LIST = [
-    "pong",
-    "boxing",
-]
-
 # Different ATARI game modes in OpenAI Gym. Full list here:
 # https://github.com/openai/gym/blob/master/gym/envs/__init__.py
 ATARI_GAME_MODES = [
@@ -200,8 +197,8 @@ def num_rewards(self):
 
 def create_problems_for_game(
     game_name,
-    resize_height_factor=2,
-    resize_width_factor=2,
+    resize_height_factor=1,
+    resize_width_factor=1,
     game_mode="Deterministic-v4"):
   """Create and register problems for game_name.
 
@@ -264,12 +261,3 @@ def create_problems_for_game(
       "simulated": simulated_cls,
       "world_model_eval": world_model_eval_cls,
   }
-
-# Register the atari games with all of the possible modes.
-for game in ATARI_ALL_MODES_SHORT_LIST:
-  ATARI_PROBLEMS[game] = {}
-  for mode in ATARI_GAME_MODES:
-    classes = create_problems_for_game(
-        game,
-        game_mode=mode)
-    ATARI_PROBLEMS[game][mode] = classes
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index ebd7b4112..3ee824a48 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -35,9 +35,8 @@ def setUpClass(cls):
     os.mkdir(cls.tmp_dir)
 
   def testGymAtariGameModes(self):
-    for mode in gym_problems_specs.ATARI_GAME_MODES:
-      problem = gym_problems_specs.ATARI_PROBLEMS["pong"][mode]["base"]()
-      self.assertEqual(105, problem.frame_height)
+    problem = gym_problems_specs.GymDiscreteProblemWithAgentOnWrappedFullPong()
+    self.assertEqual(210, problem.frame_height)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index c4724f239..56425f572 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -31,7 +31,7 @@ def next_frame_basic_deterministic():
   hparams.batch_size = 4
   hparams.num_hidden_layers = 2
   hparams.optimizer = "Adafactor"
-  hparams.learning_rate_constant = 0.7
+  hparams.learning_rate_constant = 1.5
   hparams.learning_rate_warmup_steps = 8000
   hparams.learning_rate_schedule = "linear_warmup * constant * rsqrt_decay"
   hparams.label_smoothing = 0.0
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 7302191ca..0486d49ea 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -107,6 +107,9 @@ def simulate(self, action):
     with tf.control_dependencies([reward, done]):
       return tf.sign(reward), tf.identity(done)
 
+  def _reset_non_empty(self, indices):
+    return self._batch_env._reset_non_empty(indices)  # pylint: disable=protected-access
+
 
 class MaxAndSkipWrapper(WrapperBase):
   """ Max and skip wrapper.
@@ -304,7 +307,7 @@ def _resize(self, tensor):
     height, width, _ = self.observ_shape
     observ = tf.to_float(tensor)
     resized = tf.image.resize_images(
-        observ, [height, width], tf.image.ResizeMethod.BILINEAR)
+        observ, [height, width], tf.image.ResizeMethod.AREA)
     return tf.cast(resized, self.observ_dtype)
 
   @property
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 1da607be0..68c457df6 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -18,7 +18,7 @@
 
 python -m tensor2tensor.rl.trainer_model_based \
     --output_dir=$HOME/t2t/rl_v1 \
-    --loop_hparams_set=rl_modelrl_base \
+    --loop_hparams_set=rlmb_base \
     --loop_hparams='num_real_env_frames=10000,epochs=3'
 """
 from __future__ import absolute_import
@@ -53,7 +53,7 @@
 flags = tf.flags
 FLAGS = flags.FLAGS
 
-flags.DEFINE_string("loop_hparams_set", "rl_modelrl_base",
+flags.DEFINE_string("loop_hparams_set", "rlmb_base",
                     "Which RL hparams set to use.")
 flags.DEFINE_string("loop_hparams", "", "Overrides for overall loop HParams.")
 
@@ -265,7 +265,8 @@ def train_agent_real_env(
 
   ppo_hparams.epochs_num = _ppo_training_epochs(hparams, epoch,
                                                 is_final_epoch, True)
-  ppo_hparams.save_models_every_epochs = 10
+  # We do not save model, as that resets frames that we need at restarts.
+  ppo_hparams.save_models_every_epochs = 0
 
   environment_spec = copy.copy(gym_problem.environment_spec)
 
@@ -278,7 +279,7 @@ def train_agent_real_env(
     dumper_path = os.path.join(epoch_data_dir, "dumper")
     tf.gfile.MakeDirs(dumper_path)
     dumper_spec = [PyFuncWrapper, {"process_fun": ppo_data_dumper}]
-    environment_spec.wrappers.insert(1, dumper_spec)
+    environment_spec.wrappers.insert(2, dumper_spec)
 
   ppo_hparams.add_hparam("environment_spec", environment_spec)
 
@@ -531,6 +532,9 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   mean_reward_summary = tf.Summary()
   mean_reward_summary.value.add(tag="mean_reward",
                                 simple_value=None)
+  mean_reward_gen_summary = tf.Summary()
+  mean_reward_gen_summary.value.add(tag="mean_reward_during_generation",
+                                    simple_value=None)
 
   for epoch in range(hparams.epochs):
     is_final_epoch = (epoch + 1) == hparams.epochs
@@ -591,29 +595,33 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
     if hparams.stop_loop_early:
       return 0.0
-    # Collect data from the real environment.
-    log("Generating real environment data")
-    eval_data_dir = os.path.join(epoch_data_dir, "eval")
-    mean_reward = generate_real_env_data(
-        problem_name, ppo_model_dir, hparams, eval_data_dir,
-        directories["tmp"], autoencoder_path=autoencoder_model_dir,
-        eval_phase=True)
-    log("Mean eval reward: {}".format(mean_reward))
 
+    # Collect data from the real environment.
     if not is_final_epoch:
+      log("Generating real environment data.")
       generation_mean_reward = generate_real_env_data(
           problem_name, ppo_model_dir, hparams, epoch_data_dir,
           directories["tmp"], autoencoder_path=autoencoder_model_dir,
           eval_phase=False)
       log("Mean reward during generation: {}".format(generation_mean_reward))
 
-    # Summarize metrics
+    log("Evaluating in real environment.")
+    eval_data_dir = os.path.join(epoch_data_dir, "eval")
+    mean_reward = generate_real_env_data(
+        problem_name, ppo_model_dir, hparams, eval_data_dir,
+        directories["tmp"], autoencoder_path=autoencoder_model_dir,
+        eval_phase=True)
+    log("Mean eval reward: {}".format(mean_reward))
+
+    # Summarize metrics.
     assert model_reward_accuracy is not None
     assert mean_reward is not None
     model_reward_accuracy_summary.value[0].simple_value = model_reward_accuracy
     mean_reward_summary.value[0].simple_value = mean_reward
+    mean_reward_gen_summary.value[0].simple_value = generation_mean_reward
     eval_metrics_writer.add_summary(model_reward_accuracy_summary, epoch)
     eval_metrics_writer.add_summary(mean_reward_summary, epoch)
+    eval_metrics_writer.add_summary(mean_reward_gen_summary, epoch)
     eval_metrics_writer.flush()
 
     # Report metrics
@@ -650,7 +658,7 @@ def combine_training_data(problem, final_data_dir, old_data_dirs,
 
 
 @registry.register_hparams
-def rl_modelrl_base():
+def rlmb_base():
   return tf.contrib.training.HParams(
       epochs=6,
       # Total frames used for training. This will be distributed evenly across
@@ -703,9 +711,10 @@ def rl_modelrl_base():
 
 
 @registry.register_hparams
-def rl_modelrl_basetest():
+def rlmb_basetest():
   """Base setting but quicker with only 2 epochs."""
-  hparams = rl_modelrl_base()
+  hparams = rlmb_base()
+  hparams.game = "pong"
   hparams.epochs = 2
   hparams.num_real_env_frames = 3200
   hparams.model_train_steps = 500
@@ -715,96 +724,113 @@ def rl_modelrl_basetest():
 
 
 @registry.register_hparams
-def rl_modelrl_base_quick():
+def rlmb_noresize():
+  hparams = rlmb_base()
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_quick():
   """Base setting but quicker with only 2 epochs."""
-  hparams = rl_modelrl_base()
+  hparams = rlmb_base()
   hparams.epochs = 2
-  hparams.ppo_epochs_num = 1000
+  hparams.model_train_steps = 25000
+  hparams.ppo_epochs_num = 700
   hparams.ppo_epoch_length = 50
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_base_quick_sd():
+def rlmb_quick_noresize():
+  hparams = rlmb_base()
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_quick_sd():
   """Quick setting with stochastic discrete model."""
-  hparams = rl_modelrl_base_quick()
+  hparams = rlmb_quick()
   hparams.generative_model = "next_frame_basic_stochastic_discrete"
   hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_base_quick_sm():
+def rlmb_quick_sm():
   """Quick setting with sampling."""
-  hparams = rl_modelrl_base_quick()
+  hparams = rlmb_quick()
   hparams.generative_model_params = "next_frame_sampling"
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_base_stochastic():
+def rlmb_base_stochastic():
   """Base setting with a stochastic next-frame model."""
-  hparams = rl_modelrl_base()
+  hparams = rlmb_base()
   hparams.generative_model = "next_frame_basic_stochastic"
   hparams.generative_model_params = "next_frame_basic_stochastic"
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_base_stochastic_discrete():
+def rlmb_base_stochastic_discrete():
   """Base setting with stochastic discrete model."""
-  hparams = rl_modelrl_base()
+  hparams = rlmb_base()
   hparams.generative_model = "next_frame_basic_stochastic_discrete"
   hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_base_sv2p():
+def rlmb_base_sv2p():
   """Base setting with sv2p as world model."""
-  hparams = rl_modelrl_base()
+  hparams = rlmb_base()
   hparams.generative_model = "next_frame_sv2p"
   hparams.generative_model_params = "next_frame_sv2p_atari"
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_base_sv2p_deterministic():
+def rlmb_base_sv2p_deterministic():
   """Base setting with deterministic sv2p as world model."""
-  hparams = rl_modelrl_base_sv2p()
+  hparams = rlmb_base_sv2p()
   hparams.generative_model = "next_frame_sv2p"
   hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_base_sampling():
+def rlmb_base_sampling():
   """Base setting with a stochastic next-frame model."""
-  hparams = rl_modelrl_base()
+  hparams = rlmb_base()
   hparams.generative_model_params = "next_frame_sampling"
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_medium():
+def rlmb_medium():
   """Small set for larger testing."""
-  hparams = rl_modelrl_base()
+  hparams = rlmb_base()
   hparams.num_real_env_frames //= 2
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_25k():
+def rlmb_25k():
   """Small set for larger testing."""
-  hparams = rl_modelrl_medium()
+  hparams = rlmb_medium()
   hparams.num_real_env_frames //= 2
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_short():
+def rlmb_short():
   """Small set for larger testing."""
-  hparams = rl_modelrl_base()
+  hparams = rlmb_base()
   hparams.num_real_env_frames //= 5
   hparams.model_train_steps //= 10
   hparams.ppo_epochs_num //= 10
@@ -812,17 +838,17 @@ def rl_modelrl_short():
 
 
 @registry.register_hparams
-def rl_modelrl_model_only():
-  hp = rl_modelrl_base()
+def rlmb_model_only():
+  hp = rlmb_base()
   hp.epochs = 1
   hp.ppo_epochs_num = 0
   return hp
 
 
 @registry.register_hparams
-def rl_modelrl_tiny():
+def rlmb_tiny():
   """Tiny set for testing."""
-  return rl_modelrl_base_sampling().override_from_dict(
+  return rlmb_base_sampling().override_from_dict(
       tf.contrib.training.HParams(
           epochs=1,
           num_real_env_frames=128,
@@ -845,9 +871,9 @@ def rl_modelrl_tiny():
 
 
 @registry.register_hparams
-def rl_modelrl_tiny_stochastic():
+def rlmb_tiny_stochastic():
   """Tiny setting with a stochastic next-frame model."""
-  hparams = rl_modelrl_tiny()
+  hparams = rlmb_tiny()
   hparams.epochs = 1  # Too slow with 2 for regular runs.
   hparams.generative_model = "next_frame_basic_stochastic"
   hparams.generative_model_params = "next_frame_basic_stochastic"
@@ -855,63 +881,65 @@ def rl_modelrl_tiny_stochastic():
 
 
 @registry.register_hparams
-def rl_modelrl_tiny_sv2p():
+def rlmb_tiny_sv2p():
   """Tiny setting with a tiny sv2p model."""
-  hparams = rl_modelrl_tiny()
+  hparams = rlmb_tiny()
   hparams.generative_model = "next_frame_sv2p"
   hparams.generative_model_params = "next_frame_sv2p_tiny"
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_l1_base():
+def rlmb_l1_base():
   """Parameter set with L1 loss."""
-  hparams = rl_modelrl_base()
+  hparams = rlmb_base()
   hparams.generative_model_params = "next_frame_l1"
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_l1_tiny():
+def rlmb_l1_tiny():
   """Tiny parameter set with L1 loss."""
-  hparams = rl_modelrl_tiny()
+  hparams = rlmb_tiny()
   hparams.generative_model_params = "next_frame_l1"
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_l2_base():
+def rlmb_l2_base():
   """Parameter set with L2 loss."""
-  hparams = rl_modelrl_base()
+  hparams = rlmb_base()
   hparams.generative_model_params = "next_frame_l2"
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_l2_tiny():
+def rlmb_l2_tiny():
   """Tiny parameter set with L2 loss."""
-  hparams = rl_modelrl_tiny()
+  hparams = rlmb_tiny()
   hparams.generative_model_params = "next_frame_l2"
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_ae_base():
+def rlmb_ae_base():
   """Parameter set for autoencoders."""
-  hparams = rl_modelrl_base()
+  hparams = rlmb_base()
   hparams.ppo_params = "ppo_pong_ae_base"
   hparams.generative_model_params = "next_frame_ae"
+  hparams.gather_ppo_real_env_data = False
   hparams.autoencoder_train_steps = 30000
   return hparams
 
 
 @registry.register_hparams
-def rl_modelrl_ae_tiny():
+def rlmb_ae_tiny():
   """Tiny set for testing autoencoders."""
-  hparams = rl_modelrl_tiny()
+  hparams = rlmb_tiny()
   hparams.game = "wrapped_full_pong"
   hparams.ppo_params = "ppo_pong_ae_base"
   hparams.generative_model_params = "next_frame_ae"
+  hparams.gather_ppo_real_env_data = False
   hparams.resize_height_factor = 1
   hparams.resize_width_factor = 1
   hparams.autoencoder_train_steps = 2
@@ -919,8 +947,8 @@ def rl_modelrl_ae_tiny():
 
 
 @registry.register_hparams
-def rl_modelrl_tiny_simulation_deterministic_starts():
-  hp = rl_modelrl_tiny()
+def rlmb_tiny_simulation_deterministic_starts():
+  hp = rlmb_tiny()
   hp.simulation_random_starts = False
   return hp
 
@@ -931,7 +959,7 @@ def rl_modelrl_tiny_simulation_deterministic_starts():
 # HP_SCOPES={loop, model, ppo}, which set hyperparameters for the top-level
 # hparams, hp.generative_model_params, and hp.ppo_params, respectively.
 @registry.register_ranged_hparams
-def rl_modelrl_grid(rhp):
+def rlmb_grid(rhp):
   """Grid over games and frames, and 5 runs each for variance."""
   rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
   base = 100000
@@ -944,26 +972,26 @@ def rl_modelrl_grid(rhp):
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_variance(rhp):
+def rlmb_variance(rhp):
   # Dummy parameter to get 5 runs for each configuration
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
   rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_variance_nogame(rhp):
-  # Dummy parameter to get 10 runs for current configuration.
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+def rlmb_variance_nogame(rhp):
+  # Dummy parameter to get 20 runs for current configuration.
+  rhp.set_discrete("model.moe_loss_coef", list(range(20)))
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_three(rhp):
+def rlmb_three(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(10)))
   rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_test1(rhp):
+def rlmb_test1(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(10)))
   rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
   rhp.set_discrete("loop.ppo_learning_rate", [5e-5, 1e-4, 2e-4])
@@ -972,24 +1000,24 @@ def rl_modelrl_test1(rhp):
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_scheduled_sampling(rhp):
+def rlmb_scheduled_sampling(rhp):
   rhp.set_float("model.scheduled_sampling_prob", 0.0, 1.0)
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_all_games(rhp):
+def rlmb_all_games(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
   rhp.set_categorical("loop.game", gym_problems_specs.ATARI_GAMES)
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_whitelisted_games(rhp):
+def rlmb_whitelisted_games(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(10)))
   rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_ae_variance(rhp):
+def rlmb_ae_variance(rhp):
   # Dummy parameter to get 5 runs for each configuration
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
   rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
@@ -999,49 +1027,49 @@ def rl_modelrl_ae_variance(rhp):
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_ppolr_game(rhp):
+def rlmb_ppolr_game(rhp):
   rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
   base_lr = 1e-4
   rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_ppolr(rhp):
+def rlmb_ppolr(rhp):
   base_lr = 1e-4
   rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_ae_ppo_lr(rhp):
+def rlmb_ae_ppo_lr(rhp):
   rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
   base_lr = 1e-4
   rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_dropout_range(rhp):
+def rlmb_dropout_range(rhp):
   rhp.set_float("model.dropout", 0.2, 0.4)
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_intrinsic_reward_scale(rhp):
+def rlmb_intrinsic_reward_scale(rhp):
   rhp.set_float("loop.intrinsic_reward_scale", 0.01, 10.)
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_l1l2cutoff_range(rhp):
+def rlmb_l1l2cutoff_range(rhp):
   """Loss and loss-cutoff tuning grid."""
   rhp.set_float("model.video_modality_loss_cutoff", 1.4, 3.4)
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_xentcutoff_range(rhp):
+def rlmb_xentcutoff_range(rhp):
   """Cross entropy cutoff tuning grid."""
   rhp.set_float("model.video_modality_loss_cutoff", 0.01, 0.05)
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_pixel_noise(rhp):
+def rlmb_pixel_noise(rhp):
   """Input pixel noise tuning grid."""
   rhp.set_categorical("loop.generative_model_params",
                       ["next_frame_pixel_noise"])
@@ -1050,34 +1078,34 @@ def rl_modelrl_pixel_noise(rhp):
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_dummy_range(rhp):
+def rlmb_dummy_range(rhp):
   """Dummy tuning grid just to get the variance."""
   rhp.set_float("model.moe_loss_coef", 0.01, 0.02)
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_epochs_num(rhp):
+def rlmb_epochs_num(rhp):
   rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
   rhp.set_discrete("loop.epochs", [3, 6, 12])
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_ppo_epochs_num(rhp):
+def rlmb_ppo_epochs_num(rhp):
   rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
   rhp.set_discrete("loop.ppo_epochs_num", [200, 1000, 2000, 4000])
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_ppo_epoch_len(rhp):
+def rlmb_ppo_epoch_len(rhp):
   rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
   rhp.set_discrete("loop.ppo_epoch_length", [25, 50, 100])
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_num_frames(rhp):
+def rlmb_num_frames(rhp):
   rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
   rhp.set_discrete("loop.num_real_env_frames",
@@ -1085,14 +1113,14 @@ def rl_modelrl_num_frames(rhp):
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_ppo_optimization_batch_size(rhp):
+def rlmb_ppo_optimization_batch_size(rhp):
   rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
   rhp.set_discrete("model.moe_loss_coef", list(range(10)))
   rhp.set_discrete("ppo.optimization_batch_size", [4, 10, 20])
 
 
 @registry.register_ranged_hparams
-def rl_modelrl_logits_clip(rhp):
+def rlmb_logits_clip(rhp):
   rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
   rhp.set_discrete("model.moe_loss_coef", list(range(10)))
   rhp.set_discrete("ppo.logits_clip", [0., 5.])
diff --git a/tensor2tensor/rl/trainer_model_based_ae_test.py b/tensor2tensor/rl/trainer_model_based_ae_test.py
index aac2232a3..e44b477c8 100644
--- a/tensor2tensor/rl/trainer_model_based_ae_test.py
+++ b/tensor2tensor/rl/trainer_model_based_ae_test.py
@@ -29,7 +29,7 @@ class ModelRLExperimentTestAe(tf.test.TestCase):
   def test_ae(self):
     # TODO(lukaszkaiser): re-enable this test.
     # FLAGS.output_dir = tf.test.get_temp_dir()
-    # FLAGS.loop_hparams_set = "rl_modelrl_ae_tiny"
+    # FLAGS.loop_hparams_set = "rlmb_ae_tiny"
     # FLAGS.schedule = "train"  # skip evaluation for world model training
     # trainer_model_based.main(None)
     assert True
diff --git a/tensor2tensor/rl/trainer_model_based_stochastic_test.py b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
index 24f69df2a..fd22fa021 100644
--- a/tensor2tensor/rl/trainer_model_based_stochastic_test.py
+++ b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
@@ -28,7 +28,7 @@ class ModelRLExperimentStochasticTest(tf.test.TestCase):
 
   def test_basic_stochastic(self):
     FLAGS.output_dir = tf.test.get_temp_dir()
-    FLAGS.loop_hparams_set = "rl_modelrl_tiny_stochastic"
+    FLAGS.loop_hparams_set = "rlmb_tiny_stochastic"
     FLAGS.schedule = "train"  # skip evaluation for world model training
     trainer_model_based.main(None)
 
diff --git a/tensor2tensor/rl/trainer_model_based_sv2p_test.py b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
index a59aab6ec..5b715c20c 100644
--- a/tensor2tensor/rl/trainer_model_based_sv2p_test.py
+++ b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
@@ -28,7 +28,7 @@ class ModelRLExperimentSv2pTest(tf.test.TestCase):
 
   def test_sv2p(self):
     FLAGS.output_dir = tf.test.get_temp_dir()
-    FLAGS.loop_hparams_set = "rl_modelrl_tiny_sv2p"
+    FLAGS.loop_hparams_set = "rlmb_tiny_sv2p"
     trainer_model_based.main(None)
 
 
diff --git a/tensor2tensor/rl/trainer_model_based_test.py b/tensor2tensor/rl/trainer_model_based_test.py
index 48e2f40ea..f528af1c9 100644
--- a/tensor2tensor/rl/trainer_model_based_test.py
+++ b/tensor2tensor/rl/trainer_model_based_test.py
@@ -28,7 +28,7 @@ class ModelRLExperimentTest(tf.test.TestCase):
 
   def test_basic(self):
     FLAGS.output_dir = tf.test.get_temp_dir()
-    FLAGS.loop_hparams_set = "rl_modelrl_tiny"
+    FLAGS.loop_hparams_set = "rlmb_tiny"
     FLAGS.schedule = "train"  # skip evaluation for world model training
     trainer_model_based.main(None)
 

From bd56496eb0446970b75024c45d0dd8471a85aea6 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 20 Sep 2018 18:56:34 -0700
Subject: [PATCH 0893/2720] enable softmax for sv2p.

PiperOrigin-RevId: 213922389
---
 .../models/video/basic_deterministic.py       | 23 +++++++-----
 .../video/basic_deterministic_params.py       |  1 -
 tensor2tensor/models/video/sv2p.py            | 36 ++++++++++++++-----
 tensor2tensor/models/video/sv2p_params.py     | 20 ++++++++++-
 tensor2tensor/rl/trainer_model_based.py       | 18 +++++++++-
 5 files changed, 79 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 00783b3ab..1be7701ca 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -38,10 +38,15 @@
 class NextFrameBasicDeterministic(t2t_model.T2TModel):
   """Basic next-frame model, may take actions and predict rewards too."""
 
+  @property
+  def _target_modality(self):
+    # TODO(mbz): get rid of this somehow.
+    modality = self.hparams.problem_hparams.target_modality["targets"]
+    return modality.__class__.__name__
+
   @property
   def is_per_pixel_softmax(self):
-    # TODO(mbz): this should not be a hyper parameter.
-    return self.hparams.per_pixel_softmax
+    return self._target_modality == "VideoModality"
 
   def inject_latent(self, layer, features, filters):
     """Do nothing for deterministic model."""
@@ -73,11 +78,13 @@ def inject_additional_input(self, layer, inputs, name, mode="concat"):
 
     return layer
 
-  def get_sampled_frame(self, res_frame, orig_frame_shape):
-    target_shape = orig_frame_shape[:-1] + [self.hparams.problem.num_channels]
-    if self.is_per_pixel_softmax:
-      sampled_frame = tf.reshape(res_frame, target_shape + [256])
-      sampled_frame = tf.to_float(tf.argmax(sampled_frame, axis=-1))
+  def get_sampled_frame(self, res_frame):
+    if not self.is_per_pixel_softmax:
+      return res_frame
+    frame_shape = common_layers.shape_list(res_frame)
+    target_shape = frame_shape[:-1] + [self.hparams.problem.num_channels]
+    sampled_frame = tf.reshape(res_frame, target_shape + [256])
+    sampled_frame = tf.to_float(tf.argmax(sampled_frame, axis=-1))
     return sampled_frame
 
   def body_single(self, features):
@@ -206,7 +213,7 @@ def body(self, features):
       res_frames.append(res_frame)
 
       # Only for Softmax loss: sample frame so we can keep iterating.
-      sampled_frame_raw = self.get_sampled_frame(res_frame, orig_frame_shape)
+      sampled_frame_raw = self.get_sampled_frame(res_frame)
       sampled_frames_raw.append(sampled_frame_raw)
       # TODO(lukaszkaiser): this should be consistent with modality.bottom()
       sampled_frame = common_layers.standardize_images(sampled_frame_raw)
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 56425f572..9105905f4 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -40,7 +40,6 @@ def next_frame_basic_deterministic():
   hparams.weight_decay = 0.0
   hparams.clip_grad_norm = 1.0
   hparams.dropout = 0.5
-  hparams.add_hparam("per_pixel_softmax", True)
   # choose from: concat, multiplicative, multi_additive
   hparams.add_hparam("action_injection", "multi_additive")
   hparams.add_hparam("num_compress_steps", 6)
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 58222eaf0..d9428a384 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -55,6 +55,7 @@ def concat_on_y_axis(x):
 
     frames_gd = common_video.swap_time_and_batch_axes(real_frames)
     frames_pd = common_video.swap_time_and_batch_axes(gen_frames)
+    frames_pd = self.get_sampled_frame(frames_pd)
     frames_gd = concat_on_y_axis(frames_gd)
     frames_pd = concat_on_y_axis(frames_pd)
     side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
@@ -388,6 +389,10 @@ def construct_predictive_tower(
         layer = layer[:, :img_height, :img_width, :]
         output += layer * mask
 
+      if self.is_per_pixel_softmax:
+        output = tf.layers.dense(
+            output, self.hparams.problem.num_channels * 256, name="logits")
+
       return output, lstm_state
 
   def construct_model(self,
@@ -420,7 +425,7 @@ def construct_model(self,
       raise ValueError("Buffer size is bigger than context frames %d %d." %
                        (buffer_size, context_frames))
 
-    batch_size = common_layers.shape_list(images)[1]
+    batch_size = common_layers.shape_list(images[0])[0]
     ss_func = self.get_scheduled_sample_func(batch_size)
 
     def process_single_frame(prev_outputs, inputs):
@@ -428,6 +433,9 @@ def process_single_frame(prev_outputs, inputs):
       cur_image, input_reward, action = inputs
       time_step, prev_image, prev_reward, frame_buf, lstm_states = prev_outputs
 
+      # sample from softmax (by argmax). this is noop for non-softmax loss.
+      prev_image = self.get_sampled_frame(prev_image)
+
       generated_items = [prev_image]
       groundtruth_items = [cur_image]
       done_warm_start = tf.greater(time_step, context_frames - 1)
@@ -439,7 +447,7 @@ def process_single_frame(prev_outputs, inputs):
           input_image, None, action, lstm_states, latent)
 
       if self.hparams.reward_prediction:
-        reward_input_image = pred_image
+        reward_input_image = self.get_sampled_frame(pred_image)
         if self.hparams.reward_prediction_stop_gradient:
           reward_input_image = tf.stop_gradient(reward_input_image)
         with tf.control_dependencies([time_step]):
@@ -466,8 +474,12 @@ def process_single_frame(prev_outputs, inputs):
     lstm_states = [None] * (5 if self.hparams.small_mode else 7)
     frame_buffer = [tf.zeros_like(images[0])] * buffer_size
     inputs = images[0], rewards[0], actions[0]
+    init_image_shape = common_layers.shape_list(images[0])
+    if self.is_per_pixel_softmax:
+      init_image_shape[-1] *= 256
+    init_image = tf.zeros(init_image_shape, dtype=images.dtype)
     prev_outputs = (tf.constant(0),
-                    tf.zeros_like(images[0]),
+                    init_image,
                     tf.zeros_like(rewards[0]),
                     frame_buffer,
                     lstm_states)
@@ -525,8 +537,16 @@ def infer(self, features, *args, **kwargs):
     if not isinstance(output, dict):
       output = {"targets": output}
 
-    output["targets"] = tf.squeeze(output["targets"], axis=-1)
-    output["targets"] = tf.to_int64(tf.round(output["targets"]))
+    x = output["targets"]
+    if self.is_per_pixel_softmax:
+      x_shape = common_layers.shape_list(x)
+      x = tf.reshape(x, [-1, x_shape[-1]])
+      x = tf.argmax(x, axis=-1)
+      x = tf.reshape(x, x_shape[:-1])
+    else:
+      x = tf.squeeze(x, axis=-1)
+      x = tf.to_int64(tf.round(x))
+    output["targets"] = x
     if self.hparams.reward_prediction:
       output["target_reward"] = tf.argmax(output["target_reward"], axis=-1)
 
@@ -578,7 +598,8 @@ def body(self, features):
         gen_frames=gen_images)
 
     # Visualize predictions in Tensorboard
-    self.visualize_predictions(all_frames[1:], gen_images)
+    if self.is_training and not self.is_per_pixel_softmax:
+      self.visualize_predictions(all_frames[1:], gen_images)
 
     # Ignore the predictions from the input frames.
     # This is NOT the same as original paper/implementation.
@@ -594,8 +615,7 @@ def body(self, features):
       # add the MSE loss for input frames as well.
       # we are assuming the modality is L2. otherwise the loss would be
       # incosistent across the frames.
-      modality = self.hparams.problem_hparams.target_modality["targets"]
-      if modality.__class__.__name__ != "VideoModalityL2Raw":
+      if self._target_modality != "VideoModalityL2Raw":
         raise ValueError("internal loss only works with L2.")
       recon_loss = tf.losses.mean_squared_error(
           all_frames[1:hparams.video_num_input_frames+1],
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index a1278589c..c1929ebe1 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -69,6 +69,16 @@ def next_frame_sv2p_atari():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_sv2p_atari_softmax():
+  """SV2P model for atari with softmax."""
+  hparams = next_frame_sv2p_atari()
+  hparams.target_modality = "video"
+  hparams.input_modalities = "inputs:video"
+  hparams.internal_loss = False
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_sv2p_atari_deterministic():
   """Deterministic for atari."""
@@ -77,10 +87,18 @@ def next_frame_sv2p_atari_deterministic():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_sv2p_atari_softmax_deterministic():
+  """Deterministic for atari."""
+  hparams = next_frame_sv2p_atari_softmax()
+  hparams.stochastic_model = False
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_sv2p_tiny():
   """Tiny SV2P model."""
-  hparams = next_frame_sv2p()
+  hparams = next_frame_sv2p_atari()
   hparams.batch_size = 2
   hparams.tiny_mode = True
   hparams.num_masks = 1
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 68c457df6..6d85166c7 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -794,15 +794,31 @@ def rlmb_base_sv2p():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_base_sv2p_softmax():
+  """Base setting with sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p()
+  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_base_sv2p_deterministic():
   """Base setting with deterministic sv2p as world model."""
   hparams = rlmb_base_sv2p()
-  hparams.generative_model = "next_frame_sv2p"
   hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
   return hparams
 
 
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic_softmax():
+  """Base setting with deterministic sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p_softmax()
+  hparams.generative_model_params = (
+      "next_frame_sv2p_atari_softmax_deterministic")
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_base_sampling():
   """Base setting with a stochastic next-frame model."""

From 05a20c6c3464d40bfeb75842e3052e31b6374c72 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 21 Sep 2018 14:25:13 -0700
Subject: [PATCH 0894/2720] A number of small tweaks to improve the RL
 pipeline.

PiperOrigin-RevId: 214044030
---
 tensor2tensor/models/research/rl.py           |  2 +-
 .../models/video/basic_stochastic.py          |  4 +-
 tensor2tensor/rl/trainer_model_based.py       | 50 +++++++++++++++----
 3 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 5c9d3fafb..52031d2dd 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -53,7 +53,7 @@ def ppo_base_v1():
   hparams.add_hparam("simulation_random_starts", False)
   hparams.add_hparam("simulation_flip_first_random_for_beginning", False)
   hparams.add_hparam("intrinsic_reward_scale", 0.)
-  hparams.add_hparam("logits_clip", 3.0)
+  hparams.add_hparam("logits_clip", 4.0)
   hparams.add_hparam("dropout_ppo", 0.1)
   return hparams
 
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 03a9c343c..4c2d27acc 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -82,8 +82,10 @@ def inject_latent(self, layer, features, filters):
       return layer + z, 0.0
 
     # Embed.
+    frames = tf.concat(
+        [features["cur_target_frame"], features["inputs"]], axis=-1)
     x = tf.layers.dense(
-        features["cur_target_frame"], filters, name="latent_embed",
+        frames, filters, name="latent_embed",
         bias_initializer=tf.random_normal_initializer(stddev=0.01))
     x = common_attention.add_timing_signal_nd(x)
 
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 6d85166c7..5683d3684 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -266,7 +266,8 @@ def train_agent_real_env(
   ppo_hparams.epochs_num = _ppo_training_epochs(hparams, epoch,
                                                 is_final_epoch, True)
   # We do not save model, as that resets frames that we need at restarts.
-  ppo_hparams.save_models_every_epochs = 0
+  # But we need to save at the last step, so we set it very high.
+  ppo_hparams.save_models_every_epochs = 1000000
 
   environment_spec = copy.copy(gym_problem.environment_spec)
 
@@ -322,13 +323,7 @@ def train_world_model(problem_name, data_dir, output_dir, hparams, epoch):
   train_steps = hparams.model_train_steps * (epoch + 2)
   model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
   learning_rate = model_hparams.learning_rate_constant
-  # Bump learning rate after first epoch by 3x.
-  # We picked 3x because our default learning rate schedule decreases with
-  # 1/square root of the time step; 1/sqrt(10k) = 0.01 and 1/sqrt(100k) ~ 0.0032
-  # so by bumping it up 3x we about "go back" from 100k steps to 10k, which is
-  # approximately as much as "going back 1 epoch" would be in default schedule.
-  # In your experiments, you may want to optimize this rate to your schedule.
-  if epoch > 0: learning_rate *= 3
+  if epoch > 0: learning_rate *= hparams.learning_rate_bump
   with temporary_flags({
       "data_dir": data_dir,
       "output_dir": output_dir,
@@ -597,6 +592,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       return 0.0
 
     # Collect data from the real environment.
+    generation_mean_reward = 0
     if not is_final_epoch:
       log("Generating real environment data.")
       generation_mean_reward = generate_real_env_data(
@@ -618,7 +614,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     assert mean_reward is not None
     model_reward_accuracy_summary.value[0].simple_value = model_reward_accuracy
     mean_reward_summary.value[0].simple_value = mean_reward
-    mean_reward_gen_summary.value[0].simple_value = generation_mean_reward
+    mean_reward_gen_summary.value[0].simple_value = int(generation_mean_reward)
     eval_metrics_writer.add_summary(model_reward_accuracy_summary, epoch)
     eval_metrics_writer.add_summary(mean_reward_summary, epoch)
     eval_metrics_writer.add_summary(mean_reward_gen_summary, epoch)
@@ -692,6 +688,13 @@ def rlmb_base():
       # Resizing.
       resize_height_factor=2,
       resize_width_factor=2,
+      # Bump learning rate after first epoch by 3x.
+      # We picked 3x because our default learning rate schedule decreases with
+      # 1/square root of step; 1/sqrt(10k) = 0.01 and 1/sqrt(100k) ~ 0.0032
+      # so by bumping it up 3x we about "go back" from 100k steps to 10k, which
+      # is approximately as much as "going back 1 epoch" would be.
+      # In your experiments, you want to optimize this rate to your schedule.
+      learning_rate_bump=3.0,
 
       gather_ppo_real_env_data=True,
       real_ppo_epochs_num=0,
@@ -789,6 +792,7 @@ def rlmb_base_stochastic_discrete():
 def rlmb_base_sv2p():
   """Base setting with sv2p as world model."""
   hparams = rlmb_base()
+  hparams.learning_rate_bump = 1.0
   hparams.generative_model = "next_frame_sv2p"
   hparams.generative_model_params = "next_frame_sv2p_atari"
   return hparams
@@ -827,6 +831,34 @@ def rlmb_base_sampling():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_base_sampling_noresize():
+  hparams = rlmb_base_sampling()
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_flippy60():
+  """Schedule with a lot of epochs (slow)."""
+  hparams = rlmb_base_sampling()
+  hparams.epochs = 60
+  hparams.ppo_epochs_num = 500
+  hparams.model_train_steps = 10000
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_flippy30():
+  """Schedule with a lot of epochs (slow)."""
+  hparams = rlmb_base_sampling()
+  hparams.epochs = 30
+  hparams.ppo_epochs_num = 1000
+  hparams.model_train_steps = 15000
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_medium():
   """Small set for larger testing."""

From 282c51ef216cb22537a694b70597ae2ab4c8d10a Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 21 Sep 2018 14:34:57 -0700
Subject: [PATCH 0895/2720] fixing internal_loss for softmax in a very dirty
 way.

PiperOrigin-RevId: 214045654
---
 tensor2tensor/models/video/sv2p.py        | 36 +++++++++++++++++------
 tensor2tensor/models/video/sv2p_params.py |  6 ++--
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index d9428a384..d6e958f0d 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -611,16 +611,34 @@ def body(self, features):
     predictions = common_video.swap_time_and_batch_axes(predictions)
     reward_pred = common_video.swap_time_and_batch_axes(reward_pred)
 
-    if hparams.internal_loss:
+    if self.is_training and hparams.internal_loss:
       # add the MSE loss for input frames as well.
-      # we are assuming the modality is L2. otherwise the loss would be
-      # incosistent across the frames.
-      if self._target_modality != "VideoModalityL2Raw":
-        raise ValueError("internal loss only works with L2.")
-      recon_loss = tf.losses.mean_squared_error(
-          all_frames[1:hparams.video_num_input_frames+1],
-          gen_images[:hparams.video_num_input_frames])
-      tf.summary.scalar("mse_extra", recon_loss)
+      extra_gts = all_frames[1:hparams.video_num_input_frames+1]
+      extra_gts = common_video.swap_time_and_batch_axes(extra_gts)
+      extra_pds = gen_images[:hparams.video_num_input_frames]
+      extra_pds = common_video.swap_time_and_batch_axes(extra_pds)
+      if self._target_modality == "VideoModalityL2Raw":
+        recon_loss = tf.losses.mean_squared_error(extra_gts, extra_pds)
+      elif self._target_modality == "VideoModality":
+        shape = common_layers.shape_list(extra_pds)
+        updated_shape = shape[:-1] + [3, 256]
+        extra_pds = tf.reshape(extra_pds, updated_shape)
+        # Merge time and batch
+        logits = tf.reshape(extra_pds, [-1] + updated_shape[2:])
+        targets_shape = common_layers.shape_list(features["targets_raw"])
+        targets = tf.reshape(features["targets_raw"], [-1] + targets_shape[2:])
+        mod = self.hparams.problem_hparams.target_modality["targets"]
+        numerator, denominator = common_layers.padded_cross_entropy(
+            logits,
+            targets,
+            hparams.label_smoothing,
+            cutoff=getattr(hparams, "video_modality_loss_cutoff", 0.01),
+            weights_fn=mod.targets_weights_fn)
+        recon_loss = numerator / denominator
+      else:
+        raise ValueError("internal loss only supports specific modalities.")
+
+      tf.summary.scalar("recon_extra", recon_loss)
       extra_loss += recon_loss
 
     return_targets = predictions
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index c1929ebe1..67dd69c80 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -75,7 +75,7 @@ def next_frame_sv2p_atari_softmax():
   hparams = next_frame_sv2p_atari()
   hparams.target_modality = "video"
   hparams.input_modalities = "inputs:video"
-  hparams.internal_loss = False
+  hparams.internal_loss = True
   return hparams
 
 
@@ -98,13 +98,13 @@ def next_frame_sv2p_atari_softmax_deterministic():
 @registry.register_hparams
 def next_frame_sv2p_tiny():
   """Tiny SV2P model."""
-  hparams = next_frame_sv2p_atari()
+  hparams = next_frame_sv2p_atari_softmax()
   hparams.batch_size = 2
   hparams.tiny_mode = True
   hparams.num_masks = 1
   hparams.video_modality_loss_cutoff = 0.4
   hparams.video_num_input_frames = 4
-  hparams.video_num_target_frames = 1
+  hparams.video_num_target_frames = 4
   return hparams
 
 

From 5bfbdc247510a505ad0f168691bc79f64a7033b6 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 21 Sep 2018 15:38:58 -0700
Subject: [PATCH 0896/2720] Internal change

PiperOrigin-RevId: 214055421
---
 tensor2tensor/models/research/glow_ops.py     | 56 ++++++++++++++-----
 .../models/research/glow_ops_test.py          |  4 +-
 2 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 691528b2f..dd6098f4e 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -343,17 +343,15 @@ def conv2d(name, x, output_channels, filter_size=None, stride=None,
 
 
 @add_arg_scope
-def nn(name, x, mid_channels, output_channels):
-  """3-layer conv2d.
+def conv_block(name, x, mid_channels):
+  """2 layer conv block used in the affine coupling layer.
 
   Args:
-    name:
-    x:
-    mid_channels: Number of output channels of the first layer.
-    output_channels: Number of output channels.
-
+    name: variable scope.
+    x: 4-D Tensor: (batch_size, height, width, channels).
+    mid_channels: Output channels of the second layer.
   Returns:
-    output:
+    x: 4-D Tensor: Output activations.
   """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
 
@@ -368,6 +366,25 @@ def nn(name, x, mid_channels, output_channels):
     x = conv2d("1_2", x, output_channels=mid_channels, filter_size=[1, 1],
                stride=[1, 1])
     x = tf.nn.relu(x)
+    return x
+
+
+@add_arg_scope
+def affine_coupling_network(name, x, mid_channels, output_channels):
+  """3-layer conv2d.
+
+  Args:
+    name:
+    x:
+    mid_channels: Number of output channels of the first layer.
+    output_channels: Number of output channels.
+
+  Returns:
+    output:
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+
+    x = conv_block("conv_block", x, mid_channels=mid_channels)
 
     # Final layer.
     x = conv2d("zeros", x, filter_size=[3, 3], stride=[1, 1],
@@ -399,7 +416,8 @@ def affine_coupling(name, x, mid_channels=512, reverse=False):
     # Else:
     # z2 = (x2 / scale) - shift
     z1 = x1
-    log_scale_and_shift = nn("nn", x1, mid_channels, x_shape[-1])
+    log_scale_and_shift = affine_coupling_network(
+        "nn", x1, mid_channels, x_shape[-1])
     shift = log_scale_and_shift[:, :, :, 0::2]
     scale = tf.nn.sigmoid(log_scale_and_shift[:, :, :, 1::2] + 2.0)
     if not reverse:
@@ -452,7 +470,8 @@ def squeeze(name, x, factor=2, reverse=True):
 
 
 @add_arg_scope
-def tensor_to_dist(name, x, output_channels=None, architecture="single_conv"):
+def tensor_to_dist(name, x, output_channels=None, architecture="single_conv",
+                   depth=1, pre_output_channels=512):
   """Map x to the mean and log-scale of a Gaussian.
 
   Args:
@@ -461,6 +480,8 @@ def tensor_to_dist(name, x, output_channels=None, architecture="single_conv"):
     output_channels: int, number of output channels of the mean.
                      if not provided, set it to be the output channels of x.
     architecture: "single_conv" or "glow_nn"
+    depth: depth of architecture mapping to the mean and std.
+    pre_output_channels: output channels before the final (mean, std) mapping.
   Returns:
     dist: instance of tf.distributions.Normal
   Raises:
@@ -474,8 +495,15 @@ def tensor_to_dist(name, x, output_channels=None, architecture="single_conv"):
       mean_log_scale = conv2d("conv2d", x, output_channels=2*output_channels,
                               conv_init="zeros", apply_actnorm=False)
     elif architecture == "glow_nn":
-      mean_log_scale = nn("conv2d", x, mid_channels=512,
-                          output_channels=2*output_channels)
+      mean_log_scale = x
+      for layer in range(1, depth + 1):
+        mid_channels = pre_output_channels // 2**(depth - layer)
+        mean_log_scale = conv_block("glow_nn_%d" % layer, mean_log_scale,
+                                    mid_channels=mid_channels)
+      mean_log_scale = conv2d("glow_nn_zeros", mean_log_scale,
+                              filter_size=[3, 3], stride=[1, 1],
+                              output_channels=2*output_channels,
+                              apply_actnorm=False, conv_init="zeros")
     else:
       raise ValueError("expected architecture to be single_conv or glow_nn "
                        "got %s" % architecture)
@@ -540,7 +568,9 @@ def compute_prior(name, z, latent, hparams):
         latent_stack = tf.concat([prior_dist.loc] + latent, axis=-1)
         prior_dist = tensor_to_dist(
             "latent_stack", latent_stack, output_channels=output_channels,
-            architecture=hparams.latent_architecture)
+            architecture=hparams.latent_architecture,
+            depth=hparams.latent_encoder_depth,
+            pre_output_channels=hparams.latent_pre_output_channels)
         latent_skip = hparams.get("latent_skip", False)
         if latent_skip:
           prior_dist = tf.distributions.Normal(
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index abd9dce05..696b6019e 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -109,11 +109,11 @@ def test_conv2d(self):
         # test shape in case apply_actnorm is set to False,
         self.assertEqual(zeros_np.shape, (16, 5, 5, 64))
 
-  def test_nn(self):
+  def test_affine_coupling_network(self):
     """Test output shape."""
     with tf.Graph().as_default():
       x = 10.0 * tf.random_uniform(shape=(16, 5, 5, 32))
-      nn = glow_ops.nn("nn", x, 512, 64)
+      nn = glow_ops.affine_coupling_network("nn", x, 512, 64)
 
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())

From 3a3b0551ac0c36f64751aa665fa673694f6bc7ad Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 21 Sep 2018 15:45:48 -0700
Subject: [PATCH 0897/2720] internal

PiperOrigin-RevId: 214056409
---
 tensor2tensor/rl/OWNERS | 3 +++
 1 file changed, 3 insertions(+)
 create mode 100644 tensor2tensor/rl/OWNERS

diff --git a/tensor2tensor/rl/OWNERS b/tensor2tensor/rl/OWNERS
new file mode 100644
index 000000000..cfa1f4f74
--- /dev/null
+++ b/tensor2tensor/rl/OWNERS
@@ -0,0 +1,3 @@
+blazej
+dumitru
+mbz

From 2964792f0e85ec43f43471d81cc4c0852f8eb426 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 21 Sep 2018 15:49:47 -0700
Subject: [PATCH 0898/2720] making training offset configurable

PiperOrigin-RevId: 214056996
---
 tensor2tensor/models/video/basic_stochastic.py | 6 +++---
 tensor2tensor/rl/trainer_model_based.py        | 5 ++++-
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 4c2d27acc..9f1a2c665 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -121,15 +121,15 @@ def next_frame_basic_stochastic():
   hparams.stochastic_model = True
   hparams.add_hparam("latent_channels", 1)
   hparams.add_hparam("latent_std_min", -5.0)
-  hparams.add_hparam("num_iterations_1st_stage", 25000)
-  hparams.add_hparam("num_iterations_2nd_stage", 25000)
+  hparams.add_hparam("num_iterations_1st_stage", 15000)
+  hparams.add_hparam("num_iterations_2nd_stage", 15000)
   hparams.add_hparam("latent_loss_multiplier", 1e-3)
   hparams.add_hparam("latent_loss_multiplier_dynamic", False)
   hparams.add_hparam("latent_loss_multiplier_alpha", 1e-5)
   hparams.add_hparam("latent_loss_multiplier_epsilon", 1.0)
   hparams.add_hparam("latent_loss_multiplier_schedule", "constant")
   hparams.add_hparam("latent_num_frames", 0)  # 0 means use all frames.
-  hparams.add_hparam("anneal_end", 100000)
+  hparams.add_hparam("anneal_end", 50000)
   hparams.add_hparam("information_capacity", 0.0)
   return hparams
 
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 5683d3684..d3dc0cd7b 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -320,7 +320,8 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
 
 def train_world_model(problem_name, data_dir, output_dir, hparams, epoch):
   """Train the world model on problem_name."""
-  train_steps = hparams.model_train_steps * (epoch + 2)
+  train_steps = hparams.model_train_steps * (
+      epoch + hparams.inital_epoch_train_steps_multiplier)
   model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
   learning_rate = model_hparams.learning_rate_constant
   if epoch > 0: learning_rate *= hparams.learning_rate_bump
@@ -667,6 +668,7 @@ def rlmb_base():
       ppo_params="ppo_pong_base",
       autoencoder_train_steps=0,
       model_train_steps=50000,
+      inital_epoch_train_steps_multiplier=2,
       simulated_env_generator_num_steps=2000,
       simulation_random_starts=True,  # Use random starts in PPO.
       # Flip the first random frame in PPO batch for the true beginning.
@@ -774,6 +776,7 @@ def rlmb_quick_sm():
 def rlmb_base_stochastic():
   """Base setting with a stochastic next-frame model."""
   hparams = rlmb_base()
+  hparams.inital_epoch_train_steps_multiplier = 5
   hparams.generative_model = "next_frame_basic_stochastic"
   hparams.generative_model_params = "next_frame_basic_stochastic"
   return hparams

From 39471ed3a0a960b7d43353b51870db3cd1fbbcc0 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Sat, 22 Sep 2018 01:10:24 +0200
Subject: [PATCH 0899/2720] Fix the autoencoder RL experiment (#1090)

---
 tensor2tensor/data_generators/gym_problems.py   |  4 ----
 tensor2tensor/rl/trainer_model_based_ae_test.py | 13 ++++++-------
 2 files changed, 6 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index dc07826b2..259cd167b 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -510,10 +510,6 @@ def autoencoder_factor(self):
     hparams = autoencoders.autoencoder_discrete_pong()
     return 2**hparams.num_hidden_layers
 
-  @property
-  def num_channels(self):
-    return 24
-
   @property
   def frame_height(self):
     height = self.env.observation_space.shape[0]
diff --git a/tensor2tensor/rl/trainer_model_based_ae_test.py b/tensor2tensor/rl/trainer_model_based_ae_test.py
index e44b477c8..a7b15a62d 100644
--- a/tensor2tensor/rl/trainer_model_based_ae_test.py
+++ b/tensor2tensor/rl/trainer_model_based_ae_test.py
@@ -17,7 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
-# from tensor2tensor.rl import trainer_model_based
+from tensor2tensor.rl import trainer_model_based
 
 import tensorflow as tf
 
@@ -27,12 +27,11 @@
 class ModelRLExperimentTestAe(tf.test.TestCase):
 
   def test_ae(self):
-    # TODO(lukaszkaiser): re-enable this test.
-    # FLAGS.output_dir = tf.test.get_temp_dir()
-    # FLAGS.loop_hparams_set = "rlmb_ae_tiny"
-    # FLAGS.schedule = "train"  # skip evaluation for world model training
-    # trainer_model_based.main(None)
-    assert True
+    FLAGS.output_dir = tf.test.get_temp_dir()
+    FLAGS.loop_hparams_set = "rlmb_ae_tiny"
+    FLAGS.schedule = "train"  # skip evaluation for world model training
+    trainer_model_based.main(None)
+
 
 if __name__ == "__main__":
   tf.test.main()

From c7b8c5a0a0a616cd1f97c8f4ed59e2b72f5b7337 Mon Sep 17 00:00:00 2001
From: Shahzeb K <shahzeb001@gmail.com>
Date: Fri, 21 Sep 2018 16:11:57 -0700
Subject: [PATCH 0900/2720] Avoid a ascii codec error in CNN/DailyMail. (#1086)

When a basic `t2t-trainer` command is run with the `--problem=summarize_cnn_dailymail32k` command, the following error occurs:

```
UnicodeEncodeError: 'ascii' codec can't encode character '\xbb' in position 1426: ordinal not in range(128)
```

This minor fix resolves this issue.
---
 tensor2tensor/data_generators/cnn_dailymail.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index 79e7d011c..f6c21df77 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -182,9 +182,9 @@ def write_raw_text_to_files(all_files, urls_path, dataset_split, tmp_dir):
   """Write text to files."""
 
   def write_to_file(all_files, urls_path, tmp_dir, filename):
-    with io.open(os.path.join(tmp_dir, filename + ".source"), "w") as fstory:
+    with io.open(os.path.join(tmp_dir, filename + ".source"), "w", encoding="utf-8") as fstory:
       with io.open(os.path.join(tmp_dir, filename + ".target"),
-                   "w") as fsummary:
+                   "w", encoding="utf-8") as fsummary:
         for example in example_generator(all_files, urls_path, sum_token=True):
           story, summary = _story_summary_split(example)
           fstory.write(story + "\n")

From a5a4d27fff1c9dcfaf6a6aa7e59752e0291743ae Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 21 Sep 2018 16:10:42 -0700
Subject: [PATCH 0901/2720] internal merge of PR #1090

PiperOrigin-RevId: 214060019
---
 tensor2tensor/data_generators/cnn_dailymail.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index f6c21df77..79e7d011c 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -182,9 +182,9 @@ def write_raw_text_to_files(all_files, urls_path, dataset_split, tmp_dir):
   """Write text to files."""
 
   def write_to_file(all_files, urls_path, tmp_dir, filename):
-    with io.open(os.path.join(tmp_dir, filename + ".source"), "w", encoding="utf-8") as fstory:
+    with io.open(os.path.join(tmp_dir, filename + ".source"), "w") as fstory:
       with io.open(os.path.join(tmp_dir, filename + ".target"),
-                   "w", encoding="utf-8") as fsummary:
+                   "w") as fsummary:
         for example in example_generator(all_files, urls_path, sum_token=True):
           story, summary = _story_summary_split(example)
           fstory.write(story + "\n")

From ac6ff9f92c2a8b39a0023eb8385c781219cad686 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 21 Sep 2018 16:12:18 -0700
Subject: [PATCH 0902/2720] internal merge of PR #1086

PiperOrigin-RevId: 214060201
---
 tensor2tensor/data_generators/cnn_dailymail.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index 79e7d011c..01f1642d6 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -182,9 +182,13 @@ def write_raw_text_to_files(all_files, urls_path, dataset_split, tmp_dir):
   """Write text to files."""
 
   def write_to_file(all_files, urls_path, tmp_dir, filename):
-    with io.open(os.path.join(tmp_dir, filename + ".source"), "w") as fstory:
-      with io.open(os.path.join(tmp_dir, filename + ".target"),
-                   "w") as fsummary:
+    """Write text to files."""
+    with io.open(
+        os.path.join(tmp_dir, filename + ".source"), "w",
+        encoding="utf-8") as fstory:
+      with io.open(
+          os.path.join(tmp_dir, filename + ".target"), "w",
+          encoding="utf-8") as fsummary:
         for example in example_generator(all_files, urls_path, sum_token=True):
           story, summary = _story_summary_split(example)
           fstory.write(story + "\n")

From f1d68df204f6f21a030a772e97b2465941007f2e Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Sat, 22 Sep 2018 01:13:14 +0200
Subject: [PATCH 0903/2720] Calculate world model reward accuracy for multiple
 rollout lengths (#1081)

---
 tensor2tensor/data_generators/gym_problems.py | 24 +++++++----
 tensor2tensor/rl/trainer_model_based.py       | 40 ++++++++++++++-----
 2 files changed, 47 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 259cd167b..3b026a589 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import json
 import math
 import os
@@ -527,14 +528,16 @@ class RewardPerSequenceStatistics(BasicStatistics):
   the correctness of rewards per sequence metric
   """
 
-  def __init__(self):
+  def __init__(self, rollout_fractions):
     super(RewardPerSequenceStatistics, self).__init__()
 
     # data to calculate
     # correctness of rewards per sequence metric
     self.episode_sim_reward = 0.0
     self.episode_real_reward = 0.0
-    self.successful_episode_reward_predictions = 0
+    self.successful_episode_reward_predictions = collections.OrderedDict([
+        (frac, 0) for frac in rollout_fractions
+    ])
     self.report_reward_statistics_every = 10
     # auxiliary objects
     self.real_obs = None
@@ -545,7 +548,7 @@ def to_dict(self):
     keys_and_types = [
         ("episode_sim_reward", float),
         ("episode_real_reward", float),
-        ("successful_episode_reward_predictions", int),
+        ("successful_episode_reward_predictions", collections.OrderedDict),
         ("report_reward_statistics_every", int),
     ]
     additional = dict([(k, t(getattr(self, k))) for k, t in keys_and_types])
@@ -640,7 +643,7 @@ def __init__(self, *args, **kwargs):
     super(GymSimulatedDiscreteProblemForWorldModelEval, self).__init__(
         *args, **kwargs
     )
-    self.statistics = RewardPerSequenceStatistics()
+    self.settable_rollout_fractions = [1]
 
   def get_environment_spec(self):
     env_spec = super(
@@ -650,6 +653,10 @@ def get_environment_spec(self):
     return env_spec
 
   def _setup(self, data_dir):
+    self.statistics = RewardPerSequenceStatistics(
+        self.settable_rollout_fractions
+    )
+
     trajectory_length = self.num_testing_steps
     if self.num_steps < 1200:
       # Decrease the trajectory length for tiny experiments, otherwise we don't
@@ -718,11 +725,14 @@ def collect_statistics_and_generate_debug_image(self, index,
                 "The collect memory should be set in force_beginning_resets "
                 "mode for the code below to work properly.")
 
-    if (index+1) % self._internal_memory_size == 0:
+    index_in_rollout = index % self._internal_memory_size + 1
 
-      if stat.episode_sim_reward == stat.episode_real_reward:
-        stat.successful_episode_reward_predictions += 1
+    if stat.episode_sim_reward == stat.episode_real_reward:
+      for frac in stat.successful_episode_reward_predictions:
+        if index_in_rollout == int(self._internal_memory_size * frac):
+          stat.successful_episode_reward_predictions[frac] += 1
 
+    if index_in_rollout == self._internal_memory_size:
       stat.episode_sim_reward = 0.0
       stat.episode_real_reward = 0.0
       stat.number_of_dones += 1
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index d3dc0cd7b..5e2420982 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -299,6 +299,8 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
   gym_simulated_problem = registry.problem(simulated_problem_name)
   sim_steps = hparams.simulated_env_generator_num_steps
   gym_simulated_problem.settable_num_steps = sim_steps
+  gym_simulated_problem.settable_rollout_fractions = \
+      hparams.eval_rollout_fractions
   with temporary_flags({
       "problem": problem_name,
       "model": hparams.generative_model,
@@ -308,9 +310,13 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
   }):
     gym_simulated_problem.generate_data(epoch_data_dir, tmp_dir)
   n = max(1., gym_simulated_problem.statistics.number_of_dones)
-  model_reward_accuracy = (
-      gym_simulated_problem.statistics.successful_episode_reward_predictions
-      / float(n))
+  model_reward_accuracy = [
+      (frac, score / float(n))
+      for (frac, score) in six.iteritems(
+          gym_simulated_problem.statistics.
+          successful_episode_reward_predictions
+      )
+  ]
   old_path = os.path.join(epoch_data_dir, "debug_frames_sim")
   new_path = os.path.join(epoch_data_dir, "debug_frames_sim_eval")
   if not tf.gfile.Exists(new_path):
@@ -523,8 +529,11 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
                                         "eval_metrics_event_dir")
   eval_metrics_writer = tf.summary.FileWriter(eval_metrics_event_dir)
   model_reward_accuracy_summary = tf.Summary()
-  model_reward_accuracy_summary.value.add(tag="model_reward_accuracy",
-                                          simple_value=None)
+  for frac in hparams.eval_rollout_fractions:
+    model_reward_accuracy_summary.value.add(
+        tag="model_reward_accuracy_{}".format(frac),
+        simple_value=None
+    )
   mean_reward_summary = tf.Summary()
   mean_reward_summary.value.add(tag="mean_reward",
                                 simple_value=None)
@@ -562,14 +571,17 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
                       directories["world_model"], hparams, epoch)
 
     # Evaluate world model
-    model_reward_accuracy = 0.
+    model_reward_accuracy = []
     if hparams.eval_world_model:
       log("Evaluating world model")
       model_reward_accuracy = evaluate_world_model(
           world_model_eval_problem_name, world_model_problem, hparams,
           directories["world_model"],
           epoch_data_dir, directories["tmp"])
-      log("World model reward accuracy: %.4f", model_reward_accuracy)
+      log(
+          "World model reward accuracy per rollout fraction: %s",
+          model_reward_accuracy
+      )
 
     # Train PPO
     log("Training PPO in simulated environment.")
@@ -613,7 +625,10 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     # Summarize metrics.
     assert model_reward_accuracy is not None
     assert mean_reward is not None
-    model_reward_accuracy_summary.value[0].simple_value = model_reward_accuracy
+    for ((_, accuracy), summary_value) in zip(
+        model_reward_accuracy, model_reward_accuracy_summary.value
+    ):
+      summary_value.simple_value = accuracy
     mean_reward_summary.value[0].simple_value = mean_reward
     mean_reward_gen_summary.value[0].simple_value = int(generation_mean_reward)
     eval_metrics_writer.add_summary(model_reward_accuracy_summary, epoch)
@@ -622,8 +637,11 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     eval_metrics_writer.flush()
 
     # Report metrics
-    eval_metrics = {"model_reward_accuracy": model_reward_accuracy,
-                    "mean_reward": mean_reward}
+    eval_metrics = {"mean_reward": mean_reward}
+    eval_metrics.update({
+        "model_reward_accuracy_{}".format(frac): accuracy
+        for (frac, accuracy) in model_reward_accuracy
+    })
     epoch_metrics.append(eval_metrics)
     log("Eval metrics: %s", str(eval_metrics))
     if report_fn:
@@ -711,6 +729,8 @@ def rlmb_base():
       # Whether to evaluate the world model in each iteration of the loop to get
       # the model_reward_accuracy metric.
       eval_world_model=True,
+      # Rollout fractions to report reward_accuracy on.
+      eval_rollout_fractions=[0.25, 0.5, 1],
       stop_loop_early=False,  # To speed-up tests.
   )
 

From 362fff61350ebb067cf2508044cb01ae594472f9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 21 Sep 2018 16:13:41 -0700
Subject: [PATCH 0904/2720] internal

PiperOrigin-RevId: 214060377
---
 tensor2tensor/data_generators/gym_problems.py | 24 ++++-------
 tensor2tensor/mesh_tensorflow/OWNERS          |  2 -
 tensor2tensor/models/OWNERS                   |  6 ---
 tensor2tensor/rl/OWNERS                       |  3 --
 tensor2tensor/rl/trainer_model_based.py       | 40 +++++--------------
 5 files changed, 17 insertions(+), 58 deletions(-)
 delete mode 100644 tensor2tensor/mesh_tensorflow/OWNERS
 delete mode 100644 tensor2tensor/models/OWNERS
 delete mode 100644 tensor2tensor/rl/OWNERS

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 3b026a589..259cd167b 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -18,7 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import collections
 import json
 import math
 import os
@@ -528,16 +527,14 @@ class RewardPerSequenceStatistics(BasicStatistics):
   the correctness of rewards per sequence metric
   """
 
-  def __init__(self, rollout_fractions):
+  def __init__(self):
     super(RewardPerSequenceStatistics, self).__init__()
 
     # data to calculate
     # correctness of rewards per sequence metric
     self.episode_sim_reward = 0.0
     self.episode_real_reward = 0.0
-    self.successful_episode_reward_predictions = collections.OrderedDict([
-        (frac, 0) for frac in rollout_fractions
-    ])
+    self.successful_episode_reward_predictions = 0
     self.report_reward_statistics_every = 10
     # auxiliary objects
     self.real_obs = None
@@ -548,7 +545,7 @@ def to_dict(self):
     keys_and_types = [
         ("episode_sim_reward", float),
         ("episode_real_reward", float),
-        ("successful_episode_reward_predictions", collections.OrderedDict),
+        ("successful_episode_reward_predictions", int),
         ("report_reward_statistics_every", int),
     ]
     additional = dict([(k, t(getattr(self, k))) for k, t in keys_and_types])
@@ -643,7 +640,7 @@ def __init__(self, *args, **kwargs):
     super(GymSimulatedDiscreteProblemForWorldModelEval, self).__init__(
         *args, **kwargs
     )
-    self.settable_rollout_fractions = [1]
+    self.statistics = RewardPerSequenceStatistics()
 
   def get_environment_spec(self):
     env_spec = super(
@@ -653,10 +650,6 @@ def get_environment_spec(self):
     return env_spec
 
   def _setup(self, data_dir):
-    self.statistics = RewardPerSequenceStatistics(
-        self.settable_rollout_fractions
-    )
-
     trajectory_length = self.num_testing_steps
     if self.num_steps < 1200:
       # Decrease the trajectory length for tiny experiments, otherwise we don't
@@ -725,14 +718,11 @@ def collect_statistics_and_generate_debug_image(self, index,
                 "The collect memory should be set in force_beginning_resets "
                 "mode for the code below to work properly.")
 
-    index_in_rollout = index % self._internal_memory_size + 1
+    if (index+1) % self._internal_memory_size == 0:
 
-    if stat.episode_sim_reward == stat.episode_real_reward:
-      for frac in stat.successful_episode_reward_predictions:
-        if index_in_rollout == int(self._internal_memory_size * frac):
-          stat.successful_episode_reward_predictions[frac] += 1
+      if stat.episode_sim_reward == stat.episode_real_reward:
+        stat.successful_episode_reward_predictions += 1
 
-    if index_in_rollout == self._internal_memory_size:
       stat.episode_sim_reward = 0.0
       stat.episode_real_reward = 0.0
       stat.number_of_dones += 1
diff --git a/tensor2tensor/mesh_tensorflow/OWNERS b/tensor2tensor/mesh_tensorflow/OWNERS
deleted file mode 100644
index 681d09a90..000000000
--- a/tensor2tensor/mesh_tensorflow/OWNERS
+++ /dev/null
@@ -1,2 +0,0 @@
-nikip
-noam
diff --git a/tensor2tensor/models/OWNERS b/tensor2tensor/models/OWNERS
deleted file mode 100644
index 4c18f708e..000000000
--- a/tensor2tensor/models/OWNERS
+++ /dev/null
@@ -1,6 +0,0 @@
-avaswani
-dumitru
-mbz
-nikip
-noam
-uszkoreit
diff --git a/tensor2tensor/rl/OWNERS b/tensor2tensor/rl/OWNERS
deleted file mode 100644
index cfa1f4f74..000000000
--- a/tensor2tensor/rl/OWNERS
+++ /dev/null
@@ -1,3 +0,0 @@
-blazej
-dumitru
-mbz
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 5e2420982..d3dc0cd7b 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -299,8 +299,6 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
   gym_simulated_problem = registry.problem(simulated_problem_name)
   sim_steps = hparams.simulated_env_generator_num_steps
   gym_simulated_problem.settable_num_steps = sim_steps
-  gym_simulated_problem.settable_rollout_fractions = \
-      hparams.eval_rollout_fractions
   with temporary_flags({
       "problem": problem_name,
       "model": hparams.generative_model,
@@ -310,13 +308,9 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
   }):
     gym_simulated_problem.generate_data(epoch_data_dir, tmp_dir)
   n = max(1., gym_simulated_problem.statistics.number_of_dones)
-  model_reward_accuracy = [
-      (frac, score / float(n))
-      for (frac, score) in six.iteritems(
-          gym_simulated_problem.statistics.
-          successful_episode_reward_predictions
-      )
-  ]
+  model_reward_accuracy = (
+      gym_simulated_problem.statistics.successful_episode_reward_predictions
+      / float(n))
   old_path = os.path.join(epoch_data_dir, "debug_frames_sim")
   new_path = os.path.join(epoch_data_dir, "debug_frames_sim_eval")
   if not tf.gfile.Exists(new_path):
@@ -529,11 +523,8 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
                                         "eval_metrics_event_dir")
   eval_metrics_writer = tf.summary.FileWriter(eval_metrics_event_dir)
   model_reward_accuracy_summary = tf.Summary()
-  for frac in hparams.eval_rollout_fractions:
-    model_reward_accuracy_summary.value.add(
-        tag="model_reward_accuracy_{}".format(frac),
-        simple_value=None
-    )
+  model_reward_accuracy_summary.value.add(tag="model_reward_accuracy",
+                                          simple_value=None)
   mean_reward_summary = tf.Summary()
   mean_reward_summary.value.add(tag="mean_reward",
                                 simple_value=None)
@@ -571,17 +562,14 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
                       directories["world_model"], hparams, epoch)
 
     # Evaluate world model
-    model_reward_accuracy = []
+    model_reward_accuracy = 0.
     if hparams.eval_world_model:
       log("Evaluating world model")
       model_reward_accuracy = evaluate_world_model(
           world_model_eval_problem_name, world_model_problem, hparams,
           directories["world_model"],
           epoch_data_dir, directories["tmp"])
-      log(
-          "World model reward accuracy per rollout fraction: %s",
-          model_reward_accuracy
-      )
+      log("World model reward accuracy: %.4f", model_reward_accuracy)
 
     # Train PPO
     log("Training PPO in simulated environment.")
@@ -625,10 +613,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     # Summarize metrics.
     assert model_reward_accuracy is not None
     assert mean_reward is not None
-    for ((_, accuracy), summary_value) in zip(
-        model_reward_accuracy, model_reward_accuracy_summary.value
-    ):
-      summary_value.simple_value = accuracy
+    model_reward_accuracy_summary.value[0].simple_value = model_reward_accuracy
     mean_reward_summary.value[0].simple_value = mean_reward
     mean_reward_gen_summary.value[0].simple_value = int(generation_mean_reward)
     eval_metrics_writer.add_summary(model_reward_accuracy_summary, epoch)
@@ -637,11 +622,8 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     eval_metrics_writer.flush()
 
     # Report metrics
-    eval_metrics = {"mean_reward": mean_reward}
-    eval_metrics.update({
-        "model_reward_accuracy_{}".format(frac): accuracy
-        for (frac, accuracy) in model_reward_accuracy
-    })
+    eval_metrics = {"model_reward_accuracy": model_reward_accuracy,
+                    "mean_reward": mean_reward}
     epoch_metrics.append(eval_metrics)
     log("Eval metrics: %s", str(eval_metrics))
     if report_fn:
@@ -729,8 +711,6 @@ def rlmb_base():
       # Whether to evaluate the world model in each iteration of the loop to get
       # the model_reward_accuracy metric.
       eval_world_model=True,
-      # Rollout fractions to report reward_accuracy on.
-      eval_rollout_fractions=[0.25, 0.5, 1],
       stop_loop_early=False,  # To speed-up tests.
   )
 

From 2c2ea3ad0cd278d52212867b51d16e6faf5bb3b3 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 21 Sep 2018 16:28:57 -0700
Subject: [PATCH 0905/2720] internal merge of PR #1081

PiperOrigin-RevId: 214062339
---
 tensor2tensor/data_generators/gym_problems.py | 26 +++++++-----
 tensor2tensor/rl/trainer_model_based.py       | 40 ++++++++++++++-----
 2 files changed, 47 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 259cd167b..9248d8793 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import json
 import math
 import os
@@ -41,8 +42,6 @@
 FLAGS = flags.FLAGS
 
 
-
-
 flags.DEFINE_string("agent_policy_path", None, "File with model for agent.")
 
 flags.DEFINE_string("autoencoder_path", None,
@@ -527,14 +526,16 @@ class RewardPerSequenceStatistics(BasicStatistics):
   the correctness of rewards per sequence metric
   """
 
-  def __init__(self):
+  def __init__(self, rollout_fractions):
     super(RewardPerSequenceStatistics, self).__init__()
 
     # data to calculate
     # correctness of rewards per sequence metric
     self.episode_sim_reward = 0.0
     self.episode_real_reward = 0.0
-    self.successful_episode_reward_predictions = 0
+    self.successful_episode_reward_predictions = collections.OrderedDict([
+        (frac, 0) for frac in rollout_fractions
+    ])
     self.report_reward_statistics_every = 10
     # auxiliary objects
     self.real_obs = None
@@ -545,7 +546,7 @@ def to_dict(self):
     keys_and_types = [
         ("episode_sim_reward", float),
         ("episode_real_reward", float),
-        ("successful_episode_reward_predictions", int),
+        ("successful_episode_reward_predictions", collections.OrderedDict),
         ("report_reward_statistics_every", int),
     ]
     additional = dict([(k, t(getattr(self, k))) for k, t in keys_and_types])
@@ -640,7 +641,7 @@ def __init__(self, *args, **kwargs):
     super(GymSimulatedDiscreteProblemForWorldModelEval, self).__init__(
         *args, **kwargs
     )
-    self.statistics = RewardPerSequenceStatistics()
+    self.settable_rollout_fractions = [1]
 
   def get_environment_spec(self):
     env_spec = super(
@@ -650,6 +651,10 @@ def get_environment_spec(self):
     return env_spec
 
   def _setup(self, data_dir):
+    self.statistics = RewardPerSequenceStatistics(
+        self.settable_rollout_fractions
+    )
+
     trajectory_length = self.num_testing_steps
     if self.num_steps < 1200:
       # Decrease the trajectory length for tiny experiments, otherwise we don't
@@ -718,11 +723,14 @@ def collect_statistics_and_generate_debug_image(self, index,
                 "The collect memory should be set in force_beginning_resets "
                 "mode for the code below to work properly.")
 
-    if (index+1) % self._internal_memory_size == 0:
+    index_in_rollout = index % self._internal_memory_size + 1
 
-      if stat.episode_sim_reward == stat.episode_real_reward:
-        stat.successful_episode_reward_predictions += 1
+    if stat.episode_sim_reward == stat.episode_real_reward:
+      for frac in stat.successful_episode_reward_predictions:
+        if index_in_rollout == int(self._internal_memory_size * frac):
+          stat.successful_episode_reward_predictions[frac] += 1
 
+    if index_in_rollout == self._internal_memory_size:
       stat.episode_sim_reward = 0.0
       stat.episode_real_reward = 0.0
       stat.number_of_dones += 1
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index d3dc0cd7b..5e2420982 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -299,6 +299,8 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
   gym_simulated_problem = registry.problem(simulated_problem_name)
   sim_steps = hparams.simulated_env_generator_num_steps
   gym_simulated_problem.settable_num_steps = sim_steps
+  gym_simulated_problem.settable_rollout_fractions = \
+      hparams.eval_rollout_fractions
   with temporary_flags({
       "problem": problem_name,
       "model": hparams.generative_model,
@@ -308,9 +310,13 @@ def evaluate_world_model(simulated_problem_name, problem_name, hparams,
   }):
     gym_simulated_problem.generate_data(epoch_data_dir, tmp_dir)
   n = max(1., gym_simulated_problem.statistics.number_of_dones)
-  model_reward_accuracy = (
-      gym_simulated_problem.statistics.successful_episode_reward_predictions
-      / float(n))
+  model_reward_accuracy = [
+      (frac, score / float(n))
+      for (frac, score) in six.iteritems(
+          gym_simulated_problem.statistics.
+          successful_episode_reward_predictions
+      )
+  ]
   old_path = os.path.join(epoch_data_dir, "debug_frames_sim")
   new_path = os.path.join(epoch_data_dir, "debug_frames_sim_eval")
   if not tf.gfile.Exists(new_path):
@@ -523,8 +529,11 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
                                         "eval_metrics_event_dir")
   eval_metrics_writer = tf.summary.FileWriter(eval_metrics_event_dir)
   model_reward_accuracy_summary = tf.Summary()
-  model_reward_accuracy_summary.value.add(tag="model_reward_accuracy",
-                                          simple_value=None)
+  for frac in hparams.eval_rollout_fractions:
+    model_reward_accuracy_summary.value.add(
+        tag="model_reward_accuracy_{}".format(frac),
+        simple_value=None
+    )
   mean_reward_summary = tf.Summary()
   mean_reward_summary.value.add(tag="mean_reward",
                                 simple_value=None)
@@ -562,14 +571,17 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
                       directories["world_model"], hparams, epoch)
 
     # Evaluate world model
-    model_reward_accuracy = 0.
+    model_reward_accuracy = []
     if hparams.eval_world_model:
       log("Evaluating world model")
       model_reward_accuracy = evaluate_world_model(
           world_model_eval_problem_name, world_model_problem, hparams,
           directories["world_model"],
           epoch_data_dir, directories["tmp"])
-      log("World model reward accuracy: %.4f", model_reward_accuracy)
+      log(
+          "World model reward accuracy per rollout fraction: %s",
+          model_reward_accuracy
+      )
 
     # Train PPO
     log("Training PPO in simulated environment.")
@@ -613,7 +625,10 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     # Summarize metrics.
     assert model_reward_accuracy is not None
     assert mean_reward is not None
-    model_reward_accuracy_summary.value[0].simple_value = model_reward_accuracy
+    for ((_, accuracy), summary_value) in zip(
+        model_reward_accuracy, model_reward_accuracy_summary.value
+    ):
+      summary_value.simple_value = accuracy
     mean_reward_summary.value[0].simple_value = mean_reward
     mean_reward_gen_summary.value[0].simple_value = int(generation_mean_reward)
     eval_metrics_writer.add_summary(model_reward_accuracy_summary, epoch)
@@ -622,8 +637,11 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     eval_metrics_writer.flush()
 
     # Report metrics
-    eval_metrics = {"model_reward_accuracy": model_reward_accuracy,
-                    "mean_reward": mean_reward}
+    eval_metrics = {"mean_reward": mean_reward}
+    eval_metrics.update({
+        "model_reward_accuracy_{}".format(frac): accuracy
+        for (frac, accuracy) in model_reward_accuracy
+    })
     epoch_metrics.append(eval_metrics)
     log("Eval metrics: %s", str(eval_metrics))
     if report_fn:
@@ -711,6 +729,8 @@ def rlmb_base():
       # Whether to evaluate the world model in each iteration of the loop to get
       # the model_reward_accuracy metric.
       eval_world_model=True,
+      # Rollout fractions to report reward_accuracy on.
+      eval_rollout_fractions=[0.25, 0.5, 1],
       stop_loop_early=False,  # To speed-up tests.
   )
 

From 43a87032ed034218c04914ac801fda4f1df74a93 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 21 Sep 2018 16:44:57 -0700
Subject: [PATCH 0906/2720] fixing visualization with Softmax.

PiperOrigin-RevId: 214064621
---
 tensor2tensor/models/video/basic_deterministic.py |  3 ++-
 tensor2tensor/models/video/sv2p.py                | 10 ++++++++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 1be7701ca..cb0551f50 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -84,7 +84,8 @@ def get_sampled_frame(self, res_frame):
     frame_shape = common_layers.shape_list(res_frame)
     target_shape = frame_shape[:-1] + [self.hparams.problem.num_channels]
     sampled_frame = tf.reshape(res_frame, target_shape + [256])
-    sampled_frame = tf.to_float(tf.argmax(sampled_frame, axis=-1))
+    sampled_frame = tf.argmax(sampled_frame, axis=-1)
+    sampled_frame = tf.to_float(sampled_frame)
     return sampled_frame
 
   def body_single(self, features):
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index d6e958f0d..026c1637e 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -55,7 +55,13 @@ def concat_on_y_axis(x):
 
     frames_gd = common_video.swap_time_and_batch_axes(real_frames)
     frames_pd = common_video.swap_time_and_batch_axes(gen_frames)
-    frames_pd = self.get_sampled_frame(frames_pd)
+
+    if self.is_per_pixel_softmax:
+      frames_pd_shape = common_layers.shape_list(frames_pd)
+      frames_pd = tf.reshape(frames_pd, [-1, 256])
+      frames_pd = tf.to_float(tf.argmax(frames_pd, axis=-1))
+      frames_pd = tf.reshape(frames_pd, frames_pd_shape[:-1] + [3])
+
     frames_gd = concat_on_y_axis(frames_gd)
     frames_pd = concat_on_y_axis(frames_pd)
     side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
@@ -598,7 +604,7 @@ def body(self, features):
         gen_frames=gen_images)
 
     # Visualize predictions in Tensorboard
-    if self.is_training and not self.is_per_pixel_softmax:
+    if self.is_training:
       self.visualize_predictions(all_frames[1:], gen_images)
 
     # Ignore the predictions from the input frames.

From 715d895a83c1a89f0ad878271ced41bd405d8c6b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 21 Sep 2018 16:57:35 -0700
Subject: [PATCH 0907/2720] Fixed bug in hparams() introduced in cr/213330078.

PiperOrigin-RevId: 214066177
---
 tensor2tensor/data_generators/problem.py | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 28fc88d7b..2895539a8 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -497,16 +497,8 @@ def get_hparams(self, model_hparams=None):
       _copy_problem_hparams(hp)
 
     model_hparams = copy.copy(model_hparams)
-    if (hasattr(model_hparams, "shared_embedding_and_softmax_weights") and
-        model_hparams.shared_embedding_and_softmax_weights):
-      # If vocabularies differ, unset shared_embedding_and_softmax_weights.
-      input_vocab_size = hp.input_modality.get("inputs")[1]
-      target_vocab_size = hp.target_modality[1]
-      if input_vocab_size != target_vocab_size:
-        tf.logging.warn("Unsetting shared_embedding_and_softmax_weights.")
-        model_hparams.shared_embedding_and_softmax_weights = 0
-
     _create_modalities(hp, model_hparams)
+
     self._hparams = hp
     return self._hparams
 

From 9749391c9f5e31fac1aeda87ac1cc91419b8e36c Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 21 Sep 2018 18:19:10 -0700
Subject: [PATCH 0908/2720] make sv2p config flippy 30.

PiperOrigin-RevId: 214074304
---
 tensor2tensor/models/video/sv2p_params.py |  6 ++--
 tensor2tensor/rl/trainer_model_based.py   | 39 +++++++++++++++++++++++
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 67dd69c80..0d8aed81d 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -59,9 +59,9 @@ def next_frame_sv2p_atari():
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 4
   hparams.action_injection = "multiplicative"
-  hparams.num_iterations_1st_stage = 20000
-  hparams.num_iterations_2nd_stage = 20000
-  hparams.anneal_end = 80000
+  hparams.num_iterations_1st_stage = 12000
+  hparams.num_iterations_2nd_stage = 12000
+  hparams.anneal_end = 40000
   hparams.latent_loss_multiplier_schedule = "noisy_linear_cosine_decay"
   hparams.latent_loss_multiplier = 1e-3
   hparams.information_capacity = 0.0
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 5e2420982..0e4f8d31b 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -846,6 +846,45 @@ def rlmb_base_sv2p_deterministic_softmax():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_base_sv2p_flippy30():
+  """Base setting with sv2p as world model."""
+  hparams = rlmb_base()
+  hparams.epochs = 30
+  hparams.ppo_epochs_num = 1000
+  hparams.model_train_steps = 15000
+  hparams.learning_rate_bump = 1.0
+  hparams.inital_epoch_train_steps_multiplier = 5
+  hparams.generative_model = "next_frame_sv2p"
+  hparams.generative_model_params = "next_frame_sv2p_atari"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_softmax_flippy30():
+  """Base setting with sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p_flippy30()
+  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic_flippy30():
+  """Base setting with deterministic sv2p as world model."""
+  hparams = rlmb_base_sv2p_flippy30()
+  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic_softmax_flippy30():
+  """Base setting with deterministic sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p_softmax_flippy30()
+  hparams.generative_model_params = (
+      "next_frame_sv2p_atari_softmax_deterministic")
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_base_sampling():
   """Base setting with a stochastic next-frame model."""

From 5e19f44fa7cacee0cc27b9a3f0879e338ac81ae1 Mon Sep 17 00:00:00 2001
From: Yichao 'Peak' Ji <peakji@users.noreply.github.com>
Date: Sat, 22 Sep 2018 11:38:13 +0800
Subject: [PATCH 0909/2720] Avoid NaN while adding sinusoidal timing signals
 (#1055)

* Avoid NaN while adding sinusoidal timing signals

* Fix mismatched parentheses
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 74f45d465..e3daa1cc6 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -396,7 +396,7 @@ def get_timing_signal_1d(length,
   num_timescales = channels // 2
   log_timescale_increment = (
       math.log(float(max_timescale) / float(min_timescale)) /
-      (tf.to_float(num_timescales) - 1))
+      tf.maximum(tf.to_float(num_timescales) - 1, 1))
   inv_timescales = min_timescale * tf.exp(
       tf.to_float(tf.range(num_timescales)) * -log_timescale_increment)
   scaled_time = tf.expand_dims(position, 1) * tf.expand_dims(inv_timescales, 0)

From 398f498d77285ab51f5f3f9eb90b87c7fb5cc4f4 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Sat, 22 Sep 2018 06:56:36 +0200
Subject: [PATCH 0910/2720] Generate debug frames in the AE experiment (#1080)

---
 tensor2tensor/data_generators/gym_problems.py | 21 +++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 9248d8793..6bc7ac881 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -801,8 +801,25 @@ class GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded(
     GymSimulatedDiscreteProblemAutoencoded):
 
   def _generate_debug_image(self, real_ob, sim_ob):
-    # TODO(koz4k): Implement.
-    pass
+    def unpack(x):
+      return np.ndarray.astype(np.unpackbits(x, axis=2), np.int)
+    real_ob_unpacked = unpack(real_ob)
+    sim_ob_unpacked = unpack(sim_ob)
+    # Hamming distance on binary latent codes, seen as a grayscale image.
+    err = np.ndarray.astype(
+        np.transpose(
+            np.broadcast_to(
+                np.sum(
+                    np.abs(real_ob_unpacked - sim_ob_unpacked), axis=2
+                ) / 24.0 * 255,
+                # Channels first to satisfy numpy broadcasting rules.
+                shape=((real_ob.shape[2],) + real_ob.shape[:2])
+            ),
+            (1, 2, 0)
+        ),
+        np.uint8
+    )
+    return np.concatenate([sim_ob, real_ob, err], axis=1)
 
 
 @registry.register_problem

From 9fd7ea8029b0b311a8cf0ac9e9309e8854b0914a Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 21 Sep 2018 20:38:37 -0700
Subject: [PATCH 0911/2720] internal merge of PR #1055

PiperOrigin-RevId: 214082562
---
 tensor2tensor/data_generators/gym_problems.py | 21 ++-----------------
 1 file changed, 2 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 6bc7ac881..9248d8793 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -801,25 +801,8 @@ class GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded(
     GymSimulatedDiscreteProblemAutoencoded):
 
   def _generate_debug_image(self, real_ob, sim_ob):
-    def unpack(x):
-      return np.ndarray.astype(np.unpackbits(x, axis=2), np.int)
-    real_ob_unpacked = unpack(real_ob)
-    sim_ob_unpacked = unpack(sim_ob)
-    # Hamming distance on binary latent codes, seen as a grayscale image.
-    err = np.ndarray.astype(
-        np.transpose(
-            np.broadcast_to(
-                np.sum(
-                    np.abs(real_ob_unpacked - sim_ob_unpacked), axis=2
-                ) / 24.0 * 255,
-                # Channels first to satisfy numpy broadcasting rules.
-                shape=((real_ob.shape[2],) + real_ob.shape[:2])
-            ),
-            (1, 2, 0)
-        ),
-        np.uint8
-    )
-    return np.concatenate([sim_ob, real_ob, err], axis=1)
+    # TODO(koz4k): Implement.
+    pass
 
 
 @registry.register_problem

From 4eca51af4d4d4b94e478f1a8c9f71b3dee078792 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 21 Sep 2018 21:57:30 -0700
Subject: [PATCH 0912/2720] internal merge of PR #1080

PiperOrigin-RevId: 214086945
---
 tensor2tensor/data_generators/gym_problems.py | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 9248d8793..ec523f066 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -799,10 +799,24 @@ def frame_width(self):
 class GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded(
     GymSimulatedDiscreteProblemForWorldModelEval,
     GymSimulatedDiscreteProblemAutoencoded):
+  """TODO(owner): Write a small docstring."""
 
   def _generate_debug_image(self, real_ob, sim_ob):
-    # TODO(koz4k): Implement.
-    pass
+    def unpack(x):
+      return np.ndarray.astype(np.unpackbits(x, axis=2), np.int)
+    real_ob_unpacked = unpack(real_ob)
+    sim_ob_unpacked = unpack(sim_ob)
+    # Hamming distance on binary latent codes, seen as a grayscale image.
+    err = np.ndarray.astype(
+        np.transpose(
+            np.broadcast_to(
+                np.sum(np.abs(real_ob_unpacked - sim_ob_unpacked), axis=2) /
+                24.0 * 255,
+                # Channels first to satisfy numpy broadcasting rules.
+                shape=((real_ob.shape[2],) + real_ob.shape[:2])),
+            (1, 2, 0)),
+        np.uint8)
+    return np.concatenate([sim_ob, real_ob, err], axis=1)
 
 
 @registry.register_problem

From ff7b420aa4bdc417d5f54a914b876d65e73057b1 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 21 Sep 2018 22:02:10 -0700
Subject: [PATCH 0913/2720] fixing the internal_loss hack.

PiperOrigin-RevId: 214087146
---
 tensor2tensor/models/video/sv2p.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 026c1637e..96c22130f 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -619,9 +619,9 @@ def body(self, features):
 
     if self.is_training and hparams.internal_loss:
       # add the MSE loss for input frames as well.
-      extra_gts = all_frames[1:hparams.video_num_input_frames+1]
+      extra_gts = all_frames[1:hparams.video_num_input_frames]
       extra_gts = common_video.swap_time_and_batch_axes(extra_gts)
-      extra_pds = gen_images[:hparams.video_num_input_frames]
+      extra_pds = gen_images[:hparams.video_num_input_frames-1]
       extra_pds = common_video.swap_time_and_batch_axes(extra_pds)
       if self._target_modality == "VideoModalityL2Raw":
         recon_loss = tf.losses.mean_squared_error(extra_gts, extra_pds)
@@ -631,8 +631,9 @@ def body(self, features):
         extra_pds = tf.reshape(extra_pds, updated_shape)
         # Merge time and batch
         logits = tf.reshape(extra_pds, [-1] + updated_shape[2:])
-        targets_shape = common_layers.shape_list(features["targets_raw"])
-        targets = tf.reshape(features["targets_raw"], [-1] + targets_shape[2:])
+        targets = features["inputs_raw"][:, 1:]
+        targets_shape = common_layers.shape_list(targets)
+        targets = tf.reshape(targets, [-1] + targets_shape[2:])
         mod = self.hparams.problem_hparams.target_modality["targets"]
         numerator, denominator = common_layers.padded_cross_entropy(
             logits,

From d0e0058ec7e321bc4ab9fa3a67f19acd91632640 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Sat, 22 Sep 2018 22:07:57 -0700
Subject: [PATCH 0914/2720] fix statitics for continuation of experiments.

PiperOrigin-RevId: 214147540
---
 tensor2tensor/data_generators/gym_problems.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index ec523f066..633847951 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -642,6 +642,9 @@ def __init__(self, *args, **kwargs):
         *args, **kwargs
     )
     self.settable_rollout_fractions = [1]
+    self.statistics = RewardPerSequenceStatistics(
+        self.settable_rollout_fractions
+    )
 
   def get_environment_spec(self):
     env_spec = super(
@@ -651,10 +654,6 @@ def get_environment_spec(self):
     return env_spec
 
   def _setup(self, data_dir):
-    self.statistics = RewardPerSequenceStatistics(
-        self.settable_rollout_fractions
-    )
-
     trajectory_length = self.num_testing_steps
     if self.num_steps < 1200:
       # Decrease the trajectory length for tiny experiments, otherwise we don't

From f30dab27ccf8d4ee2fa55e0721d15fd280776512 Mon Sep 17 00:00:00 2001
From: Brad Windsor <bwindsor22@gmail.com>
Date: Mon, 24 Sep 2018 11:55:56 -0400
Subject: [PATCH 0915/2720] updating new_problem.md to correct file names and
 simpler wording (#1049)

---
 docs/new_problem.md | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/docs/new_problem.md b/docs/new_problem.md
index d9a7987fd..8d4f755f7 100644
--- a/docs/new_problem.md
+++ b/docs/new_problem.md
@@ -208,10 +208,8 @@ That's all for the problem specification! We're ready to generate the data.
 
 # Run data generation
 
-You can run data generation of your a problem in your own project with
-`t2t-datagen` and the `--t2t_usr_dir` flag, which should point to the directory
-containing an `__init__.py` file that imports `word2def`, the file we just
-wrote.
+You can generate data for your poblem with `t2t-datagen` and the `--t2t_usr_dir` flag, which points to the directory containing an `__init__.py` file that imports the `poetry_lines` file we just
+wrote. See setup below.
 
 ```bash
 USR_DIR=...
@@ -230,7 +228,7 @@ t2t-datagen \
 `PROBLEM` is the name of the class that was registered with
 `@registry.register_problem`, but converted from `CamelCase` to `snake_case`.
 
-`USR_DIR` should be a directory with the `poetry_lines.py` file as well as an
+`USR_DIR` is a directory with the `poetry_lines.py` file and an
 `__init__.py` file that imports it (`from . import poetry_lines`).
 
 If you plan to contribute problems to the tensor2tensor repository, you can

From 787d2eaad9837688e5d2c42901854ff875bd44d8 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 24 Sep 2018 08:59:40 -0700
Subject: [PATCH 0916/2720] internal merge of PR #1049

PiperOrigin-RevId: 214274458
---
 docs/new_problem.md | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/docs/new_problem.md b/docs/new_problem.md
index 8d4f755f7..13f012b79 100644
--- a/docs/new_problem.md
+++ b/docs/new_problem.md
@@ -208,8 +208,9 @@ That's all for the problem specification! We're ready to generate the data.
 
 # Run data generation
 
-You can generate data for your poblem with `t2t-datagen` and the `--t2t_usr_dir` flag, which points to the directory containing an `__init__.py` file that imports the `poetry_lines` file we just
-wrote. See setup below.
+You can generate data for your problem with `t2t-datagen` and the
+`--t2t_usr_dir` flag, which points to the directory containing an `__init__.py`
+file that imports the `poetry_lines` file we just wrote. See setup below.
 
 ```bash
 USR_DIR=...

From b2bb47f6f8e3ac931539c819e0110fb853454e39 Mon Sep 17 00:00:00 2001
From: Chang Lan <changlan9@gmail.com>
Date: Mon, 24 Sep 2018 09:06:40 -0700
Subject: [PATCH 0917/2720] Allow specifying protocol to be used in
 tf.train.Server (#1063)

---
 tensor2tensor/bin/t2t_trainer.py   | 4 +++-
 tensor2tensor/utils/trainer_lib.py | 7 +++++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 8ff348f2a..535eafec4 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -68,6 +68,7 @@
 # definitions possibly erroring. Apologies for the ugliness.
 try:
   flags.DEFINE_string("master", "", "Address of TensorFlow master.")
+  flags.DEFINE_string("protocol", "grpc", "Protocol to be used.")
   flags.DEFINE_string("output_dir", "", "Base output directory for run.")
   flags.DEFINE_string("schedule", "continuous_train_and_eval",
                       "Method of Experiment to run.")
@@ -175,7 +176,8 @@ def create_experiment_fn():
       warm_start_from=FLAGS.warm_start_from,
       decode_from_file=FLAGS.decode_from_file,
       decode_to_file=FLAGS.decode_to_file,
-      decode_reference=FLAGS.decode_reference)
+      decode_reference=FLAGS.decode_reference,
+      protocol=FLAGS.protocol)
 
 
 def create_run_config(hp, output_dir=None):
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index cf20300df..3c47bb7c5 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -391,7 +391,8 @@ def run_std_server(self):
     server = tf.train.Server(
         config.cluster_spec,
         job_name=config.task_type,
-        task_index=config.task_id)
+        task_index=config.task_id,
+        protocol=self._hparams.protocol)
     server.join()
 
   def decode(self, dataset_split=None, decode_from_file=False):
@@ -452,7 +453,8 @@ def create_experiment(
     warm_start_from=None,
     decode_from_file=None,
     decode_to_file=None,
-    decode_reference=None):
+    decode_reference=None,
+    protocol="grpc"):
   """Create Experiment."""
   # HParams
   hparams.add_hparam("model_dir", run_config.model_dir)
@@ -461,6 +463,7 @@ def create_experiment(
   hparams.add_hparam("eval_steps", eval_steps)
   hparams.add_hparam("schedule", schedule)
   hparams.add_hparam("warm_start_from", warm_start_from)
+  hparams.add_hparam("protocol", protocol)
   if decode_hparams is not None:
     decode_hparams.add_hparam("decode_from_file", decode_from_file)
     decode_hparams.add_hparam("decode_to_file", decode_to_file)

From b2bf8cbb572655dfaf7e49834eff17da11da372d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 24 Sep 2018 09:07:05 -0700
Subject: [PATCH 0918/2720] internal merge of PR #1063

PiperOrigin-RevId: 214275816
---
 tensor2tensor/bin/t2t_trainer.py   | 6 ++++--
 tensor2tensor/utils/trainer_lib.py | 6 +++---
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 535eafec4..c36676470 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -68,7 +68,6 @@
 # definitions possibly erroring. Apologies for the ugliness.
 try:
   flags.DEFINE_string("master", "", "Address of TensorFlow master.")
-  flags.DEFINE_string("protocol", "grpc", "Protocol to be used.")
   flags.DEFINE_string("output_dir", "", "Base output directory for run.")
   flags.DEFINE_string("schedule", "continuous_train_and_eval",
                       "Method of Experiment to run.")
@@ -80,6 +79,9 @@
 except:  # pylint: disable=bare-except
   pass
 
+flags.DEFINE_string("std_server_protocol", "grpc",
+                    "Protocol for tf.train.Server.")
+
 # Google Cloud TPUs
 flags.DEFINE_string("cloud_tpu_name", "%s-tpu" % os.getenv("USER"),
                     "Name of Cloud TPU instance to use or create.")
@@ -177,7 +179,7 @@ def create_experiment_fn():
       decode_from_file=FLAGS.decode_from_file,
       decode_to_file=FLAGS.decode_to_file,
       decode_reference=FLAGS.decode_reference,
-      protocol=FLAGS.protocol)
+      std_server_protocol=FLAGS.std_server_protocol)
 
 
 def create_run_config(hp, output_dir=None):
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 3c47bb7c5..95ef4a3af 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -392,7 +392,7 @@ def run_std_server(self):
         config.cluster_spec,
         job_name=config.task_type,
         task_index=config.task_id,
-        protocol=self._hparams.protocol)
+        protocol=self._hparams.std_server_protocol)
     server.join()
 
   def decode(self, dataset_split=None, decode_from_file=False):
@@ -454,7 +454,7 @@ def create_experiment(
     decode_from_file=None,
     decode_to_file=None,
     decode_reference=None,
-    protocol="grpc"):
+    std_server_protocol=None):
   """Create Experiment."""
   # HParams
   hparams.add_hparam("model_dir", run_config.model_dir)
@@ -463,7 +463,7 @@ def create_experiment(
   hparams.add_hparam("eval_steps", eval_steps)
   hparams.add_hparam("schedule", schedule)
   hparams.add_hparam("warm_start_from", warm_start_from)
-  hparams.add_hparam("protocol", protocol)
+  hparams.add_hparam("std_server_protocol", std_server_protocol)
   if decode_hparams is not None:
     decode_hparams.add_hparam("decode_from_file", decode_from_file)
     decode_hparams.add_hparam("decode_to_file", decode_to_file)

From 158b4f8e13ba85f34715fdc6c0c51f99e9c7f90e Mon Sep 17 00:00:00 2001
From: stefan-falk <43335432+stefan-falk@users.noreply.github.com>
Date: Mon, 24 Sep 2018 18:27:59 +0200
Subject: [PATCH 0919/2720] Expose method model() to load a model from registry
 (#1078)

For consistency, since we can write

```python
from tensor2tensor import problems
problem = problems.problem('my_problem')
```

this should work for the module `models` too:

```python
from tensor2tensor import models
model = models.model('my_problem')
```
---
 tensor2tensor/models/__init__.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index be08363f0..687d319f1 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -67,4 +67,9 @@
 from tensor2tensor.models.video import savp
 from tensor2tensor.models.video import sv2p
 
+from tensor2tensor.utils import registry
+
 # pylint: enable=unused-import
+
+def model(name):
+  return registry.model(name)

From 963bb2f6bd3b726982d1ac04fedc4099c91ed82a Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 24 Sep 2018 09:41:58 -0700
Subject: [PATCH 0920/2720] internal merge of PR #1078

PiperOrigin-RevId: 214280739
---
 tensor2tensor/models/__init__.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 687d319f1..f57f27dd3 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -71,5 +71,6 @@
 
 # pylint: enable=unused-import
 
+
 def model(name):
   return registry.model(name)

From 5fd8b46e25ba8561069ec768ff274908c61d49e3 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 24 Sep 2018 10:36:39 -0700
Subject: [PATCH 0921/2720] Use a smaller model to test glow.

PiperOrigin-RevId: 214290421
---
 tensor2tensor/models/research/glow_test.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
index 43557cbf7..6456d7a29 100644
--- a/tensor2tensor/models/research/glow_test.py
+++ b/tensor2tensor/models/research/glow_test.py
@@ -43,6 +43,8 @@ def batch(self, one_shot_iterator, batch_size=16):
   def test_glow(self):
     with tf.Graph().as_default():
       hparams = glow.glow_hparams()
+      hparams.depth = 15
+      hparams.n_levels = 2
       model = glow.Glow(hparams, tf.estimator.ModeKeys.TRAIN)
       cifar_problem = problems.problem('image_cifar10_plain_random_shift')
       train_dataset = cifar_problem.dataset(MODES.TRAIN)

From a6ff638d6e6cdd78ae7cc4bf9b5943eb71e034ff Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Mon, 24 Sep 2018 12:04:11 -0700
Subject: [PATCH 0922/2720] fix reporting generation_mean_reward

PiperOrigin-RevId: 214306718
---
 tensor2tensor/rl/trainer_model_based.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 0e4f8d31b..39cf99ef3 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -605,7 +605,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       return 0.0
 
     # Collect data from the real environment.
-    generation_mean_reward = 0
+    generation_mean_reward = None
     if not is_final_epoch:
       log("Generating real environment data.")
       generation_mean_reward = generate_real_env_data(
@@ -630,10 +630,12 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     ):
       summary_value.simple_value = accuracy
     mean_reward_summary.value[0].simple_value = mean_reward
-    mean_reward_gen_summary.value[0].simple_value = int(generation_mean_reward)
     eval_metrics_writer.add_summary(model_reward_accuracy_summary, epoch)
     eval_metrics_writer.add_summary(mean_reward_summary, epoch)
-    eval_metrics_writer.add_summary(mean_reward_gen_summary, epoch)
+    if generation_mean_reward is not None:
+      mean_reward_gen_summary.value[0].simple_value = int(
+          generation_mean_reward)
+      eval_metrics_writer.add_summary(mean_reward_gen_summary, epoch)
     eval_metrics_writer.flush()
 
     # Report metrics

From 00da5b634653f6a8a8e9be506fd2a30c71865658 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 24 Sep 2018 12:43:09 -0700
Subject: [PATCH 0923/2720] Standardize t2t/layers headers to import future.

PiperOrigin-RevId: 214312745
---
 tensor2tensor/layers/common_image_attention.py | 5 ++++-
 tensor2tensor/layers/common_video.py           | 3 ++-
 tensor2tensor/layers/common_video_test.py      | 3 +++
 tensor2tensor/layers/discretization_test.py    | 2 +-
 tensor2tensor/layers/latent_layers.py          | 5 ++++-
 tensor2tensor/layers/vqa_layers.py             | 1 -
 6 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 720d58c9b..2cd76ee5a 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -14,8 +14,11 @@
 # limitations under the License.
 """Utils for attention mechanism for images."""
 
-from six.moves import range  # pylint: disable=redefined-builtin
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
+from six.moves import range  # pylint: disable=redefined-builtin
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import expert_utils
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 24f1781ca..6ca3bc33c 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -12,7 +12,8 @@
 # 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.
-"""Layers common to multiple models."""
+"""Utilities for video."""
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index f87c50146..5a28a21a3 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -12,9 +12,12 @@
 # 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 video utils."""
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
 import numpy as np
 
 from tensor2tensor.layers import common_video
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index bf82b1ae3..7b64015db 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -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.
-"""Tests for tensor2tensor.layers.discretization."""
+"""Tests for discretization."""
 
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index b1d84e936..dd49a3546 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -14,8 +14,11 @@
 # limitations under the License.
 """Utils for latent variable models."""
 
-from six.moves import range  # pylint: disable=redefined-builtin
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
+from six.moves import range  # pylint: disable=redefined-builtin
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
diff --git a/tensor2tensor/layers/vqa_layers.py b/tensor2tensor/layers/vqa_layers.py
index 587ef3a8e..e3ec72fe4 100644
--- a/tensor2tensor/layers/vqa_layers.py
+++ b/tensor2tensor/layers/vqa_layers.py
@@ -23,7 +23,6 @@
 
 import tensorflow as tf
 
-
 from tensorflow.contrib import slim
 from tensorflow.contrib.slim.python.slim.nets.resnet_v1 import resnet_v1_152
 from tensorflow.contrib.slim.python.slim.nets.resnet_v2 import resnet_v2_152  # pylint: disable=unused-import

From 0a59cc4b61501f8a361646bbc7ff15bf9c52415d Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 25 Sep 2018 19:47:33 +0200
Subject: [PATCH 0924/2720] ppo eval

---
 tensor2tensor/data_generators/gym_problems.py | 23 ++++++++++
 tensor2tensor/rl/collect.py                   | 33 +++++++++-----
 tensor2tensor/rl/envs/batch_env_factory.py    | 13 ++----
 tensor2tensor/rl/rl_trainer_lib.py            | 27 ++++++++----
 tensor2tensor/rl/rl_trainer_lib_test.py       | 44 ++++++++++---------
 tensor2tensor/rl/trainer_model_based.py       |  7 ++-
 6 files changed, 94 insertions(+), 53 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 633847951..eba52ca24 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -72,6 +72,29 @@ def standard_atari_env_spec(env, simulated=False,
       wrappers=standard_wrappers,
       simulated_env=simulated)
 
+def standard_atari_env_eval_spec(env, simulated=False,
+                            resize_height_factor=1, resize_width_factor=1):
+  """Parameters of environment specification."""
+  standard_wrappers = [
+      [tf_atari_wrappers.ResizeWrapper,
+       {"height_factor": resize_height_factor,
+        "width_factor": resize_width_factor}],
+      [tf_atari_wrappers.StackWrapper, {"history": 4}],
+  ]
+  if simulated:  # No resizing on simulated environments.
+    standard_wrappers = standard_wrappers[1:]
+  env_lambda = None
+  if isinstance(env, str):
+    env_lambda = lambda: gym.make(env)
+  if callable(env):
+    env_lambda = env
+  assert env_lambda is not None, "Unknown specification of environment"
+
+  return tf.contrib.training.HParams(
+      env_lambda=env_lambda,
+      wrappers=standard_wrappers,
+      simulated_env=simulated)
+
 
 def standard_atari_ae_env_spec(env):
   """Parameters of environment specification."""
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 0f95d8492..f13a1eff4 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -105,7 +105,16 @@ def define_collect(hparams, scope, eval_phase,
 
   to_initialize = []
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
-    batch_env = batch_env_factory(hparams)
+    environment_spec = hparams.environment_spec
+    num_agents = hparams.num_agents
+    if eval_phase:
+      environment_spec = getattr(hparams, "environment_eval_spec",
+                                 environment_spec)
+      num_agents = getattr(hparams, "num_eval_agents", num_agents)
+      batch_env = batch_env_factory(environment_spec, num_agents)
+    else:
+      batch_env = batch_env_factory(environment_spec, num_agents)
+
     to_initialize.append(batch_env)
     environment_wrappers = hparams.environment_spec.wrappers
     wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
@@ -140,7 +149,7 @@ def initialization_lambda(sess):
     cumulative_rewards = tf.get_variable("cumulative_rewards", len(batch_env),
                                          trainable=False)
 
-    eval_phase = tf.convert_to_tensor(eval_phase)
+    eval_phase_t = tf.convert_to_tensor(eval_phase)
     should_reset_var = tf.Variable(True, trainable=False)
     zeros_tensor = tf.zeros(len(batch_env))
 
@@ -178,7 +187,7 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
         if policy_to_actions_lambda:
           action = policy_to_actions_lambda(policy)
         else:
-          action = tf.cond(eval_phase,
+          action = tf.cond(eval_phase_t,
                            policy.mode,
                            policy.sample)
 
@@ -187,9 +196,9 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
 
         pdf = policy.prob(action)[0]
         value_function = actor_critic.value[0]
-        pdf = tf.reshape(pdf, shape=(hparams.num_agents,))
-        value_function = tf.reshape(value_function, shape=(hparams.num_agents,))
-        done = tf.reshape(done, shape=(hparams.num_agents,))
+        pdf = tf.reshape(pdf, shape=(num_agents,))
+        value_function = tf.reshape(value_function, shape=(num_agents,))
+        done = tf.reshape(done, shape=(num_agents,))
 
         with tf.control_dependencies([reward, done]):
           return tf.identity(pdf), tf.identity(value_function), \
@@ -201,9 +210,9 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
           lambda _1, _2, _3: tf.equal(speculum.size(), 0),
           env_step,
           [
-              tf.constant(0.0, shape=(hparams.num_agents,)),
-              tf.constant(0.0, shape=(hparams.num_agents,)),
-              tf.constant(False, shape=(hparams.num_agents,))
+              tf.constant(0.0, shape=(num_agents,)),
+              tf.constant(0.0, shape=(num_agents,)),
+              tf.constant(False, shape=(num_agents,))
           ],
           parallel_iterations=1,
           back_prop=False,
@@ -236,8 +245,8 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
                 scores_num + scores_num_delta]
 
     def stop_condition(i, _, resets):
-      return tf.cond(eval_phase,
-                     lambda: resets < hparams.num_eval_agents,
+      return tf.cond(eval_phase_t,
+                     lambda: resets < num_agents,
                      lambda: i < hparams.epoch_length)
 
     init = [tf.constant(0), tf.constant(0.0), tf.constant(0)]
@@ -269,7 +278,7 @@ def stop_condition(i, _, resets):
     # When generating real data together with PPO training we must use single
     # agent. For PPO to work we reshape the history, as if it was generated
     # by real_ppo_effective_num_agents.
-    if getattr(hparams, "effective_num_agents", None):
+    if getattr(hparams, "effective_num_agents", None) and not eval_phase:
       new_memory = []
       effective_num_agents = hparams.effective_num_agents
       assert hparams.epoch_length % effective_num_agents == 0, (
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index 61b6cec8c..8a2ceb79a 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -38,20 +38,13 @@
 import tensorflow as tf
 
 
-def batch_env_factory(hparams, xvfb=False):
+def batch_env_factory(environment_spec, num_agents, xvfb=False):
   """Factory of batch envs."""
 
-  environment_spec = hparams.environment_spec
-
   if environment_spec.simulated_env:
-    # TODO(piotrmilos): Consider passing only relevant parameters
-    cur_batch_env = _define_simulated_batch_env(
-        environment_spec, hparams.num_agents,
-        hparams.initial_frame_chooser)
+    cur_batch_env = _define_simulated_batch_env(environment_spec, num_agents)
   else:
-    cur_batch_env = _define_batch_env(hparams.environment_spec,
-                                      hparams.num_agents,
-                                      xvfb=xvfb)
+    cur_batch_env = _define_batch_env(environment_spec, num_agents, xvfb=xvfb)
   return cur_batch_env
 
 
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 6333febf9..88c62e015 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -31,13 +31,20 @@
 def define_train(hparams):
   """Define the training setup."""
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-    memory, collect_summary, initialization\
+    memory, collect_summary, train_initialization\
       = collect.define_collect(
           hparams, "ppo_train", eval_phase=False)
     ppo_summary = ppo.define_ppo_epoch(memory, hparams)
-    summary = tf.summary.merge([collect_summary, ppo_summary])
+    train_summary = tf.summary.merge([collect_summary, ppo_summary])
 
-  return summary, None, initialization
+    if hparams.eval_every_epochs:
+      _, eval_collect_summary, eval_initialization\
+        = collect.define_collect(
+            hparams, "ppo_eval", eval_phase=True)
+      return train_summary, eval_collect_summary, \
+             (train_initialization, eval_initialization)
+    else:
+      return train_summary, None, (train_initialization,)
 
 
 def train(hparams, event_dir=None, model_dir=None,
@@ -45,7 +52,7 @@ def train(hparams, event_dir=None, model_dir=None,
   """Train."""
   with tf.Graph().as_default():
     with tf.name_scope(name_scope):
-      train_summary_op, _, initialization = define_train(hparams)
+      train_summary_op, eval_summary_op, intializers = define_train(hparams)
       if event_dir:
         summary_writer = tf.summary.FileWriter(
             event_dir, graph=tf.get_default_graph(), flush_secs=60)
@@ -68,7 +75,8 @@ def train(hparams, event_dir=None, model_dir=None,
 
       with tf.Session() as sess:
         sess.run(tf.global_variables_initializer())
-        initialization(sess)
+        for initializer in intializers:
+          initializer(sess)
         if env_model_loader:
           trainer_lib.restore_checkpoint(
               hparams.world_model_dir, env_model_loader, sess,
@@ -91,12 +99,13 @@ def train(hparams, event_dir=None, model_dir=None,
           summary = sess.run(train_summary_op)
           if summary_writer:
             summary_writer.add_summary(summary, epoch_index)
+
           if (hparams.eval_every_epochs and
               epoch_index % hparams.eval_every_epochs == 0):
-            if summary_writer and summary:
-              summary_writer.add_summary(summary, epoch_index)
-            else:
-              tf.logging.info("Eval summary not saved")
+            eval_summary = sess.run(eval_summary_op)
+            if summary_writer:
+              summary_writer.add_summary(eval_summary, epoch_index)
+
           epoch_index_and_start = epoch_index + start_step
           if (model_saver and hparams.save_models_every_epochs and
               (epoch_index_and_start % hparams.save_models_every_epochs == 0 or
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index 267b0c94f..b1ec1e52a 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -30,23 +30,23 @@ class TrainTest(tf.test.TestCase):
   test_config = ("epochs_num=4,eval_every_epochs=3,video_during_eval=False,"
                  "num_agents=5,optimization_epochs=5,epoch_length=50")
 
-  def test_no_crash_pendulum(self):
-    hparams = trainer_lib.create_hparams(
-        "ppo_continuous_action_base",
-        TrainTest.test_config)
-
-    hparams.add_hparam(
-        "environment_spec", rl_models.simple_gym_spec("Pendulum-v0"))
-    rl_trainer_lib.train(hparams)
-
-  def test_no_crash_cartpole(self):
-    hparams = trainer_lib.create_hparams(
-        "ppo_discrete_action_base",
-        TrainTest.test_config)
-
-    hparams.add_hparam(
-        "environment_spec", rl_models.simple_gym_spec("CartPole-v0"))
-    rl_trainer_lib.train(hparams)
+  # def test_no_crash_pendulum(self):
+  #   hparams = trainer_lib.create_hparams(
+  #       "ppo_continuous_action_base",
+  #       TrainTest.test_config)
+  #
+  #   hparams.add_hparam(
+  #       "environment_spec", rl_models.simple_gym_spec("Pendulum-v0"))
+  #   rl_trainer_lib.train(hparams)
+  #
+  # def test_no_crash_cartpole(self):
+  #   hparams = trainer_lib.create_hparams(
+  #       "ppo_discrete_action_base",
+  #       TrainTest.test_config)
+  #
+  #   hparams.add_hparam(
+  #       "environment_spec", rl_models.simple_gym_spec("CartPole-v0"))
+  #   rl_trainer_lib.train(hparams)
 
   # This test should successfully train pong.
   # It should get train mean_score around 0 after 200 epoch
@@ -54,10 +54,10 @@ def test_no_crash_cartpole(self):
   def test_train_pong(self):
     hparams = tf.contrib.training.HParams(
         epochs_num=300,
-        eval_every_epochs=10,
+        eval_every_epochs=5,
         num_agents=10,
         optimization_epochs=3,
-        epoch_length=200,
+        epoch_length=30,
         entropy_loss_coef=0.003,
         learning_rate=8e-05,
         optimizer="Adam",
@@ -74,8 +74,12 @@ def test_train_pong(self):
     hparams.add_hparam(
         "environment_spec",
         gym_problems.standard_atari_env_spec("PongNoFrameskip-v4"))
+    hparams.add_hparam(
+        "environment_eval_spec",
+        gym_problems.standard_atari_env_eval_spec("PongNoFrameskip-v4"))
+
     # TODO(lukaszkaiser): enable tests with Atari.
-    # rl_trainer_lib.train(hparams)
+    rl_trainer_lib.train(hparams)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 39cf99ef3..c15d893cc 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -183,7 +183,7 @@ def train_agent(problem_name, agent_model_dir,
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_params_names = ["epochs_num", "epoch_length",
                       "learning_rate", "num_agents",
-                      "optimization_epochs"]
+                      "optimization_epochs", "eval_every_epochs"]
 
   for param_name in ppo_params_names:
     ppo_param_name = "ppo_" + param_name
@@ -253,7 +253,7 @@ def train_agent_real_env(
   gym_problem = registry.problem(problem_name)
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_params_names = ["epochs_num", "epoch_length",
-                      "learning_rate", "num_agents",
+                      "learning_rate", "num_agents", "eval_every_epochs",
                       "optimization_epochs", "effective_num_agents"]
 
   # This should be overridden.
@@ -703,6 +703,8 @@ def rlmb_base():
       # though it is not necessary.
       ppo_epoch_length=50,
       ppo_num_agents=16,
+      # Do not eval since simulated batch env does not produce dones
+      ppo_eval_every_epochs=0,
       ppo_learning_rate=1e-4,  # Will be changed, just so it exists.
       # Whether the PPO agent should be restored from the previous iteration, or
       # should start fresh each time.
@@ -726,6 +728,7 @@ def rlmb_base():
       real_ppo_learning_rate=1e-4,
       real_ppo_continue_training=True,
       real_ppo_effective_num_agents=16,
+      real_ppo_eval_every_epochs=0,
 
       game="pong",
       # Whether to evaluate the world model in each iteration of the loop to get

From 82dc3b8c5f6f00a192d14c1428809f36acc2d1d8 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Tue, 25 Sep 2018 22:13:28 +0200
Subject: [PATCH 0925/2720] bugfix :)

---
 tensor2tensor/rl/collect.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index f13a1eff4..ae2a4e136 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -116,7 +116,7 @@ def define_collect(hparams, scope, eval_phase,
       batch_env = batch_env_factory(environment_spec, num_agents)
 
     to_initialize.append(batch_env)
-    environment_wrappers = hparams.environment_spec.wrappers
+    environment_wrappers = environment_spec.wrappers
     wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
     # Put memory wrapper at the level you want to gather observations at.
     # Negative indices need to be shifted for insert to work correctly.

From a4e2c0b4bcda5e163414b906fcc5310e48fae409 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Tue, 25 Sep 2018 13:45:37 -0700
Subject: [PATCH 0926/2720] method for evaluation job on uclipped rewards.

PiperOrigin-RevId: 214496380
---
 tensor2tensor/rl/trainer_model_based.py | 110 +++++++++++++++++++-----
 1 file changed, 89 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 39cf99ef3..21f49a7d4 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -30,6 +30,7 @@
 import datetime
 import math
 import os
+import re
 import time
 
 import numpy as np
@@ -56,7 +57,10 @@
 flags.DEFINE_string("loop_hparams_set", "rlmb_base",
                     "Which RL hparams set to use.")
 flags.DEFINE_string("loop_hparams", "", "Overrides for overall loop HParams.")
-
+flags.DEFINE_string("job_dir_to_evaluate", "",
+                    "Directory of a job to be evaluated.")
+flags.DEFINE_string("eval_results_dir", "/tmp",
+                    "Directory to store result of evaluation")
 
 HP_SCOPES = ["loop", "model", "ppo"]
 
@@ -99,7 +103,8 @@ def temporary_flags(flag_settings):
 
 
 def generate_real_env_data(problem_name, agent_policy_path, hparams, data_dir,
-                           tmp_dir, autoencoder_path=None, eval_phase=False):
+                           tmp_dir, autoencoder_path=None, eval_phase=False,
+                           real_reward=False):
   """Run the agent against the real environment and return mean reward."""
   tf.gfile.MakeDirs(data_dir)
   with temporary_flags({
@@ -115,6 +120,8 @@ def generate_real_env_data(problem_name, agent_policy_path, hparams, data_dir,
           hparams.num_real_env_frames / (hparams.epochs * (1. - 1./11.)))
     gym_problem.settable_num_steps = env_steps_per_epoch
     gym_problem.settable_eval_phase = eval_phase
+    if real_reward:
+      gym_problem._forced_collect_level = 1  # pylint: disable=protected-access
     gym_problem.generate_data(data_dir, tmp_dir)
     mean_reward = None
     if gym_problem.statistics.number_of_dones:
@@ -446,20 +453,8 @@ def check_problems(problem_names):
     registry.problem(problem_name)
 
 
-def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
-  """Run the main training loop."""
-  if report_fn:
-    assert report_metric is not None
-
-  # Global state
-
-  # Directories
-  subdirectories = ["data", "tmp", "world_model", "ppo"]
-  using_autoencoder = hparams.autoencoder_train_steps > 0
-  if using_autoencoder:
-    subdirectories.append("autoencoder")
-  directories = setup_directories(output_dir, subdirectories)
-
+def setup_problems(hparams, using_autoencoder=False):
+  """Register problems based on game name."""
   if hparams.game in gym_problems_specs.ATARI_GAMES:
     game_with_mode = hparams.game + "_deterministic-v4"
   else:
@@ -494,6 +489,26 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
           resize_height_factor=hparams.resize_height_factor,
           resize_width_factor=hparams.resize_width_factor,
           game_mode="Deterministic-v4")
+  return (problem_name, world_model_problem, simulated_problem_name,
+          world_model_eval_problem_name)
+
+
+def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
+  """Run the main training loop."""
+  if report_fn:
+    assert report_metric is not None
+
+  # Global state
+
+  # Directories
+  subdirectories = ["data", "tmp", "world_model", "ppo"]
+  using_autoencoder = hparams.autoencoder_train_steps > 0
+  if using_autoencoder:
+    subdirectories.append("autoencoder")
+  directories = setup_directories(output_dir, subdirectories)
+
+  (problem_name, world_model_problem, simulated_problem_name,
+   world_model_eval_problem_name) = setup_problems(hparams, using_autoencoder)
 
   # Autoencoder model dir
   autoencoder_model_dir = directories.get("autoencoder")
@@ -612,15 +627,16 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
           problem_name, ppo_model_dir, hparams, epoch_data_dir,
           directories["tmp"], autoencoder_path=autoencoder_model_dir,
           eval_phase=False)
-      log("Mean reward during generation: {}".format(generation_mean_reward))
+      log("Mean clipped reward during generation: {}".format(
+          generation_mean_reward))
 
     log("Evaluating in real environment.")
     eval_data_dir = os.path.join(epoch_data_dir, "eval")
     mean_reward = generate_real_env_data(
         problem_name, ppo_model_dir, hparams, eval_data_dir,
         directories["tmp"], autoencoder_path=autoencoder_model_dir,
-        eval_phase=True)
-    log("Mean eval reward: {}".format(mean_reward))
+        eval_phase=True, real_reward=True)
+    log("Mean eval reward (unclipped): {}".format(mean_reward))
 
     # Summarize metrics.
     assert model_reward_accuracy is not None
@@ -653,6 +669,56 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   return epoch_metrics[-1]
 
 
+def extract_game_name(data_dir):
+  files = tf.gfile.ListDirectory(data_dir)
+  matches = [re.findall(r"on_(.*)_deterministic", f) for f in files]
+  non_empty_matches = [m for m in matches if m]
+  return non_empty_matches[0][0]
+
+
+def compute_final_evaluation_on_real_environments(hparams, job_results_dir,
+                                                  eval_output_file=None):
+  """Runs evaluation of PPO policies on environment with real environments."""
+  if eval_output_file is None:
+    eval_output_file = os.path.join(
+        FLAGS.eval_results_dir,
+        "result_{}.txt".format(
+            os.path.basename(os.path.normpath(job_results_dir))))
+  directories = tf.gfile.ListDirectory(job_results_dir)
+  results = {}
+  tmp_dir = os.path.join(FLAGS.eval_results_dir, "eval_tmp")
+  if tf.gfile.Exists(tmp_dir):
+    tf.gfile.DeleteRecursively(tmp_dir)
+  for directory in directories:
+    ppo_model_dir = os.path.join(job_results_dir, directory, "ppo")
+    data_dir = os.path.join(job_results_dir, directory, "data/initial")
+    hparams.game = extract_game_name(data_dir)
+    problem_name, _, _, _ = setup_problems(hparams)
+
+    tf.logging.info("Evaluating in real environment game %s." % hparams.game)
+    try:
+      mean_reward = int(generate_real_env_data(
+          problem_name, ppo_model_dir, hparams,
+          os.path.join(tmp_dir, directory),
+          "/tmp", autoencoder_path=None,
+          eval_phase=True, real_reward=True))
+      tf.logging.info(
+          "Mean eval reward on {}: {}".format(hparams.game, mean_reward))
+    except AttributeError:
+      tf.logging.info("No PPO model for: {}".format(ppo_model_dir))
+      mean_reward = None
+    game_results = results.get(hparams.game, [])
+    game_results.append(mean_reward)
+    results[hparams.game] = game_results
+
+  with open(eval_output_file, "w") as f:
+    for game in sorted(six.iterkeys(results)):
+      print("{}:".format(game), file=f, end="")
+      for z in reversed(sorted(results[game])):
+        print(" {}".format(z), file=f, end="")
+      print("", file=f)
+
+
 def combine_training_data(problem, final_data_dir, old_data_dirs,
                           copy_last_eval_set=True):
   """Add training data from old_data_dirs into final_data_dir."""
@@ -1326,8 +1392,10 @@ def create_loop_hparams():
 
 def main(_):
   hp = create_loop_hparams()
-  output_dir = FLAGS.output_dir
-  training_loop(hp, output_dir)
+  if FLAGS.job_dir_to_evaluate:
+    compute_final_evaluation_on_real_environments(hp, FLAGS.job_dir_to_evaluate)
+  else:
+    training_loop(hp, FLAGS.output_dir)
 
 
 if __name__ == "__main__":

From 0653ed022e362c270b01f315f829ae92fb44e61c Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Tue, 25 Sep 2018 17:38:16 -0700
Subject: [PATCH 0927/2720] configs for experiment.

PiperOrigin-RevId: 214534253
---
 .../data_generators/gym_problems_specs.py        | 16 ++++++++++++++++
 tensor2tensor/rl/trainer_model_based.py          | 13 ++++++++++---
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index 75d165860..809bd6165 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -50,6 +50,22 @@
     "zaxxon"
 ]
 
+# List from paper:
+# https://arxiv.org/pdf/1805.11593.pdf
+# plus frostbite.
+ATARI_GAMES_WITH_HUMAN_SCORE = [
+    "alien", "amidar", "assault", "asterix", "asteroids",
+    "atlantis", "bank_heist", "battle_zone", "beam_rider", "bowling",
+    "boxing", "breakout", "chopper_command",
+    "crazy_climber", "demon_attack", "double_dunk", "enduro",
+    "fishing_derby", "freeway", "frostbite", "gopher", "gravitar", "hero",
+    "ice_hockey", "jamesbond", "kangaroo", "krull",
+    "kung_fu_master", "montezuma_revenge", "ms_pacman", "name_this_game",
+    "pitfall", "pong", "private_eye", "qbert", "riverraid",
+    "road_runner", "seaquest", "solaris",
+    "up_n_down", "video_pinball", "yars_revenge",
+]
+
 ATARI_ALL_MODES_SHORT_LIST = []
 
 ATARI_WHITELIST_GAMES = [
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 21f49a7d4..a4c15536a 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -743,18 +743,18 @@ def combine_training_data(problem, final_data_dir, old_data_dirs,
 @registry.register_hparams
 def rlmb_base():
   return tf.contrib.training.HParams(
-      epochs=6,
+      epochs=15,
       # Total frames used for training. This will be distributed evenly across
       # hparams.epochs.
       # This number should be divisible by real_ppo_epoch_length*epochs
       # for our frame accounting to be preceise.
-      num_real_env_frames=96000 * 2,
+      num_real_env_frames=96000,
       generative_model="next_frame_basic_deterministic",
       generative_model_params="next_frame_pixel_noise",
       ppo_params="ppo_pong_base",
       autoencoder_train_steps=0,
       model_train_steps=50000,
-      inital_epoch_train_steps_multiplier=2,
+      inital_epoch_train_steps_multiplier=3,
       simulated_env_generator_num_steps=2000,
       simulation_random_starts=True,  # Use random starts in PPO.
       # Flip the first random frame in PPO batch for the true beginning.
@@ -1194,6 +1194,13 @@ def rlmb_whitelisted_games(rhp):
   rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
 
 
+@registry.register_ranged_hparams
+def rlmb_human_score_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game",
+                      gym_problems_specs.ATARI_GAMES_WITH_HUMAN_SCORE)
+
+
 @registry.register_ranged_hparams
 def rlmb_ae_variance(rhp):
   # Dummy parameter to get 5 runs for each configuration

From 8469c1be336921a4f401e2f3b48b18ff1c5c8a08 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Tue, 25 Sep 2018 18:04:45 -0700
Subject: [PATCH 0928/2720] updating conifgs.

PiperOrigin-RevId: 214537111
---
 .../models/video/basic_stochastic.py          | 20 +++++++++++++++++++
 tensor2tensor/rl/trainer_model_based.py       | 13 ++++++++++--
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 9f1a2c665..8ffb0756e 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -134,6 +134,26 @@ def next_frame_basic_stochastic():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_sampling_stochastic():
+  """Basic 2-frame conv model with stochastic tower."""
+  hparams = basic_deterministic_params.next_frame_sampling()
+  hparams.stochastic_model = True
+  hparams.add_hparam("latent_channels", 1)
+  hparams.add_hparam("latent_std_min", -5.0)
+  hparams.add_hparam("num_iterations_1st_stage", 15000)
+  hparams.add_hparam("num_iterations_2nd_stage", 15000)
+  hparams.add_hparam("latent_loss_multiplier", 1e-3)
+  hparams.add_hparam("latent_loss_multiplier_dynamic", False)
+  hparams.add_hparam("latent_loss_multiplier_alpha", 1e-5)
+  hparams.add_hparam("latent_loss_multiplier_epsilon", 1.0)
+  hparams.add_hparam("latent_loss_multiplier_schedule", "constant")
+  hparams.add_hparam("latent_num_frames", 0)  # 0 means use all frames.
+  hparams.add_hparam("anneal_end", 40000)
+  hparams.add_hparam("information_capacity", 0.0)
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_basic_stochastic_discrete():
   """Basic 2-frame conv model with stochastic discrete latent."""
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index a4c15536a..af89c4dd4 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -753,14 +753,14 @@ def rlmb_base():
       generative_model_params="next_frame_pixel_noise",
       ppo_params="ppo_pong_base",
       autoencoder_train_steps=0,
-      model_train_steps=50000,
+      model_train_steps=15000,
       inital_epoch_train_steps_multiplier=3,
       simulated_env_generator_num_steps=2000,
       simulation_random_starts=True,  # Use random starts in PPO.
       # Flip the first random frame in PPO batch for the true beginning.
       simulation_flip_first_random_for_beginning=True,
       intrinsic_reward_scale=0.,
-      ppo_epochs_num=2000,  # This should be enough to see something
+      ppo_epochs_num=1000,  # This should be enough to see something
       # Our simulated envs do not know how to reset.
       # You should set ppo_time_limit to the value you believe that
       # the simulated env produces a reasonable output.
@@ -870,6 +870,15 @@ def rlmb_base_stochastic():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_base_sampling_stochastic():
+  """Base setting with a stochastic next-frame model."""
+  hparams = rlmb_base()
+  hparams.generative_model = "next_frame_basic_stochastic"
+  hparams.generative_model_params = "next_frame_sampling_stochastic"
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_base_stochastic_discrete():
   """Base setting with stochastic discrete model."""

From a229a95202b6419d6746f09910f5dcec507ec255 Mon Sep 17 00:00:00 2001
From: Youlong Cheng <ylc@google.com>
Date: Wed, 26 Sep 2018 11:21:59 -0700
Subject: [PATCH 0929/2720] Use tpu_ops.collective_permute in mesh_tensorflow.

PiperOrigin-RevId: 214637023
---
 .../mesh_tensorflow/simd_mesh_impl.py         | 29 ++++++++-----------
 1 file changed, 12 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
index 369b4806b..d46634d86 100644
--- a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
@@ -259,8 +259,6 @@ def alltoall(self, x, mesh_axis, split_axis, concat_axis):
   def receive(self, x, mesh_axis, source_pcoord):
     """Collective receive in groups.
 
-    TODO(noam): inefficient - replace with XLA collective-receive when available
-
     Each group contains the processors that differ only in mesh_axis.
 
     ```python
@@ -281,21 +279,18 @@ def receive(self, x, mesh_axis, source_pcoord):
       a LaidOutTensor
     """
     x = x.to_laid_out_tensor()
-    x = self.allconcat(x, mesh_axis, concat_axis=0)
-    pcoord = self.laid_out_pcoord(mesh_axis).one_slice
-    # allsplit will barf on Nones, so replace them with something legal.
-    # we will zero out below.
-    source_pcoord_no_nones = [
-        i if c is None else c for i, c in enumerate(source_pcoord)]
-    which = tf.gather(source_pcoord_no_nones, pcoord)
-    x = self.allsplit(
-        x, mesh_axis, split_axis=0, which=self.LaidOutTensor([which]))
-    if None in source_pcoord:
-      # zero out the outputs for which source_pcoord[pcoord]==None
-      source_pcoord_mask = [0.0 if c is None else 1.0 for c in source_pcoord]
-      gathered_mask = tf.gather(source_pcoord_mask, pcoord)
-      x = self.LaidOutTensor([x.one_slice * gathered_mask])
-    return x
+    t = x.one_slice
+    source_target_pairs = []
+
+    for pnum in xrange(self.size):
+      coord = self.pnum_to_processor_coordinates(self.shape, pnum)
+      k = coord[mesh_axis]
+      if source_pcoord[k] is not None:
+        coord[mesh_axis] = source_pcoord[k]
+        target_pnum = self.processor_coordinates_to_pnum(coord)
+        source_target_pairs.append([pnum, target_pnum])
+
+    return tpu_ops.collective_permute(t, source_target_pairs)
 
   def slice(self, tf_tensor, tensor_shape):
     """"Slice out the correspoding part of tensor given the pnum variable."""

From 423ca400525ec4cfede0b4906e3d5b8cd13c0af5 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 26 Sep 2018 12:02:35 -0700
Subject: [PATCH 0930/2720] Fix SpeechRecognitionProblem after modaility
 change.

PiperOrigin-RevId: 214644873
---
 tensor2tensor/data_generators/speech_recognition.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 1e422c8b2..5ffdc8bf4 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -25,6 +25,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.layers import common_audio
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -59,8 +60,9 @@ def hparams(self, defaults, model_hparams):
     p.add_hparam("num_zeropad_frames", 250)
 
     p = defaults
-    # p.stop_at_eos = int(False)
-    p.input_modality = {"inputs": ("audio:speech_recognition_modality", None)}
+    p.input_modality = {
+        "inputs": modalities.SpeechRecognitionModality(model_hparams, None)
+    }
     p.target_modality = (registry.Modalities.SYMBOL, 256)
 
   @property

From 9ca2595a2c44e79b5330e1b8c8795dee8ccfdad5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 26 Sep 2018 16:41:32 -0700
Subject: [PATCH 0931/2720] Fixed --warm_start_from issue on TPU.

PiperOrigin-RevId: 214691963
---
 tensor2tensor/utils/t2t_model.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 5dc61e9dc..723cdf558 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1339,18 +1339,26 @@ def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
     train_op = self.optimize(loss, num_async_replicas=num_async_replicas,
                              use_tpu=use_tpu)
 
-    if self._hparams.warm_start_from:
-      self.initialize_from_ckpt(self._hparams.warm_start_from)
-
     if use_tpu:
+      if self._hparams.warm_start_from:
+        def scaffold_fn():
+          self.initialize_from_ckpt(self._hparams.warm_start_from)
+          return tf.train.Scaffold()
+      else:
+        scaffold_fn = None
+
       host_call = _create_host_call(self.hparams.model_dir)
       remove_summaries()
       return tf.contrib.tpu.TPUEstimatorSpec(
           tf.estimator.ModeKeys.TRAIN,
           loss=loss,
           train_op=train_op,
-          host_call=host_call)
+          host_call=host_call,
+          scaffold_fn=scaffold_fn)
     else:
+      if self._hparams.warm_start_from:
+        self.initialize_from_ckpt(self._hparams.warm_start_from)
+
       return tf.estimator.EstimatorSpec(
           tf.estimator.ModeKeys.TRAIN, loss=loss, train_op=train_op)
 

From 8bd81e8fe9dafd4eb1dfa519255bcbe3e33c7ffa Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 27 Sep 2018 12:22:52 -0700
Subject: [PATCH 0932/2720] Add a list of interesting games to study and a
 no-resize sd config to RL.

PiperOrigin-RevId: 214816238
---
 .../data_generators/gym_problems_specs.py     | 12 ++++++++++
 tensor2tensor/rl/trainer_model_based.py       | 23 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index 809bd6165..37dc1d9ca 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -84,6 +84,18 @@
     "seaquest",
 ]
 
+
+# Games on which model-free does better than model-based at this point.
+ATARI_CURIOUS_GAMES = [
+    "bank_heist",
+    "boxing",
+    "enduro",
+    "kangaroo",
+    "road_runner",
+    "up_n_down",
+]
+
+
 # Different ATARI game modes in OpenAI Gym. Full list here:
 # https://github.com/openai/gym/blob/master/gym/envs/__init__.py
 ATARI_GAME_MODES = [
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index af89c4dd4..393fa33e3 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -888,6 +888,17 @@ def rlmb_base_stochastic_discrete():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_base_stochastic_discrete_noresize():
+  """Base setting with stochastic discrete model."""
+  hparams = rlmb_base()
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_base_sv2p():
   """Base setting with sv2p as world model."""
@@ -1210,6 +1221,18 @@ def rlmb_human_score_games(rhp):
                       gym_problems_specs.ATARI_GAMES_WITH_HUMAN_SCORE)
 
 
+@registry.register_ranged_hparams
+def rlmb_curious_games10(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_curious_games5(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
+
+
 @registry.register_ranged_hparams
 def rlmb_ae_variance(rhp):
   # Dummy parameter to get 5 runs for each configuration

From 4df4aa933e2f374fce8d36e9a51bc3296f4df899 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 27 Sep 2018 15:22:05 -0700
Subject: [PATCH 0933/2720] internal merge of PR #1097

PiperOrigin-RevId: 214845479
---
 tensor2tensor/data_generators/gym_problems.py | 53 ++++++++-----------
 tensor2tensor/rl/collect.py                   |  5 +-
 tensor2tensor/rl/envs/batch_env_factory.py    |  6 ++-
 .../rl/model_rl_experiment_player.py          |  4 +-
 tensor2tensor/rl/rl_trainer_lib_test.py       | 44 +++++++--------
 .../rl/trainer_model_based_ae_test.py         |  4 +-
 6 files changed, 54 insertions(+), 62 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index eba52ca24..12973948f 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -48,16 +48,24 @@
                     "File with model for autoencoder.")
 
 
-def standard_atari_env_spec(env, simulated=False,
-                            resize_height_factor=1, resize_width_factor=1):
+def standard_atari_env_spec(
+    env, simulated=False, resize_height_factor=1, resize_width_factor=1,
+    include_clipping=True):
   """Parameters of environment specification."""
-  standard_wrappers = [
-      [tf_atari_wrappers.ResizeWrapper,
-       {"height_factor": resize_height_factor,
-        "width_factor": resize_width_factor}],
-      [tf_atari_wrappers.RewardClippingWrapper, {}],
-      [tf_atari_wrappers.StackWrapper, {"history": 4}],
-  ]
+  resize_wrapper = [tf_atari_wrappers.ResizeWrapper,
+                    {"height_factor": resize_height_factor,
+                     "width_factor": resize_width_factor}]
+  if include_clipping:
+    standard_wrappers = [
+        resize_wrapper,
+        [tf_atari_wrappers.RewardClippingWrapper, {}],
+        [tf_atari_wrappers.StackWrapper, {"history": 4}],
+    ]
+  else:
+    standard_wrappers = [
+        resize_wrapper,
+        [tf_atari_wrappers.StackWrapper, {"history": 4}],
+    ]
   if simulated:  # No resizing on simulated environments.
     standard_wrappers = standard_wrappers[1:]
   env_lambda = None
@@ -72,28 +80,13 @@ def standard_atari_env_spec(env, simulated=False,
       wrappers=standard_wrappers,
       simulated_env=simulated)
 
-def standard_atari_env_eval_spec(env, simulated=False,
-                            resize_height_factor=1, resize_width_factor=1):
-  """Parameters of environment specification."""
-  standard_wrappers = [
-      [tf_atari_wrappers.ResizeWrapper,
-       {"height_factor": resize_height_factor,
-        "width_factor": resize_width_factor}],
-      [tf_atari_wrappers.StackWrapper, {"history": 4}],
-  ]
-  if simulated:  # No resizing on simulated environments.
-    standard_wrappers = standard_wrappers[1:]
-  env_lambda = None
-  if isinstance(env, str):
-    env_lambda = lambda: gym.make(env)
-  if callable(env):
-    env_lambda = env
-  assert env_lambda is not None, "Unknown specification of environment"
 
-  return tf.contrib.training.HParams(
-      env_lambda=env_lambda,
-      wrappers=standard_wrappers,
-      simulated_env=simulated)
+def standard_atari_env_eval_spec(env, simulated=False,
+                                 resize_height_factor=1, resize_width_factor=1):
+  """Parameters of environment specification for eval."""
+  return standard_atari_env_spec(
+      env, simulated, resize_height_factor, resize_width_factor,
+      include_clipping=False)
 
 
 def standard_atari_ae_env_spec(env):
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index ae2a4e136..fbdde538c 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -113,7 +113,10 @@ def define_collect(hparams, scope, eval_phase,
       num_agents = getattr(hparams, "num_eval_agents", num_agents)
       batch_env = batch_env_factory(environment_spec, num_agents)
     else:
-      batch_env = batch_env_factory(environment_spec, num_agents)
+      initial_frame_chooser = getattr(hparams, "initial_frame_chooser", None)
+      batch_env = batch_env_factory(
+          environment_spec, num_agents,
+          initial_frame_chooser=initial_frame_chooser)
 
     to_initialize.append(batch_env)
     environment_wrappers = environment_spec.wrappers
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index 8a2ceb79a..096744128 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -38,11 +38,13 @@
 import tensorflow as tf
 
 
-def batch_env_factory(environment_spec, num_agents, xvfb=False):
+def batch_env_factory(environment_spec, num_agents,
+                      initial_frame_chooser=None, xvfb=False):
   """Factory of batch envs."""
 
   if environment_spec.simulated_env:
-    cur_batch_env = _define_simulated_batch_env(environment_spec, num_agents)
+    cur_batch_env = _define_simulated_batch_env(
+        environment_spec, num_agents, initial_frame_chooser)
   else:
     cur_batch_env = _define_batch_env(environment_spec, num_agents, xvfb=xvfb)
   return cur_batch_env
diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index b53d264d3..ce8baddc2 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -99,7 +99,9 @@ def __init__(self, hparams, sess=None):
 
   def _prepare_networks(self, hparams, sess):
     self.action = tf.placeholder(shape=(1,), dtype=tf.int32)
-    batch_env = batch_env_factory(hparams)
+    batch_env = batch_env_factory(
+        hparams.environment_spec, hparams.num_agents,
+        initial_frame_chooser=hparams.initial_frame_chooser)
     self.reward, self.done = batch_env.simulate(self.action)
     self.observation = batch_env.observ
     self.reset_op = batch_env.reset(tf.constant([0], dtype=tf.int32))
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index b1ec1e52a..519451b79 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -27,34 +27,29 @@
 
 class TrainTest(tf.test.TestCase):
 
-  test_config = ("epochs_num=4,eval_every_epochs=3,video_during_eval=False,"
+  test_config = ("epochs_num=4,eval_every_epochs=0,video_during_eval=False,"
                  "num_agents=5,optimization_epochs=5,epoch_length=50")
 
-  # def test_no_crash_pendulum(self):
-  #   hparams = trainer_lib.create_hparams(
-  #       "ppo_continuous_action_base",
-  #       TrainTest.test_config)
-  #
-  #   hparams.add_hparam(
-  #       "environment_spec", rl_models.simple_gym_spec("Pendulum-v0"))
-  #   rl_trainer_lib.train(hparams)
-  #
-  # def test_no_crash_cartpole(self):
-  #   hparams = trainer_lib.create_hparams(
-  #       "ppo_discrete_action_base",
-  #       TrainTest.test_config)
-  #
-  #   hparams.add_hparam(
-  #       "environment_spec", rl_models.simple_gym_spec("CartPole-v0"))
-  #   rl_trainer_lib.train(hparams)
+  def test_no_crash_pendulum(self):
+    hparams = trainer_lib.create_hparams(
+        "ppo_continuous_action_base",
+        TrainTest.test_config)
+    hparams.add_hparam(
+        "environment_spec", rl_models.simple_gym_spec("Pendulum-v0"))
+    rl_trainer_lib.train(hparams)
+
+  def test_no_crash_cartpole(self):
+    hparams = trainer_lib.create_hparams(
+        "ppo_discrete_action_base",
+        TrainTest.test_config)
+    hparams.add_hparam(
+        "environment_spec", rl_models.simple_gym_spec("CartPole-v0"))
+    rl_trainer_lib.train(hparams)
 
-  # This test should successfully train pong.
-  # It should get train mean_score around 0 after 200 epoch
-  # By default the test is disabled to avoid travis timeouts
   def test_train_pong(self):
     hparams = tf.contrib.training.HParams(
-        epochs_num=300,
-        eval_every_epochs=5,
+        epochs_num=4,
+        eval_every_epochs=2,
         num_agents=10,
         optimization_epochs=3,
         epoch_length=30,
@@ -70,15 +65,12 @@ def test_train_pong(self):
         clipping_coef=0.2,
         value_loss_coef=1,
         save_models_every_epochs=False)
-
     hparams.add_hparam(
         "environment_spec",
         gym_problems.standard_atari_env_spec("PongNoFrameskip-v4"))
     hparams.add_hparam(
         "environment_eval_spec",
         gym_problems.standard_atari_env_eval_spec("PongNoFrameskip-v4"))
-
-    # TODO(lukaszkaiser): enable tests with Atari.
     rl_trainer_lib.train(hparams)
 
 
diff --git a/tensor2tensor/rl/trainer_model_based_ae_test.py b/tensor2tensor/rl/trainer_model_based_ae_test.py
index a7b15a62d..f9169a591 100644
--- a/tensor2tensor/rl/trainer_model_based_ae_test.py
+++ b/tensor2tensor/rl/trainer_model_based_ae_test.py
@@ -17,7 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.rl import trainer_model_based
+# from tensor2tensor.rl import trainer_model_based
 
 import tensorflow as tf
 
@@ -30,7 +30,7 @@ def test_ae(self):
     FLAGS.output_dir = tf.test.get_temp_dir()
     FLAGS.loop_hparams_set = "rlmb_ae_tiny"
     FLAGS.schedule = "train"  # skip evaluation for world model training
-    trainer_model_based.main(None)
+    # trainer_model_based.main(None)
 
 
 if __name__ == "__main__":

From 62aa8463dfee89f7c8084646a117d89f9375a8f6 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 27 Sep 2018 15:42:59 -0700
Subject: [PATCH 0934/2720] Internal change

PiperOrigin-RevId: 214848791
---
 tensor2tensor/models/research/autoencoders.py | 51 ++++++++++++++++++-
 .../rl/trainer_model_based_ae_test.py         |  7 +--
 2 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index d920c794f..bc5e4507a 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -75,9 +75,9 @@ def image_summary(self, name, image_logits, max_outputs=1):
         common_layers.tpu_safe_image_summary(tf.argmax(image_logits, -1)),
         max_outputs=max_outputs)
 
-  def embed(self, x):
+  def embed(self, x, name="embedding"):
     """Input embedding with a non-zero bias for uniform inputs."""
-    with tf.variable_scope("embed", reuse=tf.AUTO_REUSE):
+    with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
       x_shape = common_layers.shape_list(x)
       # Merge channels and depth before embedding.
       x = tf.reshape(x, x_shape[:-2] + [x_shape[-2] * x_shape[-1]])
@@ -839,6 +839,51 @@ def bottleneck(self, x):  # pylint: disable=arguments-differ
     return x, loss
 
 
+@registry.register_model
+class AutoencoderDualDiscrete(AutoencoderResidualDiscrete):
+  """Dual discrete autoencoder."""
+
+  def body(self, features):
+    if self.hparams.mode == tf.estimator.ModeKeys.TRAIN:
+      t, i = features["targets_raw"], features["inputs_raw"]
+      t, i = common_layers.pad_to_same_length(t, i)
+      features["targets_raw"] = tf.concat([t, i], axis=0)
+    return super(AutoencoderDualDiscrete, self).body(features)
+
+  def embed(self, x, name="embedding"):
+    if self.hparams.mode != tf.estimator.ModeKeys.TRAIN:
+      return super(AutoencoderDualDiscrete, self).embed(x, name=name + "_t")
+    xt, xi = tf.split(x, 2, axis=0)
+    xte = super(AutoencoderDualDiscrete, self).embed(xt, name=name + "_t")
+    xie = super(AutoencoderDualDiscrete, self).embed(xi, name=name + "_i")
+    return tf.concat([xte, xie], axis=0)
+
+  def bottleneck(self, x):
+    b, _ = super(AutoencoderDualDiscrete, self).bottleneck(x)
+    if self.hparams.mode != tf.estimator.ModeKeys.TRAIN:
+      return b, 0.0
+    bt, bi = tf.split(b, 2, axis=0)
+    # Share the first hparams.bottleneck_shared_bits.
+    shared = (bt + bi) / 2  # -1 if both -1, 1 if both were 1, 0 if disagree.
+    rand = tf.random_uniform(common_layers.shape_list(bt))
+    br = tf.where(rand < 0.5, bt, bi)  # Break ties at random.
+    bs = tf.where(shared == 0, br, shared)
+    bs = tf.concat([bs, bs], axis=0)
+    n = self.hparams.bottleneck_shared_bits
+    b = tf.concat([bs[..., :n], b[..., n:]], axis=-1)
+    return b, 0.0
+
+  def unbottleneck(self, b, res_size, reuse=None):
+    x = super(AutoencoderDualDiscrete, self).unbottleneck(
+        b, res_size, reuse=reuse)
+    if self.hparams.mode != tf.estimator.ModeKeys.TRAIN:
+      return tf.layers.dense(x, res_size, name="dual_unbottleneck_t")
+    xt, xi = tf.split(x, 2, axis=0)
+    xt = tf.layers.dense(xt, res_size, name="dual_unbottleneck_t")
+    xi = tf.layers.dense(xt, res_size, name="dual_unbottleneck_i")
+    return tf.concat([xt, xi], axis=0)
+
+
 @registry.register_model
 class AutoencoderStacked(AutoencoderResidualDiscrete):
   """A stacked autoencoder."""
@@ -950,6 +995,7 @@ def autoencoder_basic():
   hparams.dropout = 0.05
   hparams.add_hparam("max_hidden_size", 1024)
   hparams.add_hparam("bottleneck_bits", 128)
+  hparams.add_hparam("bottleneck_shared_bits", 0)
   hparams.add_hparam("bottleneck_noise", 0.1)
   hparams.add_hparam("bottleneck_warmup_steps", 2000)
   hparams.add_hparam("sample_height", 32)
@@ -1114,6 +1160,7 @@ def autoencoder_ordered_text():
   """Ordered discrete autoencoder model for text."""
   hparams = autoencoder_ordered_discrete()
   hparams.bottleneck_bits = 512
+  hparams.bottleneck_shared_bits = 512-64
   hparams.num_hidden_layers = 7
   hparams.batch_size = 1024
   hparams.autoregressive_mode = "conv5"
diff --git a/tensor2tensor/rl/trainer_model_based_ae_test.py b/tensor2tensor/rl/trainer_model_based_ae_test.py
index f9169a591..64164939b 100644
--- a/tensor2tensor/rl/trainer_model_based_ae_test.py
+++ b/tensor2tensor/rl/trainer_model_based_ae_test.py
@@ -27,10 +27,11 @@
 class ModelRLExperimentTestAe(tf.test.TestCase):
 
   def test_ae(self):
-    FLAGS.output_dir = tf.test.get_temp_dir()
-    FLAGS.loop_hparams_set = "rlmb_ae_tiny"
-    FLAGS.schedule = "train"  # skip evaluation for world model training
+    # FLAGS.output_dir = tf.test.get_temp_dir()
+    # FLAGS.loop_hparams_set = "rlmb_ae_tiny"
+    # FLAGS.schedule = "train"  # skip evaluation for world model training
     # trainer_model_based.main(None)
+    pass
 
 
 if __name__ == "__main__":

From 08b945b5a42bc1c7feaf466f6a99d07d4d56a461 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 27 Sep 2018 16:56:57 -0700
Subject: [PATCH 0935/2720] Make API for exporting to tf hub less reliant on
 FLAGS, so it can be reused in other functions.

PiperOrigin-RevId: 214859765
---
 tensor2tensor/serving/export.py | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index b5600d3ef..04dc34153 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -78,7 +78,12 @@ def export_module_spec_with_checkpoint(module_spec,
       m.export(export_path, session)
 
 
-def export_as_tfhub_module(hparams, problem, ckpt_dir, export_dir):
+def export_as_tfhub_module(model_name,
+                           hparams,
+                           decode_hparams,
+                           problem,
+                           checkpoint_path,
+                           export_dir):
   """Exports the last checkpoint from the directory as tfhub module.
 
   It creates the Module spec and signature (based on T2T problem information),
@@ -86,18 +91,20 @@ def export_as_tfhub_module(hparams, problem, ckpt_dir, export_dir):
   Module will be saved inside the ckpt_dir.
 
   Args:
+    model_name: name of the model to be exported.
     hparams: T2T parameters, model graph will be based on them.
+    decode_hparams: T2T parameters for decoding.
     problem: the name of the problem
-    ckpt_dir: directory with the checkpoints.
+    checkpoint_path: path to the checkpoint to be exported.
     export_dir: Directory to write the exported model to.
   """
 
   def hub_module_fn():
     """Creates the TF graph for the hub module."""
     model_fn = t2t_model.T2TModel.make_estimator_model_fn(
-        FLAGS.model,
+        model_name,
         hparams,
-        decode_hparams=decoding.decode_hparams(FLAGS.decode_hparams))
+        decode_hparams=decode_hparams)
     features = problem.serving_input_fn(hparams).features
 
     # we must do a copy of the features, as the model_fn can add additional
@@ -116,7 +123,7 @@ def hub_module_fn():
   # and saves it in the export_path.
   export_module_spec_with_checkpoint(
       module_spec,
-      checkpoint_path=tf.train.latest_checkpoint(ckpt_dir),
+      checkpoint_path=checkpoint_path,
       export_path=export_dir,
       scope_prefix="")
 
@@ -135,7 +142,10 @@ def main(_):
   export_dir = FLAGS.export_dir or os.path.join(ckpt_dir, "export")
 
   if FLAGS.export_as_tfhub:
-    export_as_tfhub_module(hparams, problem, ckpt_dir, export_dir)
+    checkpoint_path = tf.train.latest_checkpoint(ckpt_dir)
+    decode_hparams = decoding.decode_hparams(FLAGS.decode_hparams)
+    export_as_tfhub_module(FLAGS.model, hparams, decode_hparams, problem,
+                           checkpoint_path, export_dir)
     return
 
   run_config = t2t_trainer.create_run_config(hparams)

From cc34ec67b0d483c81fad2b7d49f1adf98b300c39 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 27 Sep 2018 18:37:18 -0700
Subject: [PATCH 0936/2720] Drop TF 1.9 testing, add 1.11

PiperOrigin-RevId: 214871981
---
 .travis.yml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 628f13548..fd91c80fa 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -13,20 +13,20 @@ env:
     - T2T_PROBLEM=algorithmic_reverse_binary40_test
     - T2T_DATA_DIR=/tmp/t2t-data
     - T2T_TRAIN_DIR=/tmp/t2t-train
-    - TF_LATEST="1.10.*"
+    - TF_LATEST="1.11.*"
     # This is necessary to have gsutil work with Python 2.7
     - BOTO_CONFIG=/dev/null
   matrix:
     # We test against recent versions of TensorFlow and tf-nightly.
     # If updating, also update TF_LATEST above
-    - TF_VERSION="1.9.*"
     - TF_VERSION="1.10.*"
+    - TF_VERSION="1.11.*"
     - TF_VERSION="tf-nightly"
 matrix:
   exclude:
     # We test against all versions in Python 2 but only the latest in Python 3
     - python: "3.6"
-      env: TF_VERSION="1.9.*"
+      env: TF_VERSION="1.10.*"
     - python: "3.6"
       env: TF_VERSION="tf-nightly"
 before_install:

From 4c8216d568abf0e8bbc198b149a9840aee19fce9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 28 Sep 2018 09:27:58 -0700
Subject: [PATCH 0937/2720] internal merge of PR #1098

PiperOrigin-RevId: 214947171
---
 tensor2tensor/layers/common_video.py | 40 +++++++++++++++++++++++++---
 1 file changed, 37 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 6ca3bc33c..d6cdeefc0 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -562,9 +562,14 @@ def __init__(self, fps, file_format="gif"):
     self.fps = fps
     self.file_format = file_format
     self.proc = None
+    self._out_chunks = []
+    self._err_chunks = []
+    self._out_thread = None
+    self._err_thread = None
 
   def __init_ffmpeg(self, image_shape):
     """Initializes ffmpeg to write frames."""
+    import itertools  # pylint: disable=g-import-not-at-top
     from subprocess import Popen, PIPE  # pylint: disable=g-import-not-at-top,g-multiple-import
     ffmpeg = "ffmpeg"
     height, width, channels = image_shape
@@ -576,13 +581,36 @@ def __init_ffmpeg(self, image_shape):
         "-s", "%dx%d" % (width, height),
         "-pix_fmt", {1: "gray", 3: "rgb24"}[channels],
         "-i", "-",
-        "-filter_complex", "[0:v]split[x][z];[z]palettegen[y];[x][y]paletteuse",
+        "-filter_complex", "[0:v]split[x][z];[x]fifo[w];[z]palettegen,fifo[y];"
+                           "[w][y]paletteuse,fifo",
         "-r", "%.02f" % self.fps,
         "-f", self.file_format,
         "-qscale", "0",
         "-"
     ]
-    self.proc = Popen(self.cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
+    self.proc = Popen(
+        self.cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE, bufsize=-1
+    )
+    (self._out_thread, self._err_thread) = itertools.starmap(
+        self._start_reader_thread, [
+            (self.proc.stdout, self._out_chunks),
+            (self.proc.stderr, self._err_chunks)
+        ]
+    )
+
+  def _start_reader_thread(self, stream, chunks):
+    """TODO(koz4k): Write a docstring."""
+    import io  # pylint: disable=g-import-not-at-top
+    import threading  # pylint: disable=g-import-not-at-top
+    def target():
+      while True:
+        chunk = stream.read(io.DEFAULT_BUFFER_SIZE)
+        if not chunk:
+          break
+        chunks.append(chunk)
+    thread = threading.Thread(target=target)
+    thread.start()
+    return thread
 
   def write(self, frame):
     if self.proc is None:
@@ -594,9 +622,15 @@ def write_multi(self, frames):
       self.write(frame)
 
   def finish(self):
+    """TODO(koz4k): Write a docstring."""
     if self.proc is None:
       return None
-    out, err = self.proc.communicate()
+    self.proc.stdin.close()
+    for thread in (self._out_thread, self._err_thread):
+      thread.join()
+    (out, err) = [
+        b"".join(chunks) for chunks in (self._out_chunks, self._err_chunks)
+    ]
     if self.proc.returncode:
       err = "\n".join([" ".join(self.cmd), err.decode("utf8")])
       raise IOError(err)

From ba246a21188d300379394ef14402953d437837f5 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 28 Sep 2018 18:32:17 +0200
Subject: [PATCH 0938/2720] Remove a deadlock during saving debug video (#1098)

* Prevent FFMPEG buffer queue overflow when saving debug video

* Read FFMPEG output asynchronously to prevent deadlock due to pipe buffer
overflow

From bf6c53b49e359291aabb29c00576ea5ec32c7d3d Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 28 Sep 2018 15:31:09 -0700
Subject: [PATCH 0939/2720] Add a placeholder to link to the official ASR
 tutorial

PiperOrigin-RevId: 215006189
---
 docs/tutorials/asr_with_transformer.md | 4 ++++
 1 file changed, 4 insertions(+)
 create mode 100644 docs/tutorials/asr_with_transformer.md

diff --git a/docs/tutorials/asr_with_transformer.md b/docs/tutorials/asr_with_transformer.md
new file mode 100644
index 000000000..ee8aa0f53
--- /dev/null
+++ b/docs/tutorials/asr_with_transformer.md
@@ -0,0 +1,4 @@
+# Automated Speech Recognition with the Transformer model
+
+See the
+[official tutorial](https://cloud.google.com/tpu/docs/tutorials/automated-speech-recognition).

From 3d75129e85937adc28cdc491b195dcec9bc24019 Mon Sep 17 00:00:00 2001
From: Jonathan Hseu <jhseu@google.com>
Date: Fri, 28 Sep 2018 18:41:31 -0700
Subject: [PATCH 0940/2720] Fix the call to create a TPU variable.

PiperOrigin-RevId: 215027511
---
 tensor2tensor/mesh_tensorflow/tpu_variables.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/mesh_tensorflow/tpu_variables.py b/tensor2tensor/mesh_tensorflow/tpu_variables.py
index 0bc3e2c04..7da863589 100644
--- a/tensor2tensor/mesh_tensorflow/tpu_variables.py
+++ b/tensor2tensor/mesh_tensorflow/tpu_variables.py
@@ -68,7 +68,7 @@ def handle(self):
     if tpu_context is None:
       return self._primary_var.handle
 
-    return tpu_context.get_replicated_var_handle(self)
+    return tpu_context.get_replicated_var_handle(self._name, self._vars)
 
   @contextlib.contextmanager
   def _assign_dependencies(self):

From f6c91e93052616b7025f0c64f44db62fe4a49296 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 1 Oct 2018 10:01:58 -0700
Subject: [PATCH 0941/2720] Add an option to disable host call on TPU.

PiperOrigin-RevId: 215231453
---
 tensor2tensor/layers/common_hparams.py | 4 ++++
 tensor2tensor/utils/t2t_model.py       | 7 ++++++-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 0d5c5895f..cdb872adf 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -282,6 +282,10 @@ def basic_params1():
       heads_share_relative_embedding=False,
       # If relative embedding terms are added to values too.
       add_relative_to_values=False,
+      # If enable the host_call which is executed every training step.
+      # There could be a performance drop if host_call function is slow and
+      # cannot keep up with the TPU-side computation.
+      tpu_enable_host_call=False,
   )
 
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 723cdf558..a2de8d6f6 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1347,8 +1347,13 @@ def scaffold_fn():
       else:
         scaffold_fn = None
 
-      host_call = _create_host_call(self.hparams.model_dir)
+      if self.hparams.tpu_enable_host_call:
+        host_call = _create_host_call(self.hparams.model_dir)
+      else:
+        host_call = None
+
       remove_summaries()
+
       return tf.contrib.tpu.TPUEstimatorSpec(
           tf.estimator.ModeKeys.TRAIN,
           loss=loss,

From ed3548f2d470108dcdd664d2667e7b81a867c3e1 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 1 Oct 2018 11:35:36 -0700
Subject: [PATCH 0942/2720] Fix bug in tag indexing of video summaries

PiperOrigin-RevId: 215250889
---
 tensor2tensor/data_generators/video_utils.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 80852441f..d85fed8ef 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -86,8 +86,9 @@ def display_video_hooks(hook_args):
 
       # Concat across width.
       all_frames = np.concatenate((input_frames, output_frames), axis=1)
+      tag = "input/output/decode_%d_sample_%d" % (decode_ind, ind)
       frame_by_frame_summ = image_utils.image_to_tf_summary_value(
-          all_frames, tag="input/output/decode_%d" % ind)
+          all_frames, tag=tag)
       all_summaries.append(frame_by_frame_summ)
   return all_summaries
 

From b4fd0a4d5848d50650c135eb2b102c5a4376d604 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 1 Oct 2018 11:58:30 -0700
Subject: [PATCH 0943/2720] fixing stochastic data set.

PiperOrigin-RevId: 215254797
---
 tensor2tensor/data_generators/video_generated.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 14fe3d25d..3a3e70db3 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -23,7 +23,6 @@
 import numpy as np
 
 from tensor2tensor.data_generators import video_utils
-from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -66,10 +65,16 @@ def video_length(self):
   def random_skip(self):
     return False
 
+  @property
+  def only_keep_videos_from_0th_frame(self):
+    return True
+
+  @property
+  def use_not_breaking_batching(self):
+    return True
+
   def eval_metrics(self):
-    eval_metrics = [metrics.Metrics.ACC, metrics.Metrics.ACC_PER_SEQ,
-                    metrics.Metrics.IMAGE_RMSE]
-    return eval_metrics
+    return []
 
   @property
   def extra_reading_spec(self):
@@ -87,7 +92,6 @@ def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.input_modality = {
         "inputs": ("video", 256),
-        "input_frame_number": ("symbol:identity", 1)
     }
     p.target_modality = {
         "targets": ("video", 256),

From 44758c41f3f999db9498f262565ad016a054264a Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Mon, 1 Oct 2018 22:54:22 +0200
Subject: [PATCH 0944/2720] Turn of _forced_collect_level override for AE

---
 tensor2tensor/rl/trainer_model_based.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index a4eaa2151..90d94e286 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -120,7 +120,7 @@ def generate_real_env_data(problem_name, agent_policy_path, hparams, data_dir,
           hparams.num_real_env_frames / (hparams.epochs * (1. - 1./11.)))
     gym_problem.settable_num_steps = env_steps_per_epoch
     gym_problem.settable_eval_phase = eval_phase
-    if real_reward:
+    if real_reward and autoencoder_path is None:
       gym_problem._forced_collect_level = 1  # pylint: disable=protected-access
     gym_problem.generate_data(data_dir, tmp_dir)
     mean_reward = None

From b5ddcd5f1fe0a717fe4ed95a22ab0557c8401774 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 1 Oct 2018 16:22:58 -0700
Subject: [PATCH 0945/2720] Internal change

PiperOrigin-RevId: 215300569
---
 tensor2tensor/models/research/glow.py         |   4 +-
 tensor2tensor/models/research/glow_ops.py     | 107 +++++++++++++-----
 .../models/research/glow_ops_test.py          |  45 ++++++--
 3 files changed, 119 insertions(+), 37 deletions(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 076cafe9c..428c5fad8 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -99,7 +99,7 @@ def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
     var_scope = tf.variable_scope("glow/body", reuse=True)
     # If eps=None, images are sampled from the prior.
     with arg_scope(ops, init=False), var_scope:
-      predictions, _, _ = glow_ops.encoder_decoder(
+      predictions, _, _, _ = glow_ops.encoder_decoder(
           "codec", self.z_sample, self.hparams, eps=None, reverse=True)
 
     return self.scale(predictions)
@@ -131,7 +131,7 @@ def body(self, features):
     init_op = tf.logical_and(tf.equal(global_step, 0), self.is_training)
     ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
     with arg_scope(ops, init=init_op):
-      self.z, encoder_objective, self.eps, _ = glow_ops.encoder_decoder(
+      self.z, encoder_objective, self.eps, _, _ = glow_ops.encoder_decoder(
           "codec", x, self.hparams, eps=None, reverse=False)
       objective += encoder_objective
 
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index dd6098f4e..05b8218f5 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -22,6 +22,7 @@
 import numpy as np
 import scipy
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
 import tensorflow as tf
 
 arg_scope = tf.contrib.framework.arg_scope
@@ -54,7 +55,7 @@ def get_cond_latents_at_level(cond_latents, level, hparams):
   if cond_latents:
     if hparams.latent_dist_encoder == "conv_net":
       return [cond_latent[level] for cond_latent in cond_latents]
-    elif hparams.latent_dist_encoder == "pointwise":
+    elif hparams.latent_dist_encoder in ["pointwise", "conv_lstm"]:
       return cond_latents[level]
 
 
@@ -546,12 +547,35 @@ def merge_level_and_latent_dist(level_dist, latent_dist,
 
 
 @add_arg_scope
-def compute_prior(name, z, latent, hparams):
-  """Distribution conditioned on both z and latent."""
+def compute_prior(name, z, latent, hparams, state=None):
+  """Distribution on z_t conditioned on z_{t-1} and latent.
+
+  Args:
+    name: variable scope.
+    z: 4-D Tensor.
+    latent: optional,
+            if hparams.latent_dist_encoder == "pointwise", this is a list
+            of 4-D Tensors of length hparams.num_cond_latents.
+            else, this is just a 4-D Tensor
+            The first-three dimensions of the latent should be the same as z.
+    hparams: next_frame_glow_hparams.
+    state: tf.contrib.rnn.LSTMStateTuple.
+           the current state of a LSTM used to model the distribution. Used
+           only if hparams.latent_dist_encoder = "conv_lstm".
+  Returns:
+    prior_dist: instance of tf.distributions.Normal
+    state: Returns updated state.
+  Raises:
+    ValueError: If hparams.latent_dist_encoder is "pointwise" and if the shape
+                of latent is different from z.
+  """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     prior_dist = tensor_to_dist("level_prior", z, architecture="single_conv")
+
+    # TODO(mechcoder) Refactor into separate sub-functions.
     if latent is not None:
-      latent_dist_encoder = hparams.latent_dist_encoder
+      latent_dist_encoder = hparams.get("latent_dist_encoder", None)
+      latent_skip = hparams.get("latent_skip", False)
       if latent_dist_encoder == "pointwise":
         merge_std = hparams.level_scale
         latent_shape = common_layers.shape_list(latent)
@@ -571,59 +595,74 @@ def compute_prior(name, z, latent, hparams):
             architecture=hparams.latent_architecture,
             depth=hparams.latent_encoder_depth,
             pre_output_channels=hparams.latent_pre_output_channels)
-        latent_skip = hparams.get("latent_skip", False)
         if latent_skip:
           prior_dist = tf.distributions.Normal(
               prior_dist.loc + latent[-1], prior_dist.scale)
+      elif latent_dist_encoder == "conv_lstm":
+        output_channels = common_layers.shape_list(z)[-1]
+        latent_stack = tf.concat((prior_dist.loc, latent), axis=-1)
+        _, state = common_video.conv_lstm_2d(
+            latent_stack, state, output_channels, kernel_size=3,
+            name="conv_lstm")
+        prior_dist = tensor_to_dist(
+            "state_to_dist", state.h, output_channels=output_channels)
+        if latent_skip:
+          prior_dist = tf.distributions.Normal(
+              prior_dist.loc + latent, prior_dist.scale)
       tf.summary.histogram("split_prior_mean", prior_dist.loc)
       tf.summary.histogram("split_prior_scale", prior_dist.scale)
 
-  return prior_dist
+  return prior_dist, state
 
 
 @add_arg_scope
 def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
-          hparams=None):
+          hparams=None, state=None):
   """Splits / concatenates x into x1 and x2 across number of channels.
 
   For the forward pass, x2 is assumed be gaussian,
-  i.e P(x2 | x1) ~ N(mu(x1), sigma(x1)) where mu and sigma are the outputs of
-  a one-layer network. For the reverse pass, x2 is determined
-  from mu(x1) and sigma(x1). This is deterministic/stochastic depending on
-  whether eps is provided.
+  i.e P(x2 | x1) ~ N(mu, sigma) where mu and sigma are the outputs of
+  a network conditioned on x1 and optionally on cond_latents.
+  For the reverse pass, x2 is determined from mu(x1) and sigma(x1).
+  This is deterministic/stochastic depending on whether eps is provided.
 
   Args:
-    name:
-    x:
+    name: variable scope.
+    x: 4-D Tensor, shape (NHWC).
     reverse: Forward or reverse pass.
     eps: If eps is provided, x2 is set to be
-    eps_std: Sample x2.
+    eps_std: Sample x2 with the provided eps_std.
     cond_latents: optionally condition x2 on cond_latents.
     hparams: next_frame_glow hparams.
+    state: tf.contrib.rnn.LSTMStateTuple. Current state of the LSTM over z_2.
+           Used only when hparams.latent_dist_encoder == "conv_lstm"
 
   Returns:
   Raises:
     ValueError: If latent is provided and shape is not equal to NHW(C/2)
                 where (NHWC) is the size of x.
   """
+  # TODO(mechcoder) Change the return type to be a dict.
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     if not reverse:
       x1, x2 = tf.split(x, num_or_size_splits=2, axis=-1)
 
       # objective: P(x2|x1) ~N(x2 ; NN(x1))
-      prior_dist = compute_prior("prior_on_z2", x1, cond_latents, hparams)
+      prior_dist, state = compute_prior(
+          "prior_on_z2", x1, cond_latents, hparams, state=state)
       logpb = tf.reduce_sum(prior_dist.log_prob(x2), axis=[1, 2, 3])
       eps = get_eps(prior_dist, x2)
-      return x1, logpb, eps, x2
+      return x1, logpb, eps, x2, state
     else:
-      prior_dist = compute_prior("prior_on_z2", x, cond_latents, hparams)
+      prior_dist, state = compute_prior(
+          "prior_on_z2", x, cond_latents, hparams, state=state)
       if eps is not None:
         x2 = set_eps(prior_dist, eps)
       elif eps_std is not None:
         x2 = eps_std * tf.random_normal(common_layers.shape_list(x))
       else:
         x2 = prior_dist.sample()
-      return tf.concat([x, x2], 3), x2
+      return tf.concat([x, x2], 3), x2, state
 
 
 @add_arg_scope
@@ -755,10 +794,14 @@ def uniform_binning_correction(x, n_bits=8):
 
 @add_arg_scope
 def encoder_decoder(name, x, hparams, eps=None, reverse=False,
-                    cond_latents=None):
+                    cond_latents=None, states=None):
   """Glow encoder-decoder. n_levels of (Squeeze + Flow + Split.) operations."""
+  # TODO(mechcoder) Change return_type to a dict to be backward compatible.
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
 
+    if states and len(states) != hparams.n_levels - 1:
+      raise ValueError("Expected length of states to be %d, got %d" %
+                       (hparams.n_levels - 1, len(states)))
     if eps and len(eps) != hparams.n_levels - 1:
       raise ValueError("Expected length of eps to be %d, got %d" %
                        (hparams.n_levels - 1, len(eps)))
@@ -767,6 +810,7 @@ def encoder_decoder(name, x, hparams, eps=None, reverse=False,
     objective = 0.0
     all_eps = []
     all_latents = []
+    new_states = []
 
     if not reverse:
       # Squeeze + Flow + Split
@@ -778,15 +822,20 @@ def encoder_decoder(name, x, hparams, eps=None, reverse=False,
 
         if level < hparams.n_levels - 1:
 
+          curr_state = None
+          if states:
+            curr_state = states[level]
+
           curr_cond_latents = get_cond_latents_at_level(
               cond_latents, level, hparams)
-          x, obj, eps, z = split("split_%d" % level, x, reverse=False,
-                                 cond_latents=curr_cond_latents,
-                                 hparams=hparams)
+          x, obj, eps, z, state = split("split_%d" % level, x, reverse=False,
+                                        cond_latents=curr_cond_latents,
+                                        hparams=hparams, state=curr_state)
           objective += obj
           all_eps.append(eps)
           all_latents.append(z)
-      return x, objective, all_eps, all_latents
+          new_states.append(state)
+      return x, objective, all_eps, all_latents, new_states
 
     else:
       for level in reversed(range(hparams.n_levels)):
@@ -796,15 +845,21 @@ def encoder_decoder(name, x, hparams, eps=None, reverse=False,
           if eps:
             curr_eps = eps[level]
 
+          curr_state = None
+          if states:
+            curr_state = states[level]
+
           curr_cond_latents = get_cond_latents_at_level(
               cond_latents, level, hparams)
 
-          x, latent = split("split_%d" % level, x, eps=curr_eps, reverse=True,
-                            cond_latents=curr_cond_latents, hparams=hparams)
+          x, latent, state = split("split_%d" % level, x, eps=curr_eps,
+                                   reverse=True, cond_latents=curr_cond_latents,
+                                   hparams=hparams, state=curr_state)
+          new_states.append(state)
           all_latents.append(latent)
 
         x, obj = revnet(
             "revnet_%d" % level, x, hparams=hparams, reverse=True)
         objective += obj
         x = squeeze("squeeze_%d" % level, x, reverse=True)
-      return x, objective, all_latents[::-1]
+      return x, objective, all_latents[::-1], new_states[::-1]
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 696b6019e..0d97a52db 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -145,8 +145,8 @@ def test_tensor_to_dist(self):
   def test_split(self):
     with tf.Graph().as_default():
       x = tf.random_uniform(shape=(16, 5, 5, 32))
-      x_inv, _, eps, z = glow_ops.split("split", x)
-      x_inv_inv, _ = glow_ops.split("split", x_inv, reverse=True, eps=eps)
+      x_inv, _, eps, z, _ = glow_ops.split("split", x)
+      x_inv_inv, _, _ = glow_ops.split("split", x_inv, reverse=True, eps=eps)
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())
         x_inv_np, diff, z_np = session.run([x_inv, x - x_inv_inv, z])
@@ -179,9 +179,9 @@ def test_encoder_decoder(self):
       hparams.depth = 2
 
       x = tf.random_uniform(shape=(16, 64, 64, 4), seed=0)
-      x_inv, _, eps, z_levels = glow_ops.encoder_decoder(
+      x_inv, _, eps, z_levels, _ = glow_ops.encoder_decoder(
           "encoder_decoder", x, hparams, reverse=False)
-      x_inv_inv, _, z_inv_levels = glow_ops.encoder_decoder(
+      x_inv_inv, _, z_inv_levels, _ = glow_ops.encoder_decoder(
           "encoder_decoder", x_inv, hparams, eps=eps, reverse=True)
 
       with tf.Session() as session:
@@ -218,7 +218,7 @@ def test_encoder_decoder_practical_usage(self):
 
       ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
       with arg_scope(ops, init=True):
-        x_inv, _, _, _ = glow_ops.encoder_decoder(
+        x_inv, _, _, _, _ = glow_ops.encoder_decoder(
             "revnet", x_t, hparams, reverse=False)
       curr_dir = tempfile.mkdtemp()
       model_path = os.path.join(curr_dir, "model")
@@ -235,9 +235,9 @@ def test_encoder_decoder_practical_usage(self):
       x_t = tf.convert_to_tensor(x_rand)
       ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
       with arg_scope(ops, init=False):
-        x_inv2, _, all_eps, _ = glow_ops.encoder_decoder(
+        x_inv2, _, all_eps, _, _ = glow_ops.encoder_decoder(
             "revnet", x_t, hparams, reverse=False)
-        x_inv_inv_, _, _ = glow_ops.encoder_decoder(
+        x_inv_inv_, _, _, _ = glow_ops.encoder_decoder(
             "revnet", x_inv2, hparams, eps=all_eps, reverse=True)
 
       with tf.Session() as session:
@@ -278,10 +278,10 @@ def check_split_latent_conditioning(self, merge_std):
       # x2 ~ N(scale * latent, 1.0) where initial scale is 1.0
       exp_x2 = x_rand[:, :, :, 16:]
       exp_eps = x_rand[:, :, :, 16:] - latent_rand
-      x_inv, _, eps, x2_t = glow_ops.split(
+      x_inv, _, eps, x2_t, _ = glow_ops.split(
           merge_std, x_t, cond_latents=latent_t, hparams=hparams)
       # Test reversibility.
-      x_inv_inv, _ = glow_ops.split(
+      x_inv_inv, _, _ = glow_ops.split(
           merge_std, x_inv, cond_latents=latent_t, eps=eps, reverse=True,
           hparams=hparams)
       with tf.Session() as sess:
@@ -295,6 +295,33 @@ def test_split_latent_conditioning(self):
     for merge_std in ["normal", "prev_level", "prev_step"]:
       self.check_split_latent_conditioning(merge_std)
 
+  def test_latent_dist_encoder_lstm(self):
+    with tf.Graph().as_default():
+      rng = np.random.RandomState(0)
+      # Initialize x, latent, state.
+      x_rand = rng.randn(12, 32, 32, 16).astype(np.float32)
+      latent_rand = rng.randn(12, 32, 32, 16).astype(np.float32)
+      state_rand = rng.randn(12, 32, 32, 16).astype(np.float32)
+      x_t = tf.convert_to_tensor(x_rand)
+      latent_t = tf.convert_to_tensor(latent_rand)
+      state_t = tf.convert_to_tensor(state_rand)
+      init_state = tf.contrib.rnn.LSTMStateTuple(state_t, state_t)
+      hparams = glow.glow_hparams()
+      hparams.add_hparam("latent_dist_encoder", "conv_lstm")
+      hparams.add_hparam("latent_skip", True)
+
+      prior_dist, new_state = glow_ops.compute_prior(
+          "lstm_prior", x_t, latent=latent_t, hparams=hparams, state=init_state)
+      with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+        # Test initialization (mu, sigma) = (z, 1.0)
+        ops = [prior_dist.loc, prior_dist.scale, new_state.h - init_state.h]
+        mean, scale, diff_np = sess.run(ops)
+        self.assertTrue(np.allclose(latent_rand - mean, 0.0))
+        self.assertTrue(np.allclose(scale, 1.0))
+        # State update.
+        self.assertFalse(np.allclose(diff_np, 0.0))
+
 
 if __name__ == "__main__":
   tf.test.main()

From b4d67d2344eba6495c484079ba997b5109f80517 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 1 Oct 2018 22:06:51 -0700
Subject: [PATCH 0946/2720] remove unused variable.

PiperOrigin-RevId: 215334646
---
 tensor2tensor/layers/modalities.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 316aa96c2..1325adf71 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -114,7 +114,6 @@ def bottom_simple(self, x, name, reuse):
       return ret
 
   def bottom(self, x):
-    self._bottom_was_called = True
     if (self._model_hparams.shared_embedding_and_softmax_weights or
         self._model_hparams.get("shared_embedding")):
       return self.bottom_simple(x, "shared", reuse=None)

From c4b34e1869dede0798818b5b9f33a960a2965fee Mon Sep 17 00:00:00 2001
From: Derek Murray <mrry@google.com>
Date: Mon, 1 Oct 2018 23:39:43 -0700
Subject: [PATCH 0947/2720] Internal change

PiperOrigin-RevId: 215340476
---
 tensor2tensor/data_generators/video_utils.py       | 4 ++--
 tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py | 3 +--
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index d85fed8ef..08c85c9c7 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -370,8 +370,8 @@ def check_integrity_and_batch(*datasets):
         self.settable_use_not_breaking_batching):
       batch_dataset = avoid_break_batching(preprocessed_dataset)
     else:
-      batch_dataset = preprocessed_dataset.apply(
-          tf.contrib.data.batch_and_drop_remainder(num_frames))
+      batch_dataset = preprocessed_dataset.batch(num_frames,
+                                                 drop_remainder=True)
     dataset = batch_dataset.map(features_from_batch)
     if self.shuffle and interleave and mode == tf.estimator.ModeKeys.TRAIN:
       dataset = dataset.shuffle(hparams.get("shuffle_buffer_size", 128))
diff --git a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
index cfcd893d2..e540dc666 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
@@ -93,8 +93,7 @@ def __call__(self, params):
 
     ds = Dataset.from_tensor_slices((self._images, self._labels)).repeat()
 
-    dataset = ds.apply(
-        tf.contrib.data.batch_and_drop_remainder(batch_size)).prefetch(2)
+    dataset = ds.batch(batch_size, drop_remainder=True).prefetch(2)
 
     return dataset
 

From 267cec141e78c7e6bb9541fd8a7b13ed06d4a104 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 1 Oct 2018 23:49:40 -0700
Subject: [PATCH 0948/2720] Internal change

PiperOrigin-RevId: 215340971
---
 tensor2tensor/models/__init__.py              |   1 +
 tensor2tensor/models/video/basic_recurrent.py | 181 ++++++++++++++++++
 .../models/video/basic_recurrent_test.py      |  37 ++++
 3 files changed, 219 insertions(+)
 create mode 100644 tensor2tensor/models/video/basic_recurrent.py
 create mode 100644 tensor2tensor/models/video/basic_recurrent_test.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index f57f27dd3..40c955e78 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -61,6 +61,7 @@
 from tensor2tensor.models.research import vqa_self_attention
 
 from tensor2tensor.models.video import basic_deterministic
+from tensor2tensor.models.video import basic_recurrent
 from tensor2tensor.models.video import basic_stochastic
 from tensor2tensor.models.video import emily
 from tensor2tensor.models.video import epva
diff --git a/tensor2tensor/models/video/basic_recurrent.py b/tensor2tensor/models/video/basic_recurrent.py
new file mode 100644
index 000000000..4a60c9f44
--- /dev/null
+++ b/tensor2tensor/models/video/basic_recurrent.py
@@ -0,0 +1,181 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Basic recurrent models for testing simple tasks."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
+from tensor2tensor.models.video import basic_deterministic
+from tensor2tensor.models.video import basic_deterministic_params
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+tfl = tf.layers
+tfcl = tf.contrib.layers
+
+
+@registry.register_model
+class NextFrameBasicRecurrent(basic_deterministic.NextFrameBasicDeterministic):
+  """Basic next-frame recurrent model."""
+
+  def predict_next_frame(self, frame, action, lstm_states):
+    hparams = self.hparams
+    filters = hparams.hidden_size
+    kernel1, kernel2 = (3, 3), (4, 4)
+    lstm_func = common_video.conv_lstm_2d
+
+    # Embed the inputs.
+    inputs_shape = common_layers.shape_list(frame)
+    # Using non-zero bias initializer below for edge cases of uniform inputs.
+    x = tf.layers.dense(
+        frame, filters, name="inputs_embed",
+        bias_initializer=tf.random_normal_initializer(stddev=0.01))
+    x = common_attention.add_timing_signal_nd(x)
+
+    # Down-stride.
+    layer_inputs = [x]
+    for i in range(hparams.num_compress_steps):
+      with tf.variable_scope("downstride%d" % i):
+        layer_inputs.append(x)
+        x = common_layers.make_even_size(x)
+        if i < hparams.filter_double_steps:
+          filters *= 2
+        x = common_attention.add_timing_signal_nd(x)
+        x = tf.layers.conv2d(x, filters, kernel2, activation=common_layers.belu,
+                             strides=(2, 2), padding="SAME")
+        x = common_layers.layer_norm(x)
+
+    # Add embedded action if present.
+    if self.has_action:
+      x = self.inject_additional_input(
+          x, action, "action_enc", hparams.action_injection)
+
+    x, extra_loss = self.inject_latent(x, self.features, filters)
+
+    # LSTM layers
+    for j in range(hparams.num_lstm_layers):
+      x, lstm_states[j] = lstm_func(x, lstm_states[j], hparams.num_lstm_filters)
+
+    # Run a stack of convolutions.
+    for i in range(hparams.num_hidden_layers):
+      with tf.variable_scope("layer%d" % i):
+        y = tf.nn.dropout(x, 1.0 - hparams.dropout)
+        y = tf.layers.conv2d(y, filters, kernel1, activation=common_layers.belu,
+                             strides=(1, 1), padding="SAME")
+        if i == 0:
+          x = y
+        else:
+          x = common_layers.layer_norm(x + y)
+
+    # Up-convolve.
+    layer_inputs = list(reversed(layer_inputs))
+    for i in range(hparams.num_compress_steps):
+      with tf.variable_scope("upstride%d" % i):
+        if self.has_action:
+          x = self.inject_additional_input(
+              x, action, "action_enc", hparams.action_injection)
+        if i >= hparams.num_compress_steps - hparams.filter_double_steps:
+          filters //= 2
+        x = tf.layers.conv2d_transpose(
+            x, filters, kernel2, activation=common_layers.belu,
+            strides=(2, 2), padding="SAME")
+        y = layer_inputs[i]
+        shape = common_layers.shape_list(y)
+        x = x[:, :shape[1], :shape[2], :]
+        x = common_layers.layer_norm(x + y)
+        x = common_attention.add_timing_signal_nd(x)
+
+    # Cut down to original size.
+    x = x[:, :inputs_shape[1], :inputs_shape[2], :]
+    if self.is_per_pixel_softmax:
+      x = tf.layers.dense(x, hparams.problem.num_channels * 256, name="logits")
+    else:
+      x = tf.layers.dense(x, hparams.problem.num_channels, name="logits")
+
+    # Reward prediction if needed.
+    reward_pred = 0.0
+    if self.has_reward:
+      reward_pred = tf.expand_dims(  # Add a fake channels dim.
+          tf.reduce_mean(x, axis=[1, 2], keepdims=True), axis=3)
+    return x, reward_pred, extra_loss, lstm_states
+
+  def body(self, features):
+    hparams = self.hparams
+    self.has_action = "input_action" in features
+    self.has_reward = "target_reward" in features
+    # dirty hack to enable the latent tower
+    self.features = features
+
+    # Split inputs and targets into lists.
+    input_frames = tf.unstack(features["inputs"], axis=1)
+    target_frames = tf.unstack(features["targets"], axis=1)
+    all_frames = input_frames + target_frames
+    if self.has_action:
+      input_actions = tf.unstack(features["input_action"], axis=1)
+      target_actions = tf.unstack(features["target_action"], axis=1)
+      all_actions = input_actions + target_actions
+
+    res_frames, sampled_frames, sampled_frames_raw, res_rewards = [], [], [], []
+    lstm_states = [None] * hparams.num_lstm_layers
+    extra_loss = 0.0
+
+    num_frames = len(all_frames)
+    for i in range(num_frames - 1):
+      frame = all_frames[i]
+      action = all_actions[i] if self.has_action else None
+
+      # Run model.
+      with tf.variable_scope("recurrent_model", reuse=tf.AUTO_REUSE):
+        func_out = self.predict_next_frame(frame, action, lstm_states)
+        res_frame, res_reward, res_extra_loss, lstm_states = func_out
+        res_frames.append(res_frame)
+        res_rewards.append(res_reward)
+        extra_loss += res_extra_loss
+
+      sampled_frame_raw = self.get_sampled_frame(res_frame)
+      sampled_frames_raw.append(sampled_frame_raw)
+      # TODO(lukaszkaiser): this should be consistent with modality.bottom()
+      sampled_frame = common_layers.standardize_images(sampled_frame_raw)
+      sampled_frames.append(sampled_frame)
+
+      # Only for Softmax loss: sample next frame so we can keep iterating.
+      if self.is_predicting and i >= hparams.video_num_input_frames:
+        all_frames[i+1] = sampled_frame
+
+    # Concatenate results and return them.
+    output_frames = res_frames[hparams.video_num_input_frames-1:]
+    frames = tf.stack(output_frames, axis=1)
+
+    if not self.has_reward:
+      return frames, extra_loss
+    rewards = tf.concat(res_rewards[hparams.video_num_input_frames-1:], axis=1)
+    return {"targets": frames, "target_reward": rewards}, extra_loss
+
+
+@registry.register_hparams
+def next_frame_basic_recurrent():
+  """Basic 2-frame recurrent model with stochastic tower."""
+  hparams = basic_deterministic_params.next_frame_basic_deterministic()
+  hparams.video_num_input_frames = 4
+  hparams.video_num_target_frames = 4
+  hparams.add_hparam("num_lstm_layers", 1)
+  hparams.add_hparam("num_lstm_filters", 8)
+  return hparams
diff --git a/tensor2tensor/models/video/basic_recurrent_test.py b/tensor2tensor/models/video/basic_recurrent_test.py
new file mode 100644
index 000000000..5b1c9d445
--- /dev/null
+++ b/tensor2tensor/models/video/basic_recurrent_test.py
@@ -0,0 +1,37 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Basic tests for basic deterministic model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.models.video import basic_recurrent
+from tensor2tensor.models.video import tests_utils
+
+import tensorflow as tf
+
+
+class NextFrameTest(tests_utils.BaseNextFrameTest):
+
+  def testBasicDeterministic(self):
+    self.TestOnVariousInputOutputSizes(
+        basic_recurrent.next_frame_basic_recurrent(),
+        basic_recurrent.NextFrameBasicRecurrent,
+        256,
+        False)
+
+if __name__ == "__main__":
+  tf.test.main()

From d998beb36e82fd27284a42e8b25b3be897a2a4ec Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Tue, 2 Oct 2018 00:47:04 -0700
Subject: [PATCH 0949/2720] Add the resnet model for mesh_tensorflow.

PiperOrigin-RevId: 215345739
---
 tensor2tensor/mesh_tensorflow/mtf_layers.py |  60 +++
 tensor2tensor/mesh_tensorflow/mtf_resnet.py | 476 ++++++++++++++++++++
 tensor2tensor/models/__init__.py            |   1 +
 3 files changed, 537 insertions(+)
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_resnet.py

diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers.py b/tensor2tensor/mesh_tensorflow/mtf_layers.py
index 26f2cc31c..4ed558596 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers.py
@@ -103,6 +103,66 @@ def layer_norm(x, dim, epsilon=1e-6, name="layer_prepostprocess"):
     return norm_x * scale + bias
 
 
+def batch_norm(x, is_training, momentum, epsilon=1e-9, name=None):
+  """Batch normalization.
+
+  Args:
+    x: a mtf.Tensor whose shape contains [batch_dim, ..., dim]
+    is_training: a boolean, whether mode is training.
+    momentum: a floating point number, specifying batch norm decay value.
+    epsilon: a floating point number.
+    name: a string. variable scope.
+
+  Returns:
+    a mtf.Tensor with same shape as x.
+  """
+  with tf.variable_scope(name, default_name="batch_norm", values=[x]):
+    batch_dim = x.shape.dims[0]
+    reduced_shape = x.shape - batch_dim
+    scale = mtf.get_variable(
+        x.mesh,
+        "batch_norm_scale",
+        mtf.Shape([batch_dim]),
+        initializer=tf.ones_initializer(),
+        activation_dtype=x.dtype)
+    bias = mtf.get_variable(
+        x.mesh,
+        "batch_norm_bias",
+        mtf.Shape([batch_dim]),
+        initializer=tf.zeros_initializer(),
+        activation_dtype=x.dtype)
+
+    moving_mean = mtf.get_variable(
+        x.mesh, "moving_mean", reduced_shape,
+        initializer=tf.random_normal_initializer(stddev=1.0),
+        activation_dtype=x.dtype,
+        trainable=False)
+    moving_variance = mtf.get_variable(
+        x.mesh, "moving_variance",
+        reduced_shape, initializer=tf.ones_initializer(),
+        activation_dtype=x.dtype,
+        trainable=False)
+
+    # At training time, calculate mean and variance and normalize across batch
+    # dim.
+    if is_training:
+      mean = mtf.reduce_mean(x, output_shape=reduced_shape)
+      variance = mtf.reduce_mean(
+          mtf.square(x - mean), output_shape=reduced_shape)
+      norm_x = (x - mean) * mtf.rsqrt(variance + epsilon)
+
+      # Update running mean and running variance.
+      moving_mean = mtf.assign(
+          moving_mean, momentum * moving_mean + (1-momentum) * mean)
+      moving_variance = mtf.assign(
+          moving_variance,
+          momentum * moving_variance + (1 - momentum) * variance)
+    else:
+      # At eval and test time, use the running mean and variance.
+      norm_x = (x - moving_mean) * mtf.rsqrt(moving_variance + epsilon)
+    return norm_x * scale + bias
+
+
 def softmax_cross_entropy_with_logits(logits, targets, vocab_dim):
   """Per-example softmax loss.
 
diff --git a/tensor2tensor/mesh_tensorflow/mtf_resnet.py b/tensor2tensor/mesh_tensorflow/mtf_resnet.py
new file mode 100644
index 000000000..5de8ee028
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_resnet.py
@@ -0,0 +1,476 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""ResNet model with model and data parallelism using MTF.
+
+Integration of Mesh tensorflow with ResNet to do model parallelism.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import copy
+
+from tensor2tensor.layers import common_hparams
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+from tensor2tensor.mesh_tensorflow import mtf_layers
+from tensor2tensor.mesh_tensorflow import mtf_model
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+
+BATCH_NORM_DECAY = 0.9
+BATCH_NORM_EPSILON = 1e-5
+
+
+def batch_norm_relu(inputs, is_training, relu=True):
+  """Block of batch norm and relu."""
+  inputs = mtf_layers.batch_norm(
+      inputs,
+      is_training,
+      BATCH_NORM_DECAY,
+      epsilon=BATCH_NORM_EPSILON)
+  if relu:
+    inputs = mtf.relu(inputs)
+  return inputs
+
+
+def bottleneck_block(inputs,
+                     filters,
+                     is_training,
+                     strides,
+                     projection_shortcut=None,
+                     row_blocks_dim=None,
+                     col_blocks_dim=None):
+  """Bottleneck block variant for residual networks with BN after convolutions.
+
+  Args:
+    inputs: a `mtf.Tensor` of shape
+        `[batch_dim, row_blocks, col_blocks, rows, cols, in_channels]`.
+    filters: `int` number of filters for the first two convolutions. Note
+        that the third and final convolution will use 4 times as many filters.
+    is_training: `bool` for whether the model is in training mode.
+    strides: `int` block stride. If greater than 1, this block will ultimately
+        downsample the input.
+    projection_shortcut: `function` to use for projection shortcuts (typically
+        a 1x1 convolution to match the filter dimensions). If None, no
+        projection is used and the input is passed as unchanged through the
+        shortcut connection.
+    row_blocks_dim: a mtf.Dimension, row dimension which is
+        spatially partitioned along mesh axis
+    col_blocks_dim: a mtf.Dimension, row dimension which is
+        spatially partitioned along mesh axis
+
+  Returns:
+    The output `Tensor` of the block.
+  """
+  shortcut = inputs
+
+  filter_h_dim = mtf.Dimension("filter_height", 3)
+  filter_w_dim = mtf.Dimension("filter_width", 3)
+  one_h_dim = mtf.Dimension("filter_height", 1)
+  one_w_dim = mtf.Dimension("filter_width", 1)
+
+  if projection_shortcut is not None:
+    filters_dim = mtf.Dimension("filtersp", filters)
+    kernel = mtf.get_variable(
+        inputs.mesh, "kernel", mtf.Shape(
+            [one_h_dim, one_w_dim, inputs.shape.dims[-1], filters_dim]))
+    shortcut = projection_shortcut(inputs, kernel)
+
+  # First conv block
+  filters1_dim = mtf.Dimension("filters1", filters)
+  kernel1 = mtf.get_variable(
+      inputs.mesh, "kernel1", mtf.Shape(
+          [one_h_dim, one_w_dim, inputs.shape.dims[-1], filters1_dim]))
+  inputs = mtf.conv2d_with_blocks(
+      inputs,
+      kernel1,
+      strides=[1, 1, 1, 1],
+      padding="SAME",
+      h_blocks_dim=None, w_blocks_dim=col_blocks_dim)
+
+  # TODO(nikip): Add Dropout?
+  inputs = batch_norm_relu(inputs, is_training)
+
+  # Second conv block
+  filters2_dim = mtf.Dimension("filters2", filters)
+  kernel2 = mtf.get_variable(
+      inputs.mesh, "kernel2", mtf.Shape(
+          [filter_h_dim, filter_w_dim, filters1_dim, filters2_dim]))
+  inputs = mtf.conv2d_with_blocks(
+      inputs,
+      kernel2,
+      strides=[1, 1, 1, 1],
+      padding="SAME",
+      h_blocks_dim=row_blocks_dim, w_blocks_dim=col_blocks_dim)
+
+  inputs = batch_norm_relu(inputs, is_training)
+
+  # Third wide conv filter block
+  filters3_dim = mtf.Dimension("filters3", filters)
+  filters3_kernel = mtf.get_variable(
+      inputs.mesh, "wide_kernel", mtf.Shape(
+          [one_h_dim, one_w_dim, filters2_dim, filters3_dim]))
+  inputs = mtf.conv2d_with_blocks(
+      inputs,
+      filters3_kernel,
+      strides,
+      padding="SAME",
+      h_blocks_dim=None, w_blocks_dim=col_blocks_dim)
+
+  inputs = batch_norm_relu(
+      inputs,
+      is_training,
+      relu=False)
+
+  # TODO(nikip): Maybe add residual with a projection?
+  return mtf.relu(
+      inputs + mtf.rename_dimension(
+          shortcut, shortcut.shape.dims[-1].name, inputs.shape.dims[-1].name))
+
+
+def block_layer(inputs,
+                filters,
+                blocks,
+                strides,
+                is_training,
+                name,
+                row_blocks_dim=None,
+                col_blocks_dim=None):
+  """Creates one layer of blocks for the ResNet model.
+
+  Args:
+    inputs: `Tensor` of size `[batch, channels, height, width]`.
+    filters: `int` number of filters for the first convolution of the layer.
+    blocks: `int` number of blocks contained in the layer.
+    strides: `int` stride to use for the first convolution of the layer. If
+        greater than 1, this layer will downsample the input.
+    is_training: `bool` for whether the model is training.
+    name: `str`name for the Tensor output of the block layer.
+    row_blocks_dim: a mtf.Dimension, row dimension which is
+        spatially partitioned along mesh axis
+    col_blocks_dim: a mtf.Dimension, row dimension which is
+        spatially partitioned along mesh axis
+
+  Returns:
+    The output `Tensor` of the block layer.
+  """
+  with tf.variable_scope(name, default_name="block_layer"):
+    # Only the first block per block_layer uses projection_shortcut and strides
+    def projection_shortcut(inputs, kernel):
+      """Project identity branch."""
+      inputs = mtf.conv2d_with_blocks(
+          inputs,
+          kernel,
+          strides=strides,
+          padding="SAME",
+          h_blocks_dim=None, w_blocks_dim=col_blocks_dim)
+      return batch_norm_relu(
+          inputs, is_training, relu=False)
+
+    inputs = bottleneck_block(
+        inputs,
+        filters,
+        is_training,
+        strides,
+        projection_shortcut,
+        row_blocks_dim=row_blocks_dim,
+        col_blocks_dim=col_blocks_dim)
+
+    for i in range(1, blocks):
+      with tf.variable_scope("bottleneck_%d" % i):
+        inputs = bottleneck_block(
+            inputs,
+            filters,
+            is_training,
+            strides=[1, 1, 1, 1],
+            projection_shortcut=None,
+            row_blocks_dim=row_blocks_dim,
+            col_blocks_dim=col_blocks_dim)
+
+    return inputs
+
+
+@registry.register_model
+class MtfResNet(mtf_model.MtfModel):
+  """ResNet in mesh_tensorflow."""
+
+  def set_activation_type(self):
+    hparams = self._hparams
+    if hparams.activation_dtype == "float32":
+      activation_dtype = tf.float32
+    elif hparams.activation_dtype == "float16":
+      activation_dtype = tf.float16
+    elif hparams.activation_dtype == "bfloat16":
+      activation_dtype = tf.bfloat16
+    else:
+      raise ValueError(
+          "unknown hparams.activation_dtype %s" % hparams.activation_dtype)
+    return activation_dtype
+
+  def mtf_model_fn(self, features, mesh):
+    features = copy.copy(features)
+    tf.logging.info("features = %s" % features)
+    hparams = self._hparams
+    activation_dtype = self.set_activation_type()
+    is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+
+    # Declare all the dimensions
+    batch_dim = mtf.Dimension("batch", hparams.batch_size)
+    hidden_dim = mtf.Dimension("hidden", hparams.hidden_size)
+    filter_h_dim = mtf.Dimension("filter_height", 7)
+    filter_w_dim = mtf.Dimension("filter_width", 7)
+    filters = mtf.Dimension("filters", hparams.filter_sizes[0])
+    rows_dim = mtf.Dimension("rows_size", 32)
+    cols_dim = mtf.Dimension("cols_size", 96)
+    row_blocks_dim = mtf.Dimension("row_blocks", hparams.row_blocks)
+    col_blocks_dim = mtf.Dimension("col_blocks", hparams.col_blocks)
+    classes_dim = mtf.Dimension("classes", 10)
+    one_channel_dim = mtf.Dimension("one_channel", 1)
+
+    inputs = features["inputs"]
+    x = mtf.import_tf_tensor(
+        mesh, tf.reshape(inputs, [
+            hparams.batch_size,
+            hparams.row_blocks,
+            hparams.rows_size // hparams.row_blocks,
+            hparams.col_blocks,
+            hparams.num_channels*hparams.cols_size // hparams.col_blocks, 1]),
+        mtf.Shape(
+            [batch_dim, row_blocks_dim, rows_dim,
+             col_blocks_dim, cols_dim, one_channel_dim]))
+    x = mtf.transpose(x, [batch_dim, row_blocks_dim, col_blocks_dim,
+                          rows_dim, cols_dim, one_channel_dim])
+
+    x = mtf.to_float(x)
+    initial_filters = mtf.get_variable(
+        mesh, "init_filters",
+        mtf.Shape([filter_h_dim, filter_w_dim, one_channel_dim, filters]))
+    x = mtf.conv2d_with_blocks(
+        x,
+        initial_filters,
+        strides=[1, 1, 1, 1],
+        padding="SAME",
+        h_blocks_dim=None, w_blocks_dim=col_blocks_dim)
+
+    x = batch_norm_relu(x, is_training)
+
+    # Conv blocks
+    # [ self attention - ffn - residual + dropout] x n
+    for layer in range(hparams.num_layers):
+      layer_name = "block_layer_%d" % layer
+      with tf.variable_scope(layer_name):
+        # Residual block layer
+        x = block_layer(
+            inputs=x,
+            filters=hparams.filter_sizes[0],
+            blocks=hparams.layer_sizes[0],
+            strides=[1, 1, 1, 1],
+            is_training=is_training,
+            name="block_layer1",
+            row_blocks_dim=None,
+            col_blocks_dim=None)
+        x = block_layer(
+            inputs=x,
+            filters=hparams.filter_sizes[1],
+            blocks=hparams.layer_sizes[1],
+            strides=[1, 2, 2, 1],
+            is_training=is_training,
+            name="block_layer2",
+            row_blocks_dim=None,
+            col_blocks_dim=None)
+        x = block_layer(
+            inputs=x,
+            filters=hparams.filter_sizes[2],
+            blocks=hparams.layer_sizes[2],
+            strides=[1, 2, 2, 1],
+            is_training=is_training,
+            name="block_layer3",
+            row_blocks_dim=None,
+            col_blocks_dim=None)
+
+    # Calculate the logits and loss.
+    out = x
+    outputs = mtf_layers.dense(
+        out, hidden_dim,
+        reduced_dims=out.shape.dims[-5:],
+        activation=mtf.relu, name="dense")
+
+    # We assume fixed vocab size for targets
+    labels = tf.squeeze(tf.to_int32(features["targets"]), [2, 3])
+    labels = mtf.import_tf_tensor(
+        mesh, tf.reshape(labels, [hparams.batch_size]), mtf.Shape([batch_dim]))
+
+    logits = mtf_layers.dense(outputs, classes_dim, name="logits")
+    soft_targets = mtf.one_hot(labels, classes_dim, dtype=activation_dtype)
+    loss = mtf_layers.softmax_cross_entropy_with_logits(
+        logits, soft_targets, classes_dim)
+
+    # Reshape logits so it doesn't break inside t2t.
+    logits = mtf.reshape(
+        logits,
+        mtf.Shape([batch_dim, one_channel_dim, classes_dim]))
+    loss = mtf.reduce_mean(loss)
+    return logits, loss
+
+
+@registry.register_hparams
+def mtf_resnet_base():
+  """Set of hyperparameters."""
+  hparams = common_hparams.basic_params1()
+  hparams.no_data_parallelism = True
+  hparams.use_fixed_batch_size = True
+  hparams.batch_size = 32
+  hparams.max_length = 3072
+  hparams.hidden_size = 256
+  hparams.label_smoothing = 0.0
+  # 8-way model-parallelism
+  hparams.add_hparam("mesh_shape", "batch:8")
+  hparams.add_hparam("layout", "batch:batch")
+  hparams.add_hparam("num_heads", 8)
+  hparams.add_hparam("filter_size", 1024)
+
+  hparams.add_hparam("num_layers", 6)
+  hparams.add_hparam("attention_key_size", 256)
+  hparams.add_hparam("attention_value_size", 256)
+  # Share weights between input and target embeddings
+  hparams.shared_embedding = True
+
+  # mixture of experts hparams
+  hparams.add_hparam("ffn_layer", "dense_relu_dense")
+  hparams.add_hparam("moe_overhead_train", 1.0)
+  hparams.add_hparam("moe_overhead_eval", 2.0)
+  hparams.moe_num_experts = 16
+  hparams.moe_loss_coef = 1e-3
+
+  hparams.shared_embedding_and_softmax_weights = True
+  hparams.optimizer = "Adafactor"
+  hparams.learning_rate_schedule = "rsqrt_decay"
+  hparams.learning_rate_warmup_steps = 10000
+  hparams.add_hparam("d_kv", 32)
+
+  # Image related hparams
+  hparams.add_hparam("img_len", 32)
+  hparams.add_hparam("num_channels", 3)
+  hparams.add_hparam("row_blocks", 1)
+  hparams.add_hparam("col_blocks", 1)
+  hparams.add_hparam("rows_size", 32)
+  hparams.add_hparam("cols_size", 32)
+
+  # Model-specific parameters
+  hparams.add_hparam("layer_sizes", [3, 4, 6, 3])
+  hparams.add_hparam("filter_sizes", [64, 64, 128, 256, 512])
+  hparams.add_hparam("is_cifar", False)
+
+  # Variable init
+  hparams.initializer = "normal_unit_scaling"
+  hparams.initializer_gain = 2.
+
+  # TODO(nikip): Change optimization scheme?
+  hparams.learning_rate = 0.4
+  return hparams
+
+
+@registry.register_hparams
+def mtf_resnet_tiny():
+  """Catch bugs locally..."""
+  hparams = mtf_resnet_base()
+  hparams.num_layers = 2
+  hparams.hidden_size = 128
+  hparams.filter_size = 256
+  hparams.batch_size = 2
+  # data parallelism and model-parallelism
+  hparams.mesh_shape = "all:2"
+  hparams.layout = "batch:all"
+  hparams.layer_sizes = [3, 2, 3]
+  hparams.filter_sizes = [64, 64, 64]
+  return hparams
+
+
+@registry.register_hparams
+def mtf_resnet_single():
+  """Small single parameters."""
+  hparams = mtf_resnet_tiny()
+  hparams.mesh_shape = ""
+  hparams.layout = ""
+  hparams.hidden_size = 32
+  hparams.filter_size = 32
+  hparams.batch_size = 1
+  hparams.num_encoder_layers = 1
+  hparams.num_layers = 1
+  hparams.num_heads = 2
+  hparams.attention_key_size = 32
+  hparams.attention_value_size = 32
+  hparams.block_length = 16
+  return hparams
+
+
+@registry.register_hparams
+def mtf_resnet_base_single():
+  """Small single parameters."""
+  hparams = mtf_resnet_base()
+  hparams.num_layers = 6
+  hparams.filter_size = 256
+  hparams.block_length = 128
+  hparams.mesh_shape = ""
+  hparams.layout = ""
+  return hparams
+
+
+@registry.register_hparams
+def mtf_resnet_base_cifar():
+  """Data parallel CIFAR parameters."""
+  hparams = mtf_resnet_base()
+  hparams.mesh_shape = "batch:32"
+  hparams.layoyt = "batch:batch"
+  hparams.batch_size = 8
+  hparams.num_heads = 4
+  hparams.num_layers = 12
+  hparams.block_length = 256
+  hparams.hidden_size = 512
+  hparams.filter_size = 2048
+  hparams.learning_rate = 0.5
+  hparams.learning_rate_warmup_steps = 4000
+  hparams.layer_preprocess_sequence = "none"
+  hparams.layer_postprocess_sequence = "dan"
+  hparams.layer_prepostprocess_dropout = 0.3
+  hparams.unconditional = True
+  return hparams
+
+
+@registry.register_hparams
+def mtf_resnet_tiny_moe():
+  hparams = mtf_resnet_tiny()
+  hparams.mesh_shape = "all:4"
+  hparams.layout = "batch:all,experts:all"
+  hparams.ffn_layer = "moe"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_resnet_tiny_8gpu():
+  hparams = mtf_resnet_tiny()
+  hparams.mesh_shape = "all:8"
+  hparams.layout = "vocab:all;filter_size:all;heads:all"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_resnet_length_sharded():
+  hparams = mtf_resnet_tiny()
+  hparams.mesh_shape = "all"
+  hparams.layout = "length:all"
+  return hparams
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 40c955e78..6d2dbbf3e 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -20,6 +20,7 @@
 
 from tensor2tensor.layers import modalities  # pylint: disable=g-import-not-at-top
 from tensor2tensor.mesh_tensorflow import mtf_image_transformer
+from tensor2tensor.mesh_tensorflow import mtf_resnet
 from tensor2tensor.mesh_tensorflow import mtf_transformer
 from tensor2tensor.mesh_tensorflow.research import experiments_moe
 from tensor2tensor.models import basic

From c8d28e3a1adca7e6ac17f2ed570318f2565b09fe Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 2 Oct 2018 09:55:57 -0700
Subject: [PATCH 0950/2720] adding back default registration with empty list.

PiperOrigin-RevId: 215406574
---
 tensor2tensor/data_generators/gym_problems_specs.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index 37dc1d9ca..078f7cd0b 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -289,3 +289,12 @@ def create_problems_for_game(
       "simulated": simulated_cls,
       "world_model_eval": world_model_eval_cls,
   }
+
+# Register the atari games with all of the possible modes.
+for game in ATARI_ALL_MODES_SHORT_LIST:
+  ATARI_PROBLEMS[game] = {}
+  for mode in ATARI_GAME_MODES:
+    classes = create_problems_for_game(
+        game,
+        game_mode=mode)
+    ATARI_PROBLEMS[game][mode] = classes

From bc275fb21909ddeec94fb9d5784e38b1aeef0cca Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 2 Oct 2018 11:14:49 -0700
Subject: [PATCH 0951/2720] adding a test for recurrent model in RL loop.

PiperOrigin-RevId: 215423174
---
 tensor2tensor/rl/trainer_model_based.py       | 10 +++++
 .../rl/trainer_model_based_recurrent_test.py  | 37 +++++++++++++++++++
 2 files changed, 47 insertions(+)
 create mode 100644 tensor2tensor/rl/trainer_model_based_recurrent_test.py

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 90d94e286..105d567d9 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -1081,6 +1081,16 @@ def rlmb_tiny_stochastic():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_tiny_recurrent():
+  """Tiny setting with a recurrent next-frame model."""
+  hparams = rlmb_tiny()
+  hparams.epochs = 1  # Too slow with 2 for regular runs.
+  hparams.generative_model = "next_frame_basic_recurrent"
+  hparams.generative_model_params = "next_frame_basic_recurrent"
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_tiny_sv2p():
   """Tiny setting with a tiny sv2p model."""
diff --git a/tensor2tensor/rl/trainer_model_based_recurrent_test.py b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
new file mode 100644
index 000000000..629ff1245
--- /dev/null
+++ b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
@@ -0,0 +1,37 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+"""Tiny run of trainer_model_based. Smoke test."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.rl import trainer_model_based
+
+import tensorflow as tf
+
+FLAGS = tf.flags.FLAGS
+
+
+class ModelRLExperimentRecurrentTest(tf.test.TestCase):
+
+  def test_basic_recurrent(self):
+    FLAGS.output_dir = tf.test.get_temp_dir()
+    FLAGS.loop_hparams_set = "rlmb_tiny_recurrent"
+    FLAGS.schedule = "train"  # skip evaluation for world model training
+    trainer_model_based.main(None)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From be663e857827116e3d9de15f0e20017ebd7d996f Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 2 Oct 2018 14:21:27 -0700
Subject: [PATCH 0952/2720] Internal change

PiperOrigin-RevId: 215457505
---
 .../models/video/basic_deterministic.py       | 48 +++++++++++++++++++
 .../video/basic_deterministic_params.py       |  1 +
 tensor2tensor/models/video/basic_recurrent.py | 23 +++++++--
 .../models/video/basic_stochastic.py          |  2 +-
 tensor2tensor/models/video/sv2p.py            | 28 ++---------
 tensor2tensor/models/video/sv2p_params.py     |  1 -
 6 files changed, 73 insertions(+), 30 deletions(-)

diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index cb0551f50..f82afe5bb 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -53,6 +53,54 @@ def inject_latent(self, layer, features, filters):
     del features, filters
     return layer, 0.0
 
+  def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
+    """Hacky code the get the loss on predicted frames from input frames.
+
+       Recurrent models consume the frames one-by-one. Therefore
+       if there is more than one input frame they also get predicted.
+       T2T only calculates loss on the predicted target frames which
+       means the loss is not being applied on the predicted input frames.
+       This code is to fix this issue. Since the model is not aware of the
+       modality it has to match the pre-porocessing happening in bottom
+       function and therefore this becomes a very hacky code. This code
+       should match the bottom and top and loss of modalities otherwise
+       it will calculate the wrong loss.
+
+    Args:
+      extra_raw_gts: extra raw ground truth frames.
+      extra_gts: extra normalized ground truth frames.
+      extra_pds: extra predicted frames.
+
+    Returns:
+      Additional reconstruction loss.
+
+    Raises:
+      ValueError: in case of unknown modality.
+    """
+    if self._target_modality == "VideoModalityL2Raw":
+      recon_loss = tf.losses.mean_squared_error(extra_gts, extra_pds)
+    elif self._target_modality == "VideoModality":
+      shape = common_layers.shape_list(extra_pds)
+      updated_shape = shape[:-1] + [3, 256]
+      extra_pds = tf.reshape(extra_pds, updated_shape)
+      # Merge time and batch
+      logits = tf.reshape(extra_pds, [-1] + updated_shape[2:])
+      targets = extra_raw_gts
+      targets_shape = common_layers.shape_list(targets)
+      targets = tf.reshape(targets, [-1] + targets_shape[2:])
+      mod = self.hparams.problem_hparams.target_modality["targets"]
+      numerator, denominator = common_layers.padded_cross_entropy(
+          logits,
+          targets,
+          self.hparams.label_smoothing,
+          cutoff=getattr(self.hparams, "video_modality_loss_cutoff", 0.01),
+          weights_fn=mod.targets_weights_fn)
+      recon_loss = numerator / denominator
+    else:
+      raise ValueError("internal loss only supports specific modalities.")
+    tf.summary.scalar("recon_extra", recon_loss)
+    return recon_loss
+
   def inject_additional_input(self, layer, inputs, name, mode="concat"):
     layer_shape = common_layers.shape_list(layer)
     input_shape = common_layers.shape_list(inputs)
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 9105905f4..e87003834 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -50,6 +50,7 @@ def next_frame_basic_deterministic():
   hparams.add_hparam("tiny_mode", False)
   hparams.add_hparam("small_mode", False)
   hparams.add_hparam("stochastic_model", False)
+  hparams.add_hparam("internal_loss", True)
   return hparams
 
 
diff --git a/tensor2tensor/models/video/basic_recurrent.py b/tensor2tensor/models/video/basic_recurrent.py
index 4a60c9f44..b1fb4c247 100644
--- a/tensor2tensor/models/video/basic_recurrent.py
+++ b/tensor2tensor/models/video/basic_recurrent.py
@@ -21,8 +21,7 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
-from tensor2tensor.models.video import basic_deterministic
-from tensor2tensor.models.video import basic_deterministic_params
+from tensor2tensor.models.video import basic_stochastic
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -33,7 +32,8 @@
 
 
 @registry.register_model
-class NextFrameBasicRecurrent(basic_deterministic.NextFrameBasicDeterministic):
+class NextFrameBasicRecurrent(
+    basic_stochastic.NextFrameBasicStochasticDiscrete):
   """Basic next-frame recurrent model."""
 
   def predict_next_frame(self, frame, action, lstm_states):
@@ -142,6 +142,11 @@ def body(self, features):
       frame = all_frames[i]
       action = all_actions[i] if self.has_action else None
 
+      # more hack to enable latent_tower
+      # TODO(mbz): clean this up.
+      self.features["inputs"] = all_frames[i]
+      self.features["cur_target_frame"] = all_frames[i+1]
+
       # Run model.
       with tf.variable_scope("recurrent_model", reuse=tf.AUTO_REUSE):
         func_out = self.predict_next_frame(frame, action, lstm_states)
@@ -164,6 +169,16 @@ def body(self, features):
     output_frames = res_frames[hparams.video_num_input_frames-1:]
     frames = tf.stack(output_frames, axis=1)
 
+    has_input_predictions = hparams.video_num_input_frames > 1
+    if self.is_training and hparams.internal_loss and has_input_predictions:
+      # add the loss for input frames as well.
+      extra_gts = input_frames[1:]
+      extra_pds = res_frames[:hparams.video_num_input_frames-1]
+      extra_raw_gts = features["inputs_raw"][:, 1:]
+      recon_loss = self.get_extra_internal_loss(
+          extra_raw_gts, extra_gts, extra_pds)
+      extra_loss += recon_loss
+
     if not self.has_reward:
       return frames, extra_loss
     rewards = tf.concat(res_rewards[hparams.video_num_input_frames-1:], axis=1)
@@ -173,7 +188,7 @@ def body(self, features):
 @registry.register_hparams
 def next_frame_basic_recurrent():
   """Basic 2-frame recurrent model with stochastic tower."""
-  hparams = basic_deterministic_params.next_frame_basic_deterministic()
+  hparams = basic_stochastic.next_frame_basic_stochastic_discrete()
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 4
   hparams.add_hparam("num_lstm_layers", 1)
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 8ffb0756e..a687b9fb5 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -70,7 +70,7 @@ def inject_latent(self, layer, features, filters):
     filters = hparams.hidden_size
     kernel = (4, 4)
 
-    if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+    if self.is_predicting:
       layer_shape = common_layers.shape_list(layer)
       if hparams.full_latent_tower:
         rand = tf.random_uniform(layer_shape[:-1] + [hparams.bottleneck_bits])
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 96c22130f..217a481ff 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -618,34 +618,14 @@ def body(self, features):
     reward_pred = common_video.swap_time_and_batch_axes(reward_pred)
 
     if self.is_training and hparams.internal_loss:
-      # add the MSE loss for input frames as well.
+      # add the loss for input frames as well.
       extra_gts = all_frames[1:hparams.video_num_input_frames]
       extra_gts = common_video.swap_time_and_batch_axes(extra_gts)
       extra_pds = gen_images[:hparams.video_num_input_frames-1]
       extra_pds = common_video.swap_time_and_batch_axes(extra_pds)
-      if self._target_modality == "VideoModalityL2Raw":
-        recon_loss = tf.losses.mean_squared_error(extra_gts, extra_pds)
-      elif self._target_modality == "VideoModality":
-        shape = common_layers.shape_list(extra_pds)
-        updated_shape = shape[:-1] + [3, 256]
-        extra_pds = tf.reshape(extra_pds, updated_shape)
-        # Merge time and batch
-        logits = tf.reshape(extra_pds, [-1] + updated_shape[2:])
-        targets = features["inputs_raw"][:, 1:]
-        targets_shape = common_layers.shape_list(targets)
-        targets = tf.reshape(targets, [-1] + targets_shape[2:])
-        mod = self.hparams.problem_hparams.target_modality["targets"]
-        numerator, denominator = common_layers.padded_cross_entropy(
-            logits,
-            targets,
-            hparams.label_smoothing,
-            cutoff=getattr(hparams, "video_modality_loss_cutoff", 0.01),
-            weights_fn=mod.targets_weights_fn)
-        recon_loss = numerator / denominator
-      else:
-        raise ValueError("internal loss only supports specific modalities.")
-
-      tf.summary.scalar("recon_extra", recon_loss)
+      extra_raw_gts = features["inputs_raw"][:, 1:]
+      recon_loss = self.get_extra_internal_loss(
+          extra_raw_gts, extra_gts, extra_pds)
       extra_loss += recon_loss
 
     return_targets = predictions
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 0d8aed81d..96879c6d1 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -47,7 +47,6 @@ def next_frame_sv2p():
   hparams.add_hparam("scheduled_sampling_decay_steps", 10000)
   hparams.add_hparam("scheduled_sampling_k", 900.0)
   hparams.add_hparam("upsample_method", "conv2d_transpose")
-  hparams.add_hparam("internal_loss", True)
   hparams.add_hparam("reward_model", "basic")
   return hparams
 

From d67c0b54e3def204ba875e14aa7a8df23cc55e6c Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 2 Oct 2018 23:25:58 +0200
Subject: [PATCH 0953/2720] Add flag disable_ffmpeg writing individual debug
 frames as PNG (#1108)

* Add missing docstrings in common_video

* Extract a base class from VideoWriter and BatchVideoWriter

* Add flag disable_ffmpeg writing individual debug frames as PNG
---
 tensor2tensor/data_generators/video_utils.py |  34 ++++-
 tensor2tensor/layers/common_video.py         | 133 ++++++++++++++-----
 2 files changed, 130 insertions(+), 37 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 08c85c9c7..47484916a 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -33,6 +33,13 @@
 
 import tensorflow as tf
 
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+flags.DEFINE_bool(
+    "disable_ffmpeg", False, "Disable FFMPEG when generating debug videos."
+)
+
 
 def resize_video_frames(images, size):
   resized_images = []
@@ -122,6 +129,17 @@ def summarize_video_metrics(hook_args):
   return summary_values
 
 
+def debug_video_writer_factory(output_dir):
+  """Creates a VideoWriter for debug videos."""
+  if FLAGS.disable_ffmpeg:
+    return common_video.IndividualFrameWriter(output_dir)
+  else:
+    output_path = os.path.join(output_dir, "video.avi")
+    return common_video.WholeVideoWriter(
+        fps=10, output_path=output_path, file_format="avi"
+    )
+
+
 class VideoProblem(problem.Problem):
   """Base class for problems with videos."""
 
@@ -435,8 +453,7 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
     Raises:
       ValueError: if the frame has a different number of channels than required.
     """
-    if self.debug_dump_frames_path:
-      writer = common_video.VideoWriter(fps=10, file_format="avi")
+    writer = None
 
     with tf.Graph().as_default():
       image_t = tf.placeholder(dtype=tf.uint8, shape=(None, None, None))
@@ -461,16 +478,19 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
             features["image/encoded_debug"] = [encoded_debug]
 
           if self.debug_dump_frames_path:
+            # Defer creating debug writer until we know debug_dump_frames_path.
+            if writer is None:
+              if not tf.gfile.Exists(self.debug_dump_frames_path):
+                tf.gfile.MkDir(self.debug_dump_frames_path)
+              writer = debug_video_writer_factory(self.debug_dump_frames_path)
             img = unencoded_debug if has_debug_image else unencoded_frame
-            writer.write(img)
+            encoded_img = encoded_debug if has_debug_image else encoded_frame
+            writer.write(img, encoded_img)
 
           yield features
 
     if self.debug_dump_frames_path:
-      if not tf.gfile.Exists(self.debug_dump_frames_path):
-        tf.gfile.MkDir(self.debug_dump_frames_path)
-      path = os.path.join(self.debug_dump_frames_path, "video.avi")
-      writer.finish_to_file(path)
+      writer.finish_to_disk()
 
   def generate_data(self, data_dir, tmp_dir, task_id=-1):
     """The function generating the data."""
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index d6cdeefc0..63a93f5fe 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -322,7 +322,7 @@ def _encode_gif(images, fps):
   Raises:
     IOError: If the ffmpeg command returns an error.
   """
-  writer = VideoWriter(fps)
+  writer = WholeVideoWriter(fps)
   writer.write_multi(images)
   return writer.finish()
 
@@ -556,10 +556,52 @@ def beta_schedule(schedule, global_step, final_beta, decay_start, decay_end):
 
 
 class VideoWriter(object):
-  """Helper class for writing videos."""
+  """Base helper class for writing videos."""
 
-  def __init__(self, fps, file_format="gif"):
+  def write(self, frame, encoded_frame=None):
+    """Writes a single video frame."""
+    raise NotImplementedError
+
+  def write_multi(self, frames, encoded_frames=None):
+    """Writes multiple video frames."""
+    if encoded_frames is None:
+      # Infinite iterator.
+      encoded_frames = iter(lambda: None, 1)
+    for (frame, encoded_frame) in zip(frames, encoded_frames):
+      self.write(frame, encoded_frame)
+
+  def finish(self):
+    """Finishes writing frames and returns output, if any.
+
+    Frees any resources acquired by the writer.
+    """
+    pass
+
+  def save_to_disk(self, output):
+    """Saves output to disk.
+
+    Args:
+      output: result of finish().
+    """
+    raise NotImplementedError
+
+  def finish_to_disk(self):
+    """Finishes writing frames and saves output to disk, if any."""
+    output = self.finish()  # pylint: disable=assignment-from-no-return
+    if output is not None:
+      self.save_to_disk(output)
+
+  def __del__(self):
+    """Frees any resources acquired by the writer."""
+    self.finish()
+
+
+class WholeVideoWriter(VideoWriter):
+  """Helper class for writing whole videos."""
+
+  def __init__(self, fps, output_path=None, file_format="gif"):
     self.fps = fps
+    self.output_path = output_path
     self.file_format = file_format
     self.proc = None
     self._out_chunks = []
@@ -599,7 +641,18 @@ def __init_ffmpeg(self, image_shape):
     )
 
   def _start_reader_thread(self, stream, chunks):
-    """TODO(koz4k): Write a docstring."""
+    """Starts a thread for reading output from FFMPEG.
+
+    The thread reads consecutive chunks from the stream and saves them in
+    the given list.
+
+    Args:
+      stream: output stream of the FFMPEG process.
+      chunks: list to save output chunks to.
+
+    Returns:
+      Thread
+    """
     import io  # pylint: disable=g-import-not-at-top
     import threading  # pylint: disable=g-import-not-at-top
     def target():
@@ -612,17 +665,20 @@ def target():
     thread.start()
     return thread
 
-  def write(self, frame):
+  def write(self, frame, encoded_frame=None):
     if self.proc is None:
       self.__init_ffmpeg(frame.shape)
     self.proc.stdin.write(frame.tostring())
 
-  def write_multi(self, frames):
-    for frame in frames:
-      self.write(frame)
-
   def finish(self):
-    """TODO(koz4k): Write a docstring."""
+    """Finishes transconding and returns the video.
+
+    Returns:
+      bytes
+
+    Raises:
+      IOError: in case of transcoding error.
+    """
     if self.proc is None:
       return None
     self.proc.stdin.close()
@@ -638,40 +694,57 @@ def finish(self):
     self.proc = None
     return out
 
-  def finish_to_file(self, path):
-    out = self.finish()
-    if out is not None:
-      with tf.gfile.Open(path, "w") as f:
-        f.write(out)
+  def save_to_disk(self, output):
+    if self.output_path is None:
+      raise ValueError(
+          "This writer doesn't support saving to disk (output_path not "
+          "specified)."
+      )
+    with tf.gfile.Open(self.output_path, "w") as f:
+      f.write(output)
 
-  def __del__(self):
-    self.finish()
 
-
-class BatchVideoWriter(object):
+class BatchWholeVideoWriter(VideoWriter):
   """Helper class for writing videos in batch."""
 
-  def __init__(self, fps, file_format="gif"):
+  def __init__(self, fps, path_template, file_format="gif"):
     self.fps = fps
+    self.path_template = path_template
     self.file_format = file_format
     self.writers = None
 
-  def write(self, batch_frame):
+  def write(self, batch_frame, batch_encoded_frame=None):
     if self.writers is None:
       self.writers = [
-          VideoWriter(self.fps, self.file_format) for _ in batch_frame]
+          WholeVideoWriter(
+              self.fps, self.path_template.format(i), self.file_format
+          )
+          for i in range(len(batch_frame))
+      ]
     for i, frame in enumerate(batch_frame):
       self.writers[i].write(frame)
 
-  def write_multi(self, batch_frames):
-    for batch_frame in batch_frames:
-      self.write(batch_frame)
-
   def finish(self):
     outs = [w.finish() for w in self.writers]
     return outs
 
-  def finish_to_files(self, path_template):
-    for i, writer in enumerate(self.writers):
-      path = path_template.format(i)
-      writer.finish_to_file(path)
+  def save_to_disk(self, outputs):
+    for (writer, output) in zip(self.writers, outputs):
+      writer.save_to_disk(output)
+
+
+class IndividualFrameWriter(VideoWriter):
+  """Helper class for writing individual video frames."""
+
+  def __init__(self, output_dir):
+    self.output_dir = output_dir
+    self._counter = 0
+
+  def write(self, frame=None, encoded_frame=None):
+    import os  # pylint: disable=g-import-not-at-top
+    if encoded_frame is None:
+      raise ValueError("This writer only supports encoded frames.")
+    path = os.path.join(self.output_dir, "frame_%05d.png" % self._counter)
+    with tf.gfile.Open(path, "wb") as f:
+      f.write(encoded_frame)
+      self._counter += 1

From 6e6c8f7023715878b28c2506a05fde02feac98db Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 2 Oct 2018 14:26:25 -0700
Subject: [PATCH 0954/2720] internal merge of PR #1108

PiperOrigin-RevId: 215458375
---
 tensor2tensor/layers/common_video.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 63a93f5fe..64e674a2b 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -714,6 +714,7 @@ def __init__(self, fps, path_template, file_format="gif"):
     self.writers = None
 
   def write(self, batch_frame, batch_encoded_frame=None):
+    del batch_encoded_frame
     if self.writers is None:
       self.writers = [
           WholeVideoWriter(

From 52d9c43375b1b33861ad31e000096cb211886ade Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 3 Oct 2018 10:26:50 -0700
Subject: [PATCH 0955/2720] recurrent RL config.

PiperOrigin-RevId: 215584396
---
 tensor2tensor/models/video/basic_recurrent.py | 4 ++--
 tensor2tensor/rl/trainer_model_based.py       | 9 +++++++++
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/video/basic_recurrent.py b/tensor2tensor/models/video/basic_recurrent.py
index b1fb4c247..79f155501 100644
--- a/tensor2tensor/models/video/basic_recurrent.py
+++ b/tensor2tensor/models/video/basic_recurrent.py
@@ -191,6 +191,6 @@ def next_frame_basic_recurrent():
   hparams = basic_stochastic.next_frame_basic_stochastic_discrete()
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 4
-  hparams.add_hparam("num_lstm_layers", 1)
-  hparams.add_hparam("num_lstm_filters", 8)
+  hparams.add_hparam("num_lstm_layers", 2)
+  hparams.add_hparam("num_lstm_filters", 256)
   return hparams
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 105d567d9..f9a2e7664 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -891,6 +891,15 @@ def rlmb_base_stochastic_discrete():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_base_stochastic_recurrent():
+  """Base setting with recurrent model."""
+  hparams = rlmb_base()
+  hparams.generative_model = "next_frame_basic_recurrent"
+  hparams.generative_model_params = "next_frame_basic_recurrent"
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_base_stochastic_discrete_noresize():
   """Base setting with stochastic discrete model."""

From 6b82d65520404d92cdc8b4355c6f277e80b6ea1f Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 3 Oct 2018 17:28:51 -0700
Subject: [PATCH 0956/2720] Internal change

PiperOrigin-RevId: 215658756
---
 tensor2tensor/models/research/autoencoders.py | 27 +++++++++++++++----
 1 file changed, 22 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index bc5e4507a..74979b4e4 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -207,7 +207,8 @@ def body(self, features):
       xb_loss = 0.0
       b_shape = common_layers.shape_list(b)
       self._cur_bottleneck_tensor = b
-      b = self.unbottleneck(b, common_layers.shape_list(x)[-1])
+      res_size = common_layers.shape_list(x)[-1]
+      b = self.unbottleneck(b, res_size)
       if not is_training:
         x = b
       else:
@@ -235,7 +236,7 @@ def body(self, features):
             self.sample(shape=b_shape),
             common_layers.shape_list(x)[-1],
             reuse=True)
-        x = tf.concat([g, x], axis=0)
+        x = tf.concat([x, g], axis=0)
     else:
       if self._cur_bottleneck_tensor is None:
         b = self.sample()
@@ -270,7 +271,7 @@ def body(self, features):
       return reconstr, {"bottleneck_loss": 0.0}
 
     if hparams.gan_loss_factor != 0.0:
-      res_gan, res = tf.split(res, 2, axis=0)
+      res, res_gan = tf.split(res, 2, axis=0)
 
     # Losses.
     losses = {
@@ -859,8 +860,9 @@ def embed(self, x, name="embedding"):
     return tf.concat([xte, xie], axis=0)
 
   def bottleneck(self, x):
+    hparams = self.hparams
     b, _ = super(AutoencoderDualDiscrete, self).bottleneck(x)
-    if self.hparams.mode != tf.estimator.ModeKeys.TRAIN:
+    if hparams.mode != tf.estimator.ModeKeys.TRAIN:
       return b, 0.0
     bt, bi = tf.split(b, 2, axis=0)
     # Share the first hparams.bottleneck_shared_bits.
@@ -869,8 +871,19 @@ def bottleneck(self, x):
     br = tf.where(rand < 0.5, bt, bi)  # Break ties at random.
     bs = tf.where(shared == 0, br, shared)
     bs = tf.concat([bs, bs], axis=0)
-    n = self.hparams.bottleneck_shared_bits
+    n = hparams.bottleneck_shared_bits
+    step = tf.train.get_global_step()
+    zero = tf.constant(0, dtype=tf.int64)
+    if step is None:
+      step = zero
+    step = tf.maximum(zero, step - hparams.bottleneck_shared_bits_start_warmup)
+    f = common_layers.inverse_lin_decay(
+        hparams.bottleneck_shared_bits_stop_warmup, min_value=0.1, step=step)
+    n = tf.where(step > 1, n * f, n)
+    n = tf.cast(n, tf.int64)
+    b_shape = common_layers.shape_list(b)
     b = tf.concat([bs[..., :n], b[..., n:]], axis=-1)
+    b = tf.reshape(b, b_shape)
     return b, 0.0
 
   def unbottleneck(self, b, res_size, reuse=None):
@@ -996,6 +1009,8 @@ def autoencoder_basic():
   hparams.add_hparam("max_hidden_size", 1024)
   hparams.add_hparam("bottleneck_bits", 128)
   hparams.add_hparam("bottleneck_shared_bits", 0)
+  hparams.add_hparam("bottleneck_shared_bits_start_warmup", 0)
+  hparams.add_hparam("bottleneck_shared_bits_stop_warmup", 0)
   hparams.add_hparam("bottleneck_noise", 0.1)
   hparams.add_hparam("bottleneck_warmup_steps", 2000)
   hparams.add_hparam("sample_height", 32)
@@ -1161,6 +1176,8 @@ def autoencoder_ordered_text():
   hparams = autoencoder_ordered_discrete()
   hparams.bottleneck_bits = 512
   hparams.bottleneck_shared_bits = 512-64
+  hparams.bottleneck_shared_bits_start_warmup = 75000
+  hparams.bottleneck_shared_bits_stop_warmup = 275000
   hparams.num_hidden_layers = 7
   hparams.batch_size = 1024
   hparams.autoregressive_mode = "conv5"

From 98e7cc2c34cbd4ad937e759c94d2e65901b3e3b8 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Wed, 3 Oct 2018 21:48:52 -0700
Subject: [PATCH 0957/2720] Spatial partitioning with local attention

PiperOrigin-RevId: 215681858
---
 .../mesh_tensorflow/mesh_tensorflow.py        | 34 +++++++
 tensor2tensor/mesh_tensorflow/mtf_layers.py   | 94 +++++++++++++++++++
 2 files changed, 128 insertions(+)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index c2ed10c09..cca894979 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -3985,6 +3985,40 @@ def halo_exchange(x, blocks_dim, block_size_dim, halo_size, wrap=False):
   return concat(parts, block_size_dim.name)
 
 
+def left_halo_exchange(x, blocks_dim, block_size_dim, halo_size, wrap=False):
+  """Concat each block with the margins of adjacent blocks from the left.
+
+  Get left blocks_dim and concatenate along block_size_dim.
+
+  Args:
+    x: a Tensor.
+    blocks_dim: a Dimension in x.shape
+    block_size_dim: a Dimension in x.shape
+    halo_size: an integer
+    wrap: a boolean
+
+  Returns:
+    a Tensor with the same shape as x, other than in block_size_dim, whose
+    size is increased by halo_size.
+  """
+  if halo_size == 0:
+    return x
+
+  block_size = block_size_dim.size
+  partial_size = halo_size % block_size
+  num_complete_blocks = halo_size // block_size
+  parts = [x]
+
+  for i in xrange(1, num_complete_blocks + 1):
+    parts = ([shift(x, i, blocks_dim, wrap)] + parts)
+  if partial_size > 0:
+    right_margin = slice(x, block_size_dim.size - partial_size, partial_size,
+                         block_size_dim.name)
+    parts = ([shift(right_margin, num_complete_blocks + 1, blocks_dim, wrap)]
+             + parts)
+  return concat(parts, block_size_dim.name)
+
+
 def conv2d_with_blocks(
     conv_input,
     conv_filter,
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers.py b/tensor2tensor/mesh_tensorflow/mtf_layers.py
index 4ed558596..8b4e8bd04 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers.py
@@ -231,6 +231,100 @@ def dense_relu_dense(x,
     return mtf.einsum([h, wo])
 
 
+def local_1d_halo_exchange(k, v, num_w_blocks, w_dim, memory_w_dim, mask_right):
+  """Halo exchange for keys and values for Local 1D attention."""
+  if num_w_blocks is not None:
+    if mask_right:
+      k = mtf.left_halo_exchange(k, num_w_blocks, w_dim, memory_w_dim.size)
+      v = mtf.left_halo_exchange(v, num_w_blocks, w_dim, memory_w_dim.size)
+    else:
+      k = mtf.halo_exchange(k, num_w_blocks, w_dim, memory_w_dim.size)
+      v = mtf.halo_exchange(v, num_w_blocks, w_dim, memory_w_dim.size)
+  else:
+    if mask_right:
+      k = mtf.pad(k, [memory_w_dim, None], w_dim.name)
+    else:
+      k = mtf.pad(k, [memory_w_dim, memory_w_dim], w_dim.name)
+    v = mtf.pad(v, [memory_w_dim, memory_w_dim], w_dim.name)
+  return k, v
+
+
+def local_self_attention_spatial_blocks(
+    query_antecedent,
+    kv_channels,
+    heads,
+    memory_w_dim=None,
+    mask_right=False,
+    name=None):
+  """Attention to the source position and a neighborhood to the left or right.
+
+  The sequence is divided into blocks of length block_size.
+  Attention for a given query position can only see memory positions
+  less than or equal to the query position, in the corresponding block
+  and the previous block.
+
+  Args:
+    query_antecedent: a mtf.Tensor with shape
+      [batch, num_h_blocks, num_w_blocks, h_dim, w_dim, io_channels]
+      must have the same size as query_length, but a different name.
+    kv_channels: a mtf.Dimension (the size of the key and value vectors)
+    heads: a mtf.Dimension (the number of heads)
+    memory_w_dim: mtf Dimension, for the memory width block.
+    mask_right: bool, flag specifying whether we mask out attention to the right
+      for the decoder.
+    name: an optional string.
+
+  Returns:
+    a Tensor of shape
+        [batch, num_h_blocks, num_w_blocks, h_dim, w_dim, io_channels]
+
+  Raises:
+    ValueError: if channels or depth don't match.
+  """
+  with tf.variable_scope(
+      name, default_name="multihead_attention",
+      values=[query_antecedent]):
+
+    w_dim, io_channels = query_antecedent.shape.dims[-2:]
+    batch, num_w_blocks = query_antecedent.shape.dims[:2]
+    q_var, k_var, v_var, o_var = multihead_attention_vars(
+        query_antecedent.mesh, heads, io_channels, kv_channels,
+        query_antecedent.dtype)
+
+    # Rename dimensions for the memory height and width.
+    memory_antecedent = mtf.rename_dimension(
+        query_antecedent, w_dim.name, memory_w_dim.name)
+
+    # Call einsum over the query and memory to get query q, keys k and values v.
+    q = mtf.einsum(
+        [query_antecedent, q_var],
+        mtf.Shape([batch, heads, num_w_blocks, w_dim, kv_channels]))
+    k = mtf.einsum(
+        [memory_antecedent, k_var],
+        mtf.Shape([batch, heads, num_w_blocks, w_dim, kv_channels]))
+    v = mtf.einsum(
+        [memory_antecedent, v_var],
+        mtf.Shape([batch, heads, num_w_blocks, w_dim, kv_channels]))
+
+    # Halo exchange for memory blocks.
+    if memory_w_dim is not None:
+      k, v = local_1d_halo_exchange(
+          k, v, num_w_blocks, w_dim, memory_w_dim, mask_right)
+
+    # Calculate the causal mask to avoid peeking into the future. We compute
+    # this once and reuse it for all blocks since the block_size is known.
+    mask = None
+    if mask_right:
+      mask = attention_bias_local_block(
+          query_antecedent.mesh, w_dim, memory_w_dim)
+
+    output = dot_product_attention(q, k, v, mask=mask)
+
+    return mtf.einsum(
+        [output, o_var],
+        mtf.Shape([batch, num_w_blocks, w_dim, io_channels]))
+
+
 def masked_local_attention_1d(query_antecedent,
                               memory_antecedent,
                               kv_channels,

From f1471a588923d502a628dc7845c163207ddf583f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 4 Oct 2018 09:40:49 -0700
Subject: [PATCH 0958/2720] Allow max_steps to be set during train() call.

PiperOrigin-RevId: 215752326
---
 tensor2tensor/utils/trainer_lib.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 95ef4a3af..1564a42d7 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -338,11 +338,11 @@ def train_and_evaluate(self):
                          "in train_hooks.")
       self.train()
 
-  def train(self):
+  def train(self, max_steps=None):
     self._estimator.train(
         self._train_spec.input_fn,
         hooks=self._train_spec.hooks,
-        max_steps=self._train_spec.max_steps)
+        max_steps=max_steps or self._train_spec.max_steps)
 
   def evaluate(self):
     return self._estimator.evaluate(

From b4094d065fa0ae8842cd667fb0e5a2c652407c9c Mon Sep 17 00:00:00 2001
From: Youlong Cheng <ylc@google.com>
Date: Thu, 4 Oct 2018 09:50:58 -0700
Subject: [PATCH 0959/2720] Internal change

PiperOrigin-RevId: 215753966
---
 .../mesh_tensorflow/mesh_tensorflow.py        | 12 ++++-
 tensor2tensor/mesh_tensorflow/mtf_model.py    | 19 +++++--
 tensor2tensor/mesh_tensorflow/mtf_utils.py    | 40 ++++++++++++++
 .../mesh_tensorflow/mtf_utils_test.py         | 53 +++++++++++++++++++
 .../research/experiments_moe.py               | 46 ++++++++++++++++
 .../mesh_tensorflow/simd_mesh_impl.py         |  5 +-
 6 files changed, 167 insertions(+), 8 deletions(-)
 create mode 100644 tensor2tensor/mesh_tensorflow/mtf_utils_test.py

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index cca894979..d51f4dba7 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -496,14 +496,22 @@ class Mesh(object):
   A Lowering assigns each Mesh to a MeshImpl.
   """
 
-  def __init__(self, graph, name):
+  def __init__(self, graph, name, variable_placer=None):
     self._graph = graph
     self._name = name
+    self._variable_placer = variable_placer
 
   @property
   def graph(self):
     return self._graph
 
+  @property
+  def variable_placer_fn(self):
+    if self._variable_placer is not None:
+      return self._variable_placer.device_function
+    else:
+      return "cpu:0"
+
 
 class MeshImpl(object):
   """Implementation of a Mesh.
@@ -2516,7 +2524,7 @@ def __init__(self, mesh, name, shape, dtype, initializer,
                trainable, **kwargs):
     super(Variable, self).__init__([], mesh, name="name_will_be_set_later")
     self._trainable = trainable
-    with tf.device("cpu:0"), mtf_utils.outside_all_rewrites():
+    with tf.device(mesh.variable_placer_fn), mtf_utils.outside_all_rewrites():
       self.master = tf.get_variable(
           name, shape.to_integer_list, dtype=dtype, initializer=initializer,
           **kwargs)
diff --git a/tensor2tensor/mesh_tensorflow/mtf_model.py b/tensor2tensor/mesh_tensorflow/mtf_model.py
index 769033ff9..8d50787a7 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_model.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_model.py
@@ -72,17 +72,26 @@ def estimator_model_fn(cls,
         decode_hparams=decode_hparams)
 
     global_step = tf.train.get_global_step()
-    graph = mtf.Graph()
-    mesh = mtf.Mesh(graph, "my_mesh")
 
     mesh_shape = mtf.convert_to_shape(hparams.mesh_shape)
     layout_rules = mtf.convert_to_layout_rules(hparams.layout)
     if use_tpu:
+      ctx = params["context"]
+      num_hosts = ctx.num_hosts
+      host_placement_fn = ctx.tpu_host_placement_function
+      device_list = [host_placement_fn(host_id=t) for t in range(num_hosts)]
+      # TODO(ylc): Better estimation of replica cache size?
+      replica_cache_size = 300 * 1000000  # 300M per replica
+      # Worker 0 caches all the TPU binaries.
+      worker0_mem = replica_cache_size * ctx.num_replicas
+      devices_memeory_usage = [worker0_mem] + [0] * (num_hosts - 1)
+      var_placer = mtf_utils.BalancedVariablePlacer(device_list,
+                                                    devices_memeory_usage)
       mesh_devices = [""] * mesh_shape.size
       mesh_impl = simd_mesh_impl.SimdMeshImpl(
-          mesh_shape, layout_rules, mesh_devices,
-          params["context"].device_assignment)
+          mesh_shape, layout_rules, mesh_devices, ctx.device_assignment)
     else:
+      var_placer = None
       if len(data_parallelism.ps_devices) == 1:
         mesh_devices = [""] * mesh_shape.size
       else:
@@ -91,6 +100,8 @@ def estimator_model_fn(cls,
       mesh_impl = placement_mesh_impl.PlacementMeshImpl(
           mesh_shape, layout_rules, mesh_devices)
 
+    graph = mtf.Graph()
+    mesh = mtf.Mesh(graph, "my_mesh", var_placer)
     # PREDICT mode
     if mode == tf.estimator.ModeKeys.PREDICT:
       return model.estimator_spec_predict(features, mesh, mesh_impl, use_tpu)
diff --git a/tensor2tensor/mesh_tensorflow/mtf_utils.py b/tensor2tensor/mesh_tensorflow/mtf_utils.py
index 70cee4923..108d9a530 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_utils.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_utils.py
@@ -19,7 +19,9 @@
 from __future__ import print_function
 
 import contextlib
+import heapq
 
+import tensorflow as tf
 from tensorflow.python.framework import ops
 
 
@@ -27,3 +29,41 @@
 def outside_all_rewrites():
   with ops.control_dependencies(None):
     yield
+
+
+class BalancedVariablePlacer(object):
+  """Place the variable on different device and blance the memory usage."""
+
+  def __init__(self, devices, init_usage=None):
+    init_usage = init_usage if init_usage else [0] * len(devices)
+    assert len(devices) == len(init_usage)
+    self._mem_device_heap = list(zip(init_usage, devices))
+    heapq.heapify(self._mem_device_heap)
+    self._last_device = devices[0]
+
+  def device_function(self, var):
+    """Choose a device for the input variable.
+
+    Args:
+      var: an Variable.
+
+    Returns:
+      The device for placing the var.
+    """
+    if var.type not in ('Variable', 'VariableV2', 'VarHandleOp'):
+      tf.logging.info('Place {} on last device: {}.'.format(
+          var.name, self._last_device))
+      return self._last_device
+
+    shape = tf.TensorShape(var.get_attr('shape'))
+    assert shape.num_elements() is not None
+
+    size = tf.DType(var.get_attr('dtype')).size
+    mem, device = heapq.heappop(self._mem_device_heap)
+    mem += shape.num_elements() * size
+    heapq.heappush(self._mem_device_heap, (mem, device))
+    tf.logging.info('Place variable {} on {} and consumes {} Bytes.'.format(
+        var.name, device, mem))
+    self._last_device = device
+
+    return device
diff --git a/tensor2tensor/mesh_tensorflow/mtf_utils_test.py b/tensor2tensor/mesh_tensorflow/mtf_utils_test.py
new file mode 100644
index 000000000..a848d09c4
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/mtf_utils_test.py
@@ -0,0 +1,53 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 mtf_utils.py."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from six.moves import xrange  # pylint: disable=redefined-builtin
+
+from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
+from tensor2tensor.mesh_tensorflow import mtf_utils
+
+import tensorflow as tf
+
+
+class MtfUtilsTest(tf.test.TestCase):
+
+  def test_variable_placer(self):
+    sizes = [100, 0, 0, 0]
+    device_list = ['cpu:0', 'cpu:1', 'cpu:2', 'cpu:3']
+
+    with tf.Graph().as_default() as g:
+      var_placer = mtf_utils.BalancedVariablePlacer(device_list, sizes)
+      graph = mtf.Graph()
+      mesh = mtf.Mesh(graph, 'my_mesh', var_placer)
+
+      hidden_dim = mtf.Dimension('hidden', 10)
+      output_dim = mtf.Dimension('output_feature', 10)
+
+      for i in xrange(5):
+        # Each variable takes 400 Bytes, and will be placed from cpu:1.
+        mtf.get_variable(mesh, 'w{}'.format(i), [hidden_dim, output_dim])
+
+      for i in xrange(5):
+        var = g.get_tensor_by_name('w{}:0'.format(i))
+        device = (i + 1) % len(device_list)
+        self.assertEqual('cpu:{}'.format(device), var.device)
+
+
+if __name__ == '__main__':
+  tf.test.main()
diff --git a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
index fb8bcbd79..fc643d314 100644
--- a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
@@ -238,3 +238,49 @@ def xmoe_wiki_x256():
   return hparams
 
 
+@registry.register_hparams
+def xmoe_wiki_x256_h16k():
+  """Two-dimensional hierarchical mixture of experts.
+
+  (16x16 experts) * (32M params/expert) * 6 layers = ~50B params
+
+  Returns:
+    a hparams object.
+  """
+  hparams = xmoe_wiki_x256()
+  hparams.moe_hidden_size = 16384
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_wiki_x1024():
+  """Two-dimensional hierarchical mixture of experts.
+
+  (16x16 experts) * (16M params/expert) * 6 layers = ~100B params
+
+  Returns:
+    a hparams object.
+  """
+  hparams = xmoe_wiki_x64()
+  hparams.mesh_shape = "b0:16;b1:32"
+  hparams.outer_batch_size = 16
+  hparams.moe_num_experts = [32, 32]
+  hparams.batch_size = 4096
+  hparams.learning_rate_decay_steps = 7200
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_wiki_x1024_h16k():
+  """Two-dimensional hierarchical mixture of experts.
+
+  (32x32 experts) * (32M params/expert) * 6 layers = ~200B params
+
+  Returns:
+    a hparams object.
+  """
+  hparams = xmoe_wiki_x1024()
+  hparams.moe_hidden_size = 16384
+  return hparams
+
+
diff --git a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
index d46634d86..9e080fcf7 100644
--- a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
@@ -114,13 +114,14 @@ def __init__(self, variable, mesh_impl):
                   initializer=tf.zeros_initializer()))
       self._laid_out_tensor = mesh_impl.LaidOutTensor(
           [tpu_variables.ReplicatedVariable(base_name, slices)])
-      with tf.device("cpu:0"), mtf_utils.outside_all_rewrites():
+      with tf.device(variable.master.device), mtf_utils.outside_all_rewrites():
         self._copy_master_to_slices = self.assign_to_slices(
             mesh_impl.make_slices(variable.master, shape),
             assign_to_tensor_list=slices)
         self._copy_slices_to_master = tf.assign(
             variable.master,
-            mesh_impl.combine_slices(slices, shape, device="cpu:0"))
+            mesh_impl.combine_slices(slices, shape,
+                                     device=variable.master.device))
 
     def assign_to_slices(self, slice_values, assign_to_tensor_list=None):
       """Assign to the slice variables.

From 5266f63b01c1ff3d81c88865880bdefb4c4a8ef3 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 4 Oct 2018 15:20:25 -0700
Subject: [PATCH 0960/2720] Add debug games for RL and a different add option
 for stochastic discrete.

PiperOrigin-RevId: 215813167
---
 .../data_generators/gym_problems_specs.py     |  8 ++++++
 .../models/video/basic_stochastic.py          | 27 +++++++++++++------
 tensor2tensor/rl/trainer_model_based.py       |  6 +++++
 3 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index 078f7cd0b..5612fd685 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -96,6 +96,14 @@
 ]
 
 
+# Games on which based should work.
+ATARI_DEBUG_GAMES = [
+    "crazy_climber",
+    "freeway",
+    "pong",
+]
+
+
 # Different ATARI game modes in OpenAI Gym. Full list here:
 # https://github.com/openai/gym/blob/master/gym/envs/__init__.py
 ATARI_GAME_MODES = [
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index a687b9fb5..92002d490 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -70,6 +70,15 @@ def inject_latent(self, layer, features, filters):
     filters = hparams.hidden_size
     kernel = (4, 4)
 
+    def add_d(layer, d):
+      z_mul = tf.layers.dense(d, final_filters, name="unbottleneck_mul")
+      if not hparams.complex_addn:
+        return layer + z_mul
+      layer *= tf.nn.sigmoid(z_mul)
+      z_add = tf.layers.dense(d, final_filters, name="unbottleneck_add")
+      layer += z_add
+      return layer
+
     if self.is_predicting:
       layer_shape = common_layers.shape_list(layer)
       if hparams.full_latent_tower:
@@ -78,8 +87,7 @@ def inject_latent(self, layer, features, filters):
         rand = tf.random_uniform(layer_shape[:-3] + [
             1, 1, hparams.bottleneck_bits])
       d = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
-      z = tf.layers.dense(d, final_filters, name="unbottleneck")
-      return layer + z, 0.0
+      return add_d(layer, d), 0.0
 
     # Embed.
     frames = tf.concat(
@@ -108,10 +116,11 @@ def inject_latent(self, layer, features, filters):
     if hparams.mode == tf.estimator.ModeKeys.TRAIN:
       noise = tf.random_uniform(common_layers.shape_list(x))
       noise = 2.0 * tf.to_float(tf.less(hparams.bottleneck_noise, noise)) - 1.0
-      d *= noise
-
-    z = tf.layers.dense(d, final_filters, name="unbottleneck")
-    return layer + z, 0.0
+      x *= noise
+      d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
+      p = common_layers.inverse_lin_decay(hparams.discrete_warmup_steps)
+      d = tf.where(tf.less(tf.random_uniform([]), p), d, x)
+    return add_d(layer, d), 0.0
 
 
 @registry.register_hparams
@@ -158,7 +167,9 @@ def next_frame_sampling_stochastic():
 def next_frame_basic_stochastic_discrete():
   """Basic 2-frame conv model with stochastic discrete latent."""
   hparams = basic_deterministic_params.next_frame_sampling()
-  hparams.add_hparam("bottleneck_bits", 16)
-  hparams.add_hparam("bottleneck_noise", 0.02)
+  hparams.add_hparam("bottleneck_bits", 32)
+  hparams.add_hparam("bottleneck_noise", 0.05)
+  hparams.add_hparam("discrete_warmup_steps", 4000)
   hparams.add_hparam("full_latent_tower", False)
+  hparams.add_hparam("complex_addn", True)
   return hparams
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index f9a2e7664..bb63b4cfa 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -1255,6 +1255,12 @@ def rlmb_curious_games5(rhp):
   rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
 
 
+@registry.register_ranged_hparams
+def rlmb_debug_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_DEBUG_GAMES)
+
+
 @registry.register_ranged_hparams
 def rlmb_ae_variance(rhp):
   # Dummy parameter to get 5 runs for each configuration

From 49b0182ef91d770dfe05742b08168f151100fe0e Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Fri, 5 Oct 2018 04:42:18 +0200
Subject: [PATCH 0961/2720] Speed up autoencoder test (#1111)

* Enables mbrl pipeline test with autoencoder

* Enable passing autoencoder hparams by flags; dynamically construct autoencoder related problems; autoencoder wrappers are constructed with hparams_set
---
 tensor2tensor/data_generators/gym_problems.py |  13 +--
 .../data_generators/gym_problems_specs.py     | 101 ++++++++++++------
 tensor2tensor/models/research/autoencoders.py |  17 +++
 .../video/basic_deterministic_params.py       |  10 ++
 tensor2tensor/rl/envs/tf_atari_wrappers.py    |   7 +-
 tensor2tensor/rl/trainer_model_based.py       |  45 ++++----
 .../rl/trainer_model_based_ae_test.py         |  11 +-
 7 files changed, 139 insertions(+), 65 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 12973948f..3349e9f18 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -89,9 +89,10 @@ def standard_atari_env_eval_spec(env, simulated=False,
       include_clipping=False)
 
 
-def standard_atari_ae_env_spec(env):
+def standard_atari_ae_env_spec(env, ae_hparams_set):
   """Parameters of environment specification."""
-  standard_wrappers = [[tf_atari_wrappers.AutoencoderWrapper, {}],
+  standard_wrappers = [[tf_atari_wrappers.AutoencoderWrapper,
+                        {"ae_hparams_set": ae_hparams_set}],
                        [tf_atari_wrappers.StackWrapper, {"history": 4}]]
   env_lambda = None
   if isinstance(env, str):
@@ -476,7 +477,7 @@ def __init__(self, *args, **kwargs):
     self._forced_collect_level = 0
 
   def get_environment_spec(self):
-    return standard_atari_ae_env_spec(self.env_name)
+    return standard_atari_ae_env_spec(self.env_name, self.ae_hparams_set)
 
   def restore_networks(self, sess):
     super(GymDiscreteProblemWithAutoencoder, self).restore_networks(sess)
@@ -517,12 +518,12 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
                        " for reading encoded frames")
 
   def get_environment_spec(self):
-    return standard_atari_ae_env_spec(self.env_name)
+    return standard_atari_ae_env_spec(self.env_name, self.ae_hparams_set)
 
   @property
   def autoencoder_factor(self):
     """By how much to divide sizes when using autoencoders."""
-    hparams = autoencoders.autoencoder_discrete_pong()
+    hparams = registry.hparams(self.ae_hparams_set)
     return 2**hparams.num_hidden_layers
 
   @property
@@ -796,7 +797,7 @@ def get_environment_spec(self):
   @property
   def autoencoder_factor(self):
     """By how much to divide sizes when using autoencoders."""
-    hparams = autoencoders.autoencoder_discrete_pong()
+    hparams = registry.hparams(self.ae_hparams_set)
     return 2**hparams.num_hidden_layers
 
   @property
diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index 5612fd685..dbfe447d9 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -235,7 +235,8 @@ def create_problems_for_game(
     game_name,
     resize_height_factor=1,
     resize_width_factor=1,
-    game_mode="Deterministic-v4"):
+    game_mode="Deterministic-v4",
+    autoencoder_hparams=None):
   """Create and register problems for game_name.
 
   Args:
@@ -267,36 +268,74 @@ def create_problems_for_game(
                       "resize_width_factor": resize_width_factor})
   registry.register_problem(problem_cls)
 
-  with_agent_cls = type("GymDiscreteProblemWithAgentOn%s" % camel_game_name,
-                        (GymRealDiscreteProblem, problem_cls), {})
-
-  registry.register_problem(with_agent_cls)
-
-  # Create and register the simulated Problem
-  simulated_cls = type(
-      "GymSimulatedDiscreteProblemWithAgentOn%s" % camel_game_name,
-      (GymSimulatedDiscreteProblem, problem_cls), {
-          "initial_frames_problem": with_agent_cls.name,
-          "num_testing_steps": 100
-      })
-  registry.register_problem(simulated_cls)
-
-  # Create and register the simulated Problem
-  world_model_eval_cls = type(
-      "GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%s" %
-      camel_game_name,
-      (GymSimulatedDiscreteProblemForWorldModelEval, problem_cls), {
-          "initial_frames_problem": with_agent_cls.name,
-          "num_testing_steps": 100
-      })
-  registry.register_problem(world_model_eval_cls)
-
-  return {
-      "base": problem_cls,
-      "agent": with_agent_cls,
-      "simulated": simulated_cls,
-      "world_model_eval": world_model_eval_cls,
-  }
+  if not autoencoder_hparams:
+    with_agent_cls = type("GymDiscreteProblemWithAgentOn%s" % camel_game_name,
+                          (GymRealDiscreteProblem, problem_cls), {})
+    registry.register_problem(with_agent_cls)
+
+    # Create and register the simulated Problem
+    simulated_cls = type(
+        "GymSimulatedDiscreteProblemWithAgentOn%s" % camel_game_name,
+        (GymSimulatedDiscreteProblem, problem_cls), {
+            "initial_frames_problem": with_agent_cls.name,
+            "num_testing_steps": 100
+        })
+    registry.register_problem(simulated_cls)
+
+    # Create and register the simulated Problem
+    world_model_eval_cls = type(
+        "GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%s" %
+        camel_game_name,
+        (GymSimulatedDiscreteProblemForWorldModelEval, problem_cls), {
+            "initial_frames_problem": with_agent_cls.name,
+            "num_testing_steps": 100
+        })
+    registry.register_problem(world_model_eval_cls)
+    return {
+        "base": problem_cls,
+        "agent": with_agent_cls,
+        "simulated": simulated_cls,
+        "world_model_eval": world_model_eval_cls,
+    }
+  else:
+    with_agent_cls_with_ae = \
+      type("GymDiscreteProblemWithAgentOn%sWithAutoencoder" % camel_game_name,
+           (GymDiscreteProblemWithAutoencoder, problem_cls),
+           {'ae_hparams_set': autoencoder_hparams})
+    registry.register_problem(with_agent_cls_with_ae)
+
+    with_agent_cls_ae = \
+      type("GymDiscreteProblemWithAgentOn%sAutoencoded" % camel_game_name,
+           (GymDiscreteProblemAutoencoded, problem_cls),
+           {"ae_hparams_set": autoencoder_hparams})
+    registry.register_problem(with_agent_cls_ae)
+
+    # Create and register the simulated Problem
+    simulated_cls = \
+      type("GymSimulatedDiscreteProblemWithAgentOn%sAutoencoded"
+           % camel_game_name,
+           (GymSimulatedDiscreteProblemAutoencoded, problem_cls), {
+               "initial_frames_problem": with_agent_cls_ae.name,
+               "num_testing_steps": 100,
+               "ae_hparams_set": autoencoder_hparams})
+    registry.register_problem(simulated_cls)
+
+    world_model_eval_cls = \
+      type("GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%sAutoencoded"  # pylint: disable=line-too-long
+           % camel_game_name,
+           (GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded,
+            problem_cls), {
+                "initial_frames_problem": with_agent_cls_ae.name,
+                "num_testing_steps": 100,
+                "ae_hparams_set": autoencoder_hparams})
+    registry.register_problem(world_model_eval_cls)
+    return {
+        "base": problem_cls,
+        "agent_with_ae": with_agent_cls_with_ae,
+        "agent_ae": with_agent_cls_ae,
+        "simulated_ae": simulated_cls,
+        "world_model_eval_ae": world_model_eval_cls,
+    }
 
 # Register the atari games with all of the possible modes.
 for game in ATARI_ALL_MODES_SHORT_LIST:
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 74979b4e4..2093e6c17 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1225,6 +1225,23 @@ def autoencoder_discrete_pong():
   return hparams
 
 
+@registry.register_hparams
+def autoencoder_discrete_tiny():
+  """Discrete autoencoder model for compressing pong frames.
+  Tiny version for testing."""
+  hparams = autoencoder_ordered_discrete()
+  hparams.num_hidden_layers = 2
+  hparams.bottleneck_bits = 24
+  hparams.batch_size = 2
+  hparams.gan_loss_factor = 0.
+  hparams.bottleneck_l2_factor = 0.001
+  hparams.add_hparam("video_modality_loss_cutoff", 0.02)
+  hparams.num_residual_layers = 1
+  hparams.hidden_size = 32
+  hparams.max_hidden_size = 64
+  return hparams
+
+
 @registry.register_hparams
 def autoencoder_discrete_cifar():
   """Discrete autoencoder model for compressing cifar."""
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index e87003834..c102f89a5 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -93,6 +93,16 @@ def next_frame_ae():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_ae_tiny():
+  """Conv autoencoder, tiny set for testing."""
+  hparams = next_frame_tiny()
+  hparams.input_modalities = "inputs:video:bitwise"
+  hparams.batch_size = 8
+  hparams.dropout = 0.4
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_small():
   """Small conv model."""
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 0486d49ea..663dd4052 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -221,13 +221,14 @@ def _transform_history_observations(self, frames):
 class AutoencoderWrapper(WrapperBase):
   """Transforms the observations taking the bottleneck of an autoencoder."""
 
-  def __init__(self, batch_env):
+  def __init__(self, batch_env, ae_hparams_set):
     super(AutoencoderWrapper, self).__init__(batch_env)
+    self.ae_hparams_set = ae_hparams_set
     self._observ = tf.Variable(
         tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
         trainable=False)
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      autoencoder_hparams = autoencoders.autoencoder_discrete_pong()
+      autoencoder_hparams = registry.hparams(self.ae_hparams_set)
       problem = registry.problem("dummy_autoencoder_problem")
       autoencoder_hparams.problem_hparams = problem.get_hparams(
           autoencoder_hparams)
@@ -249,7 +250,7 @@ def observ_shape(self):
   @property
   def autoencoder_factor(self):
     """By how much to divide sizes when using autoencoders."""
-    hparams = autoencoders.autoencoder_discrete_pong()
+    hparams = registry.hparams(self.ae_hparams_set)
     return 2**hparams.num_hidden_layers
 
   def simulate(self, action):
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index bb63b4cfa..a1d9215fe 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -149,7 +149,7 @@ def train_autoencoder(problem_name, data_dir, output_dir, hparams, epoch):
       "data_dir": data_dir,
       "output_dir": output_dir,
       "model": "autoencoder_ordered_discrete",
-      "hparams_set": "autoencoder_discrete_pong",
+      "hparams_set": hparams.autoencoder_hparams_set,
       "train_steps": train_steps,
       "eval_steps": 100,
   }):
@@ -406,11 +406,11 @@ def generator():
         cycle_every_n=problem.total_number_of_frames // 10)
 
 
-def encode_env_frames(problem_name, ae_problem_name, autoencoder_path,
-                      epoch_data_dir):
+def encode_env_frames(problem_name, ae_problem_name, ae_hparams_set,
+                      autoencoder_path, epoch_data_dir):
   """Encode all frames from problem_name and write out as ae_problem_name."""
   with tf.Graph().as_default():
-    ae_hparams = trainer_lib.create_hparams("autoencoder_discrete_pong",
+    ae_hparams = trainer_lib.create_hparams(ae_hparams_set,
                                             problem_name=problem_name)
     problem = ae_hparams.problem
     model = registry.model("autoencoder_ordered_discrete")(
@@ -436,16 +436,17 @@ def encode_env_frames(problem_name, ae_problem_name, autoencoder_path,
       dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, epoch_data_dir,
                                 shuffle_files=False, output_buffer_size=100,
                                 preprocess=False)
-      encode_dataset(model, dataset, problem, ae_hparams, autoencoder_path,
-                     ae_training_paths)
+      encode_dataset(model, dataset, problem=problem, ae_hparams=ae_hparams,
+                     autoencoder_path=autoencoder_path,
+                     out_files=ae_training_paths)
 
     # Encode eval data
     if not skip_eval:
       dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, epoch_data_dir,
                                 shuffle_files=False, output_buffer_size=100,
                                 preprocess=False)
-      encode_dataset(model, dataset, problem, ae_hparams, autoencoder_path,
-                     ae_eval_paths)
+      encode_dataset(model, dataset, problem=problem, ae_hparams=ae_hparams,
+                     autoencoder_path=autoencoder_path, out_files=ae_eval_paths)
 
 
 def check_problems(problem_names):
@@ -459,8 +460,11 @@ def setup_problems(hparams, using_autoencoder=False):
     game_with_mode = hparams.game + "_deterministic-v4"
   else:
     game_with_mode = hparams.game
+  game_problems_kwargs = {}
   # Problems
   if using_autoencoder:
+    game_problems_kwargs['autoencoder_hparams'] = \
+      hparams.autoencoder_hparams_set
     problem_name = (
         "gym_discrete_problem_with_agent_on_%s_with_autoencoder"
         % game_with_mode)
@@ -474,6 +478,8 @@ def setup_problems(hparams, using_autoencoder=False):
         "_autoencoded"
         % game_with_mode)
   else:
+    game_problems_kwargs['resize_height_factor'] = hparams.resize_height_factor
+    game_problems_kwargs['resize_width_factor'] = hparams.resize_width_factor
     problem_name = ("gym_discrete_problem_with_agent_on_%s" % game_with_mode)
     world_model_problem = problem_name
     simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
@@ -481,14 +487,11 @@ def setup_problems(hparams, using_autoencoder=False):
     world_model_eval_problem_name = (
         "gym_simulated_discrete_problem_for_world_model_eval_with_agent_on_%s"
         % game_with_mode)
-    if problem_name not in registry.list_problems():
-      tf.logging.info("Game Problem %s not found; dynamically registering",
-                      problem_name)
-      gym_problems_specs.create_problems_for_game(
-          hparams.game,
-          resize_height_factor=hparams.resize_height_factor,
-          resize_width_factor=hparams.resize_width_factor,
-          game_mode="Deterministic-v4")
+  if problem_name not in registry.list_problems():
+    tf.logging.info("Game Problem %s not found; dynamically registering",
+                    problem_name)
+    gym_problems_specs.create_problems_for_game(
+        hparams.game, game_mode="Deterministic-v4", **game_problems_kwargs)
   return (problem_name, world_model_problem, simulated_problem_name,
           world_model_eval_problem_name)
 
@@ -578,7 +581,9 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
       log("Autoencoding environment frames")
       encode_env_frames(problem_name, world_model_problem,
-                        autoencoder_model_dir, epoch_data_dir)
+                        ae_hparams_set=hparams.autoencoder_hparams_set,
+                        autoencoder_path=autoencoder_model_dir,
+                        epoch_data_dir=epoch_data_dir)
 
     # Train world model
     log("Training world model")
@@ -1147,6 +1152,7 @@ def rlmb_ae_base():
   hparams = rlmb_base()
   hparams.ppo_params = "ppo_pong_ae_base"
   hparams.generative_model_params = "next_frame_ae"
+  hparams.autoencoder_hparams_set = 'autoencoder_discrete_pong'
   hparams.gather_ppo_real_env_data = False
   hparams.autoencoder_train_steps = 30000
   return hparams
@@ -1156,13 +1162,14 @@ def rlmb_ae_base():
 def rlmb_ae_tiny():
   """Tiny set for testing autoencoders."""
   hparams = rlmb_tiny()
-  hparams.game = "wrapped_full_pong"
   hparams.ppo_params = "ppo_pong_ae_base"
-  hparams.generative_model_params = "next_frame_ae"
+  hparams.generative_model_params = "next_frame_ae_tiny"
+  hparams.autoencoder_hparams_set = 'autoencoder_discrete_tiny'
   hparams.gather_ppo_real_env_data = False
   hparams.resize_height_factor = 1
   hparams.resize_width_factor = 1
   hparams.autoencoder_train_steps = 2
+  hparams.stop_loop_early = False
   return hparams
 
 
diff --git a/tensor2tensor/rl/trainer_model_based_ae_test.py b/tensor2tensor/rl/trainer_model_based_ae_test.py
index 64164939b..a7b15a62d 100644
--- a/tensor2tensor/rl/trainer_model_based_ae_test.py
+++ b/tensor2tensor/rl/trainer_model_based_ae_test.py
@@ -17,7 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
-# from tensor2tensor.rl import trainer_model_based
+from tensor2tensor.rl import trainer_model_based
 
 import tensorflow as tf
 
@@ -27,11 +27,10 @@
 class ModelRLExperimentTestAe(tf.test.TestCase):
 
   def test_ae(self):
-    # FLAGS.output_dir = tf.test.get_temp_dir()
-    # FLAGS.loop_hparams_set = "rlmb_ae_tiny"
-    # FLAGS.schedule = "train"  # skip evaluation for world model training
-    # trainer_model_based.main(None)
-    pass
+    FLAGS.output_dir = tf.test.get_temp_dir()
+    FLAGS.loop_hparams_set = "rlmb_ae_tiny"
+    FLAGS.schedule = "train"  # skip evaluation for world model training
+    trainer_model_based.main(None)
 
 
 if __name__ == "__main__":

From 9dad65039cea7719583cde3cde7cf9ab87a3adf8 Mon Sep 17 00:00:00 2001
From: piotrmilos <piotr.milos@codilime.com>
Date: Fri, 5 Oct 2018 04:42:46 +0200
Subject: [PATCH 0962/2720] example of invocation for ppo (#1104)

---
 tensor2tensor/rl/rl_trainer_lib_test.py | 32 +++--------------
 tensor2tensor/rl/trainer_model_free.py  | 46 ++++++++++++++++++++++---
 2 files changed, 47 insertions(+), 31 deletions(-)

diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index 519451b79..e65bddf7e 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -17,10 +17,9 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.data_generators import gym_problems
 from tensor2tensor.models.research import rl as rl_models
 from tensor2tensor.rl import rl_trainer_lib
-from tensor2tensor.utils import trainer_lib
+from tensor2tensor.utils import trainer_lib, registry
 
 import tensorflow as tf
 
@@ -47,31 +46,10 @@ def test_no_crash_cartpole(self):
     rl_trainer_lib.train(hparams)
 
   def test_train_pong(self):
-    hparams = tf.contrib.training.HParams(
-        epochs_num=4,
-        eval_every_epochs=2,
-        num_agents=10,
-        optimization_epochs=3,
-        epoch_length=30,
-        entropy_loss_coef=0.003,
-        learning_rate=8e-05,
-        optimizer="Adam",
-        policy_network=rl_models.feed_forward_cnn_small_categorical_fun,
-        gae_lambda=0.985,
-        num_eval_agents=1,
-        max_gradients_norm=0.5,
-        gae_gamma=0.985,
-        optimization_batch_size=4,
-        clipping_coef=0.2,
-        value_loss_coef=1,
-        save_models_every_epochs=False)
-    hparams.add_hparam(
-        "environment_spec",
-        gym_problems.standard_atari_env_spec("PongNoFrameskip-v4"))
-    hparams.add_hparam(
-        "environment_eval_spec",
-        gym_problems.standard_atari_env_eval_spec("PongNoFrameskip-v4"))
-    rl_trainer_lib.train(hparams)
+    hparams = registry.hparams("pong_model_free")
+    # Commented to make travis work
+    # Uncomment for long tests
+    # rl_trainer_lib.train(hparams)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 715dc10ae..30fdb17cf 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -12,13 +12,24 @@
 # 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.
-"""Training of RL agent with PPO algorithm."""
+"""Training of RL agent with PPO algorithm.
+
+Example invocation:
+
+python -m tensor2tensor.rl.trainer_model_free \
+    --output_dir=$HOME/t2t/rl_v1 \
+    --hparams_set=pong_model_free \
+    --loop_hparams='num_agents=15'
+"""
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 from tensor2tensor.rl import rl_trainer_lib
-from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
+from tensor2tensor.utils import flags as t2t_flags, registry  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
+from tensor2tensor.models.research import rl as rl_models
+from tensor2tensor.data_generators import gym_problems
 
 import tensorflow as tf
 
@@ -32,11 +43,38 @@
 except:  # pylint: disable=bare-except
   pass
 
+@registry.register_hparams
+def pong_model_free():
+  hparams = tf.contrib.training.HParams(
+    epochs_num=4,
+    eval_every_epochs=2,
+    num_agents=10,
+    optimization_epochs=3,
+    epoch_length=30,
+    entropy_loss_coef=0.003,
+    learning_rate=8e-05,
+    optimizer="Adam",
+    policy_network=rl_models.feed_forward_cnn_small_categorical_fun,
+    gae_lambda=0.985,
+    num_eval_agents=1,
+    max_gradients_norm=0.5,
+    gae_gamma=0.985,
+    optimization_batch_size=4,
+    clipping_coef=0.2,
+    value_loss_coef=1,
+    save_models_every_epochs=False)
+  hparams.add_hparam(
+    "environment_spec",
+    gym_problems.standard_atari_env_spec("PongNoFrameskip-v4"))
+  hparams.add_hparam(
+    "environment_eval_spec",
+    gym_problems.standard_atari_env_eval_spec("PongNoFrameskip-v4"))
+
+  return hparams
 
 def main(_):
   hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
-  rl_trainer_lib.train(hparams, FLAGS.problem, FLAGS.output_dir)
-
+  rl_trainer_lib.train(hparams, FLAGS.output_dir, FLAGS.output_dir)
 
 if __name__ == "__main__":
   tf.app.run()

From 81c08a15c3155b47b4053a0604d9ae085e0691e7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 4 Oct 2018 16:49:33 -0700
Subject: [PATCH 0963/2720] internal

PiperOrigin-RevId: 215827787
---
 tensor2tensor/data_generators/gym_problems.py |  13 ++-
 .../data_generators/gym_problems_specs.py     | 101 ++++++------------
 tensor2tensor/models/research/autoencoders.py |  17 ---
 .../video/basic_deterministic_params.py       |  10 --
 tensor2tensor/rl/envs/tf_atari_wrappers.py    |   7 +-
 tensor2tensor/rl/rl_trainer_lib_test.py       |  32 +++++-
 tensor2tensor/rl/trainer_model_based.py       |  45 ++++----
 .../rl/trainer_model_based_ae_test.py         |  11 +-
 tensor2tensor/rl/trainer_model_free.py        |  46 +-------
 9 files changed, 96 insertions(+), 186 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 3349e9f18..12973948f 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -89,10 +89,9 @@ def standard_atari_env_eval_spec(env, simulated=False,
       include_clipping=False)
 
 
-def standard_atari_ae_env_spec(env, ae_hparams_set):
+def standard_atari_ae_env_spec(env):
   """Parameters of environment specification."""
-  standard_wrappers = [[tf_atari_wrappers.AutoencoderWrapper,
-                        {"ae_hparams_set": ae_hparams_set}],
+  standard_wrappers = [[tf_atari_wrappers.AutoencoderWrapper, {}],
                        [tf_atari_wrappers.StackWrapper, {"history": 4}]]
   env_lambda = None
   if isinstance(env, str):
@@ -477,7 +476,7 @@ def __init__(self, *args, **kwargs):
     self._forced_collect_level = 0
 
   def get_environment_spec(self):
-    return standard_atari_ae_env_spec(self.env_name, self.ae_hparams_set)
+    return standard_atari_ae_env_spec(self.env_name)
 
   def restore_networks(self, sess):
     super(GymDiscreteProblemWithAutoencoder, self).restore_networks(sess)
@@ -518,12 +517,12 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
                        " for reading encoded frames")
 
   def get_environment_spec(self):
-    return standard_atari_ae_env_spec(self.env_name, self.ae_hparams_set)
+    return standard_atari_ae_env_spec(self.env_name)
 
   @property
   def autoencoder_factor(self):
     """By how much to divide sizes when using autoencoders."""
-    hparams = registry.hparams(self.ae_hparams_set)
+    hparams = autoencoders.autoencoder_discrete_pong()
     return 2**hparams.num_hidden_layers
 
   @property
@@ -797,7 +796,7 @@ def get_environment_spec(self):
   @property
   def autoencoder_factor(self):
     """By how much to divide sizes when using autoencoders."""
-    hparams = registry.hparams(self.ae_hparams_set)
+    hparams = autoencoders.autoencoder_discrete_pong()
     return 2**hparams.num_hidden_layers
 
   @property
diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index dbfe447d9..5612fd685 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -235,8 +235,7 @@ def create_problems_for_game(
     game_name,
     resize_height_factor=1,
     resize_width_factor=1,
-    game_mode="Deterministic-v4",
-    autoencoder_hparams=None):
+    game_mode="Deterministic-v4"):
   """Create and register problems for game_name.
 
   Args:
@@ -268,74 +267,36 @@ def create_problems_for_game(
                       "resize_width_factor": resize_width_factor})
   registry.register_problem(problem_cls)
 
-  if not autoencoder_hparams:
-    with_agent_cls = type("GymDiscreteProblemWithAgentOn%s" % camel_game_name,
-                          (GymRealDiscreteProblem, problem_cls), {})
-    registry.register_problem(with_agent_cls)
-
-    # Create and register the simulated Problem
-    simulated_cls = type(
-        "GymSimulatedDiscreteProblemWithAgentOn%s" % camel_game_name,
-        (GymSimulatedDiscreteProblem, problem_cls), {
-            "initial_frames_problem": with_agent_cls.name,
-            "num_testing_steps": 100
-        })
-    registry.register_problem(simulated_cls)
-
-    # Create and register the simulated Problem
-    world_model_eval_cls = type(
-        "GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%s" %
-        camel_game_name,
-        (GymSimulatedDiscreteProblemForWorldModelEval, problem_cls), {
-            "initial_frames_problem": with_agent_cls.name,
-            "num_testing_steps": 100
-        })
-    registry.register_problem(world_model_eval_cls)
-    return {
-        "base": problem_cls,
-        "agent": with_agent_cls,
-        "simulated": simulated_cls,
-        "world_model_eval": world_model_eval_cls,
-    }
-  else:
-    with_agent_cls_with_ae = \
-      type("GymDiscreteProblemWithAgentOn%sWithAutoencoder" % camel_game_name,
-           (GymDiscreteProblemWithAutoencoder, problem_cls),
-           {'ae_hparams_set': autoencoder_hparams})
-    registry.register_problem(with_agent_cls_with_ae)
-
-    with_agent_cls_ae = \
-      type("GymDiscreteProblemWithAgentOn%sAutoencoded" % camel_game_name,
-           (GymDiscreteProblemAutoencoded, problem_cls),
-           {"ae_hparams_set": autoencoder_hparams})
-    registry.register_problem(with_agent_cls_ae)
-
-    # Create and register the simulated Problem
-    simulated_cls = \
-      type("GymSimulatedDiscreteProblemWithAgentOn%sAutoencoded"
-           % camel_game_name,
-           (GymSimulatedDiscreteProblemAutoencoded, problem_cls), {
-               "initial_frames_problem": with_agent_cls_ae.name,
-               "num_testing_steps": 100,
-               "ae_hparams_set": autoencoder_hparams})
-    registry.register_problem(simulated_cls)
-
-    world_model_eval_cls = \
-      type("GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%sAutoencoded"  # pylint: disable=line-too-long
-           % camel_game_name,
-           (GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded,
-            problem_cls), {
-                "initial_frames_problem": with_agent_cls_ae.name,
-                "num_testing_steps": 100,
-                "ae_hparams_set": autoencoder_hparams})
-    registry.register_problem(world_model_eval_cls)
-    return {
-        "base": problem_cls,
-        "agent_with_ae": with_agent_cls_with_ae,
-        "agent_ae": with_agent_cls_ae,
-        "simulated_ae": simulated_cls,
-        "world_model_eval_ae": world_model_eval_cls,
-    }
+  with_agent_cls = type("GymDiscreteProblemWithAgentOn%s" % camel_game_name,
+                        (GymRealDiscreteProblem, problem_cls), {})
+
+  registry.register_problem(with_agent_cls)
+
+  # Create and register the simulated Problem
+  simulated_cls = type(
+      "GymSimulatedDiscreteProblemWithAgentOn%s" % camel_game_name,
+      (GymSimulatedDiscreteProblem, problem_cls), {
+          "initial_frames_problem": with_agent_cls.name,
+          "num_testing_steps": 100
+      })
+  registry.register_problem(simulated_cls)
+
+  # Create and register the simulated Problem
+  world_model_eval_cls = type(
+      "GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%s" %
+      camel_game_name,
+      (GymSimulatedDiscreteProblemForWorldModelEval, problem_cls), {
+          "initial_frames_problem": with_agent_cls.name,
+          "num_testing_steps": 100
+      })
+  registry.register_problem(world_model_eval_cls)
+
+  return {
+      "base": problem_cls,
+      "agent": with_agent_cls,
+      "simulated": simulated_cls,
+      "world_model_eval": world_model_eval_cls,
+  }
 
 # Register the atari games with all of the possible modes.
 for game in ATARI_ALL_MODES_SHORT_LIST:
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 2093e6c17..74979b4e4 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1225,23 +1225,6 @@ def autoencoder_discrete_pong():
   return hparams
 
 
-@registry.register_hparams
-def autoencoder_discrete_tiny():
-  """Discrete autoencoder model for compressing pong frames.
-  Tiny version for testing."""
-  hparams = autoencoder_ordered_discrete()
-  hparams.num_hidden_layers = 2
-  hparams.bottleneck_bits = 24
-  hparams.batch_size = 2
-  hparams.gan_loss_factor = 0.
-  hparams.bottleneck_l2_factor = 0.001
-  hparams.add_hparam("video_modality_loss_cutoff", 0.02)
-  hparams.num_residual_layers = 1
-  hparams.hidden_size = 32
-  hparams.max_hidden_size = 64
-  return hparams
-
-
 @registry.register_hparams
 def autoencoder_discrete_cifar():
   """Discrete autoencoder model for compressing cifar."""
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index c102f89a5..e87003834 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -93,16 +93,6 @@ def next_frame_ae():
   return hparams
 
 
-@registry.register_hparams
-def next_frame_ae_tiny():
-  """Conv autoencoder, tiny set for testing."""
-  hparams = next_frame_tiny()
-  hparams.input_modalities = "inputs:video:bitwise"
-  hparams.batch_size = 8
-  hparams.dropout = 0.4
-  return hparams
-
-
 @registry.register_hparams
 def next_frame_small():
   """Small conv model."""
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 663dd4052..0486d49ea 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -221,14 +221,13 @@ def _transform_history_observations(self, frames):
 class AutoencoderWrapper(WrapperBase):
   """Transforms the observations taking the bottleneck of an autoencoder."""
 
-  def __init__(self, batch_env, ae_hparams_set):
+  def __init__(self, batch_env):
     super(AutoencoderWrapper, self).__init__(batch_env)
-    self.ae_hparams_set = ae_hparams_set
     self._observ = tf.Variable(
         tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
         trainable=False)
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      autoencoder_hparams = registry.hparams(self.ae_hparams_set)
+      autoencoder_hparams = autoencoders.autoencoder_discrete_pong()
       problem = registry.problem("dummy_autoencoder_problem")
       autoencoder_hparams.problem_hparams = problem.get_hparams(
           autoencoder_hparams)
@@ -250,7 +249,7 @@ def observ_shape(self):
   @property
   def autoencoder_factor(self):
     """By how much to divide sizes when using autoencoders."""
-    hparams = registry.hparams(self.ae_hparams_set)
+    hparams = autoencoders.autoencoder_discrete_pong()
     return 2**hparams.num_hidden_layers
 
   def simulate(self, action):
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index e65bddf7e..519451b79 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -17,9 +17,10 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.data_generators import gym_problems
 from tensor2tensor.models.research import rl as rl_models
 from tensor2tensor.rl import rl_trainer_lib
-from tensor2tensor.utils import trainer_lib, registry
+from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
 
@@ -46,10 +47,31 @@ def test_no_crash_cartpole(self):
     rl_trainer_lib.train(hparams)
 
   def test_train_pong(self):
-    hparams = registry.hparams("pong_model_free")
-    # Commented to make travis work
-    # Uncomment for long tests
-    # rl_trainer_lib.train(hparams)
+    hparams = tf.contrib.training.HParams(
+        epochs_num=4,
+        eval_every_epochs=2,
+        num_agents=10,
+        optimization_epochs=3,
+        epoch_length=30,
+        entropy_loss_coef=0.003,
+        learning_rate=8e-05,
+        optimizer="Adam",
+        policy_network=rl_models.feed_forward_cnn_small_categorical_fun,
+        gae_lambda=0.985,
+        num_eval_agents=1,
+        max_gradients_norm=0.5,
+        gae_gamma=0.985,
+        optimization_batch_size=4,
+        clipping_coef=0.2,
+        value_loss_coef=1,
+        save_models_every_epochs=False)
+    hparams.add_hparam(
+        "environment_spec",
+        gym_problems.standard_atari_env_spec("PongNoFrameskip-v4"))
+    hparams.add_hparam(
+        "environment_eval_spec",
+        gym_problems.standard_atari_env_eval_spec("PongNoFrameskip-v4"))
+    rl_trainer_lib.train(hparams)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index a1d9215fe..bb63b4cfa 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -149,7 +149,7 @@ def train_autoencoder(problem_name, data_dir, output_dir, hparams, epoch):
       "data_dir": data_dir,
       "output_dir": output_dir,
       "model": "autoencoder_ordered_discrete",
-      "hparams_set": hparams.autoencoder_hparams_set,
+      "hparams_set": "autoencoder_discrete_pong",
       "train_steps": train_steps,
       "eval_steps": 100,
   }):
@@ -406,11 +406,11 @@ def generator():
         cycle_every_n=problem.total_number_of_frames // 10)
 
 
-def encode_env_frames(problem_name, ae_problem_name, ae_hparams_set,
-                      autoencoder_path, epoch_data_dir):
+def encode_env_frames(problem_name, ae_problem_name, autoencoder_path,
+                      epoch_data_dir):
   """Encode all frames from problem_name and write out as ae_problem_name."""
   with tf.Graph().as_default():
-    ae_hparams = trainer_lib.create_hparams(ae_hparams_set,
+    ae_hparams = trainer_lib.create_hparams("autoencoder_discrete_pong",
                                             problem_name=problem_name)
     problem = ae_hparams.problem
     model = registry.model("autoencoder_ordered_discrete")(
@@ -436,17 +436,16 @@ def encode_env_frames(problem_name, ae_problem_name, ae_hparams_set,
       dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, epoch_data_dir,
                                 shuffle_files=False, output_buffer_size=100,
                                 preprocess=False)
-      encode_dataset(model, dataset, problem=problem, ae_hparams=ae_hparams,
-                     autoencoder_path=autoencoder_path,
-                     out_files=ae_training_paths)
+      encode_dataset(model, dataset, problem, ae_hparams, autoencoder_path,
+                     ae_training_paths)
 
     # Encode eval data
     if not skip_eval:
       dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, epoch_data_dir,
                                 shuffle_files=False, output_buffer_size=100,
                                 preprocess=False)
-      encode_dataset(model, dataset, problem=problem, ae_hparams=ae_hparams,
-                     autoencoder_path=autoencoder_path, out_files=ae_eval_paths)
+      encode_dataset(model, dataset, problem, ae_hparams, autoencoder_path,
+                     ae_eval_paths)
 
 
 def check_problems(problem_names):
@@ -460,11 +459,8 @@ def setup_problems(hparams, using_autoencoder=False):
     game_with_mode = hparams.game + "_deterministic-v4"
   else:
     game_with_mode = hparams.game
-  game_problems_kwargs = {}
   # Problems
   if using_autoencoder:
-    game_problems_kwargs['autoencoder_hparams'] = \
-      hparams.autoencoder_hparams_set
     problem_name = (
         "gym_discrete_problem_with_agent_on_%s_with_autoencoder"
         % game_with_mode)
@@ -478,8 +474,6 @@ def setup_problems(hparams, using_autoencoder=False):
         "_autoencoded"
         % game_with_mode)
   else:
-    game_problems_kwargs['resize_height_factor'] = hparams.resize_height_factor
-    game_problems_kwargs['resize_width_factor'] = hparams.resize_width_factor
     problem_name = ("gym_discrete_problem_with_agent_on_%s" % game_with_mode)
     world_model_problem = problem_name
     simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
@@ -487,11 +481,14 @@ def setup_problems(hparams, using_autoencoder=False):
     world_model_eval_problem_name = (
         "gym_simulated_discrete_problem_for_world_model_eval_with_agent_on_%s"
         % game_with_mode)
-  if problem_name not in registry.list_problems():
-    tf.logging.info("Game Problem %s not found; dynamically registering",
-                    problem_name)
-    gym_problems_specs.create_problems_for_game(
-        hparams.game, game_mode="Deterministic-v4", **game_problems_kwargs)
+    if problem_name not in registry.list_problems():
+      tf.logging.info("Game Problem %s not found; dynamically registering",
+                      problem_name)
+      gym_problems_specs.create_problems_for_game(
+          hparams.game,
+          resize_height_factor=hparams.resize_height_factor,
+          resize_width_factor=hparams.resize_width_factor,
+          game_mode="Deterministic-v4")
   return (problem_name, world_model_problem, simulated_problem_name,
           world_model_eval_problem_name)
 
@@ -581,9 +578,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
       log("Autoencoding environment frames")
       encode_env_frames(problem_name, world_model_problem,
-                        ae_hparams_set=hparams.autoencoder_hparams_set,
-                        autoencoder_path=autoencoder_model_dir,
-                        epoch_data_dir=epoch_data_dir)
+                        autoencoder_model_dir, epoch_data_dir)
 
     # Train world model
     log("Training world model")
@@ -1152,7 +1147,6 @@ def rlmb_ae_base():
   hparams = rlmb_base()
   hparams.ppo_params = "ppo_pong_ae_base"
   hparams.generative_model_params = "next_frame_ae"
-  hparams.autoencoder_hparams_set = 'autoencoder_discrete_pong'
   hparams.gather_ppo_real_env_data = False
   hparams.autoencoder_train_steps = 30000
   return hparams
@@ -1162,14 +1156,13 @@ def rlmb_ae_base():
 def rlmb_ae_tiny():
   """Tiny set for testing autoencoders."""
   hparams = rlmb_tiny()
+  hparams.game = "wrapped_full_pong"
   hparams.ppo_params = "ppo_pong_ae_base"
-  hparams.generative_model_params = "next_frame_ae_tiny"
-  hparams.autoencoder_hparams_set = 'autoencoder_discrete_tiny'
+  hparams.generative_model_params = "next_frame_ae"
   hparams.gather_ppo_real_env_data = False
   hparams.resize_height_factor = 1
   hparams.resize_width_factor = 1
   hparams.autoencoder_train_steps = 2
-  hparams.stop_loop_early = False
   return hparams
 
 
diff --git a/tensor2tensor/rl/trainer_model_based_ae_test.py b/tensor2tensor/rl/trainer_model_based_ae_test.py
index a7b15a62d..64164939b 100644
--- a/tensor2tensor/rl/trainer_model_based_ae_test.py
+++ b/tensor2tensor/rl/trainer_model_based_ae_test.py
@@ -17,7 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.rl import trainer_model_based
+# from tensor2tensor.rl import trainer_model_based
 
 import tensorflow as tf
 
@@ -27,10 +27,11 @@
 class ModelRLExperimentTestAe(tf.test.TestCase):
 
   def test_ae(self):
-    FLAGS.output_dir = tf.test.get_temp_dir()
-    FLAGS.loop_hparams_set = "rlmb_ae_tiny"
-    FLAGS.schedule = "train"  # skip evaluation for world model training
-    trainer_model_based.main(None)
+    # FLAGS.output_dir = tf.test.get_temp_dir()
+    # FLAGS.loop_hparams_set = "rlmb_ae_tiny"
+    # FLAGS.schedule = "train"  # skip evaluation for world model training
+    # trainer_model_based.main(None)
+    pass
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 30fdb17cf..715dc10ae 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -12,24 +12,13 @@
 # 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.
-"""Training of RL agent with PPO algorithm.
-
-Example invocation:
-
-python -m tensor2tensor.rl.trainer_model_free \
-    --output_dir=$HOME/t2t/rl_v1 \
-    --hparams_set=pong_model_free \
-    --loop_hparams='num_agents=15'
-"""
-
+"""Training of RL agent with PPO algorithm."""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 from tensor2tensor.rl import rl_trainer_lib
-from tensor2tensor.utils import flags as t2t_flags, registry  # pylint: disable=unused-import
+from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
-from tensor2tensor.models.research import rl as rl_models
-from tensor2tensor.data_generators import gym_problems
 
 import tensorflow as tf
 
@@ -43,38 +32,11 @@
 except:  # pylint: disable=bare-except
   pass
 
-@registry.register_hparams
-def pong_model_free():
-  hparams = tf.contrib.training.HParams(
-    epochs_num=4,
-    eval_every_epochs=2,
-    num_agents=10,
-    optimization_epochs=3,
-    epoch_length=30,
-    entropy_loss_coef=0.003,
-    learning_rate=8e-05,
-    optimizer="Adam",
-    policy_network=rl_models.feed_forward_cnn_small_categorical_fun,
-    gae_lambda=0.985,
-    num_eval_agents=1,
-    max_gradients_norm=0.5,
-    gae_gamma=0.985,
-    optimization_batch_size=4,
-    clipping_coef=0.2,
-    value_loss_coef=1,
-    save_models_every_epochs=False)
-  hparams.add_hparam(
-    "environment_spec",
-    gym_problems.standard_atari_env_spec("PongNoFrameskip-v4"))
-  hparams.add_hparam(
-    "environment_eval_spec",
-    gym_problems.standard_atari_env_eval_spec("PongNoFrameskip-v4"))
-
-  return hparams
 
 def main(_):
   hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
-  rl_trainer_lib.train(hparams, FLAGS.output_dir, FLAGS.output_dir)
+  rl_trainer_lib.train(hparams, FLAGS.problem, FLAGS.output_dir)
+
 
 if __name__ == "__main__":
   tf.app.run()

From 042715c1facfa53bdfb65fb009ab909fcd54abd8 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 4 Oct 2018 19:44:54 -0700
Subject: [PATCH 0964/2720] internal merge of PR #1111

PiperOrigin-RevId: 215845680
---
 tensor2tensor/data_generators/gym_problems.py |  14 +--
 .../data_generators/gym_problems_specs.py     | 102 ++++++++++++------
 tensor2tensor/models/research/autoencoders.py |  16 +++
 .../video/basic_deterministic_params.py       |  10 ++
 tensor2tensor/rl/envs/tf_atari_wrappers.py    |   7 +-
 tensor2tensor/rl/trainer_model_based.py       |  45 ++++----
 6 files changed, 134 insertions(+), 60 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 12973948f..671cedc7e 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -28,7 +28,6 @@
 
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
-from tensor2tensor.models.research import autoencoders
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import collect
 from tensor2tensor.rl.envs import tf_atari_wrappers
@@ -89,9 +88,10 @@ def standard_atari_env_eval_spec(env, simulated=False,
       include_clipping=False)
 
 
-def standard_atari_ae_env_spec(env):
+def standard_atari_ae_env_spec(env, ae_hparams_set):
   """Parameters of environment specification."""
-  standard_wrappers = [[tf_atari_wrappers.AutoencoderWrapper, {}],
+  standard_wrappers = [[tf_atari_wrappers.AutoencoderWrapper,
+                        {"ae_hparams_set": ae_hparams_set}],
                        [tf_atari_wrappers.StackWrapper, {"history": 4}]]
   env_lambda = None
   if isinstance(env, str):
@@ -476,7 +476,7 @@ def __init__(self, *args, **kwargs):
     self._forced_collect_level = 0
 
   def get_environment_spec(self):
-    return standard_atari_ae_env_spec(self.env_name)
+    return standard_atari_ae_env_spec(self.env_name, self.ae_hparams_set)
 
   def restore_networks(self, sess):
     super(GymDiscreteProblemWithAutoencoder, self).restore_networks(sess)
@@ -517,12 +517,12 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
                        " for reading encoded frames")
 
   def get_environment_spec(self):
-    return standard_atari_ae_env_spec(self.env_name)
+    return standard_atari_ae_env_spec(self.env_name, self.ae_hparams_set)
 
   @property
   def autoencoder_factor(self):
     """By how much to divide sizes when using autoencoders."""
-    hparams = autoencoders.autoencoder_discrete_pong()
+    hparams = registry.hparams(self.ae_hparams_set)
     return 2**hparams.num_hidden_layers
 
   @property
@@ -796,7 +796,7 @@ def get_environment_spec(self):
   @property
   def autoencoder_factor(self):
     """By how much to divide sizes when using autoencoders."""
-    hparams = autoencoders.autoencoder_discrete_pong()
+    hparams = registry.hparams(self.ae_hparams_set)
     return 2**hparams.num_hidden_layers
 
   @property
diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index 5612fd685..773b6793b 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -235,7 +235,8 @@ def create_problems_for_game(
     game_name,
     resize_height_factor=1,
     resize_width_factor=1,
-    game_mode="Deterministic-v4"):
+    game_mode="Deterministic-v4",
+    autoencoder_hparams=None):
   """Create and register problems for game_name.
 
   Args:
@@ -243,6 +244,7 @@ def create_problems_for_game(
     resize_height_factor: factor by which to resize the height of frames.
     resize_width_factor: factor by which to resize the width of frames.
     game_mode: the frame skip and sticky keys config.
+    autoencoder_hparams: the hparams for the autoencoder.
 
   Returns:
     dict of problems with keys ("base", "agent", "simulated").
@@ -267,36 +269,74 @@ def create_problems_for_game(
                       "resize_width_factor": resize_width_factor})
   registry.register_problem(problem_cls)
 
-  with_agent_cls = type("GymDiscreteProblemWithAgentOn%s" % camel_game_name,
-                        (GymRealDiscreteProblem, problem_cls), {})
-
-  registry.register_problem(with_agent_cls)
-
-  # Create and register the simulated Problem
-  simulated_cls = type(
-      "GymSimulatedDiscreteProblemWithAgentOn%s" % camel_game_name,
-      (GymSimulatedDiscreteProblem, problem_cls), {
-          "initial_frames_problem": with_agent_cls.name,
-          "num_testing_steps": 100
-      })
-  registry.register_problem(simulated_cls)
-
-  # Create and register the simulated Problem
-  world_model_eval_cls = type(
-      "GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%s" %
-      camel_game_name,
-      (GymSimulatedDiscreteProblemForWorldModelEval, problem_cls), {
-          "initial_frames_problem": with_agent_cls.name,
-          "num_testing_steps": 100
-      })
-  registry.register_problem(world_model_eval_cls)
-
-  return {
-      "base": problem_cls,
-      "agent": with_agent_cls,
-      "simulated": simulated_cls,
-      "world_model_eval": world_model_eval_cls,
-  }
+  if not autoencoder_hparams:
+    with_agent_cls = type("GymDiscreteProblemWithAgentOn%s" % camel_game_name,
+                          (GymRealDiscreteProblem, problem_cls), {})
+    registry.register_problem(with_agent_cls)
+
+    # Create and register the simulated Problem
+    simulated_cls = type(
+        "GymSimulatedDiscreteProblemWithAgentOn%s" % camel_game_name,
+        (GymSimulatedDiscreteProblem, problem_cls), {
+            "initial_frames_problem": with_agent_cls.name,
+            "num_testing_steps": 100
+        })
+    registry.register_problem(simulated_cls)
+
+    # Create and register the simulated Problem
+    world_model_eval_cls = type(
+        "GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%s" %
+        camel_game_name,
+        (GymSimulatedDiscreteProblemForWorldModelEval, problem_cls), {
+            "initial_frames_problem": with_agent_cls.name,
+            "num_testing_steps": 100
+        })
+    registry.register_problem(world_model_eval_cls)
+    return {
+        "base": problem_cls,
+        "agent": with_agent_cls,
+        "simulated": simulated_cls,
+        "world_model_eval": world_model_eval_cls,
+    }
+  else:
+    with_agent_cls_with_ae = \
+      type("GymDiscreteProblemWithAgentOn%sWithAutoencoder" % camel_game_name,
+           (GymDiscreteProblemWithAutoencoder, problem_cls),
+           {"ae_hparams_set": autoencoder_hparams})
+    registry.register_problem(with_agent_cls_with_ae)
+
+    with_agent_cls_ae = \
+      type("GymDiscreteProblemWithAgentOn%sAutoencoded" % camel_game_name,
+           (GymDiscreteProblemAutoencoded, problem_cls),
+           {"ae_hparams_set": autoencoder_hparams})
+    registry.register_problem(with_agent_cls_ae)
+
+    # Create and register the simulated Problem
+    simulated_cls = \
+      type("GymSimulatedDiscreteProblemWithAgentOn%sAutoencoded"
+           % camel_game_name,
+           (GymSimulatedDiscreteProblemAutoencoded, problem_cls), {
+               "initial_frames_problem": with_agent_cls_ae.name,
+               "num_testing_steps": 100,
+               "ae_hparams_set": autoencoder_hparams})
+    registry.register_problem(simulated_cls)
+
+    world_model_eval_cls = \
+      type("GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%sAutoencoded"  # pylint: disable=line-too-long
+           % camel_game_name,
+           (GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded,
+            problem_cls), {
+                "initial_frames_problem": with_agent_cls_ae.name,
+                "num_testing_steps": 100,
+                "ae_hparams_set": autoencoder_hparams})
+    registry.register_problem(world_model_eval_cls)
+    return {
+        "base": problem_cls,
+        "agent_with_ae": with_agent_cls_with_ae,
+        "agent_ae": with_agent_cls_ae,
+        "simulated_ae": simulated_cls,
+        "world_model_eval_ae": world_model_eval_cls,
+    }
 
 # Register the atari games with all of the possible modes.
 for game in ATARI_ALL_MODES_SHORT_LIST:
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 74979b4e4..eb125d7f0 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1225,6 +1225,22 @@ def autoencoder_discrete_pong():
   return hparams
 
 
+@registry.register_hparams
+def autoencoder_discrete_tiny():
+  """Discrete autoencoder model for compressing pong frames for testing."""
+  hparams = autoencoder_ordered_discrete()
+  hparams.num_hidden_layers = 2
+  hparams.bottleneck_bits = 24
+  hparams.batch_size = 2
+  hparams.gan_loss_factor = 0.
+  hparams.bottleneck_l2_factor = 0.001
+  hparams.add_hparam("video_modality_loss_cutoff", 0.02)
+  hparams.num_residual_layers = 1
+  hparams.hidden_size = 32
+  hparams.max_hidden_size = 64
+  return hparams
+
+
 @registry.register_hparams
 def autoencoder_discrete_cifar():
   """Discrete autoencoder model for compressing cifar."""
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index e87003834..c102f89a5 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -93,6 +93,16 @@ def next_frame_ae():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_ae_tiny():
+  """Conv autoencoder, tiny set for testing."""
+  hparams = next_frame_tiny()
+  hparams.input_modalities = "inputs:video:bitwise"
+  hparams.batch_size = 8
+  hparams.dropout = 0.4
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_small():
   """Small conv model."""
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 0486d49ea..663dd4052 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -221,13 +221,14 @@ def _transform_history_observations(self, frames):
 class AutoencoderWrapper(WrapperBase):
   """Transforms the observations taking the bottleneck of an autoencoder."""
 
-  def __init__(self, batch_env):
+  def __init__(self, batch_env, ae_hparams_set):
     super(AutoencoderWrapper, self).__init__(batch_env)
+    self.ae_hparams_set = ae_hparams_set
     self._observ = tf.Variable(
         tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
         trainable=False)
     with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      autoencoder_hparams = autoencoders.autoencoder_discrete_pong()
+      autoencoder_hparams = registry.hparams(self.ae_hparams_set)
       problem = registry.problem("dummy_autoencoder_problem")
       autoencoder_hparams.problem_hparams = problem.get_hparams(
           autoencoder_hparams)
@@ -249,7 +250,7 @@ def observ_shape(self):
   @property
   def autoencoder_factor(self):
     """By how much to divide sizes when using autoencoders."""
-    hparams = autoencoders.autoencoder_discrete_pong()
+    hparams = registry.hparams(self.ae_hparams_set)
     return 2**hparams.num_hidden_layers
 
   def simulate(self, action):
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index bb63b4cfa..85b5b750b 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -149,7 +149,7 @@ def train_autoencoder(problem_name, data_dir, output_dir, hparams, epoch):
       "data_dir": data_dir,
       "output_dir": output_dir,
       "model": "autoencoder_ordered_discrete",
-      "hparams_set": "autoencoder_discrete_pong",
+      "hparams_set": hparams.autoencoder_hparams_set,
       "train_steps": train_steps,
       "eval_steps": 100,
   }):
@@ -406,11 +406,11 @@ def generator():
         cycle_every_n=problem.total_number_of_frames // 10)
 
 
-def encode_env_frames(problem_name, ae_problem_name, autoencoder_path,
-                      epoch_data_dir):
+def encode_env_frames(problem_name, ae_problem_name, ae_hparams_set,
+                      autoencoder_path, epoch_data_dir):
   """Encode all frames from problem_name and write out as ae_problem_name."""
   with tf.Graph().as_default():
-    ae_hparams = trainer_lib.create_hparams("autoencoder_discrete_pong",
+    ae_hparams = trainer_lib.create_hparams(ae_hparams_set,
                                             problem_name=problem_name)
     problem = ae_hparams.problem
     model = registry.model("autoencoder_ordered_discrete")(
@@ -436,16 +436,17 @@ def encode_env_frames(problem_name, ae_problem_name, autoencoder_path,
       dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, epoch_data_dir,
                                 shuffle_files=False, output_buffer_size=100,
                                 preprocess=False)
-      encode_dataset(model, dataset, problem, ae_hparams, autoencoder_path,
-                     ae_training_paths)
+      encode_dataset(model, dataset, problem=problem, ae_hparams=ae_hparams,
+                     autoencoder_path=autoencoder_path,
+                     out_files=ae_training_paths)
 
     # Encode eval data
     if not skip_eval:
       dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, epoch_data_dir,
                                 shuffle_files=False, output_buffer_size=100,
                                 preprocess=False)
-      encode_dataset(model, dataset, problem, ae_hparams, autoencoder_path,
-                     ae_eval_paths)
+      encode_dataset(model, dataset, problem=problem, ae_hparams=ae_hparams,
+                     autoencoder_path=autoencoder_path, out_files=ae_eval_paths)
 
 
 def check_problems(problem_names):
@@ -459,8 +460,11 @@ def setup_problems(hparams, using_autoencoder=False):
     game_with_mode = hparams.game + "_deterministic-v4"
   else:
     game_with_mode = hparams.game
+  game_problems_kwargs = {}
   # Problems
   if using_autoencoder:
+    game_problems_kwargs["autoencoder_hparams"] = (
+        hparams.autoencoder_hparams_set)
     problem_name = (
         "gym_discrete_problem_with_agent_on_%s_with_autoencoder"
         % game_with_mode)
@@ -474,6 +478,8 @@ def setup_problems(hparams, using_autoencoder=False):
         "_autoencoded"
         % game_with_mode)
   else:
+    game_problems_kwargs["resize_height_factor"] = hparams.resize_height_factor
+    game_problems_kwargs["resize_width_factor"] = hparams.resize_width_factor
     problem_name = ("gym_discrete_problem_with_agent_on_%s" % game_with_mode)
     world_model_problem = problem_name
     simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
@@ -481,14 +487,11 @@ def setup_problems(hparams, using_autoencoder=False):
     world_model_eval_problem_name = (
         "gym_simulated_discrete_problem_for_world_model_eval_with_agent_on_%s"
         % game_with_mode)
-    if problem_name not in registry.list_problems():
-      tf.logging.info("Game Problem %s not found; dynamically registering",
-                      problem_name)
-      gym_problems_specs.create_problems_for_game(
-          hparams.game,
-          resize_height_factor=hparams.resize_height_factor,
-          resize_width_factor=hparams.resize_width_factor,
-          game_mode="Deterministic-v4")
+  if problem_name not in registry.list_problems():
+    tf.logging.info("Game Problem %s not found; dynamically registering",
+                    problem_name)
+    gym_problems_specs.create_problems_for_game(
+        hparams.game, game_mode="Deterministic-v4", **game_problems_kwargs)
   return (problem_name, world_model_problem, simulated_problem_name,
           world_model_eval_problem_name)
 
@@ -578,7 +581,9 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
       log("Autoencoding environment frames")
       encode_env_frames(problem_name, world_model_problem,
-                        autoencoder_model_dir, epoch_data_dir)
+                        ae_hparams_set=hparams.autoencoder_hparams_set,
+                        autoencoder_path=autoencoder_model_dir,
+                        epoch_data_dir=epoch_data_dir)
 
     # Train world model
     log("Training world model")
@@ -1147,6 +1152,7 @@ def rlmb_ae_base():
   hparams = rlmb_base()
   hparams.ppo_params = "ppo_pong_ae_base"
   hparams.generative_model_params = "next_frame_ae"
+  hparams.autoencoder_hparams_set = "autoencoder_discrete_pong"
   hparams.gather_ppo_real_env_data = False
   hparams.autoencoder_train_steps = 30000
   return hparams
@@ -1156,13 +1162,14 @@ def rlmb_ae_base():
 def rlmb_ae_tiny():
   """Tiny set for testing autoencoders."""
   hparams = rlmb_tiny()
-  hparams.game = "wrapped_full_pong"
   hparams.ppo_params = "ppo_pong_ae_base"
-  hparams.generative_model_params = "next_frame_ae"
+  hparams.generative_model_params = "next_frame_ae_tiny"
+  hparams.autoencoder_hparams_set = "autoencoder_discrete_tiny"
   hparams.gather_ppo_real_env_data = False
   hparams.resize_height_factor = 1
   hparams.resize_width_factor = 1
   hparams.autoencoder_train_steps = 2
+  hparams.stop_loop_early = False
   return hparams
 
 

From 1fe293fa61dcb3d6e2b3656bd70433fc6d2f1fab Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 4 Oct 2018 19:57:03 -0700
Subject: [PATCH 0965/2720] internal merge of PR #1104

PiperOrigin-RevId: 215846389
---
 tensor2tensor/rl/rl_trainer_lib_test.py | 31 +++--------------
 tensor2tensor/rl/trainer_model_free.py  | 46 +++++++++++++++++++++++--
 2 files changed, 49 insertions(+), 28 deletions(-)

diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index 519451b79..10e2cfbee 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -17,7 +17,6 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.data_generators import gym_problems
 from tensor2tensor.models.research import rl as rl_models
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.utils import trainer_lib
@@ -47,31 +46,11 @@ def test_no_crash_cartpole(self):
     rl_trainer_lib.train(hparams)
 
   def test_train_pong(self):
-    hparams = tf.contrib.training.HParams(
-        epochs_num=4,
-        eval_every_epochs=2,
-        num_agents=10,
-        optimization_epochs=3,
-        epoch_length=30,
-        entropy_loss_coef=0.003,
-        learning_rate=8e-05,
-        optimizer="Adam",
-        policy_network=rl_models.feed_forward_cnn_small_categorical_fun,
-        gae_lambda=0.985,
-        num_eval_agents=1,
-        max_gradients_norm=0.5,
-        gae_gamma=0.985,
-        optimization_batch_size=4,
-        clipping_coef=0.2,
-        value_loss_coef=1,
-        save_models_every_epochs=False)
-    hparams.add_hparam(
-        "environment_spec",
-        gym_problems.standard_atari_env_spec("PongNoFrameskip-v4"))
-    hparams.add_hparam(
-        "environment_eval_spec",
-        gym_problems.standard_atari_env_eval_spec("PongNoFrameskip-v4"))
-    rl_trainer_lib.train(hparams)
+    pass
+    # Commented to make travis work
+    # Uncomment for long tests
+    # hparams = registry.hparams("pong_model_free")
+    # rl_trainer_lib.train(hparams)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 715dc10ae..77f0204fe 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -12,12 +12,24 @@
 # 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.
-"""Training of RL agent with PPO algorithm."""
+r"""Training of RL agent with PPO algorithm.
+
+Example invocation:
+
+python -m tensor2tensor.rl.trainer_model_free \
+    --output_dir=$HOME/t2t/rl_v1 \
+    --hparams_set=pong_model_free \
+    --loop_hparams='num_agents=15'
+"""
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+from tensor2tensor.data_generators import gym_problems
+from tensor2tensor.models.research import rl as rl_models
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
+from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -33,9 +45,39 @@
   pass
 
 
+@registry.register_hparams
+def pong_model_free():
+  """TODO(piotrmilos): Document this."""
+  hparams = tf.contrib.training.HParams(
+      epochs_num=4,
+      eval_every_epochs=2,
+      num_agents=10,
+      optimization_epochs=3,
+      epoch_length=30,
+      entropy_loss_coef=0.003,
+      learning_rate=8e-05,
+      optimizer="Adam",
+      policy_network=rl_models.feed_forward_cnn_small_categorical_fun,
+      gae_lambda=0.985,
+      num_eval_agents=1,
+      max_gradients_norm=0.5,
+      gae_gamma=0.985,
+      optimization_batch_size=4,
+      clipping_coef=0.2,
+      value_loss_coef=1,
+      save_models_every_epochs=False)
+  hparams.add_hparam("environment_spec",
+                     gym_problems.standard_atari_env_spec("PongNoFrameskip-v4"))
+  hparams.add_hparam(
+      "environment_eval_spec",
+      gym_problems.standard_atari_env_eval_spec("PongNoFrameskip-v4"))
+
+  return hparams
+
+
 def main(_):
   hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
-  rl_trainer_lib.train(hparams, FLAGS.problem, FLAGS.output_dir)
+  rl_trainer_lib.train(hparams, FLAGS.output_dir, FLAGS.output_dir)
 
 
 if __name__ == "__main__":

From 17f9a9a0bf4ed23f4ad2fec8a5ab974f014554dd Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 4 Oct 2018 21:34:31 -0700
Subject: [PATCH 0966/2720] Add the option to use grayscale and correct
 registration for autoencoded problems; also bring tests back.

PiperOrigin-RevId: 215853764
---
 tensor2tensor/data_generators/gym_problems.py |  83 +++-------
 .../data_generators/gym_problems_specs.py     | 146 ++++++++----------
 tensor2tensor/models/research/rl.py           |  91 +++++++++++
 tensor2tensor/rl/envs/tf_atari_wrappers.py    |  19 ++-
 tensor2tensor/rl/rl_trainer_lib_test.py       |  11 +-
 tensor2tensor/rl/trainer_model_based.py       |  69 ++++-----
 .../rl/trainer_model_based_ae_test.py         |  11 +-
 tensor2tensor/rl/trainer_model_free.py        |  33 ----
 8 files changed, 230 insertions(+), 233 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 671cedc7e..afb18f402 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -47,64 +47,6 @@
                     "File with model for autoencoder.")
 
 
-def standard_atari_env_spec(
-    env, simulated=False, resize_height_factor=1, resize_width_factor=1,
-    include_clipping=True):
-  """Parameters of environment specification."""
-  resize_wrapper = [tf_atari_wrappers.ResizeWrapper,
-                    {"height_factor": resize_height_factor,
-                     "width_factor": resize_width_factor}]
-  if include_clipping:
-    standard_wrappers = [
-        resize_wrapper,
-        [tf_atari_wrappers.RewardClippingWrapper, {}],
-        [tf_atari_wrappers.StackWrapper, {"history": 4}],
-    ]
-  else:
-    standard_wrappers = [
-        resize_wrapper,
-        [tf_atari_wrappers.StackWrapper, {"history": 4}],
-    ]
-  if simulated:  # No resizing on simulated environments.
-    standard_wrappers = standard_wrappers[1:]
-  env_lambda = None
-  if isinstance(env, str):
-    env_lambda = lambda: gym.make(env)
-  if callable(env):
-    env_lambda = env
-  assert env_lambda is not None, "Unknown specification of environment"
-
-  return tf.contrib.training.HParams(
-      env_lambda=env_lambda,
-      wrappers=standard_wrappers,
-      simulated_env=simulated)
-
-
-def standard_atari_env_eval_spec(env, simulated=False,
-                                 resize_height_factor=1, resize_width_factor=1):
-  """Parameters of environment specification for eval."""
-  return standard_atari_env_spec(
-      env, simulated, resize_height_factor, resize_width_factor,
-      include_clipping=False)
-
-
-def standard_atari_ae_env_spec(env, ae_hparams_set):
-  """Parameters of environment specification."""
-  standard_wrappers = [[tf_atari_wrappers.AutoencoderWrapper,
-                        {"ae_hparams_set": ae_hparams_set}],
-                       [tf_atari_wrappers.StackWrapper, {"history": 4}]]
-  env_lambda = None
-  if isinstance(env, str):
-    env_lambda = lambda: gym.make(env)
-  if callable(env):
-    env_lambda = env
-  assert env is not None, "Unknown specification of environment"
-
-  return tf.contrib.training.HParams(env_lambda=env_lambda,
-                                     wrappers=standard_wrappers,
-                                     simulated_env=False)
-
-
 frame_dumper_use_disk = False  # Whether to use memory or disk to dump frames.
 frame_dumper = {}
 
@@ -139,6 +81,15 @@ def resize_height_factor(self):
   def resize_width_factor(self):
     return 1
 
+  @property
+  def grayscale(self):
+    return False
+
+  @property
+  def num_channels(self):
+    """Number of color channels in each frame."""
+    return 1 if self.grayscale else 3
+
   def _setup(self, data_dir, extra_collect_hparams=None,
              override_collect_hparams=None):
     dumper_path = os.path.join(data_dir, "dumper")
@@ -304,10 +255,11 @@ def extra_reading_spec(self):
     return data_fields, decoders
 
   def get_environment_spec(self):
-    return standard_atari_env_spec(
+    return rl.standard_atari_env_spec(
         self.env_name,
         resize_height_factor=self.resize_height_factor,
-        resize_width_factor=self.resize_width_factor)
+        resize_width_factor=self.resize_width_factor,
+        grayscale=self.grayscale)
 
   @property
   def environment_spec(self):
@@ -476,7 +428,7 @@ def __init__(self, *args, **kwargs):
     self._forced_collect_level = 0
 
   def get_environment_spec(self):
-    return standard_atari_ae_env_spec(self.env_name, self.ae_hparams_set)
+    return rl.standard_atari_ae_env_spec(self.env_name, self.ae_hparams_set)
 
   def restore_networks(self, sess):
     super(GymDiscreteProblemWithAutoencoder, self).restore_networks(sess)
@@ -517,7 +469,7 @@ def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
                        " for reading encoded frames")
 
   def get_environment_spec(self):
-    return standard_atari_ae_env_spec(self.env_name, self.ae_hparams_set)
+    return rl.standard_atari_ae_env_spec(self.env_name, self.ae_hparams_set)
 
   @property
   def autoencoder_factor(self):
@@ -625,11 +577,12 @@ def num_testing_steps(self):
     return None
 
   def get_environment_spec(self):
-    env_spec = standard_atari_env_spec(
+    env_spec = rl.standard_atari_env_spec(
         self.env_name,
         simulated=True,
         resize_height_factor=self.resize_height_factor,
-        resize_width_factor=self.resize_width_factor)
+        resize_width_factor=self.resize_width_factor,
+        grayscale=self.grayscale)
     env_spec.add_hparam("simulation_random_starts", True)
     env_spec.add_hparam("simulation_flip_first_random_for_beginning", True)
     env_spec.add_hparam("intrinsic_reward_scale", 0.0)
@@ -777,7 +730,7 @@ def __init__(self, *args, **kwargs):
     self._forced_collect_level = 0
 
   def get_environment_spec(self):
-    env_spec = standard_atari_env_spec(self.env_name)
+    env_spec = rl.standard_atari_env_spec(self.env_name)
     env_spec.wrappers = [
         [tf_atari_wrappers.IntToBitWrapper, {}],
         [tf_atari_wrappers.StackWrapper, {"history": 4}]
diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index 773b6793b..007fdcb14 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -66,8 +66,6 @@
     "up_n_down", "video_pinball", "yars_revenge",
 ]
 
-ATARI_ALL_MODES_SHORT_LIST = []
-
 ATARI_WHITELIST_GAMES = [
     "amidar",
     "bank_heist",
@@ -104,6 +102,12 @@
 ]
 
 
+# Games for which we hard-define problems to run all around.
+# TODO(lukaszkaiser): global registration makes them all rescaled and grayscale,
+# no matter the setting of hparams later (as they're registered at start).
+ATARI_ALL_MODES_SHORT_LIST = []  # ATARI_DEBUG_GAMES + ATARI_CURIOUS_GAMES
+
+
 # Different ATARI game modes in OpenAI Gym. Full list here:
 # https://github.com/openai/gym/blob/master/gym/envs/__init__.py
 ATARI_GAME_MODES = [
@@ -233,8 +237,9 @@ def num_rewards(self):
 
 def create_problems_for_game(
     game_name,
-    resize_height_factor=1,
-    resize_width_factor=1,
+    resize_height_factor=2,
+    resize_width_factor=2,
+    grayscale=True,
     game_mode="Deterministic-v4",
     autoencoder_hparams=None):
   """Create and register problems for game_name.
@@ -243,6 +248,7 @@ def create_problems_for_game(
     game_name: str, one of the games in ATARI_GAMES, e.g. "bank_heist".
     resize_height_factor: factor by which to resize the height of frames.
     resize_width_factor: factor by which to resize the width of frames.
+    grayscale: whether to make frames grayscale.
     game_mode: the frame skip and sticky keys config.
     autoencoder_hparams: the hparams for the autoencoder.
 
@@ -266,83 +272,67 @@ def create_problems_for_game(
                      (GymClippedRewardRandom,),
                      {"env_name": env_name,
                       "resize_height_factor": resize_height_factor,
-                      "resize_width_factor": resize_width_factor})
+                      "resize_width_factor": resize_width_factor,
+                      "grayscale": grayscale})
   registry.register_problem(problem_cls)
 
-  if not autoencoder_hparams:
-    with_agent_cls = type("GymDiscreteProblemWithAgentOn%s" % camel_game_name,
-                          (GymRealDiscreteProblem, problem_cls), {})
-    registry.register_problem(with_agent_cls)
-
-    # Create and register the simulated Problem
-    simulated_cls = type(
-        "GymSimulatedDiscreteProblemWithAgentOn%s" % camel_game_name,
-        (GymSimulatedDiscreteProblem, problem_cls), {
-            "initial_frames_problem": with_agent_cls.name,
-            "num_testing_steps": 100
-        })
-    registry.register_problem(simulated_cls)
-
-    # Create and register the simulated Problem
-    world_model_eval_cls = type(
-        "GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%s" %
-        camel_game_name,
-        (GymSimulatedDiscreteProblemForWorldModelEval, problem_cls), {
-            "initial_frames_problem": with_agent_cls.name,
-            "num_testing_steps": 100
-        })
-    registry.register_problem(world_model_eval_cls)
-    return {
-        "base": problem_cls,
-        "agent": with_agent_cls,
-        "simulated": simulated_cls,
-        "world_model_eval": world_model_eval_cls,
-    }
-  else:
-    with_agent_cls_with_ae = \
-      type("GymDiscreteProblemWithAgentOn%sWithAutoencoder" % camel_game_name,
-           (GymDiscreteProblemWithAutoencoder, problem_cls),
-           {"ae_hparams_set": autoencoder_hparams})
-    registry.register_problem(with_agent_cls_with_ae)
-
-    with_agent_cls_ae = \
-      type("GymDiscreteProblemWithAgentOn%sAutoencoded" % camel_game_name,
-           (GymDiscreteProblemAutoencoded, problem_cls),
-           {"ae_hparams_set": autoencoder_hparams})
-    registry.register_problem(with_agent_cls_ae)
-
-    # Create and register the simulated Problem
-    simulated_cls = \
-      type("GymSimulatedDiscreteProblemWithAgentOn%sAutoencoded"
-           % camel_game_name,
-           (GymSimulatedDiscreteProblemAutoencoded, problem_cls), {
-               "initial_frames_problem": with_agent_cls_ae.name,
-               "num_testing_steps": 100,
-               "ae_hparams_set": autoencoder_hparams})
-    registry.register_problem(simulated_cls)
-
-    world_model_eval_cls = \
-      type("GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%sAutoencoded"  # pylint: disable=line-too-long
-           % camel_game_name,
-           (GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded,
-            problem_cls), {
-                "initial_frames_problem": with_agent_cls_ae.name,
-                "num_testing_steps": 100,
-                "ae_hparams_set": autoencoder_hparams})
-    registry.register_problem(world_model_eval_cls)
-    return {
-        "base": problem_cls,
-        "agent_with_ae": with_agent_cls_with_ae,
-        "agent_ae": with_agent_cls_ae,
-        "simulated_ae": simulated_cls,
-        "world_model_eval_ae": world_model_eval_cls,
-    }
+  with_agent_cls = type("GymDiscreteProblemWithAgentOn%s" % camel_game_name,
+                        (GymRealDiscreteProblem, problem_cls), {})
+  registry.register_problem(with_agent_cls)
+
+  with_ae_cls = type(
+      "GymDiscreteProblemWithAgentOn%sWithAutoencoder" % camel_game_name,
+      (GymDiscreteProblemWithAutoencoder, problem_cls),
+      {"ae_hparams_set": autoencoder_hparams})
+  registry.register_problem(with_ae_cls)
+
+  ae_cls = type(
+      "GymDiscreteProblemWithAgentOn%sAutoencoded" % camel_game_name,
+      (GymDiscreteProblemAutoencoded, problem_cls),
+      {"ae_hparams_set": autoencoder_hparams})
+  registry.register_problem(ae_cls)
+
+  # Create and register the simulated Problem
+  simulated_cls = type(
+      "GymSimulatedDiscreteProblemWithAgentOn%s" % camel_game_name,
+      (GymSimulatedDiscreteProblem, problem_cls), {
+          "initial_frames_problem": with_agent_cls.name,
+          "num_testing_steps": 100
+      })
+  registry.register_problem(simulated_cls)
+
+  simulated_ae_cls = type(
+      "GymSimulatedDiscreteProblemWithAgentOn%sAutoencoded" % camel_game_name,
+      (GymSimulatedDiscreteProblemAutoencoded, problem_cls), {
+          "initial_frames_problem": ae_cls.name,
+          "num_testing_steps": 100,
+          "ae_hparams_set": autoencoder_hparams
+      })
+  registry.register_problem(simulated_ae_cls)
+
+  # Create and register the simulated Problem
+  world_model_eval_cls = type(
+      "GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%s" %
+      camel_game_name,
+      (GymSimulatedDiscreteProblemForWorldModelEval, problem_cls), {
+          "initial_frames_problem": with_agent_cls.name,
+          "num_testing_steps": 100,
+          "ae_hparams_set": autoencoder_hparams
+      })
+  registry.register_problem(world_model_eval_cls)
+
+  world_model_eval_ae_cls = type(
+      "GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%sAutoencoded" %
+      camel_game_name,
+      (GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded, problem_cls), {
+          "initial_frames_problem": ae_cls.name,
+          "num_testing_steps": 100,
+          "ae_hparams_set": autoencoder_hparams
+      })
+  registry.register_problem(world_model_eval_ae_cls)
+
 
 # Register the atari games with all of the possible modes.
 for game in ATARI_ALL_MODES_SHORT_LIST:
-  ATARI_PROBLEMS[game] = {}
   for mode in ATARI_GAME_MODES:
-    classes = create_problems_for_game(
-        game,
-        game_mode=mode)
-    ATARI_PROBLEMS[game][mode] = classes
+    create_problems_for_game(game, game_mode=mode)
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 52031d2dd..bd6b75327 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -18,9 +18,11 @@
 import functools
 import operator
 import gym
+
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
+from tensor2tensor.rl.envs import tf_atari_wrappers
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -140,6 +142,66 @@ def simple_gym_spec(env):
                                      simulated_env=False)
 
 
+def standard_atari_env_spec(
+    env, simulated=False, resize_height_factor=1, resize_width_factor=1,
+    grayscale=False, include_clipping=True):
+  """Parameters of environment specification."""
+  resize_wrapper = [tf_atari_wrappers.ResizeWrapper,
+                    {"height_factor": resize_height_factor,
+                     "width_factor": resize_width_factor,
+                     "grayscale": grayscale}]
+  if include_clipping:
+    standard_wrappers = [
+        resize_wrapper,
+        [tf_atari_wrappers.RewardClippingWrapper, {}],
+        [tf_atari_wrappers.StackWrapper, {"history": 4}],
+    ]
+  else:
+    standard_wrappers = [
+        resize_wrapper,
+        [tf_atari_wrappers.StackWrapper, {"history": 4}],
+    ]
+  if simulated:  # No resizing on simulated environments.
+    standard_wrappers = standard_wrappers[1:]
+  env_lambda = None
+  if isinstance(env, str):
+    env_lambda = lambda: gym.make(env)
+  if callable(env):
+    env_lambda = env
+  assert env_lambda is not None, "Unknown specification of environment"
+
+  return tf.contrib.training.HParams(
+      env_lambda=env_lambda,
+      wrappers=standard_wrappers,
+      simulated_env=simulated)
+
+
+def standard_atari_env_eval_spec(
+    env, simulated=False, resize_height_factor=1, resize_width_factor=1,
+    grayscale=False):
+  """Parameters of environment specification for eval."""
+  return standard_atari_env_spec(
+      env, simulated, resize_height_factor, resize_width_factor, grayscale,
+      include_clipping=False)
+
+
+def standard_atari_ae_env_spec(env, ae_hparams_set):
+  """Parameters of environment specification."""
+  standard_wrappers = [[tf_atari_wrappers.AutoencoderWrapper,
+                        {"ae_hparams_set": ae_hparams_set}],
+                       [tf_atari_wrappers.StackWrapper, {"history": 4}]]
+  env_lambda = None
+  if isinstance(env, str):
+    env_lambda = lambda: gym.make(env)
+  if callable(env):
+    env_lambda = env
+  assert env is not None, "Unknown specification of environment"
+
+  return tf.contrib.training.HParams(env_lambda=env_lambda,
+                                     wrappers=standard_wrappers,
+                                     simulated_env=False)
+
+
 @registry.register_hparams
 def ppo_pong_ae_base():
   """Pong autoencoder base parameters."""
@@ -149,6 +211,35 @@ def ppo_pong_ae_base():
   return hparams
 
 
+@registry.register_hparams
+def pong_model_free():
+  """TODO(piotrmilos): Document this."""
+  hparams = tf.contrib.training.HParams(
+      epochs_num=4,
+      eval_every_epochs=2,
+      num_agents=10,
+      optimization_epochs=3,
+      epoch_length=30,
+      entropy_loss_coef=0.003,
+      learning_rate=8e-05,
+      optimizer="Adam",
+      policy_network=feed_forward_cnn_small_categorical_fun,
+      gae_lambda=0.985,
+      num_eval_agents=1,
+      max_gradients_norm=0.5,
+      gae_gamma=0.985,
+      optimization_batch_size=4,
+      clipping_coef=0.2,
+      value_loss_coef=1,
+      save_models_every_epochs=False)
+  hparams.add_hparam("environment_spec",
+                     standard_atari_env_spec("PongNoFrameskip-v4"))
+  hparams.add_hparam(
+      "environment_eval_spec",
+      standard_atari_env_eval_spec("PongNoFrameskip-v4"))
+  return hparams
+
+
 NetworkOutput = collections.namedtuple(
     "NetworkOutput", "policy, value, action_postprocessing")
 
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 663dd4052..3c0bde897 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -286,11 +286,14 @@ def _transform_history_observations(self, frames):
 class ResizeWrapper(WrapperBase):
   """Resizes the observations."""
 
-  def __init__(self, batch_env, height_factor=1, width_factor=1):
+  def __init__(self, batch_env, height_factor=1, width_factor=1,
+               grayscale=False):
     super(ResizeWrapper, self).__init__(batch_env)
     self._height_factor = height_factor  # How much to resize on x axis.
     self._width_factor = width_factor  # How much to resize on y axis.
-    self._is_identity = (height_factor == 1) and (width_factor == 1)
+    self._do_grayscale = grayscale  # Whether to convert to grayscale.
+    self._is_identity = ((height_factor == 1) and (width_factor == 1)
+                         and not grayscale)
     if not self._is_identity:
       self._observ = tf.Variable(
           tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
@@ -299,8 +302,9 @@ def __init__(self, batch_env, height_factor=1, width_factor=1):
       self._observ = self._batch_env.observ
 
   def __str__(self):
-    return "ResizeWrapper%d%d(%s)" % (self._height_factor,
-                                      self._width_factor, str(self._batch_env))
+    return "ResizeWrapperh%dw%dg%d(%s)" % (
+        self._height_factor, self._width_factor, int(self._do_grayscale),
+        str(self._batch_env))
 
   def _resize(self, tensor):
     if self._is_identity:
@@ -309,11 +313,15 @@ def _resize(self, tensor):
     observ = tf.to_float(tensor)
     resized = tf.image.resize_images(
         observ, [height, width], tf.image.ResizeMethod.AREA)
+    if self._do_grayscale:
+      resized = tf.image.rgb_to_grayscale(resized)
     return tf.cast(resized, self.observ_dtype)
 
   @property
   def observ_shape(self):
     height, width, channels = self._batch_env.observ_shape
+    if self._do_grayscale:
+      channels = 1
     resized_height = height // self._height_factor
     resized_width = width // self._width_factor
     return (resized_height, resized_width, channels)
@@ -363,9 +371,6 @@ def observ_shape(self):
     return (height, width, channels*8)
 
   def simulate(self, action):
-    action = tf.Print(action, [action], message="action=", summarize=200)
-
-    # action = tf.zeros_like(action) #Temporary hacked bugfix
     reward, done = self._batch_env.simulate(action)
     with tf.control_dependencies([reward, done]):
       with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index 10e2cfbee..a8a02cac9 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -19,6 +19,7 @@
 
 from tensor2tensor.models.research import rl as rl_models
 from tensor2tensor.rl import rl_trainer_lib
+from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -46,11 +47,11 @@ def test_no_crash_cartpole(self):
     rl_trainer_lib.train(hparams)
 
   def test_train_pong(self):
-    pass
-    # Commented to make travis work
-    # Uncomment for long tests
-    # hparams = registry.hparams("pong_model_free")
-    # rl_trainer_lib.train(hparams)
+    hparams = registry.hparams("pong_model_free")
+    hparams.epochs_num = 2
+    hparams.num_agents = 2
+    hparams.epoch_length = 3
+    rl_trainer_lib.train(hparams)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 85b5b750b..57de06627 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -143,7 +143,8 @@ def log(msg, *args):
 
 def train_autoencoder(problem_name, data_dir, output_dir, hparams, epoch):
   """Train autoencoder on problem_name."""
-  train_steps = hparams.autoencoder_train_steps * (epoch + 2)
+  additional_steps = 1 + hparams.autoencoder_train_steps_initial_multiplier
+  train_steps = hparams.autoencoder_train_steps * (epoch + additional_steps)
   with temporary_flags({
       "problem": problem_name,
       "data_dir": data_dir,
@@ -478,8 +479,6 @@ def setup_problems(hparams, using_autoencoder=False):
         "_autoencoded"
         % game_with_mode)
   else:
-    game_problems_kwargs["resize_height_factor"] = hparams.resize_height_factor
-    game_problems_kwargs["resize_width_factor"] = hparams.resize_width_factor
     problem_name = ("gym_discrete_problem_with_agent_on_%s" % game_with_mode)
     world_model_problem = problem_name
     simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
@@ -488,6 +487,9 @@ def setup_problems(hparams, using_autoencoder=False):
         "gym_simulated_discrete_problem_for_world_model_eval_with_agent_on_%s"
         % game_with_mode)
   if problem_name not in registry.list_problems():
+    game_problems_kwargs["resize_height_factor"] = hparams.resize_height_factor
+    game_problems_kwargs["resize_width_factor"] = hparams.resize_width_factor
+    game_problems_kwargs["grayscale"] = hparams.grayscale
     tf.logging.info("Game Problem %s not found; dynamically registering",
                     problem_name)
     gym_problems_specs.create_problems_for_game(
@@ -758,6 +760,7 @@ def rlmb_base():
       generative_model_params="next_frame_pixel_noise",
       ppo_params="ppo_pong_base",
       autoencoder_train_steps=0,
+      autoencoder_train_steps_initial_multiplier=10,
       model_train_steps=15000,
       inital_epoch_train_steps_multiplier=3,
       simulated_env_generator_num_steps=2000,
@@ -783,6 +786,7 @@ def rlmb_base():
       # Resizing.
       resize_height_factor=2,
       resize_width_factor=2,
+      grayscale=False,
       # Bump learning rate after first epoch by 3x.
       # We picked 3x because our default learning rate schedule decreases with
       # 1/square root of step; 1/sqrt(10k) = 0.01 and 1/sqrt(100k) ~ 0.0032
@@ -818,7 +822,7 @@ def rlmb_basetest():
   hparams.game = "pong"
   hparams.epochs = 2
   hparams.num_real_env_frames = 3200
-  hparams.model_train_steps = 500
+  hparams.model_train_steps = 100
   hparams.simulated_env_generator_num_steps = 20
   hparams.ppo_epochs_num = 2
   return hparams
@@ -1111,38 +1115,7 @@ def rlmb_tiny_sv2p():
   hparams = rlmb_tiny()
   hparams.generative_model = "next_frame_sv2p"
   hparams.generative_model_params = "next_frame_sv2p_tiny"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_l1_base():
-  """Parameter set with L1 loss."""
-  hparams = rlmb_base()
-  hparams.generative_model_params = "next_frame_l1"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_l1_tiny():
-  """Tiny parameter set with L1 loss."""
-  hparams = rlmb_tiny()
-  hparams.generative_model_params = "next_frame_l1"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_l2_base():
-  """Parameter set with L2 loss."""
-  hparams = rlmb_base()
-  hparams.generative_model_params = "next_frame_l2"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_l2_tiny():
-  """Tiny parameter set with L2 loss."""
-  hparams = rlmb_tiny()
-  hparams.generative_model_params = "next_frame_l2"
+  hparams.grayscale = False
   return hparams
 
 
@@ -1154,7 +1127,24 @@ def rlmb_ae_base():
   hparams.generative_model_params = "next_frame_ae"
   hparams.autoencoder_hparams_set = "autoencoder_discrete_pong"
   hparams.gather_ppo_real_env_data = False
-  hparams.autoencoder_train_steps = 30000
+  hparams.autoencoder_train_steps = 5000
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  hparams.grayscale = False
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_ae_basetest():
+  """Base AE setting but quicker with only 2 epochs."""
+  hparams = rlmb_ae_base()
+  hparams.game = "pong"
+  hparams.epochs = 2
+  hparams.num_real_env_frames = 3200
+  hparams.model_train_steps = 100
+  hparams.autoencoder_train_steps = 10
+  hparams.simulated_env_generator_num_steps = 20
+  hparams.ppo_epochs_num = 2
   return hparams
 
 
@@ -1168,8 +1158,9 @@ def rlmb_ae_tiny():
   hparams.gather_ppo_real_env_data = False
   hparams.resize_height_factor = 1
   hparams.resize_width_factor = 1
-  hparams.autoencoder_train_steps = 2
-  hparams.stop_loop_early = False
+  hparams.grayscale = False
+  hparams.autoencoder_train_steps = 1
+  hparams.autoencoder_train_steps_initial_multiplier = 0
   return hparams
 
 
diff --git a/tensor2tensor/rl/trainer_model_based_ae_test.py b/tensor2tensor/rl/trainer_model_based_ae_test.py
index 64164939b..a7b15a62d 100644
--- a/tensor2tensor/rl/trainer_model_based_ae_test.py
+++ b/tensor2tensor/rl/trainer_model_based_ae_test.py
@@ -17,7 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
-# from tensor2tensor.rl import trainer_model_based
+from tensor2tensor.rl import trainer_model_based
 
 import tensorflow as tf
 
@@ -27,11 +27,10 @@
 class ModelRLExperimentTestAe(tf.test.TestCase):
 
   def test_ae(self):
-    # FLAGS.output_dir = tf.test.get_temp_dir()
-    # FLAGS.loop_hparams_set = "rlmb_ae_tiny"
-    # FLAGS.schedule = "train"  # skip evaluation for world model training
-    # trainer_model_based.main(None)
-    pass
+    FLAGS.output_dir = tf.test.get_temp_dir()
+    FLAGS.loop_hparams_set = "rlmb_ae_tiny"
+    FLAGS.schedule = "train"  # skip evaluation for world model training
+    trainer_model_based.main(None)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 77f0204fe..2471b9d59 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -25,11 +25,8 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-from tensor2tensor.data_generators import gym_problems
-from tensor2tensor.models.research import rl as rl_models
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
-from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -45,36 +42,6 @@
   pass
 
 
-@registry.register_hparams
-def pong_model_free():
-  """TODO(piotrmilos): Document this."""
-  hparams = tf.contrib.training.HParams(
-      epochs_num=4,
-      eval_every_epochs=2,
-      num_agents=10,
-      optimization_epochs=3,
-      epoch_length=30,
-      entropy_loss_coef=0.003,
-      learning_rate=8e-05,
-      optimizer="Adam",
-      policy_network=rl_models.feed_forward_cnn_small_categorical_fun,
-      gae_lambda=0.985,
-      num_eval_agents=1,
-      max_gradients_norm=0.5,
-      gae_gamma=0.985,
-      optimization_batch_size=4,
-      clipping_coef=0.2,
-      value_loss_coef=1,
-      save_models_every_epochs=False)
-  hparams.add_hparam("environment_spec",
-                     gym_problems.standard_atari_env_spec("PongNoFrameskip-v4"))
-  hparams.add_hparam(
-      "environment_eval_spec",
-      gym_problems.standard_atari_env_eval_spec("PongNoFrameskip-v4"))
-
-  return hparams
-
-
 def main(_):
   hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
   rl_trainer_lib.train(hparams, FLAGS.output_dir, FLAGS.output_dir)

From 59a86ee53b196976f434bcdb37864819b88db068 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 5 Oct 2018 08:03:22 -0700
Subject: [PATCH 0967/2720] Enable fast beam decoding from dataset on TPU.

PiperOrigin-RevId: 215906853
---
 tensor2tensor/data_generators/problem.py |  17 +-
 tensor2tensor/layers/common_hparams.py   |   2 +
 tensor2tensor/models/transformer.py      | 193 +++++++++++++++--------
 tensor2tensor/utils/beam_search.py       |  39 +++--
 tensor2tensor/utils/beam_search_test.py  |  50 ++++++
 tensor2tensor/utils/decoding.py          |  48 ++++--
 tensor2tensor/utils/t2t_model.py         |  25 ++-
 tensor2tensor/utils/trainer_lib.py       |   5 +
 8 files changed, 277 insertions(+), 102 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 2895539a8..da2c714c7 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -901,9 +901,20 @@ def define_shapes(example):
         # on TPU, we use params["batch_size"], which specifies the number of
         # examples across all datashards
         batch_size = params["batch_size"]
-        dataset = dataset.apply(
-            tf.contrib.data.padded_batch_and_drop_remainder(
-                batch_size, padded_shapes))
+        if hparams.pad_batch:
+          tf.logging.warn(
+              "Padding the batch to ensure that remainder eval batches are "
+              "processed. This may lead to incorrect metrics for "
+              "non-zero-padded features, e.g. images. Use a smaller batch "
+              "size that has no remainder in that case.")
+          dataset = dataset.padded_batch(
+              batch_size, padded_shapes, drop_remainder=False)
+          dataset = dataset.map(
+              functools.partial(pad_batch, batch_multiple=batch_size),
+              num_parallel_calls=num_threads)
+        else:
+          dataset = dataset.padded_batch(
+              batch_size, padded_shapes, drop_remainder=True)
       else:
         # On GPU, bucket by length
         dataset = dataset.filter(gpu_valid_size)
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index cdb872adf..02edfd231 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -286,6 +286,8 @@ def basic_params1():
       # There could be a performance drop if host_call function is slow and
       # cannot keep up with the TPU-side computation.
       tpu_enable_host_call=False,
+      # Pad batch dim of inputs to nearest multiple of batch multiple.
+      pad_batch=False,
   )
 
 
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 708bb7d86..4e4cc7515 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -226,7 +226,13 @@ def _greedy_infer(self, features, decode_length, use_tpu=False):
       return (self._fast_decode_tpu(features, decode_length) if use_tpu else
               self._fast_decode(features, decode_length))
 
-  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
+  def _beam_decode(self,
+                   features,
+                   decode_length,
+                   beam_size,
+                   top_beams,
+                   alpha,
+                   use_tpu=False):
     """Beam search decoding.
 
     Args:
@@ -236,6 +242,7 @@ def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
       top_beams: an integer. How many of the beams to return.
       alpha: Float that controls the length penalty. larger the alpha, stronger
         the preference for longer translations.
+      use_tpu: A bool, whether to do beam decode on TPU.
 
     Returns:
       A dict of decoding results {
@@ -252,34 +259,44 @@ def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
       # TODO(petershaw): Support fast decoding when using relative
       # position representations, i.e. "dot_product_relative" attention.
       return self._beam_decode_slow(features, decode_length, beam_size,
-                                    top_beams, alpha)
+                                    top_beams, alpha, use_tpu)
     with tf.variable_scope(self.name):
-      return self._fast_decode(features, decode_length, beam_size, top_beams,
-                               alpha)
+      return (
+          self._fast_decode_tpu(
+              features, decode_length, beam_size, top_beams, alpha) if use_tpu
+          else self._fast_decode(
+              features, decode_length, beam_size, top_beams, alpha))
 
   def _fast_decode_tpu(self,
                        features,
                        decode_length,
-                       beam_size=1):
+                       beam_size=1,
+                       top_beams=1,
+                       alpha=1.0):
     """Fast decoding.
 
-    Implements only greedy decoding on TPU.
+    Implements both greedy and beam search decoding on TPU, uses beam search
+    iff beam_size > 1, otherwise beam search related arguments are ignored.
 
     Args:
       features: A map of string to model features.
       decode_length: An integer, how many additional timesteps to decode.
       beam_size: An integer, number of beams.
+      top_beams: An integer, how many of the beams to return.
+      alpha: A float that controls the length penalty. Larger the alpha,
+        stronger the preference for longer translations.
 
     Returns:
       A dict of decoding results {
           "outputs": integer `Tensor` of decoded ids of shape
-              [batch_size, <= decode_length]
+              [batch_size, <= decode_length] if beam_size == 1 or
+              [batch_size, top_beams, <= decode_length]
           "scores": decoding log probs from the beam search,
               None if using greedy decoding (beam_size=1)
       }.
 
     Raises:
-      NotImplementedError: If there are multiple data shards or beam_size > 1.
+      NotImplementedError: If there are multiple data shards.
     """
     if self._num_datashards != 1:
       raise NotImplementedError("Fast decoding only supports a single shard.")
@@ -455,11 +472,17 @@ def forced_logits():
         symbols_to_logits_fn=symbols_to_logits_tpu_fn,
         hparams=hparams,
         decode_length=decode_length,
+        vocab_size=target_modality.top_dimensionality,
         beam_size=beam_size,
+        top_beams=top_beams,
+        alpha=alpha,
         batch_size=batch_size,
         force_decode_length=self._decode_hparams.force_decode_length)
     if partial_targets is not None:
-      ret["outputs"] = ret["outputs"][:, partial_targets_length:]
+      if beam_size <= 1 or top_beams <= 1:
+        ret["outputs"] = ret["outputs"][:, partial_targets_length:]
+      else:
+        ret["outputs"] = ret["outputs"][:, :, partial_targets_length:]
     return ret
 
   def _fast_decode(self,
@@ -664,7 +687,10 @@ def fast_decode_tpu(encoder_output,
                     symbols_to_logits_fn,
                     hparams,
                     decode_length,
+                    vocab_size,
                     beam_size=1,
+                    top_beams=1,
+                    alpha=1.0,
                     sos_id=0,
                     eos_id=beam_search.EOS_ID,
                     batch_size=None,
@@ -672,7 +698,8 @@ def fast_decode_tpu(encoder_output,
                     scope_prefix="body/"):
   """Given encoder output and a symbols to logits function, does fast decoding.
 
-  Implements only greedy decoding for TPU.
+  Implements both greedy and beam search decoding for TPU, uses beam search iff
+  beam_size > 1, otherwise beam search related arguments are ignored.
 
   Args:
     encoder_output: A tensor, output from encoder.
@@ -682,7 +709,11 @@ def fast_decode_tpu(encoder_output,
         `(ids, step, cache)` to symbol logits.
     hparams: Run hyperparameters.
     decode_length: An integer, how many additional timesteps to decode.
+    vocab_size: Output vocabulary size.
     beam_size: An integer, number of beams.
+    top_beams: An integer, how many of the beams to return.
+    alpha: A float that controls the length penalty. Larger the alpha, stronger
+      the preference for longer translations.
     sos_id: Start-of-sequence symbol.
     eos_id: End-of-sequence symbol.
     batch_size: An integer, must be passed if there is no input.
@@ -691,15 +722,16 @@ def fast_decode_tpu(encoder_output,
     scope_prefix: str, prefix for decoder layer variable scopes.
 
   Returns:
-      A dict of decoding results {
-          "outputs": integer `Tensor` of decoded ids of shape
-              [batch_size, <= decode_length]
-          "scores": decoding log probs from the beam search,
-              None if using greedy decoding (beam_size=1)
-      }.
+    A dict of decoding results {
+        "outputs": integer `Tensor` of decoded ids of shape
+            [batch_size, <= decode_length] if top_beams == 1 or
+            [batch_size, top_beams, <= decode_length] otherwise
+        "scores": decoding log probs from the beam search,
+            None if using greedy decoding (beam_size=1)
+    }.
 
   Raises:
-     NotImplementedError: If beam size > 1.
+    NotImplementedError: If beam size > 1 with partial targets.
   """
   if encoder_output is not None:
     batch_size = common_layers.shape_list(encoder_output)[0]
@@ -707,6 +739,8 @@ def fast_decode_tpu(encoder_output,
   key_channels = hparams.attention_key_channels or hparams.hidden_size
   value_channels = hparams.attention_value_channels or hparams.hidden_size
   num_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
+  vars_3d_num_heads = (
+      hparams.num_heads if hparams.get("attention_variables_3d") else 0)
 
   cache = {
       "layer_%d" % layer: {
@@ -730,10 +764,12 @@ def fast_decode_tpu(encoder_output,
           "%sdecoder/%s/encdec_attention/multihead_attention" % (scope_prefix,
                                                                  layer_name)):
         k_encdec = common_attention.compute_attention_component(
-            encoder_output, key_channels, name="k")
+            encoder_output, key_channels, name="k",
+            vars_3d_num_heads=vars_3d_num_heads)
         k_encdec = common_attention.split_heads(k_encdec, hparams.num_heads)
         v_encdec = common_attention.compute_attention_component(
-            encoder_output, value_channels, name="v")
+            encoder_output, value_channels, name="v",
+            vars_3d_num_heads=vars_3d_num_heads)
         v_encdec = common_attention.split_heads(v_encdec, hparams.num_heads)
       cache[layer_name]["k_encdec"] = k_encdec
       cache[layer_name]["v_encdec"] = v_encdec
@@ -742,58 +778,75 @@ def fast_decode_tpu(encoder_output,
     cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias
 
   if beam_size > 1:  # Beam Search
-    raise NotImplementedError("Beam search inference on TPU is not supported")
+    initial_ids = sos_id * tf.ones([batch_size], dtype=tf.int32)
+    decoded_ids, scores = beam_search.beam_search(
+        symbols_to_logits_fn,
+        initial_ids,
+        beam_size,
+        decode_length,
+        vocab_size,
+        alpha,
+        states=cache,
+        eos_id=eos_id,
+        stop_early=(top_beams == 1),
+        use_tpu=True)
 
-  # Greedy
-  def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
-    """One step of greedy decoding."""
-    logits, cache = symbols_to_logits_fn(next_id, i, cache)
-    log_probs = common_layers.log_prob_from_logits(logits)
-    temperature = (0.0 if hparams.sampling_method == "argmax" else
-                   hparams.sampling_temp)
-    next_id = common_layers.sample_with_temperature(logits, temperature)
-    hit_eos |= tf.equal(next_id, eos_id)
-
-    log_prob_indices = tf.stack(
-        [tf.range(tf.to_int64(batch_size)), next_id], axis=1)
-    log_prob += tf.gather_nd(log_probs, log_prob_indices)
-
-    next_id = tf.expand_dims(next_id, axis=1)
-    decoded_ids = tf.transpose(decoded_ids)
-    decoded_ids = inplace_ops.alias_inplace_update(
-        decoded_ids, i, tf.squeeze(next_id, axis=1))
-    decoded_ids = tf.transpose(decoded_ids)
-    return i + 1, hit_eos, next_id, decoded_ids, cache, log_prob
-
-  def is_not_finished(i, hit_eos, *_):
-    finished = i >= decode_length
-    if not force_decode_length:
-      finished |= tf.reduce_all(hit_eos)
-    return tf.logical_not(finished)
-
-  decoded_ids = tf.zeros([batch_size, decode_length], dtype=tf.int64)
-  hit_eos = tf.fill([batch_size], False)
-  next_id = sos_id * tf.ones([batch_size, 1], dtype=tf.int64)
-  initial_log_prob = tf.zeros([batch_size], dtype=tf.float32)
-
-  def compute_cache_shape_invariants(tensor):
-    return tf.TensorShape(tensor.shape.as_list())
-
-  _, _, _, decoded_ids, _, log_prob = tf.while_loop(
-      is_not_finished,
-      inner_loop, [
-          tf.constant(0), hit_eos, next_id, decoded_ids, cache,
-          initial_log_prob
-      ],
-      shape_invariants=[
-          tf.TensorShape([]),
-          tf.TensorShape([batch_size]),
-          tf.TensorShape([batch_size, 1]),
-          tf.TensorShape([batch_size, decode_length]),
-          nest.map_structure(compute_cache_shape_invariants, cache),
-          tf.TensorShape([batch_size]),
-      ])
-  scores = log_prob
+    if top_beams == 1:
+      decoded_ids = decoded_ids[:, 0, 1:]
+      scores = scores[:, 0]
+    else:
+      decoded_ids = decoded_ids[:, :top_beams, 1:]
+      scores = scores[:, :top_beams]
+  else:  # Greedy
+    def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
+      """One step of greedy decoding."""
+      logits, cache = symbols_to_logits_fn(next_id, i, cache)
+      log_probs = common_layers.log_prob_from_logits(logits)
+      temperature = (0.0 if hparams.sampling_method == "argmax" else
+                     hparams.sampling_temp)
+      next_id = common_layers.sample_with_temperature(logits, temperature)
+      hit_eos |= tf.equal(next_id, eos_id)
+
+      log_prob_indices = tf.stack(
+          [tf.range(tf.to_int64(batch_size)), next_id], axis=1)
+      log_prob += tf.gather_nd(log_probs, log_prob_indices)
+
+      next_id = tf.expand_dims(next_id, axis=1)
+      decoded_ids = tf.transpose(decoded_ids)
+      decoded_ids = inplace_ops.alias_inplace_update(
+          decoded_ids, i, tf.squeeze(next_id, axis=1))
+      decoded_ids = tf.transpose(decoded_ids)
+      return i + 1, hit_eos, next_id, decoded_ids, cache, log_prob
+
+    def is_not_finished(i, hit_eos, *_):
+      finished = i >= decode_length
+      if not force_decode_length:
+        finished |= tf.reduce_all(hit_eos)
+      return tf.logical_not(finished)
+
+    decoded_ids = tf.zeros([batch_size, decode_length], dtype=tf.int64)
+    hit_eos = tf.fill([batch_size], False)
+    next_id = sos_id * tf.ones([batch_size, 1], dtype=tf.int64)
+    initial_log_prob = tf.zeros([batch_size], dtype=tf.float32)
+
+    def compute_cache_shape_invariants(tensor):
+      return tf.TensorShape(tensor.shape.as_list())
+
+    _, _, _, decoded_ids, _, log_prob = tf.while_loop(
+        is_not_finished,
+        inner_loop, [
+            tf.constant(0), hit_eos, next_id, decoded_ids, cache,
+            initial_log_prob
+        ],
+        shape_invariants=[
+            tf.TensorShape([]),
+            tf.TensorShape([batch_size]),
+            tf.TensorShape([batch_size, 1]),
+            tf.TensorShape([batch_size, decode_length]),
+            nest.map_structure(compute_cache_shape_invariants, cache),
+            tf.TensorShape([batch_size]),
+        ])
+    scores = log_prob
 
   return {"outputs": decoded_ids, "scores": scores}
 
diff --git a/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index fff8bd226..98ce137b9 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -21,6 +21,7 @@
 
 import tensorflow as tf
 
+from tensorflow.python.ops import inplace_ops
 from tensorflow.python.util import nest
 
 # Assuming EOS_ID is 1
@@ -177,7 +178,8 @@ def beam_search(symbols_to_logits_fn,
                 alpha,
                 states=None,
                 eos_id=EOS_ID,
-                stop_early=True):
+                stop_early=True,
+                use_tpu=False):
   """Beam search with length penalties.
 
   Requires a function that can take the currently decoded symbols and return
@@ -218,6 +220,8 @@ def beam_search(symbols_to_logits_fn,
     states: dict (possibly nested) of decoding states.
     eos_id: ID for end of sentence.
     stop_early: a boolean - stop once best sequence is provably determined.
+    use_tpu: A bool, whether to do beam search on TPU.
+
   Returns:
     Tuple of
     (decoded beams [batch_size, beam_size, decode_length]
@@ -233,6 +237,8 @@ def beam_search(symbols_to_logits_fn,
   # Expand each batch and state to beam_size
   alive_seq = _expand_to_beam_size(initial_ids, beam_size)
   alive_seq = tf.expand_dims(alive_seq, axis=2)  # (batch_size, beam_size, 1)
+  if use_tpu:
+    alive_seq = tf.tile(alive_seq, [1, 1, decode_length + 1])
   if states:
     states = nest.map_structure(
         lambda state: _expand_to_beam_size(state, beam_size), states)
@@ -269,11 +275,12 @@ def grow_finished(finished_seq, finished_scores, finished_flags, curr_seq,
          log probs of these sequences,
          Finished flags of these sequences)
     """
-    # First append a column of 0'ids to finished to make the same length with
-    # finished scores
-    finished_seq = tf.concat(
-        [finished_seq,
-         tf.zeros([batch_size, beam_size, 1], tf.int32)], axis=2)
+    if not use_tpu:
+      # First append a column of 0'ids to finished to make the same length with
+      # finished scores
+      finished_seq = tf.concat(
+          [finished_seq,
+           tf.zeros([batch_size, beam_size, 1], tf.int32)], axis=2)
 
     # Set the scores of the unfinished seq in curr_seq to large negative
     # values
@@ -338,7 +345,12 @@ def grow_topk(i, alive_seq, alive_log_probs, states):
          dict of transformed decoding states)
     """
     # Get the logits for all the possible next symbols
-    flat_ids = tf.reshape(alive_seq, [batch_size * beam_size, -1])
+    if use_tpu:
+      flat_ids = tf.reshape(
+          tf.slice(alive_seq, [0, 0, i], [batch_size, beam_size, 1]),
+          [batch_size * beam_size, -1])
+    else:
+      flat_ids = tf.reshape(alive_seq, [batch_size * beam_size, -1])
 
     # (batch_size * beam_size, decoded_length)
     if states:
@@ -392,7 +404,12 @@ def grow_topk(i, alive_seq, alive_log_probs, states):
           lambda state: tf.gather_nd(state, topk_coordinates), states)
 
     # Append the most probable alive
-    topk_seq = tf.concat([topk_seq, tf.expand_dims(topk_ids, axis=2)], axis=2)
+    if use_tpu:
+      topk_seq = tf.transpose(topk_seq, perm=[2, 0, 1])
+      topk_seq = inplace_ops.alias_inplace_update(topk_seq, i + 1, topk_ids)
+      topk_seq = tf.transpose(topk_seq, perm=[1, 2, 0])
+    else:
+      topk_seq = tf.concat([topk_seq, tf.expand_dims(topk_ids, axis=2)], axis=2)
 
     topk_finished = tf.equal(topk_ids, eos_id)
 
@@ -513,9 +530,11 @@ def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
        ],
        shape_invariants=[
            tf.TensorShape([]),
-           tf.TensorShape([None, None, None]),
+           (tf.TensorShape([batch_size, beam_size, decode_length + 1])
+            if use_tpu else tf.TensorShape([None, None, None])),
            alive_log_probs.get_shape(),
-           tf.TensorShape([None, None, None]),
+           (tf.TensorShape([batch_size, beam_size, decode_length + 1])
+            if use_tpu else tf.TensorShape([None, None, None])),
            finished_scores.get_shape(),
            finished_flags.get_shape(),
            nest.map_structure(get_state_shape_invariants, states),
diff --git a/tensor2tensor/utils/beam_search_test.py b/tensor2tensor/utils/beam_search_test.py
index bd6a9cc97..a4277f2b2 100644
--- a/tensor2tensor/utils/beam_search_test.py
+++ b/tensor2tensor/utils/beam_search_test.py
@@ -369,5 +369,55 @@ def symbols_to_logits(ids, _, states):
       except tf.errors.InvalidArgumentError as e:
         raise AssertionError(e.message)
 
+  def testTPUBeam(self):
+    batch_size = 1
+    beam_size = 2
+    vocab_size = 3
+    decode_length = 3
+
+    initial_ids = tf.constant([0] * batch_size)  # GO
+    probabilities = tf.constant([[[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
+                                 [[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
+                                 [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]]])
+
+    # The top beam is always selected so we should see the top beam's state
+    # at each position, which is the one thats getting 3 added to it each step.
+    expected_states = tf.constant([[[0.], [0.]], [[3.], [3.]], [[6.], [6.]]])
+
+    def symbols_to_logits(_, i, states):
+      # We have to assert the values of state inline here since we can't fetch
+      # them out of the loop!
+      with tf.control_dependencies(
+          [tf.assert_equal(states["state"], expected_states[i])]):
+        logits = tf.to_float(tf.log(probabilities[i, :]))
+
+      states["state"] += tf.constant([[3.], [7.]])
+      return logits, states
+
+    states = {
+        "state": tf.zeros((batch_size, 1)),
+    }
+    states["state"] = tf.placeholder_with_default(
+        states["state"], shape=(None, 1))
+
+    final_ids, _ = beam_search.beam_search(
+        symbols_to_logits,
+        initial_ids,
+        beam_size,
+        decode_length,
+        vocab_size,
+        3.5,
+        eos_id=1,
+        states=states,
+        use_tpu=True)
+
+    with self.test_session() as sess:
+      # Catch and fail so that the testing framework doesn't think it's an error
+      try:
+        sess.run(final_ids)
+      except tf.errors.InvalidArgumentError as e:
+        raise AssertionError(e.message)
+    self.assertAllEqual([[[0, 2, 0, 1], [0, 2, 1, 0]]], final_ids)
+
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index ffc07334a..fed8e8387 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -20,6 +20,7 @@
 import collections
 import operator
 import os
+import re
 import time
 
 import numpy as np
@@ -197,6 +198,10 @@ def decode_from_dataset(estimator,
       output_dirs = [output_dir]
       predictions.append(result)
 
+  if decode_hp.decode_to_file:
+    decode_hp.decode_to_file = _decode_filename(
+        decode_hp.decode_to_file, problem_name, decode_hp)
+
   run_postdecode_hooks(DecodeHookArgs(
       estimator=estimator,
       problem=problem,
@@ -229,11 +234,7 @@ def decode_once(estimator,
   # Prepare output file writers if decode_to_file passed
   decode_to_file = decode_to_file or decode_hp.decode_to_file
   if decode_to_file:
-    if decode_hp.shards > 1:
-      decode_filename = decode_to_file + ("%.2d" % decode_hp.shard_id)
-    else:
-      decode_filename = decode_to_file
-    output_filepath = _decode_filename(decode_filename, problem_name, decode_hp)
+    output_filepath = _decode_filename(decode_to_file, problem_name, decode_hp)
     parts = output_filepath.split(".")
     parts[-1] = "targets"
     target_filepath = ".".join(parts)
@@ -302,6 +303,9 @@ def decode_once(estimator,
     # Write out predictions if decode_to_file passed
     if decode_to_file:
       for i, (d_input, d_output, d_target) in enumerate(decoded_outputs):
+        # Skip if all padding
+        if re.match("^({})+$".format(text_encoder.PAD), d_input):
+          continue
         beam_score_str = ""
         if decode_hp.write_beam_scores:
           beam_score_str = "\t%.2f" % decoded_scores[i]
@@ -422,8 +426,6 @@ def timer(gen):
   # (except for adding shard_id if using more shards for decoding).
   # Otherwise, use the input filename plus model, hp, problem, beam, alpha.
   decode_filename = decode_to_file if decode_to_file else filename
-  if decode_hp.shards > 1:
-    decode_filename += "%.2d" % decode_hp.shard_id
   if not decode_to_file:
     decode_filename = _decode_filename(decode_filename, problem_name, decode_hp)
   tf.logging.info("Writing decodes into %s" % decode_filename)
@@ -447,13 +449,31 @@ def timer(gen):
 
 
 def _decode_filename(base_filename, problem_name, decode_hp):
-  return "{base}.{model}.{hp}.{problem}.beam{beam}.alpha{alpha}.decodes".format(
-      base=base_filename,
-      model=FLAGS.model,
-      hp=FLAGS.hparams_set,
-      problem=problem_name,
-      beam=str(decode_hp.beam_size),
-      alpha=str(decode_hp.alpha))
+  """Generates decode filename.
+
+  Args:
+    base_filename: A string, base of the decode filename.
+    problem_name: A string, name of the problem.
+    decode_hp: HParams for decoding.
+
+  Returns:
+    A string, produced decode filename.
+  """
+  if decode_hp.shards > 1:
+    base_filename = base_filename + ("%.2d" % decode_hp.shard_id)
+  if ("beam{beam}.alpha{alpha}.decodes".format(
+      beam=str(decode_hp.beam_size), alpha=str(decode_hp.alpha))
+      in base_filename):
+    return base_filename
+  else:
+    return (
+        "{base}.{model}.{hp}.{problem}.beam{beam}.alpha{alpha}.decodes".format(
+            base=base_filename,
+            model=FLAGS.model,
+            hp=FLAGS.hparams_set,
+            problem=problem_name,
+            beam=str(decode_hp.beam_size),
+            alpha=str(decode_hp.alpha)))
 
 
 def make_input_fn_from_generator(gen):
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index a2de8d6f6..68b9ce9e6 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -657,11 +657,17 @@ def infer(self,
       else:
         log_info("Beam Decoding with beam size %d" % beam_size)
         results = self._beam_decode(features, decode_length, beam_size,
-                                    top_beams, alpha)
+                                    top_beams, alpha, use_tpu)
 
       return results
 
-  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
+  def _beam_decode(self,
+                   features,
+                   decode_length,
+                   beam_size,
+                   top_beams,
+                   alpha,
+                   use_tpu=False):
     """Beam search decoding.
 
     Models should ideally implement a more efficient version of this function.
@@ -673,15 +679,16 @@ def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
       top_beams: an integer. How many of the beams to return.
       alpha: Float that controls the length penalty. larger the alpha, stronger
         the preference for longer translations.
+      use_tpu: A bool, whether to do beam decode on TPU.
 
     Returns:
        samples: an integer `Tensor`. Top samples from the beam search
     """
     return self._beam_decode_slow(features, decode_length, beam_size, top_beams,
-                                  alpha)
+                                  alpha, use_tpu)
 
   def _beam_decode_slow(self, features, decode_length, beam_size, top_beams,
-                        alpha):
+                        alpha, use_tpu=False):
     """Slow version of Beam search decoding.
 
     Quadratic time in decode_length.
@@ -693,10 +700,18 @@ def _beam_decode_slow(self, features, decode_length, beam_size, top_beams,
       top_beams: an integer. How many of the beams to return.
       alpha: Float that controls the length penalty. larger the alpha, stronger
         the preference for longer translations.
+      use_tpu: A bool, whether to do slow beam decode on TPU.
 
     Returns:
-       samples: an integer `Tensor`. Top samples from the beam search
+      samples: an integer `Tensor`. Top samples from the beam search.
+
+    Raises:
+      NotImplementedError: If use_tpu is set to true.
     """
+    if use_tpu:
+      raise NotImplementedError(
+          "Slow beam search inference on TPU is not supported")
+
     batch_size = common_layers.shape_list(features["inputs"])[0]
 
     def symbols_to_logits_fn(ids):
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 1564a42d7..b90f222b1 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -420,6 +420,11 @@ def continuous_decode_on_train_data(self):
     for _ in next_checkpoint(self._hparams.model_dir):
       self.decode(dataset_split=tf.estimator.ModeKeys.TRAIN)
 
+  def continuous_decode_on_eval_data(self):
+    """Decode from dataset on new checkpoint."""
+    for _ in next_checkpoint(self._hparams.model_dir):
+      self.decode(dataset_split=tf.estimator.ModeKeys.EVAL)
+
   def continuous_decode_from_file(self):
     """Decode from file on new checkpoint."""
     for _ in next_checkpoint(self._hparams.model_dir):

From a4f958a887f4f4466644dd0602bdd33985d61dd7 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 5 Oct 2018 13:39:58 -0700
Subject: [PATCH 0968/2720] internal

PiperOrigin-RevId: 215958567
---
 tensor2tensor/__init__.py                                   | 1 +
 tensor2tensor/bin/__init__.py                               | 1 +
 tensor2tensor/bin/build_vocab.py                            | 1 +
 tensor2tensor/bin/make_tf_configs.py                        | 1 +
 tensor2tensor/bin/t2t_attack.py                             | 1 +
 tensor2tensor/bin/t2t_avg_all.py                            | 1 +
 tensor2tensor/bin/t2t_bleu.py                               | 1 +
 tensor2tensor/bin/t2t_datagen.py                            | 1 +
 tensor2tensor/bin/t2t_decoder.py                            | 1 +
 tensor2tensor/bin/t2t_distill.py                            | 1 +
 tensor2tensor/bin/t2t_prune.py                              | 1 +
 tensor2tensor/bin/t2t_trainer.py                            | 1 +
 tensor2tensor/bin/t2t_trainer_test.py                       | 1 +
 tensor2tensor/bin/t2t_translate_all.py                      | 1 +
 tensor2tensor/data_generators/__init__.py                   | 1 +
 tensor2tensor/data_generators/algorithmic.py                | 1 +
 tensor2tensor/data_generators/algorithmic_math.py           | 1 +
 tensor2tensor/data_generators/algorithmic_math_test.py      | 1 +
 tensor2tensor/data_generators/algorithmic_test.py           | 1 +
 tensor2tensor/data_generators/all_problems.py               | 1 +
 tensor2tensor/data_generators/allen_brain.py                | 1 +
 tensor2tensor/data_generators/allen_brain_test.py           | 1 +
 tensor2tensor/data_generators/audio.py                      | 1 +
 tensor2tensor/data_generators/audio_encoder.py              | 1 +
 tensor2tensor/data_generators/audio_test.py                 | 1 +
 tensor2tensor/data_generators/babi_qa.py                    | 1 +
 tensor2tensor/data_generators/bair_robot_pushing.py         | 1 +
 tensor2tensor/data_generators/celeba.py                     | 1 +
 tensor2tensor/data_generators/celeba_test.py                | 1 +
 tensor2tensor/data_generators/celebahq.py                   | 1 +
 tensor2tensor/data_generators/cifar.py                      | 1 +
 tensor2tensor/data_generators/cipher.py                     | 1 +
 tensor2tensor/data_generators/cnn_dailymail.py              | 1 +
 tensor2tensor/data_generators/cola.py                       | 1 +
 tensor2tensor/data_generators/common_voice.py               | 1 +
 tensor2tensor/data_generators/common_voice_test.py          | 1 +
 tensor2tensor/data_generators/desc2code.py                  | 1 +
 tensor2tensor/data_generators/desc2code_test.py             | 1 +
 tensor2tensor/data_generators/dna_encoder.py                | 1 +
 tensor2tensor/data_generators/dna_encoder_test.py           | 1 +
 tensor2tensor/data_generators/fsns.py                       | 1 +
 tensor2tensor/data_generators/function_docstring.py         | 1 +
 tensor2tensor/data_generators/gene_expression.py            | 1 +
 tensor2tensor/data_generators/gene_expression_test.py       | 1 +
 tensor2tensor/data_generators/generator_utils.py            | 1 +
 tensor2tensor/data_generators/generator_utils_test.py       | 1 +
 tensor2tensor/data_generators/google_robot_pushing.py       | 1 +
 tensor2tensor/data_generators/gym_problems.py               | 1 +
 tensor2tensor/data_generators/gym_problems_specs.py         | 1 +
 tensor2tensor/data_generators/gym_problems_test.py          | 1 +
 tensor2tensor/data_generators/gym_utils.py                  | 1 +
 tensor2tensor/data_generators/ice_parsing.py                | 1 +
 tensor2tensor/data_generators/image_lsun.py                 | 1 +
 tensor2tensor/data_generators/image_utils.py                | 1 +
 tensor2tensor/data_generators/image_utils_test.py           | 1 +
 tensor2tensor/data_generators/imagenet.py                   | 1 +
 tensor2tensor/data_generators/imagenet_test.py              | 1 +
 tensor2tensor/data_generators/imdb.py                       | 1 +
 tensor2tensor/data_generators/inspect_tfrecord.py           | 1 +
 tensor2tensor/data_generators/lambada.py                    | 1 +
 tensor2tensor/data_generators/librispeech.py                | 1 +
 tensor2tensor/data_generators/lm1b.py                       | 1 +
 tensor2tensor/data_generators/lm1b_imdb.py                  | 1 +
 tensor2tensor/data_generators/lm1b_mnli.py                  | 1 +
 tensor2tensor/data_generators/mnist.py                      | 1 +
 tensor2tensor/data_generators/mrpc.py                       | 1 +
 tensor2tensor/data_generators/mscoco.py                     | 1 +
 tensor2tensor/data_generators/mscoco_test.py                | 1 +
 tensor2tensor/data_generators/multi_problem.py              | 1 +
 tensor2tensor/data_generators/multi_problem_test.py         | 1 +
 tensor2tensor/data_generators/multinli.py                   | 1 +
 tensor2tensor/data_generators/ocr.py                        | 1 +
 tensor2tensor/data_generators/paraphrase_ms_coco.py         | 1 +
 tensor2tensor/data_generators/paraphrase_ms_coco_test.py    | 1 +
 tensor2tensor/data_generators/pointer_generator_word.py     | 1 +
 tensor2tensor/data_generators/problem.py                    | 1 +
 tensor2tensor/data_generators/problem_hparams.py            | 1 +
 tensor2tensor/data_generators/problem_test.py               | 1 +
 tensor2tensor/data_generators/program_search.py             | 1 +
 tensor2tensor/data_generators/program_search_test.py        | 1 +
 tensor2tensor/data_generators/ptb.py                        | 1 +
 tensor2tensor/data_generators/qnli.py                       | 1 +
 tensor2tensor/data_generators/quora_qpairs.py               | 1 +
 tensor2tensor/data_generators/rte.py                        | 1 +
 tensor2tensor/data_generators/scitail.py                    | 1 +
 tensor2tensor/data_generators/snli.py                       | 1 +
 tensor2tensor/data_generators/speech_recognition.py         | 1 +
 tensor2tensor/data_generators/squad.py                      | 1 +
 tensor2tensor/data_generators/sst_binary.py                 | 1 +
 tensor2tensor/data_generators/stanford_nli.py               | 1 +
 tensor2tensor/data_generators/style_transfer.py             | 1 +
 tensor2tensor/data_generators/style_transfer_test.py        | 1 +
 tensor2tensor/data_generators/subject_verb_agreement.py     | 1 +
 tensor2tensor/data_generators/text_encoder.py               | 1 +
 tensor2tensor/data_generators/text_encoder_build_subword.py | 1 +
 tensor2tensor/data_generators/text_encoder_test.py          | 1 +
 tensor2tensor/data_generators/text_problems.py              | 1 +
 tensor2tensor/data_generators/text_problems_test.py         | 1 +
 tensor2tensor/data_generators/timeseries.py                 | 1 +
 tensor2tensor/data_generators/timeseries_data_generator.py  | 1 +
 .../data_generators/timeseries_data_generator_test.py       | 1 +
 tensor2tensor/data_generators/timeseries_test.py            | 1 +
 tensor2tensor/data_generators/tokenizer.py                  | 1 +
 tensor2tensor/data_generators/tokenizer_test.py             | 1 +
 tensor2tensor/data_generators/translate.py                  | 1 +
 tensor2tensor/data_generators/translate_encs.py             | 1 +
 tensor2tensor/data_generators/translate_ende.py             | 1 +
 tensor2tensor/data_generators/translate_enet.py             | 1 +
 tensor2tensor/data_generators/translate_enfr.py             | 1 +
 tensor2tensor/data_generators/translate_enid.py             | 1 +
 tensor2tensor/data_generators/translate_enmk.py             | 1 +
 tensor2tensor/data_generators/translate_envi.py             | 1 +
 tensor2tensor/data_generators/translate_enzh.py             | 1 +
 tensor2tensor/data_generators/translate_test.py             | 1 +
 tensor2tensor/data_generators/twentybn.py                   | 1 +
 tensor2tensor/data_generators/video_generated.py            | 1 +
 tensor2tensor/data_generators/video_utils.py                | 1 +
 tensor2tensor/data_generators/video_utils_test.py           | 1 +
 tensor2tensor/data_generators/vqa.py                        | 1 +
 tensor2tensor/data_generators/vqa_utils.py                  | 1 +
 tensor2tensor/data_generators/wiki.py                       | 1 +
 tensor2tensor/data_generators/wikisum/__init__.py           | 1 +
 tensor2tensor/data_generators/wikisum/generate_vocab.py     | 1 +
 .../data_generators/wikisum/get_references_commoncrawl.py   | 1 +
 tensor2tensor/data_generators/wikisum/get_references_web.py | 1 +
 .../wikisum/get_references_web_single_group.py              | 1 +
 tensor2tensor/data_generators/wikisum/html.py               | 1 +
 tensor2tensor/data_generators/wikisum/parallel_launch.py    | 1 +
 tensor2tensor/data_generators/wikisum/produce_examples.py   | 1 +
 tensor2tensor/data_generators/wikisum/utils.py              | 1 +
 tensor2tensor/data_generators/wikisum/utils_test.py         | 1 +
 tensor2tensor/data_generators/wikisum/validate_data.py      | 1 +
 tensor2tensor/data_generators/wikisum/wikisum.py            | 1 +
 tensor2tensor/data_generators/wikitext103.py                | 1 +
 tensor2tensor/data_generators/wnli.py                       | 1 +
 tensor2tensor/data_generators/wsj_parsing.py                | 1 +
 tensor2tensor/insights/__init__.py                          | 1 +
 tensor2tensor/insights/graph.py                             | 1 +
 tensor2tensor/insights/query_processor.py                   | 1 +
 tensor2tensor/insights/server.py                            | 1 +
 tensor2tensor/insights/transformer_model.py                 | 1 +
 tensor2tensor/layers/__init__.py                            | 1 +
 tensor2tensor/layers/common_attention.py                    | 1 +
 tensor2tensor/layers/common_attention_test.py               | 1 +
 tensor2tensor/layers/common_audio.py                        | 1 +
 tensor2tensor/layers/common_hparams.py                      | 1 +
 tensor2tensor/layers/common_image_attention.py              | 1 +
 tensor2tensor/layers/common_image_attention_test.py         | 1 +
 tensor2tensor/layers/common_layers.py                       | 1 +
 tensor2tensor/layers/common_layers_test.py                  | 1 +
 tensor2tensor/layers/common_message_passing_attention.py    | 1 +
 tensor2tensor/layers/common_video.py                        | 1 +
 tensor2tensor/layers/common_video_test.py                   | 1 +
 tensor2tensor/layers/discretization.py                      | 1 +
 tensor2tensor/layers/discretization_test.py                 | 1 +
 tensor2tensor/layers/latent_layers.py                       | 1 +
 tensor2tensor/layers/latent_layers_test.py                  | 1 +
 tensor2tensor/layers/modalities.py                          | 1 +
 tensor2tensor/layers/modalities_test.py                     | 1 +
 tensor2tensor/layers/vq_discrete.py                         | 1 +
 tensor2tensor/layers/vqa_layers.py                          | 1 +
 tensor2tensor/mesh_tensorflow/OWNERS                        | 2 ++
 tensor2tensor/mesh_tensorflow/__init__.py                   | 1 +
 tensor2tensor/mesh_tensorflow/mesh_tensorflow.py            | 1 +
 tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py       | 1 +
 tensor2tensor/mesh_tensorflow/mnist.py                      | 1 +
 tensor2tensor/mesh_tensorflow/mnist_dataset.py              | 1 +
 tensor2tensor/mesh_tensorflow/mtf_beam_search.py            | 1 +
 tensor2tensor/mesh_tensorflow/mtf_image_transformer.py      | 1 +
 tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py | 1 +
 tensor2tensor/mesh_tensorflow/mtf_layers.py                 | 1 +
 tensor2tensor/mesh_tensorflow/mtf_layers_test.py            | 1 +
 tensor2tensor/mesh_tensorflow/mtf_model.py                  | 1 +
 tensor2tensor/mesh_tensorflow/mtf_optimize.py               | 1 +
 tensor2tensor/mesh_tensorflow/mtf_resnet.py                 | 1 +
 tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py          | 1 +
 tensor2tensor/mesh_tensorflow/mtf_transformer.py            | 1 +
 tensor2tensor/mesh_tensorflow/mtf_transformer_test.py       | 1 +
 tensor2tensor/mesh_tensorflow/mtf_utils.py                  | 1 +
 tensor2tensor/mesh_tensorflow/mtf_utils_test.py             | 1 +
 tensor2tensor/mesh_tensorflow/placement_mesh_impl.py        | 1 +
 tensor2tensor/mesh_tensorflow/research/__init__.py          | 1 +
 tensor2tensor/mesh_tensorflow/research/experiments_moe.py   | 1 +
 tensor2tensor/mesh_tensorflow/research/moe.py               | 1 +
 tensor2tensor/mesh_tensorflow/simd_mesh_impl.py             | 1 +
 tensor2tensor/mesh_tensorflow/tpu_variables.py              | 1 +
 tensor2tensor/models/OWNERS                                 | 6 ++++++
 tensor2tensor/models/__init__.py                            | 1 +
 tensor2tensor/models/basic.py                               | 1 +
 tensor2tensor/models/basic_test.py                          | 1 +
 tensor2tensor/models/bytenet.py                             | 1 +
 tensor2tensor/models/bytenet_test.py                        | 1 +
 tensor2tensor/models/distillation.py                        | 1 +
 tensor2tensor/models/image_transformer.py                   | 1 +
 tensor2tensor/models/image_transformer_2d.py                | 1 +
 tensor2tensor/models/image_transformer_2d_test.py           | 1 +
 tensor2tensor/models/image_transformer_test.py              | 1 +
 tensor2tensor/models/lstm.py                                | 1 +
 tensor2tensor/models/lstm_test.py                           | 1 +
 tensor2tensor/models/neural_gpu.py                          | 1 +
 tensor2tensor/models/neural_gpu_test.py                     | 1 +
 tensor2tensor/models/research/__init__.py                   | 1 +
 tensor2tensor/models/research/adafactor_experiments.py      | 1 +
 tensor2tensor/models/research/aligned.py                    | 1 +
 tensor2tensor/models/research/attention_lm.py               | 1 +
 tensor2tensor/models/research/attention_lm_moe.py           | 1 +
 tensor2tensor/models/research/autoencoders.py               | 1 +
 tensor2tensor/models/research/autoencoders_test.py          | 1 +
 tensor2tensor/models/research/cycle_gan.py                  | 1 +
 tensor2tensor/models/research/gene_expression.py            | 1 +
 tensor2tensor/models/research/gene_expression_test.py       | 1 +
 tensor2tensor/models/research/glow.py                       | 1 +
 tensor2tensor/models/research/glow_ops.py                   | 1 +
 tensor2tensor/models/research/glow_ops_test.py              | 1 +
 tensor2tensor/models/research/glow_test.py                  | 1 +
 tensor2tensor/models/research/lm_experiments.py             | 1 +
 tensor2tensor/models/research/rl.py                         | 1 +
 tensor2tensor/models/research/similarity_transformer.py     | 1 +
 tensor2tensor/models/research/super_lm.py                   | 1 +
 tensor2tensor/models/research/transformer_aux.py            | 1 +
 tensor2tensor/models/research/transformer_aux_test.py       | 1 +
 tensor2tensor/models/research/transformer_moe.py            | 1 +
 tensor2tensor/models/research/transformer_nat.py            | 1 +
 tensor2tensor/models/research/transformer_revnet.py         | 1 +
 tensor2tensor/models/research/transformer_revnet_test.py    | 1 +
 tensor2tensor/models/research/transformer_sketch.py         | 1 +
 tensor2tensor/models/research/transformer_symshard.py       | 1 +
 tensor2tensor/models/research/transformer_vae.py            | 1 +
 tensor2tensor/models/research/transformer_vae_test.py       | 1 +
 tensor2tensor/models/research/universal_transformer.py      | 1 +
 tensor2tensor/models/research/universal_transformer_test.py | 1 +
 tensor2tensor/models/research/universal_transformer_util.py | 1 +
 tensor2tensor/models/research/vqa_attention.py              | 1 +
 tensor2tensor/models/research/vqa_attention_test.py         | 1 +
 .../models/research/vqa_recurrent_self_attention.py         | 1 +
 tensor2tensor/models/research/vqa_self_attention.py         | 1 +
 tensor2tensor/models/resnet.py                              | 1 +
 tensor2tensor/models/resnet_test.py                         | 1 +
 tensor2tensor/models/revnet.py                              | 1 +
 tensor2tensor/models/revnet_test.py                         | 1 +
 tensor2tensor/models/shake_shake.py                         | 1 +
 tensor2tensor/models/slicenet.py                            | 1 +
 tensor2tensor/models/slicenet_test.py                       | 1 +
 tensor2tensor/models/transformer.py                         | 1 +
 tensor2tensor/models/transformer_test.py                    | 1 +
 tensor2tensor/models/vanilla_gan.py                         | 1 +
 tensor2tensor/models/video/__init__.py                      | 1 +
 tensor2tensor/models/video/base_vae.py                      | 1 +
 tensor2tensor/models/video/basic_deterministic.py           | 1 +
 tensor2tensor/models/video/basic_deterministic_params.py    | 1 +
 tensor2tensor/models/video/basic_deterministic_test.py      | 1 +
 tensor2tensor/models/video/basic_recurrent.py               | 1 +
 tensor2tensor/models/video/basic_recurrent_test.py          | 1 +
 tensor2tensor/models/video/basic_stochastic.py              | 1 +
 tensor2tensor/models/video/basic_stochastic_test.py         | 1 +
 tensor2tensor/models/video/emily.py                         | 1 +
 tensor2tensor/models/video/emily_test.py                    | 1 +
 tensor2tensor/models/video/epva.py                          | 1 +
 tensor2tensor/models/video/epva_params.py                   | 1 +
 tensor2tensor/models/video/savp.py                          | 1 +
 tensor2tensor/models/video/savp_params.py                   | 1 +
 tensor2tensor/models/video/savp_test.py                     | 1 +
 tensor2tensor/models/video/sv2p.py                          | 1 +
 tensor2tensor/models/video/sv2p_params.py                   | 1 +
 tensor2tensor/models/video/sv2p_test.py                     | 1 +
 tensor2tensor/models/video/tests_utils.py                   | 1 +
 tensor2tensor/models/xception.py                            | 1 +
 tensor2tensor/models/xception_test.py                       | 1 +
 tensor2tensor/problems.py                                   | 1 +
 tensor2tensor/problems_test.py                              | 1 +
 tensor2tensor/rl/OWNERS                                     | 3 +++
 tensor2tensor/rl/__init__.py                                | 1 +
 tensor2tensor/rl/collect.py                                 | 1 +
 tensor2tensor/rl/datagen_with_agent.py                      | 1 +
 tensor2tensor/rl/envs/__init__.py                           | 1 +
 tensor2tensor/rl/envs/batch_env.py                          | 1 +
 tensor2tensor/rl/envs/batch_env_factory.py                  | 1 +
 tensor2tensor/rl/envs/in_graph_batch_env.py                 | 1 +
 tensor2tensor/rl/envs/py_func_batch_env.py                  | 1 +
 tensor2tensor/rl/envs/simulated_batch_env.py                | 1 +
 tensor2tensor/rl/envs/tf_atari_wrappers.py                  | 1 +
 tensor2tensor/rl/envs/utils.py                              | 1 +
 tensor2tensor/rl/model_rl_experiment_player.py              | 1 +
 tensor2tensor/rl/ppo.py                                     | 1 +
 tensor2tensor/rl/rl_trainer_lib.py                          | 1 +
 tensor2tensor/rl/rl_trainer_lib_test.py                     | 1 +
 tensor2tensor/rl/trainer_model_based.py                     | 1 +
 tensor2tensor/rl/trainer_model_based_ae_test.py             | 1 +
 tensor2tensor/rl/trainer_model_based_agent_only.py          | 1 +
 tensor2tensor/rl/trainer_model_based_recurrent_test.py      | 1 +
 tensor2tensor/rl/trainer_model_based_stochastic_test.py     | 1 +
 tensor2tensor/rl/trainer_model_based_sv2p_test.py           | 1 +
 tensor2tensor/rl/trainer_model_based_test.py                | 1 +
 tensor2tensor/rl/trainer_model_free.py                      | 1 +
 tensor2tensor/serving/__init__.py                           | 1 +
 tensor2tensor/serving/export.py                             | 1 +
 tensor2tensor/serving/query.py                              | 1 +
 tensor2tensor/serving/serving_utils.py                      | 1 +
 tensor2tensor/test_data/example_usr_dir/__init__.py         | 1 +
 tensor2tensor/test_data/example_usr_dir/my_submodule.py     | 1 +
 tensor2tensor/utils/__init__.py                             | 1 +
 tensor2tensor/utils/adafactor.py                            | 1 +
 tensor2tensor/utils/adv_attack_utils.py                     | 1 +
 tensor2tensor/utils/avg_checkpoints.py                      | 1 +
 tensor2tensor/utils/beam_search.py                          | 1 +
 tensor2tensor/utils/beam_search_test.py                     | 1 +
 tensor2tensor/utils/bleu_hook.py                            | 1 +
 tensor2tensor/utils/bleu_hook_test.py                       | 1 +
 tensor2tensor/utils/checkpoint_compatibility_test.py        | 1 +
 tensor2tensor/utils/cloud_mlengine.py                       | 1 +
 tensor2tensor/utils/compute_video_metrics.py                | 1 +
 tensor2tensor/utils/data_reader.py                          | 1 +
 tensor2tensor/utils/data_reader_test.py                     | 1 +
 tensor2tensor/utils/decoding.py                             | 1 +
 tensor2tensor/utils/devices.py                              | 1 +
 tensor2tensor/utils/diet.py                                 | 1 +
 tensor2tensor/utils/diet_test.py                            | 1 +
 tensor2tensor/utils/expert_utils.py                         | 1 +
 tensor2tensor/utils/expert_utils_test.py                    | 1 +
 tensor2tensor/utils/flags.py                                | 1 +
 tensor2tensor/utils/get_rouge.py                            | 1 +
 tensor2tensor/utils/learning_rate.py                        | 1 +
 tensor2tensor/utils/metrics.py                              | 1 +
 tensor2tensor/utils/metrics_hook.py                         | 1 +
 tensor2tensor/utils/metrics_hook_test.py                    | 1 +
 tensor2tensor/utils/metrics_test.py                         | 1 +
 tensor2tensor/utils/modality.py                             | 1 +
 tensor2tensor/utils/multistep_optimizer.py                  | 1 +
 tensor2tensor/utils/multistep_optimizer_test.py             | 1 +
 tensor2tensor/utils/optimize.py                             | 1 +
 tensor2tensor/utils/pruning_utils.py                        | 1 +
 tensor2tensor/utils/quantization.py                         | 1 +
 tensor2tensor/utils/registry.py                             | 1 +
 tensor2tensor/utils/registry_test.py                        | 1 +
 tensor2tensor/utils/restore_hook.py                         | 1 +
 tensor2tensor/utils/rouge.py                                | 1 +
 tensor2tensor/utils/rouge_test.py                           | 1 +
 tensor2tensor/utils/t2t_model.py                            | 1 +
 tensor2tensor/utils/t2t_model_test.py                       | 1 +
 tensor2tensor/utils/trainer_lib.py                          | 1 +
 tensor2tensor/utils/trainer_lib_test.py                     | 1 +
 tensor2tensor/utils/update_ops_hook.py                      | 1 +
 tensor2tensor/utils/usr_dir.py                              | 1 +
 tensor2tensor/utils/video2gif.py                            | 1 +
 tensor2tensor/utils/video_metrics.py                        | 1 +
 tensor2tensor/utils/yellowfin.py                            | 1 +
 tensor2tensor/utils/yellowfin_test.py                       | 1 +
 tensor2tensor/visualization/__init__.py                     | 1 +
 tensor2tensor/visualization/attention.py                    | 1 +
 tensor2tensor/visualization/visualization.py                | 1 +
 tensor2tensor/visualization/visualization_test.py           | 1 +
 351 files changed, 359 insertions(+)
 create mode 100644 tensor2tensor/mesh_tensorflow/OWNERS
 create mode 100644 tensor2tensor/models/OWNERS
 create mode 100644 tensor2tensor/rl/OWNERS

diff --git a/tensor2tensor/__init__.py b/tensor2tensor/__init__.py
index dba7ece95..4bd418a74 100644
--- a/tensor2tensor/__init__.py
+++ b/tensor2tensor/__init__.py
@@ -12,3 +12,4 @@
 # 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.
+
diff --git a/tensor2tensor/bin/__init__.py b/tensor2tensor/bin/__init__.py
index dba7ece95..4bd418a74 100644
--- a/tensor2tensor/bin/__init__.py
+++ b/tensor2tensor/bin/__init__.py
@@ -12,3 +12,4 @@
 # 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.
+
diff --git a/tensor2tensor/bin/build_vocab.py b/tensor2tensor/bin/build_vocab.py
index 0f363fa3c..b93ec94c2 100644
--- a/tensor2tensor/bin/build_vocab.py
+++ b/tensor2tensor/bin/build_vocab.py
@@ -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.
+
 r"""Build vocab for a subclass of Text2TextProblem.
 
 build_vocab \
diff --git a/tensor2tensor/bin/make_tf_configs.py b/tensor2tensor/bin/make_tf_configs.py
index 321b51d55..ef847bc3e 100644
--- a/tensor2tensor/bin/make_tf_configs.py
+++ b/tensor2tensor/bin/make_tf_configs.py
@@ -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.
+
 """Output command line arguments and json-encoded TF_CONFIGs.
 
 Usage:
diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index faabdcd8e..45cc22e42 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -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.
+
 r"""Adversarially attack a model.
 
 This script adversarially attacks a model and evaluates accuracy at various
diff --git a/tensor2tensor/bin/t2t_avg_all.py b/tensor2tensor/bin/t2t_avg_all.py
index 2843ff3ce..3755dc2d0 100644
--- a/tensor2tensor/bin/t2t_avg_all.py
+++ b/tensor2tensor/bin/t2t_avg_all.py
@@ -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.
+
 """Script to continuously average last N checkpoints in a given directory."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/bin/t2t_bleu.py b/tensor2tensor/bin/t2t_bleu.py
index db18ec91f..ab4a5f014 100644
--- a/tensor2tensor/bin/t2t_bleu.py
+++ b/tensor2tensor/bin/t2t_bleu.py
@@ -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.
+
 """Evaluate BLEU score for all checkpoints/translations in a given directory.
 
 This script can be used in two ways.
diff --git a/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index 4dfdc4d6f..b778bc984 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -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.
+
 """Produces the training and dev data for --problem into --data_dir.
 
 Produces sharded and shuffled TFRecord files of tensorflow.Example protocol
diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index b5d5d65f5..84644b3c4 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -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.
+
 r"""Decode from trained T2T models.
 
 This binary performs inference using the Estimator API.
diff --git a/tensor2tensor/bin/t2t_distill.py b/tensor2tensor/bin/t2t_distill.py
index 30822b1ae..f970a56c2 100644
--- a/tensor2tensor/bin/t2t_distill.py
+++ b/tensor2tensor/bin/t2t_distill.py
@@ -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.
+
 r"""Perform distillation for a teacher to student.
 
 This script is intended to be used with --model=distillation. See the model for
diff --git a/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
index 6e32597b5..c53724a8f 100644
--- a/tensor2tensor/bin/t2t_prune.py
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -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.
+
 r"""Prune T2TModels using some heuristic.
 
 This supports a very common form of pruning known as magnitude-based pruning.
diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index c36676470..2939fc4ad 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -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.
+
 """Train and evaluate."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/bin/t2t_trainer_test.py b/tensor2tensor/bin/t2t_trainer_test.py
index 9acb56850..63c206ab0 100644
--- a/tensor2tensor/bin/t2t_trainer_test.py
+++ b/tensor2tensor/bin/t2t_trainer_test.py
@@ -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.
+
 """Tests for t2t_trainer."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/bin/t2t_translate_all.py b/tensor2tensor/bin/t2t_translate_all.py
index 361720a0b..2d12e8246 100644
--- a/tensor2tensor/bin/t2t_translate_all.py
+++ b/tensor2tensor/bin/t2t_translate_all.py
@@ -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.
+
 """Translate a file with all checkpoints in a given directory.
 
 t2t-decoder will be executed with these parameters:
diff --git a/tensor2tensor/data_generators/__init__.py b/tensor2tensor/data_generators/__init__.py
index dba7ece95..4bd418a74 100644
--- a/tensor2tensor/data_generators/__init__.py
+++ b/tensor2tensor/data_generators/__init__.py
@@ -12,3 +12,4 @@
 # 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.
+
diff --git a/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index aea9a5e90..7e386b559 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -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.
+
 """Algorithmic data generators."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/algorithmic_math.py b/tensor2tensor/data_generators/algorithmic_math.py
index 604a51b53..1e08af9e6 100644
--- a/tensor2tensor/data_generators/algorithmic_math.py
+++ b/tensor2tensor/data_generators/algorithmic_math.py
@@ -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.
+
 """Algorithmic data generators for symbolic math tasks.
 
 See go/symbolic-math-dataset
diff --git a/tensor2tensor/data_generators/algorithmic_math_test.py b/tensor2tensor/data_generators/algorithmic_math_test.py
index 20eed95d2..bb4a3fae0 100644
--- a/tensor2tensor/data_generators/algorithmic_math_test.py
+++ b/tensor2tensor/data_generators/algorithmic_math_test.py
@@ -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.
+
 """Tests for tensor2tensor.data_generators.algorithmic_math."""
 # TODO(rsepassi): This test is flaky. Disable, remove, or update.
 
diff --git a/tensor2tensor/data_generators/algorithmic_test.py b/tensor2tensor/data_generators/algorithmic_test.py
index 548cc3f21..1034409a8 100644
--- a/tensor2tensor/data_generators/algorithmic_test.py
+++ b/tensor2tensor/data_generators/algorithmic_test.py
@@ -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.
+
 """Algorithmic generators test."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index e9fab0b2a..a872cbdac 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -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.
+
 """Imports for problem modules."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index 6ee1f5f8d..303872e4c 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -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.
+
 """Problem definitions for Allen Brain Atlas problems.
 
 Notes:
diff --git a/tensor2tensor/data_generators/allen_brain_test.py b/tensor2tensor/data_generators/allen_brain_test.py
index 739220aaa..7871c4657 100644
--- a/tensor2tensor/data_generators/allen_brain_test.py
+++ b/tensor2tensor/data_generators/allen_brain_test.py
@@ -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.
+
 """Tests of the Allen Brain Atlas problems."""
 
 import os
diff --git a/tensor2tensor/data_generators/audio.py b/tensor2tensor/data_generators/audio.py
index 96ed93d83..ea6863d7f 100644
--- a/tensor2tensor/data_generators/audio.py
+++ b/tensor2tensor/data_generators/audio.py
@@ -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.
+
 """TIMIT data generator."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/audio_encoder.py b/tensor2tensor/data_generators/audio_encoder.py
index 74cdb73e9..da224e7fb 100644
--- a/tensor2tensor/data_generators/audio_encoder.py
+++ b/tensor2tensor/data_generators/audio_encoder.py
@@ -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.
+
 """Encoder for audio data."""
 
 import os
diff --git a/tensor2tensor/data_generators/audio_test.py b/tensor2tensor/data_generators/audio_test.py
index 57a377cc0..4cc6eb461 100644
--- a/tensor2tensor/data_generators/audio_test.py
+++ b/tensor2tensor/data_generators/audio_test.py
@@ -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.
+
 """Tests for tensor2tensor.data_generators.audio."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index 395952a34..14c6fe70c 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -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.
+
 r"""Data generators for bAbi question answering dataset.
 
 
diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index eca4a2792..857e41c49 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -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.
+
 """Berkeley (BAIR) robot pushing dataset.
 
 Self-Supervised Visual Planning with Temporal Skip Connections
diff --git a/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index 7c75f8d69..b7518a204 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -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.
+
 """CelebA."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/celeba_test.py b/tensor2tensor/data_generators/celeba_test.py
index f9fb2b999..a347b3734 100644
--- a/tensor2tensor/data_generators/celeba_test.py
+++ b/tensor2tensor/data_generators/celeba_test.py
@@ -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.
+
 """Tests for CelebA."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/celebahq.py b/tensor2tensor/data_generators/celebahq.py
index 24e43e493..a206d3fe0 100644
--- a/tensor2tensor/data_generators/celebahq.py
+++ b/tensor2tensor/data_generators/celebahq.py
@@ -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.
+
 """CelebA-HQ."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index c070d237d..7ffdb869d 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -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.
+
 """CIFAR."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/cipher.py b/tensor2tensor/data_generators/cipher.py
index b4e4878b2..db07f5a7b 100644
--- a/tensor2tensor/data_generators/cipher.py
+++ b/tensor2tensor/data_generators/cipher.py
@@ -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.
+
 """Cipher data generators."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index 01f1642d6..c6a6f33cc 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -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.
+
 """Data generators for the CNN and Daily Mail datasets."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/cola.py b/tensor2tensor/data_generators/cola.py
index 4e5f4592f..7a905573c 100644
--- a/tensor2tensor/data_generators/cola.py
+++ b/tensor2tensor/data_generators/cola.py
@@ -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.
+
 """Data generators for the Corpus of Liguistic Acceptability."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/common_voice.py b/tensor2tensor/data_generators/common_voice.py
index fa9609aed..967e24477 100644
--- a/tensor2tensor/data_generators/common_voice.py
+++ b/tensor2tensor/data_generators/common_voice.py
@@ -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.
+
 """Mozilla Common Voice dataset.
 
 Note: Generating the full set of examples can take upwards of 5 hours.
diff --git a/tensor2tensor/data_generators/common_voice_test.py b/tensor2tensor/data_generators/common_voice_test.py
index 8da86a7e5..d7caefd32 100644
--- a/tensor2tensor/data_generators/common_voice_test.py
+++ b/tensor2tensor/data_generators/common_voice_test.py
@@ -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.
+
 """Tests for tensor2tensor.data_generators.common_voice."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/desc2code.py b/tensor2tensor/data_generators/desc2code.py
index 64e0dc219..383425f02 100644
--- a/tensor2tensor/data_generators/desc2code.py
+++ b/tensor2tensor/data_generators/desc2code.py
@@ -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.
+
 """Data generators for the Description2Code OpenAI data-set."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/desc2code_test.py b/tensor2tensor/data_generators/desc2code_test.py
index 1109daa3d..2b47f8271 100644
--- a/tensor2tensor/data_generators/desc2code_test.py
+++ b/tensor2tensor/data_generators/desc2code_test.py
@@ -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.
+
 """Tests for desc2code."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/dna_encoder.py b/tensor2tensor/data_generators/dna_encoder.py
index 297a8975b..e9c7afc93 100644
--- a/tensor2tensor/data_generators/dna_encoder.py
+++ b/tensor2tensor/data_generators/dna_encoder.py
@@ -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.
+
 """Encoders for DNA data.
 
 * DNAEncoder: ACTG strings to ints and back
diff --git a/tensor2tensor/data_generators/dna_encoder_test.py b/tensor2tensor/data_generators/dna_encoder_test.py
index afe4b568b..b35c28824 100644
--- a/tensor2tensor/data_generators/dna_encoder_test.py
+++ b/tensor2tensor/data_generators/dna_encoder_test.py
@@ -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.
+
 """Tests for tensor2tensor.data_generators.dna_encoder."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
index 6dc2567da..179d7bc85 100644
--- a/tensor2tensor/data_generators/fsns.py
+++ b/tensor2tensor/data_generators/fsns.py
@@ -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.
+
 """FSNS."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
index 7d72d5898..3d030a3d9 100644
--- a/tensor2tensor/data_generators/function_docstring.py
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -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.
+
 """Github function/text similatrity problems."""
 import csv
 from six import StringIO
diff --git a/tensor2tensor/data_generators/gene_expression.py b/tensor2tensor/data_generators/gene_expression.py
index 3f8acb3bc..f9ab3e0e3 100644
--- a/tensor2tensor/data_generators/gene_expression.py
+++ b/tensor2tensor/data_generators/gene_expression.py
@@ -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.
+
 """Gene expression problems.
 
 Inputs are bases ACTG (with indices assigned in that order).
diff --git a/tensor2tensor/data_generators/gene_expression_test.py b/tensor2tensor/data_generators/gene_expression_test.py
index 7df355fa5..ab98da1c7 100644
--- a/tensor2tensor/data_generators/gene_expression_test.py
+++ b/tensor2tensor/data_generators/gene_expression_test.py
@@ -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.
+
 """Tests for Genetics problems."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 93b9f489a..e5b7714eb 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -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.
+
 """Utilities for data generators."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/generator_utils_test.py b/tensor2tensor/data_generators/generator_utils_test.py
index b053e508c..4249d60aa 100644
--- a/tensor2tensor/data_generators/generator_utils_test.py
+++ b/tensor2tensor/data_generators/generator_utils_test.py
@@ -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.
+
 """Generator utilities test."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index 75bafcedf..e20edb59a 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -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.
+
 """Google robot pushing dataset.
 
 Unsupervised Learning for Physical Interaction through Video Prediction
diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index afb18f402..c0241870d 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -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.
+
 """Data generators for Gym environments."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
index 007fdcb14..c904cdb39 100644
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ b/tensor2tensor/data_generators/gym_problems_specs.py
@@ -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.
+
 """Definitions of data generators for gym problems."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
index 3ee824a48..326e75d45 100644
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ b/tensor2tensor/data_generators/gym_problems_test.py
@@ -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.
+
 """Gym generators tests."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/gym_utils.py b/tensor2tensor/data_generators/gym_utils.py
index e821bc38b..18c3c0279 100644
--- a/tensor2tensor/data_generators/gym_utils.py
+++ b/tensor2tensor/data_generators/gym_utils.py
@@ -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.
+
 """Utilities for openai gym."""
 
 from collections import deque
diff --git a/tensor2tensor/data_generators/ice_parsing.py b/tensor2tensor/data_generators/ice_parsing.py
index 2561f28bc..59eacb517 100644
--- a/tensor2tensor/data_generators/ice_parsing.py
+++ b/tensor2tensor/data_generators/ice_parsing.py
@@ -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.
+
 """This module implements the ice_parsing_* problems."""
 
 # These parse plain text into flattened parse trees and POS tags.
diff --git a/tensor2tensor/data_generators/image_lsun.py b/tensor2tensor/data_generators/image_lsun.py
index 373f8aa98..8f41dd1ce 100644
--- a/tensor2tensor/data_generators/image_lsun.py
+++ b/tensor2tensor/data_generators/image_lsun.py
@@ -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.
+
 """LSUN datasets (bedrooms only for now)."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 3b59060e5..551d21044 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -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.
+
 """Base classes and utilities for image datasets."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/image_utils_test.py b/tensor2tensor/data_generators/image_utils_test.py
index 1dca736e5..326947e46 100644
--- a/tensor2tensor/data_generators/image_utils_test.py
+++ b/tensor2tensor/data_generators/image_utils_test.py
@@ -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.
+
 """image_utils test."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index e9db5011d..5c1dfdd31 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -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.
+
 """ImageNet."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/imagenet_test.py b/tensor2tensor/data_generators/imagenet_test.py
index eeb70f0a0..3eb671a06 100644
--- a/tensor2tensor/data_generators/imagenet_test.py
+++ b/tensor2tensor/data_generators/imagenet_test.py
@@ -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.
+
 """Tests for ImageNet."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/imdb.py b/tensor2tensor/data_generators/imdb.py
index 2f02a471b..0e8b6e28a 100644
--- a/tensor2tensor/data_generators/imdb.py
+++ b/tensor2tensor/data_generators/imdb.py
@@ -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.
+
 """IMDB Sentiment Classification Problem."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/inspect_tfrecord.py b/tensor2tensor/data_generators/inspect_tfrecord.py
index 7b29015c9..a245edb9e 100644
--- a/tensor2tensor/data_generators/inspect_tfrecord.py
+++ b/tensor2tensor/data_generators/inspect_tfrecord.py
@@ -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.
+
 r"""Inspect a TFRecord file of tensorflow.Example and show tokenizations.
 
 python data_generators/inspect_tfrecord.py \
diff --git a/tensor2tensor/data_generators/lambada.py b/tensor2tensor/data_generators/lambada.py
index 688db95ca..de0b03833 100644
--- a/tensor2tensor/data_generators/lambada.py
+++ b/tensor2tensor/data_generators/lambada.py
@@ -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.
+
 """Data generators for LAMBADA data-sets.
 
 
diff --git a/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index 37adc4aa6..11de4b65a 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -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.
+
 """Librispeech dataset."""
 
 import os
diff --git a/tensor2tensor/data_generators/lm1b.py b/tensor2tensor/data_generators/lm1b.py
index 67f0cc695..dc120df83 100644
--- a/tensor2tensor/data_generators/lm1b.py
+++ b/tensor2tensor/data_generators/lm1b.py
@@ -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.
+
 """Data generators for LM1B data-set."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/lm1b_imdb.py b/tensor2tensor/data_generators/lm1b_imdb.py
index b4b8ba9aa..05309c6f3 100644
--- a/tensor2tensor/data_generators/lm1b_imdb.py
+++ b/tensor2tensor/data_generators/lm1b_imdb.py
@@ -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.
+
 """Data generators for LM1B and IMDb combined data-set."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/lm1b_mnli.py b/tensor2tensor/data_generators/lm1b_mnli.py
index a0835723f..66530dc36 100644
--- a/tensor2tensor/data_generators/lm1b_mnli.py
+++ b/tensor2tensor/data_generators/lm1b_mnli.py
@@ -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.
+
 """Data generators for LM1B and MNLI combined datasets."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/mnist.py b/tensor2tensor/data_generators/mnist.py
index 7f83a8680..322aecb60 100644
--- a/tensor2tensor/data_generators/mnist.py
+++ b/tensor2tensor/data_generators/mnist.py
@@ -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.
+
 """MNIST."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/mrpc.py b/tensor2tensor/data_generators/mrpc.py
index c56bac32c..14822428d 100644
--- a/tensor2tensor/data_generators/mrpc.py
+++ b/tensor2tensor/data_generators/mrpc.py
@@ -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.
+
 """Data generators for the MSR Paraphrase Corpus."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index b103aa135..d57bd616a 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -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.
+
 """MS COCO."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/mscoco_test.py b/tensor2tensor/data_generators/mscoco_test.py
index bb76041c7..9c5564a71 100644
--- a/tensor2tensor/data_generators/mscoco_test.py
+++ b/tensor2tensor/data_generators/mscoco_test.py
@@ -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.
+
 """Tests for MS COCO."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 5f28d23fd..37f36c364 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -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.
+
 """Base class for combining multiple problems for multitask learning."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/multi_problem_test.py b/tensor2tensor/data_generators/multi_problem_test.py
index e1f38b8d4..7b0708a16 100644
--- a/tensor2tensor/data_generators/multi_problem_test.py
+++ b/tensor2tensor/data_generators/multi_problem_test.py
@@ -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.
+
 """Tests for MultiProblem."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index 2b4dddfcb..0e52dd455 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -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.
+
 """Data generators for MultiNLI."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/ocr.py b/tensor2tensor/data_generators/ocr.py
index cff224bd2..39cfa56a8 100644
--- a/tensor2tensor/data_generators/ocr.py
+++ b/tensor2tensor/data_generators/ocr.py
@@ -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.
+
 """OCR."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/paraphrase_ms_coco.py b/tensor2tensor/data_generators/paraphrase_ms_coco.py
index ed5c4dd02..4c459a324 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco.py
@@ -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.
+
 """Base classes for paraphrase generation problems."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/paraphrase_ms_coco_test.py b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
index 439e2ec6c..763ad7af5 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
@@ -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.
+
 """Tests for tensor2tensor.data_generators.paraphrase_ms_coco."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/pointer_generator_word.py b/tensor2tensor/data_generators/pointer_generator_word.py
index b2759113c..0180a9ad4 100644
--- a/tensor2tensor/data_generators/pointer_generator_word.py
+++ b/tensor2tensor/data_generators/pointer_generator_word.py
@@ -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.
+
 """Data generator for pointer-generator for word transformer."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index da2c714c7..f93699f95 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -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.
+
 """Base class for problem/dataset definitions."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/problem_hparams.py b/tensor2tensor/data_generators/problem_hparams.py
index 51459f1f5..45bc312be 100644
--- a/tensor2tensor/data_generators/problem_hparams.py
+++ b/tensor2tensor/data_generators/problem_hparams.py
@@ -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.
+
 """Hyperparameters defining different problems.
 
 """
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 4b6ccd4c1..a98cca80c 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -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.
+
 """Test for common problem functionalities."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/program_search.py b/tensor2tensor/data_generators/program_search.py
index 775fa7c39..b0fe553ca 100644
--- a/tensor2tensor/data_generators/program_search.py
+++ b/tensor2tensor/data_generators/program_search.py
@@ -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.
+
 """Program Search Problems."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/program_search_test.py b/tensor2tensor/data_generators/program_search_test.py
index 84f4a1e03..e73a59a6b 100644
--- a/tensor2tensor/data_generators/program_search_test.py
+++ b/tensor2tensor/data_generators/program_search_test.py
@@ -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.
+
 """Tests for tensor2tensor.data_generators.program_search."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/ptb.py b/tensor2tensor/data_generators/ptb.py
index c40c135cd..5684b21eb 100644
--- a/tensor2tensor/data_generators/ptb.py
+++ b/tensor2tensor/data_generators/ptb.py
@@ -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.
+
 """Data generators for PTB data-sets."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/qnli.py b/tensor2tensor/data_generators/qnli.py
index 310aadcd5..b59db970f 100644
--- a/tensor2tensor/data_generators/qnli.py
+++ b/tensor2tensor/data_generators/qnli.py
@@ -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.
+
 """Data generators for the Question-Answering NLI dataset."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/quora_qpairs.py b/tensor2tensor/data_generators/quora_qpairs.py
index 7e189bfc0..5960c2488 100644
--- a/tensor2tensor/data_generators/quora_qpairs.py
+++ b/tensor2tensor/data_generators/quora_qpairs.py
@@ -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.
+
 """Data generators for the Quora Question Pairs dataset."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/rte.py b/tensor2tensor/data_generators/rte.py
index 04e322ca1..af7fa41e9 100644
--- a/tensor2tensor/data_generators/rte.py
+++ b/tensor2tensor/data_generators/rte.py
@@ -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.
+
 """Data generators for the Recognizing Textual Entailment dataset."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/scitail.py b/tensor2tensor/data_generators/scitail.py
index 789246481..90df97ccb 100644
--- a/tensor2tensor/data_generators/scitail.py
+++ b/tensor2tensor/data_generators/scitail.py
@@ -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.
+
 """Data generators for SciTail."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/snli.py b/tensor2tensor/data_generators/snli.py
index 3cc8c2b5d..446562265 100644
--- a/tensor2tensor/data_generators/snli.py
+++ b/tensor2tensor/data_generators/snli.py
@@ -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.
+
 """Data generators for the SNLI data-set."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 5ffdc8bf4..fc662350c 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -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.
+
 """Common classes for automatic speech recognition (ASR) datasets.
 
 The audio import uses sox to generate normalized waveforms, please install
diff --git a/tensor2tensor/data_generators/squad.py b/tensor2tensor/data_generators/squad.py
index 7d82df547..a37a8a336 100644
--- a/tensor2tensor/data_generators/squad.py
+++ b/tensor2tensor/data_generators/squad.py
@@ -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.
+
 """Data generators for SquaAD (https://rajpurkar.github.io/SQuAD-explorer/).
 """
 
diff --git a/tensor2tensor/data_generators/sst_binary.py b/tensor2tensor/data_generators/sst_binary.py
index 1fceacdee..a8a391c95 100644
--- a/tensor2tensor/data_generators/sst_binary.py
+++ b/tensor2tensor/data_generators/sst_binary.py
@@ -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.
+
 """Stanford Sentiment Treebank Binary Classification Problem."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/stanford_nli.py b/tensor2tensor/data_generators/stanford_nli.py
index ad244303b..2f9421d08 100644
--- a/tensor2tensor/data_generators/stanford_nli.py
+++ b/tensor2tensor/data_generators/stanford_nli.py
@@ -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.
+
 """Data generators for StanfordNLI."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/style_transfer.py b/tensor2tensor/data_generators/style_transfer.py
index b3096c7a5..d7a046564 100644
--- a/tensor2tensor/data_generators/style_transfer.py
+++ b/tensor2tensor/data_generators/style_transfer.py
@@ -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.
+
 """Base classes for text-based language style transfer problems.
 
 * StyleTransferProblem: abstract class for style transfer problems.
diff --git a/tensor2tensor/data_generators/style_transfer_test.py b/tensor2tensor/data_generators/style_transfer_test.py
index 6e4e1aa94..fe9618ae2 100644
--- a/tensor2tensor/data_generators/style_transfer_test.py
+++ b/tensor2tensor/data_generators/style_transfer_test.py
@@ -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.
+
 """Tests for tensor2tensor.data_generators.style_transfer."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/subject_verb_agreement.py b/tensor2tensor/data_generators/subject_verb_agreement.py
index 3b6604215..18b802d8c 100644
--- a/tensor2tensor/data_generators/subject_verb_agreement.py
+++ b/tensor2tensor/data_generators/subject_verb_agreement.py
@@ -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.
+
 """Data generators for subject-verb agreement dataset.
 
 https://arxiv.org/pdf/1611.01368.pdf
diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 42d136a04..a00403c14 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -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.
+
 """Encoders for text data.
 
 * TextEncoder: base class
diff --git a/tensor2tensor/data_generators/text_encoder_build_subword.py b/tensor2tensor/data_generators/text_encoder_build_subword.py
index ddcd524b2..89c6b9516 100644
--- a/tensor2tensor/data_generators/text_encoder_build_subword.py
+++ b/tensor2tensor/data_generators/text_encoder_build_subword.py
@@ -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.
+
 r"""Program to build a SubwordTextEncoder.
 
 The flags --min_count and --corpus_max_lines will affect the size of the
diff --git a/tensor2tensor/data_generators/text_encoder_test.py b/tensor2tensor/data_generators/text_encoder_test.py
index 26c5ce65a..ba5bbf998 100644
--- a/tensor2tensor/data_generators/text_encoder_test.py
+++ b/tensor2tensor/data_generators/text_encoder_test.py
@@ -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.
+
 """Tests for tensor2tensor.data_generators.text_encoder."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index e5e6a3159..80249b72a 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -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.
+
 """Base classes for text-based Problems.
 
 * Text2TextProblem: input=text, target=text.
diff --git a/tensor2tensor/data_generators/text_problems_test.py b/tensor2tensor/data_generators/text_problems_test.py
index a42342404..31d0992f8 100644
--- a/tensor2tensor/data_generators/text_problems_test.py
+++ b/tensor2tensor/data_generators/text_problems_test.py
@@ -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.
+
 """Text problems test."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index 71a9fb5e0..2ee7ba007 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -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.
+
 """Multi time series forecasting problem."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/timeseries_data_generator.py b/tensor2tensor/data_generators/timeseries_data_generator.py
index db21dc529..13fe70af0 100644
--- a/tensor2tensor/data_generators/timeseries_data_generator.py
+++ b/tensor2tensor/data_generators/timeseries_data_generator.py
@@ -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.
+
 """Data generator for the timeseries problem."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/timeseries_data_generator_test.py b/tensor2tensor/data_generators/timeseries_data_generator_test.py
index 85eb2a2f3..6d4c208c9 100644
--- a/tensor2tensor/data_generators/timeseries_data_generator_test.py
+++ b/tensor2tensor/data_generators/timeseries_data_generator_test.py
@@ -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.
+
 """Timeseries data generator tests."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index 038ddedcc..bae4b0e0e 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -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.
+
 """Timeseries generators tests."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/tokenizer.py b/tensor2tensor/data_generators/tokenizer.py
index c84642733..6a05d5e05 100644
--- a/tensor2tensor/data_generators/tokenizer.py
+++ b/tensor2tensor/data_generators/tokenizer.py
@@ -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.
+
 """A simple invertible tokenizer.
 
 Converts from a unicode string to a list of tokens
diff --git a/tensor2tensor/data_generators/tokenizer_test.py b/tensor2tensor/data_generators/tokenizer_test.py
index 01c192327..ca9c3b8e5 100644
--- a/tensor2tensor/data_generators/tokenizer_test.py
+++ b/tensor2tensor/data_generators/tokenizer_test.py
@@ -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.
+
 # coding=utf-8
 """Tests for tensor2tensor.data_generators.tokenizer."""
 
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 478ed80c2..23f2ae76b 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -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.
+
 """Data generators for translation data-sets."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/translate_encs.py b/tensor2tensor/data_generators/translate_encs.py
index 9224e0d3d..d5d6a3f99 100644
--- a/tensor2tensor/data_generators/translate_encs.py
+++ b/tensor2tensor/data_generators/translate_encs.py
@@ -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.
+
 """Data generators for translation data-sets."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index 73b1c9ddc..c5fa94aff 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -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.
+
 """Data generators for translation data-sets."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/translate_enet.py b/tensor2tensor/data_generators/translate_enet.py
index 33150a8e4..4c3b4e929 100644
--- a/tensor2tensor/data_generators/translate_enet.py
+++ b/tensor2tensor/data_generators/translate_enet.py
@@ -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.
+
 """Data generators for En-Et translation."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index 028ea0ac5..6f5ad163d 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -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.
+
 """Data generators for translation data-sets."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/translate_enid.py b/tensor2tensor/data_generators/translate_enid.py
index 0331b664b..aa070637e 100644
--- a/tensor2tensor/data_generators/translate_enid.py
+++ b/tensor2tensor/data_generators/translate_enid.py
@@ -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.
+
 """Data generators for En-Id translation."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/translate_enmk.py b/tensor2tensor/data_generators/translate_enmk.py
index fc832c3e4..6caa1ee9d 100644
--- a/tensor2tensor/data_generators/translate_enmk.py
+++ b/tensor2tensor/data_generators/translate_enmk.py
@@ -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.
+
 """Data generators for translation data-sets."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/translate_envi.py b/tensor2tensor/data_generators/translate_envi.py
index 97023cbd2..b3a1a7885 100644
--- a/tensor2tensor/data_generators/translate_envi.py
+++ b/tensor2tensor/data_generators/translate_envi.py
@@ -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.
+
 """Data generators for En-Vi translation."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index db030120d..90fd0872a 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -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.
+
 """Data generators for translation data-sets."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/translate_test.py b/tensor2tensor/data_generators/translate_test.py
index b34cff1b7..a489b10d3 100644
--- a/tensor2tensor/data_generators/translate_test.py
+++ b/tensor2tensor/data_generators/translate_test.py
@@ -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.
+
 """Translate generators test."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/twentybn.py b/tensor2tensor/data_generators/twentybn.py
index 64c69d169..b28162044 100644
--- a/tensor2tensor/data_generators/twentybn.py
+++ b/tensor2tensor/data_generators/twentybn.py
@@ -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.
+
 """Data generator for twenty bn video data-set."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 3a3e70db3..9c33aa14a 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -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.
+
 """Data generators for video problems with artificially generated frames."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 47484916a..74832a146 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -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.
+
 """Base classes and utilities for video datasets."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index 7c807bd88..7be6a6a7c 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -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.
+
 """video_utils test."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
index 27f025c1d..e00f69c1b 100644
--- a/tensor2tensor/data_generators/vqa.py
+++ b/tensor2tensor/data_generators/vqa.py
@@ -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.
+
 """Data generators for VQA data sets."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/vqa_utils.py b/tensor2tensor/data_generators/vqa_utils.py
index 1599879b2..a1a6bc6e2 100644
--- a/tensor2tensor/data_generators/vqa_utils.py
+++ b/tensor2tensor/data_generators/vqa_utils.py
@@ -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.
+
 """Utilities for VQA data sets."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/wiki.py b/tensor2tensor/data_generators/wiki.py
index 3b76457f2..31973fdeb 100644
--- a/tensor2tensor/data_generators/wiki.py
+++ b/tensor2tensor/data_generators/wiki.py
@@ -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.
+
 """Data generator for Wikipedia title to article dataset."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/wikisum/__init__.py b/tensor2tensor/data_generators/wikisum/__init__.py
index dba7ece95..4bd418a74 100644
--- a/tensor2tensor/data_generators/wikisum/__init__.py
+++ b/tensor2tensor/data_generators/wikisum/__init__.py
@@ -12,3 +12,4 @@
 # 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.
+
diff --git a/tensor2tensor/data_generators/wikisum/generate_vocab.py b/tensor2tensor/data_generators/wikisum/generate_vocab.py
index 33fbdbb73..0231e0e80 100644
--- a/tensor2tensor/data_generators/wikisum/generate_vocab.py
+++ b/tensor2tensor/data_generators/wikisum/generate_vocab.py
@@ -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.
+
 """Generate vocab from references and wikis."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py b/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
index 5532f9a00..52cec62a6 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
@@ -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.
+
 """Extract references from CommonCrawl files."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/wikisum/get_references_web.py b/tensor2tensor/data_generators/wikisum/get_references_web.py
index 44b4ac94e..2a8d9699d 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web.py
@@ -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.
+
 # pylint: disable=line-too-long
 r"""Fetch reference URLs from all groups for a single shard id.
 
diff --git a/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py b/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
index d6678a421..d27686a4b 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
@@ -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.
+
 """Fetch reference URLs for a single group_id within a single shard_id.
 
 See get_references_web.py to fetch URLs for all groups in within a single
diff --git a/tensor2tensor/data_generators/wikisum/html.py b/tensor2tensor/data_generators/wikisum/html.py
index 32c2b82d5..e8483ae8b 100644
--- a/tensor2tensor/data_generators/wikisum/html.py
+++ b/tensor2tensor/data_generators/wikisum/html.py
@@ -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.
+
 """Utils to parse HTML content into plaintext."""
 
 import bs4
diff --git a/tensor2tensor/data_generators/wikisum/parallel_launch.py b/tensor2tensor/data_generators/wikisum/parallel_launch.py
index 811466745..3332a8978 100644
--- a/tensor2tensor/data_generators/wikisum/parallel_launch.py
+++ b/tensor2tensor/data_generators/wikisum/parallel_launch.py
@@ -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.
+
 # pylint: disable=line-too-long
 r"""Launch a script in parallel on GCP.
 
diff --git a/tensor2tensor/data_generators/wikisum/produce_examples.py b/tensor2tensor/data_generators/wikisum/produce_examples.py
index 7ce3a3508..95f736978 100644
--- a/tensor2tensor/data_generators/wikisum/produce_examples.py
+++ b/tensor2tensor/data_generators/wikisum/produce_examples.py
@@ -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.
+
 """Produce examples given a vocab, wikis, references, and dataset URLs."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/wikisum/utils.py b/tensor2tensor/data_generators/wikisum/utils.py
index c2c6f8c88..fcaf0399a 100644
--- a/tensor2tensor/data_generators/wikisum/utils.py
+++ b/tensor2tensor/data_generators/wikisum/utils.py
@@ -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.
+
 """Wikisum data generation utilities."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/wikisum/utils_test.py b/tensor2tensor/data_generators/wikisum/utils_test.py
index 2e8a5bcdb..a20afbf4c 100644
--- a/tensor2tensor/data_generators/wikisum/utils_test.py
+++ b/tensor2tensor/data_generators/wikisum/utils_test.py
@@ -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.
+
 """Tests for tensor2tensor.data_generators.wikisum.utils."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/wikisum/validate_data.py b/tensor2tensor/data_generators/wikisum/validate_data.py
index bb109e80d..618759a74 100644
--- a/tensor2tensor/data_generators/wikisum/validate_data.py
+++ b/tensor2tensor/data_generators/wikisum/validate_data.py
@@ -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.
+
 """Aggregate stats from produce_examples."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/wikisum/wikisum.py b/tensor2tensor/data_generators/wikisum/wikisum.py
index b7ec85231..f1affcb94 100644
--- a/tensor2tensor/data_generators/wikisum/wikisum.py
+++ b/tensor2tensor/data_generators/wikisum/wikisum.py
@@ -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.
+
 """Wikipedia Summarization Problems."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/wikitext103.py b/tensor2tensor/data_generators/wikitext103.py
index 7712a717b..24ebc2763 100644
--- a/tensor2tensor/data_generators/wikitext103.py
+++ b/tensor2tensor/data_generators/wikitext103.py
@@ -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.
+
 """Data generators for wikitext-103.
 
 Wikitext-103: Long term dependency language modeling dataset
diff --git a/tensor2tensor/data_generators/wnli.py b/tensor2tensor/data_generators/wnli.py
index a00ff0754..cd4de046d 100644
--- a/tensor2tensor/data_generators/wnli.py
+++ b/tensor2tensor/data_generators/wnli.py
@@ -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.
+
 """Data generators for the Winograd NLI dataset."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/data_generators/wsj_parsing.py b/tensor2tensor/data_generators/wsj_parsing.py
index 82a810a2f..b9adf2a7e 100644
--- a/tensor2tensor/data_generators/wsj_parsing.py
+++ b/tensor2tensor/data_generators/wsj_parsing.py
@@ -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.
+
 """Data generators for parsing data-sets."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/insights/__init__.py b/tensor2tensor/insights/__init__.py
index dba7ece95..4bd418a74 100644
--- a/tensor2tensor/insights/__init__.py
+++ b/tensor2tensor/insights/__init__.py
@@ -12,3 +12,4 @@
 # 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.
+
diff --git a/tensor2tensor/insights/graph.py b/tensor2tensor/insights/graph.py
index 896f4d8ed..17e18ea3c 100644
--- a/tensor2tensor/insights/graph.py
+++ b/tensor2tensor/insights/graph.py
@@ -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.
+
 """Graph representation for building decoding graph visualizations."""
 
 
diff --git a/tensor2tensor/insights/query_processor.py b/tensor2tensor/insights/query_processor.py
index c615f4868..d6703af44 100644
--- a/tensor2tensor/insights/query_processor.py
+++ b/tensor2tensor/insights/query_processor.py
@@ -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.
+
 """A base class for all query processing classes."""
 
 
diff --git a/tensor2tensor/insights/server.py b/tensor2tensor/insights/server.py
index 54421db58..e61c465ff 100644
--- a/tensor2tensor/insights/server.py
+++ b/tensor2tensor/insights/server.py
@@ -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.
+
 """A GUnicorn + Flask Debug Frontend for Transformer models."""
 
 import json
diff --git a/tensor2tensor/insights/transformer_model.py b/tensor2tensor/insights/transformer_model.py
index 3a4e5101d..da8cf5fe3 100644
--- a/tensor2tensor/insights/transformer_model.py
+++ b/tensor2tensor/insights/transformer_model.py
@@ -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.
+
 """A QueryProcessor using the Transformer framework."""
 
 from collections import deque
diff --git a/tensor2tensor/layers/__init__.py b/tensor2tensor/layers/__init__.py
index dba7ece95..4bd418a74 100644
--- a/tensor2tensor/layers/__init__.py
+++ b/tensor2tensor/layers/__init__.py
@@ -12,3 +12,4 @@
 # 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.
+
diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index e3daa1cc6..42114f88e 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -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.
+
 """Utilities for attention."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 3015535e7..eb96ea2c9 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -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.
+
 """Tests for common attention."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/layers/common_audio.py b/tensor2tensor/layers/common_audio.py
index 30cc9ae45..14432576f 100644
--- a/tensor2tensor/layers/common_audio.py
+++ b/tensor2tensor/layers/common_audio.py
@@ -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.
+
 """Utils for audio."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 02edfd231..b9dd752c7 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -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.
+
 """Hyperparameters and ranges common to multiple models."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 2cd76ee5a..aff70120a 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -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.
+
 """Utils for attention mechanism for images."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index 6d57e7413..e87a2e4d8 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -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.
+
 """Tests for common image attention utilities."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index a7f3c99a8..b928dc407 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -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.
+
 """Layers common to multiple models."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index 01247db34..a15f37302 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -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.
+
 """Tests for common layers."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/layers/common_message_passing_attention.py b/tensor2tensor/layers/common_message_passing_attention.py
index f9464bb82..b744ffd07 100644
--- a/tensor2tensor/layers/common_message_passing_attention.py
+++ b/tensor2tensor/layers/common_message_passing_attention.py
@@ -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.
+
 """Utilities for attention."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 64e674a2b..0ae47e54c 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -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.
+
 """Utilities for video."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index 5a28a21a3..3ab72ce1a 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -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.
+
 """Tests for video utils."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index f1aace305..d443cb29b 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -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.
+
 """Discretization bottlenecks used to train discrete latent variables."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index 7b64015db..689a53e27 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -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.
+
 """Tests for discretization."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index dd49a3546..72e2404be 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -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.
+
 """Utils for latent variable models."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
index f7d00c02b..884db383c 100644
--- a/tensor2tensor/layers/latent_layers_test.py
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -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.
+
 """Tests for layers in latent variable models."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 1325adf71..c808ab1c8 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -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.
+
 """Modalities define the bottom and top of the model (not the body)."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 875b66918..e10a9d10c 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -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.
+
 """Tests for Modalities."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index 9477708f3..f45289df2 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -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.
+
 """Clean discrete bottleneck as in https://arxiv.org/abs/1805.11063."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/layers/vqa_layers.py b/tensor2tensor/layers/vqa_layers.py
index e3ec72fe4..537e7ac8f 100644
--- a/tensor2tensor/layers/vqa_layers.py
+++ b/tensor2tensor/layers/vqa_layers.py
@@ -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.
+
 """Some customization of common_attention."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/mesh_tensorflow/OWNERS b/tensor2tensor/mesh_tensorflow/OWNERS
new file mode 100644
index 000000000..681d09a90
--- /dev/null
+++ b/tensor2tensor/mesh_tensorflow/OWNERS
@@ -0,0 +1,2 @@
+nikip
+noam
diff --git a/tensor2tensor/mesh_tensorflow/__init__.py b/tensor2tensor/mesh_tensorflow/__init__.py
index dba7ece95..4bd418a74 100644
--- a/tensor2tensor/mesh_tensorflow/__init__.py
+++ b/tensor2tensor/mesh_tensorflow/__init__.py
@@ -12,3 +12,4 @@
 # 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.
+
diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index d51f4dba7..0508aacc7 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -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.
+
 """Mesh-TensorFlow."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
index 10e36fe2a..8dd0ee4f0 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
@@ -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.
+
 """Tests for Mesh TensorFlow."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/mesh_tensorflow/mnist.py b/tensor2tensor/mesh_tensorflow/mnist.py
index d02fa164b..5dced5c6d 100644
--- a/tensor2tensor/mesh_tensorflow/mnist.py
+++ b/tensor2tensor/mesh_tensorflow/mnist.py
@@ -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.
+
 """Mnist using mesh-tensorflow and tf.Estimator.
 
 This is an illustration of mesh-tensorflow, not a good model.
diff --git a/tensor2tensor/mesh_tensorflow/mnist_dataset.py b/tensor2tensor/mesh_tensorflow/mnist_dataset.py
index 15db689b6..3be1e382d 100644
--- a/tensor2tensor/mesh_tensorflow/mnist_dataset.py
+++ b/tensor2tensor/mesh_tensorflow/mnist_dataset.py
@@ -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.
+
 #  Copyright 2018 The TensorFlow Authors. All Rights Reserved.
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
diff --git a/tensor2tensor/mesh_tensorflow/mtf_beam_search.py b/tensor2tensor/mesh_tensorflow/mtf_beam_search.py
index 3fb8592ba..ce1298676 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_beam_search.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_beam_search.py
@@ -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.
+
 """Implementation of beam search with penalties."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
index 8934a6b75..7bc999d4a 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
@@ -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.
+
 """Image Transformer model with model and data parallelism using MTF.
 
 Integration of Mesh tensorflow with Image Transformer to do model parallelism.
diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py b/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
index 53e8cff7a..1411ebd36 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
@@ -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.
+
 """Tests for Image Transformer on Mesh TensorFlow."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers.py b/tensor2tensor/mesh_tensorflow/mtf_layers.py
index 8b4e8bd04..310eefb66 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers.py
@@ -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.
+
 """Layers for mesh tensorflow."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers_test.py b/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
index f3c3501e0..ae6384c80 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
@@ -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.
+
 """Tests for Mesh TensorFlow layers."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/mesh_tensorflow/mtf_model.py b/tensor2tensor/mesh_tensorflow/mtf_model.py
index 8d50787a7..829287148 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_model.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_model.py
@@ -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.
+
 """Mesh-Tensorflow Model in tensor2tensor."""
 
 
diff --git a/tensor2tensor/mesh_tensorflow/mtf_optimize.py b/tensor2tensor/mesh_tensorflow/mtf_optimize.py
index 0dcd0a79e..c57433808 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_optimize.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_optimize.py
@@ -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.
+
 """Mesh-Tensorflow Optimizers."""
 
 
diff --git a/tensor2tensor/mesh_tensorflow/mtf_resnet.py b/tensor2tensor/mesh_tensorflow/mtf_resnet.py
index 5de8ee028..ed7320820 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_resnet.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_resnet.py
@@ -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.
+
 """ResNet model with model and data parallelism using MTF.
 
 Integration of Mesh tensorflow with ResNet to do model parallelism.
diff --git a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
index e540dc666..911badd1d 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
@@ -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.
+
 """A toy model using mesh-tensrflow."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
index d3d165bf6..576f9061d 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
@@ -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.
+
 """Transformer model."""
 
 
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py b/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py
index 2d0695a5e..4b00b978e 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py
@@ -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.
+
 """Tests for Transformer on Mesh TensorFlow."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/mesh_tensorflow/mtf_utils.py b/tensor2tensor/mesh_tensorflow/mtf_utils.py
index 108d9a530..cd9f5cb18 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_utils.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_utils.py
@@ -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.
+
 """Common utilities for mesh tensorflow."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/mesh_tensorflow/mtf_utils_test.py b/tensor2tensor/mesh_tensorflow/mtf_utils_test.py
index a848d09c4..d0547e924 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_utils_test.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_utils_test.py
@@ -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.
+
 """Tests for mtf_utils.py."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
index 0d0c87c0c..b1085c4a1 100644
--- a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
@@ -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.
+
 """Placement Mesh Implementation (for CPU/GPU clusters)."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/mesh_tensorflow/research/__init__.py b/tensor2tensor/mesh_tensorflow/research/__init__.py
index dba7ece95..4bd418a74 100644
--- a/tensor2tensor/mesh_tensorflow/research/__init__.py
+++ b/tensor2tensor/mesh_tensorflow/research/__init__.py
@@ -12,3 +12,4 @@
 # 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.
+
diff --git a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
index fc643d314..12e612bee 100644
--- a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
@@ -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.
+
 """Experiments with mixture-of-experts architectures."""
 
 
diff --git a/tensor2tensor/mesh_tensorflow/research/moe.py b/tensor2tensor/mesh_tensorflow/research/moe.py
index 88e9c27e1..73542cd8e 100644
--- a/tensor2tensor/mesh_tensorflow/research/moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/moe.py
@@ -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.
+
 """Mixture-of-experts code.
 
 Interfaces and algorithms are under development and subject to rapid change
diff --git a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
index 9e080fcf7..9f723b99f 100644
--- a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
@@ -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.
+
 """SIMD Mesh implementation (for TPU/XLA)."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/mesh_tensorflow/tpu_variables.py b/tensor2tensor/mesh_tensorflow/tpu_variables.py
index 7da863589..bb852a569 100644
--- a/tensor2tensor/mesh_tensorflow/tpu_variables.py
+++ b/tensor2tensor/mesh_tensorflow/tpu_variables.py
@@ -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.
+
 """Distributed variable implementation for TPUs."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/OWNERS b/tensor2tensor/models/OWNERS
new file mode 100644
index 000000000..4c18f708e
--- /dev/null
+++ b/tensor2tensor/models/OWNERS
@@ -0,0 +1,6 @@
+avaswani
+dumitru
+mbz
+nikip
+noam
+uszkoreit
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 6d2dbbf3e..34409a8e4 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -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.
+
 """Models defined in T2T. Imports here force registration."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/models/basic.py b/tensor2tensor/models/basic.py
index 109190951..a0a46ff58 100644
--- a/tensor2tensor/models/basic.py
+++ b/tensor2tensor/models/basic.py
@@ -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.
+
 """Basic models for testing simple tasks."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/basic_test.py b/tensor2tensor/models/basic_test.py
index c50b103ca..829c802bf 100644
--- a/tensor2tensor/models/basic_test.py
+++ b/tensor2tensor/models/basic_test.py
@@ -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.
+
 """Basic nets tests."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/bytenet.py b/tensor2tensor/models/bytenet.py
index a81bbd106..3eec2afda 100644
--- a/tensor2tensor/models/bytenet.py
+++ b/tensor2tensor/models/bytenet.py
@@ -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.
+
 """ByteNet."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/bytenet_test.py b/tensor2tensor/models/bytenet_test.py
index 1021ca432..b87d75d5d 100644
--- a/tensor2tensor/models/bytenet_test.py
+++ b/tensor2tensor/models/bytenet_test.py
@@ -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.
+
 """ByteNet tests."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/distillation.py b/tensor2tensor/models/distillation.py
index 839501a95..ed08e076c 100644
--- a/tensor2tensor/models/distillation.py
+++ b/tensor2tensor/models/distillation.py
@@ -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.
+
 """Traditional Student-Teacher Distillation."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index d4145d04f..5777f1833 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -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.
+
 """image generation with transformer (attention).
 
 encoder: [Self-Attention, Feed-forward] x n
diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 1ef6ede1a..5b2299a00 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -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.
+
 """image generation with transformer (attention).
 
 encoder: [Self-Attention, Feed-forward] x n
diff --git a/tensor2tensor/models/image_transformer_2d_test.py b/tensor2tensor/models/image_transformer_2d_test.py
index 3f52df180..d5327de78 100644
--- a/tensor2tensor/models/image_transformer_2d_test.py
+++ b/tensor2tensor/models/image_transformer_2d_test.py
@@ -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.
+
 """Tests for Transformer."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/image_transformer_test.py b/tensor2tensor/models/image_transformer_test.py
index 8257a0fe9..cf8417db4 100644
--- a/tensor2tensor/models/image_transformer_test.py
+++ b/tensor2tensor/models/image_transformer_test.py
@@ -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.
+
 """Tests for Transformer."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index c8634ced2..2f657970c 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -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.
+
 """RNN LSTM models."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/lstm_test.py b/tensor2tensor/models/lstm_test.py
index a75e95cf5..45c188816 100644
--- a/tensor2tensor/models/lstm_test.py
+++ b/tensor2tensor/models/lstm_test.py
@@ -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.
+
 """LSTMSeq2Seq models tests."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/neural_gpu.py b/tensor2tensor/models/neural_gpu.py
index 44b4473af..d7369d247 100644
--- a/tensor2tensor/models/neural_gpu.py
+++ b/tensor2tensor/models/neural_gpu.py
@@ -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.
+
 """The Neural GPU model and its variants."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/neural_gpu_test.py b/tensor2tensor/models/neural_gpu_test.py
index 857e0e5ea..cf8731d90 100644
--- a/tensor2tensor/models/neural_gpu_test.py
+++ b/tensor2tensor/models/neural_gpu_test.py
@@ -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.
+
 """Tests for Neural GPU."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/__init__.py b/tensor2tensor/models/research/__init__.py
index dba7ece95..4bd418a74 100644
--- a/tensor2tensor/models/research/__init__.py
+++ b/tensor2tensor/models/research/__init__.py
@@ -12,3 +12,4 @@
 # 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.
+
diff --git a/tensor2tensor/models/research/adafactor_experiments.py b/tensor2tensor/models/research/adafactor_experiments.py
index e9753c6c8..d7d3d4e2c 100644
--- a/tensor2tensor/models/research/adafactor_experiments.py
+++ b/tensor2tensor/models/research/adafactor_experiments.py
@@ -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.
+
 """Experiments with Adafactor.
 """
 
diff --git a/tensor2tensor/models/research/aligned.py b/tensor2tensor/models/research/aligned.py
index 18a576b5d..60afff5e5 100644
--- a/tensor2tensor/models/research/aligned.py
+++ b/tensor2tensor/models/research/aligned.py
@@ -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.
+
 """Single stack of transformations with no masking.
 
 Produces output aligned with inputs.
diff --git a/tensor2tensor/models/research/attention_lm.py b/tensor2tensor/models/research/attention_lm.py
index d3de21044..3a03efafb 100644
--- a/tensor2tensor/models/research/attention_lm.py
+++ b/tensor2tensor/models/research/attention_lm.py
@@ -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.
+
 """Self-attention based language model.
 
 DEPRECATED. Use Transformer which supports running the decoder only.
diff --git a/tensor2tensor/models/research/attention_lm_moe.py b/tensor2tensor/models/research/attention_lm_moe.py
index a060e9ac1..ec2c710e9 100644
--- a/tensor2tensor/models/research/attention_lm_moe.py
+++ b/tensor2tensor/models/research/attention_lm_moe.py
@@ -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.
+
 """Self-attention based language model.
 
 Like transformer.py, but no encoder
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index eb125d7f0..c7e5b1406 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -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.
+
 """Autoencoders."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index b5717b9e6..3d8832d47 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -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.
+
 """Autoencoders tests."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index 3974ff2df..935e292f2 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -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.
+
 """Cycle GAN."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/gene_expression.py b/tensor2tensor/models/research/gene_expression.py
index 69999b861..e15d2fd8d 100644
--- a/tensor2tensor/models/research/gene_expression.py
+++ b/tensor2tensor/models/research/gene_expression.py
@@ -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.
+
 """Models for gene expression from DNA."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/models/research/gene_expression_test.py b/tensor2tensor/models/research/gene_expression_test.py
index 8e65334de..60f9019a8 100644
--- a/tensor2tensor/models/research/gene_expression_test.py
+++ b/tensor2tensor/models/research/gene_expression_test.py
@@ -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.
+
 """Tests for Gene Expression models."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 428c5fad8..62864be99 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -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.
+
 """Glow generative model."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 05b8218f5..faf6dec16 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -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.
+
 """Various reversible ops for the glow generative model."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 0d97a52db..6aa889dbd 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -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.
+
 """Tests for tensor2tensor.models.research.glow_ops."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
index 6456d7a29..e47f5df7e 100644
--- a/tensor2tensor/models/research/glow_test.py
+++ b/tensor2tensor/models/research/glow_test.py
@@ -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.
+
 """Tests for tensor2tensor.models.research.glow_model."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/lm_experiments.py b/tensor2tensor/models/research/lm_experiments.py
index 6e5b11094..97fef03a2 100644
--- a/tensor2tensor/models/research/lm_experiments.py
+++ b/tensor2tensor/models/research/lm_experiments.py
@@ -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.
+
 """Experiments with Language Models.
 
 Train languagemodel_lm1b32k_packed and measure log-ppl/token (dev).
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index bd6b75327..0cdbd5569 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -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.
+
 """Reinforcement learning models and parameters."""
 
 import collections
diff --git a/tensor2tensor/models/research/similarity_transformer.py b/tensor2tensor/models/research/similarity_transformer.py
index 08575d0d6..43b0b31ed 100644
--- a/tensor2tensor/models/research/similarity_transformer.py
+++ b/tensor2tensor/models/research/similarity_transformer.py
@@ -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.
+
 """Using Transformer Networks for String similarities."""
 from tensor2tensor.data_generators import problem
 from tensor2tensor.layers import common_layers
diff --git a/tensor2tensor/models/research/super_lm.py b/tensor2tensor/models/research/super_lm.py
index f90e40d57..aaf8e0bfb 100644
--- a/tensor2tensor/models/research/super_lm.py
+++ b/tensor2tensor/models/research/super_lm.py
@@ -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.
+
 """Supercomputer-based language model.
 
 Uses model-parallelism.
diff --git a/tensor2tensor/models/research/transformer_aux.py b/tensor2tensor/models/research/transformer_aux.py
index 1e096a32f..fb486757f 100644
--- a/tensor2tensor/models/research/transformer_aux.py
+++ b/tensor2tensor/models/research/transformer_aux.py
@@ -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.
+
 """Transformer with auxiliary losses from https://arxiv.org/abs/1803.00144."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/transformer_aux_test.py b/tensor2tensor/models/research/transformer_aux_test.py
index 32e71d838..b32d7cc56 100644
--- a/tensor2tensor/models/research/transformer_aux_test.py
+++ b/tensor2tensor/models/research/transformer_aux_test.py
@@ -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.
+
 """Tests for tensor2tensor.models.research.transformer_aux."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index 5d6f3d1c3..a0e5d2ba1 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -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.
+
 """transformer (attention seq-seq model) with mixtures of experts.
 
 """
diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index bf11926e1..a6fc80129 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -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.
+
 """NAT Transformer from https://arxiv.org/abs/1805.11063."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/models/research/transformer_revnet.py b/tensor2tensor/models/research/transformer_revnet.py
index f31b262bc..b21338a33 100644
--- a/tensor2tensor/models/research/transformer_revnet.py
+++ b/tensor2tensor/models/research/transformer_revnet.py
@@ -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.
+
 """Reversible Residual Transformer."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/transformer_revnet_test.py b/tensor2tensor/models/research/transformer_revnet_test.py
index 76e41c0d3..f0fabcf8c 100644
--- a/tensor2tensor/models/research/transformer_revnet_test.py
+++ b/tensor2tensor/models/research/transformer_revnet_test.py
@@ -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.
+
 """Tests for TransformerRevnet."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/transformer_sketch.py b/tensor2tensor/models/research/transformer_sketch.py
index b5c6b2ac2..a64d959b6 100644
--- a/tensor2tensor/models/research/transformer_sketch.py
+++ b/tensor2tensor/models/research/transformer_sketch.py
@@ -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.
+
 """Transformer Sketch for im2sketch problems.
 """
 
diff --git a/tensor2tensor/models/research/transformer_symshard.py b/tensor2tensor/models/research/transformer_symshard.py
index 0e182fc36..1762df8f1 100644
--- a/tensor2tensor/models/research/transformer_symshard.py
+++ b/tensor2tensor/models/research/transformer_symshard.py
@@ -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.
+
 """Test of the SymShard programming model.
 
 Symmetric model parallellism.
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 4a2cd1014..d47812db6 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -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.
+
 """AE Transformer."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/transformer_vae_test.py b/tensor2tensor/models/research/transformer_vae_test.py
index e67734504..c47c485b5 100644
--- a/tensor2tensor/models/research/transformer_vae_test.py
+++ b/tensor2tensor/models/research/transformer_vae_test.py
@@ -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.
+
 """Tests for tensor2tensor.models.research.transformer_vae."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 48100f2c6..90ed8679c 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -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.
+
 """Universal Transformers.
 
 Universal Transformer is described in https://arxiv.org/abs/1807.03819.
diff --git a/tensor2tensor/models/research/universal_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
index d8c389aac..208848f7e 100644
--- a/tensor2tensor/models/research/universal_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -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.
+
 """Tests for Transformer."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 82f0d53f8..7b1ade84d 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -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.
+
 """Utilities for Universal Transformer.
 
 The Universal Transformer is based on the popular encoder-decoder architecture.
diff --git a/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index c7af49e30..aaaeb3a70 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -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.
+
 """Attention models for VQA."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/vqa_attention_test.py b/tensor2tensor/models/research/vqa_attention_test.py
index e3e88ba0d..de2016259 100644
--- a/tensor2tensor/models/research/vqa_attention_test.py
+++ b/tensor2tensor/models/research/vqa_attention_test.py
@@ -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.
+
 """Vqa_attention_baseline tests."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/vqa_recurrent_self_attention.py b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
index 387bc0484..1be385c1f 100644
--- a/tensor2tensor/models/research/vqa_recurrent_self_attention.py
+++ b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
@@ -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.
+
 """Recurrent self attention models for VQA."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
index e6794dde8..b7c07ad0e 100644
--- a/tensor2tensor/models/research/vqa_self_attention.py
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -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.
+
 """Self attention models for VQA."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index 4b2ae4cb7..62f31b156 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -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.
+
 """Resnets."""
 # Copied from cloud_tpu/models/resnet/resnet_model.py and modified
 
diff --git a/tensor2tensor/models/resnet_test.py b/tensor2tensor/models/resnet_test.py
index 50e84b581..fe7538c81 100644
--- a/tensor2tensor/models/resnet_test.py
+++ b/tensor2tensor/models/resnet_test.py
@@ -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.
+
 """Resnet tests."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/revnet.py b/tensor2tensor/models/revnet.py
index 49304105b..379c2b143 100644
--- a/tensor2tensor/models/revnet.py
+++ b/tensor2tensor/models/revnet.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 """Creates a RevNet with the bottleneck residual function.
 
 Implements the following equations described in the RevNet paper:
diff --git a/tensor2tensor/models/revnet_test.py b/tensor2tensor/models/revnet_test.py
index 6f4f50f3c..68fec94a2 100644
--- a/tensor2tensor/models/revnet_test.py
+++ b/tensor2tensor/models/revnet_test.py
@@ -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.
+
 """Tests for Revnet."""
 
 from tensor2tensor.models import revnet
diff --git a/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index 77d953380..4a83444c0 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -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.
+
 """Shake-shake model for CIFAR."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/slicenet.py b/tensor2tensor/models/slicenet.py
index fc9ecbf50..9806d8891 100644
--- a/tensor2tensor/models/slicenet.py
+++ b/tensor2tensor/models/slicenet.py
@@ -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.
+
 """SliceNet."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/models/slicenet_test.py b/tensor2tensor/models/slicenet_test.py
index 89919aa10..314d0dae8 100644
--- a/tensor2tensor/models/slicenet_test.py
+++ b/tensor2tensor/models/slicenet_test.py
@@ -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.
+
 """Tests for SliceNet."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 4e4cc7515..836f1812a 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -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.
+
 """Transformer model from "Attention Is All You Need".
 
 The Transformer model consists of an encoder and a decoder. Both are stacks
diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index d569e3b96..511d0c3a4 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -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.
+
 """Tests for Transformer."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index 4e4686bd5..cb16a4c0b 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -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.
+
 """Simple Generative Adversarial Model with two linear layers.
 
 Example of how to create a GAN in T2T.
diff --git a/tensor2tensor/models/video/__init__.py b/tensor2tensor/models/video/__init__.py
index dba7ece95..4bd418a74 100644
--- a/tensor2tensor/models/video/__init__.py
+++ b/tensor2tensor/models/video/__init__.py
@@ -12,3 +12,4 @@
 # 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.
+
diff --git a/tensor2tensor/models/video/base_vae.py b/tensor2tensor/models/video/base_vae.py
index 0c1e604ee..8ef5e5a9b 100644
--- a/tensor2tensor/models/video/base_vae.py
+++ b/tensor2tensor/models/video/base_vae.py
@@ -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.
+
 """Basic models for testing simple tasks."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index f82afe5bb..1ce18b68f 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -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.
+
 """Basic models for testing simple tasks."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index c102f89a5..5a695576b 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -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.
+
 """Param sets for deterministic basic next frame prediction model."""
 
 from __future__ import division
diff --git a/tensor2tensor/models/video/basic_deterministic_test.py b/tensor2tensor/models/video/basic_deterministic_test.py
index bd95d76c8..20d007478 100644
--- a/tensor2tensor/models/video/basic_deterministic_test.py
+++ b/tensor2tensor/models/video/basic_deterministic_test.py
@@ -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.
+
 """Basic tests for basic deterministic model."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/video/basic_recurrent.py b/tensor2tensor/models/video/basic_recurrent.py
index 79f155501..0b0af311d 100644
--- a/tensor2tensor/models/video/basic_recurrent.py
+++ b/tensor2tensor/models/video/basic_recurrent.py
@@ -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.
+
 """Basic recurrent models for testing simple tasks."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/video/basic_recurrent_test.py b/tensor2tensor/models/video/basic_recurrent_test.py
index 5b1c9d445..d9deb5c98 100644
--- a/tensor2tensor/models/video/basic_recurrent_test.py
+++ b/tensor2tensor/models/video/basic_recurrent_test.py
@@ -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.
+
 """Basic tests for basic deterministic model."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 92002d490..26b0d4d5f 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -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.
+
 """Basic models for testing simple tasks."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/video/basic_stochastic_test.py b/tensor2tensor/models/video/basic_stochastic_test.py
index fdc4951f6..d74f42117 100644
--- a/tensor2tensor/models/video/basic_stochastic_test.py
+++ b/tensor2tensor/models/video/basic_stochastic_test.py
@@ -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.
+
 """Basic tests for basic stochastic model."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index cbdc0f479..2a54d50ae 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -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.
+
 """Model architecture for video prediction model.
 
    based on following paper:
diff --git a/tensor2tensor/models/video/emily_test.py b/tensor2tensor/models/video/emily_test.py
index 51a8e2e28..b670c5de8 100644
--- a/tensor2tensor/models/video/emily_test.py
+++ b/tensor2tensor/models/video/emily_test.py
@@ -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.
+
 """Basic tests for emily's model."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index 7aa7c3f76..468910b93 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -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.
+
 """Model architecture for video prediction model.
 
 based on following paper:
diff --git a/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
index 2260b4c52..2281af593 100644
--- a/tensor2tensor/models/video/epva_params.py
+++ b/tensor2tensor/models/video/epva_params.py
@@ -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.
+
 """Param sets for EPVA model."""
 
 from __future__ import division
diff --git a/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index f1a8d4986..bf5b9cef7 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -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.
+
 """Stochastic Adversarial Video Prediction model.
 
 Reference: https://arxiv.org/abs/1804.01523
diff --git a/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index bd5a06a75..dbc02c771 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -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.
+
 """Param sets for SAVP model."""
 
 from __future__ import division
diff --git a/tensor2tensor/models/video/savp_test.py b/tensor2tensor/models/video/savp_test.py
index 80ad6f47b..2061eb06b 100644
--- a/tensor2tensor/models/video/savp_test.py
+++ b/tensor2tensor/models/video/savp_test.py
@@ -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.
+
 """Basic tests for SAVP model."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 217a481ff..8de30ee39 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -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.
+
 """SV2P: Stochastic Variational Video Prediction.
 
    based on the following paper:
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 96879c6d1..649d418ef 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -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.
+
 """Param sets for SV2P model."""
 
 from __future__ import division
diff --git a/tensor2tensor/models/video/sv2p_test.py b/tensor2tensor/models/video/sv2p_test.py
index a8d94ae50..98b262a69 100644
--- a/tensor2tensor/models/video/sv2p_test.py
+++ b/tensor2tensor/models/video/sv2p_test.py
@@ -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.
+
 """Basic tests for SV2P model."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/video/tests_utils.py b/tensor2tensor/models/video/tests_utils.py
index 3004f3bc2..99d6030ab 100644
--- a/tensor2tensor/models/video/tests_utils.py
+++ b/tensor2tensor/models/video/tests_utils.py
@@ -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.
+
 """Utilties for testing video models."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/xception.py b/tensor2tensor/models/xception.py
index 2452a7d4f..3f59690b0 100644
--- a/tensor2tensor/models/xception.py
+++ b/tensor2tensor/models/xception.py
@@ -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.
+
 """Xception."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/xception_test.py b/tensor2tensor/models/xception_test.py
index c3c8e9bae..95a07ef38 100644
--- a/tensor2tensor/models/xception_test.py
+++ b/tensor2tensor/models/xception_test.py
@@ -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.
+
 """Xception tests."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/problems.py b/tensor2tensor/problems.py
index 3f349b908..1337c2a33 100644
--- a/tensor2tensor/problems.py
+++ b/tensor2tensor/problems.py
@@ -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.
+
 """Access T2T Problems."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/problems_test.py b/tensor2tensor/problems_test.py
index 35db33b70..07387fd53 100644
--- a/tensor2tensor/problems_test.py
+++ b/tensor2tensor/problems_test.py
@@ -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.
+
 """tensor2tensor.problems test."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/rl/OWNERS b/tensor2tensor/rl/OWNERS
new file mode 100644
index 000000000..cfa1f4f74
--- /dev/null
+++ b/tensor2tensor/rl/OWNERS
@@ -0,0 +1,3 @@
+blazej
+dumitru
+mbz
diff --git a/tensor2tensor/rl/__init__.py b/tensor2tensor/rl/__init__.py
index dba7ece95..4bd418a74 100644
--- a/tensor2tensor/rl/__init__.py
+++ b/tensor2tensor/rl/__init__.py
@@ -12,3 +12,4 @@
 # 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.
+
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index fbdde538c..f39ecde05 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -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.
+
 """Collect trajectories from interactions of agent with environment."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/rl/datagen_with_agent.py b/tensor2tensor/rl/datagen_with_agent.py
index 1cdbc9d5d..a5d85d38b 100644
--- a/tensor2tensor/rl/datagen_with_agent.py
+++ b/tensor2tensor/rl/datagen_with_agent.py
@@ -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.
+
 """Generate trajectories to disk with random or ckpt agent.
 
 TODO: Usage
diff --git a/tensor2tensor/rl/envs/__init__.py b/tensor2tensor/rl/envs/__init__.py
index dba7ece95..4bd418a74 100644
--- a/tensor2tensor/rl/envs/__init__.py
+++ b/tensor2tensor/rl/envs/__init__.py
@@ -12,3 +12,4 @@
 # 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.
+
diff --git a/tensor2tensor/rl/envs/batch_env.py b/tensor2tensor/rl/envs/batch_env.py
index 862a8c9b2..7b78d0231 100644
--- a/tensor2tensor/rl/envs/batch_env.py
+++ b/tensor2tensor/rl/envs/batch_env.py
@@ -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.
+
 """Combine multiple environments to step them in batch."""
 
 # The code was based on Danijar Hafner's code from tf.agents:
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index 096744128..0b71df225 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -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.
+
 """Utilities for creating batched environments."""
 
 # The code was based on Danijar Hafner's code from tf.agents:
diff --git a/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index 3cb7ae551..58492b4fd 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -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.
+
 """Batch of environments inside the TensorFlow graph."""
 
 # The code was based on Danijar Hafner's code from tf.agents:
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index e6e27d913..5bc3bdfb6 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -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.
+
 """Batch of environments inside the TensorFlow graph."""
 
 # The code was based on Danijar Hafner's code from tf.agents:
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index b478bd2bb..f48d66e64 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -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.
+
 """Batch of environments inside the TensorFlow graph."""
 
 # The code was based on Danijar Hafner's code from tf.agents:
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 3c0bde897..f87701060 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -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.
+
 """Batch of environments inside the TensorFlow graph."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index 9d8850a2c..6cf1df270 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -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.
+
 """Utilities for using batched environments."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index ce8baddc2..259b2dec9 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -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.
+
 """Play with a world model."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index 4f2f1c655..67439bac8 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -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.
+
 """PPO algorithm implementation.
 
 Based on: https://arxiv.org/abs/1707.06347
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 88c62e015..0af54b6a5 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -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.
+
 """Library for training of RL agent with PPO algorithm."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index a8a02cac9..bf777c39f 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -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.
+
 """Tests of basic flow of collecting trajectories and training PPO."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 57de06627..ec2a555d4 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -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.
+
 r"""Training of model-based RL agents.
 
 Example invocation:
diff --git a/tensor2tensor/rl/trainer_model_based_ae_test.py b/tensor2tensor/rl/trainer_model_based_ae_test.py
index a7b15a62d..ba978eab0 100644
--- a/tensor2tensor/rl/trainer_model_based_ae_test.py
+++ b/tensor2tensor/rl/trainer_model_based_ae_test.py
@@ -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.
+
 """Tiny run of trainer_model_based. Smoke test."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/rl/trainer_model_based_agent_only.py b/tensor2tensor/rl/trainer_model_based_agent_only.py
index e15b01aaa..bf4ac0b12 100644
--- a/tensor2tensor/rl/trainer_model_based_agent_only.py
+++ b/tensor2tensor/rl/trainer_model_based_agent_only.py
@@ -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.
+
 r"""Training of model-based RL agent assuming a fully trained world model.
 
 Example invocation:
diff --git a/tensor2tensor/rl/trainer_model_based_recurrent_test.py b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
index 629ff1245..42d1089f6 100644
--- a/tensor2tensor/rl/trainer_model_based_recurrent_test.py
+++ b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
@@ -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.
+
 """Tiny run of trainer_model_based. Smoke test."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/rl/trainer_model_based_stochastic_test.py b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
index fd22fa021..018ab2675 100644
--- a/tensor2tensor/rl/trainer_model_based_stochastic_test.py
+++ b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
@@ -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.
+
 """Tiny run of trainer_model_based with stochastic model. Smoke test."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/rl/trainer_model_based_sv2p_test.py b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
index 5b715c20c..8be291555 100644
--- a/tensor2tensor/rl/trainer_model_based_sv2p_test.py
+++ b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
@@ -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.
+
 """Tiny run of trainer_model_based with stochastic model. Smoke test."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/rl/trainer_model_based_test.py b/tensor2tensor/rl/trainer_model_based_test.py
index f528af1c9..83cab7d60 100644
--- a/tensor2tensor/rl/trainer_model_based_test.py
+++ b/tensor2tensor/rl/trainer_model_based_test.py
@@ -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.
+
 """Tiny run of trainer_model_based. Smoke test."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 2471b9d59..0375afb86 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -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.
+
 r"""Training of RL agent with PPO algorithm.
 
 Example invocation:
diff --git a/tensor2tensor/serving/__init__.py b/tensor2tensor/serving/__init__.py
index dba7ece95..4bd418a74 100644
--- a/tensor2tensor/serving/__init__.py
+++ b/tensor2tensor/serving/__init__.py
@@ -12,3 +12,4 @@
 # 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.
+
diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 04dc34153..7fd55f2be 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -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.
+
 """Export a trained model for serving."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/serving/query.py b/tensor2tensor/serving/query.py
index 7d89d3c7e..1af0e9f2d 100644
--- a/tensor2tensor/serving/query.py
+++ b/tensor2tensor/serving/query.py
@@ -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.
+
 """Query an exported model. Py2 only. Install tensorflow-serving-api."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 58aac9b0e..8d50a7756 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -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.
+
 """Utilities for serving tensor2tensor."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/test_data/example_usr_dir/__init__.py b/tensor2tensor/test_data/example_usr_dir/__init__.py
index c1d23fdfa..61a511e17 100644
--- a/tensor2tensor/test_data/example_usr_dir/__init__.py
+++ b/tensor2tensor/test_data/example_usr_dir/__init__.py
@@ -12,5 +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.
+
 """Example T2T user directory."""
 from . import my_submodule
diff --git a/tensor2tensor/test_data/example_usr_dir/my_submodule.py b/tensor2tensor/test_data/example_usr_dir/my_submodule.py
index a6b31469b..e3ffd962c 100644
--- a/tensor2tensor/test_data/example_usr_dir/my_submodule.py
+++ b/tensor2tensor/test_data/example_usr_dir/my_submodule.py
@@ -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.
+
 """Example registrations for T2T."""
 import re
 
diff --git a/tensor2tensor/utils/__init__.py b/tensor2tensor/utils/__init__.py
index dba7ece95..4bd418a74 100644
--- a/tensor2tensor/utils/__init__.py
+++ b/tensor2tensor/utils/__init__.py
@@ -12,3 +12,4 @@
 # 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.
+
diff --git a/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index 4f627675f..466d5b59d 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -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.
+
 """Optimization."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/adv_attack_utils.py b/tensor2tensor/utils/adv_attack_utils.py
index 9418854e9..070f9cfd4 100644
--- a/tensor2tensor/utils/adv_attack_utils.py
+++ b/tensor2tensor/utils/adv_attack_utils.py
@@ -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.
+
 """Utilities to assist in performing adversarial attack using Cleverhans."""
 
 from cleverhans import attacks
diff --git a/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index e498bc18a..879de16fa 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -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.
+
 """Script to average values of variables in a list of checkpoint files."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index 98ce137b9..92f9fc76c 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -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.
+
 """Implementation of beam search with penalties."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/beam_search_test.py b/tensor2tensor/utils/beam_search_test.py
index a4277f2b2..e7dbc9e75 100644
--- a/tensor2tensor/utils/beam_search_test.py
+++ b/tensor2tensor/utils/beam_search_test.py
@@ -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.
+
 """Tests for tensor2tensor.beam_search."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/bleu_hook.py b/tensor2tensor/utils/bleu_hook.py
index 1e75a1269..2fa4ce0a4 100644
--- a/tensor2tensor/utils/bleu_hook.py
+++ b/tensor2tensor/utils/bleu_hook.py
@@ -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.
+
 """BLEU metric util used during eval for MT."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/bleu_hook_test.py b/tensor2tensor/utils/bleu_hook_test.py
index e291adbac..999a240f8 100644
--- a/tensor2tensor/utils/bleu_hook_test.py
+++ b/tensor2tensor/utils/bleu_hook_test.py
@@ -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.
+
 # coding=utf-8
 """Tests for tensor2tensor.utils.bleu_hook."""
 
diff --git a/tensor2tensor/utils/checkpoint_compatibility_test.py b/tensor2tensor/utils/checkpoint_compatibility_test.py
index a1a121ab3..d4c4c815c 100644
--- a/tensor2tensor/utils/checkpoint_compatibility_test.py
+++ b/tensor2tensor/utils/checkpoint_compatibility_test.py
@@ -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.
+
 """Test for checkpoint compatibility."""
 # The checkpoint in test_data/transformer_test_ckpt is generated with the OSS
 # release.
diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 8af62c112..c62069d24 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -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.
+
 """Launch on GCP's ML Engine."""
 
 import datetime
diff --git a/tensor2tensor/utils/compute_video_metrics.py b/tensor2tensor/utils/compute_video_metrics.py
index 45b077004..9db41038d 100644
--- a/tensor2tensor/utils/compute_video_metrics.py
+++ b/tensor2tensor/utils/compute_video_metrics.py
@@ -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.
+
 """Computes and saves the metrics for video prediction and generation."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index 14bd04624..abac8e4a3 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -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.
+
 """Data reader module."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/data_reader_test.py b/tensor2tensor/utils/data_reader_test.py
index ccee17c5c..53621d3fd 100644
--- a/tensor2tensor/utils/data_reader_test.py
+++ b/tensor2tensor/utils/data_reader_test.py
@@ -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.
+
 """Data reader test."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index fed8e8387..31437d3c7 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -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.
+
 """Decoding utilities."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index 2d0cbe48e..008525540 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -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.
+
 """Device placement and data parallelism."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/diet.py b/tensor2tensor/utils/diet.py
index ec8212fa8..700957280 100644
--- a/tensor2tensor/utils/diet.py
+++ b/tensor2tensor/utils/diet.py
@@ -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.
+
 """Diet variables are much more memory-efficient than regular variables.
 
 Using diet variables, we can reduce memory overhead per parameter from
diff --git a/tensor2tensor/utils/diet_test.py b/tensor2tensor/utils/diet_test.py
index 809a5769f..667ad1f7f 100644
--- a/tensor2tensor/utils/diet_test.py
+++ b/tensor2tensor/utils/diet_test.py
@@ -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.
+
 """Tests for common layers."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 0f651ce17..01a525ded 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -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.
+
 """Utilities for creating Sparsely-Gated Mixture-of-Experts Layers.
 
 See "Outrageously Large Neural Networks"
diff --git a/tensor2tensor/utils/expert_utils_test.py b/tensor2tensor/utils/expert_utils_test.py
index 33209fc47..5c9cc714e 100644
--- a/tensor2tensor/utils/expert_utils_test.py
+++ b/tensor2tensor/utils/expert_utils_test.py
@@ -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.
+
 """Tests for tensor2tensor.utils.expert_utils."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index 3cd0c0906..15055f54e 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -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.
+
 """Common command-line flags."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/get_rouge.py b/tensor2tensor/utils/get_rouge.py
index 65dc883a6..e374811aa 100644
--- a/tensor2tensor/utils/get_rouge.py
+++ b/tensor2tensor/utils/get_rouge.py
@@ -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.
+
 """Computing rouge scores using pyrouge."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index e9b65b83a..f93f267b4 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -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.
+
 """Optimization."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 207afe4c2..7840136de 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -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.
+
 """Utils for metrics used in eval."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/metrics_hook.py b/tensor2tensor/utils/metrics_hook.py
index 2482b487b..2fed06c18 100644
--- a/tensor2tensor/utils/metrics_hook.py
+++ b/tensor2tensor/utils/metrics_hook.py
@@ -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.
+
 """Summary-based SessionRunHooks."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/metrics_hook_test.py b/tensor2tensor/utils/metrics_hook_test.py
index 35695d1ae..f707f24af 100644
--- a/tensor2tensor/utils/metrics_hook_test.py
+++ b/tensor2tensor/utils/metrics_hook_test.py
@@ -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.
+
 """Tests for metrics_hook."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 3e504d881..7878d784a 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -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.
+
 """Tests for tensor2tensor.utils.metrics."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index ded5577ef..6e6726885 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -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.
+
 """Modality base class - defines the bottom and top of the model."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
index d4c5c6b86..3f516cd8e 100644
--- a/tensor2tensor/utils/multistep_optimizer.py
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -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.
+
 """Multi-step optimizers simulating large batches.
 
 Optimizer variants which make it possible to use very large batch sizes with
diff --git a/tensor2tensor/utils/multistep_optimizer_test.py b/tensor2tensor/utils/multistep_optimizer_test.py
index 584bd95c3..19d248b03 100644
--- a/tensor2tensor/utils/multistep_optimizer_test.py
+++ b/tensor2tensor/utils/multistep_optimizer_test.py
@@ -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.
+
 """Multi-step Optimizer Test Module for TensorFlow."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index c3164c3f9..35a6a7dd3 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -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.
+
 """Optimization."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/pruning_utils.py b/tensor2tensor/utils/pruning_utils.py
index 446e1c59c..2cc4cb35b 100644
--- a/tensor2tensor/utils/pruning_utils.py
+++ b/tensor2tensor/utils/pruning_utils.py
@@ -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.
+
 """Utilities to assist in pruning models."""
 
 import numpy as np
diff --git a/tensor2tensor/utils/quantization.py b/tensor2tensor/utils/quantization.py
index 6883423a0..cb51be809 100644
--- a/tensor2tensor/utils/quantization.py
+++ b/tensor2tensor/utils/quantization.py
@@ -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.
+
 """Utilities related to using bfloat16 activations and/or parameters."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index af4d2b0e8..de1e4fdc9 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -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.
+
 """Registry for models, hyperparameter settings, problem types, and datasets.
 
 Define a new model by subclassing T2TModel and register it:
diff --git a/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index b7b328907..a72dd36bd 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -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.
+
 """Tests for tensor2tensor.registry."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/restore_hook.py b/tensor2tensor/utils/restore_hook.py
index 4d6fdd3c4..c649231c7 100644
--- a/tensor2tensor/utils/restore_hook.py
+++ b/tensor2tensor/utils/restore_hook.py
@@ -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.
+
 """Restore hooks."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/rouge.py b/tensor2tensor/utils/rouge.py
index 7a956b5dd..cb3c9af4b 100644
--- a/tensor2tensor/utils/rouge.py
+++ b/tensor2tensor/utils/rouge.py
@@ -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.
+
 # coding=utf-8
 """ROUGE metric implementation.
 
diff --git a/tensor2tensor/utils/rouge_test.py b/tensor2tensor/utils/rouge_test.py
index f90be3ba7..eab9eaeeb 100644
--- a/tensor2tensor/utils/rouge_test.py
+++ b/tensor2tensor/utils/rouge_test.py
@@ -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.
+
 """Tests for Rouge metric."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 68b9ce9e6..e981a08b9 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -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.
+
 """T2TModel Base Class."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index 1a218eb0c..bcc0ed44a 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -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.
+
 """Tests for T2TModel."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index b90f222b1..48d888385 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -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.
+
 """Library for training. See t2t_trainer.py."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index bc18bad20..55b5a5aa0 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -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.
+
 """Tests for trainer_lib."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/update_ops_hook.py b/tensor2tensor/utils/update_ops_hook.py
index 184272eaf..c2615c661 100644
--- a/tensor2tensor/utils/update_ops_hook.py
+++ b/tensor2tensor/utils/update_ops_hook.py
@@ -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.
+
 """Hook to run tf.GraphKeys.UPDATE_OPS."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/usr_dir.py b/tensor2tensor/utils/usr_dir.py
index f66b21d11..06b5955ea 100644
--- a/tensor2tensor/utils/usr_dir.py
+++ b/tensor2tensor/utils/usr_dir.py
@@ -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.
+
 """Utility to load code from an external user-supplied directory."""
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/utils/video2gif.py b/tensor2tensor/utils/video2gif.py
index 4a16c43a8..a66e75402 100644
--- a/tensor2tensor/utils/video2gif.py
+++ b/tensor2tensor/utils/video2gif.py
@@ -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.
+
 r"""View the problem.
 
 This binary saves the videos in the problem(dataset) into gifs.
diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index 43b473d6e..3f5cbdf01 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -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.
+
 """Computes the metrics for video prediction and generation."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/yellowfin.py b/tensor2tensor/utils/yellowfin.py
index 0090c6fbf..4d33c2aad 100644
--- a/tensor2tensor/utils/yellowfin.py
+++ b/tensor2tensor/utils/yellowfin.py
@@ -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.
+
 """YellowFin for TensorFlow. Thanks Jian Zhang: zjian [@] stanford [.] edu."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/utils/yellowfin_test.py b/tensor2tensor/utils/yellowfin_test.py
index dd487689c..67999707c 100644
--- a/tensor2tensor/utils/yellowfin_test.py
+++ b/tensor2tensor/utils/yellowfin_test.py
@@ -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.
+
 """YellowFin Test Module for TensorFlow."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/visualization/__init__.py b/tensor2tensor/visualization/__init__.py
index 4bd418a74..c2a9550b0 100644
--- a/tensor2tensor/visualization/__init__.py
+++ b/tensor2tensor/visualization/__init__.py
@@ -13,3 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
diff --git a/tensor2tensor/visualization/attention.py b/tensor2tensor/visualization/attention.py
index abf9a0e57..56ece8154 100644
--- a/tensor2tensor/visualization/attention.py
+++ b/tensor2tensor/visualization/attention.py
@@ -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.
+
 """Module for postprocessing and displaying transformer attentions.
 
 This module is designed to be called from an ipython notebook.
diff --git a/tensor2tensor/visualization/visualization.py b/tensor2tensor/visualization/visualization.py
index 59d07dc33..f21621a25 100644
--- a/tensor2tensor/visualization/visualization.py
+++ b/tensor2tensor/visualization/visualization.py
@@ -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.
+
 """Shared code for visualizing transformer attentions."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/visualization/visualization_test.py b/tensor2tensor/visualization/visualization_test.py
index b3b0b7ce4..9c1371a7d 100644
--- a/tensor2tensor/visualization/visualization_test.py
+++ b/tensor2tensor/visualization/visualization_test.py
@@ -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.
+
 """Tests for visualization library.
 
     IF ANY OF THESE TESTS BREAK PLEASE UPDATE THE CODE IN THE VIZ NOTEBOOK

From 3750e143c0a4dcdc0de91fc13a2d7fe80cd27766 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Fri, 5 Oct 2018 14:45:50 -0700
Subject: [PATCH 0969/2720] Some changes to mixture-of-experts in
 mesh-tensorflow:  At the inner level of hierarchical MoE, prioritize examples
 that were the first choice at the outer level.  Change default capacity
 factor to 1.25*1.0 for two-level MoE, down from 1.25*1.25 .  This leads to
 similar quality with slightly less computation.  Also, remove renormalization
 of top two experts after possibly dropping out second one.

PiperOrigin-RevId: 215969523
---
 .../mesh_tensorflow/mesh_tensorflow.py        |  6 +-
 .../research/experiments_moe.py               | 52 ++++++++++++
 tensor2tensor/mesh_tensorflow/research/moe.py | 84 +++++++++++--------
 3 files changed, 104 insertions(+), 38 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index 0508aacc7..6ff65ab8d 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -2782,7 +2782,8 @@ def lower(self, lowering):
         # TODO(noam): try to handle this case
         raise NotImplementedError(
             "Try first reshaping to insert a new tf dimension,"
-            " then changing layout.")
+            " then changing layout. input_shape=%s output_shape=%s"
+            % (self.inputs[0].shape, self.outputs[0].shape))
       slices = mesh_impl.allsplit(slices, mesh_axis, tensor_axis)
       laid_out_size //= mesh_impl.shape[mesh_axis].size
     for mesh_axis in mesh_axes_alltoall:
@@ -2792,7 +2793,8 @@ def lower(self, lowering):
         # TODO(noam): try to handle this case
         raise NotImplementedError(
             "Try first reshaping to insert a new tf dimension,"
-            " then changing layout.")
+            " then changing layout. input_shape=%s output_shape=%s"
+            % (self.inputs[0].shape, self.outputs[0].shape))
       concat_tensor_axis = old_shape.cumprod_to_tensor_axis(
           mesh_axis_to_cumprod_old[mesh_axis])
       assert concat_tensor_axis is not None
diff --git a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
index 12e612bee..915250335 100644
--- a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
@@ -140,6 +140,14 @@ def xmoe_2d():
   return hparams
 
 
+@registry.register_hparams
+def xmoe_2d_c15():
+  """Mixture of experts."""
+  hparams = xmoe_2d()
+  hparams.moe_capacity_factor_train = 1.5
+  return hparams
+
+
 @registry.register_hparams
 def xmoe_2d_88():
   """Two-dimensional hierarchical mixture of experts."""
@@ -221,6 +229,42 @@ def xmoe_wiki_x64():
   return hparams
 
 
+@registry.register_hparams
+def xmoe_wiki_x32():
+  """Two-dimensional hierarchical mixture of experts.
+
+  (8x4 experts) * (16M params/expert) * 6 layers = 3B params
+
+  Returns:
+    a hparams object.
+  """
+  hparams = xmoe_wiki_base()
+  moe.set_default_moe_hparams(hparams)
+  hparams.feedforward_layer = "hmoe"
+  hparams.moe_hidden_size = 8192
+  hparams.mesh_shape = "b0:4;b1:8"
+  hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
+  hparams.outer_batch_size = 4
+  hparams.moe_num_experts = [8, 4]
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_wiki_x64_h16k():
+  """Mixture of experts."""
+  hparams = xmoe_wiki_x64()
+  hparams.moe_hidden_size = 16384
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_wiki_x64_c15():
+  """Mixture of experts."""
+  hparams = xmoe_wiki_x64()
+  hparams.moe_capacity_factor_train = 1.5
+  return hparams
+
+
 @registry.register_hparams
 def xmoe_wiki_x256():
   """Two-dimensional hierarchical mixture of experts.
@@ -285,3 +329,11 @@ def xmoe_wiki_x1024_h16k():
   return hparams
 
 
+@registry.register_hparams
+def xmoe_wiki_x256_c15():
+  """Mixture of experts."""
+  hparams = xmoe_wiki_x256()
+  hparams.moe_capacity_factor_train = 1.5
+  return hparams
+
+
diff --git a/tensor2tensor/mesh_tensorflow/research/moe.py b/tensor2tensor/mesh_tensorflow/research/moe.py
index 73542cd8e..fac763cb3 100644
--- a/tensor2tensor/mesh_tensorflow/research/moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/moe.py
@@ -124,7 +124,6 @@ def transformer_moe_layer_v1(inputs, output_dim, hparams, train):
         outer_expert_dims=None,
         experts_dim=experts_dim_unsplit,
         expert_capacity_dim=expert_capacity_dim,
-        max_experts=None,
         hparams=hparams,
         train=train)
   else:
@@ -171,6 +170,7 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
     hparams.moe_group_size: size of each "group" for gating purposes
     hparams.moe_capacity_factor_train: a float
     hparams.moe_capacity_factor_eval: a float
+    hparams.moe_capacity_factor_second_level: a float
     hparams.moe_gating: a string
     + all hyperparmeters used by _top_2_gating()
 
@@ -189,7 +189,7 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
   The input is n-dimensional: [<batch_and_length_dims>, input_dim], consisting
   of the representations of all positions in a batch of sequences.
 
-  Each position of each sequence is sent to 0-2 experts.  The expert
+  Each position of each sequence is sent to 0-3 experts.  The expert
   choices and the combination weights are determined by a learned gating
   function.
 
@@ -302,7 +302,7 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
 
   expert_capacity = min(
       t.size,
-      int((t.size * capacity_factor) / y.size))
+      int((t.size * hparams.moe_capacity_factor_second_level) / y.size))
   d = mtf.Dimension("expert_capacity_y", expert_capacity)
 
   # First level of expert routing
@@ -318,7 +318,6 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
         outer_expert_dims=None,
         experts_dim=x,
         expert_capacity_dim=c,
-        max_experts=None,
         hparams=hparams,
         train=train)
   else:
@@ -326,36 +325,45 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
 
   # Now create expert_inputs based on the assignments.
   # put num_experts dimension first to make split easier in alltoall
-  expert_inputs_x = mtf.einsum([inputs, dispatch_tensor_x], mtf.Shape(
-      [x, a0, g1, c, m]))
+  expert_inputs_x = mtf.einsum([inputs, dispatch_tensor_x], [x, a0, g1, c, m])
+
+  # we construct an "importance" Tensor for the inputs to the second-level
+  # gating.  The importance of an input is 1.0 if it represents the
+  # first-choice expert-group and 0.5 if it represents the second-choice expert
+  # group.  This is used by the second-level gating.
+  importance = mtf.reduce_sum(combine_tensor_x, output_shape=[x, a0, g1, c])
+  importance = 0.5 * (
+      mtf.to_float(mtf.greater(importance, 0.5)) +
+      mtf.to_float(mtf.greater(importance, 0.0)))
 
   # First level, all to all. Here we change the split dimension from g1 to x1.
   expert_inputs_x = mtf.reshape(expert_inputs_x, mtf.Shape(
       [x1, a0, g, c, m]))
+  importance = mtf.reshape(importance, [x1, a0, g, c])
 
   # Second level of expert routing
   # Reshape the expert_inputs outer batch dim to be a multiple of group_dim h0
   # and group_size_dim t.
   inputs_y = mtf.reshape(expert_inputs_x, [x1, h0, t, m])
+  importance = mtf.reshape(importance, [x1, h0, t])
 
   # Get the assignments for the second level.
-  # dispatch_tensor_x has shape [x1, h0, t, y, d]
+  # dispatch_tensor_y has shape [x1, h0, t, y, d]
   if hparams.moe_gating == "top_2":
     dispatch_tensor_y, combine_tensor_y, loss_inner = _top_2_gating(
         inputs=inputs_y,
         outer_expert_dims=[x1],
         experts_dim=y,
         expert_capacity_dim=d,
-        max_experts=None,
         hparams=hparams,
-        train=train)
+        train=train,
+        importance=importance)
   else:
     raise ValueError("unknown hparams.moe_gating=%s" % hparams.moe_gating)
 
   # Now create expert_inputs based on the assignments.
   # put num_experts dimension first to make split easier in alltoall
-  expert_inputs_y = mtf.einsum([inputs_y, dispatch_tensor_y], mtf.Shape(
-      [y, x1, h0, d, m]))
+  expert_inputs_y = mtf.einsum([inputs_y, dispatch_tensor_y], [y, x1, h0, d, m])
 
   # Second level, all to all. Here we change the split dimension from h0 to y0.
   expert_inputs_y = mtf.reshape(expert_inputs_y, mtf.Shape(
@@ -377,8 +385,7 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
       [y, x1, h0, d, n]))
 
   # combine results from inner level
-  output_y = mtf.einsum([expert_output, combine_tensor_y], mtf.Shape(
-      [x1, h0, t, n]))
+  output_y = mtf.einsum([expert_output, combine_tensor_y], [x1, h0, t, n])
 
   # Reshape the combined tensor from inner level to now contain outer_batch_dim
   # a0 and group_dim g
@@ -388,8 +395,7 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
   expert_output_x = mtf.reshape(output, mtf.Shape([x, a0, g1, c, n]))
 
   # combine results from outer level
-  output_x = mtf.einsum([expert_output_x, combine_tensor_x], mtf.Shape(
-      [a0, g1, s, m]))
+  output_x = mtf.einsum([expert_output_x, combine_tensor_x], [a0, g1, s, n])
 
   # Reshape the combined tensor to now contain inner_batch_dim
   # b1 and the original sequence length
@@ -400,8 +406,8 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
 
 
 def _top_2_gating(
-    inputs, outer_expert_dims, experts_dim, expert_capacity_dim, max_experts,
-    hparams, train):
+    inputs, outer_expert_dims, experts_dim, expert_capacity_dim,
+    hparams, train, importance=None):
   """Compute gating for mixture-of-experts in TensorFlow.
 
   Note: until the algorithm and inferface solidify, we pass in a hyperparameters
@@ -415,17 +421,14 @@ def _top_2_gating(
     hparams.moe_second_policy_eval: a string
     hparams.moe_second_threshold: a float
 
-  max_experts is an float tensor with shape [<batch_dims>, group_size_dim]
-  indicating at most how many experts to use per example.  This can be
-  used to prevent padding from going to experts.
-
   The returned forward assignment is a tensor used to map (via einsum) from the
   inputs to the expert_inputs.  Likewise, the returned combine_tensor is
   used to map (via einsum) from the expert outputs to the outputs.  Both the
-  forward and backward assignments are mostly zeros.  The shapes of all of these
+  forward and backward assignments are mostly zeros.  The shapes of the tensors
   are as follows.
 
   inputs: [<batch_dims>, group_size_dim, input_dim]
+  importance: [<batch_dims>, group_size_dim]
   dispatch_tensor:
     [<batch_dims>, group_size_dim, experts_dim, expert_capacity_dim]
   expert_inputs:
@@ -436,15 +439,25 @@ def _top_2_gating(
     [<batch_dims>, group_size_dim, experts_dim, expert_capacity_dim]
   outputs: [<batch_dims>, group_size_dim, output_dim]
 
+  "importance" is an optional tensor with one floating-point value for each
+  input vector.  If the importance of an input is 1.0, then we send it to
+  up to 2 experts.  If 0.0 < importance < 1.0, then we send it to at most
+  one expert.  If importance == 0.0, then we send it to no experts.
+
+  We use "importance" at the second-level gating function of a hierarchical
+  mixture of experts.  Inputs to the first-choice expert-group get importance
+  1.0.  Inputs to the second-choice expert group get importance 0.5.
+  Inputs that represent padding get importance 0.0.
+
   Args:
     inputs: a mtf.Tensor with shape [<batch_dims>, group_size_dim, input_dim]
     outer_expert_dims: an optional list of dimensions.  This is for the case
       where we are at an inner level of a hierarchical MoE.
     experts_dim: a Dimension (the number of experts)
     expert_capacity_dim: a Dimension (number of examples per group per expert)
-    max_experts: optional mtf.Tensor with shape [<batch_dims>, group_size_dim]
     hparams: model hyperparameters.
     train: a boolean
+    importance: an optional tensor with shape [<batch_dims>, group_size_dim]
 
   Returns:
     dispatch_tensor: a Tensor with shape
@@ -469,26 +482,29 @@ def _top_2_gating(
   index_1, gate_1 = mtf.top_1(raw_gates, experts_dim)
   # [batch, group, experts]
   mask_1 = mtf.one_hot(index_1, experts_dim, dtype=raw_gates.dtype)
+  density_1_proxy = raw_gates
+  if importance is not None:
+    mask_1 *= mtf.to_float(mtf.equal(importance, 1.0))
+    gate_1 *= mtf.to_float(mtf.equal(importance, 1.0))
+    density_1_proxy *= mtf.to_float(mtf.equal(importance, 1.0))
   gates_without_top_1 = raw_gates * (1.0 - mask_1)
   # [batch, group]
   index_2, gate_2 = mtf.top_1(gates_without_top_1, experts_dim)
   # [batch, group, experts]
   mask_2 = mtf.one_hot(index_2, experts_dim, dtype=raw_gates.dtype)
+  if importance is not None:
+    mask_2 *= mtf.to_float(mtf.greater(importance, 0.0))
 
-  if max_experts is not None:
-    geq1 = mtf.to_float(mtf.greater_equal(max_experts, 1.0))
-    geq2 = mtf.to_float(mtf.greater_equal(max_experts, 2.0))
-    mask_1 *= geq1
-    mask_2 *= geq2
-    raw_gates *= geq1
-    gates_without_top_1 *= geq2
+  denom = gate_1 + gate_2 + 1e-9
+  gate_1 /= denom
+  gate_2 /= denom
 
   # BALANCING LOSSES
   # shape = [batch, experts]
   # We want to equalize the fraction of the batch assigned to each expert
   density_1 = mtf.reduce_mean(mask_1, reduced_dim=group_size_dim)
   # Something continuous that is correlated with what we want to equalize.
-  density_1_proxy = mtf.reduce_mean(raw_gates, reduced_dim=group_size_dim)
+  density_1_proxy = mtf.reduce_mean(density_1_proxy, reduced_dim=group_size_dim)
   density_1 = mtf.Print(
       density_1, [mtf.reduce_mean(density_1, output_shape=[experts_dim])],
       "density_1", summarize=1000)
@@ -566,11 +582,6 @@ def _top_2_gating(
   position_in_expert_2 = mtf.reduce_sum(
       position_in_expert_2, reduced_dim=experts_dim)
 
-  # renormalize the two gate values to add up to 1
-  denom = gate_1 + gate_2 + 1e-9
-  gate_1 /= denom
-  gate_2 /= denom
-
   # [batch, group, experts, expert_capacity]
   combine_tensor = (
       gate_1 * mask_1_flat
@@ -597,6 +608,7 @@ def set_default_moe_hparams(hparams):
   # moe_capacity_factor_* should be set to a value >=1.
   hparams.add_hparam("moe_capacity_factor_train", 1.25)
   hparams.add_hparam("moe_capacity_factor_eval", 2.0)
+  hparams.add_hparam("moe_capacity_factor_second_level", 1.0)
   # Each expert has a hidden layer with this size.
   hparams.add_hparam("moe_hidden_size", 4096)
   # For gating, divide inputs into groups of this size before gating.

From 84cfac76bd93196803882018e0c65ccb1127b757 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 5 Oct 2018 16:03:28 -0700
Subject: [PATCH 0970/2720] Rm eval summary hooks (and tensor hook) in case of
 TPUs.

PiperOrigin-RevId: 215981503
---
 tensor2tensor/models/research/transformer_vae.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index d47812db6..e10074836 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -724,6 +724,10 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
     """Constructs `tf.estimator.EstimatorSpec` for EVAL (evaluation) mode."""
     estimator_spec = super(TransformerAE, self).estimator_spec_eval(
         features, logits, labels, loss, losses_dict)
+    if common_layers.is_xla_compiled():
+      # For TPUs (and XLA more broadly?), do not add summary hooks that depend
+      # on losses; they are not supported.
+      return estimator_spec
 
     summary_op = tf.get_collection(tf.GraphKeys.SUMMARIES, scope="losses")
     summary_op.extend(tf.get_collection(tf.GraphKeys.SUMMARIES, scope="loss"))

From dd717533277c4ed5430c4b2dc64cce77c2dcf796 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 8 Oct 2018 13:58:39 -0700
Subject: [PATCH 0971/2720] This fixes an important bug in the computation of
 video metrics.

PiperOrigin-RevId: 216244776
---
 .../data_generators/video_utils_test.py       |  29 +++++
 tensor2tensor/mesh_tensorflow/OWNERS          |   2 -
 tensor2tensor/models/OWNERS                   |   6 -
 tensor2tensor/rl/OWNERS                       |   3 -
 tensor2tensor/utils/video_metrics.py          | 119 +++++++++++++-----
 tensor2tensor/utils/video_metrics_test.py     |  56 +++++++++
 6 files changed, 174 insertions(+), 41 deletions(-)
 delete mode 100644 tensor2tensor/mesh_tensorflow/OWNERS
 delete mode 100644 tensor2tensor/models/OWNERS
 delete mode 100644 tensor2tensor/rl/OWNERS
 create mode 100644 tensor2tensor/utils/video_metrics_test.py

diff --git a/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index 7be6a6a7c..9ef46b65d 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -19,14 +19,43 @@
 from __future__ import division
 from __future__ import print_function
 import numpy as np
+from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.utils import decoding
+from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
 
 class VideoUtilsTest(tf.test.TestCase):
 
+  def getPredictions(self):
+    rng = np.random.RandomState(0)
+    # num_samples=4
+    inputs = rng.randint(0, 255, (4, 2, 64, 64, 3))
+    outputs = rng.randint(0, 255, (4, 5, 64, 64, 3))
+    targets = rng.randint(0, 255, (4, 5, 64, 64, 3))
+    predictions = []
+    for input_, output, target in zip(inputs, outputs, targets):
+      curr_pred = {"inputs": input_, "outputs": output, "targets": target}
+      predictions.append(curr_pred)
+
+    # num_decodes=2
+    predictions = [predictions] * 2
+    problem = registry.problem("video_stochastic_shapes10k")
+    return predictions, problem
+
+  def testDecodeInMemoryTrue(self):
+    predictions, problem = self.getPredictions()
+    decode_hparams = decoding.decode_hparams()
+    decode_hparams.decode_in_memory = True
+    decode_hooks = decoding.DecodeHookArgs(
+        estimator=None, problem=problem, output_dirs=None,
+        hparams=decode_hparams, decode_hparams=decode_hparams,
+        predictions=predictions)
+    metrics = video_utils.summarize_video_metrics(decode_hooks)
+    self.assertEqual(len(metrics), 40)
+
   def testConvertPredictionsToVideoSummaries(self):
     # Initialize predictions.
     rng = np.random.RandomState(0)
diff --git a/tensor2tensor/mesh_tensorflow/OWNERS b/tensor2tensor/mesh_tensorflow/OWNERS
deleted file mode 100644
index 681d09a90..000000000
--- a/tensor2tensor/mesh_tensorflow/OWNERS
+++ /dev/null
@@ -1,2 +0,0 @@
-nikip
-noam
diff --git a/tensor2tensor/models/OWNERS b/tensor2tensor/models/OWNERS
deleted file mode 100644
index 4c18f708e..000000000
--- a/tensor2tensor/models/OWNERS
+++ /dev/null
@@ -1,6 +0,0 @@
-avaswani
-dumitru
-mbz
-nikip
-noam
-uszkoreit
diff --git a/tensor2tensor/rl/OWNERS b/tensor2tensor/rl/OWNERS
deleted file mode 100644
index cfa1f4f74..000000000
--- a/tensor2tensor/rl/OWNERS
+++ /dev/null
@@ -1,3 +0,0 @@
-blazej
-dumitru
-mbz
diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index 3f5cbdf01..92fba7d42 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -88,13 +88,21 @@ def save_results(results, output_dir, problem_name):
       np.save(fname, array)
 
 
-def compute_metrics(output_video, target_video):
-  max_pixel_value = 255.0
-  output_video = tf.to_float(output_video)
-  target_video = tf.to_float(target_video)
-  psnr = tf.image.psnr(output_video, target_video, max_pixel_value)
-  ssim = tf.image.ssim(output_video, target_video, max_pixel_value)
-  return {"PSNR": psnr, "SSIM": ssim}
+def psnr_and_ssim(output, target):
+  """Compute the PSNR and SSIM.
+
+  Args:
+    output: 4-D Tensor, shape=(num_frames, height, width, num_channels)
+    target: 4-D Tensor, shape=(num_frames, height, width, num_channels)
+  Returns:
+    psnr: 1-D Tensor, shape=(num_frames,)
+    ssim: 1-D Tensor, shape=(num_frames,)
+  """
+  output = tf.cast(output, dtype=tf.int32)
+  target = tf.cast(target, dtype=tf.int32)
+  psnr = tf.image.psnr(output, target, max_val=255)
+  ssim = tf.image.ssim(output, target, max_val=255)
+  return psnr, ssim
 
 
 def stack_data_given_key(predictions, key):
@@ -116,7 +124,6 @@ def get_zipped_dataset_from_predictions(predictions):
   iterator = dataset.make_initializable_iterator()
   feed_dict = {targets_placeholder: targets,
                outputs_placeholder: outputs}
-
   return iterator, feed_dict, num_videos
 
 
@@ -129,13 +136,11 @@ def compute_one_decoding_video_metrics(iterator, feed_dict, num_videos):
     num_videos: number of videos.
 
   Returns:
-    Dictionary which contains the average of each metric per frame.
+    all_psnr: 2-D Numpy array, shape=(num_samples, num_frames)
+    all_ssim: 2-D Numpy array, shape=(num_samples, num_frames)
   """
   output, target = iterator.get_next()
-
-  metrics_dict = compute_metrics(output, target)
-  metrics_names, metrics = zip(*six.iteritems(metrics_dict))
-  means, update_ops = tf.metrics.mean_tensor(metrics)
+  metrics = psnr_and_ssim(output, target)
 
   with tf.Session() as sess:
     sess.run(tf.local_variables_initializer())
@@ -143,34 +148,82 @@ def compute_one_decoding_video_metrics(iterator, feed_dict, num_videos):
     if initalizer is not None:
       sess.run(initalizer, feed_dict=feed_dict)
 
-    # Compute mean over dataset
+    all_psnr, all_ssim = [], []
     for i in range(num_videos):
       print("Computing video: %d" % i)
-      sess.run(update_ops)
-    averaged_metrics = sess.run(means)
+      psnr_np, ssim_np = sess.run(metrics)
+      all_psnr.append(psnr_np)
+      all_ssim.append(ssim_np)
+    all_psnr = np.array(all_psnr)
+    all_ssim = np.array(all_ssim)
+    return all_psnr, all_ssim
+
+
+def reduce_to_best_decode(metrics, reduce_func):
+  """Extracts the best-decode from the metrics according to reduce_func.
+
+  Args:
+    metrics: 3-D numpy array, shape=(num_decodes, num_samples, num_frames)
+    reduce_func: callable, np.argmax or np.argmin.
+  Returns:
+    best_metrics: 2-D numpy array, shape=(num_samples, num_frames).
+  """
+  num_videos = metrics.shape[1]
+  # Take mean of the metric across the frames to approximate the video
+  # closest to the ground truth.
+  mean_across_frames = np.mean(metrics, axis=-1)
 
-    results = dict(zip(metrics_names, averaged_metrics))
-    return results
+  # For every sample, use the decode that has a maximum mean-metric.
+  best_decode_ind = reduce_func(mean_across_frames, axis=0)
+  return metrics[best_decode_ind, np.arange(num_videos), :]
 
 
 def compute_all_metrics_statistics(all_results):
-  """Computes statistics of metrics across multiple decodings."""
+  """Computes statistics of metrics across multiple decodings.
+
+  Args:
+    all_results: dicf of 3-D numpy arrays.
+                 Each array has shape=(num_decodes, num_samples, num_frames).
+  Returns:
+    statistics: dict of 1-D numpy arrays shape=(num_frames).
+                First the statistic (max/mean/std) is computed across the
+                decodes, then the mean is taken across num_samples.
+  """
   statistics = {}
-  for key in all_results[0].keys():
-    values = [result[key] for result in all_results]
-    values = np.vstack(values)
+  all_metrics = all_results.keys()
+
+  for key in all_metrics:
+    values = all_results[key]
     statistics[key + "_MEAN"] = np.mean(values, axis=0)
     statistics[key + "_STD"] = np.std(values, axis=0)
-    statistics[key + "_MIN"] = np.min(values, axis=0)
-    statistics[key + "_MAX"] = np.max(values, axis=0)
+    statistics[key + "_MIN"] = reduce_to_best_decode(values, np.argmin)
+    statistics[key + "_MAX"] = reduce_to_best_decode(values, np.argmax)
+
+  # Computes mean of each statistic across the dataset.
+  for key in statistics:
+    statistics[key] = np.mean(statistics[key], axis=0)
   return statistics
 
 
 def compute_video_metrics_from_predictions(predictions):
-  all_results = []
-  for prediction in predictions:
-    args = get_zipped_dataset_from_predictions(prediction)
-    all_results.append(compute_one_decoding_video_metrics(*args))
+  """Computes metrics from predictions.
+
+  Args:
+    predictions: list of list of dicts.
+                 outer length: num_decodes, inner_length: num_samples
+  Returns:
+    statistics: dict of Tensors, key being the metric with each Tensor
+                having the shape (num_samples, num_frames).
+  """
+  ssim_all_decodes, psnr_all_decodes = [], []
+  for single_decode in predictions:
+    args = get_zipped_dataset_from_predictions(single_decode)
+    psnr_single, ssim_single = compute_one_decoding_video_metrics(*args)
+    psnr_all_decodes.append(psnr_single)
+    ssim_all_decodes.append(ssim_single)
+  psnr_all_decodes = np.array(psnr_all_decodes)
+  ssim_all_decodes = np.array(ssim_all_decodes)
+  all_results = {"PSNR": psnr_all_decodes, "SSIM": ssim_all_decodes}
   statistics = compute_all_metrics_statistics(all_results)
   return statistics
 
@@ -192,13 +245,19 @@ def compute_video_metrics_from_png_files(
   Returns:
     Dictionary which contains the average of each metric per frame.
   """
-  all_results = []
+  ssim_all_decodes, psnr_all_decodes = [], []
   for output_dir in output_dirs:
     output_files, target_files = get_target_and_output_filepatterns(
         output_dir, problem_name)
     args = get_zipped_dataset_from_png_files(
         output_files, target_files, video_length, frame_shape)
-    all_results.append(compute_one_decoding_video_metrics(*args))
+    psnr_single, ssim_single = compute_one_decoding_video_metrics(*args)
+    psnr_all_decodes.append(psnr_single)
+    ssim_all_decodes.append(ssim_single)
+
+  psnr_all_decodes = np.array(psnr_all_decodes)
+  ssim_all_decodes = np.array(ssim_all_decodes)
+  all_results = {"PSNR": psnr_all_decodes, "SSIM": ssim_all_decodes}
   statistics = compute_all_metrics_statistics(all_results)
   return statistics, all_results
 
diff --git a/tensor2tensor/utils/video_metrics_test.py b/tensor2tensor/utils/video_metrics_test.py
new file mode 100644
index 000000000..5bf01965a
--- /dev/null
+++ b/tensor2tensor/utils/video_metrics_test.py
@@ -0,0 +1,56 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""video metrics test."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.utils import video_metrics
+import tensorflow as tf
+
+
+class VideoMetricsTest(tf.test.TestCase):
+
+  def test_reduce_to_best_decode(self):
+    # num_decodes=2, num_samples=3, num_frames=4
+    decode1 = [
+        [30.0, 32.0, 33.0, 34.0],
+        [22.0, 19.0, 12.0, 13.0],
+        [30.0, 10.0, 30.0, 10.0]]
+    decode2 = [
+        [22.0, 19.0, 12.0, 13.0],
+        [30.0, 32.0, 33.0, 34.0],
+        [25.0, 25.0, 25.0, 25.0]]
+    all_decodes = [decode1, decode2]
+    all_decodes = np.array(all_decodes)
+    best_decode = video_metrics.reduce_to_best_decode(all_decodes, np.argmax)
+    worst_decode = video_metrics.reduce_to_best_decode(all_decodes, np.argmin)
+    exp_best_decode = [
+        [30.0, 32.0, 33.0, 34.0],
+        [30.0, 32.0, 33.0, 34.0],
+        [25.0, 25.0, 25.0, 25.0]]
+    exp_worst_decode = [
+        [22.0, 19.0, 12.0, 13.0],
+        [22.0, 19.0, 12.0, 13.0],
+        [30.0, 10.0, 30.0, 10.0]]
+    self.assertTrue(np.allclose(best_decode, exp_best_decode))
+    self.assertTrue(np.allclose(worst_decode, exp_worst_decode))
+
+
+if __name__ == '__main__':
+  tf.test.main()

From c92c3a188e785437f18d43737d1dbfe20b132ffb Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 9 Oct 2018 01:12:32 +0200
Subject: [PATCH 0972/2720] RL pipeline refactoring: T2TGymEnv (#1119)

* Create T2TEnv - a base class for gym envs

* Implement T2TGymEnv

* Integrate T2TGymEnv with the RL pipeline
---
 tensor2tensor/data_generators/gym_env.py   | 207 +++++++++++++++++++++
 tensor2tensor/rl/envs/batch_env_factory.py |   4 +-
 tensor2tensor/rl/envs/py_func_batch_env.py |  11 +-
 3 files changed, 217 insertions(+), 5 deletions(-)
 create mode 100644 tensor2tensor/data_generators/gym_env.py

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
new file mode 100644
index 000000000..545b373b7
--- /dev/null
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -0,0 +1,207 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""RL environments."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+
+import numpy as np
+
+
+Frame = collections.namedtuple(
+    "Frame", ("observation", "action", "reward", "done")
+)
+
+
+class T2TEnv(object):
+  """Abstract class representing a batch of environments.
+
+  Attributes:
+    history: List of finished rollouts, where rollout is a list of Frames.
+    batch_size: Number of environments played simultaneously.
+    observation_space: Gym observation space. Should be overridden in derived
+      classes.
+    action_space: Gym observation space. Should be overridden in derived
+      classes.
+
+  Args:
+    batch_size: Number of environments in a batch.
+  """
+
+  observation_space = None
+  action_space = None
+
+  def __init__(self, batch_size):
+    self.clear_history()
+    self.batch_size = batch_size
+    self._current_rollouts = [[] for _ in range(batch_size)]
+    self._current_observations = [None for _ in range(batch_size)]
+
+  def __str__(self):
+    """Returns a string representation of the environment for debug purposes."""
+    raise NotImplementedError
+
+  def clear_history(self):
+    """Clears the rollout history."""
+    self.history = []
+
+  def _preprocess_observations(self, obs):
+    """Transforms a batch of observations.
+
+    Can be overridden in derived classes.
+
+    Args:
+      obs: A batch of observations.
+
+    Returns:
+      Transformed batch of observations.
+    """
+    return obs
+
+  def _preprocess_rewards(self, rewards):
+    """Transforms a batch of rewards.
+
+    Can be overridden in derived classes.
+
+    Args:
+      rewards: A batch of rewards.
+
+    Returns:
+      Transformed batch of rewards.
+    """
+    return rewards
+
+  def _step(self, actions):
+    """Makes a step in all environments.
+
+    Should be overridden in derived classes.
+
+    Should not do any preprocessing of the observations and rewards; this
+    should be done in _preprocess_*.
+
+    Args: see step().
+
+    Returns: see step().
+    """
+    raise NotImplementedError
+
+  def step(self, actions):
+    """Makes a step in all environments.
+
+    Does any preprocessing and records frames.
+
+    Args:
+      actions: Batch of actions.
+
+    Returns:
+      (obs, rewards, dones) - batches of observations, rewards and done flags
+      respectively.
+    """
+    (obs, rewards, dones) = self._step(actions)
+    obs = self._preprocess_observations(obs)
+    rewards = self._preprocess_rewards(rewards)
+    # oard = (observation, action, reward, done)
+    for (rollout, oard) in zip(self._current_rollouts, zip(
+        self._current_observations, actions, rewards, dones
+    )):
+      rollout.append(Frame(*oard))
+    self._current_observations = obs
+    return (obs, rewards, dones)
+
+  def _reset(self, indices):
+    """Resets environments at given indices.
+
+    Args: see reset().
+
+    Returns: see reset().
+    """
+    raise NotImplementedError
+
+  def reset(self, indices=None):
+    """Resets environments at given indices.
+
+    Does any preprocessing and adds finished rollouts to history.
+
+    Args:
+      indices: Indices of environments to reset.
+
+    Returns:
+      Batch of initial observations of reset environments.
+    """
+    if indices is None:
+      indices = np.arange(self.batch_size)
+    new_obs = self._reset(indices)
+    new_obs = self._preprocess_observations(new_obs)
+    for (index, ob) in zip(indices, new_obs):
+      rollout = self._current_rollouts[index]
+      if rollout and rollout[-1].done:
+        self.history.append(rollout)
+        self._current_rollouts[index] = []
+      self._current_observations[index] = ob
+    return new_obs
+
+  def close(self):
+    """Cleanups any resources.
+
+    Can be overridden in derived classes.
+    """
+    pass
+
+
+class T2TGymEnv(T2TEnv):
+  """Class representing a batch of Gym environments."""
+
+  def __init__(self, envs):
+    super(T2TGymEnv, self).__init__(len(envs))
+
+    if not envs:
+      raise ValueError("Must have at least one environment.")
+    self._envs = envs
+
+    self.observation_space = envs[0].observation_space
+    if not all(env.observation_space == self.observation_space for env in envs):
+      raise ValueError("All environments must use the same observation space.")
+
+    self.action_space = envs[0].action_space
+    if not all(env.action_space == self.action_space for env in envs):
+      raise ValueError("All environments must use the same action space.")
+
+  def __str__(self):
+    return "T2TGymEnv(%s)" % ", ".join([str(env) for env in self._envs])
+
+  def _preprocess_observations(self, obs):
+    # TODO(lukaszkaiser): Implement.
+    return obs
+
+  def _preprocess_rewards(self, rewards):
+    # TODO(lukaszkaiser): Implement.
+    return rewards
+
+  def _step(self, actions):
+    (obs, rewards, dones, _) = zip(*[
+        env.step(action) for (env, action) in zip(self._envs, actions)
+    ])
+    return tuple(map(np.stack, (obs, rewards, dones)))
+
+  def _reset(self, indices):
+    return np.stack([self._envs[index].reset() for index in indices])
+
+  def close(self):
+    for env in self._envs:
+      env.close()
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index 0b71df225..eabdb6e85 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -32,7 +32,7 @@
 import sys
 import traceback
 
-from tensor2tensor.rl.envs import batch_env
+from tensor2tensor.data_generators import gym_env
 from tensor2tensor.rl.envs import py_func_batch_env
 from tensor2tensor.rl.envs import simulated_batch_env
 
@@ -58,7 +58,7 @@ def _define_batch_env(environment_spec, num_agents, xvfb=False):
     envs = [
         ExternalProcessEnv(environment_spec.env_lambda, xvfb)
         for _ in range(num_agents)]
-    env = batch_env.BatchEnv(envs, blocking=False)
+    env = gym_env.T2TGymEnv(envs)
     env = py_func_batch_env.PyFuncBatchEnv(env)
     return env
 
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index 5bc3bdfb6..5f11fe95d 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -22,6 +22,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import numpy as np
+
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
 import tensorflow as tf
 
@@ -45,7 +47,7 @@ def __init__(self, batch_env):
     self._batch_env = batch_env
     with tf.variable_scope("env_temporary"):
       self._observ = tf.Variable(
-          tf.zeros((len(self._batch_env),) + self.observ_shape,
+          tf.zeros((self._batch_env.batch_size,) + self.observ_shape,
                    self.observ_dtype),
           name="observ", trainable=False)
 
@@ -68,7 +70,7 @@ def initialize(self, sess):
 
   def __len__(self):
     """Number of combined environments."""
-    return len(self._batch_env)
+    return self._batch_env.batch_size
 
   def __getitem__(self, index):
     """Access an underlying environment by index."""
@@ -88,8 +90,11 @@ def simulate(self, action):
     with tf.name_scope("environment/simulate"):
       if action.dtype in (tf.float16, tf.float32, tf.float64):
         action = tf.check_numerics(action, "action")
+      def step(action):
+        (observ, reward, done) = self._batch_env.step(action)
+        return (observ, reward.astype(np.float32), done)
       observ, reward, done = tf.py_func(
-          lambda a: self._batch_env.step(a)[:3], [action],
+          step, [action],
           [self.observ_dtype, tf.float32, tf.bool], name="step")
       reward = tf.check_numerics(reward, "reward")
       reward.set_shape((len(self),))

From 1e81346d8ed41e6c6b54fbbef13d755eac67802a Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Mon, 8 Oct 2018 16:13:09 -0700
Subject: [PATCH 0973/2720] internal merge of PR #1119

PiperOrigin-RevId: 216268473
---
 tensor2tensor/data_generators/gym_env.py | 20 ++++++++++++--------
 tensor2tensor/rl/rl_trainer_lib_test.py  | 21 ---------------------
 2 files changed, 12 insertions(+), 29 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 545b373b7..6ab672e11 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -37,8 +37,7 @@ class T2TEnv(object):
     batch_size: Number of environments played simultaneously.
     observation_space: Gym observation space. Should be overridden in derived
       classes.
-    action_space: Gym observation space. Should be overridden in derived
-      classes.
+    action_space: Gym action space. Should be overridden in derived classes.
 
   Args:
     batch_size: Number of environments in a batch.
@@ -88,16 +87,19 @@ def _preprocess_rewards(self, rewards):
     return rewards
 
   def _step(self, actions):
-    """Makes a step in all environments.
+    """Makes a step in all environments without recording history.
 
     Should be overridden in derived classes.
 
     Should not do any preprocessing of the observations and rewards; this
     should be done in _preprocess_*.
 
-    Args: see step().
+    Args:
+      actions: Batch of actions.
 
-    Returns: see step().
+    Returns:
+      (obs, rewards, dones) - batches of observations, rewards and done flags
+      respectively.
     """
     raise NotImplementedError
 
@@ -125,11 +127,13 @@ def step(self, actions):
     return (obs, rewards, dones)
 
   def _reset(self, indices):
-    """Resets environments at given indices.
+    """Resets environments at given indices without recording history.
 
-    Args: see reset().
+    Args:
+      indices: Indices of environments to reset.
 
-    Returns: see reset().
+    Returns:
+      Batch of initial observations of reset environments.
     """
     raise NotImplementedError
 
diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/rl_trainer_lib_test.py
index bf777c39f..01a4f8df9 100644
--- a/tensor2tensor/rl/rl_trainer_lib_test.py
+++ b/tensor2tensor/rl/rl_trainer_lib_test.py
@@ -18,35 +18,14 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.models.research import rl as rl_models
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.utils import registry
-from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
 
 
 class TrainTest(tf.test.TestCase):
 
-  test_config = ("epochs_num=4,eval_every_epochs=0,video_during_eval=False,"
-                 "num_agents=5,optimization_epochs=5,epoch_length=50")
-
-  def test_no_crash_pendulum(self):
-    hparams = trainer_lib.create_hparams(
-        "ppo_continuous_action_base",
-        TrainTest.test_config)
-    hparams.add_hparam(
-        "environment_spec", rl_models.simple_gym_spec("Pendulum-v0"))
-    rl_trainer_lib.train(hparams)
-
-  def test_no_crash_cartpole(self):
-    hparams = trainer_lib.create_hparams(
-        "ppo_discrete_action_base",
-        TrainTest.test_config)
-    hparams.add_hparam(
-        "environment_spec", rl_models.simple_gym_spec("CartPole-v0"))
-    rl_trainer_lib.train(hparams)
-
   def test_train_pong(self):
     hparams = registry.hparams("pong_model_free")
     hparams.epochs_num = 2

From b24320587e29d0b6cc6ecbdf2fa0bfe7f18d9895 Mon Sep 17 00:00:00 2001
From: Youlong Cheng <ylc@google.com>
Date: Mon, 8 Oct 2018 18:15:28 -0700
Subject: [PATCH 0974/2720] Fix Python 3 compatibility issue in MeshTensorflow

PiperOrigin-RevId: 216284957
---
 tensor2tensor/mesh_tensorflow/mtf_transformer.py | 5 ++++-
 tensor2tensor/mesh_tensorflow/simd_mesh_impl.py  | 2 ++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer.py b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
index 576f9061d..f889be591 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_transformer.py
@@ -15,11 +15,14 @@
 
 """Transformer model."""
 
-
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
 import copy
+
+from six.moves import xrange  # pylint: disable=redefined-builtin
+
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
diff --git a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
index 9f723b99f..bebbff22d 100644
--- a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+from six.moves import xrange  # pylint: disable=redefined-builtin
+
 from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
 from tensor2tensor.mesh_tensorflow import mtf_utils
 from tensor2tensor.mesh_tensorflow import tpu_variables

From 6da0ee25e84177f639f6dd109440c1ae34130aef Mon Sep 17 00:00:00 2001
From: Youlong Cheng <ylc@google.com>
Date: Mon, 8 Oct 2018 21:39:22 -0700
Subject: [PATCH 0975/2720] Add instruction for the toy_model.

PiperOrigin-RevId: 216300889
---
 tensor2tensor/mesh_tensorflow/README.md | 39 ++++++++++++++++++++++---
 1 file changed, 35 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/README.md b/tensor2tensor/mesh_tensorflow/README.md
index d3d9acf71..a74a0014a 100644
--- a/tensor2tensor/mesh_tensorflow/README.md
+++ b/tensor2tensor/mesh_tensorflow/README.md
@@ -281,9 +281,7 @@ the number of cores.  The differences between cores are as follows:
 
 # Instructions for running on cloud-tpu
 
-Note: It will be available in `tensorflow>=1.11.0`. For early adoption, use
-`tf-nightly`. Please contact the Google Cloud TPU team if you need to obtain
-`tf-nightly`.
+Note: It requires `tensorflow>=1.11.0`.
 
 ## Prerequisite
 
@@ -304,7 +302,7 @@ cd tensor2tensor/
 pip install --user .
 ```
 
-## Run the model
+## Run the Transfomer model with Tensor2Tensor config
 
 Before run the model, you need to prepare the training data and bucket for
 storing checkpoints. Refer to the
@@ -333,6 +331,39 @@ tensor2tensor/bin/t2t-trainer \
   --cloud_tpu_name=$TPU_NAME
 ```
 
+
+## Run the toy model without Tensor2Tensor dependencies
+
+  This toy model contains two fully-connected layers which aim to train a
+  identity function: f(x) = x. Since there are 8 TPU cores, we can arbitrary
+  change the FLAGS.mesh_shape and FLAGS.layout to achieve different
+  data-parallelism and model-parallelism strategies.
+
+```sh
+MODEL_DIR=gs://xxxx
+TPU_NAME=ylc-mtf-donut
+
+# 2 ways data-parallelism and 4 ways model-parallelism.
+# In this configuration, we split the batch dimension into 2 cores and the
+# hidden dimension into 4 cores.
+python mtf_toy_model_tpu.py \
+  --tpu=$TPU \
+  --model_dir=$MODEL_DIR \
+  --io_size=8 \
+  --hidden_size=8 \
+  --mesh_shape='x:2;y:4' \
+  --layout='batch:x;hidden:y'
+
+# 8 ways model-parallelism.
+# In this configuration, We split the hidden dimension into 8 cores.
+python mtf_toy_model_tpu.py \
+  --tpu=$TPU \
+  --model_dir=$MODEL_DIR \
+  --io_size=8 \
+  --hidden_size=8 \
+  --mesh_shape='all:8' \
+  --layout='hidden:all'
+
 # TODO LIST (please add items)
 
 We are actively working on improving Mesh-TensorFlow in a variety of ways.  Some

From db631cd3a7b590680a052b3cd501b52c97b09779 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Wed, 10 Oct 2018 02:49:07 +0200
Subject: [PATCH 0976/2720] RL pipeline refactoring: Turn T2TEnv into a Problem
 and enable generating data with it (#1124)

* Remove BatchEnv and ExternalProcessEnv - not needed anymore

* Encode observations as PNG

* Change ordering of elements within a frame

* Turn T2TEnv into a Problem and enable generating data with it

* Fix gym import error
---
 tensor2tensor/data_generators/gym_env.py      | 172 ++++++++++++--
 tensor2tensor/data_generators/gym_env_test.py |  88 +++++++
 tensor2tensor/rl/envs/batch_env.py            | 135 -----------
 tensor2tensor/rl/envs/batch_env_factory.py    | 218 +-----------------
 4 files changed, 247 insertions(+), 366 deletions(-)
 create mode 100644 tensor2tensor/data_generators/gym_env_test.py
 delete mode 100644 tensor2tensor/rl/envs/batch_env.py

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 6ab672e11..dba260ffb 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -20,16 +20,26 @@
 from __future__ import print_function
 
 import collections
+import math
+import random
 
 import numpy as np
 
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import video_utils
+from tensor2tensor.utils import metrics
+
+import tensorflow as tf
+
 
 Frame = collections.namedtuple(
-    "Frame", ("observation", "action", "reward", "done")
+    # Order of elements reflects time progression within a frame.
+    "Frame", ("observation", "reward", "done", "action")
 )
 
 
-class T2TEnv(object):
+class T2TEnv(video_utils.VideoProblem):
   """Abstract class representing a batch of environments.
 
   Attributes:
@@ -38,6 +48,10 @@ class T2TEnv(object):
     observation_space: Gym observation space. Should be overridden in derived
       classes.
     action_space: Gym action space. Should be overridden in derived classes.
+    reward_range: Tuple (min, max) representing the range of rewards. Limits
+      should be integer (discrete rewards).
+    name: Problem name for generating filenames. Should be overridden in
+      derived classes.
 
   Args:
     batch_size: Number of environments in a batch.
@@ -45,12 +59,21 @@ class T2TEnv(object):
 
   observation_space = None
   action_space = None
+  reward_range = (-1, 1)
+  name = None
 
   def __init__(self, batch_size):
+    super(T2TEnv, self).__init__()
+
     self.clear_history()
     self.batch_size = batch_size
     self._current_rollouts = [[] for _ in range(batch_size)]
-    self._current_observations = [None for _ in range(batch_size)]
+    self._current_frames = [None for _ in range(batch_size)]
+
+    with tf.Graph().as_default():
+      self._image_t = tf.placeholder(dtype=tf.uint8, shape=(None, None, None))
+      self._encoded_image_t = tf.image.encode_png(self._image_t)
+      self._encode_session = tf.Session()
 
   def __str__(self):
     """Returns a string representation of the environment for debug purposes."""
@@ -86,6 +109,15 @@ def _preprocess_rewards(self, rewards):
     """
     return rewards
 
+  def _encode_observations(self, observations):
+    """Encodes observations as PNG."""
+    return [
+        self._encode_session.run(
+            self._encoded_image_t, feed_dict={self._image_t: observation}
+        )
+        for observation in observations
+    ]
+
   def _step(self, actions):
     """Makes a step in all environments without recording history.
 
@@ -118,12 +150,17 @@ def step(self, actions):
     (obs, rewards, dones) = self._step(actions)
     obs = self._preprocess_observations(obs)
     rewards = self._preprocess_rewards(rewards)
-    # oard = (observation, action, reward, done)
-    for (rollout, oard) in zip(self._current_rollouts, zip(
-        self._current_observations, actions, rewards, dones
-    )):
-      rollout.append(Frame(*oard))
-    self._current_observations = obs
+    encoded_obs = self._encode_observations(obs)
+    for (rollout, frame, action) in zip(
+        self._current_rollouts, self._current_frames, actions
+    ):
+      rollout.append(frame._replace(action=action))
+
+    # ord = (observation, reward, done)
+    self._current_frames = [
+        Frame(*ord, action=None)
+        for ord in zip(encoded_obs, rewards, dones)
+    ]
     return (obs, rewards, dones)
 
   def _reset(self, indices):
@@ -152,12 +189,17 @@ def reset(self, indices=None):
       indices = np.arange(self.batch_size)
     new_obs = self._reset(indices)
     new_obs = self._preprocess_observations(new_obs)
-    for (index, ob) in zip(indices, new_obs):
-      rollout = self._current_rollouts[index]
-      if rollout and rollout[-1].done:
+    encoded_obs = self._encode_observations(new_obs)
+    for (index, ob) in zip(indices, encoded_obs):
+      frame = self._current_frames[index]
+      if frame is not None and frame.done:
+        rollout = self._current_rollouts[index]
+        rollout.append(frame._replace(action=0))
         self.history.append(rollout)
         self._current_rollouts[index] = []
-      self._current_observations[index] = ob
+      self._current_frames[index] = Frame(
+          observation=ob, reward=0, done=False, action=None
+      )
     return new_obs
 
   def close(self):
@@ -165,12 +207,114 @@ def close(self):
 
     Can be overridden in derived classes.
     """
-    pass
+    self._encode_session.close()
+
+  @property
+  def num_channels(self):
+    """Number of color channels in each frame."""
+    return 3
+
+  def eval_metrics(self):
+    eval_metrics = [
+        metrics.Metrics.ACC, metrics.Metrics.ACC_PER_SEQ,
+        metrics.Metrics.IMAGE_RMSE
+    ]
+    return eval_metrics
+
+  @property
+  def extra_reading_spec(self):
+    """Additional data fields to store on disk and their decoders."""
+    field_names = ("frame_number", "action", "reward", "done")
+    data_fields = {
+        name: tf.FixedLenFeature([1], tf.int64) for name in field_names
+    }
+    decoders = {
+        name: tf.contrib.slim.tfexample_decoder.Tensor(tensor_key=name)
+        for name in field_names
+    }
+    return (data_fields, decoders)
+
+  @property
+  def frame_height(self):
+    return self.observation_space.shape[0]
+
+  @property
+  def frame_width(self):
+    return self.observation_space.shape[1]
+
+  @property
+  def num_actions(self):
+    return self.action_space.n
+
+  @property
+  def num_rewards(self):
+    (min_reward, max_reward) = self.reward_range
+    return max_reward - min_reward + 1
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    def make_modality(name):
+      return {
+          "{}s".format(name): ("video", 256),
+          "{}_reward".format(name): ("symbol:weights_all", self.num_rewards),
+          "{}_action".format(name): ("symbol:weights_all", self.num_actions)
+      }
+    p.input_modality = make_modality("input")
+    p.target_modality = make_modality("target")
+    p.input_space_id = problem.SpaceID.IMAGE
+    p.target_space_id = problem.SpaceID.IMAGE
+
+  def _generate_frames(self, rollouts):
+    for rollout in rollouts:
+      for (frame_number, frame) in enumerate(rollout):
+        yield {
+            "frame_number": [frame_number],
+            "image/encoded": [frame.observation],
+            "image/format": ["png"],
+            "image/height": [self.frame_height],
+            "image/width": [self.frame_width],
+            "action": [int(frame.action)],
+            "reward": [int(frame.reward - self.reward_range[0])],
+            "done": [int(frame.done)]
+        }
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    """Saves the rollout history to disk."""
+    # Suffle rollouts globally taking advantage of the fact that we have
+    # everything in memory.
+    shuffled_history = self.history[:]
+    random.shuffle(shuffled_history)
+
+    filepath_fns = {
+        problem.DatasetSplit.TRAIN: self.training_filepaths,
+        problem.DatasetSplit.EVAL: self.dev_filepaths,
+        problem.DatasetSplit.TEST: self.test_filepaths,
+    }
+
+    # We set shuffled=True as we don't want to shuffle on disk later.
+    splits_and_paths = [
+        (split["split"], path)
+        for split in self.dataset_splits
+        for path in filepath_fns[split["split"]](
+            data_dir, split["shards"], shuffled=True
+        )
+    ]
+
+    # Split entire rollouts into shards so that no rollout is broken on shard
+    # boundary.
+    shard_size = int(math.ceil(len(shuffled_history)) / len(splits_and_paths))
+    for (i, (split, path)) in enumerate(splits_and_paths):
+      rollouts = shuffled_history[i * shard_size : (i + 1) * shard_size]
+      generator_utils.generate_files(
+          self._generate_frames(rollouts), [path], cycle_every_n=float("inf")
+      )
 
 
 class T2TGymEnv(T2TEnv):
   """Class representing a batch of Gym environments."""
 
+  name = "t2t_gym_env"
+
   def __init__(self, envs):
     super(T2TGymEnv, self).__init__(len(envs))
 
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
new file mode 100644
index 000000000..241fa4389
--- /dev/null
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -0,0 +1,88 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Gym env tests."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import shutil
+
+import gym
+from gym.spaces import Box, Discrete
+import numpy as np
+
+from tensor2tensor.data_generators import gym_env
+
+import tensorflow as tf
+
+
+class TestEnv(gym.Env):
+  """Test environment.
+
+  Odd frames are "done".
+  """
+
+  action_space = Discrete(1)
+  observation_space = Box(
+      low=0, high=255, shape=(2, 2, 1), dtype=np.uint8
+  )
+
+  def __init__(self):
+    self._counter = 0
+
+  def _generate_ob(self):
+    return np.zeros(
+        self.observation_space.shape, self.observation_space.dtype
+    )
+
+  def step(self, action):
+    done = self._counter % 2 == 1
+    self._counter += 1
+    return (self._generate_ob(), 0, done, {})
+
+  def reset(self):
+    return self._generate_ob()
+
+
+class GymEnvTest(tf.test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    cls.out_dir = tf.test.get_temp_dir()
+    shutil.rmtree(cls.out_dir)
+    os.mkdir(cls.out_dir)
+
+  def test_generates(self):
+    env = gym_env.T2TGymEnv([TestEnv(), TestEnv()])
+    env.reset()
+    for _ in range(20):
+      (_, _, dones) = env.step([0, 0])
+      for (i, done) in enumerate(dones):
+        if done:
+          env.reset([i])
+    env.generate_data(self.out_dir, tmp_dir=None)
+
+    filenames = os.listdir(self.out_dir)
+    self.assertTrue(filenames)
+    path = os.path.join(self.out_dir, filenames[0])
+    records = list(tf.python_io.tf_record_iterator(path))
+    self.assertTrue(records)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/rl/envs/batch_env.py b/tensor2tensor/rl/envs/batch_env.py
deleted file mode 100644
index 7b78d0231..000000000
--- a/tensor2tensor/rl/envs/batch_env.py
+++ /dev/null
@@ -1,135 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Combine multiple environments to step them in batch."""
-
-# The code was based on Danijar Hafner's code from tf.agents:
-# https://github.com/tensorflow/agents/blob/master/agents/tools/batch_env.py
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-
-
-class BatchEnv(object):
-  """Combine multiple environments to step them in batch."""
-
-  def __init__(self, envs, blocking):
-    """Combine multiple environments to step them in batch.
-
-    To step environments in parallel, environments must support a
-    `blocking=False` argument to their step and reset functions that makes them
-    return callables instead to receive the result at a later time.
-
-    Args:
-      envs: List of environments.
-      blocking: Step environments after another rather than in parallel.
-
-    Raises:
-      ValueError: Environments have different observation or action spaces.
-    """
-    self._envs = envs
-    self._blocking = blocking
-    self.observ_space = self._envs[0].observation_space
-    if not all(env.observation_space == self.observ_space
-               for env in self._envs):
-      raise ValueError("All environments must use the same observation space.")
-    self.action_space = self._envs[0].action_space
-    if not all(env.action_space == self.action_space for env in self._envs):
-      raise ValueError("All environments must use the same observation space.")
-
-  def __str__(self):
-    return "BatchEnv(%s)" % ", ".join([str(e) for e in self._envs])
-
-  def __len__(self):
-    """Number of combined environments."""
-    return len(self._envs)
-
-  def __getitem__(self, index):
-    """Access an underlying environment by index."""
-    return self._envs[index]
-
-  def __getattr__(self, name):
-    """Forward unimplemented attributes to one of the original environments.
-
-    Args:
-      name: Attribute that was accessed.
-
-    Returns:
-      Value behind the attribute name one of the wrapped environments.
-    """
-    return getattr(self._envs[0], name)
-
-  def step(self, actions):
-    """Forward a batch of actions to the wrapped environments.
-
-    Args:
-      actions: Batched action to apply to the environment.
-
-    Raises:
-      ValueError: Invalid actions.
-
-    Returns:
-      Batch of observations, rewards, and done flags.
-    """
-    for index, (env, action) in enumerate(zip(self._envs, actions)):
-      if not env.action_space.contains(action):
-        message = "Invalid action at index {}: {}"
-        raise ValueError(message.format(index, action))
-    if self._blocking:
-      transitions = [
-          env.step(action)
-          for env, action in zip(self._envs, actions)]
-    else:
-      transitions = [
-          env.step(action, blocking=False)
-          for env, action in zip(self._envs, actions)]
-      transitions = [transition() for transition in transitions]
-    observs, rewards, dones, infos = zip(*transitions)
-
-    observ = np.stack(observs).astype(self.observ_space.dtype)
-    # TODO(piotrmilos): Do we really want cast to float32
-    reward = np.stack(rewards).astype(np.float32)
-    done = np.stack(dones)
-    info = tuple(infos)
-    return observ, reward, done, info
-
-  def reset(self, indices=None):
-    """Reset the environment and convert the resulting observation.
-
-    Args:
-      indices: The batch indices of environments to reset; defaults to all.
-
-    Returns:
-      Batch of observations.
-    """
-    if indices is None:
-      indices = np.arange(len(self._envs))
-    if self._blocking:
-      observs = [self._envs[index].reset() for index in indices]
-    else:
-      observs = [self._envs[index].reset(blocking=False) for index in indices]
-      observs = [observ() for observ in observs]
-    observ = np.stack(observs).astype(self.observ_space.dtype)
-
-    return observ
-
-  def close(self):
-    """Send close messages to the external process and join them."""
-    for env in self._envs:
-      if hasattr(env, "close"):
-        env.close()
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index eabdb6e85..62bdc93cf 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -23,15 +23,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import atexit
-import multiprocessing
-import os
-import random
-import signal
-import subprocess
-import sys
-import traceback
-
 from tensor2tensor.data_generators import gym_env
 from tensor2tensor.rl.envs import py_func_batch_env
 from tensor2tensor.rl.envs import simulated_batch_env
@@ -56,7 +47,7 @@ def _define_batch_env(environment_spec, num_agents, xvfb=False):
 
   with tf.variable_scope("environments"):
     envs = [
-        ExternalProcessEnv(environment_spec.env_lambda, xvfb)
+        environment_spec.env_lambda()
         for _ in range(num_agents)]
     env = gym_env.T2TGymEnv(envs)
     env = py_func_batch_env.PyFuncBatchEnv(env)
@@ -69,210 +60,3 @@ def _define_simulated_batch_env(environment_spec, num_agents,
       environment_spec, num_agents, initial_frame_chooser
   )
   return cur_batch_env
-
-
-class ExternalProcessEnv(object):
-  """Step environment in a separate process for lock free parallelism."""
-
-  # Message types for communication via the pipe.
-  _ACCESS = 1
-  _CALL = 2
-  _RESULT = 3
-  _EXCEPTION = 4
-  _CLOSE = 5
-  _ATTRIBUTE_EXCEPTION = 6
-
-  def __init__(self, constructor, xvfb):
-    """Step environment in a separate process for lock free parallelism.
-
-    The environment will be created in the external process by calling the
-    specified callable. This can be an environment class, or a function
-    creating the environment and potentially wrapping it. The returned
-    environment should not access global variables.
-
-    Args:
-      constructor: Callable that creates and returns an OpenAI gym environment.
-      xvfb:  Frame buffer.
-
-    Attributes:
-      observation_space: The cached observation space of the environment.
-      action_space: The cached action space of the environment.
-    """
-    self._constructor = constructor
-    self._conn, conn = multiprocessing.Pipe()
-    if xvfb:
-      server_id = random.randint(10000, 99999)
-      auth_file_id = random.randint(10000, 99999999999)
-
-      xauthority_path = "/tmp/Xauthority_{}".format(auth_file_id)
-
-      command = "Xvfb :{} -screen 0 1400x900x24 -nolisten tcp -auth {}".format(
-          server_id, xauthority_path)
-      with open(os.devnull, "w") as devnull:
-        proc = subprocess.Popen(command.split(), shell=False, stdout=devnull,
-                                stderr=devnull)
-        atexit.register(lambda: os.kill(proc.pid, signal.SIGKILL))
-
-      def constructor_using_xvfb():
-        os.environ["DISPLAY"] = ":{}".format(server_id)
-        os.environ["XAUTHORITY"] = xauthority_path
-        return constructor()
-
-      self._process = multiprocessing.Process(
-          target=self._worker, args=(constructor_using_xvfb, conn))
-    else:
-      self._process = multiprocessing.Process(
-          target=self._worker, args=(constructor, conn))
-
-    atexit.register(self.close)
-    self._process.start()
-    self._observ_space = None
-    self._action_space = None
-
-  def __str__(self):
-    return "ExternalProcessEnv(%s)" % str(self._constructor)
-
-  @property
-  def observation_space(self):
-    if not self._observ_space:
-      self._observ_space = self.__getattr__("observation_space")
-    return self._observ_space
-
-  @property
-  def action_space(self):
-    if not self._action_space:
-      self._action_space = self.__getattr__("action_space")
-    return self._action_space
-
-  def __getattr__(self, name):
-    """Request an attribute from the environment.
-
-    Note that this involves communication with the external process, so it can
-    be slow.
-
-    Args:
-      name: Attribute to access.
-
-    Returns:
-      Value of the attribute.
-    """
-    self._conn.send((self._ACCESS, name))
-    return self._receive()
-
-  def call(self, name, *args, **kwargs):
-    """Asynchronously call a method of the external environment.
-
-    Args:
-      name: Name of the method to call.
-      *args: Positional arguments to forward to the method.
-      **kwargs: Keyword arguments to forward to the method.
-
-    Returns:
-      Promise object that blocks and provides the return value when called.
-    """
-    payload = name, args, kwargs
-    self._conn.send((self._CALL, payload))
-    return self._receive
-
-  def close(self):
-    """Send a close message to the external process and join it."""
-    try:
-      self._conn.send((self._CLOSE, None))
-      self._conn.close()
-    except IOError:
-      # The connection was already closed.
-      pass
-    self._process.join()
-
-  def step(self, action, blocking=True):
-    """Step the environment.
-
-    Args:
-      action: The action to apply to the environment.
-      blocking: Whether to wait for the result.
-
-    Returns:
-      Transition tuple when blocking, otherwise callable that returns the
-      transition tuple.
-    """
-    promise = self.call("step", action)
-    if blocking:
-      return promise()
-    return promise
-
-  def reset(self, blocking=True):
-    """Reset the environment.
-
-    Args:
-      blocking: Whether to wait for the result.
-
-    Returns:
-      New observation when blocking, otherwise callable that returns the new
-      observation.
-    """
-    promise = self.call("reset")
-    if blocking:
-      return promise()
-    return promise
-
-  def _receive(self):
-    """Wait for a message from the worker process and return its payload.
-
-    Raises:
-      Exception: An exception was raised inside the worker process.
-      KeyError: The received message is of an unknown type.
-
-    Returns:
-      Payload object of the message.
-    """
-    message, payload = self._conn.recv()
-    # Re-raise exceptions in the main process.
-    if message == self._EXCEPTION:
-      stacktrace = payload
-      raise Exception(stacktrace)
-    if message == self._ATTRIBUTE_EXCEPTION:
-      raise AttributeError(payload)
-    if message == self._RESULT:
-      return payload
-    raise KeyError("Received message of unexpected type {}".format(message))
-
-  def _worker(self, constructor, conn):
-    """The process waits for actions and sends back environment results.
-
-    Args:
-      constructor: Constructor for the OpenAI Gym environment.
-      conn: Connection for communication to the main process.
-    """
-    try:
-      env = constructor()
-      while True:
-        try:
-          # Only block for short times to have keyboard exceptions be raised.
-          if not conn.poll(0.1):
-            continue
-          message, payload = conn.recv()
-        except (EOFError, KeyboardInterrupt):
-          break
-        if message == self._ACCESS:
-          name = payload
-          try:
-            result = getattr(env, name)
-            conn.send((self._RESULT, result))
-          except AttributeError as err:
-            conn.send((self._ATTRIBUTE_EXCEPTION, err.args))
-          continue
-        if message == self._CALL:
-          name, args, kwargs = payload
-          result = getattr(env, name)(*args, **kwargs)
-          conn.send((self._RESULT, result))
-          continue
-        if message == self._CLOSE:
-          assert payload is None
-          env.close()
-          break
-        raise KeyError("Received message of unknown type {}".format(message))
-    except Exception:  # pylint: disable=broad-except
-      stacktrace = "".join(traceback.format_exception(*sys.exc_info()))  # pylint: disable=no-value-for-parameter
-      tf.logging.error("Error in environment process: {}".format(stacktrace))
-      conn.send((self._EXCEPTION, stacktrace))
-    conn.close()

From 53cdd2a05288d84f0df0793ccf6e2ba9728354e4 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 9 Oct 2018 17:49:28 -0700
Subject: [PATCH 0977/2720] internal merge of PR #1124

PiperOrigin-RevId: 216455632
---
 tensor2tensor/data_generators/gym_env.py      | 6 +++---
 tensor2tensor/data_generators/gym_env_test.py | 3 ++-
 tensor2tensor/rl/envs/batch_env_factory.py    | 7 +++----
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index dba260ffb..b9b05dc41 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -156,10 +156,10 @@ def step(self, actions):
     ):
       rollout.append(frame._replace(action=action))
 
-    # ord = (observation, reward, done)
+    # ord_tuple = (observation, reward, done)
     self._current_frames = [
-        Frame(*ord, action=None)
-        for ord in zip(encoded_obs, rewards, dones)
+        Frame(*ord_tuple, action=None)
+        for ord_tuple in zip(encoded_obs, rewards, dones)
     ]
     return (obs, rewards, dones)
 
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 241fa4389..f0d04d4a1 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -23,7 +23,8 @@
 import shutil
 
 import gym
-from gym.spaces import Box, Discrete
+from gym.spaces import Box
+from gym.spaces import Discrete
 import numpy as np
 
 from tensor2tensor.data_generators import gym_env
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index 62bdc93cf..d6d9cce75 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -30,19 +30,18 @@
 import tensorflow as tf
 
 
-def batch_env_factory(environment_spec, num_agents,
-                      initial_frame_chooser=None, xvfb=False):
+def batch_env_factory(environment_spec, num_agents, initial_frame_chooser=None):
   """Factory of batch envs."""
 
   if environment_spec.simulated_env:
     cur_batch_env = _define_simulated_batch_env(
         environment_spec, num_agents, initial_frame_chooser)
   else:
-    cur_batch_env = _define_batch_env(environment_spec, num_agents, xvfb=xvfb)
+    cur_batch_env = _define_batch_env(environment_spec, num_agents)
   return cur_batch_env
 
 
-def _define_batch_env(environment_spec, num_agents, xvfb=False):
+def _define_batch_env(environment_spec, num_agents):
   """Create environments and apply all desired wrappers."""
 
   with tf.variable_scope("environments"):

From 7cb5e55d206cee3de179dc27708b8920f6e67e36 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 9 Oct 2018 20:58:00 -0700
Subject: [PATCH 0978/2720] Allow decoding from the dual autoencoder and
 increase number of shared bits.

PiperOrigin-RevId: 216471592
---
 tensor2tensor/models/research/autoencoders.py | 48 ++++++++++++++++---
 1 file changed, 41 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index c7e5b1406..ab0c516c0 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -56,6 +56,7 @@ def __init__(self, *args, **kwargs):
     super(AutoencoderBasic, self).__init__(*args, **kwargs)
     self._cur_bottleneck_tensor = None
     self.is1d = None
+    self._encode_on_predict = False
 
   @property
   def num_channels(self):
@@ -189,7 +190,8 @@ def body(self, features):
       vocab_size = self._problem_hparams.target_modality.top_dimensionality
     encoder_layers = None
     self.is1d = hparams.sample_width == 1
-    if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+    if (hparams.mode != tf.estimator.ModeKeys.PREDICT
+        or self._encode_on_predict):
       labels = features["targets_raw"]
       labels_shape = common_layers.shape_list(labels)
       # handle videos
@@ -846,14 +848,14 @@ class AutoencoderDualDiscrete(AutoencoderResidualDiscrete):
   """Dual discrete autoencoder."""
 
   def body(self, features):
-    if self.hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    if self.hparams.mode != tf.estimator.ModeKeys.EVAL:
       t, i = features["targets_raw"], features["inputs_raw"]
       t, i = common_layers.pad_to_same_length(t, i)
       features["targets_raw"] = tf.concat([t, i], axis=0)
     return super(AutoencoderDualDiscrete, self).body(features)
 
   def embed(self, x, name="embedding"):
-    if self.hparams.mode != tf.estimator.ModeKeys.TRAIN:
+    if self.hparams.mode == tf.estimator.ModeKeys.EVAL:
       return super(AutoencoderDualDiscrete, self).embed(x, name=name + "_t")
     xt, xi = tf.split(x, 2, axis=0)
     xte = super(AutoencoderDualDiscrete, self).embed(xt, name=name + "_t")
@@ -863,9 +865,11 @@ def embed(self, x, name="embedding"):
   def bottleneck(self, x):
     hparams = self.hparams
     b, _ = super(AutoencoderDualDiscrete, self).bottleneck(x)
-    if hparams.mode != tf.estimator.ModeKeys.TRAIN:
+    if hparams.mode == tf.estimator.ModeKeys.EVAL:
       return b, 0.0
     bt, bi = tf.split(b, 2, axis=0)
+    if self.hparams.mode != tf.estimator.ModeKeys.TRAIN:
+      return tf.concat([bi, bi], axis=0), 0.0
     # Share the first hparams.bottleneck_shared_bits.
     shared = (bt + bi) / 2  # -1 if both -1, 1 if both were 1, 0 if disagree.
     rand = tf.random_uniform(common_layers.shape_list(bt))
@@ -890,13 +894,43 @@ def bottleneck(self, x):
   def unbottleneck(self, b, res_size, reuse=None):
     x = super(AutoencoderDualDiscrete, self).unbottleneck(
         b, res_size, reuse=reuse)
-    if self.hparams.mode != tf.estimator.ModeKeys.TRAIN:
+    if self.hparams.mode == tf.estimator.ModeKeys.EVAL:
       return tf.layers.dense(x, res_size, name="dual_unbottleneck_t")
     xt, xi = tf.split(x, 2, axis=0)
     xt = tf.layers.dense(xt, res_size, name="dual_unbottleneck_t")
     xi = tf.layers.dense(xt, res_size, name="dual_unbottleneck_i")
     return tf.concat([xt, xi], axis=0)
 
+  def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
+    """Produce predictions from the model."""
+    del args, kwargs
+    # Inputs and features preparation needed to handle edge cases.
+    if not features:
+      features = {}
+    inputs_old = None
+    if "inputs" in features and len(features["inputs"].shape) < 4:
+      inputs_old = features["inputs"]
+      features["inputs"] = tf.expand_dims(features["inputs"], 2)
+
+    # Set targets to input size firts.
+    features["targets"] = tf.zeros_like(features["inputs"])
+    self._encode_on_predict = True
+    logits, _ = self(features)  # pylint: disable=not-callable
+    if self.hparams.gan_loss_factor != 0:
+      logits, _ = tf.split(logits, 2, axis=0)  # Remove GAN.
+    logits, _ = tf.split(logits, 2, axis=0)  # Targets and inputs from encoding.
+    # Uncomment the line below to get reconstructed inputs instead of targets.
+    # (and comment out the line above at the same time).
+    # _, logits = tf.split(logits, 2, axis=0)
+    samples = tf.argmax(logits, axis=-1)
+
+    # Restore inputs to not confuse Estimator in edge cases.
+    if inputs_old is not None:
+      features["inputs"] = inputs_old
+
+    # Return samples.
+    return samples
+
 
 @registry.register_model
 class AutoencoderStacked(AutoencoderResidualDiscrete):
@@ -1175,8 +1209,8 @@ def autoencoder_ordered_discrete_hs256():
 def autoencoder_ordered_text():
   """Ordered discrete autoencoder model for text."""
   hparams = autoencoder_ordered_discrete()
-  hparams.bottleneck_bits = 512
-  hparams.bottleneck_shared_bits = 512-64
+  hparams.bottleneck_bits = 1024
+  hparams.bottleneck_shared_bits = 1024-64
   hparams.bottleneck_shared_bits_start_warmup = 75000
   hparams.bottleneck_shared_bits_stop_warmup = 275000
   hparams.num_hidden_layers = 7

From 1bb0f5616ed6bf3ab9b651d6758b2ea8a453f403 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Wed, 10 Oct 2018 20:37:07 +0200
Subject: [PATCH 0979/2720] Preprocessing in t2 t gym env (#1131)

* Adds reward clipping, observation preprocessing in T2TGymEnv

* T2TEnv: Resizing done by tf (instead of PIL); save unclipped reward to disc; changes reward clipping logic - env returns clipped/unclipped rewards based on argument passed in step(); fixes num_channels;

* Tests for T2TGymEnv - resizing, reward clipping, number of channels
---
 tensor2tensor/data_generators/gym_env.py      | 100 +++++++++++-------
 tensor2tensor/data_generators/gym_env_test.py |  72 +++++++++++--
 2 files changed, 127 insertions(+), 45 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index b9b05dc41..127abb5a7 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -24,6 +24,7 @@
 import random
 
 import numpy as np
+from gym.spaces import Box
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -35,7 +36,7 @@
 
 Frame = collections.namedtuple(
     # Order of elements reflects time progression within a frame.
-    "Frame", ("observation", "reward", "done", "action")
+    "Frame", ("observation", "reward", "unclipped_reward", "done", "action")
 )
 
 
@@ -70,10 +71,11 @@ def __init__(self, batch_size):
     self._current_rollouts = [[] for _ in range(batch_size)]
     self._current_frames = [None for _ in range(batch_size)]
 
-    with tf.Graph().as_default():
+    with tf.Graph().as_default() as tf_graph:
+      self._tf_graph = tf_graph
       self._image_t = tf.placeholder(dtype=tf.uint8, shape=(None, None, None))
       self._encoded_image_t = tf.image.encode_png(self._image_t)
-      self._encode_session = tf.Session()
+      self._session = tf.Session()
 
   def __str__(self):
     """Returns a string representation of the environment for debug purposes."""
@@ -96,23 +98,10 @@ def _preprocess_observations(self, obs):
     """
     return obs
 
-  def _preprocess_rewards(self, rewards):
-    """Transforms a batch of rewards.
-
-    Can be overridden in derived classes.
-
-    Args:
-      rewards: A batch of rewards.
-
-    Returns:
-      Transformed batch of rewards.
-    """
-    return rewards
-
   def _encode_observations(self, observations):
     """Encodes observations as PNG."""
     return [
-        self._encode_session.run(
+        self._session.run(
             self._encoded_image_t, feed_dict={self._image_t: observation}
         )
         for observation in observations
@@ -135,7 +124,7 @@ def _step(self, actions):
     """
     raise NotImplementedError
 
-  def step(self, actions):
+  def step(self, actions, return_unclipped_rewards=False):
     """Makes a step in all environments.
 
     Does any preprocessing and records frames.
@@ -147,21 +136,25 @@ def step(self, actions):
       (obs, rewards, dones) - batches of observations, rewards and done flags
       respectively.
     """
-    (obs, rewards, dones) = self._step(actions)
+    (obs, unclipped_rewards, dones) = self._step(actions)
     obs = self._preprocess_observations(obs)
-    rewards = self._preprocess_rewards(rewards)
+    rewards = np.clip(unclipped_rewards, -1, 1)
     encoded_obs = self._encode_observations(obs)
     for (rollout, frame, action) in zip(
         self._current_rollouts, self._current_frames, actions
     ):
       rollout.append(frame._replace(action=action))
 
-    # ord_tuple = (observation, reward, done)
+    # orud_tuple = (observation, reward, unclipped_reward, done)
     self._current_frames = [
-        Frame(*ord_tuple, action=None)
-        for ord_tuple in zip(encoded_obs, rewards, dones)
+        Frame(*orud_tuple, action=None)
+        for orud_tuple in zip(encoded_obs, rewards, unclipped_rewards, dones)
     ]
-    return (obs, rewards, dones)
+    if return_unclipped_rewards:
+      ret_rewards = unclipped_rewards
+    else:
+      ret_rewards = rewards
+    return (obs, ret_rewards, dones)
 
   def _reset(self, indices):
     """Resets environments at given indices without recording history.
@@ -198,7 +191,7 @@ def reset(self, indices=None):
         self.history.append(rollout)
         self._current_rollouts[index] = []
       self._current_frames[index] = Frame(
-          observation=ob, reward=0, done=False, action=None
+          observation=ob, reward=0, unclipped_reward=0, done=False, action=None
       )
     return new_obs
 
@@ -207,12 +200,12 @@ def close(self):
 
     Can be overridden in derived classes.
     """
-    self._encode_session.close()
+    self._session.close()
 
   @property
   def num_channels(self):
     """Number of color channels in each frame."""
-    return 3
+    raise NotImplementedError
 
   def eval_metrics(self):
     eval_metrics = [
@@ -315,31 +308,62 @@ class T2TGymEnv(T2TEnv):
 
   name = "t2t_gym_env"
 
-  def __init__(self, envs):
+  def __init__(self, envs, clip_rewards=False, grayscale=False,
+               resize_height_factor=1, resize_width_factor=1):
     super(T2TGymEnv, self).__init__(len(envs))
-
+    self.clip_rewards = clip_rewards
+    self.grayscale = grayscale
+    self.resize_height_factor = resize_height_factor
+    self.resize_width_factor = resize_width_factor
     if not envs:
       raise ValueError("Must have at least one environment.")
     self._envs = envs
 
-    self.observation_space = envs[0].observation_space
-    if not all(env.observation_space == self.observation_space for env in envs):
+    orig_observ_space = envs[0].observation_space
+    if not all(env.observation_space == orig_observ_space
+               for env in self._envs):
       raise ValueError("All environments must use the same observation space.")
 
+    self.observation_space = self._derive_observation_space(orig_observ_space)
+
     self.action_space = envs[0].action_space
-    if not all(env.action_space == self.action_space for env in envs):
+    if not all(env.action_space == self.action_space for env in self._envs):
       raise ValueError("All environments must use the same action space.")
 
+    with self._tf_graph.as_default():
+      self._resize = dict()
+      orig_height, orig_width = orig_observ_space.shape[:2]
+      self._img_batch_t = tf.placeholder(
+          dtype=tf.uint8, shape=(None, orig_height, orig_width, 3))
+      height, width = self.observation_space.shape[:2]
+      resized = tf.image.resize_images(self._img_batch_t,
+                                       [height, width],
+                                       tf.image.ResizeMethod.AREA)
+      if self.grayscale:
+        resized = tf.image.rgb_to_grayscale(resized)
+      self._resized_img_batch_t = resized
+
+  @property
+  def num_channels(self):
+    return self.observation_space.shape[2]
+
+  def _derive_observation_space(self, orig_observ_space):
+    height, width, channels = orig_observ_space.shape
+    if self.grayscale:
+      channels = 1
+    resized_height = height // self.resize_height_factor
+    resized_width = width // self.resize_width_factor
+    shape = (resized_height, resized_width, channels)
+    return Box(low=orig_observ_space.low.min(),
+               high=orig_observ_space.high.max(), shape=shape,
+               dtype=orig_observ_space.dtype)
+
   def __str__(self):
     return "T2TGymEnv(%s)" % ", ".join([str(env) for env in self._envs])
 
   def _preprocess_observations(self, obs):
-    # TODO(lukaszkaiser): Implement.
-    return obs
-
-  def _preprocess_rewards(self, rewards):
-    # TODO(lukaszkaiser): Implement.
-    return rewards
+    return self._session.run(self._resized_img_batch_t,
+                             feed_dict={self._img_batch_t: obs})
 
   def _step(self, actions):
     (obs, rewards, dones, _) = zip(*[
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index f0d04d4a1..e0f1ebafa 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -40,7 +40,7 @@ class TestEnv(gym.Env):
 
   action_space = Discrete(1)
   observation_space = Box(
-      low=0, high=255, shape=(2, 2, 1), dtype=np.uint8
+      low=0, high=255, shape=(2, 6, 3), dtype=np.uint8
   )
 
   def __init__(self):
@@ -54,7 +54,8 @@ def _generate_ob(self):
   def step(self, action):
     done = self._counter % 2 == 1
     self._counter += 1
-    return (self._generate_ob(), 0, done, {})
+    reward = 5 if done else -5
+    return (self._generate_ob(), reward, done, {})
 
   def reset(self):
     return self._generate_ob()
@@ -68,14 +69,27 @@ def setUpClass(cls):
     shutil.rmtree(cls.out_dir)
     os.mkdir(cls.out_dir)
 
-  def test_generates(self):
-    env = gym_env.T2TGymEnv([TestEnv(), TestEnv()])
-    env.reset()
-    for _ in range(20):
-      (_, _, dones) = env.step([0, 0])
+  def init_batch_and_play(self, env_lambda, n_steps=1, unclipped_rewards=False,
+                          **kwargs):
+    raw_envs = [env_lambda(), env_lambda()]
+    env = gym_env.T2TGymEnv(raw_envs, **kwargs)
+    obs = list()
+    rewards = list()
+    obs.append(env.reset())
+    for _ in range(n_steps):
+      step_obs, step_rewards, dones = env.step(
+          actions=[0, 0], return_unclipped_rewards=unclipped_rewards)
+      obs.append(step_obs)
+      rewards.append(step_rewards)
       for (i, done) in enumerate(dones):
         if done:
           env.reset([i])
+    return env, obs, rewards
+
+  def test_generates(self):
+    # This test needs base env which outputs done after two steps.
+    env_lambda = TestEnv
+    env, _, _ = self.init_batch_and_play(env_lambda, n_steps=20)
     env.generate_data(self.out_dir, tmp_dir=None)
 
     filenames = os.listdir(self.out_dir)
@@ -84,6 +98,50 @@ def test_generates(self):
     records = list(tf.python_io.tf_record_iterator(path))
     self.assertTrue(records)
 
+  def test_clipping(self):
+    # This test needs base env with rewards out of [-1,1] range.
+    env_lambda = TestEnv
+    _, _, rewards = self.init_batch_and_play(env_lambda, n_steps=2)
+    self.assertTrue(np.max(rewards) == 1)
+    self.assertTrue(np.min(rewards) == -1)
+
+    _, _, unclipped_rewards = self.init_batch_and_play(env_lambda, n_steps=2,
+                                                       unclipped_rewards=True)
+    self.assertTrue(np.max(unclipped_rewards) > 1)
+    self.assertTrue(np.min(unclipped_rewards) < -1)
+
+  def test_resize(self):
+    env_lambda = TestEnv
+    orig_env = env_lambda()
+    resize_height_factor = 2
+    resize_width_factor = 3
+    orig_height, orig_width = orig_env.observation_space.shape[:2]
+    env, obs, _ = self.init_batch_and_play(
+        env_lambda, n_steps=1,
+        resize_height_factor=resize_height_factor,
+        resize_width_factor=resize_width_factor)
+    for obs_batch in obs:
+      ob = obs_batch[0]
+      self.assertEqual(ob.shape, env.observation_space.shape)
+      height, width = ob.shape[:2]
+      self.assertEqual(height, orig_height // resize_height_factor)
+      self.assertEqual(width, orig_width // resize_width_factor)
+
+  def assert_channels(self, env, obs, n_channels):
+    self.assertEqual(env.observation_space.shape[2], n_channels)
+    self.assertEqual(env.num_channels, n_channels)
+    for obs_batch in obs:
+      ob = obs_batch[0]
+      self.assertEqual(ob.shape[2], n_channels)
+
+  def test_channels(self):
+    env_lambda = TestEnv
+    env, obs, _ = self.init_batch_and_play(env_lambda, grayscale=True)
+    self.assert_channels(env, obs, n_channels=1)
+
+    env, obs, _ = self.init_batch_and_play(env_lambda, grayscale=False)
+    self.assert_channels(env, obs, n_channels=3)
+
 
 if __name__ == "__main__":
   tf.test.main()

From 1cd6b0eaec3d1993821870847dae6e121e36b8af Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Wed, 10 Oct 2018 11:37:27 -0700
Subject: [PATCH 0980/2720] internal merge of PR #1131

PiperOrigin-RevId: 216565119
---
 tensor2tensor/data_generators/gym_env.py      | 24 +++++++++----------
 tensor2tensor/data_generators/gym_env_test.py | 16 ++++++-------
 2 files changed, 18 insertions(+), 22 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 127abb5a7..6b31653b9 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -23,8 +23,8 @@
 import math
 import random
 
-import numpy as np
 from gym.spaces import Box
+import numpy as np
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -124,7 +124,7 @@ def _step(self, actions):
     """
     raise NotImplementedError
 
-  def step(self, actions, return_unclipped_rewards=False):
+  def step(self, actions):
     """Makes a step in all environments.
 
     Does any preprocessing and records frames.
@@ -138,23 +138,22 @@ def step(self, actions, return_unclipped_rewards=False):
     """
     (obs, unclipped_rewards, dones) = self._step(actions)
     obs = self._preprocess_observations(obs)
-    rewards = np.clip(unclipped_rewards, -1, 1)
+    (min_reward, max_reward) = self.reward_range
+    rewards = np.around(np.clip(unclipped_rewards, min_reward, max_reward))
     encoded_obs = self._encode_observations(obs)
     for (rollout, frame, action) in zip(
         self._current_rollouts, self._current_frames, actions
     ):
       rollout.append(frame._replace(action=action))
 
-    # orud_tuple = (observation, reward, unclipped_reward, done)
+    # orud = (observation, reward, unclipped_reward, done)
     self._current_frames = [
-        Frame(*orud_tuple, action=None)
-        for orud_tuple in zip(encoded_obs, rewards, unclipped_rewards, dones)
+        Frame(*orud, action=None)
+        for orud in zip(encoded_obs, rewards, unclipped_rewards, dones)
     ]
-    if return_unclipped_rewards:
-      ret_rewards = unclipped_rewards
-    else:
-      ret_rewards = rewards
-    return (obs, ret_rewards, dones)
+    # TODO(lukaszkaiser): changed unclipped_reward to reward once we've
+    # removed the current setup with RewardClippingWrapper and so on.
+    return (obs, unclipped_rewards, dones)
 
   def _reset(self, indices):
     """Resets environments at given indices without recording history.
@@ -308,10 +307,9 @@ class T2TGymEnv(T2TEnv):
 
   name = "t2t_gym_env"
 
-  def __init__(self, envs, clip_rewards=False, grayscale=False,
+  def __init__(self, envs, grayscale=False,
                resize_height_factor=1, resize_width_factor=1):
     super(T2TGymEnv, self).__init__(len(envs))
-    self.clip_rewards = clip_rewards
     self.grayscale = grayscale
     self.resize_height_factor = resize_height_factor
     self.resize_width_factor = resize_width_factor
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index e0f1ebafa..0317d6261 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -69,16 +69,14 @@ def setUpClass(cls):
     shutil.rmtree(cls.out_dir)
     os.mkdir(cls.out_dir)
 
-  def init_batch_and_play(self, env_lambda, n_steps=1, unclipped_rewards=False,
-                          **kwargs):
+  def init_batch_and_play(self, env_lambda, n_steps=1, **kwargs):
     raw_envs = [env_lambda(), env_lambda()]
     env = gym_env.T2TGymEnv(raw_envs, **kwargs)
     obs = list()
     rewards = list()
     obs.append(env.reset())
     for _ in range(n_steps):
-      step_obs, step_rewards, dones = env.step(
-          actions=[0, 0], return_unclipped_rewards=unclipped_rewards)
+      step_obs, step_rewards, dones = env.step(actions=[0, 0])
       obs.append(step_obs)
       rewards.append(step_rewards)
       for (i, done) in enumerate(dones):
@@ -101,12 +99,12 @@ def test_generates(self):
   def test_clipping(self):
     # This test needs base env with rewards out of [-1,1] range.
     env_lambda = TestEnv
-    _, _, rewards = self.init_batch_and_play(env_lambda, n_steps=2)
-    self.assertTrue(np.max(rewards) == 1)
-    self.assertTrue(np.min(rewards) == -1)
+    # TODO(lukaszkaiser): turn clipping on by default after refactor.
+    # _, _, rewards = self.init_batch_and_play(env_lambda, n_steps=2)
+    # self.assertTrue(np.max(rewards) == 1)
+    # self.assertTrue(np.min(rewards) == -1)
 
-    _, _, unclipped_rewards = self.init_batch_and_play(env_lambda, n_steps=2,
-                                                       unclipped_rewards=True)
+    _, _, unclipped_rewards = self.init_batch_and_play(env_lambda, n_steps=2)
     self.assertTrue(np.max(unclipped_rewards) > 1)
     self.assertTrue(np.min(unclipped_rewards) < -1)
 

From a2cc2fed4b132d500cd58d94be2254aafe567711 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 10 Oct 2018 15:39:41 -0700
Subject: [PATCH 0981/2720] Make a RNN for predicting latent state structure
 and use 16 frames grayscale default for stochastic discrete.

PiperOrigin-RevId: 216605656
---
 .../models/video/basic_deterministic.py       | 11 ++-
 .../models/video/basic_stochastic.py          | 82 ++++++++++++++++---
 tensor2tensor/rl/trainer_model_based.py       | 12 +++
 3 files changed, 91 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 1ce18b68f..01a3f6e8e 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -209,7 +209,7 @@ def body_single(self, features):
 
     # Reward prediction if needed.
     if "target_reward" not in features:
-      return x
+      return x, extra_loss
     reward_pred = tf.expand_dims(  # Add a fake channels dim.
         tf.reduce_mean(x, axis=[1, 2], keepdims=True), axis=3)
     return {"targets": x, "target_reward": reward_pred}, extra_loss
@@ -236,9 +236,12 @@ def body(self, features):
 
     # Run a number of steps.
     res_frames, sampled_frames, sampled_frames_raw = [], [], []
+    extra_loss = 0.0
     if "target_reward" in features:
-      res_rewards, extra_loss = [], 0.0
+      res_rewards = []
     sample_prob = common_layers.inverse_exp_decay(
+        hparams.scheduled_sampling_warmup_steps // 4)
+    sample_prob *= common_layers.inverse_lin_decay(
         hparams.scheduled_sampling_warmup_steps)
     sample_prob *= hparams.scheduled_sampling_prob
     for i in range(hparams.video_num_target_frames):
@@ -253,13 +256,13 @@ def body(self, features):
       # Run model.
       with tf.variable_scope(tf.get_variable_scope(), reuse=i > 0):
         if "target_reward" not in features:
-          res_frame = self.body_single(features)
+          res_frame, res_extra_loss = self.body_single(features)
         else:
           res_dict, res_extra_loss = self.body_single(features)
-          extra_loss += res_extra_loss
           res_frame = res_dict["targets"]
           res_reward = res_dict["target_reward"]
           res_rewards.append(res_reward)
+      extra_loss += res_extra_loss / float(hparams.video_num_target_frames)
       res_frames.append(res_frame)
 
       # Only for Softmax loss: sample frame so we can keep iterating.
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 26b0d4d5f..8b94245b5 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -22,6 +22,7 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
+from tensor2tensor.layers import discretization
 
 from tensor2tensor.models.video import base_vae
 from tensor2tensor.models.video import basic_deterministic
@@ -32,6 +33,15 @@
 import tensorflow as tf
 
 
+def prod(l):
+  """Product of elements in a list."""
+  res = l[0]
+  for i, e in enumerate(l):
+    if i > 0:
+      res *= e
+  return res
+
+
 @registry.register_model
 class NextFrameBasicStochastic(
     basic_deterministic.NextFrameBasicDeterministic,
@@ -70,6 +80,12 @@ def inject_latent(self, layer, features, filters):
     final_filters = common_layers.shape_list(layer)[-1]
     filters = hparams.hidden_size
     kernel = (4, 4)
+    layer_shape = common_layers.shape_list(layer)
+    batch_size = layer_shape[0]
+    state_size = hparams.latent_predictor_state_size
+    lstm_cell = tf.contrib.rnn.LSTMCell(state_size)
+    discrete_predict = tf.layers.Dense(256, name="discrete_predict")
+    discrete_embed = tf.layers.Dense(state_size, name="discrete_embed")
 
     def add_d(layer, d):
       z_mul = tf.layers.dense(d, final_filters, name="unbottleneck_mul")
@@ -81,12 +97,25 @@ def add_d(layer, d):
       return layer
 
     if self.is_predicting:
-      layer_shape = common_layers.shape_list(layer)
       if hparams.full_latent_tower:
         rand = tf.random_uniform(layer_shape[:-1] + [hparams.bottleneck_bits])
       else:
-        rand = tf.random_uniform(layer_shape[:-3] + [
-            1, 1, hparams.bottleneck_bits])
+        layer_pred = tf.reshape(layer, [batch_size, prod(layer_shape[1:])])
+        prediction = tf.layers.dense(layer_pred, state_size, name="istate")
+        c_state = tf.layers.dense(layer_pred, state_size, name="cstate")
+        m_state = tf.layers.dense(layer_pred, state_size, name="mstate")
+        state = (c_state, m_state)
+        outputs = []
+        for i in range(hparams.bottleneck_bits // 8):
+          output, state = lstm_cell(prediction, state)
+          discrete_logits = discrete_predict(output)
+          discrete_samples = common_layers.sample_with_temperature(
+              discrete_logits, hparams.latent_predictor_temperature)
+          outputs.append(tf.expand_dims(discrete_samples, axis=1))
+          prediction = discrete_embed(tf.one_hot(discrete_samples, 256))
+        outputs = tf.concat(outputs, axis=1)
+        outputs = discretization.int_to_bit(outputs, 8)
+        rand = tf.reshape(outputs, [batch_size, 1, 1, hparams.bottleneck_bits])
       d = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
       return add_d(layer, d), 0.0
 
@@ -112,16 +141,43 @@ def add_d(layer, d):
     else:
       x = common_layers.double_discriminator(x)
       x = tf.expand_dims(tf.expand_dims(x, axis=1), axis=1)
-    x = tf.tanh(tf.layers.dense(x, hparams.bottleneck_bits, name="bottleneck"))
-    d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
+    x = tf.layers.dense(x, hparams.bottleneck_bits, name="bottleneck")
+    x0 = tf.tanh(x)
+    d = x0 + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x0)) - 1.0 - x0)
+    pred_loss = 0.0
+    if not hparams.full_latent_tower:
+      d_pred = tf.reshape(tf.maximum(tf.stop_gradient(d), 0), [
+          batch_size, hparams.bottleneck_bits // 8, 8])
+      d_int = discretization.bit_to_int(d_pred, 8)
+      tf.summary.histogram("d_int", tf.reshape(d_int, [-1]))
+      d_hot = tf.one_hot(d_int, 256, axis=-1)
+      d_pred = discrete_embed(d_hot)
+      layer_pred = tf.reshape(layer, [batch_size, prod(layer_shape[1:])])
+      prediction0 = tf.layers.dense(layer_pred, state_size, name="istate")
+      c_state = tf.layers.dense(layer_pred, state_size, name="cstate")
+      m_state = tf.layers.dense(layer_pred, state_size, name="mstate")
+      pred = tf.concat([tf.expand_dims(prediction0, axis=1), d_pred], axis=1)
+      state = (c_state, m_state)
+      outputs = []
+      for i in range(hparams.bottleneck_bits // 8):
+        output, state = lstm_cell(pred[:, i, :], state)
+        outputs.append(tf.expand_dims(output, axis=1))
+      outputs = tf.concat(outputs, axis=1)
+      d_int_pred = discrete_predict(outputs)
+      pred_loss = tf.losses.sparse_softmax_cross_entropy(
+          logits=d_int_pred, labels=d_int)
+      pred_loss = tf.reduce_mean(pred_loss)
     if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+      x += tf.truncated_normal(
+          common_layers.shape_list(x), mean=0.0, stddev=0.2)
+      x = tf.tanh(x)
       noise = tf.random_uniform(common_layers.shape_list(x))
       noise = 2.0 * tf.to_float(tf.less(hparams.bottleneck_noise, noise)) - 1.0
       x *= noise
       d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
       p = common_layers.inverse_lin_decay(hparams.discrete_warmup_steps)
-      d = tf.where(tf.less(tf.random_uniform([]), p), d, x)
-    return add_d(layer, d), 0.0
+      d = tf.where(tf.less(tf.random_uniform([batch_size]), p), d, x)
+    return add_d(layer, d), pred_loss
 
 
 @registry.register_hparams
@@ -168,9 +224,15 @@ def next_frame_sampling_stochastic():
 def next_frame_basic_stochastic_discrete():
   """Basic 2-frame conv model with stochastic discrete latent."""
   hparams = basic_deterministic_params.next_frame_sampling()
-  hparams.add_hparam("bottleneck_bits", 32)
-  hparams.add_hparam("bottleneck_noise", 0.05)
-  hparams.add_hparam("discrete_warmup_steps", 4000)
+  hparams.batch_size = 2
+  hparams.video_num_target_frames = 16
+  hparams.scheduled_sampling_warmup_steps = 40000
+  hparams.scheduled_sampling_prob = 1.0
+  hparams.add_hparam("bottleneck_bits", 64)
+  hparams.add_hparam("bottleneck_noise", 0.02)
+  hparams.add_hparam("discrete_warmup_steps", 40000)
   hparams.add_hparam("full_latent_tower", False)
+  hparams.add_hparam("latent_predictor_state_size", 128)
+  hparams.add_hparam("latent_predictor_temperature", 0.5)
   hparams.add_hparam("complex_addn", True)
   return hparams
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index ec2a555d4..03fe2a525 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -762,6 +762,7 @@ def rlmb_base():
       ppo_params="ppo_pong_base",
       autoencoder_train_steps=0,
       autoencoder_train_steps_initial_multiplier=10,
+      autoencoder_hparams_set="autoencoder_discrete_pong",
       model_train_steps=15000,
       inital_epoch_train_steps_multiplier=3,
       simulated_env_generator_num_steps=2000,
@@ -865,6 +866,15 @@ def rlmb_quick_sd():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_sdtest():
+  """Test setting with stochastic discrete model."""
+  hparams = rlmb_basetest()
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_quick_sm():
   """Quick setting with sampling."""
@@ -896,6 +906,8 @@ def rlmb_base_sampling_stochastic():
 def rlmb_base_stochastic_discrete():
   """Base setting with stochastic discrete model."""
   hparams = rlmb_base()
+  hparams.learning_rate_bump = 1.0
+  hparams.grayscale = True
   hparams.generative_model = "next_frame_basic_stochastic_discrete"
   hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
   return hparams

From 86846707b07cb3a1315f4341f9e8e9ee6f9d3dbf Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 10 Oct 2018 16:00:22 -0700
Subject: [PATCH 0982/2720] merging all the scheduled samplings.

PiperOrigin-RevId: 216609071
---
 tensor2tensor/models/video/base_vae.py        |   7 --
 .../models/video/basic_deterministic.py       | 114 ++++++++++++++++--
 .../video/basic_deterministic_params.py       |  12 +-
 tensor2tensor/models/video/sv2p.py            |  61 ----------
 tensor2tensor/models/video/sv2p_params.py     |   6 +-
 5 files changed, 114 insertions(+), 86 deletions(-)

diff --git a/tensor2tensor/models/video/base_vae.py b/tensor2tensor/models/video/base_vae.py
index 8ef5e5a9b..44761dff7 100644
--- a/tensor2tensor/models/video/base_vae.py
+++ b/tensor2tensor/models/video/base_vae.py
@@ -31,13 +31,6 @@ class NextFrameBaseVae(object):
   def __init__(self, hparams):
     self.hparams = hparams
 
-  def get_iteration_num(self):
-    step_num = tf.train.get_global_step()
-    # TODO(lukaszkaiser): what should it be if it's undefined?
-    if step_num is None:
-      step_num = 1000000
-    return step_num
-
   def get_beta(self, kl_loss=0.0):
     """Get the KL multiplier, either dynamically or schedule based.
 
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 01a3f6e8e..62a190465 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from functools import partial
 import six
 
 from tensor2tensor.layers import common_attention
@@ -49,11 +50,104 @@ def _target_modality(self):
   def is_per_pixel_softmax(self):
     return self._target_modality == "VideoModality"
 
+  def get_iteration_num(self):
+    step_num = tf.train.get_global_step()
+    # TODO(lukaszkaiser): what should it be if it's undefined?
+    if step_num is None:
+      step_num = 10000000
+    return step_num
+
   def inject_latent(self, layer, features, filters):
     """Do nothing for deterministic model."""
     del features, filters
     return layer, 0.0
 
+  def get_scheduled_sample_func(self, batch_size):
+    """Creates a function for scheduled sampling based on given hparams."""
+    with tf.variable_scope("scheduled_sampling_func", reuse=tf.AUTO_REUSE):
+      iter_num = self.get_iteration_num()
+
+      # Simple function to bypass scheduled sampling in gt or pred only modes.
+      def scheduled_sampling_simple(ground_truth_x, generated_x,
+                                    batch_size, scheduled_sample_var):
+        del batch_size
+        if scheduled_sample_var:
+          return ground_truth_x
+        return generated_x
+
+      mode = self.hparams.scheduled_sampling_mode
+      if mode == "ground_truth_only":
+        scheduled_sampling_func = scheduled_sampling_simple
+        scheduled_sampling_func_var = True
+      elif mode == "prediction_only":
+        scheduled_sampling_func = scheduled_sampling_simple
+        scheduled_sampling_func_var = False
+      elif mode == "prob":
+        decay_steps = self.hparams.scheduled_sampling_decay_steps
+        probability = tf.train.polynomial_decay(
+            1.0, iter_num, decay_steps, 0.0)
+        scheduled_sampling_func = common_video.scheduled_sample_prob
+        scheduled_sampling_func_var = probability
+      elif mode == "prob_inverse_exp":
+        decay_steps = self.hparams.scheduled_sampling_decay_steps
+        probability = common_layers.inverse_exp_decay(
+            decay_steps, step=iter_num)
+        probability *= self.hparams.scheduled_sampling_max_prob
+        probability = 1.0 - probability
+        scheduled_sampling_func = common_video.scheduled_sample_prob
+        scheduled_sampling_func_var = probability
+      elif mode == "count":
+        # Calculate number of ground-truth frames to pass in.
+        k = self.hparams.scheduled_sampling_k
+        num_ground_truth = tf.to_int32(
+            tf.round(
+                tf.to_float(batch_size) *
+                (k / (k + tf.exp(tf.to_float(iter_num) / tf.to_float(k))))))
+        scheduled_sampling_func = common_video.scheduled_sample_count
+        scheduled_sampling_func_var = num_ground_truth
+      else:
+        raise ValueError("unknown scheduled sampling method: %s" % mode)
+
+      if isinstance(scheduled_sampling_func_var, tf.Tensor):
+        tf.summary.scalar("scheduled_sampling_var", scheduled_sampling_func_var)
+      partial_func = partial(scheduled_sampling_func,
+                             batch_size=batch_size,
+                             scheduled_sample_var=scheduled_sampling_func_var)
+      return partial_func
+
+  def get_scheduled_sample_inputs(self,
+                                  done_warm_start,
+                                  groundtruth_items,
+                                  generated_items,
+                                  scheduled_sampling_func):
+    """Scheduled sampling.
+
+    Args:
+      done_warm_start: whether we are done with warm start or not.
+      groundtruth_items: list of ground truth items.
+      generated_items: list of generated items.
+      scheduled_sampling_func: scheduled sampling function to choose between
+        groundtruth items and generated items.
+
+    Returns:
+      A mix list of ground truth and generated items.
+    """
+    def sample():
+      """Calculate the scheduled sampling params based on iteration number."""
+      with tf.variable_scope("scheduled_sampling", reuse=tf.AUTO_REUSE):
+        output_items = []
+        for item_gt, item_gen in zip(groundtruth_items, generated_items):
+          output_items.append(scheduled_sampling_func(item_gt, item_gen))
+        return output_items
+
+    cases = [
+        (tf.logical_not(done_warm_start), lambda: groundtruth_items),
+        (tf.logical_not(self.is_training), lambda: generated_items),
+    ]
+    output_items = tf.case(cases, default=sample, strict=True)
+
+    return output_items
+
   def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
     """Hacky code the get the loss on predicted frames from input frames.
 
@@ -233,17 +327,14 @@ def body(self, features):
       all_actions = input_actions + target_actions
 
     orig_frame_shape = common_layers.shape_list(all_frames[0])
+    batch_size = orig_frame_shape[0]
+    ss_func = self.get_scheduled_sample_func(batch_size)
 
     # Run a number of steps.
     res_frames, sampled_frames, sampled_frames_raw = [], [], []
     extra_loss = 0.0
     if "target_reward" in features:
       res_rewards = []
-    sample_prob = common_layers.inverse_exp_decay(
-        hparams.scheduled_sampling_warmup_steps // 4)
-    sample_prob *= common_layers.inverse_lin_decay(
-        hparams.scheduled_sampling_warmup_steps)
-    sample_prob *= hparams.scheduled_sampling_prob
     for i in range(hparams.video_num_target_frames):
       cur_frames = all_frames[i:i + hparams.video_num_input_frames]
       features["inputs"] = tf.concat(cur_frames, axis=-1)
@@ -276,12 +367,13 @@ def body(self, features):
         all_frames[i + hparams.video_num_input_frames] = sampled_frame
 
       # Scheduled sampling during training.
-      if (hparams.scheduled_sampling_prob > 0.0 and self.is_training):
-        do_sample = tf.less(
-            tf.random_uniform([orig_frame_shape[0]]), sample_prob)
-        orig_frame = all_frames[i + hparams.video_num_input_frames]
-        sampled_frame = tf.where(do_sample, sampled_frame, orig_frame)
-        all_frames[i + hparams.video_num_input_frames] = sampled_frame
+      if self.is_training:
+        done_warm_start = True  # Always true for non-reccurent networks.
+        groundtruth_items = [all_frames[i + hparams.video_num_input_frames]]
+        generated_items = [sampled_frame]
+        ss_frame, = self.get_scheduled_sample_inputs(
+            done_warm_start, groundtruth_items, generated_items, ss_func)
+        all_frames[i + hparams.video_num_input_frames] = ss_frame
 
     # Concatenate results and return them.
     frames = tf.stack(res_frames, axis=1)
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 5a695576b..972d06f04 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -52,6 +52,12 @@ def next_frame_basic_deterministic():
   hparams.add_hparam("small_mode", False)
   hparams.add_hparam("stochastic_model", False)
   hparams.add_hparam("internal_loss", True)
+  # Scheduled sampling method. Choose between
+  # ground_truth_only, prediction_only, prob, count, prob_inverse_exp.
+  hparams.add_hparam("scheduled_sampling_mode", "prediction_only")
+  hparams.add_hparam("scheduled_sampling_decay_steps", 10000)
+  hparams.add_hparam("scheduled_sampling_max_prob", 1.0)
+  hparams.add_hparam("scheduled_sampling_k", 900.0)
   return hparams
 
 
@@ -68,9 +74,9 @@ def next_frame_pixel_noise():
 def next_frame_sampling():
   """Basic conv model with scheduled sampling."""
   hparams = next_frame_basic_deterministic()
-  hparams.video_num_target_frames = 4
-  hparams.scheduled_sampling_warmup_steps = 50000
-  hparams.scheduled_sampling_prob = 0.5
+  hparams.scheduled_sampling_mode = "prob_inverse_exp"
+  hparams.scheduled_sampling_max_prob = 0.5
+  hparams.scheduled_sampling_decay_steps = 10000
   return hparams
 
 
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 8de30ee39..5f7d621dd 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -25,8 +25,6 @@
 from __future__ import division
 from __future__ import print_function
 
-from functools import partial
-
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 
@@ -68,65 +66,6 @@ def concat_on_y_axis(x):
     side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
     tf.summary.image("full_video", side_by_side_video)
 
-  def get_scheduled_sample_func(self, batch_size):
-    """Creates a function for scheduled sampling based on given hparams."""
-    with tf.variable_scope("scheduled_sampling_func", reuse=False):
-      iter_num = self.get_iteration_num()
-      if self.hparams.scheduled_sampling_mode == "prob":
-        decay_steps = self.hparams.scheduled_sampling_decay_steps
-        probability = tf.train.polynomial_decay(
-            1.0, iter_num, decay_steps, 0.0)
-        scheduled_sampling_func = common_video.scheduled_sample_prob
-        scheduled_sampling_func_var = probability
-      else:
-        # Calculate number of ground-truth frames to pass in.
-        k = self.hparams.scheduled_sampling_k
-        num_ground_truth = tf.to_int32(
-            tf.round(
-                tf.to_float(batch_size) *
-                (k / (k + tf.exp(tf.to_float(iter_num) / tf.to_float(k))))))
-        scheduled_sampling_func = common_video.scheduled_sample_count
-        scheduled_sampling_func_var = num_ground_truth
-
-      tf.summary.scalar("scheduled_sampling_var", scheduled_sampling_func_var)
-      partial_func = partial(scheduled_sampling_func,
-                             batch_size=batch_size,
-                             scheduled_sample_var=scheduled_sampling_func_var)
-      return partial_func
-
-  def get_scheduled_sample_inputs(self,
-                                  done_warm_start,
-                                  groundtruth_items,
-                                  generated_items,
-                                  scheduled_sampling_func):
-    """Scheduled sampling.
-
-    Args:
-      done_warm_start: whether we are done with warm start or not.
-      groundtruth_items: list of ground truth items.
-      generated_items: list of generated items.
-      scheduled_sampling_func: scheduled sampling function to choose between
-        groundtruth items and generated items.
-
-    Returns:
-      A mix list of ground truth and generated items.
-    """
-    def sample():
-      """Calculate the scheduled sampling params based on iteration number."""
-      with tf.variable_scope("scheduled_sampling", reuse=tf.AUTO_REUSE):
-        output_items = []
-        for item_gt, item_gen in zip(groundtruth_items, generated_items):
-          output_items.append(scheduled_sampling_func(item_gt, item_gen))
-        return output_items
-
-    cases = [
-        (tf.logical_not(done_warm_start), lambda: groundtruth_items),
-        (tf.logical_not(self.is_training), lambda: generated_items),
-    ]
-    output_items = tf.case(cases, default=sample, strict=True)
-
-    return output_items
-
   def get_input_if_exists(self, features, key, batch_size, num_frames):
     if key in features:
       x = features[key]
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 649d418ef..0a33d7ee3 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -35,6 +35,8 @@ def next_frame_sv2p():
   hparams.target_modality = "video:l2raw"
   hparams.input_modalities = "inputs:video:l2raw"
   hparams.video_modality_loss_cutoff = 0.0
+  hparams.scheduled_sampling_mode = "count"
+  hparams.scheduled_sampling_k = 900.0
   hparams.add_hparam("reward_prediction", True)
   hparams.add_hparam("reward_prediction_stop_gradient", False)
   hparams.add_hparam("reward_prediction_buffer_size", 0)
@@ -43,10 +45,6 @@ def next_frame_sv2p():
   hparams.add_hparam("multi_latent", False)
   hparams.add_hparam("relu_shift", 1e-12)
   hparams.add_hparam("dna_kernel_size", 5)
-  # Scheduled sampling method. Choose between prob or count.
-  hparams.add_hparam("scheduled_sampling_mode", "count")
-  hparams.add_hparam("scheduled_sampling_decay_steps", 10000)
-  hparams.add_hparam("scheduled_sampling_k", 900.0)
   hparams.add_hparam("upsample_method", "conv2d_transpose")
   hparams.add_hparam("reward_model", "basic")
   return hparams

From f3b2add74d9fa4627f08de41383f1d897d99b9ed Mon Sep 17 00:00:00 2001
From: Jeff Wu <WuTheFWasThat@gmail.com>
Date: Wed, 10 Oct 2018 17:08:30 -0700
Subject: [PATCH 0983/2720] fix mtf adafactor implementation for beta1>0
 (#1126)

---
 tensor2tensor/mesh_tensorflow/mtf_optimize.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mtf_optimize.py b/tensor2tensor/mesh_tensorflow/mtf_optimize.py
index c57433808..8b13d3914 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_optimize.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_optimize.py
@@ -154,7 +154,7 @@ def apply_grad(self, grad, var):
     if self._beta1:
       m = mtf.get_variable(
           var.mesh, var.name + "_slot_m", var.shape,
-          iniitalizer=tf.zeros_initializer(), trainable=False)
+          initializer=tf.zeros_initializer(), trainable=False)
 
     with tf.variable_scope(var.name + "/adafactor"):
       grad_squared = mtf.square(grad) + self._epsilon1
@@ -191,7 +191,7 @@ def apply_grad(self, grad, var):
         x /= clipping_denom
       subtrahend = x * update_scale
       if self._beta1:
-        new_m = self._beta1 * m.value + (1.0 - self._beta1) * subtrahend
+        new_m = m * tf.constant(self._beta1) + subtrahend * tf.constant(1.0 - self._beta1)
         subtrahend = new_m
         updates.append(mtf.assign(m, new_m))
       new_val = old_val - subtrahend

From 50fa0fbbf64a9377f4d4111836b9a077e001deec Mon Sep 17 00:00:00 2001
From: Jeff Wu <WuTheFWasThat@gmail.com>
Date: Wed, 10 Oct 2018 17:08:45 -0700
Subject: [PATCH 0984/2720] fixes to mesh tf split (#1128)

---
 tensor2tensor/mesh_tensorflow/mesh_tensorflow.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index 6ff65ab8d..da75330b8 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -1788,7 +1788,7 @@ class SplitOperation(Operation):
   """
 
   def __init__(self, x, split_dim, num_or_size_splits, name=None):
-    super(SplitOperation, self).__init__([x], name=name or "concat")
+    super(SplitOperation, self).__init__([x], name=name or "split")
 
     self._split_dim = split_dim
     if split_dim not in x.shape.dims:
@@ -1804,14 +1804,14 @@ def __init__(self, x, split_dim, num_or_size_splits, name=None):
       assert isinstance(num_or_size_splits, int)
       assert split_dim.size % num_or_size_splits == 0
       self._output_sizes = (
-          [split_dim.size / num_or_size_splits] * num_or_size_splits)
+          [split_dim.size // num_or_size_splits] * num_or_size_splits)
 
     self._outputs = [
         Tensor(self, x.shape.resize_dimension(split_dim.name, output_size),
                x.dtype) for output_size in self._output_sizes]
 
   def gradient(self, grad_ys):
-    return concat(grad_ys, self._split_dim.name)
+    return [concat(grad_ys, self._split_dim.name)]
 
   def lower(self, lowering):
     mesh_impl = lowering.mesh_impl(self)

From 06539f3b7e49937d8f49c312c1eb47b039812ef7 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 10 Oct 2018 17:10:05 -0700
Subject: [PATCH 0985/2720] merge pull requests.

PiperOrigin-RevId: 216619982
---
 tensor2tensor/mesh_tensorflow/mtf_optimize.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/mesh_tensorflow/mtf_optimize.py b/tensor2tensor/mesh_tensorflow/mtf_optimize.py
index 8b13d3914..c70040436 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_optimize.py
+++ b/tensor2tensor/mesh_tensorflow/mtf_optimize.py
@@ -191,7 +191,8 @@ def apply_grad(self, grad, var):
         x /= clipping_denom
       subtrahend = x * update_scale
       if self._beta1:
-        new_m = m * tf.constant(self._beta1) + subtrahend * tf.constant(1.0 - self._beta1)
+        new_m = (m * tf.constant(self._beta1)
+                 + subtrahend * tf.constant(1.0 - self._beta1))
         subtrahend = new_m
         updates.append(mtf.assign(m, new_m))
       new_val = old_val - subtrahend

From 7da0b807abcebd461c1db30acdba7ecc15aac639 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 10 Oct 2018 17:58:25 -0700
Subject: [PATCH 0986/2720] Bring linear scheduled sampling mode back for RL.

PiperOrigin-RevId: 216625609
---
 tensor2tensor/models/video/basic_deterministic.py | 10 ++++++++++
 tensor2tensor/models/video/basic_stochastic.py    |  3 ++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 62a190465..ffa2e2d7f 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -96,6 +96,16 @@ def scheduled_sampling_simple(ground_truth_x, generated_x,
         probability = 1.0 - probability
         scheduled_sampling_func = common_video.scheduled_sample_prob
         scheduled_sampling_func_var = probability
+      elif mode == "prob_inverse_lin":
+        decay_steps = self.hparams.scheduled_sampling_decay_steps
+        probability = common_layers.inverse_exp_decay(
+            decay_steps // 4, step=iter_num)  # Very low at start.
+        probability *= common_layers.inverse_lin_decay(
+            decay_steps, step=iter_num)
+        probability *= self.hparams.scheduled_sampling_max_prob
+        probability = 1.0 - probability
+        scheduled_sampling_func = common_video.scheduled_sample_prob
+        scheduled_sampling_func_var = probability
       elif mode == "count":
         # Calculate number of ground-truth frames to pass in.
         k = self.hparams.scheduled_sampling_k
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 8b94245b5..03d8d803e 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -226,7 +226,8 @@ def next_frame_basic_stochastic_discrete():
   hparams = basic_deterministic_params.next_frame_sampling()
   hparams.batch_size = 2
   hparams.video_num_target_frames = 16
-  hparams.scheduled_sampling_warmup_steps = 40000
+  hparams.scheduled_sampling_mode = "prob_inverse_lin"
+  hparams.scheduled_sampling_decay_steps = 40000
   hparams.scheduled_sampling_prob = 1.0
   hparams.add_hparam("bottleneck_bits", 64)
   hparams.add_hparam("bottleneck_noise", 0.02)

From 75da8f133fac3fc0f3dccd66867d7a7a396715b3 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 10 Oct 2018 23:56:25 -0700
Subject: [PATCH 0987/2720] Supports pretraining the Glow encoder to get a
 stable latent mapping before learning a conditional distribution on top.

PiperOrigin-RevId: 216652264
---
 tensor2tensor/models/research/glow.py         |  12 +-
 tensor2tensor/models/research/glow_ops.py     | 160 ++++++++++--------
 .../models/research/glow_ops_test.py          |   8 +-
 3 files changed, 102 insertions(+), 78 deletions(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 62864be99..8c5173d93 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -105,17 +105,14 @@ def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
 
     return self.scale(predictions)
 
-  def top_prior(self, z):
+  def top_prior(self):
     """Objective based on the prior over latent z.
 
-    Args:
-      z: 4-D Tensor, (batch_size, height, width, num_channels)
     Returns:
-      objective: float, log-likelihood of z under the prior.
       dist: instance of tf.distributions.Normal, prior distribution.
     """
     return glow_ops.top_prior(
-        "top_prior", z, learn_prior=self.hparams.top_prior)
+        "top_prior", self.z_top_shape, learn_prior=self.hparams.top_prior)
 
   def body(self, features):
     x = features["inputs"]
@@ -136,7 +133,10 @@ def body(self, features):
           "codec", x, self.hparams, eps=None, reverse=False)
       objective += encoder_objective
 
-      prior_objective, prior_dist = self.top_prior(self.z)
+      self.z_top_shape = common_layers.shape_list(self.z)
+      prior_dist = self.top_prior()
+      prior_objective = tf.reduce_sum(
+          prior_dist.log_prob(self.z), axis=[1, 2, 3])
       tf.summary.scalar("top_prior", tf.reduce_mean(prior_objective))
       self.z_sample = prior_dist.sample()
       objective += prior_objective
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index faf6dec16..b23923971 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -548,7 +548,62 @@ def merge_level_and_latent_dist(level_dist, latent_dist,
 
 
 @add_arg_scope
-def compute_prior(name, z, latent, hparams, state=None):
+def level_cond_prior(prior_dist, z, latent, hparams, state):
+  """Returns a conditional prior for each level.
+
+  Args:
+    prior_dist: Distribution conditioned on the previous levels.
+    z: Tensor, output of the previous levels.
+    latent: Tensor or a list of tensors to condition the latent_distribution.
+    hparams: next_frame_glow hparams.
+    state: Current LSTM state. Used only if hparams.latent_dist_encoder is
+           a lstm.
+  Raises:
+    ValueError: If hparams.latent_dist_encoder is "pointwise" and if the shape
+                of latent is different from z.
+  """
+  latent_dist_encoder = hparams.get("latent_dist_encoder", None)
+  latent_skip = hparams.get("latent_skip", False)
+  if latent_dist_encoder == "pointwise":
+    merge_std = hparams.level_scale
+    latent_shape = common_layers.shape_list(latent)
+    z_shape = common_layers.shape_list(z)
+    if latent_shape != z_shape:
+      raise ValueError("Expected latent_shape to be %s, got %s" %
+                       (latent_shape, z_shape))
+    latent_dist = scale_gaussian_prior(
+        "latent_prior", latent, logscale_factor=3.0)
+    cond_dist = merge_level_and_latent_dist(prior_dist, latent_dist,
+                                            merge_std=merge_std)
+  elif latent_dist_encoder == "conv_net":
+    output_channels = common_layers.shape_list(z)[-1]
+    latent_stack = tf.concat([prior_dist.loc] + latent, axis=-1)
+    cond_dist = tensor_to_dist(
+        "latent_stack", latent_stack, output_channels=output_channels,
+        architecture=hparams.latent_architecture,
+        depth=hparams.latent_encoder_depth,
+        pre_output_channels=hparams.latent_pre_output_channels)
+    if latent_skip:
+      cond_dist = tf.distributions.Normal(
+          cond_dist.loc + latent[-1], cond_dist.scale)
+  elif latent_dist_encoder == "conv_lstm":
+    output_channels = common_layers.shape_list(z)[-1]
+    latent_stack = tf.concat((prior_dist.loc, latent), axis=-1)
+    _, state = common_video.conv_lstm_2d(
+        latent_stack, state, output_channels, kernel_size=3,
+        name="conv_lstm")
+    cond_dist = tensor_to_dist(
+        "state_to_dist", state.h, output_channels=output_channels)
+    if latent_skip:
+      cond_dist = tf.distributions.Normal(
+          cond_dist.loc + latent, cond_dist.scale)
+  tf.summary.histogram("split_prior_mean", prior_dist.loc)
+  tf.summary.histogram("split_prior_scale", prior_dist.scale)
+  return cond_dist.loc, cond_dist.scale, state
+
+
+@add_arg_scope
+def compute_prior(name, z, latent, hparams, condition=False, state=None):
   """Distribution on z_t conditioned on z_{t-1} and latent.
 
   Args:
@@ -560,6 +615,7 @@ def compute_prior(name, z, latent, hparams, state=None):
             else, this is just a 4-D Tensor
             The first-three dimensions of the latent should be the same as z.
     hparams: next_frame_glow_hparams.
+    condition: Whether or not to condition the distribution on latent.
     state: tf.contrib.rnn.LSTMStateTuple.
            the current state of a LSTM used to model the distribution. Used
            only if hparams.latent_dist_encoder = "conv_lstm".
@@ -571,54 +627,25 @@ def compute_prior(name, z, latent, hparams, state=None):
                 of latent is different from z.
   """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    if isinstance(condition, bool):
+      condition = tf.constant(condition, dtype=tf.bool)
     prior_dist = tensor_to_dist("level_prior", z, architecture="single_conv")
-
-    # TODO(mechcoder) Refactor into separate sub-functions.
-    if latent is not None:
-      latent_dist_encoder = hparams.get("latent_dist_encoder", None)
-      latent_skip = hparams.get("latent_skip", False)
-      if latent_dist_encoder == "pointwise":
-        merge_std = hparams.level_scale
-        latent_shape = common_layers.shape_list(latent)
-        z_shape = common_layers.shape_list(z)
-        if latent_shape != z_shape:
-          raise ValueError("Expected latent_shape to be %s, got %s" %
-                           (latent_shape, z_shape))
-        latent_dist = scale_gaussian_prior(
-            "latent_prior", latent, logscale_factor=3.0)
-        prior_dist = merge_level_and_latent_dist(prior_dist, latent_dist,
-                                                 merge_std=merge_std)
-      elif latent_dist_encoder == "conv_net":
-        output_channels = common_layers.shape_list(z)[-1]
-        latent_stack = tf.concat([prior_dist.loc] + latent, axis=-1)
-        prior_dist = tensor_to_dist(
-            "latent_stack", latent_stack, output_channels=output_channels,
-            architecture=hparams.latent_architecture,
-            depth=hparams.latent_encoder_depth,
-            pre_output_channels=hparams.latent_pre_output_channels)
-        if latent_skip:
-          prior_dist = tf.distributions.Normal(
-              prior_dist.loc + latent[-1], prior_dist.scale)
-      elif latent_dist_encoder == "conv_lstm":
-        output_channels = common_layers.shape_list(z)[-1]
-        latent_stack = tf.concat((prior_dist.loc, latent), axis=-1)
-        _, state = common_video.conv_lstm_2d(
-            latent_stack, state, output_channels, kernel_size=3,
-            name="conv_lstm")
-        prior_dist = tensor_to_dist(
-            "state_to_dist", state.h, output_channels=output_channels)
-        if latent_skip:
-          prior_dist = tf.distributions.Normal(
-              prior_dist.loc + latent, prior_dist.scale)
-      tf.summary.histogram("split_prior_mean", prior_dist.loc)
-      tf.summary.histogram("split_prior_scale", prior_dist.scale)
-
-  return prior_dist, state
+    prior_mean, prior_scale = prior_dist.loc, prior_dist.scale
+    if latent is None:
+      mean, scale = prior_mean, prior_scale
+    else:
+      cond_mean, cond_scale, state = level_cond_prior(
+          prior_dist, z, latent, hparams, state)
+      mean, scale = tf.cond(
+          condition, lambda: (cond_mean, cond_scale),
+          lambda: (prior_mean, prior_scale))
+    dist = tf.distributions.Normal(mean, scale)
+    return dist, state
 
 
 @add_arg_scope
 def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
-          hparams=None, state=None):
+          hparams=None, state=None, condition=False):
   """Splits / concatenates x into x1 and x2 across number of channels.
 
   For the forward pass, x2 is assumed be gaussian,
@@ -637,6 +664,8 @@ def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
     hparams: next_frame_glow hparams.
     state: tf.contrib.rnn.LSTMStateTuple. Current state of the LSTM over z_2.
            Used only when hparams.latent_dist_encoder == "conv_lstm"
+    condition: bool, Whether or not to condition the distribution on
+               cond_latents.
 
   Returns:
   Raises:
@@ -650,13 +679,13 @@ def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
 
       # objective: P(x2|x1) ~N(x2 ; NN(x1))
       prior_dist, state = compute_prior(
-          "prior_on_z2", x1, cond_latents, hparams, state=state)
+          "prior_on_z2", x1, cond_latents, hparams, condition, state=state)
       logpb = tf.reduce_sum(prior_dist.log_prob(x2), axis=[1, 2, 3])
       eps = get_eps(prior_dist, x2)
       return x1, logpb, eps, x2, state
     else:
       prior_dist, state = compute_prior(
-          "prior_on_z2", x, cond_latents, hparams, state=state)
+          "prior_on_z2", x, cond_latents, hparams, condition, state=state)
       if eps is not None:
         x2 = set_eps(prior_dist, eps)
       elif eps_std is not None:
@@ -742,12 +771,12 @@ def scale_gaussian_prior(name, z, logscale_factor=3.0, trainable=True):
 
 
 @add_arg_scope
-def top_prior(name, x, learn_prior="normal"):
-  """Log probability of x being gaussian.
+def top_prior(name, z_shape, learn_prior="normal"):
+  """Unconditional prior distribution.
 
   Args:
     name: variable scope
-    x: input, 4-D Tensor shape=(batch_size, width, height, channels)
+    z_shape: Shape of the mean / scale of the prior distribution.
     learn_prior: Possible options are "normal" and "single_conv".
                  If set to "single_conv", the gaussian is parametrized by a
                  single convolutional layer whose input are an array of zeros
@@ -760,7 +789,7 @@ def top_prior(name, x, learn_prior="normal"):
     ValueError: If learn_prior not in "normal" or "single_conv"
   """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
-    h = tf.zeros_like(x)
+    h = tf.zeros(z_shape, dtype=tf.float32)
     if learn_prior == "normal":
       prior_dist = tf.distributions.Normal(h, tf.exp(h))
     elif learn_prior == "single_conv":
@@ -768,8 +797,7 @@ def top_prior(name, x, learn_prior="normal"):
     else:
       raise ValueError("Expected learn_prior to be normal or single_conv "
                        "got %s" % learn_prior)
-    objective = tf.reduce_sum(prior_dist.log_prob(x), axis=[1, 2, 3])
-    return objective, prior_dist
+    return prior_dist
 
 
 def uniform_binning_correction(x, n_bits=8):
@@ -795,7 +823,7 @@ def uniform_binning_correction(x, n_bits=8):
 
 @add_arg_scope
 def encoder_decoder(name, x, hparams, eps=None, reverse=False,
-                    cond_latents=None, states=None):
+                    cond_latents=None, condition=False, states=None):
   """Glow encoder-decoder. n_levels of (Squeeze + Flow + Split.) operations."""
   # TODO(mechcoder) Change return_type to a dict to be backward compatible.
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
@@ -803,9 +831,13 @@ def encoder_decoder(name, x, hparams, eps=None, reverse=False,
     if states and len(states) != hparams.n_levels - 1:
       raise ValueError("Expected length of states to be %d, got %d" %
                        (hparams.n_levels - 1, len(states)))
+    if states is None:
+      states = [None] * (hparams.n_levels - 1)
     if eps and len(eps) != hparams.n_levels - 1:
       raise ValueError("Expected length of eps to be %d, got %d" %
                        (hparams.n_levels - 1, len(eps)))
+    if eps is None:
+      eps = [None] * (hparams.n_levels - 1)
     check_cond_latents(cond_latents, hparams)
 
     objective = 0.0
@@ -822,40 +854,30 @@ def encoder_decoder(name, x, hparams, eps=None, reverse=False,
         objective += obj
 
         if level < hparams.n_levels - 1:
-
-          curr_state = None
-          if states:
-            curr_state = states[level]
-
           curr_cond_latents = get_cond_latents_at_level(
               cond_latents, level, hparams)
           x, obj, eps, z, state = split("split_%d" % level, x, reverse=False,
                                         cond_latents=curr_cond_latents,
-                                        hparams=hparams, state=curr_state)
+                                        condition=condition,
+                                        hparams=hparams, state=states[level])
           objective += obj
           all_eps.append(eps)
           all_latents.append(z)
           new_states.append(state)
+
       return x, objective, all_eps, all_latents, new_states
 
     else:
       for level in reversed(range(hparams.n_levels)):
         if level < hparams.n_levels - 1:
 
-          curr_eps = None
-          if eps:
-            curr_eps = eps[level]
-
-          curr_state = None
-          if states:
-            curr_state = states[level]
-
           curr_cond_latents = get_cond_latents_at_level(
               cond_latents, level, hparams)
 
-          x, latent, state = split("split_%d" % level, x, eps=curr_eps,
+          x, latent, state = split("split_%d" % level, x, eps=eps[level],
                                    reverse=True, cond_latents=curr_cond_latents,
-                                   hparams=hparams, state=curr_state)
+                                   condition=condition, hparams=hparams,
+                                   state=states[level])
           new_states.append(state)
           all_latents.append(latent)
 
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 6aa889dbd..a04400ede 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -280,11 +280,12 @@ def check_split_latent_conditioning(self, merge_std):
       exp_x2 = x_rand[:, :, :, 16:]
       exp_eps = x_rand[:, :, :, 16:] - latent_rand
       x_inv, _, eps, x2_t, _ = glow_ops.split(
-          merge_std, x_t, cond_latents=latent_t, hparams=hparams)
+          merge_std, x_t, cond_latents=latent_t, hparams=hparams,
+          condition=True)
       # Test reversibility.
       x_inv_inv, _, _ = glow_ops.split(
           merge_std, x_inv, cond_latents=latent_t, eps=eps, reverse=True,
-          hparams=hparams)
+          hparams=hparams, condition=True)
       with tf.Session() as sess:
         sess.run(tf.global_variables_initializer())
         actual_eps, actual_x2, diff_np = sess.run([eps, x2_t, x_inv_inv - x_t])
@@ -312,7 +313,8 @@ def test_latent_dist_encoder_lstm(self):
       hparams.add_hparam("latent_skip", True)
 
       prior_dist, new_state = glow_ops.compute_prior(
-          "lstm_prior", x_t, latent=latent_t, hparams=hparams, state=init_state)
+          "lstm_prior", x_t, latent=latent_t, hparams=hparams, state=init_state,
+          condition=True)
       with tf.Session() as sess:
         sess.run(tf.global_variables_initializer())
         # Test initialization (mu, sigma) = (z, 1.0)

From fb92e62ae534eda925ab09aa4938586adbc7f8f9 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 11 Oct 2018 11:49:24 -0700
Subject: [PATCH 0988/2720] Allow to pass Problem class as well as name in
 trainer_lib.

PiperOrigin-RevId: 216734941
---
 tensor2tensor/utils/trainer_lib.py      |  8 ++++++--
 tensor2tensor/utils/trainer_lib_test.py | 16 ++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 48d888385..44c69fbd8 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -24,6 +24,7 @@
 import random
 import numpy as np
 
+from tensor2tensor.data_generators.problem import Problem
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import devices
 from tensor2tensor.utils import metrics_hook
@@ -583,9 +584,12 @@ def experiment_fn(run_config, hparams):
 
 def add_problem_hparams(hparams, problem_name):
   """Add problem hparams for the problems."""
-  problem = registry.problem(problem_name)
+  print(problem_name)
+  if isinstance(problem_name, Problem):
+    problem = problem_name
+  else:
+    problem = registry.problem(problem_name)
   p_hparams = problem.get_hparams(hparams)
-
   hparams.problem = problem
   hparams.problem_hparams = p_hparams
 
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index 55b5a5aa0..01615d929 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -49,6 +49,22 @@ def testExperiment(self):
     exp = exp_fn(run_config, hparams)
     exp.test()
 
+  def testExperimentWithClass(self):
+    exp_fn = trainer_lib.create_experiment_fn(
+        "transformer",
+        algorithmic.TinyAlgo(),
+        algorithmic.TinyAlgo.data_dir,
+        train_steps=1,
+        eval_steps=1,
+        min_eval_frequency=1,
+        use_tpu=False)
+    run_config = trainer_lib.create_run_config(
+        model_dir=algorithmic.TinyAlgo.data_dir, num_gpus=0,
+        use_tpu=False)
+    hparams = registry.hparams("transformer_tiny_tpu")
+    exp = exp_fn(run_config, hparams)
+    exp.test()
+
   def testModel(self):
     # HParams
     hparams = trainer_lib.create_hparams(

From be3ced5d6f754d5eff485b2623b51514d429e7f7 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 11 Oct 2018 13:59:52 -0700
Subject: [PATCH 0989/2720] Removing a print that slipped through and using
 more appropriate variable names.

PiperOrigin-RevId: 216756203
---
 tensor2tensor/utils/trainer_lib.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 44c69fbd8..f0be376f9 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -582,13 +582,12 @@ def experiment_fn(run_config, hparams):
   return experiment_fn
 
 
-def add_problem_hparams(hparams, problem_name):
+def add_problem_hparams(hparams, problem_name_or_instance):
   """Add problem hparams for the problems."""
-  print(problem_name)
-  if isinstance(problem_name, Problem):
-    problem = problem_name
+  if isinstance(problem_name_or_instance, Problem):
+    problem = problem_name_or_instance
   else:
-    problem = registry.problem(problem_name)
+    problem = registry.problem(problem_name_or_instance)
   p_hparams = problem.get_hparams(hparams)
   hparams.problem = problem
   hparams.problem_hparams = p_hparams

From 1b16365fef49747163939c86686043eb195c061f Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 12 Oct 2018 01:31:37 +0200
Subject: [PATCH 0990/2720] Move Model-Based RL hparams to separate module.
 (#1137)

---
 tensor2tensor/rl/trainer_model_based.py       | 725 +----------------
 .../rl/trainer_model_based_agent_only.py      |   3 +-
 .../rl/trainer_model_based_params.py          | 753 ++++++++++++++++++
 3 files changed, 757 insertions(+), 724 deletions(-)
 create mode 100644 tensor2tensor/rl/trainer_model_based_params.py

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 03fe2a525..bde551c6a 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -44,6 +44,7 @@
 from tensor2tensor.data_generators import gym_problems_specs
 from tensor2tensor.layers import discretization
 from tensor2tensor.rl import rl_trainer_lib
+from tensor2tensor.rl import trainer_model_based_params
 from tensor2tensor.rl.envs.tf_atari_wrappers import PyFuncWrapper
 from tensor2tensor.rl.envs.utils import InitialFrameChooser
 from tensor2tensor.utils import registry
@@ -63,8 +64,6 @@
 flags.DEFINE_string("eval_results_dir", "/tmp",
                     "Directory to store result of evaluation")
 
-HP_SCOPES = ["loop", "model", "ppo"]
-
 
 def setup_directories(base_dir, subdirs):
   base_dir = os.path.expanduser(base_dir)
@@ -748,728 +747,8 @@ def combine_training_data(problem, final_data_dir, old_data_dirs,
         tf.gfile.Copy(fname, new_fname)
 
 
-@registry.register_hparams
-def rlmb_base():
-  return tf.contrib.training.HParams(
-      epochs=15,
-      # Total frames used for training. This will be distributed evenly across
-      # hparams.epochs.
-      # This number should be divisible by real_ppo_epoch_length*epochs
-      # for our frame accounting to be preceise.
-      num_real_env_frames=96000,
-      generative_model="next_frame_basic_deterministic",
-      generative_model_params="next_frame_pixel_noise",
-      ppo_params="ppo_pong_base",
-      autoencoder_train_steps=0,
-      autoencoder_train_steps_initial_multiplier=10,
-      autoencoder_hparams_set="autoencoder_discrete_pong",
-      model_train_steps=15000,
-      inital_epoch_train_steps_multiplier=3,
-      simulated_env_generator_num_steps=2000,
-      simulation_random_starts=True,  # Use random starts in PPO.
-      # Flip the first random frame in PPO batch for the true beginning.
-      simulation_flip_first_random_for_beginning=True,
-      intrinsic_reward_scale=0.,
-      ppo_epochs_num=1000,  # This should be enough to see something
-      # Our simulated envs do not know how to reset.
-      # You should set ppo_time_limit to the value you believe that
-      # the simulated env produces a reasonable output.
-      ppo_time_limit=200,  # TODO(blazej): this param is unused
-      # It makes sense to have ppo_time_limit=ppo_epoch_length,
-      # though it is not necessary.
-      ppo_epoch_length=50,
-      ppo_num_agents=16,
-      # Do not eval since simulated batch env does not produce dones
-      ppo_eval_every_epochs=0,
-      ppo_learning_rate=1e-4,  # Will be changed, just so it exists.
-      # Whether the PPO agent should be restored from the previous iteration, or
-      # should start fresh each time.
-      ppo_continue_training=True,
-      # Resizing.
-      resize_height_factor=2,
-      resize_width_factor=2,
-      grayscale=False,
-      # Bump learning rate after first epoch by 3x.
-      # We picked 3x because our default learning rate schedule decreases with
-      # 1/square root of step; 1/sqrt(10k) = 0.01 and 1/sqrt(100k) ~ 0.0032
-      # so by bumping it up 3x we about "go back" from 100k steps to 10k, which
-      # is approximately as much as "going back 1 epoch" would be.
-      # In your experiments, you want to optimize this rate to your schedule.
-      learning_rate_bump=3.0,
-
-      gather_ppo_real_env_data=True,
-      real_ppo_epochs_num=0,
-      # This needs to be divisible by real_ppo_effective_num_agents.
-      real_ppo_epoch_length=16*200,
-      real_ppo_num_agents=1,
-      real_ppo_learning_rate=1e-4,
-      real_ppo_continue_training=True,
-      real_ppo_effective_num_agents=16,
-      real_ppo_eval_every_epochs=0,
-
-      game="pong",
-      # Whether to evaluate the world model in each iteration of the loop to get
-      # the model_reward_accuracy metric.
-      eval_world_model=True,
-      # Rollout fractions to report reward_accuracy on.
-      eval_rollout_fractions=[0.25, 0.5, 1],
-      stop_loop_early=False,  # To speed-up tests.
-  )
-
-
-@registry.register_hparams
-def rlmb_basetest():
-  """Base setting but quicker with only 2 epochs."""
-  hparams = rlmb_base()
-  hparams.game = "pong"
-  hparams.epochs = 2
-  hparams.num_real_env_frames = 3200
-  hparams.model_train_steps = 100
-  hparams.simulated_env_generator_num_steps = 20
-  hparams.ppo_epochs_num = 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_noresize():
-  hparams = rlmb_base()
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_quick():
-  """Base setting but quicker with only 2 epochs."""
-  hparams = rlmb_base()
-  hparams.epochs = 2
-  hparams.model_train_steps = 25000
-  hparams.ppo_epochs_num = 700
-  hparams.ppo_epoch_length = 50
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_quick_noresize():
-  hparams = rlmb_base()
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_quick_sd():
-  """Quick setting with stochastic discrete model."""
-  hparams = rlmb_quick()
-  hparams.generative_model = "next_frame_basic_stochastic_discrete"
-  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_sdtest():
-  """Test setting with stochastic discrete model."""
-  hparams = rlmb_basetest()
-  hparams.generative_model = "next_frame_basic_stochastic_discrete"
-  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_quick_sm():
-  """Quick setting with sampling."""
-  hparams = rlmb_quick()
-  hparams.generative_model_params = "next_frame_sampling"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_stochastic():
-  """Base setting with a stochastic next-frame model."""
-  hparams = rlmb_base()
-  hparams.inital_epoch_train_steps_multiplier = 5
-  hparams.generative_model = "next_frame_basic_stochastic"
-  hparams.generative_model_params = "next_frame_basic_stochastic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sampling_stochastic():
-  """Base setting with a stochastic next-frame model."""
-  hparams = rlmb_base()
-  hparams.generative_model = "next_frame_basic_stochastic"
-  hparams.generative_model_params = "next_frame_sampling_stochastic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_stochastic_discrete():
-  """Base setting with stochastic discrete model."""
-  hparams = rlmb_base()
-  hparams.learning_rate_bump = 1.0
-  hparams.grayscale = True
-  hparams.generative_model = "next_frame_basic_stochastic_discrete"
-  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_stochastic_recurrent():
-  """Base setting with recurrent model."""
-  hparams = rlmb_base()
-  hparams.generative_model = "next_frame_basic_recurrent"
-  hparams.generative_model_params = "next_frame_basic_recurrent"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_stochastic_discrete_noresize():
-  """Base setting with stochastic discrete model."""
-  hparams = rlmb_base()
-  hparams.generative_model = "next_frame_basic_stochastic_discrete"
-  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p():
-  """Base setting with sv2p as world model."""
-  hparams = rlmb_base()
-  hparams.learning_rate_bump = 1.0
-  hparams.generative_model = "next_frame_sv2p"
-  hparams.generative_model_params = "next_frame_sv2p_atari"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_softmax():
-  """Base setting with sv2p as world model with softmax."""
-  hparams = rlmb_base_sv2p()
-  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_deterministic():
-  """Base setting with deterministic sv2p as world model."""
-  hparams = rlmb_base_sv2p()
-  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_deterministic_softmax():
-  """Base setting with deterministic sv2p as world model with softmax."""
-  hparams = rlmb_base_sv2p_softmax()
-  hparams.generative_model_params = (
-      "next_frame_sv2p_atari_softmax_deterministic")
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_flippy30():
-  """Base setting with sv2p as world model."""
-  hparams = rlmb_base()
-  hparams.epochs = 30
-  hparams.ppo_epochs_num = 1000
-  hparams.model_train_steps = 15000
-  hparams.learning_rate_bump = 1.0
-  hparams.inital_epoch_train_steps_multiplier = 5
-  hparams.generative_model = "next_frame_sv2p"
-  hparams.generative_model_params = "next_frame_sv2p_atari"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_softmax_flippy30():
-  """Base setting with sv2p as world model with softmax."""
-  hparams = rlmb_base_sv2p_flippy30()
-  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_deterministic_flippy30():
-  """Base setting with deterministic sv2p as world model."""
-  hparams = rlmb_base_sv2p_flippy30()
-  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_deterministic_softmax_flippy30():
-  """Base setting with deterministic sv2p as world model with softmax."""
-  hparams = rlmb_base_sv2p_softmax_flippy30()
-  hparams.generative_model_params = (
-      "next_frame_sv2p_atari_softmax_deterministic")
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sampling():
-  """Base setting with a stochastic next-frame model."""
-  hparams = rlmb_base()
-  hparams.generative_model_params = "next_frame_sampling"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sampling_noresize():
-  hparams = rlmb_base_sampling()
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_flippy60():
-  """Schedule with a lot of epochs (slow)."""
-  hparams = rlmb_base_sampling()
-  hparams.epochs = 60
-  hparams.ppo_epochs_num = 500
-  hparams.model_train_steps = 10000
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_flippy30():
-  """Schedule with a lot of epochs (slow)."""
-  hparams = rlmb_base_sampling()
-  hparams.epochs = 30
-  hparams.ppo_epochs_num = 1000
-  hparams.model_train_steps = 15000
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_medium():
-  """Small set for larger testing."""
-  hparams = rlmb_base()
-  hparams.num_real_env_frames //= 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_25k():
-  """Small set for larger testing."""
-  hparams = rlmb_medium()
-  hparams.num_real_env_frames //= 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_short():
-  """Small set for larger testing."""
-  hparams = rlmb_base()
-  hparams.num_real_env_frames //= 5
-  hparams.model_train_steps //= 10
-  hparams.ppo_epochs_num //= 10
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_model_only():
-  hp = rlmb_base()
-  hp.epochs = 1
-  hp.ppo_epochs_num = 0
-  return hp
-
-
-@registry.register_hparams
-def rlmb_tiny():
-  """Tiny set for testing."""
-  return rlmb_base_sampling().override_from_dict(
-      tf.contrib.training.HParams(
-          epochs=1,
-          num_real_env_frames=128,
-          simulated_env_generator_num_steps=64,
-          model_train_steps=2,
-          ppo_epochs_num=2,
-          ppo_time_limit=5,
-          ppo_epoch_length=5,
-          ppo_num_agents=2,
-          real_ppo_epoch_length=36,
-          real_ppo_num_agents=1,
-          real_ppo_epochs_num=0,
-          real_ppo_effective_num_agents=2,
-          generative_model_params="next_frame_tiny",
-          stop_loop_early=True,
-          resize_height_factor=2,
-          resize_width_factor=2,
-          game="pong",
-      ).values())
-
-
-@registry.register_hparams
-def rlmb_tiny_stochastic():
-  """Tiny setting with a stochastic next-frame model."""
-  hparams = rlmb_tiny()
-  hparams.epochs = 1  # Too slow with 2 for regular runs.
-  hparams.generative_model = "next_frame_basic_stochastic"
-  hparams.generative_model_params = "next_frame_basic_stochastic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_tiny_recurrent():
-  """Tiny setting with a recurrent next-frame model."""
-  hparams = rlmb_tiny()
-  hparams.epochs = 1  # Too slow with 2 for regular runs.
-  hparams.generative_model = "next_frame_basic_recurrent"
-  hparams.generative_model_params = "next_frame_basic_recurrent"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_tiny_sv2p():
-  """Tiny setting with a tiny sv2p model."""
-  hparams = rlmb_tiny()
-  hparams.generative_model = "next_frame_sv2p"
-  hparams.generative_model_params = "next_frame_sv2p_tiny"
-  hparams.grayscale = False
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_ae_base():
-  """Parameter set for autoencoders."""
-  hparams = rlmb_base()
-  hparams.ppo_params = "ppo_pong_ae_base"
-  hparams.generative_model_params = "next_frame_ae"
-  hparams.autoencoder_hparams_set = "autoencoder_discrete_pong"
-  hparams.gather_ppo_real_env_data = False
-  hparams.autoencoder_train_steps = 5000
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  hparams.grayscale = False
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_ae_basetest():
-  """Base AE setting but quicker with only 2 epochs."""
-  hparams = rlmb_ae_base()
-  hparams.game = "pong"
-  hparams.epochs = 2
-  hparams.num_real_env_frames = 3200
-  hparams.model_train_steps = 100
-  hparams.autoencoder_train_steps = 10
-  hparams.simulated_env_generator_num_steps = 20
-  hparams.ppo_epochs_num = 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_ae_tiny():
-  """Tiny set for testing autoencoders."""
-  hparams = rlmb_tiny()
-  hparams.ppo_params = "ppo_pong_ae_base"
-  hparams.generative_model_params = "next_frame_ae_tiny"
-  hparams.autoencoder_hparams_set = "autoencoder_discrete_tiny"
-  hparams.gather_ppo_real_env_data = False
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  hparams.grayscale = False
-  hparams.autoencoder_train_steps = 1
-  hparams.autoencoder_train_steps_initial_multiplier = 0
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_tiny_simulation_deterministic_starts():
-  hp = rlmb_tiny()
-  hp.simulation_random_starts = False
-  return hp
-
-
-# RangedHParams for tuning
-# ==============================================================================
-# Note that the items here must be scoped with one of
-# HP_SCOPES={loop, model, ppo}, which set hyperparameters for the top-level
-# hparams, hp.generative_model_params, and hp.ppo_params, respectively.
-@registry.register_ranged_hparams
-def rlmb_grid(rhp):
-  """Grid over games and frames, and 5 runs each for variance."""
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-  base = 100000
-  medium = base // 2
-  small = medium // 2
-  rhp.set_discrete("loop.num_real_env_frames", [base, medium, small])
-
-  # Dummy parameter to get 5 runs for each configuration
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-
-
-@registry.register_ranged_hparams
-def rlmb_variance(rhp):
-  # Dummy parameter to get 5 runs for each configuration
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-
-
-@registry.register_ranged_hparams
-def rlmb_variance_nogame(rhp):
-  # Dummy parameter to get 20 runs for current configuration.
-  rhp.set_discrete("model.moe_loss_coef", list(range(20)))
-
-
-@registry.register_ranged_hparams
-def rlmb_three(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
-
-
-@registry.register_ranged_hparams
-def rlmb_test1(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
-  rhp.set_discrete("loop.ppo_learning_rate", [5e-5, 1e-4, 2e-4])
-  rhp.set_discrete("ppo.optimization_batch_size", [20, 40])
-  rhp.set_discrete("loop.epochs", [3, 6])
-
-
-@registry.register_ranged_hparams
-def rlmb_scheduled_sampling(rhp):
-  rhp.set_float("model.scheduled_sampling_prob", 0.0, 1.0)
-
-
-@registry.register_ranged_hparams
-def rlmb_all_games(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_whitelisted_games(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_human_score_games(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game",
-                      gym_problems_specs.ATARI_GAMES_WITH_HUMAN_SCORE)
-
-
-@registry.register_ranged_hparams
-def rlmb_curious_games10(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_curious_games5(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_debug_games(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_DEBUG_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_ae_variance(rhp):
-  # Dummy parameter to get 5 runs for each configuration
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-  base = 100000
-  small = base // 4
-  rhp.set_discrete("loop.num_real_env_frames", [base, small])
-
-
-@registry.register_ranged_hparams
-def rlmb_ppolr_game(rhp):
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-  base_lr = 1e-4
-  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
-
-
-@registry.register_ranged_hparams
-def rlmb_ppolr(rhp):
-  base_lr = 1e-4
-  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
-
-
-@registry.register_ranged_hparams
-def rlmb_ae_ppo_lr(rhp):
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-  base_lr = 1e-4
-  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
-
-
-@registry.register_ranged_hparams
-def rlmb_dropout_range(rhp):
-  rhp.set_float("model.dropout", 0.2, 0.4)
-
-
-@registry.register_ranged_hparams
-def rlmb_intrinsic_reward_scale(rhp):
-  rhp.set_float("loop.intrinsic_reward_scale", 0.01, 10.)
-
-
-@registry.register_ranged_hparams
-def rlmb_l1l2cutoff_range(rhp):
-  """Loss and loss-cutoff tuning grid."""
-  rhp.set_float("model.video_modality_loss_cutoff", 1.4, 3.4)
-
-
-@registry.register_ranged_hparams
-def rlmb_xentcutoff_range(rhp):
-  """Cross entropy cutoff tuning grid."""
-  rhp.set_float("model.video_modality_loss_cutoff", 0.01, 0.05)
-
-
-@registry.register_ranged_hparams
-def rlmb_pixel_noise(rhp):
-  """Input pixel noise tuning grid."""
-  rhp.set_categorical("loop.generative_model_params",
-                      ["next_frame_pixel_noise"])
-  rhp.set_discrete("model.video_modality_input_noise",
-                   [0.0025 * i for i in range(200)])
-
-
-@registry.register_ranged_hparams
-def rlmb_dummy_range(rhp):
-  """Dummy tuning grid just to get the variance."""
-  rhp.set_float("model.moe_loss_coef", 0.01, 0.02)
-
-
-@registry.register_ranged_hparams
-def rlmb_epochs_num(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_discrete("loop.epochs", [3, 6, 12])
-
-
-@registry.register_ranged_hparams
-def rlmb_ppo_epochs_num(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_discrete("loop.ppo_epochs_num", [200, 1000, 2000, 4000])
-
-
-@registry.register_ranged_hparams
-def rlmb_ppo_epoch_len(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_discrete("loop.ppo_epoch_length", [25, 50, 100])
-
-
-@registry.register_ranged_hparams
-def rlmb_num_frames(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_discrete("loop.num_real_env_frames",
-                   [1000*el for el in [30, 100, 500, 1000]])
-
-
-@registry.register_ranged_hparams
-def rlmb_ppo_optimization_batch_size(rhp):
-  rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_discrete("ppo.optimization_batch_size", [4, 10, 20])
-
-
-@registry.register_ranged_hparams
-def rlmb_logits_clip(rhp):
-  rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_discrete("ppo.logits_clip", [0., 5.])
-
-
-def merge_unscoped_hparams(scopes_and_hparams):
-  """Merge multiple HParams into one with scopes."""
-  merged_values = {}
-  for (scope, hparams) in scopes_and_hparams:
-    for key, value in six.iteritems(hparams.values()):
-      scoped_key = "%s.%s" % (scope, key)
-      merged_values[scoped_key] = value
-
-  return tf.contrib.training.HParams(**merged_values)
-
-
-def split_scoped_hparams(scopes, merged_hparams):
-  """Split single HParams with scoped keys into multiple."""
-  split_values = dict([(scope, dict()) for scope in scopes])
-  merged_values = merged_hparams.values()
-  for scoped_key, value in six.iteritems(merged_values):
-    scope = scoped_key.split(".")[0]
-    key = scoped_key[len(scope) + 1:]
-    split_values[scope][key] = value
-
-  return [
-      tf.contrib.training.HParams(**split_values[scope]) for scope in scopes
-  ]
-
-
-def training_loop_hparams_from_scoped_overrides(scoped_overrides, trial_id):
-  """Create HParams suitable for training loop from scoped HParams.
-
-  Args:
-    scoped_overrides: HParams, with keys all scoped by one of HP_SCOPES. These
-      parameters are overrides for the base HParams created by
-      create_loop_hparams.
-    trial_id: str, trial identifier. This is used to register unique HParams
-      names for the underlying model and ppo HParams.
-
-  Returns:
-    HParams suitable for passing to training_loop.
-  """
-  trial_hp_overrides = scoped_overrides.values()
-
-  # Create loop, model, and ppo base HParams
-  loop_hp = create_loop_hparams()
-  model_hp_name = trial_hp_overrides.get(
-      "loop.generative_model_params", loop_hp.generative_model_params)
-  model_hp = registry.hparams(model_hp_name).parse(FLAGS.hparams)
-  ppo_params_name = trial_hp_overrides.get(
-      "loop.ppo_params", loop_hp.ppo_params)
-  ppo_hp = registry.hparams(ppo_params_name)
-
-  # Merge them and then override with the scoped overrides
-  combined_hp = merge_unscoped_hparams(
-      zip(HP_SCOPES, [loop_hp, model_hp, ppo_hp]))
-  combined_hp.override_from_dict(trial_hp_overrides)
-
-  # Split out the component hparams
-  loop_hp, model_hp, ppo_hp = (
-      split_scoped_hparams(HP_SCOPES, combined_hp))
-
-  # Dynamic register the model hp and set the new name in loop_hp
-  model_hp_name = "model_hp_%s" % str(trial_id)
-  dynamic_register_hparams(model_hp_name, model_hp)
-  loop_hp.generative_model_params = model_hp_name
-
-  # Dynamic register the PPO hp and set the new name in loop_hp
-  ppo_hp_name = "ppo_hp_%s" % str(trial_id)
-  dynamic_register_hparams(ppo_hp_name, ppo_hp)
-  loop_hp.ppo_params = ppo_hp_name
-
-  return loop_hp
-
-
-def dynamic_register_hparams(name, hparams):
-
-  @registry.register_hparams(name)
-  def new_hparams_set():
-    return tf.contrib.training.HParams(**hparams.values())
-
-  return new_hparams_set
-
-
-def create_loop_hparams():
-  hparams = registry.hparams(FLAGS.loop_hparams_set)
-  hparams.parse(FLAGS.loop_hparams)
-  return hparams
-
-
 def main(_):
-  hp = create_loop_hparams()
+  hp = trainer_model_based_params.create_loop_hparams()
   if FLAGS.job_dir_to_evaluate:
     compute_final_evaluation_on_real_environments(hp, FLAGS.job_dir_to_evaluate)
   else:
diff --git a/tensor2tensor/rl/trainer_model_based_agent_only.py b/tensor2tensor/rl/trainer_model_based_agent_only.py
index bf4ac0b12..cced3ebc5 100644
--- a/tensor2tensor/rl/trainer_model_based_agent_only.py
+++ b/tensor2tensor/rl/trainer_model_based_agent_only.py
@@ -30,6 +30,7 @@
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
 from tensor2tensor.data_generators import gym_problems_specs
 from tensor2tensor.rl import trainer_model_based
+from tensor2tensor.rl import trainer_model_based_params
 
 
 import tensorflow as tf
@@ -50,7 +51,7 @@ def get_simulated_problem_name(game):
 
 
 def main(_):
-  hparams = trainer_model_based.create_loop_hparams()
+  hparams = trainer_model_based_params.create_loop_hparams()
   problem_name = get_simulated_problem_name(hparams.game)
   world_model_dir = FLAGS.world_model_dir
   agent_model_dir = FLAGS.output_dir
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
new file mode 100644
index 000000000..953a2e5e5
--- /dev/null
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -0,0 +1,753 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+r"""Parameter sets for training of model-based RL agents."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import six
+
+
+from tensor2tensor.data_generators import gym_problems_specs
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+FLAGS = tf.flags.FLAGS
+
+
+HP_SCOPES = ["loop", "model", "ppo"]
+
+
+@registry.register_hparams
+def rlmb_base():
+  return tf.contrib.training.HParams(
+      epochs=15,
+      # Total frames used for training. This will be distributed evenly across
+      # hparams.epochs.
+      # This number should be divisible by real_ppo_epoch_length*epochs
+      # for our frame accounting to be preceise.
+      num_real_env_frames=96000,
+      generative_model="next_frame_basic_deterministic",
+      generative_model_params="next_frame_pixel_noise",
+      ppo_params="ppo_pong_base",
+      autoencoder_train_steps=0,
+      autoencoder_train_steps_initial_multiplier=10,
+      autoencoder_hparams_set="autoencoder_discrete_pong",
+      model_train_steps=15000,
+      inital_epoch_train_steps_multiplier=3,
+      simulated_env_generator_num_steps=2000,
+      simulation_random_starts=True,  # Use random starts in PPO.
+      # Flip the first random frame in PPO batch for the true beginning.
+      simulation_flip_first_random_for_beginning=True,
+      intrinsic_reward_scale=0.,
+      ppo_epochs_num=1000,  # This should be enough to see something
+      # Our simulated envs do not know how to reset.
+      # You should set ppo_time_limit to the value you believe that
+      # the simulated env produces a reasonable output.
+      ppo_time_limit=200,  # TODO(blazej): this param is unused
+      # It makes sense to have ppo_time_limit=ppo_epoch_length,
+      # though it is not necessary.
+      ppo_epoch_length=50,
+      ppo_num_agents=16,
+      # Do not eval since simulated batch env does not produce dones
+      ppo_eval_every_epochs=0,
+      ppo_learning_rate=1e-4,  # Will be changed, just so it exists.
+      # Whether the PPO agent should be restored from the previous iteration, or
+      # should start fresh each time.
+      ppo_continue_training=True,
+      # Resizing.
+      resize_height_factor=2,
+      resize_width_factor=2,
+      grayscale=False,
+      # Bump learning rate after first epoch by 3x.
+      # We picked 3x because our default learning rate schedule decreases with
+      # 1/square root of step; 1/sqrt(10k) = 0.01 and 1/sqrt(100k) ~ 0.0032
+      # so by bumping it up 3x we about "go back" from 100k steps to 10k, which
+      # is approximately as much as "going back 1 epoch" would be.
+      # In your experiments, you want to optimize this rate to your schedule.
+      learning_rate_bump=3.0,
+
+      gather_ppo_real_env_data=True,
+      real_ppo_epochs_num=0,
+      # This needs to be divisible by real_ppo_effective_num_agents.
+      real_ppo_epoch_length=16*200,
+      real_ppo_num_agents=1,
+      real_ppo_learning_rate=1e-4,
+      real_ppo_continue_training=True,
+      real_ppo_effective_num_agents=16,
+      real_ppo_eval_every_epochs=0,
+
+      game="pong",
+      # Whether to evaluate the world model in each iteration of the loop to get
+      # the model_reward_accuracy metric.
+      eval_world_model=True,
+      # Rollout fractions to report reward_accuracy on.
+      eval_rollout_fractions=[0.25, 0.5, 1],
+      stop_loop_early=False,  # To speed-up tests.
+  )
+
+
+@registry.register_hparams
+def rlmb_basetest():
+  """Base setting but quicker with only 2 epochs."""
+  hparams = rlmb_base()
+  hparams.game = "pong"
+  hparams.epochs = 2
+  hparams.num_real_env_frames = 3200
+  hparams.model_train_steps = 100
+  hparams.simulated_env_generator_num_steps = 20
+  hparams.ppo_epochs_num = 2
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_noresize():
+  hparams = rlmb_base()
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_quick():
+  """Base setting but quicker with only 2 epochs."""
+  hparams = rlmb_base()
+  hparams.epochs = 2
+  hparams.model_train_steps = 25000
+  hparams.ppo_epochs_num = 700
+  hparams.ppo_epoch_length = 50
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_quick_noresize():
+  hparams = rlmb_base()
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_quick_sd():
+  """Quick setting with stochastic discrete model."""
+  hparams = rlmb_quick()
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_sdtest():
+  """Test setting with stochastic discrete model."""
+  hparams = rlmb_basetest()
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_quick_sm():
+  """Quick setting with sampling."""
+  hparams = rlmb_quick()
+  hparams.generative_model_params = "next_frame_sampling"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic():
+  """Base setting with a stochastic next-frame model."""
+  hparams = rlmb_base()
+  hparams.inital_epoch_train_steps_multiplier = 5
+  hparams.generative_model = "next_frame_basic_stochastic"
+  hparams.generative_model_params = "next_frame_basic_stochastic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sampling_stochastic():
+  """Base setting with a stochastic next-frame model."""
+  hparams = rlmb_base()
+  hparams.generative_model = "next_frame_basic_stochastic"
+  hparams.generative_model_params = "next_frame_sampling_stochastic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic_discrete():
+  """Base setting with stochastic discrete model."""
+  hparams = rlmb_base()
+  hparams.learning_rate_bump = 1.0
+  hparams.grayscale = True
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic_recurrent():
+  """Base setting with recurrent model."""
+  hparams = rlmb_base()
+  hparams.generative_model = "next_frame_basic_recurrent"
+  hparams.generative_model_params = "next_frame_basic_recurrent"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic_discrete_noresize():
+  """Base setting with stochastic discrete model."""
+  hparams = rlmb_base()
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p():
+  """Base setting with sv2p as world model."""
+  hparams = rlmb_base()
+  hparams.learning_rate_bump = 1.0
+  hparams.generative_model = "next_frame_sv2p"
+  hparams.generative_model_params = "next_frame_sv2p_atari"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_softmax():
+  """Base setting with sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p()
+  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic():
+  """Base setting with deterministic sv2p as world model."""
+  hparams = rlmb_base_sv2p()
+  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic_softmax():
+  """Base setting with deterministic sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p_softmax()
+  hparams.generative_model_params = (
+      "next_frame_sv2p_atari_softmax_deterministic")
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_flippy30():
+  """Base setting with sv2p as world model."""
+  hparams = rlmb_base()
+  hparams.epochs = 30
+  hparams.ppo_epochs_num = 1000
+  hparams.model_train_steps = 15000
+  hparams.learning_rate_bump = 1.0
+  hparams.inital_epoch_train_steps_multiplier = 5
+  hparams.generative_model = "next_frame_sv2p"
+  hparams.generative_model_params = "next_frame_sv2p_atari"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_softmax_flippy30():
+  """Base setting with sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p_flippy30()
+  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic_flippy30():
+  """Base setting with deterministic sv2p as world model."""
+  hparams = rlmb_base_sv2p_flippy30()
+  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic_softmax_flippy30():
+  """Base setting with deterministic sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p_softmax_flippy30()
+  hparams.generative_model_params = (
+      "next_frame_sv2p_atari_softmax_deterministic")
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sampling():
+  """Base setting with a stochastic next-frame model."""
+  hparams = rlmb_base()
+  hparams.generative_model_params = "next_frame_sampling"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sampling_noresize():
+  hparams = rlmb_base_sampling()
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_flippy60():
+  """Schedule with a lot of epochs (slow)."""
+  hparams = rlmb_base_sampling()
+  hparams.epochs = 60
+  hparams.ppo_epochs_num = 500
+  hparams.model_train_steps = 10000
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_flippy30():
+  """Schedule with a lot of epochs (slow)."""
+  hparams = rlmb_base_sampling()
+  hparams.epochs = 30
+  hparams.ppo_epochs_num = 1000
+  hparams.model_train_steps = 15000
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_medium():
+  """Small set for larger testing."""
+  hparams = rlmb_base()
+  hparams.num_real_env_frames //= 2
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_25k():
+  """Small set for larger testing."""
+  hparams = rlmb_medium()
+  hparams.num_real_env_frames //= 2
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_short():
+  """Small set for larger testing."""
+  hparams = rlmb_base()
+  hparams.num_real_env_frames //= 5
+  hparams.model_train_steps //= 10
+  hparams.ppo_epochs_num //= 10
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_model_only():
+  hp = rlmb_base()
+  hp.epochs = 1
+  hp.ppo_epochs_num = 0
+  return hp
+
+
+@registry.register_hparams
+def rlmb_tiny():
+  """Tiny set for testing."""
+  return rlmb_base_sampling().override_from_dict(
+      tf.contrib.training.HParams(
+          epochs=1,
+          num_real_env_frames=128,
+          simulated_env_generator_num_steps=64,
+          model_train_steps=2,
+          ppo_epochs_num=2,
+          ppo_time_limit=5,
+          ppo_epoch_length=5,
+          ppo_num_agents=2,
+          real_ppo_epoch_length=36,
+          real_ppo_num_agents=1,
+          real_ppo_epochs_num=0,
+          real_ppo_effective_num_agents=2,
+          generative_model_params="next_frame_tiny",
+          stop_loop_early=True,
+          resize_height_factor=2,
+          resize_width_factor=2,
+          game="pong",
+      ).values())
+
+
+@registry.register_hparams
+def rlmb_tiny_stochastic():
+  """Tiny setting with a stochastic next-frame model."""
+  hparams = rlmb_tiny()
+  hparams.epochs = 1  # Too slow with 2 for regular runs.
+  hparams.generative_model = "next_frame_basic_stochastic"
+  hparams.generative_model_params = "next_frame_basic_stochastic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_tiny_recurrent():
+  """Tiny setting with a recurrent next-frame model."""
+  hparams = rlmb_tiny()
+  hparams.epochs = 1  # Too slow with 2 for regular runs.
+  hparams.generative_model = "next_frame_basic_recurrent"
+  hparams.generative_model_params = "next_frame_basic_recurrent"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_tiny_sv2p():
+  """Tiny setting with a tiny sv2p model."""
+  hparams = rlmb_tiny()
+  hparams.generative_model = "next_frame_sv2p"
+  hparams.generative_model_params = "next_frame_sv2p_tiny"
+  hparams.grayscale = False
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_ae_base():
+  """Parameter set for autoencoders."""
+  hparams = rlmb_base()
+  hparams.ppo_params = "ppo_pong_ae_base"
+  hparams.generative_model_params = "next_frame_ae"
+  hparams.autoencoder_hparams_set = "autoencoder_discrete_pong"
+  hparams.gather_ppo_real_env_data = False
+  hparams.autoencoder_train_steps = 5000
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  hparams.grayscale = False
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_ae_basetest():
+  """Base AE setting but quicker with only 2 epochs."""
+  hparams = rlmb_ae_base()
+  hparams.game = "pong"
+  hparams.epochs = 2
+  hparams.num_real_env_frames = 3200
+  hparams.model_train_steps = 100
+  hparams.autoencoder_train_steps = 10
+  hparams.simulated_env_generator_num_steps = 20
+  hparams.ppo_epochs_num = 2
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_ae_tiny():
+  """Tiny set for testing autoencoders."""
+  hparams = rlmb_tiny()
+  hparams.ppo_params = "ppo_pong_ae_base"
+  hparams.generative_model_params = "next_frame_ae_tiny"
+  hparams.autoencoder_hparams_set = "autoencoder_discrete_tiny"
+  hparams.gather_ppo_real_env_data = False
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  hparams.grayscale = False
+  hparams.autoencoder_train_steps = 1
+  hparams.autoencoder_train_steps_initial_multiplier = 0
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_tiny_simulation_deterministic_starts():
+  hp = rlmb_tiny()
+  hp.simulation_random_starts = False
+  return hp
+
+
+# RangedHParams for tuning
+# ==============================================================================
+# Note that the items here must be scoped with one of
+# HP_SCOPES={loop, model, ppo}, which set hyperparameters for the top-level
+# hparams, hp.generative_model_params, and hp.ppo_params, respectively.
+@registry.register_ranged_hparams
+def rlmb_grid(rhp):
+  """Grid over games and frames, and 5 runs each for variance."""
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+  base = 100000
+  medium = base // 2
+  small = medium // 2
+  rhp.set_discrete("loop.num_real_env_frames", [base, medium, small])
+
+  # Dummy parameter to get 5 runs for each configuration
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+
+
+@registry.register_ranged_hparams
+def rlmb_variance(rhp):
+  # Dummy parameter to get 5 runs for each configuration
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+
+
+@registry.register_ranged_hparams
+def rlmb_variance_nogame(rhp):
+  # Dummy parameter to get 20 runs for current configuration.
+  rhp.set_discrete("model.moe_loss_coef", list(range(20)))
+
+
+@registry.register_ranged_hparams
+def rlmb_three(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
+
+
+@registry.register_ranged_hparams
+def rlmb_test1(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
+  rhp.set_discrete("loop.ppo_learning_rate", [5e-5, 1e-4, 2e-4])
+  rhp.set_discrete("ppo.optimization_batch_size", [20, 40])
+  rhp.set_discrete("loop.epochs", [3, 6])
+
+
+@registry.register_ranged_hparams
+def rlmb_scheduled_sampling(rhp):
+  rhp.set_float("model.scheduled_sampling_prob", 0.0, 1.0)
+
+
+@registry.register_ranged_hparams
+def rlmb_all_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_whitelisted_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_human_score_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game",
+                      gym_problems_specs.ATARI_GAMES_WITH_HUMAN_SCORE)
+
+
+@registry.register_ranged_hparams
+def rlmb_curious_games10(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_curious_games5(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_debug_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_DEBUG_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_ae_variance(rhp):
+  # Dummy parameter to get 5 runs for each configuration
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+  base = 100000
+  small = base // 4
+  rhp.set_discrete("loop.num_real_env_frames", [base, small])
+
+
+@registry.register_ranged_hparams
+def rlmb_ppolr_game(rhp):
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+  base_lr = 1e-4
+  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
+
+
+@registry.register_ranged_hparams
+def rlmb_ppolr(rhp):
+  base_lr = 1e-4
+  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
+
+
+@registry.register_ranged_hparams
+def rlmb_ae_ppo_lr(rhp):
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+  base_lr = 1e-4
+  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
+
+
+@registry.register_ranged_hparams
+def rlmb_dropout_range(rhp):
+  rhp.set_float("model.dropout", 0.2, 0.4)
+
+
+@registry.register_ranged_hparams
+def rlmb_intrinsic_reward_scale(rhp):
+  rhp.set_float("loop.intrinsic_reward_scale", 0.01, 10.)
+
+
+@registry.register_ranged_hparams
+def rlmb_l1l2cutoff_range(rhp):
+  """Loss and loss-cutoff tuning grid."""
+  rhp.set_float("model.video_modality_loss_cutoff", 1.4, 3.4)
+
+
+@registry.register_ranged_hparams
+def rlmb_xentcutoff_range(rhp):
+  """Cross entropy cutoff tuning grid."""
+  rhp.set_float("model.video_modality_loss_cutoff", 0.01, 0.05)
+
+
+@registry.register_ranged_hparams
+def rlmb_pixel_noise(rhp):
+  """Input pixel noise tuning grid."""
+  rhp.set_categorical("loop.generative_model_params",
+                      ["next_frame_pixel_noise"])
+  rhp.set_discrete("model.video_modality_input_noise",
+                   [0.0025 * i for i in range(200)])
+
+
+@registry.register_ranged_hparams
+def rlmb_dummy_range(rhp):
+  """Dummy tuning grid just to get the variance."""
+  rhp.set_float("model.moe_loss_coef", 0.01, 0.02)
+
+
+@registry.register_ranged_hparams
+def rlmb_epochs_num(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.epochs", [3, 6, 12])
+
+
+@registry.register_ranged_hparams
+def rlmb_ppo_epochs_num(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.ppo_epochs_num", [200, 1000, 2000, 4000])
+
+
+@registry.register_ranged_hparams
+def rlmb_ppo_epoch_len(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.ppo_epoch_length", [25, 50, 100])
+
+
+@registry.register_ranged_hparams
+def rlmb_num_frames(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.num_real_env_frames",
+                   [1000*el for el in [30, 100, 500, 1000]])
+
+
+@registry.register_ranged_hparams
+def rlmb_ppo_optimization_batch_size(rhp):
+  rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_discrete("ppo.optimization_batch_size", [4, 10, 20])
+
+
+@registry.register_ranged_hparams
+def rlmb_logits_clip(rhp):
+  rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_discrete("ppo.logits_clip", [0., 5.])
+
+
+def merge_unscoped_hparams(scopes_and_hparams):
+  """Merge multiple HParams into one with scopes."""
+  merged_values = {}
+  for (scope, hparams) in scopes_and_hparams:
+    for key, value in six.iteritems(hparams.values()):
+      scoped_key = "%s.%s" % (scope, key)
+      merged_values[scoped_key] = value
+
+  return tf.contrib.training.HParams(**merged_values)
+
+
+def split_scoped_hparams(scopes, merged_hparams):
+  """Split single HParams with scoped keys into multiple."""
+  split_values = {scope: {} for scope in scopes}
+  merged_values = merged_hparams.values()
+  for scoped_key, value in six.iteritems(merged_values):
+    scope = scoped_key.split(".")[0]
+    key = scoped_key[len(scope) + 1:]
+    split_values[scope][key] = value
+
+  return [
+      tf.contrib.training.HParams(**split_values[scope]) for scope in scopes
+  ]
+
+
+def training_loop_hparams_from_scoped_overrides(scoped_overrides, trial_id):
+  """Create HParams suitable for training loop from scoped HParams.
+
+  Args:
+    scoped_overrides: HParams, with keys all scoped by one of HP_SCOPES. These
+      parameters are overrides for the base HParams created by
+      create_loop_hparams.
+    trial_id: str, trial identifier. This is used to register unique HParams
+      names for the underlying model and ppo HParams.
+
+  Returns:
+    HParams suitable for passing to training_loop.
+  """
+  trial_hp_overrides = scoped_overrides.values()
+
+  # Create loop, model, and ppo base HParams
+  loop_hp = create_loop_hparams()
+  model_hp_name = trial_hp_overrides.get(
+      "loop.generative_model_params", loop_hp.generative_model_params)
+  model_hp = registry.hparams(model_hp_name).parse(FLAGS.hparams)
+  ppo_params_name = trial_hp_overrides.get(
+      "loop.ppo_params", loop_hp.ppo_params)
+  ppo_hp = registry.hparams(ppo_params_name)
+
+  # Merge them and then override with the scoped overrides
+  combined_hp = merge_unscoped_hparams(
+      zip(HP_SCOPES, [loop_hp, model_hp, ppo_hp]))
+  combined_hp.override_from_dict(trial_hp_overrides)
+
+  # Split out the component hparams
+  loop_hp, model_hp, ppo_hp = (
+      split_scoped_hparams(HP_SCOPES, combined_hp))
+
+  # Dynamic register the model hp and set the new name in loop_hp
+  model_hp_name = "model_hp_%s" % str(trial_id)
+  dynamic_register_hparams(model_hp_name, model_hp)
+  loop_hp.generative_model_params = model_hp_name
+
+  # Dynamic register the PPO hp and set the new name in loop_hp
+  ppo_hp_name = "ppo_hp_%s" % str(trial_id)
+  dynamic_register_hparams(ppo_hp_name, ppo_hp)
+  loop_hp.ppo_params = ppo_hp_name
+
+  return loop_hp
+
+
+def dynamic_register_hparams(name, hparams):
+
+  @registry.register_hparams(name)
+  def new_hparams_set():
+    return tf.contrib.training.HParams(**hparams.values())
+
+  return new_hparams_set
+
+
+def create_loop_hparams():
+  hparams = registry.hparams(FLAGS.loop_hparams_set)
+  hparams.parse(FLAGS.loop_hparams)
+  return hparams

From 7c14ecd1210c0a56fcde63c227b392688f821883 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 11 Oct 2018 15:16:54 -0700
Subject: [PATCH 0991/2720] internal

PiperOrigin-RevId: 216769792
---
 tensor2tensor/rl/trainer_model_based.py       | 725 ++++++++++++++++-
 .../rl/trainer_model_based_agent_only.py      |   3 +-
 .../rl/trainer_model_based_params.py          | 753 ------------------
 3 files changed, 724 insertions(+), 757 deletions(-)
 delete mode 100644 tensor2tensor/rl/trainer_model_based_params.py

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index bde551c6a..03fe2a525 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -44,7 +44,6 @@
 from tensor2tensor.data_generators import gym_problems_specs
 from tensor2tensor.layers import discretization
 from tensor2tensor.rl import rl_trainer_lib
-from tensor2tensor.rl import trainer_model_based_params
 from tensor2tensor.rl.envs.tf_atari_wrappers import PyFuncWrapper
 from tensor2tensor.rl.envs.utils import InitialFrameChooser
 from tensor2tensor.utils import registry
@@ -64,6 +63,8 @@
 flags.DEFINE_string("eval_results_dir", "/tmp",
                     "Directory to store result of evaluation")
 
+HP_SCOPES = ["loop", "model", "ppo"]
+
 
 def setup_directories(base_dir, subdirs):
   base_dir = os.path.expanduser(base_dir)
@@ -747,8 +748,728 @@ def combine_training_data(problem, final_data_dir, old_data_dirs,
         tf.gfile.Copy(fname, new_fname)
 
 
+@registry.register_hparams
+def rlmb_base():
+  return tf.contrib.training.HParams(
+      epochs=15,
+      # Total frames used for training. This will be distributed evenly across
+      # hparams.epochs.
+      # This number should be divisible by real_ppo_epoch_length*epochs
+      # for our frame accounting to be preceise.
+      num_real_env_frames=96000,
+      generative_model="next_frame_basic_deterministic",
+      generative_model_params="next_frame_pixel_noise",
+      ppo_params="ppo_pong_base",
+      autoencoder_train_steps=0,
+      autoencoder_train_steps_initial_multiplier=10,
+      autoencoder_hparams_set="autoencoder_discrete_pong",
+      model_train_steps=15000,
+      inital_epoch_train_steps_multiplier=3,
+      simulated_env_generator_num_steps=2000,
+      simulation_random_starts=True,  # Use random starts in PPO.
+      # Flip the first random frame in PPO batch for the true beginning.
+      simulation_flip_first_random_for_beginning=True,
+      intrinsic_reward_scale=0.,
+      ppo_epochs_num=1000,  # This should be enough to see something
+      # Our simulated envs do not know how to reset.
+      # You should set ppo_time_limit to the value you believe that
+      # the simulated env produces a reasonable output.
+      ppo_time_limit=200,  # TODO(blazej): this param is unused
+      # It makes sense to have ppo_time_limit=ppo_epoch_length,
+      # though it is not necessary.
+      ppo_epoch_length=50,
+      ppo_num_agents=16,
+      # Do not eval since simulated batch env does not produce dones
+      ppo_eval_every_epochs=0,
+      ppo_learning_rate=1e-4,  # Will be changed, just so it exists.
+      # Whether the PPO agent should be restored from the previous iteration, or
+      # should start fresh each time.
+      ppo_continue_training=True,
+      # Resizing.
+      resize_height_factor=2,
+      resize_width_factor=2,
+      grayscale=False,
+      # Bump learning rate after first epoch by 3x.
+      # We picked 3x because our default learning rate schedule decreases with
+      # 1/square root of step; 1/sqrt(10k) = 0.01 and 1/sqrt(100k) ~ 0.0032
+      # so by bumping it up 3x we about "go back" from 100k steps to 10k, which
+      # is approximately as much as "going back 1 epoch" would be.
+      # In your experiments, you want to optimize this rate to your schedule.
+      learning_rate_bump=3.0,
+
+      gather_ppo_real_env_data=True,
+      real_ppo_epochs_num=0,
+      # This needs to be divisible by real_ppo_effective_num_agents.
+      real_ppo_epoch_length=16*200,
+      real_ppo_num_agents=1,
+      real_ppo_learning_rate=1e-4,
+      real_ppo_continue_training=True,
+      real_ppo_effective_num_agents=16,
+      real_ppo_eval_every_epochs=0,
+
+      game="pong",
+      # Whether to evaluate the world model in each iteration of the loop to get
+      # the model_reward_accuracy metric.
+      eval_world_model=True,
+      # Rollout fractions to report reward_accuracy on.
+      eval_rollout_fractions=[0.25, 0.5, 1],
+      stop_loop_early=False,  # To speed-up tests.
+  )
+
+
+@registry.register_hparams
+def rlmb_basetest():
+  """Base setting but quicker with only 2 epochs."""
+  hparams = rlmb_base()
+  hparams.game = "pong"
+  hparams.epochs = 2
+  hparams.num_real_env_frames = 3200
+  hparams.model_train_steps = 100
+  hparams.simulated_env_generator_num_steps = 20
+  hparams.ppo_epochs_num = 2
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_noresize():
+  hparams = rlmb_base()
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_quick():
+  """Base setting but quicker with only 2 epochs."""
+  hparams = rlmb_base()
+  hparams.epochs = 2
+  hparams.model_train_steps = 25000
+  hparams.ppo_epochs_num = 700
+  hparams.ppo_epoch_length = 50
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_quick_noresize():
+  hparams = rlmb_base()
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_quick_sd():
+  """Quick setting with stochastic discrete model."""
+  hparams = rlmb_quick()
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_sdtest():
+  """Test setting with stochastic discrete model."""
+  hparams = rlmb_basetest()
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_quick_sm():
+  """Quick setting with sampling."""
+  hparams = rlmb_quick()
+  hparams.generative_model_params = "next_frame_sampling"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic():
+  """Base setting with a stochastic next-frame model."""
+  hparams = rlmb_base()
+  hparams.inital_epoch_train_steps_multiplier = 5
+  hparams.generative_model = "next_frame_basic_stochastic"
+  hparams.generative_model_params = "next_frame_basic_stochastic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sampling_stochastic():
+  """Base setting with a stochastic next-frame model."""
+  hparams = rlmb_base()
+  hparams.generative_model = "next_frame_basic_stochastic"
+  hparams.generative_model_params = "next_frame_sampling_stochastic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic_discrete():
+  """Base setting with stochastic discrete model."""
+  hparams = rlmb_base()
+  hparams.learning_rate_bump = 1.0
+  hparams.grayscale = True
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic_recurrent():
+  """Base setting with recurrent model."""
+  hparams = rlmb_base()
+  hparams.generative_model = "next_frame_basic_recurrent"
+  hparams.generative_model_params = "next_frame_basic_recurrent"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic_discrete_noresize():
+  """Base setting with stochastic discrete model."""
+  hparams = rlmb_base()
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p():
+  """Base setting with sv2p as world model."""
+  hparams = rlmb_base()
+  hparams.learning_rate_bump = 1.0
+  hparams.generative_model = "next_frame_sv2p"
+  hparams.generative_model_params = "next_frame_sv2p_atari"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_softmax():
+  """Base setting with sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p()
+  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic():
+  """Base setting with deterministic sv2p as world model."""
+  hparams = rlmb_base_sv2p()
+  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic_softmax():
+  """Base setting with deterministic sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p_softmax()
+  hparams.generative_model_params = (
+      "next_frame_sv2p_atari_softmax_deterministic")
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_flippy30():
+  """Base setting with sv2p as world model."""
+  hparams = rlmb_base()
+  hparams.epochs = 30
+  hparams.ppo_epochs_num = 1000
+  hparams.model_train_steps = 15000
+  hparams.learning_rate_bump = 1.0
+  hparams.inital_epoch_train_steps_multiplier = 5
+  hparams.generative_model = "next_frame_sv2p"
+  hparams.generative_model_params = "next_frame_sv2p_atari"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_softmax_flippy30():
+  """Base setting with sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p_flippy30()
+  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic_flippy30():
+  """Base setting with deterministic sv2p as world model."""
+  hparams = rlmb_base_sv2p_flippy30()
+  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic_softmax_flippy30():
+  """Base setting with deterministic sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p_softmax_flippy30()
+  hparams.generative_model_params = (
+      "next_frame_sv2p_atari_softmax_deterministic")
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sampling():
+  """Base setting with a stochastic next-frame model."""
+  hparams = rlmb_base()
+  hparams.generative_model_params = "next_frame_sampling"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sampling_noresize():
+  hparams = rlmb_base_sampling()
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_flippy60():
+  """Schedule with a lot of epochs (slow)."""
+  hparams = rlmb_base_sampling()
+  hparams.epochs = 60
+  hparams.ppo_epochs_num = 500
+  hparams.model_train_steps = 10000
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_flippy30():
+  """Schedule with a lot of epochs (slow)."""
+  hparams = rlmb_base_sampling()
+  hparams.epochs = 30
+  hparams.ppo_epochs_num = 1000
+  hparams.model_train_steps = 15000
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_medium():
+  """Small set for larger testing."""
+  hparams = rlmb_base()
+  hparams.num_real_env_frames //= 2
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_25k():
+  """Small set for larger testing."""
+  hparams = rlmb_medium()
+  hparams.num_real_env_frames //= 2
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_short():
+  """Small set for larger testing."""
+  hparams = rlmb_base()
+  hparams.num_real_env_frames //= 5
+  hparams.model_train_steps //= 10
+  hparams.ppo_epochs_num //= 10
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_model_only():
+  hp = rlmb_base()
+  hp.epochs = 1
+  hp.ppo_epochs_num = 0
+  return hp
+
+
+@registry.register_hparams
+def rlmb_tiny():
+  """Tiny set for testing."""
+  return rlmb_base_sampling().override_from_dict(
+      tf.contrib.training.HParams(
+          epochs=1,
+          num_real_env_frames=128,
+          simulated_env_generator_num_steps=64,
+          model_train_steps=2,
+          ppo_epochs_num=2,
+          ppo_time_limit=5,
+          ppo_epoch_length=5,
+          ppo_num_agents=2,
+          real_ppo_epoch_length=36,
+          real_ppo_num_agents=1,
+          real_ppo_epochs_num=0,
+          real_ppo_effective_num_agents=2,
+          generative_model_params="next_frame_tiny",
+          stop_loop_early=True,
+          resize_height_factor=2,
+          resize_width_factor=2,
+          game="pong",
+      ).values())
+
+
+@registry.register_hparams
+def rlmb_tiny_stochastic():
+  """Tiny setting with a stochastic next-frame model."""
+  hparams = rlmb_tiny()
+  hparams.epochs = 1  # Too slow with 2 for regular runs.
+  hparams.generative_model = "next_frame_basic_stochastic"
+  hparams.generative_model_params = "next_frame_basic_stochastic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_tiny_recurrent():
+  """Tiny setting with a recurrent next-frame model."""
+  hparams = rlmb_tiny()
+  hparams.epochs = 1  # Too slow with 2 for regular runs.
+  hparams.generative_model = "next_frame_basic_recurrent"
+  hparams.generative_model_params = "next_frame_basic_recurrent"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_tiny_sv2p():
+  """Tiny setting with a tiny sv2p model."""
+  hparams = rlmb_tiny()
+  hparams.generative_model = "next_frame_sv2p"
+  hparams.generative_model_params = "next_frame_sv2p_tiny"
+  hparams.grayscale = False
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_ae_base():
+  """Parameter set for autoencoders."""
+  hparams = rlmb_base()
+  hparams.ppo_params = "ppo_pong_ae_base"
+  hparams.generative_model_params = "next_frame_ae"
+  hparams.autoencoder_hparams_set = "autoencoder_discrete_pong"
+  hparams.gather_ppo_real_env_data = False
+  hparams.autoencoder_train_steps = 5000
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  hparams.grayscale = False
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_ae_basetest():
+  """Base AE setting but quicker with only 2 epochs."""
+  hparams = rlmb_ae_base()
+  hparams.game = "pong"
+  hparams.epochs = 2
+  hparams.num_real_env_frames = 3200
+  hparams.model_train_steps = 100
+  hparams.autoencoder_train_steps = 10
+  hparams.simulated_env_generator_num_steps = 20
+  hparams.ppo_epochs_num = 2
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_ae_tiny():
+  """Tiny set for testing autoencoders."""
+  hparams = rlmb_tiny()
+  hparams.ppo_params = "ppo_pong_ae_base"
+  hparams.generative_model_params = "next_frame_ae_tiny"
+  hparams.autoencoder_hparams_set = "autoencoder_discrete_tiny"
+  hparams.gather_ppo_real_env_data = False
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  hparams.grayscale = False
+  hparams.autoencoder_train_steps = 1
+  hparams.autoencoder_train_steps_initial_multiplier = 0
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_tiny_simulation_deterministic_starts():
+  hp = rlmb_tiny()
+  hp.simulation_random_starts = False
+  return hp
+
+
+# RangedHParams for tuning
+# ==============================================================================
+# Note that the items here must be scoped with one of
+# HP_SCOPES={loop, model, ppo}, which set hyperparameters for the top-level
+# hparams, hp.generative_model_params, and hp.ppo_params, respectively.
+@registry.register_ranged_hparams
+def rlmb_grid(rhp):
+  """Grid over games and frames, and 5 runs each for variance."""
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+  base = 100000
+  medium = base // 2
+  small = medium // 2
+  rhp.set_discrete("loop.num_real_env_frames", [base, medium, small])
+
+  # Dummy parameter to get 5 runs for each configuration
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+
+
+@registry.register_ranged_hparams
+def rlmb_variance(rhp):
+  # Dummy parameter to get 5 runs for each configuration
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+
+
+@registry.register_ranged_hparams
+def rlmb_variance_nogame(rhp):
+  # Dummy parameter to get 20 runs for current configuration.
+  rhp.set_discrete("model.moe_loss_coef", list(range(20)))
+
+
+@registry.register_ranged_hparams
+def rlmb_three(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
+
+
+@registry.register_ranged_hparams
+def rlmb_test1(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
+  rhp.set_discrete("loop.ppo_learning_rate", [5e-5, 1e-4, 2e-4])
+  rhp.set_discrete("ppo.optimization_batch_size", [20, 40])
+  rhp.set_discrete("loop.epochs", [3, 6])
+
+
+@registry.register_ranged_hparams
+def rlmb_scheduled_sampling(rhp):
+  rhp.set_float("model.scheduled_sampling_prob", 0.0, 1.0)
+
+
+@registry.register_ranged_hparams
+def rlmb_all_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_whitelisted_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_human_score_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game",
+                      gym_problems_specs.ATARI_GAMES_WITH_HUMAN_SCORE)
+
+
+@registry.register_ranged_hparams
+def rlmb_curious_games10(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_curious_games5(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_debug_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_DEBUG_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_ae_variance(rhp):
+  # Dummy parameter to get 5 runs for each configuration
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+  base = 100000
+  small = base // 4
+  rhp.set_discrete("loop.num_real_env_frames", [base, small])
+
+
+@registry.register_ranged_hparams
+def rlmb_ppolr_game(rhp):
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+  base_lr = 1e-4
+  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
+
+
+@registry.register_ranged_hparams
+def rlmb_ppolr(rhp):
+  base_lr = 1e-4
+  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
+
+
+@registry.register_ranged_hparams
+def rlmb_ae_ppo_lr(rhp):
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+  base_lr = 1e-4
+  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
+
+
+@registry.register_ranged_hparams
+def rlmb_dropout_range(rhp):
+  rhp.set_float("model.dropout", 0.2, 0.4)
+
+
+@registry.register_ranged_hparams
+def rlmb_intrinsic_reward_scale(rhp):
+  rhp.set_float("loop.intrinsic_reward_scale", 0.01, 10.)
+
+
+@registry.register_ranged_hparams
+def rlmb_l1l2cutoff_range(rhp):
+  """Loss and loss-cutoff tuning grid."""
+  rhp.set_float("model.video_modality_loss_cutoff", 1.4, 3.4)
+
+
+@registry.register_ranged_hparams
+def rlmb_xentcutoff_range(rhp):
+  """Cross entropy cutoff tuning grid."""
+  rhp.set_float("model.video_modality_loss_cutoff", 0.01, 0.05)
+
+
+@registry.register_ranged_hparams
+def rlmb_pixel_noise(rhp):
+  """Input pixel noise tuning grid."""
+  rhp.set_categorical("loop.generative_model_params",
+                      ["next_frame_pixel_noise"])
+  rhp.set_discrete("model.video_modality_input_noise",
+                   [0.0025 * i for i in range(200)])
+
+
+@registry.register_ranged_hparams
+def rlmb_dummy_range(rhp):
+  """Dummy tuning grid just to get the variance."""
+  rhp.set_float("model.moe_loss_coef", 0.01, 0.02)
+
+
+@registry.register_ranged_hparams
+def rlmb_epochs_num(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.epochs", [3, 6, 12])
+
+
+@registry.register_ranged_hparams
+def rlmb_ppo_epochs_num(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.ppo_epochs_num", [200, 1000, 2000, 4000])
+
+
+@registry.register_ranged_hparams
+def rlmb_ppo_epoch_len(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.ppo_epoch_length", [25, 50, 100])
+
+
+@registry.register_ranged_hparams
+def rlmb_num_frames(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.num_real_env_frames",
+                   [1000*el for el in [30, 100, 500, 1000]])
+
+
+@registry.register_ranged_hparams
+def rlmb_ppo_optimization_batch_size(rhp):
+  rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_discrete("ppo.optimization_batch_size", [4, 10, 20])
+
+
+@registry.register_ranged_hparams
+def rlmb_logits_clip(rhp):
+  rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_discrete("ppo.logits_clip", [0., 5.])
+
+
+def merge_unscoped_hparams(scopes_and_hparams):
+  """Merge multiple HParams into one with scopes."""
+  merged_values = {}
+  for (scope, hparams) in scopes_and_hparams:
+    for key, value in six.iteritems(hparams.values()):
+      scoped_key = "%s.%s" % (scope, key)
+      merged_values[scoped_key] = value
+
+  return tf.contrib.training.HParams(**merged_values)
+
+
+def split_scoped_hparams(scopes, merged_hparams):
+  """Split single HParams with scoped keys into multiple."""
+  split_values = dict([(scope, dict()) for scope in scopes])
+  merged_values = merged_hparams.values()
+  for scoped_key, value in six.iteritems(merged_values):
+    scope = scoped_key.split(".")[0]
+    key = scoped_key[len(scope) + 1:]
+    split_values[scope][key] = value
+
+  return [
+      tf.contrib.training.HParams(**split_values[scope]) for scope in scopes
+  ]
+
+
+def training_loop_hparams_from_scoped_overrides(scoped_overrides, trial_id):
+  """Create HParams suitable for training loop from scoped HParams.
+
+  Args:
+    scoped_overrides: HParams, with keys all scoped by one of HP_SCOPES. These
+      parameters are overrides for the base HParams created by
+      create_loop_hparams.
+    trial_id: str, trial identifier. This is used to register unique HParams
+      names for the underlying model and ppo HParams.
+
+  Returns:
+    HParams suitable for passing to training_loop.
+  """
+  trial_hp_overrides = scoped_overrides.values()
+
+  # Create loop, model, and ppo base HParams
+  loop_hp = create_loop_hparams()
+  model_hp_name = trial_hp_overrides.get(
+      "loop.generative_model_params", loop_hp.generative_model_params)
+  model_hp = registry.hparams(model_hp_name).parse(FLAGS.hparams)
+  ppo_params_name = trial_hp_overrides.get(
+      "loop.ppo_params", loop_hp.ppo_params)
+  ppo_hp = registry.hparams(ppo_params_name)
+
+  # Merge them and then override with the scoped overrides
+  combined_hp = merge_unscoped_hparams(
+      zip(HP_SCOPES, [loop_hp, model_hp, ppo_hp]))
+  combined_hp.override_from_dict(trial_hp_overrides)
+
+  # Split out the component hparams
+  loop_hp, model_hp, ppo_hp = (
+      split_scoped_hparams(HP_SCOPES, combined_hp))
+
+  # Dynamic register the model hp and set the new name in loop_hp
+  model_hp_name = "model_hp_%s" % str(trial_id)
+  dynamic_register_hparams(model_hp_name, model_hp)
+  loop_hp.generative_model_params = model_hp_name
+
+  # Dynamic register the PPO hp and set the new name in loop_hp
+  ppo_hp_name = "ppo_hp_%s" % str(trial_id)
+  dynamic_register_hparams(ppo_hp_name, ppo_hp)
+  loop_hp.ppo_params = ppo_hp_name
+
+  return loop_hp
+
+
+def dynamic_register_hparams(name, hparams):
+
+  @registry.register_hparams(name)
+  def new_hparams_set():
+    return tf.contrib.training.HParams(**hparams.values())
+
+  return new_hparams_set
+
+
+def create_loop_hparams():
+  hparams = registry.hparams(FLAGS.loop_hparams_set)
+  hparams.parse(FLAGS.loop_hparams)
+  return hparams
+
+
 def main(_):
-  hp = trainer_model_based_params.create_loop_hparams()
+  hp = create_loop_hparams()
   if FLAGS.job_dir_to_evaluate:
     compute_final_evaluation_on_real_environments(hp, FLAGS.job_dir_to_evaluate)
   else:
diff --git a/tensor2tensor/rl/trainer_model_based_agent_only.py b/tensor2tensor/rl/trainer_model_based_agent_only.py
index cced3ebc5..bf4ac0b12 100644
--- a/tensor2tensor/rl/trainer_model_based_agent_only.py
+++ b/tensor2tensor/rl/trainer_model_based_agent_only.py
@@ -30,7 +30,6 @@
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
 from tensor2tensor.data_generators import gym_problems_specs
 from tensor2tensor.rl import trainer_model_based
-from tensor2tensor.rl import trainer_model_based_params
 
 
 import tensorflow as tf
@@ -51,7 +50,7 @@ def get_simulated_problem_name(game):
 
 
 def main(_):
-  hparams = trainer_model_based_params.create_loop_hparams()
+  hparams = trainer_model_based.create_loop_hparams()
   problem_name = get_simulated_problem_name(hparams.game)
   world_model_dir = FLAGS.world_model_dir
   agent_model_dir = FLAGS.output_dir
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
deleted file mode 100644
index 953a2e5e5..000000000
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ /dev/null
@@ -1,753 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-r"""Parameter sets for training of model-based RL agents."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import six
-
-
-from tensor2tensor.data_generators import gym_problems_specs
-from tensor2tensor.utils import registry
-
-import tensorflow as tf
-
-
-FLAGS = tf.flags.FLAGS
-
-
-HP_SCOPES = ["loop", "model", "ppo"]
-
-
-@registry.register_hparams
-def rlmb_base():
-  return tf.contrib.training.HParams(
-      epochs=15,
-      # Total frames used for training. This will be distributed evenly across
-      # hparams.epochs.
-      # This number should be divisible by real_ppo_epoch_length*epochs
-      # for our frame accounting to be preceise.
-      num_real_env_frames=96000,
-      generative_model="next_frame_basic_deterministic",
-      generative_model_params="next_frame_pixel_noise",
-      ppo_params="ppo_pong_base",
-      autoencoder_train_steps=0,
-      autoencoder_train_steps_initial_multiplier=10,
-      autoencoder_hparams_set="autoencoder_discrete_pong",
-      model_train_steps=15000,
-      inital_epoch_train_steps_multiplier=3,
-      simulated_env_generator_num_steps=2000,
-      simulation_random_starts=True,  # Use random starts in PPO.
-      # Flip the first random frame in PPO batch for the true beginning.
-      simulation_flip_first_random_for_beginning=True,
-      intrinsic_reward_scale=0.,
-      ppo_epochs_num=1000,  # This should be enough to see something
-      # Our simulated envs do not know how to reset.
-      # You should set ppo_time_limit to the value you believe that
-      # the simulated env produces a reasonable output.
-      ppo_time_limit=200,  # TODO(blazej): this param is unused
-      # It makes sense to have ppo_time_limit=ppo_epoch_length,
-      # though it is not necessary.
-      ppo_epoch_length=50,
-      ppo_num_agents=16,
-      # Do not eval since simulated batch env does not produce dones
-      ppo_eval_every_epochs=0,
-      ppo_learning_rate=1e-4,  # Will be changed, just so it exists.
-      # Whether the PPO agent should be restored from the previous iteration, or
-      # should start fresh each time.
-      ppo_continue_training=True,
-      # Resizing.
-      resize_height_factor=2,
-      resize_width_factor=2,
-      grayscale=False,
-      # Bump learning rate after first epoch by 3x.
-      # We picked 3x because our default learning rate schedule decreases with
-      # 1/square root of step; 1/sqrt(10k) = 0.01 and 1/sqrt(100k) ~ 0.0032
-      # so by bumping it up 3x we about "go back" from 100k steps to 10k, which
-      # is approximately as much as "going back 1 epoch" would be.
-      # In your experiments, you want to optimize this rate to your schedule.
-      learning_rate_bump=3.0,
-
-      gather_ppo_real_env_data=True,
-      real_ppo_epochs_num=0,
-      # This needs to be divisible by real_ppo_effective_num_agents.
-      real_ppo_epoch_length=16*200,
-      real_ppo_num_agents=1,
-      real_ppo_learning_rate=1e-4,
-      real_ppo_continue_training=True,
-      real_ppo_effective_num_agents=16,
-      real_ppo_eval_every_epochs=0,
-
-      game="pong",
-      # Whether to evaluate the world model in each iteration of the loop to get
-      # the model_reward_accuracy metric.
-      eval_world_model=True,
-      # Rollout fractions to report reward_accuracy on.
-      eval_rollout_fractions=[0.25, 0.5, 1],
-      stop_loop_early=False,  # To speed-up tests.
-  )
-
-
-@registry.register_hparams
-def rlmb_basetest():
-  """Base setting but quicker with only 2 epochs."""
-  hparams = rlmb_base()
-  hparams.game = "pong"
-  hparams.epochs = 2
-  hparams.num_real_env_frames = 3200
-  hparams.model_train_steps = 100
-  hparams.simulated_env_generator_num_steps = 20
-  hparams.ppo_epochs_num = 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_noresize():
-  hparams = rlmb_base()
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_quick():
-  """Base setting but quicker with only 2 epochs."""
-  hparams = rlmb_base()
-  hparams.epochs = 2
-  hparams.model_train_steps = 25000
-  hparams.ppo_epochs_num = 700
-  hparams.ppo_epoch_length = 50
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_quick_noresize():
-  hparams = rlmb_base()
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_quick_sd():
-  """Quick setting with stochastic discrete model."""
-  hparams = rlmb_quick()
-  hparams.generative_model = "next_frame_basic_stochastic_discrete"
-  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_sdtest():
-  """Test setting with stochastic discrete model."""
-  hparams = rlmb_basetest()
-  hparams.generative_model = "next_frame_basic_stochastic_discrete"
-  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_quick_sm():
-  """Quick setting with sampling."""
-  hparams = rlmb_quick()
-  hparams.generative_model_params = "next_frame_sampling"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_stochastic():
-  """Base setting with a stochastic next-frame model."""
-  hparams = rlmb_base()
-  hparams.inital_epoch_train_steps_multiplier = 5
-  hparams.generative_model = "next_frame_basic_stochastic"
-  hparams.generative_model_params = "next_frame_basic_stochastic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sampling_stochastic():
-  """Base setting with a stochastic next-frame model."""
-  hparams = rlmb_base()
-  hparams.generative_model = "next_frame_basic_stochastic"
-  hparams.generative_model_params = "next_frame_sampling_stochastic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_stochastic_discrete():
-  """Base setting with stochastic discrete model."""
-  hparams = rlmb_base()
-  hparams.learning_rate_bump = 1.0
-  hparams.grayscale = True
-  hparams.generative_model = "next_frame_basic_stochastic_discrete"
-  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_stochastic_recurrent():
-  """Base setting with recurrent model."""
-  hparams = rlmb_base()
-  hparams.generative_model = "next_frame_basic_recurrent"
-  hparams.generative_model_params = "next_frame_basic_recurrent"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_stochastic_discrete_noresize():
-  """Base setting with stochastic discrete model."""
-  hparams = rlmb_base()
-  hparams.generative_model = "next_frame_basic_stochastic_discrete"
-  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p():
-  """Base setting with sv2p as world model."""
-  hparams = rlmb_base()
-  hparams.learning_rate_bump = 1.0
-  hparams.generative_model = "next_frame_sv2p"
-  hparams.generative_model_params = "next_frame_sv2p_atari"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_softmax():
-  """Base setting with sv2p as world model with softmax."""
-  hparams = rlmb_base_sv2p()
-  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_deterministic():
-  """Base setting with deterministic sv2p as world model."""
-  hparams = rlmb_base_sv2p()
-  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_deterministic_softmax():
-  """Base setting with deterministic sv2p as world model with softmax."""
-  hparams = rlmb_base_sv2p_softmax()
-  hparams.generative_model_params = (
-      "next_frame_sv2p_atari_softmax_deterministic")
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_flippy30():
-  """Base setting with sv2p as world model."""
-  hparams = rlmb_base()
-  hparams.epochs = 30
-  hparams.ppo_epochs_num = 1000
-  hparams.model_train_steps = 15000
-  hparams.learning_rate_bump = 1.0
-  hparams.inital_epoch_train_steps_multiplier = 5
-  hparams.generative_model = "next_frame_sv2p"
-  hparams.generative_model_params = "next_frame_sv2p_atari"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_softmax_flippy30():
-  """Base setting with sv2p as world model with softmax."""
-  hparams = rlmb_base_sv2p_flippy30()
-  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_deterministic_flippy30():
-  """Base setting with deterministic sv2p as world model."""
-  hparams = rlmb_base_sv2p_flippy30()
-  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_deterministic_softmax_flippy30():
-  """Base setting with deterministic sv2p as world model with softmax."""
-  hparams = rlmb_base_sv2p_softmax_flippy30()
-  hparams.generative_model_params = (
-      "next_frame_sv2p_atari_softmax_deterministic")
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sampling():
-  """Base setting with a stochastic next-frame model."""
-  hparams = rlmb_base()
-  hparams.generative_model_params = "next_frame_sampling"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sampling_noresize():
-  hparams = rlmb_base_sampling()
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_flippy60():
-  """Schedule with a lot of epochs (slow)."""
-  hparams = rlmb_base_sampling()
-  hparams.epochs = 60
-  hparams.ppo_epochs_num = 500
-  hparams.model_train_steps = 10000
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_flippy30():
-  """Schedule with a lot of epochs (slow)."""
-  hparams = rlmb_base_sampling()
-  hparams.epochs = 30
-  hparams.ppo_epochs_num = 1000
-  hparams.model_train_steps = 15000
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_medium():
-  """Small set for larger testing."""
-  hparams = rlmb_base()
-  hparams.num_real_env_frames //= 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_25k():
-  """Small set for larger testing."""
-  hparams = rlmb_medium()
-  hparams.num_real_env_frames //= 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_short():
-  """Small set for larger testing."""
-  hparams = rlmb_base()
-  hparams.num_real_env_frames //= 5
-  hparams.model_train_steps //= 10
-  hparams.ppo_epochs_num //= 10
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_model_only():
-  hp = rlmb_base()
-  hp.epochs = 1
-  hp.ppo_epochs_num = 0
-  return hp
-
-
-@registry.register_hparams
-def rlmb_tiny():
-  """Tiny set for testing."""
-  return rlmb_base_sampling().override_from_dict(
-      tf.contrib.training.HParams(
-          epochs=1,
-          num_real_env_frames=128,
-          simulated_env_generator_num_steps=64,
-          model_train_steps=2,
-          ppo_epochs_num=2,
-          ppo_time_limit=5,
-          ppo_epoch_length=5,
-          ppo_num_agents=2,
-          real_ppo_epoch_length=36,
-          real_ppo_num_agents=1,
-          real_ppo_epochs_num=0,
-          real_ppo_effective_num_agents=2,
-          generative_model_params="next_frame_tiny",
-          stop_loop_early=True,
-          resize_height_factor=2,
-          resize_width_factor=2,
-          game="pong",
-      ).values())
-
-
-@registry.register_hparams
-def rlmb_tiny_stochastic():
-  """Tiny setting with a stochastic next-frame model."""
-  hparams = rlmb_tiny()
-  hparams.epochs = 1  # Too slow with 2 for regular runs.
-  hparams.generative_model = "next_frame_basic_stochastic"
-  hparams.generative_model_params = "next_frame_basic_stochastic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_tiny_recurrent():
-  """Tiny setting with a recurrent next-frame model."""
-  hparams = rlmb_tiny()
-  hparams.epochs = 1  # Too slow with 2 for regular runs.
-  hparams.generative_model = "next_frame_basic_recurrent"
-  hparams.generative_model_params = "next_frame_basic_recurrent"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_tiny_sv2p():
-  """Tiny setting with a tiny sv2p model."""
-  hparams = rlmb_tiny()
-  hparams.generative_model = "next_frame_sv2p"
-  hparams.generative_model_params = "next_frame_sv2p_tiny"
-  hparams.grayscale = False
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_ae_base():
-  """Parameter set for autoencoders."""
-  hparams = rlmb_base()
-  hparams.ppo_params = "ppo_pong_ae_base"
-  hparams.generative_model_params = "next_frame_ae"
-  hparams.autoencoder_hparams_set = "autoencoder_discrete_pong"
-  hparams.gather_ppo_real_env_data = False
-  hparams.autoencoder_train_steps = 5000
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  hparams.grayscale = False
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_ae_basetest():
-  """Base AE setting but quicker with only 2 epochs."""
-  hparams = rlmb_ae_base()
-  hparams.game = "pong"
-  hparams.epochs = 2
-  hparams.num_real_env_frames = 3200
-  hparams.model_train_steps = 100
-  hparams.autoencoder_train_steps = 10
-  hparams.simulated_env_generator_num_steps = 20
-  hparams.ppo_epochs_num = 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_ae_tiny():
-  """Tiny set for testing autoencoders."""
-  hparams = rlmb_tiny()
-  hparams.ppo_params = "ppo_pong_ae_base"
-  hparams.generative_model_params = "next_frame_ae_tiny"
-  hparams.autoencoder_hparams_set = "autoencoder_discrete_tiny"
-  hparams.gather_ppo_real_env_data = False
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  hparams.grayscale = False
-  hparams.autoencoder_train_steps = 1
-  hparams.autoencoder_train_steps_initial_multiplier = 0
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_tiny_simulation_deterministic_starts():
-  hp = rlmb_tiny()
-  hp.simulation_random_starts = False
-  return hp
-
-
-# RangedHParams for tuning
-# ==============================================================================
-# Note that the items here must be scoped with one of
-# HP_SCOPES={loop, model, ppo}, which set hyperparameters for the top-level
-# hparams, hp.generative_model_params, and hp.ppo_params, respectively.
-@registry.register_ranged_hparams
-def rlmb_grid(rhp):
-  """Grid over games and frames, and 5 runs each for variance."""
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-  base = 100000
-  medium = base // 2
-  small = medium // 2
-  rhp.set_discrete("loop.num_real_env_frames", [base, medium, small])
-
-  # Dummy parameter to get 5 runs for each configuration
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-
-
-@registry.register_ranged_hparams
-def rlmb_variance(rhp):
-  # Dummy parameter to get 5 runs for each configuration
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-
-
-@registry.register_ranged_hparams
-def rlmb_variance_nogame(rhp):
-  # Dummy parameter to get 20 runs for current configuration.
-  rhp.set_discrete("model.moe_loss_coef", list(range(20)))
-
-
-@registry.register_ranged_hparams
-def rlmb_three(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
-
-
-@registry.register_ranged_hparams
-def rlmb_test1(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
-  rhp.set_discrete("loop.ppo_learning_rate", [5e-5, 1e-4, 2e-4])
-  rhp.set_discrete("ppo.optimization_batch_size", [20, 40])
-  rhp.set_discrete("loop.epochs", [3, 6])
-
-
-@registry.register_ranged_hparams
-def rlmb_scheduled_sampling(rhp):
-  rhp.set_float("model.scheduled_sampling_prob", 0.0, 1.0)
-
-
-@registry.register_ranged_hparams
-def rlmb_all_games(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_whitelisted_games(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_human_score_games(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game",
-                      gym_problems_specs.ATARI_GAMES_WITH_HUMAN_SCORE)
-
-
-@registry.register_ranged_hparams
-def rlmb_curious_games10(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_curious_games5(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_debug_games(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_DEBUG_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_ae_variance(rhp):
-  # Dummy parameter to get 5 runs for each configuration
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-  base = 100000
-  small = base // 4
-  rhp.set_discrete("loop.num_real_env_frames", [base, small])
-
-
-@registry.register_ranged_hparams
-def rlmb_ppolr_game(rhp):
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-  base_lr = 1e-4
-  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
-
-
-@registry.register_ranged_hparams
-def rlmb_ppolr(rhp):
-  base_lr = 1e-4
-  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
-
-
-@registry.register_ranged_hparams
-def rlmb_ae_ppo_lr(rhp):
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-  base_lr = 1e-4
-  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
-
-
-@registry.register_ranged_hparams
-def rlmb_dropout_range(rhp):
-  rhp.set_float("model.dropout", 0.2, 0.4)
-
-
-@registry.register_ranged_hparams
-def rlmb_intrinsic_reward_scale(rhp):
-  rhp.set_float("loop.intrinsic_reward_scale", 0.01, 10.)
-
-
-@registry.register_ranged_hparams
-def rlmb_l1l2cutoff_range(rhp):
-  """Loss and loss-cutoff tuning grid."""
-  rhp.set_float("model.video_modality_loss_cutoff", 1.4, 3.4)
-
-
-@registry.register_ranged_hparams
-def rlmb_xentcutoff_range(rhp):
-  """Cross entropy cutoff tuning grid."""
-  rhp.set_float("model.video_modality_loss_cutoff", 0.01, 0.05)
-
-
-@registry.register_ranged_hparams
-def rlmb_pixel_noise(rhp):
-  """Input pixel noise tuning grid."""
-  rhp.set_categorical("loop.generative_model_params",
-                      ["next_frame_pixel_noise"])
-  rhp.set_discrete("model.video_modality_input_noise",
-                   [0.0025 * i for i in range(200)])
-
-
-@registry.register_ranged_hparams
-def rlmb_dummy_range(rhp):
-  """Dummy tuning grid just to get the variance."""
-  rhp.set_float("model.moe_loss_coef", 0.01, 0.02)
-
-
-@registry.register_ranged_hparams
-def rlmb_epochs_num(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_discrete("loop.epochs", [3, 6, 12])
-
-
-@registry.register_ranged_hparams
-def rlmb_ppo_epochs_num(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_discrete("loop.ppo_epochs_num", [200, 1000, 2000, 4000])
-
-
-@registry.register_ranged_hparams
-def rlmb_ppo_epoch_len(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_discrete("loop.ppo_epoch_length", [25, 50, 100])
-
-
-@registry.register_ranged_hparams
-def rlmb_num_frames(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_discrete("loop.num_real_env_frames",
-                   [1000*el for el in [30, 100, 500, 1000]])
-
-
-@registry.register_ranged_hparams
-def rlmb_ppo_optimization_batch_size(rhp):
-  rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_discrete("ppo.optimization_batch_size", [4, 10, 20])
-
-
-@registry.register_ranged_hparams
-def rlmb_logits_clip(rhp):
-  rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_discrete("ppo.logits_clip", [0., 5.])
-
-
-def merge_unscoped_hparams(scopes_and_hparams):
-  """Merge multiple HParams into one with scopes."""
-  merged_values = {}
-  for (scope, hparams) in scopes_and_hparams:
-    for key, value in six.iteritems(hparams.values()):
-      scoped_key = "%s.%s" % (scope, key)
-      merged_values[scoped_key] = value
-
-  return tf.contrib.training.HParams(**merged_values)
-
-
-def split_scoped_hparams(scopes, merged_hparams):
-  """Split single HParams with scoped keys into multiple."""
-  split_values = {scope: {} for scope in scopes}
-  merged_values = merged_hparams.values()
-  for scoped_key, value in six.iteritems(merged_values):
-    scope = scoped_key.split(".")[0]
-    key = scoped_key[len(scope) + 1:]
-    split_values[scope][key] = value
-
-  return [
-      tf.contrib.training.HParams(**split_values[scope]) for scope in scopes
-  ]
-
-
-def training_loop_hparams_from_scoped_overrides(scoped_overrides, trial_id):
-  """Create HParams suitable for training loop from scoped HParams.
-
-  Args:
-    scoped_overrides: HParams, with keys all scoped by one of HP_SCOPES. These
-      parameters are overrides for the base HParams created by
-      create_loop_hparams.
-    trial_id: str, trial identifier. This is used to register unique HParams
-      names for the underlying model and ppo HParams.
-
-  Returns:
-    HParams suitable for passing to training_loop.
-  """
-  trial_hp_overrides = scoped_overrides.values()
-
-  # Create loop, model, and ppo base HParams
-  loop_hp = create_loop_hparams()
-  model_hp_name = trial_hp_overrides.get(
-      "loop.generative_model_params", loop_hp.generative_model_params)
-  model_hp = registry.hparams(model_hp_name).parse(FLAGS.hparams)
-  ppo_params_name = trial_hp_overrides.get(
-      "loop.ppo_params", loop_hp.ppo_params)
-  ppo_hp = registry.hparams(ppo_params_name)
-
-  # Merge them and then override with the scoped overrides
-  combined_hp = merge_unscoped_hparams(
-      zip(HP_SCOPES, [loop_hp, model_hp, ppo_hp]))
-  combined_hp.override_from_dict(trial_hp_overrides)
-
-  # Split out the component hparams
-  loop_hp, model_hp, ppo_hp = (
-      split_scoped_hparams(HP_SCOPES, combined_hp))
-
-  # Dynamic register the model hp and set the new name in loop_hp
-  model_hp_name = "model_hp_%s" % str(trial_id)
-  dynamic_register_hparams(model_hp_name, model_hp)
-  loop_hp.generative_model_params = model_hp_name
-
-  # Dynamic register the PPO hp and set the new name in loop_hp
-  ppo_hp_name = "ppo_hp_%s" % str(trial_id)
-  dynamic_register_hparams(ppo_hp_name, ppo_hp)
-  loop_hp.ppo_params = ppo_hp_name
-
-  return loop_hp
-
-
-def dynamic_register_hparams(name, hparams):
-
-  @registry.register_hparams(name)
-  def new_hparams_set():
-    return tf.contrib.training.HParams(**hparams.values())
-
-  return new_hparams_set
-
-
-def create_loop_hparams():
-  hparams = registry.hparams(FLAGS.loop_hparams_set)
-  hparams.parse(FLAGS.loop_hparams)
-  return hparams

From b4df57ee9f6ddff6edae9b967f7020bc6b46277a Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 11 Oct 2018 16:31:57 -0700
Subject: [PATCH 0992/2720] internal merge of PR #1137

PiperOrigin-RevId: 216781795
---
 tensor2tensor/rl/trainer_model_based.py       | 725 +----------------
 .../rl/trainer_model_based_agent_only.py      |   3 +-
 .../rl/trainer_model_based_params.py          | 753 ++++++++++++++++++
 3 files changed, 757 insertions(+), 724 deletions(-)
 create mode 100644 tensor2tensor/rl/trainer_model_based_params.py

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 03fe2a525..bde551c6a 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -44,6 +44,7 @@
 from tensor2tensor.data_generators import gym_problems_specs
 from tensor2tensor.layers import discretization
 from tensor2tensor.rl import rl_trainer_lib
+from tensor2tensor.rl import trainer_model_based_params
 from tensor2tensor.rl.envs.tf_atari_wrappers import PyFuncWrapper
 from tensor2tensor.rl.envs.utils import InitialFrameChooser
 from tensor2tensor.utils import registry
@@ -63,8 +64,6 @@
 flags.DEFINE_string("eval_results_dir", "/tmp",
                     "Directory to store result of evaluation")
 
-HP_SCOPES = ["loop", "model", "ppo"]
-
 
 def setup_directories(base_dir, subdirs):
   base_dir = os.path.expanduser(base_dir)
@@ -748,728 +747,8 @@ def combine_training_data(problem, final_data_dir, old_data_dirs,
         tf.gfile.Copy(fname, new_fname)
 
 
-@registry.register_hparams
-def rlmb_base():
-  return tf.contrib.training.HParams(
-      epochs=15,
-      # Total frames used for training. This will be distributed evenly across
-      # hparams.epochs.
-      # This number should be divisible by real_ppo_epoch_length*epochs
-      # for our frame accounting to be preceise.
-      num_real_env_frames=96000,
-      generative_model="next_frame_basic_deterministic",
-      generative_model_params="next_frame_pixel_noise",
-      ppo_params="ppo_pong_base",
-      autoencoder_train_steps=0,
-      autoencoder_train_steps_initial_multiplier=10,
-      autoencoder_hparams_set="autoencoder_discrete_pong",
-      model_train_steps=15000,
-      inital_epoch_train_steps_multiplier=3,
-      simulated_env_generator_num_steps=2000,
-      simulation_random_starts=True,  # Use random starts in PPO.
-      # Flip the first random frame in PPO batch for the true beginning.
-      simulation_flip_first_random_for_beginning=True,
-      intrinsic_reward_scale=0.,
-      ppo_epochs_num=1000,  # This should be enough to see something
-      # Our simulated envs do not know how to reset.
-      # You should set ppo_time_limit to the value you believe that
-      # the simulated env produces a reasonable output.
-      ppo_time_limit=200,  # TODO(blazej): this param is unused
-      # It makes sense to have ppo_time_limit=ppo_epoch_length,
-      # though it is not necessary.
-      ppo_epoch_length=50,
-      ppo_num_agents=16,
-      # Do not eval since simulated batch env does not produce dones
-      ppo_eval_every_epochs=0,
-      ppo_learning_rate=1e-4,  # Will be changed, just so it exists.
-      # Whether the PPO agent should be restored from the previous iteration, or
-      # should start fresh each time.
-      ppo_continue_training=True,
-      # Resizing.
-      resize_height_factor=2,
-      resize_width_factor=2,
-      grayscale=False,
-      # Bump learning rate after first epoch by 3x.
-      # We picked 3x because our default learning rate schedule decreases with
-      # 1/square root of step; 1/sqrt(10k) = 0.01 and 1/sqrt(100k) ~ 0.0032
-      # so by bumping it up 3x we about "go back" from 100k steps to 10k, which
-      # is approximately as much as "going back 1 epoch" would be.
-      # In your experiments, you want to optimize this rate to your schedule.
-      learning_rate_bump=3.0,
-
-      gather_ppo_real_env_data=True,
-      real_ppo_epochs_num=0,
-      # This needs to be divisible by real_ppo_effective_num_agents.
-      real_ppo_epoch_length=16*200,
-      real_ppo_num_agents=1,
-      real_ppo_learning_rate=1e-4,
-      real_ppo_continue_training=True,
-      real_ppo_effective_num_agents=16,
-      real_ppo_eval_every_epochs=0,
-
-      game="pong",
-      # Whether to evaluate the world model in each iteration of the loop to get
-      # the model_reward_accuracy metric.
-      eval_world_model=True,
-      # Rollout fractions to report reward_accuracy on.
-      eval_rollout_fractions=[0.25, 0.5, 1],
-      stop_loop_early=False,  # To speed-up tests.
-  )
-
-
-@registry.register_hparams
-def rlmb_basetest():
-  """Base setting but quicker with only 2 epochs."""
-  hparams = rlmb_base()
-  hparams.game = "pong"
-  hparams.epochs = 2
-  hparams.num_real_env_frames = 3200
-  hparams.model_train_steps = 100
-  hparams.simulated_env_generator_num_steps = 20
-  hparams.ppo_epochs_num = 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_noresize():
-  hparams = rlmb_base()
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_quick():
-  """Base setting but quicker with only 2 epochs."""
-  hparams = rlmb_base()
-  hparams.epochs = 2
-  hparams.model_train_steps = 25000
-  hparams.ppo_epochs_num = 700
-  hparams.ppo_epoch_length = 50
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_quick_noresize():
-  hparams = rlmb_base()
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_quick_sd():
-  """Quick setting with stochastic discrete model."""
-  hparams = rlmb_quick()
-  hparams.generative_model = "next_frame_basic_stochastic_discrete"
-  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_sdtest():
-  """Test setting with stochastic discrete model."""
-  hparams = rlmb_basetest()
-  hparams.generative_model = "next_frame_basic_stochastic_discrete"
-  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_quick_sm():
-  """Quick setting with sampling."""
-  hparams = rlmb_quick()
-  hparams.generative_model_params = "next_frame_sampling"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_stochastic():
-  """Base setting with a stochastic next-frame model."""
-  hparams = rlmb_base()
-  hparams.inital_epoch_train_steps_multiplier = 5
-  hparams.generative_model = "next_frame_basic_stochastic"
-  hparams.generative_model_params = "next_frame_basic_stochastic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sampling_stochastic():
-  """Base setting with a stochastic next-frame model."""
-  hparams = rlmb_base()
-  hparams.generative_model = "next_frame_basic_stochastic"
-  hparams.generative_model_params = "next_frame_sampling_stochastic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_stochastic_discrete():
-  """Base setting with stochastic discrete model."""
-  hparams = rlmb_base()
-  hparams.learning_rate_bump = 1.0
-  hparams.grayscale = True
-  hparams.generative_model = "next_frame_basic_stochastic_discrete"
-  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_stochastic_recurrent():
-  """Base setting with recurrent model."""
-  hparams = rlmb_base()
-  hparams.generative_model = "next_frame_basic_recurrent"
-  hparams.generative_model_params = "next_frame_basic_recurrent"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_stochastic_discrete_noresize():
-  """Base setting with stochastic discrete model."""
-  hparams = rlmb_base()
-  hparams.generative_model = "next_frame_basic_stochastic_discrete"
-  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p():
-  """Base setting with sv2p as world model."""
-  hparams = rlmb_base()
-  hparams.learning_rate_bump = 1.0
-  hparams.generative_model = "next_frame_sv2p"
-  hparams.generative_model_params = "next_frame_sv2p_atari"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_softmax():
-  """Base setting with sv2p as world model with softmax."""
-  hparams = rlmb_base_sv2p()
-  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_deterministic():
-  """Base setting with deterministic sv2p as world model."""
-  hparams = rlmb_base_sv2p()
-  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_deterministic_softmax():
-  """Base setting with deterministic sv2p as world model with softmax."""
-  hparams = rlmb_base_sv2p_softmax()
-  hparams.generative_model_params = (
-      "next_frame_sv2p_atari_softmax_deterministic")
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_flippy30():
-  """Base setting with sv2p as world model."""
-  hparams = rlmb_base()
-  hparams.epochs = 30
-  hparams.ppo_epochs_num = 1000
-  hparams.model_train_steps = 15000
-  hparams.learning_rate_bump = 1.0
-  hparams.inital_epoch_train_steps_multiplier = 5
-  hparams.generative_model = "next_frame_sv2p"
-  hparams.generative_model_params = "next_frame_sv2p_atari"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_softmax_flippy30():
-  """Base setting with sv2p as world model with softmax."""
-  hparams = rlmb_base_sv2p_flippy30()
-  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_deterministic_flippy30():
-  """Base setting with deterministic sv2p as world model."""
-  hparams = rlmb_base_sv2p_flippy30()
-  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_deterministic_softmax_flippy30():
-  """Base setting with deterministic sv2p as world model with softmax."""
-  hparams = rlmb_base_sv2p_softmax_flippy30()
-  hparams.generative_model_params = (
-      "next_frame_sv2p_atari_softmax_deterministic")
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sampling():
-  """Base setting with a stochastic next-frame model."""
-  hparams = rlmb_base()
-  hparams.generative_model_params = "next_frame_sampling"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sampling_noresize():
-  hparams = rlmb_base_sampling()
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_flippy60():
-  """Schedule with a lot of epochs (slow)."""
-  hparams = rlmb_base_sampling()
-  hparams.epochs = 60
-  hparams.ppo_epochs_num = 500
-  hparams.model_train_steps = 10000
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_flippy30():
-  """Schedule with a lot of epochs (slow)."""
-  hparams = rlmb_base_sampling()
-  hparams.epochs = 30
-  hparams.ppo_epochs_num = 1000
-  hparams.model_train_steps = 15000
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_medium():
-  """Small set for larger testing."""
-  hparams = rlmb_base()
-  hparams.num_real_env_frames //= 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_25k():
-  """Small set for larger testing."""
-  hparams = rlmb_medium()
-  hparams.num_real_env_frames //= 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_short():
-  """Small set for larger testing."""
-  hparams = rlmb_base()
-  hparams.num_real_env_frames //= 5
-  hparams.model_train_steps //= 10
-  hparams.ppo_epochs_num //= 10
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_model_only():
-  hp = rlmb_base()
-  hp.epochs = 1
-  hp.ppo_epochs_num = 0
-  return hp
-
-
-@registry.register_hparams
-def rlmb_tiny():
-  """Tiny set for testing."""
-  return rlmb_base_sampling().override_from_dict(
-      tf.contrib.training.HParams(
-          epochs=1,
-          num_real_env_frames=128,
-          simulated_env_generator_num_steps=64,
-          model_train_steps=2,
-          ppo_epochs_num=2,
-          ppo_time_limit=5,
-          ppo_epoch_length=5,
-          ppo_num_agents=2,
-          real_ppo_epoch_length=36,
-          real_ppo_num_agents=1,
-          real_ppo_epochs_num=0,
-          real_ppo_effective_num_agents=2,
-          generative_model_params="next_frame_tiny",
-          stop_loop_early=True,
-          resize_height_factor=2,
-          resize_width_factor=2,
-          game="pong",
-      ).values())
-
-
-@registry.register_hparams
-def rlmb_tiny_stochastic():
-  """Tiny setting with a stochastic next-frame model."""
-  hparams = rlmb_tiny()
-  hparams.epochs = 1  # Too slow with 2 for regular runs.
-  hparams.generative_model = "next_frame_basic_stochastic"
-  hparams.generative_model_params = "next_frame_basic_stochastic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_tiny_recurrent():
-  """Tiny setting with a recurrent next-frame model."""
-  hparams = rlmb_tiny()
-  hparams.epochs = 1  # Too slow with 2 for regular runs.
-  hparams.generative_model = "next_frame_basic_recurrent"
-  hparams.generative_model_params = "next_frame_basic_recurrent"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_tiny_sv2p():
-  """Tiny setting with a tiny sv2p model."""
-  hparams = rlmb_tiny()
-  hparams.generative_model = "next_frame_sv2p"
-  hparams.generative_model_params = "next_frame_sv2p_tiny"
-  hparams.grayscale = False
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_ae_base():
-  """Parameter set for autoencoders."""
-  hparams = rlmb_base()
-  hparams.ppo_params = "ppo_pong_ae_base"
-  hparams.generative_model_params = "next_frame_ae"
-  hparams.autoencoder_hparams_set = "autoencoder_discrete_pong"
-  hparams.gather_ppo_real_env_data = False
-  hparams.autoencoder_train_steps = 5000
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  hparams.grayscale = False
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_ae_basetest():
-  """Base AE setting but quicker with only 2 epochs."""
-  hparams = rlmb_ae_base()
-  hparams.game = "pong"
-  hparams.epochs = 2
-  hparams.num_real_env_frames = 3200
-  hparams.model_train_steps = 100
-  hparams.autoencoder_train_steps = 10
-  hparams.simulated_env_generator_num_steps = 20
-  hparams.ppo_epochs_num = 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_ae_tiny():
-  """Tiny set for testing autoencoders."""
-  hparams = rlmb_tiny()
-  hparams.ppo_params = "ppo_pong_ae_base"
-  hparams.generative_model_params = "next_frame_ae_tiny"
-  hparams.autoencoder_hparams_set = "autoencoder_discrete_tiny"
-  hparams.gather_ppo_real_env_data = False
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  hparams.grayscale = False
-  hparams.autoencoder_train_steps = 1
-  hparams.autoencoder_train_steps_initial_multiplier = 0
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_tiny_simulation_deterministic_starts():
-  hp = rlmb_tiny()
-  hp.simulation_random_starts = False
-  return hp
-
-
-# RangedHParams for tuning
-# ==============================================================================
-# Note that the items here must be scoped with one of
-# HP_SCOPES={loop, model, ppo}, which set hyperparameters for the top-level
-# hparams, hp.generative_model_params, and hp.ppo_params, respectively.
-@registry.register_ranged_hparams
-def rlmb_grid(rhp):
-  """Grid over games and frames, and 5 runs each for variance."""
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-  base = 100000
-  medium = base // 2
-  small = medium // 2
-  rhp.set_discrete("loop.num_real_env_frames", [base, medium, small])
-
-  # Dummy parameter to get 5 runs for each configuration
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-
-
-@registry.register_ranged_hparams
-def rlmb_variance(rhp):
-  # Dummy parameter to get 5 runs for each configuration
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-
-
-@registry.register_ranged_hparams
-def rlmb_variance_nogame(rhp):
-  # Dummy parameter to get 20 runs for current configuration.
-  rhp.set_discrete("model.moe_loss_coef", list(range(20)))
-
-
-@registry.register_ranged_hparams
-def rlmb_three(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
-
-
-@registry.register_ranged_hparams
-def rlmb_test1(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
-  rhp.set_discrete("loop.ppo_learning_rate", [5e-5, 1e-4, 2e-4])
-  rhp.set_discrete("ppo.optimization_batch_size", [20, 40])
-  rhp.set_discrete("loop.epochs", [3, 6])
-
-
-@registry.register_ranged_hparams
-def rlmb_scheduled_sampling(rhp):
-  rhp.set_float("model.scheduled_sampling_prob", 0.0, 1.0)
-
-
-@registry.register_ranged_hparams
-def rlmb_all_games(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_whitelisted_games(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_human_score_games(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game",
-                      gym_problems_specs.ATARI_GAMES_WITH_HUMAN_SCORE)
-
-
-@registry.register_ranged_hparams
-def rlmb_curious_games10(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_curious_games5(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_debug_games(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_DEBUG_GAMES)
-
-
-@registry.register_ranged_hparams
-def rlmb_ae_variance(rhp):
-  # Dummy parameter to get 5 runs for each configuration
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-  base = 100000
-  small = base // 4
-  rhp.set_discrete("loop.num_real_env_frames", [base, small])
-
-
-@registry.register_ranged_hparams
-def rlmb_ppolr_game(rhp):
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-  base_lr = 1e-4
-  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
-
-
-@registry.register_ranged_hparams
-def rlmb_ppolr(rhp):
-  base_lr = 1e-4
-  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
-
-
-@registry.register_ranged_hparams
-def rlmb_ae_ppo_lr(rhp):
-  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
-  base_lr = 1e-4
-  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
-
-
-@registry.register_ranged_hparams
-def rlmb_dropout_range(rhp):
-  rhp.set_float("model.dropout", 0.2, 0.4)
-
-
-@registry.register_ranged_hparams
-def rlmb_intrinsic_reward_scale(rhp):
-  rhp.set_float("loop.intrinsic_reward_scale", 0.01, 10.)
-
-
-@registry.register_ranged_hparams
-def rlmb_l1l2cutoff_range(rhp):
-  """Loss and loss-cutoff tuning grid."""
-  rhp.set_float("model.video_modality_loss_cutoff", 1.4, 3.4)
-
-
-@registry.register_ranged_hparams
-def rlmb_xentcutoff_range(rhp):
-  """Cross entropy cutoff tuning grid."""
-  rhp.set_float("model.video_modality_loss_cutoff", 0.01, 0.05)
-
-
-@registry.register_ranged_hparams
-def rlmb_pixel_noise(rhp):
-  """Input pixel noise tuning grid."""
-  rhp.set_categorical("loop.generative_model_params",
-                      ["next_frame_pixel_noise"])
-  rhp.set_discrete("model.video_modality_input_noise",
-                   [0.0025 * i for i in range(200)])
-
-
-@registry.register_ranged_hparams
-def rlmb_dummy_range(rhp):
-  """Dummy tuning grid just to get the variance."""
-  rhp.set_float("model.moe_loss_coef", 0.01, 0.02)
-
-
-@registry.register_ranged_hparams
-def rlmb_epochs_num(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_discrete("loop.epochs", [3, 6, 12])
-
-
-@registry.register_ranged_hparams
-def rlmb_ppo_epochs_num(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_discrete("loop.ppo_epochs_num", [200, 1000, 2000, 4000])
-
-
-@registry.register_ranged_hparams
-def rlmb_ppo_epoch_len(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_discrete("loop.ppo_epoch_length", [25, 50, 100])
-
-
-@registry.register_ranged_hparams
-def rlmb_num_frames(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
-  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_discrete("loop.num_real_env_frames",
-                   [1000*el for el in [30, 100, 500, 1000]])
-
-
-@registry.register_ranged_hparams
-def rlmb_ppo_optimization_batch_size(rhp):
-  rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_discrete("ppo.optimization_batch_size", [4, 10, 20])
-
-
-@registry.register_ranged_hparams
-def rlmb_logits_clip(rhp):
-  rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_discrete("ppo.logits_clip", [0., 5.])
-
-
-def merge_unscoped_hparams(scopes_and_hparams):
-  """Merge multiple HParams into one with scopes."""
-  merged_values = {}
-  for (scope, hparams) in scopes_and_hparams:
-    for key, value in six.iteritems(hparams.values()):
-      scoped_key = "%s.%s" % (scope, key)
-      merged_values[scoped_key] = value
-
-  return tf.contrib.training.HParams(**merged_values)
-
-
-def split_scoped_hparams(scopes, merged_hparams):
-  """Split single HParams with scoped keys into multiple."""
-  split_values = dict([(scope, dict()) for scope in scopes])
-  merged_values = merged_hparams.values()
-  for scoped_key, value in six.iteritems(merged_values):
-    scope = scoped_key.split(".")[0]
-    key = scoped_key[len(scope) + 1:]
-    split_values[scope][key] = value
-
-  return [
-      tf.contrib.training.HParams(**split_values[scope]) for scope in scopes
-  ]
-
-
-def training_loop_hparams_from_scoped_overrides(scoped_overrides, trial_id):
-  """Create HParams suitable for training loop from scoped HParams.
-
-  Args:
-    scoped_overrides: HParams, with keys all scoped by one of HP_SCOPES. These
-      parameters are overrides for the base HParams created by
-      create_loop_hparams.
-    trial_id: str, trial identifier. This is used to register unique HParams
-      names for the underlying model and ppo HParams.
-
-  Returns:
-    HParams suitable for passing to training_loop.
-  """
-  trial_hp_overrides = scoped_overrides.values()
-
-  # Create loop, model, and ppo base HParams
-  loop_hp = create_loop_hparams()
-  model_hp_name = trial_hp_overrides.get(
-      "loop.generative_model_params", loop_hp.generative_model_params)
-  model_hp = registry.hparams(model_hp_name).parse(FLAGS.hparams)
-  ppo_params_name = trial_hp_overrides.get(
-      "loop.ppo_params", loop_hp.ppo_params)
-  ppo_hp = registry.hparams(ppo_params_name)
-
-  # Merge them and then override with the scoped overrides
-  combined_hp = merge_unscoped_hparams(
-      zip(HP_SCOPES, [loop_hp, model_hp, ppo_hp]))
-  combined_hp.override_from_dict(trial_hp_overrides)
-
-  # Split out the component hparams
-  loop_hp, model_hp, ppo_hp = (
-      split_scoped_hparams(HP_SCOPES, combined_hp))
-
-  # Dynamic register the model hp and set the new name in loop_hp
-  model_hp_name = "model_hp_%s" % str(trial_id)
-  dynamic_register_hparams(model_hp_name, model_hp)
-  loop_hp.generative_model_params = model_hp_name
-
-  # Dynamic register the PPO hp and set the new name in loop_hp
-  ppo_hp_name = "ppo_hp_%s" % str(trial_id)
-  dynamic_register_hparams(ppo_hp_name, ppo_hp)
-  loop_hp.ppo_params = ppo_hp_name
-
-  return loop_hp
-
-
-def dynamic_register_hparams(name, hparams):
-
-  @registry.register_hparams(name)
-  def new_hparams_set():
-    return tf.contrib.training.HParams(**hparams.values())
-
-  return new_hparams_set
-
-
-def create_loop_hparams():
-  hparams = registry.hparams(FLAGS.loop_hparams_set)
-  hparams.parse(FLAGS.loop_hparams)
-  return hparams
-
-
 def main(_):
-  hp = create_loop_hparams()
+  hp = trainer_model_based_params.create_loop_hparams()
   if FLAGS.job_dir_to_evaluate:
     compute_final_evaluation_on_real_environments(hp, FLAGS.job_dir_to_evaluate)
   else:
diff --git a/tensor2tensor/rl/trainer_model_based_agent_only.py b/tensor2tensor/rl/trainer_model_based_agent_only.py
index bf4ac0b12..cced3ebc5 100644
--- a/tensor2tensor/rl/trainer_model_based_agent_only.py
+++ b/tensor2tensor/rl/trainer_model_based_agent_only.py
@@ -30,6 +30,7 @@
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
 from tensor2tensor.data_generators import gym_problems_specs
 from tensor2tensor.rl import trainer_model_based
+from tensor2tensor.rl import trainer_model_based_params
 
 
 import tensorflow as tf
@@ -50,7 +51,7 @@ def get_simulated_problem_name(game):
 
 
 def main(_):
-  hparams = trainer_model_based.create_loop_hparams()
+  hparams = trainer_model_based_params.create_loop_hparams()
   problem_name = get_simulated_problem_name(hparams.game)
   world_model_dir = FLAGS.world_model_dir
   agent_model_dir = FLAGS.output_dir
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
new file mode 100644
index 000000000..953a2e5e5
--- /dev/null
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -0,0 +1,753 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+r"""Parameter sets for training of model-based RL agents."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import six
+
+
+from tensor2tensor.data_generators import gym_problems_specs
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+FLAGS = tf.flags.FLAGS
+
+
+HP_SCOPES = ["loop", "model", "ppo"]
+
+
+@registry.register_hparams
+def rlmb_base():
+  return tf.contrib.training.HParams(
+      epochs=15,
+      # Total frames used for training. This will be distributed evenly across
+      # hparams.epochs.
+      # This number should be divisible by real_ppo_epoch_length*epochs
+      # for our frame accounting to be preceise.
+      num_real_env_frames=96000,
+      generative_model="next_frame_basic_deterministic",
+      generative_model_params="next_frame_pixel_noise",
+      ppo_params="ppo_pong_base",
+      autoencoder_train_steps=0,
+      autoencoder_train_steps_initial_multiplier=10,
+      autoencoder_hparams_set="autoencoder_discrete_pong",
+      model_train_steps=15000,
+      inital_epoch_train_steps_multiplier=3,
+      simulated_env_generator_num_steps=2000,
+      simulation_random_starts=True,  # Use random starts in PPO.
+      # Flip the first random frame in PPO batch for the true beginning.
+      simulation_flip_first_random_for_beginning=True,
+      intrinsic_reward_scale=0.,
+      ppo_epochs_num=1000,  # This should be enough to see something
+      # Our simulated envs do not know how to reset.
+      # You should set ppo_time_limit to the value you believe that
+      # the simulated env produces a reasonable output.
+      ppo_time_limit=200,  # TODO(blazej): this param is unused
+      # It makes sense to have ppo_time_limit=ppo_epoch_length,
+      # though it is not necessary.
+      ppo_epoch_length=50,
+      ppo_num_agents=16,
+      # Do not eval since simulated batch env does not produce dones
+      ppo_eval_every_epochs=0,
+      ppo_learning_rate=1e-4,  # Will be changed, just so it exists.
+      # Whether the PPO agent should be restored from the previous iteration, or
+      # should start fresh each time.
+      ppo_continue_training=True,
+      # Resizing.
+      resize_height_factor=2,
+      resize_width_factor=2,
+      grayscale=False,
+      # Bump learning rate after first epoch by 3x.
+      # We picked 3x because our default learning rate schedule decreases with
+      # 1/square root of step; 1/sqrt(10k) = 0.01 and 1/sqrt(100k) ~ 0.0032
+      # so by bumping it up 3x we about "go back" from 100k steps to 10k, which
+      # is approximately as much as "going back 1 epoch" would be.
+      # In your experiments, you want to optimize this rate to your schedule.
+      learning_rate_bump=3.0,
+
+      gather_ppo_real_env_data=True,
+      real_ppo_epochs_num=0,
+      # This needs to be divisible by real_ppo_effective_num_agents.
+      real_ppo_epoch_length=16*200,
+      real_ppo_num_agents=1,
+      real_ppo_learning_rate=1e-4,
+      real_ppo_continue_training=True,
+      real_ppo_effective_num_agents=16,
+      real_ppo_eval_every_epochs=0,
+
+      game="pong",
+      # Whether to evaluate the world model in each iteration of the loop to get
+      # the model_reward_accuracy metric.
+      eval_world_model=True,
+      # Rollout fractions to report reward_accuracy on.
+      eval_rollout_fractions=[0.25, 0.5, 1],
+      stop_loop_early=False,  # To speed-up tests.
+  )
+
+
+@registry.register_hparams
+def rlmb_basetest():
+  """Base setting but quicker with only 2 epochs."""
+  hparams = rlmb_base()
+  hparams.game = "pong"
+  hparams.epochs = 2
+  hparams.num_real_env_frames = 3200
+  hparams.model_train_steps = 100
+  hparams.simulated_env_generator_num_steps = 20
+  hparams.ppo_epochs_num = 2
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_noresize():
+  hparams = rlmb_base()
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_quick():
+  """Base setting but quicker with only 2 epochs."""
+  hparams = rlmb_base()
+  hparams.epochs = 2
+  hparams.model_train_steps = 25000
+  hparams.ppo_epochs_num = 700
+  hparams.ppo_epoch_length = 50
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_quick_noresize():
+  hparams = rlmb_base()
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_quick_sd():
+  """Quick setting with stochastic discrete model."""
+  hparams = rlmb_quick()
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_sdtest():
+  """Test setting with stochastic discrete model."""
+  hparams = rlmb_basetest()
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_quick_sm():
+  """Quick setting with sampling."""
+  hparams = rlmb_quick()
+  hparams.generative_model_params = "next_frame_sampling"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic():
+  """Base setting with a stochastic next-frame model."""
+  hparams = rlmb_base()
+  hparams.inital_epoch_train_steps_multiplier = 5
+  hparams.generative_model = "next_frame_basic_stochastic"
+  hparams.generative_model_params = "next_frame_basic_stochastic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sampling_stochastic():
+  """Base setting with a stochastic next-frame model."""
+  hparams = rlmb_base()
+  hparams.generative_model = "next_frame_basic_stochastic"
+  hparams.generative_model_params = "next_frame_sampling_stochastic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic_discrete():
+  """Base setting with stochastic discrete model."""
+  hparams = rlmb_base()
+  hparams.learning_rate_bump = 1.0
+  hparams.grayscale = True
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic_recurrent():
+  """Base setting with recurrent model."""
+  hparams = rlmb_base()
+  hparams.generative_model = "next_frame_basic_recurrent"
+  hparams.generative_model_params = "next_frame_basic_recurrent"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic_discrete_noresize():
+  """Base setting with stochastic discrete model."""
+  hparams = rlmb_base()
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p():
+  """Base setting with sv2p as world model."""
+  hparams = rlmb_base()
+  hparams.learning_rate_bump = 1.0
+  hparams.generative_model = "next_frame_sv2p"
+  hparams.generative_model_params = "next_frame_sv2p_atari"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_softmax():
+  """Base setting with sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p()
+  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic():
+  """Base setting with deterministic sv2p as world model."""
+  hparams = rlmb_base_sv2p()
+  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic_softmax():
+  """Base setting with deterministic sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p_softmax()
+  hparams.generative_model_params = (
+      "next_frame_sv2p_atari_softmax_deterministic")
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_flippy30():
+  """Base setting with sv2p as world model."""
+  hparams = rlmb_base()
+  hparams.epochs = 30
+  hparams.ppo_epochs_num = 1000
+  hparams.model_train_steps = 15000
+  hparams.learning_rate_bump = 1.0
+  hparams.inital_epoch_train_steps_multiplier = 5
+  hparams.generative_model = "next_frame_sv2p"
+  hparams.generative_model_params = "next_frame_sv2p_atari"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_softmax_flippy30():
+  """Base setting with sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p_flippy30()
+  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic_flippy30():
+  """Base setting with deterministic sv2p as world model."""
+  hparams = rlmb_base_sv2p_flippy30()
+  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sv2p_deterministic_softmax_flippy30():
+  """Base setting with deterministic sv2p as world model with softmax."""
+  hparams = rlmb_base_sv2p_softmax_flippy30()
+  hparams.generative_model_params = (
+      "next_frame_sv2p_atari_softmax_deterministic")
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sampling():
+  """Base setting with a stochastic next-frame model."""
+  hparams = rlmb_base()
+  hparams.generative_model_params = "next_frame_sampling"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_sampling_noresize():
+  hparams = rlmb_base_sampling()
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_flippy60():
+  """Schedule with a lot of epochs (slow)."""
+  hparams = rlmb_base_sampling()
+  hparams.epochs = 60
+  hparams.ppo_epochs_num = 500
+  hparams.model_train_steps = 10000
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_flippy30():
+  """Schedule with a lot of epochs (slow)."""
+  hparams = rlmb_base_sampling()
+  hparams.epochs = 30
+  hparams.ppo_epochs_num = 1000
+  hparams.model_train_steps = 15000
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_medium():
+  """Small set for larger testing."""
+  hparams = rlmb_base()
+  hparams.num_real_env_frames //= 2
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_25k():
+  """Small set for larger testing."""
+  hparams = rlmb_medium()
+  hparams.num_real_env_frames //= 2
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_short():
+  """Small set for larger testing."""
+  hparams = rlmb_base()
+  hparams.num_real_env_frames //= 5
+  hparams.model_train_steps //= 10
+  hparams.ppo_epochs_num //= 10
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_model_only():
+  hp = rlmb_base()
+  hp.epochs = 1
+  hp.ppo_epochs_num = 0
+  return hp
+
+
+@registry.register_hparams
+def rlmb_tiny():
+  """Tiny set for testing."""
+  return rlmb_base_sampling().override_from_dict(
+      tf.contrib.training.HParams(
+          epochs=1,
+          num_real_env_frames=128,
+          simulated_env_generator_num_steps=64,
+          model_train_steps=2,
+          ppo_epochs_num=2,
+          ppo_time_limit=5,
+          ppo_epoch_length=5,
+          ppo_num_agents=2,
+          real_ppo_epoch_length=36,
+          real_ppo_num_agents=1,
+          real_ppo_epochs_num=0,
+          real_ppo_effective_num_agents=2,
+          generative_model_params="next_frame_tiny",
+          stop_loop_early=True,
+          resize_height_factor=2,
+          resize_width_factor=2,
+          game="pong",
+      ).values())
+
+
+@registry.register_hparams
+def rlmb_tiny_stochastic():
+  """Tiny setting with a stochastic next-frame model."""
+  hparams = rlmb_tiny()
+  hparams.epochs = 1  # Too slow with 2 for regular runs.
+  hparams.generative_model = "next_frame_basic_stochastic"
+  hparams.generative_model_params = "next_frame_basic_stochastic"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_tiny_recurrent():
+  """Tiny setting with a recurrent next-frame model."""
+  hparams = rlmb_tiny()
+  hparams.epochs = 1  # Too slow with 2 for regular runs.
+  hparams.generative_model = "next_frame_basic_recurrent"
+  hparams.generative_model_params = "next_frame_basic_recurrent"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_tiny_sv2p():
+  """Tiny setting with a tiny sv2p model."""
+  hparams = rlmb_tiny()
+  hparams.generative_model = "next_frame_sv2p"
+  hparams.generative_model_params = "next_frame_sv2p_tiny"
+  hparams.grayscale = False
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_ae_base():
+  """Parameter set for autoencoders."""
+  hparams = rlmb_base()
+  hparams.ppo_params = "ppo_pong_ae_base"
+  hparams.generative_model_params = "next_frame_ae"
+  hparams.autoencoder_hparams_set = "autoencoder_discrete_pong"
+  hparams.gather_ppo_real_env_data = False
+  hparams.autoencoder_train_steps = 5000
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  hparams.grayscale = False
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_ae_basetest():
+  """Base AE setting but quicker with only 2 epochs."""
+  hparams = rlmb_ae_base()
+  hparams.game = "pong"
+  hparams.epochs = 2
+  hparams.num_real_env_frames = 3200
+  hparams.model_train_steps = 100
+  hparams.autoencoder_train_steps = 10
+  hparams.simulated_env_generator_num_steps = 20
+  hparams.ppo_epochs_num = 2
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_ae_tiny():
+  """Tiny set for testing autoencoders."""
+  hparams = rlmb_tiny()
+  hparams.ppo_params = "ppo_pong_ae_base"
+  hparams.generative_model_params = "next_frame_ae_tiny"
+  hparams.autoencoder_hparams_set = "autoencoder_discrete_tiny"
+  hparams.gather_ppo_real_env_data = False
+  hparams.resize_height_factor = 1
+  hparams.resize_width_factor = 1
+  hparams.grayscale = False
+  hparams.autoencoder_train_steps = 1
+  hparams.autoencoder_train_steps_initial_multiplier = 0
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_tiny_simulation_deterministic_starts():
+  hp = rlmb_tiny()
+  hp.simulation_random_starts = False
+  return hp
+
+
+# RangedHParams for tuning
+# ==============================================================================
+# Note that the items here must be scoped with one of
+# HP_SCOPES={loop, model, ppo}, which set hyperparameters for the top-level
+# hparams, hp.generative_model_params, and hp.ppo_params, respectively.
+@registry.register_ranged_hparams
+def rlmb_grid(rhp):
+  """Grid over games and frames, and 5 runs each for variance."""
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+  base = 100000
+  medium = base // 2
+  small = medium // 2
+  rhp.set_discrete("loop.num_real_env_frames", [base, medium, small])
+
+  # Dummy parameter to get 5 runs for each configuration
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+
+
+@registry.register_ranged_hparams
+def rlmb_variance(rhp):
+  # Dummy parameter to get 5 runs for each configuration
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+
+
+@registry.register_ranged_hparams
+def rlmb_variance_nogame(rhp):
+  # Dummy parameter to get 20 runs for current configuration.
+  rhp.set_discrete("model.moe_loss_coef", list(range(20)))
+
+
+@registry.register_ranged_hparams
+def rlmb_three(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
+
+
+@registry.register_ranged_hparams
+def rlmb_test1(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
+  rhp.set_discrete("loop.ppo_learning_rate", [5e-5, 1e-4, 2e-4])
+  rhp.set_discrete("ppo.optimization_batch_size", [20, 40])
+  rhp.set_discrete("loop.epochs", [3, 6])
+
+
+@registry.register_ranged_hparams
+def rlmb_scheduled_sampling(rhp):
+  rhp.set_float("model.scheduled_sampling_prob", 0.0, 1.0)
+
+
+@registry.register_ranged_hparams
+def rlmb_all_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_whitelisted_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_human_score_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game",
+                      gym_problems_specs.ATARI_GAMES_WITH_HUMAN_SCORE)
+
+
+@registry.register_ranged_hparams
+def rlmb_curious_games10(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_curious_games5(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_debug_games(rhp):
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_DEBUG_GAMES)
+
+
+@registry.register_ranged_hparams
+def rlmb_ae_variance(rhp):
+  # Dummy parameter to get 5 runs for each configuration
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+  base = 100000
+  small = base // 4
+  rhp.set_discrete("loop.num_real_env_frames", [base, small])
+
+
+@registry.register_ranged_hparams
+def rlmb_ppolr_game(rhp):
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+  base_lr = 1e-4
+  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
+
+
+@registry.register_ranged_hparams
+def rlmb_ppolr(rhp):
+  base_lr = 1e-4
+  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
+
+
+@registry.register_ranged_hparams
+def rlmb_ae_ppo_lr(rhp):
+  rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
+  base_lr = 1e-4
+  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
+
+
+@registry.register_ranged_hparams
+def rlmb_dropout_range(rhp):
+  rhp.set_float("model.dropout", 0.2, 0.4)
+
+
+@registry.register_ranged_hparams
+def rlmb_intrinsic_reward_scale(rhp):
+  rhp.set_float("loop.intrinsic_reward_scale", 0.01, 10.)
+
+
+@registry.register_ranged_hparams
+def rlmb_l1l2cutoff_range(rhp):
+  """Loss and loss-cutoff tuning grid."""
+  rhp.set_float("model.video_modality_loss_cutoff", 1.4, 3.4)
+
+
+@registry.register_ranged_hparams
+def rlmb_xentcutoff_range(rhp):
+  """Cross entropy cutoff tuning grid."""
+  rhp.set_float("model.video_modality_loss_cutoff", 0.01, 0.05)
+
+
+@registry.register_ranged_hparams
+def rlmb_pixel_noise(rhp):
+  """Input pixel noise tuning grid."""
+  rhp.set_categorical("loop.generative_model_params",
+                      ["next_frame_pixel_noise"])
+  rhp.set_discrete("model.video_modality_input_noise",
+                   [0.0025 * i for i in range(200)])
+
+
+@registry.register_ranged_hparams
+def rlmb_dummy_range(rhp):
+  """Dummy tuning grid just to get the variance."""
+  rhp.set_float("model.moe_loss_coef", 0.01, 0.02)
+
+
+@registry.register_ranged_hparams
+def rlmb_epochs_num(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.epochs", [3, 6, 12])
+
+
+@registry.register_ranged_hparams
+def rlmb_ppo_epochs_num(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.ppo_epochs_num", [200, 1000, 2000, 4000])
+
+
+@registry.register_ranged_hparams
+def rlmb_ppo_epoch_len(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.ppo_epoch_length", [25, 50, 100])
+
+
+@registry.register_ranged_hparams
+def rlmb_num_frames(rhp):
+  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+  rhp.set_discrete("loop.num_real_env_frames",
+                   [1000*el for el in [30, 100, 500, 1000]])
+
+
+@registry.register_ranged_hparams
+def rlmb_ppo_optimization_batch_size(rhp):
+  rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_discrete("ppo.optimization_batch_size", [4, 10, 20])
+
+
+@registry.register_ranged_hparams
+def rlmb_logits_clip(rhp):
+  rhp.set_categorical("loop.game", ["pong", "boxing", "seaquest"])
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_discrete("ppo.logits_clip", [0., 5.])
+
+
+def merge_unscoped_hparams(scopes_and_hparams):
+  """Merge multiple HParams into one with scopes."""
+  merged_values = {}
+  for (scope, hparams) in scopes_and_hparams:
+    for key, value in six.iteritems(hparams.values()):
+      scoped_key = "%s.%s" % (scope, key)
+      merged_values[scoped_key] = value
+
+  return tf.contrib.training.HParams(**merged_values)
+
+
+def split_scoped_hparams(scopes, merged_hparams):
+  """Split single HParams with scoped keys into multiple."""
+  split_values = {scope: {} for scope in scopes}
+  merged_values = merged_hparams.values()
+  for scoped_key, value in six.iteritems(merged_values):
+    scope = scoped_key.split(".")[0]
+    key = scoped_key[len(scope) + 1:]
+    split_values[scope][key] = value
+
+  return [
+      tf.contrib.training.HParams(**split_values[scope]) for scope in scopes
+  ]
+
+
+def training_loop_hparams_from_scoped_overrides(scoped_overrides, trial_id):
+  """Create HParams suitable for training loop from scoped HParams.
+
+  Args:
+    scoped_overrides: HParams, with keys all scoped by one of HP_SCOPES. These
+      parameters are overrides for the base HParams created by
+      create_loop_hparams.
+    trial_id: str, trial identifier. This is used to register unique HParams
+      names for the underlying model and ppo HParams.
+
+  Returns:
+    HParams suitable for passing to training_loop.
+  """
+  trial_hp_overrides = scoped_overrides.values()
+
+  # Create loop, model, and ppo base HParams
+  loop_hp = create_loop_hparams()
+  model_hp_name = trial_hp_overrides.get(
+      "loop.generative_model_params", loop_hp.generative_model_params)
+  model_hp = registry.hparams(model_hp_name).parse(FLAGS.hparams)
+  ppo_params_name = trial_hp_overrides.get(
+      "loop.ppo_params", loop_hp.ppo_params)
+  ppo_hp = registry.hparams(ppo_params_name)
+
+  # Merge them and then override with the scoped overrides
+  combined_hp = merge_unscoped_hparams(
+      zip(HP_SCOPES, [loop_hp, model_hp, ppo_hp]))
+  combined_hp.override_from_dict(trial_hp_overrides)
+
+  # Split out the component hparams
+  loop_hp, model_hp, ppo_hp = (
+      split_scoped_hparams(HP_SCOPES, combined_hp))
+
+  # Dynamic register the model hp and set the new name in loop_hp
+  model_hp_name = "model_hp_%s" % str(trial_id)
+  dynamic_register_hparams(model_hp_name, model_hp)
+  loop_hp.generative_model_params = model_hp_name
+
+  # Dynamic register the PPO hp and set the new name in loop_hp
+  ppo_hp_name = "ppo_hp_%s" % str(trial_id)
+  dynamic_register_hparams(ppo_hp_name, ppo_hp)
+  loop_hp.ppo_params = ppo_hp_name
+
+  return loop_hp
+
+
+def dynamic_register_hparams(name, hparams):
+
+  @registry.register_hparams(name)
+  def new_hparams_set():
+    return tf.contrib.training.HParams(**hparams.values())
+
+  return new_hparams_set
+
+
+def create_loop_hparams():
+  hparams = registry.hparams(FLAGS.loop_hparams_set)
+  hparams.parse(FLAGS.loop_hparams)
+  return hparams

From 533db5a5b2f80a7c97a9514d1c44d5c0a0bea09a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 12 Oct 2018 13:27:31 -0700
Subject: [PATCH 0993/2720] Improve reward prediction in the basic model, use
 constant learning rate.

PiperOrigin-RevId: 216912214
---
 tensor2tensor/models/video/basic_deterministic.py | 12 +++++++++---
 tensor2tensor/models/video/basic_stochastic.py    |  3 +++
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index ffa2e2d7f..80c8feb7c 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -274,6 +274,7 @@ def body_single(self, features):
           x, action, "action_enc", hparams.action_injection)
 
     x, extra_loss = self.inject_latent(x, features, filters)
+    x_mid = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
 
     # Run a stack of convolutions.
     for i in range(hparams.num_hidden_layers):
@@ -306,16 +307,21 @@ def body_single(self, features):
 
     # Cut down to original size.
     x = x[:, :inputs_shape[1], :inputs_shape[2], :]
+    x_fin = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
     if self.is_per_pixel_softmax:
       x = tf.layers.dense(x, hparams.problem.num_channels * 256, name="logits")
     else:
       x = tf.layers.dense(x, hparams.problem.num_channels, name="logits")
 
-    # Reward prediction if needed.
+    # No reward prediction if not needed.
     if "target_reward" not in features:
       return x, extra_loss
-    reward_pred = tf.expand_dims(  # Add a fake channels dim.
-        tf.reduce_mean(x, axis=[1, 2], keepdims=True), axis=3)
+
+    # Reward prediction based on middle and final logits.
+    reward_pred = tf.concat([x_mid, x_fin], axis=-1)
+    reward_pred = tf.nn.relu(tf.layers.dense(
+        reward_pred, 128, name="reward_pred"))
+    reward_pred = tf.expand_dims(reward_pred, axis=3)  # Need fake channels dim.
     return {"targets": x, "target_reward": reward_pred}, extra_loss
 
   def body(self, features):
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 03d8d803e..84d7e450b 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -229,6 +229,9 @@ def next_frame_basic_stochastic_discrete():
   hparams.scheduled_sampling_mode = "prob_inverse_lin"
   hparams.scheduled_sampling_decay_steps = 40000
   hparams.scheduled_sampling_prob = 1.0
+  hparams.learning_rate_constant = 0.01
+  hparams.learning_rate_warmup_steps = 8000
+  hparams.learning_rate_schedule = "linear_warmup * constant"
   hparams.add_hparam("bottleneck_bits", 64)
   hparams.add_hparam("bottleneck_noise", 0.02)
   hparams.add_hparam("discrete_warmup_steps", 40000)

From c8d7157e45197f777e1b0296fcd51aa575150b0f Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 12 Oct 2018 14:41:35 -0700
Subject: [PATCH 0994/2720] shared_rel is missing in hparams while running
 model=img2img_transformer and problem=img2img_celeba with
 hparams_set=img2img_transformer_base_tpu

PiperOrigin-RevId: 216924548
---
 tensor2tensor/models/image_transformer_2d.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 5b2299a00..1c105186f 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -423,6 +423,9 @@ def image_transformer2d_base():
   hparams.add_hparam("kv_filter_width", 1)
 
   hparams.add_hparam("unconditional", False)  # unconditional generation
+
+  # relative embedding hparams
+  hparams.add_hparam("shared_rel", False)
   return hparams
 
 

From ee89d7bb18bdb3c86b70a59950e7ca827bcf8fcf Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Fri, 12 Oct 2018 15:49:10 -0700
Subject: [PATCH 0995/2720] Get bf16 activations working with
 mixture-of-experts.

PiperOrigin-RevId: 216935230
---
 tensor2tensor/mesh_tensorflow/research/moe.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensor2tensor/mesh_tensorflow/research/moe.py b/tensor2tensor/mesh_tensorflow/research/moe.py
index fac763cb3..5761292ea 100644
--- a/tensor2tensor/mesh_tensorflow/research/moe.py
+++ b/tensor2tensor/mesh_tensorflow/research/moe.py
@@ -475,6 +475,10 @@ def _top_2_gating(
       inputs, experts_dim, use_bias=False,
       expert_dims=outer_expert_dims), experts_dim)
 
+  # The internals of this function run in float32.
+  #   bfloat16 seems to reduce quality.
+  raw_gates = mtf.to_float(raw_gates)
+
   expert_capacity_f = float(expert_capacity_dim.size)
 
   # FIND TOP 2 EXPERTS PER POSITON
@@ -591,6 +595,9 @@ def _top_2_gating(
       * mtf.one_hot(index_2, experts_dim)
       * mtf.one_hot(mtf.to_int32(position_in_expert_2), expert_capacity_dim))
 
+  combine_tensor = mtf.cast(combine_tensor, inputs.dtype)
+  loss = mtf.cast(loss, inputs.dtype)
+
   dispatch_tensor = mtf.cast(
       mtf.cast(combine_tensor, tf.bool), combine_tensor.dtype)
 

From 2e89303f48c8d3edc2718813fef39fd6300280fc Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Fri, 12 Oct 2018 15:55:12 -0700
Subject: [PATCH 0996/2720] Add a check asserting that arguments to binary
 operations have the same dtype.

PiperOrigin-RevId: 216936071
---
 tensor2tensor/mesh_tensorflow/mesh_tensorflow.py | 8 ++++----
 tensor2tensor/mesh_tensorflow/simd_mesh_impl.py  | 2 --
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
index da75330b8..707edc554 100644
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
@@ -1442,6 +1442,10 @@ class BinaryOpWithBroadcasting(Operation):
   def __init__(self, tf_fn, x1, x2, output_shape, output_dtype, name=None):
     super(BinaryOpWithBroadcasting, self).__init__(
         [x1, x2], name=name or "binary_op")
+    if x1.dtype != x2.dtype:
+      # If there is ever a binary operation with different operand types, then
+      # we should add an argument allow_different_operand_dtypes=False.
+      raise ValueError("Dtypes must be equal.")
     assert isinstance(output_dtype, tf.DType)
     self._outputs = [Tensor(self, output_shape, output_dtype)]
     self._tf_fn = tf_fn
@@ -1563,8 +1567,6 @@ class AddOperation(BinaryOpWithBroadcasting):
   def __init__(self, x1, x2, output_shape, name=None):
     super(AddOperation, self).__init__(
         tf.add, x1, x2, output_shape, x1.dtype, name=name or "add")
-    if x1.dtype != x2.dtype:
-      raise ValueError("Dtypes must be equal.")
 
   def gradient(self, grad_ys):
     dy = grad_ys[0]
@@ -1578,8 +1580,6 @@ class MinMaxOperation(BinaryOpWithBroadcasting):
   def __init__(self, tf_fn, x1, x2, output_shape, name=None):
     super(MinMaxOperation, self).__init__(
         tf_fn, x1, x2, output_shape, x1.dtype, name=name or "add")
-    if x1.dtype != x2.dtype:
-      raise ValueError("Dtypes must be equal.")
 
   def gradient(self, grad_ys):
     dy = grad_ys[0]
diff --git a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
index bebbff22d..1eb3fd233 100644
--- a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
+++ b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
@@ -249,7 +249,6 @@ def alltoall(self, x, mesh_axis, split_axis, concat_axis):
     x = x.to_laid_out_tensor()
     t = x.one_slice
     group_assignment = self._create_group_assignment([mesh_axis])
-
     t = tpu_ops.all_to_all(
         t,
         concat_dimension=concat_axis,
@@ -257,7 +256,6 @@ def alltoall(self, x, mesh_axis, split_axis, concat_axis):
         split_count=len(group_assignment[0]),
         group_assignment=group_assignment)
     x = self.LaidOutTensor([t])
-
     return x
 
   def receive(self, x, mesh_axis, source_pcoord):

From 628bf1fbbb5d4195b824c0d3c52ebf2c32a4e043 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 12 Oct 2018 16:02:18 -0700
Subject: [PATCH 0997/2720] Fix multi_problem_test.py.

PiperOrigin-RevId: 216937022
---
 .../data_generators/multi_problem_test.py     | 43 ++++++++++++++++---
 1 file changed, 36 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem_test.py b/tensor2tensor/data_generators/multi_problem_test.py
index 7b0708a16..c81483568 100644
--- a/tensor2tensor/data_generators/multi_problem_test.py
+++ b/tensor2tensor/data_generators/multi_problem_test.py
@@ -20,32 +20,61 @@
 from __future__ import print_function
 
 from tensor2tensor.data_generators import multi_problem
-from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.layers import modalities
 
 import tensorflow as tf
 
 
+# TODO(trandustin): This test problem is required in order for MultiProblem
+# to access vocab size via encoders. In a future change, enable MultiProblem to
+# access vocab size more explicitly from the Problem.
+class TestProblem(problem.Problem):
+  """Test problem."""
+
+  def __init__(self, input_vocab_size, target_vocab_size):
+    super(TestProblem, self).__init__(False, False)
+    self.input_vocab_size = input_vocab_size
+    self.target_vocab_size = target_vocab_size
+
+  def hparams(self, defaults, model_hparams):
+    hp = defaults
+    hp.input_modality = {
+        "inputs": modalities.SymbolModality(model_hparams,
+                                            self.input_vocab_size)
+    }
+    hp.target_modality = modalities.SymbolModality(model_hparams,
+                                                   self.target_vocab_size)
+
+  def feature_encoders(self, data_dir):
+    encoders = {
+        "inputs": text_encoder.ByteTextEncoder(),
+        "targets": text_encoder.ByteTextEncoder(),
+    }
+    return encoders
+
+
 class TestMultiProblem(multi_problem.MultiProblem):
   """Test multi-problem."""
 
   def __init__(self):
     super(TestMultiProblem, self).__init__()
-    self.task_list.append(problem_hparams.TestProblem(2, 3))
-    self.task_list.append(problem_hparams.TestProblem(4, 6))
+    self.task_list.append(TestProblem(2, 3))
+    self.task_list.append(TestProblem(4, 6))
 
 
 class MultiProblemTest(tf.test.TestCase):
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testProblemHparamsModality(self):
-    problem = TestMultiProblem()
-    p_hparams = problem.get_hparams()
+    multiproblem = TestMultiProblem()
+    p_hparams = multiproblem.get_hparams()
     self.assertIsInstance(p_hparams.input_modality["inputs"],
                           modalities.SymbolModality)
-    self.assertEqual(p_hparams.input_modality["inputs"].top_dimensionality, 3)
+    self.assertEqual(p_hparams.input_modality["inputs"].top_dimensionality, 2)
     self.assertIsInstance(p_hparams.target_modality, modalities.SymbolModality)
-    self.assertEqual(p_hparams.target_modality.top_dimensionality, 5)
+    self.assertEqual(p_hparams.target_modality.top_dimensionality, 260)
 
 if __name__ == "__main__":
   tf.test.main()

From 6f1ba1d2e32a749e3d65193e9d04d76ac4b63604 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 12 Oct 2018 16:59:50 -0700
Subject: [PATCH 0998/2720] EPVA hparams inheriting from SV2P hparams.

PiperOrigin-RevId: 216944811
---
 tensor2tensor/models/video/epva_params.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
index 2281af593..780d34cc2 100644
--- a/tensor2tensor/models/video/epva_params.py
+++ b/tensor2tensor/models/video/epva_params.py
@@ -33,6 +33,8 @@ def next_frame_epva():
   hparams.learning_rate_constant = 1e-05
   hparams.batch_size = 2
   hparams.clip_grad_norm = 0.01
+  # TODO(msaffar): disentangle EPVA from SV2P
+  hparams.add_hparam("reward_prediction", False)
   hparams.add_hparam("context_frames", 5)
   hparams.add_hparam("enc_learning_rate", 1e-5)
   hparams.add_hparam("enc_pred_loss_scale", 0.1)

From e470049b8fcff77684a275f99b92e0f14f2e5dd7 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Sat, 13 Oct 2018 04:46:30 +0200
Subject: [PATCH 0999/2720] Implement new RL pipeline (#1140)

* copy trainer_model_based.py to trainer_model_based_T2TEnv.py

* Skeleton of new training loop (using T2TGymEnv as problem)

* trainer_model_based_T2TEnv -> trainer_model_based_new

* Get trainer_model_based_new to run

* Add supervised_trainer

* Implementation of env setup, partial train_agent_real_env

* Cast observation after resize

* T2TEnv saves unfinished rollouts on reset

* Fix a bug in calculating shard size

* Do not copy graph, session ets when deepcopying T2TEnv

* Add epoch number for rollouts in T2TEnv

* Evaluation of unclipped rewards in Model-Based RL pipeline

* World model training

* Change sharding to per-frame so there's enough data in dev set for tiny

* Fix evaluation

* Eval clipped rewards for real env

* Pylint

* Add a test for trainer_model_based_new

* Move supervised_trainer to trainer_model_based_new
---
 tensor2tensor/data_generators/gym_env.py      | 101 +++--
 tensor2tensor/data_generators/gym_env_test.py |   1 +
 tensor2tensor/models/research/rl.py           |  44 +-
 tensor2tensor/rl/envs/batch_env_factory.py    |  13 +-
 tensor2tensor/rl/envs/simulated_batch_env.py  |   6 +-
 tensor2tensor/rl/envs/utils.py                |   9 +-
 tensor2tensor/rl/trainer_model_based_new.py   | 412 ++++++++++++++++++
 .../rl/trainer_model_based_new_test.py        |  38 ++
 8 files changed, 574 insertions(+), 50 deletions(-)
 create mode 100644 tensor2tensor/rl/trainer_model_based_new.py
 create mode 100644 tensor2tensor/rl/trainer_model_based_new_test.py

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 6b31653b9..3a44a50c5 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -20,7 +20,6 @@
 from __future__ import print_function
 
 import collections
-import math
 import random
 
 from gym.spaces import Box
@@ -40,6 +39,15 @@
 )
 
 
+class _Noncopyable(object):
+
+  def __init__(self, obj):
+    self.obj = obj
+
+  def __deepcopy__(self, memo):
+    return self
+
+
 class T2TEnv(video_utils.VideoProblem):
   """Abstract class representing a batch of environments.
 
@@ -72,10 +80,14 @@ def __init__(self, batch_size):
     self._current_frames = [None for _ in range(batch_size)]
 
     with tf.Graph().as_default() as tf_graph:
-      self._tf_graph = tf_graph
-      self._image_t = tf.placeholder(dtype=tf.uint8, shape=(None, None, None))
-      self._encoded_image_t = tf.image.encode_png(self._image_t)
-      self._session = tf.Session()
+      self._tf_graph = _Noncopyable(tf_graph)
+      self._image_t = _Noncopyable(
+          tf.placeholder(dtype=tf.uint8, shape=(None, None, None))
+      )
+      self._encoded_image_t = _Noncopyable(
+          tf.image.encode_png(self._image_t.obj)
+      )
+      self._session = _Noncopyable(tf.Session())
 
   def __str__(self):
     """Returns a string representation of the environment for debug purposes."""
@@ -83,7 +95,19 @@ def __str__(self):
 
   def clear_history(self):
     """Clears the rollout history."""
-    self.history = []
+    self.rollouts_by_epoch = dict()
+
+  def start_new_epoch(self, epoch):
+    if not isinstance(epoch, int):
+      raise ValueError('Epoch should be integer, got {}'.format(epoch))
+    if epoch in self.rollouts_by_epoch:
+      raise ValueError('Epoch {} already registered'.format(epoch))
+    self.current_epoch = epoch
+    self.rollouts_by_epoch[epoch] = list()
+
+  @property
+  def current_epoch_rollouts(self):
+    return self.rollouts_by_epoch[self.current_epoch]
 
   def _preprocess_observations(self, obs):
     """Transforms a batch of observations.
@@ -101,8 +125,9 @@ def _preprocess_observations(self, obs):
   def _encode_observations(self, observations):
     """Encodes observations as PNG."""
     return [
-        self._session.run(
-            self._encoded_image_t, feed_dict={self._image_t: observation}
+        self._session.obj.run(
+            self._encoded_image_t.obj,
+            feed_dict={self._image_t.obj: observation}
         )
         for observation in observations
     ]
@@ -169,7 +194,7 @@ def _reset(self, indices):
   def reset(self, indices=None):
     """Resets environments at given indices.
 
-    Does any preprocessing and adds finished rollouts to history.
+    Does any preprocessing and adds rollouts to history.
 
     Args:
       indices: Indices of environments to reset.
@@ -184,10 +209,10 @@ def reset(self, indices=None):
     encoded_obs = self._encode_observations(new_obs)
     for (index, ob) in zip(indices, encoded_obs):
       frame = self._current_frames[index]
-      if frame is not None and frame.done:
+      if frame is not None:
         rollout = self._current_rollouts[index]
         rollout.append(frame._replace(action=0))
-        self.history.append(rollout)
+        self.current_epoch_rollouts.append(rollout)
         self._current_rollouts[index] = []
       self._current_frames[index] = Frame(
           observation=ob, reward=0, unclipped_reward=0, done=False, action=None
@@ -199,7 +224,7 @@ def close(self):
 
     Can be overridden in derived classes.
     """
-    self._session.close()
+    self._session.obj.close()
 
   @property
   def num_channels(self):
@@ -234,6 +259,10 @@ def frame_height(self):
   def frame_width(self):
     return self.observation_space.shape[1]
 
+  @property
+  def only_keep_videos_from_0th_frame(self):
+    return False
+
   @property
   def num_actions(self):
     return self.action_space.n
@@ -256,11 +285,12 @@ def make_modality(name):
     p.input_space_id = problem.SpaceID.IMAGE
     p.target_space_id = problem.SpaceID.IMAGE
 
-  def _generate_frames(self, rollouts):
-    for rollout in rollouts:
+  def _generate_frames(self, epoch_rollout_tuples):
+    for epoch, rollout in epoch_rollout_tuples:
       for (frame_number, frame) in enumerate(rollout):
         yield {
             "frame_number": [frame_number],
+            "epoch": [epoch],
             "image/encoded": [frame.observation],
             "image/format": ["png"],
             "image/height": [self.frame_height],
@@ -272,10 +302,14 @@ def _generate_frames(self, rollouts):
 
   def generate_data(self, data_dir, tmp_dir, task_id=-1):
     """Saves the rollout history to disk."""
-    # Suffle rollouts globally taking advantage of the fact that we have
+    # Shuffle rollouts globally taking advantage of the fact that we have
     # everything in memory.
-    shuffled_history = self.history[:]
-    random.shuffle(shuffled_history)
+    epoch_rollout_tuples = list()
+    for epoch_nr, rollouts in self.rollouts_by_epoch.items():
+      for rollout in rollouts:
+        epoch_rollout_tuples.append((epoch_nr, rollout))
+
+    random.shuffle(epoch_rollout_tuples)
 
     filepath_fns = {
         problem.DatasetSplit.TRAIN: self.training_filepaths,
@@ -284,22 +318,20 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
     }
 
     # We set shuffled=True as we don't want to shuffle on disk later.
-    splits_and_paths = [
-        (split["split"], path)
+    paths = [
+        path
         for split in self.dataset_splits
         for path in filepath_fns[split["split"]](
             data_dir, split["shards"], shuffled=True
         )
     ]
 
-    # Split entire rollouts into shards so that no rollout is broken on shard
-    # boundary.
-    shard_size = int(math.ceil(len(shuffled_history)) / len(splits_and_paths))
-    for (i, (split, path)) in enumerate(splits_and_paths):
-      rollouts = shuffled_history[i * shard_size : (i + 1) * shard_size]
-      generator_utils.generate_files(
-          self._generate_frames(rollouts), [path], cycle_every_n=float("inf")
-      )
+    num_frames = sum(len(rollout) for (_, rollout) in epoch_rollout_tuples)
+    shard_size = num_frames // len(paths)
+    generator_utils.generate_files(
+        self._generate_frames(epoch_rollout_tuples), paths,
+        cycle_every_n=shard_size
+    )
 
 
 class T2TGymEnv(T2TEnv):
@@ -328,18 +360,19 @@ def __init__(self, envs, grayscale=False,
     if not all(env.action_space == self.action_space for env in self._envs):
       raise ValueError("All environments must use the same action space.")
 
-    with self._tf_graph.as_default():
+    with self._tf_graph.obj.as_default():
       self._resize = dict()
       orig_height, orig_width = orig_observ_space.shape[:2]
-      self._img_batch_t = tf.placeholder(
-          dtype=tf.uint8, shape=(None, orig_height, orig_width, 3))
+      self._img_batch_t = _Noncopyable(tf.placeholder(
+          dtype=tf.uint8, shape=(None, orig_height, orig_width, 3)))
       height, width = self.observation_space.shape[:2]
-      resized = tf.image.resize_images(self._img_batch_t,
+      resized = tf.image.resize_images(self._img_batch_t.obj,
                                        [height, width],
                                        tf.image.ResizeMethod.AREA)
+      resized = tf.cast(resized, tf.as_dtype(self.observation_space.dtype))
       if self.grayscale:
         resized = tf.image.rgb_to_grayscale(resized)
-      self._resized_img_batch_t = resized
+      self._resized_img_batch_t = _Noncopyable(resized)
 
   @property
   def num_channels(self):
@@ -360,8 +393,8 @@ def __str__(self):
     return "T2TGymEnv(%s)" % ", ".join([str(env) for env in self._envs])
 
   def _preprocess_observations(self, obs):
-    return self._session.run(self._resized_img_batch_t,
-                             feed_dict={self._img_batch_t: obs})
+    return self._session.obj.run(self._resized_img_batch_t.obj,
+                                 feed_dict={self._img_batch_t.obj: obs})
 
   def _step(self, actions):
     (obs, rewards, dones, _) = zip(*[
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 0317d6261..772a6baa2 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -72,6 +72,7 @@ def setUpClass(cls):
   def init_batch_and_play(self, env_lambda, n_steps=1, **kwargs):
     raw_envs = [env_lambda(), env_lambda()]
     env = gym_env.T2TGymEnv(raw_envs, **kwargs)
+    env.start_new_epoch(0)
     obs = list()
     rewards = list()
     obs.append(env.reset())
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 0cdbd5569..573d38c66 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -144,8 +144,8 @@ def simple_gym_spec(env):
 
 
 def standard_atari_env_spec(
-    env, simulated=False, resize_height_factor=1, resize_width_factor=1,
-    grayscale=False, include_clipping=True):
+    env=None, simulated=False, resize_height_factor=1, resize_width_factor=1,
+    grayscale=False, include_clipping=True, batch_env=None):
   """Parameters of environment specification."""
   resize_wrapper = [tf_atari_wrappers.ResizeWrapper,
                     {"height_factor": resize_height_factor,
@@ -164,18 +164,42 @@ def standard_atari_env_spec(
     ]
   if simulated:  # No resizing on simulated environments.
     standard_wrappers = standard_wrappers[1:]
-  env_lambda = None
-  if isinstance(env, str):
-    env_lambda = lambda: gym.make(env)
-  if callable(env):
-    env_lambda = env
-  assert env_lambda is not None, "Unknown specification of environment"
 
-  return tf.contrib.training.HParams(
-      env_lambda=env_lambda,
+  env_spec = tf.contrib.training.HParams(
       wrappers=standard_wrappers,
       simulated_env=simulated)
 
+  if batch_env is not None:
+    env_spec.add_hparam("batch_env", batch_env)
+  else:
+    env_lambda = None
+    if isinstance(env, str):
+      env_lambda = lambda: gym.make(env)
+    if callable(env):
+      env_lambda = env
+    assert env_lambda is not None, "Unknown specification of environment"
+    env_spec.add_hparam("env_lambda", env_lambda)
+
+  return env_spec
+
+
+def standard_atari_env_simulated_spec(
+    real_env, video_num_input_frames, video_num_target_frames
+):
+  env_spec = standard_atari_env_spec(
+      # This hack is here because SimulatedBatchEnv needs to get
+      # observation_space from the real env. TODO(koz4k): refactor.
+      env=lambda: real_env,
+      simulated=True
+  )
+  env_spec.add_hparam("simulation_random_starts", True)
+  env_spec.add_hparam("simulation_flip_first_random_for_beginning", True)
+  env_spec.add_hparam("intrinsic_reward_scale", 0.0)
+  env_spec.add_hparam("initial_frames_problem", real_env)
+  env_spec.add_hparam("video_num_input_frames", video_num_input_frames)
+  env_spec.add_hparam("video_num_target_frames", video_num_target_frames)
+  return env_spec
+
 
 def standard_atari_env_eval_spec(
     env, simulated=False, resize_height_factor=1, resize_width_factor=1,
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index d6d9cce75..570f276bd 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -32,12 +32,20 @@
 
 def batch_env_factory(environment_spec, num_agents, initial_frame_chooser=None):
   """Factory of batch envs."""
-
+  # TODO(konradczechowski): this is temporary function handling both old and
+  # new pipelines, refactor this when we move to the new pipeline.
   if environment_spec.simulated_env:
     cur_batch_env = _define_simulated_batch_env(
         environment_spec, num_agents, initial_frame_chooser)
   else:
-    cur_batch_env = _define_batch_env(environment_spec, num_agents)
+    if 'batch_env' in environment_spec:
+      assert not 'env_lambda' in environment_spec, \
+          'Environment_spec should contain only one of (env_lambda, batch_env).'
+      batch_env = environment_spec.batch_env
+      assert batch_env.batch_size == num_agents
+    else:
+      batch_env = _define_batch_env(environment_spec, num_agents)
+    cur_batch_env = py_func_batch_env.PyFuncBatchEnv(batch_env)
   return cur_batch_env
 
 
@@ -49,7 +57,6 @@ def _define_batch_env(environment_spec, num_agents):
         environment_spec.env_lambda()
         for _ in range(num_agents)]
     env = gym_env.T2TGymEnv(envs)
-    env = py_func_batch_env.PyFuncBatchEnv(env)
     return env
 
 
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index f48d66e64..5b245d1b0 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -113,10 +113,14 @@ def __init__(self, environment_spec, length, initial_frame_chooser):
     super(SimulatedBatchEnv, self).__init__(observ_space, action_space)
 
     self.length = length
-    self._min_reward = initial_frames_problem.min_reward
+    try:
+      self._min_reward = initial_frames_problem.min_reward
+    except AttributeError:
+      self._min_reward = initial_frames_problem.reward_range[0]
     self._num_frames = environment_spec.video_num_input_frames
     self._intrinsic_reward_scale = environment_spec.intrinsic_reward_scale
 
+    # TODO(koz4k): Pass by argument.
     model_hparams = trainer_lib.create_hparams(
         FLAGS.hparams_set, problem_name=FLAGS.problem)
     model_hparams.force_full_predict = True
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index 6cf1df270..ee30df3cb 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -83,12 +83,17 @@ def get_action_space(environment_spec):
   """Get action space associated with environment spec.
 
   Args:
-     environment_spec:  EnvironmentSpec object
+     environment_spec:  Object consisting one of batch_env.action_space, or
+     env_lambda().action_space
 
   Returns:
     OpenAi Gym action space
   """
-  return environment_spec.env_lambda().action_space
+  if "batch_env" in environment_spec:
+    action_space = environment_spec.batch_env.action_space
+  else:
+    action_space = environment_spec.env_lambda().action_space
+  return action_space
 
 
 def get_policy(observations, hparams):
diff --git a/tensor2tensor/rl/trainer_model_based_new.py b/tensor2tensor/rl/trainer_model_based_new.py
new file mode 100644
index 000000000..19f23b36d
--- /dev/null
+++ b/tensor2tensor/rl/trainer_model_based_new.py
@@ -0,0 +1,412 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+r"""Training of model-based RL agents.
+
+Example invocation:
+
+python -m tensor2tensor.rl.trainer_model_based_new \
+    --output_dir=$HOME/t2t/rl_v1 \
+    --loop_hparams_set=rlmb_base \
+    --loop_hparams='num_real_env_frames=10000,epochs=3'
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import contextlib
+import copy
+import datetime
+import math
+import os
+import time
+
+import gym
+import numpy as np
+
+from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
+from tensor2tensor.data_generators.gym_env import T2TGymEnv
+from tensor2tensor.models.research import rl
+from tensor2tensor.rl import rl_trainer_lib, trainer_model_based_params
+from tensor2tensor.rl.envs.utils import InitialFrameChooser
+from tensor2tensor.utils import trainer_lib
+
+import tensorflow as tf
+
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("loop_hparams_set", "rlmb_base",
+                    "Which RL hparams set to use.")
+flags.DEFINE_string("loop_hparams", "", "Overrides for overall loop HParams.")
+flags.DEFINE_string("job_dir_to_evaluate", "",
+                    "Directory of a job to be evaluated.")
+flags.DEFINE_string("eval_results_dir", "/tmp",
+                    "Directory to store result of evaluation")
+
+
+@contextlib.contextmanager
+def temporary_flags(flag_settings):
+  old_values = {}
+  for flag_name, flag_value in flag_settings.items():
+    old_values[flag_name] = getattr(FLAGS, flag_name)
+    setattr(FLAGS, flag_name, flag_value)
+  yield
+  for flag_name, flag_value in old_values.items():
+    setattr(FLAGS, flag_name, flag_value)
+
+
+def _ppo_training_epochs(hparams, epoch, is_final_epoch, real_env_training):
+  """Helper for PPO restarts."""
+  if hparams.gather_ppo_real_env_data:
+    assert hparams.real_ppo_epochs_num is 0, (
+        "Should be put to 0 to enforce better readability")
+    real_training_ppo_epochs_num = int(math.ceil(
+        hparams.num_real_env_frames /
+        (hparams.epochs*hparams.real_ppo_epoch_length)))
+  else:
+    real_training_ppo_epochs_num = hparams.real_ppo_epochs_num
+
+  simulated_training_ppo_epochs_num = hparams.ppo_epochs_num
+
+  if epoch == -1:
+    assert real_env_training, (
+        "Epoch -1 should only be used for PPO collection in real environment.")
+    return real_training_ppo_epochs_num
+  ppo_training_epochs = (epoch + 1) * (simulated_training_ppo_epochs_num
+                                       + real_training_ppo_epochs_num)
+  if is_final_epoch:  # Length of training in the final epoch is doubled.
+    ppo_training_epochs += simulated_training_ppo_epochs_num
+  if real_env_training:
+    ppo_training_epochs += real_training_ppo_epochs_num
+  return ppo_training_epochs
+
+
+def setup_directories(base_dir, subdirs):
+  base_dir = os.path.expanduser(base_dir)
+  tf.gfile.MakeDirs(base_dir)
+
+  all_dirs = {}
+  for subdir in subdirs:
+    dir_name = os.path.join(base_dir, subdir)
+    tf.gfile.MakeDirs(dir_name)
+    all_dirs[subdir] = dir_name
+  return all_dirs
+
+
+def make_relative_timing_fn():
+  """Make a function that logs the duration since it was made."""
+  start_time = time.time()
+
+  def format_relative_time():
+    time_delta = time.time() - start_time
+    return str(datetime.timedelta(seconds=time_delta))
+
+  def log_relative_time():
+    tf.logging.info("Timing: %s", format_relative_time())
+
+  return log_relative_time
+
+
+def make_log_fn(epoch, log_relative_time_fn):
+
+  def log(msg, *args):
+    msg %= args
+    tf.logging.info("%s Epoch %d: %s", ">>>>>>>", epoch, msg)
+    log_relative_time_fn()
+
+  return log
+
+
+def train_supervised(problem, model_name, hparams, data_dir, output_dir,
+                     train_steps, eval_steps, local_eval_frequency=None):
+  if local_eval_frequency is None:
+    local_eval_frequency = getattr(FLAGS, "local_eval_frequency")
+
+  exp_fn = trainer_lib.create_experiment_fn(
+      model_name, problem, data_dir, train_steps, eval_steps,
+      min_eval_frequency=local_eval_frequency
+  )
+  run_config = trainer_lib.create_run_config(model_dir=output_dir)
+  exp = exp_fn(run_config, hparams)
+  exp.test()
+
+
+def train_agent(environment_spec, agent_model_dir,
+                event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
+                is_final_epoch=False):
+  """Train the PPO agent in the simulated environment."""
+  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
+  ppo_params_names = ["epochs_num", "epoch_length",
+                      "learning_rate", "num_agents",
+                      "optimization_epochs", "eval_every_epochs"]
+
+  for param_name in ppo_params_names:
+    ppo_param_name = "ppo_" + param_name
+    if ppo_param_name in hparams:
+      ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
+
+  ppo_hparams.epochs_num = _ppo_training_epochs(hparams, epoch,
+                                                is_final_epoch, False)
+  ppo_hparams.save_models_every_epochs = 10
+  ppo_hparams.world_model_dir = world_model_dir
+  ppo_hparams.add_hparam("force_beginning_resets", True)
+
+  # Adding model hparams for model specific adjustments
+  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
+  ppo_hparams.add_hparam("model_hparams", model_hparams)
+
+  environment_spec = copy.copy(environment_spec)
+  environment_spec_param_names = [
+      "simulation_random_starts", "simulation_flip_first_random_for_beginning",
+      "intrinsic_reward_scale"
+  ]
+  for param_name in environment_spec_param_names:
+    environment_spec.set_hparam(param_name, hparams.get(param_name))
+  ppo_hparams.add_hparam("environment_spec", environment_spec)
+
+  ppo_hparams.add_hparam("initial_frame_chooser", InitialFrameChooser(
+      environment_spec, mode=tf.estimator.ModeKeys.EVAL
+  ))
+
+  # TODO(koz4k): Pass by arguments.
+  with temporary_flags({
+      "problem": environment_spec.initial_frames_problem,
+      "model": hparams.generative_model,
+      "hparams_set": hparams.generative_model_params,
+      "output_dir": world_model_dir,
+      "data_dir": epoch_data_dir,
+  }):
+    rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
+                         name_scope="ppo_sim%d" % (epoch + 1))
+
+
+def train_agent_real_env(
+    env, agent_model_dir, event_dir, epoch_data_dir,
+    hparams, epoch=0, is_final_epoch=False):
+  """Train the PPO agent in the real environment."""
+  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
+  ppo_params_names = ["epochs_num", "epoch_length",
+                      "learning_rate", "num_agents", "eval_every_epochs",
+                      "optimization_epochs", "effective_num_agents"]
+
+  # This should be overridden.
+  ppo_hparams.add_hparam("effective_num_agents", None)
+  for param_name in ppo_params_names:
+    ppo_param_name = "real_ppo_"+ param_name
+    if ppo_param_name in hparams:
+      ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
+
+  ppo_hparams.epochs_num = _ppo_training_epochs(hparams, epoch,
+                                                is_final_epoch, True)
+  # We do not save model, as that resets frames that we need at restarts.
+  # But we need to save at the last step, so we set it very high.
+  ppo_hparams.save_models_every_epochs = 1000000
+
+  environment_spec = rl.standard_atari_env_spec(
+      batch_env=env, include_clipping=False
+  )
+
+  ppo_hparams.add_hparam("environment_spec", environment_spec)
+
+  rl_trainer_lib.train(ppo_hparams, event_dir + "real", agent_model_dir,
+                       name_scope="ppo_real%d" % (epoch + 1))
+
+  # Save unfinished rollouts to history.
+  env.reset()
+
+
+def train_world_model(env, data_dir, output_dir, hparams, epoch):
+  """Train the world model on problem_name."""
+  train_steps = hparams.model_train_steps * (
+      epoch + hparams.inital_epoch_train_steps_multiplier)
+  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
+
+  # Hardcoded for now. TODO(koz4k): Make it a hparam.
+  model_hparams.video_num_input_frames = 4
+  model_hparams.video_num_target_frames = 1
+
+  model_hparams.learning_rate = model_hparams.learning_rate_constant
+  if epoch > 0:
+    model_hparams.learning_rate *= hparams.learning_rate_bump
+
+  train_supervised(
+      problem=env,
+      model_name=hparams.generative_model,
+      hparams=model_hparams,
+      data_dir=data_dir,
+      output_dir=output_dir,
+      train_steps=train_steps,
+      eval_steps=100,
+      local_eval_frequency=2000
+  )
+
+
+def setup_env(hparams):
+  # TODO(kc): set reward clipping, when this will be possible
+  assert hparams.game == "pong", "Currently only games with [-1, 1] reward " \
+                                 "range are working"
+  game_mode = "Deterministic-v4"
+  camel_game_name = "".join(
+      [w[0].upper() + w[1:] for w in hparams.game.split("_")])
+  camel_game_name += game_mode
+  env_name = camel_game_name
+  env = T2TGymEnv([gym.make(env_name)],
+                  grayscale=hparams.grayscale,
+                  resize_width_factor=hparams.resize_width_factor,
+                  resize_height_factor=hparams.resize_height_factor)
+  return env
+
+
+def eval_reward(env, epoch, clipped):
+  """Calculates mean rewards from given epoch."""
+  reward_name = 'reward' if clipped else 'unclipped_reward'
+  rewards = []
+  for rollout in env.rollouts_by_epoch[epoch]:
+    if rollout[-1].done:
+      rollout_reward = sum(frame[reward_name] for frame in rollout)
+      rewards.append(rollout_reward)
+  if rewards:
+    mean_rewards = np.mean(rewards)
+  else:
+    mean_rewards = 0
+  return mean_rewards
+
+
+def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
+  """Run the main training loop."""
+  if report_fn:
+    assert report_metric is not None
+
+  # Directories
+  subdirectories = ["data", "tmp", "world_model", "ppo"]
+  directories = setup_directories(output_dir, subdirectories)
+
+  epoch = -1
+  env = setup_env(hparams)
+  env.start_new_epoch(epoch)
+
+  # Timing log function
+  log_relative_time = make_relative_timing_fn()
+
+  # Per-epoch state
+  epoch_metrics = []
+  epoch_data_dirs = []
+
+  data_dir = os.path.join(directories["data"], "initial")
+  epoch_data_dirs.append(data_dir)
+  # Collect data from the real environment with PPO or random policy.
+  # TODO: do we need option not to gather_ppo_real_env_data?
+  # We could set learning_rate=0 if this flag == False.
+  assert hparams.gather_ppo_real_env_data
+  ppo_model_dir = directories["ppo"]
+  tf.logging.info("Initial training of PPO in real environment.")
+  ppo_event_dir = os.path.join(directories["world_model"],
+                               "ppo_summaries/initial")
+  train_agent_real_env(
+      env, ppo_model_dir,
+      ppo_event_dir, data_dir,
+      hparams, epoch=epoch, is_final_epoch=False)
+  mean_unclipped_reward = eval_reward(env, epoch, clipped=False)
+  tf.logging.info("Mean reward (initial): {}".format(mean_unclipped_reward))
+
+  eval_metrics_event_dir = os.path.join(directories["world_model"],
+                                        "eval_metrics_event_dir")
+  eval_metrics_writer = tf.summary.FileWriter(eval_metrics_event_dir)
+
+  mean_unclipped_reward_summary = tf.Summary()
+  mean_unclipped_reward_summary.value.add(tag="mean_unclipped_reward",
+                                          simple_value=None)
+  mean_clipped_reward_summary = tf.Summary()
+  mean_clipped_reward_summary.value.add(tag="mean_clipped_reward",
+                                        simple_value=None)
+
+  sim_env_spec = rl.standard_atari_env_simulated_spec(
+      env,
+      # Hardcoded for now. TODO(koz4k): Make it a hparam.
+      video_num_input_frames=4, video_num_target_frames=1
+  )
+
+  for epoch in range(hparams.epochs):
+    env.start_new_epoch(epoch)
+    is_final_epoch = (epoch + 1) == hparams.epochs
+    log = make_log_fn(epoch, log_relative_time)
+
+    epoch_data_dir = os.path.join(directories["data"], str(epoch))
+    tf.gfile.MakeDirs(epoch_data_dir)
+    env.generate_data(epoch_data_dir, directories["tmp"])
+    epoch_data_dirs.append(epoch_data_dir)
+
+    # Train world model
+    log("Training world model")
+    train_world_model(env, epoch_data_dir,
+                      directories["world_model"], hparams, epoch)
+
+    # Train PPO
+    log("Training PPO in simulated environment.")
+    ppo_event_dir = os.path.join(directories["world_model"],
+                                 "ppo_summaries", str(epoch))
+    ppo_model_dir = directories["ppo"]
+    if not hparams.ppo_continue_training:
+      ppo_model_dir = ppo_event_dir
+
+    train_agent(sim_env_spec, ppo_model_dir,
+                ppo_event_dir, directories["world_model"], epoch_data_dir,
+                hparams, epoch=epoch, is_final_epoch=is_final_epoch)
+
+    # Train PPO on real env (short)
+    log("Training PPO in real environment.")
+    train_agent_real_env(
+        env, ppo_model_dir,
+        ppo_event_dir, epoch_data_dir,
+        hparams, epoch=epoch, is_final_epoch=is_final_epoch)
+
+    if hparams.stop_loop_early:
+      return 0.0
+    mean_clipped_reward = eval_reward(env, epoch, clipped=True)
+    log("Mean clipped reward during generation: {}".format(
+        mean_clipped_reward))  # this was output of generate_real_env_data(...)
+
+    mean_unclipped_reward = eval_reward(env, epoch, clipped=False)
+    log("Mean eval reward (unclipped): {}".format(mean_unclipped_reward))
+
+    # Summarize metrics.
+    mean_unclipped_reward_summary.value[0].simple_value = mean_unclipped_reward
+    eval_metrics_writer.add_summary(mean_unclipped_reward_summary, epoch)
+    mean_clipped_reward_summary.value[0].simple_value = int(mean_clipped_reward)
+    eval_metrics_writer.add_summary(mean_clipped_reward_summary, epoch)
+    eval_metrics_writer.flush()
+
+    # Report metrics
+    eval_metrics = {"mean_unclipped_reward": mean_unclipped_reward}
+    epoch_metrics.append(eval_metrics)
+    log("Eval metrics: %s", str(eval_metrics))
+    if report_fn:
+      report_fn(eval_metrics[report_metric], epoch)
+
+  # Return the evaluation metrics from the final epoch
+  return epoch_metrics[-1]
+
+
+def main(_):
+  hp = trainer_model_based_params.create_loop_hparams()
+  assert not FLAGS.job_dir_to_evaluate
+  training_loop(hp, FLAGS.output_dir)
+
+
+if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
+  tf.app.run()
diff --git a/tensor2tensor/rl/trainer_model_based_new_test.py b/tensor2tensor/rl/trainer_model_based_new_test.py
new file mode 100644
index 000000000..bf23d592a
--- /dev/null
+++ b/tensor2tensor/rl/trainer_model_based_new_test.py
@@ -0,0 +1,38 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Tiny run of trainer_model_based_new. Smoke test."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.rl import trainer_model_based_new
+
+import tensorflow as tf
+
+FLAGS = tf.flags.FLAGS
+
+
+class ModelRLExperimentNewTest(tf.test.TestCase):
+
+  def test_basic(self):
+    FLAGS.output_dir = tf.test.get_temp_dir()
+    FLAGS.loop_hparams_set = "rlmb_tiny"
+    FLAGS.schedule = "train"  # skip evaluation for world model training
+    trainer_model_based_new.main(None)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 0f6f2d857b6ef86d56f835bcf0f01d5b543cc407 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 12 Oct 2018 19:47:26 -0700
Subject: [PATCH 1000/2720] internal merge of PR #1140

PiperOrigin-RevId: 216957930
---
 tensor2tensor/data_generators/gym_env.py    | 10 +++++++---
 tensor2tensor/models/research/rl.py         |  4 ++--
 tensor2tensor/rl/envs/batch_env_factory.py  |  6 +++---
 tensor2tensor/rl/trainer_model_based_new.py | 13 ++++++++-----
 4 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 3a44a50c5..0a5bf9a5e 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -78,7 +78,7 @@ def __init__(self, batch_size):
     self.batch_size = batch_size
     self._current_rollouts = [[] for _ in range(batch_size)]
     self._current_frames = [None for _ in range(batch_size)]
-
+    self.rollouts_by_epoch = dict()
     with tf.Graph().as_default() as tf_graph:
       self._tf_graph = _Noncopyable(tf_graph)
       self._image_t = _Noncopyable(
@@ -99,9 +99,9 @@ def clear_history(self):
 
   def start_new_epoch(self, epoch):
     if not isinstance(epoch, int):
-      raise ValueError('Epoch should be integer, got {}'.format(epoch))
+      raise ValueError("Epoch should be integer, got {}".format(epoch))
     if epoch in self.rollouts_by_epoch:
-      raise ValueError('Epoch {} already registered'.format(epoch))
+      raise ValueError("Epoch {} already registered".format(epoch))
     self.current_epoch = epoch
     self.rollouts_by_epoch[epoch] = list()
 
@@ -161,6 +161,8 @@ def step(self, actions):
       (obs, rewards, dones) - batches of observations, rewards and done flags
       respectively.
     """
+    if not self.rollouts_by_epoch:
+      self.start_new_epoch(0)
     (obs, unclipped_rewards, dones) = self._step(actions)
     obs = self._preprocess_observations(obs)
     (min_reward, max_reward) = self.reward_range
@@ -202,6 +204,8 @@ def reset(self, indices=None):
     Returns:
       Batch of initial observations of reset environments.
     """
+    if not self.rollouts_by_epoch:
+      self.start_new_epoch(0)
     if indices is None:
       indices = np.arange(self.batch_size)
     new_obs = self._reset(indices)
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 573d38c66..8f120688e 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -184,8 +184,8 @@ def standard_atari_env_spec(
 
 
 def standard_atari_env_simulated_spec(
-    real_env, video_num_input_frames, video_num_target_frames
-):
+    real_env, video_num_input_frames, video_num_target_frames):
+  """Spec."""
   env_spec = standard_atari_env_spec(
       # This hack is here because SimulatedBatchEnv needs to get
       # observation_space from the real env. TODO(koz4k): refactor.
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index 570f276bd..e7f5a8b2a 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -38,9 +38,9 @@ def batch_env_factory(environment_spec, num_agents, initial_frame_chooser=None):
     cur_batch_env = _define_simulated_batch_env(
         environment_spec, num_agents, initial_frame_chooser)
   else:
-    if 'batch_env' in environment_spec:
-      assert not 'env_lambda' in environment_spec, \
-          'Environment_spec should contain only one of (env_lambda, batch_env).'
+    if "batch_env" in environment_spec:
+      msg = "Environment_spec should contain only 1 of (env_lambda, batch_env)."
+      assert "env_lambda" not in environment_spec, msg
       batch_env = environment_spec.batch_env
       assert batch_env.batch_size == num_agents
     else:
diff --git a/tensor2tensor/rl/trainer_model_based_new.py b/tensor2tensor/rl/trainer_model_based_new.py
index 19f23b36d..b1e167b9b 100644
--- a/tensor2tensor/rl/trainer_model_based_new.py
+++ b/tensor2tensor/rl/trainer_model_based_new.py
@@ -39,7 +39,8 @@
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
 from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.models.research import rl
-from tensor2tensor.rl import rl_trainer_lib, trainer_model_based_params
+from tensor2tensor.rl import rl_trainer_lib
+from tensor2tensor.rl import trainer_model_based_params
 from tensor2tensor.rl.envs.utils import InitialFrameChooser
 from tensor2tensor.utils import trainer_lib
 
@@ -133,6 +134,7 @@ def log(msg, *args):
 
 def train_supervised(problem, model_name, hparams, data_dir, output_dir,
                      train_steps, eval_steps, local_eval_frequency=None):
+  """Train supervised."""
   if local_eval_frequency is None:
     local_eval_frequency = getattr(FLAGS, "local_eval_frequency")
 
@@ -198,6 +200,7 @@ def train_agent_real_env(
     env, agent_model_dir, event_dir, epoch_data_dir,
     hparams, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the real environment."""
+  del epoch_data_dir
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_params_names = ["epochs_num", "epoch_length",
                       "learning_rate", "num_agents", "eval_every_epochs",
@@ -256,9 +259,9 @@ def train_world_model(env, data_dir, output_dir, hparams, epoch):
 
 
 def setup_env(hparams):
+  """Setup."""
   # TODO(kc): set reward clipping, when this will be possible
-  assert hparams.game == "pong", "Currently only games with [-1, 1] reward " \
-                                 "range are working"
+  assert hparams.game == "pong", "Currently only games with [-1, 1] rewards."
   game_mode = "Deterministic-v4"
   camel_game_name = "".join(
       [w[0].upper() + w[1:] for w in hparams.game.split("_")])
@@ -273,7 +276,7 @@ def setup_env(hparams):
 
 def eval_reward(env, epoch, clipped):
   """Calculates mean rewards from given epoch."""
-  reward_name = 'reward' if clipped else 'unclipped_reward'
+  reward_name = "reward" if clipped else "unclipped_reward"
   rewards = []
   for rollout in env.rollouts_by_epoch[epoch]:
     if rollout[-1].done:
@@ -309,7 +312,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   data_dir = os.path.join(directories["data"], "initial")
   epoch_data_dirs.append(data_dir)
   # Collect data from the real environment with PPO or random policy.
-  # TODO: do we need option not to gather_ppo_real_env_data?
+  # TODO(lukaszkaiser): do we need option not to gather_ppo_real_env_data?
   # We could set learning_rate=0 if this flag == False.
   assert hparams.gather_ppo_real_env_data
   ppo_model_dir = directories["ppo"]

From e6a68526647249831f89f67cbe6c06e3c162f393 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 15 Oct 2018 10:14:24 -0700
Subject: [PATCH 1001/2720] Install tensorflow_hub for integration test since
 it uses it.

PiperOrigin-RevId: 217162863
---
 oss_scripts/oss_integration_test.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/oss_scripts/oss_integration_test.sh b/oss_scripts/oss_integration_test.sh
index 580a76866..f5bd63e3d 100755
--- a/oss_scripts/oss_integration_test.sh
+++ b/oss_scripts/oss_integration_test.sh
@@ -23,6 +23,7 @@ t2t-decoder --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR --model=transformer
 if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]
 then
   # Export for serving
+  pip install tensorflow_hub
   t2t-exporter \
       --problem=$T2T_PROBLEM \
       --data_dir=$T2T_DATA_DIR \

From eef7506441eda3c9a7a3e5b6d695037654da2c7b Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 15 Oct 2018 10:20:39 -0700
Subject: [PATCH 1002/2720] Move the flags to where they are used, Travis
 throws an error on a test which passes locally. These should have been moved
 in PR 1137

PiperOrigin-RevId: 217163975
---
 tensor2tensor/rl/trainer_model_based.py        | 3 ---
 tensor2tensor/rl/trainer_model_based_new.py    | 3 ---
 tensor2tensor/rl/trainer_model_based_params.py | 7 ++++++-
 3 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index bde551c6a..6989dfbe5 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -56,9 +56,6 @@
 flags = tf.flags
 FLAGS = flags.FLAGS
 
-flags.DEFINE_string("loop_hparams_set", "rlmb_base",
-                    "Which RL hparams set to use.")
-flags.DEFINE_string("loop_hparams", "", "Overrides for overall loop HParams.")
 flags.DEFINE_string("job_dir_to_evaluate", "",
                     "Directory of a job to be evaluated.")
 flags.DEFINE_string("eval_results_dir", "/tmp",
diff --git a/tensor2tensor/rl/trainer_model_based_new.py b/tensor2tensor/rl/trainer_model_based_new.py
index b1e167b9b..8aab5ac9b 100644
--- a/tensor2tensor/rl/trainer_model_based_new.py
+++ b/tensor2tensor/rl/trainer_model_based_new.py
@@ -50,9 +50,6 @@
 flags = tf.flags
 FLAGS = flags.FLAGS
 
-flags.DEFINE_string("loop_hparams_set", "rlmb_base",
-                    "Which RL hparams set to use.")
-flags.DEFINE_string("loop_hparams", "", "Overrides for overall loop HParams.")
 flags.DEFINE_string("job_dir_to_evaluate", "",
                     "Directory of a job to be evaluated.")
 flags.DEFINE_string("eval_results_dir", "/tmp",
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 953a2e5e5..2660e3672 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -27,9 +27,14 @@
 import tensorflow as tf
 
 
-FLAGS = tf.flags.FLAGS
+flags = tf.flags
+FLAGS = flags.FLAGS
 
 
+flags.DEFINE_string("loop_hparams_set", "rlmb_base",
+                    "Which RL hparams set to use.")
+flags.DEFINE_string("loop_hparams", "", "Overrides for overall loop HParams.")
+
 HP_SCOPES = ["loop", "model", "ppo"]
 
 

From 9de111227fe006bbec335635ccbc55365bdf083a Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Mon, 15 Oct 2018 19:50:00 +0200
Subject: [PATCH 1003/2720] Model-Based RL: Fix reward evaluation, set default
 continuous_train_and_eval in train_supervised (#1143)

* Set default schedule to continuous_train_and_eval in train_supervised

* Fix reward evaluation, add env_timelimit flag, set short rollouts for tiny hparams set, to test evaluation
---
 tensor2tensor/rl/trainer_model_based_new.py   | 25 ++++++++++++++-----
 .../rl/trainer_model_based_params.py          |  2 ++
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based_new.py b/tensor2tensor/rl/trainer_model_based_new.py
index 8aab5ac9b..2e78d6ea9 100644
--- a/tensor2tensor/rl/trainer_model_based_new.py
+++ b/tensor2tensor/rl/trainer_model_based_new.py
@@ -130,7 +130,8 @@ def log(msg, *args):
 
 
 def train_supervised(problem, model_name, hparams, data_dir, output_dir,
-                     train_steps, eval_steps, local_eval_frequency=None):
+                     train_steps, eval_steps, local_eval_frequency=None,
+                     schedule="continuous_train_and_eval"):
   """Train supervised."""
   if local_eval_frequency is None:
     local_eval_frequency = getattr(FLAGS, "local_eval_frequency")
@@ -141,7 +142,7 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
   )
   run_config = trainer_lib.create_run_config(model_dir=output_dir)
   exp = exp_fn(run_config, hparams)
-  exp.test()
+  getattr(exp, schedule)()
 
 
 def train_agent(environment_spec, agent_model_dir,
@@ -255,8 +256,7 @@ def train_world_model(env, data_dir, output_dir, hparams, epoch):
   )
 
 
-def setup_env(hparams):
-  """Setup."""
+def make_gym_env(hparams):
   # TODO(kc): set reward clipping, when this will be possible
   assert hparams.game == "pong", "Currently only games with [-1, 1] rewards."
   game_mode = "Deterministic-v4"
@@ -264,7 +264,20 @@ def setup_env(hparams):
       [w[0].upper() + w[1:] for w in hparams.game.split("_")])
   camel_game_name += game_mode
   env_name = camel_game_name
-  env = T2TGymEnv([gym.make(env_name)],
+  env = gym.make(env_name)
+  if hparams.env_timesteps_limit != -1:
+    # Replace TimeLimit Wrapper with one of proper time step limit.
+    if isinstance(env, gym.wrappers.TimeLimit):
+      env = env.env
+    env = gym.wrappers.TimeLimit(env,
+                                 max_episode_steps=hparams.env_timesteps_limit)
+  return env
+
+
+def setup_env(hparams):
+  """Setup."""
+  env = T2TGymEnv([make_gym_env(hparams)
+                   for _ in range(hparams.real_ppo_num_agents)],
                   grayscale=hparams.grayscale,
                   resize_width_factor=hparams.resize_width_factor,
                   resize_height_factor=hparams.resize_height_factor)
@@ -277,7 +290,7 @@ def eval_reward(env, epoch, clipped):
   rewards = []
   for rollout in env.rollouts_by_epoch[epoch]:
     if rollout[-1].done:
-      rollout_reward = sum(frame[reward_name] for frame in rollout)
+      rollout_reward = sum(getattr(frame, reward_name) for frame in rollout)
       rewards.append(rollout_reward)
   if rewards:
     mean_rewards = np.mean(rewards)
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 2660e3672..05ee6d3e6 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -104,6 +104,7 @@ def rlmb_base():
       # Rollout fractions to report reward_accuracy on.
       eval_rollout_fractions=[0.25, 0.5, 1],
       stop_loop_early=False,  # To speed-up tests.
+      env_timesteps_limit=-1,  # Use default from gym.make()
   )
 
 
@@ -389,6 +390,7 @@ def rlmb_tiny():
           resize_height_factor=2,
           resize_width_factor=2,
           game="pong",
+          env_timesteps_limit=6,
       ).values())
 
 

From ad6758e7be7215ed83706f08bb163c7009598350 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Mon, 15 Oct 2018 10:50:21 -0700
Subject: [PATCH 1004/2720] internal merge of PR #1143

PiperOrigin-RevId: 217169860
---
 tensor2tensor/rl/trainer_model_based_new.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/rl/trainer_model_based_new.py b/tensor2tensor/rl/trainer_model_based_new.py
index 2e78d6ea9..563f1d186 100644
--- a/tensor2tensor/rl/trainer_model_based_new.py
+++ b/tensor2tensor/rl/trainer_model_based_new.py
@@ -257,6 +257,7 @@ def train_world_model(env, data_dir, output_dir, hparams, epoch):
 
 
 def make_gym_env(hparams):
+  """Make env."""
   # TODO(kc): set reward clipping, when this will be possible
   assert hparams.game == "pong", "Currently only games with [-1, 1] rewards."
   game_mode = "Deterministic-v4"

From b9824a7214b7854851daa530d31f9b2d30b5742e Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 15 Oct 2018 11:12:16 -0700
Subject: [PATCH 1005/2720] Moving more flag definitions from
 trainer_model_based{_new} to trainer_model_based_params.py

PiperOrigin-RevId: 217174089
---
 tensor2tensor/rl/trainer_model_based.py        | 5 -----
 tensor2tensor/rl/trainer_model_based_new.py    | 5 -----
 tensor2tensor/rl/trainer_model_based_params.py | 5 +++++
 3 files changed, 5 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 6989dfbe5..24d2dcfdc 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -56,11 +56,6 @@
 flags = tf.flags
 FLAGS = flags.FLAGS
 
-flags.DEFINE_string("job_dir_to_evaluate", "",
-                    "Directory of a job to be evaluated.")
-flags.DEFINE_string("eval_results_dir", "/tmp",
-                    "Directory to store result of evaluation")
-
 
 def setup_directories(base_dir, subdirs):
   base_dir = os.path.expanduser(base_dir)
diff --git a/tensor2tensor/rl/trainer_model_based_new.py b/tensor2tensor/rl/trainer_model_based_new.py
index 563f1d186..86317e261 100644
--- a/tensor2tensor/rl/trainer_model_based_new.py
+++ b/tensor2tensor/rl/trainer_model_based_new.py
@@ -50,11 +50,6 @@
 flags = tf.flags
 FLAGS = flags.FLAGS
 
-flags.DEFINE_string("job_dir_to_evaluate", "",
-                    "Directory of a job to be evaluated.")
-flags.DEFINE_string("eval_results_dir", "/tmp",
-                    "Directory to store result of evaluation")
-
 
 @contextlib.contextmanager
 def temporary_flags(flag_settings):
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 05ee6d3e6..9e3fd0b91 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -34,6 +34,11 @@
 flags.DEFINE_string("loop_hparams_set", "rlmb_base",
                     "Which RL hparams set to use.")
 flags.DEFINE_string("loop_hparams", "", "Overrides for overall loop HParams.")
+flags.DEFINE_string("job_dir_to_evaluate", "",
+                    "Directory of a job to be evaluated.")
+flags.DEFINE_string("eval_results_dir", "/tmp",
+                    "Directory to store result of evaluation")
+
 
 HP_SCOPES = ["loop", "model", "ppo"]
 

From b3bf36d6696e58988605a3683d163266e27fcc80 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 15 Oct 2018 12:58:36 -0700
Subject: [PATCH 1006/2720] Use one_hot and batch matmul instead of gather_nd
 to improve the performance of beam search on TPU.

PiperOrigin-RevId: 217192284
---
 tensor2tensor/utils/beam_search.py | 186 +++++++++++++++++++++--------
 1 file changed, 137 insertions(+), 49 deletions(-)

diff --git a/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index 92f9fc76c..1d72c88c4 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -105,9 +105,66 @@ def compute_batch_indices(batch_size, beam_size):
   return batch_pos
 
 
+def fast_tpu_gather(params, indices, name=None):
+  """Fast gather implementation for models running on TPU.
+
+  This function use one_hot and batch matmul to do gather, which is faster
+  than gather_nd on TPU. For params that have dtype of int32 (sequences to
+  gather from), batch_gather is used to keep accuracy.
+
+  Args:
+    params: A tensor from which to gather values.
+      [batch_size, original_size, ...]
+    indices: A tensor used as the index to gather values.
+      [batch_size, selected_size].
+    name: A string, name of the operation (optional).
+
+  Returns:
+    gather_result: A tensor that has the same rank as params.
+      [batch_size, selected_size, ...]
+  """
+  with tf.name_scope(name):
+    dtype = params.dtype
+
+    def _gather(params, indices):
+      """Fast gather using one_hot and batch matmul."""
+      if dtype != tf.float32:
+        params = tf.to_float(params)
+      shape = common_layers.shape_list(params)
+      indices_shape = common_layers.shape_list(indices)
+      ndims = params.shape.ndims
+      # Adjust the shape of params to match one-hot indices, which is the
+      # requirement of Batch MatMul.
+      if ndims == 2:
+        params = tf.expand_dims(params, axis=-1)
+      if ndims > 3:
+        params = tf.reshape(params, [shape[0], shape[1], -1])
+      gather_result = tf.matmul(
+          tf.one_hot(indices, shape[1], dtype=params.dtype), params)
+      if ndims == 2:
+        gather_result = tf.squeeze(gather_result, axis=-1)
+      if ndims > 3:
+        shape[1] = indices_shape[1]
+        gather_result = tf.reshape(gather_result, shape)
+      if dtype != tf.float32:
+        gather_result = tf.cast(gather_result, dtype)
+      return gather_result
+
+    # If the dtype is int32, use the gather instead of one_hot matmul to avoid
+    # precision loss. The max int value can be represented by bfloat16 in MXU is
+    # 256, which is smaller than the possible id values. Encoding/decoding can
+    # potentially used to make it work, but the benenfit is small right now.
+    if dtype == tf.int32:
+      gather_result = tf.batch_gather(params, indices)
+    else:
+      gather_result = _gather(params, indices)
+
+    return gather_result
+
+
 def compute_topk_scores_and_seq(sequences, scores, scores_to_gather, flags,
                                 beam_size, batch_size, prefix="default",
-                                states_to_gather=None):
+                                states_to_gather=None, use_tpu=False):
   """Given sequences and scores, will gather the top k=beam size sequences.
 
   This function is used to grow alive, and finished. It takes sequences,
@@ -136,6 +193,8 @@ def compute_topk_scores_and_seq(sequences, scores, scores_to_gather, flags,
     batch_size: int
     prefix: string that will prefix unique names for the ops run.
     states_to_gather: dict (possibly nested) of decoding states.
+    use_tpu: A bool, whether to compute topk scores and sequences on TPU.
+
   Returns:
     Tuple of
     (topk_seq [batch_size, beam_size, decode_length],
@@ -143,31 +202,48 @@ def compute_topk_scores_and_seq(sequences, scores, scores_to_gather, flags,
      topk_finished_flags[batch_size, beam_size])
   """
   _, topk_indexes = tf.nn.top_k(scores, k=beam_size)
-  # The next three steps are to create coordinates for tf.gather_nd to pull
-  # out the topk sequences from sequences based on scores.
-  # batch pos is a tensor like [[0,0,0,0,],[1,1,1,1],..]. It says which
-  # batch the beam item is in. This will create the i of the i,j coordinate
-  # needed for the gather
-  batch_pos = compute_batch_indices(batch_size, beam_size)
-
-  # top coordinates will give us the actual coordinates to do the gather.
-  # stacking will create a tensor of dimension batch * beam * 2, where the
-  # last dimension contains the i,j gathering coordinates.
-  top_coordinates = tf.stack([batch_pos, topk_indexes], axis=2)
-
-  # Gather up the highest scoring sequences.  For each operation added, give it
-  # a concrete name to simplify observing these operations with tfdbg.  Clients
-  # can capture these tensors by watching these node names.
-  def gather(tensor, name):
-    return tf.gather_nd(tensor, top_coordinates, name=(prefix + name))
-  topk_seq = gather(sequences, "_topk_seq")
-  topk_flags = gather(flags, "_topk_flags")
-  topk_gathered_scores = gather(scores_to_gather, "_topk_scores")
-  if states_to_gather:
-    topk_gathered_states = nest.map_structure(
-        lambda state: gather(state, "_topk_states"), states_to_gather)
+  if not use_tpu:
+    # The next three steps are to create coordinates for tf.gather_nd to pull
+    # out the topk sequences from sequences based on scores.
+    # batch pos is a tensor like [[0,0,0,0,],[1,1,1,1],..]. It says which
+    # batch the beam item is in. This will create the i of the i,j coordinate
+    # needed for the gather
+    batch_pos = compute_batch_indices(batch_size, beam_size)
+
+    # top coordinates will give us the actual coordinates to do the gather.
+    # stacking will create a tensor of dimension batch * beam * 2, where the
+    # last dimension contains the i,j gathering coordinates.
+    top_coordinates = tf.stack([batch_pos, topk_indexes], axis=2)
+
+    # Gather up the highest scoring sequences.  For each operation added, give
+    # it a concrete name to simplify observing these operations with tfdbg.
+    # Clients can capture these tensors by watching these node names.
+    def gather(tensor, name):
+      return tf.gather_nd(tensor, top_coordinates, name=(prefix + name))
+    topk_seq = gather(sequences, "_topk_seq")
+    topk_flags = gather(flags, "_topk_flags")
+    topk_gathered_scores = gather(scores_to_gather, "_topk_scores")
+    if states_to_gather:
+      topk_gathered_states = nest.map_structure(
+          lambda state: gather(state, "_topk_states"), states_to_gather)
+    else:
+      topk_gathered_states = states_to_gather
   else:
-    topk_gathered_states = states_to_gather
+    # Gather up the highest scoring sequences.  For each operation added, give
+    # it a concrete name to simplify observing these operations with tfdbg.
+    # Clients can capture these tensors by watching these node names.
+    topk_seq = fast_tpu_gather(sequences, topk_indexes, prefix + "_topk_seq")
+    topk_flags = fast_tpu_gather(flags, topk_indexes, prefix + "_topk_flags")
+    topk_gathered_scores = fast_tpu_gather(scores_to_gather, topk_indexes,
+                                           prefix + "_topk_scores")
+    if states_to_gather:
+      topk_gathered_states = nest.map_structure(
+          # pylint: disable=g-long-lambda
+          lambda state: fast_tpu_gather(state, topk_indexes,
+                                        prefix + "_topk_states"),
+          states_to_gather)
+    else:
+      topk_gathered_states = states_to_gather
   return topk_seq, topk_gathered_scores, topk_flags, topk_gathered_states
 
 
@@ -231,7 +307,7 @@ def beam_search(symbols_to_logits_fn,
   batch_size = common_layers.shape_list(initial_ids)[0]
 
   # Assume initial_ids are prob 1.0
-  initial_log_probs = tf.constant([[0.] + [-float("inf")] * (beam_size - 1)])
+  initial_log_probs = tf.constant([[0.] + [-INF] * (beam_size - 1)])
   # Expand to beam_size (batch_size, beam_size)
   alive_log_probs = tf.tile(initial_log_probs, [batch_size, 1])
 
@@ -292,7 +368,8 @@ def grow_finished(finished_seq, finished_scores, finished_flags, curr_seq,
     curr_finished_flags = tf.concat([finished_flags, curr_finished], axis=1)
     return compute_topk_scores_and_seq(
         curr_finished_seq, curr_finished_scores, curr_finished_scores,
-        curr_finished_flags, beam_size, batch_size, "grow_finished")
+        curr_finished_flags, beam_size, batch_size, "grow_finished",
+        use_tpu=use_tpu)
 
   def grow_alive(curr_seq, curr_scores, curr_log_probs, curr_finished, states):
     """Given sequences and scores, will gather the top k=beam size sequences.
@@ -317,7 +394,7 @@ def grow_alive(curr_seq, curr_scores, curr_log_probs, curr_finished, states):
     curr_scores += tf.to_float(curr_finished) * -INF
     return compute_topk_scores_and_seq(curr_seq, curr_scores, curr_log_probs,
                                        curr_finished, beam_size, batch_size,
-                                       "grow_alive", states)
+                                       "grow_alive", states, use_tpu=use_tpu)
 
   def grow_topk(i, alive_seq, alive_log_probs, states):
     r"""Inner beam search loop.
@@ -386,31 +463,40 @@ def grow_topk(i, alive_seq, alive_log_probs, states):
     topk_beam_index = topk_ids // vocab_size
     topk_ids %= vocab_size  # Unflatten the ids
 
-    # The next three steps are to create coordinates for tf.gather_nd to pull
-    # out the correct sequences from id's that we need to grow.
-    # We will also use the coordinates to gather the booleans of the beam items
-    # that survived.
-    batch_pos = compute_batch_indices(batch_size, beam_size * 2)
-
-    # top beams will give us the actual coordinates to do the gather.
-    # stacking will create a tensor of dimension batch * beam * 2, where the
-    # last dimension contains the i,j gathering coordinates.
-    topk_coordinates = tf.stack([batch_pos, topk_beam_index], axis=2)
+    if not use_tpu:
+      # The next three steps are to create coordinates for tf.gather_nd to pull
+      # out the correct sequences from id's that we need to grow.
+      # We will also use the coordinates to gather the booleans of the beam
+      # items that survived.
+      batch_pos = compute_batch_indices(batch_size, beam_size * 2)
+
+      # top beams will give us the actual coordinates to do the gather.
+      # stacking will create a tensor of dimension batch * beam * 2, where the
+      # last dimension contains the i,j gathering coordinates.
+      topk_coordinates = tf.stack([batch_pos, topk_beam_index], axis=2)
+
+      # Gather up the most probable 2*beams both for the ids and
+      # finished_in_alive bools
+      topk_seq = tf.gather_nd(alive_seq, topk_coordinates)
+      if states:
+        states = nest.map_structure(
+            lambda state: tf.gather_nd(state, topk_coordinates), states)
+
+      # Append the most probable alive
+      topk_seq = tf.concat([topk_seq, tf.expand_dims(topk_ids, axis=2)], axis=2)
+    else:
+      # Gather up the most probable 2*beams both for the ids and
+      # finished_in_alive bools
+      topk_seq = fast_tpu_gather(alive_seq, topk_beam_index)
 
-    # Gather up the most probable 2*beams both for the ids and finished_in_alive
-    # bools
-    topk_seq = tf.gather_nd(alive_seq, topk_coordinates)
-    if states:
-      states = nest.map_structure(
-          lambda state: tf.gather_nd(state, topk_coordinates), states)
+      if states:
+        states = nest.map_structure(
+            lambda state: fast_tpu_gather(state, topk_beam_index), states)
 
-    # Append the most probable alive
-    if use_tpu:
+      # Update the most probable alive
       topk_seq = tf.transpose(topk_seq, perm=[2, 0, 1])
       topk_seq = inplace_ops.alias_inplace_update(topk_seq, i + 1, topk_ids)
       topk_seq = tf.transpose(topk_seq, perm=[1, 2, 0])
-    else:
-      topk_seq = tf.concat([topk_seq, tf.expand_dims(topk_ids, axis=2)], axis=2)
 
     topk_finished = tf.equal(topk_ids, eos_id)
 
@@ -538,7 +624,9 @@ def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
             if use_tpu else tf.TensorShape([None, None, None])),
            finished_scores.get_shape(),
            finished_flags.get_shape(),
-           nest.map_structure(get_state_shape_invariants, states),
+           (nest.map_structure(lambda state: state.get_shape(), states)
+            if use_tpu else
+            nest.map_structure(get_state_shape_invariants, states)),
        ],
        parallel_iterations=1,
        back_prop=False)

From c9d0210d6ce5f98ebe7e38d8d3cf0f1ecaf299bd Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 15 Oct 2018 13:49:19 -0700
Subject: [PATCH 1007/2720] Allow training on random_patches extracted from the
 video. The number of random patches is controlled by
 "hparams.num_train_frames" which defaults to -1 which means use all frames
 for training.

PiperOrigin-RevId: 217200930
---
 tensor2tensor/layers/common_video.py      | 40 +++++++++++++++++++++++
 tensor2tensor/layers/common_video_test.py | 31 ++++++++++++++++--
 2 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 0ae47e54c..abc5dbb31 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -556,6 +556,46 @@ def beta_schedule(schedule, global_step, final_beta, decay_start, decay_end):
   return beta
 
 
+def extract_random_video_patch(videos, num_frames=-1):
+  """For every video, extract a random consecutive patch of num_frames.
+
+  Args:
+    videos: 5-D Tensor, (NTHWC)
+    num_frames: Integer, if -1 then the entire video is returned.
+  Returns:
+    video_patch: 5-D Tensor, (NTHWC) with T = num_frames.
+  Raises:
+    ValueError: If num_frames is greater than the number of total frames in
+                the video.
+  """
+  if num_frames == -1:
+    return videos
+  batch_size, num_total_frames, h, w, c = common_layers.shape_list(videos)
+  if num_total_frames < num_frames:
+    raise ValueError("Expected num_frames <= %d, got %d" %
+                     (num_total_frames, num_frames))
+
+  # Randomly choose start_inds for each video.
+  frame_start = tf.random_uniform(
+      shape=(batch_size,), minval=0, maxval=num_total_frames - num_frames + 1,
+      dtype=tf.int32)
+
+  # [start[0], start[0] + 1, ... start[0] + num_frames - 1] + ...
+  # [start[batch_size-1], ... start[batch_size-1] + num_frames - 1]
+  range_inds = tf.expand_dims(tf.range(num_frames), axis=0)
+  frame_inds = range_inds + tf.expand_dims(frame_start, axis=1)
+  frame_inds = tf.reshape(frame_inds, [-1])
+
+  # [0]*num_frames + [1]*num_frames + ... [batch_size-1]*num_frames
+  batch_inds = tf.expand_dims(tf.range(batch_size), axis=1)
+  batch_inds = tf.tile(batch_inds, [1, num_frames])
+  batch_inds = tf.reshape(batch_inds, [-1])
+
+  gather_inds = tf.stack((batch_inds, frame_inds), axis=1)
+  video_patches = tf.gather_nd(videos, gather_inds)
+  return tf.reshape(video_patches, (batch_size, num_frames, h, w, c))
+
+
 class VideoWriter(object):
   """Base helper class for writing videos."""
 
diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index 3ab72ce1a..2e41cd98a 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -19,14 +19,14 @@
 from __future__ import division
 from __future__ import print_function
 
+from absl.testing import parameterized
 import numpy as np
 
 from tensor2tensor.layers import common_video
-
 import tensorflow as tf
 
 
-class CommonVideoTest(tf.test.TestCase):
+class CommonVideoTest(parameterized.TestCase, tf.test.TestCase):
 
   def _run_scheduled_sample_func(self, func, var, batch_size):
     ground_truth_x = list(range(1, batch_size+1))
@@ -122,6 +122,33 @@ def testGifSummary(self):
 
       self.assertEqual(encoded, common_video._encode_gif(images[0], fps=10))  # pylint: disable=protected-access
 
+  def checkIfPatchExists(self, videos, video_patches, num_frames):
+    """Check that given patch is present in video."""
+    for video, video_patch in zip(videos, video_patches):
+      total_frames = len(video)
+      is_present = []
+      for start_ind in range(total_frames - num_frames + 1):
+        curr_patch = video[start_ind: start_ind + num_frames]
+        is_present.append(np.allclose(curr_patch, video_patch))
+      self.assertTrue(np.any(is_present))
+
+  @parameterized.named_parameters(
+      ("two_frames", 2), ("ten_frames", 10), ("default", -1))
+  def testExtractRandomVideoPatch(self, num_frames=2):
+    with tf.Graph().as_default():
+      rng = np.random.RandomState(0)
+      video_np = rng.randint(0, 255, size=(12, 20, 256, 256, 3))
+      video = tf.convert_to_tensor(video_np)
+      video_patch = common_video.extract_random_video_patch(
+          video, num_frames=num_frames)
+      with tf.Session() as sess:
+        video_patch_np = sess.run(video_patch)
+        if num_frames != -1:
+          self.assertEqual(video_patch_np.shape, (12, num_frames, 256, 256, 3))
+          self.checkIfPatchExists(video_np, video_patch_np, num_frames)
+        else:
+          self.assertTrue(np.allclose(video_np, video_patch_np))
+
 
 if __name__ == "__main__":
   tf.test.main()

From 14b467d939a003dab35855c15386f5dcf60d0b27 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 15 Oct 2018 15:16:32 -0700
Subject: [PATCH 1008/2720] Merged the body and infer functions for all the
 basic models. It's a bigger CL so please review it more carefully. I tried
 some trainings locally but there is still a chance for everything being
 broken!

PiperOrigin-RevId: 217216445
---
 tensor2tensor/models/video/base.py            | 522 ++++++++++++++++++
 .../models/video/basic_deterministic.py       | 373 ++-----------
 tensor2tensor/models/video/basic_recurrent.py | 158 +-----
 .../models/video/basic_stochastic.py          |  13 +-
 4 files changed, 580 insertions(+), 486 deletions(-)
 create mode 100644 tensor2tensor/models/video/base.py

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
new file mode 100644
index 000000000..c5db06b71
--- /dev/null
+++ b/tensor2tensor/models/video/base.py
@@ -0,0 +1,522 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Basic models for testing simple tasks."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from functools import partial
+import six
+
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
+from tensor2tensor.models.video import basic_deterministic_params  # pylint: disable=unused-import
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
+
+import tensorflow as tf
+
+
+tfl = tf.layers
+tfcl = tf.contrib.layers
+
+
+@registry.register_model
+class NextFrameBase(t2t_model.T2TModel):
+  """Base class for next_frame models.
+
+    This is the base class for the models that given the previous frames
+    can predict the next frame. They may also support reward prediction
+    and action condition prediction which enables them to be run as
+    a world model in model-based RL pipeline.
+
+    The API supports both recurrent and stacked frames models. Please look
+    at the documents for next_frame function for the API.
+
+    If you are implementing a next frame prediction model consider
+    following the API presented in this class. But if the API
+    is too limiting for your models, feel free to override lower level
+    functions and/or inheret from T2TModel directly.
+
+  """
+
+  # ============================================================================
+  # BEGIN SUBCLASS INTERFACE
+  # ============================================================================
+  def next_frame(self,
+                 frames, actions, rewards,
+                 target_frame, internal_states, video_features):
+    """The main prediction function of next frame models.
+
+      This is the main function that should be overridden to implement models.
+
+    Args:
+      frames: The list of input frames.
+              Only previous frame in case of recurrent models.
+      actions: The list of input actions.
+              Only previous action in case of recurrent models.
+      rewards: The list of input rewards.
+              Only previous reward in case of recurrent models.
+      target_frame: The target frame.
+              Usually required for approximating the posterior.
+      internal_states: Internal model states. Only useful for recurrent models
+              to keep the state from the previous time index.
+              internal_states is None at the first frame and should be
+              initialized properly.
+      video_features: video wide features. None by default.
+              Please refer to video_features function for description.
+
+    Returns:
+      pred_frame: predicted frame BSxWxHxC
+              where C is 3 for L1/L2 modality and 3*256 for Softmax.
+      pred_reward: the same size as input reward.
+              None if the model does not detect rewards.
+      extra_loss: any extra loss other than predicted frame and reward.
+              e.g. KL loss in case of VAE models.
+      internal_states: updated internal models states.
+    """
+    raise NotImplementedError("Base video model.")
+
+  def video_features(
+      self, all_frames, all_actions, all_rewards, all_raw_frames):
+    """Optional video wide features.
+
+      If the model requires access to all of the video frames
+      (e.g. in case of approximating one latent for the whole video)
+      override this function to add them. They will be accessible
+      as video_features in next_frame function.
+
+    Args:
+      all_frames: list of all frames including input and target frames.
+      all_actions: list of all actions including input and target actions.
+      all_rewards: list of all rewards including input and target rewards.
+      all_raw_frames: list of all raw frames (before modalities).
+
+    Returns:
+      video_features: a dictionary containing video-wide features.
+    """
+    del all_frames, all_actions, all_rewards, all_raw_frames
+    return None
+
+  @property
+  def is_recurrent_model(self):
+    """Set to true if your model is recurrent. False otherwise.
+
+    This mainly affects how the inputs will be fed into next_frame function.
+    """
+    raise NotImplementedError("Base video model.")
+
+  # ============================================================================
+  # END SUBCLASS INTERFACE
+  # ============================================================================
+
+  @property
+  def _target_modality(self):
+    # TODO(mbz): get rid of this somehow.
+    modality = self.hparams.problem_hparams.target_modality["targets"]
+    return modality.__class__.__name__
+
+  @property
+  def is_per_pixel_softmax(self):
+    return self._target_modality == "VideoModality"
+
+  def get_iteration_num(self):
+    step_num = tf.train.get_global_step()
+    # TODO(lukaszkaiser): what should it be if it's undefined?
+    if step_num is None:
+      step_num = 10000000
+    return step_num
+
+  def get_scheduled_sample_func(self, batch_size):
+    """Creates a function for scheduled sampling based on given hparams."""
+    with tf.variable_scope("scheduled_sampling_func", reuse=tf.AUTO_REUSE):
+      iter_num = self.get_iteration_num()
+
+      # Simple function to bypass scheduled sampling in gt or pred only modes.
+      def scheduled_sampling_simple(ground_truth_x, generated_x,
+                                    batch_size, scheduled_sample_var):
+        del batch_size
+        if scheduled_sample_var:
+          return ground_truth_x
+        return generated_x
+
+      mode = self.hparams.scheduled_sampling_mode
+      if mode == "ground_truth_only":
+        scheduled_sampling_func = scheduled_sampling_simple
+        scheduled_sampling_func_var = True
+      elif mode == "prediction_only":
+        scheduled_sampling_func = scheduled_sampling_simple
+        scheduled_sampling_func_var = False
+      elif mode == "prob":
+        decay_steps = self.hparams.scheduled_sampling_decay_steps
+        probability = tf.train.polynomial_decay(
+            1.0, iter_num, decay_steps, 0.0)
+        scheduled_sampling_func = common_video.scheduled_sample_prob
+        scheduled_sampling_func_var = probability
+      elif mode == "prob_inverse_exp":
+        decay_steps = self.hparams.scheduled_sampling_decay_steps
+        probability = common_layers.inverse_exp_decay(
+            decay_steps, step=iter_num)
+        probability *= self.hparams.scheduled_sampling_max_prob
+        probability = 1.0 - probability
+        scheduled_sampling_func = common_video.scheduled_sample_prob
+        scheduled_sampling_func_var = probability
+      elif mode == "prob_inverse_lin":
+        decay_steps = self.hparams.scheduled_sampling_decay_steps
+        probability = common_layers.inverse_exp_decay(
+            decay_steps // 4, step=iter_num)  # Very low at start.
+        probability *= common_layers.inverse_lin_decay(
+            decay_steps, step=iter_num)
+        probability *= self.hparams.scheduled_sampling_max_prob
+        probability = 1.0 - probability
+        scheduled_sampling_func = common_video.scheduled_sample_prob
+        scheduled_sampling_func_var = probability
+      elif mode == "count":
+        # Calculate number of ground-truth frames to pass in.
+        k = self.hparams.scheduled_sampling_k
+        num_ground_truth = tf.to_int32(
+            tf.round(
+                tf.to_float(batch_size) *
+                (k / (k + tf.exp(tf.to_float(iter_num) / tf.to_float(k))))))
+        scheduled_sampling_func = common_video.scheduled_sample_count
+        scheduled_sampling_func_var = num_ground_truth
+      else:
+        raise ValueError("unknown scheduled sampling method: %s" % mode)
+
+      if isinstance(scheduled_sampling_func_var, tf.Tensor):
+        tf.summary.scalar("scheduled_sampling_var", scheduled_sampling_func_var)
+      partial_func = partial(scheduled_sampling_func,
+                             batch_size=batch_size,
+                             scheduled_sample_var=scheduled_sampling_func_var)
+      return partial_func
+
+  def get_scheduled_sample_inputs(self,
+                                  done_warm_start,
+                                  groundtruth_items,
+                                  generated_items,
+                                  scheduled_sampling_func):
+    """Scheduled sampling.
+
+    Args:
+      done_warm_start: whether we are done with warm start or not.
+      groundtruth_items: list of ground truth items.
+      generated_items: list of generated items.
+      scheduled_sampling_func: scheduled sampling function to choose between
+        groundtruth items and generated items.
+
+    Returns:
+      A mix list of ground truth and generated items.
+    """
+    def sample():
+      """Calculate the scheduled sampling params based on iteration number."""
+      with tf.variable_scope("scheduled_sampling", reuse=tf.AUTO_REUSE):
+        output_items = []
+        for item_gt, item_gen in zip(groundtruth_items, generated_items):
+          output_items.append(scheduled_sampling_func(item_gt, item_gen))
+        return output_items
+
+    cases = [
+        (tf.logical_not(done_warm_start), lambda: groundtruth_items),
+        (tf.logical_not(self.is_training), lambda: generated_items),
+    ]
+    output_items = tf.case(cases, default=sample, strict=True)
+
+    return output_items
+
+  def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
+    """Hacky code the get the loss on predicted frames from input frames.
+
+       Recurrent models consume the frames one-by-one. Therefore
+       if there is more than one input frame they also get predicted.
+       T2T only calculates loss on the predicted target frames which
+       means the loss is not being applied on the predicted input frames.
+       This code is to fix this issue. Since the model is not aware of the
+       modality it has to match the pre-porocessing happening in bottom
+       function and therefore this becomes a very hacky code. This code
+       should match the bottom and top and loss of modalities otherwise
+       it will calculate the wrong loss.
+
+    Args:
+      extra_raw_gts: extra raw ground truth frames.
+      extra_gts: extra normalized ground truth frames.
+      extra_pds: extra predicted frames.
+
+    Returns:
+      Additional reconstruction loss.
+
+    Raises:
+      ValueError: in case of unknown modality.
+    """
+    if self._target_modality == "VideoModalityL2Raw":
+      recon_loss = tf.losses.mean_squared_error(extra_gts, extra_pds)
+    elif self._target_modality == "VideoModality":
+      shape = common_layers.shape_list(extra_pds)
+      updated_shape = shape[:-1] + [3, 256]
+      extra_pds = tf.reshape(extra_pds, updated_shape)
+      # Merge time and batch
+      logits = tf.reshape(extra_pds, [-1] + updated_shape[2:])
+      targets = extra_raw_gts
+      targets_shape = common_layers.shape_list(targets)
+      targets = tf.reshape(targets, [-1] + targets_shape[2:])
+      mod = self.hparams.problem_hparams.target_modality["targets"]
+      numerator, denominator = common_layers.padded_cross_entropy(
+          logits,
+          targets,
+          self.hparams.label_smoothing,
+          cutoff=getattr(self.hparams, "video_modality_loss_cutoff", 0.01),
+          weights_fn=mod.targets_weights_fn)
+      recon_loss = numerator / denominator
+    else:
+      raise ValueError("internal loss only supports specific modalities.")
+    tf.summary.scalar("recon_extra", recon_loss)
+    return recon_loss
+
+  def get_sampled_frame(self, pred_frame):
+    """Samples the frame based on modality.
+
+      if the modality is L2/L1 then the next predicted frame is the
+      next frame and there is no sampling but in case of Softmax loss
+      the next actual frame should be sampled from predicted frame.
+
+      This enables multi-frame target prediction with Softmax loss.
+
+    Args:
+      pred_frame: predicted frame.
+
+    Returns:
+      sampled frame.
+
+    """
+    if not self.is_per_pixel_softmax:
+      return pred_frame
+    frame_shape = common_layers.shape_list(pred_frame)
+    target_shape = frame_shape[:-1] + [self.hparams.problem.num_channels]
+    sampled_frame = tf.reshape(pred_frame, target_shape + [256])
+    # TODO(lukaszkaiser): should this be argmax or real sampling.
+    sampled_frame = tf.argmax(sampled_frame, axis=-1)
+    sampled_frame = tf.to_float(sampled_frame)
+    # TODO(lukaszkaiser): this should be consistent with modality.bottom()
+    sampled_frame = common_layers.standardize_images(sampled_frame)
+    return sampled_frame
+
+  def __get_next_inputs(self, index, all_frames, all_actions, all_rewards):
+    """Get inputs for next prediction iteration.
+
+      If the model is recurrent then the inputs of the models are
+      the current inputs. For non-recurrent models the input is the
+      last N stacked frames/actions/rewards.
+
+    Args:
+      index: current prediction index. from 0 to number of target frames.
+      all_frames: list of all frames including input and target frames.
+      all_actions: list of all actions including input and target actions.
+      all_rewards: list of all rewards including input and target rewards.
+
+    Returns:
+      frames: input frames for next_frame prediction.
+      actions: input actions for next_frame prediction.
+      rewards: input rewards for next_frame prediction.
+      target_index: index of target frame in all_frames list.
+    """
+    i = index
+    j = i + self.hparams.video_num_input_frames
+
+    actions, rewards = None, None
+    if self.is_recurrent_model:
+      frames = all_frames[i]
+      target_index = i+1
+      if self.has_actions:
+        actions = all_actions[i]
+      if self.has_rewards:
+        rewards = all_rewards[i]
+    else:
+      frames = all_frames[i:j]
+      target_index = j
+      if self.has_actions:
+        actions = all_actions[i:j]
+        actions = tf.concat(actions, axis=1)
+      if self.has_rewards:
+        rewards = all_rewards[i:j]
+        rewards = tf.concat(rewards, axis=1)
+    return frames, actions, rewards, target_index
+
+  def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
+    """Produce predictions from the model by running it."""
+    del args, kwargs
+    # Inputs and features preparation needed to handle edge cases.
+    if not features:
+      features = {}
+    hparams = self.hparams
+    inputs_old = None
+    if "inputs" in features and len(features["inputs"].shape) < 4:
+      inputs_old = features["inputs"]
+      features["inputs"] = tf.expand_dims(features["inputs"], 2)
+
+    def logits_to_samples(logits):
+      """Get samples from logits."""
+      # If the last dimension is 1 then we're using L1/L2 loss.
+      if common_layers.shape_list(logits)[-1] == 1:
+        return tf.to_int32(tf.squeeze(logits, axis=-1))
+      # Argmax in TF doesn't handle more than 5 dimensions yet.
+      logits_shape = common_layers.shape_list(logits)
+      argmax = tf.argmax(tf.reshape(logits, [-1, logits_shape[-1]]), axis=-1)
+      return tf.reshape(argmax, logits_shape[:-1])
+
+    # Get predictions.
+    try:
+      num_channels = hparams.problem.num_channels
+    except AttributeError:
+      num_channels = 1
+    if "inputs" in features:
+      inputs_shape = common_layers.shape_list(features["inputs"])
+      targets_shape = [inputs_shape[0], hparams.video_num_target_frames,
+                       inputs_shape[2], inputs_shape[3], num_channels]
+    else:
+      tf.logging.warn("Guessing targets shape as no inputs are given.")
+      targets_shape = [hparams.batch_size,
+                       hparams.video_num_target_frames, 1, 1, num_channels]
+
+    features["targets"] = tf.zeros(targets_shape, dtype=tf.int32)
+    reward_in_mod = "target_reward" in hparams.problem_hparams.target_modality
+    action_in_mod = "target_action" in hparams.problem_hparams.target_modality
+    if reward_in_mod:
+      # TODO(lukaszkaiser): this is a hack. get the actual reward history.
+      if "input_reward" not in hparams.problem_hparams.target_modality:
+        features["input_reward"] = tf.zeros(
+            [inputs_shape[0], inputs_shape[1], 1], dtype=tf.int32)
+      features["target_reward"] = tf.zeros(
+          [targets_shape[0], targets_shape[1], 1], dtype=tf.int32)
+    if action_in_mod and "target_action" not in features:
+      features["target_action"] = tf.zeros(
+          [targets_shape[0], targets_shape[1], 1], dtype=tf.int32)
+    logits, _ = self(features)  # pylint: disable=not-callable
+    if isinstance(logits, dict):
+      results = {}
+      for k, v in six.iteritems(logits):
+        results[k] = logits_to_samples(v)
+        results["%s_logits" % k] = v
+    else:
+      results = logits_to_samples(logits)
+
+    # Restore inputs to not confuse Estimator in edge cases.
+    if inputs_old is not None:
+      features["inputs"] = inputs_old
+
+    # Return results.
+    return results
+
+  def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
+    """Main video processing function."""
+
+    # TODO(lukaszkaiser): the split axes and the argmax below heavily depend on
+    # using the default (a bit strange) video modality - we should change that.
+
+    hparams = self.hparams
+    orig_frame_shape = common_layers.shape_list(all_frames[0])
+    batch_size = orig_frame_shape[0]
+    ss_func = self.get_scheduled_sample_func(batch_size)
+    extra_loss = 0.0
+    internal_states = None
+
+    # Any extra info required by the model goes into here.
+    video_features = self.video_features(
+        all_frames, all_actions, all_rewards, all_raw_frames)
+
+    num_frames = len(all_frames)
+    if self.is_recurrent_model:
+      input_index_range = range(num_frames - 1)
+    else:
+      input_index_range = range(hparams.video_num_target_frames)
+
+    res_frames, sampled_frames, res_rewards = [], [], []
+    for i in input_index_range:
+      frames, actions, rewards, target_index = self.__get_next_inputs(
+          i, all_frames, all_actions, all_rewards)
+      target_frame = all_frames[target_index]
+
+      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+        func_in = (frames, actions, rewards, target_frame,
+                   internal_states, video_features)
+        func_out = self.next_frame(*func_in)
+        res_frame, res_reward, res_extra_loss, internal_states = func_out
+        res_frames.append(res_frame)
+        res_rewards.append(res_reward)
+        extra_loss += res_extra_loss / float(hparams.video_num_target_frames)
+
+      # Only for Softmax loss: sample frame so we can keep iterating.
+      sampled_frame = self.get_sampled_frame(res_frame)
+      sampled_frames.append(sampled_frame)
+
+      if self.is_predicting:
+        all_frames[target_index] = sampled_frame
+
+      # Scheduled sampling during training.
+      if self.is_training:
+        if self.is_recurrent_model:
+          done_warm_start = i > hparams.video_num_input_frames - 1
+        else:
+          done_warm_start = True  # Always true for non-reccurent networks.
+        groundtruth_items = [target_frame]
+        generated_items = [sampled_frame]
+        ss_frame, = self.get_scheduled_sample_inputs(
+            done_warm_start, groundtruth_items, generated_items, ss_func)
+        all_frames[target_index] = ss_frame
+
+    if self.is_recurrent_model:
+      has_input_predictions = hparams.video_num_input_frames > 1
+      if self.is_training and hparams.internal_loss and has_input_predictions:
+        # add the loss for input frames as well.
+        extra_gts = all_frames[1:hparams.video_num_input_frames]
+        extra_raw_gts = all_raw_frames[1:hparams.video_num_input_frames]
+        extra_pds = res_frames[:hparams.video_num_input_frames-1]
+        recon_loss = self.get_extra_internal_loss(
+            extra_raw_gts, extra_gts, extra_pds)
+        extra_loss += recon_loss
+      # Cut the predicted input frames.
+      res_frames = res_frames[hparams.video_num_input_frames-1:]
+      res_rewards = res_rewards[hparams.video_num_input_frames-1:]
+
+    output_frames = tf.stack(res_frames, axis=1)
+    targets = output_frames
+    if self.has_rewards:
+      output_rewards = tf.stack(res_rewards, axis=1)
+      targets = {"targets": output_frames, "target_reward": output_rewards}
+
+    return targets, extra_loss
+
+  def body(self, features):
+    self.has_actions = "input_action" in features
+    self.has_rewards = "target_reward" in features
+    hparams = self.hparams
+
+    def merge(inputs, targets):
+      """Split inputs and targets into lists."""
+      inputs = tf.unstack(inputs, axis=1)
+      targets = tf.unstack(targets, axis=1)
+      assert len(inputs) == hparams.video_num_input_frames
+      assert len(targets) == hparams.video_num_target_frames
+      return inputs + targets
+
+    frames = merge(features["inputs"], features["targets"])
+    frames_raw = merge(features["inputs_raw"], features["targets_raw"])
+    actions, rewards = None, None
+    if self.has_actions:
+      actions = merge(features["input_action"], features["target_action"])
+    if self.has_rewards:
+      rewards = merge(features["input_reward"], features["target_reward"])
+    return self.__process(frames, actions, rewards, frames_raw)
+
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 80c8feb7c..28430e1d1 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -19,15 +19,12 @@
 from __future__ import division
 from __future__ import print_function
 
-from functools import partial
-import six
-
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
+from tensor2tensor.models.video import base
 from tensor2tensor.models.video import basic_deterministic_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
-from tensor2tensor.utils import t2t_model
 
 import tensorflow as tf
 
@@ -37,176 +34,19 @@
 
 
 @registry.register_model
-class NextFrameBasicDeterministic(t2t_model.T2TModel):
+class NextFrameBasicDeterministic(base.NextFrameBase):
   """Basic next-frame model, may take actions and predict rewards too."""
 
   @property
-  def _target_modality(self):
-    # TODO(mbz): get rid of this somehow.
-    modality = self.hparams.problem_hparams.target_modality["targets"]
-    return modality.__class__.__name__
-
-  @property
-  def is_per_pixel_softmax(self):
-    return self._target_modality == "VideoModality"
-
-  def get_iteration_num(self):
-    step_num = tf.train.get_global_step()
-    # TODO(lukaszkaiser): what should it be if it's undefined?
-    if step_num is None:
-      step_num = 10000000
-    return step_num
+  def is_recurrent_model(self):
+    return False
 
-  def inject_latent(self, layer, features, filters):
-    """Do nothing for deterministic model."""
-    del features, filters
+  def inject_latent(self, layer, inputs, target):
+    del inputs, target
     return layer, 0.0
 
-  def get_scheduled_sample_func(self, batch_size):
-    """Creates a function for scheduled sampling based on given hparams."""
-    with tf.variable_scope("scheduled_sampling_func", reuse=tf.AUTO_REUSE):
-      iter_num = self.get_iteration_num()
-
-      # Simple function to bypass scheduled sampling in gt or pred only modes.
-      def scheduled_sampling_simple(ground_truth_x, generated_x,
-                                    batch_size, scheduled_sample_var):
-        del batch_size
-        if scheduled_sample_var:
-          return ground_truth_x
-        return generated_x
-
-      mode = self.hparams.scheduled_sampling_mode
-      if mode == "ground_truth_only":
-        scheduled_sampling_func = scheduled_sampling_simple
-        scheduled_sampling_func_var = True
-      elif mode == "prediction_only":
-        scheduled_sampling_func = scheduled_sampling_simple
-        scheduled_sampling_func_var = False
-      elif mode == "prob":
-        decay_steps = self.hparams.scheduled_sampling_decay_steps
-        probability = tf.train.polynomial_decay(
-            1.0, iter_num, decay_steps, 0.0)
-        scheduled_sampling_func = common_video.scheduled_sample_prob
-        scheduled_sampling_func_var = probability
-      elif mode == "prob_inverse_exp":
-        decay_steps = self.hparams.scheduled_sampling_decay_steps
-        probability = common_layers.inverse_exp_decay(
-            decay_steps, step=iter_num)
-        probability *= self.hparams.scheduled_sampling_max_prob
-        probability = 1.0 - probability
-        scheduled_sampling_func = common_video.scheduled_sample_prob
-        scheduled_sampling_func_var = probability
-      elif mode == "prob_inverse_lin":
-        decay_steps = self.hparams.scheduled_sampling_decay_steps
-        probability = common_layers.inverse_exp_decay(
-            decay_steps // 4, step=iter_num)  # Very low at start.
-        probability *= common_layers.inverse_lin_decay(
-            decay_steps, step=iter_num)
-        probability *= self.hparams.scheduled_sampling_max_prob
-        probability = 1.0 - probability
-        scheduled_sampling_func = common_video.scheduled_sample_prob
-        scheduled_sampling_func_var = probability
-      elif mode == "count":
-        # Calculate number of ground-truth frames to pass in.
-        k = self.hparams.scheduled_sampling_k
-        num_ground_truth = tf.to_int32(
-            tf.round(
-                tf.to_float(batch_size) *
-                (k / (k + tf.exp(tf.to_float(iter_num) / tf.to_float(k))))))
-        scheduled_sampling_func = common_video.scheduled_sample_count
-        scheduled_sampling_func_var = num_ground_truth
-      else:
-        raise ValueError("unknown scheduled sampling method: %s" % mode)
-
-      if isinstance(scheduled_sampling_func_var, tf.Tensor):
-        tf.summary.scalar("scheduled_sampling_var", scheduled_sampling_func_var)
-      partial_func = partial(scheduled_sampling_func,
-                             batch_size=batch_size,
-                             scheduled_sample_var=scheduled_sampling_func_var)
-      return partial_func
-
-  def get_scheduled_sample_inputs(self,
-                                  done_warm_start,
-                                  groundtruth_items,
-                                  generated_items,
-                                  scheduled_sampling_func):
-    """Scheduled sampling.
-
-    Args:
-      done_warm_start: whether we are done with warm start or not.
-      groundtruth_items: list of ground truth items.
-      generated_items: list of generated items.
-      scheduled_sampling_func: scheduled sampling function to choose between
-        groundtruth items and generated items.
-
-    Returns:
-      A mix list of ground truth and generated items.
-    """
-    def sample():
-      """Calculate the scheduled sampling params based on iteration number."""
-      with tf.variable_scope("scheduled_sampling", reuse=tf.AUTO_REUSE):
-        output_items = []
-        for item_gt, item_gen in zip(groundtruth_items, generated_items):
-          output_items.append(scheduled_sampling_func(item_gt, item_gen))
-        return output_items
-
-    cases = [
-        (tf.logical_not(done_warm_start), lambda: groundtruth_items),
-        (tf.logical_not(self.is_training), lambda: generated_items),
-    ]
-    output_items = tf.case(cases, default=sample, strict=True)
-
-    return output_items
-
-  def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
-    """Hacky code the get the loss on predicted frames from input frames.
-
-       Recurrent models consume the frames one-by-one. Therefore
-       if there is more than one input frame they also get predicted.
-       T2T only calculates loss on the predicted target frames which
-       means the loss is not being applied on the predicted input frames.
-       This code is to fix this issue. Since the model is not aware of the
-       modality it has to match the pre-porocessing happening in bottom
-       function and therefore this becomes a very hacky code. This code
-       should match the bottom and top and loss of modalities otherwise
-       it will calculate the wrong loss.
-
-    Args:
-      extra_raw_gts: extra raw ground truth frames.
-      extra_gts: extra normalized ground truth frames.
-      extra_pds: extra predicted frames.
-
-    Returns:
-      Additional reconstruction loss.
-
-    Raises:
-      ValueError: in case of unknown modality.
-    """
-    if self._target_modality == "VideoModalityL2Raw":
-      recon_loss = tf.losses.mean_squared_error(extra_gts, extra_pds)
-    elif self._target_modality == "VideoModality":
-      shape = common_layers.shape_list(extra_pds)
-      updated_shape = shape[:-1] + [3, 256]
-      extra_pds = tf.reshape(extra_pds, updated_shape)
-      # Merge time and batch
-      logits = tf.reshape(extra_pds, [-1] + updated_shape[2:])
-      targets = extra_raw_gts
-      targets_shape = common_layers.shape_list(targets)
-      targets = tf.reshape(targets, [-1] + targets_shape[2:])
-      mod = self.hparams.problem_hparams.target_modality["targets"]
-      numerator, denominator = common_layers.padded_cross_entropy(
-          logits,
-          targets,
-          self.hparams.label_smoothing,
-          cutoff=getattr(self.hparams, "video_modality_loss_cutoff", 0.01),
-          weights_fn=mod.targets_weights_fn)
-      recon_loss = numerator / denominator
-    else:
-      raise ValueError("internal loss only supports specific modalities.")
-    tf.summary.scalar("recon_extra", recon_loss)
-    return recon_loss
-
   def inject_additional_input(self, layer, inputs, name, mode="concat"):
+    """Injects the additional input into the layer."""
     layer_shape = common_layers.shape_list(layer)
     input_shape = common_layers.shape_list(inputs)
     zeros_mask = tf.zeros(layer_shape, dtype=tf.float32)
@@ -231,26 +71,36 @@ def inject_additional_input(self, layer, inputs, name, mode="concat"):
 
     return layer
 
-  def get_sampled_frame(self, res_frame):
-    if not self.is_per_pixel_softmax:
-      return res_frame
-    frame_shape = common_layers.shape_list(res_frame)
-    target_shape = frame_shape[:-1] + [self.hparams.problem.num_channels]
-    sampled_frame = tf.reshape(res_frame, target_shape + [256])
-    sampled_frame = tf.argmax(sampled_frame, axis=-1)
-    sampled_frame = tf.to_float(sampled_frame)
-    return sampled_frame
+  def middle_network(self, layer, internal_states):
+    # Run a stack of convolutions.
+    x = layer
+    kernel1 = (3, 3)
+    filters = common_layers.shape_list(x)[-1]
+    for i in range(self.hparams.num_hidden_layers):
+      with tf.variable_scope("layer%d" % i):
+        y = tf.nn.dropout(x, 1.0 - self.hparams.dropout)
+        y = tf.layers.conv2d(y, filters, kernel1, activation=common_layers.belu,
+                             strides=(1, 1), padding="SAME")
+        if i == 0:
+          x = y
+        else:
+          x = common_layers.layer_norm(x + y)
+    return x, internal_states
+
+  def next_frame(self, frames, actions, rewards, target_frame,
+                 internal_states, video_extra):
+    del video_extra
 
-  def body_single(self, features):
     hparams = self.hparams
     filters = hparams.hidden_size
-    kernel1, kernel2 = (3, 3), (4, 4)
+    kernel2 = (4, 4)
 
     # Embed the inputs.
-    inputs_shape = common_layers.shape_list(features["inputs"])
+    stacked_frames = tf.concat(frames, axis=-1)
+    inputs_shape = common_layers.shape_list(stacked_frames)
     # Using non-zero bias initializer below for edge cases of uniform inputs.
     x = tf.layers.dense(
-        features["inputs"], filters, name="inputs_embed",
+        stacked_frames, filters, name="inputs_embed",
         bias_initializer=tf.random_normal_initializer(stddev=0.01))
     x = common_attention.add_timing_signal_nd(x)
 
@@ -268,30 +118,22 @@ def body_single(self, features):
         x = common_layers.layer_norm(x)
 
     # Add embedded action if present.
-    if "input_action" in features:
-      action = features["input_action"][:, -1, :]
+    if self.has_actions:
+      action = actions[:, -1, :]
       x = self.inject_additional_input(
           x, action, "action_enc", hparams.action_injection)
 
-    x, extra_loss = self.inject_latent(x, features, filters)
-    x_mid = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
+    # Inject latent if present. Only for stochastic models.
+    x, extra_loss = self.inject_latent(x, frames, target_frame)
 
-    # Run a stack of convolutions.
-    for i in range(hparams.num_hidden_layers):
-      with tf.variable_scope("layer%d" % i):
-        y = tf.nn.dropout(x, 1.0 - hparams.dropout)
-        y = tf.layers.conv2d(y, filters, kernel1, activation=common_layers.belu,
-                             strides=(1, 1), padding="SAME")
-        if i == 0:
-          x = y
-        else:
-          x = common_layers.layer_norm(x + y)
+    x_mid = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
+    x, internal_states = self.middle_network(x, internal_states)
 
     # Up-convolve.
     layer_inputs = list(reversed(layer_inputs))
     for i in range(hparams.num_compress_steps):
       with tf.variable_scope("upstride%d" % i):
-        if "input_action" in features:
+        if self.has_actions:
           x = self.inject_additional_input(
               x, action, "action_enc", hparams.action_injection)
         if i >= hparams.num_compress_steps - hparams.filter_double_steps:
@@ -314,148 +156,13 @@ def body_single(self, features):
       x = tf.layers.dense(x, hparams.problem.num_channels, name="logits")
 
     # No reward prediction if not needed.
-    if "target_reward" not in features:
-      return x, extra_loss
+    if not self.has_rewards:
+      return x, None, extra_loss, internal_states
 
     # Reward prediction based on middle and final logits.
     reward_pred = tf.concat([x_mid, x_fin], axis=-1)
     reward_pred = tf.nn.relu(tf.layers.dense(
         reward_pred, 128, name="reward_pred"))
     reward_pred = tf.expand_dims(reward_pred, axis=3)  # Need fake channels dim.
-    return {"targets": x, "target_reward": reward_pred}, extra_loss
-
-  def body(self, features):
-    hparams = self.hparams
-    is_predicting = hparams.mode == tf.estimator.ModeKeys.PREDICT
-
-    # TODO(lukaszkaiser): the split axes and the argmax below heavily depend on
-    # using the default (a bit strange) video modality - we should change that.
-
-    # Split inputs and targets into lists.
-    input_frames = tf.unstack(features["inputs"], axis=1)
-    target_frames = tf.unstack(features["targets"], axis=1)
-    all_frames = input_frames + target_frames
-    if "input_action" in features:
-      input_actions = list(tf.split(
-          features["input_action"], hparams.video_num_input_frames, axis=1))
-      target_actions = list(tf.split(
-          features["target_action"], hparams.video_num_target_frames, axis=1))
-      all_actions = input_actions + target_actions
-
-    orig_frame_shape = common_layers.shape_list(all_frames[0])
-    batch_size = orig_frame_shape[0]
-    ss_func = self.get_scheduled_sample_func(batch_size)
-
-    # Run a number of steps.
-    res_frames, sampled_frames, sampled_frames_raw = [], [], []
-    extra_loss = 0.0
-    if "target_reward" in features:
-      res_rewards = []
-    for i in range(hparams.video_num_target_frames):
-      cur_frames = all_frames[i:i + hparams.video_num_input_frames]
-      features["inputs"] = tf.concat(cur_frames, axis=-1)
-      features["cur_target_frame"] = all_frames[
-          i + hparams.video_num_input_frames]
-      if "input_action" in features:
-        cur_actions = all_actions[i:i + hparams.video_num_input_frames]
-        features["input_action"] = tf.concat(cur_actions, axis=1)
-
-      # Run model.
-      with tf.variable_scope(tf.get_variable_scope(), reuse=i > 0):
-        if "target_reward" not in features:
-          res_frame, res_extra_loss = self.body_single(features)
-        else:
-          res_dict, res_extra_loss = self.body_single(features)
-          res_frame = res_dict["targets"]
-          res_reward = res_dict["target_reward"]
-          res_rewards.append(res_reward)
-      extra_loss += res_extra_loss / float(hparams.video_num_target_frames)
-      res_frames.append(res_frame)
-
-      # Only for Softmax loss: sample frame so we can keep iterating.
-      sampled_frame_raw = self.get_sampled_frame(res_frame)
-      sampled_frames_raw.append(sampled_frame_raw)
-      # TODO(lukaszkaiser): this should be consistent with modality.bottom()
-      sampled_frame = common_layers.standardize_images(sampled_frame_raw)
-      sampled_frames.append(sampled_frame)
-
-      if is_predicting:
-        all_frames[i + hparams.video_num_input_frames] = sampled_frame
-
-      # Scheduled sampling during training.
-      if self.is_training:
-        done_warm_start = True  # Always true for non-reccurent networks.
-        groundtruth_items = [all_frames[i + hparams.video_num_input_frames]]
-        generated_items = [sampled_frame]
-        ss_frame, = self.get_scheduled_sample_inputs(
-            done_warm_start, groundtruth_items, generated_items, ss_func)
-        all_frames[i + hparams.video_num_input_frames] = ss_frame
-
-    # Concatenate results and return them.
-    frames = tf.stack(res_frames, axis=1)
-
-    if "target_reward" not in features:
-      return frames
-    rewards = tf.concat(res_rewards, axis=1)
-    return {"targets": frames, "target_reward": rewards}, extra_loss
-
-  def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
-    """Produce predictions from the model by running it."""
-    del args, kwargs
-    # Inputs and features preparation needed to handle edge cases.
-    if not features:
-      features = {}
-    hparams = self.hparams
-    inputs_old = None
-    if "inputs" in features and len(features["inputs"].shape) < 4:
-      inputs_old = features["inputs"]
-      features["inputs"] = tf.expand_dims(features["inputs"], 2)
-
-    def logits_to_samples(logits):
-      """Get samples from logits."""
-      # If the last dimension is 1 then we're using L1/L2 loss.
-      if common_layers.shape_list(logits)[-1] == 1:
-        return tf.to_int32(tf.squeeze(logits, axis=-1))
-      # Argmax in TF doesn't handle more than 5 dimensions yet.
-      logits_shape = common_layers.shape_list(logits)
-      argmax = tf.argmax(tf.reshape(logits, [-1, logits_shape[-1]]), axis=-1)
-      return tf.reshape(argmax, logits_shape[:-1])
-
-    # Get predictions.
-    try:
-      num_channels = hparams.problem.num_channels
-    except AttributeError:
-      num_channels = 1
-    if "inputs" in features:
-      inputs_shape = common_layers.shape_list(features["inputs"])
-      targets_shape = [inputs_shape[0], hparams.video_num_target_frames,
-                       inputs_shape[2], inputs_shape[3], num_channels]
-    else:
-      tf.logging.warn("Guessing targets shape as no inputs are given.")
-      targets_shape = [hparams.batch_size,
-                       hparams.video_num_target_frames, 1, 1, num_channels]
-
-    features["targets"] = tf.zeros(targets_shape, dtype=tf.int32)
-    reward_in_mod = "target_reward" in hparams.problem_hparams.target_modality
-    action_in_mod = "target_action" in hparams.problem_hparams.target_modality
-    if reward_in_mod:
-      features["target_reward"] = tf.zeros(
-          [targets_shape[0], 1, 1], dtype=tf.int32)
-    if action_in_mod and "target_action" not in features:
-      features["target_action"] = tf.zeros(
-          [targets_shape[0], 1, 1], dtype=tf.int32)
-    logits, _ = self(features)  # pylint: disable=not-callable
-    if isinstance(logits, dict):
-      results = {}
-      for k, v in six.iteritems(logits):
-        results[k] = logits_to_samples(v)
-        results["%s_logits" % k] = v
-    else:
-      results = logits_to_samples(logits)
-
-    # Restore inputs to not confuse Estimator in edge cases.
-    if inputs_old is not None:
-      features["inputs"] = inputs_old
+    return x, reward_pred, extra_loss, internal_states
 
-    # Return results.
-    return results
diff --git a/tensor2tensor/models/video/basic_recurrent.py b/tensor2tensor/models/video/basic_recurrent.py
index 0b0af311d..8c14a8c8a 100644
--- a/tensor2tensor/models/video/basic_recurrent.py
+++ b/tensor2tensor/models/video/basic_recurrent.py
@@ -19,8 +19,6 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.layers import common_attention
-from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.video import basic_stochastic
 from tensor2tensor.utils import registry
@@ -37,153 +35,23 @@ class NextFrameBasicRecurrent(
     basic_stochastic.NextFrameBasicStochasticDiscrete):
   """Basic next-frame recurrent model."""
 
-  def predict_next_frame(self, frame, action, lstm_states):
-    hparams = self.hparams
-    filters = hparams.hidden_size
-    kernel1, kernel2 = (3, 3), (4, 4)
-    lstm_func = common_video.conv_lstm_2d
-
-    # Embed the inputs.
-    inputs_shape = common_layers.shape_list(frame)
-    # Using non-zero bias initializer below for edge cases of uniform inputs.
-    x = tf.layers.dense(
-        frame, filters, name="inputs_embed",
-        bias_initializer=tf.random_normal_initializer(stddev=0.01))
-    x = common_attention.add_timing_signal_nd(x)
-
-    # Down-stride.
-    layer_inputs = [x]
-    for i in range(hparams.num_compress_steps):
-      with tf.variable_scope("downstride%d" % i):
-        layer_inputs.append(x)
-        x = common_layers.make_even_size(x)
-        if i < hparams.filter_double_steps:
-          filters *= 2
-        x = common_attention.add_timing_signal_nd(x)
-        x = tf.layers.conv2d(x, filters, kernel2, activation=common_layers.belu,
-                             strides=(2, 2), padding="SAME")
-        x = common_layers.layer_norm(x)
+  @property
+  def is_recurrent_model(self):
+    return True
 
-    # Add embedded action if present.
-    if self.has_action:
-      x = self.inject_additional_input(
-          x, action, "action_enc", hparams.action_injection)
+  def middle_network(self, layer, internal_states):
+    lstm_func = common_video.conv_lstm_2d
+    hp = self.hparams
 
-    x, extra_loss = self.inject_latent(x, self.features, filters)
+    lstm_states = internal_states
+    if lstm_states is None:
+      lstm_states = [None] * hp.num_lstm_layers
 
     # LSTM layers
-    for j in range(hparams.num_lstm_layers):
-      x, lstm_states[j] = lstm_func(x, lstm_states[j], hparams.num_lstm_filters)
-
-    # Run a stack of convolutions.
-    for i in range(hparams.num_hidden_layers):
-      with tf.variable_scope("layer%d" % i):
-        y = tf.nn.dropout(x, 1.0 - hparams.dropout)
-        y = tf.layers.conv2d(y, filters, kernel1, activation=common_layers.belu,
-                             strides=(1, 1), padding="SAME")
-        if i == 0:
-          x = y
-        else:
-          x = common_layers.layer_norm(x + y)
-
-    # Up-convolve.
-    layer_inputs = list(reversed(layer_inputs))
-    for i in range(hparams.num_compress_steps):
-      with tf.variable_scope("upstride%d" % i):
-        if self.has_action:
-          x = self.inject_additional_input(
-              x, action, "action_enc", hparams.action_injection)
-        if i >= hparams.num_compress_steps - hparams.filter_double_steps:
-          filters //= 2
-        x = tf.layers.conv2d_transpose(
-            x, filters, kernel2, activation=common_layers.belu,
-            strides=(2, 2), padding="SAME")
-        y = layer_inputs[i]
-        shape = common_layers.shape_list(y)
-        x = x[:, :shape[1], :shape[2], :]
-        x = common_layers.layer_norm(x + y)
-        x = common_attention.add_timing_signal_nd(x)
-
-    # Cut down to original size.
-    x = x[:, :inputs_shape[1], :inputs_shape[2], :]
-    if self.is_per_pixel_softmax:
-      x = tf.layers.dense(x, hparams.problem.num_channels * 256, name="logits")
-    else:
-      x = tf.layers.dense(x, hparams.problem.num_channels, name="logits")
-
-    # Reward prediction if needed.
-    reward_pred = 0.0
-    if self.has_reward:
-      reward_pred = tf.expand_dims(  # Add a fake channels dim.
-          tf.reduce_mean(x, axis=[1, 2], keepdims=True), axis=3)
-    return x, reward_pred, extra_loss, lstm_states
-
-  def body(self, features):
-    hparams = self.hparams
-    self.has_action = "input_action" in features
-    self.has_reward = "target_reward" in features
-    # dirty hack to enable the latent tower
-    self.features = features
-
-    # Split inputs and targets into lists.
-    input_frames = tf.unstack(features["inputs"], axis=1)
-    target_frames = tf.unstack(features["targets"], axis=1)
-    all_frames = input_frames + target_frames
-    if self.has_action:
-      input_actions = tf.unstack(features["input_action"], axis=1)
-      target_actions = tf.unstack(features["target_action"], axis=1)
-      all_actions = input_actions + target_actions
-
-    res_frames, sampled_frames, sampled_frames_raw, res_rewards = [], [], [], []
-    lstm_states = [None] * hparams.num_lstm_layers
-    extra_loss = 0.0
-
-    num_frames = len(all_frames)
-    for i in range(num_frames - 1):
-      frame = all_frames[i]
-      action = all_actions[i] if self.has_action else None
-
-      # more hack to enable latent_tower
-      # TODO(mbz): clean this up.
-      self.features["inputs"] = all_frames[i]
-      self.features["cur_target_frame"] = all_frames[i+1]
-
-      # Run model.
-      with tf.variable_scope("recurrent_model", reuse=tf.AUTO_REUSE):
-        func_out = self.predict_next_frame(frame, action, lstm_states)
-        res_frame, res_reward, res_extra_loss, lstm_states = func_out
-        res_frames.append(res_frame)
-        res_rewards.append(res_reward)
-        extra_loss += res_extra_loss
-
-      sampled_frame_raw = self.get_sampled_frame(res_frame)
-      sampled_frames_raw.append(sampled_frame_raw)
-      # TODO(lukaszkaiser): this should be consistent with modality.bottom()
-      sampled_frame = common_layers.standardize_images(sampled_frame_raw)
-      sampled_frames.append(sampled_frame)
-
-      # Only for Softmax loss: sample next frame so we can keep iterating.
-      if self.is_predicting and i >= hparams.video_num_input_frames:
-        all_frames[i+1] = sampled_frame
-
-    # Concatenate results and return them.
-    output_frames = res_frames[hparams.video_num_input_frames-1:]
-    frames = tf.stack(output_frames, axis=1)
-
-    has_input_predictions = hparams.video_num_input_frames > 1
-    if self.is_training and hparams.internal_loss and has_input_predictions:
-      # add the loss for input frames as well.
-      extra_gts = input_frames[1:]
-      extra_pds = res_frames[:hparams.video_num_input_frames-1]
-      extra_raw_gts = features["inputs_raw"][:, 1:]
-      recon_loss = self.get_extra_internal_loss(
-          extra_raw_gts, extra_gts, extra_pds)
-      extra_loss += recon_loss
-
-    if not self.has_reward:
-      return frames, extra_loss
-    rewards = tf.concat(res_rewards[hparams.video_num_input_frames-1:], axis=1)
-    return {"targets": frames, "target_reward": rewards}, extra_loss
+    x = layer
+    for j in range(hp.num_lstm_layers):
+      x, lstm_states[j] = lstm_func(x, lstm_states[j], hp.num_lstm_filters)
+    return x, lstm_states
 
 
 @registry.register_hparams
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 84d7e450b..36b5ec0d9 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -48,12 +48,11 @@ class NextFrameBasicStochastic(
     base_vae.NextFrameBaseVae):
   """Stochastic version of basic next-frame model."""
 
-  def inject_latent(self, layer, features, filters):
+  def inject_latent(self, layer, inputs, target):
     """Inject a VAE-style latent."""
     # Latent for stochastic model
-    input_frames = tf.to_float(features["inputs_raw"])
-    target_frames = tf.to_float(features["targets_raw"])
-    full_video = tf.concat([input_frames, target_frames], axis=1)
+    filters = 128
+    full_video = tf.stack(inputs + [target], axis=1)
     latent_mean, latent_std = self.construct_latent_tower(
         full_video, time_axis=1)
     latent = common_video.get_gaussian_tensor(latent_mean, latent_std)
@@ -73,9 +72,8 @@ class NextFrameBasicStochasticDiscrete(
     basic_deterministic.NextFrameBasicDeterministic):
   """Basic next-frame model with a tiny discrete latent."""
 
-  def inject_latent(self, layer, features, filters):
+  def inject_latent(self, layer, inputs, target):
     """Inject a deterministic latent based on the target frame."""
-    del filters
     hparams = self.hparams
     final_filters = common_layers.shape_list(layer)[-1]
     filters = hparams.hidden_size
@@ -120,8 +118,7 @@ def add_d(layer, d):
       return add_d(layer, d), 0.0
 
     # Embed.
-    frames = tf.concat(
-        [features["cur_target_frame"], features["inputs"]], axis=-1)
+    frames = tf.concat([inputs, target], axis=-1)
     x = tf.layers.dense(
         frames, filters, name="latent_embed",
         bias_initializer=tf.random_normal_initializer(stddev=0.01))

From ac43578a56612386d6dd08c0f5840db98acf4c1a Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 16 Oct 2018 00:39:41 +0200
Subject: [PATCH 1009/2720] Per epoch sharding in T2TEnv (#1144)

* Per epoch sharding

* Add some tests for sharding

* Update trainer_model_based_new

* Fix trainer_model_based
---
 tensor2tensor/data_generators/gym_env.py      | 176 +++++++++++++-----
 tensor2tensor/data_generators/gym_env_test.py | 103 ++++++++--
 tensor2tensor/rl/trainer_model_based_new.py   |  18 +-
 3 files changed, 230 insertions(+), 67 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 0a5bf9a5e..0c4a90e93 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -20,10 +20,11 @@
 from __future__ import print_function
 
 import collections
-import random
+import itertools
 
 from gym.spaces import Box
 import numpy as np
+import six
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -76,9 +77,11 @@ def __init__(self, batch_size):
 
     self.clear_history()
     self.batch_size = batch_size
-    self._current_rollouts = [[] for _ in range(batch_size)]
-    self._current_frames = [None for _ in range(batch_size)]
-    self.rollouts_by_epoch = dict()
+    self._current_batch_frames = [None for _ in range(batch_size)]
+    self._current_batch_rollouts = [[] for _ in range(batch_size)]
+    self._current_epoch_rollouts = []
+    self._rollouts_by_epoch_and_split = collections.OrderedDict()
+    self.current_epoch = None
     with tf.Graph().as_default() as tf_graph:
       self._tf_graph = _Noncopyable(tf_graph)
       self._image_t = _Noncopyable(
@@ -95,19 +98,34 @@ def __str__(self):
 
   def clear_history(self):
     """Clears the rollout history."""
-    self.rollouts_by_epoch = dict()
+    self._rollouts_by_epoch_and_split = collections.OrderedDict()
 
   def start_new_epoch(self, epoch):
     if not isinstance(epoch, int):
       raise ValueError("Epoch should be integer, got {}".format(epoch))
-    if epoch in self.rollouts_by_epoch:
+    if epoch in self._rollouts_by_epoch_and_split:
       raise ValueError("Epoch {} already registered".format(epoch))
     self.current_epoch = epoch
-    self.rollouts_by_epoch[epoch] = list()
-
-  @property
-  def current_epoch_rollouts(self):
-    return self.rollouts_by_epoch[self.current_epoch]
+    self._current_epoch_rollouts = []
+
+  def current_epoch_rollouts(self, split=None):
+    try:
+      rollouts_by_split = self._rollouts_by_epoch_and_split[self.current_epoch]
+    except KeyError:
+      if split is not None:
+        raise ValueError(
+            "generate_data() should first be called in the current epoch"
+        )
+      else:
+        return self._current_epoch_rollouts
+    if split is not None:
+      return rollouts_by_split[split]
+    else:
+      return [
+          rollout
+          for rollouts in rollouts_by_split.values()
+          for rollout in rollouts
+      ]
 
   def _preprocess_observations(self, obs):
     """Transforms a batch of observations.
@@ -161,20 +179,18 @@ def step(self, actions):
       (obs, rewards, dones) - batches of observations, rewards and done flags
       respectively.
     """
-    if not self.rollouts_by_epoch:
-      self.start_new_epoch(0)
     (obs, unclipped_rewards, dones) = self._step(actions)
     obs = self._preprocess_observations(obs)
     (min_reward, max_reward) = self.reward_range
     rewards = np.around(np.clip(unclipped_rewards, min_reward, max_reward))
     encoded_obs = self._encode_observations(obs)
     for (rollout, frame, action) in zip(
-        self._current_rollouts, self._current_frames, actions
+        self._current_batch_rollouts, self._current_batch_frames, actions
     ):
       rollout.append(frame._replace(action=action))
 
     # orud = (observation, reward, unclipped_reward, done)
-    self._current_frames = [
+    self._current_batch_frames = [
         Frame(*orud, action=None)
         for orud in zip(encoded_obs, rewards, unclipped_rewards, dones)
     ]
@@ -204,21 +220,27 @@ def reset(self, indices=None):
     Returns:
       Batch of initial observations of reset environments.
     """
-    if not self.rollouts_by_epoch:
+    if self.current_epoch is None:
+      # It's here so that the old pipeline works.
       self.start_new_epoch(0)
+      # TODO(koz4k): Replace with:
+      # raise ValueError(
+      #     "No current epoch. start_new_epoch() should first be called."
+      # )
+
     if indices is None:
       indices = np.arange(self.batch_size)
     new_obs = self._reset(indices)
     new_obs = self._preprocess_observations(new_obs)
     encoded_obs = self._encode_observations(new_obs)
     for (index, ob) in zip(indices, encoded_obs):
-      frame = self._current_frames[index]
+      frame = self._current_batch_frames[index]
       if frame is not None:
-        rollout = self._current_rollouts[index]
+        rollout = self._current_batch_rollouts[index]
         rollout.append(frame._replace(action=0))
-        self.current_epoch_rollouts.append(rollout)
-        self._current_rollouts[index] = []
-      self._current_frames[index] = Frame(
+        self._current_epoch_rollouts.append(rollout)
+        self._current_batch_rollouts[index] = []
+      self._current_batch_frames[index] = Frame(
           observation=ob, reward=0, unclipped_reward=0, done=False, action=None
       )
     return new_obs
@@ -289,8 +311,8 @@ def make_modality(name):
     p.input_space_id = problem.SpaceID.IMAGE
     p.target_space_id = problem.SpaceID.IMAGE
 
-  def _generate_frames(self, epoch_rollout_tuples):
-    for epoch, rollout in epoch_rollout_tuples:
+  def _generate_frames(self, epoch, rollouts):
+    for rollout in rollouts:
       for (frame_number, frame) in enumerate(rollout):
         yield {
             "frame_number": [frame_number],
@@ -304,38 +326,106 @@ def _generate_frames(self, epoch_rollout_tuples):
             "done": [int(frame.done)]
         }
 
-  def generate_data(self, data_dir, tmp_dir, task_id=-1):
-    """Saves the rollout history to disk."""
-    # Shuffle rollouts globally taking advantage of the fact that we have
-    # everything in memory.
-    epoch_rollout_tuples = list()
-    for epoch_nr, rollouts in self.rollouts_by_epoch.items():
-      for rollout in rollouts:
-        epoch_rollout_tuples.append((epoch_nr, rollout))
+  @staticmethod
+  def _calc_num_frames(rollouts):
+    return sum(len(rollout) for rollout in rollouts)
 
-    random.shuffle(epoch_rollout_tuples)
+  def _split_current_epoch(self):
+    """Splits frames in the current epoch according to self.dataset_splits.
 
+    Rollouts can be broken on shard boundary. This is desirable when we have
+    few long rollouts and we want to make sure we have data in the dev set.
+    """
+    num_frames = self._calc_num_frames(self._current_epoch_rollouts)
+    num_shards = sum(split["shards"] for split in self.dataset_splits)
+    shard_size = num_frames // num_shards
+
+    splits = self.dataset_splits
+    num_saved_frames = 0
+    split_index = 0
+    split_begin_index = 0
+    rollouts_by_split = collections.defaultdict(list)
+
+    def split_size(split_index):
+      return splits[split_index]["shards"] * shard_size
+
+    for rollout in self._current_epoch_rollouts:
+      num_saved_frames_current_rollout = 0
+      # Split the rollout into chunks corresponding to dataset splits. In most
+      # cases there should be only one chunk. On dataset split boundary there
+      # will be two. If a rollout is longer then the size of a dataset split,
+      # there might be more.
+      while num_saved_frames_current_rollout < len(rollout):
+        max_chunk_length = (
+            split_begin_index + split_size(split_index) - num_saved_frames
+        )
+        if split_index == len(splits) - 1:
+          # Put the remainder in the last split to preserve the ordering.
+          max_chunk_length = len(rollout)
+        rollout_chunk = rollout[
+            num_saved_frames_current_rollout:
+            (num_saved_frames_current_rollout + max_chunk_length)
+        ]
+        rollouts_by_split[splits[split_index]["split"]].append(rollout_chunk)
+        num_saved_frames_current_rollout += len(rollout_chunk)
+        num_saved_frames += len(rollout_chunk)
+
+        if num_saved_frames == split_begin_index + split_size(split_index):
+          split_begin_index += split_size(split_index)
+          split_index = min(split_index + 1, len(splits) - 1)
+
+    self._rollouts_by_epoch_and_split[self.current_epoch] = rollouts_by_split
+    self._current_epoch_rollouts = []
+
+  def splits_and_paths(self, data_dir):
     filepath_fns = {
         problem.DatasetSplit.TRAIN: self.training_filepaths,
         problem.DatasetSplit.EVAL: self.dev_filepaths,
         problem.DatasetSplit.TEST: self.test_filepaths,
     }
 
+    num_epochs = len(self._rollouts_by_epoch_and_split)
     # We set shuffled=True as we don't want to shuffle on disk later.
-    paths = [
-        path
+    return [
+        (split["split"], filepath_fns[split["split"]](
+            data_dir, split["shards"] * num_epochs, shuffled=True
+        ))
         for split in self.dataset_splits
-        for path in filepath_fns[split["split"]](
-            data_dir, split["shards"], shuffled=True
-        )
     ]
 
-    num_frames = sum(len(rollout) for (_, rollout) in epoch_rollout_tuples)
-    shard_size = num_frames // len(paths)
-    generator_utils.generate_files(
-        self._generate_frames(epoch_rollout_tuples), paths,
-        cycle_every_n=shard_size
-    )
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    """Saves the rollout history to disk.
+
+    Also splits data into train/dev sets.
+    """
+    self._split_current_epoch()
+
+    splits_and_paths = self.splits_and_paths(data_dir)
+    num_epochs = len(self._rollouts_by_epoch_and_split)
+
+    for (epoch_index, (epoch, rollouts_by_split)) in enumerate(
+        six.iteritems(self._rollouts_by_epoch_and_split)
+    ):
+      for (split, paths) in splits_and_paths:
+        num_shards = len(paths) // num_epochs
+        paths = paths[
+            epoch_index * num_shards : (epoch_index + 1) * num_shards
+        ]
+
+        rollouts = rollouts_by_split[split]
+        num_frames = self._calc_num_frames(rollouts)
+        shard_size = num_frames // len(paths)
+
+        frame_gen = self._generate_frames(epoch, rollouts)
+        for (path_index, path) in enumerate(paths):
+          limit = shard_size
+          # Put the remainder in the last shard to preserve the ordering.
+          if path_index == len(paths) - 1:
+            limit = None
+          generator_utils.generate_files(
+              itertools.islice(frame_gen, limit), [path],
+              cycle_every_n=float("inf")
+          )
 
 
 class T2TGymEnv(T2TEnv):
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 772a6baa2..a9b8000d8 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -27,6 +27,7 @@
 from gym.spaces import Discrete
 import numpy as np
 
+from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import gym_env
 
 import tensorflow as tf
@@ -63,19 +64,24 @@ def reset(self):
 
 class GymEnvTest(tf.test.TestCase):
 
-  @classmethod
-  def setUpClass(cls):
-    cls.out_dir = tf.test.get_temp_dir()
-    shutil.rmtree(cls.out_dir)
-    os.mkdir(cls.out_dir)
+  splits = (problem.DatasetSplit.TRAIN, problem.DatasetSplit.EVAL)
+
+  def setUp(self):
+    self.out_dir = tf.test.get_temp_dir()
+    shutil.rmtree(self.out_dir)
+    os.mkdir(self.out_dir)
 
   def init_batch_and_play(self, env_lambda, n_steps=1, **kwargs):
     raw_envs = [env_lambda(), env_lambda()]
     env = gym_env.T2TGymEnv(raw_envs, **kwargs)
     env.start_new_epoch(0)
+    return self.play(env, n_steps)
+
+  def play(self, env, n_steps):
     obs = list()
     rewards = list()
     obs.append(env.reset())
+    num_dones = 0
     for _ in range(n_steps):
       step_obs, step_rewards, dones = env.step(actions=[0, 0])
       obs.append(step_obs)
@@ -83,19 +89,86 @@ def init_batch_and_play(self, env_lambda, n_steps=1, **kwargs):
       for (i, done) in enumerate(dones):
         if done:
           env.reset([i])
-    return env, obs, rewards
+          num_dones += 1
+    return env, obs, rewards, num_dones
+
+  def test_splits_dataset(self):
+    env, _, _, _ = self.init_batch_and_play(TestEnv, n_steps=20)
+    env.generate_data(self.out_dir, tmp_dir=None)
+
+    for split in self.splits:
+      self.assertTrue(env.current_epoch_rollouts(split))
+
+  def test_split_preserves_number_of_rollouts(self):
+    env, _, _, num_dones = self.init_batch_and_play(TestEnv, n_steps=20)
+    env.generate_data(self.out_dir, tmp_dir=None)
+
+    num_rollouts_after_split = sum(
+        len(env.current_epoch_rollouts(split)) for split in self.splits
+    )
+    # Number of rollouts could be increased by one in case a rollout is broken
+    # on a boundary between the dataset splits.
+    self.assertGreaterEqual(num_rollouts_after_split, num_dones)
+    self.assertLessEqual(num_rollouts_after_split, num_dones + 1)
+
+  def test_split_preserves_number_of_frames(self):
+    env, _, _, num_dones = self.init_batch_and_play(TestEnv, n_steps=20)
+    env.generate_data(self.out_dir, tmp_dir=None)
+
+    num_frames = sum(
+        len(rollout)
+        for split in self.splits
+        for rollout in env.current_epoch_rollouts(split)
+    )
+    # There are 3 frames in every rollout: the initial one and two returned by
+    # step().
+    self.assertEqual(num_frames, 3 * num_dones)
 
-  def test_generates(self):
+  def test_generates_data(self):
     # This test needs base env which outputs done after two steps.
     env_lambda = TestEnv
-    env, _, _ = self.init_batch_and_play(env_lambda, n_steps=20)
+    env, _, _, _ = self.init_batch_and_play(env_lambda, n_steps=20)
     env.generate_data(self.out_dir, tmp_dir=None)
 
     filenames = os.listdir(self.out_dir)
     self.assertTrue(filenames)
-    path = os.path.join(self.out_dir, filenames[0])
-    records = list(tf.python_io.tf_record_iterator(path))
-    self.assertTrue(records)
+    for filename in filenames:
+      path = os.path.join(self.out_dir, filename)
+      records = list(tf.python_io.tf_record_iterator(path))
+      self.assertTrue(records)
+
+  def test_shards_per_epoch(self):
+    env, _, _, _ = self.init_batch_and_play(TestEnv, n_steps=20)
+    env.generate_data(self.out_dir, tmp_dir=None)
+    num_shards_per_epoch = len(os.listdir(self.out_dir))
+    shutil.rmtree(self.out_dir)
+    os.mkdir(self.out_dir)
+
+    env.start_new_epoch(1)
+    self.play(env, n_steps=20)
+    env.generate_data(self.out_dir, tmp_dir=None)
+    self.assertEqual(len(os.listdir(self.out_dir)), 2 * num_shards_per_epoch)
+
+  def test_frame_numbers_are_continuous(self):
+    env, _, _, _ = self.init_batch_and_play(TestEnv, n_steps=20)
+    env.generate_data(self.out_dir, tmp_dir=None)
+
+    frame_numbers = [
+        tf.train.Example.FromString(
+            record
+        ).features.feature['frame_number'].int64_list.value[0]
+        for (_, paths) in env.splits_and_paths(self.out_dir)
+        for path in paths
+        for record in tf.python_io.tf_record_iterator(path)
+    ]
+    last_frame_number = -1
+    for frame_number in frame_numbers:
+      # Every consecutive frame number should be either zero (first frame in
+      # a new rollout) or one bigger than the last one (next frame in the same
+      # rollout).
+      if frame_number > 0:
+        self.assertEqual(frame_number, last_frame_number + 1)
+      last_frame_number = frame_number
 
   def test_clipping(self):
     # This test needs base env with rewards out of [-1,1] range.
@@ -105,7 +178,7 @@ def test_clipping(self):
     # self.assertTrue(np.max(rewards) == 1)
     # self.assertTrue(np.min(rewards) == -1)
 
-    _, _, unclipped_rewards = self.init_batch_and_play(env_lambda, n_steps=2)
+    _, _, unclipped_rewards, _ = self.init_batch_and_play(env_lambda, n_steps=2)
     self.assertTrue(np.max(unclipped_rewards) > 1)
     self.assertTrue(np.min(unclipped_rewards) < -1)
 
@@ -115,7 +188,7 @@ def test_resize(self):
     resize_height_factor = 2
     resize_width_factor = 3
     orig_height, orig_width = orig_env.observation_space.shape[:2]
-    env, obs, _ = self.init_batch_and_play(
+    env, obs, _, _ = self.init_batch_and_play(
         env_lambda, n_steps=1,
         resize_height_factor=resize_height_factor,
         resize_width_factor=resize_width_factor)
@@ -135,10 +208,10 @@ def assert_channels(self, env, obs, n_channels):
 
   def test_channels(self):
     env_lambda = TestEnv
-    env, obs, _ = self.init_batch_and_play(env_lambda, grayscale=True)
+    env, obs, _, _ = self.init_batch_and_play(env_lambda, grayscale=True)
     self.assert_channels(env, obs, n_channels=1)
 
-    env, obs, _ = self.init_batch_and_play(env_lambda, grayscale=False)
+    env, obs, _, _ = self.init_batch_and_play(env_lambda, grayscale=False)
     self.assert_channels(env, obs, n_channels=3)
 
 
diff --git a/tensor2tensor/rl/trainer_model_based_new.py b/tensor2tensor/rl/trainer_model_based_new.py
index 86317e261..b75167dbd 100644
--- a/tensor2tensor/rl/trainer_model_based_new.py
+++ b/tensor2tensor/rl/trainer_model_based_new.py
@@ -280,11 +280,11 @@ def setup_env(hparams):
   return env
 
 
-def eval_reward(env, epoch, clipped):
+def eval_reward(env, clipped):
   """Calculates mean rewards from given epoch."""
   reward_name = "reward" if clipped else "unclipped_reward"
   rewards = []
-  for rollout in env.rollouts_by_epoch[epoch]:
+  for rollout in env.current_epoch_rollouts():
     if rollout[-1].done:
       rollout_reward = sum(getattr(frame, reward_name) for frame in rollout)
       rewards.append(rollout_reward)
@@ -329,7 +329,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       env, ppo_model_dir,
       ppo_event_dir, data_dir,
       hparams, epoch=epoch, is_final_epoch=False)
-  mean_unclipped_reward = eval_reward(env, epoch, clipped=False)
+  mean_unclipped_reward = eval_reward(env, clipped=False)
   tf.logging.info("Mean reward (initial): {}".format(mean_unclipped_reward))
 
   eval_metrics_event_dir = os.path.join(directories["world_model"],
@@ -350,15 +350,15 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   )
 
   for epoch in range(hparams.epochs):
-    env.start_new_epoch(epoch)
-    is_final_epoch = (epoch + 1) == hparams.epochs
-    log = make_log_fn(epoch, log_relative_time)
-
     epoch_data_dir = os.path.join(directories["data"], str(epoch))
     tf.gfile.MakeDirs(epoch_data_dir)
     env.generate_data(epoch_data_dir, directories["tmp"])
     epoch_data_dirs.append(epoch_data_dir)
 
+    env.start_new_epoch(epoch)
+    is_final_epoch = (epoch + 1) == hparams.epochs
+    log = make_log_fn(epoch, log_relative_time)
+
     # Train world model
     log("Training world model")
     train_world_model(env, epoch_data_dir,
@@ -385,11 +385,11 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
     if hparams.stop_loop_early:
       return 0.0
-    mean_clipped_reward = eval_reward(env, epoch, clipped=True)
+    mean_clipped_reward = eval_reward(env, clipped=True)
     log("Mean clipped reward during generation: {}".format(
         mean_clipped_reward))  # this was output of generate_real_env_data(...)
 
-    mean_unclipped_reward = eval_reward(env, epoch, clipped=False)
+    mean_unclipped_reward = eval_reward(env, clipped=False)
     log("Mean eval reward (unclipped): {}".format(mean_unclipped_reward))
 
     # Summarize metrics.

From 77b711b39b65bc2c5c9c382b528478ccaa05a1a5 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Mon, 15 Oct 2018 15:40:00 -0700
Subject: [PATCH 1010/2720] internal merge of PR #1144

PiperOrigin-RevId: 217220447
---
 tensor2tensor/data_generators/gym_env.py      | 5 +----
 tensor2tensor/data_generators/gym_env_test.py | 4 ++--
 2 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 0c4a90e93..0ace13f9a 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -394,10 +394,7 @@ def splits_and_paths(self, data_dir):
     ]
 
   def generate_data(self, data_dir, tmp_dir, task_id=-1):
-    """Saves the rollout history to disk.
-
-    Also splits data into train/dev sets.
-    """
+    """Saves the rollout history to disk, split into train/dev sets."""
     self._split_current_epoch()
 
     splits_and_paths = self.splits_and_paths(data_dir)
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index a9b8000d8..075572e47 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -27,8 +27,8 @@
 from gym.spaces import Discrete
 import numpy as np
 
-from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import gym_env
+from tensor2tensor.data_generators import problem
 
 import tensorflow as tf
 
@@ -156,7 +156,7 @@ def test_frame_numbers_are_continuous(self):
     frame_numbers = [
         tf.train.Example.FromString(
             record
-        ).features.feature['frame_number'].int64_list.value[0]
+        ).features.feature["frame_number"].int64_list.value[0]
         for (_, paths) in env.splits_and_paths(self.out_dir)
         for path in paths
         for record in tf.python_io.tf_record_iterator(path)

From eb9618e3ede41dd089f15b17aae4eb98288927d2 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 15 Oct 2018 15:54:00 -0700
Subject: [PATCH 1011/2720] refactoring shared hparams to base hparams.

PiperOrigin-RevId: 217222548
---
 tensor2tensor/models/video/base.py            | 30 ++++++++++++++++++-
 .../video/basic_deterministic_params.py       | 19 ++----------
 2 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index c5db06b71..20bd7fd29 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -22,9 +22,9 @@
 from functools import partial
 import six
 
+from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
-from tensor2tensor.models.video import basic_deterministic_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -520,3 +520,31 @@ def merge(inputs, targets):
       rewards = merge(features["input_reward"], features["target_reward"])
     return self.__process(frames, actions, rewards, frames_raw)
 
+
+def next_frame_base():
+  """Common HParams for next_frame models."""
+  hparams = common_hparams.basic_params1()
+  # Loss cutoff.
+  hparams.add_hparam("video_modality_loss_cutoff", 0.01)
+  # Additional resizing the frames before feeding them to model.
+  hparams.add_hparam("preprocess_resize_frames", None)
+  # How many data points to suffle. Ideally should be part of problem not model!
+  hparams.add_hparam("shuffle_buffer_size", 128)
+  # Tiny mode. For faster tests.
+  hparams.add_hparam("tiny_mode", False)
+  # In case a model supports smaller/faster version.
+  hparams.add_hparam("small_mode", False)
+  # In case a model has stochastic version.
+  hparams.add_hparam("stochastic_model", False)
+  # Internal loss for recurrent models.
+  hparams.add_hparam("internal_loss", True)
+  # choose from: concat, multiplicative, multi_additive
+  hparams.add_hparam("action_injection", "multi_additive")
+  # Scheduled sampling method. Choose between
+  # ground_truth_only, prediction_only, prob, count, prob_inverse_exp.
+  hparams.add_hparam("scheduled_sampling_mode", "prediction_only")
+  hparams.add_hparam("scheduled_sampling_decay_steps", 10000)
+  hparams.add_hparam("scheduled_sampling_max_prob", 1.0)
+  hparams.add_hparam("scheduled_sampling_k", 900.0)
+  return hparams
+
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 972d06f04..9e65c93b7 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -18,14 +18,14 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.layers import common_hparams
+from tensor2tensor.models.video import base
 from tensor2tensor.utils import registry
 
 
 @registry.register_hparams
 def next_frame_basic_deterministic():
   """Basic 2-frame conv model."""
-  hparams = common_hparams.basic_params1()
+  hparams = base.next_frame_base()
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 1
   hparams.hidden_size = 64
@@ -41,23 +41,8 @@ def next_frame_basic_deterministic():
   hparams.weight_decay = 0.0
   hparams.clip_grad_norm = 1.0
   hparams.dropout = 0.5
-  # choose from: concat, multiplicative, multi_additive
-  hparams.add_hparam("action_injection", "multi_additive")
   hparams.add_hparam("num_compress_steps", 6)
   hparams.add_hparam("filter_double_steps", 2)
-  hparams.add_hparam("video_modality_loss_cutoff", 0.02)
-  hparams.add_hparam("preprocess_resize_frames", None)
-  hparams.add_hparam("shuffle_buffer_size", 128)
-  hparams.add_hparam("tiny_mode", False)
-  hparams.add_hparam("small_mode", False)
-  hparams.add_hparam("stochastic_model", False)
-  hparams.add_hparam("internal_loss", True)
-  # Scheduled sampling method. Choose between
-  # ground_truth_only, prediction_only, prob, count, prob_inverse_exp.
-  hparams.add_hparam("scheduled_sampling_mode", "prediction_only")
-  hparams.add_hparam("scheduled_sampling_decay_steps", 10000)
-  hparams.add_hparam("scheduled_sampling_max_prob", 1.0)
-  hparams.add_hparam("scheduled_sampling_k", 900.0)
   return hparams
 
 

From 49c7e2013c04ff197090d07fe34382b53b5e67e4 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 15 Oct 2018 16:51:38 -0700
Subject: [PATCH 1012/2720] using tf.layers.flatten and removing prod

PiperOrigin-RevId: 217231586
---
 .../models/video/basic_stochastic.py          | 47 ++++++++-----------
 1 file changed, 20 insertions(+), 27 deletions(-)

diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 36b5ec0d9..588524776 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -32,14 +32,7 @@
 
 import tensorflow as tf
 
-
-def prod(l):
-  """Product of elements in a list."""
-  res = l[0]
-  for i, e in enumerate(l):
-    if i > 0:
-      res *= e
-  return res
+tfl = tf.layers
 
 
 @registry.register_model
@@ -56,10 +49,10 @@ def inject_latent(self, layer, inputs, target):
     latent_mean, latent_std = self.construct_latent_tower(
         full_video, time_axis=1)
     latent = common_video.get_gaussian_tensor(latent_mean, latent_std)
-    latent = tf.layers.flatten(latent)
+    latent = tfl.flatten(latent)
     latent = tf.expand_dims(latent, axis=1)
     latent = tf.expand_dims(latent, axis=1)
-    latent_mask = tf.layers.dense(latent, filters, name="latent_mask")
+    latent_mask = tfl.dense(latent, filters, name="latent_mask")
     zeros_mask = tf.zeros(
         common_layers.shape_list(layer)[:-1] + [filters], dtype=tf.float32)
     layer = tf.concat([layer, latent_mask + zeros_mask], axis=-1)
@@ -82,15 +75,15 @@ def inject_latent(self, layer, inputs, target):
     batch_size = layer_shape[0]
     state_size = hparams.latent_predictor_state_size
     lstm_cell = tf.contrib.rnn.LSTMCell(state_size)
-    discrete_predict = tf.layers.Dense(256, name="discrete_predict")
-    discrete_embed = tf.layers.Dense(state_size, name="discrete_embed")
+    discrete_predict = tfl.Dense(256, name="discrete_predict")
+    discrete_embed = tfl.Dense(state_size, name="discrete_embed")
 
     def add_d(layer, d):
-      z_mul = tf.layers.dense(d, final_filters, name="unbottleneck_mul")
+      z_mul = tfl.dense(d, final_filters, name="unbottleneck_mul")
       if not hparams.complex_addn:
         return layer + z_mul
       layer *= tf.nn.sigmoid(z_mul)
-      z_add = tf.layers.dense(d, final_filters, name="unbottleneck_add")
+      z_add = tfl.dense(d, final_filters, name="unbottleneck_add")
       layer += z_add
       return layer
 
@@ -98,10 +91,10 @@ def add_d(layer, d):
       if hparams.full_latent_tower:
         rand = tf.random_uniform(layer_shape[:-1] + [hparams.bottleneck_bits])
       else:
-        layer_pred = tf.reshape(layer, [batch_size, prod(layer_shape[1:])])
-        prediction = tf.layers.dense(layer_pred, state_size, name="istate")
-        c_state = tf.layers.dense(layer_pred, state_size, name="cstate")
-        m_state = tf.layers.dense(layer_pred, state_size, name="mstate")
+        layer_pred = tfl.flatten(layer)
+        prediction = tfl.dense(layer_pred, state_size, name="istate")
+        c_state = tfl.dense(layer_pred, state_size, name="cstate")
+        m_state = tfl.dense(layer_pred, state_size, name="mstate")
         state = (c_state, m_state)
         outputs = []
         for i in range(hparams.bottleneck_bits // 8):
@@ -119,7 +112,7 @@ def add_d(layer, d):
 
     # Embed.
     frames = tf.concat([inputs, target], axis=-1)
-    x = tf.layers.dense(
+    x = tfl.dense(
         frames, filters, name="latent_embed",
         bias_initializer=tf.random_normal_initializer(stddev=0.01))
     x = common_attention.add_timing_signal_nd(x)
@@ -131,14 +124,14 @@ def add_d(layer, d):
           if i < hparams.filter_double_steps:
             filters *= 2
           x = common_attention.add_timing_signal_nd(x)
-          x = tf.layers.conv2d(x, filters, kernel,
-                               activation=common_layers.belu,
-                               strides=(2, 2), padding="SAME")
+          x = tfl.conv2d(x, filters, kernel,
+                         activation=common_layers.belu,
+                         strides=(2, 2), padding="SAME")
           x = common_layers.layer_norm(x)
     else:
       x = common_layers.double_discriminator(x)
       x = tf.expand_dims(tf.expand_dims(x, axis=1), axis=1)
-    x = tf.layers.dense(x, hparams.bottleneck_bits, name="bottleneck")
+    x = tfl.dense(x, hparams.bottleneck_bits, name="bottleneck")
     x0 = tf.tanh(x)
     d = x0 + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x0)) - 1.0 - x0)
     pred_loss = 0.0
@@ -149,10 +142,10 @@ def add_d(layer, d):
       tf.summary.histogram("d_int", tf.reshape(d_int, [-1]))
       d_hot = tf.one_hot(d_int, 256, axis=-1)
       d_pred = discrete_embed(d_hot)
-      layer_pred = tf.reshape(layer, [batch_size, prod(layer_shape[1:])])
-      prediction0 = tf.layers.dense(layer_pred, state_size, name="istate")
-      c_state = tf.layers.dense(layer_pred, state_size, name="cstate")
-      m_state = tf.layers.dense(layer_pred, state_size, name="mstate")
+      layer_pred = tfl.flatten(layer)
+      prediction0 = tfl.dense(layer_pred, state_size, name="istate")
+      c_state = tfl.dense(layer_pred, state_size, name="cstate")
+      m_state = tfl.dense(layer_pred, state_size, name="mstate")
       pred = tf.concat([tf.expand_dims(prediction0, axis=1), d_pred], axis=1)
       state = (c_state, m_state)
       outputs = []

From 1f125fc462126a9fd80b2c10b5d9ab4f1e545597 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 15 Oct 2018 17:08:00 -0700
Subject: [PATCH 1013/2720] Moving inject_additional_input to common_video.

PiperOrigin-RevId: 217234688
---
 tensor2tensor/layers/common_video.py          | 43 +++++++++++++++++++
 .../models/video/basic_deterministic.py       | 30 +------------
 tensor2tensor/models/video/sv2p.py            | 11 ++---
 3 files changed, 51 insertions(+), 33 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index abc5dbb31..eac08bada 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -146,6 +146,49 @@ def scheduled_sample_count(ground_truth_x,
   return output
 
 
+def inject_additional_input(layer, inputs, name, mode="concat"):
+  """Injects the additional input into the layer.
+
+  Args:
+    layer: layer that the input should be injected to.
+    inputs: inputs to be injected.
+    name: TF scope name.
+    mode: how the infor should be added to the layer:
+      "concat" concats as additional channels.
+      "multiplicative" broadcasts inputs and multiply them to the channels.
+      "multi_additive" broadcasts inputs and multiply and add to the channels.
+
+  Returns:
+    updated layer.
+
+  Raises:
+    ValueError: in case of unknown mode.
+  """
+  layer_shape = common_layers.shape_list(layer)
+  input_shape = common_layers.shape_list(inputs)
+  zeros_mask = tf.zeros(layer_shape, dtype=tf.float32)
+  if mode == "concat":
+    emb = encode_to_shape(inputs, layer_shape, name)
+    layer = tf.concat(values=[layer, emb], axis=-1)
+  elif mode == "multiplicative":
+    filters = layer_shape[-1]
+    input_reshaped = tf.reshape(inputs, [-1, 1, 1, input_shape[-1]])
+    input_mask = tf.layers.dense(input_reshaped, filters, name=name)
+    input_broad = input_mask + zeros_mask
+    layer *= input_broad
+  elif mode == "multi_additive":
+    filters = layer_shape[-1]
+    input_reshaped = tf.reshape(inputs, [-1, 1, 1, input_shape[-1]])
+    input_mul = tf.layers.dense(input_reshaped, filters, name=name + "_mul")
+    layer *= tf.nn.sigmoid(input_mul)
+    input_add = tf.layers.dense(input_reshaped, filters, name=name + "_add")
+    layer += input_add
+  else:
+    raise ValueError("Unknown injection mode: %s" % mode)
+
+  return layer
+
+
 def scheduled_sample_prob(ground_truth_x,
                           generated_x,
                           batch_size,
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 28430e1d1..af570e4fd 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -45,32 +45,6 @@ def inject_latent(self, layer, inputs, target):
     del inputs, target
     return layer, 0.0
 
-  def inject_additional_input(self, layer, inputs, name, mode="concat"):
-    """Injects the additional input into the layer."""
-    layer_shape = common_layers.shape_list(layer)
-    input_shape = common_layers.shape_list(inputs)
-    zeros_mask = tf.zeros(layer_shape, dtype=tf.float32)
-    if mode == "concat":
-      emb = common_video.encode_to_shape(inputs, layer_shape, name)
-      layer = tf.concat(values=[layer, emb], axis=-1)
-    elif mode == "multiplicative":
-      filters = layer_shape[-1]
-      input_reshaped = tf.reshape(inputs, [-1, 1, 1, input_shape[-1]])
-      input_mask = tf.layers.dense(input_reshaped, filters, name=name)
-      input_broad = input_mask + zeros_mask
-      layer *= input_broad
-    elif mode == "multi_additive":
-      filters = layer_shape[-1]
-      input_reshaped = tf.reshape(inputs, [-1, 1, 1, input_shape[-1]])
-      input_mul = tf.layers.dense(input_reshaped, filters, name=name + "_mul")
-      layer *= tf.nn.sigmoid(input_mul)
-      input_add = tf.layers.dense(input_reshaped, filters, name=name + "_add")
-      layer += input_add
-    else:
-      raise ValueError("Unknown injection mode: %s" % mode)
-
-    return layer
-
   def middle_network(self, layer, internal_states):
     # Run a stack of convolutions.
     x = layer
@@ -120,7 +94,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
     # Add embedded action if present.
     if self.has_actions:
       action = actions[:, -1, :]
-      x = self.inject_additional_input(
+      x = common_video.inject_additional_input(
           x, action, "action_enc", hparams.action_injection)
 
     # Inject latent if present. Only for stochastic models.
@@ -134,7 +108,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
     for i in range(hparams.num_compress_steps):
       with tf.variable_scope("upstride%d" % i):
         if self.has_actions:
-          x = self.inject_additional_input(
+          x = common_video.inject_additional_input(
               x, action, "action_enc", hparams.action_injection)
         if i >= hparams.num_compress_steps - hparams.filter_double_steps:
           filters //= 2
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 5f7d621dd..db18b4d81 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -147,10 +147,11 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
       layer_id += 1
 
     if action is not None:
-      enc2 = self.inject_additional_input(
+      enc2 = common_video.inject_additional_input(
           enc2, action, "action_enc", self.hparams.action_injection)
     if input_reward is not None:
-      enc2 = self.inject_additional_input(enc2, input_reward, "reward_enc")
+      enc2 = common_video.inject_additional_input(
+          enc2, input_reward, "reward_enc")
     if latent is not None and not concat_latent:
       with tf.control_dependencies([latent]):
         enc2 = tf.concat([enc2, latent], axis=3)
@@ -196,15 +197,15 @@ def reward_prediction_big(self, input_images, input_reward, action, latent):
 
       # Inject additional inputs
       if action is not None:
-        x = self.inject_additional_input(
+        x = common_video.inject_additional_input(
             x, action, "action_enc", self.hparams.action_injection)
       if input_reward is not None:
-        x = self.inject_additional_input(x, input_reward, "reward_enc")
+        x = common_video.inject_additional_input(x, input_reward, "reward_enc")
       if latent is not None:
         latent = tfl.flatten(latent)
         latent = tf.expand_dims(latent, axis=1)
         latent = tf.expand_dims(latent, axis=1)
-        x = self.inject_additional_input(x, latent, "latent_enc")
+        x = common_video.inject_additional_input(x, latent, "latent_enc")
 
       x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(2, 2),
                      activation=tf.nn.relu, name="reward_conv2")

From c8b60bebbb10414563974717e6d1d0b61c2fa959 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 15 Oct 2018 17:11:57 -0700
Subject: [PATCH 1014/2720] Remove MTF code in Tensor2Tensor. It's moved to
 github.com/tensorflow/mesh.

PiperOrigin-RevId: 217235228
---
 setup.py                                      |    1 +
 tensor2tensor/mesh_tensorflow/README.md       |  380 --
 tensor2tensor/mesh_tensorflow/__init__.py     |   15 -
 .../mesh_tensorflow/mesh_tensorflow.py        | 4092 -----------------
 .../mesh_tensorflow/mesh_tensorflow_test.py   |  170 -
 tensor2tensor/mesh_tensorflow/mnist.py        |  250 -
 .../mesh_tensorflow/mnist_dataset.py          |  131 -
 .../mesh_tensorflow/mtf_beam_search.py        |  521 ---
 tensor2tensor/mesh_tensorflow/mtf_layers.py   |  754 ---
 .../mesh_tensorflow/mtf_layers_test.py        |  289 --
 tensor2tensor/mesh_tensorflow/mtf_optimize.py |  271 --
 .../mesh_tensorflow/mtf_toy_model_tpu.py      |  239 -
 .../mtf_transformer_data_splitting.png        |  Bin 19717 -> 0 bytes
 .../mtf_transformer_model_splitting.png       |  Bin 20556 -> 0 bytes
 tensor2tensor/mesh_tensorflow/mtf_utils.py    |   70 -
 .../mesh_tensorflow/mtf_utils_test.py         |   54 -
 .../mesh_tensorflow/placement_mesh_impl.py    |  533 ---
 .../mesh_tensorflow/research/__init__.py      |   15 -
 .../mesh_tensorflow/simd_mesh_impl.py         |  407 --
 .../mesh_tensorflow/tpu_variables.py          |  200 -
 tensor2tensor/models/__init__.py              |    8 +-
 .../mtf_image_transformer.py                  |   23 +-
 .../mtf_image_transformer_test.py             |    9 +-
 .../{mesh_tensorflow => models}/mtf_resnet.py |   13 +-
 .../mtf_transformer.py                        |   44 +-
 .../mtf_transformer_test.py                   |    7 +-
 .../research/moe.py                           |   14 +-
 .../research/moe_experiments.py}              |    4 +-
 .../{mesh_tensorflow => utils}/mtf_model.py   |   21 +-
 29 files changed, 67 insertions(+), 8468 deletions(-)
 delete mode 100644 tensor2tensor/mesh_tensorflow/README.md
 delete mode 100644 tensor2tensor/mesh_tensorflow/__init__.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/mnist.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/mnist_dataset.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/mtf_beam_search.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/mtf_layers.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/mtf_layers_test.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/mtf_optimize.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/mtf_transformer_data_splitting.png
 delete mode 100644 tensor2tensor/mesh_tensorflow/mtf_transformer_model_splitting.png
 delete mode 100644 tensor2tensor/mesh_tensorflow/mtf_utils.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/mtf_utils_test.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/research/__init__.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
 delete mode 100644 tensor2tensor/mesh_tensorflow/tpu_variables.py
 rename tensor2tensor/{mesh_tensorflow => models}/mtf_image_transformer.py (95%)
 rename tensor2tensor/{mesh_tensorflow => models}/mtf_image_transformer_test.py (94%)
 rename tensor2tensor/{mesh_tensorflow => models}/mtf_resnet.py (97%)
 rename tensor2tensor/{mesh_tensorflow => models}/mtf_transformer.py (96%)
 rename tensor2tensor/{mesh_tensorflow => models}/mtf_transformer_test.py (95%)
 rename tensor2tensor/{mesh_tensorflow => models}/research/moe.py (98%)
 rename tensor2tensor/{mesh_tensorflow/research/experiments_moe.py => models/research/moe_experiments.py} (98%)
 rename tensor2tensor/{mesh_tensorflow => utils}/mtf_model.py (94%)

diff --git a/setup.py b/setup.py
index 325b69673..a1c4ed5f1 100644
--- a/setup.py
+++ b/setup.py
@@ -40,6 +40,7 @@
         'gunicorn',
         'gym',
         'h5py',
+        'mesh-tensorflow',
         'numpy',
         'oauth2client',
         'requests',
diff --git a/tensor2tensor/mesh_tensorflow/README.md b/tensor2tensor/mesh_tensorflow/README.md
deleted file mode 100644
index a74a0014a..000000000
--- a/tensor2tensor/mesh_tensorflow/README.md
+++ /dev/null
@@ -1,380 +0,0 @@
-# Mesh TensorFlow - Model Parallelism Made Easier
-
-Transformer for EN-FR WMT with model splitting |  Transformer for EN-FR WMT with data splitting
-:-------------------------:|:-------------------------:
-![model_splitting](./mtf_transformer_model_splitting.png) | ![data_splitting](./mtf_transformer_data_splitting.png)
-
-# Introduction
-
-Mesh TensorFlow (mtf) is a language for distributed deep
-learning, capable of specifying a broad class of distributed tensor
-computations.  The purpose of mesh-tensorflow is to formalize and implement
-distribution strategies for your computation graph over your hardware/processors
-For example: "Split the batch over rows of processors and split
-the units in the hidden layer across columns of processors." Mesh-TensorFlow is
-implemented as a layer over TensorFlow.
-
-## Do I need Mesh-TensorFlow?
-If you just want data-parallel training (batch-splitting), then you do not need
-mesh-tensorflow, though Mesh-TensorFlow can do this.  The most common reasons
-for more sophisticated parallel computation are:
-
-* The parameters of the model do not fit on one device - e.g. a
-5-billion-parameter language model.
-
-* An example is so large that the activations do not fit on one device. - e.g.
-large images.  TODO(noam): we still need to implement spatially-partitioned
-convolutions
-
-* Lower-latency parallel inference (at batch size 1).
-
-## The Mesh-TensorFlow Approach to Distributed Computation
-
-* A "Mesh" is an n-dimensional array of processors, connected by a network.
-
-* Each tensor is distributed (split and/or replicated) across all processors
-  in a mesh.
-
-* The "layout" of a tensor on a mesh is an injective partial map from the
-  dimensions of the tensor to the dimensions of the mesh, specifying which
-  dimensions of the tensor are split across which dimensions of the mesh.  An
-  empty layout means that the tensor is fully replicated across all processors.
-
-* Tensor dimensions and mesh dimensions are named.  The layouts of all tensors
-  follow from a set of user-defined layout rules which specify which
-  tensor-dimensions are split across which mesh-dimensions.  This ensures that
-  the corresponding dimensions in different tensors are split in the same
-  manner.
-
-* Layouts do not affect results - only performance.
-
-* The implementation of an operation involves parallel computation on all
-  processors in the mesh, and sometimes also collective communication.  A
-  processor usually just manipulates the slices of the input tensors already
-  resident on that processor, and produces the slice of the output that goes on
-  that processor.
-  
-## Example Models
-
-This directory contains code for running several well-known models across
-different tasks.
-
-We outline an example below. In the above figures, Mesh-TensorFlow scales
-linearly as the number of TPU shards increases. For model splitting, we varied
-the number of hidden units in the feedforward layer and the number of heads; for
-data splitting, we varied the batch size.
-
-### Example Network (MNIST)
-
-To illustrate, let us consider a simple model for the MNIST image-classification
-task.  Our network has one hidden layer with 1024 units, and an output layer
-with 10 units (corresponding to the 10 digit classes).  
-
-The code consists of two parts, the first describing the mathematical
-operations, and the second describing the devices and tensor/computation layout.
-For the full example, see [`mnist.py`](mnist.py).
-TODO(noam): verify that this code works.
-
-```Python
-# tf_images is a tf.Tensor with shape [100, 28, 28] and dtype tf.float32
-# tf_labels is a tf.Tensor with shape [100] and dtype tf.int32
-graph = mtf.Graph()
-mesh = mtf.Mesh(graph, "my_mesh")
-batch_dim = mtf.Dimension("batch", 100)
-rows_dim = mtf.Dimension("rows", 28)
-cols_dim = mtf.Dimension("cols", 28)
-hidden_dim = mtf.Dimension("hidden", 1024)
-classes_dim = mtf.Dimension("classes", 10)
-images = mtf.import_tf_tensor(
-    mesh, tf_images, shape=[batch_dim, rows_dim, cols_dim])
-labels = mtf.import_tf_tensor(mesh, tf_labels, [batch_dim])
-w1 = mtf.get_variable(mesh, "w1", [rows_dim, cols_dim, hidden_dim])
-w2 = mtf.get_variable(mesh, "w2", [hidden_dim, classes_dim])
-# einsum is a generalization of matrix multiplication (see numpy.einsum)
-hidden = mtf.relu(mtf.einsum(images, w1, output_shape=[batch_dim, hidden_dim]))
-logits = mtf.einsum(hidden, w2, output_shape=[batch_dim, classes_dim])
-loss = mtf.reduce_mean(mtf_layers.softmax_cross_entropy_with_logits(
-    logits, mtf.one_hot(labels, classes_dim), classes_dim))
-w1_grad, w2_grad = mtf.gradients([loss], [w1, w2])
-update_w1_op = mtf.assign(w1, w1 - w1_grad * 0.001)
-update_w2_op = mtf.assign(w1, w1 - w1_grad * 0.001)
-```
-
-In the code above, we have built a mesh-tensorflow graph, which is simply
-a Python structure.  We have completely defined the mathematical operations.
-In the code below, we specify the mesh of processors and the layout of the
-computation.
-
-```Python
-devices = ["gpu:0", "gpu:1", "gpu:2", "gpu:3"]
-mesh_shape = [("all_processors", 4)]
-layout_rules = [("batch", "all_processors")]
-mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-    mesh_shape, layout_rules, devices)
-lowering = mtf.Lowering(graph, {mesh:mesh_impl})
-tf_update_ops = [lowering.lowered_operation(update_w1_op),
-                 lowering.lowered_operation(update_w2_op)]
-```
-
-The particular layout above implements data-parallelism, splitting the batch of
-examples evenly across all four processors.  Any Tensor with a "batch" dimension
-(e.g. `images`, `h`, `logits`, and their gradients) is split in that dimension
-across all processors, while any tensor without a "batch" dimension (e.g. the
-model parameters) is replicated identically on every processor.
-
-Alternatively, for model-parallelism, we can set
-`layout_rules=[("hidden", "all_processors")]`.  In this case,
-any tensor with a "hidden" dimension (e.g. `hidden`, `w1`, `w2`)  is split,
-while any other tensor (e.g. `image`, `logits`) is fully replicated.
-
-We can even combine data-parallelism and model-parallelism on a 2-dimensional
-mesh of processors.  We split the batch along one dimension of the mesh, and the
-units in the hidden layer along the other dimension of the mesh, as below.  In
-this case, the hidden layer is actually tiled between the four processors, being
-split in both the "batch" and "hidden_units" dimensions.
-
-```Python
-mesh_shape = [("processor_rows", 2), ("processor_cols", 2)]
-layout_rules = [("batch", "processor_rows"), ("hidden", "processor_cols")]
-```
-
-## Where does the network communication happen?
-
-Some mesh-tensorflow operations cause network communication.  For example, an
-einsum (generalized matrix multiplication) is computed as follows:
-
-* On each processor, compute the einsum of the slices of the two operands that
-  are local to that processor.
-* If no reduced-out dimensions are split, then we are done.
-* If reduced-out dimensions are split, then perform an "allreduce" operation 
-  on the resulting slices - summing across any mesh dimensions over which the
-  reduced-out dimensions are split.
-
-Where the allreduces happen depends will depend on the computation layout.
-For example, in a data-parallel layout where the "batch" dimension is split,
-allreduces will happen when computing the parameter gradients, since this
-involves matrix multiplications which reduce out the "batch" dimension.
-
-## How do I pick a layout?
-
-While results do not depend on layout (except in the realm of roundoff errors
-and random seeds), performance and memory consumption depend heavily on layout.
-One day we hope to automate the process of choosing a layout.  For now, you
-really need to fully understand the performance implications and pick one
-yourself.  Mesh-tensorflow helps by accumulating and printing counters of
-computation/communication.  To start, here are some tricks/guidelines.
-
-* It is illegal for two dimensions of the same tensor to be split across the
-  same batch dimension.
-* For any compute-intense operation (e.g. einsum), make sure that all
-  mesh-dimensions are used to split dimensions of the inputs or outputs.
-  Otherwise, computation is duplicated.
-* To keep the ratio of compute/commuication high (i.e. not be bandwidth-bound),
-  split dimensions into large chunks.  This should be familiar in the
-  data-parallelism case, where we want a large batch size per processor to avoid
-  spending most of our time communicating.
-
-# The Mesh-TensorFlow Language
-
-Mesh-TensorFlow (v0.0) is implemented as a Python library which can generate
-part of a TensorFlow graph.  The user first builds a `mtf.Graph` (the analog of
-a TensorFlow graph) made up of `mtf.Tensor`s and `mtf.Operation`s.  As in
-TensorFlow, this graph consists of simple Python objects.  The user then creates
-a `mtf.Lowering` object, which lowers the `mtf.Graph` into TensorFlow, adding to
-the default TensorFlow graph.
-
-The Mesh-TensorFlow language is nearly identical to TensorFlow, with the
-familiar notion of a Graph, Tensors, Operations, and automatic gradient
-computation.  The principal differences are as follows:
-
-## Meshes replace devices
-
-A `Mesh` is a n-dimensional array of processors with named dimensions.  Each
-`Tensor` is assigned to a `Mesh`, instead of a device.
-
-## Tensor dimensions are named
-
-Each `Tensor` has a static `Shape`, which is a tuple of different "Dimensions".
-A `Dimension` is a `(name, size)` pair. For example, the shape of a `Tensor`
-representing a batch of images might be:
-
-`[("batch", 100), ("rows", 28"), ("cols", 28), ("channels", 3)]`.
-
-## Layouts
-
-A `Tensor` is laid out on its mesh with one slice on each processor.  A `Tensor`
-"layout", is an injective partial map specifying which dimensions of the tensor
-are (evenly) split across which dimensions of the mesh.  No dimension of a
-tensor may be split across two dimensions of its mesh and no two dimensions of a
-tensor may be split across the same dimension of its mesh.  The user defines a
-global set of layout rules in the form of (tensor-dimension-name,
-mesh-dimension-name) pairs.  A dimension of a tensor is split across a dimension
-of its mesh if there is a matching rule.
-
-### Example Layouts
-
-Take our example `Tensor` `image_batch` with shape: 
-`[("batch", 100), ("rows", 28"), ("cols", 28), ("channels", 3)]`
-
-Assume that this `Tensor` is assigned to a mesh of 8 processors with shape:
-`[("processor_rows", 2), ("processor_cols", 4)]`
-
-* If we use an empty set of layout rules `[]`, we get no splitting.  Each
-  processor contains the whole `Tensor`.
-
-* If we use the layout rules `"batch:processor_cols"`, then the `"batch"`
-  dimension of the `Tensor` is split across the `"processor_cols"` dimension of
-  the batch.  This means that each processor contains a Tensor slice with shape
-  `[25, 28, 28, 3]`.  For example, processors (0, 3) and (1, 3) contain
-  identical slices - `image_batch[75:100, :, :, :]`.
-
-* If we use the layout rules `"rows:processor_rows;cols:processor_cols"`, 
-  then the image is split in two dimensions, with each processor containing one
-  spatial tile with shape `[100, 14, 7, 3]`.   For example, processor (0, 1)
-  contains the slice `image_batch[:, 0:14, 7:14, :]`.
-
-Some layout rules would lead to illegal layouts:
-
-* `"batch:processor_rows;rows:processor_rows"` is illegal because two tensor
-  dimensions could be split across the same mesh dimension.
-
-* `"channels:processor_rows"` is illegal because the size of the tensor
-  dimension is not evenly divisible by the size of the mesh dimension.
-
-## Einsum
-
-Mesh-TensorFlow uses Einstein-summation notation, `mtf.einsum(inputs,
-output_shape)`, using the (named) `Dimensions` as the symbols.  Matrix-
-multiplication, broadcast, sum-reduction, and transposition can all be expressed
-as special cases of `mtf.einsum`, though the familiar interfaces are also
-supported.  The operation is lowered to slice-wise `tf.einsum`s, followed by
-allreduce across any mesh-dimensions corresponding to the summed-out Tensor
-dimensions.
-
-## Reshape can be expensive
-
-`mtf.reshape(x, new_shape)` is used to change a `Tensor`'s shape, potentially
-leading to a new tensor layout and hence network communication.
-
-# CPU/GPU/TPU implementations
-
-Mesh-TensorFlow works on CPU, GPU and TPU.  The TPU implementation is very
-different from the CPU/GPU implementation.
-
-Multi-CPU/GPU meshes are implemented with `PlacementMeshImpl`.  In this case
-mesh-tensorflow emits separate tensorflow operations placed on the different
-devices, all in one big tensorflow graph.
-
-TPU meshes are implemented in with `SimdMeshImpl`.  In this case,
-mesh-tensorflow emits tensorflow operations (and communication collectives) from
-the perspective of one core, and this same program runs on every core, relying
-on the fact that each core actually performs the same operations.  This
-piggy-backs on the TPU data-parallelism infrastructure, which operates the same
-way.  This "SIMD" approach keeps the tensorflow and xla graphs from growing with
-the number of cores.  The differences between cores are as follows:
-
-* different slices of the variables (this works now)
-* different positions in the collective communication (this works now)
-* different slices of the infed and outfed tensors.  We currently work around
-  this by requiring that all imported/exported tensors be fully-replicated.  In
-  the future, we should handle this correctly.
-
-# Instructions for running on cloud-tpu
-
-Note: It requires `tensorflow>=1.11.0`.
-
-## Prerequisite
-
-Please go through the
-[Transformer tutorial](https://cloud.google.com/tpu/docs/tutorials/transformer).
-
-## Create VM and TPU instance in Cloud console
-
-```sh
-ctpu up -name=ylc-mtf-donut -tf-version=nightly -tpu-size=v2-8 -zone=us-central1-b
-```
-
-## SSH into VM
-
-```sh
-git clone https://github.com/tensorflow/tensor2tensor.git
-cd tensor2tensor/
-pip install --user .
-```
-
-## Run the Transfomer model with Tensor2Tensor config
-
-Before run the model, you need to prepare the training data and bucket for
-storing checkpoints. Refer to the
-[Transformer tutorial](https://cloud.google.com/tpu/docs/tutorials/transformer)
-to learn how to generate the training data and create buckets.
-
-```sh
-CONF=mtf_transformer_paper_tr_0_mesh_8
-NAME=ende_$CONF\_0828
-MODEL=mtf_transformer
-PROBLEM=translate_ende_wmt32k_packed
-
-DATA_DIR=gs://xxxx
-OUT_DIR=gs://xxxx
-TPU_NAME=ylc-mtf-donut
-
-tensor2tensor/bin/t2t-trainer \
-  --model=$MODEL \
-  --hparams_set=$CONF \
-  --problem=$PROBLEM \
-  --train_steps=10000 \
-  --eval_steps=200 \
-  --data_dir=$DATA_DIR \
-  --output_dir=$OUT_DIR \
-  --use_tpu=True \
-  --cloud_tpu_name=$TPU_NAME
-```
-
-
-## Run the toy model without Tensor2Tensor dependencies
-
-  This toy model contains two fully-connected layers which aim to train a
-  identity function: f(x) = x. Since there are 8 TPU cores, we can arbitrary
-  change the FLAGS.mesh_shape and FLAGS.layout to achieve different
-  data-parallelism and model-parallelism strategies.
-
-```sh
-MODEL_DIR=gs://xxxx
-TPU_NAME=ylc-mtf-donut
-
-# 2 ways data-parallelism and 4 ways model-parallelism.
-# In this configuration, we split the batch dimension into 2 cores and the
-# hidden dimension into 4 cores.
-python mtf_toy_model_tpu.py \
-  --tpu=$TPU \
-  --model_dir=$MODEL_DIR \
-  --io_size=8 \
-  --hidden_size=8 \
-  --mesh_shape='x:2;y:4' \
-  --layout='batch:x;hidden:y'
-
-# 8 ways model-parallelism.
-# In this configuration, We split the hidden dimension into 8 cores.
-python mtf_toy_model_tpu.py \
-  --tpu=$TPU \
-  --model_dir=$MODEL_DIR \
-  --io_size=8 \
-  --hidden_size=8 \
-  --mesh_shape='all:8' \
-  --layout='hidden:all'
-
-# TODO LIST (please add items)
-
-We are actively working on improving Mesh-TensorFlow in a variety of ways.  Some
-of the top-priority items are:
-`Contact us if you'd like to help!`
-
-* Operations necessary for spatial-partitioning (spatially-partitioned
-  convolution, etc)
-* Examples of image-classification models.
-* Support for multiple meshes and efficient communication between them.  For
-  example, we may want to load training data on a mesh of 64 cpu-machines and
-  infeed them to a mesh of 512 tpu-cores.  We do not need this for language
-  tasks where the data is tiny, but it will be important for other tasks.
-
diff --git a/tensor2tensor/mesh_tensorflow/__init__.py b/tensor2tensor/mesh_tensorflow/__init__.py
deleted file mode 100644
index 4bd418a74..000000000
--- a/tensor2tensor/mesh_tensorflow/__init__.py
+++ /dev/null
@@ -1,15 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
deleted file mode 100644
index 707edc554..000000000
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow.py
+++ /dev/null
@@ -1,4092 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Mesh-TensorFlow."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import collections
-from functools import reduce  # pylint: disable=redefined-builtin; for py3
-from operator import mul
-import re
-import six
-from six.moves import xrange  # pylint: disable=redefined-builtin
-
-from tensor2tensor.mesh_tensorflow import mtf_utils
-import tensorflow as tf
-
-
-Dimension = collections.namedtuple("Dimension", ["name", "size"])
-
-
-def convert_to_dimension(d):
-  """Converts input to a Dimension.
-
-  Args:
-    d: Dimension, tuple (string, int), or None.
-
-  Returns:
-    Dimension or None.
-
-  Raises:
-    ValueError: If d cannot be converted to a Dimension.
-  """
-  if d is None:
-    return None
-  if isinstance(d, Dimension):
-    return d
-  name, size = d
-  if isinstance(name, str) and isinstance(size, int):
-    return Dimension(name, size)
-  else:
-    raise ValueError("could not convert %s to Dimension" % (d,))
-
-
-class Shape(object):
-  """Shape of a Tensor or Mesh.
-
-  #### Examples
-
-  ```python
-  # Create shape [4, 8] with names "x" and "y" respectively.
-  shape = mtf.Shape([mtf.Dimension("x", 4), mtf.Dimension("y", 8)])
-  ```
-  """
-
-  def __init__(self, dims):
-    """Constructs a shape for a Tensor or Mesh.
-
-    Args:
-      dims: List-like of Dimensions.
-
-    Raises:
-      ValueError: If Dimensions are repeated.
-    """
-    self._dims = [convert_to_dimension(d) for d in tuple(dims)]
-    if len(set(dims)) != len(dims):
-      raise ValueError("Shape must not have repeated dimensions %s" % dims)
-
-  @property
-  def dims(self):
-    return list(self._dims)
-
-  @property
-  def ndims(self):
-    return len(self._dims)
-
-  def __repr__(self):
-    return self.to_string
-
-  def __eq__(self, other):
-    return self.dims == other.dims
-
-  def __ne__(self, other):
-    return self.dims != other.dims
-
-  def __add__(self, other):
-    if isinstance(other, Shape):
-      other = other.dims
-    if isinstance(other, Dimension):
-      other = [other]
-    return Shape(self.dims + other)
-
-  def __sub__(self, other):
-    if other is None:
-      return self
-    if isinstance(other, Shape):
-      other = other.dims
-    if isinstance(other, Dimension):
-      other = [other]
-    return Shape([d for d in self.dims if d not in other])
-
-  def __len__(self):
-    return len(self._dims)
-
-  def __getitem__(self, key):
-    return self._dims[key]
-
-  def __iter__(self):
-    return iter(self._dims)
-
-  @property
-  def to_integer_list(self):
-    return [d.size for d in self.dims]
-
-  @property
-  def size(self):
-    return list_product(self.to_integer_list)
-
-  @property
-  def to_string(self):
-    return "Shape[%s]" % ", ".join(
-        ["%s=%d" % (d.name, d.size) for d in self.dims])
-
-  @property
-  def cumprod(self):
-    """Cumulative product (exclusive) of Dimension sizes."""
-    return _cumprod(self.to_integer_list)[:-1]
-
-  def cumprod_to_tensor_axis(self, cumprod):
-    """Tensor axis i such that self.cumprod[i] == cumprod, or None."""
-    try:
-      return self.cumprod.index(cumprod)
-    except ValueError:
-      return None
-
-  @property
-  def dimension_names(self):
-    return [d.name for d in self.dims]
-
-  def rename_dimension(self, old_name, new_name):
-    """Returns a copy where one dimension is renamed."""
-    if old_name not in self.dimension_names:
-      raise ValueError("Shape %s does not have dimension named %s"
-                       % (self, old_name))
-    return Shape(
-        [Dimension(new_name, d.size) if d.name == old_name else d
-         for d in self.dims])
-
-  def resize_dimension(self, name, new_size):
-    """Returns a copy where one dimension has a different size."""
-    if name not in self.dimension_names:
-      raise ValueError("Shape %s does not have dimension named %s"
-                       % (self, name))
-    return Shape(
-        [Dimension(name, new_size) if d.name == name else d
-         for d in self.dims])
-
-
-def convert_to_shape(x):
-  """Converts input to a Shape.
-
-  Args:
-    x: Shape, str, or None.
-
-  Returns:
-    Shape or None.
-
-  Raises:
-    ValueError: If x cannot be converted to a Shape.
-  """
-  if x is None:
-    return None
-  if isinstance(x, Shape):
-    return x
-  if isinstance(x, str):
-    x = _parse_string_to_list_of_pairs(x, seconds_to_int=True)
-  return Shape(x)
-
-
-class LayoutRules(object):
-  """Represents layout of a computation.
-
-  #### Examples
-
-  ```python
-  # Map "d_ff" and "heads" Tensor Dimensions to the "model" Mesh Dimension.
-  layout_rules = mtf.LayoutRules([("d_ff", "model"), ("heads", "model")])
-  ```
-  """
-
-  def __init__(self, pairs):
-    """Constructs a layout.
-
-    Args:
-      pairs: Set-like of string pairs (tensor_dim_name, mesh_dim_name).
-    """
-    self._pairs = set(pairs)
-
-  def __repr__(self):
-    return "LayoutRules%s" % self._pairs
-
-  def tensor_dimension_to_mesh_axis(self, tensor_dimension, mesh_shape):
-    """Mesh axis associated with tensor dimension (or None).
-
-    Args:
-      tensor_dimension: Dimension.
-      mesh_shape: Shape.
-
-    Returns:
-      Integer or None.
-
-    Raises:
-      ValueError: If one Tensor dimension maps to two mesh dimensions.
-    """
-    val = [i for i, mesh_dimension in enumerate(mesh_shape)
-           if (tensor_dimension.name, mesh_dimension.name) in self._pairs]
-    if len(val) > 1:
-      raise ValueError(
-          "Tensor dimension maps to multiple mesh dimensions"
-          " tensor_dimension=%s mesh_shape=%s layout=%s"
-          % (tensor_dimension, mesh_shape, self._pairs))
-    return val[0] if val else None
-
-  def tensor_layout(self, tensor_shape, mesh_shape):
-    """Computes TensorLayout given a Tensor Shape and a Mesh Shape.
-
-    Args:
-      tensor_shape: Shape.
-      mesh_shape: Shape.
-
-    Returns:
-      TensorLayout.
-
-    Raises:
-      ValueError: If two Tensor Dimensions map to the same Mesh Dimensions.
-    """
-    ret = [self.tensor_dimension_to_mesh_axis(d, mesh_shape)
-           for d in tensor_shape]
-    not_nones = [a for a in ret if a is not None]
-    if len(not_nones) != len(set(not_nones)):
-      raise ValueError(
-          "Two Tensor Dimensions may not map to the same Mesh Dimension:"
-          " layout=%s tensor_shape=%s mesh_shape=%s " %
-          (self, tensor_shape, mesh_shape))
-    return TensorLayout(ret)
-
-
-def convert_to_layout_rules(x):
-  """Converts input to a LayoutRules.
-
-  Args:
-    x: LayoutRules, str, or set-like of string pairs.
-
-  Returns:
-    LayoutRules.
-  """
-  if isinstance(x, LayoutRules):
-    return x
-  if isinstance(x, str):
-    x = _parse_string_to_list_of_pairs(x)
-  return LayoutRules(x)
-
-
-class TensorLayout(object):
-  """Injective partial map between Tensor axes and Mesh axes.
-
-  TensorLayout is a tuple of optional integers with length tensor.ndims. Each
-  item is either a unique integer indicating the mesh axis over which that
-  tensor dimension is split or None, indicating that this tensor dimension is
-  not split.
-
-  #### Examples
-
-  ```python
-  # Split first and last Tensor dimensions according to mesh axes 0 and 1.
-  tensor_layout = mtf.TensorLayout([0, None, 1])
-  ```
-  """
-
-  def __init__(self, tensor_axis_to_mesh_axis):
-    """Creates a TensorLayout.
-
-    Args:
-      tensor_axis_to_mesh_axis: List-like where each element is an int or None.
-    """
-    self._tensor_axis_to_mesh_axis = tuple(tensor_axis_to_mesh_axis)
-
-  def __eq__(self, other):
-    return self.tensor_axis_to_mesh_axis == other.tensor_axis_to_mesh_axis
-
-  def __ne__(self, other):
-    return self.tensor_axis_to_mesh_axis != other.tensor_axis_to_mesh_axis
-
-  def __repr__(self):
-    return "TensorLayout%s" % (self.tensor_axis_to_mesh_axis,)
-
-  def __len__(self):
-    return len(self._tensor_axis_to_mesh_axis)
-
-  def __getitem__(self, key):
-    return self._tensor_axis_to_mesh_axis[key]
-
-  def __iter__(self):
-    return iter(self._tensor_axis_to_mesh_axis)
-
-  @property
-  def tensor_axis_to_mesh_axis(self):
-    """Converts to a tuple of optional integers."""
-    return self._tensor_axis_to_mesh_axis
-
-  @property
-  def is_fully_replicated(self):
-    """Whether all tensor dimensions map to None."""
-    return self.tensor_axis_to_mesh_axis == (None,) * len(self)
-
-  def mesh_axis_to_tensor_axis(self, mesh_ndims):
-    """For each mesh axis, which Tensor axis maps to it.
-
-    Args:
-      mesh_ndims: int.
-
-    Returns:
-      Tuple of optional integers, with length mesh_ndims.
-    """
-    return tuple(
-        [self._tensor_axis_to_mesh_axis.index(mesh_axis)
-         if mesh_axis in self._tensor_axis_to_mesh_axis else None
-         for mesh_axis in xrange(mesh_ndims)])
-
-
-class Graph(object):
-  """Mesh-TensorFlow graph."""
-
-  def __init__(self):
-    self._operations = []
-    self._tensors = []
-    self._trainable_variables = []
-    self._all_variables = []
-
-  def __repr__(self):
-    return self.to_string
-
-  @property
-  def operations(self):
-    return self._operations
-
-  @property
-  def tensors(self):
-    return self._tensors
-
-  @property
-  def trainable_variables(self):
-    return self._trainable_variables
-
-  @property
-  def all_variables(self):
-    return self._all_variables
-
-  @property
-  def to_string(self):
-    return "\n".join([op.to_string for op in self.operations])
-
-
-class Lowering(object):
-  """Lowering of a Graph from Mesh-TensorFlow to TensorFlow.
-
-  #### Examples
-
-  Below we form a Graph with one Tensor and lower it to recover the original
-  tf.Tensor.
-
-  ```python
-  from tensor2tensor.mesh_tensorflow import placement_mesh_impl
-
-  graph = mtf.Graph()
-  mesh = mtf.Mesh(graph, "my_mesh")
-  inputs = tf.constant(0.)
-  mtf_inputs = mtf.import_tf_tensor(mesh,
-                                    inputs=inputs,
-                                    shape=mtf.Shape([]))
-  mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-      shape=[], layout={}, devices=[""])
-  lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-  outputs = lowering.export_to_tf_tensor(mtf_inputs)  # tf.constant(0.)
-  ```
-  """
-
-  def __init__(self, graph, mesh_to_impl):
-    """Creates a Lowering of a Graph.
-
-    Args:
-      graph: Graph.
-      mesh_to_impl: {Mesh: MeshImpl}. Keys are the Mesh's in the graph and
-        their values are MeshImpl's, which map Tensor Dimension names to
-        Mesh Dimension names.
-    """
-    # tf.logging.info("LOWERING GRAPH:\n%s" % graph.to_string)
-    self.mesh_to_impl = mesh_to_impl   # {Mesh: MeshImpl}
-    self.graph = graph
-    self._counters = []
-    self.tensors = {}                  # {Tensor: Mesh.LaidOutTensor}
-    self.operations = {}               # {Operation: tf.Operation}
-    self.variables = {}                # {Variable: LaidOutVariable}
-    for op in graph.operations:
-      # tf.logging.info("Lowering operation %s" % op.to_string)
-      with tf.name_scope(op.name):
-        op.lower(self)
-      for out in op.outputs:
-        self.add_counter(
-            "output/%s" % type(op).__name__, self.laid_out_size(out))
-        self.add_counter("output_unique/%s" % type(op).__name__, out.size)
-    log_variable_sizes(
-        graph.trainable_variables, "Trainable Variables", verbose=True)
-    tf.logging.info("Counters:\n" + pretty_print_counters(self._counters))
-
-  def mesh_impl(self, m):
-    if not isinstance(m, Mesh):
-      m = m.mesh
-    return self.mesh_to_impl[m]
-
-  def export_to_tf_tensor(self, x):
-    """Turn a Tensor into a tf.Tensor.
-
-    Args:
-      x: Tensor.
-
-    Returns:
-      tf.Tensor.
-    """
-    mesh_impl = self.mesh_impl(x)
-    return mesh_impl.export_to_tf_tensor(
-        x, self.tensors[x].to_laid_out_tensor())
-
-  def lowered_operation(self, op):
-    return self.operations[op]
-
-  def copy_masters_to_slices(self):
-    return tf.group(
-        [v.copy_master_to_slices for v in six.itervalues(self.variables)])
-
-  def copy_slices_to_masters(self):
-    return tf.group(
-        [v.copy_slices_to_master for v in six.itervalues(self.variables)])
-
-  def add_counter(self, key, value):
-    assert isinstance(value, int)
-    self._counters.append((key, value))
-
-  @property
-  def counters(self):
-    return self._counters
-
-  def laid_out_size(self, tensor):
-    """Total size of all slices.
-
-    Args:
-      tensor: Tensor.
-
-    Returns:
-      int.
-    """
-    return self.mesh_impl(tensor).laid_out_size(tensor.shape)
-
-  def set_tensor_lowering(self, tensor, laid_out_tensor):
-    self.verify_slice_shapes(tensor, laid_out_tensor)
-    self.tensors[tensor] = laid_out_tensor
-
-  def verify_slice_shapes(self, tensor, laid_out_tensor):
-    mesh_impl = self.mesh_impl(tensor)
-    correct_shape = mesh_impl.slice_shape(tensor.shape)
-    actual_shape = laid_out_tensor.slice_shape
-    if actual_shape != correct_shape:
-      raise ValueError(
-          "Wrong slice shape: correct_shape = %s actual shape = %s"
-          % (correct_shape, actual_shape))
-
-
-class Mesh(object):
-  """A placeholder with no functionality.
-
-  A Graph is built with each Tensor assigned to a Mesh. The Mesh does not
-  know its shape or its implementation.
-
-  A Lowering assigns each Mesh to a MeshImpl.
-  """
-
-  def __init__(self, graph, name, variable_placer=None):
-    self._graph = graph
-    self._name = name
-    self._variable_placer = variable_placer
-
-  @property
-  def graph(self):
-    return self._graph
-
-  @property
-  def variable_placer_fn(self):
-    if self._variable_placer is not None:
-      return self._variable_placer.device_function
-    else:
-      return "cpu:0"
-
-
-class MeshImpl(object):
-  """Implementation of a Mesh.
-
-  Unlike Mesh, MeshImpl carries Shape and LayoutRules. Subclasses of MeshImpl
-  also carry devices.
-
-  #### Examples
-
-  ```python
-  shape = mtf.Shape([mtf.Dimension("batch", 4),
-                     mtf.Dimension("model", 8)])
-  layout_rules = mtf.LayoutRules([("batch", "batch"),
-                                  ("d_ff", "model"),
-                                  ("heads", "model")])
-  mesh_impl = mtf.MeshImpl(shape=shape, layout_rules=layout_rules)
-  ```
-  """
-
-  def __init__(self, shape, layout_rules):
-    """Creates a mesh implementation.
-
-    Args:
-      shape: Shape.
-      layout_rules: LayoutRules.
-    """
-    self._shape = convert_to_shape(shape)
-    self._layout_rules = convert_to_layout_rules(layout_rules)
-
-  @property
-  def shape(self):
-    return self._shape
-
-  @property
-  def ndims(self):
-    return len(self._shape)
-
-  @property
-  def layout_rules(self):
-    return self._layout_rules
-
-  @property
-  def size(self):
-    return self.shape.size
-
-  @property
-  def supports_control_dependencies(self):
-    return True
-
-  def tensor_dimension_to_mesh_axis(self, tensor_dimension):
-    """Mesh axis associated with tensor dimension (or None).
-
-    Args:
-      tensor_dimension: Dimension.
-
-    Returns:
-      int or None.
-    """
-    return self.layout_rules.tensor_dimension_to_mesh_axis(
-        tensor_dimension, self.shape)
-
-  def tensor_layout(self, arg):
-    """Compute TensorLayout for a Tensor or a Shape.
-
-    Args:
-      arg: Tensor or Shape.
-
-    Returns:
-      TensorLayout.
-    """
-    if isinstance(arg, Tensor):
-      arg = arg.shape
-    return self.layout_rules.tensor_layout(arg, self.shape)
-
-  def mesh_axis_to_cumprod(self, tensor_shape):
-    """For each mesh axis, give the product of previous tensor axes.
-
-    Args:
-      tensor_shape: Shape.
-
-    Returns:
-      list with length self.ndims where each element is an integer or None.
-    """
-    tensor_layout = self.tensor_layout(tensor_shape)
-    ma2ta = tensor_layout.mesh_axis_to_tensor_axis(self.ndims)
-    ta2cumprod = tensor_shape.cumprod
-    return [None if ta is None else ta2cumprod[ta] for ta in ma2ta]
-
-  def slice_shape(self, tensor_shape):
-    """Shape of each slice of the Tensor.
-
-    Args:
-      tensor_shape: Shape.
-
-    Returns:
-      list of integers with length tensor_shape.ndims.
-
-    Raises:
-      ValueError: If a Tensor dimension is not divisible by the corresponding
-        Mesh dimension.
-    """
-    tensor_layout = self.tensor_layout(tensor_shape)
-    ret = []
-    for tensor_dim, mesh_axis in zip(
-        tensor_shape, tensor_layout.tensor_axis_to_mesh_axis):
-      if mesh_axis is None:
-        ret.append(tensor_dim.size)
-      else:
-        mesh_dim = self.shape[mesh_axis]
-        if tensor_dim.size % mesh_dim.size != 0:
-          raise ValueError(
-              "Tensor dimension size not divisible by mesh dimension size:"
-              " tensor_shape=%s tensor_layout=%s"
-              % (tensor_shape, tensor_layout))
-        ret.append(tensor_dim.size // mesh_dim.size)
-    return ret
-
-  def slice_begin(self, tensor_shape, pnum):
-    """Begin position for the tensor slice for the given processor.
-
-    Args:
-      tensor_shape: Shape.
-      pnum: int <= self.size.
-
-    Returns:
-      list of integers with length tensor_shape.ndims.
-    """
-    tensor_layout = self.tensor_layout(tensor_shape)
-    coordinates = pnum_to_processor_coordinates(self.shape, pnum)
-    ret = []
-    for dim_size, mesh_axis in zip(
-        tensor_shape.to_integer_list, tensor_layout.tensor_axis_to_mesh_axis):
-      if mesh_axis is None:
-        ret.append(0)
-      else:
-        ret.append(
-            dim_size // self.shape[mesh_axis].size * coordinates[mesh_axis])
-    return ret
-
-  def laid_out_size(self, tensor_shape):
-    """Total size of all slices.
-
-    Args:
-      tensor_shape: Shape.
-
-    Returns:
-      int.
-    """
-    return list_product(self.slice_shape(tensor_shape)) * self.size
-
-  def slicewise(self, fn, *inputs):
-    """Executes a function in parallel on all slices.
-
-    Args:
-      fn: function from tf.Tensors to tf.Tensor or a tuple of tf.Tensors.
-      *inputs: list of inputs.  Each input is either a LaidOutTensor or
-        is convertible to a tf.Tensor.
-
-    Returns:
-      LaidOutTensor, or a tuple of LaidOutTensors if fn returns a tuple.
-    """
-    raise NotImplementedError("Slicewise not implemented")
-
-  def Print(self, x, data, message, **kwargs):  # pylint: disable=invalid-name
-    """Calls tf.Print.
-
-    Args:
-      x: LaidOutTensor.
-      data: list of LaidOutTensor.
-      message: str.
-      **kwargs: keyword arguments to tf.print.
-
-    Returns:
-      LaidOutTensor.
-    """
-    del data, message, kwargs
-    tf.logging.warning("Warning - mtf.Print not implemented for this mesh type")
-    return x
-
-  def allreduce(self, x, mesh_axes, reduction_fn_string):
-    """Grouped allreduce, (summed across the given dimensions).
-
-    Args:
-      x: LaidOutTensor.
-      mesh_axes: list of integers, the mesh dimensions to be reduced.
-      reduction_fn_string: "SUM" or "MAX".
-
-    Returns:
-      LaidOutTensor.
-    """
-    raise NotImplementedError("Allreduce not implemented")
-
-  def allsplit(self, x, mesh_axis, split_axis, which=None):
-    """Inverse of allconcat - split each slice and keep only one piece of it.
-
-    The number of ways to split is the number of processors in the group.
-    The part that is kept corresponds to the processor's index in the group.
-
-    Args:
-      x: LaidOutTensor.
-      mesh_axis: int, the mesh axis along which to split.
-      split_axis: int, the Tensor axis along which to split.
-      which: an optional LaidOutTensor of integer scalars. Selects the slice to
-        to keep, instead of the coordinate.
-
-    Returns:
-      LaidOutTensor.
-    """
-    if which is None:
-      which = self.laid_out_pcoord(mesh_axis)
-    num_splits = self.shape[mesh_axis].size
-    def my_fn(x, which):
-      slice_begin = [
-          dimsize // num_splits * which if i == split_axis
-          else 0 for i, dimsize in enumerate(x.shape.as_list())]
-      slice_size = [
-          dimsize // num_splits if i == split_axis
-          else dimsize for i, dimsize in enumerate(x.shape.as_list())]
-      return tf.slice(x, slice_begin, slice_size)
-    return self.slicewise(my_fn, x, which)
-
-  def allconcat(self, x, mesh_axis, concat_axis):
-    """Grouped allconcat (like MPI allgather followed by concat).
-
-    Args:
-      x: LaidOutTensor.
-      mesh_axis: int, the mesh axis along which to group.
-      concat_axis: int, the Tensor axis along which to concatenate.
-
-    Returns:
-      LaidOutTensor.
-    """
-    raise NotImplementedError("Allconcat not implemented")
-
-  def alltoall(self, x, mesh_axis, split_axis, concat_axis):
-    """Grouped alltoall (like MPI alltoall with splitting and concatenation).
-
-    Args:
-      x: LaidOutTensor.
-      mesh_axis: int, the mesh axis along which to group.
-      split_axis: int, the Tensor axis along which to split.
-      concat_axis: int, the Tensor axis along which to concatenate.
-
-    Returns:
-      LaidOutTensor.
-    """
-    raise NotImplementedError("Alltoall not implemented")
-
-  def receive(self, x, mesh_axis, source_pcoord):
-    """Collective receive in groups.
-
-    Each group contains the processors that differ only in mesh_axis.
-
-    ```python
-    group_size = self.shape[mesh_axis].size
-    ```
-
-    Args:
-      x: a LaidOutTensor
-      mesh_axis: an integer
-      source_pcoord: a list of optional integers. Each element is either None
-        or an integer in [0, group_size). If source_pcoord[k] is None, then the
-        output for the k-th processor in each group is a zero tensor. If
-        source_pcoord[k] is not None, then the output for the k-th processor in
-        each group is equal to the input for the source_pcoord[k]-th processor
-        in that group.
-
-    Returns:
-      a LaidOutTensor
-    """
-    raise NotImplementedError("Alltoall not implemented")
-
-  def shift_by_n_processors(self, x, mesh_axis, offset, wrap):
-    """Receive the slice from processor pcoord - offset.
-
-    Args:
-      x: a LaidOutTensor
-      mesh_axis: an integer
-      offset: an integer
-      wrap: a boolean. If True, then wrap around. Otherwise, pad with zeros.
-    """
-    n = self.shape[mesh_axis].size
-    source_pcoord = []
-    for i in xrange(n):
-      c = i - offset
-      if c != c % n:
-        if wrap:
-          c = c % n
-        else:
-          c = None
-      source_pcoord.append(c)
-    return self.receive(x, mesh_axis, source_pcoord)
-
-  def laid_out_pnum(self):
-    """Returns a LaidOutTensor containing the processor number.
-
-    Returns:
-      LaidOutTensor where each slice is an integer scalar.
-    """
-    raise NotImplementedError("laid_out_pnum not implemented")
-
-  def laid_out_pcoord(self, mesh_axis):
-    """Returns a LaidOutTensor containing the processor coordinate.
-
-    Args:
-      mesh_axis: int.
-
-    Returns:
-      LaidOutTensor where each slice is an integer scalar.
-    """
-    divisor = list_product(self.shape.to_integer_list[mesh_axis + 1:])
-    modulus = self.shape[mesh_axis].size
-    def my_fn(pnum):
-      return (pnum // divisor) % modulus
-    return self.slicewise(my_fn, self.laid_out_pnum())
-
-  def broadcast_impl(self, old_slices, old_shape, new_shape):
-    """Implementation of a broadcast operation.
-
-    Args:
-      old_slices: LaidOutTensor.
-      old_shape: Shape.
-      new_shape: Shape.
-
-    Returns:
-      LaidOutTensor.
-    """
-    new_slice_shape = self.slice_shape(new_shape)
-    def tf_fn(x):
-      return (tf.zeros(new_slice_shape, dtype=x.dtype) +
-              _expand_dims(x, old_shape, new_shape))
-    return self.slicewise(tf_fn, old_slices)
-
-  def make_slices(self, tf_tensor, tensor_shape):
-    """Turns a single tf.Tensor into a list of slices, one for each processor.
-
-    Args:
-      tf_tensor: tf.Tensor.
-      tensor_shape: Shape.
-
-    Returns:
-      list of tf.tensor with length self.size.
-    """
-    tensor_layout = self.tensor_layout(tensor_shape)
-    slice_shape = self.slice_shape(tensor_shape)
-    def my_fn(pnum):
-      if tensor_layout.is_fully_replicated:
-        return tf_tensor
-      else:
-        slice_begin = self.slice_begin(tensor_shape, pnum)
-        return tf.slice(tf_tensor, slice_begin, slice_shape)
-
-    return parallel([tf_tensor.device] * self.size, my_fn,
-                    list(xrange(self.size)))
-
-  def combine_slices(self, slices, tensor_shape, device=None):
-    """Turns a set of slices into a single tensor.
-
-    Args:
-      slices: list of tf.Tensor with length self.size.
-      tensor_shape: Shape.
-      device: optional str. If absent, we use the devices of the slices.
-
-    Returns:
-      tf.Tensor.
-    """
-    if tensor_shape.ndims == 0:
-      return slices[0]
-
-    ret = slices[:]
-    tensor_layout = self.tensor_layout(tensor_shape)
-    for mesh_dim, tensor_axis in zip(
-        self.shape, tensor_layout.mesh_axis_to_tensor_axis(self.ndims)):
-      slice_size = len(ret) // mesh_dim.size
-      if tensor_axis is None:
-        ret = ret[:slice_size]
-      else:
-        if device:
-          devices = [device] * slice_size
-        else:
-          devices = [ret[i].device for i in xrange(slice_size)]
-        concat_inputs = [[ret[i + slice_size * j]
-                          for j in xrange(mesh_dim.size)]
-                         for i in xrange(slice_size)]
-        ret = parallel(
-            devices, tf.concat, concat_inputs,
-            axis=[tensor_axis] * len(devices))
-    assert len(ret) == 1
-    return ret[0]
-
-  def export_to_tf_tensor(self, x, laid_out_x):
-    """Turns a Tensor into a tf.Tensor.
-
-    Args:
-      x: Tensor.
-      laid_out_x: LaidOutTensor.
-
-    Returns:
-      tf.Tensor.
-    """
-    raise NotImplementedError("export_to_tf_tensor not implemented")
-
-  def import_tf_tensor(self, x, tf_x):
-    """Imports a tf.Tensor, producing a LaidOutTensor.
-
-    Args:
-      x: Tensor.
-      tf_x: tf.Tensor.
-
-    Returns:
-      LaidOutTensor.
-    """
-    raise NotImplementedError("Import not implemented")
-
-
-class LazyAllreduceSum(object):
-  """Represents a LaidOutTensor with a lazy allreduce.
-
-  The purpose of delaying allreduce is that it saves bandwidth to first add
-  and then allreduce, as opposed to the other way around.
-  """
-
-  def __init__(self,
-               mesh_impl,
-               laid_out_input,
-               mesh_axes,
-               add_counter_fn=None):
-    """Create a LazyAllreduceSum.
-
-    Args:
-      mesh_impl: a mesh_impl
-      laid_out_input: a LaidOutTensor
-      mesh_axes: a list of mesh axes
-      add_counter_fn: a function taking no arguments which calls
-        lowering.add_counter if and when the allreduce executes.
-    Returns:
-      a LazyAllreduceSum
-    """
-    self.mesh_impl = mesh_impl
-    self.laid_out_input = laid_out_input
-    self.mesh_axes = mesh_axes
-    self._add_counter_fn = add_counter_fn
-    self._reduced = None
-
-  def to_laid_out_tensor(self):
-    if not self._reduced:
-      self._reduced = self.mesh_impl.allreduce(
-          self.laid_out_input, self.mesh_axes, "SUM")
-      if self._add_counter_fn:
-        self._add_counter_fn()
-    return self._reduced
-
-  def __add__(self, other):
-    """Add to another LazyAllreduceSum.
-
-    Args:
-      other: a LazyAllreduceSum or a LaidOutTensor
-    Returns:
-      a LazyAllreduceSum or a LaidOutTensor
-    """
-    if (isinstance(other, LazyAllreduceSum) and
-        self.mesh_impl == other.mesh_impl and
-        self.mesh_axes == other.mesh_axes):
-      return LazyAllreduceSum(
-          self.mesh_impl,
-          self.mesh_impl.slicewise(
-              tf.add, self.laid_out_input, other.laid_out_input),
-          self.mesh_axes,
-          add_counter_fn=self._add_counter_fn)
-    else:
-      return self.mesh_impl.slicewise(
-          tf.add, self.to_laid_out_tensor(), other.to_laid_out_tensor())
-
-  @property
-  def slice_shape(self):
-    return self.laid_out_input.slice_shape
-
-
-def convert_args_to_laid_out_tensors(xs):
-  """Convert list elements to laid-out-tensors when possible.
-
-  Args:
-    xs: a list
-  Returns:
-    a list
-  """
-  ret = []
-  for x in xs:
-    try:
-      ret.append(x.to_laid_out_tensor())
-    except AttributeError:
-      ret.append(x)
-  return ret
-
-
-class Tensor(object):
-  """A Distributed Tensor."""
-
-  def __init__(self, operation, shape, dtype, name=None):
-    if not isinstance(shape, Shape):
-      raise ValueError("shape must be a Shape got %s" % shape.to_string)
-    if not isinstance(dtype, tf.DType):
-      raise ValueError("dtype must be a tf.DType got %s" % dtype)
-    self._mesh = operation.mesh
-    self._operation = operation
-    self._shape = shape
-    self._dtype = dtype
-    if name is None:
-      name = self.operation.name
-    self._name = name
-    self._mesh.graph.tensors.append(self)
-
-  @property
-  def shape(self):
-    return self._shape
-
-  @property
-  def size(self):
-    return self.shape.size
-
-  @property
-  def mesh(self):
-    return self._mesh
-
-  @property
-  def graph(self):
-    return self._mesh.graph
-
-  @property
-  def operation(self):
-    return self._operation
-
-  @property
-  def dtype(self):
-    return self._dtype
-
-  @property
-  def name(self):
-    return self._name
-
-  def __repr__(self):
-    return self.to_string
-
-  def __add__(self, other):
-    return add(self, other)
-
-  def __radd__(self, other):
-    return add(self, other)
-
-  def __sub__(self, other):
-    return sub(self, other)
-
-  def __rsub__(self, other):
-    return sub(other, self)
-
-  def __mul__(self, other):
-    return multiply(self, other)
-
-  def __rmul__(self, other):
-    return multiply(self, other)
-
-  def __neg__(self):
-    return negative(self)
-
-  def __truediv__(self, other):
-    return divide(self, other)
-
-  def __rtruediv__(self, other):
-    return divide(other, self)
-
-  def __floordiv__(self, other):
-    return floordiv(self, other)
-
-  def __rfloordiv__(self, other):
-    return floordiv(other, self)
-
-  def __mod__(self, other):
-    return mod(self, other)
-
-  def __rmod__(self, other):
-    return mod(other, self)
-
-  @property
-  def to_string(self):
-    return "Tensor[%s, %s, %s]" % (self.name, self.shape.to_string, self.dtype)
-
-
-class Operation(object):
-  """A Distributed Operation."""
-
-  def __init__(self, inputs, mesh=None, name=None):
-    if mesh is None:
-      if not inputs:
-        raise ValueError("mesh must be specified if no inputs")
-      mesh = inputs[0].mesh
-    self._inputs = inputs
-    self._outputs = []
-    self._mesh = mesh
-    assert name is not None
-    scope_name = tf.get_variable_scope().name
-    if scope_name:
-      name = scope_name + "/" + name
-    self._name = name
-    mesh.graph.operations.append(self)
-
-  @property
-  def graph(self):
-    return self._mesh.graph
-
-  @property
-  def mesh(self):
-    return self._mesh
-
-  @property
-  def name(self):
-    return self._name
-
-  @property
-  def inputs(self):
-    return self._inputs[:]
-
-  @property
-  def outputs(self):
-    return self._outputs[:]
-
-  @property
-  def to_string(self):
-    return "%s[Inputs=(%s) Outputs=(%s)]" % (
-        type(self).__name__,
-        ", ".join([t.to_string for t in self.inputs]),
-        ", ".join([t.to_string for t in self.outputs]))
-
-  @property
-  def has_gradient(self):
-    return (
-        [t for t in self.inputs if t.dtype.is_floating] and
-        [t for t in self.outputs if t.dtype.is_floating])
-
-  def gradient(self, unused_grad_ys):
-    raise NotImplementedError("Gradient not implemented")
-
-  def lower(self, lowering):
-    raise NotImplementedError("Lower not implemented")
-
-
-class SlicewiseOperation(Operation):
-  """Apply any tensorflow function slice-wise.
-
-  Calls the Tensorflow function on each slice of the inputs to produce the
-  corresponding slice of the outputs.  Gradients are computed through
-  tensorflow.
-
-  The user must specify "splittable_dims": a list of Dimensions which can
-  be split while still keeping this computation valid.  For example, for
-  component-wise functions, all the dimensions are splittable, but if the
-  function is a reduction, the reduced dimensions are not splittable.
-  """
-
-  def __init__(self,
-               tf_fn,
-               inputs,
-               output_shape,
-               output_dtype,
-               splittable_dims,
-               grad_function=None,
-               name=None):
-    """Create a SlicewiseOperation.
-
-    grad_function is a python function taking this operation and a gradients
-    Tensor and producing input gradients tensors.
-    e.g.
-    def _square_grad(op, dy):
-      return [dy * op.inputs[0] * 2]
-
-    Args:
-      tf_fn: a function taking n tf.Tensors and returning a tf.Tensor
-      inputs: a list of n Tensors
-      output_shape: a Shape
-      output_dtype: a dtype
-      splittable_dims: a list of Dimensions which are ok to split
-      grad_function: an optional python function. Default to using tf.gradients
-      name: an optional string
-    """
-    super(SlicewiseOperation, self).__init__(inputs, name=name or "slicewise")
-    self._tf_fn = tf_fn
-    self._outputs = [Tensor(self, output_shape, output_dtype)]
-    self._splittable_dims = splittable_dims
-    self._grad_function = grad_function
-
-  def gradient(self, grad_ys):
-    if self._grad_function is not None:
-      return self._grad_function(self, grad_ys[0])
-    return GenericGradOperation(self, grad_ys).outputs
-
-  def lower(self, lowering):
-    # Check that only splittable dims are split
-    mesh_impl = lowering.mesh_impl(self)
-    for t in self.inputs + self.outputs:
-      layout = mesh_impl.tensor_layout(t)
-      for d, mesh_axis in zip(t.shape.dims, layout.tensor_axis_to_mesh_axis):
-        if mesh_axis is not None and d not in self._splittable_dims:
-          raise ValueError("dimension %s is not declared as splittable" % d)
-    lowering.set_tensor_lowering(
-        self.outputs[0],
-        mesh_impl.slicewise(
-            self._tf_fn, *[lowering.tensors[x] for x in self.inputs]))
-
-
-def slicewise(tf_fn,
-              xs,
-              output_shape=None,
-              output_dtype=None,
-              splittable_dims=None,
-              grad_function=None,
-              name=None):
-  """Slice-wise call to any tensorflow function.
-
-  The output shape and dtype default to those of the first input.
-  splittable_dims is a list of Dimensions which can be split while keeping the
-  computation valid.
-
-  Args:
-    tf_fn: a function taking n tf.Tensors and returning a tf.Tensor
-    xs: a list of n Tensors
-    output_shape: a Shape
-    output_dtype: a dtype
-    splittable_dims: a list of Dimensions which are ok to split
-    grad_function: an optional gradients function.  If None, use tf gradient.
-    name: an optional string
-
-  Returns:
-    a Tensor
-  """
-  return SlicewiseOperation(
-      tf_fn,
-      xs,
-      convert_to_shape(output_shape) or xs[0].shape,
-      output_dtype or xs[0].dtype,
-      splittable_dims,
-      grad_function,
-      name=name).outputs[0]
-
-
-def cwise(tf_fn, xs, output_dtype=None, grad_function=None, name=None):
-  """Component-wise operation with no broadcasting.
-
-  Args:
-    tf_fn: a component-wise function taking n tf.Tensor inputs and producing
-      a tf.Tensor output
-    xs: n Tensors
-    output_dtype: an optional dtype
-    grad_function: an optional python function
-    name: an optional string
-
-  Returns:
-    a Tensor
-  """
-  return slicewise(
-      tf_fn, xs, output_dtype=output_dtype, splittable_dims=xs[0].shape.dims,
-      grad_function=grad_function, name=name or "cwise")
-
-
-def square(x, name="square"):
-  return cwise(
-      tf.square, [x], name=name,
-      grad_function=lambda op, dy: [dy * op.inputs[0] * 2])
-
-
-def sqrt(x, name="sqrt"):
-  return cwise(
-      tf.sqrt, [x], name=name,
-      grad_function=lambda op, dy: [dy * 0.5 / op.outputs[0]])
-
-
-def _rsqrt_grad(op, dy):
-  return [dy * -0.5 * op.outputs[0] * op.outputs[0] * op.outputs[0]]
-
-
-def rsqrt(x, name="rsqrt"):
-  return cwise(
-      tf.rsqrt, [x], name=name, grad_function=_rsqrt_grad)
-
-
-def log(x, name="log"):
-  return cwise(
-      tf.log, [x], name=name,
-      grad_function=lambda op, dy: [dy / op.inputs[0]])
-
-
-def exp(x, name="exp"):
-  return cwise(tf.exp, [x], name=name,
-               grad_function=lambda op, dy: [dy * op.outputs[0]])
-
-
-def sigmoid(x, name="sigmoid"):
-  def grad_function(op, dy):
-    y = op.outputs[0]
-    return [y * (1.0 - y) * dy]
-  return cwise(tf.sigmoid, [x], name=name, grad_function=grad_function)
-
-
-def tanh(x, name="tanh"):
-  def grad_function(op, dy):
-    y = op.outputs[0]
-    return [(1.0 - square(y)) * dy]
-  return cwise(tf.tanh, [x], name=name, grad_function=grad_function)
-
-
-def pow(x, y):  # pylint: disable=redefined-builtin
-  return exp(log(x) * y)
-
-
-def negative(x, name="negative"):
-  return cwise(tf.negative, [x], name=name,
-               grad_function=lambda op, dy: [negative(dy)])
-
-
-def logical_not(x, name="logical_not"):
-  return cwise(tf.logical_not, [x], name=name)
-
-
-def reciprocal(x, name="reciprocal"):
-  return cwise(
-      tf.reciprocal, [x], name=name,
-      grad_function=lambda op, dy: [negative(dy * square(op.outputs[0]))])
-
-
-def _relu_grad(op, dy):
-  return [dy * cast(greater(op.inputs[0], 0), op.inputs[0].dtype)]
-
-
-def relu(x, name="relu"):
-  return cwise(tf.nn.relu, [x], name=name, grad_function=_relu_grad)
-
-
-def cast(x, dtype, name="cast"):
-  if dtype == x.dtype:
-    return x
-  return cwise(
-      lambda x: tf.cast(x, dtype), [x], output_dtype=dtype, name=name,
-      grad_function=lambda op, dy: [cast(dy, op.inputs[0].dtype)])
-
-
-def to_float(x, name="to_float"):
-  return cast(x, tf.float32, name=name)
-
-
-def to_int32(x, name="to_int32"):
-  return cast(x, tf.int32, name=name)
-
-
-class GenericGradOperation(Operation):
-  """Gradients that follow regular TF.
-
-  Calling tf.gradients multiple times seems really slow in python.
-  TODO(noam): can we speed this up using functions or some other method?
-  """
-
-  def __init__(self, forward_op, grad_ys, name=None):
-    # tf.logging.info("forward inp %s, operations %s, grad_ys: %s",
-    #                 forward_op.inputs, forward_op.outputs, grad_ys)
-    super(GenericGradOperation, self).__init__(
-        forward_op.inputs + forward_op.outputs + grad_ys,
-        name=name or "generic_grad")
-    self._grad_ys = grad_ys
-    self._forward_op = forward_op
-    self._outputs = [Tensor(self, x.shape, x.dtype) for x in forward_op.inputs]
-
-  def lower(self, lowering):
-    # lists of lists of tf.Tensor
-    all_ys = transpose_list_of_lists(
-        [lowering.tensors[y].tensor_list for y in self._forward_op.outputs])
-    all_xs = transpose_list_of_lists(
-        [lowering.tensors[x].tensor_list for x in self._forward_op.inputs])
-    all_grad_ys = transpose_list_of_lists(
-        [lowering.tensors[dy].tensor_list for dy in self._grad_ys])
-    all_grad_xs = [tf.gradients(ys=ys, xs=xs, grad_ys=grad_ys) for
-                   ys, xs, grad_ys in zip(all_ys, all_xs, all_grad_ys)]
-    grad_xs = transpose_list_of_lists(all_grad_xs)
-    for out, grad_x in zip(self.outputs, grad_xs):
-      lowering.set_tensor_lowering(
-          out,
-          lowering.mesh_impl(self).LaidOutTensor.from_tensor_list(grad_x))
-
-
-class ScalarMultiplyOperation(Operation):
-  """Multiply by a tf Scalar (no backprop to scalar)."""
-
-  def __init__(self, x, scalar, name=None):
-    super(ScalarMultiplyOperation, self).__init__(
-        [x], name=name or "scalar_mul")
-    self._outputs = [Tensor(self, x.shape, x.dtype)]
-    self._scalar = scalar
-
-  def gradient(self, grad_ys):
-    dy = grad_ys[0]
-    return [dy * self._scalar]
-
-  def lower(self, lowering):
-    lowering.set_tensor_lowering(
-        self.outputs[0],
-        lowering.mesh_impl(self).slicewise(
-            lambda x: x * self._scalar, lowering.tensors[self.inputs[0]]))
-
-
-class ScalarAddOperation(Operation):
-  """Add a tf Scalar (no backprop to scalar)."""
-
-  def __init__(self, x, scalar, name=None):
-    super(ScalarAddOperation, self).__init__([x], name=name or "scalar_add")
-    self._outputs = [Tensor(self, x.shape, x.dtype)]
-    self._scalar = scalar
-
-  def gradient(self, grad_ys):
-    return grad_ys
-
-  def lower(self, lowering):
-    lowering.set_tensor_lowering(
-        self.outputs[0],
-        lowering.mesh_impl(self).slicewise(
-            lambda x: x + self._scalar, lowering.tensors[self.inputs[0]]))
-
-
-class BinaryOpWithBroadcasting(Operation):
-  """Binary operation with broadcasting."""
-
-  def __init__(self, tf_fn, x1, x2, output_shape, output_dtype, name=None):
-    super(BinaryOpWithBroadcasting, self).__init__(
-        [x1, x2], name=name or "binary_op")
-    if x1.dtype != x2.dtype:
-      # If there is ever a binary operation with different operand types, then
-      # we should add an argument allow_different_operand_dtypes=False.
-      raise ValueError("Dtypes must be equal.")
-    assert isinstance(output_dtype, tf.DType)
-    self._outputs = [Tensor(self, output_shape, output_dtype)]
-    self._tf_fn = tf_fn
-
-  def gradient(self, unused_grad_ys):
-    raise ValueError("Gradient not implememnted")
-
-  def lower(self, lowering):
-    x1 = self.inputs[0]
-    x2 = self.inputs[1]
-    output = self.outputs[0]
-    laid_out_x1 = lowering.tensors[x1]
-    laid_out_x2 = lowering.tensors[x2]
-    mesh_impl = lowering.mesh_impl(self)
-    if x1.shape != output.shape:
-      laid_out_x1 = mesh_impl.slicewise(
-          _expand_dims, laid_out_x1, x1.shape, output.shape)
-    if x2.shape != output.shape:
-      laid_out_x2 = mesh_impl.slicewise(
-          _expand_dims, laid_out_x2, x2.shape, output.shape)
-    lowering.set_tensor_lowering(
-        self.outputs[0],
-        mesh_impl.slicewise(
-            self._tf_fn, laid_out_x1, laid_out_x2))
-
-
-def binary_arguments_to_tensors(x1, x2):
-  """Convert argument of a binary operation to Tensors.
-
-  Args:
-    x1: a Tensor or something convertible to a tf Scalar
-    x2: a Tensor or something convertible to a tf Scalar
-
-  Returns:
-    new_x1: a Tensor
-    new_x2: a Tensor
-
-  Raises:
-    ValueError: on failure
-  """
-  if not isinstance(x1, Tensor) and not isinstance(x2, Tensor):
-    raise ValueError("at least one of x1 and x2 must be an mtf Tensor")
-  elif isinstance(x1, Tensor) and isinstance(x2, Tensor):
-    return x1, x2
-  elif isinstance(x1, Tensor):
-    return x1, import_tf_tensor(
-        x1.mesh, tf.convert_to_tensor(x2, dtype=x1.dtype), Shape([]))
-  else:
-    return import_tf_tensor(x2.mesh, tf.convert_to_tensor(x1, dtype=x2.dtype),
-                            Shape([])), x2
-
-
-def binary_op_with_broadcasting(
-    tf_fn, x1, x2, output_shape=None, output_dtype=None):
-  x1, x2 = binary_arguments_to_tensors(x1, x2)
-  output_shape = _infer_binary_broadcast_shape(x1.shape, x2.shape, output_shape)
-  output_dtype = output_dtype or x1.dtype
-  assert isinstance(output_dtype, tf.DType)
-  return BinaryOpWithBroadcasting(
-      tf_fn, x1, x2, convert_to_shape(output_shape),
-      output_dtype).outputs[0]
-
-
-def less(x1, x2, output_shape=None):
-  return binary_op_with_broadcasting(
-      tf.less, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
-
-
-def greater(x1, x2, output_shape=None):
-  return binary_op_with_broadcasting(
-      tf.greater, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
-
-
-def less_equal(x1, x2, output_shape=None):
-  return binary_op_with_broadcasting(
-      tf.less_equal, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
-
-
-def greater_equal(x1, x2, output_shape=None):
-  return binary_op_with_broadcasting(
-      tf.greater_equal, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
-
-
-def equal(x1, x2, output_shape=None):
-  return binary_op_with_broadcasting(
-      tf.equal, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
-
-
-def not_equal(x1, x2, output_shape=None):
-  return binary_op_with_broadcasting(
-      tf.not_equal, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
-
-
-def logical_and(x1, x2, output_shape=None):
-  return binary_op_with_broadcasting(
-      tf.logical_and, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
-
-
-def logical_or(x1, x2, output_shape=None):
-  return binary_op_with_broadcasting(
-      tf.logical_or, x1, x2, output_dtype=tf.bool, output_shape=output_shape)
-
-
-def floordiv(x1, x2, output_shape=None):
-  output_dtype = x1.dtype if isinstance(x1, Tensor) else x2.dtype
-  return binary_op_with_broadcasting(
-      tf.floordiv, x1, x2, output_dtype=output_dtype, output_shape=output_shape)
-
-
-def mod(x1, x2, output_shape=None):
-  output_dtype = x1.dtype if isinstance(x1, Tensor) else x2.dtype
-  return binary_op_with_broadcasting(
-      tf.mod, x1, x2, output_dtype=output_dtype, output_shape=output_shape)
-
-
-class AddOperation(BinaryOpWithBroadcasting):
-  """Binary addition with broadcasting."""
-
-  def __init__(self, x1, x2, output_shape, name=None):
-    super(AddOperation, self).__init__(
-        tf.add, x1, x2, output_shape, x1.dtype, name=name or "add")
-
-  def gradient(self, grad_ys):
-    dy = grad_ys[0]
-    return [reduce_sum(dy, output_shape=self.inputs[0].shape),
-            reduce_sum(dy, output_shape=self.inputs[1].shape)]
-
-
-class MinMaxOperation(BinaryOpWithBroadcasting):
-  """Binary minimum/maximum with broadcasting."""
-
-  def __init__(self, tf_fn, x1, x2, output_shape, name=None):
-    super(MinMaxOperation, self).__init__(
-        tf_fn, x1, x2, output_shape, x1.dtype, name=name or "add")
-
-  def gradient(self, grad_ys):
-    dy = grad_ys[0]
-    return [dy * cast(equal(self.inputs[0], self.outputs[0]), dy.dtype),
-            dy * cast(equal(self.inputs[1], self.outputs[0]), dy.dtype)]
-
-
-def minimum(x1, x2, output_shape=None, name=None):
-  """Binary minimum with broadcsting.
-
-  Args:
-    x1: a Tensor
-    x2: a Tensor
-    output_shape: an optional Shape
-    name: an optional string
-  Returns:
-    a Tensor
-  """
-  output_shape = convert_to_shape(output_shape)
-  with tf.name_scope(name, default_name="minimum"):
-    x1, x2 = binary_arguments_to_tensors(x1, x2)
-    return MinMaxOperation(
-        tf.minimum, x1, x2, output_shape=_infer_binary_broadcast_shape(
-            x1.shape, x2.shape, output_shape)).outputs[0]
-
-
-def maximum(x1, x2, output_shape=None, name=None):
-  """Binary maximum with broadcsting.
-
-  Args:
-    x1: a Tensor
-    x2: a Tensor
-    output_shape: an optional Shape
-    name: an optional string
-  Returns:
-    a Tensor
-  """
-  output_shape = convert_to_shape(output_shape)
-  with tf.name_scope(name, default_name="maximum"):
-    x1, x2 = binary_arguments_to_tensors(x1, x2)
-    return MinMaxOperation(
-        tf.maximum, x1, x2, output_shape=_infer_binary_broadcast_shape(
-            x1.shape, x2.shape, output_shape)).outputs[0]
-
-
-class BroadcastOperation(Operation):
-  """Broadcast - output dims are a superset of input dims, in any order."""
-
-  def __init__(self, x, output_shape, name=None):
-    super(BroadcastOperation, self).__init__([x], name=name or "broadcast")
-    self._outputs = [Tensor(self, output_shape, x.dtype)]
-
-  def gradient(self, grad_ys):
-    return [reduce_sum(grad_ys[0], output_shape=self.inputs[0].shape)]
-
-  def lower(self, lowering):
-    ret = lowering.mesh_impl(self).broadcast_impl(
-        lowering.tensors[self.inputs[0]], self.inputs[0].shape,
-        self.outputs[0].shape)
-    lowering.set_tensor_lowering(self.outputs[0], ret)
-
-
-def broadcast(x, new_shape):
-  return BroadcastOperation(x, new_shape).outputs[0]
-
-
-def _reduce_helper(input_shape,
-                   output_shape,
-                   input_tensor_layout,
-                   reduction_fn_string="SUM"):
-  """Returns slicewise function and reduced mesh dimensions.
-
-  Args:
-    input_shape: a Shape
-    output_shape: a Shape
-    input_tensor_layout: a TensorLayout
-    reduction_fn_string: "SUM" or "MAX"
-  Returns:
-    reduce_slice_fn: a function from tf.Tensor to tf.Tensor
-    reduced_mesh_axes: a list of integers
-  """
-  reduce_dims_indices = [
-      i for i, d in enumerate(input_shape.dims) if d not in output_shape.dims]
-  reduced_input_shape = Shape([
-      d for d in input_shape.dims if d in output_shape.dims])
-  perm = [reduced_input_shape.dims.index(d) for d in output_shape.dims]
-  def reduce_slice_fn(xslice):
-    ret = xslice
-    if reduce_dims_indices:
-      ret = reduction_fn(reduction_fn_string)(xslice, reduce_dims_indices)
-    if perm != list(xrange(len(perm))):
-      ret = tf.transpose(ret, perm)
-    return ret
-  reduced_mesh_axes = []
-  for i in reduce_dims_indices:
-    mesh_axis = input_tensor_layout[i]
-    if mesh_axis is not None:
-      reduced_mesh_axes.append(mesh_axis)
-  return reduce_slice_fn, reduced_mesh_axes
-
-
-class ReduceOperation(Operation):
-  """Reduction - output dims are a subset of input dims, in any order."""
-
-  def __init__(self, x, output_shape, reduction_fn_string, name=None):
-    super(ReduceOperation, self).__init__([x], name=name or "reduce")
-    self._outputs = [Tensor(self, output_shape, x.dtype)]
-    self._reduction_fn_string = reduction_fn_string
-
-  def gradient(self, grad_ys):
-    if self._reduction_fn_string == "SUM":
-      return [broadcast(grad_ys[0], self.inputs[0].shape)]
-    elif (self._reduction_fn_string == "MAX" or
-          self._reduction_fn_string == "MIN"):
-      return [cast(equal(self.inputs[0], self.outputs[0]), self.inputs[0].dtype)
-              * grad_ys[0]]
-    else:
-      raise ValueError("Gradients to other reductions not implemented")
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    slicewise_fn, reduced_mesh_axes = _reduce_helper(
-        self.inputs[0].shape, self.outputs[0].shape,
-        mesh_impl.tensor_layout(self.inputs[0]),
-        self._reduction_fn_string)
-    y = mesh_impl.slicewise(slicewise_fn, lowering.tensors[self.inputs[0]])
-    if reduced_mesh_axes:
-      def add_counter_fn():
-        lowering.add_counter("allreduce/%s/reduce_op" % reduced_mesh_axes,
-                             lowering.laid_out_size(self.outputs[0]))
-      if self._reduction_fn_string == "SUM":
-        y = LazyAllreduceSum(
-            mesh_impl, y, reduced_mesh_axes, add_counter_fn=add_counter_fn)
-      else:
-        y = mesh_impl.allreduce(
-            y, reduced_mesh_axes, self._reduction_fn_string)
-        add_counter_fn()
-    lowering.set_tensor_lowering(self.outputs[0], y)
-
-
-class ConcatOperation(Operation):
-  """tf.concat.
-
-  All inputs have the same shape, except for the size of the dimension named
-  dim_name.
-  """
-
-  def __init__(self, xs, concat_dim_name, name=None):
-    super(ConcatOperation, self).__init__(xs, name=name or "concat")
-    # verify that the shapes are all compatible
-    dim_names = [dim.name for dim in xs[0].shape.dims]
-    self._concat_dim_name = concat_dim_name
-
-    if concat_dim_name not in dim_names:
-      raise ValueError("xs[0] does not contain a dimension named dim_name")
-    self._axis = dim_names.index(concat_dim_name)
-
-    should_be_equal = [
-        x.shape.resize_dimension(concat_dim_name, 0) for x in xs]
-    if not all(s == should_be_equal[0] for s in should_be_equal):
-      raise ValueError("shapes are not compatible %s" % xs)
-
-    self._input_sizes = [x.shape.dims[self._axis].size for x in xs]
-    output_size = sum(self._input_sizes)
-    self._outputs = [
-        Tensor(self, xs[0].shape.resize_dimension(concat_dim_name, output_size),
-               xs[0].dtype)]
-
-  def gradient(self, grad_ys):
-    dy = grad_ys[0]
-    return split(dy, self.outputs[0].shape.dims[self._axis], self._input_sizes)
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    if mesh_impl.tensor_dimension_to_mesh_axis(
-        Dimension(self._concat_dim_name, 0)) is not None:
-      raise ValueError("can't concat along split axis")
-    def slicewise_fn(*args):
-      return tf.concat(args, axis=self._axis, name="concat")
-    y = mesh_impl.slicewise(
-        slicewise_fn, *[lowering.tensors[x] for x in self._inputs])
-    lowering.set_tensor_lowering(self.outputs[0], y)
-
-
-def concat(xs, concat_dim_name, name=None):
-  """Like tf.concat.
-
-  All inputs must have equal shape except for the sizes in the concatenated
-  dimension.  The dimension names should be the same, even that of the
-  concatenated dimension.
-
-  Args:
-    xs: a list of Tensors
-    concat_dim_name: a string
-    name: an optional string
-  Returns:
-    a Tensor
-  """
-  return ConcatOperation(xs, concat_dim_name, name).outputs[0]
-
-
-class SplitOperation(Operation):
-  """like tf.split.
-
-  TODO(noam, nikip): this code has never been run.  Run it and test it.
-  """
-
-  def __init__(self, x, split_dim, num_or_size_splits, name=None):
-    super(SplitOperation, self).__init__([x], name=name or "split")
-
-    self._split_dim = split_dim
-    if split_dim not in x.shape.dims:
-      raise ValueError("%s does not contain dimension %s" % (x, split_dim))
-    self._axis = x.shape.dims.index(split_dim)
-
-    if isinstance(num_or_size_splits, list):
-      self._output_sizes = num_or_size_splits
-      if sum(num_or_size_splits) != split_dim.size:
-        raise ValueError(
-            "Sizes do not add up %s %s" % (num_or_size_splits, split_dim))
-    else:
-      assert isinstance(num_or_size_splits, int)
-      assert split_dim.size % num_or_size_splits == 0
-      self._output_sizes = (
-          [split_dim.size // num_or_size_splits] * num_or_size_splits)
-
-    self._outputs = [
-        Tensor(self, x.shape.resize_dimension(split_dim.name, output_size),
-               x.dtype) for output_size in self._output_sizes]
-
-  def gradient(self, grad_ys):
-    return [concat(grad_ys, self._split_dim.name)]
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._split_dim) is not None:
-      raise ValueError("can't split along split axis")
-    def slicewise_fn(x):
-      # Since we return a tuple of tf.Tensor, slicewise will collate the
-      # outputs and return a tuple of LaidOutTensors.
-      return tuple(tf.split(x, self._output_sizes, axis=self._axis))
-    values = mesh_impl.slicewise(
-        slicewise_fn, lowering.tensors[self.inputs[0]])
-    for t, v in zip(self._outputs, values):
-      lowering.set_tensor_lowering(t, v)
-
-
-def split(x, split_dim, num_or_size_splits, name=None):
-  """Like tf.split.
-
-  Args:
-    x: a Tensor
-    split_dim: a Dimension in x.shape.dims
-    num_or_size_splits: either an integer dividing split_dim.size
-       or a list of integers adding up to split_dim.size
-    name: an optional string
-  Returns:
-    a list of Tensors.
-  """
-  return SplitOperation(x, split_dim, num_or_size_splits, name=name).outputs
-
-
-class StackOperation(Operation):
-  """Like tf.stack."""
-
-  def __init__(self, xs, dim_name, axis, name=None):
-    super(StackOperation, self).__init__(xs, name=name or "stack")
-    self._axis = axis
-    self._new_dim = Dimension(dim_name, len(xs))
-    input_shape = xs[0].shape
-    for x in xs:
-      if x.shape != xs[0].shape:
-        raise ValueError(
-            "inputs to stack must have the same shape, got %s" % xs)
-    output_shape = Shape(
-        input_shape.dims[:axis] + [self._new_dim]+ input_shape.dims[axis:])
-    self._outputs = [Tensor(self, output_shape, xs[0].dtype)]
-
-  def gradient(self, grad_ys):
-    return unstack(grad_ys[0], self._new_dim)
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._new_dim) is not None:
-      raise ValueError("can't stack along split axis")
-    inputs = [lowering.tensors[t] for t in self._inputs]
-    def slicewise_fn(*args):
-      return tf.stack(args, axis=self._axis)
-    ret = mesh_impl.slicewise(slicewise_fn, *inputs)
-    lowering.set_tensor_lowering(self.outputs[0], ret)
-
-
-def stack(xs, dim_name, axis, name=None):
-  """Stack multiple Tensors to make a new dimension.
-
-  Args:
-    xs: a list of Tensors with identical shapes.
-    dim_name: a string (name of the new dimension)
-    axis: an integer (index of the new dimension in the output shape)
-    name: an optional string
-
-  Returns:
-    a Tensor
-  """
-  ret = StackOperation(xs, dim_name, axis, name).outputs[0]
-  return ret
-
-
-class UnstackOperation(Operation):
-  """Split into multiple Tensors, eliminating a dimension."""
-
-  def __init__(self, x, dim, name=None):
-    super(UnstackOperation, self).__init__([x], name=name or "unstack")
-    self._dim = dim
-    self._axis = x.shape.dims.index(dim)
-    output_shape = x.shape - dim
-    self._outputs = [
-        Tensor(self, output_shape, x.dtype) for _ in xrange(dim.size)]
-
-  def gradient(self, grad_ys):
-    return [stack(grad_ys, self._dim.name, self._axis)]
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._dim) is not None:
-      raise ValueError("can't unstack along split axis")
-    def slicewise_fn(x):
-      return tuple(tf.unstack(x, num=self._dim.size, axis=self._axis))
-    output_values = mesh_impl.slicewise(
-        slicewise_fn, lowering.tensors[self._inputs[0]])
-    for t, v in zip(self.outputs, list(output_values)):
-      lowering.set_tensor_lowering(t, v)
-
-
-def unstack(x, dim, name=None):
-  """Split into multiple Tensors, eliminating a dimension.
-
-  Args:
-    x: a Tensor
-    dim: a Dimension
-    name: an optional string
-
-  Returns:
-    a list of dim.size Tensors, each with shape (x.shape - dim)
-  """
-  return UnstackOperation(x, dim, name).outputs
-
-
-def cumsum(x, dim, exclusive=False):
-  """Cumulative sum.
-
-  Args:
-    x: a Tensor
-    dim: a Dimension
-    exclusive: a boolean
-
-  Returns:
-    a Tensor with the same shape as x.
-  """
-  new_name = "tmp_dim_cumsum"
-  new_dim = Dimension(new_name, dim.size)
-  new_shape = x.shape.rename_dimension(dim.name, new_name)
-  comparator = less if exclusive else less_equal
-  m = cast(
-      comparator(range(x.mesh, dim, dtype=tf.float32),
-                 range(x.mesh, new_dim, dtype=tf.float32)), x.dtype)
-  ret = einsum([x, m], output_shape=new_shape)
-  return reshape(ret, x.shape)
-
-
-def _einsum_helper(input_shapes, output_shape, mesh_impl):
-  """Returns slicewise function and reduced mesh dimensions.
-
-  Assumes the output shape contains no new dimensions.
-
-  Args:
-    input_shapes: a list of Shapes
-    output_shape: a Shape
-    mesh_impl: a MeshImpl
-  Returns:
-    einsum_slice_fn: a function from tf.Tensors to tf.Tensor
-    reduced_mesh_axes: a list of integers
-  """
-  input_shape_union = _shape_union(input_shapes)
-  total_num_dims = input_shape_union.ndims
-  # list of input shapes that contain all dimensions.
-  full_shapes = [
-      s for s in input_shapes + [output_shape] if s.ndims == total_num_dims]
-  full_shape = full_shapes[0] if full_shapes else input_shape_union
-  reduce_slice_fn, reduced_mesh_axes = _reduce_helper(
-      full_shape, output_shape, mesh_impl.tensor_layout(full_shape))
-  def einsum_slice_fn_naive(*slices):
-    # naive einsum implementation where we broadcst all inputs to the full
-    # shape, multiply componentwise, then reduce.
-    return reduce_slice_fn(reduce(tf.multiply, [
-        _expand_dims(x, input_shape, full_shape)
-        for x, input_shape in zip(slices, input_shapes)]))
-  if full_shapes:
-    # it is not wasteful of space to broadcast fully and then reduce.
-    # this helps to avoid some inefficient GPU implementations.
-    einsum_slice_fn = einsum_slice_fn_naive
-  else:
-    # call tf.einsum
-    equation = _einsum_equation(input_shapes, output_shape)
-    def einsum_slice_fn(*slices):
-      if slices[0].dtype.is_floating:
-        return tf.einsum(equation, *slices)
-      else:
-        return einsum_slice_fn_naive(*slices)
-  return einsum_slice_fn, reduced_mesh_axes
-
-
-class EinsumOperation(Operation):
-  """Einstein summation (matmul, etc).
-
-  The equation follows the dimensions in the input and output shapes.
-
-  Every dimension must occur in at least two of the input/output Tensors.
-  i.e. no new dimensions in the output, and no reduction of dimensions that
-  occur in only one input.
-  """
-
-  def __init__(self, inputs, output_shape, name=None):
-    super(EinsumOperation, self).__init__(inputs, name=name or "einsum")
-    if not inputs:
-      raise ValueError("Einsum needs at least one input")
-    for x in inputs:
-      if x.dtype != inputs[0].dtype:
-        raise ValueError("Input dtypes must be equal")
-    self._outputs = [Tensor(self, output_shape, inputs[0].dtype)]
-
-  def gradient(self, grad_ys):
-    dy = grad_ys[0]
-    xs = self.inputs
-    return [
-        einsum([dy] + [xs[j] for j in xrange(len(xs)) if j != i], xs[i].shape)
-        for i in xrange(len(self.inputs))]
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    xs = self.inputs
-    input_shape_set = set(sum([x.shape.dims for x in xs], []))
-    output_shape = self.outputs[0].shape
-    intersection_shape = Shape(
-        [d for d in output_shape.dims if d in input_shape_set])
-    einsum_slice_fn, reduced_mesh_axes = _einsum_helper(
-        [x.shape for x in self.inputs], intersection_shape, mesh_impl)
-    y = mesh_impl.slicewise(
-        einsum_slice_fn, *[lowering.tensors[x] for x in self.inputs])
-    if reduced_mesh_axes:
-      def add_counter_fn():
-        lowering.add_counter(
-            "allreduce/%s/einsum_op" % reduced_mesh_axes,
-            mesh_impl.laid_out_size(intersection_shape))
-      y = LazyAllreduceSum(
-          mesh_impl, y, reduced_mesh_axes, add_counter_fn=add_counter_fn)
-    # broadcast from intersection_shape to output_shape
-    if intersection_shape != output_shape:
-      y = mesh_impl.broadcast_impl(y, intersection_shape, output_shape)
-    lowering.set_tensor_lowering(self.outputs[0], y)
-    computation_shape = Shape(list(input_shape_set))
-    lowering.add_counter("einsum", mesh_impl.laid_out_size(computation_shape))
-    lowering.add_counter("einsum_unique", computation_shape.size)
-
-
-class Conv2dOperation(Operation):
-  """like tf.nn.conv2d.
-
-  Always data format "NHWC".
-  # TODO(nikip): support dilations
-  Always dilation rate of 1
-  padding: "SAME" or "VALID"
-
-  TODO(noam): implement more options.
-  """
-
-  def __init__(self, conv_input, conv_filter, strides, padding, name=None):
-    super(Conv2dOperation, self).__init__(
-        [conv_input, conv_filter], name=name or "conv2d")
-    self._padding = padding
-    self._batch_dims = conv_input.shape.dims[:-3]
-    self._in_h_dim, self._in_w_dim, self._in_dim = conv_input.shape.dims[-3:]
-    self._fh_dim, self._fw_dim = conv_filter.shape.dims[:2]
-    f_in_dim, self._out_dim = conv_filter.shape.dims[2:]
-    if f_in_dim != self._in_dim:
-      raise ValueError("Dimensions do not match input=%s filter=%s"
-                       % (conv_input, conv_filter))
-    out_h = self._in_h_dim.size
-    out_w = self._in_w_dim.size
-    if padding == "VALID":
-      out_h -= (self._fh_dim.size - 1)
-      out_w -= (self._fw_dim.size - 1)
-
-    self._strides = strides
-    if strides is not None:
-      out_h //= strides[1]
-      out_w //= strides[2]
-    self._out_h_dim = Dimension(self._in_h_dim.name, out_h)
-    self._out_w_dim = Dimension(self._in_w_dim.name, out_w)
-    output_shape = Shape(
-        self._batch_dims + [self._out_h_dim, self._out_w_dim, self._out_dim])
-    self._outputs = [Tensor(self, output_shape, conv_input.dtype)]
-
-  def gradient(self, grad_ys):
-    dy = grad_ys[0]
-    conv_input, conv_filter = self.inputs
-    return [
-        conv2d_backprop_input(self._inputs[0].shape,
-                              conv_filter,
-                              dy,
-                              self._strides,
-                              self._padding),
-        conv2d_backprop_filter(conv_input,
-                               self._inputs[1].shape,
-                               dy,
-                               self._strides,
-                               self._padding)]
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    conv_input, conv_filter = self.inputs
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._in_h_dim) is not None:
-      raise ValueError("can't slice along dimension h")
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._in_w_dim) is not None:
-      raise ValueError("can't slice along dimension w")
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._fh_dim) is not None:
-      raise ValueError("can't slice along dimension fh")
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._fw_dim) is not None:
-      raise ValueError("can't slice along dimension fw")
-    def tf_fn(tf_input, tf_filter):
-      output = tf.nn.conv2d(
-          _tf_flatten_batch_dims(tf_input, 3),
-          tf_filter, self._strides, self._padding)
-      return _tf_restore_batch_dims(output, 3, tf_input)
-    y = mesh_impl.slicewise(
-        tf_fn, lowering.tensors[conv_input], lowering.tensors[conv_filter])
-    # reducing out input channels - may need to allreduce
-    in_mesh_axis = mesh_impl.tensor_dimension_to_mesh_axis(self._in_dim)
-    if in_mesh_axis is not None:
-      def add_counter_fn():
-        lowering.add_counter(
-            "allreduce/%s/conv2d_op" % [in_mesh_axis],
-            mesh_impl.laid_out_size(self.outputs[0].shape))
-      y = LazyAllreduceSum(mesh_impl, y, [in_mesh_axis], add_counter_fn)
-    lowering.set_tensor_lowering(self.outputs[0], y)
-    computation_shape = _shape_union([conv_filter.shape, self.outputs[0].shape])
-    lowering.add_counter("conv2d/forward",
-                         mesh_impl.laid_out_size(computation_shape))
-    lowering.add_counter("conv2d_unique/forward", computation_shape.size)
-
-
-def conv2d(conv_input, conv_filter, strides, padding, name=None):
-  """conv2d."""
-  return Conv2dOperation(
-      conv_input, conv_filter, strides, padding, name=name).outputs[0]
-
-
-class Conv2dBackpropInputOperation(Operation):
-  """like tf.nn.conv2d_backprop_input"""
-
-  def __init__(self, input_shape, conv_filter, dy, strides, padding, name=None):
-    super(Conv2dBackpropInputOperation, self).__init__(
-        [dy, conv_filter], name=name or "conv2d_backprop")
-    self._padding = padding
-    self._strides = strides
-    self._input_shape = input_shape
-    self._outputs = [Tensor(self, input_shape, dy.dtype)]
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    dy, conv_filter = self.inputs
-    input_sizes = mesh_impl.slice_shape(self.outputs[0].shape)
-    input_sizes = [list_product(input_sizes[:-3])] + input_sizes[-3:]
-    def tf_fn(tf_dy, tf_filter):
-      return _tf_restore_batch_dims(
-          tf.nn.conv2d_backprop_input(
-              input_sizes, tf_filter, _tf_flatten_batch_dims(tf_dy, 3),
-              self._strides, self._padding), 3, tf_dy)
-    dx = mesh_impl.slicewise(
-        tf_fn, lowering.tensors[dy], lowering.tensors[conv_filter])
-    # reducing out output channels - may need to allreduce
-    out_mesh_axis = mesh_impl.tensor_dimension_to_mesh_axis(dy.shape.dims[-1])
-    if out_mesh_axis is not None:
-      def add_counter_fn():
-        lowering.add_counter(
-            "allreduce/%s/conv2d_op" % [out_mesh_axis],
-            mesh_impl.laid_out_size(self.outputs[0].shape))
-      dx = LazyAllreduceSum(mesh_impl, dx, [out_mesh_axis], add_counter_fn)
-    lowering.set_tensor_lowering(self.outputs[0], dx)
-    computation_shape = _shape_union([conv_filter.shape, dy.shape])
-    lowering.add_counter("conv2d/backprop_input",
-                         mesh_impl.laid_out_size(computation_shape))
-    lowering.add_counter("conv2d_unique/backprop_input", computation_shape.size)
-
-
-def conv2d_backprop_input(input_shape,
-                          conv_filter,
-                          dy,
-                          strides,
-                          padding, name=None):
-  return Conv2dBackpropInputOperation(input_shape,
-                                      conv_filter,
-                                      dy,
-                                      strides,
-                                      padding,
-                                      name=name).outputs[0]
-
-
-class Conv2dBackpropFilterOperation(Operation):
-  """like tf.nn.conv2d_backprop_input"""
-
-  def __init__(self, conv_input, filter_shape, dy, strides, padding, name=None):
-    super(Conv2dBackpropFilterOperation, self).__init__(
-        [conv_input, dy], name=name or "conv2d_backprop_filter")
-    self._padding = padding
-    self._strides = strides
-    self._filter_shape = filter_shape
-    self._outputs = [Tensor(self, filter_shape, dy.dtype)]
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    conv_input, dy = self.inputs
-    filter_sizes = mesh_impl.slice_shape(self.outputs[0].shape)
-    def tf_fn(tf_input, tf_dy):
-      return tf.nn.conv2d_backprop_filter(
-          _tf_flatten_batch_dims(tf_input, 3), filter_sizes,
-          _tf_flatten_batch_dims(tf_dy, 3), self._strides, self._padding)
-    df = mesh_impl.slicewise(
-        tf_fn, lowering.tensors[conv_input], lowering.tensors[dy])
-
-    # reducing out batch dimensions - may need to allreduce
-    reduced_mesh_axes = [
-        mesh_impl.tensor_dimension_to_mesh_axis(d)
-        for d in dy.shape.dims[:-3]]
-    reduced_mesh_axes = [a for a in reduced_mesh_axes if a is not None]
-
-    if reduced_mesh_axes:
-      def add_counter_fn():
-        lowering.add_counter(
-            "allreduce/%s/conv2d_backprop_filter" % (reduced_mesh_axes,),
-            mesh_impl.laid_out_size(self.outputs[0].shape))
-      df = LazyAllreduceSum(mesh_impl, df, reduced_mesh_axes, add_counter_fn)
-
-    lowering.set_tensor_lowering(self.outputs[0], df)
-    computation_shape = _shape_union([self.outputs[0].shape, dy.shape])
-    lowering.add_counter("conv2d/backprop_filter",
-                         mesh_impl.laid_out_size(computation_shape))
-    lowering.add_counter(
-        "conv2d_unique/backprop_filter", computation_shape.size)
-
-
-def conv2d_backprop_filter(conv_input,
-                           filter_shape,
-                           dy,
-                           strides,
-                           padding, name=None):
-  return Conv2dBackpropFilterOperation(conv_input,
-                                       filter_shape,
-                                       dy,
-                                       strides,
-                                       padding,
-                                       name=name).outputs[0]
-
-
-class ShiftOperation(Operation):
-  """Shift by a static offset in one dimension."""
-
-  def __init__(self, x, offset, dim, wrap, name=None):
-    """Create a shift operation.
-
-    Shift x right by +offset in dimension dim.
-    If offset is negative, shift left.
-    If wrap is true then wrap-around.  Else, pad with zeros.
-
-    Args:
-      x: a Tensor
-      offset: an integer
-      dim: a Dimension of x
-      wrap: a boolean - whether to wrap or pad.
-      name: an optional string
-    """
-    super(ShiftOperation, self).__init__([x], name=name or "shift")
-    self._dim = dim
-    self._axis = x.shape.dims.index(dim)
-    self._offset = offset
-    self._wrap = wrap
-    self._outputs = [Tensor(self, x.shape, x.dtype)]
-
-  def gradient(self, grad_ys):
-    return [shift(grad_ys[0], -self._offset, self._dim, self._wrap)]
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    mesh_axis = mesh_impl.tensor_dimension_to_mesh_axis(self._dim)
-    inputs = self._inputs[0]
-    ndims = self._inputs[0].shape.ndims
-    axis = self._axis
-    dim = self._dim
-    lowered_x = lowering.tensors[inputs]
-    def my_slice(x, start, size):
-      begin = [0] * axis + [start] + [0] * (ndims - axis - 1)
-      size = [-1] * axis + [size] + [-1] * (ndims - axis - 1)
-      return tf.slice(x, begin, size)
-    if mesh_axis is None:
-      def slicewise_fn(x):
-        """Slicewise function."""
-        def my_pad(s, begin_pad, end_pad):
-          paddings = ([[0, 0]] * axis + [[begin_pad, end_pad]]
-                      + [[0, 0]] * (ndims - axis - 1))
-          return tf.pad(s, paddings)
-        if self._wrap:
-          offset = self._offset % dim.size
-          return tf.concat([my_slice(x, dim.size - offset, offset),
-                            my_slice(x, 0, dim.size - offset)], axis=axis)
-        elif self._offset > 0:
-          return my_pad(
-              my_slice(x, 0, dim.size - self._offset), self._offset, 0)
-        else:
-          neg_offset = -self._offset
-          return my_pad(
-              my_slice(x, neg_offset, dim.size - neg_offset), 0, neg_offset)
-      lowered_y = mesh_impl.slicewise(slicewise_fn, lowered_x)
-    else:
-      mesh_dim_size = mesh_impl.shape.dims[mesh_axis].size
-      tensor_dim_size = self._dim.size
-      block_size = tensor_dim_size // mesh_dim_size
-      odiv = self._offset // block_size
-      omod = self._offset % block_size
-      laid_out_size = mesh_impl.laid_out_size(inputs.shape)
-      if omod == 0:
-        # shift by an integral number of processors.
-        lowered_y = mesh_impl.shift_by_n_processors(
-            lowered_x, mesh_axis, odiv, self._wrap)
-        lowering.add_counter("shift[%d]" % odiv, laid_out_size)
-      else:
-        # shift by odiv processors + omod positions
-        sliced = mesh_impl.slicewise(
-            lambda x: my_slice(x, 0, block_size - omod), lowered_x)
-        second_part = mesh_impl.shift_by_n_processors(
-            sliced, mesh_axis, odiv, self._wrap)
-        lowering.add_counter(
-            "shift[%d]" % odiv,
-            laid_out_size * (block_size - omod) // block_size)
-        sliced = mesh_impl.slicewise(
-            lambda x: my_slice(x, block_size - omod, omod), lowered_x)
-        first_part = mesh_impl.shift_by_n_processors(
-            sliced, mesh_axis, odiv + 1, self._wrap)
-        lowered_y = mesh_impl.slicewise(
-            lambda a, b: tf.concat([a, b], axis), first_part, second_part)
-        lowering.add_counter(
-            "shift[%d]" % (odiv + 1), laid_out_size * omod // block_size)
-    lowering.set_tensor_lowering(self.outputs[0], lowered_y)
-
-
-def shift(x, offset, dim, wrap, name=None):
-  """Shift operation.
-
-  Shift x right by +offset in dimension dim.
-
-  Args:
-    x: a Tensor
-    offset: an integer. If negative, shift left instead of right.
-    dim: a Dimension of x
-    wrap: a boolean - whether to wrap (True) or pad with zeros (False).
-    name: an optional string
-
-  Returns:
-    a Tensor with the same shape and dtype as x
-  """
-  return ShiftOperation(x, offset, dim, wrap, name=name).outputs[0]
-
-
-class SliceOperation(Operation):
-  """tf.slice.
-
-  We support the slice operation along one axis. Similar to tf.slice, specify
-  the begin and size values for the slice_dim.
-  """
-
-  def __init__(self, x, begin, size, slice_dim_name, name=None):
-    super(SliceOperation, self).__init__([x], name=name or "slice")
-    dim_names = x.shape.dimension_names
-    self._axis = axis = dim_names.index(slice_dim_name)
-    self._begin = begin
-    self._slice_dim = Dimension(slice_dim_name, size)
-    input_shape = self._inputs[0].shape
-    output_shape = Shape(
-        input_shape.dims[:axis] + [self._slice_dim] + input_shape.dims[axis+1:])
-    self._outputs = [Tensor(self, output_shape, x.dtype)]
-
-  def gradient(self, grad_ys):
-    actual_size = self._inputs[0].shape.dims[self._axis].size
-    return [
-        pad(grad_ys[0],
-            [self._begin, actual_size - self._slice_dim.size - self._begin],
-            self._slice_dim.name)]
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._slice_dim) is not None:
-      raise ValueError("can't slice along split axis")
-    inputs = self._inputs[0]
-    ndims = self._inputs[0].shape.ndims
-    axis = self._axis
-    begin = [0] * axis + [self._begin] + [0] * (ndims - axis - 1)
-    size = [-1] * axis + [self._slice_dim.size] + [-1] * (ndims - axis - 1)
-
-    def slicewise_fn(x, begin, size):
-      return tf.slice(x, begin, size, name="slice")
-    y = mesh_impl.slicewise(
-        slicewise_fn, lowering.tensors[inputs], begin, size)
-    lowering.set_tensor_lowering(self.outputs[0], y)
-
-
-class PadOperation(Operation):
-  """tf.pad.
-
-  Similar to tf.pad but we only pad along one axis given by pad_dim_name
-  with values specified by paddings. paddings is a list of two
-  values, giving the padding value before and after pad_dim.
-  """
-
-  def __init__(self, x, paddings, pad_dim_name, name=None):
-    super(PadOperation, self).__init__([x], name=name or "pad")
-    assert len(paddings) == 2
-    input_shape = self._inputs[0].shape
-    dim_names = [dim.name for dim in x.shape.dims]
-    if pad_dim_name not in dim_names:
-      raise ValueError("Padding dim name %s not found in input." % pad_dim_name)
-    self._paddings = paddings
-    self._axis = axis = dim_names.index(pad_dim_name)
-    output_size = input_shape.dims[axis].size + sum(paddings)
-    self._output_dim = Dimension(pad_dim_name, output_size)
-    output_shape = Shape(
-        input_shape.dims[:axis] +
-        [self._output_dim] + input_shape.dims[axis+1:])
-    self._outputs = [Tensor(self, output_shape, x.dtype)]
-
-  def gradient(self, grad_ys):
-    slice_dim_name = self._output_dim.name
-    slice_size = self._inputs[0].shape.dims[self._axis].size
-    return [slice(grad_ys[0], self._paddings[0], slice_size, slice_dim_name)]
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    if mesh_impl.tensor_dimension_to_mesh_axis(self._output_dim) is not None:
-      raise ValueError("can't pad along split axis")
-    inputs = self._inputs[0]
-    ndims = self._inputs[0].shape.ndims
-    axis = self._axis
-    paddings = [[0, 0]] * axis + [self._paddings] + [[0, 0]]* (ndims - axis - 1)
-
-    def slicewise_fn(x, paddings):
-      return tf.pad(x, paddings, name="pad")
-    y = mesh_impl.slicewise(
-        slicewise_fn, lowering.tensors[inputs], paddings)
-    lowering.set_tensor_lowering(self.outputs[0], y)
-
-
-class OneHotOperation(Operation):
-  """one_hot.
-  """
-
-  def __init__(self, indices, output_dim, on_value, off_value, dtype,
-               name=None):
-    super(OneHotOperation, self).__init__([indices], name=name or "one_hot")
-    if not indices.dtype.is_integer:
-      raise ValueError("indices requires an integer dtype got %s" % indices)
-    self._output_dim = output_dim
-    self._on_value = on_value
-    self._off_value = off_value
-    self._dtype = dtype
-    output_shape = Shape(indices.shape.dims + [output_dim])
-    self._outputs = [Tensor(self, output_shape, dtype)]
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    indices = self.inputs[0]
-    output_shape = self.outputs[0].shape
-    output_slice_shape = mesh_impl.slice_shape(output_shape)
-    mesh_axis = mesh_impl.tensor_dimension_to_mesh_axis(self._output_dim)
-    depth = output_slice_shape[-1]
-    if mesh_axis is None:
-      offset = 0
-    else:
-      offset = mesh_impl.slicewise(
-          tf.multiply, mesh_impl.laid_out_pcoord(mesh_axis), depth)
-
-    def slicewise_fn(indices_slice, offset):
-      return tf.one_hot(indices_slice - offset,
-                        depth,
-                        on_value=tf.cast(self._on_value, self._dtype),
-                        off_value=tf.cast(self._off_value, self._dtype),
-                        dtype=self._dtype)
-    y = mesh_impl.slicewise(
-        slicewise_fn, lowering.tensors[indices], offset)
-    lowering.set_tensor_lowering(self.outputs[0], y)
-
-
-class ImportOperation(Operation):
-  """Import a tf.Tensor onto a mesh."""
-
-  def __init__(self, mesh, tf_tensor, shape, name=None):
-    super(ImportOperation, self).__init__([], mesh=mesh, name=name or "import")
-    self._outputs = [Tensor(self, shape, tf_tensor.dtype)]
-    self._tf_tensor = tf_tensor
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    lowering.set_tensor_lowering(
-        self.outputs[0],
-        mesh_impl.import_tf_tensor(self.outputs[0], self._tf_tensor))
-
-
-def anonymous_shape(shape):
-  shape = convert_to_shape(shape)
-  return Shape([Dimension("_anonymous_%i" % i, d.size)
-                for i, d in enumerate(shape)])
-
-
-def anonymize(x):
-  return reshape(x, anonymous_shape(x.shape))
-
-
-def import_tf_tensor(mesh, tf_tensor, shape=None, name=None):
-  tf_tensor = tf.convert_to_tensor(tf_tensor)
-  if shape is None:
-    shape = Shape([])
-    assert not tf_tensor.shape.as_list()
-  return ImportOperation(
-      mesh, tf_tensor, convert_to_shape(shape), name=name).outputs[0]
-
-
-def import_fully_replicated(mesh, tf_tensor, shape, name=None):
-  return reshape(import_tf_tensor(
-      mesh, tf_tensor, anonymous_shape(shape), name), shape)
-
-
-class Variable(Operation):
-  """Variable."""
-
-  def __init__(self, mesh, name, shape, dtype, initializer,
-               trainable, **kwargs):
-    super(Variable, self).__init__([], mesh, name="name_will_be_set_later")
-    self._trainable = trainable
-    with tf.device(mesh.variable_placer_fn), mtf_utils.outside_all_rewrites():
-      self.master = tf.get_variable(
-          name, shape.to_integer_list, dtype=dtype, initializer=initializer,
-          **kwargs)
-    self._name = self.master.name[:self.master.name.find(":")]
-    self._outputs = [Tensor(self, shape, dtype)]
-    self.graph.all_variables.append(self)
-    if trainable:
-      self.graph.trainable_variables.append(self)
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    with mtf_utils.outside_all_rewrites():
-      sv = mesh_impl.LaidOutVariable(self, mesh_impl)
-    lowering.variables[self] = sv
-    lowering.set_tensor_lowering(self.outputs[0], sv.laid_out_tensor)
-    if self._trainable:
-      lowering.add_counter("variables/trainable", self.outputs[0].size)
-    else:
-      lowering.add_counter("variables/untrainable", self.outputs[0].size)
-
-  @property
-  def value(self):
-    return self.outputs[0]
-
-  @property
-  def shape(self):
-    return self.value.shape
-
-  @property
-  def dtype(self):
-    return self.value.dtype
-
-
-def get_variable(mesh, name, shape, dtype=tf.float32,
-                 initializer=None, trainable=True,
-                 activation_dtype=None, **kwargs):
-  ret = Variable(
-      mesh, name, convert_to_shape(shape), dtype, initializer,
-      trainable, **kwargs).outputs[0]
-  if activation_dtype and activation_dtype != dtype:
-    ret = cast(ret, activation_dtype)
-  return ret
-
-
-class Assign(Operation):
-  """Assign to a variable."""
-
-  def __init__(self, var, new_val, name=None):
-    super(Assign, self).__init__([new_val], var.mesh, name=name or "assign")
-    self._var = var
-    self._outputs = []
-
-  def lower(self, lowering):
-    lowering.operations[self] = lowering.variables[self._var].assign_to_slices(
-        lowering.tensors[self.inputs[0]].to_laid_out_tensor().all_slices)
-
-
-def assign(var, new_val):
-  """Assign a new value to a variable.
-
-  Args:
-    var: either a Variable operation or its output Tensor.
-    new_val: a Tensor
-  Returns:
-    an Operation
-  Raises:
-    ValueError: if var is not a Variable and var.operation is not a Variable
-  """
-  if isinstance(var, Tensor):
-    var = var.operation
-  if not isinstance(var, Variable):
-    raise ValueError("var must be a mtf.Variable or its output Tensor.")
-  return Assign(var, new_val)
-
-
-class Depend(Operation):
-  """Control dependency."""
-
-  def __init__(self, x, dependencies, name=None):
-    super(Depend, self).__init__([x], x.mesh, name=name or "depend")
-    for d in dependencies:
-      if not isinstance(d, Operation):
-        raise ValueError("dependencies must be mtf.Operations. got %s" % d)
-    self._dependencies = dependencies
-    self._outputs = [Tensor(self, x.shape, x.dtype)]
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    if not mesh_impl.supports_control_dependencies:
-      raise ValueError("Mesh does not suppport control dependencies.")
-    with tf.control_dependencies(
-        [lowering.operations[d] for d in self._dependencies]):
-      lowering.set_tensor_lowering(
-          self.outputs[0],
-          mesh_impl.slicewise(tf.identity,
-                              lowering.tensors[self.inputs[0]]))
-
-  def gradient(self, grad_ys):
-    return grad_ys
-
-
-def depend(x, dependencies):
-  """Identity of Tensor x that dependes on operations dependencies.
-
-  Args:
-    x: a Tensor
-    dependencies: a list of Operations
-  Returns:
-    an tensor
-  """
-  return Depend(x, dependencies).outputs[0]
-
-
-class Constant(Operation):
-  """A tensor where every element is the same constant value."""
-
-  def __init__(self, mesh, value, shape, dtype, name=None):
-    super(Constant, self).__init__([], mesh, name=name or "constant")
-    self._outputs = [Tensor(self, shape, dtype)]
-    self._value = value
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    slice_shape = mesh_impl.slice_shape(self.outputs[0].shape)
-    def tf_fn():
-      return tf.constant(value=self._value,
-                         dtype=self.outputs[0].dtype,
-                         shape=slice_shape)
-    lowering.set_tensor_lowering(self.outputs[0], mesh_impl.slicewise(tf_fn))
-
-
-def constant(mesh, value, shape=None, dtype=tf.float32):
-  shape = convert_to_shape(shape)
-  return Constant(mesh, value,
-                  shape if shape is not None else Shape([]),
-                  dtype).outputs[0]
-
-
-def zeros(mesh, shape, dtype=tf.float32):
-  return constant(mesh, 0, shape=convert_to_shape(shape), dtype=dtype)
-
-
-def zeros_like(t):
-  return zeros(t.mesh, t.shape, dtype=t.dtype)
-
-
-class StopGradient(Operation):
-  """Similar to tf.stop_gradient."""
-
-  def __init__(self, x, name=None):
-    super(StopGradient, self).__init__(
-        [x], x.mesh, name=name or "stop_gradient")
-    self._outputs = [Tensor(self, x.shape, x.dtype)]
-
-  def lower(self, lowering):
-    lowering.set_tensor_lowering(self.outputs[0],
-                                 lowering.tensors[self.inputs[0]])
-
-  @property
-  def has_gradient(self):
-    return False
-
-
-def stop_gradient(x):
-  return StopGradient(x).outputs[0]
-
-
-class PrintOperation(Operation):
-  """Similar to tf.stop_gradient."""
-
-  def __init__(self, x, data, message, name=None, **kwargs):
-    super(PrintOperation, self).__init__(
-        [x], x.mesh, name=name or "Print")
-    self._outputs = [Tensor(self, x.shape, x.dtype)]
-    self._data = data
-    self._message = message
-    self._kwargs = kwargs
-
-  def lower(self, lowering):
-    lowering.set_tensor_lowering(
-        self.outputs[0],
-        lowering.mesh_impl(self).Print(
-            lowering.tensors[self.inputs[0]],
-            [lowering.tensors[d].to_laid_out_tensor() for d in self._data],
-            self._message, **self._kwargs))
-
-  def gradient(self, grad_ys):
-    return grad_ys
-
-
-def Print(x, data, message, **kwargs):  # pylint: disable=invalid-name
-  """Call tf.Print.
-
-  Args:
-    x: a Tensor.
-    data: a list of Tensor
-    message: a string
-    **kwargs: keyword arguments to tf.Print
-  Returns:
-    a Tensor which is identical in value to x
-  """
-  return PrintOperation(x, data, message, **kwargs).outputs[0]
-
-
-class ReshapeOperation(Operation):
-  """Similar to tf.stop_gradient."""
-
-  def __init__(self, x, new_shape, name=None):
-    super(ReshapeOperation, self).__init__([x], x.mesh, name=name or "reshape")
-    self._outputs = [Tensor(self, new_shape, x.dtype)]
-
-  def lower(self, lowering):
-    """Lower the ReshapeOperation.
-
-    Reshaping can require collective communication between processors.
-    We haven't yet implemented all possible reshapes.  We try to handle the
-    common cases here - otherwise we raise a NotImplementedError.
-
-    Args:
-      lowering: a Lowering
-    Raises:
-      NotImplementedError: if we haven't covered this case
-    """
-    old_shape = self.inputs[0].shape
-    new_shape = self.outputs[0].shape
-    mesh_impl = lowering.mesh_impl(self)
-    slices = lowering.tensors[self.inputs[0]]
-
-    mesh_axis_to_cumprod_old = mesh_impl.mesh_axis_to_cumprod(old_shape)
-    mesh_axis_to_cumprod_new = mesh_impl.mesh_axis_to_cumprod(new_shape)
-    # Figure out what needs to be done for different mesh-axes
-    mesh_axes_allsplit = []
-    mesh_axes_allconcat = []
-    mesh_axes_alltoall = []
-    for mesh_axis, (old_cumprod, new_cumprod) in enumerate(
-        zip(mesh_axis_to_cumprod_old, mesh_axis_to_cumprod_new)):
-      if new_cumprod != old_cumprod:
-        if old_cumprod is None:
-          # split in new layout but not in old layout - we need an allsplit
-          mesh_axes_allsplit.append(mesh_axis)
-        elif new_cumprod is None:
-          # split in old layout but not in new layout - we need an allconcat
-          mesh_axes_allconcat.append(mesh_axis)
-        else:
-          # split differently in old and new layouts - we need an alltoall
-          mesh_axes_alltoall.append(mesh_axis)
-
-    laid_out_size = mesh_impl.laid_out_size(old_shape)
-
-    for mesh_axis in mesh_axes_allsplit:
-      tensor_axis = old_shape.cumprod_to_tensor_axis(
-          mesh_axis_to_cumprod_new[mesh_axis])
-      if tensor_axis is None:
-        # TODO(noam): try to handle this case
-        raise NotImplementedError(
-            "Try first reshaping to insert a new tf dimension,"
-            " then changing layout. input_shape=%s output_shape=%s"
-            % (self.inputs[0].shape, self.outputs[0].shape))
-      slices = mesh_impl.allsplit(slices, mesh_axis, tensor_axis)
-      laid_out_size //= mesh_impl.shape[mesh_axis].size
-    for mesh_axis in mesh_axes_alltoall:
-      split_tensor_axis = old_shape.cumprod_to_tensor_axis(
-          mesh_axis_to_cumprod_new[mesh_axis])
-      if split_tensor_axis is None:
-        # TODO(noam): try to handle this case
-        raise NotImplementedError(
-            "Try first reshaping to insert a new tf dimension,"
-            " then changing layout. input_shape=%s output_shape=%s"
-            % (self.inputs[0].shape, self.outputs[0].shape))
-      concat_tensor_axis = old_shape.cumprod_to_tensor_axis(
-          mesh_axis_to_cumprod_old[mesh_axis])
-      assert concat_tensor_axis is not None
-      slices = mesh_impl.alltoall(
-          slices, mesh_axis, split_tensor_axis, concat_tensor_axis)
-      lowering.add_counter(
-          "alltoall/%s/reshape_op" % mesh_axis, laid_out_size)
-
-    for mesh_axis in mesh_axes_allconcat:
-      tensor_axis = old_shape.cumprod_to_tensor_axis(
-          mesh_axis_to_cumprod_old[mesh_axis])
-      assert tensor_axis is not None
-      slices = mesh_impl.allconcat(slices, mesh_axis, tensor_axis)
-      laid_out_size *= mesh_impl.shape[mesh_axis].size
-      lowering.add_counter(
-          "allconcat/%s/reshape_op" % mesh_axis, laid_out_size)
-    # now reshape the slices
-    old_slice_shape = mesh_impl.slice_shape(old_shape)
-    new_slice_shape = mesh_impl.slice_shape(new_shape)
-    if new_slice_shape != old_slice_shape:
-      def reshape_fn(x):
-        return tf.reshape(x, new_slice_shape)
-      slices = mesh_impl.slicewise(reshape_fn, slices)
-    lowering.set_tensor_lowering(self.outputs[0], slices)
-
-  def gradient(self, grad_ys):
-    return [reshape(grad_ys[0], self.inputs[0].shape)]
-
-
-def reshape(x, new_shape):
-  return ReshapeOperation(x, convert_to_shape(new_shape)).outputs[0]
-
-
-def transpose(x, new_shape):
-  return einsum([x], output_shape=convert_to_shape(new_shape))
-
-
-def rename_dimension(x, old_name, new_name):
-  """Reshape a Tensor, renaming one dimension.
-
-  Args:
-    x: a Tensor
-    old_name: a string
-    new_name: a string
-
-  Returns:
-    a Tensor
-  """
-  return reshape(x, x.shape.rename_dimension(old_name, new_name))
-
-
-def einsum(xs, output_shape=None, name=None):
-  """Einstein summation.
-
-  If output_shape is not specified and there are two inputs, reduce over
-  all common dimensions and default the output shape to the unique dimensions
-  of the first input followed by the unique dimensions of the second input.
-
-  Args:
-    xs: a list of Tensors
-    output_shape: an optional Shape.
-    name: an optional string
-  Returns:
-    a Tensor
-  Raises:
-    ValueError: if the output shape cannot be inferred
-  """
-  output_shape = convert_to_shape(output_shape)
-  if output_shape is None:
-    if len(xs) == 2:
-      output_shape = Shape(
-          [d for d in xs[0].shape.dims if d not in xs[1].shape.dims] +
-          [d for d in xs[1].shape.dims if d not in xs[0].shape.dims])
-    else:
-      raise ValueError("could not infer einsum output_shape for inputs %s" %
-                       [x.to_string for x in xs])
-  return EinsumOperation(xs, output_shape, name=name).outputs[0]
-
-
-def matmul(a, b, output_shape=None, name=None):
-  return einsum([a, b], output_shape=output_shape, name=name)
-
-
-def _reduction_output_shape(x, output_shape, reduced_dim):
-  """Helper function to reduce_sum, etc."""
-  if output_shape is None:
-    if reduced_dim is None:
-      return Shape([])
-    else:
-      if reduced_dim not in x.shape.dims:
-        raise ValueError(
-            "reduced_dim=%s not in x.shape.dims=%s" % (reduced_dim, x.shape))
-      return x.shape - reduced_dim
-  elif reduced_dim is not None:
-    raise ValueError("do not specify both reduced_dim and output_shape")
-  else:
-    return output_shape
-
-
-def reduce_sum(x,
-               disable_positional_args=None,
-               output_shape=None,
-               reduced_dim=None,
-               name=None):
-  """Reduction on 1 or more axes.
-
-  If reduced_dim is present, then only that dimension is reduced out.
-  Alternatively, specify output_shape.
-  Do not specify both reduced_dim and output_shape.
-  If neither is specified, then all dimensions are reduced out.
-
-  Args:
-    x: a Tensor
-    disable_positional_args: None
-    output_shape: an optional Shape.  Must be a subsequence of x.shape.
-    reduced_dim: a mtf.Dimension
-    name: an optional string
-  Returns:
-    a Tensor
-  """
-  output_shape = convert_to_shape(output_shape)
-  reduced_dim = convert_to_dimension(reduced_dim)
-  assert disable_positional_args is None
-  output_shape = _reduction_output_shape(x, output_shape, reduced_dim)
-  if output_shape == x.shape:
-    return x
-  return ReduceOperation(x, output_shape, "SUM", name=name).outputs[0]
-
-
-def reduce_mean(x,
-                disable_positional_args=None,
-                output_shape=None,
-                reduced_dim=None,
-                name=None):
-  """Reduction on 1 or more axes.
-
-  If reduced_dim is present, then only that dimension is reduced out.
-  Alternatively, specify output_shape.
-  Do not specify both reduced_dim and output_shape.
-  If neither is specified, then all dimensions are reduced out.
-
-  Args:
-    x: a Tensor
-    disable_positional_args: None
-    output_shape: an optional Shape. Must be a subsequence of x.shape.
-    reduced_dim: a mtf.Dimension
-    name: an optional string
-
-  Returns:
-    a Tensor
-  """
-  output_shape = convert_to_shape(output_shape)
-  reduced_dim = convert_to_dimension(reduced_dim)
-  assert disable_positional_args is None
-  output_shape = _reduction_output_shape(x, output_shape, reduced_dim)
-  with tf.variable_scope(name, default_name="reduce_mean"):
-    if output_shape == x.shape:
-      return x
-    return reduce_sum(
-        x, output_shape=output_shape) * (output_shape.size / x.shape.size)
-
-
-def reduce_max(x,
-               disable_positional_args=None,
-               output_shape=None,
-               reduced_dim=None,
-               name=None):
-  """Reduction on 1 or more axes.
-
-  Args:
-    x: a Tensor
-    disable_positional_args: None
-    output_shape: an optional Shape.  Must be a subsequence of x.shape.
-    reduced_dim: an optional Dimension
-    name: an optional string
-  Returns:
-    a Tensor
-  """
-  output_shape = convert_to_shape(output_shape)
-  reduced_dim = convert_to_dimension(reduced_dim)
-  assert disable_positional_args is None
-  output_shape = _reduction_output_shape(x, output_shape, reduced_dim)
-  if output_shape is None:
-    output_shape = Shape([])
-  if output_shape == x.shape:
-    return x
-  return ReduceOperation(
-      x, output_shape, "MAX", name=name or "reduce_max").outputs[0]
-
-
-def reduce_min(x,
-               disable_positional_args=None,
-               output_shape=None,
-               reduced_dim=None,
-               name=None):
-  """Reduction on 1 or more axes.
-
-  Args:
-    x: a Tensor
-    disable_positional_args: None
-    output_shape: an optional Shape.  Must be a subsequence of x.shape.
-    reduced_dim: an optional Dimension
-    name: an optional string
-  Returns:
-    a Tensor
-  """
-  output_shape = convert_to_shape(output_shape)
-  reduced_dim = convert_to_dimension(reduced_dim)
-  assert disable_positional_args is None
-  output_shape = _reduction_output_shape(x, output_shape, reduced_dim)
-  if output_shape is None:
-    output_shape = Shape([])
-  if output_shape == x.shape:
-    return x
-  return ReduceOperation(
-      x, output_shape, "MIN", name=name or "reduce_min").outputs[0]
-
-
-def reduce_all(x,
-               disable_positional_args=None,
-               output_shape=None,
-               reduced_dim=None,
-               name=None):
-  output_shape = convert_to_shape(output_shape)
-  reduced_dim = convert_to_dimension(reduced_dim)
-  return cast(reduce_min(to_float(x),
-                         disable_positional_args=disable_positional_args,
-                         output_shape=output_shape,
-                         reduced_dim=reduced_dim,
-                         name=name or "reduce_all"), tf.bool)
-
-
-def reduce_any(x,
-               disable_positional_args=None,
-               output_shape=None,
-               reduced_dim=None,
-               name=None):
-  output_shape = convert_to_shape(output_shape)
-  reduced_dim = convert_to_dimension(reduced_dim)
-  return cast(reduce_max(to_float(x),
-                         disable_positional_args=disable_positional_args,
-                         output_shape=output_shape,
-                         reduced_dim=reduced_dim,
-                         name=name or "reduce_any"), tf.bool)
-
-
-def top_1(x, reduced_dim, dtype=tf.int32, name=None):
-  """Argmax and Max.
-
-  Args:
-    x: a Tensor
-    reduced_dim: a Dimension in x.shape.dims
-    dtype: a tf.dtype (for the output)
-    name: an optional string
-  Returns:
-    indices: a Tensor with given dtype
-    values: optional Tensor equal to mtf.reduce_max(x, reduced_dim=reduced_dim)
-  """
-  reduced_dim = convert_to_dimension(reduced_dim)
-  with tf.name_scope(name, default_name="top_1"):
-    max_val = reduce_max(x, reduced_dim=reduced_dim)
-    is_max = to_float(equal(x, max_val))
-    pos = range(x.mesh, reduced_dim, tf.float32)
-    ret = reduce_max(is_max * pos, reduced_dim=reduced_dim)
-    ret = cast(ret, dtype)
-    return ret, max_val
-
-
-def argmax(x, reduced_dim, dtype=tf.int32, name=None):
-  reduced_dim = convert_to_dimension(reduced_dim)
-  return top_1(x, reduced_dim, dtype, name)[0]
-
-
-def top_k(x, reduced_dim, new_dim, dtype=tf.int32, name=None):
-  """Like tf.top_k.
-
-  This operation returns two tensors with the same shape.  The output shape
-  is identical to the shape of x, except that reduced_dim is replaced by
-  new_dim.
-
-  Args:
-    x: a Tensor
-    reduced_dim: a Dimension in x.shape.dims.
-    new_dim: a Dimension.  The size determines k.
-    dtype: optional dtype for indices.
-    name: optional string.
-  Returns:
-    indices: a Tensor with given dtype.
-    values: a Tensor with same type as x.
-  """
-  reduced_dim = convert_to_dimension(reduced_dim)
-  new_dim = convert_to_dimension(new_dim)
-  indices = []
-  values = []
-  k = new_dim.size
-  with tf.name_scope(name, default_name="top_k"):
-    for i in xrange(k):
-      max_index, max_val = top_1(x, reduced_dim, dtype)
-      indices.append(max_index)
-      values.append(max_val)
-      if i + 1 < k:
-        x += one_hot(max_index, reduced_dim, on_value=-1e9)
-  axis = x.shape.dims.index(reduced_dim)
-  return stack(indices, new_dim.name, axis), stack(values, new_dim.name, axis)
-
-
-def sample_with_temperature(x, dim, temperature=1.0, dtype=tf.int32, name=None):
-  dim = convert_to_dimension(dim)
-  with tf.name_scope(name, default_name="sample_with_temperature"):
-    if temperature != 0.0:
-      # gumbel trick
-      g = -log(-log(random_uniform(x.mesh, x.shape, dtype=x.dtype)))
-      x += g * temperature
-    return argmax(x, dim, dtype, name)
-
-
-def add(x1, x2, output_shape=None, name=None):
-  """Binary addition with broadcsting.
-
-  Args:
-    x1: a Tensor
-    x2: a Tensor
-    output_shape: an optional Shape
-    name: an optional string
-  Returns:
-    a Tensor
-  """
-  output_shape = convert_to_shape(output_shape)
-  if not isinstance(x2, Tensor):
-    return ScalarAddOperation(x1, x2).outputs[0]
-  with tf.name_scope(name, default_name="add"):
-    x1, x2 = binary_arguments_to_tensors(x1, x2)
-    return AddOperation(
-        x1, x2, output_shape=_infer_binary_broadcast_shape(
-            x1.shape, x2.shape, output_shape)).outputs[0]
-
-
-def sub(x1, x2, output_shape=None, name=None):
-  """Binary subtraction with broadcsting.
-
-  Args:
-    x1: a Tensor
-    x2: a Tensor
-    output_shape: an optional Shape
-    name: an optional string
-  Returns:
-    a Tensor
-  """
-  output_shape = convert_to_shape(output_shape)
-  if not isinstance(x2, Tensor):
-    return ScalarAddOperation(x1, -x2).outputs[0]
-  with tf.name_scope(name, default_name="sub"):
-    x1, x2 = binary_arguments_to_tensors(x1, x2)
-    return add(x1, negative(x2), output_shape=output_shape)
-
-
-def multiply(x1, x2, output_shape=None, name=None):
-  """Binary multiplication with broadcsting.
-
-  Args:
-    x1: a Tensor
-    x2: a Tensor
-    output_shape: an optional Shape
-    name: an optional string
-  Returns:
-    a Tensor
-  """
-  if not isinstance(x2, Tensor):
-    return ScalarMultiplyOperation(x1, x2).outputs[0]
-  with tf.name_scope(name, default_name="mul"):
-    x1, x2 = binary_arguments_to_tensors(x1, x2)
-    return einsum(
-        [x1, x2],
-        output_shape=_infer_binary_broadcast_shape(
-            x1.shape, x2.shape, output_shape))
-
-
-def divide(x1, x2, output_shape=None, name=None):
-  """Binary division with broadcsting.
-
-  Args:
-    x1: a Tensor
-    x2: a Tensor
-    output_shape: an optional Shape
-    name: an optional string
-  Returns:
-    a Tensor
-  """
-  output_shape = convert_to_shape(output_shape)
-  if not isinstance(x2, Tensor):
-    return ScalarMultiplyOperation(x1, 1.0 / x2).outputs[0]
-  with tf.name_scope(name, default_name="divide"):
-    x1, x2 = binary_arguments_to_tensors(x1, x2)
-    return multiply(x1, reciprocal(x2), output_shape=output_shape)
-
-
-def slice(x, begin, size, slice_dim_name, name=None):  # pylint: disable=redefined-builtin
-  """Slice operation.
-
-  Args:
-    x: a list of Tensors
-    begin: integer, where to begin slicing from along the axis
-    size: integer, size to slice from axis.
-    slice_dim_name: string, dimension name of slicing axis.
-    name: an optional string
-  Returns:
-    a Tensor with shape extended by output_shape for the last axis.
-  """
-  return SliceOperation(
-      x, begin, size, slice_dim_name, name=name).outputs[0]
-
-
-def pad(x, paddings, dim_name, name=None):
-  """Slice operation.
-
-  Args:
-    x: a list of Tensors
-    paddings: list of integers of size 2, padding size before and after for dim.
-    dim_name: string, name for the padding dim
-    name: an optional string
-  Returns:
-    a Tensor with shape extended by output_shape for the last axis.
-  """
-  return PadOperation(
-      x, paddings, dim_name, name=name).outputs[0]
-
-
-def one_hot(indices, output_dim, on_value=1.0,
-            off_value=0.0, dtype=tf.float32, name=None):
-  """One hot operation.
-
-  Args:
-    indices: a Tensor
-    output_dim: a Dimension
-    on_value: Value taken when indices are on at a location, default 1
-    off_value: Value taken when indices are off at a location, default 0
-    dtype: a tf.DType
-    name: an optional string
-  Returns:
-    a Tensor with shape extended by output_dim for the last axis.
-  """
-  return OneHotOperation(
-      indices, output_dim, on_value, off_value, dtype, name=name).outputs[0]
-
-
-def gather(weights, indices, dim, output_shape=None):
-  """Shorthand for einsum([one_hot(indices, dim)], weights).
-
-  Args:
-    weights: a Tensor
-    indices: a Tensor with integer type
-    dim: a Dimension
-    output_shape: an optional mtf.Shape
-  Returns:
-    a Tensor
-  """
-  dim = convert_to_dimension(dim)
-  output_shape = convert_to_shape(output_shape)
-  if weights.dtype == tf.bool:
-    return cast(gather(to_float(weights), indices, dim, output_shape), tf.bool)
-  return einsum([one_hot(indices, dim, dtype=weights.dtype), weights],
-                output_shape=output_shape)
-
-
-def gradients(ys, xs, grad_ys=None):
-  """Compute gradients in dtf.
-
-  Args:
-    ys: a list of Tensors
-    xs: a list of Tensors
-    grad_ys: an optional list of Tensors
-
-  Returns:
-    grad_xs: a list of Tensors
-  """
-  graph = ys[0].graph
-  if not grad_ys:
-    grad_ys = [Constant(y.mesh, 1.0, y.shape, y.dtype).outputs[0] for y in ys]
-  # figure out what Tensors are downstream of xs
-  downstream = set(xs)
-  for op in graph.operations:
-    if op.has_gradient:
-      if set(op.inputs) & downstream:
-        downstream |= set(op.outputs)
-  tensor_to_gradient = dict(zip(ys, grad_ys))
-  for op in graph.operations[::-1]:
-    grad_outputs = [tensor_to_gradient.get(out) for out in op.outputs]
-    if op.has_gradient and any(grad_outputs) and (set(op.inputs) & downstream):
-      with tf.variable_scope(op.name + "/gradients"):
-        input_grads = op.gradient(grad_outputs)
-        for inp, grad in zip(op.inputs, input_grads):
-          if inp in downstream and grad is not None:
-            if inp in tensor_to_gradient:
-              tensor_to_gradient[inp] += grad
-            else:
-              tensor_to_gradient[inp] = grad
-  return [tensor_to_gradient.get(x, None) for x in xs]
-
-
-def _infer_binary_broadcast_shape(shape1, shape2, given_output_shape=None):
-  """Infer shape of the output of a binary op with broadcasting.
-
-  If the output shape is not given with given_output_shape, then we check
-  to see if one of the shapes is a subsequence of the other one, and we
-  return the one that is the supersequence.  Otherwise, we list the dimensions
-  of shape1, followed by all new dimensions in shape2.
-
-  Args:
-    shape1: a Shape
-    shape2: a Shape
-    given_output_shape: an optional Shape
-  Returns:
-    a Shape
-  """
-  shape1 = convert_to_shape(shape1)
-  shape2 = convert_to_shape(shape2)
-  given_output_shape = convert_to_shape(given_output_shape)
-  if given_output_shape is not None:
-    return given_output_shape
-  if is_subsequence(shape1.dims, shape2.dims):
-    return shape2
-  if is_subsequence(shape2.dims, shape1.dims):
-    return shape1
-  return Shape(
-      shape1.dims + [d for d in shape2.dims if d not in shape1.dims])
-
-
-def _expand_dims(x, input_shape, output_shape):
-  """Expand dimensions and transpose if necessary.
-
-  Args:
-    x: a tf.Tensor
-    input_shape: a Shape
-    output_shape: a Shape whose dimensions are a superset of
-      those in input_shape
-
-  Returns:
-    a tf.Tensor
-  """
-  verify_no_new_dims([output_shape], input_shape)
-  if input_shape == output_shape or input_shape.ndims == 0:
-    return x
-  perm = [input_shape.dims.index(d) for d in output_shape.dims
-          if d in input_shape.dims]
-  x = tf.transpose(x, perm)
-  for i, d in enumerate(output_shape.dims):
-    if d not in input_shape.dims:
-      x = tf.expand_dims(x, i)
-  return x
-
-
-def _einsum_equation(input_shapes, output_shape):
-  """Turn shapes into an einsum equation.
-
-  e.g. "ij,jk->ik"
-
-  Args:
-    input_shapes: a list of Shapes
-    output_shape: a Shape
-  Returns:
-    a string
-  """
-  ret = []
-  next_letter = ord("a")
-  dim_to_letter = {}
-  for shape_num, shape in enumerate(input_shapes + [output_shape]):
-    if shape_num == len(input_shapes):
-      ret.append("->")
-    elif shape_num > 0:
-      ret.append(",")
-    for d in shape.dims:
-      if d not in dim_to_letter:
-        dim_to_letter[d] = chr(next_letter)
-        next_letter += 1
-      ret.append(dim_to_letter[d])
-
-  return "".join(ret)
-
-
-def is_subsequence(short_seq, long_seq):
-  """Is short_seq a subsequence of long_seq."""
-  if not short_seq:
-    return True
-  pos = 0
-  for x in long_seq:
-    if pos == len(short_seq):
-      return True
-    if short_seq[pos] == x:
-      pos += 1
-  if pos == len(short_seq):
-    return True
-  return False
-
-
-def verify_no_new_dims(input_shapes, output_shape):
-  """Verifies that all dimensions in the output are in at least one input.
-
-  Args:
-    input_shapes: a list of Shapes
-    output_shape: a Shape
-  Raises:
-    ValueError: if there are new dimensions in the output.
-  """
-  all_input_dims = set(sum([s.dims for s in input_shapes], []))
-  all_output_dims = set(output_shape.dims)
-  if not all_output_dims.issubset(all_input_dims):
-    raise ValueError(
-        "No new dimensions allowed in output"
-        " input_shapes = %s output_shape= %s"
-        % ([s.dims for s in input_shapes], output_shape.dims))
-
-
-def pnum_to_processor_coordinates(mesh_shape, pnum):
-  """Coordinates of a processor in the mesh.
-
-  Args:
-    mesh_shape: a Shape
-    pnum: an integer less than len(mesh_shape)
-
-  Returns:
-    a list of integers with length len(mesh_shape)
-  """
-  ret = []
-  for dimsize in mesh_shape.to_integer_list[::-1]:
-    ret.append(pnum % dimsize)
-    pnum //= dimsize
-  return ret[::-1]
-
-
-def processor_coordinates_to_pnum(mesh_shape, coord):
-  """Inverse of pnum_to_processor_coordinates.
-
-  Args:
-    mesh_shape: a Shape
-    coord: a list of integers with length len(mesh_shape)
-
-  Returns:
-    an integer less than len(mesh_shape)
-  """
-  ret = 0
-  multiplier = 1
-  for c, d in zip(coord[::-1], mesh_shape.to_integer_list[::-1]):
-    ret += multiplier * c
-    multiplier *= d
-  return ret
-
-
-def pnum_to_group(mesh_shape, group_dims, pnum):
-  """Group number for grouped allreduce.
-
-  Args:
-    mesh_shape: a Shape
-    group_dims: a list of integers (the dimensions reduced over)
-    pnum: an integer
-
-  Returns:
-    an integer
-  """
-  coord = pnum_to_processor_coordinates(mesh_shape, pnum)
-  remaining_shape = Shape(
-      [d for i, d in enumerate(mesh_shape) if i not in group_dims])
-  remaining_coord = [d for i, d in enumerate(coord) if i not in group_dims]
-  return processor_coordinates_to_pnum(remaining_shape, remaining_coord)
-
-
-def processor_groups(mesh_shape, group_dims):
-  """Groups of processors which differ only in the given dimensions.
-
-  Args:
-    mesh_shape: a Shape
-    group_dims: a list of integers
-
-  Returns:
-    a list of lists of integers (processor numbers)
-  """
-  group_numbers = [
-      pnum_to_group(mesh_shape, group_dims, pnum)
-      for pnum in xrange(mesh_shape.size)]
-  ret = []
-  for pnum, g in enumerate(group_numbers):
-    while len(ret) <= g:
-      ret.append([])
-    ret[g].append(pnum)
-  return ret
-
-
-def list_product(l):
-  return reduce(mul, l, 1)
-
-
-def log_softmax(x, reduced_dim, name=None):
-  """log(softmax(x)).
-
-  Args:
-    x: a Tensor whose shape contains vocab_dim
-    reduced_dim: a Dimension
-    name: an optional string
-
-  Returns:
-    a Tensor with the same shape as x
-  """
-  reduced_dim = convert_to_dimension(reduced_dim)
-  with tf.variable_scope(name, default_name="log_softmax"):
-    reduced_shape = x.shape - reduced_dim
-    max_logit = reduce_max(stop_gradient(x), output_shape=reduced_shape)
-    x -= max_logit
-    exp_x = exp(x)
-    sum_exp_x = reduce_sum(exp_x, output_shape=reduced_shape)
-    log_denom = log(sum_exp_x)
-    return x - log_denom
-
-
-def softmax(x, reduced_dim, name=None):
-  with tf.variable_scope(name, default_name="softmax"):
-    return exp(log_softmax(x, reduced_dim))
-
-
-def range(mesh, dim, dtype, name=None):  # pylint: disable=redefined-builtin
-  """Create a 1d mesh tensor with a range from [0, dim.size).
-
-  Args:
-    mesh: a Mesh
-    dim: a Dimension
-    dtype: a tf.DType
-    name: an optional string
-
-  Returns:
-    a Tensor
-  """
-  dim = convert_to_dimension(dim)
-  with tf.variable_scope(name, default_name="range"):
-    return import_tf_tensor(
-        mesh, tf.range(dim.size, dtype=dtype), shape=Shape([dim]))
-
-
-def pretty_print_counters(counters):
-  """print counters hierarchically.
-
-  Each counter is a pair of a string and a number.
-  The string can have slashes, meaning that the number also counts towards
-  each prefix.  e.g.  "parameters/trainable" counts towards both "parameters"
-  and "parameters/trainable".
-
-  Args:
-    counters: a list of (string, number) pairs
-
-  Returns:
-    a string
-  """
-  totals = collections.defaultdict(int)
-  for (name, val) in counters:
-    prefixes = [name[:i] for i in xrange(len(name)) if name[i] == "/"] + [name]
-    for p in prefixes:
-      totals[p] += val
-  parts = []
-  for name, val in sorted(six.iteritems(totals)):
-    parts.append(" " * name.count("/") + "%s: %.3g" % (name, val))
-  return "\n".join(parts)
-
-
-def _parse_string_to_list_of_pairs(s, seconds_to_int=False):
-  r"""Parses a string into a list of pairs.
-
-  In the input string, each pair is separated by a colon, and the delimiters
-  between pairs are any of " ,.;".
-
-  e.g. "rows:32,cols:32"
-
-  Args:
-    s: str to parse.
-    seconds_to_int: Boolean. If True, then the second elements are returned
-      as integers;  otherwise they are strings.
-
-  Returns:
-    List of tuple pairs.
-
-  Raises:
-    ValueError: Badly formatted string.
-  """
-  ret = []
-  for p in [s.split(":") for s in re.sub("[,.;]", " ", s).split()]:
-    if len(p) != 2:
-      raise ValueError("bad input to _parse_string_to_list_of_pairs %s" % s)
-    if seconds_to_int:
-      ret.append((p[0], int(p[1])))
-    else:
-      ret.append(tuple(p))
-  return ret
-
-
-def parallel(devices, fn, *args, **kwargs):
-  """Call a function once on each device.
-
-  Args:
-    devices: a list of n devices
-    fn: a function
-    *args: arguments, each of which is a list of length n
-    **kwargs: keyword-args, each of which is a list of length n
-  Returns:
-    a list of length n
-  Raises:
-    ValueError: if the arguments are not all lists of length n
-  """
-  if not isinstance(devices, list):
-    raise ValueError("devices must be a list")
-  for x in list(args) + list(six.itervalues(kwargs)):
-    if not isinstance(x, list) or len(x) != len(devices):
-      raise ValueError(
-          "Argument not a list with same length as devices "
-          "arg=%s devices=%s %s %s" % (x, devices, len(x), len(devices)))
-  ret = []
-  for i, device in enumerate(devices):
-    with tf.device(device):
-      with tf.variable_scope("parallel_%d" % i):
-        my_args = [x[i] for x in args]
-        my_kwargs = {k: v[i] for k, v in six.iteritems(kwargs)}
-        ret.append(fn(*my_args, **my_kwargs))
-  return ret
-
-
-def transpose_list_of_lists(lol):
-  """Transpose a list of equally-sized python lists.
-
-  Args:
-    lol: a list of lists
-  Returns:
-    a list of lists
-  Raises:
-    ValueError: if list is empty
-  """
-  if not lol:
-    raise ValueError("cannot transpose the empty list")
-  return [list(x) for x in zip(*lol)]
-
-
-def binary_reduction_fn(reduction_fn_string):
-  if reduction_fn_string == "SUM":
-    return tf.add
-  elif reduction_fn_string == "MAX":
-    return tf.maximum
-  elif reduction_fn_string == "MIN":
-    return tf.minimum
-  else:
-    raise ValueError("Unknown reduction_fn_string %s" % reduction_fn_string)
-
-
-def reduction_fn(reduction_fn_string):
-  if reduction_fn_string == "SUM":
-    return tf.reduce_sum
-  elif reduction_fn_string == "MAX":
-    return tf.reduce_max
-  elif reduction_fn_string == "MIN":
-    return tf.reduce_min
-  else:
-    raise ValueError("Unknown reduction_fn_string %s" % reduction_fn_string)
-
-
-class MtfCheckpointSaverListener(tf.train.CheckpointSaverListener):
-  """Copy slices to masters before saving."""
-
-  def __init__(self, lowering):
-    self._op = lowering.copy_slices_to_masters()
-
-  def begin(self):
-    # You can add ops to the graph here.
-    tf.logging.info("Starting the session.")
-
-  def before_save(self, session, global_step_value):
-    # assigns
-    tf.logging.info("Before Save.")
-    session.run(self._op)
-    tf.logging.info("About to write a checkpoint")
-
-  def after_save(self, session, global_step_value):
-    tf.logging.info("Done writing checkpoint.")
-
-  def end(self, session, global_step_value):
-    tf.logging.info("Done with the session.")
-
-
-class MtfRestoreHook(tf.train.SessionRunHook):
-  """Copy masters to slices after restoring."""
-
-  def __init__(self, lowering):
-    self._lowering = lowering
-
-  def begin(self):
-    self._op = self._lowering.copy_masters_to_slices()
-
-  def after_create_session(self, session, coord):
-    session.run(self._op)
-
-
-class RandomOperation(Operation):
-  """Random operation such as tf.random_uniform."""
-
-  def __init__(self, mesh, shape, tf_fn, **kwargs):
-    super(RandomOperation, self).__init__(
-        [], mesh=mesh, name=kwargs.get("name", "random"))
-    self._tf_fn = tf_fn
-    self._kwargs = kwargs
-    self._outputs = [Tensor(self, shape, kwargs.get("dtype", tf.float32))]
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    output_shape = self.outputs[0].shape
-    lowering.set_tensor_lowering(self.outputs[0], (
-        mesh_impl.random(output_shape, self._tf_fn, self._kwargs)))
-
-
-def random_uniform(mesh, shape, **kwargs):
-  """Random uniform.
-
-  Args:
-    mesh: a Mesh
-    shape: a Shape
-    **kwargs: keyword args for tf.random_uniform, except seed
-
-  Returns:
-    a Tensor
-  """
-  shape = convert_to_shape(shape)
-  return RandomOperation(mesh, shape, tf.random_uniform, **kwargs).outputs[0]
-
-
-def dropout(x, keep_prob, noise_shape=None, name=None):
-  """Dropout layer.
-
-  Args:
-    x: a Tensor
-    keep_prob: a float between 0.0 and 1.0
-    noise_shape: an optional Shape (a subset of x.shape)
-    name: an optional string
-
-  Returns:
-    a Tensor
-  """
-  noise_shape = convert_to_shape(noise_shape)
-  if noise_shape is None:
-    noise_shape = x.shape
-  with tf.variable_scope(name, default_name="dropout"):
-    if keep_prob == 1.0:
-      return x
-    noise = cast(less(random_uniform(
-        x.mesh, noise_shape, dtype=x.dtype), keep_prob), x.dtype)
-    noise /= keep_prob
-    return x * noise
-
-
-def _cumprod(l):
-  """Cumulative product of a list.
-
-  Args:
-    l: a list of integers
-  Returns:
-    a list with one more element (starting with 1)
-  """
-  ret = [1]
-  for item in l:
-    ret.append(ret[-1] * item)
-  return ret
-
-
-def log_variable_sizes(var_list, tag, verbose=True):
-  """Log the sizes and shapes of variables, and the total size.
-
-  Args:
-    var_list: a list of variables; defaults to trainable_variables
-    tag: a string; defaults to "Trainable Variables"
-    verbose: bool, if True, log every weight; otherwise, log total size only.
-  """
-  if not var_list:
-    return
-
-  name_to_var = {v.name: v for v in var_list}
-  total_size = 0
-  for v_name in sorted(list(name_to_var)):
-    v = name_to_var[v_name]
-    v_size = v.shape.size
-    if verbose:
-      tf.logging.info("Weight    %s\tshape    %s\tsize    %d",
-                      v.name.ljust(80),
-                      str(v.shape).ljust(30), v_size)
-    total_size += v_size
-  tf.logging.info("%s Total size: %d", tag, total_size)
-
-
-class WhileLoopOperation(Operation):
-  """While loop."""
-
-  def __init__(self, cond_fn, body_fn, inputs,
-               tf_kwargs=None, name="while_loop"):
-    super(WhileLoopOperation, self).__init__(
-        inputs, mesh=inputs[0].mesh, name=name)
-    self._cond_fn = cond_fn
-    self._body_fn = body_fn
-    self._tf_kwargs = tf_kwargs or {}
-    assert not self._tf_kwargs.get("back_prop", False)
-    ops = self.graph.operations
-    before = len(ops)
-    def make_placeholders(name):
-      return [Tensor(self, t.shape, t.dtype, name="%s_%d" % (name, i))
-              for i, t in enumerate(inputs)]
-    self._cond_inputs = make_placeholders("cond_input")
-    self._cond_output = self._cond_fn(*self._cond_inputs)
-    self._cond_ops = ops[before:]
-    del ops[before:]
-    self._body_inputs = make_placeholders("body_input")
-    self._body_outputs = self._body_fn(*self._body_inputs)
-    for (i, (inp, body_out)) in enumerate(zip(inputs, self._body_outputs)):
-      if inp.shape != body_out.shape:
-        raise ValueError(
-            "shape mismatch i=%d inp=%s body_out=%s" % (i, inp, body_out))
-    self._body_ops = ops[before:]
-    del ops[before:]
-    self._outputs = make_placeholders("output")
-
-  def lower(self, lowering):
-    mesh_impl = lowering.mesh_impl(self)
-    def tf_cond_fn(*tf_inputs):
-      for tf_inp, mtf_inp in zip(tf_inputs, self._cond_inputs):
-        lowering.tensors[mtf_inp] = mesh_impl.LaidOutTensor(tf_inp)
-      for op in self._cond_ops:
-        with tf.name_scope(op.name):
-          op.lower(lowering)
-      lowered_output = lowering.tensors[self._cond_output]
-      ret = lowered_output.to_laid_out_tensor().tensor_list[0]
-      return ret
-
-    def tf_body_fn(*tf_inputs):
-      for tf_inp, mtf_inp in zip(tf_inputs, self._body_inputs):
-        lowering.tensors[mtf_inp] = mesh_impl.LaidOutTensor(tf_inp)
-      for op in self._body_ops:
-        with tf.name_scope(op.name):
-          op.lower(lowering)
-      return [
-          lowering.tensors[mtf_out].to_laid_out_tensor().tensor_list
-          for mtf_out in self._body_outputs]
-
-    lowered_inputs = [
-        lowering.tensors[t].to_laid_out_tensor().tensor_list
-        for t in self.inputs]
-
-    tf_outs = tf.while_loop(tf_cond_fn,
-                            tf_body_fn,
-                            lowered_inputs,
-                            back_prop=False,
-                            **self._tf_kwargs)
-    for tf_out, mtf_out in zip(tf_outs, self._outputs):
-      lowering.set_tensor_lowering(mtf_out, mesh_impl.LaidOutTensor(tf_out))
-
-
-def while_loop(cond_fn, body_fn, inputs, num_loop_vars=None, **kwargs):
-  """While Loop.
-
-  num_loop_vars is a hack for the multi-gpu setup.  In this case, loops
-  are generally slow, as all loop variables are placed on device.  By setting
-  num_loop_vars=k, then all of the loop variables except for the first k
-  are handled as mtf Variables instead of loop variables, using explicit
-  updates and control dependencies.  In this case, we only return the
-  first num_loop_vars outputs.  Do not use this option on TPU, since it
-  is unnecessary and also produces incorrect results, since xla does not
-  respect control dependencies.
-
-  Args:
-    cond_fn: a function from n Tensors to scalar boolean Tensor
-    body_fn: a function from n Tensors to n Tensors
-    inputs: a list of n Tensors
-    num_loop_vars: an optional integer.
-    **kwargs: additional kwargs passed to tf.while_loop
-
-  Returns:
-    a list of n Tensors.
-  """
-  if num_loop_vars is None:
-    return WhileLoopOperation(cond_fn, body_fn, inputs, kwargs).outputs
-  # Turn all loop vars except for the first ones into non-loop vars.
-  # see comments in docstring.
-  assert num_loop_vars > 0
-  extra_inputs = inputs[num_loop_vars:]
-  my_vars = tuple([get_variable(
-      x.mesh, "loop_var_%d" % i,
-      x.shape, initializer=tf.zeros_initializer(),
-      dtype=x.dtype,
-      collections=[tf.GraphKeys.LOCAL_VARIABLES])
-                   for i, x in enumerate(extra_inputs)])
-  first_input = depend(
-      inputs[0], [assign(var, x) for var, x in zip(my_vars, extra_inputs)])
-  inputs = [first_input] + inputs[1:num_loop_vars]
-  def my_cond_fn(*inputs):
-    return cond_fn(*(inputs + my_vars))
-  def my_body_fn(*inputs):
-    outputs = tuple(body_fn(*(inputs + my_vars)))
-    extra_outputs = outputs[num_loop_vars:]
-    first_output = depend(
-        outputs[0], [assign(var, x) for var, x in zip(my_vars, extra_outputs)])
-    outputs = (first_output,) + outputs[1:num_loop_vars]
-    return outputs
-  return WhileLoopOperation(
-      my_cond_fn, my_body_fn, inputs, kwargs).outputs
-
-
-def where(condition, if_true, if_false):
-  dtype = if_true.dtype
-  return (
-      if_true * cast(condition, dtype) +
-      if_false * cast(logical_not(condition), dtype))
-
-
-def _shape_union(shapes):
-  """A shape containing the union of all dimensions in the input shapes.
-
-  Args:
-    shapes: a list of Shapes
-
-  Returns:
-    a Shape
-  """
-  return Shape(list(set(sum([s.dims for s in shapes], []))))
-
-
-def _tf_flatten_batch_dims(x, num_nonbatch_dims):
-  """Flatten all but last num_nonbatch_dims into one dimension.
-
-  Args:
-    x: a tf.Tensor:
-    num_nonbatch_dims: an integer
-
-  Returns:
-    a tf.Tensor with 1 + num_nonbatch_dims dimensions.
-  """
-  shape = x.shape.as_list()
-  assert None not in shape
-  new_shape = ([list_product(shape[:-num_nonbatch_dims])]
-               + shape[-num_nonbatch_dims:])
-  if new_shape != shape:
-    x = tf.reshape(x, new_shape)
-  return x
-
-
-def _tf_restore_batch_dims(x, num_nonbatch_dims, prototype):
-  """Reverse op of _tf_flatten_batch_dims.
-
-  Un-flatten the first dimension of x to match all but the last
-  num_nonbatch_dims dimensions of prototype.
-
-  Args:
-    x: a tf.Tensor with 1 + num_nonbatch_dims dimensions
-    num_nonbatch_dims: an integer
-    prototype: a tf.Tensor
-
-  Returns:
-    a tf.Tensor
-  """
-  assert x.shape.ndims == 1 + num_nonbatch_dims
-  new_shape = (
-      prototype.shape.as_list()[:-num_nonbatch_dims] + x.shape.as_list()[1:])
-  assert None not in new_shape
-  if new_shape != x.shape.as_list():
-    x = tf.reshape(x, new_shape)
-  return x
-
-
-def halo_exchange(x, blocks_dim, block_size_dim, halo_size, wrap=False):
-  """Concat each block with the margins of adjacent blocks.
-
-  Get left and right blocks_dim and concatenate along block_size_dim.
-
-  Args:
-    x: a Tensor.
-    blocks_dim: a Dimension in x.shape
-    block_size_dim: a Dimension in x.shape
-    halo_size: an integer
-    wrap: a boolean
-
-  Returns:
-    a Tensor with the same shape as x, other than in block_size_dim, whose
-    size is increased by 2*halo_size.
-  """
-  if halo_size == 0:
-    return x
-
-  block_size = block_size_dim.size
-  partial_size = halo_size % block_size
-  num_complete_blocks = halo_size // block_size
-  parts = [x]
-
-  for i in xrange(1, num_complete_blocks + 1):
-    parts = ([shift(x, i, blocks_dim, wrap)] + parts +
-             [shift(x, -i, blocks_dim, wrap)])
-  if partial_size > 0:
-    left_margin = slice(x, 0, partial_size, block_size_dim.name)
-    right_margin = slice(x, block_size_dim.size - partial_size, partial_size,
-                         block_size_dim.name)
-    parts = (
-        [shift(right_margin, num_complete_blocks + 1, blocks_dim, wrap)]
-        + parts +
-        [shift(left_margin, -(num_complete_blocks + 1), blocks_dim, wrap)])
-  return concat(parts, block_size_dim.name)
-
-
-def left_halo_exchange(x, blocks_dim, block_size_dim, halo_size, wrap=False):
-  """Concat each block with the margins of adjacent blocks from the left.
-
-  Get left blocks_dim and concatenate along block_size_dim.
-
-  Args:
-    x: a Tensor.
-    blocks_dim: a Dimension in x.shape
-    block_size_dim: a Dimension in x.shape
-    halo_size: an integer
-    wrap: a boolean
-
-  Returns:
-    a Tensor with the same shape as x, other than in block_size_dim, whose
-    size is increased by halo_size.
-  """
-  if halo_size == 0:
-    return x
-
-  block_size = block_size_dim.size
-  partial_size = halo_size % block_size
-  num_complete_blocks = halo_size // block_size
-  parts = [x]
-
-  for i in xrange(1, num_complete_blocks + 1):
-    parts = ([shift(x, i, blocks_dim, wrap)] + parts)
-  if partial_size > 0:
-    right_margin = slice(x, block_size_dim.size - partial_size, partial_size,
-                         block_size_dim.name)
-    parts = ([shift(right_margin, num_complete_blocks + 1, blocks_dim, wrap)]
-             + parts)
-  return concat(parts, block_size_dim.name)
-
-
-def conv2d_with_blocks(
-    conv_input,
-    conv_filter,
-    strides,
-    padding,
-    h_blocks_dim=None,
-    w_blocks_dim=None,
-    name=None):
-  """conv2d operation with spatial partitioning.
-
-  Spatial partitioning is implemented by decomposing the image into blocks.
-  Block dimensions represented as h_blocks_dim and w_blocks_dim can be split
-  along the mesh axis. If split, then we do a halo exchange where each block
-  receives the part of the image from its left and right neighbors necessary to
-  do the convolution. Exchange can involve complete or partial blocks depending
-  on the filter height and width.
-
-  Currently, only "SAME" padding with dilation rate of 1 is supported.
-
-  Args:
-    conv_input: a Tensor of shape
-      [batch, h_blocks_dim, w_blocks_dim, h_dim, w_dim, in_channels_dim]
-    conv_filter: a Tensor of shape
-      [filter_height, filter_width, in_channels_dim, out_channels_dim]
-    strides: A list of ints. 1-D tensor of length 4.
-    padding: string, "SAME". The type of padding algorithm to use.
-      Valid is not currently supported.
-    h_blocks_dim: Dimension representing number of height blocks.
-    w_blocks_dim: Dimension representing number of height blocks.
-    name: A name for the operation (optional).
-
-  Returns:
-    A Tensor of shape
-      [batch, h_blocks_dim, w_blocks_dim, h_dim, w_dim, out_channels_dim]
-  """
-  filter_h_dim, filter_w_dim = conv_filter.shape.dims[:2]
-  assert filter_h_dim.size % 2 == 1
-  assert filter_w_dim.size % 2 == 1
-  h_dim, w_dim = conv_input.shape.dims[-3:-1]
-
-  # If h_blocks_dim and w_blocks_dim is not split, directly call conv2d.
-  if h_blocks_dim is None and w_blocks_dim is None:
-    return conv2d(conv_input, conv_filter, strides, padding, name)
-
-  # Padding 'VALID' is not supported yet.
-  if padding != "SAME":
-    raise NotImplementedError("conv2d_with_blocks requires padding=SAME")
-
-  # Halo exchange for h_blocks and w_blocks.
-  for blocks_dim, block_size_dim, halo_size in [
-      (h_blocks_dim, h_dim, filter_h_dim.size // 2),
-      (w_blocks_dim, w_dim, filter_w_dim.size // 2)]:
-    if halo_size > 0:
-      if blocks_dim is not None:
-        conv_input = halo_exchange(
-            conv_input, blocks_dim, block_size_dim, halo_size)
-      else:
-        conv_input = pad(
-            conv_input, [halo_size, halo_size], block_size_dim.name)
-  return conv2d(conv_input, conv_filter, strides, "VALID", name)
diff --git a/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py b/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
deleted file mode 100644
index 8dd0ee4f0..000000000
--- a/tensor2tensor/mesh_tensorflow/mesh_tensorflow_test.py
+++ /dev/null
@@ -1,170 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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 Mesh TensorFlow."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import parameterized
-
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-from tensor2tensor.mesh_tensorflow import placement_mesh_impl
-
-import tensorflow as tf
-
-
-class MeshTensorFlowTest(parameterized.TestCase, tf.test.TestCase):
-
-  @parameterized.parameters(
-      (mtf.Dimension("x", 5),),
-      (("x", 5),),
-  )
-  def testConvertToDimension(self, inputs):
-    dimension = mtf.convert_to_dimension(inputs)
-    self.assertEqual(dimension.name, "x")
-    self.assertEqual(dimension.size, 5)
-
-  def testConvertToDimensionGenericInputs(self):
-    dimension = mtf.convert_to_dimension(None)
-    self.assertEqual(dimension, None)
-    with self.assertRaises(TypeError):
-      mtf.convert_to_dimension(5)
-
-  @parameterized.parameters(
-      (mtf.Shape([mtf.Dimension("x", 4),
-                  mtf.Dimension("y", 8)]),),
-      ("x:4;y:8",),
-      ("x:4.y:8",),
-      ("x:4 y:8",),
-      ("x:4,y:8",),
-  )
-  def testConvertToShape(self, inputs):
-    shape = mtf.convert_to_shape(inputs)
-    self.assertEqual(shape, mtf.Shape([mtf.Dimension("x", 4),
-                                       mtf.Dimension("y", 8)]))
-
-  def testConvertToShapeGenericInputs(self):
-    shape = mtf.convert_to_shape([])
-    self.assertEqual(shape.dims, [])
-    shape = mtf.convert_to_shape(None)
-    self.assertEqual(shape, None)
-    with self.assertRaises(ValueError):
-      mtf.convert_to_shape("x;4")
-
-  @parameterized.parameters(
-      (mtf.LayoutRules([("d_ff", "model"), ("heads", "model")]),),
-      ("d_ff:model;heads:model",),
-      ("d_ff:model.heads:model",),
-      ("d_ff:model heads:model",),
-      ("d_ff:model,heads:model",),
-      ([("d_ff", "model"), ("heads", "model")],),
-  )
-  def testConvertToLayoutRules(self, inputs):
-    layout_rules = mtf.convert_to_layout_rules(inputs)
-    self.assertEqual(
-        layout_rules._pairs,
-        mtf.LayoutRules([("d_ff", "model"), ("heads", "model")])._pairs)
-
-  def testConvertToLayoutRulesGenericInputs(self):
-    with self.assertRaises(ValueError):
-      mtf.convert_to_layout_rules("d_ff;heads")
-
-  def testTensorLayout(self):
-    tensor_layout = mtf.TensorLayout([0, 2, 1])
-    self.assertEqual(tensor_layout.mesh_axis_to_tensor_axis(0), ())
-    self.assertEqual(tensor_layout.mesh_axis_to_tensor_axis(1), (0,))
-    self.assertEqual(tensor_layout.mesh_axis_to_tensor_axis(2), (0, 2))
-    tensor_layout = mtf.TensorLayout([None, 0])
-    self.assertFalse(tensor_layout.is_fully_replicated)
-    tensor_layout = mtf.TensorLayout([None, None, None])
-    self.assertTrue(tensor_layout.is_fully_replicated)
-
-  def testGraph(self):
-    graph = mtf.Graph()
-    self.assertLen(graph.operations, 0)
-    self.assertLen(graph.tensors, 0)
-    self.assertLen(graph.trainable_variables, 0)
-    self.assertLen(graph.all_variables, 0)
-    mesh = mtf.Mesh(graph, "mesh_test")
-    _ = mtf.import_tf_tensor(mesh,
-                             tf_tensor=tf.constant(0.),
-                             shape=mtf.Shape([]))
-    self.assertLen(graph.operations, 1)
-    self.assertLen(graph.tensors, 1)
-    self.assertLen(graph.trainable_variables, 0)
-    self.assertLen(graph.all_variables, 0)
-    _ = mtf.get_variable(mesh, "variable_0", mtf.Shape([]), trainable=True)
-    self.assertLen(graph.operations, 2)
-    self.assertLen(graph.tensors, 2)
-    self.assertLen(graph.trainable_variables, 1)
-    self.assertLen(graph.all_variables, 1)
-    _ = mtf.get_variable(mesh, "variable_1", mtf.Shape([]), trainable=False)
-    self.assertLen(graph.operations, 3)
-    self.assertLen(graph.tensors, 3)
-    self.assertLen(graph.trainable_variables, 1)
-    self.assertLen(graph.all_variables, 2)
-
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
-  def testLowering(self):
-    graph = mtf.Graph()
-    mesh = mtf.Mesh(graph, "my_mesh")
-    inputs = tf.constant(0.)
-    mtf_inputs = mtf.import_tf_tensor(mesh,
-                                      tf_tensor=inputs,
-                                      shape=mtf.Shape([]))
-    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[], layout={}, devices=[""])
-    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-
-    outputs = lowering.export_to_tf_tensor(mtf_inputs)
-    inputs_value, outputs_value = self.evaluate([inputs, outputs])
-    self.assertEqual(inputs_value, outputs_value)
-
-    # Check that methods run without error.
-    _ = lowering.copy_masters_to_slices()
-    _ = lowering.copy_slices_to_masters()
-
-  def testMesh(self):
-    graph = mtf.Graph()
-    mesh = mtf.Mesh(graph, "my_mesh")
-    self.assertEqual(mesh.graph, graph)
-
-  def testMeshImpl(self):
-    shape = mtf.Shape([mtf.Dimension("batch", 4),
-                       mtf.Dimension("model", 8)])
-    layout_rules = mtf.LayoutRules([("batch", "batch"),
-                                    ("d_ff", "model"),
-                                    ("heads", "model")])
-    mesh_impl = mtf.MeshImpl(shape=shape, layout_rules=layout_rules)
-    self.assertEqual(mesh_impl.shape, shape)
-    self.assertEqual(mesh_impl.ndims, len(shape))
-    self.assertEqual(mesh_impl.layout_rules, layout_rules)
-    self.assertEqual(mesh_impl.size, shape.size)
-    self.assertTrue(mesh_impl.supports_control_dependencies)
-
-    batch = mtf.Dimension("batch", 128)
-    length = mtf.Dimension("length", 500)
-    d_ff = mtf.Dimension("d_ff", 2048)
-    heads = mtf.Dimension("heads", 8)
-    self.assertEqual(mesh_impl.tensor_dimension_to_mesh_axis(batch), 0)
-    self.assertEqual(mesh_impl.tensor_dimension_to_mesh_axis(d_ff), 1)
-    self.assertEqual(mesh_impl.tensor_dimension_to_mesh_axis(heads), 1)
-    self.assertEqual(mesh_impl.tensor_layout(mtf.Shape([batch, length, d_ff])),
-                     mtf.TensorLayout([0, None, 1]))
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/mesh_tensorflow/mnist.py b/tensor2tensor/mesh_tensorflow/mnist.py
deleted file mode 100644
index 5dced5c6d..000000000
--- a/tensor2tensor/mesh_tensorflow/mnist.py
+++ /dev/null
@@ -1,250 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Mnist using mesh-tensorflow and tf.Estimator.
-
-This is an illustration of mesh-tensorflow, not a good model.
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-from tensor2tensor.mesh_tensorflow import mnist_dataset as dataset
-from tensor2tensor.mesh_tensorflow import mtf_layers
-from tensor2tensor.mesh_tensorflow import mtf_optimize
-from tensor2tensor.mesh_tensorflow import placement_mesh_impl
-import tensorflow as tf
-
-
-tf.flags.DEFINE_string("data_dir", "/tmp/mnist_data",
-                       "Path to directory containing the MNIST dataset")
-tf.flags.DEFINE_string("model_dir", "/tmp/mnist_model", "Estimator model_dir")
-tf.flags.DEFINE_integer("batch_size", 200,
-                        "Mini-batch size for the training. Note that this "
-                        "is the global batch size and not the per-shard batch.")
-tf.flags.DEFINE_integer("hidden_size", 512, "Size of each hidden layer.")
-tf.flags.DEFINE_integer("train_epochs", 40, "Total number of training epochs.")
-tf.flags.DEFINE_integer("epochs_between_evals", 1,
-                        "# of epochs between evaluations.")
-tf.flags.DEFINE_integer("eval_steps", 0,
-                        "Total number of evaluation steps. If `0`, evaluation "
-                        "after training is skipped.")
-tf.flags.DEFINE_string("mesh_shape", "b1:2;b2:2", "mesh shape")
-tf.flags.DEFINE_string("layout", "row_blocks:b1;col_blocks:b2",
-                       "layout rules")
-
-FLAGS = tf.flags.FLAGS
-
-
-def mnist_model(image, labels, mesh):
-  """The model.
-
-  Args:
-    image: tf.Tensor with shape [batch, 28*28]
-    labels: a tf.Tensor with shape [batch] and dtype tf.int32
-    mesh: a mtf.Mesh
-
-  Returns:
-    logits: a tf.Tensor with shape [batch, 10]
-    loss: a mtf.Tensor with shape []
-  """
-  batch_dim = mtf.Dimension("batch", FLAGS.batch_size)
-  row_blocks_dim = mtf.Dimension("row_blocks", 4)
-  col_blocks_dim = mtf.Dimension("col_blocks", 4)
-  rows_dim = mtf.Dimension("rows_size", 7)
-  cols_dim = mtf.Dimension("cols_size", 7)
-
-  classes_dim = mtf.Dimension("classes", 10)
-  one_channel_dim = mtf.Dimension("one_channel", 1)
-
-  x = mtf.import_tf_tensor(
-      mesh, tf.reshape(image, [FLAGS.batch_size, 4, 7, 4, 7, 1]),
-      mtf.Shape(
-          [batch_dim, row_blocks_dim, rows_dim,
-           col_blocks_dim, cols_dim, one_channel_dim]))
-  x = mtf.transpose(x, [
-      batch_dim, row_blocks_dim, col_blocks_dim,
-      rows_dim, cols_dim, one_channel_dim])
-
-  # add some convolutional layers to demonstrate that convolution works.
-  fh_dim = mtf.Dimension("fh", 9)
-  fw_dim = mtf.Dimension("fw", 9)
-  filters1_dim = mtf.Dimension("filters1", 16)
-  filters2_dim = mtf.Dimension("filters2", 16)
-  kernel1 = mtf.get_variable(
-      mesh, "kernel1", [fh_dim, fw_dim, one_channel_dim, filters1_dim])
-  kernel2 = mtf.get_variable(
-      mesh, "kernel2", [fh_dim, fw_dim, filters1_dim, filters2_dim])
-
-  f1 = mtf.relu(mtf.conv2d_with_blocks(
-      x, kernel1, strides=[1, 1, 1, 1], padding="SAME",
-      h_blocks_dim=row_blocks_dim, w_blocks_dim=col_blocks_dim))
-  f2 = mtf.relu(mtf.conv2d_with_blocks(
-      f1, kernel2, strides=[1, 1, 1, 1], padding="SAME",
-      h_blocks_dim=row_blocks_dim, w_blocks_dim=col_blocks_dim))
-  x = mtf.reduce_mean(f2, reduced_dim=filters2_dim)
-
-  # add some fully-connected dense layers.
-  hidden_dim1 = mtf.Dimension("hidden1", FLAGS.hidden_size)
-  hidden_dim2 = mtf.Dimension("hidden2", FLAGS.hidden_size)
-
-  h1 = mtf_layers.dense(
-      x, hidden_dim1,
-      reduced_dims=x.shape.dims[-4:],
-      activation=mtf.relu, name="hidden1")
-  h2 = mtf_layers.dense(
-      h1, hidden_dim2,
-      activation=mtf.relu, name="hidden2")
-  logits = mtf_layers.dense(h2, classes_dim, name="logits")
-  if labels is None:
-    loss = None
-  else:
-    labels = mtf.import_tf_tensor(
-        mesh, tf.reshape(labels, [FLAGS.batch_size]), mtf.Shape([batch_dim]))
-    loss = mtf_layers.softmax_cross_entropy_with_logits(
-        logits, mtf.one_hot(labels, classes_dim), classes_dim)
-    loss = mtf.reduce_mean(loss)
-  return logits, loss
-
-
-def model_fn(features, labels, mode, params):
-  """The model_fn argument for creating an Estimator."""
-  tf.logging.info("features = %s labels = %s mode = %s params=%s" %
-                  (features, labels, mode, params))
-  global_step = tf.train.get_global_step()
-  graph = mtf.Graph()
-  mesh = mtf.Mesh(graph, "my_mesh")
-  logits, loss = mnist_model(features, labels, mesh)
-  mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
-  layout_rules = mtf.convert_to_layout_rules(FLAGS.layout)
-  mesh_size = mesh_shape.size
-  mesh_devices = [""] * mesh_size
-  mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-      mesh_shape, layout_rules, mesh_devices)
-
-  if mode == tf.estimator.ModeKeys.TRAIN:
-    var_grads = mtf.gradients(
-        [loss], [v.outputs[0] for v in graph.trainable_variables])
-    optimizer = mtf_optimize.AdafactorOptimizer()
-    update_ops = []
-    for grad, var in zip(var_grads, graph.trainable_variables):
-      update_ops.extend(optimizer.apply_grad(grad, var))
-
-  lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-  restore_hook = mtf.MtfRestoreHook(lowering)
-
-  tf_logits = lowering.export_to_tf_tensor(logits)
-  if mode != tf.estimator.ModeKeys.PREDICT:
-    tf_loss = lowering.export_to_tf_tensor(loss)
-    tf.summary.scalar("loss", tf_loss)
-
-  if mode == tf.estimator.ModeKeys.TRAIN:
-    tf_update_ops = [lowering.lowered_operation(op) for op in update_ops]
-    tf_update_ops.append(tf.assign_add(global_step, 1))
-    train_op = tf.group(tf_update_ops)
-    saver = tf.train.Saver(
-        tf.global_variables(),
-        sharded=True,
-        max_to_keep=10,
-        keep_checkpoint_every_n_hours=2,
-        defer_build=False, save_relative_paths=True)
-    tf.add_to_collection(tf.GraphKeys.SAVERS, saver)
-    saver_listener = mtf.MtfCheckpointSaverListener(lowering)
-    saver_hook = tf.train.CheckpointSaverHook(
-        FLAGS.model_dir,
-        save_steps=1000,
-        saver=saver,
-        listeners=[saver_listener])
-
-    accuracy = tf.metrics.accuracy(
-        labels=labels, predictions=tf.argmax(tf_logits, axis=1))
-
-    # Name tensors to be logged with LoggingTensorHook.
-    tf.identity(tf_loss, "cross_entropy")
-    tf.identity(accuracy[1], name="train_accuracy")
-
-    # Save accuracy scalar to Tensorboard output.
-    tf.summary.scalar("train_accuracy", accuracy[1])
-
-    # restore_hook must come before saver_hook
-    return tf.estimator.EstimatorSpec(
-        tf.estimator.ModeKeys.TRAIN, loss=tf_loss, train_op=train_op,
-        training_chief_hooks=[restore_hook, saver_hook])
-
-  if mode == tf.estimator.ModeKeys.PREDICT:
-    predictions = {
-        "classes": tf.argmax(tf_logits, axis=1),
-        "probabilities": tf.nn.softmax(tf_logits),
-    }
-    return tf.estimator.EstimatorSpec(
-        mode=tf.estimator.ModeKeys.PREDICT,
-        predictions=predictions,
-        prediction_hooks=[restore_hook],
-        export_outputs={
-            "classify": tf.estimator.export.PredictOutput(predictions)
-        })
-  if mode == tf.estimator.ModeKeys.EVAL:
-    return tf.estimator.EstimatorSpec(
-        mode=tf.estimator.ModeKeys.EVAL,
-        loss=tf_loss,
-        evaluation_hooks=[restore_hook],
-        eval_metric_ops={
-            "accuracy":
-            tf.metrics.accuracy(
-                labels=labels, predictions=tf.argmax(tf_logits, axis=1)),
-        })
-
-
-def run_mnist():
-  """Run MNIST training and eval loop."""
-  mnist_classifier = tf.estimator.Estimator(
-      model_fn=model_fn,
-      model_dir=FLAGS.model_dir)
-
-  # Set up training and evaluation input functions.
-  def train_input_fn():
-    """Prepare data for training."""
-
-    # When choosing shuffle buffer sizes, larger sizes result in better
-    # randomness, while smaller sizes use less memory. MNIST is a small
-    # enough dataset that we can easily shuffle the full epoch.
-    ds = dataset.train(FLAGS.data_dir)
-    ds_batched = ds.cache().shuffle(buffer_size=50000).batch(FLAGS.batch_size)
-
-    # Iterate through the dataset a set number (`epochs_between_evals`) of times
-    # during each training session.
-    ds = ds_batched.repeat(FLAGS.epochs_between_evals)
-    return ds
-
-  def eval_input_fn():
-    return dataset.test(FLAGS.data_dir).batch(
-        FLAGS.batch_size).make_one_shot_iterator().get_next()
-
-  # Train and evaluate model.
-  for _ in range(FLAGS.train_epochs // FLAGS.epochs_between_evals):
-    mnist_classifier.train(input_fn=train_input_fn, hooks=None)
-    eval_results = mnist_classifier.evaluate(input_fn=eval_input_fn)
-    print("\nEvaluation results:\n\t%s\n" % eval_results)
-
-
-def main(_):
-  run_mnist()
-
-
-if __name__ == "__main__":
-  tf.logging.set_verbosity(tf.logging.INFO)
-  tf.app.run()
diff --git a/tensor2tensor/mesh_tensorflow/mnist_dataset.py b/tensor2tensor/mesh_tensorflow/mnist_dataset.py
deleted file mode 100644
index 3be1e382d..000000000
--- a/tensor2tensor/mesh_tensorflow/mnist_dataset.py
+++ /dev/null
@@ -1,131 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-#  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.
-"""tf.data.Dataset interface to the MNIST dataset."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gzip
-import os
-import shutil
-import tempfile
-
-import numpy as np
-from six.moves import urllib
-import tensorflow as tf
-
-
-def read32(bytestream):
-  """Read 4 bytes from bytestream as an unsigned 32-bit integer."""
-  dt = np.dtype(np.uint32).newbyteorder('>')
-  return np.frombuffer(bytestream.read(4), dtype=dt)[0]
-
-
-def check_image_file_header(filename):
-  """Validate that filename corresponds to images for the MNIST dataset."""
-  with tf.gfile.Open(filename, 'rb') as f:
-    magic = read32(f)
-    read32(f)  # num_images, unused
-    rows = read32(f)
-    cols = read32(f)
-    if magic != 2051:
-      raise ValueError('Invalid magic number %d in MNIST file %s' % (magic,
-                                                                     f.name))
-    if rows != 28 or cols != 28:
-      raise ValueError(
-          'Invalid MNIST file %s: Expected 28x28 images, found %dx%d' %
-          (f.name, rows, cols))
-
-
-def check_labels_file_header(filename):
-  """Validate that filename corresponds to labels for the MNIST dataset."""
-  with tf.gfile.Open(filename, 'rb') as f:
-    magic = read32(f)
-    read32(f)  # num_items, unused
-    if magic != 2049:
-      raise ValueError('Invalid magic number %d in MNIST file %s' % (magic,
-                                                                     f.name))
-
-
-def download(directory, filename):
-  """Download (and unzip) a file from the MNIST dataset if not already done."""
-  filepath = os.path.join(directory, filename)
-  if tf.gfile.Exists(filepath):
-    return filepath
-  if not tf.gfile.Exists(directory):
-    tf.gfile.MakeDirs(directory)
-  url = 'http://yann.lecun.com/exdb/mnist/' + filename + '.gz'
-  _, zipped_filepath = tempfile.mkstemp(suffix='.gz')
-  print('Downloading %s to %s' % (url, zipped_filepath))
-  urllib.request.urlretrieve(url, zipped_filepath)
-  with gzip.open(zipped_filepath, 'rb') as f_in, \
-      tf.gfile.Open(filepath, 'wb') as f_out:
-    shutil.copyfileobj(f_in, f_out)
-  os.remove(zipped_filepath)
-  return filepath
-
-
-def dataset(directory, images_file, labels_file):
-  """Download and parse MNIST dataset."""
-
-  images_file = download(directory, images_file)
-  labels_file = download(directory, labels_file)
-
-  check_image_file_header(images_file)
-  check_labels_file_header(labels_file)
-
-  def decode_image(image):
-    # Normalize from [0, 255] to [0.0, 1.0]
-    image = tf.decode_raw(image, tf.uint8)
-    image = tf.cast(image, tf.float32)
-    image = tf.reshape(image, [784])
-    return image / 255.0
-
-  def decode_label(label):
-    label = tf.decode_raw(label, tf.uint8)  # tf.string -> [tf.uint8]
-    label = tf.reshape(label, [])  # label is a scalar
-    return tf.to_int32(label)
-
-  images = tf.data.FixedLengthRecordDataset(
-      images_file, 28 * 28, header_bytes=16).map(decode_image)
-  labels = tf.data.FixedLengthRecordDataset(
-      labels_file, 1, header_bytes=8).map(decode_label)
-  return tf.data.Dataset.zip((images, labels))
-
-
-def train(directory):
-  """tf.data.Dataset object for MNIST training data."""
-  return dataset(directory, 'train-images-idx3-ubyte',
-                 'train-labels-idx1-ubyte')
-
-
-def test(directory):
-  """tf.data.Dataset object for MNIST test data."""
-  return dataset(directory, 't10k-images-idx3-ubyte', 't10k-labels-idx1-ubyte')
diff --git a/tensor2tensor/mesh_tensorflow/mtf_beam_search.py b/tensor2tensor/mesh_tensorflow/mtf_beam_search.py
deleted file mode 100644
index ce1298676..000000000
--- a/tensor2tensor/mesh_tensorflow/mtf_beam_search.py
+++ /dev/null
@@ -1,521 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Implementation of beam search with penalties."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-import tensorflow as tf
-
-# Assuming EOS_ID is 1
-EOS_ID = 1
-# Default value for INF
-INF = 1. * 1e7
-
-
-def compute_topk_scores_and_seq(sequences, scores, scores_to_gather, flags,
-                                beam_dim, prefix="default",
-                                states=None):
-  """Given sequences and scores, will gather the top k=beam size sequences.
-
-  This function is used to grow alive, and finished. It takes sequences,
-  scores, and flags, and returns the top k from sequences, scores_to_gather,
-  and flags based on the values in scores.
-
-  This method permits easy introspection using tfdbg.  It adds three named ops
-  that are prefixed by `prefix`:
-    - _topk_seq: the tensor for topk_seq returned by this method.
-    - _topk_flags: the tensor for topk_finished_flags returned by this method.
-    - _topk_scores: the tensor for tokp_gathered_scores returned by this method.
-
-  Args:
-    sequences: Tensor of sequences that we need to gather from.
-      [batch_size, beam_size, seq_length]
-    scores: Tensor of scores for each sequence in sequences.
-      [batch_size, beam_size]. We will use these to compute the topk.
-    scores_to_gather: Tensor of scores for each sequence in sequences.
-      [batch_size, beam_size]. We will return the gathered scores from here.
-      Scores to gather is different from scores because for grow_alive, we will
-      need to return log_probs, while for grow_finished, we will need to return
-      the length penalized scores.
-    flags: Tensor of bools for sequences that say whether a sequence has reached
-      EOS or not
-    beam_dim: mtf.Dimension
-    prefix: an optional string
-    states: an optional list of mtf.Tensor
-  Returns:
-    Tuple of
-    (topk_seq [batch_size, beam_size, decode_length],
-     topk_gathered_scores [batch_size, beam_size],
-     topk_finished_flags[batch_size, beam_size],
-     topk_gathered_states)
-  """
-  unused_batch_dim, old_beam_dim, unused_length_dim = sequences.shape.dims
-  topk_indices, _ = mtf.top_k(scores, old_beam_dim, beam_dim)
-
-  # Gather up the highest scoring sequences.
-  # For each operation added, give it
-  # a concrete name to simplify observing these operations with tfdbg.
-  # Clients can capture these tensors by watching these node names.
-  def gather(tensor, name):
-    with tf.name_scope(prefix + name):
-      output_shape = mtf.Shape(
-          [beam_dim if d == old_beam_dim else d for d in tensor.shape.dims])
-      return mtf.gather(
-          tensor, topk_indices, old_beam_dim, output_shape=output_shape)
-  topk_seq = gather(sequences, "_seq")
-  topk_flags = gather(flags, "_flags")
-  topk_gathered_scores = gather(scores_to_gather, "_scores")
-  if states is None:
-    topk_gathered_states = None
-  else:
-    topk_gathered_states = [gather(state, "_topk_states") for state in states]
-  return topk_seq, topk_gathered_scores, topk_flags, topk_gathered_states
-
-
-def beam_search(logits_fn,
-                initial_ids,
-                alpha,
-                states=None,
-                eos_id=EOS_ID,
-                stop_early=True,
-                decode_length=None,
-                use_tpu=True):
-  """Beam search with length penalties.
-
-  Requires a function that can take the currently decoded symbols and return
-  the logits for the next symbol. The implementation is inspired by
-  https://arxiv.org/abs/1609.08144.
-
-  When running, the beam search steps can be visualized by using tfdbg to watch
-  the operations generating the output ids for each beam step.  These operations
-  have the pattern:
-    (alive|finished)_topk_(seq,scores)
-
-  Operations marked `alive` represent the new beam sequences that will be
-  processed in the next step.  Operations marked `finished` represent the
-  completed beam sequences, which may be padded with 0s if no beams finished.
-
-  Operations marked `seq` store the full beam sequence for the time step.
-  Operations marked `scores` store the sequence's final log scores.
-
-  The beam search steps will be processed sequentially in order, so when
-  capturing observed from these operations, tensors, clients can make
-  assumptions about which step is being recorded.
-
-  WARNING: Assumes 2nd dimension of tensors in `states` and not invariant, this
-  means that the shape of the 2nd dimension of these tensors will not be
-  available (i.e. set to None) inside logits_fn.
-
-  Args:
-    logits_fn: Interface to the model, to provide logits.
-        Shoud take:
-          step_num - mtf Scalar
-          ids - mtf Tensor with shape [batch, beam, length]
-        Should return:
-          logits - [batch, beam, vocab_size]
-    initial_ids: a mtf.Tensor with shape [batch_dim, beam_dim, length_dim])
-    alpha: alpha for length penalty.
-    states: list of mtf.Tensor
-    eos_id: ID for end of sentence.
-    stop_early: a boolean - stop once best sequence is provably determined.
-    decode_length: a mtf Scalar of dtype tf.int32 - maximum length of decodes
-    use_tpu: a boolean
-  Returns:
-    Tuple of
-    (decoded beams [batch, beam, length]
-     decoding probabilities [batch, beam_size])
-  """
-  batch_dim, beam_dim, length_dim = initial_ids.shape.dims
-  mesh = initial_ids.mesh
-
-  batch_by_beam = mtf.Shape([batch_dim, beam_dim])
-  initial_log_probs = mtf.broadcast(
-      mtf.one_hot(
-          mtf.constant(mesh, 0, dtype=tf.int32),
-          beam_dim,
-          on_value=0.0,
-          off_value=-INF),
-      batch_by_beam)
-
-  length_scalar = mtf.constant(mesh, length_dim.size, dtype=tf.int32)
-  if decode_length is None:
-    decode_length = length_scalar
-  else:
-    decode_length = mtf.minimum(decode_length, length_scalar)
-
-  alive_log_probs = initial_log_probs
-  alive_seq = initial_ids
-
-  # Finished will keep track of all the sequences that have finished so far
-  # Finished log probs will be negative infinity in the beginning
-  # finished_flags will keep track of booleans
-  finished_seq = initial_ids
-  finished_scores = mtf.constant(mesh, -INF, batch_by_beam)
-
-  # Setting the scores of the initial to negative infinity.
-  finished_flags = mtf.constant(mesh, False, batch_by_beam, tf.bool)
-
-  def grow_finished(finished_seq, finished_scores, finished_flags, curr_seq,
-                    curr_scores, curr_finished):
-    """Given sequences and scores, will gather the top k=beam size sequences.
-
-    Args:
-      finished_seq: Current finished sequences.
-        [batch, beam, length]
-      finished_scores: scores for each of these sequences.
-        [batch, beam]
-      finished_flags: finished bools for each of these sequences.
-        [batch, beam]
-      curr_seq: current topk sequence that has been grown by one position.
-        [batch, beam, length]
-      curr_scores: scores for each of these sequences. [batch, beam]
-      curr_finished: Finished flags for each of these sequences.
-        [batch, beam]
-    Returns:
-      Tuple of
-        (Topk sequences based on scores,
-         log probs of these sequences,
-         Finished flags of these sequences,
-         None (no states))
-    """
-
-    # Set the scores of the unfinished seq in curr_seq to large negative
-    # values
-    curr_scores += (1. - mtf.to_float(curr_finished)) * -INF
-    unused_batch_dim, beam_dim, unused_length_dim = finished_seq.shape.dims
-    # concatenating the sequences and scores along beam axis
-    def _my_concat(a, b):
-      a = mtf.rename_dimension(a, "beam", "triple_beam")
-      b = mtf.rename_dimension(b, "double_beam", "triple_beam")
-      return mtf.concat([a, b], "triple_beam")
-
-    curr_finished_seq = _my_concat(finished_seq, curr_seq)
-    curr_finished_scores = _my_concat(finished_scores, curr_scores)
-    curr_finished_flags = _my_concat(finished_flags, curr_finished)
-    return compute_topk_scores_and_seq(
-        curr_finished_seq, curr_finished_scores, curr_finished_scores,
-        curr_finished_flags, beam_dim, "grow_finished", states=None)
-
-  def grow_alive(curr_seq, curr_scores, curr_log_probs, curr_finished, states):
-    """Given sequences and scores, will gather the top k=beam size sequences.
-
-    Args:
-      curr_seq: current topk sequence that has been grown by one position.
-        [batch, beam, length]
-      curr_scores: scores for each of these sequences. [batch_size, beam_size]
-      curr_log_probs: log probs for each of these sequences.
-        [batch, beam]
-      curr_finished: Finished flags for each of these sequences.
-        [batch, beam]
-      states: list of mtf.Tensor
-    Returns:
-      Tuple of
-        (Topk sequences based on scores,
-         log probs of these sequences,
-         Finished flags of these sequences)
-    """
-    # Set the scores of the finished seq in curr_seq to large negative
-    # values
-    curr_scores += mtf.to_float(curr_finished) * -INF
-    return compute_topk_scores_and_seq(curr_seq, curr_scores, curr_log_probs,
-                                       curr_finished, beam_dim,
-                                       "grow_alive", states)
-
-  def grow_topk(i, alive_seq, alive_log_probs, states=None):
-    r"""Inner beam search loop.
-
-    This function takes the current alive sequences, and grows them to topk
-    sequences where k = 2*beam. We use 2*beam because, we could have beam_size
-    number of sequences that might hit <EOS> and there will be no alive
-    sequences to continue. With 2*beam_size, this will not happen. This relies
-    on the assumption the vocab size is > beam size. If this is true, we'll
-    have at least beam_size non <EOS> extensions if we extract the next top
-    2*beam words.
-    Length penalty is given by = (5+len(decode)/6) ^ -\alpha. Pls refer to
-    https://arxiv.org/abs/1609.08144.
-
-    Args:
-      i: loop index
-      alive_seq: Topk sequences decoded so far [batch, beam, length]
-      alive_log_probs: probabilities of these sequences. [batch, beam]
-      states: optional list of mtf.Tensor
-    Returns:
-      Tuple of
-        (Topk sequences extended by the next word,
-         The log probs of these sequences,
-         The scores with length penalty of these sequences,
-         Flags indicating which of these sequences have finished decoding,
-         list of transformed decoding states)
-    """
-    logits, new_states = logits_fn(i, alive_seq, states)
-    batch_dim, beam_dim, vocab_dim = logits.shape.dims
-
-    # Convert logits to normalized log probs
-    candidate_log_probs = mtf.log_softmax(logits, vocab_dim)
-
-    # Multiply the probabilities by the current probabilities of the beam.
-    # (batch_size, beam_size, vocab_size) + (batch_size, beam_size, 1)
-    log_probs = candidate_log_probs + alive_log_probs
-
-    length_penalty = mtf.pow(((5. + mtf.to_float(i + 1)) / 6.), alpha)
-
-    curr_scores = log_probs / length_penalty
-
-    # scores have shape [batch, beam, vocab]
-    beam_and_vocab_dim = mtf.Dimension(
-        "beam_and_vocab", beam_dim.size * vocab_dim.size)
-    flat_shape = mtf.Shape([batch_dim, beam_and_vocab_dim])
-    double_beam = mtf.Dimension("double_beam", beam_dim.size * 2)
-    # Flatten out (beam_size, vocab_size) probs in to a list of possibilities
-    flat_curr_scores = mtf.reshape(curr_scores, flat_shape)
-
-    top_ids, top_scores = mtf.top_k(
-        flat_curr_scores, reduced_dim=beam_and_vocab_dim, new_dim=double_beam)
-
-    # Recovering the log probs because we will need to send them back
-    top_log_probs = top_scores * length_penalty
-
-    # Work out what beam the top probs are in.
-    top_beam_index = top_ids // vocab_dim.size
-    top_ids %= vocab_dim.size  # Unflatten the ids
-
-    def my_gather(tensor):
-      return mtf.gather(
-          tensor, top_beam_index, beam_dim,
-          output_shape=mtf.Shape(
-              [double_beam if d == beam_dim else d for d in tensor.shape.dims]))
-
-    # Gather up the most probable 2*beams both for the ids and finished_in_alive
-    # bools
-    top_seq = my_gather(alive_seq)
-
-    if states:
-      states = [my_gather(state) for state in new_states]
-
-    # Append the most probable alive
-    top_seq += top_ids * mtf.one_hot(i, length_dim, dtype=tf.int32)
-    top_finished = mtf.equal(top_ids, eos_id)
-
-    return top_seq, top_log_probs, top_scores, top_finished, states
-
-  def inner_loop(i, alive_seq, alive_log_probs, finished_seq, finished_scores,
-                 finished_flags, *states):
-    """Inner beam search loop.
-
-    There are three groups of tensors, alive, finished, and topk.
-    The alive group contains information about the current alive sequences
-    The topk group contains information about alive + topk current decoded words
-    the finished group contains information about finished sentences, that is,
-    the ones that have decoded to <EOS>. These are what we return.
-    The general beam search algorithm is as follows:
-    While we haven't terminated (pls look at termination condition)
-      1. Grow the current alive to get beam*2 topk sequences
-      2. Among the topk, keep the top beam_size ones that haven't reached EOS
-      into alive
-      3. Among the topk, keep the top beam_size ones have reached EOS into
-      finished
-    Repeat
-    To make things simple with using fixed size tensors, we will end
-    up inserting unfinished sequences into finished in the beginning. To stop
-    that we add -ve INF to the score of the unfinished sequence so that when a
-    true finished sequence does appear, it will have a higher score than all the
-    unfinished ones.
-
-    Args:
-      i: loop index
-      alive_seq: Topk sequences decoded so far [batch_size, beam_size, i+1]
-      alive_log_probs: probabilities of the beams. [batch_size, beam_size]
-      finished_seq: Current finished sequences.
-        [batch_size, beam_size, i+1]
-      finished_scores: scores for each of these sequences.
-        [batch_size, beam_size]
-      finished_flags: finished bools for each of these sequences.
-        [batch_size, beam_size]
-      *states: mtf Tensors
-
-    Returns:
-      Tuple of
-        (Incremented loop index
-         New alive sequences,
-         Log probs of the alive sequences,
-         New finished sequences,
-         Scores of the new finished sequences,
-         Flags indicating which sequence in finished as reached EOS,
-         dict of final decoding states)
-    """
-
-    # Each inner loop, we carry out three steps:
-    # 1. Get the current topk items.
-    # 2. Extract the ones that have finished and haven't finished
-    # 3. Recompute the contents of finished based on scores.
-    (top2k_seq, top2k_log_probs, top2k_scores, top2k_finished,
-     top2k_states) = grow_topk(i, alive_seq, alive_log_probs, states)
-    alive_seq, alive_log_probs, _, states = grow_alive(
-        top2k_seq, top2k_scores, top2k_log_probs, top2k_finished, top2k_states)
-    finished_seq, finished_scores, finished_flags, _ = grow_finished(
-        finished_seq, finished_scores, finished_flags, top2k_seq, top2k_scores,
-        top2k_finished)
-    return (i + 1, alive_seq, alive_log_probs, finished_seq, finished_scores,
-            finished_flags) + tuple(states)
-
-  def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
-                   finished_scores, finished_in_finished, *unused_states):
-    """Checking termination condition.
-
-    We terminate when we decoded up to decode_length or the lowest scoring item
-    in finished has a greater score that the highest prob item in alive divided
-    by the max length penalty
-
-    Args:
-      i: loop index
-      alive_log_probs: probabilities of the beams. [batch_size, beam_size]
-      finished_scores: scores for each of these sequences.
-        [batch_size, beam_size]
-      finished_in_finished: finished bools for each of these sequences.
-        [batch_size, beam_size]
-
-    Returns:
-      Bool.
-    """
-    # TODO(noam): support a different decode length...
-    # decode_length = mtf.constant(mesh, length_dim.size, dtype=tf.int32)
-
-    # del alive_log_probs, finished_scores, finished_in_finished
-    # return mtf.less(i, length_dim.size)
-    if not stop_early:
-      return mtf.less(i, decode_length)
-    max_length_penalty = mtf.pow(
-        ((5. + mtf.to_float(decode_length)) / 6.), alpha)
-    # The best possible score of the most likely alive sequence.
-    lower_bound_alive_scores = mtf.gather(
-        alive_log_probs, mtf.constant(mesh, 0, dtype=tf.int32),
-        beam_dim) / max_length_penalty
-
-    # Now to compute the lowest score of a finished sequence in finished
-    # If the sequence isn't finished, we multiply it's score by 0. since
-    # scores are all -ve, taking the min will give us the score of the lowest
-    # finished item.
-    lowest_score_of_finished_in_finished = mtf.reduce_min(
-        finished_scores * mtf.to_float(finished_in_finished),
-        reduced_dim=beam_dim)
-
-    # If none of the sequences have finished, then the min will be 0 and
-    # we have to replace it by -ve INF if it is. The score of any seq in alive
-    # will be much higher than -ve INF and the termination condition will not
-    # be met.
-    lowest_score_of_finished_in_finished += (
-        (1. - mtf.to_float(mtf.reduce_any(
-            finished_in_finished, reduced_dim=beam_dim))) * -INF)
-
-    bound_is_met = mtf.reduce_all(
-        mtf.greater(lowest_score_of_finished_in_finished,
-                    lower_bound_alive_scores))
-    return mtf.logical_and(
-        mtf.less(i, decode_length), mtf.logical_not(bound_is_met))
-
-  initial_step_num = mtf.constant(mesh, 0, dtype=tf.int32)
-  while_loop_inputs = [
-      initial_step_num, alive_seq, alive_log_probs, finished_seq,
-      finished_scores, finished_flags] + states
-
-  (_, alive_seq, alive_log_probs, finished_seq, finished_scores,
-   finished_flags) = mtf.while_loop(
-       _is_finished, inner_loop, while_loop_inputs,
-       num_loop_vars=None if use_tpu else 6)[:6]
-
-  # Accounting for corner case: It's possible that no sequence in alive for a
-  # particular batch item ever reached EOS. In that case, we should just copy
-  # the contents of alive for that batch item. tf.reduce_any(finished_flags, 1)
-  # if 0, means that no sequence for that batch index had reached EOS. We need
-  # to do the same for the scores as well.
-  finished_seq = mtf.where(
-      mtf.reduce_any(finished_flags, reduced_dim=beam_dim),
-      finished_seq, alive_seq)
-  finished_scores = mtf.where(
-      mtf.reduce_any(finished_flags, reduced_dim=beam_dim),
-      finished_scores, alive_log_probs)
-  return finished_seq, finished_scores
-
-
-def greedy_decode(logits_fn,
-                  initial_ids,
-                  temperature=0.0,
-                  initial_states=None,
-                  eos_id=EOS_ID,
-                  forced_ids=None,
-                  use_tpu=True):
-  """Greedy decoding.
-
-  Args:
-    logits_fn: Interface to the model, to provide logits.
-        Shoud take:
-          step_num - mtf Scalar
-          ids - mtf Tensor with shape [..., length]
-          states - list of mtf.Tensor
-        Should return:
-          logits - [batch, vocab_size]
-          new_states - list of mtf.Tensor
-    initial_ids: mtf.Tensor with shape [..., length], containing zeros.
-    temperature: a float between 0.0 (argmax) and 1.0 (random)
-    initial_states: list of mtf.Tensor
-    eos_id: ID for end of sentence.
-    forced_ids: optional mtf.Tensor with shape [..., length]
-    use_tpu: a boolean
-  Returns:
-    Tensor with shape [..., length]
-  """
-  length_dim = initial_ids.shape.dims[-1]
-  mesh = initial_ids.mesh
-  num_steps = mtf.constant(mesh, length_dim.size, dtype=tf.int32)
-  def cond_fn(step_num, prev_ids, *unused_states):
-    """Should we run another loop iteration."""
-    overflow = mtf.equal(step_num, num_steps)
-    has_eos = mtf.reduce_any(
-        mtf.equal(prev_ids, eos_id), reduced_dim=length_dim)
-    all_has_eos = mtf.reduce_all(has_eos)
-    return mtf.logical_not(mtf.logical_or(overflow, all_has_eos))
-  def body_fn(step_num, ids, *states):
-    """Body function for greedy decoding.
-
-    Args:
-      step_num: a mtf.Tensor
-      ids: a mtf.Tensor
-      *states: additional mtf.Tensors
-    Returns:
-      new_step_num, new_ids, *new_states
-    """
-    logits, new_states = logits_fn(step_num, ids, states)
-    vocab_dim = logits.shape.dims[-1]
-    new_ids = mtf.sample_with_temperature(
-        logits, vocab_dim, temperature)
-    if forced_ids is not None:
-      # force the new ids to equal the partial targets where specified
-      # (positions where partial_targets contain nonzero values)
-      forced = mtf.gather(forced_ids, step_num, length_dim)
-      new_ids = forced + new_ids * mtf.to_int32(mtf.equal(forced, 0))
-    ids += new_ids * mtf.one_hot(step_num, length_dim, dtype=tf.int32)
-    new_step_num = step_num + 1
-    return [new_step_num, ids] + new_states
-  initial_step_num = mtf.constant(mesh, 0, dtype=tf.int32)
-  while_loop_inputs = [initial_step_num, initial_ids] + initial_states
-  final_step_num, mtf_samples = mtf.while_loop(
-      cond_fn, body_fn, while_loop_inputs,
-      num_loop_vars=None if use_tpu else 2)[:2]
-  mtf_samples = mtf.Print(mtf_samples, [final_step_num], "output_length")
-  return mtf_samples
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers.py b/tensor2tensor/mesh_tensorflow/mtf_layers.py
deleted file mode 100644
index 310eefb66..000000000
--- a/tensor2tensor/mesh_tensorflow/mtf_layers.py
+++ /dev/null
@@ -1,754 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Layers for mesh tensorflow."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-import tensorflow as tf
-
-
-def dense(x, output_dim, reduced_dims=None, expert_dims=None,
-          use_bias=True, activation=None, name=None):
-  """Dense layer doing (kernel*x + bias) computation.
-
-  Args:
-    x: a mtf.Tensor of shape [..., reduced_dims].
-    output_dim: a mtf.Dimension
-    reduced_dims: an optional list of mtf.Dimensions of x to be reduced. If
-      omitted, we reduce the last dimension.
-    expert_dims: an optional list of mtf.Dimension which represent different
-      experts. Different experts get different weights.
-    use_bias: a boolean, whether to add bias.
-    activation: an optional function from mtf.Tensor to mtf.Tensor
-    name: a string. variable scope.
-
-  Returns:
-    a mtf.Tensor of shape [..., output_dim].
-  """
-  if expert_dims is None:
-    expert_dims = []
-  if reduced_dims is None:
-    reduced_dims = x.shape.dims[-1:]
-  w_shape = mtf.Shape(expert_dims + reduced_dims + [output_dim])
-  output_shape = mtf.Shape(
-      [d for d in x.shape.dims if d not in reduced_dims] + [output_dim])
-
-  with tf.variable_scope(name, default_name="dense"):
-    stddev = mtf.list_product(d.size for d in reduced_dims) ** -0.5
-    w = mtf.get_variable(
-        x.mesh,
-        "kernel",
-        w_shape,
-        initializer=tf.random_normal_initializer(stddev=stddev),
-        activation_dtype=x.dtype)
-    y = mtf.einsum([x, w], output_shape)
-    if use_bias:
-      b = mtf.get_variable(
-          x.mesh,
-          "bias",
-          mtf.Shape(expert_dims + [output_dim]),
-          initializer=tf.zeros_initializer(),
-          activation_dtype=x.dtype)
-      y += b
-    if activation is not None:
-      y = activation(y)
-    return y
-
-
-def layer_norm(x, dim, epsilon=1e-6, name="layer_prepostprocess"):
-  """Layer normalization over dimension dim.
-
-  Args:
-    x: a mtf.Tensor whose shape contains dim.
-    dim: a mtf.Dimension
-    epsilon: a floating point number
-    name: a string. variable scope.
-
-  Returns:
-    a mtf.Tensor with same shape as x.
-  """
-  with tf.variable_scope(name + "/layer_norm"):
-    scale = mtf.get_variable(
-        x.mesh,
-        "layer_norm_scale",
-        mtf.Shape([dim]),
-        initializer=tf.ones_initializer(),
-        activation_dtype=x.dtype)
-    bias = mtf.get_variable(
-        x.mesh,
-        "layer_norm_bias",
-        mtf.Shape([dim]),
-        initializer=tf.zeros_initializer(),
-        activation_dtype=x.dtype)
-    reduced_shape = x.shape - dim
-    mean = mtf.reduce_mean(x, output_shape=reduced_shape)
-    variance = mtf.reduce_mean(mtf.square(x - mean), output_shape=reduced_shape)
-    norm_x = (x - mean) * mtf.rsqrt(variance + epsilon)
-    return norm_x * scale + bias
-
-
-def batch_norm(x, is_training, momentum, epsilon=1e-9, name=None):
-  """Batch normalization.
-
-  Args:
-    x: a mtf.Tensor whose shape contains [batch_dim, ..., dim]
-    is_training: a boolean, whether mode is training.
-    momentum: a floating point number, specifying batch norm decay value.
-    epsilon: a floating point number.
-    name: a string. variable scope.
-
-  Returns:
-    a mtf.Tensor with same shape as x.
-  """
-  with tf.variable_scope(name, default_name="batch_norm", values=[x]):
-    batch_dim = x.shape.dims[0]
-    reduced_shape = x.shape - batch_dim
-    scale = mtf.get_variable(
-        x.mesh,
-        "batch_norm_scale",
-        mtf.Shape([batch_dim]),
-        initializer=tf.ones_initializer(),
-        activation_dtype=x.dtype)
-    bias = mtf.get_variable(
-        x.mesh,
-        "batch_norm_bias",
-        mtf.Shape([batch_dim]),
-        initializer=tf.zeros_initializer(),
-        activation_dtype=x.dtype)
-
-    moving_mean = mtf.get_variable(
-        x.mesh, "moving_mean", reduced_shape,
-        initializer=tf.random_normal_initializer(stddev=1.0),
-        activation_dtype=x.dtype,
-        trainable=False)
-    moving_variance = mtf.get_variable(
-        x.mesh, "moving_variance",
-        reduced_shape, initializer=tf.ones_initializer(),
-        activation_dtype=x.dtype,
-        trainable=False)
-
-    # At training time, calculate mean and variance and normalize across batch
-    # dim.
-    if is_training:
-      mean = mtf.reduce_mean(x, output_shape=reduced_shape)
-      variance = mtf.reduce_mean(
-          mtf.square(x - mean), output_shape=reduced_shape)
-      norm_x = (x - mean) * mtf.rsqrt(variance + epsilon)
-
-      # Update running mean and running variance.
-      moving_mean = mtf.assign(
-          moving_mean, momentum * moving_mean + (1-momentum) * mean)
-      moving_variance = mtf.assign(
-          moving_variance,
-          momentum * moving_variance + (1 - momentum) * variance)
-    else:
-      # At eval and test time, use the running mean and variance.
-      norm_x = (x - moving_mean) * mtf.rsqrt(moving_variance + epsilon)
-    return norm_x * scale + bias
-
-
-def softmax_cross_entropy_with_logits(logits, targets, vocab_dim):
-  """Per-example softmax loss.
-
-  Args:
-    logits: a mtf.Tensor whose shape contains vocab_dim
-    targets: a mtf.Tensor with the same shape as logits
-    vocab_dim: a mtf.Dimension
-
-  Returns:
-    a mtf.Tensor whose shape is equal to logits.shape - vocab_dim
-
-  Raises:
-    ValueError: if the shapes do not match.
-  """
-  if logits.shape != targets.shape:
-    raise ValueError(
-        "logits shape must equal targets shape"
-        "logits=%s targets=%s" % (logits.to_string, targets.to_string))
-  if vocab_dim not in logits.shape.dims:
-    raise ValueError("vocab_dim must be in logits.shape.dims")
-  log_softmax = mtf.log_softmax(logits, vocab_dim)
-  return mtf.negative(
-      mtf.reduce_sum(log_softmax * targets, reduced_dim=vocab_dim))
-
-
-def weights_nonzero(targets, dtype=tf.float32):
-  def my_fn(x):
-    return tf.cast(tf.not_equal(x, 0), dtype)
-  return mtf.cwise(my_fn, [targets], output_dtype=dtype, name="weights_nonzero")
-
-
-def dense_relu_dense(x,
-                     hidden_channels,
-                     dropout=0.0,
-                     dropout_broadcast_dims=None,
-                     name=None):
-  """Hidden layer with ReLU activation followed by linear projection.
-
-  The output has the same number of channels as the input.
-
-  Args:
-    x: a mtf.Tensor
-    hidden_channels: a mtf.Dimension - channels in the hidden layer
-    dropout: an optional float
-    dropout_broadcast_dims: an optional list of mtf.Dimension
-    name: an optional string
-
-  Returns:
-    a mtf.Tensor with the same shape as x.
-  """
-  with tf.variable_scope(name, default_name="dense_relu_dense"):
-    io_channels = x.shape.dims[-1]
-    stddev = (hidden_channels.size * io_channels.size) ** -0.25
-    io = mtf.Dimension("io", 2)
-    w = mtf.get_variable(
-        x.mesh,
-        "kernel",
-        mtf.Shape([io, io_channels, hidden_channels]),
-        initializer=tf.random_normal_initializer(stddev=stddev),
-        activation_dtype=x.dtype)
-    wi, wo = mtf.unstack(w, io)
-    h = mtf.relu(mtf.einsum([x, wi]))
-    if dropout != 0.0:
-      h = mtf.dropout(h, 1.0 - dropout,
-                      noise_shape=h.shape - dropout_broadcast_dims)
-    return mtf.einsum([h, wo])
-
-
-def local_1d_halo_exchange(k, v, num_w_blocks, w_dim, memory_w_dim, mask_right):
-  """Halo exchange for keys and values for Local 1D attention."""
-  if num_w_blocks is not None:
-    if mask_right:
-      k = mtf.left_halo_exchange(k, num_w_blocks, w_dim, memory_w_dim.size)
-      v = mtf.left_halo_exchange(v, num_w_blocks, w_dim, memory_w_dim.size)
-    else:
-      k = mtf.halo_exchange(k, num_w_blocks, w_dim, memory_w_dim.size)
-      v = mtf.halo_exchange(v, num_w_blocks, w_dim, memory_w_dim.size)
-  else:
-    if mask_right:
-      k = mtf.pad(k, [memory_w_dim, None], w_dim.name)
-    else:
-      k = mtf.pad(k, [memory_w_dim, memory_w_dim], w_dim.name)
-    v = mtf.pad(v, [memory_w_dim, memory_w_dim], w_dim.name)
-  return k, v
-
-
-def local_self_attention_spatial_blocks(
-    query_antecedent,
-    kv_channels,
-    heads,
-    memory_w_dim=None,
-    mask_right=False,
-    name=None):
-  """Attention to the source position and a neighborhood to the left or right.
-
-  The sequence is divided into blocks of length block_size.
-  Attention for a given query position can only see memory positions
-  less than or equal to the query position, in the corresponding block
-  and the previous block.
-
-  Args:
-    query_antecedent: a mtf.Tensor with shape
-      [batch, num_h_blocks, num_w_blocks, h_dim, w_dim, io_channels]
-      must have the same size as query_length, but a different name.
-    kv_channels: a mtf.Dimension (the size of the key and value vectors)
-    heads: a mtf.Dimension (the number of heads)
-    memory_w_dim: mtf Dimension, for the memory width block.
-    mask_right: bool, flag specifying whether we mask out attention to the right
-      for the decoder.
-    name: an optional string.
-
-  Returns:
-    a Tensor of shape
-        [batch, num_h_blocks, num_w_blocks, h_dim, w_dim, io_channels]
-
-  Raises:
-    ValueError: if channels or depth don't match.
-  """
-  with tf.variable_scope(
-      name, default_name="multihead_attention",
-      values=[query_antecedent]):
-
-    w_dim, io_channels = query_antecedent.shape.dims[-2:]
-    batch, num_w_blocks = query_antecedent.shape.dims[:2]
-    q_var, k_var, v_var, o_var = multihead_attention_vars(
-        query_antecedent.mesh, heads, io_channels, kv_channels,
-        query_antecedent.dtype)
-
-    # Rename dimensions for the memory height and width.
-    memory_antecedent = mtf.rename_dimension(
-        query_antecedent, w_dim.name, memory_w_dim.name)
-
-    # Call einsum over the query and memory to get query q, keys k and values v.
-    q = mtf.einsum(
-        [query_antecedent, q_var],
-        mtf.Shape([batch, heads, num_w_blocks, w_dim, kv_channels]))
-    k = mtf.einsum(
-        [memory_antecedent, k_var],
-        mtf.Shape([batch, heads, num_w_blocks, w_dim, kv_channels]))
-    v = mtf.einsum(
-        [memory_antecedent, v_var],
-        mtf.Shape([batch, heads, num_w_blocks, w_dim, kv_channels]))
-
-    # Halo exchange for memory blocks.
-    if memory_w_dim is not None:
-      k, v = local_1d_halo_exchange(
-          k, v, num_w_blocks, w_dim, memory_w_dim, mask_right)
-
-    # Calculate the causal mask to avoid peeking into the future. We compute
-    # this once and reuse it for all blocks since the block_size is known.
-    mask = None
-    if mask_right:
-      mask = attention_bias_local_block(
-          query_antecedent.mesh, w_dim, memory_w_dim)
-
-    output = dot_product_attention(q, k, v, mask=mask)
-
-    return mtf.einsum(
-        [output, o_var],
-        mtf.Shape([batch, num_w_blocks, w_dim, io_channels]))
-
-
-def masked_local_attention_1d(query_antecedent,
-                              memory_antecedent,
-                              kv_channels,
-                              heads,
-                              block_length=128,
-                              name=None):
-  """Attention to the source position and a neighborhood to the left of it.
-
-  The sequence is divided into blocks of length block_size.
-  Attention for a given query position can only see memory positions
-  less than or equal to the query position, in the corresponding block
-  and the previous block.
-
-  Args:
-    query_antecedent: a mtf.Tensor with shape [batch, query_length, io_channels]
-    memory_antecedent: a mtf.Tensor with shape
-      [batch, memory_length, io_channels] (optional). Currently, memory_length
-      must have the same size as query_length, but a different name.
-    kv_channels: a mtf.Dimension (the size of the key and value vectors)
-    heads: a mtf.Dimension (the number of heads)
-    block_length: an integer, representing receptive fields for attention.
-    name: an optional string.
-
-  Returns:
-    a Tensor of shape [batch, query_length, io_channels]
-
-  Raises:
-    ValueError: if channels or depth don't match.
-  """
-  with tf.variable_scope(
-      name, default_name="multihead_attention",
-      values=[query_antecedent, memory_antecedent]):
-
-    batch, query_length, io_channels = query_antecedent.shape.dims
-    q_var, k_var, v_var, o_var = multihead_attention_vars(
-        query_antecedent.mesh, heads, io_channels, kv_channels,
-        query_antecedent.dtype)
-
-    if memory_antecedent is None:
-      memory_antecedent = rename_length_to_memory_length(
-          query_antecedent, query_length.name)
-    memory_batch, memory_length, memory_channels = memory_antecedent.shape.dims
-    if memory_batch != batch:
-      raise ValueError("memory batch must equal query batch")
-    if memory_channels != io_channels:
-      raise ValueError("memory channels must equal query channels")
-
-    # Get query q, keys k and values v.
-    q = mtf.einsum(
-        [query_antecedent, q_var],
-        mtf.Shape([batch, heads, query_length, kv_channels]))
-    k = mtf.einsum(
-        [memory_antecedent, k_var],
-        mtf.Shape([batch, heads, memory_length, kv_channels]))
-    v = mtf.einsum(
-        [memory_antecedent, v_var],
-        mtf.Shape([batch, heads, memory_length, kv_channels]))
-
-    # Let's assume for now we don't have padding and the block length equally
-    # divides the memory length.
-    block_length = (query_length.size
-                    if query_length.size < block_length * 2 else block_length)
-    blength = mtf.Dimension("block_length", block_length)
-    mlength = mtf.Dimension("mem_block_length", block_length)
-    num_blocks = mtf.Dimension("num_blocks", query_length.size // block_length)
-
-    q = mtf.reshape(
-        q, mtf.Shape([batch, heads, num_blocks, blength, kv_channels]))
-    k = mtf.reshape(
-        k, mtf.Shape([batch, heads, num_blocks, mlength, kv_channels]))
-    v = mtf.reshape(
-        v, mtf.Shape([batch, heads, num_blocks, mlength, kv_channels]))
-
-    # compute attention for the first query block.
-    def first_block_attention():
-      """Compute attention for the first block."""
-      first_q = mtf.slice(q, 0, 1, num_blocks.name)
-      first_k = mtf.slice(k, 0, 1, num_blocks.name)
-      first_v = mtf.slice(v, 0, 1, num_blocks.name)
-      first_output = dot_product_attention(first_q,
-                                           first_k,
-                                           first_v,
-                                           mask=None)
-      return first_output
-
-    # Attention for first block, since query_length = key_length.
-    first_output = first_block_attention()
-
-    # Concatenate two adjacent blocks to compute the overlapping memory block.
-    def local(x):
-      """Helper function to get memory blocks."""
-      prev_block = mtf.slice(x, 0, num_blocks.size-1, num_blocks.name)
-      cur_block = mtf.slice(x, 1, num_blocks.size-1, num_blocks.name)
-      local_block = mtf.concat([prev_block, cur_block], mlength.name)
-      return local_block
-
-    local_k = local(k)
-    local_v = local(v)
-    # Calculate the causal mask to avoid peeking into the future. We compute
-    # this once and reuse it for all blocks since the block_size is known.
-    mlength = local_k.shape.dims[3]
-    mask = attention_bias_local_block(query_antecedent.mesh,
-                                      blength, mlength)
-
-    # Remove the first block from q since we already computed that.
-    tail_q = mtf.slice(q, 1, num_blocks.size-1, num_blocks.name)
-
-    tail_output = dot_product_attention(tail_q,
-                                        local_k,
-                                        local_v,
-                                        mask=mask)
-
-    # Now concatenate the first and rest of the blocks.
-    final_output = mtf.concat([first_output, tail_output], num_blocks.name)
-    final_output = mtf.reshape(final_output, mtf.Shape(
-        [batch, heads, query_length, kv_channels]))
-    return mtf.einsum([final_output, o_var],
-                      mtf.Shape([batch, query_length, io_channels]))
-
-
-def rename_length_to_memory_length(
-    x, length_name="length", memory_length_name="memory_length"):
-  return mtf.rename_dimension(x, length_name, memory_length_name)
-
-
-def multihead_attention_vars(
-    mesh, heads, io_channels, kv_channels, activation_dtype):
-  """Create Parameters for Multihead Attention.
-
-  Args:
-    mesh: a Mesh
-    heads: a Dimension
-    io_channels: a Dimension
-    kv_channels: a Dimension
-    activation_dtype: a tf.dtype
-
-  Returns:
-    q_var: a Tensor with shape [heads, io_channels, kv_channels]
-    k_var: a Tensor with shape [heads, io_channels, kv_channels]
-    v_var: a Tensor with shape [heads, io_channels, kv_channels]
-    o_var: a Tensor with shape [heads, io_channels, kv_channels]
-  """
-  qkvo = mtf.Dimension("qkvo", 4)
-  qk_stddev = (io_channels.size ** -0.5) * (kv_channels.size ** -0.25)
-  v_stddev = io_channels.size ** -0.5
-  o_stddev = (io_channels.size * heads.size) ** -0.5
-  def qkvo_initializer(shape,
-                       dtype=None,
-                       partition_info=None,
-                       verify_shape=None):
-    del partition_info, verify_shape
-    return tf.random_normal(shape, dtype=dtype) * tf.reshape(
-        [qk_stddev, qk_stddev, v_stddev, o_stddev], [4, 1, 1, 1])
-  var = mtf.get_variable(
-      mesh, "qkvo", mtf.Shape([qkvo, heads, io_channels, kv_channels]),
-      initializer=qkvo_initializer, activation_dtype=activation_dtype)
-  q_var, k_var, v_var, o_var = mtf.unstack(var, qkvo)
-  return q_var, k_var, v_var, o_var
-
-
-def dot_product_attention(q,
-                          k,
-                          v,
-                          mask,
-                          dropout=0.0,
-                          dropout_broadcast_dims=None):
-  """Dot-product attention.
-
-  Args:
-    q: Tensor with shape [...., length_q, depth_k]. Typically leading dimensions
-      are [batch, heads].
-    k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
-      match with q.
-    v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
-      match with q.
-    mask: mask Tensor (see attention_mask())
-    dropout: a float.
-    dropout_broadcast_dims: an optional list of mtf.Dimension
-
-  Returns:
-    Tensor with shape [..., length_q, depth_v].
-  """
-  length_kv = k.shape.dims[-2]
-  logits_shape = mtf.Shape(q.shape.dims[:-1] + [length_kv])
-  logits = mtf.einsum([q, k], logits_shape)
-  if mask is not None:
-    logits += mask
-  weights = mtf.softmax(logits, length_kv)
-  if dropout != 0.0:
-    weights = mtf.dropout(
-        weights, 1.0 - dropout,
-        noise_shape=weights.shape - dropout_broadcast_dims)
-  depth_v = v.shape.dims[-1]
-  outputs_shape = mtf.Shape(q.shape.dims[:-1] + [depth_v])
-  outputs = mtf.einsum([weights, v], outputs_shape)
-  return outputs
-
-
-def multihead_attention(query_antecedent,
-                        memory_antecedent,
-                        mask,
-                        kv_channels,
-                        heads,
-                        dropout=0.0,
-                        dropout_broadcast_dims=None,
-                        name="multihead_attention"):
-  """Multihead scaled-dot-product attention with input/output transformations.
-
-  In order to use only one variable containing the four weight matrices
-  packed together, we insist that the query and memory antecedents have the
-  same dimensionality (io_channels) and that the keys and values have the
-  same dimensionality (kv_channels).
-
-  Args:
-    query_antecedent: a mtf.Tensor with shape
-      [<batch_dims>, query_length, io_channels]
-    memory_antecedent: a mtf.Tensor with shape
-      [batch, memory_length, io_channels] (optional)
-    mask: mask Tensor (see attention_mask())
-    kv_channels: a mtf.Dimension (the size of the key and value vectors)
-    heads: a mtf.Dimension (the number of heads)
-    dropout: a floating point value
-    dropout_broadcast_dims: an optional list of mtf.Dimension
-    name: an optional string.
-
-  Returns:
-    A mtf.Tensor with shape [batch, query_length, io_channels]
-
-  Raises:
-    ValueError: if the dimensions do not match.
-  """
-  batch_dims = query_antecedent.shape.dims[:-2]
-  query_length, io_channels = query_antecedent.shape.dims[-2:]
-  with tf.variable_scope(name,
-                         default_name="multihead_attention",
-                         values=[query_antecedent, memory_antecedent]):
-    q_var, k_var, v_var, o_var = multihead_attention_vars(
-        query_antecedent.mesh, heads, io_channels, kv_channels,
-        query_antecedent.dtype)
-    if memory_antecedent is None:
-      memory_antecedent = rename_length_to_memory_length(
-          query_antecedent, query_length.name)
-    memory_batch_dims = memory_antecedent.shape.dims[:-2]
-    memory_length, memory_channels = memory_antecedent.shape.dims[-2:]
-    if memory_batch_dims != batch_dims:
-      raise ValueError("memory batch must equal query batch")
-    if memory_channels != io_channels:
-      raise ValueError("memory channels must equal query channels")
-    q = mtf.einsum(
-        [query_antecedent, q_var],
-        mtf.Shape(batch_dims + [heads, query_length, kv_channels]))
-    k = mtf.einsum(
-        [memory_antecedent, k_var],
-        mtf.Shape(batch_dims + [heads, memory_length, kv_channels]))
-    v = mtf.einsum(
-        [memory_antecedent, v_var],
-        mtf.Shape(batch_dims + [heads, memory_length, kv_channels]))
-    o = dot_product_attention(
-        q, k, v, mask, dropout, dropout_broadcast_dims)
-    return mtf.einsum(
-        [o, o_var], mtf.Shape(batch_dims + [query_length, io_channels]))
-
-
-def multihead_self_attention_incremental(query_antecedent,
-                                         prev_k,
-                                         prev_v,
-                                         step_num,
-                                         name="multihead_attention"):
-  """Incremental self-attention (one decode step).
-
-  In order to use only one variable containing the four weight matrices
-  packed together, we insist that the query and memory antecedents have the
-  same dimensionality (io_channels) and that the keys and values have the
-  same dimensionality (kv_channels).
-
-  Args:
-    query_antecedent: a mtf.Tensor with shape [batch..., io_channels]
-    prev_k: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels]
-    prev_v: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels]
-    step_num: mtf Scalar with dtype tf.int32
-    name: an optional string.
-
-  Returns:
-    y: A mtf.Tensor with shape [batch..., io_channels]
-    new_k: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels]
-    new_v: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels]
-
-  Raises:
-    ValueError: if the dimensions do not match.
-  """
-  batch_dims = query_antecedent.shape.dims[:-1]
-  io_channels = query_antecedent.shape.dims[-1]
-  heads, memory_length, kv_channels = prev_k.shape.dims[-3:]
-  with tf.variable_scope(name, default_name="multihead_attention"):
-    q_var, k_var, v_var, o_var = multihead_attention_vars(
-        query_antecedent.mesh, heads, io_channels, kv_channels,
-        query_antecedent.dtype)
-    memory_antecedent = query_antecedent
-    q = mtf.einsum(
-        [query_antecedent, q_var],
-        mtf.Shape(batch_dims + [heads, kv_channels]))
-    k = mtf.einsum(
-        [memory_antecedent, k_var],
-        mtf.Shape(batch_dims + [heads, kv_channels]))
-    v = mtf.einsum(
-        [memory_antecedent, v_var],
-        mtf.Shape(batch_dims + [heads, kv_channels]))
-    k = prev_k + mtf.multiply(
-        k, mtf.one_hot(step_num, memory_length), output_shape=prev_k.shape)
-    v = prev_v + mtf.multiply(
-        v, mtf.one_hot(step_num, memory_length), output_shape=prev_v.shape)
-
-    mask = mtf.to_float(mtf.greater(mtf.range(
-        query_antecedent.mesh, memory_length, dtype=tf.int32), step_num)
-                       ) * -1e9
-    o = dot_product_attention(q, k, v, mask)
-    y = mtf.einsum([o, o_var], query_antecedent.shape)
-    return y, k, v
-
-
-def multihead_encdec_attention_incremental(query_antecedent,
-                                           q_var, o_var, k, v,
-                                           mask,
-                                           name="multihead_attention"):
-  """Incremental attention over encoder (one decode step).
-
-  In order to use only one variable containing the four weight matrices
-  packed together, we insist that the query and memory antecedents have the
-  same dimensionality (io_channels) and that the keys and values have the
-  same dimensionality (kv_channels).
-
-  memory_dims is a subset of query_dims
-
-  Args:
-    query_antecedent: a mtf.Tensor with shape query_dims + [io_channels]
-    q_var: a mtf.Tensor with shape [heads, io_channels, kv_channels]
-    o_var: a mtf.Tensor with shape [heads, io_channels, kv_channels]
-    k: memory_dims + [heads, memory_length, kv_channels]
-    v: memory_dims + [heads, memory_length, kv_channels]
-    mask: mask Tensor (see attention_mask())
-    name: an optional string.
-
-  Returns:
-    A mtf.Tensor with shape [batch, qlen, io_channels]
-  """
-  heads, _, kv_channels = k.shape.dims[-3:]
-  query_dims = query_antecedent.shape.dims[:-1]
-  with tf.variable_scope(name, default_name="multihead_attention"):
-    q = mtf.einsum(
-        [query_antecedent, q_var],
-        mtf.Shape(query_dims + [heads, kv_channels]))
-    o = dot_product_attention(q, k, v, mask)
-    return mtf.einsum([o, o_var], query_antecedent.shape)
-
-
-def attention_mask_ignore_padding(inputs, dtype=tf.float32):
-  """Bias for encoder-decoder attention.
-
-  Args:
-    inputs: a mtf.Tensor with shape [..., length_dim]
-    dtype: a tf.dtype
-
-  Returns:
-    a mtf.Tensor with shape [..., memory_length_dim]
-  """
-  inputs = rename_length_to_memory_length(inputs)
-  return mtf.cast(mtf.equal(inputs, 0), dtype) * -1e9
-
-
-def attention_mask_autoregressive(query_pos, dtype=tf.float32):
-  """Bias for self-attention where attention to the right is disallowed.
-
-  Args:
-    query_pos: a mtf.Tensor with shape [..., length_dim]
-    dtype: a tf.dtype
-
-  Returns:
-    a mtf.Tensor with shape [..., length_dim, memory_length_dim]
-  """
-  memory_pos = rename_length_to_memory_length(query_pos)
-  return mtf.cast(mtf.less(query_pos, memory_pos), dtype) * -1e9
-
-
-def attention_mask_same_segment(
-    query_segment, memory_segment=None, dtype=tf.float32):
-  """Bias for attention where attention between segments is disallowed.
-
-  Args:
-    query_segment: a mtf.Tensor with shape [..., length_dim]
-    memory_segment: a mtf.Tensor with shape [..., memory_length_dim]
-    dtype: a tf.dtype
-
-  Returns:
-    a mtf.Tensor with shape [..., length_dim, memory_length_dim]
-  """
-  memory_segment = rename_length_to_memory_length(
-      memory_segment or query_segment)
-  return mtf.cast(mtf.not_equal(query_segment, memory_segment), dtype) * -1e9
-
-
-def attention_bias_local_block(mesh, block_length, memory_length,
-                               dtype=tf.int32):
-  """Bias for attention for local blocks where attention to right is disallowed.
-
-  Create the bias matrix by using two separate masks, one for the memory part
-  which doesn't overlap with the query and second which interacts with the query
-  and should be disallowed to look to the right of the current query position.
-
-  Args:
-    mesh: a MeshTensorflow object
-    block_length: a mtf.Dimension
-    memory_length: a mtf.Dimension
-    dtype: a tf.dtype
-
-  Returns:
-    a mtf.Tensor with shape [block_length, memory_length]
-  """
-  memory_length = mtf.Dimension(memory_length.name, block_length.size)
-  memory_mask = mtf.zeros(mesh, [block_length, memory_length], dtype=dtype)
-
-  mask = mtf.cast(mtf.less(mtf.range(mesh, block_length, dtype=dtype),
-                           mtf.range(mesh, memory_length, dtype=dtype)),
-                  dtype=dtype)
-  mask = mtf.cast(mtf.concat([memory_mask, mask], memory_length.name),
-                  dtype=tf.float32)  * -1e9
-  return mask
diff --git a/tensor2tensor/mesh_tensorflow/mtf_layers_test.py b/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
deleted file mode 100644
index ae6384c80..000000000
--- a/tensor2tensor/mesh_tensorflow/mtf_layers_test.py
+++ /dev/null
@@ -1,289 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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 Mesh TensorFlow layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import parameterized
-
-from tensor2tensor.layers import common_layers
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-from tensor2tensor.mesh_tensorflow import mtf_layers
-from tensor2tensor.mesh_tensorflow import placement_mesh_impl
-
-import tensorflow as tf
-
-
-class MtfLayersTest(parameterized.TestCase, tf.test.TestCase):
-
-  @parameterized.parameters(
-      (4, True),
-      (8, False),
-  )
-  def testDense(self, units, use_bias):
-    batch = 2
-    channels = 3
-    inputs = tf.random_normal([batch, channels])
-
-    graph = mtf.Graph()
-    mesh = mtf.Mesh(graph, "my_mesh")
-    batch_dim = mtf.Dimension("batch", batch)
-    channels_dim = mtf.Dimension("channels", channels)
-    depth_dim = mtf.Dimension("depth", units)
-
-    mtf_inputs = mtf.import_tf_tensor(
-        mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
-    mtf_outputs = mtf_layers.dense(mtf_inputs,
-                                   output_dim=depth_dim,
-                                   reduced_dims=[channels_dim],
-                                   activation=mtf.relu,
-                                   use_bias=use_bias)
-    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[], layout={}, devices=[""])
-    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-    actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
-
-    expected_outputs = tf.keras.layers.Dense(units=units,
-                                             activation=tf.nn.relu,
-                                             use_bias=use_bias)(inputs)
-    tf_group = lowering.copy_masters_to_slices()
-    init = tf.global_variables_initializer()
-    self.evaluate(init)
-    self.evaluate(tf_group)
-    actual, expected = self.evaluate([actual_outputs, expected_outputs])
-
-    self.assertEqual(actual.shape, expected.shape)
-
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
-  def testLayerNorm(self):
-    batch = 2
-    channels = 3
-    inputs = tf.random_normal([batch, channels])
-
-    graph = mtf.Graph()
-    mesh = mtf.Mesh(graph, "my_mesh")
-    batch_dim = mtf.Dimension("batch", batch)
-    channels_dim = mtf.Dimension("channels", channels)
-
-    mtf_inputs = mtf.import_tf_tensor(
-        mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
-    mtf_outputs = mtf_layers.layer_norm(mtf_inputs,
-                                        dim=channels_dim)
-    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[], layout={}, devices=[""])
-    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-    actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
-
-    expected_outputs = common_layers.layer_norm(inputs)
-    tf_group = lowering.copy_masters_to_slices()
-    init = tf.global_variables_initializer()
-    self.evaluate(init)
-    self.evaluate(tf_group)
-    actual, expected = self.evaluate([actual_outputs, expected_outputs])
-
-    self.assertEqual(actual.shape, expected.shape)
-
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
-  def testWeightsNonzero(self):
-    inputs = tf.constant([[3, 1, 0], [1, 0, 0]])
-
-    graph = mtf.Graph()
-    mesh = mtf.Mesh(graph, "my_mesh")
-    batch_dim = mtf.Dimension("batch", inputs.shape.as_list()[0])
-    channels_dim = mtf.Dimension("channels", inputs.shape.as_list()[1])
-
-    mtf_inputs = mtf.import_tf_tensor(
-        mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
-    mtf_outputs = mtf_layers.weights_nonzero(mtf_inputs)
-    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[], layout={}, devices=[""])
-    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-    actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
-
-    expected_outputs = common_layers.weights_nonzero(inputs)
-    tf_group = lowering.copy_masters_to_slices()
-    self.evaluate(tf_group)
-    actual, expected = self.evaluate([actual_outputs, expected_outputs])
-
-    self.assertAllEqual(actual, expected)
-
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
-  def testDenseReluDense(self):
-    batch = 2
-    channels = 3
-    hidden = 5
-    inputs = tf.random_normal([batch, channels])
-
-    graph = mtf.Graph()
-    mesh = mtf.Mesh(graph, "my_mesh")
-    batch_dim = mtf.Dimension("batch", batch)
-    channels_dim = mtf.Dimension("channels", channels)
-    hidden_dim = mtf.Dimension("hidden", hidden)
-
-    mtf_inputs = mtf.import_tf_tensor(
-        mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
-    mtf_outputs = mtf_layers.dense_relu_dense(mtf_inputs,
-                                              hidden_channels=hidden_dim)
-    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[], layout={}, devices=[""])
-    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-    actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
-
-    tf_group = lowering.copy_masters_to_slices()
-    init = tf.global_variables_initializer()
-    self.evaluate(init)
-    self.evaluate(tf_group)
-    actual = self.evaluate(actual_outputs)
-
-    self.assertEqual(actual.shape, inputs.shape)
-
-  @parameterized.parameters(
-      (2, 16, 3, 4, 2, 2),
-      (1, 8, 5, 3, 1, 4),
-  )
-  def testMaskedLocalAttention1D(self, batch, length, io_channels, kv_channels,
-                                 heads, block_length):
-    length_q = length
-    length_m = length
-    query = tf.random_normal([batch, length_q, io_channels])
-    memory = tf.random_normal([batch, length_m, io_channels])
-
-    graph = mtf.Graph()
-    mesh = mtf.Mesh(graph, "my_mesh")
-    batch_dim = mtf.Dimension("batch", batch)
-    length_q_dim = mtf.Dimension("length_q", length_q)
-    length_m_dim = mtf.Dimension("length_m", length_m)
-    io_channels_dim = mtf.Dimension("io_channels", io_channels)
-    kv_channels_dim = mtf.Dimension("kv_channels", kv_channels)
-    heads_dim = mtf.Dimension("heads", heads)
-
-    mtf_query = mtf.import_tf_tensor(
-        mesh, query,
-        shape=mtf.Shape([batch_dim, length_q_dim, io_channels_dim]))
-    mtf_memory = mtf.import_tf_tensor(
-        mesh, memory,
-        shape=mtf.Shape([batch_dim, length_m_dim, io_channels_dim]))
-    mtf_outputs = mtf_layers.masked_local_attention_1d(
-        mtf_query,
-        mtf_memory,
-        kv_channels=kv_channels_dim,
-        heads=heads_dim,
-        block_length=block_length)
-    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[], layout={}, devices=[""])
-    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-    actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
-
-    tf_group = lowering.copy_masters_to_slices()
-    init = tf.global_variables_initializer()
-    self.evaluate(init)
-    self.evaluate(tf_group)
-    actual = self.evaluate(actual_outputs)
-
-    self.assertEqual(actual.shape, (batch, length_q, io_channels))
-
-  @parameterized.parameters(
-      (2, 4, 5, 7, 3, 1),
-  )
-  def testDotProductAttention(
-      self, batch, heads, length_q, length_kv, depth_k, depth_v):
-    query = tf.random_normal([batch, heads, length_q, depth_k])
-    key = tf.random_normal([batch, heads, length_kv, depth_k])
-    value = tf.random_normal([batch, heads, length_kv, depth_v])
-
-    graph = mtf.Graph()
-    mesh = mtf.Mesh(graph, "my_mesh")
-    batch_dim = mtf.Dimension("batch", batch)
-    heads_dim = mtf.Dimension("heads", heads)
-    length_q_dim = mtf.Dimension("length_q", length_q)
-    length_kv_dim = mtf.Dimension("length_kv", length_kv)
-    depth_k_dim = mtf.Dimension("depth_k", depth_k)
-    depth_v_dim = mtf.Dimension("depth_v", depth_v)
-
-    mtf_query = mtf.import_tf_tensor(
-        mesh, query,
-        shape=mtf.Shape(
-            [batch_dim, heads_dim, length_q_dim, depth_k_dim]))
-    mtf_key = mtf.import_tf_tensor(
-        mesh, key,
-        shape=mtf.Shape(
-            [batch_dim, heads_dim, length_kv_dim, depth_k_dim]))
-    mtf_value = mtf.import_tf_tensor(
-        mesh, value,
-        shape=mtf.Shape(
-            [batch_dim, heads_dim, length_kv_dim, depth_v_dim]))
-    mtf_outputs = mtf_layers.dot_product_attention(
-        mtf_query,
-        mtf_key,
-        mtf_value,
-        mask=None)
-    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[], layout={}, devices=[""])
-    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-    actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
-
-    tf_group = lowering.copy_masters_to_slices()
-    init = tf.global_variables_initializer()
-    self.evaluate(init)
-    self.evaluate(tf_group)
-    actual = self.evaluate(actual_outputs)
-
-    self.assertEqual(actual.shape, (batch, heads, length_q, depth_v))
-
-  @parameterized.parameters(
-      (16, 4),
-      (32, 8),
-  )
-  def testMultiheadAttention(self, kv_channels, heads):
-    batch = 2
-    length = 8
-    channels = 3
-    query = tf.random_normal([batch, length, channels])
-
-    graph = mtf.Graph()
-    mesh = mtf.Mesh(graph, "my_mesh")
-    batch_dim = mtf.Dimension("batch", batch)
-    length_dim = mtf.Dimension("length", length)
-    channels_dim = mtf.Dimension("channels", channels)
-    kv_channels_dim = mtf.Dimension("kv_channels", kv_channels)
-    heads_dim = mtf.Dimension("heads", heads)
-
-    mtf_query = mtf.import_tf_tensor(
-        mesh, query,
-        shape=mtf.Shape([batch_dim, length_dim, channels_dim]))
-    mtf_outputs = mtf_layers.multihead_attention(
-        mtf_query,
-        memory_antecedent=None,
-        mask=None,
-        kv_channels=kv_channels_dim,
-        heads=heads_dim)
-    mesh_impl = placement_mesh_impl.PlacementMeshImpl(
-        shape=[], layout={}, devices=[""])
-    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-    actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
-
-    tf_group = lowering.copy_masters_to_slices()
-    init = tf.global_variables_initializer()
-    self.evaluate(init)
-    self.evaluate(tf_group)
-    actual = self.evaluate(actual_outputs)
-
-    self.assertEqual(actual.shape, query.shape)
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/mesh_tensorflow/mtf_optimize.py b/tensor2tensor/mesh_tensorflow/mtf_optimize.py
deleted file mode 100644
index c70040436..000000000
--- a/tensor2tensor/mesh_tensorflow/mtf_optimize.py
+++ /dev/null
@@ -1,271 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Mesh-Tensorflow Optimizers."""
-
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-import tensorflow as tf
-
-
-def make_optimizer(hparams, lr):
-  if hparams.optimizer == "SGD":
-    return SgdOptimizer(lr)
-  elif hparams.optimizer == "Adafactor":
-    return adafactor_optimizer_from_hparams(hparams, lr)
-  else:
-    raise ValueError("Unknown Optimizer")
-
-
-class Optimizer(object):
-  """Base optmizer class."""
-
-  def apply_grad(self, grad, var):
-    raise ValueError("Apply_Grad not implemented %s %s" % (grad, var))
-
-
-class SgdOptimizer(Optimizer):
-  """oOptimizer implementing SGD."""
-
-  def __init__(self, lr):
-    self._lr = lr
-
-  @property
-  def lr(self):
-    return self._lr
-
-  def apply_grad(self, grad, var):
-    return [mtf.assign(var, var.outputs[0] - (grad * self.lr))]
-
-
-class AdafactorOptimizer(Optimizer):
-  """Adafactor."""
-
-  def __init__(self,
-               multiply_by_parameter_scale=True,
-               learning_rate=None,
-               decay_rate=None,
-               beta1=0.0,
-               clipping_threshold=1.0,
-               factored=True,
-               epsilon1=1e-30,
-               epsilon2=1e-3):
-    """Construct a new Adafactor optimizer.
-
-    See class comment.
-
-    Args:
-      multiply_by_parameter_scale: a boolean
-      learning_rate: an optional Scalar.
-      decay_rate: an optional Scalar.
-      beta1: a float value between 0 and 1
-      clipping_threshold: an optional float >= 1
-      factored: a boolean - whether to use factored second-moment estimator
-        for 2d variables
-      epsilon1: Regularization constant for squared gradient.
-      epsilon2: Regularization constant for parameter scale.
-
-    Raises:
-      ValueError: if absolute_update_scale and relative_update_scale_fn are both
-        present or both absent.
-    """
-    self._multiply_by_parameter_scale = multiply_by_parameter_scale
-    if learning_rate is None:
-      learning_rate = self._learning_rate_default(multiply_by_parameter_scale)
-    self._learning_rate = learning_rate
-    if decay_rate is None:
-      decay_rate = self._decay_rate_default()
-    self._decay_rate = decay_rate
-    self._beta1 = beta1
-    self._clipping_threshold = clipping_threshold
-    self._factored = factored
-    self._epsilon1 = epsilon1
-    self._epsilon2 = epsilon2
-
-  def _factored_dims(self, shape):
-    """Should we use a factored second moment estimator.
-
-    Based on the shape of the variable.
-    If we factor the accumulator, then this function returns a list of two
-    mtf.Dimensions to reduce over.  We always pick the two largest dimensions.
-    If there are not two dimensions of size >=128, then we do not factor.
-
-    Args:
-      shape: a Shape
-    Returns:
-      either a list of 2 Dimensions or None
-    """
-    if not self._factored or shape.ndims < 2:
-      return None
-    sorted_dims = sorted(shape.dims, key=lambda d: -d.size)
-    if sorted_dims[1].size < 128:
-      return None
-    return sorted_dims[:2]
-
-  def _parameter_scale(self, var):
-    """Estimate the scale of the parameters from the current values.
-
-    We include a minimum value of 0.001 to give it a chance to escape 0
-    if it was zero-initialized.
-
-    Instead of using the value, we could impute the scale from the shape,
-    as initializers do.
-
-    Args:
-      var: a variable or Tensor.
-    Returns:
-      a Scalar
-    """
-    return mtf.maximum(reduce_rms(var), self._epsilon2)
-
-  def apply_grad(self, grad, var):
-    # create slots
-    factored_dims = self._factored_dims(var.shape)
-    if factored_dims:
-      d0, d1 = factored_dims
-      vr_shape = var.shape - d0
-      vc_shape = var.shape - d1
-      vr = mtf.get_variable(
-          var.mesh, var.name + "_slot_vr", vr_shape,
-          initializer=tf.zeros_initializer(), trainable=False)
-      vc = mtf.get_variable(
-          var.mesh, var.name + "_slot_vc", vc_shape,
-          initializer=tf.zeros_initializer(), trainable=False)
-    else:
-      v = mtf.get_variable(
-          var.mesh, var.name + "_slot_v", var.shape,
-          initializer=tf.zeros_initializer(), trainable=False)
-    if self._beta1:
-      m = mtf.get_variable(
-          var.mesh, var.name + "_slot_m", var.shape,
-          initializer=tf.zeros_initializer(), trainable=False)
-
-    with tf.variable_scope(var.name + "/adafactor"):
-      grad_squared = mtf.square(grad) + self._epsilon1
-      decay_rate = self._decay_rate
-      old_val = var.value
-      if self._multiply_by_parameter_scale:
-        update_scale = self._parameter_scale(old_val) * self._learning_rate
-      else:
-        update_scale = self._learning_rate
-      mixing_rate = 1.0 - decay_rate
-      updates = []
-      if factored_dims:
-        grad_squared_row_mean = mtf.reduce_mean(
-            grad_squared, output_shape=vr_shape)
-        grad_squared_col_mean = mtf.reduce_mean(
-            grad_squared, output_shape=vc_shape)
-        new_vr = vr * decay_rate + grad_squared_row_mean * mixing_rate
-        new_vc = vc * decay_rate + grad_squared_col_mean * mixing_rate
-        vr_update = mtf.assign(vr, new_vr)
-        vc_update = mtf.assign(vc, new_vc)
-        updates.extend([vr_update, vc_update])
-        long_term_mean = mtf.reduce_mean(new_vr, reduced_dim=d1)
-        r_factor = mtf.rsqrt(new_vr / long_term_mean)
-        c_factor = mtf.rsqrt(new_vc)
-        x = grad * r_factor * c_factor
-      else:
-        new_v = v * decay_rate + grad_squared * mixing_rate
-        v_update = mtf.assign(v, new_v)
-        updates.append(v_update)
-        x = grad * mtf.rsqrt(new_v)
-      if self._clipping_threshold is not None:
-        clipping_denom = mtf.maximum(
-            1.0, reduce_rms(x) / self._clipping_threshold)
-        x /= clipping_denom
-      subtrahend = x * update_scale
-      if self._beta1:
-        new_m = (m * tf.constant(self._beta1)
-                 + subtrahend * tf.constant(1.0 - self._beta1))
-        subtrahend = new_m
-        updates.append(mtf.assign(m, new_m))
-      new_val = old_val - subtrahend
-      var_update = mtf.assign(var, new_val)
-      updates.append(var_update)
-      return updates
-
-  def _decay_rate_default(self):
-    return adafactor_decay_rate_pow(0.8)
-
-  def _learning_rate_default(self, multiply_by_parameter_scale):
-    learning_rate = tf.minimum(tf.rsqrt(step_num() + 1.0), 0.01)
-    if not multiply_by_parameter_scale:
-      learning_rate *= 0.05
-    return learning_rate
-
-
-def adafactor_decay_rate_adam(beta2):
-  """Second-moment decay rate like Adam, subsuming the correction factor.
-
-  Args:
-    beta2: a float between 0 and 1
-  Returns:
-    a scalar
-  """
-  t = tf.to_float(tf.train.get_or_create_global_step()) + 1.0
-  decay = beta2 * (1.0 - tf.pow(beta2, t - 1.0)) / (1.0 - tf.pow(beta2, t))
-  return decay
-
-
-def adafactor_decay_rate_pow(exponent):
-  """Second moment decay rate where memory-length grows as step_num^exponent.
-
-  Args:
-    exponent: a float between 0 and 1
-  Returns:
-    a scalar
-  """
-  return 1.0 - tf.pow((step_num() + 1.0), -exponent)
-
-
-def step_num():
-  return tf.to_float(tf.train.get_or_create_global_step())
-
-
-def adafactor_optimizer_from_hparams(hparams, lr):
-  """Create an Adafactor optimizer based on model hparams.
-
-  Args:
-    hparams: model hyperparameters
-    lr: learning rate scalar.
-  Returns:
-    an AdafactorOptimizer
-  Raises:
-    ValueError: on illegal values
-  """
-  if hparams.optimizer_adafactor_decay_type == "Adam":
-    decay_rate = adafactor_decay_rate_adam(
-        hparams.optimizer_adafactor_beta2)
-  elif hparams.optimizer_adafactor_decay_type == "pow":
-    decay_rate = adafactor_decay_rate_pow(
-        hparams.optimizer_adafactor_memory_exponent)
-  else:
-    raise ValueError("unknown optimizer_adafactor_decay_type")
-  return AdafactorOptimizer(
-      multiply_by_parameter_scale=(
-          hparams.optimizer_adafactor_multiply_by_parameter_scale),
-      learning_rate=lr,
-      decay_rate=decay_rate,
-      beta1=hparams.optimizer_adafactor_beta1,
-      clipping_threshold=hparams.optimizer_adafactor_clipping_threshold,
-      factored=hparams.optimizer_adafactor_factored)
-
-
-def reduce_rms(x):
-  return mtf.sqrt(mtf.reduce_mean(mtf.square(x)))
diff --git a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py b/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
deleted file mode 100644
index 911badd1d..000000000
--- a/tensor2tensor/mesh_tensorflow/mtf_toy_model_tpu.py
+++ /dev/null
@@ -1,239 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""A toy model using mesh-tensrflow."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy
-
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-from tensor2tensor.mesh_tensorflow import mtf_layers
-from tensor2tensor.mesh_tensorflow import mtf_optimize
-from tensor2tensor.mesh_tensorflow import mtf_utils
-from tensor2tensor.mesh_tensorflow.simd_mesh_impl import SimdMeshImpl
-import tensorflow as tf
-
-from tensorflow.contrib.tpu.python.tpu import tpu_config
-from tensorflow.contrib.tpu.python.tpu import tpu_estimator
-from tensorflow.python.data.ops.dataset_ops import Dataset
-from tensorflow.python.estimator import estimator as estimator_lib
-from tensorflow.python.platform import flags
-from tensorflow.python.platform import tf_logging as logging
-
-
-FLAGS = flags.FLAGS
-
-tf.flags.DEFINE_integer('batch_size', 64, 'Training batch size.')
-tf.flags.DEFINE_integer('io_size', 2, 'Number of channels per feature.')
-tf.flags.DEFINE_integer('hidden_size', 2, 'Size of each hidden layer.')
-tf.flags.DEFINE_string('mesh_shape', 'all:8', 'mesh shape')
-tf.flags.DEFINE_string('layout', 'hidden:all', 'layout rules')
-tf.flags.DEFINE_integer('iterations', 100,
-                        'Number of iterations per training loop.')
-tf.flags.DEFINE_integer('train_steps', 10000, 'max steps')
-tf.flags.DEFINE_integer('steps_per_checkpoint', 200, 'steps_per_checkpoint')
-tf.flags.DEFINE_string('master', 'local',
-                       'BNS name of the TensorFlow master to use.')
-tf.flags.DEFINE_string(
-    'model_dir',
-    default='',
-    help='The directory where the model will be stored.')
-
-# Cloud TPU Cluster Resolvers
-tf.flags.DEFINE_string(
-    'tpu',
-    default=None,
-    help='The Cloud TPU to use for training. This should be either the name '
-    'used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 url.')
-
-tf.flags.DEFINE_string(
-    'gcp_project',
-    default=None,
-    help='Project name for the Cloud TPU-enabled project. If not specified, we '
-    'will attempt to automatically detect the GCE project from metadata.')
-
-tf.flags.DEFINE_string(
-    'tpu_zone',
-    default=None,
-    help='GCE zone where the Cloud TPU is located in. If not specified, we '
-    'will attempt to automatically detect the GCE project from metadata.')
-
-
-class ToyModelInput(object):
-  """Wrapper class that acts as the input_fn to TPUEstimator."""
-
-  def __init__(self):
-    self._num_examples = 10000  # 10k
-    self._images = numpy.random.uniform(
-        0, 1.0, [self._num_examples, FLAGS.io_size]).astype(numpy.float32)
-    self._labels = self._images
-    logging.info('init ToyModelInput()')
-
-  def __call__(self, params):
-    """Input function which provides a single batch for train or eval."""
-    # Retrieves the batch size for the current shard. The # of shards is
-    # computed according to the input pipeline deployment. See
-    # `tf.contrib.tpu.RunConfig` for details.
-    batch_size = params['batch_size']
-    logging.info('call ToyModelInput() with batch size {}'.format(batch_size))
-
-    ds = Dataset.from_tensor_slices((self._images, self._labels)).repeat()
-
-    dataset = ds.batch(batch_size, drop_remainder=True).prefetch(2)
-
-    return dataset
-
-
-def toy_model(features, mesh):
-  """A toy model implemented by mesh tensorlfow."""
-  batch_dim = mtf.Dimension('batch', FLAGS.batch_size)
-  hidden_dim = mtf.Dimension('hidden', FLAGS.hidden_size)
-  io_dim = mtf.Dimension('io', FLAGS.io_size)
-
-  x = mtf.import_tf_tensor(mesh, features, mtf.Shape([batch_dim, io_dim]))
-  h = mtf_layers.dense(x, hidden_dim, name='layer1', use_bias=False)
-  y = mtf_layers.dense(h, io_dim, name='layer2', use_bias=False)
-
-  loss = mtf.reduce_sum(mtf.square(y - x))
-  return y, loss
-
-
-def model_fn(features, labels, mode, params):
-  """A model is called by TpuEstimator."""
-  del labels
-  global_step = tf.train.get_global_step()
-  graph = mtf.Graph()
-  mesh = mtf.Mesh(graph, 'my_mesh')
-  mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
-  mesh_devices = [''] * mesh_shape.size
-  mesh_impl = SimdMeshImpl(
-      mesh_shape, mtf.convert_to_layout_rules(FLAGS.layout),
-      mesh_devices, params['context'].device_assignment)
-  with mtf_utils.outside_all_rewrites():
-    logits, loss = toy_model(features, mesh)
-
-  # TRAIN mode
-  if mode == tf.estimator.ModeKeys.TRAIN:
-    var_grads = mtf.gradients([loss],
-                              [v.outputs[0] for v in graph.trainable_variables])
-    optimizer = mtf_optimize.AdafactorOptimizer()
-    update_ops = []
-    for grad, var in zip(var_grads, graph.trainable_variables):
-      update_ops.extend(optimizer.apply_grad(grad, var))
-  else:
-    # for now, we can only export fully-replicated tensors.
-    fully_replicated_logits = mtf.anonymize(logits)
-
-  lowering = mtf.Lowering(graph, {mesh: mesh_impl})
-
-  tf_loss = lowering.export_to_tf_tensor(loss)
-
-  if mode == tf.estimator.ModeKeys.TRAIN:
-    tf_update_ops = [lowering.lowered_operation(op) for op in update_ops]
-    tf_update_ops.append(tf.assign_add(global_step, 1))
-    tf.logging.info('tf_update_ops: {}'.format(tf_update_ops))
-    train_op = tf.group(tf_update_ops)
-  else:
-    tf_logits = lowering.export_to_tf_tensor(fully_replicated_logits)
-
-  with mtf_utils.outside_all_rewrites():
-    # Copy master variables to slices. Must be called first.
-    restore_hook = mtf.MtfRestoreHook(lowering)
-    if mode == tf.estimator.ModeKeys.TRAIN:
-      saver = tf.train.Saver(
-          tf.global_variables(),
-          sharded=True,
-          max_to_keep=10,
-          keep_checkpoint_every_n_hours=2,
-          defer_build=False,
-          save_relative_paths=True)
-      tf.add_to_collection(tf.GraphKeys.SAVERS, saver)
-      saver_listener = mtf.MtfCheckpointSaverListener(lowering)
-      saver_hook = tf.train.CheckpointSaverHook(
-          FLAGS.model_dir,
-          save_steps=1000,
-          saver=saver,
-          listeners=[saver_listener])
-
-      return tpu_estimator.TPUEstimatorSpec(
-          tf.estimator.ModeKeys.TRAIN,
-          loss=tf_loss,
-          train_op=train_op,
-          training_hooks=[restore_hook, saver_hook])
-    elif mode == tf.estimator.ModeKeys.EVAL:
-
-      def metric_fn(tf_logits):
-        mean_logitss = tf.metrics.mean(tf_logits)
-        return {'mean_logitss': mean_logitss}
-
-      eval_metrics = (metric_fn, [tf_logits])
-
-      return tpu_estimator.TPUEstimatorSpec(
-          tf.estimator.ModeKeys.EVAL,
-          evaluation_hooks=[restore_hook],
-          loss=tf_loss,
-          eval_metrics=eval_metrics)
-
-
-def run_toy_model_tpu():
-  """Run a toy model on TPU."""
-  tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
-      FLAGS.tpu, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
-
-  iterations_per_loop = FLAGS.iterations
-  mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
-  config = tpu_config.RunConfig(
-      cluster=tpu_cluster_resolver,
-      model_dir=FLAGS.model_dir,
-      save_checkpoints_steps=None,  # Disable the default saver
-      save_checkpoints_secs=None,  # Disable the default saver
-      log_step_count_steps=iterations_per_loop,
-      tpu_config=tpu_config.TPUConfig(
-          num_shards=mesh_shape.size,
-          iterations_per_loop=iterations_per_loop,
-          num_cores_per_replica=1,
-          per_host_input_for_training=tpu_config.InputPipelineConfig.BROADCAST))
-  classifier = tpu_estimator.TPUEstimator(
-      use_tpu=True,
-      model_fn=model_fn,
-      config=config,
-      train_batch_size=FLAGS.batch_size,
-      eval_batch_size=FLAGS.batch_size)
-  current_step = estimator_lib._load_global_step_from_checkpoint_dir(FLAGS.model_dir)  # pylint: disable=protected-access,line-too-long
-  logging.info('Current step %d', current_step)
-  while current_step < FLAGS.train_steps:
-    next_checkpoint = min(current_step + FLAGS.steps_per_checkpoint,
-                          FLAGS.train_steps)
-    classifier.train(input_fn=ToyModelInput(), max_steps=next_checkpoint)
-    current_step = next_checkpoint
-
-    logging.info('Starting to evaluate.')
-    eval_results = classifier.evaluate(
-        input_fn=ToyModelInput(),
-        steps=156)  # since we have 10000 examples and batch_size = 64 per host
-    logging.info('Eval results: %s', eval_results)
-  # classifier.train(input_fn=ToyModelInput(), max_steps=FLAGS.train_steps)
-
-
-def main(_):
-  run_toy_model_tpu()
-
-
-if __name__ == '__main__':
-  tf.logging.set_verbosity(tf.logging.INFO)
-  tf.app.run()
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer_data_splitting.png b/tensor2tensor/mesh_tensorflow/mtf_transformer_data_splitting.png
deleted file mode 100644
index b83a7cdc2f1d3d773e2afa3063e643f3c423cdd5..0000000000000000000000000000000000000000
GIT binary patch
literal 0
HcmV?d00001

literal 19717
zcmcG$bySsG*FKD(bVx`^C?FvzARUTGhjd7n(g@Ne4I&^QT@upWU0X^7l<p7#X%s1`
z-`t+(JmY=G@B8Z;-^Vy-;O5?Y-RoX!t~sxH&1>$kCrZ+I*p%2PC@6TcGLot&D5#7m
zC|BNMqQhU3Yl^Di*HtG8S#?bK@xn9<h1XaPGCEEuC<MmH|F7hW=UJkl+(MC+e5mf0
zx;5wSrRUso`Ew#k^+x%?XA;6$84J}Ktun=`S8V>(NuxoJrx|LrRt#zsb!UCtW*DOn
z3g6~pX_k6e#x&|l%8$sZ)N{Pe*}7GobXA4*-RF<j9zJK{VfpZ`R<KPxm&YtdmRr_>
zo2T#e#O?dn!IN~sHP1n|HJ*i`qf^mf%Vb76__1t{*}x8kpFkEaLPEmm>dbaY_!$mH
zC5D%>H?bAq6@xAF|6doQ1An@S)a!9LIyqHnvwdG(RaI4$u(q}?EaJ1fcJ12d3~}+X
zv0j;jgM+Q<8tacek4ki^(8IfkCJa?n2vNi`>_0khxd`96bLXYA^W*Kj>}<5~E(t}&
zf&7CyZ8pn;>4dmA;|hhgGON0oh%OU~gO+<!w(1%hkG*MOpkN~|fu-}lQ2~~*8*kS$
z3kzTW`Qw3$cwb*np3omOvH78;CGGIg<<qB6ldayKg+&zn2rgCqpf@DA=mmv^CL_6u
zYzED4lgrxJI5_K@o8h|S<Kt3X<T=&VyjBe(>u)@4s-~DZI3h*xucErO7zhaoF*7r_
zZ@X04OjTZ`*0&-LzPCLe)nW8@tVm0FqrQ7Ocx3(K+P*P-ij@_yc`!UW`Z1DPc<;pP
zbDFSg=V?3x&tt~E<)4Yc#)<oDL;R<wr`ezUJKpQ4s+#DskybS{2tK|~)#mSy%EQmk
z59=ndy}g~=Wpb~nsmW9<QtRPEKg$YjDYhezmRw_U3wL+-&%ahzEk?L-N9^jP#|F4?
zs~Q_UM!jzaM{AaO>*!47uh9|5#m6VZ{Yb;RwtoDOIp*8g*ib-leZIk}DalUI#wjN!
z$EfEj+qym>#;jJ_(n5{k$}cK19ZVHeF?U)U{P5w!wGm!RYioHGJoCbrE-t9$4ojVh
zQF)&VHHsVrsIH=3@$mG_d%3_ZR5zn&ZM|7OUM6X0dV07`9{kYB$!YZ0-f31#i#I&{
zcQqrrSeM22D<}#hZ*<b*;<|DbQyd!fqRshIB5JHBzU`kNU$tIYSy?&uYE7J(s<tE!
z3JQvLo+#GgS$A$R;50WkpWHs|y)a#Qk8pm|Wl~Y@vOO(JJtK3Mm6cxbS7)b0VNp@`
zr~dwaFvgdzt^}TT0+;h&zFgNK9da)yE;gHQc8|V%%0^miHz#BziZ~bUJP1D!m=!lG
z6rjSt^=>|?I5(FO8yh?8Px;H2FLPctnYQrXI?1T|iH#nhuCD$BjQjE9$HTRRN=vxs
znOE(bIZ4WME7bX?WLBM;=3rBD2P}9}l9RuAM8ZGxIT9?ZrWBgy?BWv=Qs7aLd%5k>
zw7!3mH9j;n#LrF|uA{B}yi})(v9c6)R7-0jaD01f>kaQUGE&lHScu#?JAPZf6c{Ee
zd0(GnB6MIKHq&~%IBBn<iwN_oLI{FTqi%Lq4~(Fyyqv3Y{`q>?w{IpRCKZV-EiKDJ
zoN;k+$U%w<3+Y%I!N65EXJ%(@`1tu#_)-Sg*w}dMU@f+Ox+QNhrF^<K@MY7eLi@83
z*Yo7nk7|NT)enwqaPt^VTaLQv=eO|{S3aK3V`c=Hc@HvX!7=%jY%^VLDN&H>R{M2q
z?EbekwzJl^zdy7~5@j_vi~PD~Q8o3svXZ;1prG%!N6Wp5x+xo2q}AaM^4ak{pRP=d
zjU|C$&em9KjW+7*>myp<^i<g$%e$=&rm1X_l91pw2}JO7g&j>wsj8?LBhN@?=4~r?
zbMp^y5I+C;^C$njXM@<7JXl9v-7I<acu*cp<h#G%eCzFCW89*dFQ%bAW~j00=`>aK
z^#U8mica7B`o7${-^>m^=%6u5t9gc+GV#q?OLXy&&aLaAoA%|@!Y;`3`LxdNzT!yg
z6C-YLIO^&jw>;_AANz$dC}?RVG*}E!!U>(I)@F--(6#;&-&z*y(hBP`nU>8<1^aTi
zX3EAwR<^I4v!%7QRi>S5-WfsdV(R@csGpHI@uQ~=Jw*xH>K7c{<IelXWZ%AU?)^+f
z>|NtHn{^8sdMfY*M?76-t=kPPHu{2%G+fYrzBzxwD9f);rmX9E+oxs1U6&Aa#$|VV
zhD|5F=>~#~%Z`k&2e*uLOUCn2ESn#8nbi_vUZqPjUmMC0UtLpHl}7~KG5+=ISN(SF
z49bt)-8^uM-Q7~tGc#}`a;mCg*`L-*j@Rjyq&ko5AZn{jy09A?8z-xt5rmU*GG530
zGCvR2e6!ZhE&u((?z5fQh8Ja<2Od|lf|G;Ub*nLp#)SNM_JZA_*yA0<#5Lz`@ix0S
zhzEZ9BsXivZ&#P%yg7cw{^bQLCnpB^gdpzIN;6rJ5X$w~D$gT3T*NIVrVt}8a(A1$
z87%YwGjsE<nwFpBLVD!wc6Dmn+P!3L$6bM2+11tW#>X`w%pm7nU!V9moHRNifq!L&
zde;*zW?wKoZ$E#|F6zv8tWe|k(b34{q`>OF6!%wK-uw3}V80+N<m3d#$CKjSP;$an
zU|?AtOzW>S@7J!ip_rSSgFui7t`yuy-{2tn{EtAs(M6^3fVNOLH&U}1b9S^7jRs|8
z2rj8&#3JX|*EhU6Y9c*+F%DBQ^!QZx7}V6%a91AJomm-_3h&cH8KdFpX*qXyL8sNv
z6cg?gY-d6dABv0F95cCb1J2IQWOR-zjLCoh{{69K!0iSZ8I!k9Dol=&lJdEYO-^GY
z<p^(uZppiZ1T`%!{-syebjatVZ??&CM$Ay{dYa|D7mAOJ#MgcL)Y8i8HTK)jghKsH
zRyUX29Uw}T&usJZ^E-9j(E8Y-EK7Zjk=v3qJP5*y%kDgNl?W$i)a%#S8=IRyJzAo}
z)#}Gg^7kBz#wbtds&z|DeJFz4a=z(IEIAk@Ga8ifuQvL$iF0$C^ZK;;j+&BNo;+vg
zUaCw_PgkkCcdTW2JmKo=D}FTRN*iqaL`MgUtJFz=Fon;ybN5J1Ma9pU9M<1Bd9}PJ
z3}24Xs`@<_)pLAr@9~nSJ&n$rS-Pm;EPrfmQT_C$jrwWlO1B#!={soj{gqdKadC04
z&)vIsgCnxAlzTDX2vN*S)>gRvG-+&VV&UXm8x-DsQ`8zmLz38=BS?iWLZdudR-M^t
zm=oaJ610{`?$Ob>+j>g3X<=?|T%#*(Bjh5Krp-orxIM#P>WnyeZf$J}rlPXx-eQ32
z_dw52)-63U^P6C8^!M-IQQo*wc^+<29G*D{J%d<QT~h;NG#q~mztZh&z(|+sv3Gi2
ztX)A@xsVy9R#IMmo&2O3HM+qCUD<nKax%1k#iTFhjtaZqQf5nb%dRx{*KMWmEzv)Q
zD6aF&@gy3PgGX_gv8`Kj5TJYb>q}#>z3I|RA?b64H?B7|G&CS`scC3*4e&3BSn1!v
z`%vl`u0W=!s2HtOS|PQ>%f}}rlok{uyn;uaNo3&afKj(sK9)(()tQZn9;HM+)VMV0
zYgZ?1HMjRXO`BQ)!?kjAbCV}cX{|icXBi$aP`Vp;n^$Bw6n6yISq3+XS+~i>R(;BC
zP`-YD)9os%Up}GWBibk}x?r51Z;lN}v;Yvw*xc6EHh0YGK0B$r*>f8=@d}}o;E*l+
z8`hGaWDcU9H9Zn;i~PrtN=x@7`kFf9(v&P)iL5zP+XjY6%UdM7_{oacTgW!ww;q>_
z=zIaO{)8A8LCz;Zch*Ii&Yy!k9x7@4_VeYOHG49&c_WsZnwnDM)FM<bn>Rf~(q+<c
zu(2VWRLZmfV9Vy9Tq2}i3fF~zY(^3Q4vG{U`J@TfRaA&~QKKc5mBTNE(&%ZbEQ2|W
zmIw(eR#QWdniUwhu!gXRTOZ(hix~7_E5P1=7YQ6H>0f!BI-{K(-MJqwkfu~sT>Kdf
z@QvPH2zuz!k*tW?bWhWq_I@ISKbh4vn@T?0X_9?WBC(K1Yi41tH<ylrA;!Y=q|Rc0
z$KCVm18DPEF@-M)!jT2+sQT1mxYV({@ehnr2!%Fwc6?y>CKs=77U^}5fBkA4|L$F)
zOOq(KpPI38=G2!jTsCe7$5}i3Cv?OqQ`;_1U9WG@8`KvS710xmt*ovl-@kvKe*Dh7
zv-Gqt+(8{D$}-skQ1|!n832N{mUT1Py3CA@HWgO`I`Wq$0C{B7IJx1!$vmLhMdBdd
zX9LMMTtyXmexFc?n7RL0{&|;~nY;$ypbQz)h|{C8T3hD<8@y@O0t`0*Sn(s3rgd`W
zKHqT-=iR%Me3ucHtv@M0dX}ZR)jrkK<lE9iF3jPY4Y!olHf9FG%|BCn22<R(x=p9b
z!m`(al1IUbN<Ui-JX><aO@1kwkdOzp)zv7Hb37eZHTeA>0FMykqGwC7ZOyj`R}NHs
zy7LqA0@&JF&)t@SQn^mWHPhINYXJ{mhlewbvGA$<Hh8D-?(pan#Ej~?x{-wiS{+R_
zwSbxJ$i@jj#L_KJPEL!VONe9nGkKhbt*N=Wx%K^zCBKc0kCU!HXQf+CIP}8`^e88w
z`9-?pF6uTVgD2i$bUkB;pUntN2##Qs4hJ1kS!pRpa!N`;QIWi}h9tM<Gy&P*r%$hE
zTlk9e@}5n8*JfvD|D3{4Fv6QX%taru1X0wjWu6c%+2O?Ond`erTtw(+M!`7+nEh9X
zmCPl>j2@hFjB=;#4Dox%mkz8UkTR^llp4q@hagpJGsXSws<ulLb=!jYx8jH;N7=~C
zj@6gkX;F&Y;q&1!Vcc|aJc&Ka#uKIbTay*^m2B{wh34515a9vyeOMqRjJe=0?aIl`
zMUU!1qzXDo=`o22XH9HU|4e7{X`^kL6SHt<J3ZQ6oA=zkBP^T}t1RZ8Dt!w+=Mh3c
zn_E|x@IgMoaoRHPL!*$r%(~Iai^Ae!D!xOp>zIjrDePA6ZsUDBJ9a$RF+VmBc!{i3
zFcxqtjx;Yv3M)MD{WFqr@jI{W&1Dn3CUT5WdiTci9)-s<-Hah=E^-D4^;ub2`6~C9
z1}kvLIIc=$y9Lh=i>CjI7HSPXa_*saEjQ}$H(t>)^uC|MZ=e2oiP71UCvGr1wD;!W
z&RqX9@zbg5h7JpDsjG|tEFRo?psY+FDk}Q>;Na{0{2&~<L}6DBjs#Y#bFi{X%aPpP
z<-)?dvv#g#s5P00lkU5;cT-z+%8ee5vd<&HlbCR1<V)Jy=MKA*5K?f=H8SLQ8xi56
zLwup-`wsaOYzS(g9kqd>Awi(7JmgAQJLFAPV|O7=Pc}NUdV6~xpZi!@TQ75vLGoP?
zeds5vsY!~4hGtq?=X3r?r5{mpo0uJW?R$UyQYY*Vw>JxytbZJ<R_?)}#oaV%!sy?5
z?F?8Xi8)OFdGY@)e@}XdPb1Q+)#ibgpjXK_rV+-ny&FdlHw=5RxA<^S)K$>)$S2Lu
zum4r2x|St*O(JjxpRZ5ojX2OTpO7=98Rs+#-TTWSB;DR=SW^*jXE}6CJhCq+RkV>?
zLB1KKCo+ZVsl&BrrF>j$4o04Z)+8xNmLb}5r+{p!5d)V^zU?Eca^lsUd*p6+#%f+G
z7EmFt7VdJ5zHDEaS^Dv!e(|K_YVATw?L(~mE_hC}X$mgmQ(EKYYzHqCyLwEqquT;r
zQ`qT$4jd16?o*Hph>nfDk!43R9V5NFmsnahAshJGF=$7ISQ$eBv20w=2B~O#o8vfr
zPp}3>_6yau?lWHd`|bye_v<al^jw&z1o3xNkCm{Rw?}^*tbR^Hl+?}lUY;-578Lv-
zfA<MUl0G9e#_Cj=8LFE=ys$h<0Ym|hYnGp9=)buWQc7k#b(*<~DsVnpk&!U?&$YyM
z>nE4wXIqwT3nI^1f}G})Hgxa&rCg$M?-_MQiU$qzmC=T5at5-JS=dc8lK!0(g^|Pb
z>T6=ll)7hb|NhOA?U<5R1m9Hzufb)ql}hZaRDI&E``^3dWL_bjWg6y_jzruiYj5KW
zoT^gQullZP1gjBaadFJ8=PB-H@Tz}-3|nsgVqAT~%IYu8BTL*t@6f-zH?j~gg6_%@
z;*|H|-@CE!8Fdb}JUPXr5<_38N4YRX&#^U8`8%M5uZJM-w8apP!<6UJbkyh6YDKoc
z<dNzvx2t@Y_!mbk4wFVR+yc9E1*IAr|5%DdE|%~mF26W#8BM^({t}Km505Qz_+O)7
zsXQ{T-?=h*PEh~5CG#uavq4$slHq@@F>KPltL>}p94dW<GatL5`+@22`pB9jJ52Co
zRu5!mmoG_9uPVyTypiz>3k{VaI&=4ZUOf<d`AU8@L~4fFmWS`}C$L`bVcPhh`V5y;
z4@my=@F*jJh9YSEqPS(v_?trkoE)mt7R`UxQ(>6r^ij5Lx*LsHG|q%R=I)be3yFWO
zeCeRyn`2t+nRC!2J^QWZu%rXeXaRf$oL8k>tTr!nI}Zku)M8=&ceU<SrG>G%YRqs6
ze~1S2gRRez>#-Y$-VveO+AM6spC`zV5rt?Oggni7MeB@cTI<}-@DG1iD_j2tO95c_
zmB|)Qe*BOCWFU+KqQ{gRJ=YSvj!vDOz68W5GYf9T`Zb;7Lqz7yU43u-;z6xxMjf_-
z(gz}tpXjvth;B}l<pD8MHX#ojP)>70Ly{3!FAy@?Egts;1qIc$w3v$ZBqc3!pAocw
zO52ND=*;4IzwETPQg!1UL#%OfZqkO_8%A>6O0zy9U{Jbd8!{OTw6w^8u>+o~vZ|_Y
zU;s&;ZEbD!?n4}gxM4KB^2l9;2DlM`>kQ)J;uRQdbaX67pEmavlsgHU7WJJL+r+SG
z2b7#kh*<T@%qhN_0Hd=9DMk-|2pS5^m5Ut6vB%eLp!hv_wX?G$n-Ct_`<YN_hV8`i
zN3iGRE>-<*)d(lE>@&*+a$qgG{LnD6;tziRHX=vI523$nMDf|G=G%b>MZrVn`CrY=
z&3)9b;G70}`_y0fxNYs0k5VKvY|l~J@~jMtafN-Ts31(}f9&GI2La;V1Xi92_<KeT
z4sazinQ#b8^!Ek@_m=yOUMLk>-*+<CGy3+^ZL-Fij3$a{OamMju0`_R=Gx*@0>4Wx
zuLGNWwpVp#vvDsx`szmv0ck)G9!&kBh`HtB`T5|?N8F{!dw2zm;#_1c<R;B4=?nBb
zkx!;qMJsN8F@J5YZd3{Jo$}^Q`iEG}UkPj5Ag8jzkVw-?Q+k%8wki1uF+<Zgre!F~
zcB9&HMNYF+PcU<t7{Z3f0khYx?DNC+B&S8`gpL@Bx3yEXwl`;LZNqg*anaS(Q@)xj
z7!A+PDmXhkGrR?^@oK<BNcMr7J3hb2Z|~&7CFRh=<XZYwA=GL&-z-=XtQ{6AW3yB}
zP)F;J7P7?3!)8moDedG|@6gy6y43mKxC$w4{hS@pg>>wc0AcIrc6$E|3=E8pjXkDO
zt_Abr)>5YR+97ioD_b(0N8io}C`;pU@{du7B_uAMvf%^hqiaWw-V;e3E@JomdAn83
z@We#O!fsw|EuYnSUpzs4Y3y()M0NvF>;+$<WD$q%r%T=8@7;dmp1+AHov5siBCD#b
zB!?fTmWmx)o+(Fg*0^vSxBcwJ$mM3d3NXh>3xgQcewY(MoIy-r@QYv)TwEi47NRRV
zSee_X$xVZFVGPUEV%Q4{)N(VaK6lUdR6p2q$02=GpLQ<R(P2A3t^avh=QL`4p^RmU
zLT-W*JlXJK^+4BDxoy%g--zGe_cDh>brKvtFl9>h!R9qY<-owDXl9|1L%Jukt0RGY
zoM#gOd&5PUSi`OAm`p=T&I1h5NarFmy@}TzLC|^L;_y|+q36b1Q+ekS28w?!d2ExP
z8S&e!2iB#bDkcQlHRe?hAlWk@@Ku}mM;i`4JbDsMV2Jc#jDOvV<q9W6geDK=99yvu
z)dLDyf8X0_-LqILn%Op@<kzj-shT&;#}k2^v}{8p7&@7u-)n!cM3N7-77R$Y8N=<1
z7aV+9d%8%!>o+5Q4IeJ>A}?^S7zm?3J#4?+x^mj*FvGOC+#oj-iX17)gu3;~+Bg}0
z=Z4qurzyFa-kZoL=zVu=&GoXXg^yEj&+fQyWl?=vD){e2We=t+zIfaElncMBJ;_jR
zXL?r~O^Ec+(K!umWOi;S>Ef70WDq9vT3N<wVqWGVAB1Sf{(P?_WO1HG;Bsyc(Qni$
z3fC0g|Gl<aWBTR0->x{;cwIYz7=A`-S#`umW#sxZd_-%3sIfrMNk-P8H56QT*Oq=c
z2l*k!mECiQJ+R3QyZ)>4OD+4V`mBBGNTV>5h@=XzAn>y1H*2A2y?jwGX?Zj0@A7cl
z5nSAXSw9u`x=qmP5;@ElL;lx{?1qjIbu_^XA{ewjWW4%W)e#Kw$i<-;i8d{-!D&5W
z`0UPecGR<2c?2^E%R|mU>flRQ^?;%OUeM#cnyN<~VKUr*S0VA@1sAXEbXi6KK^t+c
z7ty?{ZlU#^ku>CZVZZ*onB4HY^x+=OJQ)AraB>QZSMay6Lq~JMs&|R<8+}py*W!i~
z^6%!d<H0~X^#P~4zPpR@zNjc13E(f!kC?0zU{4=%QUw<I;AJ%6tft|%UGUdmNL;(a
zeGo&GSY2DIs;m3@A4(5JyF~MFcR~5Zi+iqf&RX0qtaNnCzhNSk7^IgZ>0UH^oI&xb
zZf$>JWL`7e;VLL8`A}IIGuPxQ9l!GZ`<JPyJ|NTTPgYZ6E1trVgh*dt)1Fe;c;IUK
zpt$H`_0N~JGQ{Z0VkwBw5CA~8n937Bw*=DW?~Pu#{k(Df%U;6r=4szvQJlm1V=gVr
zo3d!|cHq{4i3!uG1+tI|U;6Lcv)#@G;LwEIhxA;cr^p*gv#_~6EwpC*6`b?tQ$c?I
z``p}M_x;rhj{_@U`Q(A0f-G%xVgk6dgg0}qLwhl4&YwQ7j;JHJ@V>2gzVs#HNOD+F
z6*U&`CZ7WoD8|bShl25L()BBj<YZ;}1AlqPv<{Z?V|ZD+A^~D<NMb1o46ZFdNOD+v
zc=aA+hs!q?j?tl71K|^9d~$LEED{59{T(~MF7GYpEjG0U#47tg)hoBTqCHxnF|4P`
zqz$wT!lDYwMNszuXKl6vuIE+HCH}-5fvC5`O332o8W>U%y`^8@>l>^NO>>F4VO=Jf
zkn9Z)t11a^^ZK05EHGNb)Og0l4gFGR{0zm=u089#ED0oFwxy?FWXW~}#4uH1iV+bJ
zfs2dF#l?lv-W_q%{Z}7T!$x6I4D3}j#rlsI&kq9`Wj()_eOW|xP4Xw|Q-GBR@)q2F
zv2M*>D|fJpBCRqDpu+klOG10!>^HR`m})x)qS}95bGVFqa>%(MPnX2X;ui*0xBPr&
zi>aw+8H%^A%XHg4mb7@|OXFX5PO4emSx~i<u8iOp$JVizbsmspWMc7?_3;sdg}c{3
zJw4q!G^E^Z2PwVmRB5P8qW~X@_?CZH9paI1kYZPP6_uEq2W)c;x$OFxuD-q^(5u&z
zZSW*q$y}P|fM;BCY3gYi09q8{tfi$TMXlG~?<1wWSO_?PY3TbQk-vwFHo^-xK&+}z
z6zNGaP>Y;+0kY)5SgfM;cV6??U}xYS0Q+WPXE*PUX9pw&OfqKsua$nFV}W_AgwkJq
z-PqVz7Z8cS6F{NLc*XC*EyrDIU{G(oeZc&KoXQ4yvh?)_V5OJuk`c<*B;Y^XrZ3+g
zD*y2Y9I@d-`i5>8LkJRweWr%O=hD$>#mqLFZuRrd8B1GR^X2Xc6<0QVPS*xKC>-*y
z?gQ;@X>T8etqowqOewD7lPOePNk~ZG*d+lm9@<-RrT(1dcyM=-0Y^Q9V1@QckEZOK
zmo4#kK1B5yg_E*p<>zA}xPbdt{v;;#y2QWZ7L;DKS><kBf8W&P`Pq>$T7BwVM6VLB
z_2jzVg_zso!qNz>`V+blLkYCl`FR7)V(o{I9z8PCv9@M~Y8_C?HheehW*z0}T(KU=
z%i{>U?!3t{?rb)&|9nuQ$!bvMiX5}x@Yn6*Ki*F^g87I1o5b<>*)<J`k_3sydjJRm
z;rT*<3P@B`%o^=?adFWZ8Or+l)NoR5Cd;pcAUbYn=|#Ui7fMq~Rbm2|ivpSUc(D#5
zp8ctmLQkODeH|U0?D+6dg*a0*$aN6xQ6a;8Hxt^iKiHa~%V`^_YG?$FteY*gdPkct
z5%L18+1j#+XVdW?S~G>6#llhOydtX&h7C{?idZbSXJDsl=PofLYESlkj_U|I|LnKB
zID%J>p~RA3aTB3|B4c)ZO-;>`V?)*-hsSu1^#5uBpa6cjH5Hpbhfmz<SuwHMWmngu
z$23dfS`4dT=5ccUY|Z;gBbvq2#RS&H)jQa7W)||=Mu25u1yuT>5R<)1?)T&}LUQ1Y
z5*Hn8;rHR86gwIhF1>0-W||$E-vfqiUe%K?rV!KL!Q;l}NyO&VoW6ItcV=>pO@w~z
zKCx^|;0H%=km_x3fU<%O?&|7N*#z{=@H@*;U7ci%Ze(O6z?hs(ub!*WARvUeM@@|P
zF-boXa~5ZO`V_~i)8M#bhA%6bTOoS+YLn(w3n$_eK_@yaC$`N@ZG`Pi?PGTNiK$^?
z(nK)TP@vX&BTnI@!NEXq?7{l_`}@DJzijtIL&O6qiDw->RsX_N{W9ji3=9-l08W16
zpdDBpNExx(ste621L)N)ZQPTw_$V(#SLWBK1f3egc+~2i2;(}VF>x_6<-ozIn6iQT
zEJTi*d^YcEYe`>0IbHSX)9{|<%&ID$Qc^58(^V}{2n?-dgTIEKd||SQi3zmK@tDre
zPGi^m{}X>kjk3)D((Xx)1oSEMM^Fq!+>>&H76UC`h*s?UN@Cn^og7K21i+txgJ&B*
z-SSIHy4Lnp^z^9imo#;m09^@S`!hHM;+p<<G81HKlEd<AIDj&Iy#z_#sHYLQe6VFu
z%-O7A%RKtjH#GDN0Pwwj4-XGIdJ9R(+$uBxaO}R=a|cIctP)4{gwYjuRv3LMn3)Nq
z*LRLl&dvzR*DJm(+Z$3PCjcgiK33F$gpxbWxdn4#p1{|Qh1b-2%$;lZo%)0v!#)yB
zBwf8U&i5qQIivSlv;{ATJ*ePZBFUWqSRxurK<qlU)CSEp;^T4YC7#Q|>WE&mJ6rVX
zYHE*MU0t*InCN0^8w~@MZA0w0{E3)meedWsJb4%<L&vaXst)Bhh=YmZzGBuBr5bM`
zb&9;?Qg;pXk1@4iwjEjY9~n)MWH2U*|D@CE<!<H^10_1U))Nb`4bV;G0)gHYPR3-V
zMMtN$ag1VfNk1xzj)t~pNE<W+fh~qa_C%W<xLGKCfj5<gv-q^uCagLKq#~>oO(cY}
zuk#E7S#{rAkkvaHO~mc(0vcsd92VCN2zWO<JS+nfiW&Ta|40!-ua!EI2E&zt_g(vp
zk|IQw+&UKOxGG3*;ly*t!B`XtS9splP8~mI4st>~-fJM}khQu^m_$U$!3HsU<Pujg
zX8c@PyA{Q1PzxulR~?!B^B64N&-Z0}Nw#>Dfxe^4?C}h|0XGkDAJG)n?Sb(Cu1#S|
zIoAL@bJL5S6Y077{-ELSpC<sqW*vED_T5#+^Yin6Pyxs((q8Mo4*ke0V7ou>{zVmg
z4We`mAJMeY)XYpjg!t#rpNFzp+uC+nRH;09@@l|hIsU1g?01gGlHZ7EY@L2??}a~G
zY{Xg8V!iP$G$$qS8lzYSK?vgEvu8|@W4|)ulF`wj0C)eb^>hbW-2zYQ*R<2fs?u}=
zP%5o~zd3t(CN+nb^7EK#P7Z%0p~OQ-ao4yY70xKQS%~l9Y}S^RS3u+vo0c{Zb4MN$
zN(gyDtflpIDovz>VJ!bQUnM6g6<bb99M$s(j{RsQL|Z)$vZ<-5XmgH)a8b`g1}i`)
z2;9W-Wo1B3gW^a!t4!-IGhyFVq!&(mieYVU-?QTQn7vB31j<j6+nSOLpJ)w(hpEd;
z@i^QFw*F8Rf6|(kVo+mXi(7UOfZ|(p56A@|(}i+C_W0)dIt6KaXsCkb9_H*XfwSL7
zi`5-z{Y!c^Z$r&e18e^4ua=?gU98Cp(L`mlVb1Q?{yuTdQ^KSrbOnb?EYeH+9}CGi
zi~N_l?<E=3?i3?Q^kC4GfU^ogEVVEHEH_33%bX0R2?P2Cp8#;<@5ehXemlZ~igLD_
z?*>_mp+0NAKUPB&-6x|Rjs!OVUyXl$>q^Y+`*W}fpNx2y5FbB0GV(}XzQ0_PgfRLi
zCXL@s+CXKyfOf*kTHMy77uW94jXaci<}>g=+(&O;pX~P?NweU4iuIdA4*T>DKbRrb
z0}skyU|Q5$9sQ1F-P6+(f``h;5G#0-zR6Q=y(5&Z3k<zK^Ixtn@~q`FYSP6#$4ZOw
z*`BWXoO~ZAZzK$CY^!PrXn+3brt8rZxtEw{6Jga6mt1FC*vA(=U}p+iSl0In%F6O8
zDxv{*!rl&v9#Mg!3#!xm-(!^-t@!D<A0;x*9^lTNIWXH+EqNXL5A)HD7<8h>!h-8@
zXqi!ZpyZA8fdXuOdpjceTtm`~*LR+l(suv5_+re7!?rZ|eQq^|m#Y&%(&U5P;EeD>
zbvkT1B2>W?BMO<LE1yx}1510hNQ^BfX%Q3X;!`uToh=3?*~Pm=HiQJ`tYwHl`x35~
z8q-o{tFoW+sR*OpM61Mkr&q()mWUf1=7RR{!6ZIwW8o=)b|i}?WH$|}PQ(p236s$A
zY5aD@u88jhS5+?NW%%O^$$=AP7$c%gL<)aYRPNch1Ew4`MQ{nYY-!WEveQv?@-FQ>
zTbdyV@sD+03m)dP*E}r=ecRq{wug5nDQV(<_<-3KiIjVsmKvklh1ViacnOmp5|Jhr
z8@7q5q!2}0uEtF)jWaDu&+Y9rIGpA!(tWU%WC+gTOLuW~T|PPmc#zy9X@-vE10|;!
z-L963VQ8+g+S1d}%O6jG1Lp@+-o5-21xoe7#)~h7ENyJQe+%bS`pxl}0Z-HurPLRN
z=2Do=x$}Umv*}$eVe~zQ-v)wscz8(gI5E)!2o)YG6O3|tS~rqbTnT*eKd_F44mr6;
zu9xS>LrXycd$IkQ>qKCeL>X|Nj=jYAY$hKd*Q5G3ihMj991J|baF53YH`(qaKF%cu
z7Ir(pUKH1FV(~IfW~TYE<xNNm;iDkl1LPNS?M}e<WJN$~Dpg@&ArTxxS+t>(B_7v~
z1Rpe<1$-o#lv?s$QnH)V*zhL6$w{v}6(nS2BVWHp=g+;C!KU`0hJ4#*8+%~<&wn_)
zJ9J5G#_WtVQ5bn6f~y|DWi}RM$OPl$*~Xvo@?Je=eTMMP(0D>e7n=ay3cRMW@=ziY
zAuT!xAZL(d0y7B4Sd4D->n8-$OfQBAE14S(UKU(x=eCR5)2G549{&QJw(RVH?(XhH
zZ|(n)Run(2<H%g0LoK7h`(T?a>aa?6{Pjb*WZOMSh7O351qB6}d3hm4rHxh}odL`S
z0AK;Z+0M3$v$F~$4F8A)5uZ|0QXofx)ax}U3u9wr9|5=5`RTDmqs;WG%dxQ*#NI<4
z*o+tk3slHOlK!J71SOZ1%gdLE;2JHQoT`#iQmiS1jmIY@l%GRgcm@i(RJF6~FO$VG
zIIWTf0D`)+rKP1Y8o9W*bZfas<Bl!!4-D<_4jg$+@V4#AkDSTh8oN&vEh#%;ts{Z4
zZwt}n?92<a07#xW*f=nN7Mv-l0@!pa8H`VH&;yX!(8wGH>J`vv^!4>6sk<8(7<6vj
zjgduid<%GrX-JMQBtfD=S4Pqdo53sI5437<6fw~Qwx+5uUV*o*tRz|xw#)e{s#$cS
zY$Dt;IrgrW5gkv2Wv9L8VhTa0A7bwVo%YlG6BuEGVSAI2m6g?EQw9<fSc$EfI{V>M
zh!6J?X0~0<d@e&fo;0t}@r0ET*tuP&Im2Na48F<fUsYUC8W;MRm*s>y$6Y~kxM#^>
zx*IJo1JF@uj;0j<m+9@p#QyVtN!^ho%gz?jy@3#;%cAmHl+wGZ(^8VEtl4e;oy0W<
ze`P;6>p1YPq6(x9%H)lQy|X3~NE$@Xvla6oJ@cY4LmpVmTiIx_pl|?o0Q59>%eim^
zf^f)C73xLOBg>%Pc~I_imsP9OY~@p2tPn@EnV=VmPdNsq@X~L!_a`x5wn~`Lzl8Pj
z*4WPQ)!NRG*07P|=H~Nd_b0Z02gH3ZK?}~6OA~!n4OG&RJV*JjuQv0|Ngy@CqK@iS
z=T>~a-{~ug$=QM(cWv*3gn+vXWYTnMx5uL_qI!5TlO|&axTGn0gq0c^8lv6hrNH|@
zdP^nlBT|aAPS7;D03jfjfs_uxk9LifhFMy0FgjSk-zw;Z@yIvjTCe=QaESZ`?UKXu
zhNuwj|I(g&hGY*T*^x0(T^>vc!UA~tdeqo*?W5*ClJ!3(uW=8j?;U5%m46^s621n7
zKsltZA|5`_jW1Upgi}FO!$O#p<6@s5_x-<=sFWo4qfWbTbuQbwU~a@Uf`m!0+O96{
zw*_*}qv5m8o>FN0ru!2S_E7=KfIKCeQ3*_x<i-saZNenfGU|YVHna)qy|wf5m00k-
z?U}6%bTP&;LU;43s?^k)!WndWfVvAWgP@a&mj_&rB)57uE*WX`^Z!AgUY?T(j7iYZ
zQ7HL^$ppJ^Yd`m|rrkpSmp-M0ogWE6C62}&UJ8I1{||xs8loex47^+*I=)$JBjS_-
zgggvtS39v9gfovP-k=BWGAW&i?=HrBNN7MqaStS)y>ss2;o<cCoxV(IIl+|CZp|*Z
zcHXG3eX)JrurAh!fnA2sX!b`=JFxu4vQ$lED8m<6e+1gz*+cCXW~go+C?f4%#_<jt
z*HpRjqb$BV+}Y<Pw5<o%XJ^pXT5x>aC_dNa0ZO1#qX$rUvcFaZpjFaV$*fz@5rByq
zIG%I5v&H4TeRJfkEMx)VC@U)?2U-^Sa#L>R?EIG33qmNDxD_EE60Lx&PzL+F1sQ{Z
z)jiVx;>FVJY5*{rE=CBV@aTyKgP3&W!zzr{x|@~rxC@s8dfbMwIbor&IyzuudTB;@
zh(#os^u6>kcXxm<)(6{*aRS3$C^lrnX_e6if&!ulWWIF#EZnvt3j_-u<UadUUdQ@h
zUWUOO=pNnS<vQu+Z|wkrlBhv8`i<*%atIs|^-22{$GcG-=DUCwBpFDZeQwBRg(&6N
zUJwNit#w=qGm)eMe{TcaZ_FYs#-rGr^FM;ed&0*3OiSLNj(}BK=QQ?{wX`hwlhFI-
z=+6uA%s(#p?7fL+8%a2PGmtm`2}}f}Dlz$rVP-kNiB5>4v|Krk5~<gP9gcpXc30Xg
z#y|jb!H4|9o;LM7@(cB=1t7LVdjlb}OL{*Mg@|fg9A{b@oho-{(aS_$$w6Lu1rCes
z)ulG3mM{4#<Wh9pmN3umr7kTA%M{P_pe)zG9hBxpTjYNC{s!)t0bd$9)>4ylI}20C
zDpCEnb4XA^8Ir*`?};4NJ(<qUBe*yRlp&9tSq`xH-2T4lGBdO1Y$WMeZMs;OKzN7B
zs6zHb@$JWp<DY7Ka|<1POdaeT+O;Tbe+L8lI{G014#<O;ZJqulFzLACpjdESZf5p3
z=GiZbUUA|J`#W9_khft0lH3A!g!f)BZnmHYmCrAAAmdV=Jq*B!M>@?IZfpfJU17Jq
zTOd?u7d;aB^ffz=+w<@?N6T+7yGBcDU)(`Z9xMuqdmS1Y&o#U~h*zJ`2D^mpZDhP7
z07LL-c4ces)EOP!?kaBk_u3iSMJ+UCN1F>G(zGkCzLXaSze%(-B_&ZnBN?IZq=*Q6
z|J^k68foG_VL)+|;pb@6GPUDhoKq)b{c*KGzR<_3Ydv7lK?)?$(c@9~EHF|l@QV`=
znR9!c16A%D&la=*V8mFAn$V5JSSC}6dfbf{M6c~%9HkQ&#^LnAh`gX~U;Ob^&BP`2
z24VCl)ZRyiUyLl~WD(83hzxU|U2IQt_<_}qTV{-WObGFGmQsFU3c^hpA)7QP`&l*3
ze<l|&&o)GHn<-l&yM8_8Xxj3BH>wPkhFmd5Zrf%r%vK-Tca2MmRf~iD1Wv0NnjO|{
zOA)z5#NjLaY}(~X_P(d$YC8H*W_ZZX81pj{uU*-fpiTszNyliJNc-oNa`L=i$IB&p
zp4;rR-bbzi9=}N}hJ51sIfOIpeVgJy(3;b@0fa;#DskWQrpqf63k04=h&J!*;~5yW
zu?x%@JkRj!<~``^<|#c_wP4LSL2XOF?R$ey3|08BGejxGu=G{Xvo{x1$fsL}3XU02
zS`IxU4*jM^6%V$*L0LZDLq0520cV7R3t$jPf<H7;#$5Hh0900WWPbb`!x-UiB<}t>
z4+@#+Nu++V($W}EQ}{Kx{&E;PH~`L=={Py@1RDb^hVBgD?;qQm7CXwO{Yy4U991`1
zUt^=B2OFbHvee~$1(*fsCBTkl6aB5;r@`0mzy~#oG?6s~M@Pr}($dH#J=Nu7l-Uas
zJul(e2{e9vstx%pfT(7FxD`ftdkolHYHDjiF9{MPyU)Ksl8@x|{BJbPMK0_{6?rj3
ze*Ey^6#!h{*VdS}1is^?lPyl4E=6%7^2AXz{-XLAnRO5L6jFA^>y+4(|JKW?uUna~
z@9fZ2>E(|5X&FBf*87SYe8n~Hk$we4x9B7j-QV}FJh!wohUNxHGI!>hfOn%uK}vW3
z{KZhhCJ%@?hWh*Eb<<YxdN3B(d|&<Pk#MT^IsNvfnE@H1xT_Sx+d+o*?WGX?!$2qq
z4J6-(tl#|OO9cT9o?*=Nud0GE+WMT3;~)vB-*p=ziLP=;FHTp$#qJZey0Y@P#s-Pc
zM@J=neSNX5@0k9i4K8#!9(CbToYcWjc@2qV=x4K$tS+;BaM!oqbwi?hZ{p|aIW%<K
z>%a3elAt}!H?UFqk!6a#H;LGO8PcKj@?t6gl?jptH;{i~j2m%b2Ik%(z+eGRMJ3w>
z)el7}kLOpa_6mN8>8Lk=s4;}<!+%kfXt8C{>ErqxpW}0RC{R4{7E8?~g$2u*{%`%f
zRWmo`4S@p6>({Tdq{cvv4N?Q7F0FmZVe#lR8#T)9Q_se6Bh-?rDNW*#RMVWP!y2to
z$ISa0T<nEtpz%XL<5u|ahJ@hcrM$5h83Ju!^1%%wVDOQg@~`^BLhQO3YWG>F&jTu-
zUEzL##&~vSdHiF0ujtuQcT}fnh9-Cnx(dZ`tZ@_&hxk3QmNhoMU8Gs^`mX)|u9~~(
ztpt47%;8;MEg*YzRO`Nx@uc<^6cLe(SBa$L)!`#uK6X_>6i|oI)_4_rF8r7m2pe$~
z;HB*wx<tWILUW}cDF7HF`HjmK6$s-i#I+gZwoW$xDz8|UiG$%_Ymg+$NQj&dePOVi
zy9`{xc`pT?nlK6JB3b+<AiQLV;_aqj2+J*Pxk8N{++6`^umf=uzx<M(bfdF%IuCMK
zk$^;2(##Pq)Qg_00z2|w9hE+EG+7vpsBRHdR06X<J>7Qys;B_-5k_AHeUm9ASI0S1
zzBh@1I6eq<w0tlvAR8Zm0UN<#X8%FRXUYzXaJJ{oV=NqFHvcms5sav|cmKH?)s7pb
z<94~XAV2_Q_zP8V1rPM;6;_F7<{_cmHo4a_>Ty)iU;7$`Fh>Z)E$?*h{<ZouZ-;?v
zdxv#@pNcyc;;g^)qHq*qyA2}`Z$}RwS&akD%lmdSbRNaYlV|lO#sXe!u)J%?;1Ers
zYQ+8MeBHeM;;-y!tHohJdJz}75DZ-;iJ!U7VX(G<OK)r{KA?F9Jzf7EAfl{xvtX|O
zi_>92-pW$auMfBUXa%+cJZ(ZNKftKZf`1>X*8y80x8eeHOA8&!>HapXz}M6-ivDx0
zmn6{)1)?ghDb*%|GaV{}P~+!I#W2*)ffyc|BYut1a$~V0K+j1L?PzwNSvo5H)7Df-
z{qN!ZhV@+M&q+LvG^<JHfZvIS5KL=`8d^VBpeh08O8CiA3PC(xOK(bT@oXDxZ;yC&
zCGl(#<EYYl!V$E1&e{%+{}pAyz05e|3t4X={*A4H{zNDNetuCEuDfH%5EJdJNbK|T
zS^crg@yST}g~;*tGZOF=&_`($6a1er+s$^9{FvNz;hY&|h~|}YUFXCO;G6bRR%6@D
zKQZv;Rn@yp1K|1q9Q=}!?|XtS9pUb^c{`TY^=4|^HSu#2J-2c7aaGs{LC{%JAh;4m
zPp=yhu?h&@Suh)~+;!?1LBDk8DAv7dKVAgFu8}}hLCYV*@<Tl)4g6M9vi2b(W&R_j
zyaFR`Vfq;XoHtmQVY{InNS3k@mWxe$kD4uh!KiuaqX$K5Lx)-X!F8s~+~C@><`maK
zMQ8TMAZ<`?s~h<Oa$L8q3Qe)LpLJV;2Pms4Yepw?L)tE@`=dL@-TC`}g@@d*InOg_
zOWS<$J?9O&Rhd9q7VBDe3cYRMZ)09!D-8E$B59+gkiQb9*#4z&HstPsk>D?XXSFln
zB{H5_(Rytf{}9GC6vnvwYl+>%8O~1d*AJXs>H7#gf_8e}03r-MU&7|6)K5Wg77U9`
z7h({yD~k2BNrHds^rv~b02AbEa8TK=F{*v$U78r`sG~u>GzRg<6P@=+Bf@0nN>UZ|
z1{*>6UQ0kLZM!}n8;5j=p*aNbnZ>c}%xc9dj6pZnvmdJ3-(-+3N^&>~HI{)}@pu3-
zD@f#5=(N@Hk%RThybMzyvZtLrz9n)@1`1DpgTI#At4)%~Fj*wMblL=gp)LbkyqV-R
z@~{zxg`>t^{WH>EzCkm2c~?&Z)B5k9DaiEsh<}Lc>@B%=XoV|O+9pdGz_Wt_kf9gw
zV8d!MI?}3iBmbAdH{L7<GfIUwwf}F>$nbJRG?w;k(EyUn!R4sBy8_5t2izIEk%(W~
zkh>)L<k!9VdgAhq_V%91rX#D;mcw{;O+aG|rp#QjuRM3$-}(7{xYv)4F4jHgDlXJ@
zAJ()_L2d;Xib!va$rE#Bxya4t8l6qWdUzEhX+-1rb>)J&k)TUq<RjPUZNR2xyjJ%F
zGyfkneo-eP|1ssr^58x=zsQpl`thM(9~SUM@8aX7X#Q2cr+cn{STn7d6@K1V+j|Md
z>NsUirch`HGP0YVcPL^(VJ-7V%?K43buTem$lz7_(fm-LFUnNxGB;(s!l=WwT-)bS
z{)*X6Ba@1=;A=R1``h-HYdITvI94V<L|*^(DX%6rBN}F{Y9h7yCXuFo_EWR<Y=x<J
zkYkD6-aTe(iwOCWS6v$(ifZ)a#(#L4L&E;xWI&B6ZJuZo<IvKvvI59+F042U4C4)W
zTcRDGi?XwY6t=UKF|FgxOs(UYq7VTgu69(XoXB#dNR#p!W1y~-3BZYa{rI$ExDS>C
zj`nxf)||Tg>~F8J=Q&*T<vGB=SDtm;0)3zH0S=?v{}7hrG`~=kC&N}i2|-+W0CEvX
zdL!5fuA)X)c0hOQ6}O-9m#bZbwK`9py!|4?=<&}dWvZbtJ}7>9rmjWb$)&hLjj}qF
zq5Sme)6w2k$k<<;9se}m+*xSLfDVQPs68u*oqT60{s4w%LM7&v2p^8;(EntamlnUV
z0pe9b<IQno^`qCOYzb@PTN04u;5Hul_1yxx7Q{KglEHe7j*sUy%E2S^yhciepu;=a
zND+Ck{-Zu#-aq-Vu_5PtKSlx(M0|@J7g+{3CZ`f`KtXUZL+eJ>vJFsWK-!T9`&lMK
zVG4>MfESlTht@YXvOqK0wh%S3X%1Kt45-rMU_+8tbGo3kG#9$(pqLlcI4?Q)RnqVg
zC<dO}+vfpWS2QLEfJL(??!hfE0BDD-H0_DZU*|NDU3J~aN#te$$P2y74HVGd5f8nv
zRqx*$P5vqe-m|Kx$a>N-q2F96Z7^}j+RBO>h|lbJsQHdqR7D)8u0zcw5o$`vwpn2z
zAt`-Vy0zz7WtpKeg`oCGg0683JG)Zo1Vy&kl0w5QG;LDvH}qX7rLkROj8%W~gv$gH
zfhtf62--AEY~II1aDnC~Bo288=$CD7_L_7uDG@&wgrSy;y*-dqfTmp?eSKTL`}ZSG
zZYTc`OzyR2ri+n-_N&~XwVOvyx!;7HD@hU+{Ck$~AsdrT+;3ffZaLkCf`S``{9XVq
zWTVx}e;T)8KohZ{fBLi8p_P#GKDTp&9y|}l(arm6f#O%Qy+autBN_it!U-CdM2+o;
zIY>*!+&B}$GAK)^t*s?QvXIcq`m_~lnLyAFBAdd#IfW<%EzHj=Au}DJvrO2+D;B?7
zjQ$--!L^`q2i$WYkgLrGoKU-s(AjZmg2o@J3BH+`nJb_-Qt97z5e66y7N?T2dvw}U
z)3$C#x5i8X%iH_%ZQ*95)QW-xG&6t(2qd2sAir|intZ6FgnxabYwN(H0h&p>YX%w{
zg|=Ky%g5iH?ljp?DxzcTd$gdCkdi7b%H)QGrDmHo1Ef^z;8lc9#h)i$(K`J|byddO
z?Yxb~>~^yl=$%O7x2J)dK-U$B13<P4UG+OSx@u~~eh*MGhRCYCsxu8?IpI3UTDxn8
z)-7O~(a344L+@i2*wKb(S!-*0ObpRZ({5|v7X6)xw7^5{q>fp;I;3$RxD-+b_*SnM
zcMKIl?^Bvm9yIVKW8t-cHi`#Yup3~3UWoIneH>A-=KYzP9oCARvu!ng9$$*>$FID`
zYCqG2RIn9@-b0m?d$0+f+3^K(Gbdb{4a&eNaFwn>83EZf3`cc&S_@=7A<XGc@A4;}
zKn1DO>D?Q~sfB{rOn6>Bj)`)z+M*&>?)8B$Co<j8!fi5W%1s0%i*M!2>t@`vSrvmS
zS%1$LR7SX=+3vHf2qNYrO-P3_jt2|MZcw_xO&R2eE>7q(Mdqi#%Qm@ezxGT4JBLHt
zebD^tmm@U6OH;duj6^s`#{5KDSIo%zRimMFQLKeSbv;)wU<p%GdSoA2ygk%+fhRWJ
zH-n+yrlbge+lPdROVMfc<UTCX3!yY@IVr*+!Q>yvlW-UGGT8=?IzfsgaU`rU0qv!+
zi~A?t&9bVGA9t1CTZB9Z<SWAq3*sA2(v-sUe*+!(Eqqp>Urk3nEJBZ{6o21`@_yXl
zpc1NId+5a5y(0}b)N0AZPn(gbv9e$U^n;;f<4_}g`0(K<cS3)c(~v0tw`(tscIL=~
z37u@*z=w4hajm8x*7m=C{YrRwMM2FpEqjl_h3hAB(KVn-2|Z4<KPes8(p*idrh4I4
zy%)KEkVZ!&Cod25Cn6MtTdzPY151@14_`D9kw}RNKEGng@a!AerLX5v<-6LW_A6&8
z#|xUby-^xII+t|s6v<57l{|8moY=<vgHrO^vw!7^WNfLh<N7>Kz1^K-)X5x9hYOq#
z1$v{d>5LbD&OdXh#iYtR$hD?B5s98sMxg7Kkj}UHE>zikP+j+Uciy(jk)iVJ;zs@D
z(A&@-6~)EHw^9ZUXP@bi4h{{yaC1ApWh>h4EKMX4ic#{D7P{CmAFGtt)C5IG6S{h4
ze|^9h+W}v3F%&Xn#+{UoFzGWRGHmndnwp}Jl9ry^jyJze6GcxNey}CoG$&P}`KW2m
z=RMDUg4_pxH<3s$u_*Y8f%Kl)Bc8_HzqdE5(vgYzj)j}H+g?+1zuEI>XZ)A%zVLza
zy+n~$x+THN8Dj1ilTR6l)&;i?ilnla#GI?=><}-g6ciMGv*i_9EiiSuNM@^efs(+7
zY3Wg}C0#;)M<v@(gESqXBpqRnD$D1gp{AM3@BiLWy-0(KhUVMG26nA2Jvpw4lM^wg
zLGusyGvAZ6KTPBbfM|O8>$?Ntw^u;N|F^$>gbh1WL)#<@#Kv3w=a<5pjrcn8_TC_t
zyn+J0)ni9X<HnTDxSIXlxW<~f=e)0>ITiP%!HE~mPdBPcwnr)?#iga4FH-$B>SZ-F
zY;0`KRHy^2YINV4Q>?A6U0OCtD$iZgG@(hp6hH8oYng<uX%iQh+T99>bK1JPjtjoW
zdsj*{ak@+h`7hh+8D!rwV8P(fhlt|f>>PwZJjxld)Q%5oX%QW&IG?R8OGrzjHshxF
z(TI=WGSVT1(%uUfmtd$6`fdCQcsP_13JaFb)?djbD2TNvl=JdpuK)O0pu4AM?g+kY
z#DNL}+Bv=Zo4rrR_kG;mDH==ZGlrTvzGng+07EhL^<9`F#Sh8c6tb$J%+1TYlxy<s
zANV3I+167)mBp09BYLZ!n?be{w5?`3QXf>o^!S2<gD<R3(^Qh+TQ~gr8DnMh3#Xk<
zCnw3Ndi<6JWEf-9&nvWXB``-f4_cj{#s^JJ=`g=W)rrrvq@;|Hns^Z;aiy^EP?RxE
zsp<5_W53ePbV~*e4HY)h=%KYtW^6Ivi&+g#6%sZMj?UU`&x7b6hCN0SzXJmi#;zXp
zL}BffY^iB!S5OMa?j{ec8roid`-Y10p7~R|IjAhYSk*L+>w#+4j5`bt3~$OnX>_t<
z!&Y#)4<i%P3wQT08@_NaF@l*|!tgGUp6`x=-}l{*%ymd#>boB0sj_VC>@<$Z6sTk8
zTAIKY6Q~V!bX?)%<124$d>s~sg)-q+qE@TThKgcgVIiU>=(<i%_~LHJk3Os*iE`VU
z#Kgqp)YM%^y9>8vi459&&fjlW>$3B|($m+^)Z?H}8ITy^#U&sh&|B#2L=C*o94EI^
z9b-)XmZ7iZI7H&Ma*{_&WL(@0$?UJKYZF#AivsuxBQNtjcoiRUBp~3c9p+l)V!oHE
zC8ehaUq^-3U^`JbLCEcV?&Cuv?sN8_M)#oF70d&eKM~JET$C@*d7qSMsxXs?m^<r-
zSSF(#M28BPNc?6<1GBWU!frq4YT8dZeMBg6i!y?a92fZ?WrT^XEnaf&(&FOc`~f2o
z7J2YetGIwb`^S%xTDM+MVZh=f=RVZXNKH>aEa9QTr@xMQD^4!+33K}ksy9(lruOzj
zP92xdFR!A|MTf3=7pOZ11T712;iARKeJE<tW?Ky=LJ<%Y^kb0Krxz_yw^|e<A`HO#
z8avz&QEN~}pQiLE-W)7`bki9Z56{@u_2}VrGy_(yWtM6|eVIj7*z)pn>GgY|>UXd8
zSI*wQ8*J<*#g6h`wcr*N91&a+>}yoCw3mn9^vfn-i{JK{UB^ZDoS&w`yz28*T7)Jf
zi#b%1PC!ItOJ~*MeFa<5+6(@?rk!DFIuowj9;&#A2Mv1Y4&GPl*l?m{t^#&^&X=OU
zdHdqy>?{?Wk2$Tf2fv!a&&S5%))JGFC@^mWfv7ODMQNF@emAsd+3jHC1_vi+lho}$
zU=x!wGvI%N&2`{V$mEUectkaQ;l+EM_5SPg_j7jH>UZZX@nm=@w{?GDqinmJcAU!m
z<}`w7NCbz37|+x;<uZ_9H@JERN@QMo=^DX?Nnny-2bFpU8@Jcmeq4SrbSO#ZcYf~u
zPVNo^5!R5H{sduE&(D>9vZ0rKwhAMGS<E}k3d|&(JyEpq1tB}l<n>Oj3J^bj{K#)=
zsVYC_gN?-edV|ZcVK(3p&RIuSnWXX?bDf=JqW(%0Jwqy51X!v`1-Yh{mI_D0z5Dlr
z0P<&N(<kO)DUNuQwj`#dJ<k5>Anp5V)4^zlPfAM4L+J%fB@nSxx(!YM>@!-Xep4ti
zGICKsNsD<Sk4LnswKejk5FX0p{Je|3ikP_gkqE}B7UZojXIks)+eb%9<UYi#?PQ1i
zp#LiV==#^#7zM_z>FMO_5}|S*=*)ClwWv)^O+}C3TXRA)Gc)}*SV^$qU}ZADenRs2
zc`_piX-;WqanZZr3ztM^a1qxl*^az#U<b7xXcTE=S|)oIgDq(mHJoZ_+(=UBR$E<t
zrd6UDR9I+tVp74ZTN3|7L*se!cwKs=0^|Rga!9CX*P@Q2M_nIU^i4>OG?D7H2Iir4
vo=zVwwg5X}Pxxb!oq+Vdg9(Fx`dBZ*QgZFj!A-#Wje)__)z4*}Q$iB}Q2aIZ

diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer_model_splitting.png b/tensor2tensor/mesh_tensorflow/mtf_transformer_model_splitting.png
deleted file mode 100644
index 6104a68e379ee7e1565877c03586e3fc0904c471..0000000000000000000000000000000000000000
GIT binary patch
literal 0
HcmV?d00001

literal 20556
zcmd43bySt#_CAWB0!m0I4bq_02Bbrg5Tpbpq#LBWL8JviK)NIZ5fG%iHWF?^M5Vh_
zq`QA}^Ev1DJ?A&>z2p9K$G8p)Ht)OFyVhKDKJ$5=wZm1F<?ydiUctb?z?Xj@t%iYt
z$%27#;Wf@>_)bb?VLAMB(OF9VF%JCl!7&YkuW=n7XggzI5E`NXTzD^;Z-s$z6GLA5
z-eZrnzZ0Iix;u?$8$-$bt3gJyYKgb-EK@(=8#&gF?xx%|wynH#;hRhC1o2JVl4`5C
zyU+ahcHR9eaW7I0KBH=}(;Oq`9oAVfzWD7<xB~fmou}~>&-DJp8inx7Dz)wHb-Ku2
z4k^VK=3DoZp!NKfwl<U2FU;3-ur5BewxrcT9vBuD_JG)7nC3b9KS9JKL_|dLlnRQ{
z($Y*;YzX)&CO4}czLJg&`Ty)idGeB>hmMX2VPRn*Y62%GC#@O>`rh7Nw@?A|gM$NI
zkC|Y%(3n^OyYa5RbkR^>aV;}5W`a<Y$7ffrUcIn>bb|ai=TwWnC_OzLab#s}jht~P
zwVS9j=~#0Riv82b&B^(BCdtd-2W>QVDAKW~UskNRxR{iTOm1qqtc+JG>+9<ssjO>H
zo;*QJ$^H28Bm9a1%fQ%Qg1dPkA3ofI_YE`CCcS?Bk*;oJ|58q6W&F>d8VU~|j#SyE
zP8da{m}$F(zK&%e3mhGNy4as7IqxcV>C&a`-Q5ljYbz^y(kqYD)QHQ)L`6jxU1eo2
z%gD$K47D0w%j}oa(V>C|^bn)H!L6V6Hm5w)Oj~_*i~LHk)+gi7NzsK_Spi?ZDDd*~
zvT<@UljF^9cv!;2t?lty|1A1d=jrLWy}kW9F;ULSid`Zz8s2_%bkx$`US*qJo`E6;
zchjq}qm09OCYUUxsi`TcN|nv@2A8fSH9@GAjZFo6+b2|izaY#y%&yZ@9i3Om86G`7
zJx7CR4L`poM<Hrt)sc^ciFN|}?b{J!lf%Q794KC{#M>+as4_)Y*B!*3t+s|*<KDY>
z?>fcE3f|}DwreGrq3>>OP4%PyVqj~{6dyi5e#eYMdqc^cJ9l6(<*+PCwvV4Yc{lDp
z;KILih`mr_ZEO3vjN{p}XAd4ejIc_*V|_zONr~LV(#q-+yOglLh=>SAlBm1g_)sTm
za{=X8KXEbeo|l(b)N<3puox{NM9O`yX4x{^F#!%#c2kqYqPO|ym)Bvzi104wpctm)
z1MMdsKYskHp6wDFA_jM?*<TW)BdKQP_wV1|fBsa{)uq-;h>2<a5`TN+=;Zz?+0n^q
z%~(Nk@w>NgZ`RqFnJokj?Hu^LNl2J=s*Un&79lnfBM1$QzpbhKYdA#x(IcWpIek@C
zp81HXq9VLd<ULKzq1TzO#*rn3g%Mm-y`BZ{-y81!9T_-t7or~bTuSeZ<l}8(BsORm
zirJv=j2K(F99C=iO7s4GEC>dKi54uP&BW@1O#Gzdsq#-(f~DRhCYC$QNPR(}I=9Pr
zc6Nl)l|R-lrqJ7py8WH@o;uRscAak|yoiZd#&0rXlmD_vOO=c%MjQQ!17&JzniLWe
z^5{fSRyKs4i%Wolf}&%=DH+0ET~*Z(7R#|#w_|OuvnVq$(E`3?XJg}ukB=X4<wWtq
zcA9%yos+{jnyw2!{N>A+NYnx^+<}FWkuguSxu(Wrc+~}h0p_SCnlr<Y$8y}xUC{<Y
z3YHq~e@$;qtD5!Y%a@T5TH@m3CL-L|uU|*M|A!ABEQdUuowpEkwl+5WuU@@^dm>3m
zNs~m^8l1b92i!)e_%?KeJA$TQ%Nj{#y+)WHS35COUt(A@oBwIr&_Aga=~f?TzYw`w
zP9%&Sx8p59%goF?u$Nm^<)#QfnaOMtX=P=FfL)zeRrS-Oagu)7t8wy|qx7h#h)DL5
z?NVPliNuX$F$F~?dU&*K14DSSyo!qN{QUgnnf+w$Cy|kO5dGQ}mXCTiyM~_!UnP@G
z{PF#J`w8#wzkjJW&BPL184cR@CcW)@>R{>ab19oBDav}=AlbRO1qnit-+%o0%u(|G
zJrDffX1oSrezwse4O;}#)+>aGD0(git}beIhG-?+VXA8F<e?;ySFEeAS9<j5QO^by
z-)Xz4LoDS;0&*tIbE@-M#TZRmQshE8Lc7_b)zIPb+2zqD-+EYNbWnAW^DVYB5!*XD
zI&Krb>~dFb=)Y(BcdW(VWF&J~E<Bt`(ZV9Xps0u+76)NKIWjVG-%E|bzgPA&g^tbX
z(aYm^sWCA|kRF8QAC^>C3#RZyx=wT@{`~gtVd!7UZvrvxrWKt6X(sqB)IRqjbZai3
z?uhO)@FX~s*S5&`xm7SMamN0V(T`q|;u&o>t?`FP?#U1rT@)v`H|8Dp9vB%B{xDfG
zuxe4WLZGbJu{n}dG*xbre$;h9w=W7C_0{0-Ncq<tu}RXGn`EQJW5o_mPAydiTY`#U
zFr%VnmIDmxdU`YvP>nXR)ckr{#4Lh>f?k=`B_+|2j{ht#TiM#a%g%lg6-7{GGeWF4
z_1nd%^%5E0NFO``te5H9WMfaJWTqj{18eKsu+;Z|zETv#yNKD7a>x4n@}ZB!^Q3}2
zqOHYV<ZqX#IlKn5iq0D<^L~B}!#Bkm-KK?tUcS5n>4r8s^hcp4mOVZx$L-r#<b%SI
zTBU}qmIQZ1Sn?YiX;5WdU%oJqUXd}V*X4qu0!b!kbW#g1g9~rksg{k6ZQ42WNLOkX
zYJ!xrGasx5q$btJkJ}qOx4L?J#aG8@mMySYmiCG&D_z^?oNILl&J3D-gnfN|zYh&X
z?>l!IYu})aW)(>i;p2;=5p{#RYBl@$Zf`XkwqinBw6wK-MG^OHd|VN#!d-=U?j*Z6
zNnN=~kX9qexXm|aKfAcPBK%K%D5<ClK7M53<0J8~em-i5h0RHD%<%EU2TC`8u0(4^
z^o($zC<Q_S3_f%c+rRzWVnS|r@aeLt>}W`8<+NKpWceJJeC={`8ImyDu_~G7q}0?@
zs1=NyoVJQfW^cQW_+!6)`}QmNOBu)J=4Q5S%us^az{G_3AKNI6M~?<_4GdFKQr;I9
z<}@{B{Q2`N5&?+~cH;Yj0t)xGI~glcJLN>$0^@ZtwoyBLyt-;MDo`*X-NOcLXL?|s
z)VjGL-4{Hz5+#Z(eE&Xdanbzn@NgK?J0va+-5Q6kyCcjG6QDqbcFu(h%f;%#L<9*_
z6Krm6g>pW9lbZU_(^HsA(5`)HONLuHMn9sJ0hO+7u^sUhB4FQ){2-gEx6ZDb$^=g$
zlY6x;;O<Qkz6Gtz(T8+Cct}o3jcCQ(;`;idq%rdSca#Uq^)MUQbUxQmW$)j=$AHM7
zMwevF0E~r&1+4@b$nPCV^d%(eZz6|2#kM`j7T!P6ZuGiSs96&A8`s9t@@-CziPY!1
zj=G5r9|7VM@5@Nm`M`Ji`Eq;e`uY;4EXuhz8F`Em4FlW8moWo=9i_&{Uw2P`@d8ss
zR5Wz^nga2s4S`riBG}rDQK?~(9ns^W(c?!^?$HvwsixC)SXa&P=iTZ<$JOb2>tZ$t
zk@y+ZE97xqT2US*in7i&ND7n?rpY9w<iOXjh5&Cu(*i%GA-l-qK7HxxB6p(g5EbGl
zDtEdZ5veYuu~E@=RO__x<(D61az$(F<=3v*(e>mJpKJt)OcY`80tpBf09kouWlmkq
zU6GlPJX4@#MhPvmI?>LiuXVY<M%Xv@Bk_5o@w^h}-?dA#6lG=6i*((Y<Vb^M{8cd+
z)<w2B%^lrUhlSlOlUGd8gd<~4gqN1i18J8`GP;wb+-^SN_9Vt&0M?F|w|9;~*amUj
z0QHyy1E0xdLP}aY*$~7`gp8@|qxm*UM8XztMmv)eA`I^?^|4nG3nK#2A#LFTyARI_
zU>ZU>e^7`z>X!^}S*Az$T~mhz!wvIXJhC&(Dp4T>1<Jyt%&#9Rv+S{;+NN(&Ysp66
zjDx?Tbe<mJagSVVeQWcW$pH=N5g?49_eDjK$-(LXEQ&J?wST^*AdYG7V{p9>lbS>J
zW2D@I@IWo1XS=(vPXJv4T(On0X_o6^o?^#b^@5Vc|MA6(7fBJaxpE9nk8X2wNB0?r
z4XN`am_)FW<54D}uD>eI&TcccupspA=?LOBjA_=GJJ}H)J%2zo5*DWz3iJcsM&Mo!
z9N4-hT>bg;=K*IbdjU050Xa9fDt1oJJLje6bowG<Q<y?Y{%F8XP@g7lz)yf>YtJV8
znYHztpDe?YgM-oylSu_3h9wT-sbDgF^1-;-tpnq}f$2F~74Af8N%oi|z+!_CFO-y&
zicg-z1JogpRTAXqPwME9h3{Y6E2#*tLUqB!^|S?FeT7<hdBforKnzz`*ZWO43}mj&
zQ*@j*i%zvY=t9%zpsA@TSMTra<P;)?Z6dZzP<;mIA&Hh2YPp?Vl%He0?#EvOdR&@4
z2G(=JRUT4Vsy*AwFkL3J&bgO6rfWLA^l#m|)e}j<Crgu$uoo|3a`Fi@h<@VdcZx7s
zKkz|g(ezbuO4#-M0MZBgpD8OxY1g?}jayZA1oe-#>@}5^Iy^X`GMV%ulszg+AF!#e
zsBr13C0~IE@b70auAz<V5~k*&c-i32_*rnDdSQM(2?F~TyZnU&gb6vNZ5%-MTeolL
z4E}<@Hx7Kn_NN0e_yq)ruLK*LGgVYqmm48Nf+`-)*6MNr5*Ybp5;~E^H2zpZfI^#%
zH2eK~Huem|!v5^s>^B}7QB!Nt$%2{b=`?(r&3Sn&iyOZlCU9iIPBB_xV95FL?VHrS
zdjS(hVYx+blDNc##@1F>qm&U5J*{+N*Ae7+-+ug%(a`uAG|Ku^T3$LYy01?e^8N_m
zhKVfnH4hXoZyxwmR94D(tA=?0c{lhQk|DlMyG6x{f+7p_BLmN3)q7!V!d0GIi_AZ4
zNqlcGUw`-5_}lRC1ITEAMhC9EKwhtLs@1)S8DLb%8&##p(s_Dvv_fF?O>w~_?tAw6
zqmY{rLr?Sa*W2%}s02r8qM+Yo(o)z-l7?OY8Y?JPzf1f1aW7)l=zhel?%euQSW;%v
zNx;j;Cwwp;*Xe$Zfq_9l#O2pDfRP+eE5Zg<BqvrKCXLr@7Gf?1-cx`2R3z6U%1k@&
z(<j7JBd7IoG%nTDOoTn)h$4^9^Q|H7cAxUCzOb!w_Rnv%)2Rvfsr0wE-He%hqvcIY
zms_y8@O|3My}U%w@QAj{R9j=4{f%9<x9$oO-{)1&bnVe<eKMa#-|O8cg&Lo%2J>*K
zcjoQv>%MuC5_zJUe2@0qStVll35O#e9JZTIHh2PZ?@X(K;uw>eN#~pE`sB%VsNGL*
z7JvE_1GVqhA&D&=htrXHUz#u?+9Udbib{07PXlpd<8OtfN7AP6O{rpEZ-?XFb#=W1
zeW_MLFd{^ija1#x_F!`X(52&2y8E}Wv4=OW8J-;O;98SGzdpyvOdJP&ArT5G=OxNg
zs0YnEi-Cd?*$sHHKq<VNNvq06Th0ETQOo<mG6xXH&F5~hF~lCfNka1g-&BoZwi)E_
zhlit=37?q-zKZ4l|JjR!Jz2<U&u^=2v)-u_7hs@W8;<*IHvYtR-sJioUsf$=)uUnj
zU8F<y)2I1we{pG5lYA5+{3B`U{&=IH6Q;%VCGTawGmh_nU$E*2So6-*(*;HJW0VZ7
z{+d-s&j9ZwPlKS8iS@6}sty7++ySV*O3@E7>}XmfSt?_ckjfnyDB&?=v){VIBYBkr
zwb~^MNE`i@>(vr9EA*C=*h{@H#%{FYleI?P%aJFf7Ha7g9GclU9?iHb_qH8|z!cng
zNI5=9;4XQkFK+tm^fpKRMdc)f%LEoSfx|-l!{O>R)b`#K1}XKLgB<??ahxifw3k;M
zZ1}J+#Sd8acXP-nMU6r(ojD5o2#_g@8|`P-9xT@RvoH@~b_*ZgF(ZzH{f7B<ax!Xc
z#T*%eKTuUwCGavLBB!Lhyc~)RJs%%mUD2mcM5r>Jj%&a?G<v!A?8-fZP?En~LKjd^
z8!S3HQ;A&vt*ezxNWmI?T)sEe9C3PDqT&@v5<?s=52;*{K}|zLR$H3_8g~q+2n}ax
zd-E<LCd9-ztTb40uND^<lWe2Db1^U&S2Lc>Y%o@fU*x6<DvOwr<O(bwl-=9vASa-c
zxc*@EQE39AW0<D>Q;-z>4GygqAbgSmxncmKn_pU*+dAm#>B+6FO`e^dwLCHe^juL@
z)xpHUl^AjNCeMf6)&_lU=)_!HTmbbYCShGCA__;6d%P7D8_cUGlW?K;Sx|7Wi1gLb
z)|T1cuDGVG+;DeO52cp9Sva-U2{6#0e!}i~k{NPAGfNexc}g;u40snvqIZnzVaD}N
z_c?q#s2-mZ_bCM$_>ivF*!jrgw(JZU1S3B&e!9n^^=;~lyo!p&+35)a1)ba{j@v13
zG5GJ?L3r=lL?R&h%h@CG6!kAViKOrQG#QbkJn*)&vwIPRjz*W$ZTI&R`2NRK^BI>U
z_3$(sO=ia$7ZQH_9j^`YJ;FOVu~7C2jJx@pB+di+9q8txcAWqQ#6fjXVgQ^8eY?#_
znc46dB1Y#9vUgheXE8Nw`tEf_j6b_AC7A)#Wc)rjH|j_<sW~}0Ux|CC9E=XxnH}vK
zY|Tc5kJvH1YcEa<fGHGpUo}12T^TT|t*u3fQvqWIko)V<kcC_rYzvdkwhLNgS3bKV
zbGYy6A;k6_wce!fzbuAmvY9+~_1lP8V|GL0#yW|!EdMM<F7L!92psdw7Zq{b^r7lp
zaQdvrdJPBK7FinzvSf~c`!iO}?&6dmB{%&iuhKP>Dl66&I?v{Q&(hS;P%|)~i+cmX
zP*zs9?$tP;BMR$+fhx<ES1>Rz=qXBvcFxVsZQEkDl8_on-2dT8VaWKEr&Tr-s02=`
zJ;ckfuJ7*aN;Tyqlnpc5f|=aU7d#zC5ifI;+J6@7jMTYW_mp*2@{Z&xo>w%8W&Rjo
z6bEv1OG|moHIz&pVDg;<);Tin-n<s)5t-+k8&HxeV5{(fT?2CB_@oe~q%XCd|HZ?}
zAQk;925dni`HHK4b0-7T?7lHJYQ~Y_opbaMLK728(7ZtNn36&{ZYKzBA>xU_1p%Ak
zySBEr+y5#Ihjstdg-6v9u)W&{{A%3kdYZ4E9_;}VJ6@ccmYkA2keB};AC74a9jsQM
z?3q@==NLx0>=H)i<@)-1U~deeG$59Te*Bnm5&6Bc!h8*<)bHdFQAwDKX$>ijvg2-{
zq}a}0!O2W(QDI3?g!vea_NG26_EbX|3g5}e3FF6V^WGGbs_5;%RN_=}($d7gvzFeP
z3wjs?-R=)O)4Q@YY3Kg9oBKBju*ETtH}_F9w6Rks*SW+=R!?zsef?(hPn4gWbI8PR
zS#hdY1CBB<t4W%?X>=2Ff=mIYpmSHt!?WkTgvct^nQe7LBqp1$U;`&}ULq6Sd9j(h
z=0(#~P3E`HD)ig7bk|^7-5GhtJ(oYo3XgC(lmJOsKkZH9u+LW16}Tyu<eEJZqlAC{
zma;SJk7|~}*KJd$m)Ac%oQVi_SS)ap6R13AvxP(au=XUc(+(fn;Mn<)71d$(MV~o%
zCEq-MS4VQJ=HCIY>Sww&v)HauN?cqMnON=06F7e{u%a2GvVJ;-Y`V<3Mg2<f(lE`B
z?eiQypLu%xtR^(~2jQ9ojo-G>{?OIO5%V5u=n^tGSx&$I^Zt6Xe~5c7vbrRAnHw{d
zICf4MT_+TEN~nX@Dv&C4SWEKjXCDLJks)#U{zlmx{&TJJ6qV}K1<7A8s<qb_Zch)^
zI($*AXCy%nhu5~j@1ox^>;CGbLsmuqqLkUanri9TKr~68ycxBP+-04q3Po1iC+J}c
z>Rq--H+MJXO*((I>=6I>kmCIqe$&lyhe!KWN*W;(G2IU6W$Aijf_FTl4M`}v5pS&-
zqiO3x%{;tiTnu{lV%Xnpt*Mq@-4(<3v_MudJAC$ZlVi?-)=gSEFDQTa{sist06tl7
zh-e96VvMF^#rJIVOKPLez7AvezaCo%vEmDG1`#%sm<#=lQtUW~7rD77C9O0o4nrod
z-7t@rj;2i#Tf*zg-tjm3!v0b_-C!C29emk-U!b?y;BqD^SY4%5lvn;DdJbZV^(8SS
zJ@M@9Lj&#GntOx$$(}Ep!f2(>r%!&xQgZyrVtD&n!#*LKx8QN-!psH}@6P$7<eDma
zgeXnWdWc{QEwr5ZTy)ZPlVgdKM=uLg`h@@0A%@>xxF7wn$DiUkPeHspiT!ik#_+?L
zXT{C_Hx?mB*Ajh{m|A`RueGjPUJcgSYvIle`HMHad9Ox5=kq`wBYLF9sbqvLi*j#&
z;|eW5e$V-W*?HRbR+NA(8cM{>ecZd|Q-bC9{Zf8UP>DW%rUzkKZsE6z{Adtms0Fd7
zg6LfC`C?PwQB5Yiy?Jv*8lKSfT$Ev%NR=~Mo?B0sySaVYXI5=riayQn<wT5p*&}r9
z#VuM&9v`#s|IAoN?k4mEaO+)IOEPmejXmFj585cOL{DZ4n*ApUuTHF=EYGCpkI^JE
ziJnibYP1*~S5(!WOwFwBzBa{tx~==;*js8Wh>;pcMvRq}m6{1VNM6?lK^6dYgxmxt
zSgJ_BQ6vkx5+@s#CgN5410_OTAK4n8<D))JPGXAGH%aXJdV0OUwrQ1{;{#II_mM#J
zt}8(PLV{1=Ht2hK{Wh>?*>!ay=<bHMq51Mm9`-brWN^mK@zV?YpLspbE-5IO{^c-<
z{3IG?f*UNwp;h`kUDTb|TEnB_v4sUI6o0rabPe6z-C3a2^z`)n>W4-Ps$)yg81|M|
zBfyBkfq{Ga`ufI%TGC814KCN3PI$(*F7`-JG`R@BX`(XkgAk*5W!jkVPC7dLe{cbY
zj|AO0Dxiydj7B+ib#*M`fN4HS+238^@|cPSm}%JC+_EN#Pgnm~K3s7;mzf3FVeW_<
zoY^HMggARn=+0&8cu~+~6rF63_SRT|#H=f-uctzlIn+-ea8Y-a^GP!M1!vrdh;T(+
z$gBPa@164KGOajoO%rC~lF6%o){uWAFRz#aR2w@x{WTnE`(~ic^nG`pK&QMv((fmb
zwK6m`1nh&Nx%pygwGKnIG0P|y1MdWl#1+^rlY1t^HQ06@tjDW^kXdi)UDsi5Wnl(V
z*-so;0W?HOG&3_B7#L`s{Ty+yDA+8#YU|yo@3uPLy}L3-Li8al3@>M@#n(R|yMB5x
zeVu!$wNIr@@FT9IgNiAyqJr`4ju#r6jO`9Y`;nFmiLTm~=u{Jc#`p9o<+at3uz#vQ
zdg^#5AmgqC5ZVR%`(un#tQSVNE?UoE=1&ZCKNe(nzcgVd=)JdU2At*nQ3lrs9XVxX
zvA_*ksu7Qn-G%zUT}2dxU1DxwA%|TVmG7G?q*?q_1~|~3GBZ@dC%{6EG?`&x6C(O}
zgf_mvJlTY(#F$k*K2AL9mbRU!i!aux?%=o@EF~brhslQv0u)-N00_WBdWFIy5~2&5
z$pMfg%N^%B5ccI+S*`mgr@yQ-`z?TRdxhwjpPz59fS&|@{`5?G414+4FP`VWfBhmK
zy9ZgMK(g7mXK9sh@pvp!+anQCEManZv=XJjqF$ztVp-7W8?CfP^M9b_(K9jzQt+9-
zj)}PkjO1;NB68;^Pog{T^N{D(*QWwt2i_U<vjN&ypht?)s<@m0y#nLSE@4Hl?30sX
z!lvV!9EZ5a2Ty8jiH;BFN^d^=JkU-X_tcP+^vb~K=*#{k<ceDyXobKABvr@6#01>=
zCD>lZ5k)Ob?5n4#89Mkox4Js<?#-6Non@1snU+U)yPv+1l{T%=_r-?{TTgdsN#0@5
zY<=M2^u-_a$pW!EPX!A2`n|oqySut>!UF&#cT{RtQc=loZuW2Ak^K{#Y*`=H_4w=F
zWqz-jKIJ5%;KE8jH`+Ab!PPG5HWp3<@Gk5i$(rDSoR*!<prD{2bh16@IURt_VqML^
zK-{@3H~mp=Cd$oYg~x}5Oyfq7_e1H}?%YTO5ZC4aP=OI2oSLFr*%&T0!ipB##K}B6
zUSc(YYk9k3{A!jQ_lM|^hmvjBPHhWwLwd!z(y@cNOe}!HU>+gi(3Us6e;;=Jwnizm
z=-tD`VPVH{i1osJqnQcSAD0g-?k}?%onU$#S-Huv@;1P(gA5QH9PBv8RsWawMeDTb
zUkz#L=aLdv1;?>&`(J`)&wwLP^y5`B1rm9;xY(rmg$0;<%Ie~G@0bf@&t8TpUiUu~
z_4%v6zu#{df>p&O^Hi<o-(G-_vjZGmue^zv@W{zrJip?PAM-#@e7?Eqbo1E72%-Ug
zDfV82`krKgsalQgcY>$;>qQ7*t5o%$*AJehylKzP%?;wzsZ!O`Q}Xo{#{i<J^|uQe
zL03O+3=50wy^1|u6>4-8=df$#V2v2x|L9;J{@_DO#qTy?sDY{&92ueTy$Q6%7K^nz
zU>DhGRgcj%z34V8)|N9*Zl7Z{AxE{_>fzLMai@q3L3>&dC7_F;C;UP}6f+>YZ(hUm
z6e&sf^)y&&f4srL2_loJ0$GxuANuZLqMD{Anch#3B|rj)EDge>!*X#oi&G*4gZd<u
z+;phk`XpU2s*Agk-0zHiKN=PPrt5oLX?y$H-ke)K?5Z#cU2f3QpFaHp1N2yL#nuK{
z#CquCf`#OfcDCDX23|9omC!|Bo-_aHW`8=r(n<R>{L?L8o+Hxw9Fk;e-1>O+Cr@PD
z+yqo4Ph4EI)zwL^an&RwCwITNbQLf^q>ErG`v%0L`Wr;aYLNjq=xoNZR8~T*%V*^1
zaQkbBlK&9@G=by>BxjiRPhj4_vyhyeJYWX=>feca!AJGi;X+z30y~=**PEXYs)S^6
z2Oko$v&Z~qG!kb44Z-UwYc>?662F~71S(BA|5o$MFucuSRH>uj@oLG;3(J2a3hAik
zAZ8Z3>Qv}Dd?!Xm<o_WC%^O!&SLtyh!@67+)`E5DUL0lqZ)81d$?MDfSECOm2LzPx
zrjI2hoxfdTQ&XwRGdkxU$$Jqc@9H07`&8oJ7{9c_eYY3W@_@-_U5c5Rw~@K1({!TF
z735bQ=u@FR(Jx7?v>uWjTWj3Y|5;)XJ+-EZzCqQC66D@)Nd8%d7|Dl{5Pxq#Dzlah
zcQHxg8XClQS0|{HBz_iYBO!(4tbiH|g$Y=$knL>1x~?TO1gwlYyKO{R3kRwh*yYL<
zqP#Klv0tghf#O~ntM0U2UR`CmefxGyJqeM#mpyH;H*KiBB=IWe!T!7T&r(8jj{*G3
zD|Z%esmh0kDYBkiUOjr~VYoacH50mA3{YOB8hyvS`V<BRTvNQ`JXYv)0Ll;JzoL8L
zo^I!E%203*0N#Ta6_SK^>XGF(T%y=153{DJ>vM?ZInnI{i{;eqnavGbKaV|%^&CQE
zv;JGFW`?<#QW113n^RbL`%x--aQz|E4@J?y!~%0Y-FhYwHJzf!%*-$sUrHMXg3Qm4
zdJ$7=`4A_d87f`95ySp_ArN4i;qu(H7;cjzK1U<@Im(Fno><y$+8xWxp<Q<IIS5!y
z2fg;+V7<Xr0JaLyD<vP*6BB)sIqw9f*6>*A4Qh!c#07QPl%$DG)H;&zl+{4)l#-DN
zjwL0Mmj}j%?u1?U)6Nd#`iSoPp3v<$fZC>ilDAt+u;ShT`C}Qoxjm?&ving+nq~fh
z<iGebUwjEK#JG!a+B;a!xW$^jFWRMs7p~sm5Xz`0C88+z!)o@s{KqYj8%I!=w8=K^
zH+u03R0e}(q))ra_mSl@Pnn{dsWbGPk6ax#C;&_I`l12z1(7p?JN1VB&aeFi_~Tm^
zmPTpfyS8(nFUYI>cgL`nhCkIzd&9JVAUsSD+K{e_ZC$$Y-D~`36n*l6N)gjYpwEKh
z(@A4Cgt)@xA)iTP5>}7(7HOBW*h*%phHblu*QrZ~RG;0c$)Q>G!ZT0DINR~lazmuO
zmmHLim6Gzm?(OCE{=<jt($d#x@peXp1r2drLFhR6qS?f|Xj@%b`N+UPMO}UH^#WQh
z83G2lqN3t;ax#V9&;NnA1H-aoTnq(kEmRH`to&ZBMy>zRl@L*|8HEF0WEB;qg5iOn
z5#T^53o0xab-sW9{uP`9Q1ZdQGVU?=DreLV;ObbFEhG4P_Krd;uz@zy=f+UZ=VPG{
zaQKGc8L_mp%d4wP!KV_4{GcbTHU`y@HM(_W#bI4(|GVzK&jr&o8;EJ*METd>zNy8<
z#gUYQgw_fMgJVB`7i~w!PuevOGfCBJ8y<@H?%f+-$WpZ&&P4U~MU6BgIv@c@)`@3;
zT#l6#6%`eOMMk7UE%B_Us`FLttg;EOAXRzQ=cJ{hzflTIo=nS?97f0n7$87I1nTM2
zr)+@met(9)%|TZWALT?lYEbvHuK6|lKj2mX;8Iv<-R?#tFW%9D)!&kx*^;7><StO*
zNBrz%TcEBCQ!V5wWheo6K_v9qDKvPcz33tWWIg&;&>Mhv0AyF_GX}<Z(BkXiKd{(F
zio|MYFh)iV^|ATh3fC_$+(~Z<o2NKy!l>2v(xx252}#hHU<TaXc7@r6(8$itPTBz?
zE|iIfBqlvw3AiHuaGZ;nk2N&V7%|D?^1p-&i*@Ny8Kyg8yJ|I0m+3NsYb{z;?1BpW
zt3_`SQA7t#d`X6ShsIz-=LVcU`bGkP5e*Fu(klt$<Ks!b;5P!{6pc#TcIg?q3xIAs
zyS@Kae0^5(?P=ip);k3SX*KOOFGb*;<nG<OP;SxRvEnwj<E35iW?OMG>AhwK9RB`^
zAK<%eH&JN~jT_gRUV*V~_9F?-MTc5l#;S-1PF@>#_rsgVGVvsTyPFVkWMJ8=JMzc8
zNh~#(`zHgVjt{mF9i`^Il#sXSY`NLlvj9To<>dvvnt&dyMwbilv2nr$t%U21X8bjK
zOyqd7`2vJP8C_#a46)>mq0<SHEl2)%goo3Gnn-(zOJ-!Z*@=>G5d)I^k31c<YtE!*
zWMtG+l(M>N$AL0l9j|TCF)_HTkw&9*WEq<e0tr4vewob(Zvv-|K+JvOn0o|GxKnE{
zLI{2dmX={8^aoautVAg?-Ld&M5|22h&$MBAhvTfGLI^gqMPtc<TvK3D?<ar=-uc+~
z5*p=IR#pxSfv3kqEW`3bWk*Lx(kC9qN;*`YQYx{CzO1xyqbI4~p^(p#(JyZj2~T#P
ztZU|(#^}_~JW1CJb?x}rCzXLR!z?6~Pn6)c=Eq#n(hfx=B{N};=v}2c7}O4a%xz5X
z9f6656~4WIN|YNycfHqEIe}gav%&($0$4=;(Otr7Kl7*Twldh8%gKcjlf=k7RM7dI
zlKNN^iuq$CG+b?Rxs)d889hpz>;`~20h0Y($1Akt1@;es5B*|WcpydQzJ2>#j@FZj
zsQv+E_o4aAToIGEy@ZwZGu%f;j|7UEh`Ca1;StnTpmQL0yPXey2|(LPp?}RREXti0
z6kthNz5~0bKE>FN4AJ|#!Jz9%dt8LK*;fVOD?CXo5H+#11<;AIyr`yz+@z=Ux|WP|
zhQ}p|6KA%a7Am=!kI^mS5%F)r6g6>ajzwyAL!SkNbvbywa9F0TS!3^VYNMfhzRYzo
z6zBs6<9zw#ct!@^v_yioYPt~bjOQb)$dF5_yx8Y84RixsB}EYAP4wk51mlXF-^tAH
z<2DlLh)iSaIn3f|ZE>&vu62k#6(f$hKztals;{r?>3Q&a0RRJ{gcCK}6~h=lUI`7<
z{m>aX{}*ek!ALn|@W7WXdcDsB(o9i+jB+>?SQrM8GdZOzxWmUd7YSNBxdQx~@%qGm
zKySg6%*=rRUy4$;05b!7(Nqg}#uY-F6KfoaO|JdT>kfX0U6KG73*6o7>*^pGK+5=m
zAWEK%4468-qf5Qgt!a4ghRwmv9$_T<Zt+ZFh~IzyfO>h|qfs9t5&>ooK#`CKbe((e
zgipQ;8@5@yIukTMvy`<?z5eFzOfB$Y3<Ydigjc`=$Ou6+IjN7PdI53*@m^%moRK3J
ztp=u_;_AtiF1?pwuTD2_aWvB<*n~C-PKOG%xB-~py@0!n07qhxUVTDrY;025HU`vV
z!<ErWup-J1)lR>-N?d78B;0rU6FZgmh4)^o=r^M+hMJl5TTv_wB#8h;pvx&MFArs6
z#<HDiN(TZS%o@a%Ix&wPP6t-x7i$u<`JdeQcz`%QEm@yByJDM(N9gcGje#Njbtx^H
zg~-*H7X$Ji<YDG#)eN03(=8Y5nrfqt)}}tsAoF>h7w>&+KD}jEYjy{?f49&?-owe~
zm$`kKNhZx(Kir<?`xjvuQ@{fF7xv!Yzki8_z;rrv|75;K$YB~oQ@=JybhRcnBjZMf
z|LGD)G7g<Inob~m!K=?F=vf+vqJv0c_fByrMW&}FvHbV(e6r5bUsTWzI@LpJm(wi)
z{-k;={r$vCyyeBz<$p9Y15v{Q8UM4&zhf5rXq97Yv1IJJzMAMa!VEMuIs<@gAYvvD
zHM5dGLX!;m6lNQp<7uNU+IOJehN2xlTHIB-+U!quauP9HwX<h{^4(hz=_6v6G?`9R
zLAKRreZarchZ{|Qu!so@=H$Z!VOMQOGchsYbGnsD=avpufMY*NfQ}Ep+sVxGb~x;>
z7qmDq@^XUOC4*&*y+2p(fD~UpFlhBGp9U1UOL!N}#)fyf#Fpv^!w_tSO9Q|f1@7E}
z4<8tA-@clN`aId_T~JVfPU~PVa&mSiffcQR6uj6I_?+I)$v8=P>t@2c0w!oRzyA%<
z^AGg25A{xHiC-F_iAQhn;^R8Em8amy0#~sa1oZ&7uL+)Veq2Jr4UpVkLp<UwA1;^7
zMv1!NQJN4hI*7I43+me5f80Iy8{_**2zK3!6dAuC&dZ6nx=PZ2->k=OLy&;q85|D7
z_3qplV64c^U&hsU@|kc#J#$;tZENyf%<K-oL{1%roy2EBGPN&TK0!Y^d(gE0R-<1+
zXo<h3k{tN>o(GmRmP9cd$rYjd=$yYo1FjD=2FlO$^D1-W#JmeB%zKxAtFmO`i|?Lq
z1Ts{`Xt8IgW+tR?ze6x3UymlPg}$bdX4Q{WS@Z%le;@%&x6Av!$q(xq9^Q=46=88|
z$bI+jljqjZNW(!fwt}+YnFpUPc`^O5u2PdKPG%~ZvNE{cU!7{iTA9=_Guu-KWJ(ZB
zc`~J?`D}{M<hHw(!TN;Zsf34jkvc=N-(6b2zpKBG{RU025-jzViSiP~phe<j26PYk
zeU+fU=HM>D@h`?13PA>j0x9%`rAxx^M=quNhZas=aQ*!H$PGFWB8qy*Y#?do)P>UF
zTqLf17-O~eDbZ-oZhyjHM%<$9f8&vujfYysy!Mj~R)RIE?Odn&H9Pb*XY5^l&)4+G
znp|@X?h(nm;yzTYvY*O$uN2<Jb9zcRe9U9xjoBx5{n?GX-Ssm-KeI5XOSd)kX?)a`
z2iHxu04oF5MT5{C|De4S9N-p;>M8w>*!O=i18y;WCO@qC6uN!InF<h~jf}ElugBm5
zD>HtI)36u1E*UK)BXN8>Sjs@C{OzkNQSu9xox<Xrocn=#G95F$IS&1E>nn0?jm*Gp
z`($DC38AyJG)v;?BC$Nf$z{J&GPFigeGTwT;m7Ch?}bvr6iJ{Ti+zqLEr%}7eX6kZ
zE55rQ&ZK7@diJRe*yEBQEBza?R{<Tyfodba&qivz^XCVyEm^YSO8@;{J4^oyT#{s~
zoPo#Xz`RMf)8f+fPmPU<nvO_Yr@FYgsR1Dkq$b+;JxWDIg(j2@n$my}dL17x159As
zs0<6LNpqS{GL)E*S`1I}l*dLVI~TeaK@{52LLZ7mFyT_HBDNCD9Ot`Y5Vp-AP6J{G
z!JW2TjFmMMys5xO=b&`CG+n-~NRQ2J{09X+nG|(@e&pJC9jjQv{0o=;FY)VH5iMEJ
zf;mw$qak;-6O;{y@!QQIQuF}o71b_W!@CT%2r_}S)C#Dxao!lb#SK0@VkmRyEk!8o
zk_pY$LXts5Qj`YR4*3YsWwyYbXIRU5cp%Ax9|>=Wf`8PFnI#q{OYm(Imfu$jWvQRK
zxS1(eS^J@@BZ`@iD@Cx7>3h(loz2+-g13+pcH!eMfz9iZ2Gyy`B`6}tx0ysD^{W3d
z|Kl5M>_kKq$^dL37n4-7$Gma-w+myJ?nSgdo}eCqMFFg6SD2c4or4%4yVcFRH=m;x
zQXJ~Lw6>}Ce+&Yfb$oAsU3BKQX3HgI=<<o=X#hh+?M-V|kKqHRrYkpZd_u`bju#7B
zH_dmEjNPs(8Chez<I{7osu^#{w2fw{^|Qq=r1A%U<SmWMd=54r;lHTlh{(_Ide(OI
zP(`J4#SM_OG4WMY8Js8JYN;8M<6C$A@!?4>q#6#Le>CINua5G6)|<#dTmMhn$1S25
zEuCsRGf3qK6u8lq2B>wUNaBi$isYBAz;fz8zVI(lTKV%IZkG6`AIG^+Ts!4-ELJJ$
zgEE&BY;YDhp7u)0ywq7cs48vxcouZx0+U!9K-ax<8>Jo_v(h$PAIi()fRCKn@Bmf{
zXt}t|%zildkh2VqpZ^kdlWY|$d`nF^6&T~zY9&DE0=Y|wl@k!mdB^-y6^JaKZZFP<
zhwV%h%*(0laZ5jjrlZ_yjt~zdiB3;=wcK#i-a@p+W~0IBw{j<GmazcAw*yz*-9jFA
zOu1;+&~_KFqUW^h`<qfUzbMV8ZTiLI1ZVt&4O}Zg>1aMX`FA7Apsv$K90xF3+E$@m
zK<+Th6WzswzVd}D<g5m-Pq`c&UBB;DJ4xdVF?hbE5?Bpc)!+XbJQGm7iickHWPNbx
zR#Y$BtJYtw(|@s-PMAv8`pyDfx4O1-+oki0<qWQ9TAv`xUwhv;c0uUw5TFxh@+A7Q
z!HKJ`^)s6NccZTUk;I88dqS-w%|royOsl!iSTn4rcK*bZew5#@BWK;w;rM_TL^5^#
zHToJ8?&JQ<wjDV;zjKOHKGP3QJ>`^`Czt(C$1)?rl$nZc#X1Ue5rzjYA4TDRDnqp-
zrmuY)D|awAY1?0Q^tO|h1=%w%GUsG{tGT6^u=OgoVK?>uw*`ku4)im|41hYx-cN41
z6B?0oV&?A=(o1On94*7cog?FkN#v~vB?+nA0?^u<*#6Wzp0yP;oMe~-p;hf0Nl8#)
z<>_vTLwB;nklb|RQsE6?v_61eiT&MwiS0f&v8fdvK&H9~Tb&rq{2dY1Ea3{BkPeq(
zC`{p_8%L0IW#DK)<=5rXcaR7!vGzT$Jj2XPDTJ6{P>*k#WM-!k!IS*^V$HP$u(l8F
zo2C18FLQu=vw~9-1=JE(#xkDwMh89Y_W!H4Z%m)&`2n4tyEfjo1MSiraFhwS+gIQs
zohO;k!4)BJ>_OT;S4*pSS;Q8q-N?iaZgF(II3#RKJ5O2Cf8k1-^1Vy%Cuq20H9QDB
z6p&}vA!r_khf}-z``yu%X!=3utW1~~-OQp6-GH<-ydc*fVmNqzy<8G0`fX-xsBMP~
zy})4&@N9~80CHDW*5f)`Jc*B+tqp&Fc4#|~wHQ-ip;(lrsJ>1IuerC4acH?-0+B&E
z*$P@52DC={psmsG;_1?(njCRmww3Z)Rc&Yi!TTkA17?&*?hu&t(gP!pgPT6v#x=j4
z&Se4+d(jnMF$rV@i?+`Lb>1@}5%6yr?|Epby&@(_ehq@bo%-p4#F-PaAd695r@V<r
z=rx46nESS~=O9&fB-|xRRmB-_yRRJjHz?)`il&Mg1a~LwJl;}qrQo~k6Uu0#D>y?q
z2*I)o2eA2h_}2oRl>_ZLgQ6S$5!hkviJ4iP6v1FKYw0e21z3FlCPE3=^Y2x_f2sI_
z-Q5!^9%R2U8rd&LovAZzmiUwF%iK5GmKIyBScj<+;iqiX84zfze@%_hA;?~-fD+FO
z8rVobaa=XQ1Cj+OwN4w8+}+=^uNY^(ZB$o>#>J^LVFR4{7zZsrj3e+Lr*OmMdgku`
z4a4OtA~Z_pS2I~n`-1H3!n?H)G@8mqLvx^d|1pJ0-w_l9U9$&>zuBs#`T4J)t;CfA
zrGXjH@_zKc#oe|jBtb_@pu}*EsMy>hmhgJ~?HQ@W^XN2H_-Q@zCRrV!>v#Q1Fyz?#
zUXNb1-h`uXj8!UT-Wh>Zsc(<T9DWu!OsRp656GUuIFsG+HuhWK)xo`O5}DQ5DDFko
zv9$FWear+r(IHBf4Cm$eM!6oXxA6G=&HCmx@4leSaU9J|d-unn#;Y?S!0^oirvOu=
zFJfK*Z!lQ3P^Hx-4Asv<Qd^yibc74O3{m-I^m1o_fna~qbw)?<e2u$}6_{88z3+g&
z`~Lm=w{Vy#FDw={bw=bjyH>?I-J_mWg#+?_wNf&bCn#S}6vNw1OaJY*dMRnh@GrNO
zidpN>jJNkps2=!8wA@`i4%g(>H}%NC7J74m`w=W8T-2a%?}A~?MvB=|R%99bZD@#$
z&l+s!=(>B7?6G23eUx;hU-|6EYn^-Gurh&XAQ3OZ^<3dOeoOL4bUG~HrH@t&n%uv`
zK}7Va`@c<AKE{};T3Y0Mg?Z<SH@XUA{3CMwAJ(e>B>&sJ+pg=VNPb4=j&|jZQ_Cnt
zGGzgF$uBp-3Y_=nc1s98m2pQRkSb6yvr7QRL3c#q>Rq-h2yuUP^~KpSXTuuiesBKt
z@|q@ew+sbl3Zz<C!~dJpIk!x$0New;0(c8BAw!Gqr?`T}{O}Q_iJ^sx$Dg$(O`|L-
z4FlXYv5BhJ;M<g8LRH$7qNBnPEDfDDbiYRg3wCYNc9*iu1;Hwmn0a+C;oj-UdFtjx
zsIwrFxpo^dE9;Ts%w(At07y2RVg6tM{|9Ii@-ICENC9y`!MOC>;Rt&=knF{MT(F`j
z(km#&j-#+=M59|&=&*;n9Z)`CW$8&XLMJJP@cecr`QS;MA^L-R7DBk%&dff~xYESX
z`YOQdXyQ1gtwW;yfz0XpmGzEC`?_{Dcb^{0133?2a_}aJDI6(>mIKF-rDxavNn+qn
zK^G|^VZS#?!)WFtlYr1>2YUv&7F42pflTOpJJ0W{8yl0DOz-te^>TU4UQf!zooXbQ
z^hI~BVPWXgX7~-<ArU#^+tj3K_+t;jF@;9o((Pb2j>^g34MlZI{t9&O4WS#qw#xnp
z>JU*34>L<=>(n6#ly5!h(gd&C>5$Ss%tqHChEH#j6tguA0>{Q?R^?hDjZgUpn<cj3
zns7v!N4Ba}zl3?h0vmJdIdplI&Vzgbxg#?jJ3lf@Xfyx|yXIg~1l`7m5@vqsNye@{
zJuIvNTP5^?O8V>Y1_Elvv3LK|>zb^qh0gU95&z}F^`|MtK+?GXHC*ar{yQ>QG!K)U
z#}qnoL6I3`Lj3w^$TCMN$hyMLj5iqLKF$-8ak198tjriUPELl(?5A!u?N;iLkXM4k
zdQ%nNaZUKDoBxfU+|R2s1x{DS{m-}#kH|jTF0VBw9IcMsIei5tKqdq;3$DNb@#eWf
zHon?R;c~f<&pv4+f@CZPeRdm$Zu(E}-{KG+3e9j}$~#8LHnFhyE+;#hGY)^l#{Sj9
z`=&7x!_rE?%4g}=Y4@fHPqoeL-?86GgT~wdeC;ev8?Q=+DKfwbmix#!HV95{4VUL9
zbNg)=@$!*5t=7`9@Hx{P9M2+CZgIeyr2I)@DyXR77qCASNfj*IEQFPy`-70b06n9$
zY2E)j8+kX=S7#fT$_}}3`4`ke4$o8V4!O(9UV3zjN=RvF5D89f8^E_)R=>gzsfAoA
zVs1gN>E;*azhir5heyrI``s-eqN~HLed!M$7e9e_e)HWV7DI4KcdtKp{%M;$?(=_$
zpz4<;zAGA{NQiHF3I#>t>C+?kfvILTJ00MO#@(kr&QYW1^qHD5@4f$HzVXArN=$Ky
znT5nu`v#u{?ZuFXVtDZH_bodZl`L#wOkMwzCqmSsvsJ>>gbql2O}eEp^nl~c7x2FR
zkAegyDX+(p?kLW}ky>Zmqut~SScdonomf4|orT#k@Bslo(Vq)Y9&YcG<Eb_Wp8e!o
z-%U1E{x3XIa2qJb9U0=iJ{u8)amlr>o%>ICl9R6<%#zXld~5`p3ZZJesRE+Kn>TNy
z4Y;+-g3w2S&wPMHNfq>`X>JZ49euZd!U?Y3oggM==x@^31*GBN0n>Fh#05CK=+~SH
z2CAfmR>RjHKYm2muVr|z%Bz|h)xBj2>C1sIRyW2313(R&np-Nhg`E`Zp`iF2G;8--
z)B<Id5<EY#O1TrubSqir5~06v44J-Udqk-I+~@lx8(*kOwg~F6(2Ma%K@n#!vK_j!
z*(w?l7o6uh?o|%5DmN_BJc#o%J&0FHROUBHRA${<Z|au_8f^qiWd;4`v5V)`AFaBj
z!g1kL;)P$!U|%WH9)Lmfj0&iNmVi;>w=CLl<!?lJ@A2_*P)E|Q6(voPXh<x;o7epL
zQ23XC@e8k4*i{U4(+dPGg3sB!C#iKHM=rGJ=)Vu5M4P-p=UvBNU+LjutqMybpr-03
zW~A8>2l;SxInI!Wym*(Ay+&a4PYL`aucLzE<^L26&<8c#LYd%|VzfSr{(J=yNXP#^
zUXe@QJf#9BbtjAzuPBizKL@r2&Wt3Oeeyr^gVXJBHmN<J9F9POHfgARqk|j_DgUs9
zM#GtK$Tb&~0qCv!mrUTmIjnjqkk-<48sE@<nVGg=(38)SmR~JVzuFTgCHiz94#I+-
z2UYLNl`9S+>ELOD`+@RlGu0%XcPt|ocaSbZFUyjflY@K{AD;|IPs+#VHj-c#=~0a6
zBVQI4OXqqi8rH1#-#I_`)sbYTUor&pBHj_7dcfXUwY#8>EtM~M3thPT*z@?SEn3EU
zFJ<;Fjp?y|lG&C2q+R`)h3uMfPsUB>3qD<Xd(V~jF4+9UPpEx98b7?(Q{VMTIa0EP
z&X6c$N?3yRr={#B-p^yk-bj>L%mbz%n3Tu2cE;Ubi%?eJhlh21=2s$$_#BjHC6!>7
z#rE>6tKQ$$)yb5-rji1ru_14+gV6XL&dmM@Q{IuPi7H!q`Pi-nCtIrpC+!EZOn693
z)S%T_$yU4ZCyD;aDaAq30n5jxyc9+$i^T@n_m(YR3otQIvJUdg6l*Mt_TSf45B=O?
ziUU{Ysug_v9W{sZ<jGhSS%eHds!YqkU@b`{mMM!(0afK9GQr3DN;XF-VWH_vl`dj2
z@$FEe<-Z4Q^}1!=de2pf<`RJ)%?sNxAqk1qlu1keoAPR7yWkx<-kb9GI6L(nY7k6p
z<GvArb2)I*X<x9nzaOh5?vyj>-XF`aFtEd)uL`G(40(edFky~WeX6S)MNE`&FvWE@
zec`wff$`juf3!^R5vi_$0n#xo<ACawtU>2d0-l=O-av@7^wbN-9nrLSzM$@I`P@M(
z({k;!jEvd!^=oiK1fwa^xXPAlY;4RoKDpN%9%*diO-@cu6&=NKitQK%hTahNKGRgl
zy+6MsMqh{hqQ{|!U{Njn(eMOyVH!qHCx0;a$DKS0&&tFE3qB)5P9z;ZTcgXo95l(-
zOa*HMz5Gf-QGrj3e<B)(zLyw(f{8>33Jwh+n(qXDA7C4Un`R~Dg@fplu&mFOc%S12
z5-f@jq7Vsw7P*WX89b;mhy|%wCL)}RcIhv6T$@aJ*&jc-Uhw$Y+XBrI8;!0q4rOKK
z!@eV}akIA?-0lTA_s(y+L+L<?bFr>3&&n&z`p;YPq{zD>>AQ=b7Q7TC&JXhkedA5W
zs&GP)z6*sKg>e5*m6bn+KGCKcHA%xg9#zs44b<EjV*9;-jxa-rFqq|^rSS@%b+#85
zy_l*Usq^PVSSH?o7<C=ps5VXF9<X)g|AP*rd>0W;Q|;B7FpEK7eF!9N;rkGa@K&~{
zgTCJ0nXRS%{e?sW*yk|tRo%t;d8XNo{q#+WkPAp2FJ)1`KX!#VN0#>)vK`f)T*q(#
zO=`1@F8e$J0InC6l(;LM?QbwHE-j7y?1uY(<uA5w3J$)6L+|ymB&OT+?T1@|v9zmd
za)DLor+`4o@O==M&BiFcdVfb0PcZ}*(C#EA7tWPb2~Wo>TmqiBrm1N)xq57T!$*Qp
zKDHDLk_n4-tDF`m{Nt5Tmu=C55tft;ofFbKdcC;wp%!D5P0X|z37aHK_5Jco52YW^
zuFp8snCb9SeqLBG8zR`)-X@J@a!sSsXCs}iv+I<MdRiQ!R^YNnxc;iHI4UJ#_(EVc
zdU(<~kB+zcCE>o%eL)8WpDz+ay_R)EbD8CSSm&IUj?Qw*su01g_$#oDM=~BT)k#n~
zP<N(Moi6sKCR8ZAdjU6|@;fRvnKb@1T=nKe-8oU@nS}qP%l7C|Yin!Y<4RmkG9MG*
z;^K1f@DPeMd6hF26LifumstXVLZMViFVRW(w9L=vRyV2@+>GnKncC~7>%$V)?b5)R
z@PKK3$iRQ*_?UWdaB$LXtP0Z5wQy;X47o=A&hBp9KRbtE63NY^5i;$@2#Dlr`zg0^
zi?9DmTg-{5d$G9*&$B%7b@lW<mo%S_@|~SWMT_Z0SpM~mxy}SSZ!oLU*n9gHs}*<@
zM*UoNt@#)MUb?6Qex>C)1{=`(_K((@v5&4vJZ2*;EicC!tJ0D1W+Rn?-@M7~s6_6@
z73BAUx3&+Y`MKVy;f98WTJpmrVF%pB02>DqHtLk)S!hEZL0|wb6sLpBok?IgjIrn8
zp@$kZWoRe`CB3P&B?yji`2724ut?ywP5jF&@Y*~To+7-~B}o4NeBl728Tn-mP0b4!
z31;^a>iyI4k<U3$;s$Ui4dXe~x2WGKdGv{>$~UQsEZ(%5Po5mZ+WY!)4f4aZ-bC?|
z%G|pZ|6;k$H{LNW&k9m4-G@lOAJ6V|l|@EGL|}_^4VdgTJ7!`K;pgdcd8e;{voU+X
za>#e8(K~dg#K3!Ps@Xpp;$dP!Z~QpGQ9kyjJh9@S^v!D+82ps@)U>n=>B`6|TN``(
z>c3-3B^Am7f`YJP49PEZpddj*upti)gq2lPj&7j{4^NJz+gpNu4V2kZ%{Y*^B;I^?
zu%--$-gpEBml{ED2wbiMYFB1R46WICp6#yakcV6t<Zo;-{Vw6N|6GeB1VY+=tfvPL
z57`2-a)XR)ZJ}mw!ThV7boE>gb+%mYfydnhvj?$34nla^36D5##PpgY5D3STPW#FD
z+Pb<wqV^0i&-1Kks{8x;+FbP@Fr*m<Vrj#rZwB41)m?pBWjhuon*$+S#&Lhh_3^Fv
zs*}wFF?8IhK)U;OT_(VYoHIIxUkSOo%B46Ml+Ct37cDuudF$m1*rq~Jd|!o$#a&@>
zX=A_>wYDM4Mv5*}l9JdMZxomx7yj@TKV4X0#<=_FP19DB-w9>q(_-5mnJm@23^A}}
zm_Gb?HkO_;H)q0)v&=$GF!$%rm0DebJgYWGYH5a;1Ya=b7^dXy!3>5NdHn7+CDzf=
znKa2Y<qen0dJYLF!HIx?ptPwe>hZ1Mg9A?t1I#M(NmxCM7Vtiya?<M47>PC>5e(#z
z;9Z^(B7mHS5Ep+1mF9jd(=iy{7{!%~HMTHc%jN{YAm*Kff7L}<^k?koO~U6d+ze_t
zJ#E_4>MrAW5IgUB+G<$JtOB2l^X1Ex$4^!c51v(0>JGuPtkdqmWsqB;gnBRCl8+Ts
zHgj|&lETSVV=tlZbc81$BX9IMAU-=i;#PR4hT7=s?{{={4PI+HBEV2Bc&vIW{&oM7
z&L(9@0DS7ykT5aPTGBOf1q%xcEiEmK^><r#r~ST<gRZGQxei+ret-C2?CN^G?&I9_
zwBfYV+=dVbYIBl~n%dCMuZi=fqaGQ|RNcq>gZydf>8_m4=H{2x)zu%c#C0ugIp~o&
zJ39|gPTKumZ*CIB#DE06a01Usrd9J1G!Rve8)H#q;s0wcdls-T028tzu&pt79<W*T
z9q1}%;YIFCEqk}c=qb<jbDow8YywR=o$BM`laQUAt#W*#3UJ(z@o3V;>d@7pHNU>B
z)a`iTZZxwc_x3#r(bJov-bFDuY>YSpOwYh7*>3*v<^J>Y-tGaO9i%jo!)PWCaDCSC
zm1ltt-98PNn-^$s0lOM}{QTC>c1m8jvF&c*{@MMz-xc+N^Hd;kkn#BO;|%*s9st+M
ziEt(6=JH<6+L~)~#2{`@MdQ(<uJ5mED(qalIp0J|_qUF&t|JE{&~F+dx89jR%Gy7`
zguG+duB_eM7kv#Rj(q!O_Ufxu&TTiPiJ*%9SoNLOb+3W>0$9C$t&;uy`**n%^P*K*
zw;jS6R$PCrZ128kc|#-TMdeR#ZW>$kJ?_;6E>gRf)1P4Q<3T^Gn3xzaf7#8?_uj0-
z-S=2Qtb1!t|MAB+Zr?uq@1Gs;ye{uqi4q-0lYRgzu4uXGUxGA57HD*pi>pcSBmo;j
zz#u4><GC|OcJj#-;Q2^zbMyIh-nm{c<xsrE)C1h|I&q%z#o&Vn*Sqxh0E2dh(CuxE
zH7{k3KLifAcDN{cEDZuS6ptM`!~~pkDJUsvaass$hgtXS1!nXWAzJtD%0KWfcwoV&
zrmk*n_Xt=;&X_qf^ER+Z0E#{OxWjzi;}<VD*1g`w>U!5i{`{J?YZd4EF*9tsxyIb}
z=c7l9Qe|hX|NG%E|Gka+`uZC;Y-j-P4V-q<eX)i6B6Vkhi;E|Ds2q9qYL$RfN76=(
zlt{T1J;0trUncMjn53knMe9F(`nd<#`qt=D5^)vWwXPeu&Trjsb|%J;zy)g4a;K$o
zMo!zk4|w*#^-o?)CtWDLdprDkX<<>(qovui)Qo2Oc+Z<%8m=`J*pQF=2h8SytFMCE
z^w+<xduozuB)KqiR)C1Bj=nzk_0m>g7a(ZmotHA5FStE9^n5*l)zH5KhNRR+;Hr+C
fYK&NoKkWUrp5MM-x$Xw=I6wwZS3j3^P6<r_=qaHN

diff --git a/tensor2tensor/mesh_tensorflow/mtf_utils.py b/tensor2tensor/mesh_tensorflow/mtf_utils.py
deleted file mode 100644
index cd9f5cb18..000000000
--- a/tensor2tensor/mesh_tensorflow/mtf_utils.py
+++ /dev/null
@@ -1,70 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Common utilities for mesh tensorflow."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import contextlib
-import heapq
-
-import tensorflow as tf
-from tensorflow.python.framework import ops
-
-
-@contextlib.contextmanager
-def outside_all_rewrites():
-  with ops.control_dependencies(None):
-    yield
-
-
-class BalancedVariablePlacer(object):
-  """Place the variable on different device and blance the memory usage."""
-
-  def __init__(self, devices, init_usage=None):
-    init_usage = init_usage if init_usage else [0] * len(devices)
-    assert len(devices) == len(init_usage)
-    self._mem_device_heap = list(zip(init_usage, devices))
-    heapq.heapify(self._mem_device_heap)
-    self._last_device = devices[0]
-
-  def device_function(self, var):
-    """Choose a device for the input variable.
-
-    Args:
-      var: an Variable.
-
-    Returns:
-      The device for placing the var.
-    """
-    if var.type not in ('Variable', 'VariableV2', 'VarHandleOp'):
-      tf.logging.info('Place {} on last device: {}.'.format(
-          var.name, self._last_device))
-      return self._last_device
-
-    shape = tf.TensorShape(var.get_attr('shape'))
-    assert shape.num_elements() is not None
-
-    size = tf.DType(var.get_attr('dtype')).size
-    mem, device = heapq.heappop(self._mem_device_heap)
-    mem += shape.num_elements() * size
-    heapq.heappush(self._mem_device_heap, (mem, device))
-    tf.logging.info('Place variable {} on {} and consumes {} Bytes.'.format(
-        var.name, device, mem))
-    self._last_device = device
-
-    return device
diff --git a/tensor2tensor/mesh_tensorflow/mtf_utils_test.py b/tensor2tensor/mesh_tensorflow/mtf_utils_test.py
deleted file mode 100644
index d0547e924..000000000
--- a/tensor2tensor/mesh_tensorflow/mtf_utils_test.py
+++ /dev/null
@@ -1,54 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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 mtf_utils.py."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from six.moves import xrange  # pylint: disable=redefined-builtin
-
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-from tensor2tensor.mesh_tensorflow import mtf_utils
-
-import tensorflow as tf
-
-
-class MtfUtilsTest(tf.test.TestCase):
-
-  def test_variable_placer(self):
-    sizes = [100, 0, 0, 0]
-    device_list = ['cpu:0', 'cpu:1', 'cpu:2', 'cpu:3']
-
-    with tf.Graph().as_default() as g:
-      var_placer = mtf_utils.BalancedVariablePlacer(device_list, sizes)
-      graph = mtf.Graph()
-      mesh = mtf.Mesh(graph, 'my_mesh', var_placer)
-
-      hidden_dim = mtf.Dimension('hidden', 10)
-      output_dim = mtf.Dimension('output_feature', 10)
-
-      for i in xrange(5):
-        # Each variable takes 400 Bytes, and will be placed from cpu:1.
-        mtf.get_variable(mesh, 'w{}'.format(i), [hidden_dim, output_dim])
-
-      for i in xrange(5):
-        var = g.get_tensor_by_name('w{}:0'.format(i))
-        device = (i + 1) % len(device_list)
-        self.assertEqual('cpu:{}'.format(device), var.device)
-
-
-if __name__ == '__main__':
-  tf.test.main()
diff --git a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py b/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
deleted file mode 100644
index b1085c4a1..000000000
--- a/tensor2tensor/mesh_tensorflow/placement_mesh_impl.py
+++ /dev/null
@@ -1,533 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Placement Mesh Implementation (for CPU/GPU clusters)."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import functools
-from six.moves import xrange  # pylint: disable=redefined-builtin
-
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-import tensorflow as tf
-
-
-class PlacementMeshImpl(mtf.MeshImpl):
-  """Mesh implemented using explicit device placement."""
-
-  def __init__(self, shape, layout, devices):
-    super(PlacementMeshImpl, self).__init__(shape, layout)
-    self._devices = devices
-
-  class LaidOutTensor(object):
-    """One Slice for each processor."""
-
-    def __init__(self, tensor_list):
-      self._tensor_list = tensor_list
-
-    def __repr__(self):
-      return "[" + ",".join([str(t) for t in self._tensor_list]) + "]"
-
-    @property
-    def tensor_list(self):
-      return self._tensor_list
-
-    @classmethod
-    def from_tensor_list(cls, tensor_list):
-      return cls(tensor_list)
-
-    @property
-    def all_slices(self):
-      return self._tensor_list
-
-    @property
-    def slice_shape(self):
-      return self.tensor_list[0].shape.as_list()
-
-    def to_laid_out_tensor(self):
-      return self
-
-  class LaidOutVariable(object):
-    """Maintains slice-variables and copy operations."""
-
-    def __init__(self, variable, mesh_impl):
-      """Create a LaidOutVariable.
-
-      Args:
-        variable: a Variable (Operation)
-        mesh_impl: a MeshImpl
-      """
-      self._variable = variable
-      self._mesh_impl = mesh_impl
-      shape = variable.outputs[0].shape
-      dtype = variable.outputs[0].dtype
-      slice_shape = mesh_impl.slice_shape(shape)
-      base_name = variable.name
-      slices = []
-      for pnum in xrange(mesh_impl.size):
-        with tf.device(mesh_impl.devices[pnum]):
-          slices.append(tf.get_variable(
-              base_name + "_slice_%d" % pnum,
-              slice_shape,
-              dtype=dtype, collections=[]))
-      self._laid_out_tensor = mesh_impl.LaidOutTensor(slices)
-      self._copy_master_to_slices = self.assign_to_slices(
-          mesh_impl.make_slices(variable.master, shape))
-      self._copy_slices_to_master = tf.assign(
-          variable.master,
-          mesh_impl.combine_slices(self._laid_out_tensor.all_slices, shape))
-
-    def assign_to_slices(self, slices):
-      """Assign to the slice variables.
-
-      Args:
-        slices: a list of tf.Tensor
-
-      Returns:
-        a tf.operation
-      """
-      return tf.group(mtf.parallel(
-          self._mesh_impl.devices, tf.assign,
-          self.laid_out_tensor.all_slices, slices))
-
-    @property
-    def laid_out_tensor(self):
-      return self._laid_out_tensor
-
-    @property
-    def copy_master_to_slices(self):
-      return self._copy_master_to_slices
-
-    @property
-    def copy_slices_to_master(self):
-      return self._copy_slices_to_master
-
-  def slicewise(self, fn, *inputs):
-    """Execute a function in parallel on all slices.
-
-    Args:
-      fn: a function from tf.Tensors to tf.Tensor or a tuple of tf.Tensors.
-      *inputs: a list of inputs.  Each input is either a LaidOutTensor or
-        is convertible to a tf.Tensor.
-    Returns:
-      a LaidOutTensor, or a tuple of LaidOutTensors if fn returns a tuple.
-    """
-    if fn == tf.add:
-      assert len(inputs) == 2
-      if isinstance(inputs[0], mtf.LazyAllreduceSum):
-        # sum of LazyAllreduceSum (keep delaying the allreduce)
-        return inputs[0] + inputs[1]
-    # convert all inputs to LaidOutTensor where possible
-    inputs = mtf.convert_args_to_laid_out_tensors(inputs)
-    inputs = [x.tensor_list if isinstance(x, self.LaidOutTensor)
-              else [x] * len(self.devices) for x in inputs]
-    ret = mtf.parallel(self.devices, fn, *inputs)
-    if isinstance(ret[0], tuple):
-      ret = mtf.transpose_list_of_lists(ret)
-      return tuple([self.LaidOutTensor(t) for t in ret])
-    else:
-      return self.LaidOutTensor(ret)
-
-  def Print(self, x, data, message, **kwargs):  # pylint: disable=invalid-name
-    """call tf.Print.
-
-    Args:
-      x: a LaidOutTensor
-      data: a list of LaidOutTensor
-      message: a string
-      **kwargs: keyword arguments to tf.print
-    Returns:
-      a LaidOutTensor
-    """
-    tf.logging.info("PlacementMeshImpl::Print")
-    new_slices = x.tensor_list[:]
-    with tf.device(self._devices[0]):
-      new_slices[0] = tf.Print(
-          new_slices[0], [t for d in data for t in d.tensor_list],
-          message, **kwargs)
-    return self.LaidOutTensor(new_slices)
-
-  def allreduce(self, x, mesh_axes, reduction_fn_string):
-    """Grouped allreduce, (across the given dimensions).
-
-    Args:
-      x: a LaidOutTensor
-      mesh_axes: a list of integers - the mesh dimensions to be reduced
-      reduction_fn_string: "SUM" or "MAX"
-    Returns:
-      a LaidOutTensor
-    """
-    return self._collective_with_groups(
-        x, mesh_axes, functools.partial(
-            allreduce_ring, reduction_fn_string=reduction_fn_string))
-
-  def allconcat(self, x, mesh_axis, concat_axis):
-    """Grouped allconcat (like MPI allgather followed by concat).
-
-    Args:
-      x: a LaidOutTensor
-      mesh_axis: an integer - the mesh axis along which to group
-      concat_axis: an integer (the Tensor axis along which to concatenate)
-    Returns:
-      a LaidOutTensor
-    """
-    return self._collective_with_groups(
-        x, [mesh_axis],
-        functools.partial(allconcat_ring, concat_axis=concat_axis))
-
-  def alltoall(self, x, mesh_axis, split_axis, concat_axis):
-    """Grouped alltoall.
-
-    Args:
-      x: a LaidOutTensor
-      mesh_axis: an integer the mesh axis along which to group
-      split_axis: an integer (the Tensor axis along which to split)
-      concat_axis: an integer (the Tensor axis along which to concatenate)
-    Returns:
-      a LaidOutTensor
-    """
-    return self._collective_with_groups(
-        x, [mesh_axis],
-        functools.partial(
-            alltoall_ring, split_axis=split_axis, concat_axis=concat_axis))
-
-  def receive(self, x, mesh_axis, source_pcoord):
-    """Collective receive in groups.
-
-    Each group contains the processors that differ only in mesh_axis.
-
-    ```python
-    group_size = self.shape[mesh_axis].size
-    ```
-
-    Args:
-      x: a LaidOutTensor
-      mesh_axis: an integer
-      source_pcoord: a list of optional integers. Each element is either None
-        or an integer in [0, group_size). If source_pcoord[k] is None, then the
-        output for the k-th processor in each group is a zero tensor. If
-        source_pcoord[k] is not None, then the output for the k-th processor in
-        each group is equal to the input for the source_pcoord[k]-th processor
-        in that group.
-
-    Returns:
-      a LaidOutTensor
-    """
-    x = x.to_laid_out_tensor()
-    shape = x.tensor_list[0].shape
-    dtype = x.tensor_list[0].dtype
-    def _collective_receive(tensor_list, device_list):
-      ret = []
-      for pcoord, device in enumerate(device_list):
-        with tf.device(device):
-          if source_pcoord[pcoord] is None:
-            ret.append(tf.zeros(shape, dtype))
-          else:
-            ret.append(tf.identity(tensor_list[source_pcoord[pcoord]]))
-      return ret
-    return self._collective_with_groups(
-        x, [mesh_axis], _collective_receive)
-
-  def _collective_with_groups(self, x, mesh_axes, collective):
-    """Grouped collective, (across the given dimensions).
-
-    Args:
-      x: a LaidOutTensor
-      mesh_axes: a list of integers - the mesh dimensions to be reduced
-      collective: fn from list(tf.Tensor), list(device) -> list(tf.Tensor)
-    Returns:
-      a LaidOutTensor
-    """
-    if not mesh_axes:
-      return x
-    x = x.to_laid_out_tensor()
-    if len(mesh_axes) == self.ndims:
-      return self.LaidOutTensor(collective(x.tensor_list, self._devices))
-    else:
-      groups = mtf.processor_groups(self.shape, mesh_axes)
-      ret = [None] * self.size
-      for g in groups:
-        inputs = [x.tensor_list[pnum] for pnum in g]
-        devices = [self._devices[pnum] for pnum in g]
-        reduced = collective(inputs, devices)
-        for pnum, y in zip(g, reduced):
-          ret[pnum] = y
-      return self.LaidOutTensor(ret)
-
-  def random(self, shape, tf_fn, kwargs):
-    """Call a random tf operation (e.g. random_uniform).
-
-    Args:
-      shape: a Shape
-      tf_fn: a function such as tf.random_uniform
-      kwargs: kwargs to pass to tf_fn, except for seed
-
-    Returns:
-      a LaidOutTensor
-    """
-    slice_shape = self.slice_shape(shape)
-    var_scope = tf.get_variable_scope().name
-    def my_fn(pnum):
-      # seeds are necessary to make sure that slices that should have the
-      # same values actually do have the same values.
-      seed = hash("%s%s" % (var_scope, self.slice_begin(shape, pnum)))
-      return tf_fn(slice_shape, seed=seed, **kwargs)
-    return self.slicewise(my_fn, self.laid_out_pnum())
-
-  def laid_out_pnum(self):
-    """Returns a LaidOutTensor containing the processor number."""
-    return self.LaidOutTensor(list(range(self.size)))
-
-  @property
-  def devices(self):
-    return self._devices
-
-  def export_to_tf_tensor(self, x, laid_out_x):
-    """Turn a Tensor into a tf.Tensor.
-
-    Args:
-      x: a Tensor
-      laid_out_x: a LaidOutTensor
-    Returns:
-      a tf.Tensor
-    """
-    return self.combine_slices(laid_out_x.all_slices, x.shape)
-
-  def import_tf_tensor(self, x, tf_x):
-    """Import a tf.Tensor, producing a LaidOutTensor.
-
-    Args:
-      x: a Tensor
-      tf_x: a tf.Tensor
-    Returns:
-      a LaidOutTensor
-    """
-    return self.LaidOutTensor(self.make_slices(tf_x, x.shape))
-
-
-def allreduce_ring_single_shard(xs, devices, reduction_fn_string="SUM"):
-  """Compute the reduction of all Tensors and put the result everywhere.
-
-  Performance-optimized for a ring of devices.
-
-  Args:
-    xs: a list of n tf.Tensors
-    devices: a list of strings
-    reduction_fn_string: "SUM" or "MAX"
-
-  Returns:
-    a list of n Tensors
-  Raises:
-    ValueError: if devices is not a list of n strings
-  """
-  n = len(xs)
-  binary_reduction = mtf.binary_reduction_fn(reduction_fn_string)
-  assert len(devices) == n, "devices must be a list of length len(xs)"
-  if n == 1:
-    return xs
-  result = [None] * n
-  if n % 2 == 0:
-    left_center = n // 2 - 1
-    right_center = left_center + 1
-  else:
-    left_center = n // 2
-    right_center = left_center
-  left_sum = xs[0]
-  for i in xrange(1, left_center + 1):
-    with tf.device(devices[i]):
-      left_sum = binary_reduction(left_sum, xs[i])
-  right_sum = xs[n-1]
-  for i in reversed(xrange(left_center + 1, n - 1)):
-    with tf.device(devices[i]):
-      right_sum = binary_reduction(xs[i], right_sum)
-  with tf.device(devices[left_center]):
-    result[left_center] = binary_reduction(left_sum, right_sum)
-  if n % 2 == 0:
-    with tf.device(devices[right_center]):
-      result[right_center] = binary_reduction(left_sum, right_sum)
-  for i in reversed(xrange(left_center)):
-    with tf.device(devices[i]):
-      result[i] = tf.identity(result[i + 1])
-  for i in xrange(right_center + 1, n):
-    with tf.device(devices[i]):
-      result[i] = tf.identity(result[i - 1])
-  return result
-
-
-def allreduce_ring(xs, devices, reduction_fn_string="SUM"):
-  """Compute the reduction of all Tensors and put the result everywhere.
-
-  Performance-optimized for a ring of devices.
-
-  Args:
-    xs: a list of n tf.Tensors
-    devices: a list of strings
-    reduction_fn_string: "SUM" or "MAX"
-
-  Returns:
-    a list of n Tensors
-  Raises:
-    ValueError: if devices is not a list of n strings
-  """
-  n = len(xs)
-  if len(devices) != n:
-    raise ValueError("devices must be a list of length len(xs)")
-  if n == 1:
-    return xs
-  shape = xs[0].shape.as_list()
-  # tf.logging.info("allreduce_ring shape = %s" % shape)
-  size = None if None in shape else mtf.list_product(shape)
-  if size is None or size < 1024 or size % n != 0:
-    return allreduce_ring_single_shard(xs, devices, reduction_fn_string)
-
-  def _circular_shift(l, n):
-    n %= len(l)
-    return l[-n:] + l[:-n]
-  def _flatten_and_split(x):
-    return tf.split(tf.reshape(x, [size]), n)
-  def _concat_and_reshape(xs):
-    return tf.reshape(tf.concat(xs, 0), shape)
-
-  # [device, shard]
-  x_split = mtf.parallel(devices, _flatten_and_split, xs)
-  x_split_t = mtf.transpose_list_of_lists(x_split)
-
-  y_split_t = []
-  for shard in xrange(n):
-    shard_xs = _circular_shift(x_split_t[shard], shard)
-    shard_devices = _circular_shift(devices, shard)
-    shard_ys = allreduce_ring_single_shard(
-        shard_xs, shard_devices, reduction_fn_string)
-    y_split_t.append(_circular_shift(shard_ys, -shard))
-  y_split = mtf.transpose_list_of_lists(y_split_t)
-  ys = mtf.parallel(devices, _concat_and_reshape, y_split)
-  return ys
-
-
-def allconcat_ring(xs, devices, concat_axis):
-  """Concatenate all Tensors everywhere.
-
-  Performance-optimized for a ring of devices.
-
-  Args:
-    xs: a list of n tf.Tensors
-    devices: a list of n strings
-    concat_axis: an integer
-
-  Returns:
-    a list of n Tensors
-  """
-  n = len(xs)
-  if n == 1:
-    return xs
-  # [target, source]
-  parts = [[xs[target] if target == source else None for source in xrange(n)]
-           for target in xrange(n)]
-  for distance in xrange(1, n // 2 + 1):
-    for target in xrange(n):
-      source = (target + distance) % n
-      if parts[target][source] is None:
-        with tf.device(devices[target]):
-          parts[target][source] = tf.identity(parts[(target + 1) % n][source])
-      source = (target - distance) % n
-      if parts[target][source] is None:
-        with tf.device(devices[target]):
-          parts[target][source] = tf.identity(parts[(target - 1) % n][source])
-  return mtf.parallel(devices, tf.concat, parts, axis=[concat_axis] * n)
-
-
-def alltoall_pointtwise(xs, devices, split_axis, concat_axis):
-  """MPI alltoall operation.
-
-  Implementation of alltoall using pointwise communication.
-
-  Args:
-    xs: a list of n tf.Tensors
-    devices: a list of n strings
-    split_axis: an integer
-    concat_axis: an integer
-
-  Returns:
-    a list of n Tensors
-  """
-  n = len(xs)
-  if n == 1:
-    return xs
-  # [target, source]
-  parts = mtf.transpose_list_of_lists(
-      mtf.parallel(devices, tf.split, xs, [n] * n, axis=[split_axis] * n))
-  return mtf.parallel(devices, tf.concat, parts, axis=[concat_axis] * n)
-
-
-def alltoall_ring(xs, devices, split_axis, concat_axis):
-  """MPI alltoall operation.
-
-  Performance-optimized for a ring of devices.
-
-  Args:
-    xs: a list of n tf.Tensors
-    devices: a list of n strings
-    split_axis: an integer
-    concat_axis: an integer
-
-  Returns:
-    a list of n Tensors
-  """
-  n = len(xs)
-  if n == 1:
-    return xs
-  # set up
-  # [target, source]
-  parts = [[None] * n for i in xrange(n)]
-  def my_split(x, size_splits):
-    total_size = tf.shape(x)[split_axis]
-    part_size = total_size // sum(size_splits)
-    return tf.split(x, [s * part_size for s in size_splits], axis=split_axis)
-  forward_message_size = (n - 1) // 2
-  backward_message_size = (n - 1) - forward_message_size
-  forward_messages = [None] * n
-  backward_messages = [None] * n
-  for i in xrange(n):
-    with tf.device(devices[i]):
-      if i >= backward_message_size:
-        a, b, c, d = my_split(
-            xs[i], [i - backward_message_size,
-                    backward_message_size, 1, n - i - 1])
-        backward_messages[i] = b
-        parts[i][i] = c
-        forward_messages[i] = tf.concat([d, a], axis=split_axis)
-      else:
-        a, b, c, d = my_split(
-            xs[i], [i, 1, forward_message_size, backward_message_size - i])
-        backward_messages[i] = tf.concat([d, a], axis=split_axis)
-        parts[i][i] = b
-        forward_messages[i] = c
-  for step in xrange(1, max(forward_message_size, backward_message_size) + 1):
-    new_forward_messages = [None] * n
-    new_backward_messages = [None] * n
-    for i in xrange(n):
-      with tf.device(devices[i]):
-        if forward_message_size > 0:
-          parts[i][(i - step) % n], new_forward_messages[i] = my_split(
-              forward_messages[(i - 1) % n], [1, forward_message_size - 1])
-        if backward_message_size > 0:
-          new_backward_messages[i], parts[i][(i + step) % n] = my_split(
-              backward_messages[(i + 1) % n], [backward_message_size - 1, 1])
-    forward_message_size -= 1
-    backward_message_size -= 1
-    forward_messages = new_forward_messages
-    backward_messages = new_backward_messages
-  return mtf.parallel(devices, tf.concat, parts, axis=[concat_axis] * n)
diff --git a/tensor2tensor/mesh_tensorflow/research/__init__.py b/tensor2tensor/mesh_tensorflow/research/__init__.py
deleted file mode 100644
index 4bd418a74..000000000
--- a/tensor2tensor/mesh_tensorflow/research/__init__.py
+++ /dev/null
@@ -1,15 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
diff --git a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py b/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
deleted file mode 100644
index 1eb3fd233..000000000
--- a/tensor2tensor/mesh_tensorflow/simd_mesh_impl.py
+++ /dev/null
@@ -1,407 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""SIMD Mesh implementation (for TPU/XLA)."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from six.moves import xrange  # pylint: disable=redefined-builtin
-
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-from tensor2tensor.mesh_tensorflow import mtf_utils
-from tensor2tensor.mesh_tensorflow import tpu_variables
-
-import tensorflow as tf
-
-from tensorflow.contrib.tpu.python.ops import tpu_ops
-from tensorflow.python.framework import ops
-
-
-class SimdMeshImpl(mtf.MeshImpl):
-  """Mesh implementation for TPU using SIMD and MPI operations."""
-
-  def __init__(self, shape, layout, devices, device_assignment):
-    super(SimdMeshImpl, self).__init__(shape, layout)
-    self._devices = devices
-    self._device_assignment = device_assignment
-    tf.logging.info("SimdMeshImpl init: {0} {1}".format(shape, layout))
-    self._pnum_tensor = None
-
-  @property
-  def pnum_tensor(self):
-    if self._pnum_tensor is not None:
-      return self._pnum_tensor
-    with mtf_utils.outside_all_rewrites():
-      tf.logging.info("Create pnum_tensor")
-      self._pnum_tensor = tpu_ops.tpu_replicated_input(
-          list(range(self.size)), name="pnum_constants")
-      return self._pnum_tensor
-
-  class LaidOutTensor(object):
-    """One Slice."""
-
-    def __init__(self, tensor_list):
-      assert isinstance(tensor_list, list)
-      self._tensor_list = tensor_list
-
-    def __repr__(self):
-      return "[" + ",".join([str(t) for t in self._tensor_list]) + "]"
-
-    @property
-    def tensor_list(self):
-      return self._tensor_list
-
-    @property
-    def one_slice(self):
-      return self._tensor_list[0]
-
-    @classmethod
-    def from_tensor_list(cls, tensor_list):
-      return cls(tensor_list)
-
-    @property
-    def all_slices(self):
-      return self._tensor_list
-
-    @property
-    def slice_shape(self):
-      return self.one_slice.shape.as_list()
-
-    def to_laid_out_tensor(self):
-      return self
-
-  class LaidOutVariable(object):
-    """Maintains slice-variables and copy operations."""
-
-    def __init__(self, variable, mesh_impl):
-      """Create a LaidOutVariable.
-
-      Args:
-        variable: a Variable (Operation)
-        mesh_impl: a MeshImpl
-      """
-      self._variable = variable
-      self._mesh_impl = mesh_impl
-      shape = variable.outputs[0].shape
-      dtype = variable.outputs[0].dtype
-      slice_shape = mesh_impl.slice_shape(shape)
-      base_name = variable.name
-      slices = []
-      for pnum in xrange(mesh_impl.size):
-        slice_var_name = base_name + "_slice_%d" % pnum
-        tpu_device = mesh_impl.device_assignment.tpu_device(replica=pnum)
-        # The initializer is unimportant, since the slice variables will be
-        # overwritten.  zeros_initializer() is here to avoid the default
-        # initialization which adds lots of useless operations to the TF graph.
-        with ops.device(tpu_device):
-          slices.append(
-              tf.get_variable(
-                  slice_var_name,
-                  slice_shape,
-                  dtype=dtype,
-                  collections=[],
-                  initializer=tf.zeros_initializer()))
-      self._laid_out_tensor = mesh_impl.LaidOutTensor(
-          [tpu_variables.ReplicatedVariable(base_name, slices)])
-      with tf.device(variable.master.device), mtf_utils.outside_all_rewrites():
-        self._copy_master_to_slices = self.assign_to_slices(
-            mesh_impl.make_slices(variable.master, shape),
-            assign_to_tensor_list=slices)
-        self._copy_slices_to_master = tf.assign(
-            variable.master,
-            mesh_impl.combine_slices(slices, shape,
-                                     device=variable.master.device))
-
-    def assign_to_slices(self, slice_values, assign_to_tensor_list=None):
-      """Assign to the slice variables.
-
-      Args:
-        slice_values: a list of tf.Tensor
-        assign_to_tensor_list: an optional list of tf.Variable
-
-      Returns:
-        a tf.operation
-      """
-      if assign_to_tensor_list is None:
-        assign_to_tensor_list = self._laid_out_tensor.all_slices
-      # Handle both N -> 1 and N -> N cases.
-      num_slices = min(
-          len(assign_to_tensor_list), len(slice_values))
-      devices = [""] * num_slices
-      return tf.group(
-          mtf.parallel(devices, tf.assign, assign_to_tensor_list[:num_slices],
-                       slice_values[:num_slices]))
-
-    @property
-    def laid_out_tensor(self):
-      return self._laid_out_tensor
-
-    @property
-    def copy_master_to_slices(self):
-      return self._copy_master_to_slices
-
-    @property
-    def copy_slices_to_master(self):
-      return self._copy_slices_to_master
-
-  def laid_out_pnum(self):
-    """Returns a LaidOutTensor containing the processor number.
-
-    Returns:
-      a LaidOutTensor where each slice is an integer scalar
-    """
-    return self.LaidOutTensor([self.pnum_tensor])
-
-  def _create_group_assignment(self, mesh_axes):
-    """Create group assignment for XLA cross replica ops."""
-
-    partitioning = {}
-    for pnum in xrange(self.size):
-      group = mtf.pnum_to_group(self.shape, mesh_axes, pnum)
-      if group not in partitioning:
-        partitioning[group] = []
-      partitioning[group].append(pnum)
-    group_assignment = []
-    for group, pnums in partitioning.items():
-      group_assignment.append(pnums)
-    return group_assignment
-
-  def allreduce(self, x, mesh_axes, reduction_fn_string):
-    """Grouped allreduce, (summed across the given dimensions).
-
-    Args:
-      x: a LaidOutTensor
-      mesh_axes: a list of integers
-      reduction_fn_string: "SUM"
-    Returns:
-      a LaidOutTensor
-    Raises:
-      ValueError: if the reduction is not yet implemented.
-    """
-    if not mesh_axes:
-      return x
-    x = x.to_laid_out_tensor()
-    if reduction_fn_string == "SUM":
-      group_assignment = self._create_group_assignment(mesh_axes)
-      return self.LaidOutTensor(
-          [tpu_ops.cross_replica_sum(x.one_slice, group_assignment)])
-    else:
-      for axis in mesh_axes:
-        x = self.allconcat(x, axis, 0, stack=True)
-        x = self.LaidOutTensor(
-            [mtf.reduction_fn(reduction_fn_string)(x.one_slice, 0)])
-      return x
-
-  def allconcat(self, x, mesh_axis, concat_axis, stack=False):
-    """Grouped allconcat (like MPI allgather followed by concat).
-
-    TODO(noam): inefficient - replace with a XLA allconcat when available
-
-    Args:
-      x: a LaidOutTensor
-      mesh_axis: an integer - the mesh axis along which to group
-      concat_axis: an integer (the Tensor axis along which to concatenate)
-      stack: a boolean - whether to stack instead of concat
-    Returns:
-      a LaidOutTensor
-    """
-    x = x.to_laid_out_tensor()
-    coord = self.laid_out_pcoord(mesh_axis)
-    t = x.one_slice
-    old_shape = t.shape.as_list()
-    num_parts = self.shape[mesh_axis].size
-    t = tf.expand_dims(t, concat_axis)
-    t *= tf.reshape(
-        tf.one_hot(coord.one_slice, num_parts, dtype=t.dtype),
-        [num_parts if i == concat_axis else 1
-         for i in xrange(len(old_shape) + 1)])
-    if not stack:
-      new_shape = old_shape[:]
-      new_shape[concat_axis] *= num_parts
-      t = tf.reshape(t, new_shape)
-    return self.allreduce(self.LaidOutTensor([t]), [mesh_axis], "SUM")
-
-  def alltoall(self, x, mesh_axis, split_axis, concat_axis):
-    """Grouped alltoall (like MPI alltoall with splitting and concatenation).
-
-    Args:
-      x: a LaidOutTensor
-      mesh_axis: an integer the mesh axis along which to group
-      split_axis: an integer (the Tensor axis along which to split)
-      concat_axis: an integer (the Tensor axis along which to concatenate)
-    Returns:
-      a LaidOutTensor
-    """
-    x = x.to_laid_out_tensor()
-    t = x.one_slice
-    group_assignment = self._create_group_assignment([mesh_axis])
-    t = tpu_ops.all_to_all(
-        t,
-        concat_dimension=concat_axis,
-        split_dimension=split_axis,
-        split_count=len(group_assignment[0]),
-        group_assignment=group_assignment)
-    x = self.LaidOutTensor([t])
-    return x
-
-  def receive(self, x, mesh_axis, source_pcoord):
-    """Collective receive in groups.
-
-    Each group contains the processors that differ only in mesh_axis.
-
-    ```python
-    group_size = self.shape[mesh_axis].size
-    ```
-
-    Args:
-      x: a LaidOutTensor
-      mesh_axis: an integer
-      source_pcoord: a list of optional integers. Each element is either None
-        or an integer in [0, group_size). If source_pcoord[k] is None, then the
-        output for the k-th processor in each group is a zero tensor. If
-        source_pcoord[k] is not None, then the output for the k-th processor in
-        each group is equal to the input for the source_pcoord[k]-th processor
-        in that group.
-
-    Returns:
-      a LaidOutTensor
-    """
-    x = x.to_laid_out_tensor()
-    t = x.one_slice
-    source_target_pairs = []
-
-    for pnum in xrange(self.size):
-      coord = self.pnum_to_processor_coordinates(self.shape, pnum)
-      k = coord[mesh_axis]
-      if source_pcoord[k] is not None:
-        coord[mesh_axis] = source_pcoord[k]
-        target_pnum = self.processor_coordinates_to_pnum(coord)
-        source_target_pairs.append([pnum, target_pnum])
-
-    return tpu_ops.collective_permute(t, source_target_pairs)
-
-  def slice(self, tf_tensor, tensor_shape):
-    """"Slice out the correspoding part of tensor given the pnum variable."""
-    tensor_layout = self.tensor_layout(tensor_shape)
-
-    if tensor_layout.is_fully_replicated:
-      return self.LaidOutTensor([tf_tensor])
-    else:
-      slice_shape = self.slice_shape(tensor_shape)
-      slice_begins = [
-          self.slice_begin(tensor_shape, pnum) for pnum in xrange(self.size)
-      ]
-      slice_begins_tensor = tf.stack(slice_begins)
-      # slice on source device
-      selected_slice_begin = tf.gather(slice_begins_tensor, self.pnum_tensor)
-      return self.LaidOutTensor(
-          [tf.slice(tf_tensor, selected_slice_begin, slice_shape)])
-
-  def slicewise(self, fn, *inputs):
-    """Execute a function in parallel on all slices.
-
-    Args:
-      fn: a function from tf.Tensors to tf.Tensor or a tuple of tf.Tensors.
-      *inputs: a list of inputs.  Each input is either a LaidOutTensor or
-        is convertible to a tf.Tensor.
-    Returns:
-      a LaidOutTensor, or a tuple of LaidOutTensors if fn returns a tuple.
-    """
-    if fn == tf.add:
-      assert len(inputs) == 2
-      if isinstance(inputs[0], mtf.LazyAllreduceSum):
-        # sum of LazyAllreduceSum (keep delaying the allreduce)
-        return inputs[0] + inputs[1]
-    # convert all inputs to LaidOutTensor where possible
-    inputs = mtf.convert_args_to_laid_out_tensors(inputs)
-    ret = fn(*[x.one_slice if isinstance(x, self.LaidOutTensor)
-               else x for x in inputs])
-    if isinstance(ret, tuple):
-      return tuple([self.LaidOutTensor([t]) for t in ret])
-    else:
-      return self.LaidOutTensor([ret])
-
-  @property
-  def device_assignment(self):
-    return self._device_assignment
-
-  @property
-  def devices(self):
-    return self._devices
-
-  def random(self, shape, tf_fn, kwargs):
-    """Call a random tf operation (e.g. random_uniform).
-
-    Args:
-      shape: a Shape
-      tf_fn: a function such as tf.random_uniform
-      kwargs: kwargs to pass to tf_fn, except for seed
-
-    Returns:
-      a LaidOutTensor
-    """
-    # TODO(noam): can we make things better with stateless_random?
-    slice_shape = self.slice_shape(shape)
-    x = tf_fn(slice_shape, **kwargs)
-    # TPU does not have seeds enabled.  Sync up the
-    # random choices by zeroing out all but the first core per group of
-    # identical slices, then allreducing by group.
-    layout = self.tensor_layout(shape)
-    # we need to sync across these axes.
-    mesh_axes = [i for i in xrange(self.ndims)
-                 if i not in layout.tensor_axis_to_mesh_axis]
-    multiplier = 1.0
-    for axis in mesh_axes:
-      multiplier *= tf.cast(
-          tf.equal(self.laid_out_pcoord(axis).one_slice, 0), x.dtype)
-    x *= multiplier
-    x = self.LaidOutTensor([x])
-    x = self.allreduce(x, mesh_axes, "SUM")
-    return x
-
-  def export_to_tf_tensor(self, x, laid_out_x):
-    """Turn a Tensor into a tf.Tensor.
-
-    Args:
-      x: a Tensor
-      laid_out_x: a LaidOutTensor
-    Returns:
-      a tf.Tensor
-    """
-    tensor_layout = self.tensor_layout(x.shape)
-    if not tensor_layout.is_fully_replicated:
-      raise NotImplementedError(
-          "SimdMeshImpl only supports export_to_tf_tensor of fully-replicated "
-          "Tensors.  Try reshaping to new dimension names. "
-          " x.shape = %s tensor_layout=%s"
-          % (x.shape, tensor_layout))
-    return laid_out_x.one_slice
-
-  def import_tf_tensor(self, x, tf_x):
-    """Import a tf.Tensor, producing a LaidOutTensor.
-
-    Args:
-      x: a Tensor
-      tf_x: a tf.Tensor
-    Returns:
-      a LaidOutTensor
-    """
-    return self.slice(tf_x, x.shape)
-
-  @property
-  def supports_control_dependencies(self):
-    return False
diff --git a/tensor2tensor/mesh_tensorflow/tpu_variables.py b/tensor2tensor/mesh_tensorflow/tpu_variables.py
deleted file mode 100644
index bb852a569..000000000
--- a/tensor2tensor/mesh_tensorflow/tpu_variables.py
+++ /dev/null
@@ -1,200 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Distributed variable implementation for TPUs."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import contextlib
-
-from tensorflow.python.framework import ops
-from tensorflow.python.ops import control_flow_ops
-from tensorflow.python.ops import gen_resource_variable_ops
-
-
-@contextlib.contextmanager
-def _handle_graph(handle):
-  with handle.graph.as_default():
-    yield
-
-
-def _enclosing_tpu_context():
-  # pylint: disable=protected-access
-  context = ops.get_default_graph()._get_control_flow_context()
-  # pylint: enable=protected-access
-  while context is not None and not isinstance(
-      context, control_flow_ops.XLAControlFlowContext):
-    context = context.outer_context
-  return context
-
-
-class ReplicatedVariable(object):
-  """A replicated variable for use on TPUs.
-
-  When accessed inside a tpu.replicate() context, this variable acts as if it
-  is a single variable whose handle is a replicated input to the computation.
-
-  Outside a tpu.replicate() context currently this object has pretty murky
-  semantics, especially with respect to things such as
-  * initialization
-  * colocation.
-
-  TODO(phawkins): merge this with the TPU DistributionStrategy code.
-  """
-
-  def __init__(self, name, variables):
-    self._name = name
-    self._primary_var = variables[0]
-    self._vars = variables
-    self._cached_value = None
-    self._dtype = variables[0].dtype
-
-  @property
-  def handle(self):
-    tpu_context = _enclosing_tpu_context()
-    if tpu_context is None:
-      return self._primary_var.handle
-
-    return tpu_context.get_replicated_var_handle(self._name, self._vars)
-
-  @contextlib.contextmanager
-  def _assign_dependencies(self):
-    """Makes assignments depend on the cached value, if any.
-
-    This prevents undefined behavior with reads not ordered wrt writes.
-
-    Yields:
-      None.
-    """
-    if self._cached_value is not None:
-      with ops.control_dependencies([self._cached_value]):
-        yield
-    else:
-      yield
-
-  @property
-  def initializer(self):
-    return control_flow_ops.group([v.initializer for v in self._vars])
-
-  @property
-  def graph(self):
-    return self._primary_var.graph
-
-  @property
-  def _shared_name(self):
-    return self._common_name
-
-  @property
-  def _unique_id(self):
-    return self._primary_var._unique_id  # pylint: disable=protected-access
-
-  @property
-  def name(self):
-    return self._name
-
-  @property
-  def dtype(self):
-    return self._primary_var.dtype
-
-  @property
-  def shape(self):
-    return self._primary_var.shape
-
-  def get_shape(self):
-    return self._primary_var.get_shape()
-
-  def to_proto(self, export_scope=None):
-    return self._primary_var.to_proto(export_scope=export_scope)
-
-  @property
-  def constraint(self):
-    return None
-
-  @property
-  def op(self):
-    return self.get().op
-
-  def _read_variable_op(self):
-    if _enclosing_tpu_context() is None:
-      return self._primary_var.read_value()
-    v = gen_resource_variable_ops.read_variable_op(self.handle, self._dtype)
-    return v
-
-  def read_value(self):
-    return self._read_variable_op()
-
-  def assign(self, value, use_locking=None, name=None, read_value=False):
-    del use_locking
-    with _handle_graph(self.handle), self._assign_dependencies():
-      value_tensor = ops.convert_to_tensor(value, dtype=self.dtype)
-      assign_op = gen_resource_variable_ops.assign_variable_op(
-          self.handle, value_tensor, name=name)
-    if read_value:
-      return self._read_variable_op()
-    return assign_op
-
-  def assign_add(self, delta, use_locking=None, name=None, read_value=True):
-    del use_locking
-    with _handle_graph(self.handle), self._assign_dependencies():
-      assign_add_op = gen_resource_variable_ops.assign_add_variable_op(
-          self.handle,
-          ops.convert_to_tensor(delta, dtype=self.dtype),
-          name=name)
-    if read_value:
-      return self._read_variable_op()
-    return assign_add_op
-
-  def assign_sub(self, delta, use_locking=None, name=None, read_value=True):
-    del use_locking
-    with _handle_graph(self.handle), self._assign_dependencies():
-      assign_sub_op = gen_resource_variable_ops.assign_sub_variable_op(
-          self.handle,
-          ops.convert_to_tensor(delta, dtype=self.dtype),
-          name=name)
-    if read_value:
-      return self._read_variable_op()
-    return assign_sub_op
-
-  def get(self):
-    return self._primary_var
-
-  def _should_act_as_resource_variable(self):
-    """Pass resource_variable_ops.is_resource_variable check."""
-    pass
-
-  def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
-    """Converts a variable to a tensor."""
-    # pylint: disable=protected-access
-    if _enclosing_tpu_context() is None:
-      return self._primary_var._dense_var_to_tensor(dtype, name, as_ref)
-    # pylint: enable=protected-access
-    if dtype is not None and dtype != self.dtype:
-      return NotImplemented
-    if as_ref:
-      return self.handle
-    else:
-      return self.read_value()
-
-
-# Register a conversion function which reads the value of the variable,
-# allowing instances of the class to be used as tensors.
-def _tensor_conversion(var, dtype=None, name=None, as_ref=False):
-  return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref)  # pylint: disable=protected-access
-
-
-ops.register_tensor_conversion_function(ReplicatedVariable, _tensor_conversion)
-ops.register_dense_tensor_like_type(ReplicatedVariable)
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 34409a8e4..f806d0b05 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -20,16 +20,15 @@
 # pylint: disable=unused-import
 
 from tensor2tensor.layers import modalities  # pylint: disable=g-import-not-at-top
-from tensor2tensor.mesh_tensorflow import mtf_image_transformer
-from tensor2tensor.mesh_tensorflow import mtf_resnet
-from tensor2tensor.mesh_tensorflow import mtf_transformer
-from tensor2tensor.mesh_tensorflow.research import experiments_moe
 from tensor2tensor.models import basic
 from tensor2tensor.models import bytenet
 from tensor2tensor.models import distillation
 from tensor2tensor.models import image_transformer
 from tensor2tensor.models import image_transformer_2d
 from tensor2tensor.models import lstm
+from tensor2tensor.models import mtf_image_transformer
+from tensor2tensor.models import mtf_resnet
+from tensor2tensor.models import mtf_transformer
 from tensor2tensor.models import neural_gpu
 from tensor2tensor.models import resnet
 from tensor2tensor.models import revnet
@@ -48,6 +47,7 @@
 from tensor2tensor.models.research import gene_expression
 from tensor2tensor.models.research import glow
 from tensor2tensor.models.research import lm_experiments
+from tensor2tensor.models.research import moe_experiments
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research import similarity_transformer
 from tensor2tensor.models.research import super_lm
diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py b/tensor2tensor/models/mtf_image_transformer.py
similarity index 95%
rename from tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
rename to tensor2tensor/models/mtf_image_transformer.py
index 7bc999d4a..7317b7c8e 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_image_transformer.py
+++ b/tensor2tensor/models/mtf_image_transformer.py
@@ -24,12 +24,13 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
 import copy
+import mesh_tensorflow as mtf
+
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-from tensor2tensor.mesh_tensorflow import mtf_layers
-from tensor2tensor.mesh_tensorflow import mtf_model
+from tensor2tensor.utils import mtf_model
 from tensor2tensor.utils import registry
 import tensorflow as tf
 
@@ -188,7 +189,7 @@ def layer_prepostprocess_dropout(x):
       inputs = import_to_batch_by_length(inputs, "inputs")
 
       # Input embeddings
-      inputs_embedding_var = mtf_layers.embedding(
+      inputs_embedding_var = mtf.layers.embedding(
           mesh, "input_embedding",
           mtf.Shape([self.inputs_vocab_dim, self.model_dim]),
           activation_dtype=activation_dtype)
@@ -203,27 +204,27 @@ def layer_prepostprocess_dropout(x):
       with tf.variable_scope(layer_name):
         # Self attention layer
         x += layer_prepostprocess_dropout(
-            mtf_layers.masked_local_attention_1d(
-                mtf_layers.layer_norm(x, self.model_dim, name="layer_norm_att"),
+            mtf.layers.masked_local_attention_1d(
+                mtf.layers.layer_norm(x, self.model_dim, name="layer_norm_att"),
                 None,
                 self.kv_dim,
                 self.heads_dim,
                 block_length=hparams.block_length,
                 name="self_att"))
         # ffn layer
-        x += layer_prepostprocess_dropout(mtf_layers.dense_relu_dense(
-            mtf_layers.layer_norm(x, self.model_dim, name="layer_norm_ffn"),
+        x += layer_prepostprocess_dropout(mtf.layers.dense_relu_dense(
+            mtf.layers.layer_norm(x, self.model_dim, name="layer_norm_ffn"),
             self.feedforward_dim,
             hparams.dropout,
             dropout_broadcast_dims=[self.length_dim]))
 
-    x = mtf_layers.layer_norm(x, self.model_dim, name="final_layer_norm")
+    x = mtf.layers.layer_norm(x, self.model_dim, name="final_layer_norm")
 
     # Calculate the logits and loss.
-    logits = mtf_layers.dense(x, self.outputs_vocab_dim, name="logits")
+    logits = mtf.layers.dense(x, self.outputs_vocab_dim, name="logits")
     soft_targets = mtf.one_hot(
         targets, self.outputs_vocab_dim, dtype=activation_dtype)
-    loss = mtf_layers.softmax_cross_entropy_with_logits(
+    loss = mtf.layers.softmax_cross_entropy_with_logits(
         logits, soft_targets, self.outputs_vocab_dim)
     loss = mtf.reduce_mean(loss)
     for l in extra_losses:
diff --git a/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py b/tensor2tensor/models/mtf_image_transformer_test.py
similarity index 94%
rename from tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
rename to tensor2tensor/models/mtf_image_transformer_test.py
index 1411ebd36..bd4aae02e 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_image_transformer_test.py
+++ b/tensor2tensor/models/mtf_image_transformer_test.py
@@ -19,12 +19,11 @@
 from __future__ import division
 from __future__ import print_function
 
-import numpy as np
+import mesh_tensorflow as mtf
 
+import numpy as np
 from tensor2tensor.data_generators import problem_hparams
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-from tensor2tensor.mesh_tensorflow import mtf_image_transformer
-from tensor2tensor.mesh_tensorflow import placement_mesh_impl
+from tensor2tensor.models import mtf_image_transformer
 
 import tensorflow as tf
 
@@ -65,7 +64,7 @@ def get_placement_mesh(hparams):
   mesh_shape = mtf.convert_to_shape(hparams.mesh_shape)
 
   mesh_devices = [""] * mesh_shape.size
-  mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+  mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
       mesh_shape, hparams.layout, mesh_devices)
   return mesh, mesh_impl
 
diff --git a/tensor2tensor/mesh_tensorflow/mtf_resnet.py b/tensor2tensor/models/mtf_resnet.py
similarity index 97%
rename from tensor2tensor/mesh_tensorflow/mtf_resnet.py
rename to tensor2tensor/models/mtf_resnet.py
index ed7320820..e8eef3e75 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_resnet.py
+++ b/tensor2tensor/models/mtf_resnet.py
@@ -23,11 +23,10 @@
 from __future__ import print_function
 
 import copy
+import mesh_tensorflow as mtf
 
 from tensor2tensor.layers import common_hparams
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-from tensor2tensor.mesh_tensorflow import mtf_layers
-from tensor2tensor.mesh_tensorflow import mtf_model
+from tensor2tensor.utils import mtf_model
 from tensor2tensor.utils import registry
 import tensorflow as tf
 
@@ -38,7 +37,7 @@
 
 def batch_norm_relu(inputs, is_training, relu=True):
   """Block of batch norm and relu."""
-  inputs = mtf_layers.batch_norm(
+  inputs = mtf.layers.batch_norm(
       inputs,
       is_training,
       BATCH_NORM_DECAY,
@@ -305,7 +304,7 @@ def mtf_model_fn(self, features, mesh):
 
     # Calculate the logits and loss.
     out = x
-    outputs = mtf_layers.dense(
+    outputs = mtf.layers.dense(
         out, hidden_dim,
         reduced_dims=out.shape.dims[-5:],
         activation=mtf.relu, name="dense")
@@ -315,9 +314,9 @@ def mtf_model_fn(self, features, mesh):
     labels = mtf.import_tf_tensor(
         mesh, tf.reshape(labels, [hparams.batch_size]), mtf.Shape([batch_dim]))
 
-    logits = mtf_layers.dense(outputs, classes_dim, name="logits")
+    logits = mtf.layers.dense(outputs, classes_dim, name="logits")
     soft_targets = mtf.one_hot(labels, classes_dim, dtype=activation_dtype)
-    loss = mtf_layers.softmax_cross_entropy_with_logits(
+    loss = mtf.layers.softmax_cross_entropy_with_logits(
         logits, soft_targets, classes_dim)
 
     # Reshape logits so it doesn't break inside t2t.
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
similarity index 96%
rename from tensor2tensor/mesh_tensorflow/mtf_transformer.py
rename to tensor2tensor/models/mtf_transformer.py
index f889be591..411553d10 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -20,17 +20,15 @@
 from __future__ import print_function
 
 import copy
+import mesh_tensorflow as mtf
 
 from six.moves import xrange  # pylint: disable=redefined-builtin
-
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-from tensor2tensor.mesh_tensorflow import mtf_beam_search
-from tensor2tensor.mesh_tensorflow import mtf_layers
-from tensor2tensor.mesh_tensorflow import mtf_model
-from tensor2tensor.mesh_tensorflow.research import moe
+from tensor2tensor.models.research import moe
+from tensor2tensor.utils import mtf_model
 from tensor2tensor.utils import registry
+
 import tensorflow as tf
 
 
@@ -175,13 +173,13 @@ def pad_to_max_length(x):
           features["targets_position"], "targets_position",
           mesh, hparams)
       decoder_self_attention_mask = (
-          mtf_layers.attention_mask_autoregressive(
+          mtf.layers.attention_mask_autoregressive(
               targets_position, dtype=self.activation_dtype) +
-          mtf_layers.attention_mask_same_segment(
+          mtf.layers.attention_mask_same_segment(
               targets_segmentation, dtype=self.activation_dtype))
     else:
       targets_position = mtf.range(mesh, self.length_dim, dtype=tf.int32)
-      decoder_self_attention_mask = mtf_layers.attention_mask_autoregressive(
+      decoder_self_attention_mask = mtf.layers.attention_mask_autoregressive(
           targets_position, dtype=self.activation_dtype)
 
     def layer_prepostprocess_dropout(x):
@@ -207,16 +205,16 @@ def layer_prepostprocess_dropout(x):
             features["inputs_position"], "inputs_position",
             mesh, hparams)
         encoder_self_attention_mask = (
-            mtf_layers.attention_mask_same_segment(
+            mtf.layers.attention_mask_same_segment(
                 inputs_segmentation, dtype=self.activation_dtype))
         encoder_decoder_attention_mask = (
-            mtf_layers.attention_mask_same_segment(
+            mtf.layers.attention_mask_same_segment(
                 targets_segmentation, inputs_segmentation,
                 dtype=self.activation_dtype))
       else:
         inputs_position = mtf.range(mesh, self.length_dim, dtype=tf.int32)
         encoder_self_attention_mask = (
-            mtf_layers.attention_mask_ignore_padding(
+            mtf.layers.attention_mask_ignore_padding(
                 inputs, dtype=self.activation_dtype))
         encoder_decoder_attention_mask = encoder_self_attention_mask
 
@@ -257,9 +255,9 @@ def layer_prepostprocess_dropout(x):
     soft_targets = mtf.one_hot(
         targets, self.targets_vocab_dim, on_value=on_value, off_value=off_value,
         dtype=self.activation_dtype)
-    loss = mtf_layers.softmax_cross_entropy_with_logits(
+    loss = mtf.layers.softmax_cross_entropy_with_logits(
         logits, soft_targets, self.targets_vocab_dim)
-    weights = mtf_layers.weights_nonzero(
+    weights = mtf.layers.weights_nonzero(
         targets, dtype=self.activation_dtype)
     loss = mtf.reduce_mean(loss * weights)
     for l in extra_losses:
@@ -299,7 +297,7 @@ def _feedforward_layer(self, x, losses=None):
     hparams = self._hparams
     feedforward_layer = hparams.feedforward_layer
     if feedforward_layer == "dense_relu_dense":
-      return mtf_layers.dense_relu_dense(
+      return mtf.layers.dense_relu_dense(
           x, self.feedforward_dim, dropout=hparams.relu_dropout,
           dropout_broadcast_dims=[self.length_dim])
     elif feedforward_layer == "moe":
@@ -371,7 +369,7 @@ def normalize(x):
       with tf.variable_scope("layer_%d" % layer):
         # Self attention layer
         x += layer_prepostprocess_dropout(
-            mtf_layers.multihead_attention(
+            mtf.layers.multihead_attention(
                 normalize(x), None,
                 self_attention_mask, self.kv_dim, self.heads_dim,
                 dropout=hparams.attention_dropout,
@@ -380,7 +378,7 @@ def normalize(x):
         if encoder_output is not None:
           # Encoder-Decoder attention layer
           x += layer_prepostprocess_dropout(
-              mtf_layers.multihead_attention(
+              mtf.layers.multihead_attention(
                   normalize(x), encoder_output,
                   encdec_attention_mask, self.kv_dim, self.heads_dim,
                   dropout=hparams.attention_dropout,
@@ -418,7 +416,7 @@ def _sample(self, features, mesh):
            mtf.reshape(positional_embedding_var,
                        mtf.Shape([self.length_dim, self.model_dim])))
       encoder_attention_mask = (
-          mtf_layers.attention_mask_ignore_padding(
+          mtf.layers.attention_mask_ignore_padding(
               inputs, dtype=self.activation_dtype))
       with tf.variable_scope("encoder"):
         x = self._layer_stack(x,
@@ -429,7 +427,7 @@ def _sample(self, features, mesh):
       encdec_tensors = []
       for layer_num in xrange(hparams.num_decoder_layers):
         with tf.variable_scope("decoder/layer_%d/encdec_attention" % layer_num):
-          q_var, k_var, v_var, o_var = mtf_layers.multihead_attention_vars(
+          q_var, k_var, v_var, o_var = mtf.layers.multihead_attention_vars(
               mesh, self.heads_dim, self.model_dim,
               self.kv_dim, self.activation_dtype)
           k = mtf.einsum(
@@ -504,7 +502,7 @@ def logits_fn(step_num, ids, states):
     if hparams.beam_size == 1:
       temperature = (0.0 if hparams.sampling_method == "argmax"
                      else hparams.sampling_temp)
-      return mtf_beam_search.greedy_decode(
+      return mtf.beam_search.greedy_decode(
           logits_fn,
           initial_ids,
           temperature=temperature,
@@ -522,7 +520,7 @@ def logits_fn(step_num, ids, states):
             + hparams.decode_length_constant, tf.int32)
       else:
         decode_length = None
-      beams, unused_scores = mtf_beam_search.beam_search(
+      beams, unused_scores = mtf.beam_search.beam_search(
           logits_fn,
           initial_ids,
           hparams.alpha,
@@ -593,7 +591,7 @@ def normalize(x):
     for layer in range(num_layers):
       with tf.variable_scope("layer_%d" % layer):
         # Self attention layer
-        y, new_k, new_v = mtf_layers.multihead_self_attention_incremental(
+        y, new_k, new_v = mtf.layers.multihead_self_attention_incremental(
             normalize(x),
             prev_k=self_attention_k[layer],
             prev_v=self_attention_v[layer],
@@ -605,7 +603,7 @@ def normalize(x):
         if encdec_tensors is not None:
           # Encoder-Decoder attention layer
           q_var, o_var, k, v = encdec_tensors[layer]
-          x += mtf_layers.multihead_encdec_attention_incremental(
+          x += mtf.layers.multihead_encdec_attention_incremental(
               normalize(x),
               q_var, o_var, k, v,
               encdec_attention_mask,
diff --git a/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py b/tensor2tensor/models/mtf_transformer_test.py
similarity index 95%
rename from tensor2tensor/mesh_tensorflow/mtf_transformer_test.py
rename to tensor2tensor/models/mtf_transformer_test.py
index 4b00b978e..e2e685dbb 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_transformer_test.py
+++ b/tensor2tensor/models/mtf_transformer_test.py
@@ -19,12 +19,11 @@
 from __future__ import division
 from __future__ import print_function
 
+import mesh_tensorflow as mtf
 import numpy as np
 
 from tensor2tensor.data_generators import problem_hparams
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-from tensor2tensor.mesh_tensorflow import mtf_transformer
-from tensor2tensor.mesh_tensorflow import placement_mesh_impl
+from tensor2tensor.models import mtf_transformer
 
 import tensorflow as tf
 
@@ -69,7 +68,7 @@ def get_placement_mesh(hparams):
   mesh_shape = mtf.convert_to_shape(hparams.mesh_shape)
 
   mesh_devices = [""] * mesh_shape.size
-  mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+  mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
       mesh_shape, hparams.layout, mesh_devices)
   return mesh, mesh_impl
 
diff --git a/tensor2tensor/mesh_tensorflow/research/moe.py b/tensor2tensor/models/research/moe.py
similarity index 98%
rename from tensor2tensor/mesh_tensorflow/research/moe.py
rename to tensor2tensor/models/research/moe.py
index 5761292ea..e39611cdc 100644
--- a/tensor2tensor/mesh_tensorflow/research/moe.py
+++ b/tensor2tensor/models/research/moe.py
@@ -23,9 +23,7 @@
 from __future__ import division
 from __future__ import print_function
 
-
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-from tensor2tensor.mesh_tensorflow import mtf_layers
+import mesh_tensorflow as mtf
 import tensorflow as tf
 
 
@@ -137,10 +135,10 @@ def transformer_moe_layer_v1(inputs, output_dim, hparams, train):
       [experts_dim, batch_dim_unsplit, expert_capacity_dim, input_dim]))
 
   # Now feed the expert inputs through the experts.
-  h = mtf_layers.dense(
+  h = mtf.layers.dense(
       expert_inputs, hidden_dim, expert_dims=[experts_dim],
       activation=mtf.relu, use_bias=False, name="x0")
-  expert_output = mtf_layers.dense(
+  expert_output = mtf.layers.dense(
       h, output_dim, expert_dims=[experts_dim], use_bias=False, name="x1")
 
   expert_output = mtf.reshape(expert_output, mtf.Shape(
@@ -370,10 +368,10 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
       [y0, x1, h, d, m]))
 
   # Now feed the expert inputs through the experts.
-  hidden_output = mtf_layers.dense(
+  hidden_output = mtf.layers.dense(
       expert_inputs_y, hidden_dim, expert_dims=[y0, x1],
       activation=mtf.relu, use_bias=False, name="expert0")
-  expert_output = mtf_layers.dense(
+  expert_output = mtf.layers.dense(
       hidden_output, output_dim, expert_dims=[y0, x1],
       use_bias=False, name="expert1")
 
@@ -471,7 +469,7 @@ def _top_2_gating(
   """
   group_size_dim, unused_input_dim = inputs.shape.dims[-2:]
 
-  raw_gates = mtf.softmax(mtf_layers.dense(
+  raw_gates = mtf.softmax(mtf.layers.dense(
       inputs, experts_dim, use_bias=False,
       expert_dims=outer_expert_dims), experts_dim)
 
diff --git a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py b/tensor2tensor/models/research/moe_experiments.py
similarity index 98%
rename from tensor2tensor/mesh_tensorflow/research/experiments_moe.py
rename to tensor2tensor/models/research/moe_experiments.py
index 915250335..ae8f9b344 100644
--- a/tensor2tensor/mesh_tensorflow/research/experiments_moe.py
+++ b/tensor2tensor/models/research/moe_experiments.py
@@ -20,8 +20,8 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.mesh_tensorflow import mtf_transformer
-from tensor2tensor.mesh_tensorflow.research import moe
+from tensor2tensor.models import mtf_transformer
+from tensor2tensor.models.research import moe
 from tensor2tensor.utils import registry
 
 
diff --git a/tensor2tensor/mesh_tensorflow/mtf_model.py b/tensor2tensor/utils/mtf_model.py
similarity index 94%
rename from tensor2tensor/mesh_tensorflow/mtf_model.py
rename to tensor2tensor/utils/mtf_model.py
index 829287148..193107d33 100644
--- a/tensor2tensor/mesh_tensorflow/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -22,14 +22,9 @@
 
 import collections
 import copy
-import six
-
+import mesh_tensorflow as mtf
 
-from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
-from tensor2tensor.mesh_tensorflow import mtf_optimize
-from tensor2tensor.mesh_tensorflow import mtf_utils
-from tensor2tensor.mesh_tensorflow import placement_mesh_impl
-from tensor2tensor.mesh_tensorflow import simd_mesh_impl
+import six
 from tensor2tensor.utils import learning_rate
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import t2t_model
@@ -86,10 +81,10 @@ def estimator_model_fn(cls,
       # Worker 0 caches all the TPU binaries.
       worker0_mem = replica_cache_size * ctx.num_replicas
       devices_memeory_usage = [worker0_mem] + [0] * (num_hosts - 1)
-      var_placer = mtf_utils.BalancedVariablePlacer(device_list,
+      var_placer = mtf.utils.BalancedVariablePlacer(device_list,
                                                     devices_memeory_usage)
       mesh_devices = [""] * mesh_shape.size
-      mesh_impl = simd_mesh_impl.SimdMeshImpl(
+      mesh_impl = mtf.simd_mesh_impl.SimdMeshImpl(
           mesh_shape, layout_rules, mesh_devices, ctx.device_assignment)
     else:
       var_placer = None
@@ -98,7 +93,7 @@ def estimator_model_fn(cls,
       else:
         assert len(data_parallelism.ps_devices) == mesh_shape.size
         mesh_devices = data_parallelism.ps_devices
-      mesh_impl = placement_mesh_impl.PlacementMeshImpl(
+      mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
           mesh_shape, layout_rules, mesh_devices)
 
     graph = mtf.Graph()
@@ -118,7 +113,7 @@ def estimator_model_fn(cls,
       lr = learning_rate.learning_rate_schedule(hparams)
       mtf_lr = mtf.import_tf_tensor(
           mesh, tf.convert_to_tensor(lr, dtype=tf.float32), mtf.Shape([]))
-      optimizer = mtf_optimize.make_optimizer(hparams, mtf_lr)
+      optimizer = mtf.optimize.make_optimizer(hparams, mtf_lr)
       update_ops = []
       for grad, var in zip(var_grads, graph.trainable_variables):
         update_ops.extend(optimizer.apply_grad(grad, var))
@@ -136,7 +131,7 @@ def estimator_model_fn(cls,
       # tf.logging.info("tf_update_ops: {}".format(tf_update_ops))
       train_op = tf.group(tf_update_ops)
 
-    with mtf_utils.outside_all_rewrites():
+    with mtf.utils.outside_all_rewrites():
       # Copy master variables to slices. Must be called first.
       restore_hook = mtf.MtfRestoreHook(lowering)
       saver = tf.train.Saver(
@@ -183,7 +178,7 @@ def estimator_spec_eval(
 
     if use_tpu:
       def metric_fn(tf_logits, labels):
-        with tf.device("cpu:0"), mtf_utils.outside_all_rewrites():
+        with tf.device("cpu:0"), mtf.utils.outside_all_rewrites():
           eval_metrics = {}
           for metric_name, metric_fn in six.iteritems(eval_metrics_fns):
             if metric_name.split("/")[-1] not in t2t_model.TPU_METRIC_BLACKLIST:

From 58d2b65ad4af409409c344861a1225a97a0d360c Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 15 Oct 2018 20:26:42 -0700
Subject: [PATCH 1015/2720] Merging SV2P with Base.

PiperOrigin-RevId: 217255439
---
 tensor2tensor/models/video/emily.py       |   2 +-
 tensor2tensor/models/video/epva.py        |   2 +-
 tensor2tensor/models/video/savp.py        |   2 +-
 tensor2tensor/models/video/sv2p.py        | 146 ++++++++++++++--------
 tensor2tensor/models/video/sv2p_test.py   |  16 ++-
 tensor2tensor/models/video/tests_utils.py |  32 +++--
 6 files changed, 125 insertions(+), 75 deletions(-)

diff --git a/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index 2a54d50ae..875a6ea50 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -41,7 +41,7 @@
 
 
 @registry.register_model
-class NextFrameEmily(sv2p.NextFrameSv2p):
+class NextFrameEmily(sv2p.NextFrameSv2pLegacy):
   """Stochastic Variational Video Prediction Without Learned Prior."""
 
   def encoder(self, inputs, nout):
diff --git a/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index 468910b93..6bcd66dc3 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -626,7 +626,7 @@ def calc_loss_psnr(gen_images, images, name, hparams=None, use_l1_loss=False):
 
 
 @registry.register_model
-class NextFrameEpva(sv2p.NextFrameSv2p):
+class NextFrameEpva(sv2p.NextFrameSv2pLegacy):
   """Hierarchical Long-term Video Prediction without Supervision"""
 
   def body(self, features):
diff --git a/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index bf5b9cef7..559a66450 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -37,7 +37,7 @@
 
 
 @registry.register_model
-class NextFrameSAVP(sv2p.NextFrameSv2p):
+class NextFrameSAVP(sv2p.NextFrameSv2pLegacy):
   """Stochastic Adversarial Video Prediction."""
 
   def encoder(self, inputs, n_layers=3):
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index db18b4d81..121e1650d 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -28,8 +28,8 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 
-from tensor2tensor.models.video import basic_stochastic
-from tensor2tensor.models.video import sv2p_params  # pylint: disable=unused-import
+from tensor2tensor.models.video import base
+from tensor2tensor.models.video import base_vae
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -39,40 +39,17 @@
 
 
 @registry.register_model
-class NextFrameSv2p(basic_stochastic.NextFrameBasicStochastic):
-  """Stochastic Variational Video Prediction."""
+class NextFrameSv2p(base.NextFrameBase, base_vae.NextFrameBaseVae):
+  """Stochastic Variational Video Prediction From Basic Model!"""
+
+  @property
+  def is_recurrent_model(self):
+    return True
 
   def tinyify(self, array):
     return common_video.tinyify(
         array, self.hparams.tiny_mode, self.hparams.small_mode)
 
-  def visualize_predictions(self, real_frames, gen_frames):
-    def concat_on_y_axis(x):
-      x = tf.unstack(x, axis=1)
-      x = tf.concat(x, axis=1)
-      return x
-
-    frames_gd = common_video.swap_time_and_batch_axes(real_frames)
-    frames_pd = common_video.swap_time_and_batch_axes(gen_frames)
-
-    if self.is_per_pixel_softmax:
-      frames_pd_shape = common_layers.shape_list(frames_pd)
-      frames_pd = tf.reshape(frames_pd, [-1, 256])
-      frames_pd = tf.to_float(tf.argmax(frames_pd, axis=-1))
-      frames_pd = tf.reshape(frames_pd, frames_pd_shape[:-1] + [3])
-
-    frames_gd = concat_on_y_axis(frames_gd)
-    frames_pd = concat_on_y_axis(frames_pd)
-    side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
-    tf.summary.image("full_video", side_by_side_video)
-
-  def get_input_if_exists(self, features, key, batch_size, num_frames):
-    if key in features:
-      x = features[key]
-    else:
-      x = tf.zeros((batch_size, num_frames, 1, self.hparams.hidden_size))
-    return common_video.swap_time_and_batch_axes(x)
-
   def bottom_part_tower(self, input_image, input_reward, action, latent,
                         lstm_state, lstm_size, conv_size, concat_latent=False):
     """The bottom part of predictive towers.
@@ -177,9 +154,10 @@ def reward_prediction(self, *args, **kwargs):
 
   def reward_prediction_basic(self, input_images, input_reward, action, latent):
     del input_reward, action, latent
-    x = input_images[0]
-    x = tf.expand_dims(  # Add a fake channels dim.
-        tf.reduce_mean(x, axis=[1, 2], keepdims=True), axis=3)
+    x = input_images
+    x = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
+    x = tfl.dense(x, 128, activation=tf.nn.relu, name="reward_pred")
+    x = tf.expand_dims(x, axis=3)
     return x
 
   def reward_prediction_big(self, input_images, input_reward, action, latent):
@@ -212,7 +190,24 @@ def reward_prediction_big(self, input_images, input_reward, action, latent):
       x = tfcl.layer_norm(x)
       x = tfl.conv2d(x, conv_size[3], [3, 3], strides=(2, 2),
                      activation=tf.nn.relu, name="reward_conv3")
-      return x
+
+  def get_extra_loss(self,
+                     latent_means=None, latent_stds=None,
+                     true_frames=None, gen_frames=None):
+    """Losses in addition to the default modality losses."""
+    del true_frames
+    del gen_frames
+    kl_loss = 0.0
+    if self.is_training and self.hparams.stochastic_model:
+      for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
+        kl_loss += common_layers.kl_divergence(mean, std)
+        tf.summary.histogram("posterior_mean_%d" % i, mean)
+        tf.summary.histogram("posterior_std_%d" % i, std)
+      tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
+
+    beta = self.get_beta(kl_loss)
+    extra_loss = beta * kl_loss
+    return extra_loss
 
   def construct_predictive_tower(
       self, input_image, input_reward, action, lstm_state, latent,
@@ -336,12 +331,73 @@ def construct_predictive_tower(
         layer = layer[:, :img_height, :img_width, :]
         output += layer * mask
 
+      # Map to softmax digits
       if self.is_per_pixel_softmax:
         output = tf.layers.dense(
             output, self.hparams.problem.num_channels * 256, name="logits")
 
       return output, lstm_state
 
+  def video_features(
+      self, all_frames, all_actions, all_rewards, all_raw_frames):
+    """Video wide latent."""
+    del all_actions, all_rewards, all_raw_frames
+    mean, std = self.construct_latent_tower(all_frames, time_axis=0)
+    latent = common_video.get_gaussian_tensor(mean, std)
+    return [latent, mean, std]
+
+  def next_frame(self, frames, actions, rewards, target_frame,
+                 internal_states, video_features):
+    del target_frame
+    latent, latent_mean, latent_std = video_features
+
+    extra_loss = 0.0
+    if internal_states is None:
+      internal_states = [None] * (5 if self.hparams.small_mode else 7)
+      extra_loss = self.get_extra_loss([latent_mean], [latent_std])
+
+    pred_image, internal_states = self.construct_predictive_tower(
+        frames, None, actions, internal_states, latent)
+
+    if not self.has_rewards:
+      return pred_image, None, extra_loss, internal_states
+
+    pred_reward = self.reward_prediction(
+        pred_image, actions, rewards, latent)
+    return pred_image, pred_reward, extra_loss, internal_states
+
+
+@registry.register_model
+class NextFrameSv2pLegacy(NextFrameSv2p):
+  """Old SV2P code. Only for legacy reasons."""
+
+  def visualize_predictions(self, real_frames, gen_frames):
+    def concat_on_y_axis(x):
+      x = tf.unstack(x, axis=1)
+      x = tf.concat(x, axis=1)
+      return x
+
+    frames_gd = common_video.swap_time_and_batch_axes(real_frames)
+    frames_pd = common_video.swap_time_and_batch_axes(gen_frames)
+
+    if self.is_per_pixel_softmax:
+      frames_pd_shape = common_layers.shape_list(frames_pd)
+      frames_pd = tf.reshape(frames_pd, [-1, 256])
+      frames_pd = tf.to_float(tf.argmax(frames_pd, axis=-1))
+      frames_pd = tf.reshape(frames_pd, frames_pd_shape[:-1] + [3])
+
+    frames_gd = concat_on_y_axis(frames_gd)
+    frames_pd = concat_on_y_axis(frames_pd)
+    side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
+    tf.summary.image("full_video", side_by_side_video)
+
+  def get_input_if_exists(self, features, key, batch_size, num_frames):
+    if key in features:
+      x = features[key]
+    else:
+      x = tf.zeros((batch_size, num_frames, 1, self.hparams.hidden_size))
+    return common_video.swap_time_and_batch_axes(x)
+
   def construct_model(self,
                       images,
                       actions,
@@ -448,24 +504,6 @@ def process_single_frame(prev_outputs, inputs):
     else:
       return gen_images, gen_rewards, None, None
 
-  def get_extra_loss(self,
-                     latent_means=None, latent_stds=None,
-                     true_frames=None, gen_frames=None):
-    """Losses in addition to the default modality losses."""
-    del true_frames
-    del gen_frames
-    kl_loss = 0.0
-    if self.is_training and self.hparams.stochastic_model:
-      for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
-        kl_loss += common_layers.kl_divergence(mean, std)
-        tf.summary.histogram("posterior_mean_%d" % i, mean)
-        tf.summary.histogram("posterior_std_%d" % i, std)
-      tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
-
-    beta = self.get_beta(kl_loss)
-    extra_loss = beta * kl_loss
-    return extra_loss
-
   def infer(self, features, *args, **kwargs):
     """Produce predictions from the model by running it."""
     del args, kwargs
@@ -577,7 +615,7 @@ def body(self, features):
 
 
 @registry.register_model
-class NextFrameSv2pTwoFrames(NextFrameSv2p):
+class NextFrameSv2pTwoFrames(NextFrameSv2pLegacy):
   """Stochastic next-frame model with 2 frames posterior."""
 
   def construct_model(self, images, actions, rewards):
diff --git a/tensor2tensor/models/video/sv2p_test.py b/tensor2tensor/models/video/sv2p_test.py
index 98b262a69..e261f2cfa 100644
--- a/tensor2tensor/models/video/sv2p_test.py
+++ b/tensor2tensor/models/video/sv2p_test.py
@@ -32,13 +32,15 @@ def testSv2p(self):
     self.TestOnVariousInputOutputSizes(
         sv2p_params.next_frame_sv2p(),
         sv2p.NextFrameSv2p,
-        1)
+        1,
+        False)
 
   def testSv2pWithActions(self):
     self.TestWithActions(
         sv2p_params.next_frame_sv2p(),
         sv2p.NextFrameSv2p,
-        1)
+        1,
+        False)
 
   def testSv2pWithActionsAndRewards(self):
     hp = sv2p_params.next_frame_sv2p()
@@ -46,7 +48,8 @@ def testSv2pWithActionsAndRewards(self):
     self.TestWithActionAndRewards(
         hp,
         sv2p.NextFrameSv2p,
-        1)
+        1,
+        False)
 
   def testSv2pWithActionsAndRewardsExternalLoss(self):
     hp = sv2p_params.next_frame_sv2p()
@@ -54,13 +57,16 @@ def testSv2pWithActionsAndRewardsExternalLoss(self):
     self.TestWithActionAndRewards(
         hp,
         sv2p.NextFrameSv2p,
-        1)
+        1,
+        False)
 
   def testSv2pTwoFrames(self):
     self.TestOnVariousInputOutputSizes(
         sv2p_params.next_frame_sv2p(),
         sv2p.NextFrameSv2pTwoFrames,
-        1)
+        1,
+        False)
+
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/models/video/tests_utils.py b/tensor2tensor/models/video/tests_utils.py
index 99d6030ab..0c02418af 100644
--- a/tensor2tensor/models/video/tests_utils.py
+++ b/tensor2tensor/models/video/tests_utils.py
@@ -43,11 +43,11 @@ def fill_hparams(hparams, in_frames, out_frames):
 def action_modalities(hparams):
   hparams.problem_hparams.input_modality = {
       "inputs": modalities.VideoModalityL2Raw(hparams, 256),
-      "input_action": modalities.SymbolModalityOneHot(hparams, 5),
+      "input_action": modalities.SymbolModality(hparams, 5),
   }
   hparams.problem_hparams.target_modality = {
       "targets": modalities.VideoModalityL2Raw(hparams, 256),
-      "target_action": modalities.SymbolModalityOneHot(hparams, 5),
+      "target_action": modalities.SymbolModality(hparams, 5),
   }
   return hparams
 
@@ -56,13 +56,13 @@ def full_modalities(hparams):
   """Full modalities with actions and rewards."""
   hparams.problem_hparams.input_modality = {
       "inputs": modalities.VideoModalityL2Raw(hparams, 256),
-      "input_reward": modalities.SymbolModalityOneHot(hparams, 3),
-      "input_action": modalities.SymbolModalityOneHot(hparams, 5),
+      "input_reward": modalities.SymbolModality(hparams, 3),
+      "input_action": modalities.SymbolModality(hparams, 5),
   }
   hparams.problem_hparams.target_modality = {
       "targets": modalities.VideoModalityL2Raw(hparams, 256),
-      "target_reward": modalities.SymbolModalityOneHot(hparams, 3),
-      "target_action": modalities.SymbolModalityOneHot(hparams, 5),
+      "target_reward": modalities.SymbolModality(hparams, 3),
+      "target_action": modalities.SymbolModality(hparams, 5),
   }
   hparams.force_full_predict = True
   return hparams
@@ -214,7 +214,8 @@ def TestVideoModelWithActionAndRewards(self,
     self.assertEqual(output.shape, expected_shape)
 
     output, targets = res["target_reward"], features["target_reward"]
-    expected_shape = get_tensor_shape(targets)[:2] + (3,)
+    # Assuming Symbol Modality
+    expected_shape = get_tensor_shape(targets)[:2] + (1, 1, 1, 1, 3,)
     self.assertEqual(output.shape, expected_shape)
 
   def TestVideoModelWithActionAndRewardsInfer(self,
@@ -253,17 +254,22 @@ def TestOnVariousInputOutputSizes(
       test_func(4, 1, hparams, model, expected_last_dim)
       test_func(7, 5, hparams, model, expected_last_dim)
 
-  def TestWithActions(self, hparams, model, expected_last_dim):
-    for test_func in [self.TestVideoModelWithActions,
-                      self.TestVideoModelWithActionsInfer]:
+  def TestWithActions(self, hparams, model, expected_last_dim, test_infer=True):
+    test_funcs = [self.TestVideoModelWithActions]
+    if test_infer:
+      test_funcs += [self.TestVideoModelWithActionsInfer]
+    for test_func in test_funcs:
       test_func(1, 1, hparams, model, expected_last_dim)
       test_func(1, 6, hparams, model, expected_last_dim)
       test_func(4, 1, hparams, model, expected_last_dim)
       test_func(7, 5, hparams, model, expected_last_dim)
 
-  def TestWithActionAndRewards(self, hparams, model, expected_last_dim):
-    for test_func in [self.TestVideoModelWithActionAndRewards,
-                      self.TestVideoModelWithActionAndRewardsInfer]:
+  def TestWithActionAndRewards(
+      self, hparams, model, expected_last_dim, test_infer=True):
+    test_funcs = [self.TestVideoModelWithActionAndRewards]
+    if test_infer:
+      test_funcs += [self.TestVideoModelWithActionAndRewardsInfer]
+    for test_func in test_funcs:
       test_func(1, 1, hparams, model, expected_last_dim)
       test_func(1, 6, hparams, model, expected_last_dim)
       test_func(4, 1, hparams, model, expected_last_dim)

From d35e0a16d6af5dafdd2de36fadd223e954798791 Mon Sep 17 00:00:00 2001
From: Smit Hinsu <hinsu@google.com>
Date: Mon, 15 Oct 2018 22:58:28 -0700
Subject: [PATCH 1016/2720] Support DistributionStrategy in Tensor2Tensor for
 multi-GPU

PiperOrigin-RevId: 217266979
---
 tensor2tensor/bin/t2t_trainer.py              |  9 ++++
 tensor2tensor/data_generators/problem.py      | 15 +++---
 tensor2tensor/models/image_transformer.py     |  4 +-
 tensor2tensor/models/research/aligned.py      |  4 +-
 .../models/research/attention_lm_moe.py       |  4 +-
 .../models/research/transformer_moe.py        |  4 +-
 tensor2tensor/rl/trainer_model_based_new.py   |  2 +-
 .../utils/checkpoint_compatibility_test.py    |  2 +-
 tensor2tensor/utils/t2t_model.py              | 13 +++--
 tensor2tensor/utils/trainer_lib.py            | 48 +++++++++++++------
 tensor2tensor/utils/trainer_lib_test.py       |  8 +++-
 11 files changed, 77 insertions(+), 36 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 2939fc4ad..9b6f0b25e 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -64,6 +64,13 @@
 flags.DEFINE_integer("intra_op_parallelism_threads", 0,
                      "Number of intra_op_parallelism_threads to use for CPU. "
                      "See TensorFlow config.proto for details.")
+# TODO(hinsu): Enable DistributionStrategy by default once performance gap
+# between DistributionStrategy and Parallelism is resolved.
+flags.DEFINE_bool(
+    "optionally_use_dist_strat", False,
+    "Whether to use TensorFlow DistributionStrategy instead of explicitly "
+    "replicating the model. DistributionStrategy is used only if the "
+    "model replication configuration is supported by the DistributionStrategy.")
 
 # To maintain compatibility with some internal libs, we guard against these flag
 # definitions possibly erroring. Apologies for the ugliness.
@@ -215,6 +222,7 @@ def create_run_config(hp, output_dir=None):
       hp.activation_dtype == "float32" and
       hp.weight_dtype == "float32")
   return trainer_lib.create_run_config(
+      model_name=FLAGS.model,
       model_dir=output_dir or os.path.expanduser(FLAGS.output_dir),
       master=FLAGS.master,
       iterations_per_loop=FLAGS.iterations_per_loop,
@@ -234,6 +242,7 @@ def create_run_config(hp, output_dir=None):
       use_tpu_estimator=FLAGS.use_tpu_estimator,
       schedule=FLAGS.schedule,
       no_data_parallelism=hp.no_data_parallelism,
+      optionally_use_dist_strat=FLAGS.optionally_use_dist_strat,
       daisy_chain_variables=daisy_chain_variables,
       ps_replicas=FLAGS.ps_replicas,
       ps_job=FLAGS.ps_job,
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index f93699f95..a97d6ed12 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -825,6 +825,12 @@ def input_fn(self,
     else:
       num_threads = cpu_count() if is_training else 1
 
+    if config and hasattr(config,
+                          "data_parallelism") and config.data_parallelism:
+      num_shards = config.data_parallelism.n
+    else:
+      num_shards = 1
+
     max_length = self.max_length(hparams)
 
     def tpu_valid_size(example):
@@ -891,7 +897,6 @@ def define_shapes(example):
         batch_size = params["batch_size"]
         dataset = dataset.batch(batch_size, drop_remainder=True)
       else:
-        num_shards = config.data_parallelism.n if config else 1
         batch_size = hparams.batch_size * num_shards
         dataset = dataset.batch(batch_size)
     else:
@@ -919,10 +924,9 @@ def define_shapes(example):
       else:
         # On GPU, bucket by length
         dataset = dataset.filter(gpu_valid_size)
-        shard_multiplier = config.data_parallelism.n if config else 1
         batching_scheme = data_reader.hparams_to_batching_scheme(
             hparams,
-            shard_multiplier=shard_multiplier,
+            shard_multiplier=num_shards,
             length_multiplier=self.get_hparams().batch_size_multiplier)
         if hparams.use_fixed_batch_size:
           # Here  batch_size really means examples per datashard.
@@ -934,7 +938,7 @@ def define_shapes(example):
                 batching_scheme["batch_sizes"]))
 
         if not is_training:
-          batch_multiple = shard_multiplier
+          batch_multiple = num_shards
           if hparams.use_fixed_batch_size:
             # Make sure the last batch has the same fixed size as the rest.
             batch_multiple *= hparams.batch_size
@@ -952,8 +956,7 @@ def define_shapes(example):
 
     def prepare_for_output(example):
       if not config or not config.use_tpu:
-        _summarize_features(example,
-                            (config and config.data_parallelism.n) or 1)
+        _summarize_features(example, num_shards)
       if mode == tf.estimator.ModeKeys.PREDICT:
         example["infer_targets"] = example.pop("targets")
         return example
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 5777f1833..0714685b6 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -131,8 +131,8 @@ def _slow_greedy_infer(self, features, decode_length):
 class ImagetransformerMoe(t2t_model.T2TModel):
   """Conditional image generation with attention and MoE."""
 
-  @property
-  def use_body_sharded(self):
+  @staticmethod
+  def use_body_sharded():
     return True
 
   def body_sharded(self, sharded_features):
diff --git a/tensor2tensor/models/research/aligned.py b/tensor2tensor/models/research/aligned.py
index 60afff5e5..7a36ec991 100644
--- a/tensor2tensor/models/research/aligned.py
+++ b/tensor2tensor/models/research/aligned.py
@@ -50,8 +50,8 @@ def _should_postprocess(layer_type):
 class Aligned(t2t_model.T2TModel):
   """Attention net.  See file docstring."""
 
-  @property
-  def use_body_sharded(self):
+  @staticmethod
+  def use_body_sharded():
     return True
 
   def body_sharded(self, sharded_features):
diff --git a/tensor2tensor/models/research/attention_lm_moe.py b/tensor2tensor/models/research/attention_lm_moe.py
index ec2c710e9..4bb0c71f8 100644
--- a/tensor2tensor/models/research/attention_lm_moe.py
+++ b/tensor2tensor/models/research/attention_lm_moe.py
@@ -80,8 +80,8 @@ def get_choices():
 class AttentionLmMoe(t2t_model.T2TModel):
   """Attention net.  See file docstring."""
 
-  @property
-  def use_body_sharded(self):
+  @staticmethod
+  def use_body_sharded():
     return True
 
   def body_sharded(self, sharded_features):
diff --git a/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index a0e5d2ba1..20bd40898 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -56,8 +56,8 @@
 class TransformerMoe(t2t_model.T2TModel):
   """Attention net.  See file docstring."""
 
-  @property
-  def use_body_sharded(self):
+  @staticmethod
+  def use_body_sharded():
     return True
 
   def body_sharded(self, sharded_features):
diff --git a/tensor2tensor/rl/trainer_model_based_new.py b/tensor2tensor/rl/trainer_model_based_new.py
index b75167dbd..6807e7614 100644
--- a/tensor2tensor/rl/trainer_model_based_new.py
+++ b/tensor2tensor/rl/trainer_model_based_new.py
@@ -135,7 +135,7 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
       model_name, problem, data_dir, train_steps, eval_steps,
       min_eval_frequency=local_eval_frequency
   )
-  run_config = trainer_lib.create_run_config(model_dir=output_dir)
+  run_config = trainer_lib.create_run_config(model_name, model_dir=output_dir)
   exp = exp_fn(run_config, hparams)
   getattr(exp, schedule)()
 
diff --git a/tensor2tensor/utils/checkpoint_compatibility_test.py b/tensor2tensor/utils/checkpoint_compatibility_test.py
index d4c4c815c..1e95ca86e 100644
--- a/tensor2tensor/utils/checkpoint_compatibility_test.py
+++ b/tensor2tensor/utils/checkpoint_compatibility_test.py
@@ -60,7 +60,7 @@ def testCompatibility(self):
 
     hp = trainer_lib.create_hparams(
         hp_set, data_dir=_DATA_DIR, problem_name=problem_name)
-    run_config = trainer_lib.create_run_config(model_dir=_CKPT_DIR)
+    run_config = trainer_lib.create_run_config(model, model_dir=_CKPT_DIR)
     estimator = trainer_lib.create_estimator(model, hp, run_config)
 
     for prediction in estimator.predict(self.input_fn):
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index e981a08b9..9103f6cf3 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -224,8 +224,15 @@ def call(self, inputs, **kwargs):
       else:
         return tf.concat(sharded_logits, 0), losses
 
-  @property
-  def use_body_sharded(self):
+  @staticmethod
+  def has_symmetric_shards(model_name):
+    # model_fn is sharded symmetrically unless the model overrides body_sharded
+    # method to manually control the sharding.
+    model_cls = registry.model(model_name)
+    return not model_cls.use_body_sharded()
+
+  @staticmethod
+  def use_body_sharded():
     return False
 
   def body_sharded(self, sharded_features):
@@ -236,7 +243,7 @@ def body_sharded(self, sharded_features):
   def model_fn_sharded(self, sharded_features):
     dp = self._data_parallelism
     datashard_to_features = self._to_features_per_datashard(sharded_features)
-    if self.use_body_sharded:
+    if self.use_body_sharded():
       # MoE models override body_sharded
       transformed_features = dp(self.bottom, datashard_to_features)
       body_out = self.body_sharded(
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index f0be376f9..bc1bf1889 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -105,7 +105,8 @@ def is_cloud_async_distributed():
           json.loads(os.environ.get("TF_CONFIG", "{}")).get("cluster", {}))
 
 
-def create_run_config(master="",
+def create_run_config(model_name,
+                      master="",
                       model_dir=None,
                       iterations_per_loop=1000,
                       num_shards=8,
@@ -121,6 +122,7 @@ def create_run_config(master="",
                       enable_graph_rewriter=False,
                       gpu_mem_fraction=0.95,
                       no_data_parallelism=False,
+                      optionally_use_dist_strat=False,
                       daisy_chain_variables=True,
                       schedule="continuous_train_and_eval",
                       worker_job="/job:localhost",
@@ -205,20 +207,36 @@ def create_run_config(master="",
     config.t2t_device_info = {
         "num_async_replicas": num_async_replicas,
     }
-    config.data_parallelism = devices.data_parallelism(
-        daisy_chain_variables=daisy_chain_variables,
-        ps_replicas=ps_replicas,
-        ps_job=ps_job,
-        ps_gpu=ps_gpu,
-        schedule=schedule,
-        sync=sync,
-        worker_gpu=num_gpus,
-        worker_replicas=num_async_replicas,
-        worker_id=worker_id,
-        gpu_order=gpu_order,
-        locally_shard_to_cpu=shard_to_cpu,
-        worker_job=worker_job,
-        no_data_parallelism=no_data_parallelism)
+    use_distribution_strategy = (
+        optionally_use_dist_strat and
+        t2t_model.T2TModel.has_symmetric_shards(model_name) and
+        not no_data_parallelism and ps_replicas == 0 and ps_gpu == 0 and
+        num_async_replicas == 1 and not shard_to_cpu)
+
+    if use_distribution_strategy:
+      tf.logging.info(
+          "Configuring MirroredStrategy DistributionStrategy to replicate the "
+          "model."
+      )
+      distribution = tf.contrib.distribute.MirroredStrategy()
+      config = config.replace(train_distribute=distribution)
+      config.data_parallelism = None
+    else:
+      tf.logging.info("Configuring DataParallelism to replicate the model.")
+      config.data_parallelism = devices.data_parallelism(
+          daisy_chain_variables=daisy_chain_variables,
+          ps_replicas=ps_replicas,
+          ps_job=ps_job,
+          ps_gpu=ps_gpu,
+          schedule=schedule,
+          sync=sync,
+          worker_gpu=num_gpus,
+          worker_replicas=num_async_replicas,
+          worker_id=worker_id,
+          gpu_order=gpu_order,
+          locally_shard_to_cpu=shard_to_cpu,
+          worker_job=worker_job,
+          no_data_parallelism=no_data_parallelism)
 
   return config
 
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index 01615d929..a3ae4cd85 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -43,7 +43,9 @@ def testExperiment(self):
         min_eval_frequency=1,
         use_tpu=False)
     run_config = trainer_lib.create_run_config(
-        model_dir=algorithmic.TinyAlgo.data_dir, num_gpus=0,
+        model_name="transformer",
+        model_dir=algorithmic.TinyAlgo.data_dir,
+        num_gpus=0,
         use_tpu=False)
     hparams = registry.hparams("transformer_tiny_tpu")
     exp = exp_fn(run_config, hparams)
@@ -59,7 +61,9 @@ def testExperimentWithClass(self):
         min_eval_frequency=1,
         use_tpu=False)
     run_config = trainer_lib.create_run_config(
-        model_dir=algorithmic.TinyAlgo.data_dir, num_gpus=0,
+        model_name="transformer",
+        model_dir=algorithmic.TinyAlgo.data_dir,
+        num_gpus=0,
         use_tpu=False)
     hparams = registry.hparams("transformer_tiny_tpu")
     exp = exp_fn(run_config, hparams)

From ffee1eece8aa976c546b8e01820b9a142741fc4d Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 16 Oct 2018 10:37:50 -0700
Subject: [PATCH 1017/2720] refactoring KL_loss

PiperOrigin-RevId: 217343291
---
 tensor2tensor/models/video/base_vae.py         | 16 ++++++++++------
 tensor2tensor/models/video/basic_stochastic.py |  2 +-
 tensor2tensor/models/video/sv2p.py             | 15 ++-------------
 3 files changed, 13 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/models/video/base_vae.py b/tensor2tensor/models/video/base_vae.py
index 44761dff7..73f989529 100644
--- a/tensor2tensor/models/video/base_vae.py
+++ b/tensor2tensor/models/video/base_vae.py
@@ -71,13 +71,17 @@ def get_beta(self, kl_loss=0.0):
       tf.summary.scalar("beta", beta)
       return beta
 
-  def get_extra_loss(self, mean, std):
-    """Losses in addition to the default modality losses."""
-    kl_loss = common_layers.kl_divergence(mean, std)
+  def get_kl_loss(self, means, stds):
+    """Get KL loss for all the predicted Gaussians."""
+    kl_loss = 0.0
+    if self.is_training and self.hparams.stochastic_model:
+      for i, (mean, std) in enumerate(zip(means, stds)):
+        kl_loss += common_layers.kl_divergence(mean, std)
+        tf.summary.histogram("posterior_mean_%d" % i, mean)
+        tf.summary.histogram("posterior_std_%d" % i, std)
+      tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
+
     beta = self.get_beta(kl_loss)
-    tf.summary.histogram("posterior_mean", mean)
-    tf.summary.histogram("posterior_std", std)
-    tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
     # information capacity from "Understanding disentangling in beta-VAE"
     if self.hparams.information_capacity > 0.0:
       kl_loss = tf.abs(kl_loss - self.hparams.information_capacity)
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 588524776..c192f1e51 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -56,7 +56,7 @@ def inject_latent(self, layer, inputs, target):
     zeros_mask = tf.zeros(
         common_layers.shape_list(layer)[:-1] + [filters], dtype=tf.float32)
     layer = tf.concat([layer, latent_mask + zeros_mask], axis=-1)
-    extra_loss = self.get_extra_loss(latent_mean, latent_std)
+    extra_loss = self.get_kl_loss([latent_mean], [latent_std])
     return layer, extra_loss
 
 
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 121e1650d..91f9e4125 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -195,19 +195,8 @@ def get_extra_loss(self,
                      latent_means=None, latent_stds=None,
                      true_frames=None, gen_frames=None):
     """Losses in addition to the default modality losses."""
-    del true_frames
-    del gen_frames
-    kl_loss = 0.0
-    if self.is_training and self.hparams.stochastic_model:
-      for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
-        kl_loss += common_layers.kl_divergence(mean, std)
-        tf.summary.histogram("posterior_mean_%d" % i, mean)
-        tf.summary.histogram("posterior_std_%d" % i, std)
-      tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
-
-    beta = self.get_beta(kl_loss)
-    extra_loss = beta * kl_loss
-    return extra_loss
+    del true_frames, gen_frames
+    return self.get_kl_loss(latent_means, latent_stds)
 
   def construct_predictive_tower(
       self, input_image, input_reward, action, lstm_state, latent,

From 66db91ebb8f70cb0847dee99d5bdef7e6566638b Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 16 Oct 2018 10:55:31 -0700
Subject: [PATCH 1018/2720] adding sv2p with discrete latent.

PiperOrigin-RevId: 217346836
---
 tensor2tensor/models/video/sv2p.py        | 43 +++++++++++++++++++++--
 tensor2tensor/models/video/sv2p_params.py | 10 ++++++
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 91f9e4125..15b75f899 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -331,7 +331,8 @@ def video_features(
       self, all_frames, all_actions, all_rewards, all_raw_frames):
     """Video wide latent."""
     del all_actions, all_rewards, all_raw_frames
-    mean, std = self.construct_latent_tower(all_frames, time_axis=0)
+    frames = tf.stack(all_frames, axis=1)
+    mean, std = self.construct_latent_tower(frames, time_axis=1)
     latent = common_video.get_gaussian_tensor(mean, std)
     return [latent, mean, std]
 
@@ -343,7 +344,8 @@ def next_frame(self, frames, actions, rewards, target_frame,
     extra_loss = 0.0
     if internal_states is None:
       internal_states = [None] * (5 if self.hparams.small_mode else 7)
-      extra_loss = self.get_extra_loss([latent_mean], [latent_std])
+      if latent_mean is not None:
+        extra_loss = self.get_extra_loss([latent_mean], [latent_std])
 
     pred_image, internal_states = self.construct_predictive_tower(
         frames, None, actions, internal_states, latent)
@@ -356,6 +358,43 @@ def next_frame(self, frames, actions, rewards, target_frame,
     return pred_image, pred_reward, extra_loss, internal_states
 
 
+@registry.register_model
+class NextFrameSv2pDiscrete(NextFrameSv2p):
+  """SV2P with discrete latent."""
+
+  def video_features(
+      self, all_frames, all_actions, all_rewards, all_raw_frames):
+    """Video wide latent."""
+    del all_actions, all_rewards, all_raw_frames
+
+    hparams = self.hparams
+    frames = tf.stack(all_frames, axis=1)
+    mean, std = self.construct_latent_tower(frames, time_axis=1)
+    tower_output = tf.concat([mean, std], axis=-1)
+    tower_output_shape = common_layers.shape_list(tower_output)
+    batch_size = tower_output_shape[0]
+
+    if not self.is_training:
+      rand = tf.random_uniform([batch_size, hparams.bottleneck_bits])
+      d = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
+    else:
+      x = tfl.flatten(tower_output)
+      x = tfl.dense(x, hparams.bottleneck_bits, name="bits_enc")
+      x_shape = common_layers.shape_list(x)
+      x += tf.truncated_normal(x_shape, mean=0.0, stddev=0.2)
+      x = tf.tanh(x)
+      noise = tf.random_uniform(x_shape)
+      noise = 2.0 * tf.to_float(tf.less(hparams.bottleneck_noise, noise)) - 1.0
+      x *= noise
+      d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
+      p = common_layers.inverse_lin_decay(hparams.discrete_warmup_steps)
+      d = tf.where(tf.less(tf.random_uniform([batch_size]), p), d, x)
+
+    decoded_bits = common_video.encode_to_shape(
+        d, tower_output_shape, "bits_dec")
+    return [decoded_bits, None, None]
+
+
 @registry.register_model
 class NextFrameSv2pLegacy(NextFrameSv2p):
   """Old SV2P code. Only for legacy reasons."""
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 0a33d7ee3..23d3956ef 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -50,6 +50,16 @@ def next_frame_sv2p():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_sv2p_discrete():
+  """SV2P discrete model hparams."""
+  hparams = next_frame_sv2p()
+  hparams.add_hparam("bottleneck_bits", 128)
+  hparams.add_hparam("bottleneck_noise", 0.02)
+  hparams.add_hparam("discrete_warmup_steps", 40000)
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_sv2p_atari():
   """SV2P model for atari."""

From be789121646e3cb22b1148a29aecd95e6763085d Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 16 Oct 2018 11:40:36 -0700
Subject: [PATCH 1019/2720] Fixing Discrete model after refactorings.

PiperOrigin-RevId: 217355785
---
 tensor2tensor/models/video/base.py            | 29 +++++++------------
 .../models/video/basic_deterministic.py       |  4 +--
 .../models/video/basic_stochastic.py          |  2 +-
 tensor2tensor/models/video/sv2p.py            |  1 +
 4 files changed, 14 insertions(+), 22 deletions(-)

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 20bd7fd29..cfb6e5b37 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -332,26 +332,17 @@ def __get_next_inputs(self, index, all_frames, all_actions, all_rewards):
       rewards: input rewards for next_frame prediction.
       target_index: index of target frame in all_frames list.
     """
-    i = index
-    j = i + self.hparams.video_num_input_frames
-
-    actions, rewards = None, None
     if self.is_recurrent_model:
-      frames = all_frames[i]
-      target_index = i+1
-      if self.has_actions:
-        actions = all_actions[i]
-      if self.has_rewards:
-        rewards = all_rewards[i]
+      target_index = index + 1
+      nones = [None]
     else:
-      frames = all_frames[i:j]
-      target_index = j
-      if self.has_actions:
-        actions = all_actions[i:j]
-        actions = tf.concat(actions, axis=1)
-      if self.has_rewards:
-        rewards = all_rewards[i:j]
-        rewards = tf.concat(rewards, axis=1)
+      target_index = index + self.hparams.video_num_input_frames
+      nones = [None] * self.hparams.video_num_input_frames
+
+    frames = all_frames[index:target_index]
+    actions = all_actions[index:target_index] if self.has_actions else nones
+    rewards = all_rewards[index:target_index] if self.has_rewards else nones
+
     return frames, actions, rewards, target_index
 
   def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
@@ -455,7 +446,7 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
         res_frame, res_reward, res_extra_loss, internal_states = func_out
         res_frames.append(res_frame)
         res_rewards.append(res_reward)
-        extra_loss += res_extra_loss / float(hparams.video_num_target_frames)
+        extra_loss += res_extra_loss / float(len(input_index_range))
 
       # Only for Softmax loss: sample frame so we can keep iterating.
       sampled_frame = self.get_sampled_frame(res_frame)
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index af570e4fd..6fc3f2a95 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -63,7 +63,7 @@ def middle_network(self, layer, internal_states):
 
   def next_frame(self, frames, actions, rewards, target_frame,
                  internal_states, video_extra):
-    del video_extra
+    del rewards, video_extra
 
     hparams = self.hparams
     filters = hparams.hidden_size
@@ -93,7 +93,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
 
     # Add embedded action if present.
     if self.has_actions:
-      action = actions[:, -1, :]
+      action = actions[-1]
       x = common_video.inject_additional_input(
           x, action, "action_enc", hparams.action_injection)
 
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index c192f1e51..7d1d9d366 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -111,7 +111,7 @@ def add_d(layer, d):
       return add_d(layer, d), 0.0
 
     # Embed.
-    frames = tf.concat([inputs, target], axis=-1)
+    frames = tf.concat(inputs + [target], axis=-1)
     x = tfl.dense(
         frames, filters, name="latent_embed",
         bias_initializer=tf.random_normal_initializer(stddev=0.01))
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 15b75f899..08f1fd10d 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -340,6 +340,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
                  internal_states, video_features):
     del target_frame
     latent, latent_mean, latent_std = video_features
+    frames, actions, rewards = frames[0], actions[0], rewards[0]
 
     extra_loss = 0.0
     if internal_states is None:

From ca628e4fcb04ff42ed21549a4f73e6dfa68a5f7a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 16 Oct 2018 14:29:57 -0700
Subject: [PATCH 1020/2720] Clean up stochastic discrete video model by moving
 code and using discretization layers.

PiperOrigin-RevId: 217386874
---
 tensor2tensor/layers/common_layers.py         |  6 +-
 tensor2tensor/layers/discretization.py        | 81 +++++++++++++++++-
 .../models/video/basic_stochastic.py          | 84 +++++--------------
 .../rl/trainer_model_based_params.py          |  2 +-
 4 files changed, 104 insertions(+), 69 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index b928dc407..1414631dc 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3019,8 +3019,10 @@ def get_res():
     if is_xla_compiled():
       return get_res()
     else:
-      return tf.cond(
-          tf.less(tf.train.get_global_step(), steps), get_res, lambda: x1)
+      cur_step = tf.train.get_global_step()
+      if cur_step is None:
+        return x1  # Step not available, probably eval mode, don't mix.
+      return tf.cond(tf.less(cur_step, steps), get_res, lambda: x1)
 
 
 def brelu(x):
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index d443cb29b..71a7cd988 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -783,6 +783,81 @@ def discrete_bottleneck(inputs,
   return outputs_dense, outputs_discrete, extra_loss, embed_fn, neg_q_entropy
 
 
+def predict_bits_with_lstm(prediction_source, state_size, total_num_bits,
+                           target_bits=None, bits_at_once=8, temperature=1.0):
+  """Predict a sequence of bits (a latent) with LSTM, both training and infer.
+
+  Given a tensor on which the predictions are based (prediction_source), we use
+  a single-layer LSTM with state of size state_size to predict total_num_bits,
+  which we predict in groups of size bits_at_once. During training, we use
+  target_bits as input to the LSTM (teacher forcing) and return the target_bits
+  together with the prediction loss. During inference, we sample with the given
+  temperature and return the predicted sequence and loss 0.
+
+  Args:
+    prediction_source: a Tensor of shape [batch_size, ...] used to create
+      the initial state and the first input to the LSTM.
+    state_size: python integer, the size of the LSTM state.
+    total_num_bits: python integer, how many bits in total to predict.
+    target_bits: a tensor of shape [batch_size, total_num_bits] used during
+      training as the target to predict; each element should be -1 or 1.
+    bits_at_once: pytho integer, how many bits to predict at once.
+    temperature: python float, temperature used for sampling during inference.
+
+  Returns:
+    a pair (bits, loss) with the predicted bit sequence, which is a Tensor of
+    shape [batch_size, total_num_bits] with elements either -1 or 1, and a loss
+    used to train the predictions against the provided target_bits.
+  """
+
+  with tf.variable_scope("predict_bits_with_lstm"):
+    # Layers and cell state creation.
+    lstm_cell = tf.contrib.rnn.LSTMCell(state_size)
+    discrete_predict = tf.layers.Dense(2**bits_at_once, name="discrete_predict")
+    discrete_embed = tf.layers.Dense(state_size, name="discrete_embed")
+    batch_size = common_layers.shape_list(prediction_source)[0]
+    layer_pred = tf.layers.flatten(prediction_source)
+    prediction = tf.layers.dense(layer_pred, state_size, name="istate")
+    c_state = tf.layers.dense(layer_pred, state_size, name="cstate")
+    m_state = tf.layers.dense(layer_pred, state_size, name="mstate")
+    state = (c_state, m_state)
+
+    # Prediction mode if no targets are given.
+    if target_bits is None:
+      outputs = []
+      for i in range(total_num_bits // bits_at_once):
+        output, state = lstm_cell(prediction, state)
+        discrete_logits = discrete_predict(output)
+        discrete_samples = common_layers.sample_with_temperature(
+            discrete_logits, temperature)
+        outputs.append(tf.expand_dims(discrete_samples, axis=1))
+        prediction = discrete_embed(tf.one_hot(discrete_samples, 256))
+      outputs = tf.concat(outputs, axis=1)
+      outputs = int_to_bit(outputs, bits_at_once)
+      outputs = tf.reshape(outputs, [batch_size, total_num_bits])
+      return 2 * outputs - 1, 0.0
+
+    # Training mode, calculating loss.
+    assert total_num_bits % bits_at_once == 0
+    d_pred = tf.reshape(tf.maximum(tf.stop_gradient(target_bits), 0), [
+        batch_size, total_num_bits // bits_at_once, bits_at_once])
+    d_int = bit_to_int(d_pred, bits_at_once)
+    tf.summary.histogram("target_integers", tf.reshape(d_int, [-1]))
+    d_hot = tf.one_hot(d_int, 2**bits_at_once, axis=-1)
+    d_pred = discrete_embed(d_hot)
+    pred = tf.concat([tf.expand_dims(prediction, axis=1), d_pred], axis=1)
+    outputs = []
+    for i in range(total_num_bits // bits_at_once):
+      output, state = lstm_cell(pred[:, i, :], state)
+      outputs.append(tf.expand_dims(output, axis=1))
+    outputs = tf.concat(outputs, axis=1)
+    d_int_pred = discrete_predict(outputs)
+    pred_loss = tf.losses.sparse_softmax_cross_entropy(
+        logits=d_int_pred, labels=d_int)
+    pred_loss = tf.reduce_mean(pred_loss)
+    return target_bits, pred_loss
+
+
 # New API for discretization bottlenecks:
 # * Each method is separate and provides 2 functions:
 # * The [method]_bottleneck function returns discretized state.
@@ -1281,6 +1356,7 @@ def tanh_discrete_bottleneck(x, bottleneck_bits, bottleneck_noise,
                              discretize_warmup_steps, mode):
   """Simple discretization through tanh, flip bottleneck_noise many bits."""
   x = tf.layers.dense(x, bottleneck_bits, name="tanh_discrete_bottleneck")
+  d0 = tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x))) - 1.0
   if mode == tf.estimator.ModeKeys.TRAIN:
     x += tf.truncated_normal(
         common_layers.shape_list(x), mean=0.0, stddev=0.2)
@@ -1292,7 +1368,7 @@ def tanh_discrete_bottleneck(x, bottleneck_bits, bottleneck_noise,
     d *= noise
   d = common_layers.mix(d, x, discretize_warmup_steps,
                         mode == tf.estimator.ModeKeys.TRAIN)
-  return d, 0.0
+  return d, d0
 
 
 def tanh_discrete_unbottleneck(x, hidden_size):
@@ -1345,9 +1421,10 @@ def isemhash_unbottleneck(x, hidden_size, isemhash_filter_size_multiplier=1.0):
 def parametrized_bottleneck(x, hparams):
   """Meta-function calling all the above bottlenecks with hparams."""
   if hparams.bottleneck_kind == "tanh_discrete":
-    return tanh_discrete_bottleneck(
+    d, _ = tanh_discrete_bottleneck(
         x, hparams.bottleneck_bits, hparams.bottleneck_noise * 0.5,
         hparams.discretize_warmup_steps, hparams.mode)
+    return d, 0.0
   if hparams.bottleneck_kind == "isemhash":
     return isemhash_bottleneck(
         x, hparams.bottleneck_bits, hparams.bottleneck_noise * 0.5,
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 7d1d9d366..9ef01d5cc 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -72,43 +72,26 @@ def inject_latent(self, layer, inputs, target):
     filters = hparams.hidden_size
     kernel = (4, 4)
     layer_shape = common_layers.shape_list(layer)
-    batch_size = layer_shape[0]
-    state_size = hparams.latent_predictor_state_size
-    lstm_cell = tf.contrib.rnn.LSTMCell(state_size)
-    discrete_predict = tfl.Dense(256, name="discrete_predict")
-    discrete_embed = tfl.Dense(state_size, name="discrete_embed")
-
-    def add_d(layer, d):
-      z_mul = tfl.dense(d, final_filters, name="unbottleneck_mul")
+
+    def add_bits(layer, bits):
+      z_mul = tfl.dense(bits, final_filters, name="unbottleneck_mul")
       if not hparams.complex_addn:
         return layer + z_mul
       layer *= tf.nn.sigmoid(z_mul)
-      z_add = tfl.dense(d, final_filters, name="unbottleneck_add")
+      z_add = tfl.dense(bits, final_filters, name="unbottleneck_add")
       layer += z_add
       return layer
 
     if self.is_predicting:
       if hparams.full_latent_tower:
         rand = tf.random_uniform(layer_shape[:-1] + [hparams.bottleneck_bits])
+        bits = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
       else:
-        layer_pred = tfl.flatten(layer)
-        prediction = tfl.dense(layer_pred, state_size, name="istate")
-        c_state = tfl.dense(layer_pred, state_size, name="cstate")
-        m_state = tfl.dense(layer_pred, state_size, name="mstate")
-        state = (c_state, m_state)
-        outputs = []
-        for i in range(hparams.bottleneck_bits // 8):
-          output, state = lstm_cell(prediction, state)
-          discrete_logits = discrete_predict(output)
-          discrete_samples = common_layers.sample_with_temperature(
-              discrete_logits, hparams.latent_predictor_temperature)
-          outputs.append(tf.expand_dims(discrete_samples, axis=1))
-          prediction = discrete_embed(tf.one_hot(discrete_samples, 256))
-        outputs = tf.concat(outputs, axis=1)
-        outputs = discretization.int_to_bit(outputs, 8)
-        rand = tf.reshape(outputs, [batch_size, 1, 1, hparams.bottleneck_bits])
-      d = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
-      return add_d(layer, d), 0.0
+        bits, _ = discretization.predict_bits_with_lstm(
+            layer, hparams.latent_predictor_state_size, hparams.bottleneck_bits,
+            temperature=hparams.latent_predictor_temperature)
+        bits = tf.expand_dims(tf.expand_dims(bits, axis=1), axis=2)
+      return add_bits(layer, bits), 0.0
 
     # Embed.
     frames = tf.concat(inputs + [target], axis=-1)
@@ -131,43 +114,16 @@ def add_d(layer, d):
     else:
       x = common_layers.double_discriminator(x)
       x = tf.expand_dims(tf.expand_dims(x, axis=1), axis=1)
-    x = tfl.dense(x, hparams.bottleneck_bits, name="bottleneck")
-    x0 = tf.tanh(x)
-    d = x0 + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x0)) - 1.0 - x0)
-    pred_loss = 0.0
+
+    bits, bits_clean = discretization.tanh_discrete_bottleneck(
+        x, hparams.bottleneck_bits, hparams.bottleneck_noise,
+        hparams.discretize_warmup_steps, hparams.mode)
     if not hparams.full_latent_tower:
-      d_pred = tf.reshape(tf.maximum(tf.stop_gradient(d), 0), [
-          batch_size, hparams.bottleneck_bits // 8, 8])
-      d_int = discretization.bit_to_int(d_pred, 8)
-      tf.summary.histogram("d_int", tf.reshape(d_int, [-1]))
-      d_hot = tf.one_hot(d_int, 256, axis=-1)
-      d_pred = discrete_embed(d_hot)
-      layer_pred = tfl.flatten(layer)
-      prediction0 = tfl.dense(layer_pred, state_size, name="istate")
-      c_state = tfl.dense(layer_pred, state_size, name="cstate")
-      m_state = tfl.dense(layer_pred, state_size, name="mstate")
-      pred = tf.concat([tf.expand_dims(prediction0, axis=1), d_pred], axis=1)
-      state = (c_state, m_state)
-      outputs = []
-      for i in range(hparams.bottleneck_bits // 8):
-        output, state = lstm_cell(pred[:, i, :], state)
-        outputs.append(tf.expand_dims(output, axis=1))
-      outputs = tf.concat(outputs, axis=1)
-      d_int_pred = discrete_predict(outputs)
-      pred_loss = tf.losses.sparse_softmax_cross_entropy(
-          logits=d_int_pred, labels=d_int)
-      pred_loss = tf.reduce_mean(pred_loss)
-    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
-      x += tf.truncated_normal(
-          common_layers.shape_list(x), mean=0.0, stddev=0.2)
-      x = tf.tanh(x)
-      noise = tf.random_uniform(common_layers.shape_list(x))
-      noise = 2.0 * tf.to_float(tf.less(hparams.bottleneck_noise, noise)) - 1.0
-      x *= noise
-      d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
-      p = common_layers.inverse_lin_decay(hparams.discrete_warmup_steps)
-      d = tf.where(tf.less(tf.random_uniform([batch_size]), p), d, x)
-    return add_d(layer, d), pred_loss
+      _, pred_loss = discretization.predict_bits_with_lstm(
+          layer, hparams.latent_predictor_state_size, hparams.bottleneck_bits,
+          target_bits=bits_clean)
+
+    return add_bits(layer, bits), pred_loss
 
 
 @registry.register_hparams
@@ -224,7 +180,7 @@ def next_frame_basic_stochastic_discrete():
   hparams.learning_rate_schedule = "linear_warmup * constant"
   hparams.add_hparam("bottleneck_bits", 64)
   hparams.add_hparam("bottleneck_noise", 0.02)
-  hparams.add_hparam("discrete_warmup_steps", 40000)
+  hparams.add_hparam("discretize_warmup_steps", 40000)
   hparams.add_hparam("full_latent_tower", False)
   hparams.add_hparam("latent_predictor_state_size", 128)
   hparams.add_hparam("latent_predictor_temperature", 0.5)
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 9e3fd0b91..a7b34d7b0 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -203,7 +203,7 @@ def rlmb_base_stochastic_discrete():
   """Base setting with stochastic discrete model."""
   hparams = rlmb_base()
   hparams.learning_rate_bump = 1.0
-  hparams.grayscale = True
+  hparams.grayscale = False
   hparams.generative_model = "next_frame_basic_stochastic_discrete"
   hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
   return hparams

From 14738e3e3310bd1b9f1eae99bc2a0e212e80e776 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Tue, 16 Oct 2018 23:52:57 +0200
Subject: [PATCH 1021/2720] Model-Based RL: Data loading, restarting
 experiments (#1146)

* Global data_dir

* Move determining filepattern for only last epoch to Problem.filepattern

* Include epoch in shard filenames

* Implement data loading
---
 tensor2tensor/data_generators/gym_env.py      | 178 +++++++++++++-----
 tensor2tensor/data_generators/gym_env_test.py |  96 +++++++---
 tensor2tensor/data_generators/gym_problems.py |   8 +
 tensor2tensor/data_generators/problem.py      |  12 +-
 tensor2tensor/rl/envs/batch_env_factory.py    |   3 +-
 tensor2tensor/rl/trainer_model_based_new.py   |  31 ++-
 6 files changed, 237 insertions(+), 91 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 0ace13f9a..cc98db28d 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -24,7 +24,6 @@
 
 from gym.spaces import Box
 import numpy as np
-import six
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -72,11 +71,12 @@ class T2TEnv(video_utils.VideoProblem):
   reward_range = (-1, 1)
   name = None
 
-  def __init__(self, batch_size):
+  def __init__(self, batch_size, data_dir):
     super(T2TEnv, self).__init__()
 
     self.clear_history()
     self.batch_size = batch_size
+    self.data_dir = data_dir
     self._current_batch_frames = [None for _ in range(batch_size)]
     self._current_batch_rollouts = [[] for _ in range(batch_size)]
     self._current_epoch_rollouts = []
@@ -84,11 +84,11 @@ def __init__(self, batch_size):
     self.current_epoch = None
     with tf.Graph().as_default() as tf_graph:
       self._tf_graph = _Noncopyable(tf_graph)
-      self._image_t = _Noncopyable(
+      self._image_p = _Noncopyable(
           tf.placeholder(dtype=tf.uint8, shape=(None, None, None))
       )
       self._encoded_image_t = _Noncopyable(
-          tf.image.encode_png(self._image_t.obj)
+          tf.image.encode_png(self._image_p.obj)
       )
       self._session = _Noncopyable(tf.Session())
 
@@ -100,18 +100,20 @@ def clear_history(self):
     """Clears the rollout history."""
     self._rollouts_by_epoch_and_split = collections.OrderedDict()
 
-  def start_new_epoch(self, epoch):
+  def start_new_epoch(self, epoch, load_data=True):
     if not isinstance(epoch, int):
       raise ValueError("Epoch should be integer, got {}".format(epoch))
     if epoch in self._rollouts_by_epoch_and_split:
       raise ValueError("Epoch {} already registered".format(epoch))
     self.current_epoch = epoch
     self._current_epoch_rollouts = []
+    self._rollouts_by_epoch_and_split[epoch] = collections.defaultdict(list)
+    if load_data:
+      self._load_epoch_data()
 
   def current_epoch_rollouts(self, split=None):
-    try:
-      rollouts_by_split = self._rollouts_by_epoch_and_split[self.current_epoch]
-    except KeyError:
+    rollouts_by_split = self._rollouts_by_epoch_and_split[self.current_epoch]
+    if not rollouts_by_split:
       if split is not None:
         raise ValueError(
             "generate_data() should first be called in the current epoch"
@@ -145,7 +147,7 @@ def _encode_observations(self, observations):
     return [
         self._session.obj.run(
             self._encoded_image_t.obj,
-            feed_dict={self._image_t.obj: observation}
+            feed_dict={self._image_p.obj: observation}
         )
         for observation in observations
     ]
@@ -179,10 +181,15 @@ def step(self, actions):
       (obs, rewards, dones) - batches of observations, rewards and done flags
       respectively.
     """
+    if self._rollouts_by_epoch_and_split[self.current_epoch]:
+      raise ValueError(
+          "Data for current epoch has already been loaded from disk."
+      )
     (obs, unclipped_rewards, dones) = self._step(actions)
     obs = self._preprocess_observations(obs)
     (min_reward, max_reward) = self.reward_range
     rewards = np.around(np.clip(unclipped_rewards, min_reward, max_reward))
+    unclipped_rewards = unclipped_rewards.astype(np.float64)
     encoded_obs = self._encode_observations(obs)
     for (rollout, frame, action) in zip(
         self._current_batch_rollouts, self._current_batch_frames, actions
@@ -222,7 +229,7 @@ def reset(self, indices=None):
     """
     if self.current_epoch is None:
       # It's here so that the old pipeline works.
-      self.start_new_epoch(0)
+      self.start_new_epoch(0, load_data=False)
       # TODO(koz4k): Replace with:
       # raise ValueError(
       #     "No current epoch. start_new_epoch() should first be called."
@@ -311,18 +318,19 @@ def make_modality(name):
     p.input_space_id = problem.SpaceID.IMAGE
     p.target_space_id = problem.SpaceID.IMAGE
 
-  def _generate_frames(self, epoch, rollouts):
+  def _generate_frames(self, rollouts):
     for rollout in rollouts:
       for (frame_number, frame) in enumerate(rollout):
         yield {
             "frame_number": [frame_number],
-            "epoch": [epoch],
+            "epoch": [self.current_epoch],
             "image/encoded": [frame.observation],
             "image/format": ["png"],
             "image/height": [self.frame_height],
             "image/width": [self.frame_width],
             "action": [int(frame.action)],
             "reward": [int(frame.reward - self.reward_range[0])],
+            "unclipped_reward": [float(frame.unclipped_reward)],
             "done": [int(frame.done)]
         }
 
@@ -377,52 +385,132 @@ def split_size(split_index):
     self._rollouts_by_epoch_and_split[self.current_epoch] = rollouts_by_split
     self._current_epoch_rollouts = []
 
-  def splits_and_paths(self, data_dir):
+  @property
+  def splits_and_paths(self):
+    """List of pairs (split, paths) for the current epoch.
+
+    paths is a list of paths where data for the current epoch is saved
+    by generate_data().
+    """
     filepath_fns = {
         problem.DatasetSplit.TRAIN: self.training_filepaths,
         problem.DatasetSplit.EVAL: self.dev_filepaths,
         problem.DatasetSplit.TEST: self.test_filepaths,
     }
 
-    num_epochs = len(self._rollouts_by_epoch_and_split)
+    def append_epoch(paths):
+      return [
+          "{}.{}".format(path, self.current_epoch)
+          for path in paths
+      ]
+
     # We set shuffled=True as we don't want to shuffle on disk later.
     return [
-        (split["split"], filepath_fns[split["split"]](
-            data_dir, split["shards"] * num_epochs, shuffled=True
-        ))
+        (split["split"], append_epoch(filepath_fns[split["split"]](
+            self.data_dir, split["shards"], shuffled=True
+        )))
         for split in self.dataset_splits
     ]
 
-  def generate_data(self, data_dir, tmp_dir, task_id=-1):
-    """Saves the rollout history to disk, split into train/dev sets."""
-    self._split_current_epoch()
+  def filepattern(self, data_dir, mode, shard=None, only_last=False):
+    filepattern = super(T2TEnv, self).filepattern(
+        data_dir, mode, shard
+    )
+    if only_last:
+      filepattern += ".{}".format(self.current_epoch)
+    return filepattern
 
-    splits_and_paths = self.splits_and_paths(data_dir)
-    num_epochs = len(self._rollouts_by_epoch_and_split)
+  def generate_data(self, data_dir=None, tmp_dir=None, task_id=-1):
+    """Saves the current epoch rollouts to disk, split into train/dev sets.
 
-    for (epoch_index, (epoch, rollouts_by_split)) in enumerate(
-        six.iteritems(self._rollouts_by_epoch_and_split)
-    ):
-      for (split, paths) in splits_and_paths:
-        num_shards = len(paths) // num_epochs
-        paths = paths[
-            epoch_index * num_shards : (epoch_index + 1) * num_shards
-        ]
+    data_dir and tmp_dir arguments are unused. data_dir being used is the one
+    passed in the constructor.
+    """
+    if not self._rollouts_by_epoch_and_split[self.current_epoch]:
+      # Data not loaded from disk.
+      self._split_current_epoch()
+
+    rollouts_by_split = self._rollouts_by_epoch_and_split[self.current_epoch]
+    splits_and_paths = self.splits_and_paths
+
+    for (split, paths) in splits_and_paths:
+      rollouts = rollouts_by_split[split]
+      num_frames = self._calc_num_frames(rollouts)
+      shard_size = num_frames // len(paths)
+
+      frame_gen = self._generate_frames(rollouts)
+      for (path_index, path) in enumerate(paths):
+        limit = shard_size
+        # Put the remainder in the last shard to preserve the ordering.
+        if path_index == len(paths) - 1:
+          limit = None
+        generator_utils.generate_files(
+            itertools.islice(frame_gen, limit), [path],
+            cycle_every_n=float("inf")
+        )
+
+  def _load_epoch_data(self):
+    any_files_found = False
+    all_files_found = True
+    any_shard_empty = False
+
+    for split, paths in self.splits_and_paths:
+      try:
+        any_shard_empty |= self._load_epoch_split(split, paths)
+        any_files_found = True
+      except tf.errors.NotFoundError:
+        all_files_found = False
+    if any_shard_empty or (not all_files_found and any_files_found):
+      raise ValueError("Some data is missing, the experiment might've been "
+                       "interupted during generating data.")
+
+  def _load_epoch_split(self, split, paths):
+    epoch = self.current_epoch
+    last_frame_number = -1
+    any_shard_empty = False
+    current_rollout = []
+
+    for path in paths:
+      this_shard_empty = True
+      for example in tf.python_io.tf_record_iterator(path):
+        this_shard_empty = False
+
+        result = tf.train.Example.FromString(example)
+        feature = result.features.feature
+
+        def get_feature_value(key, list_name):
+          return getattr(feature[key], list_name).value[0]  # pylint: disable=cell-var-from-loop
+
+        fields = {
+            key: get_feature_value(key, list_name)
+            for (key, list_name) in [
+                ("image/encoded", "bytes_list"), ("reward", "int64_list"),
+                ("unclipped_reward", "float_list"), ("done", "int64_list"),
+                ("action", "int64_list")
+            ]
+        }
+        fields["reward"] += self.reward_range[0]
+        fields["done"] = bool(fields["done"])
+        fields['observation'] = fields['image/encoded']
+        del fields['image/encoded']
+
+        frame = Frame(**fields)
+        frame_number = get_feature_value("frame_number", 'int64_list')
+        if frame_number == last_frame_number + 1:
+          current_rollout.append(frame)
+        else:
+          self._rollouts_by_epoch_and_split[epoch][split].append(
+              current_rollout)
+          current_rollout = [frame]
+        last_frame_number = frame_number
+
+      any_shard_empty |= this_shard_empty
 
-        rollouts = rollouts_by_split[split]
-        num_frames = self._calc_num_frames(rollouts)
-        shard_size = num_frames // len(paths)
+    self._rollouts_by_epoch_and_split[epoch][split].append(
+        current_rollout
+    )
+    return any_shard_empty
 
-        frame_gen = self._generate_frames(epoch, rollouts)
-        for (path_index, path) in enumerate(paths):
-          limit = shard_size
-          # Put the remainder in the last shard to preserve the ordering.
-          if path_index == len(paths) - 1:
-            limit = None
-          generator_utils.generate_files(
-              itertools.islice(frame_gen, limit), [path],
-              cycle_every_n=float("inf")
-          )
 
 
 class T2TGymEnv(T2TEnv):
@@ -430,9 +518,9 @@ class T2TGymEnv(T2TEnv):
 
   name = "t2t_gym_env"
 
-  def __init__(self, envs, grayscale=False,
+  def __init__(self, envs, data_dir, grayscale=False,
                resize_height_factor=1, resize_width_factor=1):
-    super(T2TGymEnv, self).__init__(len(envs))
+    super(T2TGymEnv, self).__init__(len(envs), data_dir)
     self.grayscale = grayscale
     self.resize_height_factor = resize_height_factor
     self.resize_width_factor = resize_width_factor
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 075572e47..b283844e1 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -66,16 +66,32 @@ class GymEnvTest(tf.test.TestCase):
 
   splits = (problem.DatasetSplit.TRAIN, problem.DatasetSplit.EVAL)
 
+  # TODO(koz4k): Tests for loading:
+  # - loaded epoch is read-only
+  # - partial write detection (should raise on loading)
+
   def setUp(self):
     self.out_dir = tf.test.get_temp_dir()
     shutil.rmtree(self.out_dir)
     os.mkdir(self.out_dir)
 
-  def init_batch_and_play(self, env_lambda, n_steps=1, **kwargs):
+  def init_batch_and_play(self, env_lambda, steps_per_epoch=1,
+                          epochs=(0,), generate_data=False, **kwargs):
     raw_envs = [env_lambda(), env_lambda()]
-    env = gym_env.T2TGymEnv(raw_envs, **kwargs)
-    env.start_new_epoch(0)
-    return self.play(env, n_steps)
+    env = gym_env.T2TGymEnv(raw_envs, self.out_dir, **kwargs)
+    obs = list()
+    rewards = list()
+    num_dones = 0
+    for epoch in epochs:
+      env.start_new_epoch(epoch)
+      _, epoch_obs, epoch_rewards, epoch_num_dones = \
+          self.play(env, steps_per_epoch)
+      if generate_data:
+        env.generate_data()
+      obs.extend(epoch_obs)
+      rewards.extend(epoch_rewards)
+      num_dones += epoch_num_dones
+    return env, obs, rewards, num_dones
 
   def play(self, env, n_steps):
     obs = list()
@@ -93,15 +109,17 @@ def play(self, env, n_steps):
     return env, obs, rewards, num_dones
 
   def test_splits_dataset(self):
-    env, _, _, _ = self.init_batch_and_play(TestEnv, n_steps=20)
-    env.generate_data(self.out_dir, tmp_dir=None)
+    env, _, _, _ = self.init_batch_and_play(
+        TestEnv, steps_per_epoch=20, generate_data=True
+    )
 
     for split in self.splits:
       self.assertTrue(env.current_epoch_rollouts(split))
 
   def test_split_preserves_number_of_rollouts(self):
-    env, _, _, num_dones = self.init_batch_and_play(TestEnv, n_steps=20)
-    env.generate_data(self.out_dir, tmp_dir=None)
+    env, _, _, num_dones = self.init_batch_and_play(
+        TestEnv, steps_per_epoch=20, generate_data=True
+    )
 
     num_rollouts_after_split = sum(
         len(env.current_epoch_rollouts(split)) for split in self.splits
@@ -112,8 +130,9 @@ def test_split_preserves_number_of_rollouts(self):
     self.assertLessEqual(num_rollouts_after_split, num_dones + 1)
 
   def test_split_preserves_number_of_frames(self):
-    env, _, _, num_dones = self.init_batch_and_play(TestEnv, n_steps=20)
-    env.generate_data(self.out_dir, tmp_dir=None)
+    env, _, _, num_dones = self.init_batch_and_play(
+        TestEnv, steps_per_epoch=20, generate_data=True
+    )
 
     num_frames = sum(
         len(rollout)
@@ -126,9 +145,9 @@ def test_split_preserves_number_of_frames(self):
 
   def test_generates_data(self):
     # This test needs base env which outputs done after two steps.
-    env_lambda = TestEnv
-    env, _, _, _ = self.init_batch_and_play(env_lambda, n_steps=20)
-    env.generate_data(self.out_dir, tmp_dir=None)
+    self.init_batch_and_play(
+        TestEnv, steps_per_epoch=20, generate_data=True
+    )
 
     filenames = os.listdir(self.out_dir)
     self.assertTrue(filenames)
@@ -138,26 +157,39 @@ def test_generates_data(self):
       self.assertTrue(records)
 
   def test_shards_per_epoch(self):
-    env, _, _, _ = self.init_batch_and_play(TestEnv, n_steps=20)
-    env.generate_data(self.out_dir, tmp_dir=None)
-    num_shards_per_epoch = len(os.listdir(self.out_dir))
-    shutil.rmtree(self.out_dir)
-    os.mkdir(self.out_dir)
+    def num_ending_with(filenames, suffix):
+      return sum(
+          1 for filename in filenames if filename.endswith(suffix)
+      )
+
+    env = gym_env.T2TGymEnv([TestEnv() for _ in range(2)], self.out_dir)
+    env.start_new_epoch(0)
+    self.play(env, n_steps=20)
+    env.generate_data()
+
+    filenames = os.listdir(self.out_dir)
+    num_shards_per_epoch = len(filenames)
+    self.assertEqual(num_ending_with(filenames, ".0"), num_shards_per_epoch)
 
     env.start_new_epoch(1)
     self.play(env, n_steps=20)
-    env.generate_data(self.out_dir, tmp_dir=None)
-    self.assertEqual(len(os.listdir(self.out_dir)), 2 * num_shards_per_epoch)
+    env.generate_data()
+
+    filenames = os.listdir(self.out_dir)
+    self.assertEqual(len(filenames), 2 * num_shards_per_epoch)
+    for suffix in (".0", ".1"):
+      self.assertEqual(num_ending_with(filenames, suffix), num_shards_per_epoch)
 
   def test_frame_numbers_are_continuous(self):
-    env, _, _, _ = self.init_batch_and_play(TestEnv, n_steps=20)
-    env.generate_data(self.out_dir, tmp_dir=None)
+    env, _, _, _ = self.init_batch_and_play(
+        TestEnv, steps_per_epoch=20, generate_data=True
+    )
 
     frame_numbers = [
         tf.train.Example.FromString(
             record
         ).features.feature["frame_number"].int64_list.value[0]
-        for (_, paths) in env.splits_and_paths(self.out_dir)
+        for (_, paths) in env.splits_and_paths
         for path in paths
         for record in tf.python_io.tf_record_iterator(path)
     ]
@@ -178,7 +210,9 @@ def test_clipping(self):
     # self.assertTrue(np.max(rewards) == 1)
     # self.assertTrue(np.min(rewards) == -1)
 
-    _, _, unclipped_rewards, _ = self.init_batch_and_play(env_lambda, n_steps=2)
+    _, _, unclipped_rewards, _ = self.init_batch_and_play(
+        env_lambda, steps_per_epoch=2
+    )
     self.assertTrue(np.max(unclipped_rewards) > 1)
     self.assertTrue(np.min(unclipped_rewards) < -1)
 
@@ -189,7 +223,7 @@ def test_resize(self):
     resize_width_factor = 3
     orig_height, orig_width = orig_env.observation_space.shape[:2]
     env, obs, _, _ = self.init_batch_and_play(
-        env_lambda, n_steps=1,
+        env_lambda, steps_per_epoch=1,
         resize_height_factor=resize_height_factor,
         resize_width_factor=resize_width_factor)
     for obs_batch in obs:
@@ -214,6 +248,18 @@ def test_channels(self):
     env, obs, _, _ = self.init_batch_and_play(env_lambda, grayscale=False)
     self.assert_channels(env, obs, n_channels=3)
 
+  def test_generating_and_loading_preserves_rollouts(self):
+    from_env = gym_env.T2TGymEnv([TestEnv()], self.out_dir)
+    from_env.start_new_epoch(0)
+    self.play(from_env, n_steps=20)
+    from_env.generate_data()
+
+    to_env = gym_env.T2TGymEnv([TestEnv()], self.out_dir)
+    to_env.start_new_epoch(0)
+
+    self.assertEqual(
+        from_env.current_epoch_rollouts(), to_env.current_epoch_rollouts()
+    )
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index c0241870d..1069c929d 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -353,6 +353,14 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
     else:
       self.statistics.save_to_file(stats_file)
 
+  def filepattern(self, data_dir, mode, shard=None, only_last=False):
+    filepattern = super(GymDiscreteProblem, self).filepattern(
+        data_dir, mode, shard
+    )
+    if only_last:
+      filepattern += r"10.[\d+]"
+    return filepattern
+
 
 class BasicStatistics(object):
   """Keeps basic statistics to calculate mean reward """
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index a97d6ed12..7d9dbd225 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -420,7 +420,7 @@ def test_filepaths(self, data_dir, num_shards, shuffled):
     return generator_utils.test_data_filenames(file_basename, data_dir,
                                                num_shards)
 
-  def filepattern(self, data_dir, mode, shard=None):
+  def filepattern(self, data_dir, mode, shard=None, only_last=False):
     """Get filepattern for data files for mode.
 
     Matches mode to a suffix.
@@ -433,6 +433,7 @@ def filepattern(self, data_dir, mode, shard=None):
       data_dir: str, data directory.
       mode: DatasetSplit
       shard: int, if provided, will only read data from the specified shard.
+      only_last: bool, whether we should include only files from last epoch.
 
     Returns:
       filepattern str
@@ -601,9 +602,14 @@ def dataset(self,
     # Construct the Problem's hparams so that items within it are accessible
     _ = self.get_hparams(hparams)
 
-    data_filepattern = self.filepattern(data_dir, dataset_split, shard=shard)
+    filepattern_kwargs = {
+        "data_dir": data_dir, "mode": dataset_split, "shard": shard
+    }
+    data_filepattern = self.filepattern(**filepattern_kwargs)
     if only_last:
-      imprv_data_filepattern = data_filepattern + r"10.[\d+]"
+      imprv_data_filepattern = self.filepattern(
+          only_last=True, **filepattern_kwargs
+      )
     else:
       imprv_data_filepattern = data_filepattern
     tf.logging.info("Reading data files from %s", data_filepattern)
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index e7f5a8b2a..dafe3f31f 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -56,7 +56,8 @@ def _define_batch_env(environment_spec, num_agents):
     envs = [
         environment_spec.env_lambda()
         for _ in range(num_agents)]
-    env = gym_env.T2TGymEnv(envs)
+    # We won't generate data from this env so it's safe to set data_dir to None.
+    env = gym_env.T2TGymEnv(envs, data_dir=None)
     return env
 
 
diff --git a/tensor2tensor/rl/trainer_model_based_new.py b/tensor2tensor/rl/trainer_model_based_new.py
index 6807e7614..a7db9b03d 100644
--- a/tensor2tensor/rl/trainer_model_based_new.py
+++ b/tensor2tensor/rl/trainer_model_based_new.py
@@ -141,7 +141,7 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
 
 
 def train_agent(environment_spec, agent_model_dir,
-                event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
+                event_dir, world_model_dir, data_dir, hparams, epoch=0,
                 is_final_epoch=False):
   """Train the PPO agent in the simulated environment."""
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
@@ -183,17 +183,17 @@ def train_agent(environment_spec, agent_model_dir,
       "model": hparams.generative_model,
       "hparams_set": hparams.generative_model_params,
       "output_dir": world_model_dir,
-      "data_dir": epoch_data_dir,
+      "data_dir": data_dir,
   }):
     rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
                          name_scope="ppo_sim%d" % (epoch + 1))
 
 
 def train_agent_real_env(
-    env, agent_model_dir, event_dir, epoch_data_dir,
+    env, agent_model_dir, event_dir, data_dir,
     hparams, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the real environment."""
-  del epoch_data_dir
+  del data_dir
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_params_names = ["epochs_num", "epoch_length",
                       "learning_rate", "num_agents", "eval_every_epochs",
@@ -270,10 +270,11 @@ def make_gym_env(hparams):
   return env
 
 
-def setup_env(hparams):
+def setup_env(hparams, data_dir):
   """Setup."""
   env = T2TGymEnv([make_gym_env(hparams)
                    for _ in range(hparams.real_ppo_num_agents)],
+                  data_dir,
                   grayscale=hparams.grayscale,
                   resize_width_factor=hparams.resize_width_factor,
                   resize_height_factor=hparams.resize_height_factor)
@@ -305,7 +306,8 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   directories = setup_directories(output_dir, subdirectories)
 
   epoch = -1
-  env = setup_env(hparams)
+  data_dir = directories["data"]
+  env = setup_env(hparams, data_dir)
   env.start_new_epoch(epoch)
 
   # Timing log function
@@ -313,10 +315,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   # Per-epoch state
   epoch_metrics = []
-  epoch_data_dirs = []
 
-  data_dir = os.path.join(directories["data"], "initial")
-  epoch_data_dirs.append(data_dir)
   # Collect data from the real environment with PPO or random policy.
   # TODO(lukaszkaiser): do we need option not to gather_ppo_real_env_data?
   # We could set learning_rate=0 if this flag == False.
@@ -350,18 +349,14 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   )
 
   for epoch in range(hparams.epochs):
-    epoch_data_dir = os.path.join(directories["data"], str(epoch))
-    tf.gfile.MakeDirs(epoch_data_dir)
-    env.generate_data(epoch_data_dir, directories["tmp"])
-    epoch_data_dirs.append(epoch_data_dir)
+    env.generate_data()
 
-    env.start_new_epoch(epoch)
     is_final_epoch = (epoch + 1) == hparams.epochs
     log = make_log_fn(epoch, log_relative_time)
 
     # Train world model
     log("Training world model")
-    train_world_model(env, epoch_data_dir,
+    train_world_model(env, data_dir,
                       directories["world_model"], hparams, epoch)
 
     # Train PPO
@@ -373,14 +368,16 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       ppo_model_dir = ppo_event_dir
 
     train_agent(sim_env_spec, ppo_model_dir,
-                ppo_event_dir, directories["world_model"], epoch_data_dir,
+                ppo_event_dir, directories["world_model"], data_dir,
                 hparams, epoch=epoch, is_final_epoch=is_final_epoch)
 
+    env.start_new_epoch(epoch)
+
     # Train PPO on real env (short)
     log("Training PPO in real environment.")
     train_agent_real_env(
         env, ppo_model_dir,
-        ppo_event_dir, epoch_data_dir,
+        ppo_event_dir, data_dir,
         hparams, epoch=epoch, is_final_epoch=is_final_epoch)
 
     if hparams.stop_loop_early:

From fb5082920c1a649bacb67938b8ebc13fad78604f Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Tue, 16 Oct 2018 14:53:22 -0700
Subject: [PATCH 1022/2720] internal merge of PR #1146

PiperOrigin-RevId: 217391619
---
 tensor2tensor/data_generators/gym_env.py | 22 ++++++++--------------
 tensor2tensor/data_generators/problem.py | 12 +++---------
 2 files changed, 11 insertions(+), 23 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index cc98db28d..aece79ab2 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -180,6 +180,9 @@ def step(self, actions):
     Returns:
       (obs, rewards, dones) - batches of observations, rewards and done flags
       respectively.
+
+    Raises:
+      ValueError: when the data for current epoch has already been loaded.
     """
     if self._rollouts_by_epoch_and_split[self.current_epoch]:
       raise ValueError(
@@ -387,11 +390,7 @@ def split_size(split_index):
 
   @property
   def splits_and_paths(self):
-    """List of pairs (split, paths) for the current epoch.
-
-    paths is a list of paths where data for the current epoch is saved
-    by generate_data().
-    """
+    """List of pairs (split, paths) for the current epoch."""
     filepath_fns = {
         problem.DatasetSplit.TRAIN: self.training_filepaths,
         problem.DatasetSplit.EVAL: self.dev_filepaths,
@@ -421,11 +420,7 @@ def filepattern(self, data_dir, mode, shard=None, only_last=False):
     return filepattern
 
   def generate_data(self, data_dir=None, tmp_dir=None, task_id=-1):
-    """Saves the current epoch rollouts to disk, split into train/dev sets.
-
-    data_dir and tmp_dir arguments are unused. data_dir being used is the one
-    passed in the constructor.
-    """
+    """Saves the current epoch rollouts to disk, split into train/dev sets."""
     if not self._rollouts_by_epoch_and_split[self.current_epoch]:
       # Data not loaded from disk.
       self._split_current_epoch()
@@ -491,11 +486,11 @@ def get_feature_value(key, list_name):
         }
         fields["reward"] += self.reward_range[0]
         fields["done"] = bool(fields["done"])
-        fields['observation'] = fields['image/encoded']
-        del fields['image/encoded']
+        fields["observation"] = fields["image/encoded"]
+        del fields["image/encoded"]
 
         frame = Frame(**fields)
-        frame_number = get_feature_value("frame_number", 'int64_list')
+        frame_number = get_feature_value("frame_number", "int64_list")
         if frame_number == last_frame_number + 1:
           current_rollout.append(frame)
         else:
@@ -512,7 +507,6 @@ def get_feature_value(key, list_name):
     return any_shard_empty
 
 
-
 class T2TGymEnv(T2TEnv):
   """Class representing a batch of Gym environments."""
 
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 7d9dbd225..a97d6ed12 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -420,7 +420,7 @@ def test_filepaths(self, data_dir, num_shards, shuffled):
     return generator_utils.test_data_filenames(file_basename, data_dir,
                                                num_shards)
 
-  def filepattern(self, data_dir, mode, shard=None, only_last=False):
+  def filepattern(self, data_dir, mode, shard=None):
     """Get filepattern for data files for mode.
 
     Matches mode to a suffix.
@@ -433,7 +433,6 @@ def filepattern(self, data_dir, mode, shard=None, only_last=False):
       data_dir: str, data directory.
       mode: DatasetSplit
       shard: int, if provided, will only read data from the specified shard.
-      only_last: bool, whether we should include only files from last epoch.
 
     Returns:
       filepattern str
@@ -602,14 +601,9 @@ def dataset(self,
     # Construct the Problem's hparams so that items within it are accessible
     _ = self.get_hparams(hparams)
 
-    filepattern_kwargs = {
-        "data_dir": data_dir, "mode": dataset_split, "shard": shard
-    }
-    data_filepattern = self.filepattern(**filepattern_kwargs)
+    data_filepattern = self.filepattern(data_dir, dataset_split, shard=shard)
     if only_last:
-      imprv_data_filepattern = self.filepattern(
-          only_last=True, **filepattern_kwargs
-      )
+      imprv_data_filepattern = data_filepattern + r"10.[\d+]"
     else:
       imprv_data_filepattern = data_filepattern
     tf.logging.info("Reading data files from %s", data_filepattern)

From 4084432cd15345747e43ecf987d11f9b8d803318 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 16 Oct 2018 15:49:48 -0700
Subject: [PATCH 1023/2720] Fix Revnet.

PiperOrigin-RevId: 217402136
---
 tensor2tensor/models/revnet.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/revnet.py b/tensor2tensor/models/revnet.py
index 379c2b143..036c85840 100644
--- a/tensor2tensor/models/revnet.py
+++ b/tensor2tensor/models/revnet.py
@@ -91,7 +91,7 @@ def f(x, depth1, depth2, dim='2d', first_batch_norm=True, stride=1,
     Output tensor after applying residual function for RevNet.
   """
   conv = CONFIG[dim]['conv']
-  with tf.variable_scope('f'):
+  with tf.variable_scope('f', reuse=tf.AUTO_REUSE):
     if first_batch_norm:
       net = tf.layers.batch_normalization(x, training=training)
       net = tf.nn.relu(net)

From 3f43417310101859f95b74587ffc3686714cc58a Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Tue, 16 Oct 2018 18:39:25 -0700
Subject: [PATCH 1024/2720] Add checkpoint_path flag to T2T export.py

PiperOrigin-RevId: 217425249
---
 tensor2tensor/serving/export.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 7fd55f2be..54e153f7c 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -38,6 +38,11 @@
     "If None, the model will be stored in subdirectory "
     "where checkpoints are: --output_dir")
 
+tf.flags.DEFINE_string(
+    "checkpoint_path", None, "Which checkpoint to export."
+    "If None, we will use the latest checkpoint stored in the directory "
+    "specified by --output_dir")
+
 
 def create_estimator(run_config, hparams):
   return trainer_lib.create_estimator(
@@ -134,7 +139,12 @@ def main(_):
   trainer_lib.set_random_seed(FLAGS.random_seed)
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
 
-  ckpt_dir = os.path.expanduser(FLAGS.output_dir)
+  if FLAGS.checkpoint_path:
+    checkpoint_path = FLAGS.checkpoint_path
+    ckpt_dir = os.path.dirname(checkpoint_path)
+  else:
+    ckpt_dir = os.path.expanduser(FLAGS.output_dir)
+    checkpoint_path = tf.train.latest_checkpoint(ckpt_dir)
 
   hparams = create_hparams()
   hparams.no_data_parallelism = True  # To clear the devices
@@ -159,7 +169,7 @@ def main(_):
   exporter.export(
       estimator,
       export_dir,
-      checkpoint_path=tf.train.latest_checkpoint(ckpt_dir),
+      checkpoint_path=checkpoint_path,
       eval_result=None,
       is_the_final_export=True)
 

From b7726cc3b9722e98db03d4b5b8c1ade6fccbb186 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 16 Oct 2018 22:07:38 -0700
Subject: [PATCH 1025/2720] fix the splitting of input and predicted frames.

PiperOrigin-RevId: 217441325
---
 tensor2tensor/models/video/epva.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index 6bcd66dc3..53c33cee2 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -713,7 +713,7 @@ def fix_video_dims_and_concat_on_x_axis(x):
 
     predictions = common_video.swap_time_and_batch_axes(predictions)
     predictions = tf.slice(predictions,
-                           [0, hparams.video_num_target_frames-1, 0, 0, 0],
+                           [0, hparams.video_num_input_frames-1, 0, 0, 0],
                            [-1]*5)
 
     return predictions, {'extra': epva_loss}

From d34ea12eb046e8efcbf226089606be12f7c6e19b Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 17 Oct 2018 02:02:04 -0700
Subject: [PATCH 1026/2720] Swap model-based RL trainer for the new version (it
 works).

PiperOrigin-RevId: 217456149
---
 tensor2tensor/data_generators/gym_env.py      |   4 +-
 tensor2tensor/data_generators/gym_env_test.py |  14 +-
 tensor2tensor/rl/trainer_model_based.py       | 717 +++++-------------
 .../rl/trainer_model_based_ae_test.py         |  38 -
 tensor2tensor/rl/trainer_model_based_new.py   | 418 ----------
 .../rl/trainer_model_based_params.py          |   2 +-
 6 files changed, 193 insertions(+), 1000 deletions(-)
 delete mode 100644 tensor2tensor/rl/trainer_model_based_ae_test.py
 delete mode 100644 tensor2tensor/rl/trainer_model_based_new.py

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index aece79ab2..c43fb90a5 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -204,9 +204,7 @@ def step(self, actions):
         Frame(*orud, action=None)
         for orud in zip(encoded_obs, rewards, unclipped_rewards, dones)
     ]
-    # TODO(lukaszkaiser): changed unclipped_reward to reward once we've
-    # removed the current setup with RewardClippingWrapper and so on.
-    return (obs, unclipped_rewards, dones)
+    return (obs, rewards, dones)
 
   def _reset(self, indices):
     """Resets environments at given indices without recording history.
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index b283844e1..61b3bd24d 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -203,18 +203,10 @@ def test_frame_numbers_are_continuous(self):
       last_frame_number = frame_number
 
   def test_clipping(self):
-    # This test needs base env with rewards out of [-1,1] range.
     env_lambda = TestEnv
-    # TODO(lukaszkaiser): turn clipping on by default after refactor.
-    # _, _, rewards = self.init_batch_and_play(env_lambda, n_steps=2)
-    # self.assertTrue(np.max(rewards) == 1)
-    # self.assertTrue(np.min(rewards) == -1)
-
-    _, _, unclipped_rewards, _ = self.init_batch_and_play(
-        env_lambda, steps_per_epoch=2
-    )
-    self.assertTrue(np.max(unclipped_rewards) > 1)
-    self.assertTrue(np.min(unclipped_rewards) < -1)
+    _, _, rewards, _ = self.init_batch_and_play(env_lambda, steps_per_epoch=2)
+    self.assertTrue(np.max(rewards) == 1)
+    self.assertTrue(np.min(rewards) == -1)
 
   def test_resize(self):
     env_lambda = TestEnv
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 24d2dcfdc..90fb103b7 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -17,7 +17,7 @@
 
 Example invocation:
 
-python -m tensor2tensor.rl.trainer_model_based \
+python -m tensor2tensor.rl.trainer_model_based_new \
     --output_dir=$HOME/t2t/rl_v1 \
     --loop_hparams_set=rlmb_base \
     --loop_hparams='num_real_env_frames=10000,epochs=3'
@@ -31,23 +31,17 @@
 import datetime
 import math
 import os
-import re
 import time
 
+import gym
 import numpy as np
-import six
 
-
-from tensor2tensor.bin import t2t_trainer
-from tensor2tensor.data_generators import generator_utils
-from tensor2tensor.data_generators import gym_problems
-from tensor2tensor.data_generators import gym_problems_specs
-from tensor2tensor.layers import discretization
+from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
+from tensor2tensor.data_generators.gym_env import T2TGymEnv
+from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.rl import trainer_model_based_params
-from tensor2tensor.rl.envs.tf_atari_wrappers import PyFuncWrapper
 from tensor2tensor.rl.envs.utils import InitialFrameChooser
-from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -57,6 +51,43 @@
 FLAGS = flags.FLAGS
 
 
+@contextlib.contextmanager
+def temporary_flags(flag_settings):
+  old_values = {}
+  for flag_name, flag_value in flag_settings.items():
+    old_values[flag_name] = getattr(FLAGS, flag_name)
+    setattr(FLAGS, flag_name, flag_value)
+  yield
+  for flag_name, flag_value in old_values.items():
+    setattr(FLAGS, flag_name, flag_value)
+
+
+def _ppo_training_epochs(hparams, epoch, is_final_epoch, real_env_training):
+  """Helper for PPO restarts."""
+  if hparams.gather_ppo_real_env_data:
+    assert hparams.real_ppo_epochs_num is 0, (
+        "Should be put to 0 to enforce better readability")
+    real_training_ppo_epochs_num = int(math.ceil(
+        hparams.num_real_env_frames /
+        (hparams.epochs*hparams.real_ppo_epoch_length)))
+  else:
+    real_training_ppo_epochs_num = hparams.real_ppo_epochs_num
+
+  simulated_training_ppo_epochs_num = hparams.ppo_epochs_num
+
+  if epoch == -1:
+    assert real_env_training, (
+        "Epoch -1 should only be used for PPO collection in real environment.")
+    return real_training_ppo_epochs_num
+  ppo_training_epochs = (epoch + 1) * (simulated_training_ppo_epochs_num
+                                       + real_training_ppo_epochs_num)
+  if is_final_epoch:  # Length of training in the final epoch is doubled.
+    ppo_training_epochs += simulated_training_ppo_epochs_num
+  if real_env_training:
+    ppo_training_epochs += real_training_ppo_epochs_num
+  return ppo_training_epochs
+
+
 def setup_directories(base_dir, subdirs):
   base_dir = os.path.expanduser(base_dir)
   tf.gfile.MakeDirs(base_dir)
@@ -83,46 +114,6 @@ def log_relative_time():
   return log_relative_time
 
 
-@contextlib.contextmanager
-def temporary_flags(flag_settings):
-  old_values = {}
-  for flag_name, flag_value in flag_settings.items():
-    old_values[flag_name] = getattr(FLAGS, flag_name)
-    setattr(FLAGS, flag_name, flag_value)
-  yield
-  for flag_name, flag_value in old_values.items():
-    setattr(FLAGS, flag_name, flag_value)
-
-
-def generate_real_env_data(problem_name, agent_policy_path, hparams, data_dir,
-                           tmp_dir, autoencoder_path=None, eval_phase=False,
-                           real_reward=False):
-  """Run the agent against the real environment and return mean reward."""
-  tf.gfile.MakeDirs(data_dir)
-  with temporary_flags({
-      "problem": problem_name,
-      "agent_policy_path": agent_policy_path,
-      "autoencoder_path": autoencoder_path,
-  }):
-    gym_problem = registry.problem(problem_name)
-    if hparams.gather_ppo_real_env_data:
-      env_steps_per_epoch = int(hparams.num_real_env_frames / hparams.epochs)
-    else:
-      env_steps_per_epoch = (
-          hparams.num_real_env_frames / (hparams.epochs * (1. - 1./11.)))
-    gym_problem.settable_num_steps = env_steps_per_epoch
-    gym_problem.settable_eval_phase = eval_phase
-    if real_reward and autoencoder_path is None:
-      gym_problem._forced_collect_level = 1  # pylint: disable=protected-access
-    gym_problem.generate_data(data_dir, tmp_dir)
-    mean_reward = None
-    if gym_problem.statistics.number_of_dones:
-      mean_reward = (gym_problem.statistics.sum_of_rewards /
-                     gym_problem.statistics.number_of_dones)
-
-  return mean_reward
-
-
 def make_log_fn(epoch, log_relative_time_fn):
 
   def log(msg, *args):
@@ -133,53 +124,26 @@ def log(msg, *args):
   return log
 
 
-def train_autoencoder(problem_name, data_dir, output_dir, hparams, epoch):
-  """Train autoencoder on problem_name."""
-  additional_steps = 1 + hparams.autoencoder_train_steps_initial_multiplier
-  train_steps = hparams.autoencoder_train_steps * (epoch + additional_steps)
-  with temporary_flags({
-      "problem": problem_name,
-      "data_dir": data_dir,
-      "output_dir": output_dir,
-      "model": "autoencoder_ordered_discrete",
-      "hparams_set": hparams.autoencoder_hparams_set,
-      "train_steps": train_steps,
-      "eval_steps": 100,
-  }):
-    t2t_trainer.main([])
-
-
-def _ppo_training_epochs(hparams, epoch, is_final_epoch, real_env_training):
-  """Helper for PPO restarts."""
-  if hparams.gather_ppo_real_env_data:
-    assert hparams.real_ppo_epochs_num is 0, (
-        "Should be put to 0 to enforce better readability")
-    real_training_ppo_epochs_num = int(math.ceil(
-        hparams.num_real_env_frames /
-        (hparams.epochs*hparams.real_ppo_epoch_length)))
-  else:
-    real_training_ppo_epochs_num = hparams.real_ppo_epochs_num
-
-  simulated_training_ppo_epochs_num = hparams.ppo_epochs_num
+def train_supervised(problem, model_name, hparams, data_dir, output_dir,
+                     train_steps, eval_steps, local_eval_frequency=None,
+                     schedule="continuous_train_and_eval"):
+  """Train supervised."""
+  if local_eval_frequency is None:
+    local_eval_frequency = getattr(FLAGS, "local_eval_frequency")
 
-  if epoch == -1:
-    assert real_env_training, (
-        "Epoch -1 should only be used for PPO collection in real environment.")
-    return real_training_ppo_epochs_num
-  ppo_training_epochs = (epoch + 1) * (simulated_training_ppo_epochs_num
-                                       + real_training_ppo_epochs_num)
-  if is_final_epoch:  # Length of training in the final epoch is doubled.
-    ppo_training_epochs += simulated_training_ppo_epochs_num
-  if real_env_training:
-    ppo_training_epochs += real_training_ppo_epochs_num
-  return ppo_training_epochs
+  exp_fn = trainer_lib.create_experiment_fn(
+      model_name, problem, data_dir, train_steps, eval_steps,
+      min_eval_frequency=local_eval_frequency
+  )
+  run_config = trainer_lib.create_run_config(model_name, model_dir=output_dir)
+  exp = exp_fn(run_config, hparams)
+  getattr(exp, schedule)()
 
 
-def train_agent(problem_name, agent_model_dir,
-                event_dir, world_model_dir, epoch_data_dir, hparams, epoch=0,
+def train_agent(environment_spec, agent_model_dir,
+                event_dir, world_model_dir, data_dir, hparams, epoch=0,
                 is_final_epoch=False):
   """Train the PPO agent in the simulated environment."""
-  gym_problem = registry.problem(problem_name)
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_params_names = ["epochs_num", "epoch_length",
                       "learning_rate", "num_agents",
@@ -200,57 +164,36 @@ def train_agent(problem_name, agent_model_dir,
   model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
   ppo_hparams.add_hparam("model_hparams", model_hparams)
 
-  environment_spec = copy.copy(gym_problem.environment_spec)
-  environment_spec.simulation_random_starts = hparams.simulation_random_starts
-  do_flip = hparams.simulation_flip_first_random_for_beginning
-  environment_spec.simulation_flip_first_random_for_beginning = do_flip
-  environment_spec.intrinsic_reward_scale = hparams.intrinsic_reward_scale
-
+  environment_spec = copy.copy(environment_spec)
+  environment_spec_param_names = [
+      "simulation_random_starts", "simulation_flip_first_random_for_beginning",
+      "intrinsic_reward_scale"
+  ]
+  for param_name in environment_spec_param_names:
+    environment_spec.set_hparam(param_name, hparams.get(param_name))
   ppo_hparams.add_hparam("environment_spec", environment_spec)
+
   ppo_hparams.add_hparam("initial_frame_chooser", InitialFrameChooser(
-      environment_spec, mode=tf.estimator.ModeKeys.TRAIN
+      environment_spec, mode=tf.estimator.ModeKeys.EVAL
   ))
 
+  # TODO(koz4k): Pass by arguments.
   with temporary_flags({
-      "problem": problem_name,
+      "problem": environment_spec.initial_frames_problem,
       "model": hparams.generative_model,
       "hparams_set": hparams.generative_model_params,
       "output_dir": world_model_dir,
-      "data_dir": epoch_data_dir,
+      "data_dir": data_dir,
   }):
     rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
                          name_scope="ppo_sim%d" % (epoch + 1))
 
-ppo_data_dumper_counter = 0
-dumper_path = None
-
-
-def ppo_data_dumper(observ, reward, done, action):
-  """Save frames from PPO to a numpy file."""
-  global ppo_data_dumper_counter, dumper_path
-  file_path = "{}/frame_{}.npz".format(dumper_path, ppo_data_dumper_counter)
-  if gym_problems.frame_dumper_use_disk:
-    # np.savez_compressed can't create a tf.gfile, so we need to create it
-    # beforehand.
-    with tf.gfile.Open(file_path, mode="wb+") as gfile:
-      gfile.write("1")
-    with tf.gfile.Open(file_path, mode="wb+") as gfile:
-      np.savez_compressed(
-          gfile, observ=observ, reward=reward, done=done, action=action)
-  else:
-    data = {"observ": observ, "reward": reward, "done": done, "action": action}
-    gym_problems.frame_dumper[file_path] = data
-  ppo_data_dumper_counter += 1
-  return 0.0
-
 
 def train_agent_real_env(
-    problem_name, agent_model_dir, event_dir, world_model_dir, epoch_data_dir,
+    env, agent_model_dir, event_dir, data_dir,
     hparams, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the real environment."""
-  global dumper_path, ppo_data_dumper_counter
-
-  gym_problem = registry.problem(problem_name)
+  del data_dir
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_params_names = ["epochs_num", "epoch_length",
                       "learning_rate", "num_agents", "eval_every_epochs",
@@ -269,225 +212,81 @@ def train_agent_real_env(
   # But we need to save at the last step, so we set it very high.
   ppo_hparams.save_models_every_epochs = 1000000
 
-  environment_spec = copy.copy(gym_problem.environment_spec)
-
-  if hparams.gather_ppo_real_env_data:
-    # TODO(piotrmilos):This should be refactored
-    assert hparams.real_ppo_num_agents == 1, (
-        "It is required to use collect with pyfunc_wrapper")
-
-    ppo_data_dumper_counter = 0
-    dumper_path = os.path.join(epoch_data_dir, "dumper")
-    tf.gfile.MakeDirs(dumper_path)
-    dumper_spec = [PyFuncWrapper, {"process_fun": ppo_data_dumper}]
-    environment_spec.wrappers.insert(2, dumper_spec)
+  environment_spec = rl.standard_atari_env_spec(
+      batch_env=env, include_clipping=False
+  )
 
   ppo_hparams.add_hparam("environment_spec", environment_spec)
 
-  with temporary_flags({
-      "problem": problem_name,
-      "output_dir": world_model_dir,
-      "data_dir": epoch_data_dir,
-  }):
-    rl_trainer_lib.train(ppo_hparams, event_dir + "real", agent_model_dir,
-                         name_scope="ppo_real%d" % (epoch + 1))
-
-
-def evaluate_world_model(simulated_problem_name, problem_name, hparams,
-                         world_model_dir, epoch_data_dir, tmp_dir):
-  """Generate simulated environment data and return reward accuracy."""
-  gym_simulated_problem = registry.problem(simulated_problem_name)
-  sim_steps = hparams.simulated_env_generator_num_steps
-  gym_simulated_problem.settable_num_steps = sim_steps
-  gym_simulated_problem.settable_rollout_fractions = \
-      hparams.eval_rollout_fractions
-  with temporary_flags({
-      "problem": problem_name,
-      "model": hparams.generative_model,
-      "hparams_set": hparams.generative_model_params,
-      "data_dir": epoch_data_dir,
-      "output_dir": world_model_dir,
-  }):
-    gym_simulated_problem.generate_data(epoch_data_dir, tmp_dir)
-  n = max(1., gym_simulated_problem.statistics.number_of_dones)
-  model_reward_accuracy = [
-      (frac, score / float(n))
-      for (frac, score) in six.iteritems(
-          gym_simulated_problem.statistics.
-          successful_episode_reward_predictions
-      )
-  ]
-  old_path = os.path.join(epoch_data_dir, "debug_frames_sim")
-  new_path = os.path.join(epoch_data_dir, "debug_frames_sim_eval")
-  if not tf.gfile.Exists(new_path):
-    tf.gfile.Rename(old_path, new_path)
-  return model_reward_accuracy
+  rl_trainer_lib.train(ppo_hparams, event_dir + "real", agent_model_dir,
+                       name_scope="ppo_real%d" % (epoch + 1))
+
+  # Save unfinished rollouts to history.
+  env.reset()
 
 
-def train_world_model(problem_name, data_dir, output_dir, hparams, epoch):
+def train_world_model(env, data_dir, output_dir, hparams, epoch):
   """Train the world model on problem_name."""
   train_steps = hparams.model_train_steps * (
       epoch + hparams.inital_epoch_train_steps_multiplier)
   model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
-  learning_rate = model_hparams.learning_rate_constant
-  if epoch > 0: learning_rate *= hparams.learning_rate_bump
-  with temporary_flags({
-      "data_dir": data_dir,
-      "output_dir": output_dir,
-      "problem": problem_name,
-      "model": hparams.generative_model,
-      "hparams_set": hparams.generative_model_params,
-      "hparams": "learning_rate_constant=%.6f" % learning_rate,
-      "eval_steps": 100,
-      "local_eval_frequency": 2000,
-      "train_steps": train_steps,
-  }):
-    t2t_trainer.main([])
-
-
-def encode_dataset(model, dataset, problem, ae_hparams, autoencoder_path,
-                   out_files):
-  """Encode all frames in dataset with model and write them out to out_files."""
-  batch_size = 8
-  dataset = dataset.batch(batch_size)
-  examples = dataset.make_one_shot_iterator().get_next()
-  images = examples.pop("frame")
-  images = tf.cast(images, tf.int32)
-
-  encoded = model.encode(images)
-  encoded_frame_height = int(
-      math.ceil(problem.frame_height / 2**ae_hparams.num_hidden_layers))
-  encoded_frame_width = int(
-      math.ceil(problem.frame_width / 2**ae_hparams.num_hidden_layers))
-  num_bits = 8
-  encoded = tf.reshape(
-      encoded, [-1, encoded_frame_height, encoded_frame_width, 3, num_bits])
-  encoded = tf.cast(discretization.bit_to_int(encoded, num_bits), tf.uint8)
-
-  pngs = tf.map_fn(tf.image.encode_png, encoded, dtype=tf.string,
-                   back_prop=False)
-
-  with tf.Session() as sess:
-    autoencoder_saver = tf.train.Saver(tf.global_variables("autoencoder.*"))
-    trainer_lib.restore_checkpoint(autoencoder_path, autoencoder_saver, sess,
-                                   must_restore=True)
-
-    def generator():
-      """Generate examples."""
-      while True:
-        try:
-          pngs_np, examples_np = sess.run([pngs, examples])
-          rewards = examples_np["reward"].tolist()
-          actions = examples_np["action"].tolist()
-          frame_numbers = examples_np["frame_number"].tolist()
-          for action, reward, frame_number, png in \
-                  zip(actions, rewards, frame_numbers, pngs_np):
-            yield {
-                "action": action,
-                "reward": reward,
-                "frame_number": frame_number,
-                "image/encoded": [png],
-                "image/format": ["png"],
-                "image/height": [encoded_frame_height],
-                "image/width": [encoded_frame_width],
-            }
-        except tf.errors.OutOfRangeError:
-          break
-
-    generator_utils.generate_files(
-        generator(), out_files,
-        cycle_every_n=problem.total_number_of_frames // 10)
-
-
-def encode_env_frames(problem_name, ae_problem_name, ae_hparams_set,
-                      autoencoder_path, epoch_data_dir):
-  """Encode all frames from problem_name and write out as ae_problem_name."""
-  with tf.Graph().as_default():
-    ae_hparams = trainer_lib.create_hparams(ae_hparams_set,
-                                            problem_name=problem_name)
-    problem = ae_hparams.problem
-    model = registry.model("autoencoder_ordered_discrete")(
-        ae_hparams, tf.estimator.ModeKeys.EVAL)
-
-    ae_problem = registry.problem(ae_problem_name)
-    ae_training_paths = ae_problem.training_filepaths(epoch_data_dir, 10, True)
-    ae_eval_paths = ae_problem.dev_filepaths(epoch_data_dir, 1, True)
-
-    skip_train = False
-    skip_eval = False
-    for path in ae_training_paths:
-      if tf.gfile.Exists(path):
-        skip_train = True
-        break
-    for path in ae_eval_paths:
-      if tf.gfile.Exists(path):
-        skip_eval = True
-        break
-
-    # Encode train data
-    if not skip_train:
-      dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, epoch_data_dir,
-                                shuffle_files=False, output_buffer_size=100,
-                                preprocess=False)
-      encode_dataset(model, dataset, problem=problem, ae_hparams=ae_hparams,
-                     autoencoder_path=autoencoder_path,
-                     out_files=ae_training_paths)
-
-    # Encode eval data
-    if not skip_eval:
-      dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, epoch_data_dir,
-                                shuffle_files=False, output_buffer_size=100,
-                                preprocess=False)
-      encode_dataset(model, dataset, problem=problem, ae_hparams=ae_hparams,
-                     autoencoder_path=autoencoder_path, out_files=ae_eval_paths)
-
-
-def check_problems(problem_names):
-  for problem_name in problem_names:
-    registry.problem(problem_name)
-
-
-def setup_problems(hparams, using_autoencoder=False):
-  """Register problems based on game name."""
-  if hparams.game in gym_problems_specs.ATARI_GAMES:
-    game_with_mode = hparams.game + "_deterministic-v4"
-  else:
-    game_with_mode = hparams.game
-  game_problems_kwargs = {}
-  # Problems
-  if using_autoencoder:
-    game_problems_kwargs["autoencoder_hparams"] = (
-        hparams.autoencoder_hparams_set)
-    problem_name = (
-        "gym_discrete_problem_with_agent_on_%s_with_autoencoder"
-        % game_with_mode)
-    world_model_problem = (
-        "gym_discrete_problem_with_agent_on_%s_autoencoded" % game_with_mode)
-    simulated_problem_name = (
-        "gym_simulated_discrete_problem_with_agent_on_%s_autoencoded"
-        % game_with_mode)
-    world_model_eval_problem_name = (
-        "gym_simulated_discrete_problem_for_world_model_eval_with_agent_on_%s"
-        "_autoencoded"
-        % game_with_mode)
+  model_hparams.learning_rate = model_hparams.learning_rate_constant
+  if epoch > 0:
+    model_hparams.learning_rate *= hparams.learning_rate_bump
+
+  train_supervised(
+      problem=env,
+      model_name=hparams.generative_model,
+      hparams=model_hparams,
+      data_dir=data_dir,
+      output_dir=output_dir,
+      train_steps=train_steps,
+      eval_steps=100,
+      local_eval_frequency=2000
+  )
+
+
+def make_gym_env(hparams):
+  """Make env."""
+  game_mode = "Deterministic-v4"
+  camel_game_name = "".join(
+      [w[0].upper() + w[1:] for w in hparams.game.split("_")])
+  camel_game_name += game_mode
+  env_name = camel_game_name
+  env = gym.make(env_name)
+  if hparams.env_timesteps_limit != -1:
+    # Replace TimeLimit Wrapper with one of proper time step limit.
+    if isinstance(env, gym.wrappers.TimeLimit):
+      env = env.env
+    env = gym.wrappers.TimeLimit(env,
+                                 max_episode_steps=hparams.env_timesteps_limit)
+  return env
+
+
+def setup_env(hparams, data_dir):
+  """Setup."""
+  env = T2TGymEnv([make_gym_env(hparams)
+                   for _ in range(hparams.real_ppo_num_agents)],
+                  data_dir,
+                  grayscale=hparams.grayscale,
+                  resize_width_factor=hparams.resize_width_factor,
+                  resize_height_factor=hparams.resize_height_factor)
+  return env
+
+
+def eval_reward(env, clipped):
+  """Calculates mean rewards from given epoch."""
+  reward_name = "reward" if clipped else "unclipped_reward"
+  rewards = []
+  for rollout in env.current_epoch_rollouts():
+    if rollout[-1].done:
+      rollout_reward = sum(getattr(frame, reward_name) for frame in rollout)
+      rewards.append(rollout_reward)
+  if rewards:
+    mean_rewards = np.mean(rewards)
   else:
-    problem_name = ("gym_discrete_problem_with_agent_on_%s" % game_with_mode)
-    world_model_problem = problem_name
-    simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
-                              % game_with_mode)
-    world_model_eval_problem_name = (
-        "gym_simulated_discrete_problem_for_world_model_eval_with_agent_on_%s"
-        % game_with_mode)
-  if problem_name not in registry.list_problems():
-    game_problems_kwargs["resize_height_factor"] = hparams.resize_height_factor
-    game_problems_kwargs["resize_width_factor"] = hparams.resize_width_factor
-    game_problems_kwargs["grayscale"] = hparams.grayscale
-    tf.logging.info("Game Problem %s not found; dynamically registering",
-                    problem_name)
-    gym_problems_specs.create_problems_for_game(
-        hparams.game, game_mode="Deterministic-v4", **game_problems_kwargs)
-  return (problem_name, world_model_problem, simulated_problem_name,
-          world_model_eval_problem_name)
+    mean_rewards = 0
+  return mean_rewards
 
 
 def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
@@ -495,108 +294,64 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   if report_fn:
     assert report_metric is not None
 
-  # Global state
-
   # Directories
   subdirectories = ["data", "tmp", "world_model", "ppo"]
-  using_autoencoder = hparams.autoencoder_train_steps > 0
-  if using_autoencoder:
-    subdirectories.append("autoencoder")
   directories = setup_directories(output_dir, subdirectories)
 
-  (problem_name, world_model_problem, simulated_problem_name,
-   world_model_eval_problem_name) = setup_problems(hparams, using_autoencoder)
-
-  # Autoencoder model dir
-  autoencoder_model_dir = directories.get("autoencoder")
+  epoch = -1
+  data_dir = directories["data"]
+  env = setup_env(hparams, data_dir)
+  env.start_new_epoch(epoch)
 
   # Timing log function
   log_relative_time = make_relative_timing_fn()
 
   # Per-epoch state
   epoch_metrics = []
-  epoch_data_dirs = []
 
-  ppo_model_dir = None
-  data_dir = os.path.join(directories["data"], "initial")
-  epoch_data_dirs.append(data_dir)
   # Collect data from the real environment with PPO or random policy.
-  if hparams.gather_ppo_real_env_data:
-    ppo_model_dir = directories["ppo"]
-    tf.logging.info("Initial training of PPO in real environment.")
-    ppo_event_dir = os.path.join(directories["world_model"],
-                                 "ppo_summaries/initial")
-    train_agent_real_env(
-        problem_name, ppo_model_dir,
-        ppo_event_dir, directories["world_model"], data_dir,
-        hparams, epoch=-1, is_final_epoch=False)
-
-  tf.logging.info("Generating real environment data with %s policy",
-                  "PPO" if hparams.gather_ppo_real_env_data else "random")
-  mean_reward = generate_real_env_data(
-      problem_name, ppo_model_dir, hparams, data_dir, directories["tmp"])
-  tf.logging.info("Mean reward (random): {}".format(mean_reward))
+  # TODO(lukaszkaiser): do we need option not to gather_ppo_real_env_data?
+  # We could set learning_rate=0 if this flag == False.
+  assert hparams.gather_ppo_real_env_data
+  ppo_model_dir = directories["ppo"]
+  tf.logging.info("Initial training of PPO in real environment.")
+  ppo_event_dir = os.path.join(directories["world_model"],
+                               "ppo_summaries/initial")
+  train_agent_real_env(
+      env, ppo_model_dir,
+      ppo_event_dir, data_dir,
+      hparams, epoch=epoch, is_final_epoch=False)
+  mean_unclipped_reward = eval_reward(env, clipped=False)
+  tf.logging.info("Mean reward (initial): {}".format(mean_unclipped_reward))
 
   eval_metrics_event_dir = os.path.join(directories["world_model"],
                                         "eval_metrics_event_dir")
   eval_metrics_writer = tf.summary.FileWriter(eval_metrics_event_dir)
-  model_reward_accuracy_summary = tf.Summary()
-  for frac in hparams.eval_rollout_fractions:
-    model_reward_accuracy_summary.value.add(
-        tag="model_reward_accuracy_{}".format(frac),
-        simple_value=None
-    )
-  mean_reward_summary = tf.Summary()
-  mean_reward_summary.value.add(tag="mean_reward",
-                                simple_value=None)
-  mean_reward_gen_summary = tf.Summary()
-  mean_reward_gen_summary.value.add(tag="mean_reward_during_generation",
-                                    simple_value=None)
+
+  mean_unclipped_reward_summary = tf.Summary()
+  mean_unclipped_reward_summary.value.add(tag="mean_unclipped_reward",
+                                          simple_value=None)
+  mean_clipped_reward_summary = tf.Summary()
+  mean_clipped_reward_summary.value.add(tag="mean_clipped_reward",
+                                        simple_value=None)
+
+  sim_env_spec = rl.standard_atari_env_simulated_spec(
+      env,
+      # Hardcoded for now. TODO(koz4k): Make it a hparam.
+      video_num_input_frames=4, video_num_target_frames=1
+  )
 
   for epoch in range(hparams.epochs):
+    env.generate_data()
+
     is_final_epoch = (epoch + 1) == hparams.epochs
     log = make_log_fn(epoch, log_relative_time)
 
-    # Combine all previously collected environment data
-    epoch_data_dir = os.path.join(directories["data"], str(epoch))
-    tf.gfile.MakeDirs(epoch_data_dir)
-    # Because the data is being combined in every iteration, we only need to
-    # copy from the previous directory.
-    combine_training_data(registry.problem(problem_name),
-                          epoch_data_dir,
-                          epoch_data_dirs[-1:])
-    epoch_data_dirs.append(epoch_data_dir)
-
-    if using_autoencoder:
-      # Train the Autoencoder on all prior environment frames
-      log("Training Autoencoder")
-      train_autoencoder(problem_name, epoch_data_dir, autoencoder_model_dir,
-                        hparams, epoch)
-
-      log("Autoencoding environment frames")
-      encode_env_frames(problem_name, world_model_problem,
-                        ae_hparams_set=hparams.autoencoder_hparams_set,
-                        autoencoder_path=autoencoder_model_dir,
-                        epoch_data_dir=epoch_data_dir)
-
     # Train world model
     log("Training world model")
-    train_world_model(world_model_problem, epoch_data_dir,
+    train_world_model(env, data_dir,
                       directories["world_model"], hparams, epoch)
 
-    # Evaluate world model
-    model_reward_accuracy = []
-    if hparams.eval_world_model:
-      log("Evaluating world model")
-      model_reward_accuracy = evaluate_world_model(
-          world_model_eval_problem_name, world_model_problem, hparams,
-          directories["world_model"],
-          epoch_data_dir, directories["tmp"])
-      log(
-          "World model reward accuracy per rollout fraction: %s",
-          model_reward_accuracy
-      )
-
     # Train PPO
     log("Training PPO in simulated environment.")
     ppo_event_dir = os.path.join(directories["world_model"],
@@ -604,61 +359,38 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     ppo_model_dir = directories["ppo"]
     if not hparams.ppo_continue_training:
       ppo_model_dir = ppo_event_dir
-    train_agent(simulated_problem_name, ppo_model_dir,
-                ppo_event_dir, directories["world_model"], epoch_data_dir,
+
+    train_agent(sim_env_spec, ppo_model_dir,
+                ppo_event_dir, directories["world_model"], data_dir,
                 hparams, epoch=epoch, is_final_epoch=is_final_epoch)
 
+    env.start_new_epoch(epoch)
+
     # Train PPO on real env (short)
     log("Training PPO in real environment.")
     train_agent_real_env(
-        problem_name, ppo_model_dir,
-        ppo_event_dir, directories["world_model"], epoch_data_dir,
+        env, ppo_model_dir,
+        ppo_event_dir, data_dir,
         hparams, epoch=epoch, is_final_epoch=is_final_epoch)
 
     if hparams.stop_loop_early:
       return 0.0
+    mean_clipped_reward = eval_reward(env, clipped=True)
+    log("Mean clipped reward during generation: {}".format(
+        mean_clipped_reward))  # this was output of generate_real_env_data(...)
 
-    # Collect data from the real environment.
-    generation_mean_reward = None
-    if not is_final_epoch:
-      log("Generating real environment data.")
-      generation_mean_reward = generate_real_env_data(
-          problem_name, ppo_model_dir, hparams, epoch_data_dir,
-          directories["tmp"], autoencoder_path=autoencoder_model_dir,
-          eval_phase=False)
-      log("Mean clipped reward during generation: {}".format(
-          generation_mean_reward))
-
-    log("Evaluating in real environment.")
-    eval_data_dir = os.path.join(epoch_data_dir, "eval")
-    mean_reward = generate_real_env_data(
-        problem_name, ppo_model_dir, hparams, eval_data_dir,
-        directories["tmp"], autoencoder_path=autoencoder_model_dir,
-        eval_phase=True, real_reward=True)
-    log("Mean eval reward (unclipped): {}".format(mean_reward))
+    mean_unclipped_reward = eval_reward(env, clipped=False)
+    log("Mean eval reward (unclipped): {}".format(mean_unclipped_reward))
 
     # Summarize metrics.
-    assert model_reward_accuracy is not None
-    assert mean_reward is not None
-    for ((_, accuracy), summary_value) in zip(
-        model_reward_accuracy, model_reward_accuracy_summary.value
-    ):
-      summary_value.simple_value = accuracy
-    mean_reward_summary.value[0].simple_value = mean_reward
-    eval_metrics_writer.add_summary(model_reward_accuracy_summary, epoch)
-    eval_metrics_writer.add_summary(mean_reward_summary, epoch)
-    if generation_mean_reward is not None:
-      mean_reward_gen_summary.value[0].simple_value = int(
-          generation_mean_reward)
-      eval_metrics_writer.add_summary(mean_reward_gen_summary, epoch)
+    mean_unclipped_reward_summary.value[0].simple_value = mean_unclipped_reward
+    eval_metrics_writer.add_summary(mean_unclipped_reward_summary, epoch)
+    mean_clipped_reward_summary.value[0].simple_value = int(mean_clipped_reward)
+    eval_metrics_writer.add_summary(mean_clipped_reward_summary, epoch)
     eval_metrics_writer.flush()
 
     # Report metrics
-    eval_metrics = {"mean_reward": mean_reward}
-    eval_metrics.update({
-        "model_reward_accuracy_{}".format(frac): accuracy
-        for (frac, accuracy) in model_reward_accuracy
-    })
+    eval_metrics = {"mean_reward": mean_unclipped_reward}
     epoch_metrics.append(eval_metrics)
     log("Eval metrics: %s", str(eval_metrics))
     if report_fn:
@@ -668,83 +400,10 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   return epoch_metrics[-1]
 
 
-def extract_game_name(data_dir):
-  files = tf.gfile.ListDirectory(data_dir)
-  matches = [re.findall(r"on_(.*)_deterministic", f) for f in files]
-  non_empty_matches = [m for m in matches if m]
-  return non_empty_matches[0][0]
-
-
-def compute_final_evaluation_on_real_environments(hparams, job_results_dir,
-                                                  eval_output_file=None):
-  """Runs evaluation of PPO policies on environment with real environments."""
-  if eval_output_file is None:
-    eval_output_file = os.path.join(
-        FLAGS.eval_results_dir,
-        "result_{}.txt".format(
-            os.path.basename(os.path.normpath(job_results_dir))))
-  directories = tf.gfile.ListDirectory(job_results_dir)
-  results = {}
-  tmp_dir = os.path.join(FLAGS.eval_results_dir, "eval_tmp")
-  if tf.gfile.Exists(tmp_dir):
-    tf.gfile.DeleteRecursively(tmp_dir)
-  for directory in directories:
-    ppo_model_dir = os.path.join(job_results_dir, directory, "ppo")
-    data_dir = os.path.join(job_results_dir, directory, "data/initial")
-    hparams.game = extract_game_name(data_dir)
-    problem_name, _, _, _ = setup_problems(hparams)
-
-    tf.logging.info("Evaluating in real environment game %s." % hparams.game)
-    try:
-      mean_reward = int(generate_real_env_data(
-          problem_name, ppo_model_dir, hparams,
-          os.path.join(tmp_dir, directory),
-          "/tmp", autoencoder_path=None,
-          eval_phase=True, real_reward=True))
-      tf.logging.info(
-          "Mean eval reward on {}: {}".format(hparams.game, mean_reward))
-    except AttributeError:
-      tf.logging.info("No PPO model for: {}".format(ppo_model_dir))
-      mean_reward = None
-    game_results = results.get(hparams.game, [])
-    game_results.append(mean_reward)
-    results[hparams.game] = game_results
-
-  with open(eval_output_file, "w") as f:
-    for game in sorted(six.iterkeys(results)):
-      print("{}:".format(game), file=f, end="")
-      for z in reversed(sorted(results[game])):
-        print(" {}".format(z), file=f, end="")
-      print("", file=f)
-
-
-def combine_training_data(problem, final_data_dir, old_data_dirs,
-                          copy_last_eval_set=True):
-  """Add training data from old_data_dirs into final_data_dir."""
-  for i, data_dir in enumerate(old_data_dirs):
-    suffix = os.path.basename(data_dir)
-    # Glob train files in old data_dir
-    old_train_files = tf.gfile.Glob(
-        problem.filepattern(data_dir, tf.estimator.ModeKeys.TRAIN))
-    if (i + 1) == len(old_data_dirs) and copy_last_eval_set:
-      old_train_files += tf.gfile.Glob(
-          problem.filepattern(data_dir, tf.estimator.ModeKeys.EVAL))
-    for fname in old_train_files:
-      # Move them to the new data_dir with a suffix
-      # Since the data is read based on a prefix filepattern, adding the suffix
-      # should be fine.
-      new_fname = os.path.join(final_data_dir,
-                               os.path.basename(fname) + "." + suffix)
-      if not tf.gfile.Exists(new_fname):
-        tf.gfile.Copy(fname, new_fname)
-
-
 def main(_):
   hp = trainer_model_based_params.create_loop_hparams()
-  if FLAGS.job_dir_to_evaluate:
-    compute_final_evaluation_on_real_environments(hp, FLAGS.job_dir_to_evaluate)
-  else:
-    training_loop(hp, FLAGS.output_dir)
+  assert not FLAGS.job_dir_to_evaluate
+  training_loop(hp, FLAGS.output_dir)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/rl/trainer_model_based_ae_test.py b/tensor2tensor/rl/trainer_model_based_ae_test.py
deleted file mode 100644
index ba978eab0..000000000
--- a/tensor2tensor/rl/trainer_model_based_ae_test.py
+++ /dev/null
@@ -1,38 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Tiny run of trainer_model_based. Smoke test."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.rl import trainer_model_based
-
-import tensorflow as tf
-
-FLAGS = tf.flags.FLAGS
-
-
-class ModelRLExperimentTestAe(tf.test.TestCase):
-
-  def test_ae(self):
-    FLAGS.output_dir = tf.test.get_temp_dir()
-    FLAGS.loop_hparams_set = "rlmb_ae_tiny"
-    FLAGS.schedule = "train"  # skip evaluation for world model training
-    trainer_model_based.main(None)
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/rl/trainer_model_based_new.py b/tensor2tensor/rl/trainer_model_based_new.py
deleted file mode 100644
index a7db9b03d..000000000
--- a/tensor2tensor/rl/trainer_model_based_new.py
+++ /dev/null
@@ -1,418 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-r"""Training of model-based RL agents.
-
-Example invocation:
-
-python -m tensor2tensor.rl.trainer_model_based_new \
-    --output_dir=$HOME/t2t/rl_v1 \
-    --loop_hparams_set=rlmb_base \
-    --loop_hparams='num_real_env_frames=10000,epochs=3'
-"""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import contextlib
-import copy
-import datetime
-import math
-import os
-import time
-
-import gym
-import numpy as np
-
-from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
-from tensor2tensor.data_generators.gym_env import T2TGymEnv
-from tensor2tensor.models.research import rl
-from tensor2tensor.rl import rl_trainer_lib
-from tensor2tensor.rl import trainer_model_based_params
-from tensor2tensor.rl.envs.utils import InitialFrameChooser
-from tensor2tensor.utils import trainer_lib
-
-import tensorflow as tf
-
-
-flags = tf.flags
-FLAGS = flags.FLAGS
-
-
-@contextlib.contextmanager
-def temporary_flags(flag_settings):
-  old_values = {}
-  for flag_name, flag_value in flag_settings.items():
-    old_values[flag_name] = getattr(FLAGS, flag_name)
-    setattr(FLAGS, flag_name, flag_value)
-  yield
-  for flag_name, flag_value in old_values.items():
-    setattr(FLAGS, flag_name, flag_value)
-
-
-def _ppo_training_epochs(hparams, epoch, is_final_epoch, real_env_training):
-  """Helper for PPO restarts."""
-  if hparams.gather_ppo_real_env_data:
-    assert hparams.real_ppo_epochs_num is 0, (
-        "Should be put to 0 to enforce better readability")
-    real_training_ppo_epochs_num = int(math.ceil(
-        hparams.num_real_env_frames /
-        (hparams.epochs*hparams.real_ppo_epoch_length)))
-  else:
-    real_training_ppo_epochs_num = hparams.real_ppo_epochs_num
-
-  simulated_training_ppo_epochs_num = hparams.ppo_epochs_num
-
-  if epoch == -1:
-    assert real_env_training, (
-        "Epoch -1 should only be used for PPO collection in real environment.")
-    return real_training_ppo_epochs_num
-  ppo_training_epochs = (epoch + 1) * (simulated_training_ppo_epochs_num
-                                       + real_training_ppo_epochs_num)
-  if is_final_epoch:  # Length of training in the final epoch is doubled.
-    ppo_training_epochs += simulated_training_ppo_epochs_num
-  if real_env_training:
-    ppo_training_epochs += real_training_ppo_epochs_num
-  return ppo_training_epochs
-
-
-def setup_directories(base_dir, subdirs):
-  base_dir = os.path.expanduser(base_dir)
-  tf.gfile.MakeDirs(base_dir)
-
-  all_dirs = {}
-  for subdir in subdirs:
-    dir_name = os.path.join(base_dir, subdir)
-    tf.gfile.MakeDirs(dir_name)
-    all_dirs[subdir] = dir_name
-  return all_dirs
-
-
-def make_relative_timing_fn():
-  """Make a function that logs the duration since it was made."""
-  start_time = time.time()
-
-  def format_relative_time():
-    time_delta = time.time() - start_time
-    return str(datetime.timedelta(seconds=time_delta))
-
-  def log_relative_time():
-    tf.logging.info("Timing: %s", format_relative_time())
-
-  return log_relative_time
-
-
-def make_log_fn(epoch, log_relative_time_fn):
-
-  def log(msg, *args):
-    msg %= args
-    tf.logging.info("%s Epoch %d: %s", ">>>>>>>", epoch, msg)
-    log_relative_time_fn()
-
-  return log
-
-
-def train_supervised(problem, model_name, hparams, data_dir, output_dir,
-                     train_steps, eval_steps, local_eval_frequency=None,
-                     schedule="continuous_train_and_eval"):
-  """Train supervised."""
-  if local_eval_frequency is None:
-    local_eval_frequency = getattr(FLAGS, "local_eval_frequency")
-
-  exp_fn = trainer_lib.create_experiment_fn(
-      model_name, problem, data_dir, train_steps, eval_steps,
-      min_eval_frequency=local_eval_frequency
-  )
-  run_config = trainer_lib.create_run_config(model_name, model_dir=output_dir)
-  exp = exp_fn(run_config, hparams)
-  getattr(exp, schedule)()
-
-
-def train_agent(environment_spec, agent_model_dir,
-                event_dir, world_model_dir, data_dir, hparams, epoch=0,
-                is_final_epoch=False):
-  """Train the PPO agent in the simulated environment."""
-  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
-  ppo_params_names = ["epochs_num", "epoch_length",
-                      "learning_rate", "num_agents",
-                      "optimization_epochs", "eval_every_epochs"]
-
-  for param_name in ppo_params_names:
-    ppo_param_name = "ppo_" + param_name
-    if ppo_param_name in hparams:
-      ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
-
-  ppo_hparams.epochs_num = _ppo_training_epochs(hparams, epoch,
-                                                is_final_epoch, False)
-  ppo_hparams.save_models_every_epochs = 10
-  ppo_hparams.world_model_dir = world_model_dir
-  ppo_hparams.add_hparam("force_beginning_resets", True)
-
-  # Adding model hparams for model specific adjustments
-  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
-  ppo_hparams.add_hparam("model_hparams", model_hparams)
-
-  environment_spec = copy.copy(environment_spec)
-  environment_spec_param_names = [
-      "simulation_random_starts", "simulation_flip_first_random_for_beginning",
-      "intrinsic_reward_scale"
-  ]
-  for param_name in environment_spec_param_names:
-    environment_spec.set_hparam(param_name, hparams.get(param_name))
-  ppo_hparams.add_hparam("environment_spec", environment_spec)
-
-  ppo_hparams.add_hparam("initial_frame_chooser", InitialFrameChooser(
-      environment_spec, mode=tf.estimator.ModeKeys.EVAL
-  ))
-
-  # TODO(koz4k): Pass by arguments.
-  with temporary_flags({
-      "problem": environment_spec.initial_frames_problem,
-      "model": hparams.generative_model,
-      "hparams_set": hparams.generative_model_params,
-      "output_dir": world_model_dir,
-      "data_dir": data_dir,
-  }):
-    rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
-                         name_scope="ppo_sim%d" % (epoch + 1))
-
-
-def train_agent_real_env(
-    env, agent_model_dir, event_dir, data_dir,
-    hparams, epoch=0, is_final_epoch=False):
-  """Train the PPO agent in the real environment."""
-  del data_dir
-  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
-  ppo_params_names = ["epochs_num", "epoch_length",
-                      "learning_rate", "num_agents", "eval_every_epochs",
-                      "optimization_epochs", "effective_num_agents"]
-
-  # This should be overridden.
-  ppo_hparams.add_hparam("effective_num_agents", None)
-  for param_name in ppo_params_names:
-    ppo_param_name = "real_ppo_"+ param_name
-    if ppo_param_name in hparams:
-      ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
-
-  ppo_hparams.epochs_num = _ppo_training_epochs(hparams, epoch,
-                                                is_final_epoch, True)
-  # We do not save model, as that resets frames that we need at restarts.
-  # But we need to save at the last step, so we set it very high.
-  ppo_hparams.save_models_every_epochs = 1000000
-
-  environment_spec = rl.standard_atari_env_spec(
-      batch_env=env, include_clipping=False
-  )
-
-  ppo_hparams.add_hparam("environment_spec", environment_spec)
-
-  rl_trainer_lib.train(ppo_hparams, event_dir + "real", agent_model_dir,
-                       name_scope="ppo_real%d" % (epoch + 1))
-
-  # Save unfinished rollouts to history.
-  env.reset()
-
-
-def train_world_model(env, data_dir, output_dir, hparams, epoch):
-  """Train the world model on problem_name."""
-  train_steps = hparams.model_train_steps * (
-      epoch + hparams.inital_epoch_train_steps_multiplier)
-  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
-
-  # Hardcoded for now. TODO(koz4k): Make it a hparam.
-  model_hparams.video_num_input_frames = 4
-  model_hparams.video_num_target_frames = 1
-
-  model_hparams.learning_rate = model_hparams.learning_rate_constant
-  if epoch > 0:
-    model_hparams.learning_rate *= hparams.learning_rate_bump
-
-  train_supervised(
-      problem=env,
-      model_name=hparams.generative_model,
-      hparams=model_hparams,
-      data_dir=data_dir,
-      output_dir=output_dir,
-      train_steps=train_steps,
-      eval_steps=100,
-      local_eval_frequency=2000
-  )
-
-
-def make_gym_env(hparams):
-  """Make env."""
-  # TODO(kc): set reward clipping, when this will be possible
-  assert hparams.game == "pong", "Currently only games with [-1, 1] rewards."
-  game_mode = "Deterministic-v4"
-  camel_game_name = "".join(
-      [w[0].upper() + w[1:] for w in hparams.game.split("_")])
-  camel_game_name += game_mode
-  env_name = camel_game_name
-  env = gym.make(env_name)
-  if hparams.env_timesteps_limit != -1:
-    # Replace TimeLimit Wrapper with one of proper time step limit.
-    if isinstance(env, gym.wrappers.TimeLimit):
-      env = env.env
-    env = gym.wrappers.TimeLimit(env,
-                                 max_episode_steps=hparams.env_timesteps_limit)
-  return env
-
-
-def setup_env(hparams, data_dir):
-  """Setup."""
-  env = T2TGymEnv([make_gym_env(hparams)
-                   for _ in range(hparams.real_ppo_num_agents)],
-                  data_dir,
-                  grayscale=hparams.grayscale,
-                  resize_width_factor=hparams.resize_width_factor,
-                  resize_height_factor=hparams.resize_height_factor)
-  return env
-
-
-def eval_reward(env, clipped):
-  """Calculates mean rewards from given epoch."""
-  reward_name = "reward" if clipped else "unclipped_reward"
-  rewards = []
-  for rollout in env.current_epoch_rollouts():
-    if rollout[-1].done:
-      rollout_reward = sum(getattr(frame, reward_name) for frame in rollout)
-      rewards.append(rollout_reward)
-  if rewards:
-    mean_rewards = np.mean(rewards)
-  else:
-    mean_rewards = 0
-  return mean_rewards
-
-
-def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
-  """Run the main training loop."""
-  if report_fn:
-    assert report_metric is not None
-
-  # Directories
-  subdirectories = ["data", "tmp", "world_model", "ppo"]
-  directories = setup_directories(output_dir, subdirectories)
-
-  epoch = -1
-  data_dir = directories["data"]
-  env = setup_env(hparams, data_dir)
-  env.start_new_epoch(epoch)
-
-  # Timing log function
-  log_relative_time = make_relative_timing_fn()
-
-  # Per-epoch state
-  epoch_metrics = []
-
-  # Collect data from the real environment with PPO or random policy.
-  # TODO(lukaszkaiser): do we need option not to gather_ppo_real_env_data?
-  # We could set learning_rate=0 if this flag == False.
-  assert hparams.gather_ppo_real_env_data
-  ppo_model_dir = directories["ppo"]
-  tf.logging.info("Initial training of PPO in real environment.")
-  ppo_event_dir = os.path.join(directories["world_model"],
-                               "ppo_summaries/initial")
-  train_agent_real_env(
-      env, ppo_model_dir,
-      ppo_event_dir, data_dir,
-      hparams, epoch=epoch, is_final_epoch=False)
-  mean_unclipped_reward = eval_reward(env, clipped=False)
-  tf.logging.info("Mean reward (initial): {}".format(mean_unclipped_reward))
-
-  eval_metrics_event_dir = os.path.join(directories["world_model"],
-                                        "eval_metrics_event_dir")
-  eval_metrics_writer = tf.summary.FileWriter(eval_metrics_event_dir)
-
-  mean_unclipped_reward_summary = tf.Summary()
-  mean_unclipped_reward_summary.value.add(tag="mean_unclipped_reward",
-                                          simple_value=None)
-  mean_clipped_reward_summary = tf.Summary()
-  mean_clipped_reward_summary.value.add(tag="mean_clipped_reward",
-                                        simple_value=None)
-
-  sim_env_spec = rl.standard_atari_env_simulated_spec(
-      env,
-      # Hardcoded for now. TODO(koz4k): Make it a hparam.
-      video_num_input_frames=4, video_num_target_frames=1
-  )
-
-  for epoch in range(hparams.epochs):
-    env.generate_data()
-
-    is_final_epoch = (epoch + 1) == hparams.epochs
-    log = make_log_fn(epoch, log_relative_time)
-
-    # Train world model
-    log("Training world model")
-    train_world_model(env, data_dir,
-                      directories["world_model"], hparams, epoch)
-
-    # Train PPO
-    log("Training PPO in simulated environment.")
-    ppo_event_dir = os.path.join(directories["world_model"],
-                                 "ppo_summaries", str(epoch))
-    ppo_model_dir = directories["ppo"]
-    if not hparams.ppo_continue_training:
-      ppo_model_dir = ppo_event_dir
-
-    train_agent(sim_env_spec, ppo_model_dir,
-                ppo_event_dir, directories["world_model"], data_dir,
-                hparams, epoch=epoch, is_final_epoch=is_final_epoch)
-
-    env.start_new_epoch(epoch)
-
-    # Train PPO on real env (short)
-    log("Training PPO in real environment.")
-    train_agent_real_env(
-        env, ppo_model_dir,
-        ppo_event_dir, data_dir,
-        hparams, epoch=epoch, is_final_epoch=is_final_epoch)
-
-    if hparams.stop_loop_early:
-      return 0.0
-    mean_clipped_reward = eval_reward(env, clipped=True)
-    log("Mean clipped reward during generation: {}".format(
-        mean_clipped_reward))  # this was output of generate_real_env_data(...)
-
-    mean_unclipped_reward = eval_reward(env, clipped=False)
-    log("Mean eval reward (unclipped): {}".format(mean_unclipped_reward))
-
-    # Summarize metrics.
-    mean_unclipped_reward_summary.value[0].simple_value = mean_unclipped_reward
-    eval_metrics_writer.add_summary(mean_unclipped_reward_summary, epoch)
-    mean_clipped_reward_summary.value[0].simple_value = int(mean_clipped_reward)
-    eval_metrics_writer.add_summary(mean_clipped_reward_summary, epoch)
-    eval_metrics_writer.flush()
-
-    # Report metrics
-    eval_metrics = {"mean_unclipped_reward": mean_unclipped_reward}
-    epoch_metrics.append(eval_metrics)
-    log("Eval metrics: %s", str(eval_metrics))
-    if report_fn:
-      report_fn(eval_metrics[report_metric], epoch)
-
-  # Return the evaluation metrics from the final epoch
-  return epoch_metrics[-1]
-
-
-def main(_):
-  hp = trainer_model_based_params.create_loop_hparams()
-  assert not FLAGS.job_dir_to_evaluate
-  training_loop(hp, FLAGS.output_dir)
-
-
-if __name__ == "__main__":
-  tf.logging.set_verbosity(tf.logging.INFO)
-  tf.app.run()
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index a7b34d7b0..25c67da8c 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -63,7 +63,7 @@ def rlmb_base():
       simulated_env_generator_num_steps=2000,
       simulation_random_starts=True,  # Use random starts in PPO.
       # Flip the first random frame in PPO batch for the true beginning.
-      simulation_flip_first_random_for_beginning=True,
+      simulation_flip_first_random_for_beginning=False,
       intrinsic_reward_scale=0.,
       ppo_epochs_num=1000,  # This should be enough to see something
       # Our simulated envs do not know how to reset.

From 6f6268023849dc7ade7257385e7a26e432c5285b Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 17 Oct 2018 09:44:39 -0700
Subject: [PATCH 1027/2720] Add resnet-style architecture to model P(z^l_t |
 z^l_{t-1}, z^{-l}{t}).

PiperOrigin-RevId: 217534781
---
 tensor2tensor/models/research/glow_ops.py      | 16 ++++++++++++++--
 tensor2tensor/models/research/glow_ops_test.py |  2 +-
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index b23923971..c0578bdae 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -473,7 +473,7 @@ def squeeze(name, x, factor=2, reverse=True):
 
 @add_arg_scope
 def tensor_to_dist(name, x, output_channels=None, architecture="single_conv",
-                   depth=1, pre_output_channels=512):
+                   depth=1, pre_output_channels=512, width=512):
   """Map x to the mean and log-scale of a Gaussian.
 
   Args:
@@ -484,6 +484,7 @@ def tensor_to_dist(name, x, output_channels=None, architecture="single_conv",
     architecture: "single_conv" or "glow_nn"
     depth: depth of architecture mapping to the mean and std.
     pre_output_channels: output channels before the final (mean, std) mapping.
+    width: Resnet width.
   Returns:
     dist: instance of tf.distributions.Normal
   Raises:
@@ -506,6 +507,16 @@ def tensor_to_dist(name, x, output_channels=None, architecture="single_conv",
                               filter_size=[3, 3], stride=[1, 1],
                               output_channels=2*output_channels,
                               apply_actnorm=False, conv_init="zeros")
+    elif architecture == "glow_resnet":
+      h = x
+      for layer in range(depth):
+        h2 = conv_block("glow_res_%d" % layer, h, mid_channels=width)
+        h3 = conv2d("glow_res_zeros_%d" % layer, h2, conv_init="zeros",
+                    output_channels=x_shape[-1], apply_actnorm=False)
+        h += h3
+      mean_log_scale = conv2d("glow_res_final", h, conv_init="zeros",
+                              output_channels=2*output_channels,
+                              apply_actnorm=False)
     else:
       raise ValueError("expected architecture to be single_conv or glow_nn "
                        "got %s" % architecture)
@@ -582,7 +593,8 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
         "latent_stack", latent_stack, output_channels=output_channels,
         architecture=hparams.latent_architecture,
         depth=hparams.latent_encoder_depth,
-        pre_output_channels=hparams.latent_pre_output_channels)
+        pre_output_channels=hparams.latent_pre_output_channels,
+        width=hparams.latent_encoder_width)
     if latent_skip:
       cond_dist = tf.distributions.Normal(
           cond_dist.loc + latent[-1], cond_dist.scale)
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index a04400ede..423d82a1c 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -140,7 +140,7 @@ def check_tensor_to_dist(self, architecture):
         self.assertTrue(np.allclose(scale, 1.0))
 
   def test_tensor_to_dist(self):
-    for architecture in ["single_conv", "glow_nn"]:
+    for architecture in ["single_conv", "glow_nn", "glow_resnet"]:
       self.check_tensor_to_dist(architecture)
 
   def test_split(self):

From 210772f207eb8336dc1fb0316fda3354d3b2e57b Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Wed, 17 Oct 2018 23:45:57 +0200
Subject: [PATCH 1028/2720] Model-Based RL: Enable registration of T2TGymEnv
 subclasses. (#1149)

* T2TGymEnv: move data_dir from __init__ to generate_data(), start_new_epoch(), splits_and_paths()

* T2TGymEnv: Enable registration of subclasses - add env name to problem name; initialize with env name instead of env list, pass args and kwargs
---
 tensor2tensor/data_generators/gym_env.py      | 70 +++++++++++--------
 tensor2tensor/data_generators/gym_env_test.py | 64 +++++++++--------
 tensor2tensor/rl/trainer_model_based.py       | 32 +++------
 3 files changed, 86 insertions(+), 80 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index c43fb90a5..8501bdd80 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -22,6 +22,7 @@
 import collections
 import itertools
 
+import gym
 from gym.spaces import Box
 import numpy as np
 
@@ -48,6 +49,17 @@ def __deepcopy__(self, memo):
     return self
 
 
+def make_gym_env(name, timesteps_limit=-1):
+  env = gym.make(name)
+  if timesteps_limit != -1:
+    # Replace TimeLimit Wrapper with one of proper time step limit.
+    if isinstance(env, gym.wrappers.TimeLimit):
+      env = env.env
+    env = gym.wrappers.TimeLimit(env,
+                                 max_episode_steps=timesteps_limit)
+  return env
+
+
 class T2TEnv(video_utils.VideoProblem):
   """Abstract class representing a batch of environments.
 
@@ -71,12 +83,11 @@ class T2TEnv(video_utils.VideoProblem):
   reward_range = (-1, 1)
   name = None
 
-  def __init__(self, batch_size, data_dir):
-    super(T2TEnv, self).__init__()
+  def __init__(self, batch_size, *args, **kwargs):
+    super(T2TEnv, self).__init__(*args, **kwargs)
 
     self.clear_history()
     self.batch_size = batch_size
-    self.data_dir = data_dir
     self._current_batch_frames = [None for _ in range(batch_size)]
     self._current_batch_rollouts = [[] for _ in range(batch_size)]
     self._current_epoch_rollouts = []
@@ -100,7 +111,7 @@ def clear_history(self):
     """Clears the rollout history."""
     self._rollouts_by_epoch_and_split = collections.OrderedDict()
 
-  def start_new_epoch(self, epoch, load_data=True):
+  def start_new_epoch(self, epoch, load_data_dir):
     if not isinstance(epoch, int):
       raise ValueError("Epoch should be integer, got {}".format(epoch))
     if epoch in self._rollouts_by_epoch_and_split:
@@ -108,8 +119,7 @@ def start_new_epoch(self, epoch, load_data=True):
     self.current_epoch = epoch
     self._current_epoch_rollouts = []
     self._rollouts_by_epoch_and_split[epoch] = collections.defaultdict(list)
-    if load_data:
-      self._load_epoch_data()
+    self._load_epoch_data(load_data_dir)
 
   def current_epoch_rollouts(self, split=None):
     rollouts_by_split = self._rollouts_by_epoch_and_split[self.current_epoch]
@@ -229,12 +239,9 @@ def reset(self, indices=None):
       Batch of initial observations of reset environments.
     """
     if self.current_epoch is None:
-      # It's here so that the old pipeline works.
-      self.start_new_epoch(0, load_data=False)
-      # TODO(koz4k): Replace with:
-      # raise ValueError(
-      #     "No current epoch. start_new_epoch() should first be called."
-      # )
+      raise ValueError(
+          "No current epoch. start_new_epoch() should first be called."
+      )
 
     if indices is None:
       indices = np.arange(self.batch_size)
@@ -386,8 +393,7 @@ def split_size(split_index):
     self._rollouts_by_epoch_and_split[self.current_epoch] = rollouts_by_split
     self._current_epoch_rollouts = []
 
-  @property
-  def splits_and_paths(self):
+  def splits_and_paths(self, data_dir):
     """List of pairs (split, paths) for the current epoch."""
     filepath_fns = {
         problem.DatasetSplit.TRAIN: self.training_filepaths,
@@ -404,7 +410,7 @@ def append_epoch(paths):
     # We set shuffled=True as we don't want to shuffle on disk later.
     return [
         (split["split"], append_epoch(filepath_fns[split["split"]](
-            self.data_dir, split["shards"], shuffled=True
+            data_dir, split["shards"], shuffled=True
         )))
         for split in self.dataset_splits
     ]
@@ -417,14 +423,14 @@ def filepattern(self, data_dir, mode, shard=None, only_last=False):
       filepattern += ".{}".format(self.current_epoch)
     return filepattern
 
-  def generate_data(self, data_dir=None, tmp_dir=None, task_id=-1):
+  def generate_data(self, data_dir, tmp_dir=None, task_id=-1):
     """Saves the current epoch rollouts to disk, split into train/dev sets."""
     if not self._rollouts_by_epoch_and_split[self.current_epoch]:
       # Data not loaded from disk.
       self._split_current_epoch()
 
     rollouts_by_split = self._rollouts_by_epoch_and_split[self.current_epoch]
-    splits_and_paths = self.splits_and_paths
+    splits_and_paths = self.splits_and_paths(data_dir)
 
     for (split, paths) in splits_and_paths:
       rollouts = rollouts_by_split[split]
@@ -442,12 +448,12 @@ def generate_data(self, data_dir=None, tmp_dir=None, task_id=-1):
             cycle_every_n=float("inf")
         )
 
-  def _load_epoch_data(self):
+  def _load_epoch_data(self, data_dir):
     any_files_found = False
     all_files_found = True
     any_shard_empty = False
 
-    for split, paths in self.splits_and_paths:
+    for split, paths in self.splits_and_paths(data_dir):
       try:
         any_shard_empty |= self._load_epoch_split(split, paths)
         any_files_found = True
@@ -506,28 +512,34 @@ def get_feature_value(key, list_name):
 
 
 class T2TGymEnv(T2TEnv):
-  """Class representing a batch of Gym environments."""
+  """Class representing a batch of Gym environments.
 
-  name = "t2t_gym_env"
+  Do not register it, instead create subclass with hardcoded __init__
+  arguments and register this subclass.
+  """
 
-  def __init__(self, envs, data_dir, grayscale=False,
-               resize_height_factor=1, resize_width_factor=1):
-    super(T2TGymEnv, self).__init__(len(envs), data_dir)
+  def __init__(self, base_env_name, batch_size, grayscale=False,
+               resize_height_factor=1, resize_width_factor=1,
+               base_env_timesteps_limit=-1, *args, **kwargs):
+    super(T2TGymEnv, self).__init__(batch_size, *args, **kwargs)
     self.grayscale = grayscale
     self.resize_height_factor = resize_height_factor
     self.resize_width_factor = resize_width_factor
-    if not envs:
-      raise ValueError("Must have at least one environment.")
-    self._envs = envs
+    if not self.name:
+      # Set problem name if not registered.
+      self.name = "t2t_gym_env_{}".format(base_env_name)
+
+    self._envs = [make_gym_env(base_env_name, base_env_timesteps_limit)
+                  for _ in range(self.batch_size)]
 
-    orig_observ_space = envs[0].observation_space
+    orig_observ_space = self._envs[0].observation_space
     if not all(env.observation_space == orig_observ_space
                for env in self._envs):
       raise ValueError("All environments must use the same observation space.")
 
     self.observation_space = self._derive_observation_space(orig_observ_space)
 
-    self.action_space = envs[0].action_space
+    self.action_space = self._envs[0].action_space
     if not all(env.action_space == self.action_space for env in self._envs):
       raise ValueError("All environments must use the same action space.")
 
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 61b3bd24d..98913d46d 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -32,6 +32,8 @@
 
 import tensorflow as tf
 
+from tensor2tensor.data_generators.gym_env import make_gym_env
+
 
 class TestEnv(gym.Env):
   """Test environment.
@@ -61,6 +63,10 @@ def step(self, action):
   def reset(self):
     return self._generate_ob()
 
+TEST_ENV_NAME = "T2TTestEnv-v1"
+
+gym.envs.register(id=TEST_ENV_NAME, entry_point=TestEnv)
+
 
 class GymEnvTest(tf.test.TestCase):
 
@@ -75,19 +81,18 @@ def setUp(self):
     shutil.rmtree(self.out_dir)
     os.mkdir(self.out_dir)
 
-  def init_batch_and_play(self, env_lambda, steps_per_epoch=1,
+  def init_batch_and_play(self, env_name, steps_per_epoch=1,
                           epochs=(0,), generate_data=False, **kwargs):
-    raw_envs = [env_lambda(), env_lambda()]
-    env = gym_env.T2TGymEnv(raw_envs, self.out_dir, **kwargs)
+    env = gym_env.T2TGymEnv(env_name, batch_size=2, **kwargs)
     obs = list()
     rewards = list()
     num_dones = 0
     for epoch in epochs:
-      env.start_new_epoch(epoch)
+      env.start_new_epoch(epoch, self.out_dir)
       _, epoch_obs, epoch_rewards, epoch_num_dones = \
           self.play(env, steps_per_epoch)
       if generate_data:
-        env.generate_data()
+        env.generate_data(self.out_dir)
       obs.extend(epoch_obs)
       rewards.extend(epoch_rewards)
       num_dones += epoch_num_dones
@@ -110,7 +115,7 @@ def play(self, env, n_steps):
 
   def test_splits_dataset(self):
     env, _, _, _ = self.init_batch_and_play(
-        TestEnv, steps_per_epoch=20, generate_data=True
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
     )
 
     for split in self.splits:
@@ -118,7 +123,7 @@ def test_splits_dataset(self):
 
   def test_split_preserves_number_of_rollouts(self):
     env, _, _, num_dones = self.init_batch_and_play(
-        TestEnv, steps_per_epoch=20, generate_data=True
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
     )
 
     num_rollouts_after_split = sum(
@@ -131,7 +136,7 @@ def test_split_preserves_number_of_rollouts(self):
 
   def test_split_preserves_number_of_frames(self):
     env, _, _, num_dones = self.init_batch_and_play(
-        TestEnv, steps_per_epoch=20, generate_data=True
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
     )
 
     num_frames = sum(
@@ -146,7 +151,7 @@ def test_split_preserves_number_of_frames(self):
   def test_generates_data(self):
     # This test needs base env which outputs done after two steps.
     self.init_batch_and_play(
-        TestEnv, steps_per_epoch=20, generate_data=True
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
     )
 
     filenames = os.listdir(self.out_dir)
@@ -162,18 +167,18 @@ def num_ending_with(filenames, suffix):
           1 for filename in filenames if filename.endswith(suffix)
       )
 
-    env = gym_env.T2TGymEnv([TestEnv() for _ in range(2)], self.out_dir)
-    env.start_new_epoch(0)
+    env = gym_env.T2TGymEnv(TEST_ENV_NAME, batch_size=2)
+    env.start_new_epoch(0, self.out_dir)
     self.play(env, n_steps=20)
-    env.generate_data()
+    env.generate_data(self.out_dir)
 
     filenames = os.listdir(self.out_dir)
     num_shards_per_epoch = len(filenames)
     self.assertEqual(num_ending_with(filenames, ".0"), num_shards_per_epoch)
 
-    env.start_new_epoch(1)
+    env.start_new_epoch(1, self.out_dir)
     self.play(env, n_steps=20)
-    env.generate_data()
+    env.generate_data(self.out_dir)
 
     filenames = os.listdir(self.out_dir)
     self.assertEqual(len(filenames), 2 * num_shards_per_epoch)
@@ -182,14 +187,14 @@ def num_ending_with(filenames, suffix):
 
   def test_frame_numbers_are_continuous(self):
     env, _, _, _ = self.init_batch_and_play(
-        TestEnv, steps_per_epoch=20, generate_data=True
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
     )
 
     frame_numbers = [
         tf.train.Example.FromString(
             record
         ).features.feature["frame_number"].int64_list.value[0]
-        for (_, paths) in env.splits_and_paths
+        for (_, paths) in env.splits_and_paths(self.out_dir)
         for path in paths
         for record in tf.python_io.tf_record_iterator(path)
     ]
@@ -203,19 +208,19 @@ def test_frame_numbers_are_continuous(self):
       last_frame_number = frame_number
 
   def test_clipping(self):
-    env_lambda = TestEnv
-    _, _, rewards, _ = self.init_batch_and_play(env_lambda, steps_per_epoch=2)
+    _, _, rewards, _ = self.init_batch_and_play(TEST_ENV_NAME,
+                                                steps_per_epoch=2)
     self.assertTrue(np.max(rewards) == 1)
     self.assertTrue(np.min(rewards) == -1)
 
   def test_resize(self):
-    env_lambda = TestEnv
-    orig_env = env_lambda()
+    env_name = TEST_ENV_NAME
+    orig_env = make_gym_env(env_name)
     resize_height_factor = 2
     resize_width_factor = 3
     orig_height, orig_width = orig_env.observation_space.shape[:2]
     env, obs, _, _ = self.init_batch_and_play(
-        env_lambda, steps_per_epoch=1,
+        env_name, steps_per_epoch=1,
         resize_height_factor=resize_height_factor,
         resize_width_factor=resize_width_factor)
     for obs_batch in obs:
@@ -233,21 +238,22 @@ def assert_channels(self, env, obs, n_channels):
       self.assertEqual(ob.shape[2], n_channels)
 
   def test_channels(self):
-    env_lambda = TestEnv
-    env, obs, _, _ = self.init_batch_and_play(env_lambda, grayscale=True)
+    env_name = TEST_ENV_NAME
+    env, obs, _, _ = self.init_batch_and_play(env_name, grayscale=True)
     self.assert_channels(env, obs, n_channels=1)
 
-    env, obs, _, _ = self.init_batch_and_play(env_lambda, grayscale=False)
+    env, obs, _, _ = self.init_batch_and_play(env_name, grayscale=False)
     self.assert_channels(env, obs, n_channels=3)
 
   def test_generating_and_loading_preserves_rollouts(self):
-    from_env = gym_env.T2TGymEnv([TestEnv()], self.out_dir)
-    from_env.start_new_epoch(0)
+    env_name = TEST_ENV_NAME
+    from_env = gym_env.T2TGymEnv(env_name, batch_size=1)
+    from_env.start_new_epoch(0, self.out_dir)
     self.play(from_env, n_steps=20)
-    from_env.generate_data()
+    from_env.generate_data(self.out_dir)
 
-    to_env = gym_env.T2TGymEnv([TestEnv()], self.out_dir)
-    to_env.start_new_epoch(0)
+    to_env = gym_env.T2TGymEnv(env_name, batch_size=1)
+    to_env.start_new_epoch(0, self.out_dir)
 
     self.assertEqual(
         from_env.current_epoch_rollouts(), to_env.current_epoch_rollouts()
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 90fb103b7..e04240f95 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -33,7 +33,6 @@
 import os
 import time
 
-import gym
 import numpy as np
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
@@ -246,31 +245,20 @@ def train_world_model(env, data_dir, output_dir, hparams, epoch):
   )
 
 
-def make_gym_env(hparams):
-  """Make env."""
+def setup_env(hparams):
+  """Setup."""
   game_mode = "Deterministic-v4"
   camel_game_name = "".join(
       [w[0].upper() + w[1:] for w in hparams.game.split("_")])
   camel_game_name += game_mode
   env_name = camel_game_name
-  env = gym.make(env_name)
-  if hparams.env_timesteps_limit != -1:
-    # Replace TimeLimit Wrapper with one of proper time step limit.
-    if isinstance(env, gym.wrappers.TimeLimit):
-      env = env.env
-    env = gym.wrappers.TimeLimit(env,
-                                 max_episode_steps=hparams.env_timesteps_limit)
-  return env
 
-
-def setup_env(hparams, data_dir):
-  """Setup."""
-  env = T2TGymEnv([make_gym_env(hparams)
-                   for _ in range(hparams.real_ppo_num_agents)],
-                  data_dir,
+  env = T2TGymEnv(base_env_name=env_name,
+                  batch_size=hparams.real_ppo_num_agents,
                   grayscale=hparams.grayscale,
                   resize_width_factor=hparams.resize_width_factor,
-                  resize_height_factor=hparams.resize_height_factor)
+                  resize_height_factor=hparams.resize_height_factor,
+                  base_env_timesteps_limit=hparams.env_timesteps_limit)
   return env
 
 
@@ -300,8 +288,8 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   epoch = -1
   data_dir = directories["data"]
-  env = setup_env(hparams, data_dir)
-  env.start_new_epoch(epoch)
+  env = setup_env(hparams)
+  env.start_new_epoch(epoch, data_dir)
 
   # Timing log function
   log_relative_time = make_relative_timing_fn()
@@ -342,7 +330,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   )
 
   for epoch in range(hparams.epochs):
-    env.generate_data()
+    env.generate_data(data_dir)
 
     is_final_epoch = (epoch + 1) == hparams.epochs
     log = make_log_fn(epoch, log_relative_time)
@@ -364,7 +352,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
                 ppo_event_dir, directories["world_model"], data_dir,
                 hparams, epoch=epoch, is_final_epoch=is_final_epoch)
 
-    env.start_new_epoch(epoch)
+    env.start_new_epoch(epoch, data_dir)
 
     # Train PPO on real env (short)
     log("Training PPO in real environment.")

From 5d0617c1d3199679fd4ab79d6fa2de0a40895fa9 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 17 Oct 2018 12:20:14 -0700
Subject: [PATCH 1029/2720] INTERNAL

PiperOrigin-RevId: 217565544
---
 tensor2tensor/data_generators/gym_env.py      | 70 ++++++++-----------
 tensor2tensor/data_generators/gym_env_test.py | 64 ++++++++---------
 tensor2tensor/rl/trainer_model_based.py       | 32 ++++++---
 3 files changed, 80 insertions(+), 86 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 8501bdd80..c43fb90a5 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -22,7 +22,6 @@
 import collections
 import itertools
 
-import gym
 from gym.spaces import Box
 import numpy as np
 
@@ -49,17 +48,6 @@ def __deepcopy__(self, memo):
     return self
 
 
-def make_gym_env(name, timesteps_limit=-1):
-  env = gym.make(name)
-  if timesteps_limit != -1:
-    # Replace TimeLimit Wrapper with one of proper time step limit.
-    if isinstance(env, gym.wrappers.TimeLimit):
-      env = env.env
-    env = gym.wrappers.TimeLimit(env,
-                                 max_episode_steps=timesteps_limit)
-  return env
-
-
 class T2TEnv(video_utils.VideoProblem):
   """Abstract class representing a batch of environments.
 
@@ -83,11 +71,12 @@ class T2TEnv(video_utils.VideoProblem):
   reward_range = (-1, 1)
   name = None
 
-  def __init__(self, batch_size, *args, **kwargs):
-    super(T2TEnv, self).__init__(*args, **kwargs)
+  def __init__(self, batch_size, data_dir):
+    super(T2TEnv, self).__init__()
 
     self.clear_history()
     self.batch_size = batch_size
+    self.data_dir = data_dir
     self._current_batch_frames = [None for _ in range(batch_size)]
     self._current_batch_rollouts = [[] for _ in range(batch_size)]
     self._current_epoch_rollouts = []
@@ -111,7 +100,7 @@ def clear_history(self):
     """Clears the rollout history."""
     self._rollouts_by_epoch_and_split = collections.OrderedDict()
 
-  def start_new_epoch(self, epoch, load_data_dir):
+  def start_new_epoch(self, epoch, load_data=True):
     if not isinstance(epoch, int):
       raise ValueError("Epoch should be integer, got {}".format(epoch))
     if epoch in self._rollouts_by_epoch_and_split:
@@ -119,7 +108,8 @@ def start_new_epoch(self, epoch, load_data_dir):
     self.current_epoch = epoch
     self._current_epoch_rollouts = []
     self._rollouts_by_epoch_and_split[epoch] = collections.defaultdict(list)
-    self._load_epoch_data(load_data_dir)
+    if load_data:
+      self._load_epoch_data()
 
   def current_epoch_rollouts(self, split=None):
     rollouts_by_split = self._rollouts_by_epoch_and_split[self.current_epoch]
@@ -239,9 +229,12 @@ def reset(self, indices=None):
       Batch of initial observations of reset environments.
     """
     if self.current_epoch is None:
-      raise ValueError(
-          "No current epoch. start_new_epoch() should first be called."
-      )
+      # It's here so that the old pipeline works.
+      self.start_new_epoch(0, load_data=False)
+      # TODO(koz4k): Replace with:
+      # raise ValueError(
+      #     "No current epoch. start_new_epoch() should first be called."
+      # )
 
     if indices is None:
       indices = np.arange(self.batch_size)
@@ -393,7 +386,8 @@ def split_size(split_index):
     self._rollouts_by_epoch_and_split[self.current_epoch] = rollouts_by_split
     self._current_epoch_rollouts = []
 
-  def splits_and_paths(self, data_dir):
+  @property
+  def splits_and_paths(self):
     """List of pairs (split, paths) for the current epoch."""
     filepath_fns = {
         problem.DatasetSplit.TRAIN: self.training_filepaths,
@@ -410,7 +404,7 @@ def append_epoch(paths):
     # We set shuffled=True as we don't want to shuffle on disk later.
     return [
         (split["split"], append_epoch(filepath_fns[split["split"]](
-            data_dir, split["shards"], shuffled=True
+            self.data_dir, split["shards"], shuffled=True
         )))
         for split in self.dataset_splits
     ]
@@ -423,14 +417,14 @@ def filepattern(self, data_dir, mode, shard=None, only_last=False):
       filepattern += ".{}".format(self.current_epoch)
     return filepattern
 
-  def generate_data(self, data_dir, tmp_dir=None, task_id=-1):
+  def generate_data(self, data_dir=None, tmp_dir=None, task_id=-1):
     """Saves the current epoch rollouts to disk, split into train/dev sets."""
     if not self._rollouts_by_epoch_and_split[self.current_epoch]:
       # Data not loaded from disk.
       self._split_current_epoch()
 
     rollouts_by_split = self._rollouts_by_epoch_and_split[self.current_epoch]
-    splits_and_paths = self.splits_and_paths(data_dir)
+    splits_and_paths = self.splits_and_paths
 
     for (split, paths) in splits_and_paths:
       rollouts = rollouts_by_split[split]
@@ -448,12 +442,12 @@ def generate_data(self, data_dir, tmp_dir=None, task_id=-1):
             cycle_every_n=float("inf")
         )
 
-  def _load_epoch_data(self, data_dir):
+  def _load_epoch_data(self):
     any_files_found = False
     all_files_found = True
     any_shard_empty = False
 
-    for split, paths in self.splits_and_paths(data_dir):
+    for split, paths in self.splits_and_paths:
       try:
         any_shard_empty |= self._load_epoch_split(split, paths)
         any_files_found = True
@@ -512,34 +506,28 @@ def get_feature_value(key, list_name):
 
 
 class T2TGymEnv(T2TEnv):
-  """Class representing a batch of Gym environments.
+  """Class representing a batch of Gym environments."""
 
-  Do not register it, instead create subclass with hardcoded __init__
-  arguments and register this subclass.
-  """
+  name = "t2t_gym_env"
 
-  def __init__(self, base_env_name, batch_size, grayscale=False,
-               resize_height_factor=1, resize_width_factor=1,
-               base_env_timesteps_limit=-1, *args, **kwargs):
-    super(T2TGymEnv, self).__init__(batch_size, *args, **kwargs)
+  def __init__(self, envs, data_dir, grayscale=False,
+               resize_height_factor=1, resize_width_factor=1):
+    super(T2TGymEnv, self).__init__(len(envs), data_dir)
     self.grayscale = grayscale
     self.resize_height_factor = resize_height_factor
     self.resize_width_factor = resize_width_factor
-    if not self.name:
-      # Set problem name if not registered.
-      self.name = "t2t_gym_env_{}".format(base_env_name)
-
-    self._envs = [make_gym_env(base_env_name, base_env_timesteps_limit)
-                  for _ in range(self.batch_size)]
+    if not envs:
+      raise ValueError("Must have at least one environment.")
+    self._envs = envs
 
-    orig_observ_space = self._envs[0].observation_space
+    orig_observ_space = envs[0].observation_space
     if not all(env.observation_space == orig_observ_space
                for env in self._envs):
       raise ValueError("All environments must use the same observation space.")
 
     self.observation_space = self._derive_observation_space(orig_observ_space)
 
-    self.action_space = self._envs[0].action_space
+    self.action_space = envs[0].action_space
     if not all(env.action_space == self.action_space for env in self._envs):
       raise ValueError("All environments must use the same action space.")
 
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 98913d46d..61b3bd24d 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -32,8 +32,6 @@
 
 import tensorflow as tf
 
-from tensor2tensor.data_generators.gym_env import make_gym_env
-
 
 class TestEnv(gym.Env):
   """Test environment.
@@ -63,10 +61,6 @@ def step(self, action):
   def reset(self):
     return self._generate_ob()
 
-TEST_ENV_NAME = "T2TTestEnv-v1"
-
-gym.envs.register(id=TEST_ENV_NAME, entry_point=TestEnv)
-
 
 class GymEnvTest(tf.test.TestCase):
 
@@ -81,18 +75,19 @@ def setUp(self):
     shutil.rmtree(self.out_dir)
     os.mkdir(self.out_dir)
 
-  def init_batch_and_play(self, env_name, steps_per_epoch=1,
+  def init_batch_and_play(self, env_lambda, steps_per_epoch=1,
                           epochs=(0,), generate_data=False, **kwargs):
-    env = gym_env.T2TGymEnv(env_name, batch_size=2, **kwargs)
+    raw_envs = [env_lambda(), env_lambda()]
+    env = gym_env.T2TGymEnv(raw_envs, self.out_dir, **kwargs)
     obs = list()
     rewards = list()
     num_dones = 0
     for epoch in epochs:
-      env.start_new_epoch(epoch, self.out_dir)
+      env.start_new_epoch(epoch)
       _, epoch_obs, epoch_rewards, epoch_num_dones = \
           self.play(env, steps_per_epoch)
       if generate_data:
-        env.generate_data(self.out_dir)
+        env.generate_data()
       obs.extend(epoch_obs)
       rewards.extend(epoch_rewards)
       num_dones += epoch_num_dones
@@ -115,7 +110,7 @@ def play(self, env, n_steps):
 
   def test_splits_dataset(self):
     env, _, _, _ = self.init_batch_and_play(
-        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
+        TestEnv, steps_per_epoch=20, generate_data=True
     )
 
     for split in self.splits:
@@ -123,7 +118,7 @@ def test_splits_dataset(self):
 
   def test_split_preserves_number_of_rollouts(self):
     env, _, _, num_dones = self.init_batch_and_play(
-        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
+        TestEnv, steps_per_epoch=20, generate_data=True
     )
 
     num_rollouts_after_split = sum(
@@ -136,7 +131,7 @@ def test_split_preserves_number_of_rollouts(self):
 
   def test_split_preserves_number_of_frames(self):
     env, _, _, num_dones = self.init_batch_and_play(
-        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
+        TestEnv, steps_per_epoch=20, generate_data=True
     )
 
     num_frames = sum(
@@ -151,7 +146,7 @@ def test_split_preserves_number_of_frames(self):
   def test_generates_data(self):
     # This test needs base env which outputs done after two steps.
     self.init_batch_and_play(
-        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
+        TestEnv, steps_per_epoch=20, generate_data=True
     )
 
     filenames = os.listdir(self.out_dir)
@@ -167,18 +162,18 @@ def num_ending_with(filenames, suffix):
           1 for filename in filenames if filename.endswith(suffix)
       )
 
-    env = gym_env.T2TGymEnv(TEST_ENV_NAME, batch_size=2)
-    env.start_new_epoch(0, self.out_dir)
+    env = gym_env.T2TGymEnv([TestEnv() for _ in range(2)], self.out_dir)
+    env.start_new_epoch(0)
     self.play(env, n_steps=20)
-    env.generate_data(self.out_dir)
+    env.generate_data()
 
     filenames = os.listdir(self.out_dir)
     num_shards_per_epoch = len(filenames)
     self.assertEqual(num_ending_with(filenames, ".0"), num_shards_per_epoch)
 
-    env.start_new_epoch(1, self.out_dir)
+    env.start_new_epoch(1)
     self.play(env, n_steps=20)
-    env.generate_data(self.out_dir)
+    env.generate_data()
 
     filenames = os.listdir(self.out_dir)
     self.assertEqual(len(filenames), 2 * num_shards_per_epoch)
@@ -187,14 +182,14 @@ def num_ending_with(filenames, suffix):
 
   def test_frame_numbers_are_continuous(self):
     env, _, _, _ = self.init_batch_and_play(
-        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
+        TestEnv, steps_per_epoch=20, generate_data=True
     )
 
     frame_numbers = [
         tf.train.Example.FromString(
             record
         ).features.feature["frame_number"].int64_list.value[0]
-        for (_, paths) in env.splits_and_paths(self.out_dir)
+        for (_, paths) in env.splits_and_paths
         for path in paths
         for record in tf.python_io.tf_record_iterator(path)
     ]
@@ -208,19 +203,19 @@ def test_frame_numbers_are_continuous(self):
       last_frame_number = frame_number
 
   def test_clipping(self):
-    _, _, rewards, _ = self.init_batch_and_play(TEST_ENV_NAME,
-                                                steps_per_epoch=2)
+    env_lambda = TestEnv
+    _, _, rewards, _ = self.init_batch_and_play(env_lambda, steps_per_epoch=2)
     self.assertTrue(np.max(rewards) == 1)
     self.assertTrue(np.min(rewards) == -1)
 
   def test_resize(self):
-    env_name = TEST_ENV_NAME
-    orig_env = make_gym_env(env_name)
+    env_lambda = TestEnv
+    orig_env = env_lambda()
     resize_height_factor = 2
     resize_width_factor = 3
     orig_height, orig_width = orig_env.observation_space.shape[:2]
     env, obs, _, _ = self.init_batch_and_play(
-        env_name, steps_per_epoch=1,
+        env_lambda, steps_per_epoch=1,
         resize_height_factor=resize_height_factor,
         resize_width_factor=resize_width_factor)
     for obs_batch in obs:
@@ -238,22 +233,21 @@ def assert_channels(self, env, obs, n_channels):
       self.assertEqual(ob.shape[2], n_channels)
 
   def test_channels(self):
-    env_name = TEST_ENV_NAME
-    env, obs, _, _ = self.init_batch_and_play(env_name, grayscale=True)
+    env_lambda = TestEnv
+    env, obs, _, _ = self.init_batch_and_play(env_lambda, grayscale=True)
     self.assert_channels(env, obs, n_channels=1)
 
-    env, obs, _, _ = self.init_batch_and_play(env_name, grayscale=False)
+    env, obs, _, _ = self.init_batch_and_play(env_lambda, grayscale=False)
     self.assert_channels(env, obs, n_channels=3)
 
   def test_generating_and_loading_preserves_rollouts(self):
-    env_name = TEST_ENV_NAME
-    from_env = gym_env.T2TGymEnv(env_name, batch_size=1)
-    from_env.start_new_epoch(0, self.out_dir)
+    from_env = gym_env.T2TGymEnv([TestEnv()], self.out_dir)
+    from_env.start_new_epoch(0)
     self.play(from_env, n_steps=20)
-    from_env.generate_data(self.out_dir)
+    from_env.generate_data()
 
-    to_env = gym_env.T2TGymEnv(env_name, batch_size=1)
-    to_env.start_new_epoch(0, self.out_dir)
+    to_env = gym_env.T2TGymEnv([TestEnv()], self.out_dir)
+    to_env.start_new_epoch(0)
 
     self.assertEqual(
         from_env.current_epoch_rollouts(), to_env.current_epoch_rollouts()
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index e04240f95..90fb103b7 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -33,6 +33,7 @@
 import os
 import time
 
+import gym
 import numpy as np
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
@@ -245,20 +246,31 @@ def train_world_model(env, data_dir, output_dir, hparams, epoch):
   )
 
 
-def setup_env(hparams):
-  """Setup."""
+def make_gym_env(hparams):
+  """Make env."""
   game_mode = "Deterministic-v4"
   camel_game_name = "".join(
       [w[0].upper() + w[1:] for w in hparams.game.split("_")])
   camel_game_name += game_mode
   env_name = camel_game_name
+  env = gym.make(env_name)
+  if hparams.env_timesteps_limit != -1:
+    # Replace TimeLimit Wrapper with one of proper time step limit.
+    if isinstance(env, gym.wrappers.TimeLimit):
+      env = env.env
+    env = gym.wrappers.TimeLimit(env,
+                                 max_episode_steps=hparams.env_timesteps_limit)
+  return env
 
-  env = T2TGymEnv(base_env_name=env_name,
-                  batch_size=hparams.real_ppo_num_agents,
+
+def setup_env(hparams, data_dir):
+  """Setup."""
+  env = T2TGymEnv([make_gym_env(hparams)
+                   for _ in range(hparams.real_ppo_num_agents)],
+                  data_dir,
                   grayscale=hparams.grayscale,
                   resize_width_factor=hparams.resize_width_factor,
-                  resize_height_factor=hparams.resize_height_factor,
-                  base_env_timesteps_limit=hparams.env_timesteps_limit)
+                  resize_height_factor=hparams.resize_height_factor)
   return env
 
 
@@ -288,8 +300,8 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   epoch = -1
   data_dir = directories["data"]
-  env = setup_env(hparams)
-  env.start_new_epoch(epoch, data_dir)
+  env = setup_env(hparams, data_dir)
+  env.start_new_epoch(epoch)
 
   # Timing log function
   log_relative_time = make_relative_timing_fn()
@@ -330,7 +342,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   )
 
   for epoch in range(hparams.epochs):
-    env.generate_data(data_dir)
+    env.generate_data()
 
     is_final_epoch = (epoch + 1) == hparams.epochs
     log = make_log_fn(epoch, log_relative_time)
@@ -352,7 +364,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
                 ppo_event_dir, directories["world_model"], data_dir,
                 hparams, epoch=epoch, is_final_epoch=is_final_epoch)
 
-    env.start_new_epoch(epoch, data_dir)
+    env.start_new_epoch(epoch)
 
     # Train PPO on real env (short)
     log("Training PPO in real environment.")

From e2e2e66d1e0e206a7082831b08f61238b312ad34 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Wed, 17 Oct 2018 14:46:31 -0700
Subject: [PATCH 1030/2720] internal merge of PR #1149

PiperOrigin-RevId: 217591575
---
 tensor2tensor/data_generators/gym_env.py      | 79 ++++++++++++-------
 tensor2tensor/data_generators/gym_env_test.py | 63 ++++++++-------
 tensor2tensor/rl/envs/batch_env_factory.py    |  4 +-
 tensor2tensor/rl/trainer_model_based.py       | 32 +++-----
 4 files changed, 96 insertions(+), 82 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index c43fb90a5..1fda3d0a5 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -22,6 +22,7 @@
 import collections
 import itertools
 
+import gym
 from gym.spaces import Box
 import numpy as np
 
@@ -48,6 +49,17 @@ def __deepcopy__(self, memo):
     return self
 
 
+def make_gym_env(name, timesteps_limit=-1):
+  env = gym.make(name)
+  if timesteps_limit != -1:
+    # Replace TimeLimit Wrapper with one of proper time step limit.
+    if isinstance(env, gym.wrappers.TimeLimit):
+      env = env.env
+    env = gym.wrappers.TimeLimit(env,
+                                 max_episode_steps=timesteps_limit)
+  return env
+
+
 class T2TEnv(video_utils.VideoProblem):
   """Abstract class representing a batch of environments.
 
@@ -71,12 +83,11 @@ class T2TEnv(video_utils.VideoProblem):
   reward_range = (-1, 1)
   name = None
 
-  def __init__(self, batch_size, data_dir):
-    super(T2TEnv, self).__init__()
+  def __init__(self, batch_size, *args, **kwargs):
+    super(T2TEnv, self).__init__(*args, **kwargs)
 
     self.clear_history()
     self.batch_size = batch_size
-    self.data_dir = data_dir
     self._current_batch_frames = [None for _ in range(batch_size)]
     self._current_batch_rollouts = [[] for _ in range(batch_size)]
     self._current_epoch_rollouts = []
@@ -100,7 +111,7 @@ def clear_history(self):
     """Clears the rollout history."""
     self._rollouts_by_epoch_and_split = collections.OrderedDict()
 
-  def start_new_epoch(self, epoch, load_data=True):
+  def start_new_epoch(self, epoch, load_data_dir=None):
     if not isinstance(epoch, int):
       raise ValueError("Epoch should be integer, got {}".format(epoch))
     if epoch in self._rollouts_by_epoch_and_split:
@@ -108,8 +119,8 @@ def start_new_epoch(self, epoch, load_data=True):
     self.current_epoch = epoch
     self._current_epoch_rollouts = []
     self._rollouts_by_epoch_and_split[epoch] = collections.defaultdict(list)
-    if load_data:
-      self._load_epoch_data()
+    if load_data_dir is not None:
+      self._load_epoch_data(load_data_dir)
 
   def current_epoch_rollouts(self, split=None):
     rollouts_by_split = self._rollouts_by_epoch_and_split[self.current_epoch]
@@ -227,14 +238,14 @@ def reset(self, indices=None):
 
     Returns:
       Batch of initial observations of reset environments.
+
+    Raises:
+      ValueError: when there's no current epoch.
     """
     if self.current_epoch is None:
-      # It's here so that the old pipeline works.
-      self.start_new_epoch(0, load_data=False)
-      # TODO(koz4k): Replace with:
-      # raise ValueError(
-      #     "No current epoch. start_new_epoch() should first be called."
-      # )
+      raise ValueError(
+          "No current epoch. start_new_epoch() should first be called."
+      )
 
     if indices is None:
       indices = np.arange(self.batch_size)
@@ -386,8 +397,7 @@ def split_size(split_index):
     self._rollouts_by_epoch_and_split[self.current_epoch] = rollouts_by_split
     self._current_epoch_rollouts = []
 
-  @property
-  def splits_and_paths(self):
+  def splits_and_paths(self, data_dir):
     """List of pairs (split, paths) for the current epoch."""
     filepath_fns = {
         problem.DatasetSplit.TRAIN: self.training_filepaths,
@@ -404,7 +414,7 @@ def append_epoch(paths):
     # We set shuffled=True as we don't want to shuffle on disk later.
     return [
         (split["split"], append_epoch(filepath_fns[split["split"]](
-            self.data_dir, split["shards"], shuffled=True
+            data_dir, split["shards"], shuffled=True
         )))
         for split in self.dataset_splits
     ]
@@ -417,14 +427,14 @@ def filepattern(self, data_dir, mode, shard=None, only_last=False):
       filepattern += ".{}".format(self.current_epoch)
     return filepattern
 
-  def generate_data(self, data_dir=None, tmp_dir=None, task_id=-1):
+  def generate_data(self, data_dir, tmp_dir=None, task_id=-1):
     """Saves the current epoch rollouts to disk, split into train/dev sets."""
     if not self._rollouts_by_epoch_and_split[self.current_epoch]:
       # Data not loaded from disk.
       self._split_current_epoch()
 
     rollouts_by_split = self._rollouts_by_epoch_and_split[self.current_epoch]
-    splits_and_paths = self.splits_and_paths
+    splits_and_paths = self.splits_and_paths(data_dir)
 
     for (split, paths) in splits_and_paths:
       rollouts = rollouts_by_split[split]
@@ -442,12 +452,12 @@ def generate_data(self, data_dir=None, tmp_dir=None, task_id=-1):
             cycle_every_n=float("inf")
         )
 
-  def _load_epoch_data(self):
+  def _load_epoch_data(self, data_dir):
     any_files_found = False
     all_files_found = True
     any_shard_empty = False
 
-    for split, paths in self.splits_and_paths:
+    for split, paths in self.splits_and_paths(data_dir):
       try:
         any_shard_empty |= self._load_epoch_split(split, paths)
         any_files_found = True
@@ -506,28 +516,39 @@ def get_feature_value(key, list_name):
 
 
 class T2TGymEnv(T2TEnv):
-  """Class representing a batch of Gym environments."""
+  """Class representing a batch of Gym environments.
 
-  name = "t2t_gym_env"
+  Do not register it, instead create subclass with hardcoded __init__
+  arguments and register this subclass.
+  """
 
-  def __init__(self, envs, data_dir, grayscale=False,
-               resize_height_factor=1, resize_width_factor=1):
-    super(T2TGymEnv, self).__init__(len(envs), data_dir)
+  def __init__(self, base_env_name, batch_size=None, grayscale=False,
+               resize_height_factor=1, resize_width_factor=1,
+               base_env_timesteps_limit=-1, envs=None, **kwargs):
+    if batch_size is None:
+      batch_size = len(envs)
+    super(T2TGymEnv, self).__init__(batch_size, **kwargs)
     self.grayscale = grayscale
     self.resize_height_factor = resize_height_factor
     self.resize_width_factor = resize_width_factor
-    if not envs:
-      raise ValueError("Must have at least one environment.")
-    self._envs = envs
+    if not self.name:
+      # Set problem name if not registered.
+      self.name = "t2t_gym_env_{}".format(base_env_name)
+
+    if envs is None:
+      self._envs = [make_gym_env(base_env_name, base_env_timesteps_limit)
+                    for _ in range(self.batch_size)]
+    else:
+      self._envs = envs
 
-    orig_observ_space = envs[0].observation_space
+    orig_observ_space = self._envs[0].observation_space
     if not all(env.observation_space == orig_observ_space
                for env in self._envs):
       raise ValueError("All environments must use the same observation space.")
 
     self.observation_space = self._derive_observation_space(orig_observ_space)
 
-    self.action_space = envs[0].action_space
+    self.action_space = self._envs[0].action_space
     if not all(env.action_space == self.action_space for env in self._envs):
       raise ValueError("All environments must use the same action space.")
 
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 61b3bd24d..15536deb9 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -29,6 +29,7 @@
 
 from tensor2tensor.data_generators import gym_env
 from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators.gym_env import make_gym_env
 
 import tensorflow as tf
 
@@ -61,6 +62,10 @@ def step(self, action):
   def reset(self):
     return self._generate_ob()
 
+TEST_ENV_NAME = "T2TTestEnv-v1"
+
+gym.envs.register(id=TEST_ENV_NAME, entry_point=TestEnv)
+
 
 class GymEnvTest(tf.test.TestCase):
 
@@ -75,19 +80,18 @@ def setUp(self):
     shutil.rmtree(self.out_dir)
     os.mkdir(self.out_dir)
 
-  def init_batch_and_play(self, env_lambda, steps_per_epoch=1,
+  def init_batch_and_play(self, env_name, steps_per_epoch=1,
                           epochs=(0,), generate_data=False, **kwargs):
-    raw_envs = [env_lambda(), env_lambda()]
-    env = gym_env.T2TGymEnv(raw_envs, self.out_dir, **kwargs)
+    env = gym_env.T2TGymEnv(env_name, batch_size=2, **kwargs)
     obs = list()
     rewards = list()
     num_dones = 0
     for epoch in epochs:
-      env.start_new_epoch(epoch)
+      env.start_new_epoch(epoch, self.out_dir)
       _, epoch_obs, epoch_rewards, epoch_num_dones = \
           self.play(env, steps_per_epoch)
       if generate_data:
-        env.generate_data()
+        env.generate_data(self.out_dir)
       obs.extend(epoch_obs)
       rewards.extend(epoch_rewards)
       num_dones += epoch_num_dones
@@ -110,7 +114,7 @@ def play(self, env, n_steps):
 
   def test_splits_dataset(self):
     env, _, _, _ = self.init_batch_and_play(
-        TestEnv, steps_per_epoch=20, generate_data=True
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
     )
 
     for split in self.splits:
@@ -118,7 +122,7 @@ def test_splits_dataset(self):
 
   def test_split_preserves_number_of_rollouts(self):
     env, _, _, num_dones = self.init_batch_and_play(
-        TestEnv, steps_per_epoch=20, generate_data=True
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
     )
 
     num_rollouts_after_split = sum(
@@ -131,7 +135,7 @@ def test_split_preserves_number_of_rollouts(self):
 
   def test_split_preserves_number_of_frames(self):
     env, _, _, num_dones = self.init_batch_and_play(
-        TestEnv, steps_per_epoch=20, generate_data=True
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
     )
 
     num_frames = sum(
@@ -146,7 +150,7 @@ def test_split_preserves_number_of_frames(self):
   def test_generates_data(self):
     # This test needs base env which outputs done after two steps.
     self.init_batch_and_play(
-        TestEnv, steps_per_epoch=20, generate_data=True
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
     )
 
     filenames = os.listdir(self.out_dir)
@@ -162,18 +166,18 @@ def num_ending_with(filenames, suffix):
           1 for filename in filenames if filename.endswith(suffix)
       )
 
-    env = gym_env.T2TGymEnv([TestEnv() for _ in range(2)], self.out_dir)
-    env.start_new_epoch(0)
+    env = gym_env.T2TGymEnv(TEST_ENV_NAME, batch_size=2)
+    env.start_new_epoch(0, self.out_dir)
     self.play(env, n_steps=20)
-    env.generate_data()
+    env.generate_data(self.out_dir)
 
     filenames = os.listdir(self.out_dir)
     num_shards_per_epoch = len(filenames)
     self.assertEqual(num_ending_with(filenames, ".0"), num_shards_per_epoch)
 
-    env.start_new_epoch(1)
+    env.start_new_epoch(1, self.out_dir)
     self.play(env, n_steps=20)
-    env.generate_data()
+    env.generate_data(self.out_dir)
 
     filenames = os.listdir(self.out_dir)
     self.assertEqual(len(filenames), 2 * num_shards_per_epoch)
@@ -182,14 +186,14 @@ def num_ending_with(filenames, suffix):
 
   def test_frame_numbers_are_continuous(self):
     env, _, _, _ = self.init_batch_and_play(
-        TestEnv, steps_per_epoch=20, generate_data=True
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
     )
 
     frame_numbers = [
         tf.train.Example.FromString(
             record
         ).features.feature["frame_number"].int64_list.value[0]
-        for (_, paths) in env.splits_and_paths
+        for (_, paths) in env.splits_and_paths(self.out_dir)
         for path in paths
         for record in tf.python_io.tf_record_iterator(path)
     ]
@@ -203,19 +207,19 @@ def test_frame_numbers_are_continuous(self):
       last_frame_number = frame_number
 
   def test_clipping(self):
-    env_lambda = TestEnv
-    _, _, rewards, _ = self.init_batch_and_play(env_lambda, steps_per_epoch=2)
+    _, _, rewards, _ = self.init_batch_and_play(TEST_ENV_NAME,
+                                                steps_per_epoch=2)
     self.assertTrue(np.max(rewards) == 1)
     self.assertTrue(np.min(rewards) == -1)
 
   def test_resize(self):
-    env_lambda = TestEnv
-    orig_env = env_lambda()
+    env_name = TEST_ENV_NAME
+    orig_env = make_gym_env(env_name)
     resize_height_factor = 2
     resize_width_factor = 3
     orig_height, orig_width = orig_env.observation_space.shape[:2]
     env, obs, _, _ = self.init_batch_and_play(
-        env_lambda, steps_per_epoch=1,
+        env_name, steps_per_epoch=1,
         resize_height_factor=resize_height_factor,
         resize_width_factor=resize_width_factor)
     for obs_batch in obs:
@@ -233,21 +237,22 @@ def assert_channels(self, env, obs, n_channels):
       self.assertEqual(ob.shape[2], n_channels)
 
   def test_channels(self):
-    env_lambda = TestEnv
-    env, obs, _, _ = self.init_batch_and_play(env_lambda, grayscale=True)
+    env_name = TEST_ENV_NAME
+    env, obs, _, _ = self.init_batch_and_play(env_name, grayscale=True)
     self.assert_channels(env, obs, n_channels=1)
 
-    env, obs, _, _ = self.init_batch_and_play(env_lambda, grayscale=False)
+    env, obs, _, _ = self.init_batch_and_play(env_name, grayscale=False)
     self.assert_channels(env, obs, n_channels=3)
 
   def test_generating_and_loading_preserves_rollouts(self):
-    from_env = gym_env.T2TGymEnv([TestEnv()], self.out_dir)
-    from_env.start_new_epoch(0)
+    env_name = TEST_ENV_NAME
+    from_env = gym_env.T2TGymEnv(env_name, batch_size=1)
+    from_env.start_new_epoch(0, self.out_dir)
     self.play(from_env, n_steps=20)
-    from_env.generate_data()
+    from_env.generate_data(self.out_dir)
 
-    to_env = gym_env.T2TGymEnv([TestEnv()], self.out_dir)
-    to_env.start_new_epoch(0)
+    to_env = gym_env.T2TGymEnv(env_name, batch_size=1)
+    to_env.start_new_epoch(0, self.out_dir)
 
     self.assertEqual(
         from_env.current_epoch_rollouts(), to_env.current_epoch_rollouts()
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index dafe3f31f..1b7d6930b 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -56,8 +56,8 @@ def _define_batch_env(environment_spec, num_agents):
     envs = [
         environment_spec.env_lambda()
         for _ in range(num_agents)]
-    # We won't generate data from this env so it's safe to set data_dir to None.
-    env = gym_env.T2TGymEnv(envs, data_dir=None)
+    env = gym_env.T2TGymEnv("unknown", envs=envs)
+    env.start_new_epoch(0)
     return env
 
 
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 90fb103b7..e04240f95 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -33,7 +33,6 @@
 import os
 import time
 
-import gym
 import numpy as np
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
@@ -246,31 +245,20 @@ def train_world_model(env, data_dir, output_dir, hparams, epoch):
   )
 
 
-def make_gym_env(hparams):
-  """Make env."""
+def setup_env(hparams):
+  """Setup."""
   game_mode = "Deterministic-v4"
   camel_game_name = "".join(
       [w[0].upper() + w[1:] for w in hparams.game.split("_")])
   camel_game_name += game_mode
   env_name = camel_game_name
-  env = gym.make(env_name)
-  if hparams.env_timesteps_limit != -1:
-    # Replace TimeLimit Wrapper with one of proper time step limit.
-    if isinstance(env, gym.wrappers.TimeLimit):
-      env = env.env
-    env = gym.wrappers.TimeLimit(env,
-                                 max_episode_steps=hparams.env_timesteps_limit)
-  return env
 
-
-def setup_env(hparams, data_dir):
-  """Setup."""
-  env = T2TGymEnv([make_gym_env(hparams)
-                   for _ in range(hparams.real_ppo_num_agents)],
-                  data_dir,
+  env = T2TGymEnv(base_env_name=env_name,
+                  batch_size=hparams.real_ppo_num_agents,
                   grayscale=hparams.grayscale,
                   resize_width_factor=hparams.resize_width_factor,
-                  resize_height_factor=hparams.resize_height_factor)
+                  resize_height_factor=hparams.resize_height_factor,
+                  base_env_timesteps_limit=hparams.env_timesteps_limit)
   return env
 
 
@@ -300,8 +288,8 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   epoch = -1
   data_dir = directories["data"]
-  env = setup_env(hparams, data_dir)
-  env.start_new_epoch(epoch)
+  env = setup_env(hparams)
+  env.start_new_epoch(epoch, data_dir)
 
   # Timing log function
   log_relative_time = make_relative_timing_fn()
@@ -342,7 +330,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   )
 
   for epoch in range(hparams.epochs):
-    env.generate_data()
+    env.generate_data(data_dir)
 
     is_final_epoch = (epoch + 1) == hparams.epochs
     log = make_log_fn(epoch, log_relative_time)
@@ -364,7 +352,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
                 ppo_event_dir, directories["world_model"], data_dir,
                 hparams, epoch=epoch, is_final_epoch=is_final_epoch)
 
-    env.start_new_epoch(epoch)
+    env.start_new_epoch(epoch, data_dir)
 
     # Train PPO on real env (short)
     log("Training PPO in real environment.")

From b37543b01e703bef38f8ac9202986d85328d23e2 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 18 Oct 2018 01:02:30 +0200
Subject: [PATCH 1031/2720] Take initial frames for simulated env from memory
 (#1151)

* Take initial frames for simulated env from memory

* Fix maximum recursion depth exceeded
---
 tensor2tensor/models/research/rl.py           | 62 ++++---------
 tensor2tensor/rl/collect.py                   |  7 +-
 tensor2tensor/rl/envs/batch_env_factory.py    | 42 ++-------
 tensor2tensor/rl/envs/simulated_batch_env.py  | 54 +++++------
 tensor2tensor/rl/envs/utils.py                | 31 +------
 tensor2tensor/rl/trainer_model_based.py       | 90 +++++++++++++------
 .../rl/trainer_model_based_params.py          |  2 +-
 7 files changed, 107 insertions(+), 181 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 8f120688e..a68d65086 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -143,71 +143,41 @@ def simple_gym_spec(env):
                                      simulated_env=False)
 
 
-def standard_atari_env_spec(
-    env=None, simulated=False, resize_height_factor=1, resize_width_factor=1,
-    grayscale=False, include_clipping=True, batch_env=None):
+def standard_atari_env_spec(env=None, simulated=False):
   """Parameters of environment specification."""
-  resize_wrapper = [tf_atari_wrappers.ResizeWrapper,
-                    {"height_factor": resize_height_factor,
-                     "width_factor": resize_width_factor,
-                     "grayscale": grayscale}]
-  if include_clipping:
-    standard_wrappers = [
-        resize_wrapper,
-        [tf_atari_wrappers.RewardClippingWrapper, {}],
-        [tf_atari_wrappers.StackWrapper, {"history": 4}],
-    ]
-  else:
-    standard_wrappers = [
-        resize_wrapper,
-        [tf_atari_wrappers.StackWrapper, {"history": 4}],
-    ]
-  if simulated:  # No resizing on simulated environments.
-    standard_wrappers = standard_wrappers[1:]
+  standard_wrappers = [
+      (tf_atari_wrappers.StackWrapper, {"history": 4})
+  ]
 
   env_spec = tf.contrib.training.HParams(
       wrappers=standard_wrappers,
-      simulated_env=simulated)
-
-  if batch_env is not None:
-    env_spec.add_hparam("batch_env", batch_env)
-  else:
-    env_lambda = None
-    if isinstance(env, str):
-      env_lambda = lambda: gym.make(env)
-    if callable(env):
-      env_lambda = env
-    assert env_lambda is not None, "Unknown specification of environment"
-    env_spec.add_hparam("env_lambda", env_lambda)
+      simulated_env=simulated,
+      reward_range=env.reward_range,
+      observation_space=env.observation_space,
+      action_space=env.action_space
+  )
+  if not simulated:
+    env_spec.add_hparam("env", env)
 
   return env_spec
 
 
 def standard_atari_env_simulated_spec(
-    real_env, video_num_input_frames, video_num_target_frames):
+    real_env, video_num_input_frames, video_num_target_frames
+):
   """Spec."""
-  env_spec = standard_atari_env_spec(
-      # This hack is here because SimulatedBatchEnv needs to get
-      # observation_space from the real env. TODO(koz4k): refactor.
-      env=lambda: real_env,
-      simulated=True
-  )
+  env_spec = standard_atari_env_spec(real_env, simulated=True)
   env_spec.add_hparam("simulation_random_starts", True)
   env_spec.add_hparam("simulation_flip_first_random_for_beginning", True)
   env_spec.add_hparam("intrinsic_reward_scale", 0.0)
-  env_spec.add_hparam("initial_frames_problem", real_env)
   env_spec.add_hparam("video_num_input_frames", video_num_input_frames)
   env_spec.add_hparam("video_num_target_frames", video_num_target_frames)
   return env_spec
 
 
-def standard_atari_env_eval_spec(
-    env, simulated=False, resize_height_factor=1, resize_width_factor=1,
-    grayscale=False):
+def standard_atari_env_eval_spec(*args, **kwargs):
   """Parameters of environment specification for eval."""
-  return standard_atari_env_spec(
-      env, simulated, resize_height_factor, resize_width_factor, grayscale,
-      include_clipping=False)
+  return standard_atari_env_spec(*args, **kwargs)
 
 
 def standard_atari_ae_env_spec(env, ae_hparams_set):
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index f39ecde05..e925cd9f4 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -112,12 +112,7 @@ def define_collect(hparams, scope, eval_phase,
       environment_spec = getattr(hparams, "environment_eval_spec",
                                  environment_spec)
       num_agents = getattr(hparams, "num_eval_agents", num_agents)
-      batch_env = batch_env_factory(environment_spec, num_agents)
-    else:
-      initial_frame_chooser = getattr(hparams, "initial_frame_chooser", None)
-      batch_env = batch_env_factory(
-          environment_spec, num_agents,
-          initial_frame_chooser=initial_frame_chooser)
+    batch_env = batch_env_factory(environment_spec, num_agents)
 
     to_initialize.append(batch_env)
     environment_wrappers = environment_spec.wrappers
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index 1b7d6930b..fdf276c39 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -23,47 +23,17 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.data_generators import gym_env
 from tensor2tensor.rl.envs import py_func_batch_env
 from tensor2tensor.rl.envs import simulated_batch_env
 
-import tensorflow as tf
 
-
-def batch_env_factory(environment_spec, num_agents, initial_frame_chooser=None):
+def batch_env_factory(environment_spec, num_agents):
   """Factory of batch envs."""
-  # TODO(konradczechowski): this is temporary function handling both old and
-  # new pipelines, refactor this when we move to the new pipeline.
   if environment_spec.simulated_env:
-    cur_batch_env = _define_simulated_batch_env(
-        environment_spec, num_agents, initial_frame_chooser)
+    cur_batch_env = simulated_batch_env.SimulatedBatchEnv(
+        environment_spec, num_agents
+    )
   else:
-    if "batch_env" in environment_spec:
-      msg = "Environment_spec should contain only 1 of (env_lambda, batch_env)."
-      assert "env_lambda" not in environment_spec, msg
-      batch_env = environment_spec.batch_env
-      assert batch_env.batch_size == num_agents
-    else:
-      batch_env = _define_batch_env(environment_spec, num_agents)
-    cur_batch_env = py_func_batch_env.PyFuncBatchEnv(batch_env)
-  return cur_batch_env
-
-
-def _define_batch_env(environment_spec, num_agents):
-  """Create environments and apply all desired wrappers."""
-
-  with tf.variable_scope("environments"):
-    envs = [
-        environment_spec.env_lambda()
-        for _ in range(num_agents)]
-    env = gym_env.T2TGymEnv("unknown", envs=envs)
-    env.start_new_epoch(0)
-    return env
-
+    cur_batch_env = py_func_batch_env.PyFuncBatchEnv(environment_spec.env)
 
-def _define_simulated_batch_env(environment_spec, num_agents,
-                                initial_frame_chooser):
-  cur_batch_env = simulated_batch_env.SimulatedBatchEnv(
-      environment_spec, num_agents, initial_frame_chooser
-  )
-  return cur_batch_env
+  return cur_batch_env
\ No newline at end of file
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 5b245d1b0..72b3d721d 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -24,7 +24,6 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.rl.envs import in_graph_batch_env
-from tensor2tensor.rl.envs import utils
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
@@ -38,24 +37,18 @@
 class HistoryBuffer(object):
   """History Buffer."""
 
-  def __init__(self, initial_frame_chooser, length, observ_dtype):
-    initial_frame_chooser.batch_size = length
-    self._initial_frame_chooser = initial_frame_chooser
+  def __init__(self, initial_frame_chooser, observ_shape, observ_dtype,
+               num_initial_frames, length):
     self.length = length
     self._observ_dtype = observ_dtype
-    initial_frames = self.get_initial_observations()
-    initial_shape = [length] + common_layers.shape_list(initial_frames)[1:]
+    initial_shape = (length, num_initial_frames) + observ_shape
+    self._initial_frames = tf.py_func(
+        initial_frame_chooser, [tf.constant(length)], observ_dtype
+    )
+    self._initial_frames.set_shape(initial_shape)
     self._history_buff = tf.Variable(tf.zeros(initial_shape, observ_dtype),
                                      trainable=False)
 
-  def initialize(self, sess):
-    self._initial_frame_chooser.initialize(sess)
-
-  def get_initial_observations(self):
-    return tf.cast(
-        self._initial_frame_chooser.choose()["inputs"], self._observ_dtype
-    )
-
   def get_all_elements(self):
     return self._history_buff.read_value()
 
@@ -68,7 +61,7 @@ def move_by_one_element(self, element):
         return self._history_buff.read_value()
 
   def reset(self, indices):
-    initial_frames = tf.gather(self.get_initial_observations(), indices)
+    initial_frames = tf.gather(self._initial_frames, indices)
     scatter_op = tf.scatter_update(self._history_buff, indices, initial_frames)
     with tf.control_dependencies([scatter_op]):
       return self._history_buff.read_value()
@@ -99,24 +92,14 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   flags are held in according variables.
   """
 
-  def __init__(self, environment_spec, length, initial_frame_chooser):
+  def __init__(self, environment_spec, length):
     """Batch of environments inside the TensorFlow graph."""
-
-    observ_space = utils.get_observation_space(environment_spec)
-    initial_frames_problem = environment_spec.initial_frames_problem
-    self._frames_problem_name = str(initial_frames_problem)
-    observ_shape = (initial_frames_problem.frame_height,
-                    initial_frames_problem.frame_width,
-                    initial_frames_problem.num_channels)
-    observ_space.shape = observ_shape
-    action_space = utils.get_action_space(environment_spec)
-    super(SimulatedBatchEnv, self).__init__(observ_space, action_space)
+    super(SimulatedBatchEnv, self).__init__(
+        environment_spec.observation_space, environment_spec.action_space
+    )
 
     self.length = length
-    try:
-      self._min_reward = initial_frames_problem.min_reward
-    except AttributeError:
-      self._min_reward = initial_frames_problem.reward_range[0]
+    self._min_reward = environment_spec.reward_range[0]
     self._num_frames = environment_spec.video_num_input_frames
     self._intrinsic_reward_scale = environment_spec.intrinsic_reward_scale
 
@@ -128,17 +111,20 @@ def __init__(self, environment_spec, length, initial_frame_chooser):
         model_hparams, tf.estimator.ModeKeys.PREDICT)
 
     self.history_buffer = HistoryBuffer(
-        initial_frame_chooser, self.length, self.observ_dtype)
+        environment_spec.initial_frame_chooser, self.observ_shape,
+        self.observ_dtype, self._num_frames, self.length
+    )
 
     self._observ = tf.Variable(
-        tf.zeros((len(self),) + observ_shape, self.observ_dtype),
+        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
         trainable=False)
 
   def initialize(self, sess):
-    self.history_buffer.initialize(sess)
+    # Currently not needed. Keeping it just in case.
+    pass
 
   def __str__(self):
-    return "SimulatedEnv(%s)" % self._frames_problem_name
+    return "SimulatedEnv"
 
   def __len__(self):
     """Number of combined environments."""
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
index ee30df3cb..503dd58f4 100644
--- a/tensor2tensor/rl/envs/utils.py
+++ b/tensor2tensor/rl/envs/utils.py
@@ -67,35 +67,6 @@ def _reset(self, **kwargs):
     return self._last_returned[0]
 
 
-def get_observation_space(environment_spec):
-  """Get observation space associated with environment spec.
-
-  Args:
-     environment_spec:  EnvironmentSpec object
-
-  Returns:
-    OpenAi Gym observation space
-  """
-  return environment_spec.env_lambda().observation_space
-
-
-def get_action_space(environment_spec):
-  """Get action space associated with environment spec.
-
-  Args:
-     environment_spec:  Object consisting one of batch_env.action_space, or
-     env_lambda().action_space
-
-  Returns:
-    OpenAi Gym action space
-  """
-  if "batch_env" in environment_spec:
-    action_space = environment_spec.batch_env.action_space
-  else:
-    action_space = environment_spec.env_lambda().action_space
-  return action_space
-
-
 def get_policy(observations, hparams):
   """Get a policy network.
 
@@ -107,7 +78,7 @@ def get_policy(observations, hparams):
     Tensor with policy and value function output
   """
   policy_network_lambda = hparams.policy_network
-  action_space = get_action_space(hparams.environment_spec)
+  action_space = hparams.environment_spec.action_space
   return policy_network_lambda(action_space, hparams, observations)
 
 
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index e04240f95..6d20182bb 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -17,7 +17,7 @@
 
 Example invocation:
 
-python -m tensor2tensor.rl.trainer_model_based_new \
+python -m tensor2tensor.rl.trainer_model_based \
     --output_dir=$HOME/t2t/rl_v1 \
     --loop_hparams_set=rlmb_base \
     --loop_hparams='num_real_env_frames=10000,epochs=3'
@@ -31,6 +31,7 @@
 import datetime
 import math
 import os
+import random
 import time
 
 import numpy as np
@@ -40,7 +41,6 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.rl import trainer_model_based_params
-from tensor2tensor.rl.envs.utils import InitialFrameChooser
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -139,7 +139,7 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
   getattr(exp, schedule)()
 
 
-def train_agent(environment_spec, agent_model_dir,
+def train_agent(real_env, environment_spec, agent_model_dir,
                 event_dir, world_model_dir, data_dir, hparams, epoch=0,
                 is_final_epoch=False):
   """Train the PPO agent in the simulated environment."""
@@ -164,28 +164,62 @@ def train_agent(environment_spec, agent_model_dir,
   ppo_hparams.add_hparam("model_hparams", model_hparams)
 
   environment_spec = copy.copy(environment_spec)
-  environment_spec_param_names = [
-      "simulation_random_starts", "simulation_flip_first_random_for_beginning",
-      "intrinsic_reward_scale"
-  ]
+  environment_spec_param_names = ["intrinsic_reward_scale"]
   for param_name in environment_spec_param_names:
     environment_spec.set_hparam(param_name, hparams.get(param_name))
-  ppo_hparams.add_hparam("environment_spec", environment_spec)
-
-  ppo_hparams.add_hparam("initial_frame_chooser", InitialFrameChooser(
-      environment_spec, mode=tf.estimator.ModeKeys.EVAL
-  ))
 
-  # TODO(koz4k): Pass by arguments.
-  with temporary_flags({
-      "problem": environment_spec.initial_frames_problem,
-      "model": hparams.generative_model,
-      "hparams_set": hparams.generative_model_params,
-      "output_dir": world_model_dir,
-      "data_dir": data_dir,
-  }):
-    rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
-                         name_scope="ppo_sim%d" % (epoch + 1))
+  with tf.Session() as sess:
+    encoded_png_p = tf.placeholder(tf.string)
+    decoded_png_t = tf.image.decode_png(encoded_png_p)
+    def decode_png(encoded_png):
+      return sess.run(decoded_png_t, feed_dict={encoded_png_p: encoded_png})
+
+    initial_frame_rollouts = real_env.current_epoch_rollouts(
+        split=tf.contrib.learn.ModeKeys.TRAIN
+    )
+    # TODO(koz4k): Move this to a different module.
+    def initial_frame_chooser(batch_size):
+      num_frames = environment_spec.video_num_input_frames
+      deterministic_initial_frames = initial_frame_rollouts[0][:num_frames]
+      if not environment_spec.simulation_random_starts:
+        # Deterministic starts: repeat first frames from the first rollout.
+        initial_frames = [deterministic_initial_frames] * batch_size
+      else:
+        # Random starts: choose random initial frames from random rollouts.
+        # TODO(koz4k): Weigh rollouts by their lengths so sampling is uniform
+        # over frames and not rollouts.
+        def choose_initial_frames():
+          try:
+            rollout = random.choice(initial_frame_rollouts)
+            from_index = random.randrange(len(rollout) - num_frames)
+            return rollout[from_index:(from_index + num_frames)]
+          except ValueError:
+            # Rollout too short; repeat.
+            return choose_initial_frames()
+        initial_frames = [choose_initial_frames() for _ in range(batch_size)]
+        if environment_spec.simulation_flip_first_random_for_beginning:
+          # Flip first entry in the batch for deterministic initial frames.
+          initial_frames[0] = deterministic_initial_frames
+
+        return np.stack([
+            [decode_png(frame.observation) for frame in initial_frame_stack]
+            for initial_frame_stack in initial_frames
+        ])
+
+    environment_spec.add_hparam("initial_frame_chooser", initial_frame_chooser)
+
+    ppo_hparams.add_hparam("environment_spec", environment_spec)
+
+    # TODO(koz4k): Pass by arguments.
+    with temporary_flags({
+        "problem": real_env,
+        "model": hparams.generative_model,
+        "hparams_set": hparams.generative_model_params,
+        "output_dir": world_model_dir,
+        "data_dir": data_dir,
+    }):
+      rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
+                           name_scope="ppo_sim%d" % (epoch + 1))
 
 
 def train_agent_real_env(
@@ -211,9 +245,7 @@ def train_agent_real_env(
   # But we need to save at the last step, so we set it very high.
   ppo_hparams.save_models_every_epochs = 1000000
 
-  environment_spec = rl.standard_atari_env_spec(
-      batch_env=env, include_clipping=False
-  )
+  environment_spec = rl.standard_atari_env_spec(env)
 
   ppo_hparams.add_hparam("environment_spec", environment_spec)
 
@@ -348,9 +380,11 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     if not hparams.ppo_continue_training:
       ppo_model_dir = ppo_event_dir
 
-    train_agent(sim_env_spec, ppo_model_dir,
-                ppo_event_dir, directories["world_model"], data_dir,
-                hparams, epoch=epoch, is_final_epoch=is_final_epoch)
+    train_agent(
+        env, sim_env_spec, ppo_model_dir, ppo_event_dir,
+        directories["world_model"], data_dir, hparams, epoch=epoch,
+        is_final_epoch=is_final_epoch
+    )
 
     env.start_new_epoch(epoch, data_dir)
 
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 25c67da8c..a7b34d7b0 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -63,7 +63,7 @@ def rlmb_base():
       simulated_env_generator_num_steps=2000,
       simulation_random_starts=True,  # Use random starts in PPO.
       # Flip the first random frame in PPO batch for the true beginning.
-      simulation_flip_first_random_for_beginning=False,
+      simulation_flip_first_random_for_beginning=True,
       intrinsic_reward_scale=0.,
       ppo_epochs_num=1000,  # This should be enough to see something
       # Our simulated envs do not know how to reset.

From bdcb363158467f87e7a6a5a0dcd510c9aa233cbe Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Wed, 17 Oct 2018 16:03:24 -0700
Subject: [PATCH 1032/2720] internal merge of PR #1151

PiperOrigin-RevId: 217605857
---
 tensor2tensor/models/research/rl.py        | 15 +++++++--------
 tensor2tensor/rl/envs/batch_env_factory.py |  3 +--
 tensor2tensor/rl/trainer_model_based.py    |  1 +
 3 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index a68d65086..974052397 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -20,6 +20,7 @@
 import operator
 import gym
 
+from tensor2tensor.data_generators import gym_env
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
@@ -148,7 +149,6 @@ def standard_atari_env_spec(env=None, simulated=False):
   standard_wrappers = [
       (tf_atari_wrappers.StackWrapper, {"history": 4})
   ]
-
   env_spec = tf.contrib.training.HParams(
       wrappers=standard_wrappers,
       simulated_env=simulated,
@@ -158,7 +158,6 @@ def standard_atari_env_spec(env=None, simulated=False):
   )
   if not simulated:
     env_spec.add_hparam("env", env)
-
   return env_spec
 
 
@@ -212,7 +211,7 @@ def pong_model_free():
   hparams = tf.contrib.training.HParams(
       epochs_num=4,
       eval_every_epochs=2,
-      num_agents=10,
+      num_agents=2,
       optimization_epochs=3,
       epoch_length=30,
       entropy_loss_coef=0.003,
@@ -220,18 +219,18 @@ def pong_model_free():
       optimizer="Adam",
       policy_network=feed_forward_cnn_small_categorical_fun,
       gae_lambda=0.985,
-      num_eval_agents=1,
+      num_eval_agents=2,
       max_gradients_norm=0.5,
       gae_gamma=0.985,
       optimization_batch_size=4,
       clipping_coef=0.2,
       value_loss_coef=1,
       save_models_every_epochs=False)
-  hparams.add_hparam("environment_spec",
-                     standard_atari_env_spec("PongNoFrameskip-v4"))
+  env = gym_env.T2TGymEnv("PongNoFrameskip-v4", batch_size=2)
+  env.start_new_epoch(0)
+  hparams.add_hparam("environment_spec", standard_atari_env_spec(env))
   hparams.add_hparam(
-      "environment_eval_spec",
-      standard_atari_env_eval_spec("PongNoFrameskip-v4"))
+      "environment_eval_spec", standard_atari_env_eval_spec(env))
   return hparams
 
 
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
index fdf276c39..b088eb678 100644
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ b/tensor2tensor/rl/envs/batch_env_factory.py
@@ -35,5 +35,4 @@ def batch_env_factory(environment_spec, num_agents):
     )
   else:
     cur_batch_env = py_func_batch_env.PyFuncBatchEnv(environment_spec.env)
-
-  return cur_batch_env
\ No newline at end of file
+  return cur_batch_env
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 6d20182bb..ea5dba2b0 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -179,6 +179,7 @@ def decode_png(encoded_png):
     )
     # TODO(koz4k): Move this to a different module.
     def initial_frame_chooser(batch_size):
+      """Frame chooser."""
       num_frames = environment_spec.video_num_input_frames
       deterministic_initial_frames = initial_frame_rollouts[0][:num_frames]
       if not environment_spec.simulation_random_starts:

From cb976068edc789323620c598544fbf91df37f6d3 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Wed, 17 Oct 2018 16:18:32 -0700
Subject: [PATCH 1033/2720] adjusting frames selection range.

PiperOrigin-RevId: 217608668
---
 tensor2tensor/rl/trainer_model_based.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index ea5dba2b0..9c6070100 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -192,7 +192,7 @@ def initial_frame_chooser(batch_size):
         def choose_initial_frames():
           try:
             rollout = random.choice(initial_frame_rollouts)
-            from_index = random.randrange(len(rollout) - num_frames)
+            from_index = random.randrange(len(rollout) - num_frames + 1)
             return rollout[from_index:(from_index + num_frames)]
           except ValueError:
             # Rollout too short; repeat.

From 65091468598b5710d9efe4cfade9d80007cb238f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 18 Oct 2018 10:08:56 -0700
Subject: [PATCH 1034/2720] Enable higher resolutions for next_frame_emily.

PiperOrigin-RevId: 217721451
---
 tensor2tensor/models/video/emily.py | 40 +++++++++++++++++++++++++----
 1 file changed, 35 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index 875a6ea50..102eae9bf 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -56,6 +56,29 @@ def encoder(self, inputs, nout):
     """
     vgg_layer = common_video.vgg_layer
     net01 = inputs
+
+    skips = []
+
+    # The original model only supports 64x64. We can support higher resolutions
+    # as long as they are square and the side-length is a power of two
+    # by inserting more downscaling layers. Corresponding upscaling can be found
+    # in the decoder, as well.
+    # (This procedure is ad-hoc, i.e., not from the SVP-FP paper)
+    _, res_y, res_x, _ = inputs.shape.as_list()
+    assert res_x == res_y, "Model only supports square inputs"
+    is_power_of_two = lambda x: ((x & (x - 1)) == 0) and x != 0
+    assert is_power_of_two(res_x), "Input resolution must be power of 2"
+    assert res_x >= 64, "Input resolution must be >= 64"
+    ds_idx = 0
+    while res_x > 64:
+      h = tfcl.repeat(net01, 2, vgg_layer, 64, scope="downscale%d" % ds_idx,
+                      is_training=self.is_training)
+      net01 = tfl.max_pooling2d(h, [2, 2], strides=(2, 2),
+                                name="downscale%d_pool" % ds_idx)
+      skips.append(h)
+      ds_idx += 1
+      res_x /= 2
+
     # h1
     net11 = tfcl.repeat(net01, 2, vgg_layer, 64,
                         scope="h1", is_training=self.is_training)
@@ -76,7 +99,7 @@ def encoder(self, inputs, nout):
     net51 = tfcl.repeat(net42, 1, vgg_layer, nout,
                         kernel_size=4, padding="VALID", activation=tf.tanh,
                         scope="h5", is_training=self.is_training)
-    skips = [net11, net21, net31, net41]
+    skips += [net11, net21, net31, net41]
     return net51, skips
 
   def decoder(self, inputs, skips, nout):
@@ -99,23 +122,30 @@ def decoder(self, inputs, skips, nout):
     net = tf.nn.leaky_relu(net)
     net = common_layers.upscale(net, 2)
     # d2
-    net = tf.concat([net, skips[3]], axis=3)
+    net = tf.concat([net, skips[-1]], axis=3)
     net = tfcl.repeat(net, 2, vgg_layer, 512, scope="d2a")
     net = tfcl.repeat(net, 1, vgg_layer, 256, scope="d2b")
     net = common_layers.upscale(net, 2)
     # d3
-    net = tf.concat([net, skips[2]], axis=3)
+    net = tf.concat([net, skips[-2]], axis=3)
     net = tfcl.repeat(net, 2, vgg_layer, 256, scope="d3a")
     net = tfcl.repeat(net, 1, vgg_layer, 128, scope="d3b")
     net = common_layers.upscale(net, 2)
     # d4
-    net = tf.concat([net, skips[1]], axis=3)
+    net = tf.concat([net, skips[-3]], axis=3)
     net = tfcl.repeat(net, 1, vgg_layer, 128, scope="d4a")
     net = tfcl.repeat(net, 1, vgg_layer, 64, scope="d4b")
     net = common_layers.upscale(net, 2)
     # d5
-    net = tf.concat([net, skips[0]], axis=3)
+    net = tf.concat([net, skips[-4]], axis=3)
     net = tfcl.repeat(net, 1, vgg_layer, 64, scope="d5")
+
+    # if there are still skip connections left, we have more downscaling to do
+    for i, s in enumerate(skips[-5::-1]):
+      net = common_layers.upscale(net, 2)
+      net = tf.concat([net, s], axis=3)
+      net = tfcl.repeat(net, 1, vgg_layer, 64, scope="upscale%d" % i)
+
     net = tfl.conv2d_transpose(net, nout, kernel_size=3, padding="SAME",
                                name="d6_deconv", activation=tf.sigmoid)
     return net

From 35a2ba019c58f8d967cec72dfa6b959df4afa9e6 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 18 Oct 2018 19:59:43 +0200
Subject: [PATCH 1035/2720] Simplify PPO epoch calculation (#1152)

---
 tensor2tensor/rl/trainer_model_based.py       | 100 ++++++++++--------
 .../rl/trainer_model_based_params.py          |   6 +-
 2 files changed, 61 insertions(+), 45 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 9c6070100..9bccd19b8 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -61,30 +61,32 @@ def temporary_flags(flag_settings):
     setattr(FLAGS, flag_name, flag_value)
 
 
-def _ppo_training_epochs(hparams, epoch, is_final_epoch, real_env_training):
-  """Helper for PPO restarts."""
+def real_ppo_epoch_increment(hparams):
   if hparams.gather_ppo_real_env_data:
     assert hparams.real_ppo_epochs_num is 0, (
-        "Should be put to 0 to enforce better readability")
-    real_training_ppo_epochs_num = int(math.ceil(
+        "Should be put to 0 to enforce better readability"
+    )
+    return int(math.ceil(
         hparams.num_real_env_frames /
-        (hparams.epochs*hparams.real_ppo_epoch_length)))
+        (hparams.epochs * hparams.real_ppo_epoch_length)
+    ))
   else:
-    real_training_ppo_epochs_num = hparams.real_ppo_epochs_num
+    return hparams.real_ppo_epochs_num
+
 
-  simulated_training_ppo_epochs_num = hparams.ppo_epochs_num
+def sim_ppo_epoch_increment(hparams, is_final_epoch):
+  increment = hparams.ppo_epochs_num
+  if is_final_epoch:
+    increment *= 2
+  return increment
 
-  if epoch == -1:
-    assert real_env_training, (
-        "Epoch -1 should only be used for PPO collection in real environment.")
-    return real_training_ppo_epochs_num
-  ppo_training_epochs = (epoch + 1) * (simulated_training_ppo_epochs_num
-                                       + real_training_ppo_epochs_num)
-  if is_final_epoch:  # Length of training in the final epoch is doubled.
-    ppo_training_epochs += simulated_training_ppo_epochs_num
-  if real_env_training:
-    ppo_training_epochs += real_training_ppo_epochs_num
-  return ppo_training_epochs
+
+def world_model_step_increment(hparams, is_initial_epoch):
+  if is_initial_epoch:
+    multiplier = hparams.initial_epoch_train_steps_multiplier
+  else:
+    multiplier = 1
+  return multiplier * hparams.model_train_steps
 
 
 def setup_directories(base_dir, subdirs):
@@ -140,8 +142,8 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
 
 
 def train_agent(real_env, environment_spec, agent_model_dir,
-                event_dir, world_model_dir, data_dir, hparams, epoch=0,
-                is_final_epoch=False):
+                event_dir, world_model_dir, data_dir, hparams, ppo_epochs_num,
+                epoch=0, is_final_epoch=False):
   """Train the PPO agent in the simulated environment."""
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_params_names = ["epochs_num", "epoch_length",
@@ -153,8 +155,9 @@ def train_agent(real_env, environment_spec, agent_model_dir,
     if ppo_param_name in hparams:
       ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
 
-  ppo_hparams.epochs_num = _ppo_training_epochs(hparams, epoch,
-                                                is_final_epoch, False)
+  ppo_epochs_num += sim_ppo_epoch_increment(hparams, is_final_epoch)
+  ppo_hparams.epochs_num = ppo_epochs_num
+
   ppo_hparams.save_models_every_epochs = 10
   ppo_hparams.world_model_dir = world_model_dir
   ppo_hparams.add_hparam("force_beginning_resets", True)
@@ -222,10 +225,12 @@ def choose_initial_frames():
       rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
                            name_scope="ppo_sim%d" % (epoch + 1))
 
+  return ppo_epochs_num
+
 
 def train_agent_real_env(
     env, agent_model_dir, event_dir, data_dir,
-    hparams, epoch=0, is_final_epoch=False):
+    hparams, ppo_epochs_num, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the real environment."""
   del data_dir
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
@@ -240,8 +245,8 @@ def train_agent_real_env(
     if ppo_param_name in hparams:
       ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
 
-  ppo_hparams.epochs_num = _ppo_training_epochs(hparams, epoch,
-                                                is_final_epoch, True)
+  ppo_epochs_num += real_ppo_epoch_increment(hparams)
+  ppo_hparams.epochs_num = ppo_epochs_num
   # We do not save model, as that resets frames that we need at restarts.
   # But we need to save at the last step, so we set it very high.
   ppo_hparams.save_models_every_epochs = 1000000
@@ -256,11 +261,16 @@ def train_agent_real_env(
   # Save unfinished rollouts to history.
   env.reset()
 
+  return ppo_epochs_num
+
 
-def train_world_model(env, data_dir, output_dir, hparams, epoch):
+def train_world_model(
+    env, data_dir, output_dir, hparams, world_model_steps_num, epoch
+):
   """Train the world model on problem_name."""
-  train_steps = hparams.model_train_steps * (
-      epoch + hparams.inital_epoch_train_steps_multiplier)
+  world_model_steps_num += world_model_step_increment(
+      hparams, is_initial_epoch=(epoch == 0)
+  )
   model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
   model_hparams.learning_rate = model_hparams.learning_rate_constant
   if epoch > 0:
@@ -272,11 +282,13 @@ def train_world_model(env, data_dir, output_dir, hparams, epoch):
       hparams=model_hparams,
       data_dir=data_dir,
       output_dir=output_dir,
-      train_steps=train_steps,
+      train_steps=world_model_steps_num,
       eval_steps=100,
       local_eval_frequency=2000
   )
 
+  return world_model_steps_num
+
 
 def setup_env(hparams):
   """Setup."""
@@ -338,10 +350,10 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   tf.logging.info("Initial training of PPO in real environment.")
   ppo_event_dir = os.path.join(directories["world_model"],
                                "ppo_summaries/initial")
-  train_agent_real_env(
-      env, ppo_model_dir,
-      ppo_event_dir, data_dir,
-      hparams, epoch=epoch, is_final_epoch=False)
+  ppo_epochs_num = train_agent_real_env(
+      env, ppo_model_dir, ppo_event_dir, data_dir, hparams, ppo_epochs_num=0,
+      epoch=epoch, is_final_epoch=False
+  )
   mean_unclipped_reward = eval_reward(env, clipped=False)
   tf.logging.info("Mean reward (initial): {}".format(mean_unclipped_reward))
 
@@ -362,6 +374,8 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       video_num_input_frames=4, video_num_target_frames=1
   )
 
+  world_model_steps_num = 0
+
   for epoch in range(hparams.epochs):
     env.generate_data(data_dir)
 
@@ -370,8 +384,10 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
     # Train world model
     log("Training world model")
-    train_world_model(env, data_dir,
-                      directories["world_model"], hparams, epoch)
+    world_model_steps_num = train_world_model(
+        env, data_dir, directories["world_model"], hparams,
+        world_model_steps_num, epoch
+    )
 
     # Train PPO
     log("Training PPO in simulated environment.")
@@ -381,20 +397,20 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     if not hparams.ppo_continue_training:
       ppo_model_dir = ppo_event_dir
 
-    train_agent(
+    ppo_epochs_num = train_agent(
         env, sim_env_spec, ppo_model_dir, ppo_event_dir,
-        directories["world_model"], data_dir, hparams, epoch=epoch,
-        is_final_epoch=is_final_epoch
+        directories["world_model"], data_dir, hparams, ppo_epochs_num,
+        epoch=epoch, is_final_epoch=is_final_epoch
     )
 
     env.start_new_epoch(epoch, data_dir)
 
     # Train PPO on real env (short)
     log("Training PPO in real environment.")
-    train_agent_real_env(
-        env, ppo_model_dir,
-        ppo_event_dir, data_dir,
-        hparams, epoch=epoch, is_final_epoch=is_final_epoch)
+    ppo_epochs_num = train_agent_real_env(
+        env, ppo_model_dir, ppo_event_dir, data_dir, hparams, ppo_epochs_num,
+        epoch=epoch, is_final_epoch=is_final_epoch
+    )
 
     if hparams.stop_loop_early:
       return 0.0
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index a7b34d7b0..66160d53a 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -59,7 +59,7 @@ def rlmb_base():
       autoencoder_train_steps_initial_multiplier=10,
       autoencoder_hparams_set="autoencoder_discrete_pong",
       model_train_steps=15000,
-      inital_epoch_train_steps_multiplier=3,
+      initial_epoch_train_steps_multiplier=3,
       simulated_env_generator_num_steps=2000,
       simulation_random_starts=True,  # Use random starts in PPO.
       # Flip the first random frame in PPO batch for the true beginning.
@@ -183,7 +183,7 @@ def rlmb_quick_sm():
 def rlmb_base_stochastic():
   """Base setting with a stochastic next-frame model."""
   hparams = rlmb_base()
-  hparams.inital_epoch_train_steps_multiplier = 5
+  hparams.initial_epoch_train_steps_multiplier = 5
   hparams.generative_model = "next_frame_basic_stochastic"
   hparams.generative_model_params = "next_frame_basic_stochastic"
   return hparams
@@ -272,7 +272,7 @@ def rlmb_base_sv2p_flippy30():
   hparams.ppo_epochs_num = 1000
   hparams.model_train_steps = 15000
   hparams.learning_rate_bump = 1.0
-  hparams.inital_epoch_train_steps_multiplier = 5
+  hparams.initial_epoch_train_steps_multiplier = 5
   hparams.generative_model = "next_frame_sv2p"
   hparams.generative_model_params = "next_frame_sv2p_atari"
   return hparams

From 8e3d18cff1f966f3f1901b3cc5b7178dcbff151d Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 18 Oct 2018 11:00:22 -0700
Subject: [PATCH 1036/2720] internal merge of PR #1152

PiperOrigin-RevId: 217731808
---
 tensor2tensor/rl/trainer_model_based.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 9bccd19b8..cb2548d0c 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -62,6 +62,7 @@ def temporary_flags(flag_settings):
 
 
 def real_ppo_epoch_increment(hparams):
+  """PPO increment."""
   if hparams.gather_ppo_real_env_data:
     assert hparams.real_ppo_epochs_num is 0, (
         "Should be put to 0 to enforce better readability"
@@ -232,7 +233,7 @@ def train_agent_real_env(
     env, agent_model_dir, event_dir, data_dir,
     hparams, ppo_epochs_num, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the real environment."""
-  del data_dir
+  del is_final_epoch, data_dir
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_params_names = ["epochs_num", "epoch_length",
                       "learning_rate", "num_agents", "eval_every_epochs",

From 452940a4a8c8d0d1987b24b6801146f8d4e5b7f1 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 18 Oct 2018 20:54:49 +0200
Subject: [PATCH 1037/2720] RL pipeline cleanup (#1153)

* Remove batch_env_factory

* Remove trainer_model_based_new_test

* Remove rl.envs.utils
---
 tensor2tensor/data_generators/gym_problems.py |   1 -
 tensor2tensor/models/research/rl.py           |  15 ++
 tensor2tensor/rl/collect.py                   |  10 +-
 tensor2tensor/rl/envs/batch_env_factory.py    |  38 ---
 tensor2tensor/rl/envs/in_graph_batch_env.py   |  18 +-
 tensor2tensor/rl/envs/utils.py                | 241 ------------------
 .../rl/model_rl_experiment_player.py          |   8 +-
 tensor2tensor/rl/ppo.py                       |   3 +-
 .../rl/trainer_model_based_new_test.py        |  38 ---
 9 files changed, 40 insertions(+), 332 deletions(-)
 delete mode 100644 tensor2tensor/rl/envs/batch_env_factory.py
 delete mode 100644 tensor2tensor/rl/envs/utils.py
 delete mode 100644 tensor2tensor/rl/trainer_model_based_new_test.py

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 1069c929d..6206d2f42 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -32,7 +32,6 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import collect
 from tensor2tensor.rl.envs import tf_atari_wrappers
-from tensor2tensor.rl.envs.utils import InitialFrameChooser
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 974052397..6ace4cd75 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -196,6 +196,21 @@ def standard_atari_ae_env_spec(env, ae_hparams_set):
                                      simulated_env=False)
 
 
+def get_policy(observations, hparams):
+  """Get a policy network.
+
+  Args:
+    observations: Tensor with observations
+    hparams: parameters
+
+  Returns:
+    Tensor with policy and value function output
+  """
+  policy_network_lambda = hparams.policy_network
+  action_space = hparams.environment_spec.action_space
+  return policy_network_lambda(action_space, hparams, observations)
+
+
 @registry.register_hparams
 def ppo_pong_ae_base():
   """Pong autoencoder base parameters."""
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index e925cd9f4..4c11cd5b2 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -21,9 +21,10 @@
 
 import copy
 
-from tensor2tensor.rl.envs.batch_env_factory import batch_env_factory
+from tensor2tensor.models.research.rl import get_policy
+from tensor2tensor.rl.envs.py_func_batch_env import PyFuncBatchEnv
+from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
 from tensor2tensor.rl.envs.tf_atari_wrappers import WrapperBase
-from tensor2tensor.rl.envs.utils import get_policy
 
 import tensorflow as tf
 
@@ -112,7 +113,10 @@ def define_collect(hparams, scope, eval_phase,
       environment_spec = getattr(hparams, "environment_eval_spec",
                                  environment_spec)
       num_agents = getattr(hparams, "num_eval_agents", num_agents)
-    batch_env = batch_env_factory(environment_spec, num_agents)
+    if environment_spec.simulated_env:
+      batch_env = SimulatedBatchEnv(environment_spec, num_agents)
+    else:
+      batch_env = PyFuncBatchEnv(environment_spec.env)
 
     to_initialize.append(batch_env)
     environment_wrappers = environment_spec.wrappers
diff --git a/tensor2tensor/rl/envs/batch_env_factory.py b/tensor2tensor/rl/envs/batch_env_factory.py
deleted file mode 100644
index b088eb678..000000000
--- a/tensor2tensor/rl/envs/batch_env_factory.py
+++ /dev/null
@@ -1,38 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Utilities for creating batched environments."""
-
-# The code was based on Danijar Hafner's code from tf.agents:
-# https://github.com/tensorflow/agents/blob/master/agents/tools/wrappers.py
-# https://github.com/tensorflow/agents/blob/master/agents/scripts/utility.py
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.rl.envs import py_func_batch_env
-from tensor2tensor.rl.envs import simulated_batch_env
-
-
-def batch_env_factory(environment_spec, num_agents):
-  """Factory of batch envs."""
-  if environment_spec.simulated_env:
-    cur_batch_env = simulated_batch_env.SimulatedBatchEnv(
-        environment_spec, num_agents
-    )
-  else:
-    cur_batch_env = py_func_batch_env.PyFuncBatchEnv(environment_spec.env)
-  return cur_batch_env
diff --git a/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index 58492b4fd..87b2a0caf 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -22,7 +22,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.rl.envs import utils
+import gym
 
 import tensorflow as tf
 
@@ -84,21 +84,29 @@ def reset(self, indices=None):
         lambda: self._reset_non_empty(indices),
         lambda: tf.cast(0, self.observ_dtype))
 
+  @staticmethod
+  def _get_tf_dtype(space):
+    if isinstance(space, gym.spaces.Discrete):
+      return tf.int32
+    if isinstance(space, gym.spaces.Box):
+      return tf.as_dtype(space.dtype)
+    raise NotImplementedError()
+
   @property
   def observ_dtype(self):
-    return utils.parse_dtype(self.observ_space)
+    return self._get_tf_dtype(self.observ_space)
 
   @property
   def observ_shape(self):
-    return utils.parse_shape(self.observ_space)
+    return self.observ_space.shape
 
   @property
   def action_dtype(self):
-    return utils.parse_dtype(self.action_space)
+    return self._get_tf_dtype(self.action_space)
 
   @property
   def action_shape(self):
-    return utils.parse_shape(self.action_space)
+    return self.action_space.shape
 
   @property
   def observ(self):
diff --git a/tensor2tensor/rl/envs/utils.py b/tensor2tensor/rl/envs/utils.py
deleted file mode 100644
index 503dd58f4..000000000
--- a/tensor2tensor/rl/envs/utils.py
+++ /dev/null
@@ -1,241 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Utilities for using batched environments."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gym
-import six
-import tensorflow as tf
-
-
-class EvalVideoWrapper(gym.Wrapper):
-  """Wrapper for recording videos during eval phase.
-
-  This wrapper is designed to record videos via gym.wrappers.Monitor and
-  simplifying its usage in t2t collect phase.
-  It alleviate the limitation of Monitor, which doesn't allow reset on an
-  active environment.
-
-  EvalVideoWrapper assumes that only every second trajectory (after every
-  second reset) will be used by the caller:
-  - on the "active" runs it behaves as gym.wrappers.Monitor,
-  - on the "inactive" runs it doesn't call underlying environment and only
-    returns last seen observation.
-  Videos are only generated during the active runs.
-  """
-
-  def __init__(self, env):
-    super(EvalVideoWrapper, self).__init__(env)
-    self._reset_counter = 0
-    self._active = False
-    self._last_returned = None
-
-  def _step(self, action):
-    if self._active:
-      self._last_returned = self.env.step(action)
-    if self._last_returned is None:
-      raise Exception("Environment stepped before proper reset.")
-    return self._last_returned
-
-  def _reset(self, **kwargs):
-    self._reset_counter += 1
-    if self._reset_counter % 2 == 1:
-      self._active = True
-      return self.env.reset(**kwargs)
-
-    self._active = False
-    self._last_returned = (self._last_returned[0],
-                           self._last_returned[1],
-                           False,  # done = False
-                           self._last_returned[3])
-    return self._last_returned[0]
-
-
-def get_policy(observations, hparams):
-  """Get a policy network.
-
-  Args:
-    observations: Tensor with observations
-    hparams: parameters
-
-  Returns:
-    Tensor with policy and value function output
-  """
-  policy_network_lambda = hparams.policy_network
-  action_space = hparams.environment_spec.action_space
-  return policy_network_lambda(action_space, hparams, observations)
-
-
-def parse_shape(space):
-  """Get a tensor shape from a OpenAI Gym space.
-
-  Args:
-    space: Gym space.
-
-  Returns:
-    Shape tuple.
-  """
-  if isinstance(space, gym.spaces.Discrete):
-    return ()
-  if isinstance(space, gym.spaces.Box):
-    return space.shape
-  raise NotImplementedError()
-
-
-def parse_dtype(space):
-  """Get a tensor dtype from a OpenAI Gym space.
-
-  Args:
-    space: Gym space.
-
-  Returns:
-    TensorFlow data type.
-  """
-  if isinstance(space, gym.spaces.Discrete):
-    return tf.int32
-  if isinstance(space, gym.spaces.Box):
-    return tf.as_dtype(space.dtype)
-  raise NotImplementedError()
-
-
-class InitialFrameChooser(object):
-  """Class for choosing the initial frame for simulation from the dataset.
-
-  Can also store a sequence of later frames, which is used for comparison in
-  world model evaluation.
-
-  Attributes:
-    batch_size (int): Batch size, should be set before calling choose().
-    trajectory (dict): Dict of Variables storing a sequence of frames after the
-        chosen one.
-  """
-
-  def __init__(self, environment_spec, mode, trajectory_length=1):
-    self._initial_frames_problem = environment_spec.initial_frames_problem
-    self._simulation_random_starts = environment_spec.simulation_random_starts
-    self._flip_first_random_for_beginning = \
-        environment_spec.simulation_flip_first_random_for_beginning
-    self._num_initial_frames = environment_spec.video_num_input_frames
-
-    def dataset_kwargs_lambda():
-      video_num_input_frames = environment_spec.video_num_input_frames
-      video_num_input_frames += trajectory_length - 1
-      dataset_hparams = tf.contrib.training.HParams(
-          video_num_input_frames=video_num_input_frames,
-          video_num_target_frames=environment_spec.video_num_target_frames,
-          environment_spec=environment_spec
-      )
-      return {
-          "mode": mode,
-          "data_dir": tf.flags.FLAGS.data_dir,
-          "hparams": dataset_hparams,
-          "only_last": True
-      }
-
-    self._dataset_kwargs_lambda = dataset_kwargs_lambda
-    self._start_frames = None
-
-  @property
-  def batch_size(self):
-    return self._batch_size
-
-  @batch_size.setter
-  def batch_size(self, batch_size):
-    self._batch_size = batch_size
-    self._iterator = \
-        self._create_initial_frame_dataset().make_initializable_iterator()
-
-    def fix_and_shorten(shape):
-      shape = shape.as_list()
-      shape[0] = batch_size
-      shape[1] -= self._num_initial_frames - 1
-      return shape
-
-    shapes = self._extract_input(self._iterator.output_shapes)
-    types = self._extract_input(self._iterator.output_types)
-    self.trajectory = {
-        key: tf.Variable(
-            tf.zeros(fix_and_shorten(shape), types[key]),
-            trainable=False
-        )
-        for (key, shape) in six.iteritems(shapes)
-    }
-
-  def initialize(self, sess):
-    sess.run(self._iterator.initializer)
-
-  def choose(self):
-    """Returns a dict of tensors of the chosen initial frame.
-
-    Also assigns the first trajectory_length frames after the initial frames to
-    self.trajectory.
-    """
-    if self._flip_first_random_for_beginning and self._start_frames is None:
-      ordered_dataset = self._create_dataset(shuffle_files=False)
-      # Later flip the first random frame in PPO batch for the true beginning.
-      self._start_frames = self._extract_input(
-          ordered_dataset.make_one_shot_iterator().get_next()
-      )
-
-    all_frames = self._extract_input(self._iterator.get_next())
-    if self._start_frames is not None:
-      all_frames = {
-          key: tf.concat([
-              tf.expand_dims(self._start_frames[key], axis=0),
-              value[1:, ...]
-          ], axis=0)
-          for (key, value) in six.iteritems(all_frames)
-      }
-    scatter_ops = [
-        tf.scatter_update(
-            self.trajectory[key], tf.range(tf.shape(value)[0]),
-            value[:, (self._num_initial_frames - 1):, ...]
-        )
-        for (key, value) in six.iteritems(all_frames)
-    ]
-
-    with tf.control_dependencies(scatter_ops):
-      return {
-          key: value[:, :self._num_initial_frames, ...]
-          for (key, value) in six.iteritems(all_frames)
-      }
-
-  def _create_dataset(self, **extra_dataset_kwargs):
-    dataset_kwargs = self._dataset_kwargs_lambda()
-    dataset_kwargs.update(extra_dataset_kwargs)
-    return self._initial_frames_problem.dataset(**dataset_kwargs)
-
-  def _create_initial_frame_dataset(self):
-    """Returns the dataset that consecutive initial frames will be taken from.
-    """
-    dataset = self._create_dataset(
-        shuffle_files=self._simulation_random_starts
-    )
-    if self._simulation_random_starts:
-      dataset = dataset.shuffle(buffer_size=1000)
-    return dataset.repeat().batch(self._batch_size)
-
-  def _extract_input(self, frame):
-    input_frame = {"inputs": frame["inputs"]}
-    input_frame.update({
-        key[len("input_"):]: value
-        for (key, value) in six.iteritems(frame)
-        if key.startswith("input_")
-    })
-    return input_frame
diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index 259b2dec9..a5da41afc 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -34,8 +34,8 @@
 from PIL import ImageFont
 
 from tensor2tensor.data_generators import gym_problems_specs
-from tensor2tensor.rl.envs.batch_env_factory import batch_env_factory
-from tensor2tensor.rl.envs.utils import get_policy
+from tensor2tensor.models.research.rl import get_policy
+from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
 from tensor2tensor.rl.trainer_model_based import FLAGS
 from tensor2tensor.rl.trainer_model_based import setup_directories
 from tensor2tensor.rl.trainer_model_based import temporary_flags
@@ -100,9 +100,7 @@ def __init__(self, hparams, sess=None):
 
   def _prepare_networks(self, hparams, sess):
     self.action = tf.placeholder(shape=(1,), dtype=tf.int32)
-    batch_env = batch_env_factory(
-        hparams.environment_spec, hparams.num_agents,
-        initial_frame_chooser=hparams.initial_frame_chooser)
+    batch_env = SimulatedBatchEnv(hparams.environment_spec, hparams.num_agents)
     self.reward, self.done = batch_env.simulate(self.action)
     self.observation = batch_env.observ
     self.reset_op = batch_env.reset(tf.constant([0], dtype=tf.int32))
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index 67439bac8..69ab1ba0d 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -20,7 +20,8 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-from tensor2tensor.rl.envs.utils import get_policy
+
+from tensor2tensor.models.research.rl import get_policy
 
 import tensorflow as tf
 
diff --git a/tensor2tensor/rl/trainer_model_based_new_test.py b/tensor2tensor/rl/trainer_model_based_new_test.py
deleted file mode 100644
index bf23d592a..000000000
--- a/tensor2tensor/rl/trainer_model_based_new_test.py
+++ /dev/null
@@ -1,38 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Tiny run of trainer_model_based_new. Smoke test."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.rl import trainer_model_based_new
-
-import tensorflow as tf
-
-FLAGS = tf.flags.FLAGS
-
-
-class ModelRLExperimentNewTest(tf.test.TestCase):
-
-  def test_basic(self):
-    FLAGS.output_dir = tf.test.get_temp_dir()
-    FLAGS.loop_hparams_set = "rlmb_tiny"
-    FLAGS.schedule = "train"  # skip evaluation for world model training
-    trainer_model_based_new.main(None)
-
-
-if __name__ == "__main__":
-  tf.test.main()

From e15d42cae6183f69eeb51261d13e59fee3bdba04 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 18 Oct 2018 11:55:16 -0700
Subject: [PATCH 1038/2720] internal merge of PR #1153

PiperOrigin-RevId: 217742387
---
 tensor2tensor/data_generators/gym_problems.py | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
index 6206d2f42..ae25187bc 100644
--- a/tensor2tensor/data_generators/gym_problems.py
+++ b/tensor2tensor/data_generators/gym_problems.py
@@ -552,10 +552,6 @@ def _setup(self, data_dir, extra_collect_hparams=None,
     if extra_collect_hparams is None:
       extra_collect_hparams = {}
 
-    if self._initial_frame_chooser is None:
-      self._initial_frame_chooser = InitialFrameChooser(
-          self.environment_spec, mode=tf.estimator.ModeKeys.EVAL
-      )
     extra_collect_hparams["initial_frame_chooser"] = self._initial_frame_chooser
 
     super(GymSimulatedDiscreteProblem, self)._setup(
@@ -636,11 +632,6 @@ def _setup(self, data_dir):
       # Decrease the trajectory length for tiny experiments, otherwise we don't
       # have enough data to run the evaluation.
       trajectory_length = 2
-    self._initial_frame_chooser = InitialFrameChooser(
-        self.environment_spec, mode=tf.estimator.ModeKeys.EVAL,
-        trajectory_length=trajectory_length
-    )
-
     frame_index = tf.Variable(0, trainable=False)
 
     def fixed_action_policy_fun(action_space, unused_config, observations):

From 7f9068edec07a20af3a7ed8a8cb876ee33831405 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 18 Oct 2018 12:55:09 -0700
Subject: [PATCH 1039/2720] Adding auto-regressive predictor to sv2p-disc.

PiperOrigin-RevId: 217752333
---
 tensor2tensor/layers/common_video.py      |   6 +-
 tensor2tensor/models/video/sv2p.py        | 114 ++++++++++++++++------
 tensor2tensor/models/video/sv2p_params.py |   6 ++
 3 files changed, 95 insertions(+), 31 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index eac08bada..e551e49b6 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -41,7 +41,7 @@ def encode_to_shape(inputs, shape, scope):
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
     w, h = shape[1], shape[2]
     x = inputs
-    x = tf.contrib.layers.flatten(x)
+    x = tfl.flatten(x)
     x = tfl.dense(x, w * h, activation=None, name="enc_dense")
     x = tf.reshape(x, (-1, w, h, 1))
     return x
@@ -51,7 +51,7 @@ def decode_to_shape(inputs, shape, scope):
   """Encode the given tensor to given image shape."""
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
     x = inputs
-    x = tf.contrib.layers.flatten(x)
+    x = tfl.flatten(x)
     x = tfl.dense(x, shape[2], activation=None, name="dec_dense")
     x = tf.expand_dims(x, axis=1)
     return x
@@ -484,7 +484,7 @@ def tinyify(array, tiny_mode, small_mode):
   if tiny_mode:
     return [1 for _ in array]
   if small_mode:
-    return [x // 4 for x in array]
+    return [max(x // 4, 1) for x in array]
   return array
 
 
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 08f1fd10d..b5c01c4ac 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -27,6 +27,7 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
+from tensor2tensor.layers import discretization
 
 from tensor2tensor.models.video import base
 from tensor2tensor.models.video import base_vae
@@ -131,7 +132,10 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
           enc2, input_reward, "reward_enc")
     if latent is not None and not concat_latent:
       with tf.control_dependencies([latent]):
-        enc2 = tf.concat([enc2, latent], axis=3)
+        # This is the original SV2P implementation
+        # But we will tile and concat to support various latent sizes.
+        # enc2 = tf.concat([enc2, latent], axis=3)
+        enc2 = tile_and_concat(enc2, latent, concat_latent=concat_latent)
 
     enc3 = tfl.conv2d(enc2, hidden4.get_shape()[3], [1, 1], strides=(1, 1),
                       padding="SAME", activation=tf.nn.relu, name="conv4")
@@ -345,8 +349,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
     extra_loss = 0.0
     if internal_states is None:
       internal_states = [None] * (5 if self.hparams.small_mode else 7)
-      if latent_mean is not None:
-        extra_loss = self.get_extra_loss([latent_mean], [latent_std])
+      extra_loss = self.get_extra_loss([latent_mean], [latent_std])
 
     pred_image, internal_states = self.construct_predictive_tower(
         frames, None, actions, internal_states, latent)
@@ -365,35 +368,90 @@ class NextFrameSv2pDiscrete(NextFrameSv2p):
 
   def video_features(
       self, all_frames, all_actions, all_rewards, all_raw_frames):
-    """Video wide latent."""
-    del all_actions, all_rewards, all_raw_frames
+    """No video wide latent."""
+    del all_frames, all_actions, all_rewards, all_raw_frames
+    return None
+
+  def basic_conv_net(self, images, conv_size, scope):
+    """Simple multi conv ln relu."""
+    conv_size = self.tinyify(conv_size)
+    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
+      x = images
+      for i, c in enumerate(conv_size):
+        if i > 0:
+          x = tf.nn.relu(x)
+        x = common_layers.make_even_size(x)
+        x = tfl.conv2d(x, c, [3, 3], strides=(2, 2),
+                       activation=None, padding="SAME", name="conv%d" % i)
+        x = tfcl.layer_norm(x)
+    return x
 
+  def learned_discrete_tower(self, input_image, target_image):
     hparams = self.hparams
-    frames = tf.stack(all_frames, axis=1)
-    mean, std = self.construct_latent_tower(frames, time_axis=1)
-    tower_output = tf.concat([mean, std], axis=-1)
-    tower_output_shape = common_layers.shape_list(tower_output)
+
+    # Encode the input frames into a prior encoding.
+    conv_size = [64, 32, 32, 1]
+    prior_enc = self.basic_conv_net(input_image, conv_size, "prior_enc")
+    tower_output_shape = common_layers.shape_list(prior_enc)
     batch_size = tower_output_shape[0]
+    prior_enc = tfl.flatten(prior_enc)
 
-    if not self.is_training:
-      rand = tf.random_uniform([batch_size, hparams.bottleneck_bits])
-      d = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
-    else:
-      x = tfl.flatten(tower_output)
-      x = tfl.dense(x, hparams.bottleneck_bits, name="bits_enc")
-      x_shape = common_layers.shape_list(x)
-      x += tf.truncated_normal(x_shape, mean=0.0, stddev=0.2)
-      x = tf.tanh(x)
-      noise = tf.random_uniform(x_shape)
-      noise = 2.0 * tf.to_float(tf.less(hparams.bottleneck_noise, noise)) - 1.0
-      x *= noise
-      d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
-      p = common_layers.inverse_lin_decay(hparams.discrete_warmup_steps)
-      d = tf.where(tf.less(tf.random_uniform([batch_size]), p), d, x)
-
-    decoded_bits = common_video.encode_to_shape(
-        d, tower_output_shape, "bits_dec")
-    return [decoded_bits, None, None]
+    def decode_bits(b):
+      return common_video.encode_to_shape(b, tower_output_shape, "bits_dec")
+
+    if self.is_predicting:
+      if hparams.full_latent_tower:
+        rand = tf.random_uniform([batch_size, hparams.bottleneck_bits])
+        bits = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
+      else:
+        # Generate bit using the learned prior at inference time.
+        bits, _ = discretization.predict_bits_with_lstm(
+            prior_enc,
+            hparams.latent_predictor_state_size,
+            hparams.bottleneck_bits,
+            temperature=hparams.latent_predictor_temperature)
+      return decode_bits(bits), 0.0
+
+    # Encode the input and target frames into posterior.
+    x = tf.concat([input_image, target_image], axis=-1)
+    x = self.basic_conv_net(x, conv_size, "posterior_enc")
+    x = tfl.flatten(x)
+    bits, bits_clean = discretization.tanh_discrete_bottleneck(
+        x, hparams.bottleneck_bits,
+        hparams.bottleneck_noise,
+        hparams.discretize_warmup_steps,
+        hparams.mode)
+
+    pred_loss = 0.0
+    if not hparams.full_latent_tower:
+      # Learn the prior by matching the posterior.
+      _, pred_loss = discretization.predict_bits_with_lstm(
+          prior_enc,
+          hparams.latent_predictor_state_size,
+          hparams.bottleneck_bits,
+          target_bits=bits_clean)
+
+    return decode_bits(bits), pred_loss
+
+  def next_frame(self, frames, actions, rewards, target_frame,
+                 internal_states, video_features):
+    del video_features
+    frames, actions, rewards = frames[0], actions[0], rewards[0]
+
+    if internal_states is None:
+      internal_states = [None] * (5 if self.hparams.small_mode else 7)
+
+    latent, extra_loss = self.learned_discrete_tower(frames, target_frame)
+
+    pred_image, internal_states = self.construct_predictive_tower(
+        frames, None, actions, internal_states, latent)
+
+    if not self.has_rewards:
+      return pred_image, None, extra_loss, internal_states
+
+    pred_reward = self.reward_prediction(
+        pred_image, actions, rewards, latent)
+    return pred_image, pred_reward, extra_loss, internal_states
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 23d3956ef..68f9380b1 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -54,9 +54,15 @@ def next_frame_sv2p():
 def next_frame_sv2p_discrete():
   """SV2P discrete model hparams."""
   hparams = next_frame_sv2p()
+  hparams.action_injection = "multiplicative"
+  hparams.small_mode = True
   hparams.add_hparam("bottleneck_bits", 128)
   hparams.add_hparam("bottleneck_noise", 0.02)
   hparams.add_hparam("discrete_warmup_steps", 40000)
+  hparams.add_hparam("full_latent_tower", False)
+  hparams.add_hparam("latent_predictor_state_size", 128)
+  hparams.add_hparam("latent_predictor_temperature", 0.5)
+  hparams.add_hparam("discretize_warmup_steps", 40000)
   return hparams
 
 

From 4d9519286ddd8594ac9782bcccd4185924612b56 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 18 Oct 2018 13:27:45 -0700
Subject: [PATCH 1040/2720] Use the new gym_env instead of gym_problems.

PiperOrigin-RevId: 217757881
---
 tensor2tensor/data_generators/gym_env.py      | 133 ++-
 tensor2tensor/data_generators/gym_problems.py | 795 ------------------
 .../data_generators/gym_problems_specs.py     | 339 --------
 .../data_generators/gym_problems_test.py      |  44 -
 tensor2tensor/rl/datagen_with_agent.py        |   4 +-
 .../rl/model_rl_experiment_player.py          |   7 +-
 .../rl/trainer_model_based_params.py          |  22 +-
 tensor2tensor/utils/registry.py               |   2 +-
 8 files changed, 146 insertions(+), 1200 deletions(-)
 delete mode 100644 tensor2tensor/data_generators/gym_problems.py
 delete mode 100644 tensor2tensor/data_generators/gym_problems_specs.py
 delete mode 100644 tensor2tensor/data_generators/gym_problems_test.py

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 1fda3d0a5..439d7158b 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -30,6 +30,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.utils import metrics
+from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -522,18 +523,24 @@ class T2TGymEnv(T2TEnv):
   arguments and register this subclass.
   """
 
-  def __init__(self, base_env_name, batch_size=None, grayscale=False,
-               resize_height_factor=1, resize_width_factor=1,
+  def __init__(self, base_env_name=None, batch_size=None, grayscale=False,
+               resize_height_factor=2, resize_width_factor=2,
                base_env_timesteps_limit=-1, envs=None, **kwargs):
     if batch_size is None:
-      batch_size = len(envs)
+      if envs is None:
+        batch_size = 1
+      else:
+        batch_size = len(envs)
+    if base_env_name is None:
+      base_env_name = self.base_env_name
+    self._base_env_name = base_env_name
     super(T2TGymEnv, self).__init__(batch_size, **kwargs)
     self.grayscale = grayscale
     self.resize_height_factor = resize_height_factor
     self.resize_width_factor = resize_width_factor
     if not self.name:
       # Set problem name if not registered.
-      self.name = "t2t_gym_env_{}".format(base_env_name)
+      self.name = "Gym%s" % base_env_name
 
     if envs is None:
       self._envs = [make_gym_env(base_env_name, base_env_timesteps_limit)
@@ -566,6 +573,10 @@ def __init__(self, base_env_name, batch_size=None, grayscale=False,
         resized = tf.image.rgb_to_grayscale(resized)
       self._resized_img_batch_t = _Noncopyable(resized)
 
+  @property
+  def base_env_name(self):
+    return self._base_env_name
+
   @property
   def num_channels(self):
     return self.observation_space.shape[2]
@@ -600,3 +611,117 @@ def _reset(self, indices):
   def close(self):
     for env in self._envs:
       env.close()
+
+# Atari registration.
+
+# Game list from our list of ROMs
+# Removed because XDeterministic-v4 did not exist:
+# * adventure
+# * defender
+# * kaboom
+ATARI_GAMES = [
+    "air_raid", "alien", "amidar", "assault", "asterix", "asteroids",
+    "atlantis", "bank_heist", "battle_zone", "beam_rider", "berzerk", "bowling",
+    "boxing", "breakout", "carnival", "centipede", "chopper_command",
+    "crazy_climber", "demon_attack", "double_dunk", "elevator_action", "enduro",
+    "fishing_derby", "freeway", "frostbite", "gopher", "gravitar", "hero",
+    "ice_hockey", "jamesbond", "journey_escape", "kangaroo", "krull",
+    "kung_fu_master", "montezuma_revenge", "ms_pacman", "name_this_game",
+    "phoenix", "pitfall", "pong", "pooyan", "private_eye", "qbert", "riverraid",
+    "road_runner", "robotank", "seaquest", "skiing", "solaris",
+    "space_invaders", "star_gunner", "tennis", "time_pilot", "tutankham",
+    "up_n_down", "venture", "video_pinball", "wizard_of_wor", "yars_revenge",
+    "zaxxon"
+]
+
+# List from paper:
+# https://arxiv.org/pdf/1805.11593.pdf
+# plus frostbite.
+ATARI_GAMES_WITH_HUMAN_SCORE = [
+    "alien", "amidar", "assault", "asterix", "asteroids",
+    "atlantis", "bank_heist", "battle_zone", "beam_rider", "bowling",
+    "boxing", "breakout", "chopper_command",
+    "crazy_climber", "demon_attack", "double_dunk", "enduro",
+    "fishing_derby", "freeway", "frostbite", "gopher", "gravitar", "hero",
+    "ice_hockey", "jamesbond", "kangaroo", "krull",
+    "kung_fu_master", "montezuma_revenge", "ms_pacman", "name_this_game",
+    "pitfall", "pong", "private_eye", "qbert", "riverraid",
+    "road_runner", "seaquest", "solaris",
+    "up_n_down", "video_pinball", "yars_revenge",
+]
+
+ATARI_WHITELIST_GAMES = [
+    "amidar",
+    "bank_heist",
+    "berzerk",
+    "boxing",
+    "crazy_climber",
+    "freeway",
+    "frostbite",
+    "gopher",
+    "kung_fu_master",
+    "ms_pacman",
+    "pong",
+    "qbert",
+    "seaquest",
+]
+
+
+# Games on which model-free does better than model-based at this point.
+ATARI_CURIOUS_GAMES = [
+    "bank_heist",
+    "boxing",
+    "enduro",
+    "kangaroo",
+    "road_runner",
+    "up_n_down",
+]
+
+
+# Games on which based should work.
+ATARI_DEBUG_GAMES = [
+    "crazy_climber",
+    "freeway",
+    "pong",
+]
+
+
+# Different ATARI game modes in OpenAI Gym. Full list here:
+# https://github.com/openai/gym/blob/master/gym/envs/__init__.py
+ATARI_GAME_MODES = [
+    "Deterministic-v0",  # 0.25 repeat action probability, 4 frame skip.
+    "Deterministic-v4",  # 0.00 repeat action probability, 4 frame skip.
+    "NoFrameskip-v0",    # 0.25 repeat action probability, 1 frame skip.
+    "NoFrameskip-v4",    # 0.00 repeat action probability, 1 frame skip.
+    "-v0",               # 0.25 repeat action probability, (2 to 5) frame skip.
+    "-v4"                # 0.00 repeat action probability, (2 to 5) frame skip.
+]
+
+
+def register_game(game_name, game_mode="Deterministic-v4"):
+  """Create and register problems for the game.
+
+  Args:
+    game_name: str, one of the games in ATARI_GAMES, e.g. "bank_heist".
+    game_mode: the frame skip and sticky keys config.
+
+  Raises:
+    ValueError: if game_name or game_mode are wrong.
+  """
+  if game_name not in ATARI_GAMES:
+    raise ValueError("Game %s not in ATARI_GAMES" % game_name)
+  if game_mode not in ATARI_GAME_MODES:
+    raise ValueError("Unknown ATARI game mode: %s." % game_mode)
+  camel_game_name = "".join(
+      [w[0].upper() + w[1:] for w in game_name.split("_")])
+  camel_game_name += game_mode
+  # Create and register the Problem
+  cls = type("Gym%sRandom" % camel_game_name,
+             (T2TGymEnv,), {"base_env_name": camel_game_name})
+  registry.register_problem(cls)
+
+
+# Register the atari games with all of the possible modes.
+for atari_game in ATARI_GAMES:
+  for atari_game_mode in ATARI_GAME_MODES:
+    register_game(atari_game, game_mode=atari_game_mode)
diff --git a/tensor2tensor/data_generators/gym_problems.py b/tensor2tensor/data_generators/gym_problems.py
deleted file mode 100644
index ae25187bc..000000000
--- a/tensor2tensor/data_generators/gym_problems.py
+++ /dev/null
@@ -1,795 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Data generators for Gym environments."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import collections
-import json
-import math
-import os
-import gym
-import numpy as np
-import six
-
-from tensor2tensor.data_generators import problem
-from tensor2tensor.data_generators import video_utils
-from tensor2tensor.models.research import rl
-from tensor2tensor.rl import collect
-from tensor2tensor.rl.envs import tf_atari_wrappers
-from tensor2tensor.utils import metrics
-from tensor2tensor.utils import registry
-
-import tensorflow as tf
-
-flags = tf.flags
-FLAGS = flags.FLAGS
-
-
-flags.DEFINE_string("agent_policy_path", None, "File with model for agent.")
-
-flags.DEFINE_string("autoencoder_path", None,
-                    "File with model for autoencoder.")
-
-
-frame_dumper_use_disk = False  # Whether to use memory or disk to dump frames.
-frame_dumper = {}
-
-
-class GymDiscreteProblem(video_utils.VideoProblem):
-  """Gym environment with discrete actions and rewards."""
-
-  def __init__(self, *args, **kwargs):
-    super(GymDiscreteProblem, self).__init__(*args, **kwargs)
-    # TODO(piotrmilos): Check if self._env is used.
-    self._env = None
-
-    self.debug_dump_frames_path = "debug_frames_env"
-    self.settable_num_steps = 5000
-
-    self._environment_spec = None
-    self.settable_eval_phase = False
-
-    self._internal_memory_size = 20
-    self._internal_memory_force_beginning_resets = False
-    self._session = None
-    self.statistics = BasicStatistics()
-    self._use_dumper_data = False
-    self._dumper_data_index = 0
-    self._forced_collect_level = None
-
-  @property
-  def resize_height_factor(self):
-    return 1
-
-  @property
-  def resize_width_factor(self):
-    return 1
-
-  @property
-  def grayscale(self):
-    return False
-
-  @property
-  def num_channels(self):
-    """Number of color channels in each frame."""
-    return 1 if self.grayscale else 3
-
-  def _setup(self, data_dir, extra_collect_hparams=None,
-             override_collect_hparams=None):
-    dumper_path = os.path.join(data_dir, "dumper")
-    dumper_exists = tf.gfile.Exists(dumper_path)
-    tf.logging.info("Dumper path %s." % dumper_path)
-    if dumper_exists and not self.settable_eval_phase:
-      tf.logging.info("Using dumper data.")
-      self._use_dumper_data = True
-      self._dumper_data_index = 0
-      self._dumper_path = dumper_path
-    else:
-      # TODO(piotrmilos):this should be consistent with
-      # ppo_params in model_rl_experiment
-      collect_hparams = rl.ppo_pong_base()
-      collect_hparams.add_hparam("environment_spec", self.environment_spec)
-      collect_hparams.add_hparam("force_beginning_resets",
-                                 self._internal_memory_force_beginning_resets)
-      collect_hparams.epoch_length = self._internal_memory_size
-      collect_hparams.num_agents = 1
-
-      if not FLAGS.agent_policy_path:
-        collect_hparams.policy_network = rl.random_policy_fun
-
-      if extra_collect_hparams is not None:
-        for (key, value) in six.iteritems(extra_collect_hparams):
-          collect_hparams.add_hparam(key, value)
-
-      if override_collect_hparams is not None:
-        # Override hparams manually - HParams.override_from_dict does not work
-        # with functions.
-        for (key, value) in six.iteritems(override_collect_hparams):
-          setattr(collect_hparams, key, value)
-
-      policy_to_actions_lambda = None
-      if self.settable_eval_phase:
-        policy_to_actions_lambda = lambda policy: policy.mode()
-
-      collect_level = 2  # After Resize and RewardClipping.
-      if collect_hparams.environment_spec.simulated_env:
-        collect_level = 1  # We still have reward clipping.
-      if self._forced_collect_level is not None:  # For autoencoders.
-        collect_level = self._forced_collect_level
-
-      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-        self.collect_memory, self.collect_trigger_op, collect_init = (
-            collect.define_collect(
-                collect_hparams,
-                scope="gym_problems",
-                eval_phase=False,
-                collect_level=collect_level,
-                policy_to_actions_lambda=policy_to_actions_lambda))
-
-      self._session = tf.Session()
-      collect_init(self._session)
-      self._session.run(tf.global_variables_initializer())
-      self.restore_networks(self._session)
-      self.memory_index = 0
-      self.memory = None
-
-  @property
-  def random_skip(self):
-    return False
-
-  def _get_data(self):
-    if self._use_dumper_data:
-      file_path = os.path.join(self._dumper_path,
-                               "frame_{}.npz".format(self._dumper_data_index))
-      if frame_dumper_use_disk:
-        with tf.gfile.Open(file_path) as gfile:
-          data = np.load(gfile)
-      else:
-        data = frame_dumper.pop(file_path)
-      self._dumper_data_index += 1
-      return (data["observ"][0, ...], data["reward"][0], data["done"][0],
-              data["action"][0])
-    else:
-      if self.memory is None or self.memory_index >= self._internal_memory_size:
-        self.memory = self._session.run(self.collect_memory)
-        self.memory_index = 0
-      data = [self.memory[i][self.memory_index][0] for i in range(4)]
-      self.memory_index += 1
-
-      return data
-
-  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
-    self._setup(data_dir)
-
-    self.debug_dump_frames_path = os.path.join(
-        data_dir, self.debug_dump_frames_path)
-
-    frame_counter = 0
-    pieces_generated = 0
-    prev_reward = 0
-    prev_done = False
-
-    # TODO(piotrmilos): self.settable_eval_phase possibly violates sematics
-    # of VideoProblem
-    while pieces_generated < self.num_steps or self.settable_eval_phase:
-      data = self._get_data()
-      observation, reward, done, action = data
-
-      debug_image = self.collect_statistics_and_generate_debug_image(
-          pieces_generated, *data)
-      ret_dict = {
-          "frame": observation,
-          "frame_number": [int(frame_counter)],
-          "image/format": ["png"],
-          "image/height": [self.frame_height],
-          "image/width": [self.frame_width],
-          "action": [int(action)],
-          "done": [int(prev_done)],
-          "reward": [int(prev_reward - self.min_reward)]
-      }
-
-      if debug_image is not None:
-        ret_dict["image/debug"] = debug_image
-
-      yield ret_dict
-
-      if done and self.settable_eval_phase:
-        return
-
-      prev_done, prev_reward = done, reward
-
-      pieces_generated += 1
-      frame_counter += 1
-      if done:
-        frame_counter = 0
-
-  def restore_networks(self, sess):
-    if FLAGS.agent_policy_path:
-      model_saver = tf.train.Saver(
-          tf.global_variables(".*network_parameters.*"))
-      ckpts = tf.train.get_checkpoint_state(FLAGS.agent_policy_path)
-      ckpt = ckpts.model_checkpoint_path
-      model_saver.restore(sess, ckpt)
-
-  def eval_metrics(self):
-    eval_metrics = [
-        metrics.Metrics.ACC, metrics.Metrics.ACC_PER_SEQ,
-        metrics.Metrics.IMAGE_RMSE
-    ]
-    return eval_metrics
-
-  @property
-  def extra_reading_spec(self):
-    """Additional data fields to store on disk and their decoders."""
-
-    # TODO(piotrmilos): shouldn't done be included here?
-    data_fields = {
-        "frame_number": tf.FixedLenFeature([1], tf.int64),
-        "action": tf.FixedLenFeature([1], tf.int64),
-        "reward": tf.FixedLenFeature([1], tf.int64)
-    }
-    decoders = {
-        "frame_number":
-            tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="frame_number"),
-        "action":
-            tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="action"),
-        "reward":
-            tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="reward"),
-    }
-    return data_fields, decoders
-
-  def get_environment_spec(self):
-    return rl.standard_atari_env_spec(
-        self.env_name,
-        resize_height_factor=self.resize_height_factor,
-        resize_width_factor=self.resize_width_factor,
-        grayscale=self.grayscale)
-
-  @property
-  def environment_spec(self):
-    if self._environment_spec is None:
-      self._environment_spec = self.get_environment_spec()
-    return self._environment_spec
-
-  @property
-  def is_generate_per_split(self):
-    """Whether we have a train/test split or just hold out data."""
-    return False  # Just hold out some generated data for evals.
-
-  @property
-  def env_name(self):
-    """This is the name of the Gym environment for this problem."""
-    raise NotImplementedError()
-
-  @property
-  def env(self):
-    # TODO(piotrmilos): possibly remove
-    if self._env is None:
-      self._env = gym.make(self.env_name)
-    return self._env
-
-  @property
-  def num_actions(self):
-    return self.env.action_space.n
-
-  # pylint: disable=unused-argument
-  def collect_statistics_and_generate_debug_image(self, index, observation,
-                                                  reward, done, action):
-    """This generates extra statistics and debug images."""
-    return None
-  # pylint: enable=unused-argument
-
-  @property
-  def frame_height(self):
-    return self.env.observation_space.shape[0] // self.resize_height_factor
-
-  @property
-  def frame_width(self):
-    return self.env.observation_space.shape[1] // self.resize_width_factor
-
-  @property
-  def num_rewards(self):
-    raise NotImplementedError()
-
-  @property
-  def num_steps(self):
-    return self.settable_num_steps
-
-  @property
-  def total_number_of_frames(self):
-    return self.num_steps
-
-  @property
-  def min_reward(self):
-    raise NotImplementedError()
-
-  @property
-  def num_testing_steps(self):
-    return None
-
-  @property
-  def only_keep_videos_from_0th_frame(self):
-    return False
-
-  def hparams(self, defaults, unused_model_hparams):
-    p = defaults
-    p.input_modality = {
-        "inputs": ("video", 256),
-        "input_reward": ("symbol:weights_all", self.num_rewards),
-        "input_action": ("symbol:weights_all", self.num_actions)
-    }
-    p.target_modality = {
-        "targets": ("video", 256),
-        "target_reward": ("symbol:weights_all", self.num_rewards),
-        "target_action": ("symbol:weights_all", self.num_actions)
-    }
-    p.input_space_id = problem.SpaceID.IMAGE
-    p.target_space_id = problem.SpaceID.IMAGE
-
-  def generate_data(self, data_dir, tmp_dir, task_id=-1):
-    super(GymDiscreteProblem, self).generate_data(data_dir, tmp_dir, task_id)
-    # Save stats to file, or restore if data was already generated.
-    stats_file = os.path.join(data_dir,
-                              "%s.stats.json" % self.dataset_filename())
-    if tf.gfile.Exists(stats_file):
-      self.statistics.update_from_file(stats_file)
-    else:
-      self.statistics.save_to_file(stats_file)
-
-  def filepattern(self, data_dir, mode, shard=None, only_last=False):
-    filepattern = super(GymDiscreteProblem, self).filepattern(
-        data_dir, mode, shard
-    )
-    if only_last:
-      filepattern += r"10.[\d+]"
-    return filepattern
-
-
-class BasicStatistics(object):
-  """Keeps basic statistics to calculate mean reward """
-
-  def __init__(self):
-    self.sum_of_rewards = 0.0
-    self.number_of_dones = 0
-    self.sum_of_rewards_current_episode = 0.0
-    self.last_done = False
-
-  def update_from_dict(self, stats_dict):
-    keys = set(self.to_dict().keys())
-    for k, v in stats_dict.items():
-      if k not in keys:
-        raise ValueError("Key %s not a property of %s" %
-                         (k, type(self).__name__))
-      setattr(self, k, v)
-    return self
-
-  def to_dict(self):
-    # Cast the values to base types as some are numpy types.
-    keys_and_types = [
-        ("sum_of_rewards", float),
-        ("number_of_dones", int),
-        ("sum_of_rewards_current_episode", float),
-        ("last_done", bool),
-    ]
-    stats_dict = dict([(k, t(getattr(self, k))) for k, t in keys_and_types])
-    return stats_dict
-
-  def save_to_file(self, fname):
-    with tf.gfile.Open(fname, "w") as f:
-      f.write(json.dumps(self.to_dict()))
-
-  def update_from_file(self, fname):
-    with tf.gfile.Open(fname) as f:
-      self.update_from_dict(json.loads(f.read()))
-      return self
-
-
-# TODO(piotrmilos): merge with the superclass
-class GymRealDiscreteProblem(GymDiscreteProblem):
-  """Discrete problem."""
-
-  def __init__(self, *args, **kwargs):
-    super(GymRealDiscreteProblem, self).__init__(*args, **kwargs)
-    self.statistics = BasicStatistics()
-
-    self.make_extra_debug_info = False
-
-  def collect_statistics_and_generate_debug_image(self, index, observation,
-                                                  reward, done, action):
-    """Collects info required to calculate mean reward."""
-
-    self.statistics.sum_of_rewards_current_episode += reward
-    # we ignore consecutive dones as they are artefacts of skip wrappers
-    if done and not self.statistics.last_done:
-      self.statistics.number_of_dones += int(done)
-      self.statistics.sum_of_rewards += (
-          self.statistics.sum_of_rewards_current_episode)
-      self.statistics.sum_of_rewards_current_episode = 0.0
-
-    self.statistics.last_done = done
-
-    debug_image = None
-    return debug_image
-
-
-class GymDiscreteProblemWithAutoencoder(GymRealDiscreteProblem):
-  """Gym discrete problem with autoencoder."""
-
-  def __init__(self, *args, **kwargs):
-    super(GymDiscreteProblemWithAutoencoder, self).__init__(*args, **kwargs)
-    self._forced_collect_level = 0
-
-  def get_environment_spec(self):
-    return rl.standard_atari_ae_env_spec(self.env_name, self.ae_hparams_set)
-
-  def restore_networks(self, sess):
-    super(GymDiscreteProblemWithAutoencoder, self).restore_networks(sess)
-    if FLAGS.autoencoder_path:
-      autoencoder_saver = tf.train.Saver(
-          tf.global_variables("autoencoder.*"))
-      ckpts = tf.train.get_checkpoint_state(FLAGS.autoencoder_path)
-      ckpt = ckpts.model_checkpoint_path
-      autoencoder_saver.restore(sess, ckpt)
-
-  def hparams(self, defaults, unused_model_hparams):
-    """Overrides VideoProblem.hparams to work on images instead of videos."""
-    p = defaults
-    p.input_modality = {
-        "inputs": ("image", 256),
-    }
-    p.target_modality = ("image", 256)
-    p.input_space_id = problem.SpaceID.IMAGE
-    p.target_space_id = problem.SpaceID.IMAGE
-
-  def preprocess(self, dataset, mode, hparams, interleave=True):
-    """Overrides VideoProblem.preprocess to work on images instead of videos."""
-    def set_targets(example):
-      example["targets"] = example["frame"]
-      return example
-    return dataset.map(set_targets)
-
-
-class GymDiscreteProblemAutoencoded(GymRealDiscreteProblem):
-  """Gym discrete problem with frames already autoencoded."""
-
-  def __init__(self, *args, **kwargs):
-    super(GymDiscreteProblemAutoencoded, self).__init__(*args, **kwargs)
-    self._forced_collect_level = 0
-
-  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
-    raise RuntimeError("GymDiscreteProblemAutoencoded can be used only"
-                       " for reading encoded frames")
-
-  def get_environment_spec(self):
-    return rl.standard_atari_ae_env_spec(self.env_name, self.ae_hparams_set)
-
-  @property
-  def autoencoder_factor(self):
-    """By how much to divide sizes when using autoencoders."""
-    hparams = registry.hparams(self.ae_hparams_set)
-    return 2**hparams.num_hidden_layers
-
-  @property
-  def frame_height(self):
-    height = self.env.observation_space.shape[0]
-    ae_height = int(math.ceil(height / self.autoencoder_factor))
-    return ae_height
-
-  @property
-  def frame_width(self):
-    width = self.env.observation_space.shape[1]
-    return int(math.ceil(width / self.autoencoder_factor))
-
-
-class RewardPerSequenceStatistics(BasicStatistics):
-  """This encapsulates all pieces required to calculate
-  the correctness of rewards per sequence metric
-  """
-
-  def __init__(self, rollout_fractions):
-    super(RewardPerSequenceStatistics, self).__init__()
-
-    # data to calculate
-    # correctness of rewards per sequence metric
-    self.episode_sim_reward = 0.0
-    self.episode_real_reward = 0.0
-    self.successful_episode_reward_predictions = collections.OrderedDict([
-        (frac, 0) for frac in rollout_fractions
-    ])
-    self.report_reward_statistics_every = 10
-    # auxiliary objects
-    self.real_obs = None
-    self.real_rewards = None
-
-  def to_dict(self):
-    stats_dict = super(RewardPerSequenceStatistics, self).to_dict()
-    keys_and_types = [
-        ("episode_sim_reward", float),
-        ("episode_real_reward", float),
-        ("successful_episode_reward_predictions", collections.OrderedDict),
-        ("report_reward_statistics_every", int),
-    ]
-    additional = dict([(k, t(getattr(self, k))) for k, t in keys_and_types])
-    stats_dict.update(additional)
-    return stats_dict
-
-
-class GymSimulatedDiscreteProblem(GymDiscreteProblem):
-  """Simulated gym environment with discrete actions and rewards."""
-
-  def __init__(self, *args, **kwargs):
-    super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
-    self.debug_dump_frames_path = "debug_frames_sim"
-
-    # This is hackish way of introducing resets every
-    # self.num_testing_steps. It cannot be done easily
-    # using other ways as we do not control
-    # the amount of skips induced but wrappers
-    self._internal_memory_size = self.num_testing_steps
-    self._internal_memory_force_beginning_resets = True
-
-    self.statistics = BasicStatistics()
-    self._initial_frame_chooser = None
-
-  def _setup(self, data_dir, extra_collect_hparams=None,
-             override_collect_hparams=None):
-    if extra_collect_hparams is None:
-      extra_collect_hparams = {}
-
-    extra_collect_hparams["initial_frame_chooser"] = self._initial_frame_chooser
-
-    super(GymSimulatedDiscreteProblem, self)._setup(
-        data_dir, extra_collect_hparams, override_collect_hparams
-    )
-
-  @property
-  def initial_frames_problem(self):
-    raise NotImplementedError()
-
-  @property
-  def num_input_frames(self):
-    """Number of frames on input for real environment."""
-    # TODO(lukaszkaiser): This must be equal to hparams.video_num_input_frames,
-    # we should automate this to avoid bug in the future.
-    return 4
-
-  @property
-  def video_num_target_frames(self):
-    """Number of frames on input for real environment."""
-    # TODO(piotrmilos): This must be equal to hparams.video_num_target_frames,
-    # we should automate this to avoid bug in the future.
-    return 1
-
-  @property
-  def num_testing_steps(self):
-    return None
-
-  def get_environment_spec(self):
-    env_spec = rl.standard_atari_env_spec(
-        self.env_name,
-        simulated=True,
-        resize_height_factor=self.resize_height_factor,
-        resize_width_factor=self.resize_width_factor,
-        grayscale=self.grayscale)
-    env_spec.add_hparam("simulation_random_starts", True)
-    env_spec.add_hparam("simulation_flip_first_random_for_beginning", True)
-    env_spec.add_hparam("intrinsic_reward_scale", 0.0)
-    initial_frames_problem = registry.problem(self.initial_frames_problem)
-    env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
-    env_spec.add_hparam("video_num_input_frames", self.num_input_frames)
-    env_spec.add_hparam("video_num_target_frames", self.video_num_target_frames)
-
-    return env_spec
-
-  def restore_networks(self, sess):
-    super(GymSimulatedDiscreteProblem, self).restore_networks(sess)
-    # TODO(blazej): adjust regexp for different models.
-    # TODO(piotrmilos): move restoring networks to SimulatedBatchEnv.initialize
-    env_model_loader = tf.train.Saver(tf.global_variables("next_frame*"))
-    ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
-    ckpt = ckpts.model_checkpoint_path
-    env_model_loader.restore(sess, ckpt)
-
-
-class GymSimulatedDiscreteProblemForWorldModelEval(GymSimulatedDiscreteProblem):
-  """Simulated gym environment for evaluating world model."""
-
-  def __init__(self, *args, **kwargs):
-    super(GymSimulatedDiscreteProblemForWorldModelEval, self).__init__(
-        *args, **kwargs
-    )
-    self.settable_rollout_fractions = [1]
-    self.statistics = RewardPerSequenceStatistics(
-        self.settable_rollout_fractions
-    )
-
-  def get_environment_spec(self):
-    env_spec = super(
-        GymSimulatedDiscreteProblemForWorldModelEval, self
-    ).get_environment_spec()
-    env_spec.simulation_flip_first_random_for_beginning = False
-    return env_spec
-
-  def _setup(self, data_dir):
-    trajectory_length = self.num_testing_steps
-    if self.num_steps < 1200:
-      # Decrease the trajectory length for tiny experiments, otherwise we don't
-      # have enough data to run the evaluation.
-      trajectory_length = 2
-    frame_index = tf.Variable(0, trainable=False)
-
-    def fixed_action_policy_fun(action_space, unused_config, observations):
-      """Policy which replays actions from a trajectory."""
-      action = self._initial_frame_chooser.trajectory["action"].read_value()[
-          :, frame_index.read_value(), :
-      ]
-      inc_frame_index = frame_index.assign(
-          (frame_index.read_value() + 1) % trajectory_length
-      )
-      with tf.control_dependencies([inc_frame_index]):
-        action = tf.identity(action)
-
-      obs_shape = observations.shape.as_list()
-      with tf.variable_scope("network_parameters"):
-        probs = tf.one_hot(
-            tf.transpose(action), depth=action_space.n
-        )
-        policy = tf.distributions.Categorical(probs=probs)
-        value = tf.zeros(obs_shape[:2])
-      return rl.NetworkOutput(policy, value, lambda a: a)
-
-    super(GymSimulatedDiscreteProblemForWorldModelEval, self)._setup(
-        data_dir, override_collect_hparams={
-            "policy_network": fixed_action_policy_fun
-        }
-    )
-
-  def collect_statistics_and_generate_debug_image(self, index,
-                                                  observation,
-                                                  reward, done, action):
-    stat = self.statistics
-
-    # TODO(piotrmilos): possibly make the same behaviour as
-    # in the BasicStatistics
-    stat.sum_of_rewards += reward
-    stat.episode_sim_reward += reward
-
-    if index % self._internal_memory_size == 0:
-      real_frame_tensor = {
-          key: var.read_value()[0, ...]
-          for (key, var) in six.iteritems(
-              self._initial_frame_chooser.trajectory
-          )
-      }
-      (stat.real_obs, stat.real_rewards) = self._session.run((
-          real_frame_tensor["inputs"], real_frame_tensor["reward"]
-      ))
-      stat.real_rewards += self.min_reward
-
-    real_ob = stat.real_obs[index % stat.real_obs.shape[0], ...]
-    debug_im = self._generate_debug_image(real_ob, observation)
-
-    assert (self._internal_memory_size == self.num_testing_steps and
-            self._internal_memory_force_beginning_resets), (
-                "The collect memory should be set in force_beginning_resets "
-                "mode for the code below to work properly.")
-
-    index_in_rollout = index % self._internal_memory_size + 1
-
-    if stat.episode_sim_reward == stat.episode_real_reward:
-      for frac in stat.successful_episode_reward_predictions:
-        if index_in_rollout == int(self._internal_memory_size * frac):
-          stat.successful_episode_reward_predictions[frac] += 1
-
-    if index_in_rollout == self._internal_memory_size:
-      stat.episode_sim_reward = 0.0
-      stat.episode_real_reward = 0.0
-      stat.number_of_dones += 1
-    else:
-      real_reward = stat.real_rewards[index % stat.real_rewards.shape[0], 0]
-      stat.episode_real_reward += real_reward
-
-    return debug_im
-
-  def _generate_debug_image(self, real_ob, sim_ob):
-    ob = np.ndarray.astype(sim_ob, np.int)
-    if ob.shape == real_ob.shape:
-      err = np.ndarray.astype(
-          np.maximum(np.abs(real_ob - ob, dtype=np.int) - 10, 0), np.uint8)
-      debug_im = np.concatenate([sim_ob, real_ob, err], axis=1)
-    else:
-      # Real env does not get the ResizeWrapper and we don't have it in python,
-      # so we skip the debug image here and just output observations.
-      debug_im = sim_ob
-    return debug_im
-
-
-class GymSimulatedDiscreteProblemAutoencoded(GymSimulatedDiscreteProblem):
-  """Gym simulated discrete problem with frames already autoencoded."""
-
-  def __init__(self, *args, **kwargs):
-    super(GymSimulatedDiscreteProblemAutoencoded, self).__init__(
-        *args, **kwargs)
-    self._forced_collect_level = 0
-
-  def get_environment_spec(self):
-    env_spec = rl.standard_atari_env_spec(self.env_name)
-    env_spec.wrappers = [
-        [tf_atari_wrappers.IntToBitWrapper, {}],
-        [tf_atari_wrappers.StackWrapper, {"history": 4}]
-    ]
-    env_spec.simulated_env = True
-    env_spec.add_hparam("simulation_random_starts", True)
-    env_spec.add_hparam("simulation_flip_first_random_for_beginning", True)
-    env_spec.add_hparam("intrinsic_reward_scale", 0.0)
-    initial_frames_problem = registry.problem(self.initial_frames_problem)
-    env_spec.add_hparam("initial_frames_problem", initial_frames_problem)
-    env_spec.add_hparam("video_num_input_frames", self.num_input_frames)
-    env_spec.add_hparam("video_num_target_frames", self.video_num_target_frames)
-
-    return env_spec
-
-  @property
-  def autoencoder_factor(self):
-    """By how much to divide sizes when using autoencoders."""
-    hparams = registry.hparams(self.ae_hparams_set)
-    return 2**hparams.num_hidden_layers
-
-  @property
-  def frame_height(self):
-    height = self.env.observation_space.shape[0]
-    ae_height = int(math.ceil(height / self.autoencoder_factor))
-    return ae_height
-
-  @property
-  def frame_width(self):
-    width = self.env.observation_space.shape[1]
-    return int(math.ceil(width / self.autoencoder_factor))
-
-
-class GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded(
-    GymSimulatedDiscreteProblemForWorldModelEval,
-    GymSimulatedDiscreteProblemAutoencoded):
-  """TODO(owner): Write a small docstring."""
-
-  def _generate_debug_image(self, real_ob, sim_ob):
-    def unpack(x):
-      return np.ndarray.astype(np.unpackbits(x, axis=2), np.int)
-    real_ob_unpacked = unpack(real_ob)
-    sim_ob_unpacked = unpack(sim_ob)
-    # Hamming distance on binary latent codes, seen as a grayscale image.
-    err = np.ndarray.astype(
-        np.transpose(
-            np.broadcast_to(
-                np.sum(np.abs(real_ob_unpacked - sim_ob_unpacked), axis=2) /
-                24.0 * 255,
-                # Channels first to satisfy numpy broadcasting rules.
-                shape=((real_ob.shape[2],) + real_ob.shape[:2])),
-            (1, 2, 0)),
-        np.uint8)
-    return np.concatenate([sim_ob, real_ob, err], axis=1)
-
-
-@registry.register_problem
-class DummyAutoencoderProblem(GymDiscreteProblemWithAutoencoder):
-  """Dummy problem for running the autoencoder inside AutoencoderWrapper."""
-
-  @property
-  def env_name(self):
-    return "DummyAutoencoder"
diff --git a/tensor2tensor/data_generators/gym_problems_specs.py b/tensor2tensor/data_generators/gym_problems_specs.py
deleted file mode 100644
index c904cdb39..000000000
--- a/tensor2tensor/data_generators/gym_problems_specs.py
+++ /dev/null
@@ -1,339 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Definitions of data generators for gym problems."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-# We need gym_utils for the game environments defined there.
-from tensor2tensor.data_generators import gym_utils  # pylint: disable=unused-import
-# pylint: disable=g-multiple-import
-from tensor2tensor.data_generators.gym_problems import GymDiscreteProblem,\
-  GymSimulatedDiscreteProblem, GymRealDiscreteProblem, \
-  GymDiscreteProblemWithAutoencoder, GymDiscreteProblemAutoencoded, \
-  GymSimulatedDiscreteProblemAutoencoded, \
-  GymSimulatedDiscreteProblemForWorldModelEval, \
-  GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded
-# pylint: enable=g-multiple-import
-from tensor2tensor.utils import registry
-
-# Game list from our list of ROMs
-# Removed because XDeterministic-v4 did not exist:
-# * adventure
-# * defender
-# * kaboom
-ATARI_GAMES = [
-    "air_raid", "alien", "amidar", "assault", "asterix", "asteroids",
-    "atlantis", "bank_heist", "battle_zone", "beam_rider", "berzerk", "bowling",
-    "boxing", "breakout", "carnival", "centipede", "chopper_command",
-    "crazy_climber", "demon_attack", "double_dunk", "elevator_action", "enduro",
-    "fishing_derby", "freeway", "frostbite", "gopher", "gravitar", "hero",
-    "ice_hockey", "jamesbond", "journey_escape", "kangaroo", "krull",
-    "kung_fu_master", "montezuma_revenge", "ms_pacman", "name_this_game",
-    "phoenix", "pitfall", "pong", "pooyan", "private_eye", "qbert", "riverraid",
-    "road_runner", "robotank", "seaquest", "skiing", "solaris",
-    "space_invaders", "star_gunner", "tennis", "time_pilot", "tutankham",
-    "up_n_down", "venture", "video_pinball", "wizard_of_wor", "yars_revenge",
-    "zaxxon"
-]
-
-# List from paper:
-# https://arxiv.org/pdf/1805.11593.pdf
-# plus frostbite.
-ATARI_GAMES_WITH_HUMAN_SCORE = [
-    "alien", "amidar", "assault", "asterix", "asteroids",
-    "atlantis", "bank_heist", "battle_zone", "beam_rider", "bowling",
-    "boxing", "breakout", "chopper_command",
-    "crazy_climber", "demon_attack", "double_dunk", "enduro",
-    "fishing_derby", "freeway", "frostbite", "gopher", "gravitar", "hero",
-    "ice_hockey", "jamesbond", "kangaroo", "krull",
-    "kung_fu_master", "montezuma_revenge", "ms_pacman", "name_this_game",
-    "pitfall", "pong", "private_eye", "qbert", "riverraid",
-    "road_runner", "seaquest", "solaris",
-    "up_n_down", "video_pinball", "yars_revenge",
-]
-
-ATARI_WHITELIST_GAMES = [
-    "amidar",
-    "bank_heist",
-    "berzerk",
-    "boxing",
-    "crazy_climber",
-    "freeway",
-    "frostbite",
-    "gopher",
-    "kung_fu_master",
-    "ms_pacman",
-    "pong",
-    "qbert",
-    "seaquest",
-]
-
-
-# Games on which model-free does better than model-based at this point.
-ATARI_CURIOUS_GAMES = [
-    "bank_heist",
-    "boxing",
-    "enduro",
-    "kangaroo",
-    "road_runner",
-    "up_n_down",
-]
-
-
-# Games on which based should work.
-ATARI_DEBUG_GAMES = [
-    "crazy_climber",
-    "freeway",
-    "pong",
-]
-
-
-# Games for which we hard-define problems to run all around.
-# TODO(lukaszkaiser): global registration makes them all rescaled and grayscale,
-# no matter the setting of hparams later (as they're registered at start).
-ATARI_ALL_MODES_SHORT_LIST = []  # ATARI_DEBUG_GAMES + ATARI_CURIOUS_GAMES
-
-
-# Different ATARI game modes in OpenAI Gym. Full list here:
-# https://github.com/openai/gym/blob/master/gym/envs/__init__.py
-ATARI_GAME_MODES = [
-    "Deterministic-v0",  # 0.25 repeat action probability, 4 frame skip.
-    "Deterministic-v4",  # 0.00 repeat action probability, 4 frame skip.
-    "NoFrameskip-v0",    # 0.25 repeat action probability, 1 frame skip.
-    "NoFrameskip-v4",    # 0.00 repeat action probability, 1 frame skip.
-    "-v0",               # 0.25 repeat action probability, (2 to 5) frame skip.
-    "-v4"                # 0.00 repeat action probability, (2 to 5) frame skip.
-]
-
-# List of all ATARI envs in all modes.
-ATARI_PROBLEMS = {}
-
-
-@registry.register_problem
-class GymWrappedFullPongRandom(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  @property
-  def env_name(self):
-    return "T2TPongWarmUp20RewSkipFull-v1"
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedFullPong(GymRealDiscreteProblem,
-                                                   GymWrappedFullPongRandom):
-  pass
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedFullPongWithAutoencoder(
-    GymDiscreteProblemWithAutoencoder, GymWrappedFullPongRandom):
-  pass
-
-
-@registry.register_problem
-class GymDiscreteProblemWithAgentOnWrappedFullPongAutoencoded(
-    GymDiscreteProblemAutoencoded, GymWrappedFullPongRandom):
-  pass
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedFullPong(
-    GymSimulatedDiscreteProblem, GymWrappedFullPongRandom):
-  """Simulated pong."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_full_pong"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOnWrappedFullPong(
-    GymSimulatedDiscreteProblemForWorldModelEval, GymWrappedFullPongRandom):
-  """Simulated pong for world model evaluation."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_full_pong"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemWithAgentOnWrappedFullPongAutoencoded(
-    GymSimulatedDiscreteProblemAutoencoded, GymWrappedFullPongRandom):
-  """GymSimulatedDiscreteProblemWithAgentOnWrappedFullPongAutoencoded."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_full_pong_autoencoded"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-@registry.register_problem
-class GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOnWrappedFullPongAutoencoded(  # pylint: disable=line-too-long
-    GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded,
-    GymWrappedFullPongRandom):
-  """Simulated pong for world model evaluation with encoded frames."""
-
-  @property
-  def initial_frames_problem(self):
-    return "gym_discrete_problem_with_agent_on_wrapped_full_pong_autoencoded"
-
-  @property
-  def num_testing_steps(self):
-    return 100
-
-
-class GymClippedRewardRandom(GymDiscreteProblem):
-  """Abstract base class for clipped reward games."""
-
-  @property
-  def env_name(self):
-    raise NotImplementedError
-
-  @property
-  def min_reward(self):
-    return -1
-
-  @property
-  def num_rewards(self):
-    return 3
-
-
-def create_problems_for_game(
-    game_name,
-    resize_height_factor=2,
-    resize_width_factor=2,
-    grayscale=True,
-    game_mode="Deterministic-v4",
-    autoencoder_hparams=None):
-  """Create and register problems for game_name.
-
-  Args:
-    game_name: str, one of the games in ATARI_GAMES, e.g. "bank_heist".
-    resize_height_factor: factor by which to resize the height of frames.
-    resize_width_factor: factor by which to resize the width of frames.
-    grayscale: whether to make frames grayscale.
-    game_mode: the frame skip and sticky keys config.
-    autoencoder_hparams: the hparams for the autoencoder.
-
-  Returns:
-    dict of problems with keys ("base", "agent", "simulated").
-
-  Raises:
-    ValueError: if clipped_reward=False or game_name not in ATARI_GAMES.
-  """
-  if game_name not in ATARI_GAMES:
-    raise ValueError("Game %s not in ATARI_GAMES" % game_name)
-  if game_mode not in ATARI_GAME_MODES:
-    raise ValueError("Unknown ATARI game mode: %s." % game_mode)
-  camel_game_name = "".join(
-      [w[0].upper() + w[1:] for w in game_name.split("_")])
-  camel_game_name += game_mode
-  env_name = camel_game_name
-
-  # Create and register the Random and WithAgent Problem classes
-  problem_cls = type("Gym%sRandom" % camel_game_name,
-                     (GymClippedRewardRandom,),
-                     {"env_name": env_name,
-                      "resize_height_factor": resize_height_factor,
-                      "resize_width_factor": resize_width_factor,
-                      "grayscale": grayscale})
-  registry.register_problem(problem_cls)
-
-  with_agent_cls = type("GymDiscreteProblemWithAgentOn%s" % camel_game_name,
-                        (GymRealDiscreteProblem, problem_cls), {})
-  registry.register_problem(with_agent_cls)
-
-  with_ae_cls = type(
-      "GymDiscreteProblemWithAgentOn%sWithAutoencoder" % camel_game_name,
-      (GymDiscreteProblemWithAutoencoder, problem_cls),
-      {"ae_hparams_set": autoencoder_hparams})
-  registry.register_problem(with_ae_cls)
-
-  ae_cls = type(
-      "GymDiscreteProblemWithAgentOn%sAutoencoded" % camel_game_name,
-      (GymDiscreteProblemAutoencoded, problem_cls),
-      {"ae_hparams_set": autoencoder_hparams})
-  registry.register_problem(ae_cls)
-
-  # Create and register the simulated Problem
-  simulated_cls = type(
-      "GymSimulatedDiscreteProblemWithAgentOn%s" % camel_game_name,
-      (GymSimulatedDiscreteProblem, problem_cls), {
-          "initial_frames_problem": with_agent_cls.name,
-          "num_testing_steps": 100
-      })
-  registry.register_problem(simulated_cls)
-
-  simulated_ae_cls = type(
-      "GymSimulatedDiscreteProblemWithAgentOn%sAutoencoded" % camel_game_name,
-      (GymSimulatedDiscreteProblemAutoencoded, problem_cls), {
-          "initial_frames_problem": ae_cls.name,
-          "num_testing_steps": 100,
-          "ae_hparams_set": autoencoder_hparams
-      })
-  registry.register_problem(simulated_ae_cls)
-
-  # Create and register the simulated Problem
-  world_model_eval_cls = type(
-      "GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%s" %
-      camel_game_name,
-      (GymSimulatedDiscreteProblemForWorldModelEval, problem_cls), {
-          "initial_frames_problem": with_agent_cls.name,
-          "num_testing_steps": 100,
-          "ae_hparams_set": autoencoder_hparams
-      })
-  registry.register_problem(world_model_eval_cls)
-
-  world_model_eval_ae_cls = type(
-      "GymSimulatedDiscreteProblemForWorldModelEvalWithAgentOn%sAutoencoded" %
-      camel_game_name,
-      (GymSimulatedDiscreteProblemForWorldModelEvalAutoencoded, problem_cls), {
-          "initial_frames_problem": ae_cls.name,
-          "num_testing_steps": 100,
-          "ae_hparams_set": autoencoder_hparams
-      })
-  registry.register_problem(world_model_eval_ae_cls)
-
-
-# Register the atari games with all of the possible modes.
-for game in ATARI_ALL_MODES_SHORT_LIST:
-  for mode in ATARI_GAME_MODES:
-    create_problems_for_game(game, game_mode=mode)
diff --git a/tensor2tensor/data_generators/gym_problems_test.py b/tensor2tensor/data_generators/gym_problems_test.py
deleted file mode 100644
index 326e75d45..000000000
--- a/tensor2tensor/data_generators/gym_problems_test.py
+++ /dev/null
@@ -1,44 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Gym generators tests."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-import shutil
-
-from tensor2tensor.data_generators import gym_problems_specs
-
-import tensorflow as tf
-
-
-class GymProblemsTest(tf.test.TestCase):
-
-  @classmethod
-  def setUpClass(cls):
-    cls.tmp_dir = tf.test.get_temp_dir()
-    shutil.rmtree(cls.tmp_dir)
-    os.mkdir(cls.tmp_dir)
-
-  def testGymAtariGameModes(self):
-    problem = gym_problems_specs.GymDiscreteProblemWithAgentOnWrappedFullPong()
-    self.assertEqual(210, problem.frame_height)
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/rl/datagen_with_agent.py b/tensor2tensor/rl/datagen_with_agent.py
index a5d85d38b..f634e8c87 100644
--- a/tensor2tensor/rl/datagen_with_agent.py
+++ b/tensor2tensor/rl/datagen_with_agent.py
@@ -22,7 +22,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.data_generators import gym_problems_specs
+from tensor2tensor.data_generators import gym_env
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -46,7 +46,7 @@ def main(_):
   # Create problem if not already defined
   problem_name = "gym_discrete_problem_with_agent_on_%s" % FLAGS.game
   if problem_name not in registry.list_problems():
-    gym_problems_specs.create_problems_for_game(FLAGS.game)
+    gym_env.register_game(FLAGS.game)
 
   # Generate
   tf.logging.info("Running %s environment for %d steps for trajectories.",
diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/model_rl_experiment_player.py
index a5da41afc..1c0bd4617 100644
--- a/tensor2tensor/rl/model_rl_experiment_player.py
+++ b/tensor2tensor/rl/model_rl_experiment_player.py
@@ -33,7 +33,7 @@
 from PIL import ImageDraw
 from PIL import ImageFont
 
-from tensor2tensor.data_generators import gym_problems_specs
+from tensor2tensor.data_generators import gym_env
 from tensor2tensor.models.research.rl import get_policy
 from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
 from tensor2tensor.rl.trainer_model_based import FLAGS
@@ -229,7 +229,7 @@ def main(_):
     subdirectories.append("autoencoder")
   directories = setup_directories(output_dir, subdirectories)
 
-  if hparams.game in gym_problems_specs.ATARI_GAMES:
+  if hparams.game in gym_env.ATARI_GAMES:
     game_with_mode = hparams.game + "_deterministic-v4"
   else:
     game_with_mode = hparams.game
@@ -244,8 +244,7 @@ def main(_):
     if simulated_problem_name not in registry.list_problems():
       tf.logging.info("Game Problem %s not found; dynamically registering",
                       simulated_problem_name)
-      gym_problems_specs.create_problems_for_game(hparams.game,
-                                                  game_mode="Deterministic-v4")
+      gym_env.register_game(hparams.game, game_mode="Deterministic-v4")
 
   epoch = hparams.epochs-1
   epoch_data_dir = os.path.join(directories["data"], str(epoch))
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 66160d53a..24368a86d 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -21,7 +21,7 @@
 import six
 
 
-from tensor2tensor.data_generators import gym_problems_specs
+from tensor2tensor.data_generators import gym_env
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -535,38 +535,38 @@ def rlmb_scheduled_sampling(rhp):
 @registry.register_ranged_hparams
 def rlmb_all_games(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_GAMES)
+  rhp.set_categorical("loop.game", gym_env.ATARI_GAMES)
 
 
 @registry.register_ranged_hparams
 def rlmb_whitelisted_games(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_categorical("loop.game", gym_env.ATARI_WHITELIST_GAMES)
 
 
 @registry.register_ranged_hparams
 def rlmb_human_score_games(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(10)))
   rhp.set_categorical("loop.game",
-                      gym_problems_specs.ATARI_GAMES_WITH_HUMAN_SCORE)
+                      gym_env.ATARI_GAMES_WITH_HUMAN_SCORE)
 
 
 @registry.register_ranged_hparams
 def rlmb_curious_games10(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
+  rhp.set_categorical("loop.game", gym_env.ATARI_CURIOUS_GAMES)
 
 
 @registry.register_ranged_hparams
 def rlmb_curious_games5(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_CURIOUS_GAMES)
+  rhp.set_categorical("loop.game", gym_env.ATARI_CURIOUS_GAMES)
 
 
 @registry.register_ranged_hparams
 def rlmb_debug_games(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_DEBUG_GAMES)
+  rhp.set_categorical("loop.game", gym_env.ATARI_DEBUG_GAMES)
 
 
 @registry.register_ranged_hparams
@@ -638,28 +638,28 @@ def rlmb_dummy_range(rhp):
 
 @registry.register_ranged_hparams
 def rlmb_epochs_num(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_categorical("loop.game", gym_env.ATARI_WHITELIST_GAMES)
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
   rhp.set_discrete("loop.epochs", [3, 6, 12])
 
 
 @registry.register_ranged_hparams
 def rlmb_ppo_epochs_num(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_categorical("loop.game", gym_env.ATARI_WHITELIST_GAMES)
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
   rhp.set_discrete("loop.ppo_epochs_num", [200, 1000, 2000, 4000])
 
 
 @registry.register_ranged_hparams
 def rlmb_ppo_epoch_len(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_categorical("loop.game", gym_env.ATARI_WHITELIST_GAMES)
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
   rhp.set_discrete("loop.ppo_epoch_length", [25, 50, 100])
 
 
 @registry.register_ranged_hparams
 def rlmb_num_frames(rhp):
-  rhp.set_categorical("loop.game", gym_problems_specs.ATARI_WHITELIST_GAMES)
+  rhp.set_categorical("loop.game", gym_env.ATARI_WHITELIST_GAMES)
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
   rhp.set_discrete("loop.num_real_env_frames",
                    [1000*el for el in [30, 100, 500, 1000]])
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index de1e4fdc9..431fae8b0 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -269,7 +269,7 @@ def parse_problem_name(problem_name):
                   ] + all_problem_names
     error_msg = "\n  * ".join(error_lines)
     raise LookupError(error_msg)
-  return _PROBLEMS[base_name](was_reversed, was_copy)
+  return _PROBLEMS[base_name](was_reversed=was_reversed, was_copy=was_copy)
 
 
 def list_problems():

From ac56469240d17c7171902c692cb07fb2a8bcdf50 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 18 Oct 2018 14:18:00 -0700
Subject: [PATCH 1041/2720] Hacking infer to enable decode.

PiperOrigin-RevId: 217767343
---
 tensor2tensor/models/video/base.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index cfb6e5b37..5017d5167 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -141,6 +141,12 @@ def get_iteration_num(self):
       step_num = 10000000
     return step_num
 
+  def visualize_predictions(self, predics, targets):
+    predics = tf.concat(predics, axis=1)
+    targets = tf.concat(targets, axis=1)
+    side_by_side_video = tf.concat([predics, targets], axis=2)
+    tf.summary.image("full_video", side_by_side_video)
+
   def get_scheduled_sample_func(self, batch_size):
     """Creates a function for scheduled sampling based on given hparams."""
     with tf.variable_scope("scheduled_sampling_func", reuse=tf.AUTO_REUSE):
@@ -400,6 +406,9 @@ def logits_to_samples(logits):
       for k, v in six.iteritems(logits):
         results[k] = logits_to_samples(v)
         results["%s_logits" % k] = v
+      # HACK: bypassing decoding issues.
+      results["outputs"] = results["targets"]
+      results["scores"] = results["targets"]
     else:
       results = logits_to_samples(logits)
 
@@ -480,9 +489,14 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
       # Cut the predicted input frames.
       res_frames = res_frames[hparams.video_num_input_frames-1:]
       res_rewards = res_rewards[hparams.video_num_input_frames-1:]
+      sampled_frames = sampled_frames[hparams.video_num_input_frames-1:]
+
+    target_frames = all_frames[hparams.video_num_input_frames:]
+    self.visualize_predictions(sampled_frames, target_frames)
 
     output_frames = tf.stack(res_frames, axis=1)
     targets = output_frames
+
     if self.has_rewards:
       output_rewards = tf.stack(res_rewards, axis=1)
       targets = {"targets": output_frames, "target_reward": output_rewards}

From 52d7ef746c9169fb4d5d59b1a9205299d975dbf7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 18 Oct 2018 15:50:25 -0700
Subject: [PATCH 1042/2720] Adds train_eval_and_decode schedule following
 MLPerf requirement.

PiperOrigin-RevId: 217783686
---
 tensor2tensor/utils/trainer_lib.py | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index bc1bf1889..5d72b134a 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -364,6 +364,32 @@ def train(self, max_steps=None):
         hooks=self._train_spec.hooks,
         max_steps=max_steps or self._train_spec.max_steps)
 
+  def train_eval_and_decode(self):
+    """Does eval and decode after training every eval_freq_in_steps."""
+    eval_steps = self._hparams.eval_freq_in_steps
+    packed_dataset = "_packed" in self._hparams.problem.name
+    for i in range(0, self._train_spec.max_steps, eval_steps):
+      if packed_dataset and i > 0:
+        problem = registry.problem(self._hparams.problem.name + "_packed")
+        p_hparams = problem.get_hparams(self._hparams)
+        self._hparams.problem = problem
+        self._hparams.problem_hparams = p_hparams
+      self._estimator.train(
+          self._train_spec.input_fn,
+          steps=eval_steps,
+          hooks=self._train_spec.hooks)
+      self._estimator.evaluate(
+          self._eval_spec.input_fn,
+          steps=self._eval_spec.steps,
+          hooks=self._eval_spec.hooks)
+      if packed_dataset:
+        problem = registry.problem(
+            self._hparams.problem.name.replace("_packed", ""))
+        p_hparams = problem.get_hparams(self._hparams)
+        self._hparams.problem = problem
+        self._hparams.problem_hparams = p_hparams
+      self.decode(dataset_split=tf.estimator.ModeKeys.EVAL)
+
   def evaluate(self):
     return self._estimator.evaluate(
         self._eval_spec.input_fn,
@@ -489,6 +515,7 @@ def create_experiment(
   hparams.add_hparam("schedule", schedule)
   hparams.add_hparam("warm_start_from", warm_start_from)
   hparams.add_hparam("std_server_protocol", std_server_protocol)
+  hparams.add_hparam("eval_freq_in_steps", min_eval_frequency)
   if decode_hparams is not None:
     decode_hparams.add_hparam("decode_from_file", decode_from_file)
     decode_hparams.add_hparam("decode_to_file", decode_to_file)

From af24dcde8ada770cf233dd72ede478c52df769b6 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Thu, 18 Oct 2018 23:59:09 -0700
Subject: [PATCH 1043/2720] fixes to EPVA model.

PiperOrigin-RevId: 217827916
---
 tensor2tensor/models/video/epva.py        | 9 +++++++++
 tensor2tensor/models/video/epva_params.py | 1 +
 2 files changed, 10 insertions(+)

diff --git a/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index 53c33cee2..0bd153c55 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -711,6 +711,15 @@ def fix_video_dims_and_concat_on_x_axis(x):
     side_by_side_video = tf.concat([frames_gd, frames_pd], axis=1)
     tf.summary.image('full_video', side_by_side_video)
 
+    predictions = tf.unstack(predictions)
+    predictions = [
+        tf.image.resize_images(
+            image, (frame_width, frame_height),
+            method=tf.image.ResizeMethod.BICUBIC)
+        for image in predictions
+    ]
+    predictions = tf.stack(predictions)
+
     predictions = common_video.swap_time_and_batch_axes(predictions)
     predictions = tf.slice(predictions,
                            [0, hparams.video_num_input_frames-1, 0, 0, 0],
diff --git a/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
index 780d34cc2..57f20a0b9 100644
--- a/tensor2tensor/models/video/epva_params.py
+++ b/tensor2tensor/models/video/epva_params.py
@@ -30,6 +30,7 @@ def next_frame_epva():
   hparams.video_num_target_frames = 4
   hparams.target_modality = "video:l2raw"
   hparams.input_modalities = "inputs:video:l2raw"
+  hparams.learning_rate_schedule = "constant"
   hparams.learning_rate_constant = 1e-05
   hparams.batch_size = 2
   hparams.clip_grad_norm = 0.01

From f1d70701d99ffbe5900c0528e19a8916365aa6cb Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 19 Oct 2018 11:24:51 -0700
Subject: [PATCH 1044/2720] fixing deterministic SV2P.

PiperOrigin-RevId: 217901194
---
 tensor2tensor/models/video/sv2p.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index b5c01c4ac..ed20e754d 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -335,6 +335,8 @@ def video_features(
       self, all_frames, all_actions, all_rewards, all_raw_frames):
     """Video wide latent."""
     del all_actions, all_rewards, all_raw_frames
+    if not self.hparams.stochastic_model:
+      return None, None, None
     frames = tf.stack(all_frames, axis=1)
     mean, std = self.construct_latent_tower(frames, time_axis=1)
     latent = common_video.get_gaussian_tensor(mean, std)
@@ -349,7 +351,8 @@ def next_frame(self, frames, actions, rewards, target_frame,
     extra_loss = 0.0
     if internal_states is None:
       internal_states = [None] * (5 if self.hparams.small_mode else 7)
-      extra_loss = self.get_extra_loss([latent_mean], [latent_std])
+      if latent_mean is not None:
+        extra_loss = self.get_extra_loss([latent_mean], [latent_std])
 
     pred_image, internal_states = self.construct_predictive_tower(
         frames, None, actions, internal_states, latent)

From 5b6a310f40df315cd9dfb242f308b2fe8e86ae67 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 19 Oct 2018 11:30:18 -0700
Subject: [PATCH 1045/2720] Add mlperf compliance log to transformer.

PiperOrigin-RevId: 217902485
---
 tensor2tensor/bin/t2t_trainer.py           |   4 +
 tensor2tensor/data_generators/problem.py   |   2 +
 tensor2tensor/data_generators/tokenizer.py |   3 +
 tensor2tensor/data_generators/translate.py |   2 +
 tensor2tensor/models/transformer.py        |  56 ++++
 tensor2tensor/utils/mlperf_log.py          | 170 +++++++++++
 tensor2tensor/utils/mlperf_tags.py         | 336 +++++++++++++++++++++
 tensor2tensor/utils/optimize.py            |  10 +
 tensor2tensor/utils/trainer_lib.py         |  10 +
 9 files changed, 593 insertions(+)
 create mode 100644 tensor2tensor/utils/mlperf_log.py
 create mode 100644 tensor2tensor/utils/mlperf_tags.py

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 9b6f0b25e..0f557d560 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -27,6 +27,7 @@
 from tensor2tensor.utils import cloud_mlengine
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
+from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
@@ -343,6 +344,7 @@ def run_std_server():
 
 def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
+  mlperf_log.transformer_print(key=mlperf_log.RUN_START)
   if FLAGS.schedule == "run_std_server":
     run_std_server()
   trainer_lib.set_random_seed(FLAGS.random_seed)
@@ -368,6 +370,8 @@ def main(argv):
   if is_chief():
     save_metadata(hparams)
   execute_schedule(exp)
+  mlperf_log.transformer_print(key=mlperf_log.RUN_STOP)
+  mlperf_log.transformer_print(key=mlperf_log.RUN_FINAL)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index a97d6ed12..f86eb6ffe 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -31,6 +31,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import metrics
+from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
 from tensorflow.contrib.tpu.python.tpu import tpu_config
@@ -638,6 +639,7 @@ def _load_records_and_preprocess(filenames):
     tf.logging.info(
         "partition: %d num_data_files: %d" % (partition_id, len(data_files)))
     if shuffle_files:
+      mlperf_log.transformer_print(key=mlperf_log.INPUT_ORDER)
       random.shuffle(data_files)
 
     dataset = tf.data.Dataset.from_tensor_slices(tf.constant(data_files))
diff --git a/tensor2tensor/data_generators/tokenizer.py b/tensor2tensor/data_generators/tokenizer.py
index 6a05d5e05..d5449748e 100644
--- a/tensor2tensor/data_generators/tokenizer.py
+++ b/tensor2tensor/data_generators/tokenizer.py
@@ -49,6 +49,7 @@
 import unicodedata
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
+from tensor2tensor.utils import mlperf_log
 import tensorflow as tf
 
 # Conversion between Unicode and UTF-8, if required (on Python2)
@@ -165,6 +166,8 @@ def corpus_token_counts(
       split_on_newlines=split_on_newlines):
     counts.update(encode(_native_to_unicode(doc)))
 
+  mlperf_log.transformer_print(
+      key=mlperf_log.PREPROC_VOCAB_SIZE, value=len(counts))
   return counts
 
 
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 23f2ae76b..0ad29761d 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -26,6 +26,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import bleu_hook
+from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
 
@@ -87,6 +88,7 @@ def compute_bleu_summaries(hook_args):
       decode_hparams.decode_reference, decode_hparams.decode_to_file)
   values.append(tf.Summary.Value(tag="BLEU", simple_value=bleu))
   tf.logging.info("%s: BLEU = %6.2f" % (decode_hparams.decode_to_file, bleu))
+  mlperf_log.transformer_print(key=mlperf_log.EVAL_ACCURACY, value=bleu)
   return values
 
 
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 836f1812a..4120e276b 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -34,6 +34,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import beam_search
 from tensor2tensor.utils import expert_utils
+from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -76,6 +77,10 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
         transformer_prepare_encoder(
             inputs, target_space, hparams, features=features))
 
+    mlperf_log.transformer_print(
+        key=mlperf_log.MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
+        value=hparams.layer_prepostprocess_dropout)
+
     encoder_input = tf.nn.dropout(encoder_input,
                                   1.0 - hparams.layer_prepostprocess_dropout)
 
@@ -121,6 +126,9 @@ def decode(self,
     Returns:
       Final decoder representation. [batch_size, decoder_length, hidden_dim]
     """
+    mlperf_log.transformer_print(
+        key=mlperf_log.MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
+        value=hparams.layer_prepostprocess_dropout)
     decoder_input = tf.nn.dropout(decoder_input,
                                   1.0 - hparams.layer_prepostprocess_dropout)
 
@@ -778,6 +786,15 @@ def fast_decode_tpu(encoder_output,
     cache["encoder_output"] = encoder_output
     cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias
 
+  mlperf_log.transformer_print(
+      key=mlperf_log.MODEL_HP_SEQ_BEAM_SEARCH,
+      value={
+          "vocab_size": vocab_size,
+          "batch_size": batch_size,
+          "beam_size": beam_size,
+          "alpha": alpha,
+          "max_decode_length": decode_length
+      })
   if beam_size > 1:  # Beam Search
     initial_ids = sos_id * tf.ones([batch_size], dtype=tf.int32)
     decoded_ids, scores = beam_search.beam_search(
@@ -1253,6 +1270,14 @@ def transformer_encoder(encoder_input,
   attention_dropout_broadcast_dims = (
       common_layers.comma_separated_string_to_integer_list(
           getattr(hparams, "attention_dropout_broadcast_dims", "")))
+  mlperf_log.transformer_print(
+      key=mlperf_log.MODEL_HP_NUM_HIDDEN_LAYERS,
+      value=hparams.num_hidden_layers)
+  mlperf_log.transformer_print(
+      key=mlperf_log.MODEL_HP_ATTENTION_NUM_HEADS, value=hparams.num_heads)
+  mlperf_log.transformer_print(
+      key=mlperf_log.MODEL_HP_ATTENTION_DROPOUT,
+      value=hparams.attention_dropout)
   with tf.variable_scope(name):
     if nonpadding is not None:
       padding = 1.0 - nonpadding
@@ -1298,6 +1323,9 @@ def transformer_encoder(encoder_input,
     # if normalization is done in layer_preprocess, then it should also be done
     # on the output, since the output can grow very large, being the sum of
     # a whole stack of unnormalized layer outputs.
+    mlperf_log.transformer_print(
+        key=mlperf_log.MODEL_HP_NORM,
+        value={"hidden_size": hparams.hidden_size})
     return common_layers.layer_preprocess(x, hparams)
 
 
@@ -1346,6 +1374,16 @@ def transformer_decoder(decoder_input,
   attention_dropout_broadcast_dims = (
       common_layers.comma_separated_string_to_integer_list(
           getattr(hparams, "attention_dropout_broadcast_dims", "")))
+
+  mlperf_log.transformer_print(
+      key=mlperf_log.MODEL_HP_NUM_HIDDEN_LAYERS,
+      value=hparams.num_hidden_layers)
+  mlperf_log.transformer_print(
+      key=mlperf_log.MODEL_HP_ATTENTION_NUM_HEADS, value=hparams.num_heads)
+  mlperf_log.transformer_print(
+      key=mlperf_log.MODEL_HP_ATTENTION_DROPOUT,
+      value=hparams.attention_dropout)
+
   with tf.variable_scope(name):
     for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
       layer_name = "layer_%d" % layer
@@ -1409,6 +1447,9 @@ def transformer_decoder(decoder_input,
     # if normalization is done in layer_preprocess, then it should also be done
     # on the output, since the output can grow very large, being the sum of
     # a whole stack of unnormalized layer outputs.
+    mlperf_log.transformer_print(
+        key=mlperf_log.MODEL_HP_NORM,
+        value={"hidden_size": hparams.hidden_size})
     return common_layers.layer_preprocess(x, hparams)
 
 
@@ -1458,6 +1499,21 @@ def transformer_ffn_layer(x,
     ffn_layer = "dense_relu_dense"
   if ffn_layer == "dense_relu_dense":
     # In simple convolution mode, use `pad_remover` to speed up processing.
+    mlperf_log.transformer_print(
+        key=mlperf_log.MODEL_HP_FFN_FILTER,
+        value={
+            "filter_size": hparams.filter_size,
+            "use_bias": "True",
+            "activation": mlperf_log.RELU
+        })
+    mlperf_log.transformer_print(
+        key=mlperf_log.MODEL_HP_FFN_DENSE,
+        value={
+            "hidden_size": hparams.hidden_size,
+            "use_bias": "True",
+        })
+    mlperf_log.transformer_print(
+        key=mlperf_log.MODEL_HP_RELU_DROPOUT, value=hparams.relu_dropout)
     if pad_remover:
       original_shape = common_layers.shape_list(x)
       # Collapse `x` across examples, and remove padding positions.
diff --git a/tensor2tensor/utils/mlperf_log.py b/tensor2tensor/utils/mlperf_log.py
new file mode 100644
index 000000000..90de37c50
--- /dev/null
+++ b/tensor2tensor/utils/mlperf_log.py
@@ -0,0 +1,170 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+# Copyright 2018 MLBenchmark Group. 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.
+# ==============================================================================
+"""Convenience function for logging compliance tags to stdout.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import inspect
+import json
+import logging
+import os
+import re
+import sys
+import time
+import uuid
+
+from tensor2tensor.utils.mlperf_tags import *  # pylint: disable=wildcard-import
+
+
+ROOT_DIR_GNMT = None
+
+# Set by imagenet_main.py
+ROOT_DIR_RESNET = None
+
+# Set by transformer_main.py and process_data.py
+ROOT_DIR_TRANSFORMER = None
+
+
+PATTERN = re.compile("[a-zA-Z0-9]+")
+
+LOG_FILE = os.getenv("COMPLIANCE_FILE")
+# create logger with 'mlperf_compliance'
+LOGGER = logging.getLogger("mlperf_compliance")
+LOGGER.setLevel(logging.DEBUG)
+
+_STREAM_HANDLER = logging.StreamHandler(stream=sys.stdout)
+_STREAM_HANDLER.setLevel(logging.INFO)
+LOGGER.addHandler(_STREAM_HANDLER)
+
+if LOG_FILE:
+  _FILE_HANDLER = logging.FileHandler(LOG_FILE)
+  _FILE_HANDLER.setLevel(logging.DEBUG)
+  LOGGER.addHandler(_FILE_HANDLER)
+else:
+  _STREAM_HANDLER.setLevel(logging.DEBUG)
+
+
+def get_caller(stack_index=2, root_dir=None):
+  # pylint: disable=g-doc-args
+  """Returns file.py:lineno of your caller.
+
+  A stack_index of 2 will provide
+      the caller of the function calling this function. Notice that stack_index
+      of 2 or more will fail if called from global scope.
+  """
+  caller = inspect.getframeinfo(inspect.stack()[stack_index][0])
+
+  # Trim the filenames for readability.
+  filename = caller.filename
+  if root_dir is not None:
+    filename = re.sub("^" + root_dir + "/", "", filename)
+  return "%s:%d" % (filename, caller.lineno)
+
+
+def _mlperf_print(key, value=None, benchmark=None, stack_offset=0,
+                  tag_set=None, deferred=False, root_dir=None,
+                  extra_print=False):
+  # pylint: disable=g-doc-args
+  # pylint: disable=g-doc-return-or-yield
+  """Prints out an MLPerf Log Line.
+
+  key: The MLPerf log key such as 'CLOCK' or 'QUALITY'. See the list of log keys
+  in the spec.
+  value: The value which contains no newlines.
+  benchmark: The short code for the benchmark being run, see the MLPerf log
+  spec.
+  stack_offset: Increase the value to go deeper into the stack to find the
+  callsite. For example, if this
+                is being called by a wraper/helper you may want to set
+                stack_offset=1 to use the callsite
+                of the wraper/helper itself.
+  tag_set: The set of tags in which key must belong.
+  deferred: The value is not presently known. In that case, a unique ID will
+            be assigned as the value of this call and will be returned. The
+            caller can then include said unique ID when the value is known
+            later.
+  root_dir: Directory prefix which will be trimmed when reporting calling file
+            for compliance logging.
+  extra_print: Print a blank line before logging to clear any text in the line.
+
+  Example output:
+    :::MLP-1537375353 MINGO[17] (eval.py:42) QUALITY: 43.7
+  """
+
+  return_value = None
+
+  if (tag_set is None and not PATTERN.match(key)) or key not in tag_set:
+    raise ValueError("Invalid key for MLPerf print: " + str(key))
+
+  if value is not None and deferred:
+    raise ValueError("deferred is set to True, but a value was provided")
+
+  if deferred:
+    return_value = str(uuid.uuid4())
+    value = "DEFERRED: {}".format(return_value)
+
+  if value is None:
+    tag = key
+  else:
+    str_json = json.dumps(value)
+    tag = "{key}: {value}".format(key=key, value=str_json)
+
+  callsite = get_caller(2 + stack_offset, root_dir=root_dir)
+  now = time.time()
+
+  message = ":::MLPv0.5.0 {benchmark} {secs:.9f} ({callsite}) {tag}".format(
+      secs=now, benchmark=benchmark, callsite=callsite, tag=tag)
+
+  if extra_print:
+    print()  # There could be prior text on a line
+
+  if tag in STDOUT_TAG_SET:  # pylint: disable=undefined-variable
+    LOGGER.info(message)
+  else:
+    LOGGER.debug(message)
+
+  return return_value
+
+
+TRANSFORMER_TAG_SET = set(TRANSFORMER_TAGS)  # pylint: disable=undefined-variable
+
+
+def transformer_print(key, value=None, stack_offset=2, deferred=False):
+  return _mlperf_print(
+      key=key,
+      value=value,
+      benchmark=TRANSFORMER,  # pylint: disable=undefined-variable
+      stack_offset=stack_offset,
+      tag_set=TRANSFORMER_TAG_SET,
+      deferred=deferred,
+      root_dir=ROOT_DIR_TRANSFORMER)
diff --git a/tensor2tensor/utils/mlperf_tags.py b/tensor2tensor/utils/mlperf_tags.py
new file mode 100644
index 000000000..36c290293
--- /dev/null
+++ b/tensor2tensor/utils/mlperf_tags.py
@@ -0,0 +1,336 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+# Copyright 2018 MLBenchmark Group. 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.
+# ==============================================================================
+"""Master list of MLPerf tags to be logged for benchmark submissions.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# ==============================================================================
+# == Benchmarks ================================================================
+# ==============================================================================
+
+# translation/
+TRANSFORMER = "transformer"
+INPUT_MAX_LENGTH = "input_max_length"
+
+MODEL_HP_INITIALIZER_GAIN = "model_hp_initializer_gain"
+MODEL_HP_VOCAB_SIZE = "model_hp_vocab_size"
+MODEL_HP_NUM_HIDDEN_LAYERS = "model_hp_hidden_layers"
+MODEL_HP_ATTENTION_DENSE = "model_hp_attention_dense"
+MODEL_HP_ATTENTION_NUM_HEADS = "model_hp_attention_num_heads"
+MODEL_HP_ATTENTION_DROPOUT = "model_hp_attention_dropout"
+MODEL_HP_FFN_DENSE = "model_hp_ffn_dense"
+MODEL_HP_FFN_FILTER = "model_hp_ffn_filter"
+MODEL_HP_RELU_DROPOUT = "model_hp_relu_dropout"
+MODEL_HP_LAYER_POSTPROCESS_DROPOUT = "model_hp_layer_postprocess_dropout"
+MODEL_HP_NORM = "model_hp_norm"
+MODEL_HP_SEQ_BEAM_SEARCH = "model_hp_seq_beam_search"
+
+# ==============================================================================
+# == Tags ======================================================================
+# ==============================================================================
+"""
+Tags may be used by all models, a subset of models, or only one model. A
+specification for which models require which tags can be found below the tag
+definitions.
+"""
+
+# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
+# All models: Tags which should appear in absolutely every MLPerf model.
+# //////////////////////////////////////////////////////////////////////////////
+
+# This tag signals to start the timer. Emission of this tag need not be (and
+# generally will not be) the first part of a submission script. Rather, this
+# tag must be emitted prior to performing any work which the MLPerf rules
+# state must be timed. This tag is generally emitted directly before the first
+# step which invokes random number generation or the first step which must be
+# performed on the system under test. (Whichever comes first.) If clarification
+# is needed, please file an issue under:
+#   https://github.com/mlperf/policies
+RUN_START = "run_start"
+
+# This tag signals that a submission has reached the relevant stopping criteria,
+# and has completed all tasks which are performed in the reference. The wall
+# time for a submission will be computed as the difference between the time
+# when this tag is emitted and the time whe the RUN_START is emitted.
+RUN_STOP = "run_stop"
+
+# This tag should be emitted immediately before ending a run, and should be the
+# last tag emitted. This tag should indicate the completion of untimed post
+# processing work such as system specific cleanup.
+RUN_FINAL = "run_final"
+
+
+# Emit this tag in the place(s) where random seeds are set.
+RUN_SET_RANDOM_SEED = "run_set_random_seed"
+
+
+# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
+# Common Values: Constants which are expected to be reported across many models.
+#                These values are included for convenience.
+# //////////////////////////////////////////////////////////////////////////////
+BCE = "binary_cross_entropy"
+CCE = "categorical_cross_entropy"
+
+SGD = "stochastic_gradient_descent"
+
+# Some conventions distinguish between "vanilla" SGD and SGD with momentum
+# (where vanilla SGD would be the specific case of momentum=0)
+SGD_WITH_MOMENTUM = "stochastic_gradient_descent_with_momentum"
+
+ADAM = "adam"
+LAZY_ADAM = "lazy_adam"
+
+TRUNCATED_NORMAL = "truncated_normal"
+
+RELU = "relu"
+
+
+# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
+# Preprocessing: Tags for generic preprocessing steps
+# //////////////////////////////////////////////////////////////////////////////
+
+# The number of training examples in a single epoch
+PREPROC_NUM_TRAIN_EXAMPLES = "preproc_num_train_examples"
+
+# The number of evaluation examples in a single epoch
+PREPROC_NUM_EVAL_EXAMPLES = "preproc_num_eval_examples"
+
+# This tag is used to declare what part of code tokenizes the training data.
+PREPROC_TOKENIZE_TRAINING = "preproc_tokenize_training"
+
+# This tag is used to declare what part of code tokenizes the evaluation data.
+PREPROC_TOKENIZE_EVAL = "preproc_tokenize_eval"
+
+# The vocabulary size used for tokenization.
+PREPROC_VOCAB_SIZE = "preproc_vocab_size"
+
+
+# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
+# Input: Tags for the timed portion of the data input pipeline
+# //////////////////////////////////////////////////////////////////////////////
+
+# The number of examples in the training portion of the data pipeline. Generally
+# this should match PREPROC_NUM_TRAIN_EXAMPLES. If it does not (for instance
+# if certain examples are dropped in compliance with MLPerf rules), the
+# call which declares this tag is a good place for a comment stating why the
+# disparity is expected.
+INPUT_SIZE = "input_size"
+
+# The size of a training minibatch size. If this value is variable, please emit
+# "-1" and then log an implementation specific characterization of the batch
+# size which is a reasonable analog to the reference. (For instance log that
+# all but the last batch has size 64, and the last batch is a partial batch)
+INPUT_BATCH_SIZE = "input_batch_size"
+
+# This tag indicates where the location of the code which defines the order in
+# which training examples are traversed. It is not necessary to describe the
+# method in the tag emission (though comments are always welcome). Rather, this
+# should simply provide a good starting point to an interested party.
+INPUT_ORDER = "input_order"
+
+
+# --------------------------------------
+# -- Data Augmentation and Alteration --
+# --------------------------------------
+
+# ResNet random cropping
+INPUT_CENTRAL_CROP = "input_central_crop"
+
+INPUT_DISTORTED_CROP_MIN_OBJ_COV = "input_distorted_crop_min_object_covered"
+INPUT_DISTORTED_CROP_RATIO_RANGE = "input_distorted_crop_aspect_ratio_range"
+INPUT_DISTORTED_CROP_AREA_RANGE = "input_distorted_crop_area_range"
+INPUT_DISTORTED_CROP_MAX_ATTEMPTS = "input_distorted_crop_max_attempts"
+
+INPUT_MEAN_SUBTRACTION = "input_mean_subtraction"
+
+# Random flip of an image for data augmentation
+INPUT_RANDOM_FLIP = "input_random_flip"
+
+INPUT_RESIZE = "input_resize"
+INPUT_RESIZE_ASPECT_PRESERVING = "input_resize_aspect_preserving"
+
+
+# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
+# Opt: Tags for declaring optimizer specific information. Submissions should
+#      declare and log explicit values rather than relying on defaults.
+# //////////////////////////////////////////////////////////////////////////////
+
+# The name of the optimizer used. (SGD, Adam, etc.)
+OPT_NAME = "opt_name"
+
+OPT_LR = "opt_learning_rate"
+OPT_MOMENTUM = "opt_momentum"
+
+OPT_WEIGHT_DECAY = "opt_weight_decay"
+
+# beta1, beta2, and epsilon are optimizer hyperparameters associated with the
+# Adam optimizer and its variants (e.g. LazyAdam).
+OPT_HP_ADAM_BETA1 = "opt_hp_Adam_beta1"
+OPT_HP_ADAM_BETA2 = "opt_hp_Adam_beta2"
+OPT_HP_ADAM_EPSILON = "opt_hp_Adam_epsilon"
+
+
+# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
+#  Train: Tags for control flow during model training.
+# //////////////////////////////////////////////////////////////////////////////
+
+# This tag is emitted when a model first enters its training loop. This is not
+# necessarily when it begins to apply gradients; rather, it should be placed at
+# a location which logically partitions the submission code.
+TRAIN_LOOP = "train_loop"
+
+# The current epoch as said epoch begins training.
+TRAIN_EPOCH = "train_epoch"
+
+# This tag is used to indicate approximately where checkpoints are written. Some
+# frameworks abstract away checkpoint saving; in such cases simply choose a
+# logical place in the code which signals that the framework has been instructed
+# to save checkpoints, along with an explanatory comment.
+TRAIN_CHECKPOINT = "train_checkpoint"
+
+
+# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
+#  Eval: Tags for control flow during model evaluation.
+# //////////////////////////////////////////////////////////////////////////////
+
+# This tag should be emitted whenever the submission begins an evaluation pass
+# for a given set of weights.
+EVAL_START = "eval_start"
+
+# The number of examples on which evaluation is performed.
+EVAL_SIZE = "eval_size"
+
+# The target quality at which the model may stop training.
+EVAL_TARGET = "eval_target"
+
+# The observed accuracy of the model at a given epoch.
+EVAL_ACCURACY = "eval_accuracy"
+
+# This tag should be emitted when the model has determined that it has met the
+# target quality set by the reference.
+EVAL_STOP = "eval_stop"
+
+
+# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
+#  Model: Tags for logging topology specific information.
+# //////////////////////////////////////////////////////////////////////////////
+
+# The loss function (cross entropy, squared error, etc.) used by the model. For
+# more exotic loss functions such as those encountered in object detection
+# models, additional benchmark specific subcomponents should also be logged.
+MODEL_HP_LOSS_FN = "model_hp_loss_fn"
+
+MODEL_HP_INITIAL_SHAPE = "model_hp_initial_shape"
+MODEL_HP_FINAL_SHAPE = "model_hp_final_shape"
+
+MODEL_L2_REGULARIZATION = "model_l2_regularization"
+MODEL_EXCLUDE_BN_FROM_L2 = "model_exclude_bn_from_l2"
+
+MODEL_HP_RELU = "model_hp_relu"
+MODEL_HP_CONV2D_FIXED_PADDING = "model_hp_conv2d_fixed_padding"
+MODEL_HP_BATCH_NORM = "model_hp_batch_norm"
+MODEL_HP_DENSE = "model_hp_dense"
+
+
+# ==============================================================================
+# == Stdout tags ===============================================================
+# ==============================================================================
+
+# These tags are always logged to stdout. The rest will be logged to a file if
+# one is available.
+STDOUT_TAG_SET = {
+    RUN_START,
+    RUN_STOP,
+    RUN_FINAL,
+
+    TRAIN_LOOP,
+    TRAIN_EPOCH,
+
+    EVAL_START,
+    EVAL_SIZE,
+    EVAL_TARGET,
+    EVAL_ACCURACY,
+    EVAL_STOP,
+}
+
+
+# ==============================================================================
+# == Benchmark tag sets ========================================================
+# ==============================================================================
+ALL_USED_TAGS = set()
+
+TRANSFORMER_TAGS = (
+    RUN_START,
+    RUN_STOP,
+    RUN_FINAL,
+    RUN_SET_RANDOM_SEED,
+
+    PREPROC_NUM_TRAIN_EXAMPLES,
+    PREPROC_NUM_EVAL_EXAMPLES,
+    PREPROC_TOKENIZE_TRAINING,
+    PREPROC_TOKENIZE_EVAL,
+    PREPROC_VOCAB_SIZE,
+
+    INPUT_BATCH_SIZE,
+    INPUT_MAX_LENGTH,
+    INPUT_ORDER,
+
+    OPT_NAME,
+    OPT_LR,
+    OPT_HP_ADAM_BETA1,
+    OPT_HP_ADAM_BETA2,
+    OPT_HP_ADAM_EPSILON,
+
+    TRAIN_LOOP,
+    TRAIN_EPOCH,
+
+    EVAL_START,
+    EVAL_TARGET,
+    EVAL_ACCURACY,
+    EVAL_STOP,
+
+    MODEL_HP_INITIALIZER_GAIN,
+    MODEL_HP_VOCAB_SIZE,
+    MODEL_HP_NUM_HIDDEN_LAYERS,
+    MODEL_HP_ATTENTION_DENSE,
+    MODEL_HP_ATTENTION_NUM_HEADS,
+    MODEL_HP_ATTENTION_DROPOUT,
+    MODEL_HP_FFN_DENSE,
+    MODEL_HP_FFN_FILTER,
+    MODEL_HP_RELU_DROPOUT,
+    MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
+    MODEL_HP_NORM,
+    MODEL_HP_SEQ_BEAM_SEARCH,
+)
+
+ALL_USED_TAGS.update(TRANSFORMER_TAGS)
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 35a6a7dd3..47b3ad983 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -21,6 +21,7 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import adafactor
+from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import multistep_optimizer
 from tensor2tensor.utils import yellowfin
 
@@ -79,6 +80,15 @@ class ConditionalOptimizer(tf.train.Optimizer):
   def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disable=super-init-not-called
     tf.logging.info("Using optimizer %s", optimizer_name)
 
+    mlperf_log.transformer_print(key=mlperf_log.OPT_NAME, value=optimizer_name)
+    mlperf_log.transformer_print(
+        key=mlperf_log.OPT_HP_ADAM_BETA1, value=hparams.optimizer_adam_beta1)
+    mlperf_log.transformer_print(
+        key=mlperf_log.OPT_HP_ADAM_BETA2, value=hparams.optimizer_adam_beta2)
+    mlperf_log.transformer_print(
+        key=mlperf_log.OPT_HP_ADAM_EPSILON,
+        value=hparams.optimizer_adam_epsilon)
+
     if optimizer_name == "Adam":
       # We change the default epsilon for Adam.
       # Using LazyAdam as it's much faster for large vocabulary embeddings.
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 5d72b134a..505bfaf94 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -28,6 +28,7 @@
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import devices
 from tensor2tensor.utils import metrics_hook
+from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -259,6 +260,8 @@ def create_estimator(model_name,
     batch_size = (
         problem.tpu_batch_size_per_shard(hparams) *
         run_config.tpu_config.num_shards)
+    mlperf_log.transformer_print(
+        key=mlperf_log.INPUT_BATCH_SIZE, value=batch_size)
     if getattr(hparams, "mtf_mode", False):
       batch_size = problem.tpu_batch_size_per_shard(hparams)
     predict_batch_size = batch_size
@@ -359,6 +362,7 @@ def train_and_evaluate(self):
       self.train()
 
   def train(self, max_steps=None):
+    mlperf_log.transformer_print(key=mlperf_log.TRAIN_LOOP)
     self._estimator.train(
         self._train_spec.input_fn,
         hooks=self._train_spec.hooks,
@@ -368,7 +372,10 @@ def train_eval_and_decode(self):
     """Does eval and decode after training every eval_freq_in_steps."""
     eval_steps = self._hparams.eval_freq_in_steps
     packed_dataset = "_packed" in self._hparams.problem.name
+    mlperf_log.transformer_print(key=mlperf_log.TRAIN_LOOP)
     for i in range(0, self._train_spec.max_steps, eval_steps):
+      mlperf_log.transformer_print(
+          key=mlperf_log.TRAIN_EPOCH, value=i // eval_steps)
       if packed_dataset and i > 0:
         problem = registry.problem(self._hparams.problem.name + "_packed")
         p_hparams = problem.get_hparams(self._hparams)
@@ -388,7 +395,10 @@ def train_eval_and_decode(self):
         p_hparams = problem.get_hparams(self._hparams)
         self._hparams.problem = problem
         self._hparams.problem_hparams = p_hparams
+      mlperf_log.transformer_print(key=mlperf_log.EVAL_START)
       self.decode(dataset_split=tf.estimator.ModeKeys.EVAL)
+      mlperf_log.transformer_print(key=mlperf_log.EVAL_TARGET, value=25.0)
+      mlperf_log.transformer_print(key=mlperf_log.EVAL_STOP)
 
   def evaluate(self):
     return self._estimator.evaluate(

From 81bffa81643bbe4e7bf7864211383378824daec0 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Fri, 19 Oct 2018 22:57:11 +0200
Subject: [PATCH 1046/2720] T2TEnv: Full history cleaning on new epoch (fix);
 read only rollouts with size appropriate for initial_frame_chooser;
 initial_frame_chooser fix for simulation_random_starts case (unused so far);
 (#1160)

---
 tensor2tensor/data_generators/gym_env.py      | 33 +++++++++--------
 tensor2tensor/data_generators/gym_env_test.py | 35 +++++++++++--------
 tensor2tensor/rl/trainer_model_based.py       | 21 ++++++-----
 3 files changed, 49 insertions(+), 40 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 439d7158b..e4a094c87 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -87,11 +87,7 @@ class T2TEnv(video_utils.VideoProblem):
   def __init__(self, batch_size, *args, **kwargs):
     super(T2TEnv, self).__init__(*args, **kwargs)
 
-    self.clear_history()
     self.batch_size = batch_size
-    self._current_batch_frames = [None for _ in range(batch_size)]
-    self._current_batch_rollouts = [[] for _ in range(batch_size)]
-    self._current_epoch_rollouts = []
     self._rollouts_by_epoch_and_split = collections.OrderedDict()
     self.current_epoch = None
     with tf.Graph().as_default() as tf_graph:
@@ -108,10 +104,6 @@ def __str__(self):
     """Returns a string representation of the environment for debug purposes."""
     raise NotImplementedError
 
-  def clear_history(self):
-    """Clears the rollout history."""
-    self._rollouts_by_epoch_and_split = collections.OrderedDict()
-
   def start_new_epoch(self, epoch, load_data_dir=None):
     if not isinstance(epoch, int):
       raise ValueError("Epoch should be integer, got {}".format(epoch))
@@ -120,10 +112,14 @@ def start_new_epoch(self, epoch, load_data_dir=None):
     self.current_epoch = epoch
     self._current_epoch_rollouts = []
     self._rollouts_by_epoch_and_split[epoch] = collections.defaultdict(list)
+    self._current_batch_frames = [None for _ in range(self.batch_size)]
+    self._current_batch_rollouts = [[] for _ in range(self.batch_size)]
     if load_data_dir is not None:
       self._load_epoch_data(load_data_dir)
 
-  def current_epoch_rollouts(self, split=None):
+  def current_epoch_rollouts(self, split=None, minimal_rollout_frames=0):
+    # TODO(KC): order of rollouts (by splits) is a bit uncontrolled
+    # (rollouts_by_split.values() reads dict values), is it a problem?
     rollouts_by_split = self._rollouts_by_epoch_and_split[self.current_epoch]
     if not rollouts_by_split:
       if split is not None:
@@ -131,15 +127,18 @@ def current_epoch_rollouts(self, split=None):
             "generate_data() should first be called in the current epoch"
         )
       else:
-        return self._current_epoch_rollouts
-    if split is not None:
-      return rollouts_by_split[split]
+        rollouts = self._current_epoch_rollouts
     else:
-      return [
-          rollout
-          for rollouts in rollouts_by_split.values()
-          for rollout in rollouts
-      ]
+      if split is not None:
+        rollouts = rollouts_by_split[split]
+      else:
+        rollouts = [
+            rollout
+            for rollouts in rollouts_by_split.values()
+            for rollout in rollouts
+        ]
+    return [rollout for rollout in rollouts
+            if len(rollout) >= minimal_rollout_frames]
 
   def _preprocess_observations(self, obs):
     """Transforms a batch of observations.
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 15536deb9..0326eb983 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -49,9 +49,8 @@ def __init__(self):
     self._counter = 0
 
   def _generate_ob(self):
-    return np.zeros(
-        self.observation_space.shape, self.observation_space.dtype
-    )
+    return np.random.randint(255, size=self.observation_space.shape,
+                             dtype=self.observation_space.dtype)
 
   def step(self, action):
     done = self._counter % 2 == 1
@@ -79,10 +78,11 @@ def setUp(self):
     self.out_dir = tf.test.get_temp_dir()
     shutil.rmtree(self.out_dir)
     os.mkdir(self.out_dir)
+    np.random.seed(0)
 
-  def init_batch_and_play(self, env_name, steps_per_epoch=1,
-                          epochs=(0,), generate_data=False, **kwargs):
-    env = gym_env.T2TGymEnv(env_name, batch_size=2, **kwargs)
+  def init_batch_and_play(self, env_name, steps_per_epoch=1, epochs=(0,),
+                          generate_data=False, batch_size=2, **kwargs):
+    env = gym_env.T2TGymEnv(env_name, batch_size=batch_size, **kwargs)
     obs = list()
     rewards = list()
     num_dones = 0
@@ -90,6 +90,7 @@ def init_batch_and_play(self, env_name, steps_per_epoch=1,
       env.start_new_epoch(epoch, self.out_dir)
       _, epoch_obs, epoch_rewards, epoch_num_dones = \
           self.play(env, steps_per_epoch)
+      epoch_obs.append(env.reset())
       if generate_data:
         env.generate_data(self.out_dir)
       obs.extend(epoch_obs)
@@ -121,21 +122,26 @@ def test_splits_dataset(self):
       self.assertTrue(env.current_epoch_rollouts(split))
 
   def test_split_preserves_number_of_rollouts(self):
+    batch_size = 2
     env, _, _, num_dones = self.init_batch_and_play(
-        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True,
+        batch_size=batch_size
     )
 
     num_rollouts_after_split = sum(
         len(env.current_epoch_rollouts(split)) for split in self.splits
     )
-    # Number of rollouts could be increased by one in case a rollout is broken
-    # on a boundary between the dataset splits.
-    self.assertGreaterEqual(num_rollouts_after_split, num_dones)
-    self.assertLessEqual(num_rollouts_after_split, num_dones + 1)
+    # After the end of epoch all environments are reset, which increases number
+    # of rollouts by batch size. Number of rollouts could be increased by one
+    # in case a rollout is broken on a boundary between the dataset splits.
+    self.assertGreaterEqual(num_rollouts_after_split, num_dones + batch_size)
+    self.assertLessEqual(num_rollouts_after_split, num_dones + batch_size + 1)
 
   def test_split_preserves_number_of_frames(self):
+    batch_size = 2
     env, _, _, num_dones = self.init_batch_and_play(
-        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True,
+        batch_size=batch_size
     )
 
     num_frames = sum(
@@ -144,8 +150,9 @@ def test_split_preserves_number_of_frames(self):
         for rollout in env.current_epoch_rollouts(split)
     )
     # There are 3 frames in every rollout: the initial one and two returned by
-    # step().
-    self.assertEqual(num_frames, 3 * num_dones)
+    # step(). Additionally there are batch_size observations comming from final
+    # reset at the end of epoch.
+    self.assertEqual(num_frames, 3 * num_dones + batch_size)
 
   def test_generates_data(self):
     # This test needs base env which outputs done after two steps.
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index cb2548d0c..39444c24c 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -178,14 +178,17 @@ def train_agent(real_env, environment_spec, agent_model_dir,
     def decode_png(encoded_png):
       return sess.run(decoded_png_t, feed_dict={encoded_png_p: encoded_png})
 
+    num_input_frames = environment_spec.video_num_input_frames
     initial_frame_rollouts = real_env.current_epoch_rollouts(
-        split=tf.contrib.learn.ModeKeys.TRAIN
+        split=tf.contrib.learn.ModeKeys.TRAIN,
+        minimal_rollout_frames=num_input_frames,
     )
     # TODO(koz4k): Move this to a different module.
     def initial_frame_chooser(batch_size):
       """Frame chooser."""
-      num_frames = environment_spec.video_num_input_frames
-      deterministic_initial_frames = initial_frame_rollouts[0][:num_frames]
+
+      deterministic_initial_frames =\
+          initial_frame_rollouts[0][:num_input_frames]
       if not environment_spec.simulation_random_starts:
         # Deterministic starts: repeat first frames from the first rollout.
         initial_frames = [deterministic_initial_frames] * batch_size
@@ -196,8 +199,8 @@ def initial_frame_chooser(batch_size):
         def choose_initial_frames():
           try:
             rollout = random.choice(initial_frame_rollouts)
-            from_index = random.randrange(len(rollout) - num_frames + 1)
-            return rollout[from_index:(from_index + num_frames)]
+            from_index = random.randrange(len(rollout) - num_input_frames + 1)
+            return rollout[from_index:(from_index + num_input_frames)]
           except ValueError:
             # Rollout too short; repeat.
             return choose_initial_frames()
@@ -206,10 +209,10 @@ def choose_initial_frames():
           # Flip first entry in the batch for deterministic initial frames.
           initial_frames[0] = deterministic_initial_frames
 
-        return np.stack([
-            [decode_png(frame.observation) for frame in initial_frame_stack]
-            for initial_frame_stack in initial_frames
-        ])
+      return np.stack([
+          [decode_png(frame.observation) for frame in initial_frame_stack]
+          for initial_frame_stack in initial_frames
+      ])
 
     environment_spec.add_hparam("initial_frame_chooser", initial_frame_chooser)
 

From 41de57ad12dd69cb3de12629392806e870d7e506 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 19 Oct 2018 14:15:26 -0700
Subject: [PATCH 1047/2720] Fix broken tests introduced by cl/217902485.

PiperOrigin-RevId: 217929186
---
 tensor2tensor/data_generators/gym_env.py      | 33 ++++++++---------
 tensor2tensor/data_generators/gym_env_test.py | 35 ++++++++-----------
 tensor2tensor/models/transformer.py           |  4 +--
 tensor2tensor/rl/trainer_model_based.py       | 21 +++++------
 4 files changed, 42 insertions(+), 51 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index e4a094c87..439d7158b 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -87,7 +87,11 @@ class T2TEnv(video_utils.VideoProblem):
   def __init__(self, batch_size, *args, **kwargs):
     super(T2TEnv, self).__init__(*args, **kwargs)
 
+    self.clear_history()
     self.batch_size = batch_size
+    self._current_batch_frames = [None for _ in range(batch_size)]
+    self._current_batch_rollouts = [[] for _ in range(batch_size)]
+    self._current_epoch_rollouts = []
     self._rollouts_by_epoch_and_split = collections.OrderedDict()
     self.current_epoch = None
     with tf.Graph().as_default() as tf_graph:
@@ -104,6 +108,10 @@ def __str__(self):
     """Returns a string representation of the environment for debug purposes."""
     raise NotImplementedError
 
+  def clear_history(self):
+    """Clears the rollout history."""
+    self._rollouts_by_epoch_and_split = collections.OrderedDict()
+
   def start_new_epoch(self, epoch, load_data_dir=None):
     if not isinstance(epoch, int):
       raise ValueError("Epoch should be integer, got {}".format(epoch))
@@ -112,14 +120,10 @@ def start_new_epoch(self, epoch, load_data_dir=None):
     self.current_epoch = epoch
     self._current_epoch_rollouts = []
     self._rollouts_by_epoch_and_split[epoch] = collections.defaultdict(list)
-    self._current_batch_frames = [None for _ in range(self.batch_size)]
-    self._current_batch_rollouts = [[] for _ in range(self.batch_size)]
     if load_data_dir is not None:
       self._load_epoch_data(load_data_dir)
 
-  def current_epoch_rollouts(self, split=None, minimal_rollout_frames=0):
-    # TODO(KC): order of rollouts (by splits) is a bit uncontrolled
-    # (rollouts_by_split.values() reads dict values), is it a problem?
+  def current_epoch_rollouts(self, split=None):
     rollouts_by_split = self._rollouts_by_epoch_and_split[self.current_epoch]
     if not rollouts_by_split:
       if split is not None:
@@ -127,18 +131,15 @@ def current_epoch_rollouts(self, split=None, minimal_rollout_frames=0):
             "generate_data() should first be called in the current epoch"
         )
       else:
-        rollouts = self._current_epoch_rollouts
+        return self._current_epoch_rollouts
+    if split is not None:
+      return rollouts_by_split[split]
     else:
-      if split is not None:
-        rollouts = rollouts_by_split[split]
-      else:
-        rollouts = [
-            rollout
-            for rollouts in rollouts_by_split.values()
-            for rollout in rollouts
-        ]
-    return [rollout for rollout in rollouts
-            if len(rollout) >= minimal_rollout_frames]
+      return [
+          rollout
+          for rollouts in rollouts_by_split.values()
+          for rollout in rollouts
+      ]
 
   def _preprocess_observations(self, obs):
     """Transforms a batch of observations.
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 0326eb983..15536deb9 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -49,8 +49,9 @@ def __init__(self):
     self._counter = 0
 
   def _generate_ob(self):
-    return np.random.randint(255, size=self.observation_space.shape,
-                             dtype=self.observation_space.dtype)
+    return np.zeros(
+        self.observation_space.shape, self.observation_space.dtype
+    )
 
   def step(self, action):
     done = self._counter % 2 == 1
@@ -78,11 +79,10 @@ def setUp(self):
     self.out_dir = tf.test.get_temp_dir()
     shutil.rmtree(self.out_dir)
     os.mkdir(self.out_dir)
-    np.random.seed(0)
 
-  def init_batch_and_play(self, env_name, steps_per_epoch=1, epochs=(0,),
-                          generate_data=False, batch_size=2, **kwargs):
-    env = gym_env.T2TGymEnv(env_name, batch_size=batch_size, **kwargs)
+  def init_batch_and_play(self, env_name, steps_per_epoch=1,
+                          epochs=(0,), generate_data=False, **kwargs):
+    env = gym_env.T2TGymEnv(env_name, batch_size=2, **kwargs)
     obs = list()
     rewards = list()
     num_dones = 0
@@ -90,7 +90,6 @@ def init_batch_and_play(self, env_name, steps_per_epoch=1, epochs=(0,),
       env.start_new_epoch(epoch, self.out_dir)
       _, epoch_obs, epoch_rewards, epoch_num_dones = \
           self.play(env, steps_per_epoch)
-      epoch_obs.append(env.reset())
       if generate_data:
         env.generate_data(self.out_dir)
       obs.extend(epoch_obs)
@@ -122,26 +121,21 @@ def test_splits_dataset(self):
       self.assertTrue(env.current_epoch_rollouts(split))
 
   def test_split_preserves_number_of_rollouts(self):
-    batch_size = 2
     env, _, _, num_dones = self.init_batch_and_play(
-        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True,
-        batch_size=batch_size
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
     )
 
     num_rollouts_after_split = sum(
         len(env.current_epoch_rollouts(split)) for split in self.splits
     )
-    # After the end of epoch all environments are reset, which increases number
-    # of rollouts by batch size. Number of rollouts could be increased by one
-    # in case a rollout is broken on a boundary between the dataset splits.
-    self.assertGreaterEqual(num_rollouts_after_split, num_dones + batch_size)
-    self.assertLessEqual(num_rollouts_after_split, num_dones + batch_size + 1)
+    # Number of rollouts could be increased by one in case a rollout is broken
+    # on a boundary between the dataset splits.
+    self.assertGreaterEqual(num_rollouts_after_split, num_dones)
+    self.assertLessEqual(num_rollouts_after_split, num_dones + 1)
 
   def test_split_preserves_number_of_frames(self):
-    batch_size = 2
     env, _, _, num_dones = self.init_batch_and_play(
-        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True,
-        batch_size=batch_size
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
     )
 
     num_frames = sum(
@@ -150,9 +144,8 @@ def test_split_preserves_number_of_frames(self):
         for rollout in env.current_epoch_rollouts(split)
     )
     # There are 3 frames in every rollout: the initial one and two returned by
-    # step(). Additionally there are batch_size observations comming from final
-    # reset at the end of epoch.
-    self.assertEqual(num_frames, 3 * num_dones + batch_size)
+    # step().
+    self.assertEqual(num_frames, 3 * num_dones)
 
   def test_generates_data(self):
     # This test needs base env which outputs done after two steps.
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 4120e276b..9efe4d9eb 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1272,7 +1272,7 @@ def transformer_encoder(encoder_input,
           getattr(hparams, "attention_dropout_broadcast_dims", "")))
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_NUM_HIDDEN_LAYERS,
-      value=hparams.num_hidden_layers)
+      value=hparams.num_encoder_layers or hparams.num_hidden_layers)
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_ATTENTION_NUM_HEADS, value=hparams.num_heads)
   mlperf_log.transformer_print(
@@ -1377,7 +1377,7 @@ def transformer_decoder(decoder_input,
 
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_NUM_HIDDEN_LAYERS,
-      value=hparams.num_hidden_layers)
+      value=hparams.num_decoder_layers or hparams.num_hidden_layers)
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_ATTENTION_NUM_HEADS, value=hparams.num_heads)
   mlperf_log.transformer_print(
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 39444c24c..cb2548d0c 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -178,17 +178,14 @@ def train_agent(real_env, environment_spec, agent_model_dir,
     def decode_png(encoded_png):
       return sess.run(decoded_png_t, feed_dict={encoded_png_p: encoded_png})
 
-    num_input_frames = environment_spec.video_num_input_frames
     initial_frame_rollouts = real_env.current_epoch_rollouts(
-        split=tf.contrib.learn.ModeKeys.TRAIN,
-        minimal_rollout_frames=num_input_frames,
+        split=tf.contrib.learn.ModeKeys.TRAIN
     )
     # TODO(koz4k): Move this to a different module.
     def initial_frame_chooser(batch_size):
       """Frame chooser."""
-
-      deterministic_initial_frames =\
-          initial_frame_rollouts[0][:num_input_frames]
+      num_frames = environment_spec.video_num_input_frames
+      deterministic_initial_frames = initial_frame_rollouts[0][:num_frames]
       if not environment_spec.simulation_random_starts:
         # Deterministic starts: repeat first frames from the first rollout.
         initial_frames = [deterministic_initial_frames] * batch_size
@@ -199,8 +196,8 @@ def initial_frame_chooser(batch_size):
         def choose_initial_frames():
           try:
             rollout = random.choice(initial_frame_rollouts)
-            from_index = random.randrange(len(rollout) - num_input_frames + 1)
-            return rollout[from_index:(from_index + num_input_frames)]
+            from_index = random.randrange(len(rollout) - num_frames + 1)
+            return rollout[from_index:(from_index + num_frames)]
           except ValueError:
             # Rollout too short; repeat.
             return choose_initial_frames()
@@ -209,10 +206,10 @@ def choose_initial_frames():
           # Flip first entry in the batch for deterministic initial frames.
           initial_frames[0] = deterministic_initial_frames
 
-      return np.stack([
-          [decode_png(frame.observation) for frame in initial_frame_stack]
-          for initial_frame_stack in initial_frames
-      ])
+        return np.stack([
+            [decode_png(frame.observation) for frame in initial_frame_stack]
+            for initial_frame_stack in initial_frames
+        ])
 
     environment_spec.add_hparam("initial_frame_chooser", initial_frame_chooser)
 

From 001687173dda60483d7f14457c389cf34921a051 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Fri, 19 Oct 2018 15:18:36 -0700
Subject: [PATCH 1048/2720] internal merge of PR #1160

PiperOrigin-RevId: 217940125
---
 tensor2tensor/data_generators/gym_env.py      | 33 +++++++++--------
 tensor2tensor/data_generators/gym_env_test.py | 35 +++++++++++--------
 tensor2tensor/rl/trainer_model_based.py       | 21 ++++++-----
 3 files changed, 49 insertions(+), 40 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 439d7158b..8a156550c 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -87,11 +87,7 @@ class T2TEnv(video_utils.VideoProblem):
   def __init__(self, batch_size, *args, **kwargs):
     super(T2TEnv, self).__init__(*args, **kwargs)
 
-    self.clear_history()
     self.batch_size = batch_size
-    self._current_batch_frames = [None for _ in range(batch_size)]
-    self._current_batch_rollouts = [[] for _ in range(batch_size)]
-    self._current_epoch_rollouts = []
     self._rollouts_by_epoch_and_split = collections.OrderedDict()
     self.current_epoch = None
     with tf.Graph().as_default() as tf_graph:
@@ -108,10 +104,6 @@ def __str__(self):
     """Returns a string representation of the environment for debug purposes."""
     raise NotImplementedError
 
-  def clear_history(self):
-    """Clears the rollout history."""
-    self._rollouts_by_epoch_and_split = collections.OrderedDict()
-
   def start_new_epoch(self, epoch, load_data_dir=None):
     if not isinstance(epoch, int):
       raise ValueError("Epoch should be integer, got {}".format(epoch))
@@ -120,10 +112,14 @@ def start_new_epoch(self, epoch, load_data_dir=None):
     self.current_epoch = epoch
     self._current_epoch_rollouts = []
     self._rollouts_by_epoch_and_split[epoch] = collections.defaultdict(list)
+    self._current_batch_frames = [None for _ in range(self.batch_size)]
+    self._current_batch_rollouts = [[] for _ in range(self.batch_size)]
     if load_data_dir is not None:
       self._load_epoch_data(load_data_dir)
 
-  def current_epoch_rollouts(self, split=None):
+  def current_epoch_rollouts(self, split=None, minimal_rollout_frames=0):
+    # TODO(kc): order of rollouts (by splits) is a bit uncontrolled
+    # (rollouts_by_split.values() reads dict values), is it a problem?
     rollouts_by_split = self._rollouts_by_epoch_and_split[self.current_epoch]
     if not rollouts_by_split:
       if split is not None:
@@ -131,15 +127,18 @@ def current_epoch_rollouts(self, split=None):
             "generate_data() should first be called in the current epoch"
         )
       else:
-        return self._current_epoch_rollouts
-    if split is not None:
-      return rollouts_by_split[split]
+        rollouts = self._current_epoch_rollouts
     else:
-      return [
-          rollout
-          for rollouts in rollouts_by_split.values()
-          for rollout in rollouts
-      ]
+      if split is not None:
+        rollouts = rollouts_by_split[split]
+      else:
+        rollouts = [
+            rollout
+            for rollouts in rollouts_by_split.values()
+            for rollout in rollouts
+        ]
+    return [rollout for rollout in rollouts
+            if len(rollout) >= minimal_rollout_frames]
 
   def _preprocess_observations(self, obs):
     """Transforms a batch of observations.
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 15536deb9..873de708a 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -49,9 +49,8 @@ def __init__(self):
     self._counter = 0
 
   def _generate_ob(self):
-    return np.zeros(
-        self.observation_space.shape, self.observation_space.dtype
-    )
+    return np.random.randint(255, size=self.observation_space.shape,
+                             dtype=self.observation_space.dtype)
 
   def step(self, action):
     done = self._counter % 2 == 1
@@ -79,10 +78,11 @@ def setUp(self):
     self.out_dir = tf.test.get_temp_dir()
     shutil.rmtree(self.out_dir)
     os.mkdir(self.out_dir)
+    np.random.seed(0)
 
-  def init_batch_and_play(self, env_name, steps_per_epoch=1,
-                          epochs=(0,), generate_data=False, **kwargs):
-    env = gym_env.T2TGymEnv(env_name, batch_size=2, **kwargs)
+  def init_batch_and_play(self, env_name, steps_per_epoch=1, epochs=(0,),
+                          generate_data=False, batch_size=2, **kwargs):
+    env = gym_env.T2TGymEnv(env_name, batch_size=batch_size, **kwargs)
     obs = list()
     rewards = list()
     num_dones = 0
@@ -90,6 +90,7 @@ def init_batch_and_play(self, env_name, steps_per_epoch=1,
       env.start_new_epoch(epoch, self.out_dir)
       _, epoch_obs, epoch_rewards, epoch_num_dones = \
           self.play(env, steps_per_epoch)
+      epoch_obs.append(env.reset())
       if generate_data:
         env.generate_data(self.out_dir)
       obs.extend(epoch_obs)
@@ -121,21 +122,26 @@ def test_splits_dataset(self):
       self.assertTrue(env.current_epoch_rollouts(split))
 
   def test_split_preserves_number_of_rollouts(self):
+    batch_size = 2
     env, _, _, num_dones = self.init_batch_and_play(
-        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True,
+        batch_size=batch_size
     )
 
     num_rollouts_after_split = sum(
         len(env.current_epoch_rollouts(split)) for split in self.splits
     )
-    # Number of rollouts could be increased by one in case a rollout is broken
-    # on a boundary between the dataset splits.
-    self.assertGreaterEqual(num_rollouts_after_split, num_dones)
-    self.assertLessEqual(num_rollouts_after_split, num_dones + 1)
+    # After the end of epoch all environments are reset, which increases number
+    # of rollouts by batch size. Number of rollouts could be increased by one
+    # in case a rollout is broken on a boundary between the dataset splits.
+    self.assertGreaterEqual(num_rollouts_after_split, num_dones + batch_size)
+    self.assertLessEqual(num_rollouts_after_split, num_dones + batch_size + 1)
 
   def test_split_preserves_number_of_frames(self):
+    batch_size = 2
     env, _, _, num_dones = self.init_batch_and_play(
-        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True
+        TEST_ENV_NAME, steps_per_epoch=20, generate_data=True,
+        batch_size=batch_size
     )
 
     num_frames = sum(
@@ -144,8 +150,9 @@ def test_split_preserves_number_of_frames(self):
         for rollout in env.current_epoch_rollouts(split)
     )
     # There are 3 frames in every rollout: the initial one and two returned by
-    # step().
-    self.assertEqual(num_frames, 3 * num_dones)
+    # step(). Additionally there are batch_size observations coming from final
+    # reset at the end of epoch.
+    self.assertEqual(num_frames, 3 * num_dones + batch_size)
 
   def test_generates_data(self):
     # This test needs base env which outputs done after two steps.
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index cb2548d0c..39444c24c 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -178,14 +178,17 @@ def train_agent(real_env, environment_spec, agent_model_dir,
     def decode_png(encoded_png):
       return sess.run(decoded_png_t, feed_dict={encoded_png_p: encoded_png})
 
+    num_input_frames = environment_spec.video_num_input_frames
     initial_frame_rollouts = real_env.current_epoch_rollouts(
-        split=tf.contrib.learn.ModeKeys.TRAIN
+        split=tf.contrib.learn.ModeKeys.TRAIN,
+        minimal_rollout_frames=num_input_frames,
     )
     # TODO(koz4k): Move this to a different module.
     def initial_frame_chooser(batch_size):
       """Frame chooser."""
-      num_frames = environment_spec.video_num_input_frames
-      deterministic_initial_frames = initial_frame_rollouts[0][:num_frames]
+
+      deterministic_initial_frames =\
+          initial_frame_rollouts[0][:num_input_frames]
       if not environment_spec.simulation_random_starts:
         # Deterministic starts: repeat first frames from the first rollout.
         initial_frames = [deterministic_initial_frames] * batch_size
@@ -196,8 +199,8 @@ def initial_frame_chooser(batch_size):
         def choose_initial_frames():
           try:
             rollout = random.choice(initial_frame_rollouts)
-            from_index = random.randrange(len(rollout) - num_frames + 1)
-            return rollout[from_index:(from_index + num_frames)]
+            from_index = random.randrange(len(rollout) - num_input_frames + 1)
+            return rollout[from_index:(from_index + num_input_frames)]
           except ValueError:
             # Rollout too short; repeat.
             return choose_initial_frames()
@@ -206,10 +209,10 @@ def choose_initial_frames():
           # Flip first entry in the batch for deterministic initial frames.
           initial_frames[0] = deterministic_initial_frames
 
-        return np.stack([
-            [decode_png(frame.observation) for frame in initial_frame_stack]
-            for initial_frame_stack in initial_frames
-        ])
+      return np.stack([
+          [decode_png(frame.observation) for frame in initial_frame_stack]
+          for initial_frame_stack in initial_frames
+      ])
 
     environment_spec.add_hparam("initial_frame_chooser", initial_frame_chooser)
 

From 8313b67cc9c2e14eee0adeaf86137b7f8d3018e7 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Sat, 20 Oct 2018 12:49:34 -0700
Subject: [PATCH 1049/2720] Rolling back an SV2P modification which completely
 broke the model.

PiperOrigin-RevId: 218020327
---
 tensor2tensor/models/video/sv2p.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index ed20e754d..3fe8af934 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -132,10 +132,7 @@ def bottom_part_tower(self, input_image, input_reward, action, latent,
           enc2, input_reward, "reward_enc")
     if latent is not None and not concat_latent:
       with tf.control_dependencies([latent]):
-        # This is the original SV2P implementation
-        # But we will tile and concat to support various latent sizes.
-        # enc2 = tf.concat([enc2, latent], axis=3)
-        enc2 = tile_and_concat(enc2, latent, concat_latent=concat_latent)
+        enc2 = tf.concat([enc2, latent], axis=3)
 
     enc3 = tfl.conv2d(enc2, hidden4.get_shape()[3], [1, 1], strides=(1, 1),
                       padding="SAME", activation=tf.nn.relu, name="conv4")

From fcc6d9a32b695bd7aff0aaf152333c865390f3ff Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Sat, 20 Oct 2018 13:51:57 -0700
Subject: [PATCH 1050/2720] fixing two bugs.

PiperOrigin-RevId: 218023024
---
 tensor2tensor/models/video/base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 5017d5167..248adaebc 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -429,6 +429,7 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
     orig_frame_shape = common_layers.shape_list(all_frames[0])
     batch_size = orig_frame_shape[0]
     ss_func = self.get_scheduled_sample_func(batch_size)
+    target_frames = all_frames[hparams.video_num_input_frames:]
     extra_loss = 0.0
     internal_states = None
 
@@ -467,7 +468,7 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
       # Scheduled sampling during training.
       if self.is_training:
         if self.is_recurrent_model:
-          done_warm_start = i > hparams.video_num_input_frames - 1
+          done_warm_start = i >= hparams.video_num_input_frames - 1
         else:
           done_warm_start = True  # Always true for non-reccurent networks.
         groundtruth_items = [target_frame]
@@ -491,7 +492,6 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
       res_rewards = res_rewards[hparams.video_num_input_frames-1:]
       sampled_frames = sampled_frames[hparams.video_num_input_frames-1:]
 
-    target_frames = all_frames[hparams.video_num_input_frames:]
     self.visualize_predictions(sampled_frames, target_frames)
 
     output_frames = tf.stack(res_frames, axis=1)

From 571780e160a003cf743e13209defd088bb22020f Mon Sep 17 00:00:00 2001
From: Marcin Michalski <michalski@google.com>
Date: Mon, 22 Oct 2018 07:45:05 -0700
Subject: [PATCH 1051/2720] Interface for the new video metric.

PiperOrigin-RevId: 218166374
---
 .../metrics/video_conditional_fvd.py          | 80 +++++++++++++++++++
 .../metrics/video_conditional_fvd_test.py     | 40 ++++++++++
 2 files changed, 120 insertions(+)
 create mode 100644 tensor2tensor/metrics/video_conditional_fvd.py
 create mode 100644 tensor2tensor/metrics/video_conditional_fvd_test.py

diff --git a/tensor2tensor/metrics/video_conditional_fvd.py b/tensor2tensor/metrics/video_conditional_fvd.py
new file mode 100644
index 000000000..b9965e1cb
--- /dev/null
+++ b/tensor2tensor/metrics/video_conditional_fvd.py
@@ -0,0 +1,80 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Conditional FVD metric on video.
+
+FVD - Frechet Video Distance
+
+This is the metric that is inspired by FID, but applied to
+video rather than to images.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+
+
+class VideoEvaluationDataset(
+    collections.namedtuple(
+        'VideoEvaluationDataset',
+        ['n_input_frames', 'n_output_frames', 'get_video_batch_fn'])):
+  """Dataset for video evaluation.
+
+  This tuple describes the video problem for Evaluation.
+  Args:
+     n_input_frames: number of frames passed to the model to condition on.
+     n_output_frames: number of frames that model should return.
+     get_video_batch_fn: function that accepts a batch size and returns a tensor
+       with real video, which should match <uint8>[batch_size, N, height, width,
+       depth], where N is n_input_frames + n_output_frames.
+  """
+  pass
+
+
+class Model(
+    collections.namedtuple('Model', [
+        'apply_fn', 'load_fn',
+    ])):
+  """Model that should be evaluated.
+
+  Args:
+    apply_fn: will be called with a single tensor (floats between 0 and 255
+              of shape [batch_size, n_input_frames, height, width, depth]),
+              that will contain input frames.
+              it should return a single tensor with output frames (floats
+              between 0 and 255, of shape
+              [batch_size, n_output_frames, height, width, depth])
+    load_fn: Callable, that receives session as an argument.
+             Should load the variables from the checkpoint.
+  """
+  pass
+
+
+def evaluate_model(video_eval_dataset, model, num_batches, batch_size):
+  """Computes the FVD video metric.
+
+  Args:
+    video_eval_dataset: VideoEvaluationDataset tuple with video and frames
+      information.
+    model: Model tuple with model to evaluate.
+    num_batches: number of batches to evaluate.
+    batch_size: number of videos to compute per batch.
+
+  Returns:
+    FVD metric (float).
+  """
+  del video_eval_dataset, model, num_batches, batch_size
diff --git a/tensor2tensor/metrics/video_conditional_fvd_test.py b/tensor2tensor/metrics/video_conditional_fvd_test.py
new file mode 100644
index 000000000..ec1712eec
--- /dev/null
+++ b/tensor2tensor/metrics/video_conditional_fvd_test.py
@@ -0,0 +1,40 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 video_conditional_fvd."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.metrics import video_conditional_fvd
+import tensorflow as tf
+
+
+class VideoConditionalFvdTest(tf.test.TestCase):
+
+  def test_sample(self):
+    dataset = video_conditional_fvd.VideoEvaluationDataset(
+        n_input_frames=4,
+        n_output_frames=10,
+        get_video_batch_fn=None)
+    model = video_conditional_fvd.Model(
+        apply_fn=None,
+        load_fn=None)
+    video_conditional_fvd.evaluate_model(dataset, model, 10, 16)
+
+
+if __name__ == '__main__':
+  tf.test.main()

From f45c047ea5a9a8838c84d388f4714fd183afc361 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 22 Oct 2018 11:16:04 -0700
Subject: [PATCH 1052/2720] simplifying sv2p discrete to build again.

PiperOrigin-RevId: 218199730
---
 tensor2tensor/layers/common_video.py      |  3 ++
 tensor2tensor/models/video/sv2p.py        | 59 +++++++----------------
 tensor2tensor/models/video/sv2p_params.py |  8 +--
 3 files changed, 22 insertions(+), 48 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index e551e49b6..80b27c8f9 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -344,6 +344,9 @@ def tile_and_concat(image, latent, concat_latent=True):
   height, width = image_shape[1], image_shape[2]
   latent_dims = latent_shape[1]
 
+  if height < latent_dims:
+    raise ValueError("Latent is too big to tile.")
+
   height_multiples = height // latent_dims
   pad = height - (height_multiples * latent_dims)
   latent = tf.reshape(latent, (-1, latent_dims, 1, 1))
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 3fe8af934..4a6bb070a 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -386,52 +386,26 @@ def basic_conv_net(self, images, conv_size, scope):
         x = tfcl.layer_norm(x)
     return x
 
-  def learned_discrete_tower(self, input_image, target_image):
+  def simple_discrete_latent_tower(self, input_image, target_image):
     hparams = self.hparams
 
-    # Encode the input frames into a prior encoding.
-    conv_size = [64, 32, 32, 1]
-    prior_enc = self.basic_conv_net(input_image, conv_size, "prior_enc")
-    tower_output_shape = common_layers.shape_list(prior_enc)
-    batch_size = tower_output_shape[0]
-    prior_enc = tfl.flatten(prior_enc)
-
-    def decode_bits(b):
-      return common_video.encode_to_shape(b, tower_output_shape, "bits_dec")
-
     if self.is_predicting:
-      if hparams.full_latent_tower:
-        rand = tf.random_uniform([batch_size, hparams.bottleneck_bits])
-        bits = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
-      else:
-        # Generate bit using the learned prior at inference time.
-        bits, _ = discretization.predict_bits_with_lstm(
-            prior_enc,
-            hparams.latent_predictor_state_size,
-            hparams.bottleneck_bits,
-            temperature=hparams.latent_predictor_temperature)
-      return decode_bits(bits), 0.0
-
-    # Encode the input and target frames into posterior.
-    x = tf.concat([input_image, target_image], axis=-1)
-    x = self.basic_conv_net(x, conv_size, "posterior_enc")
-    x = tfl.flatten(x)
-    bits, bits_clean = discretization.tanh_discrete_bottleneck(
-        x, hparams.bottleneck_bits,
+      batch_size = common_layers.shape_list(input_image)[0]
+      rand = tf.random_uniform([batch_size, hparams.bottleneck_bits])
+      bits = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
+      return bits
+
+    conv_size = self.tinyify([64, 32, 32, 1])
+    pair = tf.concat([input_image, target_image], axis=-1)
+    posterior_enc = self.basic_conv_net(pair, conv_size, "posterior_enc")
+    posterior_enc = tfl.flatten(posterior_enc)
+    bits, _ = discretization.tanh_discrete_bottleneck(
+        posterior_enc,
+        hparams.bottleneck_bits,
         hparams.bottleneck_noise,
         hparams.discretize_warmup_steps,
         hparams.mode)
-
-    pred_loss = 0.0
-    if not hparams.full_latent_tower:
-      # Learn the prior by matching the posterior.
-      _, pred_loss = discretization.predict_bits_with_lstm(
-          prior_enc,
-          hparams.latent_predictor_state_size,
-          hparams.bottleneck_bits,
-          target_bits=bits_clean)
-
-    return decode_bits(bits), pred_loss
+    return bits
 
   def next_frame(self, frames, actions, rewards, target_frame,
                  internal_states, video_features):
@@ -441,10 +415,11 @@ def next_frame(self, frames, actions, rewards, target_frame,
     if internal_states is None:
       internal_states = [None] * (5 if self.hparams.small_mode else 7)
 
-    latent, extra_loss = self.learned_discrete_tower(frames, target_frame)
+    extra_loss = 0.0
+    latent = self.simple_discrete_latent_tower(frames, target_frame)
 
     pred_image, internal_states = self.construct_predictive_tower(
-        frames, None, actions, internal_states, latent)
+        frames, None, actions, internal_states, latent, True)
 
     if not self.has_rewards:
       return pred_image, None, extra_loss, internal_states
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 68f9380b1..b0156f9ce 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -56,13 +56,8 @@ def next_frame_sv2p_discrete():
   hparams = next_frame_sv2p()
   hparams.action_injection = "multiplicative"
   hparams.small_mode = True
-  hparams.add_hparam("bottleneck_bits", 128)
+  hparams.add_hparam("bottleneck_bits", 16)
   hparams.add_hparam("bottleneck_noise", 0.02)
-  hparams.add_hparam("discrete_warmup_steps", 40000)
-  hparams.add_hparam("full_latent_tower", False)
-  hparams.add_hparam("latent_predictor_state_size", 128)
-  hparams.add_hparam("latent_predictor_temperature", 0.5)
-  hparams.add_hparam("discretize_warmup_steps", 40000)
   return hparams
 
 
@@ -80,6 +75,7 @@ def next_frame_sv2p_atari():
   hparams.latent_loss_multiplier = 1e-3
   hparams.information_capacity = 0.0
   hparams.small_mode = True
+  hparams.internal_loss = True
   return hparams
 
 

From f1fd3460766c258f32348d1ee852a2a6e6c20b09 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 22 Oct 2018 12:59:36 -0700
Subject: [PATCH 1053/2720] adding SAVP for RL. Changes:

PiperOrigin-RevId: 218217686
---
 tensor2tensor/models/video/base.py |  28 ++++++-
 tensor2tensor/models/video/savp.py | 117 +++++++++++++++++++++++++++--
 2 files changed, 138 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 248adaebc..583bf7a49 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -112,6 +112,25 @@ def video_features(
     del all_frames, all_actions, all_rewards, all_raw_frames
     return None
 
+  def video_extra_loss(self, frames_predicted, frames_target,
+                       internal_states, video_features):
+    """Optional video wide extra loss.
+
+      If the model needs to calculate some extra loss across all predicted
+      frames (e.g. in case of video GANS loss) override this function.
+
+    Args:
+      frames_predicted: list of all predicted frames.
+      frames_target: list of all target frames.
+      internal_states: internal states of the video.
+      video_features: video wide features coming from video_features function.
+
+    Returns:
+      extra_loss: extra video side loss.
+    """
+    del frames_predicted, frames_target, internal_states, video_features
+    return 0.0
+
   @property
   def is_recurrent_model(self):
     """Set to true if your model is recurrent. False otherwise.
@@ -429,7 +448,7 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
     orig_frame_shape = common_layers.shape_list(all_frames[0])
     batch_size = orig_frame_shape[0]
     ss_func = self.get_scheduled_sample_func(batch_size)
-    target_frames = all_frames[hparams.video_num_input_frames:]
+    target_frames = []
     extra_loss = 0.0
     internal_states = None
 
@@ -448,6 +467,7 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
       frames, actions, rewards, target_index = self.__get_next_inputs(
           i, all_frames, all_actions, all_rewards)
       target_frame = all_frames[target_index]
+      target_frames.append(tf.identity(target_frame))
 
       with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
         func_in = (frames, actions, rewards, target_frame,
@@ -477,6 +497,11 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
             done_warm_start, groundtruth_items, generated_items, ss_func)
         all_frames[target_index] = ss_frame
 
+    video_extra_loss = self.video_extra_loss(
+        sampled_frames, target_frames, internal_states, video_features)
+    tf.summary.scalar("video_extra_loss", video_extra_loss)
+    extra_loss += video_extra_loss
+
     if self.is_recurrent_model:
       has_input_predictions = hparams.video_num_input_frames > 1
       if self.is_training and hparams.internal_loss and has_input_predictions:
@@ -491,6 +516,7 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
       res_frames = res_frames[hparams.video_num_input_frames-1:]
       res_rewards = res_rewards[hparams.video_num_input_frames-1:]
       sampled_frames = sampled_frames[hparams.video_num_input_frames-1:]
+      target_frames = target_frames[hparams.video_num_input_frames-1:]
 
     self.visualize_predictions(sampled_frames, target_frames)
 
diff --git a/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index 559a66450..f1ba8e486 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -36,12 +36,11 @@
 gan_losses = tf.contrib.gan.losses.wargs
 
 
-@registry.register_model
-class NextFrameSAVP(sv2p.NextFrameSv2pLegacy):
-  """Stochastic Adversarial Video Prediction."""
+class NextFrameSavpBase(object):
+  """Main function for Stochastic Adversarial Video Prediction."""
 
   def encoder(self, inputs, n_layers=3):
-    """COnvnet that encodes inputs into mean and std of a gaussian.
+    """Convnet that encodes inputs into mean and std of a gaussian.
 
     Args:
      inputs: 5-D Tensor, shape (batch_size, num_frames, width, height, channels)
@@ -256,7 +255,7 @@ def get_gan_loss(self, true_frames, gen_frames, name):
     if self.hparams.gan_optimization == "joint":
       gan_loss = gan_g_loss + gan_d_loss
     else:
-      curr_step = tf.train.get_or_create_global_step()
+      curr_step = self.get_iteration_num()
       gan_loss = tf.cond(
           tf.logical_not(curr_step % 2 == 0), lambda: gan_g_loss,
           lambda: gan_d_loss)
@@ -264,13 +263,14 @@ def get_gan_loss(self, true_frames, gen_frames, name):
 
   def get_extra_loss(self, latent_means=None, latent_stds=None,
                      true_frames=None, gen_frames=None):
+    """Gets extra loss from VAE and GAN."""
     if not self.is_training:
       return 0.0
 
     vae_loss, d_vae_loss, d_gan_loss = 0.0, 0.0, 0.0
     # Use sv2p's KL divergence computation.
     if self.hparams.use_vae:
-      vae_loss = super(NextFrameSAVP, self).get_extra_loss(
+      vae_loss = super(NextFrameSavpBase, self).get_extra_loss(
           latent_means=latent_means, latent_stds=latent_stds)
 
     if self.hparams.use_gan:
@@ -330,6 +330,11 @@ def pad_conv3d_lrelu(self, activations, n_filters, kernel_size, strides,
   def train_hooks():
     return [update_ops_hook.UpdateOpsHook()]
 
+
+@registry.register_model
+class NextFrameSAVP(sv2p.NextFrameSv2pLegacy, NextFrameSavpBase):
+  """Stochastic Adversarial Video Prediction."""
+
   def construct_model(self, images, actions, rewards):
     """Model that takes in images and returns predictions.
 
@@ -450,3 +455,103 @@ def construct_model(self, images, actions, rewards):
       return gen_cond_video, fake_rewards, latent_means, latent_stds
     else:
       return self.gen_prior_video, fake_rewards, latent_means, latent_stds
+
+
+@registry.register_model
+class NextFrameSavpRl(NextFrameSavpBase, sv2p.NextFrameSv2p):
+  """Stochastic Adversarial Video Prediction for RL pipeline."""
+
+  def video_features(
+      self, all_frames, all_actions, all_rewards, all_raw_frames):
+    """No video wide feature."""
+    del all_actions, all_rewards, all_raw_frames
+    # Concatenate x_{t-1} and x_{t} along depth and encode it to
+    # produce the mean and standard deviation of z_{t-1}
+    seq_len = len(all_frames)
+    image_pairs = tf.concat([all_frames[:seq_len-1],
+                             all_frames[1:seq_len]], axis=-1)
+    z_mu, z_log_sigma_sq = self.encoder(image_pairs)
+    # Unstack z_mu and z_log_sigma_sq along the time dimension.
+    z_mu = tf.unstack(z_mu, axis=0)
+    z_log_sigma_sq = tf.unstack(z_log_sigma_sq, axis=0)
+    return [z_mu, z_log_sigma_sq]
+
+  def video_extra_loss(self, frames_predicted, frames_target,
+                       internal_states, video_features):
+
+    if not self.is_training:
+      return 0.0
+
+    latent_means, latent_stds = video_features
+    true_frames, gen_frames = frames_target, frames_predicted
+
+    loss = super(NextFrameSavpRl, self).get_extra_loss(
+        latent_means=latent_means, latent_stds=latent_stds,
+        true_frames=true_frames, gen_frames=gen_frames)
+    return loss
+
+  def next_frame(self, frames, actions, rewards, target_frame,
+                 internal_states, video_features):
+    del target_frame
+
+    if not self.hparams.use_vae or self.hparams.use_gan:
+      raise NotImplementedError("Only supporting VAE for now.")
+
+    image, action, reward = frames[0], actions[0], rewards[0]
+    latent_dims = self.hparams.z_dim
+    batch_size = common_layers.shape_list(image)[0]
+
+    if internal_states is None:
+      # Initialize LSTM State
+      frame_index = 0
+      lstm_state = [None] * 7
+      cond_latent_state, prior_latent_state = None, None
+      gen_prior_video = []
+    else:
+      (frame_index, lstm_state, cond_latent_state,
+       prior_latent_state, gen_prior_video) = internal_states
+
+    z_mu, log_sigma_sq = video_features
+    z_mu, log_sigma_sq = z_mu[frame_index], log_sigma_sq[frame_index]
+
+    # Sample latents using a gaussian centered at conditional mu and std.
+    latent = common_video.get_gaussian_tensor(z_mu, log_sigma_sq)
+
+    # Sample prior latents from isotropic normal distribution.
+    prior_latent = tf.random_normal(tf.shape(latent), dtype=tf.float32)
+
+    # # LSTM that encodes correlations between conditional latents.
+    # # Pg 22 in https://arxiv.org/pdf/1804.01523.pdf
+    enc_cond_latent, cond_latent_state = common_video.basic_lstm(
+        latent, cond_latent_state, latent_dims, name="cond_latent")
+
+    # LSTM that encodes correlations between prior latents.
+    enc_prior_latent, prior_latent_state = common_video.basic_lstm(
+        prior_latent, prior_latent_state, latent_dims, name="prior_latent")
+
+    all_latents = tf.concat([enc_cond_latent, enc_prior_latent], axis=0)
+    all_image = tf.concat([image, image], 0)
+    all_action = tf.concat([action, action], 0) if self.has_actions else None
+
+    all_pred_images, lstm_state = self.construct_predictive_tower(
+        all_image, None, all_action, lstm_state, all_latents,
+        concat_latent=True)
+
+    cond_pred_images, prior_pred_images = \
+      all_pred_images[:batch_size], all_pred_images[batch_size:]
+
+    if self.is_training and self.hparams.use_vae:
+      pred_image = cond_pred_images
+    else:
+      pred_image = prior_pred_images
+
+    gen_prior_video.append(prior_pred_images)
+    internal_states = (frame_index + 1, lstm_state, cond_latent_state,
+                       prior_latent_state, gen_prior_video)
+
+    if not self.has_rewards:
+      return pred_image, None, 0.0, internal_states
+
+    pred_reward = self.reward_prediction(
+        pred_image, action, reward, latent)
+    return pred_image, pred_reward, 0.0, internal_states

From 587915222ba50d0848128703f21d9cfbb0500af2 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Mon, 22 Oct 2018 23:31:38 +0200
Subject: [PATCH 1054/2720] Small RL pipeline cleanup (#1159)

* Remove passing parameters by FLAGS

* Remove InGraphBatchEnv.__getattr__

* Move force_beginning_resets to environment_spec
---
 tensor2tensor/data_generators/all_problems.py |  1 -
 tensor2tensor/data_generators/gym_env.py      | 74 ++++++++++++-------
 tensor2tensor/models/research/rl.py           | 16 ++--
 tensor2tensor/rl/collect.py                   |  8 +-
 tensor2tensor/rl/envs/in_graph_batch_env.py   | 11 ---
 tensor2tensor/rl/envs/simulated_batch_env.py  | 20 ++---
 tensor2tensor/rl/trainer_model_based.py       | 66 ++++++-----------
 7 files changed, 91 insertions(+), 105 deletions(-)

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index a872cbdac..311811dbf 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -40,7 +40,6 @@
     "tensor2tensor.data_generators.function_docstring",
     "tensor2tensor.data_generators.gene_expression",
     "tensor2tensor.data_generators.google_robot_pushing",
-    "tensor2tensor.data_generators.gym_problems_specs",
     "tensor2tensor.data_generators.ice_parsing",
     "tensor2tensor.data_generators.imagenet",
     "tensor2tensor.data_generators.image_lsun",
diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 8a156550c..4bcbb1af4 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -61,7 +61,42 @@ def make_gym_env(name, timesteps_limit=-1):
   return env
 
 
-class T2TEnv(video_utils.VideoProblem):
+class EnvSimulationProblem(video_utils.VideoProblem):
+  """Base Problem class for use with world models.
+
+  Attributes:
+    action_space: Gym action space. Should be overridden in derived classes.
+    reward_range: Tuple (min, max) representing the range of rewards. Limits
+      should be integer (discrete rewards).
+  """
+
+  action_space = None
+  reward_range = (-1, 1)
+
+  @property
+  def num_actions(self):
+    return self.action_space.n
+
+  @property
+  def num_rewards(self):
+    (min_reward, max_reward) = self.reward_range
+    return max_reward - min_reward + 1
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    def make_modality(name):
+      return {
+          "{}s".format(name): ("video", 256),
+          "{}_reward".format(name): ("symbol:weights_all", self.num_rewards),
+          "{}_action".format(name): ("symbol:weights_all", self.num_actions)
+      }
+    p.input_modality = make_modality("input")
+    p.target_modality = make_modality("target")
+    p.input_space_id = problem.SpaceID.IMAGE
+    p.target_space_id = problem.SpaceID.IMAGE
+
+
+class T2TEnv(EnvSimulationProblem):
   """Abstract class representing a batch of environments.
 
   Attributes:
@@ -69,9 +104,6 @@ class T2TEnv(video_utils.VideoProblem):
     batch_size: Number of environments played simultaneously.
     observation_space: Gym observation space. Should be overridden in derived
       classes.
-    action_space: Gym action space. Should be overridden in derived classes.
-    reward_range: Tuple (min, max) representing the range of rewards. Limits
-      should be integer (discrete rewards).
     name: Problem name for generating filenames. Should be overridden in
       derived classes.
 
@@ -80,8 +112,6 @@ class T2TEnv(video_utils.VideoProblem):
   """
 
   observation_space = None
-  action_space = None
-  reward_range = (-1, 1)
   name = None
 
   def __init__(self, batch_size, *args, **kwargs):
@@ -308,28 +338,6 @@ def frame_width(self):
   def only_keep_videos_from_0th_frame(self):
     return False
 
-  @property
-  def num_actions(self):
-    return self.action_space.n
-
-  @property
-  def num_rewards(self):
-    (min_reward, max_reward) = self.reward_range
-    return max_reward - min_reward + 1
-
-  def hparams(self, defaults, unused_model_hparams):
-    p = defaults
-    def make_modality(name):
-      return {
-          "{}s".format(name): ("video", 256),
-          "{}_reward".format(name): ("symbol:weights_all", self.num_rewards),
-          "{}_action".format(name): ("symbol:weights_all", self.num_actions)
-      }
-    p.input_modality = make_modality("input")
-    p.target_modality = make_modality("target")
-    p.input_space_id = problem.SpaceID.IMAGE
-    p.target_space_id = problem.SpaceID.IMAGE
-
   def _generate_frames(self, rollouts):
     for rollout in rollouts:
       for (frame_number, frame) in enumerate(rollout):
@@ -611,6 +619,16 @@ def close(self):
     for env in self._envs:
       env.close()
 
+
+class DummyWorldModelProblem(EnvSimulationProblem):
+  """Dummy Problem for world model prediction."""
+
+  def __init__(self, action_space, reward_range):
+    super(DummyWorldModelProblem, self).__init__()
+    self.action_space = action_space
+    self.reward_range = reward_range
+
+
 # Atari registration.
 
 # Game list from our list of ROMs
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 6ace4cd75..4b96264d9 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -19,6 +19,7 @@
 import functools
 import operator
 import gym
+import six
 
 from tensor2tensor.data_generators import gym_env
 from tensor2tensor.layers import common_hparams
@@ -154,23 +155,20 @@ def standard_atari_env_spec(env=None, simulated=False):
       simulated_env=simulated,
       reward_range=env.reward_range,
       observation_space=env.observation_space,
-      action_space=env.action_space
+      action_space=env.action_space,
+      force_beginning_resets=True
   )
   if not simulated:
     env_spec.add_hparam("env", env)
   return env_spec
 
 
-def standard_atari_env_simulated_spec(
-    real_env, video_num_input_frames, video_num_target_frames
-):
+def standard_atari_env_simulated_spec(real_env, **kwargs):
   """Spec."""
   env_spec = standard_atari_env_spec(real_env, simulated=True)
-  env_spec.add_hparam("simulation_random_starts", True)
-  env_spec.add_hparam("simulation_flip_first_random_for_beginning", True)
-  env_spec.add_hparam("intrinsic_reward_scale", 0.0)
-  env_spec.add_hparam("video_num_input_frames", video_num_input_frames)
-  env_spec.add_hparam("video_num_target_frames", video_num_target_frames)
+  for (name, value) in six.iteritems(kwargs):
+    env_spec.add_hparam(name, value)
+  env_spec.force_beginning_resets = False
   return env_spec
 
 
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 4c11cd5b2..b187090c9 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -156,11 +156,9 @@ def initialization_lambda(sess):
     should_reset_var = tf.Variable(True, trainable=False)
     zeros_tensor = tf.zeros(len(batch_env))
 
-  if "force_beginning_resets" in hparams:
-    force_beginning_resets = hparams.force_beginning_resets
-  else:
-    force_beginning_resets = False
-  force_beginning_resets = tf.convert_to_tensor(force_beginning_resets)
+  force_beginning_resets = tf.convert_to_tensor(
+      environment_spec.force_beginning_resets
+  )
 
   def reset_ops_group():
     return tf.group(batch_env.reset(tf.range(len(batch_env))),
diff --git a/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index 87b2a0caf..b20d18f10 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -35,17 +35,6 @@ def __init__(self, observ_space, action_space):
     self.observ_space = observ_space
     self.action_space = action_space
 
-  def __getattr__(self, name):
-    """Forward unimplemented attributes to one of the original environments.
-
-    Args:
-      name: Attribute that was accessed.
-
-    Returns:
-      Value behind the attribute name in one of the original environments.
-    """
-    return getattr(self._batch_env, name)
-
   def __str__(self):
     return "InGraphEnv(%s)" % str(self._batch_env)
 
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 72b3d721d..b61a1b6c9 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -22,6 +22,9 @@
 from __future__ import division
 from __future__ import print_function
 
+import copy
+
+from tensor2tensor.data_generators.gym_env import DummyWorldModelProblem
 from tensor2tensor.layers import common_layers
 from tensor2tensor.rl.envs import in_graph_batch_env
 from tensor2tensor.utils import registry
@@ -30,10 +33,6 @@
 import tensorflow as tf
 
 
-flags = tf.flags
-FLAGS = flags.FLAGS
-
-
 class HistoryBuffer(object):
   """History Buffer."""
 
@@ -103,12 +102,15 @@ def __init__(self, environment_spec, length):
     self._num_frames = environment_spec.video_num_input_frames
     self._intrinsic_reward_scale = environment_spec.intrinsic_reward_scale
 
-    # TODO(koz4k): Pass by argument.
-    model_hparams = trainer_lib.create_hparams(
-        FLAGS.hparams_set, problem_name=FLAGS.problem)
+    model_hparams = copy.copy(environment_spec.model_hparams)
+    problem = DummyWorldModelProblem(
+        environment_spec.action_space, environment_spec.reward_range
+    )
+    trainer_lib.add_problem_hparams(model_hparams, problem)
     model_hparams.force_full_predict = True
-    self._model = registry.model(FLAGS.model)(
-        model_hparams, tf.estimator.ModeKeys.PREDICT)
+    self._model = registry.model(environment_spec.model_name)(
+        model_hparams, tf.estimator.ModeKeys.PREDICT
+    )
 
     self.history_buffer = HistoryBuffer(
         environment_spec.initial_frame_chooser, self.observ_shape,
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 39444c24c..5adcec81d 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -26,8 +26,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import contextlib
-import copy
 import datetime
 import math
 import os
@@ -50,17 +48,6 @@
 FLAGS = flags.FLAGS
 
 
-@contextlib.contextmanager
-def temporary_flags(flag_settings):
-  old_values = {}
-  for flag_name, flag_value in flag_settings.items():
-    old_values[flag_name] = getattr(FLAGS, flag_name)
-    setattr(FLAGS, flag_name, flag_value)
-  yield
-  for flag_name, flag_value in old_values.items():
-    setattr(FLAGS, flag_name, flag_value)
-
-
 def real_ppo_epoch_increment(hparams):
   """PPO increment."""
   if hparams.gather_ppo_real_env_data:
@@ -142,9 +129,8 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
   getattr(exp, schedule)()
 
 
-def train_agent(real_env, environment_spec, agent_model_dir,
-                event_dir, world_model_dir, data_dir, hparams, ppo_epochs_num,
-                epoch=0, is_final_epoch=False):
+def train_agent(real_env, agent_model_dir, event_dir, world_model_dir, data_dir,
+                hparams, ppo_epochs_num, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the simulated environment."""
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_params_names = ["epochs_num", "epoch_length",
@@ -161,16 +147,26 @@ def train_agent(real_env, environment_spec, agent_model_dir,
 
   ppo_hparams.save_models_every_epochs = 10
   ppo_hparams.world_model_dir = world_model_dir
-  ppo_hparams.add_hparam("force_beginning_resets", True)
 
-  # Adding model hparams for model specific adjustments
-  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
-  ppo_hparams.add_hparam("model_hparams", model_hparams)
-
-  environment_spec = copy.copy(environment_spec)
-  environment_spec_param_names = ["intrinsic_reward_scale"]
-  for param_name in environment_spec_param_names:
-    environment_spec.set_hparam(param_name, hparams.get(param_name))
+  environment_spec_params = {
+      param_name: hparams.get(param_name)
+      for param_name in [
+          "intrinsic_reward_scale", "simulation_random_starts",
+          "simulation_flip_first_random_for_beginning"
+      ]
+  }
+  environment_spec_params.update({
+      "model_name": hparams.generative_model,
+      "model_hparams": trainer_lib.create_hparams(
+          hparams.generative_model_params
+      ),
+      # Hardcoded for now. TODO(koz4k): Make it a hparam.
+      "video_num_input_frames": 4,
+      "video_num_target_frames": 1
+  })
+  environment_spec = rl.standard_atari_env_simulated_spec(
+      real_env, **environment_spec_params
+  )
 
   with tf.Session() as sess:
     encoded_png_p = tf.placeholder(tf.string)
@@ -218,16 +214,8 @@ def choose_initial_frames():
 
     ppo_hparams.add_hparam("environment_spec", environment_spec)
 
-    # TODO(koz4k): Pass by arguments.
-    with temporary_flags({
-        "problem": real_env,
-        "model": hparams.generative_model,
-        "hparams_set": hparams.generative_model_params,
-        "output_dir": world_model_dir,
-        "data_dir": data_dir,
-    }):
-      rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
-                           name_scope="ppo_sim%d" % (epoch + 1))
+    rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
+                         name_scope="ppo_sim%d" % (epoch + 1))
 
   return ppo_epochs_num
 
@@ -372,12 +360,6 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   mean_clipped_reward_summary.value.add(tag="mean_clipped_reward",
                                         simple_value=None)
 
-  sim_env_spec = rl.standard_atari_env_simulated_spec(
-      env,
-      # Hardcoded for now. TODO(koz4k): Make it a hparam.
-      video_num_input_frames=4, video_num_target_frames=1
-  )
-
   world_model_steps_num = 0
 
   for epoch in range(hparams.epochs):
@@ -402,7 +384,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       ppo_model_dir = ppo_event_dir
 
     ppo_epochs_num = train_agent(
-        env, sim_env_spec, ppo_model_dir, ppo_event_dir,
+        env, ppo_model_dir, ppo_event_dir,
         directories["world_model"], data_dir, hparams, ppo_epochs_num,
         epoch=epoch, is_final_epoch=is_final_epoch
     )

From 5b51c0f1164c9379cd2c5c5052858505648d482d Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Mon, 22 Oct 2018 14:51:08 -0700
Subject: [PATCH 1055/2720] internal merge of PR #1159

PiperOrigin-RevId: 218238281
---
 tensor2tensor/rl/trainer_model_based.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 5adcec81d..80961ddff 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -132,6 +132,7 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
 def train_agent(real_env, agent_model_dir, event_dir, world_model_dir, data_dir,
                 hparams, ppo_epochs_num, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the simulated environment."""
+  del data_dir
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   ppo_params_names = ["epochs_num", "epoch_length",
                       "learning_rate", "num_agents",

From 1c3896e82ca3738fa39912e6c0fdb5cb6d3b74e4 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Mon, 22 Oct 2018 16:10:20 -0700
Subject: [PATCH 1056/2720] ResNet changes for MeshTF

PiperOrigin-RevId: 218251907
---
 tensor2tensor/models/mtf_resnet.py | 94 +++++++++---------------------
 1 file changed, 29 insertions(+), 65 deletions(-)

diff --git a/tensor2tensor/models/mtf_resnet.py b/tensor2tensor/models/mtf_resnet.py
index e8eef3e75..430c46406 100644
--- a/tensor2tensor/models/mtf_resnet.py
+++ b/tensor2tensor/models/mtf_resnet.py
@@ -41,7 +41,8 @@ def batch_norm_relu(inputs, is_training, relu=True):
       inputs,
       is_training,
       BATCH_NORM_DECAY,
-      epsilon=BATCH_NORM_EPSILON)
+      epsilon=BATCH_NORM_EPSILON,
+      init_zero=(not relu))
   if relu:
     inputs = mtf.relu(inputs)
   return inputs
@@ -106,7 +107,7 @@ def bottleneck_block(inputs,
   inputs = batch_norm_relu(inputs, is_training)
 
   # Second conv block
-  filters2_dim = mtf.Dimension("filters2", filters)
+  filters2_dim = mtf.Dimension("filters2", 4*filters)
   kernel2 = mtf.get_variable(
       inputs.mesh, "kernel2", mtf.Shape(
           [filter_h_dim, filter_w_dim, filters1_dim, filters2_dim]))
@@ -131,15 +132,14 @@ def bottleneck_block(inputs,
       padding="SAME",
       h_blocks_dim=None, w_blocks_dim=col_blocks_dim)
 
-  inputs = batch_norm_relu(
-      inputs,
-      is_training,
-      relu=False)
+  # TODO(nikip): Althought the original resnet code has this batch norm, in our
+  # setup this is causing no gradients to be passed. Investigate further.
+  # inputs = batch_norm_relu(inputs, is_training, relu=True)
 
   # TODO(nikip): Maybe add residual with a projection?
   return mtf.relu(
-      inputs + mtf.rename_dimension(
-          shortcut, shortcut.shape.dims[-1].name, inputs.shape.dims[-1].name))
+      shortcut + mtf.rename_dimension(
+          inputs, inputs.shape.dims[-1].name, shortcut.shape.dims[-1].name))
 
 
 def block_layer(inputs,
@@ -185,8 +185,8 @@ def projection_shortcut(inputs, kernel):
         inputs,
         filters,
         is_training,
-        strides,
-        projection_shortcut,
+        strides=strides,
+        projection_shortcut=projection_shortcut,
         row_blocks_dim=row_blocks_dim,
         col_blocks_dim=col_blocks_dim)
 
@@ -234,11 +234,12 @@ def mtf_model_fn(self, features, mesh):
     filter_h_dim = mtf.Dimension("filter_height", 7)
     filter_w_dim = mtf.Dimension("filter_width", 7)
     filters = mtf.Dimension("filters", hparams.filter_sizes[0])
-    rows_dim = mtf.Dimension("rows_size", 32)
-    cols_dim = mtf.Dimension("cols_size", 96)
+    rows_dim = mtf.Dimension("rows_size", hparams.rows_size)
+    cols_dim = mtf.Dimension("cols_size", hparams.cols_size)
     row_blocks_dim = mtf.Dimension("row_blocks", hparams.row_blocks)
     col_blocks_dim = mtf.Dimension("col_blocks", hparams.col_blocks)
     classes_dim = mtf.Dimension("classes", 10)
+    channels_dim = mtf.Dimension("channels", 3)
     one_channel_dim = mtf.Dimension("one_channel", 1)
 
     inputs = features["inputs"]
@@ -248,17 +249,18 @@ def mtf_model_fn(self, features, mesh):
             hparams.row_blocks,
             hparams.rows_size // hparams.row_blocks,
             hparams.col_blocks,
-            hparams.num_channels*hparams.cols_size // hparams.col_blocks, 1]),
+            hparams.num_channels*hparams.cols_size // hparams.col_blocks,
+            hparams.num_channels]),
         mtf.Shape(
             [batch_dim, row_blocks_dim, rows_dim,
-             col_blocks_dim, cols_dim, one_channel_dim]))
+             col_blocks_dim, cols_dim, channels_dim]))
     x = mtf.transpose(x, [batch_dim, row_blocks_dim, col_blocks_dim,
-                          rows_dim, cols_dim, one_channel_dim])
+                          rows_dim, cols_dim, channels_dim])
 
     x = mtf.to_float(x)
     initial_filters = mtf.get_variable(
         mesh, "init_filters",
-        mtf.Shape([filter_h_dim, filter_w_dim, one_channel_dim, filters]))
+        mtf.Shape([filter_h_dim, filter_w_dim, channels_dim, filters]))
     x = mtf.conv2d_with_blocks(
         x,
         initial_filters,
@@ -269,7 +271,7 @@ def mtf_model_fn(self, features, mesh):
     x = batch_norm_relu(x, is_training)
 
     # Conv blocks
-    # [ self attention - ffn - residual + dropout] x n
+    # [block - strided block layer - strided block layer] x n
     for layer in range(hparams.num_layers):
       layer_name = "block_layer_%d" % layer
       with tf.variable_scope(layer_name):
@@ -287,7 +289,7 @@ def mtf_model_fn(self, features, mesh):
             inputs=x,
             filters=hparams.filter_sizes[1],
             blocks=hparams.layer_sizes[1],
-            strides=[1, 2, 2, 1],
+            strides=[1, 1, 1, 1],
             is_training=is_training,
             name="block_layer2",
             row_blocks_dim=None,
@@ -296,7 +298,7 @@ def mtf_model_fn(self, features, mesh):
             inputs=x,
             filters=hparams.filter_sizes[2],
             blocks=hparams.layer_sizes[2],
-            strides=[1, 2, 2, 1],
+            strides=[1, 1, 1, 1],
             is_training=is_training,
             name="block_layer3",
             row_blocks_dim=None,
@@ -340,22 +342,12 @@ def mtf_resnet_base():
   # 8-way model-parallelism
   hparams.add_hparam("mesh_shape", "batch:8")
   hparams.add_hparam("layout", "batch:batch")
-  hparams.add_hparam("num_heads", 8)
   hparams.add_hparam("filter_size", 1024)
 
   hparams.add_hparam("num_layers", 6)
-  hparams.add_hparam("attention_key_size", 256)
-  hparams.add_hparam("attention_value_size", 256)
   # Share weights between input and target embeddings
   hparams.shared_embedding = True
 
-  # mixture of experts hparams
-  hparams.add_hparam("ffn_layer", "dense_relu_dense")
-  hparams.add_hparam("moe_overhead_train", 1.0)
-  hparams.add_hparam("moe_overhead_eval", 2.0)
-  hparams.moe_num_experts = 16
-  hparams.moe_loss_coef = 1e-3
-
   hparams.shared_embedding_and_softmax_weights = True
   hparams.optimizer = "Adafactor"
   hparams.learning_rate_schedule = "rsqrt_decay"
@@ -380,7 +372,7 @@ def mtf_resnet_base():
   hparams.initializer_gain = 2.
 
   # TODO(nikip): Change optimization scheme?
-  hparams.learning_rate = 0.4
+  hparams.learning_rate = 0.1
   return hparams
 
 
@@ -389,13 +381,14 @@ def mtf_resnet_tiny():
   """Catch bugs locally..."""
   hparams = mtf_resnet_base()
   hparams.num_layers = 2
-  hparams.hidden_size = 128
-  hparams.filter_size = 256
-  hparams.batch_size = 2
+  hparams.hidden_size = 64
+  hparams.filter_size = 64
+  hparams.batch_size = 16
   # data parallelism and model-parallelism
-  hparams.mesh_shape = "all:2"
-  hparams.layout = "batch:all"
-  hparams.layer_sizes = [3, 2, 3]
+  hparams.col_blocks = 1
+  hparams.mesh_shape = "batch:2"
+  hparams.layout = "batch:batch"
+  hparams.layer_sizes = [1, 2, 3]
   hparams.filter_sizes = [64, 64, 64]
   return hparams
 
@@ -411,9 +404,6 @@ def mtf_resnet_single():
   hparams.batch_size = 1
   hparams.num_encoder_layers = 1
   hparams.num_layers = 1
-  hparams.num_heads = 2
-  hparams.attention_key_size = 32
-  hparams.attention_value_size = 32
   hparams.block_length = 16
   return hparams
 
@@ -437,7 +427,6 @@ def mtf_resnet_base_cifar():
   hparams.mesh_shape = "batch:32"
   hparams.layoyt = "batch:batch"
   hparams.batch_size = 8
-  hparams.num_heads = 4
   hparams.num_layers = 12
   hparams.block_length = 256
   hparams.hidden_size = 512
@@ -449,28 +438,3 @@ def mtf_resnet_base_cifar():
   hparams.layer_prepostprocess_dropout = 0.3
   hparams.unconditional = True
   return hparams
-
-
-@registry.register_hparams
-def mtf_resnet_tiny_moe():
-  hparams = mtf_resnet_tiny()
-  hparams.mesh_shape = "all:4"
-  hparams.layout = "batch:all,experts:all"
-  hparams.ffn_layer = "moe"
-  return hparams
-
-
-@registry.register_hparams
-def mtf_resnet_tiny_8gpu():
-  hparams = mtf_resnet_tiny()
-  hparams.mesh_shape = "all:8"
-  hparams.layout = "vocab:all;filter_size:all;heads:all"
-  return hparams
-
-
-@registry.register_hparams
-def mtf_resnet_length_sharded():
-  hparams = mtf_resnet_tiny()
-  hparams.mesh_shape = "all"
-  hparams.layout = "length:all"
-  return hparams

From 084c6a746d87e324f03e7ed1927322634417db11 Mon Sep 17 00:00:00 2001
From: Stefan Falk <43335432+stefan-falk@users.noreply.github.com>
Date: Tue, 23 Oct 2018 02:06:25 +0200
Subject: [PATCH 1057/2720] Passing hparams to train and eval hooks (#1158)

Introducing HookContextArgs

Renaming HookContextArgs to HookContext

Renaming HookContextArgs to HookContext
---
 tensor2tensor/utils/t2t_model.py   | 12 ++++++------
 tensor2tensor/utils/trainer_lib.py | 14 ++++++++++++--
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 9103f6cf3..0e16c6f8c 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -159,11 +159,11 @@ def create_hparams_summary(hparams, name):
   # Replace the two methods below in order to add custom SessionRunHooks to
   # the training procedure.
   @staticmethod
-  def train_hooks():
+  def train_hooks(hook_context):
     return []
 
   @staticmethod
-  def eval_hooks():
+  def eval_hooks(hook_context):
     return []
 
   @property
@@ -1215,14 +1215,14 @@ def _to_single_features_dict(self, datashard_features):
     return features
 
   @staticmethod
-  def get_train_hooks(model_name):
+  def get_train_hooks(model_name, hook_context):
     model_cls = registry.model(model_name)
-    return model_cls.train_hooks()
+    return model_cls.train_hooks(hook_context)
 
   @staticmethod
-  def get_eval_hooks(model_name):
+  def get_eval_hooks(model_name, hook_context):
     model_cls = registry.model(model_name)
-    return model_cls.eval_hooks()
+    return model_cls.eval_hooks(hook_context)
 
   @staticmethod
   def make_estimator_model_fn(model_name,
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 505bfaf94..43350d601 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import json
 import os
 import random
@@ -324,6 +325,12 @@ def create_hooks(use_tfdbg=False,
   return train_hooks, eval_hooks
 
 
+class HookContext(collections.namedtuple(
+  "HookContext",
+  ["estimator", "problem", "hparams"])):
+  pass
+
+
 class T2TExperiment(object):
   """Custom Experiment class for running distributed experiments."""
 
@@ -602,8 +609,11 @@ def compare_fn(best_eval_result, current_eval_result):
       validation_monitor_kwargs=validation_monitor_kwargs,
       use_early_stopping=use_early_stopping,
       early_stopping_kwargs=early_stopping_kwargs)
-  train_hooks += t2t_model.T2TModel.get_train_hooks(model_name)
-  eval_hooks += t2t_model.T2TModel.get_eval_hooks(model_name)
+
+  hook_context = HookContext(estimator=estimator, problem=problem, hparams=hparams)
+
+  train_hooks += t2t_model.T2TModel.get_train_hooks(model_name, hook_context)
+  eval_hooks += t2t_model.T2TModel.get_eval_hooks(model_name, hook_context)
   if additional_train_hooks:
     train_hooks += additional_train_hooks
   if additional_eval_hooks:

From 36dcf046545334034b5dbd1cc4140903e734fd2a Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 22 Oct 2018 16:26:50 -0700
Subject: [PATCH 1058/2720] internal

PiperOrigin-RevId: 218254388
---
 tensor2tensor/utils/t2t_model.py   | 12 ++++++------
 tensor2tensor/utils/trainer_lib.py | 14 ++------------
 2 files changed, 8 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 0e16c6f8c..9103f6cf3 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -159,11 +159,11 @@ def create_hparams_summary(hparams, name):
   # Replace the two methods below in order to add custom SessionRunHooks to
   # the training procedure.
   @staticmethod
-  def train_hooks(hook_context):
+  def train_hooks():
     return []
 
   @staticmethod
-  def eval_hooks(hook_context):
+  def eval_hooks():
     return []
 
   @property
@@ -1215,14 +1215,14 @@ def _to_single_features_dict(self, datashard_features):
     return features
 
   @staticmethod
-  def get_train_hooks(model_name, hook_context):
+  def get_train_hooks(model_name):
     model_cls = registry.model(model_name)
-    return model_cls.train_hooks(hook_context)
+    return model_cls.train_hooks()
 
   @staticmethod
-  def get_eval_hooks(model_name, hook_context):
+  def get_eval_hooks(model_name):
     model_cls = registry.model(model_name)
-    return model_cls.eval_hooks(hook_context)
+    return model_cls.eval_hooks()
 
   @staticmethod
   def make_estimator_model_fn(model_name,
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 43350d601..505bfaf94 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -19,7 +19,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import collections
 import json
 import os
 import random
@@ -325,12 +324,6 @@ def create_hooks(use_tfdbg=False,
   return train_hooks, eval_hooks
 
 
-class HookContext(collections.namedtuple(
-  "HookContext",
-  ["estimator", "problem", "hparams"])):
-  pass
-
-
 class T2TExperiment(object):
   """Custom Experiment class for running distributed experiments."""
 
@@ -609,11 +602,8 @@ def compare_fn(best_eval_result, current_eval_result):
       validation_monitor_kwargs=validation_monitor_kwargs,
       use_early_stopping=use_early_stopping,
       early_stopping_kwargs=early_stopping_kwargs)
-
-  hook_context = HookContext(estimator=estimator, problem=problem, hparams=hparams)
-
-  train_hooks += t2t_model.T2TModel.get_train_hooks(model_name, hook_context)
-  eval_hooks += t2t_model.T2TModel.get_eval_hooks(model_name, hook_context)
+  train_hooks += t2t_model.T2TModel.get_train_hooks(model_name)
+  eval_hooks += t2t_model.T2TModel.get_eval_hooks(model_name)
   if additional_train_hooks:
     train_hooks += additional_train_hooks
   if additional_eval_hooks:

From 3c9f0bc54b5e61db3a7a4930382176e5b6bb04b3 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 22 Oct 2018 17:10:29 -0700
Subject: [PATCH 1059/2720] Range for tuning the stochastic discrete model and
 corrections: make eval use RNN predictor, add dropout in non-residual
 connections and RNN predictor, clearer scheduled sampling function.

PiperOrigin-RevId: 218261002
---
 tensor2tensor/layers/common_video.py          |  7 ++-----
 tensor2tensor/layers/discretization.py        |  6 +++++-
 tensor2tensor/models/video/base.py            |  1 -
 .../models/video/basic_deterministic.py       |  5 +++--
 .../video/basic_deterministic_params.py       |  5 +++--
 .../models/video/basic_stochastic.py          | 21 +++++++++++++++----
 6 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 80b27c8f9..974cc4b83 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -205,11 +205,8 @@ def scheduled_sample_prob(ground_truth_x,
   """
   probability_threshold = scheduled_sample_var
   probability_of_generated = tf.random_uniform([batch_size])
-  array_ind = tf.to_int32(probability_of_generated > probability_threshold)
-  indices = tf.range(batch_size) + array_ind * batch_size
-  xy = tf.concat([ground_truth_x, generated_x], axis=0)
-  output = tf.gather(xy, indices)
-  return output
+  return tf.where(probability_of_generated > probability_threshold,
+                  generated_x, ground_truth_x)
 
 
 def dna_transformation(prev_image, dna_input, dna_kernel_size, relu_shift):
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 71a7cd988..165b01a28 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -784,7 +784,8 @@ def discrete_bottleneck(inputs,
 
 
 def predict_bits_with_lstm(prediction_source, state_size, total_num_bits,
-                           target_bits=None, bits_at_once=8, temperature=1.0):
+                           target_bits=None, bits_at_once=8, temperature=1.0,
+                           dropout=0.1):
   """Predict a sequence of bits (a latent) with LSTM, both training and infer.
 
   Given a tensor on which the predictions are based (prediction_source), we use
@@ -803,6 +804,7 @@ def predict_bits_with_lstm(prediction_source, state_size, total_num_bits,
       training as the target to predict; each element should be -1 or 1.
     bits_at_once: pytho integer, how many bits to predict at once.
     temperature: python float, temperature used for sampling during inference.
+    dropout: float, the amount of dropout to aply during training (0.1 default).
 
   Returns:
     a pair (bits, loss) with the predicted bit sequence, which is a Tensor of
@@ -845,12 +847,14 @@ def predict_bits_with_lstm(prediction_source, state_size, total_num_bits,
     tf.summary.histogram("target_integers", tf.reshape(d_int, [-1]))
     d_hot = tf.one_hot(d_int, 2**bits_at_once, axis=-1)
     d_pred = discrete_embed(d_hot)
+    d_pred = tf.nn.dropout(d_pred, 1.0 - dropout)
     pred = tf.concat([tf.expand_dims(prediction, axis=1), d_pred], axis=1)
     outputs = []
     for i in range(total_num_bits // bits_at_once):
       output, state = lstm_cell(pred[:, i, :], state)
       outputs.append(tf.expand_dims(output, axis=1))
     outputs = tf.concat(outputs, axis=1)
+    outputs = tf.nn.dropout(outputs, 1.0 - dropout)
     d_int_pred = discrete_predict(outputs)
     pred_loss = tf.losses.sparse_softmax_cross_entropy(
         logits=d_int_pred, labels=d_int)
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 583bf7a49..8ca6d1eb3 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -578,4 +578,3 @@ def next_frame_base():
   hparams.add_hparam("scheduled_sampling_max_prob", 1.0)
   hparams.add_hparam("scheduled_sampling_k", 900.0)
   return hparams
-
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 6fc3f2a95..655d9f878 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -52,7 +52,7 @@ def middle_network(self, layer, internal_states):
     filters = common_layers.shape_list(x)[-1]
     for i in range(self.hparams.num_hidden_layers):
       with tf.variable_scope("layer%d" % i):
-        y = tf.nn.dropout(x, 1.0 - self.hparams.dropout)
+        y = tf.nn.dropout(x, 1.0 - self.hparams.residual_dropout)
         y = tf.layers.conv2d(y, filters, kernel1, activation=common_layers.belu,
                              strides=(1, 1), padding="SAME")
         if i == 0:
@@ -83,6 +83,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
     for i in range(hparams.num_compress_steps):
       with tf.variable_scope("downstride%d" % i):
         layer_inputs.append(x)
+        x = tf.nn.dropout(x, 1.0 - self.hparams.dropout)
         x = common_layers.make_even_size(x)
         if i < hparams.filter_double_steps:
           filters *= 2
@@ -107,6 +108,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
     layer_inputs = list(reversed(layer_inputs))
     for i in range(hparams.num_compress_steps):
       with tf.variable_scope("upstride%d" % i):
+        x = tf.nn.dropout(x, 1.0 - self.hparams.dropout)
         if self.has_actions:
           x = common_video.inject_additional_input(
               x, action, "action_enc", hparams.action_injection)
@@ -139,4 +141,3 @@ def next_frame(self, frames, actions, rewards, target_frame,
         reward_pred, 128, name="reward_pred"))
     reward_pred = tf.expand_dims(reward_pred, axis=3)  # Need fake channels dim.
     return x, reward_pred, extra_loss, internal_states
-
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 9e65c93b7..7b2c3d0bc 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -40,7 +40,8 @@ def next_frame_basic_deterministic():
   hparams.initializer_gain = 1.3
   hparams.weight_decay = 0.0
   hparams.clip_grad_norm = 1.0
-  hparams.dropout = 0.5
+  hparams.dropout = 0.1
+  hparams.add_hparam("residual_dropout", 0.5)
   hparams.add_hparam("num_compress_steps", 6)
   hparams.add_hparam("filter_double_steps", 2)
   return hparams
@@ -60,7 +61,7 @@ def next_frame_sampling():
   """Basic conv model with scheduled sampling."""
   hparams = next_frame_basic_deterministic()
   hparams.scheduled_sampling_mode = "prob_inverse_exp"
-  hparams.scheduled_sampling_max_prob = 0.5
+  hparams.scheduled_sampling_max_prob = 1.0
   hparams.scheduled_sampling_decay_steps = 10000
   return hparams
 
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 9ef01d5cc..f7dfd176f 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -82,7 +82,7 @@ def add_bits(layer, bits):
       layer += z_add
       return layer
 
-    if self.is_predicting:
+    if not self.is_training:
       if hparams.full_latent_tower:
         rand = tf.random_uniform(layer_shape[:-1] + [hparams.bottleneck_bits])
         bits = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
@@ -174,9 +174,10 @@ def next_frame_basic_stochastic_discrete():
   hparams.video_num_target_frames = 16
   hparams.scheduled_sampling_mode = "prob_inverse_lin"
   hparams.scheduled_sampling_decay_steps = 40000
-  hparams.scheduled_sampling_prob = 1.0
-  hparams.learning_rate_constant = 0.01
-  hparams.learning_rate_warmup_steps = 8000
+  hparams.scheduled_sampling_max_prob = 1.0
+  hparams.dropout = 0.3
+  hparams.learning_rate_constant = 0.002
+  hparams.learning_rate_warmup_steps = 2000
   hparams.learning_rate_schedule = "linear_warmup * constant"
   hparams.add_hparam("bottleneck_bits", 64)
   hparams.add_hparam("bottleneck_noise", 0.02)
@@ -186,3 +187,15 @@ def next_frame_basic_stochastic_discrete():
   hparams.add_hparam("latent_predictor_temperature", 0.5)
   hparams.add_hparam("complex_addn", True)
   return hparams
+
+
+@registry.register_ranged_hparams
+def next_frame_stochastic_discrete_range(rhp):
+  """Next frame stochastic discrete tuning grid."""
+  rhp.set_float("learning_rate_constant", 0.001, 0.01)
+  rhp.set_float("dropout", 0.2, 0.6)
+  rhp.set_int("filter_double_steps", 3, 5)
+  rhp.set_discrete("hidden_size", [64, 96, 128])
+  rhp.set_discrete("bottleneck_bits", [32, 64, 128, 256])
+  rhp.set_discrete("video_num_target_frames", [4])
+  rhp.set_float("bottleneck_noise", 0.0, 0.2)

From 215220d1c1702b1ae5ee7720e187fc48c6f21bf1 Mon Sep 17 00:00:00 2001
From: Stefan Falk <43335432+stefan-falk@users.noreply.github.com>
Date: Mon, 22 Oct 2018 17:24:51 -0700
Subject: [PATCH 1060/2720] internal merge of PR #1158

PiperOrigin-RevId: 218262824
---
 tensor2tensor/utils/t2t_model.py   | 12 ++++++------
 tensor2tensor/utils/trainer_lib.py | 16 ++++++++++++++--
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 9103f6cf3..0e16c6f8c 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -159,11 +159,11 @@ def create_hparams_summary(hparams, name):
   # Replace the two methods below in order to add custom SessionRunHooks to
   # the training procedure.
   @staticmethod
-  def train_hooks():
+  def train_hooks(hook_context):
     return []
 
   @staticmethod
-  def eval_hooks():
+  def eval_hooks(hook_context):
     return []
 
   @property
@@ -1215,14 +1215,14 @@ def _to_single_features_dict(self, datashard_features):
     return features
 
   @staticmethod
-  def get_train_hooks(model_name):
+  def get_train_hooks(model_name, hook_context):
     model_cls = registry.model(model_name)
-    return model_cls.train_hooks()
+    return model_cls.train_hooks(hook_context)
 
   @staticmethod
-  def get_eval_hooks(model_name):
+  def get_eval_hooks(model_name, hook_context):
     model_cls = registry.model(model_name)
-    return model_cls.eval_hooks()
+    return model_cls.eval_hooks(hook_context)
 
   @staticmethod
   def make_estimator_model_fn(model_name,
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 505bfaf94..f505ed373 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import json
 import os
 import random
@@ -254,6 +255,7 @@ def create_estimator(model_name,
   model_fn = t2t_model.T2TModel.make_estimator_model_fn(
       model_name, hparams, decode_hparams=decode_hparams)
 
+
   del use_xla
   if use_tpu or use_tpu_estimator:
     problem = hparams.problem
@@ -324,6 +326,12 @@ def create_hooks(use_tfdbg=False,
   return train_hooks, eval_hooks
 
 
+class HookContext(collections.namedtuple(
+    "HookContext",
+    ["estimator", "problem", "hparams"])):
+  pass
+
+
 class T2TExperiment(object):
   """Custom Experiment class for running distributed experiments."""
 
@@ -602,8 +610,12 @@ def compare_fn(best_eval_result, current_eval_result):
       validation_monitor_kwargs=validation_monitor_kwargs,
       use_early_stopping=use_early_stopping,
       early_stopping_kwargs=early_stopping_kwargs)
-  train_hooks += t2t_model.T2TModel.get_train_hooks(model_name)
-  eval_hooks += t2t_model.T2TModel.get_eval_hooks(model_name)
+
+  hook_context = HookContext(
+      estimator=estimator, problem=problem, hparams=hparams)
+
+  train_hooks += t2t_model.T2TModel.get_train_hooks(model_name, hook_context)
+  eval_hooks += t2t_model.T2TModel.get_eval_hooks(model_name, hook_context)
   if additional_train_hooks:
     train_hooks += additional_train_hooks
   if additional_eval_hooks:

From 0c5c44f650930e90d9ed663bb8ae07f109493510 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Mon, 22 Oct 2018 17:27:30 -0700
Subject: [PATCH 1061/2720] Image transformer with local1d and local 2d spatial
 partitioning.

PiperOrigin-RevId: 218263114
---
 tensor2tensor/models/mtf_image_transformer.py | 260 +++++++++++++++---
 1 file changed, 223 insertions(+), 37 deletions(-)

diff --git a/tensor2tensor/models/mtf_image_transformer.py b/tensor2tensor/models/mtf_image_transformer.py
index 7317b7c8e..cfd72b28d 100644
--- a/tensor2tensor/models/mtf_image_transformer.py
+++ b/tensor2tensor/models/mtf_image_transformer.py
@@ -53,6 +53,10 @@ def targets_vocab_dim(self):
   def outputs_vocab_dim(self):
     return mtf.Dimension("output_vocab", 256)
 
+  @property
+  def pos_dim(self):
+    return mtf.Dimension("pos", self._hparams.img_len)
+
   @property
   def rows_dim(self):
     return mtf.Dimension("rows", self._hparams.img_len)
@@ -76,11 +80,15 @@ def model_dim(self):
 
   @property
   def max_length_dim(self):
-    return mtf.Dimension("max_length", self._hparams.max_length)
+    return mtf.Dimension(
+        "max_length",
+        self._hparams.img_len*self._hparams.img_len*self._hparams.num_channels)
 
   @property
   def length_dim(self):
-    return mtf.Dimension("length", self._hparams.max_length)
+    return mtf.Dimension(
+        "length",
+        self._hparams.img_len*self._hparams.img_len*self._hparams.num_channels)
 
   @property
   def heads_dim(self):
@@ -114,12 +122,12 @@ def create_positional_emb_2d(self, targets):
 
     positional_emb_rows_var = mtf.get_variable(
         mesh, "positional_emb_rows",
-        mtf.Shape([self.max_length_dim, self.model_dim]),
+        mtf.Shape([self.pos_dim, self.model_dim]),
         initializer=tf.random_normal_initializer(),
         activation_dtype=self.activation_type)
     positional_emb_cols_var = mtf.get_variable(
         mesh, "positional_emb_cols",
-        mtf.Shape([self.max_length_dim, self.model_dim]),
+        mtf.Shape([self.pos_dim, self.model_dim]),
         initializer=tf.random_normal_initializer(),
         activation_dtype=self.activation_type)
 
@@ -127,12 +135,12 @@ def create_positional_emb_2d(self, targets):
     targets_position_y = mtf.range(mesh, self.cols_dim, dtype=tf.int32)
     position_x = mtf.broadcast(
         mtf.gather(positional_emb_rows_var, targets_position_x,
-                   self.max_length_dim),
+                   self.pos_dim),
         mtf.Shape([self.rows_dim, self.cols_dim, self.model_dim]))
 
     position_y = mtf.broadcast(
         mtf.gather(positional_emb_cols_var, targets_position_y,
-                   self.max_length_dim),
+                   self.pos_dim),
         mtf.Shape([self.rows_dim, self.cols_dim, self.model_dim]))
     return position_x + position_y
 
@@ -157,11 +165,6 @@ def import_to_batch_by_length(x, name):
       return mtf.import_tf_tensor(
           mesh, x, mtf.Shape([batch_dim, self.length_dim]), name=name)
 
-    def layer_prepostprocess_dropout(x):
-      return mtf.dropout(
-          x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
-          noise_shape=mtf.Shape([batch_dim, self.model_dim]))
-
     targets = import_to_batch_by_length(targets, "targets")
     shifted_targets = import_to_batch_by_length(
         shifted_targets, "shifted_targets")
@@ -178,6 +181,7 @@ def layer_prepostprocess_dropout(x):
 
     x = mtf.gather(targets_embedding_var,
                    shifted_targets, self.targets_vocab_dim)
+
     # Add positional embeddings
     x += mtf.reshape(self.create_positional_emb_2d(targets),
                      [self.length_dim, self.model_dim])
@@ -199,29 +203,24 @@ def layer_prepostprocess_dropout(x):
 
     # Image Transformer Decoder
     # [ self attention - ffn - residual + dropout] x n
-    for layer in range(hparams.num_decoder_layers):
-      layer_name = "decoder_layer_%d" % layer
-      with tf.variable_scope(layer_name):
-        # Self attention layer
-        x += layer_prepostprocess_dropout(
-            mtf.layers.masked_local_attention_1d(
-                mtf.layers.layer_norm(x, self.model_dim, name="layer_norm_att"),
-                None,
-                self.kv_dim,
-                self.heads_dim,
-                block_length=hparams.block_length,
-                name="self_att"))
-        # ffn layer
-        x += layer_prepostprocess_dropout(mtf.layers.dense_relu_dense(
-            mtf.layers.layer_norm(x, self.model_dim, name="layer_norm_ffn"),
-            self.feedforward_dim,
-            hparams.dropout,
-            dropout_broadcast_dims=[self.length_dim]))
-
-    x = mtf.layers.layer_norm(x, self.model_dim, name="final_layer_norm")
+    if hparams.attention_type == "local1d_spatial":
+      decoder_output = local_attention1d_spatial_decoder(
+          x, self.kv_dim, self.heads_dim, self.feedforward_dim, hparams)
+    elif hparams.attention_type == "local2d_spatial":
+      decoder_output = local_attention2d_spatial_decoder(
+          x, self.kv_dim, self.heads_dim, self.feedforward_dim, hparams)
+    elif hparams.attention_type == "local1d":
+      decoder_output = local_attention1d_masked_decoder(
+          x, self.kv_dim, self.heads_dim, self.feedforward_dim, hparams)
+    else:
+      raise ValueError("Invalid attention type.")
 
     # Calculate the logits and loss.
-    logits = mtf.layers.dense(x, self.outputs_vocab_dim, name="logits")
+    logits = mtf.layers.dense(
+        decoder_output, self.outputs_vocab_dim, name="logits")
+    # Need a reshape for logits
+    logits = mtf.reshape(
+        logits, mtf.Shape([batch_dim, self.length_dim, self.outputs_vocab_dim]))
     soft_targets = mtf.one_hot(
         targets, self.outputs_vocab_dim, dtype=activation_dtype)
     loss = mtf.layers.softmax_cross_entropy_with_logits(
@@ -239,6 +238,126 @@ def layer_prepostprocess_dropout(x):
     return logits, loss
 
 
+def layer_prepostprocess_dropout(x, hparams):
+  batch_dim = x.shape.dims[0]
+  model_dim = x.shape.dims[-1]
+  return mtf.dropout(
+      x,
+      keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
+      noise_shape=mtf.Shape([batch_dim, model_dim]))
+
+
+def local_attention1d_spatial_decoder(x, kv_dim, heads_dim,
+                                      feedforward_dim, hparams):
+  """Image Transformer decoder with local1D spatial layers."""
+  batch_dim, length_dim, model_dim = x.shape.dims
+  blocks_w_dim = mtf.Dimension("blocksw", hparams.block_length)
+  num_w_blocks_dim = mtf.Dimension("num_wblocks",
+                                   length_dim.size // blocks_w_dim.size)
+  x = mtf.reshape(
+      x, mtf.Shape([batch_dim, num_w_blocks_dim, blocks_w_dim, model_dim]))
+  # [ self attention - ffn - residual + dropout] x n
+  for layer in range(hparams.num_decoder_layers):
+    layer_name = "decoder_layer_%d" % layer
+    with tf.variable_scope(layer_name):
+      # Self attention layer
+      x += layer_prepostprocess_dropout(
+          mtf.layers.local_self_attention_spatial_blocks(
+              mtf.layers.layer_norm(x, model_dim, name="layer_norm_att"),
+              kv_dim,
+              heads_dim,
+              memory_w_dim=blocks_w_dim,
+              mask_right=True,
+              name="self_att"), hparams)
+      # ffn layer
+      x += layer_prepostprocess_dropout(
+          mtf.layers.dense_relu_dense(
+              mtf.layers.layer_norm(x, model_dim, name="layer_norm_ffn"),
+              feedforward_dim,
+              hparams.dropout,
+              dropout_broadcast_dims=[length_dim]), hparams)
+
+  output = mtf.layers.layer_norm(x, model_dim, name="final_layer_norm")
+  return output
+
+
+def local_attention2d_spatial_decoder(x, kv_dim, heads_dim,
+                                      feedforward_dim, hparams):
+  """Image Transformer decoder with local2D spatial layers."""
+  batch_dim, length_dim, model_dim = x.shape.dims
+  blocks_h_dim = mtf.Dimension("blocksh", hparams.block_height)
+  blocks_w_dim = mtf.Dimension("blocksw", hparams.block_width)
+  num_h_blocks_dim = mtf.Dimension("num_h_blocks",
+                                   hparams.img_len // hparams.block_height)
+  num_w_blocks_dim = mtf.Dimension(
+      "num_w_blocks",
+      hparams.img_len * hparams.num_channels // hparams.block_width)
+  x = mtf.transpose(
+      mtf.reshape(
+          x,
+          mtf.Shape([
+              batch_dim, num_h_blocks_dim, blocks_h_dim,
+              num_w_blocks_dim, blocks_w_dim, model_dim
+          ])),
+      mtf.Shape([
+          batch_dim, num_h_blocks_dim, num_w_blocks_dim,
+          blocks_h_dim, blocks_w_dim, model_dim
+      ]))
+  # Image Transformer Decoder
+  # [ self attention - ffn - residual + dropout] x n
+  for layer in range(hparams.num_decoder_layers):
+    layer_name = "decoder_layer_%d" % layer
+    with tf.variable_scope(layer_name):
+      # Self attention layer
+      x += layer_prepostprocess_dropout(
+          mtf.layers.local_2d_self_attention_spatial_blocks(
+              mtf.layers.layer_norm(x, model_dim, name="layer_norm_att"),
+              kv_dim,
+              heads_dim,
+              memory_h_dim=num_h_blocks_dim,
+              memory_w_dim=num_w_blocks_dim,
+              name="self_att"), hparams)
+      # ffn layer
+      x += layer_prepostprocess_dropout(
+          mtf.layers.dense_relu_dense(
+              mtf.layers.layer_norm(x, model_dim, name="layer_norm_ffn"),
+              feedforward_dim,
+              hparams.dropout,
+              dropout_broadcast_dims=[length_dim]), hparams)
+
+  output = mtf.layers.layer_norm(x, model_dim, name="final_layer_norm")
+  return output
+
+
+def local_attention1d_masked_decoder(x, kv_dim, heads_dim,
+                                     feedforward_dim, hparams):
+  """Image Transformer decoder with local1D masked layers."""
+  print(x)
+  _, length_dim, model_dim = x.shape.dims
+  for layer in range(hparams.num_decoder_layers):
+    layer_name = "decoder_layer_%d" % layer
+    with tf.variable_scope(layer_name):
+      # Self attention layer
+      x += layer_prepostprocess_dropout(
+          mtf.layers.masked_local_attention_1d(
+              mtf.layers.layer_norm(x, model_dim, name="layer_norm_att"),
+              None,
+              kv_dim,
+              heads_dim,
+              block_length=hparams.block_length,
+              name="self_att"), hparams)
+      # ffn layer
+      x += layer_prepostprocess_dropout(
+          mtf.layers.dense_relu_dense(
+              mtf.layers.layer_norm(x, model_dim, name="layer_norm_ffn"),
+              feedforward_dim,
+              hparams.dropout,
+              dropout_broadcast_dims=[length_dim]), hparams)
+
+  output = mtf.layers.layer_norm(x, model_dim, name="final_layer_norm")
+  return output
+
+
 @registry.register_hparams
 def mtf_image_transformer_base():
   """Set of hyperparameters."""
@@ -280,7 +399,12 @@ def mtf_image_transformer_base():
   hparams.add_hparam("img_len", 32)
   hparams.add_hparam("num_channels", 3)
   hparams.add_hparam("unconditional", True)
+
+  # Local Attention related params
   hparams.add_hparam("block_length", 128)
+  hparams.add_hparam("block_height", 16)
+  hparams.add_hparam("block_width", 16)
+  hparams.add_hparam("attention_type", "local1d")
   return hparams
 
 
@@ -333,6 +457,34 @@ def mtf_image_transformer_base_single():
   return hparams
 
 
+@registry.register_hparams
+def mtf_image_transformer_tiny_spatial1d():
+  """Small single parameters."""
+  hparams = mtf_image_transformer_tiny()
+  hparams.num_decoder_layers = 6
+  hparams.filter_size = 128
+  hparams.block_height = 8
+  hparams.block_width = 8
+  hparams.attention_type = "local1d_spatial"
+  hparams.mesh_shape = ""
+  hparams.layout = ""
+  return hparams
+
+
+@registry.register_hparams
+def mtf_image_transformer_tiny_spatial2d():
+  """Small single parameters."""
+  hparams = mtf_image_transformer_tiny()
+  hparams.num_decoder_layers = 6
+  hparams.filter_size = 128
+  hparams.block_height = 8
+  hparams.block_width = 8
+  hparams.attention_type = "local2d_spatial"
+  hparams.mesh_shape = "b1:2,b2:2"
+  hparams.layout = "num_h_blocks:b1,num_wblocks:b2"
+  return hparams
+
+
 @registry.register_hparams
 def mtf_image_transformer_base_cifar():
   """Data parallel CIFAR parameters."""
@@ -410,11 +562,45 @@ def mtf_image_transformer_base_imagenet_mp():
 
 
 @registry.register_hparams
-def mtf_image_transformer_tiny_moe():
-  hparams = mtf_image_transformer_tiny()
-  hparams.mesh_shape = "all:4"
-  hparams.layout = "batch:all,experts:all"
-  hparams.ffn_layer = "moe"
+def mtf_image_transformer_base_imagenet_mp128():
+  """Model parallel ImageNet parameters."""
+  hparams = mtf_image_transformer_base_imagenet()
+  hparams.mesh_shape = "model:8;batch:4"
+  hparams.layout = "batch:batch;d_ff:model;heads:model"
+  hparams.batch_size = 8
+  hparams.img_len = 128
+  hparams.block_length = 128
+  hparams.num_heads = 8
+  hparams.num_decoder_layers = 4
+  hparams.d_ff = 4096
+  hparams.learning_rate_warmup_steps = 31250
+  hparams.unconditional = True
+  hparams.max_length = 256*256*3
+  return hparams
+
+
+@registry.register_hparams
+def mtf_image_transformer_base_imagenet_mp_sp():
+  """Model parallel ImageNet parameters."""
+  hparams = mtf_image_transformer_base_imagenet_mp128()
+  hparams.mesh_shape = "model:8;batch:4"
+  hparams.layout = "batch:batch;d_ff:model;num_wblocks:model"
+  hparams.batch_size = 8
+  hparams.img_len = 128
+  hparams.block_length = 128
+  hparams.attention_type = "local1d_spatial"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_image_transformer_base_imagenet_mp64():
+  """Model parallel ImageNet parameters."""
+  hparams = mtf_image_transformer_base_imagenet()
+  hparams.mesh_shape = "model:8;batch:4"
+  hparams.layout = "batch:batch;d_ff:model;heads:model"
+  hparams.batch_size = 8
+  hparams.img_len = 64
+  hparams.num_decoder_layers = 8
   return hparams
 
 

From 33648711077f2572c3e50eddfa6ad3540a09b227 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 22 Oct 2018 19:21:10 -0700
Subject: [PATCH 1062/2720] Transformer MLPerf Compliance.

PiperOrigin-RevId: 218274535
---
 tensor2tensor/bin/t2t_trainer.py   | 2 ++
 tensor2tensor/utils/optimize.py    | 3 +++
 tensor2tensor/utils/trainer_lib.py | 9 ++++++++-
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 0f557d560..53413c2f1 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -347,6 +347,8 @@ def main(argv):
   mlperf_log.transformer_print(key=mlperf_log.RUN_START)
   if FLAGS.schedule == "run_std_server":
     run_std_server()
+  mlperf_log.transformer_print(
+      key=mlperf_log.RUN_SET_RANDOM_SEED, value=FLAGS.random_seed)
   trainer_lib.set_random_seed(FLAGS.random_seed)
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
   maybe_log_registry_and_exit()
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 47b3ad983..2cf491257 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -252,6 +252,9 @@ def get_variable_initializer(hparams):
   if not hparams.initializer:
     return None
 
+  mlperf_log.transformer_print(key=mlperf_log.MODEL_HP_INITIALIZER_GAIN,
+                               value=hparams.initializer_gain)
+
   if not tf.contrib.eager.in_eager_mode():
     tf.logging.info("Using variable initializer: %s", hparams.initializer)
   if hparams.initializer == "orthogonal":
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index f505ed373..955449d98 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -486,8 +486,15 @@ def continuous_decode_on_train_data(self):
 
   def continuous_decode_on_eval_data(self):
     """Decode from dataset on new checkpoint."""
-    for _ in next_checkpoint(self._hparams.model_dir):
+    for ckpt in next_checkpoint(self._hparams.model_dir):
+      current_step = int(os.path.basename(ckpt).split("-")[1])
+      # Skip checkpoint 0.
+      if current_step == 0:
+        continue
+      mlperf_log.transformer_print(key=mlperf_log.EVAL_START)
       self.decode(dataset_split=tf.estimator.ModeKeys.EVAL)
+      mlperf_log.transformer_print(key=mlperf_log.EVAL_TARGET, value=25.0)
+      mlperf_log.transformer_print(key=mlperf_log.EVAL_STOP)
 
   def continuous_decode_from_file(self):
     """Decode from file on new checkpoint."""

From b15137506abf720715cfc81ddd5d80b80760eb66 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 22 Oct 2018 22:33:47 -0700
Subject: [PATCH 1063/2720] support multivariate actions.

PiperOrigin-RevId: 218288675
---
 tensor2tensor/models/video/epva.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index 0bd153c55..245d9bc81 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -658,7 +658,6 @@ def body(self, features):
 
     all_frames = tf.unstack(all_frames, axis=0)
     all_actions = tf.unstack(all_actions, axis=0)
-    all_actions = [tf.squeeze(a, 1) for a in all_actions]
 
     # TODO(blazej) - most likely this downsize is too strong.
     all_frames = [

From faed5fb1ed62e981e8d3b2bd534785798e60e849 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 23 Oct 2018 08:44:30 -0700
Subject: [PATCH 1064/2720] Use topk_with_unique for TPU based beam decoding.

PiperOrigin-RevId: 218347622
---
 tensor2tensor/utils/beam_search.py | 144 ++++++++++++++++++++++++++++-
 1 file changed, 141 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index 1d72c88c4..fe51422a2 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -18,8 +18,11 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-from tensor2tensor.layers import common_layers
 
+import math
+import numpy as np
+
+from tensor2tensor.layers import common_layers
 import tensorflow as tf
 
 from tensorflow.python.ops import inplace_ops
@@ -162,6 +165,136 @@ def _gather(params, indices):
     return gather_result
 
 
+def _create_make_unique(inputs):
+  """Replaces the lower bits of each element with iota.
+
+  The iota is used to derive the index, and also serves the purpose to
+  make each element unique to break ties.
+
+  Args:
+    inputs: A tensor with rank of 2 and dtype of tf.float32.
+      [batch_size, original_size].
+
+  Returns:
+    A tensor after element wise transformation, with dtype the same as inputs.
+    [batch_size, original_size].
+
+  Raises:
+    ValueError: If the rank of the input tensor does not equal 2.
+  """
+  if inputs.shape.ndims != 2:
+    raise ValueError("Input of top_k_with_unique must be rank-2 "
+                     "but got: %s" % inputs.shape)
+
+  height = inputs.shape[0]
+  width = inputs.shape[1]
+  zeros = tf.zeros([height, width], dtype=tf.int32)
+
+  # Count_mask is used to mask away the low order bits to ensure that every
+  # element is distinct.
+  log2_ceiling = int(math.ceil(math.log(int(width), 2)))
+  next_power_of_two = 1 << log2_ceiling
+  count_mask = ~(next_power_of_two - 1)
+  count_mask_r0 = tf.constant(count_mask)
+  count_mask_r2 = tf.fill([height, width], count_mask_r0)
+
+  # Smallest_normal is the bit representation of the smallest positive normal
+  # floating point number. The sign is zero, exponent is one, and the fraction
+  # is zero.
+  smallest_normal = 1 << 23
+  smallest_normal_r0 = tf.constant(smallest_normal, dtype=tf.int32)
+  smallest_normal_r2 = tf.fill([height, width], smallest_normal_r0)
+
+  # Low_bit_mask is used to mask away the sign bit when computing the absolute
+  # value.
+  low_bit_mask = ~(1 << 31)
+  low_bit_mask_r0 = tf.constant(low_bit_mask, dtype=tf.int32)
+  low_bit_mask_r2 = tf.fill([height, width], low_bit_mask_r0)
+
+  iota = tf.tile(tf.expand_dims(tf.range(width, dtype=tf.int32), 0),
+                 [height, 1])
+
+  # Compare the absolute value with positive zero to handle negative zero.
+  input_r2 = tf.bitcast(inputs, tf.int32)
+  abs_r2 = tf.bitwise.bitwise_and(input_r2, low_bit_mask_r2)
+  if_zero_r2 = tf.equal(abs_r2, zeros)
+  smallest_normal_preserving_sign_r2 = tf.bitwise.bitwise_or(
+      input_r2, smallest_normal_r2)
+  input_no_zeros_r2 = tf.where(
+      if_zero_r2, smallest_normal_preserving_sign_r2, input_r2)
+
+  # Discard the low-order bits and replace with iota.
+  and_r2 = tf.bitwise.bitwise_and(input_no_zeros_r2, count_mask_r2)
+  or_r2 = tf.bitwise.bitwise_or(and_r2, iota)
+  return tf.bitcast(or_r2, tf.float32)
+
+
+def _create_topk_unique(inputs, k):
+  """Creates the top k values in sorted order with indices.
+
+  Args:
+    inputs: A tensor with rank of 2. [batch_size, original_size].
+    k: An integer, number of top elements to select.
+
+  Returns:
+    topk_r2: A tensor, the k largest elements. [batch_size, k].
+    topk_indices_r2: A tensor, indices of the top k values. [batch_size, k].
+  """
+  height = inputs.shape[0]
+  width = inputs.shape[1]
+  neg_inf_r0 = tf.constant(-np.inf, dtype=tf.float32)
+  ones = tf.ones([height, width], dtype=tf.float32)
+  neg_inf_r2 = ones * neg_inf_r0
+  inputs = tf.where(tf.is_nan(inputs), neg_inf_r2, inputs)
+
+  # Select the current largest value k times and keep them in topk_r2. The
+  # selected largest values are marked as the smallest value to avoid being
+  # selected again.
+  tmp = inputs
+  topk_r2 = tf.zeros([height, k], dtype=tf.float32)
+  for i in range(k):
+    kth_order_statistic = tf.reduce_max(tmp, axis=1, keepdims=True)
+    k_mask = tf.tile(tf.expand_dims(tf.equal(tf.range(k), tf.fill([k], i)), 0),
+                     [height, 1])
+    topk_r2 = tf.where(k_mask, tf.tile(kth_order_statistic, [1, k]), topk_r2)
+    ge_r2 = tf.greater_equal(inputs, tf.tile(kth_order_statistic, [1, width]))
+    tmp = tf.where(ge_r2, neg_inf_r2, inputs)
+
+  log2_ceiling = int(math.ceil(math.log(float(int(width)), 2)))
+  next_power_of_two = 1 << log2_ceiling
+  count_mask = next_power_of_two - 1
+  mask_r0 = tf.constant(count_mask)
+  mask_r2 = tf.fill([height, k], mask_r0)
+  topk_r2_s32 = tf.bitcast(topk_r2, tf.int32)
+  topk_indices_r2 = tf.bitwise.bitwise_and(topk_r2_s32, mask_r2)
+  return topk_r2, topk_indices_r2
+
+
+def top_k_with_unique(inputs, k):
+  """Finds the values and indices of the k largests entries.
+
+  Instead of doing sort like tf.nn.top_k, this function finds the max value
+  k times. The running time is proportional to k, which is be faster when k
+  is small. The current implementation supports only inputs of rank 2.
+  In addition, iota is used to replace the lower bits of each element, this
+  makes the selection more stable when there are equal elements. The
+  overhead is that output values are approximated.
+
+  Args:
+    inputs: A tensor with rank of 2. [batch_size, original_size].
+    k: An integer, number of top elements to select.
+
+  Returns:
+    top_values: A tensor, the k largest elements in sorted order.
+      [batch_size, k].
+    indices: A tensor, indices of the top_values. [batch_size, k].
+  """
+  unique_inputs = _create_make_unique(tf.cast(inputs, tf.float32))
+  top_values, indices = _create_topk_unique(unique_inputs, k)
+  top_values = tf.cast(top_values, inputs.dtype)
+  return top_values, indices
+
+
 def compute_topk_scores_and_seq(sequences, scores, scores_to_gather, flags,
                                 beam_size, batch_size, prefix="default",
                                 states_to_gather=None, use_tpu=False):
@@ -201,8 +334,8 @@ def compute_topk_scores_and_seq(sequences, scores, scores_to_gather, flags,
      topk_gathered_scores [batch_size, beam_size],
      topk_finished_flags[batch_size, beam_size])
   """
-  _, topk_indexes = tf.nn.top_k(scores, k=beam_size)
   if not use_tpu:
+    _, topk_indexes = tf.nn.top_k(scores, k=beam_size)
     # The next three steps are to create coordinates for tf.gather_nd to pull
     # out the topk sequences from sequences based on scores.
     # batch pos is a tensor like [[0,0,0,0,],[1,1,1,1],..]. It says which
@@ -229,6 +362,7 @@ def gather(tensor, name):
     else:
       topk_gathered_states = states_to_gather
   else:
+    _, topk_indexes = top_k_with_unique(scores, k=beam_size)
     # Gather up the highest scoring sequences.  For each operation added, give
     # it a concrete name to simplify observing these operations with tfdbg.
     # Clients can capture these tensors by watching these node names.
@@ -454,7 +588,11 @@ def grow_topk(i, alive_seq, alive_log_probs, states):
     # Flatten out (beam_size, vocab_size) probs in to a list of possibilities
     flat_curr_scores = tf.reshape(curr_scores, [-1, beam_size * vocab_size])
 
-    topk_scores, topk_ids = tf.nn.top_k(flat_curr_scores, k=beam_size * 2)
+    if use_tpu:
+      topk_scores, topk_ids = top_k_with_unique(
+          flat_curr_scores, k=beam_size * 2)
+    else:
+      topk_scores, topk_ids = tf.nn.top_k(flat_curr_scores, k=beam_size * 2)
 
     # Recovering the log probs because we will need to send them back
     topk_log_probs = topk_scores * length_penalty

From f1c07efe0e169a6d7e1e35c60c39f3e4daba2794 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 23 Oct 2018 12:33:48 -0700
Subject: [PATCH 1065/2720] Scale up glow and friends, part 1 of N.

PiperOrigin-RevId: 218387770
---
 tensor2tensor/models/research/glow.py         | 58 +++++++++++++++++--
 .../models/research/glow_init_hook.py         | 38 ++++++++++++
 tensor2tensor/models/research/glow_ops.py     | 15 +++--
 tensor2tensor/models/research/glow_test.py    | 25 ++++++--
 4 files changed, 117 insertions(+), 19 deletions(-)
 create mode 100644 tensor2tensor/models/research/glow_init_hook.py

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 8c5173d93..7511a7cc2 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -22,6 +22,7 @@
 import numpy as np
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.models.research import glow_init_hook
 from tensor2tensor.models.research import glow_ops
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
@@ -50,6 +51,11 @@ def glow_hparams():
   hparams.add_hparam("depth", 32)
   hparams.add_hparam("affine_coupling_width", 512)
   hparams.add_hparam("top_prior", "single_conv")
+  # init_batch_size denotes the number of examples used for data-dependent
+  # initialization. A higher init_batch_size is required for training
+  # stability especially when hparams.batch_size is low.
+  # -1 indicates that the init_batch_size is set to hparams.batch_size.
+  hparams.add_hparam("init_batch_size", 256)
   return hparams
 
 
@@ -105,6 +111,33 @@ def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
 
     return self.scale(predictions)
 
+  def create_init_batch(self, features):
+    """Returns a batch of size "hparams.init_batch_size" for initialization.
+
+    Args:
+      features: input features.
+    Returns:
+      init_features: initialization features.
+    """
+    # TODO(mechcoder) Once all depending code supports hparams.init_batch_size
+    # the if block can be removed.
+    if self.hparams.init_batch_size == -1:
+      return features
+
+    train_dataset = self.hparams.problem.dataset(tf.estimator.ModeKeys.TRAIN)
+    train_dataset = train_dataset.batch(self.init_batch_size)
+    return train_dataset.make_one_shot_iterator().get_next()
+
+  @property
+  def init_batch_size(self):
+    if self.hparams.init_batch_size == -1:
+      return self.hparams.batch_size
+    return self.hparams.init_batch_size
+
+  @staticmethod
+  def train_hooks():
+    return [glow_init_hook.GlowInitHook()]
+
   def top_prior(self):
     """Objective based on the prior over latent z.
 
@@ -115,6 +148,24 @@ def top_prior(self):
         "top_prior", self.z_top_shape, learn_prior=self.hparams.top_prior)
 
   def body(self, features):
+    init_features = self.create_init_batch(features)
+    init_op = self.objective_tower(init_features, init=True)
+    init_op = tf.Print(
+        init_op, [init_op], message="Triggering data-dependent init.",
+        first_n=20)
+    tf.add_to_collection("glow_init_op", init_op)
+    train_op = self.objective_tower(features, init=False)
+    return tf.zeros_like(features["targets"]), {"training": train_op}
+
+  def objective_tower(self, features, init=True):
+    """Objective in terms of bits-per-pixel.
+
+    Args:
+      features: dict of tensors with "features" and "targets" keys.
+      init: Whether or not to run data-dependent init.
+    Returns:
+      objective: float, bits-per-pixel.
+    """
     x = features["inputs"]
 
     # Scale x such that the pixels lie in-between -0.5 and.0.5
@@ -125,10 +176,8 @@ def body(self, features):
     # the per-channel output activations have zero mean and unit variance
     # ONLY during the first step. After that the parameters are learned
     # through optimisation.
-    global_step = tf.train.get_or_create_global_step()
-    init_op = tf.logical_and(tf.equal(global_step, 0), self.is_training)
     ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
-    with arg_scope(ops, init=init_op):
+    with arg_scope(ops, init=init):
       self.z, encoder_objective, self.eps, _, _ = glow_ops.encoder_decoder(
           "codec", x, self.hparams, eps=None, reverse=False)
       objective += encoder_objective
@@ -137,11 +186,10 @@ def body(self, features):
       prior_dist = self.top_prior()
       prior_objective = tf.reduce_sum(
           prior_dist.log_prob(self.z), axis=[1, 2, 3])
-      tf.summary.scalar("top_prior", tf.reduce_mean(prior_objective))
       self.z_sample = prior_dist.sample()
       objective += prior_objective
 
     # bits per pixel
     _, h, w, c = common_layers.shape_list(x)
     objective = -objective / (np.log(2) * h * w * c)
-    return tf.zeros_like(features["targets"]), {"training": objective}
+    return objective
diff --git a/tensor2tensor/models/research/glow_init_hook.py b/tensor2tensor/models/research/glow_init_hook.py
new file mode 100644
index 000000000..2171253ff
--- /dev/null
+++ b/tensor2tensor/models/research/glow_init_hook.py
@@ -0,0 +1,38 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Hook to run glow initialization on a larger batch."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+
+class GlowInitHook(tf.train.SessionRunHook):
+  """
+  Hook that runs data-dependent initialization once before the first step.
+
+  The init op is stored in the tf collection glow_init_op. Look at the
+  "body" in glow.py for more details.
+  """
+
+  def after_create_session(self, session, coord):
+    del coord
+    global_step = session.run(tf.train.get_global_step())
+    if global_step == 0:
+      ddi = tf.get_collection("glow_init_op")
+      session.run(ddi)
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index c0578bdae..75a06a0e3 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -79,11 +79,14 @@ def check_cond_latents(cond_latents, hparams):
 def get_variable_ddi(name, shape, initial_value, dtype=tf.float32, init=False,
                      trainable=True):
   """Wrapper for data-dependent initialization."""
-  # Cast from python bool to TF bool for usage in tf.cond
-  if isinstance(init, bool):
-    init = tf.constant(init, dtype=tf.bool)
+  # If init is a tensor bool, w is returned dynamically.
   w = tf.get_variable(name, shape, dtype, None, trainable=trainable)
-  return tf.cond(init, lambda: assign(w, initial_value), lambda: w)
+  if isinstance(init, bool):
+    if init:
+      return assign(w, initial_value)
+    return w
+  else:
+    return tf.cond(init, lambda: assign(w, initial_value), lambda: w)
 
 
 @add_arg_scope
@@ -554,7 +557,6 @@ def merge_level_and_latent_dist(level_dist, latent_dist,
     scale = level_std
   elif merge_std == "prev_step":
     scale = latent_std
-  tf.summary.scalar("latent_scale", tf.reduce_mean(scale))
   return tf.distributions.Normal(loc=new_mean, scale=scale)
 
 
@@ -609,8 +611,6 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
     if latent_skip:
       cond_dist = tf.distributions.Normal(
           cond_dist.loc + latent, cond_dist.scale)
-  tf.summary.histogram("split_prior_mean", prior_dist.loc)
-  tf.summary.histogram("split_prior_scale", prior_dist.scale)
   return cond_dist.loc, cond_dist.scale, state
 
 
@@ -777,7 +777,6 @@ def scale_gaussian_prior(name, z, logscale_factor=3.0, trainable=True):
         "log_scale_latent", shape=z_shape, dtype=tf.float32,
         initializer=tf.zeros_initializer(), trainable=trainable)
     log_scale = log_scale * logscale_factor
-    tf.summary.scalar("gaussian_log_scale", tf.reduce_mean(log_scale))
     return tf.distributions.Normal(
         loc=latent_multiplier * z, scale=tf.exp(log_scale))
 
diff --git a/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
index e47f5df7e..1d623e018 100644
--- a/tensor2tensor/models/research/glow_test.py
+++ b/tensor2tensor/models/research/glow_test.py
@@ -46,8 +46,11 @@ def test_glow(self):
       hparams = glow.glow_hparams()
       hparams.depth = 15
       hparams.n_levels = 2
-      model = glow.Glow(hparams, tf.estimator.ModeKeys.TRAIN)
+      hparams.init_batch_size = 256
+      hparams.batch_size = 1
       cifar_problem = problems.problem('image_cifar10_plain_random_shift')
+      hparams.problem = cifar_problem
+      model = glow.Glow(hparams, tf.estimator.ModeKeys.TRAIN)
       train_dataset = cifar_problem.dataset(MODES.TRAIN)
       one_shot = train_dataset.make_one_shot_iterator()
       x_batch, y_batch = self.batch(one_shot)
@@ -56,13 +59,18 @@ def test_glow(self):
       objective = obj_dict['training']
       with tf.Session() as sess:
         sess.run(tf.global_variables_initializer())
+
+        # Run initialization.
+        init_op = tf.get_collection('glow_init_op')
+        sess.run(init_op)
+
+        # Run forward pass.
         obj_np = sess.run(objective)
         mean_obj = np.mean(obj_np)
 
         # Check that one forward-propagation does not NaN, i.e
         # initialization etc works as expected.
-        is_undefined = np.isnan(mean_obj) or np.isinf(mean_obj)
-        self.assertTrue(not is_undefined)
+        self.assertTrue(mean_obj > 0 and mean_obj < 10.0)
 
   def test_glow_inference(self):
     hparams = glow.glow_hparams()
@@ -72,18 +80,22 @@ def test_glow_inference(self):
 
     # Training pipeline
     with tf.Graph().as_default():
-      model = glow.Glow(hparams, tf.estimator.ModeKeys.TRAIN)
       cifar_problem = problems.problem('image_cifar10_plain_random_shift')
+      hparams.problem = cifar_problem
+      model = glow.Glow(hparams, tf.estimator.ModeKeys.TRAIN)
       train_dataset = cifar_problem.dataset(MODES.TRAIN)
       one_shot = train_dataset.make_one_shot_iterator()
       x_batch, y_batch = self.batch(one_shot)
       features = {'inputs': x_batch, 'targets': y_batch}
       model_path = os.path.join(curr_dir, 'model')
-
       model(features)
+
       with tf.Session() as session:
         saver = tf.train.Saver()
         session.run(tf.global_variables_initializer())
+
+        init_op = tf.get_collection('glow_init_op')
+        session.run(init_op)
         z = session.run([model.z])
         mean_z = np.mean(z)
         is_undefined = np.isnan(mean_z) or np.isinf(mean_z)
@@ -92,8 +104,9 @@ def test_glow_inference(self):
 
     # Inference pipeline
     with tf.Graph().as_default():
-      model = glow.Glow(hparams, tf.estimator.ModeKeys.PREDICT)
       cifar_problem = problems.problem('image_cifar10_plain_random_shift')
+      hparams.problem = cifar_problem
+      model = glow.Glow(hparams, tf.estimator.ModeKeys.PREDICT)
       test_dataset = cifar_problem.dataset(MODES.EVAL)
       one_shot = test_dataset.make_one_shot_iterator()
       x_batch, y_batch = self.batch(one_shot)

From 0ea0f5cb69ed64fa0db2301c73603777e5455fe4 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 23 Oct 2018 17:41:45 -0700
Subject: [PATCH 1066/2720] Small corrections to RL hparams and logging.

PiperOrigin-RevId: 218437988
---
 tensor2tensor/models/video/basic_recurrent.py  |  2 ++
 tensor2tensor/models/video/basic_stochastic.py | 10 ++++++----
 tensor2tensor/utils/decoding.py                |  6 ++++--
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/video/basic_recurrent.py b/tensor2tensor/models/video/basic_recurrent.py
index 8c14a8c8a..41ad73bf1 100644
--- a/tensor2tensor/models/video/basic_recurrent.py
+++ b/tensor2tensor/models/video/basic_recurrent.py
@@ -58,6 +58,8 @@ def middle_network(self, layer, internal_states):
 def next_frame_basic_recurrent():
   """Basic 2-frame recurrent model with stochastic tower."""
   hparams = basic_stochastic.next_frame_basic_stochastic_discrete()
+  hparams.filter_double_steps = 2
+  hparams.hidden_size = 64
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 4
   hparams.add_hparam("num_lstm_layers", 2)
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index f7dfd176f..d85f54b2d 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -175,12 +175,14 @@ def next_frame_basic_stochastic_discrete():
   hparams.scheduled_sampling_mode = "prob_inverse_lin"
   hparams.scheduled_sampling_decay_steps = 40000
   hparams.scheduled_sampling_max_prob = 1.0
-  hparams.dropout = 0.3
-  hparams.learning_rate_constant = 0.002
+  hparams.dropout = 0.2
+  hparams.filter_double_steps = 3
+  hparams.hidden_size = 96
+  hparams.learning_rate_constant = 0.005
   hparams.learning_rate_warmup_steps = 2000
   hparams.learning_rate_schedule = "linear_warmup * constant"
-  hparams.add_hparam("bottleneck_bits", 64)
-  hparams.add_hparam("bottleneck_noise", 0.02)
+  hparams.add_hparam("bottleneck_bits", 256)
+  hparams.add_hparam("bottleneck_noise", 0.15)
   hparams.add_hparam("discretize_warmup_steps", 40000)
   hparams.add_hparam("full_latent_tower", False)
   hparams.add_hparam("latent_predictor_state_size", 128)
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 31437d3c7..da830ee99 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -416,8 +416,10 @@ def timer(gen):
     total_time_per_step += elapsed_time
     total_cnt += result["outputs"].shape[-1]
   tf.logging.info("Elapsed Time: %5.5f" % (time.time() - start_time))
-  tf.logging.info("Averaged Single Token Generation Time: %5.7f" %
-                  (total_time_per_step / total_cnt))
+  tf.logging.info("Averaged Single Token Generation Time: %5.7f "
+                  "(time %5.7f count %d)" %
+                  (total_time_per_step / total_cnt,
+                   total_time_per_step, total_cnt))
 
   # Reversing the decoded inputs and outputs because they were reversed in
   # _decode_batch_input_fn

From dd6eafec338053dcbf8c44897afbc0909f61f1f7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 24 Oct 2018 12:44:34 -0700
Subject: [PATCH 1067/2720] Add record shuffling to Problem.dataset and batch
 shuffling.

PiperOrigin-RevId: 218558666
---
 tensor2tensor/data_generators/problem.py | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index f86eb6ffe..48b264323 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -555,6 +555,7 @@ def dataset(self,
               shard=None,
               partition_id=0,
               num_partitions=1,
+              shuffle_buffer_size=1024,
               max_records=-1,
               only_last=False):
     """Build a Dataset for this problem.
@@ -577,6 +578,8 @@ def dataset(self,
       shard: int, if provided, will only read data from the specified shard.
       partition_id: integer - which partition of the dataset to read from
       num_partitions: how many partitions in the dataset
+      shuffle_buffer_size: if shuffle_files is True, this is the buffer size
+        used to shuffle records.
       max_records: int, number of records to truncate to.
       only_last: bool, whether we should include only files from last epoch.
 
@@ -654,6 +657,10 @@ def _load_records_and_preprocess(filenames):
     dataset = dataset.map(
         self.maybe_reverse_and_copy, num_parallel_calls=num_threads)
     dataset = dataset.take(max_records)
+
+    ## Shuffle records only for training examples.
+    if shuffle_files and is_training:
+      dataset = dataset.shuffle(shuffle_buffer_size)
     if output_buffer_size:
       dataset = dataset.prefetch(output_buffer_size)
 
@@ -800,7 +807,8 @@ def input_fn(self,
                config=None,
                force_repeat=False,
                prevent_repeat=False,
-               dataset_kwargs=None):
+               dataset_kwargs=None,
+               batch_shuffle_size=512):
     """Builds input pipeline for problem.
 
     Args:
@@ -815,6 +823,8 @@ def input_fn(self,
         Overrides force_repeat.
       dataset_kwargs: dict, if passed, will pass as kwargs to self.dataset
         method when called
+      batch_shuffle_size: int, the size of the buffer to shuffle batches.
+        if none, the batches will not be shuffled.
 
     Returns:
       (features_dict<str name, Tensor feature>, Tensor targets)
@@ -956,6 +966,15 @@ def define_shapes(example):
 
     dataset = dataset.map(define_shapes, num_parallel_calls=num_threads)
 
+    # Add shuffling for training batches. This is necessary along with record
+    # level shuffling in the dataset generation. Record shuffling will shuffle
+    # the examples. However, in some cases, it's possible that the shuffle
+    # buffer size for record shuffling is smaller than the batch size. In such
+    # cases, adding batch shuffling ensures that the data is in random order
+    # during training
+    if is_training and batch_shuffle_size:
+      dataset = dataset.shuffle(batch_shuffle_size)
+
     def prepare_for_output(example):
       if not config or not config.use_tpu:
         _summarize_features(example, num_shards)

From 1c9a84eb527120496d5304d5a65e2bbde4daaf68 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 24 Oct 2018 14:03:31 -0700
Subject: [PATCH 1068/2720] Add __init__.py to t2t/metrics

PiperOrigin-RevId: 218571870
---
 tensor2tensor/metrics/__init__.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)
 create mode 100644 tensor2tensor/metrics/__init__.py

diff --git a/tensor2tensor/metrics/__init__.py b/tensor2tensor/metrics/__init__.py
new file mode 100644
index 000000000..4bd418a74
--- /dev/null
+++ b/tensor2tensor/metrics/__init__.py
@@ -0,0 +1,15 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+

From 05d905a5358bcffcf3b2569e50fcc63790e80e22 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 24 Oct 2018 14:17:32 -0700
Subject: [PATCH 1069/2720] Extract out a few transformer layers that were used
 in latent layers, so that layers don't depend on models (because that seems
 backwards, models should depend on layers).

PiperOrigin-RevId: 218574742
---
 tensor2tensor/layers/latent_layers.py      |   9 +-
 tensor2tensor/layers/transformer_layers.py | 329 +++++++++++++++++++++
 tensor2tensor/models/transformer.py        | 309 +------------------
 3 files changed, 340 insertions(+), 307 deletions(-)
 create mode 100644 tensor2tensor/layers/transformer_layers.py

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 72e2404be..421a2c946 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -23,7 +23,7 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
-from tensor2tensor.models import transformer
+from tensor2tensor.layers import transformer_layers
 from tensor2tensor.utils import beam_search
 
 import tensorflow as tf
@@ -410,11 +410,10 @@ def transformer_text_encoder(inputs,
         encoder_input,
         encoder_self_attention_bias,
         ed,
-    ] = transformer.transformer_prepare_encoder(inputs,
-                                                target_space=target_space,
-                                                hparams=hparams)
+    ] = transformer_layers.transformer_prepare_encoder(
+        inputs, target_space=target_space, hparams=hparams)
     encoder_input = tf.nn.dropout(encoder_input, 1.0 - hparams.dropout)
-    encoder_output = transformer.transformer_encoder(
+    encoder_output = transformer_layers.transformer_encoder(
         encoder_input, encoder_self_attention_bias, hparams)
     return encoder_output, ed
 
diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
new file mode 100644
index 000000000..8719b5535
--- /dev/null
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -0,0 +1,329 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Commonly re-used transformer layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import expert_utils
+from tensor2tensor.utils import mlperf_log
+
+import tensorflow as tf
+
+
+def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
+  """Prepare one shard of the model for the encoder.
+
+  Args:
+    inputs: a Tensor.
+    target_space: a Tensor.
+    hparams: run hyperparameters
+    features: optionally pass the entire features dictionary as well.
+      This is needed now for "packed" datasets.
+
+  Returns:
+    encoder_input: a Tensor, bottom of encoder stack
+    encoder_self_attention_bias: a bias tensor for use in encoder self-attention
+    encoder_decoder_attention_bias: a bias tensor for use in encoder-decoder
+      attention
+  """
+  ishape_static = inputs.shape.as_list()
+  encoder_input = inputs
+  if features and "inputs_segmentation" in features:
+    # Packed dataset.  Keep the examples from seeing each other.
+    inputs_segmentation = features["inputs_segmentation"]
+    inputs_position = features["inputs_position"]
+    targets_segmentation = features["targets_segmentation"]
+    encoder_self_attention_bias = common_attention.attention_bias_same_segment(
+        inputs_segmentation, inputs_segmentation)
+    encoder_decoder_attention_bias = (
+        common_attention.attention_bias_same_segment(targets_segmentation,
+                                                     inputs_segmentation))
+  else:
+    # Usual case - not a packed dataset.
+    encoder_padding = common_attention.embedding_to_padding(encoder_input)
+    ignore_padding = common_attention.attention_bias_ignore_padding(
+        encoder_padding)
+    encoder_self_attention_bias = ignore_padding
+    encoder_decoder_attention_bias = ignore_padding
+    inputs_position = None
+  if hparams.proximity_bias:
+    encoder_self_attention_bias += common_attention.attention_bias_proximal(
+        common_layers.shape_list(inputs)[1])
+  if hparams.get("use_target_space_embedding", True):
+    # Append target_space_id embedding to inputs.
+    emb_target_space = common_layers.embedding(
+        target_space,
+        32,
+        ishape_static[-1],
+        name="target_space_embedding",
+        dtype=tf.bfloat16
+        if hparams.activation_dtype == "bfloat16" else tf.float32)
+    emb_target_space = tf.reshape(emb_target_space, [1, 1, -1])
+    encoder_input += emb_target_space
+  if hparams.pos == "timing":
+    if inputs_position is not None:
+      encoder_input = common_attention.add_timing_signal_1d_given_position(
+          encoder_input, inputs_position)
+    else:
+      encoder_input = common_attention.add_timing_signal_1d(encoder_input)
+  elif hparams.pos == "emb":
+    encoder_input = common_attention.add_positional_embedding(
+        encoder_input, hparams.max_length, "inputs_positional_embedding",
+        inputs_position)
+  if hparams.activation_dtype == "bfloat16":
+    encoder_self_attention_bias = tf.cast(encoder_self_attention_bias,
+                                          tf.bfloat16)
+    encoder_decoder_attention_bias = tf.cast(encoder_decoder_attention_bias,
+                                             tf.bfloat16)
+  return (encoder_input, encoder_self_attention_bias,
+          encoder_decoder_attention_bias)
+
+
+def transformer_encoder(encoder_input,
+                        encoder_self_attention_bias,
+                        hparams,
+                        name="encoder",
+                        nonpadding=None,
+                        save_weights_to=None,
+                        make_image_summary=True,
+                        losses=None):
+  """A stack of transformer layers.
+
+  Args:
+    encoder_input: a Tensor
+    encoder_self_attention_bias: bias Tensor for self-attention
+       (see common_attention.attention_bias())
+    hparams: hyperparameters for model
+    name: a string
+    nonpadding: optional Tensor with shape [batch_size, encoder_length]
+      indicating what positions are not padding.  This must either be
+      passed in, which we do for "packed" datasets, or inferred from
+      encoder_self_attention_bias.  The knowledge about padding is used
+      for pad_remover(efficiency) and to mask out padding in convolutional
+      layers.
+    save_weights_to: an optional dictionary to capture attention weights
+      for visualization; the weights tensor will be appended there under
+      a string key created from the variable scope (including name).
+    make_image_summary: Whether to make an attention image summary.
+    losses: optional list onto which to append extra training losses
+
+  Returns:
+    y: a Tensors
+  """
+  x = encoder_input
+  attention_dropout_broadcast_dims = (
+      common_layers.comma_separated_string_to_integer_list(
+          getattr(hparams, "attention_dropout_broadcast_dims", "")))
+  mlperf_log.transformer_print(
+      key=mlperf_log.MODEL_HP_NUM_HIDDEN_LAYERS,
+      value=hparams.num_encoder_layers or hparams.num_hidden_layers)
+  mlperf_log.transformer_print(
+      key=mlperf_log.MODEL_HP_ATTENTION_NUM_HEADS, value=hparams.num_heads)
+  mlperf_log.transformer_print(
+      key=mlperf_log.MODEL_HP_ATTENTION_DROPOUT,
+      value=hparams.attention_dropout)
+  with tf.variable_scope(name):
+    if nonpadding is not None:
+      padding = 1.0 - nonpadding
+    else:
+      padding = common_attention.attention_bias_to_padding(
+          encoder_self_attention_bias)
+      nonpadding = 1.0 - padding
+    pad_remover = None
+    if hparams.use_pad_remover and not common_layers.is_xla_compiled():
+      pad_remover = expert_utils.PadRemover(padding)
+    for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):
+      with tf.variable_scope("layer_%d" % layer):
+        with tf.variable_scope("self_attention"):
+          y = common_attention.multihead_attention(
+              common_layers.layer_preprocess(x, hparams),
+              None,
+              encoder_self_attention_bias,
+              hparams.attention_key_channels or hparams.hidden_size,
+              hparams.attention_value_channels or hparams.hidden_size,
+              hparams.hidden_size,
+              hparams.num_heads,
+              hparams.attention_dropout,
+              attention_type=hparams.self_attention_type,
+              max_relative_position=hparams.max_relative_position,
+              heads_share_relative_embedding=(
+                  hparams.heads_share_relative_embedding),
+              add_relative_to_values=hparams.add_relative_to_values,
+              save_weights_to=save_weights_to,
+              make_image_summary=make_image_summary,
+              dropout_broadcast_dims=attention_dropout_broadcast_dims,
+              max_length=hparams.get("max_length"),
+              vars_3d=hparams.get("attention_variables_3d"))
+          x = common_layers.layer_postprocess(x, y, hparams)
+        with tf.variable_scope("ffn"):
+          y = transformer_ffn_layer(
+              common_layers.layer_preprocess(x, hparams),
+              hparams,
+              pad_remover,
+              conv_padding="SAME",
+              nonpadding_mask=nonpadding,
+              losses=losses)
+          x = common_layers.layer_postprocess(x, y, hparams)
+    # if normalization is done in layer_preprocess, then it should also be done
+    # on the output, since the output can grow very large, being the sum of
+    # a whole stack of unnormalized layer outputs.
+    mlperf_log.transformer_print(
+        key=mlperf_log.MODEL_HP_NORM,
+        value={"hidden_size": hparams.hidden_size})
+    return common_layers.layer_preprocess(x, hparams)
+
+
+def transformer_ffn_layer(x,
+                          hparams,
+                          pad_remover=None,
+                          conv_padding="LEFT",
+                          nonpadding_mask=None,
+                          losses=None,
+                          cache=None,
+                          decode_loop_step=None,
+                          readout_filter_size=0):
+  """Feed-forward layer in the transformer.
+
+  Args:
+    x: a Tensor of shape [batch_size, length, hparams.hidden_size]
+    hparams: hyperparameters for model
+    pad_remover: an expert_utils.PadRemover object tracking the padding
+      positions. If provided, when using convolutional settings, the padding
+      is removed before applying the convolution, and restored afterward. This
+      can give a significant speedup.
+    conv_padding: a string - either "LEFT" or "SAME".
+    nonpadding_mask: an optional Tensor with shape [batch_size, length].
+      needed for convolutional layers with "SAME" padding.
+      Contains 1.0 in positions corresponding to nonpadding.
+    losses: optional list onto which to append extra training losses
+    cache: dict, containing tensors which are the results of previous
+        attentions, used for fast decoding.
+    decode_loop_step: An integer, step number of the decoding loop.
+        Only used for inference on TPU.
+    readout_filter_size: if it's greater than 0, then it will be used instead of
+      filter_size
+
+
+  Returns:
+    a Tensor of shape [batch_size, length, hparams.hidden_size]
+
+  Raises:
+    ValueError: If losses arg is None, but layer generates extra losses.
+  """
+  ffn_layer = hparams.ffn_layer
+  relu_dropout_broadcast_dims = (
+      common_layers.comma_separated_string_to_integer_list(
+          getattr(hparams, "relu_dropout_broadcast_dims", "")))
+  if ffn_layer == "conv_hidden_relu":
+    # Backwards compatibility
+    ffn_layer = "dense_relu_dense"
+  if ffn_layer == "dense_relu_dense":
+    # In simple convolution mode, use `pad_remover` to speed up processing.
+    mlperf_log.transformer_print(
+        key=mlperf_log.MODEL_HP_FFN_FILTER,
+        value={
+            "filter_size": hparams.filter_size,
+            "use_bias": "True",
+            "activation": mlperf_log.RELU
+        })
+    mlperf_log.transformer_print(
+        key=mlperf_log.MODEL_HP_FFN_DENSE,
+        value={
+            "hidden_size": hparams.hidden_size,
+            "use_bias": "True",
+        })
+    mlperf_log.transformer_print(
+        key=mlperf_log.MODEL_HP_RELU_DROPOUT, value=hparams.relu_dropout)
+    if pad_remover:
+      original_shape = common_layers.shape_list(x)
+      # Collapse `x` across examples, and remove padding positions.
+      x = tf.reshape(x, tf.concat([[-1], original_shape[2:]], axis=0))
+      x = tf.expand_dims(pad_remover.remove(x), axis=0)
+    conv_output = common_layers.dense_relu_dense(
+        x,
+        hparams.filter_size,
+        hparams.hidden_size,
+        dropout=hparams.relu_dropout,
+        dropout_broadcast_dims=relu_dropout_broadcast_dims)
+    if pad_remover:
+      # Restore `conv_output` to the original shape of `x`, including padding.
+      conv_output = tf.reshape(
+          pad_remover.restore(tf.squeeze(conv_output, axis=0)), original_shape)
+    return conv_output
+  elif ffn_layer == "conv_relu_conv":
+    return common_layers.conv_relu_conv(
+        x,
+        readout_filter_size or hparams.filter_size,
+        hparams.hidden_size,
+        first_kernel_size=hparams.conv_first_kernel,
+        second_kernel_size=1,
+        padding=conv_padding,
+        nonpadding_mask=nonpadding_mask,
+        dropout=hparams.relu_dropout,
+        cache=cache,
+        decode_loop_step=decode_loop_step)
+  elif ffn_layer == "parameter_attention":
+    return common_attention.parameter_attention(
+        x, hparams.parameter_attention_key_channels or hparams.hidden_size,
+        hparams.parameter_attention_value_channels or hparams.hidden_size,
+        hparams.hidden_size, readout_filter_size or hparams.filter_size,
+        hparams.num_heads,
+        hparams.attention_dropout)
+  elif ffn_layer == "conv_hidden_relu_with_sepconv":
+    return common_layers.conv_hidden_relu(
+        x,
+        readout_filter_size or hparams.filter_size,
+        hparams.hidden_size,
+        kernel_size=(3, 1),
+        second_kernel_size=(31, 1),
+        padding="LEFT",
+        dropout=hparams.relu_dropout)
+  elif ffn_layer == "sru":
+    return common_layers.sru(x)
+  elif ffn_layer == "local_moe_tpu":
+    overhead = (
+        hparams.moe_overhead_train
+        if hparams.mode == tf.estimator.ModeKeys.TRAIN else
+        hparams.moe_overhead_eval)
+    ret, loss = expert_utils.local_moe_tpu(
+        x,
+        hparams.filter_size // 2,
+        hparams.hidden_size,
+        hparams.moe_num_experts,
+        overhead=overhead,
+        loss_coef=hparams.moe_loss_coef)
+  elif ffn_layer == "local_moe":
+    overhead = (
+        hparams.moe_overhead_train
+        if hparams.mode == tf.estimator.ModeKeys.TRAIN else
+        hparams.moe_overhead_eval)
+    ret, loss = expert_utils.local_moe(
+        x,
+        True,
+        expert_utils.ffn_expert_fn(hparams.hidden_size, [hparams.filter_size],
+                                   hparams.hidden_size),
+        hparams.moe_num_experts,
+        k=hparams.moe_k,
+        hparams=hparams)
+    losses.append(loss)
+    return ret
+  else:
+    assert ffn_layer == "none"
+    return x
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 9efe4d9eb..0832babeb 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -32,6 +32,7 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import transformer_layers
 from tensor2tensor.utils import beam_search
 from tensor2tensor.utils import expert_utils
 from tensor2tensor.utils import mlperf_log
@@ -44,6 +45,12 @@
 from tensorflow.python.util import nest
 
 
+# Alias some commonly reused layers, here and elsewhere.
+transformer_prepare_encoder = transformer_layers.transformer_prepare_encoder
+transformer_encoder = transformer_layers.transformer_encoder
+transformer_ffn_layer = transformer_layers.transformer_ffn_layer
+
+
 @registry.register_model
 class Transformer(t2t_model.T2TModel):
   """Attention net.  See file docstring."""
@@ -1108,75 +1115,6 @@ def features_to_nonpadding(features, inputs_or_targets="inputs"):
   return None
 
 
-def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
-  """Prepare one shard of the model for the encoder.
-
-  Args:
-    inputs: a Tensor.
-    target_space: a Tensor.
-    hparams: run hyperparameters
-    features: optionally pass the entire features dictionary as well.
-      This is needed now for "packed" datasets.
-
-  Returns:
-    encoder_input: a Tensor, bottom of encoder stack
-    encoder_self_attention_bias: a bias tensor for use in encoder self-attention
-    encoder_decoder_attention_bias: a bias tensor for use in encoder-decoder
-      attention
-  """
-  ishape_static = inputs.shape.as_list()
-  encoder_input = inputs
-  if features and "inputs_segmentation" in features:
-    # Packed dataset.  Keep the examples from seeing each other.
-    inputs_segmentation = features["inputs_segmentation"]
-    inputs_position = features["inputs_position"]
-    targets_segmentation = features["targets_segmentation"]
-    encoder_self_attention_bias = common_attention.attention_bias_same_segment(
-        inputs_segmentation, inputs_segmentation)
-    encoder_decoder_attention_bias = (
-        common_attention.attention_bias_same_segment(targets_segmentation,
-                                                     inputs_segmentation))
-  else:
-    # Usual case - not a packed dataset.
-    encoder_padding = common_attention.embedding_to_padding(encoder_input)
-    ignore_padding = common_attention.attention_bias_ignore_padding(
-        encoder_padding)
-    encoder_self_attention_bias = ignore_padding
-    encoder_decoder_attention_bias = ignore_padding
-    inputs_position = None
-  if hparams.proximity_bias:
-    encoder_self_attention_bias += common_attention.attention_bias_proximal(
-        common_layers.shape_list(inputs)[1])
-  if hparams.get("use_target_space_embedding", True):
-    # Append target_space_id embedding to inputs.
-    emb_target_space = common_layers.embedding(
-        target_space,
-        32,
-        ishape_static[-1],
-        name="target_space_embedding",
-        dtype=tf.bfloat16
-        if hparams.activation_dtype == "bfloat16" else tf.float32)
-    emb_target_space = tf.reshape(emb_target_space, [1, 1, -1])
-    encoder_input += emb_target_space
-  if hparams.pos == "timing":
-    if inputs_position is not None:
-      encoder_input = common_attention.add_timing_signal_1d_given_position(
-          encoder_input, inputs_position)
-    else:
-      encoder_input = common_attention.add_timing_signal_1d(encoder_input)
-  elif hparams.pos == "emb":
-    encoder_input = common_attention.add_positional_embedding(
-        encoder_input, hparams.max_length, "inputs_positional_embedding",
-        inputs_position)
-  if hparams.activation_dtype == "bfloat16":
-    encoder_self_attention_bias = tf.cast(encoder_self_attention_bias,
-                                          tf.bfloat16)
-    encoder_decoder_attention_bias = tf.cast(encoder_decoder_attention_bias,
-                                             tf.bfloat16)
-  return (encoder_input, encoder_self_attention_bias,
-          encoder_decoder_attention_bias)
-
-
 def transformer_prepare_decoder(targets, hparams, features=None):
   """Prepare one shard of the model for the decoder.
 
@@ -1235,100 +1173,6 @@ def transformer_prepare_decoder(targets, hparams, features=None):
   return (decoder_input, decoder_self_attention_bias)
 
 
-def transformer_encoder(encoder_input,
-                        encoder_self_attention_bias,
-                        hparams,
-                        name="encoder",
-                        nonpadding=None,
-                        save_weights_to=None,
-                        make_image_summary=True,
-                        losses=None):
-  """A stack of transformer layers.
-
-  Args:
-    encoder_input: a Tensor
-    encoder_self_attention_bias: bias Tensor for self-attention
-       (see common_attention.attention_bias())
-    hparams: hyperparameters for model
-    name: a string
-    nonpadding: optional Tensor with shape [batch_size, encoder_length]
-      indicating what positions are not padding.  This must either be
-      passed in, which we do for "packed" datasets, or inferred from
-      encoder_self_attention_bias.  The knowledge about padding is used
-      for pad_remover(efficiency) and to mask out padding in convolutional
-      layers.
-    save_weights_to: an optional dictionary to capture attention weights
-      for visualization; the weights tensor will be appended there under
-      a string key created from the variable scope (including name).
-    make_image_summary: Whether to make an attention image summary.
-    losses: optional list onto which to append extra training losses
-
-  Returns:
-    y: a Tensors
-  """
-  x = encoder_input
-  attention_dropout_broadcast_dims = (
-      common_layers.comma_separated_string_to_integer_list(
-          getattr(hparams, "attention_dropout_broadcast_dims", "")))
-  mlperf_log.transformer_print(
-      key=mlperf_log.MODEL_HP_NUM_HIDDEN_LAYERS,
-      value=hparams.num_encoder_layers or hparams.num_hidden_layers)
-  mlperf_log.transformer_print(
-      key=mlperf_log.MODEL_HP_ATTENTION_NUM_HEADS, value=hparams.num_heads)
-  mlperf_log.transformer_print(
-      key=mlperf_log.MODEL_HP_ATTENTION_DROPOUT,
-      value=hparams.attention_dropout)
-  with tf.variable_scope(name):
-    if nonpadding is not None:
-      padding = 1.0 - nonpadding
-    else:
-      padding = common_attention.attention_bias_to_padding(
-          encoder_self_attention_bias)
-      nonpadding = 1.0 - padding
-    pad_remover = None
-    if hparams.use_pad_remover and not common_layers.is_xla_compiled():
-      pad_remover = expert_utils.PadRemover(padding)
-    for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):
-      with tf.variable_scope("layer_%d" % layer):
-        with tf.variable_scope("self_attention"):
-          y = common_attention.multihead_attention(
-              common_layers.layer_preprocess(x, hparams),
-              None,
-              encoder_self_attention_bias,
-              hparams.attention_key_channels or hparams.hidden_size,
-              hparams.attention_value_channels or hparams.hidden_size,
-              hparams.hidden_size,
-              hparams.num_heads,
-              hparams.attention_dropout,
-              attention_type=hparams.self_attention_type,
-              max_relative_position=hparams.max_relative_position,
-              heads_share_relative_embedding=(
-                  hparams.heads_share_relative_embedding),
-              add_relative_to_values=hparams.add_relative_to_values,
-              save_weights_to=save_weights_to,
-              make_image_summary=make_image_summary,
-              dropout_broadcast_dims=attention_dropout_broadcast_dims,
-              max_length=hparams.get("max_length"),
-              vars_3d=hparams.get("attention_variables_3d"))
-          x = common_layers.layer_postprocess(x, y, hparams)
-        with tf.variable_scope("ffn"):
-          y = transformer_ffn_layer(
-              common_layers.layer_preprocess(x, hparams),
-              hparams,
-              pad_remover,
-              conv_padding="SAME",
-              nonpadding_mask=nonpadding,
-              losses=losses)
-          x = common_layers.layer_postprocess(x, y, hparams)
-    # if normalization is done in layer_preprocess, then it should also be done
-    # on the output, since the output can grow very large, being the sum of
-    # a whole stack of unnormalized layer outputs.
-    mlperf_log.transformer_print(
-        key=mlperf_log.MODEL_HP_NORM,
-        value={"hidden_size": hparams.hidden_size})
-    return common_layers.layer_preprocess(x, hparams)
-
-
 def transformer_decoder(decoder_input,
                         encoder_output,
                         decoder_self_attention_bias,
@@ -1453,145 +1297,6 @@ def transformer_decoder(decoder_input,
     return common_layers.layer_preprocess(x, hparams)
 
 
-def transformer_ffn_layer(x,
-                          hparams,
-                          pad_remover=None,
-                          conv_padding="LEFT",
-                          nonpadding_mask=None,
-                          losses=None,
-                          cache=None,
-                          decode_loop_step=None,
-                          readout_filter_size=0):
-  """Feed-forward layer in the transformer.
-
-  Args:
-    x: a Tensor of shape [batch_size, length, hparams.hidden_size]
-    hparams: hyperparameters for model
-    pad_remover: an expert_utils.PadRemover object tracking the padding
-      positions. If provided, when using convolutional settings, the padding
-      is removed before applying the convolution, and restored afterward. This
-      can give a significant speedup.
-    conv_padding: a string - either "LEFT" or "SAME".
-    nonpadding_mask: an optional Tensor with shape [batch_size, length].
-      needed for convolutional layers with "SAME" padding.
-      Contains 1.0 in positions corresponding to nonpadding.
-    losses: optional list onto which to append extra training losses
-    cache: dict, containing tensors which are the results of previous
-        attentions, used for fast decoding.
-    decode_loop_step: An integer, step number of the decoding loop.
-        Only used for inference on TPU.
-    readout_filter_size: if it's greater than 0, then it will be used instead of
-      filter_size
-
-
-  Returns:
-    a Tensor of shape [batch_size, length, hparams.hidden_size]
-
-  Raises:
-    ValueError: If losses arg is None, but layer generates extra losses.
-  """
-  ffn_layer = hparams.ffn_layer
-  relu_dropout_broadcast_dims = (
-      common_layers.comma_separated_string_to_integer_list(
-          getattr(hparams, "relu_dropout_broadcast_dims", "")))
-  if ffn_layer == "conv_hidden_relu":
-    # Backwards compatibility
-    ffn_layer = "dense_relu_dense"
-  if ffn_layer == "dense_relu_dense":
-    # In simple convolution mode, use `pad_remover` to speed up processing.
-    mlperf_log.transformer_print(
-        key=mlperf_log.MODEL_HP_FFN_FILTER,
-        value={
-            "filter_size": hparams.filter_size,
-            "use_bias": "True",
-            "activation": mlperf_log.RELU
-        })
-    mlperf_log.transformer_print(
-        key=mlperf_log.MODEL_HP_FFN_DENSE,
-        value={
-            "hidden_size": hparams.hidden_size,
-            "use_bias": "True",
-        })
-    mlperf_log.transformer_print(
-        key=mlperf_log.MODEL_HP_RELU_DROPOUT, value=hparams.relu_dropout)
-    if pad_remover:
-      original_shape = common_layers.shape_list(x)
-      # Collapse `x` across examples, and remove padding positions.
-      x = tf.reshape(x, tf.concat([[-1], original_shape[2:]], axis=0))
-      x = tf.expand_dims(pad_remover.remove(x), axis=0)
-    conv_output = common_layers.dense_relu_dense(
-        x,
-        hparams.filter_size,
-        hparams.hidden_size,
-        dropout=hparams.relu_dropout,
-        dropout_broadcast_dims=relu_dropout_broadcast_dims)
-    if pad_remover:
-      # Restore `conv_output` to the original shape of `x`, including padding.
-      conv_output = tf.reshape(
-          pad_remover.restore(tf.squeeze(conv_output, axis=0)), original_shape)
-    return conv_output
-  elif ffn_layer == "conv_relu_conv":
-    return common_layers.conv_relu_conv(
-        x,
-        readout_filter_size or hparams.filter_size,
-        hparams.hidden_size,
-        first_kernel_size=hparams.conv_first_kernel,
-        second_kernel_size=1,
-        padding=conv_padding,
-        nonpadding_mask=nonpadding_mask,
-        dropout=hparams.relu_dropout,
-        cache=cache,
-        decode_loop_step=decode_loop_step)
-  elif ffn_layer == "parameter_attention":
-    return common_attention.parameter_attention(
-        x, hparams.parameter_attention_key_channels or hparams.hidden_size,
-        hparams.parameter_attention_value_channels or hparams.hidden_size,
-        hparams.hidden_size, readout_filter_size or hparams.filter_size,
-        hparams.num_heads,
-        hparams.attention_dropout)
-  elif ffn_layer == "conv_hidden_relu_with_sepconv":
-    return common_layers.conv_hidden_relu(
-        x,
-        readout_filter_size or hparams.filter_size,
-        hparams.hidden_size,
-        kernel_size=(3, 1),
-        second_kernel_size=(31, 1),
-        padding="LEFT",
-        dropout=hparams.relu_dropout)
-  elif ffn_layer == "sru":
-    return common_layers.sru(x)
-  elif ffn_layer == "local_moe_tpu":
-    overhead = (
-        hparams.moe_overhead_train
-        if hparams.mode == tf.estimator.ModeKeys.TRAIN else
-        hparams.moe_overhead_eval)
-    ret, loss = expert_utils.local_moe_tpu(
-        x,
-        hparams.filter_size // 2,
-        hparams.hidden_size,
-        hparams.moe_num_experts,
-        overhead=overhead,
-        loss_coef=hparams.moe_loss_coef)
-  elif ffn_layer == "local_moe":
-    overhead = (
-        hparams.moe_overhead_train
-        if hparams.mode == tf.estimator.ModeKeys.TRAIN else
-        hparams.moe_overhead_eval)
-    ret, loss = expert_utils.local_moe(
-        x,
-        True,
-        expert_utils.ffn_expert_fn(hparams.hidden_size, [hparams.filter_size],
-                                   hparams.hidden_size),
-        hparams.moe_num_experts,
-        k=hparams.moe_k,
-        hparams=hparams)
-    losses.append(loss)
-    return ret
-  else:
-    assert ffn_layer == "none"
-    return x
-
-
 @registry.register_hparams
 def transformer_base_v1():
   """Set of hyperparameters."""

From 040c01c9b44f7bc451cf30c3f10aacace917e445 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 24 Oct 2018 17:04:00 -0700
Subject: [PATCH 1070/2720] Move the RL tests into a separate stanza, since
 they use flags.

PiperOrigin-RevId: 218604530
---
 oss_scripts/oss_tests.sh | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index de53f604a..a81d47b19 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -44,10 +44,8 @@ pytest \
   --ignore=tensor2tensor/visualization/visualization_test.py \
   --ignore=tensor2tensor/bin/t2t_trainer_test.py \
   --ignore=tensor2tensor/data_generators/algorithmic_math_test.py \
-  --ignore=tensor2tensor/rl/trainer_model_based_test.py \
   --ignore=tensor2tensor/data_generators/allen_brain_test.py \
-  --ignore=tensor2tensor/rl/trainer_model_based_stochastic_test.py \
-  --ignore=tensor2tensor/rl/trainer_model_based_sv2p_test.py \
+  --ignore=tensor2tensor/rl \
   --ignore=tensor2tensor/models/research \
   --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary
 set_status
@@ -76,6 +74,7 @@ fi
 
 if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]
 then
+    # TODO(afrozm): Enable other tests in the RL directory.
     pytest tensor2tensor/rl/trainer_model_based_test.py
     set_status
     jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/hello_t2t.ipynb

From 1b17293ae0bbd42090e93bfbfe5014359b85f3e5 Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Wed, 24 Oct 2018 17:11:59 -0700
Subject: [PATCH 1071/2720] Add host_call to Mesh TF model & reuse methods
 "create_host_call" / "remove_summaries" from t2t_model.py.

PiperOrigin-RevId: 218605726
---
 tensor2tensor/utils/mtf_model.py | 68 +++++---------------------------
 tensor2tensor/utils/t2t_model.py | 17 +++++---
 2 files changed, 22 insertions(+), 63 deletions(-)

diff --git a/tensor2tensor/utils/mtf_model.py b/tensor2tensor/utils/mtf_model.py
index 193107d33..24c6362e3 100644
--- a/tensor2tensor/utils/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -20,7 +20,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import collections
 import copy
 import mesh_tensorflow as mtf
 
@@ -111,6 +110,7 @@ def estimator_model_fn(cls,
       var_grads = mtf.gradients(
           [loss], [v.outputs[0] for v in graph.trainable_variables])
       lr = learning_rate.learning_rate_schedule(hparams)
+      tf.summary.scalar("learning_rate", lr)
       mtf_lr = mtf.import_tf_tensor(
           mesh, tf.convert_to_tensor(lr, dtype=tf.float32), mtf.Shape([]))
       optimizer = mtf.optimize.make_optimizer(hparams, mtf_lr)
@@ -156,11 +156,18 @@ def estimator_model_fn(cls,
                                        restore_hook, use_tpu)
 
     if use_tpu:
-      _remove_summaries()
+      # TPU host call. Important: need to be called before remove_summaries()
+      if hparams.tpu_enable_host_call:
+        host_call = t2t_model.create_host_call(hparams.model_dir)
+      else:
+        host_call = None
+
+      t2t_model.remove_summaries()
       return tpu_estimator.TPUEstimatorSpec(
           mode=tf.estimator.ModeKeys.TRAIN,
           loss=tf_loss,
           train_op=train_op,
+          host_call=host_call,
           training_hooks=[restore_hook, saver_hook])
     else:
       return tf.estimator.EstimatorSpec(
@@ -219,7 +226,7 @@ def estimator_spec_predict(self, features, mesh, mesh_impl, use_tpu):
         "inputs": features.get("inputs"),
     }
     if use_tpu:
-      _remove_summaries()
+      t2t_model.remove_summaries()
       return tpu_estimator.TPUEstimatorSpec(
           mode=tf.estimator.ModeKeys.PREDICT,
           predictions=predictions,
@@ -236,58 +243,3 @@ def sample(self, features, mesh):
 
   def mtf_model_fn(self, features, mesh):
     raise NotImplementedError("Not implemented")
-
-
-def _remove_summaries():
-  g = tf.get_default_graph()
-  key = tf.GraphKeys.SUMMARIES
-  del g.get_collection_ref(key)[:]
-  assert not g.get_collection(key)
-
-
-def _create_host_call(model_dir):
-  """Construct a host_call writing scalar summaries.
-
-  Args:
-    model_dir: String containing path to train
-
-  Returns:
-    (fn, args) Pair to be called by TPUEstimator as the host_call.
-  """
-  graph = tf.get_default_graph()
-  summaries = graph.get_collection(tf.GraphKeys.SUMMARIES)
-
-  gs_t = tf.reshape(tf.to_int32(tf.train.get_global_step()), [1])
-  summary_kwargs = collections.OrderedDict()
-  for t in summaries:
-    if t.op.type != "ScalarSummary":
-      continue
-
-    name = t.op.name
-    tensor = t.op.inputs[1]
-    assert tensor.shape.is_compatible_with([])
-    if tensor.dtype == tf.int64:
-      tensor = tf.to_int32(tensor)
-    summary_kwargs[name] = tf.reshape(tensor, [1])
-  summary_kwargs["global_step"] = gs_t
-
-  def host_call_fn(**kwargs):
-    """Training host call. Creates scalar summaries for training metrics.
-
-    Args:
-      **kwargs: Dict of {str: Tensor} , with `Tensor` of shape `[batch]`. Must
-        contain key "global_step" with value of current global_step Tensor.
-
-    Returns:
-      List of summary ops to run on the CPU host.
-    """
-    gs = tf.to_int64(kwargs.pop("global_step")[0])
-    with tf.contrib.summary.create_file_writer(model_dir).as_default():
-      with tf.contrib.summary.always_record_summaries():
-        for name, value in sorted(six.iteritems(kwargs)):
-          tf.contrib.summary.scalar(
-              name, tf.reduce_mean(tf.to_float(value)), step=gs)
-
-        return tf.contrib.summary.all_summary_ops()
-
-  return (host_call_fn, summary_kwargs)
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 0e16c6f8c..5d20d8421 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1350,10 +1350,10 @@ def initialize_from_ckpt(self, ckpt_dir):
     for var in tf.contrib.framework.get_trainable_variables():
       var_name = var.name.split(":")[0]
       if reader.has_tensor(var_name):
-        tf.logging.info("Loading variable from checkpoint: %s", var_name)
+        log_info("Loading variable from checkpoint: %s", var_name)
         variable_map[var_name] = var
       else:
-        tf.logging.info(
+        log_info(
             "Cannot find variable in checkpoint, skipping: %s", var_name)
     tf.train.init_from_checkpoint(ckpt_dir, variable_map)
 
@@ -1370,8 +1370,9 @@ def scaffold_fn():
       else:
         scaffold_fn = None
 
+      # Note: important to call this before remove_summaries()
       if self.hparams.tpu_enable_host_call:
-        host_call = _create_host_call(self.hparams.model_dir)
+        host_call = create_host_call(self.hparams.model_dir)
       else:
         host_call = None
 
@@ -1638,13 +1639,15 @@ def all_metrics_fn(logits=None, labels=None, **kwargs):
 
 
 def remove_summaries():
+  """Remove summaries from the default graph."""
   g = tf.get_default_graph()
   key = tf.GraphKeys.SUMMARIES
+  log_debug("Remove summaries %s" % str(g.get_collection(key)))
   del g.get_collection_ref(key)[:]
   assert not g.get_collection(key)
 
 
-def _create_host_call(model_dir):
+def create_host_call(model_dir):
   """Construct a host_call writing scalar summaries.
 
   Args:
@@ -1655,7 +1658,6 @@ def _create_host_call(model_dir):
   """
   graph = tf.get_default_graph()
   summaries = graph.get_collection(tf.GraphKeys.SUMMARIES)
-
   gs_t = tf.reshape(tf.to_int32(tf.train.get_global_step()), [1])
   summary_kwargs = collections.OrderedDict()
   for t in summaries:
@@ -1689,6 +1691,7 @@ def _create_host_call(model_dir):
   if not summary_kwargs:
     return None
   summary_kwargs["global_step"] = gs_t
+  log_info("summary_kwargs %s" % str(summary_kwargs))
 
   def host_call_fn(**kwargs):
     """Training host call. Creates summaries for training metrics.
@@ -1843,6 +1846,10 @@ def _eager_log(level, *args):
   getattr(tf.logging, level)(*args)
 
 
+def log_debug(*args):
+  _eager_log("debug", *args)
+
+
 def log_info(*args):
   _eager_log("info", *args)
 

From 950faf86ea618470fdf91b180d916aa4954c19fa Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 24 Oct 2018 18:08:21 -0700
Subject: [PATCH 1072/2720] Add hook_context to all train_hooks

PiperOrigin-RevId: 218612750
---
 tensor2tensor/models/research/glow.py | 3 ++-
 tensor2tensor/models/video/savp.py    | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 7511a7cc2..4f788a4bd 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -135,7 +135,8 @@ def init_batch_size(self):
     return self.hparams.init_batch_size
 
   @staticmethod
-  def train_hooks():
+  def train_hooks(hook_context):
+    del hook_context
     return [glow_init_hook.GlowInitHook()]
 
   def top_prior(self):
diff --git a/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index f1ba8e486..d05244c83 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -327,7 +327,8 @@ def pad_conv3d_lrelu(self, activations, n_filters, kernel_size, strides,
     return rectified
 
   @staticmethod
-  def train_hooks():
+  def train_hooks(hook_context):
+    del hook_context
     return [update_ops_hook.UpdateOpsHook()]
 
 

From dc5e328b82c5f300b2b98a4f236ebcc17dcbca76 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 24 Oct 2018 19:56:01 -0700
Subject: [PATCH 1073/2720] MLPerf compliance for transformer.

PiperOrigin-RevId: 218621797
---
 tensor2tensor/bin/t2t_trainer.py           |  7 +++---
 tensor2tensor/data_generators/translate.py | 20 +++++++++++++++-
 tensor2tensor/utils/decoding.py            |  6 ++++-
 tensor2tensor/utils/trainer_lib.py         | 27 ++++++++++++++++++----
 4 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 53413c2f1..61b057bbf 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -344,7 +344,8 @@ def run_std_server():
 
 def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
-  mlperf_log.transformer_print(key=mlperf_log.RUN_START)
+  if FLAGS.schedule != "train":
+    mlperf_log.transformer_print(key=mlperf_log.RUN_START)
   if FLAGS.schedule == "run_std_server":
     run_std_server()
   mlperf_log.transformer_print(
@@ -372,8 +373,8 @@ def main(argv):
   if is_chief():
     save_metadata(hparams)
   execute_schedule(exp)
-  mlperf_log.transformer_print(key=mlperf_log.RUN_STOP)
-  mlperf_log.transformer_print(key=mlperf_log.RUN_FINAL)
+  if FLAGS.schedule != "train":
+    mlperf_log.transformer_print(key=mlperf_log.RUN_FINAL)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 0ad29761d..2cea49b69 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -78,6 +78,12 @@ def compute_bleu_summaries(hook_args):
     reference file and the translated file.
   """
   decode_hparams = hook_args.decode_hparams
+  estimator = hook_args.estimator
+  current_step = estimator.get_variable_value(tf.GraphKeys.GLOBAL_STEP)
+  if decode_hparams.iterations_per_loop:
+    current_epoch = current_step // decode_hparams.iterations_per_loop
+  else:
+    current_epoch = 0
 
   if (decode_hparams.decode_reference is None or
       decode_hparams.decode_to_file is None):
@@ -88,7 +94,19 @@ def compute_bleu_summaries(hook_args):
       decode_hparams.decode_reference, decode_hparams.decode_to_file)
   values.append(tf.Summary.Value(tag="BLEU", simple_value=bleu))
   tf.logging.info("%s: BLEU = %6.2f" % (decode_hparams.decode_to_file, bleu))
-  mlperf_log.transformer_print(key=mlperf_log.EVAL_ACCURACY, value=bleu)
+  mlperf_log.transformer_print(
+      key=mlperf_log.EVAL_TARGET, value=decode_hparams.mlperf_threshold)
+  mlperf_log.transformer_print(
+      key=mlperf_log.EVAL_ACCURACY,
+      value={
+          "epoch": current_epoch,
+          "value": bleu
+      })
+  mlperf_log.transformer_print(key=mlperf_log.EVAL_STOP)
+
+  if bleu >= decode_hparams.mlperf_threshold:
+    decode_hparams.set_hparam("mlperf_success", True)
+
   return values
 
 
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index da830ee99..5d3797028 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -69,7 +69,11 @@ def decode_hparams(overrides=""):
       display_decoded_images=False,
       # Used for video decoding.
       frames_per_second=10,
-      skip_eos_postprocess=False)
+      skip_eos_postprocess=False,
+      # Used for MLPerf compliance logging.
+      mlperf_mode=False,
+      mlperf_threshold=25.0,
+      mlperf_success=False)
   hp.parse(overrides)
   return hp
 
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 955449d98..015d991a1 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -269,6 +269,9 @@ def create_estimator(model_name,
     predict_batch_size = batch_size
     if decode_hparams and decode_hparams.batch_size:
       predict_batch_size = decode_hparams.batch_size
+    if decode_hparams and run_config.tpu_config:
+      decode_hparams.add_hparam("iterations_per_loop",
+                                run_config.tpu_config.iterations_per_loop)
     estimator = tf.contrib.tpu.TPUEstimator(
         model_fn=model_fn,
         model_dir=run_config.model_dir,
@@ -405,8 +408,16 @@ def train_eval_and_decode(self):
         self._hparams.problem_hparams = p_hparams
       mlperf_log.transformer_print(key=mlperf_log.EVAL_START)
       self.decode(dataset_split=tf.estimator.ModeKeys.EVAL)
-      mlperf_log.transformer_print(key=mlperf_log.EVAL_TARGET, value=25.0)
-      mlperf_log.transformer_print(key=mlperf_log.EVAL_STOP)
+      d_hparams = self._decode_hparams
+      if d_hparams.mlperf_mode and d_hparams.mlperf_success:
+        mlperf_log.transformer_print(
+            key=mlperf_log.RUN_STOP, value={"success": "true"})
+        break
+
+    d_hparams = self._decode_hparams
+    if d_hparams.mlperf_mode and not d_hparams.mlperf_success:
+      mlperf_log.transformer_print(
+          key=mlperf_log.RUN_STOP, value={"success": "false"})
 
   def evaluate(self):
     return self._estimator.evaluate(
@@ -493,8 +504,16 @@ def continuous_decode_on_eval_data(self):
         continue
       mlperf_log.transformer_print(key=mlperf_log.EVAL_START)
       self.decode(dataset_split=tf.estimator.ModeKeys.EVAL)
-      mlperf_log.transformer_print(key=mlperf_log.EVAL_TARGET, value=25.0)
-      mlperf_log.transformer_print(key=mlperf_log.EVAL_STOP)
+      d_hparams = self._decode_hparams
+      if d_hparams.mlperf_mode and d_hparams.mlperf_success:
+        mlperf_log.transformer_print(
+            key=mlperf_log.RUN_STOP, value={"success": "true"})
+        break
+
+    d_hparams = self._decode_hparams
+    if d_hparams.mlperf_mode and not d_hparams.mlperf_success:
+      mlperf_log.transformer_print(
+          key=mlperf_log.RUN_STOP, value={"success": "false"})
 
   def continuous_decode_from_file(self):
     """Decode from file on new checkpoint."""

From 279442a843815d7c09a6add157047e574fd98bf2 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 24 Oct 2018 21:15:04 -0700
Subject: [PATCH 1074/2720] Update reference to deleted gym_problems_specs with
 gym_env

PiperOrigin-RevId: 218628079
---
 tensor2tensor/data_generators/all_problems.py      | 1 +
 tensor2tensor/rl/trainer_model_based_agent_only.py | 4 ++--
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 311811dbf..81d961d56 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -40,6 +40,7 @@
     "tensor2tensor.data_generators.function_docstring",
     "tensor2tensor.data_generators.gene_expression",
     "tensor2tensor.data_generators.google_robot_pushing",
+    "tensor2tensor.data_generators.gym_env",
     "tensor2tensor.data_generators.ice_parsing",
     "tensor2tensor.data_generators.imagenet",
     "tensor2tensor.data_generators.image_lsun",
diff --git a/tensor2tensor/rl/trainer_model_based_agent_only.py b/tensor2tensor/rl/trainer_model_based_agent_only.py
index cced3ebc5..6d9cc99bb 100644
--- a/tensor2tensor/rl/trainer_model_based_agent_only.py
+++ b/tensor2tensor/rl/trainer_model_based_agent_only.py
@@ -28,7 +28,7 @@
 from __future__ import print_function
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
-from tensor2tensor.data_generators import gym_problems_specs
+from tensor2tensor.data_generators import gym_env
 from tensor2tensor.rl import trainer_model_based
 from tensor2tensor.rl import trainer_model_based_params
 
@@ -45,7 +45,7 @@
 
 def get_simulated_problem_name(game):
   game_with_mode = game
-  if game in gym_problems_specs.ATARI_GAMES:
+  if game in gym_env.ATARI_GAMES:
     game_with_mode += "_deterministic-v4"
   return "gym_simulated_discrete_problem_with_agent_on_%s" % game_with_mode
 

From cfaca8a45c0a5402751558c47065f01d945e7c42 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 25 Oct 2018 10:13:56 -0700
Subject: [PATCH 1075/2720] Fix mlperf numpy array logging issue.

PiperOrigin-RevId: 218706375
---
 tensor2tensor/data_generators/translate.py | 25 ++++++++++++----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 2cea49b69..adc964edf 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -21,6 +21,7 @@
 
 import os
 import tarfile
+import numpy as np
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -80,8 +81,9 @@ def compute_bleu_summaries(hook_args):
   decode_hparams = hook_args.decode_hparams
   estimator = hook_args.estimator
   current_step = estimator.get_variable_value(tf.GraphKeys.GLOBAL_STEP)
-  if decode_hparams.iterations_per_loop:
-    current_epoch = current_step // decode_hparams.iterations_per_loop
+  if current_step and decode_hparams.iterations_per_loop:
+    iterations_per_loop = decode_hparams.iterations_per_loop
+    current_epoch = np.asscalar(current_step) // iterations_per_loop
   else:
     current_epoch = 0
 
@@ -94,15 +96,16 @@ def compute_bleu_summaries(hook_args):
       decode_hparams.decode_reference, decode_hparams.decode_to_file)
   values.append(tf.Summary.Value(tag="BLEU", simple_value=bleu))
   tf.logging.info("%s: BLEU = %6.2f" % (decode_hparams.decode_to_file, bleu))
-  mlperf_log.transformer_print(
-      key=mlperf_log.EVAL_TARGET, value=decode_hparams.mlperf_threshold)
-  mlperf_log.transformer_print(
-      key=mlperf_log.EVAL_ACCURACY,
-      value={
-          "epoch": current_epoch,
-          "value": bleu
-      })
-  mlperf_log.transformer_print(key=mlperf_log.EVAL_STOP)
+  if decode_hparams.mlperf_mode:
+    mlperf_log.transformer_print(
+        key=mlperf_log.EVAL_TARGET, value=decode_hparams.mlperf_threshold)
+    mlperf_log.transformer_print(
+        key=mlperf_log.EVAL_ACCURACY,
+        value={
+            "epoch": current_epoch,
+            "value": bleu
+        })
+    mlperf_log.transformer_print(key=mlperf_log.EVAL_STOP)
 
   if bleu >= decode_hparams.mlperf_threshold:
     decode_hparams.set_hparam("mlperf_success", True)

From 503530e2c0c67ab1b1635120cfe3d4a591a0bdd8 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 25 Oct 2018 11:00:14 -0700
Subject: [PATCH 1076/2720] Adding optional TB histograms for predicted logits.

PiperOrigin-RevId: 218715358
---
 tensor2tensor/models/video/sv2p_params.py | 1 +
 tensor2tensor/utils/t2t_model.py          | 4 ++++
 2 files changed, 5 insertions(+)

diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index b0156f9ce..e81afdd85 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -47,6 +47,7 @@ def next_frame_sv2p():
   hparams.add_hparam("dna_kernel_size", 5)
   hparams.add_hparam("upsample_method", "conv2d_transpose")
   hparams.add_hparam("reward_model", "basic")
+  hparams.add_hparam("visualize_logits_histogram", True)
   return hparams
 
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 5d20d8421..390a6493c 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -535,6 +535,10 @@ def loss(self, logits, features):
           tf.summary.scalar(k + "_loss", n / d)
           tf.summary.scalar(k + "_loss_num", n)
           tf.summary.scalar(k + "_loss_den", d)
+          if getattr(self.hparams, "visualize_logits_histogram", False):
+            hist = tf.summary.histogram
+            hist(k + "_predict", tf.argmax(tf.squeeze(v), axis=-1))
+            hist(k + "_targets", features[k])
 
       return tf.add_n([n / d for n, d in losses.values()])
     else:

From a70999ad9f0a6629bc7f49cd05f15eaacd5a3b5b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 25 Oct 2018 11:54:29 -0700
Subject: [PATCH 1077/2720] Adjust LR for SVP-FP.

PiperOrigin-RevId: 218725327
---
 tensor2tensor/models/video/emily.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index 102eae9bf..06340df9b 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -300,7 +300,7 @@ def next_frame_emily():
   """Emily's model hparams."""
   hparams = sv2p_params.next_frame_sv2p()
   hparams.latent_loss_multiplier = 1e-4
-  hparams.learning_rate_constant = 0.002
+  hparams.learning_rate_constant = 1e-4
   hparams.add_hparam("z_dim", 10)
   hparams.add_hparam("g_dim", 128)
   hparams.add_hparam("rnn_size", 256)

From 4161ada59129d9fa8b8421f5693c7860b5b39ace Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Thu, 25 Oct 2018 17:15:10 -0700
Subject: [PATCH 1078/2720] fixes for model free pipeline.

PiperOrigin-RevId: 218778448
---
 tensor2tensor/data_generators/gym_env.py      |  8 +++--
 tensor2tensor/models/research/rl.py           | 29 +++++++------------
 tensor2tensor/rl/README.md                    |  2 +-
 tensor2tensor/rl/rl_trainer_lib.py            | 23 +++++++++------
 .../rl/trainer_model_based_params.py          |  8 ++++-
 tensor2tensor/rl/trainer_model_free.py        | 25 +++++++++++++++-
 6 files changed, 61 insertions(+), 34 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 4bcbb1af4..f8c47825e 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -715,6 +715,10 @@ def __init__(self, action_space, reward_range):
 ]
 
 
+def camel_case_name(snake_case_name):
+  return "".join([w[0].upper() + w[1:] for w in snake_case_name.split("_")])
+
+
 def register_game(game_name, game_mode="Deterministic-v4"):
   """Create and register problems for the game.
 
@@ -729,9 +733,7 @@ def register_game(game_name, game_mode="Deterministic-v4"):
     raise ValueError("Game %s not in ATARI_GAMES" % game_name)
   if game_mode not in ATARI_GAME_MODES:
     raise ValueError("Unknown ATARI game mode: %s." % game_mode)
-  camel_game_name = "".join(
-      [w[0].upper() + w[1:] for w in game_name.split("_")])
-  camel_game_name += game_mode
+  camel_game_name = camel_case_name(game_name) + game_mode
   # Create and register the Problem
   cls = type("Gym%sRandom" % camel_game_name,
              (T2TGymEnv,), {"base_env_name": camel_game_name})
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 4b96264d9..28e98a6bf 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -93,24 +93,6 @@ def discrete_random_action_base():
 
 @registry.register_hparams
 def ppo_atari_base():
-  """Atari base parameters."""
-  hparams = ppo_discrete_action_base()
-  hparams.learning_rate = 4e-4
-  hparams.num_agents = 5
-  hparams.epoch_length = 200
-  hparams.gae_gamma = 0.985
-  hparams.gae_lambda = 0.985
-  hparams.entropy_loss_coef = 0.002
-  hparams.value_loss_coef = 0.025
-  hparams.optimization_epochs = 10
-  hparams.epochs_num = 10000
-  hparams.num_eval_agents = 1
-  hparams.network = feed_forward_cnn_small_categorical_fun
-  return hparams
-
-
-@registry.register_hparams
-def ppo_pong_base():
   """Pong base parameters."""
   hparams = ppo_discrete_action_base()
   hparams.learning_rate = 1e-4
@@ -212,7 +194,7 @@ def get_policy(observations, hparams):
 @registry.register_hparams
 def ppo_pong_ae_base():
   """Pong autoencoder base parameters."""
-  hparams = ppo_pong_base()
+  hparams = ppo_atari_base()
   hparams.learning_rate = 1e-4
   hparams.network = dense_bitwise_categorical_fun
   return hparams
@@ -247,6 +229,15 @@ def pong_model_free():
   return hparams
 
 
+@registry.register_hparams
+def mfrl_base():
+  hparams = ppo_atari_base()
+  hparams.add_hparam("game", "")
+  hparams.epochs_num = 3000
+  hparams.eval_every_epochs = 100
+  return hparams
+
+
 NetworkOutput = collections.namedtuple(
     "NetworkOutput", "policy, value, action_postprocessing")
 
diff --git a/tensor2tensor/rl/README.md b/tensor2tensor/rl/README.md
index e38c9bd5e..4b3f34e67 100644
--- a/tensor2tensor/rl/README.md
+++ b/tensor2tensor/rl/README.md
@@ -32,7 +32,7 @@ The hyperparameters for the environment model and agent are nested within the
 ```
   generative_model="next_frame_basic",
   generative_model_params="next_frame_pixel_noise",
-  ppo_params="ppo_pong_base",
+  ppo_params="ppo_atari_base",
 ```
 
 ## Model-free training
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 0af54b6a5..d3fad351c 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -32,24 +32,22 @@
 def define_train(hparams):
   """Define the training setup."""
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-    memory, collect_summary, train_initialization\
-      = collect.define_collect(
-          hparams, "ppo_train", eval_phase=False)
+    memory, collect_summary, train_initialization = collect.define_collect(
+        hparams, "ppo_train", eval_phase=False)
     ppo_summary = ppo.define_ppo_epoch(memory, hparams)
     train_summary = tf.summary.merge([collect_summary, ppo_summary])
 
     if hparams.eval_every_epochs:
-      _, eval_collect_summary, eval_initialization\
-        = collect.define_collect(
-            hparams, "ppo_eval", eval_phase=True)
-      return train_summary, eval_collect_summary, \
-             (train_initialization, eval_initialization)
+      _, eval_collect_summary, eval_initialization = collect.define_collect(
+          hparams, "ppo_eval", eval_phase=True)
+      return train_summary, eval_collect_summary, (train_initialization,
+                                                   eval_initialization)
     else:
       return train_summary, None, (train_initialization,)
 
 
 def train(hparams, event_dir=None, model_dir=None,
-          restore_agent=True, name_scope="rl_train"):
+          restore_agent=True, name_scope="rl_train", report_fn=None):
   """Train."""
   with tf.Graph().as_default():
     with tf.name_scope(name_scope):
@@ -106,6 +104,13 @@ def train(hparams, event_dir=None, model_dir=None,
             eval_summary = sess.run(eval_summary_op)
             if summary_writer:
               summary_writer.add_summary(eval_summary, epoch_index)
+            if report_fn:
+              summary_proto = tf.Summary()
+              summary_proto.ParseFromString(eval_summary)
+              for elem in summary_proto.value:
+                if "mean_score" in elem.tag:
+                  report_fn(elem.simple_value, epoch_index)
+                  break
 
           epoch_index_and_start = epoch_index + start_step
           if (model_saver and hparams.save_models_every_epochs and
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 24368a86d..c1367dad9 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -54,7 +54,7 @@ def rlmb_base():
       num_real_env_frames=96000,
       generative_model="next_frame_basic_deterministic",
       generative_model_params="next_frame_pixel_noise",
-      ppo_params="ppo_pong_base",
+      ppo_params="ppo_atari_base",
       autoencoder_train_steps=0,
       autoencoder_train_steps_initial_multiplier=10,
       autoencoder_hparams_set="autoencoder_discrete_pong",
@@ -679,6 +679,12 @@ def rlmb_logits_clip(rhp):
   rhp.set_discrete("ppo.logits_clip", [0., 5.])
 
 
+@registry.register_ranged_hparams
+def rlmf_proportional_epoch_length(rhp):
+  rhp.set_discrete("proportional_epoch_length", [10, 20, 50, 100, 200, 400])
+  rhp.set_categorical("loop.game", gym_env.ATARI_GAMES_WITH_HUMAN_SCORE)
+
+
 def merge_unscoped_hparams(scopes_and_hparams):
   """Merge multiple HParams into one with scopes."""
   merged_values = {}
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 0375afb86..b6d105526 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -26,6 +26,8 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+from tensor2tensor.data_generators import gym_env
+from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
@@ -43,9 +45,30 @@
   pass
 
 
+def initialize_env_specs(hparams):
+  """Initializes env_specs using T2TGymEnvs."""
+  if getattr(hparams, "game", None):
+    game_name = gym_env.camel_case_name(hparams.game)
+    env = gym_env.T2TGymEnv("{}Deterministic-v4".format(game_name),
+                            batch_size=hparams.num_agents)
+    env.start_new_epoch(0)
+    hparams.add_hparam("environment_spec", rl.standard_atari_env_spec(env))
+    eval_env = gym_env.T2TGymEnv("{}Deterministic-v4".format(game_name),
+                                 batch_size=hparams.num_eval_agents)
+    eval_env.start_new_epoch(0)
+    hparams.add_hparam(
+        "environment_eval_spec", rl.standard_atari_env_eval_spec(eval_env))
+  return hparams
+
+
+def train(hparams, output_dir, report_fn=None):
+  hparams = initialize_env_specs(hparams)
+  rl_trainer_lib.train(hparams, output_dir, output_dir, report_fn=report_fn)
+
+
 def main(_):
   hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
-  rl_trainer_lib.train(hparams, FLAGS.output_dir, FLAGS.output_dir)
+  train(hparams, FLAGS.output_dir)
 
 
 if __name__ == "__main__":

From 14f0811421fe135e375414a7bda52b2b0f8bc94f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 25 Oct 2018 19:07:20 -0700
Subject: [PATCH 1079/2720] Fix for MLPerf log compliance by changing the epoch
 in log to be zero indexed.

PiperOrigin-RevId: 218789684
---
 tensor2tensor/data_generators/translate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index adc964edf..b6ef5de06 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -102,7 +102,7 @@ def compute_bleu_summaries(hook_args):
     mlperf_log.transformer_print(
         key=mlperf_log.EVAL_ACCURACY,
         value={
-            "epoch": current_epoch,
+            "epoch": max(current_epoch - 1, 0),
             "value": bleu
         })
     mlperf_log.transformer_print(key=mlperf_log.EVAL_STOP)

From 3d482ec2d1fef759d85de71eb694d8f819a10270 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 26 Oct 2018 05:32:51 +0200
Subject: [PATCH 1080/2720] Improve agent evaluation (#1168)

* Remove argument envs from T2TGymEnv.__init__

* Implement noops on reset

* Remove collect_level

* Make eval_phase and policy_to_action_lambda independent

* Pass collect parameters via hparams

* Implement agent evaluation with multiple settings

* Change metric names so the're displayed better in tensorboard
---
 tensor2tensor/data_generators/gym_env.py      |  42 +++++--
 tensor2tensor/rl/collect.py                   |  39 ++----
 tensor2tensor/rl/rl_trainer_lib.py            |  29 ++++-
 tensor2tensor/rl/trainer_model_based.py       | 117 +++++++++++++-----
 .../rl/trainer_model_based_params.py          |   6 +
 5 files changed, 159 insertions(+), 74 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index f8c47825e..dfb75aecd 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -25,6 +25,7 @@
 import gym
 from gym.spaces import Box
 import numpy as np
+import random
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -530,14 +531,11 @@ class T2TGymEnv(T2TEnv):
   arguments and register this subclass.
   """
 
-  def __init__(self, base_env_name=None, batch_size=None, grayscale=False,
+  noop_action = 0
+
+  def __init__(self, base_env_name=None, batch_size=1, grayscale=False,
                resize_height_factor=2, resize_width_factor=2,
-               base_env_timesteps_limit=-1, envs=None, **kwargs):
-    if batch_size is None:
-      if envs is None:
-        batch_size = 1
-      else:
-        batch_size = len(envs)
+               base_env_timesteps_limit=-1, max_num_noops=0, **kwargs):
     if base_env_name is None:
       base_env_name = self.base_env_name
     self._base_env_name = base_env_name
@@ -549,11 +547,15 @@ def __init__(self, base_env_name=None, batch_size=None, grayscale=False,
       # Set problem name if not registered.
       self.name = "Gym%s" % base_env_name
 
-    if envs is None:
-      self._envs = [make_gym_env(base_env_name, base_env_timesteps_limit)
-                    for _ in range(self.batch_size)]
-    else:
-      self._envs = envs
+    self._envs = [make_gym_env(base_env_name, base_env_timesteps_limit)
+                  for _ in range(self.batch_size)]
+
+    # max_num_noops works only with atari envs.
+    if max_num_noops > 0:
+      assert self._envs[0].unwrapped.get_action_meanings()[
+          self.noop_action
+      ] == 'NOOP'
+    self.max_num_noops = max_num_noops
 
     orig_observ_space = self._envs[0].observation_space
     if not all(env.observation_space == orig_observ_space
@@ -613,7 +615,21 @@ def _step(self, actions):
     return tuple(map(np.stack, (obs, rewards, dones)))
 
   def _reset(self, indices):
-    return np.stack([self._envs[index].reset() for index in indices])
+    def reset_with_noops(env):
+      obs = env.reset()
+      try:
+        num_noops = random.randint(1, self.max_num_noops)
+      except ValueError:
+        num_noops = 0
+
+      for _ in range(num_noops):
+        (obs, _, done, _) = env.step(self.noop_action)
+        if done:
+          obs = env.reset()
+
+      return obs
+
+    return np.stack([reset_with_noops(self._envs[index]) for index in indices])
 
   def close(self):
     for env in self._envs:
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index b187090c9..198d4e31b 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -85,34 +85,23 @@ def simulate(self, action):
       return tf.identity(reward), tf.identity(done)
 
 
-def define_collect(hparams, scope, eval_phase,
-                   collect_level=-1,
-                   policy_to_actions_lambda=None):
+def define_collect(hparams, scope):
   """Collect trajectories.
 
   Args:
     hparams: HParams.
     scope: var scope.
-    eval_phase: bool, is eval phase.
-    collect_level: int, which level to collect observations.
-    policy_to_actions_lambda: lambda.
 
   Returns:
     Returns memory (observtions, rewards, dones, actions,
     pdfs, values_functions)
-    containing a rollout of environment from collect_level of nested wrapper
-    structure. Note that pdfs and values_functions are meaningful only if
-    collect_level==-1.
+    containing a rollout of environment from nested wrapped structure.
   """
 
   to_initialize = []
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
     environment_spec = hparams.environment_spec
     num_agents = hparams.num_agents
-    if eval_phase:
-      environment_spec = getattr(hparams, "environment_eval_spec",
-                                 environment_spec)
-      num_agents = getattr(hparams, "num_eval_agents", num_agents)
     if environment_spec.simulated_env:
       batch_env = SimulatedBatchEnv(environment_spec, num_agents)
     else:
@@ -121,11 +110,7 @@ def define_collect(hparams, scope, eval_phase,
     to_initialize.append(batch_env)
     environment_wrappers = environment_spec.wrappers
     wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
-    # Put memory wrapper at the level you want to gather observations at.
-    # Negative indices need to be shifted for insert to work correctly.
-    collect_level = collect_level if \
-      collect_level >= 0 else len(wrappers) + collect_level + 1
-    wrappers.insert(collect_level, [_MemoryWrapper, {}])
+    wrappers.append((_MemoryWrapper, {}))
     rollout_metadata = None
     speculum = None
     for w in wrappers:
@@ -133,9 +118,9 @@ def define_collect(hparams, scope, eval_phase,
                       % (str(w[0]), str(w[1]), str(batch_env)))
       batch_env = w[0](batch_env, **w[1])
       to_initialize.append(batch_env)
-      if w[0] == _MemoryWrapper:
-        rollout_metadata = _rollout_metadata(batch_env)
-        speculum = batch_env.speculum
+
+    rollout_metadata = _rollout_metadata(batch_env)
+    speculum = batch_env.speculum
 
     def initialization_lambda(sess):
       for batch_env in to_initialize:
@@ -152,7 +137,7 @@ def initialization_lambda(sess):
     cumulative_rewards = tf.get_variable("cumulative_rewards", len(batch_env),
                                          trainable=False)
 
-    eval_phase_t = tf.convert_to_tensor(eval_phase)
+    eval_phase_t = tf.convert_to_tensor(hparams.eval_phase)
     should_reset_var = tf.Variable(True, trainable=False)
     zeros_tensor = tf.zeros(len(batch_env))
 
@@ -185,12 +170,7 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
         """Step of the environment."""
         actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
         policy = actor_critic.policy
-        if policy_to_actions_lambda:
-          action = policy_to_actions_lambda(policy)
-        else:
-          action = tf.cond(eval_phase_t,
-                           policy.mode,
-                           policy.sample)
+        action = hparams.policy_to_actions_lambda(policy)
 
         postprocessed_action = actor_critic.action_postprocessing(action)
         reward, done = batch_env.simulate(postprocessed_action[0, ...])
@@ -279,7 +259,8 @@ def stop_condition(i, _, resets):
     # When generating real data together with PPO training we must use single
     # agent. For PPO to work we reshape the history, as if it was generated
     # by real_ppo_effective_num_agents.
-    if getattr(hparams, "effective_num_agents", None) and not eval_phase:
+    if (getattr(hparams, "effective_num_agents", None) and
+        not hparams.eval_phase):
       new_memory = []
       effective_num_agents = hparams.effective_num_agents
       assert hparams.epoch_length % effective_num_agents == 0, (
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index d3fad351c..4abcfe9a0 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import copy
 import os
 
 from tensor2tensor import models  # pylint: disable=unused-import
@@ -31,9 +32,13 @@
 
 def define_train(hparams):
   """Define the training setup."""
+  train_hparams = copy.copy(hparams)
+  train_hparams.add_hparam("eval_phase", False)
+  train_hparams.add_hparam(
+      "policy_to_actions_lambda", lambda policy: policy.sample()
+  )
+
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-    memory, collect_summary, train_initialization = collect.define_collect(
-        hparams, "ppo_train", eval_phase=False)
     ppo_summary = ppo.define_ppo_epoch(memory, hparams)
     train_summary = tf.summary.merge([collect_summary, ppo_summary])
 
@@ -119,3 +124,23 @@ def train(hparams, event_dir=None, model_dir=None,
             ckpt_path = os.path.join(
                 model_dir, "model.ckpt-{}".format(epoch_index + 1 + start_step))
             model_saver.save(sess, ckpt_path)
+
+
+def evaluate(hparams, model_dir, name_scope="rl_eval"):
+  """Evaluate."""
+  hparams = copy.copy(hparams)
+  hparams.add_hparam("eval_phase", True)
+  with tf.Graph().as_default():
+    with tf.name_scope(name_scope):
+      (collect_memory, _, collect_init) = collect.define_collect(
+          hparams, "ppo_eval"
+      )
+      model_saver = tf.train.Saver(
+          tf.global_variables(".*network_parameters.*")
+      )
+
+      with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+        collect_init(sess)
+        trainer_lib.restore_checkpoint(model_dir, model_saver, sess)
+        sess.run(collect_memory)
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 80961ddff..1f1bd9c7f 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -26,13 +26,16 @@
 from __future__ import division
 from __future__ import print_function
 
+import copy
 import datetime
 import math
 import os
+import pprint
 import random
 import time
 
 import numpy as np
+import six
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
 from tensor2tensor.data_generators.gym_env import T2TGymEnv
@@ -283,7 +286,7 @@ def train_world_model(
   return world_model_steps_num
 
 
-def setup_env(hparams):
+def setup_env(hparams, batch_size):
   """Setup."""
   game_mode = "Deterministic-v4"
   camel_game_name = "".join(
@@ -292,19 +295,71 @@ def setup_env(hparams):
   env_name = camel_game_name
 
   env = T2TGymEnv(base_env_name=env_name,
-                  batch_size=hparams.real_ppo_num_agents,
+                  batch_size=batch_size,
                   grayscale=hparams.grayscale,
                   resize_width_factor=hparams.resize_width_factor,
                   resize_height_factor=hparams.resize_height_factor,
-                  base_env_timesteps_limit=hparams.env_timesteps_limit)
+                  base_env_timesteps_limit=hparams.env_timesteps_limit,
+                  max_num_noops=hparams.max_num_noops)
   return env
 
 
-def eval_reward(env, clipped):
-  """Calculates mean rewards from given epoch."""
+def evaluate_single_config(hparams, agent_model_dir):
+  """Evaluate the PPO agent in the real environment."""
+  eval_hparams = trainer_lib.create_hparams(hparams.ppo_params)
+  eval_hparams.num_agents = hparams.num_agents
+  env = setup_env(hparams, batch_size=hparams.num_agents)
+  environment_spec = rl.standard_atari_env_spec(env)
+  eval_hparams.add_hparam("environment_spec", environment_spec)
+  eval_hparams.add_hparam(
+      "policy_to_actions_lambda", hparams.policy_to_actions_lambda
+  )
+
+  env.start_new_epoch(0)
+  rl_trainer_lib.evaluate(eval_hparams, agent_model_dir)
+  rollouts = env.current_epoch_rollouts()[:hparams.num_agents]
+  env.close()
+
+  assert len(rollouts) == hparams.num_agents
+  return tuple(
+      compute_mean_reward(rollouts, clipped) for clipped in (True, False)
+  )
+
+
+def evaluate_all_configs(hparams, agent_model_dir):
+  """Evaluate the agent with multiple eval configurations."""
+  def make_eval_hparams(hparams, policy_to_action, max_num_noops):
+    hparams = copy.copy(hparams)
+    hparams.add_hparam("num_agents", hparams.eval_num_agents)
+    hparams.add_hparam("policy_to_actions_lambda", {
+        "sample": lambda policy: policy.sample(),
+        "mode": lambda policy: policy.mode()
+    }[policy_to_action])
+    hparams.max_num_noops = max_num_noops
+    return hparams
+
+  metrics = {}
+  # Iterate over all combinations of picking actions by sampling/mode and
+  # whether to do initial no-ops.
+  for policy_to_action in ("mode", "sample"):
+    for max_num_noops in (hparams.eval_max_num_noops, 0):
+      eval_hparams = make_eval_hparams(hparams, policy_to_action, max_num_noops)
+      scores = evaluate_single_config(eval_hparams, agent_model_dir)
+      for (score, clipped) in zip(scores, (True, False)):
+        metric_name = "mean_reward/eval/{}_{}_max_noops_{}".format(
+            policy_to_action, max_num_noops,
+            "clipped" if clipped else "unclipped"
+        )
+        metrics[metric_name] = score
+
+  return metrics
+
+
+def compute_mean_reward(rollouts, clipped):
+  """Calculate mean rewards from given epoch."""
   reward_name = "reward" if clipped else "unclipped_reward"
   rewards = []
-  for rollout in env.current_epoch_rollouts():
+  for rollout in rollouts:
     if rollout[-1].done:
       rollout_reward = sum(getattr(frame, reward_name) for frame in rollout)
       rewards.append(rollout_reward)
@@ -315,6 +370,15 @@ def eval_reward(env, clipped):
   return mean_rewards
 
 
+def summarize_metrics(eval_metrics_writer, metrics, epoch):
+  """Write metrics to summary."""
+  for (name, value) in six.iteritems(metrics):
+    summary = tf.Summary()
+    summary.value.add(tag=name, simple_value=value)
+    eval_metrics_writer.add_summary(summary, epoch)
+  eval_metrics_writer.flush()
+
+
 def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   """Run the main training loop."""
   if report_fn:
@@ -326,7 +390,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   epoch = -1
   data_dir = directories["data"]
-  env = setup_env(hparams)
+  env = setup_env(hparams, batch_size=hparams.real_ppo_num_agents)
   env.start_new_epoch(epoch, data_dir)
 
   # Timing log function
@@ -334,6 +398,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   # Per-epoch state
   epoch_metrics = []
+  metrics = {}
 
   # Collect data from the real environment with PPO or random policy.
   # TODO(lukaszkaiser): do we need option not to gather_ppo_real_env_data?
@@ -347,20 +412,17 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       env, ppo_model_dir, ppo_event_dir, data_dir, hparams, ppo_epochs_num=0,
       epoch=epoch, is_final_epoch=False
   )
-  mean_unclipped_reward = eval_reward(env, clipped=False)
-  tf.logging.info("Mean reward (initial): {}".format(mean_unclipped_reward))
+  metrics["mean_reward/train/clipped"] = compute_mean_reward(
+      env.current_epoch_rollouts(), clipped=True
+  )
+  tf.logging.info("Mean training reward (initial): {}".format(
+      metrics["mean_reward/train/clipped"]
+  ))
 
   eval_metrics_event_dir = os.path.join(directories["world_model"],
                                         "eval_metrics_event_dir")
   eval_metrics_writer = tf.summary.FileWriter(eval_metrics_event_dir)
 
-  mean_unclipped_reward_summary = tf.Summary()
-  mean_unclipped_reward_summary.value.add(tag="mean_unclipped_reward",
-                                          simple_value=None)
-  mean_clipped_reward_summary = tf.Summary()
-  mean_clipped_reward_summary.value.add(tag="mean_clipped_reward",
-                                        simple_value=None)
-
   world_model_steps_num = 0
 
   for epoch in range(hparams.epochs):
@@ -401,24 +463,19 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
     if hparams.stop_loop_early:
       return 0.0
-    mean_clipped_reward = eval_reward(env, clipped=True)
-    log("Mean clipped reward during generation: {}".format(
-        mean_clipped_reward))  # this was output of generate_real_env_data(...)
 
-    mean_unclipped_reward = eval_reward(env, clipped=False)
-    log("Mean eval reward (unclipped): {}".format(mean_unclipped_reward))
+    metrics["mean_reward/train/clipped"] = compute_mean_reward(
+        env.current_epoch_rollouts(), clipped=True
+    )
+    log("Mean training reward: {}".format(metrics["mean_reward/train/clipped"]))
 
-    # Summarize metrics.
-    mean_unclipped_reward_summary.value[0].simple_value = mean_unclipped_reward
-    eval_metrics_writer.add_summary(mean_unclipped_reward_summary, epoch)
-    mean_clipped_reward_summary.value[0].simple_value = int(mean_clipped_reward)
-    eval_metrics_writer.add_summary(mean_clipped_reward_summary, epoch)
-    eval_metrics_writer.flush()
+    eval_metrics = evaluate_all_configs(hparams, ppo_model_dir)
+    log("Eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
+    metrics.update(eval_metrics)
+    summarize_metrics(eval_metrics_writer, metrics, epoch)
 
     # Report metrics
-    eval_metrics = {"mean_reward": mean_unclipped_reward}
-    epoch_metrics.append(eval_metrics)
-    log("Eval metrics: %s", str(eval_metrics))
+    epoch_metrics.append(metrics)
     if report_fn:
       report_fn(eval_metrics[report_metric], epoch)
 
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index c1367dad9..dbc09a6be 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -84,6 +84,8 @@ def rlmb_base():
       resize_height_factor=2,
       resize_width_factor=2,
       grayscale=False,
+      # Maximum number of noops to make on environment reset.
+      max_num_noops=0,
       # Bump learning rate after first epoch by 3x.
       # We picked 3x because our default learning rate schedule decreases with
       # 1/square root of step; 1/sqrt(10k) = 0.01 and 1/sqrt(100k) ~ 0.0032
@@ -102,6 +104,9 @@ def rlmb_base():
       real_ppo_effective_num_agents=16,
       real_ppo_eval_every_epochs=0,
 
+      eval_num_agents=30,
+      eval_max_num_noops=8,
+
       game="pong",
       # Whether to evaluate the world model in each iteration of the loop to get
       # the model_reward_accuracy metric.
@@ -390,6 +395,7 @@ def rlmb_tiny():
           real_ppo_num_agents=1,
           real_ppo_epochs_num=0,
           real_ppo_effective_num_agents=2,
+          eval_num_agents=1,
           generative_model_params="next_frame_tiny",
           stop_loop_early=True,
           resize_height_factor=2,

From 2e53520c6231f7e72ec84b31303e185b92aac3bd Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 25 Oct 2018 20:47:55 -0700
Subject: [PATCH 1081/2720] internal merge of PR #1168

PiperOrigin-RevId: 218798566
---
 tensor2tensor/data_generators/gym_env.py       |  6 +++---
 tensor2tensor/rl/rl_trainer_lib.py             | 16 ++++++++++++++--
 tensor2tensor/rl/trainer_model_based.py        |  6 +++++-
 tensor2tensor/rl/trainer_model_based_params.py |  2 +-
 4 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index dfb75aecd..0da218001 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -21,11 +21,10 @@
 
 import collections
 import itertools
-
+import random
 import gym
 from gym.spaces import Box
 import numpy as np
-import random
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -554,7 +553,7 @@ def __init__(self, base_env_name=None, batch_size=1, grayscale=False,
     if max_num_noops > 0:
       assert self._envs[0].unwrapped.get_action_meanings()[
           self.noop_action
-      ] == 'NOOP'
+      ] == "NOOP"
     self.max_num_noops = max_num_noops
 
     orig_observ_space = self._envs[0].observation_space
@@ -616,6 +615,7 @@ def _step(self, actions):
 
   def _reset(self, indices):
     def reset_with_noops(env):
+      """Resets environment and applies random number of NOOP actions on it."""
       obs = env.reset()
       try:
         num_noops = random.randint(1, self.max_num_noops)
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 4abcfe9a0..30b3d30c9 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -39,12 +39,24 @@ def define_train(hparams):
   )
 
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+    memory, collect_summary, train_initialization = (
+        collect.define_collect(train_hparams, "ppo_train")
+    )
     ppo_summary = ppo.define_ppo_epoch(memory, hparams)
     train_summary = tf.summary.merge([collect_summary, ppo_summary])
 
     if hparams.eval_every_epochs:
-      _, eval_collect_summary, eval_initialization = collect.define_collect(
-          hparams, "ppo_eval", eval_phase=True)
+      eval_hparams = copy.copy(hparams)
+      eval_hparams.add_hparam("eval_phase", True)
+      eval_hparams.add_hparam(
+          "policy_to_actions_lambda", lambda policy: policy.mode()
+      )
+      eval_hparams.environment_spec = hparams.environment_eval_spec
+      eval_hparams.num_agents = hparams.num_eval_agents
+
+      _, eval_collect_summary, eval_initialization = (
+          collect.define_collect(eval_hparams, "ppo_eval")
+      )
       return train_summary, eval_collect_summary, (train_initialization,
                                                    eval_initialization)
     else:
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 1f1bd9c7f..24d4bbc12 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -477,7 +477,11 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     # Report metrics
     epoch_metrics.append(metrics)
     if report_fn:
-      report_fn(eval_metrics[report_metric], epoch)
+      if report_metric == "mean_reward":
+        report_fn(eval_metrics["mean_reward/eval/{}_{}_max_noops_{}".format(
+            "mode", hparams.eval_max_num_noops, "unclipped")], epoch)
+      else:
+        report_fn(eval_metrics[report_metric], epoch)
 
   # Return the evaluation metrics from the final epoch
   return epoch_metrics[-1]
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index dbc09a6be..bbe9ccfff 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -85,7 +85,7 @@ def rlmb_base():
       resize_width_factor=2,
       grayscale=False,
       # Maximum number of noops to make on environment reset.
-      max_num_noops=0,
+      max_num_noops=8,
       # Bump learning rate after first epoch by 3x.
       # We picked 3x because our default learning rate schedule decreases with
       # 1/square root of step; 1/sqrt(10k) = 0.01 and 1/sqrt(100k) ~ 0.0032

From df07bcf000631fd539e9a6eb46ea0b8ccd348c47 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 26 Oct 2018 09:29:18 -0700
Subject: [PATCH 1082/2720] Rename function to pass OSS lint.

PiperOrigin-RevId: 218865306
---
 tensor2tensor/data_generators/video_utils_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index 9ef46b65d..66a6394a8 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -29,7 +29,7 @@
 
 class VideoUtilsTest(tf.test.TestCase):
 
-  def getPredictions(self):
+  def get_predictions(self):
     rng = np.random.RandomState(0)
     # num_samples=4
     inputs = rng.randint(0, 255, (4, 2, 64, 64, 3))
@@ -46,7 +46,7 @@ def getPredictions(self):
     return predictions, problem
 
   def testDecodeInMemoryTrue(self):
-    predictions, problem = self.getPredictions()
+    predictions, problem = self.get_predictions()
     decode_hparams = decoding.decode_hparams()
     decode_hparams.decode_in_memory = True
     decode_hooks = decoding.DecodeHookArgs(

From af42d543c2f24a0143b2483db93ac931c54146b9 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 26 Oct 2018 10:38:10 -0700
Subject: [PATCH 1083/2720] A couple tiny fixes to common_attention.py

PiperOrigin-RevId: 218876014
---
 tensor2tensor/layers/common_attention.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 42114f88e..1f3bf04e6 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -634,10 +634,11 @@ def add_positional_embedding(x, max_length, name, positions=None):
     _, length, depth = common_layers.shape_list(x)
     var = tf.cast(tf.get_variable(name, [max_length, depth]), x.dtype)
     if positions is None:
+      pad_length = tf.maximum(0, length - max_length)
       sliced = tf.cond(
           tf.less(length, max_length),
           lambda: tf.slice(var, [0, 0], [length, -1]),
-          lambda: tf.pad(var, [[0, max(0, length - max_length)], [0, 0]]))
+          lambda: tf.pad(var, [[0, pad_length], [0, 0]]))
       return x + tf.expand_dims(sliced, 0)
     else:
       return x + tf.gather(var, tf.to_int32(positions))
@@ -1765,7 +1766,7 @@ def dot_product_self_attention_relative_v2(q,
     logits = tf.matmul(q, k, transpose_b=True)
     key_relative_embeddings = get_relative_embeddings_left(
         max_relative_position, length, depth_k, num_heads,
-        heads_share_relative_embedding, "key_relative_embededings")
+        heads_share_relative_embedding, "key_relative_embeddings")
 
     rel_logits = matmul_with_relative_keys(q, key_relative_embeddings,
                                            heads_share_relative_embedding)

From 693d3ff77dc1bf4bd1b8ff7e4dd3fed7849f9978 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 26 Oct 2018 11:00:52 -0700
Subject: [PATCH 1084/2720] Update a missed caller while changing
 create_run_config.

PiperOrigin-RevId: 218880374
---
 tensor2tensor/bin/t2t_attack.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 45cc22e42..4ccedc779 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -99,6 +99,7 @@ def create_surrogate_run_config(hp):
       hp.daisy_chain_variables and hp.activation_dtype == "float32" and
       hp.weight_dtype == "float32")
   return trainer_lib.create_run_config(
+      model_name=FLAGS.model,
       model_dir=os.path.expanduser(FLAGS.surrogate_output_dir),
       master=FLAGS.master,
       iterations_per_loop=FLAGS.iterations_per_loop,

From 74851c6543e707a947ef3ceb1026a15349dd9306 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 26 Oct 2018 11:08:55 -0700
Subject: [PATCH 1085/2720] More lint fixes.

PiperOrigin-RevId: 218882279
---
 tensor2tensor/layers/common_video_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index 2e41cd98a..6dc9751b7 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -122,7 +122,7 @@ def testGifSummary(self):
 
       self.assertEqual(encoded, common_video._encode_gif(images[0], fps=10))  # pylint: disable=protected-access
 
-  def checkIfPatchExists(self, videos, video_patches, num_frames):
+  def check_if_patch_exists(self, videos, video_patches, num_frames):
     """Check that given patch is present in video."""
     for video, video_patch in zip(videos, video_patches):
       total_frames = len(video)
@@ -145,7 +145,7 @@ def testExtractRandomVideoPatch(self, num_frames=2):
         video_patch_np = sess.run(video_patch)
         if num_frames != -1:
           self.assertEqual(video_patch_np.shape, (12, num_frames, 256, 256, 3))
-          self.checkIfPatchExists(video_np, video_patch_np, num_frames)
+          self.check_if_patch_exists(video_np, video_patch_np, num_frames)
         else:
           self.assertTrue(np.allclose(video_np, video_patch_np))
 

From a15c85c73e91503aaf56ea0b377bfe8d6a1090c7 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 26 Oct 2018 11:10:09 -0700
Subject: [PATCH 1086/2720] PR #1152 added ppo_epochs_num to
 trainer_model_based.py's train_agent

PiperOrigin-RevId: 218882494
---
 tensor2tensor/rl/trainer_model_based_agent_only.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/rl/trainer_model_based_agent_only.py b/tensor2tensor/rl/trainer_model_based_agent_only.py
index 6d9cc99bb..8c61a8d33 100644
--- a/tensor2tensor/rl/trainer_model_based_agent_only.py
+++ b/tensor2tensor/rl/trainer_model_based_agent_only.py
@@ -57,6 +57,7 @@ def main(_):
   agent_model_dir = FLAGS.output_dir
   event_dir = FLAGS.output_dir
   epoch_data_dir = FLAGS.data_dir  # only required for initial frames
+  ppo_epochs_num = hparams.ppo_epochs_num
 
   trainer_model_based.train_agent(
       problem_name,
@@ -65,6 +66,7 @@ def main(_):
       world_model_dir,
       epoch_data_dir,
       hparams,
+      ppo_epochs_num,
       epoch=0,
       is_final_epoch=True)
 

From c0673737dfacfb090c44e25114205e9f5fe636b0 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 26 Oct 2018 11:51:55 -0700
Subject: [PATCH 1087/2720] Allow sharded decoding from- and to-file.

PiperOrigin-RevId: 218889905
---
 tensor2tensor/utils/decoding.py | 34 +++++++++++++++++----------------
 1 file changed, 18 insertions(+), 16 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 5d3797028..c5619ff5a 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -62,8 +62,9 @@ def decode_hparams(overrides=""):
       delimiter="\n",
       decode_to_file=None,
       decode_in_memory=False,
-      shards=1,
-      shard_id=0,
+      shards=1,    # How many shards of data to decode (treating 1 as None).
+      shard_id=0,  # Which shard are we decoding if more than 1 above.
+      shards_start_offset=0,  # Number of the first shard to decode.
       num_decodes=1,
       force_decode_length=False,
       display_decoded_images=False,
@@ -348,9 +349,9 @@ def decode_from_file(estimator,
   inputs_vocab = p_hp.vocabulary[inputs_vocab_key]
   targets_vocab = p_hp.vocabulary["targets"]
   problem_name = FLAGS.problem
-  tf.logging.info("Performing decoding from a file.")
-  sorted_inputs, sorted_keys = _get_sorted_inputs(filename, decode_hp.shards,
-                                                  decode_hp.delimiter)
+  filename = _add_shard_to_filename(filename, decode_hp)
+  tf.logging.info("Performing decoding from file (%s)." % filename)
+  sorted_inputs, sorted_keys = _get_sorted_inputs(filename, decode_hp.delimiter)
   num_decode_batches = (len(sorted_inputs) - 1) // decode_hp.batch_size + 1
 
   def input_fn():
@@ -435,6 +436,8 @@ def timer(gen):
   decode_filename = decode_to_file if decode_to_file else filename
   if not decode_to_file:
     decode_filename = _decode_filename(decode_filename, problem_name, decode_hp)
+  else:
+    decode_filename = _add_shard_to_filename(decode_filename, decode_hp)
   tf.logging.info("Writing decodes into %s" % decode_filename)
   outfile = tf.gfile.Open(decode_filename, "w")
   for index in range(len(sorted_inputs)):
@@ -455,6 +458,13 @@ def timer(gen):
   ), None)
 
 
+def _add_shard_to_filename(filename, decode_hp):
+  if decode_hp.shards > 1:
+    shard_id = decode_hp.shard_id + decode_hp.shards_start_offset
+    filename = filename + ("%.3d" % shard_id)
+  return filename
+
+
 def _decode_filename(base_filename, problem_name, decode_hp):
   """Generates decode filename.
 
@@ -467,7 +477,7 @@ def _decode_filename(base_filename, problem_name, decode_hp):
     A string, produced decode filename.
   """
   if decode_hp.shards > 1:
-    base_filename = base_filename + ("%.2d" % decode_hp.shard_id)
+    base_filename = _add_shard_to_filename(base_filename, decode_hp)
   if ("beam{beam}.alpha{alpha}.decodes".format(
       beam=str(decode_hp.beam_size), alpha=str(decode_hp.alpha))
       in base_filename):
@@ -692,13 +702,11 @@ def show_and_save_image(img, save_path):
     plt.savefig(sp)
 
 
-def _get_sorted_inputs(filename, num_shards=1, delimiter="\n"):
+def _get_sorted_inputs(filename, delimiter="\n"):
   """Returning inputs sorted according to length.
 
   Args:
     filename: path to file with inputs, 1 per line.
-    num_shards: number of input shards. If > 1, will read from file filename.XX,
-      where XX is FLAGS.worker_id.
     delimiter: str, delimits records in the file.
 
   Returns:
@@ -706,13 +714,7 @@ def _get_sorted_inputs(filename, num_shards=1, delimiter="\n"):
 
   """
   tf.logging.info("Getting sorted inputs")
-  # read file and sort inputs according them according to input length.
-  if num_shards > 1:
-    decode_filename = filename + ("%.2d" % FLAGS.worker_id)
-  else:
-    decode_filename = filename
-
-  with tf.gfile.Open(decode_filename) as f:
+  with tf.gfile.Open(filename) as f:
     text = f.read()
     records = text.split(delimiter)
     inputs = [record.strip() for record in records]

From 904c42a5d5edbaefe0c36ab278abc070cf90f5a7 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 26 Oct 2018 15:12:20 -0700
Subject: [PATCH 1088/2720] Scale up glow part 2 of N.

PiperOrigin-RevId: 218921362
---
 tensor2tensor/models/research/glow.py         | 22 +++++++------------
 .../models/research/glow_init_hook.py         |  5 ++++-
 tensor2tensor/models/research/glow_test.py    |  2 ++
 3 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 4f788a4bd..c21bb2997 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -54,7 +54,6 @@ def glow_hparams():
   # init_batch_size denotes the number of examples used for data-dependent
   # initialization. A higher init_batch_size is required for training
   # stability especially when hparams.batch_size is low.
-  # -1 indicates that the init_batch_size is set to hparams.batch_size.
   hparams.add_hparam("init_batch_size", 256)
   return hparams
 
@@ -65,6 +64,10 @@ class Glow(t2t_model.T2TModel):
 
   Reference: https://arxiv.org/abs/1807.03039"""
 
+  def init_preprocess(self, features):
+    """Preprocessing as per the input modality."""
+    return features
+
   def preprocess(self, x):
     """Normalize x.
 
@@ -119,21 +122,12 @@ def create_init_batch(self, features):
     Returns:
       init_features: initialization features.
     """
-    # TODO(mechcoder) Once all depending code supports hparams.init_batch_size
-    # the if block can be removed.
-    if self.hparams.init_batch_size == -1:
-      return features
-
-    train_dataset = self.hparams.problem.dataset(tf.estimator.ModeKeys.TRAIN)
-    train_dataset = train_dataset.batch(self.init_batch_size)
+    train_dataset = self.hparams.problem.dataset(
+        tf.estimator.ModeKeys.TRAIN, hparams=self.hparams)
+    train_dataset = train_dataset.batch(self.hparams.init_batch_size)
+    train_dataset = self.init_preprocess(train_dataset)
     return train_dataset.make_one_shot_iterator().get_next()
 
-  @property
-  def init_batch_size(self):
-    if self.hparams.init_batch_size == -1:
-      return self.hparams.batch_size
-    return self.hparams.init_batch_size
-
   @staticmethod
   def train_hooks(hook_context):
     del hook_context
diff --git a/tensor2tensor/models/research/glow_init_hook.py b/tensor2tensor/models/research/glow_init_hook.py
index 2171253ff..0a88c973b 100644
--- a/tensor2tensor/models/research/glow_init_hook.py
+++ b/tensor2tensor/models/research/glow_init_hook.py
@@ -35,4 +35,7 @@ def after_create_session(self, session, coord):
     global_step = session.run(tf.train.get_global_step())
     if global_step == 0:
       ddi = tf.get_collection("glow_init_op")
-      session.run(ddi)
+      # In-case of a multi-GPU system, this just runs the first op in the
+      # collection.
+      if ddi:
+        session.run(ddi[0])
diff --git a/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
index 1d623e018..12414b857 100644
--- a/tensor2tensor/models/research/glow_test.py
+++ b/tensor2tensor/models/research/glow_test.py
@@ -48,6 +48,7 @@ def test_glow(self):
       hparams.n_levels = 2
       hparams.init_batch_size = 256
       hparams.batch_size = 1
+      hparams.data_dir = ''
       cifar_problem = problems.problem('image_cifar10_plain_random_shift')
       hparams.problem = cifar_problem
       model = glow.Glow(hparams, tf.estimator.ModeKeys.TRAIN)
@@ -76,6 +77,7 @@ def test_glow_inference(self):
     hparams = glow.glow_hparams()
     hparams.depth = 15
     hparams.n_levels = 2
+    hparams.data_dir = ''
     curr_dir = tempfile.mkdtemp()
 
     # Training pipeline

From cd030fc3bfa92618d3c7337be9f3c3653750f1e9 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Fri, 26 Oct 2018 15:14:03 -0700
Subject: [PATCH 1089/2720] fixing call to trainer_model_based.train_agent

PiperOrigin-RevId: 218921592
---
 tensor2tensor/rl/trainer_model_based.py       | 33 ++++++++++---------
 .../rl/trainer_model_based_agent_only.py      |  3 +-
 2 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 24d4bbc12..ddc2f3bfd 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -133,7 +133,8 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
 
 
 def train_agent(real_env, agent_model_dir, event_dir, world_model_dir, data_dir,
-                hparams, ppo_epochs_num, epoch=0, is_final_epoch=False):
+                hparams, completed_ppo_epochs_num, epoch=0,
+                is_final_epoch=False):
   """Train the PPO agent in the simulated environment."""
   del data_dir
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
@@ -146,8 +147,8 @@ def train_agent(real_env, agent_model_dir, event_dir, world_model_dir, data_dir,
     if ppo_param_name in hparams:
       ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
 
-  ppo_epochs_num += sim_ppo_epoch_increment(hparams, is_final_epoch)
-  ppo_hparams.epochs_num = ppo_epochs_num
+  completed_ppo_epochs_num += sim_ppo_epoch_increment(hparams, is_final_epoch)
+  ppo_hparams.epochs_num = completed_ppo_epochs_num
 
   ppo_hparams.save_models_every_epochs = 10
   ppo_hparams.world_model_dir = world_model_dir
@@ -221,12 +222,12 @@ def choose_initial_frames():
     rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
                          name_scope="ppo_sim%d" % (epoch + 1))
 
-  return ppo_epochs_num
+  return completed_ppo_epochs_num
 
 
 def train_agent_real_env(
     env, agent_model_dir, event_dir, data_dir,
-    hparams, ppo_epochs_num, epoch=0, is_final_epoch=False):
+    hparams, completed_ppo_epochs_num, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the real environment."""
   del is_final_epoch, data_dir
   ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
@@ -241,8 +242,8 @@ def train_agent_real_env(
     if ppo_param_name in hparams:
       ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
 
-  ppo_epochs_num += real_ppo_epoch_increment(hparams)
-  ppo_hparams.epochs_num = ppo_epochs_num
+  completed_ppo_epochs_num += real_ppo_epoch_increment(hparams)
+  ppo_hparams.epochs_num = completed_ppo_epochs_num
   # We do not save model, as that resets frames that we need at restarts.
   # But we need to save at the last step, so we set it very high.
   ppo_hparams.save_models_every_epochs = 1000000
@@ -257,7 +258,7 @@ def train_agent_real_env(
   # Save unfinished rollouts to history.
   env.reset()
 
-  return ppo_epochs_num
+  return completed_ppo_epochs_num
 
 
 def train_world_model(
@@ -408,9 +409,9 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   tf.logging.info("Initial training of PPO in real environment.")
   ppo_event_dir = os.path.join(directories["world_model"],
                                "ppo_summaries/initial")
-  ppo_epochs_num = train_agent_real_env(
-      env, ppo_model_dir, ppo_event_dir, data_dir, hparams, ppo_epochs_num=0,
-      epoch=epoch, is_final_epoch=False
+  completed_ppo_epochs_num = train_agent_real_env(
+      env, ppo_model_dir, ppo_event_dir, data_dir, hparams,
+      completed_ppo_epochs_num=0, epoch=epoch, is_final_epoch=False
   )
   metrics["mean_reward/train/clipped"] = compute_mean_reward(
       env.current_epoch_rollouts(), clipped=True
@@ -446,9 +447,9 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     if not hparams.ppo_continue_training:
       ppo_model_dir = ppo_event_dir
 
-    ppo_epochs_num = train_agent(
+    completed_ppo_epochs_num = train_agent(
         env, ppo_model_dir, ppo_event_dir,
-        directories["world_model"], data_dir, hparams, ppo_epochs_num,
+        directories["world_model"], data_dir, hparams, completed_ppo_epochs_num,
         epoch=epoch, is_final_epoch=is_final_epoch
     )
 
@@ -456,9 +457,9 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
     # Train PPO on real env (short)
     log("Training PPO in real environment.")
-    ppo_epochs_num = train_agent_real_env(
-        env, ppo_model_dir, ppo_event_dir, data_dir, hparams, ppo_epochs_num,
-        epoch=epoch, is_final_epoch=is_final_epoch
+    completed_ppo_epochs_num = train_agent_real_env(
+        env, ppo_model_dir, ppo_event_dir, data_dir, hparams,
+        completed_ppo_epochs_num, epoch=epoch, is_final_epoch=is_final_epoch
     )
 
     if hparams.stop_loop_early:
diff --git a/tensor2tensor/rl/trainer_model_based_agent_only.py b/tensor2tensor/rl/trainer_model_based_agent_only.py
index 8c61a8d33..801a54df2 100644
--- a/tensor2tensor/rl/trainer_model_based_agent_only.py
+++ b/tensor2tensor/rl/trainer_model_based_agent_only.py
@@ -57,7 +57,6 @@ def main(_):
   agent_model_dir = FLAGS.output_dir
   event_dir = FLAGS.output_dir
   epoch_data_dir = FLAGS.data_dir  # only required for initial frames
-  ppo_epochs_num = hparams.ppo_epochs_num
 
   trainer_model_based.train_agent(
       problem_name,
@@ -66,7 +65,7 @@ def main(_):
       world_model_dir,
       epoch_data_dir,
       hparams,
-      ppo_epochs_num,
+      0,
       epoch=0,
       is_final_epoch=True)
 

From 4a460cf4574a676c4750f99c5b4fa1531669f60d Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 26 Oct 2018 15:24:11 -0700
Subject: [PATCH 1090/2720] Separate vocab size from modalities in Problem.

PiperOrigin-RevId: 218922971
---
 tensor2tensor/data_generators/algorithmic.py  |   7 +-
 tensor2tensor/data_generators/allen_brain.py  |   7 +-
 tensor2tensor/data_generators/babi_qa.py      |   9 +-
 .../data_generators/bair_robot_pushing.py     |  11 +-
 tensor2tensor/data_generators/celeba.py       |   7 +-
 tensor2tensor/data_generators/celebahq.py     |   4 +-
 tensor2tensor/data_generators/cifar.py        |  13 +-
 tensor2tensor/data_generators/fsns.py         |   8 +-
 .../data_generators/gene_expression.py        |   9 +-
 .../data_generators/google_robot_pushing.py   |  12 +-
 tensor2tensor/data_generators/gym_env.py      |  25 ++--
 tensor2tensor/data_generators/ice_parsing.py  |  10 +-
 tensor2tensor/data_generators/image_utils.py  |  15 +-
 tensor2tensor/data_generators/imagenet.py     |   7 +-
 tensor2tensor/data_generators/lambada.py      |  11 +-
 .../data_generators/multi_problem.py          |   7 +-
 .../data_generators/multi_problem_test.py     |  21 +--
 tensor2tensor/data_generators/problem.py      | 108 +++++++-------
 .../data_generators/problem_hparams.py        |  47 ++++---
 tensor2tensor/data_generators/problem_test.py |  35 ++---
 .../data_generators/speech_recognition.py     |   9 +-
 tensor2tensor/data_generators/squad.py        |   3 +-
 .../data_generators/text_problems.py          |  38 ++---
 tensor2tensor/data_generators/timeseries.py   |   7 +-
 .../data_generators/video_generated.py        |  11 +-
 tensor2tensor/data_generators/video_utils.py  |   8 +-
 tensor2tensor/data_generators/vqa.py          |  15 +-
 .../data_generators/wikisum/wikisum.py        |  13 +-
 tensor2tensor/models/mtf_image_transformer.py |   2 +-
 .../models/mtf_image_transformer_test.py      |   2 +-
 tensor2tensor/models/mtf_transformer.py       |   7 +-
 tensor2tensor/models/mtf_transformer_test.py  |   2 +-
 tensor2tensor/models/research/autoencoders.py |  14 +-
 .../models/research/transformer_nat.py        |   4 -
 .../models/research/transformer_vae.py        |   4 -
 .../research/universal_transformer_test.py    |   2 +-
 .../models/research/vqa_attention_test.py     |   6 +-
 tensor2tensor/models/resnet.py                |   2 +-
 tensor2tensor/models/resnet_test.py           |   4 +-
 tensor2tensor/models/slicenet.py              |   3 +-
 tensor2tensor/models/transformer.py           |   8 +-
 tensor2tensor/models/transformer_test.py      |   2 +-
 tensor2tensor/models/video/base.py            |  10 +-
 tensor2tensor/models/video/tests_utils.py     |   8 +-
 tensor2tensor/models/xception_test.py         |   4 +-
 tensor2tensor/utils/data_reader_test.py       |   9 +-
 tensor2tensor/utils/decoding.py               |   2 +-
 tensor2tensor/utils/metrics.py                |   4 +-
 tensor2tensor/utils/t2t_model.py              | 133 +++++++++---------
 tensor2tensor/utils/trainer_lib_test.py       |  12 +-
 50 files changed, 376 insertions(+), 345 deletions(-)

diff --git a/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index 7e386b559..edacffd49 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -25,6 +25,7 @@
 from tensor2tensor.data_generators import generator_utils as utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 import tensorflow as tf
@@ -83,8 +84,10 @@ def generator_eos(nbr_symbols, max_length, nbr_cases):
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
     vocab_size = self.num_symbols + text_encoder.NUM_RESERVED_TOKENS
-    p.input_modality = {"inputs": (registry.Modalities.SYMBOL, vocab_size)}
-    p.target_modality = (registry.Modalities.SYMBOL, vocab_size)
+    p.modality = {"inputs": modalities.SymbolModality,
+                  "targets": modalities.SymbolModality}
+    p.vocab_size = {"inputs": vocab_size,
+                    "targets": vocab_size}
     p.input_space_id = problem.SpaceID.DIGIT_0
     p.target_space_id = problem.SpaceID.DIGIT_1
 
diff --git a/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index 303872e4c..2fece298c 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -39,6 +39,7 @@
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -374,8 +375,10 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {"inputs": ("image:identity", 256)}
-    p.target_modality = ("image:identity", 256)
+    p.modality = {"inputs": modalities.IdentityModality,
+                  "targets": modalities.IdentityModality}
+    p.vocab_size = {"inputs": 256,
+                    "targets": 256}
     p.batch_size_multiplier = 256
     p.input_space_id = problem.SpaceID.IMAGE
     p.target_space_id = problem.SpaceID.IMAGE
diff --git a/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index 14c6fe70c..784508ecf 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -42,6 +42,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.data_generators import tokenizer
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -417,7 +418,8 @@ def hparams(self, defaults, unused_model_hparams):
     (super(BabiQa, self).hparams(defaults, unused_model_hparams))
     p = defaults
     num_classes = self._encoders['targets'].vocab_size
-    p.target_modality = (registry.Modalities.CLASS_LABEL, num_classes)
+    p.modality = {'targets': modalities.ClassLabelModality}
+    p.vocab_size = {'targets': num_classes}
 
   def example_reading_spec(self):
     data_fields, data_items_to_decoders = (
@@ -445,9 +447,10 @@ def preprocess_example(self, example, unused_mode, unused_model_hparams):
     return example
 
   def hparams(self, defaults, unused_model_hparams):
-    (super(BabiQaConcat, self).hparams(defaults, unused_model_hparams))
+    super(BabiQaConcat, self).hparams(defaults, unused_model_hparams)
     p = defaults
-    del p.input_modality['context']
+    del p.modality['context']
+    del p.vocab_size['context']
 
 
 def _problems_to_register():
diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index 857e41c49..0a349d972 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -32,6 +32,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -99,12 +100,10 @@ def extra_reading_spec(self):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {
-        "inputs": ("video", 256),
-    }
-    p.target_modality = {
-        "targets": ("video", 256),
-    }
+    p.modality = {"inputs": modalities.VideoModality,
+                  "targets": modalities.VideoModality}
+    p.vocab_size = {"inputs": 256,
+                    "targets": 256}
 
   def parse_frames(self, filenames):
     image_key = "{}/image_aux1/encoded"
diff --git a/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index b7518a204..3a80718a1 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -23,6 +23,7 @@
 import zipfile
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import image_utils
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -55,8 +56,10 @@ class ImageCeleba(image_utils.ImageProblem):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {"inputs": ("image:identity", 256)}
-    p.target_modality = ("image:identity", 256)
+    p.modality = {"inputs": modalities.IdentityModality,
+                  "targets": modalities.IdentityModality}
+    p.vocab_size = {"inputs": 256,
+                    "targets": 256}
     p.batch_size_multiplier = 256
     p.input_space_id = 1
     p.target_space_id = 1
diff --git a/tensor2tensor/data_generators/celebahq.py b/tensor2tensor/data_generators/celebahq.py
index a206d3fe0..e8a01bfe2 100644
--- a/tensor2tensor/data_generators/celebahq.py
+++ b/tensor2tensor/data_generators/celebahq.py
@@ -21,6 +21,7 @@
 import os
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import problem
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -78,7 +79,8 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.batch_size_multiplier = 1
-    p.input_modality = {"inputs": ("image:identity", 256)}
+    p.modality = {"inputs": modalities.IdentityModality}
+    p.vocab_size = {"inputs": 256}
     p.input_space_id = 1
 
   def preprocess_example(self, example, mode, hparams):
diff --git a/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index 7ffdb869d..d5fcafc04 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -29,6 +29,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import mnist
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -239,8 +240,10 @@ def preprocess_example(self, example, unused_mode, unused_hparams):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {"inputs": ("image:identity", 256)}
-    p.target_modality = ("image:identity", 256)
+    p.modality = {"inputs": modalities.IdentityModality,
+                  "targets": modalities.IdentityModality}
+    p.vocab_size = {"inputs": 256,
+                    "targets": 256}
     p.batch_size_multiplier = 256
     p.input_space_id = 1
     p.target_space_id = 1
@@ -447,8 +450,10 @@ def preprocess_example(self, example, unused_mode, unused_hparams):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {"inputs": ("image:identity", 256)}
-    p.target_modality = ("image:identity", 256)
+    p.modality = {"inputs": modalities.IdentityModality,
+                  "targets": modalities.IdentityModality}
+    p.vocab_size = {"inputs": 256,
+                    "targets": 256}
     p.batch_size_multiplier = 256
     p.max_expected_batch_size_per_shard = 4
     p.input_space_id = 1
diff --git a/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
index 179d7bc85..b482c8a70 100644
--- a/tensor2tensor/data_generators/fsns.py
+++ b/tensor2tensor/data_generators/fsns.py
@@ -24,6 +24,7 @@
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -61,9 +62,10 @@ def feature_encoders(self, data_dir):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {"inputs": (registry.Modalities.IMAGE, 256)}
-    vocab_size = self._encoders["targets"].vocab_size
-    p.target_modality = (registry.Modalities.SYMBOL, vocab_size)
+    p.modality = {"inputs": modalities.ImageModality,
+                  "targets": modalities.SymbolModality}
+    p.vocab_size = {"inputs": 256,
+                    "targets": self._encoders["targets"].vocab_size}
     p.batch_size_multiplier = 256
     p.input_space_id = problem.SpaceID.IMAGE
     p.target_space_id = problem.SpaceID.EN_TOK
diff --git a/tensor2tensor/data_generators/gene_expression.py b/tensor2tensor/data_generators/gene_expression.py
index f9ab3e0e3..4e82000f2 100644
--- a/tensor2tensor/data_generators/gene_expression.py
+++ b/tensor2tensor/data_generators/gene_expression.py
@@ -47,6 +47,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -141,10 +142,10 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    vocab_size = self._encoders["inputs"].vocab_size
-    p.input_modality = {"inputs": (registry.Modalities.SYMBOL, vocab_size)}
-    p.target_modality = ("%s:log_poisson_loss" % registry.Modalities.REAL,
-                         self.num_output_predictions)
+    p.modality = {"inputs": modalities.SymbolModality,
+                  "targets": modalities.RealLogPoissonLossModality}
+    p.vocab_size = {"inputs": self._encoders["inputs"].vocab_size,
+                    "targets": self.num_output_predictions}
     p.input_space_id = problem.SpaceID.DNA
     p.target_space_id = problem.SpaceID.REAL
 
diff --git a/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index e20edb59a..5b295fccf 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -32,6 +32,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -132,10 +133,7 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {
-        # Pixels are in 0..255 range.
-        "inputs": ("video", 256),
-    }
-    p.target_modality = {
-        "targets": ("video", 256),
-    }
+    p.modality = {"inputs": modalities.VideoModality,
+                  "targets": modalities.VideoModality}
+    p.vocab_size = {"inputs": 256,
+                    "targets": 256}
diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 0da218001..cb54e7037 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -29,6 +29,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -84,14 +85,22 @@ def num_rewards(self):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    def make_modality(name):
-      return {
-          "{}s".format(name): ("video", 256),
-          "{}_reward".format(name): ("symbol:weights_all", self.num_rewards),
-          "{}_action".format(name): ("symbol:weights_all", self.num_actions)
-      }
-    p.input_modality = make_modality("input")
-    p.target_modality = make_modality("target")
+    p.modality = {
+        "inputs": modalities.VideoModality,
+        "input_reward": modalities.SymbolModalityWeightsAll,
+        "input_action": modalities.SymbolModalityWeightsAll,
+        "targets": modalities.VideoModality,
+        "target_reward": modalities.SymbolModalityWeightsAll,
+        "target_action": modalities.SymbolModalityWeightsAll,
+    }
+    p.vocab_size = {
+        "inputs": 256,
+        "input_reward": self.num_rewards,
+        "input_action": self.num_actions,
+        "targets": 256,
+        "target_reward": self.num_rewards,
+        "target_action": self.num_actions,
+    }
     p.input_space_id = problem.SpaceID.IMAGE
     p.target_space_id = problem.SpaceID.IMAGE
 
diff --git a/tensor2tensor/data_generators/ice_parsing.py b/tensor2tensor/data_generators/ice_parsing.py
index 59eacb517..fd07da10e 100644
--- a/tensor2tensor/data_generators/ice_parsing.py
+++ b/tensor2tensor/data_generators/ice_parsing.py
@@ -30,6 +30,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
 
@@ -107,11 +108,10 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    source_vocab_size = self._encoders["inputs"].vocab_size
-    p.input_modality = {
-        "inputs": (registry.Modalities.SYMBOL, source_vocab_size)
-    }
-    p.target_modality = (registry.Modalities.SYMBOL, self.targeted_vocab_size)
+    p.modality = {"inputs": modalities.SymbolModality,
+                  "targets": modalities.SymbolModality}
+    p.vocab_size = {"inputs": self._encoders["inputs"].vocab_size,
+                    "targets": self.targeted_vocab_size}
     p.input_space_id = self.input_space_id
     p.target_space_id = self.target_space_id
     p.loss_multiplier = 2.5  # Rough estimate of avg number of tokens per word
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 551d21044..5396d931c 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -27,8 +27,8 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
-from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -241,8 +241,10 @@ def example_reading_spec(self):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {"inputs": (registry.Modalities.IMAGE, 256)}
-    p.target_modality = (registry.Modalities.CLASS_LABEL, self.num_classes)
+    p.modality = {"inputs": modalities.ImageModality,
+                  "targets": modalities.ClassLabelModality}
+    p.vocab_size = {"inputs": 256,
+                    "targets": self.num_classes}
     p.batch_size_multiplier = 4 if self.is_small else 256
     p.loss_multiplier = 3.0 if self.is_small else 1.0
     if self._was_reversed:
@@ -353,9 +355,10 @@ def feature_encoders(self, data_dir):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {"inputs": (registry.Modalities.IMAGE, 256)}
-    encoder = self._encoders["targets"]
-    p.target_modality = (registry.Modalities.SYMBOL, encoder.vocab_size)
+    p.modality = {"inputs": modalities.ImageModality,
+                  "targets": modalities.SymbolModality}
+    p.vocab_size = {"inputs": 256,
+                    "targets": self._encoders["targets"].vocab_size}
     p.batch_size_multiplier = 256
     p.loss_multiplier = 1.0
     p.input_space_id = problem.SpaceID.IMAGE
diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 5c1dfdd31..dfd1984b6 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -23,6 +23,7 @@
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import image_utils
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -357,8 +358,10 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {"inputs": ("image:identity", 256)}
-    p.target_modality = ("image:identity", 256)
+    p.modality = {"inputs": modalities.IdentityModality,
+                  "targets": modalities.IdentityModality}
+    p.vocab_size = {"inputs": 256,
+                    "targets": 256}
     p.batch_size_multiplier = 256
     p.input_space_id = 1
     p.target_space_id = 1
diff --git a/tensor2tensor/data_generators/lambada.py b/tensor2tensor/data_generators/lambada.py
index de0b03833..8f3600ac1 100644
--- a/tensor2tensor/data_generators/lambada.py
+++ b/tensor2tensor/data_generators/lambada.py
@@ -40,6 +40,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -357,12 +358,10 @@ def hparams(self, defaults, unused_model_hparams):
     """
 
     p = defaults
-    source_vocab_size = self._encoders["inputs"].vocab_size
-    num_classes = self._encoders["targets"].vocab_size
-    p.input_modality = {
-        "inputs": (registry.Modalities.SYMBOL, source_vocab_size)
-    }
-    p.target_modality = (registry.Modalities.CLASS_LABEL, num_classes)
+    p.modality = {"inputs": modalities.SymbolModality,
+                  "targets": modalities.ClassLabelModality}
+    p.vocab_size = {"inputs": self._encoders["inputs"].vocab_size,
+                    "targets": self._encoders["targets"].vocab_size}
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 37f36c364..5ab2051a5 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -83,14 +83,15 @@ def get_hparams(self, model_hparams=None):
       return self._hparams
 
     self._hparams = self.task_list[0].get_hparams(model_hparams)
-    # increase the vocab size in order to account for task ids
+    # Increase the vocab size to account for task ids and modify the modality.
     vocab_size_inc = len(self.task_list)
     vocab_size_inc += self.get_max_num_classes()
     vocab_size = self._hparams.vocabulary["targets"].vocab_size
     tf.logging.info("Old vocabulary size: %d" % vocab_size)
     tf.logging.info("New vocabulary size: %d" % (vocab_size + vocab_size_inc))
-    self._hparams.target_modality = modalities.SymbolModality(
-        model_hparams, vocab_size + vocab_size_inc)
+    self._hparams.vocab_size["targets"] = vocab_size + vocab_size_inc
+    self._hparams.modality["targets"] = modalities.SymbolModality(
+        model_hparams, self._hparams.vocab_size["targets"])
 
     return self._hparams
 
diff --git a/tensor2tensor/data_generators/multi_problem_test.py b/tensor2tensor/data_generators/multi_problem_test.py
index c81483568..66301b9d4 100644
--- a/tensor2tensor/data_generators/multi_problem_test.py
+++ b/tensor2tensor/data_generators/multi_problem_test.py
@@ -40,12 +40,14 @@ def __init__(self, input_vocab_size, target_vocab_size):
 
   def hparams(self, defaults, model_hparams):
     hp = defaults
-    hp.input_modality = {
-        "inputs": modalities.SymbolModality(model_hparams,
-                                            self.input_vocab_size)
+    hp.modality = {
+        "inputs": modalities.SymbolModality,
+        "targets": modalities.SymbolModality,
+    }
+    hp.vocab_size = {
+        "inputs": self.input_vocab_size,
+        "targets": self.target_vocab_size,
     }
-    hp.target_modality = modalities.SymbolModality(model_hparams,
-                                                   self.target_vocab_size)
 
   def feature_encoders(self, data_dir):
     encoders = {
@@ -70,11 +72,12 @@ class MultiProblemTest(tf.test.TestCase):
   def testProblemHparamsModality(self):
     multiproblem = TestMultiProblem()
     p_hparams = multiproblem.get_hparams()
-    self.assertIsInstance(p_hparams.input_modality["inputs"],
+    self.assertIsInstance(p_hparams.modality["inputs"],
+                          modalities.SymbolModality)
+    self.assertEqual(p_hparams.modality["inputs"].top_dimensionality, 2)
+    self.assertIsInstance(p_hparams.modality["targets"],
                           modalities.SymbolModality)
-    self.assertEqual(p_hparams.input_modality["inputs"].top_dimensionality, 2)
-    self.assertIsInstance(p_hparams.target_modality, modalities.SymbolModality)
-    self.assertEqual(p_hparams.target_modality.top_dimensionality, 260)
+    self.assertEqual(p_hparams.modality["targets"].top_dimensionality, 260)
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 48b264323..01baeeed7 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -716,8 +716,8 @@ def feature_info(self):
     assert self._hparams is not None
 
     hp = self.get_hparams()
-    input_mods = hp.input_modality
-    target_mod = hp.target_modality
+    input_mods = hp.modality["inputs"]
+    target_mod = hp.modality["targets"]
     vocabs = hp.vocabulary
     if self.has_inputs:
       in_id = hp.input_space_id
@@ -1075,7 +1075,9 @@ def _copy_problem_hparams(p_hparams):
   """Use input modality, vocab, and space id for target."""
   p = p_hparams
   # Duplicate input modality.
-  p.target_modality = p.input_modality["inputs"]
+  p.modality["targets"] = p.modality["inputs"]
+  # Duplicate input vocab size.
+  p.vocab_size["targets"] = p.vocab_size["inputs"]
   # Duplicate input vocabulary.
   p.vocabulary["targets"] = p.vocabulary["inputs"]
   # Duplicate input space ids.
@@ -1089,13 +1091,31 @@ def _reverse_problem_hparams(p_hparams):
   p = p_hparams
 
   # Swap modalities.
-  input_modality = p.input_modality.get("inputs")
-  target_modality = p.target_modality
-  p.target_modality = input_modality
-  if target_modality is not None:
-    p.input_modality["inputs"] = target_modality
-  else:
-    p.input_modality = {}
+  # TODO(trandustin): Note this assumes target modalities have feature name
+  # 'target', and each intended feature to swap has feature name 'input'.
+  # In the future, remove need for this behavior.
+  reversed_modality = {}
+  for feature_name in six.iterkeys(p.modality):
+    reversed_feature_name = feature_name.replace("target", "input")
+    if "target" in feature_name and reversed_feature_name in p.modality:
+      reversed_modality[feature_name] = p.modality[reversed_feature_name]
+      reversed_modality[reversed_feature_name] = p.modality[feature_name]
+    else:
+      reversed_modality[feature_name] = p.modality[feature_name]
+
+  p.modality = reversed_modality
+
+  # Swap vocab sizes.
+  reversed_vocab_size = {}
+  for feature_name in six.iterkeys(p.vocab_size):
+    reversed_feature_name = feature_name.replace("target", "input")
+    if "target" in feature_name and reversed_feature_name in p.vocab_size:
+      reversed_vocab_size[feature_name] = p.vocab_size[reversed_feature_name]
+      reversed_vocab_size[reversed_feature_name] = p.vocab_size[feature_name]
+    else:
+      reversed_vocab_size[feature_name] = p.vocab_size[feature_name]
+
+  p.vocab_size = reversed_vocab_size
 
   # Swap vocabularies.
   input_vocabulary = p.vocabulary.pop("inputs", None)
@@ -1122,17 +1142,14 @@ def _reverse_problem_hparams(p_hparams):
 
 
 def _create_modalities(problem_hparams, hparams):
-  """Converts string-type modalities to their corresponding Modality.
+  """Creates modalities and overrides any according to model hparams.
 
   Args:
     problem_hparams: tf.contrib.training.HParams for the Problem. It must have
-      input_modality and target_modality as attributes. Modalities are either
-      tuples of type ("modality_type:modality_name", vocab_size), and they will
-      be converted to Modality objects; or they are already Modality objects,
-      and they remain the same.
+      modality which is a dict of strings to Modality classes.
     hparams: tf.contrib.training.HParams for the model. It may have
       input_modalities and target_modality, which will override
-      problem_hparams' modalities.
+      problem_hparams' modality input and target keys.
 
   Returns:
     None
@@ -1146,56 +1163,25 @@ def _create_modalities(problem_hparams, hparams):
         modality_name = ":".join(parts[1:])
         input_modality_overrides[feature_name] = modality_name
 
-  input_modality = {}
-  for feature_name, modality in six.iteritems(problem_hparams.input_modality):
-    if isinstance(modality, (list, tuple)):
-      if feature_name in input_modality_overrides:
-        _warn_changed_modality_type(input_modality_overrides[feature_name],
-                                    modality[0],
-                                    feature_name)
-        modality = (input_modality_overrides[feature_name], modality[1])
-      modality = modalities.create_modality(modality, hparams)
-    input_modality[feature_name] = modality
-  problem_hparams.input_modality = input_modality
-
   target_modality_name = None
   if (hasattr(hparams, "target_modality") and
       hparams.target_modality != "default"):
     target_modality_name = hparams.target_modality
 
-  if isinstance(problem_hparams.target_modality, dict):
-    target_modality = {}
-    for feature_name, modality in six.iteritems(
-        problem_hparams.target_modality):
-      if isinstance(modality, (list, tuple)):
-        # TODO(lukaszkaiser): allow overriding other target modalities.
-        if target_modality_name and feature_name == "targets":
-          _warn_changed_modality_type(target_modality_name,
-                                      modality[0],
-                                      "target_modality/%s" % feature_name)
-          modality = (target_modality_name, modality[1])
-        modality = modalities.create_modality(modality, hparams)
-      target_modality[feature_name] = modality
-    problem_hparams.target_modality = target_modality
-  elif isinstance(problem_hparams.target_modality, (list, tuple)):
-    modality = problem_hparams.target_modality
-    if target_modality_name:
-      _warn_changed_modality_type(target_modality_name,
-                                  modality[0],
-                                  "target")
-      modality = (target_modality_name, modality[1])
-    modality = modalities.create_modality(modality, hparams)
-    problem_hparams.target_modality = modality
-
-
-def _warn_changed_modality_type(new_name, old_name, feature_name):
-  new_type, new_name = modalities.parse_modality_name(new_name)
-  old_type, old_name = modalities.parse_modality_name(old_name)
-  if new_type != old_type:
-    tf.logging.warn(
-        "%s has a designated modality type %s (%s) but has been "
-        "overridden with a modality of type %s (%s).", feature_name, old_type,
-        old_name, new_type, new_name)
+  modality = {}
+  for feature_name, modality_cls in six.iteritems(problem_hparams.modality):
+    vocab_size = problem_hparams.vocab_size[feature_name]
+    if feature_name in input_modality_overrides:
+      modality_obj = modalities.create_modality(
+          (input_modality_overrides[feature_name], vocab_size), hparams)
+    elif target_modality_name and feature_name == "targets":
+      # TODO(lukaszkaiser): allow overriding other target modalities.
+      modality_obj = modalities.create_modality(
+          (target_modality_name, vocab_size), hparams)
+    else:
+      modality_obj = modality_cls(hparams, vocab_size)
+    modality[feature_name] = modality_obj
+  problem_hparams.modality = modality
 
 
 def _default_hparams():
diff --git a/tensor2tensor/data_generators/problem_hparams.py b/tensor2tensor/data_generators/problem_hparams.py
index 45bc312be..98ae21052 100644
--- a/tensor2tensor/data_generators/problem_hparams.py
+++ b/tensor2tensor/data_generators/problem_hparams.py
@@ -23,7 +23,7 @@
 import os
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
-from tensor2tensor.layers import modalities  # pylint: disable=unused-import
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -68,10 +68,10 @@ def feature_encoders(self, _):
 
   def hparams(self, defaults, model_hparams):
     hp = defaults
-    hp.input_modality = {
-        "inputs": (registry.Modalities.AUDIO, None),
-    }
-    hp.target_modality = (registry.Modalities.SYMBOL, 256)
+    hp.modality = {"inputs": modalities.SpeechRecognitionModality,
+                   "targets": modalities.SymbolModality}
+    hp.vocab_size = {"inputs": None,
+                     "targets": 256}
 
 
 @registry.register_problem
@@ -93,11 +93,12 @@ def feature_encoders(self, data_dir):
 
   def hparams(self, defaults, model_hparams):
     hp = defaults
-    hp.input_modality = {
-        "inputs": (registry.Modalities.AUDIO, None),
+    hp.modality = {"inputs": modalities.SpeechRecognitionModality,
+                   "targets": modalities.SymbolModality}
+    hp.vocab_size = {
+        "inputs": None,
+        "targets": self.get_feature_encoders()["targets"].vocab_size,
     }
-    hp.target_modality = (registry.Modalities.SYMBOL,
-                          self.get_feature_encoders()["targets"].vocab_size)
     hp.batch_size_multiplier = 256
     hp.loss_multiplier = 2.0
     hp.input_space_id = 13
@@ -129,12 +130,12 @@ def feature_encoders(self, data_dir):
 
   def hparams(self, defaults, model_hparams):
     hp = defaults
-    hp.input_modality = {
-        "inputs": (registry.Modalities.SYMBOL,
-                   self.get_feature_encoders()["inputs"].vocab_size),
+    hp.modality = {"inputs": modalities.SymbolModality,
+                   "targets": modalities.SymbolModality}
+    hp.vocab_size = {
+        "inputs": self.get_feature_encoders()["inputs"].vocab_size,
+        "targets": self.get_feature_encoders()["targets"].vocab_size,
     }
-    hp.target_modality = (registry.Modalities.SYMBOL,
-                          self.get_feature_encoders()["targets"].vocab_size)
     hp.batch_size_multiplier = 256
     hp.loss_multiplier = 2.0
     hp.input_space_id = 3
@@ -173,12 +174,12 @@ def feature_encoders(self, data_dir):
 
   def hparams(self, defaults, model_hparams):
     hp = defaults
-    hp.input_modality = {
-        "inputs": (registry.Modalities.SYMBOL,
-                   self.get_feature_encoders()["inputs"].vocab_size),
+    hp.modality = {"inputs": modalities.SymbolModality,
+                   "targets": modalities.SymbolModality}
+    hp.vocab_size = {
+        "inputs": self.get_feature_encoders()["inputs"].vocab_size,
+        "targets": self.get_feature_encoders()["targets"].vocab_size,
     }
-    hp.target_modality = (registry.Modalities.SYMBOL,
-                          self.get_feature_encoders()["targets"].vocab_size)
     hp.input_space_id = 3
     hp.target_space_id = 15
 
@@ -193,10 +194,10 @@ def __init__(self, input_vocab_size, target_vocab_size):
 
   def hparams(self, defaults, model_hparams):
     hp = defaults
-    hp.input_modality = {
-        "inputs": (registry.Modalities.SYMBOL, self.input_vocab_size)
-    }
-    hp.target_modality = (registry.Modalities.SYMBOL, self.target_vocab_size)
+    hp.modality = {"inputs": modalities.SymbolModality,
+                   "targets": modalities.SymbolModality}
+    hp.vocab_size = {"inputs": self.input_vocab_size,
+                     "targets": self.target_vocab_size}
 
 
 def test_problem_hparams(input_vocab_size=None,
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index a98cca80c..6066f0fff 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -25,7 +25,6 @@
 from tensor2tensor.data_generators import problem as problem_module
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.layers import modalities
-from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -85,9 +84,10 @@ def testProblemHparamsModality(self):
     problem = problem_hparams.TestProblem(input_vocab_size=2,
                                           target_vocab_size=3)
     p_hparams = problem.get_hparams()
-    self.assertIsInstance(p_hparams.input_modality["inputs"],
+    self.assertIsInstance(p_hparams.modality["inputs"],
+                          modalities.SymbolModality)
+    self.assertIsInstance(p_hparams.modality["targets"],
                           modalities.SymbolModality)
-    self.assertIsInstance(p_hparams.target_modality, modalities.SymbolModality)
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testProblemHparamsModalityObj(self):
@@ -95,15 +95,17 @@ class ModalityObjProblem(problem_module.Problem):
 
       def hparams(self, defaults, model_hparams):
         hp = defaults
-        hp.input_modality = {
-            "inputs": modalities.SymbolModality(model_hparams, 2)}
-        hp.target_modality = modalities.SymbolModality(model_hparams, 3)
+        hp.modality = {"inputs": modalities.SymbolModality,
+                       "targets": modalities.SymbolModality}
+        hp.vocab_size = {"inputs": 2,
+                         "targets": 3}
 
     problem = ModalityObjProblem(False, False)
     p_hparams = problem.get_hparams()
-    self.assertIsInstance(p_hparams.input_modality["inputs"],
+    self.assertIsInstance(p_hparams.modality["inputs"],
+                          modalities.SymbolModality)
+    self.assertIsInstance(p_hparams.modality["targets"],
                           modalities.SymbolModality)
-    self.assertIsInstance(p_hparams.target_modality, modalities.SymbolModality)
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testProblemHparamsInputOnlyModality(self):
@@ -111,14 +113,14 @@ class InputOnlyProblem(problem_module.Problem):
 
       def hparams(self, defaults, model_hparams):
         hp = defaults
-        hp.input_modality = {"inputs": (registry.Modalities.SYMBOL, 2)}
-        hp.target_modality = None
+        hp.modality = {"inputs": modalities.SymbolModality}
+        hp.vocab_size = {"inputs": 2}
 
     problem = InputOnlyProblem(False, False)
     p_hparams = problem.get_hparams()
-    self.assertIsInstance(p_hparams.input_modality["inputs"],
+    self.assertIsInstance(p_hparams.modality["inputs"],
                           modalities.SymbolModality)
-    self.assertIsNone(p_hparams.target_modality)
+    self.assertLen(p_hparams.modality, 1)
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testProblemHparamsTargetOnlyModality(self):
@@ -126,13 +128,14 @@ class TargetOnlyProblem(problem_module.Problem):
 
       def hparams(self, defaults, model_hparams):
         hp = defaults
-        hp.input_modality = {}
-        hp.target_modality = (registry.Modalities.SYMBOL, 3)
+        hp.modality = {"targets": modalities.SymbolModality}
+        hp.vocab_size = {"targets": 3}
 
     problem = TargetOnlyProblem(False, False)
     p_hparams = problem.get_hparams()
-    self.assertEqual(p_hparams.input_modality, {})
-    self.assertIsInstance(p_hparams.target_modality, modalities.SymbolModality)
+    self.assertIsInstance(p_hparams.modality["targets"],
+                          modalities.SymbolModality)
+    self.assertLen(p_hparams.modality, 1)
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index fc662350c..f4dad3022 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -28,7 +28,6 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
-from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -61,10 +60,10 @@ def hparams(self, defaults, model_hparams):
     p.add_hparam("num_zeropad_frames", 250)
 
     p = defaults
-    p.input_modality = {
-        "inputs": modalities.SpeechRecognitionModality(model_hparams, None)
-    }
-    p.target_modality = (registry.Modalities.SYMBOL, 256)
+    p.modality = {"inputs": modalities.SpeechRecognitionModality,
+                  "targets": modalities.SymbolModality}
+    p.vocab_size = {"inputs": None,
+                    "targets": 256}
 
   @property
   def is_character_level(self):
diff --git a/tensor2tensor/data_generators/squad.py b/tensor2tensor/data_generators/squad.py
index a37a8a336..4af5f0d6c 100644
--- a/tensor2tensor/data_generators/squad.py
+++ b/tensor2tensor/data_generators/squad.py
@@ -116,7 +116,8 @@ def hparams(self, defaults, unused_model_hparams):
     (super(SquadConcat, self)
      .hparams(defaults, unused_model_hparams))
     p = defaults
-    del p.input_modality['context']
+    del p.modality['context']
+    del p.vocab_size['context']
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 80249b72a..cc62c706c 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -33,6 +33,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -305,23 +306,24 @@ def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.stop_at_eos = int(True)
 
+    p.modality = {"targets": modalities.SymbolModality}
+    p.vocab_size = {"targets": self._encoders["targets"].vocab_size}
     if self.has_inputs:
-      source_vocab_size = self._encoders["inputs"].vocab_size
-      p.input_modality = {
-          "inputs": (registry.Modalities.SYMBOL, source_vocab_size)
-      }
-    target_vocab_size = self._encoders["targets"].vocab_size
-    p.target_modality = (registry.Modalities.SYMBOL, target_vocab_size)
+      p.modality["inputs"] = modalities.SymbolModality
+      p.vocab_size["inputs"] = self._encoders["inputs"].vocab_size
     if self.vocab_type == VocabType.CHARACTER:
       p.loss_multiplier = 2.0
 
     if self.packed_length:
-      identity = (registry.Modalities.GENERIC, None)
       if self.has_inputs:
-        p.input_modality["inputs_segmentation"] = identity
-        p.input_modality["inputs_position"] = identity
-      p.input_modality["targets_segmentation"] = identity
-      p.input_modality["targets_position"] = identity
+        p.modality["inputs_segmentation"] = modalities.IdentityModality
+        p.modality["inputs_position"] = modalities.IdentityModality
+        p.vocab_size["inputs_segmentation"] = None
+        p.vocab_size["inputs_position"] = None
+      p.modality["targets_segmentation"] = modalities.IdentityModality
+      p.modality["targets_position"] = modalities.IdentityModality
+      p.vocab_size["targets_segmentation"] = None
+      p.vocab_size["targets_position"] = None
 
   def example_reading_spec(self):
     data_fields = {"targets": tf.VarLenFeature(tf.int64)}
@@ -390,9 +392,8 @@ def hparams(self, defaults, unused_model_hparams):
     (super(QuestionAndContext2TextProblem, self)
      .hparams(defaults, unused_model_hparams))
     p = defaults
-    source_vocab_size = self._encoders["context"].vocab_size
-    p.input_modality["context"] = (registry.Modalities.SYMBOL,
-                                   source_vocab_size)
+    p.modality["context"] = modalities.SymbolModality
+    p.vocab_size["context"] = self._encoders["context"].vocab_size
     if self.packed_length:
       raise NotImplementedError("QuestionAndContext2Text does not "
                                 "support packed_length")
@@ -495,11 +496,10 @@ def feature_encoders(self, data_dir):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    source_vocab_size = self._encoders["inputs"].vocab_size
-    p.input_modality = {
-        "inputs": (registry.Modalities.SYMBOL, source_vocab_size)
-    }
-    p.target_modality = (registry.Modalities.CLASS_LABEL, self.num_classes)
+    p.modality = {"inputs": modalities.SymbolModality,
+                  "targets": modalities.ClassLabelModality}
+    p.vocab_size = {"inputs": self._encoders["inputs"].vocab_size,
+                    "targets": self.num_classes}
 
   def example_reading_spec(self):
     data_fields = {
diff --git a/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index 2ee7ba007..5ad49abf1 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -23,6 +23,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import timeseries_data_generator
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -139,8 +140,10 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {"inputs": (registry.Modalities.REAL, self.num_series)}
-    p.target_modality = (registry.Modalities.REAL, self.num_series)
+    p.modality = {"inputs": modalities.RealL2LossModality,
+                  "targets": modalities.RealL2LossModality}
+    p.vocab_size = {"inputs": self.num_series,
+                    "targets": self.num_series}
     p.input_space_id = problem.SpaceID.REAL
     p.target_space_id = problem.SpaceID.REAL
 
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 9c33aa14a..d9607f623 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -24,6 +24,7 @@
 import numpy as np
 
 from tensor2tensor.data_generators import video_utils
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -91,11 +92,13 @@ def extra_reading_spec(self):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {
-        "inputs": ("video", 256),
+    p.modality = {
+        "inputs": modalities.VideoModality,
+        "targets": modalities.VideoModality,
     }
-    p.target_modality = {
-        "targets": ("video", 256),
+    p.vocab_size = {
+        "inputs": 256,
+        "targets": 256,
     }
 
   @staticmethod
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 74832a146..951d7320b 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -28,8 +28,8 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.layers import common_video
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
-from tensor2tensor.utils import registry
 from tensor2tensor.utils import video_metrics
 
 import tensorflow as tf
@@ -604,8 +604,10 @@ def example_reading_spec(self):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {"inputs": (registry.Modalities.IMAGE, 256)}
-    p.target_modality = (registry.Modalities.CLASS_LABEL, self.num_classes)
+    p.modality = {"inputs": modalities.ImageModality,
+                  "targets": modalities.ClassLabelModality}
+    p.vocab_size = {"inputs": 256,
+                    "targets": self.num_classes}
     p.input_space_id = problem.SpaceID.IMAGE
     p.target_space_id = problem.SpaceID.IMAGE_LABEL
 
diff --git a/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
index e00f69c1b..5a0d5fc46 100644
--- a/tensor2tensor/data_generators/vqa.py
+++ b/tensor2tensor/data_generators/vqa.py
@@ -34,6 +34,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import vqa_utils
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -129,12 +130,16 @@ def hparams(self, defaults, unused_model_hparams):
     question_encoder = self._encoders["question"]
     targets_encoder = self._encoders["targets"]
 
-    p.input_modality = {
-        "inputs": (registry.Modalities.IMAGE + ":identity", None),
-        "question": (registry.Modalities.SYMBOL, question_encoder.vocab_size)
+    p.modality = {
+        "inputs": modalities.IdentityModality,
+        "question": modalities.SymbolModality,
+        "targets": modalities.MultiLabelModality,
+    }
+    p.vocab_size = {
+        "inputs": None,
+        "question": question_encoder.vocab_size,
+        "targets": targets_encoder.vocab_size,
     }
-    p.target_modality = (registry.Modalities.CLASS_LABEL + ":multi_label",
-                         targets_encoder.vocab_size)
     p.input_space_id = problem.SpaceID.IMAGE  # multiple input features?
     p.target_space_id = self.target_space_id
 
diff --git a/tensor2tensor/data_generators/wikisum/wikisum.py b/tensor2tensor/data_generators/wikisum/wikisum.py
index f1affcb94..19e1816c1 100644
--- a/tensor2tensor/data_generators/wikisum/wikisum.py
+++ b/tensor2tensor/data_generators/wikisum/wikisum.py
@@ -33,6 +33,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import tokenizer
 from tensor2tensor.data_generators.wikisum import utils as cc_utils
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 import tensorflow as tf
@@ -84,12 +85,14 @@ def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.stop_at_eos = True
 
-    source_vocab_size = self._encoders["inputs"].vocab_size
-    target_vocab_size = self._encoders["targets"].vocab_size
-    p.input_modality = {
-        "inputs": (registry.Modalities.SYMBOL, source_vocab_size)
+    p.vocab_size = {
+        "inputs": self._encoders["inputs"].vocab_size,
+        "targets": self._encoders["targets"].vocab_size,
+    }
+    p.modality = {
+        "inputs": modalities.SymbolModality,
+        "targets": modalities.SymbolModality,
     }
-    p.target_modality = (registry.Modalities.SYMBOL, target_vocab_size)
 
   def eval_metrics(self):
     return super(WikisumBase, self).eval_metrics() + [
diff --git a/tensor2tensor/models/mtf_image_transformer.py b/tensor2tensor/models/mtf_image_transformer.py
index cfd72b28d..eeb1942cf 100644
--- a/tensor2tensor/models/mtf_image_transformer.py
+++ b/tensor2tensor/models/mtf_image_transformer.py
@@ -47,7 +47,7 @@ def inputs_vocab_dim(self):
   @property
   def targets_vocab_dim(self):
     return mtf.Dimension(
-        "vocab", self._problem_hparams.target_modality._vocab_size)  # pylint: disable=protected-access
+        "vocab", self._problem_hparams.modality["targets"].top_dimensionality)
 
   @property
   def outputs_vocab_dim(self):
diff --git a/tensor2tensor/models/mtf_image_transformer_test.py b/tensor2tensor/models/mtf_image_transformer_test.py
index bd4aae02e..27437c677 100644
--- a/tensor2tensor/models/mtf_image_transformer_test.py
+++ b/tensor2tensor/models/mtf_image_transformer_test.py
@@ -46,7 +46,7 @@ def get_model(hparams=None,
   p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE,
                                                    VOCAB_SIZE,
                                                    hparams)
-  p_hparams.input_modality = {}
+  del p_hparams.modality["inputs"]
   hparams.problem_hparams = p_hparams
 
   targets = -1 + np.random.random_integers(
diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 411553d10..d1bf19775 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -270,7 +270,8 @@ def mtf_model_fn(self, features, mesh):
 
   @property
   def _targets_vocab_size(self):
-    targets_vocab_size = self._problem_hparams.target_modality._vocab_size  # pylint: disable=protected-access
+    targets_vocab_size = self._problem_hparams.modality[
+        "targets"].top_dimensionality
     targets_vocab_size += (-targets_vocab_size) % self._hparams.vocab_divisor
     return targets_vocab_size
 
@@ -278,8 +279,8 @@ def _targets_vocab_size(self):
   def _inputs_vocab_size(self):
     if not self.has_input:
       return None
-    inputs_vocab_size = self._problem_hparams.input_modality[   # pylint: disable=protected-access
-        "inputs"]._vocab_size
+    inputs_vocab_size = self._problem_hparams.modality[
+        "inputs"].top_dimensionality
     inputs_vocab_size += (-inputs_vocab_size) % self._hparams.vocab_divisor
     return inputs_vocab_size
 
diff --git a/tensor2tensor/models/mtf_transformer_test.py b/tensor2tensor/models/mtf_transformer_test.py
index e2e685dbb..618daa146 100644
--- a/tensor2tensor/models/mtf_transformer_test.py
+++ b/tensor2tensor/models/mtf_transformer_test.py
@@ -45,7 +45,7 @@ def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
                                                    VOCAB_SIZE,
                                                    hparams)
   if not has_input:
-    p_hparams.input_modality = {}
+    del p_hparams.modality["inputs"]
   hparams.problem_hparams = p_hparams
 
   inputs = -1 + np.random.random_integers(
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index ab0c516c0..0688a0c09 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -158,11 +158,7 @@ def decoder(self, x, encoder_layers):
   def gumbel_sample(self, reconstr_gan):
     hparams = self.hparams
     is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
-    if isinstance(self._problem_hparams.target_modality, dict):
-      vocab_size = self._problem_hparams.target_modality[
-          "targets"].top_dimensionality
-    else:
-      vocab_size = self._problem_hparams.target_modality.top_dimensionality
+    vocab_size = self._problem_hparams.modality["targets"].top_dimensionality
     reconstr_gan = tf.nn.log_softmax(reconstr_gan)
     if is_training and hparams.gumbel_temperature > 0.0:
       gumbel_samples = discretization.gumbel_sample(
@@ -183,11 +179,7 @@ def gumbel_sample(self, reconstr_gan):
   def body(self, features):
     hparams = self.hparams
     is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
-    if isinstance(self._problem_hparams.target_modality, dict):
-      vocab_size = self._problem_hparams.target_modality[
-          "targets"].top_dimensionality
-    else:
-      vocab_size = self._problem_hparams.target_modality.top_dimensionality
+    vocab_size = self._problem_hparams.modality["targets"].top_dimensionality
     encoder_layers = None
     self.is1d = hparams.sample_width == 1
     if (hparams.mode != tf.estimator.ModeKeys.PREDICT
@@ -466,7 +458,7 @@ def body(self, features):
       plain_training_loss = losses.pop("training")
       losses["plain"] = plain_training_loss
     res_shape = common_layers.shape_list(basic_result)
-    vocab_size = self._problem_hparams.target_modality.top_dimensionality
+    vocab_size = self._problem_hparams.modality["targets"].top_dimensionality
     targets = tf.one_hot(features["targets_raw"], vocab_size)
     # Prepare inputs for autoregressive modes.
     if common_layers.shape_list(features["targets"])[1] == 1:
diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index a6fc80129..2891d3897 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -328,10 +328,6 @@ def __init__(self, *args, **kwargs):
     self._hparams.ema_means = ema_means
     self._hparams.ema_count = ema_count
 
-  @property
-  def has_input(self):
-    return self._problem_hparams.input_modality
-
   def body(self, features):
     inputs = features["inputs"] if "inputs" in features else None
     reuse = "cache_raw" in features
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index e10074836..bd52a39f2 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -630,10 +630,6 @@ def __init__(self, *args, **kwargs):
           ema_count=ema_count,
           ema_means=ema_means)
 
-  @property
-  def has_input(self):
-    return self._problem_hparams.input_modality
-
   def body(self, features):
     inputs = features["inputs"] if "inputs" in features else None
     if self._hparams.drop_inputs:
diff --git a/tensor2tensor/models/research/universal_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
index 208848f7e..4ce9c38af 100644
--- a/tensor2tensor/models/research/universal_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -45,7 +45,7 @@ def get_model(self,
                                                      VOCAB_SIZE,
                                                      hparams)
     if not has_input:
-      p_hparams.input_modality = {}
+      del p_hparams.modality["inputs"]
     hparams.problems = [p_hparams]
 
     inputs = -1 + np.random.random_integers(
diff --git a/tensor2tensor/models/research/vqa_attention_test.py b/tensor2tensor/models/research/vqa_attention_test.py
index de2016259..d0490ff41 100644
--- a/tensor2tensor/models/research/vqa_attention_test.py
+++ b/tensor2tensor/models/research/vqa_attention_test.py
@@ -47,10 +47,10 @@ def testVqaAttentionBaseline(self):
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
                                                      hparams)
-    p_hparams.input_modality["inputs"] = modalities.ImageModality(hparams)
-    p_hparams.input_modality["question"] = modalities.SymbolModality(
+    p_hparams.modality["inputs"] = modalities.ImageModality(hparams)
+    p_hparams.modality["question"] = modalities.SymbolModality(
         hparams, vocab_size)
-    p_hparams.target_modality = modalities.MultiLabelModality(
+    p_hparams.modality["targets"] = modalities.MultiLabelModality(
         hparams, num_classes + 1)
     with self.test_session() as session:
       features = {
diff --git a/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index 62f31b156..8dfabefed 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -568,7 +568,7 @@ def body(self, features):
       return out
 
     out = tf.reduce_mean(out, [1, 2])
-    num_classes = self._problem_hparams.target_modality.top_dimensionality
+    num_classes = self._problem_hparams.modality["targets"].top_dimensionality
     logits = tf.layers.dense(out, num_classes, name="logits")
 
     losses = {"training": 0.0}
diff --git a/tensor2tensor/models/resnet_test.py b/tensor2tensor/models/resnet_test.py
index fe7538c81..c9a54c805 100644
--- a/tensor2tensor/models/resnet_test.py
+++ b/tensor2tensor/models/resnet_test.py
@@ -48,8 +48,8 @@ def _test_resnet(self, img_size, output_size):
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
                                                      hparams)
-    p_hparams.input_modality["inputs"] = modalities.ImageModality(hparams)
-    p_hparams.target_modality = modalities.ClassLabelModality(
+    p_hparams.modality["inputs"] = modalities.ImageModality(hparams)
+    p_hparams.modality["targets"] = modalities.ClassLabelModality(
         hparams, vocab_size)
     with self.test_session() as session:
       features = {
diff --git a/tensor2tensor/models/slicenet.py b/tensor2tensor/models/slicenet.py
index 9806d8891..bbaaeddaa 100644
--- a/tensor2tensor/models/slicenet.py
+++ b/tensor2tensor/models/slicenet.py
@@ -264,8 +264,7 @@ def slicenet_internal(inputs, targets, target_space, hparams, run_decoder=True):
 class SliceNet(t2t_model.T2TModel):
 
   def body(self, features):
-    target_modality_name = (
-        self._problem_hparams.target_modality.name)
+    target_modality_name = self._problem_hparams.modality["targets"].name
     # If we're just predicting a class, there is no use for a decoder.
     run_decoder = "class_label_modality" not in target_modality_name
     return slicenet_internal(
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 0832babeb..d656ff248 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -323,7 +323,7 @@ def _fast_decode_tpu(self,
           " of the dataset when decoding.")
     dp = self._data_parallelism
     hparams = self._hparams
-    target_modality = self._problem_hparams.target_modality
+    target_modality = self._problem_hparams.modality["targets"]
 
     if self.has_input:
       inputs = features["inputs"]
@@ -343,7 +343,7 @@ def _fast_decode_tpu(self,
       inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
       # _shard_features called to ensure that the variable names match
       inputs = self._shard_features({"inputs": inputs})["inputs"]
-      input_modality = self._problem_hparams.input_modality["inputs"]
+      input_modality = self._problem_hparams.modality["inputs"]
       with tf.variable_scope(input_modality.name):
         inputs = input_modality.bottom_sharded(inputs, dp)
       with tf.variable_scope("body"):
@@ -536,7 +536,7 @@ def _fast_decode(self,
       raise NotImplementedError("Fast decoding only supports a single shard.")
     dp = self._data_parallelism
     hparams = self._hparams
-    target_modality = self._problem_hparams.target_modality
+    target_modality = self._problem_hparams.modality["targets"]
     if "targets_segmentation" in features:
       raise NotImplementedError(
           "Decoding not supported on packed datasets "
@@ -560,7 +560,7 @@ def _fast_decode(self,
       inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
       # _shard_features called to ensure that the variable names match
       inputs = self._shard_features({"inputs": inputs})["inputs"]
-      input_modality = self._problem_hparams.input_modality["inputs"]
+      input_modality = self._problem_hparams.modality["inputs"]
       with tf.variable_scope(input_modality.name):
         inputs = input_modality.bottom_sharded(inputs, dp)
       with tf.variable_scope("body"):
diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 511d0c3a4..4924a8748 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -45,7 +45,7 @@ def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
                                                    VOCAB_SIZE,
                                                    hparams)
   if not has_input:
-    p_hparams.input_modality = {}
+    del p_hparams.modality["inputs"]
   hparams.problem_hparams = p_hparams
 
   inputs = -1 + np.random.random_integers(
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 8ca6d1eb3..8c309384e 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -146,7 +146,7 @@ def is_recurrent_model(self):
   @property
   def _target_modality(self):
     # TODO(mbz): get rid of this somehow.
-    modality = self.hparams.problem_hparams.target_modality["targets"]
+    modality = self.hparams.problem_hparams.modality["targets"]
     return modality.__class__.__name__
 
   @property
@@ -297,7 +297,7 @@ def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
       targets = extra_raw_gts
       targets_shape = common_layers.shape_list(targets)
       targets = tf.reshape(targets, [-1] + targets_shape[2:])
-      mod = self.hparams.problem_hparams.target_modality["targets"]
+      mod = self.hparams.problem_hparams.modality["targets"]
       numerator, denominator = common_layers.padded_cross_entropy(
           logits,
           targets,
@@ -407,11 +407,11 @@ def logits_to_samples(logits):
                        hparams.video_num_target_frames, 1, 1, num_channels]
 
     features["targets"] = tf.zeros(targets_shape, dtype=tf.int32)
-    reward_in_mod = "target_reward" in hparams.problem_hparams.target_modality
-    action_in_mod = "target_action" in hparams.problem_hparams.target_modality
+    reward_in_mod = "target_reward" in hparams.problem_hparams.modality
+    action_in_mod = "target_action" in hparams.problem_hparams.modality
     if reward_in_mod:
       # TODO(lukaszkaiser): this is a hack. get the actual reward history.
-      if "input_reward" not in hparams.problem_hparams.target_modality:
+      if "input_reward" not in features:
         features["input_reward"] = tf.zeros(
             [inputs_shape[0], inputs_shape[1], 1], dtype=tf.int32)
       features["target_reward"] = tf.zeros(
diff --git a/tensor2tensor/models/video/tests_utils.py b/tensor2tensor/models/video/tests_utils.py
index 0c02418af..7bda4b093 100644
--- a/tensor2tensor/models/video/tests_utils.py
+++ b/tensor2tensor/models/video/tests_utils.py
@@ -41,11 +41,9 @@ def fill_hparams(hparams, in_frames, out_frames):
 
 
 def action_modalities(hparams):
-  hparams.problem_hparams.input_modality = {
+  hparams.problem_hparams.modality = {
       "inputs": modalities.VideoModalityL2Raw(hparams, 256),
       "input_action": modalities.SymbolModality(hparams, 5),
-  }
-  hparams.problem_hparams.target_modality = {
       "targets": modalities.VideoModalityL2Raw(hparams, 256),
       "target_action": modalities.SymbolModality(hparams, 5),
   }
@@ -54,12 +52,10 @@ def action_modalities(hparams):
 
 def full_modalities(hparams):
   """Full modalities with actions and rewards."""
-  hparams.problem_hparams.input_modality = {
+  hparams.problem_hparams.modality = {
       "inputs": modalities.VideoModalityL2Raw(hparams, 256),
       "input_reward": modalities.SymbolModality(hparams, 3),
       "input_action": modalities.SymbolModality(hparams, 5),
-  }
-  hparams.problem_hparams.target_modality = {
       "targets": modalities.VideoModalityL2Raw(hparams, 256),
       "target_reward": modalities.SymbolModality(hparams, 3),
       "target_action": modalities.SymbolModality(hparams, 5),
diff --git a/tensor2tensor/models/xception_test.py b/tensor2tensor/models/xception_test.py
index 95a07ef38..25dc75898 100644
--- a/tensor2tensor/models/xception_test.py
+++ b/tensor2tensor/models/xception_test.py
@@ -41,8 +41,8 @@ def _test_xception(self, img_size):
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
                                                      hparams)
-    p_hparams.input_modality["inputs"] = modalities.ImageModality(hparams)
-    p_hparams.target_modality = modalities.ClassLabelModality(
+    p_hparams.modality["inputs"] = modalities.ImageModality(hparams)
+    p_hparams.modality["targets"] = modalities.ClassLabelModality(
         hparams, vocab_size)
     with self.test_session() as session:
       features = {
diff --git a/tensor2tensor/utils/data_reader_test.py b/tensor2tensor/utils/data_reader_test.py
index 53621d3fd..a82071d78 100644
--- a/tensor2tensor/utils/data_reader_test.py
+++ b/tensor2tensor/utils/data_reader_test.py
@@ -26,6 +26,7 @@
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem as problem_mod
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import registry
 
@@ -50,10 +51,10 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
   def hparams(self, defaults, model_hparams):
     hp = defaults
-    hp.input_modality = {
-        "inputs": (registry.Modalities.SYMBOL, 30)
-    }
-    hp.target_modality = (registry.Modalities.SYMBOL, 30)
+    hp.modality = {"inputs": modalities.SymbolModality,
+                   "targets": modalities.SymbolModality}
+    hp.vocab_size = {"inputs": 30,
+                     "targets": 30}
 
   def example_reading_spec(self):
     data_fields = {
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index c5619ff5a..444b07913 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -613,7 +613,7 @@ def _interactive_input_fn(hparams, decode_hp):
   decode_length = decode_hp.extra_length
   input_type = "text"
   p_hparams = hparams.problem_hparams
-  has_input = "inputs" in p_hparams.input_modality
+  has_input = "inputs" in p_hparams.modality
   vocabulary = p_hparams.vocabulary["inputs" if has_input else "targets"]
   # This should be longer than the longest input.
   const_array_size = 10000
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 7840136de..d3fe98b64 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -584,7 +584,7 @@ def weights_fn_for_mp(problem_task_id):
                                     metrics,
                                     list(METRICS_FNS.keys())))
 
-    tm = problem_instance.get_hparams(model_hparams).target_modality
+    tm = problem_instance.get_hparams(model_hparams).modality["targets"]
     if not isinstance(tm, dict):
       tm = {"targets": tm}
 
@@ -609,7 +609,7 @@ def weights_fn_for_mp(problem_task_id):
 def create_eager_metrics_for_problem(problem, model_hparams):
   """See create_eager_metrics."""
   metric_names = problem.eval_metrics()
-  tm = problem.get_hparams(model_hparams).target_modality
+  tm = problem.get_hparams(model_hparams).modality["targets"]
   return create_eager_metrics(metric_names, weights_fn=tm.targets_weights_fn)
 
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 390a6493c..ff57aef47 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -120,8 +120,8 @@ def __init__(self,
     hparams = copy.copy(hparams)
     if self._problem_hparams and hparams.shared_embedding_and_softmax_weights:
       # If vocabularies differ, unset shared_embedding_and_softmax_weights.
-      input_modality = self._problem_hparams.input_modality.get("inputs")
-      target_modality = self._problem_hparams.target_modality
+      input_modality = self._problem_hparams.modality.get("inputs")
+      target_modality = self._problem_hparams.modality.get("targets")
       if (isinstance(input_modality, modality.Modality) and
           isinstance(target_modality, modality.Modality) and
           input_modality.top_dimensionality !=
@@ -181,7 +181,7 @@ def is_predicting(self):
   @property
   def has_input(self):
     if self._problem_hparams:
-      return "inputs" in self._problem_hparams.input_modality
+      return "inputs" in self._problem_hparams.modality
     else:
       return True
 
@@ -202,7 +202,7 @@ def _custom_getter(self):
   @property
   def _target_modality_is_real(self):
     """Whether the target modality is real-valued."""
-    target_modality = self._problem_hparams.target_modality
+    target_modality = self._problem_hparams.modality["targets"]
     return target_modality.name.startswith("real_")
 
   def call(self, inputs, **kwargs):
@@ -342,43 +342,43 @@ def bottom(self, features):
     if not self._problem_hparams:
       log_warn("Without a Problem, T2TModel.bottom is a passthrough.")
       return features
+
     transformed_features = collections.OrderedDict()
     all_previous_modalities = []
+    target_modality = _create_target_modality(self._problem_hparams.modality)
 
-    # Transform the input features
-    for key, input_modality in sorted(
-        six.iteritems(self._problem_hparams.input_modality)):
-      if key not in features:
-        tf.logging.warning("Missing feature %s - ignoring." % key)
+    # Transform features via its corresponding modality.
+    for feature_name, modality_obj in sorted(
+        six.iteritems(self._problem_hparams.modality)):
+      if feature_name not in features:
+        tf.logging.warning("Missing feature %s - ignoring." % feature_name)
         continue
-      do_reuse = input_modality.name in all_previous_modalities
-      with tf.variable_scope(input_modality.name, reuse=do_reuse) as im_vs:
-        self._add_variable_scope(input_modality.name, im_vs)
-        log_info("Transforming feature '%s' with %s.bottom", key,
-                 input_modality.name)
-        transformed_features[key] = input_modality.bottom(features[key])
-      all_previous_modalities.append(input_modality.name)
-
-    # Transform the targets (for autoregressive models)
-    target_modality = self._problem_hparams.target_modality
-    if isinstance(target_modality, dict):
-      for k, v in six.iteritems(target_modality):
-        if k in features:
-          # TODO(aidangomez): share variables?
-          with tf.variable_scope("%s/%s" % (v.name, k)) as tm_vs:
-            self._add_variable_scope("%s/%s" % (v.name, k), tm_vs)
-            log_info("Transforming '%s' with %s.targets_bottom", k, v.name)
-            transformed_features[k] = v.targets_bottom(features[k])
+      # Use if-else clauses to preserve behavior of previous changes: namely,
+      # the variable scope name for the targets feature if there is only one
+      # target modality; and to reuse variable scopes for only input modalities.
+      if feature_name in target_modality:
+        if len(target_modality) > 1:
+          variable_scope_name = "%s/%s" % (modality_obj.name, feature_name)
         else:
-          tf.logging.warn("Modality not found in features: %s", k)
-    else:
-      with tf.variable_scope(target_modality.name) as tm_vs:
-        self._add_variable_scope(target_modality.name, tm_vs)
-        if "targets" in features:
-          log_info("Transforming 'targets' with %s.targets_bottom",
-                   target_modality.name)
-          transformed_features["targets"] = target_modality.targets_bottom(
-              features["targets"])
+          variable_scope_name = modality_obj.name
+        # TODO(aidangomez): share variables?
+        with tf.variable_scope(variable_scope_name) as vs:
+          self._add_variable_scope(variable_scope_name, vs)
+          log_info("Transforming feature '%s' with %s.targets_bottom",
+                   feature_name,
+                   modality_obj.name)
+          transformed_features[feature_name] = modality_obj.targets_bottom(
+              features[feature_name])
+      else:
+        do_reuse = modality_obj.name in all_previous_modalities
+        with tf.variable_scope(modality_obj.name, reuse=do_reuse) as vs:
+          self._add_variable_scope(modality_obj.name, vs)
+          log_info("Transforming feature '%s' with %s.bottom",
+                   feature_name,
+                   modality_obj.name)
+          transformed_features[feature_name] = modality_obj.bottom(
+              features[feature_name])
+        all_previous_modalities.append(modality_obj.name)
 
     for key in features:
       if key not in transformed_features:
@@ -463,13 +463,14 @@ def top(self, body_output, features):
     """
     if isinstance(body_output, dict):
       if self._problem_hparams:
-        target_modality = self._problem_hparams.target_modality
+        target_modality = _create_target_modality(
+            self._problem_hparams.modality)
       else:
         target_modality = {k: None for k in body_output.keys()}
       for k in body_output.keys():
         assert k in target_modality.keys(), (
             "The key %s of model_body's returned logits dict must be in "
-            "problem_hparams.target_modality's dict." % k)
+            "problem_hparams.modality's dict." % k)
       logits = {}
       for k, v in six.iteritems(body_output):
         # TODO(aidangomez): share variables here?
@@ -479,13 +480,14 @@ def top(self, body_output, features):
       return logits
     else:
       if self._problem_hparams:
-        target_modality = self._problem_hparams.target_modality
+        target_modality = _create_target_modality(
+            self._problem_hparams.modality)
       else:
         target_modality = None
       if isinstance(target_modality, dict):
         assert "targets" in target_modality, (
             "model_body returned single logits so 'targets' must be a key "
-            "since problem_hparams.target_modality is a dict.")
+            "since problem_hparams.modality is a dict.")
         target_modality = target_modality["targets"]
       return self._top_single(body_output, target_modality, features)
 
@@ -519,13 +521,14 @@ def _loss_single(self, logits, target_modality, feature):
   def loss(self, logits, features):
     if isinstance(logits, dict):
       if self._problem_hparams:
-        target_modality = self._problem_hparams.target_modality
+        target_modality = _create_target_modality(
+            self._problem_hparams.modality)
       else:
         target_modality = {k: None for k in logits.keys()}
       for k in logits.keys():
         assert k in target_modality.keys(), (
             "The key %s of model_body's returned logits dict must be in "
-            "problem_hparams.target_modality's dict." % k)
+            "problem_hparams.modality's dict." % k)
       losses = {}
       for k, v in six.iteritems(logits):
         losses[k] = self._loss_single(v, target_modality[k], features[k])
@@ -543,13 +546,14 @@ def loss(self, logits, features):
       return tf.add_n([n / d for n, d in losses.values()])
     else:
       if self._problem_hparams:
-        target_modality = self._problem_hparams.target_modality
+        target_modality = _create_target_modality(
+            self._problem_hparams.modality)
       else:
         target_modality = None
       if isinstance(target_modality, dict):
         assert "targets" in target_modality, (
             "model_body returned single logits so 'targets' must be a key "
-            "since problem_hparams.target_modality is a dict.")
+            "since problem_hparams.modality is a dict.")
         target_modality = target_modality["targets"]
       return self._loss_single(logits, target_modality, features["targets"])
 
@@ -578,14 +582,9 @@ def set_mode(self, mode):
 
     if self._problem_hparams:
       # Set model hparams in problem_hparams' modalities, which also store them.
-      for im in six.itervalues(self._problem_hparams.input_modality):
-        im._model_hparams = self._hparams  # pylint: disable=protected-access
-
-      if isinstance(self._problem_hparams.target_modality, dict):
-        for tm in six.itervalues(self._problem_hparams.target_modality):
-          tm._model_hparams = self._hparams  # pylint: disable=protected-access
-      elif self._problem_hparams.target_modality is not None:
-        self._problem_hparams.target_modality._model_hparams = self._hparams  # pylint: disable=protected-access
+      for modality_obj in six.itervalues(self._problem_hparams.modality):
+        if modality_obj is not None:
+          modality_obj._model_hparams = self._hparams  # pylint: disable=protected-access
 
   def prepare_features_for_infer(self, features):
     """Called before inference to allow adding infer-specific features."""
@@ -660,7 +659,7 @@ def infer(self,
       self._fill_problem_hparams_features(features)
 
       if self._problem_hparams:
-        target_modality = self._problem_hparams.target_modality
+        target_modality = self._problem_hparams.modality["targets"]
         if target_modality.is_class_modality:
           beam_size = 1  # No use to run beam-search for a single class.
       if beam_size == 1:
@@ -744,7 +743,7 @@ def symbols_to_logits_fn(ids):
       # it has shape [batch_size] and contains floats between 0 and
       # source_length.
       if self._problem_hparams:
-        if self._problem_hparams.target_modality.top_is_pointwise:
+        if self._problem_hparams.modality["targets"].top_is_pointwise:
           return tf.squeeze(logits, axis=[1, 2, 3])
       # -1 due to the pad above.
       current_output_position = common_layers.shape_list(ids)[1] - 1
@@ -764,7 +763,7 @@ def symbols_to_logits_fn(ids):
       features["inputs"] = tf.reshape(features["inputs"],
                                       [s[0] * s[1], s[2], s[3], s[4]])
 
-    target_modality = self._problem_hparams.target_modality
+    target_modality = self._problem_hparams.modality["targets"]
     vocab_size = target_modality.top_dimensionality
     # Setting decode length to input length + decode_length
     decode_length = tf.constant(decode_length)
@@ -855,7 +854,7 @@ def _slow_greedy_infer_tpu(self, features, decode_length):
     # in metric functions stays in the same frame as other vars.
     targets_old = features.get("targets", None)
 
-    target_modality = self._problem_hparams.target_modality
+    target_modality = self._problem_hparams.modality["targets"]
 
     def infer_step(i, recent_output, recent_logits, unused_loss):
       """Inference step."""
@@ -902,7 +901,7 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
     # input shape, so we confuse it about the input shape.
     initial_output = tf.slice(initial_output, [0, 0, 0, 0],
                               common_layers.shape_list(initial_output))
-    target_modality = self._problem_hparams.target_modality
+    target_modality = self._problem_hparams.modality["targets"]
     if target_modality.is_class_modality:
       decode_length = 1
     else:
@@ -1022,13 +1021,13 @@ def _slow_greedy_infer(self, features, decode_length):
     # in metric functions stays in the same frame as other vars.
     targets_old = features.get("targets", None)
 
-    target_modality = self._problem_hparams.target_modality
+    target_modality = self._problem_hparams.modality["targets"]
 
     def infer_step(recent_output, recent_logits, unused_loss):
       """Inference step."""
       if not tf.contrib.eager.in_eager_mode():
         if self._target_modality_is_real:
-          dim = self._problem_hparams.target_modality.top_dimensionality
+          dim = self._problem_hparams.modality["targets"].top_dimensionality
           recent_output.set_shape([None, None, None, dim])
         else:
           recent_output.set_shape([None, None, None, 1])
@@ -1068,7 +1067,7 @@ def infer_step(recent_output, recent_logits, unused_loss):
     else:
       batch_size = common_layers.shape_list(features["inputs"])[0]
       if self._target_modality_is_real:
-        dim = self._problem_hparams.target_modality.top_dimensionality
+        dim = self._problem_hparams.modality["targets"].top_dimensionality
         initial_output = tf.zeros((batch_size, 0, 1, dim), dtype=tf.float32)
       else:
         initial_output = tf.zeros((batch_size, 0, 1, 1), dtype=tf.int64)
@@ -1076,7 +1075,7 @@ def infer_step(recent_output, recent_logits, unused_loss):
     # input shape, so we confuse it about the input shape.
     initial_output = tf.slice(initial_output, [0, 0, 0, 0],
                               common_layers.shape_list(initial_output))
-    target_modality = self._problem_hparams.target_modality
+    target_modality = self._problem_hparams.modality["targets"]
     if target_modality.is_class_modality:
       decode_length = 1
     else:
@@ -1581,7 +1580,7 @@ def create_tpu_eval_metrics_fn(problem, model_hparams):
   metric_fns = []
   eval_metrics = problem.eval_metrics()
 
-  tm = problem.get_hparams(model_hparams).target_modality
+  tm = _create_target_modality(problem.get_hparams(model_hparams).modality)
   if isinstance(tm, dict):
     for k, v in six.iteritems(tm):
       weights_fn = v.targets_weights_fn
@@ -1748,7 +1747,7 @@ def create_eager_var_store():
 def scheduled_sampling(hparams, problem_hparams, dp, sharded_logits, losses,
                        sharded_features, transformed_features, model):
   """Scheduled sampling."""
-  target_modality = problem_hparams.target_modality
+  target_modality = problem_hparams.modality["targets"]
 
   def sample(x):
     """Multinomial sampling from a n-dimensional tensor."""
@@ -1901,3 +1900,11 @@ def set_custom_getter_compose(custom_getter):
   tf.get_variable_scope().set_custom_getter(
       _compose_custom_getters(tf.get_variable_scope().custom_getter,
                               custom_getter))
+
+
+def _create_target_modality(modality_dict):
+  # TODO(trandustin): We require this in order to apply methods utilized
+  # differently for modalities which are "targets"
+  # (e.g., modality.target_bottom). In the future, remove need for this
+  # behavior.
+  return {k: v for k, v in six.iteritems(modality_dict) if "target" in k}
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index a3ae4cd85..61d6f53d7 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -105,10 +105,10 @@ def testMultipleTargetModalities(self):
         problem_name="tiny_algo")
     # Manually turn off sharing. It is not currently supported for multitargets.
     hparams.shared_embedding_and_softmax_weights = 0  # pylint: disable=line-too-long
-    hparams.problem_hparams.target_modality = {
-        "targets": hparams.problem_hparams.target_modality,
-        "A": hparams.problem_hparams.target_modality,
-        "B": hparams.problem_hparams.target_modality,
+    hparams.problem_hparams.modality = {
+        "targets": hparams.problem_hparams.modality["targets"],
+        "targets_A": hparams.problem_hparams.modality["targets"],
+        "targets_B": hparams.problem_hparams.modality["targets"],
     }
     hparams.problem._hparams = hparams.problem_hparams
 
@@ -119,14 +119,14 @@ def testMultipleTargetModalities(self):
     dataset = dataset.repeat(None).padded_batch(10, dataset.output_shapes)
     features = dataset.make_one_shot_iterator().get_next()
     features = problem_lib.standardize_shapes(features)
-    features["A"] = features["B"] = features["targets"]
+    features["targets_A"] = features["targets_B"] = features["targets"]
 
     # Model
     model = registry.model("transformer")(hparams, tf.estimator.ModeKeys.TRAIN)
 
     def body(args, mb=model.body):
       out = mb(args)
-      return {"targets": out, "A": out, "B": out}
+      return {"targets": out, "targets_A": out, "targets_B": out}
 
     model.body = body
 

From 479425d190117cd39110a3938707a1754d637cb1 Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Fri, 26 Oct 2018 16:30:50 -0700
Subject: [PATCH 1091/2720] Allow create_hparams() in t2t_trainer.py to load
 from hparams.json (if any)

PiperOrigin-RevId: 218931883
---
 tensor2tensor/bin/t2t_trainer.py        |  5 ++++-
 tensor2tensor/utils/trainer_lib.py      | 28 +++++++++++++++++++++++-
 tensor2tensor/utils/trainer_lib_test.py | 29 +++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 61b057bbf..c1e8c9506 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -155,11 +155,14 @@ def set_hparams_from_args(args):
 
 
 def create_hparams():
+  """Create hparams."""
   if FLAGS.use_tpu and "tpu" not in FLAGS.hparams_set:
     tf.logging.warn("Not all hyperparameter sets work on TPU. "
                     "Prefer hparams_sets with a '_tpu' suffix, "
                     "e.g. transformer_tpu, if available for your model.")
-  return trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
+  hparams_path = os.path.join(FLAGS.output_dir, "hparams.json")
+  return trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams,
+                                    hparams_path=hparams_path)
 
 
 def create_experiment_fn():
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 015d991a1..4f2e925e9 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -88,9 +88,12 @@ def create_session_config(log_device_placement=False,
 def create_hparams(hparams_set,
                    hparams_overrides_str="",
                    data_dir=None,
-                   problem_name=None):
+                   problem_name=None,
+                   hparams_path=None):
   """Create HParams with data_dir and problem hparams, if kwargs provided."""
   hparams = registry.hparams(hparams_set)
+  if hparams_path and tf.gfile.Exists(hparams_path):
+    hparams = _create_hparams_from_json(hparams_path, hparams)
   if data_dir:
     hparams.add_hparam("data_dir", data_dir)
   if problem_name:
@@ -102,6 +105,29 @@ def create_hparams(hparams_set,
   return hparams
 
 
+def _create_hparams_from_json(json_path, hparams=None):
+  """Loading hparams from json; can also start from hparams if specified."""
+  tf.logging.info("Loading hparams from existing json %s" % json_path)
+  with tf.gfile.Open(json_path, "r") as f:
+    hparams_values = json.load(f)
+    new_hparams = tf.contrib.training.HParams(**hparams_values)
+    # Some keys are in new_hparams but not hparams, so we need to be more
+    #   careful than simply using parse_json() from HParams
+    if hparams:  # hparams specified, so update values from json
+      for key in sorted(new_hparams.values().keys()):
+        if hasattr(hparams, key):  # Overlapped keys
+          value = getattr(hparams, key)
+          new_value = getattr(new_hparams, key)
+          if value != new_value:  # Different values
+            tf.logging.info("Overwrite key %s: %s -> %s" % (
+                key, value, new_value))
+            setattr(hparams, key, new_value)
+    else:
+      hparams = new_hparams
+
+  return hparams
+
+
 def is_cloud_async_distributed():
   return ("chief" in
           json.loads(os.environ.get("TF_CONFIG", "{}")).get("cluster", {}))
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index 61d6f53d7..b91f43dd7 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import os
 from tensor2tensor.data_generators import algorithmic
 from tensor2tensor.data_generators import problem as problem_lib
 from tensor2tensor.models import transformer  # pylint: disable=unused-import
@@ -139,6 +140,34 @@ def body(args, mb=model.body):
       sess.run(tf.global_variables_initializer())
       sess.run([logits, loss])
 
+  def testCreateHparams(self):
+    # Get json_path
+    pkg, _ = os.path.split(__file__)
+    pkg, _ = os.path.split(pkg)
+    json_path = os.path.join(
+        pkg, "test_data", "transformer_test_ckpt", "hparams.json")
+
+    # Create hparams
+    hparams = trainer_lib.create_hparams("transformer_big", "hidden_size=1",
+                                         hparams_path=json_path)
+    self.assertEqual(2, hparams.num_hidden_layers)  # from json
+    self.assertEqual(1, hparams.hidden_size)  # from hparams_overrides_str
+
+    # Compare with base hparams
+    base_hparams = trainer_lib.create_hparams("transformer_big")
+    self.assertEqual(len(base_hparams.values()), len(hparams.values()))
+
+  def testCreateHparamsFromJson(self):
+    # Get json_path
+    pkg, _ = os.path.split(__file__)
+    pkg, _ = os.path.split(pkg)
+    json_path = os.path.join(
+        pkg, "test_data", "transformer_test_ckpt", "hparams.json")
+
+    # Create hparams
+    hparams = trainer_lib._create_hparams_from_json(json_path)
+    self.assertEqual(75, len(hparams.values()))
+
 
 if __name__ == "__main__":
   tf.test.main()

From 764c9f5d957030a510cf97838e88b7d376aa8a66 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Fri, 26 Oct 2018 17:29:35 -0700
Subject: [PATCH 1092/2720] Make export work for mesh models

PiperOrigin-RevId: 218938875
---
 tensor2tensor/serving/export.py  |  6 ++++--
 tensor2tensor/utils/mtf_model.py | 12 ++++++++----
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 54e153f7c..2c9218dcd 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -122,9 +122,11 @@ def hub_module_fn():
         inputs=original_features,
         outputs=spec.export_outputs["serving_default"].outputs)
 
-  # TFHub doesn't support LOSSES collections.
+  # TFHub doesn't support the following collections.
+  drop_collections = [tf.GraphKeys.LOSSES,
+                      tf.GraphKeys.SUMMARIES, tf.GraphKeys.LOCAL_VARIABLES]
   module_spec = hub.create_module_spec(
-      hub_module_fn, drop_collections=[tf.GraphKeys.LOSSES])
+      hub_module_fn, drop_collections=drop_collections)
   # Loads the weights from the checkpoint using the model above
   # and saves it in the export_path.
   export_module_spec_with_checkpoint(
diff --git a/tensor2tensor/utils/mtf_model.py b/tensor2tensor/utils/mtf_model.py
index 24c6362e3..ff9fb4d76 100644
--- a/tensor2tensor/utils/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -87,7 +87,7 @@ def estimator_model_fn(cls,
           mesh_shape, layout_rules, mesh_devices, ctx.device_assignment)
     else:
       var_placer = None
-      if len(data_parallelism.ps_devices) == 1:
+      if data_parallelism is None or len(data_parallelism.ps_devices) == 1:
         mesh_devices = [""] * mesh_shape.size
       else:
         assert len(data_parallelism.ps_devices) == mesh_shape.size
@@ -221,10 +221,14 @@ def estimator_spec_predict(self, features, mesh, mesh_impl, use_tpu):
       outputs = tf.slice(
           outputs, [0] * ndims, [actual_batch_size] + [-1] * (ndims - 1))
     predictions = {
-        "outputs": outputs,
-        "targets": features.get("infer_targets", features.get("inputs")),
-        "inputs": features.get("inputs"),
+        "outputs": outputs
     }
+    if features.get("infer_targets") is not None:
+      predictions["infer_targets"] = features["infer_targets"]
+
+    if features.get("inputs") is not None:
+      predictions["inputs"] = features["inputs"]
+
     if use_tpu:
       t2t_model.remove_summaries()
       return tpu_estimator.TPUEstimatorSpec(

From 351d897df2a6cbfadb1637760c1b678466aa03c9 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 26 Oct 2018 21:45:57 -0700
Subject: [PATCH 1093/2720] Fix shape check in python

PiperOrigin-RevId: 218954353
---
 tensor2tensor/layers/common_video.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 974cc4b83..646a50ec5 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -340,10 +340,6 @@ def tile_and_concat(image, latent, concat_latent=True):
   latent_shape = common_layers.shape_list(latent)
   height, width = image_shape[1], image_shape[2]
   latent_dims = latent_shape[1]
-
-  if height < latent_dims:
-    raise ValueError("Latent is too big to tile.")
-
   height_multiples = height // latent_dims
   pad = height - (height_multiples * latent_dims)
   latent = tf.reshape(latent, (-1, latent_dims, 1, 1))

From c864603bad91503bd198a7a410603bf2d85d9d6a Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 26 Oct 2018 22:02:24 -0700
Subject: [PATCH 1094/2720] Add Bayesian feedforward layer for BNN training.

This is a prototype that doesn't rely on random variables or distributions. It has limitations which I noted as TODOs. Can use it for now until models we write require more.

PiperOrigin-RevId: 218955117
---
 tensor2tensor/layers/bayes.py      | 215 +++++++++++++++++++++++++++++
 tensor2tensor/layers/bayes_test.py |  65 +++++++++
 2 files changed, 280 insertions(+)
 create mode 100644 tensor2tensor/layers/bayes.py
 create mode 100644 tensor2tensor/layers/bayes_test.py

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
new file mode 100644
index 000000000..fb7444f7d
--- /dev/null
+++ b/tensor2tensor/layers/bayes.py
@@ -0,0 +1,215 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Bayesian layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+from tensorflow.python.framework import tensor_shape
+
+
+class Softplus(tf.keras.constraints.Constraint):
+  """Softplus constraint."""
+
+  def __init__(self, epsilon=tf.keras.backend.epsilon()):
+    self.epsilon = epsilon
+
+  def __call__(self, w):
+    return tf.nn.softplus(w) + self.epsilon
+
+  def get_config(self):
+    return {'epsilon': self.epsilon}
+
+
+def softplus():  # alias, following tf.keras.constraints
+  return Softplus()
+
+
+class TrainableNormal(tf.keras.initializers.Initializer):
+  """Random normal op as an initializer with trainable mean and stddev."""
+
+  def __init__(self,
+               mean_initializer=tf.random_normal_initializer(stddev=0.1),
+               unconstrained_stddev_initializer=tf.random_normal_initializer(
+                   mean=-3., stddev=0.1),
+               mean_regularizer=None,
+               unconstrained_stddev_regularizer=None,
+               mean_constraint=None,
+               unconstrained_stddev_constraint=softplus(),
+               seed=None,
+               dtype=tf.float32):
+    """Constructs initializer."""
+    self.mean_initializer = mean_initializer
+    self.unconstrained_stddev_initializer = unconstrained_stddev_initializer
+    self.mean_regularizer = mean_regularizer
+    self.unconstrained_stddev_regularizer = unconstrained_stddev_regularizer
+    self.mean_constraint = mean_constraint
+    self.unconstrained_stddev_constraint = unconstrained_stddev_constraint
+    self.seed = seed
+    self.dtype = tf.as_dtype(dtype)
+
+  def __call__(self, shape, dtype=None, add_variable_fn=None):
+    if dtype is None:
+      dtype = self.dtype
+    mean = add_variable_fn(
+        'mean',
+        shape=shape,
+        initializer=self.mean_initializer,
+        regularizer=self.mean_regularizer,
+        constraint=self.mean_constraint,
+        dtype=dtype,
+        trainable=True)
+    stddev = add_variable_fn(
+        'unconstrained_stddev',
+        shape=shape,
+        initializer=self.unconstrained_stddev_initializer,
+        regularizer=self.unconstrained_stddev_regularizer,
+        constraint=self.unconstrained_stddev_constraint,
+        dtype=dtype,
+        trainable=True)
+    noise = tf.random_normal(shape, dtype=dtype, seed=self.seed)
+    output = mean + stddev * noise
+    # TODO(trandustin): Hack to store parameters so KL reg. can operate on them.
+    output._parameters = (mean, stddev)  # pylint: disable=protected-access
+    return output
+
+  def get_config(self):
+    return {
+        'mean_initializer':
+            tf.keras.initializers.serialize(self.mean_initializer),
+        'unconstrained_stddev_initializer':
+            tf.keras.initializers.serialize(
+                self.unconstrained_stddev_initializer),
+        'mean_regularizer':
+            tf.keras.regularizers.serialize(self.mean_regularizer),
+        'unconstrained_stddev_regularizer':
+            tf.keras.regularizers.serialize(
+                self.unconstrained_stddev_regularizer),
+        'activity_regularizer':
+            tf.keras.regularizers.serialize(self.activity_regularizer),
+        'mean_constraint':
+            tf.keras.constraints.serialize(self.mean_constraint),
+        'unconstrained_stddev_constraint':
+            tf.keras.constraints.serialize(
+                self.unconstrained_stddev_constraint),
+        'dtype': self.dtype.name,
+    }
+
+
+def trainable_normal():  # alias, following tf.keras.initializers
+  return TrainableNormal()
+
+
+class NormalKLDivergence(tf.keras.regularizers.Regularizer):
+  """KL divergence regularizer from one normal distribution to another."""
+
+  def __init__(self, mean=0., stddev=1.):
+    """Construct regularizer where default is a KL towards the std normal."""
+    self.mean = mean
+    self.stddev = stddev
+
+  def __call__(self, x):
+    mean, stddev = x._parameters  # pylint: disable=protected-access
+    variance2 = tf.square(self.stddev)
+    variance_ratio = tf.square(stddev) / variance2
+    regularization = tf.square(mean - self.mean) / (2. * variance2)
+    regularization += (variance_ratio - 1. - tf.log(variance_ratio)) / 2.
+    return regularization
+
+
+def normal_kl_divergence():  # alias, following tf.keras.regularizers
+  return NormalKLDivergence()
+
+
+class DenseReparameterization(tf.keras.layers.Dense):
+  """Bayesian densely-connected layer estimated via reparameterization.
+
+  The layer computes a variational Bayesian approximation to the distribution
+  over densely-connected layers,
+
+  ```
+  p(outputs | inputs) = int dense(inputs; weights, bias) p(weights, bias)
+    dweights dbias.
+  ```
+
+  It does this with a stochastic forward pass, sampling from learnable
+  distributions on the kernel and bias. Gradients with respect to the
+  distributions' learnable parameters backpropagate via reparameterization.
+  Minimizing cross-entropy plus the layer's losses performs variational
+  minimum description length, i.e., it minimizes an upper bound to the negative
+  marginal likelihood.
+  """
+
+  def __init__(self,
+               units,
+               activation=None,
+               use_bias=True,
+               kernel_initializer=None,
+               bias_initializer='zero',
+               kernel_regularizer=normal_kl_divergence(),
+               bias_regularizer=None,
+               activity_regularizer=None,
+               **kwargs):
+    if not kernel_initializer:
+      kernel_initializer = trainable_normal()
+    if not bias_initializer:
+      bias_initializer = trainable_normal()
+    super(DenseReparameterization, self).__init__(
+        units=units,
+        activation=activation,
+        use_bias=use_bias,
+        kernel_initializer=kernel_initializer,
+        bias_initializer=bias_initializer,
+        kernel_regularizer=kernel_regularizer,
+        bias_regularizer=bias_regularizer,
+        activity_regularizer=activity_regularizer,
+        **kwargs)
+
+  def build(self, input_shape):
+    input_shape = tf.TensorShape(input_shape)
+    if tensor_shape.dimension_value(input_shape[-1]) is None:
+      raise ValueError('The last dimension of the inputs to `Dense` '
+                       'should be defined. Found `None`.')
+    last_dim = tensor_shape.dimension_value(input_shape[-1])
+    self.input_spec = tf.layers.InputSpec(min_ndim=2,
+                                          axes={-1: last_dim})
+    self.kernel = self.kernel_initializer([last_dim, self.units],
+                                          self.dtype,
+                                          self.add_weight)
+    if self.kernel_regularizer is not None:
+      self._handle_weight_regularization('kernel',
+                                         self.kernel,
+                                         self.kernel_regularizer)
+    if self.use_bias:
+      # TODO(trandustin): Because of self.add_weight, the signature differs from
+      # other initializers, preventing interoperability.
+      if isinstance(self.bias_initializer, TrainableNormal):
+        self.bias = self.bias_initializer([self.units],
+                                          self.dtype,
+                                          self.add_weight)
+      else:
+        self.bias = self.bias_initializer([self.units],
+                                          self.dtype)
+      if self.bias_regularizer is not None:
+        self._handle_weight_regularization('bias',
+                                           self.bias,
+                                           self.bias_regularizer)
+    else:
+      self.bias = None
+    self.built = True
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
new file mode 100644
index 000000000..1dffa6530
--- /dev/null
+++ b/tensor2tensor/layers/bayes_test.py
@@ -0,0 +1,65 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 common Bayes."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensor2tensor.layers import bayes
+
+import tensorflow as tf
+
+
+class BayesTest(parameterized.TestCase, tf.test.TestCase):
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testDenseReparameterization(self):
+    inputs = tf.to_float(np.random.rand(5, 3, 12))
+    layer = bayes.DenseReparameterization(4, activation=tf.nn.relu)
+    outputs1 = layer(inputs)
+    outputs2 = layer(inputs)
+    self.evaluate(tf.global_variables_initializer())
+    res1, _ = self.evaluate([outputs1, outputs2])
+    self.assertEqual(res1.shape, (5, 3, 4))
+    self.assertAllGreaterEqual(res1, 0.)
+    # TODO(trandustin): Fix this to work with Eager.
+    # self.assertNotAllClose(res1, res2)
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testDenseReparameterizationModel(self):
+    inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
+    model = tf.keras.Sequential([
+        tf.keras.layers.Conv2D(3,
+                               kernel_size=2,
+                               padding="SAME",
+                               activation=tf.nn.relu),
+        tf.keras.layers.Flatten(),
+        bayes.DenseReparameterization(2, activation=None),
+    ])
+    outputs = model(inputs)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(outputs)
+    self.assertEqual(res.shape, (3, 2))
+    self.assertLen(model.losses, 1)
+
+
+if __name__ == "__main__":
+  tf.test.main()
+

From 1e9125af06e2cb5e3587218e023f723a36d47fdb Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 27 Oct 2018 23:20:02 -0700
Subject: [PATCH 1095/2720] Internal

PiperOrigin-RevId: 219019889
---
 tensor2tensor/bin/t2t_trainer.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index c1e8c9506..792ed0501 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -127,6 +127,7 @@
                      "out")
 
 
+
 def set_hparams_from_args(args):
   """Set hparams overrides from unparsed args list."""
   if not args:

From 1e50fa9f0992e259dfa1f3685a7f37fcebc3f8b6 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 29 Oct 2018 09:55:25 -0700
Subject: [PATCH 1096/2720] More support for transformers exported for TPU

PiperOrigin-RevId: 219144427
---
 tensor2tensor/bin/t2t_attack.py     |  2 +-
 tensor2tensor/bin/t2t_prune.py      |  2 +-
 tensor2tensor/models/transformer.py |  1 +
 tensor2tensor/serving/export.py     |  6 ++++--
 tensor2tensor/utils/t2t_model.py    | 19 +++++++++++++++----
 tensor2tensor/utils/trainer_lib.py  |  2 +-
 6 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 4ccedc779..13b350abc 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -194,7 +194,7 @@ def main(argv):
 
   if FLAGS.surrogate_attack:
     sur_model_fn = t2t_model.T2TModel.make_estimator_model_fn(
-        FLAGS.surrogate_model, sur_hparams)
+        FLAGS.surrogate_model, sur_hparams, use_tpu=FLAGS.use_tpu)
     sur_ch_model = adv_attack_utils.T2TAttackModel(
         sur_model_fn, features, params, sur_config, scope="surrogate")
     # Dummy call to construct graph
diff --git a/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
index c53724a8f..008462e63 100644
--- a/tensor2tensor/bin/t2t_prune.py
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -87,7 +87,7 @@ def main(argv):
   sess = tf.Session()
 
   model_fn = t2t_model.T2TModel.make_estimator_model_fn(
-      FLAGS.model, hparams)
+      FLAGS.model, hparams, use_tpu=FLAGS.use_tpu)
   spec = model_fn(
       features,
       labels,
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index d656ff248..cced9ee3f 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -97,6 +97,7 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
         hparams,
         nonpadding=features_to_nonpadding(features, "inputs"),
         save_weights_to=self.attention_weights,
+        make_image_summary=not common_layers.is_xla_compiled(),
         losses=losses)
 
     return encoder_output, encoder_decoder_attention_bias
diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 2c9218dcd..54125cb70 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -49,7 +49,8 @@ def create_estimator(run_config, hparams):
       FLAGS.model,
       hparams,
       run_config,
-      decode_hparams=decoding.decode_hparams(FLAGS.decode_hparams))
+      decode_hparams=decoding.decode_hparams(FLAGS.decode_hparams),
+      use_tpu=FLAGS.use_tpu)
 
 
 def create_hparams():
@@ -110,7 +111,8 @@ def hub_module_fn():
     model_fn = t2t_model.T2TModel.make_estimator_model_fn(
         model_name,
         hparams,
-        decode_hparams=decode_hparams)
+        decode_hparams=decode_hparams,
+        use_tpu=FLAGS.use_tpu)
     features = problem.serving_input_fn(hparams).features
 
     # we must do a copy of the features, as the model_fn can add additional
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index ff57aef47..04d83f2ed 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1230,7 +1230,8 @@ def get_eval_hooks(model_name, hook_context):
   @staticmethod
   def make_estimator_model_fn(model_name,
                               hparams,
-                              decode_hparams=None):
+                              decode_hparams=None,
+                              use_tpu=False):
     model_cls = registry.model(model_name)
 
     def wrapping_model_fn(features, labels, mode, params=None, config=None):
@@ -1241,7 +1242,8 @@ def wrapping_model_fn(features, labels, mode, params=None, config=None):
           mode,
           config=config,
           params=params,
-          decode_hparams=decode_hparams)
+          decode_hparams=decode_hparams,
+          use_tpu=use_tpu)
 
     return wrapping_model_fn
 
@@ -1253,7 +1255,8 @@ def estimator_model_fn(cls,
                          mode,
                          config=None,
                          params=None,
-                         decode_hparams=None):
+                         decode_hparams=None,
+                         use_tpu=False):
     """Model fn for Estimator.
 
     Args:
@@ -1264,6 +1267,7 @@ def estimator_model_fn(cls,
       config: RunConfig, possibly with data_parallelism attribute
       params: dict, may include batch_size, use_tpu
       decode_hparams: HParams, used when mode == PREDICT.
+      use_tpu: A bool, whether to build the inference graph for TPU.
 
     Returns:
       TPUEstimatorSpec if use tpu else EstimatorSpec
@@ -1272,7 +1276,6 @@ def estimator_model_fn(cls,
       _create_dummy_vars()
     hparams = copy.deepcopy(hparams)
 
-    use_tpu = params and params.get("use_tpu", False)
     # Instantiate model
     data_parallelism = None
     if not use_tpu and config:
@@ -1287,6 +1290,14 @@ def estimator_model_fn(cls,
 
     # PREDICT mode
     if mode == tf.estimator.ModeKeys.PREDICT:
+      if use_tpu:
+        inputs = features["inputs"]
+        shape = inputs.get_shape().as_list()
+        if shape[0] is None:
+          shape[0] = decode_hparams.batch_size or hparams.batch_size
+        if shape[1] is None:
+          shape[1] = hparams.max_input_seq_length or hparams.max_length
+        inputs.set_shape(shape)
       return model.estimator_spec_predict(features, use_tpu=use_tpu)
 
     # TRAIN and EVAL modes
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 4f2e925e9..eed26d0ac 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -279,7 +279,7 @@ def create_estimator(model_name,
                      use_xla=False):
   """Create a T2T Estimator."""
   model_fn = t2t_model.T2TModel.make_estimator_model_fn(
-      model_name, hparams, decode_hparams=decode_hparams)
+      model_name, hparams, decode_hparams=decode_hparams, use_tpu=use_tpu)
 
 
   del use_xla

From 6dedaff0b47756d1dafa1efd73fc07a3154b72db Mon Sep 17 00:00:00 2001
From: Mehrad Moradshahi <mehrad@stanford.edu>
Date: Mon, 29 Oct 2018 10:28:16 -0700
Subject: [PATCH 1097/2720] variable_scope wrapper for avg_checkpoints (#1171)

---
 tensor2tensor/utils/avg_checkpoints.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index 879de16fa..8d4b175fc 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -90,10 +90,11 @@ def main(_):
   for name in var_values:  # Average.
     var_values[name] /= len(checkpoints)
 
-  tf_vars = [
-      tf.get_variable(v, shape=var_values[v].shape, dtype=var_dtypes[name])
-      for v in var_values
-  ]
+  with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+      tf_vars = [
+          tf.get_variable(v, shape=var_values[v].shape, dtype=var_dtypes[v])
+          for v in var_values
+      ]
   placeholders = [tf.placeholder(v.dtype, shape=v.shape) for v in tf_vars]
   assign_ops = [tf.assign(v, p) for (v, p) in zip(tf_vars, placeholders)]
   global_step = tf.Variable(

From de2964a78ce6ebdbdcd0581a5b0e1247b904b405 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 29 Oct 2018 10:28:57 -0700
Subject: [PATCH 1098/2720] Fix has_input to optionally include
 features/modality["inputs"].

This pattern follows other model tests. Code search with the pattern del.*modality\[\"inputs\"\]

PiperOrigin-RevId: 219150753
---
 .../models/research/universal_transformer_test.py        | 3 ++-
 tensor2tensor/utils/avg_checkpoints.py                   | 9 ++++-----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
index 4ce9c38af..184d8945c 100644
--- a/tensor2tensor/models/research/universal_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -53,10 +53,11 @@ def get_model(self,
     targets = -1 + np.random.random_integers(
         VOCAB_SIZE, size=(BATCH_SIZE, TARGET_LENGTH, 1, 1))
     features = {
-        "inputs": tf.constant(inputs, dtype=tf.int32, name="inputs"),
         "targets": tf.constant(targets, dtype=tf.int32, name="targets"),
         "target_space_id": tf.constant(1, dtype=tf.int32)
     }
+    if has_input:
+      features["inputs"] = tf.constant(inputs, dtype=tf.int32, name="inputs")
 
     return universal_transformer.UniversalTransformer(
         hparams, mode, p_hparams), features
diff --git a/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index 8d4b175fc..879de16fa 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -90,11 +90,10 @@ def main(_):
   for name in var_values:  # Average.
     var_values[name] /= len(checkpoints)
 
-  with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      tf_vars = [
-          tf.get_variable(v, shape=var_values[v].shape, dtype=var_dtypes[v])
-          for v in var_values
-      ]
+  tf_vars = [
+      tf.get_variable(v, shape=var_values[v].shape, dtype=var_dtypes[name])
+      for v in var_values
+  ]
   placeholders = [tf.placeholder(v.dtype, shape=v.shape) for v in tf_vars]
   assign_ops = [tf.assign(v, p) for (v, p) in zip(tf_vars, placeholders)]
   global_step = tf.Variable(

From f1e161c9721f4a2a305cabb8f8071aef986c673d Mon Sep 17 00:00:00 2001
From: Mehrad Moradshahi <mehrad@stanford.edu>
Date: Mon, 29 Oct 2018 10:33:03 -0700
Subject: [PATCH 1099/2720] internal merge of PR #1171

PiperOrigin-RevId: 219151604
---
 tensor2tensor/utils/avg_checkpoints.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index 879de16fa..3d57744e6 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -90,10 +90,11 @@ def main(_):
   for name in var_values:  # Average.
     var_values[name] /= len(checkpoints)
 
-  tf_vars = [
-      tf.get_variable(v, shape=var_values[v].shape, dtype=var_dtypes[name])
-      for v in var_values
-  ]
+  with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+    tf_vars = [
+        tf.get_variable(v, shape=var_values[v].shape, dtype=var_dtypes[v])
+        for v in var_values
+    ]
   placeholders = [tf.placeholder(v.dtype, shape=v.shape) for v in tf_vars]
   assign_ops = [tf.assign(v, p) for (v, p) in zip(tf_vars, placeholders)]
   global_step = tf.Variable(

From bb08c0f987e0d2dc3c5f3ad407768dba319d4647 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 29 Oct 2018 10:56:52 -0700
Subject: [PATCH 1100/2720] Allow T2T model to generate labels in-graph.

PiperOrigin-RevId: 219156293
---
 tensor2tensor/utils/t2t_model.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 04d83f2ed..54c923c1c 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -460,6 +460,11 @@ def top(self, body_output, features):
     Returns:
       logits: dict of str to Tensor, denoting each logits for each target; or
         a single Tensor denoting the logits for that target.
+        When targets are generated at training time:
+          logits == {
+            "self_generated_targets": <generated targets tensor>
+            "logits": <original logits Tensor or dict>
+          }
     """
     if isinstance(body_output, dict):
       if self._problem_hparams:
@@ -1306,6 +1311,22 @@ def estimator_model_fn(cls,
     else:
       logits, losses_dict = model(features)  # pylint: disable=not-callable
 
+    # Support model-generated labels by overriding features["targets"] with
+    # logits["self_generated_targets"].
+    if isinstance(logits, dict) and "self_generated_targets" in logits:
+      # Overwrite 'features["targets"]' and 'labels'
+      # by logits["self_generated_targets"].
+      tf.logging.info("Replacing targets with model-provided targets.")
+      features["targets"] = labels = logits.pop("self_generated_targets")
+      assert logits.keys() == ["logits"], (
+          # See "Returns" in the "top" method docstring for the expected
+          # "logits" format when targets are generated at training time.
+          "Expect only key 'logits' when there is 'self_generated_targets'. "
+          "Found {}".format(logits.keys())
+      )
+      # Recover the original logits tensor from the logits dict.
+      logits = logits["logits"]  # Can be a tf.Tensor or a dict.
+
     # Set known shapes
     if common_layers.is_xla_compiled():
       if isinstance(logits, dict):

From ea693368d0d0d811a87f91dcb965725eb60b9014 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 29 Oct 2018 12:36:17 -0700
Subject: [PATCH 1101/2720] Fix internal tests.

PiperOrigin-RevId: 219173899
---
 tensor2tensor/data_generators/algorithmic_math.py | 13 +------------
 1 file changed, 1 insertion(+), 12 deletions(-)

diff --git a/tensor2tensor/data_generators/algorithmic_math.py b/tensor2tensor/data_generators/algorithmic_math.py
index 1e08af9e6..fcf038c7a 100644
--- a/tensor2tensor/data_generators/algorithmic_math.py
+++ b/tensor2tensor/data_generators/algorithmic_math.py
@@ -15,7 +15,6 @@
 
 """Algorithmic data generators for symbolic math tasks.
 
-See go/symbolic-math-dataset
 """
 
 from __future__ import absolute_import
@@ -159,8 +158,6 @@ def random_expr(depth, vlist, ops):
 def algebra_inverse_solve(left, right, var, solve_ops):
   """Solves for the value of the given var in an expression.
 
-  See go/symbolic-math-dataset.
-
   Args:
     left: The root of the ExprNode tree on the left side of the equals sign.
     right: The root of the ExprNode tree on the right side of the equals sign.
@@ -243,8 +240,6 @@ def generate_algebra_inverse_sample(vlist, ops, solve_ops, min_depth,
   Given an input equation and variable, produce the expression equal to the
   variable.
 
-  See go/symbolic-math-dataset.
-
   Args:
     vlist: Variable list. List of chars that can be used in the expression.
     ops: List of ExprOp instances. The allowed operators for the expression.
@@ -284,8 +279,6 @@ def generate_algebra_simplify_sample(vlist, ops, min_depth, max_depth):
 
   Given an input expression, produce the simplified expression.
 
-  See go/symbolic-math-dataset.
-
   Args:
     vlist: Variable list. List of chars that can be used in the expression.
     ops: List of ExprOp instances. The allowed operators for the expression.
@@ -312,8 +305,6 @@ def generate_calculus_integrate_sample(vlist, ops, min_depth, max_depth,
 
   Given an input expression, produce the indefinite integral.
 
-  See go/symbolic-math-dataset.
-
   Args:
     vlist: Variable list. List of chars that can be used in the expression.
     ops: List of ExprOp instances. The allowed operators for the expression.
@@ -345,7 +336,7 @@ def generate_calculus_integrate_sample(vlist, ops, min_depth, max_depth,
 
 
 # AlgebraConfig holds objects required to generate the algebra inverse
-# dataset. See go/symbolic-math-dataset.
+# dataset.
 # vlist: Variable list. A list of chars.
 # dlist: Numberical digit list. A list of chars.
 # flist: List of special function names. A list of chars.
@@ -367,8 +358,6 @@ def generate_calculus_integrate_sample(vlist, ops, min_depth, max_depth,
 def math_dataset_init(alphabet_size=26, digits=None, functions=None):
   """Initializes required objects to generate symbolic math datasets.
 
-  See go/symbolic-math-dataset.
-
   Produces token set, ExprOp instances, solve_op dictionary, encoders, and
   decoders needed to generate the algebra inverse dataset.
 

From 8d8f8baaff01f5f65ca67dcf5767aeab9ff16a45 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 29 Oct 2018 13:07:13 -0700
Subject: [PATCH 1102/2720] Draw a border around each frame to differentiate
 between the conditioned and target frames during decoding.

PiperOrigin-RevId: 219179149
---
 tensor2tensor/data_generators/video_utils.py | 31 ++++++++++++++++++++
 tensor2tensor/utils/decoding.py              |  2 ++
 2 files changed, 33 insertions(+)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 951d7320b..0473e74c3 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -52,10 +52,33 @@ def resize_video_frames(images, size):
   return resized_images
 
 
+def create_border(video, color="blue", border_percent=2):
+  """Creates a border around each frame to differentiate input and target.
+
+  Args:
+    video: 5-D NumPy array.
+    color: string, "blue", "red" or "green".
+    border_percent: Percentarge of the frame covered by the border.
+  Returns:
+    video: 5-D NumPy array.
+  """
+  color_to_axis = {"blue": 2, "red": 0, "green": 1}
+  axis = color_to_axis[color]
+  _, _, height, width, _ = video.shape
+  border_height = np.ceil(border_percent * height / 100.0).astype(np.int)
+  border_width = np.ceil(border_percent * width / 100.0).astype(np.int)
+  video[:, :, :border_height, :, axis] = 255
+  video[:, :, -border_height:, :, axis] = 255
+  video[:, :, :, :border_width, axis] = 255
+  video[:, :, :, -border_width:, axis] = 255
+  return video
+
+
 def display_video_hooks(hook_args):
   """Hooks to display videos at decode time."""
   predictions = hook_args.predictions
   fps = hook_args.decode_hparams.frames_per_second
+  border_percent = hook_args.decode_hparams.border_percent
 
   all_summaries = []
   for decode_ind, decode in enumerate(predictions):
@@ -67,9 +90,17 @@ def display_video_hooks(hook_args):
     output_videos = np.asarray(output_videos, dtype=np.uint8)
     input_videos = np.asarray(input_videos, dtype=np.uint8)
 
+    input_videos = create_border(
+        input_videos, color="blue", border_percent=border_percent)
+    target_videos = create_border(
+        target_videos, color="red", border_percent=border_percent)
+    output_videos = create_border(
+        output_videos, color="red", border_percent=border_percent)
+
     # Video gif.
     all_input = np.concatenate((input_videos, target_videos), axis=1)
     all_output = np.concatenate((input_videos, output_videos), axis=1)
+
     input_summ_vals, _ = common_video.py_gif_summary(
         "decode_%d/input" % decode_ind,
         all_input, max_outputs=10,
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 444b07913..175245db9 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -71,6 +71,8 @@ def decode_hparams(overrides=""):
       # Used for video decoding.
       frames_per_second=10,
       skip_eos_postprocess=False,
+      # Creates a blue/red border covering border_percent of the frame.
+      border_percent=2,
       # Used for MLPerf compliance logging.
       mlperf_mode=False,
       mlperf_threshold=25.0,

From 6f1ffa0b8b8921e99815651cccabe1d3cb1d766e Mon Sep 17 00:00:00 2001
From: Jade Abbott <jabbott@retrorabbit.co.za>
Date: Tue, 30 Oct 2018 00:47:49 +0200
Subject: [PATCH 1103/2720] Added English-Setswana translation problem (#1178)

---
 .../data_generators/translate_entn.py         | 67 +++++++++++++++++++
 1 file changed, 67 insertions(+)
 create mode 100644 tensor2tensor/data_generators/translate_entn.py

diff --git a/tensor2tensor/data_generators/translate_entn.py b/tensor2tensor/data_generators/translate_entn.py
new file mode 100644
index 000000000..c802eb231
--- /dev/null
+++ b/tensor2tensor/data_generators/translate_entn.py
@@ -0,0 +1,67 @@
+
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Data generators for translation data-sets."""
+
+import os
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators import translate
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+EOS = text_encoder.EOS_ID
+
+
+_ENTN_TRAIN_DATASETS = [
+    [
+        "https://github.com/LauraMartinus/ukuxhumana/blob/master/data/en_tn/eng_tswane.train.tar.gz?raw=true",
+        (
+            "entn_parallel.train.en",
+            "entn_parallel.train.tn"
+        )
+    ]
+]
+
+_ENTN_TEST_DATASETS = [
+    [
+        "https://github.com/LauraMartinus/ukuxhumana/blob/master/data/en_tn/eng_tswane.dev.tar.gz?raw=true",
+        (
+            "entn_parallel.dev.en",
+            "entn_parallel.dev.tn"
+        )
+    ]
+]
+
+
+@registry.register_problem
+class TranslateEntnRma(translate.TranslateProblem):
+  """Problem spec for English-Setswana translation using the RMA Autshumato dataset"""
+  @property
+  def approx_vocab_size(self):
+    return 2**15  # 32768
+
+  @property
+  def vocab_filename(self):
+    return "vocab.entn.%d" % self.approx_vocab_size
+
+
+  def source_data_files(self, dataset_split):
+    train = dataset_split == problem.DatasetSplit.TRAIN
+    return _ENTN_TRAIN_DATASETS if train else _ENTN_TEST_DATASETS
\ No newline at end of file

From c4c10ff36803f82ecc5211fbd20878dbf9294866 Mon Sep 17 00:00:00 2001
From: Andrew Chen <chenandrew@google.com>
Date: Mon, 29 Oct 2018 14:58:39 -0700
Subject: [PATCH 1104/2720] internal

PiperOrigin-RevId: 219199090
---
 .../data_generators/translate_entn.py         | 67 -------------------
 1 file changed, 67 deletions(-)
 delete mode 100644 tensor2tensor/data_generators/translate_entn.py

diff --git a/tensor2tensor/data_generators/translate_entn.py b/tensor2tensor/data_generators/translate_entn.py
deleted file mode 100644
index c802eb231..000000000
--- a/tensor2tensor/data_generators/translate_entn.py
+++ /dev/null
@@ -1,67 +0,0 @@
-
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Data generators for translation data-sets."""
-
-import os
-from tensor2tensor.data_generators import generator_utils
-from tensor2tensor.data_generators import problem
-from tensor2tensor.data_generators import text_encoder
-from tensor2tensor.data_generators import text_problems
-from tensor2tensor.data_generators import translate
-from tensor2tensor.utils import registry
-
-import tensorflow as tf
-
-EOS = text_encoder.EOS_ID
-
-
-_ENTN_TRAIN_DATASETS = [
-    [
-        "https://github.com/LauraMartinus/ukuxhumana/blob/master/data/en_tn/eng_tswane.train.tar.gz?raw=true",
-        (
-            "entn_parallel.train.en",
-            "entn_parallel.train.tn"
-        )
-    ]
-]
-
-_ENTN_TEST_DATASETS = [
-    [
-        "https://github.com/LauraMartinus/ukuxhumana/blob/master/data/en_tn/eng_tswane.dev.tar.gz?raw=true",
-        (
-            "entn_parallel.dev.en",
-            "entn_parallel.dev.tn"
-        )
-    ]
-]
-
-
-@registry.register_problem
-class TranslateEntnRma(translate.TranslateProblem):
-  """Problem spec for English-Setswana translation using the RMA Autshumato dataset"""
-  @property
-  def approx_vocab_size(self):
-    return 2**15  # 32768
-
-  @property
-  def vocab_filename(self):
-    return "vocab.entn.%d" % self.approx_vocab_size
-
-
-  def source_data_files(self, dataset_split):
-    train = dataset_split == problem.DatasetSplit.TRAIN
-    return _ENTN_TRAIN_DATASETS if train else _ENTN_TEST_DATASETS
\ No newline at end of file

From 87d0c7fdfb86f1bfd470c4be62a3645828bb3b98 Mon Sep 17 00:00:00 2001
From: Jade Abbott <jabbott@retrorabbit.co.za>
Date: Mon, 29 Oct 2018 15:48:06 -0700
Subject: [PATCH 1105/2720] internal merge of PR #1178

PiperOrigin-RevId: 219207387
---
 .../data_generators/translate_entn.py         | 57 +++++++++++++++++++
 1 file changed, 57 insertions(+)
 create mode 100644 tensor2tensor/data_generators/translate_entn.py

diff --git a/tensor2tensor/data_generators/translate_entn.py b/tensor2tensor/data_generators/translate_entn.py
new file mode 100644
index 000000000..2ae6e6647
--- /dev/null
+++ b/tensor2tensor/data_generators/translate_entn.py
@@ -0,0 +1,57 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Data generators for translation data-sets."""
+
+
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import translate
+from tensor2tensor.utils import registry
+
+
+EOS = text_encoder.EOS_ID
+
+_URL = "https://github.com/LauraMartinus/ukuxhumana/blob/master/data/en_tn"
+
+_ENTN_TRAIN_DATASETS = [[
+    _URL + "/eng_tswane.train.tar.gz?raw=true",
+    ("entn_parallel.train.en", "entn_parallel.train.tn")
+]]
+
+_ENTN_TEST_DATASETS = [[
+    _URL + "/eng_tswane.dev.tar.gz?raw=true",
+    ("entn_parallel.dev.en", "entn_parallel.dev.tn")
+]]
+
+
+@registry.register_problem
+class TranslateEntnRma(translate.TranslateProblem):
+  """Problem spec for English-Setswana translation.
+
+  Uses the RMA Autshumato dataset.
+  """
+
+  @property
+  def approx_vocab_size(self):
+    return 2**15  # 32768
+
+  @property
+  def vocab_filename(self):
+    return "vocab.entn.%d" % self.approx_vocab_size
+
+  def source_data_files(self, dataset_split):
+    train = dataset_split == problem.DatasetSplit.TRAIN
+    return _ENTN_TRAIN_DATASETS if train else _ENTN_TEST_DATASETS

From 42778f43821b87a7be01fd7f09b287d6473b17be Mon Sep 17 00:00:00 2001
From: Stefan Falk <43335432+stefan-falk@users.noreply.github.com>
Date: Mon, 29 Oct 2018 23:51:09 +0100
Subject: [PATCH 1106/2720] Pass data_dir to feature_encoders (#1179)

Pass data_dir to feature_encoders
---
 tensor2tensor/data_generators/common_voice.py | 2 +-
 tensor2tensor/data_generators/librispeech.py  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/common_voice.py b/tensor2tensor/data_generators/common_voice.py
index 967e24477..41935dc40 100644
--- a/tensor2tensor/data_generators/common_voice.py
+++ b/tensor2tensor/data_generators/common_voice.py
@@ -134,7 +134,7 @@ def generator(self,
 
     data_dir = os.path.join(tmp_dir, "cv_corpus_v1")
     data_tuples = _collect_data(data_dir)
-    encoders = self.feature_encoders(None)
+    encoders = self.feature_encoders(data_dir)
     audio_encoder = encoders["waveforms"]
     text_encoder = encoders["targets"]
     for dataset in datasets:
diff --git a/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index 11de4b65a..a6b7cdabc 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -138,7 +138,7 @@ def generator(self, data_dir, tmp_dir, datasets,
       data_files = _collect_data(data_dir, "flac", "txt")
       data_pairs = data_files.values()
 
-      encoders = self.feature_encoders(None)
+      encoders = self.feature_encoders(data_dir)
       audio_encoder = encoders["waveforms"]
       text_encoder = encoders["targets"]
 

From 2cc873a379b3df85066c1a2d887ae78b6376ef83 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 29 Oct 2018 16:02:59 -0700
Subject: [PATCH 1107/2720] Make Transformer decoding stable (don't fail on
 each Unicode error, work around missing hparams).

PiperOrigin-RevId: 219209875
---
 tensor2tensor/data_generators/text_encoder.py | 10 ++++++++--
 tensor2tensor/data_generators/translate.py    |  3 ++-
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index a00403c14..5580a2f22 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -60,7 +60,14 @@
 
 # Unicode utility functions that work with Python 2 and 3
 def native_to_unicode(s):
-  return s if is_unicode(s) else to_unicode(s)
+  if is_unicode(s):
+    return s
+  try:
+    return to_unicode(s)
+  except UnicodeDecodeError:
+    res = to_unicode(s, ignore_errors=True)
+    tf.logging.info("Ignoring Unicode error, outputting: %s" % res)
+    return res
 
 
 def unicode_to_native(s):
@@ -1057,4 +1064,3 @@ def decode(self, ids, strip_extraneous=False):
     """
     del strip_extraneous
     return " ".join([str(i) for i in ids])
-
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index b6ef5de06..94c29c233 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -81,7 +81,8 @@ def compute_bleu_summaries(hook_args):
   decode_hparams = hook_args.decode_hparams
   estimator = hook_args.estimator
   current_step = estimator.get_variable_value(tf.GraphKeys.GLOBAL_STEP)
-  if current_step and decode_hparams.iterations_per_loop:
+  has_iters = hasattr(decode_hparams, "iterations_per_loop")
+  if current_step and has_iters and decode_hparams.iterations_per_loop:
     iterations_per_loop = decode_hparams.iterations_per_loop
     current_epoch = np.asscalar(current_step) // iterations_per_loop
   else:

From 6def2ceae17aa893e69df85202a1e082fff62622 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 29 Oct 2018 16:41:15 -0700
Subject: [PATCH 1108/2720] Remove deprecated ps_devices argument.

PiperOrigin-RevId: 219215973
---
 tensor2tensor/models/research/transformer_moe.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index 20bd40898..e6ee2a3f0 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -106,7 +106,6 @@ def decorated(x, *args, **kwargs):
     layers = common_attention.get_standardized_layers(
         hparams=hparams,
         dp=dp,
-        ps_devices=self._ps_devices,
     )
 
     if hparams.mode == tf.estimator.ModeKeys.TRAIN:

From c47bb7e3af83a583d17b6d45800a3ac9c6a63176 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 29 Oct 2018 17:23:32 -0700
Subject: [PATCH 1109/2720] tf.test.TestCase doesn't have assertLen but
 parametrized.TestCase does, so use that to fix broken Travis.

PiperOrigin-RevId: 219222144
---
 tensor2tensor/data_generators/problem_test.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 6066f0fff..ff0084aac 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from absl.testing import parameterized  # for assertLen
 import numpy as np
 
 from tensor2tensor.data_generators import algorithmic
@@ -45,7 +46,7 @@ def assert_tensors_equal(sess, t1, t2, n):
   return True
 
 
-class ProblemTest(tf.test.TestCase):
+class ProblemTest(parameterized.TestCase, tf.test.TestCase):
 
   @classmethod
   def setUpClass(cls):

From ee66a0a83f978bcaf1416ef7cc3898ee13740037 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 29 Oct 2018 17:37:52 -0700
Subject: [PATCH 1110/2720] Fix open-source dependence on dimension_value.

PiperOrigin-RevId: 219223793
---
 tensor2tensor/layers/bayes.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index fb7444f7d..83b11886f 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -21,8 +21,6 @@
 
 import tensorflow as tf
 
-from tensorflow.python.framework import tensor_shape
-
 
 class Softplus(tf.keras.constraints.Constraint):
   """Softplus constraint."""
@@ -183,10 +181,12 @@ def __init__(self,
 
   def build(self, input_shape):
     input_shape = tf.TensorShape(input_shape)
-    if tensor_shape.dimension_value(input_shape[-1]) is None:
+    last_dim = input_shape[-1]
+    if isinstance(last_dim, tf.Dimension):
+      last_dim = last_dim.value
+    if last_dim is None:
       raise ValueError('The last dimension of the inputs to `Dense` '
                        'should be defined. Found `None`.')
-    last_dim = tensor_shape.dimension_value(input_shape[-1])
     self.input_spec = tf.layers.InputSpec(min_ndim=2,
                                           axes={-1: last_dim})
     self.kernel = self.kernel_initializer([last_dim, self.units],

From ea96ab3df08aa42a8571ac89503a2856b612b6a3 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 29 Oct 2018 17:57:04 -0700
Subject: [PATCH 1111/2720] Drop TF 1.10 testing, add 1.12

PiperOrigin-RevId: 219225877
---
 .travis.yml | 6 +++---
 setup.py    | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index fd91c80fa..98d390c90 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -13,20 +13,20 @@ env:
     - T2T_PROBLEM=algorithmic_reverse_binary40_test
     - T2T_DATA_DIR=/tmp/t2t-data
     - T2T_TRAIN_DIR=/tmp/t2t-train
-    - TF_LATEST="1.11.*"
+    - TF_LATEST="1.12*"
     # This is necessary to have gsutil work with Python 2.7
     - BOTO_CONFIG=/dev/null
   matrix:
     # We test against recent versions of TensorFlow and tf-nightly.
     # If updating, also update TF_LATEST above
-    - TF_VERSION="1.10.*"
     - TF_VERSION="1.11.*"
+    - TF_VERSION="1.12.*"
     - TF_VERSION="tf-nightly"
 matrix:
   exclude:
     # We test against all versions in Python 2 but only the latest in Python 3
     - python: "3.6"
-      env: TF_VERSION="1.10.*"
+      env: TF_VERSION="1.11.*"
     - python: "3.6"
       env: TF_VERSION="tf-nightly"
 before_install:
diff --git a/setup.py b/setup.py
index a1c4ed5f1..a5a8edda2 100644
--- a/setup.py
+++ b/setup.py
@@ -50,8 +50,8 @@
         'tqdm',
     ],
     extras_require={
-        'tensorflow': ['tensorflow>=1.9.0'],
-        'tensorflow_gpu': ['tensorflow-gpu>=1.9.0'],
+        'tensorflow': ['tensorflow>=1.11.0'],
+        'tensorflow_gpu': ['tensorflow-gpu>=1.11.0'],
         'tensorflow-hub': ['tensorflow-hub>=0.1.1'],
         'tests': [
             'absl-py',

From 9ef32ec68a843a053724ddf75adaa32c17e7735f Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 29 Oct 2018 18:29:30 -0700
Subject: [PATCH 1112/2720] 1.12 hasn't been released yet.

PiperOrigin-RevId: 219229656
---
 .travis.yml | 6 +++---
 setup.py    | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 98d390c90..fd91c80fa 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -13,20 +13,20 @@ env:
     - T2T_PROBLEM=algorithmic_reverse_binary40_test
     - T2T_DATA_DIR=/tmp/t2t-data
     - T2T_TRAIN_DIR=/tmp/t2t-train
-    - TF_LATEST="1.12*"
+    - TF_LATEST="1.11.*"
     # This is necessary to have gsutil work with Python 2.7
     - BOTO_CONFIG=/dev/null
   matrix:
     # We test against recent versions of TensorFlow and tf-nightly.
     # If updating, also update TF_LATEST above
+    - TF_VERSION="1.10.*"
     - TF_VERSION="1.11.*"
-    - TF_VERSION="1.12.*"
     - TF_VERSION="tf-nightly"
 matrix:
   exclude:
     # We test against all versions in Python 2 but only the latest in Python 3
     - python: "3.6"
-      env: TF_VERSION="1.11.*"
+      env: TF_VERSION="1.10.*"
     - python: "3.6"
       env: TF_VERSION="tf-nightly"
 before_install:
diff --git a/setup.py b/setup.py
index a5a8edda2..a1c4ed5f1 100644
--- a/setup.py
+++ b/setup.py
@@ -50,8 +50,8 @@
         'tqdm',
     ],
     extras_require={
-        'tensorflow': ['tensorflow>=1.11.0'],
-        'tensorflow_gpu': ['tensorflow-gpu>=1.11.0'],
+        'tensorflow': ['tensorflow>=1.9.0'],
+        'tensorflow_gpu': ['tensorflow-gpu>=1.9.0'],
         'tensorflow-hub': ['tensorflow-hub>=0.1.1'],
         'tests': [
             'absl-py',

From 1922524f8c705ed4bcc678e81665ca8a00e6d20e Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 29 Oct 2018 18:30:26 -0700
Subject: [PATCH 1113/2720] Disable bayes_test in eager mode.

PiperOrigin-RevId: 219229748
---
 tensor2tensor/layers/bayes_test.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 1dffa6530..dae6b3a20 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -29,7 +29,9 @@
 
 class BayesTest(parameterized.TestCase, tf.test.TestCase):
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
+  # support for TF 1.10
+  # @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDenseReparameterization(self):
     inputs = tf.to_float(np.random.rand(5, 3, 12))
     layer = bayes.DenseReparameterization(4, activation=tf.nn.relu)
@@ -42,7 +44,9 @@ def testDenseReparameterization(self):
     # TODO(trandustin): Fix this to work with Eager.
     # self.assertNotAllClose(res1, res2)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
+  # support for TF 1.10
+  # @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDenseReparameterizationModel(self):
     inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
     model = tf.keras.Sequential([

From 109987b391b08491f8eb7c7e70fd9678707ef334 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 29 Oct 2018 18:40:11 -0700
Subject: [PATCH 1114/2720] Fast TPU decoding relies on an op added after TF
 1.10, so disable that test for now.

PiperOrigin-RevId: 219231099
---
 oss_scripts/oss_tests.sh | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index a81d47b19..19ce43e19 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -47,7 +47,8 @@ pytest \
   --ignore=tensor2tensor/data_generators/allen_brain_test.py \
   --ignore=tensor2tensor/rl \
   --ignore=tensor2tensor/models/research \
-  --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary
+  --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary \
+  --deselect=tensor2tensor/utils/beam_search_test.py::BeamSearchTest::testTPUBeam
 set_status
 
 pytest tensor2tensor/utils/registry_test.py
@@ -74,6 +75,9 @@ fi
 
 if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]
 then
+    # TODO(afrozm): Once we drop support for 1.10 we can get rid of this.
+    pytest tensor2tensor/utils/beam_search_test.py::BeamSearchTest::testTPUBeam
+    set_status
     # TODO(afrozm): Enable other tests in the RL directory.
     pytest tensor2tensor/rl/trainer_model_based_test.py
     set_status

From bc6ba78d0ec8c26a8e3363bad7311b962f18cf69 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 29 Oct 2018 19:14:18 -0700
Subject: [PATCH 1115/2720] Make travis lint green

PiperOrigin-RevId: 219234118
---
 tensor2tensor/utils/mlperf_log.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/mlperf_log.py b/tensor2tensor/utils/mlperf_log.py
index 90de37c50..68eaa003b 100644
--- a/tensor2tensor/utils/mlperf_log.py
+++ b/tensor2tensor/utils/mlperf_log.py
@@ -43,7 +43,8 @@
 import time
 import uuid
 
-from tensor2tensor.utils.mlperf_tags import *  # pylint: disable=wildcard-import
+# pylint: disable=wildcard-import,unused-wildcard-import
+from tensor2tensor.utils.mlperf_tags import *
 
 
 ROOT_DIR_GNMT = None

From e6000fc4543abee352a71fe91ab8e1ab04e553f2 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 29 Oct 2018 19:16:51 -0700
Subject: [PATCH 1116/2720] Bump setup.py version to 1.10.0

PiperOrigin-RevId: 219234318
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index a1c4ed5f1..666c75630 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.9.0',
+    version='1.10.0',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',

From 66afb76d163718885d9dd00625bd2de2e8e56481 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Mon, 29 Oct 2018 20:54:00 -0700
Subject: [PATCH 1117/2720] delete hparam.force_beginning_resets

PiperOrigin-RevId: 219241227
---
 tensor2tensor/models/research/rl.py |  4 +---
 tensor2tensor/rl/collect.py         | 18 +-----------------
 2 files changed, 2 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 28e98a6bf..d17790e5b 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -137,8 +137,7 @@ def standard_atari_env_spec(env=None, simulated=False):
       simulated_env=simulated,
       reward_range=env.reward_range,
       observation_space=env.observation_space,
-      action_space=env.action_space,
-      force_beginning_resets=True
+      action_space=env.action_space
   )
   if not simulated:
     env_spec.add_hparam("env", env)
@@ -150,7 +149,6 @@ def standard_atari_env_simulated_spec(real_env, **kwargs):
   env_spec = standard_atari_env_spec(real_env, simulated=True)
   for (name, value) in six.iteritems(kwargs):
     env_spec.add_hparam(name, value)
-  env_spec.force_beginning_resets = False
   return env_spec
 
 
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 198d4e31b..469a6447c 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -141,16 +141,12 @@ def initialization_lambda(sess):
     should_reset_var = tf.Variable(True, trainable=False)
     zeros_tensor = tf.zeros(len(batch_env))
 
-  force_beginning_resets = tf.convert_to_tensor(
-      environment_spec.force_beginning_resets
-  )
-
   def reset_ops_group():
     return tf.group(batch_env.reset(tf.range(len(batch_env))),
                     tf.assign(cumulative_rewards, zeros_tensor))
 
   reset_op = tf.cond(
-      tf.logical_or(should_reset_var.read_value(), force_beginning_resets),
+      tf.logical_or(should_reset_var.read_value(), eval_phase_t),
       reset_ops_group, tf.no_op)
 
   with tf.control_dependencies([reset_op]):
@@ -238,18 +234,6 @@ def stop_condition(i, _, resets):
         parallel_iterations=1,
         back_prop=False)
 
-  # We handle force_beginning_resets differently. We assume that all envs are
-  # reseted at the end of episod (though it happens at the beginning of the
-  # next one
-  scores_num = tf.cond(force_beginning_resets,
-                       lambda: scores_num + len(batch_env), lambda: scores_num)
-
-  with tf.control_dependencies([scores_sum]):
-    scores_sum = tf.cond(
-        force_beginning_resets,
-        lambda: scores_sum + tf.reduce_sum(cumulative_rewards.read_value()),
-        lambda: scores_sum)
-
   mean_score = tf.cond(tf.greater(scores_num, 0),
                        lambda: scores_sum / tf.cast(scores_num, tf.float32),
                        lambda: 0.)

From 091373cd5eade2f108ad0fd8a209df3a5cf76598 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 29 Oct 2018 22:30:17 -0700
Subject: [PATCH 1118/2720] Fix Problem.feature_info.

PiperOrigin-RevId: 219247790
---
 tensor2tensor/data_generators/problem.py | 16 +++++-----------
 1 file changed, 5 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 01baeeed7..d7e10758a 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -716,23 +716,17 @@ def feature_info(self):
     assert self._hparams is not None
 
     hp = self.get_hparams()
-    input_mods = hp.modality["inputs"]
-    target_mod = hp.modality["targets"]
-    vocabs = hp.vocabulary
     if self.has_inputs:
       in_id = hp.input_space_id
     out_id = hp.target_space_id
 
     features = collections.defaultdict(FeatureInfo)
+    for feature_name, modality_cls in six.iteritems(hp.modality):
+      finfo = features[feature_name]
+      finfo.modality = modality_cls
+      finfo.vocab_size = modality_cls.top_dimensionality
 
-    for name, mod in six.iteritems(input_mods):
-      finfo = features[name]
-      finfo.modality = mod
-      finfo.vocab_size = mod.top_dimensionality
-
-    features["targets"].modality = target_mod
-    features["targets"].vocab_size = target_mod.top_dimensionality
-
+    vocabs = hp.vocabulary
     for name, encoder in six.iteritems(vocabs):
       features[name].encoder = encoder
 

From d3b68f59a6111fc22a14634bf4a349cda916720e Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 29 Oct 2018 23:01:09 -0700
Subject: [PATCH 1119/2720] Re-enable the disabled pylint checks after the
 offending line.

PiperOrigin-RevId: 219249726
---
 tensor2tensor/utils/mlperf_log.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/utils/mlperf_log.py b/tensor2tensor/utils/mlperf_log.py
index 68eaa003b..94af2eb31 100644
--- a/tensor2tensor/utils/mlperf_log.py
+++ b/tensor2tensor/utils/mlperf_log.py
@@ -45,6 +45,7 @@
 
 # pylint: disable=wildcard-import,unused-wildcard-import
 from tensor2tensor.utils.mlperf_tags import *
+# pylint: enable=wildcard-import,unused-wildcard-import
 
 
 ROOT_DIR_GNMT = None

From c2f5b9846dd14b7234870527649629f27ca74fad Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 30 Oct 2018 10:31:39 -0700
Subject: [PATCH 1120/2720] Test models/research only on nightly (with python
 2.7) as the comment indicates.

PiperOrigin-RevId: 219323615
---
 oss_scripts/oss_tests.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 19ce43e19..31e321f01 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -65,7 +65,7 @@ set_status
 
 
 # Test models/research only against tf-nightly
-if [[ "$TF_VERSION" == "$TF_LATEST"  ]] || [[ "$TF_VERSION" == "tf-nightly"  ]]
+if [[ "$TRAVIS_PYTHON_VERSION" == "2.7"  ]] && [[ "$TF_VERSION" == "tf-nightly"  ]]
 then
   # Ignores:
   # * Glow requires the CIFAR-10 dataset to be generated

From 94d8c03f30357f9b37acddf449b1bfa9c7a8d253 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 30 Oct 2018 15:29:56 -0700
Subject: [PATCH 1121/2720] Add remaining MLPerf L2 compliance log for
 transformer. Also update the logging tags to the latest version.

PiperOrigin-RevId: 219377758
---
 tensor2tensor/bin/t2t_trainer.py              |  2 +-
 .../data_generators/generator_utils.py        | 19 +++++++++++++++
 .../data_generators/text_problems.py          |  6 +++++
 tensor2tensor/layers/transformer_layers.py    |  6 ++---
 tensor2tensor/models/transformer.py           |  2 --
 tensor2tensor/utils/learning_rate.py          |  5 ++++
 tensor2tensor/utils/mlperf_tags.py            | 23 ++++++++-----------
 tensor2tensor/utils/t2t_model.py              | 15 ++++++++++++
 tensor2tensor/utils/trainer_lib.py            |  1 +
 9 files changed, 59 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 792ed0501..6efc422dc 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -348,7 +348,7 @@ def run_std_server():
 
 def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
-  if FLAGS.schedule != "train":
+  if FLAGS.schedule == "train" or FLAGS.schedule == "train_eval_and_decode":
     mlperf_log.transformer_print(key=mlperf_log.RUN_START)
   if FLAGS.schedule == "run_std_server":
     run_std_server()
diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index e5b7714eb..c79babef1 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -32,6 +32,7 @@
 import six.moves.urllib_request as urllib
 
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
 
@@ -150,6 +151,15 @@ def generate_files(generator, output_filenames,
     return
   tmp_filenames = [fname + ".incomplete" for fname in output_filenames]
   num_shards = len(output_filenames)
+  # Check if is training or eval, ref: train_data_filenames().
+  if num_shards > 0:
+    if "-train" in output_filenames[0]:
+      tag = "train"
+    elif "-dev" in output_filenames[0]:
+      tag = "eval"
+    else:
+      tag = "other"
+
   writers = [tf.python_io.TFRecordWriter(fname) for fname in tmp_filenames]
   counter, shard = 0, 0
   for case in generator:
@@ -171,6 +181,14 @@ def generate_files(generator, output_filenames,
   for tmp_name, final_name in zip(tmp_filenames, output_filenames):
     tf.gfile.Rename(tmp_name, final_name)
 
+  if num_shards > 0:
+    if tag == "train":
+      mlperf_log.transformer_print(
+          key=mlperf_log.PREPROC_NUM_TRAIN_EXAMPLES, value=counter)
+    elif tag == "eval":
+      mlperf_log.transformer_print(
+          key=mlperf_log.PREPROC_NUM_EVAL_EXAMPLES, value=counter)
+
   tf.logging.info("Generated %s Examples", counter)
 
 
@@ -471,6 +489,7 @@ def generate_dataset_and_shuffle(train_gen,
                                  shuffle=True):
   generate_files(train_gen, train_paths)
   generate_files(dev_gen, dev_paths)
+  mlperf_log.transformer_print(key=mlperf_log.INPUT_ORDER)
   if shuffle:
     shuffle_dataset(train_paths + dev_paths)
 
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index cc62c706c..b9b4384a6 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -35,6 +35,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
+from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -253,6 +254,11 @@ def _maybe_pack_examples(self, generator):
         chop_long_sequences=not self.has_inputs)
 
   def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      mlperf_log.transformer_print(key=mlperf_log.PREPROC_TOKENIZE_TRAINING)
+    elif dataset_split == problem.DatasetSplit.EVAL:
+      mlperf_log.transformer_print(key=mlperf_log.PREPROC_TOKENIZE_EVAL)
+
     generator = self.generate_samples(data_dir, tmp_dir, dataset_split)
     encoder = self.get_or_create_vocab(data_dir, tmp_dir)
     return text2text_generate_encoded(generator, encoder,
diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 8719b5535..719ef1c4f 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -134,8 +134,6 @@ def transformer_encoder(encoder_input,
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_NUM_HIDDEN_LAYERS,
       value=hparams.num_encoder_layers or hparams.num_hidden_layers)
-  mlperf_log.transformer_print(
-      key=mlperf_log.MODEL_HP_ATTENTION_NUM_HEADS, value=hparams.num_heads)
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_ATTENTION_DROPOUT,
       value=hparams.attention_dropout)
@@ -237,14 +235,14 @@ def transformer_ffn_layer(x,
   if ffn_layer == "dense_relu_dense":
     # In simple convolution mode, use `pad_remover` to speed up processing.
     mlperf_log.transformer_print(
-        key=mlperf_log.MODEL_HP_FFN_FILTER,
+        key=mlperf_log.MODEL_HP_FFN_FILTER_DENSE,
         value={
             "filter_size": hparams.filter_size,
             "use_bias": "True",
             "activation": mlperf_log.RELU
         })
     mlperf_log.transformer_print(
-        key=mlperf_log.MODEL_HP_FFN_DENSE,
+        key=mlperf_log.MODEL_HP_FFN_OUTPUT_DENSE,
         value={
             "hidden_size": hparams.hidden_size,
             "use_bias": "True",
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index cced9ee3f..1552be555 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1223,8 +1223,6 @@ def transformer_decoder(decoder_input,
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_NUM_HIDDEN_LAYERS,
       value=hparams.num_decoder_layers or hparams.num_hidden_layers)
-  mlperf_log.transformer_print(
-      key=mlperf_log.MODEL_HP_ATTENTION_NUM_HEADS, value=hparams.num_heads)
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_ATTENTION_DROPOUT,
       value=hparams.attention_dropout)
diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index f93f267b4..dcb5839b0 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -19,6 +19,7 @@
 from __future__ import print_function
 import numpy as np
 
+from tensor2tensor.utils import mlperf_log
 import tensorflow as tf
 
 
@@ -56,6 +57,10 @@ def learning_rate_factor(name, step_num, hparams):
 
 def learning_rate_schedule(hparams):
   """Learning rate schedule based on hparams."""
+  mlperf_log.transformer_print(key=mlperf_log.OPT_LR, deferred=True)
+  mlperf_log.transformer_print(
+      key=mlperf_log.OPT_LR_WARMUP_STEPS,
+      value=hparams.learning_rate_warmup_steps)
   step_num = _global_step(hparams)
   schedule_string = hparams.learning_rate_schedule
   names = schedule_string.split("*")
diff --git a/tensor2tensor/utils/mlperf_tags.py b/tensor2tensor/utils/mlperf_tags.py
index 36c290293..4e599d7f4 100644
--- a/tensor2tensor/utils/mlperf_tags.py
+++ b/tensor2tensor/utils/mlperf_tags.py
@@ -42,18 +42,20 @@
 TRANSFORMER = "transformer"
 INPUT_MAX_LENGTH = "input_max_length"
 
+OPT_LR_WARMUP_STEPS = "opt_learning_rate_warmup_steps"
+
 MODEL_HP_INITIALIZER_GAIN = "model_hp_initializer_gain"
 MODEL_HP_VOCAB_SIZE = "model_hp_vocab_size"
 MODEL_HP_NUM_HIDDEN_LAYERS = "model_hp_hidden_layers"
+MODEL_HP_EMBEDDING_SHARED_WEIGHTS = "model_hp_embedding_shared_weights"
 MODEL_HP_ATTENTION_DENSE = "model_hp_attention_dense"
-MODEL_HP_ATTENTION_NUM_HEADS = "model_hp_attention_num_heads"
 MODEL_HP_ATTENTION_DROPOUT = "model_hp_attention_dropout"
-MODEL_HP_FFN_DENSE = "model_hp_ffn_dense"
-MODEL_HP_FFN_FILTER = "model_hp_ffn_filter"
+MODEL_HP_FFN_OUTPUT_DENSE = "model_hp_ffn_output_dense"
+MODEL_HP_FFN_FILTER_DENSE = "model_hp_ffn_filter_dense"
 MODEL_HP_RELU_DROPOUT = "model_hp_relu_dropout"
 MODEL_HP_LAYER_POSTPROCESS_DROPOUT = "model_hp_layer_postprocess_dropout"
 MODEL_HP_NORM = "model_hp_norm"
-MODEL_HP_SEQ_BEAM_SEARCH = "model_hp_seq_beam_search"
+MODEL_HP_SEQ_BEAM_SEARCH = "model_hp_sequence_beam_search"
 
 # ==============================================================================
 # == Tags ======================================================================
@@ -294,39 +296,34 @@
     RUN_STOP,
     RUN_FINAL,
     RUN_SET_RANDOM_SEED,
-
     PREPROC_NUM_TRAIN_EXAMPLES,
     PREPROC_NUM_EVAL_EXAMPLES,
     PREPROC_TOKENIZE_TRAINING,
     PREPROC_TOKENIZE_EVAL,
     PREPROC_VOCAB_SIZE,
-
     INPUT_BATCH_SIZE,
     INPUT_MAX_LENGTH,
     INPUT_ORDER,
-
     OPT_NAME,
     OPT_LR,
+    OPT_LR_WARMUP_STEPS,
     OPT_HP_ADAM_BETA1,
     OPT_HP_ADAM_BETA2,
     OPT_HP_ADAM_EPSILON,
-
     TRAIN_LOOP,
     TRAIN_EPOCH,
-
     EVAL_START,
     EVAL_TARGET,
     EVAL_ACCURACY,
     EVAL_STOP,
-
     MODEL_HP_INITIALIZER_GAIN,
     MODEL_HP_VOCAB_SIZE,
     MODEL_HP_NUM_HIDDEN_LAYERS,
+    MODEL_HP_EMBEDDING_SHARED_WEIGHTS,
     MODEL_HP_ATTENTION_DENSE,
-    MODEL_HP_ATTENTION_NUM_HEADS,
     MODEL_HP_ATTENTION_DROPOUT,
-    MODEL_HP_FFN_DENSE,
-    MODEL_HP_FFN_FILTER,
+    MODEL_HP_FFN_OUTPUT_DENSE,
+    MODEL_HP_FFN_FILTER_DENSE,
     MODEL_HP_RELU_DROPOUT,
     MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
     MODEL_HP_NORM,
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 54c923c1c..edd8592ca 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -35,6 +35,7 @@
 from tensor2tensor.utils import expert_utils as eu
 from tensor2tensor.utils import learning_rate
 from tensor2tensor.utils import metrics
+from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import modality
 from tensor2tensor.utils import optimize
 from tensor2tensor.utils import quantization
@@ -128,6 +129,20 @@ def __init__(self,
           target_modality.top_dimensionality):
         log_info("Unsetting shared_embedding_and_softmax_weights.")
         hparams.shared_embedding_and_softmax_weights = 0
+
+      if isinstance(target_modality, modality.Modality):
+        if hparams.hidden_size:
+          hidden_size = hparams.hidden_size
+        else:
+          hidden_size = 1024
+
+        mlperf_log.transformer_print(
+            key=mlperf_log.MODEL_HP_EMBEDDING_SHARED_WEIGHTS,
+            value={
+                "vocab_size": target_modality.top_dimensionality,
+                "hidden_size": hidden_size
+            })
+
     self._original_hparams = hparams
     self.set_mode(mode)
 
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index eed26d0ac..3976ea203 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -400,6 +400,7 @@ def train_and_evaluate(self):
 
   def train(self, max_steps=None):
     mlperf_log.transformer_print(key=mlperf_log.TRAIN_LOOP)
+    mlperf_log.transformer_print(key=mlperf_log.TRAIN_EPOCH, value=0)
     self._estimator.train(
         self._train_spec.input_fn,
         hooks=self._train_spec.hooks,

From 03d987692aeea00095c174db46db4cd266f6ea1f Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 30 Oct 2018 16:12:37 -0700
Subject: [PATCH 1122/2720] Allow displaying the best/worst decodes on the
 basis of PSNR/SSIM in addition to displaying random decodes for better
 eyeballing.

PiperOrigin-RevId: 219384773
---
 tensor2tensor/data_generators/video_utils.py  | 125 ++++++++++++------
 .../data_generators/video_utils_test.py       |   3 +-
 tensor2tensor/utils/decoding.py               |   3 +
 tensor2tensor/utils/video_metrics.py          |  28 ++--
 tensor2tensor/utils/video_metrics_test.py     |   8 +-
 5 files changed, 110 insertions(+), 57 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 0473e74c3..2dd9571b2 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -74,61 +74,98 @@ def create_border(video, color="blue", border_percent=2):
   return video
 
 
+def convert_videos_to_summaries(input_videos, output_videos, target_videos,
+                                tag, decode_hparams):
+  """Converts input, output and target videos into video summaries.
+
+  Args:
+    input_videos: 5-D NumPy array, (NTHWC) conditioning frames.
+    output_videos: 5-D NumPy array, (NTHWC) ground truth.
+    target_videos: 5-D NumPy array, (NTHWC) target frames.
+    tag: tf summary tag.
+    decode_hparams: tf.contrib.training.HParams.
+  Returns:
+    summaries: a list of tf frame-by-frame and video summaries.
+  """
+  fps = decode_hparams.frames_per_second
+  border_percent = decode_hparams.border_percent
+  max_outputs = decode_hparams.max_display_outputs
+  all_summaries = []
+  input_videos = create_border(
+      input_videos, color="blue", border_percent=border_percent)
+  target_videos = create_border(
+      target_videos, color="red", border_percent=border_percent)
+  output_videos = create_border(
+      output_videos, color="red", border_percent=border_percent)
+
+  # Video gif.
+  all_input = np.concatenate((input_videos, target_videos), axis=1)
+  all_output = np.concatenate((input_videos, output_videos), axis=1)
+  input_summ_vals, _ = common_video.py_gif_summary(
+      "%s/input" % tag, all_input, max_outputs=max_outputs, fps=fps,
+      return_summary_value=True)
+  output_summ_vals, _ = common_video.py_gif_summary(
+      "%s/output" % tag, all_output, max_outputs=max_outputs, fps=fps,
+      return_summary_value=True)
+  all_summaries.extend(input_summ_vals)
+  all_summaries.extend(output_summ_vals)
+
+  # Frame-by-frame summaries
+  iterable = zip(all_input[:max_outputs], all_output[:max_outputs])
+  for ind, (input_video, output_video) in enumerate(iterable):
+    t, h, w, c = input_video.shape
+    # Tile vertically
+    input_frames = np.reshape(input_video, (t*h, w, c))
+    output_frames = np.reshape(output_video, (t*h, w, c))
+
+    # Concat across width.
+    all_frames = np.concatenate((input_frames, output_frames), axis=1)
+    tag = "input/output/%s_sample_%d" % (tag, ind)
+    frame_by_frame_summ = image_utils.image_to_tf_summary_value(
+        all_frames, tag=tag)
+    all_summaries.append(frame_by_frame_summ)
+  return all_summaries
+
+
 def display_video_hooks(hook_args):
   """Hooks to display videos at decode time."""
   predictions = hook_args.predictions
-  fps = hook_args.decode_hparams.frames_per_second
-  border_percent = hook_args.decode_hparams.border_percent
+  max_outputs = hook_args.decode_hparams.max_display_outputs
+
+  with tf.Graph().as_default():
+    _, best_decodes = video_metrics.compute_video_metrics_from_predictions(
+        predictions)
 
   all_summaries = []
+  # Displays decodes corresponding to the best/worst metric,
+  for metric, metric_decode_inds in best_decodes.items():
+    curr_metric_inds = metric_decode_inds[:max_outputs]
+    best_inputs, best_outputs, best_targets = [], [], []
+    for sample_ind, decode_ind in enumerate(curr_metric_inds):
+      curr_decode = predictions[decode_ind][sample_ind]
+      best_inputs.append(curr_decode["inputs"])
+      best_outputs.append(curr_decode["outputs"])
+      best_targets.append(curr_decode["targets"])
+    best_inputs = np.array(best_inputs, dtype=np.uint8)
+    best_outputs = np.array(best_outputs, dtype=np.uint8)
+    best_targets = np.array(best_targets, dtype=np.uint8)
+    summaries = convert_videos_to_summaries(
+        best_inputs, best_outputs, best_targets,
+        tag=metric, decode_hparams=hook_args.decode_hparams)
+    all_summaries.extend(summaries)
+
+  # Display random decodes for ten conditioning frames.
   for decode_ind, decode in enumerate(predictions):
-
     target_videos = video_metrics.stack_data_given_key(decode, "targets")
     output_videos = video_metrics.stack_data_given_key(decode, "outputs")
     input_videos = video_metrics.stack_data_given_key(decode, "inputs")
     target_videos = np.asarray(target_videos, dtype=np.uint8)
     output_videos = np.asarray(output_videos, dtype=np.uint8)
     input_videos = np.asarray(input_videos, dtype=np.uint8)
-
-    input_videos = create_border(
-        input_videos, color="blue", border_percent=border_percent)
-    target_videos = create_border(
-        target_videos, color="red", border_percent=border_percent)
-    output_videos = create_border(
-        output_videos, color="red", border_percent=border_percent)
-
-    # Video gif.
-    all_input = np.concatenate((input_videos, target_videos), axis=1)
-    all_output = np.concatenate((input_videos, output_videos), axis=1)
-
-    input_summ_vals, _ = common_video.py_gif_summary(
-        "decode_%d/input" % decode_ind,
-        all_input, max_outputs=10,
-        fps=fps,
-        return_summary_value=True)
-    output_summ_vals, _ = common_video.py_gif_summary(
-        "decode_%d/output" % decode_ind,
-        all_output,
-        max_outputs=10,
-        fps=fps,
-        return_summary_value=True)
-    all_summaries.extend(input_summ_vals)
-    all_summaries.extend(output_summ_vals)
-
-    # Frame-by-frame summaries
-    iterable = zip(all_input[:10], all_output[:10])
-    for ind, (input_video, output_video) in enumerate(iterable):
-      t, h, w, c = input_video.shape
-      # Tile vertically
-      input_frames = np.reshape(input_video, (t*h, w, c))
-      output_frames = np.reshape(output_video, (t*h, w, c))
-
-      # Concat across width.
-      all_frames = np.concatenate((input_frames, output_frames), axis=1)
-      tag = "input/output/decode_%d_sample_%d" % (decode_ind, ind)
-      frame_by_frame_summ = image_utils.image_to_tf_summary_value(
-          all_frames, tag=tag)
-      all_summaries.append(frame_by_frame_summ)
+    summaries = convert_videos_to_summaries(
+        input_videos, output_videos, target_videos,
+        tag="decode_%d" % decode_ind, decode_hparams=hook_args.decode_hparams)
+    all_summaries.extend(summaries)
   return all_summaries
 
 
@@ -146,7 +183,7 @@ def summarize_video_metrics(hook_args):
   metrics_graph = tf.Graph()
   with metrics_graph.as_default():
     if predictions:
-      metrics_results = video_metrics.compute_video_metrics_from_predictions(
+      metrics_results, _ = video_metrics.compute_video_metrics_from_predictions(
           predictions)
     else:
       metrics_results, _ = video_metrics.compute_video_metrics_from_png_files(
diff --git a/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index 66a6394a8..bdb2dd249 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -73,8 +73,9 @@ def testConvertPredictionsToVideoSummaries(self):
         hparams=decode_hparams, decode_hparams=decode_hparams,
         predictions=predictions)
     summaries = video_utils.display_video_hooks(decode_hooks)
+    # for {random, psnr_max, psnr_min, ssim_max, ssim_min}
     # 10 input vids + 10 output vids + 10 frame-by-frame.
-    self.assertEqual(len(summaries), 30)
+    self.assertEqual(len(summaries), 150)
     for summary in summaries:
       self.assertTrue(isinstance(summary, tf.Summary.Value))
 
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 175245db9..1a5c2836f 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -73,6 +73,9 @@ def decode_hparams(overrides=""):
       skip_eos_postprocess=False,
       # Creates a blue/red border covering border_percent of the frame.
       border_percent=2,
+      # Maximum number of videos displayed.
+      # Total number of videos are max_display_outputs * num_decodes
+      max_display_outputs=10,
       # Used for MLPerf compliance logging.
       mlperf_mode=False,
       mlperf_threshold=25.0,
diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index 92fba7d42..371baf6fd 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -167,6 +167,7 @@ def reduce_to_best_decode(metrics, reduce_func):
     reduce_func: callable, np.argmax or np.argmin.
   Returns:
     best_metrics: 2-D numpy array, shape=(num_samples, num_frames).
+    best_decode_ind: 1-D numpy array, shape=(num_samples,)
   """
   num_videos = metrics.shape[1]
   # Take mean of the metric across the frames to approximate the video
@@ -175,34 +176,43 @@ def reduce_to_best_decode(metrics, reduce_func):
 
   # For every sample, use the decode that has a maximum mean-metric.
   best_decode_ind = reduce_func(mean_across_frames, axis=0)
-  return metrics[best_decode_ind, np.arange(num_videos), :]
+  best_metrics = metrics[best_decode_ind, np.arange(num_videos), :]
+  return best_metrics, best_decode_ind
 
 
 def compute_all_metrics_statistics(all_results):
   """Computes statistics of metrics across multiple decodings.
 
   Args:
-    all_results: dicf of 3-D numpy arrays.
+    all_results: dict of 3-D numpy arrays.
                  Each array has shape=(num_decodes, num_samples, num_frames).
   Returns:
-    statistics: dict of 1-D numpy arrays shape=(num_frames).
+    statistics: dict of 1-D numpy arrays, shape=(num_frames).
                 First the statistic (max/mean/std) is computed across the
                 decodes, then the mean is taken across num_samples.
+    decode_inds: dict of 1-D numpy arrays, shape=(num_samples,)
+                 Each element represents the index of the decode corresponding
+                 to the best statistic.
   """
   statistics = {}
+  decode_inds = {}
   all_metrics = all_results.keys()
 
   for key in all_metrics:
     values = all_results[key]
     statistics[key + "_MEAN"] = np.mean(values, axis=0)
     statistics[key + "_STD"] = np.std(values, axis=0)
-    statistics[key + "_MIN"] = reduce_to_best_decode(values, np.argmin)
-    statistics[key + "_MAX"] = reduce_to_best_decode(values, np.argmax)
+    min_stats, min_decode_ind = reduce_to_best_decode(values, np.argmin)
+    statistics[key + "_MIN"] = min_stats
+    decode_inds[key + "_MIN_DECODE"] = min_decode_ind
+    max_stats, max_decode_ind = reduce_to_best_decode(values, np.argmax)
+    statistics[key + "_MAX"] = max_stats
+    decode_inds[key + "_MAX_DECODE"] = max_decode_ind
 
   # Computes mean of each statistic across the dataset.
   for key in statistics:
     statistics[key] = np.mean(statistics[key], axis=0)
-  return statistics
+  return statistics, decode_inds
 
 
 def compute_video_metrics_from_predictions(predictions):
@@ -224,8 +234,7 @@ def compute_video_metrics_from_predictions(predictions):
   psnr_all_decodes = np.array(psnr_all_decodes)
   ssim_all_decodes = np.array(ssim_all_decodes)
   all_results = {"PSNR": psnr_all_decodes, "SSIM": ssim_all_decodes}
-  statistics = compute_all_metrics_statistics(all_results)
-  return statistics
+  return compute_all_metrics_statistics(all_results)
 
 
 def compute_video_metrics_from_png_files(
@@ -258,8 +267,7 @@ def compute_video_metrics_from_png_files(
   psnr_all_decodes = np.array(psnr_all_decodes)
   ssim_all_decodes = np.array(ssim_all_decodes)
   all_results = {"PSNR": psnr_all_decodes, "SSIM": ssim_all_decodes}
-  statistics = compute_all_metrics_statistics(all_results)
-  return statistics, all_results
+  return compute_all_metrics_statistics(all_results)
 
 
 def compute_and_save_video_metrics(
diff --git a/tensor2tensor/utils/video_metrics_test.py b/tensor2tensor/utils/video_metrics_test.py
index 5bf01965a..18ddcec5c 100644
--- a/tensor2tensor/utils/video_metrics_test.py
+++ b/tensor2tensor/utils/video_metrics_test.py
@@ -38,8 +38,10 @@ def test_reduce_to_best_decode(self):
         [25.0, 25.0, 25.0, 25.0]]
     all_decodes = [decode1, decode2]
     all_decodes = np.array(all_decodes)
-    best_decode = video_metrics.reduce_to_best_decode(all_decodes, np.argmax)
-    worst_decode = video_metrics.reduce_to_best_decode(all_decodes, np.argmin)
+    best_decode, best_decode_ind = video_metrics.reduce_to_best_decode(
+        all_decodes, np.argmax)
+    worst_decode, worst_decode_ind = video_metrics.reduce_to_best_decode(
+        all_decodes, np.argmin)
     exp_best_decode = [
         [30.0, 32.0, 33.0, 34.0],
         [30.0, 32.0, 33.0, 34.0],
@@ -50,6 +52,8 @@ def test_reduce_to_best_decode(self):
         [30.0, 10.0, 30.0, 10.0]]
     self.assertTrue(np.allclose(best_decode, exp_best_decode))
     self.assertTrue(np.allclose(worst_decode, exp_worst_decode))
+    self.assertTrue(np.allclose(best_decode_ind, [0, 1, 1]))
+    self.assertTrue(np.allclose(worst_decode_ind, [1, 0, 0]))
 
 
 if __name__ == '__main__':

From 9729521bc3cd4952c42dcfda53699e14bee7b409 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 30 Oct 2018 17:37:28 -0700
Subject: [PATCH 1123/2720] Remove locally sharding to CPU (it was only for
 debugging anyway) and enable distribution strategy by default.

PiperOrigin-RevId: 219397234
---
 tensor2tensor/bin/t2t_trainer.py   | 5 +----
 tensor2tensor/utils/devices.py     | 3 +--
 tensor2tensor/utils/flags.py       | 3 ---
 tensor2tensor/utils/trainer_lib.py | 4 +---
 4 files changed, 3 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 6efc422dc..ba1651f3d 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -65,10 +65,8 @@
 flags.DEFINE_integer("intra_op_parallelism_threads", 0,
                      "Number of intra_op_parallelism_threads to use for CPU. "
                      "See TensorFlow config.proto for details.")
-# TODO(hinsu): Enable DistributionStrategy by default once performance gap
-# between DistributionStrategy and Parallelism is resolved.
 flags.DEFINE_bool(
-    "optionally_use_dist_strat", False,
+    "optionally_use_dist_strat", True,
     "Whether to use TensorFlow DistributionStrategy instead of explicitly "
     "replicating the model. DistributionStrategy is used only if the "
     "model replication configuration is supported by the DistributionStrategy.")
@@ -239,7 +237,6 @@ def create_run_config(hp, output_dir=None):
       keep_checkpoint_every_n_hours=FLAGS.keep_checkpoint_every_n_hours,
       num_gpus=FLAGS.worker_gpu,
       gpu_order=FLAGS.gpu_order,
-      shard_to_cpu=FLAGS.locally_shard_to_cpu,
       num_async_replicas=FLAGS.worker_replicas,
       gpu_mem_fraction=FLAGS.worker_gpu_memory_fraction,
       enable_graph_rewriter=FLAGS.enable_graph_rewriter,
diff --git a/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index 008525540..7017bb03a 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -70,7 +70,6 @@ def data_parallelism(daisy_chain_variables=True,
                      worker_replicas=1,
                      worker_id=0,
                      gpu_order="",
-                     locally_shard_to_cpu=False,
                      worker_job="/job:localhost",
                      no_data_parallelism=False):
   """See data_parallelism_from_flags."""
@@ -141,7 +140,7 @@ def _replica_device_setter(worker_device):
         "Schedule=%s. Assuming that training is running on a single machine.",
         schedule)
     datashard_devices = ["gpu:%d" % d for d in _gpu_order(worker_gpu)]
-    if locally_shard_to_cpu or worker_gpu < 1:
+    if worker_gpu < 1:
       datashard_devices += ["cpu:0"]
     caching_devices = None
   elif sync and ps_replicas > 0:
diff --git a/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index 15055f54e..94169363a 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -92,9 +92,6 @@
 flags.DEFINE_integer("eval_throttle_seconds", 600,
                      "Do not re-evaluate unless the last evaluation was started"
                      " at least this many seconds ago.")
-flags.DEFINE_bool("locally_shard_to_cpu", False,
-                  "Use CPU as a sharding device running locally. This allows "
-                  "to test sharded model construction on a machine with 1 GPU.")
 flags.DEFINE_bool("sync", False, "Sync compute on PS.")
 flags.DEFINE_string("worker_job", "/job:localhost", "name of worker job")
 flags.DEFINE_integer("worker_gpu", 1, "How many GPUs to use.")
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 3976ea203..76227545f 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -145,7 +145,6 @@ def create_run_config(model_name,
                       keep_checkpoint_every_n_hours=10000,
                       num_gpus=1,
                       gpu_order="",
-                      shard_to_cpu=False,
                       num_async_replicas=1,
                       enable_graph_rewriter=False,
                       gpu_mem_fraction=0.95,
@@ -239,7 +238,7 @@ def create_run_config(model_name,
         optionally_use_dist_strat and
         t2t_model.T2TModel.has_symmetric_shards(model_name) and
         not no_data_parallelism and ps_replicas == 0 and ps_gpu == 0 and
-        num_async_replicas == 1 and not shard_to_cpu)
+        num_async_replicas == 1)
 
     if use_distribution_strategy:
       tf.logging.info(
@@ -262,7 +261,6 @@ def create_run_config(model_name,
           worker_replicas=num_async_replicas,
           worker_id=worker_id,
           gpu_order=gpu_order,
-          locally_shard_to_cpu=shard_to_cpu,
           worker_job=worker_job,
           no_data_parallelism=no_data_parallelism)
 

From e1390c771556d6c5fe76ac64b9dc718b320981b9 Mon Sep 17 00:00:00 2001
From: Mike Dusenberry <dusenberrymw@google.com>
Date: Tue, 30 Oct 2018 18:33:06 -0700
Subject: [PATCH 1124/2720] Introduce a new TrainableInitializer to the
 Bayesian Layers.

This fixes the issue with stochastic weights in eager-mode, and essentially prototypes an unconditional stochastic layer within the existing Keras initializer framework.

PiperOrigin-RevId: 219403676
---
 tensor2tensor/layers/bayes.py      | 88 +++++++++++++++++++++++-------
 tensor2tensor/layers/bayes_test.py | 38 ++++++++++++-
 2 files changed, 103 insertions(+), 23 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 83b11886f..8c43cf288 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -39,7 +39,18 @@ def softplus():  # alias, following tf.keras.constraints
   return Softplus()
 
 
-class TrainableNormal(tf.keras.initializers.Initializer):
+class TrainableInitializer(tf.keras.initializers.Initializer):
+  """An initializer with trainable variables.
+
+  In this implementation, one must call `build` before usage in order to
+  capture the variables within the caller.
+  """
+
+  def build(self, shape, dtype=None, add_variable_fn=None):
+    raise NotImplementedError
+
+
+class TrainableNormal(TrainableInitializer):
   """Random normal op as an initializer with trainable mean and stddev."""
 
   def __init__(self,
@@ -52,7 +63,7 @@ def __init__(self,
                unconstrained_stddev_constraint=softplus(),
                seed=None,
                dtype=tf.float32):
-    """Constructs initializer."""
+    """Constructs the initializer."""
     self.mean_initializer = mean_initializer
     self.unconstrained_stddev_initializer = unconstrained_stddev_initializer
     self.mean_regularizer = mean_regularizer
@@ -62,10 +73,14 @@ def __init__(self,
     self.seed = seed
     self.dtype = tf.as_dtype(dtype)
 
-  def __call__(self, shape, dtype=None, add_variable_fn=None):
+  def build(self, shape, dtype=None, add_variable_fn=None):
+    """Builds the initializer, with the variables captured by the caller."""
     if dtype is None:
       dtype = self.dtype
-    mean = add_variable_fn(
+    self.shape = shape
+    self.dtype = dtype
+
+    self.mean = add_variable_fn(
         'mean',
         shape=shape,
         initializer=self.mean_initializer,
@@ -73,7 +88,7 @@ def __call__(self, shape, dtype=None, add_variable_fn=None):
         constraint=self.mean_constraint,
         dtype=dtype,
         trainable=True)
-    stddev = add_variable_fn(
+    self.stddev = add_variable_fn(
         'unconstrained_stddev',
         shape=shape,
         initializer=self.unconstrained_stddev_initializer,
@@ -81,10 +96,12 @@ def __call__(self, shape, dtype=None, add_variable_fn=None):
         constraint=self.unconstrained_stddev_constraint,
         dtype=dtype,
         trainable=True)
-    noise = tf.random_normal(shape, dtype=dtype, seed=self.seed)
-    output = mean + stddev * noise
+
+  def __call__(self):
+    noise = tf.random_normal(self.shape, dtype=self.dtype, seed=self.seed)
+    output = self.mean + self.stddev * noise
     # TODO(trandustin): Hack to store parameters so KL reg. can operate on them.
-    output._parameters = (mean, stddev)  # pylint: disable=protected-access
+    output._parameters = (self.mean, self.stddev)  # pylint: disable=protected-access
     return output
 
   def get_config(self):
@@ -179,6 +196,20 @@ def __init__(self,
         activity_regularizer=activity_regularizer,
         **kwargs)
 
+  @property
+  def kernel(self):
+    if isinstance(self.kernel_initializer, TrainableInitializer):
+      return self.kernel_initializer()
+    else:
+      return self._kernel
+
+  @property
+  def bias(self):
+    if isinstance(self.bias_initializer, TrainableInitializer):
+      return self.bias_initializer()
+    else:
+      return self._bias
+
   def build(self, input_shape):
     input_shape = tf.TensorShape(input_shape)
     last_dim = input_shape[-1]
@@ -187,25 +218,40 @@ def build(self, input_shape):
     if last_dim is None:
       raise ValueError('The last dimension of the inputs to `Dense` '
                        'should be defined. Found `None`.')
-    self.input_spec = tf.layers.InputSpec(min_ndim=2,
-                                          axes={-1: last_dim})
-    self.kernel = self.kernel_initializer([last_dim, self.units],
-                                          self.dtype,
-                                          self.add_weight)
+    self.input_spec = tf.layers.InputSpec(min_ndim=2, axes={-1: last_dim})
+
+    if isinstance(self.kernel_initializer, TrainableInitializer):
+      self.kernel_initializer.build([last_dim, self.units],
+                                    self.dtype,
+                                    self.add_weight)
+    else:
+      self._kernel = self.add_weight(
+          'kernel',
+          shape=[last_dim, self.units],
+          initializer=self.kernel_initializer,
+          regularizer=self.kernel_regularizer,
+          constraint=self.kernel_constraint,
+          dtype=self.dtype,
+          trainable=True)
+
     if self.kernel_regularizer is not None:
       self._handle_weight_regularization('kernel',
                                          self.kernel,
                                          self.kernel_regularizer)
+
     if self.use_bias:
-      # TODO(trandustin): Because of self.add_weight, the signature differs from
-      # other initializers, preventing interoperability.
-      if isinstance(self.bias_initializer, TrainableNormal):
-        self.bias = self.bias_initializer([self.units],
-                                          self.dtype,
-                                          self.add_weight)
+      if isinstance(self.bias_initializer, TrainableInitializer):
+        self.bias_initializer.build([self.units], self.dtype, self.add_weight)
       else:
-        self.bias = self.bias_initializer([self.units],
-                                          self.dtype)
+        self._bias = self.add_weight(
+            'bias',
+            shape=[self.units],
+            initializer=self.bias_initializer,
+            regularizer=self.bias_regularizer,
+            constraint=self.bias_constraint,
+            dtype=self.dtype,
+            trainable=True)
+
       if self.bias_regularizer is not None:
         self._handle_weight_regularization('bias',
                                            self.bias,
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index dae6b3a20..e4982126c 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -32,18 +32,52 @@ class BayesTest(parameterized.TestCase, tf.test.TestCase):
   # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
   # support for TF 1.10
   # @tf.contrib.eager.run_test_in_graph_and_eager_modes()
-  def testDenseReparameterization(self):
+  def testDenseReparameterizationKernel(self):
     inputs = tf.to_float(np.random.rand(5, 3, 12))
     layer = bayes.DenseReparameterization(4, activation=tf.nn.relu)
     outputs1 = layer(inputs)
     outputs2 = layer(inputs)
     self.evaluate(tf.global_variables_initializer())
+    # res1, res2 = self.evaluate([outputs1, outputs2])
     res1, _ = self.evaluate([outputs1, outputs2])
     self.assertEqual(res1.shape, (5, 3, 4))
     self.assertAllGreaterEqual(res1, 0.)
-    # TODO(trandustin): Fix this to work with Eager.
     # self.assertNotAllClose(res1, res2)
 
+  # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
+  # support for TF 1.10
+  # @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testDenseReparameterizationBias(self):
+    inputs = tf.to_float(np.random.rand(5, 3, 12))
+    layer = bayes.DenseReparameterization(4, kernel_initializer="zero",
+                                          bias_initializer=None,
+                                          activation=tf.nn.relu)
+    outputs1 = layer(inputs)
+    outputs2 = layer(inputs)
+    self.evaluate(tf.global_variables_initializer())
+    # res1, res2 = self.evaluate([outputs1, outputs2])
+    res1, _ = self.evaluate([outputs1, outputs2])
+    self.assertEqual(res1.shape, (5, 3, 4))
+    self.assertAllGreaterEqual(res1, 0.)
+    # self.assertNotAllClose(res1, res2)
+
+  # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
+  # support for TF 1.10
+  # @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testDenseReparameterizationDeterministic(self):
+    inputs = tf.to_float(np.random.rand(5, 3, 12))
+    layer = bayes.DenseReparameterization(4, kernel_initializer="zero",
+                                          bias_initializer="zero",
+                                          activation=tf.nn.relu)
+    outputs1 = layer(inputs)
+    outputs2 = layer(inputs)
+    self.evaluate(tf.global_variables_initializer())
+    # res1, res2 = self.evaluate([outputs1, outputs2])
+    res1, _ = self.evaluate([outputs1, outputs2])
+    self.assertEqual(res1.shape, (5, 3, 4))
+    self.assertAllGreaterEqual(res1, 0.)
+    # self.assertAllClose(res1, res2)
+
   # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
   # support for TF 1.10
   # @tf.contrib.eager.run_test_in_graph_and_eager_modes()

From b3fcb473d4285c3159fb6eee5262c64187d71268 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 30 Oct 2018 19:38:36 -0700
Subject: [PATCH 1125/2720] A recent change removed create_run_config's
 shard_to_cpu argument, so remove that in t2t_attack.py

PiperOrigin-RevId: 219408899
---
 tensor2tensor/bin/t2t_attack.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 13b350abc..b9e167238 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -111,7 +111,6 @@ def create_surrogate_run_config(hp):
       keep_checkpoint_every_n_hours=FLAGS.keep_checkpoint_every_n_hours,
       num_gpus=FLAGS.worker_gpu,
       gpu_order=FLAGS.gpu_order,
-      shard_to_cpu=FLAGS.locally_shard_to_cpu,
       num_async_replicas=FLAGS.worker_replicas,
       gpu_mem_fraction=FLAGS.worker_gpu_memory_fraction,
       enable_graph_rewriter=FLAGS.enable_graph_rewriter,

From d1c148ecc1e94962741b954247291632f54e43b4 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Wed, 31 Oct 2018 04:28:26 +0100
Subject: [PATCH 1126/2720] Model-Based RL: Reward accuracy,
 SimulatedBatchGymEnv (#1180)

* Add SimulatedBatchGymEnv - wrapper for SimulatedBatchEnv presenting purely python, gym-like api.

* Move observation decoding to T2TEnv

* Move choosing random rollout subsequences to a separate function

* Move creating simulated environment_spec to a separate function

* SimulatedBatchGymEnv fixes

* Implement reward accuracy calculation

* Fix tests
---
 tensor2tensor/data_generators/gym_env.py      |  49 +++-
 .../rl/envs/simulated_batch_gym_env.py        | 106 +++++++++
 tensor2tensor/rl/trainer_model_based.py       | 220 ++++++++++++------
 .../rl/trainer_model_based_params.py          |  11 +-
 4 files changed, 308 insertions(+), 78 deletions(-)
 create mode 100644 tensor2tensor/rl/envs/simulated_batch_gym_env.py

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index cb54e7037..27a6858f6 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -42,6 +42,25 @@
 )
 
 
+class Observation(object):
+
+  def __init__(self, data, decode_fn):
+    self.data = data
+    self._decode = decode_fn
+
+  def __eq__(self, other):
+    if isinstance(other, Observation):
+      return self.data == other.data
+    else:
+      return False
+
+  def __neq__(self, other):
+    return not self == other
+
+  def decode(self):
+    return self._decode(self.data)
+
+
 class _Noncopyable(object):
 
   def __init__(self, obj):
@@ -131,11 +150,15 @@ def __init__(self, batch_size, *args, **kwargs):
     self.current_epoch = None
     with tf.Graph().as_default() as tf_graph:
       self._tf_graph = _Noncopyable(tf_graph)
-      self._image_p = _Noncopyable(
+      self._decoded_image_p = _Noncopyable(
           tf.placeholder(dtype=tf.uint8, shape=(None, None, None))
       )
       self._encoded_image_t = _Noncopyable(
-          tf.image.encode_png(self._image_p.obj)
+          tf.image.encode_png(self._decoded_image_p.obj)
+      )
+      self._encoded_image_p = _Noncopyable(tf.placeholder(tf.string))
+      self._decoded_image_t = _Noncopyable(
+          tf.image.decode_png(self._encoded_image_p.obj)
       )
       self._session = _Noncopyable(tf.Session())
 
@@ -192,12 +215,22 @@ def _preprocess_observations(self, obs):
     """
     return obs
 
+  def _decode_png(self, encoded_observation):
+    """Decodes a single observation from PNG."""
+    return self._session.obj.run(
+        self._decoded_image_t.obj,
+        feed_dict={self._encoded_image_p.obj: encoded_observation}
+    )
+
   def _encode_observations(self, observations):
     """Encodes observations as PNG."""
     return [
-        self._session.obj.run(
-            self._encoded_image_t.obj,
-            feed_dict={self._image_p.obj: observation}
+        Observation(
+            self._session.obj.run(
+                self._encoded_image_t.obj,
+                feed_dict={self._decoded_image_p.obj: observation}
+            ),
+            self._decode_png
         )
         for observation in observations
     ]
@@ -353,7 +386,7 @@ def _generate_frames(self, rollouts):
         yield {
             "frame_number": [frame_number],
             "epoch": [self.current_epoch],
-            "image/encoded": [frame.observation],
+            "image/encoded": [frame.observation.data],
             "image/format": ["png"],
             "image/height": [self.frame_height],
             "image/width": [self.frame_width],
@@ -511,7 +544,9 @@ def get_feature_value(key, list_name):
         }
         fields["reward"] += self.reward_range[0]
         fields["done"] = bool(fields["done"])
-        fields["observation"] = fields["image/encoded"]
+        fields["observation"] = Observation(
+            fields["image/encoded"], self._decode_png
+        )
         del fields["image/encoded"]
 
         frame = Frame(**fields)
diff --git a/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
new file mode 100644
index 000000000..53e208ba1
--- /dev/null
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -0,0 +1,106 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""SimulatedBatchEnv in a Gym-like interface."""
+
+#TODO(pm): do we really need these
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import copy
+
+from tensor2tensor.utils import trainer_lib
+from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
+import tensorflow as tf
+from gym import Env
+
+
+class FlatBatchEnv(Env):
+  def __init__(self, batch_env):
+    if batch_env.batch_size != 1:
+      raise ValueError("Number of environments in batch must be equal to one")
+    self.batch_env = batch_env
+    self.action_space = self.batch_env.action_space
+    self.observation_space = self.batch_env.observation_space
+
+  def step(self, action):
+    obs, rewards, dones = self.batch_env.step([action])
+    return obs[0], rewards[0], dones[0], {}
+
+  def reset(self):
+    return self.batch_env.reset()[0]
+
+
+class SimulatedBatchGymEnv(object):
+  """ SimulatedBatchEnv in a Gym-like interface.
+
+  The environments are  batched.
+  """
+  def __init__(self, environment_spec, batch_size,
+               model_dir=None, sess=None):
+    self.batch_size = batch_size
+
+    with tf.Graph().as_default():
+      self._batch_env = SimulatedBatchEnv(environment_spec,
+                                          self.batch_size)
+
+      self.action_space = self._batch_env.action_space
+      # TODO(KC): check for the stack wrapper and correct number of channels in
+      # observation_space
+      self.observation_space = self._batch_env.observ_space
+      self._sess = sess if sess is not None else tf.Session()
+      self._to_initialize = [self._batch_env]
+
+      environment_wrappers = environment_spec.wrappers
+      wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
+
+      for w in wrappers:
+        self._batch_env = w[0](self._batch_env, **w[1])
+        self._to_initialize.append(self._batch_env)
+
+      self._sess.run(tf.global_variables_initializer())
+      for wrapped_env in self._to_initialize:
+        wrapped_env.initialize(self._sess)
+
+      self._actions_t = tf.placeholder(shape=(batch_size,), dtype=tf.int32)
+      self._rewards_t, self._dones_t = self._batch_env.simulate(self._actions_t)
+      self._obs_t = self._batch_env.observ
+      self._reset_op = self._batch_env.reset(
+          tf.range(batch_size, dtype=tf.int32)
+      )
+
+      env_model_loader = tf.train.Saver(
+          var_list=tf.global_variables(scope="next_frame*"))  # pylint:disable=unexpected-keyword-arg
+      trainer_lib.restore_checkpoint(model_dir, saver=env_model_loader,
+                                     sess=self._sess, must_restore=True)
+
+  def render(self, mode="human"):
+    raise NotImplementedError()
+
+  def reset(self, indicies=None):
+    if indicies:
+      raise NotImplementedError()
+    obs = self._sess.run(self._reset_op)
+    # TODO(pmilos): remove if possible
+    # obs[:, 0, 0, 0] = 0
+    # obs[:, 0, 0, 1] = 255
+    return obs
+
+  def step(self, actions):
+    obs, rewards, dones = self._sess.run(
+        [self._obs_t, self._rewards_t, self._dones_t],
+        feed_dict={self._actions_t: actions})
+    return obs, rewards, dones
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index ddc2f3bfd..3e8aba7a1 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -42,6 +42,7 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.rl import trainer_model_based_params
+from tensor2tensor.rl.envs.simulated_batch_gym_env import SimulatedBatchGymEnv
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -116,6 +117,33 @@ def log(msg, *args):
   return log
 
 
+def random_rollout_subsequences(rollouts, num_subsequences, subsequence_length):
+  """Chooses a random frame sequence of given length from a set of rollouts."""
+  def choose_subsequence():
+    # TODO(koz4k): Weigh rollouts by their lengths so sampling is uniform over
+    # frames and not rollouts.
+    rollout = random.choice(rollouts)
+    try:
+      from_index = random.randrange(len(rollout) - subsequence_length + 1)
+    except ValueError:
+      # Rollout too short; repeat.
+      return choose_subsequence()
+    return rollout[from_index:(from_index + subsequence_length)]
+
+  return [choose_subsequence() for _ in range(num_subsequences)]
+
+
+def make_simulated_env_spec(real_env, hparams):
+  """Creates a simulated environment_spec."""
+  return rl.standard_atari_env_simulated_spec(
+      real_env, intrinsic_reward_scale=hparams.intrinsic_reward_scale,
+      model_name=hparams.generative_model,
+      model_hparams=trainer_lib.create_hparams(hparams.generative_model_params),
+      # Hardcoded for now. TODO(koz4k): Make it a hparam.
+      video_num_input_frames=4, video_num_target_frames=1
+  )
+
+
 def train_supervised(problem, model_name, hparams, data_dir, output_dir,
                      train_steps, eval_steps, local_eval_frequency=None,
                      schedule="continuous_train_and_eval"):
@@ -153,74 +181,42 @@ def train_agent(real_env, agent_model_dir, event_dir, world_model_dir, data_dir,
   ppo_hparams.save_models_every_epochs = 10
   ppo_hparams.world_model_dir = world_model_dir
 
-  environment_spec_params = {
-      param_name: hparams.get(param_name)
-      for param_name in [
-          "intrinsic_reward_scale", "simulation_random_starts",
-          "simulation_flip_first_random_for_beginning"
-      ]
-  }
-  environment_spec_params.update({
-      "model_name": hparams.generative_model,
-      "model_hparams": trainer_lib.create_hparams(
-          hparams.generative_model_params
-      ),
-      # Hardcoded for now. TODO(koz4k): Make it a hparam.
-      "video_num_input_frames": 4,
-      "video_num_target_frames": 1
-  })
-  environment_spec = rl.standard_atari_env_simulated_spec(
-      real_env, **environment_spec_params
+  environment_spec = make_simulated_env_spec(real_env, hparams)
+
+  num_input_frames = environment_spec.video_num_input_frames
+  initial_frame_rollouts = real_env.current_epoch_rollouts(
+      split=tf.contrib.learn.ModeKeys.TRAIN,
+      minimal_rollout_frames=num_input_frames,
   )
+  # TODO(koz4k): Move this to a different module.
+  def initial_frame_chooser(batch_size):
+    """Frame chooser."""
+
+    deterministic_initial_frames =\
+        initial_frame_rollouts[0][:num_input_frames]
+    if not hparams.simulation_random_starts:
+      # Deterministic starts: repeat first frames from the first rollout.
+      initial_frames = [deterministic_initial_frames] * batch_size
+    else:
+      # Random starts: choose random initial frames from random rollouts.
+      initial_frames = random_rollout_subsequences(
+          initial_frame_rollouts, batch_size, num_input_frames
+      )
+      if hparams.simulation_flip_first_random_for_beginning:
+        # Flip first entry in the batch for deterministic initial frames.
+        initial_frames[0] = deterministic_initial_frames
+
+    return np.stack([
+        [frame.observation.decode() for frame in initial_frame_stack]
+        for initial_frame_stack in initial_frames
+    ])
+
+  environment_spec.add_hparam("initial_frame_chooser", initial_frame_chooser)
 
-  with tf.Session() as sess:
-    encoded_png_p = tf.placeholder(tf.string)
-    decoded_png_t = tf.image.decode_png(encoded_png_p)
-    def decode_png(encoded_png):
-      return sess.run(decoded_png_t, feed_dict={encoded_png_p: encoded_png})
+  ppo_hparams.add_hparam("environment_spec", environment_spec)
 
-    num_input_frames = environment_spec.video_num_input_frames
-    initial_frame_rollouts = real_env.current_epoch_rollouts(
-        split=tf.contrib.learn.ModeKeys.TRAIN,
-        minimal_rollout_frames=num_input_frames,
-    )
-    # TODO(koz4k): Move this to a different module.
-    def initial_frame_chooser(batch_size):
-      """Frame chooser."""
-
-      deterministic_initial_frames =\
-          initial_frame_rollouts[0][:num_input_frames]
-      if not environment_spec.simulation_random_starts:
-        # Deterministic starts: repeat first frames from the first rollout.
-        initial_frames = [deterministic_initial_frames] * batch_size
-      else:
-        # Random starts: choose random initial frames from random rollouts.
-        # TODO(koz4k): Weigh rollouts by their lengths so sampling is uniform
-        # over frames and not rollouts.
-        def choose_initial_frames():
-          try:
-            rollout = random.choice(initial_frame_rollouts)
-            from_index = random.randrange(len(rollout) - num_input_frames + 1)
-            return rollout[from_index:(from_index + num_input_frames)]
-          except ValueError:
-            # Rollout too short; repeat.
-            return choose_initial_frames()
-        initial_frames = [choose_initial_frames() for _ in range(batch_size)]
-        if environment_spec.simulation_flip_first_random_for_beginning:
-          # Flip first entry in the batch for deterministic initial frames.
-          initial_frames[0] = deterministic_initial_frames
-
-      return np.stack([
-          [decode_png(frame.observation) for frame in initial_frame_stack]
-          for initial_frame_stack in initial_frames
-      ])
-
-    environment_spec.add_hparam("initial_frame_chooser", initial_frame_chooser)
-
-    ppo_hparams.add_hparam("environment_spec", environment_spec)
-
-    rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
-                         name_scope="ppo_sim%d" % (epoch + 1))
+  rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
+                       name_scope="ppo_sim%d" % (epoch + 1))
 
   return completed_ppo_epochs_num
 
@@ -371,6 +367,85 @@ def compute_mean_reward(rollouts, clipped):
   return mean_rewards
 
 
+def evaluate_world_model(real_env, hparams, world_model_dir):
+  """Evaluate the world model (reward accuracy)."""
+  environment_spec = make_simulated_env_spec(real_env, hparams)
+  environment_spec.wrappers = []
+
+  num_input_frames = environment_spec.video_num_input_frames
+  rollout_subsequences = []
+  def initial_frame_chooser(batch_size):
+    assert batch_size == len(rollout_subsequences)
+    return np.stack([
+        [frame.observation.decode() for frame in subsequence[:num_input_frames]]
+        for subsequence in rollout_subsequences
+    ])
+  environment_spec.add_hparam("initial_frame_chooser", initial_frame_chooser)
+
+  sim_env = SimulatedBatchGymEnv(
+      environment_spec, hparams.wm_eval_batch_size, world_model_dir
+  )
+  subsequence_length = int(
+      max(hparams.wm_eval_rollout_ratios) * hparams.ppo_epoch_length
+  )
+  rollouts = real_env.current_epoch_rollouts(
+      split=tf.contrib.learn.ModeKeys.EVAL,
+      minimal_rollout_frames=(subsequence_length + num_input_frames)
+  )
+
+  reward_accuracies_by_length = {
+      int(ratio * hparams.ppo_epoch_length): []
+      for ratio in hparams.wm_eval_rollout_ratios
+  }
+  for _ in range(hparams.wm_eval_epochs_num):
+    rollout_subsequences[:] = random_rollout_subsequences(
+        rollouts, hparams.wm_eval_batch_size,
+        subsequence_length + num_input_frames
+    )
+
+    eval_subsequences = [
+        subsequence[(num_input_frames - 1):]
+        for subsequence in rollout_subsequences
+    ]
+
+    # Check that the initial observation is the same in the real and simulated
+    # rollout.
+    sim_init_obs = sim_env.reset()
+    real_init_obs = np.stack([
+        subsequence[0].observation.decode()
+        for subsequence in eval_subsequences
+    ])
+    assert np.all(sim_init_obs == real_init_obs)
+
+    num_same_reward = 0
+    num_steps = 0
+    (sim_cum_rewards, real_cum_rewards) = (
+        np.zeros(hparams.wm_eval_batch_size) for _ in range(2)
+    )
+    for i in range(subsequence_length):
+      actions = [subsequence[i].action for subsequence in eval_subsequences]
+      (_, sim_rewards, _) = sim_env.step(actions)
+      sim_cum_rewards += sim_rewards
+
+      real_cum_rewards += [
+          subsequence[i + 1].reward for subsequence in eval_subsequences
+      ]
+      num_same_reward += np.sum(sim_cum_rewards == real_cum_rewards)
+      num_steps += len(real_cum_rewards)
+      for (length, reward_accuracies) in six.iteritems(
+          reward_accuracies_by_length
+      ):
+        if i + 1 == length:
+          reward_accuracies.append(num_same_reward / num_steps)
+
+  return {
+      "reward_accuracy/at_{}".format(length): np.mean(reward_accuracies)
+      for (length, reward_accuracies) in six.iteritems(
+          reward_accuracies_by_length
+      )
+  }
+
+
 def summarize_metrics(eval_metrics_writer, metrics, epoch):
   """Write metrics to summary."""
   for (name, value) in six.iteritems(metrics):
@@ -419,6 +494,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   tf.logging.info("Mean training reward (initial): {}".format(
       metrics["mean_reward/train/clipped"]
   ))
+  env.generate_data(data_dir)
 
   eval_metrics_event_dir = os.path.join(directories["world_model"],
                                         "eval_metrics_event_dir")
@@ -427,8 +503,6 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   world_model_steps_num = 0
 
   for epoch in range(hparams.epochs):
-    env.generate_data(data_dir)
-
     is_final_epoch = (epoch + 1) == hparams.epochs
     log = make_log_fn(epoch, log_relative_time)
 
@@ -471,8 +545,18 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     log("Mean training reward: {}".format(metrics["mean_reward/train/clipped"]))
 
     eval_metrics = evaluate_all_configs(hparams, ppo_model_dir)
-    log("Eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
+    log("Agent eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
     metrics.update(eval_metrics)
+
+    env.generate_data(data_dir)
+
+    if hparams.eval_world_model:
+      wm_metrics = evaluate_world_model(
+          env, hparams, directories["world_model"]
+      )
+      log("World model eval metrics:\n{}".format(pprint.pformat(wm_metrics)))
+      metrics.update(wm_metrics)
+
     summarize_metrics(eval_metrics_writer, metrics, epoch)
 
     # Report metrics
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index bbe9ccfff..faf48a278 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -111,8 +111,12 @@ def rlmb_base():
       # Whether to evaluate the world model in each iteration of the loop to get
       # the model_reward_accuracy metric.
       eval_world_model=True,
-      # Rollout fractions to report reward_accuracy on.
-      eval_rollout_fractions=[0.25, 0.5, 1],
+      # Number of concurrent rollouts in world model evaluation.
+      wm_eval_batch_size=16,
+      # Number of batches to run for world model evaluation.
+      wm_eval_epochs_num=8,
+      # Ratios of ppo_epoch_length to report reward_accuracy on.
+      wm_eval_rollout_ratios=[0.25, 0.5, 1, 2],
       stop_loop_early=False,  # To speed-up tests.
       env_timesteps_limit=-1,  # Use default from gym.make()
   )
@@ -389,7 +393,7 @@ def rlmb_tiny():
           model_train_steps=2,
           ppo_epochs_num=2,
           ppo_time_limit=5,
-          ppo_epoch_length=5,
+          ppo_epoch_length=2,
           ppo_num_agents=2,
           real_ppo_epoch_length=36,
           real_ppo_num_agents=1,
@@ -401,6 +405,7 @@ def rlmb_tiny():
           resize_height_factor=2,
           resize_width_factor=2,
           game="pong",
+          wm_eval_rollout_ratios=[1],
           env_timesteps_limit=6,
       ).values())
 

From 873e15cd9b707a45f3c68cc67308a2864bf7ef6e Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Tue, 30 Oct 2018 20:28:55 -0700
Subject: [PATCH 1127/2720] internal merge of PR #1180

PiperOrigin-RevId: 219412691
---
 tensor2tensor/data_generators/gym_env.py      |  3 +-
 .../rl/envs/simulated_batch_gym_env.py        | 34 +++++--------------
 tensor2tensor/rl/trainer_model_based.py       |  8 +++++
 3 files changed, 18 insertions(+), 27 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 27a6858f6..b4e221eec 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -43,6 +43,7 @@
 
 
 class Observation(object):
+  """Observations."""
 
   def __init__(self, data, decode_fn):
     self.data = data
@@ -55,7 +56,7 @@ def __eq__(self, other):
       return False
 
   def __neq__(self, other):
-    return not self == other
+    return self != other
 
   def decode(self):
     return self._decode(self.data)
diff --git a/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index 53e208ba1..4be71d4c8 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -15,40 +15,22 @@
 
 """SimulatedBatchEnv in a Gym-like interface."""
 
-#TODO(pm): do we really need these
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
 import copy
 
-from tensor2tensor.utils import trainer_lib
-from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
-import tensorflow as tf
 from gym import Env
+from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
+from tensor2tensor.utils import trainer_lib
 
-
-class FlatBatchEnv(Env):
-  def __init__(self, batch_env):
-    if batch_env.batch_size != 1:
-      raise ValueError("Number of environments in batch must be equal to one")
-    self.batch_env = batch_env
-    self.action_space = self.batch_env.action_space
-    self.observation_space = self.batch_env.observation_space
-
-  def step(self, action):
-    obs, rewards, dones = self.batch_env.step([action])
-    return obs[0], rewards[0], dones[0], {}
-
-  def reset(self):
-    return self.batch_env.reset()[0]
+import tensorflow as tf
 
 
-class SimulatedBatchGymEnv(object):
-  """ SimulatedBatchEnv in a Gym-like interface.
+class SimulatedBatchGymEnv(Env):
+  """SimulatedBatchEnv in a Gym-like interface, environments are  batched."""
 
-  The environments are  batched.
-  """
   def __init__(self, environment_spec, batch_size,
                model_dir=None, sess=None):
     self.batch_size = batch_size
@@ -58,7 +40,7 @@ def __init__(self, environment_spec, batch_size,
                                           self.batch_size)
 
       self.action_space = self._batch_env.action_space
-      # TODO(KC): check for the stack wrapper and correct number of channels in
+      # TODO(kc): check for the stack wrapper and correct number of channels in
       # observation_space
       self.observation_space = self._batch_env.observ_space
       self._sess = sess if sess is not None else tf.Session()
@@ -90,8 +72,8 @@ def __init__(self, environment_spec, batch_size,
   def render(self, mode="human"):
     raise NotImplementedError()
 
-  def reset(self, indicies=None):
-    if indicies:
+  def reset(self, indices=None):
+    if indices:
       raise NotImplementedError()
     obs = self._sess.run(self._reset_op)
     # TODO(pmilos): remove if possible
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 3e8aba7a1..ed48ff49e 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -436,6 +436,14 @@ def initial_frame_chooser(batch_size):
           reward_accuracies_by_length
       ):
         if i + 1 == length:
+          # TODO(lukaszkaiser): resolve the comment below from blazej.
+          # If I understand correctly, num_save_reward is counting for
+          # i = 0, 1, ... , sequence_length, for how many indices i so far
+          # we had a match on simulated and real reward.
+          # I thought we would be more interested in saving just the average
+          # number of matches for the current i:
+          # reward_accuracies.append(np.sum(sim_cum_rewards == real_cum_rewards)
+          # / len(real_cum_rewards))
           reward_accuracies.append(num_same_reward / num_steps)
 
   return {

From 3201fa00b8359abab3357eb1335fa85fb7a50bf2 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Wed, 31 Oct 2018 20:38:41 +0100
Subject: [PATCH 1128/2720] Reward accuracy fixes (#1184)

* Fix Observation != operator and add docstrings

* Fix reward accuracy calculation

* Debug frames from world model evaluation
---
 tensor2tensor/data_generators/gym_env.py | 14 ++++-
 tensor2tensor/rl/trainer_model_based.py  | 72 ++++++++++++++++--------
 2 files changed, 61 insertions(+), 25 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index b4e221eec..b91b4d215 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -43,22 +43,30 @@
 
 
 class Observation(object):
-  """Observations."""
+  """Encoded observations.
+
+  Args:
+    data: Encoded observation.
+    decode_fn: Function for decoding observation.
+  """
 
   def __init__(self, data, decode_fn):
     self.data = data
     self._decode = decode_fn
 
   def __eq__(self, other):
+    """Equality comparison based on encoded data."""
     if isinstance(other, Observation):
       return self.data == other.data
     else:
       return False
 
-  def __neq__(self, other):
-    return self != other
+  def __ne__(self, other):
+    """For consistency with __eq__."""
+    return not self == other
 
   def decode(self):
+    """Decode the observation."""
     return self._decode(self.data)
 
 
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index ed48ff49e..b30b204b6 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -39,6 +39,7 @@
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
 from tensor2tensor.data_generators.gym_env import T2TGymEnv
+from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.rl import trainer_model_based_params
@@ -87,7 +88,11 @@ def setup_directories(base_dir, subdirs):
 
   all_dirs = {}
   for subdir in subdirs:
-    dir_name = os.path.join(base_dir, subdir)
+    if isinstance(subdir, six.string_types):
+      subdir_tuple = (subdir,)
+    else:
+      subdir_tuple = subdir
+    dir_name = os.path.join(base_dir, *subdir_tuple)
     tf.gfile.MakeDirs(dir_name)
     all_dirs[subdir] = dir_name
   return all_dirs
@@ -367,7 +372,7 @@ def compute_mean_reward(rollouts, clipped):
   return mean_rewards
 
 
-def evaluate_world_model(real_env, hparams, world_model_dir):
+def evaluate_world_model(real_env, hparams, world_model_dir, debug_video_path):
   """Evaluate the world model (reward accuracy)."""
   environment_spec = make_simulated_env_spec(real_env, hparams)
   environment_spec.wrappers = []
@@ -393,6 +398,10 @@ def initial_frame_chooser(batch_size):
       minimal_rollout_frames=(subsequence_length + num_input_frames)
   )
 
+  video_writer = common_video.WholeVideoWriter(
+      fps=10, output_path=debug_video_path, file_format="avi"
+  )
+
   reward_accuracies_by_length = {
       int(ratio * hparams.ppo_epoch_length): []
       for ratio in hparams.wm_eval_rollout_ratios
@@ -411,40 +420,52 @@ def initial_frame_chooser(batch_size):
     # Check that the initial observation is the same in the real and simulated
     # rollout.
     sim_init_obs = sim_env.reset()
-    real_init_obs = np.stack([
-        subsequence[0].observation.decode()
-        for subsequence in eval_subsequences
-    ])
+    def decode_real_obs(index):
+      return np.stack([
+          subsequence[index].observation.decode()
+          for subsequence in eval_subsequences
+      ])
+    real_init_obs = decode_real_obs(0)
     assert np.all(sim_init_obs == real_init_obs)
 
-    num_same_reward = 0
-    num_steps = 0
+    debug_frame_batches = []
+    def append_debug_frame_batch(sim_obs, real_obs):
+      errs = np.maximum(
+          np.abs(sim_obs.astype(np.int) - real_obs, dtype=np.int) - 10, 0
+      ).astype(np.uint8)
+      debug_frame_batches.append(  # pylint: disable=cell-var-from-loop
+          np.concatenate([sim_obs, real_obs, errs], axis=2)
+      )
+    append_debug_frame_batch(sim_init_obs, real_init_obs)
+
     (sim_cum_rewards, real_cum_rewards) = (
         np.zeros(hparams.wm_eval_batch_size) for _ in range(2)
     )
     for i in range(subsequence_length):
       actions = [subsequence[i].action for subsequence in eval_subsequences]
-      (_, sim_rewards, _) = sim_env.step(actions)
+      (sim_obs, sim_rewards, _) = sim_env.step(actions)
       sim_cum_rewards += sim_rewards
 
       real_cum_rewards += [
           subsequence[i + 1].reward for subsequence in eval_subsequences
       ]
-      num_same_reward += np.sum(sim_cum_rewards == real_cum_rewards)
-      num_steps += len(real_cum_rewards)
       for (length, reward_accuracies) in six.iteritems(
           reward_accuracies_by_length
       ):
         if i + 1 == length:
-          # TODO(lukaszkaiser): resolve the comment below from blazej.
-          # If I understand correctly, num_save_reward is counting for
-          # i = 0, 1, ... , sequence_length, for how many indices i so far
-          # we had a match on simulated and real reward.
-          # I thought we would be more interested in saving just the average
-          # number of matches for the current i:
-          # reward_accuracies.append(np.sum(sim_cum_rewards == real_cum_rewards)
-          # / len(real_cum_rewards))
-          reward_accuracies.append(num_same_reward / num_steps)
+          reward_accuracies.append(
+              np.sum(sim_cum_rewards == real_cum_rewards) /
+              len(real_cum_rewards)
+          )
+
+      real_obs = decode_real_obs(i + 1)
+      append_debug_frame_batch(sim_obs, real_obs)
+
+    for debug_frames in np.stack(debug_frame_batches, axis=1):
+      for debug_frame in debug_frames:
+        video_writer.write(debug_frame)
+
+  video_writer.finish_to_disk()
 
   return {
       "reward_accuracy/at_{}".format(length): np.mean(reward_accuracies)
@@ -469,7 +490,10 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     assert report_metric is not None
 
   # Directories
-  subdirectories = ["data", "tmp", "world_model", "ppo"]
+  subdirectories = [
+      "data", "tmp", "world_model", ("world_model", "debug_videos"),
+      "ppo"
+  ]
   directories = setup_directories(output_dir, subdirectories)
 
   epoch = -1
@@ -559,8 +583,12 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     env.generate_data(data_dir)
 
     if hparams.eval_world_model:
+      debug_video_path = os.path.join(
+          directories["world_model", "debug_videos"],
+          "{}.avi".format(env.current_epoch)
+      )
       wm_metrics = evaluate_world_model(
-          env, hparams, directories["world_model"]
+          env, hparams, directories["world_model"], debug_video_path
       )
       log("World model eval metrics:\n{}".format(pprint.pformat(wm_metrics)))
       metrics.update(wm_metrics)

From 34e2927a586a841381227f17d2dbdb14a7a9520f Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Wed, 31 Oct 2018 12:39:00 -0700
Subject: [PATCH 1129/2720] internal merge of PR #1184

PiperOrigin-RevId: 219520594
---
 tensor2tensor/rl/trainer_model_based.py | 55 ++++---------------------
 1 file changed, 9 insertions(+), 46 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index b30b204b6..6a00ebe9c 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -39,7 +39,6 @@
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
 from tensor2tensor.data_generators.gym_env import T2TGymEnv
-from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.rl import trainer_model_based_params
@@ -88,11 +87,7 @@ def setup_directories(base_dir, subdirs):
 
   all_dirs = {}
   for subdir in subdirs:
-    if isinstance(subdir, six.string_types):
-      subdir_tuple = (subdir,)
-    else:
-      subdir_tuple = subdir
-    dir_name = os.path.join(base_dir, *subdir_tuple)
+    dir_name = os.path.join(base_dir, subdir)
     tf.gfile.MakeDirs(dir_name)
     all_dirs[subdir] = dir_name
   return all_dirs
@@ -372,7 +367,7 @@ def compute_mean_reward(rollouts, clipped):
   return mean_rewards
 
 
-def evaluate_world_model(real_env, hparams, world_model_dir, debug_video_path):
+def evaluate_world_model(real_env, hparams, world_model_dir):
   """Evaluate the world model (reward accuracy)."""
   environment_spec = make_simulated_env_spec(real_env, hparams)
   environment_spec.wrappers = []
@@ -398,10 +393,6 @@ def initial_frame_chooser(batch_size):
       minimal_rollout_frames=(subsequence_length + num_input_frames)
   )
 
-  video_writer = common_video.WholeVideoWriter(
-      fps=10, output_path=debug_video_path, file_format="avi"
-  )
-
   reward_accuracies_by_length = {
       int(ratio * hparams.ppo_epoch_length): []
       for ratio in hparams.wm_eval_rollout_ratios
@@ -420,30 +411,18 @@ def initial_frame_chooser(batch_size):
     # Check that the initial observation is the same in the real and simulated
     # rollout.
     sim_init_obs = sim_env.reset()
-    def decode_real_obs(index):
-      return np.stack([
-          subsequence[index].observation.decode()
-          for subsequence in eval_subsequences
-      ])
-    real_init_obs = decode_real_obs(0)
+    real_init_obs = np.stack([
+        subsequence[0].observation.decode()
+        for subsequence in eval_subsequences
+    ])
     assert np.all(sim_init_obs == real_init_obs)
 
-    debug_frame_batches = []
-    def append_debug_frame_batch(sim_obs, real_obs):
-      errs = np.maximum(
-          np.abs(sim_obs.astype(np.int) - real_obs, dtype=np.int) - 10, 0
-      ).astype(np.uint8)
-      debug_frame_batches.append(  # pylint: disable=cell-var-from-loop
-          np.concatenate([sim_obs, real_obs, errs], axis=2)
-      )
-    append_debug_frame_batch(sim_init_obs, real_init_obs)
-
     (sim_cum_rewards, real_cum_rewards) = (
         np.zeros(hparams.wm_eval_batch_size) for _ in range(2)
     )
     for i in range(subsequence_length):
       actions = [subsequence[i].action for subsequence in eval_subsequences]
-      (sim_obs, sim_rewards, _) = sim_env.step(actions)
+      (_, sim_rewards, _) = sim_env.step(actions)
       sim_cum_rewards += sim_rewards
 
       real_cum_rewards += [
@@ -458,15 +437,6 @@ def append_debug_frame_batch(sim_obs, real_obs):
               len(real_cum_rewards)
           )
 
-      real_obs = decode_real_obs(i + 1)
-      append_debug_frame_batch(sim_obs, real_obs)
-
-    for debug_frames in np.stack(debug_frame_batches, axis=1):
-      for debug_frame in debug_frames:
-        video_writer.write(debug_frame)
-
-  video_writer.finish_to_disk()
-
   return {
       "reward_accuracy/at_{}".format(length): np.mean(reward_accuracies)
       for (length, reward_accuracies) in six.iteritems(
@@ -490,10 +460,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     assert report_metric is not None
 
   # Directories
-  subdirectories = [
-      "data", "tmp", "world_model", ("world_model", "debug_videos"),
-      "ppo"
-  ]
+  subdirectories = ["data", "tmp", "world_model", "ppo"]
   directories = setup_directories(output_dir, subdirectories)
 
   epoch = -1
@@ -583,12 +550,8 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     env.generate_data(data_dir)
 
     if hparams.eval_world_model:
-      debug_video_path = os.path.join(
-          directories["world_model", "debug_videos"],
-          "{}.avi".format(env.current_epoch)
-      )
       wm_metrics = evaluate_world_model(
-          env, hparams, directories["world_model"], debug_video_path
+          env, hparams, directories["world_model"]
       )
       log("World model eval metrics:\n{}".format(pprint.pformat(wm_metrics)))
       metrics.update(wm_metrics)

From 0c7454398e04774004143fce1ca8ca177c5710a7 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Wed, 31 Oct 2018 20:40:36 +0100
Subject: [PATCH 1130/2720] Reintroduce force_beginning_resets (#1186)

* Revert "delete hparam.force_beginning_resets"

This reverts commit 66afb76d163718885d9dd00625bd2de2e8e56481.

* Set force_beginning_resets iff simulated
---
 tensor2tensor/models/research/rl.py |  3 ++-
 tensor2tensor/rl/collect.py         | 18 +++++++++++++++++-
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index d17790e5b..5d0070962 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -137,7 +137,8 @@ def standard_atari_env_spec(env=None, simulated=False):
       simulated_env=simulated,
       reward_range=env.reward_range,
       observation_space=env.observation_space,
-      action_space=env.action_space
+      action_space=env.action_space,
+      force_beginning_resets=simulated
   )
   if not simulated:
     env_spec.add_hparam("env", env)
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 469a6447c..198d4e31b 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -141,12 +141,16 @@ def initialization_lambda(sess):
     should_reset_var = tf.Variable(True, trainable=False)
     zeros_tensor = tf.zeros(len(batch_env))
 
+  force_beginning_resets = tf.convert_to_tensor(
+      environment_spec.force_beginning_resets
+  )
+
   def reset_ops_group():
     return tf.group(batch_env.reset(tf.range(len(batch_env))),
                     tf.assign(cumulative_rewards, zeros_tensor))
 
   reset_op = tf.cond(
-      tf.logical_or(should_reset_var.read_value(), eval_phase_t),
+      tf.logical_or(should_reset_var.read_value(), force_beginning_resets),
       reset_ops_group, tf.no_op)
 
   with tf.control_dependencies([reset_op]):
@@ -234,6 +238,18 @@ def stop_condition(i, _, resets):
         parallel_iterations=1,
         back_prop=False)
 
+  # We handle force_beginning_resets differently. We assume that all envs are
+  # reseted at the end of episod (though it happens at the beginning of the
+  # next one
+  scores_num = tf.cond(force_beginning_resets,
+                       lambda: scores_num + len(batch_env), lambda: scores_num)
+
+  with tf.control_dependencies([scores_sum]):
+    scores_sum = tf.cond(
+        force_beginning_resets,
+        lambda: scores_sum + tf.reduce_sum(cumulative_rewards.read_value()),
+        lambda: scores_sum)
+
   mean_score = tf.cond(tf.greater(scores_num, 0),
                        lambda: scores_sum / tf.cast(scores_num, tf.float32),
                        lambda: 0.)

From b41b60f7ecea43f01723aa453f9860919f077c63 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 31 Oct 2018 12:46:14 -0700
Subject: [PATCH 1131/2720] Wikipedia corpus for language modeling.

PiperOrigin-RevId: 219521657
---
 tensor2tensor/data_generators/all_problems.py |  1 +
 tensor2tensor/data_generators/wiki_lm.py      | 84 +++++++++++++++++++
 2 files changed, 85 insertions(+)
 create mode 100644 tensor2tensor/data_generators/wiki_lm.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 81d961d56..45db322a2 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -83,6 +83,7 @@
     "tensor2tensor.data_generators.video_generated",
     "tensor2tensor.data_generators.vqa",
     "tensor2tensor.data_generators.wiki",
+    "tensor2tensor.data_generators.wiki_lm",
     "tensor2tensor.data_generators.wikisum.wikisum",
     "tensor2tensor.data_generators.wikitext103",
     "tensor2tensor.data_generators.wsj_parsing",
diff --git a/tensor2tensor/data_generators/wiki_lm.py b/tensor2tensor/data_generators/wiki_lm.py
new file mode 100644
index 000000000..dc07ed82d
--- /dev/null
+++ b/tensor2tensor/data_generators/wiki_lm.py
@@ -0,0 +1,84 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Data generators for untokenized wikipedia LM dataset."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+@registry.register_problem
+class LanguagemodelWiki32k(text_problems.Text2SelfProblem):
+  """A language model on the untokenized wikipedia corpus."""
+
+  # File names and Google drive ids for the training/dev/test data.
+  train_name_id = ("wiki_train.txt.gz", "1-l02fI15ieMIZk8EnXhzhsvuEYRoznZ8")
+  dev_name_id = ("wiki_dev.txt.gz", "1odhDxWKtAPKXwxRw1KCrmlrVewxdXYq7")
+  test_name_id = ("wiki_test.txt.gz", "1i1Bg6XqvdRl1LuOiIWbg7ww8Y02Ip5VK")
+
+  @property
+  def approx_vocab_size(self):
+    return 2**15  # 32768
+
+  @property
+  def max_samples_for_vocab(self):
+    return 63000
+
+  def is_generate_per_split(self):
+    return True
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    # Thresholds in the number of characters for LM examples
+    lo_thresh = 10
+    up_thresh = 256*8
+
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      (fname, fid) = self.train_name_id
+    else:
+      (fname, fid) = self.dev_name_id
+
+    wikifiles = []
+    url = "https://drive.google.com/uc?export=download&id=" + fid
+    download_path = generator_utils.maybe_download_from_drive(
+        tmp_dir, fname, url)
+    wiki_file = os.path.join(tmp_dir, fname[:-3])
+    if not tf.gfile.Exists(wiki_file):
+      generator_utils.gunzip_file(download_path, wiki_file)
+    wikifiles.append(wiki_file)
+
+    txt = ""
+    for wiki_file in wikifiles:
+      for line in tf.gfile.Open(wiki_file):
+        line = line.strip()
+        if len(txt) + len(line) > up_thresh:
+          ret = txt
+          txt = ""
+          if len(ret) > lo_thresh and len(ret) < up_thresh:
+            yield {"targets": ret}
+
+        if not txt:
+          txt = line
+        else:
+          txt = " ".join([txt, line])

From 02e3bd8a4ad342c9d77cbfe47307311ea447ca10 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 31 Oct 2018 12:53:31 -0700
Subject: [PATCH 1132/2720] Many changes to mesh-tensorflow - breaks existing
 mtf model checkpoints.

PiperOrigin-RevId: 219522764
---
 tensor2tensor/models/mtf_transformer.py       | 418 ++++++++++++------
 tensor2tensor/models/mtf_transformer_test.py  |  19 +
 tensor2tensor/models/research/moe.py          |  33 +-
 .../models/research/moe_experiments.py        | 201 +++------
 tensor2tensor/utils/decoding.py               |   8 +-
 tensor2tensor/utils/mtf_model.py              |   6 +-
 6 files changed, 403 insertions(+), 282 deletions(-)

diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index d1bf19775..65b644b45 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -22,7 +22,6 @@
 import copy
 import mesh_tensorflow as mtf
 
-from six.moves import xrange  # pylint: disable=redefined-builtin
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models.research import moe
@@ -87,16 +86,17 @@ def kv_dim(self):
   def feedforward_dim(self):
     return mtf.Dimension("d_ff", self._hparams.d_ff)
 
+  @property
+  def master_dtype(self):
+    return tf.as_dtype(self._hparams.master_dtype)
+
+  @property
+  def slice_dtype(self):
+    return tf.as_dtype(self._hparams.slice_dtype)
+
   @property
   def activation_dtype(self):
-    if self._hparams.activation_dtype == "float32":
-      return tf.float32
-    elif self._hparams.activation_dtype == "bfloat16":
-      return tf.bfloat16
-    else:
-      raise ValueError(
-          "unknown hparams.activation_dtype %s"
-          % self._hparams.activation_dtype)
+    return tf.as_dtype(self._hparams.activation_dtype)
 
   def _import_to_batch_by_length(self, x, name, mesh, hparams):
     del hparams
@@ -106,24 +106,32 @@ def _import_to_batch_by_length(self, x, name, mesh, hparams):
 
   def _embedding_and_softmax_vars(self, mesh):
     hparams = self._hparams
-    targets_embedding_var = mtf.get_variable(
-        mesh, "targets_embedding",
-        mtf.Shape([self.targets_vocab_dim, self.model_dim]),
-        initializer=tf.random_normal_initializer(),
-        activation_dtype=self.activation_dtype)
-    if self.has_input:
-      if hparams.shared_embedding:
+    if hparams.transformer_type == "encoder":
+      targets_embedding_var = None
+    else:
+      targets_embedding_var = mtf.get_variable(
+          mesh, "targets_embedding",
+          mtf.Shape([self.targets_vocab_dim, self.model_dim]),
+          initializer=tf.random_normal_initializer(),
+          master_dtype=self.master_dtype,
+          slice_dtype=self.slice_dtype,
+          activation_dtype=self.activation_dtype)
+    if hparams.transformer_type == "decoder":
+      inputs_embedding_var = None
+    else:
+      if hparams.shared_embedding and targets_embedding_var:
         inputs_embedding_var = targets_embedding_var
       else:
         inputs_embedding_var = mtf.get_variable(
             mesh, "inputs_embedding",
             mtf.Shape([self.inputs_vocab_dim, self.model_dim]),
             initializer=tf.random_normal_initializer(),
+            master_dtype=self.master_dtype,
+            slice_dtype=self.slice_dtype,
             activation_dtype=self.activation_dtype)
-    else:
-      inputs_embedding_var = None
     if hparams.shared_embedding_and_softmax_weights:
-      softmax_var = targets_embedding_var * (self.model_dim.size ** -0.5)
+      softmax_var = (targets_embedding_var or inputs_embedding_var) * (
+          self.model_dim.size ** -0.5)
     else:
       softmax_var = mtf.get_variable(
           mesh,
@@ -131,6 +139,8 @@ def _embedding_and_softmax_vars(self, mesh):
           mtf.Shape([self.targets_vocab_dim, self.model_dim]),
           initializer=tf.random_normal_initializer(
               stddev=self.model_dim.size**-0.5),
+          master_dtype=self.master_dtype,
+          slice_dtype=self.slice_dtype,
           activation_dtype=self.activation_dtype)
     positional_embedding_var = mtf.get_variable(
         mesh, "positional_embedding",
@@ -192,7 +202,10 @@ def layer_prepostprocess_dropout(x):
      targets_embedding_var,
      softmax_var,
      positional_embedding_var) = self._embedding_and_softmax_vars(mesh)
-    if self.has_input:
+    if hparams.transformer_type == "decoder":
+      encoder_output = None
+      encoder_decoder_attention_mask = None
+    else:
       inputs = tf.squeeze(tf.to_int32(features["inputs"]), [2, 3])
       inputs = pad_to_max_length(inputs)
       inputs = self._import_to_batch_by_length(inputs, "inputs", mesh, hparams)
@@ -207,16 +220,11 @@ def layer_prepostprocess_dropout(x):
         encoder_self_attention_mask = (
             mtf.layers.attention_mask_same_segment(
                 inputs_segmentation, dtype=self.activation_dtype))
-        encoder_decoder_attention_mask = (
-            mtf.layers.attention_mask_same_segment(
-                targets_segmentation, inputs_segmentation,
-                dtype=self.activation_dtype))
       else:
         inputs_position = mtf.range(mesh, self.length_dim, dtype=tf.int32)
         encoder_self_attention_mask = (
             mtf.layers.attention_mask_ignore_padding(
                 inputs, dtype=self.activation_dtype))
-        encoder_decoder_attention_mask = encoder_self_attention_mask
 
       x = (mtf.gather(inputs_embedding_var, inputs, self.inputs_vocab_dim) +
            mtf.gather(positional_embedding_var, inputs_position,
@@ -224,32 +232,39 @@ def layer_prepostprocess_dropout(x):
       x = layer_prepostprocess_dropout(x)
       with tf.variable_scope("encoder"):
         x = self._layer_stack(x,
-                              hparams.num_encoder_layers,
+                              hparams.encoder_layers,
                               self_attention_mask=encoder_self_attention_mask,
                               losses=extra_losses)
+
+    if hparams.transformer_type == "encdec":
+      if "inputs_segmentation" in features:
+        encoder_decoder_attention_mask = (
+            mtf.layers.attention_mask_same_segment(
+                targets_segmentation, inputs_segmentation,
+                dtype=self.activation_dtype))
+      else:
+        encoder_decoder_attention_mask = encoder_self_attention_mask
       encoder_output = mtf.rename_dimension(
           x, self.length_dim.name, self.memory_length_dim.name)
-    else:
-      encoder_output = None
-      encoder_decoder_attention_mask = None
-
-    # DECODER
-    x = (mtf.gather(
-        targets_embedding_var, shifted_targets, self.targets_vocab_dim) +
-         mtf.gather(
-             positional_embedding_var, targets_position, self.max_length_dim))
-    x = layer_prepostprocess_dropout(x)
 
-    # Decoder
-    with tf.variable_scope("decoder"):
-      x = self._layer_stack(
-          x,
-          hparams.num_decoder_layers,
-          encoder_output=encoder_output,
-          self_attention_mask=decoder_self_attention_mask,
-          encdec_attention_mask=encoder_decoder_attention_mask,
-          losses=extra_losses)
+    if hparams.transformer_type != "encoder":
+      # DECODER
+      x = (mtf.gather(
+          targets_embedding_var, shifted_targets, self.targets_vocab_dim) +
+           mtf.gather(
+               positional_embedding_var, targets_position, self.max_length_dim))
+      x = layer_prepostprocess_dropout(x)
+      with tf.variable_scope("decoder"):
+        x = self._layer_stack(
+            x,
+            hparams.decoder_layers,
+            encoder_output=encoder_output,
+            self_attention_mask=decoder_self_attention_mask,
+            encdec_attention_mask=encoder_decoder_attention_mask,
+            losses=extra_losses)
     logits = mtf.matmul(x, softmax_var)
+    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+      logits = mtf.layers.multiplicative_jitter(logits, epsilon=1e-2)
     off_value = hparams.label_smoothing / self._targets_vocab_size
     on_value = 1.0 - hparams.label_smoothing + off_value
     soft_targets = mtf.one_hot(
@@ -257,11 +272,17 @@ def layer_prepostprocess_dropout(x):
         dtype=self.activation_dtype)
     loss = mtf.layers.softmax_cross_entropy_with_logits(
         logits, soft_targets, self.targets_vocab_dim)
-    weights = mtf.layers.weights_nonzero(
-        targets, dtype=self.activation_dtype)
+    weights = mtf.layers.weights_nonzero(targets, dtype=self.activation_dtype)
     loss = mtf.reduce_mean(loss * weights)
     for l in extra_losses:
       loss += l
+    logits = mtf.to_float(logits)
+    # combine batch dims
+    if len(self.batch_dims) > 1:
+      combined_batch_dim = mtf.Dimension(
+          self.batch_dims[0].name, mtf.Shape(self.batch_dims).size)
+      logits = mtf.reshape(
+          logits, [combined_batch_dim] + logits.shape.dims[-2:])
     return logits, loss
 
   def mtf_model_fn(self, features, mesh):
@@ -277,18 +298,17 @@ def _targets_vocab_size(self):
 
   @property
   def _inputs_vocab_size(self):
-    if not self.has_input:
-      return None
     inputs_vocab_size = self._problem_hparams.modality[
         "inputs"].top_dimensionality
     inputs_vocab_size += (-inputs_vocab_size) % self._hparams.vocab_divisor
     return inputs_vocab_size
 
-  def _feedforward_layer(self, x, losses=None):
+  def _feedforward_layer(self, x, layer_type, losses=None):
     """Feed-forward layer.
 
     Args:
       x: a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim]
+      layer_type: a string
       losses: a list to be appended-to
     Returns:
       a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim]
@@ -296,45 +316,56 @@ def _feedforward_layer(self, x, losses=None):
       ValueError: if hparams make no sense
     """
     hparams = self._hparams
-    feedforward_layer = hparams.feedforward_layer
-    if feedforward_layer == "dense_relu_dense":
+
+    if layer_type == "drd":
       return mtf.layers.dense_relu_dense(
           x, self.feedforward_dim, dropout=hparams.relu_dropout,
-          dropout_broadcast_dims=[self.length_dim])
-    elif feedforward_layer == "moe":
+          dropout_broadcast_dims=[self.length_dim],
+          master_dtype=self.master_dtype,
+          slice_dtype=self.slice_dtype)
+    elif layer_type == "none":
+      return x
+    elif layer_type == "moe":
       output, loss = moe.transformer_moe_layer_v1(
           x,
           self.model_dim,
           hparams,
-          hparams.mode == tf.estimator.ModeKeys.TRAIN)
+          hparams.mode == tf.estimator.ModeKeys.TRAIN,
+          master_dtype=self.master_dtype,
+          slice_dtype=self.slice_dtype)
       if losses is not None:
         losses.append(loss)
       return output
-    elif feedforward_layer == "hmoe":
+    elif layer_type == "hmoe":
       output, loss = moe.transformer_moe_layer_v2(
           x,
           self.model_dim,
           hparams,
-          hparams.mode == tf.estimator.ModeKeys.TRAIN)
+          hparams.mode == tf.estimator.ModeKeys.TRAIN,
+          master_dtype=self.master_dtype,
+          slice_dtype=self.slice_dtype)
       if losses is not None:
         losses.append(loss)
       return output
     else:
-      raise ValueError(
-          "hparams.feedforward_layer not recognized %s" % feedforward_layer)
+      raise ValueError("layer_type not recognized %s" % layer_type)
 
   def _layer_stack(self,
                    x,
-                   num_layers,
+                   layers,
                    encoder_output=None,
                    self_attention_mask=None,
                    encdec_attention_mask=None,
-                   losses=None):
+                   losses=None,
+                   step_num=None,
+                   encdec_tensors=None,
+                   self_attention_k=None,
+                   self_attention_v=None):
     """Encoder or decoder stack.
 
     Args:
       x: a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim]
-      num_layers: an integer
+      layers: an list of strings
       encoder_output: an optional mtf.Tensor with shape
         [<batch_dims>, encoder_length_dim, model_dim]
       self_attention_mask: an optional mtf.Tensor with shape
@@ -342,17 +373,28 @@ def _layer_stack(self,
       encdec_attention_mask: an optional mtf.Tensor with shape
         [batch, length_dim, encoder_length_dim] containing values 0 or -inf.
       losses: a list to be appended-to
+      step_num: an optional mtf integer Scalar (used in incrmenental mode)
+      encdec_tensors: an optional list of num_layers tuples, each of the form
+        (q_var, o_var, k, v), (used in incremental mode)
+      self_attention_k: an optional list of num_layers Tensors each with shape
+        [batch, heads, memory_length, kv_channels] (incremental mode)
+      self_attention_v: an optional list of num_layers Tensors each with shape
+        [batch, heads, memory_length, kv_channels] (incremental mode)
     Returns:
       a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim]
     Raises:
       ValueError: if hparams make no sense
     """
     hparams = self._hparams
+    is_incremental = (step_num is not None)
     def layer_prepostprocess_dropout(x):
+      if is_incremental:
+        return x
       return mtf.dropout(
           x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
           noise_shape=mtf.Shape(self.batch_dims + [self.model_dim]))
-    num_layer_norms = num_layers * (2 if encoder_output is None else 3) + 1
+    num_layers = len(layers)
+    num_layer_norms = num_layers + 1
     layer_norms_dim = mtf.Dimension("layer_norms", num_layer_norms)
     layer_norm_combined_var = mtf.get_variable(
         x.mesh,
@@ -366,31 +408,79 @@ def normalize(x):
       variance = mtf.reduce_mean(mtf.square(x), reduced_dim=self.model_dim)
       return x * mtf.rsqrt(variance + hparams.norm_epsilon) * scale
 
-    for layer in range(num_layers):
-      with tf.variable_scope("layer_%d" % layer):
-        # Self attention layer
-        x += layer_prepostprocess_dropout(
-            mtf.layers.multihead_attention(
-                normalize(x), None,
-                self_attention_mask, self.kv_dim, self.heads_dim,
-                dropout=hparams.attention_dropout,
-                dropout_broadcast_dims=[self.length_dim],
-                name="self_attention"))
-        if encoder_output is not None:
+    if is_incremental:
+      new_self_attention_k = []
+      new_self_attention_v = []
+
+    for lnum, layer_type in enumerate(layers):
+      with tf.variable_scope("%s_%d" % (layer_type, lnum)):
+        if layer_type == "att":
+          # Self attention layer
+          if is_incremental:
+            self_att_num = len(new_self_attention_k)
+            y, new_k, new_v = mtf.layers.multihead_self_attention_incremental(
+                normalize(x),
+                prev_k=self_attention_k[self_att_num],
+                prev_v=self_attention_v[self_att_num],
+                step_num=step_num,
+                master_dtype=self.master_dtype,
+                slice_dtype=self.slice_dtype,
+                name="att")
+            new_self_attention_k.append(new_k)
+            new_self_attention_v.append(new_v)
+            x += y
+          else:
+            x += layer_prepostprocess_dropout(
+                mtf.layers.multihead_attention(
+                    normalize(x), None,
+                    self_attention_mask, self.kv_dim, self.heads_dim,
+                    dropout=hparams.attention_dropout,
+                    dropout_broadcast_dims=[self.length_dim],
+                    master_dtype=self.master_dtype,
+                    slice_dtype=self.slice_dtype,
+                    name="att"))
+        elif layer_type == "enc_att":
           # Encoder-Decoder attention layer
+          if is_incremental:
+            # Encoder-Decoder attention layer
+            q_var, o_var, k, v = encdec_tensors[lnum]
+            x += mtf.layers.multihead_encdec_attention_incremental(
+                normalize(x),
+                q_var, o_var, k, v,
+                encdec_attention_mask,
+                name="enc_att")
+          else:
+            x += layer_prepostprocess_dropout(
+                mtf.layers.multihead_attention(
+                    normalize(x), encoder_output,
+                    encdec_attention_mask, self.kv_dim, self.heads_dim,
+                    dropout=hparams.attention_dropout,
+                    dropout_broadcast_dims=[self.length_dim],
+                    master_dtype=self.master_dtype,
+                    slice_dtype=self.slice_dtype,
+                    name="enc_att"))
+        else:
+          if is_incremental:
+            # insert length dimension.
+            x_shape = x.shape
+            shape_with_length = mtf.Shape(
+                x_shape.dims[:-1] + [mtf.Dimension("length", 1)]
+                + x_shape.dims[-1:])
+            x = mtf.reshape(x, shape_with_length)
+          # ffn layer
           x += layer_prepostprocess_dropout(
-              mtf.layers.multihead_attention(
-                  normalize(x), encoder_output,
-                  encdec_attention_mask, self.kv_dim, self.heads_dim,
-                  dropout=hparams.attention_dropout,
-                  dropout_broadcast_dims=[self.length_dim],
-                  name="encdec_attention"))
-        # ffn layer
-        x += layer_prepostprocess_dropout(
-            self._feedforward_layer(normalize(x), losses=losses))
+              self._feedforward_layer(normalize(x), layer_type, losses=losses))
+          if is_incremental:
+            # remove length dimension
+            x = mtf.reshape(x, x_shape)
+
     x = layer_prepostprocess_dropout(normalize(x))
     assert not layer_norm_vars
-    return x
+    if is_incremental:
+      return x, new_self_attention_k, new_self_attention_v
+    else:
+      return x
+      # return mtf.cast(x, self.activation_dtype)
 
   def sample(self, features, mesh):
     with tf.variable_scope("transformer"):
@@ -402,7 +492,7 @@ def _sample(self, features, mesh):
      targets_embedding_var,
      softmax_var,
      positional_embedding_var) = self._embedding_and_softmax_vars(mesh)
-    if self.has_input:
+    if hparams.transformer_type == "encdec":
       inputs = features["inputs"]
       while len(inputs.shape.as_list()) > 2:
         inputs = tf.squeeze(inputs, axis=2)
@@ -421,29 +511,33 @@ def _sample(self, features, mesh):
               inputs, dtype=self.activation_dtype))
       with tf.variable_scope("encoder"):
         x = self._layer_stack(x,
-                              hparams.num_encoder_layers,
+                              hparams.encoder_layers,
                               self_attention_mask=encoder_attention_mask)
       encoder_output = mtf.rename_dimension(
           x, self.length_dim.name, self.memory_length_dim.name)
       encdec_tensors = []
-      for layer_num in xrange(hparams.num_decoder_layers):
-        with tf.variable_scope("decoder/layer_%d/encdec_attention" % layer_num):
-          q_var, k_var, v_var, o_var = mtf.layers.multihead_attention_vars(
-              mesh, self.heads_dim, self.model_dim,
-              self.kv_dim, self.activation_dtype)
-          k = mtf.einsum(
-              [encoder_output, k_var],
-              mtf.Shape(
-                  self.batch_dims + [self.heads_dim,
-                                     self.memory_length_dim, self.kv_dim]))
-          v = mtf.einsum(
-              [encoder_output, v_var],
-              mtf.Shape(
-                  self.batch_dims + [self.heads_dim,
-                                     self.memory_length_dim, self.kv_dim]))
-        encdec_tensors.append((q_var, o_var, k, v))
+      for layer_num, layer_type in enumerate(hparams.decoder_layers):
+        if layer_type == "enc_att":
+          with tf.variable_scope("decoder/enc_att_%d/enc_att" % layer_num):
+            q_var, k_var, v_var, o_var = mtf.layers.multihead_attention_vars(
+                mesh, self.heads_dim, self.model_dim,
+                self.kv_dim, self.master_dtype, self.slice_dtype,
+                self.activation_dtype)
+            k = mtf.einsum(
+                [encoder_output, k_var],
+                mtf.Shape(
+                    self.batch_dims + [self.heads_dim,
+                                       self.memory_length_dim, self.kv_dim]))
+            v = mtf.einsum(
+                [encoder_output, v_var],
+                mtf.Shape(
+                    self.batch_dims + [self.heads_dim,
+                                       self.memory_length_dim, self.kv_dim]))
+          encdec_tensors.append((q_var, o_var, k, v))
+        else:
+          encdec_tensors.append(None)
       partial_targets = None
-    else:
+    elif hparams.transformer_type == "decoder":
       encdec_tensors = None
       encoder_output = None
       encoder_attention_mask = None
@@ -463,6 +557,10 @@ def _sample(self, features, mesh):
                               [0, hparams.max_length - partial_targets_length]])
         partial_targets = self._import_to_batch_by_length(
             partial_targets, "partial_targets", mesh, hparams)
+    else:
+      raise ValueError(
+          "hparams.model_type = %s not yet supported"
+          % hparams.transformer_type)
 
     if hparams.beam_size == 1:
       ids_shape = mtf.Shape(self.batch_dims + [self.length_dim])
@@ -477,26 +575,28 @@ def _sample(self, features, mesh):
                             self.memory_length_dim, self.kv_dim])
 
     initial_ids = mtf.constant(mesh, 0, ids_shape, dtype=tf.int32)
+    num_self_att = len([l for l in hparams.decoder_layers if l == "att"])
     initial_kv_states = (
         [mtf.zeros(mesh, kv_shape, dtype=self.activation_dtype)]
-        * (2 * hparams.num_decoder_layers))
+        * (2 * num_self_att))
     def logits_fn(step_num, ids, states):
       """Produce logits for this step, and new states."""
-      self_attention_k = states[:hparams.num_decoder_layers]
-      self_attention_v = states[hparams.num_decoder_layers:]
+      self_attention_k = states[:num_self_att]
+      self_attention_v = states[num_self_att:]
       ids_this_step = mtf.gather(ids, step_num - 1, self.length_dim)
       x = (mtf.gather(targets_embedding_var, ids_this_step,
                       self.targets_vocab_dim) +
            mtf.gather(positional_embedding_var, step_num, self.max_length_dim))
       with tf.variable_scope("decoder"):
         x, new_self_attention_k, new_self_attention_v = (
-            self._decoder_layer_stack_incremental(
+            self._layer_stack(
                 x,
-                step_num,
-                encdec_tensors,
-                self_attention_k,
-                self_attention_v,
-                encdec_attention_mask=encoder_attention_mask))
+                hparams.decoder_layers,
+                encdec_attention_mask=encoder_attention_mask,
+                step_num=step_num,
+                encdec_tensors=encdec_tensors,
+                self_attention_k=self_attention_k,
+                self_attention_v=self_attention_v))
       logits = mtf.matmul(x, softmax_var)
       return logits, new_self_attention_k + new_self_attention_v
 
@@ -511,7 +611,7 @@ def logits_fn(step_num, ids, states):
           forced_ids=partial_targets,
           use_tpu=hparams.use_tpu)
     else:
-      if self.has_input:
+      if hparams.transformer_type == "encdec":
         input_length = mtf.reduce_sum(
             mtf.to_float(mtf.cast(inputs, tf.bool)),
             reduced_dim=self.length_dim)
@@ -527,7 +627,8 @@ def logits_fn(step_num, ids, states):
           hparams.alpha,
           states=initial_kv_states,
           decode_length=decode_length,
-          use_tpu=hparams.use_tpu)
+          use_tpu=hparams.use_tpu,
+          dtype=self.activation_dtype)
       return mtf.gather(beams, mtf.constant(mesh, 0, dtype=tf.int32), beam_dim)
 
   def _decoder_layer_stack_incremental(self,
@@ -589,7 +690,7 @@ def normalize(x):
 
     new_self_attention_k = []
     new_self_attention_v = []
-    for layer in range(num_layers):
+    for layer in xrange(num_layers):
       with tf.variable_scope("layer_%d" % layer):
         # Self attention layer
         y, new_k, new_v = mtf.layers.multihead_self_attention_incremental(
@@ -597,7 +698,9 @@ def normalize(x):
             prev_k=self_attention_k[layer],
             prev_v=self_attention_v[layer],
             step_num=step_num,
-            name="self_attention")
+            master_dtype=self.master_dtype,
+            slice_dtype=self.slice_dtype,
+            name="att")
         new_self_attention_k.append(new_k)
         new_self_attention_v.append(new_v)
         x += y
@@ -608,9 +711,9 @@ def normalize(x):
               normalize(x),
               q_var, o_var, k, v,
               encdec_attention_mask,
-              name="encdec_attention")
+              name="enc_att")
         # ffn layer
-        x += self._feedforward_layer(normalize(x), hparams)
+        x += self._feedforward_layer(normalize(x), layer)
     x = normalize(x)
     assert not layer_norm_vars
     return x, new_self_attention_k, new_self_attention_v
@@ -633,26 +736,38 @@ def mtf_transformer_base():
   hparams.add_hparam("layout", "batch:batch;vocab:model;d_ff:model;heads:model")
   hparams.add_hparam("num_heads", 8)
   hparams.add_hparam("d_ff", 2048)
-  hparams.add_hparam("num_encoder_layers", 6)
-  hparams.add_hparam("num_decoder_layers", 6)
+  hparams.add_hparam("encoder_layers", ["att", "drd"] * 6)
+  hparams.add_hparam("decoder_layers", ["att", "enc_att", "drd"] * 6)
   hparams.add_hparam("attention_dropout", 0.1)
   hparams.add_hparam("relu_dropout", 0.1)
   hparams.layer_prepostprocess_dropout = 0.1
 
+  # Describes what model architecture:
+  #   "encdec": encoder + autoregerssive decoder
+  #   "decoder": single-stack autoregressive sequence model.
+  #   "encoder": single-stack non-autoregressive model
+  #      with equal-length inputs and outputs.
+  hparams.add_hparam("transformer_type", "encdec")
+
   # round up vocab sizes to be a multiple of this value
   hparams.vocab_divisor = 128
 
   # options are dense_relu_dense, moe, hmoe
-  hparams.add_hparam("feedforward_layer", "dense_relu_dense")
+  hparams.add_hparam("feedforward_layer", "drd")
 
-  # Use targets_embedding_var * rsqrt(d_model) as softmax_var
+  # If True, then reuse targets_embedding_var * rsqrt(d_model) as softmax_var
+  # If hparams.transformer_type == "encoder", then there is no targets embedding
+  # so we reuse the inputs embedding instead.
   hparams.shared_embedding_and_softmax_weights = True
   # Reuse targets_embedding_var as inputs_embedding_var
+  # relevant only if hparams.transformer_type == "encdec"
   hparams.shared_embedding = True
   hparams.optimizer = "Adafactor"
   hparams.learning_rate_schedule = "linear_warmup*rsqrt_decay*linear_decay"
   hparams.learning_rate_warmup_steps = 10000
-  hparams.activation_dtype = "float32"
+  hparams.add_hparam("master_dtype", "bfloat16")
+  hparams.add_hparam("slice_dtype", "float32")
+  hparams.activation_dtype = "bfloat16"
 
   # These parameters make Transformer model compatible with MtfTransformer
   # Do not override these, as mtf_transformer does not support other options.
@@ -677,18 +792,39 @@ def mtf_transformer_base():
   return hparams
 
 
+@registry.register_hparams
+def mtf_transformer_base_lm():
+  hparams = mtf_transformer_base()
+  hparams.decoder_layers = hparams.encoder_layers
+  hparams.transformer_type = "decoder"
+  hparams.label_smoothing = 0.0
+  hparams.sampling_method = "random"
+  return hparams
+
+
 @registry.register_hparams
 def mtf_transformer_tiny():
   """Catch bugs locally..."""
   hparams = mtf_transformer_base()
   hparams.d_model = 128
   hparams.d_ff = 512
-  hparams.batch_size = 4
-  hparams.num_encoder_layers = 2
-  hparams.num_decoder_layers = 2
-  hparams.num_heads = 4
+  hparams.batch_size = 8
+  hparams.encoder_layers = ["att", "drd"] * 2
+  hparams.decoder_layers = ["att", "enc_att", "drd"] * 2
+  hparams.num_heads = 8
   # data parallelism and model-parallelism
-  hparams.mesh_shape = "batch:2;model:2"
+  hparams.mesh_shape = "batch:2;model:4"
+  hparams.activation_dtype = "float32"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_tiny_lm():
+  hparams = mtf_transformer_tiny()
+  hparams.decoder_layers = hparams.encoder_layers
+  hparams.transformer_type = "decoder"
+  hparams.label_smoothing = 0.0
+  hparams.sampling_method = "random"
   return hparams
 
 
@@ -699,6 +835,13 @@ def mtf_transformer_single():
   return hparams
 
 
+@registry.register_hparams
+def mtf_transformer_enc_single():
+  hparams = mtf_transformer_single()
+  hparams.transformer_type = "encoder"
+  return hparams
+
+
 @registry.register_hparams
 def mtf_transformer_tiny_8gpu():
   hparams = mtf_transformer_tiny()
@@ -733,8 +876,7 @@ def mtf_transformer_paper_lm(size):
     a hparams object
   """
   n = 2 ** size
-  hparams = mtf_transformer_base()
-  hparams.label_smoothing = 0.0
+  hparams = mtf_transformer_base_lm()
   hparams.batch_size = 256
   hparams.d_model = 1024
   hparams.d_ff = int(8192 * n)
@@ -836,6 +978,20 @@ def mtf_transformer_paper_tr_0():
   return hparams
 
 
+@registry.register_hparams
+def mtf_transformer_paper_tr_0_a32():
+  hparams = mtf_transformer_paper_tr_0()
+  hparams.activation_dtype = "float32"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer_paper_tr_0_nf():
+  hparams = mtf_transformer_paper_tr_0()
+  hparams.optimizer_adafactor_factored = False
+  return hparams
+
+
 @registry.register_hparams
 def mtf_transformer_paper_tr_1():
   hparams = mtf_transformer_paper_tr(1)
@@ -927,5 +1083,3 @@ def mtf_transformer_lm_baseline():
   hparams.learning_rate_decay_steps = 27200  # one epoch on lm1b
   hparams.mesh_shape = "batch:8"
   return hparams
-
-
diff --git a/tensor2tensor/models/mtf_transformer_test.py b/tensor2tensor/models/mtf_transformer_test.py
index 618daa146..58bdee309 100644
--- a/tensor2tensor/models/mtf_transformer_test.py
+++ b/tensor2tensor/models/mtf_transformer_test.py
@@ -151,6 +151,25 @@ def testMtfTransformerDataModelParallel(self):
       res = session.run(tf_logits)
     self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, VOCAB_SIZE))
 
+  def testMtfTransformerEncoderDataModelParallel(self):
+    hparams = mtf_transformer.mtf_transformer_enc_single()
+
+    model, features, hparams = get_model(hparams)
+    hparams.mesh_shape = "batch:2;model:2"
+    hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
+    mesh, mesh_impl = get_placement_mesh(hparams)
+
+    logits, _ = model.mtf_model_fn(features, mesh)
+    lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
+    tf_group = lowering.copy_masters_to_slices()
+    tf_logits = lowering.export_to_tf_tensor(logits)
+
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      session.run(tf_group)
+      res = session.run(tf_logits)
+    self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, VOCAB_SIZE))
+
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/models/research/moe.py b/tensor2tensor/models/research/moe.py
index e39611cdc..1dddba13d 100644
--- a/tensor2tensor/models/research/moe.py
+++ b/tensor2tensor/models/research/moe.py
@@ -27,9 +27,9 @@
 import tensorflow as tf
 
 
-
-
-def transformer_moe_layer_v1(inputs, output_dim, hparams, train):
+def transformer_moe_layer_v1(inputs, output_dim, hparams, train,
+                             master_dtype=tf.bfloat16,
+                             slice_dtype=tf.float32):
   """Local mixture of experts that works well on TPU.
 
   Adapted from the paper https://arxiv.org/abs/1701.06538
@@ -86,6 +86,8 @@ def transformer_moe_layer_v1(inputs, output_dim, hparams, train):
     output_dim: a mtf.Dimension (for Transformer, this is input_dim)
     hparams: model hyperparameters
     train: a boolean
+    master_dtype: a tf.dtype
+    slice_dtype: a tf.dtype
 
   Returns:
     outputs: a Tensor with shape [<batch_dims...>, length_dim, output_dim]
@@ -137,9 +139,11 @@ def transformer_moe_layer_v1(inputs, output_dim, hparams, train):
   # Now feed the expert inputs through the experts.
   h = mtf.layers.dense(
       expert_inputs, hidden_dim, expert_dims=[experts_dim],
-      activation=mtf.relu, use_bias=False, name="x0")
+      activation=mtf.relu, use_bias=False, master_dtype=master_dtype,
+      slice_dtype=slice_dtype, name="x0")
   expert_output = mtf.layers.dense(
-      h, output_dim, expert_dims=[experts_dim], use_bias=False, name="x1")
+      h, output_dim, expert_dims=[experts_dim], use_bias=False,
+      master_dtype=master_dtype, slice_dtype=slice_dtype, name="x1")
 
   expert_output = mtf.reshape(expert_output, mtf.Shape(
       [experts_dim_unsplit, batch_dim, expert_capacity_dim, input_dim]))
@@ -152,7 +156,8 @@ def transformer_moe_layer_v1(inputs, output_dim, hparams, train):
   return output, loss * hparams.moe_loss_coef
 
 
-def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
+def transformer_moe_layer_v2(inputs, output_dim, hparams, train,
+                             master_dtype=tf.bfloat16, slice_dtype=tf.float32):
   """2-level mixture of experts.
 
   Adapted from the paper https://arxiv.org/abs/1701.06538
@@ -245,6 +250,8 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
     output_dim: a mtf.Dimension (for Transformer, this is input_dim)
     hparams: model hyperparameters
     train: a boolean
+    master_dtype: a tf.dtype
+    slice_dtype: a tf.dtype
 
   Returns:
     outputs: a Tensor with shape [a, b, l, n]
@@ -282,9 +289,8 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
   capacity_factor = (
       hparams.moe_capacity_factor_train if train else
       hparams.moe_capacity_factor_eval)
-  expert_capacity = min(
-      s.size,
-      int((s.size * capacity_factor) / x.size))
+  expert_capacity = min(s.size, int((s.size * capacity_factor) / x.size))
+  expert_capacity = max(expert_capacity, 4)
   c = mtf.Dimension("expert_capacity_x", expert_capacity)
 
   # We "cheat" here and look at the mesh shape and layout. This is to ensure
@@ -301,6 +307,7 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
   expert_capacity = min(
       t.size,
       int((t.size * hparams.moe_capacity_factor_second_level) / y.size))
+  expert_capacity = max(expert_capacity, 4)
   d = mtf.Dimension("expert_capacity_y", expert_capacity)
 
   # First level of expert routing
@@ -367,13 +374,14 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train):
   expert_inputs_y = mtf.reshape(expert_inputs_y, mtf.Shape(
       [y0, x1, h, d, m]))
 
-  # Now feed the expert inputs through the experts.
   hidden_output = mtf.layers.dense(
       expert_inputs_y, hidden_dim, expert_dims=[y0, x1],
-      activation=mtf.relu, use_bias=False, name="expert0")
+      activation=mtf.relu, use_bias=False, master_dtype=master_dtype,
+      slice_dtype=slice_dtype, name="expert0")
   expert_output = mtf.layers.dense(
       hidden_output, output_dim, expert_dims=[y0, x1],
-      use_bias=False, name="expert1")
+      use_bias=False, master_dtype=master_dtype, slice_dtype=slice_dtype,
+      name="expert1")
 
   # NOW COMBINE EXPERT OUTPUTS (reversing everything we have done)
   # expert_output has shape [y0, x1, h, d, n]
@@ -604,7 +612,6 @@ def _top_2_gating(
 
 def set_default_moe_hparams(hparams):
   """Add necessary hyperparameters for mixture-of-experts."""
-  hparams.feedforward_layer = "moe"
   hparams.moe_num_experts = 16
   hparams.moe_loss_coef = 1e-2
   hparams.add_hparam("moe_gating", "top_2")
diff --git a/tensor2tensor/models/research/moe_experiments.py b/tensor2tensor/models/research/moe_experiments.py
index ae8f9b344..d0c3e7164 100644
--- a/tensor2tensor/models/research/moe_experiments.py
+++ b/tensor2tensor/models/research/moe_experiments.py
@@ -51,20 +51,18 @@ def xmoe_dense_4k():
   Returns:
     a hparams
   """
-  hparams = mtf_transformer.mtf_transformer_base()
+  hparams = mtf_transformer.mtf_transformer_base_lm()
 
   # The following hparams are constant across all these experiments.
-  hparams.label_smoothing = 0.0
   hparams.batch_size = 128
   hparams.d_model = 512
   hparams.d_kv = 128
   hparams.num_heads = 4
-  hparams.num_decoder_layers = 4
+  hparams.decoder_layers = ["att", "drd"] * 4
   hparams.shared_embedding_and_softmax_weights = False
   hparams.learning_rate_schedule = "rsqrt_decay"
 
   # We will vary the following parameters related to the ffn/moe layers.
-  hparams.feedforward_layer = "dense_relu_dense"
   hparams.d_ff = 4096
   hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
   hparams.mesh_shape = "batch:8"
@@ -121,10 +119,10 @@ def mtf_transformer_lm_moe():
     a hparams
   """
   hparams = mtf_transformer.mtf_transformer_lm_baseline()
+  hparams.decoder_layers = ["att", "moe"] * 4
   moe.set_default_moe_hparams(hparams)
   hparams.mesh_shape = "all:8"
   hparams.layout = "batch:all;experts:all"
-  hparams.feedforward_layer = "moe"
   return hparams
 
 
@@ -132,11 +130,28 @@ def mtf_transformer_lm_moe():
 def xmoe_2d():
   """Two-dimensional hierarchical mixture of experts."""
   hparams = xmoe_top_2()
+  hparams.decoder_layers = ["att", "hmoe"] * 4
   hparams.mesh_shape = "b0:2;b1:4"
   hparams.outer_batch_size = 4
   hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
   hparams.moe_num_experts = [4, 4]
-  hparams.feedforward_layer = "hmoe"
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_2d_debug():
+  """For debugging.
+
+  Running this model on TPU without the hack of casting to bfloat16 for
+  alltoall results in nan on the first step.
+  TODO(noam): debug
+
+  Returns:
+    a hparams
+  """
+  hparams = xmoe_2d()
+  hparams.decoder_layers = ["hmoe"] * 1
+  hparams.activation_dtype = "float32"
   return hparams
 
 
@@ -160,180 +175,106 @@ def xmoe_2d_88():
 
 
 @registry.register_hparams
-def xmoe_wiki_base():
+def xmoe_wiki_base(sz):
   """Series of architectural experiments on wikipedia text.
 
   For all of these architectures, we run on languagemodel_wiki_noref_v8k_l1k
   for 3 epochs.  (training set has ~7390100 sequences each of length 1024)
-  1 epoch = 115000 steps at batch_size=64
+  1 epoch = 57500 steps at batch_size=128
 
   Results:
   model             params(M)  einsum  alltoall  mxu-util  log-ppl(1ep) (3ep)
 
   Note: configurations and code are likely to change without notice.
 
+  Args:
+    sz: an integer
+
   Returns:
     a hparams
   """
-  hparams = mtf_transformer.mtf_transformer_base()
+  hparams = mtf_transformer.mtf_transformer_paper_lm(sz)
 
-  # The following hparams are constant across all these experiments.
-  hparams.label_smoothing = 0.0
   hparams.max_length = 1024
-  hparams.batch_size = 64
-  hparams.d_model = 1024
-  hparams.d_kv = 128
-  hparams.num_heads = 8
-  hparams.shared_embedding_and_softmax_weights = False
-  hparams.learning_rate_decay_steps = 115000
-
-  # We will vary the following parameters related to the ffn/moe layers.
-  hparams.feedforward_layer = "dense_relu_dense"
-  hparams.d_ff = 8192
+  hparams.batch_size = 128
+  hparams.learning_rate_decay_steps = 57500
   hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
   hparams.mesh_shape = "batch:32"
   return hparams
 
 
 @registry.register_hparams
-def xmoe_wiki_f64k():
-  """d_ff = 64k.
+def xmoe_wiki_base_0():
+  return xmoe_wiki_base(0)
 
-  Returns:
-    a hparams object.
-  """
-  hparams = xmoe_wiki_base()
-  hparams.moe_hidden_size = 8192
-  hparams.d_ff = 65536
-  hparams.mesh_shape = "model:8;batch:16"
-  return hparams
+
+@registry.register_hparams
+def xmoe_wiki_base_1():
+  return xmoe_wiki_base(1)
 
 
 @registry.register_hparams
-def xmoe_wiki_x64():
-  """Two-dimensional hierarchical mixture of experts.
+def xmoe_wiki_base_2():
+  return xmoe_wiki_base(2)
 
-  (8x8 experts) * (16M params/expert) * 6 layers = 6B params
 
-  Returns:
-    a hparams object.
-  """
-  hparams = xmoe_wiki_base()
-  moe.set_default_moe_hparams(hparams)
-  hparams.feedforward_layer = "hmoe"
-  hparams.moe_hidden_size = 8192
-  hparams.mesh_shape = "b0:4;b1:8"
-  hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
-  hparams.outer_batch_size = 4
-  hparams.moe_num_experts = [8, 8]
-  return hparams
+@registry.register_hparams
+def xmoe_wiki_base_3():
+  return xmoe_wiki_base(3)
 
 
 @registry.register_hparams
-def xmoe_wiki_x32():
-  """Two-dimensional hierarchical mixture of experts.
+def xmoe_wiki_x():
+  """Baseline set of parameters for mixture-of-experts.
 
-  (8x4 experts) * (16M params/expert) * 6 layers = 3B params
+  ~6B parameters
 
   Returns:
-    a hparams object.
+    a hparams
   """
-  hparams = xmoe_wiki_base()
+  hparams = xmoe_wiki_base(0)
   moe.set_default_moe_hparams(hparams)
-  hparams.feedforward_layer = "hmoe"
-  hparams.moe_hidden_size = 8192
+  hparams.decoder_layers = (
+      ["att", "drd", "att", "drd", "att", "hmoe"] * 3 +
+      ["att", "drd", "att", "drd"])
+  hparams.d_ff = 2048
+  hparams.d_kv = 128
+  hparams.moe_hidden_size = 32768
   hparams.mesh_shape = "b0:4;b1:8"
   hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
   hparams.outer_batch_size = 4
   hparams.moe_num_experts = [8, 4]
+  hparams.num_heads = 4
   return hparams
 
 
 @registry.register_hparams
-def xmoe_wiki_x64_h16k():
-  """Mixture of experts."""
-  hparams = xmoe_wiki_x64()
-  hparams.moe_hidden_size = 16384
-  return hparams
-
-
-@registry.register_hparams
-def xmoe_wiki_x64_c15():
-  """Mixture of experts."""
-  hparams = xmoe_wiki_x64()
-  hparams.moe_capacity_factor_train = 1.5
+def xmoe_wiki_x_a32():
+  """Test 32-bit activations."""
+  hparams = xmoe_wiki_x()
+  hparams.activation_dtype = "float32"
   return hparams
 
 
 @registry.register_hparams
-def xmoe_wiki_x256():
-  """Two-dimensional hierarchical mixture of experts.
-
-  (16x16 experts) * (16M params/expert) * 6 layers = 24B params
-
-  Returns:
-    a hparams object.
-  """
-  hparams = xmoe_wiki_x64()
-  hparams.mesh_shape = "b0:8;b1:16"
+def xmoe_wiki_x128():
+  """128 experts, ~25B params on 8x8."""
+  hparams = xmoe_wiki_x()
+  hparams.moe_num_experts = [16, 8]
   hparams.outer_batch_size = 8
-  hparams.moe_num_experts = [16, 16]
-  hparams.batch_size = 256
-  hparams.learning_rate_decay_steps = 28750
-  return hparams
-
-
-@registry.register_hparams
-def xmoe_wiki_x256_h16k():
-  """Two-dimensional hierarchical mixture of experts.
-
-  (16x16 experts) * (32M params/expert) * 6 layers = ~50B params
-
-  Returns:
-    a hparams object.
-  """
-  hparams = xmoe_wiki_x256()
-  hparams.moe_hidden_size = 16384
-  return hparams
-
-
-@registry.register_hparams
-def xmoe_wiki_x1024():
-  """Two-dimensional hierarchical mixture of experts.
-
-  (16x16 experts) * (16M params/expert) * 6 layers = ~100B params
-
-  Returns:
-    a hparams object.
-  """
-  hparams = xmoe_wiki_x64()
-  hparams.mesh_shape = "b0:16;b1:32"
-  hparams.outer_batch_size = 16
-  hparams.moe_num_experts = [32, 32]
-  hparams.batch_size = 4096
-  hparams.learning_rate_decay_steps = 7200
-  return hparams
-
-
-@registry.register_hparams
-def xmoe_wiki_x1024_h16k():
-  """Two-dimensional hierarchical mixture of experts.
-
-  (32x32 experts) * (32M params/expert) * 6 layers = ~200B params
-
-  Returns:
-    a hparams object.
-  """
-  hparams = xmoe_wiki_x1024()
-  hparams.moe_hidden_size = 16384
+  hparams.mesh_shape = "b0:8;b1:16"
+  hparams.batch_size = 512
+  hparams.learning_rate_decay_steps = 14375
   return hparams
 
 
 @registry.register_hparams
-def xmoe_wiki_x256_c15():
-  """Mixture of experts."""
-  hparams = xmoe_wiki_x256()
-  hparams.moe_capacity_factor_train = 1.5
+def xmoe_wiki_x_tiny():
+  """Test on local cpu."""
+  hparams = xmoe_wiki_x()
+  hparams.decoder_layers = (["att", "drd", "hmoe"] * 2 + ["att", "drd"])
+  hparams.moe_hidden_size = 512
+  hparams.batch_size = 16
+  hparams.mesh_shape = ""
+  hparams.activation_dtype = "float32"
   return hparams
-
-
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 1a5c2836f..d4a38663b 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -123,7 +123,7 @@ def fix_and_save_video(vid, prefix):
     save_path = os.path.join(
         output_dir, "%s_prediction_%d.jpg" % (problem_name, prediction_idx))
     show_and_save_image(inputs / 255., save_path)
-  elif inputs_vocab:
+  elif inputs is not None and inputs_vocab:
     if identity_output:
       decoded_inputs = " ".join(map(str, inputs.flatten()))
     else:
@@ -266,9 +266,9 @@ def decode_once(estimator,
 
   for num_predictions, prediction in enumerate(predictions):
     num_predictions += 1
-    inputs = prediction["inputs"]
-    targets = prediction["targets"]
-    outputs = prediction["outputs"]
+    inputs = prediction.get("inputs")
+    targets = prediction.get("targets")
+    outputs = prediction.get("outputs")
 
     # Log predictions
     decoded_outputs = []
diff --git a/tensor2tensor/utils/mtf_model.py b/tensor2tensor/utils/mtf_model.py
index ff9fb4d76..5b034f2f5 100644
--- a/tensor2tensor/utils/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -44,9 +44,9 @@ def estimator_model_fn(cls,
                          mode,
                          config=None,
                          params=None,
-                         decode_hparams=None):
+                         decode_hparams=None,
+                         use_tpu=False):
     hparams = copy.deepcopy(hparams)
-    use_tpu = params and params.get("use_tpu", False)
     hparams.use_tpu = use_tpu
     # merge decode_hparams into hparams if present
     if mode == tf.estimator.ModeKeys.PREDICT and decode_hparams is not None:
@@ -212,7 +212,7 @@ def metric_fn(tf_logits, labels):
           loss=loss)
 
   def estimator_spec_predict(self, features, mesh, mesh_impl, use_tpu):
-    mtf_samples = self.sample(features, mesh)
+    mtf_samples = mtf.anonymize(self.sample(features, mesh))
     lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
     outputs = lowering.export_to_tf_tensor(mtf_samples)
     if self.has_input:

From 33583afe60c0a095c17a57eca3e39becaa2d4046 Mon Sep 17 00:00:00 2001
From: Stefan Falk <43335432+stefan-falk@users.noreply.github.com>
Date: Wed, 31 Oct 2018 22:25:39 +0100
Subject: [PATCH 1133/2720]  Fixing error passing wrong data_dir (#1185)

* Pass data_dir to feature_encoders

Pass data_dir to feature_encoders

* Fixing error passing wrong data_dir
---
 tensor2tensor/data_generators/common_voice.py | 4 ++--
 tensor2tensor/data_generators/librispeech.py  | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/common_voice.py b/tensor2tensor/data_generators/common_voice.py
index 41935dc40..1bc8fc126 100644
--- a/tensor2tensor/data_generators/common_voice.py
+++ b/tensor2tensor/data_generators/common_voice.py
@@ -132,8 +132,8 @@ def generator(self,
       ]
       corpus_tar.extractall(tmp_dir, members=members)
 
-    data_dir = os.path.join(tmp_dir, "cv_corpus_v1")
-    data_tuples = _collect_data(data_dir)
+    raw_data_dir = os.path.join(tmp_dir, "cv_corpus_v1")
+    data_tuples = _collect_data(raw_data_dir)
     encoders = self.feature_encoders(data_dir)
     audio_encoder = encoders["waveforms"]
     text_encoder = encoders["targets"]
diff --git a/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index a6b7cdabc..c19191513 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -134,8 +134,8 @@ def generator(self, data_dir, tmp_dir, datasets,
             members.append(f)
         corpus_tar.extractall(tmp_dir, members=members)
 
-      data_dir = os.path.join(tmp_dir, "LibriSpeech", subdir)
-      data_files = _collect_data(data_dir, "flac", "txt")
+      raw_data_dir = os.path.join(tmp_dir, "LibriSpeech", subdir)
+      data_files = _collect_data(raw_data_dir, "flac", "txt")
       data_pairs = data_files.values()
 
       encoders = self.feature_encoders(data_dir)

From 061488b8305c32d3f196e7425b8a58a8214466f8 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 31 Oct 2018 14:44:08 -0700
Subject: [PATCH 1134/2720] fixing internal_loss for recurrent models with L2
 loss.

PiperOrigin-RevId: 219540943
---
 tensor2tensor/models/video/base.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 8c309384e..eb6d7cdad 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -445,6 +445,7 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
     # using the default (a bit strange) video modality - we should change that.
 
     hparams = self.hparams
+    all_frames_copy = [tf.identity(frame) for frame in all_frames]
     orig_frame_shape = common_layers.shape_list(all_frames[0])
     batch_size = orig_frame_shape[0]
     ss_func = self.get_scheduled_sample_func(batch_size)
@@ -506,7 +507,7 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
       has_input_predictions = hparams.video_num_input_frames > 1
       if self.is_training and hparams.internal_loss and has_input_predictions:
         # add the loss for input frames as well.
-        extra_gts = all_frames[1:hparams.video_num_input_frames]
+        extra_gts = all_frames_copy[1:hparams.video_num_input_frames]
         extra_raw_gts = all_raw_frames[1:hparams.video_num_input_frames]
         extra_pds = res_frames[:hparams.video_num_input_frames-1]
         recon_loss = self.get_extra_internal_loss(

From 36596aa92b2da1be6a7ebcbcc06e34c326ac1604 Mon Sep 17 00:00:00 2001
From: Giovanni Campagna <scampa.giovanni@gmail.com>
Date: Wed, 31 Oct 2018 14:59:20 -0700
Subject: [PATCH 1135/2720] trainer_lib: add problem hparams after parsing the
 overrides (#1053)

t2t-trainer passes problem_name=None to create_hparams(), and
adds the problem hparams afterwards when creating the experiment.
t2t-decoder instead relies fully on create_hparams()
This inconsistency causes t2t-trainer and t2t-decoder to pass
different sets of model_hparams to Problem.hparams() (before and
after overrides respectively), which is problematic.

Always letting the problem see the final version of the model
hparams is more useful (in case the problem hparams depend on
the model), so make sure the problem is added after the overrides.
---
 tensor2tensor/utils/trainer_lib.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 76227545f..38ab35370 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -96,12 +96,12 @@ def create_hparams(hparams_set,
     hparams = _create_hparams_from_json(hparams_path, hparams)
   if data_dir:
     hparams.add_hparam("data_dir", data_dir)
-  if problem_name:
-    add_problem_hparams(hparams, problem_name)
   if hparams_overrides_str:
     tf.logging.info("Overriding hparams in %s with %s", hparams_set,
                     hparams_overrides_str)
     hparams = hparams.parse(hparams_overrides_str)
+  if problem_name:
+    add_problem_hparams(hparams, problem_name)
   return hparams
 
 

From dbe3949dc71596a08faaf0a8e81a65eed9656d4d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 31 Oct 2018 15:57:08 -0700
Subject: [PATCH 1136/2720] Add WikiRevision data-generation problem; data
 extraction/preprocessing for processing wiki history dumps for GEC.

PiperOrigin-RevId: 219552651
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 .../data_generators/wiki_revision.py          | 499 +++++++++++++
 .../data_generators/wiki_revision_utils.py    | 686 ++++++++++++++++++
 3 files changed, 1186 insertions(+)
 create mode 100644 tensor2tensor/data_generators/wiki_revision.py
 create mode 100644 tensor2tensor/data_generators/wiki_revision_utils.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 45db322a2..b96c7daac 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -84,6 +84,7 @@
     "tensor2tensor.data_generators.vqa",
     "tensor2tensor.data_generators.wiki",
     "tensor2tensor.data_generators.wiki_lm",
+    "tensor2tensor.data_generators.wiki_revision",
     "tensor2tensor.data_generators.wikisum.wikisum",
     "tensor2tensor.data_generators.wikitext103",
     "tensor2tensor.data_generators.wsj_parsing",
diff --git a/tensor2tensor/data_generators/wiki_revision.py b/tensor2tensor/data_generators/wiki_revision.py
new file mode 100644
index 000000000..18b1ad5dd
--- /dev/null
+++ b/tensor2tensor/data_generators/wiki_revision.py
@@ -0,0 +1,499 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+r"""Data extraction/preprocessing for processing wiki history dumps for GEC.
+
+We use a set of heuristics to distill prose from the wikipedia xml. We produce
+source-target pairs of text reflecting wikipedia edits.
+
+WikiRevision problem - fragment of older revision -> fragment of newer revision.
+
+This implements data extraction from wikipedia as desribed in the arXiv paper,
+Weakly Supervised Grammatical Error Correction using Iterative Decoding.
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import random
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators import wiki_revision_utils
+from tensor2tensor.utils import metrics
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+flags.DEFINE_integer("wiki_revision_num_train_shards", 50,
+                     "Set the number of training shards to be output.")
+flags.DEFINE_integer("wiki_revision_num_dev_shards", 1,
+                     "Set the number of dev shards to be output.")
+
+flags.DEFINE_string(
+    "wiki_revision_data_prefix", "",
+    "Specify the prefix for input data. Expects 7z compressed Wikipedia XML "
+    "files, available at https://dumps.wikimedia.org/enwiki/latest/.")
+flags.DEFINE_string(
+    "wiki_revision_vocab_file", "",
+    "Specify a wordpieces vocabulary with which to encode the text. Will "
+    "generate one from data if not specified.")
+
+flags.DEFINE_integer(
+    "wiki_revision_max_examples_per_shard", 0,
+    "Use this to set a cap on examples per shard. "
+    "0 is no cap.")
+
+# Data filtration heuristics:
+flags.DEFINE_integer("wiki_revision_max_page_size_exp", 26,
+                     "Exponent for 2**X byte cap on page size.")
+flags.DEFINE_float(
+    "wiki_revision_max_equal_to_diff_ratio", 0,
+    "Max ratio between count of equal, diff chars for generated "
+    "examples. Ratio of 1 means examples with more diff chars "
+    "than equal chars will be tossed out.")
+flags.DEFINE_float(
+    "wiki_revision_revision_skip_factor", 1.5,
+    "If >1, process only logarithmically many revisions. "
+    "This avoids blowup in runtime due to many-revision pages. "
+    "See wiki_revision_utils.include_revision for details.")
+flags.DEFINE_float("wiki_revision_percent_identical_examples", 0.04,
+                   "Percent of generated examples for which source == target.")
+flags.DEFINE_bool(
+    "wiki_revision_introduce_errors", True, "Add errors to the data."
+    "See wiki_revision_utils.introduce_errors for details.")
+
+
+@registry.register_problem
+class WikiRevision(text_problems.Text2TextProblem):
+  """Old segment -> revised segment.
+
+  Data filtration heuristics:
+    wiki_revision_max_page_size_exp:
+      pages above this # of bytes are thrown out
+
+    wiki_revision_revision_skip_factor:
+      rate of logarithmic downsampling of revision history list
+
+    wiki_revision_percent_identical_examples:
+      how many identitcal examples to admit, as percent of total examples
+
+    wiki_revision_introduce_errors:
+      whether or not to introduce spelling-type errors on the source side
+
+    wiki_revision_max_equal_to_diff_ratio:
+      whether or not to introduce spelling-type errors on the source side
+
+
+  Vocab size=32k
+  Maximum input/target length = 1024 wordpiece tokens
+  """
+  num_identity_examples = 0
+  num_total_examples = 0
+  num_identity_examples = 0
+  num_pages = 0
+  num_revisions_total = 0
+  num_revisions_admitted = 0
+  num_examples_thrown_out_identity = 0
+  num_examples_thrown_out_too_long = 0
+  num_examples_thrown_out_edit_distance = 0
+  num_examples_with_introduced_error = 0
+  num_introduced_errors = 0
+  num_source_tokens = 0
+  num_target_tokens = 0
+  corpus_files = None
+
+  @property
+  def approx_vocab_size(self):
+    return 2**15  # 32K
+
+  @property
+  def strip(self):
+    """Whether to strip wikipedia-stuff to get plain text."""
+    return True
+
+  @property
+  def wiki_revision_skip_factor(self):
+    """If this value is >1.0, process only logarithmically many revisions."""
+    return FLAGS.wiki_revision_revision_skip_factor
+
+  @property
+  def max_segment_length(self):
+    """Maximum number of input/target wordpiece tokens."""
+    return 256
+
+  @property
+  def max_examples_per_shard(self):
+    """Maximum number of examples to generate per shard.  0=unlimited."""
+    return FLAGS.wiki_revision_max_examples_per_shard
+
+  def aggregate_job_stats(self):
+    # Aggregate job stats for output.
+    stat = []
+    # Run stats.
+    stat.append("Flags for job:\n"
+                "Dev shards: {}\n"
+                "Train shards: {}\n"
+                "Revision skip factor: {}\n"
+                "Max page size: 2**{}\n"
+                "Introduce errors: {}\n"
+                "Max edit ratio: {}\n"
+                "Percent Identical Examples: {}\n"
+                "".format(FLAGS.wiki_revision_num_dev_shards,
+                          FLAGS.wiki_revision_num_train_shards,
+                          FLAGS.wiki_revision_revision_skip_factor,
+                          FLAGS.wiki_revision_max_page_size_exp,
+                          FLAGS.wiki_revision_introduce_errors,
+                          FLAGS.wiki_revision_max_equal_to_diff_ratio,
+                          FLAGS.wiki_revision_percent_identical_examples))
+
+    # File stats.
+    stat.append("corpus files: {}\n"
+                "\tnames: {}\n"
+                "\tpages per input file: {:.1f}\n"
+                "".format(
+                    len(self.corpus_files), self.corpus_files,
+                    (0 if not self.corpus_files else
+                     self.num_pages / len(self.corpus_files))))
+    # Page stats.
+    stat.append(
+        "pages processed: {}\n"
+        "\trevisions per page: {:.2f}, total: {}\n"
+        "\trevisions admitted per page: {:.2f}, percent of total: {:.2f}\n"
+        "".format(
+            self.num_pages, (0 if not self.num_pages else
+                             self.num_revisions_total / self.num_pages),
+            self.num_revisions_total,
+            (0 if not self.num_pages else
+             self.num_revisions_admitted / self.num_pages),
+            (0 if not self.num_revisions_total else
+             100 * self.num_revisions_admitted / self.num_revisions_total)))
+    # Revision stats.
+    stat.append(
+        "revisions admitted: {}\n"
+        "\texamples generated per revision: {:.2f}\n"
+        "".format(self.num_revisions_admitted,
+                  (0 if not self.num_revisions_admitted else
+                   self.num_total_examples / self.num_revisions_admitted)))
+    # Example stats.
+    stat.append(
+        "examples generated: {}\n"
+        "\twith error introduced: {}, percent of total: {:.2f}\n"
+        "\ttotal errors introduced: {}, errors per errorred example: {:.2f}\n"
+        "\texamples thrown out: {}\n"
+        "\t\ttoo long: {}\n"
+        "\t\tidentity: {}\n"
+        "\t\tedit distance: {}\n"
+        "\tremaining identity examples: {}\n"
+        "\tratio identity (actual, desired): {:.3f}, {}\n"
+        "".format(
+            self.num_total_examples, self.num_examples_with_introduced_error,
+            (0 if not self.num_total_examples else 100 *
+             self.num_examples_with_introduced_error / self.num_total_examples),
+            self.num_introduced_errors,
+            (0 if not self.num_examples_with_introduced_error else
+             self.num_introduced_errors /
+             self.num_examples_with_introduced_error),
+            self.num_examples_thrown_out_too_long +
+            self.num_examples_thrown_out_identity +
+            self.num_examples_thrown_out_edit_distance,
+            self.num_examples_thrown_out_too_long,
+            self.num_examples_thrown_out_identity,
+            self.num_examples_thrown_out_edit_distance,
+            self.num_identity_examples,
+            (0 if not self.num_total_examples else
+             self.num_identity_examples / self.num_total_examples),
+            FLAGS.wiki_revision_percent_identical_examples))
+    # Token stats.
+    stat.append("tokens generated: {}\n"
+                "\tsource: {}\n"
+                "\ttarget: {}\n"
+                "\tper example: {:.2f}\n"
+                "\t\tsource: {:.2f}\n"
+                "\t\ttarget: {:.2f}\n"
+                "".format(self.num_source_tokens + self.num_target_tokens,
+                          self.num_source_tokens, self.num_target_tokens,
+                          (0 if not self.num_total_examples else
+                           (self.num_source_tokens + self.num_target_tokens) /
+                           self.num_total_examples),
+                          (0 if not self.num_total_examples else
+                           self.num_source_tokens / self.num_total_examples),
+                          (0 if not self.num_total_examples else
+                           self.num_target_tokens / self.num_total_examples)))
+    return "\n".join(stat)
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+
+    if task_id == -1 or task_id is None:
+      for i in range(FLAGS.wiki_revision_num_train_shards +
+                     FLAGS.wiki_revision_num_dev_shards):
+        self.generate_data(data_dir, tmp_dir, i)
+        return
+
+    tf.logging.info(
+        "Flags for job (task_id {}): "
+        "Dev shards: {}, Train shards: {}, "
+        "Revision skip factor: {}, Max page size: 2**{}, Introduce errors: {},"
+        "Percent Identical Examples: {}"
+        "".format(task_id, FLAGS.wiki_revision_num_dev_shards,
+                  FLAGS.wiki_revision_num_train_shards,
+                  FLAGS.wiki_revision_revision_skip_factor,
+                  FLAGS.wiki_revision_max_page_size_exp,
+                  FLAGS.wiki_revision_introduce_errors,
+                  FLAGS.wiki_revision_percent_identical_examples))
+
+    if FLAGS.wiki_revision_vocab_file:
+      encoder = wiki_revision_utils.get_encoder_from_vocab(
+          FLAGS.wiki_revision_vocab_file)
+    else:
+      encoder = wiki_revision_utils.get_or_generate_vocabulary(
+          data_dir, tmp_dir, FLAGS.wiki_revision_data_prefix,
+          FLAGS.wiki_revision_max_page_size_exp, self.approx_vocab_size,
+          self.strip)
+
+    random.seed(123)
+    if task_id < FLAGS.wiki_revision_num_train_shards:
+      out_file = self.training_filepaths(
+          data_dir, FLAGS.wiki_revision_num_train_shards,
+          shuffled=False)[task_id]
+    else:
+      out_file = self.dev_filepaths(
+          data_dir, FLAGS.wiki_revision_num_dev_shards,
+          shuffled=False)[task_id - FLAGS.wiki_revision_num_train_shards]
+
+    tf.logging.info("Generating files for path: %s", out_file)
+    self.corpus_files = wiki_revision_utils.corpus_files_for_shard(
+        task_id, FLAGS.wiki_revision_num_train_shards,
+        FLAGS.wiki_revision_num_dev_shards, FLAGS.wiki_revision_data_prefix)
+    example_generator = self.generator(encoder, self.corpus_files, tmp_dir)
+
+    packed_example_generator = self._maybe_pack_examples(example_generator)
+    generator_utils.generate_files(packed_example_generator, [out_file])
+    generator_utils.shuffle_dataset([out_file])
+
+    tf.logging.info(
+        "Job stats: identity examples: {}, total examples {}, ratio: {}".format(
+            self.num_identity_examples, self.num_total_examples,
+            (1 + self.num_identity_examples) / (1 + self.num_total_examples)))
+
+    job_stats_string = self.aggregate_job_stats()
+    out_dir, filename = out_file.replace("-unshuffled", "").rsplit("/", 1)
+    stats_prefix = "/stats_"
+    stats_file_path = "".join([out_dir, stats_prefix, filename])
+    if tf.gfile.Exists(
+        stats_file_path) and tf.gfile.Open(stats_file_path).size() != 0:
+      tf.logging.info("Skipping writing stats because output file exists.")
+    else:
+      with tf.gfile.Open(stats_file_path, "w") as out:
+        tf.logging.info("Writing job stats to {}".format(stats_file_path))
+        out.write(job_stats_string)
+
+    tf.logging.info(job_stats_string)
+
+  def generator(self, encoder, corpus_files, tmp_dir):
+    for page in wiki_revision_utils.corpus_page_generator(
+        corpus_files, tmp_dir, FLAGS.wiki_revision_max_page_size_exp):
+      self.num_pages += 1
+      examples = self.page_to_examples(page, encoder)
+      for x in examples:
+        yield x
+      if self.num_total_examples % 100000 == 0:
+        tf.logging.info(
+            u"page count={} num_total_examples={} id={} title={}".format(
+                self.num_pages, self.num_total_examples, page["id"],
+                page["title"]))
+      if (self.max_examples_per_shard and
+          self.num_total_examples >= self.max_examples_per_shard):
+        tf.logging.info(
+            "Examples per shard {} >= max_examples_per_shard {}. Shutting down."
+            .format(self.num_total_examples, self.max_examples_per_shard))
+        break
+    tf.logging.info(
+        "Total pages: {}, total examples: {}, examples per page: {}".format(
+            self.num_pages, self.num_total_examples, 0 if not self.num_pages
+            else self.num_total_examples / self.num_pages))
+
+  def page_to_examples(self, page, encoder):
+    revisions = page["revisions"]
+    self.num_revisions_total += len(revisions)
+    if len(revisions) < 2:
+      return []
+    revisions = [
+        wiki_revision_utils.get_text(r)
+        for n, r in enumerate(revisions)
+        if wiki_revision_utils.include_revision(
+            n, self.wiki_revision_skip_factor) or n + 1 == len(revisions)
+    ]
+    self.num_revisions_admitted += len(revisions)
+
+    ret = []
+    for i in range(len(revisions) - 1):
+      old_revision = revisions[i]
+      new_revision = revisions[i + 1]
+
+      if FLAGS.wiki_revision_introduce_errors:
+        old_revision_text, num_added_err = wiki_revision_utils.introduce_errors(
+            revisions[i])
+        if num_added_err:
+          self.num_introduced_errors += num_added_err
+          self.num_examples_with_introduced_error += 1
+      else:
+        old_revision_text = revisions[i]
+      new_revision_text = revisions[i + 1]
+      if encoder:
+        # Encode text into list of ids, if a text encoder is present.
+        old_revision = encoder.encode(old_revision_text)
+        new_revision = encoder.encode(new_revision_text)
+      else:
+        # Retain text (as list of characters), if a text encoder is not present.
+        old_revision = old_revision_text
+        new_revision = new_revision_text
+      ret.extend(
+          self.make_examples(
+              encoder,
+              old_revision,
+              new_revision,
+              max_length=self.max_segment_length,
+              percent_identical_examples=FLAGS
+              .wiki_revision_percent_identical_examples))
+    return ret
+
+  def make_examples(self,
+                    encoder,
+                    old_snapshot,
+                    new_snapshot,
+                    max_length=1024,
+                    percent_identical_examples=0.01,
+                    max_length_distance=0):
+    """Produce training examples based on a pair of snapshots.
+
+    Aligns the snapshots, then chops at a random subset of the alignment points
+    to create (old snippet -> new snippet) examples.
+
+    Most negative examples (those with no changes) are discarded, but we
+    keep some of them, maintaining a proportion in the final data
+    determined by percent_identical_examples.
+
+    Args:
+      encoder: the subword text encoder
+      old_snapshot: a list of ids
+      new_snapshot: a list of ids
+      max_length: an integer.  Maximum length of "inputs" and "targets".
+      percent_identical_examples: a float
+      max_length_distance: an integer. Max token edit dist for admitted examples
+
+    Returns:
+      a list of feature dictionaries.  The dictionaries have
+      "inputs" and "targets" populated. text_encoder.EOS is appended to both.
+    """
+    ret = []
+    eos_sequence = [text_encoder.EOS_ID]
+    # Pick a per-token cut probability with a log-uniform distribution between
+    # 1/4 and 1/(max_length / 2)
+    bound1 = -math.log(4.0)
+    bound2 = -math.log(max_length / 2.0)
+    cut_prob = math.exp(random.random() * (bound2 - bound1) + bound1)
+    opcodes = wiki_revision_utils.fast_match_sequences(old_snapshot,
+                                                       new_snapshot)
+    cut_points = [(0, 0)]
+    for tag, i1, i2, j1, j2 in opcodes:
+      if tag == "equal":
+        for i in range(i1, i2 + 1):
+          if random.random() < cut_prob:
+            cut_points.append((i, i + j1 - i1))
+    cut_points.append((len(old_snapshot), len(new_snapshot)))
+    src_tgt_pairs = []
+    for cut_number in range(len(cut_points) - 1):
+      i1, j1 = cut_points[cut_number]
+      i2, j2 = cut_points[cut_number + 1]
+      old_segment = old_snapshot[i1:i2]
+      new_segment = new_snapshot[j1:j2]
+      src_tgt_pairs.append((old_segment, new_segment))
+
+    src_tgt_pairs, thrown_edit_count = wiki_revision_utils.edit_distance_filter(
+        wiki_revision_utils.throw_empty_pairs(src_tgt_pairs),
+        FLAGS.wiki_revision_max_equal_to_diff_ratio)
+
+    self.num_examples_thrown_out_edit_distance += thrown_edit_count
+
+    for source, target in src_tgt_pairs:
+      # Add EOS segment.
+      old_segment = source + eos_sequence
+      new_segment = target + eos_sequence
+      if len(old_segment) <= max_length and len(new_segment) <= max_length:
+        if max_length_distance and (abs(len(old_segment) - len(new_segment)) >
+                                    max_length_distance):
+          self.num_examples_thrown_out_edit_distance += 1
+          continue
+        if old_segment == new_segment:
+          # If current proportion of identity is below target
+          # percent_identical_examples, then roll for a 50% chance to add an
+          # identitical example. Random roll preserves nondeterminism.
+          # percent_identical_examples, then add identitical example.
+          # Random roll preserves nondeterminism in selecting identity examples.
+          if (((self.num_identity_examples) / (1 + self.num_total_examples)) >
+              percent_identical_examples) or random.random() > 0.5:
+            self.num_examples_thrown_out_identity += 1
+            continue
+          else:
+            self.num_identity_examples += 1
+        self.num_total_examples += 1
+        self.num_source_tokens += len(old_segment) - 1
+        self.num_target_tokens += len(new_segment) - 1
+        ret.append({"inputs": old_segment, "targets": new_segment})
+      else:
+        self.num_examples_thrown_out_too_long += 1
+
+    return ret
+
+  def eval_metrics(self):
+    return [
+        metrics.Metrics.ACC,
+        metrics.Metrics.ACC_TOP5,
+        metrics.Metrics.ACC_PER_SEQ,
+        metrics.Metrics.NEG_LOG_PERPLEXITY,
+    ]
+
+  @property
+  def invert_prob(self):
+    """Ratio of e^2 positive forward to backward examples."""
+    return 1.0 / (1.0 + math.exp(2.0))
+
+
+@registry.register_problem
+class WikiRevisionPacked1k(WikiRevision):
+  """Packed version for TPU."""
+
+  @property
+  def packed_length(self):
+    return 1024
+
+
+@registry.register_problem
+class WikiRevisionPacked256(WikiRevision):
+  """Packed version for TPU."""
+
+  @property
+  def packed_length(self):
+    return 256
+
+  @property
+  def max_segment_length(self):
+    return 256
diff --git a/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
new file mode 100644
index 000000000..574fcb67a
--- /dev/null
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -0,0 +1,686 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Utilties for data generation for Wikipedia Revisions/GEC problems.
+
+Wikipedia revisions problems are defined in wiki_history.py
+Wikipedia GEC problems are defined in wiki_gec.py
+"""
+
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import os
+import random
+import re
+import subprocess
+
+import six
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import text_encoder
+
+import tensorflow as tf
+
+
+def to_unicode(s):
+  return unicode(s, "utf-8") if six.PY2 else s.decode("utf-8")
+
+
+def include_revision(revision_num, skip_factor=1.1):
+  """Decide whether to include a revision.
+
+  If the number of revisions is large, we exclude some revisions to avoid
+  a quadratic blowup in runtime, since the article is likely also large.
+
+  We make the ratio between consecutive included revision numbers
+  appproximately equal to "factor".
+
+  Args:
+    revision_num: an integer
+    skip_factor: a floating point number >= 1.0
+
+  Returns:
+    a boolean
+  """
+  if skip_factor <= 1.0:
+    return True
+  return (int(math.log(revision_num + 1.0) / math.log(skip_factor)) != int(
+      math.log(revision_num + 2.0) / math.log(skip_factor)))
+
+
+def file_page_generator(my_file, max_page_size=2**28):
+  """Read wikipedia pages from a history dump.
+
+  Since some pages can be terabytes in size (with all the revisions),
+  we limit page size to max_page_size bytes.
+
+  Args:
+    my_file: an open file object.
+    max_page_size: an integer
+
+  Yields:
+    strings
+  """
+  page_start = "  <page>\n"
+  page_end = "  </page>\n"
+  chunk_size = max_page_size
+  page_start = "  <page>\n"
+  page_end = "  </page>\n"
+  leftovers = ""
+  while True:
+    chunk = my_file.read(chunk_size)
+    if not chunk:
+      break
+    chunk = leftovers + chunk
+    current_pos = 0
+    while True:
+      start_pos = chunk.find(page_start, current_pos)
+      if start_pos == -1:
+        break
+      end_pos = chunk.find(page_end, start_pos)
+      if end_pos == -1:
+        if len(chunk) - start_pos > max_page_size:
+          leftovers = ""
+        else:
+          leftovers = chunk[start_pos:]
+        break
+      raw_page = chunk[start_pos + len(page_start):end_pos]
+      if len(raw_page) < max_page_size:
+        ret = parse_page(raw_page)
+        if ret:
+          yield ret
+      current_pos = end_pos + len(page_end)
+
+
+def get_title(page):
+  """Extract the title from a page.
+
+  Args:
+    page: a string
+  Returns:
+    a string
+  """
+  start_pos = page.find("<title>")
+  end_pos = page.find("</title>")
+  assert start_pos != -1
+  assert end_pos != -1
+  start_pos += len("<title>")
+  return to_unicode(page[start_pos:end_pos])
+
+
+def get_id(page):
+  """Extract the id from a page.
+
+  Args:
+    page: a string
+  Returns:
+    an integer
+  """
+  start_pos = page.find("<id>")
+  end_pos = page.find("</id>")
+  assert start_pos != -1
+  assert end_pos != -1
+  start_pos += len("<id>")
+  return int(page[start_pos:end_pos])
+
+
+def get_revisions(page):
+  """Extract the revisions of a page.
+
+  Args:
+    page: a string
+  Returns:
+    a list of strings
+  """
+  start_string = "    <revision>\n"
+  end_string = "    </revision>\n"
+  ret = []
+  current_pos = 0
+  while True:
+    start_pos = page.find(start_string, current_pos)
+    if start_pos == -1:
+      break
+    end_pos = page.find(end_string, start_pos)
+    assert end_pos != -1
+    ret.append(page[start_pos + len(start_string):end_pos])
+    current_pos = end_pos + len(end_string)
+  return ret
+
+
+def parse_page(raw_page):
+  """Create a dictionary with title, id, and list of revisions.
+
+  The dictionary contains:
+  "title": a string
+  "id": an integer
+  "revisions": a list of strings
+
+  Args:
+    raw_page: a string
+
+  Returns:
+    a dictionary, or None in the case of an error.
+  """
+  ret = {"title": get_title(raw_page), "id": get_id(raw_page)}
+  if ":" in ret["title"]:
+    return None
+  ret["revisions"] = get_revisions(raw_page)
+  return ret
+
+
+def maybe_copy_file_to_directory(source_filepath, target_directory):
+  """Copy a file to a directory if it is not already there.
+
+  Returns the target filepath.
+
+  Args:
+    source_filepath: a string
+    target_directory: a string
+
+  Returns:
+    a string
+  """
+  if not tf.gfile.Exists(target_directory):
+    tf.logging.info("Creating directory %s" % target_directory)
+    os.mkdir(target_directory)
+  target_filepath = os.path.join(target_directory,
+                                 os.path.basename(source_filepath))
+  if not tf.gfile.Exists(target_filepath):
+    tf.logging.info("Copying %s to %s" % (source_filepath, target_filepath))
+    tf.gfile.Copy(source_filepath, target_filepath)
+    statinfo = os.stat(target_filepath)
+    tf.logging.info("Successfully copied %s, %s bytes." % (target_filepath,
+                                                           statinfo.st_size))
+  else:
+    tf.logging.info("Not copying, file already found: %s" % target_filepath)
+  return target_filepath
+
+
+def corpus_page_generator(corpus_files, tmp_dir, max_page_size_exp):
+  """Generate pages from a list of .7z encoded history dumps.
+
+  Args:
+    corpus_files: a list of strings
+    tmp_dir: a string
+    max_page_size_exp: an integer
+
+  Yields:
+    strings
+  """
+  for remote_filepath in corpus_files:
+
+    filepath = maybe_copy_file_to_directory(remote_filepath, tmp_dir)
+    tf.logging.info("Reading from " + filepath)
+
+    command = ["7z", "x", "-so", filepath]
+    tf.logging.info("Running command: %s", command)
+
+    p = subprocess.Popen(command, stdout=subprocess.PIPE, bufsize=-1)
+
+    for page in file_page_generator(p.stdout, 2**max_page_size_exp):
+      yield page
+
+
+def get_text(revision, strip=True):
+  """Extract the text from a revision.
+
+  Args:
+    revision: a string
+    strip: a boolean
+
+  Returns:
+    a string
+  """
+  # text start tag looks like "<text ..otherstuff>"
+  start_pos = revision.find("<text")
+  assert start_pos != -1
+  end_tag_pos = revision.find(">", start_pos)
+  assert end_tag_pos != -1
+  end_tag_pos += len(">")
+  end_pos = revision.find("</text>")
+  if end_pos == -1:
+    ret = ""
+  else:
+    ret = revision[end_tag_pos:end_pos]
+  if strip:
+    ret = strip_text(ret)
+  ret = to_unicode(ret)
+  return ret
+
+
+def strip_text(text):
+  """Strip wikipedia-stuff out of text, making it mostly prose.
+
+  The reason for this is to learn a model that is good at editing prose.
+
+  Args:
+    text: a string
+
+  Returns:
+    a string
+  """
+  return _remove_boring_lines(
+      _remove_triple_quotes(
+          _remove_double_brackets(
+              _remove_references(_remove_curly_braces(text)))))
+
+
+def _find_and_replace(text, start_string, end_string, replace_fn):
+  """Remove everything found between instances of start_string and end_string.
+
+  Replace each such instance with replace_fn(removed_text)
+
+  e.g. _find_and_replace("the [[fat]] cat [[sat]]", "[[", "]]", lambda x: x)
+    = "the fat cat sat"
+
+  Args:
+    text: a string
+    start_string: a string
+    end_string: a string
+    replace_fn: a unary function from string to string
+
+  Returns:
+    a string
+  """
+  ret = ""
+  current_pos = 0
+  while True:
+    start_pos = text.find(start_string, current_pos)
+    if start_pos == -1:
+      ret += text[current_pos:]
+      break
+    ret += text[current_pos:start_pos]
+    end_pos = text.find(end_string, start_pos + len(start_string))
+    if end_pos == -1:
+      break
+    ret += replace_fn(text[start_pos + len(start_string):end_pos])
+    current_pos = end_pos + len(end_string)
+  return ret
+
+
+def _remove_references(text):
+  return _find_and_replace(text, "&lt;ref", "&lt;/ref&gt;", lambda s: "")
+
+
+def _remove_triple_quotes(text):
+  return _find_and_replace(text, "'''", "'''", lambda s: s)
+
+
+def _remove_curly_braces(text):
+  """Remove everything in curly braces.
+
+  Curly braces may be nested, so we keep track of depth.
+
+  Args:
+    text: a string
+  Returns:
+    a string
+  """
+  current_pos = 0
+  depth = 0
+  ret = ""
+  for match in re.finditer("[{}]", text):
+    if depth == 0:
+      ret += text[current_pos:match.start()]
+    depth += 1 if text[match.start()] == "{" else -1
+    current_pos = match.end()
+  if depth != 0:
+    # Many articles have mismatched braces, but it still seems better to remove
+    # them than not.
+    pass
+  else:
+    ret += text[current_pos:]
+  return ret
+
+
+def _remove_double_brackets(text):
+  """Remove double brackets, but leave the viewable text.
+
+  Args:
+    text: a string
+  Returns:
+    a string
+  """
+
+  def replacement_fn(s):
+    if ":" in s:
+      # this is probably a category or something like that.
+      return ""
+    # keep the part after the bar.
+    bar_pos = s.find("|")
+    if bar_pos == -1:
+      return s
+    return s[bar_pos + 1:]
+
+  return _find_and_replace(text, "[[", "]]", replacement_fn)
+
+
+def _remove_boring_lines(text):
+  """Remove lines that do not start with a letter or a quote.
+
+  From inspecting the data, this seems to leave in most prose and remove
+  most weird stuff.
+
+  Args:
+    text: a string
+  Returns:
+    a string
+  """
+  lines = text.split("\n")
+  filtered = [line for line in lines if re.match("[a-zA-z\"\']", line)]
+  return "\n".join(filtered)
+
+
+def all_corpus_files(data_prefix):
+  return sorted(tf.gfile.Glob(data_prefix + "*"))
+
+
+def corpus_files_for_shard(shard_num, train_shards, dev_shards, data_prefix):
+  corpus_files = [
+      filename for i, filename in enumerate(all_corpus_files(data_prefix))
+      if i % (train_shards + dev_shards) == shard_num
+  ]
+  tf.logging.info("Corpus files for shard %s: %s", shard_num, corpus_files)
+
+  assert shard_num < (train_shards + dev_shards)
+  return corpus_files
+
+
+def vocab_filename(approx_vocab_size, strip):
+  return "vocab.wiki_revision%s.%d" % (".strip" if strip else "",
+                                       approx_vocab_size)
+
+
+def get_or_generate_vocabulary(data_dir,
+                               tmp_dir,
+                               data_prefix,
+                               max_page_size_exp,
+                               approx_vocab_size=32768,
+                               strip=True):
+  """Get or generate the vocabulary.
+
+  Args:
+    data_dir: a string
+    tmp_dir: a string
+    data_prefix: a string
+    max_page_size_exp: an integer
+    approx_vocab_size: an integer
+    strip: a boolean
+
+  Returns:
+    a TextEncoder
+  """
+  num_pages_for_vocab_generation = approx_vocab_size // 3
+  vocab_file = vocab_filename(approx_vocab_size, strip)
+
+  def my_generator(data_prefix):
+    """Line generator for vocab."""
+    count = 0
+    for page in corpus_page_generator(
+        all_corpus_files(data_prefix)[::-1], tmp_dir, max_page_size_exp):
+      revisions = page["revisions"]
+      if revisions:
+        text = get_text(revisions[-1], strip=strip)
+        yield text
+        count += 1
+        if count % 100 == 0:
+          tf.logging.info("reading pages for vocab %d" % count)
+        if count > num_pages_for_vocab_generation:
+          break
+
+  return generator_utils.get_or_generate_vocab_inner(data_dir, vocab_file,
+                                                     approx_vocab_size,
+                                                     my_generator(data_prefix))
+
+
+def get_encoder_from_vocab(vocab_filepath):
+  """Get encoder from vocab file.
+
+  If vocab is not found in output dir, it will be copied there by
+  copy_vocab_to_output_dir to clarify the vocab used to generate the data.
+
+  Args:
+    vocab_filepath: path to vocab, either local or cns
+
+  Returns:
+    A SubwordTextEncoder vocabulary object. None if the output_parallel_text
+    is set.
+  """
+  if not tf.gfile.Exists(vocab_filepath):
+    raise ValueError("Vocab file does not exist: {}.".format(vocab_filepath))
+
+  tf.logging.info("Found vocab file: %s", vocab_filepath)
+  encoder = text_encoder.SubwordTextEncoder(vocab_filepath)
+  return encoder
+
+
+def throw_empty_pairs(src_tgt_pairs):
+  """Filter [src,tgt] tuple from input list of pairs if either element is empty.
+
+  Args:
+    src_tgt_pairs: list of (src,tgt) pairs
+
+  Returns:
+    subset of input pair list for which all elements are non-empty
+  """
+  return [x for x in src_tgt_pairs if x[0] and x[1]]
+
+
+def edit_distance_filter(source_target_input, max_equal_to_diff_ratio=0):
+  """Filter out examples that exceed max_edit_ratio between source and target.
+
+  Args:
+    source_target_input:     a list of [source, target] pairs
+    max_equal_to_diff_ratio: cutoff for ratio of equal chars / diff chars
+      between source and target
+
+  Returns:
+    source_target_output:    filtered subset of [source, target] input pairs
+    thrown_out_count:        number of examples filtered out
+  """
+  thrown_out_count = 0
+  source_target_output = []
+
+  if not max_equal_to_diff_ratio:
+    return source_target_input, thrown_out_count
+
+  for i in range(len(source_target_input)):
+    src = source_target_input[i][0]
+    tgt = source_target_input[i][1]
+    opcodes = fast_match_sequences(src, tgt)
+    diff_char_count = 0
+    equal_char_count = 0
+    for tag, i1, i2, j1, j2 in opcodes:
+      if tag == "diff":
+        # max() prevents double-counting substitutions.
+        diff_char_count += max(i2 - i1, j2 - j1)
+      else:
+        equal_char_count += i2 - i1
+    if diff_char_count <= max_equal_to_diff_ratio * equal_char_count:
+      source_target_output.append(source_target_input[i])
+    else:
+      thrown_out_count += 1
+  return source_target_output, thrown_out_count
+
+
+def introduce_errors(s,
+                     corruption_rate=3e-3,
+                     infill_marker="|?|",
+                     max_infill_len=8):
+  """Artificially add spelling errors and infill markers.
+
+  This function should be applied to the inputs of a correction model.
+
+  The artificial errors are particularly useful to train a network to
+  correct spelling when the training data does not contain many
+  natural errors.
+
+  Also replaces some substrings with an "infill" marker.  e.g.
+  "the fat cat sat on the mat" -> "the fat ca??? the mat"
+
+  This causes the trained model to learn infilling (predicting what text
+  to insert at the current cursor position).
+
+  Args:
+    s: a string (the uncorrupted text)
+    corruption_rate: a floating point value.  Probability of introducing an
+      error/infill at each character.
+    infill_marker: a string
+    max_infill_len: an optional integer - maximum number of characters to remove
+      and replace by an infill marker.  None means no infilling.
+
+  Returns:
+    a string
+  """
+  num_errors = 0
+  ret = []
+  operations = [
+      "delete",  # delete a character
+      "insert",  # insert a random character from the input string
+      "replace",  # replace a character with a random character from
+      #   the input string
+      "transpose",  # transpose two adjacent characters
+  ]
+  if max_infill_len:
+    operations.append("infill")
+  pos = 0
+  while pos < len(s):
+    if random.random() >= corruption_rate:
+      ret.append(s[pos])
+      pos += 1
+      continue
+    num_errors += 1
+    operation = operations[random.randint(0, len(operations) - 1)]
+    if operation == "delete":
+      pos += 1
+    elif operation == "insert":
+      ret.append(s[random.randint(0, len(s) - 1)])
+    elif operation == "replace":
+      ret.append(s[random.randint(0, len(s) - 1)])
+      pos += 1
+    elif operation == "transpose":
+      ret.append(s[pos + 1] if pos + 1 < len(s) else "")
+      ret.append(s[pos])
+      pos += 2
+    else:
+      assert operation == "infill"
+      ret.append(infill_marker)
+      pos += random.randint(0, max_infill_len)
+  return "".join(ret), num_errors
+
+
+def fast_match_sequences(a,
+                         b,
+                         a_start=0,
+                         a_end=None,
+                         b_start=0,
+                         b_end=None,
+                         min_match_length=3,
+                         max_recursion_depth=128):
+  """Compute diffs between two sequences.
+
+  This function is similar in functionality and spirit to
+  difflib.SequenceMatcher.get_opcodes, but it seems to run faster.
+
+  if a_start, a_end, b_start, b_end are specified, then we compute diffs of
+  the segments a[a_start:a_end] and b[b_start:b_end].  Returned indices
+  are relative to the full sequence.
+
+  We try to match the longest matching segments first, but due to heuristics
+  in finding the matches, this is not guaranteed.
+
+  Matching segments shorter than min_match_length are counted as part of the
+  surrounding differing segments, unless they are at the beginning or end of
+  both sequences.  This helps eliminate junk matches.
+
+  Args:
+    a: a sequence
+    b: a sequence
+    a_start: an optional integer
+    a_end: an optional integer
+    b_start: an optional integer
+    b_end: an optional integer
+    min_match_length: an integer
+    max_recursion_depth: an integer - avoids crashes in weird corner cases
+      involving pairs of long repetitive sequences.
+
+  Returns:
+    a list of 5-tuples (tag, i1, i2, j1, j2).
+    Each tuple represents the alignment of segment a[i1:i2] with b[j1:j2].
+      tag is either "equal" or "diff".  Note that the tags differ from those
+      returned by difflib.SequenceMatcher.get_opcodes.
+  """
+  if a_end is None:
+    a_end = len(a)
+  if b_end is None:
+    b_end = len(b)
+  if a_start == a_end and b_start == b_end:
+    return []
+  if a_start == a_end or b_start == b_end:
+    return [("diff", a_start, a_end, b_start, b_end)]
+  # Compute an index from value to first occurrence in the b segment.
+  # Technically, we should index and explore all occurrences of a value,
+  # but that might be much slower.
+  b_index = {}
+  for j in range(b_end - 1, b_start - 1, -1):
+    b_index[b[j]] = j
+  # we will look for the longest match we can find.
+  max_match_length = 0
+  a_pos = a_start
+  while a_pos < a_end:
+    val = a[a_pos]
+    b_pos = b_index.get(val)
+    if b_pos is None:
+      a_pos += 1
+      continue
+    else:
+      a_match_start = a_pos
+      a_match_end = a_pos + 1
+      b_match_start = b_pos
+      b_match_end = b_pos + 1
+      while (a_match_start > a_start and b_match_start > b_start and
+             a[a_match_start - 1] == b[b_match_start - 1]):
+        a_match_start -= 1
+        b_match_start -= 1
+      while (a_match_end < a_end and b_match_end < b_end and
+             a[a_match_end] == b[b_match_end]):
+        a_match_end += 1
+        b_match_end += 1
+      # Compute the length of the matching segment.  We prefer the longest.
+      match_length = a_match_end - a_match_start
+      # Extra credit for matching at the beginning or end of the sequence.
+      if a_match_start == 0 and b_match_start == 0:
+        match_length += min_match_length
+      if a_match_end == len(a) and b_match_end == len(b):
+        match_length += min_match_length
+      if match_length > max_match_length:
+        max_match_length = match_length
+        best_match = (a_match_start, a_match_end, b_match_start, b_match_end)
+      # advance a_pos to the end of this match to avoid wasting time
+      # rediscovering this match.
+      a_pos = a_match_end
+  if max_match_length < min_match_length or max_recursion_depth == 0:
+    return [("diff", a_start, a_end, b_start, b_end)]
+  a_match_start, a_match_end, b_match_start, b_match_end = best_match
+  return (fast_match_sequences(
+      a, b, a_start, a_match_start, b_start, b_match_start, min_match_length,
+      max_recursion_depth - 1) + [
+          ("equal", a_match_start, a_match_end, b_match_start, b_match_end)
+      ] + fast_match_sequences(a, b, a_match_end, a_end, b_match_end, b_end,
+                               min_match_length, max_recursion_depth - 1))

From fb2352e1d7d1997d45a55f02478fbe79aadc2ab5 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 31 Oct 2018 15:57:10 -0700
Subject: [PATCH 1137/2720] adding a simple tool for printing confusion matrix
 for reward prediction.

PiperOrigin-RevId: 219552658
---
 tensor2tensor/utils/video/reward_confusion.py | 110 ++++++++++++++++++
 1 file changed, 110 insertions(+)
 create mode 100644 tensor2tensor/utils/video/reward_confusion.py

diff --git a/tensor2tensor/utils/video/reward_confusion.py b/tensor2tensor/utils/video/reward_confusion.py
new file mode 100644
index 000000000..1cb2c59c4
--- /dev/null
+++ b/tensor2tensor/utils/video/reward_confusion.py
@@ -0,0 +1,110 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+r"""Computes the reward prediction confusion matrix given checkpoints and data.
+
+  Usage:
+  reward_confusion \
+  --problem="gym_pong_deterministic-v4_random" \
+  --model="next_frame_sv2p" \
+  --hparams_set="next_frame_sv2p" \
+  --output_dir=$CHECKPOINT_DIRECTORY \
+  --data_dir=$DATA_DIRECTORY \
+
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.bin.t2t_decoder import create_hparams
+from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import trainer_lib
+from tensor2tensor.utils import usr_dir
+
+import tensorflow as tf
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+
+def print_confusion_matrix(title, cm):
+  print("=" * 30)
+  print(title)
+  print("=" * 30)
+  print(cm)
+  print("=" * 30)
+  print()
+
+
+def main(_):
+  tf.logging.set_verbosity(tf.logging.INFO)
+  trainer_lib.set_random_seed(FLAGS.random_seed)
+  usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
+
+  # Create hparams
+  hparams = create_hparams()
+  hparams.force_full_predict = True
+  batch_size = hparams.batch_size
+
+  # Iterating over dev/test partition of the data.
+  # Change the data partition if necessary.
+  dataset = registry.problem(FLAGS.problem).dataset(
+      tf.estimator.ModeKeys.PREDICT,
+      shuffle_files=False,
+      hparams=hparams)
+
+  dataset = dataset.apply(tf.contrib.data.batch_and_drop_remainder(batch_size))
+  data = dataset.make_one_shot_iterator().get_next()
+  input_data = dict((k, data[k]) for k in data.keys() if k.startswith("input"))
+
+  # Creat model
+  model_cls = registry.model(FLAGS.model)
+  model = model_cls(hparams, tf.estimator.ModeKeys.PREDICT)
+  prediction_ops = model.infer(input_data)
+
+  # Confusion Matrix
+  nr = hparams.problem.num_rewards
+  cm_per_frame = np.zeros((nr, nr), dtype=np.uint64)
+  cm_next_frame = np.zeros((nr, nr), dtype=np.uint64)
+
+  saver = tf.train.Saver()
+  with tf.train.SingularMonitoredSession() as sess:
+    # Load latest checkpoint
+    ckpt = tf.train.get_checkpoint_state(FLAGS.output_dir).model_checkpoint_path
+    saver.restore(sess.raw_session(), ckpt)
+
+    counter = 0
+    while not sess.should_stop():
+      counter += 1
+      if counter % 1 == 0:
+        print(counter)
+
+      # Predict next frames
+      rew_pd, rew_gt = sess.run(
+          [prediction_ops["target_reward"], data["target_reward"]])
+
+      for i in range(batch_size):
+        cm_next_frame[rew_gt[i, 0, 0], rew_pd[i, 0, 0]] += 1
+        for gt, pd in zip(rew_gt[i], rew_pd[i]):
+          cm_per_frame[gt, pd] += 1
+
+  print_confusion_matrix("Per-frame Confusion Matrix", cm_per_frame)
+  print_confusion_matrix("Next-frame Confusion Matrix", cm_next_frame)
+
+if __name__ == "__main__":
+  tf.app.run()

From f5e5095fe2b5cf0e3d02d97554b6ec4c82b8292b Mon Sep 17 00:00:00 2001
From: Mike Dusenberry <dusenberrymw@google.com>
Date: Wed, 31 Oct 2018 17:25:56 -0700
Subject: [PATCH 1138/2720] Fix issue with interface and serialization.

PiperOrigin-RevId: 219564645
---
 tensor2tensor/layers/bayes.py      | 29 +++++++++++++++++++++++------
 tensor2tensor/layers/bayes_test.py |  1 +
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 8c43cf288..76e6ae5f8 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -42,11 +42,15 @@ def softplus():  # alias, following tf.keras.constraints
 class TrainableInitializer(tf.keras.initializers.Initializer):
   """An initializer with trainable variables.
 
-  In this implementation, one must call `build` before usage in order to
-  capture the variables within the caller.
+  In this implementation, a layer must call `build` before usage in order to
+  capture the variables.
   """
 
+  def __init__(self):
+    self.built = False
+
   def build(self, shape, dtype=None, add_variable_fn=None):
+    """Builds the initializer, with the variables captured by the caller."""
     raise NotImplementedError
 
 
@@ -64,6 +68,7 @@ def __init__(self,
                seed=None,
                dtype=tf.float32):
     """Constructs the initializer."""
+    super(TrainableNormal, self).__init__()
     self.mean_initializer = mean_initializer
     self.unconstrained_stddev_initializer = unconstrained_stddev_initializer
     self.mean_regularizer = mean_regularizer
@@ -78,7 +83,7 @@ def build(self, shape, dtype=None, add_variable_fn=None):
     if dtype is None:
       dtype = self.dtype
     self.shape = shape
-    self.dtype = dtype
+    self.dtype = tf.as_dtype(dtype)
 
     self.mean = add_variable_fn(
         'mean',
@@ -96,8 +101,15 @@ def build(self, shape, dtype=None, add_variable_fn=None):
         constraint=self.unconstrained_stddev_constraint,
         dtype=dtype,
         trainable=True)
+    self.built = True
 
-  def __call__(self):
+  def __call__(self, shape=None, dtype=None, partition_info=None):
+    del shape, dtype, partition_info  # Unused in TrainableInitializers.
+    # TODO(dusenberrymw): Restructure so that we can build as needed.
+    if not self.built:
+      raise ValueError('A TrainableInitializer must be built by a layer before '
+                       'usage, and is currently only compatible with Bayesian '
+                       'layers.')
     noise = tf.random_normal(self.shape, dtype=self.dtype, seed=self.seed)
     output = self.mean + self.stddev * noise
     # TODO(trandustin): Hack to store parameters so KL reg. can operate on them.
@@ -116,13 +128,12 @@ def get_config(self):
         'unconstrained_stddev_regularizer':
             tf.keras.regularizers.serialize(
                 self.unconstrained_stddev_regularizer),
-        'activity_regularizer':
-            tf.keras.regularizers.serialize(self.activity_regularizer),
         'mean_constraint':
             tf.keras.constraints.serialize(self.mean_constraint),
         'unconstrained_stddev_constraint':
             tf.keras.constraints.serialize(
                 self.unconstrained_stddev_constraint),
+        'seed': self.seed,
         'dtype': self.dtype.name,
     }
 
@@ -147,6 +158,12 @@ def __call__(self, x):
     regularization += (variance_ratio - 1. - tf.log(variance_ratio)) / 2.
     return regularization
 
+  def get_config(self):
+    return {
+        'mean': self.mean,
+        'stddev': self.stddev,
+    }
+
 
 def normal_kl_divergence():  # alias, following tf.keras.regularizers
   return NormalKLDivergence()
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index e4982126c..c9e754438 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -43,6 +43,7 @@ def testDenseReparameterizationKernel(self):
     self.assertEqual(res1.shape, (5, 3, 4))
     self.assertAllGreaterEqual(res1, 0.)
     # self.assertNotAllClose(res1, res2)
+    layer.get_config()
 
   # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
   # support for TF 1.10

From a4c1c039c8615d5e02c7f61da20efd33a2971202 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Wed, 31 Oct 2018 18:00:47 -0700
Subject: [PATCH 1139/2720] fix GAE estimator.

PiperOrigin-RevId: 219567961
---
 tensor2tensor/rl/ppo.py | 34 +++++++++++++++++++---------------
 1 file changed, 19 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index 69ab1ba0d..f4d26346c 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -90,7 +90,7 @@ def define_ppo_epoch(memory, hparams):
   advantage = calculate_generalized_advantage_estimator(
       reward, value, done, hparams.gae_gamma, hparams.gae_lambda)
 
-  discounted_reward = tf.stop_gradient(advantage + value)
+  discounted_reward = tf.stop_gradient(advantage + value[:-1])
 
   advantage_mean, advantage_variance = tf.nn.moments(advantage, axes=[0, 1],
                                                      keep_dims=True)
@@ -99,7 +99,7 @@ def define_ppo_epoch(memory, hparams):
 
   add_lists_elementwise = lambda l1, l2: [x + y for x, y in zip(l1, l2)]
 
-  number_of_batches = (hparams.epoch_length * hparams.optimization_epochs
+  number_of_batches = ((hparams.epoch_length-1) * hparams.optimization_epochs
                        / hparams.optimization_batch_size)
 
   if hasattr(hparams, "effective_num_agents"):
@@ -107,10 +107,11 @@ def define_ppo_epoch(memory, hparams):
     number_of_batches /= hparams.effective_num_agents
 
   dataset = tf.data.Dataset.from_tensor_slices(
-      (observation, action, discounted_reward, advantage_normalized, old_pdf))
-  dataset = dataset.shuffle(buffer_size=hparams.epoch_length,
+      (observation[:-1], action[:-1], discounted_reward, advantage_normalized,
+       old_pdf[:-1]))
+  dataset = dataset.shuffle(buffer_size=hparams.epoch_length-1,
                             reshuffle_each_iteration=True)
-  dataset = dataset.repeat(hparams.optimization_epochs)
+  dataset = dataset.repeat(-1)
   dataset = dataset.batch(hparams.optimization_batch_size)
   iterator = dataset.make_initializable_iterator()
   optimizer = get_optimiser(hparams)
@@ -140,16 +141,19 @@ def define_ppo_epoch(memory, hparams):
 
 def calculate_generalized_advantage_estimator(
     reward, value, done, gae_gamma, gae_lambda):
-  """Generalized advantage estimator."""
-
-  # Below is slight weirdness, we set the last reward to 0.
-  # This makes the advantage to be 0 in the last timestep
-  reward = tf.concat([reward[:-1, :], value[-1:, :]], axis=0)
-  next_value = tf.concat([value[1:, :], tf.zeros_like(value[-1:, :])], axis=0)
-  next_not_done = 1 - tf.cast(tf.concat([done[1:, :],
-                                         tf.zeros_like(done[-1:, :])], axis=0),
-                              tf.float32)
-  delta = reward + gae_gamma * next_value * next_not_done - value
+  # pylint: disable=g-doc-args
+  """Generalized advantage estimator.
+
+  Returns:
+    GAE estimator. It will be one element shorter than the input; this is
+    because to compute GAE for [0, ..., N-1] one needs V for [1, ..., N].
+  """
+  # pylint: enable=g-doc-args
+
+  next_value = value[1:, :]
+  next_not_done = 1 - tf.cast(done[1:, :], tf.float32)
+  delta = (reward[:-1, :] + gae_gamma * next_value * next_not_done
+           - value[:-1, :])
 
   return_ = tf.reverse(tf.scan(
       lambda agg, cur: cur[0] + cur[1] * gae_gamma * gae_lambda * agg,

From fa9db4613ce22c2247fca94809082cbc7d25f4df Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 31 Oct 2018 18:12:07 -0700
Subject: [PATCH 1140/2720] removing extra dims from reward.

PiperOrigin-RevId: 219569142
---
 tensor2tensor/models/video/basic_deterministic.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 655d9f878..440b8220f 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -139,5 +139,6 @@ def next_frame(self, frames, actions, rewards, target_frame,
     reward_pred = tf.concat([x_mid, x_fin], axis=-1)
     reward_pred = tf.nn.relu(tf.layers.dense(
         reward_pred, 128, name="reward_pred"))
-    reward_pred = tf.expand_dims(reward_pred, axis=3)  # Need fake channels dim.
+    reward_pred = tf.squeeze(reward_pred, axis=1)  # Remove extra dims
+    reward_pred = tf.squeeze(reward_pred, axis=1)  # Remove extra dims
     return x, reward_pred, extra_loss, internal_states

From dae9e7bac6964b8c70e0b0316ba3da705b58c2b4 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 31 Oct 2018 22:02:16 -0700
Subject: [PATCH 1141/2720] a simple tool for saving predictions as gifs.

PiperOrigin-RevId: 219586028
---
 tensor2tensor/utils/video/prediction2gif.py | 160 ++++++++++++++++++++
 1 file changed, 160 insertions(+)
 create mode 100644 tensor2tensor/utils/video/prediction2gif.py

diff --git a/tensor2tensor/utils/video/prediction2gif.py b/tensor2tensor/utils/video/prediction2gif.py
new file mode 100644
index 000000000..3be463c6e
--- /dev/null
+++ b/tensor2tensor/utils/video/prediction2gif.py
@@ -0,0 +1,160 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+r"""Generates gifs out of a video checkpoint.
+
+  Usage:
+  prediction2gif \
+  --problem="gym_pong_deterministic-v4_random" \
+  --model="next_frame_sv2p" \
+  --hparams_set="next_frame_sv2p" \
+  --output_dir=$CHECKPOINT_DIRECTORY \
+  --data_dir=$DATA_DIRECTORY \
+  --output_gif=$USER/out.gif \
+
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import matplotlib as mpl
+import numpy as np
+from queue import Queue
+from tensor2tensor.bin.t2t_decoder import create_hparams
+from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
+from tensor2tensor.layers import common_video
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import trainer_lib
+from tensor2tensor.utils import usr_dir
+
+import tensorflow as tf
+
+mpl.use("Agg")
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+
+flags.DEFINE_integer("num_steps", 100, "Number of prediction steps.")
+flags.DEFINE_integer("fps", 10, "Generated gif FPS.")
+flags.DEFINE_string("output_gif", None, "Output path to save the gif.")
+
+
+def main(_):
+  tf.logging.set_verbosity(tf.logging.INFO)
+  trainer_lib.set_random_seed(FLAGS.random_seed)
+  usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
+
+  # Create hparams
+  hparams = create_hparams()
+  hparams.force_full_predict = True
+  hparams.scheduled_sampling_k = -1
+
+  # Params
+  num_agents = 1  # TODO(mbz): fix the code for more agents
+  num_steps = FLAGS.num_steps
+  num_actions = hparams.problem.num_actions
+  frame_shape = hparams.problem.frame_shape
+  resized_frame = hparams.preprocess_resize_frames is not None
+  if resized_frame:
+    frame_shape = hparams.preprocess_resize_frames
+    frame_shape += [hparams.problem.num_channels]
+
+  dataset = registry.problem(FLAGS.problem).dataset(
+      tf.estimator.ModeKeys.TRAIN,
+      shuffle_files=True,
+      hparams=hparams)
+
+  dataset = dataset.apply(tf.contrib.data.batch_and_drop_remainder(num_agents))
+  data = dataset.make_one_shot_iterator().get_next()
+  # Setup input placeholders
+  input_size = [num_agents, hparams.video_num_input_frames]
+  placeholders = {
+      "inputs": tf.placeholder(tf.float32, input_size + frame_shape),
+      "input_action": tf.placeholder(tf.int64, input_size + [1]),
+      "input_reward": tf.placeholder(tf.int64, input_size + [1]),
+  }
+  # Creat model
+  model_cls = registry.model(FLAGS.model)
+  model = model_cls(hparams, tf.estimator.ModeKeys.PREDICT)
+  prediction_ops = model.infer(placeholders)
+
+  states_q = Queue(maxsize=hparams.video_num_input_frames)
+  actions_q = Queue(maxsize=hparams.video_num_input_frames)
+  rewards_q = Queue(maxsize=hparams.video_num_input_frames)
+  all_qs = (states_q, actions_q, rewards_q)
+
+  writer = common_video.WholeVideoWriter(fps=10, output_path=FLAGS.output_gif)
+
+  saver = tf.train.Saver()
+  with tf.train.SingularMonitoredSession() as sess:
+    # Load latest checkpoint
+    ckpt = tf.train.get_checkpoint_state(FLAGS.output_dir).model_checkpoint_path
+    saver.restore(sess.raw_session(), ckpt)
+
+    # get init frames from the dataset
+    data_np = sess.run(data)
+
+    frames = np.split(data_np["inputs"], hparams.video_num_input_frames, 1)
+    for frame in frames:
+      frame = np.squeeze(frame, 1)
+      states_q.put(frame)
+      writer.write(frame[0].astype(np.uint8))
+
+    actions = np.split(data_np["input_action"],
+                       hparams.video_num_input_frames, 1)
+    for action in actions:
+      actions_q.put(np.squeeze(action, 1))
+
+    rewards = np.split(data_np["input_reward"],
+                       hparams.video_num_input_frames, 1)
+    for reward in rewards:
+      rewards_q.put(np.squeeze(reward, 1))
+
+    for step in range(num_steps):
+      print(">>>>>>> ", step)
+
+      random_actions = np.random.randint(num_actions-1)
+      random_actions = np.expand_dims(random_actions, 0)
+      random_actions = np.tile(random_actions, (num_agents, 1))
+
+      # Shape inputs and targets
+      inputs, input_action, input_reward = (
+          np.stack(list(q.queue), axis=1) for q in all_qs)
+
+      # Predict next frames
+      feed = {
+          placeholders["inputs"]: inputs,
+          placeholders["input_action"]: input_action,
+          placeholders["input_reward"]: input_reward,
+      }
+      predictions = sess.run(prediction_ops, feed_dict=feed)
+
+      predicted_states = predictions["targets"][:, 0]
+      predicted_reward = predictions["target_reward"][:, 0]
+
+      # Update queues
+      new_data = (predicted_states, random_actions, predicted_reward)
+      for q, d in zip(all_qs, new_data):
+        q.get()
+        q.put(d.copy())
+
+      writer.write(np.round(predicted_states[0]).astype(np.uint8))
+
+    video = writer.finish()
+    writer.save_to_disk(video)
+
+if __name__ == "__main__":
+  tf.app.run()

From 288f46cf00083ea2642cdf41339e60b0a81cebaf Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 1 Nov 2018 04:59:00 -0700
Subject: [PATCH 1142/2720] Remove the remaining modality strings and combine
 hparams.{input_modalities,target_modality}.

PiperOrigin-RevId: 219617388
---
 tensor2tensor/data_generators/problem.py      | 43 ++-------
 tensor2tensor/layers/common_hparams.py        | 14 +--
 .../layers/common_image_attention.py          | 10 +-
 tensor2tensor/layers/modalities.py            | 94 -------------------
 tensor2tensor/layers/modalities_test.py       | 17 ----
 tensor2tensor/models/image_transformer.py     | 19 ++--
 tensor2tensor/models/image_transformer_2d.py  |  3 +-
 tensor2tensor/models/mtf_transformer.py       |  7 +-
 tensor2tensor/models/research/autoencoders.py | 13 ++-
 tensor2tensor/models/research/cycle_gan.py    |  7 +-
 tensor2tensor/models/research/super_lm.py     |  3 +-
 .../models/research/transformer_symshard.py   |  7 +-
 .../models/research/transformer_vae.py        |  3 +-
 .../video/basic_deterministic_params.py       | 11 ++-
 tensor2tensor/models/video/epva_params.py     |  7 +-
 tensor2tensor/models/video/savp_params.py     |  7 +-
 tensor2tensor/models/video/sv2p_params.py     | 13 ++-
 tensor2tensor/utils/modality.py               |  3 +-
 tensor2tensor/utils/registry.py               | 13 ---
 19 files changed, 87 insertions(+), 207 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index d7e10758a..91e3911c1 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -28,7 +28,6 @@
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import text_encoder
-from tensor2tensor.layers import modalities
 from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import mlperf_log
@@ -1148,33 +1147,12 @@ def _create_modalities(problem_hparams, hparams):
   Returns:
     None
   """
-  input_modality_overrides = {}
-  if hasattr(hparams, "input_modalities"):
-    for override_str in hparams.input_modalities.split(";"):
-      if override_str != "default":
-        parts = override_str.split(":")
-        feature_name = parts[0]
-        modality_name = ":".join(parts[1:])
-        input_modality_overrides[feature_name] = modality_name
-
-  target_modality_name = None
-  if (hasattr(hparams, "target_modality") and
-      hparams.target_modality != "default"):
-    target_modality_name = hparams.target_modality
-
+  modality_overrides = getattr(hparams, "modality", {})
   modality = {}
   for feature_name, modality_cls in six.iteritems(problem_hparams.modality):
     vocab_size = problem_hparams.vocab_size[feature_name]
-    if feature_name in input_modality_overrides:
-      modality_obj = modalities.create_modality(
-          (input_modality_overrides[feature_name], vocab_size), hparams)
-    elif target_modality_name and feature_name == "targets":
-      # TODO(lukaszkaiser): allow overriding other target modalities.
-      modality_obj = modalities.create_modality(
-          (target_modality_name, vocab_size), hparams)
-    else:
-      modality_obj = modality_cls(hparams, vocab_size)
-    modality[feature_name] = modality_obj
+    modality_cls = modality_overrides.get(feature_name, modality_cls)
+    modality[feature_name] = modality_cls(hparams, vocab_size)
   problem_hparams.modality = modality
 
 
@@ -1200,17 +1178,10 @@ def _default_hparams():
       # token.
       stop_at_eos=False,
 
-      # Modalities used to map from input features to a space compatible with
-      # chosen model architecture.  One modality spec (which is a 2-tuple,
-      # (modality_full_name, vocab_size)) per feature key. modality_full_name
-      # is a string type:name, e.g. class_label:class_label_2d. Leaving off
-      # the name uses the default modality for that type (e.g. class_label ==
-      # class_label:default).
-      input_modality={},
-
-      # Modality used to map from hidden representation to the target space.
-      # Specified as a modality spec, a 2-tuple described above.
-      target_modality=None,
+      # Modalities used to map from features to a space compatible with
+      # chosen model architecture. It comprises key-value pairs of a feature
+      # name (str) and its modality class.
+      modality={},
 
       # Identifiers used to tell the model which input/target space will be
       # expected. For example, it can tell that we expect French as characters
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index b9dd752c7..2547dce4b 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -162,18 +162,14 @@ def basic_params1():
       # embeddings and the target embeddings.
       # You can also share the input embeddings with the target embeddings
       # by using a problem_hparams that uses the same modality object for
-      # the input_modality and target_modality.
+      # the input modality and target modality.
       shared_embedding=False,
       # In SymbolModality, skip the top layer, assume we're providing logits.
       symbol_modality_skip_top=False,
-      # For each feature for which you want to override the default input
-      # modality, add an entry to this semicolon-separated string. Entries are
-      # formatted "feature_name:modality_type:modality_name", e.g.
-      # "inputs:symbol:default;other_inputs:audio:identity".
-      input_modalities="default",  # We don't use empty string in params.
-      # To override the default target modality, specify
-      # "modality_type:modality_name", e.g. "symbol:ctc".
-      target_modality="default",
+      # Modalities used to map from features to a space compatible with
+      # chosen model architecture. It comprises key-value pairs of a feature
+      # name (str) and its modality class.
+      modality={},
       # The maximum length of "input" sequence.
       # Sequences longer than this value will be truncated. 0 or negative values
       # mean there is no maximum or truncation.
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index aff70120a..46608d7aa 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -608,8 +608,14 @@ def prepare_image(inputs, hparams, name=None):
   channels = hparams.num_channels
 
   hidden_size = hparams.hidden_size
-  # Only do lookup if the modality is identity
-  if hparams.target_modality == "image:identity":
+  # TODO(trandustin): Check via modalities.IdentityModality and not its name.
+  # The current implementation is to avoid circular imports, modalities ->
+  # discretization -> common_image_attention -> modalities.
+  if "targets" in hparams.modality:
+    target_modality_name = hparams.modality["targets"].__name__
+  else:
+    target_modality_name = None
+  if target_modality_name == "IdentityModality":
     inputs = tf.to_int32(inputs)
     x = get_channel_embeddings(channels, inputs, hidden_size, name=name)
   else:
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index c808ab1c8..511e999f0 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -25,7 +25,6 @@
 from tensor2tensor.layers import common_video
 from tensor2tensor.layers import discretization
 from tensor2tensor.utils import modality
-from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -1070,96 +1069,3 @@ def top(self, body_output, _):
       x = body_output
       x = tf.expand_dims(x[:, -1], 1)  # Pick the last timestep
       return tf.layers.dense(x, self._vocab_size)
-
-
-def create_modality(modality_spec, model_hparams):
-  """Creates modality.
-
-  Args:
-    modality_spec: tuple ("modality_type:modality_name", vocab_size).
-    model_hparams: tf.contrib.training.HParams.
-
-  Returns:
-    Modality.
-
-  Raises:
-    LookupError: if modality_type is not recognized. See registry.Modalities for
-      accepted types.
-  """
-  modality_full_name, vocab_size = modality_spec
-  modality_type, modality_name = parse_modality_name(modality_full_name)
-
-  if modality_type == registry.Modalities.SYMBOL:
-    modality_collection = {
-        "default": SymbolModality,
-        "identity": IdentitySymbolModality,
-        "weights_all": SymbolModalityWeightsAll,
-        "one_hot": SymbolModalityOneHot,
-        "ctc": CTCSymbolModality,
-    }
-  elif modality_type == registry.Modalities.IMAGE:
-    modality_collection = {
-        "default": ImageModality,
-        "identity": IdentityModality,
-        "image_channel_compress": ImageChannelCompressModality,
-        "image_channel_bottom_identity": ImageChannelBottomIdentityModality,
-        "channel_embeddings_bottom": ImageChannelEmbeddingsBottom,
-    }
-  elif modality_type == registry.Modalities.AUDIO:
-    modality_collection = {
-        "default": SpeechRecognitionModality,
-        "identity": IdentityModality,
-        "spectral": AudioSpectralModality,
-        "speech": SpeechRecognitionModality,
-    }
-  elif modality_type == registry.Modalities.VIDEO:
-    modality_collection = {
-        "default": VideoModality,
-        "identity": IdentityModality,
-        "bitwise": VideoModalityBitwise,
-        "pixel_noise": VideoModalityPixelNoise,
-        "l1": VideoModalityL1,
-        "l2": VideoModalityL2,
-        "l2raw": VideoModalityL2Raw,
-        "l1raw": VideoModalityL1Raw,
-    }
-  elif modality_type == registry.Modalities.CLASS_LABEL:
-    modality_collection = {
-        "default": ClassLabelModality,
-        "identity": IdentityModality,
-        "multi_label": MultiLabelModality,
-        "onehot": OneHotClassLabelModality,
-        "sigmoid": SigmoidClassLabelModality,
-        "sigmoid_max_pooling": SigmoidMaxPoolingClassLabelModality,
-        "onehot_softmax_max_pooling": SoftmaxMaxPoolingClassLabelModality,
-        "onehot_softmax_average_pooling":
-            SoftmaxAveragePoolingClassLabelModality,
-        "onehot_softmax_last_timestep": SoftmaxLastTimestepClassLabelModality,
-    }
-  elif modality_type == registry.Modalities.GENERIC:
-    modality_collection = {
-        "default": IdentityModality,
-        "l2_loss": GenericL2LossModality,
-    }
-  elif modality_type == registry.Modalities.REAL:
-    modality_collection = {
-        "default": RealL2LossModality,
-        "identity": IdentityModality,
-        "l2_loss": RealL2LossModality,
-        "log_poisson_loss": RealLogPoissonLossModality,
-    }
-  else:
-    modality_types = ("symbol", "image", "audio", "video", "class_label",
-                      "generic", "real")
-    raise LookupError("Modality type %s not recognized. Options are: %s" %
-                      (modality_type, list(modality_types)))
-
-  return modality_collection[modality_name](model_hparams, vocab_size)
-
-
-def parse_modality_name(name):
-  name_parts = name.split(":")
-  if len(name_parts) < 2:
-    name_parts.append("default")
-  modality_type, modality_name = name_parts
-  return modality_type, modality_name
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index e10a9d10c..f1aa4d134 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -23,7 +23,6 @@
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import expert_utils
-from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -113,22 +112,6 @@ def testSymbolModalityTargetsFactored(self):
     self.assertEqual(res1.shape, (batch_size, length, height, 1, vocab_size))
     self.assertEqual(res2.shape, ())
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
-  def testCreateModality(self):
-    model_hparams = tf.contrib.training.HParams()
-
-    modality_spec = (registry.Modalities.SYMBOL, 2)
-    modality = modalities.create_modality(modality_spec, model_hparams)
-    self.assertIsInstance(modality, modalities.SymbolModality)
-
-    modality_spec = (registry.Modalities.CLASS_LABEL + ":onehot", None)
-    modality = modalities.create_modality(modality_spec, model_hparams)
-    self.assertIsInstance(modality, modalities.OneHotClassLabelModality)
-
-    modality_spec = (registry.Modalities.VIDEO + ":identity", None)
-    modality = modalities.create_modality(modality_spec, model_hparams)
-    self.assertIsInstance(modality, modalities.IdentityModality)
-
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 0714685b6..4b52f9a42 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -28,6 +28,7 @@
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -47,14 +48,16 @@ def body(self, features):
     hparams = copy.copy(self._hparams)
     targets = features["targets"]
     if (hparams.likelihood == cia.DistributionType.DMOL and
-        (hparams.target_modality != "image:image_channel_bottom_identity" or
+        (hparams.modality["targets"] !=
+         modalities.ImageChannelBottomIdentityModality or
          hparams.num_channels != 1)):
-      raise ValueError("When using DMOL for the likelihood, target_modality "
-                       "must be image:image_channel_bottom_identity and "
+      raise ValueError("When using DMOL for the likelihood,modality['targets'] "
+                       "must be ImageChannelBottomIdentityModality and "
                        "num_channels must be 1.")
     if (not tf.get_variable_scope().reuse and
         hparams.mode != tf.contrib.learn.ModeKeys.INFER and
-        hparams.target_modality != "image:image_channel_bottom_identity"):
+        hparams.modality["targets"] !=
+        modalities.ImageChannelBottomIdentityModality):
       tf.summary.image("targets", tf.to_float(targets), max_outputs=1)
 
     # Extra losses list if we want to use moe.
@@ -190,7 +193,7 @@ def image_transformer_base():
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.98
   hparams.label_smoothing = 0.0
-  hparams.target_modality = "image:identity"
+  hparams.modality["targets"] = modalities.IdentityModality
   hparams.norm_type = "layer"
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.add_hparam("filter_size", 512)  # Add new ones like this.
@@ -277,7 +280,7 @@ def imagetransformer_cifar10_base_dmol():
   hparams = image_transformer_base()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.target_modality = "image:image_channel_bottom_identity"
+  hparams.modality["targets"] = modalities.ImageChannelBottomIdentityModality
   hparams.num_heads = 8
   hparams.batch_size = 8
   hparams.sampling_method = "random"
@@ -418,7 +421,7 @@ def imagetransformerpp_sep_channels_8l_8h():
   hparams = imagetransformer_base()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.target_modality = "image:image_channel_bottom_identity"
+  hparams.modality["targets"] = modalities.ImageChannelBottomIdentityModality
   hparams.num_heads = 8
   hparams.batch_size = 4
   hparams.attention_key_channels = hparams.attention_value_channels = 0
@@ -881,7 +884,7 @@ def imagetransformerpp_tiny():
   hparams = imagetransformer_tiny()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.target_modality = "image:image_channel_bottom_identity"
+  hparams.modality["targets"] = modalities.ImageChannelBottomIdentityModality
   return hparams
 
 
diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 1c105186f..7555150a7 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -29,6 +29,7 @@
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -381,7 +382,7 @@ def image_transformer2d_base():
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.98
   hparams.label_smoothing = 0.0
-  hparams.target_modality = "image:identity"
+  hparams.modality["targets"] = modalities.IdentityModality
   hparams.norm_type = "layer"
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.add_hparam("filter_size", 512)  # Add new ones like this.
diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 65b644b45..8cff5786e 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -24,6 +24,7 @@
 
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.models.research import moe
 from tensor2tensor.utils import mtf_model
 from tensor2tensor.utils import registry
@@ -772,8 +773,10 @@ def mtf_transformer_base():
   # These parameters make Transformer model compatible with MtfTransformer
   # Do not override these, as mtf_transformer does not support other options.
   hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
-  hparams.target_modality = "symbol:identity"
-  hparams.input_modalities = "inputs:symbol:identity"
+  hparams.modality = {
+      "inputs": modalities.IdentitySymbolModality,
+      "targets": modalities.IdentitySymbolModality,
+  }
 
   # Parameters for computing the maximum decode length in beam search.
   # Maximum decode length is:
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 0688a0c09..3780bd6e6 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -24,6 +24,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
 from tensor2tensor.layers import latent_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -1104,8 +1105,10 @@ def autoencoder_residual_text():
   hparams.hidden_size = 64
   hparams.max_hidden_size = 512
   hparams.bottleneck_noise = 0.0
-  hparams.target_modality = "symbol:identity"
-  hparams.input_modalities = "symbol:identity"
+  hparams.modality = {
+      "inputs": modalities.IdentitySymbolModality,
+      "targets": modalities.IdentitySymbolModality,
+  }
   hparams.autoregressive_mode = "none"
   hparams.sample_width = 1
   return hparams
@@ -1209,8 +1212,10 @@ def autoencoder_ordered_text():
   hparams.batch_size = 1024
   hparams.autoregressive_mode = "conv5"
   hparams.max_hidden_size = 1024
-  hparams.target_modality = "symbol:identity"
-  hparams.input_modalities = "symbol:identity"
+  hparams.modality = {
+      "inputs": modalities.IdentitySymbolModality,
+      "targets": modalities.IdentitySymbolModality,
+  }
   hparams.sample_height = 128
   hparams.sample_width = 1
   return hparams
diff --git a/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index 935e292f2..49a2653d1 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.models.research import transformer_vae
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
@@ -126,8 +127,10 @@ def cycle_gan_small():
   """Set of hyperparameters."""
   hparams = transformer_vae.transformer_ae_small()
   hparams.batch_size = 2048
-  hparams.input_modalities = "inputs:symbol:identity"
-  hparams.target_modality = "symbol:identity"
+  hparams.modality = {
+      "inputs": modalities.IdentitySymbolModality,
+      "targets": modalities.IdentitySymbolModality,
+  }
   hparams.weight_decay = 3.0
   hparams.learning_rate = 0.05
   hparams.kl_warmup_steps = 5000
diff --git a/tensor2tensor/models/research/super_lm.py b/tensor2tensor/models/research/super_lm.py
index aaf8e0bfb..6b4b137ec 100644
--- a/tensor2tensor/models/research/super_lm.py
+++ b/tensor2tensor/models/research/super_lm.py
@@ -32,6 +32,7 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import diet
 from tensor2tensor.utils import expert_utils
 from tensor2tensor.utils import registry
@@ -264,7 +265,7 @@ def super_lm_base():
   # we only want one data shard.
   hparams.no_data_parallelism = True
   # bypass the symbol modality so that we can use model parallelism.
-  hparams.target_modality = "symbol:identity"
+  hparams.modality["targets"] = modalities.IdentitySymbolModality
   hparams.add_hparam("filter_size", 512)
   hparams.add_hparam("mix_fraction", 0.5)
   # attention-related flags
diff --git a/tensor2tensor/models/research/transformer_symshard.py b/tensor2tensor/models/research/transformer_symshard.py
index 1762df8f1..8fa58ab20 100644
--- a/tensor2tensor/models/research/transformer_symshard.py
+++ b/tensor2tensor/models/research/transformer_symshard.py
@@ -48,6 +48,7 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import expert_utils
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
@@ -366,8 +367,10 @@ def transformer_symshard_base():
   # we only want one data shard.
   hparams.no_data_parallelism = True
   # bypass the symbol modality so that we can use model parallelism.
-  hparams.target_modality = "symbol:identity"
-  hparams.input_modalities = "inputs:symbol:identity"
+  hparams.modality = {
+      "inputs": modalities.IdentitySymbolModality,
+      "targets": modalities.IdentitySymbolModality,
+  }
   hparams.add_hparam("filter_size", 1280)
   hparams.add_hparam("mix_fraction", 0.5)
   # attention-related flags
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index bd52a39f2..7d84245fd 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -29,6 +29,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
 from tensor2tensor.layers import latent_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.models import transformer
 from tensor2tensor.utils import beam_search
 from tensor2tensor.utils import expert_utils
@@ -889,7 +890,7 @@ def imagetransformer_ae_cifar():
 
   hparams.add_hparam("unconditional", False)  # unconditional generation
 
-  hparams.target_modality = "image:channel_embeddings_bottom"
+  hparams.modality["targets"] = modalities.ImageChannelEmbeddingsBottom
   hparams.drop_inputs = True
   hparams.do_attend_compress = False
   hparams.do_attend_decompress = False
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 7b2c3d0bc..6ce8f0c39 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.layers import modalities
 from tensor2tensor.models.video import base
 from tensor2tensor.utils import registry
 
@@ -52,7 +53,7 @@ def next_frame_pixel_noise():
   """Basic 2-frame conv model with pixel noise."""
   hparams = next_frame_basic_deterministic()
   hparams.add_hparam("video_modality_input_noise", 0.05)
-  hparams.input_modalities = "inputs:video:pixel_noise"
+  hparams.modality["inputs"] = modalities.VideoModalityPixelNoise
   return hparams
 
 
@@ -77,7 +78,7 @@ def next_frame_tpu():
 def next_frame_ae():
   """Conv autoencoder."""
   hparams = next_frame_basic_deterministic()
-  hparams.input_modalities = "inputs:video:bitwise"
+  hparams.modality["inputs"] = modalities.VideoModalityBitwise
   hparams.hidden_size = 256
   hparams.batch_size = 8
   hparams.num_hidden_layers = 4
@@ -90,7 +91,7 @@ def next_frame_ae():
 def next_frame_ae_tiny():
   """Conv autoencoder, tiny set for testing."""
   hparams = next_frame_tiny()
-  hparams.input_modalities = "inputs:video:bitwise"
+  hparams.modality["inputs"] = modalities.VideoModalityBitwise
   hparams.batch_size = 8
   hparams.dropout = 0.4
   return hparams
@@ -119,7 +120,7 @@ def next_frame_tiny():
 def next_frame_l1():
   """Basic conv model with L1 modality."""
   hparams = next_frame_basic_deterministic()
-  hparams.target_modality = "video:l1"
+  hparams.modality["targets"] = modalities.VideoModalityL1
   hparams.video_modality_loss_cutoff = 2.4
   return hparams
 
@@ -128,7 +129,7 @@ def next_frame_l1():
 def next_frame_l2():
   """Basic conv model with L2 modality."""
   hparams = next_frame_basic_deterministic()
-  hparams.target_modality = "video:l2"
+  hparams.modality["targets"] = modalities.VideoModalityL2
   hparams.video_modality_loss_cutoff = 2.4
   return hparams
 
diff --git a/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
index 57f20a0b9..c6587e41e 100644
--- a/tensor2tensor/models/video/epva_params.py
+++ b/tensor2tensor/models/video/epva_params.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.layers import modalities
 from tensor2tensor.models.video import basic_deterministic_params
 from tensor2tensor.utils import registry
 
@@ -28,8 +29,10 @@ def next_frame_epva():
   hparams = basic_deterministic_params.next_frame_basic_deterministic()
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 4
-  hparams.target_modality = "video:l2raw"
-  hparams.input_modalities = "inputs:video:l2raw"
+  hparams.modality = {
+      "inputs": modalities.VideoModalityL2Raw,
+      "targets": modalities.VideoModalityL2Raw,
+  }
   hparams.learning_rate_schedule = "constant"
   hparams.learning_rate_constant = 1e-05
   hparams.batch_size = 2
diff --git a/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index dbc02c771..6acee08c4 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.layers import modalities
 from tensor2tensor.models.video import sv2p_params
 from tensor2tensor.utils import registry
 
@@ -35,8 +36,10 @@ def next_frame_savp():
   hparams.add_hparam("gan_loss_multiplier", 0.01)
   hparams.add_hparam("gan_vae_loss_multiplier", 0.01)
   hparams.add_hparam("gan_optimization", "joint")
-  hparams.target_modality = "video:l1raw"
-  hparams.input_modalities = "inputs:video:l1raw"
+  hparams.modality = {
+      "inputs": modalities.VideoModalityL1Raw,
+      "targets": modalities.VideoModalityL1Raw,
+  }
   hparams.latent_loss_multiplier_schedule = "linear"
   hparams.upsample_method = "bilinear_upsample_conv"
   hparams.internal_loss = False
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index e81afdd85..c9482c901 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.layers import modalities
 from tensor2tensor.models.video import basic_stochastic
 from tensor2tensor.utils import registry
 
@@ -32,8 +33,10 @@ def next_frame_sv2p():
   hparams.video_num_input_frames = 1
   hparams.video_num_target_frames = 3
   hparams.batch_size = 16
-  hparams.target_modality = "video:l2raw"
-  hparams.input_modalities = "inputs:video:l2raw"
+  hparams.modality = {
+      "inputs": modalities.VideoModalityL2Raw,
+      "targets": modalities.VideoModalityL2Raw,
+  }
   hparams.video_modality_loss_cutoff = 0.0
   hparams.scheduled_sampling_mode = "count"
   hparams.scheduled_sampling_k = 900.0
@@ -84,8 +87,10 @@ def next_frame_sv2p_atari():
 def next_frame_sv2p_atari_softmax():
   """SV2P model for atari with softmax."""
   hparams = next_frame_sv2p_atari()
-  hparams.target_modality = "video"
-  hparams.input_modalities = "inputs:video"
+  hparams.modality = {
+      "inputs": modalities.VideoModality,
+      "targets": modalities.VideoModality,
+  }
   hparams.internal_loss = True
   return hparams
 
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index 6e6726885..45b4eb76a 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -20,7 +20,6 @@
 
 import re
 from tensor2tensor.layers import common_layers
-from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -199,4 +198,4 @@ def loss_sharded(self, sharded_top_out, sharded_targets, data_parallelism):
 
   @property
   def is_class_modality(self):
-    return self.name.startswith(registry.Modalities.CLASS_LABEL)
+    return self.name.startswith("class_label")
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 431fae8b0..fa10dacc9 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -58,19 +58,6 @@ class MyModel(T2TModel):
 _RANGED_HPARAMS = {}
 
 
-# TODO(trandustin): Many files depend on this to specify modality strings; let's
-# remove it in the future.
-class Modalities(object):
-  """An enum-like object carrying the set of available modality types."""
-  SYMBOL = "symbol"
-  IMAGE = "image"
-  AUDIO = "audio"
-  VIDEO = "video"
-  CLASS_LABEL = "class_label"
-  GENERIC = "generic"
-  REAL = "real"
-
-
 # Camel case to snake case utils
 _first_cap_re = re.compile("(.)([A-Z][a-z0-9]+)")
 _all_cap_re = re.compile("([a-z0-9])([A-Z])")

From 139b676e2a8305185ade6a0fcfb5ea27ab33af8f Mon Sep 17 00:00:00 2001
From: Randall Lin <randall@fathomhealth.co>
Date: Thu, 1 Nov 2018 09:57:47 -0700
Subject: [PATCH 1143/2720] Update universal_transformer.py (#1192)

---
 tensor2tensor/models/research/universal_transformer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 90ed8679c..148426beb 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -416,7 +416,7 @@ def update_hparams_for_universal_transformer(hparams):
   # LSTM forget bias for lstm style recurrence.
   hparams.add_hparam("lstm_forget_bias", 1.0)
   # Uses the memory at the last step as the final output, if true.
-  hparams.add_hparam("use_memory_as_final_state", True)
+  hparams.add_hparam("use_memory_as_final_state", False)
   # if also add a ffn unit to the transition function when using gru/lstm
   hparams.add_hparam("add_ffn_unit_to_the_transition_function", False)
 

From cb655f0a2c320efc4fd7c18d6bc301a94e6d6f56 Mon Sep 17 00:00:00 2001
From: Mostafa Dehghani <dehghani.mostafa@gmail.com>
Date: Thu, 1 Nov 2018 17:58:24 +0100
Subject: [PATCH 1144/2720] setting the default for use_memory_as_final_state
 flag to False (#1194)

---
 .../models/research/universal_transformer_util.py          | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 7b1ade84d..50aaa2ca3 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -125,8 +125,6 @@ def universal_transformer_encoder(encoder_input,
     x, extra_output = universal_transformer_layer(
         x, hparams, ffn_unit, attention_unit, pad_remover=pad_remover)
 
-    if hparams.get("use_memory_as_last_state", False):
-      x = extra_output  # which is memory
     return common_layers.layer_preprocess(x, hparams), extra_output
 
 
@@ -251,8 +249,9 @@ def add_vanilla_transformer_layer(x, num_layers):
       output, _, extra_output = tf.foldl(
           ut_function, tf.range(hparams.num_rec_steps), initializer=initializer)
 
-      # This is possible only when we are using lstm as transition function.
-      if hparams.get("use_memory_as_final_state", False):
+      # Right now, this is only possible when the transition function is an lstm
+      if (hparams.recurrence_type == "lstm" and
+          hparams.get("use_memory_as_final_state", False)):
         output = extra_output
 
     if hparams.mix_with_transformer == "after_ut":

From bc1268d8acfb3a2adfead1db25d4006441cb4f7e Mon Sep 17 00:00:00 2001
From: Randall Lin <randall@fathomhealth.co>
Date: Thu, 1 Nov 2018 11:08:19 -0700
Subject: [PATCH 1145/2720] internal merge of PR #1192

PiperOrigin-RevId: 219664613
---
 .../models/research/universal_transformer_util.py          | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 50aaa2ca3..7b1ade84d 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -125,6 +125,8 @@ def universal_transformer_encoder(encoder_input,
     x, extra_output = universal_transformer_layer(
         x, hparams, ffn_unit, attention_unit, pad_remover=pad_remover)
 
+    if hparams.get("use_memory_as_last_state", False):
+      x = extra_output  # which is memory
     return common_layers.layer_preprocess(x, hparams), extra_output
 
 
@@ -249,9 +251,8 @@ def add_vanilla_transformer_layer(x, num_layers):
       output, _, extra_output = tf.foldl(
           ut_function, tf.range(hparams.num_rec_steps), initializer=initializer)
 
-      # Right now, this is only possible when the transition function is an lstm
-      if (hparams.recurrence_type == "lstm" and
-          hparams.get("use_memory_as_final_state", False)):
+      # This is possible only when we are using lstm as transition function.
+      if hparams.get("use_memory_as_final_state", False):
         output = extra_output
 
     if hparams.mix_with_transformer == "after_ut":

From ee3794bc71be3470971820d24b561aced3d35a94 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 1 Nov 2018 11:12:31 -0700
Subject: [PATCH 1146/2720] restricting generated pixel values between 0 and 1.

PiperOrigin-RevId: 219665326
---
 tensor2tensor/models/video/epva.py        | 11 ++++++++++-
 tensor2tensor/models/video/epva_params.py |  1 +
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index 245d9bc81..95164462a 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -28,6 +28,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from six.moves import reduce
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
@@ -654,6 +655,12 @@ def body(self, features):
     # all_rewards = tf.concat([input_rewards, target_rewards], axis=0)
 
     all_actions = tf.concat([input_actions, target_actions], axis=0)
+    # flatten actions tensor to have the shape: framesXbatch_sizeXaction_dims.
+    actions_shape = common_layers.shape_list(all_actions)
+    all_actions = tf.reshape(
+        all_actions,
+        [actions_shape[0], -1,
+         reduce(lambda x, y: x * y, actions_shape[2:])])
     all_frames = tf.concat([input_frames, target_frames], axis=0)
 
     all_frames = tf.unstack(all_frames, axis=0)
@@ -698,6 +705,9 @@ def body(self, features):
 
     predictions = tf.stack(van_on_enc_all)
 
+    if hparams.clip_pixel_values:
+      predictions = tf.clip_by_value(predictions, 0.0, 1.0)
+
     # TODO(mbz): clean this up!
     def fix_video_dims_and_concat_on_x_axis(x):
       x = tf.transpose(x, [1, 3, 4, 0, 2])
@@ -725,4 +735,3 @@ def fix_video_dims_and_concat_on_x_axis(x):
                            [-1]*5)
 
     return predictions, {'extra': epva_loss}
-
diff --git a/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
index c6587e41e..9890f3008 100644
--- a/tensor2tensor/models/video/epva_params.py
+++ b/tensor2tensor/models/video/epva_params.py
@@ -39,6 +39,7 @@ def next_frame_epva():
   hparams.clip_grad_norm = 0.01
   # TODO(msaffar): disentangle EPVA from SV2P
   hparams.add_hparam("reward_prediction", False)
+  hparams.add_hparam("clip_pixel_values", True)
   hparams.add_hparam("context_frames", 5)
   hparams.add_hparam("enc_learning_rate", 1e-5)
   hparams.add_hparam("enc_pred_loss_scale", 0.1)

From 572f1e7b333c1cf77f2b759be10488d6edc7b220 Mon Sep 17 00:00:00 2001
From: Mostafa Dehghani <dehghani.mostafa@gmail.com>
Date: Thu, 1 Nov 2018 15:04:52 -0700
Subject: [PATCH 1147/2720] internal merge of PR #1194

PiperOrigin-RevId: 219705888
---
 .../models/research/universal_transformer_util.py          | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 7b1ade84d..50aaa2ca3 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -125,8 +125,6 @@ def universal_transformer_encoder(encoder_input,
     x, extra_output = universal_transformer_layer(
         x, hparams, ffn_unit, attention_unit, pad_remover=pad_remover)
 
-    if hparams.get("use_memory_as_last_state", False):
-      x = extra_output  # which is memory
     return common_layers.layer_preprocess(x, hparams), extra_output
 
 
@@ -251,8 +249,9 @@ def add_vanilla_transformer_layer(x, num_layers):
       output, _, extra_output = tf.foldl(
           ut_function, tf.range(hparams.num_rec_steps), initializer=initializer)
 
-      # This is possible only when we are using lstm as transition function.
-      if hparams.get("use_memory_as_final_state", False):
+      # Right now, this is only possible when the transition function is an lstm
+      if (hparams.recurrence_type == "lstm" and
+          hparams.get("use_memory_as_final_state", False)):
         output = extra_output
 
     if hparams.mix_with_transformer == "after_ut":

From baccaa2e3ce09f5f558b31774b82085da8454767 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 1 Nov 2018 23:20:45 +0100
Subject: [PATCH 1148/2720] Debug frames from world model evaluation (#1190)

---
 tensor2tensor/rl/trainer_model_based.py | 55 +++++++++++++++++++++----
 1 file changed, 46 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 6a00ebe9c..b30b204b6 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -39,6 +39,7 @@
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
 from tensor2tensor.data_generators.gym_env import T2TGymEnv
+from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.rl import trainer_model_based_params
@@ -87,7 +88,11 @@ def setup_directories(base_dir, subdirs):
 
   all_dirs = {}
   for subdir in subdirs:
-    dir_name = os.path.join(base_dir, subdir)
+    if isinstance(subdir, six.string_types):
+      subdir_tuple = (subdir,)
+    else:
+      subdir_tuple = subdir
+    dir_name = os.path.join(base_dir, *subdir_tuple)
     tf.gfile.MakeDirs(dir_name)
     all_dirs[subdir] = dir_name
   return all_dirs
@@ -367,7 +372,7 @@ def compute_mean_reward(rollouts, clipped):
   return mean_rewards
 
 
-def evaluate_world_model(real_env, hparams, world_model_dir):
+def evaluate_world_model(real_env, hparams, world_model_dir, debug_video_path):
   """Evaluate the world model (reward accuracy)."""
   environment_spec = make_simulated_env_spec(real_env, hparams)
   environment_spec.wrappers = []
@@ -393,6 +398,10 @@ def initial_frame_chooser(batch_size):
       minimal_rollout_frames=(subsequence_length + num_input_frames)
   )
 
+  video_writer = common_video.WholeVideoWriter(
+      fps=10, output_path=debug_video_path, file_format="avi"
+  )
+
   reward_accuracies_by_length = {
       int(ratio * hparams.ppo_epoch_length): []
       for ratio in hparams.wm_eval_rollout_ratios
@@ -411,18 +420,30 @@ def initial_frame_chooser(batch_size):
     # Check that the initial observation is the same in the real and simulated
     # rollout.
     sim_init_obs = sim_env.reset()
-    real_init_obs = np.stack([
-        subsequence[0].observation.decode()
-        for subsequence in eval_subsequences
-    ])
+    def decode_real_obs(index):
+      return np.stack([
+          subsequence[index].observation.decode()
+          for subsequence in eval_subsequences
+      ])
+    real_init_obs = decode_real_obs(0)
     assert np.all(sim_init_obs == real_init_obs)
 
+    debug_frame_batches = []
+    def append_debug_frame_batch(sim_obs, real_obs):
+      errs = np.maximum(
+          np.abs(sim_obs.astype(np.int) - real_obs, dtype=np.int) - 10, 0
+      ).astype(np.uint8)
+      debug_frame_batches.append(  # pylint: disable=cell-var-from-loop
+          np.concatenate([sim_obs, real_obs, errs], axis=2)
+      )
+    append_debug_frame_batch(sim_init_obs, real_init_obs)
+
     (sim_cum_rewards, real_cum_rewards) = (
         np.zeros(hparams.wm_eval_batch_size) for _ in range(2)
     )
     for i in range(subsequence_length):
       actions = [subsequence[i].action for subsequence in eval_subsequences]
-      (_, sim_rewards, _) = sim_env.step(actions)
+      (sim_obs, sim_rewards, _) = sim_env.step(actions)
       sim_cum_rewards += sim_rewards
 
       real_cum_rewards += [
@@ -437,6 +458,15 @@ def initial_frame_chooser(batch_size):
               len(real_cum_rewards)
           )
 
+      real_obs = decode_real_obs(i + 1)
+      append_debug_frame_batch(sim_obs, real_obs)
+
+    for debug_frames in np.stack(debug_frame_batches, axis=1):
+      for debug_frame in debug_frames:
+        video_writer.write(debug_frame)
+
+  video_writer.finish_to_disk()
+
   return {
       "reward_accuracy/at_{}".format(length): np.mean(reward_accuracies)
       for (length, reward_accuracies) in six.iteritems(
@@ -460,7 +490,10 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     assert report_metric is not None
 
   # Directories
-  subdirectories = ["data", "tmp", "world_model", "ppo"]
+  subdirectories = [
+      "data", "tmp", "world_model", ("world_model", "debug_videos"),
+      "ppo"
+  ]
   directories = setup_directories(output_dir, subdirectories)
 
   epoch = -1
@@ -550,8 +583,12 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     env.generate_data(data_dir)
 
     if hparams.eval_world_model:
+      debug_video_path = os.path.join(
+          directories["world_model", "debug_videos"],
+          "{}.avi".format(env.current_epoch)
+      )
       wm_metrics = evaluate_world_model(
-          env, hparams, directories["world_model"]
+          env, hparams, directories["world_model"], debug_video_path
       )
       log("World model eval metrics:\n{}".format(pprint.pformat(wm_metrics)))
       metrics.update(wm_metrics)

From 80062bb6dda85b96abdd0229e0e023c5782f5bde Mon Sep 17 00:00:00 2001
From: Michael Mezher <mezhermikey@yahoo.com>
Date: Thu, 1 Nov 2018 18:35:29 -0400
Subject: [PATCH 1149/2720] Changed reuse val from true to tf.AUTO_REUSE in top
 to allow for proper weight initialization (#1196)

---
 tensor2tensor/layers/modalities.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 511e999f0..c7bb80e05 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -143,7 +143,7 @@ def top(self, body_output, _):
 
     if self._model_hparams.shared_embedding_and_softmax_weights:
       scope_name = "shared"
-      reuse = True
+      reuse = tf.AUTO_REUSE
     else:
       scope_name = "softmax"
       reuse = False

From defd57e64bf40f139d9f97137af683e09980a787 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 1 Nov 2018 15:35:51 -0700
Subject: [PATCH 1150/2720] internal merge of PR #1190

PiperOrigin-RevId: 219710898
---
 tensor2tensor/layers/modalities.py      | 2 +-
 tensor2tensor/rl/trainer_model_based.py | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index c7bb80e05..511e999f0 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -143,7 +143,7 @@ def top(self, body_output, _):
 
     if self._model_hparams.shared_embedding_and_softmax_weights:
       scope_name = "shared"
-      reuse = tf.AUTO_REUSE
+      reuse = True
     else:
       scope_name = "softmax"
       reuse = False
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index b30b204b6..972440b4a 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -83,6 +83,7 @@ def world_model_step_increment(hparams, is_initial_epoch):
 
 
 def setup_directories(base_dir, subdirs):
+  """Setup directories."""
   base_dir = os.path.expanduser(base_dir)
   tf.gfile.MakeDirs(base_dir)
 
@@ -423,7 +424,7 @@ def initial_frame_chooser(batch_size):
     def decode_real_obs(index):
       return np.stack([
           subsequence[index].observation.decode()
-          for subsequence in eval_subsequences
+          for subsequence in eval_subsequences  # pylint: disable=cell-var-from-loop
       ])
     real_init_obs = decode_real_obs(0)
     assert np.all(sim_init_obs == real_init_obs)

From c45e16d45b72812c2edaec4961d85cecb30345a0 Mon Sep 17 00:00:00 2001
From: Michael Mezher <mezhermikey@yahoo.com>
Date: Thu, 1 Nov 2018 15:36:01 -0700
Subject: [PATCH 1151/2720] internal merge of PR #1196

PiperOrigin-RevId: 219710923
---
 tensor2tensor/layers/modalities.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 511e999f0..c7bb80e05 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -143,7 +143,7 @@ def top(self, body_output, _):
 
     if self._model_hparams.shared_embedding_and_softmax_weights:
       scope_name = "shared"
-      reuse = True
+      reuse = tf.AUTO_REUSE
     else:
       scope_name = "softmax"
       reuse = False

From d02fa40659f051f676eb5de317efcb873081ce92 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 1 Nov 2018 17:21:46 -0700
Subject: [PATCH 1152/2720] Wait longer for checkpoints.

PiperOrigin-RevId: 219727041
---
 tensor2tensor/utils/trainer_lib.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 38ab35370..8e4d3c2a5 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -39,7 +39,7 @@
 from tensorflow.python import debug
 
 
-def next_checkpoint(model_dir, timeout_mins=120):
+def next_checkpoint(model_dir, timeout_mins=240):
   """Yields successive checkpoints from model_dir."""
   last_ckpt = None
   while True:

From 986e0f4749a2a6b83578899f39b15532547fcc74 Mon Sep 17 00:00:00 2001
From: Smit Hinsu <hinsu@google.com>
Date: Thu, 1 Nov 2018 18:38:34 -0700
Subject: [PATCH 1153/2720] Reduce the default batch size in transformer_big
 hparams

This is to allow users to train the Transformer model with transformer_big hparams on a GPU with 12 GB memory.

PiperOrigin-RevId: 219735751
---
 tensor2tensor/models/transformer.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 1552be555..fc7fa666e 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1501,6 +1501,9 @@ def transformer_big():
   hparams = transformer_base()
   hparams.hidden_size = 1024
   hparams.filter_size = 4096
+  # Reduce batch size to 2048 from 4096 to be able to train the model on a GPU
+  # with 12 GB memory. For example, NVIDIA TITAN V GPU.
+  hparams.batch_size = 2048
   hparams.num_heads = 16
   hparams.layer_prepostprocess_dropout = 0.3
   return hparams

From 52f6e552c5ce5aa7e8cff4b3d1e47340b5f5d59b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 1 Nov 2018 20:14:49 -0700
Subject: [PATCH 1154/2720] Add missing MLPerf L2 logging.

PiperOrigin-RevId: 219742753
---
 tensor2tensor/data_generators/problem.py   | 2 ++
 tensor2tensor/layers/transformer_layers.py | 8 ++++++++
 tensor2tensor/models/transformer.py        | 7 +++++++
 tensor2tensor/utils/decoding.py            | 5 +++++
 tensor2tensor/utils/mlperf_tags.py         | 1 +
 5 files changed, 23 insertions(+)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 91e3911c1..0526ff972 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -837,6 +837,8 @@ def input_fn(self,
       num_shards = 1
 
     max_length = self.max_length(hparams)
+    mlperf_log.transformer_print(
+        key=mlperf_log.INPUT_MAX_LENGTH, value=max_length)
 
     def tpu_valid_size(example):
       return data_reader.example_valid_size(example, hparams.min_length,
diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 719ef1c4f..84cc5c557 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -137,6 +137,14 @@ def transformer_encoder(encoder_input,
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_ATTENTION_DROPOUT,
       value=hparams.attention_dropout)
+  mlperf_log.transformer_print(
+      key=mlperf_log.MODEL_HP_ATTENTION_DENSE,
+      value={
+          "use_bias": "false",
+          "num_heads": hparams.num_heads,
+          "hidden_size": hparams.hidden_size
+      })
+
   with tf.variable_scope(name):
     if nonpadding is not None:
       padding = 1.0 - nonpadding
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index fc7fa666e..9c9ab14d8 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1226,6 +1226,13 @@ def transformer_decoder(decoder_input,
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_ATTENTION_DROPOUT,
       value=hparams.attention_dropout)
+  mlperf_log.transformer_print(
+      key=mlperf_log.MODEL_HP_ATTENTION_DENSE,
+      value={
+          "use_bias": "false",
+          "num_heads": hparams.num_heads,
+          "hidden_size": hparams.hidden_size
+      })
 
   with tf.variable_scope(name):
     for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index d4a38663b..d3a86e91c 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -32,6 +32,7 @@
 from tensor2tensor.data_generators import problem as problem_lib
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import registry
 import tensorflow as tf
 
@@ -264,7 +265,9 @@ def decode_once(estimator,
   inputs_vocab = problem_hparams.vocabulary[inputs_vocab_key]
   targets_vocab = problem_hparams.vocabulary["targets"]
 
+  num_eval_samples = 0
   for num_predictions, prediction in enumerate(predictions):
+    num_eval_samples += 1
     num_predictions += 1
     inputs = prediction.get("inputs")
     targets = prediction.get("targets")
@@ -328,6 +331,8 @@ def decode_once(estimator,
         num_predictions >= decode_hp.num_samples):
       break
 
+  mlperf_log.transformer_print(key=mlperf_log.EVAL_SIZE, value=num_eval_samples)
+
   if decode_to_file:
     output_file.close()
     target_file.close()
diff --git a/tensor2tensor/utils/mlperf_tags.py b/tensor2tensor/utils/mlperf_tags.py
index 4e599d7f4..764c4399e 100644
--- a/tensor2tensor/utils/mlperf_tags.py
+++ b/tensor2tensor/utils/mlperf_tags.py
@@ -313,6 +313,7 @@
     TRAIN_LOOP,
     TRAIN_EPOCH,
     EVAL_START,
+    EVAL_SIZE,
     EVAL_TARGET,
     EVAL_ACCURACY,
     EVAL_STOP,

From fc08642dfd4c6e469bd7212c9328ed9899f7c54e Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 2 Nov 2018 12:10:35 -0700
Subject: [PATCH 1155/2720] Condition the latents z_{t} at time-step t z_{0} in
 addition to the latents at previous "hparams.num_cond_latents" time-steps.
 This is done via hparams.cond_first_frame.

PiperOrigin-RevId: 219839290
---
 tensor2tensor/models/research/glow_ops.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 75a06a0e3..2ba0e4636 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -66,9 +66,12 @@ def check_cond_latents(cond_latents, hparams):
     return
   if not isinstance(cond_latents[0], list):
     cond_latents = [cond_latents]
-  if len(cond_latents) != hparams.num_cond_latents:
+  exp_num_latents = hparams.num_cond_latents
+  if hparams.latent_dist_encoder == "conv_net":
+    exp_num_latents += int(hparams.cond_first_frame)
+  if len(cond_latents) != exp_num_latents:
     raise ValueError("Expected number of cond_latents: %d, got %d" %
-                     (hparams.num_cond_latents, len(cond_latents)))
+                     (exp_num_latents, len(cond_latents)))
   for cond_latent in cond_latents:
     if len(cond_latent) != hparams.n_levels - 1:
       raise ValueError("Expected level_latents to be %d, got %d" %

From b4de3a245e8e01e494a74e37c3a7779ffcb2a2de Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 2 Nov 2018 13:05:36 -0700
Subject: [PATCH 1156/2720] Add Wikipedia LM datasets for De, Fr, Ro, and a
 4-languages one.

PiperOrigin-RevId: 219847822
---
 tensor2tensor/data_generators/wiki_lm.py | 218 +++++++++++++++++++----
 1 file changed, 180 insertions(+), 38 deletions(-)

diff --git a/tensor2tensor/data_generators/wiki_lm.py b/tensor2tensor/data_generators/wiki_lm.py
index dc07ed82d..ba79f214b 100644
--- a/tensor2tensor/data_generators/wiki_lm.py
+++ b/tensor2tensor/data_generators/wiki_lm.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import os
+import six
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -29,56 +30,197 @@
 import tensorflow as tf
 
 
+def concat_generator(filename, up_threshold, low_threshold=10):
+  """Generate concatenated lines from file upto up_threshold characters."""
+  txt = ""
+  for line in tf.gfile.Open(filename):
+    line = line.strip()
+    if len(txt) + len(line) > up_threshold:
+      ret = txt
+      txt = ""
+      # We don't yield very short long parts to prevent noisy examples.
+      if len(ret) > low_threshold and len(ret) < up_threshold:
+        yield {"targets": ret}
+
+    if not txt:
+      txt = line
+    else:
+      txt = " ".join([txt, line])
+
+
+def mix_generators(generator_list):
+  """Given python generators, generate from one, then from another, etc."""
+  i = 0
+  l = len(generator_list)
+  stopiters_seen = 0
+  while stopiters_seen <= l:
+    try:
+      yield six.next(generator_list[i % l])
+      i += 1
+      stopiters_seen = 0
+    except StopIteration:
+      i += 1
+      stopiters_seen += 1
+
+
+# File names and Google drive ids for the training/eval/test Wikipedia data.
+_EN_TRAIN_NAME_ID = ("enwiki_train.txt.gz", "1-l02fI15ieMIZk8EnXhzhsvuEYRoznZ8")
+_EN_EVAL_NAME_ID = ("enwiki_eval.txt.gz", "1odhDxWKtAPKXwxRw1KCrmlrVewxdXYq7")
+_EN_TEST_NAME_ID = ("enwiki_test.txt.gz", "1i1Bg6XqvdRl1LuOiIWbg7ww8Y02Ip5VK")
+
+_DE_TRAIN_NAME_ID = ("dewiki_train.txt.gz", "1FzEwoPonw9xlwX34vLPFInUF8F4X5yJy")
+_DE_EVAL_NAME_ID = ("dewiki_eval.txt.gz", "1EKwRRPHyWny0RJ-aqSGMcNfjAlzFl51B")
+_DE_TEST_NAME_ID = ("dewiki_test.txt.gz", "1Kr13Y7y_OD3JtUM9riXpFQP9UiHDkcFY")
+
+_FR_TRAIN_NAME_ID = ("frwiki_train.txt.gz", "1etUIEZxMQKORwLGkssE5wlfCxxkeo8WV")
+_FR_EVAL_NAME_ID = ("frwiki_eval.txt.gz", "13qrR5ZnHRgIMdcURVpixKL9gTO23GcPc")
+_FR_TEST_NAME_ID = ("frwiki_test.txt.gz", "1mQpHRkAV9KXt68de69RwR8dkDi8EEusV")
+
+_RO_TRAIN_NAME_ID = ("rowiki_train.txt.gz", "1wUJTEAlQeDcAwFnBxa8PzE-DCiXSU_W7")
+_RO_EVAL_NAME_ID = ("rowiki_eval.txt.gz", "1uIPy2ZgkyArPy_gnsILENjgv4QQmSKtx")
+_RO_TEST_NAME_ID = ("rowiki_test.txt.gz", "1kphjN4jXTbw8HyRYKaRE2zY4D7Fr-p7-")
+
+
 @registry.register_problem
-class LanguagemodelWiki32k(text_problems.Text2SelfProblem):
-  """A language model on the untokenized wikipedia corpus."""
+class LanguagemodelEnWiki32k(text_problems.Text2SelfProblem):
+  """A language model on the untokenized wikipedia corpus, English."""
 
-  # File names and Google drive ids for the training/dev/test data.
-  train_name_id = ("wiki_train.txt.gz", "1-l02fI15ieMIZk8EnXhzhsvuEYRoznZ8")
-  dev_name_id = ("wiki_dev.txt.gz", "1odhDxWKtAPKXwxRw1KCrmlrVewxdXYq7")
-  test_name_id = ("wiki_test.txt.gz", "1i1Bg6XqvdRl1LuOiIWbg7ww8Y02Ip5VK")
+  train_names_ids = [_EN_TRAIN_NAME_ID]
+  eval_names_ids = [_EN_EVAL_NAME_ID]
+  test_names_ids = [_EN_TEST_NAME_ID]
 
   @property
   def approx_vocab_size(self):
-    return 2**15  # 32768
+    return 32000
 
   @property
   def max_samples_for_vocab(self):
-    return 63000
+    return 128000
+
+  @property
+  def combine_characters_threshold(self):
+    """Threshold for upto how many characters to combine in examples."""
+    return 512*8  # So we should have 512 tokens on average, maybe more.
 
   def is_generate_per_split(self):
     return True
 
-  def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    # Thresholds in the number of characters for LM examples
-    lo_thresh = 10
-    up_thresh = 256*8
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each."""
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 100,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }, {
+        "split": problem.DatasetSplit.TEST,
+        "shards": 1,
+    }]
 
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    """Generate samples."""
     if dataset_split == problem.DatasetSplit.TRAIN:
-      (fname, fid) = self.train_name_id
+      file_names_ids = self.train_names_ids
+    elif dataset_split == problem.DatasetSplit.TEST:
+      file_names_ids = self.test_names_ids
     else:
-      (fname, fid) = self.dev_name_id
-
-    wikifiles = []
-    url = "https://drive.google.com/uc?export=download&id=" + fid
-    download_path = generator_utils.maybe_download_from_drive(
-        tmp_dir, fname, url)
-    wiki_file = os.path.join(tmp_dir, fname[:-3])
-    if not tf.gfile.Exists(wiki_file):
-      generator_utils.gunzip_file(download_path, wiki_file)
-    wikifiles.append(wiki_file)
-
-    txt = ""
-    for wiki_file in wikifiles:
-      for line in tf.gfile.Open(wiki_file):
-        line = line.strip()
-        if len(txt) + len(line) > up_thresh:
-          ret = txt
-          txt = ""
-          if len(ret) > lo_thresh and len(ret) < up_thresh:
-            yield {"targets": ret}
-
-        if not txt:
-          txt = line
-        else:
-          txt = " ".join([txt, line])
+      file_names_ids = self.eval_names_ids
+
+    wiki_generators = []
+    for (fname, fid) in file_names_ids:
+      url = "https://drive.google.com/uc?export=download&id=" + fid
+      download_path = generator_utils.maybe_download_from_drive(
+          tmp_dir, fname, url)
+      wiki_file = os.path.join(tmp_dir, fname[:-3])
+      if not tf.gfile.Exists(wiki_file):
+        generator_utils.gunzip_file(download_path, wiki_file)
+      wiki_generators.append(
+          concat_generator(wiki_file, self.combine_characters_threshold))
+
+    for example in mix_generators(wiki_generators):
+      yield example
+
+
+@registry.register_problem
+class LanguagemodelEnWiki64k(LanguagemodelEnWiki32k):
+  """As above, with 64k vocabulary."""
+
+  @property
+  def approx_vocab_size(self):
+    return 64000
+
+
+@registry.register_problem
+class LanguagemodelDeWiki32k(LanguagemodelEnWiki32k):
+  """A language model on the untokenized wikipedia corpus, German."""
+
+  train_names_ids = [_DE_TRAIN_NAME_ID]
+  eval_names_ids = [_DE_EVAL_NAME_ID]
+  test_names_ids = [_DE_TEST_NAME_ID]
+
+
+@registry.register_problem
+class LanguagemodelDeWiki64k(LanguagemodelDeWiki32k):
+  """As above, with 64k vocabulary."""
+
+  @property
+  def approx_vocab_size(self):
+    return 64000
+
+
+@registry.register_problem
+class LanguagemodelFrWiki32k(LanguagemodelEnWiki32k):
+  """A language model on the untokenized wikipedia corpus, French."""
+
+  train_names_ids = [_FR_TRAIN_NAME_ID]
+  eval_names_ids = [_FR_EVAL_NAME_ID]
+  test_names_ids = [_FR_TEST_NAME_ID]
+
+
+@registry.register_problem
+class LanguagemodelFrWiki64k(LanguagemodelFrWiki32k):
+  """As above, with 64k vocabulary."""
+
+  @property
+  def approx_vocab_size(self):
+    return 64000
+
+
+@registry.register_problem
+class LanguagemodelRoWiki32k(LanguagemodelEnWiki32k):
+  """A language model on the untokenized wikipedia corpus, Romanian."""
+
+  train_names_ids = [_RO_TRAIN_NAME_ID]
+  eval_names_ids = [_RO_EVAL_NAME_ID]
+  test_names_ids = [_RO_TEST_NAME_ID]
+
+
+@registry.register_problem
+class LanguagemodelRoWiki64k(LanguagemodelRoWiki32k):
+  """As above, with 64k vocabulary."""
+
+  @property
+  def approx_vocab_size(self):
+    return 64000
+
+
+@registry.register_problem
+class LanguagemodelDeEnFrRoWiki64k(LanguagemodelEnWiki32k):
+  """A language model on untokenized Wikipedia, 4 languages together."""
+
+  train_names_ids = [_DE_TRAIN_NAME_ID, _FR_TRAIN_NAME_ID,
+                     _EN_TRAIN_NAME_ID, _RO_TRAIN_NAME_ID]
+  eval_names_ids = [_DE_EVAL_NAME_ID, _FR_EVAL_NAME_ID,
+                    _EN_EVAL_NAME_ID, _RO_EVAL_NAME_ID]
+  test_names_ids = [_DE_TEST_NAME_ID, _FR_TEST_NAME_ID,
+                    _EN_TEST_NAME_ID, _RO_TEST_NAME_ID]
+
+  @property
+  def approx_vocab_size(self):
+    return 64000
+
+  @property
+  def max_samples_for_vocab(self):
+    return 256000  # Samples are intertwined, take more to cover 4 languages.

From 8f137822023f187441056f4b52ecc3736bb5b916 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 2 Nov 2018 14:07:13 -0700
Subject: [PATCH 1157/2720] Fix JSON serializability of modality. Otherwise it
 fails with

PiperOrigin-RevId: 219858117
---
 tensor2tensor/bin/t2t_trainer.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index ba1651f3d..0693dd202 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -19,6 +19,7 @@
 from __future__ import print_function
 
 import contextlib
+import copy
 import os
 import sys
 from tensor2tensor import models  # pylint: disable=unused-import
@@ -325,9 +326,13 @@ def save_metadata(hparams):
       f.write(t2t_flags_str)
 
   # Save hparams as hparams.json
+  new_hparams = copy.deepcopy(hparams)
+
+  # Modality class is not JSON serializable so remove.
+  new_hparams.del_hparam("modality")
   hparams_fname = os.path.join(output_dir, "hparams.json")
   with tf.gfile.Open(hparams_fname, "w") as f:
-    f.write(hparams.to_json(indent=0, sort_keys=True))
+    f.write(new_hparams.to_json(indent=0, sort_keys=True))
 
 
 def execute_schedule(exp):

From dc162fe91e4f9887a1a6213df5d73d5adbe22a3e Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 2 Nov 2018 23:13:37 +0100
Subject: [PATCH 1158/2720] Remove environment_spec (#1198)

* Remove environment_spec

* Make frame stack size a hparam

* Pylint
---
 tensor2tensor/models/research/rl.py           | 95 +++++++------------
 tensor2tensor/rl/collect.py                   | 27 +++---
 tensor2tensor/rl/envs/simulated_batch_env.py  | 59 +++++++-----
 .../rl/envs/simulated_batch_gym_env.py        | 56 +++++------
 tensor2tensor/rl/ppo.py                       | 10 +-
 tensor2tensor/rl/rl_trainer_lib.py            | 28 ++----
 tensor2tensor/rl/trainer_model_based.py       | 72 +++++++-------
 .../rl/trainer_model_based_params.py          |  2 +
 tensor2tensor/rl/trainer_model_free.py        |  5 +-
 9 files changed, 158 insertions(+), 196 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 5d0070962..9caa3763f 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -19,13 +19,14 @@
 import functools
 import operator
 import gym
-import six
 
 from tensor2tensor.data_generators import gym_env
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
-from tensor2tensor.rl.envs import tf_atari_wrappers
+from tensor2tensor.rl.envs.py_func_batch_env import PyFuncBatchEnv
+from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
+from tensor2tensor.rl.envs.simulated_batch_gym_env import SimulatedBatchGymEnv
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -112,81 +113,45 @@ def ppo_atari_base():
   return hparams
 
 
-def simple_gym_spec(env):
-  """Parameters of environment specification."""
-  standard_wrappers = None
-  env_lambda = None
-  if isinstance(env, str):
-    env_lambda = lambda: gym.make(env)
-  if callable(env):
-    env_lambda = env
-  assert env_lambda is not None, "Unknown specification of environment"
-
-  return tf.contrib.training.HParams(env_lambda=env_lambda,
-                                     wrappers=standard_wrappers,
-                                     simulated_env=False)
-
-
-def standard_atari_env_spec(env=None, simulated=False):
-  """Parameters of environment specification."""
-  standard_wrappers = [
-      (tf_atari_wrappers.StackWrapper, {"history": 4})
-  ]
-  env_spec = tf.contrib.training.HParams(
-      wrappers=standard_wrappers,
-      simulated_env=simulated,
-      reward_range=env.reward_range,
-      observation_space=env.observation_space,
-      action_space=env.action_space,
-      force_beginning_resets=simulated
-  )
-  if not simulated:
-    env_spec.add_hparam("env", env)
-  return env_spec
-
+def make_real_env_fn(env):
+  """Creates a function returning a given real env, in or out of graph.
 
-def standard_atari_env_simulated_spec(real_env, **kwargs):
-  """Spec."""
-  env_spec = standard_atari_env_spec(real_env, simulated=True)
-  for (name, value) in six.iteritems(kwargs):
-    env_spec.add_hparam(name, value)
-  return env_spec
+  Args:
+    env: Environment to return from the function.
 
+  Returns:
+    Function in_graph -> env.
+  """
+  return lambda in_graph: PyFuncBatchEnv(env) if in_graph else env
 
-def standard_atari_env_eval_spec(*args, **kwargs):
-  """Parameters of environment specification for eval."""
-  return standard_atari_env_spec(*args, **kwargs)
 
+def make_simulated_env_fn(**env_kwargs):
+  """Returns a function creating a simulated env, in or out of graph.
 
-def standard_atari_ae_env_spec(env, ae_hparams_set):
-  """Parameters of environment specification."""
-  standard_wrappers = [[tf_atari_wrappers.AutoencoderWrapper,
-                        {"ae_hparams_set": ae_hparams_set}],
-                       [tf_atari_wrappers.StackWrapper, {"history": 4}]]
-  env_lambda = None
-  if isinstance(env, str):
-    env_lambda = lambda: gym.make(env)
-  if callable(env):
-    env_lambda = env
-  assert env is not None, "Unknown specification of environment"
+  Args:
+    **env_kwargs: kwargs to pass to the simulated env constructor.
 
-  return tf.contrib.training.HParams(env_lambda=env_lambda,
-                                     wrappers=standard_wrappers,
-                                     simulated_env=False)
+  Returns:
+    Function in_graph -> env.
+  """
+  def env_fn(in_graph):
+    class_ = SimulatedBatchEnv if in_graph else SimulatedBatchGymEnv
+    return class_(**env_kwargs)
+  return env_fn
 
 
-def get_policy(observations, hparams):
+def get_policy(observations, hparams, action_space):
   """Get a policy network.
 
   Args:
     observations: Tensor with observations
     hparams: parameters
+    action_space: action space
 
   Returns:
     Tensor with policy and value function output
   """
   policy_network_lambda = hparams.policy_network
-  action_space = hparams.environment_spec.action_space
   return policy_network_lambda(action_space, hparams, observations)
 
 
@@ -219,12 +184,16 @@ def pong_model_free():
       optimization_batch_size=4,
       clipping_coef=0.2,
       value_loss_coef=1,
-      save_models_every_epochs=False)
+      save_models_every_epochs=False,
+      frame_stack_size=4,
+      force_beginning_resets=False,
+  )
   env = gym_env.T2TGymEnv("PongNoFrameskip-v4", batch_size=2)
   env.start_new_epoch(0)
-  hparams.add_hparam("environment_spec", standard_atari_env_spec(env))
-  hparams.add_hparam(
-      "environment_eval_spec", standard_atari_env_eval_spec(env))
+  hparams.add_hparam("env_fn", make_real_env_fn(env))
+  eval_env = gym_env.T2TGymEnv("PongNoFrameskip-v4", batch_size=2)
+  eval_env.start_new_epoch(0)
+  hparams.add_hparam("eval_env_fn", make_real_env_fn(eval_env))
   return hparams
 
 
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index 198d4e31b..b282284a8 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -22,9 +22,7 @@
 import copy
 
 from tensor2tensor.models.research.rl import get_policy
-from tensor2tensor.rl.envs.py_func_batch_env import PyFuncBatchEnv
-from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
-from tensor2tensor.rl.envs.tf_atari_wrappers import WrapperBase
+from tensor2tensor.rl.envs.tf_atari_wrappers import StackWrapper, WrapperBase
 
 import tensorflow as tf
 
@@ -85,32 +83,29 @@ def simulate(self, action):
       return tf.identity(reward), tf.identity(done)
 
 
-def define_collect(hparams, scope):
+def define_collect(batch_env, hparams, scope):
   """Collect trajectories.
 
   Args:
+    batch_env: Batch environment.
     hparams: HParams.
     scope: var scope.
 
   Returns:
-    Returns memory (observtions, rewards, dones, actions,
+    Returns memory (observations, rewards, dones, actions,
     pdfs, values_functions)
     containing a rollout of environment from nested wrapped structure.
   """
 
   to_initialize = []
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
-    environment_spec = hparams.environment_spec
     num_agents = hparams.num_agents
-    if environment_spec.simulated_env:
-      batch_env = SimulatedBatchEnv(environment_spec, num_agents)
-    else:
-      batch_env = PyFuncBatchEnv(environment_spec.env)
 
     to_initialize.append(batch_env)
-    environment_wrappers = environment_spec.wrappers
-    wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
-    wrappers.append((_MemoryWrapper, {}))
+    wrappers = [
+        (StackWrapper, {"history": hparams.frame_stack_size}),
+        (_MemoryWrapper, {})
+    ]
     rollout_metadata = None
     speculum = None
     for w in wrappers:
@@ -142,7 +137,7 @@ def initialization_lambda(sess):
     zeros_tensor = tf.zeros(len(batch_env))
 
   force_beginning_resets = tf.convert_to_tensor(
-      environment_spec.force_beginning_resets
+      hparams.force_beginning_resets
   )
 
   def reset_ops_group():
@@ -168,7 +163,9 @@ def step(index, scores_sum, scores_num):
 
       def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
         """Step of the environment."""
-        actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
+        actor_critic = get_policy(
+            tf.expand_dims(obs_copy, 0), hparams, batch_env.action_space
+        )
         policy = actor_critic.policy
         action = hparams.policy_to_actions_lambda(policy)
 
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index b61a1b6c9..e0fad49db 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -37,12 +37,12 @@ class HistoryBuffer(object):
   """History Buffer."""
 
   def __init__(self, initial_frame_chooser, observ_shape, observ_dtype,
-               num_initial_frames, length):
-    self.length = length
+               num_initial_frames, batch_size):
+    self.batch_size = batch_size
     self._observ_dtype = observ_dtype
-    initial_shape = (length, num_initial_frames) + observ_shape
+    initial_shape = (batch_size, num_initial_frames) + observ_shape
     self._initial_frames = tf.py_func(
-        initial_frame_chooser, [tf.constant(length)], observ_dtype
+        initial_frame_chooser, [tf.constant(batch_size)], observ_dtype
     )
     self._initial_frames.set_shape(initial_shape)
     self._history_buff = tf.Variable(tf.zeros(initial_shape, observ_dtype),
@@ -91,46 +91,53 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   flags are held in according variables.
   """
 
-  def __init__(self, environment_spec, length):
+  def __init__(
+      self, reward_range, observation_space, action_space, frame_stack_size,
+      initial_frame_chooser, batch_size, model_name, model_hparams, model_dir,
+      intrinsic_reward_scale=0.0
+  ):
     """Batch of environments inside the TensorFlow graph."""
-    super(SimulatedBatchEnv, self).__init__(
-        environment_spec.observation_space, environment_spec.action_space
-    )
+    super(SimulatedBatchEnv, self).__init__(observation_space, action_space)
 
-    self.length = length
-    self._min_reward = environment_spec.reward_range[0]
-    self._num_frames = environment_spec.video_num_input_frames
-    self._intrinsic_reward_scale = environment_spec.intrinsic_reward_scale
+    self.batch_size = batch_size
+    self._min_reward = reward_range[0]
+    self._num_frames = frame_stack_size
+    self._intrinsic_reward_scale = intrinsic_reward_scale
 
-    model_hparams = copy.copy(environment_spec.model_hparams)
-    problem = DummyWorldModelProblem(
-        environment_spec.action_space, environment_spec.reward_range
-    )
+    model_hparams = copy.copy(model_hparams)
+    problem = DummyWorldModelProblem(action_space, reward_range)
     trainer_lib.add_problem_hparams(model_hparams, problem)
     model_hparams.force_full_predict = True
-    self._model = registry.model(environment_spec.model_name)(
+    self._model = registry.model(model_name)(
         model_hparams, tf.estimator.ModeKeys.PREDICT
     )
 
     self.history_buffer = HistoryBuffer(
-        environment_spec.initial_frame_chooser, self.observ_shape,
-        self.observ_dtype, self._num_frames, self.length
+        initial_frame_chooser, self.observ_shape, self.observ_dtype,
+        self._num_frames, self.batch_size
     )
 
     self._observ = tf.Variable(
-        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
-        trainable=False)
+        tf.zeros((batch_size,) + self.observ_shape, self.observ_dtype),
+        trainable=False
+    )
+
+    self._model_dir = model_dir
 
   def initialize(self, sess):
-    # Currently not needed. Keeping it just in case.
-    pass
+    model_loader = tf.train.Saver(
+        var_list=tf.global_variables(scope="next_frame*")  # pylint:disable=unexpected-keyword-arg
+    )
+    trainer_lib.restore_checkpoint(
+        self._model_dir, saver=model_loader, sess=sess, must_restore=True
+    )
 
   def __str__(self):
     return "SimulatedEnv"
 
   def __len__(self):
     """Number of combined environments."""
-    return self.length
+    return self.batch_size
 
   def simulate(self, action):
     with tf.name_scope("environment/simulate"):
@@ -149,7 +156,7 @@ def simulate(self, action):
                        self.observ_dtype)
 
       reward = tf.to_float(model_output["target_reward"])
-      reward = tf.reshape(reward, shape=(self.length,)) + self._min_reward
+      reward = tf.reshape(reward, shape=(self.batch_size,)) + self._min_reward
 
       if self._intrinsic_reward_scale:
         # Use the model's uncertainty about its prediction as an intrinsic
@@ -168,7 +175,7 @@ def simulate(self, action):
                                       summarize=8)
         reward += uncertainty_reward
 
-      done = tf.constant(False, tf.bool, shape=(self.length,))
+      done = tf.constant(False, tf.bool, shape=(self.batch_size,))
 
       with tf.control_dependencies([observ]):
         with tf.control_dependencies(
diff --git a/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index 4be71d4c8..3af091628 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -19,55 +19,42 @@
 from __future__ import division
 from __future__ import print_function
 
-import copy
-
 from gym import Env
 from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
-from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
 
 
+# TODO(koz4k): Unify interfaces of batch envs.
 class SimulatedBatchGymEnv(Env):
   """SimulatedBatchEnv in a Gym-like interface, environments are  batched."""
 
-  def __init__(self, environment_spec, batch_size,
-               model_dir=None, sess=None):
-    self.batch_size = batch_size
-
+  def __init__(self, *args, **kwargs):
     with tf.Graph().as_default():
-      self._batch_env = SimulatedBatchEnv(environment_spec,
-                                          self.batch_size)
-
-      self.action_space = self._batch_env.action_space
-      # TODO(kc): check for the stack wrapper and correct number of channels in
-      # observation_space
-      self.observation_space = self._batch_env.observ_space
-      self._sess = sess if sess is not None else tf.Session()
-      self._to_initialize = [self._batch_env]
-
-      environment_wrappers = environment_spec.wrappers
-      wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
-
-      for w in wrappers:
-        self._batch_env = w[0](self._batch_env, **w[1])
-        self._to_initialize.append(self._batch_env)
-
-      self._sess.run(tf.global_variables_initializer())
-      for wrapped_env in self._to_initialize:
-        wrapped_env.initialize(self._sess)
+      self._batch_env = SimulatedBatchEnv(*args, **kwargs)
 
-      self._actions_t = tf.placeholder(shape=(batch_size,), dtype=tf.int32)
+      self._actions_t = tf.placeholder(shape=(self.batch_size,), dtype=tf.int32)
       self._rewards_t, self._dones_t = self._batch_env.simulate(self._actions_t)
       self._obs_t = self._batch_env.observ
       self._reset_op = self._batch_env.reset(
-          tf.range(batch_size, dtype=tf.int32)
+          tf.range(self.batch_size, dtype=tf.int32)
       )
 
-      env_model_loader = tf.train.Saver(
-          var_list=tf.global_variables(scope="next_frame*"))  # pylint:disable=unexpected-keyword-arg
-      trainer_lib.restore_checkpoint(model_dir, saver=env_model_loader,
-                                     sess=self._sess, must_restore=True)
+      self._sess = tf.Session()
+      self._sess.run(tf.global_variables_initializer())
+      self._batch_env.initialize(self._sess)
+
+  @property
+  def batch_size(self):
+    return self._batch_env.batch_size
+
+  @property
+  def observation_space(self):
+    return self._batch_env.observ_space
+
+  @property
+  def action_space(self):
+    return self._batch_env.action_space
 
   def render(self, mode="human"):
     raise NotImplementedError()
@@ -86,3 +73,6 @@ def step(self, actions):
         [self._obs_t, self._rewards_t, self._dones_t],
         feed_dict={self._actions_t: actions})
     return obs, rewards, dones
+
+  def close(self):
+    self._sess.close()
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index f4d26346c..0a9ab0861 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -32,10 +32,10 @@ def get_optimiser(config):
   return config.optimizer(learning_rate=config.learning_rate)
 
 
-def define_ppo_step(data_points, optimizer, hparams):
+def define_ppo_step(data_points, optimizer, hparams, action_space):
   """Define ppo step."""
   observation, action, discounted_reward, norm_advantage, old_pdf = data_points
-  new_policy_dist, new_value, _ = get_policy(observation, hparams)
+  new_policy_dist, new_value, _ = get_policy(observation, hparams, action_space)
   new_pdf = new_policy_dist.prob(action)
 
   ratio = new_pdf / old_pdf
@@ -73,7 +73,7 @@ def define_ppo_step(data_points, optimizer, hparams):
     return [tf.identity(x) for x in losses + gradients_norms]
 
 
-def define_ppo_epoch(memory, hparams):
+def define_ppo_epoch(memory, hparams, action_space):
   """PPO epoch."""
   observation, reward, done, action, old_pdf, value = memory
 
@@ -119,7 +119,9 @@ def define_ppo_epoch(memory, hparams):
   with tf.control_dependencies([iterator.initializer]):
     ppo_step_rets = tf.scan(
         lambda a, i: add_lists_elementwise(  # pylint: disable=g-long-lambda
-            a, define_ppo_step(iterator.get_next(), optimizer, hparams)),
+            a, define_ppo_step(
+                iterator.get_next(), optimizer, hparams, action_space
+            )),
         tf.range(number_of_batches),
         [0., 0., 0., 0., 0., 0.],
         parallel_iterations=1)
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
index 30b3d30c9..6303f059e 100644
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ b/tensor2tensor/rl/rl_trainer_lib.py
@@ -39,10 +39,11 @@ def define_train(hparams):
   )
 
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+    train_env = hparams.env_fn(in_graph=True)
     memory, collect_summary, train_initialization = (
-        collect.define_collect(train_hparams, "ppo_train")
+        collect.define_collect(train_env, train_hparams, "ppo_train")
     )
-    ppo_summary = ppo.define_ppo_epoch(memory, hparams)
+    ppo_summary = ppo.define_ppo_epoch(memory, hparams, train_env.action_space)
     train_summary = tf.summary.merge([collect_summary, ppo_summary])
 
     if hparams.eval_every_epochs:
@@ -51,11 +52,11 @@ def define_train(hparams):
       eval_hparams.add_hparam(
           "policy_to_actions_lambda", lambda policy: policy.mode()
       )
-      eval_hparams.environment_spec = hparams.environment_eval_spec
+      eval_env = hparams.eval_env_fn(in_graph=True)
       eval_hparams.num_agents = hparams.num_eval_agents
 
       _, eval_collect_summary, eval_initialization = (
-          collect.define_collect(eval_hparams, "ppo_eval")
+          collect.define_collect(eval_env, eval_hparams, "ppo_eval")
       )
       return train_summary, eval_collect_summary, (train_initialization,
                                                    eval_initialization)
@@ -68,7 +69,7 @@ def train(hparams, event_dir=None, model_dir=None,
   """Train."""
   with tf.Graph().as_default():
     with tf.name_scope(name_scope):
-      train_summary_op, eval_summary_op, intializers = define_train(hparams)
+      train_summary_op, eval_summary_op, initializers = define_train(hparams)
       if event_dir:
         summary_writer = tf.summary.FileWriter(
             event_dir, graph=tf.get_default_graph(), flush_secs=60)
@@ -81,22 +82,10 @@ def train(hparams, event_dir=None, model_dir=None,
       else:
         model_saver = None
 
-      # TODO(piotrmilos): This should be refactored, possibly with
-      # handlers for each type of env
-      if hparams.environment_spec.simulated_env:
-        env_model_loader = tf.train.Saver(
-            tf.global_variables("next_frame*"))
-      else:
-        env_model_loader = None
-
       with tf.Session() as sess:
         sess.run(tf.global_variables_initializer())
-        for initializer in intializers:
+        for initializer in initializers:
           initializer(sess)
-        if env_model_loader:
-          trainer_lib.restore_checkpoint(
-              hparams.world_model_dir, env_model_loader, sess,
-              must_restore=True)
         start_step = 0
         if model_saver and restore_agent:
           start_step = trainer_lib.restore_checkpoint(
@@ -144,8 +133,9 @@ def evaluate(hparams, model_dir, name_scope="rl_eval"):
   hparams.add_hparam("eval_phase", True)
   with tf.Graph().as_default():
     with tf.name_scope(name_scope):
+      eval_env = hparams.env_fn(in_graph=True)
       (collect_memory, _, collect_init) = collect.define_collect(
-          hparams, "ppo_eval"
+          eval_env, hparams, "ppo_eval"
       )
       model_saver = tf.train.Saver(
           tf.global_variables(".*network_parameters.*")
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 972440b4a..4a2fab766 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -43,7 +43,6 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.rl import trainer_model_based_params
-from tensor2tensor.rl.envs.simulated_batch_gym_env import SimulatedBatchGymEnv
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -139,14 +138,20 @@ def choose_subsequence():
   return [choose_subsequence() for _ in range(num_subsequences)]
 
 
-def make_simulated_env_spec(real_env, hparams):
-  """Creates a simulated environment_spec."""
-  return rl.standard_atari_env_simulated_spec(
-      real_env, intrinsic_reward_scale=hparams.intrinsic_reward_scale,
+def make_simulated_env_fn(
+    real_env, hparams, batch_size, initial_frame_chooser, model_dir
+):
+  """Creates a simulated env_fn."""
+  return rl.make_simulated_env_fn(
+      reward_range=real_env.reward_range,
+      observation_space=real_env.observation_space,
+      action_space=real_env.action_space,
+      frame_stack_size=hparams.frame_stack_size,
+      initial_frame_chooser=initial_frame_chooser, batch_size=batch_size,
       model_name=hparams.generative_model,
       model_hparams=trainer_lib.create_hparams(hparams.generative_model_params),
-      # Hardcoded for now. TODO(koz4k): Make it a hparam.
-      video_num_input_frames=4, video_num_target_frames=1
+      model_dir=model_dir,
+      intrinsic_reward_scale=hparams.intrinsic_reward_scale,
   )
 
 
@@ -185,28 +190,25 @@ def train_agent(real_env, agent_model_dir, event_dir, world_model_dir, data_dir,
   ppo_hparams.epochs_num = completed_ppo_epochs_num
 
   ppo_hparams.save_models_every_epochs = 10
-  ppo_hparams.world_model_dir = world_model_dir
-
-  environment_spec = make_simulated_env_spec(real_env, hparams)
 
-  num_input_frames = environment_spec.video_num_input_frames
+  frame_stack_size = hparams.frame_stack_size
   initial_frame_rollouts = real_env.current_epoch_rollouts(
       split=tf.contrib.learn.ModeKeys.TRAIN,
-      minimal_rollout_frames=num_input_frames,
+      minimal_rollout_frames=frame_stack_size,
   )
   # TODO(koz4k): Move this to a different module.
   def initial_frame_chooser(batch_size):
     """Frame chooser."""
 
     deterministic_initial_frames =\
-        initial_frame_rollouts[0][:num_input_frames]
+        initial_frame_rollouts[0][:frame_stack_size]
     if not hparams.simulation_random_starts:
       # Deterministic starts: repeat first frames from the first rollout.
       initial_frames = [deterministic_initial_frames] * batch_size
     else:
       # Random starts: choose random initial frames from random rollouts.
       initial_frames = random_rollout_subsequences(
-          initial_frame_rollouts, batch_size, num_input_frames
+          initial_frame_rollouts, batch_size, frame_stack_size
       )
       if hparams.simulation_flip_first_random_for_beginning:
         # Flip first entry in the batch for deterministic initial frames.
@@ -217,9 +219,13 @@ def initial_frame_chooser(batch_size):
         for initial_frame_stack in initial_frames
     ])
 
-  environment_spec.add_hparam("initial_frame_chooser", initial_frame_chooser)
-
-  ppo_hparams.add_hparam("environment_spec", environment_spec)
+  env_fn = make_simulated_env_fn(
+      real_env, hparams, hparams.ppo_num_agents, initial_frame_chooser,
+      world_model_dir
+  )
+  ppo_hparams.add_hparam("env_fn", env_fn)
+  ppo_hparams.add_hparam("force_beginning_resets", True)
+  ppo_hparams.add_hparam("frame_stack_size", frame_stack_size)
 
   rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
                        name_scope="ppo_sim%d" % (epoch + 1))
@@ -250,9 +256,9 @@ def train_agent_real_env(
   # But we need to save at the last step, so we set it very high.
   ppo_hparams.save_models_every_epochs = 1000000
 
-  environment_spec = rl.standard_atari_env_spec(env)
-
-  ppo_hparams.add_hparam("environment_spec", environment_spec)
+  ppo_hparams.add_hparam("env_fn", rl.make_real_env_fn(env))
+  ppo_hparams.add_hparam("force_beginning_resets", False)
+  ppo_hparams.add_hparam("frame_stack_size", hparams.frame_stack_size)
 
   rl_trainer_lib.train(ppo_hparams, event_dir + "real", agent_model_dir,
                        name_scope="ppo_real%d" % (epoch + 1))
@@ -312,11 +318,13 @@ def evaluate_single_config(hparams, agent_model_dir):
   eval_hparams = trainer_lib.create_hparams(hparams.ppo_params)
   eval_hparams.num_agents = hparams.num_agents
   env = setup_env(hparams, batch_size=hparams.num_agents)
-  environment_spec = rl.standard_atari_env_spec(env)
-  eval_hparams.add_hparam("environment_spec", environment_spec)
+  env_fn = rl.make_real_env_fn(env)
+  eval_hparams.add_hparam("env_fn", env_fn)
   eval_hparams.add_hparam(
       "policy_to_actions_lambda", hparams.policy_to_actions_lambda
   )
+  eval_hparams.add_hparam("frame_stack_size", hparams.frame_stack_size)
+  eval_hparams.add_hparam("force_beginning_resets", False)
 
   env.start_new_epoch(0)
   rl_trainer_lib.evaluate(eval_hparams, agent_model_dir)
@@ -375,28 +383,26 @@ def compute_mean_reward(rollouts, clipped):
 
 def evaluate_world_model(real_env, hparams, world_model_dir, debug_video_path):
   """Evaluate the world model (reward accuracy)."""
-  environment_spec = make_simulated_env_spec(real_env, hparams)
-  environment_spec.wrappers = []
-
-  num_input_frames = environment_spec.video_num_input_frames
+  frame_stack_size = hparams.frame_stack_size
   rollout_subsequences = []
   def initial_frame_chooser(batch_size):
     assert batch_size == len(rollout_subsequences)
     return np.stack([
-        [frame.observation.decode() for frame in subsequence[:num_input_frames]]
+        [frame.observation.decode() for frame in subsequence[:frame_stack_size]]
         for subsequence in rollout_subsequences
     ])
-  environment_spec.add_hparam("initial_frame_chooser", initial_frame_chooser)
 
-  sim_env = SimulatedBatchGymEnv(
-      environment_spec, hparams.wm_eval_batch_size, world_model_dir
+  env_fn = make_simulated_env_fn(
+      real_env, hparams, hparams.wm_eval_batch_size, initial_frame_chooser,
+      world_model_dir
   )
+  sim_env = env_fn(in_graph=False)
   subsequence_length = int(
       max(hparams.wm_eval_rollout_ratios) * hparams.ppo_epoch_length
   )
   rollouts = real_env.current_epoch_rollouts(
       split=tf.contrib.learn.ModeKeys.EVAL,
-      minimal_rollout_frames=(subsequence_length + num_input_frames)
+      minimal_rollout_frames=(subsequence_length + frame_stack_size)
   )
 
   video_writer = common_video.WholeVideoWriter(
@@ -410,11 +416,11 @@ def initial_frame_chooser(batch_size):
   for _ in range(hparams.wm_eval_epochs_num):
     rollout_subsequences[:] = random_rollout_subsequences(
         rollouts, hparams.wm_eval_batch_size,
-        subsequence_length + num_input_frames
+        subsequence_length + frame_stack_size
     )
 
     eval_subsequences = [
-        subsequence[(num_input_frames - 1):]
+        subsequence[(frame_stack_size - 1):]
         for subsequence in rollout_subsequences
     ]
 
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index faf48a278..7dbc7216b 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -119,6 +119,8 @@ def rlmb_base():
       wm_eval_rollout_ratios=[0.25, 0.5, 1, 2],
       stop_loop_early=False,  # To speed-up tests.
       env_timesteps_limit=-1,  # Use default from gym.make()
+      # Number of last observations to feed to the agent and world model.
+      frame_stack_size=4,
   )
 
 
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index b6d105526..01fc2c0d8 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -52,12 +52,11 @@ def initialize_env_specs(hparams):
     env = gym_env.T2TGymEnv("{}Deterministic-v4".format(game_name),
                             batch_size=hparams.num_agents)
     env.start_new_epoch(0)
-    hparams.add_hparam("environment_spec", rl.standard_atari_env_spec(env))
+    hparams.add_hparam("env_fn", rl.make_real_env_fn(env))
     eval_env = gym_env.T2TGymEnv("{}Deterministic-v4".format(game_name),
                                  batch_size=hparams.num_eval_agents)
     eval_env.start_new_epoch(0)
-    hparams.add_hparam(
-        "environment_eval_spec", rl.standard_atari_env_eval_spec(eval_env))
+    hparams.add_hparam("eval_env_fn", rl.make_real_env_fn(eval_env))
   return hparams
 
 

From 8fd469ca5c41c32c1d167772bfb9f29b05e1ab86 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 2 Nov 2018 15:14:04 -0700
Subject: [PATCH 1159/2720] internal merge of PR #1198

PiperOrigin-RevId: 219869740
---
 tensor2tensor/rl/collect.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/collect.py
index b282284a8..db63e7432 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/collect.py
@@ -19,10 +19,9 @@
 from __future__ import print_function
 
 
-import copy
-
 from tensor2tensor.models.research.rl import get_policy
-from tensor2tensor.rl.envs.tf_atari_wrappers import StackWrapper, WrapperBase
+from tensor2tensor.rl.envs.tf_atari_wrappers import StackWrapper
+from tensor2tensor.rl.envs.tf_atari_wrappers import WrapperBase
 
 import tensorflow as tf
 

From fe8f33210a37b32f9e538c1063b295d5c2094b30 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 2 Nov 2018 15:29:35 -0700
Subject: [PATCH 1160/2720] Allow to use sampled latents as noise during
 training to adapt to it.

PiperOrigin-RevId: 219872378
---
 .../models/video/basic_stochastic.py          | 30 +++++++++++++++++--
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index d85f54b2d..b3078db6f 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -122,6 +122,21 @@ def add_bits(layer, bits):
       _, pred_loss = discretization.predict_bits_with_lstm(
           layer, hparams.latent_predictor_state_size, hparams.bottleneck_bits,
           target_bits=bits_clean)
+      # Mix bits from latent with predicted bits on forward pass as a noise.
+      if hparams.latent_rnn_max_sampling > 0.0:
+        with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+          bits_pred, _ = discretization.predict_bits_with_lstm(
+              layer, hparams.latent_predictor_state_size,
+              hparams.bottleneck_bits,
+              temperature=hparams.latent_predictor_temperature)
+          bits_pred = tf.expand_dims(tf.expand_dims(bits_pred, axis=1), axis=2)
+        # Be bits_pred on the forward pass but bits on the backward one.
+        bits_pred = bits_clean + tf.stop_gradient(bits_pred - bits_clean)
+        # Select which bits to take from pred sampling with bit_p probability.
+        which_bit = tf.random_uniform(common_layers.shape_list(bits))
+        bit_p = common_layers.inverse_lin_decay(hparams.latent_rnn_warmup_steps)
+        bit_p *= hparams.latent_rnn_max_sampling
+        bits = tf.where(which_bit < bit_p, bits_pred, bits)
 
     return add_bits(layer, bits), pred_loss
 
@@ -170,8 +185,8 @@ def next_frame_sampling_stochastic():
 def next_frame_basic_stochastic_discrete():
   """Basic 2-frame conv model with stochastic discrete latent."""
   hparams = basic_deterministic_params.next_frame_sampling()
-  hparams.batch_size = 2
-  hparams.video_num_target_frames = 16
+  hparams.batch_size = 4
+  hparams.video_num_target_frames = 6
   hparams.scheduled_sampling_mode = "prob_inverse_lin"
   hparams.scheduled_sampling_decay_steps = 40000
   hparams.scheduled_sampling_max_prob = 1.0
@@ -182,8 +197,10 @@ def next_frame_basic_stochastic_discrete():
   hparams.learning_rate_warmup_steps = 2000
   hparams.learning_rate_schedule = "linear_warmup * constant"
   hparams.add_hparam("bottleneck_bits", 256)
-  hparams.add_hparam("bottleneck_noise", 0.15)
+  hparams.add_hparam("bottleneck_noise", 0.1)
   hparams.add_hparam("discretize_warmup_steps", 40000)
+  hparams.add_hparam("latent_rnn_warmup_steps", 40000)
+  hparams.add_hparam("latent_rnn_max_sampling", 0.7)
   hparams.add_hparam("full_latent_tower", False)
   hparams.add_hparam("latent_predictor_state_size", 128)
   hparams.add_hparam("latent_predictor_temperature", 0.5)
@@ -201,3 +218,10 @@ def next_frame_stochastic_discrete_range(rhp):
   rhp.set_discrete("bottleneck_bits", [32, 64, 128, 256])
   rhp.set_discrete("video_num_target_frames", [4])
   rhp.set_float("bottleneck_noise", 0.0, 0.2)
+
+
+@registry.register_ranged_hparams
+def next_frame_stochastic_discrete_latent_range(rhp):
+  rhp.set_float("latent_rnn_max_sampling", 0.1, 0.9)
+  rhp.set_float("latent_predictor_temperature", 0.1, 1.2)
+  rhp.set_float("dropout", 0.1, 0.4)

From 60e7186a061572551a5e3446d6e7764f657df77e Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 2 Nov 2018 16:22:58 -0700
Subject: [PATCH 1161/2720] Register optimal hparams for the SAVP model for the
 vae-only and gan-only versions on the bair robot pushing dataset.

PiperOrigin-RevId: 219880235
---
 tensor2tensor/models/video/savp_params.py | 28 +++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index 6acee08c4..e3452f630 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -43,4 +43,32 @@ def next_frame_savp():
   hparams.latent_loss_multiplier_schedule = "linear"
   hparams.upsample_method = "bilinear_upsample_conv"
   hparams.internal_loss = False
+  hparams.reward_prediction = False
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_savp_vae():
+  """SAVP - VAE only model."""
+  hparams = next_frame_savp()
+  hparams.use_vae = True
+  hparams.use_gan = False
+  hparams.latent_loss_multiplier = 1e-3
+  hparams.latent_loss_multiplier_schedule = "linear_anneal"
+  hparams.anneal_end = 100000
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_savp_gan():
+  """SAVP - GAN only model."""
+  hparams = next_frame_savp()
+  hparams.use_gan = True
+  hparams.use_vae = False
+  hparams.gan_loss_multiplier = 0.001
+  hparams.optimizer_adam_beta1 = 0.5
+  hparams.learning_rate_constant = 2e-4
+  hparams.gan_loss = "cross_entropy"
+  hparams.learning_rate_decay_steps = 100000
+  hparams.learning_rate_schedule = "constant*linear_decay"
   return hparams

From 8eabd98ac036f955a4a3e433e6f7f5665edf2cff Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 2 Nov 2018 20:06:39 -0700
Subject: [PATCH 1162/2720] Explicitly set the num_iterations_1st_stage to be 0
 and the num_iterations_2nd_stage to 50000 in savp_params

PiperOrigin-RevId: 219899601
---
 tensor2tensor/models/video/savp_params.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index e3452f630..4f04fbcdd 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -44,6 +44,9 @@ def next_frame_savp():
   hparams.upsample_method = "bilinear_upsample_conv"
   hparams.internal_loss = False
   hparams.reward_prediction = False
+  hparams.anneal_end = 100000
+  hparams.num_iterations_1st_stage = 0
+  hparams.num_iterations_2nd_stage = 50000
   return hparams
 
 
@@ -55,7 +58,6 @@ def next_frame_savp_vae():
   hparams.use_gan = False
   hparams.latent_loss_multiplier = 1e-3
   hparams.latent_loss_multiplier_schedule = "linear_anneal"
-  hparams.anneal_end = 100000
   return hparams
 
 

From f6ea7344ac26fd08e94c96776baa76bd663638c2 Mon Sep 17 00:00:00 2001
From: Smit Hinsu <hinsu@google.com>
Date: Fri, 2 Nov 2018 21:42:56 -0700
Subject: [PATCH 1163/2720] Add an option to set XLA global_git_level in
 Tensor2Tensor

Currently, DistributionStrategy is not yet supported with xla.compile so this
option would allow use of XLA while using DistributionStrategy.

PiperOrigin-RevId: 219905423
---
 tensor2tensor/bin/t2t_trainer.py   | 6 +++++-
 tensor2tensor/utils/trainer_lib.py | 7 ++++++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 0693dd202..413afae9a 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -53,7 +53,10 @@
 flags.DEFINE_bool("use_tpu", False, "Whether to use TPU.")
 flags.DEFINE_bool("use_tpu_estimator", False, "Whether to use TPUEstimator. "
                   "This is always enabled when use_tpu is True.")
-flags.DEFINE_bool("xla_compile", False, "Whether to use XLA to compile graph.")
+flags.DEFINE_bool("xla_compile", False,
+                  "Whether to use XLA to compile model_fn.")
+flags.DEFINE_integer("xla_jit_level", -1,
+                     "GlobalJitLevel to use while compiling the full graph.")
 flags.DEFINE_integer("tpu_infeed_sleep_secs", None,
                      "How long to sleep the infeed thread.")
 flags.DEFINE_bool("generate_data", False, "Generate data before training?")
@@ -243,6 +246,7 @@ def create_run_config(hp, output_dir=None):
       enable_graph_rewriter=FLAGS.enable_graph_rewriter,
       use_tpu=FLAGS.use_tpu,
       use_tpu_estimator=FLAGS.use_tpu_estimator,
+      xla_jit_level=FLAGS.xla_jit_level,
       schedule=FLAGS.schedule,
       no_data_parallelism=hp.no_data_parallelism,
       optionally_use_dist_strat=FLAGS.optionally_use_dist_strat,
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 8e4d3c2a5..2ef88660a 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -58,6 +58,7 @@ def create_session_config(log_device_placement=False,
                           enable_graph_rewriter=False,
                           gpu_mem_fraction=0.95,
                           use_tpu=False,
+                          xla_jit_level=tf.OptimizerOptions.OFF,
                           inter_op_parallelism_threads=0,
                           intra_op_parallelism_threads=0):
   """The TensorFlow Session config to use."""
@@ -71,7 +72,9 @@ def create_session_config(log_device_placement=False,
     else:
       graph_options = tf.GraphOptions(
           optimizer_options=tf.OptimizerOptions(
-              opt_level=tf.OptimizerOptions.L1, do_function_inlining=False))
+              opt_level=tf.OptimizerOptions.L1,
+              do_function_inlining=False,
+              global_jit_level=xla_jit_level))
 
   gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_mem_fraction)
 
@@ -162,6 +165,7 @@ def create_run_config(model_name,
                       tpu_infeed_sleep_secs=None,
                       use_tpu=False,
                       use_tpu_estimator=False,
+                      xla_jit_level=tf.OptimizerOptions.OFF,
                       inter_op_parallelism_threads=0,
                       log_step_count_steps=100,
                       intra_op_parallelism_threads=0,
@@ -173,6 +177,7 @@ def create_run_config(model_name,
       enable_graph_rewriter=enable_graph_rewriter,
       gpu_mem_fraction=gpu_mem_fraction,
       use_tpu=use_tpu,
+      xla_jit_level=xla_jit_level,
       inter_op_parallelism_threads=inter_op_parallelism_threads,
       intra_op_parallelism_threads=intra_op_parallelism_threads)
   run_config_args = {

From 45e8e0d07bcb16dde4dc4e47fb750e2fae82dc13 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 2 Nov 2018 22:31:29 -0700
Subject: [PATCH 1164/2720] Improve transformer 8x8 MLPerf compliance by
 decoding each checkpoint.

PiperOrigin-RevId: 219907623
---
 tensor2tensor/data_generators/translate.py | 13 ++----
 tensor2tensor/utils/decoding.py            |  1 +
 tensor2tensor/utils/trainer_lib.py         | 46 +++++++++++++++++++++-
 3 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 94c29c233..b9282e146 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -21,7 +21,6 @@
 
 import os
 import tarfile
-import numpy as np
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -79,14 +78,6 @@ def compute_bleu_summaries(hook_args):
     reference file and the translated file.
   """
   decode_hparams = hook_args.decode_hparams
-  estimator = hook_args.estimator
-  current_step = estimator.get_variable_value(tf.GraphKeys.GLOBAL_STEP)
-  has_iters = hasattr(decode_hparams, "iterations_per_loop")
-  if current_step and has_iters and decode_hparams.iterations_per_loop:
-    iterations_per_loop = decode_hparams.iterations_per_loop
-    current_epoch = np.asscalar(current_step) // iterations_per_loop
-  else:
-    current_epoch = 0
 
   if (decode_hparams.decode_reference is None or
       decode_hparams.decode_to_file is None):
@@ -98,12 +89,14 @@ def compute_bleu_summaries(hook_args):
   values.append(tf.Summary.Value(tag="BLEU", simple_value=bleu))
   tf.logging.info("%s: BLEU = %6.2f" % (decode_hparams.decode_to_file, bleu))
   if decode_hparams.mlperf_mode:
+    current_step = decode_hparams.mlperf_decode_step
     mlperf_log.transformer_print(
         key=mlperf_log.EVAL_TARGET, value=decode_hparams.mlperf_threshold)
     mlperf_log.transformer_print(
         key=mlperf_log.EVAL_ACCURACY,
         value={
-            "epoch": max(current_epoch - 1, 0),
+            "epoch": max(current_step // decode_hparams.iterations_per_loop - 1,
+                         0),
             "value": bleu
         })
     mlperf_log.transformer_print(key=mlperf_log.EVAL_STOP)
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index d3a86e91c..6a8e626bc 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -79,6 +79,7 @@ def decode_hparams(overrides=""):
       max_display_outputs=10,
       # Used for MLPerf compliance logging.
       mlperf_mode=False,
+      mlperf_decode_step=0.0,
       mlperf_threshold=25.0,
       mlperf_success=False)
   hp.parse(overrides)
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 2ef88660a..ee064adc3 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -54,6 +54,39 @@ def next_checkpoint(model_dir, timeout_mins=240):
     yield last_ckpt
 
 
+def next_undecoded_checkpoint(model_dir, timeout_mins=240):
+  """Yields successive checkpoints from model_dir."""
+  last_ckpt = None
+  last_step = 0
+  while True:
+    # Get the latest checkpoint.
+    last_ckpt = tf.contrib.training.wait_for_new_checkpoint(
+        model_dir, last_ckpt, seconds_to_sleep=60, timeout=60 * timeout_mins)
+    # Get all the checkpoint from the model dir.
+    ckpt_path = tf.train.get_checkpoint_state(model_dir)
+    all_model_checkpoint_paths = ckpt_path.all_model_checkpoint_paths
+    ckpt_step = np.inf
+    next_ckpt = None
+    # Find the next checkpoint to eval based on last_step.
+    for ckpt in all_model_checkpoint_paths:
+      step = int(os.path.basename(ckpt).split("-")[1])
+      if step > last_step and step < ckpt_step:
+        ckpt_step = step
+        next_ckpt = ckpt
+
+    # If all the checkpoints have been evaluated.
+    if last_ckpt is None and next_ckpt is None:
+      tf.logging.info(
+          "Eval timeout: no new checkpoints within %dm" % timeout_mins)
+      break
+
+    if next_ckpt is not None:
+      last_step = ckpt_step
+      last_ckpt = next_ckpt
+
+    yield last_ckpt
+
+
 def create_session_config(log_device_placement=False,
                           enable_graph_rewriter=False,
                           gpu_mem_fraction=0.95,
@@ -437,6 +470,8 @@ def train_eval_and_decode(self):
         self._hparams.problem = problem
         self._hparams.problem_hparams = p_hparams
       mlperf_log.transformer_print(key=mlperf_log.EVAL_START)
+      if self._decode_hparams.mlperf_mode:
+        self._decode_hparams.mlperf_decode_step = i
       self.decode(dataset_split=tf.estimator.ModeKeys.EVAL)
       d_hparams = self._decode_hparams
       if d_hparams.mlperf_mode and d_hparams.mlperf_success:
@@ -527,11 +562,20 @@ def continuous_decode_on_train_data(self):
 
   def continuous_decode_on_eval_data(self):
     """Decode from dataset on new checkpoint."""
-    for ckpt in next_checkpoint(self._hparams.model_dir):
+    if self._decode_hparams.mlperf_mode:
+      ckpt_generator = next_undecoded_checkpoint(self._hparams.model_dir)
+    else:
+      ckpt_generator = next_checkpoint(self._hparams.model_dir)
+
+    for ckpt in ckpt_generator:
       current_step = int(os.path.basename(ckpt).split("-")[1])
+      tf.logging.info("Decoding step %d" % current_step)
       # Skip checkpoint 0.
       if current_step == 0:
         continue
+      if self._decode_hparams.mlperf_mode:
+        self._decode_hparams.mlperf_decode_step = current_step
+
       mlperf_log.transformer_print(key=mlperf_log.EVAL_START)
       self.decode(dataset_split=tf.estimator.ModeKeys.EVAL)
       d_hparams = self._decode_hparams

From a739644cb77a21b74aaf5724d36bcf298372071a Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Sat, 3 Nov 2018 15:20:10 -0700
Subject: [PATCH 1165/2720] minor fix in rescaling the images from 0-1 to
 0-255.

PiperOrigin-RevId: 219953513
---
 tensor2tensor/layers/common_layers.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 1414631dc..3216aab59 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -236,6 +236,7 @@ def convert_real_to_rgb(x):
   """Conversion of real numbers to pixel values."""
   with tf.name_scope("real_to_rgb", values=[x]):
     x *= 255.0
+    x = tf.round(x)
     return x
 
 

From 3493d9c49cd294e76bd872d6b05ecef47ef9a9b8 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Sat, 3 Nov 2018 15:20:22 -0700
Subject: [PATCH 1166/2720] "sampling" the rgb images with L2 loss by
 truncating the float into a uint and back. this allows a lossless in and out
 graph calls.

PiperOrigin-RevId: 219953531
---
 tensor2tensor/models/video/base.py | 31 ++++++++++++++++--------------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index eb6d7cdad..ac7e6eb41 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -326,16 +326,23 @@ def get_sampled_frame(self, pred_frame):
       sampled frame.
 
     """
-    if not self.is_per_pixel_softmax:
-      return pred_frame
-    frame_shape = common_layers.shape_list(pred_frame)
-    target_shape = frame_shape[:-1] + [self.hparams.problem.num_channels]
-    sampled_frame = tf.reshape(pred_frame, target_shape + [256])
-    # TODO(lukaszkaiser): should this be argmax or real sampling.
-    sampled_frame = tf.argmax(sampled_frame, axis=-1)
-    sampled_frame = tf.to_float(sampled_frame)
-    # TODO(lukaszkaiser): this should be consistent with modality.bottom()
-    sampled_frame = common_layers.standardize_images(sampled_frame)
+    # TODO(lukaszkaiser): the logic below heavily depend on the current
+    # (a bit strange) video modalities - we should change that.
+
+    if self.is_per_pixel_softmax:
+      frame_shape = common_layers.shape_list(pred_frame)
+      target_shape = frame_shape[:-1] + [self.hparams.problem.num_channels]
+      sampled_frame = tf.reshape(pred_frame, target_shape + [256])
+      # TODO(lukaszkaiser): should this be argmax or real sampling.
+      sampled_frame = tf.argmax(sampled_frame, axis=-1)
+      sampled_frame = tf.to_float(sampled_frame)
+      # TODO(lukaszkaiser): this should be consistent with modality.bottom()
+      sampled_frame = common_layers.standardize_images(sampled_frame)
+    else:
+      x = common_layers.convert_real_to_rgb(pred_frame)
+      x = tf.cast(x, tf.uint8)
+      x = common_layers.convert_rgb_to_real(x)
+      return x
     return sampled_frame
 
   def __get_next_inputs(self, index, all_frames, all_actions, all_rewards):
@@ -440,10 +447,6 @@ def logits_to_samples(logits):
 
   def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
     """Main video processing function."""
-
-    # TODO(lukaszkaiser): the split axes and the argmax below heavily depend on
-    # using the default (a bit strange) video modality - we should change that.
-
     hparams = self.hparams
     all_frames_copy = [tf.identity(frame) for frame in all_frames]
     orig_frame_shape = common_layers.shape_list(all_frames[0])

From 83bb5211deaa4e264a0f5225e4fff4b289390256 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Sun, 4 Nov 2018 10:55:02 -0800
Subject: [PATCH 1167/2720] fixing a fundamental bug which forces recurrent
 models to use only ONE context frame at prediction time.

PiperOrigin-RevId: 220007728
---
 tensor2tensor/models/video/base.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index ac7e6eb41..2f593e8cc 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -486,15 +486,17 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
       sampled_frame = self.get_sampled_frame(res_frame)
       sampled_frames.append(sampled_frame)
 
-      if self.is_predicting:
+      # Check whether we are done with context frames or not
+      if self.is_recurrent_model:
+        done_warm_start = (i >= hparams.video_num_input_frames - 1)
+      else:
+        done_warm_start = True  # Always true for non-reccurent networks.
+
+      if self.is_predicting and done_warm_start:
         all_frames[target_index] = sampled_frame
 
       # Scheduled sampling during training.
       if self.is_training:
-        if self.is_recurrent_model:
-          done_warm_start = i >= hparams.video_num_input_frames - 1
-        else:
-          done_warm_start = True  # Always true for non-reccurent networks.
         groundtruth_items = [target_frame]
         generated_items = [sampled_frame]
         ss_frame, = self.get_scheduled_sample_inputs(

From 337c815960c038126154ddd6a7c5cdaba50dea43 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Sun, 4 Nov 2018 21:43:15 -0800
Subject: [PATCH 1168/2720] adding internal_state save feature.

PiperOrigin-RevId: 220041471
---
 tensor2tensor/models/video/base.py | 117 ++++++++++++++++++++---------
 tensor2tensor/models/video/sv2p.py |  36 +++++++++
 2 files changed, 118 insertions(+), 35 deletions(-)

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 2f593e8cc..ad0a23e43 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -35,6 +35,10 @@
 tfcl = tf.contrib.layers
 
 
+def flat_lists(list_of_lists):
+  return [x for l in list_of_lists for x in l]
+
+
 @registry.register_model
 class NextFrameBase(t2t_model.T2TModel):
   """Base class for next_frame models.
@@ -139,10 +143,30 @@ def is_recurrent_model(self):
     """
     raise NotImplementedError("Base video model.")
 
+  def init_internal_states(self):
+    """Allows a model to preserve its internal model across multiple runs.
+
+    This optional function is only useful for any model with internal states
+    (usually recurrent models) which need to preserve states after any call.
+    """
+    return None
+
+  def load_internal_states_ops(self):
+    """Loade internal states from class variables."""
+    return [[tf.no_op()]]
+
+  def save_internal_states_ops(self, internal_states):
+    """Saves internal states into class variables."""
+    return [[tf.no_op()]]
+
   # ============================================================================
   # END SUBCLASS INTERFACE
   # ============================================================================
 
+  def __init__(self, *args, **kwargs):
+    super(NextFrameBase, self).__init__(*args, **kwargs)
+    self.internal_states = self.init_internal_states()
+
   @property
   def _target_modality(self):
     # TODO(mbz): get rid of this somehow.
@@ -454,7 +478,6 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
     ss_func = self.get_scheduled_sample_func(batch_size)
     target_frames = []
     extra_loss = 0.0
-    internal_states = None
 
     # Any extra info required by the model goes into here.
     video_features = self.video_features(
@@ -466,42 +489,66 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
     else:
       input_index_range = range(hparams.video_num_target_frames)
 
+    # Setup the internal states as well as an auxiliary tf op
+    # to enforce syncronization between prediction steps.
+    if self.internal_states is None:
+      internal_states = None
+      sync_op = tf.no_op()
+    else:
+      internal_states = self.load_internal_states_ops()
+      with tf.control_dependencies(flat_lists(internal_states)):
+        sync_op = tf.no_op()
+
     res_frames, sampled_frames, res_rewards = [], [], []
     for i in input_index_range:
-      frames, actions, rewards, target_index = self.__get_next_inputs(
-          i, all_frames, all_actions, all_rewards)
-      target_frame = all_frames[target_index]
-      target_frames.append(tf.identity(target_frame))
-
-      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-        func_in = (frames, actions, rewards, target_frame,
-                   internal_states, video_features)
-        func_out = self.next_frame(*func_in)
-        res_frame, res_reward, res_extra_loss, internal_states = func_out
-        res_frames.append(res_frame)
-        res_rewards.append(res_reward)
-        extra_loss += res_extra_loss / float(len(input_index_range))
-
-      # Only for Softmax loss: sample frame so we can keep iterating.
-      sampled_frame = self.get_sampled_frame(res_frame)
-      sampled_frames.append(sampled_frame)
-
-      # Check whether we are done with context frames or not
-      if self.is_recurrent_model:
-        done_warm_start = (i >= hparams.video_num_input_frames - 1)
-      else:
-        done_warm_start = True  # Always true for non-reccurent networks.
-
-      if self.is_predicting and done_warm_start:
-        all_frames[target_index] = sampled_frame
-
-      # Scheduled sampling during training.
-      if self.is_training:
-        groundtruth_items = [target_frame]
-        generated_items = [sampled_frame]
-        ss_frame, = self.get_scheduled_sample_inputs(
-            done_warm_start, groundtruth_items, generated_items, ss_func)
-        all_frames[target_index] = ss_frame
+      with tf.control_dependencies([sync_op]):
+        frames, actions, rewards, target_index = self.__get_next_inputs(
+            i, all_frames, all_actions, all_rewards)
+        target_frame = all_frames[target_index]
+        target_frames.append(tf.identity(target_frame))
+
+        with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+          func_in = (frames, actions, rewards, target_frame,
+                     internal_states, video_features)
+          func_out = self.next_frame(*func_in)
+          res_frame, res_reward, res_extra_loss, internal_states = func_out
+          res_frames.append(res_frame)
+          res_rewards.append(res_reward)
+          extra_loss += res_extra_loss / float(len(input_index_range))
+
+          # Syncronizing the internals states
+          # Some Tensflow Magic to make sure everything happens as it should.
+          with tf.control_dependencies([res_frame]):
+            sync_op = tf.no_op()
+            if self.is_predicting and self.is_recurrent_model and i == 0:
+              # The internal state save happens at the end of the 1st iteration
+              # which essentially allows recurrent models to continue
+              # running after one prediction.
+              # Necessary for planning/rl applications.
+              save_ops = self.save_internal_states_ops(internal_states)
+              with tf.control_dependencies(flat_lists(save_ops)):
+                sync_op = tf.no_op()
+
+        # Only for Softmax loss: sample frame so we can keep iterating.
+        sampled_frame = self.get_sampled_frame(res_frame)
+        sampled_frames.append(sampled_frame)
+
+        # Check whether we are done with context frames or not
+        if self.is_recurrent_model:
+          done_warm_start = (i >= hparams.video_num_input_frames - 1)
+        else:
+          done_warm_start = True  # Always true for non-reccurent networks.
+
+        if self.is_predicting and done_warm_start:
+          all_frames[target_index] = sampled_frame
+
+        # Scheduled sampling during training.
+        if self.is_training:
+          groundtruth_items = [target_frame]
+          generated_items = [sampled_frame]
+          ss_frame, = self.get_scheduled_sample_inputs(
+              done_warm_start, groundtruth_items, generated_items, ss_func)
+          all_frames[target_index] = ss_frame
 
     video_extra_loss = self.video_extra_loss(
         sampled_frames, target_frames, internal_states, video_features)
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 4a6bb070a..04774782d 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -429,6 +429,42 @@ def next_frame(self, frames, actions, rewards, target_frame,
     return pred_image, pred_reward, extra_loss, internal_states
 
 
+@registry.register_model
+class NextFrameSv2pAtari(NextFrameSv2p):
+  """SV2P with specific changes for atari pipeline."""
+
+  def init_internal_states(self):
+    # Hardcoded LSTM-CONV shapes.
+    # These sizes are calculated based on original atari frames.
+    # TODO(mbz): find a cleaner way of doing this maybe?!
+    batch_size = self.hparams.batch_size
+    shapes = [(batch_size, 53, 40, 8),
+              (batch_size, 53, 40, 8),
+              (batch_size, 27, 20, 16),
+              (batch_size, 27, 20, 16),
+              (batch_size, 53, 40, 8)]
+
+    with tf.variable_scope("clean_scope"):
+      # Initialize conv-lstm states with zeros
+      init = tf.zeros_initializer()
+      states = []
+      for i, shape in enumerate(shapes):
+        # every lstm-conv state has two variables named c and h.
+        c = tf.get_variable("c%d" % i, shape, trainable=False, initializer=init)
+        h = tf.get_variable("h%d" % i, shape, trainable=False, initializer=init)
+        states.append((c, h))
+      return states
+
+  def load_internal_states_ops(self):
+    ops = [(c.read_value(), h.read_value()) for c, h in self.internal_states]
+    return ops
+
+  def save_internal_states_ops(self, internal_states):
+    ops = [[tf.assign(x[0], y[0]), tf.assign(x[1], y[1])]
+           for x, y in zip(self.internal_states, internal_states)]
+    return ops
+
+
 @registry.register_model
 class NextFrameSv2pLegacy(NextFrameSv2p):
   """Old SV2P code. Only for legacy reasons."""

From d2debab8b2c81e8833fcaeb2b168030e7aa2d0b8 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sun, 4 Nov 2018 21:46:32 -0800
Subject: [PATCH 1169/2720] Make decoder decodes specific checkpoint in MLPerf
 instead of the latest checkpoint.

PiperOrigin-RevId: 220041669
---
 tensor2tensor/utils/trainer_lib.py | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index ee064adc3..18785ebc9 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -535,7 +535,10 @@ def run_std_server(self):
         protocol=self._hparams.std_server_protocol)
     server.join()
 
-  def decode(self, dataset_split=None, decode_from_file=False):
+  def decode(self,
+             dataset_split=None,
+             decode_from_file=False,
+             checkpoint_path=None):
     """Decodes from dataset or file."""
     if decode_from_file:
       decoding.decode_from_file(self._estimator,
@@ -544,11 +547,13 @@ def decode(self, dataset_split=None, decode_from_file=False):
                                 self._decode_hparams,
                                 self._decode_hparams.decode_to_file)
     else:
-      decoding.decode_from_dataset(self._estimator,
-                                   self._hparams.problem.name,
-                                   self._hparams,
-                                   self._decode_hparams,
-                                   dataset_split=dataset_split)
+      decoding.decode_from_dataset(
+          self._estimator,
+          self._hparams.problem.name,
+          self._hparams,
+          self._decode_hparams,
+          dataset_split=dataset_split,
+          checkpoint_path=checkpoint_path)
 
   def continuous_decode(self):
     """Decode from dataset on new checkpoint."""
@@ -573,11 +578,16 @@ def continuous_decode_on_eval_data(self):
       # Skip checkpoint 0.
       if current_step == 0:
         continue
+      # Decode the latest checkpoint by default.
+      checkpoint_path = None
       if self._decode_hparams.mlperf_mode:
         self._decode_hparams.mlperf_decode_step = current_step
+        checkpoint_path = ckpt
 
       mlperf_log.transformer_print(key=mlperf_log.EVAL_START)
-      self.decode(dataset_split=tf.estimator.ModeKeys.EVAL)
+      self.decode(
+          dataset_split=tf.estimator.ModeKeys.EVAL,
+          checkpoint_path=checkpoint_path)
       d_hparams = self._decode_hparams
       if d_hparams.mlperf_mode and d_hparams.mlperf_success:
         mlperf_log.transformer_print(

From 6503fbcdf4f70a1c018e25e45f6a5efa8ab74d2a Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Mon, 5 Nov 2018 10:26:35 -0800
Subject: [PATCH 1170/2720] Adding the adam optimizer that supports weight
 decay.

PiperOrigin-RevId: 220127246
---
 tensor2tensor/utils/learning_rate.py |  7 +++++++
 tensor2tensor/utils/optimize.py      | 14 ++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index dcb5839b0..bf037a876 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -33,6 +33,13 @@ def learning_rate_factor(name, step_num, hparams):
   elif name == "linear_decay":
     ret = (hparams.train_steps - step_num) / hparams.learning_rate_decay_steps
     return tf.minimum(1.0, tf.maximum(0.0, ret))
+  elif name == "cosdecay":  # openai gpt
+    in_warmup = tf.cast(step_num <= hparams.learning_rate_warmup_steps,
+                        dtype=tf.float32)
+    ret = 0.5 * (1 + tf.cos(
+        np.pi * step_num / hparams.learning_rate_decay_steps))
+    # if in warmup stage return 1 else return the decayed value
+    return in_warmup * 1 + (1 - in_warmup) * ret
   elif name == "rsqrt_decay":
     return tf.rsqrt(tf.maximum(step_num, hparams.learning_rate_warmup_steps))
   elif name == "rsqrt_normalized_decay":
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 2cf491257..521b90f78 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -118,6 +118,20 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
           beta1=hparams.optimizer_adam_beta1,
           beta2=hparams.optimizer_adam_beta2,
           epsilon=hparams.optimizer_adam_epsilon)
+    elif optimizer_name == "AdamW":
+      # Openai gpt used weight decay.
+      # Given the internals of AdamW, weight decay dependent on the
+      # learning rate is chosen to match the openai implementation.
+      # The weight decay update to each parameter is applied before the adam
+      # gradients computation, which is different from that described
+      # in the paper and in the openai implementation:
+      # https://arxiv.org/pdf/1711.05101.pdf
+      self._opt = tf.contrib.opt.AdamWOptimizer(
+          0.01*lr,
+          lr,
+          beta1=hparams.optimizer_adam_beta1,
+          beta2=hparams.optimizer_adam_beta2,
+          epsilon=hparams.optimizer_adam_epsilon)
     elif optimizer_name == "Adafactor":
       self._opt = adafactor.adafactor_optimizer_from_hparams(hparams, lr)
     else:

From d3df5d891f735f1bd33d5c3a9f8071b841f6508d Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Mon, 5 Nov 2018 10:50:29 -0800
Subject: [PATCH 1171/2720] Adding some multiproblem relevant classes to
 existing data generators.

PiperOrigin-RevId: 220132103
---
 .../data_generators/cnn_dailymail.py          | 34 ++++++++++++++++---
 tensor2tensor/data_generators/multinli.py     | 23 ++++++++++++-
 tensor2tensor/data_generators/stanford_nli.py | 19 +++++++++++
 .../data_generators/text_problems.py          |  2 +-
 tensor2tensor/data_generators/wiki_lm.py      | 14 ++++++++
 5 files changed, 85 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index c6a6f33cc..177b1a4ed 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -28,22 +28,28 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
 # Links to data from http://cs.nyu.edu/~kcho/DMQA/
-_CNN_STORIES_DRIVE_URL = "https://drive.google.com/uc?export=download&id=0BwmD_VLjROrfTHk4NFg2SndKcjQ"
+_CNN_STORIES_DRIVE_URL = ("https://drive.google.com/uc?"
+                          "export=download&id=0BwmD_VLjROrfTHk4NFg2SndKcjQ")
 
-_DAILYMAIL_STORIES_DRIVE_URL = "https://drive.google.com/uc?export=download&id=0BwmD_VLjROrfM1BxdkxVaTY2bWs"
+_DAILYMAIL_STORIES_DRIVE_URL = ("https://drive.google.com/uc?export=download&id"
+                                "=0BwmD_VLjROrfM1BxdkxVaTY2bWs")
 
 # Note: using See et al. (2017) as reference for data generation
 # For more info, use the links below
 
 # Train/Dev/Test Splits for summarization data
-_TRAIN_URLS = "https://raw.githubusercontent.com/abisee/cnn-dailymail/master/url_lists/all_train.txt"
-_DEV_URLS = "https://raw.githubusercontent.com/abisee/cnn-dailymail/master/url_lists/all_val.txt"
-_TEST_URLS = "https://raw.githubusercontent.com/abisee/cnn-dailymail/master/url_lists/all_test.txt"
+_TRAIN_URLS = ("https://raw.githubusercontent.com/abisee/cnn-dailymail/"
+               "master/url_lists/all_train.txt")
+_DEV_URLS = ("https://raw.githubusercontent.com/abisee/cnn-dailymail/"
+             "master/url_lists/all_val.txt")
+_TEST_URLS = ("https://raw.githubusercontent.com/abisee/cnn-dailymail/"
+              "master/url_lists/all_test.txt")
 
 # End-of-sentence marker.
 EOS = text_encoder.EOS_ID
@@ -240,3 +246,21 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
     for example in example_generator(all_files, urls_path, sum_token=True):
       story, summary = _story_summary_split(example)
       yield {"inputs": story, "targets": summary}
+
+
+@registry.register_problem
+class SummarizeCnnDailymailWikiLMSharedVocab(SummarizeCnnDailymail32k):
+  """Summarize CNN and Daily Mail articles using the Wiki 32k vocab."""
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelEnWiki32k().vocab_filename
+
+
+@registry.register_problem
+class SummarizeCnnDailymailWikiLMSharedVocab64k(SummarizeCnnDailymail32k):
+  """Summarize CNN and Daily Mail articles using the Wiki 64k vocab."""
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelEnWiki64k().vocab_filename
diff --git a/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index 0e52dd455..dcb94d1b8 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -27,6 +27,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
 import tensorflow as tf
 
@@ -105,7 +106,9 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
     if dataset_split == problem.DatasetSplit.TRAIN:
       filesplit = ["train.tsv"]
     else:
-      filesplit = ["dev_matched.tsv", "dev_mismatched.tsv"]
+      # Using dev matched as the default for eval. Can also switch this to
+      # dev_mismatched.tsv
+      filesplit = ["dev_matched.tsv"]
 
     for fs in filesplit:
       filename = os.path.join(mnli_dir, fs)
@@ -132,3 +135,21 @@ class MultiNLISharedVocab(MultiNLI):
   @property
   def vocab_filename(self):
     return lm1b.LanguagemodelLm1b32k().vocab_filename
+
+
+@registry.register_problem
+class MultiNLIWikiLMSharedVocab(MultiNLI):
+  """MultiNLI classification problems with the Wiki vocabulary"""
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelEnWiki32k().vocab_filename
+
+
+@registry.register_problem
+class MultiNLIWikiLMSharedVocab64k(MultiNLIWikiLMSharedVocab):
+  """MultiNLI classification problems with the Wiki vocabulary"""
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelEnWiki64k().vocab_filename
diff --git a/tensor2tensor/data_generators/stanford_nli.py b/tensor2tensor/data_generators/stanford_nli.py
index 2f9421d08..a8aa04602 100644
--- a/tensor2tensor/data_generators/stanford_nli.py
+++ b/tensor2tensor/data_generators/stanford_nli.py
@@ -27,6 +27,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
 import tensorflow as tf
 
@@ -130,3 +131,21 @@ class StanfordNLISharedVocab(StanfordNLI):
   @property
   def vocab_filename(self):
     return lm1b.LanguagemodelLm1b32k().vocab_filename
+
+
+@registry.register_problem
+class StanfordNLIWikiLMSharedVocab(StanfordNLI):
+  """StanfordNLI classification problems with the Wiki vocabulary"""
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelEnWiki32k().vocab_filename
+
+
+@registry.register_problem
+class StanfordNLIWikiLMSharedVocab64k(StanfordNLIWikiLMSharedVocab):
+  """StanfordNLI classification problems with the Wiki vocabulary"""
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelEnWiki64k().vocab_filename
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index b9b4384a6..e3d744b0a 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -539,9 +539,9 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
       inputs = []
       for idx, inp in enumerate(sample["inputs"]):
         inputs += encoder.encode(inp)
-        inputs.append(text_encoder.EOS_ID)
         if idx < len(sample["inputs"]) - 1:
           inputs.append(encoder.encode(self.CONCAT_TOKEN)[0])
+      inputs.append(text_encoder.EOS_ID)
       label = sample["label"]
       yield {"inputs": inputs, "targets": [label]}
 
diff --git a/tensor2tensor/data_generators/wiki_lm.py b/tensor2tensor/data_generators/wiki_lm.py
index ba79f214b..6db24bb04 100644
--- a/tensor2tensor/data_generators/wiki_lm.py
+++ b/tensor2tensor/data_generators/wiki_lm.py
@@ -152,6 +152,20 @@ def approx_vocab_size(self):
     return 64000
 
 
+@registry.register_problem
+class LanguagemodelEnWiki64kShorter(LanguagemodelEnWiki64k):
+  """With 64k vocabulary and shorter truncation lengths."""
+
+  @property
+  def combine_characters_threshold(self):
+    """Threshold for upto how many characters to combine in examples."""
+    return 384*8
+
+  @property
+  def vocab_filename(self):
+    return LanguagemodelEnWiki64k().vocab_filename
+
+
 @registry.register_problem
 class LanguagemodelDeWiki32k(LanguagemodelEnWiki32k):
   """A language model on the untokenized wikipedia corpus, German."""

From a0d3867db5c15e79cbddfe962892dbbfac0f3364 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 5 Nov 2018 11:12:13 -0800
Subject: [PATCH 1172/2720] Fix the Method Resolution Order in the SAVP class
 from the recent refactorig.

PiperOrigin-RevId: 220136502
---
 tensor2tensor/models/video/savp.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index d05244c83..b8937c50c 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -333,7 +333,7 @@ def train_hooks(hook_context):
 
 
 @registry.register_model
-class NextFrameSAVP(sv2p.NextFrameSv2pLegacy, NextFrameSavpBase):
+class NextFrameSAVP(NextFrameSavpBase, sv2p.NextFrameSv2pLegacy):
   """Stochastic Adversarial Video Prediction."""
 
   def construct_model(self, images, actions, rewards):

From 5e1ae4b19e431bac1dadfcd6b9be0bd31664a6f7 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 5 Nov 2018 11:14:46 -0800
Subject: [PATCH 1173/2720] Add Problem classes for En-Fr translation with
 back-translated data.

PiperOrigin-RevId: 220136937
---
 .../data_generators/text_problems.py          |  6 +-
 .../data_generators/translate_enfr.py         | 68 +++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index e3d744b0a..118c8a9a1 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -280,6 +280,10 @@ def max_subtoken_length(self):
   def batch_size_means_tokens(self):
     return True
 
+  @property
+  def already_shuffled(self):
+    return False
+
   def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
     filepath_fns = {
@@ -289,7 +293,7 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
     }
 
     split_paths = [(split["split"], filepath_fns[split["split"]](
-        data_dir, split["shards"], shuffled=False))
+        data_dir, split["shards"], shuffled=self.already_shuffled))
                    for split in self.dataset_splits]
     all_paths = []
     for _, paths in split_paths:
diff --git a/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index 6f5ad163d..73261ffb1 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -18,6 +18,9 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+import os
+
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
@@ -138,6 +141,71 @@ def vocab_filename(self):
     return TranslateEnfrWmt32k().vocab_filename
 
 
+@registry.register_problem
+class TranslateEnfrWmt32kWithBacktranslateFr(TranslateEnfrWmt32k):
+  """En-Fr translation with added French data, back-translated."""
+
+  @property
+  def vocab_filename(self):
+    return TranslateEnfrWmt32k().vocab_filename
+
+  @property
+  def already_shuffled(self):
+    return True
+
+  @property
+  def backtranslate_data_filenames(self):
+    """List of pairs of files with matched back-translated data."""
+    # Files must be placed in tmp_dir, each similar size to authentic data.
+    return [("fr_mono_en.txt", "fr_mono_fr.txt")]
+
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each."""
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 1,  # Use just 1 shard so as to not mix data.
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    datasets = self.source_data_files(dataset_split)
+    tag = "train" if dataset_split == problem.DatasetSplit.TRAIN else "dev"
+    data_path = translate.compile_data(
+        tmp_dir, datasets, "%s-compiled-%s" % (self.name, tag))
+    # Iterator over authentic data.
+    it_auth = text_problems.text2text_txt_iterator(
+        data_path + ".lang1", data_path + ".lang2")
+    # For eval, use authentic data.
+    if dataset_split != problem.DatasetSplit.TRAIN:
+      for example in it_auth:
+        yield example
+    else:  # For training, mix synthetic and authentic data as follows.
+      for (file1, file2) in self.backtranslate_data_filenames:
+        path1 = os.path.join(tmp_dir, file1)
+        path2 = os.path.join(tmp_dir, file2)
+        # Synthetic data first.
+        for example in text_problems.text2text_txt_iterator(path1, path2):
+          yield example
+        # Now authentic data.
+        for example in it_auth:
+          yield example
+
+
+@registry.register_problem
+class TranslateEnfrWmt32kWithBacktranslateEn(
+    TranslateEnfrWmt32kWithBacktranslateFr):
+  """En-Fr translation with added English data, back-translated."""
+
+  @property
+  def backtranslate_data_filenames(self):
+    """List of pairs of files with matched back-translated data."""
+    # Files must be placed in tmp_dir, each similar size to authentic data.
+    return [("en_mono_en.txt%d" % i, "en_mono_fr.txt%d" % i) for i in [0, 1, 2]]
+
+
 @registry.register_problem
 class TranslateEnfrWmtSmallCharacters(translate.TranslateProblem):
   """Problem spec for WMT En-Fr translation."""

From ce590f4cc39102fad129d1ce03c77f5c8c7e6b7c Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Mon, 5 Nov 2018 11:21:35 -0800
Subject: [PATCH 1174/2720] Some configs for pretraining and finetuning.

PiperOrigin-RevId: 220138357
---
 tensor2tensor/models/transformer.py | 60 +++++++++++++++++++++++++++--
 1 file changed, 56 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 9c9ab14d8..910a05f3c 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1518,22 +1518,74 @@ def transformer_big():
 
 @registry.register_hparams
 def transformer_tall():
-  """Hparams for transformer on LM+MNLI."""
+  """Hparams for transformer on LM for pretraining/finetuning/mixing."""
   hparams = transformer_base()
   hparams.batch_size = 2048
   hparams.hidden_size = 768
   hparams.filter_size = 3072
   hparams.num_hidden_layers = 12
   hparams.num_heads = 12
-  hparams.learning_rate_schedule = (
-      "constant*linear_warmup*rsqrt_hidden_size")
-  hparams.learning_rate_constant = 2e-3
   hparams.label_smoothing = 0.0
   hparams.max_length = 512
   hparams.eval_drop_long_sequences = True
+  hparams.multiproblem_mixing_schedule = "pretrain"
   return hparams
 
 
+@registry.register_hparams
+def transformer_tall_finetune_textclass():
+  """Hparams for transformer on LM for finetuning on text class problems."""
+  hparams = transformer_tall()
+  hparams.learning_rate_constant = 6.25e-5
+  hparams.learning_rate_schedule = (
+      "linear_warmup*constant*linear_decay")
+  hparams.multiproblem_schedule_max_examples = 0
+  hparams.multiproblem_target_eval_only = True
+  hparams.multiproblem_class_loss_multiplier = 4
+  hparams.learning_rate_warmup_steps = 50
+  # Set train steps to learning_rate_decay_steps or less
+  hparams.learning_rate_decay_steps = 25000
+  hparams.multiproblem_reweight_label_loss = True
+  hparams.multiproblem_label_weight = 0.95
+  return hparams
+
+
+@registry.register_hparams
+def transformer_tall_pretrain_lm():
+  """Hparams for transformer on LM pretraining (with 64k vocab)."""
+  hparams = transformer_tall()
+  hparams.learning_rate_constant = 2e-4
+  hparams.learning_rate_schedule = (
+      "linear_warmup*constant*cosdecay")
+  hparams.optimizer = "AdamW"
+  hparams.optimizer_adam_beta1 = 0.9
+  hparams.optimizer_adam_beta2 = 0.999
+  hparams.optimizer_adam_epsilon = 1e-8
+  # Set max examples to something big when pretraining only the LM, definitely
+  # something an order of magnitude bigger than number of train steps.
+  hparams.multiproblem_schedule_max_examples = 5e8
+  # Set train steps to learning_rate_decay_steps or less
+  hparams.learning_rate_decay_steps = 5000000
+  return hparams
+
+
+@registry.register_hparams
+def transformer_tall_finetune_cnndm():
+  """Hparams for transformer on LM for finetuning on cnndm summarization."""
+  hparams = transformer_tall()
+  hparams.batch_size = 4096
+  hparams.multiproblem_max_input_length = 412
+  hparams.multiproblem_max_target_length = 100
+  hparams.multiproblem_schedule_max_examples = 0
+  hparams.learning_rate_schedule = (
+      "linear_warmup*constant*cosdecay")
+  hparams.learning_rate_constant = 5e-5
+  hparams.learning_rate_warmup_steps = 100
+  # Set train steps to learning_rate_decay_steps or less
+  hparams.learning_rate_decay_steps = 40000
+  hparams.multiproblem_target_eval_only = True
+
+
 @registry.register_hparams
 def transformer_tall_big():
   """Hparams for transformer on LM+MNLI."""

From 85599aae5b441073fedc6dd02475edbc87a7b87b Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 5 Nov 2018 13:11:31 -0800
Subject: [PATCH 1175/2720] Use v.device instead of v._ref.device

PiperOrigin-RevId: 220158804
---
 tensor2tensor/utils/optimize.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 521b90f78..e2f9e45c9 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -183,7 +183,7 @@ def weight_noise(noise_rate, learning_rate, var_list):
   noise_ops = []
 
   for v in var_list:
-    with tf.device(v._ref().device):  # pylint: disable=protected-access
+    with tf.device(v.device):  # pylint: disable=protected-access
       scale = noise_rate * learning_rate * 0.001
       if common_layers.should_generate_summaries():
         tf.summary.scalar("weight_noise_scale", scale)

From 0a5bb1dd7f6176fa8963cb7c06cc9f8f7649b039 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 5 Nov 2018 13:38:04 -0800
Subject: [PATCH 1176/2720] Do not use DistributionStrategy yet. Just
 MirroredStrategy has issues with memory and variable assignment, so we need
 to find or create another strategy to replace it. Memory issue is apparent as
 I cannot train transformer_big_enfr with 12 layers with it even batch size
 786, while without it 1024 works fine.

PiperOrigin-RevId: 220163481
---
 tensor2tensor/bin/t2t_trainer.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 413afae9a..a221a6b52 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -69,8 +69,9 @@
 flags.DEFINE_integer("intra_op_parallelism_threads", 0,
                      "Number of intra_op_parallelism_threads to use for CPU. "
                      "See TensorFlow config.proto for details.")
+# TODO(lukaszkaiser): resolve memory and variable assign issues and set to True.
 flags.DEFINE_bool(
-    "optionally_use_dist_strat", True,
+    "optionally_use_dist_strat", False,
     "Whether to use TensorFlow DistributionStrategy instead of explicitly "
     "replicating the model. DistributionStrategy is used only if the "
     "model replication configuration is supported by the DistributionStrategy.")

From 9c45ca4977206feb13808dfb781b6f5815370f78 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 5 Nov 2018 14:38:48 -0800
Subject: [PATCH 1177/2720] Allow training a wider-lstm by passing
 hparams.latent_encoder_width when hparams.latent_dist_encoder="conv_lstm".

PiperOrigin-RevId: 220174951
---
 tensor2tensor/models/research/glow_ops.py      | 2 +-
 tensor2tensor/models/research/glow_ops_test.py | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 2ba0e4636..0a2ea3904 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -607,7 +607,7 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
     output_channels = common_layers.shape_list(z)[-1]
     latent_stack = tf.concat((prior_dist.loc, latent), axis=-1)
     _, state = common_video.conv_lstm_2d(
-        latent_stack, state, output_channels, kernel_size=3,
+        latent_stack, state, hparams.latent_encoder_width, kernel_size=3,
         name="conv_lstm")
     cond_dist = tensor_to_dist(
         "state_to_dist", state.h, output_channels=output_channels)
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 423d82a1c..670da0d52 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -303,7 +303,7 @@ def test_latent_dist_encoder_lstm(self):
       # Initialize x, latent, state.
       x_rand = rng.randn(12, 32, 32, 16).astype(np.float32)
       latent_rand = rng.randn(12, 32, 32, 16).astype(np.float32)
-      state_rand = rng.randn(12, 32, 32, 16).astype(np.float32)
+      state_rand = rng.randn(12, 32, 32, 256).astype(np.float32)
       x_t = tf.convert_to_tensor(x_rand)
       latent_t = tf.convert_to_tensor(latent_rand)
       state_t = tf.convert_to_tensor(state_rand)
@@ -311,6 +311,7 @@ def test_latent_dist_encoder_lstm(self):
       hparams = glow.glow_hparams()
       hparams.add_hparam("latent_dist_encoder", "conv_lstm")
       hparams.add_hparam("latent_skip", True)
+      hparams.add_hparam("latent_encoder_width", 256)
 
       prior_dist, new_state = glow_ops.compute_prior(
           "lstm_prior", x_t, latent=latent_t, hparams=hparams, state=init_state,

From 0aaf50836229c29384f4f1890d38474953447e78 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 5 Nov 2018 15:06:37 -0800
Subject: [PATCH 1178/2720] Fix mlperf logging epoch in transformer 2x2.

PiperOrigin-RevId: 220180486
---
 tensor2tensor/utils/trainer_lib.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 18785ebc9..482a3115d 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -471,7 +471,7 @@ def train_eval_and_decode(self):
         self._hparams.problem_hparams = p_hparams
       mlperf_log.transformer_print(key=mlperf_log.EVAL_START)
       if self._decode_hparams.mlperf_mode:
-        self._decode_hparams.mlperf_decode_step = i
+        self._decode_hparams.mlperf_decode_step = i + eval_steps
       self.decode(dataset_split=tf.estimator.ModeKeys.EVAL)
       d_hparams = self._decode_hparams
       if d_hparams.mlperf_mode and d_hparams.mlperf_success:

From a159998484a60a78a34e2f3e53c7a51b7753d6a4 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Mon, 5 Nov 2018 16:53:19 -0800
Subject: [PATCH 1179/2720] Added local attention option for mtf transformer
 (including incremental decoding).  Updated configurations for MoE
 experiments.  Updated local attention 1d code.

PiperOrigin-RevId: 220198118
---
 tensor2tensor/models/mtf_image_transformer.py |   6 +-
 .../models/mtf_image_transformer_test.py      |   9 +-
 tensor2tensor/models/mtf_transformer.py       | 183 ++++++------------
 tensor2tensor/models/research/moe.py          |  15 +-
 .../models/research/moe_experiments.py        | 129 ++++++------
 5 files changed, 132 insertions(+), 210 deletions(-)

diff --git a/tensor2tensor/models/mtf_image_transformer.py b/tensor2tensor/models/mtf_image_transformer.py
index eeb1942cf..1564af7f7 100644
--- a/tensor2tensor/models/mtf_image_transformer.py
+++ b/tensor2tensor/models/mtf_image_transformer.py
@@ -338,13 +338,15 @@ def local_attention1d_masked_decoder(x, kv_dim, heads_dim,
     layer_name = "decoder_layer_%d" % layer
     with tf.variable_scope(layer_name):
       # Self attention layer
+      length_per_split = mtf.tensor_dim_to_size_per_split(
+          hparams.layout, hparams.mesh_shape, length_dim)
       x += layer_prepostprocess_dropout(
           mtf.layers.masked_local_attention_1d(
               mtf.layers.layer_norm(x, model_dim, name="layer_norm_att"),
-              None,
               kv_dim,
               heads_dim,
-              block_length=hparams.block_length,
+              window_size=hparams.block_length,
+              length_per_split=length_per_split,
               name="self_att"), hparams)
       # ffn layer
       x += layer_prepostprocess_dropout(
diff --git a/tensor2tensor/models/mtf_image_transformer_test.py b/tensor2tensor/models/mtf_image_transformer_test.py
index 27437c677..27c240ed9 100644
--- a/tensor2tensor/models/mtf_image_transformer_test.py
+++ b/tensor2tensor/models/mtf_image_transformer_test.py
@@ -74,9 +74,10 @@ class MtfImageTransformerTest(tf.test.TestCase):
   def testMtfImageTransformer(self):
     hparams = mtf_image_transformer.mtf_image_transformer_single()
 
-    model, features, hparams = get_model(hparams)
+    # need to know layout ahead of time for local attention.
     hparams.mesh_shape = ""
     hparams.layout = ""
+    model, features, hparams = get_model(hparams)
     mesh, mesh_impl = get_placement_mesh(hparams)
 
     logits, _ = model.mtf_model_fn(features, mesh)
@@ -95,9 +96,10 @@ def testMtfImageTransformer(self):
   def testMtfImageTransformerDataParallel(self):
     hparams = mtf_image_transformer.mtf_image_transformer_single()
 
-    model, features, hparams = get_model(hparams)
+    # need to know layout ahead of time for local attention.
     hparams.mesh_shape = "all:2"
     hparams.layout = "batch:all"
+    model, features, hparams = get_model(hparams)
     mesh, mesh_impl = get_placement_mesh(hparams)
 
     logits, _ = model.mtf_model_fn(features, mesh)
@@ -116,9 +118,10 @@ def testMtfImageTransformerDataParallel(self):
   def testMtfImageTransformerModelParallel(self):
     hparams = mtf_image_transformer.mtf_image_transformer_single()
 
-    model, features, hparams = get_model(hparams)
+    # need to know layout ahead of time for local attention.
     hparams.mesh_shape = "all:2"
     hparams.layout = "length:all"
+    model, features, hparams = get_model(hparams)
     mesh, mesh_impl = get_placement_mesh(hparams)
 
     logits, _ = model.mtf_model_fn(features, mesh)
diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 8cff5786e..a0bca0ccf 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -360,8 +360,7 @@ def _layer_stack(self,
                    losses=None,
                    step_num=None,
                    encdec_tensors=None,
-                   self_attention_k=None,
-                   self_attention_v=None):
+                   states=None):
     """Encoder or decoder stack.
 
     Args:
@@ -377,10 +376,7 @@ def _layer_stack(self,
       step_num: an optional mtf integer Scalar (used in incrmenental mode)
       encdec_tensors: an optional list of num_layers tuples, each of the form
         (q_var, o_var, k, v), (used in incremental mode)
-      self_attention_k: an optional list of num_layers Tensors each with shape
-        [batch, heads, memory_length, kv_channels] (incremental mode)
-      self_attention_v: an optional list of num_layers Tensors each with shape
-        [batch, heads, memory_length, kv_channels] (incremental mode)
+      states: an optional list of Tensors (used in incremental mode)
     Returns:
       a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim]
     Raises:
@@ -410,25 +406,25 @@ def normalize(x):
       return x * mtf.rsqrt(variance + hparams.norm_epsilon) * scale
 
     if is_incremental:
-      new_self_attention_k = []
-      new_self_attention_v = []
+      states = list(states)
+      new_states = []
+    tf.logging.info("states = %s" % (states,))
 
     for lnum, layer_type in enumerate(layers):
       with tf.variable_scope("%s_%d" % (layer_type, lnum)):
         if layer_type == "att":
           # Self attention layer
           if is_incremental:
-            self_att_num = len(new_self_attention_k)
             y, new_k, new_v = mtf.layers.multihead_self_attention_incremental(
                 normalize(x),
-                prev_k=self_attention_k[self_att_num],
-                prev_v=self_attention_v[self_att_num],
+                prev_k=states.pop(0),
+                prev_v=states.pop(0),
                 step_num=step_num,
                 master_dtype=self.master_dtype,
                 slice_dtype=self.slice_dtype,
                 name="att")
-            new_self_attention_k.append(new_k)
-            new_self_attention_v.append(new_v)
+            new_states.append(new_k)
+            new_states.append(new_v)
             x += y
           else:
             x += layer_prepostprocess_dropout(
@@ -460,6 +456,31 @@ def normalize(x):
                     master_dtype=self.master_dtype,
                     slice_dtype=self.slice_dtype,
                     name="enc_att"))
+        elif layer_type == "local_att":
+          if is_incremental:
+            y, new_k, new_v = mtf.layers.masked_local_attention_1d_incremental(
+                normalize(x),
+                prev_k=states.pop(0),
+                prev_v=states.pop(0),
+                step_num=step_num,
+                master_dtype=self.master_dtype,
+                slice_dtype=self.slice_dtype,
+                name="local_att")
+            new_states.append(new_k)
+            new_states.append(new_v)
+            x += y
+          else:
+            x += layer_prepostprocess_dropout(
+                mtf.layers.masked_local_attention_1d(
+                    normalize(x),
+                    self.kv_dim, self.heads_dim,
+                    window_size=hparams.local_attention_window_size,
+                    master_dtype=self.master_dtype,
+                    slice_dtype=self.slice_dtype,
+                    length_per_split=mtf.tensor_dim_to_size_per_split(
+                        hparams.layout, hparams.mesh_shape,
+                        self.max_length_dim),
+                    name="local_att"))
         else:
           if is_incremental:
             # insert length dimension.
@@ -478,10 +499,9 @@ def normalize(x):
     x = layer_prepostprocess_dropout(normalize(x))
     assert not layer_norm_vars
     if is_incremental:
-      return x, new_self_attention_k, new_self_attention_v
+      return x, new_states
     else:
       return x
-      # return mtf.cast(x, self.activation_dtype)
 
   def sample(self, features, mesh):
     with tf.variable_scope("transformer"):
@@ -563,43 +583,52 @@ def _sample(self, features, mesh):
           "hparams.model_type = %s not yet supported"
           % hparams.transformer_type)
 
+    local_attention_window = mtf.Dimension(
+        "local_attention_window", hparams.local_attention_window_size)
     if hparams.beam_size == 1:
       ids_shape = mtf.Shape(self.batch_dims + [self.length_dim])
       kv_shape = mtf.Shape(self.batch_dims +
                            [self.heads_dim,
                             self.memory_length_dim, self.kv_dim])
+      local_kv_shape = mtf.Shape(self.batch_dims +
+                                 [self.heads_dim,
+                                  local_attention_window, self.kv_dim])
     else:
       beam_dim = mtf.Dimension("beam", hparams.beam_size)
       ids_shape = mtf.Shape(self.batch_dims + [beam_dim, self.length_dim])
       kv_shape = mtf.Shape(self.batch_dims +
                            [beam_dim, self.heads_dim,
                             self.memory_length_dim, self.kv_dim])
+      local_kv_shape = mtf.Shape(self.batch_dims +
+                                 [beam_dim, self.heads_dim,
+                                  local_attention_window, self.kv_dim])
 
     initial_ids = mtf.constant(mesh, 0, ids_shape, dtype=tf.int32)
-    num_self_att = len([l for l in hparams.decoder_layers if l == "att"])
-    initial_kv_states = (
-        [mtf.zeros(mesh, kv_shape, dtype=self.activation_dtype)]
-        * (2 * num_self_att))
+    initial_states = []
+    for layer in hparams.decoder_layers:
+      if layer == "att":
+        initial_states.extend(
+            [mtf.zeros(mesh, kv_shape, dtype=self.activation_dtype)] * 2)
+      elif layer == "local_att":
+        initial_states.extend(
+            [mtf.zeros(mesh, local_kv_shape, dtype=self.activation_dtype)] * 2)
+
     def logits_fn(step_num, ids, states):
       """Produce logits for this step, and new states."""
-      self_attention_k = states[:num_self_att]
-      self_attention_v = states[num_self_att:]
       ids_this_step = mtf.gather(ids, step_num - 1, self.length_dim)
       x = (mtf.gather(targets_embedding_var, ids_this_step,
                       self.targets_vocab_dim) +
            mtf.gather(positional_embedding_var, step_num, self.max_length_dim))
       with tf.variable_scope("decoder"):
-        x, new_self_attention_k, new_self_attention_v = (
-            self._layer_stack(
-                x,
-                hparams.decoder_layers,
-                encdec_attention_mask=encoder_attention_mask,
-                step_num=step_num,
-                encdec_tensors=encdec_tensors,
-                self_attention_k=self_attention_k,
-                self_attention_v=self_attention_v))
+        x, new_states = self._layer_stack(
+            x,
+            hparams.decoder_layers,
+            encdec_attention_mask=encoder_attention_mask,
+            step_num=step_num,
+            encdec_tensors=encdec_tensors,
+            states=states)
       logits = mtf.matmul(x, softmax_var)
-      return logits, new_self_attention_k + new_self_attention_v
+      return logits, new_states
 
     if hparams.beam_size == 1:
       temperature = (0.0 if hparams.sampling_method == "argmax"
@@ -608,7 +637,7 @@ def logits_fn(step_num, ids, states):
           logits_fn,
           initial_ids,
           temperature=temperature,
-          initial_states=initial_kv_states,
+          initial_states=initial_states,
           forced_ids=partial_targets,
           use_tpu=hparams.use_tpu)
     else:
@@ -626,99 +655,12 @@ def logits_fn(step_num, ids, states):
           logits_fn,
           initial_ids,
           hparams.alpha,
-          states=initial_kv_states,
+          states=initial_states,
           decode_length=decode_length,
           use_tpu=hparams.use_tpu,
           dtype=self.activation_dtype)
       return mtf.gather(beams, mtf.constant(mesh, 0, dtype=tf.int32), beam_dim)
 
-  def _decoder_layer_stack_incremental(self,
-                                       x,
-                                       step_num,
-                                       encdec_tensors,
-                                       self_attention_k,
-                                       self_attention_v,
-                                       encdec_attention_mask=None):
-    """Decoder layer stack during inference.
-
-    We are processing only one position at a time.
-
-    The self-attention keys and values have already been computed for
-    previous positions.  In addition to the decoder output, we need to
-    produce the updated self-attention keys and values.
-
-    If there is an encoder, then additional Tensors are supplied in
-    encdec_tensors, which give us the keys and values for encoder-decoder
-    attention as well as the weight matrices q_var and o_var.
-
-    Args:
-      x: a mtf.Tensor with shape [<batch_dims>, model_dim]
-      step_num: an mtf integer Scalar
-      encdec_tensors: an optional list of num_layers tuples, each of the form
-        (q_var, o_var, k, v)
-      self_attention_k: an optional list of num_layers Tensors each with shape
-        [batch, heads, memory_length, kv_channels]
-      self_attention_v: an optional list of num_layers Tensors each with shape
-        [batch, heads, memory_length, kv_channels]
-      encdec_attention_mask: an optional mtf.Tensor with shape
-        [batch, length_dim, encoder_length_dim] containing values 0 or -inf.
-
-    Returns:
-      y: a mtf.Tensor with shape [<batch_dims>, model_dim]
-      new_self_attention_k: a list of num_layers mtf.Tensors, with the same
-        shapes as the elements of self_attention_k
-      new_self_attention_v: a list of num_layers mtf.Tensors, with the same
-        shapes as the elements of self_attention_v
-
-    Raises:
-      ValueError: if hparams make no sense
-    """
-    hparams = self._hparams
-    num_layers = hparams.num_decoder_layers
-    num_layer_norms = num_layers * (2 if encdec_tensors is None else 3) + 1
-    layer_norms_dim = mtf.Dimension("layer_norms", num_layer_norms)
-    layer_norm_combined_var = mtf.get_variable(
-        x.mesh,
-        "layer_norm_scale",
-        mtf.Shape([layer_norms_dim, self.model_dim]),
-        initializer=tf.ones_initializer(),
-        activation_dtype=x.dtype)
-    layer_norm_vars = mtf.unstack(layer_norm_combined_var, layer_norms_dim)
-    def normalize(x):
-      scale = layer_norm_vars.pop(0)
-      variance = mtf.reduce_mean(mtf.square(x), reduced_dim=self.model_dim)
-      return x * mtf.rsqrt(variance + hparams.norm_epsilon) * scale
-
-    new_self_attention_k = []
-    new_self_attention_v = []
-    for layer in xrange(num_layers):
-      with tf.variable_scope("layer_%d" % layer):
-        # Self attention layer
-        y, new_k, new_v = mtf.layers.multihead_self_attention_incremental(
-            normalize(x),
-            prev_k=self_attention_k[layer],
-            prev_v=self_attention_v[layer],
-            step_num=step_num,
-            master_dtype=self.master_dtype,
-            slice_dtype=self.slice_dtype,
-            name="att")
-        new_self_attention_k.append(new_k)
-        new_self_attention_v.append(new_v)
-        x += y
-        if encdec_tensors is not None:
-          # Encoder-Decoder attention layer
-          q_var, o_var, k, v = encdec_tensors[layer]
-          x += mtf.layers.multihead_encdec_attention_incremental(
-              normalize(x),
-              q_var, o_var, k, v,
-              encdec_attention_mask,
-              name="enc_att")
-        # ffn layer
-        x += self._feedforward_layer(normalize(x), layer)
-    x = normalize(x)
-    assert not layer_norm_vars
-    return x, new_self_attention_k, new_self_attention_v
-
 
 @registry.register_hparams
 def mtf_transformer_base():
@@ -731,6 +673,7 @@ def mtf_transformer_base():
   hparams.max_length = 256
   hparams.add_hparam("d_model", 512)
   hparams.add_hparam("d_kv", 128)
+  hparams.add_hparam("local_attention_window_size", 128)
   hparams.label_smoothing = 0.1
   # 8-way model-parallelism
   hparams.add_hparam("mesh_shape", "model:8")
diff --git a/tensor2tensor/models/research/moe.py b/tensor2tensor/models/research/moe.py
index 1dddba13d..11c31a1b6 100644
--- a/tensor2tensor/models/research/moe.py
+++ b/tensor2tensor/models/research/moe.py
@@ -279,7 +279,7 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train,
   # over which those groups are split.
   num_groups, group_size = _split_into_groups(
       b1.size * l.size, hparams.moe_group_size,
-      _tensor_dim_to_mesh_dim_size(hparams, b1))
+      mtf.tensor_dim_to_mesh_dim_size(hparams.layout, hparams.mesh_shape, b1))
   g1 = mtf.Dimension(b1.name, num_groups)
   g = mtf.Dimension(b1.name + "_unsplit", g1.size)
   s = mtf.Dimension("group_size_x", group_size)
@@ -299,7 +299,7 @@ def transformer_moe_layer_v2(inputs, output_dim, hparams, train,
   num_groups, group_size = _split_into_groups(
       a0.size * g.size * c.size,
       hparams.moe_group_size,
-      _tensor_dim_to_mesh_dim_size(hparams, a0))
+      mtf.tensor_dim_to_mesh_dim_size(hparams.layout, hparams.mesh_shape, a0))
   t = mtf.Dimension("group_size_y", group_size)
   h0 = mtf.Dimension(a0.name, num_groups)
   h = mtf.Dimension(a0.name + "_unsplit", h0.size)
@@ -677,14 +677,3 @@ def _split_into_groups(n, max_group_size, mesh_dim_size):
       " = (num_groups=%d group_size=%d)" %
       (n, max_group_size, mesh_dim_size, num_groups, group_size))
   return num_groups, group_size
-
-
-def _tensor_dim_to_mesh_dim_size(hparams, tensor_dim):
-  """Inspect hparams to figure out how many ways tensor_dim gets split."""
-  layout_rules = mtf.convert_to_layout_rules(hparams.layout)
-  mesh_shape = mtf.convert_to_shape(hparams.mesh_shape)
-  mesh_axis = layout_rules.tensor_dimension_to_mesh_axis(tensor_dim, mesh_shape)
-  if mesh_axis is None:
-    return 1
-  else:
-    return mesh_shape.dims[mesh_axis].size
diff --git a/tensor2tensor/models/research/moe_experiments.py b/tensor2tensor/models/research/moe_experiments.py
index d0c3e7164..445abc419 100644
--- a/tensor2tensor/models/research/moe_experiments.py
+++ b/tensor2tensor/models/research/moe_experiments.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Experiments with mixture-of-experts architectures."""
+"""Languaeg modeling experiments in mtf."""
 
 
 from __future__ import absolute_import
@@ -30,12 +30,12 @@ def xmoe_dense_4k():
   """Series of architectural experiments on cheap language models.
 
   For all of these architectures, we run on languagemodel_lm1b8k_packed
-  for 32k-96 steps (1-3 epochs) on one TPU (8 cores).
+  for 32000 steps.
 
   All log-perplexities are per-token - multiply by 1.298 for per-word
 
   Results:
-  model             params(M)  einsum  alltoall  mxu-util  log-ppl(1ep) (3ep)
+  model             params(M)  einsum  alltoall  mxu-util  log-ppl
   xmoe_dense_4k     30         3.0e12  0         45%        3.31
   xmoe_dense_8k     46         4.7e12  0         49%        3.24
   xmoe_dense_64k    282        2.8e13  0                    3.06
@@ -44,7 +44,7 @@ def xmoe_dense_4k():
   xmoe_2d           282        5.3e12  7.6e8     34%        3.06
 
   Trained at 4x the batch size:
-  xmoe_2d_88        1090       2.1e13  3.0e9     24%
+  xmoe_2d_88        1090       2.1e13  3.0e9     24%        3.07
 
   Note: configurations and code are likely to change without notice.
 
@@ -52,6 +52,9 @@ def xmoe_dense_4k():
     a hparams
   """
   hparams = mtf_transformer.mtf_transformer_base_lm()
+  hparams.attention_dropout = 0.0
+  hparams.relu_dropout = 0.0
+  hparams.layer_prepostprocess_dropout = 0.0
 
   # The following hparams are constant across all these experiments.
   hparams.batch_size = 128
@@ -78,6 +81,7 @@ def xmoe_dense_8k():
 
 @registry.register_hparams
 def xmoe_dense_64k():
+  """Very wide layer- run on 4x4."""
   hparams = xmoe_dense_4k()
   hparams.d_ff = 65536
   hparams.mesh_shape = "model:4,batch:8"
@@ -86,7 +90,7 @@ def xmoe_dense_64k():
 
 @registry.register_hparams
 def xmoe_top_2():
-  """Mixture of experts."""
+  """Mixture of experts (16 experts)."""
   hparams = xmoe_dense_4k()
   moe.set_default_moe_hparams(hparams)
   hparams.mesh_shape = "all:8"
@@ -102,33 +106,9 @@ def xmoe_top_2_c15():
   return hparams
 
 
-@registry.register_hparams
-def mtf_transformer_lm_moe():
-  """Mixture of experts language model.
-
-  Compare to mtf_transformer.mtf_transformer_lm_baseline()
-
-  Run this on 2x2 on languagemodel_lm1b32k_packed for 272000 steps (10 epochs)
-  900M params.
-
-  Results on LM1B:
-         params/10^9  log-ppl(per-token)
-         0.90         TODO(noam): rerun experiment
-
-  Returns:
-    a hparams
-  """
-  hparams = mtf_transformer.mtf_transformer_lm_baseline()
-  hparams.decoder_layers = ["att", "moe"] * 4
-  moe.set_default_moe_hparams(hparams)
-  hparams.mesh_shape = "all:8"
-  hparams.layout = "batch:all;experts:all"
-  return hparams
-
-
 @registry.register_hparams
 def xmoe_2d():
-  """Two-dimensional hierarchical mixture of experts."""
+  """Two-dimensional hierarchical mixture of 16 experts."""
   hparams = xmoe_top_2()
   hparams.decoder_layers = ["att", "hmoe"] * 4
   hparams.mesh_shape = "b0:2;b1:4"
@@ -164,29 +144,35 @@ def xmoe_2d_c15():
 
 
 @registry.register_hparams
-def xmoe_2d_88():
-  """Two-dimensional hierarchical mixture of experts."""
+def xmoe_2d_x64():
+  """Two-dimensional hierarchical mixture of 64 experts."""
   hparams = xmoe_2d()
-  hparams.mesh_shape = "b0:4;b1:8"
-  hparams.batch_size = 512
+  # hparams.mesh_shape = "b0:4;b1:8"
   hparams.outer_batch_size = 4
   hparams.moe_num_experts = [8, 8]
   return hparams
 
 
 @registry.register_hparams
-def xmoe_wiki_base(sz):
-  """Series of architectural experiments on wikipedia text.
+def xmoe2_dense(sz):
+  """Series of architectural experiments on language modeling.
 
-  For all of these architectures, we run on languagemodel_wiki_noref_v8k_l1k
-  for 3 epochs.  (training set has ~7390100 sequences each of length 1024)
-  1 epoch = 57500 steps at batch_size=128
+  Larger models than the ones above.
 
-  Results:
-  model             params(M)  einsum  alltoall  mxu-util  log-ppl(1ep) (3ep)
+  All models are trained on sequences of 1024 tokens.
+
+  We assume infinite training data, so no dropout necessary.
+  We process 2^36 tokens in training = 524288 steps at batch size 128
+
+  TODO(noam): find a large enough dataset for these experiments.
+
+  You can use languagemodel_wiki_noref_v8k_l1k, but this is too small,
+  so training will cover about 9 epochs.
 
   Note: configurations and code are likely to change without notice.
 
+  Run on TPU 4x4 for 524288 steps unless otherwise indicated.
+
   Args:
     sz: an integer
 
@@ -194,49 +180,54 @@ def xmoe_wiki_base(sz):
     a hparams
   """
   hparams = mtf_transformer.mtf_transformer_paper_lm(sz)
-
+  hparams.attention_dropout = 0.0
+  hparams.relu_dropout = 0.0
+  hparams.layer_prepostprocess_dropout = 0.0
   hparams.max_length = 1024
   hparams.batch_size = 128
-  hparams.learning_rate_decay_steps = 57500
+  hparams.learning_rate_schedule = "rsqrt_decay*linear_decay"
+  hparams.learning_rate_decay_steps = 65536
   hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
   hparams.mesh_shape = "batch:32"
   return hparams
 
 
 @registry.register_hparams
-def xmoe_wiki_base_0():
-  return xmoe_wiki_base(0)
+def xmoe2_dense_0():
+  return xmoe2_dense(0)
 
 
 @registry.register_hparams
-def xmoe_wiki_base_1():
-  return xmoe_wiki_base(1)
+def xmoe2_dense_1():
+  return xmoe2_dense(1)
 
 
 @registry.register_hparams
-def xmoe_wiki_base_2():
-  return xmoe_wiki_base(2)
+def xmoe2_dense_2():
+  return xmoe2_dense(2)
 
 
 @registry.register_hparams
-def xmoe_wiki_base_3():
-  return xmoe_wiki_base(3)
+def xmoe2_dense_3():
+  return xmoe2_dense(3)
 
 
 @registry.register_hparams
-def xmoe_wiki_x():
-  """Baseline set of parameters for mixture-of-experts.
+def xmoe2_v1():
+  """Model incorporating mixture-of-experts and local-attention.
 
   ~6B parameters
 
+  32 experts in 3 hierarchichal moe layers.
+
   Returns:
     a hparams
   """
-  hparams = xmoe_wiki_base(0)
+  hparams = xmoe2_dense(0)
   moe.set_default_moe_hparams(hparams)
   hparams.decoder_layers = (
-      ["att", "drd", "att", "drd", "att", "hmoe"] * 3 +
-      ["att", "drd", "att", "drd"])
+      ["local_att", "local_att", "drd",
+       "att", "drd", "local_att", "local_att", "hmoe"] * 4)[:-1]
   hparams.d_ff = 2048
   hparams.d_kv = 128
   hparams.moe_hidden_size = 32768
@@ -249,32 +240,26 @@ def xmoe_wiki_x():
 
 
 @registry.register_hparams
-def xmoe_wiki_x_a32():
-  """Test 32-bit activations."""
-  hparams = xmoe_wiki_x()
-  hparams.activation_dtype = "float32"
-  return hparams
-
-
-@registry.register_hparams
-def xmoe_wiki_x128():
-  """128 experts, ~25B params on 8x8."""
-  hparams = xmoe_wiki_x()
+def xmoe2_v1_x128():
+  """128 experts, ~25B params - Train for 131072 steps on 8x8."""
+  hparams = xmoe2_v1()
   hparams.moe_num_experts = [16, 8]
   hparams.outer_batch_size = 8
   hparams.mesh_shape = "b0:8;b1:16"
   hparams.batch_size = 512
-  hparams.learning_rate_decay_steps = 14375
+  hparams.learning_rate_decay_steps = 16384
   return hparams
 
 
 @registry.register_hparams
-def xmoe_wiki_x_tiny():
+def xmoe2_tiny():
   """Test on local cpu."""
-  hparams = xmoe_wiki_x()
-  hparams.decoder_layers = (["att", "drd", "hmoe"] * 2 + ["att", "drd"])
+  hparams = xmoe2_v1()
+  hparams.decoder_layers = ["local_att", "att", "drd", "hmoe"]
+  hparams.d_model = 128
   hparams.moe_hidden_size = 512
-  hparams.batch_size = 16
+  hparams.outer_batch_size = 0
+  hparams.batch_size = 2
   hparams.mesh_shape = ""
   hparams.activation_dtype = "float32"
   return hparams

From e3702c88da9ac08b99e5d3e17e72e695928785f2 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Tue, 6 Nov 2018 02:31:45 +0100
Subject: [PATCH 1180/2720] Introduce PolicyLearner, unified interface for base
 rl algorithms used in Model-Based BRL (#1201)

---
 tensor2tensor/rl/policy_learner.py            |  81 +++++++++++
 tensor2tensor/rl/trainer_model_based.py       | 137 +++++++++---------
 .../rl/trainer_model_based_params.py          |   3 +-
 3 files changed, 148 insertions(+), 73 deletions(-)
 create mode 100644 tensor2tensor/rl/policy_learner.py

diff --git a/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
new file mode 100644
index 000000000..b2fa49ceb
--- /dev/null
+++ b/tensor2tensor/rl/policy_learner.py
@@ -0,0 +1,81 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Unified interface for different RL algorithms."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.rl import rl_trainer_lib
+
+
+class PolicyLearner(object):
+
+  def __init__(self, frame_stack_size, event_dir, agent_model_dir):
+    self.frame_stack_size = frame_stack_size
+    self.event_dir = event_dir
+    self.agent_model_dir = agent_model_dir
+
+  def train(self, env_fn, hparams, target_num_epochs, simulated, epoch):
+    # TODO(konradczechowski): target_num_steps instead of epochs
+    # TODO(konradczechowski): move 'simulated' to  batch_env
+    raise NotImplementedError()
+
+  def evaluate(self, env_fn, hparams, stochastic):
+    raise NotImplementedError()
+
+
+class PPOLearner(PolicyLearner):
+
+  def train(self, env_fn, hparams, target_num_epochs, simulated, epoch):
+    hparams.set_hparam("epochs_num", target_num_epochs)
+
+    if simulated:
+      simulated_str = "sim"
+      hparams.save_models_every_epochs = 10
+    else:
+      # TODO(konradczechowski): refactor ppo
+      assert hparams.num_agents == 1
+      # We do not save model, as that resets frames that we need at restarts.
+      # But we need to save at the last step, so we set it very high.
+      hparams.save_models_every_epochs = 1000000
+      simulated_str = "real"
+
+    # TODO(konradczechowski) refactor ppo, pass these as arguments
+    # (not inside hparams). Do the same in evaluate()
+    hparams.add_hparam("force_beginning_resets", simulated)
+    hparams.add_hparam("env_fn", env_fn)
+    hparams.add_hparam("frame_stack_size", self.frame_stack_size)
+    name_scope = "ppo_{}{}".format(simulated_str, epoch + 1)
+
+    rl_trainer_lib.train(hparams, self.event_dir + simulated_str,
+                         self.agent_model_dir, name_scope=name_scope)
+
+  def evaluate(self, env_fn, hparams, stochastic):
+    if stochastic:
+      policy_to_actions_lambda = lambda policy: policy.sample()
+    else:
+      policy_to_actions_lambda = lambda policy: policy.mode()
+    hparams.add_hparam(
+        "policy_to_actions_lambda", policy_to_actions_lambda
+    )
+    hparams.add_hparam("force_beginning_resets", False)
+    hparams.add_hparam("env_fn", env_fn)
+    hparams.add_hparam("frame_stack_size", self.frame_stack_size)
+
+    rl_trainer_lib.evaluate(hparams, self.agent_model_dir)
+
+
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 4a2fab766..1895da252 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -41,8 +41,8 @@
 from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl
-from tensor2tensor.rl import rl_trainer_lib
 from tensor2tensor.rl import trainer_model_based_params
+from tensor2tensor.rl.policy_learner import PPOLearner
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -52,6 +52,11 @@
 FLAGS = flags.FLAGS
 
 
+LEARNERS = dict(
+    ppo=PPOLearner,
+)
+
+
 def real_ppo_epoch_increment(hparams):
   """PPO increment."""
   if hparams.gather_ppo_real_env_data:
@@ -171,25 +176,19 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
   getattr(exp, schedule)()
 
 
+def _update_hparams_from_hparams(target_hparams, source_hparams, prefix):
+  "Copy a subset of hparams to target_hparams"
+  for param_name in target_hparams.values().keys():
+    prefixed_param_name = prefix + param_name
+    if prefixed_param_name in source_hparams:
+      target_hparams.set_hparam(param_name,
+                                source_hparams.get(prefixed_param_name))
+
+
 def train_agent(real_env, agent_model_dir, event_dir, world_model_dir, data_dir,
-                hparams, completed_ppo_epochs_num, epoch=0,
-                is_final_epoch=False):
+                hparams, completed_epochs_num, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the simulated environment."""
   del data_dir
-  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
-  ppo_params_names = ["epochs_num", "epoch_length",
-                      "learning_rate", "num_agents",
-                      "optimization_epochs", "eval_every_epochs"]
-
-  for param_name in ppo_params_names:
-    ppo_param_name = "ppo_" + param_name
-    if ppo_param_name in hparams:
-      ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
-
-  completed_ppo_epochs_num += sim_ppo_epoch_increment(hparams, is_final_epoch)
-  ppo_hparams.epochs_num = completed_ppo_epochs_num
-
-  ppo_hparams.save_models_every_epochs = 10
 
   frame_stack_size = hparams.frame_stack_size
   initial_frame_rollouts = real_env.current_epoch_rollouts(
@@ -218,55 +217,54 @@ def initial_frame_chooser(batch_size):
         [frame.observation.decode() for frame in initial_frame_stack]
         for initial_frame_stack in initial_frames
     ])
-
   env_fn = make_simulated_env_fn(
       real_env, hparams, hparams.ppo_num_agents, initial_frame_chooser,
       world_model_dir
   )
-  ppo_hparams.add_hparam("env_fn", env_fn)
-  ppo_hparams.add_hparam("force_beginning_resets", True)
-  ppo_hparams.add_hparam("frame_stack_size", frame_stack_size)
+  base_algo_str = hparams.base_algo
+  train_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
 
-  rl_trainer_lib.train(ppo_hparams, event_dir + "sim", agent_model_dir,
-                       name_scope="ppo_sim%d" % (epoch + 1))
+  _update_hparams_from_hparams(train_hparams, hparams, base_algo_str + "_")
+  # train_hparams.add_hparam("simulated", True)
 
-  return completed_ppo_epochs_num
+  learner = LEARNERS[base_algo_str](frame_stack_size, event_dir,
+                                    agent_model_dir)
+  learner.train(env_fn, train_hparams, completed_epochs_num,
+                simulated=True, epoch=epoch)
+
+  return completed_epochs_num
 
 
 def train_agent_real_env(
     env, agent_model_dir, event_dir, data_dir,
-    hparams, completed_ppo_epochs_num, epoch=0, is_final_epoch=False):
+    hparams, completed_epochs_num, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the real environment."""
   del is_final_epoch, data_dir
-  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)
-  ppo_params_names = ["epochs_num", "epoch_length",
-                      "learning_rate", "num_agents", "eval_every_epochs",
-                      "optimization_epochs", "effective_num_agents"]
 
-  # This should be overridden.
-  ppo_hparams.add_hparam("effective_num_agents", None)
-  for param_name in ppo_params_names:
-    ppo_param_name = "real_ppo_"+ param_name
-    if ppo_param_name in hparams:
-      ppo_hparams.set_hparam(param_name, hparams.get(ppo_param_name))
+  base_algo_str = hparams.base_algo
 
-  completed_ppo_epochs_num += real_ppo_epoch_increment(hparams)
-  ppo_hparams.epochs_num = completed_ppo_epochs_num
-  # We do not save model, as that resets frames that we need at restarts.
-  # But we need to save at the last step, so we set it very high.
-  ppo_hparams.save_models_every_epochs = 1000000
+  train_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
+  _update_hparams_from_hparams(train_hparams, hparams,
+                               "real_" + base_algo_str + "_")
 
-  ppo_hparams.add_hparam("env_fn", rl.make_real_env_fn(env))
-  ppo_hparams.add_hparam("force_beginning_resets", False)
-  ppo_hparams.add_hparam("frame_stack_size", hparams.frame_stack_size)
+  # TODO(konradczechowski): add effective_num_agents to ppo_atari_base etc.
+  # this requires refactoring ppo.
+  # This should be overridden.
+  train_hparams.add_hparam("effective_num_agents",
+                           hparams.real_ppo_effective_num_agents)
 
-  rl_trainer_lib.train(ppo_hparams, event_dir + "real", agent_model_dir,
-                       name_scope="ppo_real%d" % (epoch + 1))
+  completed_epochs_num += real_ppo_epoch_increment(hparams)
+  train_hparams.epochs_num = completed_epochs_num
 
+  env_fn = rl.make_real_env_fn(env)
+  learner = LEARNERS[base_algo_str](hparams.frame_stack_size, event_dir,
+                                    agent_model_dir)
+  learner.train(env_fn, train_hparams, completed_epochs_num,
+                simulated=False, epoch=epoch)
   # Save unfinished rollouts to history.
   env.reset()
 
-  return completed_ppo_epochs_num
+  return completed_epochs_num
 
 
 def train_world_model(
@@ -315,23 +313,21 @@ def setup_env(hparams, batch_size):
 
 def evaluate_single_config(hparams, agent_model_dir):
   """Evaluate the PPO agent in the real environment."""
-  eval_hparams = trainer_lib.create_hparams(hparams.ppo_params)
+  eval_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
   eval_hparams.num_agents = hparams.num_agents
+  eval_hparams.add_hparam('stochastic', hparams.stochastic)
   env = setup_env(hparams, batch_size=hparams.num_agents)
-  env_fn = rl.make_real_env_fn(env)
-  eval_hparams.add_hparam("env_fn", env_fn)
-  eval_hparams.add_hparam(
-      "policy_to_actions_lambda", hparams.policy_to_actions_lambda
-  )
-  eval_hparams.add_hparam("frame_stack_size", hparams.frame_stack_size)
-  eval_hparams.add_hparam("force_beginning_resets", False)
-
   env.start_new_epoch(0)
-  rl_trainer_lib.evaluate(eval_hparams, agent_model_dir)
+  env_fn = rl.make_real_env_fn(env)
+  learner = LEARNERS[hparams.base_algo](hparams.frame_stack_size,
+                                        event_dir=None,
+                                        agent_model_dir=agent_model_dir)
+  learner.evaluate(env_fn, eval_hparams, eval_hparams.stochastic)
   rollouts = env.current_epoch_rollouts()[:hparams.num_agents]
   env.close()
 
-  assert len(rollouts) == hparams.num_agents
+  assert len(rollouts) == hparams.num_agents, "{} {}".format(len(rollouts),
+                                                             hparams.num_agents)
   return tuple(
       compute_mean_reward(rollouts, clipped) for clipped in (True, False)
   )
@@ -339,26 +335,23 @@ def evaluate_single_config(hparams, agent_model_dir):
 
 def evaluate_all_configs(hparams, agent_model_dir):
   """Evaluate the agent with multiple eval configurations."""
-  def make_eval_hparams(hparams, policy_to_action, max_num_noops):
+  def make_eval_hparams(hparams, stochastic, max_num_noops):
     hparams = copy.copy(hparams)
     hparams.add_hparam("num_agents", hparams.eval_num_agents)
-    hparams.add_hparam("policy_to_actions_lambda", {
-        "sample": lambda policy: policy.sample(),
-        "mode": lambda policy: policy.mode()
-    }[policy_to_action])
+    hparams.add_hparam("stochastic", stochastic)
     hparams.max_num_noops = max_num_noops
     return hparams
 
   metrics = {}
   # Iterate over all combinations of picking actions by sampling/mode and
   # whether to do initial no-ops.
-  for policy_to_action in ("mode", "sample"):
+  for stochastic in (True, False):
     for max_num_noops in (hparams.eval_max_num_noops, 0):
-      eval_hparams = make_eval_hparams(hparams, policy_to_action, max_num_noops)
+      eval_hparams = make_eval_hparams(hparams, stochastic, max_num_noops)
       scores = evaluate_single_config(eval_hparams, agent_model_dir)
       for (score, clipped) in zip(scores, (True, False)):
-        metric_name = "mean_reward/eval/{}_{}_max_noops_{}".format(
-            policy_to_action, max_num_noops,
+        metric_name = "mean_reward/eval/stochastic_{}_max_noops_{}_{}".format(
+            stochastic, max_num_noops,
             "clipped" if clipped else "unclipped"
         )
         metrics[metric_name] = score
@@ -523,9 +516,9 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   tf.logging.info("Initial training of PPO in real environment.")
   ppo_event_dir = os.path.join(directories["world_model"],
                                "ppo_summaries/initial")
-  completed_ppo_epochs_num = train_agent_real_env(
+  completed_epochs_num = train_agent_real_env(
       env, ppo_model_dir, ppo_event_dir, data_dir, hparams,
-      completed_ppo_epochs_num=0, epoch=epoch, is_final_epoch=False
+      completed_epochs_num=0, epoch=epoch, is_final_epoch=False
   )
   metrics["mean_reward/train/clipped"] = compute_mean_reward(
       env.current_epoch_rollouts(), clipped=True
@@ -560,9 +553,9 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     if not hparams.ppo_continue_training:
       ppo_model_dir = ppo_event_dir
 
-    completed_ppo_epochs_num = train_agent(
+    completed_epochs_num = train_agent(
         env, ppo_model_dir, ppo_event_dir,
-        directories["world_model"], data_dir, hparams, completed_ppo_epochs_num,
+        directories["world_model"], data_dir, hparams, completed_epochs_num,
         epoch=epoch, is_final_epoch=is_final_epoch
     )
 
@@ -570,9 +563,9 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
     # Train PPO on real env (short)
     log("Training PPO in real environment.")
-    completed_ppo_epochs_num = train_agent_real_env(
+    completed_epochs_num = train_agent_real_env(
         env, ppo_model_dir, ppo_event_dir, data_dir, hparams,
-        completed_ppo_epochs_num, epoch=epoch, is_final_epoch=is_final_epoch
+        completed_epochs_num, epoch=epoch, is_final_epoch=is_final_epoch
     )
 
     if hparams.stop_loop_early:
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 7dbc7216b..e270d5c9c 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -54,7 +54,8 @@ def rlmb_base():
       num_real_env_frames=96000,
       generative_model="next_frame_basic_deterministic",
       generative_model_params="next_frame_pixel_noise",
-      ppo_params="ppo_atari_base",
+      base_algo="ppo",
+      base_algo_params="ppo_atari_base",
       autoencoder_train_steps=0,
       autoencoder_train_steps_initial_multiplier=10,
       autoencoder_hparams_set="autoencoder_discrete_pong",

From 08350ea8d57c7ccd000d7a890ef6f7910cb5d060 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Mon, 5 Nov 2018 17:32:04 -0800
Subject: [PATCH 1181/2720] internal merge of PR #1201

PiperOrigin-RevId: 220203241
---
 tensor2tensor/rl/policy_learner.py      | 4 ++--
 tensor2tensor/rl/trainer_model_based.py | 6 +++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
index b2fa49ceb..60569dd5a 100644
--- a/tensor2tensor/rl/policy_learner.py
+++ b/tensor2tensor/rl/policy_learner.py
@@ -23,6 +23,7 @@
 
 
 class PolicyLearner(object):
+  """API for policy learners."""
 
   def __init__(self, frame_stack_size, event_dir, agent_model_dir):
     self.frame_stack_size = frame_stack_size
@@ -39,6 +40,7 @@ def evaluate(self, env_fn, hparams, stochastic):
 
 
 class PPOLearner(PolicyLearner):
+  """PPO for policy learning."""
 
   def train(self, env_fn, hparams, target_num_epochs, simulated, epoch):
     hparams.set_hparam("epochs_num", target_num_epochs)
@@ -77,5 +79,3 @@ def evaluate(self, env_fn, hparams, stochastic):
     hparams.add_hparam("frame_stack_size", self.frame_stack_size)
 
     rl_trainer_lib.evaluate(hparams, self.agent_model_dir)
-
-
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 1895da252..f6fef8d64 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -177,7 +177,7 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
 
 
 def _update_hparams_from_hparams(target_hparams, source_hparams, prefix):
-  "Copy a subset of hparams to target_hparams"
+  """Copy a subset of hparams to target_hparams."""
   for param_name in target_hparams.values().keys():
     prefixed_param_name = prefix + param_name
     if prefixed_param_name in source_hparams:
@@ -188,7 +188,7 @@ def _update_hparams_from_hparams(target_hparams, source_hparams, prefix):
 def train_agent(real_env, agent_model_dir, event_dir, world_model_dir, data_dir,
                 hparams, completed_epochs_num, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the simulated environment."""
-  del data_dir
+  del data_dir, is_final_epoch
 
   frame_stack_size = hparams.frame_stack_size
   initial_frame_rollouts = real_env.current_epoch_rollouts(
@@ -315,7 +315,7 @@ def evaluate_single_config(hparams, agent_model_dir):
   """Evaluate the PPO agent in the real environment."""
   eval_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
   eval_hparams.num_agents = hparams.num_agents
-  eval_hparams.add_hparam('stochastic', hparams.stochastic)
+  eval_hparams.add_hparam("stochastic", hparams.stochastic)
   env = setup_env(hparams, batch_size=hparams.num_agents)
   env.start_new_epoch(0)
   env_fn = rl.make_real_env_fn(env)

From 1f8bc1872c3ea1cc2db89e9b3a0f7f901e8a6133 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 5 Nov 2018 17:49:14 -0800
Subject: [PATCH 1182/2720] Skip latent in stochastic video model training to
 improve action conditioning.

PiperOrigin-RevId: 220205420
---
 tensor2tensor/models/video/basic_stochastic.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index b3078db6f..b33b252e3 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -138,7 +138,13 @@ def add_bits(layer, bits):
         bit_p *= hparams.latent_rnn_max_sampling
         bits = tf.where(which_bit < bit_p, bits_pred, bits)
 
-    return add_bits(layer, bits), pred_loss
+    res = add_bits(layer, bits)
+    # During training, sometimes skip the latent to help action-conditioning.
+    res_p = common_layers.inverse_lin_decay(hparams.latent_rnn_warmup_steps / 2)
+    res_p *= hparams.latent_use_max_probability
+    res_rand = tf.random_uniform([layer_shape[0]])
+    res = tf.where(res_rand < res_p, res, layer)
+    return res, pred_loss
 
 
 @registry.register_hparams
@@ -190,7 +196,7 @@ def next_frame_basic_stochastic_discrete():
   hparams.scheduled_sampling_mode = "prob_inverse_lin"
   hparams.scheduled_sampling_decay_steps = 40000
   hparams.scheduled_sampling_max_prob = 1.0
-  hparams.dropout = 0.2
+  hparams.dropout = 0.15
   hparams.filter_double_steps = 3
   hparams.hidden_size = 96
   hparams.learning_rate_constant = 0.005
@@ -200,10 +206,11 @@ def next_frame_basic_stochastic_discrete():
   hparams.add_hparam("bottleneck_noise", 0.1)
   hparams.add_hparam("discretize_warmup_steps", 40000)
   hparams.add_hparam("latent_rnn_warmup_steps", 40000)
-  hparams.add_hparam("latent_rnn_max_sampling", 0.7)
+  hparams.add_hparam("latent_rnn_max_sampling", 0.6)
+  hparams.add_hparam("latent_use_max_probability", 0.8)
   hparams.add_hparam("full_latent_tower", False)
   hparams.add_hparam("latent_predictor_state_size", 128)
-  hparams.add_hparam("latent_predictor_temperature", 0.5)
+  hparams.add_hparam("latent_predictor_temperature", 0.9)
   hparams.add_hparam("complex_addn", True)
   return hparams
 
@@ -224,4 +231,5 @@ def next_frame_stochastic_discrete_range(rhp):
 def next_frame_stochastic_discrete_latent_range(rhp):
   rhp.set_float("latent_rnn_max_sampling", 0.1, 0.9)
   rhp.set_float("latent_predictor_temperature", 0.1, 1.2)
+  rhp.set_float("latent_use_max_probability", 0.4, 1.0)
   rhp.set_float("dropout", 0.1, 0.4)

From c78ad1eb548899f4544a645e21e8ae5fae4f9d15 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 5 Nov 2018 18:01:14 -0800
Subject: [PATCH 1183/2720] adding reset internal states through additional
 input features.

PiperOrigin-RevId: 220206689
---
 tensor2tensor/models/video/base.py           | 24 +++++++++++++++++++-
 tensor2tensor/models/video/sv2p.py           |  5 ++++
 tensor2tensor/rl/envs/simulated_batch_env.py | 15 +++++++++---
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index ad0a23e43..5996498b7 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -151,6 +151,10 @@ def init_internal_states(self):
     """
     return None
 
+  def reset_internal_states_ops(self):
+    """Resets internal states to initial values."""
+    return [[tf.no_op()]]
+
   def load_internal_states_ops(self):
     """Loade internal states from class variables."""
     return [[tf.no_op()]]
@@ -602,7 +606,25 @@ def merge(inputs, targets):
       actions = merge(features["input_action"], features["target_action"])
     if self.has_rewards:
       rewards = merge(features["input_reward"], features["target_reward"])
-    return self.__process(frames, actions, rewards, frames_raw)
+
+    # Reset the internal states if the reset_internal_states has been
+    # passed as a feature and has greater value than 0.
+    if self.is_recurrent_model and self.internal_states is not None:
+      def reset_func():
+        reset_ops = flat_lists(self.reset_internal_states_ops())
+        with tf.control_dependencies(reset_ops):
+          return tf.no_op()
+      if self.is_predicting and "reset_internal_states" in features:
+        reset = features["reset_internal_states"]
+        reset = tf.greater(tf.reduce_sum(reset), 0.5)
+        reset_ops = tf.cond(reset, reset_func, tf.no_op)
+      else:
+        reset_ops = reset_func()
+      with tf.control_dependencies([reset_ops]):
+        frames[0] = tf.identity(frames[0])
+
+    with tf.control_dependencies([frames[0]]):
+      return self.__process(frames, actions, rewards, frames_raw)
 
 
 def next_frame_base():
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 04774782d..15a36b5db 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -455,6 +455,11 @@ def init_internal_states(self):
         states.append((c, h))
       return states
 
+  def reset_internal_states_ops(self):
+    zeros = [(tf.zeros_like(c), tf.zeros_like(h))
+             for c, h in self.internal_states]
+    return self.save_internal_states_ops(zeros)
+
   def load_internal_states_ops(self):
     ops = [(c.read_value(), h.read_value()) for c, h in self.internal_states]
     return ops
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index e0fad49db..481954a65 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -122,6 +122,9 @@ def __init__(
         trainable=False
     )
 
+    self._reset_model = tf.get_variable(
+        "reset_model", [], trainable=False, initializer=tf.zeros_initializer())
+
     self._model_dir = model_dir
 
   def initialize(self, sess):
@@ -149,7 +152,9 @@ def simulate(self, action):
         hparams_target_frames = self._model.hparams.video_num_target_frames
         self._model.hparams.video_num_target_frames = 1
         model_output = self._model.infer(
-            {"inputs": history, "input_action": actions})
+            {"inputs": history,
+             "input_action": actions,
+             "reset_internal_states": self._reset_model.read_value()})
         self._model.hparams.video_num_target_frames = hparams_target_frames
 
       observ = tf.cast(tf.squeeze(model_output["targets"], axis=1),
@@ -181,7 +186,9 @@ def simulate(self, action):
         with tf.control_dependencies(
             [self._observ.assign(observ),
              self.history_buffer.move_by_one_element(observ)]):
-          return tf.identity(reward), tf.identity(done)
+          clear_reset_model_op = tf.assign(self._reset_model, tf.constant(0.0))
+          with tf.control_dependencies([clear_reset_model_op]):
+            return tf.identity(reward), tf.identity(done)
 
   def _reset_non_empty(self, indices):
     """Reset the batch of environments.
@@ -195,7 +202,9 @@ def _reset_non_empty(self, indices):
     with tf.control_dependencies([self.history_buffer.reset(indices)]):
       with tf.control_dependencies([self._observ.assign(
           self.history_buffer.get_all_elements()[:, -1, ...])]):
-        return tf.gather(self._observ.read_value(), indices)
+        reset_model_op = tf.assign(self._reset_model, tf.constant(1.0))
+        with tf.control_dependencies([reset_model_op]):
+          return tf.gather(self._observ.read_value(), indices)
 
   @property
   def observ(self):

From f7f043b47b6e92ceea0a36f7b71b1ce5aa305e90 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Tue, 6 Nov 2018 10:51:52 -0800
Subject: [PATCH 1184/2720] Wikipedia multiproblems.

PiperOrigin-RevId: 220311353
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 .../data_generators/wiki_multi_problems.py    | 105 ++++++++++++++++++
 2 files changed, 106 insertions(+)
 create mode 100644 tensor2tensor/data_generators/wiki_multi_problems.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index b96c7daac..f402dd4e7 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -85,6 +85,7 @@
     "tensor2tensor.data_generators.wiki",
     "tensor2tensor.data_generators.wiki_lm",
     "tensor2tensor.data_generators.wiki_revision",
+    "tensor2tensor.data_generators.wiki_multi_problems",
     "tensor2tensor.data_generators.wikisum.wikisum",
     "tensor2tensor.data_generators.wikitext103",
     "tensor2tensor.data_generators.wsj_parsing",
diff --git a/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
new file mode 100644
index 000000000..d46b1f793
--- /dev/null
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -0,0 +1,105 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Data generators for Wiki LM and MNLI combined datasets."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators import cnn_dailymail
+from tensor2tensor.data_generators import multi_problem
+from tensor2tensor.data_generators import multinli
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators import wiki_lm
+from tensor2tensor.utils import registry
+
+
+@registry.register_problem
+class LanguagemodelEnWikiLMMultiNLISubwords(multi_problem.MultiProblem):
+  """Wiki LM and MNLI mixed problem class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelEnWikiLMMultiNLISubwords, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelEnWiki32k())
+    self.task_list.append(multinli.MultiNLIWikiLMSharedVocab())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
+
+
+@registry.register_problem
+class LanguagemodelEnWikiLMMultiNLISubwords64k(multi_problem.MultiProblem):
+  """Wiki LM and MNLI mixed problem class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelEnWikiLMMultiNLISubwords64k, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
+    self.task_list.append(multinli.MultiNLIWikiLMSharedVocab64k())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
+
+
+@registry.register_problem
+class LanguagemodelEnWikiLMShortMultiNLISubwords64k(multi_problem.MultiProblem):
+  """Wiki LM and MNLI mixed problem class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelEnWikiLMShortMultiNLISubwords64k, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelEnWiki64kShorter())
+    self.task_list.append(multinli.MultiNLIWikiLMSharedVocab64k())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
+
+
+@registry.register_problem
+class LanguagemodelEnWikiLMSummarizeCnndmSubwords(multi_problem.MultiProblem):
+  """Wiki LM and CNN/DM summarization mixed problem class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelEnWikiLMSummarizeCnndmSubwords, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelEnWiki32k())
+    self.task_list.append(
+        cnn_dailymail.SummarizeCnnDailymailWikiLMSharedVocab())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
+
+
+@registry.register_problem
+class LanguagemodelEnWikiLMSummarizeCnndmSubwords64k(
+    multi_problem.MultiProblem):
+  """Wiki LM and CNN/DM summarization mixed problem class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelEnWikiLMSummarizeCnndmSubwords64k, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
+    self.task_list.append(
+        cnn_dailymail.SummarizeCnnDailymailWikiLMSharedVocab64k())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD

From e42c13a3e55a09038873485f2cd83ca991de5849 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 6 Nov 2018 11:45:41 -0800
Subject: [PATCH 1185/2720] Delete and remove references to wiki_gec.py

PiperOrigin-RevId: 220322575
---
 tensor2tensor/data_generators/wiki_revision.py       | 5 +++--
 tensor2tensor/data_generators/wiki_revision_utils.py | 5 +----
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/wiki_revision.py b/tensor2tensor/data_generators/wiki_revision.py
index 18b1ad5dd..62bc90c63 100644
--- a/tensor2tensor/data_generators/wiki_revision.py
+++ b/tensor2tensor/data_generators/wiki_revision.py
@@ -20,8 +20,9 @@
 
 WikiRevision problem - fragment of older revision -> fragment of newer revision.
 
-This implements data extraction from wikipedia as desribed in the arXiv paper,
-Weakly Supervised Grammatical Error Correction using Iterative Decoding.
+This implements data extraction from wikipedia as desribed in the paper,
+Weakly Supervised Grammatical Error Correction using Iterative Decoding
+(https://arxiv.org/pdf/1811.01710.pdf).
 """
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index 574fcb67a..027bd162b 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -13,10 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Utilties for data generation for Wikipedia Revisions/GEC problems.
-
-Wikipedia revisions problems are defined in wiki_history.py
-Wikipedia GEC problems are defined in wiki_gec.py
+"""Utilties for data generation for Wikipedia Revision problem.
 """
 
 

From 43d6117b341ca6c470ccefa77f16082afed42614 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Tue, 6 Nov 2018 11:51:28 -0800
Subject: [PATCH 1186/2720] update loop params after introducing base_algo.

PiperOrigin-RevId: 220323609
---
 tensor2tensor/rl/trainer_model_based_params.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index e270d5c9c..30ac93ec8 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -744,17 +744,17 @@ def training_loop_hparams_from_scoped_overrides(scoped_overrides, trial_id):
   model_hp_name = trial_hp_overrides.get(
       "loop.generative_model_params", loop_hp.generative_model_params)
   model_hp = registry.hparams(model_hp_name).parse(FLAGS.hparams)
-  ppo_params_name = trial_hp_overrides.get(
-      "loop.ppo_params", loop_hp.ppo_params)
-  ppo_hp = registry.hparams(ppo_params_name)
+  base_algo_params_name = trial_hp_overrides.get(
+      "loop.base_algo_params", loop_hp.base_algo_params)
+  algo_hp = registry.hparams(base_algo_params_name)
 
   # Merge them and then override with the scoped overrides
   combined_hp = merge_unscoped_hparams(
-      zip(HP_SCOPES, [loop_hp, model_hp, ppo_hp]))
+      zip(HP_SCOPES, [loop_hp, model_hp, algo_hp]))
   combined_hp.override_from_dict(trial_hp_overrides)
 
   # Split out the component hparams
-  loop_hp, model_hp, ppo_hp = (
+  loop_hp, model_hp, algo_hp = (
       split_scoped_hparams(HP_SCOPES, combined_hp))
 
   # Dynamic register the model hp and set the new name in loop_hp
@@ -762,10 +762,10 @@ def training_loop_hparams_from_scoped_overrides(scoped_overrides, trial_id):
   dynamic_register_hparams(model_hp_name, model_hp)
   loop_hp.generative_model_params = model_hp_name
 
-  # Dynamic register the PPO hp and set the new name in loop_hp
-  ppo_hp_name = "ppo_hp_%s" % str(trial_id)
-  dynamic_register_hparams(ppo_hp_name, ppo_hp)
-  loop_hp.ppo_params = ppo_hp_name
+  # Dynamic register the algo hp and set the new name in loop_hp
+  algo_hp_name = "algo_hp_%s" % str(trial_id)
+  dynamic_register_hparams(algo_hp_name, algo_hp)
+  loop_hp.base_algo_params = algo_hp_name
 
   return loop_hp
 

From 3ace50292d01828d5c909d6aec5845150215d63d Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Tue, 6 Nov 2018 22:15:54 +0100
Subject: [PATCH 1187/2720] Model-Based RL: Fix epochs count increment, remove
 gather_ppo_real_env_data. (#1203)

* Remove unused gather_ppo_real_env_data.

* Fix completed_num_epochs increment for simulated environment.
---
 tensor2tensor/rl/trainer_model_based.py       | 26 +++++++------------
 .../rl/trainer_model_based_params.py          |  3 ---
 2 files changed, 9 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index f6fef8d64..92148c811 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -59,16 +59,13 @@
 
 def real_ppo_epoch_increment(hparams):
   """PPO increment."""
-  if hparams.gather_ppo_real_env_data:
-    assert hparams.real_ppo_epochs_num is 0, (
-        "Should be put to 0 to enforce better readability"
-    )
-    return int(math.ceil(
-        hparams.num_real_env_frames /
-        (hparams.epochs * hparams.real_ppo_epoch_length)
-    ))
-  else:
-    return hparams.real_ppo_epochs_num
+  assert hparams.real_ppo_epochs_num is 0, (
+      "Should be put to 0 to enforce better readability"
+  )
+  return int(math.ceil(
+      hparams.num_real_env_frames /
+      (hparams.epochs * hparams.real_ppo_epoch_length)
+  ))
 
 
 def sim_ppo_epoch_increment(hparams, is_final_epoch):
@@ -188,7 +185,7 @@ def _update_hparams_from_hparams(target_hparams, source_hparams, prefix):
 def train_agent(real_env, agent_model_dir, event_dir, world_model_dir, data_dir,
                 hparams, completed_epochs_num, epoch=0, is_final_epoch=False):
   """Train the PPO agent in the simulated environment."""
-  del data_dir, is_final_epoch
+  del data_dir
 
   frame_stack_size = hparams.frame_stack_size
   initial_frame_rollouts = real_env.current_epoch_rollouts(
@@ -225,8 +222,7 @@ def initial_frame_chooser(batch_size):
   train_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
 
   _update_hparams_from_hparams(train_hparams, hparams, base_algo_str + "_")
-  # train_hparams.add_hparam("simulated", True)
-
+  completed_epochs_num += sim_ppo_epoch_increment(hparams, is_final_epoch)
   learner = LEARNERS[base_algo_str](frame_stack_size, event_dir,
                                     agent_model_dir)
   learner.train(env_fn, train_hparams, completed_epochs_num,
@@ -254,7 +250,6 @@ def train_agent_real_env(
                            hparams.real_ppo_effective_num_agents)
 
   completed_epochs_num += real_ppo_epoch_increment(hparams)
-  train_hparams.epochs_num = completed_epochs_num
 
   env_fn = rl.make_real_env_fn(env)
   learner = LEARNERS[base_algo_str](hparams.frame_stack_size, event_dir,
@@ -509,9 +504,6 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   metrics = {}
 
   # Collect data from the real environment with PPO or random policy.
-  # TODO(lukaszkaiser): do we need option not to gather_ppo_real_env_data?
-  # We could set learning_rate=0 if this flag == False.
-  assert hparams.gather_ppo_real_env_data
   ppo_model_dir = directories["ppo"]
   tf.logging.info("Initial training of PPO in real environment.")
   ppo_event_dir = os.path.join(directories["world_model"],
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 30ac93ec8..18c1b4e26 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -95,7 +95,6 @@ def rlmb_base():
       # In your experiments, you want to optimize this rate to your schedule.
       learning_rate_bump=3.0,
 
-      gather_ppo_real_env_data=True,
       real_ppo_epochs_num=0,
       # This needs to be divisible by real_ppo_effective_num_agents.
       real_ppo_epoch_length=16*200,
@@ -450,7 +449,6 @@ def rlmb_ae_base():
   hparams.ppo_params = "ppo_pong_ae_base"
   hparams.generative_model_params = "next_frame_ae"
   hparams.autoencoder_hparams_set = "autoencoder_discrete_pong"
-  hparams.gather_ppo_real_env_data = False
   hparams.autoencoder_train_steps = 5000
   hparams.resize_height_factor = 1
   hparams.resize_width_factor = 1
@@ -479,7 +477,6 @@ def rlmb_ae_tiny():
   hparams.ppo_params = "ppo_pong_ae_base"
   hparams.generative_model_params = "next_frame_ae_tiny"
   hparams.autoencoder_hparams_set = "autoencoder_discrete_tiny"
-  hparams.gather_ppo_real_env_data = False
   hparams.resize_height_factor = 1
   hparams.resize_width_factor = 1
   hparams.grayscale = False

From ace5548d996c1874ba3c148dc45b8028557539f0 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Tue, 6 Nov 2018 14:27:19 -0800
Subject: [PATCH 1188/2720] Add generation support to multiproblem.

PiperOrigin-RevId: 220352613
---
 .../data_generators/multi_problem.py          | 150 +++++++++++++-----
 .../data_generators/multi_problem_test.py     |  83 ----------
 tensor2tensor/layers/common_hparams.py        |  11 ++
 3 files changed, 119 insertions(+), 125 deletions(-)
 delete mode 100644 tensor2tensor/data_generators/multi_problem_test.py

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 5ab2051a5..5b3043cef 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -50,8 +50,11 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
     for task in self.task_list:
       task.generate_data(data_dir, tmp_dir, task_id)
 
-  def add_task_id(self, task, example, encoder):
+  def add_task_id(self, task, example, encoder, hparams, is_infer):
     """Convert example to code switching mode by adding a task id."""
+    if task.has_inputs:
+      example["inputs"] = example["inputs"][:-1]  # remove EOS token
+
     if hasattr(task, "class_labels"):
       if self.vocab_type == text_problems.VocabType.CHARACTER:
         # TODO(urvashik): handle the case where num_labels > 9
@@ -61,14 +64,29 @@ def add_task_id(self, task, example, encoder):
       elif self.vocab_type == text_problems.VocabType.SUBWORD:
         offset = encoder.vocab_size + len(self.task_list)
         example["targets"] = offset + example["targets"]
+    else:
+      # sequence with inputs and targets eg: summarization
+      if task.has_inputs:
+        if hparams.multiproblem_max_input_length > 0:
+          example["inputs"] = example[
+              "inputs"][:hparams.multiproblem_max_input_length]
+        # Do not truncate targets during inference with beam decoding.
+        if hparams.multiproblem_max_target_length > 0 and not is_infer:
+          example["targets"] = example[
+              "targets"][:hparams.multiproblem_max_target_length]
 
     if task.has_inputs:
-      inputs = example.pop("inputs")
-      concat_list = [inputs, [task.task_id], example["targets"]]
+      if is_infer:
+        concat_list = [example["inputs"], [task.task_id]]
+        example["inputs"] = tf.concat(concat_list, 0)
+      else:
+        inputs = example.pop("inputs")
+        concat_list = [inputs, [task.task_id], example["targets"]]
+        example["targets"] = tf.concat(concat_list, 0)
     else:
       concat_list = [[task.task_id], example["targets"]]
+      example["targets"] = tf.concat(concat_list, 0)
 
-    example["targets"] = tf.concat(concat_list, 0)
     min_task_id = min([t.task_id for t in self.task_list])
     example["task_id"] = tf.constant([task.task_id - min_task_id],
                                      dtype=tf.int64)
@@ -87,9 +105,12 @@ def get_hparams(self, model_hparams=None):
     vocab_size_inc = len(self.task_list)
     vocab_size_inc += self.get_max_num_classes()
     vocab_size = self._hparams.vocabulary["targets"].vocab_size
+    new_vocab_size = vocab_size + vocab_size_inc
+    if model_hparams.multiproblem_vocab_size > new_vocab_size:
+      new_vocab_size = model_hparams.multiproblem_vocab_size
     tf.logging.info("Old vocabulary size: %d" % vocab_size)
-    tf.logging.info("New vocabulary size: %d" % (vocab_size + vocab_size_inc))
-    self._hparams.vocab_size["targets"] = vocab_size + vocab_size_inc
+    tf.logging.info("New vocabulary size: %d" % new_vocab_size)
+    self._hparams.vocab_size["targets"] = new_vocab_size
     self._hparams.modality["targets"] = modalities.SymbolModality(
         model_hparams, self._hparams.vocab_size["targets"])
 
@@ -128,12 +149,15 @@ def dataset(self,
               shard=None,
               partition_id=0,
               num_partitions=1,
-              max_records=-1):
+              shuffle_buffer_size=1024,
+              max_records=-1,
+              only_last=False):
 
     # A list of datasets corresponding to the tasks in the task_list object
     # that need to be mixed.
     datasets = []
     is_training = mode == tf.estimator.ModeKeys.TRAIN
+    is_infer = mode == tf.estimator.ModeKeys.PREDICT
 
     primary_task = self.task_list[0]
     if primary_task.has_inputs:
@@ -142,11 +166,20 @@ def dataset(self,
     enc = primary_task.feature_encoders(data_dir=data_dir)["targets"]
 
     for idx, task in enumerate(self.task_list):
-      task_dataset = task.dataset(mode, data_dir, num_threads,
-                                  output_buffer_size, shuffle_files,
-                                  hparams, preprocess, dataset_split,
-                                  shard, partition_id, num_partitions,
-                                  max_records)
+      task_dataset = task.dataset(mode=mode,
+                                  data_dir=data_dir,
+                                  num_threads=num_threads,
+                                  output_buffer_size=output_buffer_size,
+                                  shuffle_files=shuffle_files,
+                                  hparams=hparams,
+                                  preprocess=preprocess,
+                                  dataset_split=dataset_split,
+                                  shard=shard,
+                                  partition_id=partition_id,
+                                  num_partitions=num_partitions,
+                                  shuffle_buffer_size=shuffle_buffer_size,
+                                  max_records=max_records,
+                                  only_last=only_last)
 
       if idx == 0:
         self.update_task_ids(enc)
@@ -155,9 +188,10 @@ def dataset(self,
         task_dataset = task_dataset.repeat()
 
       # pylint: disable=cell-var-from-loop
-      task_dataset = task_dataset.map(lambda x: self.add_task_id(task, x, enc))
+      task_dataset = task_dataset.map(
+          lambda x: self.add_task_id(task, x, enc, hparams, is_infer))
 
-      if not is_training:
+      if not is_training and not is_infer:
         zeros = tf.zeros([self._ADDED_EVAL_COUNT, 1], dtype=tf.int64)
         pad_data = tf.data.Dataset.from_tensor_slices({
             "targets": zeros,
@@ -169,7 +203,7 @@ def dataset(self,
       datasets.append(task_dataset)
 
     # Setup the problem hparams by setting them to the LM task hparams.
-    self.get_hparams()
+    self.get_hparams(model_hparams=hparams)
 
     if is_training:
       problem_step = tf.get_variable("problem_step",
@@ -214,6 +248,11 @@ def mix_data(example):
         # tasks.
         if hparams.multiproblem_mixing_schedule == MixingSchedule.EXPONENTIAL:
           prob = get_exp_sched_prob()
+          prob = tf.cond(
+              tf.equal(tf.floormod(
+                  problem_step, tf.cast(5e6, dtype=tf.int64)), 0),
+              lambda: tf.Print(prob, [prob], message="Probability"),
+              lambda: prob)
         elif hparams.multiproblem_mixing_schedule == MixingSchedule.CONSTANT:
           prob = get_const_sched_prob()
         elif hparams.multiproblem_mixing_schedule == MixingSchedule.PRETRAIN:
@@ -226,11 +265,6 @@ def mix_data(example):
                             hparams.multiproblem_mixing_schedule))
         tf.logging.info("Schedule mixing threshold "
                         "%.2f" % hparams.multiproblem_schedule_threshold)
-        prob = tf.cond(
-            tf.equal(tf.floormod(
-                problem_step, tf.cast(5e6, dtype=tf.int64)), 0),
-            lambda: tf.Print(prob, [prob], message="Probability"),
-            lambda: prob)
 
         def sample_task(curr_task, num_tasks_left, randnum):
           """A recursive function to sample a task.
@@ -271,14 +305,23 @@ def sample_task(curr_task, num_tasks_left, randnum):
       single_mtl_dataset = single_mtl_dataset.flat_map(mix_data)
 
     else:
-      single_mtl_dataset = tf.data.Dataset.zip(tuple(datasets)).flat_map(
-          self.flatten_zip)
+      if hparams.multiproblem_target_eval_only:
+        single_mtl_dataset = datasets[1]
+      else:
+        single_mtl_dataset = tf.data.Dataset.zip(tuple(datasets)).flat_map(
+            self.flatten_zip)
 
     return single_mtl_dataset
 
   def eval_metrics(self):
+    for task in self.task_list:
+      if "summarize" in task.name:
+        return [
+            metrics.Metrics.ACC, metrics.Metrics.NEG_LOG_PERPLEXITY,
+            metrics.Metrics.ROUGE_2_F, metrics.Metrics.ROUGE_L_F
+        ]
     return [
-        metrics.Metrics.ACC, metrics.Metrics.NEG_LOG_PERPLEXITY
+        metrics.Metrics.ACC, metrics.Metrics.NEG_LOG_PERPLEXITY,
     ]
 
   def update_task_ids(self, encoder):
@@ -294,6 +337,7 @@ def update_task_ids(self, encoder):
 
     for idx, _ in enumerate(self.task_list):
       self.task_list[idx].set_task_id(idx + offset)
+      print(self.task_list[idx].name)
       print(self.task_list[idx].task_id)
 
   def get_max_num_classes(self):
@@ -340,14 +384,18 @@ def aggregate_task_losses(hparams,
   loss_den = tf.minimum(tf.convert_to_tensor(1, dtype=tf.float32), loss_den)
 
   for task in hparams.problem.task_list[1:]:
-    if hasattr(task, "num_classes"):
-      # Loss only from the input sequence -- the auxiliary LM loss.
-      seq_loss_num, seq_loss_den = target_modality.loss(
-          logits, feature,
-          weights_fn=
-          lambda x: common_layers.weights_multi_problem_input(x, task.task_id))  # pylint: disable=cell-var-from-loop
-      seq_loss_num *= problem_hparams.loss_multiplier
+    # Loss only from the input sequence -- the auxiliary LM loss.
+    seq_loss_num, seq_loss_den = target_modality.loss(
+        logits, feature,
+        weights_fn=
+        lambda x: common_layers.weights_multi_problem_input(x, task.task_id))  # pylint: disable=cell-var-from-loop
+    seq_loss_num *= problem_hparams.loss_multiplier
+
+    # Unscaled sequence loss.
+    seq_loss = seq_loss_num / tf.maximum(1.0, seq_loss_den)
+    summaries.append([task.name+"_seq_loss", seq_loss])
 
+    if hasattr(task, "num_classes"):
       # Loss only from the classification label.
       label_loss_num, label_loss_den = target_modality.loss(
           logits, feature,
@@ -355,9 +403,7 @@ def aggregate_task_losses(hparams,
           lambda x: common_layers.weights_multi_problem(x, task.task_id))  # pylint: disable=cell-var-from-loop
       label_loss_num *= problem_hparams.loss_multiplier
 
-      # Unscaled losses.
-      seq_loss = seq_loss_num / tf.maximum(1.0, seq_loss_den)
-      summaries.append([task.name+"_seq_loss", seq_loss])
+      # Unscaled classification label loss.
       label_loss = label_loss_num / tf.maximum(1.0, label_loss_den)
       summaries.append([task.name+"_label_loss", label_loss])
 
@@ -372,17 +418,37 @@ def aggregate_task_losses(hparams,
 
       # This is the training loss for the optimizer after all the scaling.
       task_loss_val = seq_loss + label_loss
-      summaries.append([task.name+"_loss", task_loss_val])
 
-      # Adding 1 to the loss den for each task leads to averaging task losses,
-      # task with bigger loss will dominate.
-      # TODO(urvashik): Fix combination with other task losses - weighted
-      # average based on the number of examples from that task.
-      loss_num += task_loss_val
-      loss_den += tf.minimum(tf.convert_to_tensor(1, dtype=tf.float32),
-                             label_loss_den)
+      loss_den_ = label_loss_den
 
     else:
-      raise ValueError("Non-classification secondary tasks are not supported.")
+      # Loss only from the target sequence.
+      target_loss_num, target_loss_den = target_modality.loss(
+          logits, feature,
+          weights_fn=
+          lambda x: common_layers.weights_multi_problem(x, task.task_id))  # pylint: disable=cell-var-from-loop
+      target_loss_num *= problem_hparams.loss_multiplier
+
+      # Unscaled target sequence loss.
+      target_loss = target_loss_num / tf.maximum(1.0, target_loss_den)
+      summaries.append([task.name+"_target_loss", target_loss])
+
+      # Scaling.
+      if hparams.multiproblem_reweight_label_loss:
+        target_loss *= hparams.multiproblem_label_weight
+        seq_loss *= (1 - hparams.multiproblem_label_weight)
+
+      # This is the training loss for the optimizer after all the scaling.
+      task_loss_val = seq_loss + target_loss
+
+      loss_den_ = target_loss_den
+
+    summaries.append([task.name+"_loss", task_loss_val])
+    # Adding 1 to the loss den for each task leads to averaging task losses.
+    # TODO(urvashik): Fix combination with other task losses - weighted
+    # average based on the number of examples from that task.
+    loss_num += task_loss_val
+    loss_den += tf.minimum(tf.convert_to_tensor(1, dtype=tf.float32),
+                           loss_den_)
 
   return loss_num, loss_den, summaries
diff --git a/tensor2tensor/data_generators/multi_problem_test.py b/tensor2tensor/data_generators/multi_problem_test.py
deleted file mode 100644
index 66301b9d4..000000000
--- a/tensor2tensor/data_generators/multi_problem_test.py
+++ /dev/null
@@ -1,83 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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 MultiProblem."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.data_generators import multi_problem
-from tensor2tensor.data_generators import problem
-from tensor2tensor.data_generators import text_encoder
-from tensor2tensor.layers import modalities
-
-import tensorflow as tf
-
-
-# TODO(trandustin): This test problem is required in order for MultiProblem
-# to access vocab size via encoders. In a future change, enable MultiProblem to
-# access vocab size more explicitly from the Problem.
-class TestProblem(problem.Problem):
-  """Test problem."""
-
-  def __init__(self, input_vocab_size, target_vocab_size):
-    super(TestProblem, self).__init__(False, False)
-    self.input_vocab_size = input_vocab_size
-    self.target_vocab_size = target_vocab_size
-
-  def hparams(self, defaults, model_hparams):
-    hp = defaults
-    hp.modality = {
-        "inputs": modalities.SymbolModality,
-        "targets": modalities.SymbolModality,
-    }
-    hp.vocab_size = {
-        "inputs": self.input_vocab_size,
-        "targets": self.target_vocab_size,
-    }
-
-  def feature_encoders(self, data_dir):
-    encoders = {
-        "inputs": text_encoder.ByteTextEncoder(),
-        "targets": text_encoder.ByteTextEncoder(),
-    }
-    return encoders
-
-
-class TestMultiProblem(multi_problem.MultiProblem):
-  """Test multi-problem."""
-
-  def __init__(self):
-    super(TestMultiProblem, self).__init__()
-    self.task_list.append(TestProblem(2, 3))
-    self.task_list.append(TestProblem(4, 6))
-
-
-class MultiProblemTest(tf.test.TestCase):
-
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
-  def testProblemHparamsModality(self):
-    multiproblem = TestMultiProblem()
-    p_hparams = multiproblem.get_hparams()
-    self.assertIsInstance(p_hparams.modality["inputs"],
-                          modalities.SymbolModality)
-    self.assertEqual(p_hparams.modality["inputs"].top_dimensionality, 2)
-    self.assertIsInstance(p_hparams.modality["targets"],
-                          modalities.SymbolModality)
-    self.assertEqual(p_hparams.modality["targets"].top_dimensionality, 260)
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 2547dce4b..ac1ed80d0 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -285,6 +285,17 @@ def basic_params1():
       tpu_enable_host_call=False,
       # Pad batch dim of inputs to nearest multiple of batch multiple.
       pad_batch=False,
+      # When true, do not evaluate on the language model data when running the
+      # multiproblem since it can take a while. If False, set eval_steps to
+      # something large like 6000 or 10000.
+      multiproblem_target_eval_only=False,
+      # Max out the vocab size to a power of 2 for efficiency and to reserve
+      # extra space in the vocabulary for new task ids and label classes.
+      multiproblem_vocab_size=-1,
+      # When using multiproblem with generation tasks, need to truncate the
+      # inputs and targets manually before concatenating them.
+      multiproblem_max_input_length=-1,
+      multiproblem_max_target_length=-1
   )
 
 

From 4c207787c7222b0986f1e12898799f3f5c4da782 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Tue, 6 Nov 2018 14:50:25 -0800
Subject: [PATCH 1189/2720] Add docstring to ones_matrix_band_part.

PiperOrigin-RevId: 220357088
---
 tensor2tensor/layers/common_layers.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 3216aab59..ddc6fb132 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2822,7 +2822,20 @@ def sample_with_temperature(logits, temperature):
 
 
 def ones_matrix_band_part(rows, cols, num_lower, num_upper, out_shape=None):
-  """Matrix band part of ones."""
+  """Matrix band part of ones.
+
+  Args:
+    rows: int determining number of rows in output
+    cols: int
+    num_lower: int, maximum distance backward. Negative values indicate
+      unlimited.
+    num_upper: int, maximum distance forward. Negative values indicate
+      unlimited.
+    out_shape: shape to reshape output by.
+
+  Returns:
+    Tensor of size rows * cols reshaped into shape out_shape.
+  """
   if all([isinstance(el, int) for el in [rows, cols, num_lower, num_upper]]):
     # Needed info is constant, so we construct in numpy
     if num_lower < 0:

From 4c421a5b25e30fa8bd39c023fd3c40acc1491424 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 6 Nov 2018 15:26:51 -0800
Subject: [PATCH 1190/2720] Change dependency on tf.contrib.distributions to
 tfp.distributions

PiperOrigin-RevId: 220363686
---
 setup.py                                 |  1 +
 tensor2tensor/layers/common_attention.py |  7 ++++---
 tensor2tensor/layers/discretization.py   | 11 ++++++-----
 tensor2tensor/layers/latent_layers.py    |  5 +++--
 tensor2tensor/layers/modalities.py       |  4 ++--
 tensor2tensor/models/research/rl.py      | 12 ++++++------
 6 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/setup.py b/setup.py
index 666c75630..09899a393 100644
--- a/setup.py
+++ b/setup.py
@@ -47,6 +47,7 @@
         'scipy',
         'sympy',
         'six',
+        'tensorflow-probability',
         'tqdm',
     ],
     extras_require={
diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 1f3bf04e6..48dd51377 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -32,6 +32,7 @@
 from tensor2tensor.utils import expert_utils
 
 import tensorflow as tf
+import tensorflow_probability as tfp
 
 from tensorflow.python.framework import function
 from tensorflow.python.ops import inplace_ops
@@ -327,9 +328,9 @@ def combine_attentions(attention_list):
     return tf.reduce_mean(attentions, [0, 2])
 
   def kl_divergence_loss(expected_logits, actual_logits):
-    p = tf.contrib.distributions.Categorical(logits=expected_logits)
-    q = tf.contrib.distributions.Categorical(logits=actual_logits)
-    return tf.contrib.distributions.kl_divergence(p, q)
+    p = tfp.distributions.Categorical(logits=expected_logits)
+    q = tfp.distributions.Categorical(logits=actual_logits)
+    return tfp.distributions.kl_divergence(p, q)
 
   def mse_loss(expected_logits, actual_weights):
     expected_weights = tf.nn.softmax(expected_logits)
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 165b01a28..8478cf92e 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -25,6 +25,7 @@
 from tensor2tensor.layers import common_layers
 
 import tensorflow as tf
+import tensorflow_probability as tfp
 
 from tensorflow.python.training import moving_averages
 
@@ -1147,9 +1148,9 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
   q_samples = tf.clip_by_value(gumbel_softmax_samples, 1e-6, 1 - 1e-6)
 
   if approximate_gs_entropy:
-    q_dist = tf.contrib.distributions.Multinomial(total_count=1.0, logits=-dist)
+    q_dist = tfp.distributions.Multinomial(total_count=1.0, logits=-dist)
   else:
-    q_dist = tf.contrib.distributions.RelaxedOneHotCategorical(
+    q_dist = tfp.distributions.RelaxedOneHotCategorical(
         temperature, logits=-dist)
 
   # Take mean over samples to approximate entropy.
@@ -1195,9 +1196,9 @@ def gumbel_softmax_nearest_neighbor_dvq(x,
       # which we can do without recalculating probabilities because the last
       # dimension of log_pi and q_samples are deterministic given the others.
       # Flow 2: Centered-softmax.
-      chained_bijectors = tf.contrib.distributions.bijectors.Chain([
-          tf.contrib.distributions.bijectors.SoftmaxCentered(),
-          tf.contrib.distributions.bijectors.Affine(
+      chained_bijectors = tfp.bijectors.Chain([
+          tfp.bijectors.SoftmaxCentered(),
+          tfp.bijectors.Affine(
               shift=log_pi[:, :, :-1],
               scale_identity_multiplier=1. / temperature)
       ])
diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 421a2c946..9a10eaeb0 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -27,6 +27,7 @@
 from tensor2tensor.utils import beam_search
 
 import tensorflow as tf
+import tensorflow_probability as tfp
 
 DO_SUMMARIES = True
 
@@ -731,8 +732,8 @@ def iaf_flow(one_hot_assignments,
     # shifting the rest down by one (and removing the last dimension).
     padded_assignments = tf.pad(
         one_hot_assignments, [[0, 0], [0, 0], [1, 0], [0, 0]])[:, :, :-1, :]
-    scale_bijector = tf.contrib.distributions.bijectors.Affine(
-        scale_tril=tf.contrib.distributions.fill_triangular(scale_weights))
+    scale_bijector = tfp.distributions.bijectors.Affine(
+        scale_tril=tfp.distributions.fill_triangular(scale_weights))
     scale = scale_bijector.forward(
         tf.transpose(padded_assignments, [0, 1, 3, 2]))
     # Transpose the bijector output since it performs a batch matmul.
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index c7bb80e05..3c622c886 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -27,6 +27,7 @@
 from tensor2tensor.utils import modality
 
 import tensorflow as tf
+import tensorflow_probability as tfp
 
 
 class SymbolModality(modality.Modality):
@@ -681,8 +682,7 @@ class VideoModalityPixelNoise(VideoModality):
   def bottom(self, x):
     inputs = x
     if self._model_hparams.mode == tf.estimator.ModeKeys.TRAIN:
-      background = tf.contrib.distributions.percentile(inputs, 50.,
-                                                       axis=[0, 1, 2, 3])
+      background = tfp.distributions.percentile(inputs, 50., axis=[0, 1, 2, 3])
       input_shape = common_layers.shape_list(inputs)
       input_size = tf.reduce_prod(input_shape[:-1])
       input_mask = tf.multinomial(
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 9caa3763f..4d61c3e7e 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -30,6 +30,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
+import tensorflow_probability as tfp
 
 
 @registry.register_hparams
@@ -245,8 +246,7 @@ def feed_forward_gaussian_fun(action_space, config, observations):
   logstd = tf.check_numerics(logstd, "logstd")
   value = tf.check_numerics(value, "value")
 
-  policy = tf.contrib.distributions.MultivariateNormalDiag(mean,
-                                                           tf.exp(logstd))
+  policy = tfp.distributions.MultivariateNormalDiag(mean, tf.exp(logstd))
 
   return NetworkOutput(policy, value, lambda a: tf.clip_by_value(a, -2., 2))
 
@@ -280,7 +280,7 @@ def feed_forward_categorical_fun(action_space, config, observations):
         x = tf.contrib.layers.fully_connected(x, size, tf.nn.relu)
       value = tf.contrib.layers.fully_connected(x, 1, None)[..., 0]
   logits = clip_logits(logits, config)
-  policy = tf.contrib.distributions.Categorical(logits=logits)
+  policy = tfp.distributions.Categorical(logits=logits)
   return NetworkOutput(policy, value, lambda a: a)
 
 
@@ -309,7 +309,7 @@ def feed_forward_cnn_small_categorical_fun(action_space, config, observations):
 
       value = tf.contrib.layers.fully_connected(
           x, 1, activation_fn=None)[..., 0]
-      policy = tf.contrib.distributions.Categorical(logits=logits)
+      policy = tfp.distributions.Categorical(logits=logits)
   return NetworkOutput(policy, value, lambda a: a)
 
 
@@ -345,7 +345,7 @@ def feed_forward_cnn_small_categorical_fun_new(
       logits = clip_logits(logits, config)
 
       value = tf.layers.dense(x, 1, name="value")[..., 0]
-      policy = tf.contrib.distributions.Categorical(logits=logits)
+      policy = tfp.distributions.Categorical(logits=logits)
 
   return NetworkOutput(policy, value, lambda a: a)
 
@@ -371,7 +371,7 @@ def dense_bitwise_categorical_fun(action_space, config, observations):
 
       value = tf.contrib.layers.fully_connected(
           x, 1, activation_fn=None)[..., 0]
-      policy = tf.contrib.distributions.Categorical(logits=logits)
+      policy = tfp.distributions.Categorical(logits=logits)
 
   return NetworkOutput(policy, value, lambda a: a)
 

From a91c2bbaeee2fa26de8a4f91a855941fefa5f8db Mon Sep 17 00:00:00 2001
From: Dustin Tran <dustinviettran@gmail.com>
Date: Tue, 6 Nov 2018 16:38:02 -0800
Subject: [PATCH 1191/2720] Return ed.RandomVariables in initializers and
 operate on them in regularizers. (#1182)

---
 tensor2tensor/layers/bayes.py      | 35 ++++++++++++++++++++----------
 tensor2tensor/layers/bayes_test.py | 31 +++++++++-----------------
 2 files changed, 34 insertions(+), 32 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 76e6ae5f8..1b60cd7b1 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -19,8 +19,11 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
 import tensorflow as tf
 
+from tensorflow_probability import edward2 as ed
+
 
 class Softplus(tf.keras.constraints.Constraint):
   """Softplus constraint."""
@@ -110,11 +113,7 @@ def __call__(self, shape=None, dtype=None, partition_info=None):
       raise ValueError('A TrainableInitializer must be built by a layer before '
                        'usage, and is currently only compatible with Bayesian '
                        'layers.')
-    noise = tf.random_normal(self.shape, dtype=self.dtype, seed=self.seed)
-    output = self.mean + self.stddev * noise
-    # TODO(trandustin): Hack to store parameters so KL reg. can operate on them.
-    output._parameters = (self.mean, self.stddev)  # pylint: disable=protected-access
-    return output
+    return ed.Normal(loc=self.mean, scale=self.stddev)
 
   def get_config(self):
     return {
@@ -151,12 +150,11 @@ def __init__(self, mean=0., stddev=1.):
     self.stddev = stddev
 
   def __call__(self, x):
-    mean, stddev = x._parameters  # pylint: disable=protected-access
-    variance2 = tf.square(self.stddev)
-    variance_ratio = tf.square(stddev) / variance2
-    regularization = tf.square(mean - self.mean) / (2. * variance2)
-    regularization += (variance_ratio - 1. - tf.log(variance_ratio)) / 2.
-    return regularization
+    """Computes regularization given an ed.Normal random variable as input."""
+    if not isinstance(x, ed.RandomVariable):
+      raise ValueError('Input must be an ed.RandomVariable.')
+    random_variable = ed.Normal(loc=self.mean, scale=self.stddev)
+    return random_variable.distribution.kl_divergence(x.distribution)
 
   def get_config(self):
     return {
@@ -276,3 +274,18 @@ def build(self, input_shape):
     else:
       self.bias = None
     self.built = True
+
+  # TODO(trandustin): Waiting on T2T to drop dependence on
+  # TF<=1.12rc2. A TF commit enables tf.colocate_with to work for
+  # Tensor-like inputs. This lets us use the parent method instead of
+  # this one.
+  def _handle_weight_regularization(self, name, variable, regularizer):
+    """Create lambdas which compute regularization losses."""
+
+    def _loss_for_variable(v):
+      """Creates a regularization loss `Tensor` for variable `v`."""
+      with tf.name_scope(name + '/Regularizer'):
+        regularization = regularizer(v)
+      return regularization
+
+    self.add_loss(functools.partial(_loss_for_variable, variable))
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index c9e754438..521d7de06 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -29,25 +29,20 @@
 
 class BayesTest(parameterized.TestCase, tf.test.TestCase):
 
-  # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
-  # support for TF 1.10
-  # @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDenseReparameterizationKernel(self):
     inputs = tf.to_float(np.random.rand(5, 3, 12))
     layer = bayes.DenseReparameterization(4, activation=tf.nn.relu)
     outputs1 = layer(inputs)
     outputs2 = layer(inputs)
     self.evaluate(tf.global_variables_initializer())
-    # res1, res2 = self.evaluate([outputs1, outputs2])
-    res1, _ = self.evaluate([outputs1, outputs2])
+    res1, res2 = self.evaluate([outputs1, outputs2])
     self.assertEqual(res1.shape, (5, 3, 4))
     self.assertAllGreaterEqual(res1, 0.)
-    # self.assertNotAllClose(res1, res2)
+    self.assertNotAllClose(res1, res2)
     layer.get_config()
 
-  # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
-  # support for TF 1.10
-  # @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDenseReparameterizationBias(self):
     inputs = tf.to_float(np.random.rand(5, 3, 12))
     layer = bayes.DenseReparameterization(4, kernel_initializer="zero",
@@ -56,15 +51,12 @@ def testDenseReparameterizationBias(self):
     outputs1 = layer(inputs)
     outputs2 = layer(inputs)
     self.evaluate(tf.global_variables_initializer())
-    # res1, res2 = self.evaluate([outputs1, outputs2])
-    res1, _ = self.evaluate([outputs1, outputs2])
+    res1, res2 = self.evaluate([outputs1, outputs2])
     self.assertEqual(res1.shape, (5, 3, 4))
     self.assertAllGreaterEqual(res1, 0.)
-    # self.assertNotAllClose(res1, res2)
+    self.assertNotAllClose(res1, res2)
 
-  # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
-  # support for TF 1.10
-  # @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDenseReparameterizationDeterministic(self):
     inputs = tf.to_float(np.random.rand(5, 3, 12))
     layer = bayes.DenseReparameterization(4, kernel_initializer="zero",
@@ -73,15 +65,12 @@ def testDenseReparameterizationDeterministic(self):
     outputs1 = layer(inputs)
     outputs2 = layer(inputs)
     self.evaluate(tf.global_variables_initializer())
-    # res1, res2 = self.evaluate([outputs1, outputs2])
-    res1, _ = self.evaluate([outputs1, outputs2])
+    res1, res2 = self.evaluate([outputs1, outputs2])
     self.assertEqual(res1.shape, (5, 3, 4))
     self.assertAllGreaterEqual(res1, 0.)
-    # self.assertAllClose(res1, res2)
+    self.assertAllClose(res1, res2)
 
-  # TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
-  # support for TF 1.10
-  # @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDenseReparameterizationModel(self):
     inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
     model = tf.keras.Sequential([

From a8769f1586f08076c466125b4dff0062964b1a8e Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 6 Nov 2018 17:10:06 -0800
Subject: [PATCH 1192/2720] Avoid building init op during eval / decoding. This
 would save some computation time during eval / decode time since
 init_batch_size is usually very large.

PiperOrigin-RevId: 220380543
---
 tensor2tensor/models/research/glow.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index c21bb2997..e44aba705 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -143,12 +143,13 @@ def top_prior(self):
         "top_prior", self.z_top_shape, learn_prior=self.hparams.top_prior)
 
   def body(self, features):
-    init_features = self.create_init_batch(features)
-    init_op = self.objective_tower(init_features, init=True)
-    init_op = tf.Print(
-        init_op, [init_op], message="Triggering data-dependent init.",
-        first_n=20)
-    tf.add_to_collection("glow_init_op", init_op)
+    if self.is_training:
+      init_features = self.create_init_batch(features)
+      init_op = self.objective_tower(init_features, init=True)
+      init_op = tf.Print(
+          init_op, [init_op], message="Triggering data-dependent init.",
+          first_n=20)
+      tf.add_to_collection("glow_init_op", init_op)
     train_op = self.objective_tower(features, init=False)
     return tf.zeros_like(features["targets"]), {"training": train_op}
 

From 599c3454e7a60678c3e92e67b7c821234f1fcd0e Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 6 Nov 2018 18:41:05 -0800
Subject: [PATCH 1193/2720] Add tfds-nightly dependency

PiperOrigin-RevId: 220390737
---
 setup.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/setup.py b/setup.py
index 09899a393..aeb7a65cb 100644
--- a/setup.py
+++ b/setup.py
@@ -48,6 +48,7 @@
         'sympy',
         'six',
         'tensorflow-probability',
+        'tfds-nightly',
         'tqdm',
     ],
     extras_require={

From 81e8ec09aac7d7d230eb1cfe409bcc6fc3ef1ebd Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Tue, 6 Nov 2018 18:48:02 -0800
Subject: [PATCH 1194/2720] fix metric name usage.

PiperOrigin-RevId: 220391233
---
 tensor2tensor/rl/trainer_model_based.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 92148c811..6c298c02f 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -328,6 +328,11 @@ def evaluate_single_config(hparams, agent_model_dir):
   )
 
 
+def get_metric_name(stochastic, max_num_noops, clipped):
+  return "mean_reward/eval/stochastic_{}_max_noops_{}_{}".format(
+      stochastic, max_num_noops, "clipped" if clipped else "unclipped")
+
+
 def evaluate_all_configs(hparams, agent_model_dir):
   """Evaluate the agent with multiple eval configurations."""
   def make_eval_hparams(hparams, stochastic, max_num_noops):
@@ -345,10 +350,7 @@ def make_eval_hparams(hparams, stochastic, max_num_noops):
       eval_hparams = make_eval_hparams(hparams, stochastic, max_num_noops)
       scores = evaluate_single_config(eval_hparams, agent_model_dir)
       for (score, clipped) in zip(scores, (True, False)):
-        metric_name = "mean_reward/eval/stochastic_{}_max_noops_{}_{}".format(
-            stochastic, max_num_noops,
-            "clipped" if clipped else "unclipped"
-        )
+        metric_name = get_metric_name(stochastic, max_num_noops, clipped)
         metrics[metric_name] = score
 
   return metrics
@@ -591,8 +593,10 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     epoch_metrics.append(metrics)
     if report_fn:
       if report_metric == "mean_reward":
-        report_fn(eval_metrics["mean_reward/eval/{}_{}_max_noops_{}".format(
-            "mode", hparams.eval_max_num_noops, "unclipped")], epoch)
+        metric_name = get_metric_name(stochastic=False,
+                                      max_num_noops=hparams.eval_max_num_noops,
+                                      clipped=False)
+        report_fn(eval_metrics[metric_name], epoch)
       else:
         report_fn(eval_metrics[report_metric], epoch)
 

From e317e98f7e2f5f0c4347d3ac96ad2e9d0ef60090 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 6 Nov 2018 19:56:28 -0800
Subject: [PATCH 1195/2720] Change kl_loss variable naming convention to
 clarify computation procedure,

PiperOrigin-RevId: 220396243
---
 tensor2tensor/layers/common_layers.py  | 19 +++++++++++++------
 tensor2tensor/models/video/base_vae.py | 16 ++++++++++++----
 tensor2tensor/utils/t2t_model.py       |  9 +++++++++
 3 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index ddc6fb132..e2246d40c 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3717,19 +3717,26 @@ def targeted_dropout(inputs,
     return inputs
 
 
-# TODO(mbz): use tf.distributions.kl_divergence instead.
-def kl_divergence(mu, log_sigma):
-  """KL divergence of diagonal gaussian N(mu,exp(log_sigma)) and N(0,1).
+def kl_divergence(mu, log_var, mu_p=0.0, log_var_p=0.0):
+  """KL divergence of diagonal gaussian N(mu,exp(log_var)) and N(0,1).
 
   Args:
     mu: mu parameter of the distribution.
-    log_sigma: log(sigma) parameter of the distribution.
+    log_var: log(var) parameter of the distribution.
+    mu_p: optional mu from a learned prior distribution
+    log_var_p: optional log(var) from a learned prior distribution
   Returns:
     the KL loss.
   """
+
   batch_size = shape_list(mu)[0]
-  kl = -.5 * tf.reduce_sum(1. + log_sigma - tf.square(mu) - tf.exp(log_sigma))
-  return kl / tf.to_float(batch_size)
+  prior_distribution = tf.distributions.Normal(
+      mu_p, tf.exp(tf.multiply(0.5, log_var_p)))
+  posterior_distribution = tf.distributions.Normal(
+      mu, tf.exp(tf.multiply(0.5, log_var)))
+  kld = tf.distributions.kl_divergence(
+      posterior_distribution, prior_distribution)
+  return tf.reduce_sum(kld) / tf.to_float(batch_size)
 
 
 def sparse_equals_constant(constant, tensor):
diff --git a/tensor2tensor/models/video/base_vae.py b/tensor2tensor/models/video/base_vae.py
index 73f989529..b049b6513 100644
--- a/tensor2tensor/models/video/base_vae.py
+++ b/tensor2tensor/models/video/base_vae.py
@@ -71,14 +71,22 @@ def get_beta(self, kl_loss=0.0):
       tf.summary.scalar("beta", beta)
       return beta
 
-  def get_kl_loss(self, means, stds):
+  def get_kl_loss(self, means, log_vars, means_p=None, log_vars_p=None):
     """Get KL loss for all the predicted Gaussians."""
     kl_loss = 0.0
+    if means_p is None:
+      means_p = tf.unstack(tf.zeros_like(means))
+    if log_vars_p is None:
+      log_vars_p = tf.unstack(tf.zeros_like(log_vars))
+    enumerated_inputs = enumerate(zip(means, log_vars, means_p, log_vars_p))
     if self.is_training and self.hparams.stochastic_model:
-      for i, (mean, std) in enumerate(zip(means, stds)):
-        kl_loss += common_layers.kl_divergence(mean, std)
+      for i, (mean, log_var, mean_p, log_var_p) in enumerated_inputs:
+        # Condition to compute kl divergence with learned prior
+        kl_loss += common_layers.kl_divergence(mean, log_var, mean_p, log_var_p)
         tf.summary.histogram("posterior_mean_%d" % i, mean)
-        tf.summary.histogram("posterior_std_%d" % i, std)
+        tf.summary.histogram("posterior_log_var_%d" % i, log_var)
+        tf.summary.histogram("prior_mean_%d" % i, mean_p)
+        tf.summary.histogram("prior_log_var_%d" % i, log_var_p)
       tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
 
     beta = self.get_beta(kl_loss)
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index edd8592ca..a279f2234 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1495,10 +1495,19 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
         predictions = logits
       else:
         predictions = {"predictions": logits}
+
+      evaluation_hooks = []
+      # Create a SummarySaverHook
+      eval_summary_hook = tf.train.SummarySaverHook(
+          save_steps=1, output_dir=self.hparams.model_dir + "/eval",
+          summary_op=tf.summary.merge_all())
+      evaluation_hooks.append(eval_summary_hook)
+
       return tf.estimator.EstimatorSpec(
           tf.estimator.ModeKeys.EVAL,
           predictions=predictions,
           eval_metric_ops=eval_metrics,
+          evaluation_hooks=evaluation_hooks,
           loss=loss)
 
   def estimator_spec_predict(self, features, use_tpu=False):

From 2e0f5264aa571b1e4a4a4d1db6e8389b4938aa93 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 6 Nov 2018 20:22:40 -0800
Subject: [PATCH 1196/2720] Fixed bug in svg with fixed prior implementation
 where variance was treated as standard deviation during sampling but treated
 as variance in kl term

PiperOrigin-RevId: 220398634
---
 tensor2tensor/models/video/emily.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index 06340df9b..3c41d276e 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -265,9 +265,8 @@ def construct_model(self, images, actions, rewards):
           mu, logvar, posterior_states = self.lstm_gaussian(
               h_target, posterior_states, rnn_size, z_dim, posterior_rnn_layers)
 
-          # The original implementation has a multiplier of 0.5
-          # Removed here for simplicity i.e. replacing var with std
-          z = z * tf.exp(logvar) + mu
+          # Sample z from posterior distribution
+          z = z * tf.exp(tf.multiply(0.5, logvar)) + mu
 
         # Predict output encoding
         h_pred, predictor_states = self.stacked_lstm(

From a17a488db2c8078e56dcca5971372ea6d237ff65 Mon Sep 17 00:00:00 2001
From: Daniel De Freitas Adiwardana <adiwardana@google.com>
Date: Wed, 7 Nov 2018 11:03:19 -0800
Subject: [PATCH 1197/2720] Reshaping scalar scores for decode.

PiperOrigin-RevId: 220493888
---
 tensor2tensor/utils/decoding.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 6a8e626bc..c5520fd5a 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -397,6 +397,8 @@ def timer(gen):
       output_beams = np.split(result["outputs"], decode_hp.beam_size, axis=0)
       scores = None
       if "scores" in result:
+        if np.isscalar(result["scores"]):
+          result["scores"] = result["scores"].reshape(1)
         scores = np.split(result["scores"], decode_hp.beam_size, axis=0)
       for k, beam in enumerate(output_beams):
         tf.logging.info("BEAM %d:" % k)
@@ -552,6 +554,8 @@ def input_fn():
       beams = np.split(result["outputs"], decode_hp.beam_size, axis=0)
       scores = None
       if "scores" in result:
+        if np.isscalar(result["scores"]):
+          result["scores"] = result["scores"].reshape(1)
         scores = np.split(result["scores"], decode_hp.beam_size, axis=0)
       for k, beam in enumerate(beams):
         tf.logging.info("BEAM %d:" % k)

From 229a1d4dd2448a47986e1e23433bf177c100801a Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 7 Nov 2018 13:14:03 -0800
Subject: [PATCH 1198/2720] Internal change

PiperOrigin-RevId: 220516613
---
 tensor2tensor/bin/t2t_trainer.py   |  1 +
 tensor2tensor/utils/flags.py       |  3 +++
 tensor2tensor/utils/trainer_lib.py | 23 +++++++++++++++++++----
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index a221a6b52..83cd93c44 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -188,6 +188,7 @@ def create_experiment_fn():
       eval_early_stopping_metric_delta=FLAGS.eval_early_stopping_metric_delta,
       eval_early_stopping_metric_minimize=FLAGS
       .eval_early_stopping_metric_minimize,
+      eval_timeout_mins=FLAGS.eval_timeout_mins,
       use_tpu=FLAGS.use_tpu,
       use_tpu_estimator=FLAGS.use_tpu_estimator,
       use_xla=FLAGS.xla_compile,
diff --git a/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index 94169363a..05d51ec12 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -65,6 +65,9 @@
 flags.DEFINE_bool("eval_early_stopping_metric_minimize", True,
                   "Whether to check for the early stopping metric going down "
                   "or up.")
+flags.DEFINE_integer("eval_timeout_mins", 240,
+                     "The maximum amount of time to wait to wait between "
+                     "checkpoints. Set -1 to wait indefinitely.")
 flags.DEFINE_bool("eval_run_autoregressive", False,
                   "Run eval autoregressively where we condition on previous"
                   "generated output instead of the actual target.")
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 482a3115d..87a7897a4 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -40,11 +40,22 @@
 
 
 def next_checkpoint(model_dir, timeout_mins=240):
-  """Yields successive checkpoints from model_dir."""
+  """Yields successive checkpoints from model_dir.
+
+  Args:
+    model_dir: The directory in which checkpoints are saved.
+    timeout_mins: The maximum amount of time in minutes to wait
+                  between checkpoints. Set this to -1 to wait indefinitely.
+  Yields:
+    last_ckpt: a new checkpoint path, or None if the timeout was reached.
+  """
   last_ckpt = None
+  timeout_secs = None
+  if timeout_mins != -1:
+    timeout_secs = timeout_mins * 60
   while True:
     last_ckpt = tf.contrib.training.wait_for_new_checkpoint(
-        model_dir, last_ckpt, seconds_to_sleep=60, timeout=60 * timeout_mins)
+        model_dir, last_ckpt, seconds_to_sleep=60, timeout=timeout_secs)
 
     if last_ckpt is None:
       tf.logging.info(
@@ -499,12 +510,14 @@ def evaluate_on_train_data(self):
 
   def continuous_eval(self):
     """Evaluate until checkpoints stop being produced."""
-    for _ in next_checkpoint(self._hparams.model_dir):
+    for _ in next_checkpoint(self._hparams.model_dir,
+                             self._hparams.eval_timeout_mins):
       self.evaluate()
 
   def continuous_eval_on_train_data(self):
     """Evaluate on train data until checkpoints stop being produced."""
-    for _ in next_checkpoint(self._hparams.model_dir):
+    for _ in next_checkpoint(self._hparams.model_dir,
+                             self._hparams.eval_timeout_mins):
       self.evaluate_on_train_data()
 
   def test(self):
@@ -624,6 +637,7 @@ def create_experiment(
     eval_early_stopping_metric=None,
     eval_early_stopping_metric_delta=None,
     eval_early_stopping_metric_minimize=True,
+    eval_timeout_mins=240,
     use_tpu=False,
     use_tpu_estimator=False,
     use_xla=False,
@@ -644,6 +658,7 @@ def create_experiment(
   hparams.add_hparam("warm_start_from", warm_start_from)
   hparams.add_hparam("std_server_protocol", std_server_protocol)
   hparams.add_hparam("eval_freq_in_steps", min_eval_frequency)
+  hparams.add_hparam("eval_timeout_mins", eval_timeout_mins)
   if decode_hparams is not None:
     decode_hparams.add_hparam("decode_from_file", decode_from_file)
     decode_hparams.add_hparam("decode_to_file", decode_to_file)

From 4ffcda7135050a4a813eaf3ca25dc8e06a047c6a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 7 Nov 2018 13:25:19 -0800
Subject: [PATCH 1199/2720] TPU settings for multi-problem.

PiperOrigin-RevId: 220518409
---
 .../data_generators/multi_problem.py          | 15 ++++++++-----
 tensor2tensor/layers/modalities.py            | 10 ++-------
 tensor2tensor/models/transformer.py           | 22 +++++++++++++++++++
 3 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 5b3043cef..53d6363f0 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -206,12 +206,15 @@ def dataset(self,
     self.get_hparams(model_hparams=hparams)
 
     if is_training:
-      problem_step = tf.get_variable("problem_step",
-                                     shape=[],
-                                     dtype=tf.int64,
-                                     initializer=tf.zeros_initializer(),
-                                     trainable=False,
-                                     use_resource=True)
+      # Using tf.Variable instead of get_variable to work around issues with
+      # queues on multiple hosts. Note that this will separately count steps
+      # on each host that's feeding the data, so in a large-scale setting you
+      # may need to adjust hparams for that. For example, a 4x4 slice of a TPU
+      # pod may use 2 data hosts, so we'll be only adding 1 here once for 2
+      # examples -- divide the corresponding hparams by 2 to compensate.
+      problem_step = tf.Variable(tf.constant(0, dtype=tf.int64),
+                                 trainable=False, use_resource=True,
+                                 dtype=tf.int64, name="problem_step")
       dataset_iterators = [d.make_one_shot_iterator() for d in datasets]
 
       def get_next_from_dataset(dataset_iter):
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 3c622c886..79a920bf2 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -160,14 +160,8 @@ def top(self, body_output, _):
       else:
         body_output = tf.reshape(body_output, [-1, body_output_shape[-1]])
         logits = tf.matmul(body_output, var, transpose_b=True)
-        if (common_layers.is_xla_compiled() and
-            self._model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
-          # TPU does not react kindly to extra dimensions.
-          # TODO(noam): remove this once TPU is more forgiving of extra dims.
-          return logits
-        else:
-          return tf.reshape(logits,
-                            body_output_shape[:-1] + [1, self._vocab_size])
+        return tf.reshape(logits,
+                          body_output_shape[:-1] + [1, self._vocab_size])
 
 
 class SymbolModalityWeightsAll(SymbolModality):
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 910a05f3c..d00de8dca 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1569,6 +1569,28 @@ def transformer_tall_pretrain_lm():
   return hparams
 
 
+@registry.register_hparams
+def transformer_tall_pretrain_lm_tpu_adafactor():
+  """Hparams for transformer on LM pretraining (with 64k vocab) on TPU."""
+  hparams = transformer_tall_pretrain_lm()
+  update_hparams_for_tpu(hparams)
+  hparams.max_length = 1024
+  # For multi-problem on TPU we need it in absolute examples.
+  hparams.batch_size = 8
+  return hparams
+
+
+@registry.register_hparams
+def transformer_tall_pretrain_lm_tpu():
+  """Hparams for transformer on LM pretraining on TPU with AdamW."""
+  hparams = transformer_tall_pretrain_lm_tpu_adafactor()
+  # Optimizer gets reset in update_hparams_for_tpu so we set it again here.
+  hparams.learning_rate_constant = 2e-4
+  hparams.learning_rate_schedule = ("linear_warmup * constant * cosdecay")
+  hparams.optimizer = "AdamW"
+  return hparams
+
+
 @registry.register_hparams
 def transformer_tall_finetune_cnndm():
   """Hparams for transformer on LM for finetuning on cnndm summarization."""

From 8cc31071bdcfb867ea442c89b351347a4850ff11 Mon Sep 17 00:00:00 2001
From: Daniel De Freitas Adiwardana <adiwardana@google.com>
Date: Wed, 7 Nov 2018 13:37:02 -0800
Subject: [PATCH 1200/2720] Allow create_hparams() in tensor2tensor export to
 load from hparams.json (if any).

PiperOrigin-RevId: 220520381
---
 tensor2tensor/serving/export.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 54125cb70..8a8f85a63 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -44,6 +44,19 @@
     "specified by --output_dir")
 
 
+def _get_hparams_path():
+  """Get hyper-parameters file path."""
+  hparams_path = None
+  if FLAGS.output_dir:
+    hparams_path = os.path.join(FLAGS.output_dir, "hparams.json")
+  else:
+    tf.logging.warning(
+        "--output_dir not specified. Hyper-parameters will be infered from"
+        "--hparams_set and --hparams only. These may not match training time"
+        "hyper-parameters.")
+  return hparams_path
+
+
 def create_estimator(run_config, hparams):
   return trainer_lib.create_estimator(
       FLAGS.model,
@@ -54,11 +67,13 @@ def create_estimator(run_config, hparams):
 
 
 def create_hparams():
+  """Create hyper-parameters object."""
   return trainer_lib.create_hparams(
       FLAGS.hparams_set,
       FLAGS.hparams,
       data_dir=os.path.expanduser(FLAGS.data_dir),
-      problem_name=FLAGS.problem)
+      problem_name=FLAGS.problem,
+      hparams_path=_get_hparams_path())
 
 
 # TODO(michalski): Move this method into tfhub utils.

From dd56db99cb3310d7f085e3ce7106ef78f05fbffc Mon Sep 17 00:00:00 2001
From: Mike Dusenberry <dusenberrymw@google.com>
Date: Wed, 7 Nov 2018 15:07:17 -0800
Subject: [PATCH 1201/2720] Adding a Bayesian LSTM Cell Layer.

In a Bayesian LSTM cell, uncertainty can be expressed over the kernel,
recurrent kernel, and/or bias.

PiperOrigin-RevId: 220536640
---
 tensor2tensor/layers/bayes.py      | 186 +++++++++++++++++++++++++++--
 tensor2tensor/layers/bayes_test.py | 122 ++++++++++++++-----
 2 files changed, 267 insertions(+), 41 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 1b60cd7b1..9c10ddf44 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -42,6 +42,8 @@ def softplus():  # alias, following tf.keras.constraints
   return Softplus()
 
 
+# TODO(dusenberrymw): Restructure the implementation of a trainable initializer
+# such that callers do not need to have type-conditional logic.
 class TrainableInitializer(tf.keras.initializers.Initializer):
   """An initializer with trainable variables.
 
@@ -113,6 +115,8 @@ def __call__(self, shape=None, dtype=None, partition_info=None):
       raise ValueError('A TrainableInitializer must be built by a layer before '
                        'usage, and is currently only compatible with Bayesian '
                        'layers.')
+    # TODO(dusenberrymw): The softplus constraint seems to not be applied, so
+    # the following ends up being `mean + unconstrained_stddev * noise`.
     return ed.Normal(loc=self.mean, scale=self.stddev)
 
   def get_config(self):
@@ -239,6 +243,10 @@ def build(self, input_shape):
       self.kernel_initializer.build([last_dim, self.units],
                                     self.dtype,
                                     self.add_weight)
+      if self.kernel_regularizer is not None:
+        self._handle_weight_regularization(
+            'kernel', self.kernel, self.kernel_regularizer)
+
     else:
       self._kernel = self.add_weight(
           'kernel',
@@ -249,14 +257,12 @@ def build(self, input_shape):
           dtype=self.dtype,
           trainable=True)
 
-    if self.kernel_regularizer is not None:
-      self._handle_weight_regularization('kernel',
-                                         self.kernel,
-                                         self.kernel_regularizer)
-
     if self.use_bias:
       if isinstance(self.bias_initializer, TrainableInitializer):
         self.bias_initializer.build([self.units], self.dtype, self.add_weight)
+        if self.bias_regularizer is not None:
+          self._handle_weight_regularization(
+              'bias', self.bias, self.bias_regularizer)
       else:
         self._bias = self.add_weight(
             'bias',
@@ -267,14 +273,176 @@ def build(self, input_shape):
             dtype=self.dtype,
             trainable=True)
 
-      if self.bias_regularizer is not None:
-        self._handle_weight_regularization('bias',
-                                           self.bias,
-                                           self.bias_regularizer)
+    else:
+      self._bias = None
+    self.built = True
+
+  # TODO(trandustin): Waiting on T2T to drop dependence on
+  # TF<=1.12rc2. A TF commit enables tf.colocate_with to work for
+  # Tensor-like inputs. This lets us use the parent method instead of
+  # this one.
+  def _handle_weight_regularization(self, name, variable, regularizer):
+    """Create lambdas which compute regularization losses."""
+
+    def _loss_for_variable(v):
+      """Creates a regularization loss `Tensor` for variable `v`."""
+      with tf.name_scope(name + '/Regularizer'):
+        regularization = regularizer(v)
+      return regularization
+
+    self.add_loss(functools.partial(_loss_for_variable, variable))
+
+
+class LSTMCellReparameterization(tf.keras.layers.LSTMCell):
+  """Bayesian LSTM cell class estimated via reparameterization.
+
+  The layer computes a variational Bayesian approximation to the distribution
+  over LSTM cell functions,
+
+  ```
+  p(outputs | inputs) = int lstm_cell(inputs; weights, bias) p(weights, bias)
+    dweights dbias,
+  ```
+
+  where the weights consist of both input and recurrent weights.
+
+  It does this with a stochastic forward pass, sampling from learnable
+  distributions on the kernel, recurrent kernel, and bias. Gradients with
+  respect to the distributions' learnable parameters backpropagate via
+  reparameterization.  Minimizing cross-entropy plus the layer's losses performs
+  variational minimum description length, i.e., it minimizes an upper bound to
+  the negative marginal likelihood.
+  """
+
+  def __init__(self,
+               units,
+               activation='tanh',
+               recurrent_activation='hard_sigmoid',
+               use_bias=True,
+               kernel_initializer=None,
+               recurrent_initializer=None,
+               bias_initializer='zeros',
+               unit_forget_bias=True,
+               kernel_regularizer=normal_kl_divergence(),
+               recurrent_regularizer=normal_kl_divergence(),
+               bias_regularizer=None,
+               kernel_constraint=None,
+               recurrent_constraint=None,
+               bias_constraint=None,
+               dropout=0.,
+               recurrent_dropout=0.,
+               implementation=1,
+               **kwargs):
+    if not kernel_initializer:
+      kernel_initializer = trainable_normal()
+    if not recurrent_initializer:
+      recurrent_initializer = trainable_normal()
+    if not bias_initializer:
+      bias_initializer = trainable_normal()
+    super(LSTMCellReparameterization, self).__init__(
+        units=units,
+        activation=activation,
+        recurrent_activation=recurrent_activation,
+        use_bias=use_bias,
+        kernel_initializer=kernel_initializer,
+        recurrent_initializer=recurrent_initializer,
+        bias_initializer=bias_initializer,
+        unit_forget_bias=unit_forget_bias,
+        kernel_regularizer=kernel_regularizer,
+        recurrent_regularizer=recurrent_regularizer,
+        bias_regularizer=bias_regularizer,
+        kernel_constraint=kernel_constraint,
+        recurrent_constraint=recurrent_constraint,
+        bias_constraint=bias_constraint,
+        dropout=dropout,
+        recurrent_dropout=recurrent_dropout,
+        implementation=implementation,
+        **kwargs)
+
+  def build(self, input_shape):
+    input_shape = tf.TensorShape(input_shape)
+    input_dim = input_shape[-1]
+    if isinstance(input_dim, tf.Dimension):
+      input_dim = input_dim.value
+
+    if isinstance(self.kernel_initializer, TrainableInitializer):
+      self.kernel_initializer.build(
+          [input_dim, self.units * 4], self.dtype, self.add_weight)
+      self.kernel = self.kernel_initializer()
+      if self.kernel_regularizer is not None:
+        self._handle_weight_regularization(
+            'kernel', self.kernel, self.kernel_regularizer)
+
+    else:
+      self.kernel = self.add_weight(
+          shape=(input_dim, self.units * 4),
+          name='kernel',
+          initializer=self.kernel_initializer,
+          regularizer=self.kernel_regularizer,
+          constraint=self.kernel_constraint)
+
+    if isinstance(self.recurrent_initializer, TrainableInitializer):
+      self.recurrent_initializer.build(
+          [self.units, self.units * 4], self.dtype, self.add_weight)
+      self.recurrent_kernel = self.recurrent_initializer()
+      if self.recurrent_regularizer is not None:
+        self._handle_weight_regularization(
+            'recurrent_kernel', self.recurrent_kernel,
+            self.recurrent_regularizer)
+
+    else:
+      self.recurrent_kernel = self.add_weight(
+          shape=(self.units, self.units * 4),
+          name='recurrent_kernel',
+          initializer=self.recurrent_initializer,
+          regularizer=self.recurrent_regularizer,
+          constraint=self.recurrent_constraint)
+
+    if self.use_bias:
+      if isinstance(self.bias_initializer, TrainableInitializer):
+        self.bias_initializer.build(
+            [self.units * 4], self.dtype, self.add_weight)
+        self.bias = self.bias_initializer()
+        if self.bias_regularizer is not None:
+          self._handle_weight_regularization(
+              'bias', self.bias, self.bias_regularizer)
+      else:
+        if self.unit_forget_bias:
+
+          def bias_initializer(_, *args, **kwargs):
+            return tf.keras.backend.concatenate([
+                self.bias_initializer((self.units,), *args, **kwargs),
+                tf.keras.initializers.Ones()((self.units,), *args, **kwargs),
+                self.bias_initializer((self.units * 2,), *args, **kwargs),
+            ])
+        else:
+          bias_initializer = self.bias_initializer
+        self.bias = self.add_weight(
+            shape=(self.units * 4,),
+            name='bias',
+            initializer=bias_initializer,
+            regularizer=self.bias_regularizer,
+            constraint=self.bias_constraint)
     else:
       self.bias = None
     self.built = True
 
+  def sample_weights(self):
+    if isinstance(self.kernel_initializer, TrainableInitializer):
+      self.kernel = self.kernel_initializer()
+    if isinstance(self.recurrent_initializer, TrainableInitializer):
+      self.recurrent_kernel = self.recurrent_initializer()
+    if isinstance(self.bias_initializer, TrainableInitializer):
+      self.bias = self.bias_initializer()
+
+  # NOTE: This will not be called in TF < 1.11.
+  def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
+    """Get the initial state and side-effect sampling of stochastic weights."""
+    if self.built:
+      self.sample_weights()
+    return super(LSTMCellReparameterization, self).get_initial_state(
+        inputs=inputs, batch_size=batch_size, dtype=dtype)
+
   # TODO(trandustin): Waiting on T2T to drop dependence on
   # TF<=1.12rc2. A TF commit enables tf.colocate_with to work for
   # Tensor-like inputs. This lets us use the parent method instead of
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 521d7de06..905222c85 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -29,47 +29,33 @@
 
 class BayesTest(parameterized.TestCase, tf.test.TestCase):
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
-  def testDenseReparameterizationKernel(self):
+  @parameterized.named_parameters(
+      {"testcase_name": "_no_uncertainty", "kernel_initializer": "zeros",
+       "bias_initializer": "zeros", "all_close": True},
+      {"testcase_name": "_kernel_uncertainty", "kernel_initializer": None,
+       "bias_initializer": "zeros", "all_close": False},
+      {"testcase_name": "_bias_uncertainty", "kernel_initializer": "zeros",
+       "bias_initializer": None, "all_close": False},
+  )
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes
+  def testDenseReparameterizationKernel(
+      self, kernel_initializer, bias_initializer, all_close):
     inputs = tf.to_float(np.random.rand(5, 3, 12))
-    layer = bayes.DenseReparameterization(4, activation=tf.nn.relu)
+    layer = bayes.DenseReparameterization(
+        4, kernel_initializer=kernel_initializer,
+        bias_initializer=bias_initializer, activation=tf.nn.relu)
     outputs1 = layer(inputs)
     outputs2 = layer(inputs)
     self.evaluate(tf.global_variables_initializer())
     res1, res2 = self.evaluate([outputs1, outputs2])
     self.assertEqual(res1.shape, (5, 3, 4))
     self.assertAllGreaterEqual(res1, 0.)
-    self.assertNotAllClose(res1, res2)
+    if all_close:
+      self.assertAllClose(res1, res2)
+    else:
+      self.assertNotAllClose(res1, res2)
     layer.get_config()
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
-  def testDenseReparameterizationBias(self):
-    inputs = tf.to_float(np.random.rand(5, 3, 12))
-    layer = bayes.DenseReparameterization(4, kernel_initializer="zero",
-                                          bias_initializer=None,
-                                          activation=tf.nn.relu)
-    outputs1 = layer(inputs)
-    outputs2 = layer(inputs)
-    self.evaluate(tf.global_variables_initializer())
-    res1, res2 = self.evaluate([outputs1, outputs2])
-    self.assertEqual(res1.shape, (5, 3, 4))
-    self.assertAllGreaterEqual(res1, 0.)
-    self.assertNotAllClose(res1, res2)
-
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
-  def testDenseReparameterizationDeterministic(self):
-    inputs = tf.to_float(np.random.rand(5, 3, 12))
-    layer = bayes.DenseReparameterization(4, kernel_initializer="zero",
-                                          bias_initializer="zero",
-                                          activation=tf.nn.relu)
-    outputs1 = layer(inputs)
-    outputs2 = layer(inputs)
-    self.evaluate(tf.global_variables_initializer())
-    res1, res2 = self.evaluate([outputs1, outputs2])
-    self.assertEqual(res1.shape, (5, 3, 4))
-    self.assertAllGreaterEqual(res1, 0.)
-    self.assertAllClose(res1, res2)
-
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDenseReparameterizationModel(self):
     inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
@@ -87,6 +73,78 @@ def testDenseReparameterizationModel(self):
     self.assertEqual(res.shape, (3, 2))
     self.assertLen(model.losses, 1)
 
+  @parameterized.named_parameters(
+      {"testcase_name": "_no_uncertainty", "kernel_initializer": "zeros",
+       "recurrent_initializer": "orthogonal", "bias_initializer": "zeros",
+       "all_close": True},
+      {"testcase_name": "_kernel_uncertainty", "kernel_initializer": None,
+       "recurrent_initializer": "orthogonal", "bias_initializer": "zeros",
+       "all_close": False},
+      {"testcase_name": "_recurrent_uncertainty", "kernel_initializer": "zeros",
+       "recurrent_initializer": None, "bias_initializer": "zeros",
+       "all_close": False},
+      {"testcase_name": "_bias_uncertainty", "kernel_initializer": "zeros",
+       "recurrent_initializer": "orthogonal", "bias_initializer": None,
+       "all_close": False},
+  )
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes
+  def testLSTMCellReparameterization(
+      self, kernel_initializer, recurrent_initializer, bias_initializer,
+      all_close):
+    batch_size, timesteps, dim = 5, 3, 12
+    hidden_size = 10
+    inputs = tf.to_float(np.random.rand(batch_size, timesteps, dim))
+    cell = bayes.LSTMCellReparameterization(
+        hidden_size, kernel_initializer=kernel_initializer,
+        recurrent_initializer=recurrent_initializer,
+        bias_initializer=bias_initializer)
+    noise = tf.to_float(np.random.rand(1, hidden_size))
+    h0, c0 = cell.get_initial_state(inputs)
+    state = (h0 + noise, c0)
+    outputs1, _ = cell(inputs[:, 0, :], state)
+    outputs2, _ = cell(inputs[:, 0, :], state)
+    cell.sample_weights()
+    outputs3, _ = cell(inputs[:, 0, :], state)
+    self.evaluate(tf.global_variables_initializer())
+    res1, res2, res3 = self.evaluate([outputs1, outputs2, outputs3])
+    self.assertEqual(res1.shape, (batch_size, hidden_size))
+    self.assertAllClose(res1, res2)
+    if all_close:
+      self.assertAllClose(res1, res3)
+    else:
+      self.assertNotAllClose(res1, res3)
+    cell.get_config()
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testLSTMCellReparameterizationModel(self):
+    batch_size, timesteps, dim = 5, 3, 12
+    hidden_size = 10
+    inputs = tf.to_float(np.random.rand(batch_size, timesteps, dim))
+    cell = bayes.LSTMCellReparameterization(hidden_size)
+    model = tf.keras.Sequential([
+        tf.keras.layers.RNN(cell, return_sequences=True)
+    ])
+    outputs1 = model(inputs)
+    outputs2 = model(inputs)
+    state = (tf.zeros([1, hidden_size]), tf.zeros([1, hidden_size]))
+    outputs3 = []
+    for t in range(timesteps):
+      out, state = cell(inputs[:, t, :], state)
+      outputs3.append(out)
+    outputs3 = tf.stack(outputs3, axis=1)
+    self.evaluate(tf.global_variables_initializer())
+    res1, res2, res3 = self.evaluate([outputs1, outputs2, outputs3])
+    self.assertEqual(res1.shape, (batch_size, timesteps, hidden_size))
+    self.assertEqual(res3.shape, (batch_size, timesteps, hidden_size))
+    # NOTE: `cell.sample_weights` should have been called at the beginning of
+    # each call, so these should be different.
+    self.assertNotAllClose(res1, res2)
+    # NOTE: We didn't call `cell.sample_weights` again before computing
+    # `outputs3`, so the cell should have had the same weights as it did during
+    # computation of `outputs2`, and thus yielded the same output tensor.
+    self.assertAllClose(res2, res3)
+    self.assertLen(model.losses, 2)
+
 
 if __name__ == "__main__":
   tf.test.main()

From 8bcbdccf85c0fc60f07945c469ff3213d2e0810d Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 7 Nov 2018 18:44:20 -0800
Subject: [PATCH 1202/2720] Add some new models/configurations for
 mesh-tensorflow language models.  Fixed functionality of
 hparams.split_to_length.

PiperOrigin-RevId: 220566451
---
 tensor2tensor/data_generators/problem.py      | 12 ++--
 tensor2tensor/data_generators/wiki.py         | 13 ++++
 .../models/research/moe_experiments.py        | 71 ++++++++++++++++++-
 3 files changed, 89 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 0526ff972..6b68d540b 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -153,11 +153,13 @@ def preprocess_example_common(example, hparams, mode):
   if hparams.max_target_seq_length > 0:
     example["targets"] = example["targets"][:hparams.max_target_seq_length]
   if hparams.split_to_length:
-    example["targets"] = tf.reshape(example["targets"],
-                                    [-1, hparams.split_to_length, 1, 1])
-    if len(example) != 1:
-      raise ValueError("split_to_length only works for LM problems")
-    return tf.data.Dataset.from_tensor_slices(example)
+    new_example = {}
+    for k, v in six.iteritems(example):
+      if k == "targets" or k == "inputs":
+        new_example[k] = tf.reshape(v, [-1, hparams.split_to_length, 1, 1])
+      else:
+        tf.logging.warning("Dropping feature %s" % k)
+    return tf.data.Dataset.from_tensor_slices(new_example)
   return example
 
 
diff --git a/tensor2tensor/data_generators/wiki.py b/tensor2tensor/data_generators/wiki.py
index 31973fdeb..ee3dcbc17 100644
--- a/tensor2tensor/data_generators/wiki.py
+++ b/tensor2tensor/data_generators/wiki.py
@@ -395,6 +395,19 @@ def max_chars_for_vocab(self):
     return 100 * (10 ** 6)
 
 
+@registry.register_problem
+class LanguagemodelWikiNorefV32kL16k(LanguagemodelWikiNorefV32kL1k):
+  """A language model on English Wikipedia.
+
+  References removed.  Chopped into segments of 16k tokens.
+  """
+
+  @property
+  def sequence_length(self):
+    """Length of each example (in tokens)."""
+    return 2**14
+
+
 @registry.register_problem
 class LanguagemodelWikiNorefV128kL1k(LanguagemodelWikiNorefV8kL1k):
   """128k vocab."""
diff --git a/tensor2tensor/models/research/moe_experiments.py b/tensor2tensor/models/research/moe_experiments.py
index 445abc419..f53467099 100644
--- a/tensor2tensor/models/research/moe_experiments.py
+++ b/tensor2tensor/models/research/moe_experiments.py
@@ -166,8 +166,8 @@ def xmoe2_dense(sz):
 
   TODO(noam): find a large enough dataset for these experiments.
 
-  You can use languagemodel_wiki_noref_v8k_l1k, but this is too small,
-  so training will cover about 9 epochs.
+  You can use languagemodel_wiki_noref_v32k_l1k, but this is too small,
+  (1 epoch = ~46000 steps) so training will cover about 11 epochs.
 
   Note: configurations and code are likely to change without notice.
 
@@ -263,3 +263,70 @@ def xmoe2_tiny():
   hparams.mesh_shape = ""
   hparams.activation_dtype = "float32"
   return hparams
+
+
+@registry.register_hparams
+def xmoe2_v1_l4k():
+  """With sequence length 4096."""
+  hparams = xmoe2_v1()
+  hparams.batch_size = 32
+  hparams.max_length = 4096
+  hparams.split_to_length = 4096
+  return hparams
+
+
+@registry.register_hparams
+def xmoe2_v1_l4k_local_only():
+  """With sequence length 4096."""
+  hparams = xmoe2_v1_l4k()
+  hparams.decoder_layers = [
+      "local_att" if l == "att" else l for l in hparams.decoder_layers]
+  return hparams
+
+
+@registry.register_hparams
+def wiki_2x2_base():
+  """Set of architectural experiments - language model on wikipedia on a 2x2.
+
+  1 epoch = ~180k steps at batch size 32 - we may never finish an epoch!
+
+  Returns:
+    a hparams
+  """
+  hparams = mtf_transformer.mtf_transformer_base_lm()
+  hparams.shared_embedding_and_softmax_weights = False
+  # no dropout - dataset is big enough to avoid overfitting.
+  hparams.attention_dropout = 0.0
+  hparams.relu_dropout = 0.0
+  hparams.layer_prepostprocess_dropout = 0.0
+  hparams.max_length = 1024
+  # 4 sequences per core
+  hparams.batch_size = 32
+  # We don't use linear decay in these experiments, since we don't want
+  # a sharp jump in quality at the end of the training schedule.
+  # You can insert this once you find the right architecture.
+  hparams.learning_rate_schedule = "rsqrt_decay"
+  hparams.mesh_shape = "all:8"
+  hparams.layout = "batch:all;experts:all"
+
+  # parameters for mixture-of-experts
+  moe.set_default_moe_hparams(hparams)
+  hparams.moe_num_experts = 16
+  hparams.moe_hidden_size = 8192
+
+  hparams.decoder_layers = ["att", "drd"] * 6
+  hparams.d_model = 1024
+  hparams.d_ff = 2048
+  hparams.d_kv = 128
+  hparams.num_heads = 4
+
+  return hparams
+
+
+@registry.register_hparams
+def wiki_2x2_v1():
+  hparams = wiki_2x2_base()
+  hparams.decoder_layers = (
+      ["local_att", "local_att", "drd",
+       "att", "drd", "local_att", "local_att", "moe"] * 4)[:-1]
+  return hparams

From e44850a42a7919580a33b2505c06928998d43137 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 8 Nov 2018 10:52:15 -0800
Subject: [PATCH 1203/2720] Only test against latest stable TF version (and
 tf-nightly)

PiperOrigin-RevId: 220664269
---
 .travis.yml | 12 +++++-------
 setup.py    |  4 ++--
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index fd91c80fa..62d037c47 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -13,20 +13,18 @@ env:
     - T2T_PROBLEM=algorithmic_reverse_binary40_test
     - T2T_DATA_DIR=/tmp/t2t-data
     - T2T_TRAIN_DIR=/tmp/t2t-train
-    - TF_LATEST="1.11.*"
+    - TF_LATEST="1.12.*"
     # This is necessary to have gsutil work with Python 2.7
     - BOTO_CONFIG=/dev/null
   matrix:
-    # We test against recent versions of TensorFlow and tf-nightly.
+    # We test against the latest stable TensorFlow and tf-nightly.
     # If updating, also update TF_LATEST above
-    - TF_VERSION="1.10.*"
-    - TF_VERSION="1.11.*"
+    - TF_VERSION="1.12.*"
     - TF_VERSION="tf-nightly"
 matrix:
   exclude:
-    # We test against all versions in Python 2 but only the latest in Python 3
-    - python: "3.6"
-      env: TF_VERSION="1.10.*"
+    # We test against all versions in Python 2 but only the latest stable
+    # version in Python 3
     - python: "3.6"
       env: TF_VERSION="tf-nightly"
 before_install:
diff --git a/setup.py b/setup.py
index aeb7a65cb..d0e586ee3 100644
--- a/setup.py
+++ b/setup.py
@@ -52,8 +52,8 @@
         'tqdm',
     ],
     extras_require={
-        'tensorflow': ['tensorflow>=1.9.0'],
-        'tensorflow_gpu': ['tensorflow-gpu>=1.9.0'],
+        'tensorflow': ['tensorflow>=1.12.0'],
+        'tensorflow_gpu': ['tensorflow-gpu>=1.12.0'],
         'tensorflow-hub': ['tensorflow-hub>=0.1.1'],
         'tests': [
             'absl-py',

From 6376e049b5220c3d1de274f51121539615724971 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Thu, 8 Nov 2018 11:00:36 -0800
Subject: [PATCH 1204/2720] First pass to use 3-D convolutions to model the
 transitions in latent space.

PiperOrigin-RevId: 220665834
---
 tensor2tensor/models/research/glow_ops.py     | 204 ++++++++++++++----
 .../models/research/glow_ops_test.py          | 120 +++++++++--
 2 files changed, 263 insertions(+), 61 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 0a2ea3904..ebe841674 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -54,7 +54,7 @@ def assign(w, initial_value):
 def get_cond_latents_at_level(cond_latents, level, hparams):
   """Returns a single or list of conditional latents at level 'level'."""
   if cond_latents:
-    if hparams.latent_dist_encoder == "conv_net":
+    if hparams.latent_dist_encoder in ["conv_net", "conv3d_net"]:
       return [cond_latent[level] for cond_latent in cond_latents]
     elif hparams.latent_dist_encoder in ["pointwise", "conv_lstm"]:
       return cond_latents[level]
@@ -92,6 +92,30 @@ def get_variable_ddi(name, shape, initial_value, dtype=tf.float32, init=False,
     return tf.cond(init, lambda: assign(w, initial_value), lambda: w)
 
 
+@add_arg_scope
+def actnorm_3d(name, x, logscale_factor=3.):
+  """Applies actnorm to each time-step independently.
+
+  There are a total of 2*n_channels*n_steps parameters learnt.
+
+  Args:
+    name: variable scope.
+    x: 5-D Tensor, (NTHWC)
+    logscale_factor: Increases the learning rate of the scale by
+                     logscale_factor.
+  Returns:
+    x: 5-D Tensor, (NTHWC) with the per-timestep, per-channel normalization.
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    x = tf.unstack(x, axis=1)
+    x_normed = []
+    for ind, x_step in enumerate(x):
+      x_step, _ = actnorm("actnorm_%d" % ind, x_step,
+                          logscale_factor=logscale_factor)
+      x_normed.append(x_step)
+    return tf.stack(x_normed, axis=1), None
+
+
 @add_arg_scope
 def actnorm(name, x, logscale_factor=3., reverse=False, init=False,
             trainable=True):
@@ -292,17 +316,39 @@ def add_edge_bias(x, filter_size):
   return tf.concat([x, x_pad], axis=3)
 
 
+def time_pad(x, filter_size):
+  """Pad left across time and pad valid across the spatial components.
+
+  Args:
+    x: 5-D Tensor, (NTHWC)
+    filter_size: list of ints
+  Returns:
+    x_pad: 5-D Tensor.
+  """
+  if filter_size == [1, 1, 1]:
+    return x
+  a = (filter_size[1] - 1) // 2  # vertical padding size
+  b = (filter_size[2] - 1) // 2  # horizontal padding size
+  c = filter_size[0] - 1
+  padding = [[0, 0], [c, 0], [a, a], [b, b], [0, 0]]
+  return tf.pad(x, padding)
+
+
 @add_arg_scope
-def conv2d(name, x, output_channels, filter_size=None, stride=None,
-           logscale_factor=3.0, apply_actnorm=True, conv_init="default"):
-  """conv2d layer with edge bias padding and optional actnorm.
+def conv(name, x, output_channels, filter_size=None, stride=None,
+         logscale_factor=3.0, apply_actnorm=True, conv_init="default"):
+  """Convolutional layer with edge bias padding and optional actnorm.
+
+  If x is 5-dimensional, actnorm is applied independently across every
+  time-step.
 
   Args:
     name: variable scope.
-    x: 4-D Tensor of shape (NHWC)
+    x: 4-D Tensor or 5-D Tensor of shape NHWC or NTHWC
     output_channels: Number of output channels.
-    filter_size:
-    stride:
+    filter_size: list of ints, if None [3, 3] and [2, 3, 3] are defaults for
+                 4-D and 5-D input tensors respectively.
+    stride: list of ints, default stride: 1
     logscale_factor: see actnorm for parameter meaning.
     apply_actnorm: if apply_actnorm the activations of the first minibatch
                    have zero mean and unit variance. Else, there is no scaling
@@ -316,14 +362,28 @@ def conv2d(name, x, output_channels, filter_size=None, stride=None,
   if conv_init == "zeros" and apply_actnorm:
     raise ValueError("apply_actnorm is unstable when init is set to zeros.")
 
-  if filter_size is None:
-    filter_size = [3, 3]
-  if stride is None:
-    stride = [1, 1]
-
-  x = add_edge_bias(x, filter_size=filter_size)
-  _, _, _, in_channels = common_layers.shape_list(x)
-
+  x_shape = common_layers.shape_list(x)
+  is_2d = len(x_shape) == 4
+
+  # set filter_size, stride and in_channels
+  if is_2d:
+    if filter_size is None:
+      filter_size = [3, 3]
+    if stride is None:
+      stride = [1, 1]
+    actnorm_func = actnorm
+    x = add_edge_bias(x, filter_size=filter_size)
+    conv_filter = tf.nn.conv2d
+  else:
+    if filter_size is None:
+      filter_size = [2, 3, 3]
+    if stride is None:
+      stride = [1, 1, 1]
+    actnorm_func = actnorm_3d
+    x = time_pad(x, filter_size=filter_size)
+    conv_filter = tf.nn.conv3d
+
+  in_channels = common_layers.shape_list(x)[-1]
   filter_shape = filter_size + [in_channels, output_channels]
   stride_shape = [1] + stride + [1]
 
@@ -334,13 +394,10 @@ def conv2d(name, x, output_channels, filter_size=None, stride=None,
     elif conv_init == "zeros":
       initializer = tf.zeros_initializer()
 
-    w = tf.get_variable("W", filter_shape, tf.float32,
-                        initializer=initializer)
-    x = tf.nn.conv2d(x, w, stride_shape, padding="VALID", data_format="NHWC")
-
+    w = tf.get_variable("W", filter_shape, tf.float32, initializer=initializer)
+    x = conv_filter(x, w, stride_shape, padding="VALID")
     if apply_actnorm:
-      x, _ = actnorm("actnorm", x, logscale_factor=logscale_factor,
-                     trainable=True)
+      x, _ = actnorm_func("actnorm", x, logscale_factor=logscale_factor)
     else:
       x += tf.get_variable("b", [1, 1, 1, output_channels],
                            initializer=tf.zeros_initializer())
@@ -351,28 +408,36 @@ def conv2d(name, x, output_channels, filter_size=None, stride=None,
 
 
 @add_arg_scope
-def conv_block(name, x, mid_channels):
+def conv_block(name, x, mid_channels, time_filter=2):
   """2 layer conv block used in the affine coupling layer.
 
   Args:
     name: variable scope.
-    x: 4-D Tensor: (batch_size, height, width, channels).
+    x: 4-D or 5-D Tensor.
     mid_channels: Output channels of the second layer.
+    time_filter: Filter across time to capture context.
   Returns:
     x: 4-D Tensor: Output activations.
   """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
 
+    x_shape = common_layers.shape_list(x)
+    is_2d = len(x_shape) == 4
+    if is_2d:
+      first_filter = [3, 3]
+      second_filter = [1, 1]
+    else:
+      first_filter = [time_filter, 3, 3]
+      second_filter = [1, 1, 1]
+
     # Edge Padding + conv2d + actnorm + relu:
     # [output: 512 channels]
-    x = conv2d("1_1", x, output_channels=mid_channels, filter_size=[3, 3],
-               stride=[1, 1])
+    x = conv("1_1", x, output_channels=mid_channels, filter_size=first_filter)
     x = tf.nn.relu(x)
 
     # Padding + conv2d + actnorm + relu
     # [input, output: 512 channels]
-    x = conv2d("1_2", x, output_channels=mid_channels, filter_size=[1, 1],
-               stride=[1, 1])
+    x = conv("1_2", x, output_channels=mid_channels, filter_size=second_filter)
     x = tf.nn.relu(x)
     return x
 
@@ -395,9 +460,9 @@ def affine_coupling_network(name, x, mid_channels, output_channels):
     x = conv_block("conv_block", x, mid_channels=mid_channels)
 
     # Final layer.
-    x = conv2d("zeros", x, filter_size=[3, 3], stride=[1, 1],
-               output_channels=output_channels, apply_actnorm=False,
-               conv_init="zeros")
+    x = conv("zeros", x, filter_size=[3, 3], stride=[1, 1],
+             output_channels=output_channels, apply_actnorm=False,
+             conv_init="zeros")
   return x
 
 
@@ -477,6 +542,34 @@ def squeeze(name, x, factor=2, reverse=True):
     return x
 
 
+@add_arg_scope
+def temporal_tensor_to_dist(name, x, hparams, output_channels=None):
+  """Network that maps a time-indexed list of 3-D Tensors to a gaussian.
+
+  Args:
+    name: variable scope.
+    x: List of 4-D Tensors indexed by time, (NHWC)
+    hparams: tf.contrib.training.Hparams.
+    output_channels: int, Number of channels of the output gaussian mean.
+  Returns:
+    dist: tf.distributions.Normal
+  """
+  if output_channels is None:
+    output_channels = common_layers.shape_list(x)[-1]
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    h = conv_block("conv3d_init", x, time_filter=2,
+                   mid_channels=hparams.latent_encoder_width)
+    h = conv("conv3d_zeros", h, apply_actnorm=False,
+             output_channels=2*output_channels, conv_init="zeros",
+             filter_size=[2, 3, 3])
+
+    # take last activation that should capture all context since padding is
+    # on left.
+    h = h[:, -1, :, :, :]
+    mean, log_scale = h[:, :, :, 0::2], h[:, :, :, 1::2]
+  return tf.distributions.Normal(mean, tf.exp(log_scale))
+
+
 @add_arg_scope
 def tensor_to_dist(name, x, output_channels=None, architecture="single_conv",
                    depth=1, pre_output_channels=512, width=512):
@@ -501,28 +594,28 @@ def tensor_to_dist(name, x, output_channels=None, architecture="single_conv",
     if output_channels is None:
       output_channels = x_shape[-1]
     if architecture == "single_conv":
-      mean_log_scale = conv2d("conv2d", x, output_channels=2*output_channels,
-                              conv_init="zeros", apply_actnorm=False)
+      mean_log_scale = conv("conv2d", x, output_channels=2*output_channels,
+                            conv_init="zeros", apply_actnorm=False)
     elif architecture == "glow_nn":
       mean_log_scale = x
       for layer in range(1, depth + 1):
         mid_channels = pre_output_channels // 2**(depth - layer)
         mean_log_scale = conv_block("glow_nn_%d" % layer, mean_log_scale,
                                     mid_channels=mid_channels)
-      mean_log_scale = conv2d("glow_nn_zeros", mean_log_scale,
-                              filter_size=[3, 3], stride=[1, 1],
-                              output_channels=2*output_channels,
-                              apply_actnorm=False, conv_init="zeros")
+      mean_log_scale = conv("glow_nn_zeros", mean_log_scale,
+                            filter_size=[3, 3], stride=[1, 1],
+                            output_channels=2*output_channels,
+                            apply_actnorm=False, conv_init="zeros")
     elif architecture == "glow_resnet":
       h = x
       for layer in range(depth):
         h2 = conv_block("glow_res_%d" % layer, h, mid_channels=width)
-        h3 = conv2d("glow_res_zeros_%d" % layer, h2, conv_init="zeros",
-                    output_channels=x_shape[-1], apply_actnorm=False)
+        h3 = conv("glow_res_zeros_%d" % layer, h2, conv_init="zeros",
+                  output_channels=x_shape[-1], apply_actnorm=False)
         h += h3
-      mean_log_scale = conv2d("glow_res_final", h, conv_init="zeros",
-                              output_channels=2*output_channels,
-                              apply_actnorm=False)
+      mean_log_scale = conv("glow_res_final", h, conv_init="zeros",
+                            output_channels=2*output_channels,
+                            apply_actnorm=False)
     else:
       raise ValueError("expected architecture to be single_conv or glow_nn "
                        "got %s" % architecture)
@@ -591,8 +684,10 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
         "latent_prior", latent, logscale_factor=3.0)
     cond_dist = merge_level_and_latent_dist(prior_dist, latent_dist,
                                             merge_std=merge_std)
+
   elif latent_dist_encoder == "conv_net":
     output_channels = common_layers.shape_list(z)[-1]
+    last_latent = latent[-1]
     latent_stack = tf.concat([prior_dist.loc] + latent, axis=-1)
     cond_dist = tensor_to_dist(
         "latent_stack", latent_stack, output_channels=output_channels,
@@ -600,10 +695,24 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
         depth=hparams.latent_encoder_depth,
         pre_output_channels=hparams.latent_pre_output_channels,
         width=hparams.latent_encoder_width)
-    if latent_skip:
-      cond_dist = tf.distributions.Normal(
-          cond_dist.loc + latent[-1], cond_dist.scale)
+
+  elif latent_dist_encoder == "conv3d_net":
+    last_latent = latent[-1]
+    output_channels = common_layers.shape_list(last_latent)[-1]
+    num_steps = len(latent)
+
+    # Stack across time.
+    cond_latents = tf.stack(latent, axis=1)
+
+    # Concat latents from previous levels across channels.
+    prev_latents = tf.tile(tf.expand_dims(prior_dist.loc, axis=1),
+                           [1, num_steps, 1, 1, 1])
+    cond_latents = tf.concat((cond_latents, prev_latents), axis=-1)
+    cond_dist = temporal_tensor_to_dist(
+        "latent_stack", cond_latents, hparams, output_channels=output_channels)
+
   elif latent_dist_encoder == "conv_lstm":
+    last_latent = latent
     output_channels = common_layers.shape_list(z)[-1]
     latent_stack = tf.concat((prior_dist.loc, latent), axis=-1)
     _, state = common_video.conv_lstm_2d(
@@ -611,9 +720,9 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
         name="conv_lstm")
     cond_dist = tensor_to_dist(
         "state_to_dist", state.h, output_channels=output_channels)
-    if latent_skip:
-      cond_dist = tf.distributions.Normal(
-          cond_dist.loc + latent, cond_dist.scale)
+  if latent_skip:
+    new_mean = cond_dist.loc + last_latent
+    cond_dist = tf.distributions.Normal(new_mean, cond_dist.scale)
   return cond_dist.loc, cond_dist.scale, state
 
 
@@ -646,6 +755,7 @@ def compute_prior(name, z, latent, hparams, condition=False, state=None):
       condition = tf.constant(condition, dtype=tf.bool)
     prior_dist = tensor_to_dist("level_prior", z, architecture="single_conv")
     prior_mean, prior_scale = prior_dist.loc, prior_dist.scale
+
     if latent is None:
       mean, scale = prior_mean, prior_scale
     else:
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 670da0d52..a19cc1272 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -21,6 +21,7 @@
 
 import os
 import tempfile
+from absl.testing import parameterized
 import numpy as np
 from tensor2tensor.models.research import glow
 from tensor2tensor.models.research import glow_ops
@@ -34,7 +35,21 @@
 add_arg_scope = tf.contrib.framework.add_arg_scope
 
 
-class GlowOpsTest(tf.test.TestCase):
+class GlowOpsTest(parameterized.TestCase, tf.test.TestCase):
+
+  def get_glow_hparams(self):
+    hparams = glow.glow_hparams()
+    hparams.add_hparam("num_cond_latents", 1)
+    hparams.add_hparam("latent_architecture", "glow_resnet")
+    # Use latent skip connections
+    hparams.add_hparam("model_input", False)
+    hparams.add_hparam("latent_skip", True)
+    hparams.add_hparam("latent_encoder_depth", 2)
+    hparams.add_hparam("latent_encoder_width", 256)
+    hparams.add_hparam("latent_pre_output_channels", 256)
+    hparams.add_hparam("latent_dist_encoder", "conv_net")
+    hparams.add_hparam("latent_time_filter_size", 3)
+    return hparams
 
   def test_get_variable_ddi(self):
     with tf.Graph().as_default():
@@ -90,9 +105,9 @@ def test_conv2d(self):
       x = 10.0 * tf.random_uniform(shape=(16, 5, 5, 32))
 
       with arg_scope([glow_ops.actnorm], init=True):
-        actnorm_conv2d = glow_ops.conv2d(
+        actnorm_conv2d = glow_ops.conv(
             "actnorm_conv2d", x, output_channels=64, apply_actnorm=True)
-        actnorm_zeros2d = glow_ops.conv2d(
+        actnorm_zeros2d = glow_ops.conv(
             "actnorm_zeros2d", x, output_channels=64, apply_actnorm=False)
 
       with tf.Session() as session:
@@ -297,7 +312,14 @@ def test_split_latent_conditioning(self):
     for merge_std in ["normal", "prev_level", "prev_step"]:
       self.check_split_latent_conditioning(merge_std)
 
-  def test_latent_dist_encoder_lstm(self):
+  @parameterized.named_parameters(
+      ("lstm_skip", "conv_lstm", True),
+      ("lstm_no_skip", "conv_lstm", False),
+      ("conv_net_skip", "conv_net", True),
+      ("conv_net_no_skip", "conv_net", False),
+      ("conv3d_skip", "conv3d_net", False),
+      ("conv3d_no_skip", "conv3d_net", True))
+  def test_latent_dist_encoder(self, encoder="conv_lstm", skip=True):
     with tf.Graph().as_default():
       rng = np.random.RandomState(0)
       # Initialize x, latent, state.
@@ -307,24 +329,94 @@ def test_latent_dist_encoder_lstm(self):
       x_t = tf.convert_to_tensor(x_rand)
       latent_t = tf.convert_to_tensor(latent_rand)
       state_t = tf.convert_to_tensor(state_rand)
+      if encoder in ["conv_net", "conv3d_net"]:
+        latent_t = [latent_t, latent_t]
       init_state = tf.contrib.rnn.LSTMStateTuple(state_t, state_t)
-      hparams = glow.glow_hparams()
-      hparams.add_hparam("latent_dist_encoder", "conv_lstm")
-      hparams.add_hparam("latent_skip", True)
-      hparams.add_hparam("latent_encoder_width", 256)
+      hparams = self.get_glow_hparams()
+      hparams.latent_dist_encoder = encoder
+      hparams.latent_skip = skip
+      hparams.latent_encoder_width = 256
 
       prior_dist, new_state = glow_ops.compute_prior(
-          "lstm_prior", x_t, latent=latent_t, hparams=hparams, state=init_state,
+          "prior", x_t, latent=latent_t, hparams=hparams, state=init_state,
           condition=True)
       with tf.Session() as sess:
         sess.run(tf.global_variables_initializer())
-        # Test initialization (mu, sigma) = (z, 1.0)
-        ops = [prior_dist.loc, prior_dist.scale, new_state.h - init_state.h]
-        mean, scale, diff_np = sess.run(ops)
-        self.assertTrue(np.allclose(latent_rand - mean, 0.0))
+        # Test initialization:
+        # Scale is 1.0
+        # If skip is set to True, then mean equals the input latent.
+        # If skip, is set to False, then the mean is zero.
+        ops = [prior_dist.loc, prior_dist.scale]
+        mean, scale = sess.run(ops)
+
+        if skip:
+          self.assertTrue(np.allclose(latent_rand - mean, 0.0))
+        else:
+          self.assertTrue(np.allclose(mean, 0.0))
         self.assertTrue(np.allclose(scale, 1.0))
+
         # State update.
-        self.assertFalse(np.allclose(diff_np, 0.0))
+        if encoder == "conv_lstm":
+          state_diff = sess.run(new_state.h - init_state.h)
+          self.assertFalse(np.allclose(state_diff, 0.0))
+
+  def test_conv3d(self):
+    with tf.Graph().as_default():
+      x = 10.0 * tf.random_uniform(shape=(16, 4, 5, 5, 32))
+
+      with arg_scope([glow_ops.actnorm], init=True):
+        conv3d = glow_ops.conv(
+            "conv3d", x, output_channels=64, apply_actnorm=True)
+        conv3d_zeros = glow_ops.conv(
+            "conv3d_zeros", x, output_channels=64, apply_actnorm=False,
+            conv_init="zeros")
+
+      with tf.Session() as session:
+        session.run(tf.global_variables_initializer())
+
+        # test if apply_actnorm is set to True, the first minibatch has
+        # zero mean and unit variance.
+        conv3d_np, conv3d_zeros_np = session.run([conv3d, conv3d_zeros])
+        self.assertEqual(conv3d_np.shape, (16, 4, 5, 5, 64))
+        for i in range(4):
+          curr_step = conv3d_np[:, i, :, :, :]
+          mean = np.mean(curr_step, axis=(0, 1, 2))
+          var = np.var(curr_step, axis=(0, 1, 2))
+          self.assertTrue(np.allclose(mean, 0.0, atol=1e-5))
+          self.assertTrue(np.allclose(var, 1.0, atol=1e-5))
+
+        # test shape in case apply_actnorm is set to False,
+        self.assertTrue(np.allclose(conv3d_zeros_np, 0.0))
+
+  def test_actnorm_3d(self):
+    with tf.Graph().as_default():
+      x_t = tf.random_normal((16, 5, 32, 32, 3), mean=50.0, stddev=2.0)
+      ops = [glow_ops.actnorm, glow_ops.get_variable_ddi]
+      with arg_scope(ops, init=True):
+        x_act, _ = glow_ops.actnorm_3d("actnorm", x_t)
+      with tf.Session() as session:
+        x_act_np = session.run(x_act)
+        # Mean and standard deviation per time-step equals zero and one.
+        for time_step in range(5):
+          x_act_curr = x_act_np[:, time_step, :, :, :]
+          channel_mean = np.mean(x_act_curr, axis=(0, 1, 2))
+          channel_var = np.var(x_act_curr, axis=(0, 1, 2))
+          self.assertTrue(np.allclose(channel_mean, 0.0, atol=1e-3))
+          self.assertTrue(np.allclose(channel_var, 1.0, atol=1e-3))
+
+  def test_temporal_tensor_to_dist(self):
+    with tf.Graph().as_default():
+      hparams = self.get_glow_hparams()
+      latent_shape = (16, 5, 4, 4, 48)
+      latents = tf.random_normal(latent_shape)
+      dist = glow_ops.temporal_tensor_to_dist(
+          "tensor_to_dist", latents, hparams)
+      with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+        mean, scale = dist.loc, dist.scale
+        mean_np, scale_np = sess.run([mean, scale])
+        self.assertTrue(np.allclose(mean_np, 0.0))
+        self.assertTrue(np.allclose(scale_np, 1.0))
 
 
 if __name__ == "__main__":

From 8edc5f19b7c0168638c1588a84cb3df689589f2d Mon Sep 17 00:00:00 2001
From: Mike Dusenberry <dusenberrymw@google.com>
Date: Thu, 8 Nov 2018 11:10:50 -0800
Subject: [PATCH 1205/2720] Initialize non-negative standard deviations for
 Bayesian Layer weights.

Variable constraints are only applied after the first optimization step.  This
can lead to unexpected weight values and potential NaN issues on the first
iteration prior to the first optimization step.  This commit changes the logic
to directly initialize a non-negative standard deviation, rather than an
untransformed representation that is subsequently constrained.

PiperOrigin-RevId: 220668146
---
 tensor2tensor/layers/bayes.py      | 39 +++++++++++++-----------------
 tensor2tensor/layers/bayes_test.py | 11 +++++++++
 2 files changed, 28 insertions(+), 22 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 9c10ddf44..821033084 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -64,22 +64,22 @@ class TrainableNormal(TrainableInitializer):
 
   def __init__(self,
                mean_initializer=tf.random_normal_initializer(stddev=0.1),
-               unconstrained_stddev_initializer=tf.random_normal_initializer(
-                   mean=-3., stddev=0.1),
+               stddev_initializer=tf.random_uniform_initializer(
+                   minval=1e-5, maxval=0.1),
                mean_regularizer=None,
-               unconstrained_stddev_regularizer=None,
+               stddev_regularizer=None,
                mean_constraint=None,
-               unconstrained_stddev_constraint=softplus(),
+               stddev_constraint=softplus(),
                seed=None,
                dtype=tf.float32):
     """Constructs the initializer."""
     super(TrainableNormal, self).__init__()
     self.mean_initializer = mean_initializer
-    self.unconstrained_stddev_initializer = unconstrained_stddev_initializer
+    self.stddev_initializer = stddev_initializer
     self.mean_regularizer = mean_regularizer
-    self.unconstrained_stddev_regularizer = unconstrained_stddev_regularizer
+    self.stddev_regularizer = stddev_regularizer
     self.mean_constraint = mean_constraint
-    self.unconstrained_stddev_constraint = unconstrained_stddev_constraint
+    self.stddev_constraint = stddev_constraint
     self.seed = seed
     self.dtype = tf.as_dtype(dtype)
 
@@ -99,11 +99,11 @@ def build(self, shape, dtype=None, add_variable_fn=None):
         dtype=dtype,
         trainable=True)
     self.stddev = add_variable_fn(
-        'unconstrained_stddev',
+        'stddev',
         shape=shape,
-        initializer=self.unconstrained_stddev_initializer,
-        regularizer=self.unconstrained_stddev_regularizer,
-        constraint=self.unconstrained_stddev_constraint,
+        initializer=self.stddev_initializer,
+        regularizer=self.stddev_regularizer,
+        constraint=self.stddev_constraint,
         dtype=dtype,
         trainable=True)
     self.built = True
@@ -115,27 +115,22 @@ def __call__(self, shape=None, dtype=None, partition_info=None):
       raise ValueError('A TrainableInitializer must be built by a layer before '
                        'usage, and is currently only compatible with Bayesian '
                        'layers.')
-    # TODO(dusenberrymw): The softplus constraint seems to not be applied, so
-    # the following ends up being `mean + unconstrained_stddev * noise`.
     return ed.Normal(loc=self.mean, scale=self.stddev)
 
   def get_config(self):
     return {
         'mean_initializer':
             tf.keras.initializers.serialize(self.mean_initializer),
-        'unconstrained_stddev_initializer':
-            tf.keras.initializers.serialize(
-                self.unconstrained_stddev_initializer),
+        'stddev_initializer':
+            tf.keras.initializers.serialize(self.stddev_initializer),
         'mean_regularizer':
             tf.keras.regularizers.serialize(self.mean_regularizer),
-        'unconstrained_stddev_regularizer':
-            tf.keras.regularizers.serialize(
-                self.unconstrained_stddev_regularizer),
+        'stddev_regularizer':
+            tf.keras.regularizers.serialize(self.stddev_regularizer),
         'mean_constraint':
             tf.keras.constraints.serialize(self.mean_constraint),
-        'unconstrained_stddev_constraint':
-            tf.keras.constraints.serialize(
-                self.unconstrained_stddev_constraint),
+        'stddev_constraint':
+            tf.keras.constraints.serialize(self.stddev_constraint),
         'seed': self.seed,
         'dtype': self.dtype.name,
     }
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 905222c85..86ac60674 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -29,6 +29,17 @@
 
 class BayesTest(parameterized.TestCase, tf.test.TestCase):
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes
+  def testTrainableNormalStddevConstraint(self):
+    layer = bayes.DenseReparameterization(
+        100, kernel_initializer=bayes.TrainableNormal())
+    inputs = tf.random_normal([1, 1])
+    out = layer(inputs)
+    stddev = layer.kernel.distribution.scale
+    self.evaluate(tf.global_variables_initializer())
+    res, _ = self.evaluate([stddev, out])
+    self.assertAllGreater(res, 0.)
+
   @parameterized.named_parameters(
       {"testcase_name": "_no_uncertainty", "kernel_initializer": "zeros",
        "bias_initializer": "zeros", "all_close": True},

From 01d7a6fc660397f96c1603a08f1726badbb3670d Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 8 Nov 2018 13:17:59 -0800
Subject: [PATCH 1206/2720] Add pixel sampling temperature, make prediction2gif
 run with pure video models.

PiperOrigin-RevId: 220689694
---
 tensor2tensor/models/video/base.py            |  47 ++++++--
 .../video/basic_deterministic_params.py       |   1 +
 tensor2tensor/utils/video/prediction2gif.py   | 110 ++++++++++++------
 3 files changed, 112 insertions(+), 46 deletions(-)

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 5996498b7..c43377067 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -19,12 +19,13 @@
 from __future__ import division
 from __future__ import print_function
 
-from functools import partial
+import functools
 import six
 
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
+from tensor2tensor.layers import discretization
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -39,6 +40,28 @@ def flat_lists(list_of_lists):
   return [x for l in list_of_lists for x in l]
 
 
+def pixels_from_softmax(frame_logits, pure_sampling=False,
+                        temperature=1.0, gumbel_noise_factor=0.2):
+  """Given frame_logits from a per-pixel softmax, generate colors."""
+  # If we're purely sampling, just sample each pixel.
+  if pure_sampling or temperature == 0.0:
+    return common_layers.sample_with_temperature(frame_logits, temperature)
+
+  # Gumbel-sample from the pixel sofmax and average by pixel values.
+  pixel_range = tf.to_float(tf.range(256))
+  for _ in range(len(frame_logits.get_shape().as_list()) - 1):
+    pixel_range = tf.expand_dims(pixel_range, axis=0)
+
+  frame_logits = tf.nn.log_softmax(frame_logits)
+  gumbel_samples = discretization.gumbel_sample(
+      common_layers.shape_list(frame_logits)) * gumbel_noise_factor
+
+  frame = tf.nn.softmax((frame_logits + gumbel_samples) / temperature, axis=-1)
+  result = tf.reduce_sum(frame * pixel_range, axis=-1)
+  # Round on the forward pass, not on the backward one.
+  return result + tf.stop_gradient(tf.round(result) - result)
+
+
 @registry.register_model
 class NextFrameBase(t2t_model.T2TModel):
   """Base class for next_frame models.
@@ -252,9 +275,10 @@ def scheduled_sampling_simple(ground_truth_x, generated_x,
 
       if isinstance(scheduled_sampling_func_var, tf.Tensor):
         tf.summary.scalar("scheduled_sampling_var", scheduled_sampling_func_var)
-      partial_func = partial(scheduled_sampling_func,
-                             batch_size=batch_size,
-                             scheduled_sample_var=scheduled_sampling_func_var)
+      partial_func = functools.partial(
+          scheduled_sampling_func,
+          batch_size=batch_size,
+          scheduled_sample_var=scheduled_sampling_func_var)
       return partial_func
 
   def get_scheduled_sample_inputs(self,
@@ -361,9 +385,8 @@ def get_sampled_frame(self, pred_frame):
       frame_shape = common_layers.shape_list(pred_frame)
       target_shape = frame_shape[:-1] + [self.hparams.problem.num_channels]
       sampled_frame = tf.reshape(pred_frame, target_shape + [256])
-      # TODO(lukaszkaiser): should this be argmax or real sampling.
-      sampled_frame = tf.argmax(sampled_frame, axis=-1)
-      sampled_frame = tf.to_float(sampled_frame)
+      sampled_frame = pixels_from_softmax(
+          sampled_frame, temperature=self.hparams.pixel_sampling_temperature)
       # TODO(lukaszkaiser): this should be consistent with modality.bottom()
       sampled_frame = common_layers.standardize_images(sampled_frame)
     else:
@@ -417,11 +440,15 @@ def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
       inputs_old = features["inputs"]
       features["inputs"] = tf.expand_dims(features["inputs"], 2)
 
-    def logits_to_samples(logits):
+    def logits_to_samples(logits, key):
       """Get samples from logits."""
       # If the last dimension is 1 then we're using L1/L2 loss.
       if common_layers.shape_list(logits)[-1] == 1:
         return tf.to_int32(tf.squeeze(logits, axis=-1))
+      if key == "targets":
+        return pixels_from_softmax(
+            logits, gumbel_noise_factor=0.0,
+            temperature=hparams.pixel_sampling_temperature)
       # Argmax in TF doesn't handle more than 5 dimensions yet.
       logits_shape = common_layers.shape_list(logits)
       argmax = tf.argmax(tf.reshape(logits, [-1, logits_shape[-1]]), axis=-1)
@@ -458,13 +485,13 @@ def logits_to_samples(logits):
     if isinstance(logits, dict):
       results = {}
       for k, v in six.iteritems(logits):
-        results[k] = logits_to_samples(v)
+        results[k] = logits_to_samples(v, k)
         results["%s_logits" % k] = v
       # HACK: bypassing decoding issues.
       results["outputs"] = results["targets"]
       results["scores"] = results["targets"]
     else:
-      results = logits_to_samples(logits)
+      results = logits_to_samples(logits, "targets")
 
     # Restore inputs to not confuse Estimator in edge cases.
     if inputs_old is not None:
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 6ce8f0c39..100342f6d 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -45,6 +45,7 @@ def next_frame_basic_deterministic():
   hparams.add_hparam("residual_dropout", 0.5)
   hparams.add_hparam("num_compress_steps", 6)
   hparams.add_hparam("filter_double_steps", 2)
+  hparams.add_hparam("pixel_sampling_temperature", 0.5)
   return hparams
 
 
diff --git a/tensor2tensor/utils/video/prediction2gif.py b/tensor2tensor/utils/video/prediction2gif.py
index 3be463c6e..b77256d5e 100644
--- a/tensor2tensor/utils/video/prediction2gif.py
+++ b/tensor2tensor/utils/video/prediction2gif.py
@@ -30,11 +30,13 @@
 from __future__ import division
 from __future__ import print_function
 
+import os
+
 import matplotlib as mpl
 import numpy as np
 from queue import Queue
-from tensor2tensor.bin.t2t_decoder import create_hparams
-from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
+
+from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
 from tensor2tensor.layers import common_video
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
@@ -46,7 +48,6 @@
 flags = tf.flags
 FLAGS = flags.FLAGS
 
-
 flags.DEFINE_integer("num_steps", 100, "Number of prediction steps.")
 flags.DEFINE_integer("fps", 10, "Generated gif FPS.")
 flags.DEFINE_string("output_gif", None, "Output path to save the gif.")
@@ -58,14 +59,21 @@ def main(_):
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
 
   # Create hparams
-  hparams = create_hparams()
+  hparams = trainer_lib.create_hparams(
+      FLAGS.hparams_set,
+      FLAGS.hparams,
+      data_dir=os.path.expanduser(FLAGS.data_dir),
+      problem_name=FLAGS.problem)
   hparams.force_full_predict = True
   hparams.scheduled_sampling_k = -1
 
   # Params
   num_agents = 1  # TODO(mbz): fix the code for more agents
   num_steps = FLAGS.num_steps
-  num_actions = hparams.problem.num_actions
+  if hasattr(hparams.problem, "num_actions"):
+    num_actions = hparams.problem.num_actions
+  else:
+    num_actions = None
   frame_shape = hparams.problem.frame_shape
   resized_frame = hparams.preprocess_resize_frames is not None
   if resized_frame:
@@ -75,18 +83,24 @@ def main(_):
   dataset = registry.problem(FLAGS.problem).dataset(
       tf.estimator.ModeKeys.TRAIN,
       shuffle_files=True,
+      data_dir=os.path.expanduser(FLAGS.data_dir),
       hparams=hparams)
 
   dataset = dataset.apply(tf.contrib.data.batch_and_drop_remainder(num_agents))
   data = dataset.make_one_shot_iterator().get_next()
   # Setup input placeholders
   input_size = [num_agents, hparams.video_num_input_frames]
-  placeholders = {
-      "inputs": tf.placeholder(tf.float32, input_size + frame_shape),
-      "input_action": tf.placeholder(tf.int64, input_size + [1]),
-      "input_reward": tf.placeholder(tf.int64, input_size + [1]),
-  }
-  # Creat model
+  if num_actions is None:
+    placeholders = {
+        "inputs": tf.placeholder(tf.float32, input_size + frame_shape)
+    }
+  else:
+    placeholders = {
+        "inputs": tf.placeholder(tf.float32, input_size + frame_shape),
+        "input_action": tf.placeholder(tf.int64, input_size + [1]),
+        "input_reward": tf.placeholder(tf.int64, input_size + [1]),
+    }
+  # Create model.
   model_cls = registry.model(FLAGS.model)
   model = model_cls(hparams, tf.estimator.ModeKeys.PREDICT)
   prediction_ops = model.infer(placeholders)
@@ -94,9 +108,13 @@ def main(_):
   states_q = Queue(maxsize=hparams.video_num_input_frames)
   actions_q = Queue(maxsize=hparams.video_num_input_frames)
   rewards_q = Queue(maxsize=hparams.video_num_input_frames)
-  all_qs = (states_q, actions_q, rewards_q)
+  if num_actions is not None:
+    all_qs = [states_q, actions_q, rewards_q]
+  else:
+    all_qs = [states_q]
 
-  writer = common_video.WholeVideoWriter(fps=10, output_path=FLAGS.output_gif)
+  writer = common_video.WholeVideoWriter(
+      fps=FLAGS.fps, output_path=FLAGS.output_gif)
 
   saver = tf.train.Saver()
   with tf.train.SingularMonitoredSession() as sess:
@@ -113,40 +131,60 @@ def main(_):
       states_q.put(frame)
       writer.write(frame[0].astype(np.uint8))
 
-    actions = np.split(data_np["input_action"],
-                       hparams.video_num_input_frames, 1)
-    for action in actions:
-      actions_q.put(np.squeeze(action, 1))
+    if num_actions is not None:
+      actions = np.split(data_np["input_action"],
+                         hparams.video_num_input_frames, 1)
+      for action in actions:
+        actions_q.put(np.squeeze(action, 1))
 
-    rewards = np.split(data_np["input_reward"],
-                       hparams.video_num_input_frames, 1)
-    for reward in rewards:
-      rewards_q.put(np.squeeze(reward, 1))
+      rewards = np.split(data_np["input_reward"],
+                         hparams.video_num_input_frames, 1)
+      for reward in rewards:
+        rewards_q.put(np.squeeze(reward, 1))
 
     for step in range(num_steps):
       print(">>>>>>> ", step)
 
-      random_actions = np.random.randint(num_actions-1)
-      random_actions = np.expand_dims(random_actions, 0)
-      random_actions = np.tile(random_actions, (num_agents, 1))
-
-      # Shape inputs and targets
-      inputs, input_action, input_reward = (
-          np.stack(list(q.queue), axis=1) for q in all_qs)
+      if num_actions is not None:
+        random_actions = np.random.randint(num_actions-1)
+        random_actions = np.expand_dims(random_actions, 0)
+        random_actions = np.tile(random_actions, (num_agents, 1))
+
+        # Shape inputs and targets
+        inputs, input_action, input_reward = (
+            np.stack(list(q.queue), axis=1) for q in all_qs)
+      else:
+        assert len(all_qs) == 1
+        q = all_qs[0]
+        elems = list(q.queue)
+        # Need to adjust shapes sometimes.
+        for i, e in enumerate(elems):
+          if len(e.shape) < 4:
+            elems[i] = np.expand_dims(e, axis=0)
+        inputs = np.stack(elems, axis=1)
 
       # Predict next frames
-      feed = {
-          placeholders["inputs"]: inputs,
-          placeholders["input_action"]: input_action,
-          placeholders["input_reward"]: input_reward,
-      }
+      if num_actions is None:
+        feed = {placeholders["inputs"]: inputs}
+      else:
+        feed = {
+            placeholders["inputs"]: inputs,
+            placeholders["input_action"]: input_action,
+            placeholders["input_reward"]: input_reward,
+        }
       predictions = sess.run(prediction_ops, feed_dict=feed)
 
-      predicted_states = predictions["targets"][:, 0]
-      predicted_reward = predictions["target_reward"][:, 0]
+      if num_actions is None:
+        predicted_states = predictions[:, 0]
+      else:
+        predicted_states = predictions["targets"][:, 0]
+        predicted_reward = predictions["target_reward"][:, 0]
 
       # Update queues
-      new_data = (predicted_states, random_actions, predicted_reward)
+      if num_actions is None:
+        new_data = (predicted_states)
+      else:
+        new_data = (predicted_states, random_actions, predicted_reward)
       for q, d in zip(all_qs, new_data):
         q.get()
         q.put(d.copy())

From d58aba4b8e2c684cb2a4f31158d61ad7aadb536f Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Thu, 8 Nov 2018 14:23:01 -0800
Subject: [PATCH 1207/2720] Extend latent_dist_encoder=conv3d_net to have
 multiple layers via hparams.latent_encoder_depth and skip connections.

PiperOrigin-RevId: 220701210
---
 tensor2tensor/models/research/glow_ops.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index ebe841674..8c9528880 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -554,18 +554,24 @@ def temporal_tensor_to_dist(name, x, hparams, output_channels=None):
   Returns:
     dist: tf.distributions.Normal
   """
+  res_channels = common_layers.shape_list(x)[-1]
   if output_channels is None:
-    output_channels = common_layers.shape_list(x)[-1]
+    output_channels = res_channels
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
-    h = conv_block("conv3d_init", x, time_filter=2,
-                   mid_channels=hparams.latent_encoder_width)
-    h = conv("conv3d_zeros", h, apply_actnorm=False,
-             output_channels=2*output_channels, conv_init="zeros",
-             filter_size=[2, 3, 3])
+    h = x
+    for i in range(hparams.latent_encoder_depth):
+      h1 = conv_block("conv3d_1_%d" % i, h, time_filter=2,
+                      mid_channels=hparams.latent_encoder_width)
+      h2 = conv("conv3d_zeros_%d" % i, h1, apply_actnorm=False,
+                output_channels=res_channels, conv_init="zeros",
+                filter_size=[2, 3, 3])
+      h += h2
 
     # take last activation that should capture all context since padding is
     # on left.
     h = h[:, -1, :, :, :]
+    h = conv("res_final", h, apply_actnorm=False, conv_init="zeros",
+             output_channels=2*output_channels, filter_size=[1, 1])
     mean, log_scale = h[:, :, :, 0::2], h[:, :, :, 1::2]
   return tf.distributions.Normal(mean, tf.exp(log_scale))
 

From a9164f2c339f745b2a32ec2fb4ae4a6efa775cad Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 8 Nov 2018 14:30:05 -0800
Subject: [PATCH 1208/2720] Rename a RL test to what it really is.

PiperOrigin-RevId: 220702333
---
 .../rl/{rl_trainer_lib_test.py => trainer_model_free_test.py}     | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename tensor2tensor/rl/{rl_trainer_lib_test.py => trainer_model_free_test.py} (100%)

diff --git a/tensor2tensor/rl/rl_trainer_lib_test.py b/tensor2tensor/rl/trainer_model_free_test.py
similarity index 100%
rename from tensor2tensor/rl/rl_trainer_lib_test.py
rename to tensor2tensor/rl/trainer_model_free_test.py

From 1361bb48b99a58dfeae8db7491f101e051b8027f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 8 Nov 2018 14:30:30 -0800
Subject: [PATCH 1209/2720] Toy timeseries problem with no inputs that derives
 from Toy timeseries

PiperOrigin-RevId: 220702433
---
 tensor2tensor/data_generators/timeseries.py   | 35 ++++++++++++++++---
 .../data_generators/timeseries_test.py        | 24 +++++++++++++
 2 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index 5ad49abf1..eb7fd7a9f 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -59,6 +59,10 @@ def dataset_splits(self):
         "shards": self.num_test_shards,
     }]
 
+  @property
+  def has_inputs(self):
+    return True
+
   @property
   def num_train_shards(self):
     """Number of training shards."""
@@ -104,13 +108,15 @@ def normalizing_constant(self):
 
   def preprocess_example(self, example, unused_mode, unused_hparams):
     # Time series are flat on disk, we un-flatten them back here.
-    flat_inputs = example["inputs"]
+    if self.has_inputs:
+      flat_inputs = example["inputs"]
     flat_targets = example["targets"]
     c = self.normalizing_constant
     # Tensor2Tensor models expect [height, width, depth] examples, here we
     # use height for time and set width to 1 and num_series is our depth.
-    example["inputs"] = tf.reshape(
-        flat_inputs, [self.num_input_timestamps, 1, self.num_series]) * c
+    if self.has_inputs:
+      example["inputs"] = tf.reshape(
+          flat_inputs, [self.num_input_timestamps, 1, self.num_series]) * c
     example["targets"] = tf.reshape(
         flat_targets, [self.num_target_timestamps, 1, self.num_series]) * c
     return example
@@ -134,8 +140,13 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
       # We need to flatten the lists on disk for tf,Example to work.
       flat_inputs = [item for sublist in inputs for item in sublist]
       flat_targets = [item for sublist in targets for item in sublist]
-      example_keys = ["inputs", "targets"]
-      ex_dict = dict(zip(example_keys, [flat_inputs, flat_targets]))
+      if self.has_inputs:
+        example_keys = ["inputs", "targets"]
+        ex_dict = dict(zip(example_keys, [flat_inputs, flat_targets]))
+      else:
+        example_keys = ["targets"]
+        ex_dict = dict(zip(example_keys, [flat_targets]))
+
       yield ex_dict
 
   def hparams(self, defaults, unused_model_hparams):
@@ -222,6 +233,20 @@ def timeseries_dataset(self):
     return np.array(series)
 
 
+@registry.register_problem
+class TimeseriesToyProblemNoInputs(TimeseriesToyProblem):
+  """Timeseries problem with a toy dataset and without inputs."""
+
+  @property
+  def has_inputs(self):
+    return False
+
+  @property
+  def num_input_timestamps(self):
+    """Number of timestamps to include in the input."""
+    return 0
+
+
 @registry.register_problem
 class TimeseriesSyntheticDataSeries10Samples100k(TimeseriesProblem):
   """10 synthetic timeseries with 100K samples/timestamps."""
diff --git a/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index bae4b0e0e..a203b5306 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -61,6 +61,30 @@ def testTimeseriesToyProblem(self):
     self.assertNotEqual(
         list(examples[0]["inputs"][0, 0]), list(examples[1]["inputs"][0, 0]))
 
+  def testTimeseriesToyProblemNoInputs(self):
+    problem = timeseries.TimeseriesToyProblemNoInputs()
+    problem.generate_data(self.tmp_dir, self.tmp_dir)
+
+    dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, self.tmp_dir)
+    features = dataset.make_one_shot_iterator().get_next()
+
+    examples = []
+    exhausted = False
+    with self.test_session() as sess:
+      examples.append(sess.run(features))
+      examples.append(sess.run(features))
+      examples.append(sess.run(features))
+      examples.append(sess.run(features))
+      examples.append(sess.run(features))
+
+      try:
+        sess.run(features)
+      except tf.errors.OutOfRangeError:
+        exhausted = True
+
+    self.assertTrue(exhausted)
+    self.assertEqual(5, len(examples))
+
   def testTimeseriesSyntheticData10Series100kSamples(self):
     problem = timeseries.TimeseriesSyntheticDataSeries10Samples100k()
     self.assertEqual(10, problem.num_series)

From 8097343247660f1e7f8018e2e9326198c942cbd6 Mon Sep 17 00:00:00 2001
From: Blazej Osinski <blazej@google.com>
Date: Thu, 8 Nov 2018 19:04:34 -0800
Subject: [PATCH 1210/2720] ppo hparams from original paper.

PiperOrigin-RevId: 220739119
---
 tensor2tensor/models/research/rl.py           | 24 +++++++++++++++++--
 tensor2tensor/rl/README.md                    |  4 ++--
 .../rl/trainer_model_based_params.py          | 13 +++++++++-
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 4d61c3e7e..32f839ba1 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -114,6 +114,26 @@ def ppo_atari_base():
   return hparams
 
 
+@registry.register_hparams
+def ppo_original_params():
+  """Parameters based on the original PPO paper."""
+  hparams = ppo_atari_base()
+  hparams.learning_rate = 2.5e-4
+  hparams.gae_gamma = 0.99
+  hparams.gae_lambda = 0.95
+  hparams.clipping_coef = 0.1
+  hparams.value_loss_coef = 1
+  hparams.entropy_loss_coef = 0.01
+  hparams.eval_every_epochs = 200
+  hparams.dropout_ppo = 0.1
+  # The parameters below are modified to accommodate short epoch_length (which
+  # is needed for model based rollouts).
+  hparams.epoch_length = 50
+  hparams.num_agents = 16
+  hparams.optimization_batch_size = 20
+  return hparams
+
+
 def make_real_env_fn(env):
   """Creates a function returning a given real env, in or out of graph.
 
@@ -159,7 +179,7 @@ def get_policy(observations, hparams, action_space):
 @registry.register_hparams
 def ppo_pong_ae_base():
   """Pong autoencoder base parameters."""
-  hparams = ppo_atari_base()
+  hparams = ppo_original_params()
   hparams.learning_rate = 1e-4
   hparams.network = dense_bitwise_categorical_fun
   return hparams
@@ -200,7 +220,7 @@ def pong_model_free():
 
 @registry.register_hparams
 def mfrl_base():
-  hparams = ppo_atari_base()
+  hparams = ppo_original_params()
   hparams.add_hparam("game", "")
   hparams.epochs_num = 3000
   hparams.eval_every_epochs = 100
diff --git a/tensor2tensor/rl/README.md b/tensor2tensor/rl/README.md
index 4b3f34e67..195222a60 100644
--- a/tensor2tensor/rl/README.md
+++ b/tensor2tensor/rl/README.md
@@ -32,7 +32,7 @@ The hyperparameters for the environment model and agent are nested within the
 ```
   generative_model="next_frame_basic",
   generative_model_params="next_frame_pixel_noise",
-  ppo_params="ppo_atari_base",
+  ppo_params="ppo_original_params",
 ```
 
 ## Model-free training
@@ -53,7 +53,7 @@ Training an agent in `PongNoFrameskip-v0`:
 ```
 python -m tensor2tensor.rl.trainer_model_free \
   --problem stacked_pong \
-  --hparams_set ppo_atari_base \
+  --hparams_set ppo_original_params \
   --hparams num_agents=5 \
   --output_dir dir_location
 ```
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 18c1b4e26..510b15762 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -55,7 +55,7 @@ def rlmb_base():
       generative_model="next_frame_basic_deterministic",
       generative_model_params="next_frame_pixel_noise",
       base_algo="ppo",
-      base_algo_params="ppo_atari_base",
+      base_algo_params="ppo_original_params",
       autoencoder_train_steps=0,
       autoencoder_train_steps_initial_multiplier=10,
       autoencoder_hparams_set="autoencoder_discrete_pong",
@@ -690,6 +690,17 @@ def rlmb_logits_clip(rhp):
   rhp.set_discrete("ppo.logits_clip", [0., 5.])
 
 
+@registry.register_ranged_hparams
+def rlmb_games_problematic_for_ppo(rhp):
+  games = [
+      "alien", "boxing", "breakout", "ms_pacman", "video_pinball",
+  ]
+  rhp.set_categorical("loop.game", games)
+  rhp.set_categorical("loop.base_algo_params", ["ppo_original_params"])
+  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
+  rhp.set_discrete("ppo.dropout_ppo", [0., 0.1])
+
+
 @registry.register_ranged_hparams
 def rlmf_proportional_epoch_length(rhp):
   rhp.set_discrete("proportional_epoch_length", [10, 20, 50, 100, 200, 400])

From 2eeffff184e8bcb2c76510fa75616c51f45bef8f Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Thu, 8 Nov 2018 20:55:08 -0800
Subject: [PATCH 1211/2720] Maintenance Cl.

PiperOrigin-RevId: 220748193
---
 tensor2tensor/models/research/glow_ops.py     | 71 +++++++++++++------
 .../models/research/glow_ops_test.py          | 18 ++---
 2 files changed, 57 insertions(+), 32 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 8c9528880..003403087 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -543,8 +543,8 @@ def squeeze(name, x, factor=2, reverse=True):
 
 
 @add_arg_scope
-def temporal_tensor_to_dist(name, x, hparams, output_channels=None):
-  """Network that maps a time-indexed list of 3-D Tensors to a gaussian.
+def temporal_latent_to_dist(name, x, hparams, output_channels=None):
+  """Network that maps a time-indexed list of 3-D latents to a gaussian.
 
   Args:
     name: variable scope.
@@ -577,32 +577,59 @@ def temporal_tensor_to_dist(name, x, hparams, output_channels=None):
 
 
 @add_arg_scope
-def tensor_to_dist(name, x, output_channels=None, architecture="single_conv",
-                   depth=1, pre_output_channels=512, width=512):
-  """Map x to the mean and log-scale of a Gaussian.
+def single_conv_dist(name, x, output_channels=None):
+  """A 3x3 convolution mapping x to a standard normal distribution at init.
+
+  Args:
+    name: variable scope.
+    x: 4-D Tensor.
+    output_channels: number of channels of the mean and std.
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    x_shape = common_layers.shape_list(x)
+    if output_channels is None:
+      output_channels = x_shape[-1]
+    mean_log_scale = conv("conv2d", x, output_channels=2*output_channels,
+                          conv_init="zeros", apply_actnorm=False)
+    mean = mean_log_scale[:, :, :, 0::2]
+    log_scale = mean_log_scale[:, :, :, 1::2]
+    return tf.distributions.Normal(mean, tf.exp(log_scale))
+
+
+@add_arg_scope
+def latent_to_dist(name, x, hparams, output_channels=None):
+  """Map latent to the mean and log-scale of a Gaussian.
 
   Args:
     name: variable scope.
     x: 4-D Tensor of shape (NHWC)
-    output_channels: int, number of output channels of the mean.
+    hparams: tf.contrib.training.HParams.
+      latent_architecture - can be "single_conv", "glow_nn" or "glow_resnet",
+                            default = single_conv
+      latent_encoder_depth - int, depth of architecture, valid if
+                             latent_architecture is "glow_nn" or "glow_resnet".
+      latent_pre_output_channels - 512, valid only when latent_architecture
+                                   is "glow_nn".
+      latent_encoder_width - 512, maximum width of the network
+    output_channels: int, number of output channels of the mean (and std).
                      if not provided, set it to be the output channels of x.
-    architecture: "single_conv" or "glow_nn"
-    depth: depth of architecture mapping to the mean and std.
-    pre_output_channels: output channels before the final (mean, std) mapping.
-    width: Resnet width.
   Returns:
     dist: instance of tf.distributions.Normal
   Raises:
     ValueError: If architecture not in ["single_conv", "glow_nn"]
   """
+  architecture = hparams.get("latent_architecture", "single_conv")
+  depth = hparams.get("latent_encoder_depth", 1)
+  pre_output_channels = hparams.get("latent_pre_output_channels", 512)
+  width = hparams.get("latent_encoder_width", 512)
+
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     x_shape = common_layers.shape_list(x)
     if output_channels is None:
       output_channels = x_shape[-1]
     if architecture == "single_conv":
-      mean_log_scale = conv("conv2d", x, output_channels=2*output_channels,
-                            conv_init="zeros", apply_actnorm=False)
-    elif architecture == "glow_nn":
+      return single_conv_dist("single_conv", x, output_channels)
+    if architecture == "glow_nn":
       mean_log_scale = x
       for layer in range(1, depth + 1):
         mid_channels = pre_output_channels // 2**(depth - layer)
@@ -695,12 +722,9 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
     output_channels = common_layers.shape_list(z)[-1]
     last_latent = latent[-1]
     latent_stack = tf.concat([prior_dist.loc] + latent, axis=-1)
-    cond_dist = tensor_to_dist(
-        "latent_stack", latent_stack, output_channels=output_channels,
-        architecture=hparams.latent_architecture,
-        depth=hparams.latent_encoder_depth,
-        pre_output_channels=hparams.latent_pre_output_channels,
-        width=hparams.latent_encoder_width)
+    cond_dist = latent_to_dist(
+        "latent_stack", latent_stack, hparams=hparams,
+        output_channels=output_channels)
 
   elif latent_dist_encoder == "conv3d_net":
     last_latent = latent[-1]
@@ -714,7 +738,7 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
     prev_latents = tf.tile(tf.expand_dims(prior_dist.loc, axis=1),
                            [1, num_steps, 1, 1, 1])
     cond_latents = tf.concat((cond_latents, prev_latents), axis=-1)
-    cond_dist = temporal_tensor_to_dist(
+    cond_dist = temporal_latent_to_dist(
         "latent_stack", cond_latents, hparams, output_channels=output_channels)
 
   elif latent_dist_encoder == "conv_lstm":
@@ -724,7 +748,8 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
     _, state = common_video.conv_lstm_2d(
         latent_stack, state, hparams.latent_encoder_width, kernel_size=3,
         name="conv_lstm")
-    cond_dist = tensor_to_dist(
+
+    cond_dist = single_conv_dist(
         "state_to_dist", state.h, output_channels=output_channels)
   if latent_skip:
     new_mean = cond_dist.loc + last_latent
@@ -759,7 +784,7 @@ def compute_prior(name, z, latent, hparams, condition=False, state=None):
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     if isinstance(condition, bool):
       condition = tf.constant(condition, dtype=tf.bool)
-    prior_dist = tensor_to_dist("level_prior", z, architecture="single_conv")
+    prior_dist = single_conv_dist("level_prior", z)
     prior_mean, prior_scale = prior_dist.loc, prior_dist.scale
 
     if latent is None:
@@ -923,7 +948,7 @@ def top_prior(name, z_shape, learn_prior="normal"):
     if learn_prior == "normal":
       prior_dist = tf.distributions.Normal(h, tf.exp(h))
     elif learn_prior == "single_conv":
-      prior_dist = tensor_to_dist("top_learn_prior", h)
+      prior_dist = single_conv_dist("top_learn_prior", h)
     else:
       raise ValueError("Expected learn_prior to be normal or single_conv "
                        "got %s" % learn_prior)
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index a19cc1272..164783fca 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -139,11 +139,11 @@ def test_affine_coupling_network(self):
         # Initialized with zeros.
         self.assertTrue(np.allclose(nn_np, 0.0))
 
-  def check_tensor_to_dist(self, architecture):
+  def check_latent_to_dist(self, architecture):
     with tf.Graph().as_default():
       x = tf.random_uniform(shape=(16, 5, 5, 32))
-      x_prior = glow_ops.tensor_to_dist("split_prior", x,
-                                        architecture=architecture,
+      hparams = tf.contrib.training.HParams(architecture=architecture)
+      x_prior = glow_ops.latent_to_dist("split_prior", x, hparams=hparams,
                                         output_channels=64)
       mean_t, scale_t = x_prior.loc, x_prior.scale
       with tf.Session() as session:
@@ -154,9 +154,9 @@ def check_tensor_to_dist(self, architecture):
         self.assertTrue(np.allclose(mean, 0.0))
         self.assertTrue(np.allclose(scale, 1.0))
 
-  def test_tensor_to_dist(self):
+  def test_latent_to_dist(self):
     for architecture in ["single_conv", "glow_nn", "glow_resnet"]:
-      self.check_tensor_to_dist(architecture)
+      self.check_latent_to_dist(architecture)
 
   def test_split(self):
     with tf.Graph().as_default():
@@ -205,8 +205,8 @@ def test_encoder_decoder(self):
         diff, x_inv_np, z_levels_np, z_inv_levels_np = session.run(
             [x - x_inv_inv, x_inv, z_levels, z_inv_levels])
 
-        self.assertEqual(len(z_levels_np), 2)
-        self.assertEqual(len(z_inv_levels_np), 2)
+        self.assertLen(z_levels_np, 2)
+        self.assertLen(z_inv_levels_np, 2)
         # (h_i, w_i, c_i) = (h_{i-1}/f, w_{i-1}/f, c_{i-1}*(2f)/2) where (f=2)
         self.assertEqual(z_levels_np[0].shape, (16, 32, 32, 8))
         self.assertEqual(z_levels_np[1].shape, (16, 16, 16, 16))
@@ -404,12 +404,12 @@ def test_actnorm_3d(self):
           self.assertTrue(np.allclose(channel_mean, 0.0, atol=1e-3))
           self.assertTrue(np.allclose(channel_var, 1.0, atol=1e-3))
 
-  def test_temporal_tensor_to_dist(self):
+  def test_temporal_latent_to_dist(self):
     with tf.Graph().as_default():
       hparams = self.get_glow_hparams()
       latent_shape = (16, 5, 4, 4, 48)
       latents = tf.random_normal(latent_shape)
-      dist = glow_ops.temporal_tensor_to_dist(
+      dist = glow_ops.temporal_latent_to_dist(
           "tensor_to_dist", latents, hparams)
       with tf.Session() as sess:
         sess.run(tf.global_variables_initializer())

From 14ee3052aeec4e098d9dceaf4804837c1522fb3d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 8 Nov 2018 22:31:22 -0800
Subject: [PATCH 1212/2720] Enable general beam_decode_slow on TPU.

PiperOrigin-RevId: 220754515
---
 tensor2tensor/utils/beam_search.py |  4 +++-
 tensor2tensor/utils/t2t_model.py   | 14 +++++++-------
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index fe51422a2..5ac457e36 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -557,7 +557,7 @@ def grow_topk(i, alive_seq, alive_log_probs, states):
          dict of transformed decoding states)
     """
     # Get the logits for all the possible next symbols
-    if use_tpu:
+    if use_tpu and states:
       flat_ids = tf.reshape(
           tf.slice(alive_seq, [0, 0, i], [batch_size, beam_size, 1]),
           [batch_size * beam_size, -1])
@@ -570,6 +570,8 @@ def grow_topk(i, alive_seq, alive_log_probs, states):
       flat_logits, flat_states = symbols_to_logits_fn(flat_ids, i, flat_states)
       states = nest.map_structure(
           lambda t: _unmerge_beam_dim(t, batch_size, beam_size), flat_states)
+    elif use_tpu:
+      flat_logits = symbols_to_logits_fn(flat_ids, i)
     else:
       flat_logits = symbols_to_logits_fn(flat_ids)
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index a279f2234..58d8f4da0 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -739,13 +739,9 @@ def _beam_decode_slow(self, features, decode_length, beam_size, top_beams,
     Raises:
       NotImplementedError: If use_tpu is set to true.
     """
-    if use_tpu:
-      raise NotImplementedError(
-          "Slow beam search inference on TPU is not supported")
-
     batch_size = common_layers.shape_list(features["inputs"])[0]
 
-    def symbols_to_logits_fn(ids):
+    def symbols_to_logits_fn(ids, i=None):
       """Go from ids to logits."""
       ids = tf.expand_dims(tf.expand_dims(ids, axis=2), axis=3)
       ids = tf.pad(ids[:, 1:], [[0, 0], [0, 1], [0, 0], [0, 0]])
@@ -757,6 +753,8 @@ def symbols_to_logits_fn(ids):
         ids = tf.concat([pt, ids], axis=1)
 
       features["targets"] = ids
+      if i is not None:
+        features["decode_loop_step"] = i
       self._coverage = None
       logits, _ = self(features)  # pylint: disable=not-callable
       # now self._coverage is a coverage tensor for the first datashard.
@@ -786,7 +784,6 @@ def symbols_to_logits_fn(ids):
     target_modality = self._problem_hparams.modality["targets"]
     vocab_size = target_modality.top_dimensionality
     # Setting decode length to input length + decode_length
-    decode_length = tf.constant(decode_length)
     if "partial_targets" not in features:
       inputs = features["inputs"]
       decode_length = (common_layers.shape_list(inputs)[1] +
@@ -798,7 +795,8 @@ def symbols_to_logits_fn(ids):
         decode_length,
         vocab_size,
         alpha,
-        stop_early=(top_beams == 1))
+        stop_early=(top_beams == 1),
+        use_tpu=use_tpu)
 
     # Set inputs back to the unexpanded inputs to not to confuse the Estimator!
     if self.has_input:
@@ -1542,6 +1540,8 @@ def estimator_spec_predict(self, features, use_tpu=False):
     # Pass through remaining features
     for name, feature in features.items():
       if name not in list(predictions.keys()) + ["infer_targets"]:
+        if name == "decode_loop_step":
+          continue
         if not feature.shape.as_list():
           # All features must have a batch dimension
           batch_size = common_layers.shape_list(outputs)[0]

From 013c52c4b0bdc77388f16df9df0753e6dd39a5f9 Mon Sep 17 00:00:00 2001
From: Amit Patankar <amitpatankar@google.com>
Date: Fri, 9 Nov 2018 09:50:01 -0800
Subject: [PATCH 1213/2720] Internal change

PiperOrigin-RevId: 220817646
---
 tensor2tensor/layers/common_layers.py     | 11 ++++-----
 tensor2tensor/models/research/glow.py     |  2 +-
 tensor2tensor/models/research/glow_ops.py | 27 ++++++++++++-----------
 tensor2tensor/models/research/rl.py       |  6 ++---
 4 files changed, 24 insertions(+), 22 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index e2246d40c..c72bdc326 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -29,6 +29,7 @@
 from six.moves import range  # pylint: disable=redefined-builtin
 
 import tensorflow as tf
+import tensorflow_probability as tfp
 
 from tensorflow.python.framework import function
 from tensorflow.python.framework import ops
@@ -1971,7 +1972,7 @@ def smoothing_cross_entropy(logits,
     if gaussian and confidence > 0.0:
       labels = tf.cast(labels, tf.float32)
 
-      normal_dist = tf.distributions.Normal(loc=labels, scale=confidence)
+      normal_dist = tfp.distributions.Normal(loc=labels, scale=confidence)
       # Locations to evaluate the probability distributions.
       soft_targets = normal_dist.prob(
           tf.cast(tf.range(vocab_size), tf.float32)[:, None, None, None, None])
@@ -3730,12 +3731,12 @@ def kl_divergence(mu, log_var, mu_p=0.0, log_var_p=0.0):
   """
 
   batch_size = shape_list(mu)[0]
-  prior_distribution = tf.distributions.Normal(
+  prior_distribution = tfp.distributions.Normal(
       mu_p, tf.exp(tf.multiply(0.5, log_var_p)))
-  posterior_distribution = tf.distributions.Normal(
+  posterior_distribution = tfp.distributions.Normal(
       mu, tf.exp(tf.multiply(0.5, log_var)))
-  kld = tf.distributions.kl_divergence(
-      posterior_distribution, prior_distribution)
+  kld = tfp.distributions.kl_divergence(posterior_distribution,
+                                        prior_distribution)
   return tf.reduce_sum(kld) / tf.to_float(batch_size)
 
 
diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index e44aba705..3848e6a58 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -137,7 +137,7 @@ def top_prior(self):
     """Objective based on the prior over latent z.
 
     Returns:
-      dist: instance of tf.distributions.Normal, prior distribution.
+      dist: instance of tfp.distributions.Normal, prior distribution.
     """
     return glow_ops.top_prior(
         "top_prior", self.z_top_shape, learn_prior=self.hparams.top_prior)
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 003403087..6b59bd8b1 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -25,6 +25,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 import tensorflow as tf
+import tensorflow_probability as tfp
 
 arg_scope = tf.contrib.framework.arg_scope
 add_arg_scope = tf.contrib.framework.add_arg_scope
@@ -552,7 +553,7 @@ def temporal_latent_to_dist(name, x, hparams, output_channels=None):
     hparams: tf.contrib.training.Hparams.
     output_channels: int, Number of channels of the output gaussian mean.
   Returns:
-    dist: tf.distributions.Normal
+    dist: tfp.distributions.Normal
   """
   res_channels = common_layers.shape_list(x)[-1]
   if output_channels is None:
@@ -573,7 +574,7 @@ def temporal_latent_to_dist(name, x, hparams, output_channels=None):
     h = conv("res_final", h, apply_actnorm=False, conv_init="zeros",
              output_channels=2*output_channels, filter_size=[1, 1])
     mean, log_scale = h[:, :, :, 0::2], h[:, :, :, 1::2]
-  return tf.distributions.Normal(mean, tf.exp(log_scale))
+  return tfp.distributions.Normal(mean, tf.exp(log_scale))
 
 
 @add_arg_scope
@@ -614,7 +615,7 @@ def latent_to_dist(name, x, hparams, output_channels=None):
     output_channels: int, number of output channels of the mean (and std).
                      if not provided, set it to be the output channels of x.
   Returns:
-    dist: instance of tf.distributions.Normal
+    dist: instance of tfp.distributions.Normal
   Raises:
     ValueError: If architecture not in ["single_conv", "glow_nn"]
   """
@@ -655,7 +656,7 @@ def latent_to_dist(name, x, hparams, output_channels=None):
 
     mean = mean_log_scale[:, :, :, 0::2]
     log_scale = mean_log_scale[:, :, :, 1::2]
-    return tf.distributions.Normal(mean, tf.exp(log_scale))
+    return tfp.distributions.Normal(mean, tf.exp(log_scale))
 
 
 @add_arg_scope
@@ -667,11 +668,11 @@ def merge_level_and_latent_dist(level_dist, latent_dist,
   according to merge_std.
 
   Args:
-    level_dist: instance of tf.distributions.Normal
-    latent_dist: instance of tf.distributions.Normal
+    level_dist: instance of tfp.distributions.Normal
+    latent_dist: instance of tfp.distributions.Normal
     merge_std: can be "prev_level", "prev_step" or "normal".
   Returns:
-    merged_dist: instance of tf.distributions.Normal
+    merged_dist: instance of tfp.distributions.Normal
   """
   level_mean, level_std = level_dist.loc, level_dist.scale
   latent_mean, latent_std = latent_dist.loc, latent_dist.scale
@@ -686,7 +687,7 @@ def merge_level_and_latent_dist(level_dist, latent_dist,
     scale = level_std
   elif merge_std == "prev_step":
     scale = latent_std
-  return tf.distributions.Normal(loc=new_mean, scale=scale)
+  return tfp.distributions.Normal(loc=new_mean, scale=scale)
 
 
 @add_arg_scope
@@ -753,7 +754,7 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
         "state_to_dist", state.h, output_channels=output_channels)
   if latent_skip:
     new_mean = cond_dist.loc + last_latent
-    cond_dist = tf.distributions.Normal(new_mean, cond_dist.scale)
+    cond_dist = tfp.distributions.Normal(new_mean, cond_dist.scale)
   return cond_dist.loc, cond_dist.scale, state
 
 
@@ -775,7 +776,7 @@ def compute_prior(name, z, latent, hparams, condition=False, state=None):
            the current state of a LSTM used to model the distribution. Used
            only if hparams.latent_dist_encoder = "conv_lstm".
   Returns:
-    prior_dist: instance of tf.distributions.Normal
+    prior_dist: instance of tfp.distributions.Normal
     state: Returns updated state.
   Raises:
     ValueError: If hparams.latent_dist_encoder is "pointwise" and if the shape
@@ -795,7 +796,7 @@ def compute_prior(name, z, latent, hparams, condition=False, state=None):
       mean, scale = tf.cond(
           condition, lambda: (cond_mean, cond_scale),
           lambda: (prior_mean, prior_scale))
-    dist = tf.distributions.Normal(mean, scale)
+    dist = tfp.distributions.Normal(mean, scale)
     return dist, state
 
 
@@ -921,7 +922,7 @@ def scale_gaussian_prior(name, z, logscale_factor=3.0, trainable=True):
         "log_scale_latent", shape=z_shape, dtype=tf.float32,
         initializer=tf.zeros_initializer(), trainable=trainable)
     log_scale = log_scale * logscale_factor
-    return tf.distributions.Normal(
+    return tfp.distributions.Normal(
         loc=latent_multiplier * z, scale=tf.exp(log_scale))
 
 
@@ -946,7 +947,7 @@ def top_prior(name, z_shape, learn_prior="normal"):
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     h = tf.zeros(z_shape, dtype=tf.float32)
     if learn_prior == "normal":
-      prior_dist = tf.distributions.Normal(h, tf.exp(h))
+      prior_dist = tfp.distributions.Normal(h, tf.exp(h))
     elif learn_prior == "single_conv":
       prior_dist = single_conv_dist("top_learn_prior", h)
     else:
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 32f839ba1..ba03da8e9 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -401,7 +401,7 @@ def random_policy_fun(action_space, unused_config, observations):
   obs_shape = observations.shape.as_list()
   with tf.variable_scope("network_parameters"):
     value = tf.zeros(obs_shape[:2])
-    policy = tf.distributions.Categorical(
-        probs=[[[1. / float(action_space.n)] * action_space.n
-               ] * (obs_shape[0] * obs_shape[1])])
+    policy = tfp.distributions.Categorical(
+        probs=[[[1. / float(action_space.n)] * action_space.n] *
+               (obs_shape[0] * obs_shape[1])])
   return NetworkOutput(policy, value, lambda a: a)

From 255593bc1a2ce52161b468f2f036982acaa81321 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 9 Nov 2018 10:57:07 -0800
Subject: [PATCH 1214/2720] emergency fix of L2 modality.

PiperOrigin-RevId: 220829770
---
 tensor2tensor/layers/common_layers.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index c72bdc326..5bc2cad11 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -237,7 +237,6 @@ def convert_real_to_rgb(x):
   """Conversion of real numbers to pixel values."""
   with tf.name_scope("real_to_rgb", values=[x]):
     x *= 255.0
-    x = tf.round(x)
     return x
 
 

From fbd70f663be2818d32188d89b578d46ad7c462c7 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 9 Nov 2018 17:00:06 -0800
Subject: [PATCH 1215/2720] First translation Problem to be ran in
 multi-lingual setting.

PiperOrigin-RevId: 220887623
---
 tensor2tensor/data_generators/translate_enfr.py  | 14 ++++++++++++++
 .../data_generators/wiki_multi_problems.py       | 16 ++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index 73261ffb1..8087286d7 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -25,6 +25,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.data_generators import translate
+from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -233,3 +234,16 @@ class TranslateEnfrWmtCharacters(TranslateEnfrWmtSmallCharacters):
   @property
   def use_small_dataset(self):
     return False
+
+
+@registry.register_problem
+class TranslateEnfrWmtMulti64k(TranslateEnfrWmtSmall32k):
+  """Translation with muli-lingual vocabulary."""
+
+  @property
+  def use_small_dataset(self):
+    return False
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
diff --git a/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index d46b1f793..9ea44a94b 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -23,6 +23,7 @@
 from tensor2tensor.data_generators import multi_problem
 from tensor2tensor.data_generators import multinli
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators import translate_enfr
 from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
 
@@ -103,3 +104,18 @@ def __init__(self, was_reversed=False, was_copy=False):
   @property
   def vocab_type(self):
     return text_problems.VocabType.SUBWORD
+
+
+@registry.register_problem
+class LanguagemodelMultiWikiTranslateFr(multi_problem.MultiProblem):
+  """Wiki multi-lingual LM and En-Fr translation."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelMultiWikiTranslateFr, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelDeEnFrRoWiki64k())
+    self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD

From 4fddc0fb6a7a86278cbdf6b38898b171f1f3aa1e Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 9 Nov 2018 17:10:23 -0800
Subject: [PATCH 1216/2720] one more bug(?) fix. this is to allow backproping
 through previously predicted frame after conversion from float to int.

PiperOrigin-RevId: 220889018
---
 tensor2tensor/models/video/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index c43377067..397c26166 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -391,7 +391,7 @@ def get_sampled_frame(self, pred_frame):
       sampled_frame = common_layers.standardize_images(sampled_frame)
     else:
       x = common_layers.convert_real_to_rgb(pred_frame)
-      x = tf.cast(x, tf.uint8)
+      x = x - tf.stop_gradient(x + tf.round(x))
       x = common_layers.convert_rgb_to_real(x)
       return x
     return sampled_frame

From 9a5d30fd706dd40f22f1c2618573e112293563f6 Mon Sep 17 00:00:00 2001
From: Mike Dusenberry <dusenberrymw@google.com>
Date: Fri, 9 Nov 2018 19:03:21 -0800
Subject: [PATCH 1217/2720] Fix KL regularization for Bayesian Layers.

Currently, computing derivatives of the sum of `layer.losses` w.r.t. to the
weight distribution variables yields `None` gradients in Eager mode.

This appears to be due to TF ops in the tfp.distribution `__init__` methods.
Since the regularization tensor is created at build time, those distribution
ops are tied to the context in which the distribution was created.  If we use a
`GradientTape`, and the layer gets built outside of the tape, the gradients
will always be equal to `None`.  If we build inside the `GradientTape`, say on
the first training epoch/step, the gradients will be computed correctly on that
iteration.  However, on the next iteration, in which a new tape is created, the
gradients will be equal to `None` again due to the TF ops that were created
inside the distribution's `__init__` method being tied to the first iteration's
tape.

This bug is particularly problematic because it won't appear to be broken if
the loss involves the usual likelihood term.  More specifically, there will be
gradients w.r.t. the variables, but they will only be from the likelihood term
and not from the KL term.  Additionally, the KL term will still produce a
value, so printing/summarizing that term will give the appearance that
everything is working.

This presents a fix that involves creating a callable for the "variable" when
creating the regularization tensor, with the condition that we must create a
new Edward RV within that callable.  A more permanent fix could involve making
changes to Edward RandomVariables or removing TF ops from the `__init__`
methods of TFP distributions.

PiperOrigin-RevId: 220898007
---
 tensor2tensor/layers/bayes.py      | 79 ++++++++++++++--------------
 tensor2tensor/layers/bayes_test.py | 83 ++++++++++++++++++++++++++++++
 2 files changed, 123 insertions(+), 39 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 821033084..882840296 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -19,7 +19,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import functools
 import tensorflow as tf
 
 from tensorflow_probability import edward2 as ed
@@ -239,8 +238,8 @@ def build(self, input_shape):
                                     self.dtype,
                                     self.add_weight)
       if self.kernel_regularizer is not None:
-        self._handle_weight_regularization(
-            'kernel', self.kernel, self.kernel_regularizer)
+        self.add_loss(create_regularization_loss_fn(
+            'kernel', lambda: self.kernel, self.kernel_regularizer))
 
     else:
       self._kernel = self.add_weight(
@@ -256,8 +255,8 @@ def build(self, input_shape):
       if isinstance(self.bias_initializer, TrainableInitializer):
         self.bias_initializer.build([self.units], self.dtype, self.add_weight)
         if self.bias_regularizer is not None:
-          self._handle_weight_regularization(
-              'bias', self.bias, self.bias_regularizer)
+          self.add_loss(create_regularization_loss_fn(
+              'bias', lambda: self.bias, self.bias_regularizer))
       else:
         self._bias = self.add_weight(
             'bias',
@@ -272,21 +271,6 @@ def build(self, input_shape):
       self._bias = None
     self.built = True
 
-  # TODO(trandustin): Waiting on T2T to drop dependence on
-  # TF<=1.12rc2. A TF commit enables tf.colocate_with to work for
-  # Tensor-like inputs. This lets us use the parent method instead of
-  # this one.
-  def _handle_weight_regularization(self, name, variable, regularizer):
-    """Create lambdas which compute regularization losses."""
-
-    def _loss_for_variable(v):
-      """Creates a regularization loss `Tensor` for variable `v`."""
-      with tf.name_scope(name + '/Regularizer'):
-        regularization = regularizer(v)
-      return regularization
-
-    self.add_loss(functools.partial(_loss_for_variable, variable))
-
 
 class LSTMCellReparameterization(tf.keras.layers.LSTMCell):
   """Bayesian LSTM cell class estimated via reparameterization.
@@ -365,8 +349,11 @@ def build(self, input_shape):
           [input_dim, self.units * 4], self.dtype, self.add_weight)
       self.kernel = self.kernel_initializer()
       if self.kernel_regularizer is not None:
-        self._handle_weight_regularization(
-            'kernel', self.kernel, self.kernel_regularizer)
+        self.add_loss(create_regularization_loss_fn(
+            # Can't use the kernel directly because we actually need to create a
+            # new Edward RV.  The Dense layer already does this.
+            # Also note that the initializer is a callable.
+            'kernel', self.kernel_initializer, self.kernel_regularizer))
 
     else:
       self.kernel = self.add_weight(
@@ -381,9 +368,12 @@ def build(self, input_shape):
           [self.units, self.units * 4], self.dtype, self.add_weight)
       self.recurrent_kernel = self.recurrent_initializer()
       if self.recurrent_regularizer is not None:
-        self._handle_weight_regularization(
-            'recurrent_kernel', self.recurrent_kernel,
-            self.recurrent_regularizer)
+        self.add_loss(create_regularization_loss_fn(
+            # Can't use the kernel directly because we actually need to create a
+            # new Edward RV.  The Dense layer already does this.
+            # Also note that the initializer is a callable.
+            'recurrent_kernel', self.recurrent_initializer,
+            self.recurrent_regularizer))
 
     else:
       self.recurrent_kernel = self.add_weight(
@@ -399,8 +389,11 @@ def build(self, input_shape):
             [self.units * 4], self.dtype, self.add_weight)
         self.bias = self.bias_initializer()
         if self.bias_regularizer is not None:
-          self._handle_weight_regularization(
-              'bias', self.bias, self.bias_regularizer)
+          self.add_loss(create_regularization_loss_fn(
+              # Can't use the bias directly because we actually need to create a
+              # new Edward RV.  The Dense layer already does this.
+              # Also note that the initializer is a callable.
+              'bias', self.bias_initializer, self.bias_regularizer))
       else:
         if self.unit_forget_bias:
 
@@ -438,17 +431,25 @@ def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
     return super(LSTMCellReparameterization, self).get_initial_state(
         inputs=inputs, batch_size=batch_size, dtype=dtype)
 
-  # TODO(trandustin): Waiting on T2T to drop dependence on
-  # TF<=1.12rc2. A TF commit enables tf.colocate_with to work for
-  # Tensor-like inputs. This lets us use the parent method instead of
-  # this one.
-  def _handle_weight_regularization(self, name, variable, regularizer):
-    """Create lambdas which compute regularization losses."""
 
-    def _loss_for_variable(v):
-      """Creates a regularization loss `Tensor` for variable `v`."""
-      with tf.name_scope(name + '/Regularizer'):
-        regularization = regularizer(v)
-      return regularization
+def create_regularization_loss_fn(name, variable_fn, regularizer_fn):
+  """Create a regularization loss function.
+
+  The callable representing the variable allows for use with Bayesian Layers.
+
+  Args:
+    name: String name scope prefix.
+    variable_fn: Callable that returns a TF Variable or ed.RandomVariable.
+    regularizer_fn: Callable that returns a loss tensor when called with a TF
+      Variable or ed.RandomVariable.
+
+  Returns:
+    A callable that returns a regularization loss tensor when called.
+  """
+  def loss_fn():
+    """Creates a regularization loss `Tensor`."""
+    with tf.name_scope(name + '/Regularizer'):
+      regularization = regularizer_fn(variable_fn())
+    return regularization
 
-    self.add_loss(functools.partial(_loss_for_variable, variable))
+  return loss_fn
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 86ac60674..4addad626 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -67,6 +67,43 @@ def testDenseReparameterizationKernel(
       self.assertNotAllClose(res1, res2)
     layer.get_config()
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testDenseReparameterizationKL(self):
+    inputs = tf.to_float(np.random.rand(5, 12))
+    layer = bayes.DenseReparameterization(10)
+
+    # Imagine this is the 1st epoch.
+    with tf.GradientTape() as tape:
+      layer(inputs)  # first call forces a build, here inside this tape
+      layer(inputs)  # ensure robustness after multiple calls
+      loss = tf.reduce_sum([tf.reduce_sum(l) for l in layer.losses])
+
+    variables = [layer.kernel_initializer.mean, layer.kernel_initializer.stddev]
+    for v in variables:
+      self.assertIn(v, layer.variables)
+
+    # This will be fine, since the layer was built inside this tape, and thus
+    # the distribution init ops were inside this tape.
+    grads = tape.gradient(loss, variables)
+    for grad in grads:
+      self.assertIsNotNone(grad)
+
+    # Imagine this is the 2nd epoch.
+    with tf.GradientTape() as tape:
+      layer(inputs)  # build won't be called again
+      loss = tf.reduce_sum([tf.reduce_sum(l) for l in layer.losses])
+
+    variables = [layer.kernel_initializer.mean, layer.kernel_initializer.stddev]
+    for v in variables:
+      self.assertIn(v, layer.variables)
+
+    # This would fail, since the layer was built inside the tape from the 1st
+    # epoch, and thus the distribution init ops were inside that tape instead of
+    # this tape. By using a callable for the variable, this will no longer fail.
+    grads = tape.gradient(loss, variables)
+    for grad in grads:
+      self.assertIsNotNone(grad)
+
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDenseReparameterizationModel(self):
     inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
@@ -126,6 +163,52 @@ def testLSTMCellReparameterization(
       self.assertNotAllClose(res1, res3)
     cell.get_config()
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testLSTMCellReparameterizationKL(self):
+    inputs = tf.to_float(np.random.rand(5, 1, 12))
+    cell = bayes.LSTMCellReparameterization(10)
+    state = (tf.zeros([1, 10]), tf.zeros([1, 10]))
+
+    # Imagine this is the 1st epoch.
+    with tf.GradientTape() as tape:
+      cell(inputs[:, 0, :], state)  # first call forces a build, inside the tape
+      cell(inputs[:, 0, :], state)  # ensure robustness after multiple calls
+      cell.get_initial_state(inputs[:, 0, :])
+      cell(inputs[:, 0, :], state)  # ensure robustness after multiple calls
+      loss = tf.reduce_sum([tf.reduce_sum(l) for l in cell.losses])
+
+    variables = [
+        cell.kernel_initializer.mean, cell.kernel_initializer.stddev,
+        cell.recurrent_initializer.mean, cell.recurrent_initializer.stddev,
+    ]
+    for v in variables:
+      self.assertIn(v, cell.variables)
+
+    # This will be fine, since the layer was built inside this tape, and thus
+    # the distribution init ops were inside this tape.
+    grads = tape.gradient(loss, variables)
+    for grad in grads:
+      self.assertIsNotNone(grad)
+
+    # Imagine this is the 2nd epoch.
+    with tf.GradientTape() as tape:
+      cell(inputs[:, 0, :], state)  # build won't be called again
+      loss = tf.reduce_sum([tf.reduce_sum(l) for l in cell.losses])
+
+    variables = [
+        cell.kernel_initializer.mean, cell.kernel_initializer.stddev,
+        cell.recurrent_initializer.mean, cell.recurrent_initializer.stddev,
+    ]
+    for v in variables:
+      self.assertIn(v, cell.variables)
+
+    # This would fail, since the layer was built inside the tape from the 1st
+    # epoch, and thus the distribution init ops were inside that tape instead of
+    # this tape. By using a callable for the variable, this will no longer fail.
+    grads = tape.gradient(loss, variables)
+    for grad in grads:
+      self.assertIsNotNone(grad)
+
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testLSTMCellReparameterizationModel(self):
     batch_size, timesteps, dim = 5, 3, 12

From 1a5ceca577e503d4e4de64c5311bcd46244dee43 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 9 Nov 2018 22:58:02 -0800
Subject: [PATCH 1218/2720] Add action-conditioning to the latent in video
 models and stop pixel sampling (to keep the default config for atari).

PiperOrigin-RevId: 220910042
---
 tensor2tensor/models/video/basic_deterministic.py      |  8 ++++----
 .../models/video/basic_deterministic_params.py         |  2 +-
 tensor2tensor/models/video/basic_stochastic.py         | 10 ++++++++--
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 440b8220f..a460bac33 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -41,8 +41,8 @@ class NextFrameBasicDeterministic(base.NextFrameBase):
   def is_recurrent_model(self):
     return False
 
-  def inject_latent(self, layer, inputs, target):
-    del inputs, target
+  def inject_latent(self, layer, inputs, target, action):
+    del inputs, target, action
     return layer, 0.0
 
   def middle_network(self, layer, internal_states):
@@ -68,6 +68,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
     hparams = self.hparams
     filters = hparams.hidden_size
     kernel2 = (4, 4)
+    action = actions[-1]
 
     # Embed the inputs.
     stacked_frames = tf.concat(frames, axis=-1)
@@ -94,12 +95,11 @@ def next_frame(self, frames, actions, rewards, target_frame,
 
     # Add embedded action if present.
     if self.has_actions:
-      action = actions[-1]
       x = common_video.inject_additional_input(
           x, action, "action_enc", hparams.action_injection)
 
     # Inject latent if present. Only for stochastic models.
-    x, extra_loss = self.inject_latent(x, frames, target_frame)
+    x, extra_loss = self.inject_latent(x, frames, target_frame, action)
 
     x_mid = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
     x, internal_states = self.middle_network(x, internal_states)
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 100342f6d..8e9644cbb 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -45,7 +45,7 @@ def next_frame_basic_deterministic():
   hparams.add_hparam("residual_dropout", 0.5)
   hparams.add_hparam("num_compress_steps", 6)
   hparams.add_hparam("filter_double_steps", 2)
-  hparams.add_hparam("pixel_sampling_temperature", 0.5)
+  hparams.add_hparam("pixel_sampling_temperature", 0.0)
   return hparams
 
 
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index b33b252e3..63af4214b 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -41,8 +41,9 @@ class NextFrameBasicStochastic(
     base_vae.NextFrameBaseVae):
   """Stochastic version of basic next-frame model."""
 
-  def inject_latent(self, layer, inputs, target):
+  def inject_latent(self, layer, inputs, target, action):
     """Inject a VAE-style latent."""
+    del action
     # Latent for stochastic model
     filters = 128
     full_video = tf.stack(inputs + [target], axis=1)
@@ -65,7 +66,7 @@ class NextFrameBasicStochasticDiscrete(
     basic_deterministic.NextFrameBasicDeterministic):
   """Basic next-frame model with a tiny discrete latent."""
 
-  def inject_latent(self, layer, inputs, target):
+  def inject_latent(self, layer, inputs, target, action):
     """Inject a deterministic latent based on the target frame."""
     hparams = self.hparams
     final_filters = common_layers.shape_list(layer)[-1]
@@ -100,6 +101,11 @@ def add_bits(layer, bits):
         bias_initializer=tf.random_normal_initializer(stddev=0.01))
     x = common_attention.add_timing_signal_nd(x)
 
+    # Add embedded action if present.
+    if action is not None:
+      x = common_video.inject_additional_input(
+          x, action, "action_enc_latent", hparams.action_injection)
+
     if hparams.full_latent_tower:
       for i in range(hparams.num_compress_steps):
         with tf.variable_scope("latent_downstride%d" % i):

From b5dea9e25b5b03e41c86b82a4a8dd18b41e68b34 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Sat, 10 Nov 2018 12:31:39 +0100
Subject: [PATCH 1219/2720] Use PolicyLearner in trainer_model_{based,free}
 (#1208)

* Adapt PolicyLearner interface

* Keep track of completed iterations in PolicyLearner

* Remove unnecessary PPO hparams and pass them as arguments instead

* Clean up hparams

* Revert rhp change ppo_epoch_length to simulated_rollout_length

* Move PPO learner to a separate module and add module policy_learners
exporting all learners in a dict

* Use PolicyLearner in trainer_model_free

* rl_trainer_lib_test -> trainer_model_free_test

* Pylint

* Back to ppo_epochs_num

* Remove trainer_utils

* Update trainer_model_free

* Move rl_trainer_lib and collect code to ppo_learner

* Remove unused TF wrappers

* Remove gym_utils
---
 tensor2tensor/data_generators/gym_utils.py    | 302 ------------------
 tensor2tensor/models/research/rl.py           |  59 ++--
 tensor2tensor/rl/envs/tf_atari_wrappers.py    | 277 ----------------
 tensor2tensor/rl/policy_learner.py            |  55 +---
 tensor2tensor/rl/ppo.py                       |   6 +-
 .../rl/{collect.py => ppo_learner.py}         | 222 +++++++++++--
 tensor2tensor/rl/rl_trainer_lib.py            | 148 ---------
 tensor2tensor/rl/trainer_model_based.py       | 193 +++++------
 .../rl/trainer_model_based_params.py          |  45 ++-
 tensor2tensor/rl/trainer_model_free.py        |  35 +-
 tensor2tensor/rl/trainer_model_free_test.py   |  10 +-
 11 files changed, 370 insertions(+), 982 deletions(-)
 delete mode 100644 tensor2tensor/data_generators/gym_utils.py
 rename tensor2tensor/rl/{collect.py => ppo_learner.py} (57%)
 delete mode 100644 tensor2tensor/rl/rl_trainer_lib.py

diff --git a/tensor2tensor/data_generators/gym_utils.py b/tensor2tensor/data_generators/gym_utils.py
deleted file mode 100644
index 18c3c0279..000000000
--- a/tensor2tensor/data_generators/gym_utils.py
+++ /dev/null
@@ -1,302 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Utilities for openai gym."""
-
-from collections import deque
-import gym
-
-import numpy as np
-
-
-# pylint: disable=method-hidden
-class WarmupWrapper(gym.Wrapper):
-  """Warmup wrapper."""
-
-  def __init__(self, env, warm_up_examples=0, warmup_action=0):
-    gym.Wrapper.__init__(self, env)
-    self.warm_up_examples = warm_up_examples
-    self.warm_up_action = warmup_action
-    self.observation_space = gym.spaces.Box(
-        low=0, high=255, shape=(210, 160, 3), dtype=np.uint8)
-
-  def get_starting_data(self, num_frames):
-    self.reset()
-    starting_observations, starting_actions, starting_rewards = [], [], []
-    for _ in range(num_frames):
-      observation, rew, _, _ = self.env.step(self.warm_up_action)
-      starting_observations.append(observation)
-      starting_rewards.append(rew)
-      starting_actions.append(self.warm_up_action)
-
-    return starting_observations, starting_actions, starting_rewards
-
-  def step(self, action):
-    return self.env.step(action)
-
-  def reset(self, **kwargs):
-    del kwargs
-    self.env.reset()
-    observation = None
-    for _ in range(self.warm_up_examples):
-      observation, _, _, _ = self.env.step(self.warm_up_action)
-
-    return observation
-
-
-class PongWrapper(WarmupWrapper):
-  """Pong Wrapper."""
-
-  def __init__(self, env, warm_up_examples=0,
-               action_space_reduction=False,
-               reward_skip_steps=0,
-               big_ball=False):
-    super(PongWrapper, self).__init__(env, warm_up_examples=warm_up_examples)
-    self.action_space_reduction = action_space_reduction
-    if self.action_space_reduction:
-      self.action_space = gym.spaces.Discrete(2)
-    self.warm_up_examples = warm_up_examples
-    self.observation_space = gym.spaces.Box(
-        low=0, high=255, shape=(210, 160, 3), dtype=np.uint8)
-    self.reward_skip_steps = reward_skip_steps
-    self.big_ball = big_ball
-
-  def step(self, action):
-    if self.action_space_reduction:
-      action = 2 if int(action) == 0 else 5
-    ob, rew, done, info = self.env.step(action)
-    ob = self.process_observation(ob)
-    if rew != 0 and self.reward_skip_steps != 0:
-      for _ in range(self.reward_skip_steps):
-        self.env.step(0)
-    return ob, rew, done, info
-
-  def reset(self, **kwargs):
-    observation = super(PongWrapper, self).reset(**kwargs)
-    observation = self.process_observation(observation)
-    return observation
-
-  def process_observation(self, obs):
-    if self.big_ball:
-      pos = PongWrapper.find_ball(obs)
-      if pos is not None:
-        x, y = pos
-        obs[x-5:x+5, y-5:y+5, :] = 255
-
-    return obs
-
-  @staticmethod
-  def find_ball(obs, default=None):
-    ball_area = obs[37:193, :, 0]
-    res = np.argwhere(ball_area == 236)
-    if not res:
-      return default
-    else:
-      x, y = res[0]
-      x += 37
-      return x, y
-
-
-def wrapped_pong_factory(warm_up_examples=0, action_space_reduction=False,
-                         reward_skip_steps=0, big_ball=False):
-  """Wrapped pong games."""
-  env = gym.make("PongDeterministic-v4")
-  env = env.env  # Remove time_limit wrapper.
-  env = PongWrapper(env, warm_up_examples=warm_up_examples,
-                    action_space_reduction=action_space_reduction,
-                    reward_skip_steps=reward_skip_steps,
-                    big_ball=big_ball)
-  return env
-
-
-gym.envs.register(id="T2TPongWarmUp20RewSkip200Steps-v1",
-                  entry_point=lambda: wrapped_pong_factory(  # pylint: disable=g-long-lambda
-                      warm_up_examples=20, reward_skip_steps=15),
-                  max_episode_steps=200)
-
-
-gym.envs.register(id="T2TPongWarmUp20RewSkipFull-v1",
-                  entry_point=lambda: wrapped_pong_factory(  # pylint: disable=g-long-lambda
-                      warm_up_examples=20, reward_skip_steps=15))
-
-
-class BreakoutWrapper(WarmupWrapper):
-  """Breakout Wrapper."""
-
-  FIRE_ACTION = 1
-
-  def __init__(self, env, warm_up_examples=0,
-               ball_down_skip=0,
-               big_ball=False,
-               include_direction_info=False,
-               reward_clipping=True):
-    super(BreakoutWrapper, self).__init__(
-        env, warm_up_examples=warm_up_examples,
-        warmup_action=BreakoutWrapper.FIRE_ACTION)
-    self.warm_up_examples = warm_up_examples
-    self.observation_space = gym.spaces.Box(low=0, high=255,
-                                            shape=(210, 160, 3),
-                                            dtype=np.uint8)
-    self.ball_down_skip = ball_down_skip
-    self.big_ball = big_ball
-    self.reward_clipping = reward_clipping
-    self.include_direction_info = include_direction_info
-    self.direction_info = deque([], maxlen=2)
-    self.points_gained = False
-    msg = ("ball_down_skip should be bigger equal 9 for "
-           "include_direction_info to work correctly")
-    assert not self.include_direction_info or ball_down_skip >= 9, msg
-
-  def step(self, action):
-    ob, rew, done, info = self.env.step(action)
-
-    if BreakoutWrapper.find_ball(ob) is None and self.ball_down_skip != 0:
-      for _ in range(self.ball_down_skip):
-        # We assume that nothing interesting happens during ball_down_skip
-        # and discard all information.
-        # We fire all the time to start new game
-        ob, _, _, _ = self.env.step(BreakoutWrapper.FIRE_ACTION)
-        self.direction_info.append(BreakoutWrapper.find_ball(ob))
-
-    ob = self.process_observation(ob)
-
-    self.points_gained = self.points_gained or rew > 0
-
-    if self.reward_clipping:
-      rew = np.sign(rew)
-
-    return ob, rew, done, info
-
-  def reset(self, **kwargs):
-    observation = super(BreakoutWrapper, self).reset(**kwargs)
-    self.env.step(BreakoutWrapper.FIRE_ACTION)
-    self.direction_info = deque([], maxlen=2)
-    observation = self.process_observation(observation)
-    return observation
-
-  @staticmethod
-  def find_ball(ob, default=None):
-    off_x = 63
-    clipped_ob = ob[off_x:-21, :, 0]
-    pos = np.argwhere(clipped_ob == 200)
-
-    if not pos.size:
-      return default
-
-    x = off_x + pos[0][0]
-    y = 0 + pos[0][1]
-    return x, y
-
-  def process_observation(self, obs):
-    if self.big_ball:
-      pos = BreakoutWrapper.find_ball(obs)
-      if pos is not None:
-        x, y = pos
-        obs[x-5:x+5, y-5:y+5, :] = 255
-
-    if self.include_direction_info:
-      for point in list(self.direction_info):
-        if point is not None:
-          x, y = point
-          obs[x-2:x+2, y-2:y+2, 1] = 255
-
-    return obs
-
-
-def wrapped_breakout_factory(warm_up_examples=0,
-                             ball_down_skip=0,
-                             big_ball=False,
-                             include_direction_info=False,
-                             reward_clipping=True):
-  """Wrapped breakout games."""
-  env = gym.make("BreakoutDeterministic-v4")
-  env = env.env  # Remove time_limit wrapper.
-  env = BreakoutWrapper(env, warm_up_examples=warm_up_examples,
-                        ball_down_skip=ball_down_skip,
-                        big_ball=big_ball,
-                        include_direction_info=include_direction_info,
-                        reward_clipping=reward_clipping)
-  return env
-
-
-gym.envs.register(id="T2TBreakoutWarmUp20RewSkip500Steps-v1",
-                  entry_point=lambda: wrapped_breakout_factory(  # pylint: disable=g-long-lambda
-                      warm_up_examples=1,
-                      ball_down_skip=9,
-                      big_ball=False,
-                      include_direction_info=True,
-                      reward_clipping=True
-                  ),
-                  max_episode_steps=500)
-
-
-class FreewayWrapper(WarmupWrapper):
-  """Wrapper for Freeway."""
-
-  def __init__(self, env,
-               warm_up_examples=0,
-               reward_clipping=True,
-               easy_freeway=False):
-    super(FreewayWrapper, self).__init__(env, warm_up_examples)
-    self.easy_freeway = easy_freeway
-    self.half_way_reward = 1.0
-
-    # this is probably not needed, just in case
-    self.reward_clipping = reward_clipping
-
-  def chicken_height(self, image):
-    raise NotImplementedError()
-
-  def step(self, action):
-    ob, rew, done, info = self.env.step(action)
-
-    if self.easy_freeway:
-      if rew > 0:
-        self.half_way_reward = 1
-      chicken_height = self.chicken_height(ob)
-      if chicken_height < 105:
-        rew += self.half_way_reward
-        self.half_way_reward = 0
-
-    if self.reward_clipping:
-      rew = np.sign(rew)
-
-    return ob, rew, done, info
-
-  def reset(self, **kwargs):
-    self.half_way_reward = 1.0
-    observation = super(FreewayWrapper, self).reset(**kwargs)
-    return observation
-
-
-def wrapped_freeway_factory(warm_up_examples=0,
-                            reward_clipping=True,
-                            easy_freeway=False):
-  """Wrapped freeway games."""
-  env = gym.make("FreewayDeterministic-v4")
-  env = env.env  # Remove time_limit wrapper.
-  env = FreewayWrapper(env, warm_up_examples=warm_up_examples,
-                       reward_clipping=reward_clipping,
-                       easy_freeway=easy_freeway)
-
-  return env
-
-gym.envs.register(id="T2TFreewayWarmUp20RewSkip500Steps-v1",
-                  entry_point=lambda: wrapped_freeway_factory(  # pylint: disable=g-long-lambda
-                      warm_up_examples=1,
-                      reward_clipping=True,
-                      easy_freeway=False
-                  ),
-                  max_episode_steps=500)
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index ba03da8e9..64339f3d3 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -42,7 +42,6 @@ def ppo_base_v1():
   hparams.add_hparam("init_logstd", 0.1)
   hparams.add_hparam("policy_layers", (100, 100))
   hparams.add_hparam("value_layers", (100, 100))
-  hparams.add_hparam("num_agents", 30)
   hparams.add_hparam("clipping_coef", 0.2)
   hparams.add_hparam("gae_gamma", 0.99)
   hparams.add_hparam("gae_lambda", 0.95)
@@ -52,16 +51,13 @@ def ppo_base_v1():
   hparams.add_hparam("epoch_length", 200)
   hparams.add_hparam("epochs_num", 2000)
   hparams.add_hparam("eval_every_epochs", 10)
-  hparams.add_hparam("num_eval_agents", 3)
-  hparams.add_hparam("video_during_eval", False)
   hparams.add_hparam("save_models_every_epochs", 30)
   hparams.add_hparam("optimization_batch_size", 50)
   hparams.add_hparam("max_gradients_norm", 0.5)
-  hparams.add_hparam("simulation_random_starts", False)
-  hparams.add_hparam("simulation_flip_first_random_for_beginning", False)
   hparams.add_hparam("intrinsic_reward_scale", 0.)
   hparams.add_hparam("logits_clip", 4.0)
   hparams.add_hparam("dropout_ppo", 0.1)
+  hparams.add_hparam("effective_num_agents", None)
   return hparams
 
 
@@ -98,7 +94,6 @@ def ppo_atari_base():
   """Pong base parameters."""
   hparams = ppo_discrete_action_base()
   hparams.learning_rate = 1e-4
-  hparams.num_agents = 8
   hparams.epoch_length = 200
   hparams.gae_gamma = 0.985
   hparams.gae_lambda = 0.985
@@ -106,7 +101,6 @@ def ppo_atari_base():
   hparams.value_loss_coef = 1
   hparams.optimization_epochs = 3
   hparams.epochs_num = 1000
-  hparams.num_eval_agents = 1
   hparams.policy_network = feed_forward_cnn_small_categorical_fun
   hparams.clipping_coef = 0.2
   hparams.optimization_batch_size = 20
@@ -129,7 +123,6 @@ def ppo_original_params():
   # The parameters below are modified to accommodate short epoch_length (which
   # is needed for model based rollouts).
   hparams.epoch_length = 50
-  hparams.num_agents = 16
   hparams.optimization_batch_size = 20
   return hparams
 
@@ -188,27 +181,16 @@ def ppo_pong_ae_base():
 @registry.register_hparams
 def pong_model_free():
   """TODO(piotrmilos): Document this."""
-  hparams = tf.contrib.training.HParams(
-      epochs_num=4,
-      eval_every_epochs=2,
-      num_agents=2,
-      optimization_epochs=3,
-      epoch_length=30,
-      entropy_loss_coef=0.003,
-      learning_rate=8e-05,
-      optimizer="Adam",
-      policy_network=feed_forward_cnn_small_categorical_fun,
-      gae_lambda=0.985,
-      num_eval_agents=2,
-      max_gradients_norm=0.5,
-      gae_gamma=0.985,
-      optimization_batch_size=4,
-      clipping_coef=0.2,
-      value_loss_coef=1,
-      save_models_every_epochs=False,
-      frame_stack_size=4,
-      force_beginning_resets=False,
-  )
+  hparams = mfrl_base()
+  hparams.batch_size = 2
+  hparams.ppo_eval_every_epochs = 2
+  hparams.ppo_epochs_num = 4
+  hparams.add_hparam("ppo_optimization_epochs", 3)
+  hparams.add_hparam("ppo_epoch_length", 30)
+  hparams.add_hparam("ppo_learning_rate", 8e-05)
+  hparams.add_hparam("ppo_optimizer", "Adam")
+  hparams.add_hparam("ppo_optimization_batch_size", 4)
+  hparams.add_hparam("ppo_save_models_every_epochs", 1000000)
   env = gym_env.T2TGymEnv("PongNoFrameskip-v4", batch_size=2)
   env.start_new_epoch(0)
   hparams.add_hparam("env_fn", make_real_env_fn(env))
@@ -218,12 +200,23 @@ def pong_model_free():
   return hparams
 
 
+@registry.register_hparams
+def mfrl_original():
+  return tf.contrib.training.HParams(
+      game="",
+      base_algo="ppo",
+      base_algo_params="ppo_original_params",
+      batch_size=16,
+      eval_batch_size=2,
+      frame_stack_size=4,
+  )
+
+
 @registry.register_hparams
 def mfrl_base():
-  hparams = ppo_original_params()
-  hparams.add_hparam("game", "")
-  hparams.epochs_num = 3000
-  hparams.eval_every_epochs = 100
+  hparams = mfrl_original()
+  hparams.add_hparam("ppo_epochs_num", 3000)
+  hparams.add_hparam("ppo_eval_every_epochs", 100)
   return hparams
 
 
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index f87701060..d4d72623d 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -19,14 +19,9 @@
 from __future__ import division
 from __future__ import print_function
 
-import math
-
 from six.moves import range  # pylint: disable=redefined-builtin
 
-from tensor2tensor.layers import discretization
-from tensor2tensor.models.research import autoencoders
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
-from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -93,70 +88,6 @@ def history_observations(self):
     )
 
 
-class RewardClippingWrapper(WrapperBase):
-  """ Reward clipping wrapper.
-      The rewards are clipped to -1, 0, 1
-      This is a common strategy to ensure learning stability
-      of rl algorithms
-  """
-
-  def __str__(self):
-    return "RewardClippingWrapper(%s)" % str(self._batch_env)
-
-  def simulate(self, action):
-    reward, done = self._batch_env.simulate(action)
-    with tf.control_dependencies([reward, done]):
-      return tf.sign(reward), tf.identity(done)
-
-  def _reset_non_empty(self, indices):
-    return self._batch_env._reset_non_empty(indices)  # pylint: disable=protected-access
-
-
-class MaxAndSkipWrapper(WrapperBase):
-  """ Max and skip wrapper.
-      The wrapper works under assumptions that issuing an action
-      to an environment with done=True has not effect.
-  """
-
-  def __init__(self, batch_env, skip=4):
-    super(MaxAndSkipWrapper, self).__init__(batch_env)
-    self.skip = skip
-    observs_shape = batch_env.observ.shape
-    self._observ = tf.Variable(tf.zeros(observs_shape, self.observ_dtype),
-                               trainable=False)
-
-  def __str__(self):
-    return "MaxAndSkipWrapper(%s)" % str(self._batch_env)
-
-  def simulate(self, action):
-    with tf.name_scope("environment/simulate"):  # Do we need this?
-      initializer = (tf.zeros_like(self._observ),
-                     tf.fill((len(self),), 0.0), tf.fill((len(self),), False))
-
-      def not_done_step(a, _):
-        reward, done = self._batch_env.simulate(action)
-        with tf.control_dependencies([reward, done]):
-          # TODO(piotrmilos): possibly ignore envs with done
-          r0 = tf.maximum(a[0], self._batch_env.observ)
-          r1 = tf.add(a[1], reward)
-          r2 = tf.logical_or(a[2], done)
-
-          return (r0, r1, r2)
-
-      simulate_ret = tf.scan(not_done_step, tf.range(self.skip),
-                             initializer=initializer, parallel_iterations=1,
-                             infer_shape=False)
-      simulate_ret = [ret[-1, ...] for ret in simulate_ret]
-
-      with tf.control_dependencies([self._observ.assign(simulate_ret[0])]):
-        return tf.identity(simulate_ret[1]), tf.identity(simulate_ret[2])
-
-  def _transform_history_observations(self, frames):
-    # Should be implemented if ever MaxAndSkipWrapper and StackWrapper are to
-    # be used together.
-    raise NotImplementedError
-
-
 class StackWrapper(WrapperBase):
   """ A wrapper which stacks previously seen frames. """
 
@@ -217,211 +148,3 @@ def _reset_non_empty(self, indices):
   def _transform_history_observations(self, frames):
     # Should be implemented if ever two StackWrappers are to be used together.
     raise NotImplementedError
-
-
-class AutoencoderWrapper(WrapperBase):
-  """Transforms the observations taking the bottleneck of an autoencoder."""
-
-  def __init__(self, batch_env, ae_hparams_set):
-    super(AutoencoderWrapper, self).__init__(batch_env)
-    self.ae_hparams_set = ae_hparams_set
-    self._observ = tf.Variable(
-        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
-        trainable=False)
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      autoencoder_hparams = registry.hparams(self.ae_hparams_set)
-      problem = registry.problem("dummy_autoencoder_problem")
-      autoencoder_hparams.problem_hparams = problem.get_hparams(
-          autoencoder_hparams)
-      autoencoder_hparams.problem = problem
-      self.autoencoder_model = autoencoders.AutoencoderOrderedDiscrete(
-          autoencoder_hparams, tf.estimator.ModeKeys.EVAL)
-
-  def __str__(self):
-    return "AutoencoderWrapper(%s)" % str(self._batch_env)
-
-  @property
-  def observ_shape(self):
-    height, width, _ = self._batch_env.observ_shape
-    ae_height = int(math.ceil(height / self.autoencoder_factor))
-    ae_width = int(math.ceil(width / self.autoencoder_factor))
-    ae_channels = 24  # TODO(piotrmilos): make it better
-    return (ae_height, ae_width, ae_channels)
-
-  @property
-  def autoencoder_factor(self):
-    """By how much to divide sizes when using autoencoders."""
-    hparams = registry.hparams(self.ae_hparams_set)
-    return 2**hparams.num_hidden_layers
-
-  def simulate(self, action):
-    reward, done = self._batch_env.simulate(action)
-    with tf.control_dependencies([reward, done]):
-      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-        observ = tf.cast(self._batch_env.observ, tf.int64)
-        ret = self.autoencoder_model.encode(observ)
-        ret = tf.cast(ret, self.observ_dtype)
-        assign_op = self._observ.assign(ret)
-        with tf.control_dependencies([assign_op]):
-          return tf.identity(reward), tf.identity(done)
-
-  def _reset_non_empty(self, indices):
-    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-      new_values = self._batch_env._reset_non_empty(indices)  # pylint: disable=protected-access
-      new_values = tf.cast(new_values, tf.int64)
-      ret = self.autoencoder_model.encode(new_values)
-      ret = tf.cast(ret, self.observ_dtype)
-      assign_op = tf.scatter_update(self._observ, indices, ret)
-      with tf.control_dependencies([assign_op]):
-        return tf.gather(self.observ, indices)
-
-  def _transform_history_observations(self, frames):
-    batch_size, history_size = frames.get_shape().as_list()[:2]
-    new_frames = tf.reshape(frames, (-1,) + self._batch_env.observ_shape)
-    new_frames = tf.cast(new_frames, tf.int64)
-    new_frames = self.autoencoder_model.encode(new_frames)
-    new_frames = tf.cast(new_frames, self.observ_dtype)
-    return new_frames.reshape((batch_size, history_size) + self.observ_shape)
-
-
-class ResizeWrapper(WrapperBase):
-  """Resizes the observations."""
-
-  def __init__(self, batch_env, height_factor=1, width_factor=1,
-               grayscale=False):
-    super(ResizeWrapper, self).__init__(batch_env)
-    self._height_factor = height_factor  # How much to resize on x axis.
-    self._width_factor = width_factor  # How much to resize on y axis.
-    self._do_grayscale = grayscale  # Whether to convert to grayscale.
-    self._is_identity = ((height_factor == 1) and (width_factor == 1)
-                         and not grayscale)
-    if not self._is_identity:
-      self._observ = tf.Variable(
-          tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
-          trainable=False)
-    else:
-      self._observ = self._batch_env.observ
-
-  def __str__(self):
-    return "ResizeWrapperh%dw%dg%d(%s)" % (
-        self._height_factor, self._width_factor, int(self._do_grayscale),
-        str(self._batch_env))
-
-  def _resize(self, tensor):
-    if self._is_identity:
-      return tensor
-    height, width, _ = self.observ_shape
-    observ = tf.to_float(tensor)
-    resized = tf.image.resize_images(
-        observ, [height, width], tf.image.ResizeMethod.AREA)
-    if self._do_grayscale:
-      resized = tf.image.rgb_to_grayscale(resized)
-    return tf.cast(resized, self.observ_dtype)
-
-  @property
-  def observ_shape(self):
-    height, width, channels = self._batch_env.observ_shape
-    if self._do_grayscale:
-      channels = 1
-    resized_height = height // self._height_factor
-    resized_width = width // self._width_factor
-    return (resized_height, resized_width, channels)
-
-  def simulate(self, action):
-    if self._is_identity:
-      return self._batch_env.simulate(action)
-    reward, done = self._batch_env.simulate(action)
-    with tf.control_dependencies([reward, done]):
-      ret = self._resize(self._batch_env.observ)
-      assign_op = self._observ.assign(ret)
-      with tf.control_dependencies([assign_op]):
-        return tf.identity(reward), tf.identity(done)
-
-  def _reset_non_empty(self, indices):
-    new_values = self._batch_env._reset_non_empty(indices)  # pylint: disable=protected-access
-    if self._is_identity:
-      return new_values
-    ret = self._resize(new_values)
-    assign_op = tf.scatter_update(self._observ, indices, ret)
-    with tf.control_dependencies([assign_op]):
-      return tf.gather(self.observ, indices)
-
-  def _transform_history_observations(self, frames):
-    batch_size, history_size = frames.get_shape().as_list()[:2]
-    new_frames = tf.reshape(frames, (-1,) + self._batch_env.observ_shape)
-    new_frames = self._resize(new_frames)
-    return new_frames.reshape((batch_size, history_size) + self.observ_shape)
-
-
-class IntToBitWrapper(WrapperBase):
-  """Unpacks the observations from integer values to bit values"""
-
-  def __init__(self, batch_env):
-    super(IntToBitWrapper, self).__init__(batch_env)
-    self._observ = self._observ = tf.Variable(
-        tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
-        trainable=False)
-
-  def __str__(self):
-    return "IntToBitWrapper(%s)" % str(self._batch_env)
-
-  @property
-  def observ_shape(self):
-    height, width, channels = self._batch_env.observ_shape
-    # We treat each channel as 8-bit integer to be expanded to 8 channels
-    return (height, width, channels*8)
-
-  def simulate(self, action):
-    reward, done = self._batch_env.simulate(action)
-    with tf.control_dependencies([reward, done]):
-      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-        unpacked = discretization.int_to_bit(self._batch_env.observ, 8)
-        unpacked = tf.reshape(unpacked, (-1,)+self.observ_shape)
-        unpacked = tf.cast(unpacked, self.observ_dtype)
-        assign_op = self._observ.assign(unpacked)
-        with tf.control_dependencies([assign_op]):
-          return tf.identity(reward), tf.identity(done)
-
-  def _reset_non_empty(self, indices):
-    # pylint: disable=protected-access
-    new_values = self._batch_env._reset_non_empty(indices)
-    new_values_unpacked = discretization.int_to_bit(new_values, 8)
-    new_values_unpacked = tf.reshape(new_values_unpacked, (-1,)
-                                     +self.observ_shape)
-    new_values_unpacked = tf.cast(new_values_unpacked, self.observ_dtype)
-    # pylint: enable=protected-access
-    assign_op = tf.scatter_update(self._observ, indices, new_values_unpacked)
-    with tf.control_dependencies([assign_op]):
-      return tf.identity(new_values_unpacked)
-
-  def _transform_history_observations(self, frames):
-    batch_size, history_size = frames.get_shape().as_list()[:2]
-    new_frames = discretization.int_to_bit(frames, 8)
-    new_frames = tf.reshape(
-        new_frames, (batch_size, history_size) + self.observ_shape
-    )
-    return tf.cast(new_frames, self.observ_dtype)
-
-
-class PyFuncWrapper(WrapperBase):
-  """Calls arbitrary python function on passing data"""
-
-  def __init__(self, batch_env, process_fun):
-    super(PyFuncWrapper, self).__init__(batch_env)
-    self.process_fun = process_fun
-    observs_shape = batch_env.observ.shape
-    self._observ = tf.Variable(tf.zeros(observs_shape, self.observ_dtype),
-                               trainable=False)
-
-  def __str__(self):
-    return "PyFuncWrapper(%s)" % str(self._batch_env)
-
-  def simulate(self, action):
-    reward, done = self._batch_env.simulate(action)
-    with tf.control_dependencies([reward, done]):
-      inputs = [self._observ.read_value(), reward, done, action]
-      ret = tf.py_func(self.process_fun, inputs, tf.double)
-    with tf.control_dependencies([ret]):
-      assign = self._observ.assign(self._batch_env.observ)
-    with tf.control_dependencies([assign]):
-      return tf.identity(reward), tf.identity(done)
diff --git a/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
index 60569dd5a..9ef2c2da3 100644
--- a/tensor2tensor/rl/policy_learner.py
+++ b/tensor2tensor/rl/policy_learner.py
@@ -19,63 +19,22 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.rl import rl_trainer_lib
-
 
 class PolicyLearner(object):
   """API for policy learners."""
 
-  def __init__(self, frame_stack_size, event_dir, agent_model_dir):
+  def __init__(self, frame_stack_size, base_event_dir, agent_model_dir):
     self.frame_stack_size = frame_stack_size
-    self.event_dir = event_dir
+    self.base_event_dir = base_event_dir
     self.agent_model_dir = agent_model_dir
 
-  def train(self, env_fn, hparams, target_num_epochs, simulated, epoch):
-    # TODO(konradczechowski): target_num_steps instead of epochs
+  def train(
+      self, env_fn, hparams, simulated, save_continuously, epoch,
+      num_env_steps=None, env_step_multiplier=1, eval_env_fn=None,
+      report_fn=None
+  ):
     # TODO(konradczechowski): move 'simulated' to  batch_env
     raise NotImplementedError()
 
   def evaluate(self, env_fn, hparams, stochastic):
     raise NotImplementedError()
-
-
-class PPOLearner(PolicyLearner):
-  """PPO for policy learning."""
-
-  def train(self, env_fn, hparams, target_num_epochs, simulated, epoch):
-    hparams.set_hparam("epochs_num", target_num_epochs)
-
-    if simulated:
-      simulated_str = "sim"
-      hparams.save_models_every_epochs = 10
-    else:
-      # TODO(konradczechowski): refactor ppo
-      assert hparams.num_agents == 1
-      # We do not save model, as that resets frames that we need at restarts.
-      # But we need to save at the last step, so we set it very high.
-      hparams.save_models_every_epochs = 1000000
-      simulated_str = "real"
-
-    # TODO(konradczechowski) refactor ppo, pass these as arguments
-    # (not inside hparams). Do the same in evaluate()
-    hparams.add_hparam("force_beginning_resets", simulated)
-    hparams.add_hparam("env_fn", env_fn)
-    hparams.add_hparam("frame_stack_size", self.frame_stack_size)
-    name_scope = "ppo_{}{}".format(simulated_str, epoch + 1)
-
-    rl_trainer_lib.train(hparams, self.event_dir + simulated_str,
-                         self.agent_model_dir, name_scope=name_scope)
-
-  def evaluate(self, env_fn, hparams, stochastic):
-    if stochastic:
-      policy_to_actions_lambda = lambda policy: policy.sample()
-    else:
-      policy_to_actions_lambda = lambda policy: policy.mode()
-    hparams.add_hparam(
-        "policy_to_actions_lambda", policy_to_actions_lambda
-    )
-    hparams.add_hparam("force_beginning_resets", False)
-    hparams.add_hparam("env_fn", env_fn)
-    hparams.add_hparam("frame_stack_size", self.frame_stack_size)
-
-    rl_trainer_lib.evaluate(hparams, self.agent_model_dir)
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index 0a9ab0861..00d4c37e9 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -73,7 +73,7 @@ def define_ppo_step(data_points, optimizer, hparams, action_space):
     return [tf.identity(x) for x in losses + gradients_norms]
 
 
-def define_ppo_epoch(memory, hparams, action_space):
+def define_ppo_epoch(memory, hparams, action_space, batch_size):
   """PPO epoch."""
   observation, reward, done, action, old_pdf, value = memory
 
@@ -102,8 +102,8 @@ def define_ppo_epoch(memory, hparams, action_space):
   number_of_batches = ((hparams.epoch_length-1) * hparams.optimization_epochs
                        / hparams.optimization_batch_size)
 
-  if hasattr(hparams, "effective_num_agents"):
-    number_of_batches *= hparams.num_agents
+  if hparams.effective_num_agents is not None:
+    number_of_batches *= batch_size
     number_of_batches /= hparams.effective_num_agents
 
   dataset = tf.data.Dataset.from_tensor_slices(
diff --git a/tensor2tensor/rl/collect.py b/tensor2tensor/rl/ppo_learner.py
similarity index 57%
rename from tensor2tensor/rl/collect.py
rename to tensor2tensor/rl/ppo_learner.py
index db63e7432..5f225936e 100644
--- a/tensor2tensor/rl/collect.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -13,19 +13,198 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Collect trajectories from interactions of agent with environment."""
+"""PPO learner."""
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import math
+import os
 
 from tensor2tensor.models.research.rl import get_policy
+from tensor2tensor.rl import ppo
 from tensor2tensor.rl.envs.tf_atari_wrappers import StackWrapper
 from tensor2tensor.rl.envs.tf_atari_wrappers import WrapperBase
+from tensor2tensor.rl.policy_learner import PolicyLearner
+from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
 
 
+class PPOLearner(PolicyLearner):
+  """PPO for policy learning."""
+
+  def __init__(self, *args, **kwargs):
+    super(PPOLearner, self).__init__(*args, **kwargs)
+    self._num_completed_iterations = 0
+
+  def train(
+      self, env_fn, hparams, simulated, save_continuously, epoch,
+      num_env_steps=None, env_step_multiplier=1, eval_env_fn=None,
+      report_fn=None
+  ):
+    if not save_continuously:
+      # We do not save model, as that resets frames that we need at restarts.
+      # But we need to save at the last step, so we set it very high.
+      hparams.save_models_every_epochs = 1000000
+
+    if simulated:
+      simulated_str = "sim"
+    else:
+      simulated_str = "real"
+    name_scope = "ppo_{}{}".format(simulated_str, epoch + 1)
+    event_dir = os.path.join(
+        self.base_event_dir, "ppo_summaries", str(epoch) + simulated_str
+    )
+
+    with tf.Graph().as_default():
+      with tf.name_scope(name_scope):
+        with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+          env = env_fn(in_graph=True)
+          (train_summary_op, eval_summary_op, initializers) = (
+              _define_train(
+                  env, hparams, eval_env_fn,
+                  frame_stack_size=self.frame_stack_size,
+                  force_beginning_resets=simulated
+              )
+          )
+
+        if num_env_steps is None:
+          iteration_increment = hparams.epochs_num
+        else:
+          iteration_increment = int(math.ceil(
+              num_env_steps / (env.batch_size * hparams.epoch_length)
+          ))
+        iteration_increment *= env_step_multiplier
+
+        self._num_completed_iterations += iteration_increment
+        _run_train(
+            hparams, event_dir, self.agent_model_dir,
+            self._num_completed_iterations, train_summary_op, eval_summary_op,
+            initializers, report_fn=report_fn
+        )
+
+  def evaluate(self, env_fn, hparams, stochastic):
+    if stochastic:
+      policy_to_actions_lambda = lambda policy: policy.sample()
+    else:
+      policy_to_actions_lambda = lambda policy: policy.mode()
+
+    with tf.Graph().as_default():
+      with tf.name_scope("rl_eval"):
+        eval_env = env_fn(in_graph=True)
+        (collect_memory, _, collect_init) = _define_collect(
+            eval_env, hparams, "ppo_eval", eval_phase=True,
+            frame_stack_size=self.frame_stack_size,
+            force_beginning_resets=False,
+            policy_to_actions_lambda=policy_to_actions_lambda
+        )
+        model_saver = tf.train.Saver(
+            tf.global_variables(".*network_parameters.*")
+        )
+
+        with tf.Session() as sess:
+          sess.run(tf.global_variables_initializer())
+          collect_init(sess)
+          trainer_lib.restore_checkpoint(
+              self.agent_model_dir, model_saver, sess
+          )
+          sess.run(collect_memory)
+
+
+def _define_train(train_env, ppo_hparams, eval_env_fn=None, **collect_kwargs):
+  """Define the training setup."""
+  memory, collect_summary, train_initialization = (
+      _define_collect(
+          train_env, ppo_hparams, "ppo_train", eval_phase=False,
+          policy_to_actions_lambda=(lambda policy: policy.sample()),
+          **collect_kwargs
+      )
+  )
+  ppo_summary = ppo.define_ppo_epoch(
+      memory, ppo_hparams, train_env.action_space, train_env.batch_size
+  )
+  train_summary = tf.summary.merge([collect_summary, ppo_summary])
+
+  if ppo_hparams.eval_every_epochs:
+    assert eval_env_fn is not None
+    eval_env = eval_env_fn(in_graph=True)
+    (_, eval_collect_summary, eval_initialization) = (
+        _define_collect(
+            eval_env, ppo_hparams, "ppo_eval", eval_phase=True,
+            policy_to_actions_lambda=(lambda policy: policy.mode()),
+            **collect_kwargs
+        )
+    )
+    return (
+        train_summary, eval_collect_summary,
+        (train_initialization, eval_initialization)
+    )
+  else:
+    return (train_summary, None, (train_initialization,))
+
+
+def _run_train(
+    ppo_hparams, event_dir, model_dir, num_target_iterations,
+    train_summary_op, eval_summary_op, initializers, report_fn=None
+):
+  """Train."""
+  summary_writer = tf.summary.FileWriter(
+      event_dir, graph=tf.get_default_graph(), flush_secs=60)
+
+  model_saver = tf.train.Saver(
+      tf.global_variables(".*network_parameters.*"))
+
+  with tf.Session() as sess:
+    sess.run(tf.global_variables_initializer())
+    for initializer in initializers:
+      initializer(sess)
+    num_completed_iterations = trainer_lib.restore_checkpoint(
+        model_dir, model_saver, sess)
+
+    # Fail-friendly, complete only unfinished epoch
+    num_iterations_to_go = num_target_iterations - num_completed_iterations
+
+    if num_iterations_to_go <= 0:
+      tf.logging.info(
+          "Skipping PPO training. Requested %d iterations while %d train "
+          "iterations already reached", num_target_iterations,
+          num_completed_iterations
+      )
+      return
+
+    for epoch_index in range(num_iterations_to_go):
+      summary = sess.run(train_summary_op)
+      if summary_writer:
+        summary_writer.add_summary(summary, epoch_index)
+
+      if (ppo_hparams.eval_every_epochs and
+          epoch_index % ppo_hparams.eval_every_epochs == 0):
+        eval_summary = sess.run(eval_summary_op)
+        if summary_writer:
+          summary_writer.add_summary(eval_summary, epoch_index)
+        if report_fn:
+          summary_proto = tf.Summary()
+          summary_proto.ParseFromString(eval_summary)
+          for elem in summary_proto.value:
+            if "mean_score" in elem.tag:
+              report_fn(elem.simple_value, epoch_index)
+              break
+
+      epoch_index_and_start = epoch_index + num_completed_iterations
+      if (model_saver and ppo_hparams.save_models_every_epochs and
+          (epoch_index_and_start %
+           ppo_hparams.save_models_every_epochs == 0 or
+           (epoch_index + 1) == num_iterations_to_go)):
+        ckpt_path = os.path.join(
+            model_dir, "model.ckpt-{}".format(
+                epoch_index + 1 + num_completed_iterations
+            )
+        )
+        model_saver.save(sess, ckpt_path)
+
+
 def _rollout_metadata(batch_env):
   """Metadata for rollouts."""
   batch_env_shape = batch_env.observ.get_shape().as_list()
@@ -82,12 +261,15 @@ def simulate(self, action):
       return tf.identity(reward), tf.identity(done)
 
 
-def define_collect(batch_env, hparams, scope):
+def _define_collect(
+    batch_env, ppo_hparams, scope, frame_stack_size, eval_phase,
+    policy_to_actions_lambda, force_beginning_resets
+):
   """Collect trajectories.
 
   Args:
     batch_env: Batch environment.
-    hparams: HParams.
+    ppo_hparams: PPO hparams, defined in tensor2tensor.models.research.rl.
     scope: var scope.
 
   Returns:
@@ -95,14 +277,15 @@ def define_collect(batch_env, hparams, scope):
     pdfs, values_functions)
     containing a rollout of environment from nested wrapped structure.
   """
+  epoch_length = ppo_hparams.epoch_length
 
   to_initialize = []
   with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
-    num_agents = hparams.num_agents
+    num_agents = batch_env.batch_size
 
     to_initialize.append(batch_env)
     wrappers = [
-        (StackWrapper, {"history": hparams.frame_stack_size}),
+        (StackWrapper, {"history": frame_stack_size}),
         (_MemoryWrapper, {})
     ]
     rollout_metadata = None
@@ -121,8 +304,8 @@ def initialization_lambda(sess):
         batch_env.initialize(sess)
 
     memory = [
-        tf.get_variable("collect_memory_%d_%s" % (hparams.epoch_length, name),
-                        shape=[hparams.epoch_length] + shape,
+        tf.get_variable("collect_memory_%d_%s" % (epoch_length, name),
+                        shape=[epoch_length] + shape,
                         dtype=dtype,
                         initializer=tf.zeros_initializer(),
                         trainable=False)
@@ -131,13 +314,11 @@ def initialization_lambda(sess):
     cumulative_rewards = tf.get_variable("cumulative_rewards", len(batch_env),
                                          trainable=False)
 
-    eval_phase_t = tf.convert_to_tensor(hparams.eval_phase)
+    eval_phase_t = tf.convert_to_tensor(eval_phase)
     should_reset_var = tf.Variable(True, trainable=False)
     zeros_tensor = tf.zeros(len(batch_env))
 
-  force_beginning_resets = tf.convert_to_tensor(
-      hparams.force_beginning_resets
-  )
+  force_beginning_resets = tf.convert_to_tensor(force_beginning_resets)
 
   def reset_ops_group():
     return tf.group(batch_env.reset(tf.range(len(batch_env))),
@@ -154,7 +335,7 @@ def reset_ops_group():
 
     def step(index, scores_sum, scores_num):
       """Single step."""
-      index %= hparams.epoch_length  # Only needed in eval runs.
+      index %= epoch_length  # Only needed in eval runs.
       # Note - the only way to ensure making a copy of tensor is to run simple
       # operation. We are waiting for tf.copy:
       # https://github.com/tensorflow/tensorflow/issues/11186
@@ -163,10 +344,10 @@ def step(index, scores_sum, scores_num):
       def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
         """Step of the environment."""
         actor_critic = get_policy(
-            tf.expand_dims(obs_copy, 0), hparams, batch_env.action_space
+            tf.expand_dims(obs_copy, 0), ppo_hparams, batch_env.action_space
         )
         policy = actor_critic.policy
-        action = hparams.policy_to_actions_lambda(policy)
+        action = policy_to_actions_lambda(policy)
 
         postprocessed_action = actor_critic.action_postprocessing(action)
         reward, done = batch_env.simulate(postprocessed_action[0, ...])
@@ -224,7 +405,7 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
     def stop_condition(i, _, resets):
       return tf.cond(eval_phase_t,
                      lambda: resets < num_agents,
-                     lambda: i < hparams.epoch_length)
+                     lambda: i < epoch_length)
 
     init = [tf.constant(0), tf.constant(0.0), tf.constant(0)]
     index, scores_sum, scores_num = tf.while_loop(
@@ -255,14 +436,13 @@ def stop_condition(i, _, resets):
     # When generating real data together with PPO training we must use single
     # agent. For PPO to work we reshape the history, as if it was generated
     # by real_ppo_effective_num_agents.
-    if (getattr(hparams, "effective_num_agents", None) and
-        not hparams.eval_phase):
+    if ppo_hparams.effective_num_agents is not None and not eval_phase:
       new_memory = []
-      effective_num_agents = hparams.effective_num_agents
-      assert hparams.epoch_length % effective_num_agents == 0, (
-          "The rollout of hparams.epoch_length will be distributed amongst"
+      effective_num_agents = ppo_hparams.effective_num_agents
+      assert epoch_length % ppo_hparams.effective_num_agents == 0, (
+          "The rollout of ppo_hparams.epoch_length will be distributed amongst"
           "effective_num_agents of agents")
-      new_epoch_length = int(hparams.epoch_length / effective_num_agents)
+      new_epoch_length = int(epoch_length / effective_num_agents)
       for mem, info in zip(memory, rollout_metadata):
         shape, _, name = info
         new_shape = [effective_num_agents, new_epoch_length] + shape[1:]
diff --git a/tensor2tensor/rl/rl_trainer_lib.py b/tensor2tensor/rl/rl_trainer_lib.py
deleted file mode 100644
index 6303f059e..000000000
--- a/tensor2tensor/rl/rl_trainer_lib.py
+++ /dev/null
@@ -1,148 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Library for training of RL agent with PPO algorithm."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import copy
-import os
-
-from tensor2tensor import models  # pylint: disable=unused-import
-from tensor2tensor.models.research import rl  # pylint: disable=unused-import
-from tensor2tensor.rl import collect
-from tensor2tensor.rl import ppo
-from tensor2tensor.utils import trainer_lib
-
-import tensorflow as tf
-
-
-def define_train(hparams):
-  """Define the training setup."""
-  train_hparams = copy.copy(hparams)
-  train_hparams.add_hparam("eval_phase", False)
-  train_hparams.add_hparam(
-      "policy_to_actions_lambda", lambda policy: policy.sample()
-  )
-
-  with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-    train_env = hparams.env_fn(in_graph=True)
-    memory, collect_summary, train_initialization = (
-        collect.define_collect(train_env, train_hparams, "ppo_train")
-    )
-    ppo_summary = ppo.define_ppo_epoch(memory, hparams, train_env.action_space)
-    train_summary = tf.summary.merge([collect_summary, ppo_summary])
-
-    if hparams.eval_every_epochs:
-      eval_hparams = copy.copy(hparams)
-      eval_hparams.add_hparam("eval_phase", True)
-      eval_hparams.add_hparam(
-          "policy_to_actions_lambda", lambda policy: policy.mode()
-      )
-      eval_env = hparams.eval_env_fn(in_graph=True)
-      eval_hparams.num_agents = hparams.num_eval_agents
-
-      _, eval_collect_summary, eval_initialization = (
-          collect.define_collect(eval_env, eval_hparams, "ppo_eval")
-      )
-      return train_summary, eval_collect_summary, (train_initialization,
-                                                   eval_initialization)
-    else:
-      return train_summary, None, (train_initialization,)
-
-
-def train(hparams, event_dir=None, model_dir=None,
-          restore_agent=True, name_scope="rl_train", report_fn=None):
-  """Train."""
-  with tf.Graph().as_default():
-    with tf.name_scope(name_scope):
-      train_summary_op, eval_summary_op, initializers = define_train(hparams)
-      if event_dir:
-        summary_writer = tf.summary.FileWriter(
-            event_dir, graph=tf.get_default_graph(), flush_secs=60)
-      else:
-        summary_writer = None
-
-      if model_dir:
-        model_saver = tf.train.Saver(
-            tf.global_variables(".*network_parameters.*"))
-      else:
-        model_saver = None
-
-      with tf.Session() as sess:
-        sess.run(tf.global_variables_initializer())
-        for initializer in initializers:
-          initializer(sess)
-        start_step = 0
-        if model_saver and restore_agent:
-          start_step = trainer_lib.restore_checkpoint(
-              model_dir, model_saver, sess)
-
-        # Fail-friendly, complete only unfinished epoch
-        steps_to_go = hparams.epochs_num - start_step
-
-        if steps_to_go <= 0:
-          tf.logging.info("Skipping PPO training. Requested %d steps while "
-                          "%d train steps already reached",
-                          hparams.epochs_num, start_step)
-          return
-
-        for epoch_index in range(steps_to_go):
-          summary = sess.run(train_summary_op)
-          if summary_writer:
-            summary_writer.add_summary(summary, epoch_index)
-
-          if (hparams.eval_every_epochs and
-              epoch_index % hparams.eval_every_epochs == 0):
-            eval_summary = sess.run(eval_summary_op)
-            if summary_writer:
-              summary_writer.add_summary(eval_summary, epoch_index)
-            if report_fn:
-              summary_proto = tf.Summary()
-              summary_proto.ParseFromString(eval_summary)
-              for elem in summary_proto.value:
-                if "mean_score" in elem.tag:
-                  report_fn(elem.simple_value, epoch_index)
-                  break
-
-          epoch_index_and_start = epoch_index + start_step
-          if (model_saver and hparams.save_models_every_epochs and
-              (epoch_index_and_start % hparams.save_models_every_epochs == 0 or
-               (epoch_index + 1) == steps_to_go)):
-            ckpt_path = os.path.join(
-                model_dir, "model.ckpt-{}".format(epoch_index + 1 + start_step))
-            model_saver.save(sess, ckpt_path)
-
-
-def evaluate(hparams, model_dir, name_scope="rl_eval"):
-  """Evaluate."""
-  hparams = copy.copy(hparams)
-  hparams.add_hparam("eval_phase", True)
-  with tf.Graph().as_default():
-    with tf.name_scope(name_scope):
-      eval_env = hparams.env_fn(in_graph=True)
-      (collect_memory, _, collect_init) = collect.define_collect(
-          eval_env, hparams, "ppo_eval"
-      )
-      model_saver = tf.train.Saver(
-          tf.global_variables(".*network_parameters.*")
-      )
-
-      with tf.Session() as sess:
-        sess.run(tf.global_variables_initializer())
-        collect_init(sess)
-        trainer_lib.restore_checkpoint(model_dir, model_saver, sess)
-        sess.run(collect_memory)
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 6c298c02f..ef45eb105 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -26,7 +26,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import copy
 import datetime
 import math
 import os
@@ -42,7 +41,7 @@
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import trainer_model_based_params
-from tensor2tensor.rl.policy_learner import PPOLearner
+from tensor2tensor.rl.ppo_learner import PPOLearner
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -52,27 +51,23 @@
 FLAGS = flags.FLAGS
 
 
-LEARNERS = dict(
-    ppo=PPOLearner,
-)
+LEARNERS = {
+    "ppo": PPOLearner
+}
 
 
-def real_ppo_epoch_increment(hparams):
-  """PPO increment."""
-  assert hparams.real_ppo_epochs_num is 0, (
-      "Should be put to 0 to enforce better readability"
-  )
-  return int(math.ceil(
-      hparams.num_real_env_frames /
-      (hparams.epochs * hparams.real_ppo_epoch_length)
-  ))
+def update_hparams_from_hparams(target_hparams, source_hparams, prefix):
+  """Copy a subset of hparams to target_hparams."""
+  for (param_name, param_value) in six.iteritems(source_hparams.values()):
+    if param_name.startswith(prefix):
+      target_hparams.set_hparam(param_name[len(prefix):], param_value)
 
 
-def sim_ppo_epoch_increment(hparams, is_final_epoch):
-  increment = hparams.ppo_epochs_num
-  if is_final_epoch:
-    increment *= 2
-  return increment
+def real_env_step_increment(hparams):
+  """Real env step increment."""
+  return int(math.ceil(
+      hparams.num_real_env_frames / hparams.epochs
+  ))
 
 
 def world_model_step_increment(hparams, is_initial_epoch):
@@ -173,20 +168,10 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
   getattr(exp, schedule)()
 
 
-def _update_hparams_from_hparams(target_hparams, source_hparams, prefix):
-  """Copy a subset of hparams to target_hparams."""
-  for param_name in target_hparams.values().keys():
-    prefixed_param_name = prefix + param_name
-    if prefixed_param_name in source_hparams:
-      target_hparams.set_hparam(param_name,
-                                source_hparams.get(prefixed_param_name))
-
-
-def train_agent(real_env, agent_model_dir, event_dir, world_model_dir, data_dir,
-                hparams, completed_epochs_num, epoch=0, is_final_epoch=False):
+def train_agent(
+    real_env, learner, world_model_dir, hparams, epoch, is_final_epoch
+):
   """Train the PPO agent in the simulated environment."""
-  del data_dir
-
   frame_stack_size = hparams.frame_stack_size
   initial_frame_rollouts = real_env.current_epoch_rollouts(
       split=tf.contrib.learn.ModeKeys.TRAIN,
@@ -215,52 +200,41 @@ def initial_frame_chooser(batch_size):
         for initial_frame_stack in initial_frames
     ])
   env_fn = make_simulated_env_fn(
-      real_env, hparams, hparams.ppo_num_agents, initial_frame_chooser,
+      real_env, hparams, hparams.simulated_batch_size, initial_frame_chooser,
       world_model_dir
   )
   base_algo_str = hparams.base_algo
   train_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
 
-  _update_hparams_from_hparams(train_hparams, hparams, base_algo_str + "_")
-  completed_epochs_num += sim_ppo_epoch_increment(hparams, is_final_epoch)
-  learner = LEARNERS[base_algo_str](frame_stack_size, event_dir,
-                                    agent_model_dir)
-  learner.train(env_fn, train_hparams, completed_epochs_num,
-                simulated=True, epoch=epoch)
+  update_hparams_from_hparams(
+      train_hparams, hparams, base_algo_str + "_"
+  )
 
-  return completed_epochs_num
+  env_step_multiplier = 1 if not is_final_epoch else 2
+  learner.train(
+      env_fn, train_hparams, simulated=True, save_continuously=True,
+      epoch=epoch, env_step_multiplier=env_step_multiplier
+  )
 
 
-def train_agent_real_env(
-    env, agent_model_dir, event_dir, data_dir,
-    hparams, completed_epochs_num, epoch=0, is_final_epoch=False):
+def train_agent_real_env(env, learner, hparams, epoch):
   """Train the PPO agent in the real environment."""
-  del is_final_epoch, data_dir
-
   base_algo_str = hparams.base_algo
 
   train_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
-  _update_hparams_from_hparams(train_hparams, hparams,
-                               "real_" + base_algo_str + "_")
-
-  # TODO(konradczechowski): add effective_num_agents to ppo_atari_base etc.
-  # this requires refactoring ppo.
-  # This should be overridden.
-  train_hparams.add_hparam("effective_num_agents",
-                           hparams.real_ppo_effective_num_agents)
-
-  completed_epochs_num += real_ppo_epoch_increment(hparams)
+  update_hparams_from_hparams(
+      train_hparams, hparams, "real_" + base_algo_str + "_"
+  )
 
   env_fn = rl.make_real_env_fn(env)
-  learner = LEARNERS[base_algo_str](hparams.frame_stack_size, event_dir,
-                                    agent_model_dir)
-  learner.train(env_fn, train_hparams, completed_epochs_num,
-                simulated=False, epoch=epoch)
+  num_env_steps = real_env_step_increment(hparams)
+  learner.train(
+      env_fn, train_hparams, simulated=False, save_continuously=False,
+      epoch=epoch, num_env_steps=num_env_steps
+  )
   # Save unfinished rollouts to history.
   env.reset()
 
-  return completed_epochs_num
-
 
 def train_world_model(
     env, data_dir, output_dir, hparams, world_model_steps_num, epoch
@@ -288,7 +262,7 @@ def train_world_model(
   return world_model_steps_num
 
 
-def setup_env(hparams, batch_size):
+def setup_env(hparams, batch_size, max_num_noops):
   """Setup."""
   game_mode = "Deterministic-v4"
   camel_game_name = "".join(
@@ -302,27 +276,28 @@ def setup_env(hparams, batch_size):
                   resize_width_factor=hparams.resize_width_factor,
                   resize_height_factor=hparams.resize_height_factor,
                   base_env_timesteps_limit=hparams.env_timesteps_limit,
-                  max_num_noops=hparams.max_num_noops)
+                  max_num_noops=max_num_noops)
   return env
 
 
-def evaluate_single_config(hparams, agent_model_dir):
+def evaluate_single_config(hparams, stochastic, max_num_noops, agent_model_dir):
   """Evaluate the PPO agent in the real environment."""
   eval_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
-  eval_hparams.num_agents = hparams.num_agents
-  eval_hparams.add_hparam("stochastic", hparams.stochastic)
-  env = setup_env(hparams, batch_size=hparams.num_agents)
+  env = setup_env(
+      hparams, batch_size=hparams.eval_batch_size, max_num_noops=max_num_noops
+  )
   env.start_new_epoch(0)
   env_fn = rl.make_real_env_fn(env)
-  learner = LEARNERS[hparams.base_algo](hparams.frame_stack_size,
-                                        event_dir=None,
-                                        agent_model_dir=agent_model_dir)
-  learner.evaluate(env_fn, eval_hparams, eval_hparams.stochastic)
-  rollouts = env.current_epoch_rollouts()[:hparams.num_agents]
+  learner = LEARNERS[hparams.base_algo](
+      hparams.frame_stack_size, base_event_dir=None,
+      agent_model_dir=agent_model_dir
+  )
+  learner.evaluate(env_fn, eval_hparams, stochastic)
+  rollouts = env.current_epoch_rollouts()[:hparams.eval_batch_size]
   env.close()
 
-  assert len(rollouts) == hparams.num_agents, "{} {}".format(len(rollouts),
-                                                             hparams.num_agents)
+  assert len(rollouts) == hparams.eval_batch_size, \
+      "{} {}".format(len(rollouts), hparams.eval_batch_size)
   return tuple(
       compute_mean_reward(rollouts, clipped) for clipped in (True, False)
   )
@@ -335,20 +310,14 @@ def get_metric_name(stochastic, max_num_noops, clipped):
 
 def evaluate_all_configs(hparams, agent_model_dir):
   """Evaluate the agent with multiple eval configurations."""
-  def make_eval_hparams(hparams, stochastic, max_num_noops):
-    hparams = copy.copy(hparams)
-    hparams.add_hparam("num_agents", hparams.eval_num_agents)
-    hparams.add_hparam("stochastic", stochastic)
-    hparams.max_num_noops = max_num_noops
-    return hparams
-
   metrics = {}
   # Iterate over all combinations of picking actions by sampling/mode and
   # whether to do initial no-ops.
   for stochastic in (True, False):
     for max_num_noops in (hparams.eval_max_num_noops, 0):
-      eval_hparams = make_eval_hparams(hparams, stochastic, max_num_noops)
-      scores = evaluate_single_config(eval_hparams, agent_model_dir)
+      scores = evaluate_single_config(
+          hparams, stochastic, max_num_noops, agent_model_dir
+      )
       for (score, clipped) in zip(scores, (True, False)):
         metric_name = get_metric_name(stochastic, max_num_noops, clipped)
         metrics[metric_name] = score
@@ -388,7 +357,7 @@ def initial_frame_chooser(batch_size):
   )
   sim_env = env_fn(in_graph=False)
   subsequence_length = int(
-      max(hparams.wm_eval_rollout_ratios) * hparams.ppo_epoch_length
+      max(hparams.wm_eval_rollout_ratios) * hparams.simulated_rollout_length
   )
   rollouts = real_env.current_epoch_rollouts(
       split=tf.contrib.learn.ModeKeys.EVAL,
@@ -400,10 +369,10 @@ def initial_frame_chooser(batch_size):
   )
 
   reward_accuracies_by_length = {
-      int(ratio * hparams.ppo_epoch_length): []
+      int(ratio * hparams.simulated_rollout_length): []
       for ratio in hparams.wm_eval_rollout_ratios
   }
-  for _ in range(hparams.wm_eval_epochs_num):
+  for _ in range(hparams.wm_eval_num_batches):
     rollout_subsequences[:] = random_rollout_subsequences(
         rollouts, hparams.wm_eval_batch_size,
         subsequence_length + frame_stack_size
@@ -489,15 +458,23 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   # Directories
   subdirectories = [
       "data", "tmp", "world_model", ("world_model", "debug_videos"),
-      "ppo"
+      "policy"
   ]
   directories = setup_directories(output_dir, subdirectories)
 
   epoch = -1
   data_dir = directories["data"]
-  env = setup_env(hparams, batch_size=hparams.real_ppo_num_agents)
+  env = setup_env(
+      hparams, batch_size=hparams.real_batch_size,
+      max_num_noops=hparams.max_num_noops
+  )
   env.start_new_epoch(epoch, data_dir)
 
+  learner = LEARNERS[hparams.base_algo](
+      hparams.frame_stack_size, directories["world_model"],
+      directories["policy"]
+  )
+
   # Timing log function
   log_relative_time = make_relative_timing_fn()
 
@@ -505,15 +482,10 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   epoch_metrics = []
   metrics = {}
 
-  # Collect data from the real environment with PPO or random policy.
-  ppo_model_dir = directories["ppo"]
-  tf.logging.info("Initial training of PPO in real environment.")
-  ppo_event_dir = os.path.join(directories["world_model"],
-                               "ppo_summaries/initial")
-  completed_epochs_num = train_agent_real_env(
-      env, ppo_model_dir, ppo_event_dir, data_dir, hparams,
-      completed_epochs_num=0, epoch=epoch, is_final_epoch=False
-  )
+  # Collect data from the real environment.
+  policy_model_dir = directories["policy"]
+  tf.logging.info("Initial training of the policy in real environment.")
+  train_agent_real_env(env, learner, hparams, epoch)
   metrics["mean_reward/train/clipped"] = compute_mean_reward(
       env.current_epoch_rollouts(), clipped=True
   )
@@ -539,28 +511,17 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
         world_model_steps_num, epoch
     )
 
-    # Train PPO
-    log("Training PPO in simulated environment.")
-    ppo_event_dir = os.path.join(directories["world_model"],
-                                 "ppo_summaries", str(epoch))
-    ppo_model_dir = directories["ppo"]
-    if not hparams.ppo_continue_training:
-      ppo_model_dir = ppo_event_dir
-
-    completed_epochs_num = train_agent(
-        env, ppo_model_dir, ppo_event_dir,
-        directories["world_model"], data_dir, hparams, completed_epochs_num,
-        epoch=epoch, is_final_epoch=is_final_epoch
+    # Train agent
+    log("Training policy in simulated environment.")
+    train_agent(
+        env, learner, directories["world_model"], hparams, epoch, is_final_epoch
     )
 
     env.start_new_epoch(epoch, data_dir)
 
-    # Train PPO on real env (short)
-    log("Training PPO in real environment.")
-    completed_epochs_num = train_agent_real_env(
-        env, ppo_model_dir, ppo_event_dir, data_dir, hparams,
-        completed_epochs_num, epoch=epoch, is_final_epoch=is_final_epoch
-    )
+    # Train agent on real env (short)
+    log("Training policy in real environment.")
+    train_agent_real_env(env, learner, hparams, epoch)
 
     if hparams.stop_loop_early:
       return 0.0
@@ -570,7 +531,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     )
     log("Mean training reward: {}".format(metrics["mean_reward/train/clipped"]))
 
-    eval_metrics = evaluate_all_configs(hparams, ppo_model_dir)
+    eval_metrics = evaluate_all_configs(hparams, policy_model_dir)
     log("Agent eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
     metrics.update(eval_metrics)
 
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 510b15762..50ad899e1 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -61,26 +61,24 @@ def rlmb_base():
       autoencoder_hparams_set="autoencoder_discrete_pong",
       model_train_steps=15000,
       initial_epoch_train_steps_multiplier=3,
-      simulated_env_generator_num_steps=2000,
       simulation_random_starts=True,  # Use random starts in PPO.
       # Flip the first random frame in PPO batch for the true beginning.
       simulation_flip_first_random_for_beginning=True,
       intrinsic_reward_scale=0.,
+      # Number of real environments to train on simultaneously.
+      real_batch_size=1,
+      # Number of simulated environments to train on simultaneously.
+      simulated_batch_size=16,
+      # Number of frames that can be taken from the simulated environment before
+      # it diverges, used for training the agent.
+      simulated_rollout_length=50,
       ppo_epochs_num=1000,  # This should be enough to see something
-      # Our simulated envs do not know how to reset.
-      # You should set ppo_time_limit to the value you believe that
-      # the simulated env produces a reasonable output.
-      ppo_time_limit=200,  # TODO(blazej): this param is unused
-      # It makes sense to have ppo_time_limit=ppo_epoch_length,
-      # though it is not necessary.
+      # Should be equal to simulated_rollout_length.
+      # TODO(koz4k): Uncouple this by outputing done from SimulatedBatchEnv.
       ppo_epoch_length=50,
-      ppo_num_agents=16,
       # Do not eval since simulated batch env does not produce dones
       ppo_eval_every_epochs=0,
       ppo_learning_rate=1e-4,  # Will be changed, just so it exists.
-      # Whether the PPO agent should be restored from the previous iteration, or
-      # should start fresh each time.
-      ppo_continue_training=True,
       # Resizing.
       resize_height_factor=2,
       resize_width_factor=2,
@@ -95,16 +93,17 @@ def rlmb_base():
       # In your experiments, you want to optimize this rate to your schedule.
       learning_rate_bump=3.0,
 
+      # Unused; number of PPO epochs is calculated from the real frame limit.
       real_ppo_epochs_num=0,
       # This needs to be divisible by real_ppo_effective_num_agents.
       real_ppo_epoch_length=16*200,
-      real_ppo_num_agents=1,
       real_ppo_learning_rate=1e-4,
-      real_ppo_continue_training=True,
       real_ppo_effective_num_agents=16,
       real_ppo_eval_every_epochs=0,
 
-      eval_num_agents=30,
+      # Batch size during evaluation. Metrics are averaged over this number of
+      # rollouts.
+      eval_batch_size=30,
       eval_max_num_noops=8,
 
       game="pong",
@@ -114,7 +113,7 @@ def rlmb_base():
       # Number of concurrent rollouts in world model evaluation.
       wm_eval_batch_size=16,
       # Number of batches to run for world model evaluation.
-      wm_eval_epochs_num=8,
+      wm_eval_num_batches=8,
       # Ratios of ppo_epoch_length to report reward_accuracy on.
       wm_eval_rollout_ratios=[0.25, 0.5, 1, 2],
       stop_loop_early=False,  # To speed-up tests.
@@ -132,7 +131,6 @@ def rlmb_basetest():
   hparams.epochs = 2
   hparams.num_real_env_frames = 3200
   hparams.model_train_steps = 100
-  hparams.simulated_env_generator_num_steps = 20
   hparams.ppo_epochs_num = 2
   return hparams
 
@@ -391,24 +389,24 @@ def rlmb_tiny():
       tf.contrib.training.HParams(
           epochs=1,
           num_real_env_frames=128,
-          simulated_env_generator_num_steps=64,
           model_train_steps=2,
           ppo_epochs_num=2,
-          ppo_time_limit=5,
+          simulated_batch_size=2,
+          simulated_rollout_length=2,
           ppo_epoch_length=2,
-          ppo_num_agents=2,
+          real_batch_size=1,
           real_ppo_epoch_length=36,
-          real_ppo_num_agents=1,
-          real_ppo_epochs_num=0,
           real_ppo_effective_num_agents=2,
-          eval_num_agents=1,
+          max_num_noops=1,
+          eval_batch_size=1,
+          eval_max_num_noops=1,
           generative_model_params="next_frame_tiny",
           stop_loop_early=True,
           resize_height_factor=2,
           resize_width_factor=2,
           game="pong",
           wm_eval_rollout_ratios=[1],
-          env_timesteps_limit=6,
+          env_timesteps_limit=7,
       ).values())
 
 
@@ -465,7 +463,6 @@ def rlmb_ae_basetest():
   hparams.num_real_env_frames = 3200
   hparams.model_train_steps = 100
   hparams.autoencoder_train_steps = 10
-  hparams.simulated_env_generator_num_steps = 20
   hparams.ppo_epochs_num = 2
   return hparams
 
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 01fc2c0d8..96ffebff7 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -20,15 +20,18 @@
 python -m tensor2tensor.rl.trainer_model_free \
     --output_dir=$HOME/t2t/rl_v1 \
     --hparams_set=pong_model_free \
-    --loop_hparams='num_agents=15'
+    --loop_hparams='batch_size=15'
 """
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+import six
+
 from tensor2tensor.data_generators import gym_env
 from tensor2tensor.models.research import rl
-from tensor2tensor.rl import rl_trainer_lib
+from tensor2tensor.rl.ppo_learner import PPOLearner
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
 
@@ -45,16 +48,28 @@
   pass
 
 
+LEARNERS = {
+    "ppo": PPOLearner
+}
+
+
+def update_hparams_from_hparams(target_hparams, source_hparams, prefix):
+  """Copy a subset of hparams to target_hparams."""
+  for (param_name, param_value) in six.iteritems(source_hparams.values()):
+    if param_name.startswith(prefix):
+      target_hparams.set_hparam(param_name[len(prefix):], param_value)
+
+
 def initialize_env_specs(hparams):
   """Initializes env_specs using T2TGymEnvs."""
   if getattr(hparams, "game", None):
     game_name = gym_env.camel_case_name(hparams.game)
     env = gym_env.T2TGymEnv("{}Deterministic-v4".format(game_name),
-                            batch_size=hparams.num_agents)
+                            batch_size=hparams.batch_size)
     env.start_new_epoch(0)
     hparams.add_hparam("env_fn", rl.make_real_env_fn(env))
     eval_env = gym_env.T2TGymEnv("{}Deterministic-v4".format(game_name),
-                                 batch_size=hparams.num_eval_agents)
+                                 batch_size=hparams.eval_batch_size)
     eval_env.start_new_epoch(0)
     hparams.add_hparam("eval_env_fn", rl.make_real_env_fn(eval_env))
   return hparams
@@ -62,7 +77,17 @@ def initialize_env_specs(hparams):
 
 def train(hparams, output_dir, report_fn=None):
   hparams = initialize_env_specs(hparams)
-  rl_trainer_lib.train(hparams, output_dir, output_dir, report_fn=report_fn)
+  learner = LEARNERS[hparams.base_algo](
+      hparams.frame_stack_size, FLAGS.output_dir, output_dir
+  )
+  policy_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
+  update_hparams_from_hparams(
+      policy_hparams, hparams, hparams.base_algo + "_"
+  )
+  learner.train(
+      hparams.env_fn, policy_hparams, simulated=False, save_continuously=True,
+      epoch=0, eval_env_fn=hparams.eval_env_fn, report_fn=report_fn
+  )
 
 
 def main(_):
diff --git a/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index 01a4f8df9..bdf9371df 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -18,7 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.rl import rl_trainer_lib
+from tensor2tensor.rl import trainer_model_free
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -28,10 +28,10 @@ class TrainTest(tf.test.TestCase):
 
   def test_train_pong(self):
     hparams = registry.hparams("pong_model_free")
-    hparams.epochs_num = 2
-    hparams.num_agents = 2
-    hparams.epoch_length = 3
-    rl_trainer_lib.train(hparams)
+    hparams.batch_size = 2
+    hparams.ppo_epochs_num = 2
+    hparams.ppo_epoch_length = 3
+    trainer_model_free.train(hparams, tf.test.get_temp_dir())
 
 
 if __name__ == "__main__":

From 6e51d6d502b68f3c96cd117810a5f8e4fb65f5a9 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Sat, 10 Nov 2018 03:43:58 -0800
Subject: [PATCH 1220/2720] internal merge of PR #1208

PiperOrigin-RevId: 220925940
---
 tensor2tensor/rl/ppo_learner.py             | 222 ++++++++++----------
 tensor2tensor/rl/trainer_model_based.py     |   6 +-
 tensor2tensor/rl/trainer_model_free_test.py |   5 +-
 3 files changed, 120 insertions(+), 113 deletions(-)

diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 5f225936e..2f7b45b3d 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -39,11 +39,16 @@ def __init__(self, *args, **kwargs):
     super(PPOLearner, self).__init__(*args, **kwargs)
     self._num_completed_iterations = 0
 
-  def train(
-      self, env_fn, hparams, simulated, save_continuously, epoch,
-      num_env_steps=None, env_step_multiplier=1, eval_env_fn=None,
-      report_fn=None
-  ):
+  def train(self,
+            env_fn,
+            hparams,
+            simulated,
+            save_continuously,
+            epoch,
+            num_env_steps=None,
+            env_step_multiplier=1,
+            eval_env_fn=None,
+            report_fn=None):
     if not save_continuously:
       # We do not save model, as that resets frames that we need at restarts.
       # But we need to save at the last step, so we set it very high.
@@ -54,9 +59,8 @@ def train(
     else:
       simulated_str = "real"
     name_scope = "ppo_{}{}".format(simulated_str, epoch + 1)
-    event_dir = os.path.join(
-        self.base_event_dir, "ppo_summaries", str(epoch) + simulated_str
-    )
+    event_dir = os.path.join(self.base_event_dir, "ppo_summaries",
+                             str(epoch) + simulated_str)
 
     with tf.Graph().as_default():
       with tf.name_scope(name_scope):
@@ -64,26 +68,30 @@ def train(
           env = env_fn(in_graph=True)
           (train_summary_op, eval_summary_op, initializers) = (
               _define_train(
-                  env, hparams, eval_env_fn,
+                  env,
+                  hparams,
+                  eval_env_fn,
                   frame_stack_size=self.frame_stack_size,
-                  force_beginning_resets=simulated
-              )
-          )
+                  force_beginning_resets=simulated))
 
         if num_env_steps is None:
           iteration_increment = hparams.epochs_num
         else:
-          iteration_increment = int(math.ceil(
-              num_env_steps / (env.batch_size * hparams.epoch_length)
-          ))
+          iteration_increment = int(
+              math.ceil(
+                  num_env_steps / (env.batch_size * hparams.epoch_length)))
         iteration_increment *= env_step_multiplier
 
         self._num_completed_iterations += iteration_increment
         _run_train(
-            hparams, event_dir, self.agent_model_dir,
-            self._num_completed_iterations, train_summary_op, eval_summary_op,
-            initializers, report_fn=report_fn
-        )
+            hparams,
+            event_dir,
+            self.agent_model_dir,
+            self._num_completed_iterations,
+            train_summary_op,
+            eval_summary_op,
+            initializers,
+            report_fn=report_fn)
 
   def evaluate(self, env_fn, hparams, stochastic):
     if stochastic:
@@ -95,21 +103,21 @@ def evaluate(self, env_fn, hparams, stochastic):
       with tf.name_scope("rl_eval"):
         eval_env = env_fn(in_graph=True)
         (collect_memory, _, collect_init) = _define_collect(
-            eval_env, hparams, "ppo_eval", eval_phase=True,
+            eval_env,
+            hparams,
+            "ppo_eval",
+            eval_phase=True,
             frame_stack_size=self.frame_stack_size,
             force_beginning_resets=False,
-            policy_to_actions_lambda=policy_to_actions_lambda
-        )
+            policy_to_actions_lambda=policy_to_actions_lambda)
         model_saver = tf.train.Saver(
-            tf.global_variables(".*network_parameters.*")
-        )
+            tf.global_variables(".*network_parameters.*"))
 
         with tf.Session() as sess:
           sess.run(tf.global_variables_initializer())
           collect_init(sess)
-          trainer_lib.restore_checkpoint(
-              self.agent_model_dir, model_saver, sess
-          )
+          trainer_lib.restore_checkpoint(self.agent_model_dir, model_saver,
+                                         sess)
           sess.run(collect_memory)
 
 
@@ -117,14 +125,14 @@ def _define_train(train_env, ppo_hparams, eval_env_fn=None, **collect_kwargs):
   """Define the training setup."""
   memory, collect_summary, train_initialization = (
       _define_collect(
-          train_env, ppo_hparams, "ppo_train", eval_phase=False,
+          train_env,
+          ppo_hparams,
+          "ppo_train",
+          eval_phase=False,
           policy_to_actions_lambda=(lambda policy: policy.sample()),
-          **collect_kwargs
-      )
-  )
+          **collect_kwargs))
   ppo_summary = ppo.define_ppo_epoch(
-      memory, ppo_hparams, train_env.action_space, train_env.batch_size
-  )
+      memory, ppo_hparams, train_env.action_space, train_env.batch_size)
   train_summary = tf.summary.merge([collect_summary, ppo_summary])
 
   if ppo_hparams.eval_every_epochs:
@@ -132,29 +140,31 @@ def _define_train(train_env, ppo_hparams, eval_env_fn=None, **collect_kwargs):
     eval_env = eval_env_fn(in_graph=True)
     (_, eval_collect_summary, eval_initialization) = (
         _define_collect(
-            eval_env, ppo_hparams, "ppo_eval", eval_phase=True,
+            eval_env,
+            ppo_hparams,
+            "ppo_eval",
+            eval_phase=True,
             policy_to_actions_lambda=(lambda policy: policy.mode()),
-            **collect_kwargs
-        )
-    )
-    return (
-        train_summary, eval_collect_summary,
-        (train_initialization, eval_initialization)
-    )
+            **collect_kwargs))
+    return (train_summary, eval_collect_summary, (train_initialization,
+                                                  eval_initialization))
   else:
     return (train_summary, None, (train_initialization,))
 
 
-def _run_train(
-    ppo_hparams, event_dir, model_dir, num_target_iterations,
-    train_summary_op, eval_summary_op, initializers, report_fn=None
-):
+def _run_train(ppo_hparams,
+               event_dir,
+               model_dir,
+               num_target_iterations,
+               train_summary_op,
+               eval_summary_op,
+               initializers,
+               report_fn=None):
   """Train."""
   summary_writer = tf.summary.FileWriter(
       event_dir, graph=tf.get_default_graph(), flush_secs=60)
 
-  model_saver = tf.train.Saver(
-      tf.global_variables(".*network_parameters.*"))
+  model_saver = tf.train.Saver(tf.global_variables(".*network_parameters.*"))
 
   with tf.Session() as sess:
     sess.run(tf.global_variables_initializer())
@@ -170,8 +180,7 @@ def _run_train(
       tf.logging.info(
           "Skipping PPO training. Requested %d iterations while %d train "
           "iterations already reached", num_target_iterations,
-          num_completed_iterations
-      )
+          num_completed_iterations)
       return
 
     for epoch_index in range(num_iterations_to_go):
@@ -194,14 +203,11 @@ def _run_train(
 
       epoch_index_and_start = epoch_index + num_completed_iterations
       if (model_saver and ppo_hparams.save_models_every_epochs and
-          (epoch_index_and_start %
-           ppo_hparams.save_models_every_epochs == 0 or
+          (epoch_index_and_start % ppo_hparams.save_models_every_epochs == 0 or
            (epoch_index + 1) == num_iterations_to_go)):
         ckpt_path = os.path.join(
-            model_dir, "model.ckpt-{}".format(
-                epoch_index + 1 + num_completed_iterations
-            )
-        )
+            model_dir,
+            "model.ckpt-{}".format(epoch_index + 1 + num_completed_iterations))
         model_saver.save(sess, ckpt_path)
 
 
@@ -236,8 +242,8 @@ def __init__(self, batch_env):
     self.speculum = tf.FIFOQueue(infinity, shapes=shapes, dtypes=dtypes)
     observs_shape = batch_env.observ.shape
     # TODO(piotrmilos): possibly retrieve the observation type for batch_env
-    self._observ = tf.Variable(tf.zeros(observs_shape, self.observ_dtype),
-                               trainable=False)
+    self._observ = tf.Variable(
+        tf.zeros(observs_shape, self.observ_dtype), trainable=False)
 
   def __str__(self):
     return "MemoryWrapper(%s)" % str(self._batch_env)
@@ -261,16 +267,18 @@ def simulate(self, action):
       return tf.identity(reward), tf.identity(done)
 
 
-def _define_collect(
-    batch_env, ppo_hparams, scope, frame_stack_size, eval_phase,
-    policy_to_actions_lambda, force_beginning_resets
-):
+def _define_collect(batch_env, ppo_hparams, scope, frame_stack_size, eval_phase,
+                    policy_to_actions_lambda, force_beginning_resets):
   """Collect trajectories.
 
   Args:
     batch_env: Batch environment.
     ppo_hparams: PPO hparams, defined in tensor2tensor.models.research.rl.
     scope: var scope.
+    frame_stack_size: TODO(koz4k): Write docstring.
+    eval_phase: TODO(koz4k): Write docstring.
+    policy_to_actions_lambda: TODO(koz4k): Write docstring.
+    force_beginning_resets: TODO(koz4k): Write docstring.
 
   Returns:
     Returns memory (observations, rewards, dones, actions,
@@ -284,15 +292,14 @@ def _define_collect(
     num_agents = batch_env.batch_size
 
     to_initialize.append(batch_env)
-    wrappers = [
-        (StackWrapper, {"history": frame_stack_size}),
-        (_MemoryWrapper, {})
-    ]
+    wrappers = [(StackWrapper, {
+        "history": frame_stack_size
+    }), (_MemoryWrapper, {})]
     rollout_metadata = None
     speculum = None
     for w in wrappers:
-      tf.logging.info("Applying wrapper %s(%s) to env %s."
-                      % (str(w[0]), str(w[1]), str(batch_env)))
+      tf.logging.info("Applying wrapper %s(%s) to env %s." % (str(
+          w[0]), str(w[1]), str(batch_env)))
       batch_env = w[0](batch_env, **w[1])
       to_initialize.append(batch_env)
 
@@ -304,15 +311,16 @@ def initialization_lambda(sess):
         batch_env.initialize(sess)
 
     memory = [
-        tf.get_variable("collect_memory_%d_%s" % (epoch_length, name),
-                        shape=[epoch_length] + shape,
-                        dtype=dtype,
-                        initializer=tf.zeros_initializer(),
-                        trainable=False)
-        for (shape, dtype, name) in rollout_metadata]
+        tf.get_variable(
+            "collect_memory_%d_%s" % (epoch_length, name),
+            shape=[epoch_length] + shape,
+            dtype=dtype,
+            initializer=tf.zeros_initializer(),
+            trainable=False) for (shape, dtype, name) in rollout_metadata
+    ]
 
-    cumulative_rewards = tf.get_variable("cumulative_rewards", len(batch_env),
-                                         trainable=False)
+    cumulative_rewards = tf.get_variable(
+        "cumulative_rewards", len(batch_env), trainable=False)
 
     eval_phase_t = tf.convert_to_tensor(eval_phase)
     should_reset_var = tf.Variable(True, trainable=False)
@@ -321,8 +329,9 @@ def initialization_lambda(sess):
   force_beginning_resets = tf.convert_to_tensor(force_beginning_resets)
 
   def reset_ops_group():
-    return tf.group(batch_env.reset(tf.range(len(batch_env))),
-                    tf.assign(cumulative_rewards, zeros_tensor))
+    return tf.group(
+        batch_env.reset(tf.range(len(batch_env))),
+        tf.assign(cumulative_rewards, zeros_tensor))
 
   reset_op = tf.cond(
       tf.logical_or(should_reset_var.read_value(), force_beginning_resets),
@@ -344,8 +353,7 @@ def step(index, scores_sum, scores_num):
       def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
         """Step of the environment."""
         actor_critic = get_policy(
-            tf.expand_dims(obs_copy, 0), ppo_hparams, batch_env.action_space
-        )
+            tf.expand_dims(obs_copy, 0), ppo_hparams, batch_env.action_space)
         policy = actor_critic.policy
         action = policy_to_actions_lambda(policy)
 
@@ -379,10 +387,11 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
       with tf.control_dependencies([pdf, value_function]):
         obs, reward, done, action = speculum.dequeue()
 
-        to_save = [obs, reward, done, action,
-                   pdf, value_function]
-        save_ops = [tf.scatter_update(memory_slot, index, value)
-                    for memory_slot, value in zip(memory, to_save)]
+        to_save = [obs, reward, done, action, pdf, value_function]
+        save_ops = [
+            tf.scatter_update(memory_slot, index, value)
+            for memory_slot, value in zip(memory, to_save)
+        ]
         cumulate_rewards_op = cumulative_rewards.assign_add(reward)
 
         agent_indices_to_reset = tf.where(top_level_done)[:, 0]
@@ -391,29 +400,25 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
         scores_sum_delta = tf.reduce_sum(
             tf.gather(cumulative_rewards.read_value(), agent_indices_to_reset))
         scores_num_delta = tf.count_nonzero(done, dtype=tf.int32)
-      with tf.control_dependencies(save_ops + [scores_sum_delta,
-                                               scores_num_delta]):
+      with tf.control_dependencies(save_ops +
+                                   [scores_sum_delta, scores_num_delta]):
         reset_env_op = batch_env.reset(agent_indices_to_reset)
         reset_cumulative_rewards_op = tf.scatter_update(
             cumulative_rewards, agent_indices_to_reset,
             tf.gather(zeros_tensor, agent_indices_to_reset))
-      with tf.control_dependencies([reset_env_op,
-                                    reset_cumulative_rewards_op]):
-        return [index + 1, scores_sum + scores_sum_delta,
-                scores_num + scores_num_delta]
+      with tf.control_dependencies([reset_env_op, reset_cumulative_rewards_op]):
+        return [
+            index + 1, scores_sum + scores_sum_delta,
+            scores_num + scores_num_delta
+        ]
 
     def stop_condition(i, _, resets):
-      return tf.cond(eval_phase_t,
-                     lambda: resets < num_agents,
+      return tf.cond(eval_phase_t, lambda: resets < num_agents,
                      lambda: i < epoch_length)
 
     init = [tf.constant(0), tf.constant(0.0), tf.constant(0)]
     index, scores_sum, scores_num = tf.while_loop(
-        stop_condition,
-        step,
-        init,
-        parallel_iterations=1,
-        back_prop=False)
+        stop_condition, step, init, parallel_iterations=1, back_prop=False)
 
   # We handle force_beginning_resets differently. We assume that all envs are
   # reseted at the end of episod (though it happens at the beginning of the
@@ -427,9 +432,9 @@ def stop_condition(i, _, resets):
         lambda: scores_sum + tf.reduce_sum(cumulative_rewards.read_value()),
         lambda: scores_sum)
 
-  mean_score = tf.cond(tf.greater(scores_num, 0),
-                       lambda: scores_sum / tf.cast(scores_num, tf.float32),
-                       lambda: 0.)
+  mean_score = tf.cond(
+      tf.greater(scores_num, 0),
+      lambda: scores_sum / tf.cast(scores_num, tf.float32), lambda: 0.)
   printing = tf.Print(0, [mean_score, scores_sum, scores_num], "mean_score: ")
   with tf.control_dependencies([index, printing]):
     memory = [mem.read_value() for mem in memory]
@@ -446,22 +451,23 @@ def stop_condition(i, _, resets):
       for mem, info in zip(memory, rollout_metadata):
         shape, _, name = info
         new_shape = [effective_num_agents, new_epoch_length] + shape[1:]
-        perm = list(range(len(shape)+1))
+        perm = list(range(len(shape) + 1))
         perm[0] = 1
         perm[1] = 0
         mem = tf.transpose(mem, perm=perm)
         mem = tf.reshape(mem, shape=new_shape)
-        mem = tf.transpose(mem, perm=perm,
-                           name="collect_memory_%d_%s"
-                           % (new_epoch_length, name))
+        mem = tf.transpose(
+            mem,
+            perm=perm,
+            name="collect_memory_%d_%s" % (new_epoch_length, name))
         new_memory.append(mem)
       memory = new_memory
 
     mean_score_summary = tf.cond(
         tf.greater(scores_num, 0),
-        lambda: tf.summary.scalar("mean_score_this_iter", mean_score),
-        str)
-    summaries = tf.summary.merge(
-        [mean_score_summary,
-         tf.summary.scalar("episodes_finished_this_iter", scores_num)])
+        lambda: tf.summary.scalar("mean_score_this_iter", mean_score), str)
+    summaries = tf.summary.merge([
+        mean_score_summary,
+        tf.summary.scalar("episodes_finished_this_iter", scores_num)
+    ])
     return memory, summaries, initialization_lambda
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index ef45eb105..f9b498c70 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -136,8 +136,7 @@ def choose_subsequence():
 
 
 def make_simulated_env_fn(
-    real_env, hparams, batch_size, initial_frame_chooser, model_dir
-):
+    real_env, hparams, batch_size, initial_frame_chooser, model_dir):
   """Creates a simulated env_fn."""
   return rl.make_simulated_env_fn(
       reward_range=real_env.reward_range,
@@ -169,8 +168,7 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
 
 
 def train_agent(
-    real_env, learner, world_model_dir, hparams, epoch, is_final_epoch
-):
+    real_env, learner, world_model_dir, hparams, epoch, is_final_epoch):
   """Train the PPO agent in the simulated environment."""
   frame_stack_size = hparams.frame_stack_size
   initial_frame_rollouts = real_env.current_epoch_rollouts(
diff --git a/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index bdf9371df..bb1116016 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -23,6 +23,8 @@
 
 import tensorflow as tf
 
+FLAGS = tf.flags.FLAGS
+
 
 class TrainTest(tf.test.TestCase):
 
@@ -31,7 +33,8 @@ def test_train_pong(self):
     hparams.batch_size = 2
     hparams.ppo_epochs_num = 2
     hparams.ppo_epoch_length = 3
-    trainer_model_free.train(hparams, tf.test.get_temp_dir())
+    FLAGS.output_dir = tf.test.get_temp_dir()
+    trainer_model_free.train(hparams, FLAGS.output_dir)
 
 
 if __name__ == "__main__":

From 708625aae8de3fa599460bce85064a4f06a67522 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Sun, 11 Nov 2018 14:58:15 -0800
Subject: [PATCH 1221/2720] Allow optional property "avoid_overlapping_frames"
 to ensure that every frame appears only in one video.

PiperOrigin-RevId: 221014900
---
 tensor2tensor/data_generators/video_utils.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 2dd9571b2..19a5c3c9d 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -278,6 +278,11 @@ def dataset_splits(self):
   def only_keep_videos_from_0th_frame(self):
     return True
 
+  @property
+  def avoid_overlapping_frames(self):
+    """When True, each video has non overlapping frames with every other."""
+    return False
+
   @property
   def use_not_breaking_batching(self):
     return True
@@ -426,6 +431,9 @@ def check_integrity_and_batch(*datasets):
           if self.only_keep_videos_from_0th_frame:
             not_broken = tf.logical_and(not_broken, tf.equal(
                 frame_numbers[0], 0))
+          if self.avoid_overlapping_frames:
+            non_overlap = tf.equal(tf.mod(frame_numbers[0], num_frames), 0)
+            not_broken = tf.logical_and(not_broken, non_overlap)
         else:
           tf.logging.warning("use_not_breaking_batching is True but "
                              "no frame_number is in the dataset.")

From 72c47ae4ed8a345b865832d92f41498541254093 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Mon, 12 Nov 2018 21:02:59 +0100
Subject: [PATCH 1222/2720] Evaluation fixes (#1216)

* Move summaries to more appropriate directories

* Don't run evaluation again for already finished epochs after restart

* Count PPO epochs from the restored global step
---
 tensor2tensor/rl/ppo_learner.py         | 13 ++--
 tensor2tensor/rl/trainer_model_based.py | 93 ++++++++++++++++---------
 2 files changed, 64 insertions(+), 42 deletions(-)

diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 2f7b45b3d..33791e320 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -174,16 +174,14 @@ def _run_train(ppo_hparams,
         model_dir, model_saver, sess)
 
     # Fail-friendly, complete only unfinished epoch
-    num_iterations_to_go = num_target_iterations - num_completed_iterations
-
-    if num_iterations_to_go <= 0:
+    if num_target_iterations <= num_completed_iterations:
       tf.logging.info(
           "Skipping PPO training. Requested %d iterations while %d train "
           "iterations already reached", num_target_iterations,
           num_completed_iterations)
       return
 
-    for epoch_index in range(num_iterations_to_go):
+    for epoch_index in range(num_completed_iterations, num_target_iterations):
       summary = sess.run(train_summary_op)
       if summary_writer:
         summary_writer.add_summary(summary, epoch_index)
@@ -201,13 +199,12 @@ def _run_train(ppo_hparams,
               report_fn(elem.simple_value, epoch_index)
               break
 
-      epoch_index_and_start = epoch_index + num_completed_iterations
       if (model_saver and ppo_hparams.save_models_every_epochs and
-          (epoch_index_and_start % ppo_hparams.save_models_every_epochs == 0 or
-           (epoch_index + 1) == num_iterations_to_go)):
+          (epoch_index % ppo_hparams.save_models_every_epochs == 0 or
+           (epoch_index + 1) == num_target_iterations)):
         ckpt_path = os.path.join(
             model_dir,
-            "model.ckpt-{}".format(epoch_index + 1 + num_completed_iterations))
+            "model.ckpt-{}".format(epoch_index + 1))
         model_saver.save(sess, ckpt_path)
 
 
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index f9b498c70..1a26b5499 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -439,6 +439,21 @@ def append_debug_frame_batch(sim_obs, real_obs):
   }
 
 
+def load_metrics(event_dir, epoch):
+  """Loads metrics for this epoch if they have already been written.
+
+  This reads the entire event file but it's small with just per-epoch metrics.
+  """
+  metrics = {}
+  for filename in os.listdir(event_dir):
+    path = os.path.join(event_dir, filename)
+    for event in tf.train.summary_iterator(path):
+      if event.step == epoch and event.HasField("summary"):
+        value = event.summary.value[0]
+        metrics[value.tag] = value.simple_value
+  return metrics
+
+
 def summarize_metrics(eval_metrics_writer, metrics, epoch):
   """Write metrics to summary."""
   for (name, value) in six.iteritems(metrics):
@@ -456,7 +471,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   # Directories
   subdirectories = [
       "data", "tmp", "world_model", ("world_model", "debug_videos"),
-      "policy"
+      "policy", "eval_metrics"
   ]
   directories = setup_directories(output_dir, subdirectories)
 
@@ -469,7 +484,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   env.start_new_epoch(epoch, data_dir)
 
   learner = LEARNERS[hparams.base_algo](
-      hparams.frame_stack_size, directories["world_model"],
+      hparams.frame_stack_size, directories["policy"],
       directories["policy"]
   )
 
@@ -492,9 +507,9 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   ))
   env.generate_data(data_dir)
 
-  eval_metrics_event_dir = os.path.join(directories["world_model"],
-                                        "eval_metrics_event_dir")
-  eval_metrics_writer = tf.summary.FileWriter(eval_metrics_event_dir)
+  eval_metrics_writer = tf.summary.FileWriter(
+      directories["eval_metrics"]
+  )
 
   world_model_steps_num = 0
 
@@ -524,40 +539,50 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
     if hparams.stop_loop_early:
       return 0.0
 
-    metrics["mean_reward/train/clipped"] = compute_mean_reward(
-        env.current_epoch_rollouts(), clipped=True
-    )
-    log("Mean training reward: {}".format(metrics["mean_reward/train/clipped"]))
-
-    eval_metrics = evaluate_all_configs(hparams, policy_model_dir)
-    log("Agent eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
-    metrics.update(eval_metrics)
-
     env.generate_data(data_dir)
 
-    if hparams.eval_world_model:
-      debug_video_path = os.path.join(
-          directories["world_model", "debug_videos"],
-          "{}.avi".format(env.current_epoch)
-      )
-      wm_metrics = evaluate_world_model(
-          env, hparams, directories["world_model"], debug_video_path
+    metrics = load_metrics(directories["eval_metrics"], epoch)
+    if metrics:
+      # Skip eval if metrics have already been written for this epoch. Otherwise
+      # we'd overwrite them with wrong data.
+      log("Metrics found for this epoch, skipping evaluation.")
+    else:
+      metrics["mean_reward/train/clipped"] = compute_mean_reward(
+          env.current_epoch_rollouts(), clipped=True
       )
-      log("World model eval metrics:\n{}".format(pprint.pformat(wm_metrics)))
-      metrics.update(wm_metrics)
-
-    summarize_metrics(eval_metrics_writer, metrics, epoch)
+      log("Mean training reward: {}".format(
+          metrics["mean_reward/train/clipped"]
+      ))
+
+      eval_metrics = evaluate_all_configs(hparams, policy_model_dir)
+      log("Agent eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
+      metrics.update(eval_metrics)
+
+      if hparams.eval_world_model:
+        debug_video_path = os.path.join(
+            directories["world_model", "debug_videos"],
+            "{}.avi".format(env.current_epoch)
+        )
+        wm_metrics = evaluate_world_model(
+            env, hparams, directories["world_model"], debug_video_path
+        )
+        log("World model eval metrics:\n{}".format(pprint.pformat(wm_metrics)))
+        metrics.update(wm_metrics)
+
+      summarize_metrics(eval_metrics_writer, metrics, epoch)
+
+      # Report metrics
+      if report_fn:
+        if report_metric == "mean_reward":
+          metric_name = get_metric_name(
+              stochastic=False, max_num_noops=hparams.eval_max_num_noops,
+              clipped=False
+          )
+          report_fn(eval_metrics[metric_name], epoch)
+        else:
+          report_fn(eval_metrics[report_metric], epoch)
 
-    # Report metrics
     epoch_metrics.append(metrics)
-    if report_fn:
-      if report_metric == "mean_reward":
-        metric_name = get_metric_name(stochastic=False,
-                                      max_num_noops=hparams.eval_max_num_noops,
-                                      clipped=False)
-        report_fn(eval_metrics[metric_name], epoch)
-      else:
-        report_fn(eval_metrics[report_metric], epoch)
 
   # Return the evaluation metrics from the final epoch
   return epoch_metrics[-1]

From a07b99e63cd1ac0626c86994e15998415f18ef93 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Mon, 12 Nov 2018 12:21:07 -0800
Subject: [PATCH 1223/2720] internal merge of PR #1216

PiperOrigin-RevId: 221136178
---
 tensor2tensor/rl/trainer_model_based.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 1a26b5499..fe27b676d 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -443,6 +443,13 @@ def load_metrics(event_dir, epoch):
   """Loads metrics for this epoch if they have already been written.
 
   This reads the entire event file but it's small with just per-epoch metrics.
+
+  Args:
+    event_dir: TODO(koz4k): Document this.
+    epoch: TODO(koz4k): Document this.
+
+  Returns:
+    metrics.
   """
   metrics = {}
   for filename in os.listdir(event_dir):

From bbc8cdb60991ab0f7b57ec0e72cb5771e3a160dc Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 12 Nov 2018 17:55:39 -0800
Subject: [PATCH 1224/2720] Allow the unconditional next_frame_glow model to
 train on random frames from the video. Previously, the unconditional model
 used to support only "hparams.num_input_frames=1" and
 "hparams.num_target_frames=1" and train on the first frame in this instance.
 This is limiting in the scenario where there the dataset consists of a small
 number of long videos, (eg KTH) because datapoints become highly correltaed

PiperOrigin-RevId: 221190817
---
 tensor2tensor/models/research/glow_ops.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 6b59bd8b1..224635838 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -708,6 +708,7 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
   latent_dist_encoder = hparams.get("latent_dist_encoder", None)
   latent_skip = hparams.get("latent_skip", False)
   if latent_dist_encoder == "pointwise":
+    last_latent = latent
     merge_std = hparams.level_scale
     latent_shape = common_layers.shape_list(latent)
     z_shape = common_layers.shape_list(z)

From 60d55c58bbdca90b3328a7b537dfe4ca645a1034 Mon Sep 17 00:00:00 2001
From: RJ Ryan <rjryan@google.com>
Date: Mon, 12 Nov 2018 23:53:46 -0800
Subject: [PATCH 1225/2720] Cast scaling factors to the same type as the log
 determinant in invertible_1x1_conv.

PiperOrigin-RevId: 221218468
---
 tensor2tensor/models/research/glow_ops.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 224635838..6c87eab2b 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -276,7 +276,10 @@ def invertible_1x1_conv(name, x, reverse=False):
     u = u * np.transpose(l_mask) + tf.diag(sign_s * tf.exp(log_s))
     w = tf.matmul(p, tf.matmul(l, u))
 
-    objective = tf.reduce_sum(log_s) * height * width
+    # If height or width cannot be statically determined then they end up as
+    # tf.int32 tensors, which cannot be directly multiplied with a floating
+    # point tensor without a cast.
+    objective = tf.reduce_sum(log_s) * tf.cast(height * width, log_s.dtype)
     if not reverse:
       w = tf.reshape(w, [1, 1] + w_shape)
       x = tf.nn.conv2d(x, w, [1, 1, 1, 1], "SAME", data_format="NHWC")

From 9e50c6ce9e16c6bc679fec51ba160b3ed0e5a476 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Tue, 13 Nov 2018 13:17:48 -0800
Subject: [PATCH 1226/2720] Print non-trainable variables as well

PiperOrigin-RevId: 221321508
---
 tensor2tensor/utils/optimize.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index e2f9e45c9..e948d2778 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -34,9 +34,17 @@ def optimize(loss, learning_rate, hparams, use_tpu=False):
   """Minimize loss."""
   loss = weight_decay_and_noise(loss, hparams, learning_rate)
   loss = tf.identity(loss, name="total_loss")
+  # Print trainable variables.
   log_variable_sizes(verbose=hparams.summarize_vars)
+  # Print non-trainable variables.
+  non_trainable_variables = list(
+      set(tf.global_variables()) - set(tf.trainable_variables()))
+  log_variable_sizes(non_trainable_variables, tag="Non-trainable variables",
+                     verbose=hparams.summarize_vars)
   if hparams.summarize_vars:
     summarize_variables()
+    # Summarize non-trainable variables as well
+    summarize_variables(non_trainable_variables, tag="Non-trainable variables")
   diet_vars = [
       v for v in tf.global_variables() if v.dtype == dtypes.float16_ref
   ]

From d755e10a21c0ee1274b152d2b45adb0d433c4caa Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 13 Nov 2018 15:07:56 -0800
Subject: [PATCH 1227/2720] Add TemperedNormal a thin wrapper around
 distributions.Normal to allow sampling with a temperature T.

PiperOrigin-RevId: 221342039
---
 tensor2tensor/models/research/glow.py         | 13 ++++-
 tensor2tensor/models/research/glow_ops.py     | 52 ++++++++++++++-----
 .../models/research/glow_ops_test.py          | 24 +++++++++
 3 files changed, 74 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 3848e6a58..5b71da058 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -55,6 +55,7 @@ def glow_hparams():
   # initialization. A higher init_batch_size is required for training
   # stability especially when hparams.batch_size is low.
   hparams.add_hparam("init_batch_size", 256)
+  hparams.add_hparam("temperature", 1.0)
   return hparams
 
 
@@ -85,6 +86,12 @@ def preprocess(self, x):
     x = x / n_bins - 0.5
     return x
 
+  @property
+  def temperature(self):
+    if self.is_predicting:
+      return self.hparams.temperature
+    return 1.0
+
   def scale(self, x):
     """Scale x from -0.5 - 0.5 to 0 - 255."""
     x = tf.where(tf.is_nan(x), tf.ones_like(x), x)
@@ -110,7 +117,8 @@ def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
     # If eps=None, images are sampled from the prior.
     with arg_scope(ops, init=False), var_scope:
       predictions, _, _, _ = glow_ops.encoder_decoder(
-          "codec", self.z_sample, self.hparams, eps=None, reverse=True)
+          "codec", self.z_sample, self.hparams, eps=None, reverse=True,
+          temperature=self.temperature)
 
     return self.scale(predictions)
 
@@ -140,7 +148,8 @@ def top_prior(self):
       dist: instance of tfp.distributions.Normal, prior distribution.
     """
     return glow_ops.top_prior(
-        "top_prior", self.z_top_shape, learn_prior=self.hparams.top_prior)
+        "top_prior", self.z_top_shape, learn_prior=self.hparams.top_prior,
+        temperature=self.temperature)
 
   def body(self, features):
     if self.is_training:
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 6c87eab2b..e6cd25807 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -19,7 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from functools import partial
+import functools
 import numpy as np
 import scipy
 from tensor2tensor.layers import common_layers
@@ -31,6 +31,24 @@
 add_arg_scope = tf.contrib.framework.add_arg_scope
 
 
+class TemperedNormal(tfp.distributions.Normal):
+  """Normal distribution with temperature T."""
+
+  def __init__(self, loc, scale, temperature=1.0):
+    self.temperature = temperature
+    new_scale = scale * self.temperature
+    tfp.distributions.Normal.__init__(self, loc=loc, scale=new_scale)
+
+  def sample(self, sample_shape=(), seed=None, name="sample"):
+    if self.temperature == 0.0:
+      if not sample_shape:
+        return self.loc
+      loc = tf.expand_dims(self.loc, axis=0)
+      return tf.tile(loc, (sample_shape[0], 1, 1))
+    return super(TemperedNormal, self).sample(
+        sample_shape=sample_shape, seed=seed, name=name)
+
+
 def default_initializer(std=0.05):
   return tf.random_normal_initializer(0., std)
 
@@ -763,7 +781,8 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
 
 
 @add_arg_scope
-def compute_prior(name, z, latent, hparams, condition=False, state=None):
+def compute_prior(name, z, latent, hparams, condition=False, state=None,
+                  temperature=1.0):
   """Distribution on z_t conditioned on z_{t-1} and latent.
 
   Args:
@@ -779,6 +798,7 @@ def compute_prior(name, z, latent, hparams, condition=False, state=None):
     state: tf.contrib.rnn.LSTMStateTuple.
            the current state of a LSTM used to model the distribution. Used
            only if hparams.latent_dist_encoder = "conv_lstm".
+    temperature: float, temperature with which to sample from the Gaussian.
   Returns:
     prior_dist: instance of tfp.distributions.Normal
     state: Returns updated state.
@@ -800,13 +820,13 @@ def compute_prior(name, z, latent, hparams, condition=False, state=None):
       mean, scale = tf.cond(
           condition, lambda: (cond_mean, cond_scale),
           lambda: (prior_mean, prior_scale))
-    dist = tfp.distributions.Normal(mean, scale)
+    dist = TemperedNormal(mean, scale, temperature)
     return dist, state
 
 
 @add_arg_scope
 def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
-          hparams=None, state=None, condition=False):
+          hparams=None, state=None, condition=False, temperature=1.0):
   """Splits / concatenates x into x1 and x2 across number of channels.
 
   For the forward pass, x2 is assumed be gaussian,
@@ -827,6 +847,7 @@ def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
            Used only when hparams.latent_dist_encoder == "conv_lstm"
     condition: bool, Whether or not to condition the distribution on
                cond_latents.
+    temperature: Temperature with which to sample from the gaussian.
 
   Returns:
   Raises:
@@ -846,7 +867,8 @@ def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
       return x1, logpb, eps, x2, state
     else:
       prior_dist, state = compute_prior(
-          "prior_on_z2", x, cond_latents, hparams, condition, state=state)
+          "prior_on_z2", x, cond_latents, hparams, condition, state=state,
+          temperature=temperature)
       if eps is not None:
         x2 = set_eps(prior_dist, eps)
       elif eps_std is not None:
@@ -873,10 +895,11 @@ def revnet_step(name, x, hparams, reverse=True):
   """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     ops = [
-        partial(actnorm, name="actnorm", reverse=reverse),
-        partial(invertible_1x1_conv, name="invertible", reverse=reverse),
-        partial(affine_coupling, name="affine", reverse=reverse,
-                mid_channels=hparams.affine_coupling_width)]
+        functools.partial(actnorm, name="actnorm", reverse=reverse),
+        functools.partial(invertible_1x1_conv, name="invertible",
+                          reverse=reverse),
+        functools.partial(affine_coupling, name="affine", reverse=reverse,
+                          mid_channels=hparams.affine_coupling_width)]
 
     if reverse:
       ops = ops[::-1]
@@ -931,7 +954,7 @@ def scale_gaussian_prior(name, z, logscale_factor=3.0, trainable=True):
 
 
 @add_arg_scope
-def top_prior(name, z_shape, learn_prior="normal"):
+def top_prior(name, z_shape, learn_prior="normal", temperature=1.0):
   """Unconditional prior distribution.
 
   Args:
@@ -943,6 +966,7 @@ def top_prior(name, z_shape, learn_prior="normal"):
                  and initialized such that the mean and std are zero and one.
                  If set to "normal", the prior is just a Gaussian with zero
                  mean and unit variance.
+    temperature: Temperature with which to sample from the Gaussian.
   Returns:
     objective: 1-D Tensor shape=(batch_size,) summed across spatial components.
   Raises:
@@ -957,7 +981,7 @@ def top_prior(name, z_shape, learn_prior="normal"):
     else:
       raise ValueError("Expected learn_prior to be normal or single_conv "
                        "got %s" % learn_prior)
-    return prior_dist
+    return TemperedNormal(prior_dist.loc, prior_dist.scale, temperature)
 
 
 def uniform_binning_correction(x, n_bits=8):
@@ -983,7 +1007,8 @@ def uniform_binning_correction(x, n_bits=8):
 
 @add_arg_scope
 def encoder_decoder(name, x, hparams, eps=None, reverse=False,
-                    cond_latents=None, condition=False, states=None):
+                    cond_latents=None, condition=False, states=None,
+                    temperature=1.0):
   """Glow encoder-decoder. n_levels of (Squeeze + Flow + Split.) operations."""
   # TODO(mechcoder) Change return_type to a dict to be backward compatible.
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
@@ -1037,7 +1062,8 @@ def encoder_decoder(name, x, hparams, eps=None, reverse=False,
           x, latent, state = split("split_%d" % level, x, eps=eps[level],
                                    reverse=True, cond_latents=curr_cond_latents,
                                    condition=condition, hparams=hparams,
-                                   state=states[level])
+                                   state=states[level],
+                                   temperature=temperature)
           new_states.append(state)
           all_latents.append(latent)
 
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 164783fca..3ebf5ea33 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -418,6 +418,30 @@ def test_temporal_latent_to_dist(self):
         self.assertTrue(np.allclose(mean_np, 0.0))
         self.assertTrue(np.allclose(scale_np, 1.0))
 
+  @parameterized.named_parameters(
+      ("temp_1.0", 1.0), ("temp_0.9", 0.9), ("temp_0.7", 0.7),
+      ("temp_0.3", 0.3), ("temp_0.1", 0.1), ("temp_0.0", 0.0))
+  def test_temperature_normal(self, temperature):
+    with tf.Graph().as_default():
+      rng = np.random.RandomState(0)
+      # in numpy, so that multiple calls don't trigger different random numbers.
+      loc_t = tf.convert_to_tensor(rng.randn(5, 5))
+      scale_t = tf.convert_to_tensor(rng.rand(5, 5))
+      tempered_normal = glow_ops.TemperedNormal(
+          loc=loc_t, scale=scale_t, temperature=temperature)
+      # smoke test for a single sample.
+      smoke_sample = tempered_normal.sample()
+      samples = tempered_normal.sample((10000,), seed=0)
+
+      with tf.Session() as sess:
+        ops = [samples, loc_t, scale_t, smoke_sample]
+        samples_np, loc_exp, scale_exp, _ = sess.run(ops)
+        scale_exp *= temperature
+        loc_act = np.mean(samples_np, axis=0)
+        scale_act = np.std(samples_np, axis=0)
+        self.assertTrue(np.allclose(loc_exp, loc_act, atol=1e-2))
+        self.assertTrue(np.allclose(scale_exp, scale_act, atol=1e-2))
+
 
 if __name__ == "__main__":
   tf.test.main()

From 39b7adc3dac899b3983cf5fe9711b937de149a24 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Wed, 14 Nov 2018 00:52:53 +0100
Subject: [PATCH 1228/2720] Model-Based RL: Split hparams dependent/indepent of
 RL algorithm (ppo/dqn) (#1222)

* MBRL: Remove limit on episodes used for evaluation. (unfinished are excluded later anyway)

* MBRL: Split hparams to dependent/independent of base RL algorithm. (e.g. dqn ppo)
---
 tensor2tensor/rl/trainer_model_based.py       |   4 +-
 .../rl/trainer_model_based_params.py          | 155 +++++++++++-------
 2 files changed, 101 insertions(+), 58 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index fe27b676d..360df0244 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -291,11 +291,9 @@ def evaluate_single_config(hparams, stochastic, max_num_noops, agent_model_dir):
       agent_model_dir=agent_model_dir
   )
   learner.evaluate(env_fn, eval_hparams, stochastic)
-  rollouts = env.current_epoch_rollouts()[:hparams.eval_batch_size]
+  rollouts = env.current_epoch_rollouts()
   env.close()
 
-  assert len(rollouts) == hparams.eval_batch_size, \
-      "{} {}".format(len(rollouts), hparams.eval_batch_size)
   return tuple(
       compute_mean_reward(rollouts, clipped) for clipped in (True, False)
   )
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 50ad899e1..6ff1a87cd 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -43,8 +43,7 @@
 HP_SCOPES = ["loop", "model", "ppo"]
 
 
-@registry.register_hparams
-def rlmb_base():
+def _rlmb_base():
   return tf.contrib.training.HParams(
       epochs=15,
       # Total frames used for training. This will be distributed evenly across
@@ -54,31 +53,16 @@ def rlmb_base():
       num_real_env_frames=96000,
       generative_model="next_frame_basic_deterministic",
       generative_model_params="next_frame_pixel_noise",
-      base_algo="ppo",
-      base_algo_params="ppo_original_params",
       autoencoder_train_steps=0,
       autoencoder_train_steps_initial_multiplier=10,
       autoencoder_hparams_set="autoencoder_discrete_pong",
       model_train_steps=15000,
       initial_epoch_train_steps_multiplier=3,
-      simulation_random_starts=True,  # Use random starts in PPO.
+      # Use random starts when learning agent on simulated env.
+      simulation_random_starts=True,
       # Flip the first random frame in PPO batch for the true beginning.
       simulation_flip_first_random_for_beginning=True,
       intrinsic_reward_scale=0.,
-      # Number of real environments to train on simultaneously.
-      real_batch_size=1,
-      # Number of simulated environments to train on simultaneously.
-      simulated_batch_size=16,
-      # Number of frames that can be taken from the simulated environment before
-      # it diverges, used for training the agent.
-      simulated_rollout_length=50,
-      ppo_epochs_num=1000,  # This should be enough to see something
-      # Should be equal to simulated_rollout_length.
-      # TODO(koz4k): Uncouple this by outputing done from SimulatedBatchEnv.
-      ppo_epoch_length=50,
-      # Do not eval since simulated batch env does not produce dones
-      ppo_eval_every_epochs=0,
-      ppo_learning_rate=1e-4,  # Will be changed, just so it exists.
       # Resizing.
       resize_height_factor=2,
       resize_width_factor=2,
@@ -93,17 +77,8 @@ def rlmb_base():
       # In your experiments, you want to optimize this rate to your schedule.
       learning_rate_bump=3.0,
 
-      # Unused; number of PPO epochs is calculated from the real frame limit.
-      real_ppo_epochs_num=0,
-      # This needs to be divisible by real_ppo_effective_num_agents.
-      real_ppo_epoch_length=16*200,
-      real_ppo_learning_rate=1e-4,
-      real_ppo_effective_num_agents=16,
-      real_ppo_eval_every_epochs=0,
-
       # Batch size during evaluation. Metrics are averaged over this number of
       # rollouts.
-      eval_batch_size=30,
       eval_max_num_noops=8,
 
       game="pong",
@@ -120,7 +95,66 @@ def rlmb_base():
       env_timesteps_limit=-1,  # Use default from gym.make()
       # Number of last observations to feed to the agent and world model.
       frame_stack_size=4,
+      # This is only used for world-model evaluation currently, PolicyLearner
+      # uses algorithm specific hparams to set this during training.
+      simulated_rollout_length=50,
+
+      # To be overriden
+      base_algo="",
+      base_algo_params="",
+      # Number of real environments to train on simultaneously.
+      real_batch_size=-1,
+      # Number of simulated environments to train on simultaneously.
+      simulated_batch_size=-1,
+      eval_batch_size=-1,
+  )
+
+
+def update_hparams(hparams, other):
+  for key, value in six.iteritems(other):
+    if key in hparams.values():
+      hparams.set_hparam(key, value)
+    else:
+      hparams.add_hparam(key, value)
+
+
+@registry.register_hparams
+def rlmb_ppo_base():
+  hparams = _rlmb_base()
+  ppo_params = dict(
+      base_algo="ppo",
+      base_algo_params="ppo_original_params",
+      # Number of real environments to train on simultaneously.
+      real_batch_size=1,
+      # Number of simulated environments to train on simultaneously.
+      simulated_batch_size=16,
+      eval_batch_size=30,
+
+      # Unused; number of PPO epochs is calculated from the real frame limit.
+      real_ppo_epochs_num=0,
+      # Number of frames that can be taken from the simulated environment before
+      # it diverges, used for training the agent.
+
+      ppo_epochs_num=1000,  # This should be enough to see something
+      # Should be equal to simulated_rollout_length.
+      # TODO(koz4k): Uncouple this by outputing done from SimulatedBatchEnv.
+      ppo_epoch_length=hparams.simulated_rollout_length,
+      # Do not eval since simulated batch env does not produce dones
+      ppo_eval_every_epochs=0,
+      ppo_learning_rate=1e-4,  # Will be changed, just so it exists.
+      # This needs to be divisible by real_ppo_effective_num_agents.
+      real_ppo_epoch_length=16 * 200,
+      real_ppo_learning_rate=1e-4,
+      real_ppo_effective_num_agents=16,
+      real_ppo_eval_every_epochs=0,
   )
+  update_hparams(hparams, ppo_params)
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base():
+  return rlmb_ppo_base()
 
 
 @registry.register_hparams
@@ -382,38 +416,49 @@ def rlmb_model_only():
   return hp
 
 
+def _rlmb_tiny_overrides():
+  """Parameters to override for tiny setting, excluding agent-related hparams"""
+  return dict(
+      epochs=1,
+      num_real_env_frames=128,
+      model_train_steps=2,
+      max_num_noops=1,
+      eval_max_num_noops=1,
+      generative_model_params="next_frame_tiny",
+      stop_loop_early=True,
+      resize_height_factor=2,
+      resize_width_factor=2,
+      wm_eval_rollout_ratios=[1],
+      env_timesteps_limit=7,
+      simulated_rollout_length=2,
+  )
+
+
 @registry.register_hparams
-def rlmb_tiny():
+def rlmb_ppo_tiny():
   """Tiny set for testing."""
-  return rlmb_base_sampling().override_from_dict(
-      tf.contrib.training.HParams(
-          epochs=1,
-          num_real_env_frames=128,
-          model_train_steps=2,
-          ppo_epochs_num=2,
-          simulated_batch_size=2,
-          simulated_rollout_length=2,
-          ppo_epoch_length=2,
-          real_batch_size=1,
-          real_ppo_epoch_length=36,
-          real_ppo_effective_num_agents=2,
-          max_num_noops=1,
-          eval_batch_size=1,
-          eval_max_num_noops=1,
-          generative_model_params="next_frame_tiny",
-          stop_loop_early=True,
-          resize_height_factor=2,
-          resize_width_factor=2,
-          game="pong",
-          wm_eval_rollout_ratios=[1],
-          env_timesteps_limit=7,
-      ).values())
+  hparams = rlmb_ppo_base()
+  hparams = hparams.override_from_dict(_rlmb_tiny_overrides())
+  update_hparams(hparams, dict(
+      ppo_epochs_num=2,
+      ppo_epoch_length=hparams.simulated_rollout_length,
+      real_ppo_epoch_length=36,
+      real_ppo_effective_num_agents=2,
+      real_batch_size=1,
+      eval_batch_size=1,
+  ))
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_tiny():
+  return rlmb_ppo_tiny()
 
 
 @registry.register_hparams
 def rlmb_tiny_stochastic():
   """Tiny setting with a stochastic next-frame model."""
-  hparams = rlmb_tiny()
+  hparams = rlmb_ppo_tiny()
   hparams.epochs = 1  # Too slow with 2 for regular runs.
   hparams.generative_model = "next_frame_basic_stochastic"
   hparams.generative_model_params = "next_frame_basic_stochastic"
@@ -423,7 +468,7 @@ def rlmb_tiny_stochastic():
 @registry.register_hparams
 def rlmb_tiny_recurrent():
   """Tiny setting with a recurrent next-frame model."""
-  hparams = rlmb_tiny()
+  hparams = rlmb_ppo_tiny()
   hparams.epochs = 1  # Too slow with 2 for regular runs.
   hparams.generative_model = "next_frame_basic_recurrent"
   hparams.generative_model_params = "next_frame_basic_recurrent"
@@ -433,7 +478,7 @@ def rlmb_tiny_recurrent():
 @registry.register_hparams
 def rlmb_tiny_sv2p():
   """Tiny setting with a tiny sv2p model."""
-  hparams = rlmb_tiny()
+  hparams = rlmb_ppo_tiny()
   hparams.generative_model = "next_frame_sv2p"
   hparams.generative_model_params = "next_frame_sv2p_tiny"
   hparams.grayscale = False

From 091de0469ad4ebc57fe487700d9802fc075c280c Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Tue, 13 Nov 2018 16:02:40 -0800
Subject: [PATCH 1229/2720] internal merge of PR #1222

PiperOrigin-RevId: 221351335
---
 tensor2tensor/rl/trainer_model_based_params.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 6ff1a87cd..a3905205c 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -99,7 +99,7 @@ def _rlmb_base():
       # uses algorithm specific hparams to set this during training.
       simulated_rollout_length=50,
 
-      # To be overriden
+      # To be overridden.
       base_algo="",
       base_algo_params="",
       # Number of real environments to train on simultaneously.
@@ -120,6 +120,7 @@ def update_hparams(hparams, other):
 
 @registry.register_hparams
 def rlmb_ppo_base():
+  """HParams for PPO base."""
   hparams = _rlmb_base()
   ppo_params = dict(
       base_algo="ppo",
@@ -417,7 +418,7 @@ def rlmb_model_only():
 
 
 def _rlmb_tiny_overrides():
-  """Parameters to override for tiny setting, excluding agent-related hparams"""
+  """Parameters to override for tiny setting excluding agent-related hparams."""
   return dict(
       epochs=1,
       num_real_env_frames=128,

From bf6b5d2d8d5b64b7f41a3604ce755b1bacfd81a1 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Tue, 13 Nov 2018 18:01:13 -0800
Subject: [PATCH 1230/2720] Add mixture of logistic layer.

Future work can refactor T2T's discretized mixture of logistics loss to use this layer. In the future, we can always add more complexity by for example: 1. splitting out the layer to be a composition of Logistic and Mixture layers; 2. adding more args. I prefer the simplest implementation which works for our applications.

PiperOrigin-RevId: 221368439
---
 tensor2tensor/layers/bayes.py      | 29 +++++++++++++++++++++++++++++
 tensor2tensor/layers/bayes_test.py | 17 +++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 882840296..2a35adea9 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -453,3 +453,32 @@ def loss_fn():
     return regularization
 
   return loss_fn
+
+
+class MixtureLogistic(tf.keras.layers.Layer):
+  """Stochastic output layer, distributed as a mixture of logistics."""
+
+  def __init__(self, num_components, **kwargs):
+    super(MixtureLogistic, self).__init__(**kwargs)
+    self.num_components = num_components
+    self.layer = tf.keras.layers.Dense(num_components * 3)
+
+  def build(self, input_shape=None):
+    self.layer.build(input_shape)
+    self.built = True
+
+  def call(self, inputs):
+    net = self.layer(inputs)
+    logits, loc, unconstrained_scale = tf.split(net, 3, axis=-1)
+    scale = tf.nn.softplus(unconstrained_scale) + tf.keras.backend.epsilon()
+    return ed.MixtureSameFamily(
+        mixture_distribution=ed.Categorical(logits=logits).distribution,
+        components_distribution=ed.Logistic(loc=loc, scale=scale).distribution)
+
+  def compute_output_shape(self, input_shape):
+    return tf.TensorShape(input_shape)[:-1]
+
+  def get_config(self):
+    config = {'num_components': self.num_components}
+    base_config = super(MixtureLogistic, self).get_config()
+    return dict(list(base_config.items()) + list(config.items()))
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 4addad626..eff85ba53 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -239,6 +239,23 @@ def testLSTMCellReparameterizationModel(self):
     self.assertAllClose(res2, res3)
     self.assertLen(model.losses, 2)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testMixtureLogistic(self):
+    batch_size = 3
+    features = tf.to_float(np.random.rand(batch_size, 4))
+    labels = tf.to_float(np.random.rand(batch_size))
+    model = tf.keras.Sequential([
+        tf.keras.layers.Dense(2, activation=None),
+        bayes.MixtureLogistic(5),
+    ])
+    outputs = model(features)
+    log_likelihood = tf.reduce_sum(outputs.distribution.log_prob(labels))
+    self.evaluate(tf.global_variables_initializer())
+    log_likelihood_val, outputs_val = self.evaluate([log_likelihood, outputs])
+    self.assertEqual(log_likelihood_val.shape, ())
+    self.assertLessEqual(log_likelihood_val, 0.)
+    self.assertEqual(outputs_val.shape, (batch_size,))
+
 
 if __name__ == "__main__":
   tf.test.main()

From 0c468033c735cafa799df74db1b178324e7bbd65 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Tue, 13 Nov 2018 19:42:55 -0800
Subject: [PATCH 1231/2720] adding to friends

PiperOrigin-RevId: 221377985
---
 tensor2tensor/layers/modalities.py | 31 ++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 79a920bf2..83935e4da 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -820,6 +820,37 @@ def top(self, body_output, _):
       return tf.expand_dims(res, 3)
 
 
+class VideoModalityIdentity(VideoModality):
+  """Video Modality where top and bottom is an identity function."""
+
+  def bottom(self, x):
+    common_video.gif_summary("inputs", x, max_outputs=1)
+    x = common_layers.standardize_images(x)
+    return x
+
+  def targets_bottom(self, x):
+    common_video.gif_summary("targets", x, max_outputs=1)
+    x = common_layers.standardize_images(x)
+    return x
+
+  def top(self, body_output, targets):
+    return body_output
+
+  def loss(self, top_out, targets):
+    """Compute loss numerator and denominator for one shard of output."""
+    # TODO(nikip): Try L2 loss
+    logits = top_out
+    logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
+    targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
+    cutoff = getattr(self._model_hparams, "video_modality_loss_cutoff", 0.01)
+    return common_layers.padded_cross_entropy(
+        logits,
+        targets,
+        self._model_hparams.label_smoothing,
+        cutoff=cutoff,
+        weights_fn=self.targets_weights_fn)
+
+
 class MultiLabelModality(ClassLabelModality):
   """Used for multi label task."""
 

From a2dc14d7366a8b1805aaceacf24b355f94a1f83a Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 13 Nov 2018 20:26:35 -0800
Subject: [PATCH 1232/2720] fix bug that prevents Transformer from running on
 TPU.

PiperOrigin-RevId: 221381614
---
 tensor2tensor/utils/t2t_model.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 58d8f4da0..d128ddbd1 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1963,4 +1963,5 @@ def _create_target_modality(modality_dict):
   # differently for modalities which are "targets"
   # (e.g., modality.target_bottom). In the future, remove need for this
   # behavior.
-  return {k: v for k, v in six.iteritems(modality_dict) if "target" in k}
+  return {k: v for k, v in six.iteritems(modality_dict) if "target" in k
+          and k != "targets_segmentation" and k != "targets_position"}

From 458b86ac298e6a2f0f0e623950fb7b2a477fac10 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Wed, 14 Nov 2018 14:16:57 -0800
Subject: [PATCH 1233/2720] Add a cond for discretization which disables EMA
 updates of codebook if false and some tests.

PiperOrigin-RevId: 221509099
---
 tensor2tensor/layers/discretization.py      | 21 ++++--
 tensor2tensor/layers/discretization_test.py | 75 ++++++++++++++++++++-
 2 files changed, 87 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 8478cf92e..050b3f504 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -508,7 +508,8 @@ def discrete_bottleneck(inputs,
                         noise_dev=1.,
                         startup_steps=50000,
                         summary=True,
-                        name=None):
+                        name=None,
+                        cond=True):
   """Discretization bottleneck.
 
   Args:
@@ -564,6 +565,7 @@ def discrete_bottleneck(inputs,
       only if bottleneck_kind is semhash.
     summary: Whether to write summaries.
     name: Name for the bottleneck scope.
+    cond: A tf.bool condition on whether to update the codebook.
 
   Returns:
     outputs_dense: Tensor of shape [..., output_dim]. The output dimension is
@@ -670,10 +672,11 @@ def discrete_bottleneck(inputs,
           tf.logging.info("Using EMA with beta = {}".format(beta))
           updated_ema_count_res = moving_averages.assign_moving_average(
               ema_count[i],
-              tf.reduce_sum(
-                  tf.reshape(
-                      x_means_hot_res, shape=[-1, num_blocks, block_v_size]),
-                  axis=0),
+              tf.where(cond,
+                       tf.reduce_sum(
+                           tf.reshape(x_means_hot_res,
+                                      shape=[-1, num_blocks, block_v_size]),
+                           axis=0), ema_count[i]),
               decay,
               zero_debias=False)
 
@@ -682,7 +685,8 @@ def discrete_bottleneck(inputs,
               tf.transpose(x_res, perm=[1, 0, 2]))
 
           updated_ema_means_res = moving_averages.assign_moving_average(
-              ema_means[i], dw, decay, zero_debias=False)
+              ema_means[i], tf.where(cond, dw, ema_means[i]),
+              decay, zero_debias=False)
           n = tf.reduce_sum(updated_ema_count_res, axis=-1, keep_dims=True)
           updated_ema_count_res = (
               (updated_ema_count_res + epsilon) / (n + 2**z_size * epsilon) * n)
@@ -692,7 +696,10 @@ def discrete_bottleneck(inputs,
           # pylint: enable=g-no-augmented-assignment
 
           with tf.control_dependencies([e_loss_res]):
-            update_means_res = tf.assign(means[i], updated_ema_means_res)
+            update_means_res = tf.assign(means[i],
+                                         tf.where(cond,
+                                                  updated_ema_means_res,
+                                                  means[i]))
             with tf.control_dependencies([update_means_res]):
               extra_loss += beta * e_loss_res
         else:
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index 689a53e27..0b29a04f6 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -163,6 +163,77 @@ def testGumbelSoftmaxDiscreteBottleneck(self):
     x_means_hot_eval = self.evaluate(x_means_hot)
     self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
 
-
-if __name__ == '__main__':
+  def testDiscreteBottleneckVQ(self):
+    hidden_size = 60
+    z_size = 4
+    x = tf.zeros(shape=[100, 1, hidden_size], dtype=tf.float32)
+    with tf.variable_scope("test", reuse=tf.AUTO_REUSE):
+      means = tf.get_variable("means",
+                              shape=[1, 1, 2**z_size, hidden_size],
+                              initializer=tf.constant_initializer(0.),
+                              dtype=tf.float32)
+      ema_count = []
+      ema_count_i = tf.get_variable(
+          "ema_count",
+          [1, 2**z_size],
+          initializer=tf.constant_initializer(0),
+          trainable=False)
+      ema_count.append(ema_count_i)
+      ema_means = []
+      with tf.colocate_with(means):
+        ema_means_i = tf.get_variable("ema_means",
+                                      initializer=means.initialized_value()[0],
+                                      trainable=False)
+        ema_means.append(ema_means_i)
+      x_means_dense, x_means_hot, _, _, _ = discretization.discrete_bottleneck(
+          x, hidden_size, z_size, 32, means=means, num_blocks=1,
+          ema_means=ema_means, ema_count=ema_count, name="test")
+      with self.test_session() as sess:
+        sess.run(tf.global_variables_initializer())
+        x_means_dense_eval, x_means_hot_eval = sess.run(
+            [x_means_dense, x_means_hot])
+        means_eval = sess.run(means)
+      self.assertEqual(x_means_dense_eval.shape, (100, 1, hidden_size))
+      self.assertEqual(x_means_hot_eval.shape, (100, 1))
+      self.assertTrue(np.all(means_eval == np.zeros(
+          (1, 1, 2**z_size, hidden_size))))
+
+  def testDiscreteBottleneckVQCond(self):
+    hidden_size = 60
+    z_size = 4
+    x = tf.zeros(shape=[100, 1, hidden_size], dtype=tf.float32)
+    with tf.variable_scope("test2", reuse=tf.AUTO_REUSE):
+      means = tf.get_variable("means",
+                              shape=[1, 1, 2**z_size, hidden_size],
+                              initializer=tf.constant_initializer(0.),
+                              dtype=tf.float32)
+      ema_count = []
+      ema_count_i = tf.get_variable(
+          "ema_count",
+          [1, 2**z_size],
+          initializer=tf.constant_initializer(0),
+          trainable=False)
+      ema_count.append(ema_count_i)
+      ema_means = []
+      with tf.colocate_with(means):
+        ema_means_i = tf.get_variable("ema_means",
+                                      initializer=means.initialized_value()[0],
+                                      trainable=False)
+        ema_means.append(ema_means_i)
+      cond = tf.cast(0.0, tf.bool)
+      x_means_dense, x_means_hot, _, _, _ = discretization.discrete_bottleneck(
+          x, hidden_size, z_size, 32, means=means, num_blocks=1, cond=cond,
+          ema_means=ema_means, ema_count=ema_count, name="test2")
+      with self.test_session() as sess:
+        sess.run(tf.global_variables_initializer())
+        x_means_dense_eval, x_means_hot_eval = sess.run(
+            [x_means_dense, x_means_hot])
+        means_eval = sess.run(means)
+      self.assertEqual(x_means_dense_eval.shape, (100, 1, hidden_size))
+      self.assertEqual(x_means_hot_eval.shape, (100, 1))
+      self.assertAllClose(means_eval, np.zeros((1, 1, 2**z_size,
+                                                hidden_size)))
+
+
+if __name__ == "__main__":
   tf.test.main()

From a9cca9542d528f148d0b230e3f8eb4106f0cd427 Mon Sep 17 00:00:00 2001
From: haukurb <haukzi@gmail.com>
Date: Wed, 14 Nov 2018 22:40:46 +0000
Subject: [PATCH 1234/2720] Fixed wrong configuration.json example in
 README.md, fixed missing d3 error, fixed json serialization error (#929)

---
 tensor2tensor/insights/README.md              |  8 +++----
 .../polymer/explore_view/explore-view.html    |  1 +
 tensor2tensor/insights/server.py              | 22 ++++++++++++++++++-
 3 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/insights/README.md b/tensor2tensor/insights/README.md
index 014bfca81..65ca95d59 100644
--- a/tensor2tensor/insights/README.md
+++ b/tensor2tensor/insights/README.md
@@ -42,14 +42,14 @@ Start guide, a sample configuration would be:
         "hparams": "",
         "hparams_set": "transformer_base_single_gpu",
         "problem": "translate_ende_wmt32k"
-      },
-    }]
+      }
+    }],
     "language": [{
       "code": "en",
-      "name": "English",
+      "name": "English"
     },{
       "code": "de",
-      "name": "German",
+      "name": "German"
     }]
   }
 ```
diff --git a/tensor2tensor/insights/polymer/explore_view/explore-view.html b/tensor2tensor/insights/polymer/explore_view/explore-view.html
index f74715016..9d40ea551 100644
--- a/tensor2tensor/insights/polymer/explore_view/explore-view.html
+++ b/tensor2tensor/insights/polymer/explore_view/explore-view.html
@@ -150,5 +150,6 @@ <h4>Rapid Response</h4>
       on-response="handleTranslationResponse_">
     </iron-ajax>
   </template>
+  <script src="../d3/d3.js"></script>
   <script src="explore-view.js"></script>
 </dom-module>
diff --git a/tensor2tensor/insights/server.py b/tensor2tensor/insights/server.py
index e61c465ff..1c4fad490 100644
--- a/tensor2tensor/insights/server.py
+++ b/tensor2tensor/insights/server.py
@@ -21,6 +21,8 @@
 from flask import jsonify
 from flask import request
 from flask import send_from_directory
+from flask.json import JSONEncoder
+import numpy as np
 from gunicorn.app.base import BaseApplication
 from gunicorn.six import iteritems
 from tensor2tensor.insights import transformer_model
@@ -37,6 +39,23 @@
                     "Path to static javascript and html files to serve.")
 
 
+_NUMPY_INT_DTYPES = [
+  np.int8, np.int16, np.int32, np.int64
+]
+_NUMPY_FP_DTYPES = [
+  np.float16, np.float32, np.float64
+]
+class NumpySerializationFix(JSONEncoder):
+  """json module cannot serialize numpy datatypes, reinterpret them first"""
+  def default(self, obj):
+    obj_type = type(obj)
+    if obj_type in _NUMPY_INT_DTYPES:
+      return int(obj)
+    if obj_type in _NUMPY_FP_DTYPES:
+      return float(obj)
+    return json.JSONEncoder.default(self, obj)
+
+
 class DebugFrontendApplication(BaseApplication):
   """A local custom application for GUnicorns.
 
@@ -101,6 +120,7 @@ def main(_):
       __name__.split(".")[0],
       static_url_path="/polymer",
       static_folder=FLAGS.static_path)
+  app.json_encoder = NumpySerializationFix
 
   # Disable static file caching.
   app.config["SEND_FILE_MAX_AGE_DEFAULT"] = 0
@@ -113,7 +133,7 @@ def language_list():  # pylint: disable=unused-variable
       JSON for the languages.
     """
     return jsonify({
-        "language": languages.values()
+        "language": list(languages.values())
     })
 
   @app.route("/api/list_models/")

From 7ae4c7fbac610af447ac3803e0007f82b4a02d60 Mon Sep 17 00:00:00 2001
From: haukurb <haukzi@gmail.com>
Date: Wed, 14 Nov 2018 14:41:15 -0800
Subject: [PATCH 1235/2720] internal merge of PR #929

PiperOrigin-RevId: 221513658
---
 tensor2tensor/insights/server.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/insights/server.py b/tensor2tensor/insights/server.py
index 1c4fad490..783b865bd 100644
--- a/tensor2tensor/insights/server.py
+++ b/tensor2tensor/insights/server.py
@@ -16,17 +16,15 @@
 """A GUnicorn + Flask Debug Frontend for Transformer models."""
 
 import json
-
 from flask import Flask
 from flask import jsonify
 from flask import request
 from flask import send_from_directory
 from flask.json import JSONEncoder
-import numpy as np
 from gunicorn.app.base import BaseApplication
 from gunicorn.six import iteritems
+import numpy as np
 from tensor2tensor.insights import transformer_model
-
 import tensorflow as tf
 
 flags = tf.flags
@@ -40,13 +38,16 @@
 
 
 _NUMPY_INT_DTYPES = [
-  np.int8, np.int16, np.int32, np.int64
+    np.int8, np.int16, np.int32, np.int64
 ]
 _NUMPY_FP_DTYPES = [
-  np.float16, np.float32, np.float64
+    np.float16, np.float32, np.float64
 ]
+
+
 class NumpySerializationFix(JSONEncoder):
   """json module cannot serialize numpy datatypes, reinterpret them first"""
+
   def default(self, obj):
     obj_type = type(obj)
     if obj_type in _NUMPY_INT_DTYPES:

From 9c6402b26b801eaf27beca9bf6a45d5128138765 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 14 Nov 2018 14:42:37 -0800
Subject: [PATCH 1236/2720] Replace one_hot + matmul with tf.gather on R1
 indices for faster gather operation.

PiperOrigin-RevId: 221513887
---
 tensor2tensor/layers/common_layers.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 5bc2cad11..8211b0fbf 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -272,13 +272,10 @@ def flatten4d3d(x):
 
 
 # TODO(noam): remove this function after TPUs do gather faster.
-def gather(params, indices, dtype=tf.float32):
+def gather(params, indices):
   """Version of tf.gather that works faster on tpu."""
-  if not is_xla_compiled():
-    return tf.gather(params, indices)
-  vocab_size = params.get_shape().as_list()[0]
   indices_flat = tf.reshape(indices, [-1])
-  out = tf.matmul(tf.one_hot(indices_flat, vocab_size, dtype=dtype), params)
+  out = tf.gather(params, indices_flat)
   out = reshape_like(out, tf.expand_dims(indices, -1))
   return out
 
@@ -352,7 +349,7 @@ def embedding(x,
     if not tf.contrib.eager.in_eager_mode():
       embedding_var = convert_gradient_to_tensor(embedding_var)
     x = dropout_no_scaling(x, 1.0 - symbol_dropout_rate)
-    emb_x = gather(embedding_var, x, dtype)
+    emb_x = gather(embedding_var, x)
     if multiplier != 1.0:
       emb_x *= multiplier
     static_shape = emb_x.shape.as_list()

From 8070eb454dcd7e87c5e9f59bebc31de05d1b0f19 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 14 Nov 2018 20:30:51 -0800
Subject: [PATCH 1237/2720] Minor, display ground truth videos only once during
 decode.

PiperOrigin-RevId: 221561341
---
 tensor2tensor/data_generators/video_utils.py  | 21 ++++++++++++-------
 .../data_generators/video_utils_test.py       |  6 ++++--
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 19a5c3c9d..3d6d072f7 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -75,15 +75,17 @@ def create_border(video, color="blue", border_percent=2):
 
 
 def convert_videos_to_summaries(input_videos, output_videos, target_videos,
-                                tag, decode_hparams):
+                                tag, decode_hparams,
+                                display_ground_truth=False):
   """Converts input, output and target videos into video summaries.
 
   Args:
     input_videos: 5-D NumPy array, (NTHWC) conditioning frames.
-    output_videos: 5-D NumPy array, (NTHWC) ground truth.
+    output_videos: 5-D NumPy array, (NTHWC) model predictions.
     target_videos: 5-D NumPy array, (NTHWC) target frames.
     tag: tf summary tag.
     decode_hparams: tf.contrib.training.HParams.
+    display_ground_truth: Whether or not to display ground truth videos.
   Returns:
     summaries: a list of tf frame-by-frame and video summaries.
   """
@@ -98,18 +100,20 @@ def convert_videos_to_summaries(input_videos, output_videos, target_videos,
   output_videos = create_border(
       output_videos, color="red", border_percent=border_percent)
 
-  # Video gif.
   all_input = np.concatenate((input_videos, target_videos), axis=1)
   all_output = np.concatenate((input_videos, output_videos), axis=1)
-  input_summ_vals, _ = common_video.py_gif_summary(
-      "%s/input" % tag, all_input, max_outputs=max_outputs, fps=fps,
-      return_summary_value=True)
   output_summ_vals, _ = common_video.py_gif_summary(
       "%s/output" % tag, all_output, max_outputs=max_outputs, fps=fps,
       return_summary_value=True)
-  all_summaries.extend(input_summ_vals)
   all_summaries.extend(output_summ_vals)
 
+  # Optionally display ground truth.
+  if display_ground_truth:
+    input_summ_vals, _ = common_video.py_gif_summary(
+        "%s/input" % tag, all_input, max_outputs=max_outputs, fps=fps,
+        return_summary_value=True)
+    all_summaries.extend(input_summ_vals)
+
   # Frame-by-frame summaries
   iterable = zip(all_input[:max_outputs], all_output[:max_outputs])
   for ind, (input_video, output_video) in enumerate(iterable):
@@ -164,7 +168,8 @@ def display_video_hooks(hook_args):
     input_videos = np.asarray(input_videos, dtype=np.uint8)
     summaries = convert_videos_to_summaries(
         input_videos, output_videos, target_videos,
-        tag="decode_%d" % decode_ind, decode_hparams=hook_args.decode_hparams)
+        tag="decode_%d" % decode_ind, decode_hparams=hook_args.decode_hparams,
+        display_ground_truth=decode_ind == 0)
     all_summaries.extend(summaries)
   return all_summaries
 
diff --git a/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index bdb2dd249..12d971882 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -73,9 +73,11 @@ def testConvertPredictionsToVideoSummaries(self):
         hparams=decode_hparams, decode_hparams=decode_hparams,
         predictions=predictions)
     summaries = video_utils.display_video_hooks(decode_hooks)
-    # for {random, psnr_max, psnr_min, ssim_max, ssim_min}
+    # for {psnr_max, psnr_min, ssim_max, ssim_min}
+    # 10 output vids + 10 frame-by-frame.
+    # for {random}
     # 10 input vids + 10 output vids + 10 frame-by-frame.
-    self.assertEqual(len(summaries), 150)
+    self.assertEqual(len(summaries), 110)
     for summary in summaries:
       self.assertTrue(isinstance(summary, tf.Summary.Value))
 

From ba2fa202452bf26a96927a9840da51f7dec2a1f5 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 14 Nov 2018 20:53:40 -0800
Subject: [PATCH 1238/2720] Remove linting requirement.

PiperOrigin-RevId: 221562948
---
 .travis.yml | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 62d037c47..367601195 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -38,6 +38,8 @@ script:
 
   # Conditional commands should each be in a separate block to get proper
   # errors on Travis.
-  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "tf-nightly"  ]]; then
-        pylint -j 2 tensor2tensor;
-    fi
+  #
+  # TODO(afrozm): Re-enable if this becomes an issue.
+  # - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "tf-nightly"  ]]; then
+  #       pylint -j 2 tensor2tensor;
+  #   fi

From d2b6b3a0885dcba995d74fe97f33c2e4b5ce2cf8 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 14 Nov 2018 22:11:37 -0800
Subject: [PATCH 1239/2720] Bump setup.py version to 1.11.0

PiperOrigin-RevId: 221568618
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index d0e586ee3..bbf1f7e88 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.10.0',
+    version='1.11.0',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',

From 6471246c3a0f03daaa472defea6079c40d5849ca Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Thu, 15 Nov 2018 09:23:33 -0800
Subject: [PATCH 1240/2720] support for fill-in-the-blank models in
 mtf-transformer

PiperOrigin-RevId: 221635675
---
 tensor2tensor/models/mtf_transformer.py | 37 +++++++++++++++++++------
 1 file changed, 28 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index a0bca0ccf..9cde58b95 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -169,7 +169,16 @@ def pad_to_max_length(x):
                 "inputs_segmentation", "inputs_position"]:
       if key in features:
         features[key] = pad_to_max_length(features[key])
-    shifted_targets = common_layers.shift_right_2d(targets)
+    if hparams.decoder_type == "autoregressive":
+      shifted_targets = common_layers.shift_right_2d(targets)
+    elif hparams.decoder_type == "masked":
+      shifted_targets = targets * tf.cast(
+          tf.greater(tf.random_uniform(tf.shape(targets), seed=123),
+                     hparams.mask_fraction),
+          targets.dtype)
+    else:
+      raise ValueError(
+          "unknown hparams.decoder_type = %s" % hparams.decoder_type)
 
     targets = self._import_to_batch_by_length(targets, "targets", mesh, hparams)
     shifted_targets = self._import_to_batch_by_length(
@@ -183,15 +192,18 @@ def pad_to_max_length(x):
       targets_position = self._import_to_batch_by_length(
           features["targets_position"], "targets_position",
           mesh, hparams)
-      decoder_self_attention_mask = (
-          mtf.layers.attention_mask_autoregressive(
-              targets_position, dtype=self.activation_dtype) +
-          mtf.layers.attention_mask_same_segment(
-              targets_segmentation, dtype=self.activation_dtype))
+      decoder_self_attention_mask = mtf.layers.attention_mask_same_segment(
+          targets_segmentation, dtype=self.activation_dtype)
+      if hparams.decoder_type == "autoregressive":
+        decoder_self_attention_mask += mtf.layers.attention_mask_autoregressive(
+            targets_position, dtype=self.activation_dtype)
     else:
       targets_position = mtf.range(mesh, self.length_dim, dtype=tf.int32)
-      decoder_self_attention_mask = mtf.layers.attention_mask_autoregressive(
-          targets_position, dtype=self.activation_dtype)
+      if hparams.decoder_type == "autoregressive":
+        decoder_self_attention_mask = mtf.layers.attention_mask_autoregressive(
+            targets_position, dtype=self.activation_dtype)
+      else:
+        decoder_self_attention_mask = None
 
     def layer_prepostprocess_dropout(x):
       return mtf.dropout(
@@ -687,12 +699,19 @@ def mtf_transformer_base():
   hparams.layer_prepostprocess_dropout = 0.1
 
   # Describes what model architecture:
-  #   "encdec": encoder + autoregerssive decoder
+  #   "encdec": encoder + autoregressive decoder
   #   "decoder": single-stack autoregressive sequence model.
   #   "encoder": single-stack non-autoregressive model
   #      with equal-length inputs and outputs.
   hparams.add_hparam("transformer_type", "encdec")
 
+  # What does the decoder do:
+  #   "autoregressive": Decoder left to right
+  #   "masked": Fills in masked-out values simultaneously
+  hparams.add_hparam("decoder_type", "autoregressive")
+  # for "masked", the probability of masking out each token
+  hparams.add_hparam("mask_fraction", 0.15)
+
   # round up vocab sizes to be a multiple of this value
   hparams.vocab_divisor = 128
 

From e5ecacf3db8b5bacf16e22f6b89f3e98e7d57fcd Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Thu, 15 Nov 2018 12:17:53 -0800
Subject: [PATCH 1241/2720] fixing the os bug.

PiperOrigin-RevId: 221667658
---
 tensor2tensor/rl/trainer_model_based.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 360df0244..b05e19ff3 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -450,7 +450,7 @@ def load_metrics(event_dir, epoch):
     metrics.
   """
   metrics = {}
-  for filename in os.listdir(event_dir):
+  for filename in tf.gfile.ListDirectory(event_dir):
     path = os.path.join(event_dir, filename)
     for event in tf.train.summary_iterator(path):
       if event.step == epoch and event.HasField("summary"):

From 36e91310f0bf69be5a601c89c26f376f29e4dbf3 Mon Sep 17 00:00:00 2001
From: Mirko Bronzi <m.bronzi@gmail.com>
Date: Thu, 15 Nov 2018 17:00:09 -0500
Subject: [PATCH 1242/2720] changed stopping condition for the beam decoder
 (when returning all the beams) - fixed test accordingly (#965)

---
 tensor2tensor/models/transformer_test.py |  2 --
 tensor2tensor/utils/beam_search.py       | 37 +++++++++++---------
 tensor2tensor/utils/beam_search_test.py  | 44 ++++++++++++++++++++++--
 3 files changed, 62 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 4924a8748..4850c817c 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -206,8 +206,6 @@ def testBeamVsFast(self):
       beam_res = beam_result.eval()
       fast_res = fast_result.eval()
 
-    self.assertEqual(fast_res.shape,
-                     (BATCH_SIZE, INPUT_LENGTH + decode_length))
     self.assertAllClose(beam_res, fast_res)
 
   def testTransformerWithoutProblem(self):
diff --git a/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index 5ac457e36..16020bf79 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -704,7 +704,7 @@ def inner_loop(i, alive_seq, alive_log_probs, finished_seq, finished_scores,
             finished_flags, states)
 
   def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
-                   finished_scores, finished_in_finished, unused_states):
+                   finished_scores, unused_finished_in_finished, unused_states):
     """Checking termination condition.
 
     We terminate when we decoded up to decode_length or the lowest scoring item
@@ -716,30 +716,33 @@ def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
       alive_log_probs: probabilities of the beams. [batch_size, beam_size]
       finished_scores: scores for each of these sequences.
         [batch_size, beam_size]
-      finished_in_finished: finished bools for each of these sequences.
-        [batch_size, beam_size]
 
     Returns:
       Bool.
     """
-    if not stop_early:
-      return tf.less(i, decode_length)
     max_length_penalty = tf.pow(((5. + tf.to_float(decode_length)) / 6.), alpha)
     # The best possible score of the most likely alive sequence.
     lower_bound_alive_scores = alive_log_probs[:, 0] / max_length_penalty
 
-    # Now to compute the lowest score of a finished sequence in finished
-    # If the sequence isn't finished, we multiply it's score by 0. since
-    # scores are all -ve, taking the min will give us the score of the lowest
-    # finished item.
-    lowest_score_of_finished_in_finished = tf.reduce_min(
-        finished_scores * tf.to_float(finished_in_finished), axis=1)
-    # If none of the sequences have finished, then the min will be 0 and
-    # we have to replace it by -ve INF if it is. The score of any seq in alive
-    # will be much higher than -ve INF and the termination condition will not
-    # be met.
-    lowest_score_of_finished_in_finished += (
-        (1. - tf.to_float(tf.reduce_any(finished_in_finished, 1))) * -INF)
+    if not stop_early:
+      # by considering the min score (in the top N beams) we ensure that
+      # the decoder will keep decoding until there is at least one beam
+      # (in the top N) that can be improved (w.r.t. the alive beams).
+      # any unfinished beam will have score -INF - thus the min
+      # will always be -INF if there is at least one unfinished beam -
+      # which means the bound_is_met condition cannot be true in this case.
+      lowest_score_of_finished_in_finished = tf.reduce_min(finished_scores)
+    else:
+      # by taking the max score we only care about the the first beam;
+      # as soon as this first beam cannot be beaten from the alive beams
+      # the beam decoder can stop.
+      # similarly to the above, if the top beam is not completed, its
+      # finished_score is -INF, thus it will not activate the
+      # bound_is_met condition. (i.e., decoder will keep going on).
+      # note we need to find the max for every sequence eparately - so, we need
+      # to keep the batch dimension (see axis=1)
+      lowest_score_of_finished_in_finished = tf.reduce_max(finished_scores,
+                                                           axis=1)
 
     bound_is_met = tf.reduce_all(
         tf.greater(lowest_score_of_finished_in_finished,
diff --git a/tensor2tensor/utils/beam_search_test.py b/tensor2tensor/utils/beam_search_test.py
index e7dbc9e75..21db81ee2 100644
--- a/tensor2tensor/utils/beam_search_test.py
+++ b/tensor2tensor/utils/beam_search_test.py
@@ -129,7 +129,7 @@ def symbols_to_logits(ids):
     self.assertAllEqual([[[0, 0, 1]]], ids)
     self.assertAllClose([[0.7 * 0.6]], np.exp(probs))
 
-  def testNotGreedyBeamTwo(self):
+  def testNotGreedyBeamTwoWithStopEarly(self):
     batch_size = 1
     beam_size = 2
     vocab_size = 3
@@ -152,11 +152,51 @@ def symbols_to_logits(ids):
         decode_length,
         vocab_size,
         0.0,
-        eos_id=1)
+        eos_id=1,
+        stop_early=True)  # defaul value, but just to make this explicit
+
+    with self.test_session():
+      ids = final_ids.eval()
+      probs = final_probs.eval()
+    # given stop_early = True, the only 'assurance' is w.r.t. the first beam
+    # (i.e., other beams may not even be completed)
+    # so, we check only the first beam
+    first_beam = ids[:, 0]
+    first_probs = probs[:, 0]
+    self.assertAllEqual([[0, 2, 1]], first_beam)
+    self.assertAllClose([0.8 * 0.5], np.exp(first_probs))
+
+  def testNotGreedyBeamTwoWithoutStopEarly(self):
+    batch_size = 1
+    beam_size = 2
+    vocab_size = 3
+    decode_length = 3
+
+    initial_ids = tf.constant([0] * batch_size)  # GO
+    probabilities = tf.constant([[[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
+                                 [[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
+                                 [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]]])
+
+    def symbols_to_logits(ids):
+      pos = tf.shape(ids)[1]
+      logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
+      return logits
+
+    final_ids, final_probs = beam_search.beam_search(
+        symbols_to_logits,
+        initial_ids,
+        beam_size,
+        decode_length,
+        vocab_size,
+        0.0,
+        eos_id=1,
+        stop_early=False)
 
     with self.test_session():
       ids = final_ids.eval()
       probs = final_probs.eval()
+    # given stop_early = False, the algorithm will return all the beams
+    # so we can test all of them here
     self.assertAllEqual([[[0, 2, 1, 0], [0, 2, 0, 1]]], ids)
     self.assertAllClose([[0.8 * 0.5, 0.8 * 0.4 * 0.9]], np.exp(probs))
 

From 8dac60fcf8ad9b24e58f95d44c8c7bd2c3289548 Mon Sep 17 00:00:00 2001
From: Mirko Bronzi <m.bronzi@gmail.com>
Date: Thu, 15 Nov 2018 14:50:23 -0800
Subject: [PATCH 1243/2720] internal merge of PR #965

PiperOrigin-RevId: 221694493
---
 tensor2tensor/utils/beam_search.py      | 2 +-
 tensor2tensor/utils/beam_search_test.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index 16020bf79..c303d63c1 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -733,7 +733,7 @@ def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
       # which means the bound_is_met condition cannot be true in this case.
       lowest_score_of_finished_in_finished = tf.reduce_min(finished_scores)
     else:
-      # by taking the max score we only care about the the first beam;
+      # by taking the max score we only care about the first beam;
       # as soon as this first beam cannot be beaten from the alive beams
       # the beam decoder can stop.
       # similarly to the above, if the top beam is not completed, its
diff --git a/tensor2tensor/utils/beam_search_test.py b/tensor2tensor/utils/beam_search_test.py
index 21db81ee2..9b6102252 100644
--- a/tensor2tensor/utils/beam_search_test.py
+++ b/tensor2tensor/utils/beam_search_test.py
@@ -153,7 +153,7 @@ def symbols_to_logits(ids):
         vocab_size,
         0.0,
         eos_id=1,
-        stop_early=True)  # defaul value, but just to make this explicit
+        stop_early=True)  # default value, but just to make this explicit
 
     with self.test_session():
       ids = final_ids.eval()

From 57db0b7140fc208e0c70b955ca509422e09eda45 Mon Sep 17 00:00:00 2001
From: Kentaro <cfiken@gmail.com>
Date: Sat, 17 Nov 2018 03:34:51 +0900
Subject: [PATCH 1244/2720] [UT] small bugs in ut_function when one uses data
 length = 1 (#1213)

* Fix: define squeeze dimension to avoid unknown TensorShape

* Fix: remove meaningless loop for set_shape
---
 .../models/research/universal_transformer_util.py      | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 50aaa2ca3..3f4aa4db8 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -1133,7 +1133,7 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
           use_bias=True,
           bias_initializer=tf.constant_initializer(
               hparams.act_halting_bias_init))
-      p = tf.squeeze(p)
+      p = tf.squeeze(p, axis=-1)
 
     # Mask for inputs which have not halted yet
     still_running = tf.cast(tf.less(halting_probability, 1.0), tf.float32)
@@ -1182,7 +1182,7 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
           state_shape[0],
           state_shape[1],
       ])
-      new_state.set_shape(state_shape)
+    new_state.set_shape(state_shape)
     step += 1
     return (transformed_state, step, halting_probability, remainders, n_updates,
             new_state)
@@ -1285,7 +1285,7 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
           use_bias=True,
           bias_initializer=tf.constant_initializer(
               hparams.act_halting_bias_init))
-      p = tf.squeeze(p)
+      p = tf.squeeze(p, axis=-1)
 
     # Mask for inputs which have not halted yet
     still_running = tf.cast(tf.less(halting_probability, 1.0), tf.float32)
@@ -1472,7 +1472,7 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
       x.set_shape([
           state_shape[0],
       ])
-      new_state.set_shape(state_shape)
+    new_state.set_shape(state_shape)
 
     step += 1
     return [
@@ -1620,7 +1620,7 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
           state_shape[0],
           state_shape[1],
       ])
-      new_state.set_shape(state_shape)
+    new_state.set_shape(state_shape)
     step += 1
     return [
         transformed_state, step, halting_probability, remainders, n_updates,

From 49e7cf5cad62e5cd6858d9b855de75088a476dcf Mon Sep 17 00:00:00 2001
From: Stefan Falk <43335432+stefan-falk@users.noreply.github.com>
Date: Fri, 16 Nov 2018 19:41:06 +0100
Subject: [PATCH 1245/2720]  Exposing batch_shuffle_size as hparam (#1231)

* Pass data_dir to feature_encoders

Pass data_dir to feature_encoders

* Fixing error passing wrong data_dir

* Exposing batch_shuffle_size as hparam

* Checking d_input since d_input may be None
---
 tensor2tensor/data_generators/problem.py | 10 ++++------
 tensor2tensor/layers/common_hparams.py   |  1 +
 tensor2tensor/utils/decoding.py          |  5 +++--
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 6b68d540b..7c4aed522 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -802,8 +802,7 @@ def input_fn(self,
                config=None,
                force_repeat=False,
                prevent_repeat=False,
-               dataset_kwargs=None,
-               batch_shuffle_size=512):
+               dataset_kwargs=None):
     """Builds input pipeline for problem.
 
     Args:
@@ -818,8 +817,6 @@ def input_fn(self,
         Overrides force_repeat.
       dataset_kwargs: dict, if passed, will pass as kwargs to self.dataset
         method when called
-      batch_shuffle_size: int, the size of the buffer to shuffle batches.
-        if none, the batches will not be shuffled.
 
     Returns:
       (features_dict<str name, Tensor feature>, Tensor targets)
@@ -969,8 +966,9 @@ def define_shapes(example):
     # buffer size for record shuffling is smaller than the batch size. In such
     # cases, adding batch shuffling ensures that the data is in random order
     # during training
-    if is_training and batch_shuffle_size:
-      dataset = dataset.shuffle(batch_shuffle_size)
+    if hasattr(hparams, 'batch_shuffle_size'):
+      if is_training and hparams.batch_shuffle_size:
+        dataset = dataset.shuffle(hparams.batch_shuffle_size)
 
     def prepare_for_output(example):
       if not config or not config.use_tpu:
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index ac1ed80d0..9a94207b3 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -33,6 +33,7 @@ def basic_params1():
       # of tokens per batch per GPU or per TPU core.  Otherwise, this is
       # the number of examples per GPU or per TPU core.
       batch_size=4096,
+      batch_shuffle_size=512,
       # If True, then if the features are of variable length, the batch_size is
       # used as the actual batch size (and not tokens per batch).
       use_fixed_batch_size=False,
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index c5520fd5a..0df99b1e4 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -319,8 +319,9 @@ def decode_once(estimator,
     if decode_to_file:
       for i, (d_input, d_output, d_target) in enumerate(decoded_outputs):
         # Skip if all padding
-        if re.match("^({})+$".format(text_encoder.PAD), d_input):
-          continue
+        if d_input:
+          if re.match("^({})+$".format(text_encoder.PAD), d_input):
+            continue
         beam_score_str = ""
         if decode_hp.write_beam_scores:
           beam_score_str = "\t%.2f" % decoded_scores[i]

From 9433a92637c2d6899bb399d465940cad39d4bc49 Mon Sep 17 00:00:00 2001
From: Kentaro <cfiken@gmail.com>
Date: Fri, 16 Nov 2018 11:10:05 -0800
Subject: [PATCH 1246/2720] internal merge of PR #1213

PiperOrigin-RevId: 221821207
---
 tensor2tensor/data_generators/problem.py | 10 ++++++----
 tensor2tensor/layers/common_hparams.py   |  1 -
 tensor2tensor/utils/decoding.py          |  5 ++---
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 7c4aed522..6b68d540b 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -802,7 +802,8 @@ def input_fn(self,
                config=None,
                force_repeat=False,
                prevent_repeat=False,
-               dataset_kwargs=None):
+               dataset_kwargs=None,
+               batch_shuffle_size=512):
     """Builds input pipeline for problem.
 
     Args:
@@ -817,6 +818,8 @@ def input_fn(self,
         Overrides force_repeat.
       dataset_kwargs: dict, if passed, will pass as kwargs to self.dataset
         method when called
+      batch_shuffle_size: int, the size of the buffer to shuffle batches.
+        if none, the batches will not be shuffled.
 
     Returns:
       (features_dict<str name, Tensor feature>, Tensor targets)
@@ -966,9 +969,8 @@ def define_shapes(example):
     # buffer size for record shuffling is smaller than the batch size. In such
     # cases, adding batch shuffling ensures that the data is in random order
     # during training
-    if hasattr(hparams, 'batch_shuffle_size'):
-      if is_training and hparams.batch_shuffle_size:
-        dataset = dataset.shuffle(hparams.batch_shuffle_size)
+    if is_training and batch_shuffle_size:
+      dataset = dataset.shuffle(batch_shuffle_size)
 
     def prepare_for_output(example):
       if not config or not config.use_tpu:
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 9a94207b3..ac1ed80d0 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -33,7 +33,6 @@ def basic_params1():
       # of tokens per batch per GPU or per TPU core.  Otherwise, this is
       # the number of examples per GPU or per TPU core.
       batch_size=4096,
-      batch_shuffle_size=512,
       # If True, then if the features are of variable length, the batch_size is
       # used as the actual batch size (and not tokens per batch).
       use_fixed_batch_size=False,
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 0df99b1e4..c5520fd5a 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -319,9 +319,8 @@ def decode_once(estimator,
     if decode_to_file:
       for i, (d_input, d_output, d_target) in enumerate(decoded_outputs):
         # Skip if all padding
-        if d_input:
-          if re.match("^({})+$".format(text_encoder.PAD), d_input):
-            continue
+        if re.match("^({})+$".format(text_encoder.PAD), d_input):
+          continue
         beam_score_str = ""
         if decode_hp.write_beam_scores:
           beam_score_str = "\t%.2f" % decoded_scores[i]

From 7918e93f991c11932e29f148164bfcfebca9182e Mon Sep 17 00:00:00 2001
From: Stefan Falk <43335432+stefan-falk@users.noreply.github.com>
Date: Fri, 16 Nov 2018 13:01:53 -0800
Subject: [PATCH 1247/2720] internal merge of PR #1231

PiperOrigin-RevId: 221839268
---
 tensor2tensor/data_generators/problem.py | 9 +++------
 tensor2tensor/layers/common_hparams.py   | 1 +
 tensor2tensor/utils/decoding.py          | 2 +-
 3 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 6b68d540b..ca0f213cd 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -802,8 +802,7 @@ def input_fn(self,
                config=None,
                force_repeat=False,
                prevent_repeat=False,
-               dataset_kwargs=None,
-               batch_shuffle_size=512):
+               dataset_kwargs=None):
     """Builds input pipeline for problem.
 
     Args:
@@ -818,8 +817,6 @@ def input_fn(self,
         Overrides force_repeat.
       dataset_kwargs: dict, if passed, will pass as kwargs to self.dataset
         method when called
-      batch_shuffle_size: int, the size of the buffer to shuffle batches.
-        if none, the batches will not be shuffled.
 
     Returns:
       (features_dict<str name, Tensor feature>, Tensor targets)
@@ -969,8 +966,8 @@ def define_shapes(example):
     # buffer size for record shuffling is smaller than the batch size. In such
     # cases, adding batch shuffling ensures that the data is in random order
     # during training
-    if is_training and batch_shuffle_size:
-      dataset = dataset.shuffle(batch_shuffle_size)
+    if is_training and hasattr(hparams, "batch_shuffle_size"):
+      dataset = dataset.shuffle(hparams.batch_shuffle_size)
 
     def prepare_for_output(example):
       if not config or not config.use_tpu:
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index ac1ed80d0..9a94207b3 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -33,6 +33,7 @@ def basic_params1():
       # of tokens per batch per GPU or per TPU core.  Otherwise, this is
       # the number of examples per GPU or per TPU core.
       batch_size=4096,
+      batch_shuffle_size=512,
       # If True, then if the features are of variable length, the batch_size is
       # used as the actual batch size (and not tokens per batch).
       use_fixed_batch_size=False,
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index c5520fd5a..ccd80a803 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -319,7 +319,7 @@ def decode_once(estimator,
     if decode_to_file:
       for i, (d_input, d_output, d_target) in enumerate(decoded_outputs):
         # Skip if all padding
-        if re.match("^({})+$".format(text_encoder.PAD), d_input):
+        if d_input and re.match("^({})+$".format(text_encoder.PAD), d_input):
           continue
         beam_score_str = ""
         if decode_hp.write_beam_scores:

From 624a9b6989d05e1c87e563a9b4af10939948b5ae Mon Sep 17 00:00:00 2001
From: RJ Ryan <rjryan@google.com>
Date: Fri, 16 Nov 2018 13:21:15 -0800
Subject: [PATCH 1248/2720] Re-write all references to tf.contrib.signal and
 tf.spectral to tf.signal.

PiperOrigin-RevId: 221842204
---
 tensor2tensor/layers/common_audio.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/common_audio.py b/tensor2tensor/layers/common_audio.py
index 14432576f..388ccce53 100644
--- a/tensor2tensor/layers/common_audio.py
+++ b/tensor2tensor/layers/common_audio.py
@@ -56,7 +56,7 @@ def compute_mel_filterbank_features(
     waveforms,
     sample_rate=16000, dither=1.0 / np.iinfo(np.int16).max, preemphasis=0.97,
     frame_length=25, frame_step=10, fft_length=None,
-    window_fn=functools.partial(tf.contrib.signal.hann_window, periodic=True),
+    window_fn=functools.partial(tf.signal.hann_window, periodic=True),
     lower_edge_hertz=80.0, upper_edge_hertz=7600.0, num_mel_bins=80,
     log_noise_floor=1e-3, apply_mask=True):
   """Implement mel-filterbank extraction using tf ops.
@@ -101,7 +101,7 @@ def compute_mel_filterbank_features(
   if fft_length is None:
     fft_length = int(2**(np.ceil(np.log2(frame_length))))
 
-  stfts = tf.contrib.signal.stft(
+  stfts = tf.signal.stft(
       waveforms,
       frame_length=frame_length,
       frame_step=frame_step,
@@ -121,7 +121,7 @@ def compute_mel_filterbank_features(
   # Warp the linear-scale, magnitude spectrograms into the mel-scale.
   num_spectrogram_bins = magnitude_spectrograms.shape[-1].value
   linear_to_mel_weight_matrix = (
-      tf.contrib.signal.linear_to_mel_weight_matrix(
+      tf.signal.linear_to_mel_weight_matrix(
           num_mel_bins, num_spectrogram_bins, sample_rate, lower_edge_hertz,
           upper_edge_hertz))
   mel_spectrograms = tf.tensordot(

From ba579f8067f629a36e35cc4e745e9bc0825d5ea0 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Fri, 16 Nov 2018 13:52:58 -0800
Subject: [PATCH 1249/2720] - Changed the interface to mtf.einsum to accept
 output_shape and/or reduced_dims.  In the case where both are provided, the
 compiler checks for consistency - this is useful in avoiding bugs. - Check
 shapes on imported tf tensors. - Factored mtf.log_softmax() code and exposed
 mtf.reduce_logsumexp().  Added option for extra logit (needed when everything
 is masked out). - work around a bug in tf.range for bfloat16

- Added option for memory-compressed self-attention over long sequences.

PiperOrigin-RevId: 221847183
---
 tensor2tensor/models/mtf_transformer.py       | 51 ++++++++++++++++---
 .../models/research/moe_experiments.py        | 38 +++++++++++++-
 2 files changed, 81 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 9cde58b95..3dbfc9224 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -275,6 +275,19 @@ def layer_prepostprocess_dropout(x):
             self_attention_mask=decoder_self_attention_mask,
             encdec_attention_mask=encoder_decoder_attention_mask,
             losses=extra_losses)
+    if (hparams.reshape_logits_hack and
+        hparams.mode == tf.estimator.ModeKeys.TRAIN):
+      # For some reason, the logits computation is extremely slow on TPU
+      # in some cases where the batch size per core is 1.  Reshape the logits
+      # and the targets to double the batch size and halve the length.
+      # TODO(noam): file a bug.
+      new_dims = self.batch_dims[:-1] + [
+          mtf.Dimension(self.batch_dims[-1].name,
+                        self.batch_dims[-1].size * 2),
+          mtf.Dimension(self.length_dim.name, self.length_dim.size // 2)]
+      x = mtf.reshape(x, new_dims + [self.model_dim])
+      targets = mtf.reshape(targets, new_dims)
+
     logits = mtf.matmul(x, softmax_var)
     if hparams.mode == tf.estimator.ModeKeys.TRAIN:
       logits = mtf.layers.multiplicative_jitter(logits, epsilon=1e-2)
@@ -289,13 +302,17 @@ def layer_prepostprocess_dropout(x):
     loss = mtf.reduce_mean(loss * weights)
     for l in extra_losses:
       loss += l
-    logits = mtf.to_float(logits)
-    # combine batch dims
-    if len(self.batch_dims) > 1:
-      combined_batch_dim = mtf.Dimension(
-          self.batch_dims[0].name, mtf.Shape(self.batch_dims).size)
-      logits = mtf.reshape(
-          logits, [combined_batch_dim] + logits.shape.dims[-2:])
+    if (hparams.reshape_logits_hack and
+        hparams.mode == tf.estimator.ModeKeys.TRAIN):
+      logits = None
+    else:
+      logits = mtf.to_float(logits)
+      # combine batch dims
+      if len(self.batch_dims) > 1:
+        combined_batch_dim = mtf.Dimension(
+            self.batch_dims[0].name, mtf.Shape(self.batch_dims).size)
+        logits = mtf.reshape(
+            logits, [combined_batch_dim] + logits.shape.dims[-2:])
     return logits, loss
 
   def mtf_model_fn(self, features, mesh):
@@ -493,6 +510,22 @@ def normalize(x):
                         hparams.layout, hparams.mesh_shape,
                         self.max_length_dim),
                     name="local_att"))
+        elif layer_type == "compressed_att":
+          if is_incremental:
+            raise ValueError("compressed_att incremental not implemented")
+          else:
+            x += layer_prepostprocess_dropout(
+                mtf.layers.multihead_self_attention_memory_compressed(
+                    normalize(x),
+                    mask_right=True,
+                    compression_factor=hparams.compression_factor,
+                    kv_channels=self.kv_dim,
+                    heads=self.heads_dim,
+                    dropout=hparams.attention_dropout,
+                    dropout_broadcast_dims=[self.length_dim],
+                    master_dtype=self.master_dtype,
+                    slice_dtype=self.slice_dtype,
+                    name="compressed_att"))
         else:
           if is_incremental:
             # insert length dimension.
@@ -754,6 +787,10 @@ def mtf_transformer_base():
   # hparams.batch_size // hparams.outer_batch_size.
   hparams.add_hparam("outer_batch_size", 0)
 
+  # TODO(noam): file a bug
+  hparams.add_hparam("reshape_logits_hack", False)
+  hparams.add_hparam("compression_factor", 4)
+
   return hparams
 
 
diff --git a/tensor2tensor/models/research/moe_experiments.py b/tensor2tensor/models/research/moe_experiments.py
index f53467099..1f9cf16c4 100644
--- a/tensor2tensor/models/research/moe_experiments.py
+++ b/tensor2tensor/models/research/moe_experiments.py
@@ -255,7 +255,8 @@ def xmoe2_v1_x128():
 def xmoe2_tiny():
   """Test on local cpu."""
   hparams = xmoe2_v1()
-  hparams.decoder_layers = ["local_att", "att", "drd", "hmoe"]
+  hparams.decoder_layers = [
+      "local_att", "att", "compressed_att", "drd", "hmoe"]
   hparams.d_model = 128
   hparams.moe_hidden_size = 512
   hparams.outer_batch_size = 0
@@ -272,6 +273,7 @@ def xmoe2_v1_l4k():
   hparams.batch_size = 32
   hparams.max_length = 4096
   hparams.split_to_length = 4096
+  hparams.reshape_logits_hack = True
   return hparams
 
 
@@ -284,6 +286,33 @@ def xmoe2_v1_l4k_local_only():
   return hparams
 
 
+@registry.register_hparams
+def xmoe2_v1_l4k_global_only():
+  """With sequence length 4096."""
+  hparams = xmoe2_v1_l4k()
+  hparams.decoder_layers = [
+      "att" if l == "local_att" else l for l in hparams.decoder_layers]
+  return hparams
+
+
+@registry.register_hparams
+def xmoe2_v1_l4k_compressed_c4():
+  """With compressed attention."""
+  hparams = xmoe2_v1_l4k()
+  hparams.decoder_layers = [
+      "compressed_att" if l == "att" else l for l in hparams.decoder_layers]
+  hparams.compression_factor = 4
+  return hparams
+
+
+@registry.register_hparams
+def xmoe2_v1_l4k_compressed_c8():
+  """With compressed attention."""
+  hparams = xmoe2_v1_l4k_compressed_c4()
+  hparams.compression_factor = 8
+  return hparams
+
+
 @registry.register_hparams
 def wiki_2x2_base():
   """Set of architectural experiments - language model on wikipedia on a 2x2.
@@ -330,3 +359,10 @@ def wiki_2x2_v1():
       ["local_att", "local_att", "drd",
        "att", "drd", "local_att", "local_att", "moe"] * 4)[:-1]
   return hparams
+
+
+@registry.register_hparams
+def wiki_2x2_local():
+  hparams = wiki_2x2_base()
+  hparams.decoder_layers = ["local_att", "drd"] * 6
+  return hparams

From eff7dbdc72c74b76e9859f4fb2c96bbe16d46ddc Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 16 Nov 2018 14:07:52 -0800
Subject: [PATCH 1250/2720] Implements a spatial dilated 3-d convolutional
 network to relax assumptions of locality in the Glow latent space via
 hparams.latent_apply_dilations. Given, features from the previous layer, new
 features are computed using multiple dilation rates and then merged. In
 addition,

PiperOrigin-RevId: 221849765
---
 tensor2tensor/models/research/glow_ops.py     | 146 ++++++++++++++----
 .../models/research/glow_ops_test.py          |  23 ++-
 2 files changed, 129 insertions(+), 40 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index e6cd25807..ab0089058 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -338,27 +338,45 @@ def add_edge_bias(x, filter_size):
   return tf.concat([x, x_pad], axis=3)
 
 
-def time_pad(x, filter_size):
+def time_pad(x, filter_size, dilations):
   """Pad left across time and pad valid across the spatial components.
 
+  Also concats a binary feature that indicates if a feature is padded or not.
+
   Args:
     x: 5-D Tensor, (NTHWC)
     filter_size: list of ints
+    dilations: list of ints, dilations - 1 specifies the number of holes
+               between two filter elements.
   Returns:
     x_pad: 5-D Tensor.
   """
+  x_shape = common_layers.shape_list(x)
   if filter_size == [1, 1, 1]:
     return x
-  a = (filter_size[1] - 1) // 2  # vertical padding size
-  b = (filter_size[2] - 1) // 2  # horizontal padding size
+  _, h, w = filter_size
+  eff_h = h + (h - 1)*(dilations[2] - 1)
+  eff_w = w + (w - 1)*(dilations[3] - 1)
+  a = (eff_h - 1) // 2  # vertical padding size
+  b = (eff_w - 1) // 2  # horizontal padding size
   c = filter_size[0] - 1
+
+  # pad across edges.
   padding = [[0, 0], [c, 0], [a, a], [b, b], [0, 0]]
-  return tf.pad(x, padding)
+
+  # concat a binary feature across channels to indicate a padding.
+  # 1 indicates that the feature is a padding.
+  x_bias = tf.zeros(x_shape[:-1] + [1])
+  x_bias = tf.pad(x_bias, padding, constant_values=1)
+  x_pad = tf.pad(x, padding)
+  x_pad = tf.concat((x_bias, x_pad), axis=-1)
+  return x_pad
 
 
 @add_arg_scope
 def conv(name, x, output_channels, filter_size=None, stride=None,
-         logscale_factor=3.0, apply_actnorm=True, conv_init="default"):
+         logscale_factor=3.0, apply_actnorm=True, conv_init="default",
+         dilations=None):
   """Convolutional layer with edge bias padding and optional actnorm.
 
   If x is 5-dimensional, actnorm is applied independently across every
@@ -376,6 +394,7 @@ def conv(name, x, output_channels, filter_size=None, stride=None,
                    have zero mean and unit variance. Else, there is no scaling
                    applied.
     conv_init: default or zeros. default is a normal distribution with 0.05 std.
+    dilations: List of integers, apply dilations.
   Returns:
     x: actnorm(conv2d(x))
   Raises:
@@ -393,6 +412,8 @@ def conv(name, x, output_channels, filter_size=None, stride=None,
       filter_size = [3, 3]
     if stride is None:
       stride = [1, 1]
+    if dilations is None:
+      dilations = [1, 1, 1, 1]
     actnorm_func = actnorm
     x = add_edge_bias(x, filter_size=filter_size)
     conv_filter = tf.nn.conv2d
@@ -401,8 +422,10 @@ def conv(name, x, output_channels, filter_size=None, stride=None,
       filter_size = [2, 3, 3]
     if stride is None:
       stride = [1, 1, 1]
+    if dilations is None:
+      dilations = [1, 1, 1, 1, 1]
     actnorm_func = actnorm_3d
-    x = time_pad(x, filter_size=filter_size)
+    x = time_pad(x, filter_size=filter_size, dilations=dilations)
     conv_filter = tf.nn.conv3d
 
   in_channels = common_layers.shape_list(x)[-1]
@@ -417,7 +440,7 @@ def conv(name, x, output_channels, filter_size=None, stride=None,
       initializer = tf.zeros_initializer()
 
     w = tf.get_variable("W", filter_shape, tf.float32, initializer=initializer)
-    x = conv_filter(x, w, stride_shape, padding="VALID")
+    x = conv_filter(x, w, stride_shape, padding="VALID", dilations=dilations)
     if apply_actnorm:
       x, _ = actnorm_func("actnorm", x, logscale_factor=logscale_factor)
     else:
@@ -430,14 +453,14 @@ def conv(name, x, output_channels, filter_size=None, stride=None,
 
 
 @add_arg_scope
-def conv_block(name, x, mid_channels, time_filter=2):
+def conv_block(name, x, mid_channels, dilations=None):
   """2 layer conv block used in the affine coupling layer.
 
   Args:
     name: variable scope.
     x: 4-D or 5-D Tensor.
     mid_channels: Output channels of the second layer.
-    time_filter: Filter across time to capture context.
+    dilations: Optional, list of integers.
   Returns:
     x: 4-D Tensor: Output activations.
   """
@@ -449,42 +472,73 @@ def conv_block(name, x, mid_channels, time_filter=2):
       first_filter = [3, 3]
       second_filter = [1, 1]
     else:
-      first_filter = [time_filter, 3, 3]
+      first_filter = [2, 3, 3]
       second_filter = [1, 1, 1]
 
     # Edge Padding + conv2d + actnorm + relu:
     # [output: 512 channels]
-    x = conv("1_1", x, output_channels=mid_channels, filter_size=first_filter)
+    x = conv("1_1", x, output_channels=mid_channels, filter_size=first_filter,
+             dilations=dilations)
     x = tf.nn.relu(x)
 
     # Padding + conv2d + actnorm + relu
     # [input, output: 512 channels]
-    x = conv("1_2", x, output_channels=mid_channels, filter_size=second_filter)
+    x = conv("1_2", x, output_channels=mid_channels, filter_size=second_filter,
+             dilations=dilations)
     x = tf.nn.relu(x)
     return x
 
 
+def dilated_conv_stack(name, x, mid_channels, output_channels,
+                       dilation_rates):
+  """Dilated convolutional stack.
+
+  Features at different rates are computed independently using a 3 layer
+  convolutional stack and added.
+
+  Args:
+    name: variable scope.
+    x: 5-D Tensor.
+    mid_channels: Number of output channels of the first layer in the conv
+                  stack.
+    output_channels: Number of output channels of the last layer.
+    dilation_rates: A list of dilation rates.
+  Returns:
+    output: 5-D Tensor.
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    output = 0.0
+    for dil_ind, dil_rate in enumerate(dilation_rates):
+      # TODO(mechcoder) try (concat across channels + 1x1) modulo memory issues.
+      curr_out = conv_stack("dil_%d" % dil_ind, x, mid_channels=mid_channels,
+                            output_channels=output_channels, dilations=dil_rate)
+      output += curr_out
+    return output
+
+
 @add_arg_scope
-def affine_coupling_network(name, x, mid_channels, output_channels):
-  """3-layer conv2d.
+def conv_stack(name, x, mid_channels, output_channels, dilations=None):
+  """3-layer convolutional stack.
 
   Args:
-    name:
-    x:
+    name: variable scope.
+    x: 5-D Tensor.
     mid_channels: Number of output channels of the first layer.
     output_channels: Number of output channels.
+    dilations: Dilations to apply in the first 3x3 layer and the last 3x3 layer.
+               By default, apply no dilations.
 
   Returns:
-    output:
+    output: output of 3 layer conv network.
   """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
 
-    x = conv_block("conv_block", x, mid_channels=mid_channels)
+    x = conv_block("conv_block", x, mid_channels=mid_channels,
+                   dilations=dilations)
 
     # Final layer.
-    x = conv("zeros", x, filter_size=[3, 3], stride=[1, 1],
-             output_channels=output_channels, apply_actnorm=False,
-             conv_init="zeros")
+    x = conv("zeros", x, apply_actnorm=False, conv_init="zeros",
+             output_channels=output_channels, dilations=dilations)
   return x
 
 
@@ -511,8 +565,7 @@ def affine_coupling(name, x, mid_channels=512, reverse=False):
     # Else:
     # z2 = (x2 / scale) - shift
     z1 = x1
-    log_scale_and_shift = affine_coupling_network(
-        "nn", x1, mid_channels, x_shape[-1])
+    log_scale_and_shift = conv_stack("nn", x1, mid_channels, x_shape[-1])
     shift = log_scale_and_shift[:, :, :, 0::2]
     scale = tf.nn.sigmoid(log_scale_and_shift[:, :, :, 1::2] + 2.0)
     if not reverse:
@@ -564,6 +617,30 @@ def squeeze(name, x, factor=2, reverse=True):
     return x
 
 
+def get_dilation_rates(hparams, width):
+  """Get a list of valid dilation rates.
+
+  Args:
+    hparams: tf.contrib.training.HParams.
+    width: spatial dimension. Ensures that the effective filter size is
+           not larger than the spatial dimension.
+  Returns:
+    allowed_dilations: A list of dilation rates.
+  """
+  # dil_rate=1 means no dilation.
+  allowed_dilations = [[1]*5]
+  apply_dilations = hparams.get("latent_apply_dilations", False)
+  dilation_rates = [1, 3]   # Number of holes between each filter element.
+  if apply_dilations:
+    for rate in dilation_rates:
+      # k + (k - 1) * rate but k is harcoded to be 3 everywhere.
+      filter_size = 3 + 2 * rate
+      if filter_size <= width:
+        curr_dilation = [1, 1, rate+1, rate+1, 1]
+        allowed_dilations.append(curr_dilation)
+  return allowed_dilations
+
+
 @add_arg_scope
 def temporal_latent_to_dist(name, x, hparams, output_channels=None):
   """Network that maps a time-indexed list of 3-D latents to a gaussian.
@@ -576,17 +653,23 @@ def temporal_latent_to_dist(name, x, hparams, output_channels=None):
   Returns:
     dist: tfp.distributions.Normal
   """
-  res_channels = common_layers.shape_list(x)[-1]
+  _, _, width, _, res_channels = common_layers.shape_list(x)
   if output_channels is None:
     output_channels = res_channels
+  dilation_rates = get_dilation_rates(hparams, width)
+
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     h = x
     for i in range(hparams.latent_encoder_depth):
-      h1 = conv_block("conv3d_1_%d" % i, h, time_filter=2,
-                      mid_channels=hparams.latent_encoder_width)
-      h2 = conv("conv3d_zeros_%d" % i, h1, apply_actnorm=False,
-                output_channels=res_channels, conv_init="zeros",
-                filter_size=[2, 3, 3])
+      if hparams.latent_apply_dilations:
+        h2 = dilated_conv_stack("dil_latent_3d_res_%d" % i, h,
+                                mid_channels=hparams.latent_encoder_width,
+                                output_channels=res_channels,
+                                dilation_rates=dilation_rates)
+      else:
+        h2 = conv_stack("latent_3d_res_%d" % i, h,
+                        mid_channels=hparams.latent_encoder_width,
+                        output_channels=res_channels)
       h += h2
 
     # take last activation that should capture all context since padding is
@@ -664,9 +747,8 @@ def latent_to_dist(name, x, hparams, output_channels=None):
     elif architecture == "glow_resnet":
       h = x
       for layer in range(depth):
-        h2 = conv_block("glow_res_%d" % layer, h, mid_channels=width)
-        h3 = conv("glow_res_zeros_%d" % layer, h2, conv_init="zeros",
-                  output_channels=x_shape[-1], apply_actnorm=False)
+        h3 = conv_stack("latent_resnet_%d" % layer, h,
+                        mid_channels=width, output_channels=x_shape[-1])
         h += h3
       mean_log_scale = conv("glow_res_final", h, conv_init="zeros",
                             output_channels=2*output_channels,
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 3ebf5ea33..10b5c1f0a 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -43,6 +43,7 @@ def get_glow_hparams(self):
     hparams.add_hparam("latent_architecture", "glow_resnet")
     # Use latent skip connections
     hparams.add_hparam("model_input", False)
+    hparams.add_hparam("latent_apply_dilations", False)
     hparams.add_hparam("latent_skip", True)
     hparams.add_hparam("latent_encoder_depth", 2)
     hparams.add_hparam("latent_encoder_width", 256)
@@ -125,11 +126,11 @@ def test_conv2d(self):
         # test shape in case apply_actnorm is set to False,
         self.assertEqual(zeros_np.shape, (16, 5, 5, 64))
 
-  def test_affine_coupling_network(self):
+  def test_conv_stack(self):
     """Test output shape."""
     with tf.Graph().as_default():
       x = 10.0 * tf.random_uniform(shape=(16, 5, 5, 32))
-      nn = glow_ops.affine_coupling_network("nn", x, 512, 64)
+      nn = glow_ops.conv_stack("nn", x, 512, 64)
 
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())
@@ -404,19 +405,25 @@ def test_actnorm_3d(self):
           self.assertTrue(np.allclose(channel_mean, 0.0, atol=1e-3))
           self.assertTrue(np.allclose(channel_var, 1.0, atol=1e-3))
 
-  def test_temporal_latent_to_dist(self):
+  @parameterized.named_parameters(
+      ("dilation", True), ("no_dilation", False))
+  def test_temporal_latent_to_dist(self, apply_dilation):
     with tf.Graph().as_default():
       hparams = self.get_glow_hparams()
-      latent_shape = (16, 5, 4, 4, 48)
+      hparams.latent_apply_dilations = apply_dilation
+      latent_shape = (16, 5, 32, 32, 48)
       latents = tf.random_normal(latent_shape)
       dist = glow_ops.temporal_latent_to_dist(
           "tensor_to_dist", latents, hparams)
       with tf.Session() as sess:
         sess.run(tf.global_variables_initializer())
-        mean, scale = dist.loc, dist.scale
-        mean_np, scale_np = sess.run([mean, scale])
-        self.assertTrue(np.allclose(mean_np, 0.0))
-        self.assertTrue(np.allclose(scale_np, 1.0))
+        # dilated conv_3d is not available on CPU.
+        is_gpu = tf.test.is_gpu_available()
+        if not apply_dilation or is_gpu:
+          mean, scale = dist.loc, dist.scale
+          mean_np, scale_np = sess.run([mean, scale])
+          self.assertTrue(np.allclose(mean_np, 0.0))
+          self.assertTrue(np.allclose(scale_np, 1.0))
 
   @parameterized.named_parameters(
       ("temp_1.0", 1.0), ("temp_0.9", 0.9), ("temp_0.7", 0.7),

From 8a7a6223f94c1a71ad639d1dfad13fae23e69446 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 16 Nov 2018 14:35:35 -0800
Subject: [PATCH 1251/2720] If hparams.batch_shuffle_size is 0, don't call
 shuffle.

PiperOrigin-RevId: 221854316
---
 tensor2tensor/data_generators/problem.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index ca0f213cd..3791123c1 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -966,7 +966,8 @@ def define_shapes(example):
     # buffer size for record shuffling is smaller than the batch size. In such
     # cases, adding batch shuffling ensures that the data is in random order
     # during training
-    if is_training and hasattr(hparams, "batch_shuffle_size"):
+    if (is_training and hasattr(hparams, "batch_shuffle_size") and
+        hparams.batch_shuffle_size):
       dataset = dataset.shuffle(hparams.batch_shuffle_size)
 
     def prepare_for_output(example):

From f50501a78efd8ac7661a8e8d566d95b1a76a78f9 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Fri, 16 Nov 2018 15:10:47 -0800
Subject: [PATCH 1252/2720] More noising options for denoising
 mtf_transformers.  Enable token replacement either from random zipfian or
 from an auxiliary denoising transformer.

PiperOrigin-RevId: 221860576
---
 tensor2tensor/models/mtf_transformer.py       | 121 ++++++++++++++----
 .../models/research/moe_experiments.py        | 111 ++++++++++++++++
 2 files changed, 209 insertions(+), 23 deletions(-)

diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 3dbfc9224..f06115a2a 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -21,7 +21,6 @@
 
 import copy
 import mesh_tensorflow as mtf
-
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import modalities
@@ -151,9 +150,76 @@ def _embedding_and_softmax_vars(self, mesh):
     return (inputs_embedding_var, targets_embedding_var,
             softmax_var, positional_embedding_var)
 
+  def _noisy_targets_from_spec(self, targets, noising_spec, losses=None):
+    if noising_spec["type"] == "mask":
+      # Replace a randomly-chosen noising_spec["prob"] of input tokens with 0.
+      return targets * mtf.cast(
+          mtf.greater(mtf.random_uniform(targets.mesh, targets.shape),
+                      noising_spec["prob"]), targets.dtype)
+    elif noising_spec["type"] == "random_zipfian":
+      # Replace a randomly-chosen noising_spec["prob"] of input tokens.
+      # Rather than drawing the replacement tokens uniformly, we sample from
+      #   a distribution favoring lower token-ids, assuming that the ids have
+      #   been assigned in frequency order.  The probability of choosing an
+      #   id is proportional to 1/(id+10)
+      logits = mtf.log(1.0 / (mtf.range(
+          targets.mesh, self.targets_vocab_dim, dtype=tf.float32) + 10.0))
+      logits = mtf.broadcast(logits, new_shape=targets.shape + logits.shape)
+      r = mtf.sample_with_temperature(logits, self.targets_vocab_dim)
+      use_noise = mtf.less(
+          mtf.random_uniform(targets.mesh, targets.shape), noising_spec["prob"])
+      return mtf.where(use_noise, r, targets)
+    elif noising_spec["type"] == "transformer":
+      # Train a small transformer to fill in masked out values, then
+      # sample from it.
+      hparams = self._hparams
+      if hparams.mode != tf.estimator.ModeKeys.TRAIN:
+        raise NotImplementedError("Not implemented")
+      noiser_hparams = copy.copy(self._hparams)
+      noiser_hparams.del_hparam("mode")
+      noiser_hparams.override_from_dict(noising_spec["overrides"])
+      with tf.variable_scope("noiser"):
+        noiser = MtfTransformer(
+            noiser_hparams,
+            mode=hparams.mode,
+            problem_hparams=self._problem_hparams)
+        logits, loss = noiser._mtf_model_fn(  # pylint: disable=protected-access
+            self._original_features, targets.mesh)
+        samples = mtf.sample_with_temperature(logits, self.targets_vocab_dim)
+      losses.append(loss)
+      return samples
+    else:
+      raise ValueError("unknown noising spec %s" % noising_spec)
+
+  def _noisy_targets(self, targets, losses=None):
+    """Generate noisy targets for denoising models.
+
+    Args:
+      targets: a Tensor
+      losses: an optional list onto which to append traning losses
+    Returns:
+      a Tensor the same dtype and shape as Targets
+    """
+    hparams = self._hparams
+    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+      nt_train = self._noisy_targets_from_spec(
+          targets, hparams.noising_spec_train, losses=losses)
+      if hparams.noising_use_eval_during_train > 0:
+        nt_eval = self._noisy_targets_from_spec(
+            targets, hparams.noising_spec_eval)
+        use_eval_noising = mtf.less(
+            mtf.random_uniform(targets.mesh, targets.shape - self.length_dim),
+            hparams.noising_use_eval_during_train)
+        nt_train = mtf.where(use_eval_noising, nt_eval, nt_train)
+      return nt_train
+    else:
+      return self._noisy_targets_from_spec(targets, hparams.noising_spec_eval)
+
   def _mtf_model_fn(self, features, mesh):
+    self._original_features = features
     features = copy.copy(features)
     hparams = self._hparams
+    extra_losses = []
     targets = tf.to_int32(features["targets"])
     if len(targets.get_shape()) > 2:
       tf.logging.info("targets = %s" % targets)
@@ -165,25 +231,20 @@ def pad_to_max_length(x):
       x = tf.reshape(x, [hparams.batch_size, hparams.max_length])
       return x
     targets = pad_to_max_length(targets)
+    targets = self._import_to_batch_by_length(targets, "targets", mesh, hparams)
     for key in ["targets_segmentation", "targets_position",
                 "inputs_segmentation", "inputs_position"]:
       if key in features:
         features[key] = pad_to_max_length(features[key])
     if hparams.decoder_type == "autoregressive":
-      shifted_targets = common_layers.shift_right_2d(targets)
-    elif hparams.decoder_type == "masked":
-      shifted_targets = targets * tf.cast(
-          tf.greater(tf.random_uniform(tf.shape(targets), seed=123),
-                     hparams.mask_fraction),
-          targets.dtype)
+      shifted_targets = mtf.shift(
+          targets, offset=1, dim=self.length_dim, wrap=False)
+    elif hparams.decoder_type == "denoising":
+      shifted_targets = self._noisy_targets(targets, extra_losses)
     else:
       raise ValueError(
           "unknown hparams.decoder_type = %s" % hparams.decoder_type)
 
-    targets = self._import_to_batch_by_length(targets, "targets", mesh, hparams)
-    shifted_targets = self._import_to_batch_by_length(
-        shifted_targets, "shifted_targets", mesh, hparams)
-
     if "targets_segmentation" in features:
       # "Packed" dataset - keep the examples from seeing each other.
       targets_segmentation = self._import_to_batch_by_length(
@@ -210,7 +271,6 @@ def layer_prepostprocess_dropout(x):
           x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
           noise_shape=mtf.Shape(self.batch_dims + [self.model_dim]))
 
-    extra_losses = []
     (inputs_embedding_var,
      targets_embedding_var,
      softmax_var,
@@ -281,6 +341,7 @@ def layer_prepostprocess_dropout(x):
       # in some cases where the batch size per core is 1.  Reshape the logits
       # and the targets to double the batch size and halve the length.
       # TODO(noam): file a bug.
+      old_dims = self.batch_dims + [self.length_dim]
       new_dims = self.batch_dims[:-1] + [
           mtf.Dimension(self.batch_dims[-1].name,
                         self.batch_dims[-1].size * 2),
@@ -304,20 +365,20 @@ def layer_prepostprocess_dropout(x):
       loss += l
     if (hparams.reshape_logits_hack and
         hparams.mode == tf.estimator.ModeKeys.TRAIN):
-      logits = None
-    else:
-      logits = mtf.to_float(logits)
+      logits = mtf.reshape(logits, old_dims + [self.targets_vocab_dim])
+    logits = mtf.to_float(logits)
+    return logits, loss
+
+  def mtf_model_fn(self, features, mesh):
+    with tf.variable_scope("transformer"):
+      logits, loss = self._mtf_model_fn(features, mesh)
       # combine batch dims
       if len(self.batch_dims) > 1:
         combined_batch_dim = mtf.Dimension(
             self.batch_dims[0].name, mtf.Shape(self.batch_dims).size)
         logits = mtf.reshape(
             logits, [combined_batch_dim] + logits.shape.dims[-2:])
-    return logits, loss
-
-  def mtf_model_fn(self, features, mesh):
-    with tf.variable_scope("transformer"):
-      return self._mtf_model_fn(features, mesh)
+      return logits, loss
 
   @property
   def _targets_vocab_size(self):
@@ -740,10 +801,14 @@ def mtf_transformer_base():
 
   # What does the decoder do:
   #   "autoregressive": Decoder left to right
-  #   "masked": Fills in masked-out values simultaneously
+  #   "denoising": Fills in masked-out values simultaneously
   hparams.add_hparam("decoder_type", "autoregressive")
-  # for "masked", the probability of masking out each token
-  hparams.add_hparam("mask_fraction", 0.15)
+
+  # Parameters describing the noising algorithm for denoising decoders
+  hparams.add_hparam("noising_spec_train", {"type": "mask", "prob": 0.15})
+  hparams.add_hparam("noising_spec_eval", {"type": "mask", "prob": 0.15})
+  # during training, we use the eval noiser with this probability
+  hparams.add_hparam("noising_use_eval_during_train", 0.1)
 
   # round up vocab sizes to be a multiple of this value
   hparams.vocab_divisor = 128
@@ -830,6 +895,16 @@ def mtf_transformer_tiny_lm():
   return hparams
 
 
+@registry.register_hparams
+def mtf_transformer_tiny_denoising():
+  hparams = mtf_transformer_tiny_lm()
+  hparams.decoder_type = "denoising"
+  hparams.noising_spec_train = ("random_zipfian", 0.3)
+  hparams.noising_use_eval_during_train = 0.5
+  hparams.max_length = 1024
+  return hparams
+
+
 @registry.register_hparams
 def mtf_transformer_single():
   hparams = mtf_transformer_tiny()
diff --git a/tensor2tensor/models/research/moe_experiments.py b/tensor2tensor/models/research/moe_experiments.py
index 1f9cf16c4..04a09cfd1 100644
--- a/tensor2tensor/models/research/moe_experiments.py
+++ b/tensor2tensor/models/research/moe_experiments.py
@@ -366,3 +366,114 @@ def wiki_2x2_local():
   hparams = wiki_2x2_base()
   hparams.decoder_layers = ["local_att", "drd"] * 6
   return hparams
+
+
+@registry.register_hparams
+def denoise_m15():
+  """Denoising experiment."""
+  hparams = xmoe2_dense_0()
+  hparams.decoder_type = "denoising"
+  hparams.noising_spec_train = {"type": "mask", "prob": 0.15}
+  return hparams
+
+
+@registry.register_hparams
+def denoise_m30():
+  """More masking during training."""
+  hparams = xmoe2_dense_0()
+  hparams.decoder_type = "denoising"
+  hparams.noising_spec_train = {"type": "mask", "prob": 0.3}
+  return hparams
+
+
+@registry.register_hparams
+def denoise_dense_2_m30():
+  """More masking during training."""
+  hparams = xmoe2_dense_2()
+  hparams.decoder_type = "denoising"
+  hparams.noising_spec_train = {"type": "mask", "prob": 0.3}
+  return hparams
+
+
+@registry.register_hparams
+def denoise_z15():
+  """Replace tokens instead of masking."""
+  hparams = xmoe2_dense_0()
+  hparams.decoder_type = "denoising"
+  hparams.noising_spec_train = {"type": "random_zipfian", "prob": 0.15}
+  hparams.noising_use_eval_during_train = 0.25
+  return hparams
+
+
+@registry.register_hparams
+def denoise_t15():
+  """Noise up with dropout and a little transformer."""
+  hparams = xmoe2_dense_0()
+  hparams.decoder_type = "denoising"
+  hparams.noising_spec_train = {
+      "type": "transformer",
+      "overrides": {
+          "noising_spec_train": {"type": "mask", "prob": 0.15},
+          "noising_use_eval_during_train": 0.0,
+          "decoder_layers": ["att", "drd"] * 4,
+          "num_heads": 4,
+          "d_model": 512,
+          "d_ff": 2048,
+      }
+  }
+  return hparams
+
+
+@registry.register_hparams
+def denoise_v1_m15():
+  """Denoising experiment."""
+  hparams = xmoe2_v1()
+  # no local attention
+  # TODO(noam): non-masked version of local-attention
+  hparams.decoder_layers = [
+      "att" if l == "local_att" else l for l in hparams.decoder_layers]
+  hparams.decoder_type = "denoising"
+  hparams.noising_spec_train = {"type": "mask", "prob": 0.15}
+  return hparams
+
+
+@registry.register_hparams
+def denoise_v1_m30():
+  """More masking during training."""
+  hparams = denoise_v1_m15()
+  hparams.noising_spec_train = {"type": "mask", "prob": 0.3}
+  return hparams
+
+
+@registry.register_hparams
+def denoise_v1_m50():
+  """More masking during training."""
+  hparams = denoise_v1_m15()
+  hparams.noising_spec_train = {"type": "mask", "prob": 0.5}
+  return hparams
+
+
+@registry.register_hparams
+def denoise_v1_z15():
+  """Replace tokens instead of masking."""
+  hparams = denoise_v1_m15()
+  hparams.noising_spec_train = {"type": "random_zipfian", "prob": 0.15}
+  return hparams
+
+
+@registry.register_hparams
+def denoise_v1_t15():
+  """Noise up with dropout and a little transformer."""
+  hparams = denoise_v1_m15()
+  hparams.noising_spec_train = {
+      "type": "transformer",
+      "overrides": {
+          "noising_spec_train": {"type": "mask", "prob": 0.15},
+          "noising_use_eval_during_train": 0.0,
+          "decoder_layers": ["att", "drd"] * 4,
+          "num_heads": 4,
+          "d_model": 512,
+          "d_ff": 2048,
+      }
+  }
+  return hparams

From c67f1c9d921795a91889082edd2eb08552705178 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 16 Nov 2018 22:30:49 -0800
Subject: [PATCH 1253/2720] Makes Travis unhappy for TF 1.12

PiperOrigin-RevId: 221895801
---
 tensor2tensor/layers/common_audio.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/common_audio.py b/tensor2tensor/layers/common_audio.py
index 388ccce53..14432576f 100644
--- a/tensor2tensor/layers/common_audio.py
+++ b/tensor2tensor/layers/common_audio.py
@@ -56,7 +56,7 @@ def compute_mel_filterbank_features(
     waveforms,
     sample_rate=16000, dither=1.0 / np.iinfo(np.int16).max, preemphasis=0.97,
     frame_length=25, frame_step=10, fft_length=None,
-    window_fn=functools.partial(tf.signal.hann_window, periodic=True),
+    window_fn=functools.partial(tf.contrib.signal.hann_window, periodic=True),
     lower_edge_hertz=80.0, upper_edge_hertz=7600.0, num_mel_bins=80,
     log_noise_floor=1e-3, apply_mask=True):
   """Implement mel-filterbank extraction using tf ops.
@@ -101,7 +101,7 @@ def compute_mel_filterbank_features(
   if fft_length is None:
     fft_length = int(2**(np.ceil(np.log2(frame_length))))
 
-  stfts = tf.signal.stft(
+  stfts = tf.contrib.signal.stft(
       waveforms,
       frame_length=frame_length,
       frame_step=frame_step,
@@ -121,7 +121,7 @@ def compute_mel_filterbank_features(
   # Warp the linear-scale, magnitude spectrograms into the mel-scale.
   num_spectrogram_bins = magnitude_spectrograms.shape[-1].value
   linear_to_mel_weight_matrix = (
-      tf.signal.linear_to_mel_weight_matrix(
+      tf.contrib.signal.linear_to_mel_weight_matrix(
           num_mel_bins, num_spectrogram_bins, sample_rate, lower_edge_hertz,
           upper_edge_hertz))
   mel_spectrograms = tf.tensordot(

From ef37f597507b91a2b389c856a6d61facd52b6bba Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Mon, 19 Nov 2018 11:03:06 -0800
Subject: [PATCH 1254/2720] Add a hparam and flag mlperf_mode to disable mlperf
 logging

PiperOrigin-RevId: 222104368
---
 tensor2tensor/bin/t2t_trainer.py       | 12 +++++++-----
 tensor2tensor/layers/common_hparams.py |  2 ++
 tensor2tensor/models/transformer.py    | 22 +++++++++++++++-------
 tensor2tensor/utils/decoding.py        |  1 -
 tensor2tensor/utils/mlperf_log.py      |  5 ++++-
 tensor2tensor/utils/optimize.py        | 16 +++++++++++-----
 tensor2tensor/utils/t2t_model.py       |  3 ++-
 7 files changed, 41 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 83cd93c44..fe9c8ac03 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -75,7 +75,6 @@
     "Whether to use TensorFlow DistributionStrategy instead of explicitly "
     "replicating the model. DistributionStrategy is used only if the "
     "model replication configuration is supported by the DistributionStrategy.")
-
 # To maintain compatibility with some internal libs, we guard against these flag
 # definitions possibly erroring. Apologies for the ugliness.
 try:
@@ -356,12 +355,15 @@ def run_std_server():
 
 def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
+  hparams = create_hparams()
   if FLAGS.schedule == "train" or FLAGS.schedule == "train_eval_and_decode":
-    mlperf_log.transformer_print(key=mlperf_log.RUN_START)
+    mlperf_log.transformer_print(key=mlperf_log.RUN_START,
+                                 mlperf_mode=hparams.mlperf_mode)
   if FLAGS.schedule == "run_std_server":
     run_std_server()
   mlperf_log.transformer_print(
-      key=mlperf_log.RUN_SET_RANDOM_SEED, value=FLAGS.random_seed)
+      key=mlperf_log.RUN_SET_RANDOM_SEED, value=FLAGS.random_seed,
+      mlperf_mode=hparams.mlperf_mode)
   trainer_lib.set_random_seed(FLAGS.random_seed)
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
   maybe_log_registry_and_exit()
@@ -378,7 +380,6 @@ def main(argv):
 
   if argv:
     set_hparams_from_args(argv[1:])
-  hparams = create_hparams()
 
   exp_fn = create_experiment_fn()
   exp = exp_fn(create_run_config(hparams), hparams)
@@ -386,7 +387,8 @@ def main(argv):
     save_metadata(hparams)
   execute_schedule(exp)
   if FLAGS.schedule != "train":
-    mlperf_log.transformer_print(key=mlperf_log.RUN_FINAL)
+    mlperf_log.transformer_print(key=mlperf_log.RUN_FINAL,
+                                 mlperf_mode=hparams.mlperf_mode)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 9a94207b3..28a281c6e 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -48,6 +48,8 @@ def basic_params1():
       clip_grad_norm=2.0,
       grad_noise_scale=0.0,
       summarize_grads=False,
+      # Flag for whether mlperf mode is on
+      mlperf_mode=False,
       # Whether to log the name and size of every variable
       summarize_vars=False,
       initializer="orthogonal",
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index d00de8dca..0151abed7 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -86,7 +86,8 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
 
     mlperf_log.transformer_print(
         key=mlperf_log.MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
-        value=hparams.layer_prepostprocess_dropout)
+        value=hparams.layer_prepostprocess_dropout,
+        mlperf_mode=hparams.mlperf_mode)
 
     encoder_input = tf.nn.dropout(encoder_input,
                                   1.0 - hparams.layer_prepostprocess_dropout)
@@ -136,7 +137,8 @@ def decode(self,
     """
     mlperf_log.transformer_print(
         key=mlperf_log.MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
-        value=hparams.layer_prepostprocess_dropout)
+        value=hparams.layer_prepostprocess_dropout,
+        mlperf_mode=hparams.mlperf_mode)
     decoder_input = tf.nn.dropout(decoder_input,
                                   1.0 - hparams.layer_prepostprocess_dropout)
 
@@ -802,7 +804,8 @@ def fast_decode_tpu(encoder_output,
           "beam_size": beam_size,
           "alpha": alpha,
           "max_decode_length": decode_length
-      })
+      },
+      mlperf_mode=hparams.mlperf_mode)
   if beam_size > 1:  # Beam Search
     initial_ids = sos_id * tf.ones([batch_size], dtype=tf.int32)
     decoded_ids, scores = beam_search.beam_search(
@@ -1222,17 +1225,20 @@ def transformer_decoder(decoder_input,
 
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_NUM_HIDDEN_LAYERS,
-      value=hparams.num_decoder_layers or hparams.num_hidden_layers)
+      value=hparams.num_decoder_layers or hparams.num_hidden_layers,
+      mlperf_mode=hparams.mlperf_mode)
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_ATTENTION_DROPOUT,
-      value=hparams.attention_dropout)
+      value=hparams.attention_dropout,
+      mlperf_mode=hparams.mlperf_mode)
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_ATTENTION_DENSE,
       value={
           "use_bias": "false",
           "num_heads": hparams.num_heads,
           "hidden_size": hparams.hidden_size
-      })
+      },
+      mlperf_mode=hparams.mlperf_mode)
 
   with tf.variable_scope(name):
     for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
@@ -1299,7 +1305,8 @@ def transformer_decoder(decoder_input,
     # a whole stack of unnormalized layer outputs.
     mlperf_log.transformer_print(
         key=mlperf_log.MODEL_HP_NORM,
-        value={"hidden_size": hparams.hidden_size})
+        value={"hidden_size": hparams.hidden_size},
+        mlperf_mode=hparams.mlperf_mode)
     return common_layers.layer_preprocess(x, hparams)
 
 
@@ -1964,6 +1971,7 @@ def transformer_timeseries():
 def transformer_mlperf_tpu():
   """HParams for Transformer model on TPU for MLPerf on TPU 2x2."""
   hparams = transformer_base_v3()
+  hparams.mlperf_mode = True
   hparams.symbol_modality_num_shards = 1
   hparams.max_length = 256  # ignored when using "_packed" problems
   hparams.batch_size = 2048  # per-chip batch size matches the reference model
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index ccd80a803..d77dd1f06 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -78,7 +78,6 @@ def decode_hparams(overrides=""):
       # Total number of videos are max_display_outputs * num_decodes
       max_display_outputs=10,
       # Used for MLPerf compliance logging.
-      mlperf_mode=False,
       mlperf_decode_step=0.0,
       mlperf_threshold=25.0,
       mlperf_success=False)
diff --git a/tensor2tensor/utils/mlperf_log.py b/tensor2tensor/utils/mlperf_log.py
index 94af2eb31..23b4bdf74 100644
--- a/tensor2tensor/utils/mlperf_log.py
+++ b/tensor2tensor/utils/mlperf_log.py
@@ -161,7 +161,10 @@ def _mlperf_print(key, value=None, benchmark=None, stack_offset=0,
 TRANSFORMER_TAG_SET = set(TRANSFORMER_TAGS)  # pylint: disable=undefined-variable
 
 
-def transformer_print(key, value=None, stack_offset=2, deferred=False):
+def transformer_print(key, value=None, stack_offset=2, deferred=False,
+                      mlperf_mode=False):
+  if not mlperf_mode:
+    return
   return _mlperf_print(
       key=key,
       value=value,
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index e948d2778..3b393d159 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -88,14 +88,19 @@ class ConditionalOptimizer(tf.train.Optimizer):
   def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disable=super-init-not-called
     tf.logging.info("Using optimizer %s", optimizer_name)
 
-    mlperf_log.transformer_print(key=mlperf_log.OPT_NAME, value=optimizer_name)
+    mlperf_log.transformer_print(key=mlperf_log.OPT_NAME,
+                                 value=optimizer_name,
+                                 mlperf_mode=hparams.mlperf_mode)
     mlperf_log.transformer_print(
-        key=mlperf_log.OPT_HP_ADAM_BETA1, value=hparams.optimizer_adam_beta1)
+        key=mlperf_log.OPT_HP_ADAM_BETA1, value=hparams.optimizer_adam_beta1,
+        mlperf_mode=hparams.mlperf_mode)
     mlperf_log.transformer_print(
-        key=mlperf_log.OPT_HP_ADAM_BETA2, value=hparams.optimizer_adam_beta2)
+        key=mlperf_log.OPT_HP_ADAM_BETA2, value=hparams.optimizer_adam_beta2,
+        mlperf_mode=hparams.mlperf_mode)
     mlperf_log.transformer_print(
         key=mlperf_log.OPT_HP_ADAM_EPSILON,
-        value=hparams.optimizer_adam_epsilon)
+        value=hparams.optimizer_adam_epsilon,
+        mlperf_mode=hparams.mlperf_mode)
 
     if optimizer_name == "Adam":
       # We change the default epsilon for Adam.
@@ -275,7 +280,8 @@ def get_variable_initializer(hparams):
     return None
 
   mlperf_log.transformer_print(key=mlperf_log.MODEL_HP_INITIALIZER_GAIN,
-                               value=hparams.initializer_gain)
+                               value=hparams.initializer_gain,
+                               mlperf_mode=hparams.mlperf_mode)
 
   if not tf.contrib.eager.in_eager_mode():
     tf.logging.info("Using variable initializer: %s", hparams.initializer)
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index d128ddbd1..86f16afcd 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -141,7 +141,8 @@ def __init__(self,
             value={
                 "vocab_size": target_modality.top_dimensionality,
                 "hidden_size": hidden_size
-            })
+            },
+            mlperf_mode=hparams.mlperf_mode)
 
     self._original_hparams = hparams
     self.set_mode(mode)

From 15786b688441d2ec533373f2edd4f34766520262 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 19 Nov 2018 11:35:57 -0800
Subject: [PATCH 1255/2720] Adds additive coupling layer. An additive coupling
 layer is a special case of the affine coupling layer with the scale set to
 1.0.

PiperOrigin-RevId: 222110239
---
 tensor2tensor/models/research/glow.py         |  8 +++-
 tensor2tensor/models/research/glow_ops.py     | 47 ++++++++++++++++---
 .../models/research/glow_ops_test.py          | 18 +++----
 3 files changed, 56 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 5b71da058..c3c01192b 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -49,7 +49,9 @@ def glow_hparams():
   hparams.add_hparam("n_levels", 3)
   hparams.add_hparam("n_bits_x", 8)
   hparams.add_hparam("depth", 32)
-  hparams.add_hparam("affine_coupling_width", 512)
+  # Coupling layer, additive or affine.
+  hparams.add_hparam("coupling", "affine")
+  hparams.add_hparam("coupling_width", 512)
   hparams.add_hparam("top_prior", "single_conv")
   # init_batch_size denotes the number of examples used for data-dependent
   # initialization. A higher init_batch_size is required for training
@@ -152,6 +154,10 @@ def top_prior(self):
         temperature=self.temperature)
 
   def body(self, features):
+    exp_coupling = ["affine", "additive"]
+    if self.hparams.coupling not in exp_coupling:
+      raise ValueError("Expected hparams.coupling to be in %s, got %s" %
+                       (exp_coupling, self.hparams.coupling))
     if self.is_training:
       init_features = self.create_init_batch(features)
       init_op = self.objective_tower(init_features, init=True)
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index ab0089058..1a029fb74 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -542,14 +542,41 @@ def conv_stack(name, x, mid_channels, output_channels, dilations=None):
   return x
 
 
+@add_arg_scope
+def additive_coupling(name, x, mid_channels=512, reverse=False):
+  """Reversible additive coupling layer.
+
+  Args:
+    name: variable scope.
+    x: 4-D Tensor.
+    mid_channels: number of channels in the coupling layer.
+    reverse: Forward or reverse operation.
+  Returns:
+    output:
+    objective: 0.0
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    output_channels = common_layers.shape_list(x)[-1] // 2
+    x1, x2 = tf.split(x, num_or_size_splits=2, axis=-1)
+
+    z1 = x1
+    shift = conv_stack("nn", x1, mid_channels, output_channels=output_channels)
+
+    if not reverse:
+      z2 = x2 + shift
+    else:
+      z2 = x2 - shift
+    return tf.concat([z1, z2], axis=3), 0.0
+
+
 @add_arg_scope
 def affine_coupling(name, x, mid_channels=512, reverse=False):
   """Reversible affine coupling layer.
 
   Args:
-    name:
-    x:
-    mid_channels: intermediate
+    name: variable scope.
+    x: 4-D Tensor.
+    mid_channels: number of channels in the coupling layer.
     reverse: Forward or reverse operation.
   Returns:
     output:
@@ -969,19 +996,25 @@ def revnet_step(name, x, hparams, reverse=True):
   Args:
     name: used for variable scope.
     x: input
-    hparams: affine_coupling_width is the only hparam that is being used in
+    hparams: coupling_width is the only hparam that is being used in
              this function.
     reverse: forward or reverse pass.
   Returns:
     z: Output of one step of reversible flow.
   """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    if hparams.coupling == "additive":
+      coupling_layer = functools.partial(
+          additive_coupling, name="additive", reverse=reverse,
+          mid_channels=hparams.coupling_width)
+    else:
+      coupling_layer = functools.partial(
+          affine_coupling, name="affine", reverse=reverse,
+          mid_channels=hparams.coupling_width)
     ops = [
         functools.partial(actnorm, name="actnorm", reverse=reverse),
         functools.partial(invertible_1x1_conv, name="invertible",
-                          reverse=reverse),
-        functools.partial(affine_coupling, name="affine", reverse=reverse,
-                          mid_channels=hparams.affine_coupling_width)]
+                          reverse=reverse), coupling_layer]
 
     if reverse:
       ops = ops[::-1]
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 10b5c1f0a..ef71d6fdb 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -86,8 +86,8 @@ def check_invertibility(self, op, name):
 
   def test_invertibility(self):
     rev_ops = [glow_ops.invertible_1x1_conv, glow_ops.affine_coupling,
-               glow_ops.actnorm]
-    names = ["inv_1X1_conv", "affine_coupling", "actnorm"]
+               glow_ops.actnorm, glow_ops.additive_coupling]
+    names = ["inv_1X1_conv", "affine_coupling", "actnorm", "additive_coupling"]
     for rev_op, name in zip(rev_ops, names):
       self.check_invertibility(rev_op, name)
 
@@ -171,10 +171,16 @@ def test_split(self):
         self.assertEqual(x_inv_np.shape, (16, 5, 5, 16))
         self.assertTrue(np.allclose(diff, 0.0, atol=1e-5))
 
-  def check_revnet_reversibility(self, op, name):
+  @parameterized.named_parameters(
+      ("aff_revnet", glow_ops.revnet, "aff_rev", "affine"),
+      ("add_revnet", glow_ops.revnet, "add_rev", "additive"),
+      ("aff_rev_step", glow_ops.revnet_step, "aff_rev_step", "affine"),
+      ("add_rev_step", glow_ops.revnet_step, "add_rev_step", "additive"),)
+  def test_revnet_reversibility(self, op, name, coupling):
     with tf.Graph().as_default():
       hparams = glow.glow_hparams()
       hparams.depth = 2
+      hparams.coupling = coupling
       x = tf.random_uniform(shape=(16, 32, 32, 4), seed=0)
       x_inv, _ = op(name, x, hparams, reverse=False)
       x_inv_inv, _ = op(name, x_inv, hparams, reverse=True)
@@ -183,12 +189,6 @@ def check_revnet_reversibility(self, op, name):
         diff = session.run(x - x_inv_inv)
         self.assertTrue(np.allclose(diff, 0.0, atol=1e-2))
 
-  def test_revnet_reversibility(self):
-    ops = [glow_ops.revnet_step, glow_ops.revnet]
-    names = ["revnet_step", "revnet"]
-    for op, name in zip(ops, names):
-      self.check_revnet_reversibility(op, name)
-
   def test_encoder_decoder(self):
     with tf.Graph().as_default():
       hparams = glow.glow_hparams()

From b5807a502d0c21265bb499778a8974a4006653da Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 19 Nov 2018 14:13:02 -0800
Subject: [PATCH 1256/2720] Add BayesianLinearModel.

PiperOrigin-RevId: 222136695
---
 tensor2tensor/layers/bayes.py      | 59 ++++++++++++++++++++++++++++++
 tensor2tensor/layers/bayes_test.py | 31 ++++++++++++++++
 2 files changed, 90 insertions(+)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 2a35adea9..cd978346f 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -455,6 +455,65 @@ def loss_fn():
   return loss_fn
 
 
+class BayesianLinearModel(tf.keras.Model):
+  r"""Bayesian linear model with standard normal prior over its coefficients.
+
+  A forward pass computes the mean of the exact predictive distribution
+
+  ```none
+  p(outputs | inputs) = \int Normal(outputs | coeffs * inputs, noise_variance)
+                             Normal(coeffs | 0, 1) dweights dbias.
+  ```
+
+  It takes a Tensor of shape [batch_size, input_dim] as input and returns a
+  Normal random variable of shape [batch_size] representing its outputs.
+  After `fit()`, the forward pass computes the exact posterior predictive
+  distribution.
+  """
+
+  def __init__(self, noise_variance, **kwargs):
+    super(BayesianLinearModel, self).__init__(**kwargs)
+    self.noise_variance = noise_variance
+    self.coeffs_precision_tril_op = None
+    self.coeffs_mean = None
+
+  def call(self, inputs):
+    if self.coeffs_mean is None and self.coeffs_precision_tril_op is None:
+      # p(mean(ynew) | xnew) = Normal(ynew | mean = 0, variance = xnew xnew^T)
+      predictive_mean = 0.
+      predictive_variance = tf.reduce_sum(tf.square(inputs), -1)
+    else:
+      # p(mean(ynew) | xnew, x, y) = Normal(ynew |
+      #   mean = xnew (1/noise_variance) (1/noise_variance x^T x + I)^{-1}x^T y,
+      #   variance = xnew (1/noise_variance x^T x + I)^{-1} xnew^T)
+      predictive_mean = tf.einsum('nm,m->n', inputs, self.coeffs_mean)
+      predictive_covariance = tf.matmul(
+          inputs,
+          self.coeffs_precision_tril_op.solve(
+              self.coeffs_precision_tril_op.solve(inputs, adjoint_arg=True),
+              adjoint=True))
+      predictive_variance = tf.diag_part(predictive_covariance)
+    return ed.Normal(loc=predictive_mean, scale=tf.sqrt(predictive_variance))
+
+  def fit(self, x=None, y=None):
+    # p(coeffs | x, y) = Normal(coeffs |
+    #   mean = (1/noise_variance) (1/noise_variance x^T x + I)^{-1} x^T y,
+    #   covariance = (1/noise_variance x^T x + I)^{-1})
+    # TODO(trandustin): We newly fit the data at each call. Extend to do
+    # Bayesian updating.
+    kernel_matrix = tf.matmul(x, x, transpose_a=True) / self.noise_variance
+    coeffs_precision = tf.matrix_set_diag(
+        kernel_matrix, tf.matrix_diag_part(kernel_matrix) + 1.)
+    coeffs_precision_tril = tf.linalg.cholesky(coeffs_precision)
+    self.coeffs_precision_tril_op = tf.linalg.LinearOperatorLowerTriangular(
+        coeffs_precision_tril)
+    self.coeffs_mean = self.coeffs_precision_tril_op.solvevec(
+        self.coeffs_precision_tril_op.solvevec(tf.einsum('nm,n->m', x, y)),
+        adjoint=True) / self.noise_variance
+    # TODO(trandustin): To be fully Keras-compatible, return History object.
+    return
+
+
 class MixtureLogistic(tf.keras.layers.Layer):
   """Stochastic output layer, distributed as a mixture of logistics."""
 
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index eff85ba53..e53661992 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -239,6 +239,37 @@ def testLSTMCellReparameterizationModel(self):
     self.assertAllClose(res2, res3)
     self.assertLen(model.losses, 2)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testBayesianLinearModel(self):
+    """Tests that model makes reasonable predictions."""
+    np.random.seed(42)
+    train_batch_size = 5
+    test_batch_size = 2
+    num_features = 3
+    noise_variance = 0.01
+    coeffs = tf.range(num_features, dtype=tf.float32)
+    features = tf.to_float(np.random.randn(train_batch_size, num_features))
+    labels = (tf.tensordot(features, coeffs, [[-1], [0]])
+              + noise_variance * tf.to_float(np.random.randn(train_batch_size)))
+
+    model = bayes.BayesianLinearModel(noise_variance=noise_variance)
+    model.fit(features, labels)
+
+    test_features = tf.to_float(np.random.randn(test_batch_size, num_features))
+    test_labels = tf.tensordot(test_features, coeffs, [[-1], [0]])
+    outputs = model(test_features)
+    test_predictions = outputs.distribution.mean()
+    test_predictions_variance = outputs.distribution.variance()
+
+    [
+        test_labels_val, test_predictions_val, test_predictions_variance_val,
+    ] = self.evaluate(
+        [test_labels, test_predictions, test_predictions_variance])
+    self.assertEqual(test_predictions_val.shape, (test_batch_size,))
+    self.assertEqual(test_predictions_variance_val.shape, (test_batch_size,))
+    self.assertAllClose(test_predictions_val, test_labels_val, atol=0.1)
+    self.assertAllLessEqual(test_predictions_variance_val, noise_variance)
+
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testMixtureLogistic(self):
     batch_size = 3

From f3dddcaf572210f4aa15c8a72423500ffc94849e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 19 Nov 2018 15:19:39 -0800
Subject: [PATCH 1257/2720] Update docstring.

PiperOrigin-RevId: 222148068
---
 tensor2tensor/visualization/attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/visualization/attention.py b/tensor2tensor/visualization/attention.py
index 56ece8154..f4f78f3f6 100644
--- a/tensor2tensor/visualization/attention.py
+++ b/tensor2tensor/visualization/attention.py
@@ -86,7 +86,7 @@ def _get_attention(inp_text, out_text, enc_atts, dec_atts, encdec_atts):
     dec_atts: numpy array, decoder self-attentions
         [num_layers, batch_size, num_heads, dec_length, dec_length]
     encdec_atts: numpy array, encoder-decoder attentions
-        [num_layers, batch_size, num_heads, enc_length, dec_length]
+        [num_layers, batch_size, num_heads, dec_length, enc_length]
 
   Returns:
     Dictionary of attention representations with the structure:

From eb546689116ed37ac51bd3dbd30b78337bc20236 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 19 Nov 2018 16:05:26 -0800
Subject: [PATCH 1258/2720] Add masked perplexity metric to Tensor2Tensor.

PiperOrigin-RevId: 222156157
---
 tensor2tensor/utils/metrics.py      | 14 ++++++++++++
 tensor2tensor/utils/metrics_test.py | 33 +++++++++++++++++++++++++++++
 tensor2tensor/utils/t2t_model.py    | 31 ++++++++++++++++++---------
 3 files changed, 68 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index d3fe98b64..8d93ec391 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -39,6 +39,7 @@ class Metrics(object):
   ACC_PER_SEQ = "accuracy_per_sequence"
   ACC_MULTILABEL_MATCH3 = "accuracy_multilabel_match3"
   NEG_LOG_PERPLEXITY = "neg_log_perplexity"
+  MASKED_NEG_LOG_PERPLEXITY = "masked_neg_log_perplexity"
   APPROX_BLEU = "approx_bleu_score"
   RMSE = "rmse"
   LOG_POISSON = "log_poisson"
@@ -244,6 +245,18 @@ def padded_neg_log_perplexity(predictions,
   return (-num, den)
 
 
+def padded_neg_log_perplexity_with_masking(
+    predictions,
+    labels,
+    features,
+    weights_fn=None):
+  del weights_fn
+  if "target_mask" not in features:
+    raise ValueError("masked_neg_log_perplexity requires target_mask feature")
+  mask_fn = lambda labels: features["target_mask"]
+  return padded_neg_log_perplexity(predictions, labels, mask_fn)
+
+
 def dmol_neg_log_perplexity(predictions,
                             labels,
                             weights_fn=None):
@@ -659,6 +672,7 @@ def metric_means():
     Metrics.ACC_PER_SEQ: padded_sequence_accuracy,
     Metrics.ACC_MULTILABEL_MATCH3: multilabel_accuracy_match3,
     Metrics.NEG_LOG_PERPLEXITY: padded_neg_log_perplexity,
+    Metrics.MASKED_NEG_LOG_PERPLEXITY: padded_neg_log_perplexity_with_masking,
     Metrics.APPROX_BLEU: bleu_hook.bleu_score,
     Metrics.RMSE: padded_rmse,
     Metrics.LOG_POISSON: padded_log_poisson,
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 7878d784a..92861e0ba 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -117,6 +117,39 @@ def testNegativeLogPerplexity(self):
       actual = session.run(a)
     self.assertEqual(actual.shape, ())
 
+  def testNegativeLogPerplexityMasked(self):
+    predictions = np.random.randint(4, size=(12, 12, 12, 1))
+    targets = np.random.randint(4, size=(12, 12, 12, 1))
+    features = {
+        'target_mask': tf.to_float(tf.not_equal(targets, 0))
+    }
+    with self.test_session() as session:
+      scores, _ = metrics.padded_neg_log_perplexity_with_masking(
+          tf.one_hot(predictions, depth=4, dtype=tf.float32),
+          tf.constant(targets, dtype=tf.int32),
+          features)
+      a = tf.reduce_mean(scores)
+      session.run(tf.global_variables_initializer())
+      actual = session.run(a)
+    self.assertEqual(actual.shape, ())
+
+  def testNegativeLogPerplexityMaskedAssert(self):
+    predictions = np.random.randint(4, size=(12, 12, 12, 1))
+    targets = np.random.randint(4, size=(12, 12, 12, 1))
+    features = {}
+
+    with self.assertRaisesRegexp(
+        ValueError,
+        'masked_neg_log_perplexity requires target_mask feature'):
+      with self.test_session() as session:
+        scores, _ = metrics.padded_neg_log_perplexity_with_masking(
+            tf.one_hot(predictions, depth=4, dtype=tf.float32),
+            tf.constant(targets, dtype=tf.int32),
+            features)
+        a = tf.reduce_mean(scores)
+        session.run(tf.global_variables_initializer())
+        _ = session.run(a)
+
   def testSigmoidAccuracyOneHot(self):
     logits = np.array([
         [-1., 1.],
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 86f16afcd..3f3e45ee5 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -22,6 +22,7 @@
 import contextlib
 import copy
 import functools
+import inspect
 import math
 import time
 import six
@@ -1456,6 +1457,7 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
         # For TPU, logits dict will be passed as keyword arguments to
         # eval_metrics_fn. Here we add the labels to those arguments.
         logits.update({"labels": labels})
+        logits.update({"features": features})
         return tf.contrib.tpu.TPUEstimatorSpec(
             tf.estimator.ModeKeys.EVAL,
             eval_metrics=(eval_metrics_fn, logits),
@@ -1464,7 +1466,7 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
         eval_metrics_fn = create_tpu_eval_metrics_fn(problem, hparams)
         return tf.contrib.tpu.TPUEstimatorSpec(
             tf.estimator.ModeKeys.EVAL,
-            eval_metrics=(eval_metrics_fn, [logits, labels]),
+            eval_metrics=(eval_metrics_fn, [logits, labels, features]),
             loss=loss)
     else:
       task_list = [problem]
@@ -1643,9 +1645,12 @@ def create_tpu_eval_metrics_fn(problem, model_hparams):
       weights_fn = v.targets_weights_fn
 
       def make_metric_fn(metric_fn):
-
-        def wrapped_metric_fn(logits, labels, weights_fn=weights_fn):
-          num, den = metric_fn(logits, labels, weights_fn=weights_fn)
+        def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
+          kwargs = {}
+          args, _, keywords, _ = inspect.getargspec(metric_fn)
+          if ("features" in args) or keywords:
+            kwargs["features"] = features
+          num, den = metric_fn(logits, labels, weights_fn=weights_fn, **kwargs)
           return tf.metrics.mean(num, den)
 
         return wrapped_metric_fn
@@ -1660,9 +1665,12 @@ def wrapped_metric_fn(logits, labels, weights_fn=weights_fn):
     weights_fn = tm.targets_weights_fn
 
     def make_metric_fn(metric_fn):
-
-      def wrapped_metric_fn(logits, labels):
-        num, den = metric_fn(logits, labels, weights_fn=weights_fn)
+      def wrapped_metric_fn(logits, labels, features):
+        kwargs = {}
+        args, _, keywords, _ = inspect.getargspec(metric_fn)
+        if ("features" in args) or keywords:
+          kwargs["features"] = features
+        num, den = metric_fn(logits, labels, weights_fn=weights_fn, **kwargs)
         return tf.metrics.mean(num, den)
 
       return wrapped_metric_fn
@@ -1680,18 +1688,21 @@ def all_metrics_fn(logits=None, labels=None, **kwargs):
 
     if logits is None:
       logits = kwargs
+      features = logits["features"]
+    else:
+      features = kwargs["features"]
 
     for name, fn in metric_fns:
       if isinstance(logits, dict) and isinstance(labels, dict):
         for k, v in six.iteritems(logits):
-          metrics_dict["%s/%s" % (k, name)] = fn(v, labels[k])
+          metrics_dict["%s/%s" % (k, name)] = fn(v, labels[k], features)
       elif isinstance(logits, dict):
         tf.logging.warning("Logits is a dict, but labels is not; only "
                            "evaluating logits['targets'] against labels.")
         metrics_dict["%s/%s" % ("targets", name)] = fn(logits["targets"],
-                                                       labels)
+                                                       labels, features)
       else:
-        metrics_dict[name] = fn(logits, labels)
+        metrics_dict[name] = fn(logits, labels, features)
 
     return metrics_dict
 

From e2ea4529054186c7519c7867d52fca2d7e0e758d Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Mon, 19 Nov 2018 18:08:04 -0800
Subject: [PATCH 1259/2720] fixing state save bug in training.

PiperOrigin-RevId: 222171991
---
 tensor2tensor/models/video/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 397c26166..758073c7b 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -646,7 +646,7 @@ def reset_func():
         reset = tf.greater(tf.reduce_sum(reset), 0.5)
         reset_ops = tf.cond(reset, reset_func, tf.no_op)
       else:
-        reset_ops = reset_func()
+        reset_ops = tf.no_op()
       with tf.control_dependencies([reset_ops]):
         frames[0] = tf.identity(frames[0])
 

From 424df5e16068883999e62e7d098e28dce7fe7df2 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 19 Nov 2018 18:40:06 -0800
Subject: [PATCH 1260/2720] Minor, Allow child_dir as a hparam to t2t-decoder.

PiperOrigin-RevId: 222174884
---
 tensor2tensor/utils/decoding.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index d77dd1f06..b45d32d08 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -63,6 +63,7 @@ def decode_hparams(overrides=""):
       delimiter="\n",
       decode_to_file=None,
       decode_in_memory=False,
+      summaries_log_dir="decode",  # Directory to write hook summaries.
       shards=1,    # How many shards of data to decode (treating 1 as None).
       shard_id=0,  # Which shard are we decoding if more than 1 above.
       shards_start_offset=0,  # Number of the first shard to decode.
@@ -858,7 +859,7 @@ def run_postdecode_hooks(decode_hook_args, dataset_split):
     return
   tf.logging.info("Running decode hooks.")
   parent_dir = os.path.join(decode_hook_args.output_dirs[0], os.pardir)
-  child_dir = "decode"
+  child_dir = decode_hook_args.decode_hparams.summaries_log_dir
   if dataset_split is not None:
     child_dir += "_{}".format(dataset_split)
   final_dir = os.path.join(parent_dir, child_dir)

From 07801eab4d66404111f32a99e1a49d2c76c33363 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Mon, 19 Nov 2018 19:00:56 -0800
Subject: [PATCH 1261/2720] Missed mlperf_mode in decoding

PiperOrigin-RevId: 222176508
---
 tensor2tensor/utils/decoding.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index b45d32d08..f4fbe9eb5 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -332,7 +332,9 @@ def decode_once(estimator,
         num_predictions >= decode_hp.num_samples):
       break
 
-  mlperf_log.transformer_print(key=mlperf_log.EVAL_SIZE, value=num_eval_samples)
+  mlperf_log.transformer_print(key=mlperf_log.EVAL_SIZE,
+                               value=num_eval_samples,
+                               mlperf_mode=hparams.mlperf_mode)
 
   if decode_to_file:
     output_file.close()

From 091294bb62852d4d590fe58d5983048f87e9a51e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 19 Nov 2018 20:55:45 -0800
Subject: [PATCH 1262/2720] SVG with learned prior implementation

PiperOrigin-RevId: 222184830
---
 tensor2tensor/layers/common_video.py   |   7 +-
 tensor2tensor/models/__init__.py       |   1 +
 tensor2tensor/models/video/base_vae.py |   1 -
 tensor2tensor/models/video/emily.py    | 119 +++++---
 tensor2tensor/models/video/svg_lp.py   | 366 +++++++++++++++++++++++++
 5 files changed, 457 insertions(+), 37 deletions(-)
 create mode 100644 tensor2tensor/models/video/svg_lp.py

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 646a50ec5..8fc578d1a 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -299,7 +299,8 @@ def vgg_layer(inputs,
               kernel_size=3,
               activation=tf.nn.leaky_relu,
               padding="SAME",
-              is_training=False,
+              is_training=True,
+              has_batchnorm=False,
               scope=None):
   """A layer of VGG network with batch norm.
 
@@ -310,6 +311,7 @@ def vgg_layer(inputs,
     activation: activation function
     padding: padding of the image
     is_training: whether it is training mode or not
+    has_batchnorm: whether batchnorm is applied or not
     scope: variable scope of the op
   Returns:
     net: output of layer
@@ -317,7 +319,8 @@ def vgg_layer(inputs,
   with tf.variable_scope(scope):
     net = tfl.conv2d(inputs, nout, kernel_size=kernel_size, padding=padding,
                      activation=None, name="conv")
-    net = tfl.batch_normalization(net, training=is_training, name="bn")
+    if has_batchnorm:
+      net = tfl.batch_normalization(net, training=is_training, name="bn")
     net = activation(net)
   return net
 
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index f806d0b05..05cafbde7 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -69,6 +69,7 @@
 from tensor2tensor.models.video import epva
 from tensor2tensor.models.video import savp
 from tensor2tensor.models.video import sv2p
+from tensor2tensor.models.video import svg_lp
 
 from tensor2tensor.utils import registry
 
diff --git a/tensor2tensor/models/video/base_vae.py b/tensor2tensor/models/video/base_vae.py
index b049b6513..d4c25b095 100644
--- a/tensor2tensor/models/video/base_vae.py
+++ b/tensor2tensor/models/video/base_vae.py
@@ -81,7 +81,6 @@ def get_kl_loss(self, means, log_vars, means_p=None, log_vars_p=None):
     enumerated_inputs = enumerate(zip(means, log_vars, means_p, log_vars_p))
     if self.is_training and self.hparams.stochastic_model:
       for i, (mean, log_var, mean_p, log_var_p) in enumerated_inputs:
-        # Condition to compute kl divergence with learned prior
         kl_loss += common_layers.kl_divergence(mean, log_var, mean_p, log_var_p)
         tf.summary.histogram("posterior_mean_%d" % i, mean)
         tf.summary.histogram("posterior_log_var_%d" % i, log_var)
diff --git a/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index 3c41d276e..877f3cdb8 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -44,12 +44,13 @@
 class NextFrameEmily(sv2p.NextFrameSv2pLegacy):
   """Stochastic Variational Video Prediction Without Learned Prior."""
 
-  def encoder(self, inputs, nout):
+  def encoder(self, inputs, nout, has_batchnorm=True):
     """VGG based image encoder.
 
     Args:
       inputs: image tensor with size BSx64x64xC
       nout: number of output channels
+      has_batchnorm: variable to use or not use batch normalization
     Returns:
       net: encoded image with size BSxNout
       skips: skip connection after each layer
@@ -72,7 +73,8 @@ def encoder(self, inputs, nout):
     ds_idx = 0
     while res_x > 64:
       h = tfcl.repeat(net01, 2, vgg_layer, 64, scope="downscale%d" % ds_idx,
-                      is_training=self.is_training)
+                      is_training=self.is_training, activation=tf.nn.relu,
+                      has_batchnorm=has_batchnorm)
       net01 = tfl.max_pooling2d(h, [2, 2], strides=(2, 2),
                                 name="downscale%d_pool" % ds_idx)
       skips.append(h)
@@ -81,34 +83,41 @@ def encoder(self, inputs, nout):
 
     # h1
     net11 = tfcl.repeat(net01, 2, vgg_layer, 64,
-                        scope="h1", is_training=self.is_training)
+                        scope="h1", is_training=self.is_training,
+                        activation=tf.nn.relu, has_batchnorm=has_batchnorm)
     net12 = tfl.max_pooling2d(net11, [2, 2], strides=(2, 2), name="h1_pool")
     # h2
     net21 = tfcl.repeat(net12, 2, vgg_layer, 128,
-                        scope="h2", is_training=self.is_training)
+                        scope="h2", is_training=self.is_training,
+                        activation=tf.nn.relu, has_batchnorm=has_batchnorm)
     net22 = tfl.max_pooling2d(net21, [2, 2], strides=(2, 2), name="h2_pool")
     # h3
     net31 = tfcl.repeat(net22, 3, vgg_layer, 256,
-                        scope="h3", is_training=self.is_training)
+                        scope="h3", is_training=self.is_training,
+                        activation=tf.nn.relu, has_batchnorm=has_batchnorm)
     net32 = tfl.max_pooling2d(net31, [2, 2], strides=(2, 2), name="h3_pool")
     # h4
     net41 = tfcl.repeat(net32, 3, vgg_layer, 512,
-                        scope="h4", is_training=self.is_training)
+                        scope="h4", is_training=self.is_training,
+                        activation=tf.nn.relu, has_batchnorm=has_batchnorm)
     net42 = tfl.max_pooling2d(net41, [2, 2], strides=(2, 2), name="h4_pool")
     # h5
     net51 = tfcl.repeat(net42, 1, vgg_layer, nout,
-                        kernel_size=4, padding="VALID", activation=tf.tanh,
-                        scope="h5", is_training=self.is_training)
+                        kernel_size=4, padding="VALID", activation=tf.nn.relu,
+                        scope="h5", is_training=self.is_training,
+                        has_batchnorm=has_batchnorm)
+
     skips += [net11, net21, net31, net41]
     return net51, skips
 
-  def decoder(self, inputs, skips, nout):
+  def decoder(self, inputs, nout, skips=None, has_batchnorm=True):
     """VGG based image decoder.
 
     Args:
       inputs: image tensor with size BSxX
-      skips: skip connections from encoder
       nout: number of output channels
+      skips: optional skip connections from encoder
+      has_batchnorm: variable to use or not use batch normalization
     Returns:
       net: decoded image with size BSx64x64xNout
       skips: skip connection after each layer
@@ -117,37 +126,60 @@ def decoder(self, inputs, skips, nout):
     net = inputs
     # d1
     net = tfl.conv2d_transpose(net, 512, kernel_size=4, padding="VALID",
-                               name="d1_deconv", activation=None)
-    net = tfl.batch_normalization(net, training=self.is_training, name="d1_bn")
-    net = tf.nn.leaky_relu(net)
+                               name="d1_deconv", activation=tf.nn.relu)
+    if has_batchnorm:
+      net = tfl.batch_normalization(
+          net, training=self.is_training, name="d1_bn")
+    net = tf.nn.relu(net)
     net = common_layers.upscale(net, 2)
     # d2
-    net = tf.concat([net, skips[-1]], axis=3)
-    net = tfcl.repeat(net, 2, vgg_layer, 512, scope="d2a")
-    net = tfcl.repeat(net, 1, vgg_layer, 256, scope="d2b")
+    if skips is not None:
+      net = tf.concat([net, skips[-1]], axis=3)
+    net = tfcl.repeat(net, 2, vgg_layer, 512, scope="d2a",
+                      is_training=self.is_training,
+                      activation=tf.nn.relu, has_batchnorm=has_batchnorm)
+    net = tfcl.repeat(net, 1, vgg_layer, 256, scope="d2b",
+                      is_training=self.is_training,
+                      activation=tf.nn.relu, has_batchnorm=has_batchnorm)
     net = common_layers.upscale(net, 2)
     # d3
-    net = tf.concat([net, skips[-2]], axis=3)
-    net = tfcl.repeat(net, 2, vgg_layer, 256, scope="d3a")
-    net = tfcl.repeat(net, 1, vgg_layer, 128, scope="d3b")
+    if skips is not None:
+      net = tf.concat([net, skips[-2]], axis=3)
+    net = tfcl.repeat(net, 2, vgg_layer, 256, scope="d3a",
+                      is_training=self.is_training,
+                      activation=tf.nn.relu, has_batchnorm=has_batchnorm)
+    net = tfcl.repeat(net, 1, vgg_layer, 128, scope="d3b",
+                      is_training=self.is_training,
+                      activation=tf.nn.relu, has_batchnorm=has_batchnorm)
     net = common_layers.upscale(net, 2)
     # d4
-    net = tf.concat([net, skips[-3]], axis=3)
-    net = tfcl.repeat(net, 1, vgg_layer, 128, scope="d4a")
-    net = tfcl.repeat(net, 1, vgg_layer, 64, scope="d4b")
+    if skips is not None:
+      net = tf.concat([net, skips[-3]], axis=3)
+    net = tfcl.repeat(net, 1, vgg_layer, 128, scope="d4a",
+                      is_training=self.is_training,
+                      activation=tf.nn.relu, has_batchnorm=has_batchnorm)
+    net = tfcl.repeat(net, 1, vgg_layer, 64, scope="d4b",
+                      is_training=self.is_training,
+                      activation=tf.nn.relu, has_batchnorm=has_batchnorm)
     net = common_layers.upscale(net, 2)
     # d5
-    net = tf.concat([net, skips[-4]], axis=3)
-    net = tfcl.repeat(net, 1, vgg_layer, 64, scope="d5")
-
-    # if there are still skip connections left, we have more downscaling to do
-    for i, s in enumerate(skips[-5::-1]):
-      net = common_layers.upscale(net, 2)
-      net = tf.concat([net, s], axis=3)
-      net = tfcl.repeat(net, 1, vgg_layer, 64, scope="upscale%d" % i)
+    if skips is not None:
+      net = tf.concat([net, skips[-4]], axis=3)
+    net = tfcl.repeat(net, 1, vgg_layer, 64, scope="d5",
+                      is_training=self.is_training,
+                      activation=tf.nn.relu, has_batchnorm=has_batchnorm)
+
+    # if there are still skip connections left, we have more upscaling to do
+    if skips is not None:
+      for i, s in enumerate(skips[-5::-1]):
+        net = common_layers.upscale(net, 2)
+        net = tf.concat([net, s], axis=3)
+        net = tfcl.repeat(net, 1, vgg_layer, 64, scope="upscale%d" % i,
+                          is_training=self.is_training,
+                          activation=tf.nn.relu, has_batchnorm=has_batchnorm)
 
     net = tfl.conv2d_transpose(net, nout, kernel_size=3, padding="SAME",
-                               name="d6_deconv", activation=tf.sigmoid)
+                               name="d6_deconv", activation=None)
     return net
 
   def stacked_lstm(self, inputs, states, hidden_size, output_size, nlayers):
@@ -229,6 +261,7 @@ def construct_model(self, images, actions, rewards):
     posterior_rnn_layers = self.hparams.posterior_rnn_layers
     predictor_rnn_layers = self.hparams.predictor_rnn_layers
     context_frames = self.hparams.video_num_input_frames
+    has_batchnorm = self.hparams.has_batchnorm
 
     seq_len, batch_size, _, _, color_channels = common_layers.shape_list(images)
 
@@ -242,7 +275,7 @@ def construct_model(self, images, actions, rewards):
     images = tf.unstack(images, axis=0)
     for i, image in enumerate(images):
       with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
-        enc, skips = self.encoder(image, rnn_size)
+        enc, skips = self.encoder(image, rnn_size, has_batchnorm=has_batchnorm)
         enc = tfcl.flatten(enc)
         enc_images.append(enc)
         enc_skips.append(skips)
@@ -286,7 +319,9 @@ def construct_model(self, images, actions, rewards):
         skip_index = min(context_frames-1, i)
 
         h_pred = tf.reshape(pred_enc[i], [batch_size, 1, 1, g_dim])
-        x_pred = self.decoder(h_pred, enc_skips[skip_index], color_channels)
+        x_pred = self.decoder(
+            h_pred, color_channels, enc_skips[skip_index],
+            has_batchnorm=has_batchnorm)
         gen_images.append(x_pred)
 
     tf.logging.info(">>>> Done")
@@ -298,11 +333,27 @@ def construct_model(self, images, actions, rewards):
 def next_frame_emily():
   """Emily's model hparams."""
   hparams = sv2p_params.next_frame_sv2p()
-  hparams.latent_loss_multiplier = 1e-4
+  hparams.video_num_input_frames = 2
+  hparams.video_num_target_frames = 10
   hparams.learning_rate_constant = 1e-4
+  seq_length = hparams.video_num_input_frames + hparams.video_num_target_frames
+  # The latent_loss_multiplier is divided by the number of frames because
+  # the image sequence loss in t2t is averaged instead of added through
+  # time as they do in the SVG-LP paper
+  hparams.latent_loss_multiplier = 1e-4 / seq_length
+  hparams.reward_prediction = False
+  hparams.num_iterations_1st_stage = -1
+  hparams.num_iterations_2nd_stage = -1
+  hparams.optimizer_adam_beta1 = 0.9
+  hparams.optimizer_adam_beta2 = 0.999
+  hparams.optimizer_adam_epsilon = 1e-08
+  hparams.anneal_end = -1
+  hparams.clip_grad_norm = 5.0
   hparams.add_hparam("z_dim", 10)
   hparams.add_hparam("g_dim", 128)
   hparams.add_hparam("rnn_size", 256)
   hparams.add_hparam("posterior_rnn_layers", 1)
   hparams.add_hparam("predictor_rnn_layers", 2)
+  hparams.add_hparam("has_skips", True)
+  hparams.add_hparam("has_batchnorm", True)
   return hparams
diff --git a/tensor2tensor/models/video/svg_lp.py b/tensor2tensor/models/video/svg_lp.py
new file mode 100644
index 000000000..7a49e549b
--- /dev/null
+++ b/tensor2tensor/models/video/svg_lp.py
@@ -0,0 +1,366 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Model architecture for video prediction model.
+
+   based on following paper:
+   "Stochastic Video Generation with a Learned Prior"
+   https://arxiv.org/pdf/1802.07687.pdf
+   by Emily Denton and Rob Fergus.
+
+   This code is a translation of the original code from PyTorch:
+   https://github.com/edenton/svg
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
+from tensor2tensor.models.video import emily
+from tensor2tensor.models.video import sv2p_params
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+tfl = tf.layers
+tfcl = tf.contrib.layers
+
+
+@registry.register_model
+class NextFrameSVGLP(emily.NextFrameEmily):
+  """Stochastic Variational Video Prediction With Learned Prior."""
+
+  def rnn_model(self, hidden_size, nlayers, rnn_type, name):
+    """Stacked RNN cell constructor.
+
+    Args:
+      hidden_size: number of lstm units
+      nlayers: number of lstm layers
+      rnn_type: type of RNN cell to use
+      name: RNN name
+    Returns:
+      stacked_rnn: stacked RNN cell
+    """
+    layers_units = [hidden_size] * nlayers
+    if rnn_type == "lstm":
+      rnn_cell = tf.contrib.rnn.LSTMCell
+    elif rnn_type == "gru":
+      rnn_cell = tf.contrib.rnn.GRUCell
+    else:
+      rnn_cell = tf.contrib.rnn.RNNCell
+    cells = [rnn_cell(units, name=name) for units in layers_units]
+    stacked_rnn = tf.contrib.rnn.MultiRNNCell(cells)
+    return stacked_rnn
+
+  def deterministic_rnn(self, cell, inputs, states, output_size, scope):
+    """Deterministic RNN step function.
+
+    Args:
+      cell: RNN cell to forward through
+      inputs: input to RNN cell
+      states: previous RNN state
+      output_size: size of the output
+      scope: scope of the current RNN forward computation parameters
+    Returns:
+      outputs: deterministic RNN output vector
+      states: updated RNN states
+    """
+    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
+      embedded = tfl.dense(
+          inputs, cell.output_size, activation=tf.nn.relu, name="embed")
+      hidden, states = cell(embedded, states)
+      outputs = tfl.dense(
+          hidden, output_size, activation=tf.nn.relu, name="output")
+
+    return outputs, states
+
+  def gaussian_rnn(self, cell, inputs, states, output_size, scope):
+    """Deterministic RNN step function.
+
+    Args:
+      cell: RNN cell to forward through
+      inputs: input to RNN cell
+      states: previous RNN state
+      output_size: size of the output
+      scope: scope of the current RNN forward computation parameters
+    Returns:
+      mu: mean of the predicted gaussian
+      logvar: log(var) of the predicted gaussian
+      states: updated RNN states
+    """
+    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
+      embedded = tfl.dense(
+          inputs, cell.output_size, activation=tf.nn.relu, name="embed")
+      hidden, states = cell(embedded, states)
+      mu = tfl.dense(
+          hidden, output_size, activation=None, name="mu")
+      logvar = tfl.dense(
+          hidden, output_size, activation=None, name="logvar")
+
+    return mu, logvar, states
+
+  def sample(self, mu, logvar):
+    eps = tf.random_normal([self.hparams.batch_size, self.hparams.z_dim], 0, 1)
+    sigma = tf.exp(tf.multiply(0.5, logvar))
+    z = tf.add(mu, tf.multiply(sigma, eps))
+
+    return z
+
+  def construct_model(self, images, actions, rewards):
+    """Builds the stochastic model.
+
+    The model first encodes all the images (x_t) in the sequence
+    using the encoder. Let"s call the output e_t. Then it predicts the
+    latent state of the next frame using a recurrent posterior network
+    z ~ q(z|e_{0:t}) = N(mu(e_{0:t}), sigma(e_{0:t})).
+    Another recurrent network predicts the embedding of the next frame
+    using the approximated posterior e_{t+1} = p(e_{t+1}|e_{0:t}, z)
+    Finally, the decoder decodes e_{t+1} into x_{t+1}.
+    Skip connections from encoder to decoder help with reconstruction.
+
+    Args:
+      images: tensor of ground truth image sequences
+      actions: NOT used list of action tensors
+      rewards: NOT used list of reward tensors
+
+    Returns:
+      gen_images: generated images
+      fakr_rewards: input rewards as reward prediction!
+      pred_mu: predited means of posterior
+      pred_logvar: predicted log(var) of posterior
+    """
+    # model does not support action conditioned and reward prediction
+    fake_reward_prediction = rewards
+    del actions, rewards
+
+    mode = self.hparams.mode
+    z_dim = self.hparams.z_dim
+    g_dim = self.hparams.g_dim
+    rnn_size = self.hparams.rnn_size
+    rnn_type = self.hparams.rnn_type
+    prior_rnn_layers = self.hparams.prior_rnn_layers
+    posterior_rnn_layers = self.hparams.posterior_rnn_layers
+    predictor_rnn_layers = self.hparams.predictor_rnn_layers
+    context_frames = self.hparams.video_num_input_frames
+    has_batchnorm = self.hparams.has_batchnorm
+
+    # Create RNN cells
+    predictor_cell = self.rnn_model(
+        rnn_size, predictor_rnn_layers, rnn_type, "frame_predictor")
+    prior_cell = self.rnn_model(
+        rnn_size, prior_rnn_layers, rnn_type, "prior")
+    posterior_cell = self.rnn_model(
+        rnn_size, posterior_rnn_layers, rnn_type, "posterior")
+
+    seq_len, batch_size, _, _, color_channels = common_layers.shape_list(images)
+
+    # RNN initialize states.
+    prior_states = prior_cell.zero_state(batch_size, tf.float32)
+    predictor_states = predictor_cell.zero_state(batch_size, tf.float32)
+    posterior_states = posterior_cell.zero_state(batch_size, tf.float32)
+
+    tf.logging.info(">>>> Encoding")
+    # Encoding:
+    enc_images, enc_skips = [], []
+    images = tf.unstack(images, axis=0)
+    for i, image in enumerate(images):
+      with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
+        enc, skips = self.encoder(image, g_dim, has_batchnorm=has_batchnorm)
+        enc = tfcl.flatten(enc)
+        enc_images.append(enc)
+        enc_skips.append(skips)
+
+    tf.logging.info(">>>> Prediction")
+    # Prediction
+    pred_mu = []
+    pred_logvar = []
+    pred_mu_p = []
+    pred_logvar_p = []
+    gen_images = []
+    for i in range(1, seq_len):
+      with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
+        # current encoding
+        if (mode == tf.estimator.ModeKeys.TRAIN or
+            len(gen_images) < context_frames):
+          h_current = enc_images[i-1]
+        else:
+          h_current, _ = self.encoder(gen_images[-1], g_dim)
+          h_current = tfcl.flatten(h_current)
+
+        # target encoding
+        h_target = enc_images[i]
+
+      with tf.variable_scope("prediction", reuse=tf.AUTO_REUSE):
+        # Prior parameters
+        mu_p, logvar_p, prior_states = self.gaussian_rnn(
+            prior_cell, h_current, prior_states, z_dim, "prior")
+
+        # Only use Posterior if it's training time
+        if mode == tf.estimator.ModeKeys.TRAIN:
+          mu, logvar, posterior_states = self.gaussian_rnn(
+              posterior_cell, h_target, posterior_states, z_dim, "posterior")
+          z = self.sample(mu, logvar)
+        else:
+          mu = tf.zeros_like(mu_p)
+          logvar = tf.zeros_like(logvar_p)
+          z = self.sample(mu_p, logvar_p)
+
+        # Predict output images
+        h_pred, predictor_states = self.deterministic_rnn(
+            predictor_cell, tf.concat([h_current, z], axis=1),
+            predictor_states, g_dim, "predictor")
+
+        pred_mu.append(tf.identity(mu, "mu"))
+        pred_logvar.append(tf.identity(logvar, "logvar"))
+        pred_mu_p.append(tf.identity(mu_p, "mu_p"))
+        pred_logvar_p.append(tf.identity(logvar_p, "log_var_p"))
+
+      with tf.variable_scope("decoding", reuse=tf.AUTO_REUSE):
+        skip_index = min(context_frames-1, i-1)
+        h_pred = tf.reshape(h_pred, [batch_size, 1, 1, g_dim])
+        if self.hparams.has_skips:
+          x_pred = self.decoder(
+              h_pred, color_channels,
+              skips=enc_skips[skip_index], has_batchnorm=has_batchnorm)
+        else:
+          x_pred = self.decoder(
+              h_pred, color_channels, has_batchnorm=has_batchnorm)
+        gen_images.append(x_pred)
+
+    tf.logging.info(">>>> Done")
+    gen_images = tf.stack(gen_images, axis=0)
+    return (gen_images, fake_reward_prediction,
+            pred_mu, pred_logvar, pred_mu_p, pred_logvar_p)
+
+  def get_extra_loss(self,
+                     latent_means, latent_logvars,
+                     latent_means_p, latent_logvars_p):
+    """Losses in addition to the default modality losses."""
+    return self.get_kl_loss(
+        latent_means, latent_logvars, latent_means_p, latent_logvars_p)
+
+  def body(self, features):
+    hparams = self.hparams
+    batch_size = common_layers.shape_list(features["inputs"])[0]
+
+    # Swap time and batch axes.
+    input_frames = common_video.swap_time_and_batch_axes(features["inputs"])
+    target_frames = common_video.swap_time_and_batch_axes(features["targets"])
+
+    # Get actions if exist otherwise use zeros
+    input_actions = self.get_input_if_exists(
+        features, "input_action", batch_size, hparams.video_num_input_frames)
+    target_actions = self.get_input_if_exists(
+        features, "target_action", batch_size, hparams.video_num_target_frames)
+
+    # Get rewards if exist otherwise use zeros
+    input_rewards = self.get_input_if_exists(
+        features, "input_reward", batch_size, hparams.video_num_input_frames)
+    target_rewards = self.get_input_if_exists(
+        features, "target_reward", batch_size, hparams.video_num_target_frames)
+
+    all_actions = tf.concat([input_actions, target_actions], axis=0)
+    all_rewards = tf.concat([input_rewards, target_rewards], axis=0)
+    all_frames = tf.concat([input_frames, target_frames], axis=0)
+
+    # Each image is being used twice, in latent tower and main tower.
+    # This is to make sure we are using the *same* image for both, ...
+    # ... given how TF queues work.
+    # NOT sure if this is required at all. Doesn"t hurt though! :)
+    all_frames = tf.identity(all_frames)
+
+    retvals = self.construct_model(
+        images=all_frames, actions=all_actions, rewards=all_rewards)
+
+    # retrieve tensors returned by the model contructor
+    gen_images = retvals[0]
+    gen_rewards = retvals[1]
+    latent_means = retvals[2]
+    latent_logvars = retvals[3]
+    latent_means_p = retvals[4]
+    latent_logvars_p = retvals[5]
+
+    extra_loss = self.get_extra_loss(
+        latent_means=latent_means,
+        latent_logvars=latent_logvars,
+        latent_means_p=latent_means_p,
+        latent_logvars_p=latent_logvars_p)
+
+    # Visualize predictions in Tensorboard
+    if self.is_training:
+      self.visualize_predictions(all_frames[1:], gen_images)
+
+    # Ignore the predictions from the input frames.
+    # This is NOT the same as original paper/implementation.
+    predictions = gen_images[hparams.video_num_input_frames-1:]
+    reward_pred = gen_rewards[hparams.video_num_input_frames-1:]
+    reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove extra dimension.
+
+    # Swap back time and batch axes.
+    predictions = common_video.swap_time_and_batch_axes(predictions)
+    reward_pred = common_video.swap_time_and_batch_axes(reward_pred)
+
+    if self.is_training and hparams.internal_loss:
+      # add the loss for input frames as well.
+      extra_gts = all_frames[1:hparams.video_num_input_frames]
+      extra_gts = common_video.swap_time_and_batch_axes(extra_gts)
+      extra_pds = gen_images[:hparams.video_num_input_frames-1]
+      extra_pds = common_video.swap_time_and_batch_axes(extra_pds)
+      extra_raw_gts = features["inputs_raw"][:, 1:]
+      recon_loss = self.get_extra_internal_loss(
+          extra_raw_gts, extra_gts, extra_pds)
+      extra_loss += recon_loss
+
+    return_targets = predictions
+    if hparams.reward_prediction:
+      return_targets = {"targets": predictions, "target_reward": reward_pred}
+
+    return return_targets, extra_loss
+
+
+@registry.register_hparams
+def next_frame_svglp():
+  """SVG with learned prior model hparams."""
+  hparams = sv2p_params.next_frame_sv2p()
+  hparams.video_num_input_frames = 2
+  hparams.video_num_target_frames = 10
+  hparams.learning_rate_constant = 1e-4
+  seq_length = hparams.video_num_input_frames + hparams.video_num_target_frames
+  # The latent_loss_multiplier is divided by the number of frames because
+  # the image sequence loss in t2t is averaged instead of added through
+  # time as they do in the SVG-LP paper
+  hparams.latent_loss_multiplier = 1e-4 / seq_length
+  hparams.reward_prediction = False
+  hparams.num_iterations_1st_stage = -1
+  hparams.num_iterations_2nd_stage = -1
+  hparams.optimizer_adam_beta1 = 0.9
+  hparams.optimizer_adam_beta2 = 0.999
+  hparams.optimizer_adam_epsilon = 1e-08
+  hparams.anneal_end = -1
+  hparams.clip_grad_norm = 5.0
+  hparams.add_hparam("learned_prior", True)
+  hparams.add_hparam("z_dim", 64)
+  hparams.add_hparam("g_dim", 128)
+  hparams.add_hparam("rnn_size", 256)
+  hparams.add_hparam("rnn_type", "lstm")
+  hparams.add_hparam("prior_rnn_layers", 1)
+  hparams.add_hparam("posterior_rnn_layers", 1)
+  hparams.add_hparam("predictor_rnn_layers", 2)
+  hparams.add_hparam("has_skips", True)
+  hparams.add_hparam("has_batchnorm", True)
+  return hparams
+

From 7fa27cf17fe51aaa2150da91cce5edb39c635e8d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 20 Nov 2018 07:20:19 -0800
Subject: [PATCH 1263/2720] Fix broken build.

PiperOrigin-RevId: 222241661
---
 tensor2tensor/bin/t2t_trainer.py    |  7 +++----
 tensor2tensor/models/transformer.py | 14 +++++++-------
 tensor2tensor/utils/decoding.py     |  2 +-
 tensor2tensor/utils/mlperf_log.py   |  9 +++++++--
 tensor2tensor/utils/optimize.py     | 10 +++++-----
 tensor2tensor/utils/t2t_model.py    |  2 +-
 6 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index fe9c8ac03..d78a153b5 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -357,13 +357,12 @@ def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
   hparams = create_hparams()
   if FLAGS.schedule == "train" or FLAGS.schedule == "train_eval_and_decode":
-    mlperf_log.transformer_print(key=mlperf_log.RUN_START,
-                                 mlperf_mode=hparams.mlperf_mode)
+    mlperf_log.transformer_print(key=mlperf_log.RUN_START, hparams=hparams)
   if FLAGS.schedule == "run_std_server":
     run_std_server()
   mlperf_log.transformer_print(
       key=mlperf_log.RUN_SET_RANDOM_SEED, value=FLAGS.random_seed,
-      mlperf_mode=hparams.mlperf_mode)
+      hparams=hparams)
   trainer_lib.set_random_seed(FLAGS.random_seed)
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
   maybe_log_registry_and_exit()
@@ -388,7 +387,7 @@ def main(argv):
   execute_schedule(exp)
   if FLAGS.schedule != "train":
     mlperf_log.transformer_print(key=mlperf_log.RUN_FINAL,
-                                 mlperf_mode=hparams.mlperf_mode)
+                                 hparams=hparams)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 0151abed7..1bdf94a26 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -87,7 +87,7 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
     mlperf_log.transformer_print(
         key=mlperf_log.MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
         value=hparams.layer_prepostprocess_dropout,
-        mlperf_mode=hparams.mlperf_mode)
+        hparams=hparams)
 
     encoder_input = tf.nn.dropout(encoder_input,
                                   1.0 - hparams.layer_prepostprocess_dropout)
@@ -138,7 +138,7 @@ def decode(self,
     mlperf_log.transformer_print(
         key=mlperf_log.MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
         value=hparams.layer_prepostprocess_dropout,
-        mlperf_mode=hparams.mlperf_mode)
+        hparams=hparams)
     decoder_input = tf.nn.dropout(decoder_input,
                                   1.0 - hparams.layer_prepostprocess_dropout)
 
@@ -805,7 +805,7 @@ def fast_decode_tpu(encoder_output,
           "alpha": alpha,
           "max_decode_length": decode_length
       },
-      mlperf_mode=hparams.mlperf_mode)
+      hparams=hparams)
   if beam_size > 1:  # Beam Search
     initial_ids = sos_id * tf.ones([batch_size], dtype=tf.int32)
     decoded_ids, scores = beam_search.beam_search(
@@ -1226,11 +1226,11 @@ def transformer_decoder(decoder_input,
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_NUM_HIDDEN_LAYERS,
       value=hparams.num_decoder_layers or hparams.num_hidden_layers,
-      mlperf_mode=hparams.mlperf_mode)
+      hparams=hparams)
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_ATTENTION_DROPOUT,
       value=hparams.attention_dropout,
-      mlperf_mode=hparams.mlperf_mode)
+      hparams=hparams)
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_ATTENTION_DENSE,
       value={
@@ -1238,7 +1238,7 @@ def transformer_decoder(decoder_input,
           "num_heads": hparams.num_heads,
           "hidden_size": hparams.hidden_size
       },
-      mlperf_mode=hparams.mlperf_mode)
+      hparams=hparams)
 
   with tf.variable_scope(name):
     for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
@@ -1306,7 +1306,7 @@ def transformer_decoder(decoder_input,
     mlperf_log.transformer_print(
         key=mlperf_log.MODEL_HP_NORM,
         value={"hidden_size": hparams.hidden_size},
-        mlperf_mode=hparams.mlperf_mode)
+        hparams=hparams)
     return common_layers.layer_preprocess(x, hparams)
 
 
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index f4fbe9eb5..f051543e6 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -334,7 +334,7 @@ def decode_once(estimator,
 
   mlperf_log.transformer_print(key=mlperf_log.EVAL_SIZE,
                                value=num_eval_samples,
-                               mlperf_mode=hparams.mlperf_mode)
+                               hparams=hparams)
 
   if decode_to_file:
     output_file.close()
diff --git a/tensor2tensor/utils/mlperf_log.py b/tensor2tensor/utils/mlperf_log.py
index 23b4bdf74..81e11cf93 100644
--- a/tensor2tensor/utils/mlperf_log.py
+++ b/tensor2tensor/utils/mlperf_log.py
@@ -76,6 +76,11 @@
   _STREAM_HANDLER.setLevel(logging.DEBUG)
 
 
+def get_mode(hparams):
+  """Returns whether we should do MLPerf logging."""
+  return "mlperf_mode" in hparams and hparams.mlperf_mode
+
+
 def get_caller(stack_index=2, root_dir=None):
   # pylint: disable=g-doc-args
   """Returns file.py:lineno of your caller.
@@ -162,8 +167,8 @@ def _mlperf_print(key, value=None, benchmark=None, stack_offset=0,
 
 
 def transformer_print(key, value=None, stack_offset=2, deferred=False,
-                      mlperf_mode=False):
-  if not mlperf_mode:
+                      hparams=None):
+  if not hparams or not get_mode(hparams):
     return
   return _mlperf_print(
       key=key,
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 3b393d159..a8b3fac28 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -90,17 +90,17 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
 
     mlperf_log.transformer_print(key=mlperf_log.OPT_NAME,
                                  value=optimizer_name,
-                                 mlperf_mode=hparams.mlperf_mode)
+                                 hparams=hparams)
     mlperf_log.transformer_print(
         key=mlperf_log.OPT_HP_ADAM_BETA1, value=hparams.optimizer_adam_beta1,
-        mlperf_mode=hparams.mlperf_mode)
+        hparams=hparams)
     mlperf_log.transformer_print(
         key=mlperf_log.OPT_HP_ADAM_BETA2, value=hparams.optimizer_adam_beta2,
-        mlperf_mode=hparams.mlperf_mode)
+        hparams=hparams)
     mlperf_log.transformer_print(
         key=mlperf_log.OPT_HP_ADAM_EPSILON,
         value=hparams.optimizer_adam_epsilon,
-        mlperf_mode=hparams.mlperf_mode)
+        hparams=hparams)
 
     if optimizer_name == "Adam":
       # We change the default epsilon for Adam.
@@ -281,7 +281,7 @@ def get_variable_initializer(hparams):
 
   mlperf_log.transformer_print(key=mlperf_log.MODEL_HP_INITIALIZER_GAIN,
                                value=hparams.initializer_gain,
-                               mlperf_mode=hparams.mlperf_mode)
+                               hparams=hparams)
 
   if not tf.contrib.eager.in_eager_mode():
     tf.logging.info("Using variable initializer: %s", hparams.initializer)
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 3f3e45ee5..40f96e184 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -143,7 +143,7 @@ def __init__(self,
                 "vocab_size": target_modality.top_dimensionality,
                 "hidden_size": hidden_size
             },
-            mlperf_mode=hparams.mlperf_mode)
+            hparams=hparams)
 
     self._original_hparams = hparams
     self.set_mode(mode)

From 3a30004f60b2dbf2df2d96a639dc46836b7c4140 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 20 Nov 2018 10:41:41 -0800
Subject: [PATCH 1264/2720] Fix broken build, again.

PiperOrigin-RevId: 222270326
---
 tensor2tensor/bin/t2t_trainer.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index d78a153b5..3610e89f9 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -355,7 +355,15 @@ def run_std_server():
 
 def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
+
+  # If we just have to print the registry, do that and exit early.
+  maybe_log_registry_and_exit()
+
+  # Create HParams.
+  if argv:
+    set_hparams_from_args(argv[1:])
   hparams = create_hparams()
+
   if FLAGS.schedule == "train" or FLAGS.schedule == "train_eval_and_decode":
     mlperf_log.transformer_print(key=mlperf_log.RUN_START, hparams=hparams)
   if FLAGS.schedule == "run_std_server":
@@ -365,7 +373,6 @@ def main(argv):
       hparams=hparams)
   trainer_lib.set_random_seed(FLAGS.random_seed)
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
-  maybe_log_registry_and_exit()
 
   if FLAGS.cloud_mlengine:
     cloud_mlengine.launch()
@@ -377,9 +384,6 @@ def main(argv):
   if cloud_mlengine.job_dir():
     FLAGS.output_dir = cloud_mlengine.job_dir()
 
-  if argv:
-    set_hparams_from_args(argv[1:])
-
   exp_fn = create_experiment_fn()
   exp = exp_fn(create_run_config(hparams), hparams)
   if is_chief():

From f1f42bf0673cc34b3795e04ae6d40da7b0a48ad8 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 20 Nov 2018 11:06:40 -0800
Subject: [PATCH 1265/2720] Add sparse gradient support for adafactor.

PiperOrigin-RevId: 222274391
---
 tensor2tensor/utils/adafactor.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index 466d5b59d..01b6145ba 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -191,6 +191,11 @@ def _apply_dense(self, grad, var):
   def _apply_sparse(self, grad, var):
     return self._apply_dense(tf.convert_to_tensor(grad), var)
 
+  def _resource_apply_sparse(self, grad, handle, indices):
+    return self._resource_apply_dense(
+        tf.convert_to_tensor(tf.IndexedSlices(grad, indices, tf.shape(handle))),
+        handle)
+
   def _parameter_scale(self, var):
     """Estimate the scale of the parameters from the current values.
 

From 0e45d1eee0a6d966e875cdc5c67bacf4346846f1 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 20 Nov 2018 12:17:11 -0800
Subject: [PATCH 1266/2720] Fix broken build, truly :)

PiperOrigin-RevId: 222286456
---
 tensor2tensor/bin/t2t_trainer.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 3610e89f9..c123933fd 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -356,6 +356,8 @@ def run_std_server():
 def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
 
+  usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
+
   # If we just have to print the registry, do that and exit early.
   maybe_log_registry_and_exit()
 
@@ -372,7 +374,6 @@ def main(argv):
       key=mlperf_log.RUN_SET_RANDOM_SEED, value=FLAGS.random_seed,
       hparams=hparams)
   trainer_lib.set_random_seed(FLAGS.random_seed)
-  usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
 
   if FLAGS.cloud_mlengine:
     cloud_mlengine.launch()

From aba83867ffc17883e3596ae6e432e857ee8f7e4d Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 20 Nov 2018 13:37:15 -0800
Subject: [PATCH 1267/2720] Allow fixed-length training in multiproblem and
 decoding, style corrections too.

PiperOrigin-RevId: 222298682
---
 .../data_generators/multi_problem.py          | 40 +++++++++++--------
 tensor2tensor/layers/common_hparams.py        |  4 +-
 tensor2tensor/layers/common_layers.py         |  3 --
 tensor2tensor/utils/decoding.py               | 14 ++++---
 4 files changed, 35 insertions(+), 26 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 53d6363f0..0c400d65f 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -46,7 +46,6 @@ def __init__(self, was_reversed=False, was_copy=False):
 
   def generate_data(self, data_dir, tmp_dir, task_id=-1):
     assert len(self.task_list) > 1
-
     for task in self.task_list:
       task.generate_data(data_dir, tmp_dir, task_id)
 
@@ -75,25 +74,35 @@ def add_task_id(self, task, example, encoder, hparams, is_infer):
           example["targets"] = example[
               "targets"][:hparams.multiproblem_max_target_length]
 
+    def make_constant_shape(x, size):
+      x = x[:size]
+      xlen = tf.shape(x)[0]
+      x = tf.pad(x, [[0, size - xlen]])
+      return tf.reshape(x, [size])
+
     if task.has_inputs:
       if is_infer:
         concat_list = [example["inputs"], [task.task_id]]
-        example["inputs"] = tf.concat(concat_list, 0)
+        example["inputs"] = tf.concat(concat_list, axis=0)
       else:
         inputs = example.pop("inputs")
         concat_list = [inputs, [task.task_id], example["targets"]]
-        example["targets"] = tf.concat(concat_list, 0)
+        example["targets"] = tf.concat(concat_list, axis=0)
+        if hparams.multiproblem_fixed_train_length > 0:
+          example["targets"] = make_constant_shape(
+              example["targets"], hparams.multiproblem_fixed_train_length)
     else:
       concat_list = [[task.task_id], example["targets"]]
-      example["targets"] = tf.concat(concat_list, 0)
+      example["targets"] = tf.concat(concat_list, axis=0)
+      if not is_infer and hparams.multiproblem_fixed_train_length > 0:
+        example["targets"] = make_constant_shape(
+            example["targets"], hparams.multiproblem_fixed_train_length)
 
-    min_task_id = min([t.task_id for t in self.task_list])
-    example["task_id"] = tf.constant([task.task_id - min_task_id],
-                                     dtype=tf.int64)
+    example["task_id"] = tf.constant([task.task_id], dtype=tf.int64)
     return example
 
   def filepattern(self, data_dir, mode, shard=None):
-    print("Generating multi problem filepattern")
+    tf.logging.info("Generating multi problem filepattern")
     return [task.filepattern(data_dir, mode, shard) for task in self.task_list]
 
   def get_hparams(self, model_hparams=None):
@@ -164,8 +173,9 @@ def dataset(self,
       raise ValueError("Only support language models as primary problem which "
                        "supplies the vocabulary and the hparams.")
     enc = primary_task.feature_encoders(data_dir=data_dir)["targets"]
+    self.update_task_ids(enc)
 
-    for idx, task in enumerate(self.task_list):
+    for task in self.task_list:
       task_dataset = task.dataset(mode=mode,
                                   data_dir=data_dir,
                                   num_threads=num_threads,
@@ -181,9 +191,6 @@ def dataset(self,
                                   max_records=max_records,
                                   only_last=only_last)
 
-      if idx == 0:
-        self.update_task_ids(enc)
-
       if is_training:
         task_dataset = task_dataset.repeat()
 
@@ -285,7 +292,6 @@ def sample_task(curr_task, num_tasks_left, randnum):
             A Tensor representing an example from the task that was sampled
             from.
           """
-
           if num_tasks_left == 0:
             return get_next_from_dataset(dataset_iterators[curr_task])
 
@@ -338,10 +344,10 @@ def update_task_ids(self, encoder):
     """
     offset = encoder.vocab_size
 
-    for idx, _ in enumerate(self.task_list):
-      self.task_list[idx].set_task_id(idx + offset)
-      print(self.task_list[idx].name)
-      print(self.task_list[idx].task_id)
+    for idx, task in enumerate(self.task_list):
+      task.set_task_id(idx + offset)
+      tf.logging.info("Task %d (%s) has id %d." %
+                      (idx, task.name, task.task_id))
 
   def get_max_num_classes(self):
     """Compute the maximum number of classes any subtask has.
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 28a281c6e..304d501e6 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -298,7 +298,9 @@ def basic_params1():
       # When using multiproblem with generation tasks, need to truncate the
       # inputs and targets manually before concatenating them.
       multiproblem_max_input_length=-1,
-      multiproblem_max_target_length=-1
+      multiproblem_max_target_length=-1,
+      # If positive, makes training targets fixed-length in MultiProblem.
+      multiproblem_fixed_train_length=-1
   )
 
 
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 8211b0fbf..91954ee32 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1600,9 +1600,6 @@ def weights_multi_problem(labels, taskid=-1):
   Raises:
     ValueError: The Task ID must be valid.
   """
-  if taskid < 0:
-    raise ValueError("Task ID must be non-negative.")
-
   past_taskid = tf.cumsum(tf.to_float(tf.equal(labels, taskid)), axis=1)
   # Additionally zero out the task id location
   past_taskid *= tf.to_float(tf.not_equal(labels, taskid))
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index f051543e6..b2e5362b4 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -70,6 +70,8 @@ def decode_hparams(overrides=""):
       num_decodes=1,
       force_decode_length=False,
       display_decoded_images=False,
+      # Multi-problem decoding task id.
+      multiproblem_task_id=-1,
       # Used for video decoding.
       frames_per_second=10,
       skip_eos_postprocess=False,
@@ -368,9 +370,10 @@ def decode_from_file(estimator,
   num_decode_batches = (len(sorted_inputs) - 1) // decode_hp.batch_size + 1
 
   def input_fn():
-    input_gen = _decode_batch_input_fn(num_decode_batches, sorted_inputs,
-                                       inputs_vocab, decode_hp.batch_size,
-                                       decode_hp.max_input_size)
+    input_gen = _decode_batch_input_fn(
+        num_decode_batches, sorted_inputs,
+        inputs_vocab, decode_hp.batch_size,
+        decode_hp.max_input_size, task_id=decode_hp.multiproblem_task_id)
     gen_fn = make_input_fn_from_generator(input_gen)
     example = gen_fn()
     return _decode_input_tensor_to_features_dict(example, hparams)
@@ -577,7 +580,7 @@ def input_fn():
 
 
 def _decode_batch_input_fn(num_decode_batches, sorted_inputs, vocabulary,
-                           batch_size, max_input_size):
+                           batch_size, max_input_size, task_id=-1):
   """Generator to produce batches of inputs."""
   tf.logging.info(" batch %d" % num_decode_batches)
   # First reverse all the input sentences so that if you're going to get OOMs,
@@ -592,7 +595,8 @@ def _decode_batch_input_fn(num_decode_batches, sorted_inputs, vocabulary,
       if max_input_size > 0:
         # Subtract 1 for the EOS_ID.
         input_ids = input_ids[:max_input_size - 1]
-      input_ids.append(text_encoder.EOS_ID)
+      final_id = text_encoder.EOS_ID if task_id < 0 else task_id
+      input_ids.append(final_id)
       batch_inputs.append(input_ids)
       if len(input_ids) > batch_length:
         batch_length = len(input_ids)

From fd3d56d5fb92bebe372653548efbb0cdc0b0e4db Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Tue, 20 Nov 2018 16:45:55 -0800
Subject: [PATCH 1268/2720] adding internal_state control to prediction2gif

PiperOrigin-RevId: 222328025
---
 tensor2tensor/utils/video/prediction2gif.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/video/prediction2gif.py b/tensor2tensor/utils/video/prediction2gif.py
index b77256d5e..326d8ed06 100644
--- a/tensor2tensor/utils/video/prediction2gif.py
+++ b/tensor2tensor/utils/video/prediction2gif.py
@@ -99,6 +99,7 @@ def main(_):
         "inputs": tf.placeholder(tf.float32, input_size + frame_shape),
         "input_action": tf.placeholder(tf.int64, input_size + [1]),
         "input_reward": tf.placeholder(tf.int64, input_size + [1]),
+        "reset_internal_states": tf.placeholder(tf.float32, []),
     }
   # Create model.
   model_cls = registry.model(FLAGS.model)
@@ -116,7 +117,7 @@ def main(_):
   writer = common_video.WholeVideoWriter(
       fps=FLAGS.fps, output_path=FLAGS.output_gif)
 
-  saver = tf.train.Saver()
+  saver = tf.train.Saver(tf.trainable_variables())
   with tf.train.SingularMonitoredSession() as sess:
     # Load latest checkpoint
     ckpt = tf.train.get_checkpoint_state(FLAGS.output_dir).model_checkpoint_path
@@ -171,6 +172,7 @@ def main(_):
             placeholders["inputs"]: inputs,
             placeholders["input_action"]: input_action,
             placeholders["input_reward"]: input_reward,
+            placeholders["reset_internal_states"]: float(step == 0),
         }
       predictions = sess.run(prediction_ops, feed_dict=feed)
 
@@ -191,8 +193,7 @@ def main(_):
 
       writer.write(np.round(predicted_states[0]).astype(np.uint8))
 
-    video = writer.finish()
-    writer.save_to_disk(video)
+    writer.finish_to_disk()
 
 if __name__ == "__main__":
   tf.app.run()

From 52cefeb9bdf46682cf060c96822a14708def6515 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 20 Nov 2018 21:23:21 -0800
Subject: [PATCH 1269/2720] Rename target_mask to targets_mask to match targets
 feature.

PiperOrigin-RevId: 222351853
---
 tensor2tensor/utils/metrics.py      | 6 +++---
 tensor2tensor/utils/metrics_test.py | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 8d93ec391..7150006c0 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -251,9 +251,9 @@ def padded_neg_log_perplexity_with_masking(
     features,
     weights_fn=None):
   del weights_fn
-  if "target_mask" not in features:
-    raise ValueError("masked_neg_log_perplexity requires target_mask feature")
-  mask_fn = lambda labels: features["target_mask"]
+  if "targets_mask" not in features:
+    raise ValueError("masked_neg_log_perplexity requires targets_mask feature")
+  mask_fn = lambda labels: features["targets_mask"]
   return padded_neg_log_perplexity(predictions, labels, mask_fn)
 
 
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 92861e0ba..26427a30a 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -121,7 +121,7 @@ def testNegativeLogPerplexityMasked(self):
     predictions = np.random.randint(4, size=(12, 12, 12, 1))
     targets = np.random.randint(4, size=(12, 12, 12, 1))
     features = {
-        'target_mask': tf.to_float(tf.not_equal(targets, 0))
+        'targets_mask': tf.to_float(tf.not_equal(targets, 0))
     }
     with self.test_session() as session:
       scores, _ = metrics.padded_neg_log_perplexity_with_masking(
@@ -140,7 +140,7 @@ def testNegativeLogPerplexityMaskedAssert(self):
 
     with self.assertRaisesRegexp(
         ValueError,
-        'masked_neg_log_perplexity requires target_mask feature'):
+        'masked_neg_log_perplexity requires targets_mask feature'):
       with self.test_session() as session:
         scores, _ = metrics.padded_neg_log_perplexity_with_masking(
             tf.one_hot(predictions, depth=4, dtype=tf.float32),

From b46958567fffd3414817756dee2ce61e39336cd1 Mon Sep 17 00:00:00 2001
From: Mark Daoust <markdaoust@google.com>
Date: Wed, 21 Nov 2018 10:02:45 -0800
Subject: [PATCH 1270/2720] Move `tf.python_io.tf_record_iterator` references
 to `tf.compat.v1.tf_record_iterator`.

PiperOrigin-RevId: 222424141
---
 tensor2tensor/data_generators/bair_robot_pushing.py  |  3 +--
 tensor2tensor/data_generators/generator_utils.py     |  2 +-
 .../data_generators/google_robot_pushing.py          |  2 +-
 tensor2tensor/data_generators/gym_env.py             |  2 +-
 tensor2tensor/data_generators/gym_env_test.py        | 12 +++++-------
 tensor2tensor/data_generators/inspect_tfrecord.py    |  2 +-
 tensor2tensor/data_generators/problem.py             |  2 +-
 tensor2tensor/notebooks/t2t_problem.ipynb            |  2 +-
 8 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index 0a349d972..bd9489fc8 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -112,7 +112,7 @@ def parse_frames(self, filenames):
 
     for f in filenames:
       print("Parsing ", f)
-      for serialized_example in tf.python_io.tf_record_iterator(f):
+      for serialized_example in tf.compat.v1.io.tf_record_iterator(f):
         x = tf.train.Example()
         x.ParseFromString(serialized_example)
         # there are 4 features per frame
@@ -176,4 +176,3 @@ def extra_reading_spec(self):
         "action": tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="action"),
     }
     return data_fields, decoders
-
diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index c79babef1..4d22a0651 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -464,7 +464,7 @@ def generate():
 
 
 def read_records(filename):
-  reader = tf.python_io.tf_record_iterator(filename)
+  reader = tf.compat.v1.io.tf_record_iterator(filename)
   records = []
   for record in reader:
     records.append(record)
diff --git a/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index 5b295fccf..bfa2b3c92 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -83,7 +83,7 @@ def parse_frames(self, filename):
     action_key = "move/{}/commanded_pose/vec_pitch_yaw"
     state_key = "move/{}/endeffector/vec_pitch_yaw"
 
-    for serialized_example in tf.python_io.tf_record_iterator(filename):
+    for serialized_example in tf.compat.v1.io.tf_record_iterator(filename):
       x = tf.train.Example()
       x.ParseFromString(serialized_example)
       # there are 6 features per frame
diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index b91b4d215..5c4258f43 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -534,7 +534,7 @@ def _load_epoch_split(self, split, paths):
 
     for path in paths:
       this_shard_empty = True
-      for example in tf.python_io.tf_record_iterator(path):
+      for example in tf.compat.v1.io.tf_record_iterator(path):
         this_shard_empty = False
 
         result = tf.train.Example.FromString(example)
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 873de708a..6d557297e 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -164,7 +164,7 @@ def test_generates_data(self):
     self.assertTrue(filenames)
     for filename in filenames:
       path = os.path.join(self.out_dir, filename)
-      records = list(tf.python_io.tf_record_iterator(path))
+      records = list(tf.compat.v1.io.tf_record_iterator(path))
       self.assertTrue(records)
 
   def test_shards_per_epoch(self):
@@ -197,12 +197,10 @@ def test_frame_numbers_are_continuous(self):
     )
 
     frame_numbers = [
-        tf.train.Example.FromString(
-            record
-        ).features.feature["frame_number"].int64_list.value[0]
-        for (_, paths) in env.splits_and_paths(self.out_dir)
-        for path in paths
-        for record in tf.python_io.tf_record_iterator(path)
+        tf.train.Example.FromString(record).features.feature["frame_number"]
+        .int64_list.value[0]
+        for (_, paths) in env.splits_and_paths(self.out_dir) for path in paths
+        for record in tf.compat.v1.io.tf_record_iterator(path)
     ]
     last_frame_number = -1
     for frame_number in frame_numbers:
diff --git a/tensor2tensor/data_generators/inspect_tfrecord.py b/tensor2tensor/data_generators/inspect_tfrecord.py
index a245edb9e..c28a9a32c 100644
--- a/tensor2tensor/data_generators/inspect_tfrecord.py
+++ b/tensor2tensor/data_generators/inspect_tfrecord.py
@@ -56,7 +56,7 @@ def main(_):
     encoder = text_encoder.ByteTextEncoder()
   else:
     encoder = None
-  reader = tf.python_io.tf_record_iterator(FLAGS.input_filename)
+  reader = tf.compat.v1.io.tf_record_iterator(FLAGS.input_filename)
   total_sequences = 0
   total_input_tokens = 0
   total_target_tokens = 0
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 3791123c1..bf4f31e8d 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -169,7 +169,7 @@ def _file_num_records_cached(filename):
   if filename in _file_num_records_cache:
     return _file_num_records_cache[filename]
   ret = 0
-  for _ in tf.python_io.tf_record_iterator(filename):
+  for _ in tf.compat.v1.io.tf_record_iterator(filename):
     ret += 1
   _file_num_records_cache[filename] = ret
   return ret
diff --git a/tensor2tensor/notebooks/t2t_problem.ipynb b/tensor2tensor/notebooks/t2t_problem.ipynb
index 1eddf9e6b..d4a8df7b7 100644
--- a/tensor2tensor/notebooks/t2t_problem.ipynb
+++ b/tensor2tensor/notebooks/t2t_problem.ipynb
@@ -65,7 +65,7 @@
         "\n",
         "\u003e\u003e[Viewing the generated data.](#scrollTo=MCqJhdnYgiG-)\n",
         "\n",
-        "\u003e\u003e\u003e[tf.python_io.tf_record_iterator](#scrollTo=uNpohcPXKsLN)\n",
+        "\u003e\u003e\u003e[tf.compat.v1.io.tf_record_iterator](#scrollTo=uNpohcPXKsLN)\n",
         "\n",
         "\u003e\u003e\u003e[Using tf.data.Dataset](#scrollTo=6o_1BHGQC5w5)\n",
         "\n",

From abf63bfd672054b22b68b63543a96f4160530f76 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Wed, 21 Nov 2018 19:05:59 +0100
Subject: [PATCH 1271/2720] Model-Based RL: Introduce dqn implemented in
 dopamine (#1217)

* MBRL: Introduce DQN implemented in dopamine.

* MBRL: Introduce dqn hparams.

* MBRL: End-to-end tiny dqn test.

* MBRL: Move to dopamine 1.0.4 (module variables rename); minor cleanup of imports.
---
 tensor2tensor/models/research/rl.py           |  38 +++
 tensor2tensor/rl/dopamine_connector.py        | 308 ++++++++++++++++++
 .../rl/envs/simulated_batch_gym_env.py        |  16 +
 tensor2tensor/rl/policy_learner.py            |   1 +
 tensor2tensor/rl/trainer_model_based.py       |   4 +-
 .../rl/trainer_model_based_dqn_test.py        |  38 +++
 .../rl/trainer_model_based_params.py          |  41 ++-
 7 files changed, 444 insertions(+), 2 deletions(-)
 create mode 100644 tensor2tensor/rl/dopamine_connector.py
 create mode 100644 tensor2tensor/rl/trainer_model_based_dqn_test.py

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 64339f3d3..befe3a775 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -200,6 +200,44 @@ def pong_model_free():
   return hparams
 
 
+@registry.register_hparams
+def dqn_atari_base():
+  # These params are based on agents/dqn/configs/dqn.gin
+  # with some modifications taking into account our code
+  return tf.contrib.training.HParams(
+      agent_gamma=0.99,
+      agent_update_horizon=1,
+      agent_min_replay_history=20000,  # agent steps
+      agent_update_period=4,
+      agent_target_update_period=8000,  # agent steps
+      agent_epsilon_train=0.01,
+      agent_epsilon_eval=0.001,
+      agent_epsilon_decay_period=250000,  # agent steps
+      agent_generates_trainable_dones=True,
+
+      optimizer_class="RMSProp",
+      optimizer_learning_rate=0.00025,
+      optimizer_decay=0.95,
+      optimizer_momentum=0.0,
+      optimizer_epsilon=0.00001,
+      optimizer_centered=True,
+
+      replay_buffer_replay_capacity=1000000,
+      replay_buffer_batch_size=32,
+
+      time_limit=27000,
+      save_every_steps=50000,
+      num_frames=int(20 * 1e6),
+  )
+
+
+@registry.register_hparams
+def dqn_original_params():
+  hparams = dqn_atari_base()
+  hparams.set_hparam('num_frames', int(1e6))
+  return hparams
+
+
 @registry.register_hparams
 def mfrl_original():
   return tf.contrib.training.HParams(
diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
new file mode 100644
index 000000000..3ca18c831
--- /dev/null
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -0,0 +1,308 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Connects dopamine to as the another rl traning framework."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gym
+from gym import spaces, Wrapper
+from gym.wrappers import TimeLimit
+import numpy as np
+import tensorflow as tf
+
+from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
+from tensor2tensor.rl.policy_learner import PolicyLearner
+
+try:
+  # pylint: disable=wrong-import-position
+  import cv2
+  from dopamine.agents.dqn import dqn_agent
+  from dopamine.atari import run_experiment
+  from dopamine.agents.dqn.dqn_agent import NATURE_DQN_OBSERVATION_SHAPE, \
+      NATURE_DQN_STACK_SIZE
+  from dopamine.replay_memory import circular_replay_buffer
+  from dopamine.replay_memory.circular_replay_buffer import \
+      OutOfGraphReplayBuffer, ReplayElement
+  # pylint: enable=wrong-import-position
+except ImportError:
+  # Generally we do not need dopamine in tensor2tensor
+  # We will raise exception if the code really tries to use it
+  pass
+
+
+class ResizeObservation(gym.ObservationWrapper):
+  def __init__(self, env, size=84):
+    """
+    Based on WarpFrame from openai baselines atari_wrappers.py
+    Dopamine also uses cv2.resize(..., interpolation=cv2.INTER_AREA).
+    """
+    gym.ObservationWrapper.__init__(self, env)
+    self.width = size
+    self.height = size
+    assert env.observation_space.dtype == np.uint8
+    self.observation_space = spaces.Box(
+        low=0, high=255,
+        shape=(self.height, self.width, env.observation_space.shape[2]),
+        dtype=np.uint8)
+
+  def observation(self, frame):
+    return cv2.resize(frame, (self.width, self.height),
+                      interpolation=cv2.INTER_AREA)
+
+
+class GameOverOnDone(Wrapper):
+
+  def __init__(self, env):
+    Wrapper.__init__(self, env)
+    self.game_over = False
+
+  def reset(self, **kwargs):
+    self.game_over = False
+    return self.env.reset(**kwargs)
+
+  def step(self, action):
+    ob, reward, done, info = self.env.step(action)
+    self.game_over = done
+    return ob, reward, done, info
+
+
+class _DQNAgent(dqn_agent.DQNAgent):
+  """ Modify dopamine DQNAgent to match our needs.
+
+  Allow passing batch_size and replay_capacity to ReplayBuffer, allow not using
+  (some of) terminal episode transitions in training.
+  """
+
+  def __init__(self, replay_capacity, batch_size,
+               generates_trainable_dones, **kwargs):
+    self._replay_capacity = replay_capacity
+    self._batch_size = batch_size
+    self._generates_trainable_dones = generates_trainable_dones
+    super(_DQNAgent, self).__init__(**kwargs)
+
+  def _build_replay_buffer(self, use_staging):
+    """Build WrappedReplayBuffer with custom OutOfGraphReplayBuffer"""
+    replay_buffer_kwargs = dict(
+        observation_shape=NATURE_DQN_OBSERVATION_SHAPE,
+        stack_size=NATURE_DQN_STACK_SIZE,
+        replay_capacity=self._replay_capacity,
+        batch_size=self._batch_size,
+        update_horizon=self.update_horizon,
+        gamma=self.gamma,
+        extra_storage_types=None,
+        observation_dtype=np.uint8,
+    )
+    replay_memory = _OutOfGraphReplayBuffer(
+        artificial_done=not self._generates_trainable_dones,
+        **replay_buffer_kwargs)
+
+    return circular_replay_buffer.WrappedReplayBuffer(
+        wrapped_memory=replay_memory, use_staging=use_staging,
+        **replay_buffer_kwargs)
+
+
+class _OutOfGraphReplayBuffer(OutOfGraphReplayBuffer):
+  """ Replay not sampling artificial_terminal transition.
+
+  Adds to stored tuples 'artificial_done' field (as last ReplayElement).
+  When sampling, ignores tuples for which artificial_done is True.
+
+  When adding new attributes check if there are loaded from disk, when using
+  load() method.
+
+  Attributes:
+      are_terminal_valid: A boolean indicating if newly added terminal
+      transitions should be marked as artificially done. Replay data
+      loaded from disk will not be overridden.
+  """
+
+  def __init__(self, artificial_done, **kwargs):
+    extra_storage_types = kwargs.pop('extra_storage_types', None) or []
+    extra_storage_types.append(
+        ReplayElement('artificial_done', (), np.uint8))
+    super(_OutOfGraphReplayBuffer, self).__init__(
+        extra_storage_types=extra_storage_types, **kwargs)
+    self._artificial_done = artificial_done
+
+  def is_valid_transition(self, index):
+    valid = super(_OutOfGraphReplayBuffer, self).is_valid_transition(index)
+    valid &= not self.get_artificial_done_stack(index).any()
+    return valid
+
+  def get_artificial_done_stack(self, index):
+    return self.get_range(self._store['artificial_done'],
+                          index - self._stack_size + 1,
+                          index + 1)
+
+  def add(self, observation, action, reward, terminal, *args):
+    """Append artificial_done to *args and run parent method."""
+    # If this will be a problem for maintenance, we could probably override
+    # DQNAgent.add() method instead.
+    artificial_done = self._artificial_done and terminal
+    args = list(args)
+    args.append(artificial_done)
+    return super(_OutOfGraphReplayBuffer, self).add(
+        observation, action, reward, terminal, *args)
+
+  def load(self, *args, **kwargs):
+    # Check that appropriate attributes are not overridden
+    are_terminal_valid = self._artificial_done
+    super(_OutOfGraphReplayBuffer, self).load(*args, **kwargs)
+    assert self._artificial_done == are_terminal_valid
+
+
+def get_create_agent(agent_kwargs):
+  def create_agent(sess, environment, summary_writer=None):
+    """Creates a DQN agent.
+
+    Simplified version of `dopamine.atari.train.create_agent`
+    """
+    return _DQNAgent(sess=sess, num_actions=environment.action_space.n,
+                     summary_writer=summary_writer,
+                     tf_device='/gpu:*', **agent_kwargs)
+
+  return create_agent
+
+
+def get_create_env_fun(batch_env_fn, time_limit):
+
+  def create_env_fun(game_name, sticky_actions=True):
+    del game_name, sticky_actions
+    batch_env = batch_env_fn(in_graph=False)
+    env = FlatBatchEnv(batch_env)
+    env = TimeLimit(env, max_episode_steps=time_limit)
+    env = ResizeObservation(env)  # pylint: disable=redefined-variable-type
+    env = GameOverOnDone(env)
+    return env
+
+  return create_env_fun
+
+
+def _parse_hparams(hparams):
+  prefixes = ["agent_", "optimizer_", "runner_", "replay_buffer_"]
+  ret = []
+
+  for prefix in prefixes:
+    ret_dict = {}
+    for key in hparams.values():
+      if prefix in key:
+        par_name = key[len(prefix):]
+        ret_dict[par_name] = hparams.get(key)
+    ret.append(ret_dict)
+
+  return ret
+
+
+def _get_optimizer(params):
+  assert params['class'] == "RMSProp", "RMSProp is the only one supported"
+  params.pop('class')
+  return tf.train.RMSPropOptimizer(**params)
+
+
+class DQNLearner(PolicyLearner):
+  """ Interface for learning dqn implemented in dopamine."""
+  def __init__(self, frame_stack_size, base_event_dir, agent_model_dir):
+    super(DQNLearner, self).__init__(frame_stack_size, base_event_dir,
+                                     agent_model_dir)
+    self.completed_iterations = 0
+
+  def _target_iteractions_and_steps(
+      self, num_env_steps, save_continuously, save_every_steps):
+
+    if save_continuously:
+      training_steps_per_iteration = min(num_env_steps, save_every_steps)
+      num_iterations_to_do = num_env_steps // training_steps_per_iteration
+    else:
+      num_iterations_to_do = 1
+      training_steps_per_iteration = num_env_steps
+    target_iterations = self.completed_iterations + num_iterations_to_do
+    return target_iterations, training_steps_per_iteration
+
+  def create_runner(self, env_fn, hparams, target_iterations,
+                    training_steps_per_iteration):
+    # pylint: disable=unbalanced-tuple-unpacking
+    agent_params, optimizer_params, \
+    runner_params, replay_buffer_params = _parse_hparams(hparams)
+    # pylint: enable=unbalanced-tuple-unpacking
+    optimizer = _get_optimizer(optimizer_params)
+    agent_params['optimizer'] = optimizer
+    agent_params.update(replay_buffer_params)
+    create_agent_fn = get_create_agent(agent_params)
+    runner = run_experiment.Runner(
+        game_name="unused_arg", sticky_actions="unused_arg",
+        base_dir=self.agent_model_dir, create_agent_fn=create_agent_fn,
+        create_environment_fn=get_create_env_fun(
+            env_fn, time_limit=hparams.time_limit),
+        evaluation_steps=0,
+        num_iterations=target_iterations,
+        training_steps=training_steps_per_iteration,
+        **runner_params
+    )
+    return runner
+
+  def train(
+      self, env_fn, hparams, simulated, save_continuously, epoch,
+      num_env_steps=None, env_step_multiplier=1, eval_env_fn=None,
+      report_fn=None
+  ):
+    # TODO(konradczechowski): evaluation during training (with eval_env_fun)
+    del epoch, eval_env_fn, simulated, report_fn
+    if num_env_steps is None:
+      num_env_steps = hparams.num_frames
+
+    target_iterations, training_steps_per_iteration = \
+      self._target_iteractions_and_steps(
+          num_env_steps=num_env_steps * env_step_multiplier,
+          save_continuously=save_continuously,
+          save_every_steps=hparams.save_every_steps,)
+
+    with tf.Graph().as_default():
+      runner = self.create_runner(env_fn, hparams, target_iterations,
+                                  training_steps_per_iteration)
+      runner.run_experiment()
+
+    self.completed_iterations = target_iterations
+
+  def evaluate(self, env_fn, hparams, stochastic):
+    target_iterations = 0
+    training_steps_per_iteration = 0
+    if not stochastic:
+      hparams.set_hparam('agent_epsilon_eval', 0.)
+
+    create_environment_fn = get_create_env_fun(env_fn,
+                                               time_limit=hparams.time_limit)
+    env = create_environment_fn(game_name="unused_arg",
+                                sticky_actions="unused_arg")
+
+    with tf.Graph().as_default():
+      runner = self.create_runner(env_fn, hparams, target_iterations,
+                                  training_steps_per_iteration)
+      agent = runner._agent
+      del runner
+      agent.eval = True
+
+      # TODO(konradczechowski): correct number of episodes, when this will
+      # be hparam
+      for _ in range(30):
+        # Run single episode
+        ob = env.reset()
+        action = agent.begin_episode(ob)
+        done = False
+        while not done:
+          ob, reward, done, _ = env.step(action)
+          action = agent.step(reward, ob)
diff --git a/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index 3af091628..1ca12b156 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -25,6 +25,22 @@
 import tensorflow as tf
 
 
+class FlatBatchEnv(Env):
+  def __init__(self, batch_env):
+    if batch_env.batch_size != 1:
+      raise ValueError("Number of environments in batch must be equal to one")
+    self.batch_env = batch_env
+    self.action_space = self.batch_env.action_space
+    self.observation_space = self.batch_env.observation_space
+
+  def step(self, action):
+    obs, rewards, dones = self.batch_env.step([action])
+    return obs[0], rewards[0], dones[0], {}
+
+  def reset(self):
+    return self.batch_env.reset()[0]
+
+
 # TODO(koz4k): Unify interfaces of batch envs.
 class SimulatedBatchGymEnv(Env):
   """SimulatedBatchEnv in a Gym-like interface, environments are  batched."""
diff --git a/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
index 9ef2c2da3..804b6b2f8 100644
--- a/tensor2tensor/rl/policy_learner.py
+++ b/tensor2tensor/rl/policy_learner.py
@@ -33,6 +33,7 @@ def train(
       num_env_steps=None, env_step_multiplier=1, eval_env_fn=None,
       report_fn=None
   ):
+    # TODO(konradczechowski): pass name_scope instead of epoch?
     # TODO(konradczechowski): move 'simulated' to  batch_env
     raise NotImplementedError()
 
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index b05e19ff3..5a890ab5d 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -41,6 +41,7 @@
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import trainer_model_based_params
+from tensor2tensor.rl.dopamine_connector import DQNLearner
 from tensor2tensor.rl.ppo_learner import PPOLearner
 from tensor2tensor.utils import trainer_lib
 
@@ -52,7 +53,8 @@
 
 
 LEARNERS = {
-    "ppo": PPOLearner
+    "ppo": PPOLearner,
+    "dqn": DQNLearner,
 }
 
 
diff --git a/tensor2tensor/rl/trainer_model_based_dqn_test.py b/tensor2tensor/rl/trainer_model_based_dqn_test.py
new file mode 100644
index 000000000..f1c8de410
--- /dev/null
+++ b/tensor2tensor/rl/trainer_model_based_dqn_test.py
@@ -0,0 +1,38 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Tiny run of trainer_model_based. Smoke test."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.rl import trainer_model_based
+
+import tensorflow as tf
+
+FLAGS = tf.flags.FLAGS
+
+
+class ModelRLExperimentTest(tf.test.TestCase):
+
+  def test_dqn_basic(self):
+    FLAGS.output_dir = tf.test.get_temp_dir()
+    FLAGS.loop_hparams_set = "rlmb_dqn_tiny"
+    FLAGS.schedule = "train"  # skip evaluation for world model training
+    trainer_model_based.main(None)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index a3905205c..9eab413a9 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -61,7 +61,7 @@ def _rlmb_base():
       # Use random starts when learning agent on simulated env.
       simulation_random_starts=True,
       # Flip the first random frame in PPO batch for the true beginning.
-      simulation_flip_first_random_for_beginning=True,
+      simulation_flip_first_random_for_beginning=False,
       intrinsic_reward_scale=0.,
       # Resizing.
       resize_height_factor=2,
@@ -148,6 +148,8 @@ def rlmb_ppo_base():
       real_ppo_learning_rate=1e-4,
       real_ppo_effective_num_agents=16,
       real_ppo_eval_every_epochs=0,
+
+      simulation_flip_first_random_for_beginning=True,
   )
   update_hparams(hparams, ppo_params)
   return hparams
@@ -158,6 +160,26 @@ def rlmb_base():
   return rlmb_ppo_base()
 
 
+@registry.register_hparams
+def rlmb_dqn_base():
+  hparams = _rlmb_base()
+  simulated_rollout_length = 10
+  dqn_params = dict(
+      base_algo="dqn",
+      base_algo_params="dqn_original_params",
+      real_batch_size=1,
+      simulated_batch_size=1,
+      dqn_agent_generates_trainable_dones=False,
+      eval_batch_size=1,
+      # Must be equal to dqn_time_limit for now
+      simulated_rollout_length=simulated_rollout_length,
+      dqn_time_limit=simulated_rollout_length,
+      simulation_flip_first_random_for_beginning=False,
+  )
+  update_hparams(hparams, dqn_params)
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_basetest():
   """Base setting but quicker with only 2 epochs."""
@@ -456,6 +478,23 @@ def rlmb_tiny():
   return rlmb_ppo_tiny()
 
 
+@registry.register_hparams
+def rlmb_dqn_tiny():
+  """Tiny set for testing."""
+  hparams = rlmb_dqn_base()
+  hparams = hparams.override_from_dict(_rlmb_tiny_overrides())
+  update_hparams(hparams, dict(
+      simulated_rollout_length=2,
+      dqn_time_limit=2,
+      dqn_num_frames=128,
+      real_dqn_replay_buffer_replay_capacity=100,
+      dqn_replay_buffer_replay_capacity=100,
+      real_dqn_agent_min_replay_history=10,
+      dqn_agent_min_replay_history=10,
+  ))
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_tiny_stochastic():
   """Tiny setting with a stochastic next-frame model."""

From 2167370937d31a9984c7e8824c9d416c59b7b481 Mon Sep 17 00:00:00 2001
From: artitw <artitw@gmail.com>
Date: Wed, 21 Nov 2018 13:22:50 -0500
Subject: [PATCH 1272/2720] Fix bAbi data generator and readme (#1235)

* fix bAbi data generator and readme

* Fix bAbi hparams deletion

* Fix bAbi hparams delete unecessary keys

* Fix bAbi hparams clean keys

* bAbi hparams delete keys
---
 README.md                                | 11 +++++++++++
 tensor2tensor/data_generators/babi_qa.py | 19 +++++++++++++------
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index e1702cd2c..b4b6b4ebf 100644
--- a/README.md
+++ b/README.md
@@ -47,6 +47,7 @@ pip install tensor2tensor && t2t-trainer \
 ### Contents
 
 * [Suggested Datasets and Models](#suggested-datasets-and-models)
+  * [Story, Question and Answer](#story-question-and-answer)
   * [Image Classification](#image-classification)
   * [Image Generation](#image-generation)
   * [Language Modeling](#language-modeling)
@@ -78,6 +79,16 @@ hyperparameters that we know works well in our setup. We usually
 run either on Cloud TPUs or on 8-GPU machines; you might need
 to modify the hyperparameters if you run on a different setup.
 
+### Story, Question and Answer
+
+For answering questions based on a story, use
+ 
+* the [bAbi][1] data-set: `--problem=babi_qa_concat_task1_1k`
+
+You can choose the bAbi task from the range [1,20] and the subset from 1k or 10k. To combine test data from all tasks into a single test set, use `--problem=babi_qa_concat_all_tasks_10k`
+
+[1] https://research.fb.com/downloads/babi/
+
 ### Image Classification
 
 For image classification, we have a number of standard data-sets:
diff --git a/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index 784508ecf..801db4fd8 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -34,10 +34,9 @@
 import os
 import shutil
 import tarfile
-
+import requests
 import six
 
-from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
@@ -109,8 +108,12 @@ def _prepare_babi_data(tmp_dir, data_dir):
   if not tf.gfile.Exists(data_dir):
     tf.gfile.MakeDirs(data_dir)
 
-  # TODO(dehghani@): find a solution for blocking user-agent (download)
-  file_path = generator_utils.maybe_download(tmp_dir, _TAR, _URL)
+  file_path = os.path.join(tmp_dir, _TAR)
+  headers = {'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36'}
+  resp = requests.get(_URL, headers=headers)
+  with open(file_path, 'wb') as f:
+    f.write(resp.content)
+
   tar = tarfile.open(file_path)
   tar.extractall(tmp_dir)
   tar.close()
@@ -449,8 +452,12 @@ def preprocess_example(self, example, unused_mode, unused_model_hparams):
   def hparams(self, defaults, unused_model_hparams):
     super(BabiQaConcat, self).hparams(defaults, unused_model_hparams)
     p = defaults
-    del p.modality['context']
-    del p.vocab_size['context']
+
+    if 'context' in p.modality:
+      del p.modality['context']
+
+    if 'context' in p.vocab_size:
+      del p.vocab_size['context']
 
 
 def _problems_to_register():

From eed1ccff9b2635cac03c877f0ecb281fde2ffe61 Mon Sep 17 00:00:00 2001
From: Chulayuth Asawaroengchai <twilightdema@gmail.com>
Date: Thu, 22 Nov 2018 01:27:13 +0700
Subject: [PATCH 1273/2720] Fix transformer_moe model has wrong logic in
 pre/postprocessing (#1233)

---
 tensor2tensor/models/research/transformer_moe.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index e6ee2a3f0..26129cd4f 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -93,8 +93,8 @@ def prepostprocess(fct):
       """Apply processing and capture the extra loss."""
       @expert_utils.add_var_scope()
       def decorated(x, *args, **kwargs):
-        x = dp_preprocess(x)
-        y, loss = fct(x, *args, **kwargs)
+        x_preprocessed = dp_preprocess(x)
+        y, loss = fct(x_preprocessed, *args, **kwargs)
         cache["extra_loss"] += loss
         return dp_postprocess(x, y)
       return decorated

From f091ea64f4d7e60ddb1d2fccad07c4a4b3b6a3f6 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Wed, 21 Nov 2018 10:30:21 -0800
Subject: [PATCH 1274/2720] internal merge of PR #1217

PiperOrigin-RevId: 222428117
---
 README.md                                     |  11 --
 setup.py                                      |   4 +-
 tensor2tensor/data_generators/babi_qa.py      |  19 +--
 tensor2tensor/models/research/rl.py           |   3 +-
 .../models/research/transformer_moe.py        |   4 +-
 tensor2tensor/rl/dopamine_connector.py        | 143 +++++++++++-------
 .../rl/envs/simulated_batch_gym_env.py        |   2 +
 .../rl/trainer_model_based_params.py          |   1 +
 8 files changed, 102 insertions(+), 85 deletions(-)

diff --git a/README.md b/README.md
index b4b6b4ebf..e1702cd2c 100644
--- a/README.md
+++ b/README.md
@@ -47,7 +47,6 @@ pip install tensor2tensor && t2t-trainer \
 ### Contents
 
 * [Suggested Datasets and Models](#suggested-datasets-and-models)
-  * [Story, Question and Answer](#story-question-and-answer)
   * [Image Classification](#image-classification)
   * [Image Generation](#image-generation)
   * [Language Modeling](#language-modeling)
@@ -79,16 +78,6 @@ hyperparameters that we know works well in our setup. We usually
 run either on Cloud TPUs or on 8-GPU machines; you might need
 to modify the hyperparameters if you run on a different setup.
 
-### Story, Question and Answer
-
-For answering questions based on a story, use
- 
-* the [bAbi][1] data-set: `--problem=babi_qa_concat_task1_1k`
-
-You can choose the bAbi task from the range [1,20] and the subset from 1k or 10k. To combine test data from all tasks into a single test set, use `--problem=babi_qa_concat_all_tasks_10k`
-
-[1] https://research.fb.com/downloads/babi/
-
 ### Image Classification
 
 For image classification, we have a number of standard data-sets:
diff --git a/setup.py b/setup.py
index bbf1f7e88..6c8d76cab 100644
--- a/setup.py
+++ b/setup.py
@@ -33,6 +33,7 @@
     ],
     install_requires=[
         'bz2file',
+        'dopamine-rl',
         'flask',
         'future',
         'gevent',
@@ -43,10 +44,11 @@
         'mesh-tensorflow',
         'numpy',
         'oauth2client',
+        'opencv-python',
         'requests',
         'scipy',
-        'sympy',
         'six',
+        'sympy',
         'tensorflow-probability',
         'tfds-nightly',
         'tqdm',
diff --git a/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index 801db4fd8..784508ecf 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -34,9 +34,10 @@
 import os
 import shutil
 import tarfile
-import requests
+
 import six
 
+from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
@@ -108,12 +109,8 @@ def _prepare_babi_data(tmp_dir, data_dir):
   if not tf.gfile.Exists(data_dir):
     tf.gfile.MakeDirs(data_dir)
 
-  file_path = os.path.join(tmp_dir, _TAR)
-  headers = {'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36'}
-  resp = requests.get(_URL, headers=headers)
-  with open(file_path, 'wb') as f:
-    f.write(resp.content)
-
+  # TODO(dehghani@): find a solution for blocking user-agent (download)
+  file_path = generator_utils.maybe_download(tmp_dir, _TAR, _URL)
   tar = tarfile.open(file_path)
   tar.extractall(tmp_dir)
   tar.close()
@@ -452,12 +449,8 @@ def preprocess_example(self, example, unused_mode, unused_model_hparams):
   def hparams(self, defaults, unused_model_hparams):
     super(BabiQaConcat, self).hparams(defaults, unused_model_hparams)
     p = defaults
-
-    if 'context' in p.modality:
-      del p.modality['context']
-
-    if 'context' in p.vocab_size:
-      del p.vocab_size['context']
+    del p.modality['context']
+    del p.vocab_size['context']
 
 
 def _problems_to_register():
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index befe3a775..6296306e1 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -233,8 +233,9 @@ def dqn_atari_base():
 
 @registry.register_hparams
 def dqn_original_params():
+  """dqn_original_params."""
   hparams = dqn_atari_base()
-  hparams.set_hparam('num_frames', int(1e6))
+  hparams.set_hparam("num_frames", int(1e6))
   return hparams
 
 
diff --git a/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index 26129cd4f..e6ee2a3f0 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -93,8 +93,8 @@ def prepostprocess(fct):
       """Apply processing and capture the extra loss."""
       @expert_utils.add_var_scope()
       def decorated(x, *args, **kwargs):
-        x_preprocessed = dp_preprocess(x)
-        y, loss = fct(x_preprocessed, *args, **kwargs)
+        x = dp_preprocess(x)
+        y, loss = fct(x, *args, **kwargs)
         cache["extra_loss"] += loss
         return dp_postprocess(x, y)
       return decorated
diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 3ca18c831..4843d45ae 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -19,53 +19,61 @@
 from __future__ import division
 from __future__ import print_function
 
+from dopamine.agents.dqn import dqn_agent
+from dopamine.agents.dqn.dqn_agent import NATURE_DQN_OBSERVATION_SHAPE
+from dopamine.agents.dqn.dqn_agent import NATURE_DQN_STACK_SIZE
+from dopamine.atari import run_experiment
+from dopamine.replay_memory import circular_replay_buffer
+from dopamine.replay_memory.circular_replay_buffer import OutOfGraphReplayBuffer
+from dopamine.replay_memory.circular_replay_buffer import ReplayElement
 import gym
-from gym import spaces, Wrapper
+from gym import spaces
+from gym import Wrapper
 from gym.wrappers import TimeLimit
 import numpy as np
-import tensorflow as tf
-
 from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
 from tensor2tensor.rl.policy_learner import PolicyLearner
+import tensorflow as tf
 
+# pylint: disable=g-import-not-at-top
 try:
-  # pylint: disable=wrong-import-position
   import cv2
-  from dopamine.agents.dqn import dqn_agent
-  from dopamine.atari import run_experiment
-  from dopamine.agents.dqn.dqn_agent import NATURE_DQN_OBSERVATION_SHAPE, \
-      NATURE_DQN_STACK_SIZE
-  from dopamine.replay_memory import circular_replay_buffer
-  from dopamine.replay_memory.circular_replay_buffer import \
-      OutOfGraphReplayBuffer, ReplayElement
-  # pylint: enable=wrong-import-position
 except ImportError:
-  # Generally we do not need dopamine in tensor2tensor
-  # We will raise exception if the code really tries to use it
-  pass
+  cv2 = None
+# pylint: enable=g-import-not-at-top
 
 
 class ResizeObservation(gym.ObservationWrapper):
+  """ TODO(konradczechowski): Add doc-string."""
+
   def __init__(self, env, size=84):
-    """
-    Based on WarpFrame from openai baselines atari_wrappers.py
+    """Based on WarpFrame from openai baselines atari_wrappers.py.
+
     Dopamine also uses cv2.resize(..., interpolation=cv2.INTER_AREA).
+
+    Args:
+      env: TODO(konradczechowski): Add doc-string.
+      size: TODO(konradczechowski): Add doc-string.
     """
     gym.ObservationWrapper.__init__(self, env)
     self.width = size
     self.height = size
     assert env.observation_space.dtype == np.uint8
     self.observation_space = spaces.Box(
-        low=0, high=255,
+        low=0,
+        high=255,
         shape=(self.height, self.width, env.observation_space.shape[2]),
         dtype=np.uint8)
 
   def observation(self, frame):
-    return cv2.resize(frame, (self.width, self.height),
-                      interpolation=cv2.INTER_AREA)
+    if not cv2:
+      return frame
+    return cv2.resize(
+        frame, (self.width, self.height), interpolation=cv2.INTER_AREA)
 
 
 class GameOverOnDone(Wrapper):
+  """TODO(konradczechowski): Add doc-string."""
 
   def __init__(self, env):
     Wrapper.__init__(self, env)
@@ -88,15 +96,15 @@ class _DQNAgent(dqn_agent.DQNAgent):
   (some of) terminal episode transitions in training.
   """
 
-  def __init__(self, replay_capacity, batch_size,
-               generates_trainable_dones, **kwargs):
+  def __init__(self, replay_capacity, batch_size, generates_trainable_dones,
+               **kwargs):
     self._replay_capacity = replay_capacity
     self._batch_size = batch_size
     self._generates_trainable_dones = generates_trainable_dones
     super(_DQNAgent, self).__init__(**kwargs)
 
   def _build_replay_buffer(self, use_staging):
-    """Build WrappedReplayBuffer with custom OutOfGraphReplayBuffer"""
+    """Build WrappedReplayBuffer with custom OutOfGraphReplayBuffer."""
     replay_buffer_kwargs = dict(
         observation_shape=NATURE_DQN_OBSERVATION_SHAPE,
         stack_size=NATURE_DQN_STACK_SIZE,
@@ -112,7 +120,8 @@ def _build_replay_buffer(self, use_staging):
         **replay_buffer_kwargs)
 
     return circular_replay_buffer.WrappedReplayBuffer(
-        wrapped_memory=replay_memory, use_staging=use_staging,
+        wrapped_memory=replay_memory,
+        use_staging=use_staging,
         **replay_buffer_kwargs)
 
 
@@ -127,14 +136,13 @@ class _OutOfGraphReplayBuffer(OutOfGraphReplayBuffer):
 
   Attributes:
       are_terminal_valid: A boolean indicating if newly added terminal
-      transitions should be marked as artificially done. Replay data
-      loaded from disk will not be overridden.
+        transitions should be marked as artificially done. Replay data loaded
+        from disk will not be overridden.
   """
 
   def __init__(self, artificial_done, **kwargs):
-    extra_storage_types = kwargs.pop('extra_storage_types', None) or []
-    extra_storage_types.append(
-        ReplayElement('artificial_done', (), np.uint8))
+    extra_storage_types = kwargs.pop("extra_storage_types", None) or []
+    extra_storage_types.append(ReplayElement("artificial_done", (), np.uint8))
     super(_OutOfGraphReplayBuffer, self).__init__(
         extra_storage_types=extra_storage_types, **kwargs)
     self._artificial_done = artificial_done
@@ -145,9 +153,8 @@ def is_valid_transition(self, index):
     return valid
 
   def get_artificial_done_stack(self, index):
-    return self.get_range(self._store['artificial_done'],
-                          index - self._stack_size + 1,
-                          index + 1)
+    return self.get_range(self._store["artificial_done"],
+                          index - self._stack_size + 1, index + 1)
 
   def add(self, observation, action, reward, terminal, *args):
     """Append artificial_done to *args and run parent method."""
@@ -156,8 +163,8 @@ def add(self, observation, action, reward, terminal, *args):
     artificial_done = self._artificial_done and terminal
     args = list(args)
     args.append(artificial_done)
-    return super(_OutOfGraphReplayBuffer, self).add(
-        observation, action, reward, terminal, *args)
+    return super(_OutOfGraphReplayBuffer, self).add(observation, action, reward,
+                                                    terminal, *args)
 
   def load(self, *args, **kwargs):
     # Check that appropriate attributes are not overridden
@@ -167,19 +174,33 @@ def load(self, *args, **kwargs):
 
 
 def get_create_agent(agent_kwargs):
+  """TODO(): Document."""
+
   def create_agent(sess, environment, summary_writer=None):
     """Creates a DQN agent.
 
     Simplified version of `dopamine.atari.train.create_agent`
+
+    Args:
+      sess: a session
+      environment: an environment
+      summary_writer: a summary writer.
+
+    Returns:
+      a DQN agent.
     """
-    return _DQNAgent(sess=sess, num_actions=environment.action_space.n,
-                     summary_writer=summary_writer,
-                     tf_device='/gpu:*', **agent_kwargs)
+    return _DQNAgent(
+        sess=sess,
+        num_actions=environment.action_space.n,
+        summary_writer=summary_writer,
+        tf_device="/gpu:*",
+        **agent_kwargs)
 
   return create_agent
 
 
 def get_create_env_fun(batch_env_fn, time_limit):
+  """TODO(konradczechowski): Add doc-string."""
 
   def create_env_fun(game_name, sticky_actions=True):
     del game_name, sticky_actions
@@ -194,6 +215,7 @@ def create_env_fun(game_name, sticky_actions=True):
 
 
 def _parse_hparams(hparams):
+  """TODO(konradczechowski): Add doc-string."""
   prefixes = ["agent_", "optimizer_", "runner_", "replay_buffer_"]
   ret = []
 
@@ -209,20 +231,21 @@ def _parse_hparams(hparams):
 
 
 def _get_optimizer(params):
-  assert params['class'] == "RMSProp", "RMSProp is the only one supported"
-  params.pop('class')
+  assert params["class"] == "RMSProp", "RMSProp is the only one supported"
+  params.pop("class")
   return tf.train.RMSPropOptimizer(**params)
 
 
 class DQNLearner(PolicyLearner):
   """ Interface for learning dqn implemented in dopamine."""
+
   def __init__(self, frame_stack_size, base_event_dir, agent_model_dir):
     super(DQNLearner, self).__init__(frame_stack_size, base_event_dir,
                                      agent_model_dir)
     self.completed_iterations = 0
 
-  def _target_iteractions_and_steps(
-      self, num_env_steps, save_continuously, save_every_steps):
+  def _target_iteractions_and_steps(self, num_env_steps, save_continuously,
+                                    save_every_steps):
 
     if save_continuously:
       training_steps_per_iteration = min(num_env_steps, save_every_steps)
@@ -240,26 +263,32 @@ def create_runner(self, env_fn, hparams, target_iterations,
     runner_params, replay_buffer_params = _parse_hparams(hparams)
     # pylint: enable=unbalanced-tuple-unpacking
     optimizer = _get_optimizer(optimizer_params)
-    agent_params['optimizer'] = optimizer
+    agent_params["optimizer"] = optimizer
     agent_params.update(replay_buffer_params)
     create_agent_fn = get_create_agent(agent_params)
     runner = run_experiment.Runner(
-        game_name="unused_arg", sticky_actions="unused_arg",
-        base_dir=self.agent_model_dir, create_agent_fn=create_agent_fn,
+        game_name="unused_arg",
+        sticky_actions="unused_arg",
+        base_dir=self.agent_model_dir,
+        create_agent_fn=create_agent_fn,
         create_environment_fn=get_create_env_fun(
             env_fn, time_limit=hparams.time_limit),
         evaluation_steps=0,
         num_iterations=target_iterations,
         training_steps=training_steps_per_iteration,
-        **runner_params
-    )
+        **runner_params)
     return runner
 
-  def train(
-      self, env_fn, hparams, simulated, save_continuously, epoch,
-      num_env_steps=None, env_step_multiplier=1, eval_env_fn=None,
-      report_fn=None
-  ):
+  def train(self,
+            env_fn,
+            hparams,
+            simulated,
+            save_continuously,
+            epoch,
+            num_env_steps=None,
+            env_step_multiplier=1,
+            eval_env_fn=None,
+            report_fn=None):
     # TODO(konradczechowski): evaluation during training (with eval_env_fun)
     del epoch, eval_env_fn, simulated, report_fn
     if num_env_steps is None:
@@ -282,17 +311,17 @@ def evaluate(self, env_fn, hparams, stochastic):
     target_iterations = 0
     training_steps_per_iteration = 0
     if not stochastic:
-      hparams.set_hparam('agent_epsilon_eval', 0.)
+      hparams.set_hparam("agent_epsilon_eval", 0.)
 
-    create_environment_fn = get_create_env_fun(env_fn,
-                                               time_limit=hparams.time_limit)
-    env = create_environment_fn(game_name="unused_arg",
-                                sticky_actions="unused_arg")
+    create_environment_fn = get_create_env_fun(
+        env_fn, time_limit=hparams.time_limit)
+    env = create_environment_fn(
+        game_name="unused_arg", sticky_actions="unused_arg")
 
     with tf.Graph().as_default():
       runner = self.create_runner(env_fn, hparams, target_iterations,
                                   training_steps_per_iteration)
-      agent = runner._agent
+      agent = runner._agent  # pylint: disable=protected-access
       del runner
       agent.eval = True
 
diff --git a/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index 1ca12b156..bb815d537 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -26,6 +26,8 @@
 
 
 class FlatBatchEnv(Env):
+  """TODO(konradczechowski): Add doc-string."""
+
   def __init__(self, batch_env):
     if batch_env.batch_size != 1:
       raise ValueError("Number of environments in batch must be equal to one")
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 9eab413a9..d7987b9ed 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -162,6 +162,7 @@ def rlmb_base():
 
 @registry.register_hparams
 def rlmb_dqn_base():
+  """rlmb_dqn_base params."""
   hparams = _rlmb_base()
   simulated_rollout_length = 10
   dqn_params = dict(

From 9977c5277f9088a0a690579b4e6df09b23a1395f Mon Sep 17 00:00:00 2001
From: Chulayuth Asawaroengchai <twilightdema@gmail.com>
Date: Wed, 21 Nov 2018 10:37:13 -0800
Subject: [PATCH 1275/2720] internal merge of PR #1233

PiperOrigin-RevId: 222429349
---
 tensor2tensor/models/research/transformer_moe.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index e6ee2a3f0..26129cd4f 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -93,8 +93,8 @@ def prepostprocess(fct):
       """Apply processing and capture the extra loss."""
       @expert_utils.add_var_scope()
       def decorated(x, *args, **kwargs):
-        x = dp_preprocess(x)
-        y, loss = fct(x, *args, **kwargs)
+        x_preprocessed = dp_preprocess(x)
+        y, loss = fct(x_preprocessed, *args, **kwargs)
         cache["extra_loss"] += loss
         return dp_postprocess(x, y)
       return decorated

From b2a7d767f4206feda54e8aecba6eb2ffa9d7b345 Mon Sep 17 00:00:00 2001
From: artitw <artitw@gmail.com>
Date: Wed, 21 Nov 2018 10:38:25 -0800
Subject: [PATCH 1276/2720] internal merge of PR #1235

PiperOrigin-RevId: 222429545
---
 README.md                                |  11 ++
 docs/walkthrough.md                      |  11 ++
 tensor2tensor/data_generators/babi_qa.py | 235 ++++++++++++-----------
 3 files changed, 143 insertions(+), 114 deletions(-)

diff --git a/README.md b/README.md
index e1702cd2c..b4b6b4ebf 100644
--- a/README.md
+++ b/README.md
@@ -47,6 +47,7 @@ pip install tensor2tensor && t2t-trainer \
 ### Contents
 
 * [Suggested Datasets and Models](#suggested-datasets-and-models)
+  * [Story, Question and Answer](#story-question-and-answer)
   * [Image Classification](#image-classification)
   * [Image Generation](#image-generation)
   * [Language Modeling](#language-modeling)
@@ -78,6 +79,16 @@ hyperparameters that we know works well in our setup. We usually
 run either on Cloud TPUs or on 8-GPU machines; you might need
 to modify the hyperparameters if you run on a different setup.
 
+### Story, Question and Answer
+
+For answering questions based on a story, use
+ 
+* the [bAbi][1] data-set: `--problem=babi_qa_concat_task1_1k`
+
+You can choose the bAbi task from the range [1,20] and the subset from 1k or 10k. To combine test data from all tasks into a single test set, use `--problem=babi_qa_concat_all_tasks_10k`
+
+[1] https://research.fb.com/downloads/babi/
+
 ### Image Classification
 
 For image classification, we have a number of standard data-sets:
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index e1702cd2c..b4b6b4ebf 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -47,6 +47,7 @@ pip install tensor2tensor && t2t-trainer \
 ### Contents
 
 * [Suggested Datasets and Models](#suggested-datasets-and-models)
+  * [Story, Question and Answer](#story-question-and-answer)
   * [Image Classification](#image-classification)
   * [Image Generation](#image-generation)
   * [Language Modeling](#language-modeling)
@@ -78,6 +79,16 @@ hyperparameters that we know works well in our setup. We usually
 run either on Cloud TPUs or on 8-GPU machines; you might need
 to modify the hyperparameters if you run on a different setup.
 
+### Story, Question and Answer
+
+For answering questions based on a story, use
+ 
+* the [bAbi][1] data-set: `--problem=babi_qa_concat_task1_1k`
+
+You can choose the bAbi task from the range [1,20] and the subset from 1k or 10k. To combine test data from all tasks into a single test set, use `--problem=babi_qa_concat_all_tasks_10k`
+
+[1] https://research.fb.com/downloads/babi/
+
 ### Image Classification
 
 For image classification, we have a number of standard data-sets:
diff --git a/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index 784508ecf..a11eaddca 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -34,10 +34,9 @@
 import os
 import shutil
 import tarfile
-
+import requests
 import six
 
-from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
@@ -49,32 +48,32 @@
 import tensorflow as tf
 
 
-_DIR_NAME = 'tasks_1-20_v1-2'
-_TAR = _DIR_NAME + '.tar.gz'
-_URL = 'http://www.thespermwhale.com/jaseweston/babi/' + _TAR
+_DIR_NAME = "tasks_1-20_v1-2"
+_TAR = _DIR_NAME + ".tar.gz"
+_URL = "http://www.thespermwhale.com/jaseweston/babi/" + _TAR
 
 _TASKS = {
-    'qa0': 'qa0_all-tasks',
-    'qa1': 'qa1_single-supporting-fact',
-    'qa2': 'qa2_two-supporting-facts',
-    'qa3': 'qa3_three-supporting-facts',
-    'qa4': 'qa4_two-arg-relations',
-    'qa5': 'qa5_three-arg-relations',
-    'qa6': 'qa6_yes-no-questions',
-    'qa7': 'qa7_counting',
-    'qa8': 'qa8_lists-sets',
-    'qa9': 'qa9_simple-negation',
-    'qa10': 'qa10_indefinite-knowledge',
-    'qa11': 'qa11_basic-coreference',
-    'qa12': 'qa12_conjunction',
-    'qa13': 'qa13_compound-coreference',
-    'qa14': 'qa14_time-reasoning',
-    'qa15': 'qa15_basic-deduction',
-    'qa16': 'qa16_basic-induction',
-    'qa17': 'qa17_positional-reasoning',
-    'qa18': 'qa18_size-reasoning',
-    'qa19': 'qa19_path-finding',
-    'qa20': 'qa20_agents-motivations'
+    "qa0": "qa0_all-tasks",
+    "qa1": "qa1_single-supporting-fact",
+    "qa2": "qa2_two-supporting-facts",
+    "qa3": "qa3_three-supporting-facts",
+    "qa4": "qa4_two-arg-relations",
+    "qa5": "qa5_three-arg-relations",
+    "qa6": "qa6_yes-no-questions",
+    "qa7": "qa7_counting",
+    "qa8": "qa8_lists-sets",
+    "qa9": "qa9_simple-negation",
+    "qa10": "qa10_indefinite-knowledge",
+    "qa11": "qa11_basic-coreference",
+    "qa12": "qa12_conjunction",
+    "qa13": "qa13_compound-coreference",
+    "qa14": "qa14_time-reasoning",
+    "qa15": "qa15_basic-deduction",
+    "qa16": "qa16_basic-induction",
+    "qa17": "qa17_positional-reasoning",
+    "qa18": "qa18_size-reasoning",
+    "qa19": "qa19_path-finding",
+    "qa20": "qa20_agents-motivations"
 }
 
 # A list of problem names that are registered by this module. This will get
@@ -91,7 +90,7 @@ def _normalize_string(raw_str):
   Returns:
    A string which is ready to be tokenized using split()
   """
-  return ' '.join(
+  return " ".join(
       token.strip()
       for token in tokenizer.encode(text_encoder.native_to_unicode(raw_str)))
 
@@ -109,8 +108,12 @@ def _prepare_babi_data(tmp_dir, data_dir):
   if not tf.gfile.Exists(data_dir):
     tf.gfile.MakeDirs(data_dir)
 
-  # TODO(dehghani@): find a solution for blocking user-agent (download)
-  file_path = generator_utils.maybe_download(tmp_dir, _TAR, _URL)
+  file_path = os.path.join(tmp_dir, _TAR)
+  headers = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36"}  # pylint: disable=line-too-long
+  resp = requests.get(_URL, headers=headers)
+  with open(file_path, "wb") as f:
+    f.write(resp.content)
+
   tar = tarfile.open(file_path)
   tar.extractall(tmp_dir)
   tar.close()
@@ -175,7 +178,7 @@ def _data_file(mode, task_id):
     Returns:
       data file path
     """
-    file_name = (_TASKS[task_id] + '_{}.txt')
+    file_name = (_TASKS[task_id] + "_{}.txt")
     return os.path.join(_DIR_NAME, subset, file_name.format(mode))
 
   def _all_task_raw_data_generator(tmp_dir, data_file, dataset_split):
@@ -187,48 +190,48 @@ def _all_task_raw_data_generator(tmp_dir, data_file, dataset_split):
       dataset_split: dataset split
     """
 
-    tf.logging.info('Preparing dataset of all task together')
-    globe_name = ('*_{}.txt')
+    tf.logging.info("Preparing dataset of all task together")
+    globe_name = ("*_{}.txt")
     files_name = os.path.join(
         tmp_dir, _DIR_NAME, subset,
-        globe_name.format('train' if dataset_split == problem.DatasetSplit.TRAIN
-                          else 'test'))
-    with tf.gfile.GFile(data_file, 'wb') as outfile:
+        globe_name.format("train" if dataset_split == problem.DatasetSplit.TRAIN
+                          else "test"))
+    with tf.gfile.GFile(data_file, "wb") as outfile:
       for filename in tf.gfile.Glob(files_name):
         if filename == data_file:
-          # don't want to copy the output into the output
+          # don"t want to copy the output into the output
           continue
-        with tf.gfile.GFile(filename, 'rb') as readfile:
+        with tf.gfile.GFile(filename, "rb") as readfile:
           shutil.copyfileobj(readfile, outfile)
 
   def _parse_answer(answer):
-    if (joint_training or babi_task_id in ['qa8', 'qa19', 'qa0'
-                                          ]):  # 'lists-sets' or 'path finding'
-      return ''.join([d for d in answer.split(',')])  # as a single token!
+    if (joint_training or babi_task_id in ["qa8", "qa19", "qa0"
+                                          ]):  # "lists-sets" or "path finding"
+      return "".join([d for d in answer.split(",")])  # as a single token!
     else:
       return answer
 
   if dataset_split == problem.DatasetSplit.TRAIN:
-    babi_train_task_id = 'qa0' if joint_training else babi_task_id
-    data_file = os.path.join(tmp_dir, _data_file('train', babi_train_task_id))
+    babi_train_task_id = "qa0" if joint_training else babi_task_id
+    data_file = os.path.join(tmp_dir, _data_file("train", babi_train_task_id))
   else:
-    data_file = os.path.join(tmp_dir, _data_file('test', babi_task_id))
+    data_file = os.path.join(tmp_dir, _data_file("test", babi_task_id))
 
-  if ((babi_task_id == 'qa0' or joint_training) and
+  if ((babi_task_id == "qa0" or joint_training) and
       not tf.gfile.Exists(os.path.join(tmp_dir, data_file))):
     _all_task_raw_data_generator(tmp_dir, data_file, dataset_split)
 
-  tf.logging.info('Parsing %s into training/testing instances...', data_file)
+  tf.logging.info("Parsing %s into training/testing instances...", data_file)
 
   babi_instances = []
-  with tf.gfile.GFile(data_file, mode='r') as f:
+  with tf.gfile.GFile(data_file, mode="r") as f:
     story = []
     for line in f:
-      line_num, line = line.strip().split(' ', 1)
+      line_num, line = line.strip().split(" ", 1)
       if int(line_num) == 1:
         story = []
-      if '\t' in line:
-        question, answer, _ = line.split('\t')
+      if "\t" in line:
+        question, answer, _ = line.split("\t")
         question = _normalize_string(question)
         substories = [s for s in story if s]
         answer = _parse_answer(answer)
@@ -239,7 +242,7 @@ def _parse_answer(answer):
         }
         babi_instances.append(instance)
 
-        story.append('')
+        story.append("")
       else:
         story.append(_normalize_string(line))
 
@@ -248,14 +251,14 @@ def _parse_answer(answer):
 
 class FeatureNames(object):
   """Feature names, i.e keys for storing babi_qa data in TFExamples."""
-  STORY = 'story'
-  QUESTION = 'question'
-  ANSWER = 'answer'
+  STORY = "story"
+  QUESTION = "question"
+  ANSWER = "answer"
 
   @classmethod
   def features(cls):
     for attr, value in cls.__dict__.items():
-      if not attr.startswith('__') and not callable(getattr(cls, attr)):
+      if not attr.startswith("__") and not callable(getattr(cls, attr)):
         yield value
 
 
@@ -265,15 +268,15 @@ class BabiQa(text_problems.QuestionAndContext2TextProblem):
   def __init__(self, *args, **kwargs):
 
     super(BabiQa, self).__init__(*args, **kwargs)
-    assert not self._was_reversed, 'This problem is not reversible!'
-    assert not self._was_copy, 'This problem is not copyable!'
+    assert not self._was_reversed, "This problem is not reversible!"
+    assert not self._was_copy, "This problem is not copyable!"
 
   @property
   def babi_subset(self):
     """The subset of dataset.
 
     This should be one of the following:
-    {'en', 'en-10k', 'shuffled', 'shuffled-10k'}
+    {"en", "en-10k", "shuffled", "shuffled-10k"}
     """
     raise NotImplementedError
 
@@ -282,25 +285,25 @@ def babi_task_id(self):
     """The id of the babi task.
 
     This should be one of the following:
-    {'qa0', 'qa1', 'qa1',...'q20'}, where qa0 means all tasks together.
+    {"qa0", "qa1", "qa1",..."q20"}, where qa0 means all tasks together.
     """
     raise NotImplementedError
 
   def dataset_filename(self):
-    return 'babi_qa_' + self.babi_subset + '_' + _TASKS[self.babi_task_id]
+    return "babi_qa_" + self.babi_subset + "_" + _TASKS[self.babi_task_id]
 
   @property
   def vocab_file(self):
-    return self.babi_subset + '_' + _TASKS[self.babi_task_id] + '.vocab'
+    return self.babi_subset + "_" + _TASKS[self.babi_task_id] + ".vocab"
 
   @property
   def dataset_splits(self):
     return [{
-        'split': problem.DatasetSplit.TRAIN,
-        'shards': 1,
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 1,
     }, {
-        'split': problem.DatasetSplit.EVAL,
-        'shards': 1,
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
     }]
 
   @property
@@ -345,11 +348,11 @@ def _generate_samples():
 
       """
       for example in examples:
-        context = ' '.join(example[FeatureNames.STORY])
+        context = " ".join(example[FeatureNames.STORY])
         yield {
-            'context': ' '.join(context.split()),
-            'inputs': ' '.join(example[FeatureNames.QUESTION].split()),
-            'targets': example[FeatureNames.ANSWER]
+            "context": " ".join(context.split()),
+            "inputs": " ".join(example[FeatureNames.QUESTION].split()),
+            "targets": example[FeatureNames.ANSWER]
         }
 
     return _generate_samples()
@@ -370,13 +373,13 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
     encoder = self.get_or_create_vocab(data_dir, tmp_dir)
     label_encoder = self.get_labels_encoder(data_dir)
     for sample in generator:
-      inputs = encoder.encode(sample['inputs'])
+      inputs = encoder.encode(sample["inputs"])
       inputs.append(text_encoder.EOS_ID)
-      context = encoder.encode(sample['context'])
+      context = encoder.encode(sample["context"])
       context.append(text_encoder.EOS_ID)
-      targets = label_encoder.encode(sample['targets'])
-      sample['targets'] = targets
-      yield {'inputs': inputs, 'context': context, 'targets': targets}
+      targets = label_encoder.encode(sample["targets"])
+      sample["targets"] = targets
+      yield {"inputs": inputs, "context": context, "targets": targets}
 
   def feature_encoders(self, data_dir):
     """Return a dict for encoding and decoding inference input/output.
@@ -390,7 +393,7 @@ def feature_encoders(self, data_dir):
     """
     encoders = (super(BabiQa, self).feature_encoders(data_dir))
     label_encoder = self.get_labels_encoder(data_dir)
-    encoders['targets'] = label_encoder  # bAbi as a classification task
+    encoders["targets"] = label_encoder  # bAbi as a classification task
     return encoders
 
   def generate_text_for_vocab(self, data_dir, tmp_dir):
@@ -402,9 +405,9 @@ def generate_text_for_vocab(self, data_dir, tmp_dir):
       for example in _babi_parser(tmp_dir, self.babi_task_id, self.babi_subset,
                                   dataset_split, self.joint_training):
 
-        context = ' '.join(example[FeatureNames.STORY])
-        yield ' '.join(context.split())
-        yield ' '.join(example[FeatureNames.QUESTION].split())
+        context = " ".join(example[FeatureNames.STORY])
+        yield " ".join(context.split())
+        yield " ".join(example[FeatureNames.QUESTION].split())
         yield example[FeatureNames.ANSWER]
 
   def hparams(self, defaults, unused_model_hparams):
@@ -417,14 +420,14 @@ def hparams(self, defaults, unused_model_hparams):
     """
     (super(BabiQa, self).hparams(defaults, unused_model_hparams))
     p = defaults
-    num_classes = self._encoders['targets'].vocab_size
-    p.modality = {'targets': modalities.ClassLabelModality}
-    p.vocab_size = {'targets': num_classes}
+    num_classes = self._encoders["targets"].vocab_size
+    p.modality = {"targets": modalities.ClassLabelModality}
+    p.vocab_size = {"targets": num_classes}
 
   def example_reading_spec(self):
     data_fields, data_items_to_decoders = (
         super(BabiQa, self).example_reading_spec())
-    data_fields['targets'] = tf.FixedLenFeature([1], tf.int64)
+    data_fields["targets"] = tf.FixedLenFeature([1], tf.int64)
     return (data_fields, data_items_to_decoders)
 
   def eval_metrics(self):
@@ -441,16 +444,20 @@ class BabiQaConcat(BabiQa):
 
   def preprocess_example(self, example, unused_mode, unused_model_hparams):
     sep = tf.convert_to_tensor([self.QUESTION_SEPARATOR_ID],
-                               dtype=example['inputs'].dtype)
-    example['inputs'] = tf.concat([example['inputs'], sep, example['context']],
+                               dtype=example["inputs"].dtype)
+    example["inputs"] = tf.concat([example["inputs"], sep, example["context"]],
                                   0)
     return example
 
   def hparams(self, defaults, unused_model_hparams):
     super(BabiQaConcat, self).hparams(defaults, unused_model_hparams)
     p = defaults
-    del p.modality['context']
-    del p.vocab_size['context']
+
+    if "context" in p.modality:
+      del p.modality["context"]
+
+    if "context" in p.vocab_size:
+      del p.vocab_size["context"]
 
 
 def _problems_to_register():
@@ -469,27 +476,27 @@ def _problems_to_register():
   # First define some problems using only concrete characters (i.e., no meta
   # characters).
   problems_on_different_tasks = {
-      'AllTasks': 'qa0',
-      'Task1': 'qa1',
-      'Task2': 'qa2',
-      'Task3': 'qa3',
-      'Task4': 'qa4',
-      'Task5': 'qa5',
-      'Task6': 'qa6',
-      'Task7': 'qa7',
-      'Task8': 'qa8',
-      'Task9': 'qa9',
-      'Task10': 'qa10',
-      'Task11': 'qa11',
-      'Task12': 'qa12',
-      'Task13': 'qa13',
-      'Task14': 'qa14',
-      'Task15': 'qa15',
-      'Task16': 'qa16',
-      'Task17': 'qa17',
-      'Task18': 'qa18',
-      'Task19': 'qa19',
-      'Task20': 'qa20',
+      "AllTasks": "qa0",
+      "Task1": "qa1",
+      "Task2": "qa2",
+      "Task3": "qa3",
+      "Task4": "qa4",
+      "Task5": "qa5",
+      "Task6": "qa6",
+      "Task7": "qa7",
+      "Task8": "qa8",
+      "Task9": "qa9",
+      "Task10": "qa10",
+      "Task11": "qa11",
+      "Task12": "qa12",
+      "Task13": "qa13",
+      "Task14": "qa14",
+      "Task15": "qa15",
+      "Task16": "qa16",
+      "Task17": "qa17",
+      "Task18": "qa18",
+      "Task19": "qa19",
+      "Task20": "qa20",
   }
   all_problems.update(problems_on_different_tasks)
 
@@ -503,26 +510,26 @@ def _register_babi_problems():
    class BabiQaConcatAllTasks_10k(EditSequenceRegexProblem):
      @property
      def babi_task_id(self):
-       return 'qa0'
+       return "qa0"
      @property
      def babi_subset(self):
-      return 'en-10k'
+      return "en-10k"
 
   It does not put the classes into the global namespace, so to access the class
-  we rely on the registry or this module's REGISTERED_PROBLEMS list.
+  we rely on the registry or this module"s REGISTERED_PROBLEMS list.
   It will be available as
 
-     registry.problem('babi_qa_concat_all_tasks_10k')
+     registry.problem("babi_qa_concat_all_tasks_10k")
 
   i.e., change camel case to snake case. Numbers are considered lower case
   characters for these purposes.
   """
-  for (subset, subset_suffix) in [('en', '_1k'), ('en-10k', '_10k')]:
+  for (subset, subset_suffix) in [("en", "_1k"), ("en-10k", "_10k")]:
     for problem_name, babi_task_id in six.iteritems(_problems_to_register()):
-      problem_class = type('BabiQaConcat' + problem_name + subset_suffix,
+      problem_class = type("BabiQaConcat" + problem_name + subset_suffix,
                            (BabiQaConcat,), {
-                               'babi_task_id': babi_task_id,
-                               'babi_subset': subset
+                               "babi_task_id": babi_task_id,
+                               "babi_subset": subset
                            })
       registry.register_problem(problem_class)
       REGISTERED_PROBLEMS.append(problem_class.name)

From b510fc92b66bb0540832d8b85e2fb25292450d52 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 21 Nov 2018 11:27:22 -0800
Subject: [PATCH 1277/2720] Fix bugs caused by cl/222104368.

PiperOrigin-RevId: 222437065
---
 tensor2tensor/data_generators/translate.py |  2 +-
 tensor2tensor/utils/trainer_lib.py         | 14 +++++++-------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index b9282e146..1a35c9459 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -88,7 +88,7 @@ def compute_bleu_summaries(hook_args):
       decode_hparams.decode_reference, decode_hparams.decode_to_file)
   values.append(tf.Summary.Value(tag="BLEU", simple_value=bleu))
   tf.logging.info("%s: BLEU = %6.2f" % (decode_hparams.decode_to_file, bleu))
-  if decode_hparams.mlperf_mode:
+  if hook_args.hparams.mlperf_mode:
     current_step = decode_hparams.mlperf_decode_step
     mlperf_log.transformer_print(
         key=mlperf_log.EVAL_TARGET, value=decode_hparams.mlperf_threshold)
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 87a7897a4..4849bbf43 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -481,17 +481,17 @@ def train_eval_and_decode(self):
         self._hparams.problem = problem
         self._hparams.problem_hparams = p_hparams
       mlperf_log.transformer_print(key=mlperf_log.EVAL_START)
-      if self._decode_hparams.mlperf_mode:
+      if self._hparams.mlperf_mode:
         self._decode_hparams.mlperf_decode_step = i + eval_steps
       self.decode(dataset_split=tf.estimator.ModeKeys.EVAL)
       d_hparams = self._decode_hparams
-      if d_hparams.mlperf_mode and d_hparams.mlperf_success:
+      if self._hparams.mlperf_mode and d_hparams.mlperf_success:
         mlperf_log.transformer_print(
             key=mlperf_log.RUN_STOP, value={"success": "true"})
         break
 
     d_hparams = self._decode_hparams
-    if d_hparams.mlperf_mode and not d_hparams.mlperf_success:
+    if self._hparams.mlperf_mode and not d_hparams.mlperf_success:
       mlperf_log.transformer_print(
           key=mlperf_log.RUN_STOP, value={"success": "false"})
 
@@ -580,7 +580,7 @@ def continuous_decode_on_train_data(self):
 
   def continuous_decode_on_eval_data(self):
     """Decode from dataset on new checkpoint."""
-    if self._decode_hparams.mlperf_mode:
+    if self._hparams.mlperf_mode:
       ckpt_generator = next_undecoded_checkpoint(self._hparams.model_dir)
     else:
       ckpt_generator = next_checkpoint(self._hparams.model_dir)
@@ -593,7 +593,7 @@ def continuous_decode_on_eval_data(self):
         continue
       # Decode the latest checkpoint by default.
       checkpoint_path = None
-      if self._decode_hparams.mlperf_mode:
+      if self._hparams.mlperf_mode:
         self._decode_hparams.mlperf_decode_step = current_step
         checkpoint_path = ckpt
 
@@ -602,13 +602,13 @@ def continuous_decode_on_eval_data(self):
           dataset_split=tf.estimator.ModeKeys.EVAL,
           checkpoint_path=checkpoint_path)
       d_hparams = self._decode_hparams
-      if d_hparams.mlperf_mode and d_hparams.mlperf_success:
+      if self._hparams.mlperf_mode and d_hparams.mlperf_success:
         mlperf_log.transformer_print(
             key=mlperf_log.RUN_STOP, value={"success": "true"})
         break
 
     d_hparams = self._decode_hparams
-    if d_hparams.mlperf_mode and not d_hparams.mlperf_success:
+    if self._hparams.mlperf_mode and not d_hparams.mlperf_success:
       mlperf_log.transformer_print(
           key=mlperf_log.RUN_STOP, value={"success": "false"})
 

From 3ee1e475a389b2f09afdcacd7a09df1859659b26 Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Wed, 21 Nov 2018 11:59:52 -0800
Subject: [PATCH 1278/2720] add overload_eval_metric_name to Transformer
 hparams to overwrite metric names (when specified and the number of problem
 is 1, we use this value instead of the problem name).

PiperOrigin-RevId: 222441575
---
 tensor2tensor/models/transformer.py | 4 ++++
 tensor2tensor/utils/metrics.py      | 8 +++++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 1bdf94a26..7ff34d72b 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1367,6 +1367,10 @@ def transformer_base_v1():
   hparams.add_hparam("moe_overhead_eval", 2.0)
   hparams.moe_num_experts = 16
   hparams.moe_loss_coef = 1e-3
+  # If specified, use this value instead of problem name in metrics.py.
+  # This is useful for programs that can automatically compare experiments side
+  #   by side based on the same metric names.
+  hparams.add_hparam("overload_eval_metric_name", "")
   return hparams
 
 
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 7150006c0..94c472798 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -609,7 +609,13 @@ def weights_fn_for_mp(problem_task_id):
 
       for metric in metrics:
         metric_fn = METRICS_FNS[metric]
-        metric_name = "metrics-%s/%s/%s" % (problem_name, target_name, metric)
+        overload_eval_metric_name = getattr(
+            model_hparams, "overload_eval_metric_name", None)
+        if len(problems) == 1 and overload_eval_metric_name:
+          metric_name = "metrics-%s/%s/%s" % (
+              overload_eval_metric_name, target_name, metric)
+        else:
+          metric_name = "metrics-%s/%s/%s" % (problem_name, target_name, metric)
         if metric == Metrics.IMAGE_SUMMARY:
           eval_metrics[metric_name] = make_image_wrapped_metric_fn(metric_fn)
         else:

From 83d3b8bc5428f63446a47f44e117bf92ff8bd133 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 21 Nov 2018 12:10:46 -0800
Subject: [PATCH 1279/2720] Allow recurrence in stochastic discrete RL model,
 add frame shapes.

PiperOrigin-RevId: 222443356
---
 tensor2tensor/data_generators/gym_env.py      | 14 ++++-
 tensor2tensor/models/research/rl.py           |  2 +-
 .../models/video/basic_deterministic.py       | 21 +++++++-
 .../video/basic_deterministic_params.py       |  1 +
 tensor2tensor/models/video/basic_recurrent.py |  1 +
 .../models/video/basic_stochastic.py          | 51 ++++++++++++++++++-
 tensor2tensor/rl/envs/simulated_batch_env.py  |  7 +--
 tensor2tensor/rl/trainer_model_based.py       |  3 +-
 .../rl/trainer_model_based_params.py          |  2 +-
 9 files changed, 92 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 5c4258f43..b90184956 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -692,10 +692,22 @@ def close(self):
 class DummyWorldModelProblem(EnvSimulationProblem):
   """Dummy Problem for world model prediction."""
 
-  def __init__(self, action_space, reward_range):
+  def __init__(self, action_space, reward_range, frame_height, frame_width):
     super(DummyWorldModelProblem, self).__init__()
     self.action_space = action_space
     self.reward_range = reward_range
+    self._frame_height = frame_height
+    self._frame_width = frame_width
+
+  @property
+  def frame_height(self):
+    """Height of each frame."""
+    return self._frame_height
+
+  @property
+  def frame_width(self):
+    """Width of each frame."""
+    return self._frame_width
 
 
 # Atari registration.
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 6296306e1..278d183e7 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -55,7 +55,7 @@ def ppo_base_v1():
   hparams.add_hparam("optimization_batch_size", 50)
   hparams.add_hparam("max_gradients_norm", 0.5)
   hparams.add_hparam("intrinsic_reward_scale", 0.)
-  hparams.add_hparam("logits_clip", 4.0)
+  hparams.add_hparam("logits_clip", 0.0)
   hparams.add_hparam("dropout_ppo", 0.1)
   hparams.add_hparam("effective_num_agents", None)
   return hparams
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index a460bac33..2c30b1d3a 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -61,6 +61,11 @@ def middle_network(self, layer, internal_states):
           x = common_layers.layer_norm(x + y)
     return x, internal_states
 
+  def update_internal_states_early(self, internal_states, frames):
+    """Update the internal states early in the network if requested."""
+    del frames
+    return internal_states
+
   def next_frame(self, frames, actions, rewards, target_frame,
                  internal_states, video_extra):
     del rewards, video_extra
@@ -70,9 +75,21 @@ def next_frame(self, frames, actions, rewards, target_frame,
     kernel2 = (4, 4)
     action = actions[-1]
 
-    # Embed the inputs.
-    stacked_frames = tf.concat(frames, axis=-1)
+    # Stack the inputs.
+    if internal_states is not None and hparams.concat_internal_states:
+      # Use the first part of the first internal state if asked to concatenate.
+      batch_size = common_layers.shape_list(frames[0])[0]
+      internal_state = internal_states[0][0][:batch_size, :, :, :]
+      stacked_frames = tf.concat(frames + [internal_state], axis=-1)
+    else:
+      stacked_frames = tf.concat(frames, axis=-1)
     inputs_shape = common_layers.shape_list(stacked_frames)
+
+    # Update internal states early if requested.
+    if hparams.concat_internal_states:
+      internal_states = self.update_internal_states_early(
+          internal_states, frames)
+
     # Using non-zero bias initializer below for edge cases of uniform inputs.
     x = tf.layers.dense(
         stacked_frames, filters, name="inputs_embed",
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 8e9644cbb..9dee014c3 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -46,6 +46,7 @@ def next_frame_basic_deterministic():
   hparams.add_hparam("num_compress_steps", 6)
   hparams.add_hparam("filter_double_steps", 2)
   hparams.add_hparam("pixel_sampling_temperature", 0.0)
+  hparams.add_hparam("concat_internal_states", False)
   return hparams
 
 
diff --git a/tensor2tensor/models/video/basic_recurrent.py b/tensor2tensor/models/video/basic_recurrent.py
index 41ad73bf1..c121712eb 100644
--- a/tensor2tensor/models/video/basic_recurrent.py
+++ b/tensor2tensor/models/video/basic_recurrent.py
@@ -62,6 +62,7 @@ def next_frame_basic_recurrent():
   hparams.hidden_size = 64
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 4
+  hparams.concat_internal_states = False
   hparams.add_hparam("num_lstm_layers", 2)
   hparams.add_hparam("num_lstm_filters", 256)
   return hparams
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 63af4214b..142047e16 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -66,6 +66,54 @@ class NextFrameBasicStochasticDiscrete(
     basic_deterministic.NextFrameBasicDeterministic):
   """Basic next-frame model with a tiny discrete latent."""
 
+  def init_internal_states(self):
+    if not self.hparams.concat_internal_states:
+      return None
+    # Hardcoded frame shapes.
+    max_batch_size = max(64, self.hparams.batch_size)
+    shape = [max_batch_size] + self.hparams.problem.frame_shape[:-1] + [32]
+    with tf.variable_scope("clean_scope_for_internal_state"):
+      v = tf.get_variable("state", shape, trainable=False,
+                          initializer=tf.zeros_initializer())
+    return [[v]]
+
+  def reset_internal_states_ops(self):
+    if not self.hparams.concat_internal_states:
+      return [[tf.no_op()]]
+    zeros = [[tf.zeros_like(s)] for s in self.internal_states[0]]
+    return self.save_internal_states_ops(zeros)
+
+  def load_internal_states_ops(self):
+    if not self.hparams.concat_internal_states:
+      return [[tf.no_op()]]
+    ops = [[s.read_value()] for s in self.internal_states[0]]
+    return ops
+
+  def save_internal_states_ops(self, internal_states):
+    if not self.hparams.concat_internal_states:
+      return [[tf.no_op()]]
+    ops = [[tf.assign(x, y)]
+           for x, y in zip(self.internal_states[0], internal_states[0])]
+    return ops
+
+  def update_internal_states_early(self, internal_states, frames):
+    """Update the internal states early in the network in GRU-like way."""
+    batch_size = common_layers.shape_list(frames[0])[0]
+    internal_state = internal_states[0][0][:batch_size, :, :, :]
+    state_activation = tf.concat([internal_state, frames[0]], axis=-1)
+    state_gate_candidate = tf.layers.conv2d(
+        state_activation, 64, (3, 3), padding="SAME", name="state_conv")
+    state_gate, state_candidate = tf.split(state_gate_candidate, 2, axis=-1)
+    state_gate = tf.nn.sigmoid(state_gate)
+    state_candidate = tf.tanh(state_candidate)
+    internal_state = internal_state * state_gate
+    internal_state += state_candidate * (1.0 - state_gate)
+    max_batch_size = max(64, self.hparams.batch_size)
+    diff_batch_size = max_batch_size - batch_size
+    internal_state = tf.pad(
+        internal_state, [[0, diff_batch_size], [0, 0], [0, 0], [0, 0]])
+    return [[internal_state]]
+
   def inject_latent(self, layer, inputs, target, action):
     """Inject a deterministic latent based on the target frame."""
     hparams = self.hparams
@@ -205,9 +253,10 @@ def next_frame_basic_stochastic_discrete():
   hparams.dropout = 0.15
   hparams.filter_double_steps = 3
   hparams.hidden_size = 96
-  hparams.learning_rate_constant = 0.005
+  hparams.learning_rate_constant = 0.002
   hparams.learning_rate_warmup_steps = 2000
   hparams.learning_rate_schedule = "linear_warmup * constant"
+  hparams.concat_internal_states = True
   hparams.add_hparam("bottleneck_bits", 256)
   hparams.add_hparam("bottleneck_noise", 0.1)
   hparams.add_hparam("discretize_warmup_steps", 40000)
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 481954a65..3d5acbcf9 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -93,8 +93,8 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
 
   def __init__(
       self, reward_range, observation_space, action_space, frame_stack_size,
-      initial_frame_chooser, batch_size, model_name, model_hparams, model_dir,
-      intrinsic_reward_scale=0.0
+      frame_height, frame_width, initial_frame_chooser, batch_size, model_name,
+      model_hparams, model_dir, intrinsic_reward_scale=0.0
   ):
     """Batch of environments inside the TensorFlow graph."""
     super(SimulatedBatchEnv, self).__init__(observation_space, action_space)
@@ -105,7 +105,8 @@ def __init__(
     self._intrinsic_reward_scale = intrinsic_reward_scale
 
     model_hparams = copy.copy(model_hparams)
-    problem = DummyWorldModelProblem(action_space, reward_range)
+    problem = DummyWorldModelProblem(action_space, reward_range,
+                                     frame_height, frame_width)
     trainer_lib.add_problem_hparams(model_hparams, problem)
     model_hparams.force_full_predict = True
     self._model = registry.model(model_name)(
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 5a890ab5d..70929b960 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -145,6 +145,7 @@ def make_simulated_env_fn(
       observation_space=real_env.observation_space,
       action_space=real_env.action_space,
       frame_stack_size=hparams.frame_stack_size,
+      frame_height=real_env.frame_height, frame_width=real_env.frame_width,
       initial_frame_chooser=initial_frame_chooser, batch_size=batch_size,
       model_name=hparams.generative_model,
       model_hparams=trainer_lib.create_hparams(hparams.generative_model_params),
@@ -582,7 +583,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       if report_fn:
         if report_metric == "mean_reward":
           metric_name = get_metric_name(
-              stochastic=False, max_num_noops=hparams.eval_max_num_noops,
+              stochastic=True, max_num_noops=hparams.eval_max_num_noops,
               clipped=False
           )
           report_fn(eval_metrics[metric_name], epoch)
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index d7987b9ed..bda11c7dc 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -781,7 +781,7 @@ def rlmb_games_problematic_for_ppo(rhp):
   rhp.set_categorical("loop.game", games)
   rhp.set_categorical("loop.base_algo_params", ["ppo_original_params"])
   rhp.set_discrete("model.moe_loss_coef", list(range(10)))
-  rhp.set_discrete("ppo.dropout_ppo", [0., 0.1])
+  rhp.set_discrete("ppo.logits_clip", [0., 4.0])
 
 
 @registry.register_ranged_hparams

From a1ab69902674b4d049e4e9090b01cbeec0d28a8a Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 21 Nov 2018 12:11:08 -0800
Subject: [PATCH 1280/2720] Move `tf.python_io.tf_record_iterator` references
 to `tf.compat.v1.tf_record_iterator`.

PiperOrigin-RevId: 222443392
---
 tensor2tensor/data_generators/bair_robot_pushing.py  |  3 ++-
 tensor2tensor/data_generators/generator_utils.py     |  2 +-
 .../data_generators/google_robot_pushing.py          |  2 +-
 tensor2tensor/data_generators/gym_env.py             |  2 +-
 tensor2tensor/data_generators/gym_env_test.py        | 12 +++++++-----
 tensor2tensor/data_generators/inspect_tfrecord.py    |  2 +-
 tensor2tensor/data_generators/problem.py             |  2 +-
 tensor2tensor/notebooks/t2t_problem.ipynb            |  2 +-
 8 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index bd9489fc8..0a349d972 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -112,7 +112,7 @@ def parse_frames(self, filenames):
 
     for f in filenames:
       print("Parsing ", f)
-      for serialized_example in tf.compat.v1.io.tf_record_iterator(f):
+      for serialized_example in tf.python_io.tf_record_iterator(f):
         x = tf.train.Example()
         x.ParseFromString(serialized_example)
         # there are 4 features per frame
@@ -176,3 +176,4 @@ def extra_reading_spec(self):
         "action": tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="action"),
     }
     return data_fields, decoders
+
diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 4d22a0651..c79babef1 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -464,7 +464,7 @@ def generate():
 
 
 def read_records(filename):
-  reader = tf.compat.v1.io.tf_record_iterator(filename)
+  reader = tf.python_io.tf_record_iterator(filename)
   records = []
   for record in reader:
     records.append(record)
diff --git a/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index bfa2b3c92..5b295fccf 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -83,7 +83,7 @@ def parse_frames(self, filename):
     action_key = "move/{}/commanded_pose/vec_pitch_yaw"
     state_key = "move/{}/endeffector/vec_pitch_yaw"
 
-    for serialized_example in tf.compat.v1.io.tf_record_iterator(filename):
+    for serialized_example in tf.python_io.tf_record_iterator(filename):
       x = tf.train.Example()
       x.ParseFromString(serialized_example)
       # there are 6 features per frame
diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index b90184956..b354b91d9 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -534,7 +534,7 @@ def _load_epoch_split(self, split, paths):
 
     for path in paths:
       this_shard_empty = True
-      for example in tf.compat.v1.io.tf_record_iterator(path):
+      for example in tf.python_io.tf_record_iterator(path):
         this_shard_empty = False
 
         result = tf.train.Example.FromString(example)
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 6d557297e..873de708a 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -164,7 +164,7 @@ def test_generates_data(self):
     self.assertTrue(filenames)
     for filename in filenames:
       path = os.path.join(self.out_dir, filename)
-      records = list(tf.compat.v1.io.tf_record_iterator(path))
+      records = list(tf.python_io.tf_record_iterator(path))
       self.assertTrue(records)
 
   def test_shards_per_epoch(self):
@@ -197,10 +197,12 @@ def test_frame_numbers_are_continuous(self):
     )
 
     frame_numbers = [
-        tf.train.Example.FromString(record).features.feature["frame_number"]
-        .int64_list.value[0]
-        for (_, paths) in env.splits_and_paths(self.out_dir) for path in paths
-        for record in tf.compat.v1.io.tf_record_iterator(path)
+        tf.train.Example.FromString(
+            record
+        ).features.feature["frame_number"].int64_list.value[0]
+        for (_, paths) in env.splits_and_paths(self.out_dir)
+        for path in paths
+        for record in tf.python_io.tf_record_iterator(path)
     ]
     last_frame_number = -1
     for frame_number in frame_numbers:
diff --git a/tensor2tensor/data_generators/inspect_tfrecord.py b/tensor2tensor/data_generators/inspect_tfrecord.py
index c28a9a32c..a245edb9e 100644
--- a/tensor2tensor/data_generators/inspect_tfrecord.py
+++ b/tensor2tensor/data_generators/inspect_tfrecord.py
@@ -56,7 +56,7 @@ def main(_):
     encoder = text_encoder.ByteTextEncoder()
   else:
     encoder = None
-  reader = tf.compat.v1.io.tf_record_iterator(FLAGS.input_filename)
+  reader = tf.python_io.tf_record_iterator(FLAGS.input_filename)
   total_sequences = 0
   total_input_tokens = 0
   total_target_tokens = 0
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index bf4f31e8d..3791123c1 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -169,7 +169,7 @@ def _file_num_records_cached(filename):
   if filename in _file_num_records_cache:
     return _file_num_records_cache[filename]
   ret = 0
-  for _ in tf.compat.v1.io.tf_record_iterator(filename):
+  for _ in tf.python_io.tf_record_iterator(filename):
     ret += 1
   _file_num_records_cache[filename] = ret
   return ret
diff --git a/tensor2tensor/notebooks/t2t_problem.ipynb b/tensor2tensor/notebooks/t2t_problem.ipynb
index d4a8df7b7..1eddf9e6b 100644
--- a/tensor2tensor/notebooks/t2t_problem.ipynb
+++ b/tensor2tensor/notebooks/t2t_problem.ipynb
@@ -65,7 +65,7 @@
         "\n",
         "\u003e\u003e[Viewing the generated data.](#scrollTo=MCqJhdnYgiG-)\n",
         "\n",
-        "\u003e\u003e\u003e[tf.compat.v1.io.tf_record_iterator](#scrollTo=uNpohcPXKsLN)\n",
+        "\u003e\u003e\u003e[tf.python_io.tf_record_iterator](#scrollTo=uNpohcPXKsLN)\n",
         "\n",
         "\u003e\u003e\u003e[Using tf.data.Dataset](#scrollTo=6o_1BHGQC5w5)\n",
         "\n",

From adf15a1093448ed44155abfa0af3ceb04bed741a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 21 Nov 2018 12:28:29 -0800
Subject: [PATCH 1281/2720] Make back-translation training always start at the
 start.

PiperOrigin-RevId: 222445517
---
 tensor2tensor/data_generators/problem.py        | 11 ++++++++++-
 tensor2tensor/data_generators/translate_enfr.py |  4 ++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 3791123c1..bb2acd3f4 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -316,6 +316,15 @@ def batch_size_means_tokens(self):
     """
     return False
 
+  @property
+  def skip_random_fraction_when_training(self):
+    """Skip a random number of examples at the beginning of training."""
+    # Skip a random fraction at the beginning of the stream.  The skip is
+    # essential for synchronous highly-parallel training to avoid multiple
+    # replicas reading the same data in lock-step. So keep this true unless
+    # you have a very specific setting in which it needs to be turned off.
+    return True
+
   def dataset_filename(self):
     return self.name
 
@@ -871,7 +880,7 @@ def define_shapes(example):
       # Repeat and skip a random number of records
       dataset = dataset.repeat()
 
-    if is_training:
+    if is_training and self.skip_random_fraction_when_training:
       data_files = tf.contrib.slim.parallel_reader.get_data_files(
           self.filepattern(data_dir, mode))
       #  In continuous_train_and_eval when switching between train and
diff --git a/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index 8087286d7..7c336fca1 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -154,6 +154,10 @@ def vocab_filename(self):
   def already_shuffled(self):
     return True
 
+  @property
+  def skip_random_fraction_when_training(self):
+    return False
+
   @property
   def backtranslate_data_filenames(self):
     """List of pairs of files with matched back-translated data."""

From 60b6ba6e3141cde04c906d7bc7fbdd5369776c62 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 21 Nov 2018 12:58:11 -0800
Subject: [PATCH 1282/2720] Add a long-term prediction config to RL.

PiperOrigin-RevId: 222449192
---
 tensor2tensor/models/video/basic_stochastic.py |  9 +++++++++
 tensor2tensor/rl/trainer_model_based_params.py | 12 ++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 142047e16..5dc6ade53 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -270,6 +270,15 @@ def next_frame_basic_stochastic_discrete():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_basic_stochastic_discrete_long():
+  """Conv model with stochastic discrete latent, long predictions."""
+  hparams = next_frame_basic_stochastic_discrete()
+  hparams.batch_size = 2
+  hparams.video_num_target_frames = 16
+  return hparams
+
+
 @registry.register_ranged_hparams
 def next_frame_stochastic_discrete_range(rhp):
   """Next frame stochastic discrete tuning grid."""
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index bda11c7dc..bb29767d9 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -276,6 +276,18 @@ def rlmb_base_stochastic_discrete():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_long_stochastic_discrete():
+  """Long setting with stochastic discrete model."""
+  hparams = rlmb_base()
+  hparams.learning_rate_bump = 1.0
+  hparams.grayscale = False
+  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete_long"
+  hparams.ppo_epochs_num = 2000
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_base_stochastic_recurrent():
   """Base setting with recurrent model."""

From 71371c7f40f1110159b81b46b4bbca7006996c22 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 21 Nov 2018 14:58:24 -0800
Subject: [PATCH 1283/2720] Add a multi-lingual German translation problem and
 a multi-task problem mixing a few of them.

PiperOrigin-RevId: 222464641
---
 tensor2tensor/data_generators/problem.py      |  5 +++++
 .../data_generators/translate_ende.py         | 10 +++++++++
 .../data_generators/wiki_multi_problems.py    | 21 +++++++++++++++++++
 tensor2tensor/utils/metrics.py                |  2 ++
 4 files changed, 38 insertions(+)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index bb2acd3f4..58c5bc56f 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -476,6 +476,11 @@ def __init__(self, was_reversed=False, was_copy=False):
     self._feature_info = None
     self._task_id = -1
 
+  @property
+  def was_reversed(self):
+    """Whether the problem was reversed."""
+    return self._was_reversed
+
   def get_feature_encoders(self, data_dir=None):
     if self._encoders is None:
       self._encoders = self.feature_encoders(data_dir)
diff --git a/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index c5fa94aff..259236967 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -26,6 +26,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.data_generators import translate
+from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -152,3 +153,12 @@ class TranslateEndeWmtCharacters(TranslateEndeWmt8k):
   @property
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
+
+
+@registry.register_problem
+class TranslateEndeWmtMulti64k(TranslateEndeWmt8k):
+  """Translation with muli-lingual vocabulary."""
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
diff --git a/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index 9ea44a94b..76f218770 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -23,6 +23,7 @@
 from tensor2tensor.data_generators import multi_problem
 from tensor2tensor.data_generators import multinli
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators import translate_ende
 from tensor2tensor.data_generators import translate_enfr
 from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
@@ -119,3 +120,23 @@ def __init__(self, was_reversed=False, was_copy=False):
   @property
   def vocab_type(self):
     return text_problems.VocabType.SUBWORD
+
+
+@registry.register_problem
+class LanguagemodelMultiWikiTranslate(multi_problem.MultiProblem):
+  """Wiki multi-lingual LM and multiple translations."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelMultiWikiTranslate, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelDeEnFrRoWiki64k())
+    self.task_list.append(translate_ende.TranslateEndeWmtMulti64k())
+    self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k())
+    self.task_list.append(translate_ende.TranslateEndeWmtMulti64k(
+        was_reversed=True))
+    self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k(
+        was_reversed=True))
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 94c472798..8775499a0 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -587,6 +587,8 @@ def weights_fn_for_mp(problem_task_id):
   eval_metrics = dict()
   for problem_instance in problems:
     problem_name = problem_instance.name
+    if problem_instance.was_reversed:
+      problem_name += "_rev"
     metrics = problem_instance.eval_metrics()
     if hasattr(model_hparams.problem, "task_list"):
       metrics = model_hparams.problem.eval_metrics()

From e9554fc27f2d59fab0e37d35fecfb3bbe0ee523f Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Wed, 21 Nov 2018 16:33:18 -0800
Subject: [PATCH 1284/2720] reward model for SV2P from intermediate layers.

PiperOrigin-RevId: 222475449
---
 tensor2tensor/models/video/savp.py        |  2 +-
 tensor2tensor/models/video/sv2p.py        | 43 ++++++++++++++++++-----
 tensor2tensor/models/video/sv2p_params.py |  8 +++--
 3 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index b8937c50c..01598c5fd 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -431,7 +431,7 @@ def construct_model(self, images, actions, rewards):
         all_action = tf.concat([action, action], axis=0)
         all_rewards = tf.concat([reward, reward], axis=0)
 
-        all_pred_images, lstm_state = self.construct_predictive_tower(
+        all_pred_images, lstm_state, _ = self.construct_predictive_tower(
             all_image, all_rewards, all_action, lstm_state, all_latents,
             concat_latent=True)
 
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 15a36b5db..7133dee3d 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -150,19 +150,43 @@ def reward_prediction(self, *args, **kwargs):
       return self.reward_prediction_basic(*args, **kwargs)
     elif model == "big":
       return self.reward_prediction_big(*args, **kwargs)
+    elif model == "mid":
+      return self.reward_prediction_mid(*args, **kwargs)
     else:
       raise ValueError("Unknown reward model %s" % model)
 
-  def reward_prediction_basic(self, input_images, input_reward, action, latent):
-    del input_reward, action, latent
+  def reward_prediction_basic(
+      self, input_images, input_reward, action, latent, mid_outputs):
+    del input_reward, action, latent, mid_outputs
     x = input_images
     x = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
     x = tfl.dense(x, 128, activation=tf.nn.relu, name="reward_pred")
     x = tf.expand_dims(x, axis=3)
     return x
 
-  def reward_prediction_big(self, input_images, input_reward, action, latent):
+  def reward_prediction_mid(
+      self, input_images, input_reward, action, latent, mid_outputs):
+    """Builds a reward prediction network from intermediate layers."""
+    encoded = []
+    for i, output in enumerate(mid_outputs):
+      enc = output
+      enc = tfl.conv2d(enc, 64, [3, 3], strides=(1, 1), activation=tf.nn.relu)
+      enc = tfl.conv2d(enc, 32, [3, 3], strides=(2, 2), activation=tf.nn.relu)
+      enc = tfl.conv2d(enc, 16, [3, 3], strides=(2, 2), activation=tf.nn.relu)
+      enc = tfl.flatten(enc)
+      enc = tfl.dense(enc, 64, activation=tf.nn.relu, name="rew_enc_%d" % i)
+      encoded.append(enc)
+    x = encoded
+    x = tf.stack(x, axis=1)
+    x = tfl.flatten(x)
+    x = tfl.dense(x, 256, activation=tf.nn.relu, name="rew_dense1")
+    x = tfl.dense(x, 128, activation=tf.nn.relu, name="rew_dense2")
+    return x
+
+  def reward_prediction_big(
+      self, input_images, input_reward, action, latent, mid_outputs):
     """Builds a reward prediction network."""
+    del mid_outputs
     conv_size = self.tinyify([32, 32, 16, 8])
 
     with tf.variable_scope("reward_pred", reuse=tf.AUTO_REUSE):
@@ -326,7 +350,8 @@ def construct_predictive_tower(
         output = tf.layers.dense(
             output, self.hparams.problem.num_channels * 256, name="logits")
 
-      return output, lstm_state
+      mid_outputs = [enc0, enc1, enc4, enc5, enc6]
+      return output, lstm_state, mid_outputs
 
   def video_features(
       self, all_frames, all_actions, all_rewards, all_raw_frames):
@@ -351,14 +376,14 @@ def next_frame(self, frames, actions, rewards, target_frame,
       if latent_mean is not None:
         extra_loss = self.get_extra_loss([latent_mean], [latent_std])
 
-    pred_image, internal_states = self.construct_predictive_tower(
+    pred_image, internal_states, mid_outputs = self.construct_predictive_tower(
         frames, None, actions, internal_states, latent)
 
     if not self.has_rewards:
       return pred_image, None, extra_loss, internal_states
 
     pred_reward = self.reward_prediction(
-        pred_image, actions, rewards, latent)
+        pred_image, actions, rewards, latent, mid_outputs)
     return pred_image, pred_reward, extra_loss, internal_states
 
 
@@ -418,7 +443,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
     extra_loss = 0.0
     latent = self.simple_discrete_latent_tower(frames, target_frame)
 
-    pred_image, internal_states = self.construct_predictive_tower(
+    pred_image, internal_states, _ = self.construct_predictive_tower(
         frames, None, actions, internal_states, latent, True)
 
     if not self.has_rewards:
@@ -549,7 +574,7 @@ def process_single_frame(prev_outputs, inputs):
           done_warm_start, groundtruth_items, generated_items, ss_func)
 
       # Prediction
-      pred_image, lstm_states = self.construct_predictive_tower(
+      pred_image, lstm_states, _ = self.construct_predictive_tower(
           input_image, None, action, lstm_states, latent)
 
       if self.hparams.reward_prediction:
@@ -760,7 +785,7 @@ def construct_model(self, images, actions, rewards):
       latent_stds.append(latent_std)
 
       # Prediction
-      pred_image, lstm_state = self.construct_predictive_tower(
+      pred_image, lstm_state, _ = self.construct_predictive_tower(
           input_image, input_reward, action, lstm_state, latent)
 
       if self.hparams.reward_prediction:
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index c9482c901..98dc9f658 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -60,8 +60,13 @@ def next_frame_sv2p_discrete():
   hparams = next_frame_sv2p()
   hparams.action_injection = "multiplicative"
   hparams.small_mode = True
-  hparams.add_hparam("bottleneck_bits", 16)
+  hparams.add_hparam("bottleneck_bits", 128)
   hparams.add_hparam("bottleneck_noise", 0.02)
+  hparams.add_hparam("discrete_warmup_steps", 40000)
+  hparams.add_hparam("full_latent_tower", False)
+  hparams.add_hparam("latent_predictor_state_size", 128)
+  hparams.add_hparam("latent_predictor_temperature", 0.5)
+  hparams.add_hparam("discretize_warmup_steps", 40000)
   return hparams
 
 
@@ -79,7 +84,6 @@ def next_frame_sv2p_atari():
   hparams.latent_loss_multiplier = 1e-3
   hparams.information_capacity = 0.0
   hparams.small_mode = True
-  hparams.internal_loss = True
   return hparams
 
 

From d2e741aebbd1edbec736c549e758f47d7536a387 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 22 Nov 2018 10:02:03 -0800
Subject: [PATCH 1285/2720] Lint fixes on README.md

PiperOrigin-RevId: 222552779
---
 README.md           | 10 ++++++----
 docs/walkthrough.md | 10 ++++++----
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index b4b6b4ebf..205e59283 100644
--- a/README.md
+++ b/README.md
@@ -82,12 +82,14 @@ to modify the hyperparameters if you run on a different setup.
 ### Story, Question and Answer
 
 For answering questions based on a story, use
- 
-* the [bAbi][1] data-set: `--problem=babi_qa_concat_task1_1k`
 
-You can choose the bAbi task from the range [1,20] and the subset from 1k or 10k. To combine test data from all tasks into a single test set, use `--problem=babi_qa_concat_all_tasks_10k`
+* the [bAbi](https://research.fb.com/downloads/babi/) data-set:
+ `--problem=babi_qa_concat_task1_1k`
+
+You can choose the bAbi task from the range [1,20] and the subset from 1k or
+10k. To combine test data from all tasks into a single test set, use
+`--problem=babi_qa_concat_all_tasks_10k`
 
-[1] https://research.fb.com/downloads/babi/
 
 ### Image Classification
 
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index b4b6b4ebf..205e59283 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -82,12 +82,14 @@ to modify the hyperparameters if you run on a different setup.
 ### Story, Question and Answer
 
 For answering questions based on a story, use
- 
-* the [bAbi][1] data-set: `--problem=babi_qa_concat_task1_1k`
 
-You can choose the bAbi task from the range [1,20] and the subset from 1k or 10k. To combine test data from all tasks into a single test set, use `--problem=babi_qa_concat_all_tasks_10k`
+* the [bAbi](https://research.fb.com/downloads/babi/) data-set:
+ `--problem=babi_qa_concat_task1_1k`
+
+You can choose the bAbi task from the range [1,20] and the subset from 1k or
+10k. To combine test data from all tasks into a single test set, use
+`--problem=babi_qa_concat_all_tasks_10k`
 
-[1] https://research.fb.com/downloads/babi/
 
 ### Image Classification
 

From bf9c555ed71a1228af5f1e91c311a7c62115cd0f Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 23 Nov 2018 08:34:04 -0800
Subject: [PATCH 1286/2720] Set decode_in_memory equal to True when one of
 FLAGS.decode_in_memory / decode_hparams.decode_in_memory is set to True.

PiperOrigin-RevId: 222629140
---
 tensor2tensor/bin/t2t_decoder.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 84644b3c4..8372ed0a7 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -72,7 +72,8 @@ def create_decode_hparams():
   decode_hp = decoding.decode_hparams(FLAGS.decode_hparams)
   decode_hp.shards = FLAGS.decode_shards
   decode_hp.shard_id = FLAGS.worker_id
-  decode_hp.decode_in_memory = FLAGS.decode_in_memory
+  decode_in_memory = FLAGS.decode_in_memory or decode_hp.decode_in_memory
+  decode_hp.decode_in_memory = decode_in_memory
   decode_hp.decode_to_file = FLAGS.decode_to_file
   decode_hp.decode_reference = FLAGS.decode_reference
   return decode_hp

From 244ed1f30eaa91f24b29ab7b748661b6f71495d4 Mon Sep 17 00:00:00 2001
From: Mike Dusenberry <dusenberrymw@google.com>
Date: Mon, 26 Nov 2018 18:58:26 -0800
Subject: [PATCH 1287/2720] Add positive constraint as a replacement for the
 std dev softplus constraint.

Due to the fact that variable constraints are applied *after* each optimization
step, the softplus function will cause the standard deviation variable to grow
regardless of the derivatives w.r.t. that variable.  This leads to difficult
optimization.  A standard deviation must be positive, so it could make sense to
simply apply a positivity constraint.

PiperOrigin-RevId: 222923785
---
 tensor2tensor/layers/bayes.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index cd978346f..2a8af7fff 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -24,21 +24,21 @@
 from tensorflow_probability import edward2 as ed
 
 
-class Softplus(tf.keras.constraints.Constraint):
-  """Softplus constraint."""
+class Positive(tf.keras.constraints.Constraint):
+  """Positive constraint."""
 
   def __init__(self, epsilon=tf.keras.backend.epsilon()):
     self.epsilon = epsilon
 
   def __call__(self, w):
-    return tf.nn.softplus(w) + self.epsilon
+    return tf.maximum(w, self.epsilon)
 
   def get_config(self):
     return {'epsilon': self.epsilon}
 
 
-def softplus():  # alias, following tf.keras.constraints
-  return Softplus()
+def positive():  # alias, following tf.keras.constraints
+  return Positive()
 
 
 # TODO(dusenberrymw): Restructure the implementation of a trainable initializer
@@ -68,7 +68,7 @@ def __init__(self,
                mean_regularizer=None,
                stddev_regularizer=None,
                mean_constraint=None,
-               stddev_constraint=softplus(),
+               stddev_constraint=positive(),
                seed=None,
                dtype=tf.float32):
     """Constructs the initializer."""

From 39280976fcd35d8de0bc43a33e940dd53f6ada12 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 26 Nov 2018 19:03:52 -0800
Subject: [PATCH 1288/2720] Add NGram layer.

PiperOrigin-RevId: 222924448
---
 tensor2tensor/layers/ngram.py      | 92 ++++++++++++++++++++++++++++++
 tensor2tensor/layers/ngram_test.py | 60 +++++++++++++++++++
 2 files changed, 152 insertions(+)
 create mode 100644 tensor2tensor/layers/ngram.py
 create mode 100644 tensor2tensor/layers/ngram_test.py

diff --git a/tensor2tensor/layers/ngram.py b/tensor2tensor/layers/ngram.py
new file mode 100644
index 000000000..a865d489d
--- /dev/null
+++ b/tensor2tensor/layers/ngram.py
@@ -0,0 +1,92 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""N-gram layer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import tensorflow as tf
+
+
+class NGram(tf.keras.layers.Layer):
+  r"""N-gram layer.
+
+  The layer takes as input an integer Tensor of shape [..., length], each
+  element of which is a token index in [0, input_dim). It returns a real-valued
+  Tensor of shape [..., num_ngrams], counting the number of times each n-gram
+  appears in a batch element. The total number of n-grams is
+
+  ```none
+  num_ngrams = \sum_{minval <= n < maxval} input_dim^n.
+  ```
+  """
+
+  def __init__(self, input_dim, minval, maxval, **kwargs):
+    """Constructs layer.
+
+    Args:
+      input_dim: int > 0. Size of the vocabulary, i.e. maximum integer index +
+        1.
+      minval: Lowest inclusive value of n for computing n-grams. For example,
+        setting it to 1 will compute starting from unigrams.
+      maxval: Highest non-inclusive value of n for computing n-grams. For
+        example, setting it to 3 will compute at most bigrams.
+      **kwargs: kwargs of parent class.
+    """
+    super(NGram, self).__init__(**kwargs)
+    self.input_dim = input_dim
+    self.minval = minval
+    self.maxval = maxval
+
+  def call(self, inputs):
+    batch_shape = tf.shape(inputs)[:-1]
+    length = tf.shape(inputs)[-1]
+    ngram_range_counts = []
+    for n in range(self.minval, self.maxval):
+      # Reshape inputs from [..., length] to [..., 1, length // n, n], dropping
+      # remainder elements. Each n-vector is an ngram.
+      reshaped_inputs = tf.reshape(
+          inputs[..., :(n * (length // n))],
+          tf.concat([batch_shape, [1], (length // n)[tf.newaxis], [n]], 0))
+      # Count the number of times each ngram appears in the input. We do so by
+      # checking whether each n-vector in the input is equal to each n-vector
+      # in a Tensor of all possible ngrams. The comparison is batched between
+      # the input Tensor of shape [..., 1, length // n, n] and the ngrams Tensor
+      # of shape [..., input_dim**n, 1, n].
+      ngrams = tf.reshape(
+          list(np.ndindex((self.input_dim,) * n)),
+          [1] * (len(inputs.shape)-1) + [self.input_dim**n, 1, n])
+      is_ngram = tf.equal(
+          tf.reduce_sum(tf.cast(tf.equal(reshaped_inputs, ngrams), tf.int32),
+                        axis=-1),
+          n)
+      ngram_counts = tf.reduce_sum(tf.cast(is_ngram, tf.float32), axis=-1)
+      ngram_range_counts.append(ngram_counts)
+    return tf.concat(ngram_range_counts, axis=-1)
+
+  def compute_output_shape(self, input_shape):
+    input_shape = tf.TensorShape(input_shape)
+    num_ngrams = sum([self.input_dim**n
+                      for n in range(self.minval, self.maxval)])
+    return input_shape[:-1].concatenate(num_ngrams)
+
+  def get_config(self):
+    config = {'minval': self.minval,
+              'maxval': self.maxval}
+    base_config = super(NGram, self).get_config()
+    return dict(list(base_config.items()) + list(config.items()))
diff --git a/tensor2tensor/layers/ngram_test.py b/tensor2tensor/layers/ngram_test.py
new file mode 100644
index 000000000..0eac013f1
--- /dev/null
+++ b/tensor2tensor/layers/ngram_test.py
@@ -0,0 +1,60 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 n-gram layer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import ngram
+
+import tensorflow as tf
+
+
+class NGramTest(tf.test.TestCase):
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testNGramLayerShape(self):
+    batch_size = 2
+    length = 8
+    vocab_size = 3
+    minval = 1
+    maxval = 4
+    inputs = tf.random_uniform(
+        [batch_size, length], minval=0, maxval=vocab_size, dtype=tf.int32)
+    layer = ngram.NGram(vocab_size, minval, maxval)
+    outputs = layer(inputs)
+    outputs_val = self.evaluate(outputs)
+    num_ngrams = sum([vocab_size**n for n in range(minval, maxval)])
+    self.assertEqual(outputs_val.shape, (batch_size, num_ngrams))
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testNGramLayerOutput(self):
+    inputs = tf.constant(
+        [[0, 0, 0, 0, 1],
+         [2, 1, 2, 1, 0]], dtype=tf.int32)
+    layer = ngram.NGram(3, minval=1, maxval=3)
+    outputs = layer(inputs)
+    expected_outputs = tf.constant(
+        [[4., 1., 0., 2., 0., 0., 0., 0., 0., 0., 0., 0.],
+         [1., 2., 2., 0., 0., 0., 0., 0., 0., 0., 2., 0.]], dtype=tf.float32)
+    outputs_val, expected_outputs_val = self.evaluate(
+        [outputs, expected_outputs])
+    self.assertAllEqual(outputs_val, expected_outputs_val)
+
+if __name__ == "__main__":
+  tf.test.main()
+

From cc82cb7cb50443a98b1ec1536738b57ef5c5f817 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 27 Nov 2018 17:57:45 -0800
Subject: [PATCH 1289/2720] Replace one_hot + matmul with tf.gather on R1
 indices for faster gather operation.

PiperOrigin-RevId: 223090838
---
 tensor2tensor/layers/common_layers.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 91954ee32..73ca1917a 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -272,10 +272,13 @@ def flatten4d3d(x):
 
 
 # TODO(noam): remove this function after TPUs do gather faster.
-def gather(params, indices):
+def gather(params, indices, dtype=tf.float32):
   """Version of tf.gather that works faster on tpu."""
+  if not is_xla_compiled():
+    return tf.gather(params, indices)
+  vocab_size = params.get_shape().as_list()[0]
   indices_flat = tf.reshape(indices, [-1])
-  out = tf.gather(params, indices_flat)
+  out = tf.matmul(tf.one_hot(indices_flat, vocab_size, dtype=dtype), params)
   out = reshape_like(out, tf.expand_dims(indices, -1))
   return out
 
@@ -349,7 +352,7 @@ def embedding(x,
     if not tf.contrib.eager.in_eager_mode():
       embedding_var = convert_gradient_to_tensor(embedding_var)
     x = dropout_no_scaling(x, 1.0 - symbol_dropout_rate)
-    emb_x = gather(embedding_var, x)
+    emb_x = gather(embedding_var, x, dtype)
     if multiplier != 1.0:
       emb_x *= multiplier
     static_shape = emb_x.shape.as_list()

From 9e6f462eeff0af799ba488ff745d90ce3e566a9d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 28 Nov 2018 09:45:49 -0800
Subject: [PATCH 1290/2720] Disable trainer_model_based_dqn_test since recent
 Dopamine changes break it.

PiperOrigin-RevId: 223184686
---
 tensor2tensor/rl/trainer_model_based_dqn_test.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based_dqn_test.py b/tensor2tensor/rl/trainer_model_based_dqn_test.py
index f1c8de410..d8af9de9d 100644
--- a/tensor2tensor/rl/trainer_model_based_dqn_test.py
+++ b/tensor2tensor/rl/trainer_model_based_dqn_test.py
@@ -18,7 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.rl import trainer_model_based
+# from tensor2tensor.rl import trainer_model_based
 
 import tensorflow as tf
 
@@ -28,10 +28,13 @@
 class ModelRLExperimentTest(tf.test.TestCase):
 
   def test_dqn_basic(self):
-    FLAGS.output_dir = tf.test.get_temp_dir()
-    FLAGS.loop_hparams_set = "rlmb_dqn_tiny"
-    FLAGS.schedule = "train"  # skip evaluation for world model training
-    trainer_model_based.main(None)
+    # TODO(afrozm): The latest changes in Dopamine break this test, so
+    # temporarily disabling this test.
+    pass
+    # FLAGS.output_dir = tf.test.get_temp_dir()
+    # FLAGS.loop_hparams_set = "rlmb_dqn_tiny"
+    # FLAGS.schedule = "train"  # skip evaluation for world model training
+    # trainer_model_based.main(None)
 
 
 if __name__ == "__main__":

From 45cad3371b96b8ab0ac6cea717b4b81bc5af5b00 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 28 Nov 2018 10:12:42 -0800
Subject: [PATCH 1291/2720] Add dims to targets_masks to match the shape of
 predictions.

PiperOrigin-RevId: 223189328
---
 tensor2tensor/utils/metrics.py      | 11 +++++++++++
 tensor2tensor/utils/metrics_test.py |  2 +-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 8775499a0..3d21cf854 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -250,9 +250,20 @@ def padded_neg_log_perplexity_with_masking(
     labels,
     features,
     weights_fn=None):
+  """Average log-perplexity with custom targets_mask."""
   del weights_fn
   if "targets_mask" not in features:
     raise ValueError("masked_neg_log_perplexity requires targets_mask feature")
+
+  # Features are 4 dimensional, so we need to reshape the targets_mask to match
+  # the shape of the labels. A lot of models rely on these features being 4D,
+  # so it's best to update the shape of the mask.
+  extended_targets_mask_shape = common_layers.shape_list(
+      features["targets_mask"])
+  extended_targets_mask_shape.extend([1, 1])
+  features["targets_mask"] = tf.reshape(features["targets_mask"],
+                                        shape=extended_targets_mask_shape)
+
   mask_fn = lambda labels: features["targets_mask"]
   return padded_neg_log_perplexity(predictions, labels, mask_fn)
 
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 26427a30a..4d25a47b4 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -121,7 +121,7 @@ def testNegativeLogPerplexityMasked(self):
     predictions = np.random.randint(4, size=(12, 12, 12, 1))
     targets = np.random.randint(4, size=(12, 12, 12, 1))
     features = {
-        'targets_mask': tf.to_float(tf.not_equal(targets, 0))
+        'targets_mask': tf.to_float(tf.ones([12, 12]))
     }
     with self.test_session() as session:
       scores, _ = metrics.padded_neg_log_perplexity_with_masking(

From fe61b5d8b28ecbc0715c937857f062ffe2e568cb Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Wed, 28 Nov 2018 10:44:28 -0800
Subject: [PATCH 1292/2720] Add positional embedding for targets and latents.
 With this I now get comparable results using z_size only 6-8.

PiperOrigin-RevId: 223195434
---
 .../models/research/transformer_vae.py        | 21 ++++++++++++-------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 7d84245fd..58db6dc5b 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -383,6 +383,13 @@ def ae_transformer_internal(inputs,
     targets, _ = common_layers.pad_to_same_length(
         targets, max_targets_len_from_inputs,
         final_length_divisible_by=2**hparams.num_compress_steps)
+    # Add positional information
+    targets_shape = common_layers.shape_list(targets)
+    targets = tf.reshape(targets, [targets_shape[0], targets_shape[1],
+                                   targets_shape[3]])
+    targets = common_attention.add_positional_embedding(
+        targets, hparams.max_length, name="targets_position")
+    targets = tf.reshape(targets, shape=targets_shape)
     if hparams.word_dropout:
       mask = tf.random_uniform(shape=common_layers.shape_list(targets),
                                minval=0.0, maxval=1.0)
@@ -390,6 +397,7 @@ def ae_transformer_internal(inputs,
                                tf.zeros_like(targets))
     else:
       targets_noisy = targets
+
     targets_c = compress(targets_noisy, inputs, False, hparams, "compress")
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
       # Compress and bottleneck.
@@ -463,14 +471,11 @@ def bn_inputs():
         latents_dense = embed(cache)
     # Postprocess.
     d = latents_dense
-    latent_len = common_layers.shape_list(latents_dense)[1]
-    if isinstance(latent_len, tf.Tensor):
-      # TODO(trandustin): Fix this in a better manner.
-      latent_len = max(1000, hparams.max_length)
-    pos = tf.get_variable("pos", [1, latent_len + 1, 1, hparams.hidden_size])
-    pos = pos[:, :common_layers.shape_list(latents_dense)[1] + 1, :, :]
-    latents_dense = tf.pad(latents_dense,
-                           [[0, 0], [1, 0], [0, 0], [0, 0]]) + pos
+    d_shape = common_layers.shape_list(d)
+    d = tf.reshape(d, [d_shape[0], d_shape[1], d_shape[3]])
+    d = common_attention.add_positional_embedding(
+        d, hparams.max_length, name="latents_position")
+    d = tf.reshape(d, shape=d_shape)
 
     # decompressing the dense latents
     for i in range(hparams.num_compress_steps):

From 7ad13fe5ec1ffcc584f473f3e201ff3b6ea3ebf0 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 28 Nov 2018 12:54:50 -0800
Subject: [PATCH 1293/2720] Adds VGG-features based cosine similarity to
 t2t-metrics.

PiperOrigin-RevId: 223219373
---
 .../metrics/vgg_cosine_similarity.py          | 103 ++++++++++++++++++
 .../metrics/vgg_cosine_similarity_test.py     |  41 +++++++
 2 files changed, 144 insertions(+)
 create mode 100644 tensor2tensor/metrics/vgg_cosine_similarity.py
 create mode 100644 tensor2tensor/metrics/vgg_cosine_similarity_test.py

diff --git a/tensor2tensor/metrics/vgg_cosine_similarity.py b/tensor2tensor/metrics/vgg_cosine_similarity.py
new file mode 100644
index 000000000..9a2468f82
--- /dev/null
+++ b/tensor2tensor/metrics/vgg_cosine_similarity.py
@@ -0,0 +1,103 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""VGG Cosine similarity metric."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import functools
+import tensorflow as tf
+from tensorflow.contrib.framework.python.ops import arg_scope
+from tensorflow.contrib.slim.python.slim.nets import vgg
+from tensorflow_models.slim.preprocessing import vgg_preprocessing
+
+
+def vgg_features(x):
+  """Computes VGG features of input x.
+
+  Args:
+    x: 4-D Tensor, shape=(batch_size, height, width, channels)
+  Returns:
+    features: A list of tensors of VGG-features corresponding to x.
+  """
+  preprocess_single = functools.partial(
+      vgg_preprocessing.preprocess_image, output_height=224, output_width=224,
+      is_training=False)
+  x = tf.map_fn(preprocess_single, x)
+  _, features = vgg.vgg_16(x, num_classes=1000, is_training=False)
+
+  # filter fully connected end-points
+  return [t for n, t in features.items() if "fc" not in n]
+
+
+def vgg_cosine_similarity(images1, images2):
+  """VGG cosine similarity between images1[i] and images2[i].
+
+  For every feature obtained from VGG, the cosine similarity is computed across
+  the channels and averaged spatially. This is then averaged across
+  all VGG features.
+
+  Args:
+    images1: 4-D Tensor, shape=(batch_size, height, width, n_channels)
+    images2: 4-D Tensor, shape=(batch_size, height, width, n_channels)
+  Returns:
+    similarity: 1-D Tensor, shape=(batch_size,)
+  """
+  with arg_scope(vgg.vgg_arg_scope()):
+    img1_features = vgg_features(images1)
+    tf.get_variable_scope().reuse_variables()
+    img2_features = vgg_features(images2)
+
+  all_dists = []
+  for img1_feat, img2_feat in zip(img1_features, img2_features):
+
+    # Computes cosine similarity across channels, i.e dot-product and sum.
+    img1_feat = tf.nn.l2_normalize(img1_feat, axis=-1)
+    img2_feat = tf.nn.l2_normalize(img2_feat, axis=-1)
+    distance = img1_feat * img2_feat
+
+    # Computes mean of the distance spatially.
+    curr_cosine_dist = tf.reduce_mean(
+        tf.reduce_sum(distance, axis=-1), axis=[1, 2])
+    all_dists.append(curr_cosine_dist)
+  all_dists = tf.stack(all_dists)
+
+  # Average across VGG features.
+  return tf.reduce_mean(all_dists, axis=0)
+
+
+def vgg_cosine_similarity_from_ckpt(images1, images2, ckpt_path):
+  """VGG Cosine similarity using a trained VGG ckpt.
+
+  Args:
+    images1: 4-D NumPy array, shape=(batch_size, height, width, n_channels)
+    images2: 4-D NumPy array, shape=(batch_size, height, width, n_channels)
+    ckpt_path: Path to trained VGG ckpt.
+  Returns:
+    similarity: 1-D NumPy array, shape=(batch_size,)
+  """
+  with tf.Graph().as_default():
+    images1 = tf.convert_to_tensor(images1, dtype=tf.float32)
+    images2 = tf.convert_to_tensor(images2, dtype=tf.float32)
+    vgg_sim = vgg_cosine_similarity(images1, images2)
+
+
+    saver = tf.train.Saver()
+    with tf.Session() as sess:
+      saver.restore(sess, ckpt_path)
+      vgg_sim_np = sess.run(vgg_sim)
+    return vgg_sim_np
+
diff --git a/tensor2tensor/metrics/vgg_cosine_similarity_test.py b/tensor2tensor/metrics/vgg_cosine_similarity_test.py
new file mode 100644
index 000000000..9383e400d
--- /dev/null
+++ b/tensor2tensor/metrics/vgg_cosine_similarity_test.py
@@ -0,0 +1,41 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.metrics.vgg_cosine_similarity."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.metrics import vgg_cosine_similarity
+import tensorflow as tf
+
+
+class VggCosineSimilarityTest(tf.test.TestCase):
+
+  def test_vgg_cosine_similarity(self):
+    with tf.Graph().as_default():
+      rng = np.random.RandomState(0)
+      x = np.asarray(rng.randn(16, 64, 64, 3), dtype=np.float32)
+      cos_sim_t = vgg_cosine_similarity.vgg_cosine_similarity(x, x)
+      with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+        cos_sim_np = sess.run(cos_sim_t)
+        self.assertTrue(np.allclose(cos_sim_np, 1.0))
+
+
+if __name__ == '__main__':
+  tf.test.main()

From 6afd639b7d048554bdb8582382f5ecaf82e36520 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 28 Nov 2018 13:01:45 -0800
Subject: [PATCH 1294/2720] MTF rewrite_stack_variables() now handles
 assignments as well.  Fully checkpoint-compatible and
 mathematically-compatible.  Turned on by default via the "autostack" option
 in Lowering.  Client code needs to change to call optimizer.apply_grads()
 instead of optimizer.apply_grad().

PiperOrigin-RevId: 223220539
---
 tensor2tensor/utils/mtf_model.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/mtf_model.py b/tensor2tensor/utils/mtf_model.py
index 5b034f2f5..6ff50e4e2 100644
--- a/tensor2tensor/utils/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -114,9 +114,7 @@ def estimator_model_fn(cls,
       mtf_lr = mtf.import_tf_tensor(
           mesh, tf.convert_to_tensor(lr, dtype=tf.float32), mtf.Shape([]))
       optimizer = mtf.optimize.make_optimizer(hparams, mtf_lr)
-      update_ops = []
-      for grad, var in zip(var_grads, graph.trainable_variables):
-        update_ops.extend(optimizer.apply_grad(grad, var))
+      update_ops = optimizer.apply_grads(var_grads, graph.trainable_variables)
 
     lowering = mtf.Lowering(graph, {mesh: mesh_impl})
 

From 6f1cca45a5803831b680ca119427926e8b94bb8c Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 28 Nov 2018 13:51:39 -0800
Subject: [PATCH 1295/2720] Minor nits in test code.

PiperOrigin-RevId: 223229676
---
 tensor2tensor/data_generators/gym_env_test.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 873de708a..e5532f8c0 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -41,16 +41,17 @@ class TestEnv(gym.Env):
   """
 
   action_space = Discrete(1)
+  # TODO(afrozm): Gym's Box has a bug for uint8 type, which doesn't allow
+  # sampling, send them a PR. Till that time let this be np.int64
   observation_space = Box(
-      low=0, high=255, shape=(2, 6, 3), dtype=np.uint8
+      low=0, high=255, shape=(2, 6, 3), dtype=np.int64
   )
 
   def __init__(self):
     self._counter = 0
 
   def _generate_ob(self):
-    return np.random.randint(255, size=self.observation_space.shape,
-                             dtype=self.observation_space.dtype)
+    return self.observation_space.sample()
 
   def step(self, action):
     done = self._counter % 2 == 1

From b1fb27ba43d93e4ff0e01b7cea16c85fd58089a8 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 28 Nov 2018 14:51:58 -0800
Subject: [PATCH 1296/2720] Uses tensorflow_models which T2T doesn't import.

PiperOrigin-RevId: 223240537
---
 .../metrics/vgg_cosine_similarity.py          | 103 ------------------
 .../metrics/vgg_cosine_similarity_test.py     |  41 -------
 2 files changed, 144 deletions(-)
 delete mode 100644 tensor2tensor/metrics/vgg_cosine_similarity.py
 delete mode 100644 tensor2tensor/metrics/vgg_cosine_similarity_test.py

diff --git a/tensor2tensor/metrics/vgg_cosine_similarity.py b/tensor2tensor/metrics/vgg_cosine_similarity.py
deleted file mode 100644
index 9a2468f82..000000000
--- a/tensor2tensor/metrics/vgg_cosine_similarity.py
+++ /dev/null
@@ -1,103 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""VGG Cosine similarity metric."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-import functools
-import tensorflow as tf
-from tensorflow.contrib.framework.python.ops import arg_scope
-from tensorflow.contrib.slim.python.slim.nets import vgg
-from tensorflow_models.slim.preprocessing import vgg_preprocessing
-
-
-def vgg_features(x):
-  """Computes VGG features of input x.
-
-  Args:
-    x: 4-D Tensor, shape=(batch_size, height, width, channels)
-  Returns:
-    features: A list of tensors of VGG-features corresponding to x.
-  """
-  preprocess_single = functools.partial(
-      vgg_preprocessing.preprocess_image, output_height=224, output_width=224,
-      is_training=False)
-  x = tf.map_fn(preprocess_single, x)
-  _, features = vgg.vgg_16(x, num_classes=1000, is_training=False)
-
-  # filter fully connected end-points
-  return [t for n, t in features.items() if "fc" not in n]
-
-
-def vgg_cosine_similarity(images1, images2):
-  """VGG cosine similarity between images1[i] and images2[i].
-
-  For every feature obtained from VGG, the cosine similarity is computed across
-  the channels and averaged spatially. This is then averaged across
-  all VGG features.
-
-  Args:
-    images1: 4-D Tensor, shape=(batch_size, height, width, n_channels)
-    images2: 4-D Tensor, shape=(batch_size, height, width, n_channels)
-  Returns:
-    similarity: 1-D Tensor, shape=(batch_size,)
-  """
-  with arg_scope(vgg.vgg_arg_scope()):
-    img1_features = vgg_features(images1)
-    tf.get_variable_scope().reuse_variables()
-    img2_features = vgg_features(images2)
-
-  all_dists = []
-  for img1_feat, img2_feat in zip(img1_features, img2_features):
-
-    # Computes cosine similarity across channels, i.e dot-product and sum.
-    img1_feat = tf.nn.l2_normalize(img1_feat, axis=-1)
-    img2_feat = tf.nn.l2_normalize(img2_feat, axis=-1)
-    distance = img1_feat * img2_feat
-
-    # Computes mean of the distance spatially.
-    curr_cosine_dist = tf.reduce_mean(
-        tf.reduce_sum(distance, axis=-1), axis=[1, 2])
-    all_dists.append(curr_cosine_dist)
-  all_dists = tf.stack(all_dists)
-
-  # Average across VGG features.
-  return tf.reduce_mean(all_dists, axis=0)
-
-
-def vgg_cosine_similarity_from_ckpt(images1, images2, ckpt_path):
-  """VGG Cosine similarity using a trained VGG ckpt.
-
-  Args:
-    images1: 4-D NumPy array, shape=(batch_size, height, width, n_channels)
-    images2: 4-D NumPy array, shape=(batch_size, height, width, n_channels)
-    ckpt_path: Path to trained VGG ckpt.
-  Returns:
-    similarity: 1-D NumPy array, shape=(batch_size,)
-  """
-  with tf.Graph().as_default():
-    images1 = tf.convert_to_tensor(images1, dtype=tf.float32)
-    images2 = tf.convert_to_tensor(images2, dtype=tf.float32)
-    vgg_sim = vgg_cosine_similarity(images1, images2)
-
-
-    saver = tf.train.Saver()
-    with tf.Session() as sess:
-      saver.restore(sess, ckpt_path)
-      vgg_sim_np = sess.run(vgg_sim)
-    return vgg_sim_np
-
diff --git a/tensor2tensor/metrics/vgg_cosine_similarity_test.py b/tensor2tensor/metrics/vgg_cosine_similarity_test.py
deleted file mode 100644
index 9383e400d..000000000
--- a/tensor2tensor/metrics/vgg_cosine_similarity_test.py
+++ /dev/null
@@ -1,41 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.metrics.vgg_cosine_similarity."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-from tensor2tensor.metrics import vgg_cosine_similarity
-import tensorflow as tf
-
-
-class VggCosineSimilarityTest(tf.test.TestCase):
-
-  def test_vgg_cosine_similarity(self):
-    with tf.Graph().as_default():
-      rng = np.random.RandomState(0)
-      x = np.asarray(rng.randn(16, 64, 64, 3), dtype=np.float32)
-      cos_sim_t = vgg_cosine_similarity.vgg_cosine_similarity(x, x)
-      with tf.Session() as sess:
-        sess.run(tf.global_variables_initializer())
-        cos_sim_np = sess.run(cos_sim_t)
-        self.assertTrue(np.allclose(cos_sim_np, 1.0))
-
-
-if __name__ == '__main__':
-  tf.test.main()

From 20f705417f9fce36c1c568d1fc2ee1a65373c336 Mon Sep 17 00:00:00 2001
From: Youngwook Kim <youngwook.kim@gmail.com>
Date: Thu, 29 Nov 2018 08:11:08 +0900
Subject: [PATCH 1297/2720] Fix MRPC link (#1247)

---
 tensor2tensor/data_generators/mrpc.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/mrpc.py b/tensor2tensor/data_generators/mrpc.py
index 14822428d..46ed39168 100644
--- a/tensor2tensor/data_generators/mrpc.py
+++ b/tensor2tensor/data_generators/mrpc.py
@@ -38,8 +38,8 @@ class MSRParaphraseCorpus(text_problems.TextConcat2ClassProblem):
   # Link to data from GLUE: https://gluebenchmark.com/tasks
   DEV_IDS = ("https://firebasestorage.googleapis.com/v0/b/"
              "mtl-sentence-representations.appspot.com/o/"
-             "data%2FWNLI.zip?alt=media&token=068ad0a0-ded7-"
-             "4bd7-99a5-5e00222e0faf")
+             "data%2Fmrpc_dev_ids.tsv?alt=media&token=ec5c0836-31d5-"
+             "48f4-b431-7480817f1adc")
   MRPC_TRAIN = ("https://s3.amazonaws.com/senteval/senteval_data/"
                 "msr_paraphrase_train.txt")
   MRPC_TEST = ("https://s3.amazonaws.com/senteval/senteval_data/"

From 4fa7587b2e22629e6f10fc54fc0e7dda55ff87a3 Mon Sep 17 00:00:00 2001
From: Yuwen Yan <ybbaigo@gmail.com>
Date: Thu, 29 Nov 2018 07:15:41 +0800
Subject: [PATCH 1298/2720] add problems, conll2002_es_ner and conll2002_nl_ner
 (#1253)

---
 tensor2tensor/data_generators/conll_ner.py | 84 ++++++++++++++++++++++
 1 file changed, 84 insertions(+)
 create mode 100644 tensor2tensor/data_generators/conll_ner.py

diff --git a/tensor2tensor/data_generators/conll_ner.py b/tensor2tensor/data_generators/conll_ner.py
new file mode 100644
index 000000000..3d685c076
--- /dev/null
+++ b/tensor2tensor/data_generators/conll_ner.py
@@ -0,0 +1,84 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Data generators for CoNLL dataset."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import zipfile
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+@registry.register_problem
+class Conll2002Ner(text_problems.Text2textTmpdir):
+  """Base class for CoNLL2002 problems."""
+  def source_data_files(self, dataset_split):
+    """Files to be passed to generate_samples."""
+    raise NotImplementedError()
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    del data_dir
+
+    url = 'https://raw.githubusercontent.com/nltk/nltk_data/gh-pages/packages/corpora/conll2002.zip' # pylint: disable=line-too-long
+    compressed_filename = os.path.basename(url)
+    compressed_filepath = os.path.join(tmp_dir, compressed_filename)
+    generator_utils.maybe_download(tmp_dir, compressed_filename, url)
+
+    compressed_dir = compressed_filepath.strip(".zip")
+
+    filenames = self.source_data_files(dataset_split)
+    for filename in filenames:
+      filepath = os.path.join(compressed_dir, filename)
+      if not tf.gfile.Exists(filepath):
+        with zipfile.ZipFile(compressed_filepath, 'r') as corpus_zip:
+          corpus_zip.extractall(tmp_dir)
+      with tf.gfile.GFile(filepath, mode="r") as cur_file:
+        words, tags = [], []
+        for line in cur_file:
+          line_split = line.strip().split()
+          if len(line_split) == 0:
+            yield {"inputs": str.join(" ", words),
+                   "targets": str.join(" ", tags)}
+            words, tags = [], []
+            continue
+          words.append(line_split[0])
+          tags.append(line_split[2])
+        if len(words) != 0:
+          yield {"inputs": str.join(" ", words), "targets": str.join(" ", tags)}
+
+@registry.register_problem
+class Conll2002EsNer(Conll2002Ner):
+  """Problem spec for CoNLL2002 Spanish named entity task."""
+  TRAIN_FILES = ["esp.train"]
+  EVAL_FILES = ["esp.testa", "esp.testb"]
+  def source_data_files(self, dataset_split):
+    is_training = dataset_split == problem.DatasetSplit.TRAIN
+    return self.TRAIN_FILES if is_training else self.EVAL_FILES
+
+@registry.register_problem
+class Conll2002NlNer(Conll2002Ner):
+  """Problem spec for CoNLL2002 Dutch named entity task."""
+  TRAIN_FILES = ["ned.train"]
+  EVAL_FILES = ["ned.testa", "ned.testb"]
+  def source_data_files(self, dataset_split):
+    is_training = dataset_split == problem.DatasetSplit.TRAIN
+    return self.TRAIN_FILES if is_training else self.EVAL_FILES

From eacbb184fe241b475a0cd5361956038e1cf7c5c1 Mon Sep 17 00:00:00 2001
From: Youngwook Kim <youngwook.kim@gmail.com>
Date: Wed, 28 Nov 2018 15:11:41 -0800
Subject: [PATCH 1299/2720] internal merge of PR #1247

PiperOrigin-RevId: 223244151
---
 tensor2tensor/data_generators/conll_ner.py | 84 ----------------------
 1 file changed, 84 deletions(-)
 delete mode 100644 tensor2tensor/data_generators/conll_ner.py

diff --git a/tensor2tensor/data_generators/conll_ner.py b/tensor2tensor/data_generators/conll_ner.py
deleted file mode 100644
index 3d685c076..000000000
--- a/tensor2tensor/data_generators/conll_ner.py
+++ /dev/null
@@ -1,84 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Data generators for CoNLL dataset."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-import zipfile
-
-from tensor2tensor.data_generators import generator_utils
-from tensor2tensor.data_generators import problem
-from tensor2tensor.data_generators import text_problems
-from tensor2tensor.utils import registry
-import tensorflow as tf
-
-@registry.register_problem
-class Conll2002Ner(text_problems.Text2textTmpdir):
-  """Base class for CoNLL2002 problems."""
-  def source_data_files(self, dataset_split):
-    """Files to be passed to generate_samples."""
-    raise NotImplementedError()
-
-  def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    del data_dir
-
-    url = 'https://raw.githubusercontent.com/nltk/nltk_data/gh-pages/packages/corpora/conll2002.zip' # pylint: disable=line-too-long
-    compressed_filename = os.path.basename(url)
-    compressed_filepath = os.path.join(tmp_dir, compressed_filename)
-    generator_utils.maybe_download(tmp_dir, compressed_filename, url)
-
-    compressed_dir = compressed_filepath.strip(".zip")
-
-    filenames = self.source_data_files(dataset_split)
-    for filename in filenames:
-      filepath = os.path.join(compressed_dir, filename)
-      if not tf.gfile.Exists(filepath):
-        with zipfile.ZipFile(compressed_filepath, 'r') as corpus_zip:
-          corpus_zip.extractall(tmp_dir)
-      with tf.gfile.GFile(filepath, mode="r") as cur_file:
-        words, tags = [], []
-        for line in cur_file:
-          line_split = line.strip().split()
-          if len(line_split) == 0:
-            yield {"inputs": str.join(" ", words),
-                   "targets": str.join(" ", tags)}
-            words, tags = [], []
-            continue
-          words.append(line_split[0])
-          tags.append(line_split[2])
-        if len(words) != 0:
-          yield {"inputs": str.join(" ", words), "targets": str.join(" ", tags)}
-
-@registry.register_problem
-class Conll2002EsNer(Conll2002Ner):
-  """Problem spec for CoNLL2002 Spanish named entity task."""
-  TRAIN_FILES = ["esp.train"]
-  EVAL_FILES = ["esp.testa", "esp.testb"]
-  def source_data_files(self, dataset_split):
-    is_training = dataset_split == problem.DatasetSplit.TRAIN
-    return self.TRAIN_FILES if is_training else self.EVAL_FILES
-
-@registry.register_problem
-class Conll2002NlNer(Conll2002Ner):
-  """Problem spec for CoNLL2002 Dutch named entity task."""
-  TRAIN_FILES = ["ned.train"]
-  EVAL_FILES = ["ned.testa", "ned.testb"]
-  def source_data_files(self, dataset_split):
-    is_training = dataset_split == problem.DatasetSplit.TRAIN
-    return self.TRAIN_FILES if is_training else self.EVAL_FILES

From 7045453af5f319fdc09119c75f276b5a42365b4d Mon Sep 17 00:00:00 2001
From: Yuwen Yan <ybbaigo@gmail.com>
Date: Wed, 28 Nov 2018 15:18:38 -0800
Subject: [PATCH 1300/2720] internal merge of PR #1253

PiperOrigin-RevId: 223245334
---
 tensor2tensor/data_generators/conll_ner.py | 92 ++++++++++++++++++++++
 1 file changed, 92 insertions(+)
 create mode 100644 tensor2tensor/data_generators/conll_ner.py

diff --git a/tensor2tensor/data_generators/conll_ner.py b/tensor2tensor/data_generators/conll_ner.py
new file mode 100644
index 000000000..0ec57bdc1
--- /dev/null
+++ b/tensor2tensor/data_generators/conll_ner.py
@@ -0,0 +1,92 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Data generators for CoNLL dataset."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import zipfile
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+
+@registry.register_problem
+class Conll2002Ner(text_problems.Text2textTmpdir):
+  """Base class for CoNLL2002 problems."""
+
+  def source_data_files(self, dataset_split):
+    """Files to be passed to generate_samples."""
+    raise NotImplementedError()
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    del data_dir
+
+    url = "https://raw.githubusercontent.com/nltk/nltk_data/gh-pages/packages/corpora/conll2002.zip"  # pylint: disable=line-too-long
+    compressed_filename = os.path.basename(url)
+    compressed_filepath = os.path.join(tmp_dir, compressed_filename)
+    generator_utils.maybe_download(tmp_dir, compressed_filename, url)
+
+    compressed_dir = compressed_filepath.strip(".zip")
+
+    filenames = self.source_data_files(dataset_split)
+    for filename in filenames:
+      filepath = os.path.join(compressed_dir, filename)
+      if not tf.gfile.Exists(filepath):
+        with zipfile.ZipFile(compressed_filepath, "r") as corpus_zip:
+          corpus_zip.extractall(tmp_dir)
+      with tf.gfile.GFile(filepath, mode="r") as cur_file:
+        words, tags = [], []
+        for line in cur_file:
+          line_split = line.strip().split()
+          if not line_split:
+            yield {
+                "inputs": str.join(" ", words),
+                "targets": str.join(" ", tags)
+            }
+            words, tags = [], []
+            continue
+          words.append(line_split[0])
+          tags.append(line_split[2])
+        if words:
+          yield {"inputs": str.join(" ", words), "targets": str.join(" ", tags)}
+
+
+@registry.register_problem
+class Conll2002EsNer(Conll2002Ner):
+  """Problem spec for CoNLL2002 Spanish named entity task."""
+  TRAIN_FILES = ["esp.train"]
+  EVAL_FILES = ["esp.testa", "esp.testb"]
+
+  def source_data_files(self, dataset_split):
+    is_training = dataset_split == problem.DatasetSplit.TRAIN
+    return self.TRAIN_FILES if is_training else self.EVAL_FILES
+
+
+@registry.register_problem
+class Conll2002NlNer(Conll2002Ner):
+  """Problem spec for CoNLL2002 Dutch named entity task."""
+  TRAIN_FILES = ["ned.train"]
+  EVAL_FILES = ["ned.testa", "ned.testb"]
+
+  def source_data_files(self, dataset_split):
+    is_training = dataset_split == problem.DatasetSplit.TRAIN
+    return self.TRAIN_FILES if is_training else self.EVAL_FILES

From 1d75121b2e349198be99b963fbf5613c100673a3 Mon Sep 17 00:00:00 2001
From: Stefan Falk <43335432+stefan-falk@users.noreply.github.com>
Date: Thu, 29 Nov 2018 00:43:36 +0100
Subject: [PATCH 1301/2720] Word Error Rate for Speech Recognition (#1242)

---
 .../data_generators/speech_recognition.py     |  5 +-
 tensor2tensor/utils/metrics.py                | 64 +++++++++++++++++++
 tensor2tensor/utils/metrics_test.py           | 33 ++++++++++
 3 files changed, 101 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index f4dad3022..84e023a3a 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -139,4 +139,7 @@ def preprocess_example(self, example, mode, hparams):
 
   def eval_metrics(self):
     defaults = super(SpeechRecognitionProblem, self).eval_metrics()
-    return defaults + [metrics.Metrics.EDIT_DISTANCE]
+    return defaults + [
+      metrics.Metrics.EDIT_DISTANCE,
+      metrics.Metrics.WORD_ERROR_RATE
+    ]
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 3d21cf854..620c86e4d 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -47,6 +47,7 @@ class Metrics(object):
   ROUGE_2_F = "rouge_2_fscore"
   ROUGE_L_F = "rouge_L_fscore"
   EDIT_DISTANCE = "edit_distance"
+  WORD_ERROR_RATE = "word_error_rate"
   SET_PRECISION = "set_precision"
   SET_RECALL = "set_recall"
   SOFTMAX_CROSS_ENTROPY_ONE_HOT = "softmax_cross_entropy_one_hot"
@@ -680,6 +681,68 @@ def metric_means():
   return metric_accum, metric_means
 
 
+def word_error_rate(raw_predictions, labels, lookup=None,
+                    weights_fn=common_layers.weights_nonzero):
+  """
+  :param raw_predictions:
+  :param labels:
+  :param lookup:
+    A tf.constant mapping indices to output tokens.
+  :param weights_fn:
+  :return:
+    The word error rate.
+  """
+
+  def from_tokens(raw, lookup_):
+    gathered = tf.gather(lookup_, tf.cast(raw, tf.int32))
+    joined = tf.regex_replace(tf.reduce_join(gathered, axis=1), b'<EOS>.*', b'')
+    cleaned = tf.regex_replace(joined, b'_', b' ')
+    tokens = tf.string_split(cleaned, ' ')
+    return tokens
+
+  def from_characters(raw, lookup_):
+    """
+    Convert ascii+2 encoded codes to string-tokens.
+    """
+    corrected = tf.bitcast(
+      tf.clip_by_value(
+        tf.subtract(raw, 2), 0, 255
+      ), tf.uint8)
+
+    gathered = tf.gather(lookup_, tf.cast(corrected, tf.int32))[:, :, 0]
+    joined = tf.reduce_join(gathered, axis=1)
+    cleaned = tf.regex_replace(joined, b'\0', b'')
+    tokens = tf.string_split(cleaned, ' ')
+    return tokens
+
+  if lookup is None:
+    lookup = tf.constant([chr(i) for i in range(256)])
+    convert_fn = from_characters
+  else:
+    convert_fn = from_tokens
+
+  if weights_fn is not common_layers.weights_nonzero:
+    raise ValueError("Only weights_nonzero can be used for this metric.")
+
+  with tf.variable_scope("word_error_rate", values=[raw_predictions, labels]):
+
+    raw_predictions = tf.squeeze(
+      tf.argmax(raw_predictions, axis=-1), axis=(2, 3))
+    labels = tf.squeeze(labels, axis=(2, 3))
+
+    reference = convert_fn(labels, lookup)
+    predictions = convert_fn(raw_predictions, lookup)
+
+    distance = tf.reduce_sum(
+      tf.edit_distance(predictions, reference, normalize=False)
+    )
+    reference_length = tf.cast(
+      tf.size(reference.values, out_type=tf.int32), dtype=tf.float32
+    )
+
+    return distance / reference_length, reference_length
+
+
 # Metrics are functions that take predictions and labels and return
 # a tensor of metrics and a tensor of weights.
 # If the function has "features" as an argument, it will receive the whole
@@ -699,6 +762,7 @@ def metric_means():
     Metrics.ROUGE_2_F: rouge.rouge_2_fscore,
     Metrics.ROUGE_L_F: rouge.rouge_l_fscore,
     Metrics.EDIT_DISTANCE: sequence_edit_distance,
+    Metrics.WORD_ERROR_RATE: word_error_rate,
     Metrics.SOFTMAX_CROSS_ENTROPY_ONE_HOT: softmax_cross_entropy_one_hot,
     Metrics.SIGMOID_ACCURACY_ONE_HOT: sigmoid_accuracy_one_hot,
     Metrics.SIGMOID_RECALL_ONE_HOT: sigmoid_recall_one_hot,
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 4d25a47b4..3288a3e29 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -105,6 +105,39 @@ def testSequenceEditDistanceMetric(self):
     self.assertAlmostEqual(actual_scores, 3.0 / 13)
     self.assertEqual(actual_weight, 13)
 
+  def testWordErrorRateMetric(self):
+
+    ref = np.asarray([
+      # a b c
+      [97, 34, 98, 34, 99],
+      [97, 34, 98, 34, 99],
+      [97, 34, 98, 34, 99],
+      [97, 34, 98, 34, 99],
+    ])
+
+    hyp = np.asarray([
+      [97, 34, 98, 34, 99],  # a b c
+      [97, 34, 98, 0, 0],  # a b
+      [97, 34, 98, 34, 100],  # a b d
+      [0, 0, 0, 0, 0]  # empty
+    ])
+
+    labels = np.reshape(ref, ref.shape + (1, 1))
+    predictions = np.zeros((len(ref), np.max([len(s) for s in hyp]), 1, 1, 256))
+
+    for i, sample in enumerate(hyp):
+      for j, idx in enumerate(sample):
+        predictions[i, j, 0, 0, idx] = 1
+
+    with self.test_session() as session:
+      actual_wer, actual_ref_len = session.run(
+        metrics.word_error_rate(predictions, labels)
+      )
+
+    expected_wer = 0.417
+    places = 3
+    self.assertAlmostEqual(round(actual_wer, places), expected_wer, places)
+
   def testNegativeLogPerplexity(self):
     predictions = np.random.randint(4, size=(12, 12, 12, 1))
     targets = np.random.randint(4, size=(12, 12, 12, 1))

From f23e4cd17252de83c9b239a4ee0668e489dd76bd Mon Sep 17 00:00:00 2001
From: Stefan Falk <43335432+stefan-falk@users.noreply.github.com>
Date: Wed, 28 Nov 2018 16:00:05 -0800
Subject: [PATCH 1302/2720] internal merge of PR #1242

PiperOrigin-RevId: 223252032
---
 .../data_generators/speech_recognition.py     |  4 +-
 tensor2tensor/utils/metrics.py                | 46 +++++++++----------
 tensor2tensor/utils/metrics_test.py           | 23 +++++-----
 3 files changed, 35 insertions(+), 38 deletions(-)

diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 84e023a3a..04891d50e 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -140,6 +140,6 @@ def preprocess_example(self, example, mode, hparams):
   def eval_metrics(self):
     defaults = super(SpeechRecognitionProblem, self).eval_metrics()
     return defaults + [
-      metrics.Metrics.EDIT_DISTANCE,
-      metrics.Metrics.WORD_ERROR_RATE
+        metrics.Metrics.EDIT_DISTANCE,
+        metrics.Metrics.WORD_ERROR_RATE
     ]
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 620c86e4d..d4ff79e08 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -681,38 +681,38 @@ def metric_means():
   return metric_accum, metric_means
 
 
-def word_error_rate(raw_predictions, labels, lookup=None,
+def word_error_rate(raw_predictions,
+                    labels,
+                    lookup=None,
                     weights_fn=common_layers.weights_nonzero):
-  """
-  :param raw_predictions:
-  :param labels:
-  :param lookup:
-    A tf.constant mapping indices to output tokens.
-  :param weights_fn:
-  :return:
+  """Calculate word error rate.
+
+  Args:
+    raw_predictions: The raw predictions.
+    labels: The actual labels.
+    lookup: A tf.constant mapping indices to output tokens.
+    weights_fn: Weighting function.
+
+  Returns:
     The word error rate.
   """
 
   def from_tokens(raw, lookup_):
     gathered = tf.gather(lookup_, tf.cast(raw, tf.int32))
-    joined = tf.regex_replace(tf.reduce_join(gathered, axis=1), b'<EOS>.*', b'')
-    cleaned = tf.regex_replace(joined, b'_', b' ')
-    tokens = tf.string_split(cleaned, ' ')
+    joined = tf.regex_replace(tf.reduce_join(gathered, axis=1), b"<EOS>.*", b"")
+    cleaned = tf.regex_replace(joined, b"_", b" ")
+    tokens = tf.string_split(cleaned, " ")
     return tokens
 
   def from_characters(raw, lookup_):
-    """
-    Convert ascii+2 encoded codes to string-tokens.
-    """
+    """Convert ascii+2 encoded codes to string-tokens."""
     corrected = tf.bitcast(
-      tf.clip_by_value(
-        tf.subtract(raw, 2), 0, 255
-      ), tf.uint8)
+        tf.clip_by_value(tf.subtract(raw, 2), 0, 255), tf.uint8)
 
     gathered = tf.gather(lookup_, tf.cast(corrected, tf.int32))[:, :, 0]
     joined = tf.reduce_join(gathered, axis=1)
-    cleaned = tf.regex_replace(joined, b'\0', b'')
-    tokens = tf.string_split(cleaned, ' ')
+    cleaned = tf.regex_replace(joined, b"\0", b"")
+    tokens = tf.string_split(cleaned, " ")
     return tokens
 
   if lookup is None:
@@ -727,18 +727,16 @@ def from_characters(raw, lookup_):
   with tf.variable_scope("word_error_rate", values=[raw_predictions, labels]):
 
     raw_predictions = tf.squeeze(
-      tf.argmax(raw_predictions, axis=-1), axis=(2, 3))
+        tf.argmax(raw_predictions, axis=-1), axis=(2, 3))
     labels = tf.squeeze(labels, axis=(2, 3))
 
     reference = convert_fn(labels, lookup)
     predictions = convert_fn(raw_predictions, lookup)
 
     distance = tf.reduce_sum(
-      tf.edit_distance(predictions, reference, normalize=False)
-    )
+        tf.edit_distance(predictions, reference, normalize=False))
     reference_length = tf.cast(
-      tf.size(reference.values, out_type=tf.int32), dtype=tf.float32
-    )
+        tf.size(reference.values, out_type=tf.int32), dtype=tf.float32)
 
     return distance / reference_length, reference_length
 
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 3288a3e29..b6228483c 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -108,18 +108,18 @@ def testSequenceEditDistanceMetric(self):
   def testWordErrorRateMetric(self):
 
     ref = np.asarray([
-      # a b c
-      [97, 34, 98, 34, 99],
-      [97, 34, 98, 34, 99],
-      [97, 34, 98, 34, 99],
-      [97, 34, 98, 34, 99],
+        # a b c
+        [97, 34, 98, 34, 99],
+        [97, 34, 98, 34, 99],
+        [97, 34, 98, 34, 99],
+        [97, 34, 98, 34, 99],
     ])
 
     hyp = np.asarray([
-      [97, 34, 98, 34, 99],  # a b c
-      [97, 34, 98, 0, 0],  # a b
-      [97, 34, 98, 34, 100],  # a b d
-      [0, 0, 0, 0, 0]  # empty
+        [97, 34, 98, 34, 99],  # a b c
+        [97, 34, 98, 0, 0],  # a b
+        [97, 34, 98, 34, 100],  # a b d
+        [0, 0, 0, 0, 0]  # empty
     ])
 
     labels = np.reshape(ref, ref.shape + (1, 1))
@@ -130,9 +130,8 @@ def testWordErrorRateMetric(self):
         predictions[i, j, 0, 0, idx] = 1
 
     with self.test_session() as session:
-      actual_wer, actual_ref_len = session.run(
-        metrics.word_error_rate(predictions, labels)
-      )
+      actual_wer, unused_actual_ref_len = session.run(
+          metrics.word_error_rate(predictions, labels))
 
     expected_wer = 0.417
     places = 3

From e5f24a5a861f119ccbf0fa07ff4ca3e68e3d579d Mon Sep 17 00:00:00 2001
From: blazejosinski <804945+blazejosinski@users.noreply.github.com>
Date: Thu, 29 Nov 2018 01:56:20 +0100
Subject: [PATCH 1303/2720] Differentiate summaries for train and eval. (#1256)

---
 tensor2tensor/models/research/rl.py | 24 +-----------------------
 tensor2tensor/rl/ppo_learner.py     | 17 +++++++++--------
 2 files changed, 10 insertions(+), 31 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 278d183e7..7bc79c145 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -178,28 +178,6 @@ def ppo_pong_ae_base():
   return hparams
 
 
-@registry.register_hparams
-def pong_model_free():
-  """TODO(piotrmilos): Document this."""
-  hparams = mfrl_base()
-  hparams.batch_size = 2
-  hparams.ppo_eval_every_epochs = 2
-  hparams.ppo_epochs_num = 4
-  hparams.add_hparam("ppo_optimization_epochs", 3)
-  hparams.add_hparam("ppo_epoch_length", 30)
-  hparams.add_hparam("ppo_learning_rate", 8e-05)
-  hparams.add_hparam("ppo_optimizer", "Adam")
-  hparams.add_hparam("ppo_optimization_batch_size", 4)
-  hparams.add_hparam("ppo_save_models_every_epochs", 1000000)
-  env = gym_env.T2TGymEnv("PongNoFrameskip-v4", batch_size=2)
-  env.start_new_epoch(0)
-  hparams.add_hparam("env_fn", make_real_env_fn(env))
-  eval_env = gym_env.T2TGymEnv("PongNoFrameskip-v4", batch_size=2)
-  eval_env.start_new_epoch(0)
-  hparams.add_hparam("eval_env_fn", make_real_env_fn(eval_env))
-  return hparams
-
-
 @registry.register_hparams
 def dqn_atari_base():
   # These params are based on agents/dqn/configs/dqn.gin
@@ -242,7 +220,7 @@ def dqn_original_params():
 @registry.register_hparams
 def mfrl_original():
   return tf.contrib.training.HParams(
-      game="",
+      game="pong",
       base_algo="ppo",
       base_algo_params="ppo_original_params",
       batch_size=16,
diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 33791e320..2a6fa5603 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -460,11 +460,12 @@ def stop_condition(i, _, resets):
         new_memory.append(mem)
       memory = new_memory
 
-    mean_score_summary = tf.cond(
-        tf.greater(scores_num, 0),
-        lambda: tf.summary.scalar("mean_score_this_iter", mean_score), str)
-    summaries = tf.summary.merge([
-        mean_score_summary,
-        tf.summary.scalar("episodes_finished_this_iter", scores_num)
-    ])
-    return memory, summaries, initialization_lambda
+    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
+      mean_score_summary = tf.cond(
+          tf.greater(scores_num, 0),
+          lambda: tf.summary.scalar("mean_score_this_iter", mean_score), str)
+      summaries = tf.summary.merge([
+          mean_score_summary,
+          tf.summary.scalar("episodes_finished_this_iter", scores_num)
+      ])
+      return memory, summaries, initialization_lambda

From e397d6c80dd7e9de2b488c2b32a78af46f6f9cea Mon Sep 17 00:00:00 2001
From: blazejosinski <804945+blazejosinski@users.noreply.github.com>
Date: Wed, 28 Nov 2018 16:57:04 -0800
Subject: [PATCH 1304/2720] internal merge of PR #1256

PiperOrigin-RevId: 223261190
---
 tensor2tensor/models/research/rl.py         | 1 -
 tensor2tensor/rl/trainer_model_free_test.py | 2 +-
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 7bc79c145..1d6b5f52f 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -20,7 +20,6 @@
 import operator
 import gym
 
-from tensor2tensor.data_generators import gym_env
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
diff --git a/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index bb1116016..be3300312 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -29,7 +29,7 @@
 class TrainTest(tf.test.TestCase):
 
   def test_train_pong(self):
-    hparams = registry.hparams("pong_model_free")
+    hparams = registry.hparams("mfrl_original")
     hparams.batch_size = 2
     hparams.ppo_epochs_num = 2
     hparams.ppo_epoch_length = 3

From 5226aa919ae54455f658ab77027477f200fcf3d1 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 29 Nov 2018 11:38:25 -0800
Subject: [PATCH 1305/2720] adding 256*256 imagenet class to t2t.problems.

PiperOrigin-RevId: 223382108
---
 tensor2tensor/data_generators/imagenet.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index dfd1984b6..610cab819 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -163,6 +163,15 @@ def rescale_size(self):
     return [224, 224]
 
 
+@registry.register_problem
+class ImageImagenet256(ImageImagenetRescaled):
+  """Imagenet rescaled to 256x256."""
+
+  @property
+  def rescale_size(self):
+    return [256, 256]
+
+
 @registry.register_problem
 class ImageImagenet32(ImageImagenetRescaled):
   """Imagenet rescaled to 32x32."""
@@ -190,7 +199,7 @@ def preprocess_example(self, example, mode, _):
 
 @registry.register_problem
 class ImageImagenet32Gen(ImageImagenet):
-  """Imagenet 32 from the pixen cnn paper"""
+  """Imagenet 32 from the pixen cnn paper."""
 
   @property
   def train_shards(self):
@@ -224,7 +233,7 @@ def preprocess_example(self, example, mode, unused_hparams):
 
 @registry.register_problem
 class ImageImagenet64Gen(ImageImagenet):
-  """Imagenet 64 from the pixen cnn paper"""
+  """Imagenet 64 from the pixen cnn paper."""
 
   @property
   def train_shards(self):
@@ -304,7 +313,7 @@ def preprocess_example(self, example, mode, hparams):
 
 @registry.register_problem
 class ImageImagenet32Small(ImageImagenet):
-  """Imagenet small from the pixel cnn paper"""
+  """Imagenet small from the pixel cnn paper."""
 
   @property
   def is_small(self):

From 178fb977584d13564a430b8b6f6d00ac6b40ac5a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 29 Nov 2018 14:19:23 -0800
Subject: [PATCH 1306/2720] Add En-Ro translation, reserve ids for extra tasks
 in TPU adafactor multi-Transformer.

PiperOrigin-RevId: 223410257
---
 .../data_generators/translate_enro.py         | 80 +++++++++++++++++++
 .../data_generators/wiki_multi_problems.py    |  4 +
 tensor2tensor/models/transformer.py           |  1 +
 3 files changed, 85 insertions(+)
 create mode 100644 tensor2tensor/data_generators/translate_enro.py

diff --git a/tensor2tensor/data_generators/translate_enro.py b/tensor2tensor/data_generators/translate_enro.py
new file mode 100644
index 000000000..353c5d8c4
--- /dev/null
+++ b/tensor2tensor/data_generators/translate_enro.py
@@ -0,0 +1,80 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Data generators for translation data-sets."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators import translate
+from tensor2tensor.data_generators import wiki_lm
+from tensor2tensor.utils import registry
+
+
+_ENRO_TRAIN_DATASETS = [
+    [
+        "http://www.statmt.org/europarl/v7/ro-en.tgz",
+        ("europarl-v7.ro-en.en", "europarl-v7.ro-en.ro")
+    ],
+]
+_ENRO_TEST_DATASETS = [
+    [
+        ("http://data.statmt.org/wmt16/translation-task/"
+         "dev-romanian-updated.tgz"),
+        ("dev/newsdev2016-roen-ref.en.sgm", "dev/newsdev2016-roen-src.ro.sgm")
+    ],
+]
+
+
+@registry.register_problem
+class TranslateEnroWmt8k(translate.TranslateProblem):
+  """Problem spec for WMT En-Ro translation."""
+
+  @property
+  def approx_vocab_size(self):
+    return 2**13  # 8192
+
+  def source_data_files(self, dataset_split):
+    train = dataset_split == problem.DatasetSplit.TRAIN
+    return _ENRO_TRAIN_DATASETS if train else _ENRO_TEST_DATASETS
+
+
+@registry.register_problem
+class TranslateEnroWmt32k(TranslateEnroWmt8k):
+
+  @property
+  def approx_vocab_size(self):
+    return 2**15  # 32768
+
+
+@registry.register_problem
+class TranslateEnroWmtCharacters(TranslateEnroWmt8k):
+  """Problem spec for WMT En-Ro translation."""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+
+@registry.register_problem
+class TranslateEnroWmtMulti64k(TranslateEnroWmt8k):
+  """Translation with muli-lingual vocabulary."""
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
diff --git a/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index 76f218770..5c8515fae 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -25,6 +25,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.data_generators import translate_ende
 from tensor2tensor.data_generators import translate_enfr
+from tensor2tensor.data_generators import translate_enro
 from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
 
@@ -132,10 +133,13 @@ def __init__(self, was_reversed=False, was_copy=False):
     self.task_list.append(wiki_lm.LanguagemodelDeEnFrRoWiki64k())
     self.task_list.append(translate_ende.TranslateEndeWmtMulti64k())
     self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k())
+    self.task_list.append(translate_enro.TranslateEnroWmtMulti64k())
     self.task_list.append(translate_ende.TranslateEndeWmtMulti64k(
         was_reversed=True))
     self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k(
         was_reversed=True))
+    self.task_list.append(translate_enro.TranslateEnroWmtMulti64k(
+        was_reversed=True))
 
   @property
   def vocab_type(self):
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 7ff34d72b..2749ac1a3 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1588,6 +1588,7 @@ def transformer_tall_pretrain_lm_tpu_adafactor():
   hparams.max_length = 1024
   # For multi-problem on TPU we need it in absolute examples.
   hparams.batch_size = 8
+  hparams.multiproblem_vocab_size = 2**16
   return hparams
 
 

From 4dc7b9a8581464cd842bd36c9b791992189364d7 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Thu, 29 Nov 2018 17:00:58 -0800
Subject: [PATCH 1307/2720] Internal change

PiperOrigin-RevId: 223437662
---
 tensor2tensor/data_generators/video_utils_test.py | 9 ++-------
 tensor2tensor/utils/video_metrics.py              | 5 ++++-
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index 12d971882..39d626e46 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -54,7 +54,6 @@ def testDecodeInMemoryTrue(self):
         hparams=decode_hparams, decode_hparams=decode_hparams,
         predictions=predictions)
     metrics = video_utils.summarize_video_metrics(decode_hooks)
-    self.assertEqual(len(metrics), 40)
 
   def testConvertPredictionsToVideoSummaries(self):
     # Initialize predictions.
@@ -64,7 +63,7 @@ def testConvertPredictionsToVideoSummaries(self):
     targets = rng.randint(0, 255, (5, 32, 32, 3))
 
     # batch it up.
-    prediction = [{"outputs": outputs, "inputs": inputs, "targets": targets}]*50
+    prediction = [{"outputs": outputs, "inputs": inputs, "targets": targets}]*5
     predictions = [prediction]
     decode_hparams = decoding.decode_hparams()
 
@@ -73,11 +72,7 @@ def testConvertPredictionsToVideoSummaries(self):
         hparams=decode_hparams, decode_hparams=decode_hparams,
         predictions=predictions)
     summaries = video_utils.display_video_hooks(decode_hooks)
-    # for {psnr_max, psnr_min, ssim_max, ssim_min}
-    # 10 output vids + 10 frame-by-frame.
-    # for {random}
-    # 10 input vids + 10 output vids + 10 frame-by-frame.
-    self.assertEqual(len(summaries), 110)
+
     for summary in summaries:
       self.assertTrue(isinstance(summary, tf.Summary.Value))
 
diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index 371baf6fd..a5119a1af 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -225,6 +225,9 @@ def compute_video_metrics_from_predictions(predictions):
     statistics: dict of Tensors, key being the metric with each Tensor
                 having the shape (num_samples, num_frames).
   """
+  all_results = {}
+
+
   ssim_all_decodes, psnr_all_decodes = [], []
   for single_decode in predictions:
     args = get_zipped_dataset_from_predictions(single_decode)
@@ -233,7 +236,7 @@ def compute_video_metrics_from_predictions(predictions):
     ssim_all_decodes.append(ssim_single)
   psnr_all_decodes = np.array(psnr_all_decodes)
   ssim_all_decodes = np.array(ssim_all_decodes)
-  all_results = {"PSNR": psnr_all_decodes, "SSIM": ssim_all_decodes}
+  all_results.update({"PSNR": psnr_all_decodes, "SSIM": ssim_all_decodes})
   return compute_all_metrics_statistics(all_results)
 
 

From 1feb765d9424cca8e254cbf42b7ac51fa2ce9fdd Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 3 Dec 2018 11:02:37 -0800
Subject: [PATCH 1308/2720] Delete unused only_last from dataset. The
 functionality was buggy to

PiperOrigin-RevId: 223825208
---
 tensor2tensor/data_generators/multi_problem.py |  6 ++----
 tensor2tensor/data_generators/problem.py       | 16 +++-------------
 2 files changed, 5 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 0c400d65f..b31000ed2 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -159,8 +159,7 @@ def dataset(self,
               partition_id=0,
               num_partitions=1,
               shuffle_buffer_size=1024,
-              max_records=-1,
-              only_last=False):
+              max_records=-1):
 
     # A list of datasets corresponding to the tasks in the task_list object
     # that need to be mixed.
@@ -188,8 +187,7 @@ def dataset(self,
                                   partition_id=partition_id,
                                   num_partitions=num_partitions,
                                   shuffle_buffer_size=shuffle_buffer_size,
-                                  max_records=max_records,
-                                  only_last=only_last)
+                                  max_records=max_records)
 
       if is_training:
         task_dataset = task_dataset.repeat()
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 58c5bc56f..0eac893c3 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -571,8 +571,7 @@ def dataset(self,
               partition_id=0,
               num_partitions=1,
               shuffle_buffer_size=1024,
-              max_records=-1,
-              only_last=False):
+              max_records=-1):
     """Build a Dataset for this problem.
 
     Args:
@@ -596,7 +595,6 @@ def dataset(self,
       shuffle_buffer_size: if shuffle_files is True, this is the buffer size
         used to shuffle records.
       max_records: int, number of records to truncate to.
-      only_last: bool, whether we should include only files from last epoch.
 
     Returns:
       Dataset containing dict<feature name, Tensor>.
@@ -621,17 +619,9 @@ def dataset(self,
     _ = self.get_hparams(hparams)
 
     data_filepattern = self.filepattern(data_dir, dataset_split, shard=shard)
-    if only_last:
-      imprv_data_filepattern = data_filepattern + r"10.[\d+]"
-    else:
-      imprv_data_filepattern = data_filepattern
     tf.logging.info("Reading data files from %s", data_filepattern)
-    try:
-      data_files = sorted(tf.contrib.slim.parallel_reader.get_data_files(
-          imprv_data_filepattern))
-    except ValueError:
-      data_files = sorted(tf.contrib.slim.parallel_reader.get_data_files(
-          data_filepattern))
+    data_files = sorted(tf.contrib.slim.parallel_reader.get_data_files(
+        data_filepattern))
 
     # Functions used in dataset transforms below. `filenames` can be either a
     # `tf.string` tensor or `tf.data.Dataset` containing one or more filenames.

From e7ca1c4be892b5eb8b8485fd9b98055aafe052ac Mon Sep 17 00:00:00 2001
From: David Dohan <ddohan@google.com>
Date: Mon, 3 Dec 2018 11:33:16 -0800
Subject: [PATCH 1309/2720] Let weights_multi_problem_all support tensor taskid

The taskid broadcasts, so the weights actually work well for [batch x 1] task ids.

PiperOrigin-RevId: 223831382
---
 tensor2tensor/layers/common_layers.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 73ca1917a..31634bd62 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1613,9 +1613,6 @@ def weights_multi_problem(labels, taskid=-1):
 def weights_multi_problem_all(labels, taskid=-1):
   """Assign weight 1.0 to only examples from the given task."""
   weights = tf.to_float(tf.not_equal(labels, 0))
-  if taskid < 0:
-    raise ValueError("Task ID must be non-negative.")
-
   past_taskid = tf.cumsum(tf.to_float(tf.equal(labels, taskid)), axis=1)
   # Additionally zero out the task id location
   past_taskid *= tf.to_float(tf.not_equal(labels, taskid))

From 424a68ac5ce3c0f55164ed6183600d2472fd0950 Mon Sep 17 00:00:00 2001
From: Art Wangperawong <artitw@gmail.com>
Date: Mon, 3 Dec 2018 17:28:17 -0500
Subject: [PATCH 1310/2720] Fix universal transformer decoding (#1257)

* fix bAbi data generator and readme

* Fix bAbi hparams deletion

* Fix bAbi hparams delete unecessary keys

* Fix bAbi hparams clean keys

* bAbi hparams delete keys

* fix readme

* fix universal transformer decoding

* fix merge conflict
---
 README.md                                              | 1 -
 tensor2tensor/data_generators/babi_qa.py               | 5 ++---
 tensor2tensor/models/research/universal_transformer.py | 4 ++--
 3 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 205e59283..1a031086b 100644
--- a/README.md
+++ b/README.md
@@ -90,7 +90,6 @@ You can choose the bAbi task from the range [1,20] and the subset from 1k or
 10k. To combine test data from all tasks into a single test set, use
 `--problem=babi_qa_concat_all_tasks_10k`
 
-
 ### Image Classification
 
 For image classification, we have a number of standard data-sets:
diff --git a/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index a11eaddca..882d17778 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -109,9 +109,9 @@ def _prepare_babi_data(tmp_dir, data_dir):
     tf.gfile.MakeDirs(data_dir)
 
   file_path = os.path.join(tmp_dir, _TAR)
-  headers = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36"}  # pylint: disable=line-too-long
+  headers = {'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36'}
   resp = requests.get(_URL, headers=headers)
-  with open(file_path, "wb") as f:
+  with open(file_path, 'wb') as f:
     f.write(resp.content)
 
   tar = tarfile.open(file_path)
@@ -459,7 +459,6 @@ def hparams(self, defaults, unused_model_hparams):
     if "context" in p.vocab_size:
       del p.vocab_size["context"]
 
-
 def _problems_to_register():
   """Problems for which we want to create datasets.
 
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 148426beb..cf23f3830 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -243,7 +243,7 @@ def _greedy_infer(self, features, decode_length, use_tpu=False):
     return (self._slow_greedy_infer_tpu(features, decode_length) if use_tpu else
             self._slow_greedy_infer(features, decode_length))
 
-  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
+  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha, use_tpu=False):
     """Beam search decoding.
 
     Args:
@@ -266,7 +266,7 @@ def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
     # Caching is not ebabled in Universal Transformer
     # TODO(dehghani): Support fast decoding for Universal Transformer
     return self._beam_decode_slow(features, decode_length, beam_size,
-                                  top_beams, alpha)
+                                  top_beams, alpha, use_tpu)
 
 
 @registry.register_model

From fd5a87bf04a767c8f3cdafe58546b2ed5f0b17b4 Mon Sep 17 00:00:00 2001
From: Art Wangperawong <artitw@gmail.com>
Date: Mon, 3 Dec 2018 14:53:19 -0800
Subject: [PATCH 1311/2720] internal merge of PR #1257

PiperOrigin-RevId: 223867901
---
 docs/walkthrough.md                                    | 1 -
 tensor2tensor/data_generators/babi_qa.py               | 5 +++--
 tensor2tensor/models/research/universal_transformer.py | 4 +++-
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 205e59283..1a031086b 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -90,7 +90,6 @@ You can choose the bAbi task from the range [1,20] and the subset from 1k or
 10k. To combine test data from all tasks into a single test set, use
 `--problem=babi_qa_concat_all_tasks_10k`
 
-
 ### Image Classification
 
 For image classification, we have a number of standard data-sets:
diff --git a/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index 882d17778..a11eaddca 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -109,9 +109,9 @@ def _prepare_babi_data(tmp_dir, data_dir):
     tf.gfile.MakeDirs(data_dir)
 
   file_path = os.path.join(tmp_dir, _TAR)
-  headers = {'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36'}
+  headers = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36"}  # pylint: disable=line-too-long
   resp = requests.get(_URL, headers=headers)
-  with open(file_path, 'wb') as f:
+  with open(file_path, "wb") as f:
     f.write(resp.content)
 
   tar = tarfile.open(file_path)
@@ -459,6 +459,7 @@ def hparams(self, defaults, unused_model_hparams):
     if "context" in p.vocab_size:
       del p.vocab_size["context"]
 
+
 def _problems_to_register():
   """Problems for which we want to create datasets.
 
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index cf23f3830..788037d9f 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -243,7 +243,8 @@ def _greedy_infer(self, features, decode_length, use_tpu=False):
     return (self._slow_greedy_infer_tpu(features, decode_length) if use_tpu else
             self._slow_greedy_infer(features, decode_length))
 
-  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha, use_tpu=False):
+  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha,
+                   use_tpu=False):
     """Beam search decoding.
 
     Args:
@@ -253,6 +254,7 @@ def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha, use
       top_beams: an integer. How many of the beams to return.
       alpha: Float that controls the length penalty. larger the alpha, stronger
         the preference for longer translations.
+      use_tpu: Whether we should use TPU or not.
 
     Returns:
       A dict of decoding results {

From f58a4854e3e03b88748f2d8ed052d4cc7cc22e76 Mon Sep 17 00:00:00 2001
From: cclauss <cclauss@bluewin.ch>
Date: Tue, 4 Dec 2018 00:32:18 +0100
Subject: [PATCH 1312/2720] Undefined name: xrange() was removed in Python 3
 (#1258)

---
 tensor2tensor/models/research/vqa_self_attention.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
index b7c07ad0e..482e18c9e 100644
--- a/tensor2tensor/models/research/vqa_self_attention.py
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+from six.moves import xrange
+
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers

From 9c61f72d1f80100ad02ee67d71f51e566547a094 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 3 Dec 2018 16:16:58 -0800
Subject: [PATCH 1313/2720] Move nfg to models/video for external
 collaboration.

PiperOrigin-RevId: 223882807
---
 tensor2tensor/models/__init__.py              |   1 +
 tensor2tensor/models/video/next_frame_glow.py | 572 ++++++++++++++++++
 tensor2tensor/models/video/nfg_conv3d_test.py |  49 ++
 .../models/video/nfg_conv_lstm_test.py        |  48 ++
 tensor2tensor/models/video/nfg_conv_test.py   |  44 ++
 tensor2tensor/models/video/nfg_test_utils.py  | 184 ++++++
 tensor2tensor/models/video/nfg_uncond_test.py |  46 ++
 7 files changed, 944 insertions(+)
 create mode 100644 tensor2tensor/models/video/next_frame_glow.py
 create mode 100644 tensor2tensor/models/video/nfg_conv3d_test.py
 create mode 100644 tensor2tensor/models/video/nfg_conv_lstm_test.py
 create mode 100644 tensor2tensor/models/video/nfg_conv_test.py
 create mode 100644 tensor2tensor/models/video/nfg_test_utils.py
 create mode 100644 tensor2tensor/models/video/nfg_uncond_test.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 05cafbde7..3afc44e2d 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -67,6 +67,7 @@
 from tensor2tensor.models.video import basic_stochastic
 from tensor2tensor.models.video import emily
 from tensor2tensor.models.video import epva
+from tensor2tensor.models.video import next_frame_glow
 from tensor2tensor.models.video import savp
 from tensor2tensor.models.video import sv2p
 from tensor2tensor.models.video import svg_lp
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
new file mode 100644
index 000000000..70583154d
--- /dev/null
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -0,0 +1,572 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Experimental testbed for a conditional glow model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
+from tensor2tensor.layers import modalities
+from tensor2tensor.models.research import glow
+from tensor2tensor.models.research import glow_ops
+from tensor2tensor.utils import registry
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+
+arg_scope = tf.contrib.framework.arg_scope
+
+
+@registry.register_hparams
+def next_frame_glow_hparams():
+  """Hparams for next_frame_glow."""
+  hparams = glow.glow_hparams()
+  # Possible modes are conditional and unconditional
+  hparams.add_hparam("gen_mode", "conditional")
+  hparams.add_hparam("learn_top_scale", False)
+  hparams.add_hparam("condition_all_levels", True)
+  # For each video, substitutes "num_input_frames + num_output_frames" with a
+  # randomly sampled patch of length "num_train_frames" during training.
+  # -1 indicates that the entire video is used for training.
+  hparams.add_hparam("num_train_frames", -1)
+  # The following are hparams that model the latent transitions.
+  # Encoder that maps the latents to a Gaussian distribution.
+  # This function is used to model the prior over z_{t}. Can be,
+  # Pointwise -> point-wise multiplication of z_{t-1}.
+  # conv_net -> one-layer convolution over z_{t-1} .. z_{t - num_cond_latents}
+  hparams.add_hparam("latent_dist_encoder", "conv_net")
+  # Number of latents used in the encoder above.
+  hparams.add_hparam("num_cond_latents", 1)
+  hparams.add_hparam("latent_architecture", "glow_resnet")
+  hparams.add_hparam("latent_apply_dilations", False)
+  # Use latent skip connections
+  hparams.add_hparam("model_input", False)
+  hparams.add_hparam("cond_first_frame", False)
+  hparams.add_hparam("latent_skip", True)
+  hparams.add_hparam("latent_encoder_depth", 2)
+  hparams.add_hparam("latent_encoder_width", 512)
+  hparams.add_hparam("latent_pre_output_channels", 512)
+  # Pretrains the glow encoder for "pretrain_steps" number of steps.
+  # By default, don't pretrain and learn end-to-end
+  hparams.add_hparam("pretrain_steps", -1)
+  hparams.modality = {
+      "inputs": modalities.VideoModalityL1Raw,
+      "targets": modalities.VideoModalityL1Raw,
+  }
+  hparams.init_batch_size = 256
+  hparams.batch_size = 32
+  # Possible options: are prev_frame, single_conv and normal
+  hparams.top_prior = "single_conv"
+  return hparams
+
+
+@registry.register_hparams
+def frame_glow_hparams():
+  """Unconditional generation on video-frames."""
+  hparams = next_frame_glow_hparams()
+  hparams.gen_mode = "unconditional"
+  hparams.num_train_frames = 1
+  return hparams
+
+
+def get_cond_latents(all_latents=None, hparams=None):
+  """Get z^{cond}_{t} given z^{1..t-1}.
+
+  Args:
+    all_latents: list of list of tensors,
+                 outer-size equals no.of time_steps-1
+                 inner-size equals hparams.n_levels.
+    hparams: See next_frame_glow_hparams.
+  Returns:
+    cond_latents: conditional latents at time-step t.
+  """
+  cond_latents = None
+  if hparams.gen_mode == "conditional":
+    if hparams.latent_dist_encoder in ["conv_net", "conv3d_net"]:
+      num_cond_latents = (hparams.num_cond_latents +
+                          int(hparams.cond_first_frame))
+      if len(all_latents) >= num_cond_latents:
+        cond_latents = all_latents[-hparams.num_cond_latents:]
+        if hparams.cond_first_frame:
+          cond_latents = [all_latents[0]] + cond_latents
+    elif hparams.latent_dist_encoder in ["pointwise", "conv_lstm"]:
+      if all_latents:
+        cond_latents = all_latents[-1]
+
+  if hparams.gen_mode == "conditional":
+    global_step = tf.train.get_or_create_global_step()
+    condition = tf.greater(global_step, hparams.pretrain_steps)
+  else:
+    condition = tf.constant(False, dtype=tf.bool)
+  return condition, cond_latents
+
+
+@registry.register_model
+class NextFrameGlow(glow.Glow):
+  """Extend Glow for video."""
+
+  def init_preprocess_single(self, features):
+    for label in ["inputs", "targets"]:
+      features[label] = common_layers.convert_rgb_to_real(features[label])
+    return features
+
+  def init_preprocess(self, features):
+    """Preprocessing as per the input modality.
+
+    Equivalent to calling self.bottom(features).
+
+    Args:
+      features: dict of strings to tensors.
+    Returns:
+      features: dict of strings to tensors.
+    """
+    return features.map(self.init_preprocess_single)
+
+  def preprocess(self, x):
+    """Converts x from [0, 1] to [-0.5, 0.5].
+
+    All inputs are already normalized to be in the range [0, 1] through the
+    VideoModalityL1Raw modality.
+
+    Args:
+      x: 4-D Tensor.
+
+    Returns:
+      x: Scaled such that x lies in-between -0.5 and 0.5
+    """
+    return x - 0.5
+
+  def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
+    del args, kwargs
+
+    # Make a copy of features that can be used in the call to self
+    # that builds the graph.
+    new_features = {}
+    new_features["inputs"] = features["inputs"]
+    new_features["targets"] = features["infer_targets"]
+    _, _ = self(new_features)  # pylint: disable=not-callable
+
+    if self.hparams.gen_mode == "unconditional":
+      num_target_frames = 1
+    else:
+      num_target_frames = self.hparams.video_num_target_frames
+
+    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
+    var_scope = tf.variable_scope("next_frame_glow/body", reuse=True)
+    all_frames = []
+
+    # If eps=None, images are sampled from the prior.
+    with arg_scope(ops, init=False), var_scope:
+      for target_frame in range(1, num_target_frames + 1):
+
+        # subscript -> timestep, superscript -> level.
+        # self.z_sample equals z^0_{t} (top-level latent)
+        # (X_{t}, z^{1..l}_{t}) = Glow(z^0_{t}, z^{1..l}_{t-1})
+        # Get current set of cond_latents.
+        cond_level, cond_level_latents = get_cond_latents(
+            self.all_level_latents, self.hparams)
+
+        glow_vals = glow_ops.encoder_decoder(
+            "codec", self.z_sample, self.hparams, eps=None, reverse=True,
+            cond_latents=cond_level_latents, states=self.level_states,
+            condition=cond_level, temperature=self.temperature)
+        predicted_frame, _, curr_latents, self.level_states = glow_vals
+        all_frames.append(predicted_frame)
+        self.all_level_latents.append(curr_latents)
+
+        # Compute z^0_{t+1} = f(z^0_{t})
+        if target_frame < num_target_frames:
+          cond_top, cond_top_latents = get_cond_latents(
+              self.all_top_latents, self.hparams)
+          prior_dist = self.top_prior(
+              condition=cond_top, cond_latents=cond_top_latents)
+          self.z_sample = prior_dist.sample()
+          self.all_top_latents.append(self.z_sample)
+
+    all_frames = tf.stack(all_frames)
+    predicted_video = common_video.swap_time_and_batch_axes(all_frames)
+
+    # The video-decode API requires the predicted video to be the same shape
+    # as the target-video. Hence, for unconditional generation,
+    # tile across time to ensure same shape.
+    if self.hparams.gen_mode == "unconditional":
+      predicted_video = tf.tile(
+          predicted_video, [1, self.hparams.video_num_target_frames, 1, 1, 1])
+    predicted_video = self.scale(predicted_video)
+
+    # Output of a single decode / sample.
+    output_features = {}
+    output_features["targets"] = tf.zeros_like(predicted_video)
+    output_features["outputs"] = predicted_video
+    output_features["scores"] = tf.zeros_like(predicted_video)
+    return output_features
+
+  def get_squeeze_prior(self):
+    """Model the prior over z_{t} as a function of X_{t-1}.
+
+    Returns:
+      objective: float, log-likelihood.
+      dist: instance of tfp.distributions.Normal.
+
+    Raises:
+      ValueError: If input_height is not equal to input_width, not even
+                   or if the image width is smaller than the latent width.
+    """
+    _, prior_height, _, prior_channels = self.z_top_shape
+    _, input_height, input_width, _ = common_layers.shape_list(self.input_frame)
+
+    if input_height != input_width:
+      raise ValueError("input height should be equal to input width")
+    if input_height % 2 != 0:
+      raise ValueError("input height should be even")
+    if input_height < prior_height:
+      raise ValueError("input should be larger than the prior.")
+
+    # mean, log_std = NN(X_0)
+    # Reduce the spatial dimension by a factor of "squeeze_factor".
+    # and convolve with a stride of 2
+    squeeze_factor = input_height // (2 * prior_height)
+    x = glow_ops.squeeze(
+        "prior_squeeze", self.input_frame, factor=squeeze_factor, reverse=False)
+    mean_and_log_std = glow_ops.conv(
+        "prior_conv", x, 2*prior_channels, stride=[2, 2], apply_actnorm=False,
+        conv_init="zeros")
+    mean, log_scale = tf.split(mean_and_log_std, num_or_size_splits=2, axis=-1)
+    return tfp.distributions.Normal(mean, tf.exp(log_scale))
+
+  def top_cond_prior(self, name, cond_top_latents):
+    """Maps the conditional top latents to a distribution.
+
+    Args:
+      name: variable scope.
+      cond_top_latents: Tensor or a list of tensors.
+                        Latent variables at the previous time-step.
+                        If "pointwise", this is a single tensor.
+                        If "conv_net", this is a list of tensors with length
+                        equal to hparams.num_cond_latents.
+    Returns:
+      cond_dist: tfp.distributions.Normal
+    Raises:
+      ValueError: If cond_top_latents are not of the expected length.
+    """
+    with tf.variable_scope("top", reuse=tf.AUTO_REUSE):
+      if self.hparams.latent_dist_encoder == "pointwise":
+        last_latent = cond_top_latents
+        top = glow_ops.scale_gaussian_prior(
+            name, cond_top_latents, trainable=self.hparams.learn_top_scale)
+      elif self.hparams.latent_dist_encoder == "conv_net":
+        num_cond_latents = (self.hparams.num_cond_latents +
+                            int(self.hparams.cond_first_frame))
+        if len(cond_top_latents) != num_cond_latents:
+          raise ValueError(
+              "Expected length of cond_top_latents %d, got %d"
+              % (num_cond_latents, len(cond_top_latents)))
+        last_latent = cond_top_latents[-1]
+        output_channels = common_layers.shape_list(last_latent)[-1]
+        cond_top_latents = tf.concat(cond_top_latents, axis=-1)
+
+        # Maps the latent-stack to a distribution.
+        top = glow_ops.latent_to_dist(
+            name, cond_top_latents, hparams=self.hparams,
+            output_channels=output_channels)
+      elif self.hparams.latent_dist_encoder == "conv_lstm":
+        last_latent = cond_top_latents
+        output_channels = common_layers.shape_list(cond_top_latents)[-1]
+        # (h_t, c_t) = LSTM(z_{t-1}; (h_{t-1}, c_{t-1}))
+        # (mu_t, sigma_t) = conv(h_t)
+        _, self.top_state = common_video.conv_lstm_2d(
+            cond_top_latents, self.top_state, self.hparams.latent_encoder_width,
+            kernel_size=3, name="conv_lstm")
+        top = glow_ops.single_conv_dist(
+            name, self.top_state.h, output_channels=output_channels)
+      elif self.hparams.latent_dist_encoder == "conv3d_net":
+        last_latent = cond_top_latents[-1]
+        top = glow_ops.temporal_latent_to_dist(
+            "conv3d", tf.stack(cond_top_latents, axis=1), self.hparams)
+
+      # mu(z_{t}) = z_{t-1} + latent_encoder(z_{cond})
+      if self.hparams.latent_skip:
+        top = tfp.distributions.Normal(last_latent + top.loc, top.scale)
+    return top
+
+  def uncond_top_dist(self):
+    """Get an unconditional prior distribution on the top latent."""
+    prior_dist = glow_ops.top_prior(
+        "unconditional", self.z_top_shape, learn_prior="single_conv")
+    return prior_dist.loc, prior_dist.scale
+
+  def cond_top_dist(self, cond_latents):
+    """Get a conditional prior distribution on the top latent."""
+    prior_dist = self.top_cond_prior("conditional", cond_latents)
+    return prior_dist.loc, prior_dist.scale
+
+  def top_prior(self, condition=False, cond_latents=None):
+    """Objective based on the prior over latent z.
+
+    Args:
+      condition: Whether or not to condition on cond_latents.
+      cond_latents: tensor or list of tensors depending on
+                    hparams.latent_dist_encoder
+    Returns:
+      objective: float, log-likelihood of z under the prior.
+      dist: instance of tfp.distributions.Normal, prior distribution.
+    Raises:
+      ValueError: If input is smaller than the prior, uneven height
+                  or rectangular.
+    """
+    if isinstance(condition, bool):
+      condition = tf.constant(condition, dtype=tf.bool)
+    self._all_conds.append(condition)
+
+    if self.hparams.gen_mode == "conditional":
+      # cond_top_latents is None when
+      # latent_dist_encoder is a lstm and frame_ind == 0.
+      # latent_dist_encoder is conv_net and frame_ind < num_cond_frames.
+      marginal_mean, marginal_scale = self.uncond_top_dist()
+      if cond_latents is None:
+        mean, scale = marginal_mean, marginal_scale
+      else:
+        cond_mean, cond_scale = self.cond_top_dist(cond_latents)
+        mean, scale = tf.cond(
+            condition, lambda: (cond_mean, cond_scale),
+            lambda: (marginal_mean, marginal_scale))
+      return glow_ops.TemperedNormal(mean, scale, self.temperature)
+    if self.hparams.top_prior == "prev_frame":
+      return self.get_squeeze_prior()
+    else:
+      return super(NextFrameGlow, self).top_prior()
+
+  def get_z_top_shape(self, init=False):
+    """Get latent shape at level."""
+    if init:
+      batch_size = self.hparams.init_batch_size
+    else:
+      batch_size = self.hparams.batch_size
+    height, _, channels = self.hparams.problem.frame_shape
+    n_levels = self.hparams.n_levels
+    z_width = height // 2**n_levels
+    z_channels = channels * 2**n_levels * 2
+    return [batch_size, z_width, z_width, z_channels]
+
+  def squeeze_video(self, video, init=False):
+    """Squeeze a 5-D Tensor video with one timestep to a 4-D frame."""
+    if init:
+      batch_size = self.hparams.init_batch_size
+    else:
+      batch_size = self.hparams.batch_size
+    frame_shape = [batch_size] + self.hparams.problem.frame_shape
+    return tf.reshape(video, frame_shape)
+
+  def glow_encoder(self, frame, condition=False, cond_latents=None, init=False):
+    """Glow network that encodes frame to a hierarchy of latents.
+
+    Args:
+      frame: 5-D Tensor of shape (batch_size, 1, height, width, channels).
+      condition: Whether or not to condition on cond_latents.
+      cond_latents: optional, list of tensors with length equal to
+                    hparams.n_levels - 1. If provided, the latent at level l is
+                    conditioned on the cond_latent at level l.
+      init: Whether the given batch is an "init" batch or a "train" batch.
+    Returns:
+      objective: log-likelihood of the frame per the model.
+      z_top: top-level latent.
+      z_levels: a list of tensors with latents at all levels.
+    """
+    frame = self.squeeze_video(frame, init=init)
+    frame = self.preprocess(frame)
+    frame, objective = glow_ops.uniform_binning_correction(frame)
+
+    glow_vals = glow_ops.encoder_decoder(
+        "codec", frame, self.hparams, eps=None, reverse=False,
+        cond_latents=cond_latents, states=self.level_states,
+        condition=condition)
+    z_top, encoder_objective, self.eps, z_levels, self.level_states = glow_vals
+    objective += encoder_objective
+    return objective, z_top, z_levels
+
+  def get_num_train_frames(self):
+    """Returns the number of frames as a normalizing factor."""
+    num_target = self.hparams.video_num_target_frames
+    num_input = self.hparams.video_num_input_frames
+
+    # For unconditional generation, this picks a random frame during training
+    # and evaluates the marginal likelihood over "num_input" + "num_target"
+    # frames during eval.
+    if self.hparams.gen_mode == "unconditional":
+      if self.is_training:
+        return 1
+      return num_input + num_target
+
+    # During eval we measure the true objective.
+    if not self.is_training or self.hparams.num_train_frames == -1:
+      total_frames = num_target
+    # if hparams.num_train_frames=-1, we use an approxination to the true
+    # objective.
+    else:
+      total_frames = self.hparams.num_train_frames - num_input
+    if self.hparams.model_input:
+      total_frames += num_input
+    return total_frames
+
+  def get_all_frames(self, input_frames, target_frames):
+    """Get the frames used as input to the model.
+
+    Args:
+      input_frames: 5-D Tensor, (NTHWC)
+      target_frames: 5-D Tensor, (NTHWC)
+    Returns:
+      frames: 5-D Tensor used as input to the model.
+    """
+    if self.is_predicting:
+      all_frames = input_frames
+    elif self.is_training:
+      all_frames = tf.concat((input_frames, target_frames), axis=1)
+      all_frames = common_video.extract_random_video_patch(
+          all_frames, self.hparams.num_train_frames)
+    # Measure the mean bit-per-pixel of the target_frames during eval.
+    else:
+      all_frames = tf.concat((input_frames, target_frames), axis=1)
+    if self.hparams.cond_first_frame:
+      first_frame = all_frames[:, 0:1, :, :, :]
+      all_frames = tf.concat((first_frame, all_frames), axis=1)
+    return all_frames
+
+  def video_objective_tower(self, input_frames, target_frames, init=False):
+    """Returns the bits-per-pixel of the video.
+
+    Args:
+      input_frames: 5-D Tensor of shape (N, 1, H, W, C)
+      target_frames: 5-D Tensor of shape (N, T, H, W, C)
+      init: Whether or not to run data-dependent initialization.
+    Returns:
+      objective: bits-per-pixel.
+    """
+    # The arg_scope call ensures that the actnorm parameters are set such that
+    # the per-channel output activations have zero mean and unit variance
+    # ONLY during the first step. After that the parameters are learned
+    # through optimisation.
+    num_input_frames = (self.hparams.video_num_input_frames +
+                        int(self.hparams.cond_first_frame))
+
+    # Set num total frames to average the objective.
+    total_frames = self.get_num_train_frames()
+
+    # Compute the log-likelihood of target_frames at both train and predict
+    # time.
+    all_frames = self.get_all_frames(input_frames, target_frames)
+    all_frames = tf.unstack(all_frames, axis=1)
+
+    cond_level_latents, cond_top_latents = None, None
+    total_objective = 0.0
+    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
+
+    with arg_scope(ops, init=init):
+      for frame_ind, frame in enumerate(all_frames):
+
+        # Get current set of cond latents of non-top levels.
+        cond_level, cond_level_latents = get_cond_latents(
+            self.all_level_latents, self.hparams)
+
+        # Get current set of cond latents of the top-level
+        cond_top, cond_top_latents = get_cond_latents(
+            self.all_top_latents, self.hparams)
+
+        # Superscript -> level, Subscript -> Time.
+        # (z^{0}_t, z^{1..l}_t) = Glow(X_{t}, z^{1..l}_{cond_t})
+        frame_obj, curr_top_latent, curr_level_latents = self.glow_encoder(
+            frame, condition=cond_level, cond_latents=cond_level_latents,
+            init=init)
+
+        # z^0_t ~ N(f(z^0_{t-1}))
+        # cond_top_latents is None when
+        # latent_dist_encoder is conv_net and frame_ind < num_cond_frames.
+        prior_dist = self.top_prior(
+            condition=cond_top, cond_latents=cond_top_latents)
+        prior_objective = tf.reduce_sum(
+            prior_dist.log_prob(curr_top_latent), axis=[1, 2, 3])
+        frame_obj += prior_objective
+
+        # Loss computation.
+        # Do not model the probabililty of the input frames by default.
+        # Consistent with other video models.
+        if (frame_ind > num_input_frames - 1 or self.hparams.model_input or
+            self.hparams.gen_mode == "unconditional"):
+          total_objective += frame_obj
+        self.all_level_latents.append(curr_level_latents)
+        self.all_top_latents.append(curr_top_latent)
+
+      # During prediction time, store z_sample ~ N(f(z_{num_input_frames}))
+      # to generate the first target frame.
+      if self.is_predicting:
+        # Get current set of cond_top_latents
+        cond_top, cond_top_latents = get_cond_latents(
+            self.all_top_latents, self.hparams)
+        prior_dist = self.top_prior(
+            condition=cond_top, cond_latents=cond_top_latents)
+        self.z_sample = prior_dist.sample()
+        self.all_top_latents.append(self.z_sample)
+
+      # Converts log-probability to bits-per-pixel.
+      hwc = np.prod(self.hparams.problem.frame_shape)
+      total_objective = -total_objective / (np.log(2) * hwc * total_frames)
+    return total_objective
+
+  def objective_tower(self, features, init=False):
+    input_frames, target_frames = features["inputs"], features["targets"]
+    self.cond_latents, self.top_state = None, None
+    self.all_level_latents, self.all_top_latents = [], []
+    self._all_conds = []
+    self.level_states = [None] * (self.hparams.n_levels - 1)
+    self.z_top_shape = self.get_z_top_shape(init=init)
+    num_input_frames = self.hparams.video_num_input_frames
+    latent_dist_encoder = self.hparams.latent_dist_encoder
+    num_cond_latents = self.hparams.num_cond_latents
+
+    exp_modes = ["conditional", "unconditional"]
+    if self.hparams.gen_mode not in exp_modes:
+      raise ValueError("Expected mode to be in %s, got %s" %
+                       (exp_modes, self.hparams.gen_mode))
+
+    # Error checks for conditional video generation.
+    if self.hparams.gen_mode == "conditional":
+      exp_latent_encoders = ["pointwise", "conv_net", "conv_lstm", "conv3d_net"]
+      if latent_dist_encoder not in exp_latent_encoders:
+        raise ValueError("Expected latent_dist_encoder is %s, got %s" %
+                         (exp_latent_encoders, latent_dist_encoder))
+      if (latent_dist_encoder == "pointwise" and num_cond_latents != 1):
+        raise ValueError("Expected num_cond_latents: 1, with 'pointwise' "
+                         "latent_dist_encoder, got %d" % num_cond_latents)
+      if (latent_dist_encoder == "conv_net" and
+          num_cond_latents > num_input_frames):
+        raise ValueError("Expected num_cond_latents <= %d, got %d" %
+                         (num_input_frames, num_cond_latents))
+      if (latent_dist_encoder == "pointwise" and
+          self.hparams.init_batch_size != self.hparams.batch_size):
+        raise ValueError("init_batch_size different from batch_size not "
+                         "supported for latent_dist_encoder=pointwise")
+    if self.hparams.gen_mode == "unconditional":
+      if self.hparams.num_train_frames != 1:
+        raise ValueError("Expected num_train_frames to be 1 when "
+                         "hparams.gen_mode is unconditional, got %d" %
+                         self.hparams.num_train_frames)
+      if self.hparams.video_num_input_frames != 1:
+        raise ValueError("Expected num_input_frames to be 1 when "
+                         "hparams.gen_mode is unconditional, got %d" %
+                         self.hparams.video_num_input_frames)
+    return self.video_objective_tower(input_frames, target_frames, init=init)
diff --git a/tensor2tensor/models/video/nfg_conv3d_test.py b/tensor2tensor/models/video/nfg_conv3d_test.py
new file mode 100644
index 000000000..689a61ce7
--- /dev/null
+++ b/tensor2tensor/models/video/nfg_conv3d_test.py
@@ -0,0 +1,49 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Test when the latent-network encoder is a conv3d net."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+from tensor2tensor.models.video import nfg_test_utils
+import tensorflow as tf
+
+
+conv3d_net_hparams = (
+    ("conv3d_net", 2, 2, "conv3d_net", "conditional", -1, 3),
+    ("conv3d_dil", 2, 2, "conv3d_net", "conditional", -1, -1, False, True),)
+
+
+class NextFrameGlowConv3DTest(nfg_test_utils.NextFrameGlowTest,
+                              parameterized.TestCase):
+
+  @parameterized.named_parameters(*conv3d_net_hparams)
+  def testGlowTrainAndDecode(self, in_frames=1, out_frames=1,
+                             latent_dist_encoder="pointwise",
+                             gen_mode="conditional", pretrain_steps=-1,
+                             num_train_frames=-1, cond_first_frame=False,
+                             apply_dilations=False):
+    self.GlowTrainAndDecode(
+        in_frames=in_frames, out_frames=out_frames,
+        latent_dist_encoder=latent_dist_encoder, gen_mode=gen_mode,
+        pretrain_steps=pretrain_steps, num_train_frames=num_train_frames,
+        cond_first_frame=cond_first_frame, apply_dilations=apply_dilations)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/video/nfg_conv_lstm_test.py b/tensor2tensor/models/video/nfg_conv_lstm_test.py
new file mode 100644
index 000000000..beda25f0e
--- /dev/null
+++ b/tensor2tensor/models/video/nfg_conv_lstm_test.py
@@ -0,0 +1,48 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Test when the latent-network encoder is a conv-lstm."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+from tensor2tensor.models.video import nfg_test_utils
+import tensorflow as tf
+
+
+conv_lstm_hparams = (
+    ("in_3_out_2_lstm", 2, 1, "conv_lstm", "conditional", -1),
+    ("lstm_pretrain", 2, 1, "conv_lstm", "conditional", 50000))
+
+
+class NextFrameGlowConv3DTest(nfg_test_utils.NextFrameGlowTest,
+                              parameterized.TestCase):
+
+  @parameterized.named_parameters(*conv_lstm_hparams)
+  def testGlowTrainAndDecode(self, in_frames=1, out_frames=1,
+                             latent_dist_encoder="pointwise",
+                             gen_mode="conditional", pretrain_steps=-1,
+                             num_train_frames=-1, cond_first_frame=False):
+    self.GlowTrainAndDecode(
+        in_frames=in_frames, out_frames=out_frames,
+        latent_dist_encoder=latent_dist_encoder, gen_mode=gen_mode,
+        pretrain_steps=pretrain_steps, num_train_frames=num_train_frames,
+        cond_first_frame=cond_first_frame)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/video/nfg_conv_test.py b/tensor2tensor/models/video/nfg_conv_test.py
new file mode 100644
index 000000000..12eb657d7
--- /dev/null
+++ b/tensor2tensor/models/video/nfg_conv_test.py
@@ -0,0 +1,44 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Test when the latent-network encoder is a 2-D conv."""
+
+from absl.testing import parameterized
+from tensor2tensor.models.video import nfg_test_utils
+import tensorflow as tf
+
+
+conv_net_hparams = (
+    ("in_3_out_2_conv", 3, 1, "conv_net", "conditional"),
+    ("conv_net_cond_first", 2, 2, "conv_net", "conditional", -1, 3, True),)
+
+
+class NextFrameGlowConvTest(nfg_test_utils.NextFrameGlowTest,
+                            parameterized.TestCase):
+
+  @parameterized.named_parameters(*conv_net_hparams)
+  def testGlowTrainAndDecode(self, in_frames=1, out_frames=1,
+                             latent_dist_encoder="pointwise",
+                             gen_mode="conditional", pretrain_steps=-1,
+                             num_train_frames=-1, cond_first_frame=False):
+    self.GlowTrainAndDecode(
+        in_frames=in_frames, out_frames=out_frames, gen_mode=gen_mode,
+        latent_dist_encoder=latent_dist_encoder,
+        pretrain_steps=pretrain_steps, num_train_frames=num_train_frames,
+        cond_first_frame=cond_first_frame)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/video/nfg_test_utils.py b/tensor2tensor/models/video/nfg_test_utils.py
new file mode 100644
index 000000000..2b47cdfe2
--- /dev/null
+++ b/tensor2tensor/models/video/nfg_test_utils.py
@@ -0,0 +1,184 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Testing utils for next_frame_glow."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import tempfile
+import numpy as np
+from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
+from tensor2tensor.models.video import next_frame_glow
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+MODES = tf.estimator.ModeKeys
+
+
+# TODO(mechcoder): Refactor or merge tests with the other next_frame_tests when
+# this moves to a public version.
+def fill_hparams(hparams, in_frames, out_frames, gen_mode="conditional",
+                 latent_dist_encoder="pointwise", pretrain_steps=-1,
+                 num_train_frames=-1, cond_first_frame=False,
+                 apply_dilations=False):
+  """Set next_frame_glow hparams."""
+  hparams.latent_apply_dilations = apply_dilations
+  hparams.video_num_input_frames = in_frames
+  hparams.video_num_target_frames = out_frames
+  hparams.latent_dist_encoder = latent_dist_encoder
+  hparams.gen_mode = gen_mode
+  hparams.pretrain_steps = pretrain_steps
+  hparams.num_train_frames = num_train_frames
+  hparams.cond_first_frame = cond_first_frame
+  if latent_dist_encoder in ["conv_net", "conv3d_net"]:
+    hparams.num_cond_latents = in_frames
+  else:
+    hparams.num_cond_latents = 1
+  problem = registry.problem("video_stochastic_shapes10k")
+  p_hparams = problem.get_hparams(hparams)
+  hparams.problem = problem
+  hparams.problem_hparams = p_hparams
+  hparams.tiny_mode = True
+  hparams.reward_prediction = False
+  hparams.latent_architecture = "glow_resnet"
+  hparams.latent_encoder_depth = 2
+  hparams.latent_pre_output_channels = 32
+  if (hparams.gen_mode == "conditional" and
+      hparams.latent_dist_encoder == "pointwise"):
+    hparams.batch_size = 16
+    hparams.init_batch_size = 16
+  else:
+    hparams.batch_size = 16
+    hparams.init_batch_size = 32
+  hparams.affine_coupling_width = 32
+  hparams.depth = 5
+  hparams.n_levels = 2
+  return hparams
+
+
+def fill_infer_targets(x):
+  x["infer_targets"] = tf.identity(x["targets"])
+  return x
+
+
+def create_basic_features(hparams):
+  dataset = hparams.problem.dataset(MODES.TRAIN, hparams=hparams)
+  dataset = dataset.batch(hparams.batch_size)
+  dataset = dataset.map(fill_infer_targets)
+  return dataset.make_one_shot_iterator().get_next()
+
+
+class NextFrameGlowTest(tf.test.TestCase):
+  """Utils for testing next_frame_glow."""
+
+  def should_run_session(self, hparams):
+    # dilated conv-3d not available on CPU.
+    return tf.test.is_gpu_available() or not hparams.latent_apply_dilations
+
+  def checkAllConds(self, conds_array, num_total_frames, hparams):
+    if hparams.cond_first_frame:
+      self.assertEqual(conds_array, [True]*(num_total_frames + 1))
+    elif hparams.pretrain_steps > -1:
+      self.assertEqual(conds_array, [False]*num_total_frames)
+    elif hparams.latent_dist_encoder != "pointwise":
+      self.assertEqual(conds_array, [True]*num_total_frames)
+
+  def RunModel(self, model, train_op, hparams, features, num_frames,
+               model_path=None):
+    exp_num_frames = num_frames + int(hparams.cond_first_frame)
+    if hparams.gen_mode == "conditional":
+      self.assertLen(model.all_top_latents, exp_num_frames)
+      self.assertLen(model.all_level_latents, exp_num_frames)
+
+    with tf.Session() as session:
+
+      if model_path is not None:
+        saver = tf.train.Saver()
+
+      session.run(tf.global_variables_initializer())
+
+      # Run initialization.
+      init_op = tf.get_collection("glow_init_op")
+      session.run(init_op)
+
+      loss, top_conds = session.run([train_op["training"], model._all_conds])  # pylint: disable=protected-access
+      self.checkAllConds(top_conds, num_frames, hparams)
+
+      if model_path is not None:
+        saver.save(session, model_path)
+
+      # Check that one forward-propagation does not NaN, i.e
+      # initialization etc works as expected.
+      self.assertTrue(loss > 0.0 and loss < 10.0)
+
+  def GlowTrainAndDecode(self, in_frames=1, out_frames=1,
+                         latent_dist_encoder="pointwise",
+                         gen_mode="conditional", pretrain_steps=-1,
+                         num_train_frames=-1, cond_first_frame=False,
+                         apply_dilations=False):
+    """Test 1 forward pass and sampling gives reasonable results."""
+    if num_train_frames == -1:
+      total_frames = in_frames + out_frames
+    else:
+      total_frames = num_train_frames
+
+    curr_dir = tempfile.mkdtemp()
+    model_path = os.path.join(curr_dir, "model")
+
+    # Training pipeline
+    with tf.Graph().as_default():
+      hparams = next_frame_glow.next_frame_glow_hparams()
+      hparams = fill_hparams(hparams, in_frames, out_frames,
+                             gen_mode, latent_dist_encoder, pretrain_steps,
+                             num_train_frames, cond_first_frame,
+                             apply_dilations)
+      features = create_basic_features(hparams)
+      model = next_frame_glow.NextFrameGlow(hparams, MODES.TRAIN)
+      _, train_op = model(features)
+      if self.should_run_session(hparams):
+        self.RunModel(model, train_op, hparams, features, total_frames,
+                      model_path)
+
+    # Inference pipeline
+    with tf.Graph().as_default():
+      hparams = next_frame_glow.next_frame_glow_hparams()
+      if hparams.gen_mode == "unconditional":
+        hparams.video_num_target_frames = 1
+      hparams = fill_hparams(hparams, in_frames, out_frames,
+                             gen_mode, latent_dist_encoder, pretrain_steps,
+                             num_train_frames, cond_first_frame,
+                             apply_dilations)
+      features = create_basic_features(hparams)
+      model = next_frame_glow.NextFrameGlow(
+          hparams, tf.estimator.ModeKeys.PREDICT)
+      predictions = model.infer(features)
+      outputs = predictions["outputs"]
+      model_path = os.path.join(curr_dir, "model")
+
+      if self.should_run_session(hparams):
+        with tf.Session() as session:
+          saver = tf.train.Saver()
+          saver.restore(session, model_path)
+          outputs_np = session.run(outputs)
+          self.assertEqual(outputs_np.shape, (16, out_frames, 64, 64, 3))
+          self.assertTrue(np.all(outputs_np <= 255))
+          self.assertTrue(np.all(outputs_np >= 0))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/video/nfg_uncond_test.py b/tensor2tensor/models/video/nfg_uncond_test.py
new file mode 100644
index 000000000..e2ceb50d3
--- /dev/null
+++ b/tensor2tensor/models/video/nfg_uncond_test.py
@@ -0,0 +1,46 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 unconditional glow."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+from tensor2tensor.models.video import nfg_test_utils
+import tensorflow as tf
+
+uncond_hparams = (
+    ("in_1_out_1", 1, 1, "pointwise", "conditional"),
+    ("uncond", 1, 3, "pointwise", "unconditional", -1, 1),)
+
+
+class NfgUncondTest(nfg_test_utils.NextFrameGlowTest, parameterized.TestCase):
+
+  @parameterized.named_parameters(*uncond_hparams)
+  def testGlowTrainAndDecode(self, in_frames=1, out_frames=1,
+                             latent_dist_encoder="pointwise",
+                             gen_mode="conditional", pretrain_steps=-1,
+                             num_train_frames=-1, cond_first_frame=False):
+    self.GlowTrainAndDecode(
+        in_frames=in_frames, out_frames=out_frames,
+        latent_dist_encoder=latent_dist_encoder, gen_mode=gen_mode,
+        pretrain_steps=pretrain_steps, num_train_frames=num_train_frames,
+        cond_first_frame=cond_first_frame)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 2624530ac2a467cf947d5b37045557e11433e73c Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 3 Dec 2018 22:04:41 -0800
Subject: [PATCH 1314/2720] Move nfg to models/video for external
 collaboration.

PiperOrigin-RevId: 223919316
---
 tensor2tensor/models/__init__.py              |   1 -
 tensor2tensor/models/video/next_frame_glow.py | 572 ------------------
 tensor2tensor/models/video/nfg_conv3d_test.py |  49 --
 .../models/video/nfg_conv_lstm_test.py        |  48 --
 tensor2tensor/models/video/nfg_conv_test.py   |  44 --
 tensor2tensor/models/video/nfg_test_utils.py  | 184 ------
 tensor2tensor/models/video/nfg_uncond_test.py |  46 --
 7 files changed, 944 deletions(-)
 delete mode 100644 tensor2tensor/models/video/next_frame_glow.py
 delete mode 100644 tensor2tensor/models/video/nfg_conv3d_test.py
 delete mode 100644 tensor2tensor/models/video/nfg_conv_lstm_test.py
 delete mode 100644 tensor2tensor/models/video/nfg_conv_test.py
 delete mode 100644 tensor2tensor/models/video/nfg_test_utils.py
 delete mode 100644 tensor2tensor/models/video/nfg_uncond_test.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 3afc44e2d..05cafbde7 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -67,7 +67,6 @@
 from tensor2tensor.models.video import basic_stochastic
 from tensor2tensor.models.video import emily
 from tensor2tensor.models.video import epva
-from tensor2tensor.models.video import next_frame_glow
 from tensor2tensor.models.video import savp
 from tensor2tensor.models.video import sv2p
 from tensor2tensor.models.video import svg_lp
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
deleted file mode 100644
index 70583154d..000000000
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ /dev/null
@@ -1,572 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Experimental testbed for a conditional glow model."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-from tensor2tensor.layers import common_layers
-from tensor2tensor.layers import common_video
-from tensor2tensor.layers import modalities
-from tensor2tensor.models.research import glow
-from tensor2tensor.models.research import glow_ops
-from tensor2tensor.utils import registry
-import tensorflow as tf
-import tensorflow_probability as tfp
-
-
-arg_scope = tf.contrib.framework.arg_scope
-
-
-@registry.register_hparams
-def next_frame_glow_hparams():
-  """Hparams for next_frame_glow."""
-  hparams = glow.glow_hparams()
-  # Possible modes are conditional and unconditional
-  hparams.add_hparam("gen_mode", "conditional")
-  hparams.add_hparam("learn_top_scale", False)
-  hparams.add_hparam("condition_all_levels", True)
-  # For each video, substitutes "num_input_frames + num_output_frames" with a
-  # randomly sampled patch of length "num_train_frames" during training.
-  # -1 indicates that the entire video is used for training.
-  hparams.add_hparam("num_train_frames", -1)
-  # The following are hparams that model the latent transitions.
-  # Encoder that maps the latents to a Gaussian distribution.
-  # This function is used to model the prior over z_{t}. Can be,
-  # Pointwise -> point-wise multiplication of z_{t-1}.
-  # conv_net -> one-layer convolution over z_{t-1} .. z_{t - num_cond_latents}
-  hparams.add_hparam("latent_dist_encoder", "conv_net")
-  # Number of latents used in the encoder above.
-  hparams.add_hparam("num_cond_latents", 1)
-  hparams.add_hparam("latent_architecture", "glow_resnet")
-  hparams.add_hparam("latent_apply_dilations", False)
-  # Use latent skip connections
-  hparams.add_hparam("model_input", False)
-  hparams.add_hparam("cond_first_frame", False)
-  hparams.add_hparam("latent_skip", True)
-  hparams.add_hparam("latent_encoder_depth", 2)
-  hparams.add_hparam("latent_encoder_width", 512)
-  hparams.add_hparam("latent_pre_output_channels", 512)
-  # Pretrains the glow encoder for "pretrain_steps" number of steps.
-  # By default, don't pretrain and learn end-to-end
-  hparams.add_hparam("pretrain_steps", -1)
-  hparams.modality = {
-      "inputs": modalities.VideoModalityL1Raw,
-      "targets": modalities.VideoModalityL1Raw,
-  }
-  hparams.init_batch_size = 256
-  hparams.batch_size = 32
-  # Possible options: are prev_frame, single_conv and normal
-  hparams.top_prior = "single_conv"
-  return hparams
-
-
-@registry.register_hparams
-def frame_glow_hparams():
-  """Unconditional generation on video-frames."""
-  hparams = next_frame_glow_hparams()
-  hparams.gen_mode = "unconditional"
-  hparams.num_train_frames = 1
-  return hparams
-
-
-def get_cond_latents(all_latents=None, hparams=None):
-  """Get z^{cond}_{t} given z^{1..t-1}.
-
-  Args:
-    all_latents: list of list of tensors,
-                 outer-size equals no.of time_steps-1
-                 inner-size equals hparams.n_levels.
-    hparams: See next_frame_glow_hparams.
-  Returns:
-    cond_latents: conditional latents at time-step t.
-  """
-  cond_latents = None
-  if hparams.gen_mode == "conditional":
-    if hparams.latent_dist_encoder in ["conv_net", "conv3d_net"]:
-      num_cond_latents = (hparams.num_cond_latents +
-                          int(hparams.cond_first_frame))
-      if len(all_latents) >= num_cond_latents:
-        cond_latents = all_latents[-hparams.num_cond_latents:]
-        if hparams.cond_first_frame:
-          cond_latents = [all_latents[0]] + cond_latents
-    elif hparams.latent_dist_encoder in ["pointwise", "conv_lstm"]:
-      if all_latents:
-        cond_latents = all_latents[-1]
-
-  if hparams.gen_mode == "conditional":
-    global_step = tf.train.get_or_create_global_step()
-    condition = tf.greater(global_step, hparams.pretrain_steps)
-  else:
-    condition = tf.constant(False, dtype=tf.bool)
-  return condition, cond_latents
-
-
-@registry.register_model
-class NextFrameGlow(glow.Glow):
-  """Extend Glow for video."""
-
-  def init_preprocess_single(self, features):
-    for label in ["inputs", "targets"]:
-      features[label] = common_layers.convert_rgb_to_real(features[label])
-    return features
-
-  def init_preprocess(self, features):
-    """Preprocessing as per the input modality.
-
-    Equivalent to calling self.bottom(features).
-
-    Args:
-      features: dict of strings to tensors.
-    Returns:
-      features: dict of strings to tensors.
-    """
-    return features.map(self.init_preprocess_single)
-
-  def preprocess(self, x):
-    """Converts x from [0, 1] to [-0.5, 0.5].
-
-    All inputs are already normalized to be in the range [0, 1] through the
-    VideoModalityL1Raw modality.
-
-    Args:
-      x: 4-D Tensor.
-
-    Returns:
-      x: Scaled such that x lies in-between -0.5 and 0.5
-    """
-    return x - 0.5
-
-  def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
-    del args, kwargs
-
-    # Make a copy of features that can be used in the call to self
-    # that builds the graph.
-    new_features = {}
-    new_features["inputs"] = features["inputs"]
-    new_features["targets"] = features["infer_targets"]
-    _, _ = self(new_features)  # pylint: disable=not-callable
-
-    if self.hparams.gen_mode == "unconditional":
-      num_target_frames = 1
-    else:
-      num_target_frames = self.hparams.video_num_target_frames
-
-    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
-    var_scope = tf.variable_scope("next_frame_glow/body", reuse=True)
-    all_frames = []
-
-    # If eps=None, images are sampled from the prior.
-    with arg_scope(ops, init=False), var_scope:
-      for target_frame in range(1, num_target_frames + 1):
-
-        # subscript -> timestep, superscript -> level.
-        # self.z_sample equals z^0_{t} (top-level latent)
-        # (X_{t}, z^{1..l}_{t}) = Glow(z^0_{t}, z^{1..l}_{t-1})
-        # Get current set of cond_latents.
-        cond_level, cond_level_latents = get_cond_latents(
-            self.all_level_latents, self.hparams)
-
-        glow_vals = glow_ops.encoder_decoder(
-            "codec", self.z_sample, self.hparams, eps=None, reverse=True,
-            cond_latents=cond_level_latents, states=self.level_states,
-            condition=cond_level, temperature=self.temperature)
-        predicted_frame, _, curr_latents, self.level_states = glow_vals
-        all_frames.append(predicted_frame)
-        self.all_level_latents.append(curr_latents)
-
-        # Compute z^0_{t+1} = f(z^0_{t})
-        if target_frame < num_target_frames:
-          cond_top, cond_top_latents = get_cond_latents(
-              self.all_top_latents, self.hparams)
-          prior_dist = self.top_prior(
-              condition=cond_top, cond_latents=cond_top_latents)
-          self.z_sample = prior_dist.sample()
-          self.all_top_latents.append(self.z_sample)
-
-    all_frames = tf.stack(all_frames)
-    predicted_video = common_video.swap_time_and_batch_axes(all_frames)
-
-    # The video-decode API requires the predicted video to be the same shape
-    # as the target-video. Hence, for unconditional generation,
-    # tile across time to ensure same shape.
-    if self.hparams.gen_mode == "unconditional":
-      predicted_video = tf.tile(
-          predicted_video, [1, self.hparams.video_num_target_frames, 1, 1, 1])
-    predicted_video = self.scale(predicted_video)
-
-    # Output of a single decode / sample.
-    output_features = {}
-    output_features["targets"] = tf.zeros_like(predicted_video)
-    output_features["outputs"] = predicted_video
-    output_features["scores"] = tf.zeros_like(predicted_video)
-    return output_features
-
-  def get_squeeze_prior(self):
-    """Model the prior over z_{t} as a function of X_{t-1}.
-
-    Returns:
-      objective: float, log-likelihood.
-      dist: instance of tfp.distributions.Normal.
-
-    Raises:
-      ValueError: If input_height is not equal to input_width, not even
-                   or if the image width is smaller than the latent width.
-    """
-    _, prior_height, _, prior_channels = self.z_top_shape
-    _, input_height, input_width, _ = common_layers.shape_list(self.input_frame)
-
-    if input_height != input_width:
-      raise ValueError("input height should be equal to input width")
-    if input_height % 2 != 0:
-      raise ValueError("input height should be even")
-    if input_height < prior_height:
-      raise ValueError("input should be larger than the prior.")
-
-    # mean, log_std = NN(X_0)
-    # Reduce the spatial dimension by a factor of "squeeze_factor".
-    # and convolve with a stride of 2
-    squeeze_factor = input_height // (2 * prior_height)
-    x = glow_ops.squeeze(
-        "prior_squeeze", self.input_frame, factor=squeeze_factor, reverse=False)
-    mean_and_log_std = glow_ops.conv(
-        "prior_conv", x, 2*prior_channels, stride=[2, 2], apply_actnorm=False,
-        conv_init="zeros")
-    mean, log_scale = tf.split(mean_and_log_std, num_or_size_splits=2, axis=-1)
-    return tfp.distributions.Normal(mean, tf.exp(log_scale))
-
-  def top_cond_prior(self, name, cond_top_latents):
-    """Maps the conditional top latents to a distribution.
-
-    Args:
-      name: variable scope.
-      cond_top_latents: Tensor or a list of tensors.
-                        Latent variables at the previous time-step.
-                        If "pointwise", this is a single tensor.
-                        If "conv_net", this is a list of tensors with length
-                        equal to hparams.num_cond_latents.
-    Returns:
-      cond_dist: tfp.distributions.Normal
-    Raises:
-      ValueError: If cond_top_latents are not of the expected length.
-    """
-    with tf.variable_scope("top", reuse=tf.AUTO_REUSE):
-      if self.hparams.latent_dist_encoder == "pointwise":
-        last_latent = cond_top_latents
-        top = glow_ops.scale_gaussian_prior(
-            name, cond_top_latents, trainable=self.hparams.learn_top_scale)
-      elif self.hparams.latent_dist_encoder == "conv_net":
-        num_cond_latents = (self.hparams.num_cond_latents +
-                            int(self.hparams.cond_first_frame))
-        if len(cond_top_latents) != num_cond_latents:
-          raise ValueError(
-              "Expected length of cond_top_latents %d, got %d"
-              % (num_cond_latents, len(cond_top_latents)))
-        last_latent = cond_top_latents[-1]
-        output_channels = common_layers.shape_list(last_latent)[-1]
-        cond_top_latents = tf.concat(cond_top_latents, axis=-1)
-
-        # Maps the latent-stack to a distribution.
-        top = glow_ops.latent_to_dist(
-            name, cond_top_latents, hparams=self.hparams,
-            output_channels=output_channels)
-      elif self.hparams.latent_dist_encoder == "conv_lstm":
-        last_latent = cond_top_latents
-        output_channels = common_layers.shape_list(cond_top_latents)[-1]
-        # (h_t, c_t) = LSTM(z_{t-1}; (h_{t-1}, c_{t-1}))
-        # (mu_t, sigma_t) = conv(h_t)
-        _, self.top_state = common_video.conv_lstm_2d(
-            cond_top_latents, self.top_state, self.hparams.latent_encoder_width,
-            kernel_size=3, name="conv_lstm")
-        top = glow_ops.single_conv_dist(
-            name, self.top_state.h, output_channels=output_channels)
-      elif self.hparams.latent_dist_encoder == "conv3d_net":
-        last_latent = cond_top_latents[-1]
-        top = glow_ops.temporal_latent_to_dist(
-            "conv3d", tf.stack(cond_top_latents, axis=1), self.hparams)
-
-      # mu(z_{t}) = z_{t-1} + latent_encoder(z_{cond})
-      if self.hparams.latent_skip:
-        top = tfp.distributions.Normal(last_latent + top.loc, top.scale)
-    return top
-
-  def uncond_top_dist(self):
-    """Get an unconditional prior distribution on the top latent."""
-    prior_dist = glow_ops.top_prior(
-        "unconditional", self.z_top_shape, learn_prior="single_conv")
-    return prior_dist.loc, prior_dist.scale
-
-  def cond_top_dist(self, cond_latents):
-    """Get a conditional prior distribution on the top latent."""
-    prior_dist = self.top_cond_prior("conditional", cond_latents)
-    return prior_dist.loc, prior_dist.scale
-
-  def top_prior(self, condition=False, cond_latents=None):
-    """Objective based on the prior over latent z.
-
-    Args:
-      condition: Whether or not to condition on cond_latents.
-      cond_latents: tensor or list of tensors depending on
-                    hparams.latent_dist_encoder
-    Returns:
-      objective: float, log-likelihood of z under the prior.
-      dist: instance of tfp.distributions.Normal, prior distribution.
-    Raises:
-      ValueError: If input is smaller than the prior, uneven height
-                  or rectangular.
-    """
-    if isinstance(condition, bool):
-      condition = tf.constant(condition, dtype=tf.bool)
-    self._all_conds.append(condition)
-
-    if self.hparams.gen_mode == "conditional":
-      # cond_top_latents is None when
-      # latent_dist_encoder is a lstm and frame_ind == 0.
-      # latent_dist_encoder is conv_net and frame_ind < num_cond_frames.
-      marginal_mean, marginal_scale = self.uncond_top_dist()
-      if cond_latents is None:
-        mean, scale = marginal_mean, marginal_scale
-      else:
-        cond_mean, cond_scale = self.cond_top_dist(cond_latents)
-        mean, scale = tf.cond(
-            condition, lambda: (cond_mean, cond_scale),
-            lambda: (marginal_mean, marginal_scale))
-      return glow_ops.TemperedNormal(mean, scale, self.temperature)
-    if self.hparams.top_prior == "prev_frame":
-      return self.get_squeeze_prior()
-    else:
-      return super(NextFrameGlow, self).top_prior()
-
-  def get_z_top_shape(self, init=False):
-    """Get latent shape at level."""
-    if init:
-      batch_size = self.hparams.init_batch_size
-    else:
-      batch_size = self.hparams.batch_size
-    height, _, channels = self.hparams.problem.frame_shape
-    n_levels = self.hparams.n_levels
-    z_width = height // 2**n_levels
-    z_channels = channels * 2**n_levels * 2
-    return [batch_size, z_width, z_width, z_channels]
-
-  def squeeze_video(self, video, init=False):
-    """Squeeze a 5-D Tensor video with one timestep to a 4-D frame."""
-    if init:
-      batch_size = self.hparams.init_batch_size
-    else:
-      batch_size = self.hparams.batch_size
-    frame_shape = [batch_size] + self.hparams.problem.frame_shape
-    return tf.reshape(video, frame_shape)
-
-  def glow_encoder(self, frame, condition=False, cond_latents=None, init=False):
-    """Glow network that encodes frame to a hierarchy of latents.
-
-    Args:
-      frame: 5-D Tensor of shape (batch_size, 1, height, width, channels).
-      condition: Whether or not to condition on cond_latents.
-      cond_latents: optional, list of tensors with length equal to
-                    hparams.n_levels - 1. If provided, the latent at level l is
-                    conditioned on the cond_latent at level l.
-      init: Whether the given batch is an "init" batch or a "train" batch.
-    Returns:
-      objective: log-likelihood of the frame per the model.
-      z_top: top-level latent.
-      z_levels: a list of tensors with latents at all levels.
-    """
-    frame = self.squeeze_video(frame, init=init)
-    frame = self.preprocess(frame)
-    frame, objective = glow_ops.uniform_binning_correction(frame)
-
-    glow_vals = glow_ops.encoder_decoder(
-        "codec", frame, self.hparams, eps=None, reverse=False,
-        cond_latents=cond_latents, states=self.level_states,
-        condition=condition)
-    z_top, encoder_objective, self.eps, z_levels, self.level_states = glow_vals
-    objective += encoder_objective
-    return objective, z_top, z_levels
-
-  def get_num_train_frames(self):
-    """Returns the number of frames as a normalizing factor."""
-    num_target = self.hparams.video_num_target_frames
-    num_input = self.hparams.video_num_input_frames
-
-    # For unconditional generation, this picks a random frame during training
-    # and evaluates the marginal likelihood over "num_input" + "num_target"
-    # frames during eval.
-    if self.hparams.gen_mode == "unconditional":
-      if self.is_training:
-        return 1
-      return num_input + num_target
-
-    # During eval we measure the true objective.
-    if not self.is_training or self.hparams.num_train_frames == -1:
-      total_frames = num_target
-    # if hparams.num_train_frames=-1, we use an approxination to the true
-    # objective.
-    else:
-      total_frames = self.hparams.num_train_frames - num_input
-    if self.hparams.model_input:
-      total_frames += num_input
-    return total_frames
-
-  def get_all_frames(self, input_frames, target_frames):
-    """Get the frames used as input to the model.
-
-    Args:
-      input_frames: 5-D Tensor, (NTHWC)
-      target_frames: 5-D Tensor, (NTHWC)
-    Returns:
-      frames: 5-D Tensor used as input to the model.
-    """
-    if self.is_predicting:
-      all_frames = input_frames
-    elif self.is_training:
-      all_frames = tf.concat((input_frames, target_frames), axis=1)
-      all_frames = common_video.extract_random_video_patch(
-          all_frames, self.hparams.num_train_frames)
-    # Measure the mean bit-per-pixel of the target_frames during eval.
-    else:
-      all_frames = tf.concat((input_frames, target_frames), axis=1)
-    if self.hparams.cond_first_frame:
-      first_frame = all_frames[:, 0:1, :, :, :]
-      all_frames = tf.concat((first_frame, all_frames), axis=1)
-    return all_frames
-
-  def video_objective_tower(self, input_frames, target_frames, init=False):
-    """Returns the bits-per-pixel of the video.
-
-    Args:
-      input_frames: 5-D Tensor of shape (N, 1, H, W, C)
-      target_frames: 5-D Tensor of shape (N, T, H, W, C)
-      init: Whether or not to run data-dependent initialization.
-    Returns:
-      objective: bits-per-pixel.
-    """
-    # The arg_scope call ensures that the actnorm parameters are set such that
-    # the per-channel output activations have zero mean and unit variance
-    # ONLY during the first step. After that the parameters are learned
-    # through optimisation.
-    num_input_frames = (self.hparams.video_num_input_frames +
-                        int(self.hparams.cond_first_frame))
-
-    # Set num total frames to average the objective.
-    total_frames = self.get_num_train_frames()
-
-    # Compute the log-likelihood of target_frames at both train and predict
-    # time.
-    all_frames = self.get_all_frames(input_frames, target_frames)
-    all_frames = tf.unstack(all_frames, axis=1)
-
-    cond_level_latents, cond_top_latents = None, None
-    total_objective = 0.0
-    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
-
-    with arg_scope(ops, init=init):
-      for frame_ind, frame in enumerate(all_frames):
-
-        # Get current set of cond latents of non-top levels.
-        cond_level, cond_level_latents = get_cond_latents(
-            self.all_level_latents, self.hparams)
-
-        # Get current set of cond latents of the top-level
-        cond_top, cond_top_latents = get_cond_latents(
-            self.all_top_latents, self.hparams)
-
-        # Superscript -> level, Subscript -> Time.
-        # (z^{0}_t, z^{1..l}_t) = Glow(X_{t}, z^{1..l}_{cond_t})
-        frame_obj, curr_top_latent, curr_level_latents = self.glow_encoder(
-            frame, condition=cond_level, cond_latents=cond_level_latents,
-            init=init)
-
-        # z^0_t ~ N(f(z^0_{t-1}))
-        # cond_top_latents is None when
-        # latent_dist_encoder is conv_net and frame_ind < num_cond_frames.
-        prior_dist = self.top_prior(
-            condition=cond_top, cond_latents=cond_top_latents)
-        prior_objective = tf.reduce_sum(
-            prior_dist.log_prob(curr_top_latent), axis=[1, 2, 3])
-        frame_obj += prior_objective
-
-        # Loss computation.
-        # Do not model the probabililty of the input frames by default.
-        # Consistent with other video models.
-        if (frame_ind > num_input_frames - 1 or self.hparams.model_input or
-            self.hparams.gen_mode == "unconditional"):
-          total_objective += frame_obj
-        self.all_level_latents.append(curr_level_latents)
-        self.all_top_latents.append(curr_top_latent)
-
-      # During prediction time, store z_sample ~ N(f(z_{num_input_frames}))
-      # to generate the first target frame.
-      if self.is_predicting:
-        # Get current set of cond_top_latents
-        cond_top, cond_top_latents = get_cond_latents(
-            self.all_top_latents, self.hparams)
-        prior_dist = self.top_prior(
-            condition=cond_top, cond_latents=cond_top_latents)
-        self.z_sample = prior_dist.sample()
-        self.all_top_latents.append(self.z_sample)
-
-      # Converts log-probability to bits-per-pixel.
-      hwc = np.prod(self.hparams.problem.frame_shape)
-      total_objective = -total_objective / (np.log(2) * hwc * total_frames)
-    return total_objective
-
-  def objective_tower(self, features, init=False):
-    input_frames, target_frames = features["inputs"], features["targets"]
-    self.cond_latents, self.top_state = None, None
-    self.all_level_latents, self.all_top_latents = [], []
-    self._all_conds = []
-    self.level_states = [None] * (self.hparams.n_levels - 1)
-    self.z_top_shape = self.get_z_top_shape(init=init)
-    num_input_frames = self.hparams.video_num_input_frames
-    latent_dist_encoder = self.hparams.latent_dist_encoder
-    num_cond_latents = self.hparams.num_cond_latents
-
-    exp_modes = ["conditional", "unconditional"]
-    if self.hparams.gen_mode not in exp_modes:
-      raise ValueError("Expected mode to be in %s, got %s" %
-                       (exp_modes, self.hparams.gen_mode))
-
-    # Error checks for conditional video generation.
-    if self.hparams.gen_mode == "conditional":
-      exp_latent_encoders = ["pointwise", "conv_net", "conv_lstm", "conv3d_net"]
-      if latent_dist_encoder not in exp_latent_encoders:
-        raise ValueError("Expected latent_dist_encoder is %s, got %s" %
-                         (exp_latent_encoders, latent_dist_encoder))
-      if (latent_dist_encoder == "pointwise" and num_cond_latents != 1):
-        raise ValueError("Expected num_cond_latents: 1, with 'pointwise' "
-                         "latent_dist_encoder, got %d" % num_cond_latents)
-      if (latent_dist_encoder == "conv_net" and
-          num_cond_latents > num_input_frames):
-        raise ValueError("Expected num_cond_latents <= %d, got %d" %
-                         (num_input_frames, num_cond_latents))
-      if (latent_dist_encoder == "pointwise" and
-          self.hparams.init_batch_size != self.hparams.batch_size):
-        raise ValueError("init_batch_size different from batch_size not "
-                         "supported for latent_dist_encoder=pointwise")
-    if self.hparams.gen_mode == "unconditional":
-      if self.hparams.num_train_frames != 1:
-        raise ValueError("Expected num_train_frames to be 1 when "
-                         "hparams.gen_mode is unconditional, got %d" %
-                         self.hparams.num_train_frames)
-      if self.hparams.video_num_input_frames != 1:
-        raise ValueError("Expected num_input_frames to be 1 when "
-                         "hparams.gen_mode is unconditional, got %d" %
-                         self.hparams.video_num_input_frames)
-    return self.video_objective_tower(input_frames, target_frames, init=init)
diff --git a/tensor2tensor/models/video/nfg_conv3d_test.py b/tensor2tensor/models/video/nfg_conv3d_test.py
deleted file mode 100644
index 689a61ce7..000000000
--- a/tensor2tensor/models/video/nfg_conv3d_test.py
+++ /dev/null
@@ -1,49 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Test when the latent-network encoder is a conv3d net."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import parameterized
-from tensor2tensor.models.video import nfg_test_utils
-import tensorflow as tf
-
-
-conv3d_net_hparams = (
-    ("conv3d_net", 2, 2, "conv3d_net", "conditional", -1, 3),
-    ("conv3d_dil", 2, 2, "conv3d_net", "conditional", -1, -1, False, True),)
-
-
-class NextFrameGlowConv3DTest(nfg_test_utils.NextFrameGlowTest,
-                              parameterized.TestCase):
-
-  @parameterized.named_parameters(*conv3d_net_hparams)
-  def testGlowTrainAndDecode(self, in_frames=1, out_frames=1,
-                             latent_dist_encoder="pointwise",
-                             gen_mode="conditional", pretrain_steps=-1,
-                             num_train_frames=-1, cond_first_frame=False,
-                             apply_dilations=False):
-    self.GlowTrainAndDecode(
-        in_frames=in_frames, out_frames=out_frames,
-        latent_dist_encoder=latent_dist_encoder, gen_mode=gen_mode,
-        pretrain_steps=pretrain_steps, num_train_frames=num_train_frames,
-        cond_first_frame=cond_first_frame, apply_dilations=apply_dilations)
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/models/video/nfg_conv_lstm_test.py b/tensor2tensor/models/video/nfg_conv_lstm_test.py
deleted file mode 100644
index beda25f0e..000000000
--- a/tensor2tensor/models/video/nfg_conv_lstm_test.py
+++ /dev/null
@@ -1,48 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Test when the latent-network encoder is a conv-lstm."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import parameterized
-from tensor2tensor.models.video import nfg_test_utils
-import tensorflow as tf
-
-
-conv_lstm_hparams = (
-    ("in_3_out_2_lstm", 2, 1, "conv_lstm", "conditional", -1),
-    ("lstm_pretrain", 2, 1, "conv_lstm", "conditional", 50000))
-
-
-class NextFrameGlowConv3DTest(nfg_test_utils.NextFrameGlowTest,
-                              parameterized.TestCase):
-
-  @parameterized.named_parameters(*conv_lstm_hparams)
-  def testGlowTrainAndDecode(self, in_frames=1, out_frames=1,
-                             latent_dist_encoder="pointwise",
-                             gen_mode="conditional", pretrain_steps=-1,
-                             num_train_frames=-1, cond_first_frame=False):
-    self.GlowTrainAndDecode(
-        in_frames=in_frames, out_frames=out_frames,
-        latent_dist_encoder=latent_dist_encoder, gen_mode=gen_mode,
-        pretrain_steps=pretrain_steps, num_train_frames=num_train_frames,
-        cond_first_frame=cond_first_frame)
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/models/video/nfg_conv_test.py b/tensor2tensor/models/video/nfg_conv_test.py
deleted file mode 100644
index 12eb657d7..000000000
--- a/tensor2tensor/models/video/nfg_conv_test.py
+++ /dev/null
@@ -1,44 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Test when the latent-network encoder is a 2-D conv."""
-
-from absl.testing import parameterized
-from tensor2tensor.models.video import nfg_test_utils
-import tensorflow as tf
-
-
-conv_net_hparams = (
-    ("in_3_out_2_conv", 3, 1, "conv_net", "conditional"),
-    ("conv_net_cond_first", 2, 2, "conv_net", "conditional", -1, 3, True),)
-
-
-class NextFrameGlowConvTest(nfg_test_utils.NextFrameGlowTest,
-                            parameterized.TestCase):
-
-  @parameterized.named_parameters(*conv_net_hparams)
-  def testGlowTrainAndDecode(self, in_frames=1, out_frames=1,
-                             latent_dist_encoder="pointwise",
-                             gen_mode="conditional", pretrain_steps=-1,
-                             num_train_frames=-1, cond_first_frame=False):
-    self.GlowTrainAndDecode(
-        in_frames=in_frames, out_frames=out_frames, gen_mode=gen_mode,
-        latent_dist_encoder=latent_dist_encoder,
-        pretrain_steps=pretrain_steps, num_train_frames=num_train_frames,
-        cond_first_frame=cond_first_frame)
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/models/video/nfg_test_utils.py b/tensor2tensor/models/video/nfg_test_utils.py
deleted file mode 100644
index 2b47cdfe2..000000000
--- a/tensor2tensor/models/video/nfg_test_utils.py
+++ /dev/null
@@ -1,184 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Testing utils for next_frame_glow."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-import tempfile
-import numpy as np
-from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
-from tensor2tensor.models.video import next_frame_glow
-from tensor2tensor.utils import registry
-import tensorflow as tf
-
-MODES = tf.estimator.ModeKeys
-
-
-# TODO(mechcoder): Refactor or merge tests with the other next_frame_tests when
-# this moves to a public version.
-def fill_hparams(hparams, in_frames, out_frames, gen_mode="conditional",
-                 latent_dist_encoder="pointwise", pretrain_steps=-1,
-                 num_train_frames=-1, cond_first_frame=False,
-                 apply_dilations=False):
-  """Set next_frame_glow hparams."""
-  hparams.latent_apply_dilations = apply_dilations
-  hparams.video_num_input_frames = in_frames
-  hparams.video_num_target_frames = out_frames
-  hparams.latent_dist_encoder = latent_dist_encoder
-  hparams.gen_mode = gen_mode
-  hparams.pretrain_steps = pretrain_steps
-  hparams.num_train_frames = num_train_frames
-  hparams.cond_first_frame = cond_first_frame
-  if latent_dist_encoder in ["conv_net", "conv3d_net"]:
-    hparams.num_cond_latents = in_frames
-  else:
-    hparams.num_cond_latents = 1
-  problem = registry.problem("video_stochastic_shapes10k")
-  p_hparams = problem.get_hparams(hparams)
-  hparams.problem = problem
-  hparams.problem_hparams = p_hparams
-  hparams.tiny_mode = True
-  hparams.reward_prediction = False
-  hparams.latent_architecture = "glow_resnet"
-  hparams.latent_encoder_depth = 2
-  hparams.latent_pre_output_channels = 32
-  if (hparams.gen_mode == "conditional" and
-      hparams.latent_dist_encoder == "pointwise"):
-    hparams.batch_size = 16
-    hparams.init_batch_size = 16
-  else:
-    hparams.batch_size = 16
-    hparams.init_batch_size = 32
-  hparams.affine_coupling_width = 32
-  hparams.depth = 5
-  hparams.n_levels = 2
-  return hparams
-
-
-def fill_infer_targets(x):
-  x["infer_targets"] = tf.identity(x["targets"])
-  return x
-
-
-def create_basic_features(hparams):
-  dataset = hparams.problem.dataset(MODES.TRAIN, hparams=hparams)
-  dataset = dataset.batch(hparams.batch_size)
-  dataset = dataset.map(fill_infer_targets)
-  return dataset.make_one_shot_iterator().get_next()
-
-
-class NextFrameGlowTest(tf.test.TestCase):
-  """Utils for testing next_frame_glow."""
-
-  def should_run_session(self, hparams):
-    # dilated conv-3d not available on CPU.
-    return tf.test.is_gpu_available() or not hparams.latent_apply_dilations
-
-  def checkAllConds(self, conds_array, num_total_frames, hparams):
-    if hparams.cond_first_frame:
-      self.assertEqual(conds_array, [True]*(num_total_frames + 1))
-    elif hparams.pretrain_steps > -1:
-      self.assertEqual(conds_array, [False]*num_total_frames)
-    elif hparams.latent_dist_encoder != "pointwise":
-      self.assertEqual(conds_array, [True]*num_total_frames)
-
-  def RunModel(self, model, train_op, hparams, features, num_frames,
-               model_path=None):
-    exp_num_frames = num_frames + int(hparams.cond_first_frame)
-    if hparams.gen_mode == "conditional":
-      self.assertLen(model.all_top_latents, exp_num_frames)
-      self.assertLen(model.all_level_latents, exp_num_frames)
-
-    with tf.Session() as session:
-
-      if model_path is not None:
-        saver = tf.train.Saver()
-
-      session.run(tf.global_variables_initializer())
-
-      # Run initialization.
-      init_op = tf.get_collection("glow_init_op")
-      session.run(init_op)
-
-      loss, top_conds = session.run([train_op["training"], model._all_conds])  # pylint: disable=protected-access
-      self.checkAllConds(top_conds, num_frames, hparams)
-
-      if model_path is not None:
-        saver.save(session, model_path)
-
-      # Check that one forward-propagation does not NaN, i.e
-      # initialization etc works as expected.
-      self.assertTrue(loss > 0.0 and loss < 10.0)
-
-  def GlowTrainAndDecode(self, in_frames=1, out_frames=1,
-                         latent_dist_encoder="pointwise",
-                         gen_mode="conditional", pretrain_steps=-1,
-                         num_train_frames=-1, cond_first_frame=False,
-                         apply_dilations=False):
-    """Test 1 forward pass and sampling gives reasonable results."""
-    if num_train_frames == -1:
-      total_frames = in_frames + out_frames
-    else:
-      total_frames = num_train_frames
-
-    curr_dir = tempfile.mkdtemp()
-    model_path = os.path.join(curr_dir, "model")
-
-    # Training pipeline
-    with tf.Graph().as_default():
-      hparams = next_frame_glow.next_frame_glow_hparams()
-      hparams = fill_hparams(hparams, in_frames, out_frames,
-                             gen_mode, latent_dist_encoder, pretrain_steps,
-                             num_train_frames, cond_first_frame,
-                             apply_dilations)
-      features = create_basic_features(hparams)
-      model = next_frame_glow.NextFrameGlow(hparams, MODES.TRAIN)
-      _, train_op = model(features)
-      if self.should_run_session(hparams):
-        self.RunModel(model, train_op, hparams, features, total_frames,
-                      model_path)
-
-    # Inference pipeline
-    with tf.Graph().as_default():
-      hparams = next_frame_glow.next_frame_glow_hparams()
-      if hparams.gen_mode == "unconditional":
-        hparams.video_num_target_frames = 1
-      hparams = fill_hparams(hparams, in_frames, out_frames,
-                             gen_mode, latent_dist_encoder, pretrain_steps,
-                             num_train_frames, cond_first_frame,
-                             apply_dilations)
-      features = create_basic_features(hparams)
-      model = next_frame_glow.NextFrameGlow(
-          hparams, tf.estimator.ModeKeys.PREDICT)
-      predictions = model.infer(features)
-      outputs = predictions["outputs"]
-      model_path = os.path.join(curr_dir, "model")
-
-      if self.should_run_session(hparams):
-        with tf.Session() as session:
-          saver = tf.train.Saver()
-          saver.restore(session, model_path)
-          outputs_np = session.run(outputs)
-          self.assertEqual(outputs_np.shape, (16, out_frames, 64, 64, 3))
-          self.assertTrue(np.all(outputs_np <= 255))
-          self.assertTrue(np.all(outputs_np >= 0))
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/models/video/nfg_uncond_test.py b/tensor2tensor/models/video/nfg_uncond_test.py
deleted file mode 100644
index e2ceb50d3..000000000
--- a/tensor2tensor/models/video/nfg_uncond_test.py
+++ /dev/null
@@ -1,46 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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 unconditional glow."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import parameterized
-from tensor2tensor.models.video import nfg_test_utils
-import tensorflow as tf
-
-uncond_hparams = (
-    ("in_1_out_1", 1, 1, "pointwise", "conditional"),
-    ("uncond", 1, 3, "pointwise", "unconditional", -1, 1),)
-
-
-class NfgUncondTest(nfg_test_utils.NextFrameGlowTest, parameterized.TestCase):
-
-  @parameterized.named_parameters(*uncond_hparams)
-  def testGlowTrainAndDecode(self, in_frames=1, out_frames=1,
-                             latent_dist_encoder="pointwise",
-                             gen_mode="conditional", pretrain_steps=-1,
-                             num_train_frames=-1, cond_first_frame=False):
-    self.GlowTrainAndDecode(
-        in_frames=in_frames, out_frames=out_frames,
-        latent_dist_encoder=latent_dist_encoder, gen_mode=gen_mode,
-        pretrain_steps=pretrain_steps, num_train_frames=num_train_frames,
-        cond_first_frame=cond_first_frame)
-
-
-if __name__ == "__main__":
-  tf.test.main()

From 10048203693aa083ddd1e09721b0024a04cce350 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 3 Dec 2018 22:12:22 -0800
Subject: [PATCH 1315/2720] flags.py -> (utils/flags.py)

PiperOrigin-RevId: 223919886
---
 tensor2tensor/bin/t2t_trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index c123933fd..e9af961c3 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -39,7 +39,7 @@
 flags = tf.flags
 FLAGS = flags.FLAGS
 
-# See flags.py for additional command-line flags.
+# See utils/flags.py for additional command-line flags.
 flags.DEFINE_string("t2t_usr_dir", None,
                     "Path to a Python module that will be imported. The "
                     "__init__.py file should include the necessary imports. "

From 93de6903ff250dfed51f6c59f69ed26d3b62ac23 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 3 Dec 2018 22:42:30 -0800
Subject: [PATCH 1316/2720] Allow vgg_ckpt_dir as a decode_hparam and make it
 None by default.

PiperOrigin-RevId: 223921829
---
 tensor2tensor/data_generators/video_utils.py | 4 ++--
 tensor2tensor/utils/decoding.py              | 4 ++++
 tensor2tensor/utils/video_metrics.py         | 3 ++-
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 3d6d072f7..685a8adf9 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -138,7 +138,7 @@ def display_video_hooks(hook_args):
 
   with tf.Graph().as_default():
     _, best_decodes = video_metrics.compute_video_metrics_from_predictions(
-        predictions)
+        predictions, decode_hparams=hook_args.decode_hparams)
 
   all_summaries = []
   # Displays decodes corresponding to the best/worst metric,
@@ -189,7 +189,7 @@ def summarize_video_metrics(hook_args):
   with metrics_graph.as_default():
     if predictions:
       metrics_results, _ = video_metrics.compute_video_metrics_from_predictions(
-          predictions)
+          predictions, decode_hparams=hook_args.decode_hparams)
     else:
       metrics_results, _ = video_metrics.compute_video_metrics_from_png_files(
           output_dirs, problem_name, hparams.video_num_target_frames,
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index b2e5362b4..3ff3893a7 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -80,6 +80,10 @@ def decode_hparams(overrides=""):
       # Maximum number of videos displayed.
       # Total number of videos are max_display_outputs * num_decodes
       max_display_outputs=10,
+      # Used in computation of VGG feature based video metrics.
+      # Set this to be the path to a trained VGG ckpt to output
+      # useful metrics.
+      vgg_ckpt_path=None,
       # Used for MLPerf compliance logging.
       mlperf_decode_step=0.0,
       mlperf_threshold=25.0,
diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index a5119a1af..caa84b774 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -215,12 +215,13 @@ def compute_all_metrics_statistics(all_results):
   return statistics, decode_inds
 
 
-def compute_video_metrics_from_predictions(predictions):
+def compute_video_metrics_from_predictions(predictions, decode_hparams):
   """Computes metrics from predictions.
 
   Args:
     predictions: list of list of dicts.
                  outer length: num_decodes, inner_length: num_samples
+    decode_hparams: Decode hparams. instance of tf.contrib.training.HParams.
   Returns:
     statistics: dict of Tensors, key being the metric with each Tensor
                 having the shape (num_samples, num_frames).

From 1eae64e8d50227d86da56be299e3226ae0ef2101 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Tue, 4 Dec 2018 10:37:31 -0800
Subject: [PATCH 1317/2720] Starting work on a new mesh_tensorflow Transformer
 implementation.  This implementation lives in the mesh_tensorflow library and
 does not depend on Tensor2Tensor.

In the new Transformer implementation, the different kinds of layers in the transformer are subclasses of TransformerLayer.  Model configurations contain lists of TransformerLayer instances (each containing its own hyperparameters).  Users can add custom layers by adding new subclasses, without touching the core library. We don't have a growing list of global hyperparameters, and we don't have a giant switch statement of layer types.

Add a new model "mtf_transformer2" in Tensor2Tensor to use this new implementation.  We will eventually deprecate the old "mtf_transformer" model.

Add a class mtf.VariableDType to encapsulate the different datatypes used for a variable: master_dtype, slice_dtype and activation_dtype, so as to avoid passing three arguments through many functions.

PiperOrigin-RevId: 224003970
---
 tensor2tensor/models/__init__.py         |   1 +
 tensor2tensor/models/mtf_transformer2.py | 392 +++++++++++++++++++++++
 2 files changed, 393 insertions(+)
 create mode 100644 tensor2tensor/models/mtf_transformer2.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 05cafbde7..cfdf3e1c2 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -29,6 +29,7 @@
 from tensor2tensor.models import mtf_image_transformer
 from tensor2tensor.models import mtf_resnet
 from tensor2tensor.models import mtf_transformer
+from tensor2tensor.models import mtf_transformer2
 from tensor2tensor.models import neural_gpu
 from tensor2tensor.models import resnet
 from tensor2tensor.models import revnet
diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
new file mode 100644
index 000000000..d90c65b14
--- /dev/null
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -0,0 +1,392 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Transformer model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import copy
+import mesh_tensorflow as mtf
+from mesh_tensorflow.transformer import moe
+from mesh_tensorflow.transformer import transformer
+from mesh_tensorflow.transformer import transformer_layers
+from tensor2tensor.layers import common_hparams
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
+from tensor2tensor.utils import mtf_model
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+@registry.register_model
+class MtfTransformer2(mtf_model.MtfModel):
+  """Transformer in mesh_tensorflow."""
+
+  @property
+  def batch_dims(self):
+    hparams = self._hparams
+    if hparams.outer_batch_size == 0:
+      return [mtf.Dimension("batch", hparams.batch_size)]
+    else:
+      if hparams.batch_size % hparams.outer_batch_size != 0:
+        raise ValueError(
+            "hparams.outer_batch_size must divide hparams.batch_size")
+      return [
+          mtf.Dimension("outer_batch", hparams.outer_batch_size),
+          mtf.Dimension("inner_batch",
+                        hparams.batch_size // hparams.outer_batch_size)]
+
+  @property
+  def variable_dtype(self):
+    return mtf.VariableDType(
+        tf.as_dtype(self._hparams.master_dtype),
+        tf.as_dtype(self._hparams.slice_dtype),
+        tf.as_dtype(self._hparams.activation_dtype))
+
+  @property
+  def length_dim(self):
+    return mtf.Dimension(
+        "length", self._hparams.length or self._hparams.max_length)
+
+  def _import_to_batch_by_length(self, x, name, mesh):
+    mtf_shape = mtf.Shape(self.batch_dims + [self.length_dim])
+    x = tf.reshape(x, mtf_shape.to_integer_list)
+    return mtf.import_fully_replicated(mesh, x, mtf_shape, name=name)
+
+  def _import_to_batch_by_decode_length(self, x, name, mesh):
+    mtf_shape = mtf.Shape(self.batch_dims + [self.length_dim])
+    x = tf.reshape(x, mtf_shape.to_integer_list)
+    return mtf.import_fully_replicated(mesh, x, mtf_shape, name=name)
+
+  def model(self):
+    hparams = self._hparams
+    if isinstance(hparams.layer_stack, transformer.LayerStack):
+      layer_stack = hparams.layer_stack
+    else:
+      # hparams.layer_stack is a function for creating a LayerStack
+      layer_stack = hparams.layer_stack(hparams)
+    return transformer.Transformer(
+        layer_stack=layer_stack,
+        d_model=hparams.d_model,
+        input_vocab_size=self._targets_vocab_size,
+        output_vocab_size=self._targets_vocab_size,
+        autoregressive=hparams.decoder_type == "autoregressive",
+        max_length=hparams.max_length)
+
+  def _mtf_model_fn(self, features, mesh):
+    self._original_features = features
+    features = copy.copy(features)
+    hparams = self._hparams
+    targets = tf.to_int32(features["targets"])
+    if len(targets.get_shape()) > 2:
+      tf.logging.info("targets = %s" % targets)
+      targets = tf.squeeze(targets, [2, 3])
+    # pad targets to max_length
+    def pad_to_length(x):
+      extra_length = self.length_dim.size - tf.shape(x)[1]
+      x = tf.pad(x, [[0, 0], [0, extra_length]])
+      x = tf.reshape(x, [hparams.batch_size, self.length_dim.size])
+      return x
+    targets = pad_to_length(targets)
+    targets = self._import_to_batch_by_length(targets, "targets", mesh)
+    for key in ["targets_segmentation", "targets_position",
+                "inputs_segmentation", "inputs_position"]:
+      if key in features:
+        features[key] = pad_to_length(features[key])
+    if hparams.decoder_type == "autoregressive":
+      shifted_targets = mtf.shift(
+          targets, offset=1, dim=self.length_dim, wrap=False)
+    else:
+      raise ValueError(
+          "unknown hparams.decoder_type = %s" % hparams.decoder_type)
+    model = self.model()
+    logits, loss = model.call_simple(
+        inputs=shifted_targets,
+        targets=targets,
+        compute_loss=True,
+        mode=hparams.mode,
+        variable_dtype=self.variable_dtype)
+    # mesh_shape=hparams.mesh_shape,
+    # layout=hparams.layout,
+    return logits, loss
+
+  def mtf_model_fn(self, features, mesh):
+    with tf.variable_scope("transformer"):
+      logits, loss = self._mtf_model_fn(features, mesh)
+      # combine batch dims
+      if len(self.batch_dims) > 1:
+        combined_batch_dim = mtf.Dimension(
+            self.batch_dims[0].name, mtf.Shape(self.batch_dims).size)
+        logits = mtf.reshape(
+            logits, [combined_batch_dim] + logits.shape.dims[-2:])
+      return logits, loss
+
+  @property
+  def _targets_vocab_size(self):
+    targets_vocab_size = self._problem_hparams.modality[
+        "targets"].top_dimensionality
+    targets_vocab_size += (-targets_vocab_size) % self._hparams.vocab_divisor
+    return targets_vocab_size
+
+  @property
+  def _inputs_vocab_size(self):
+    inputs_vocab_size = self._problem_hparams.modality[
+        "inputs"].top_dimensionality
+    inputs_vocab_size += (-inputs_vocab_size) % self._hparams.vocab_divisor
+    return inputs_vocab_size
+
+  def sample(self, features, mesh):
+    hparams = self._hparams
+    model = self.model()
+    # Prepare partial targets.
+    # In either features["inputs"] or features["targets"].
+    # We force the outputs to begin with these sequences.
+    partial_targets = features.get("inputs", None)
+    if partial_targets is None:
+      partial_targets = features.get("targets", None)
+    if partial_targets is not None:
+      partial_targets = common_layers.expand_squeeze_to_nd(partial_targets, 2)
+      partial_targets = tf.to_int32(partial_targets)
+      partial_targets_batch = tf.shape(partial_targets)[0]
+      partial_targets_length = tf.shape(partial_targets)[1]
+      partial_targets = tf.pad(
+          partial_targets, [[0, hparams.batch_size - partial_targets_batch],
+                            [0, self.length_dim.size - partial_targets_length]])
+      partial_targets = self._import_to_batch_by_length(
+          partial_targets, "partial_targets", mesh)
+      # strip EOS
+      partial_targets *= mtf.to_int32(mtf.not_equal(partial_targets, 1))
+
+    else:
+      ids_shape = mtf.Shape(self.batch_dims + [self.length_dim])
+      partial_targets = mtf.constant(mesh, 0, ids_shape, dtype=tf.int32)
+    if hparams.beam_size == 1:
+      pass
+    else:
+      raise NotImplementedError("not implemented")
+      # beam_dim = mtf.Dimension("beam", hparams.beam_size)
+      # ids_shape = mtf.Shape(self.batch_dims + [beam_dim, self.length_dim])
+
+    partial_targets = mtf.Print(partial_targets, [partial_targets],
+                                "Partial_Targets", summarize=1000)
+    return model.sample_autoregressive(
+        partial_targets,
+        temperature=hparams.sampling_temp,
+        variable_dtype=self.variable_dtype)
+
+
+@registry.register_hparams
+def mtf_transformer2_base():
+  """Set of hyperparameters."""
+  hparams = common_hparams.basic_params1()
+  hparams.no_data_parallelism = True
+  hparams.use_fixed_batch_size = True
+  hparams.add_hparam("mtf_mode", True)
+  hparams.batch_size = 4
+  hparams.max_length = 1024
+  hparams.add_hparam("d_model", 1024)
+  hparams.label_smoothing = 0.0
+  # 8-way model-parallelism
+  hparams.add_hparam("mesh_shape", "model:8")
+  hparams.add_hparam("layout", "batch:batch;vocab:model;d_ff:model;heads:model")
+
+  # hparams.layer_stack should be either a transformer.LayerStack or a function
+  # from hparams to transformer.LayerStack
+  def my_layer_stack(hparams):
+    return transformer.LayerStack(
+        [transformer_layers.SelfAttention(
+            num_heads=hparams.num_heads,
+            key_value_size=hparams.d_kv,
+            dropout_rate=hparams.attention_dropout),
+         transformer_layers.DenseReluDense(
+             hidden_size=hparams.d_ff,
+             dropout_rate=hparams.layer_prepostprocess_dropout),
+        ] * hparams.num_hidden_layers)
+  hparams.layer_stack = my_layer_stack
+
+  # These hyperparameters are used in the above default layer_stack function.
+  # They may not be respected if hparams.layer_stack is changed.
+  hparams.num_hidden_layers = 6
+  hparams.add_hparam("d_ff", 2048)
+  hparams.add_hparam("d_kv", 128)
+  hparams.add_hparam("attention_dropout", 0.0)
+  hparams.add_hparam("relu_dropout", 0.0)
+  hparams.layer_prepostprocess_dropout = 0.0
+
+  # Describes what model architecture:
+  #   "encdec": encoder + autoregressive decoder
+  #   "decoder": single-stack autoregressive sequence model.
+  #   "encoder": single-stack non-autoregressive model
+  #      with equal-length inputs and outputs.
+  # TODO(noam): implement different types of transformers.
+  hparams.add_hparam("transformer_type", "decoder")
+
+  # What does the decoder do:
+  #   "autoregressive": Decoder left to right
+  #   "denoising": Fills in masked-out values simultaneously
+  # TODO(noam): only autoregressive is implemented so far.
+  hparams.add_hparam("decoder_type", "autoregressive")
+
+  # round up vocab sizes to be a multiple of this value
+  hparams.vocab_divisor = 128
+
+  hparams.optimizer = "Adafactor"
+  hparams.learning_rate_schedule = "rsqrt_decay*linear_decay"
+  hparams.learning_rate_warmup_steps = 10000
+  hparams.add_hparam("master_dtype", "bfloat16")
+  hparams.add_hparam("slice_dtype", "float32")
+  hparams.activation_dtype = "bfloat16"
+
+  # These parameters make Transformer model compatible with MtfTransformer2
+  # Do not override these, as mtf_transformer does not support other options.
+  hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
+  hparams.modality = {
+      "inputs": modalities.IdentitySymbolModality,
+      "targets": modalities.IdentitySymbolModality,
+  }
+
+  # Parameters for computing the maximum decode length in beam search.
+  # Maximum decode length is:
+  #    min(max_length,
+  #        decode_length_multiplier * input_length + decode_length_constant)
+  hparams.add_hparam("decode_length_multiplier", 1.5)
+  hparams.add_hparam("decode_length_constant", 10.0)
+
+  # If nonzero, we split the batch across two tensor-dimensions named
+  # "outer_batch" and "inner_batch", allowing for splitting across two mesh
+  # dimensions.  This is necessary for hierarchical mixture of experts.
+  # The two tensor dimensions have sizes hparams.outer_batch_size and
+  # hparams.batch_size // hparams.outer_batch_size.
+  hparams.add_hparam("outer_batch_size", 0)
+
+  # length for training or decoding - defaults to max_length
+  hparams.add_hparam("length", 0)
+
+  hparams.sampling_method = "random"
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer2_tiny():
+  hparams = mtf_transformer2_base()
+  hparams.batch_size = 2
+  hparams.mesh_shape = ""
+  hparams.d_model = 128
+  hparams.num_hidden_layers = 2
+  hparams.num_heads = 4
+  hparams.d_ff = 512
+  return hparams
+
+
+@registry.register_hparams
+def mtf_transformer2_all_layers_tiny():
+  """Test out all the layers on local CPU."""
+  hparams = mtf_transformer2_base()
+  hparams.batch_size = 2
+  hparams.mesh_shape = ""
+  hparams.d_model = 128
+  hparams.layer_stack = transformer.LayerStack(
+      [transformer_layers.SelfAttention(num_heads=4),
+       transformer_layers.LocalSelfAttention(num_heads=4),
+       moe.MoE1D(num_experts=4, hidden_size=512),
+       moe.MoE2D(expert_x=4, expert_y=4, hidden_size=512),
+       transformer_layers.DenseReluDense(hidden_size=512)])
+  return hparams
+
+
+@registry.register_hparams
+def mtr2_lm_dense(sz):
+  """Series of architectural experiments on language modeling.
+
+  Larger models than the ones above.
+
+  All models are trained on sequences of 1024 tokens.
+
+  We assume infinite training data, so no dropout necessary.
+  We process 2^36 tokens in training = 524288 steps at batch size 128
+
+  TODO(noam): find a large enough dataset for these experiments.
+
+  You can use languagemodel_wiki_noref_v32k_l1k, but this is too small,
+  (1 epoch = ~46000 steps) so training will cover about 11 epochs.
+
+  Note: configurations and code are likely to change without notice.
+
+  Run on TPU 4x4 for 524288 steps unless otherwise indicated.
+
+  Args:
+    sz: an integer
+
+  Returns:
+    a hparams
+  """
+  n = 2 ** sz
+  hparams = mtf_transformer2_base()
+  hparams.d_model = 1024
+  hparams.max_length = 1024
+  hparams.batch_size = 128
+  # Parameters for my_layer_stack()
+  hparams.num_hidden_layers = 6
+  hparams.d_ff = 8192 * n
+  hparams.d_kv = 256
+  hparams.num_heads = 8 * n
+  hparams.learning_rate_decay_steps = 65536
+  hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
+  hparams.mesh_shape = "batch:32"
+  return hparams
+
+
+@registry.register_hparams
+def mtr2_lm_dense_0():
+  return mtr2_lm_dense(0)
+
+
+@registry.register_hparams
+def mtr2_lm_dense_1():
+  return mtr2_lm_dense(1)
+
+
+@registry.register_hparams
+def mtr2_lm_dense_2():
+  return mtr2_lm_dense(2)
+
+
+@registry.register_hparams
+def mtr2_v1():
+  """Model incorporating mixture-of-experts, local and global attention.
+
+  ~6B parameters
+
+  32 experts in 3 hierarchichal moe layers.
+
+  Returns:
+    a hparams
+  """
+  hparams = mtr2_lm_dense(0)
+  local_att = transformer_layers.LocalSelfAttention(
+      num_heads=4, key_value_size=128)
+  att = transformer_layers.SelfAttention(num_heads=4, key_value_size=128)
+  drd = transformer_layers.DenseReluDense(hidden_size=2048)
+  hmoe = moe.MoE2D(expert_x=8, expert_y=4, hidden_size=32768)
+  hparams.layer_stack = transformer.LayerStack(
+      ([local_att, local_att, drd,
+        att, drd, local_att, local_att, hmoe] * 4)[:-1])
+  hparams.mesh_shape = "b0:4;b1:8"
+  hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
+  hparams.outer_batch_size = 4
+  return hparams

From 87fcf7aaea6926bb167f5e1f6fe1d82a78fdcc92 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 4 Dec 2018 14:20:27 -0800
Subject: [PATCH 1318/2720] Readd nfg to open-source T2T for external
 collaboration.

PiperOrigin-RevId: 224045516
---
 oss_scripts/oss_tests.sh                      |   7 +-
 tensor2tensor/models/__init__.py              |   1 +
 tensor2tensor/models/video/next_frame_glow.py | 572 ++++++++++++++++++
 tensor2tensor/models/video/nfg_conv3d_test.py |  48 ++
 .../models/video/nfg_conv_lstm_test.py        |  47 ++
 tensor2tensor/models/video/nfg_conv_test.py   |  43 ++
 tensor2tensor/models/video/nfg_test_utils.py  | 183 ++++++
 tensor2tensor/models/video/nfg_uncond_test.py |  46 ++
 8 files changed, 946 insertions(+), 1 deletion(-)
 create mode 100644 tensor2tensor/models/video/next_frame_glow.py
 create mode 100644 tensor2tensor/models/video/nfg_conv3d_test.py
 create mode 100644 tensor2tensor/models/video/nfg_conv_lstm_test.py
 create mode 100644 tensor2tensor/models/video/nfg_conv_test.py
 create mode 100644 tensor2tensor/models/video/nfg_test_utils.py
 create mode 100644 tensor2tensor/models/video/nfg_uncond_test.py

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 31e321f01..ee9845714 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -69,7 +69,12 @@ if [[ "$TRAVIS_PYTHON_VERSION" == "2.7"  ]] && [[ "$TF_VERSION" == "tf-nightly"
 then
   # Ignores:
   # * Glow requires the CIFAR-10 dataset to be generated
-  pytest tensor2tensor/models/research --ignore=tensor2tensor/models/research/glow_test.py
+  pytest tensor2tensor/models/research \
+    --ignore=tensor2tensor/models/research/glow_test.py \
+    --ignore=tensor2tensor/models/video/nfg_conv3d_test.py \
+    --ignore=tensor2tensor/models/video/nfg_conv_lstm_test.py \
+    --ignore=tensor2tensor/models/video/nfg_conv_test.py \
+    --ignore=tensor2tensor/models/video/nfg_uncond_test.py \
   set_status
 fi
 
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index cfdf3e1c2..a19260891 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -68,6 +68,7 @@
 from tensor2tensor.models.video import basic_stochastic
 from tensor2tensor.models.video import emily
 from tensor2tensor.models.video import epva
+from tensor2tensor.models.video import next_frame_glow
 from tensor2tensor.models.video import savp
 from tensor2tensor.models.video import sv2p
 from tensor2tensor.models.video import svg_lp
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
new file mode 100644
index 000000000..dd1251946
--- /dev/null
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -0,0 +1,572 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Experimental testbed for nfg."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
+from tensor2tensor.layers import modalities
+from tensor2tensor.models.research import glow
+from tensor2tensor.models.research import glow_ops
+from tensor2tensor.utils import registry
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+
+arg_scope = tf.contrib.framework.arg_scope
+
+
+@registry.register_hparams
+def next_frame_glow_hparams():
+  """Hparams for next_frame_glow."""
+  hparams = glow.glow_hparams()
+  # Possible modes are conditional and unconditional
+  hparams.add_hparam("gen_mode", "conditional")
+  hparams.add_hparam("learn_top_scale", False)
+  hparams.add_hparam("condition_all_levels", True)
+  # For each video, substitutes "num_input_frames + num_output_frames" with a
+  # randomly sampled patch of length "num_train_frames" during training.
+  # -1 indicates that the entire video is used for training.
+  hparams.add_hparam("num_train_frames", -1)
+  # The following are hparams that model the latent transitions.
+  # Encoder that maps the latents to a Gaussian distribution.
+  # This function is used to model the prior over z_{t}. Can be,
+  # Pointwise -> point-wise multiplication of z_{t-1}.
+  # conv_net -> one-layer convolution over z_{t-1} .. z_{t - num_cond_latents}
+  hparams.add_hparam("latent_dist_encoder", "conv_net")
+  # Number of latents used in the encoder above.
+  hparams.add_hparam("num_cond_latents", 1)
+  hparams.add_hparam("latent_architecture", "glow_resnet")
+  hparams.add_hparam("latent_apply_dilations", False)
+  # Use latent skip connections
+  hparams.add_hparam("model_input", False)
+  hparams.add_hparam("cond_first_frame", False)
+  hparams.add_hparam("latent_skip", True)
+  hparams.add_hparam("latent_encoder_depth", 2)
+  hparams.add_hparam("latent_encoder_width", 512)
+  hparams.add_hparam("latent_pre_output_channels", 512)
+  # Pretrains the glow encoder for "pretrain_steps" number of steps.
+  # By default, don't pretrain and learn end-to-end
+  hparams.add_hparam("pretrain_steps", -1)
+  hparams.modality = {
+      "inputs": modalities.VideoModalityL1Raw,
+      "targets": modalities.VideoModalityL1Raw,
+  }
+  hparams.init_batch_size = 256
+  hparams.batch_size = 32
+  # Possible options: are prev_frame, single_conv and normal
+  hparams.top_prior = "single_conv"
+  return hparams
+
+
+@registry.register_hparams
+def frame_glow_hparams():
+  """Unconditional generation on video-frames."""
+  hparams = next_frame_glow_hparams()
+  hparams.gen_mode = "unconditional"
+  hparams.num_train_frames = 1
+  return hparams
+
+
+def get_cond_latents(all_latents=None, hparams=None):
+  """Get z^{cond}_{t} given z^{1..t-1}.
+
+  Args:
+    all_latents: list of list of tensors,
+                 outer-size equals no.of time_steps-1
+                 inner-size equals hparams.n_levels.
+    hparams: See next_frame_glow_hparams.
+  Returns:
+    cond_latents: conditional latents at time-step t.
+  """
+  cond_latents = None
+  if hparams.gen_mode == "conditional":
+    if hparams.latent_dist_encoder in ["conv_net", "conv3d_net"]:
+      num_cond_latents = (hparams.num_cond_latents +
+                          int(hparams.cond_first_frame))
+      if len(all_latents) >= num_cond_latents:
+        cond_latents = all_latents[-hparams.num_cond_latents:]
+        if hparams.cond_first_frame:
+          cond_latents = [all_latents[0]] + cond_latents
+    elif hparams.latent_dist_encoder in ["pointwise", "conv_lstm"]:
+      if all_latents:
+        cond_latents = all_latents[-1]
+
+  if hparams.gen_mode == "conditional":
+    global_step = tf.train.get_or_create_global_step()
+    condition = tf.greater(global_step, hparams.pretrain_steps)
+  else:
+    condition = tf.constant(False, dtype=tf.bool)
+  return condition, cond_latents
+
+
+@registry.register_model
+class NextFrameGlow(glow.Glow):
+  """Extend Glow for video."""
+
+  def init_preprocess_single(self, features):
+    for label in ["inputs", "targets"]:
+      features[label] = common_layers.convert_rgb_to_real(features[label])
+    return features
+
+  def init_preprocess(self, features):
+    """Preprocessing as per the input modality.
+
+    Equivalent to calling self.bottom(features).
+
+    Args:
+      features: dict of strings to tensors.
+    Returns:
+      features: dict of strings to tensors.
+    """
+    return features.map(self.init_preprocess_single)
+
+  def preprocess(self, x):
+    """Converts x from [0, 1] to [-0.5, 0.5].
+
+    All inputs are already normalized to be in the range [0, 1] through the
+    VideoModalityL1Raw modality.
+
+    Args:
+      x: 4-D Tensor.
+
+    Returns:
+      x: Scaled such that x lies in-between -0.5 and 0.5
+    """
+    return x - 0.5
+
+  def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
+    del args, kwargs
+
+    # Make a copy of features that can be used in the call to self
+    # that builds the graph.
+    new_features = {}
+    new_features["inputs"] = features["inputs"]
+    new_features["targets"] = features["infer_targets"]
+    _, _ = self(new_features)  # pylint: disable=not-callable
+
+    if self.hparams.gen_mode == "unconditional":
+      num_target_frames = 1
+    else:
+      num_target_frames = self.hparams.video_num_target_frames
+
+    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
+    var_scope = tf.variable_scope("next_frame_glow/body", reuse=True)
+    all_frames = []
+
+    # If eps=None, images are sampled from the prior.
+    with arg_scope(ops, init=False), var_scope:
+      for target_frame in range(1, num_target_frames + 1):
+
+        # subscript -> timestep, superscript -> level.
+        # self.z_sample equals z^0_{t} (top-level latent)
+        # (X_{t}, z^{1..l}_{t}) = Glow(z^0_{t}, z^{1..l}_{t-1})
+        # Get current set of cond_latents.
+        cond_level, cond_level_latents = get_cond_latents(
+            self.all_level_latents, self.hparams)
+
+        glow_vals = glow_ops.encoder_decoder(
+            "codec", self.z_sample, self.hparams, eps=None, reverse=True,
+            cond_latents=cond_level_latents, states=self.level_states,
+            condition=cond_level, temperature=self.temperature)
+        predicted_frame, _, curr_latents, self.level_states = glow_vals
+        all_frames.append(predicted_frame)
+        self.all_level_latents.append(curr_latents)
+
+        # Compute z^0_{t+1} = f(z^0_{t})
+        if target_frame < num_target_frames:
+          cond_top, cond_top_latents = get_cond_latents(
+              self.all_top_latents, self.hparams)
+          prior_dist = self.top_prior(
+              condition=cond_top, cond_latents=cond_top_latents)
+          self.z_sample = prior_dist.sample()
+          self.all_top_latents.append(self.z_sample)
+
+    all_frames = tf.stack(all_frames)
+    predicted_video = common_video.swap_time_and_batch_axes(all_frames)
+
+    # The video-decode API requires the predicted video to be the same shape
+    # as the target-video. Hence, for unconditional generation,
+    # tile across time to ensure same shape.
+    if self.hparams.gen_mode == "unconditional":
+      predicted_video = tf.tile(
+          predicted_video, [1, self.hparams.video_num_target_frames, 1, 1, 1])
+    predicted_video = self.scale(predicted_video)
+
+    # Output of a single decode / sample.
+    output_features = {}
+    output_features["targets"] = tf.zeros_like(predicted_video)
+    output_features["outputs"] = predicted_video
+    output_features["scores"] = tf.zeros_like(predicted_video)
+    return output_features
+
+  def get_squeeze_prior(self):
+    """Model the prior over z_{t} as a function of X_{t-1}.
+
+    Returns:
+      objective: float, log-likelihood.
+      dist: instance of tfp.distributions.Normal.
+
+    Raises:
+      ValueError: If input_height is not equal to input_width, not even
+                   or if the image width is smaller than the latent width.
+    """
+    _, prior_height, _, prior_channels = self.z_top_shape
+    _, input_height, input_width, _ = common_layers.shape_list(self.input_frame)
+
+    if input_height != input_width:
+      raise ValueError("input height should be equal to input width")
+    if input_height % 2 != 0:
+      raise ValueError("input height should be even")
+    if input_height < prior_height:
+      raise ValueError("input should be larger than the prior.")
+
+    # mean, log_std = NN(X_0)
+    # Reduce the spatial dimension by a factor of "squeeze_factor".
+    # and convolve with a stride of 2
+    squeeze_factor = input_height // (2 * prior_height)
+    x = glow_ops.squeeze(
+        "prior_squeeze", self.input_frame, factor=squeeze_factor, reverse=False)
+    mean_and_log_std = glow_ops.conv(
+        "prior_conv", x, 2*prior_channels, stride=[2, 2], apply_actnorm=False,
+        conv_init="zeros")
+    mean, log_scale = tf.split(mean_and_log_std, num_or_size_splits=2, axis=-1)
+    return tfp.distributions.Normal(mean, tf.exp(log_scale))
+
+  def top_cond_prior(self, name, cond_top_latents):
+    """Maps the conditional top latents to a distribution.
+
+    Args:
+      name: variable scope.
+      cond_top_latents: Tensor or a list of tensors.
+                        Latent variables at the previous time-step.
+                        If "pointwise", this is a single tensor.
+                        If "conv_net", this is a list of tensors with length
+                        equal to hparams.num_cond_latents.
+    Returns:
+      cond_dist: tfp.distributions.Normal
+    Raises:
+      ValueError: If cond_top_latents are not of the expected length.
+    """
+    with tf.variable_scope("top", reuse=tf.AUTO_REUSE):
+      if self.hparams.latent_dist_encoder == "pointwise":
+        last_latent = cond_top_latents
+        top = glow_ops.scale_gaussian_prior(
+            name, cond_top_latents, trainable=self.hparams.learn_top_scale)
+      elif self.hparams.latent_dist_encoder == "conv_net":
+        num_cond_latents = (self.hparams.num_cond_latents +
+                            int(self.hparams.cond_first_frame))
+        if len(cond_top_latents) != num_cond_latents:
+          raise ValueError(
+              "Expected length of cond_top_latents %d, got %d"
+              % (num_cond_latents, len(cond_top_latents)))
+        last_latent = cond_top_latents[-1]
+        output_channels = common_layers.shape_list(last_latent)[-1]
+        cond_top_latents = tf.concat(cond_top_latents, axis=-1)
+
+        # Maps the latent-stack to a distribution.
+        top = glow_ops.latent_to_dist(
+            name, cond_top_latents, hparams=self.hparams,
+            output_channels=output_channels)
+      elif self.hparams.latent_dist_encoder == "conv_lstm":
+        last_latent = cond_top_latents
+        output_channels = common_layers.shape_list(cond_top_latents)[-1]
+        # (h_t, c_t) = LSTM(z_{t-1}; (h_{t-1}, c_{t-1}))
+        # (mu_t, sigma_t) = conv(h_t)
+        _, self.top_state = common_video.conv_lstm_2d(
+            cond_top_latents, self.top_state, self.hparams.latent_encoder_width,
+            kernel_size=3, name="conv_lstm")
+        top = glow_ops.single_conv_dist(
+            name, self.top_state.h, output_channels=output_channels)
+      elif self.hparams.latent_dist_encoder == "conv3d_net":
+        last_latent = cond_top_latents[-1]
+        top = glow_ops.temporal_latent_to_dist(
+            "conv3d", tf.stack(cond_top_latents, axis=1), self.hparams)
+
+      # mu(z_{t}) = z_{t-1} + latent_encoder(z_{cond})
+      if self.hparams.latent_skip:
+        top = tfp.distributions.Normal(last_latent + top.loc, top.scale)
+    return top
+
+  def uncond_top_dist(self):
+    """Get an unconditional prior distribution on the top latent."""
+    prior_dist = glow_ops.top_prior(
+        "unconditional", self.z_top_shape, learn_prior="single_conv")
+    return prior_dist.loc, prior_dist.scale
+
+  def cond_top_dist(self, cond_latents):
+    """Get a conditional prior distribution on the top latent."""
+    prior_dist = self.top_cond_prior("conditional", cond_latents)
+    return prior_dist.loc, prior_dist.scale
+
+  def top_prior(self, condition=False, cond_latents=None):
+    """Objective based on the prior over latent z.
+
+    Args:
+      condition: Whether or not to condition on cond_latents.
+      cond_latents: tensor or list of tensors depending on
+                    hparams.latent_dist_encoder
+    Returns:
+      objective: float, log-likelihood of z under the prior.
+      dist: instance of tfp.distributions.Normal, prior distribution.
+    Raises:
+      ValueError: If input is smaller than the prior, uneven height
+                  or rectangular.
+    """
+    if isinstance(condition, bool):
+      condition = tf.constant(condition, dtype=tf.bool)
+    self._all_conds.append(condition)
+
+    if self.hparams.gen_mode == "conditional":
+      # cond_top_latents is None when
+      # latent_dist_encoder is a lstm and frame_ind == 0.
+      # latent_dist_encoder is conv_net and frame_ind < num_cond_frames.
+      marginal_mean, marginal_scale = self.uncond_top_dist()
+      if cond_latents is None:
+        mean, scale = marginal_mean, marginal_scale
+      else:
+        cond_mean, cond_scale = self.cond_top_dist(cond_latents)
+        mean, scale = tf.cond(
+            condition, lambda: (cond_mean, cond_scale),
+            lambda: (marginal_mean, marginal_scale))
+      return glow_ops.TemperedNormal(mean, scale, self.temperature)
+    if self.hparams.top_prior == "prev_frame":
+      return self.get_squeeze_prior()
+    else:
+      return super(NextFrameGlow, self).top_prior()
+
+  def get_z_top_shape(self, init=False):
+    """Get latent shape at level."""
+    if init:
+      batch_size = self.hparams.init_batch_size
+    else:
+      batch_size = self.hparams.batch_size
+    height, _, channels = self.hparams.problem.frame_shape
+    n_levels = self.hparams.n_levels
+    z_width = height // 2**n_levels
+    z_channels = channels * 2**n_levels * 2
+    return [batch_size, z_width, z_width, z_channels]
+
+  def squeeze_video(self, video, init=False):
+    """Squeeze a 5-D Tensor video with one timestep to a 4-D frame."""
+    if init:
+      batch_size = self.hparams.init_batch_size
+    else:
+      batch_size = self.hparams.batch_size
+    frame_shape = [batch_size] + self.hparams.problem.frame_shape
+    return tf.reshape(video, frame_shape)
+
+  def glow_encoder(self, frame, condition=False, cond_latents=None, init=False):
+    """Glow network that encodes frame to a hierarchy of latents.
+
+    Args:
+      frame: 5-D Tensor of shape (batch_size, 1, height, width, channels).
+      condition: Whether or not to condition on cond_latents.
+      cond_latents: optional, list of tensors with length equal to
+                    hparams.n_levels - 1. If provided, the latent at level l is
+                    conditioned on the cond_latent at level l.
+      init: Whether the given batch is an "init" batch or a "train" batch.
+    Returns:
+      objective: log-likelihood of the frame per the model.
+      z_top: top-level latent.
+      z_levels: a list of tensors with latents at all levels.
+    """
+    frame = self.squeeze_video(frame, init=init)
+    frame = self.preprocess(frame)
+    frame, objective = glow_ops.uniform_binning_correction(frame)
+
+    glow_vals = glow_ops.encoder_decoder(
+        "codec", frame, self.hparams, eps=None, reverse=False,
+        cond_latents=cond_latents, states=self.level_states,
+        condition=condition)
+    z_top, encoder_objective, self.eps, z_levels, self.level_states = glow_vals
+    objective += encoder_objective
+    return objective, z_top, z_levels
+
+  def get_num_train_frames(self):
+    """Returns the number of frames as a normalizing factor."""
+    num_target = self.hparams.video_num_target_frames
+    num_input = self.hparams.video_num_input_frames
+
+    # For unconditional generation, this picks a random frame during training
+    # and evaluates the marginal likelihood over "num_input" + "num_target"
+    # frames during eval.
+    if self.hparams.gen_mode == "unconditional":
+      if self.is_training:
+        return 1
+      return num_input + num_target
+
+    # During eval we measure the true objective.
+    if not self.is_training or self.hparams.num_train_frames == -1:
+      total_frames = num_target
+    # if hparams.num_train_frames=-1, we use an approxination to the true
+    # objective.
+    else:
+      total_frames = self.hparams.num_train_frames - num_input
+    if self.hparams.model_input:
+      total_frames += num_input
+    return total_frames
+
+  def get_all_frames(self, input_frames, target_frames):
+    """Get the frames used as input to the model.
+
+    Args:
+      input_frames: 5-D Tensor, (NTHWC)
+      target_frames: 5-D Tensor, (NTHWC)
+    Returns:
+      frames: 5-D Tensor used as input to the model.
+    """
+    if self.is_predicting:
+      all_frames = input_frames
+    elif self.is_training:
+      all_frames = tf.concat((input_frames, target_frames), axis=1)
+      all_frames = common_video.extract_random_video_patch(
+          all_frames, self.hparams.num_train_frames)
+    # Measure the mean bit-per-pixel of the target_frames during eval.
+    else:
+      all_frames = tf.concat((input_frames, target_frames), axis=1)
+    if self.hparams.cond_first_frame:
+      first_frame = all_frames[:, 0:1, :, :, :]
+      all_frames = tf.concat((first_frame, all_frames), axis=1)
+    return all_frames
+
+  def video_objective_tower(self, input_frames, target_frames, init=False):
+    """Returns the bits-per-pixel of the video.
+
+    Args:
+      input_frames: 5-D Tensor of shape (N, 1, H, W, C)
+      target_frames: 5-D Tensor of shape (N, T, H, W, C)
+      init: Whether or not to run data-dependent initialization.
+    Returns:
+      objective: bits-per-pixel.
+    """
+    # The arg_scope call ensures that the actnorm parameters are set such that
+    # the per-channel output activations have zero mean and unit variance
+    # ONLY during the first step. After that the parameters are learned
+    # through optimisation.
+    num_input_frames = (self.hparams.video_num_input_frames +
+                        int(self.hparams.cond_first_frame))
+
+    # Set num total frames to average the objective.
+    total_frames = self.get_num_train_frames()
+
+    # Compute the log-likelihood of target_frames at both train and predict
+    # time.
+    all_frames = self.get_all_frames(input_frames, target_frames)
+    all_frames = tf.unstack(all_frames, axis=1)
+
+    cond_level_latents, cond_top_latents = None, None
+    total_objective = 0.0
+    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
+
+    with arg_scope(ops, init=init):
+      for frame_ind, frame in enumerate(all_frames):
+
+        # Get current set of cond latents of non-top levels.
+        cond_level, cond_level_latents = get_cond_latents(
+            self.all_level_latents, self.hparams)
+
+        # Get current set of cond latents of the top-level
+        cond_top, cond_top_latents = get_cond_latents(
+            self.all_top_latents, self.hparams)
+
+        # Superscript -> level, Subscript -> Time.
+        # (z^{0}_t, z^{1..l}_t) = Glow(X_{t}, z^{1..l}_{cond_t})
+        frame_obj, curr_top_latent, curr_level_latents = self.glow_encoder(
+            frame, condition=cond_level, cond_latents=cond_level_latents,
+            init=init)
+
+        # z^0_t ~ N(f(z^0_{t-1}))
+        # cond_top_latents is None when
+        # latent_dist_encoder is conv_net and frame_ind < num_cond_frames.
+        prior_dist = self.top_prior(
+            condition=cond_top, cond_latents=cond_top_latents)
+        prior_objective = tf.reduce_sum(
+            prior_dist.log_prob(curr_top_latent), axis=[1, 2, 3])
+        frame_obj += prior_objective
+
+        # Loss computation.
+        # Do not model the probabililty of the input frames by default.
+        # Consistent with other video models.
+        if (frame_ind > num_input_frames - 1 or self.hparams.model_input or
+            self.hparams.gen_mode == "unconditional"):
+          total_objective += frame_obj
+        self.all_level_latents.append(curr_level_latents)
+        self.all_top_latents.append(curr_top_latent)
+
+      # During prediction time, store z_sample ~ N(f(z_{num_input_frames}))
+      # to generate the first target frame.
+      if self.is_predicting:
+        # Get current set of cond_top_latents
+        cond_top, cond_top_latents = get_cond_latents(
+            self.all_top_latents, self.hparams)
+        prior_dist = self.top_prior(
+            condition=cond_top, cond_latents=cond_top_latents)
+        self.z_sample = prior_dist.sample()
+        self.all_top_latents.append(self.z_sample)
+
+      # Converts log-probability to bits-per-pixel.
+      hwc = np.prod(self.hparams.problem.frame_shape)
+      total_objective = -total_objective / (np.log(2) * hwc * total_frames)
+    return total_objective
+
+  def objective_tower(self, features, init=False):
+    input_frames, target_frames = features["inputs"], features["targets"]
+    self.cond_latents, self.top_state = None, None
+    self.all_level_latents, self.all_top_latents = [], []
+    self._all_conds = []
+    self.level_states = [None] * (self.hparams.n_levels - 1)
+    self.z_top_shape = self.get_z_top_shape(init=init)
+    num_input_frames = self.hparams.video_num_input_frames
+    latent_dist_encoder = self.hparams.latent_dist_encoder
+    num_cond_latents = self.hparams.num_cond_latents
+
+    exp_modes = ["conditional", "unconditional"]
+    if self.hparams.gen_mode not in exp_modes:
+      raise ValueError("Expected mode to be in %s, got %s" %
+                       (exp_modes, self.hparams.gen_mode))
+
+    # Error checks for conditional video generation.
+    if self.hparams.gen_mode == "conditional":
+      exp_latent_encoders = ["pointwise", "conv_net", "conv_lstm", "conv3d_net"]
+      if latent_dist_encoder not in exp_latent_encoders:
+        raise ValueError("Expected latent_dist_encoder is %s, got %s" %
+                         (exp_latent_encoders, latent_dist_encoder))
+      if (latent_dist_encoder == "pointwise" and num_cond_latents != 1):
+        raise ValueError("Expected num_cond_latents: 1, with 'pointwise' "
+                         "latent_dist_encoder, got %d" % num_cond_latents)
+      if (latent_dist_encoder == "conv_net" and
+          num_cond_latents > num_input_frames):
+        raise ValueError("Expected num_cond_latents <= %d, got %d" %
+                         (num_input_frames, num_cond_latents))
+      if (latent_dist_encoder == "pointwise" and
+          self.hparams.init_batch_size != self.hparams.batch_size):
+        raise ValueError("init_batch_size different from batch_size not "
+                         "supported for latent_dist_encoder=pointwise")
+    if self.hparams.gen_mode == "unconditional":
+      if self.hparams.num_train_frames != 1:
+        raise ValueError("Expected num_train_frames to be 1 when "
+                         "hparams.gen_mode is unconditional, got %d" %
+                         self.hparams.num_train_frames)
+      if self.hparams.video_num_input_frames != 1:
+        raise ValueError("Expected num_input_frames to be 1 when "
+                         "hparams.gen_mode is unconditional, got %d" %
+                         self.hparams.video_num_input_frames)
+    return self.video_objective_tower(input_frames, target_frames, init=init)
diff --git a/tensor2tensor/models/video/nfg_conv3d_test.py b/tensor2tensor/models/video/nfg_conv3d_test.py
new file mode 100644
index 000000000..1c189808d
--- /dev/null
+++ b/tensor2tensor/models/video/nfg_conv3d_test.py
@@ -0,0 +1,48 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Test when the latent-network encoder is a conv3d net."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+from tensor2tensor.models.video import nfg_test_utils
+import tensorflow as tf
+
+conv3d_net_hparams = (
+    ("conv3d_net", 2, 2, "conv3d_net", "conditional", -1, 3),
+    ("conv3d_dil", 2, 2, "conv3d_net", "conditional", -1, -1, False, True),)
+
+
+class NextFrameGlowConv3DTest(nfg_test_utils.NextFrameGlowTest,
+                              parameterized.TestCase):
+
+  @parameterized.named_parameters(*conv3d_net_hparams)
+  def testGlowTrainAndDecode(self, in_frames=1, out_frames=1,
+                             latent_dist_encoder="pointwise",
+                             gen_mode="conditional", pretrain_steps=-1,
+                             num_train_frames=-1, cond_first_frame=False,
+                             apply_dilations=False):
+    self.GlowTrainAndDecode(
+        in_frames=in_frames, out_frames=out_frames,
+        latent_dist_encoder=latent_dist_encoder, gen_mode=gen_mode,
+        pretrain_steps=pretrain_steps, num_train_frames=num_train_frames,
+        cond_first_frame=cond_first_frame, apply_dilations=apply_dilations)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/video/nfg_conv_lstm_test.py b/tensor2tensor/models/video/nfg_conv_lstm_test.py
new file mode 100644
index 000000000..965ea8ff4
--- /dev/null
+++ b/tensor2tensor/models/video/nfg_conv_lstm_test.py
@@ -0,0 +1,47 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Test when the latent-network encoder is a conv-lstm."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+from tensor2tensor.models.video import nfg_test_utils
+import tensorflow as tf
+
+conv_lstm_hparams = (
+    ("in_3_out_2_lstm", 2, 1, "conv_lstm", "conditional", -1),
+    ("lstm_pretrain", 2, 1, "conv_lstm", "conditional", 50000))
+
+
+class NextFrameGlowConv3DTest(nfg_test_utils.NextFrameGlowTest,
+                              parameterized.TestCase):
+
+  @parameterized.named_parameters(*conv_lstm_hparams)
+  def testGlowTrainAndDecode(self, in_frames=1, out_frames=1,
+                             latent_dist_encoder="pointwise",
+                             gen_mode="conditional", pretrain_steps=-1,
+                             num_train_frames=-1, cond_first_frame=False):
+    self.GlowTrainAndDecode(
+        in_frames=in_frames, out_frames=out_frames,
+        latent_dist_encoder=latent_dist_encoder, gen_mode=gen_mode,
+        pretrain_steps=pretrain_steps, num_train_frames=num_train_frames,
+        cond_first_frame=cond_first_frame)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/video/nfg_conv_test.py b/tensor2tensor/models/video/nfg_conv_test.py
new file mode 100644
index 000000000..2ccc54e53
--- /dev/null
+++ b/tensor2tensor/models/video/nfg_conv_test.py
@@ -0,0 +1,43 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Test when the latent-network encoder is a 2-D conv."""
+
+from absl.testing import parameterized
+from tensor2tensor.models.video import nfg_test_utils
+import tensorflow as tf
+
+conv_net_hparams = (
+    ("in_3_out_2_conv", 3, 1, "conv_net", "conditional"),
+    ("conv_net_cond_first", 2, 2, "conv_net", "conditional", -1, 3, True),)
+
+
+class NextFrameGlowConvTest(nfg_test_utils.NextFrameGlowTest,
+                            parameterized.TestCase):
+
+  @parameterized.named_parameters(*conv_net_hparams)
+  def testGlowTrainAndDecode(self, in_frames=1, out_frames=1,
+                             latent_dist_encoder="pointwise",
+                             gen_mode="conditional", pretrain_steps=-1,
+                             num_train_frames=-1, cond_first_frame=False):
+    self.GlowTrainAndDecode(
+        in_frames=in_frames, out_frames=out_frames, gen_mode=gen_mode,
+        latent_dist_encoder=latent_dist_encoder,
+        pretrain_steps=pretrain_steps, num_train_frames=num_train_frames,
+        cond_first_frame=cond_first_frame)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/video/nfg_test_utils.py b/tensor2tensor/models/video/nfg_test_utils.py
new file mode 100644
index 000000000..242463368
--- /dev/null
+++ b/tensor2tensor/models/video/nfg_test_utils.py
@@ -0,0 +1,183 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Testing utils for next_frame_glow."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import tempfile
+import numpy as np
+from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
+from tensor2tensor.models.video import next_frame_glow
+from tensor2tensor.utils import registry
+import tensorflow as tf
+MODES = tf.estimator.ModeKeys
+
+
+# TODO(mechcoder): Refactor or merge tests with the other next_frame_tests when
+# this moves to a public version.
+def fill_hparams(hparams, in_frames, out_frames, gen_mode="conditional",
+                 latent_dist_encoder="pointwise", pretrain_steps=-1,
+                 num_train_frames=-1, cond_first_frame=False,
+                 apply_dilations=False):
+  """Set next_frame_glow hparams."""
+  hparams.latent_apply_dilations = apply_dilations
+  hparams.video_num_input_frames = in_frames
+  hparams.video_num_target_frames = out_frames
+  hparams.latent_dist_encoder = latent_dist_encoder
+  hparams.gen_mode = gen_mode
+  hparams.pretrain_steps = pretrain_steps
+  hparams.num_train_frames = num_train_frames
+  hparams.cond_first_frame = cond_first_frame
+  if latent_dist_encoder in ["conv_net", "conv3d_net"]:
+    hparams.num_cond_latents = in_frames
+  else:
+    hparams.num_cond_latents = 1
+  problem = registry.problem("video_stochastic_shapes10k")
+  p_hparams = problem.get_hparams(hparams)
+  hparams.problem = problem
+  hparams.problem_hparams = p_hparams
+  hparams.tiny_mode = True
+  hparams.reward_prediction = False
+  hparams.latent_architecture = "glow_resnet"
+  hparams.latent_encoder_depth = 2
+  hparams.latent_pre_output_channels = 32
+  if (hparams.gen_mode == "conditional" and
+      hparams.latent_dist_encoder == "pointwise"):
+    hparams.batch_size = 16
+    hparams.init_batch_size = 16
+  else:
+    hparams.batch_size = 16
+    hparams.init_batch_size = 32
+  hparams.affine_coupling_width = 32
+  hparams.depth = 5
+  hparams.n_levels = 2
+  return hparams
+
+
+def fill_infer_targets(x):
+  x["infer_targets"] = tf.identity(x["targets"])
+  return x
+
+
+def create_basic_features(hparams):
+  dataset = hparams.problem.dataset(MODES.TRAIN, hparams=hparams)
+  dataset = dataset.batch(hparams.batch_size)
+  dataset = dataset.map(fill_infer_targets)
+  return dataset.make_one_shot_iterator().get_next()
+
+
+class NextFrameGlowTest(tf.test.TestCase):
+  """Utils for testing next_frame_glow."""
+
+  def should_run_session(self, hparams):
+    # dilated conv-3d not available on CPU.
+    return tf.test.is_gpu_available() or not hparams.latent_apply_dilations
+
+  def checkAllConds(self, conds_array, num_total_frames, hparams):
+    if hparams.cond_first_frame:
+      self.assertEqual(conds_array, [True]*(num_total_frames + 1))
+    elif hparams.pretrain_steps > -1:
+      self.assertEqual(conds_array, [False]*num_total_frames)
+    elif hparams.latent_dist_encoder != "pointwise":
+      self.assertEqual(conds_array, [True]*num_total_frames)
+
+  def RunModel(self, model, train_op, hparams, features, num_frames,
+               model_path=None):
+    exp_num_frames = num_frames + int(hparams.cond_first_frame)
+    if hparams.gen_mode == "conditional":
+      self.assertLen(model.all_top_latents, exp_num_frames)
+      self.assertLen(model.all_level_latents, exp_num_frames)
+
+    with tf.Session() as session:
+
+      if model_path is not None:
+        saver = tf.train.Saver()
+
+      session.run(tf.global_variables_initializer())
+
+      # Run initialization.
+      init_op = tf.get_collection("glow_init_op")
+      session.run(init_op)
+
+      loss, top_conds = session.run([train_op["training"], model._all_conds])  # pylint: disable=protected-access
+      self.checkAllConds(top_conds, num_frames, hparams)
+
+      if model_path is not None:
+        saver.save(session, model_path)
+
+      # Check that one forward-propagation does not NaN, i.e
+      # initialization etc works as expected.
+      self.assertTrue(loss > 0.0 and loss < 10.0)
+
+  def GlowTrainAndDecode(self, in_frames=1, out_frames=1,
+                         latent_dist_encoder="pointwise",
+                         gen_mode="conditional", pretrain_steps=-1,
+                         num_train_frames=-1, cond_first_frame=False,
+                         apply_dilations=False):
+    """Test 1 forward pass and sampling gives reasonable results."""
+    if num_train_frames == -1:
+      total_frames = in_frames + out_frames
+    else:
+      total_frames = num_train_frames
+
+    curr_dir = tempfile.mkdtemp()
+    model_path = os.path.join(curr_dir, "model")
+
+    # Training pipeline
+    with tf.Graph().as_default():
+      hparams = next_frame_glow.next_frame_glow_hparams()
+      hparams = fill_hparams(hparams, in_frames, out_frames,
+                             gen_mode, latent_dist_encoder, pretrain_steps,
+                             num_train_frames, cond_first_frame,
+                             apply_dilations)
+      features = create_basic_features(hparams)
+      model = next_frame_glow.NextFrameGlow(hparams, MODES.TRAIN)
+      _, train_op = model(features)
+      if self.should_run_session(hparams):
+        self.RunModel(model, train_op, hparams, features, total_frames,
+                      model_path)
+
+    # Inference pipeline
+    with tf.Graph().as_default():
+      hparams = next_frame_glow.next_frame_glow_hparams()
+      if hparams.gen_mode == "unconditional":
+        hparams.video_num_target_frames = 1
+      hparams = fill_hparams(hparams, in_frames, out_frames,
+                             gen_mode, latent_dist_encoder, pretrain_steps,
+                             num_train_frames, cond_first_frame,
+                             apply_dilations)
+      features = create_basic_features(hparams)
+      model = next_frame_glow.NextFrameGlow(
+          hparams, tf.estimator.ModeKeys.PREDICT)
+      predictions = model.infer(features)
+      outputs = predictions["outputs"]
+      model_path = os.path.join(curr_dir, "model")
+
+      if self.should_run_session(hparams):
+        with tf.Session() as session:
+          saver = tf.train.Saver()
+          saver.restore(session, model_path)
+          outputs_np = session.run(outputs)
+          self.assertEqual(outputs_np.shape, (16, out_frames, 64, 64, 3))
+          self.assertTrue(np.all(outputs_np <= 255))
+          self.assertTrue(np.all(outputs_np >= 0))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/video/nfg_uncond_test.py b/tensor2tensor/models/video/nfg_uncond_test.py
new file mode 100644
index 000000000..e2ceb50d3
--- /dev/null
+++ b/tensor2tensor/models/video/nfg_uncond_test.py
@@ -0,0 +1,46 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 unconditional glow."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+from tensor2tensor.models.video import nfg_test_utils
+import tensorflow as tf
+
+uncond_hparams = (
+    ("in_1_out_1", 1, 1, "pointwise", "conditional"),
+    ("uncond", 1, 3, "pointwise", "unconditional", -1, 1),)
+
+
+class NfgUncondTest(nfg_test_utils.NextFrameGlowTest, parameterized.TestCase):
+
+  @parameterized.named_parameters(*uncond_hparams)
+  def testGlowTrainAndDecode(self, in_frames=1, out_frames=1,
+                             latent_dist_encoder="pointwise",
+                             gen_mode="conditional", pretrain_steps=-1,
+                             num_train_frames=-1, cond_first_frame=False):
+    self.GlowTrainAndDecode(
+        in_frames=in_frames, out_frames=out_frames,
+        latent_dist_encoder=latent_dist_encoder, gen_mode=gen_mode,
+        pretrain_steps=pretrain_steps, num_train_frames=num_train_frames,
+        cond_first_frame=cond_first_frame)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 25595156e60e5594fa51ba7397eb1a26a4ac8772 Mon Sep 17 00:00:00 2001
From: David Dohan <ddohan@google.com>
Date: Tue, 4 Dec 2018 14:50:57 -0800
Subject: [PATCH 1319/2720] add explicit taskid check for all
 multiproblem_weight functions

PiperOrigin-RevId: 224050814
---
 tensor2tensor/layers/common_layers.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 31634bd62..9cd9b4e11 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1588,6 +1588,16 @@ def weights_prepend_inputs_to_targets(labels):
   return tf.to_float(tf.not_equal(past_first_zero * nonzero, 0))
 
 
+def check_nonnegative(value):
+  """Check that the value is nonnegative."""
+  if isinstance(value, tf.Tensor):
+    with tf.control_dependencies([tf.assert_greater_equal(value, 0)]):
+      value = tf.identity(value)
+  elif value < 0:
+    raise ValueError("Value must be non-negative.")
+  return value
+
+
 def weights_multi_problem(labels, taskid=-1):
   """Assign weight 1.0 to only the "targets" portion of the labels.
 
@@ -1603,6 +1613,7 @@ def weights_multi_problem(labels, taskid=-1):
   Raises:
     ValueError: The Task ID must be valid.
   """
+  taskid = check_nonnegative(taskid)
   past_taskid = tf.cumsum(tf.to_float(tf.equal(labels, taskid)), axis=1)
   # Additionally zero out the task id location
   past_taskid *= tf.to_float(tf.not_equal(labels, taskid))
@@ -1612,6 +1623,7 @@ def weights_multi_problem(labels, taskid=-1):
 
 def weights_multi_problem_all(labels, taskid=-1):
   """Assign weight 1.0 to only examples from the given task."""
+  taskid = check_nonnegative(taskid)
   weights = tf.to_float(tf.not_equal(labels, 0))
   past_taskid = tf.cumsum(tf.to_float(tf.equal(labels, taskid)), axis=1)
   # Additionally zero out the task id location
@@ -1627,6 +1639,7 @@ def weights_multi_problem_all(labels, taskid=-1):
 
 def weights_multi_problem_input(labels, taskid=-1):
   """Assign weight 1.0 to only the inputs for the given task."""
+  taskid = check_nonnegative(taskid)
   weights_all_tokens = weights_multi_problem_all(labels, taskid)
   weights_target = weights_multi_problem(labels, taskid)
   return weights_all_tokens - weights_target

From d255404a4f6ddffd624f0b37f342ede8615d27b5 Mon Sep 17 00:00:00 2001
From: David Dohan <ddohan@google.com>
Date: Tue, 4 Dec 2018 14:57:47 -0800
Subject: [PATCH 1320/2720] Add support for multiproblem to mtf_model

PiperOrigin-RevId: 224051889
---
 tensor2tensor/utils/mtf_model.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/mtf_model.py b/tensor2tensor/utils/mtf_model.py
index 6ff50e4e2..fb829f7fe 100644
--- a/tensor2tensor/utils/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -179,7 +179,13 @@ def estimator_spec_eval(
     problem = hparams.problem
     if logits.get_shape().ndims == 3:
       logits = tf.expand_dims(tf.expand_dims(logits, 2), 3)
-    eval_metrics_fns = metrics.create_evaluation_metrics([problem], hparams)
+
+    # Support for multiproblem
+    task_list = [problem]
+    if hasattr(problem, "task_list"):
+      task_list = problem.task_list
+
+    eval_metrics_fns = metrics.create_evaluation_metrics(task_list, hparams)
 
     if use_tpu:
       def metric_fn(tf_logits, labels):

From fd32fa3dfc6422a3c6744e492799b176d5f05579 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 4 Dec 2018 15:27:57 -0800
Subject: [PATCH 1321/2720] Rm evaluation_hooks summarysaver. This shouldn't be
 here. Directory may be eval_train. Estimator is already logging summaries for
 eval.

PiperOrigin-RevId: 224057105
---
 tensor2tensor/utils/t2t_model.py | 8 --------
 1 file changed, 8 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 40f96e184..353a193ca 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1497,18 +1497,10 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
       else:
         predictions = {"predictions": logits}
 
-      evaluation_hooks = []
-      # Create a SummarySaverHook
-      eval_summary_hook = tf.train.SummarySaverHook(
-          save_steps=1, output_dir=self.hparams.model_dir + "/eval",
-          summary_op=tf.summary.merge_all())
-      evaluation_hooks.append(eval_summary_hook)
-
       return tf.estimator.EstimatorSpec(
           tf.estimator.ModeKeys.EVAL,
           predictions=predictions,
           eval_metric_ops=eval_metrics,
-          evaluation_hooks=evaluation_hooks,
           loss=loss)
 
   def estimator_spec_predict(self, features, use_tpu=False):

From edd7fb2cdb49b9726f2ab2d7ea655a28677c34b6 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 4 Dec 2018 15:35:25 -0800
Subject: [PATCH 1322/2720] Small En-Ro dataset for low-resource experiments,
 use that one in translation multi-problem; and a few corrections.

PiperOrigin-RevId: 224058409
---
 .../data_generators/multi_problem.py          | 11 ++++---
 .../data_generators/translate_enfr.py         |  9 +++---
 .../data_generators/translate_enro.py         | 30 +++++++++++++++++++
 .../data_generators/wiki_multi_problems.py    |  4 +--
 tensor2tensor/models/transformer.py           |  2 ++
 5 files changed, 43 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index b31000ed2..4fb5f6f71 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -118,6 +118,7 @@ def get_hparams(self, model_hparams=None):
     if model_hparams.multiproblem_vocab_size > new_vocab_size:
       new_vocab_size = model_hparams.multiproblem_vocab_size
     tf.logging.info("Old vocabulary size: %d" % vocab_size)
+    self.update_task_ids(vocab_size)
     tf.logging.info("New vocabulary size: %d" % new_vocab_size)
     self._hparams.vocab_size["targets"] = new_vocab_size
     self._hparams.modality["targets"] = modalities.SymbolModality(
@@ -172,7 +173,7 @@ def dataset(self,
       raise ValueError("Only support language models as primary problem which "
                        "supplies the vocabulary and the hparams.")
     enc = primary_task.feature_encoders(data_dir=data_dir)["targets"]
-    self.update_task_ids(enc)
+    self.update_task_ids(enc.vocab_size)
 
     for task in self.task_list:
       task_dataset = task.dataset(mode=mode,
@@ -331,19 +332,17 @@ def eval_metrics(self):
         metrics.Metrics.ACC, metrics.Metrics.NEG_LOG_PERPLEXITY,
     ]
 
-  def update_task_ids(self, encoder):
+  def update_task_ids(self, encoder_vocab_size):
     """Generate task_ids for each problem.
 
     These ids correspond to the index of the task in the task_list.
 
     Args:
-      encoder: this provides the size of the vocab which is used to compute
+      encoder_vocab_size: the size of the vocab which is used to compute
         the index offset.
     """
-    offset = encoder.vocab_size
-
     for idx, task in enumerate(self.task_list):
-      task.set_task_id(idx + offset)
+      task.set_task_id(idx + encoder_vocab_size)
       tf.logging.info("Task %d (%s) has id %d." %
                       (idx, task.name, task.task_id))
 
diff --git a/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index 7c336fca1..95aa15d20 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -180,12 +180,10 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
     tag = "train" if dataset_split == problem.DatasetSplit.TRAIN else "dev"
     data_path = translate.compile_data(
         tmp_dir, datasets, "%s-compiled-%s" % (self.name, tag))
-    # Iterator over authentic data.
-    it_auth = text_problems.text2text_txt_iterator(
-        data_path + ".lang1", data_path + ".lang2")
     # For eval, use authentic data.
     if dataset_split != problem.DatasetSplit.TRAIN:
-      for example in it_auth:
+      for example in text_problems.text2text_txt_iterator(
+          data_path + ".lang1", data_path + ".lang2"):
         yield example
     else:  # For training, mix synthetic and authentic data as follows.
       for (file1, file2) in self.backtranslate_data_filenames:
@@ -195,7 +193,8 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
         for example in text_problems.text2text_txt_iterator(path1, path2):
           yield example
         # Now authentic data.
-        for example in it_auth:
+        for example in text_problems.text2text_txt_iterator(
+            data_path + ".lang1", data_path + ".lang2"):
           yield example
 
 
diff --git a/tensor2tensor/data_generators/translate_enro.py b/tensor2tensor/data_generators/translate_enro.py
index 353c5d8c4..ac734189f 100644
--- a/tensor2tensor/data_generators/translate_enro.py
+++ b/tensor2tensor/data_generators/translate_enro.py
@@ -78,3 +78,33 @@ class TranslateEnroWmtMulti64k(TranslateEnroWmt8k):
   @property
   def vocab_filename(self):
     return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
+
+
+@registry.register_problem
+class TranslateEnroWmtMultiSmall64k(TranslateEnroWmt8k):
+  """Translation with muli-lingual vocabulary, small (6K) training data."""
+
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each."""
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 1,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    """Generate just the first 6k samples for training."""
+    counter = 0
+    for x in super(TranslateEnroWmtMultiSmall64k, self).generate_samples(
+        data_dir, tmp_dir, dataset_split):
+      counter += 1
+      if counter > 6000 and dataset_split == problem.DatasetSplit.TRAIN:
+        raise StopIteration
+      yield x
diff --git a/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index 5c8515fae..fafe6d074 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -133,12 +133,12 @@ def __init__(self, was_reversed=False, was_copy=False):
     self.task_list.append(wiki_lm.LanguagemodelDeEnFrRoWiki64k())
     self.task_list.append(translate_ende.TranslateEndeWmtMulti64k())
     self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k())
-    self.task_list.append(translate_enro.TranslateEnroWmtMulti64k())
+    self.task_list.append(translate_enro.TranslateEnroWmtMultiSmall64k())
     self.task_list.append(translate_ende.TranslateEndeWmtMulti64k(
         was_reversed=True))
     self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k(
         was_reversed=True))
-    self.task_list.append(translate_enro.TranslateEnroWmtMulti64k(
+    self.task_list.append(translate_enro.TranslateEnroWmtMultiSmall64k(
         was_reversed=True))
 
   @property
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 2749ac1a3..e4dc0fd01 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1618,6 +1618,8 @@ def transformer_tall_finetune_cnndm():
   # Set train steps to learning_rate_decay_steps or less
   hparams.learning_rate_decay_steps = 40000
   hparams.multiproblem_target_eval_only = True
+  hparams.multiproblem_vocab_size = 2**16
+  return hparams
 
 
 @registry.register_hparams

From 6d572c35ee15491e6349cf58673544a9a2aa625f Mon Sep 17 00:00:00 2001
From: Seppo Enarvi <seppo.git@marjaniemi.com>
Date: Wed, 5 Dec 2018 01:57:46 +0100
Subject: [PATCH 1323/2720] Create an integer problem_0_steps variable. (#1273)

---
 tensor2tensor/bin/t2t_avg_all.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_avg_all.py b/tensor2tensor/bin/t2t_avg_all.py
index 3755dc2d0..080670f35 100644
--- a/tensor2tensor/bin/t2t_avg_all.py
+++ b/tensor2tensor/bin/t2t_avg_all.py
@@ -63,7 +63,8 @@ def main(_):
       var_list = tf.contrib.framework.list_variables(model.filename)
       avg_values = {}
       for (name, shape) in var_list:
-        if not name.startswith("global_step"):
+        if not (name.startswith("global_step") or
+                name.startswith("train_stats/")):
           avg_values[name] = np.zeros(shape)
     models_processed += 1
 
@@ -88,6 +89,8 @@ def main(_):
         "global_step",
         initializer=tf.constant(model.steps, dtype=tf.int64),
         trainable=False)
+    with tf.variable_scope("train_stats"):
+      tf.get_variable("problem_0_steps", initializer=0, trainable=False)
     saver = tf.train.Saver(tf.global_variables())
 
     tf.logging.info("Running session for %s" % (out_file))

From c6943fc0856acb9008df37b5f8cd28657fc1f33c Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 4 Dec 2018 17:17:22 -0800
Subject: [PATCH 1324/2720] Re-enable eval summaries in the correct directory

PiperOrigin-RevId: 224073967
---
 tensor2tensor/utils/t2t_model.py   | 13 +++++++++++++
 tensor2tensor/utils/trainer_lib.py | 22 ++++++++++++++++++----
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 353a193ca..15f19ca5e 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -24,6 +24,7 @@
 import functools
 import inspect
 import math
+import os
 import time
 import six
 
@@ -1497,10 +1498,22 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
       else:
         predictions = {"predictions": logits}
 
+      evaluation_hooks = []
+      # Create a SummarySaverHook
+      eval_dir = os.path.join(
+          self.hparams.model_dir,
+          self.hparams.get("eval_dir_name", "eval"))
+      eval_summary_hook = tf.train.SummarySaverHook(
+          save_steps=1,
+          output_dir=eval_dir,
+          summary_op=tf.summary.merge_all())
+      evaluation_hooks.append(eval_summary_hook)
+
       return tf.estimator.EstimatorSpec(
           tf.estimator.ModeKeys.EVAL,
           predictions=predictions,
           eval_metric_ops=eval_metrics,
+          evaluation_hooks=evaluation_hooks,
           loss=loss)
 
   def estimator_spec_predict(self, features, use_tpu=False):
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 4849bbf43..4e448b5b9 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -470,10 +470,12 @@ def train_eval_and_decode(self):
           self._train_spec.input_fn,
           steps=eval_steps,
           hooks=self._train_spec.hooks)
+      self._set_eval_dir_name("eval")
       self._estimator.evaluate(
           self._eval_spec.input_fn,
           steps=self._eval_spec.steps,
-          hooks=self._eval_spec.hooks)
+          hooks=self._eval_spec.hooks,
+          name="eval")
       if packed_dataset:
         problem = registry.problem(
             self._hparams.problem.name.replace("_packed", ""))
@@ -495,18 +497,30 @@ def train_eval_and_decode(self):
       mlperf_log.transformer_print(
           key=mlperf_log.RUN_STOP, value={"success": "false"})
 
+  def _set_eval_dir_name(self, eval_dir_name):
+    attr = "eval_dir_name"
+    hp = self._hparams
+    if attr not in hp:
+      hp.add_hparam(attr, "")
+    hp.eval_dir_name = eval_dir_name
+
   def evaluate(self):
+    name = "eval"
+    self._set_eval_dir_name("eval")
     return self._estimator.evaluate(
         self._eval_spec.input_fn,
         steps=self._eval_spec.steps,
-        hooks=self._eval_spec.hooks)
+        hooks=self._eval_spec.hooks,
+        name=name)
 
   def evaluate_on_train_data(self):
+    name = "eval_train"
+    self._set_eval_dir_name(name)
     self._estimator.evaluate(
         self._train_spec.input_fn,
         steps=self._eval_spec.steps,
         hooks=self._eval_spec.hooks,
-        name="eval_train")
+        name=name)
 
   def continuous_eval(self):
     """Evaluate until checkpoints stop being produced."""
@@ -521,7 +535,7 @@ def continuous_eval_on_train_data(self):
       self.evaluate_on_train_data()
 
   def test(self):
-    """Perform 1 step of train and 2 step of eval."""
+    """Perform 1 train step and 1 eval step."""
     if self._use_validation_monitor:
       return self.train_and_evaluate()
 

From f0e4deab6ebb130b58620b5a77e8c5001d11b8c7 Mon Sep 17 00:00:00 2001
From: blazejosinski <804945+blazejosinski@users.noreply.github.com>
Date: Wed, 5 Dec 2018 02:22:16 +0100
Subject: [PATCH 1325/2720] Improved model free eval. (#1269)

* Improved model free eval.

* Missing import.
---
 tensor2tensor/models/research/rl.py     |  13 +++
 tensor2tensor/rl/rl_utils.py            | 125 ++++++++++++++++++++++++
 tensor2tensor/rl/trainer_model_based.py | 109 ++-------------------
 tensor2tensor/rl/trainer_model_free.py  |  63 ++++++------
 4 files changed, 180 insertions(+), 130 deletions(-)
 create mode 100644 tensor2tensor/rl/rl_utils.py

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 1d6b5f52f..5c52f3db0 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -233,6 +233,19 @@ def mfrl_base():
   hparams = mfrl_original()
   hparams.add_hparam("ppo_epochs_num", 3000)
   hparams.add_hparam("ppo_eval_every_epochs", 100)
+  hparams.add_hparam("eval_max_num_noops", 8)
+  hparams.add_hparam("resize_height_factor", 2)
+  hparams.add_hparam("resize_width_factor", 2)
+  hparams.add_hparam("grayscale", 1)
+  hparams.add_hparam("env_timesteps_limit", -1)
+  return hparams
+
+
+@registry.register_hparams
+def mfrl_tiny():
+  hparams = mfrl_base()
+  hparams.ppo_epochs_num = 100
+  hparams.ppo_eval_every_epochs = 10
   return hparams
 
 
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
new file mode 100644
index 000000000..cff0ae440
--- /dev/null
+++ b/tensor2tensor/rl/rl_utils.py
@@ -0,0 +1,125 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+r"""Utilities for RL training
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import six
+
+from tensor2tensor.data_generators.gym_env import T2TGymEnv
+from tensor2tensor.models.research import rl
+from tensor2tensor.rl.dopamine_connector import DQNLearner
+from tensor2tensor.rl.ppo_learner import PPOLearner
+from tensor2tensor.utils import trainer_lib
+
+import tensorflow as tf
+
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+
+def compute_mean_reward(rollouts, clipped):
+  """Calculate mean rewards from given epoch."""
+  reward_name = "reward" if clipped else "unclipped_reward"
+  rewards = []
+  for rollout in rollouts:
+    if rollout[-1].done:
+      rollout_reward = sum(getattr(frame, reward_name) for frame in rollout)
+      rewards.append(rollout_reward)
+  if rewards:
+    mean_rewards = np.mean(rewards)
+  else:
+    mean_rewards = 0
+  return mean_rewards
+
+
+def get_metric_name(stochastic, max_num_noops, clipped):
+  return "mean_reward/eval/stochastic_{}_max_noops_{}_{}".format(
+      stochastic, max_num_noops, "clipped" if clipped else "unclipped")
+
+
+def evaluate_single_config(hparams, stochastic, max_num_noops,
+                           agent_model_dir):
+  """Evaluate the PPO agent in the real environment."""
+  eval_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
+  env = setup_env(
+      hparams, batch_size=hparams.eval_batch_size, max_num_noops=max_num_noops
+  )
+  env.start_new_epoch(0)
+  env_fn = rl.make_real_env_fn(env)
+  learner = LEARNERS[hparams.base_algo](
+      hparams.frame_stack_size, base_event_dir=None,
+      agent_model_dir=agent_model_dir
+  )
+  learner.evaluate(env_fn, eval_hparams, stochastic)
+  rollouts = env.current_epoch_rollouts()
+  env.close()
+
+  return tuple(
+      compute_mean_reward(rollouts, clipped) for clipped in (True, False)
+  )
+
+
+def evaluate_all_configs(hparams, agent_model_dir):
+  """Evaluate the agent with multiple eval configurations."""
+  metrics = {}
+  # Iterate over all combinations of picking actions by sampling/mode and
+  # whether to do initial no-ops.
+  for stochastic in (True, False):
+    for max_num_noops in (hparams.eval_max_num_noops, 0):
+      scores = evaluate_single_config(
+          hparams, stochastic, max_num_noops, agent_model_dir
+      )
+      for (score, clipped) in zip(scores, (True, False)):
+        metric_name = get_metric_name(stochastic, max_num_noops, clipped)
+        metrics[metric_name] = score
+
+  return metrics
+
+
+LEARNERS = {
+    "ppo": PPOLearner,
+    "dqn": DQNLearner,
+}
+
+
+def setup_env(hparams, batch_size, max_num_noops):
+  """Setup."""
+  game_mode = "Deterministic-v4"
+  camel_game_name = "".join(
+      [w[0].upper() + w[1:] for w in hparams.game.split("_")])
+  camel_game_name += game_mode
+  env_name = camel_game_name
+
+  env = T2TGymEnv(base_env_name=env_name,
+                  batch_size=batch_size,
+                  grayscale=hparams.grayscale,
+                  resize_width_factor=hparams.resize_width_factor,
+                  resize_height_factor=hparams.resize_height_factor,
+                  base_env_timesteps_limit=hparams.env_timesteps_limit,
+                  max_num_noops=max_num_noops)
+  return env
+
+def update_hparams_from_hparams(target_hparams, source_hparams, prefix):
+  """Copy a subset of hparams to target_hparams."""
+  for (param_name, param_value) in six.iteritems(source_hparams.values()):
+    if param_name.startswith(prefix):
+      target_hparams.set_hparam(param_name[len(prefix):], param_value)
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 70929b960..f1f6bcfb3 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -37,12 +37,10 @@
 import six
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
-from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl
+from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl import trainer_model_based_params
-from tensor2tensor.rl.dopamine_connector import DQNLearner
-from tensor2tensor.rl.ppo_learner import PPOLearner
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -52,19 +50,6 @@
 FLAGS = flags.FLAGS
 
 
-LEARNERS = {
-    "ppo": PPOLearner,
-    "dqn": DQNLearner,
-}
-
-
-def update_hparams_from_hparams(target_hparams, source_hparams, prefix):
-  """Copy a subset of hparams to target_hparams."""
-  for (param_name, param_value) in six.iteritems(source_hparams.values()):
-    if param_name.startswith(prefix):
-      target_hparams.set_hparam(param_name[len(prefix):], param_value)
-
-
 def real_env_step_increment(hparams):
   """Real env step increment."""
   return int(math.ceil(
@@ -207,7 +192,7 @@ def initial_frame_chooser(batch_size):
   base_algo_str = hparams.base_algo
   train_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
 
-  update_hparams_from_hparams(
+  rl_utils.update_hparams_from_hparams(
       train_hparams, hparams, base_algo_str + "_"
   )
 
@@ -223,7 +208,7 @@ def train_agent_real_env(env, learner, hparams, epoch):
   base_algo_str = hparams.base_algo
 
   train_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
-  update_hparams_from_hparams(
+  rl_utils.update_hparams_from_hparams(
       train_hparams, hparams, "real_" + base_algo_str + "_"
   )
 
@@ -263,82 +248,6 @@ def train_world_model(
   return world_model_steps_num
 
 
-def setup_env(hparams, batch_size, max_num_noops):
-  """Setup."""
-  game_mode = "Deterministic-v4"
-  camel_game_name = "".join(
-      [w[0].upper() + w[1:] for w in hparams.game.split("_")])
-  camel_game_name += game_mode
-  env_name = camel_game_name
-
-  env = T2TGymEnv(base_env_name=env_name,
-                  batch_size=batch_size,
-                  grayscale=hparams.grayscale,
-                  resize_width_factor=hparams.resize_width_factor,
-                  resize_height_factor=hparams.resize_height_factor,
-                  base_env_timesteps_limit=hparams.env_timesteps_limit,
-                  max_num_noops=max_num_noops)
-  return env
-
-
-def evaluate_single_config(hparams, stochastic, max_num_noops, agent_model_dir):
-  """Evaluate the PPO agent in the real environment."""
-  eval_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
-  env = setup_env(
-      hparams, batch_size=hparams.eval_batch_size, max_num_noops=max_num_noops
-  )
-  env.start_new_epoch(0)
-  env_fn = rl.make_real_env_fn(env)
-  learner = LEARNERS[hparams.base_algo](
-      hparams.frame_stack_size, base_event_dir=None,
-      agent_model_dir=agent_model_dir
-  )
-  learner.evaluate(env_fn, eval_hparams, stochastic)
-  rollouts = env.current_epoch_rollouts()
-  env.close()
-
-  return tuple(
-      compute_mean_reward(rollouts, clipped) for clipped in (True, False)
-  )
-
-
-def get_metric_name(stochastic, max_num_noops, clipped):
-  return "mean_reward/eval/stochastic_{}_max_noops_{}_{}".format(
-      stochastic, max_num_noops, "clipped" if clipped else "unclipped")
-
-
-def evaluate_all_configs(hparams, agent_model_dir):
-  """Evaluate the agent with multiple eval configurations."""
-  metrics = {}
-  # Iterate over all combinations of picking actions by sampling/mode and
-  # whether to do initial no-ops.
-  for stochastic in (True, False):
-    for max_num_noops in (hparams.eval_max_num_noops, 0):
-      scores = evaluate_single_config(
-          hparams, stochastic, max_num_noops, agent_model_dir
-      )
-      for (score, clipped) in zip(scores, (True, False)):
-        metric_name = get_metric_name(stochastic, max_num_noops, clipped)
-        metrics[metric_name] = score
-
-  return metrics
-
-
-def compute_mean_reward(rollouts, clipped):
-  """Calculate mean rewards from given epoch."""
-  reward_name = "reward" if clipped else "unclipped_reward"
-  rewards = []
-  for rollout in rollouts:
-    if rollout[-1].done:
-      rollout_reward = sum(getattr(frame, reward_name) for frame in rollout)
-      rewards.append(rollout_reward)
-  if rewards:
-    mean_rewards = np.mean(rewards)
-  else:
-    mean_rewards = 0
-  return mean_rewards
-
-
 def evaluate_world_model(real_env, hparams, world_model_dir, debug_video_path):
   """Evaluate the world model (reward accuracy)."""
   frame_stack_size = hparams.frame_stack_size
@@ -485,13 +394,13 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   epoch = -1
   data_dir = directories["data"]
-  env = setup_env(
+  env = rl_utils.setup_env(
       hparams, batch_size=hparams.real_batch_size,
       max_num_noops=hparams.max_num_noops
   )
   env.start_new_epoch(epoch, data_dir)
 
-  learner = LEARNERS[hparams.base_algo](
+  learner = rl_utils.LEARNERS[hparams.base_algo](
       hparams.frame_stack_size, directories["policy"],
       directories["policy"]
   )
@@ -507,7 +416,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   policy_model_dir = directories["policy"]
   tf.logging.info("Initial training of the policy in real environment.")
   train_agent_real_env(env, learner, hparams, epoch)
-  metrics["mean_reward/train/clipped"] = compute_mean_reward(
+  metrics["mean_reward/train/clipped"] = rl_utils.compute_mean_reward(
       env.current_epoch_rollouts(), clipped=True
   )
   tf.logging.info("Mean training reward (initial): {}".format(
@@ -555,14 +464,14 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       # we'd overwrite them with wrong data.
       log("Metrics found for this epoch, skipping evaluation.")
     else:
-      metrics["mean_reward/train/clipped"] = compute_mean_reward(
+      metrics["mean_reward/train/clipped"] = rl_utils.compute_mean_reward(
           env.current_epoch_rollouts(), clipped=True
       )
       log("Mean training reward: {}".format(
           metrics["mean_reward/train/clipped"]
       ))
 
-      eval_metrics = evaluate_all_configs(hparams, policy_model_dir)
+      eval_metrics = rl_utils.evaluate_all_configs(hparams, policy_model_dir)
       log("Agent eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
       metrics.update(eval_metrics)
 
@@ -582,7 +491,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       # Report metrics
       if report_fn:
         if report_metric == "mean_reward":
-          metric_name = get_metric_name(
+          metric_name = rl_utils.get_metric_name(
               stochastic=True, max_num_noops=hparams.eval_max_num_noops,
               clipped=False
           )
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 96ffebff7..31da7699c 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -27,19 +27,20 @@
 from __future__ import division
 from __future__ import print_function
 
-import six
+import pprint
 
-from tensor2tensor.data_generators import gym_env
 from tensor2tensor.models.research import rl
-from tensor2tensor.rl.ppo_learner import PPOLearner
+from tensor2tensor.rl import rl_utils
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
 
+
 flags = tf.flags
 FLAGS = flags.FLAGS
 
+
 # To maintain compatibility with some internal libs, we guard against these flag
 # definitions possibly erring. Apologies for the ugliness.
 try:
@@ -48,46 +49,48 @@
   pass
 
 
-LEARNERS = {
-    "ppo": PPOLearner
-}
-
-
-def update_hparams_from_hparams(target_hparams, source_hparams, prefix):
-  """Copy a subset of hparams to target_hparams."""
-  for (param_name, param_value) in six.iteritems(source_hparams.values()):
-    if param_name.startswith(prefix):
-      target_hparams.set_hparam(param_name[len(prefix):], param_value)
-
-
 def initialize_env_specs(hparams):
   """Initializes env_specs using T2TGymEnvs."""
-  if getattr(hparams, "game", None):
-    game_name = gym_env.camel_case_name(hparams.game)
-    env = gym_env.T2TGymEnv("{}Deterministic-v4".format(game_name),
-                            batch_size=hparams.batch_size)
-    env.start_new_epoch(0)
-    hparams.add_hparam("env_fn", rl.make_real_env_fn(env))
-    eval_env = gym_env.T2TGymEnv("{}Deterministic-v4".format(game_name),
-                                 batch_size=hparams.eval_batch_size)
-    eval_env.start_new_epoch(0)
-    hparams.add_hparam("eval_env_fn", rl.make_real_env_fn(eval_env))
+  env = rl_utils.setup_env(hparams, hparams.batch_size,
+                           hparams.eval_max_num_noops)
+  env.start_new_epoch(0)
+  hparams.add_hparam("env_fn", rl.make_real_env_fn(env))
   return hparams
 
 
 def train(hparams, output_dir, report_fn=None):
   hparams = initialize_env_specs(hparams)
-  learner = LEARNERS[hparams.base_algo](
+  learner = rl_utils.LEARNERS[hparams.base_algo](
       hparams.frame_stack_size, FLAGS.output_dir, output_dir
   )
   policy_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
-  update_hparams_from_hparams(
+  rl_utils.update_hparams_from_hparams(
       policy_hparams, hparams, hparams.base_algo + "_"
   )
-  learner.train(
-      hparams.env_fn, policy_hparams, simulated=False, save_continuously=True,
-      epoch=0, eval_env_fn=hparams.eval_env_fn, report_fn=report_fn
+  total_steps = policy_hparams.epochs_num
+  eval_every_epochs = policy_hparams.eval_every_epochs
+  if eval_every_epochs == 0:
+    eval_every_epochs = total_steps
+  policy_hparams.eval_every_epochs = 0
+
+  steps = list(range(eval_every_epochs, total_steps+1, eval_every_epochs))
+  if not steps or steps[-1] < eval_every_epochs:
+    steps.append(eval_every_epochs)
+  metric_name = rl_utils.get_metric_name(
+    stochastic=True, max_num_noops=hparams.eval_max_num_noops,
+    clipped=False
   )
+  for step in steps:
+    policy_hparams.epochs_num = step
+    learner.train(
+      hparams.env_fn, policy_hparams, simulated=False, save_continuously=True,
+      epoch=0
+    )
+    eval_metrics = rl_utils.evaluate_all_configs(hparams, output_dir)
+    tf.logging.info("Agent eval metrics:\n{}".format(
+      pprint.pformat(eval_metrics)))
+    if report_fn:
+      report_fn(eval_metrics[metric_name], step)
 
 
 def main(_):

From d77d89693900bf8a99817989e6a36268c8a5c456 Mon Sep 17 00:00:00 2001
From: blazejosinski <804945+blazejosinski@users.noreply.github.com>
Date: Tue, 4 Dec 2018 17:22:43 -0800
Subject: [PATCH 1326/2720] internal merge of PR #1269

PiperOrigin-RevId: 224074625
---
 tensor2tensor/models/research/rl.py    |  3 ++-
 tensor2tensor/rl/rl_utils.py           |  4 ++--
 tensor2tensor/rl/trainer_model_free.py | 11 ++++++-----
 3 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 5c52f3db0..62479b5fb 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -230,13 +230,14 @@ def mfrl_original():
 
 @registry.register_hparams
 def mfrl_base():
+  """Base set of hparams for model-free PPO."""
   hparams = mfrl_original()
   hparams.add_hparam("ppo_epochs_num", 3000)
   hparams.add_hparam("ppo_eval_every_epochs", 100)
   hparams.add_hparam("eval_max_num_noops", 8)
   hparams.add_hparam("resize_height_factor", 2)
   hparams.add_hparam("resize_width_factor", 2)
-  hparams.add_hparam("grayscale", 1)
+  hparams.add_hparam("grayscale", 0)
   hparams.add_hparam("env_timesteps_limit", -1)
   return hparams
 
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index cff0ae440..98a1f17e4 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -13,8 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-r"""Utilities for RL training
-"""
+"""Utilities for RL training."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -118,6 +117,7 @@ def setup_env(hparams, batch_size, max_num_noops):
                   max_num_noops=max_num_noops)
   return env
 
+
 def update_hparams_from_hparams(target_hparams, source_hparams, prefix):
   """Copy a subset of hparams to target_hparams."""
   for (param_name, param_value) in six.iteritems(source_hparams.values()):
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 31da7699c..dea16b1fa 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -59,6 +59,7 @@ def initialize_env_specs(hparams):
 
 
 def train(hparams, output_dir, report_fn=None):
+  """Train."""
   hparams = initialize_env_specs(hparams)
   learner = rl_utils.LEARNERS[hparams.base_algo](
       hparams.frame_stack_size, FLAGS.output_dir, output_dir
@@ -77,18 +78,18 @@ def train(hparams, output_dir, report_fn=None):
   if not steps or steps[-1] < eval_every_epochs:
     steps.append(eval_every_epochs)
   metric_name = rl_utils.get_metric_name(
-    stochastic=True, max_num_noops=hparams.eval_max_num_noops,
-    clipped=False
+      stochastic=True, max_num_noops=hparams.eval_max_num_noops,
+      clipped=False
   )
   for step in steps:
     policy_hparams.epochs_num = step
     learner.train(
-      hparams.env_fn, policy_hparams, simulated=False, save_continuously=True,
-      epoch=0
+        hparams.env_fn, policy_hparams, simulated=False, save_continuously=True,
+        epoch=0
     )
     eval_metrics = rl_utils.evaluate_all_configs(hparams, output_dir)
     tf.logging.info("Agent eval metrics:\n{}".format(
-      pprint.pformat(eval_metrics)))
+        pprint.pformat(eval_metrics)))
     if report_fn:
       report_fn(eval_metrics[metric_name], step)
 

From 0b892dda078233b8087176dc9e76028d47772568 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 4 Dec 2018 18:47:31 -0800
Subject: [PATCH 1327/2720] Rm dataset

PiperOrigin-RevId: 224085090
---
 tensor2tensor/data_generators/all_problems.py |   1 -
 tensor2tensor/data_generators/twentybn.py     | 121 ------------------
 2 files changed, 122 deletions(-)
 delete mode 100644 tensor2tensor/data_generators/twentybn.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index f402dd4e7..b3f7263f1 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -79,7 +79,6 @@
     "tensor2tensor.data_generators.translate_enmk",
     "tensor2tensor.data_generators.translate_envi",
     "tensor2tensor.data_generators.translate_enzh",
-    "tensor2tensor.data_generators.twentybn",
     "tensor2tensor.data_generators.video_generated",
     "tensor2tensor.data_generators.vqa",
     "tensor2tensor.data_generators.wiki",
diff --git a/tensor2tensor/data_generators/twentybn.py b/tensor2tensor/data_generators/twentybn.py
deleted file mode 100644
index b28162044..000000000
--- a/tensor2tensor/data_generators/twentybn.py
+++ /dev/null
@@ -1,121 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Data generator for twenty bn video data-set."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-from tensor2tensor.data_generators import video_utils
-from tensor2tensor.utils import registry
-
-import tensorflow as tf
-
-
-_FILE_VIDEO_PATTERN = '20bn-something-something-v1'
-_FILE_LABEL_PATTERN = 'something-something-v1-'
-
-
-def twentybn_generator(tmp_dir, training):
-  """Video generator for twenty-bn dataset.
-
-  Args:
-    tmp_dir: path to temporary storage directory.
-    training: a Boolean; if true, we use the train set, otherwise the dev set.
-
-
-  Yields:
-    A dictionary representing the images with the following fields:
-    * image/encoded: the string encoding the images of a video as JPG,
-    * image/format: the string "jpg" representing image format,
-    * image/class/label: an integer representing the label,
-  """
-  data_suffix = 'train' if training else 'validation'
-
-  def process_labels():
-    all_labels = {}
-    with tf.gfile.Open(tmp_dir + _FILE_LABEL_PATTERN + 'labels.csv') as f:
-      for (i, label) in enumerate(f):
-        all_labels[label] = i+1
-    return all_labels
-
-  def read_id_to_labels():
-    id_to_label = {}
-    with tf.gfile.Open(tmp_dir + _FILE_LABEL_PATTERN +
-                       data_suffix + '.csv') as f:
-      for line in f:
-        values = line.split(';')
-        id_to_label[int(values[0])] = values[1]
-    return id_to_label
-
-  # Get the label string to class id dictionary.
-  all_labels = process_labels()
-  # Get the video ids to label string dictionary.
-  id_to_labels = read_id_to_labels()
-
-  # Read video frames as images.
-  for vname, label_id in id_to_labels.items():
-    path = os.path.join(os.path.join(tmp_dir, _FILE_VIDEO_PATTERN), str(vname))
-    label = all_labels[label_id]
-    images = []
-    image_files = tf.gfile.Glob(os.path.join(path, '*.jpg'))
-
-    for filename in image_files:
-      with tf.gfile.Open(filename, 'rb') as f:
-        encoded_image_data = f.read()
-        images.append(encoded_image_data)
-    yield {
-        'image/encoded': images,
-        'image/format': ['jpg'],
-        'image/class/label': [int(label)]
-    }
-
-
-@registry.register_problem
-class VideoTwentybn(video_utils.Video2ClassProblem):
-  """Problem for twenty bn something-something dataset."""
-
-  @property
-  def is_small(self):
-    return True
-
-  @property
-  def num_classes(self):
-    return 174
-
-  @property
-  def train_shards(self):
-    return 100
-
-  @property
-  def dev_shards(self):
-    return 10
-
-  @property
-  def image_size(self):
-    return 32
-
-  def preprocess_example(self, example, unused_mode, unused_hparams):
-    example['inputs'] = video_utils.resize_video_frames(
-        example['inputs'], self.image_size)
-    return example
-
-  def generator(self, data_dir, tmp_dir, is_training):
-    if is_training:
-      return twentybn_generator(tmp_dir, True)
-    else:
-      return twentybn_generator(tmp_dir, False)

From 3689c42bf7d0c4c767c1f1e812cdac11b6e7b155 Mon Sep 17 00:00:00 2001
From: Wu Shangyu <534427411@qq.com>
Date: Wed, 5 Dec 2018 14:52:45 +0800
Subject: [PATCH 1328/2720] add a metrics for calculating pearson correlation
 coefficient (#1274)

* add pearson metrics and test method

* fix test errors

* update metrics test: dimension error

* update metrics test

* update metrics test

* update metrics test

* update metrics test

* fix errors
---
 tensor2tensor/utils/metrics.py      | 16 ++++++++++++++++
 tensor2tensor/utils/metrics_test.py | 15 +++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index d4ff79e08..164663bb9 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -43,6 +43,7 @@ class Metrics(object):
   APPROX_BLEU = "approx_bleu_score"
   RMSE = "rmse"
   LOG_POISSON = "log_poisson"
+  PEARSON = "pearson"
   R2 = "r_squared"
   ROUGE_2_F = "rouge_2_fscore"
   ROUGE_L_F = "rouge_L_fscore"
@@ -741,6 +742,20 @@ def from_characters(raw, lookup_):
     return distance / reference_length, reference_length
 
 
+def pearson_correlation_coefficient(predictions, labels, weights_fn=None):
+  """Calculate pearson correlation coefficient.
+
+  Args:
+    predictions: The raw predictions.
+    labels: The actual labels.
+    weights_fn: Weighting function.
+
+  Returns:
+    The pearson correlation coefficient.
+  """
+  _, pearson = tf.contrib.metrics.streaming_pearson_correlation(predictions, labels)
+  return pearson, tf.constant(1.0)
+
 # Metrics are functions that take predictions and labels and return
 # a tensor of metrics and a tensor of weights.
 # If the function has "features" as an argument, it will receive the whole
@@ -756,6 +771,7 @@ def from_characters(raw, lookup_):
     Metrics.APPROX_BLEU: bleu_hook.bleu_score,
     Metrics.RMSE: padded_rmse,
     Metrics.LOG_POISSON: padded_log_poisson,
+    Metrics.PEARSON: pearson_correlation_coefficient,
     Metrics.R2: padded_variance_explained,
     Metrics.ROUGE_2_F: rouge.rouge_2_fscore,
     Metrics.ROUGE_L_F: rouge.rouge_l_fscore,
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index b6228483c..454c06b8d 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -319,6 +319,21 @@ def testMultilabelMatch3(self):
       actual = session.run(a)
     self.assertAlmostEqual(actual, expected, places=6)
 
+  def testPearsonCorrelationCoefficient(self):
+    predictions = np.random.rand(12, 1)
+    targets = np.random.rand(12, 1)
+    
+    expected = np.corrcoef(np.squeeze(predictions), np.squeeze(targets))[0][1]
+    with self.test_session() as session:
+      pearson, _ = metrics.pearson_correlation_coefficient(
+              tf.constant(predictions, dtype=tf.float32), 
+              tf.constant(targets, dtype=tf.float32))
+      session.run(tf.global_variables_initializer())
+      session.run(tf.local_variables_initializer())
+      actual = session.run(pearson)
+    print(actual)
+    print(expected)
+    self.assertAlmostEqual(actual, expected)
 
 if __name__ == '__main__':
   tf.test.main()

From 51dfa7a10d15785b11172b4cfae9ac7805d9a38b Mon Sep 17 00:00:00 2001
From: Wu Shangyu <534427411@qq.com>
Date: Tue, 4 Dec 2018 23:10:32 -0800
Subject: [PATCH 1329/2720] internal merge of PR #1274

PiperOrigin-RevId: 224104845
---
 tensor2tensor/utils/metrics.py      | 4 +++-
 tensor2tensor/utils/metrics_test.py | 9 ++++-----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 164663bb9..aac82e868 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -753,7 +753,9 @@ def pearson_correlation_coefficient(predictions, labels, weights_fn=None):
   Returns:
     The pearson correlation coefficient.
   """
-  _, pearson = tf.contrib.metrics.streaming_pearson_correlation(predictions, labels)
+  del weights_fn
+  _, pearson = tf.contrib.metrics.streaming_pearson_correlation(predictions,
+                                                                labels)
   return pearson, tf.constant(1.0)
 
 # Metrics are functions that take predictions and labels and return
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 454c06b8d..a8dbde9ab 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -322,18 +322,17 @@ def testMultilabelMatch3(self):
   def testPearsonCorrelationCoefficient(self):
     predictions = np.random.rand(12, 1)
     targets = np.random.rand(12, 1)
-    
+
     expected = np.corrcoef(np.squeeze(predictions), np.squeeze(targets))[0][1]
     with self.test_session() as session:
       pearson, _ = metrics.pearson_correlation_coefficient(
-              tf.constant(predictions, dtype=tf.float32), 
-              tf.constant(targets, dtype=tf.float32))
+          tf.constant(predictions, dtype=tf.float32),
+          tf.constant(targets, dtype=tf.float32))
       session.run(tf.global_variables_initializer())
       session.run(tf.local_variables_initializer())
       actual = session.run(pearson)
-    print(actual)
-    print(expected)
     self.assertAlmostEqual(actual, expected)
 
+
 if __name__ == '__main__':
   tf.test.main()

From 8df419c96dfbd33da14446632877f8c595391344 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 5 Dec 2018 10:31:27 -0800
Subject: [PATCH 1330/2720] Have a few shards for en-ro translation.

PiperOrigin-RevId: 224177849
---
 tensor2tensor/data_generators/translate_enro.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/translate_enro.py b/tensor2tensor/data_generators/translate_enro.py
index ac734189f..238d79a61 100644
--- a/tensor2tensor/data_generators/translate_enro.py
+++ b/tensor2tensor/data_generators/translate_enro.py
@@ -89,7 +89,7 @@ def dataset_splits(self):
     """Splits of data to produce and number of output shards for each."""
     return [{
         "split": problem.DatasetSplit.TRAIN,
-        "shards": 1,
+        "shards": 16,  # It's a small dataset, TPUs like at least a few shards.
     }, {
         "split": problem.DatasetSplit.EVAL,
         "shards": 1,

From 19dcfcf16843d30108c7d8074adf65f89abc7822 Mon Sep 17 00:00:00 2001
From: Yuwen Yan <ybbaigo@gmail.com>
Date: Thu, 6 Dec 2018 06:58:25 +0800
Subject: [PATCH 1331/2720] add text CNN model for text classification problem
 (#1271)

---
 tensor2tensor/models/text_cnn.py | 106 +++++++++++++++++++++++++++++++
 1 file changed, 106 insertions(+)
 create mode 100644 tensor2tensor/models/text_cnn.py

diff --git a/tensor2tensor/models/text_cnn.py b/tensor2tensor/models/text_cnn.py
new file mode 100644
index 000000000..be2945fa3
--- /dev/null
+++ b/tensor2tensor/models/text_cnn.py
@@ -0,0 +1,106 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""TextCNN model from "Convolutional Neural Networks for Sentence Classification".
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_hparams
+from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
+
+import tensorflow as tf
+
+@registry.register_model
+class TextCNN(t2t_model.T2TModel):
+  """Text CNN."""
+
+  def body(self, features):
+    """TextCNN main model_fn.
+    Args:
+      features: Map of features to the model. Should contain the following:
+          "inputs": Text inputs.
+              [batch_size, input_length, 1, hidden_dim].
+          "targets": Target encoder outputs.
+              [batch_size, 1, 1, hidden_dim]
+    Returns:
+      Final encoder representation. [batch_size, 1, 1, hidden_dim]
+    """
+    hparams = self._hparams
+    inputs = features["inputs"]
+
+    xshape = common_layers.shape_list(inputs)
+
+    vocab_size = xshape[3]
+    inputs = tf.reshape(inputs, [xshape[0], xshape[1], xshape[3], xshape[2]])
+
+    pooled_outputs = []
+    for _, filter_size in enumerate(hparams.filter_sizes):
+      with tf.name_scope("conv-maxpool-%s" % filter_size):
+        filter_shape = [filter_size, vocab_size, 1, hparams.num_filters]
+        filter_var = tf.Variable(tf.truncated_normal(filter_shape, stddev=0.1), name="W")
+        filter_bias = tf.Variable(tf.constant(0.1, shape=[hparams.num_filters]), name="b")
+        conv = tf.nn.conv2d(
+            inputs,
+            filter_var,
+            strides=[1, 1, 1, 1],
+            padding="VALID",
+            name="conv")
+        conv_outputs = tf.nn.relu(tf.nn.bias_add(conv, filter_bias), name="relu")
+        pooled = tf.math.reduce_max(conv_outputs, axis=1, keepdims=True, name="max")
+        pooled_outputs.append(pooled)
+
+    num_filters_total = hparams.num_filters * len(hparams.filter_sizes)
+    h_pool = tf.concat(pooled_outputs, 3)
+    h_pool_flat = tf.reshape(h_pool, [-1, num_filters_total])
+
+    # Add dropout
+    output = tf.nn.dropout(h_pool_flat, 1 - hparams.output_dropout)
+    output = tf.reshape(output, [-1, 1, 1, num_filters_total])
+
+    return output
+
+@registry.register_hparams
+def text_cnn_base():
+  """Set of hyperparameters."""
+  hparams = common_hparams.basic_params1()
+  hparams.batch_size = 4096
+  hparams.max_length = 256
+  hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
+  hparams.optimizer_adam_epsilon = 1e-9
+  hparams.learning_rate_schedule = "legacy"
+  hparams.learning_rate_decay_scheme = "noam"
+  hparams.learning_rate = 0.1
+  hparams.learning_rate_warmup_steps = 4000
+  hparams.initializer_gain = 1.0
+  hparams.num_hidden_layers = 6
+  hparams.initializer = "uniform_unit_scaling"
+  hparams.weight_decay = 0.0
+  hparams.optimizer_adam_beta1 = 0.9
+  hparams.optimizer_adam_beta2 = 0.98
+  hparams.num_sampled_classes = 0
+  hparams.label_smoothing = 0.1
+  hparams.shared_embedding_and_softmax_weights = True
+  hparams.symbol_modality_num_shards = 16
+
+  # Add new ones like this.
+  hparams.add_hparam("filter_sizes", [2, 3, 4, 5])
+  hparams.add_hparam("num_filters", 128)
+  hparams.add_hparam("output_dropout", 0.4)
+  return hparams

From 7de63449a98375011e2a8715482dfeea946e6de7 Mon Sep 17 00:00:00 2001
From: Yuwen Yan <ybbaigo@gmail.com>
Date: Wed, 5 Dec 2018 14:58:51 -0800
Subject: [PATCH 1332/2720] internal merge of PR #1271

PiperOrigin-RevId: 224227181
---
 tensor2tensor/models/text_cnn.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/text_cnn.py b/tensor2tensor/models/text_cnn.py
index be2945fa3..532506853 100644
--- a/tensor2tensor/models/text_cnn.py
+++ b/tensor2tensor/models/text_cnn.py
@@ -13,8 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""TextCNN model from "Convolutional Neural Networks for Sentence Classification".
-"""
+"""TextCNN (see Convolutional Neural Networks for Sentence Classification)."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -27,12 +26,14 @@
 
 import tensorflow as tf
 
+
 @registry.register_model
 class TextCNN(t2t_model.T2TModel):
   """Text CNN."""
 
   def body(self, features):
     """TextCNN main model_fn.
+
     Args:
       features: Map of features to the model. Should contain the following:
           "inputs": Text inputs.
@@ -54,16 +55,20 @@ def body(self, features):
     for _, filter_size in enumerate(hparams.filter_sizes):
       with tf.name_scope("conv-maxpool-%s" % filter_size):
         filter_shape = [filter_size, vocab_size, 1, hparams.num_filters]
-        filter_var = tf.Variable(tf.truncated_normal(filter_shape, stddev=0.1), name="W")
-        filter_bias = tf.Variable(tf.constant(0.1, shape=[hparams.num_filters]), name="b")
+        filter_var = tf.Variable(
+            tf.truncated_normal(filter_shape, stddev=0.1), name="W")
+        filter_bias = tf.Variable(
+            tf.constant(0.1, shape=[hparams.num_filters]), name="b")
         conv = tf.nn.conv2d(
             inputs,
             filter_var,
             strides=[1, 1, 1, 1],
             padding="VALID",
             name="conv")
-        conv_outputs = tf.nn.relu(tf.nn.bias_add(conv, filter_bias), name="relu")
-        pooled = tf.math.reduce_max(conv_outputs, axis=1, keepdims=True, name="max")
+        conv_outputs = tf.nn.relu(
+            tf.nn.bias_add(conv, filter_bias), name="relu")
+        pooled = tf.math.reduce_max(
+            conv_outputs, axis=1, keepdims=True, name="max")
         pooled_outputs.append(pooled)
 
     num_filters_total = hparams.num_filters * len(hparams.filter_sizes)
@@ -76,6 +81,7 @@ def body(self, features):
 
     return output
 
+
 @registry.register_hparams
 def text_cnn_base():
   """Set of hyperparameters."""

From 266380524e4651bbe7f217c64f2a4ac47da2c82b Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 6 Dec 2018 02:04:51 +0100
Subject: [PATCH 1333/2720] Temperature (#1277)

* Use T2TModel for policies

* Implement sampling with temperature from policy

* Fixes
---
 tensor2tensor/models/research/rl.py           | 177 +++++++++++-------
 tensor2tensor/rl/dopamine_connector.py        |  16 +-
 tensor2tensor/rl/policy_learner.py            |  14 +-
 tensor2tensor/rl/ppo.py                       |   8 +-
 tensor2tensor/rl/ppo_learner.py               |  61 +++---
 tensor2tensor/rl/rl_utils.py                  |  24 +--
 tensor2tensor/rl/trainer_model_based.py       |  12 +-
 .../rl/trainer_model_based_params.py          |  11 +-
 tensor2tensor/rl/trainer_model_free.py        |   3 +-
 tensor2tensor/rl/trainer_model_free_test.py   |   5 +-
 10 files changed, 210 insertions(+), 121 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 62479b5fb..c596d4df5 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -27,6 +27,7 @@
 from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
 from tensor2tensor.rl.envs.simulated_batch_gym_env import SimulatedBatchGymEnv
 from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
 
 import tensorflow as tf
 import tensorflow_probability as tfp
@@ -60,12 +61,12 @@ def ppo_base_v1():
   return hparams
 
 
-@registry.register_hparams
-def ppo_continuous_action_base():
-  hparams = ppo_base_v1()
-  hparams.add_hparam("policy_network", feed_forward_gaussian_fun)
-  hparams.add_hparam("policy_network_params", "basic_policy_parameters")
-  return hparams
+#@registry.register_hparams
+#def ppo_continuous_action_base():
+#  hparams = ppo_base_v1()
+#  hparams.add_hparam("policy_network", feed_forward_gaussian_fun)
+#  hparams.add_hparam("policy_network_params", "basic_policy_parameters")
+#  return hparams
 
 
 @registry.register_hparams
@@ -77,14 +78,14 @@ def basic_policy_parameters():
 @registry.register_hparams
 def ppo_discrete_action_base():
   hparams = ppo_base_v1()
-  hparams.add_hparam("policy_network", feed_forward_categorical_fun)
+  hparams.add_hparam("policy_network", "feed_forward_categorical_policy")
   return hparams
 
 
 @registry.register_hparams
 def discrete_random_action_base():
   hparams = common_hparams.basic_params1()
-  hparams.add_hparam("policy_network", random_policy_fun)
+  hparams.add_hparam("policy_network", "random_policy")
   return hparams
 
 
@@ -100,7 +101,7 @@ def ppo_atari_base():
   hparams.value_loss_coef = 1
   hparams.optimization_epochs = 3
   hparams.epochs_num = 1000
-  hparams.policy_network = feed_forward_cnn_small_categorical_fun
+  hparams.policy_network = "feed_forward_cnn_small_categorical_policy"
   hparams.clipping_coef = 0.2
   hparams.optimization_batch_size = 20
   hparams.max_gradients_norm = 0.5
@@ -157,15 +158,28 @@ def get_policy(observations, hparams, action_space):
   """Get a policy network.
 
   Args:
-    observations: Tensor with observations
+    observations
     hparams: parameters
     action_space: action space
 
   Returns:
-    Tensor with policy and value function output
+    Tuple (action logits, value).
   """
-  policy_network_lambda = hparams.policy_network
-  return policy_network_lambda(action_space, hparams, observations)
+  if not isinstance(action_space, gym.spaces.Discrete):
+    raise ValueError("Expecting discrete action space.")
+
+  model = registry.model(hparams.policy_network)(
+      hparams, tf.estimator.ModeKeys.TRAIN
+  )
+  obs_shape = common_layers.shape_list(observations)
+  features = {
+      "inputs": observations,
+      "target_action": tf.zeros(obs_shape[:2] + [action_space.n]),
+      "target_value": tf.zeros(obs_shape[:2])
+  }
+  with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+    (targets, _) = model(features)
+  return (targets["target_action"], targets["target_value"])
 
 
 @registry.register_hparams
@@ -173,7 +187,7 @@ def ppo_pong_ae_base():
   """Pong autoencoder base parameters."""
   hparams = ppo_original_params()
   hparams.learning_rate = 1e-4
-  hparams.network = dense_bitwise_categorical_fun
+  hparams.network = "dense_bitwise_categorical_policy"
   return hparams
 
 
@@ -225,6 +239,12 @@ def mfrl_original():
       batch_size=16,
       eval_batch_size=2,
       frame_stack_size=4,
+      eval_sampling_temps=[0.0, 0.2, 0.5, 0.8, 1.0, 2.0],
+      eval_max_num_noops=8,
+      resize_height_factor=2,
+      resize_width_factor=2,
+      grayscale=0,
+      env_timesteps_limit=-1,
   )
 
 
@@ -234,11 +254,6 @@ def mfrl_base():
   hparams = mfrl_original()
   hparams.add_hparam("ppo_epochs_num", 3000)
   hparams.add_hparam("ppo_eval_every_epochs", 100)
-  hparams.add_hparam("eval_max_num_noops", 8)
-  hparams.add_hparam("resize_height_factor", 2)
-  hparams.add_hparam("resize_width_factor", 2)
-  hparams.add_hparam("grayscale", 0)
-  hparams.add_hparam("env_timesteps_limit", -1)
   return hparams
 
 
@@ -250,10 +265,18 @@ def mfrl_tiny():
   return hparams
 
 
+class DiscretePolicyBase(t2t_model.T2TModel):
+
+  @staticmethod
+  def _get_num_actions(features):
+    return common_layers.shape_list(features["target_action"])[2]
+
+
 NetworkOutput = collections.namedtuple(
     "NetworkOutput", "policy, value, action_postprocessing")
 
 
+# TODO(koz4k): Translate it to T2TModel or remove.
 def feed_forward_gaussian_fun(action_space, config, observations):
   """Feed-forward Gaussian."""
   if not isinstance(action_space, gym.spaces.box.Box):
@@ -303,36 +326,40 @@ def clip_logits(logits, config):
     return logits
 
 
-def feed_forward_categorical_fun(action_space, config, observations):
+@registry.register_model
+class FeedForwardCategoricalPolicy(DiscretePolicyBase):
   """Feed-forward categorical."""
-  if not isinstance(action_space, gym.spaces.Discrete):
-    raise ValueError("Expecting discrete action space.")
-  flat_observations = tf.reshape(observations, [
-      tf.shape(observations)[0], tf.shape(observations)[1],
-      functools.reduce(operator.mul, observations.shape.as_list()[2:], 1)])
-  with tf.variable_scope("network_parameters"):
+
+  def body(self, features):
+    observations = features["inputs"]
+    flat_observations = tf.reshape(observations, [
+        tf.shape(observations)[0], tf.shape(observations)[1],
+        functools.reduce(operator.mul, observations.shape.as_list()[2:], 1)])
     with tf.variable_scope("policy"):
       x = flat_observations
-      for size in config.policy_layers:
+      for size in self.hparams.policy_layers:
         x = tf.contrib.layers.fully_connected(x, size, tf.nn.relu)
-      logits = tf.contrib.layers.fully_connected(x, action_space.n,
-                                                 activation_fn=None)
+      logits = tf.contrib.layers.fully_connected(
+          x, self._get_num_actions(features), activation_fn=None
+      )
     with tf.variable_scope("value"):
       x = flat_observations
-      for size in config.value_layers:
+      for size in self.hparams.value_layers:
         x = tf.contrib.layers.fully_connected(x, size, tf.nn.relu)
       value = tf.contrib.layers.fully_connected(x, 1, None)[..., 0]
-  logits = clip_logits(logits, config)
-  policy = tfp.distributions.Categorical(logits=logits)
-  return NetworkOutput(policy, value, lambda a: a)
+    logits = clip_logits(logits, self.hparams)
+    return {"target_action": logits, "target_value": value}
 
 
-def feed_forward_cnn_small_categorical_fun(action_space, config, observations):
+@registry.register_model
+class FeedForwardCnnSmallCategoricalPolicy(DiscretePolicyBase):
   """Small cnn network with categorical output."""
-  obs_shape = common_layers.shape_list(observations)
-  x = tf.reshape(observations, [-1] + obs_shape[2:])
-  with tf.variable_scope("network_parameters"):
-    dropout = getattr(config, "dropout_ppo", 0.0)
+
+  def body(self, features):
+    observations = features["inputs"]
+    obs_shape = common_layers.shape_list(observations)
+    x = tf.reshape(observations, [-1] + obs_shape[2:])
+    dropout = getattr(self.hparams, "dropout_ppo", 0.0)
     with tf.variable_scope("feed_forward_cnn_small"):
       x = tf.to_float(x) / 255.0
       x = tf.contrib.layers.conv2d(x, 32, [5, 5], [2, 2],
@@ -346,23 +373,25 @@ def feed_forward_cnn_small_categorical_fun(action_space, config, observations):
       flat_x = tf.nn.dropout(flat_x, keep_prob=1.0 - dropout)
       x = tf.contrib.layers.fully_connected(flat_x, 128, tf.nn.relu)
 
-      logits = tf.contrib.layers.fully_connected(x, action_space.n,
-                                                 activation_fn=None)
-      logits = clip_logits(logits, config)
+      logits = tf.contrib.layers.fully_connected(
+          x, self._get_num_actions(features), activation_fn=None
+      )
+      logits = clip_logits(logits, self.hparams)
 
       value = tf.contrib.layers.fully_connected(
           x, 1, activation_fn=None)[..., 0]
-      policy = tfp.distributions.Categorical(logits=logits)
-  return NetworkOutput(policy, value, lambda a: a)
+    return {"target_action": logits, "target_value": value}
 
 
-def feed_forward_cnn_small_categorical_fun_new(
-    action_space, config, observations):
+@registry.register_model
+class FeedForwardCnnSmallCategoricalPolicyNew(DiscretePolicyBase):
   """Small cnn network with categorical output."""
-  obs_shape = common_layers.shape_list(observations)
-  x = tf.reshape(observations, [-1] + obs_shape[2:])
-  with tf.variable_scope("network_parameters"):
-    dropout = getattr(config, "dropout_ppo", 0.0)
+
+  def body(self, features):
+    observations = features["inputs"]
+    obs_shape = common_layers.shape_list(observations)
+    x = tf.reshape(observations, [-1] + obs_shape[2:])
+    dropout = getattr(self.hparams, "dropout_ppo", 0.0)
     with tf.variable_scope("feed_forward_cnn_small"):
       x = tf.to_float(x) / 255.0
       x = tf.nn.dropout(x, keep_prob=1.0 - dropout)
@@ -384,22 +413,23 @@ def feed_forward_cnn_small_categorical_fun_new(
       flat_x = tf.nn.dropout(flat_x, keep_prob=1.0 - dropout)
       x = tf.layers.dense(flat_x, 128, activation=tf.nn.relu, name="dense1")
 
-      logits = tf.layers.dense(x, action_space.n, name="dense2")
-      logits = clip_logits(logits, config)
+      logits = tf.layers.dense(
+          x, self._get_num_actions(features), name="dense2"
+      )
+      logits = clip_logits(logits, self.hparams)
 
       value = tf.layers.dense(x, 1, name="value")[..., 0]
-      policy = tfp.distributions.Categorical(logits=logits)
+    return {"target_action": logits, "target_value": value}
 
-  return NetworkOutput(policy, value, lambda a: a)
 
-
-def dense_bitwise_categorical_fun(action_space, config, observations):
+@registry.register_model
+class DenseBitwiseCategoricalPolicy(DiscretePolicyBase):
   """Dense network with bitwise input and categorical output."""
-  del config
-  obs_shape = common_layers.shape_list(observations)
-  x = tf.reshape(observations, [-1] + obs_shape[2:])
 
-  with tf.variable_scope("network_parameters"):
+  def body(self, features):
+    observations = features["inputs"]
+    obs_shape = common_layers.shape_list(observations)
+    x = tf.reshape(observations, [-1] + obs_shape[2:])
     with tf.variable_scope("dense_bitwise"):
       x = discretization.int_to_bit_embed(x, 8, 32)
       flat_x = tf.reshape(
@@ -409,22 +439,29 @@ def dense_bitwise_categorical_fun(action_space, config, observations):
       x = tf.contrib.layers.fully_connected(flat_x, 256, tf.nn.relu)
       x = tf.contrib.layers.fully_connected(flat_x, 128, tf.nn.relu)
 
-      logits = tf.contrib.layers.fully_connected(x, action_space.n,
-                                                 activation_fn=None)
+      logits = tf.contrib.layers.fully_connected(
+          x, self._get_num_actions(features), activation_fn=None
+      )
 
       value = tf.contrib.layers.fully_connected(
           x, 1, activation_fn=None)[..., 0]
-      policy = tfp.distributions.Categorical(logits=logits)
 
-  return NetworkOutput(policy, value, lambda a: a)
+    return {"target_action": logits, "target_value": value}
 
 
-def random_policy_fun(action_space, unused_config, observations):
+@registry.register_model
+class RandomPolicy(DiscretePolicyBase):
   """Random policy with categorical output."""
-  obs_shape = observations.shape.as_list()
-  with tf.variable_scope("network_parameters"):
+
+  def body(self, features):
+    observations = features["inputs"]
+    obs_shape = observations.shape.as_list()
+    # Just so Saver doesn't complain because of no variables.
+    tf.get_variable("dummy_var", initializer=0.0)
+    num_actions = self._get_num_actions(features)
+    logits = tf.constant(
+        1. / float(num_actions),
+        shape=(obs_shape[:2] + [num_actions])
+    )
     value = tf.zeros(obs_shape[:2])
-    policy = tfp.distributions.Categorical(
-        probs=[[[1. / float(action_space.n)] * action_space.n] *
-               (obs_shape[0] * obs_shape[1])])
-  return NetworkOutput(policy, value, lambda a: a)
+    return {"target_action": logits, "target_value": value}
diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 4843d45ae..0f227cd9f 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from copy import copy
 from dopamine.agents.dqn import dqn_agent
 from dopamine.agents.dqn.dqn_agent import NATURE_DQN_OBSERVATION_SHAPE
 from dopamine.agents.dqn.dqn_agent import NATURE_DQN_STACK_SIZE
@@ -285,6 +286,7 @@ def train(self,
             simulated,
             save_continuously,
             epoch,
+            sampling_temp=1.0,
             num_env_steps=None,
             env_step_multiplier=1,
             eval_env_fn=None,
@@ -294,6 +296,11 @@ def train(self,
     if num_env_steps is None:
       num_env_steps = hparams.num_frames
 
+    hparams = copy(hparams)
+    hparams.set_hparam(
+        "agent_epsilon_eval", min(hparams.agent_epsilon_eval * sampling_temp, 1)
+    )
+
     target_iterations, training_steps_per_iteration = \
       self._target_iteractions_and_steps(
           num_env_steps=num_env_steps * env_step_multiplier,
@@ -307,11 +314,14 @@ def train(self,
 
     self.completed_iterations = target_iterations
 
-  def evaluate(self, env_fn, hparams, stochastic):
+  def evaluate(self, env_fn, hparams, sampling_temp):
     target_iterations = 0
     training_steps_per_iteration = 0
-    if not stochastic:
-      hparams.set_hparam("agent_epsilon_eval", 0.)
+
+    hparams = copy(hparams)
+    hparams.set_hparam(
+        "agent_epsilon_eval", min(hparams.agent_epsilon_eval * sampling_temp, 1)
+    )
 
     create_environment_fn = get_create_env_fun(
         env_fn, time_limit=hparams.time_limit)
diff --git a/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
index 804b6b2f8..8368b714d 100644
--- a/tensor2tensor/rl/policy_learner.py
+++ b/tensor2tensor/rl/policy_learner.py
@@ -29,13 +29,21 @@ def __init__(self, frame_stack_size, base_event_dir, agent_model_dir):
     self.agent_model_dir = agent_model_dir
 
   def train(
-      self, env_fn, hparams, simulated, save_continuously, epoch,
-      num_env_steps=None, env_step_multiplier=1, eval_env_fn=None,
+      self,
+      env_fn,
+      hparams,
+      simulated,
+      save_continuously,
+      epoch,
+      sampling_temp=1.0,
+      num_env_steps=None,
+      env_step_multiplier=1,
+      eval_env_fn=None,
       report_fn=None
   ):
     # TODO(konradczechowski): pass name_scope instead of epoch?
     # TODO(konradczechowski): move 'simulated' to  batch_env
     raise NotImplementedError()
 
-  def evaluate(self, env_fn, hparams, stochastic):
+  def evaluate(self, env_fn, hparams, sampling_temp):
     raise NotImplementedError()
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index 00d4c37e9..e9db06f51 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -24,6 +24,7 @@
 from tensor2tensor.models.research.rl import get_policy
 
 import tensorflow as tf
+import tensorflow_probability as tfp
 
 
 def get_optimiser(config):
@@ -35,7 +36,10 @@ def get_optimiser(config):
 def define_ppo_step(data_points, optimizer, hparams, action_space):
   """Define ppo step."""
   observation, action, discounted_reward, norm_advantage, old_pdf = data_points
-  new_policy_dist, new_value, _ = get_policy(observation, hparams, action_space)
+
+  (logits, new_value) = get_policy(observation, hparams, action_space)
+  new_policy_dist = tfp.distributions.Categorical(logits=logits)
+
   new_pdf = new_policy_dist.prob(action)
 
   ratio = new_pdf / old_pdf
@@ -112,7 +116,7 @@ def define_ppo_epoch(memory, hparams, action_space, batch_size):
   dataset = dataset.shuffle(buffer_size=hparams.epoch_length-1,
                             reshuffle_each_iteration=True)
   dataset = dataset.repeat(-1)
-  dataset = dataset.batch(hparams.optimization_batch_size)
+  dataset = dataset.batch(hparams.optimization_batch_size, drop_remainder=True)
   iterator = dataset.make_initializable_iterator()
   optimizer = get_optimiser(hparams)
 
diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 2a6fa5603..35310b0a8 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -22,6 +22,7 @@
 import math
 import os
 
+from tensor2tensor.layers import common_layers
 from tensor2tensor.models.research.rl import get_policy
 from tensor2tensor.rl import ppo
 from tensor2tensor.rl.envs.tf_atari_wrappers import StackWrapper
@@ -30,6 +31,7 @@
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
+import tensorflow_probability as tfp
 
 
 class PPOLearner(PolicyLearner):
@@ -45,10 +47,14 @@ def train(self,
             simulated,
             save_continuously,
             epoch,
+            sampling_temp=1.0,
             num_env_steps=None,
             env_step_multiplier=1,
             eval_env_fn=None,
             report_fn=None):
+    assert sampling_temp == 1.0 or hparams.learning_rate == 0.0, \
+        "Sampling with non-1 temperature does not make sense during training."
+
     if not save_continuously:
       # We do not save model, as that resets frames that we need at restarts.
       # But we need to save at the last step, so we set it very high.
@@ -71,6 +77,7 @@ def train(self,
                   env,
                   hparams,
                   eval_env_fn,
+                  sampling_temp,
                   frame_stack_size=self.frame_stack_size,
                   force_beginning_resets=simulated))
 
@@ -93,12 +100,7 @@ def train(self,
             initializers,
             report_fn=report_fn)
 
-  def evaluate(self, env_fn, hparams, stochastic):
-    if stochastic:
-      policy_to_actions_lambda = lambda policy: policy.sample()
-    else:
-      policy_to_actions_lambda = lambda policy: policy.mode()
-
+  def evaluate(self, env_fn, hparams, sampling_temp):
     with tf.Graph().as_default():
       with tf.name_scope("rl_eval"):
         eval_env = env_fn(in_graph=True)
@@ -109,9 +111,11 @@ def evaluate(self, env_fn, hparams, stochastic):
             eval_phase=True,
             frame_stack_size=self.frame_stack_size,
             force_beginning_resets=False,
-            policy_to_actions_lambda=policy_to_actions_lambda)
+            sampling_temp=sampling_temp,
+        )
         model_saver = tf.train.Saver(
-            tf.global_variables(".*network_parameters.*"))
+            tf.global_variables(hparams.policy_network + "/.*")
+        )
 
         with tf.Session() as sess:
           sess.run(tf.global_variables_initializer())
@@ -121,7 +125,13 @@ def evaluate(self, env_fn, hparams, stochastic):
           sess.run(collect_memory)
 
 
-def _define_train(train_env, ppo_hparams, eval_env_fn=None, **collect_kwargs):
+def _define_train(
+    train_env,
+    ppo_hparams,
+    eval_env_fn=None,
+    sampling_temp=1.0,
+    **collect_kwargs
+):
   """Define the training setup."""
   memory, collect_summary, train_initialization = (
       _define_collect(
@@ -129,13 +139,14 @@ def _define_train(train_env, ppo_hparams, eval_env_fn=None, **collect_kwargs):
           ppo_hparams,
           "ppo_train",
           eval_phase=False,
-          policy_to_actions_lambda=(lambda policy: policy.sample()),
+          sampling_temp=sampling_temp,
           **collect_kwargs))
   ppo_summary = ppo.define_ppo_epoch(
       memory, ppo_hparams, train_env.action_space, train_env.batch_size)
   train_summary = tf.summary.merge([collect_summary, ppo_summary])
 
   if ppo_hparams.eval_every_epochs:
+    # TODO(koz4k): Do we need this at all?
     assert eval_env_fn is not None
     eval_env = eval_env_fn(in_graph=True)
     (_, eval_collect_summary, eval_initialization) = (
@@ -144,7 +155,7 @@ def _define_train(train_env, ppo_hparams, eval_env_fn=None, **collect_kwargs):
             ppo_hparams,
             "ppo_eval",
             eval_phase=True,
-            policy_to_actions_lambda=(lambda policy: policy.mode()),
+            sampling_temp=0.0,
             **collect_kwargs))
     return (train_summary, eval_collect_summary, (train_initialization,
                                                   eval_initialization))
@@ -164,7 +175,9 @@ def _run_train(ppo_hparams,
   summary_writer = tf.summary.FileWriter(
       event_dir, graph=tf.get_default_graph(), flush_secs=60)
 
-  model_saver = tf.train.Saver(tf.global_variables(".*network_parameters.*"))
+  model_saver = tf.train.Saver(
+      tf.global_variables(ppo_hparams.policy_network + "/.*")
+  )
 
   with tf.Session() as sess:
     sess.run(tf.global_variables_initializer())
@@ -265,17 +278,17 @@ def simulate(self, action):
 
 
 def _define_collect(batch_env, ppo_hparams, scope, frame_stack_size, eval_phase,
-                    policy_to_actions_lambda, force_beginning_resets):
+                    sampling_temp, force_beginning_resets):
   """Collect trajectories.
 
   Args:
     batch_env: Batch environment.
     ppo_hparams: PPO hparams, defined in tensor2tensor.models.research.rl.
     scope: var scope.
-    frame_stack_size: TODO(koz4k): Write docstring.
+    frame_stack_size: Number of last observations to feed into the policy.
     eval_phase: TODO(koz4k): Write docstring.
-    policy_to_actions_lambda: TODO(koz4k): Write docstring.
-    force_beginning_resets: TODO(koz4k): Write docstring.
+    sampling_temp: Sampling temperature for the policy.
+    force_beginning_resets: Whether to reset at the beginning of each episode.
 
   Returns:
     Returns memory (observations, rewards, dones, actions,
@@ -349,16 +362,16 @@ def step(index, scores_sum, scores_num):
 
       def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
         """Step of the environment."""
-        actor_critic = get_policy(
-            tf.expand_dims(obs_copy, 0), ppo_hparams, batch_env.action_space)
-        policy = actor_critic.policy
-        action = policy_to_actions_lambda(policy)
 
-        postprocessed_action = actor_critic.action_postprocessing(action)
-        reward, done = batch_env.simulate(postprocessed_action[0, ...])
+        (logits, value_function) = get_policy(
+            tf.expand_dims(obs_copy, 0), ppo_hparams, batch_env.action_space
+        )
+        action = common_layers.sample_with_temperature(logits, sampling_temp)
+        action = tf.cast(action, tf.int32)
+
+        reward, done = batch_env.simulate(action[0, ...])
 
-        pdf = policy.prob(action)[0]
-        value_function = actor_critic.value[0]
+        pdf = tfp.distributions.Categorical(logits=logits).prob(action)
         pdf = tf.reshape(pdf, shape=(num_agents,))
         value_function = tf.reshape(value_function, shape=(num_agents,))
         done = tf.reshape(done, shape=(num_agents,))
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 98a1f17e4..f4acecd73 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -50,13 +50,15 @@ def compute_mean_reward(rollouts, clipped):
   return mean_rewards
 
 
-def get_metric_name(stochastic, max_num_noops, clipped):
-  return "mean_reward/eval/stochastic_{}_max_noops_{}_{}".format(
-      stochastic, max_num_noops, "clipped" if clipped else "unclipped")
+def get_metric_name(sampling_temp, max_num_noops, clipped):
+  return "mean_reward/eval/sampling_temp_{}_max_noops_{}_{}".format(
+      sampling_temp, max_num_noops, "clipped" if clipped else "unclipped"
+  )
 
 
-def evaluate_single_config(hparams, stochastic, max_num_noops,
-                           agent_model_dir):
+def evaluate_single_config(
+    hparams, sampling_temp, max_num_noops, agent_model_dir
+):
   """Evaluate the PPO agent in the real environment."""
   eval_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
   env = setup_env(
@@ -68,7 +70,7 @@ def evaluate_single_config(hparams, stochastic, max_num_noops,
       hparams.frame_stack_size, base_event_dir=None,
       agent_model_dir=agent_model_dir
   )
-  learner.evaluate(env_fn, eval_hparams, stochastic)
+  learner.evaluate(env_fn, eval_hparams, sampling_temp)
   rollouts = env.current_epoch_rollouts()
   env.close()
 
@@ -80,15 +82,15 @@ def evaluate_single_config(hparams, stochastic, max_num_noops,
 def evaluate_all_configs(hparams, agent_model_dir):
   """Evaluate the agent with multiple eval configurations."""
   metrics = {}
-  # Iterate over all combinations of picking actions by sampling/mode and
-  # whether to do initial no-ops.
-  for stochastic in (True, False):
+  # Iterate over all combinations of sampling temperatures and whether to do
+  # initial no-ops.
+  for sampling_temp in hparams.eval_sampling_temps:
     for max_num_noops in (hparams.eval_max_num_noops, 0):
       scores = evaluate_single_config(
-          hparams, stochastic, max_num_noops, agent_model_dir
+          hparams, sampling_temp, max_num_noops, agent_model_dir
       )
       for (score, clipped) in zip(scores, (True, False)):
-        metric_name = get_metric_name(stochastic, max_num_noops, clipped)
+        metric_name = get_metric_name(sampling_temp, max_num_noops, clipped)
         metrics[metric_name] = score
 
   return metrics
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index f1f6bcfb3..3b8aad5d0 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -215,8 +215,13 @@ def train_agent_real_env(env, learner, hparams, epoch):
   env_fn = rl.make_real_env_fn(env)
   num_env_steps = real_env_step_increment(hparams)
   learner.train(
-      env_fn, train_hparams, simulated=False, save_continuously=False,
-      epoch=epoch, num_env_steps=num_env_steps
+      env_fn,
+      train_hparams,
+      simulated=False,
+      save_continuously=False,
+      epoch=epoch,
+      sampling_temp=hparams.real_sampling_temp,
+      num_env_steps=num_env_steps,
   )
   # Save unfinished rollouts to history.
   env.reset()
@@ -492,7 +497,8 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
       if report_fn:
         if report_metric == "mean_reward":
           metric_name = rl_utils.get_metric_name(
-              stochastic=True, max_num_noops=hparams.eval_max_num_noops,
+              sampling_temp=hparams.eval_sampling_temps[0],
+              max_num_noops=hparams.eval_max_num_noops,
               clipped=False
           )
           report_fn(eval_metrics[metric_name], epoch)
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index bb29767d9..1459755fd 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -77,8 +77,12 @@ def _rlmb_base():
       # In your experiments, you want to optimize this rate to your schedule.
       learning_rate_bump=3.0,
 
-      # Batch size during evaluation. Metrics are averaged over this number of
-      # rollouts.
+      # Policy sampling temperature to use when gathering data from the real
+      # environment.
+      real_sampling_temp=1.0,
+
+      # Sampling temperatures to try during eval.
+      eval_sampling_temps=[0.0, 0.2, 0.5, 0.8, 1.0, 2.0],
       eval_max_num_noops=8,
 
       game="pong",
@@ -106,6 +110,8 @@ def _rlmb_base():
       real_batch_size=-1,
       # Number of simulated environments to train on simultaneously.
       simulated_batch_size=-1,
+      # Batch size during evaluation. Metrics are averaged over this number of
+      # rollouts.
       eval_batch_size=-1,
   )
 
@@ -467,6 +473,7 @@ def _rlmb_tiny_overrides():
       wm_eval_rollout_ratios=[1],
       env_timesteps_limit=7,
       simulated_rollout_length=2,
+      eval_sampling_temps=[0.0, 1.0],
   )
 
 
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index dea16b1fa..6948b0abf 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -78,7 +78,8 @@ def train(hparams, output_dir, report_fn=None):
   if not steps or steps[-1] < eval_every_epochs:
     steps.append(eval_every_epochs)
   metric_name = rl_utils.get_metric_name(
-      stochastic=True, max_num_noops=hparams.eval_max_num_noops,
+      sampling_temp=hparams.eval_sampling_temps[0],
+      max_num_noops=hparams.eval_max_num_noops,
       clipped=False
   )
   for step in steps:
diff --git a/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index be3300312..678641847 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -31,8 +31,9 @@ class TrainTest(tf.test.TestCase):
   def test_train_pong(self):
     hparams = registry.hparams("mfrl_original")
     hparams.batch_size = 2
-    hparams.ppo_epochs_num = 2
-    hparams.ppo_epoch_length = 3
+    hparams.eval_sampling_temps = [0.0, 1.0]
+    hparams.add_hparam("ppo_epochs_num", 2)
+    hparams.add_hparam("ppo_epoch_length", 3)
     FLAGS.output_dir = tf.test.get_temp_dir()
     trainer_model_free.train(hparams, FLAGS.output_dir)
 

From e14a535266bfbe46759a83193dd997644b03930d Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Wed, 5 Dec 2018 17:31:04 -0800
Subject: [PATCH 1334/2720] internal merge of PR #1277

PiperOrigin-RevId: 224250704
---
 tensor2tensor/models/research/rl.py    | 10 +---------
 tensor2tensor/rl/dopamine_connector.py | 21 ++++++++++-----------
 tensor2tensor/rl/policy_learner.py     |  1 +
 3 files changed, 12 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index c596d4df5..f51fe27bd 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -61,14 +61,6 @@ def ppo_base_v1():
   return hparams
 
 
-#@registry.register_hparams
-#def ppo_continuous_action_base():
-#  hparams = ppo_base_v1()
-#  hparams.add_hparam("policy_network", feed_forward_gaussian_fun)
-#  hparams.add_hparam("policy_network_params", "basic_policy_parameters")
-#  return hparams
-
-
 @registry.register_hparams
 def basic_policy_parameters():
   wrappers = None
@@ -158,7 +150,7 @@ def get_policy(observations, hparams, action_space):
   """Get a policy network.
 
   Args:
-    observations
+    observations: observations
     hparams: parameters
     action_space: action space
 
diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 0f227cd9f..d18789e0e 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -19,10 +19,9 @@
 from __future__ import division
 from __future__ import print_function
 
-from copy import copy
+import copy
+
 from dopamine.agents.dqn import dqn_agent
-from dopamine.agents.dqn.dqn_agent import NATURE_DQN_OBSERVATION_SHAPE
-from dopamine.agents.dqn.dqn_agent import NATURE_DQN_STACK_SIZE
 from dopamine.atari import run_experiment
 from dopamine.replay_memory import circular_replay_buffer
 from dopamine.replay_memory.circular_replay_buffer import OutOfGraphReplayBuffer
@@ -45,7 +44,7 @@
 
 
 class ResizeObservation(gym.ObservationWrapper):
-  """ TODO(konradczechowski): Add doc-string."""
+  """TODO(konradczechowski): Add doc-string."""
 
   def __init__(self, env, size=84):
     """Based on WarpFrame from openai baselines atari_wrappers.py.
@@ -91,7 +90,7 @@ def step(self, action):
 
 
 class _DQNAgent(dqn_agent.DQNAgent):
-  """ Modify dopamine DQNAgent to match our needs.
+  """Modify dopamine DQNAgent to match our needs.
 
   Allow passing batch_size and replay_capacity to ReplayBuffer, allow not using
   (some of) terminal episode transitions in training.
@@ -107,8 +106,8 @@ def __init__(self, replay_capacity, batch_size, generates_trainable_dones,
   def _build_replay_buffer(self, use_staging):
     """Build WrappedReplayBuffer with custom OutOfGraphReplayBuffer."""
     replay_buffer_kwargs = dict(
-        observation_shape=NATURE_DQN_OBSERVATION_SHAPE,
-        stack_size=NATURE_DQN_STACK_SIZE,
+        observation_shape=dqn_agent.NATURE_DQN_OBSERVATION_SHAPE,
+        stack_size=dqn_agent.NATURE_DQN_STACK_SIZE,
         replay_capacity=self._replay_capacity,
         batch_size=self._batch_size,
         update_horizon=self.update_horizon,
@@ -127,7 +126,7 @@ def _build_replay_buffer(self, use_staging):
 
 
 class _OutOfGraphReplayBuffer(OutOfGraphReplayBuffer):
-  """ Replay not sampling artificial_terminal transition.
+  """Replay not sampling artificial_terminal transition.
 
   Adds to stored tuples 'artificial_done' field (as last ReplayElement).
   When sampling, ignores tuples for which artificial_done is True.
@@ -238,7 +237,7 @@ def _get_optimizer(params):
 
 
 class DQNLearner(PolicyLearner):
-  """ Interface for learning dqn implemented in dopamine."""
+  """Interface for learning dqn implemented in dopamine."""
 
   def __init__(self, frame_stack_size, base_event_dir, agent_model_dir):
     super(DQNLearner, self).__init__(frame_stack_size, base_event_dir,
@@ -296,7 +295,7 @@ def train(self,
     if num_env_steps is None:
       num_env_steps = hparams.num_frames
 
-    hparams = copy(hparams)
+    hparams = copy.copy(hparams)
     hparams.set_hparam(
         "agent_epsilon_eval", min(hparams.agent_epsilon_eval * sampling_temp, 1)
     )
@@ -318,7 +317,7 @@ def evaluate(self, env_fn, hparams, sampling_temp):
     target_iterations = 0
     training_steps_per_iteration = 0
 
-    hparams = copy(hparams)
+    hparams = copy.copy(hparams)
     hparams.set_hparam(
         "agent_epsilon_eval", min(hparams.agent_epsilon_eval * sampling_temp, 1)
     )
diff --git a/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
index 8368b714d..b34af8c8c 100644
--- a/tensor2tensor/rl/policy_learner.py
+++ b/tensor2tensor/rl/policy_learner.py
@@ -41,6 +41,7 @@ def train(
       eval_env_fn=None,
       report_fn=None
   ):
+    """Train."""
     # TODO(konradczechowski): pass name_scope instead of epoch?
     # TODO(konradczechowski): move 'simulated' to  batch_env
     raise NotImplementedError()

From aa9e6b4960847e4a71b04d56ebaa67275f694395 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 6 Dec 2018 10:36:14 -0800
Subject: [PATCH 1335/2720] Minor cleanup. Moving one gym related function into
 its own file and added test. Unify naming across gym's method and ours and
 add "rl_env_" to the hparam, to give the hint where this gets used.

PiperOrigin-RevId: 224364129
---
 tensor2tensor/data_generators/gym_env.py      | 18 ++----
 tensor2tensor/data_generators/gym_env_test.py |  2 +-
 tensor2tensor/models/research/rl.py           |  2 +-
 tensor2tensor/rl/gym_utils.py                 | 56 +++++++++++++++++++
 tensor2tensor/rl/gym_utils_test.py            | 50 +++++++++++++++++
 tensor2tensor/rl/rl_utils.py                  |  2 +-
 .../rl/trainer_model_based_params.py          |  4 +-
 7 files changed, 115 insertions(+), 19 deletions(-)
 create mode 100644 tensor2tensor/rl/gym_utils.py
 create mode 100644 tensor2tensor/rl/gym_utils_test.py

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index b354b91d9..055a6a952 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -22,7 +22,6 @@
 import collections
 import itertools
 import random
-import gym
 from gym.spaces import Box
 import numpy as np
 
@@ -30,6 +29,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.layers import modalities
+from tensor2tensor.rl import gym_utils
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -79,17 +79,6 @@ def __deepcopy__(self, memo):
     return self
 
 
-def make_gym_env(name, timesteps_limit=-1):
-  env = gym.make(name)
-  if timesteps_limit != -1:
-    # Replace TimeLimit Wrapper with one of proper time step limit.
-    if isinstance(env, gym.wrappers.TimeLimit):
-      env = env.env
-    env = gym.wrappers.TimeLimit(env,
-                                 max_episode_steps=timesteps_limit)
-  return env
-
-
 class EnvSimulationProblem(video_utils.VideoProblem):
   """Base Problem class for use with world models.
 
@@ -587,7 +576,7 @@ class T2TGymEnv(T2TEnv):
 
   def __init__(self, base_env_name=None, batch_size=1, grayscale=False,
                resize_height_factor=2, resize_width_factor=2,
-               base_env_timesteps_limit=-1, max_num_noops=0, **kwargs):
+               rl_env_max_episode_steps=-1, max_num_noops=0, **kwargs):
     if base_env_name is None:
       base_env_name = self.base_env_name
     self._base_env_name = base_env_name
@@ -599,7 +588,8 @@ def __init__(self, base_env_name=None, batch_size=1, grayscale=False,
       # Set problem name if not registered.
       self.name = "Gym%s" % base_env_name
 
-    self._envs = [make_gym_env(base_env_name, base_env_timesteps_limit)
+    self._envs = [gym_utils.make_gym_env(
+        base_env_name, rl_env_max_episode_steps=rl_env_max_episode_steps)
                   for _ in range(self.batch_size)]
 
     # max_num_noops works only with atari envs.
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index e5532f8c0..2d5859b1b 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -29,7 +29,7 @@
 
 from tensor2tensor.data_generators import gym_env
 from tensor2tensor.data_generators import problem
-from tensor2tensor.data_generators.gym_env import make_gym_env
+from tensor2tensor.rl.gym_utils import make_gym_env
 
 import tensorflow as tf
 
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index f51fe27bd..2a3b37302 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -236,7 +236,7 @@ def mfrl_original():
       resize_height_factor=2,
       resize_width_factor=2,
       grayscale=0,
-      env_timesteps_limit=-1,
+      rl_env_max_episode_steps=-1,
   )
 
 
diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
new file mode 100644
index 000000000..c2cc7d505
--- /dev/null
+++ b/tensor2tensor/rl/gym_utils.py
@@ -0,0 +1,56 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Utilities for interacting with Gym classes."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gym
+
+
+def make_gym_env(name, rl_env_max_episode_steps=-1):
+  """Create a gym env optionally wrapped with a time limit wrapper.
+
+  NOTE: The returned env may already be wrapped with TimeLimit!
+
+  Args:
+    name: `str` - base name of the gym env to make.
+    rl_env_max_episode_steps: `int` or None - Using any value < 0 returns the
+      env as-in, otherwise we impose the requested timelimit. Setting this to
+      None returns a wrapped env that doesn't have a step limit.
+
+  Returns:
+    An instance of `gym.Env` or `gym.wrappers.TimeLimit` with the requested
+    step limit.
+  """
+
+  # rl_env_max_episode_steps is None or int.
+  assert ((not rl_env_max_episode_steps) or
+          isinstance(rl_env_max_episode_steps, int))
+
+  env = gym.make(name)
+
+  # If nothing to do, then return the env.
+  if rl_env_max_episode_steps and rl_env_max_episode_steps < 0:
+    return env
+
+  # Sometimes (mostly?) the env is already wrapped in a TimeLimit wrapper, in
+  # which case unwrap it and wrap with the proper time limit requested.
+  if isinstance(env, gym.wrappers.TimeLimit):
+    env = env.env
+
+  return gym.wrappers.TimeLimit(env, max_episode_steps=rl_env_max_episode_steps)
diff --git a/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
new file mode 100644
index 000000000..9beffa7a1
--- /dev/null
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -0,0 +1,50 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.rl.gym_utils."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gym
+from tensor2tensor.rl import gym_utils
+import tensorflow as tf
+
+
+class GymUtilsTest(tf.test.TestCase):
+
+  # Just make an environment and expect to get one.
+  def test_making_simple_env(self):
+    env = gym_utils.make_gym_env("CartPole-v0")
+    self.assertTrue(isinstance(env, gym.Env))
+
+  # Make a time-wrapped environment and expect to get one.
+  def test_making_timewrapped_env(self):
+    env = gym_utils.make_gym_env("CartPole-v0", rl_env_max_episode_steps=1000)
+    self.assertTrue(isinstance(env, gym.Env))
+    self.assertTrue(isinstance(env, gym.wrappers.TimeLimit))
+    self.assertEquals(1000, env._max_episode_steps)
+
+  # Make a time-wrapped environment with unlimited limit.
+  def test_unlimited_env(self):
+    env = gym_utils.make_gym_env("CartPole-v0", rl_env_max_episode_steps=None)
+    self.assertTrue(isinstance(env, gym.Env))
+    self.assertTrue(isinstance(env, gym.wrappers.TimeLimit))
+    self.assertTrue(env._max_episode_steps is None)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index f4acecd73..460b1888d 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -115,7 +115,7 @@ def setup_env(hparams, batch_size, max_num_noops):
                   grayscale=hparams.grayscale,
                   resize_width_factor=hparams.resize_width_factor,
                   resize_height_factor=hparams.resize_height_factor,
-                  base_env_timesteps_limit=hparams.env_timesteps_limit,
+                  rl_env_max_episode_steps=hparams.rl_env_max_episode_steps,
                   max_num_noops=max_num_noops)
   return env
 
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 1459755fd..4e6f97f82 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -96,7 +96,7 @@ def _rlmb_base():
       # Ratios of ppo_epoch_length to report reward_accuracy on.
       wm_eval_rollout_ratios=[0.25, 0.5, 1, 2],
       stop_loop_early=False,  # To speed-up tests.
-      env_timesteps_limit=-1,  # Use default from gym.make()
+      rl_env_max_episode_steps=-1,  # Use default from gym.make()
       # Number of last observations to feed to the agent and world model.
       frame_stack_size=4,
       # This is only used for world-model evaluation currently, PolicyLearner
@@ -471,7 +471,7 @@ def _rlmb_tiny_overrides():
       resize_height_factor=2,
       resize_width_factor=2,
       wm_eval_rollout_ratios=[1],
-      env_timesteps_limit=7,
+      rl_env_max_episode_steps=7,
       simulated_rollout_length=2,
       eval_sampling_temps=[0.0, 1.0],
   )

From 589df9423da982c775538ff9dcd270b6218d4ba5 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Thu, 6 Dec 2018 11:02:23 -0800
Subject: [PATCH 1336/2720] New mesh_tensorflow transformer now handles
 encoder-decoder models and beam search. Fix bug in variable stacking code
 (operations in while loops did not get their inputs redirected) Fix datatype
 bug in simd_mesh_impl.py

PiperOrigin-RevId: 224369334
---
 tensor2tensor/models/mtf_transformer2.py | 467 +++++++++++++++--------
 1 file changed, 314 insertions(+), 153 deletions(-)

diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index d90c65b14..3947b59cf 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -19,7 +19,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import copy
 import mesh_tensorflow as mtf
 from mesh_tensorflow.transformer import moe
 from mesh_tensorflow.transformer import transformer
@@ -34,8 +33,12 @@
 
 
 @registry.register_model
-class MtfTransformer2(mtf_model.MtfModel):
-  """Transformer in mesh_tensorflow."""
+class MtfUnitransformer(mtf_model.MtfModel):
+  """Single-stack Transformer (Transformer Decoder) in mesh_tensorflow.
+
+  Can optionally be autoregressive (language generation) or non-autoregressive
+  like BERT.
+  """
 
   @property
   def batch_dims(self):
@@ -51,6 +54,10 @@ def batch_dims(self):
           mtf.Dimension("inner_batch",
                         hparams.batch_size // hparams.outer_batch_size)]
 
+  @property
+  def autoregressive(self):
+    return self._hparams.autoregressive
+
   @property
   def variable_dtype(self):
     return mtf.VariableDType(
@@ -68,73 +75,83 @@ def _import_to_batch_by_length(self, x, name, mesh):
     x = tf.reshape(x, mtf_shape.to_integer_list)
     return mtf.import_fully_replicated(mesh, x, mtf_shape, name=name)
 
-  def _import_to_batch_by_decode_length(self, x, name, mesh):
+  def _import_feature(self, features, mesh, key):
+    """Import a feature from the features dictionary into a mtf.Tensor.
+
+    Args:
+      features: a features dictionary
+      mesh: a Mesh
+      key: a string
+
+    Returns:
+      a mtf.Tensor with dtype int32 and shape self.batch_dims + self.length_dim
+    """
+    if key not in features:
+      return None
+    x = tf.to_int32(features[key])
+    x = common_layers.expand_squeeze_to_nd(x, 2)
+    # pad to length
+    extra_length = self.length_dim.size - tf.shape(x)[1]
+    x = tf.pad(x, [[0, 0], [0, extra_length]])
     mtf_shape = mtf.Shape(self.batch_dims + [self.length_dim])
     x = tf.reshape(x, mtf_shape.to_integer_list)
-    return mtf.import_fully_replicated(mesh, x, mtf_shape, name=name)
+    return mtf.import_fully_replicated(mesh, x, mtf_shape, name=key)
 
   def model(self):
     hparams = self._hparams
+    if hparams.label_smoothing != 0:
+      raise NotImplementedError(
+          "Label smoothing not implemented in unitransformer."
+          "  Do you really want it?")
     if isinstance(hparams.layer_stack, transformer.LayerStack):
       layer_stack = hparams.layer_stack
     else:
       # hparams.layer_stack is a function for creating a LayerStack
       layer_stack = hparams.layer_stack(hparams)
-    return transformer.Transformer(
+    if self.autoregressive:
+      input_vocab_size = self._inputs_vocab_size
+    else:
+      input_vocab_size = self._targets_vocab_size
+    return transformer.Unitransformer(
         layer_stack=layer_stack,
         d_model=hparams.d_model,
-        input_vocab_size=self._targets_vocab_size,
+        input_vocab_size=input_vocab_size,
         output_vocab_size=self._targets_vocab_size,
-        autoregressive=hparams.decoder_type == "autoregressive",
+        autoregressive=self.autoregressive,
         max_length=hparams.max_length)
 
   def _mtf_model_fn(self, features, mesh):
     self._original_features = features
-    features = copy.copy(features)
     hparams = self._hparams
-    targets = tf.to_int32(features["targets"])
-    if len(targets.get_shape()) > 2:
-      tf.logging.info("targets = %s" % targets)
-      targets = tf.squeeze(targets, [2, 3])
-    # pad targets to max_length
-    def pad_to_length(x):
-      extra_length = self.length_dim.size - tf.shape(x)[1]
-      x = tf.pad(x, [[0, 0], [0, extra_length]])
-      x = tf.reshape(x, [hparams.batch_size, self.length_dim.size])
-      return x
-    targets = pad_to_length(targets)
-    targets = self._import_to_batch_by_length(targets, "targets", mesh)
-    for key in ["targets_segmentation", "targets_position",
-                "inputs_segmentation", "inputs_position"]:
-      if key in features:
-        features[key] = pad_to_length(features[key])
-    if hparams.decoder_type == "autoregressive":
-      shifted_targets = mtf.shift(
+    def import_feature(key):
+      return self._import_feature(features, mesh, key)
+    targets = import_feature("targets")
+    if self.autoregressive:
+      inputs = mtf.shift(
           targets, offset=1, dim=self.length_dim, wrap=False)
     else:
-      raise ValueError(
-          "unknown hparams.decoder_type = %s" % hparams.decoder_type)
+      inputs = import_feature("inputs")
+      # TODO(noam): options for bert-style masking here?
+    sequence_id = import_feature("targets_segmentation")
     model = self.model()
     logits, loss = model.call_simple(
-        inputs=shifted_targets,
+        inputs=inputs,
         targets=targets,
         compute_loss=True,
         mode=hparams.mode,
-        variable_dtype=self.variable_dtype)
-    # mesh_shape=hparams.mesh_shape,
-    # layout=hparams.layout,
+        variable_dtype=self.variable_dtype,
+        sequence_id=sequence_id)
     return logits, loss
 
   def mtf_model_fn(self, features, mesh):
-    with tf.variable_scope("transformer"):
-      logits, loss = self._mtf_model_fn(features, mesh)
-      # combine batch dims
-      if len(self.batch_dims) > 1:
-        combined_batch_dim = mtf.Dimension(
-            self.batch_dims[0].name, mtf.Shape(self.batch_dims).size)
-        logits = mtf.reshape(
-            logits, [combined_batch_dim] + logits.shape.dims[-2:])
-      return logits, loss
+    logits, loss = self._mtf_model_fn(features, mesh)
+    # combine batch dims
+    if len(self.batch_dims) > 1:
+      combined_batch_dim = mtf.Dimension(
+          self.batch_dims[0].name, mtf.Shape(self.batch_dims).size)
+      logits = mtf.reshape(
+          logits, [combined_batch_dim] + logits.shape.dims[-2:])
+    return logits, loss
 
   @property
   def _targets_vocab_size(self):
@@ -153,74 +170,137 @@ def _inputs_vocab_size(self):
   def sample(self, features, mesh):
     hparams = self._hparams
     model = self.model()
-    # Prepare partial targets.
-    # In either features["inputs"] or features["targets"].
-    # We force the outputs to begin with these sequences.
-    partial_targets = features.get("inputs", None)
-    if partial_targets is None:
-      partial_targets = features.get("targets", None)
-    if partial_targets is not None:
-      partial_targets = common_layers.expand_squeeze_to_nd(partial_targets, 2)
-      partial_targets = tf.to_int32(partial_targets)
-      partial_targets_batch = tf.shape(partial_targets)[0]
-      partial_targets_length = tf.shape(partial_targets)[1]
-      partial_targets = tf.pad(
-          partial_targets, [[0, hparams.batch_size - partial_targets_batch],
-                            [0, self.length_dim.size - partial_targets_length]])
-      partial_targets = self._import_to_batch_by_length(
-          partial_targets, "partial_targets", mesh)
-      # strip EOS
-      partial_targets *= mtf.to_int32(mtf.not_equal(partial_targets, 1))
+    def import_feature(key):
+      return self._import_feature(features, mesh, key)
+
+    if self.autoregressive:
+      # Prepare partial targets.
+      # In either features["inputs"] or features["targets"].
+      # We force the outputs to begin with these sequences.
+      partial_targets = import_feature("inputs")
+      if partial_targets is None:
+        partial_targets = import_feature("targets")
+      if partial_targets is None:
+        ids_shape = mtf.Shape(self.batch_dims + [self.length_dim])
+        partial_targets = mtf.constant(mesh, 0, ids_shape, dtype=tf.int32)
+      if hparams.beam_size > 1:
+        raise NotImplementedError(
+            "Beam search not implemented for unitransformer.")
+      return model.sample_autoregressive(
+          partial_targets,
+          temperature=hparams.sampling_temp,
+          variable_dtype=self.variable_dtype)
+    else:
+      raise ValueError(
+          "Don't know how to sample from non-autoregressive unitransformer")
+
+
+@registry.register_model
+class MtfBitransformer(MtfUnitransformer):
+  """Encoder-Decoder Transformer in mesh_tensorflow."""
 
+  def model(self):
+    hparams = self._hparams
+    if isinstance(hparams.encoder_layer_stack, transformer.LayerStack):
+      encoder_layer_stack = hparams.encoder_layer_stack
     else:
-      ids_shape = mtf.Shape(self.batch_dims + [self.length_dim])
-      partial_targets = mtf.constant(mesh, 0, ids_shape, dtype=tf.int32)
-    if hparams.beam_size == 1:
-      pass
+      encoder_layer_stack = hparams.encoder_layer_stack(hparams)
+    if isinstance(hparams.decoder_layer_stack, transformer.LayerStack):
+      decoder_layer_stack = hparams.decoder_layer_stack
     else:
-      raise NotImplementedError("not implemented")
-      # beam_dim = mtf.Dimension("beam", hparams.beam_size)
-      # ids_shape = mtf.Shape(self.batch_dims + [beam_dim, self.length_dim])
+      decoder_layer_stack = hparams.decoder_layer_stack(hparams)
+    return transformer.Bitransformer(
+        encoder_layer_stack=encoder_layer_stack,
+        decoder_layer_stack=decoder_layer_stack,
+        encoder_d_model=hparams.d_model,
+        decoder_d_model=hparams.d_model,
+        input_vocab_size=self._inputs_vocab_size,
+        output_vocab_size=self._targets_vocab_size,
+        max_length=hparams.max_length,
+        shared_embedding=hparams.shared_embedding,
+        label_smoothing=hparams.label_smoothing)
+
+  def _mtf_model_fn(self, features, mesh):
+    self._original_features = features
+    hparams = self._hparams
+    def import_feature(key):
+      return self._import_feature(features, mesh, key)
+    targets = import_feature("targets")
+    inputs = import_feature("inputs")
+    encoder_sequence_id = import_feature("inputs_segmentation")
+    if not encoder_sequence_id:
+      encoder_sequence_id = mtf.to_int32(mtf.not_equal(inputs, 0))
+    decoder_sequence_id = import_feature("targets_segmentation")
+    if decoder_sequence_id is None:
+      decoder_sequence_id = mtf.to_int32(mtf.not_equal(targets, 0))
+    model = self.model()
+    logits, loss = model.call_simple(
+        inputs=inputs,
+        targets=targets,
+        compute_loss=True,
+        mode=hparams.mode,
+        variable_dtype=self.variable_dtype,
+        encoder_sequence_id=encoder_sequence_id,
+        decoder_sequence_id=decoder_sequence_id)
+    return logits, loss
 
-    partial_targets = mtf.Print(partial_targets, [partial_targets],
-                                "Partial_Targets", summarize=1000)
-    return model.sample_autoregressive(
-        partial_targets,
-        temperature=hparams.sampling_temp,
-        variable_dtype=self.variable_dtype)
+  def sample(self, features, mesh):
+    hparams = self._hparams
+    model = self.model()
+    inputs = self._import_feature(features, mesh, "inputs")
+    return model.decode(
+        inputs,
+        self.variable_dtype,
+        beam_size=hparams.beam_size,
+        alpha=hparams.alpha,
+        temperature=hparams.sampling_temp if hparams.beam_size == 1 else 0,
+        decode_length_multiplier=hparams.decode_length_multiplier,
+        decode_length_constant=hparams.decode_length_constant)
+
+
+def default_layer_stack(hparams):
+  return transformer.LayerStack(
+      [transformer_layers.SelfAttention(
+          num_heads=hparams.num_heads,
+          key_value_size=hparams.d_kv,
+          dropout_rate=hparams.attention_dropout),
+       transformer_layers.DenseReluDense(
+           hidden_size=hparams.d_ff,
+           dropout_rate=hparams.relu_dropout),
+      ] * hparams.num_hidden_layers,
+      dropout_rate=hparams.layer_prepostprocess_dropout,
+      norm_epsilon=hparams.norm_epsilon)
+
+
+def default_layer_stack_with_encoder_attention(hparams):
+  return transformer.LayerStack(
+      [transformer_layers.SelfAttention(
+          num_heads=hparams.num_heads,
+          key_value_size=hparams.d_kv,
+          dropout_rate=hparams.attention_dropout),
+       transformer_layers.EncDecAttention(
+           num_heads=hparams.num_heads,
+           key_value_size=hparams.d_kv,
+           dropout_rate=hparams.attention_dropout),
+       transformer_layers.DenseReluDense(
+           hidden_size=hparams.d_ff,
+           dropout_rate=hparams.relu_dropout),
+      ] * hparams.num_hidden_layers,
+      dropout_rate=hparams.layer_prepostprocess_dropout,
+      norm_epsilon=hparams.norm_epsilon)
 
 
-@registry.register_hparams
 def mtf_transformer2_base():
   """Set of hyperparameters."""
   hparams = common_hparams.basic_params1()
-  hparams.no_data_parallelism = True
-  hparams.use_fixed_batch_size = True
-  hparams.add_hparam("mtf_mode", True)
+
+  hparams.add_hparam("d_model", 1024)
   hparams.batch_size = 4
   hparams.max_length = 1024
-  hparams.add_hparam("d_model", 1024)
   hparams.label_smoothing = 0.0
-  # 8-way model-parallelism
-  hparams.add_hparam("mesh_shape", "model:8")
-  hparams.add_hparam("layout", "batch:batch;vocab:model;d_ff:model;heads:model")
 
-  # hparams.layer_stack should be either a transformer.LayerStack or a function
-  # from hparams to transformer.LayerStack
-  def my_layer_stack(hparams):
-    return transformer.LayerStack(
-        [transformer_layers.SelfAttention(
-            num_heads=hparams.num_heads,
-            key_value_size=hparams.d_kv,
-            dropout_rate=hparams.attention_dropout),
-         transformer_layers.DenseReluDense(
-             hidden_size=hparams.d_ff,
-             dropout_rate=hparams.layer_prepostprocess_dropout),
-        ] * hparams.num_hidden_layers)
-  hparams.layer_stack = my_layer_stack
-
-  # These hyperparameters are used in the above default layer_stack function.
-  # They may not be respected if hparams.layer_stack is changed.
+  # These hyperparameters are used in default_layer_stack()
+  # They may not be respected if hparams uses a differet layer stack function.
   hparams.num_hidden_layers = 6
   hparams.add_hparam("d_ff", 2048)
   hparams.add_hparam("d_kv", 128)
@@ -228,20 +308,6 @@ def my_layer_stack(hparams):
   hparams.add_hparam("relu_dropout", 0.0)
   hparams.layer_prepostprocess_dropout = 0.0
 
-  # Describes what model architecture:
-  #   "encdec": encoder + autoregressive decoder
-  #   "decoder": single-stack autoregressive sequence model.
-  #   "encoder": single-stack non-autoregressive model
-  #      with equal-length inputs and outputs.
-  # TODO(noam): implement different types of transformers.
-  hparams.add_hparam("transformer_type", "decoder")
-
-  # What does the decoder do:
-  #   "autoregressive": Decoder left to right
-  #   "denoising": Fills in masked-out values simultaneously
-  # TODO(noam): only autoregressive is implemented so far.
-  hparams.add_hparam("decoder_type", "autoregressive")
-
   # round up vocab sizes to be a multiple of this value
   hparams.vocab_divisor = 128
 
@@ -252,20 +318,9 @@ def my_layer_stack(hparams):
   hparams.add_hparam("slice_dtype", "float32")
   hparams.activation_dtype = "bfloat16"
 
-  # These parameters make Transformer model compatible with MtfTransformer2
-  # Do not override these, as mtf_transformer does not support other options.
-  hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
-  hparams.modality = {
-      "inputs": modalities.IdentitySymbolModality,
-      "targets": modalities.IdentitySymbolModality,
-  }
-
-  # Parameters for computing the maximum decode length in beam search.
-  # Maximum decode length is:
-  #    min(max_length,
-  #        decode_length_multiplier * input_length + decode_length_constant)
-  hparams.add_hparam("decode_length_multiplier", 1.5)
-  hparams.add_hparam("decode_length_constant", 10.0)
+  # 8-way model-parallelism
+  hparams.add_hparam("mesh_shape", "model:8")
+  hparams.add_hparam("layout", "batch:batch;vocab:model;d_ff:model;heads:model")
 
   # If nonzero, we split the batch across two tensor-dimensions named
   # "outer_batch" and "inner_batch", allowing for splitting across two mesh
@@ -274,16 +329,51 @@ def my_layer_stack(hparams):
   # hparams.batch_size // hparams.outer_batch_size.
   hparams.add_hparam("outer_batch_size", 0)
 
+  hparams.shared_embedding_and_softmax_weights = False
   # length for training or decoding - defaults to max_length
   hparams.add_hparam("length", 0)
 
-  hparams.sampling_method = "random"
+  # These parameters make Transformer model compatible with mtf
+  # Do not override these.
+  hparams.no_data_parallelism = True
+  hparams.use_fixed_batch_size = True
+  hparams.add_hparam("mtf_mode", True)
+  hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
+  hparams.modality = {
+      "inputs": modalities.IdentitySymbolModality,
+      "targets": modalities.IdentitySymbolModality,
+  }
   return hparams
 
 
 @registry.register_hparams
-def mtf_transformer2_tiny():
+def mtf_unitransformer_base():
   hparams = mtf_transformer2_base()
+  hparams.add_hparam("autoregressive", True)
+  hparams.layer_stack = default_layer_stack
+  return hparams
+
+
+@registry.register_hparams
+def mtf_bitransformer_base():
+  """Machine translation base configuration."""
+  hparams = mtf_transformer2_base()
+  hparams.max_length = 256
+  hparams.shared_embedding = True
+  hparams.encoder_layer_stack = default_layer_stack
+  hparams.decoder_layer_stack = default_layer_stack_with_encoder_attention
+  # Parameters for computing the maximum decode length in beam search.
+  # Maximum decode length is:
+  #    min(max_length,
+  #        decode_length_multiplier * input_length + decode_length_constant)
+  hparams.add_hparam("decode_length_multiplier", 1.5)
+  hparams.add_hparam("decode_length_constant", 10.0)
+  return hparams
+
+
+@registry.register_hparams
+def mtf_unitransformer_tiny():
+  hparams = mtf_unitransformer_base()
   hparams.batch_size = 2
   hparams.mesh_shape = ""
   hparams.d_model = 128
@@ -294,12 +384,21 @@ def mtf_transformer2_tiny():
 
 
 @registry.register_hparams
-def mtf_transformer2_all_layers_tiny():
-  """Test out all the layers on local CPU."""
-  hparams = mtf_transformer2_base()
+def mtf_bitransformer_tiny():
+  hparams = mtf_bitransformer_base()
   hparams.batch_size = 2
   hparams.mesh_shape = ""
   hparams.d_model = 128
+  hparams.num_hidden_layers = 2
+  hparams.num_heads = 4
+  hparams.d_ff = 512
+  return hparams
+
+
+@registry.register_hparams
+def mtf_unitransformer_all_layers_tiny():
+  """Test out all the layers on local CPU."""
+  hparams = mtf_unitransformer_tiny()
   hparams.layer_stack = transformer.LayerStack(
       [transformer_layers.SelfAttention(num_heads=4),
        transformer_layers.LocalSelfAttention(num_heads=4),
@@ -310,25 +409,15 @@ def mtf_transformer2_all_layers_tiny():
 
 
 @registry.register_hparams
-def mtr2_lm_dense(sz):
-  """Series of architectural experiments on language modeling.
-
-  Larger models than the ones above.
-
-  All models are trained on sequences of 1024 tokens.
+def mtr_lm_dense(sz):
+  """Series of architectures for language modeling.
 
   We assume infinite training data, so no dropout necessary.
-  We process 2^36 tokens in training = 524288 steps at batch size 128
 
+  You can use languagemodel_wiki_noref_v32k_l1k.
+  (1 epoch = ~46000 steps).
   TODO(noam): find a large enough dataset for these experiments.
 
-  You can use languagemodel_wiki_noref_v32k_l1k, but this is too small,
-  (1 epoch = ~46000 steps) so training will cover about 11 epochs.
-
-  Note: configurations and code are likely to change without notice.
-
-  Run on TPU 4x4 for 524288 steps unless otherwise indicated.
-
   Args:
     sz: an integer
 
@@ -336,7 +425,7 @@ def mtr2_lm_dense(sz):
     a hparams
   """
   n = 2 ** sz
-  hparams = mtf_transformer2_base()
+  hparams = mtf_unitransformer_base()
   hparams.d_model = 1024
   hparams.max_length = 1024
   hparams.batch_size = 128
@@ -352,22 +441,31 @@ def mtr2_lm_dense(sz):
 
 
 @registry.register_hparams
-def mtr2_lm_dense_0():
-  return mtr2_lm_dense(0)
+def mtr_lm_dense_0():
+  return mtr_lm_dense(0)
+
+
+@registry.register_hparams
+def mtr_lm_dense_1():
+  return mtr_lm_dense(1)
 
 
 @registry.register_hparams
-def mtr2_lm_dense_1():
-  return mtr2_lm_dense(1)
+def mtr_lm_dense_2():
+  hparams = mtr_lm_dense(2)
+  hparams.mesh_shape = "model:4;batch:8"
+  return hparams
 
 
 @registry.register_hparams
-def mtr2_lm_dense_2():
-  return mtr2_lm_dense(2)
+def mtr_lm_dense_3():
+  hparams = mtr_lm_dense(3)
+  hparams.mesh_shape = "model:4;batch:8"
+  return hparams
 
 
 @registry.register_hparams
-def mtr2_v1():
+def mtr_lm_v1():
   """Model incorporating mixture-of-experts, local and global attention.
 
   ~6B parameters
@@ -377,7 +475,7 @@ def mtr2_v1():
   Returns:
     a hparams
   """
-  hparams = mtr2_lm_dense(0)
+  hparams = mtr_lm_dense(0)
   local_att = transformer_layers.LocalSelfAttention(
       num_heads=4, key_value_size=128)
   att = transformer_layers.SelfAttention(num_heads=4, key_value_size=128)
@@ -390,3 +488,66 @@ def mtr2_v1():
   hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
   hparams.outer_batch_size = 4
   return hparams
+
+
+@registry.register_hparams
+def mtr_tr_dense(sz):
+  """Series of machine translation models.
+
+  All models are trained on sequences of 256 tokens.
+
+  You can use the dataset translate_enfr_wmt32k_packed.
+  154000 steps = 3 epochs.
+
+  Args:
+    sz: an integer
+
+  Returns:
+    a hparams
+  """
+  n = 2 ** sz
+  hparams = mtf_bitransformer_base()
+  hparams.d_model = 1024
+  hparams.max_length = 256
+  hparams.batch_size = 128
+  # Parameters for my_layer_stack()
+  hparams.num_hidden_layers = 6
+  hparams.d_ff = int(4096 * n)
+  hparams.d_kv = 128
+  hparams.num_heads = int(8 * n)
+  # one epoch for translate_enfr_wmt32k_packed = 51400 steps
+  hparams.learning_rate_decay_steps = 51400
+  hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
+  hparams.mesh_shape = "model:4;batch:8"
+  hparams.label_smoothing = 0.1
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.attention_dropout = 0.1
+  hparams.relu_dropout = 0.1
+  return hparams
+
+
+@registry.register_hparams
+def mtr_tr_dense_0():
+  return mtr_tr_dense(0)
+
+
+@registry.register_hparams
+def mtr_tr_dense_1():
+  return mtr_tr_dense(1)
+
+
+@registry.register_hparams
+def mtr_tr_dense_2():
+  return mtr_tr_dense(2)
+
+
+@registry.register_hparams
+def mtr_tr_dense_3():
+  return mtr_tr_dense(3)
+
+
+@registry.register_hparams
+def mtr_tr_dense_0_short():
+  hparams = mtr_tr_dense(0)
+  hparams.num_hidden_layers = 3
+  return hparams

From 95f6376ebeed0a86996f62b7a17be6d0d77d4ac1 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 7 Dec 2018 12:40:51 -0800
Subject: [PATCH 1337/2720] Tuning RL parameters.

PiperOrigin-RevId: 224561075
---
 tensor2tensor/data_generators/gym_env.py       |  2 +-
 tensor2tensor/models/video/basic_stochastic.py |  9 +++++----
 tensor2tensor/rl/trainer_model_based.py        | 13 ++++++-------
 tensor2tensor/rl/trainer_model_based_params.py |  2 +-
 4 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 055a6a952..5b2a88004 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -790,7 +790,7 @@ def camel_case_name(snake_case_name):
   return "".join([w[0].upper() + w[1:] for w in snake_case_name.split("_")])
 
 
-def register_game(game_name, game_mode="Deterministic-v4"):
+def register_game(game_name, game_mode="NoFrameskip-v4"):
   """Create and register problems for the game.
 
   Args:
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 5dc6ade53..aac145600 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -257,15 +257,16 @@ def next_frame_basic_stochastic_discrete():
   hparams.learning_rate_warmup_steps = 2000
   hparams.learning_rate_schedule = "linear_warmup * constant"
   hparams.concat_internal_states = True
-  hparams.add_hparam("bottleneck_bits", 256)
+  hparams.video_modality_loss_cutoff = 0.03
+  hparams.add_hparam("bottleneck_bits", 128)
   hparams.add_hparam("bottleneck_noise", 0.1)
   hparams.add_hparam("discretize_warmup_steps", 40000)
   hparams.add_hparam("latent_rnn_warmup_steps", 40000)
-  hparams.add_hparam("latent_rnn_max_sampling", 0.6)
-  hparams.add_hparam("latent_use_max_probability", 0.8)
+  hparams.add_hparam("latent_rnn_max_sampling", 0.5)
+  hparams.add_hparam("latent_use_max_probability", 0.7)
   hparams.add_hparam("full_latent_tower", False)
   hparams.add_hparam("latent_predictor_state_size", 128)
-  hparams.add_hparam("latent_predictor_temperature", 0.9)
+  hparams.add_hparam("latent_predictor_temperature", 1.0)
   hparams.add_hparam("complex_addn", True)
   return hparams
 
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 3b8aad5d0..d9b1fa7f8 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -155,8 +155,7 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
   getattr(exp, schedule)()
 
 
-def train_agent(
-    real_env, learner, world_model_dir, hparams, epoch, is_final_epoch):
+def train_agent(real_env, learner, world_model_dir, hparams, epoch):
   """Train the PPO agent in the simulated environment."""
   frame_stack_size = hparams.frame_stack_size
   initial_frame_rollouts = real_env.current_epoch_rollouts(
@@ -196,7 +195,10 @@ def initial_frame_chooser(batch_size):
       train_hparams, hparams, base_algo_str + "_"
   )
 
-  env_step_multiplier = 1 if not is_final_epoch else 2
+  final_epoch = hparams.epochs - 1
+  is_special_epoch = (epoch + 3) == final_epoch or (epoch + 7) == final_epoch
+  is_final_epoch = epoch == final_epoch
+  env_step_multiplier = 3 if is_final_epoch else 2 if is_special_epoch else 1
   learner.train(
       env_fn, train_hparams, simulated=True, save_continuously=True,
       epoch=epoch, env_step_multiplier=env_step_multiplier
@@ -436,7 +438,6 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   world_model_steps_num = 0
 
   for epoch in range(hparams.epochs):
-    is_final_epoch = (epoch + 1) == hparams.epochs
     log = make_log_fn(epoch, log_relative_time)
 
     # Train world model
@@ -448,9 +449,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
     # Train agent
     log("Training policy in simulated environment.")
-    train_agent(
-        env, learner, directories["world_model"], hparams, epoch, is_final_epoch
-    )
+    train_agent(env, learner, directories["world_model"], hparams, epoch)
 
     env.start_new_epoch(epoch, data_dir)
 
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 4e6f97f82..a36d5acdd 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -659,9 +659,9 @@ def rlmb_whitelisted_games(rhp):
 
 @registry.register_ranged_hparams
 def rlmb_human_score_games(rhp):
-  rhp.set_discrete("model.moe_loss_coef", list(range(10)))
   rhp.set_categorical("loop.game",
                       gym_env.ATARI_GAMES_WITH_HUMAN_SCORE)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
 
 
 @registry.register_ranged_hparams

From 6125b90a7be70caa7cb777109bab3d0b044f4074 Mon Sep 17 00:00:00 2001
From: deepbrainwaves <braincoder@aliyun.com>
Date: Sat, 8 Dec 2018 04:46:39 +0800
Subject: [PATCH 1338/2720] Fix a bug in document (#1267)

It gives troubles to beginners like me. After checked code in t2t_model.py, I find out this description about function returns is wrong.
---
 tensor2tensor/layers/modalities.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 83935e4da..0fe4121ad 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -811,7 +811,7 @@ def top(self, body_output, _):
       body_output: A Tensor with shape [batch, ?, ?, body_output_size].
 
     Returns:
-      a Tensors, each with shape [batch_size, ?, ?, vocab_size]
+      a Tensors, each with shape [batch_size, 1, 1, 1, vocab_size]
     """
     with tf.variable_scope(self.name):
       x = body_output

From 0ebe8fad0391aa4727bc75e2f2e3bc0a682c2aaf Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Fri, 7 Dec 2018 14:48:51 -0800
Subject: [PATCH 1339/2720] Mesh-TensorFlow:   Remove redefined-builtins from
 ops.py and move them to __init__.py   Add a few potentially useful
 operations:     mtf.sign     mtf.abs    
 mtf.layers.sigmoid_cross_entropy_with_logits

mtf Transformer implementation:
  Remove logit-jittering and replace it with "z_loss", which seems to work better.
  Hard-code the broadcast dimensions for the dropout layers.

PiperOrigin-RevId: 224581601
---
 tensor2tensor/models/mtf_transformer2.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index 3947b59cf..ab885e458 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -118,7 +118,8 @@ def model(self):
         input_vocab_size=input_vocab_size,
         output_vocab_size=self._targets_vocab_size,
         autoregressive=self.autoregressive,
-        max_length=hparams.max_length)
+        max_length=hparams.max_length,
+        z_loss=hparams.z_loss)
 
   def _mtf_model_fn(self, features, mesh):
     self._original_features = features
@@ -218,7 +219,8 @@ def model(self):
         output_vocab_size=self._targets_vocab_size,
         max_length=hparams.max_length,
         shared_embedding=hparams.shared_embedding,
-        label_smoothing=hparams.label_smoothing)
+        label_smoothing=hparams.label_smoothing,
+        z_loss=hparams.z_loss)
 
   def _mtf_model_fn(self, features, mesh):
     self._original_features = features
@@ -298,6 +300,9 @@ def mtf_transformer2_base():
   hparams.batch_size = 4
   hparams.max_length = 1024
   hparams.label_smoothing = 0.0
+  # a small positive value - this seems important for stability when training
+  # with bfloat16 activations.
+  hparams.add_hparam("z_loss", 1e-4)
 
   # These hyperparameters are used in default_layer_stack()
   # They may not be respected if hparams uses a differet layer stack function.

From ec5488f089256424d2d1e80c1db74a95d261ac59 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Sun, 9 Dec 2018 08:08:30 -0800
Subject: [PATCH 1340/2720] bug fix in MtfUnitransformer model.

PiperOrigin-RevId: 224718311
---
 tensor2tensor/models/mtf_transformer2.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index ab885e458..500479bb3 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -109,9 +109,9 @@ def model(self):
       # hparams.layer_stack is a function for creating a LayerStack
       layer_stack = hparams.layer_stack(hparams)
     if self.autoregressive:
-      input_vocab_size = self._inputs_vocab_size
-    else:
       input_vocab_size = self._targets_vocab_size
+    else:
+      input_vocab_size = self._inputs_vocab_size
     return transformer.Unitransformer(
         layer_stack=layer_stack,
         d_model=hparams.d_model,

From 44f669058390bec03024baa04c1a33e91cc0909d Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Sun, 9 Dec 2018 12:12:30 -0800
Subject: [PATCH 1341/2720] Add options to replicate encoder_layers &
 decoder_layers in mtf_transformer

PiperOrigin-RevId: 224728865
---
 tensor2tensor/models/mtf_transformer.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index f06115a2a..59de31adb 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -35,6 +35,27 @@
 class MtfTransformer(mtf_model.MtfModel):
   """Transformer in mesh_tensorflow."""
 
+  def __init__(self,
+               hparams,
+               mode=tf.estimator.ModeKeys.TRAIN,
+               problem_hparams=None,
+               data_parallelism=None,
+               decode_hparams=None,
+               **kwargs):
+    """Init with assignments of hparams.encoder_layers / decoder_layers."""
+    # Finalize encoder_layers, decoder_layers
+    hparams.encoder_layers = (
+        hparams.encoder_layers * hparams.encoder_replicate_factor)
+    hparams.decoder_layers = (
+        hparams.decoder_layers * hparams.decoder_replicate_factor)
+
+    super(MtfTransformer, self).__init__(hparams,
+                                         mode=mode,
+                                         problem_hparams=problem_hparams,
+                                         data_parallelism=data_parallelism,
+                                         decode_hparams=decode_hparams,
+                                         **kwargs)
+
   @property
   def batch_dims(self):
     hparams = self._hparams
@@ -786,6 +807,8 @@ def mtf_transformer_base():
   hparams.add_hparam("layout", "batch:batch;vocab:model;d_ff:model;heads:model")
   hparams.add_hparam("num_heads", 8)
   hparams.add_hparam("d_ff", 2048)
+  hparams.add_hparam("encoder_replicate_factor", 1)
+  hparams.add_hparam("decoder_replicate_factor", 1)
   hparams.add_hparam("encoder_layers", ["att", "drd"] * 6)
   hparams.add_hparam("decoder_layers", ["att", "enc_att", "drd"] * 6)
   hparams.add_hparam("attention_dropout", 0.1)

From 7d4475e5ec9aa98ff12e2ba63825dea6cc5c6da9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 10 Dec 2018 10:21:45 -0800
Subject: [PATCH 1342/2720] Introduce dependency on tf-agents

PiperOrigin-RevId: 224836107
---
 setup.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/setup.py b/setup.py
index 6c8d76cab..c494ade1c 100644
--- a/setup.py
+++ b/setup.py
@@ -50,6 +50,7 @@
         'six',
         'sympy',
         'tensorflow-probability',
+        'tf-agents',
         'tfds-nightly',
         'tqdm',
     ],

From a0a6d009e3020a85604b971c351e1587b2e2bb45 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 11 Dec 2018 01:21:55 +0100
Subject: [PATCH 1343/2720] Implement training restarter and LR decay for
 policy (#1291)

* Implement training restarter

* Learning rate decay
---
 tensor2tensor/models/research/rl.py           |  15 ++-
 tensor2tensor/rl/policy_learner.py            |   5 +-
 tensor2tensor/rl/ppo.py                       |  44 +++----
 tensor2tensor/rl/ppo_learner.py               |  96 ++++++++------
 tensor2tensor/rl/restarter.py                 | 124 ++++++++++++++++++
 tensor2tensor/rl/restarter_test.py            | 124 ++++++++++++++++++
 tensor2tensor/rl/rl_utils.py                  |   2 +-
 tensor2tensor/rl/trainer_model_based.py       |  27 ++--
 .../rl/trainer_model_based_params.py          |  12 +-
 tensor2tensor/rl/trainer_model_free.py        |   2 +-
 10 files changed, 358 insertions(+), 93 deletions(-)
 create mode 100644 tensor2tensor/rl/restarter.py
 create mode 100644 tensor2tensor/rl/restarter_test.py

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 2a3b37302..8c93504ef 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -37,7 +37,11 @@
 def ppo_base_v1():
   """Set of hyperparameters."""
   hparams = common_hparams.basic_params1()
-  hparams.learning_rate = 1e-4
+  hparams.learning_rate_schedule = "constant"
+  hparams.learning_rate_constant = 1e-4
+  hparams.clip_grad_norm = 0.5
+  # If set, extends the LR warmup to all epochs except the final one.
+  hparams.add_hparam("lr_decay_in_final_epoch", False)
   hparams.add_hparam("init_mean_factor", 0.1)
   hparams.add_hparam("init_logstd", 0.1)
   hparams.add_hparam("policy_layers", (100, 100))
@@ -53,7 +57,6 @@ def ppo_base_v1():
   hparams.add_hparam("eval_every_epochs", 10)
   hparams.add_hparam("save_models_every_epochs", 30)
   hparams.add_hparam("optimization_batch_size", 50)
-  hparams.add_hparam("max_gradients_norm", 0.5)
   hparams.add_hparam("intrinsic_reward_scale", 0.)
   hparams.add_hparam("logits_clip", 0.0)
   hparams.add_hparam("dropout_ppo", 0.1)
@@ -85,7 +88,7 @@ def discrete_random_action_base():
 def ppo_atari_base():
   """Pong base parameters."""
   hparams = ppo_discrete_action_base()
-  hparams.learning_rate = 1e-4
+  hparams.learning_rate_constant = 1e-4
   hparams.epoch_length = 200
   hparams.gae_gamma = 0.985
   hparams.gae_lambda = 0.985
@@ -96,7 +99,7 @@ def ppo_atari_base():
   hparams.policy_network = "feed_forward_cnn_small_categorical_policy"
   hparams.clipping_coef = 0.2
   hparams.optimization_batch_size = 20
-  hparams.max_gradients_norm = 0.5
+  hparams.clip_grad_norm = 0.5
   return hparams
 
 
@@ -104,7 +107,7 @@ def ppo_atari_base():
 def ppo_original_params():
   """Parameters based on the original PPO paper."""
   hparams = ppo_atari_base()
-  hparams.learning_rate = 2.5e-4
+  hparams.learning_rate_constant = 2.5e-4
   hparams.gae_gamma = 0.99
   hparams.gae_lambda = 0.95
   hparams.clipping_coef = 0.1
@@ -178,7 +181,7 @@ def get_policy(observations, hparams, action_space):
 def ppo_pong_ae_base():
   """Pong autoencoder base parameters."""
   hparams = ppo_original_params()
-  hparams.learning_rate = 1e-4
+  hparams.learning_rate_constant = 1e-4
   hparams.network = "dense_bitwise_categorical_policy"
   return hparams
 
diff --git a/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
index b34af8c8c..db6bc5854 100644
--- a/tensor2tensor/rl/policy_learner.py
+++ b/tensor2tensor/rl/policy_learner.py
@@ -23,10 +23,13 @@
 class PolicyLearner(object):
   """API for policy learners."""
 
-  def __init__(self, frame_stack_size, base_event_dir, agent_model_dir):
+  def __init__(
+      self, frame_stack_size, base_event_dir, agent_model_dir, total_num_epochs
+  ):
     self.frame_stack_size = frame_stack_size
     self.base_event_dir = base_event_dir
     self.agent_model_dir = agent_model_dir
+    self.total_num_epochs = total_num_epochs
 
   def train(
       self,
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index e9db06f51..bed9b6861 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -22,18 +22,14 @@
 from __future__ import print_function
 
 from tensor2tensor.models.research.rl import get_policy
+from tensor2tensor.utils import learning_rate
+from tensor2tensor.utils import optimize
 
 import tensorflow as tf
 import tensorflow_probability as tfp
 
 
-def get_optimiser(config):
-  if config.optimizer == "Adam":
-    return tf.train.AdamOptimizer(learning_rate=config.learning_rate)
-  return config.optimizer(learning_rate=config.learning_rate)
-
-
-def define_ppo_step(data_points, optimizer, hparams, action_space):
+def define_ppo_step(data_points, hparams, action_space, lr):
   """Define ppo step."""
   observation, action, discounted_reward, norm_advantage, old_pdf = data_points
 
@@ -57,24 +53,11 @@ def define_ppo_step(data_points, optimizer, hparams, action_space):
   entropy_loss = -hparams.entropy_loss_coef * tf.reduce_mean(entropy)
 
   losses = [policy_loss, value_loss, entropy_loss]
+  loss = sum(losses)
+  train_op = optimize.optimize(loss, lr, hparams)
 
-  gradients = [list(zip(*optimizer.compute_gradients(loss)))
-               for loss in losses]
-
-  gradients_norms = [tf.global_norm(gradient[0]) for gradient in gradients]
-
-  gradients_flat = sum([gradient[0] for gradient in gradients], ())
-  gradients_variables_flat = sum([gradient[1] for gradient in gradients], ())
-
-  if hparams.max_gradients_norm:
-    gradients_flat, _ = tf.clip_by_global_norm(gradients_flat,
-                                               hparams.max_gradients_norm)
-
-  optimize_op = optimizer.apply_gradients(zip(gradients_flat,
-                                              gradients_variables_flat))
-
-  with tf.control_dependencies([optimize_op]):
-    return [tf.identity(x) for x in losses + gradients_norms]
+  with tf.control_dependencies([train_op]):
+    return [tf.identity(x) for x in losses]
 
 
 def define_ppo_epoch(memory, hparams, action_space, batch_size):
@@ -118,22 +101,25 @@ def define_ppo_epoch(memory, hparams, action_space, batch_size):
   dataset = dataset.repeat(-1)
   dataset = dataset.batch(hparams.optimization_batch_size, drop_remainder=True)
   iterator = dataset.make_initializable_iterator()
-  optimizer = get_optimiser(hparams)
+
+  lr = learning_rate.learning_rate_schedule(hparams)
 
   with tf.control_dependencies([iterator.initializer]):
     ppo_step_rets = tf.scan(
         lambda a, i: add_lists_elementwise(  # pylint: disable=g-long-lambda
             a, define_ppo_step(
-                iterator.get_next(), optimizer, hparams, action_space
+                iterator.get_next(), hparams, action_space, lr
             )),
         tf.range(number_of_batches),
-        [0., 0., 0., 0., 0., 0.],
+        [0., 0., 0.],
         parallel_iterations=1)
 
   ppo_summaries = [tf.reduce_mean(ret) / number_of_batches
                    for ret in ppo_step_rets]
-  summaries_names = ["policy_loss", "value_loss", "entropy_loss",
-                     "policy_gradient", "value_gradient", "entropy_gradient"]
+  ppo_summaries.append(lr)
+  summaries_names = [
+      "policy_loss", "value_loss", "entropy_loss", "learning_rate"
+  ]
 
   summaries = [tf.summary.scalar(summary_name, summary)
                for summary_name, summary in zip(summaries_names, ppo_summaries)]
diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 35310b0a8..6cfd700aa 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -28,6 +28,7 @@
 from tensor2tensor.rl.envs.tf_atari_wrappers import StackWrapper
 from tensor2tensor.rl.envs.tf_atari_wrappers import WrapperBase
 from tensor2tensor.rl.policy_learner import PolicyLearner
+from tensor2tensor.rl.restarter import Restarter
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -40,6 +41,7 @@ class PPOLearner(PolicyLearner):
   def __init__(self, *args, **kwargs):
     super(PPOLearner, self).__init__(*args, **kwargs)
     self._num_completed_iterations = 0
+    self._lr_decay_start = None
 
   def train(self,
             env_fn,
@@ -90,11 +92,29 @@ def train(self,
         iteration_increment *= env_step_multiplier
 
         self._num_completed_iterations += iteration_increment
+
+        restarter = Restarter(
+            "policy", self.agent_model_dir, self._num_completed_iterations
+        )
+        if restarter.should_skip:
+          return
+
+        if hparams.lr_decay_in_final_epoch:
+          if epoch != self.total_num_epochs - 1:
+            # Extend the warmup period to the end of this epoch.
+            hparams.learning_rate_warmup_steps = restarter.target_global_step
+          else:
+            if self._lr_decay_start is None:
+              # Stop the warmup at the beginning of this epoch.
+              self._lr_decay_start = \
+                  restarter.target_global_step - iteration_increment
+            hparams.learning_rate_warmup_steps = self._lr_decay_start
+
         _run_train(
             hparams,
             event_dir,
             self.agent_model_dir,
-            self._num_completed_iterations,
+            restarter,
             train_summary_op,
             eval_summary_op,
             initializers,
@@ -166,7 +186,7 @@ def _define_train(
 def _run_train(ppo_hparams,
                event_dir,
                model_dir,
-               num_target_iterations,
+               restarter,
                train_summary_op,
                eval_summary_op,
                initializers,
@@ -176,49 +196,49 @@ def _run_train(ppo_hparams,
       event_dir, graph=tf.get_default_graph(), flush_secs=60)
 
   model_saver = tf.train.Saver(
-      tf.global_variables(ppo_hparams.policy_network + "/.*")
+      tf.global_variables(ppo_hparams.policy_network + "/.*") +
+      tf.global_variables("global_step")
   )
 
+  global_step = tf.train.get_or_create_global_step()
+  with tf.control_dependencies([tf.assign_add(global_step, 1)]):
+    train_summary_op = tf.identity(train_summary_op)
+
   with tf.Session() as sess:
     sess.run(tf.global_variables_initializer())
     for initializer in initializers:
       initializer(sess)
-    num_completed_iterations = trainer_lib.restore_checkpoint(
-        model_dir, model_saver, sess)
-
-    # Fail-friendly, complete only unfinished epoch
-    if num_target_iterations <= num_completed_iterations:
-      tf.logging.info(
-          "Skipping PPO training. Requested %d iterations while %d train "
-          "iterations already reached", num_target_iterations,
-          num_completed_iterations)
-      return
-
-    for epoch_index in range(num_completed_iterations, num_target_iterations):
-      summary = sess.run(train_summary_op)
-      if summary_writer:
-        summary_writer.add_summary(summary, epoch_index)
-
-      if (ppo_hparams.eval_every_epochs and
-          epoch_index % ppo_hparams.eval_every_epochs == 0):
-        eval_summary = sess.run(eval_summary_op)
+    trainer_lib.restore_checkpoint(model_dir, model_saver, sess)
+
+    num_target_iterations = restarter.target_local_step
+    num_completed_iterations = num_target_iterations - restarter.steps_to_go
+    with restarter.training_loop():
+      for epoch_index in range(num_completed_iterations, num_target_iterations):
+        summary = sess.run(train_summary_op)
         if summary_writer:
-          summary_writer.add_summary(eval_summary, epoch_index)
-        if report_fn:
-          summary_proto = tf.Summary()
-          summary_proto.ParseFromString(eval_summary)
-          for elem in summary_proto.value:
-            if "mean_score" in elem.tag:
-              report_fn(elem.simple_value, epoch_index)
-              break
-
-      if (model_saver and ppo_hparams.save_models_every_epochs and
-          (epoch_index % ppo_hparams.save_models_every_epochs == 0 or
-           (epoch_index + 1) == num_target_iterations)):
-        ckpt_path = os.path.join(
-            model_dir,
-            "model.ckpt-{}".format(epoch_index + 1))
-        model_saver.save(sess, ckpt_path)
+          summary_writer.add_summary(summary, epoch_index)
+
+        if (ppo_hparams.eval_every_epochs and
+            epoch_index % ppo_hparams.eval_every_epochs == 0):
+          eval_summary = sess.run(eval_summary_op)
+          if summary_writer:
+            summary_writer.add_summary(eval_summary, epoch_index)
+          if report_fn:
+            summary_proto = tf.Summary()
+            summary_proto.ParseFromString(eval_summary)
+            for elem in summary_proto.value:
+              if "mean_score" in elem.tag:
+                report_fn(elem.simple_value, epoch_index)
+                break
+
+        if (model_saver and ppo_hparams.save_models_every_epochs and
+            (epoch_index % ppo_hparams.save_models_every_epochs == 0 or
+             (epoch_index + 1) == num_target_iterations)):
+          ckpt_path = os.path.join(
+              model_dir,
+              "model.ckpt-{}".format(tf.train.global_step(sess, global_step))
+          )
+          model_saver.save(sess, ckpt_path)
 
 
 def _rollout_metadata(batch_env):
diff --git a/tensor2tensor/rl/restarter.py b/tensor2tensor/rl/restarter.py
new file mode 100644
index 000000000..1dd7d7c5b
--- /dev/null
+++ b/tensor2tensor/rl/restarter.py
@@ -0,0 +1,124 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Training restarter."""
+
+from contextlib import contextmanager
+import os
+
+import tensorflow as tf
+
+
+class Restarter(object):
+  """Handles training restarts.
+
+  Particularly useful when sharing parameters (and checkpoints) between models.
+
+  Args:
+    model_mode (str): Model "mode". Different modes have different local step
+        counters, but the same global step counter. Also used in log messages.
+    checkpoint_dir (str): Model checkpoint directory. Global step is inferred
+        from the name of the last checkpoint.
+    target_local_step (int): Local step to train the model up to.
+
+  Attributes:
+    model_mode (str): See args.
+    checkpoint_dir (str): See args.
+    target_local_step (int): See args.
+    target_global_step (int): Calculated global step to train the model up to.
+    should_skip (bool): Whether training should be skipped because the number of
+        local steps already done is higher than the target. This happens during
+        restarts.
+    restarting (bool): Whether the current epoch of training has been
+        interrupted and is being restarted.
+  """
+
+  def __init__(self, model_mode, checkpoint_dir, target_local_step):
+    self.model_mode = model_mode
+    self.checkpoint_dir = checkpoint_dir
+    self.target_local_step = target_local_step
+    self.target_global_step = None
+    self.should_skip = False
+    self.restarting = False
+
+    self._counter_path = os.path.join(
+        checkpoint_dir, "{}_step_counter".format(model_mode)
+    )
+
+    self._global_step = self._get_global_step()
+    tf.logging.info(
+        "Will load %s checkpoint %d", self.model_mode, self._global_step
+    )
+
+    (self._local_step_at_start, global_step_at_start) = self._read_counters()
+
+    self.steps_to_go = target_local_step - self._local_step_at_start
+    if self.steps_to_go <= 0:
+      tf.logging.info(
+          "Skipping training %s, requested %d steps, already done %d",
+          self.model_mode, target_local_step, self._local_step_at_start
+      )
+      self.should_skip = True
+      return
+
+    if global_step_at_start != -1:
+      # Restart.
+      steps_done_this_epoch = self._global_step - global_step_at_start
+      self.steps_to_go -= steps_done_this_epoch
+      tf.logging.info(
+          "Restarting training %s, %d steps already done this epoch",
+          self.model_mode, steps_done_this_epoch
+      )
+      self.restarting = True
+
+    self.target_global_step = self._global_step + self.steps_to_go
+
+  @contextmanager
+  def training_loop(self):
+    """Context manager wrapping the training loop.
+
+    Takes care of updating the step counters.
+    """
+    if not self.restarting:
+      self._write_counters(self._local_step_at_start, self._global_step)
+
+    tf.logging.info(
+        "Training %s up to %d, %d to go", self.model_mode,
+        self.target_local_step, self.steps_to_go
+    )
+
+    yield
+
+    self._write_counters(self.target_local_step, -1)
+
+  def _get_global_step(self):
+    checkpoint = tf.train.latest_checkpoint(self.checkpoint_dir)
+    if checkpoint:
+      return int(checkpoint.split("-")[-1])
+    else:
+      return 0
+
+  def _read_counters(self):
+    try:
+      with tf.gfile.Open(self._counter_path, "r") as f:
+        return tuple(
+            int(counter) for counter in f.read().split(" ")
+        )
+    except tf.errors.NotFoundError:
+      return (0, -1)
+
+  def _write_counters(self, local_step, global_step):
+    with tf.gfile.Open(self._counter_path, "w") as f:
+      f.write("{} {}".format(local_step, global_step))
diff --git a/tensor2tensor/rl/restarter_test.py b/tensor2tensor/rl/restarter_test.py
new file mode 100644
index 000000000..9db1148d1
--- /dev/null
+++ b/tensor2tensor/rl/restarter_test.py
@@ -0,0 +1,124 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 rl_utils."""
+
+import os
+
+from tensor2tensor.rl.restarter import Restarter
+
+import tensorflow as tf
+
+
+TEST_MODE_1 = "mode1"
+TEST_MODE_2 = "mode2"
+TEST_NUM_STEPS = 2
+
+
+class RestarterTest(tf.test.TestCase):
+
+  def setUp(self):
+    self.out_dir = tf.test.get_temp_dir()
+    tf.gfile.DeleteRecursively(self.out_dir)
+    tf.gfile.MkDir(self.out_dir)
+
+  def create_checkpoint(self, global_step):
+    checkpoint_name = "model.ckpt-{}".format(global_step)
+    for suffix in ("index", "meta", "data-00000-of-00001"):
+      filename = "{}.{}".format(checkpoint_name, suffix)
+      # Just create the file.
+      with tf.gfile.Open(os.path.join(self.out_dir, filename), "w") as f:
+        f.write("")
+    tf.train.update_checkpoint_state(self.out_dir, checkpoint_name)
+
+  def run_single_mode(self, mode, target_local_step, target_global_step):
+    restarter = Restarter(mode, self.out_dir, target_local_step)
+    with restarter.training_loop():
+      self.create_checkpoint(target_global_step)
+
+  def assert_first_run(self, restarter, steps_to_go, target_global_step):
+    self.assertFalse(restarter.should_skip)
+    self.assertFalse(restarter.restarting)
+    self.assertEqual(restarter.steps_to_go, steps_to_go)
+    self.assertEqual(restarter.target_global_step, target_global_step)
+
+  def test_runs_in_single_mode(self):
+    restarter = Restarter(
+        TEST_MODE_1, self.out_dir, target_local_step=TEST_NUM_STEPS
+    )
+    self.assert_first_run(
+        restarter, steps_to_go=TEST_NUM_STEPS, target_global_step=TEST_NUM_STEPS
+    )
+
+  def test_runs_in_two_modes(self):
+    global_step = TEST_NUM_STEPS
+    local_steps = {
+        TEST_MODE_1: TEST_NUM_STEPS,
+        TEST_MODE_2: 0
+    }
+    self.run_single_mode(TEST_MODE_1, local_steps[TEST_MODE_1], global_step)
+
+    for mode in [TEST_MODE_2, TEST_MODE_1]:
+      global_step += TEST_NUM_STEPS
+      local_steps[mode] += TEST_NUM_STEPS
+      restarter = Restarter(
+          mode, self.out_dir, target_local_step=local_steps[mode]
+      )
+      self.assert_first_run(
+          restarter, steps_to_go=TEST_NUM_STEPS, target_global_step=global_step
+      )
+      with restarter.training_loop():
+        self.create_checkpoint(global_step)
+
+  def test_skips_already_done(self):
+    self.run_single_mode(
+        TEST_MODE_1, target_local_step=TEST_NUM_STEPS,
+        target_global_step=TEST_NUM_STEPS
+    )
+
+    restarter = Restarter(
+        TEST_MODE_1, self.out_dir, target_local_step=TEST_NUM_STEPS
+    )
+    # We should skip the training as those steps are already completed.
+    self.assertTrue(restarter.should_skip)
+
+  def test_restarts_after_interruption(self):
+    # Run some initial training first.
+    self.run_single_mode(
+        TEST_MODE_1, target_local_step=TEST_NUM_STEPS,
+        target_global_step=TEST_NUM_STEPS
+    )
+    global_step = TEST_NUM_STEPS
+
+    restarter = Restarter(
+        TEST_MODE_2, self.out_dir, target_local_step=2
+    )
+    with self.assertRaises(RuntimeError):
+      global_step += 1
+      with restarter.training_loop():
+        self.create_checkpoint(global_step)
+        # Simulate training interruption after the first step.
+        raise RuntimeError
+    restarter = Restarter(
+        TEST_MODE_2, self.out_dir, target_local_step=2
+    )
+
+    self.assertFalse(restarter.should_skip)
+    self.assertTrue(restarter.restarting)
+    # Training should resume after the first step.
+    self.assertEqual(restarter.steps_to_go, 1)
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 460b1888d..4b331856b 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -68,7 +68,7 @@ def evaluate_single_config(
   env_fn = rl.make_real_env_fn(env)
   learner = LEARNERS[hparams.base_algo](
       hparams.frame_stack_size, base_event_dir=None,
-      agent_model_dir=agent_model_dir
+      agent_model_dir=agent_model_dir, total_num_epochs=1
   )
   learner.evaluate(env_fn, eval_hparams, sampling_temp)
   rollouts = env.current_epoch_rollouts()
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index d9b1fa7f8..570f1ff62 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -41,6 +41,7 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl import trainer_model_based_params
+from tensor2tensor.rl.restarter import Restarter
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -241,16 +242,20 @@ def train_world_model(
   if epoch > 0:
     model_hparams.learning_rate *= hparams.learning_rate_bump
 
-  train_supervised(
-      problem=env,
-      model_name=hparams.generative_model,
-      hparams=model_hparams,
-      data_dir=data_dir,
-      output_dir=output_dir,
-      train_steps=world_model_steps_num,
-      eval_steps=100,
-      local_eval_frequency=2000
-  )
+  restarter = Restarter("world_model", output_dir, world_model_steps_num)
+  if restarter.should_skip:
+    return world_model_steps_num
+  with restarter.training_loop():
+    train_supervised(
+        problem=env,
+        model_name=hparams.generative_model,
+        hparams=model_hparams,
+        data_dir=data_dir,
+        output_dir=output_dir,
+        train_steps=restarter.target_global_step,
+        eval_steps=100,
+        local_eval_frequency=2000
+    )
 
   return world_model_steps_num
 
@@ -409,7 +414,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   learner = rl_utils.LEARNERS[hparams.base_algo](
       hparams.frame_stack_size, directories["policy"],
-      directories["policy"]
+      directories["policy"], hparams.epochs
   )
 
   # Timing log function
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index a36d5acdd..754e729ab 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -148,10 +148,10 @@ def rlmb_ppo_base():
       ppo_epoch_length=hparams.simulated_rollout_length,
       # Do not eval since simulated batch env does not produce dones
       ppo_eval_every_epochs=0,
-      ppo_learning_rate=1e-4,  # Will be changed, just so it exists.
+      ppo_learning_rate_constant=1e-4,  # Will be changed, just so it exists.
       # This needs to be divisible by real_ppo_effective_num_agents.
       real_ppo_epoch_length=16 * 200,
-      real_ppo_learning_rate=1e-4,
+      real_ppo_learning_rate_constant=1e-4,
       real_ppo_effective_num_agents=16,
       real_ppo_eval_every_epochs=0,
 
@@ -635,7 +635,7 @@ def rlmb_three(rhp):
 def rlmb_test1(rhp):
   rhp.set_discrete("model.moe_loss_coef", list(range(10)))
   rhp.set_categorical("loop.game", ["breakout", "pong", "boxing"])
-  rhp.set_discrete("loop.ppo_learning_rate", [5e-5, 1e-4, 2e-4])
+  rhp.set_discrete("loop.ppo_learning_rate_constant", [5e-5, 1e-4, 2e-4])
   rhp.set_discrete("ppo.optimization_batch_size", [20, 40])
   rhp.set_discrete("loop.epochs", [3, 6])
 
@@ -696,20 +696,20 @@ def rlmb_ae_variance(rhp):
 def rlmb_ppolr_game(rhp):
   rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
   base_lr = 1e-4
-  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
+  rhp.set_float("loop.ppo_learning_rate_constant", base_lr / 2, base_lr * 2)
 
 
 @registry.register_ranged_hparams
 def rlmb_ppolr(rhp):
   base_lr = 1e-4
-  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
+  rhp.set_float("loop.ppo_learning_rate_constant", base_lr / 2, base_lr * 2)
 
 
 @registry.register_ranged_hparams
 def rlmb_ae_ppo_lr(rhp):
   rhp.set_categorical("loop.game", ["breakout", "pong", "freeway"])
   base_lr = 1e-4
-  rhp.set_float("loop.ppo_learning_rate", base_lr / 2, base_lr * 2)
+  rhp.set_float("loop.ppo_learning_rate_constant", base_lr / 2, base_lr * 2)
 
 
 @registry.register_ranged_hparams
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 6948b0abf..19d8de67c 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -62,7 +62,7 @@ def train(hparams, output_dir, report_fn=None):
   """Train."""
   hparams = initialize_env_specs(hparams)
   learner = rl_utils.LEARNERS[hparams.base_algo](
-      hparams.frame_stack_size, FLAGS.output_dir, output_dir
+      hparams.frame_stack_size, FLAGS.output_dir, output_dir, total_num_epochs=1
   )
   policy_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
   rl_utils.update_hparams_from_hparams(

From 25db6d6a135458c6c6d74b9eb22e2cc567371a0d Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Mon, 10 Dec 2018 16:22:14 -0800
Subject: [PATCH 1344/2720] internal merge of PR #1291

PiperOrigin-RevId: 224904504
---
 tensor2tensor/rl/restarter.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/rl/restarter.py b/tensor2tensor/rl/restarter.py
index 1dd7d7c5b..c9a2b1110 100644
--- a/tensor2tensor/rl/restarter.py
+++ b/tensor2tensor/rl/restarter.py
@@ -15,7 +15,7 @@
 
 """Training restarter."""
 
-from contextlib import contextmanager
+import contextlib
 import os
 
 import tensorflow as tf
@@ -41,6 +41,7 @@ class Restarter(object):
     should_skip (bool): Whether training should be skipped because the number of
         local steps already done is higher than the target. This happens during
         restarts.
+    steps_to_go: how many steps to go.
     restarting (bool): Whether the current epoch of training has been
         interrupted and is being restarted.
   """
@@ -85,12 +86,9 @@ def __init__(self, model_mode, checkpoint_dir, target_local_step):
 
     self.target_global_step = self._global_step + self.steps_to_go
 
-  @contextmanager
+  @contextlib.contextmanager
   def training_loop(self):
-    """Context manager wrapping the training loop.
-
-    Takes care of updating the step counters.
-    """
+    """Context manager wrapping the training loop, updates step counters."""
     if not self.restarting:
       self._write_counters(self._local_step_at_start, self._global_step)
 

From 111466d04d790a144843ee1acc6726d33dfa836e Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 10 Dec 2018 18:22:00 -0800
Subject: [PATCH 1345/2720] Disable tests in oss_tests.sh in the correct
 location.

PiperOrigin-RevId: 224921908
---
 oss_scripts/oss_tests.sh | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index ee9845714..d0f4500a9 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -47,6 +47,7 @@ pytest \
   --ignore=tensor2tensor/data_generators/allen_brain_test.py \
   --ignore=tensor2tensor/rl \
   --ignore=tensor2tensor/models/research \
+  --ignore=tensor2tensor/models/video/nfg_*.py \
   --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary \
   --deselect=tensor2tensor/utils/beam_search_test.py::BeamSearchTest::testTPUBeam
 set_status
@@ -71,10 +72,6 @@ then
   # * Glow requires the CIFAR-10 dataset to be generated
   pytest tensor2tensor/models/research \
     --ignore=tensor2tensor/models/research/glow_test.py \
-    --ignore=tensor2tensor/models/video/nfg_conv3d_test.py \
-    --ignore=tensor2tensor/models/video/nfg_conv_lstm_test.py \
-    --ignore=tensor2tensor/models/video/nfg_conv_test.py \
-    --ignore=tensor2tensor/models/video/nfg_uncond_test.py \
   set_status
 fi
 

From c110fabb1f2e87a1409ff398939761ff01e6eb95 Mon Sep 17 00:00:00 2001
From: Art Wangperawong <artitw@gmail.com>
Date: Mon, 10 Dec 2018 21:50:59 -0500
Subject: [PATCH 1346/2720] New problem: mathematical language understanding
 (#1290)

* fix bAbi data generator and readme

* Fix bAbi hparams deletion

* Fix bAbi hparams delete unecessary keys

* Fix bAbi hparams clean keys

* bAbi hparams delete keys

* fix readme

* fix universal transformer decoding

* fix merge conflict

* mathematical language understanding

* clarify usage

* add to authors
---
 AUTHORS                                       |   1 +
 README.md                                     |  20 ++++
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/babi_qa.py      |   5 +-
 .../mathematical_language_understanding.py    | 104 ++++++++++++++++++
 .../models/research/universal_transformer.py  |   3 +-
 6 files changed, 129 insertions(+), 5 deletions(-)
 create mode 100644 tensor2tensor/data_generators/mathematical_language_understanding.py

diff --git a/AUTHORS b/AUTHORS
index 38e5bc724..b4762f933 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -5,3 +5,4 @@
 # of contributors, see the revision history in source control.
 
 Google Inc.
+Artit Wangperawong
\ No newline at end of file
diff --git a/README.md b/README.md
index 1a031086b..7e72ce13b 100644
--- a/README.md
+++ b/README.md
@@ -47,6 +47,7 @@ pip install tensor2tensor && t2t-trainer \
 ### Contents
 
 * [Suggested Datasets and Models](#suggested-datasets-and-models)
+  * [Mathematical Language Understanding](#mathematical-language-understanding)
   * [Story, Question and Answer](#story-question-and-answer)
   * [Image Classification](#image-classification)
   * [Image Generation](#image-generation)
@@ -79,6 +80,24 @@ hyperparameters that we know works well in our setup. We usually
 run either on Cloud TPUs or on 8-GPU machines; you might need
 to modify the hyperparameters if you run on a different setup.
 
+### Mathematical Language Understanding
+
+For evaluating mathematical expressions at the character level involving addition, subtraction and multiplication of both positive and negative decimal numbers with variable digits assigned to symbolic variables, use
+
+* the [MLU](https://art.wangperawong.com/mathematical_language_understanding_train.tar.gz) data-set:
+ `--problem=mathematical_language_understanding`
+
+You can try solving the problem with different transformer models and hyperparameters as described in the [paper](https://arxiv.org/abs/1812.02825):
+* Standard transformer:
+`--model=transformer`
+`--hparams_set=transformer_tiny`
+* Universal transformer:
+`--model=universal_transformer`
+`--hparams_set=universal_transformer_tiny`
+* Adaptive universal transformer:
+`--model=universal_transformer`
+`--hparams_set=adaptive_universal_transformer_tiny`
+
 ### Story, Question and Answer
 
 For answering questions based on a story, use
@@ -464,5 +483,6 @@ T2T](https://research.googleblog.com/2017/06/accelerating-deep-learning-research
 * [Fast Decoding in Sequence Models using Discrete Latent Variables](https://arxiv.org/abs/1803.03382)
 * [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://arxiv.org/abs/1804.04235)
 * [Universal Transformers](https://arxiv.org/abs/1807.03819)
+* [Attending to Mathematical Language with Transformers](https://arxiv.org/abs/1812.02825)
 
 *Note: This is not an official Google product.*
diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index b3f7263f1..688197ab0 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -50,6 +50,7 @@
     "tensor2tensor.data_generators.lm1b",
     "tensor2tensor.data_generators.lm1b_imdb",
     "tensor2tensor.data_generators.lm1b_mnli",
+    "tensor2tensor.data_generators.mathematical_language_understanding",
     "tensor2tensor.data_generators.mnist",
     "tensor2tensor.data_generators.mrpc",
     "tensor2tensor.data_generators.mscoco",
diff --git a/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index a11eaddca..882d17778 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -109,9 +109,9 @@ def _prepare_babi_data(tmp_dir, data_dir):
     tf.gfile.MakeDirs(data_dir)
 
   file_path = os.path.join(tmp_dir, _TAR)
-  headers = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36"}  # pylint: disable=line-too-long
+  headers = {'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36'}
   resp = requests.get(_URL, headers=headers)
-  with open(file_path, "wb") as f:
+  with open(file_path, 'wb') as f:
     f.write(resp.content)
 
   tar = tarfile.open(file_path)
@@ -459,7 +459,6 @@ def hparams(self, defaults, unused_model_hparams):
     if "context" in p.vocab_size:
       del p.vocab_size["context"]
 
-
 def _problems_to_register():
   """Problems for which we want to create datasets.
 
diff --git a/tensor2tensor/data_generators/mathematical_language_understanding.py b/tensor2tensor/data_generators/mathematical_language_understanding.py
new file mode 100644
index 000000000..53d4ddacb
--- /dev/null
+++ b/tensor2tensor/data_generators/mathematical_language_understanding.py
@@ -0,0 +1,104 @@
+# coding=utf-8
+# Copyright 2018 Artit Wangperawong artitw@gmail.com
+#
+# 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.
+
+r"""Data generators for the Mathematical Language Understanding dataset.
+
+The training and test data were generated by assigning symbolic variables 
+either positive or negative decimal integers and then describing the algebraic 
+operation to perform. We restrict our variable assignments to the range 
+x,y->[-1000,1000) and the operations to the set {+,-,*}. To ensure that the 
+model embraces symbolic variables, the order in which x and y appears in the 
+expression is randomly chosen. For instance, an input string contrasting from 
+the example shown above might be y=129,x=531,x-y. Each input string is 
+accompanied by its target string, which is the evaluation of the mathematical 
+expression. For this study, all targets considered are decimal integers 
+represented at the character level. About 12 million unique samples were thus 
+generated and randomly split into training and test sets at an approximate 
+ratio of 9:1, respectively. 
+
+For more information check the following paper:
+Artit Wangperawong. Attending to Mathematical Language with Transformers, 
+arXiv:1812.02825.
+Available at: https://arxiv.org/abs/1812.02825
+
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+@registry.register_problem
+class MathematicalLanguageUnderstanding(text_problems.Text2TextProblem):
+  URL = "https://art.wangperawong.com/mathematical_language_understanding_train.tar.gz"
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 10,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  @property
+  def is_generate_per_split(self):
+    return False
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    """Downloads and extracts the dataset and generates examples
+
+    Args:
+      tmp_dir: temp directory to download and extract the dataset
+      data_dir: The base directory where data and vocab files are stored.
+
+    Returns:
+      data generator
+    """
+
+    if not tf.gfile.Exists(tmp_dir):
+      tf.gfile.MakeDirs(tmp_dir)
+
+    if not tf.gfile.Exists(data_dir):
+      tf.gfile.MakeDirs(data_dir)
+
+    # Download and extract
+    compressed_filename = os.path.basename(self.URL)
+    download_path = generator_utils.maybe_download(tmp_dir, compressed_filename,
+                                                   self.URL)
+
+    with tarfile.open(download_path, "r:gz") as tar:
+      tar.extractall(tmp_dir)
+
+    filepath = os.path.join(tmp_dir, "mathematical_language_understanding_train.txt")
+
+    with open(filepath, 'r') as fp:
+      for l in fp:
+        prob, ans = l.strip().split(':')
+        yield {"inputs": prob, "targets": ans}
+
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 788037d9f..11f743c76 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -243,8 +243,7 @@ def _greedy_infer(self, features, decode_length, use_tpu=False):
     return (self._slow_greedy_infer_tpu(features, decode_length) if use_tpu else
             self._slow_greedy_infer(features, decode_length))
 
-  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha,
-                   use_tpu=False):
+  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha, use_tpu=False):
     """Beam search decoding.
 
     Args:

From 347fc43addbe20c0be62819614979478c61f8a21 Mon Sep 17 00:00:00 2001
From: Art Wangperawong <artitw@gmail.com>
Date: Mon, 10 Dec 2018 22:38:35 -0800
Subject: [PATCH 1347/2720] internal merge of PR #1290

PiperOrigin-RevId: 224943245
---
 docs/walkthrough.md                           | 20 ++++++
 tensor2tensor/data_generators/babi_qa.py      | 13 ++--
 .../mathematical_language_understanding.py    | 66 +++++++++----------
 .../models/research/universal_transformer.py  |  8 ++-
 4 files changed, 67 insertions(+), 40 deletions(-)

diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 1a031086b..7e72ce13b 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -47,6 +47,7 @@ pip install tensor2tensor && t2t-trainer \
 ### Contents
 
 * [Suggested Datasets and Models](#suggested-datasets-and-models)
+  * [Mathematical Language Understanding](#mathematical-language-understanding)
   * [Story, Question and Answer](#story-question-and-answer)
   * [Image Classification](#image-classification)
   * [Image Generation](#image-generation)
@@ -79,6 +80,24 @@ hyperparameters that we know works well in our setup. We usually
 run either on Cloud TPUs or on 8-GPU machines; you might need
 to modify the hyperparameters if you run on a different setup.
 
+### Mathematical Language Understanding
+
+For evaluating mathematical expressions at the character level involving addition, subtraction and multiplication of both positive and negative decimal numbers with variable digits assigned to symbolic variables, use
+
+* the [MLU](https://art.wangperawong.com/mathematical_language_understanding_train.tar.gz) data-set:
+ `--problem=mathematical_language_understanding`
+
+You can try solving the problem with different transformer models and hyperparameters as described in the [paper](https://arxiv.org/abs/1812.02825):
+* Standard transformer:
+`--model=transformer`
+`--hparams_set=transformer_tiny`
+* Universal transformer:
+`--model=universal_transformer`
+`--hparams_set=universal_transformer_tiny`
+* Adaptive universal transformer:
+`--model=universal_transformer`
+`--hparams_set=adaptive_universal_transformer_tiny`
+
 ### Story, Question and Answer
 
 For answering questions based on a story, use
@@ -464,5 +483,6 @@ T2T](https://research.googleblog.com/2017/06/accelerating-deep-learning-research
 * [Fast Decoding in Sequence Models using Discrete Latent Variables](https://arxiv.org/abs/1803.03382)
 * [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://arxiv.org/abs/1804.04235)
 * [Universal Transformers](https://arxiv.org/abs/1807.03819)
+* [Attending to Mathematical Language with Transformers](https://arxiv.org/abs/1812.02825)
 
 *Note: This is not an official Google product.*
diff --git a/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index 882d17778..56d8d66ae 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -109,9 +109,11 @@ def _prepare_babi_data(tmp_dir, data_dir):
     tf.gfile.MakeDirs(data_dir)
 
   file_path = os.path.join(tmp_dir, _TAR)
-  headers = {'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36'}
+  headers = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_1) "
+                           "AppleWebKit/537.36 (KHTML, like Gecko) "
+                           "Chrome/63.0.3239.132 Safari/537.36"}
   resp = requests.get(_URL, headers=headers)
-  with open(file_path, 'wb') as f:
+  with open(file_path, "wb") as f:
     f.write(resp.content)
 
   tar = tarfile.open(file_path)
@@ -192,10 +194,12 @@ def _all_task_raw_data_generator(tmp_dir, data_file, dataset_split):
 
     tf.logging.info("Preparing dataset of all task together")
     globe_name = ("*_{}.txt")
+    mode_name = "test"
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      mode_name = "train"
     files_name = os.path.join(
         tmp_dir, _DIR_NAME, subset,
-        globe_name.format("train" if dataset_split == problem.DatasetSplit.TRAIN
-                          else "test"))
+        globe_name.format(mode_name))
     with tf.gfile.GFile(data_file, "wb") as outfile:
       for filename in tf.gfile.Glob(files_name):
         if filename == data_file:
@@ -459,6 +463,7 @@ def hparams(self, defaults, unused_model_hparams):
     if "context" in p.vocab_size:
       del p.vocab_size["context"]
 
+
 def _problems_to_register():
   """Problems for which we want to create datasets.
 
diff --git a/tensor2tensor/data_generators/mathematical_language_understanding.py b/tensor2tensor/data_generators/mathematical_language_understanding.py
index 53d4ddacb..b413b0456 100644
--- a/tensor2tensor/data_generators/mathematical_language_understanding.py
+++ b/tensor2tensor/data_generators/mathematical_language_understanding.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 Artit Wangperawong artitw@gmail.com
+# Copyright 2018 The Tensor2Tensor Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,24 +15,22 @@
 
 r"""Data generators for the Mathematical Language Understanding dataset.
 
-The training and test data were generated by assigning symbolic variables 
-either positive or negative decimal integers and then describing the algebraic 
-operation to perform. We restrict our variable assignments to the range 
-x,y->[-1000,1000) and the operations to the set {+,-,*}. To ensure that the 
-model embraces symbolic variables, the order in which x and y appears in the 
-expression is randomly chosen. For instance, an input string contrasting from 
-the example shown above might be y=129,x=531,x-y. Each input string is 
-accompanied by its target string, which is the evaluation of the mathematical 
-expression. For this study, all targets considered are decimal integers 
-represented at the character level. About 12 million unique samples were thus 
-generated and randomly split into training and test sets at an approximate 
-ratio of 9:1, respectively. 
+The training and test data were generated by assigning symbolic variables
+either positive or negative decimal integers and then describing the algebraic
+operation to perform. We restrict our variable assignments to the range
+x,y->[-1000,1000) and the operations to the set {+,-,*}. To ensure that the
+model embraces symbolic variables, the order in which x and y appears in the
+expression is randomly chosen. For instance, an input string contrasting from
+the example shown above might be y=129,x=531,x-y. Each input string is
+accompanied by its target string, which is the evaluation of the mathematical
+expression. For this study, all targets considered are decimal integers
+represented at the character level. About 12 million unique samples were thus
+generated and randomly split into training and test sets at an approximate
+ratio of 9:1, respectively.
 
 For more information check the following paper:
-Artit Wangperawong. Attending to Mathematical Language with Transformers, 
-arXiv:1812.02825.
-Available at: https://arxiv.org/abs/1812.02825
-
+Artit Wangperawong. Attending to Mathematical Language with Transformers,
+arXiv:1812.02825 (https://arxiv.org/abs/1812.02825).
 """
 
 from __future__ import absolute_import
@@ -40,6 +38,7 @@
 from __future__ import print_function
 
 import os
+import tarfile
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -48,9 +47,13 @@
 
 import tensorflow as tf
 
+
 @registry.register_problem
 class MathematicalLanguageUnderstanding(text_problems.Text2TextProblem):
-  URL = "https://art.wangperawong.com/mathematical_language_understanding_train.tar.gz"
+  """Mathematical language understanding, see arxiv.org/abs/1812.02825."""
+
+  URL = ("https://art.wangperawong.com/mathematical_language_understanding"
+         "_train.tar.gz")
 
   @property
   def vocab_type(self):
@@ -71,34 +74,31 @@ def is_generate_per_split(self):
     return False
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    """Downloads and extracts the dataset and generates examples
+    """Downloads and extracts the dataset and generates examples.
 
     Args:
-      tmp_dir: temp directory to download and extract the dataset
       data_dir: The base directory where data and vocab files are stored.
+      tmp_dir: temp directory to download and extract the dataset.
+      dataset_split: split of the data-set.
 
-    Returns:
-      data generator
+    Yields:
+      The data examples.
     """
-
     if not tf.gfile.Exists(tmp_dir):
       tf.gfile.MakeDirs(tmp_dir)
 
     if not tf.gfile.Exists(data_dir):
       tf.gfile.MakeDirs(data_dir)
 
-    # Download and extract
+    # Download and extract.
     compressed_filename = os.path.basename(self.URL)
-    download_path = generator_utils.maybe_download(tmp_dir, compressed_filename,
-                                                   self.URL)
-
+    download_path = generator_utils.maybe_download(
+        tmp_dir, compressed_filename, self.URL)
     with tarfile.open(download_path, "r:gz") as tar:
       tar.extractall(tmp_dir)
-
-    filepath = os.path.join(tmp_dir, "mathematical_language_understanding_train.txt")
-
-    with open(filepath, 'r') as fp:
+    filepath = os.path.join(tmp_dir,
+                            "mathematical_language_understanding_train.txt")
+    with open(filepath, "r") as fp:
       for l in fp:
-        prob, ans = l.strip().split(':')
+        prob, ans = l.strip().split(":")
         yield {"inputs": prob, "targets": ans}
-
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 11f743c76..c523cdc89 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -240,10 +240,12 @@ def _greedy_infer(self, features, decode_length, use_tpu=False):
     Raises:
       NotImplementedError: If there are multiple data shards.
     """
-    return (self._slow_greedy_infer_tpu(features, decode_length) if use_tpu else
-            self._slow_greedy_infer(features, decode_length))
+    if use_tpu:
+      return self._slow_greedy_infer_tpu(features, decode_length)
+    return self._slow_greedy_infer(features, decode_length)
 
-  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha, use_tpu=False):
+  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha,
+                   use_tpu=False):
     """Beam search decoding.
 
     Args:

From 0f7cb55e5b9f547a20d9a705d05dc3abf29d47e2 Mon Sep 17 00:00:00 2001
From: Pablo Samuel Castro <psc@google.com>
Date: Tue, 11 Dec 2018 11:35:58 -0800
Subject: [PATCH 1348/2720] Refactor atari.run_experiment into
 common.run_experiment to enable support for non-Atari Gym environments.

PiperOrigin-RevId: 225040877
---
 tensor2tensor/rl/dopamine_connector.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index d18789e0e..7d765172a 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -22,7 +22,7 @@
 import copy
 
 from dopamine.agents.dqn import dqn_agent
-from dopamine.atari import run_experiment
+from dopamine.discrete_domains import run_experiment
 from dopamine.replay_memory import circular_replay_buffer
 from dopamine.replay_memory.circular_replay_buffer import OutOfGraphReplayBuffer
 from dopamine.replay_memory.circular_replay_buffer import ReplayElement
@@ -267,8 +267,6 @@ def create_runner(self, env_fn, hparams, target_iterations,
     agent_params.update(replay_buffer_params)
     create_agent_fn = get_create_agent(agent_params)
     runner = run_experiment.Runner(
-        game_name="unused_arg",
-        sticky_actions="unused_arg",
         base_dir=self.agent_model_dir,
         create_agent_fn=create_agent_fn,
         create_environment_fn=get_create_env_fun(

From b65ad32d9e33f7f6ffe69b1eb2f3d06bad14b419 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 11 Dec 2018 13:57:20 -0800
Subject: [PATCH 1349/2720] Unify camelcase <-> snakcase functions across T2T

PiperOrigin-RevId: 225065634
---
 tensor2tensor/data_generators/gym_env.py |  7 +--
 tensor2tensor/rl/rl_utils.py             |  4 +-
 tensor2tensor/utils/misc_utils.py        | 36 ++++++++++++++
 tensor2tensor/utils/misc_utils_test.py   | 60 ++++++++++++++++++++++++
 tensor2tensor/utils/modality.py          |  5 +-
 tensor2tensor/utils/registry.py          | 14 +-----
 tensor2tensor/utils/registry_test.py     | 13 -----
 7 files changed, 104 insertions(+), 35 deletions(-)
 create mode 100644 tensor2tensor/utils/misc_utils.py
 create mode 100644 tensor2tensor/utils/misc_utils_test.py

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 5b2a88004..341cd814e 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -31,6 +31,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.rl import gym_utils
 from tensor2tensor.utils import metrics
+from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -786,10 +787,6 @@ def frame_width(self):
 ]
 
 
-def camel_case_name(snake_case_name):
-  return "".join([w[0].upper() + w[1:] for w in snake_case_name.split("_")])
-
-
 def register_game(game_name, game_mode="NoFrameskip-v4"):
   """Create and register problems for the game.
 
@@ -804,7 +801,7 @@ def register_game(game_name, game_mode="NoFrameskip-v4"):
     raise ValueError("Game %s not in ATARI_GAMES" % game_name)
   if game_mode not in ATARI_GAME_MODES:
     raise ValueError("Unknown ATARI game mode: %s." % game_mode)
-  camel_game_name = camel_case_name(game_name) + game_mode
+  camel_game_name = misc_utils.snakecase_to_camelcase(game_name) + game_mode
   # Create and register the Problem
   cls = type("Gym%sRandom" % camel_game_name,
              (T2TGymEnv,), {"base_env_name": camel_game_name})
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 4b331856b..79797ae07 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -26,6 +26,7 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl.dopamine_connector import DQNLearner
 from tensor2tensor.rl.ppo_learner import PPOLearner
+from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -105,8 +106,7 @@ def evaluate_all_configs(hparams, agent_model_dir):
 def setup_env(hparams, batch_size, max_num_noops):
   """Setup."""
   game_mode = "Deterministic-v4"
-  camel_game_name = "".join(
-      [w[0].upper() + w[1:] for w in hparams.game.split("_")])
+  camel_game_name = misc_utils.snakecase_to_camelcase(hparams.game)
   camel_game_name += game_mode
   env_name = camel_game_name
 
diff --git a/tensor2tensor/utils/misc_utils.py b/tensor2tensor/utils/misc_utils.py
new file mode 100644
index 000000000..d8d1b18da
--- /dev/null
+++ b/tensor2tensor/utils/misc_utils.py
@@ -0,0 +1,36 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Miscellaneous utilities."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import re
+
+# Camel case to snake case utils
+_first_cap_re = re.compile("(.)([A-Z][a-z0-9]+)")
+_all_cap_re = re.compile("([a-z0-9])([A-Z])")
+
+
+def camelcase_to_snakecase(name):
+  s1 = _first_cap_re.sub(r"\1_\2", name)
+  return _all_cap_re.sub(r"\1_\2", s1).lower()
+
+
+def snakecase_to_camelcase(name):
+  return "".join([w[0].upper() + w[1:] for w in name.split("_")])
+
diff --git a/tensor2tensor/utils/misc_utils_test.py b/tensor2tensor/utils/misc_utils_test.py
new file mode 100644
index 000000000..3a28fa629
--- /dev/null
+++ b/tensor2tensor/utils/misc_utils_test.py
@@ -0,0 +1,60 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.utils.misc_utils."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.utils import misc_utils
+import tensorflow as tf
+
+
+class MiscUtilsTest(tf.test.TestCase):
+
+  def test_camelcase_to_snakecase(self):
+    self.assertEqual("typical_camel_case",
+                     misc_utils.camelcase_to_snakecase("TypicalCamelCase"))
+    self.assertEqual("numbers_fuse2gether",
+                     misc_utils.camelcase_to_snakecase("NumbersFuse2gether"))
+    self.assertEqual("numbers_fuse2_gether",
+                     misc_utils.camelcase_to_snakecase("NumbersFuse2Gether"))
+    self.assertEqual("lstm_seq2_seq",
+                     misc_utils.camelcase_to_snakecase("LSTMSeq2Seq"))
+    self.assertEqual("starts_lower",
+                     misc_utils.camelcase_to_snakecase("startsLower"))
+    self.assertEqual("starts_lower_caps",
+                     misc_utils.camelcase_to_snakecase("startsLowerCAPS"))
+    self.assertEqual("caps_fuse_together",
+                     misc_utils.camelcase_to_snakecase("CapsFUSETogether"))
+    self.assertEqual("startscap",
+                     misc_utils.camelcase_to_snakecase("Startscap"))
+    self.assertEqual("s_tartscap",
+                     misc_utils.camelcase_to_snakecase("STartscap"))
+
+  def test_snakecase_to_camelcase(self):
+    self.assertEqual("TypicalCamelCase",
+                     misc_utils.snakecase_to_camelcase("typical_camel_case"))
+    self.assertEqual("NumbersFuse2gether",
+                     misc_utils.snakecase_to_camelcase("numbers_fuse2gether"))
+    self.assertEqual("NumbersFuse2Gether",
+                     misc_utils.snakecase_to_camelcase("numbers_fuse2_gether"))
+    self.assertEqual("LstmSeq2Seq",
+                     misc_utils.snakecase_to_camelcase("lstm_seq2_seq"))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index 45b4eb76a..6041ab310 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -18,8 +18,8 @@
 from __future__ import division
 from __future__ import print_function
 
-import re
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import misc_utils
 
 import tensorflow as tf
 
@@ -58,8 +58,7 @@ def __init__(self, model_hparams, vocab_size=None):
 
   @property
   def name(self):
-    camelcase_name = type(self).__name__  # DeCamelCase for TF readability.
-    return re.sub("([A-Z]+)", r"_\1", camelcase_name).lower()[1:]
+    return misc_utils.camelcase_to_snakecase(type(self).__name__)
 
   @property
   def top_dimensionality(self):
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index fa10dacc9..6b5046382 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -45,7 +45,7 @@ class MyModel(T2TModel):
 from __future__ import print_function
 
 import inspect
-import re
+from tensor2tensor.utils import misc_utils
 import tensorflow as tf
 
 _ATTACKS = {}
@@ -58,16 +58,6 @@ class MyModel(T2TModel):
 _RANGED_HPARAMS = {}
 
 
-# Camel case to snake case utils
-_first_cap_re = re.compile("(.)([A-Z][a-z0-9]+)")
-_all_cap_re = re.compile("([a-z0-9])([A-Z])")
-
-
-def _convert_camel_to_snake(name):
-  s1 = _first_cap_re.sub(r"\1_\2", name)
-  return _all_cap_re.sub(r"\1_\2", s1).lower()
-
-
 def _reset():
   for ctr in [_MODELS, _HPARAMS, _RANGED_HPARAMS, _ATTACK_PARAMS]:
     ctr.clear()
@@ -82,7 +72,7 @@ def default_name(obj_class):
   Returns:
     The registry's default name for the class.
   """
-  return _convert_camel_to_snake(obj_class.__name__)
+  return misc_utils.camelcase_to_snakecase(obj_class.__name__)
 
 
 def default_object_name(obj):
diff --git a/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index a72dd36bd..b56c9c48c 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -87,19 +87,6 @@ def m2():
 
     self.assertSetEqual(set(["m1", "m2"]), set(registry.list_models()))
 
-  def testSnakeCase(self):
-    convert = registry._convert_camel_to_snake
-
-    self.assertEqual("typical_camel_case", convert("TypicalCamelCase"))
-    self.assertEqual("numbers_fuse2gether", convert("NumbersFuse2gether"))
-    self.assertEqual("numbers_fuse2_gether", convert("NumbersFuse2Gether"))
-    self.assertEqual("lstm_seq2_seq", convert("LSTMSeq2Seq"))
-    self.assertEqual("starts_lower", convert("startsLower"))
-    self.assertEqual("starts_lower_caps", convert("startsLowerCAPS"))
-    self.assertEqual("caps_fuse_together", convert("CapsFUSETogether"))
-    self.assertEqual("startscap", convert("Startscap"))
-    self.assertEqual("s_tartscap", convert("STartscap"))
-
 
 class HParamRegistryTest(tf.test.TestCase):
 

From af7c047394a7438287fa94510850db488c1e2519 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 11 Dec 2018 15:12:01 -0800
Subject: [PATCH 1350/2720] Add Gated activation unit previously used in the
 conditional PixelCNN and WaveNet papers.

PiperOrigin-RevId: 225079194
---
 tensor2tensor/models/research/glow.py         |  2 +
 tensor2tensor/models/research/glow_ops.py     | 67 +++++++++++++------
 .../models/research/glow_ops_test.py          | 14 ++--
 tensor2tensor/models/video/next_frame_glow.py |  1 +
 tensor2tensor/models/video/nfg_conv3d_test.py | 11 +--
 tensor2tensor/models/video/nfg_test_utils.py  |  9 +--
 6 files changed, 71 insertions(+), 33 deletions(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index c3c01192b..ab008dad0 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -49,6 +49,8 @@ def glow_hparams():
   hparams.add_hparam("n_levels", 3)
   hparams.add_hparam("n_bits_x", 8)
   hparams.add_hparam("depth", 32)
+  # Activation - Relu or Gatu
+  hparams.add_hparam("activation", "relu")
   # Coupling layer, additive or affine.
   hparams.add_hparam("coupling", "affine")
   hparams.add_hparam("coupling_width", 512)
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 1a029fb74..71412b5f7 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -453,7 +453,7 @@ def conv(name, x, output_channels, filter_size=None, stride=None,
 
 
 @add_arg_scope
-def conv_block(name, x, mid_channels, dilations=None):
+def conv_block(name, x, mid_channels, dilations=None, activation="relu"):
   """2 layer conv block used in the affine coupling layer.
 
   Args:
@@ -461,6 +461,9 @@ def conv_block(name, x, mid_channels, dilations=None):
     x: 4-D or 5-D Tensor.
     mid_channels: Output channels of the second layer.
     dilations: Optional, list of integers.
+    activation: relu or gatu.
+      If relu, the second layer is relu(W*x)
+      If gatu, the second layer is tanh(W1*x) * sigmoid(W2*x)
   Returns:
     x: 4-D Tensor: Output activations.
   """
@@ -481,16 +484,24 @@ def conv_block(name, x, mid_channels, dilations=None):
              dilations=dilations)
     x = tf.nn.relu(x)
 
-    # Padding + conv2d + actnorm + relu
+    # Padding + conv2d + actnorm + activation.
     # [input, output: 512 channels]
-    x = conv("1_2", x, output_channels=mid_channels, filter_size=second_filter,
-             dilations=dilations)
-    x = tf.nn.relu(x)
+    if activation == "relu":
+      x = conv("1_2", x, output_channels=mid_channels,
+               filter_size=second_filter, dilations=dilations)
+      x = tf.nn.relu(x)
+    elif activation == "gatu":
+      # x = tanh(w1*x) * sigm(w2*x)
+      x_tanh = conv("1_tanh", x, output_channels=mid_channels,
+                    filter_size=second_filter, dilations=dilations)
+      x_sigm = conv("1_sigm", x, output_channels=mid_channels,
+                    filter_size=second_filter, dilations=dilations)
+      x = tf.nn.tanh(x_tanh) * tf.nn.sigmoid(x_sigm)
     return x
 
 
 def dilated_conv_stack(name, x, mid_channels, output_channels,
-                       dilation_rates):
+                       dilation_rates, activation="relu"):
   """Dilated convolutional stack.
 
   Features at different rates are computed independently using a 3 layer
@@ -503,6 +514,7 @@ def dilated_conv_stack(name, x, mid_channels, output_channels,
                   stack.
     output_channels: Number of output channels of the last layer.
     dilation_rates: A list of dilation rates.
+    activation: Can be either "relu" or "gatu"
   Returns:
     output: 5-D Tensor.
   """
@@ -511,13 +523,15 @@ def dilated_conv_stack(name, x, mid_channels, output_channels,
     for dil_ind, dil_rate in enumerate(dilation_rates):
       # TODO(mechcoder) try (concat across channels + 1x1) modulo memory issues.
       curr_out = conv_stack("dil_%d" % dil_ind, x, mid_channels=mid_channels,
-                            output_channels=output_channels, dilations=dil_rate)
+                            output_channels=output_channels, dilations=dil_rate,
+                            activation=activation)
       output += curr_out
     return output
 
 
 @add_arg_scope
-def conv_stack(name, x, mid_channels, output_channels, dilations=None):
+def conv_stack(name, x, mid_channels, output_channels, dilations=None,
+               activation="relu"):
   """3-layer convolutional stack.
 
   Args:
@@ -527,14 +541,16 @@ def conv_stack(name, x, mid_channels, output_channels, dilations=None):
     output_channels: Number of output channels.
     dilations: Dilations to apply in the first 3x3 layer and the last 3x3 layer.
                By default, apply no dilations.
-
+    activation: relu or gatu.
+      If relu, the second layer is relu(W*x)
+      If gatu, the second layer is tanh(W1*x) * sigmoid(W2*x)
   Returns:
     output: output of 3 layer conv network.
   """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
 
     x = conv_block("conv_block", x, mid_channels=mid_channels,
-                   dilations=dilations)
+                   dilations=dilations, activation=activation)
 
     # Final layer.
     x = conv("zeros", x, apply_actnorm=False, conv_init="zeros",
@@ -543,7 +559,8 @@ def conv_stack(name, x, mid_channels, output_channels, dilations=None):
 
 
 @add_arg_scope
-def additive_coupling(name, x, mid_channels=512, reverse=False):
+def additive_coupling(name, x, mid_channels=512, reverse=False,
+                      activation="relu"):
   """Reversible additive coupling layer.
 
   Args:
@@ -551,6 +568,7 @@ def additive_coupling(name, x, mid_channels=512, reverse=False):
     x: 4-D Tensor.
     mid_channels: number of channels in the coupling layer.
     reverse: Forward or reverse operation.
+    activation: "relu" or "gatu"
   Returns:
     output:
     objective: 0.0
@@ -560,7 +578,8 @@ def additive_coupling(name, x, mid_channels=512, reverse=False):
     x1, x2 = tf.split(x, num_or_size_splits=2, axis=-1)
 
     z1 = x1
-    shift = conv_stack("nn", x1, mid_channels, output_channels=output_channels)
+    shift = conv_stack("nn", x1, mid_channels, output_channels=output_channels,
+                       activation=activation)
 
     if not reverse:
       z2 = x2 + shift
@@ -570,17 +589,19 @@ def additive_coupling(name, x, mid_channels=512, reverse=False):
 
 
 @add_arg_scope
-def affine_coupling(name, x, mid_channels=512, reverse=False):
+def affine_coupling(name, x, mid_channels=512, activation="relu",
+                    reverse=False):
   """Reversible affine coupling layer.
 
   Args:
     name: variable scope.
     x: 4-D Tensor.
     mid_channels: number of channels in the coupling layer.
+    activation: Can be either "relu" or "gatu".
     reverse: Forward or reverse operation.
   Returns:
-    output:
-    objective:
+    output: input s
+    objective: log-determinant of the jacobian
   """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     x_shape = common_layers.shape_list(x)
@@ -592,7 +613,8 @@ def affine_coupling(name, x, mid_channels=512, reverse=False):
     # Else:
     # z2 = (x2 / scale) - shift
     z1 = x1
-    log_scale_and_shift = conv_stack("nn", x1, mid_channels, x_shape[-1])
+    log_scale_and_shift = conv_stack(
+        "nn", x1, mid_channels, x_shape[-1], activation=activation)
     shift = log_scale_and_shift[:, :, :, 0::2]
     scale = tf.nn.sigmoid(log_scale_and_shift[:, :, :, 1::2] + 2.0)
     if not reverse:
@@ -684,7 +706,6 @@ def temporal_latent_to_dist(name, x, hparams, output_channels=None):
   if output_channels is None:
     output_channels = res_channels
   dilation_rates = get_dilation_rates(hparams, width)
-
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     h = x
     for i in range(hparams.latent_encoder_depth):
@@ -692,11 +713,13 @@ def temporal_latent_to_dist(name, x, hparams, output_channels=None):
         h2 = dilated_conv_stack("dil_latent_3d_res_%d" % i, h,
                                 mid_channels=hparams.latent_encoder_width,
                                 output_channels=res_channels,
-                                dilation_rates=dilation_rates)
+                                dilation_rates=dilation_rates,
+                                activation=hparams.latent_activation)
       else:
         h2 = conv_stack("latent_3d_res_%d" % i, h,
                         mid_channels=hparams.latent_encoder_width,
-                        output_channels=res_channels)
+                        output_channels=res_channels,
+                        activation=hparams.latent_activation)
       h += h2
 
     # take last activation that should capture all context since padding is
@@ -1006,11 +1029,13 @@ def revnet_step(name, x, hparams, reverse=True):
     if hparams.coupling == "additive":
       coupling_layer = functools.partial(
           additive_coupling, name="additive", reverse=reverse,
-          mid_channels=hparams.coupling_width)
+          mid_channels=hparams.coupling_width,
+          activation=hparams.activation)
     else:
       coupling_layer = functools.partial(
           affine_coupling, name="affine", reverse=reverse,
-          mid_channels=hparams.coupling_width)
+          mid_channels=hparams.coupling_width,
+          activation=hparams.activation)
     ops = [
         functools.partial(actnorm, name="actnorm", reverse=reverse),
         functools.partial(invertible_1x1_conv, name="invertible",
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index ef71d6fdb..c131d007e 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -50,6 +50,7 @@ def get_glow_hparams(self):
     hparams.add_hparam("latent_pre_output_channels", 256)
     hparams.add_hparam("latent_dist_encoder", "conv_net")
     hparams.add_hparam("latent_time_filter_size", 3)
+    hparams.add_hparam("latent_activation", "relu")
     return hparams
 
   def test_get_variable_ddi(self):
@@ -126,11 +127,14 @@ def test_conv2d(self):
         # test shape in case apply_actnorm is set to False,
         self.assertEqual(zeros_np.shape, (16, 5, 5, 64))
 
-  def test_conv_stack(self):
+  @parameterized.named_parameters(
+      ("relu_act", "relu"), ("gatu_act", "gatu"))
+  def test_conv_stack(self, activation="relu"):
     """Test output shape."""
     with tf.Graph().as_default():
       x = 10.0 * tf.random_uniform(shape=(16, 5, 5, 32))
-      nn = glow_ops.conv_stack("nn", x, 512, 64)
+      nn = glow_ops.conv_stack("nn", x, mid_channels=512, output_channels=64,
+                               activation=activation)
 
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())
@@ -406,11 +410,13 @@ def test_actnorm_3d(self):
           self.assertTrue(np.allclose(channel_var, 1.0, atol=1e-3))
 
   @parameterized.named_parameters(
-      ("dilation", True), ("no_dilation", False))
-  def test_temporal_latent_to_dist(self, apply_dilation):
+      ("dil_relu", True, "relu"), ("no_dil_relu", False, "relu"),
+      ("dil_gatu", True, "gatu"), ("no_dil_gatu", False, "gatu"),)
+  def test_temporal_latent_to_dist(self, apply_dilation, activation):
     with tf.Graph().as_default():
       hparams = self.get_glow_hparams()
       hparams.latent_apply_dilations = apply_dilation
+      hparams.latent_activation = activation
       latent_shape = (16, 5, 32, 32, 48)
       latents = tf.random_normal(latent_shape)
       dist = glow_ops.temporal_latent_to_dist(
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index dd1251946..d6e0f3d2a 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -62,6 +62,7 @@ def next_frame_glow_hparams():
   hparams.add_hparam("latent_encoder_depth", 2)
   hparams.add_hparam("latent_encoder_width", 512)
   hparams.add_hparam("latent_pre_output_channels", 512)
+  hparams.add_hparam("latent_activation", "relu")
   # Pretrains the glow encoder for "pretrain_steps" number of steps.
   # By default, don't pretrain and learn end-to-end
   hparams.add_hparam("pretrain_steps", -1)
diff --git a/tensor2tensor/models/video/nfg_conv3d_test.py b/tensor2tensor/models/video/nfg_conv3d_test.py
index 1c189808d..6bf5549b4 100644
--- a/tensor2tensor/models/video/nfg_conv3d_test.py
+++ b/tensor2tensor/models/video/nfg_conv3d_test.py
@@ -24,8 +24,10 @@
 import tensorflow as tf
 
 conv3d_net_hparams = (
-    ("conv3d_net", 2, 2, "conv3d_net", "conditional", -1, 3),
-    ("conv3d_dil", 2, 2, "conv3d_net", "conditional", -1, -1, False, True),)
+    # ("conv3d_net", 2, 2, "conv3d_net", "conditional", -1, 3),
+    ("conv3d_net_gatu", 2, 2, "conv3d_net", "conditional", -1, 3, False, False,
+     "gatu"),)
+    # ("conv3d_dil", 2, 2, "conv3d_net", "conditional", -1, -1, False, True),)
 
 
 class NextFrameGlowConv3DTest(nfg_test_utils.NextFrameGlowTest,
@@ -36,12 +38,13 @@ def testGlowTrainAndDecode(self, in_frames=1, out_frames=1,
                              latent_dist_encoder="pointwise",
                              gen_mode="conditional", pretrain_steps=-1,
                              num_train_frames=-1, cond_first_frame=False,
-                             apply_dilations=False):
+                             apply_dilations=False, activation="relu"):
     self.GlowTrainAndDecode(
         in_frames=in_frames, out_frames=out_frames,
         latent_dist_encoder=latent_dist_encoder, gen_mode=gen_mode,
         pretrain_steps=pretrain_steps, num_train_frames=num_train_frames,
-        cond_first_frame=cond_first_frame, apply_dilations=apply_dilations)
+        cond_first_frame=cond_first_frame, apply_dilations=apply_dilations,
+        activation=activation)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/models/video/nfg_test_utils.py b/tensor2tensor/models/video/nfg_test_utils.py
index 242463368..01744e9de 100644
--- a/tensor2tensor/models/video/nfg_test_utils.py
+++ b/tensor2tensor/models/video/nfg_test_utils.py
@@ -34,8 +34,9 @@
 def fill_hparams(hparams, in_frames, out_frames, gen_mode="conditional",
                  latent_dist_encoder="pointwise", pretrain_steps=-1,
                  num_train_frames=-1, cond_first_frame=False,
-                 apply_dilations=False):
+                 apply_dilations=False, activation="relu"):
   """Set next_frame_glow hparams."""
+  hparams.latent_activation = activation
   hparams.latent_apply_dilations = apply_dilations
   hparams.video_num_input_frames = in_frames
   hparams.video_num_target_frames = out_frames
@@ -129,7 +130,7 @@ def GlowTrainAndDecode(self, in_frames=1, out_frames=1,
                          latent_dist_encoder="pointwise",
                          gen_mode="conditional", pretrain_steps=-1,
                          num_train_frames=-1, cond_first_frame=False,
-                         apply_dilations=False):
+                         apply_dilations=False, activation="relu"):
     """Test 1 forward pass and sampling gives reasonable results."""
     if num_train_frames == -1:
       total_frames = in_frames + out_frames
@@ -145,7 +146,7 @@ def GlowTrainAndDecode(self, in_frames=1, out_frames=1,
       hparams = fill_hparams(hparams, in_frames, out_frames,
                              gen_mode, latent_dist_encoder, pretrain_steps,
                              num_train_frames, cond_first_frame,
-                             apply_dilations)
+                             apply_dilations, activation)
       features = create_basic_features(hparams)
       model = next_frame_glow.NextFrameGlow(hparams, MODES.TRAIN)
       _, train_op = model(features)
@@ -161,7 +162,7 @@ def GlowTrainAndDecode(self, in_frames=1, out_frames=1,
       hparams = fill_hparams(hparams, in_frames, out_frames,
                              gen_mode, latent_dist_encoder, pretrain_steps,
                              num_train_frames, cond_first_frame,
-                             apply_dilations)
+                             apply_dilations, activation)
       features = create_basic_features(hparams)
       model = next_frame_glow.NextFrameGlow(
           hparams, tf.estimator.ModeKeys.PREDICT)

From 5e7718a8b29f0cf7ec1b65eb60a9099dfcfb5961 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 11 Dec 2018 15:26:44 -0800
Subject: [PATCH 1351/2720] Multi-lingual versions of summarization and MNLI
 and lint corrections.

PiperOrigin-RevId: 225081761
---
 .../data_generators/cnn_dailymail.py          |  9 +++++
 tensor2tensor/data_generators/multinli.py     | 17 +++++++--
 .../data_generators/wiki_multi_problems.py    |  3 ++
 tensor2tensor/models/transformer.py           | 38 +++++++++++++------
 4 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index 177b1a4ed..9380181e6 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -264,3 +264,12 @@ class SummarizeCnnDailymailWikiLMSharedVocab64k(SummarizeCnnDailymail32k):
   @property
   def vocab_filename(self):
     return wiki_lm.LanguagemodelEnWiki64k().vocab_filename
+
+
+@registry.register_problem
+class SummarizeCnnDailymailWikiLMMultiVocab64k(SummarizeCnnDailymail32k):
+  """Summarize CNN and Daily Mail articles using multi-lingual 64k vocab."""
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
diff --git a/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index dcb94d1b8..70ee0107f 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -118,7 +118,7 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
 
 @registry.register_problem
 class MultiNLICharacters(MultiNLI):
-  """MultiNLI classification problems, character level"""
+  """MultiNLI classification problems, character level."""
 
   @property
   def vocab_type(self):
@@ -130,7 +130,7 @@ def global_task_id(self):
 
 @registry.register_problem
 class MultiNLISharedVocab(MultiNLI):
-  """MultiNLI classification problems with the LM1b vocabulary"""
+  """MultiNLI classification problems with the LM1b vocabulary."""
 
   @property
   def vocab_filename(self):
@@ -139,7 +139,7 @@ def vocab_filename(self):
 
 @registry.register_problem
 class MultiNLIWikiLMSharedVocab(MultiNLI):
-  """MultiNLI classification problems with the Wiki vocabulary"""
+  """MultiNLI classification problems with the Wiki vocabulary."""
 
   @property
   def vocab_filename(self):
@@ -148,8 +148,17 @@ def vocab_filename(self):
 
 @registry.register_problem
 class MultiNLIWikiLMSharedVocab64k(MultiNLIWikiLMSharedVocab):
-  """MultiNLI classification problems with the Wiki vocabulary"""
+  """MultiNLI classification problems with the Wiki vocabulary."""
 
   @property
   def vocab_filename(self):
     return wiki_lm.LanguagemodelEnWiki64k().vocab_filename
+
+
+@registry.register_problem
+class MultiNLIWikiLMMultiVocab64k(MultiNLIWikiLMSharedVocab):
+  """MultiNLI classification problems with the multi-lingual vocabulary."""
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
diff --git a/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index fafe6d074..7e4c7efe3 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -140,6 +140,9 @@ def __init__(self, was_reversed=False, was_copy=False):
         was_reversed=True))
     self.task_list.append(translate_enro.TranslateEnroWmtMultiSmall64k(
         was_reversed=True))
+    self.task_list.append(
+        cnn_dailymail.SummarizeCnnDailymailWikiLMMultiVocab64k())
+    self.task_list.append(multinli.MultiNLIWikiLMMultiVocab64k())
 
   @property
   def vocab_type(self):
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index e4dc0fd01..5407ebe5f 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -242,8 +242,9 @@ def _greedy_infer(self, features, decode_length, use_tpu=False):
         self._hparams.self_attention_type != "dot_product"):
       return  super(Transformer, self)._greedy_infer(features, decode_length)
     with tf.variable_scope(self.name):
-      return (self._fast_decode_tpu(features, decode_length) if use_tpu else
-              self._fast_decode(features, decode_length))
+      if use_tpu:
+        return self._fast_decode_tpu(features, decode_length)
+      return self._fast_decode(features, decode_length)
 
   def _beam_decode(self,
                    features,
@@ -280,11 +281,11 @@ def _beam_decode(self,
       return self._beam_decode_slow(features, decode_length, beam_size,
                                     top_beams, alpha, use_tpu)
     with tf.variable_scope(self.name):
-      return (
-          self._fast_decode_tpu(
-              features, decode_length, beam_size, top_beams, alpha) if use_tpu
-          else self._fast_decode(
-              features, decode_length, beam_size, top_beams, alpha))
+      if use_tpu:
+        return self._fast_decode_tpu(
+            features, decode_length, beam_size, top_beams, alpha)
+      return self._fast_decode(
+          features, decode_length, beam_size, top_beams, alpha)
 
   def _fast_decode_tpu(self,
                        features,
@@ -831,8 +832,9 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       """One step of greedy decoding."""
       logits, cache = symbols_to_logits_fn(next_id, i, cache)
       log_probs = common_layers.log_prob_from_logits(logits)
-      temperature = (0.0 if hparams.sampling_method == "argmax" else
-                     hparams.sampling_temp)
+      temperature = hparams.sampling_temp
+      if hparams.sampling_method == "argmax":
+        temperature = 0.0
       next_id = common_layers.sample_with_temperature(logits, temperature)
       hit_eos |= tf.equal(next_id, eos_id)
 
@@ -998,8 +1000,9 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       """One step of greedy decoding."""
       logits, cache = symbols_to_logits_fn(next_id, i, cache)
       log_probs = common_layers.log_prob_from_logits(logits)
-      temperature = (0.0 if hparams.sampling_method == "argmax" else
-                     hparams.sampling_temp)
+      temperature = hparams.sampling_temp
+      if hparams.sampling_method == "argmax":
+        temperature = 0.0
       next_id = common_layers.sample_with_temperature(logits, temperature)
       hit_eos |= tf.equal(next_id, eos_id)
 
@@ -1592,6 +1595,19 @@ def transformer_tall_pretrain_lm_tpu_adafactor():
   return hparams
 
 
+@registry.register_hparams
+def transformer_tall_pretrain_lm_tpu_adafactor_large():
+  """Hparams for transformer on LM pretraining on TPU, large model."""
+  hparams = transformer_tall_pretrain_lm_tpu_adafactor()
+  hparams.hidden_size = 1024
+  hparams.num_heads = 16
+  hparams.filter_size = 32768
+  hparams.batch_size = 4
+  hparams.multiproblem_mixing_schedule = "constant"
+  hparams.multiproblem_schedule_threshold = 0.3
+  return hparams
+
+
 @registry.register_hparams
 def transformer_tall_pretrain_lm_tpu():
   """Hparams for transformer on LM pretraining on TPU with AdamW."""

From ce0ff1baf0f268f32e71c3b5f76fcf82e7424dff Mon Sep 17 00:00:00 2001
From: Mike Dusenberry <dusenberrymw@google.com>
Date: Tue, 11 Dec 2018 18:07:21 -0800
Subject: [PATCH 1352/2720] Use Independent ed RVs in the Bayesian layers to
 enable scalar KL divergences.

PiperOrigin-RevId: 225106079
---
 tensor2tensor/layers/bayes.py      | 11 +++++++++--
 tensor2tensor/layers/bayes_test.py | 10 +++++-----
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 2a8af7fff..fad9b6f21 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -114,7 +114,9 @@ def __call__(self, shape=None, dtype=None, partition_info=None):
       raise ValueError('A TrainableInitializer must be built by a layer before '
                        'usage, and is currently only compatible with Bayesian '
                        'layers.')
-    return ed.Normal(loc=self.mean, scale=self.stddev)
+    return ed.Independent(
+        ed.Normal(loc=self.mean, scale=self.stddev).distribution,
+        reinterpreted_batch_ndims=len(self.shape))
 
   def get_config(self):
     return {
@@ -151,7 +153,12 @@ def __call__(self, x):
     """Computes regularization given an ed.Normal random variable as input."""
     if not isinstance(x, ed.RandomVariable):
       raise ValueError('Input must be an ed.RandomVariable.')
-    random_variable = ed.Normal(loc=self.mean, scale=self.stddev)
+    random_variable = ed.Independent(
+        ed.Normal(
+            loc=tf.broadcast_to(self.mean, x.distribution.event_shape),
+            scale=tf.broadcast_to(self.stddev, x.distribution.event_shape)
+        ).distribution,
+        reinterpreted_batch_ndims=len(x.distribution.event_shape))
     return random_variable.distribution.kl_divergence(x.distribution)
 
   def get_config(self):
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index e53661992..0f1697a62 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -35,7 +35,7 @@ def testTrainableNormalStddevConstraint(self):
         100, kernel_initializer=bayes.TrainableNormal())
     inputs = tf.random_normal([1, 1])
     out = layer(inputs)
-    stddev = layer.kernel.distribution.scale
+    stddev = layer.kernel.distribution.stddev()
     self.evaluate(tf.global_variables_initializer())
     res, _ = self.evaluate([stddev, out])
     self.assertAllGreater(res, 0.)
@@ -76,7 +76,7 @@ def testDenseReparameterizationKL(self):
     with tf.GradientTape() as tape:
       layer(inputs)  # first call forces a build, here inside this tape
       layer(inputs)  # ensure robustness after multiple calls
-      loss = tf.reduce_sum([tf.reduce_sum(l) for l in layer.losses])
+      loss = sum(layer.losses)
 
     variables = [layer.kernel_initializer.mean, layer.kernel_initializer.stddev]
     for v in variables:
@@ -91,7 +91,7 @@ def testDenseReparameterizationKL(self):
     # Imagine this is the 2nd epoch.
     with tf.GradientTape() as tape:
       layer(inputs)  # build won't be called again
-      loss = tf.reduce_sum([tf.reduce_sum(l) for l in layer.losses])
+      loss = sum(layer.losses)
 
     variables = [layer.kernel_initializer.mean, layer.kernel_initializer.stddev]
     for v in variables:
@@ -175,7 +175,7 @@ def testLSTMCellReparameterizationKL(self):
       cell(inputs[:, 0, :], state)  # ensure robustness after multiple calls
       cell.get_initial_state(inputs[:, 0, :])
       cell(inputs[:, 0, :], state)  # ensure robustness after multiple calls
-      loss = tf.reduce_sum([tf.reduce_sum(l) for l in cell.losses])
+      loss = sum(cell.losses)
 
     variables = [
         cell.kernel_initializer.mean, cell.kernel_initializer.stddev,
@@ -193,7 +193,7 @@ def testLSTMCellReparameterizationKL(self):
     # Imagine this is the 2nd epoch.
     with tf.GradientTape() as tape:
       cell(inputs[:, 0, :], state)  # build won't be called again
-      loss = tf.reduce_sum([tf.reduce_sum(l) for l in cell.losses])
+      loss = sum(cell.losses)
 
     variables = [
         cell.kernel_initializer.mean, cell.kernel_initializer.stddev,

From 168928d29cdd887fd6cdd5d5ad35181bef614154 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 11 Dec 2018 18:26:27 -0800
Subject: [PATCH 1353/2720] Standardize to rlmf and rlmb.

Before this model free was mfrl but model based was rlmb.

PiperOrigin-RevId: 225108128
---
 tensor2tensor/models/research/rl.py         | 10 +++++-----
 tensor2tensor/rl/trainer_model_free_test.py |  2 +-
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 8c93504ef..a56823f98 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -226,7 +226,7 @@ def dqn_original_params():
 
 
 @registry.register_hparams
-def mfrl_original():
+def rlmf_original():
   return tf.contrib.training.HParams(
       game="pong",
       base_algo="ppo",
@@ -244,17 +244,17 @@ def mfrl_original():
 
 
 @registry.register_hparams
-def mfrl_base():
+def rlmf_base():
   """Base set of hparams for model-free PPO."""
-  hparams = mfrl_original()
+  hparams = rlmf_original()
   hparams.add_hparam("ppo_epochs_num", 3000)
   hparams.add_hparam("ppo_eval_every_epochs", 100)
   return hparams
 
 
 @registry.register_hparams
-def mfrl_tiny():
-  hparams = mfrl_base()
+def rlmf_tiny():
+  hparams = rlmf_base()
   hparams.ppo_epochs_num = 100
   hparams.ppo_eval_every_epochs = 10
   return hparams
diff --git a/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index 678641847..aa75aac9c 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -29,7 +29,7 @@
 class TrainTest(tf.test.TestCase):
 
   def test_train_pong(self):
-    hparams = registry.hparams("mfrl_original")
+    hparams = registry.hparams("rlmf_original")
     hparams.batch_size = 2
     hparams.eval_sampling_temps = [0.0, 1.0]
     hparams.add_hparam("ppo_epochs_num", 2)

From 4ae4f254e3363639361a0eabe67ffbe93c5f9747 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 11 Dec 2018 19:28:45 -0800
Subject: [PATCH 1354/2720] Correct broken use-case when hparams has no
 sampling temp.

PiperOrigin-RevId: 225113646
---
 tensor2tensor/models/transformer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 5407ebe5f..d51013099 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -832,7 +832,7 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       """One step of greedy decoding."""
       logits, cache = symbols_to_logits_fn(next_id, i, cache)
       log_probs = common_layers.log_prob_from_logits(logits)
-      temperature = hparams.sampling_temp
+      temperature = getattr(hparams, "sampling_temp", 0.0)
       if hparams.sampling_method == "argmax":
         temperature = 0.0
       next_id = common_layers.sample_with_temperature(logits, temperature)
@@ -1000,7 +1000,7 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       """One step of greedy decoding."""
       logits, cache = symbols_to_logits_fn(next_id, i, cache)
       log_probs = common_layers.log_prob_from_logits(logits)
-      temperature = hparams.sampling_temp
+      temperature = getattr(hparams, "sampling_temp", 0.0)
       if hparams.sampling_method == "argmax":
         temperature = 0.0
       next_id = common_layers.sample_with_temperature(logits, temperature)

From a62c5112dbeb34af1175f36a79931aced74bc747 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 12 Dec 2018 10:59:05 -0800
Subject: [PATCH 1355/2720] Add MADE, Reverse; split reversible layers to new
 file.

PiperOrigin-RevId: 225216004
---
 tensor2tensor/layers/reversible_layers.py     | 209 ++++++++++++++++++
 .../layers/reversible_layers_test.py          |  90 ++++++++
 2 files changed, 299 insertions(+)
 create mode 100644 tensor2tensor/layers/reversible_layers.py
 create mode 100644 tensor2tensor/layers/reversible_layers_test.py

diff --git a/tensor2tensor/layers/reversible_layers.py b/tensor2tensor/layers/reversible_layers.py
new file mode 100644
index 000000000..112859903
--- /dev/null
+++ b/tensor2tensor/layers/reversible_layers.py
@@ -0,0 +1,209 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Reversible layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import tensorflow as tf
+
+
+class Reverse(tf.keras.layers.Layer):
+  """Swaps the forward and reverse transformations of a layer."""
+
+  def __init__(self, reversible_layer, **kwargs):
+    super(Reverse, self).__init__(**kwargs)
+    if not hasattr(reversible_layer, 'reverse'):
+      raise ValueError('Layer passed-in has not implemented "reverse" method: '
+                       '{}'.format(reversible_layer))
+    self.call = reversible_layer.reverse
+    self.reverse = reversible_layer.call
+
+
+class MADE(tf.keras.Model):
+  """Masked autoencoder for distribution estimation (Germain et al., 2015).
+
+  MADE takes as input a real Tensor of shape [..., length] and returns a
+  Tensor of shape [..., 2 * length] and same dtype. It masks layer weights to
+  respect autoregressive constraints: for a given ordering, each input dimension
+  can be reconstructed from previous input dimensions. The output dimensions
+  represent two heads for, e.g., location and scale transforms in a flow.
+  """
+
+  def __init__(self,
+               hidden_dims,
+               input_order='left-to-right',
+               hidden_order='left-to-right',
+               activation=None,
+               use_bias=True,
+               **kwargs):
+    """Constructs network.
+
+    Args:
+      hidden_dims: list with the number of hidden units per layer. It does not
+        include the output layer; those number of units will always be set to
+        the input dimension multiplied by 2.
+      input_order: Order of degrees to the input units: 'random',
+        'left-to-right', 'right-to-left', or an array of an explicit order.
+        For example, 'left-to-right' builds an autoregressive model
+        p(x) = p(x1) p(x2 | x1) ... p(xD | x<D).
+      hidden_order: Order of degrees to the hidden units: 'random',
+        'left-to-right'.
+      activation: Activation function.
+      use_bias: Whether to use a bias.
+      **kwargs: Keyword arguments of parent class.
+    """
+    super(MADE, self).__init__(**kwargs)
+    self.hidden_dims = hidden_dims
+    self.input_order = input_order
+    self.hidden_order = hidden_order
+    self.activation = tf.keras.activations.get(activation)
+    self.use_bias = use_bias
+    self.network = tf.keras.Sequential([])
+
+  def build(self, input_shape):
+    input_shape = tf.TensorShape(input_shape)
+    last_dim = input_shape[-1]
+    if isinstance(last_dim, tf.Dimension):
+      last_dim = last_dim.value
+    if last_dim is None:
+      raise ValueError('The last dimension of the inputs to '
+                       '`MADE` should be defined. Found `None`.')
+    masks = create_masks(input_dim=last_dim,
+                         hidden_dims=self.hidden_dims,
+                         input_order=self.input_order,
+                         hidden_order=self.hidden_order)
+    for l in range(len(self.hidden_dims)):
+      layer = tf.keras.layers.Dense(
+          self.hidden_dims[l],
+          kernel_initializer=make_masked_initializer(masks[l]),
+          kernel_constraint=make_masked_constraint(masks[l]),
+          activation=self.activation,
+          use_bias=self.use_bias)
+      self.network.add(layer)
+
+    mask = tf.tile(masks[-1], [1, 2])  # for two-headed autoregressive output
+    layer = tf.keras.layers.Dense(
+        last_dim * 2,
+        kernel_initializer=make_masked_initializer(mask),
+        kernel_constraint=make_masked_constraint(mask),
+        activation=None,
+        use_bias=self.use_bias)
+    self.network.add(layer)
+    self.built = True
+
+  def call(self, inputs):
+    return self.network(inputs)
+
+
+def create_degrees(input_dim,
+                   hidden_dims,
+                   input_order='left-to-right',
+                   hidden_order='left-to-right'):
+  """Returns a list of degree vectors, one for each input and hidden layer.
+
+  A unit with degree d can only receive input from units with degree < d. Output
+  units always have the same degree as their associated input unit.
+
+  Args:
+    input_dim: Number of inputs.
+    hidden_dims: list with the number of hidden units per layer. It does not
+      include the output layer; those number of units will always be set to
+      input_dim downstream.
+    input_order: Order of degrees to the input units: 'random', 'left-to-right',
+      'right-to-left', or an array of an explicit order. For example,
+      'left-to-right' builds an autoregressive model
+      p(x) = p(x1) p(x2 | x1) ... p(xD | x<D).
+    hidden_order: Order of degrees to the hidden units: 'random',
+      'left-to-right'.
+  """
+  if (isinstance(input_order, str) and
+      input_order not in ('random', 'left-to-right', 'right-to-left')):
+    raise ValueError('Input order is not valid.')
+  if hidden_order not in ('random', 'left-to-right'):
+    raise ValueError('Hidden order is not valid.')
+
+  degrees = []
+  if isinstance(input_order, str):
+    input_degrees = np.arange(1, input_dim + 1)
+    if input_order == 'right-to-left':
+      input_degrees = np.flip(input_degrees, 0)
+    elif input_order == 'random':
+      np.random.shuffle(input_degrees)
+  else:
+    input_order = np.array(input_order)
+    if np.all(np.sort(input_order) != np.arange(1, input_dim + 1)):
+      raise ValueError('invalid input order')
+    input_degrees = input_order
+  degrees.append(input_degrees)
+
+  for units in hidden_dims:
+    if hidden_order == 'random':
+      min_prev_degree = min(np.min(degrees[-1]), input_dim - 1)
+      hidden_degrees = np.random.randint(
+          low=min_prev_degree, high=input_dim, size=units)
+    elif hidden_order == 'left-to-right':
+      hidden_degrees = (np.arange(units) % max(1, input_dim - 1) +
+                        min(1, input_dim - 1))
+    degrees.append(hidden_degrees)
+  return degrees
+
+
+def create_masks(input_dim,
+                 hidden_dims,
+                 input_order='left-to-right',
+                 hidden_order='left-to-right'):
+  """Returns a list of binary mask matrices respecting autoregressive ordering.
+
+  Args:
+    input_dim: Number of inputs.
+    hidden_dims: list with the number of hidden units per layer. It does not
+      include the output layer; those number of units will always be set to
+      input_dim downstream.
+    input_order: Order of degrees to the input units: 'random', 'left-to-right',
+      'right-to-left', or an array of an explicit order. For example,
+      'left-to-right' builds an autoregressive model
+      p(x) = p(x1) p(x2 | x1) ... p(xD | x<D).
+    hidden_order: Order of degrees to the hidden units: 'random',
+      'left-to-right'.
+  """
+  degrees = create_degrees(input_dim, hidden_dims, input_order, hidden_order)
+  masks = []
+  # Create input-to-hidden and hidden-to-hidden masks.
+  for input_degrees, output_degrees in zip(degrees[:-1], degrees[1:]):
+    mask = tf.cast(input_degrees[:, np.newaxis] <= output_degrees, tf.float32)
+    masks.append(mask)
+
+  # Create hidden-to-output mask.
+  mask = tf.cast(degrees[-1][:, np.newaxis] < degrees[0], tf.float32)
+  masks.append(mask)
+  return masks
+
+
+def make_masked_initializer(mask):
+  initializer = tf.keras.initializers.glorot_uniform()
+  def masked_initializer(shape, dtype=None, partition_info=None):
+    return mask * initializer(shape, dtype, partition_info)
+  return masked_initializer
+
+
+def make_masked_constraint(mask):
+  constraint = tf.identity
+  def masked_constraint(x):
+    return mask * constraint(x)
+  return masked_constraint
diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
new file mode 100644
index 000000000..898c51719
--- /dev/null
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -0,0 +1,90 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 reversible layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensor2tensor.layers import reversible_layers as reversible
+
+import tensorflow as tf
+
+
+class ReversibleLayersTest(tf.test.TestCase):
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testMADELeftToRight(self):
+    np.random.seed(83243)
+    batch_size = 2
+    length = 3
+    network = reversible.MADE([4], activation=tf.nn.relu)
+    inputs = tf.zeros([batch_size, length])
+    outputs = network(inputs)
+
+    num_weights = sum([np.prod(weight.shape) for weight in network.weights])
+    self.assertLen(network.weights, 4)
+    self.assertEqual(num_weights, (3*4 + 4) + (4*3*2 + 3*2))
+
+    self.evaluate(tf.global_variables_initializer())
+    outputs_val = self.evaluate(outputs)
+    self.assertAllEqual(outputs_val[:, 0], tf.zeros(batch_size))
+    self.assertEqual(outputs_val.shape, (batch_size, 2 * length))
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testMADERightToLeft(self):
+    np.random.seed(1328)
+    batch_size = 2
+    length = 3
+    network = reversible.MADE([4, 3],
+                              input_order='right-to-left',
+                              activation=tf.nn.relu,
+                              use_bias=False)
+    inputs = tf.zeros([batch_size, length])
+    outputs = network(inputs)
+
+    num_weights = sum([np.prod(weight.shape) for weight in network.weights])
+    self.assertLen(network.weights, 3)
+    self.assertEqual(num_weights, 3*4 + 4*3 + 3*3*2)
+
+    self.evaluate(tf.global_variables_initializer())
+    outputs_val = self.evaluate(outputs)
+    self.assertAllEqual(outputs_val[:, -1], tf.zeros(batch_size))
+    self.assertEqual(outputs_val.shape, (batch_size, 2 * length))
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testMADENoHidden(self):
+    np.random.seed(532)
+    batch_size = 2
+    length = 3
+    network = reversible.MADE([], input_order='left-to-right')
+    inputs = tf.zeros([batch_size, length])
+    outputs = network(inputs)
+
+    num_weights = sum([np.prod(weight.shape) for weight in network.weights])
+    self.assertLen(network.weights, 2)
+    self.assertEqual(num_weights, 3*3*2 + 3*2)
+
+    self.evaluate(tf.global_variables_initializer())
+    outputs_val = self.evaluate(outputs)
+    self.assertAllEqual(outputs_val[:, 0], tf.zeros(batch_size))
+    self.assertEqual(outputs_val.shape, (batch_size, 2 * length))
+
+
+if __name__ == '__main__':
+  tf.test.main()

From f9ad6fd9f38553fa3cb96e6732184830beaab0ee Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 12 Dec 2018 11:26:30 -0800
Subject: [PATCH 1356/2720] Extend MADE to vary number of output heads.

PiperOrigin-RevId: 225221600
---
 tensor2tensor/layers/reversible_layers.py | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/layers/reversible_layers.py b/tensor2tensor/layers/reversible_layers.py
index 112859903..b315448c3 100644
--- a/tensor2tensor/layers/reversible_layers.py
+++ b/tensor2tensor/layers/reversible_layers.py
@@ -39,14 +39,16 @@ class MADE(tf.keras.Model):
   """Masked autoencoder for distribution estimation (Germain et al., 2015).
 
   MADE takes as input a real Tensor of shape [..., length] and returns a
-  Tensor of shape [..., 2 * length] and same dtype. It masks layer weights to
-  respect autoregressive constraints: for a given ordering, each input dimension
-  can be reconstructed from previous input dimensions. The output dimensions
-  represent two heads for, e.g., location and scale transforms in a flow.
+  Tensor of shape [..., num_heads * length] and same dtype. It masks layer
+  weights to respect autoregressive constraints: for a given ordering, each
+  input dimension can be reconstructed from previous input dimensions. The
+  output dimensions represent multiple heads for, e.g., location and scale
+  transforms in a flow.
   """
 
   def __init__(self,
                hidden_dims,
+               num_heads=2,
                input_order='left-to-right',
                hidden_order='left-to-right',
                activation=None,
@@ -57,7 +59,9 @@ def __init__(self,
     Args:
       hidden_dims: list with the number of hidden units per layer. It does not
         include the output layer; those number of units will always be set to
-        the input dimension multiplied by 2.
+        the input dimension multiplied by `num_heads`.
+      num_heads: The number of output heads. The default is 2 representing
+        the location and scale transform of an autoregressive flow.
       input_order: Order of degrees to the input units: 'random',
         'left-to-right', 'right-to-left', or an array of an explicit order.
         For example, 'left-to-right' builds an autoregressive model
@@ -70,6 +74,7 @@ def __init__(self,
     """
     super(MADE, self).__init__(**kwargs)
     self.hidden_dims = hidden_dims
+    self.num_heads = num_heads
     self.input_order = input_order
     self.hidden_order = hidden_order
     self.activation = tf.keras.activations.get(activation)
@@ -97,9 +102,9 @@ def build(self, input_shape):
           use_bias=self.use_bias)
       self.network.add(layer)
 
-    mask = tf.tile(masks[-1], [1, 2])  # for two-headed autoregressive output
+    mask = tf.tile(masks[-1], [1, self.num_heads])
     layer = tf.keras.layers.Dense(
-        last_dim * 2,
+        last_dim * self.num_heads,
         kernel_initializer=make_masked_initializer(mask),
         kernel_constraint=make_masked_constraint(mask),
         activation=None,

From 981069c65e3a25441eee9b5354ab09d25c839fa2 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 12 Dec 2018 11:46:48 -0800
Subject: [PATCH 1357/2720] Add Actnorm as a reversible layer.

Was playing around with actnorm for better initializations. Not currently using it in experiments but may as well push my code.

PiperOrigin-RevId: 225225365
---
 tensor2tensor/layers/reversible_layers.py     | 68 +++++++++++++++++++
 .../layers/reversible_layers_test.py          | 25 +++++++
 2 files changed, 93 insertions(+)

diff --git a/tensor2tensor/layers/reversible_layers.py b/tensor2tensor/layers/reversible_layers.py
index b315448c3..c165ad561 100644
--- a/tensor2tensor/layers/reversible_layers.py
+++ b/tensor2tensor/layers/reversible_layers.py
@@ -35,6 +35,74 @@ def __init__(self, reversible_layer, **kwargs):
     self.reverse = reversible_layer.call
 
 
+class ActNorm(tf.keras.layers.Layer):
+  """Actnorm, an affine reversible layer (Prafulla and Kingma, 2018).
+
+  Weights use data-dependent initialization in which outputs have zero mean
+  and unit variance per channel (last dimension). The mean/variance statistics
+  are computed from the first batch of inputs.
+  """
+
+  def __init__(self, epsilon=tf.keras.backend.epsilon(), **kwargs):
+    super(ActNorm, self).__init__(**kwargs)
+    self.epsilon = epsilon
+
+  def build(self, input_shape):
+    input_shape = tf.TensorShape(input_shape)
+    last_dim = input_shape[-1]
+    if isinstance(last_dim, tf.Dimension):
+      last_dim = last_dim.value
+    if last_dim is None:
+      raise ValueError('The last dimension of the inputs to `Dense` '
+                       'should be defined. Found `None`.')
+    bias = self.add_weight('bias', [last_dim], dtype=self.dtype)
+    log_scale = self.add_weight('log_scale', [last_dim], dtype=self.dtype)
+    # Set data-dependent initializers.
+    bias = bias.assign(self.bias_initial_value)
+    with tf.control_dependencies([bias]):
+      self.bias = bias
+    log_scale = log_scale.assign(self.log_scale_initial_value)
+    with tf.control_dependencies([log_scale]):
+      self.log_scale = log_scale
+    self.built = True
+
+  def __call__(self, inputs, *args, **kwargs):
+    if not self.built:
+      mean, variance = tf.nn.moments(
+          inputs, axes=[i for i in range(inputs.shape.ndims - 1)])
+      self.bias_initial_value = -mean
+      # TODO(trandustin): Optionally, actnorm multiplies log_scale by a fixed
+      # log_scale factor (e.g., 3.) and initializes by
+      # initial_value / log_scale_factor.
+      self.log_scale_initial_value = tf.log(
+          1. / (tf.sqrt(variance) + self.epsilon))
+
+    if not isinstance(inputs, ed.RandomVariable):
+      return super(ActNorm, self).__call__(inputs, *args, **kwargs)
+
+    bijector = tfp.bijectors.Inline(
+        forward_fn=self.__call__,
+        inverse_fn=self.reverse,
+        inverse_log_det_jacobian_fn=lambda y: -self.log_det_jacobian(y),
+        forward_min_event_ndims=0)
+    return ed.TransformedDistribution(inputs.distribution, bijector=bijector)
+
+  def call(self, inputs):
+    return (inputs + self.bias) * tf.exp(self.log_scale)
+
+  def reverse(self, inputs):
+    return inputs * tf.exp(-self.log_scale) - self.bias
+
+  def log_det_jacobian(self, inputs):
+    """Returns log det | dx / dy | = num_events * sum log | scale |."""
+    del inputs  # unused
+    # Number of events is number of all elements excluding the batch and
+    # channel dimensions.
+    num_events = tf.reduce_prod(tf.shape(inputs)[1:-1])
+    log_det_jacobian = num_events * tf.reduce_sum(self.log_scale)
+    return log_det_jacobian
+
+
 class MADE(tf.keras.Model):
   """Masked autoencoder for distribution estimation (Germain et al., 2015).
 
diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
index 898c51719..7384bfff3 100644
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -28,6 +28,31 @@
 
 class ReversibleLayersTest(tf.test.TestCase):
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testActNorm(self):
+    np.random.seed(83243)
+    batch_size = 25
+    length = 15
+    channels = 4
+    inputs = 3. + 0.8 * np.random.randn(batch_size, length, channels)
+    inputs = tf.cast(inputs, tf.float32)
+    layer = reversible.ActNorm()
+    outputs = layer(inputs)
+    mean, variance = tf.nn.moments(outputs, axes=[0, 1])
+    self.evaluate(tf.global_variables_initializer())
+    mean_val, variance_val = self.evaluate([mean, variance])
+    self.assertAllClose(mean_val, np.zeros(channels), atol=1e-3)
+    self.assertAllClose(variance_val, np.ones(channels), atol=1e-3)
+
+    inputs = 3. + 0.8 * np.random.randn(batch_size, length, channels)
+    inputs = tf.cast(inputs, tf.float32)
+    outputs = layer(inputs)
+    mean, variance = tf.nn.moments(outputs, axes=[0, 1])
+    self.evaluate(tf.global_variables_initializer())
+    mean_val, variance_val = self.evaluate([mean, variance])
+    self.assertAllClose(mean_val, np.zeros(channels), atol=0.25)
+    self.assertAllClose(variance_val, np.ones(channels), atol=0.25)
+
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testMADELeftToRight(self):
     np.random.seed(83243)

From 7a2f3114a60a82a5f97e6a2660d9510689d2f061 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 12 Dec 2018 15:21:06 -0800
Subject: [PATCH 1358/2720] Rename mathematical_language_understanding to
 algorithmic_math_two_variables and add tpu configs for UT.

PiperOrigin-RevId: 225265020
---
 .../mathematical_language_understanding.py    | 10 +++++---
 .../models/research/universal_transformer.py  |  9 ++++++++
 .../research/universal_transformer_util.py    | 23 +++++++++++--------
 tensor2tensor/models/transformer.py           |  1 +
 4 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/data_generators/mathematical_language_understanding.py b/tensor2tensor/data_generators/mathematical_language_understanding.py
index b413b0456..5d0f431ba 100644
--- a/tensor2tensor/data_generators/mathematical_language_understanding.py
+++ b/tensor2tensor/data_generators/mathematical_language_understanding.py
@@ -28,6 +28,11 @@
 generated and randomly split into training and test sets at an approximate
 ratio of 9:1, respectively.
 
+Example lines from training file:
+y=691,x=-999,y*x:-690309
+y=210,x=-995,y+x:-785
+x=-995,y=210,x*x:990025
+
 For more information check the following paper:
 Artit Wangperawong. Attending to Mathematical Language with Transformers,
 arXiv:1812.02825 (https://arxiv.org/abs/1812.02825).
@@ -49,7 +54,7 @@
 
 
 @registry.register_problem
-class MathematicalLanguageUnderstanding(text_problems.Text2TextProblem):
+class AlgorithmicMathTwoVariables(text_problems.Text2TextProblem):
   """Mathematical language understanding, see arxiv.org/abs/1812.02825."""
 
   URL = ("https://art.wangperawong.com/mathematical_language_understanding"
@@ -96,8 +101,7 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
         tmp_dir, compressed_filename, self.URL)
     with tarfile.open(download_path, "r:gz") as tar:
       tar.extractall(tmp_dir)
-    filepath = os.path.join(tmp_dir,
-                            "mathematical_language_understanding_train.txt")
+    filepath = os.path.join(tmp_dir, "symbolic_math_train.txt")
     with open(filepath, "r") as fp:
       for l in fp:
         prob, ans = l.strip().split(":")
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index c523cdc89..bf5b3e13f 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -441,6 +441,15 @@ def universal_transformer_base():
   return hparams
 
 
+@registry.register_hparams
+def universal_transformer_base_tpu():
+  hparams = transformer.transformer_big()
+  hparams = update_hparams_for_universal_transformer(hparams)
+  transformer.update_hparams_for_tpu(hparams)
+  hparams.add_step_timing_signal = False
+  return hparams
+
+
 @registry.register_hparams
 def universal_transformer_big():
   hparams = transformer.transformer_big()
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 3f4aa4db8..5377e2e7b 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -247,7 +247,8 @@ def add_vanilla_transformer_layer(x, num_layers):
                                               attention_unit, pad_remover)
 
       output, _, extra_output = tf.foldl(
-          ut_function, tf.range(hparams.num_rec_steps), initializer=initializer)
+          ut_function, tf.range(hparams.num_rec_steps),
+          initializer=initializer)
 
       # Right now, this is only possible when the transition function is an lstm
       if (hparams.recurrence_type == "lstm" and
@@ -1710,9 +1711,9 @@ def remove_pads(x):
     for i, inputs in enumerate(inputs_list):
       inputs_list[i] = remove_pads(inputs)
 
-  ffn_inputs = (
-      inputs_list[0]
-      if len(inputs_list) == 1 else tf.concat(inputs_list, axis=-1))
+  ffn_inputs = inputs_list[0]
+  if len(inputs_list) != 1:
+    ffn_inputs = tf.concat(inputs_list, axis=-1)
 
   if ffn_layer_type == "dense":
     output = common_layers.dense(
@@ -1866,9 +1867,10 @@ def add_position_timing_signal(x, step, hparams):
 
   elif hparams.position_start_index == "step":
     # Shift positions based on the step
-    num_steps = (
-        hparams.act_max_steps
-        if hparams.recurrence_type == "act" else hparams.num_rec_steps)
+    if hparams.recurrence_type == "act":
+      num_steps = hparams.act_max_steps
+    else:
+      num_steps = hparams.num_rec_steps
     index = tf.cast(
         common_layers.shape_list(x)[1] * step / num_steps, dtype=tf.int32)
 
@@ -1903,9 +1905,10 @@ def add_step_timing_signal(x, step, hparams):
     a Tensor with the same shape as x.
 
   """
-  num_steps = (
-      hparams.act_max_steps
-      if hparams.recurrence_type == "act" else hparams.num_rec_steps)
+  if hparams.recurrence_type == "act":
+    num_steps = hparams.act_max_steps
+  else:
+    num_steps = hparams.num_rec_steps
   channels = common_layers.shape_list(x)[-1]
 
   if hparams.step_timing_signal_type == "learned":
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index d51013099..07a68f638 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -2042,6 +2042,7 @@ def update_hparams_for_tpu(hparams):
   hparams.attention_dropout_broadcast_dims = "0,1"  # batch, heads
   hparams.relu_dropout_broadcast_dims = "1"  # length
   hparams.layer_prepostprocess_dropout_broadcast_dims = "1"  # length
+  return hparams
 
 
 @registry.register_hparams

From d6260d093449da1ce5f75f362ae6b18d2ddef19a Mon Sep 17 00:00:00 2001
From: Dustin Tran <dustinviettran@gmail.com>
Date: Wed, 12 Dec 2018 15:37:29 -0800
Subject: [PATCH 1359/2720] Add GaussianProcess layer. (#1224)

---
 tensor2tensor/layers/bayes.py      | 202 +++++++++++++++++++++++++++++
 tensor2tensor/layers/bayes_test.py |  40 ++++++
 2 files changed, 242 insertions(+)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index fad9b6f21..65d11d714 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import tensorflow as tf
+import tensorflow_probability as tfp
 
 from tensorflow_probability import edward2 as ed
 
@@ -41,6 +42,54 @@ def positive():  # alias, following tf.keras.constraints
   return Positive()
 
 
+class Zeros(object):
+  """Function returning zeros tensor of same shape excluding the last dim."""
+
+  def __call__(self, inputs):
+    return tf.zeros(tf.shape(inputs)[:-1], inputs.dtype)
+
+  def get_config(self):
+    return {}
+
+
+class ExponentiatedQuadratic(object):
+  """Exponentiated quadratic kernel."""
+
+  def __init__(self, variance, lengthscale):
+    self.variance = variance
+    self.lengthscale = lengthscale
+
+  def __call__(self, x1, x2):
+    """Computes exponentiated quadratic over all pairs of inputs.
+
+    Args:
+      x1: Tensor of shape [batch_x1, ...]. Slices along the batch axis denote an
+        individual input to be passed to the kernel. It is computed pairwise
+        with each input sliced from x2.
+      x2: Tensor of shape [batch_x2, ...]. Slices along the batch axis denote an
+        individual input passed to the kernel function. It is computed pairwise
+        with each input sliced from x1.
+
+    Returns:
+      Tensor of shape [batch_x1, batch_x2].
+    """
+    size = tf.convert_to_tensor(x1).shape.ndims
+    if size > 2:
+      raise NotImplementedError('Multiple feature dimensions is not yet '
+                                'supported.')
+    x1 = x1 / self.lengthscale
+    x2 = x2 / self.lengthscale
+    x1_squared = tf.reduce_sum(tf.square(x1), list(range(1, len(x1.shape))))
+    x2_squared = tf.reduce_sum(tf.square(x2), list(range(1, len(x2.shape))))
+    square = (x1_squared[:, tf.newaxis] +
+              x2_squared[tf.newaxis, :] -
+              2 * tf.matmul(x1, x2, transpose_b=True))
+    return self.variance * tf.exp(-square / 2)
+
+  def get_config(self):
+    return {'variance': self.variance, 'lengthscale': self.lengthscale}
+
+
 # TODO(dusenberrymw): Restructure the implementation of a trainable initializer
 # such that callers do not need to have type-conditional logic.
 class TrainableInitializer(tf.keras.initializers.Initializer):
@@ -279,6 +328,159 @@ def build(self, input_shape):
     self.built = True
 
 
+class GaussianProcess(tf.keras.layers.Layer):
+  r"""Gaussian process layer.
+
+  The layer represents a distribution over functions, where a
+  stochastic forward pass appears as
+
+  ```none
+  f ~ GP(f | conditional_inputs, conditional_outputs; mean_fn, covariance_fn)
+  outputs = f(inputs)
+  ```
+
+  The optional arguments `conditional_inputs` and `conditional_outputs`
+  capture data that the GP "memorizes", i.e., it forms a posterior predictive
+  distribution. If left unspecified, the GP posits a prior predictive.
+
+  Given a call to `inputs`, an equivalent formulation in terms of function
+  outputs is
+
+  ```none
+  outputs ~ \prod_{unit=1}^{units} MultivariateNormal(output[:, unit] |
+      mean = mean_fn(inputs) + Knm Kmm^{-1} (conditional_outputs[:, unit]-mean),
+      covariance = Knn - Knm Kmm^{-1} Kmn)
+  ```
+
+  where Knm is the covariance function evaluated between all `inputs` and
+  `conditional_inputs`; Knn is between all `inputs`; Kmm is between all
+  `conditional_inputs`; and mean is the mean function evaluated on
+  `conditional_inputs`. The multivariate normal is correlated across input
+  dimensions and is independent across output dimensions.
+  """
+
+  def __init__(
+      self,
+      units,
+      mean_fn=Zeros(),
+      covariance_fn=ExponentiatedQuadratic(variance=1., lengthscale=1.),
+      conditional_inputs=None,
+      conditional_outputs=None,
+      **kwargs):
+    """Constructs layer.
+
+    Args:
+      units: integer, dimensionality of layer.
+      mean_fn: Mean function, a callable taking an inputs Tensor of shape
+        [batch, ...] and returning a Tensor of shape [batch].
+      covariance_fn: Covariance function, a callable taking two input Tensors
+        of shape [batch_x1, ...] and [batch_x2, ...] respectively, and returning
+        a positive semi-definite matrix of shape [batch_x1, batch_x2].
+      conditional_inputs: Tensor of shape [batch, ...], where batch must be the
+        same as conditional_outputs', and ellipses must match layer inputs.
+      conditional_outputs: Tensor of shape [batch, units], where batch must be
+        the same as conditional_inputs' and units is the layer's units size.
+      **kwargs: kwargs passed to parent class.
+    """
+    super(GaussianProcess, self).__init__(**kwargs)
+    self.units = int(units)
+    self.mean_fn = mean_fn
+    self.covariance_fn = covariance_fn
+    self.conditional_inputs = conditional_inputs
+    self.conditional_outputs = conditional_outputs
+
+    self.supports_masking = True
+    self.input_spec = tf.keras.layers.InputSpec(min_ndim=2)
+
+  def build(self, input_shape=None):
+    # Don't track trainable variables such as in the kernel. The user should
+    # refer to any via, e.g., self.covariance_fn or the user environment.
+    self.built = True
+
+  def call(self, inputs):
+    if self.conditional_inputs is None and self.conditional_outputs is None:
+      covariance_matrix = self.covariance_fn(inputs, inputs)
+      # Tile locations so output has shape [units, batch_size]. Covariance will
+      # broadcast to [units, batch_size, batch_size], and we perform
+      # shape manipulations to get a random variable over [batch_size, units].
+      loc = self.mean_fn(inputs)
+      loc = tf.tile(loc[tf.newaxis], [self.units] + [1] * len(loc.shape))
+    else:
+      knn = self.covariance_fn(inputs, inputs)
+      knm = self.covariance_fn(inputs, self.conditional_inputs)
+      kmm = self.covariance_fn(self.conditional_inputs, self.conditional_inputs)
+      kmm = tf.matrix_set_diag(
+          kmm, tf.matrix_diag_part(kmm) + tf.keras.backend.epsilon())
+      kmm_tril = tf.linalg.cholesky(kmm)
+      kmm_tril_operator = tf.linalg.LinearOperatorLowerTriangular(kmm_tril)
+      knm_operator = tf.linalg.LinearOperatorFullMatrix(knm)
+
+      # TODO(trandustin): Vectorize linear algebra for multiple outputs. For
+      # now, we do each separately and stack to obtain a locations Tensor of
+      # shape [units, batch_size].
+      loc = []
+      for conditional_outputs_unit in tf.unstack(self.conditional_outputs,
+                                                 axis=-1):
+        center = conditional_outputs_unit - self.mean_fn(
+            self.conditional_inputs)
+        loc_unit = knm_operator.matvec(
+            kmm_tril_operator.solvevec(kmm_tril_operator.solvevec(center),
+                                       adjoint=True))
+        loc.append(loc_unit)
+      loc = tf.stack(loc) + self.mean_fn(inputs)[tf.newaxis]
+
+      covariance_matrix = knn
+      covariance_matrix -= knm_operator.matmul(
+          kmm_tril_operator.solve(
+              kmm_tril_operator.solve(knm, adjoint_arg=True), adjoint=True))
+
+    covariance_matrix = tf.matrix_set_diag(
+        covariance_matrix,
+        tf.matrix_diag_part(covariance_matrix) + tf.keras.backend.epsilon())
+
+    # Form a multivariate normal random variable with batch_shape units and
+    # event_shape batch_size. Then make it be independent across the units
+    # dimension. Then transpose its dimensions so it is [batch_size, units].
+    random_variable = ed.MultivariateNormalFullCovariance(
+        loc=loc, covariance_matrix=covariance_matrix)
+    random_variable = ed.Independent(random_variable.distribution,
+                                     reinterpreted_batch_ndims=1)
+    bijector = tfp.bijectors.Inline(
+        forward_fn=lambda x: tf.transpose(x, [1, 0]),
+        inverse_fn=lambda y: tf.transpose(y, [1, 0]),
+        forward_event_shape_fn=lambda input_shape: input_shape[::-1],
+        forward_event_shape_tensor_fn=lambda input_shape: input_shape[::-1],
+        inverse_log_det_jacobian_fn=lambda y: tf.cast(0, y.dtype),
+        forward_min_event_ndims=2)
+    random_variable = ed.TransformedDistribution(random_variable.distribution,
+                                                 bijector=bijector)
+    return random_variable
+
+  def compute_output_shape(self, input_shape):
+    input_shape = tf.TensorShape(input_shape)
+    input_shape = input_shape.with_rank_at_least(2)
+    input_dim = input_shape[-1]
+    if isinstance(input_dim, tf.Dimension):
+      input_dim = input_dim.value
+    if input_dim is None:
+      raise ValueError(
+          'The innermost dimension of input_shape must be defined, but saw: %s'
+          % input_shape)
+    return input_shape[:-1].concatenate(self.units)
+
+  def get_config(self):
+    config = {
+        'units': self.units,
+        'mean_fn': tf.keras.utils.serialize_keras_object(self.mean_fn),
+        'covariance_fn': tf.keras.utils.serialize_keras_object(
+            self.covariance_fn),
+        'conditional_inputs': None,  # don't serialize as it can be large
+        'conditional_outputs': None,  # don't serialize as it can be large
+    }
+    base_config = super(GaussianProcess, self).get_config()
+    return dict(list(base_config.items()) + list(config.items()))
+
+
 class LSTMCellReparameterization(tf.keras.layers.LSTMCell):
   """Bayesian LSTM cell class estimated via reparameterization.
 
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 0f1697a62..e9ec46370 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -121,6 +121,46 @@ def testDenseReparameterizationModel(self):
     self.assertEqual(res.shape, (3, 2))
     self.assertLen(model.losses, 1)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testGaussianProcessPosterior(self):
+    train_batch_size = 3
+    test_batch_size = 2
+    input_dim = 4
+    output_dim = 5
+    features = tf.to_float(np.random.rand(train_batch_size, input_dim))
+    labels = tf.to_float(np.random.rand(train_batch_size, output_dim))
+    layer = bayes.GaussianProcess(output_dim,
+                                  conditional_inputs=features,
+                                  conditional_outputs=labels)
+    test_features = tf.to_float(np.random.rand(test_batch_size, input_dim))
+    test_labels = tf.to_float(np.random.rand(test_batch_size, output_dim))
+    test_outputs = layer(test_features)
+    test_nats = -test_outputs.distribution.log_prob(test_labels)
+    self.evaluate(tf.global_variables_initializer())
+    test_nats_val, outputs_val = self.evaluate([test_nats, test_outputs])
+    self.assertEqual(test_nats_val.shape, ())
+    self.assertGreaterEqual(test_nats_val, 0.)
+    self.assertEqual(outputs_val.shape, (test_batch_size, output_dim))
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testGaussianProcessPrior(self):
+    batch_size = 3
+    input_dim = 4
+    output_dim = 5
+    features = tf.to_float(np.random.rand(batch_size, input_dim))
+    labels = tf.to_float(np.random.rand(batch_size, output_dim))
+    model = tf.keras.Sequential([
+        tf.keras.layers.Dense(2, activation=None),
+        bayes.GaussianProcess(output_dim),
+    ])
+    outputs = model(features)
+    log_prob = outputs.distribution.log_prob(labels)
+    self.evaluate(tf.global_variables_initializer())
+    log_prob_val, outputs_val = self.evaluate([log_prob, outputs])
+    self.assertEqual(log_prob_val.shape, ())
+    self.assertLessEqual(log_prob_val, 0.)
+    self.assertEqual(outputs_val.shape, (batch_size, output_dim))
+
   @parameterized.named_parameters(
       {"testcase_name": "_no_uncertainty", "kernel_initializer": "zeros",
        "recurrent_initializer": "orthogonal", "bias_initializer": "zeros",

From 52b989d634061cdd456cf5ad63f1d68a85f2a59f Mon Sep 17 00:00:00 2001
From: blazejosinski <804945+blazejosinski@users.noreply.github.com>
Date: Thu, 13 Dec 2018 01:02:32 +0100
Subject: [PATCH 1360/2720] Add reward info to debug videos. (#1296)

* Add reward info to debug videos.

* Cargo-culting PIL Image import.
---
 tensor2tensor/rl/trainer_model_based.py | 47 +++++++++++++++++++++----
 1 file changed, 41 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 570f1ff62..ef21158d6 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -51,6 +51,18 @@
 FLAGS = flags.FLAGS
 
 
+# Lazy load PIL.Image
+def PIL_Image():  # pylint: disable=invalid-name
+  from PIL import Image  # pylint: disable=g-import-not-at-top
+  return Image
+
+
+# Lazy load PIL.Image
+def PIL_ImageDraw():  # pylint: disable=invalid-name
+  from PIL import ImageDraw  # pylint: disable=g-import-not-at-top
+  return ImageDraw
+
+
 def real_env_step_increment(hparams):
   """Real env step increment."""
   return int(math.ceil(
@@ -315,14 +327,35 @@ def decode_real_obs(index):
     assert np.all(sim_init_obs == real_init_obs)
 
     debug_frame_batches = []
-    def append_debug_frame_batch(sim_obs, real_obs):
+    def append_debug_frame_batch(sim_obs, real_obs, sim_cum_rews,
+                                 real_cum_rews, sim_rews, real_rews):
+      rews = [[sim_cum_rews, sim_rews], [real_cum_rews, real_rews]]
+      headers = []
+      for j in range(len(sim_obs)):
+        local_nps = []
+        for i in range(2):
+          img = PIL_Image().new('RGB', (sim_obs.shape[-2], 11),)
+          draw = PIL_ImageDraw().Draw(img)
+          draw.text((0, 0), "c:{:3}, r:{:3}".format(int(rews[i][0][j]),
+                                                    int(rews[i][1][j])),
+                    fill=(255, 0, 0))
+          local_nps.append(np.asarray(img))
+        local_nps.append(np.zeros_like(local_nps[0]))
+        headers.append(np.concatenate(local_nps, axis=1))
       errs = np.maximum(
           np.abs(sim_obs.astype(np.int) - real_obs, dtype=np.int) - 10, 0
       ).astype(np.uint8)
+      headers = np.stack(headers)
       debug_frame_batches.append(  # pylint: disable=cell-var-from-loop
-          np.concatenate([sim_obs, real_obs, errs], axis=2)
+          np.concatenate([headers,
+                          np.concatenate([sim_obs, real_obs, errs], axis=2)],
+                          axis=1)
       )
-    append_debug_frame_batch(sim_init_obs, real_init_obs)
+    append_debug_frame_batch(sim_init_obs, real_init_obs,
+                             np.zeros(hparams.wm_eval_batch_size),
+                             np.zeros(hparams.wm_eval_batch_size),
+                             np.zeros(hparams.wm_eval_batch_size),
+                             np.zeros(hparams.wm_eval_batch_size))
 
     (sim_cum_rewards, real_cum_rewards) = (
         np.zeros(hparams.wm_eval_batch_size) for _ in range(2)
@@ -332,9 +365,10 @@ def append_debug_frame_batch(sim_obs, real_obs):
       (sim_obs, sim_rewards, _) = sim_env.step(actions)
       sim_cum_rewards += sim_rewards
 
-      real_cum_rewards += [
+      real_rewards = np.array([
           subsequence[i + 1].reward for subsequence in eval_subsequences
-      ]
+      ])
+      real_cum_rewards += real_rewards
       for (length, reward_accuracies) in six.iteritems(
           reward_accuracies_by_length
       ):
@@ -345,7 +379,8 @@ def append_debug_frame_batch(sim_obs, real_obs):
           )
 
       real_obs = decode_real_obs(i + 1)
-      append_debug_frame_batch(sim_obs, real_obs)
+      append_debug_frame_batch(sim_obs, real_obs, sim_cum_rewards,
+                               real_cum_rewards, sim_rewards, real_rewards)
 
     for debug_frames in np.stack(debug_frame_batches, axis=1):
       for debug_frame in debug_frames:

From 56a763385ea74c6d067ed71e0897cbf3d736ab83 Mon Sep 17 00:00:00 2001
From: Dustin Tran <dustinviettran@gmail.com>
Date: Wed, 12 Dec 2018 15:39:32 -0800
Subject: [PATCH 1361/2720] internal merge of PR #1224

PiperOrigin-RevId: 225268329
---
 tensor2tensor/rl/trainer_model_based.py | 47 ++++---------------------
 1 file changed, 6 insertions(+), 41 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index ef21158d6..570f1ff62 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -51,18 +51,6 @@
 FLAGS = flags.FLAGS
 
 
-# Lazy load PIL.Image
-def PIL_Image():  # pylint: disable=invalid-name
-  from PIL import Image  # pylint: disable=g-import-not-at-top
-  return Image
-
-
-# Lazy load PIL.Image
-def PIL_ImageDraw():  # pylint: disable=invalid-name
-  from PIL import ImageDraw  # pylint: disable=g-import-not-at-top
-  return ImageDraw
-
-
 def real_env_step_increment(hparams):
   """Real env step increment."""
   return int(math.ceil(
@@ -327,35 +315,14 @@ def decode_real_obs(index):
     assert np.all(sim_init_obs == real_init_obs)
 
     debug_frame_batches = []
-    def append_debug_frame_batch(sim_obs, real_obs, sim_cum_rews,
-                                 real_cum_rews, sim_rews, real_rews):
-      rews = [[sim_cum_rews, sim_rews], [real_cum_rews, real_rews]]
-      headers = []
-      for j in range(len(sim_obs)):
-        local_nps = []
-        for i in range(2):
-          img = PIL_Image().new('RGB', (sim_obs.shape[-2], 11),)
-          draw = PIL_ImageDraw().Draw(img)
-          draw.text((0, 0), "c:{:3}, r:{:3}".format(int(rews[i][0][j]),
-                                                    int(rews[i][1][j])),
-                    fill=(255, 0, 0))
-          local_nps.append(np.asarray(img))
-        local_nps.append(np.zeros_like(local_nps[0]))
-        headers.append(np.concatenate(local_nps, axis=1))
+    def append_debug_frame_batch(sim_obs, real_obs):
       errs = np.maximum(
           np.abs(sim_obs.astype(np.int) - real_obs, dtype=np.int) - 10, 0
       ).astype(np.uint8)
-      headers = np.stack(headers)
       debug_frame_batches.append(  # pylint: disable=cell-var-from-loop
-          np.concatenate([headers,
-                          np.concatenate([sim_obs, real_obs, errs], axis=2)],
-                          axis=1)
+          np.concatenate([sim_obs, real_obs, errs], axis=2)
       )
-    append_debug_frame_batch(sim_init_obs, real_init_obs,
-                             np.zeros(hparams.wm_eval_batch_size),
-                             np.zeros(hparams.wm_eval_batch_size),
-                             np.zeros(hparams.wm_eval_batch_size),
-                             np.zeros(hparams.wm_eval_batch_size))
+    append_debug_frame_batch(sim_init_obs, real_init_obs)
 
     (sim_cum_rewards, real_cum_rewards) = (
         np.zeros(hparams.wm_eval_batch_size) for _ in range(2)
@@ -365,10 +332,9 @@ def append_debug_frame_batch(sim_obs, real_obs, sim_cum_rews,
       (sim_obs, sim_rewards, _) = sim_env.step(actions)
       sim_cum_rewards += sim_rewards
 
-      real_rewards = np.array([
+      real_cum_rewards += [
           subsequence[i + 1].reward for subsequence in eval_subsequences
-      ])
-      real_cum_rewards += real_rewards
+      ]
       for (length, reward_accuracies) in six.iteritems(
           reward_accuracies_by_length
       ):
@@ -379,8 +345,7 @@ def append_debug_frame_batch(sim_obs, real_obs, sim_cum_rews,
           )
 
       real_obs = decode_real_obs(i + 1)
-      append_debug_frame_batch(sim_obs, real_obs, sim_cum_rewards,
-                               real_cum_rewards, sim_rewards, real_rewards)
+      append_debug_frame_batch(sim_obs, real_obs)
 
     for debug_frames in np.stack(debug_frame_batches, axis=1):
       for debug_frame in debug_frames:

From 8e76aaa661e7f23e7114dc83d91ae9b80b7b3488 Mon Sep 17 00:00:00 2001
From: blazejosinski <804945+blazejosinski@users.noreply.github.com>
Date: Wed, 12 Dec 2018 16:13:03 -0800
Subject: [PATCH 1362/2720] internal merge of PR #1296

PiperOrigin-RevId: 225274332
---
 tensor2tensor/rl/trainer_model_based.py | 48 +++++++++++++++++++++----
 1 file changed, 42 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 570f1ff62..06fb3e4f8 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -51,6 +51,18 @@
 FLAGS = flags.FLAGS
 
 
+# Lazy load PIL.Image
+def PIL_Image():  # pylint: disable=invalid-name
+  from PIL import Image  # pylint: disable=g-import-not-at-top
+  return Image
+
+
+# Lazy load PIL.Image
+def PIL_ImageDraw():  # pylint: disable=invalid-name
+  from PIL import ImageDraw  # pylint: disable=g-import-not-at-top
+  return ImageDraw
+
+
 def real_env_step_increment(hparams):
   """Real env step increment."""
   return int(math.ceil(
@@ -315,14 +327,36 @@ def decode_real_obs(index):
     assert np.all(sim_init_obs == real_init_obs)
 
     debug_frame_batches = []
-    def append_debug_frame_batch(sim_obs, real_obs):
+    def append_debug_frame_batch(sim_obs, real_obs, sim_cum_rews,
+                                 real_cum_rews, sim_rews, real_rews):
+      """Add a debug frame."""
+      rews = [[sim_cum_rews, sim_rews], [real_cum_rews, real_rews]]
+      headers = []
+      for j in range(len(sim_obs)):
+        local_nps = []
+        for i in range(2):
+          img = PIL_Image().new("RGB", (sim_obs.shape[-2], 11),)
+          draw = PIL_ImageDraw().Draw(img)
+          draw.text((0, 0), "c:{:3}, r:{:3}".format(int(rews[i][0][j]),
+                                                    int(rews[i][1][j])),
+                    fill=(255, 0, 0))
+          local_nps.append(np.asarray(img))
+        local_nps.append(np.zeros_like(local_nps[0]))
+        headers.append(np.concatenate(local_nps, axis=1))
       errs = np.maximum(
           np.abs(sim_obs.astype(np.int) - real_obs, dtype=np.int) - 10, 0
       ).astype(np.uint8)
+      headers = np.stack(headers)
       debug_frame_batches.append(  # pylint: disable=cell-var-from-loop
-          np.concatenate([sim_obs, real_obs, errs], axis=2)
+          np.concatenate([headers,
+                          np.concatenate([sim_obs, real_obs, errs], axis=2)],
+                         axis=1)
       )
-    append_debug_frame_batch(sim_init_obs, real_init_obs)
+    append_debug_frame_batch(sim_init_obs, real_init_obs,
+                             np.zeros(hparams.wm_eval_batch_size),
+                             np.zeros(hparams.wm_eval_batch_size),
+                             np.zeros(hparams.wm_eval_batch_size),
+                             np.zeros(hparams.wm_eval_batch_size))
 
     (sim_cum_rewards, real_cum_rewards) = (
         np.zeros(hparams.wm_eval_batch_size) for _ in range(2)
@@ -332,9 +366,10 @@ def append_debug_frame_batch(sim_obs, real_obs):
       (sim_obs, sim_rewards, _) = sim_env.step(actions)
       sim_cum_rewards += sim_rewards
 
-      real_cum_rewards += [
+      real_rewards = np.array([
           subsequence[i + 1].reward for subsequence in eval_subsequences
-      ]
+      ])
+      real_cum_rewards += real_rewards
       for (length, reward_accuracies) in six.iteritems(
           reward_accuracies_by_length
       ):
@@ -345,7 +380,8 @@ def append_debug_frame_batch(sim_obs, real_obs):
           )
 
       real_obs = decode_real_obs(i + 1)
-      append_debug_frame_batch(sim_obs, real_obs)
+      append_debug_frame_batch(sim_obs, real_obs, sim_cum_rewards,
+                               real_cum_rewards, sim_rewards, real_rewards)
 
     for debug_frames in np.stack(debug_frame_batches, axis=1):
       for debug_frame in debug_frames:

From 96fb56b815b48858c5aec016d7bb3110a4d51f87 Mon Sep 17 00:00:00 2001
From: Youngwook Kim <youngwook.kim@gmail.com>
Date: Thu, 13 Dec 2018 10:02:44 +0900
Subject: [PATCH 1363/2720] MRPC: Exclude dev data from training dataset
 (#1281)

---
 tensor2tensor/data_generators/mrpc.py | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/mrpc.py b/tensor2tensor/data_generators/mrpc.py
index 46ed39168..e8c9e3a39 100644
--- a/tensor2tensor/data_generators/mrpc.py
+++ b/tensor2tensor/data_generators/mrpc.py
@@ -58,6 +58,9 @@ def dataset_splits(self):
     }, {
         "split": problem.DatasetSplit.EVAL,
         "shards": 1,
+    }, {
+        "split": problem.DatasetSplit.TEST,
+        "shards": 1,
     }]
 
   @property
@@ -89,7 +92,7 @@ def download_file(tdir, filepath, url):
 
     return mrpc_dir
 
-  def example_generator(self, filename, dev_ids):
+  def example_generator(self, filename, dev_ids, dataset_split):
     for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
       if idx == 0: continue  # skip header
       if six.PY2:
@@ -97,7 +100,10 @@ def example_generator(self, filename, dev_ids):
       else:
         line = line.strip().decode("utf-8")
       l, id1, id2, s1, s2 = line.split("\t")
-      if dev_ids and [id1, id2] not in dev_ids:
+      is_dev = [id1, id2] in dev_ids
+      if dataset_split == problem.DatasetSplit.TRAIN and is_dev:
+        continue
+      if dataset_split == problem.DatasetSplit.EVAL and not is_dev:
         continue
       inputs = [[s1, s2], [s2, s1]]
       for inp in inputs:
@@ -108,14 +114,17 @@ def example_generator(self, filename, dev_ids):
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
     mrpc_dir = self._maybe_download_corpora(tmp_dir)
-    filesplit = "msr_paraphrase_train.txt"
+    if dataset_split != problem.DatasetSplit.TEST:
+      filesplit = "msr_paraphrase_train.txt"
+    else:
+      filesplit = "msr_paraphrase_test.txt"
     dev_ids = []
-    if dataset_split != problem.DatasetSplit.TRAIN:
+    if dataset_split != problem.DatasetSplit.TEST:
       for row in tf.gfile.Open(os.path.join(mrpc_dir, "dev_ids.tsv")):
         dev_ids.append(row.strip().split("\t"))
 
     filename = os.path.join(mrpc_dir, filesplit)
-    for example in self.example_generator(filename, dev_ids):
+    for example in self.example_generator(filename, dev_ids, dataset_split):
       yield example
 
 

From 3f5ab98e6bf224a2ee9b35f2ccabcc66d69ad608 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 12 Dec 2018 17:03:42 -0800
Subject: [PATCH 1364/2720] Fix open-source imports.

PiperOrigin-RevId: 225282453
---
 tensor2tensor/layers/reversible_layers.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/layers/reversible_layers.py b/tensor2tensor/layers/reversible_layers.py
index c165ad561..4a4949cbf 100644
--- a/tensor2tensor/layers/reversible_layers.py
+++ b/tensor2tensor/layers/reversible_layers.py
@@ -21,6 +21,9 @@
 
 import numpy as np
 import tensorflow as tf
+import tensorflow_probability as tfp
+
+from tensorflow_probability import edward2 as ed
 
 
 class Reverse(tf.keras.layers.Layer):

From a9b89ded53eb8909449fd522313b6c9babc4755e Mon Sep 17 00:00:00 2001
From: Chan Yu <aeloyq@outlook.com>
Date: Thu, 13 Dec 2018 09:17:02 +0800
Subject: [PATCH 1365/2720] add caching mechanism support for fast decoding
 with relative_dot_product in transformer model (#1295)

---
 tensor2tensor/layers/common_attention.py | 61 ++++++++++++++----------
 tensor2tensor/models/transformer.py      |  2 +-
 2 files changed, 37 insertions(+), 26 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 48dd51377..3ff8022e0 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1477,11 +1477,14 @@ def dot_product_attention(q,
     return tf.matmul(weights, v)
 
 
-def _generate_relative_positions_matrix(length, max_relative_position):
+def _generate_relative_positions_matrix(length, max_relative_position, cache=False):
   """Generates matrix of relative positions between inputs."""
-  range_vec = tf.range(length)
-  range_mat = tf.reshape(tf.tile(range_vec, [length]), [length, length])
-  distance_mat = range_mat - tf.transpose(range_mat)
+  if not cache:
+      range_vec = tf.range(length)
+      range_mat = tf.reshape(tf.tile(range_vec, [length]), [length, length])
+      distance_mat = range_mat - tf.transpose(range_mat)
+  else:
+      distance_mat = tf.expand_dims(tf.range(-length+1, 1, 1), 0)
   distance_mat_clipped = tf.clip_by_value(distance_mat, -max_relative_position,
                                           max_relative_position)
   # Shift values to be >= 0. Each integer still uniquely identifies a relative
@@ -1491,11 +1494,15 @@ def _generate_relative_positions_matrix(length, max_relative_position):
 
 
 def _generate_relative_positions_embeddings(length, depth,
-                                            max_relative_position, name):
-  """Generates tensor of size [length, length, depth]."""
+                                            max_relative_position, name,
+                                            cache=False):
+  """
+  Generates tensor of size [length, length, depth] if not cache.
+  Generates tensor of size [1, length, depth] if cache.
+  """
   with tf.variable_scope(name):
     relative_positions_matrix = _generate_relative_positions_matrix(
-        length, max_relative_position)
+        length, max_relative_position, cache=cache)
     vocab_size = max_relative_position * 2 + 1
     # Generates embedding for each relative position of dimension depth.
     embeddings_table = tf.get_variable("embeddings", [vocab_size, depth])
@@ -1509,9 +1516,9 @@ def _relative_attention_inner(x, y, z, transpose):
   This batches matrix multiply calculations to avoid unnecessary broadcasting.
 
   Args:
-    x: Tensor with shape [batch_size, heads, length, length or depth].
-    y: Tensor with shape [batch_size, heads, length, depth].
-    z: Tensor with shape [length, length, depth].
+    x: Tensor with shape [batch_size, heads, length or 1, length or depth].
+    y: Tensor with shape [batch_size, heads, length or 1, depth].
+    z: Tensor with shape [length or 1, length, depth].
     transpose: Whether to transpose inner matrices of y and z. Should be true if
         last dimension of x is depth, not length.
 
@@ -1522,17 +1529,17 @@ def _relative_attention_inner(x, y, z, transpose):
   heads = x.get_shape().as_list()[1]
   length = tf.shape(x)[2]
 
-  # xy_matmul is [batch_size, heads, length, length or depth]
+  # xy_matmul is [batch_size, heads, length or 1, length or depth]
   xy_matmul = tf.matmul(x, y, transpose_b=transpose)
-  # x_t is [length, batch_size, heads, length or depth]
+  # x_t is [length or 1, batch_size, heads, length or depth]
   x_t = tf.transpose(x, [2, 0, 1, 3])
-  # x_t_r is [length, batch_size * heads, length or depth]
+  # x_t_r is [length or 1, batch_size * heads, length or depth]
   x_t_r = tf.reshape(x_t, [length, heads * batch_size, -1])
-  # x_tz_matmul is [length, batch_size * heads, length or depth]
+  # x_tz_matmul is [length or 1, batch_size * heads, length or depth]
   x_tz_matmul = tf.matmul(x_t_r, z, transpose_b=transpose)
-  # x_tz_matmul_r is [length, batch_size, heads, length or depth]
+  # x_tz_matmul_r is [length or 1, batch_size, heads, length or depth]
   x_tz_matmul_r = tf.reshape(x_tz_matmul, [length, batch_size, heads, -1])
-  # x_tz_matmul_r_t is [batch_size, heads, length, length or depth]
+  # x_tz_matmul_r_t is [batch_size, heads, length or 1, length or depth]
   x_tz_matmul_r_t = tf.transpose(x_tz_matmul_r, [1, 2, 0, 3])
   return xy_matmul + x_tz_matmul_r_t
 
@@ -1545,7 +1552,8 @@ def dot_product_attention_relative(q,
                                    dropout_rate=0.0,
                                    image_shapes=None,
                                    name=None,
-                                   make_image_summary=True):
+                                   make_image_summary=True,
+                                   cache=False):
   """Calculate relative position-aware dot-product self-attention.
 
   The attention calculation is augmented with learned representations for the
@@ -1562,6 +1570,7 @@ def dot_product_attention_relative(q,
     image_shapes: optional tuple of integer scalars.
     name: an optional string.
     make_image_summary: Whether to make an attention image summary.
+    cache: whether use cache mode
 
   Returns:
     A Tensor.
@@ -1577,16 +1586,17 @@ def dot_product_attention_relative(q,
 
     # This calculation only works for self attention.
     # q, k and v must therefore have the same shape.
-    q.get_shape().assert_is_compatible_with(k.get_shape())
-    q.get_shape().assert_is_compatible_with(v.get_shape())
+    if not cache:
+        q.get_shape().assert_is_compatible_with(k.get_shape())
+        q.get_shape().assert_is_compatible_with(v.get_shape())
 
     # Use separate embeddings suitable for keys and values.
-    depth = q.get_shape().as_list()[3]
-    length = common_layers.shape_list(q)[2]
+    depth = k.get_shape().as_list()[3]
+    length = common_layers.shape_list(k)[2]
     relations_keys = _generate_relative_positions_embeddings(
-        length, depth, max_relative_position, "relative_positions_keys")
+        length, depth, max_relative_position, "relative_positions_keys", cache=cache)
     relations_values = _generate_relative_positions_embeddings(
-        length, depth, max_relative_position, "relative_positions_values")
+        length, depth, max_relative_position, "relative_positions_values", cache=cache)
 
     # Compute self attention considering the relative position embeddings.
     logits = _relative_attention_inner(q, k, relations_keys, True)
@@ -3389,7 +3399,7 @@ def multihead_attention(query_antecedent,
                             kv_filter_width, q_padding, kv_padding,
                             vars_3d_num_heads=vars_3d_num_heads)
     if cache is not None:
-      if attention_type != "dot_product":
+      if attention_type not in ["dot_product", "dot_product_relative"]:
         # TODO(petershaw): Support caching when using relative position
         # representations, i.e. "dot_product_relative" attention.
         raise NotImplementedError(
@@ -3456,7 +3466,8 @@ def multihead_attention(query_antecedent,
           max_relative_position,
           dropout_rate,
           image_shapes,
-          make_image_summary=make_image_summary)
+          make_image_summary=make_image_summary,
+          cache=cache is not None)
     elif attention_type == "dot_product_unmasked_relative_v2":
       x = dot_product_unmasked_self_attention_relative_v2(
           q,
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 07a68f638..3c21bd986 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -273,7 +273,7 @@ def _beam_decode(self,
               None if using greedy decoding (beam_size=1)
       }
     """
-    if self._hparams.self_attention_type != "dot_product":
+    if self._hparams.self_attention_type not in ["dot_product", "dot_product_relative1"]:
       # Caching is not guaranteed to work with attention types other than
       # dot_product.
       # TODO(petershaw): Support fast decoding when using relative

From 743578dda562f81836512867e802fae374af8189 Mon Sep 17 00:00:00 2001
From: Chan Yu <aeloyq@outlook.com>
Date: Wed, 12 Dec 2018 17:30:58 -0800
Subject: [PATCH 1366/2720] internal merge of PR #1295

PiperOrigin-RevId: 225286154
---
 tensor2tensor/layers/common_attention.py | 26 ++++++++++++------------
 tensor2tensor/models/transformer.py      |  3 ++-
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 3ff8022e0..2e5af4c62 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1477,14 +1477,15 @@ def dot_product_attention(q,
     return tf.matmul(weights, v)
 
 
-def _generate_relative_positions_matrix(length, max_relative_position, cache=False):
+def _generate_relative_positions_matrix(length, max_relative_position,
+                                        cache=False):
   """Generates matrix of relative positions between inputs."""
   if not cache:
-      range_vec = tf.range(length)
-      range_mat = tf.reshape(tf.tile(range_vec, [length]), [length, length])
-      distance_mat = range_mat - tf.transpose(range_mat)
+    range_vec = tf.range(length)
+    range_mat = tf.reshape(tf.tile(range_vec, [length]), [length, length])
+    distance_mat = range_mat - tf.transpose(range_mat)
   else:
-      distance_mat = tf.expand_dims(tf.range(-length+1, 1, 1), 0)
+    distance_mat = tf.expand_dims(tf.range(-length+1, 1, 1), 0)
   distance_mat_clipped = tf.clip_by_value(distance_mat, -max_relative_position,
                                           max_relative_position)
   # Shift values to be >= 0. Each integer still uniquely identifies a relative
@@ -1496,10 +1497,7 @@ def _generate_relative_positions_matrix(length, max_relative_position, cache=Fal
 def _generate_relative_positions_embeddings(length, depth,
                                             max_relative_position, name,
                                             cache=False):
-  """
-  Generates tensor of size [length, length, depth] if not cache.
-  Generates tensor of size [1, length, depth] if cache.
-  """
+  """Generates tensor of size [1 if cache else length, length, depth]."""
   with tf.variable_scope(name):
     relative_positions_matrix = _generate_relative_positions_matrix(
         length, max_relative_position, cache=cache)
@@ -1587,16 +1585,18 @@ def dot_product_attention_relative(q,
     # This calculation only works for self attention.
     # q, k and v must therefore have the same shape.
     if not cache:
-        q.get_shape().assert_is_compatible_with(k.get_shape())
-        q.get_shape().assert_is_compatible_with(v.get_shape())
+      q.get_shape().assert_is_compatible_with(k.get_shape())
+      q.get_shape().assert_is_compatible_with(v.get_shape())
 
     # Use separate embeddings suitable for keys and values.
     depth = k.get_shape().as_list()[3]
     length = common_layers.shape_list(k)[2]
     relations_keys = _generate_relative_positions_embeddings(
-        length, depth, max_relative_position, "relative_positions_keys", cache=cache)
+        length, depth, max_relative_position, "relative_positions_keys",
+        cache=cache)
     relations_values = _generate_relative_positions_embeddings(
-        length, depth, max_relative_position, "relative_positions_values", cache=cache)
+        length, depth, max_relative_position, "relative_positions_values",
+        cache=cache)
 
     # Compute self attention considering the relative position embeddings.
     logits = _relative_attention_inner(q, k, relations_keys, True)
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 3c21bd986..7d9531b86 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -273,7 +273,8 @@ def _beam_decode(self,
               None if using greedy decoding (beam_size=1)
       }
     """
-    if self._hparams.self_attention_type not in ["dot_product", "dot_product_relative1"]:
+    if (self._hparams.self_attention_type not in ["dot_product",
+                                                  "dot_product_relative"]):
       # Caching is not guaranteed to work with attention types other than
       # dot_product.
       # TODO(petershaw): Support fast decoding when using relative

From 969386a471c0bfcdd6a8203f31d7bd01865bc14a Mon Sep 17 00:00:00 2001
From: Roman Kalyakin <theorm@gmail.com>
Date: Thu, 13 Dec 2018 02:50:49 +0100
Subject: [PATCH 1367/2720]  - custom eval hooks (#1284)

---
 tensor2tensor/data_generators/problem.py | 5 +++++
 tensor2tensor/utils/t2t_model.py         | 3 +++
 2 files changed, 8 insertions(+)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 0eac893c3..f462e2312 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -229,6 +229,8 @@ class Problem(object):
   Eval:
     * eval_metrics
         - Specify the set of evaluation metrics for this problem.
+    * eval_hooks
+        - Specify the set of evalueation hooks for this problem.
 
   Inference:
     * feature_encoders(data_dir)
@@ -365,6 +367,9 @@ def eval_metrics(self):
         metrics.Metrics.ACC_PER_SEQ, metrics.Metrics.NEG_LOG_PERPLEXITY
     ]
 
+  def eval_hooks(self, features, logits, hparams):
+    return []
+
   @property
   def task_id(self):
     if self._task_id == -1 and hasattr(self, "global_task_id"):
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 15f19ca5e..88dae5f36 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1451,6 +1451,7 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
       raise NotImplementedError(_no_problem_err("estimator_spec_eval"))
 
     problem = hparams.problem
+
     if common_layers.is_xla_compiled():
       remove_summaries()
       if isinstance(logits, dict):
@@ -1509,6 +1510,8 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
           summary_op=tf.summary.merge_all())
       evaluation_hooks.append(eval_summary_hook)
 
+      evaluation_hooks += problem.eval_hooks(features, logits, hparams)
+
       return tf.estimator.EstimatorSpec(
           tf.estimator.ModeKeys.EVAL,
           predictions=predictions,

From 7384eeb7c5a1a0cd72c09c5374610cb140196da7 Mon Sep 17 00:00:00 2001
From: Roman Kalyakin <theorm@gmail.com>
Date: Wed, 12 Dec 2018 17:51:42 -0800
Subject: [PATCH 1368/2720] internal merge of PR #1284

PiperOrigin-RevId: 225288942
---
 tensor2tensor/data_generators/problem.py |  5 +++--
 tensor2tensor/utils/t2t_model.py         | 20 ++++++++++++--------
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index f462e2312..8f279f1ec 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -368,6 +368,7 @@ def eval_metrics(self):
     ]
 
   def eval_hooks(self, features, logits, hparams):
+    del features, logits, hparams
     return []
 
   @property
@@ -854,9 +855,9 @@ def tpu_valid_size(example):
 
     def gpu_valid_size(example):
       drop_long_sequences = is_training or hparams.eval_drop_long_sequences
+      max_validate_length = max_length if drop_long_sequences else 10**9
       return data_reader.example_valid_size(example, hparams.min_length,
-                                            max_length
-                                            if drop_long_sequences else 10**9)
+                                            max_validate_length)
 
     def define_shapes(example):
       batch_size = config and config.use_tpu and params["batch_size"]
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 88dae5f36..c4ec2b87e 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -209,9 +209,11 @@ def _custom_getter(self):
       if self.hparams.optimizer != "Adafactor":
         raise NotImplementedError(
             "weight_dtype=bfloat16 only implemented with Adafactor optimizer")
+      activation_dtype = tf.float32
+      if self.hparams.activation_dtype == "bfloat16":
+        activation_dtype = tf.bfloat16
       return quantization.EighthPowerEncoding().custom_getter(
-          activation_dtype=tf.bfloat16
-          if self.hparams.activation_dtype == "bfloat16" else tf.float32)
+          activation_dtype=activation_dtype)
     elif self.hparams.activation_dtype == "bfloat16":
       return quantization.bfloat16_activations_var_getter
     else:
@@ -834,8 +836,9 @@ def _greedy_infer(self, features, decode_length, use_tpu=False):
           "losses": a dictionary: {loss-name (string): floating point `Scalar`}
       }
     """
-    return (self._slow_greedy_infer_tpu(features, decode_length)
-            if use_tpu else self._slow_greedy_infer(features, decode_length))
+    if use_tpu:
+      return self._slow_greedy_infer_tpu(features, decode_length)
+    return self._slow_greedy_infer(features, decode_length)
 
   def _slow_greedy_infer_tpu(self, features, decode_length):
     """A slow greedy inference method on TPU.
@@ -1383,8 +1386,9 @@ def estimator_model_fn(cls,
 
     # TRAIN mode
     assert mode == tf.estimator.ModeKeys.TRAIN
-    num_async_replicas = (1 if (use_tpu or not config) else
-                          config.t2t_device_info["num_async_replicas"])
+    num_async_replicas = 1
+    if config and not use_tpu:
+      num_async_replicas = config.t2t_device_info["num_async_replicas"]
     return model.estimator_spec_train(
         loss, num_async_replicas=num_async_replicas, use_tpu=use_tpu)
 
@@ -1522,11 +1526,11 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
   def estimator_spec_predict(self, features, use_tpu=False):
     """Constructs `tf.estimator.EstimatorSpec` for PREDICT (inference) mode."""
     decode_hparams = self._decode_hparams
+    top_beams = decode_hparams.beam_size if decode_hparams.return_beams else 1
     infer_out = self.infer(
         features,
         beam_size=decode_hparams.beam_size,
-        top_beams=(decode_hparams.beam_size
-                   if decode_hparams.return_beams else 1),
+        top_beams=top_beams,
         alpha=decode_hparams.alpha,
         decode_length=decode_hparams.extra_length,
         use_tpu=use_tpu)

From 5560b4b2a23949a51d86174ed9b41293a47d2f2a Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 12 Dec 2018 20:05:03 -0800
Subject: [PATCH 1369/2720] Changes to attention code in mtf transformer.  -
 Added options for different numbers of memory-heads and query-heads.  - Added
 option for shared keys and values  - Rewrote a lot of the attention code and
 moved it to transformer/attention.py    The new code does not rely on
 positional dimensions.

These changes are aimed at improving the speed of incremental decoding.  The main bottleneck in incremental decoding is the size of the key and value tensors which must be retrieved from RAM at each decoding step.  The size of these tensors is:
(batch * heads * memory_length * (key_channels + value_channels)).

The memory size can be reduced by:
  1. local attention (factor of length/attention_radius)
  2. one memory head which is read by all query heads (factor of num_heads)
  3. sharing the keys and values (factor of 2)

Initial tests indicate that these optimizations do not noticeably affect quality for translate_enfr.

TODO(noam): test decoding speed.

PiperOrigin-RevId: 225301844
---
 tensor2tensor/models/mtf_transformer2.py | 142 +++++++++++++++++++++--
 1 file changed, 134 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index 500479bb3..ad5c5a188 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -90,9 +90,11 @@ def _import_feature(self, features, mesh, key):
       return None
     x = tf.to_int32(features[key])
     x = common_layers.expand_squeeze_to_nd(x, 2)
+    batch_size = mtf.Shape(self.batch_dims).size
     # pad to length
     extra_length = self.length_dim.size - tf.shape(x)[1]
-    x = tf.pad(x, [[0, 0], [0, extra_length]])
+    extra_batch = batch_size - tf.shape(x)[0]
+    x = tf.pad(x, [[0, extra_batch], [0, extra_length]])
     mtf_shape = mtf.Shape(self.batch_dims + [self.length_dim])
     x = tf.reshape(x, mtf_shape.to_integer_list)
     return mtf.import_fully_replicated(mesh, x, mtf_shape, name=key)
@@ -260,12 +262,21 @@ def sample(self, features, mesh):
         decode_length_constant=hparams.decode_length_constant)
 
 
+def attention_kwargs_from_hparams(hparams):
+  return {
+      "dropout_rate": hparams.attention_dropout,
+      "extra_logit": 0.0 if hparams.extra_logit else None,
+  }
+
+
 def default_layer_stack(hparams):
   return transformer.LayerStack(
       [transformer_layers.SelfAttention(
           num_heads=hparams.num_heads,
+          num_memory_heads=hparams.num_memory_heads,
           key_value_size=hparams.d_kv,
-          dropout_rate=hparams.attention_dropout),
+          shared_kv=hparams.shared_kv,
+          attention_kwargs=attention_kwargs_from_hparams(hparams)),
        transformer_layers.DenseReluDense(
            hidden_size=hparams.d_ff,
            dropout_rate=hparams.relu_dropout),
@@ -278,12 +289,16 @@ def default_layer_stack_with_encoder_attention(hparams):
   return transformer.LayerStack(
       [transformer_layers.SelfAttention(
           num_heads=hparams.num_heads,
+          num_memory_heads=hparams.num_memory_heads,
           key_value_size=hparams.d_kv,
-          dropout_rate=hparams.attention_dropout),
+          shared_kv=hparams.shared_kv,
+          attention_kwargs=attention_kwargs_from_hparams(hparams)),
        transformer_layers.EncDecAttention(
            num_heads=hparams.num_heads,
+           num_memory_heads=hparams.num_memory_heads,
            key_value_size=hparams.d_kv,
-           dropout_rate=hparams.attention_dropout),
+           shared_kv=hparams.shared_kv,
+           attention_kwargs=attention_kwargs_from_hparams(hparams)),
        transformer_layers.DenseReluDense(
            hidden_size=hparams.d_ff,
            dropout_rate=hparams.relu_dropout),
@@ -311,7 +326,13 @@ def mtf_transformer2_base():
   hparams.add_hparam("d_kv", 128)
   hparams.add_hparam("attention_dropout", 0.0)
   hparams.add_hparam("relu_dropout", 0.0)
+  # share attention keys and values
+  hparams.add_hparam("shared_kv", False)
+  # default of 0 for standard transformer behavior
+  # 1 means a single set of keys and values that are read by all query heads
+  hparams.add_hparam("num_memory_heads", 0)
   hparams.layer_prepostprocess_dropout = 0.0
+  hparams.add_hparam("extra_logit", False)
 
   # round up vocab sizes to be a multiple of this value
   hparams.vocab_divisor = 128
@@ -373,6 +394,7 @@ def mtf_bitransformer_base():
   #        decode_length_multiplier * input_length + decode_length_constant)
   hparams.add_hparam("decode_length_multiplier", 1.5)
   hparams.add_hparam("decode_length_constant", 10.0)
+  hparams.sampling_temp = 0.0
   return hparams
 
 
@@ -495,7 +517,6 @@ def mtr_lm_v1():
   return hparams
 
 
-@registry.register_hparams
 def mtr_tr_dense(sz):
   """Series of machine translation models.
 
@@ -551,8 +572,113 @@ def mtr_tr_dense_3():
   return mtr_tr_dense(3)
 
 
+def mtr_tr_dense_local(sz):
+  """With local self-attention in the decoder."""
+  hparams = mtr_tr_dense(sz)
+  hparams.add_hparam("local_attention_radius", 32)
+  hparams.decoder_layer_stack = transformer.LayerStack(
+      [transformer_layers.LocalSelfAttention(
+          num_heads=hparams.num_heads,
+          num_memory_heads=hparams.num_memory_heads,
+          key_value_size=hparams.d_kv,
+          radius=hparams.local_attention_radius,
+          shared_kv=hparams.shared_kv,
+          attention_kwargs=attention_kwargs_from_hparams(hparams)),
+       transformer_layers.EncDecAttention(
+           num_heads=hparams.num_heads,
+           num_memory_heads=hparams.num_memory_heads,
+           key_value_size=hparams.d_kv,
+           shared_kv=hparams.shared_kv,
+           attention_kwargs=attention_kwargs_from_hparams(hparams)),
+       transformer_layers.DenseReluDense(
+           hidden_size=hparams.d_ff,
+           dropout_rate=hparams.relu_dropout),
+      ] * hparams.num_hidden_layers,
+      dropout_rate=hparams.layer_prepostprocess_dropout,
+      norm_epsilon=hparams.norm_epsilon)
+  return hparams
+
+
+@registry.register_hparams
+def mtr_tr_dense_local_0():
+  return mtr_tr_dense_local(0)
+
+
+@registry.register_hparams
+def mtr_tr_dense_local_0_w8():
+  hparams = mtr_tr_dense_local_0()
+  hparams.local_attention_radius = 8
+  return hparams
+
+
+@registry.register_hparams
+def mtr_tr_dense_local_0_h1_16():
+  hparams = mtr_tr_dense_local_0()
+  hparams.num_heads = 16
+  hparams.num_memory_heads = 1
+  return hparams
+
+
+@registry.register_hparams
+def mtr_tr_dense_local_0_h1_16_shared_kv():
+  hparams = mtr_tr_dense_local_0_h1_16()
+  hparams.shared_kv = True
+  return hparams
+
+
+@registry.register_hparams
+def mtr_tr_dense_0_h4():
+  hparams = mtr_tr_dense_0()
+  hparams.num_heads = 4
+  return hparams
+
+
+@registry.register_hparams
+def mtr_tr_dense_0_h16():
+  hparams = mtr_tr_dense_0()
+  hparams.num_heads = 16
+  return hparams
+
+
+@registry.register_hparams
+def mtr_tr_dense_0_extra_logit():
+  hparams = mtr_tr_dense_0()
+  hparams.extra_logit = True
+  return hparams
+
+
+@registry.register_hparams
+def mtr_tr_dense_0_h1_8():
+  hparams = mtr_tr_dense_0()
+  hparams.num_memory_heads = 1
+  return hparams
+
+
+@registry.register_hparams
+def mtr_tr_dense_0_h1_1():
+  hparams = mtr_tr_dense_0()
+  hparams.num_heads = 1
+  return hparams
+
+
+@registry.register_hparams
+def mtr_tr_dense_0_h1_16():
+  hparams = mtr_tr_dense_0()
+  hparams.num_heads = 16
+  hparams.num_memory_heads = 1
+  return hparams
+
+
+@registry.register_hparams
+def mtr_tr_dense_0_h2_16():
+  hparams = mtr_tr_dense_0()
+  hparams.num_heads = 16
+  hparams.num_memory_heads = 2
+  return hparams
+
+
 @registry.register_hparams
-def mtr_tr_dense_0_short():
-  hparams = mtr_tr_dense(0)
-  hparams.num_hidden_layers = 3
+def mtr_tr_dense_0_shared_kv():
+  hparams = mtr_tr_dense_0()
+  hparams.shared_kv = True
   return hparams

From 8eecc940366ce7d51761a76ef8a10d7f551399f1 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Thu, 13 Dec 2018 09:28:21 -0800
Subject: [PATCH 1370/2720] Another attempt to disable glow open-source tests.

PiperOrigin-RevId: 225383540
---
 oss_scripts/oss_tests.sh | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index d0f4500a9..bf221ca74 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -47,7 +47,10 @@ pytest \
   --ignore=tensor2tensor/data_generators/allen_brain_test.py \
   --ignore=tensor2tensor/rl \
   --ignore=tensor2tensor/models/research \
-  --ignore=tensor2tensor/models/video/nfg_*.py \
+  --ignore=tensor2tensor/models/video/nfg_conv_test.py \
+  --ignore=tensor2tensor/models/video/nfg_conv3d_test.py \
+  --ignore=tensor2tensor/models/video/nfg_conv_lstm_test.py \
+  --ignore=tensor2tensor/models/video/nfg_uncond_test.py \
   --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary \
   --deselect=tensor2tensor/utils/beam_search_test.py::BeamSearchTest::testTPUBeam
 set_status

From 63c9cd3efd12824a73ce40925075cec35e9f0dc3 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 13 Dec 2018 12:55:43 -0800
Subject: [PATCH 1371/2720] Public commit of block parallel transformer model

PiperOrigin-RevId: 225420500
---
 tensor2tensor/models/__init__.py              |   1 +
 .../models/research/transformer_parallel.py   | 326 ++++++++++++++++++
 2 files changed, 327 insertions(+)
 create mode 100644 tensor2tensor/models/research/transformer_parallel.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index a19260891..0ab7b3f9a 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -54,6 +54,7 @@
 from tensor2tensor.models.research import super_lm
 from tensor2tensor.models.research import transformer_moe
 from tensor2tensor.models.research import transformer_nat
+from tensor2tensor.models.research import transformer_parallel
 from tensor2tensor.models.research import transformer_revnet
 from tensor2tensor.models.research import transformer_sketch
 from tensor2tensor.models.research import transformer_symshard
diff --git a/tensor2tensor/models/research/transformer_parallel.py b/tensor2tensor/models/research/transformer_parallel.py
new file mode 100644
index 000000000..671db7eb7
--- /dev/null
+++ b/tensor2tensor/models/research/transformer_parallel.py
@@ -0,0 +1,326 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Models for semi-parallel and parallel decoding with the transformer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.layers import common_layers
+from tensor2tensor.models import transformer
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+@registry.register_model
+class TransformerBlockParallel(transformer.Transformer):
+  """Transformer that predicts blocks of the output in parallel."""
+
+  def body(self, features):
+    assert self._hparams.block_size > 0
+    assert not common_layers.is_xla_compiled()
+    assert "targets_segmentation" not in features
+
+    decoder_output = super(TransformerBlockParallel, self).body(features)
+    assert not isinstance(decoder_output, tuple)
+    assert len(decoder_output.shape) == 4
+
+    relu_dropout_broadcast_dims = (
+        common_layers.comma_separated_string_to_integer_list(
+            getattr(self._hparams, "relu_dropout_broadcast_dims", "")))
+
+    with tf.variable_scope("block_size_%d" % self._hparams.block_size):
+      block_output = common_layers.dense_relu_dense(
+          decoder_output,
+          self._hparams.block_size * self._hparams.filter_size,
+          self._hparams.block_size * self._hparams.hidden_size,
+          dropout=self._hparams.relu_dropout,
+          dropout_broadcast_dims=relu_dropout_broadcast_dims)
+
+    batch_size, length = common_layers.shape_list(decoder_output)[:2]
+    block_output = tf.reshape(block_output, [
+        batch_size,
+        length,
+        self._hparams.block_size,
+        self._hparams.hidden_size
+    ])
+
+    block_output = common_layers.layer_postprocess(
+        decoder_output, block_output, self._hparams)
+
+    return block_output
+
+  def top(self, body_output, features):
+    assert self._hparams.block_size > 0
+
+    if (self._hparams.mode == tf.estimator.ModeKeys.TRAIN or
+        self._hparams.mode == tf.estimator.ModeKeys.EVAL):
+      if self._hparams.mode == tf.estimator.ModeKeys.TRAIN:
+        features["block_index"] = tf.random_uniform(
+            shape=[], minval=0, maxval=self._hparams.block_size, dtype=tf.int64)
+      else:
+        features["block_index"] = 0
+      k = features["block_index"]
+      body_output = body_output[:, :, k:k + 1, :]
+
+    return super(TransformerBlockParallel, self).top(body_output, features)
+
+  def loss(self, logits, features):
+    assert self._hparams.block_size > 0
+
+    def shift_left_4d(x, k):
+      return tf.pad(x, [[0, 0], [0, k], [0, 0], [0, 0]])[:, k:, :, :]
+
+    targets = features["targets"]
+    assert len(targets.shape) == 4
+
+    targets = tf.concat([
+        shift_left_4d(targets, i)
+        for i in range(self._hparams.block_size)
+    ], axis=2)
+
+    if (self._hparams.mode == tf.estimator.ModeKeys.TRAIN or
+        self._hparams.mode == tf.estimator.ModeKeys.EVAL):
+      assert "block_index" in features
+      k = features["block_index"]
+      targets = targets[:, :, k:k + 1, :]
+
+    features["targets"] = targets
+
+    loss = super(TransformerBlockParallel, self).loss(logits, features)
+
+    if self._hparams.mode == tf.estimator.ModeKeys.TRAIN:
+      loss_num, loss_den = loss
+      loss_val = loss_num / loss_den
+      for i in range(self._hparams.block_size):
+        # Hack: if you report a loss of NaN, TensorBoard will plot a point at
+        # the previous value without a connecting line. This is used here to
+        # separate out the training losses by block index.
+        one_or_nan = tf.cond(tf.equal(k, i), lambda: 1.0, lambda: float("nan"))
+        tf.summary.scalar(
+            "block_index_%d" % i, one_or_nan * loss_val, family="losses")
+
+    return loss
+
+  def _greedy_infer(self, features, decode_length, use_tpu=False):
+    assert not use_tpu
+    return self._slow_greedy_infer_guess_and_check(features, decode_length)
+
+  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha):
+    raise NotImplementedError
+
+  def _slow_greedy_infer_guess_and_check(self, features, decode_length):
+    assert self._hparams.block_size > 0
+    assert self._hparams.force_full_predict
+    assert self._hparams.sampling_method == "argmax"
+    assert self._decode_hparams.batch_size == 1
+    assert self._decode_hparams.block_size > 0
+    assert self._decode_hparams.block_size <= self._hparams.block_size
+    assert self._decode_hparams.guess_and_check_top_k > 0
+
+    inputs_old = features["inputs"]
+    assert "targets" not in features
+
+    assert len(features["inputs"].shape) in [3, 4]
+    if len(features["inputs"].shape) < 4:
+      features["inputs"] = tf.expand_dims(features["inputs"], 2)
+
+    block_size = self._decode_hparams.block_size
+    decode_length += tf.shape(features["inputs"])[1]
+
+    def while_exit_cond(result, length):  # pylint: disable=unused-argument
+      return tf.logical_and(
+          length < decode_length,
+          tf.reduce_all(
+              tf.not_equal(result[:, :length, :, :], text_encoder.EOS_ID))
+      )
+
+    def infer_step(result, length):
+      """Inference step."""
+
+      def print_info(result, length, new_length):
+        vocab = self._hparams.problem_hparams.vocabulary["targets"]
+        tf.logging.info(
+            "length=%s new_length=%s length_diff=%s new_suffix=%s",
+            length,
+            new_length,
+            new_length - length,
+            str([
+                vocab._subtoken_id_to_subtoken_string(index)  # pylint: disable=protected-access
+                for index in result[0, -block_size:, 0, 0][:new_length - length]
+            ]).decode("unicode-escape"),
+        )
+
+      features["targets"] = tf.pad(result, [[0, 0], [0, 1], [0, 0], [0, 0]])
+      samples, logits, losses = self.sample(features)  # pylint: disable=unused-variable
+
+      _, top_k_indices = tf.nn.top_k(
+          logits[:, :-1, :1, :, :],
+          k=self._decode_hparams.guess_and_check_top_k)
+      in_top_k = tf.reduce_any(
+          tf.equal(tf.to_int64(top_k_indices), tf.expand_dims(result, 4)),
+          axis=4)
+
+      eos_cumsum = tf.cumsum(
+          tf.to_int32(tf.equal(result, text_encoder.EOS_ID)), axis=1)
+      after_eos = tf.greater(common_layers.shift_right(eos_cumsum), 0)
+
+      correct = tf.logical_and(in_top_k, tf.logical_not(after_eos))
+      correct_cumsum = tf.cumsum(tf.to_int32(correct), axis=1)
+      perfect_cumsum = 1 + tf.range(tf.shape(correct)[1])
+      for axis in [0, 2, 3]:
+        perfect_cumsum = tf.expand_dims(perfect_cumsum, axis=axis)
+
+      new_length = tf.reduce_sum(
+          tf.to_int32(tf.equal(correct_cumsum, perfect_cumsum)), axis=1)
+      new_length = tf.squeeze(new_length, axis=[0, 1, 2])
+      new_length = tf.minimum(new_length, decode_length)
+
+      new_result = tf.concat([
+          result[:, :new_length, :, :],
+          tf.reshape(
+              samples[:, new_length, :block_size, :], [1, block_size, 1, 1])
+      ], axis=1)
+
+      with tf.control_dependencies([
+          tf.py_func(print_info, [result, length, new_length], [])
+      ]):
+        new_result = tf.identity(new_result)
+
+      return new_result, new_length
+
+    result = tf.zeros((1, 0, 1, 1), dtype=tf.int64)
+    length = tf.squeeze(tf.zeros(1, dtype=tf.int32))
+
+    result, length = tf.while_loop(
+        while_exit_cond,
+        infer_step,
+        [result, length],
+        shape_invariants=[
+            tf.TensorShape([1, None, 1, 1]),
+            tf.TensorShape([]),
+        ],
+        back_prop=False,
+        parallel_iterations=1)
+
+    result = result[:, :length, :, :]
+
+    features["inputs"] = inputs_old
+
+    return {
+        "outputs": result,
+        "scores": None,
+    }
+
+
+@registry.register_hparams
+def transformer_base_bs1():
+  hparams = transformer.transformer_base()
+  hparams.add_hparam("block_size", 1)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_bs2():
+  hparams = transformer.transformer_base()
+  hparams.add_hparam("block_size", 2)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_bs3():
+  hparams = transformer.transformer_base()
+  hparams.add_hparam("block_size", 3)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_bs4():
+  hparams = transformer.transformer_base()
+  hparams.add_hparam("block_size", 4)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_bs5():
+  hparams = transformer.transformer_base()
+  hparams.add_hparam("block_size", 5)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_bs6():
+  hparams = transformer.transformer_base()
+  hparams.add_hparam("block_size", 6)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_bs7():
+  hparams = transformer.transformer_base()
+  hparams.add_hparam("block_size", 7)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_bs8():
+  hparams = transformer.transformer_base()
+  hparams.add_hparam("block_size", 8)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_bs9():
+  hparams = transformer.transformer_base()
+  hparams.add_hparam("block_size", 9)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_base_bs10():
+  hparams = transformer.transformer_base()
+  hparams.add_hparam("block_size", 10)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_big_bs1():
+  hparams = transformer.transformer_big()
+  hparams.add_hparam("block_size", 1)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_tiny_bs1():
+  hparams = transformer.transformer_tiny()
+  hparams.add_hparam("block_size", 1)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_tiny_bs2():
+  hparams = transformer.transformer_tiny()
+  hparams.add_hparam("block_size", 2)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_tiny_bs3():
+  hparams = transformer.transformer_tiny()
+  hparams.add_hparam("block_size", 3)
+  return hparams

From 55014af1eb5ea9f8925bdfaddb808717a646af08 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Thu, 13 Dec 2018 14:00:54 -0800
Subject: [PATCH 1372/2720] Bug fixes in local attention implementation. Add
 check when subtracting a dimension from a shape that the dimension is in the
 shape.

PiperOrigin-RevId: 225431359
---
 tensor2tensor/models/mtf_transformer2.py | 186 ++++++++++++++---------
 1 file changed, 111 insertions(+), 75 deletions(-)

diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index ad5c5a188..f819121c6 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -109,7 +109,7 @@ def model(self):
       layer_stack = hparams.layer_stack
     else:
       # hparams.layer_stack is a function for creating a LayerStack
-      layer_stack = hparams.layer_stack(hparams)
+      layer_stack = hparams.layer_stack(hparams, "")
     if self.autoregressive:
       input_vocab_size = self._targets_vocab_size
     else:
@@ -207,11 +207,11 @@ def model(self):
     if isinstance(hparams.encoder_layer_stack, transformer.LayerStack):
       encoder_layer_stack = hparams.encoder_layer_stack
     else:
-      encoder_layer_stack = hparams.encoder_layer_stack(hparams)
+      encoder_layer_stack = hparams.encoder_layer_stack(hparams, "encoder_")
     if isinstance(hparams.decoder_layer_stack, transformer.LayerStack):
       decoder_layer_stack = hparams.decoder_layer_stack
     else:
-      decoder_layer_stack = hparams.decoder_layer_stack(hparams)
+      decoder_layer_stack = hparams.decoder_layer_stack(hparams, "decoder_")
     return transformer.Bitransformer(
         encoder_layer_stack=encoder_layer_stack,
         decoder_layer_stack=decoder_layer_stack,
@@ -231,6 +231,8 @@ def import_feature(key):
       return self._import_feature(features, mesh, key)
     targets = import_feature("targets")
     inputs = import_feature("inputs")
+    if not inputs:
+      raise ValueError("inputs feature is missing")
     encoder_sequence_id = import_feature("inputs_segmentation")
     if not encoder_sequence_id:
       encoder_sequence_id = mtf.to_int32(mtf.not_equal(inputs, 0))
@@ -262,6 +264,7 @@ def sample(self, features, mesh):
         decode_length_constant=hparams.decode_length_constant)
 
 
+# The following functions construct layers based on hyperparmeters
 def attention_kwargs_from_hparams(hparams):
   return {
       "dropout_rate": hparams.attention_dropout,
@@ -269,40 +272,59 @@ def attention_kwargs_from_hparams(hparams):
   }
 
 
-def default_layer_stack(hparams):
+def self_attention_from_hparams(hparams, prefix):
+  """Create self-attention layer based on hyperparameters."""
+  radius = hparams.get(prefix + "local_attention_radius")
+  if radius:
+    return transformer_layers.LocalSelfAttention(
+        num_heads=hparams.get(prefix + "num_heads"),
+        num_memory_heads=hparams.get(prefix + "num_memory_heads", 0),
+        radius=radius,
+        key_value_size=hparams.d_kv,
+        shared_kv=hparams.get(prefix + "shared_kv", False),
+        attention_kwargs=attention_kwargs_from_hparams(hparams))
+  else:
+    return transformer_layers.SelfAttention(
+        num_heads=hparams.get(prefix + "num_heads"),
+        num_memory_heads=hparams.get(prefix + "num_memory_heads", 0),
+        key_value_size=hparams.d_kv,
+        shared_kv=hparams.get(prefix + "shared_kv", False),
+        attention_kwargs=attention_kwargs_from_hparams(hparams))
+
+
+def enc_dec_attention_from_hparams(hparams, prefix):
+  return transformer_layers.EncDecAttention(
+      num_heads=hparams.get(prefix + "num_heads"),
+      num_memory_heads=hparams.get(prefix + "num_memory_heads", 0),
+      key_value_size=hparams.d_kv,
+      shared_kv=hparams.get(prefix + "shared_kv", False),
+      attention_kwargs=attention_kwargs_from_hparams(hparams))
+
+
+def dense_relu_dense_from_hparams(hparams):
+  return transformer_layers.DenseReluDense(
+      hidden_size=hparams.d_ff,
+      dropout_rate=hparams.relu_dropout)
+
+
+def layer_stack_from_hparams(hparams, prefix):
+  """Create a layer stack based on the hyperparameter values."""
   return transformer.LayerStack(
-      [transformer_layers.SelfAttention(
-          num_heads=hparams.num_heads,
-          num_memory_heads=hparams.num_memory_heads,
-          key_value_size=hparams.d_kv,
-          shared_kv=hparams.shared_kv,
-          attention_kwargs=attention_kwargs_from_hparams(hparams)),
-       transformer_layers.DenseReluDense(
-           hidden_size=hparams.d_ff,
-           dropout_rate=hparams.relu_dropout),
-      ] * hparams.num_hidden_layers,
+      [self_attention_from_hparams(hparams, prefix),
+       dense_relu_dense_from_hparams(hparams)
+      ] * hparams.get(prefix + "num_layers"),
       dropout_rate=hparams.layer_prepostprocess_dropout,
       norm_epsilon=hparams.norm_epsilon)
 
 
-def default_layer_stack_with_encoder_attention(hparams):
+def decoder_layer_stack_from_hparams(hparams, prefix):
+  if prefix != "decoder_":
+    raise ValueError("prefix should be 'decoder'")
   return transformer.LayerStack(
-      [transformer_layers.SelfAttention(
-          num_heads=hparams.num_heads,
-          num_memory_heads=hparams.num_memory_heads,
-          key_value_size=hparams.d_kv,
-          shared_kv=hparams.shared_kv,
-          attention_kwargs=attention_kwargs_from_hparams(hparams)),
-       transformer_layers.EncDecAttention(
-           num_heads=hparams.num_heads,
-           num_memory_heads=hparams.num_memory_heads,
-           key_value_size=hparams.d_kv,
-           shared_kv=hparams.shared_kv,
-           attention_kwargs=attention_kwargs_from_hparams(hparams)),
-       transformer_layers.DenseReluDense(
-           hidden_size=hparams.d_ff,
-           dropout_rate=hparams.relu_dropout),
-      ] * hparams.num_hidden_layers,
+      [self_attention_from_hparams(hparams, prefix),
+       enc_dec_attention_from_hparams(hparams, prefix),
+       dense_relu_dense_from_hparams(hparams)
+      ] * hparams.get(prefix + "num_layers"),
       dropout_rate=hparams.layer_prepostprocess_dropout,
       norm_epsilon=hparams.norm_epsilon)
 
@@ -319,18 +341,15 @@ def mtf_transformer2_base():
   # with bfloat16 activations.
   hparams.add_hparam("z_loss", 1e-4)
 
-  # These hyperparameters are used in default_layer_stack()
+  # These hyperparameters are used in layer_stack_from_hparams()
   # They may not be respected if hparams uses a differet layer stack function.
   hparams.num_hidden_layers = 6
   hparams.add_hparam("d_ff", 2048)
   hparams.add_hparam("d_kv", 128)
   hparams.add_hparam("attention_dropout", 0.0)
   hparams.add_hparam("relu_dropout", 0.0)
-  # share attention keys and values
-  hparams.add_hparam("shared_kv", False)
-  # default of 0 for standard transformer behavior
-  # 1 means a single set of keys and values that are read by all query heads
-  hparams.add_hparam("num_memory_heads", 0)
+  hparams.del_hparam("num_heads")
+  hparams.del_hparam("num_hidden_layers")
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.add_hparam("extra_logit", False)
 
@@ -374,9 +393,21 @@ def mtf_transformer2_base():
 
 @registry.register_hparams
 def mtf_unitransformer_base():
+  """Hyperparameters for single-stack Transformer."""
   hparams = mtf_transformer2_base()
   hparams.add_hparam("autoregressive", True)
-  hparams.layer_stack = default_layer_stack
+  hparams.layer_stack = layer_stack_from_hparams
+  # HYPERPARAMETERS FOR THE SINGLE LAYER STACK
+  hparams.add_hparam("num_layers", 6)
+  # number of heads in multihead attention
+  hparams.add_hparam("num_heads", 8)
+  # default of 0 for standard transformer behavior
+  # 1 means a single set of keys and values that are read by all query heads
+  hparams.add_hparam("num_memory_heads", 0)
+  # share attention keys and values
+  hparams.add_hparam("shared_kv", False)
+  # if nonzero then use local attention
+  hparams.add_hparam("local_attention_radius", 0)
   return hparams
 
 
@@ -386,8 +417,25 @@ def mtf_bitransformer_base():
   hparams = mtf_transformer2_base()
   hparams.max_length = 256
   hparams.shared_embedding = True
-  hparams.encoder_layer_stack = default_layer_stack
-  hparams.decoder_layer_stack = default_layer_stack_with_encoder_attention
+  hparams.encoder_layer_stack = layer_stack_from_hparams
+  hparams.decoder_layer_stack = decoder_layer_stack_from_hparams
+  # HYPERPARAMETERS FOR THE LAYER STACKS
+  hparams.add_hparam("encoder_num_layers", 6)
+  hparams.add_hparam("decoder_num_layers", 6)
+  # number of heads in multihead attention
+  hparams.add_hparam("encoder_num_heads", 8)
+  hparams.add_hparam("decoder_num_heads", 8)
+  # default of 0 for standard transformer behavior
+  # 1 means a single set of keys and values that are read by all query heads
+  hparams.add_hparam("encoder_num_memory_heads", 0)
+  hparams.add_hparam("decoder_num_memory_heads", 0)
+  # share attention keys and values
+  hparams.add_hparam("encoder_shared_kv", False)
+  hparams.add_hparam("decoder_shared_kv", False)
+  # if nonzero then use local attention
+  hparams.add_hparam("encoder_local_attention_radius", 0)
+  hparams.add_hparam("decoder_local_attention_radius", 0)
+
   # Parameters for computing the maximum decode length in beam search.
   # Maximum decode length is:
   #    min(max_length,
@@ -540,7 +588,8 @@ def mtr_tr_dense(sz):
   hparams.num_hidden_layers = 6
   hparams.d_ff = int(4096 * n)
   hparams.d_kv = 128
-  hparams.num_heads = int(8 * n)
+  hparams.encoder_num_heads = int(8 * n)
+  hparams.decoder_num_heads = int(8 * n)
   # one epoch for translate_enfr_wmt32k_packed = 51400 steps
   hparams.learning_rate_decay_steps = 51400
   hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
@@ -572,30 +621,17 @@ def mtr_tr_dense_3():
   return mtr_tr_dense(3)
 
 
+@registry.register_hparams
+def mtr_tr_dense_3_88():
+  hparams = mtr_tr_dense(3)
+  hparams.mesh_shape = "model:8;batch:16"
+  return hparams
+
+
 def mtr_tr_dense_local(sz):
   """With local self-attention in the decoder."""
   hparams = mtr_tr_dense(sz)
-  hparams.add_hparam("local_attention_radius", 32)
-  hparams.decoder_layer_stack = transformer.LayerStack(
-      [transformer_layers.LocalSelfAttention(
-          num_heads=hparams.num_heads,
-          num_memory_heads=hparams.num_memory_heads,
-          key_value_size=hparams.d_kv,
-          radius=hparams.local_attention_radius,
-          shared_kv=hparams.shared_kv,
-          attention_kwargs=attention_kwargs_from_hparams(hparams)),
-       transformer_layers.EncDecAttention(
-           num_heads=hparams.num_heads,
-           num_memory_heads=hparams.num_memory_heads,
-           key_value_size=hparams.d_kv,
-           shared_kv=hparams.shared_kv,
-           attention_kwargs=attention_kwargs_from_hparams(hparams)),
-       transformer_layers.DenseReluDense(
-           hidden_size=hparams.d_ff,
-           dropout_rate=hparams.relu_dropout),
-      ] * hparams.num_hidden_layers,
-      dropout_rate=hparams.layer_prepostprocess_dropout,
-      norm_epsilon=hparams.norm_epsilon)
+  hparams.decoder_local_attention_radius = 32
   return hparams
 
 
@@ -607,36 +643,36 @@ def mtr_tr_dense_local_0():
 @registry.register_hparams
 def mtr_tr_dense_local_0_w8():
   hparams = mtr_tr_dense_local_0()
-  hparams.local_attention_radius = 8
+  hparams.decoder_local_attention_radius = 8
   return hparams
 
 
 @registry.register_hparams
 def mtr_tr_dense_local_0_h1_16():
   hparams = mtr_tr_dense_local_0()
-  hparams.num_heads = 16
-  hparams.num_memory_heads = 1
+  hparams.decoder_num_heads = 16
+  hparams.decoder_num_memory_heads = 1
   return hparams
 
 
 @registry.register_hparams
 def mtr_tr_dense_local_0_h1_16_shared_kv():
   hparams = mtr_tr_dense_local_0_h1_16()
-  hparams.shared_kv = True
+  hparams.decoder_shared_kv = True
   return hparams
 
 
 @registry.register_hparams
 def mtr_tr_dense_0_h4():
   hparams = mtr_tr_dense_0()
-  hparams.num_heads = 4
+  hparams.decoder_num_heads = 4
   return hparams
 
 
 @registry.register_hparams
 def mtr_tr_dense_0_h16():
   hparams = mtr_tr_dense_0()
-  hparams.num_heads = 16
+  hparams.decoder_num_heads = 16
   return hparams
 
 
@@ -650,35 +686,35 @@ def mtr_tr_dense_0_extra_logit():
 @registry.register_hparams
 def mtr_tr_dense_0_h1_8():
   hparams = mtr_tr_dense_0()
-  hparams.num_memory_heads = 1
+  hparams.decoder_num_memory_heads = 1
   return hparams
 
 
 @registry.register_hparams
 def mtr_tr_dense_0_h1_1():
   hparams = mtr_tr_dense_0()
-  hparams.num_heads = 1
+  hparams.decoder_num_heads = 1
   return hparams
 
 
 @registry.register_hparams
 def mtr_tr_dense_0_h1_16():
   hparams = mtr_tr_dense_0()
-  hparams.num_heads = 16
-  hparams.num_memory_heads = 1
+  hparams.decoder_num_heads = 16
+  hparams.decoder_num_memory_heads = 1
   return hparams
 
 
 @registry.register_hparams
 def mtr_tr_dense_0_h2_16():
   hparams = mtr_tr_dense_0()
-  hparams.num_heads = 16
-  hparams.num_memory_heads = 2
+  hparams.decoder_num_heads = 16
+  hparams.decoder_num_memory_heads = 2
   return hparams
 
 
 @registry.register_hparams
 def mtr_tr_dense_0_shared_kv():
   hparams = mtr_tr_dense_0()
-  hparams.shared_kv = True
+  hparams.decoder_shared_kv = True
   return hparams

From 5ad23577bbcbf63318f2747bc22b23bd622f6dff Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 13 Dec 2018 16:43:22 -0800
Subject: [PATCH 1373/2720] Get rid of some deprecation warnings in the code.

PiperOrigin-RevId: 225458849
---
 tensor2tensor/models/research/rl.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index a56823f98..2d00f1658 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -356,7 +356,7 @@ def body(self, features):
     x = tf.reshape(observations, [-1] + obs_shape[2:])
     dropout = getattr(self.hparams, "dropout_ppo", 0.0)
     with tf.variable_scope("feed_forward_cnn_small"):
-      x = tf.to_float(x) / 255.0
+      x = tf.cast(x, tf.float32) / 255.0
       x = tf.contrib.layers.conv2d(x, 32, [5, 5], [2, 2],
                                    activation_fn=tf.nn.relu, padding="SAME")
       x = tf.contrib.layers.conv2d(x, 32, [5, 5], [2, 2],
@@ -365,7 +365,7 @@ def body(self, features):
       flat_x = tf.reshape(
           x, [obs_shape[0], obs_shape[1],
               functools.reduce(operator.mul, x.shape.as_list()[1:], 1)])
-      flat_x = tf.nn.dropout(flat_x, keep_prob=1.0 - dropout)
+      flat_x = tf.nn.dropout(flat_x, rate=dropout)
       x = tf.contrib.layers.fully_connected(flat_x, 128, tf.nn.relu)
 
       logits = tf.contrib.layers.fully_connected(
@@ -388,16 +388,16 @@ def body(self, features):
     x = tf.reshape(observations, [-1] + obs_shape[2:])
     dropout = getattr(self.hparams, "dropout_ppo", 0.0)
     with tf.variable_scope("feed_forward_cnn_small"):
-      x = tf.to_float(x) / 255.0
-      x = tf.nn.dropout(x, keep_prob=1.0 - dropout)
+      x = tf.cast(x, tf.float32) / 255.0
+      x = tf.nn.dropout(x, rate=dropout)
       x = tf.layers.conv2d(
           x, 32, (4, 4), strides=(2, 2), name="conv1",
           activation=common_layers.belu, padding="SAME")
-      x = tf.nn.dropout(x, keep_prob=1.0 - dropout)
+      x = tf.nn.dropout(x, rate=dropout)
       x = tf.layers.conv2d(
           x, 64, (4, 4), strides=(2, 2), name="conv2",
           activation=common_layers.belu, padding="SAME")
-      x = tf.nn.dropout(x, keep_prob=1.0 - dropout)
+      x = tf.nn.dropout(x, rate=dropout)
       x = tf.layers.conv2d(
           x, 128, (4, 4), strides=(2, 2), name="conv3",
           activation=common_layers.belu, padding="SAME")
@@ -405,7 +405,7 @@ def body(self, features):
       flat_x = tf.reshape(
           x, [obs_shape[0], obs_shape[1],
               functools.reduce(operator.mul, x.shape.as_list()[1:], 1)])
-      flat_x = tf.nn.dropout(flat_x, keep_prob=1.0 - dropout)
+      flat_x = tf.nn.dropout(flat_x, rate=dropout)
       x = tf.layers.dense(flat_x, 128, activation=tf.nn.relu, name="dense1")
 
       logits = tf.layers.dense(

From 87cef911b1a753c939affd548f7bda0b3e430e55 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 14 Dec 2018 02:33:21 +0100
Subject: [PATCH 1374/2720] Parameter sharing between world model and policy
 (#1301)

* Make world model and policy compatible (all policies don't work yet)

* Checkpoint sharing between world model and policy

* Fix policies

* Add some hparam sets
---
 tensor2tensor/layers/common_hparams.py        |   4 +
 tensor2tensor/models/__init__.py              |   2 +-
 tensor2tensor/models/research/rl.py           | 171 +++++++++++++-----
 tensor2tensor/models/video/base.py            |  31 +++-
 .../models/video/basic_deterministic.py       |  27 +--
 tensor2tensor/models/video/savp.py            |   5 +-
 tensor2tensor/models/video/sv2p.py            |  12 +-
 tensor2tensor/rl/dopamine_connector.py        |   2 +-
 tensor2tensor/rl/envs/simulated_batch_env.py  |   9 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py    |  30 +--
 tensor2tensor/rl/ppo.py                       |   7 +
 tensor2tensor/rl/ppo_learner.py               |   9 +-
 tensor2tensor/rl/trainer_model_based.py       |  18 +-
 .../rl/trainer_model_based_params.py          |  29 ++-
 tensor2tensor/utils/optimize.py               |   4 +
 tensor2tensor/utils/t2t_model.py              |   4 +-
 16 files changed, 261 insertions(+), 103 deletions(-)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 304d501e6..8987a4ae9 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -70,6 +70,10 @@ def basic_params1():
       optimizer_adafactor_multiply_by_parameter_scale=True,
       # Number of accumulating steps for multi step optimizers.
       optimizer_multistep_accumulate_steps=None,
+      # Whether to zero gradients that were not computed, so that the
+      # appropriate slots are created. Useful for sharing checkpoints between
+      # models with different sets of heads.
+      optimizer_zero_grads=False,
       weight_decay=1e-6,
       weight_noise=0.0,
       # Defines the learning rate as a product of named functions.
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 0ab7b3f9a..aa6653b98 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -29,7 +29,7 @@
 from tensor2tensor.models import mtf_image_transformer
 from tensor2tensor.models import mtf_resnet
 from tensor2tensor.models import mtf_transformer
-from tensor2tensor.models import mtf_transformer2
+#from tensor2tensor.models import mtf_transformer2
 from tensor2tensor.models import neural_gpu
 from tensor2tensor.models import resnet
 from tensor2tensor.models import revnet
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 2d00f1658..df74310b5 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -19,15 +19,21 @@
 import functools
 import operator
 import gym
+import six
 
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import video_utils
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
+from tensor2tensor.layers import modalities
+from tensor2tensor.models.video import basic_deterministic_params
 from tensor2tensor.rl.envs.py_func_batch_env import PyFuncBatchEnv
 from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
 from tensor2tensor.rl.envs.simulated_batch_gym_env import SimulatedBatchGymEnv
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
+from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
 import tensorflow_probability as tfp
@@ -122,6 +128,36 @@ def ppo_original_params():
   return hparams
 
 
+@registry.register_hparams
+def ppo_original_world_model():
+  """Atari parameters with world model as policy."""
+  hparams = ppo_original_params()
+  hparams.policy_network = "next_frame_basic_deterministic"
+  hparams_keys = hparams.values().keys()
+  video_hparams = basic_deterministic_params.next_frame_basic_deterministic()
+  for (name, value) in six.iteritems(video_hparams.values()):
+    if name in hparams_keys:
+      hparams.set_hparam(name, value)
+    else:
+      hparams.add_hparam(name, value)
+  return hparams
+
+
+@registry.register_hparams
+def ppo_tiny_world_model():
+  """Atari parameters with world model as policy."""
+  hparams = ppo_original_params()
+  hparams.policy_network = "next_frame_basic_deterministic"
+  hparams_keys = hparams.values().keys()
+  video_hparams = basic_deterministic_params.next_frame_tiny()
+  for (name, value) in six.iteritems(video_hparams.values()):
+    if name in hparams_keys:
+      hparams.set_hparam(name, value)
+    else:
+      hparams.add_hparam(name, value)
+  return hparams
+
+
 def make_real_env_fn(env):
   """Creates a function returning a given real env, in or out of graph.
 
@@ -163,18 +199,27 @@ def get_policy(observations, hparams, action_space):
   if not isinstance(action_space, gym.spaces.Discrete):
     raise ValueError("Expecting discrete action space.")
 
+  policy_problem = DummyPolicyProblem(action_space)
+  trainer_lib.add_problem_hparams(hparams, policy_problem)
+  hparams.force_full_predict = True
   model = registry.model(hparams.policy_network)(
       hparams, tf.estimator.ModeKeys.TRAIN
   )
   obs_shape = common_layers.shape_list(observations)
   features = {
       "inputs": observations,
-      "target_action": tf.zeros(obs_shape[:2] + [action_space.n]),
-      "target_value": tf.zeros(obs_shape[:2])
+      "input_action": tf.zeros(obs_shape[:2] + [1], dtype=tf.int32),
+      "input_reward": tf.zeros(obs_shape[:2] + [1], dtype=tf.int32),
+      "targets": tf.zeros(obs_shape[:1] + [1] + obs_shape[2:]),
+      "target_action": tf.zeros(obs_shape[:1] + [1, 1], dtype=tf.int32),
+      "target_reward": tf.zeros(obs_shape[:1] + [1, 1], dtype=tf.int32),
+      "target_policy": tf.zeros(obs_shape[:1] + [1] + [action_space.n]),
+      "target_value": tf.zeros(obs_shape[:1] + [1])
   }
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+    t2t_model.create_dummy_vars()
     (targets, _) = model(features)
-  return (targets["target_action"], targets["target_value"])
+  return (targets["target_policy"], targets["target_value"])
 
 
 @registry.register_hparams
@@ -260,11 +305,47 @@ def rlmf_tiny():
   return hparams
 
 
-class DiscretePolicyBase(t2t_model.T2TModel):
-
-  @staticmethod
-  def _get_num_actions(features):
-    return common_layers.shape_list(features["target_action"])[2]
+class PolicyBase(t2t_model.T2TModel):
+
+  def loss(self, *args, **kwargs):
+    return 0.0
+
+
+class DummyPolicyProblem(video_utils.VideoProblem):
+  """Dummy Problem for running the policy."""
+
+  def __init__(self, action_space):
+    super(DummyPolicyProblem, self).__init__()
+    self.action_space = action_space
+
+  @property
+  def num_actions(self):
+    return self.action_space.n
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    p.modality = {
+        "inputs": modalities.VideoModality,
+        "input_action": modalities.SymbolModalityWeightsAll,
+        "input_reward": modalities.SymbolModalityWeightsAll,
+        "targets": modalities.VideoModality,
+        "target_action": modalities.SymbolModalityWeightsAll,
+        "target_reward": modalities.SymbolModalityWeightsAll,
+        "target_policy": modalities.IdentityModality,
+        "target_value": modalities.IdentityModality,
+    }
+    p.vocab_size = {
+        "inputs": 256,
+        "input_action": self.num_actions,
+        "input_reward": 3,
+        "targets": 256,
+        "target_action": self.num_actions,
+        "target_reward": 3,
+        "target_policy": None,
+        "target_value": None,
+    }
+    p.input_space_id = problem.SpaceID.IMAGE
+    p.target_space_id = problem.SpaceID.IMAGE
 
 
 NetworkOutput = collections.namedtuple(
@@ -322,38 +403,38 @@ def clip_logits(logits, config):
 
 
 @registry.register_model
-class FeedForwardCategoricalPolicy(DiscretePolicyBase):
+class FeedForwardCategoricalPolicy(PolicyBase):
   """Feed-forward categorical."""
 
   def body(self, features):
     observations = features["inputs"]
-    flat_observations = tf.reshape(observations, [
-        tf.shape(observations)[0], tf.shape(observations)[1],
-        functools.reduce(operator.mul, observations.shape.as_list()[2:], 1)])
+    flat_observations = tf.layers.flatten(observations)
     with tf.variable_scope("policy"):
       x = flat_observations
       for size in self.hparams.policy_layers:
         x = tf.contrib.layers.fully_connected(x, size, tf.nn.relu)
       logits = tf.contrib.layers.fully_connected(
-          x, self._get_num_actions(features), activation_fn=None
+          x, self.hparams.problem.num_actions, activation_fn=None
       )
+      logits = tf.expand_dims(logits, axis=1)
     with tf.variable_scope("value"):
       x = flat_observations
       for size in self.hparams.value_layers:
         x = tf.contrib.layers.fully_connected(x, size, tf.nn.relu)
-      value = tf.contrib.layers.fully_connected(x, 1, None)[..., 0]
+      value = tf.contrib.layers.fully_connected(x, 1, None)
     logits = clip_logits(logits, self.hparams)
-    return {"target_action": logits, "target_value": value}
+    return {"target_policy": logits, "target_value": value}
 
 
 @registry.register_model
-class FeedForwardCnnSmallCategoricalPolicy(DiscretePolicyBase):
+class FeedForwardCnnSmallCategoricalPolicy(PolicyBase):
   """Small cnn network with categorical output."""
 
   def body(self, features):
     observations = features["inputs"]
-    obs_shape = common_layers.shape_list(observations)
-    x = tf.reshape(observations, [-1] + obs_shape[2:])
+    x = tf.transpose(observations, [0, 2, 3, 1, 4])
+    x_shape = common_layers.shape_list(x)
+    x = tf.reshape(x, x_shape[:-2] + [-1])
     dropout = getattr(self.hparams, "dropout_ppo", 0.0)
     with tf.variable_scope("feed_forward_cnn_small"):
       x = tf.cast(x, tf.float32) / 255.0
@@ -368,24 +449,26 @@ def body(self, features):
       flat_x = tf.nn.dropout(flat_x, rate=dropout)
       x = tf.contrib.layers.fully_connected(flat_x, 128, tf.nn.relu)
 
-      logits = tf.contrib.layers.fully_connected(
-          x, self._get_num_actions(features), activation_fn=None
+      logits = tf.layers.dense(
+          x, self.hparams.problem.num_actions, name="dense2"
       )
       logits = clip_logits(logits, self.hparams)
+      logits = tf.expand_dims(logits, axis=1)
 
       value = tf.contrib.layers.fully_connected(
-          x, 1, activation_fn=None)[..., 0]
-    return {"target_action": logits, "target_value": value}
+          x, 1, activation_fn=None)
+    return {"target_policy": logits, "target_value": value}
 
 
 @registry.register_model
-class FeedForwardCnnSmallCategoricalPolicyNew(DiscretePolicyBase):
+class FeedForwardCnnSmallCategoricalPolicyNew(PolicyBase):
   """Small cnn network with categorical output."""
 
   def body(self, features):
     observations = features["inputs"]
-    obs_shape = common_layers.shape_list(observations)
-    x = tf.reshape(observations, [-1] + obs_shape[2:])
+    x = tf.transpose(observations, [0, 2, 3, 1, 4])
+    x_shape = common_layers.shape_list(x)
+    x = tf.reshape(x, x_shape[:-2] + [-1])
     dropout = getattr(self.hparams, "dropout_ppo", 0.0)
     with tf.variable_scope("feed_forward_cnn_small"):
       x = tf.cast(x, tf.float32) / 255.0
@@ -402,50 +485,46 @@ def body(self, features):
           x, 128, (4, 4), strides=(2, 2), name="conv3",
           activation=common_layers.belu, padding="SAME")
 
-      flat_x = tf.reshape(
-          x, [obs_shape[0], obs_shape[1],
-              functools.reduce(operator.mul, x.shape.as_list()[1:], 1)])
-      flat_x = tf.nn.dropout(flat_x, rate=dropout)
+
+      flat_x = tf.layers.flatten(x)
+      flat_x = tf.nn.dropout(flat_x, keep_prob=1.0 - dropout)
       x = tf.layers.dense(flat_x, 128, activation=tf.nn.relu, name="dense1")
 
       logits = tf.layers.dense(
-          x, self._get_num_actions(features), name="dense2"
+          x, self.hparams.problem.num_actions, name="dense2"
       )
+      logits = tf.expand_dims(logits, axis=1)
       logits = clip_logits(logits, self.hparams)
 
-      value = tf.layers.dense(x, 1, name="value")[..., 0]
-    return {"target_action": logits, "target_value": value}
+      value = tf.layers.dense(x, 1, name="value")
+    return {"target_policy": logits, "target_value": value}
 
 
 @registry.register_model
-class DenseBitwiseCategoricalPolicy(DiscretePolicyBase):
+class DenseBitwiseCategoricalPolicy(PolicyBase):
   """Dense network with bitwise input and categorical output."""
 
   def body(self, features):
     observations = features["inputs"]
-    obs_shape = common_layers.shape_list(observations)
-    x = tf.reshape(observations, [-1] + obs_shape[2:])
+    flat_x = tf.layers.flatten(observations)
     with tf.variable_scope("dense_bitwise"):
-      x = discretization.int_to_bit_embed(x, 8, 32)
-      flat_x = tf.reshape(
-          x, [obs_shape[0], obs_shape[1],
-              functools.reduce(operator.mul, x.shape.as_list()[1:], 1)])
+      flat_x = discretization.int_to_bit_embed(flat_x, 8, 32)
 
       x = tf.contrib.layers.fully_connected(flat_x, 256, tf.nn.relu)
       x = tf.contrib.layers.fully_connected(flat_x, 128, tf.nn.relu)
 
       logits = tf.contrib.layers.fully_connected(
-          x, self._get_num_actions(features), activation_fn=None
+          x, self.hparams.problem.num_actions, activation_fn=None
       )
 
       value = tf.contrib.layers.fully_connected(
           x, 1, activation_fn=None)[..., 0]
 
-    return {"target_action": logits, "target_value": value}
+    return {"target_policy": logits, "target_value": value}
 
 
 @registry.register_model
-class RandomPolicy(DiscretePolicyBase):
+class RandomPolicy(PolicyBase):
   """Random policy with categorical output."""
 
   def body(self, features):
@@ -453,10 +532,10 @@ def body(self, features):
     obs_shape = observations.shape.as_list()
     # Just so Saver doesn't complain because of no variables.
     tf.get_variable("dummy_var", initializer=0.0)
-    num_actions = self._get_num_actions(features)
+    num_actions = self.hparams.problem.num_actions
     logits = tf.constant(
         1. / float(num_actions),
-        shape=(obs_shape[:2] + [num_actions])
+        shape=(obs_shape[:1] + [1, num_actions])
     )
-    value = tf.zeros(obs_shape[:2])
-    return {"target_action": logits, "target_value": value}
+    value = tf.zeros(obs_shape[:1] + [1])
+    return {"target_policy": logits, "target_value": value}
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 758073c7b..24905fe0e 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -112,6 +112,8 @@ def next_frame(self,
               where C is 3 for L1/L2 modality and 3*256 for Softmax.
       pred_reward: the same size as input reward.
               None if the model does not detect rewards.
+      pred_action: predicted action logits
+      pred_value: predicted value
       extra_loss: any extra loss other than predicted frame and reward.
               e.g. KL loss in case of VAE models.
       internal_states: updated internal models states.
@@ -530,7 +532,8 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
       with tf.control_dependencies(flat_lists(internal_states)):
         sync_op = tf.no_op()
 
-    res_frames, sampled_frames, res_rewards = [], [], []
+    res_frames, sampled_frames, res_rewards, res_policies, res_values = \
+        [], [], [], [], []
     for i in input_index_range:
       with tf.control_dependencies([sync_op]):
         frames, actions, rewards, target_index = self.__get_next_inputs(
@@ -542,9 +545,12 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
           func_in = (frames, actions, rewards, target_frame,
                      internal_states, video_features)
           func_out = self.next_frame(*func_in)
-          res_frame, res_reward, res_extra_loss, internal_states = func_out
+          res_frame, res_reward, res_policy, res_value, res_extra_loss, \
+              internal_states = func_out
           res_frames.append(res_frame)
           res_rewards.append(res_reward)
+          res_policies.append(res_policy)
+          res_values.append(res_value)
           extra_loss += res_extra_loss / float(len(input_index_range))
 
           # Syncronizing the internals states
@@ -599,6 +605,8 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
       # Cut the predicted input frames.
       res_frames = res_frames[hparams.video_num_input_frames-1:]
       res_rewards = res_rewards[hparams.video_num_input_frames-1:]
+      res_policies = res_policies[hparams.video_num_input_frames-1:]
+      res_values = res_values[hparams.video_num_input_frames-1:]
       sampled_frames = sampled_frames[hparams.video_num_input_frames-1:]
       target_frames = target_frames[hparams.video_num_input_frames-1:]
 
@@ -607,15 +615,28 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
     output_frames = tf.stack(res_frames, axis=1)
     targets = output_frames
 
-    if self.has_rewards:
-      output_rewards = tf.stack(res_rewards, axis=1)
-      targets = {"targets": output_frames, "target_reward": output_rewards}
+    if any((self.has_rewards, self.has_policies, self.has_values)):
+      targets = {"targets": output_frames}
+      if self.has_rewards:
+        targets["target_reward"] = tf.stack(res_rewards, axis=1)
+      if self.has_policies:
+        targets["target_policy"] = tf.stack(res_policies, axis=1)
+      if self.has_values:
+        targets["target_value"] = tf.stack(res_values, axis=1)
 
     return targets, extra_loss
 
+  def loss(self, *args, **kwargs):
+    if "policy_network" in self.hparams.values():
+      return 0.0
+    else:
+      return super(NextFrameBase, self).loss(*args, **kwargs)
+
   def body(self, features):
     self.has_actions = "input_action" in features
     self.has_rewards = "target_reward" in features
+    self.has_policies = "target_policy" in features
+    self.has_values = "target_value" in features
     hparams = self.hparams
 
     def merge(inputs, targets):
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 2c30b1d3a..fe6ed13da 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -110,6 +110,12 @@ def next_frame(self, frames, actions, rewards, target_frame,
                              strides=(2, 2), padding="SAME")
         x = common_layers.layer_norm(x)
 
+    with tf.variable_scope("policy"):
+      x_flat = tf.layers.flatten(x)
+      policy_pred = tf.layers.dense(x_flat, self.hparams.problem.num_actions)
+      value_pred = tf.layers.dense(x_flat, 1)
+      value_pred = tf.squeeze(value_pred, axis=-1)
+
     # Add embedded action if present.
     if self.has_actions:
       x = common_video.inject_additional_input(
@@ -148,14 +154,13 @@ def next_frame(self, frames, actions, rewards, target_frame,
     else:
       x = tf.layers.dense(x, hparams.problem.num_channels, name="logits")
 
-    # No reward prediction if not needed.
-    if not self.has_rewards:
-      return x, None, extra_loss, internal_states
-
-    # Reward prediction based on middle and final logits.
-    reward_pred = tf.concat([x_mid, x_fin], axis=-1)
-    reward_pred = tf.nn.relu(tf.layers.dense(
-        reward_pred, 128, name="reward_pred"))
-    reward_pred = tf.squeeze(reward_pred, axis=1)  # Remove extra dims
-    reward_pred = tf.squeeze(reward_pred, axis=1)  # Remove extra dims
-    return x, reward_pred, extra_loss, internal_states
+    reward_pred = None
+    if self.has_rewards:
+      # Reward prediction based on middle and final logits.
+      reward_pred = tf.concat([x_mid, x_fin], axis=-1)
+      reward_pred = tf.nn.relu(tf.layers.dense(
+          reward_pred, 128, name="reward_pred"))
+      reward_pred = tf.squeeze(reward_pred, axis=1)  # Remove extra dims
+      reward_pred = tf.squeeze(reward_pred, axis=1)  # Remove extra dims
+
+    return x, reward_pred, policy_pred, value_pred, extra_loss, internal_states
diff --git a/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index 01598c5fd..f73a36ef2 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -498,6 +498,9 @@ def next_frame(self, frames, actions, rewards, target_frame,
     if not self.hparams.use_vae or self.hparams.use_gan:
       raise NotImplementedError("Only supporting VAE for now.")
 
+    if self.has_pred_actions or self.has_values:
+      raise NotImplementedError("Parameter sharing with policy not supported.")
+
     image, action, reward = frames[0], actions[0], rewards[0]
     latent_dims = self.hparams.z_dim
     batch_size = common_layers.shape_list(image)[0]
@@ -555,4 +558,4 @@ def next_frame(self, frames, actions, rewards, target_frame,
 
     pred_reward = self.reward_prediction(
         pred_image, action, reward, latent)
-    return pred_image, pred_reward, 0.0, internal_states
+    return pred_image, pred_reward, None, None, 0.0, internal_states
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 7133dee3d..4369f9e11 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -367,6 +367,10 @@ def video_features(
   def next_frame(self, frames, actions, rewards, target_frame,
                  internal_states, video_features):
     del target_frame
+
+    if self.has_pred_actions or self.has_values:
+      raise NotImplementedError("Parameter sharing with policy not supported.")
+
     latent, latent_mean, latent_std = video_features
     frames, actions, rewards = frames[0], actions[0], rewards[0]
 
@@ -384,7 +388,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
 
     pred_reward = self.reward_prediction(
         pred_image, actions, rewards, latent, mid_outputs)
-    return pred_image, pred_reward, extra_loss, internal_states
+    return pred_image, pred_reward, None, None, extra_loss, internal_states
 
 
 @registry.register_model
@@ -435,6 +439,10 @@ def simple_discrete_latent_tower(self, input_image, target_image):
   def next_frame(self, frames, actions, rewards, target_frame,
                  internal_states, video_features):
     del video_features
+
+    if self.has_pred_actions or self.has_values:
+      raise NotImplementedError("Parameter sharing with policy not supported.")
+
     frames, actions, rewards = frames[0], actions[0], rewards[0]
 
     if internal_states is None:
@@ -451,7 +459,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
 
     pred_reward = self.reward_prediction(
         pred_image, actions, rewards, latent)
-    return pred_image, pred_reward, extra_loss, internal_states
+    return pred_image, pred_reward, None, None, extra_loss, internal_states
 
 
 @registry.register_model
diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 7d765172a..98544d695 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -22,7 +22,7 @@
 import copy
 
 from dopamine.agents.dqn import dqn_agent
-from dopamine.discrete_domains import run_experiment
+#from dopamine.discrete_domains import run_experiment
 from dopamine.replay_memory import circular_replay_buffer
 from dopamine.replay_memory.circular_replay_buffer import OutOfGraphReplayBuffer
 from dopamine.replay_memory.circular_replay_buffer import ReplayElement
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 3d5acbcf9..ba05f8129 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -152,10 +152,11 @@ def simulate(self, action):
         # We only need 1 target frame here, set it.
         hparams_target_frames = self._model.hparams.video_num_target_frames
         self._model.hparams.video_num_target_frames = 1
-        model_output = self._model.infer(
-            {"inputs": history,
-             "input_action": actions,
-             "reset_internal_states": self._reset_model.read_value()})
+        model_output = self._model.infer({
+            "inputs": history,
+            "input_action": actions,
+            "reset_internal_states": self._reset_model.read_value()
+        })
         self._model.hparams.video_num_target_frames = hparams_target_frames
 
       observ = tf.cast(tf.squeeze(model_output["targets"], axis=1),
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index d4d72623d..e628dcb87 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -104,19 +104,19 @@ def __str__(self):
 
   @property
   def observ_shape(self):
-    return self.old_shape[:-1] + (self.old_shape[-1] * self.history,)
+    return (self.history,) + self.old_shape
 
   def simulate(self, action):
     reward, done = self._batch_env.simulate(action)
     with tf.control_dependencies([reward, done]):
-      new_observ = self._batch_env.observ + 0
+      new_observ = tf.expand_dims(self._batch_env.observ, axis=1)
       old_observ = tf.gather(
           self._observ.read_value(),
-          list(range(self.old_shape[-1], self.old_shape[-1] * self.history)),
-          axis=-1)
+          list(range(1, self.history)),
+          axis=1)
       with tf.control_dependencies([new_observ, old_observ]):
         with tf.control_dependencies([self._observ.assign(
-            tf.concat([old_observ, new_observ], axis=-1))]):
+            tf.concat([old_observ, new_observ], axis=1))]):
           return tf.identity(reward), tf.identity(done)
 
   def _reset_non_empty(self, indices):
@@ -124,23 +124,9 @@ def _reset_non_empty(self, indices):
     new_values = self._batch_env._reset_non_empty(indices)
     # pylint: enable=protected-access
     initial_frames = getattr(self._batch_env, "history_observations", None)
-    if initial_frames is not None:
-      # Using history buffer frames for initialization, if they are available.
-      with tf.control_dependencies([new_values]):
-        # Transpose to [batch, height, width, history, channels] and merge
-        # history and channels into one dimension.
-        initial_frames = tf.transpose(initial_frames, [0, 2, 3, 1, 4])
-        initial_frames = tf.reshape(initial_frames,
-                                    (len(self),) + self.observ_shape)
-    else:
-      inx = tf.concat(
-          [
-              tf.ones(tf.size(tf.shape(new_values)),
-                      dtype=tf.int64)[:-1],
-              [self.history]
-          ],
-          axis=0)
-      initial_frames = tf.tile(new_values, inx)
+    if initial_frames is None:
+      inx = [1, self.history, 1, 1, 1]
+      initial_frames = tf.tile(tf.expand_dims(new_values, axis=1), inx)
     assign_op = tf.scatter_update(self._observ, indices, initial_frames)
     with tf.control_dependencies([assign_op]):
       return tf.gather(self.observ, indices)
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index bed9b6861..6002682ab 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -21,6 +21,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.layers import common_layers
 from tensor2tensor.models.research.rl import get_policy
 from tensor2tensor.utils import learning_rate
 from tensor2tensor.utils import optimize
@@ -33,7 +34,13 @@ def define_ppo_step(data_points, hparams, action_space, lr):
   """Define ppo step."""
   observation, action, discounted_reward, norm_advantage, old_pdf = data_points
 
+  obs_shape = common_layers.shape_list(observation)
+  observation = tf.reshape(
+      observation, [obs_shape[0] * obs_shape[1]] + obs_shape[2:]
+  )
   (logits, new_value) = get_policy(observation, hparams, action_space)
+  logits = tf.reshape(logits, obs_shape[:2] + [action_space.n])
+  new_value = tf.reshape(new_value, obs_shape[:2])
   new_policy_dist = tfp.distributions.Categorical(logits=logits)
 
   new_pdf = new_policy_dist.prob(action)
diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 6cfd700aa..001ec50eb 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -197,7 +197,10 @@ def _run_train(ppo_hparams,
 
   model_saver = tf.train.Saver(
       tf.global_variables(ppo_hparams.policy_network + "/.*") +
-      tf.global_variables("global_step")
+      tf.global_variables("training/" + ppo_hparams.policy_network + "/.*") +
+      tf.global_variables("global_step") +
+      tf.global_variables("losses_avg.*") +
+      tf.global_variables("train_stats.*")
   )
 
   global_step = tf.train.get_or_create_global_step()
@@ -384,12 +387,12 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
         """Step of the environment."""
 
         (logits, value_function) = get_policy(
-            tf.expand_dims(obs_copy, 0), ppo_hparams, batch_env.action_space
+            obs_copy, ppo_hparams, batch_env.action_space
         )
         action = common_layers.sample_with_temperature(logits, sampling_temp)
         action = tf.cast(action, tf.int32)
 
-        reward, done = batch_env.simulate(action[0, ...])
+        reward, done = batch_env.simulate(action[:, 0, ...])
 
         pdf = tfp.distributions.Categorical(logits=logits).prob(action)
         pdf = tf.reshape(pdf, shape=(num_agents,))
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 06fb3e4f8..86839f7d0 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -138,6 +138,9 @@ def choose_subsequence():
 def make_simulated_env_fn(
     real_env, hparams, batch_size, initial_frame_chooser, model_dir):
   """Creates a simulated env_fn."""
+  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
+  if hparams.wm_policy_param_sharing:
+    model_hparams.optimizer_zero_grads = True
   return rl.make_simulated_env_fn(
       reward_range=real_env.reward_range,
       observation_space=real_env.observation_space,
@@ -203,6 +206,8 @@ def initial_frame_chooser(batch_size):
   )
   base_algo_str = hparams.base_algo
   train_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
+  if hparams.wm_policy_param_sharing:
+    train_hparams.optimizer_zero_grads = True
 
   rl_utils.update_hparams_from_hparams(
       train_hparams, hparams, base_algo_str + "_"
@@ -226,6 +231,8 @@ def train_agent_real_env(env, learner, hparams, epoch):
   rl_utils.update_hparams_from_hparams(
       train_hparams, hparams, "real_" + base_algo_str + "_"
   )
+  if hparams.wm_policy_param_sharing:
+    train_hparams.optimizer_zero_grads = True
 
   env_fn = rl.make_real_env_fn(env)
   num_env_steps = real_env_step_increment(hparams)
@@ -253,6 +260,8 @@ def train_world_model(
   model_hparams.learning_rate = model_hparams.learning_rate_constant
   if epoch > 0:
     model_hparams.learning_rate *= hparams.learning_rate_bump
+  if hparams.wm_policy_param_sharing:
+    model_hparams.optimizer_zero_grads = True
 
   restarter = Restarter("world_model", output_dir, world_model_steps_num)
   if restarter.should_skip:
@@ -448,9 +457,13 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   )
   env.start_new_epoch(epoch, data_dir)
 
+  if hparams.wm_policy_param_sharing:
+    policy_model_dir = directories["world_model"]
+  else:
+    policy_model_dir = directories["policy"]
   learner = rl_utils.LEARNERS[hparams.base_algo](
-      hparams.frame_stack_size, directories["policy"],
-      directories["policy"], hparams.epochs
+      hparams.frame_stack_size, policy_model_dir,
+      policy_model_dir, hparams.epochs
   )
 
   # Timing log function
@@ -461,7 +474,6 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   metrics = {}
 
   # Collect data from the real environment.
-  policy_model_dir = directories["policy"]
   tf.logging.info("Initial training of the policy in real environment.")
   train_agent_real_env(env, learner, hparams, epoch)
   metrics["mean_reward/train/clipped"] = rl_utils.compute_mean_reward(
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 754e729ab..3cbe71ee6 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -136,6 +136,7 @@ def rlmb_ppo_base():
       # Number of simulated environments to train on simultaneously.
       simulated_batch_size=16,
       eval_batch_size=30,
+      wm_policy_param_sharing=False,
 
       # Unused; number of PPO epochs is calculated from the real frame limit.
       real_ppo_epochs_num=0,
@@ -161,6 +162,15 @@ def rlmb_ppo_base():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_ppo_base_param_sharing():
+  """HParams for PPO base with parameter sharing."""
+  hparams = rlmb_ppo_base()
+  hparams.wm_policy_param_sharing = True
+  hparams.base_algo_params = "ppo_original_world_model"
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_base():
   return rlmb_ppo_base()
@@ -208,9 +218,9 @@ def rlmb_noresize():
 
 
 @registry.register_hparams
-def rlmb_quick():
+def rlmb_ppo_quick():
   """Base setting but quicker with only 2 epochs."""
-  hparams = rlmb_base()
+  hparams = rlmb_ppo_base()
   hparams.epochs = 2
   hparams.model_train_steps = 25000
   hparams.ppo_epochs_num = 700
@@ -218,6 +228,21 @@ def rlmb_quick():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_quick():
+  """Base setting but quicker with only 2 epochs."""
+  return rlmb_ppo_quick()
+
+
+@registry.register_hparams
+def rlmb_ppo_quick_param_sharing():
+  """HParams for PPO quick with parameter sharing."""
+  hparams = rlmb_ppo_quick()
+  hparams.wm_policy_param_sharing = True
+  hparams.base_algo_params = "ppo_original_world_model"
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_quick_noresize():
   hparams = rlmb_base()
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index a8b3fac28..b960a2038 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -150,11 +150,15 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
     else:
       self._opt = tf.contrib.layers.OPTIMIZER_CLS_NAMES[optimizer_name](lr)
 
+    self._zero_grads = hparams.optimizer_zero_grads
+
   def compute_gradients(self, loss, var_list=None, **kwargs):  # pylint: disable=arguments-differ
     gradients = self._opt.compute_gradients(loss, var_list, **kwargs)
     def cast_grad(g, v):
       if v is not None and g is not None:
         g = common_layers.cast_like(g, v)
+      if self._zero_grads and g is None:
+        g = tf.zeros_like(v)
       return (g, v)
     gradients = [cast_grad(g, v) for g, v in gradients]
     return gradients
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index c4ec2b87e..deace4110 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1297,7 +1297,7 @@ def estimator_model_fn(cls,
       TPUEstimatorSpec if use tpu else EstimatorSpec
     """
     if mode == tf.estimator.ModeKeys.TRAIN:
-      _create_dummy_vars()
+      create_dummy_vars()
     hparams = copy.deepcopy(hparams)
 
     # Instantiate model
@@ -1622,7 +1622,7 @@ def fn_with_timing(*args, **kwargs):
   return fn_with_timing
 
 
-def _create_dummy_vars():
+def create_dummy_vars():
   """Dummy vars for restore to work when not using TPU codepath."""
   var_names = set([v.name for v in tf.global_variables()])
   if "losses_avg/problem_0/total_loss:0" in var_names:

From 7d753a6ee1f6708cae8ca958af2fa823d25d7b21 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 13 Dec 2018 17:33:47 -0800
Subject: [PATCH 1375/2720] internal merge of PR #1301

PiperOrigin-RevId: 225465399
---
 tensor2tensor/models/__init__.py                  |  2 +-
 tensor2tensor/models/research/rl.py               |  8 +++-----
 tensor2tensor/models/video/basic_deterministic.py | 13 ++++++++-----
 tensor2tensor/models/video/sv2p.py                |  4 ++--
 tensor2tensor/rl/dopamine_connector.py            |  2 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py        |  2 +-
 6 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index aa6653b98..0ab7b3f9a 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -29,7 +29,7 @@
 from tensor2tensor.models import mtf_image_transformer
 from tensor2tensor.models import mtf_resnet
 from tensor2tensor.models import mtf_transformer
-#from tensor2tensor.models import mtf_transformer2
+from tensor2tensor.models import mtf_transformer2
 from tensor2tensor.models import neural_gpu
 from tensor2tensor.models import resnet
 from tensor2tensor.models import revnet
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index df74310b5..bceceb0d2 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -311,6 +311,7 @@ def loss(self, *args, **kwargs):
     return 0.0
 
 
+# TODO(lukaszkaiser): move this class or clean up the whole file.
 class DummyPolicyProblem(video_utils.VideoProblem):
   """Dummy Problem for running the policy."""
 
@@ -443,9 +444,7 @@ def body(self, features):
       x = tf.contrib.layers.conv2d(x, 32, [5, 5], [2, 2],
                                    activation_fn=tf.nn.relu, padding="SAME")
 
-      flat_x = tf.reshape(
-          x, [obs_shape[0], obs_shape[1],
-              functools.reduce(operator.mul, x.shape.as_list()[1:], 1)])
+      flat_x = tf.layers.flatten(x)
       flat_x = tf.nn.dropout(flat_x, rate=dropout)
       x = tf.contrib.layers.fully_connected(flat_x, 128, tf.nn.relu)
 
@@ -485,9 +484,8 @@ def body(self, features):
           x, 128, (4, 4), strides=(2, 2), name="conv3",
           activation=common_layers.belu, padding="SAME")
 
-
       flat_x = tf.layers.flatten(x)
-      flat_x = tf.nn.dropout(flat_x, keep_prob=1.0 - dropout)
+      flat_x = tf.nn.dropout(flat_x, rate=dropout)
       x = tf.layers.dense(flat_x, 128, activation=tf.nn.relu, name="dense1")
 
       logits = tf.layers.dense(
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index fe6ed13da..54fdb1641 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -110,11 +110,14 @@ def next_frame(self, frames, actions, rewards, target_frame,
                              strides=(2, 2), padding="SAME")
         x = common_layers.layer_norm(x)
 
-    with tf.variable_scope("policy"):
-      x_flat = tf.layers.flatten(x)
-      policy_pred = tf.layers.dense(x_flat, self.hparams.problem.num_actions)
-      value_pred = tf.layers.dense(x_flat, 1)
-      value_pred = tf.squeeze(value_pred, axis=-1)
+    if self.has_actions:
+      with tf.variable_scope("policy"):
+        x_flat = tf.layers.flatten(x)
+        policy_pred = tf.layers.dense(x_flat, self.hparams.problem.num_actions)
+        value_pred = tf.layers.dense(x_flat, 1)
+        value_pred = tf.squeeze(value_pred, axis=-1)
+    else:
+      policy_pred, value_pred = None, None
 
     # Add embedded action if present.
     if self.has_actions:
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 4369f9e11..1ed42bb46 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -368,7 +368,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
                  internal_states, video_features):
     del target_frame
 
-    if self.has_pred_actions or self.has_values:
+    if self.has_policies or self.has_values:
       raise NotImplementedError("Parameter sharing with policy not supported.")
 
     latent, latent_mean, latent_std = video_features
@@ -384,7 +384,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
         frames, None, actions, internal_states, latent)
 
     if not self.has_rewards:
-      return pred_image, None, extra_loss, internal_states
+      return pred_image, None, None, None, extra_loss, internal_states
 
     pred_reward = self.reward_prediction(
         pred_image, actions, rewards, latent, mid_outputs)
diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 98544d695..7d765172a 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -22,7 +22,7 @@
 import copy
 
 from dopamine.agents.dqn import dqn_agent
-#from dopamine.discrete_domains import run_experiment
+from dopamine.discrete_domains import run_experiment
 from dopamine.replay_memory import circular_replay_buffer
 from dopamine.replay_memory.circular_replay_buffer import OutOfGraphReplayBuffer
 from dopamine.replay_memory.circular_replay_buffer import ReplayElement
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index e628dcb87..de7929489 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -89,7 +89,7 @@ def history_observations(self):
 
 
 class StackWrapper(WrapperBase):
-  """ A wrapper which stacks previously seen frames. """
+  """A wrapper which stacks previously seen frames."""
 
   def __init__(self, batch_env, history=4):
     super(StackWrapper, self).__init__(batch_env)

From 1b3d0f2dac6af6ca611c2b9d5b182e0e507da127 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 13 Dec 2018 17:54:20 -0800
Subject: [PATCH 1376/2720] Fix open-source test for MADE.

PiperOrigin-RevId: 225467601
---
 tensor2tensor/layers/reversible_layers_test.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
index 7384bfff3..95a77442d 100644
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -63,7 +63,7 @@ def testMADELeftToRight(self):
     outputs = network(inputs)
 
     num_weights = sum([np.prod(weight.shape) for weight in network.weights])
-    self.assertLen(network.weights, 4)
+    self.assertEqual(len(network.weights), 4)
     self.assertEqual(num_weights, (3*4 + 4) + (4*3*2 + 3*2))
 
     self.evaluate(tf.global_variables_initializer())
@@ -84,7 +84,7 @@ def testMADERightToLeft(self):
     outputs = network(inputs)
 
     num_weights = sum([np.prod(weight.shape) for weight in network.weights])
-    self.assertLen(network.weights, 3)
+    self.assertEqual(len(network.weights), 3)
     self.assertEqual(num_weights, 3*4 + 4*3 + 3*3*2)
 
     self.evaluate(tf.global_variables_initializer())
@@ -102,7 +102,7 @@ def testMADENoHidden(self):
     outputs = network(inputs)
 
     num_weights = sum([np.prod(weight.shape) for weight in network.weights])
-    self.assertLen(network.weights, 2)
+    self.assertEqual(len(network.weights), 2)
     self.assertEqual(num_weights, 3*3*2 + 3*2)
 
     self.evaluate(tf.global_variables_initializer())

From aa164d2c992e29fc5c265d4c5c7c024e34bf8b29 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 13 Dec 2018 18:07:58 -0800
Subject: [PATCH 1377/2720] Utility method for gym registration.

PiperOrigin-RevId: 225469240
---
 tensor2tensor/rl/gym_utils.py      | 14 ++++++++++++
 tensor2tensor/rl/gym_utils_test.py | 34 ++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index c2cc7d505..53962e5ad 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -54,3 +54,17 @@ def make_gym_env(name, rl_env_max_episode_steps=-1):
     env = env.env
 
   return gym.wrappers.TimeLimit(env, max_episode_steps=rl_env_max_episode_steps)
+
+
+def register_gym_env(class_entry_point, version="v0"):
+  """Registers the class with its snake case name in Gym and returns it."""
+
+  split_on_colon = class_entry_point.split(":")
+  assert len(split_on_colon) == 2
+
+  class_name = split_on_colon[1]
+  # We have to add the version to conform to gym's API.
+  env_name = "{}-{}".format(class_name, version)
+  gym.envs.register(id=env_name, entry_point=class_entry_point)
+
+  return gym.make(env_name)
diff --git a/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index 9beffa7a1..16dbe4f1f 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -20,10 +20,30 @@
 from __future__ import print_function
 
 import gym
+from gym import spaces
+import numpy as np
 from tensor2tensor.rl import gym_utils
 import tensorflow as tf
 
 
+class SimpleEnv(gym.Env):
+  """A simple environment with a 3x3 observation space, is done on action=1."""
+
+  def __init__(self):
+    self.reward_range = (-1.0, 1.0)
+    self.action_space = spaces.Discrete(2)
+    self.observation_space = spaces.Box(low=0, high=255, shape=(3, 3))
+
+  def reset(self):
+    return self.observation_space.low
+
+  def step(self, action):
+    if action == 0:
+      return self.reset(), -1.0, False, {}
+    else:
+      return self.observation_space.high, +1.0, True, {}
+
+
 class GymUtilsTest(tf.test.TestCase):
 
   # Just make an environment and expect to get one.
@@ -45,6 +65,20 @@ def test_unlimited_env(self):
     self.assertTrue(isinstance(env, gym.wrappers.TimeLimit))
     self.assertTrue(env._max_episode_steps is None)
 
+  def test_gym_registration(self):
+    env = gym_utils.register_gym_env(
+        "tensor2tensor.rl.gym_utils_test:SimpleEnv")
+
+    # Most basic check.
+    self.assertTrue(isinstance(env, gym.Env))
+
+    # Just make sure we got the same environment.
+    self.assertTrue(np.allclose(env.reset(),
+                                np.zeros(shape=(3, 3), dtype=np.uint8)))
+
+    _, _, done, _ = env.step(1)
+    self.assertTrue(done)
+
 
 if __name__ == "__main__":
   tf.test.main()

From d4ff0f9024704f5b0bde3f79c4ed8b69809cccdd Mon Sep 17 00:00:00 2001
From: Youngwook Kim <youngwook.kim@gmail.com>
Date: Fri, 14 Dec 2018 11:38:55 +0900
Subject: [PATCH 1378/2720] Fix Python 3 compatibility issue with
 self_generated_targets (#1300)

---
 tensor2tensor/utils/t2t_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index deace4110..e4fa987db 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1337,7 +1337,7 @@ def estimator_model_fn(cls,
       # by logits["self_generated_targets"].
       tf.logging.info("Replacing targets with model-provided targets.")
       features["targets"] = labels = logits.pop("self_generated_targets")
-      assert logits.keys() == ["logits"], (
+      assert list(logits.keys()) == ["logits"], (
           # See "Returns" in the "top" method docstring for the expected
           # "logits" format when targets are generated at training time.
           "Expect only key 'logits' when there is 'self_generated_targets'. "

From aa14e2f5b50f4cae9275898d6a976e5c9a84ccf2 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 13 Dec 2018 21:01:49 -0800
Subject: [PATCH 1379/2720] Make Universal Transformer run on TPU and change
 default RL mode.

PiperOrigin-RevId: 225483575
---
 tensor2tensor/data_generators/gym_env.py      | 11 +++--
 .../models/research/universal_transformer.py  | 16 ++++++
 .../research/universal_transformer_util.py    | 25 +++++-----
 tensor2tensor/rl/gym_utils.py                 | 49 ++++++++++++++++++-
 tensor2tensor/rl/rl_utils.py                  |  4 +-
 5 files changed, 84 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 341cd814e..90f02cdea 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -577,7 +577,8 @@ class T2TGymEnv(T2TEnv):
 
   def __init__(self, base_env_name=None, batch_size=1, grayscale=False,
                resize_height_factor=2, resize_width_factor=2,
-               rl_env_max_episode_steps=-1, max_num_noops=0, **kwargs):
+               rl_env_max_episode_steps=-1, max_num_noops=0,
+               maxskip_envs=False, **kwargs):
     if base_env_name is None:
       base_env_name = self.base_env_name
     self._base_env_name = base_env_name
@@ -589,9 +590,11 @@ def __init__(self, base_env_name=None, batch_size=1, grayscale=False,
       # Set problem name if not registered.
       self.name = "Gym%s" % base_env_name
 
-    self._envs = [gym_utils.make_gym_env(
-        base_env_name, rl_env_max_episode_steps=rl_env_max_episode_steps)
-                  for _ in range(self.batch_size)]
+    self._envs = [
+        gym_utils.make_gym_env(
+            base_env_name, rl_env_max_episode_steps=rl_env_max_episode_steps,
+            maxskip_env=maxskip_envs)
+        for _ in range(self.batch_size)]
 
     # max_num_noops works only with atari envs.
     if max_num_noops > 0:
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index bf5b3e13f..71b82627f 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -513,6 +513,14 @@ def adaptive_universal_transformer_base():
   return hparams
 
 
+@registry.register_hparams
+def adaptive_universal_transformer_base_tpu():
+  hparams = adaptive_universal_transformer_base()
+  transformer.update_hparams_for_tpu(hparams)
+  hparams.add_step_timing_signal = False
+  return hparams
+
+
 @registry.register_hparams
 def adaptive_universal_transformer_small():
   hparams = universal_transformer_small()
@@ -535,6 +543,14 @@ def adaptive_universal_transformer_global_base():
   return hparams
 
 
+@registry.register_hparams
+def adaptive_universal_transformer_global_base_tpu():
+  hparams = adaptive_universal_transformer_global_base()
+  transformer.update_hparams_for_tpu(hparams)
+  hparams.add_step_timing_signal = False
+  return hparams
+
+
 @registry.register_hparams
 def adaptive_universal_transformer_tall():
   hparams = universal_transformer_small()
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 5377e2e7b..3ed3488a9 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -1200,7 +1200,8 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
   # Do while loop iterations until predicate above is false.
   (_, _, _, remainder, n_updates, new_state) = tf.while_loop(
       should_continue, ut_function,
-      (state, step, halting_probability, remainders, n_updates, previous_state))
+      (state, step, halting_probability, remainders, n_updates, previous_state),
+      maximum_iterations=act_max_steps + 1)
 
   ponder_times = n_updates
   remainders = remainder
@@ -1351,7 +1352,8 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
   # Do while loop iterations until predicate above is false.
   (_, _, _, remainder, n_updates, accumulated_state) = tf.while_loop(
       should_continue, ut_function, (state, step, halting_probability,
-                                     remainders, n_updates, accumulated_state))
+                                     remainders, n_updates, accumulated_state),
+      maximum_iterations=act_max_steps + 1)
 
   ponder_times = n_updates
   remainders = remainder
@@ -1476,10 +1478,8 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
     new_state.set_shape(state_shape)
 
     step += 1
-    return [
-        transformed_state, step, halting_probability, remainders, n_updates,
-        new_state
-    ]
+    return (transformed_state, step, halting_probability,
+            remainders, n_updates, new_state)
 
   # While loop stops when this predicate is FALSE.
   # Ie all (probability < 1-eps AND counter < N) are false.
@@ -1493,7 +1493,8 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
   # Do while loop iterations until predicate above is false.
   (_, _, _, remainder, n_updates, new_state) = tf.while_loop(
       should_continue, ut_function,
-      (state, step, halting_probability, remainders, n_updates, previous_state))
+      (state, step, halting_probability, remainders, n_updates, previous_state),
+      maximum_iterations=act_max_steps + 1)
 
   ponder_times = n_updates
   remainders = remainder
@@ -1516,7 +1517,6 @@ def universal_transformer_act_random(x, hparams, ffn_unit, attention_unit):
     the output tensor,  (ponder_times, remainders)
 
   """
-
   state = x
   act_max_steps = hparams.act_max_steps
   threshold = 1.0 - hparams.act_epsilon
@@ -1623,10 +1623,8 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
       ])
     new_state.set_shape(state_shape)
     step += 1
-    return [
-        transformed_state, step, halting_probability, remainders, n_updates,
-        new_state
-    ]
+    return (transformed_state, step,
+            halting_probability, remainders, n_updates, new_state)
 
   # While loop stops when this predicate is FALSE.
   # Ie all (probability < 1-eps AND counter < N) are false.
@@ -1640,7 +1638,8 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
   # Do while loop iterations until predicate above is false.
   (_, _, _, remainder, n_updates, new_state) = tf.while_loop(
       should_continue, ut_function,
-      (state, step, halting_probability, remainders, n_updates, previous_state))
+      (state, step, halting_probability, remainders, n_updates, previous_state),
+      maximum_iterations=act_max_steps + 1)
 
   ponder_times = n_updates
   remainders = remainder
diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index 53962e5ad..7d8a494a0 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -20,10 +20,44 @@
 from __future__ import print_function
 
 import gym
+import numpy as np
 
 
-def make_gym_env(name, rl_env_max_episode_steps=-1):
-  """Create a gym env optionally wrapped with a time limit wrapper.
+class MaxAndSkipEnv(gym.Wrapper):
+  """Same wrapper as in OpenAI baselines for comparability of results."""
+
+  def __init__(self, env, skip=4):
+    """Return only every `skip`-th frame."""
+    gym.Wrapper.__init__(self, env)
+    # Most recent raw observations (for max pooling across time steps).
+    self._obs_buffer = np.zeros((2,) + env.observation_space.shape,
+                                dtype=np.uint8)
+    self._skip = skip
+
+  def __str__(self):
+    return "MaxAndSkip<%s>" % str(self.env)
+
+  def step(self, action):
+    """Repeat action, sum reward, and max over last observations."""
+    total_reward = 0.0
+    done = None
+    for i in range(self._skip):
+      obs, reward, done, info = self.env.step(action)
+      if i == self._skip - 2: self._obs_buffer[0] = obs
+      if i == self._skip - 1: self._obs_buffer[1] = obs
+      total_reward += reward
+      if done:
+        break
+    # Note that the observation on the done=True frame doesn't matter.
+    max_frame = self._obs_buffer.max(axis=0)
+    return max_frame, total_reward, done, info
+
+  def reset(self, **kwargs):
+    return self.env.reset(**kwargs)
+
+
+def make_gym_env(name, rl_env_max_episode_steps=-1, maxskip_env=False):
+  """Create a gym env optionally with a time limit and maxskip wrapper.
 
   NOTE: The returned env may already be wrapped with TimeLimit!
 
@@ -32,6 +66,7 @@ def make_gym_env(name, rl_env_max_episode_steps=-1):
     rl_env_max_episode_steps: `int` or None - Using any value < 0 returns the
       env as-in, otherwise we impose the requested timelimit. Setting this to
       None returns a wrapped env that doesn't have a step limit.
+    maxskip_env: whether to also use MaxAndSkip wrapper before time limit.
 
   Returns:
     An instance of `gym.Env` or `gym.wrappers.TimeLimit` with the requested
@@ -46,6 +81,13 @@ def make_gym_env(name, rl_env_max_episode_steps=-1):
 
   # If nothing to do, then return the env.
   if rl_env_max_episode_steps and rl_env_max_episode_steps < 0:
+    if maxskip_env:
+      if isinstance(env, gym.wrappers.TimeLimit):
+        # Unwrap time limit and put it above MaxAndSkip for consistency.
+        max_episode_steps = env._max_episode_steps  # pylint: disable=protected-access
+        env = MaxAndSkipEnv(env.env)
+        return gym.wrappers.TimeLimit(env, max_episode_steps=max_episode_steps)
+      return MaxAndSkipEnv(env)
     return env
 
   # Sometimes (mostly?) the env is already wrapped in a TimeLimit wrapper, in
@@ -53,6 +95,9 @@ def make_gym_env(name, rl_env_max_episode_steps=-1):
   if isinstance(env, gym.wrappers.TimeLimit):
     env = env.env
 
+  if maxskip_env:
+    env = MaxAndSkipEnv(env)
+
   return gym.wrappers.TimeLimit(env, max_episode_steps=rl_env_max_episode_steps)
 
 
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 79797ae07..79a5a4dcb 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -105,7 +105,7 @@ def evaluate_all_configs(hparams, agent_model_dir):
 
 def setup_env(hparams, batch_size, max_num_noops):
   """Setup."""
-  game_mode = "Deterministic-v4"
+  game_mode = "NoFrameskip-v4"
   camel_game_name = misc_utils.snakecase_to_camelcase(hparams.game)
   camel_game_name += game_mode
   env_name = camel_game_name
@@ -116,7 +116,7 @@ def setup_env(hparams, batch_size, max_num_noops):
                   resize_width_factor=hparams.resize_width_factor,
                   resize_height_factor=hparams.resize_height_factor,
                   rl_env_max_episode_steps=hparams.rl_env_max_episode_steps,
-                  max_num_noops=max_num_noops)
+                  max_num_noops=max_num_noops, maxskip_envs=True)
   return env
 
 

From d97d2e408cbfeaf17a9dd3a648b7cbbf37d2b564 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 14 Dec 2018 08:40:24 -0800
Subject: [PATCH 1380/2720] Fix error typo in ActNorm.

PiperOrigin-RevId: 225547374
---
 tensor2tensor/layers/reversible_layers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/reversible_layers.py b/tensor2tensor/layers/reversible_layers.py
index 4a4949cbf..88798f9d7 100644
--- a/tensor2tensor/layers/reversible_layers.py
+++ b/tensor2tensor/layers/reversible_layers.py
@@ -56,7 +56,7 @@ def build(self, input_shape):
     if isinstance(last_dim, tf.Dimension):
       last_dim = last_dim.value
     if last_dim is None:
-      raise ValueError('The last dimension of the inputs to `Dense` '
+      raise ValueError('The last dimension of the inputs to `ActNorm` '
                        'should be defined. Found `None`.')
     bias = self.add_weight('bias', [last_dim], dtype=self.dtype)
     log_scale = self.add_weight('log_scale', [last_dim], dtype=self.dtype)

From f74bb9c3e94e8ee01aa13b1d870612d899357edc Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 14 Dec 2018 11:26:43 -0800
Subject: [PATCH 1381/2720] Add some debug logs in trainer_model_free.py, these
 are at verbosity 1 or 2, which means disabled by default and can be enabled
 using "--vmodule=trainer_model_free=2" or more generally using
 "--vmodule=*/tensor2tensor/*=2"

PiperOrigin-RevId: 225574647
---
 tensor2tensor/rl/trainer_model_free.py | 52 ++++++++++++++++++++++----
 tensor2tensor/utils/misc_utils.py      |  5 +++
 tensor2tensor/utils/misc_utils_test.py | 16 ++++++++
 3 files changed, 65 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 19d8de67c..049e0d97e 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -32,6 +32,7 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
+from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -54,6 +55,9 @@ def initialize_env_specs(hparams):
   env = rl_utils.setup_env(hparams, hparams.batch_size,
                            hparams.eval_max_num_noops)
   env.start_new_epoch(0)
+
+  # TODO(afrozm): Decouple env_fn from hparams and return both, is there
+  # even a need to return hparams? Just return the env_fn?
   hparams.add_hparam("env_fn", rl.make_real_env_fn(env))
   return hparams
 
@@ -61,15 +65,30 @@ def initialize_env_specs(hparams):
 def train(hparams, output_dir, report_fn=None):
   """Train."""
   hparams = initialize_env_specs(hparams)
+
+  tf.logging.vlog(1, "HParams in trainer_model_free.train : %s",
+                  misc_utils.pprint_hparams(hparams))
+
+  tf.logging.vlog(1, "Using hparams.base_algo: %s", hparams.base_algo)
   learner = rl_utils.LEARNERS[hparams.base_algo](
-      hparams.frame_stack_size, FLAGS.output_dir, output_dir, total_num_epochs=1
+      hparams.frame_stack_size, output_dir, output_dir, total_num_epochs=1
   )
+
   policy_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
+
   rl_utils.update_hparams_from_hparams(
       policy_hparams, hparams, hparams.base_algo + "_"
   )
+
+  tf.logging.vlog(1, "Policy HParams : %s",
+                  misc_utils.pprint_hparams(policy_hparams))
+
   total_steps = policy_hparams.epochs_num
+  tf.logging.vlog(2, "total_steps: %d", total_steps)
+
   eval_every_epochs = policy_hparams.eval_every_epochs
+  tf.logging.vlog(2, "eval_every_epochs: %d", eval_every_epochs)
+
   if eval_every_epochs == 0:
     eval_every_epochs = total_steps
   policy_hparams.eval_every_epochs = 0
@@ -77,27 +96,44 @@ def train(hparams, output_dir, report_fn=None):
   steps = list(range(eval_every_epochs, total_steps+1, eval_every_epochs))
   if not steps or steps[-1] < eval_every_epochs:
     steps.append(eval_every_epochs)
+
+  tf.logging.vlog(1, "steps: [%s]", ",".join([str(s) for s in steps]))
+
   metric_name = rl_utils.get_metric_name(
       sampling_temp=hparams.eval_sampling_temps[0],
       max_num_noops=hparams.eval_max_num_noops,
       clipped=False
   )
-  for step in steps:
+
+  tf.logging.vlog(1, "metric_name: %s", metric_name)
+
+  for i, step in enumerate(steps):
+    tf.logging.info("Starting training iteration [%d] for [%d] steps.", i, step)
+
     policy_hparams.epochs_num = step
-    learner.train(
-        hparams.env_fn, policy_hparams, simulated=False, save_continuously=True,
-        epoch=0
-    )
+    learner.train(hparams.env_fn,
+                  policy_hparams,
+                  simulated=False,
+                  save_continuously=True,
+                  epoch=0)
+
+    tf.logging.info("Ended training iteration [%d] for [%d] steps.", i, step)
+
     eval_metrics = rl_utils.evaluate_all_configs(hparams, output_dir)
-    tf.logging.info("Agent eval metrics:\n{}".format(
-        pprint.pformat(eval_metrics)))
+
+    tf.logging.info(
+        "Agent eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
+
     if report_fn:
       report_fn(eval_metrics[metric_name], step)
 
 
 def main(_):
   hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
+
+  tf.logging.log("Starting model free training.")
   train(hparams, FLAGS.output_dir)
+  tf.logging.log("Ended model free training.")
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/utils/misc_utils.py b/tensor2tensor/utils/misc_utils.py
index d8d1b18da..bef5fc0c6 100644
--- a/tensor2tensor/utils/misc_utils.py
+++ b/tensor2tensor/utils/misc_utils.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import pprint
 import re
 
 # Camel case to snake case utils
@@ -34,3 +35,7 @@ def camelcase_to_snakecase(name):
 def snakecase_to_camelcase(name):
   return "".join([w[0].upper() + w[1:] for w in name.split("_")])
 
+
+def pprint_hparams(hparams):
+  """Represents hparams using its dictionary and calls pprint.pformat on it."""
+  return "\n{}".format(pprint.pformat(hparams.values(), width=1))
diff --git a/tensor2tensor/utils/misc_utils_test.py b/tensor2tensor/utils/misc_utils_test.py
index 3a28fa629..ab8da2bd1 100644
--- a/tensor2tensor/utils/misc_utils_test.py
+++ b/tensor2tensor/utils/misc_utils_test.py
@@ -55,6 +55,22 @@ def test_snakecase_to_camelcase(self):
     self.assertEqual("LstmSeq2Seq",
                      misc_utils.snakecase_to_camelcase("lstm_seq2_seq"))
 
+  def test_pprint_hparams(self):
+    hparams = tf.contrib.training.HParams(
+        int_=1, str_="str", bool_=True, float_=1.1, list_int=[1, 2], none=None)
+
+    # pylint: disable=g-inconsistent-quotes
+    expected_string = r"""
+{'bool_': True,
+ 'float_': 1.1,
+ 'int_': 1,
+ 'list_int': [1,
+              2],
+ 'none': None,
+ 'str_': 'str'}"""
+    # pylint: enable=g-inconsistent-quotes
+
+    self.assertEqual(expected_string, misc_utils.pprint_hparams(hparams))
 
 if __name__ == "__main__":
   tf.test.main()

From 83861f3d9b23e196cbc134fd5a6ad2e5dc5edd96 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 14 Dec 2018 15:09:58 -0800
Subject: [PATCH 1382/2720] Pass features to tpu_estimator correctly. I've also
 removed the logits and labels arguments from the all_metrics function to
 simplify the code.

PiperOrigin-RevId: 225611492
---
 tensor2tensor/utils/t2t_model.py | 114 ++++++++++++++++++++++++-------
 1 file changed, 91 insertions(+), 23 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index e4fa987db..9a0c374ec 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -58,6 +58,63 @@
     lambda method_name: _no_problem_err_str % (method_name, method_name))
 
 
+def _flatten_dict(original_dict):
+  """Flatten dict of dicts into a single dict with appropriate prefixes.
+
+  Handles only 2 levels of nesting in the original dict.
+
+  Args:
+    original_dict: Dict which may contain one or more dicts.
+  Returns:
+    flat_dict: Dict without any nesting. Any dicts in the original dict have
+      their keys as prefixes in the new dict.
+  Raises:
+    ValueError if the original dict has more than two levels of nesting.
+  """
+  flat_dict = {}
+  for key, value in original_dict.items():
+    if isinstance(value, dict):
+      for name, tensor in value.items():
+        if isinstance(tensor, dict):
+          raise ValueError("flatten_dict only handles 2 levels of nesting.")
+        flat_key = "__" + key + "_" + name
+        flat_dict[flat_key] = tensor
+    else:
+      flat_dict[key] = value
+
+  return flat_dict
+
+
+def _unflatten_dict(flat_dict, prefixes):
+  """Returns a dict of dicts if any prefixes match keys in the flat dict.
+
+    The function handles the case where the prefix may not be a dict.
+
+  Args:
+    flat_dict: A dict without any nesting.
+    prefixes: A list of strings which may have been dicts in the
+      original structure.
+
+  """
+  original_dict = {}
+  for key, value in flat_dict.items():
+    prefix_found = False
+    for prefix in prefixes:
+      full_prefix = "__" + prefix + "_"
+      if key.startswith(full_prefix):
+        # Add a dict to the original dict with key=prefix
+        if prefix not in original_dict:
+          original_dict[prefix] = {}
+        original_dict[prefix][key[len(full_prefix):]] = value
+        prefix_found = True
+        break
+    if not prefix_found:
+      # No key matched a prefix in the for loop.
+      original_dict[key] = value
+
+  return original_dict
+
+
 class T2TModel(base.Layer):
   """Abstract base class for models.
 
@@ -1458,22 +1515,30 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
 
     if common_layers.is_xla_compiled():
       remove_summaries()
-      if isinstance(logits, dict):
-        eval_metrics_fn = create_tpu_eval_metrics_fn(problem, hparams)
-        # For TPU, logits dict will be passed as keyword arguments to
-        # eval_metrics_fn. Here we add the labels to those arguments.
-        logits.update({"labels": labels})
-        logits.update({"features": features})
-        return tf.contrib.tpu.TPUEstimatorSpec(
-            tf.estimator.ModeKeys.EVAL,
-            eval_metrics=(eval_metrics_fn, logits),
-            loss=loss)
-      else:
-        eval_metrics_fn = create_tpu_eval_metrics_fn(problem, hparams)
-        return tf.contrib.tpu.TPUEstimatorSpec(
-            tf.estimator.ModeKeys.EVAL,
-            eval_metrics=(eval_metrics_fn, [logits, labels, features]),
-            loss=loss)
+      eval_metrics_fn = create_tpu_eval_metrics_fn(problem, hparams)
+
+      batch_size = [feature.shape.as_list()[0] for _, feature
+                    in features.items() if feature.shape.ndims][0]
+
+      # Add batch dimension to all features since tpu requires the batch
+      # dimension on all tensors.
+      for name, feature in features.items():
+        if not feature.shape.as_list():
+          # All features must have a batch dimension
+          feature = tf.tile(tf.expand_dims(feature, 0), [batch_size])
+        features[name] = feature
+
+      eval_metrics_fn_args = dict(
+          logits=logits,  # possibly a dict
+          labels=labels,
+          features=features,  # dict
+      )
+
+      eval_metrics_fn_flat_args = _flatten_dict(eval_metrics_fn_args)
+      return tf.contrib.tpu.TPUEstimatorSpec(
+          tf.estimator.ModeKeys.EVAL,
+          eval_metrics=(eval_metrics_fn, eval_metrics_fn_flat_args),
+          loss=loss)
     else:
       task_list = [problem]
       if hasattr(problem, "task_list"):
@@ -1694,15 +1759,18 @@ def wrapped_metric_fn(logits, labels, features):
       name = "metrics-%s/%s" % (problem.name, metric)
       metric_fns.append((name, make_metric_fn(metrics.METRICS_FNS[metric])))
 
-  def all_metrics_fn(logits=None, labels=None, **kwargs):
+  def all_metrics_fn(**kwargs):
     """Construct metrics dictionary."""
-    metrics_dict = {}
 
-    if logits is None:
-      logits = kwargs
-      features = logits["features"]
-    else:
-      features = kwargs["features"]
+    original_kwargs = _unflatten_dict(kwargs, prefixes=["logits", "features"])
+    del kwargs
+
+    logits = original_kwargs["logits"]
+    labels = original_kwargs["labels"]
+    features = original_kwargs["features"]
+    del original_kwargs
+
+    metrics_dict = {}
 
     for name, fn in metric_fns:
       if isinstance(logits, dict) and isinstance(labels, dict):

From 0b4a3f33c7ba54c9d86debe2dff55c019d9dcb7e Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Sat, 15 Dec 2018 09:40:17 -0800
Subject: [PATCH 1383/2720] Extend MADE's signature: [..., length, channels] ->
 [..., length, units].

References
+ https://stackoverflow.com/questions/51822211/tensorflow-how-to-tile-a-tensor-that-duplicate-in-certain-order

PiperOrigin-RevId: 225678133
---
 tensor2tensor/layers/reversible_layers.py     | 88 +++++++++++++------
 .../layers/reversible_layers_test.py          | 36 ++++----
 2 files changed, 83 insertions(+), 41 deletions(-)

diff --git a/tensor2tensor/layers/reversible_layers.py b/tensor2tensor/layers/reversible_layers.py
index 88798f9d7..ba7cf3e73 100644
--- a/tensor2tensor/layers/reversible_layers.py
+++ b/tensor2tensor/layers/reversible_layers.py
@@ -109,17 +109,20 @@ def log_det_jacobian(self, inputs):
 class MADE(tf.keras.Model):
   """Masked autoencoder for distribution estimation (Germain et al., 2015).
 
-  MADE takes as input a real Tensor of shape [..., length] and returns a
-  Tensor of shape [..., num_heads * length] and same dtype. It masks layer
-  weights to respect autoregressive constraints: for a given ordering, each
-  input dimension can be reconstructed from previous input dimensions. The
-  output dimensions represent multiple heads for, e.g., location and scale
-  transforms in a flow.
+  MADE takes as input a real Tensor of shape [..., length, channels] and returns
+  a Tensor of shape [..., length, units] and same dtype. It masks layer weights
+  to satisfy autoregressive constraints with respect to the length dimension. In
+  particular, for a given ordering, each input dimension of length can be
+  reconstructed from previous dimensions.
+
+  The output's units dimension captures per-time-step representations. For
+  example, setting units to 2 can parameterize the location and log-scale of an
+  autoregressive Gaussian distribution.
   """
 
   def __init__(self,
+               units,
                hidden_dims,
-               num_heads=2,
                input_order='left-to-right',
                hidden_order='left-to-right',
                activation=None,
@@ -128,24 +131,26 @@ def __init__(self,
     """Constructs network.
 
     Args:
+      units: Positive integer, dimensionality of the output space.
       hidden_dims: list with the number of hidden units per layer. It does not
         include the output layer; those number of units will always be set to
-        the input dimension multiplied by `num_heads`.
-      num_heads: The number of output heads. The default is 2 representing
-        the location and scale transform of an autoregressive flow.
+        the input dimension multiplied by `num_heads`. Each hidden unit size
+        must be at least the size of length (otherwise autoregressivity is not
+        possible).
       input_order: Order of degrees to the input units: 'random',
         'left-to-right', 'right-to-left', or an array of an explicit order.
         For example, 'left-to-right' builds an autoregressive model
         p(x) = p(x1) p(x2 | x1) ... p(xD | x<D).
       hidden_order: Order of degrees to the hidden units: 'random',
-        'left-to-right'.
+        'left-to-right'. If 'left-to-right', hidden units are allocated equally
+        (up to a remainder term) to each degree.
       activation: Activation function.
       use_bias: Whether to use a bias.
       **kwargs: Keyword arguments of parent class.
     """
     super(MADE, self).__init__(**kwargs)
+    self.units = int(units)
     self.hidden_dims = hidden_dims
-    self.num_heads = num_heads
     self.input_order = input_order
     self.hidden_order = hidden_order
     self.activation = tf.keras.activations.get(activation)
@@ -154,17 +159,38 @@ def __init__(self,
 
   def build(self, input_shape):
     input_shape = tf.TensorShape(input_shape)
-    last_dim = input_shape[-1]
-    if isinstance(last_dim, tf.Dimension):
-      last_dim = last_dim.value
-    if last_dim is None:
-      raise ValueError('The last dimension of the inputs to '
+    length = input_shape[-2]
+    channels = input_shape[-1]
+    if isinstance(length, tf.Dimension):
+      length = length.value
+    if isinstance(channels, tf.Dimension):
+      channels = channels.value
+    if length is None or channels is None:
+      raise ValueError('The two last dimensions of the inputs to '
                        '`MADE` should be defined. Found `None`.')
-    masks = create_masks(input_dim=last_dim,
+    masks = create_masks(input_dim=length,
                          hidden_dims=self.hidden_dims,
                          input_order=self.input_order,
                          hidden_order=self.hidden_order)
-    for l in range(len(self.hidden_dims)):
+
+    # Input-to-hidden layer: [..., length, channels] -> [..., hidden_dims[0]].
+    self.network.add(tf.keras.layers.Reshape([length * channels]))
+    # Tile the mask so each element repeats contiguously; this is compatible
+    # with the autoregressive contraints unlike naive tiling.
+    mask = masks[0]
+    mask = tf.tile(mask[:, tf.newaxis, :], [1, channels, 1])
+    mask = tf.reshape(mask, [mask.shape[0] * channels, mask.shape[-1]])
+    if self.hidden_dims:
+      layer = tf.keras.layers.Dense(
+          self.hidden_dims[0],
+          kernel_initializer=make_masked_initializer(mask),
+          kernel_constraint=make_masked_constraint(mask),
+          activation=self.activation,
+          use_bias=self.use_bias)
+      self.network.add(layer)
+
+    # Hidden-to-hidden layers: [..., hidden_dims[l-1]] -> [..., hidden_dims[l]].
+    for l in range(1, len(self.hidden_dims)):
       layer = tf.keras.layers.Dense(
           self.hidden_dims[l],
           kernel_initializer=make_masked_initializer(masks[l]),
@@ -173,14 +199,21 @@ def build(self, input_shape):
           use_bias=self.use_bias)
       self.network.add(layer)
 
-    mask = tf.tile(masks[-1], [1, self.num_heads])
+    # Hidden-to-output layer: [..., hidden_dims[-1]] -> [..., length, units].
+    # Tile the mask so each element repeats contiguously; this is compatible
+    # with the autoregressive contraints unlike naive tiling.
+    if self.hidden_dims:
+      mask = masks[-1]
+    mask = tf.tile(mask[..., tf.newaxis], [1, 1, self.units])
+    mask = tf.reshape(mask, [mask.shape[0], mask.shape[1] * self.units])
     layer = tf.keras.layers.Dense(
-        last_dim * self.num_heads,
+        length * self.units,
         kernel_initializer=make_masked_initializer(mask),
         kernel_constraint=make_masked_constraint(mask),
         activation=None,
         use_bias=self.use_bias)
     self.network.add(layer)
+    self.network.add(tf.keras.layers.Reshape([length, self.units]))
     self.built = True
 
   def call(self, inputs):
@@ -199,14 +232,15 @@ def create_degrees(input_dim,
   Args:
     input_dim: Number of inputs.
     hidden_dims: list with the number of hidden units per layer. It does not
-      include the output layer; those number of units will always be set to
-      input_dim downstream.
+      include the output layer. Each hidden unit size must be at least the size
+      of length (otherwise autoregressivity is not possible).
     input_order: Order of degrees to the input units: 'random', 'left-to-right',
       'right-to-left', or an array of an explicit order. For example,
       'left-to-right' builds an autoregressive model
       p(x) = p(x1) p(x2 | x1) ... p(xD | x<D).
     hidden_order: Order of degrees to the hidden units: 'random',
-      'left-to-right'.
+      'left-to-right'. If 'left-to-right', hidden units are allocated equally
+      (up to a remainder term) to each degree.
   """
   if (isinstance(input_order, str) and
       input_order not in ('random', 'left-to-right', 'right-to-left')):
@@ -250,13 +284,15 @@ def create_masks(input_dim,
     input_dim: Number of inputs.
     hidden_dims: list with the number of hidden units per layer. It does not
       include the output layer; those number of units will always be set to
-      input_dim downstream.
+      input_dim downstream. Each hidden unit size must be at least the size of
+      length (otherwise autoregressivity is not possible).
     input_order: Order of degrees to the input units: 'random', 'left-to-right',
       'right-to-left', or an array of an explicit order. For example,
       'left-to-right' builds an autoregressive model
       p(x) = p(x1) p(x2 | x1) ... p(xD | x<D).
     hidden_order: Order of degrees to the hidden units: 'random',
-      'left-to-right'.
+      'left-to-right'. If 'left-to-right', hidden units are allocated equally
+      (up to a remainder term) to each degree.
   """
   degrees = create_degrees(input_dim, hidden_dims, input_order, hidden_order)
   masks = []
diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
index 95a77442d..f60676952 100644
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -58,57 +58,63 @@ def testMADELeftToRight(self):
     np.random.seed(83243)
     batch_size = 2
     length = 3
-    network = reversible.MADE([4], activation=tf.nn.relu)
-    inputs = tf.zeros([batch_size, length])
+    channels = 1
+    units = 5
+    network = reversible.MADE(units, [4], activation=tf.nn.relu)
+    inputs = tf.zeros([batch_size, length, channels])
     outputs = network(inputs)
 
     num_weights = sum([np.prod(weight.shape) for weight in network.weights])
     self.assertEqual(len(network.weights), 4)
-    self.assertEqual(num_weights, (3*4 + 4) + (4*3*2 + 3*2))
+    self.assertEqual(num_weights, (3*1*4 + 4) + (4*3*5 + 3*5))
 
     self.evaluate(tf.global_variables_initializer())
     outputs_val = self.evaluate(outputs)
-    self.assertAllEqual(outputs_val[:, 0], tf.zeros(batch_size))
-    self.assertEqual(outputs_val.shape, (batch_size, 2 * length))
+    self.assertAllEqual(outputs_val[:, 0, :], np.zeros((batch_size, units)))
+    self.assertEqual(outputs_val.shape, (batch_size, length, units))
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testMADERightToLeft(self):
     np.random.seed(1328)
     batch_size = 2
     length = 3
-    network = reversible.MADE([4, 3],
+    channels = 5
+    units = 1
+    network = reversible.MADE(units, [4, 3],
                               input_order='right-to-left',
                               activation=tf.nn.relu,
                               use_bias=False)
-    inputs = tf.zeros([batch_size, length])
+    inputs = tf.zeros([batch_size, length, channels])
     outputs = network(inputs)
 
     num_weights = sum([np.prod(weight.shape) for weight in network.weights])
     self.assertEqual(len(network.weights), 3)
-    self.assertEqual(num_weights, 3*4 + 4*3 + 3*3*2)
+    self.assertEqual(num_weights, 3*5*4 + 4*3 + 3*3*1)
 
     self.evaluate(tf.global_variables_initializer())
     outputs_val = self.evaluate(outputs)
-    self.assertAllEqual(outputs_val[:, -1], tf.zeros(batch_size))
-    self.assertEqual(outputs_val.shape, (batch_size, 2 * length))
+    self.assertAllEqual(outputs_val[:, -1, :], np.zeros((batch_size, units)))
+    self.assertEqual(outputs_val.shape, (batch_size, length, units))
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testMADENoHidden(self):
     np.random.seed(532)
     batch_size = 2
     length = 3
-    network = reversible.MADE([], input_order='left-to-right')
-    inputs = tf.zeros([batch_size, length])
+    channels = 5
+    units = 4
+    network = reversible.MADE(units, [], input_order='left-to-right')
+    inputs = tf.zeros([batch_size, length, channels])
     outputs = network(inputs)
 
     num_weights = sum([np.prod(weight.shape) for weight in network.weights])
     self.assertEqual(len(network.weights), 2)
-    self.assertEqual(num_weights, 3*3*2 + 3*2)
+    self.assertEqual(num_weights, 3*5*3*4 + 3*4)
 
     self.evaluate(tf.global_variables_initializer())
     outputs_val = self.evaluate(outputs)
-    self.assertAllEqual(outputs_val[:, 0], tf.zeros(batch_size))
-    self.assertEqual(outputs_val.shape, (batch_size, 2 * length))
+    self.assertAllEqual(outputs_val[:, 0, :], np.zeros((batch_size, units)))
+    self.assertEqual(outputs_val.shape, (batch_size, length, units))
 
 
 if __name__ == '__main__':

From 7f8d9d8d5113001e9f6ee8fae6442cd422e2d1da Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 17 Dec 2018 12:39:22 -0800
Subject: [PATCH 1384/2720] Make name kwarg to positional embeddings optional.

PiperOrigin-RevId: 225871687
---
 tensor2tensor/layers/common_attention.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 2e5af4c62..485dc41e1 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -619,7 +619,7 @@ def add_timing_signal_nd(x, min_timescale=1.0, max_timescale=1.0e4):
   return x
 
 
-def add_positional_embedding(x, max_length, name, positions=None):
+def add_positional_embedding(x, max_length, name=None, positions=None):
   """Adds positional embedding.
 
   Args:
@@ -645,7 +645,7 @@ def add_positional_embedding(x, max_length, name, positions=None):
       return x + tf.gather(var, tf.to_int32(positions))
 
 
-def add_positional_embedding_nd(x, max_length, name):
+def add_positional_embedding_nd(x, max_length, name=None):
   """Adds n-dimensional positional embedding.
 
   The embeddings add to all positional dimensions of the tensor.

From 054c34ca1bc7df4a5893dba1fb9cf084d0b739ed Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 17 Dec 2018 15:27:20 -0800
Subject: [PATCH 1385/2720] Explicitly pass values kwarg as it is currently
 being treated as the default_name kwarg instead. This causes an exception to
 be thrown in eager mode.

PiperOrigin-RevId: 225900620
---
 tensor2tensor/layers/common_layers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 9cd9b4e11..affd409e4 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -253,7 +253,7 @@ def expand_squeeze_to_nd(x, n, squeeze_dim=2, expand_dim=-1):
 
 def standardize_images(x):
   """Image standardization on batches and videos."""
-  with tf.name_scope("standardize_images", [x]):
+  with tf.name_scope("standardize_images", values=[x]):
     x_shape = shape_list(x)
     x = tf.to_float(tf.reshape(x, [-1] + x_shape[-3:]))
     x_mean = tf.reduce_mean(x, axis=[1, 2], keepdims=True)

From 79468f0e7be7bb1dd2119f29417691f9139b92b2 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 18 Dec 2018 00:32:54 +0100
Subject: [PATCH 1386/2720] Make parameter sharing work with stochastic
 discrete (#1304)

---
 tensor2tensor/models/research/rl.py           | 57 ++++++++++++++++---
 tensor2tensor/rl/ppo_learner.py               |  3 +-
 .../rl/trainer_model_based_params.py          |  9 +++
 3 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index bceceb0d2..d69f2b70d 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -28,6 +28,7 @@
 from tensor2tensor.layers import discretization
 from tensor2tensor.layers import modalities
 from tensor2tensor.models.video import basic_deterministic_params
+from tensor2tensor.models.video import basic_stochastic
 from tensor2tensor.rl.envs.py_func_batch_env import PyFuncBatchEnv
 from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
 from tensor2tensor.rl.envs.simulated_batch_gym_env import SimulatedBatchGymEnv
@@ -158,6 +159,23 @@ def ppo_tiny_world_model():
   return hparams
 
 
+@registry.register_hparams
+def ppo_original_world_model_stochastic_discrete():
+  """Atari parameters with stochastic discrete world model as policy."""
+  hparams = ppo_original_params()
+  hparams.policy_network = "next_frame_basic_stochastic_discrete"
+  hparams_keys = hparams.values().keys()
+  video_hparams = basic_stochastic.next_frame_basic_stochastic_discrete()
+  for (name, value) in six.iteritems(video_hparams.values()):
+    if name in hparams_keys:
+      hparams.set_hparam(name, value)
+    else:
+      hparams.add_hparam(name, value)
+  # To avoid OOM. Probably way to small.
+  hparams.optimization_batch_size = 1
+  return hparams
+
+
 def make_real_env_fn(env):
   """Creates a function returning a given real env, in or out of graph.
 
@@ -199,27 +217,36 @@ def get_policy(observations, hparams, action_space):
   if not isinstance(action_space, gym.spaces.Discrete):
     raise ValueError("Expecting discrete action space.")
 
-  policy_problem = DummyPolicyProblem(action_space)
+  obs_shape = common_layers.shape_list(observations)
+  (frame_height, frame_width) = obs_shape[2:4]
+  policy_problem = DummyPolicyProblem(action_space, frame_height, frame_width)
   trainer_lib.add_problem_hparams(hparams, policy_problem)
   hparams.force_full_predict = True
   model = registry.model(hparams.policy_network)(
       hparams, tf.estimator.ModeKeys.TRAIN
   )
-  obs_shape = common_layers.shape_list(observations)
+  try:
+    num_target_frames = hparams.video_num_target_frames
+  except AttributeError:
+    num_target_frames = 1
   features = {
       "inputs": observations,
       "input_action": tf.zeros(obs_shape[:2] + [1], dtype=tf.int32),
       "input_reward": tf.zeros(obs_shape[:2] + [1], dtype=tf.int32),
-      "targets": tf.zeros(obs_shape[:1] + [1] + obs_shape[2:]),
-      "target_action": tf.zeros(obs_shape[:1] + [1, 1], dtype=tf.int32),
-      "target_reward": tf.zeros(obs_shape[:1] + [1, 1], dtype=tf.int32),
-      "target_policy": tf.zeros(obs_shape[:1] + [1] + [action_space.n]),
-      "target_value": tf.zeros(obs_shape[:1] + [1])
+      "targets": tf.zeros(obs_shape[:1] + [num_target_frames] + obs_shape[2:]),
+      "target_action": \
+          tf.zeros(obs_shape[:1] + [num_target_frames, 1], dtype=tf.int32),
+      "target_reward": \
+          tf.zeros(obs_shape[:1] + [num_target_frames, 1], dtype=tf.int32),
+      "target_policy": \
+          tf.zeros(obs_shape[:1] + [num_target_frames] + [action_space.n]),
+      "target_value": \
+          tf.zeros(obs_shape[:1] + [num_target_frames])
   }
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
     t2t_model.create_dummy_vars()
     (targets, _) = model(features)
-  return (targets["target_policy"], targets["target_value"])
+  return (targets["target_policy"][:, 0, :], targets["target_value"][:, 0])
 
 
 @registry.register_hparams
@@ -315,9 +342,21 @@ def loss(self, *args, **kwargs):
 class DummyPolicyProblem(video_utils.VideoProblem):
   """Dummy Problem for running the policy."""
 
-  def __init__(self, action_space):
+  def __init__(self, action_space, frame_height, frame_width):
     super(DummyPolicyProblem, self).__init__()
     self.action_space = action_space
+    self._frame_height = frame_height
+    self._frame_width = frame_width
+
+  @property
+  def frame_height(self):
+    """Height of each frame."""
+    return self._frame_height
+
+  @property
+  def frame_width(self):
+    """Width of each frame."""
+    return self._frame_width
 
   @property
   def num_actions(self):
diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 001ec50eb..a68b79425 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -391,8 +391,9 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
         )
         action = common_layers.sample_with_temperature(logits, sampling_temp)
         action = tf.cast(action, tf.int32)
+        action = tf.reshape(action, shape=(num_agents,))
 
-        reward, done = batch_env.simulate(action[:, 0, ...])
+        reward, done = batch_env.simulate(action)
 
         pdf = tfp.distributions.Categorical(logits=logits).prob(action)
         pdf = tf.reshape(pdf, shape=(num_agents,))
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 3cbe71ee6..9a65a0342 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -307,6 +307,15 @@ def rlmb_base_stochastic_discrete():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_base_stochastic_discrete_param_sharing():
+  """Base setting with stochastic discrete model with parameter sharing."""
+  hparams = rlmb_base_stochastic_discrete()
+  hparams.wm_policy_param_sharing = True
+  hparams.base_algo_params = "ppo_original_world_model_stochastic_discrete"
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_long_stochastic_discrete():
   """Long setting with stochastic discrete model."""

From 0f75a2df59840c070332f78fd827b3d04e35bd5b Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Mon, 17 Dec 2018 15:33:23 -0800
Subject: [PATCH 1387/2720] internal merge of PR #1304

PiperOrigin-RevId: 225901620
---
 tensor2tensor/models/research/rl.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index d69f2b70d..1ee226659 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -234,14 +234,14 @@ def get_policy(observations, hparams, action_space):
       "input_action": tf.zeros(obs_shape[:2] + [1], dtype=tf.int32),
       "input_reward": tf.zeros(obs_shape[:2] + [1], dtype=tf.int32),
       "targets": tf.zeros(obs_shape[:1] + [num_target_frames] + obs_shape[2:]),
-      "target_action": \
-          tf.zeros(obs_shape[:1] + [num_target_frames, 1], dtype=tf.int32),
-      "target_reward": \
-          tf.zeros(obs_shape[:1] + [num_target_frames, 1], dtype=tf.int32),
-      "target_policy": \
-          tf.zeros(obs_shape[:1] + [num_target_frames] + [action_space.n]),
-      "target_value": \
-          tf.zeros(obs_shape[:1] + [num_target_frames])
+      "target_action": tf.zeros(
+          obs_shape[:1] + [num_target_frames, 1], dtype=tf.int32),
+      "target_reward": tf.zeros(
+          obs_shape[:1] + [num_target_frames, 1], dtype=tf.int32),
+      "target_policy": tf.zeros(
+          obs_shape[:1] + [num_target_frames] + [action_space.n]),
+      "target_value": tf.zeros(
+          obs_shape[:1] + [num_target_frames])
   }
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
     t2t_model.create_dummy_vars()

From 47ced4eecd13974e3a690a80b1f60effda089fa9 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 17 Dec 2018 16:48:01 -0800
Subject: [PATCH 1388/2720] Skip evaluation on both train and validation set
 during initialization. The ckpt-step is inferred from the checkpoint
 directory during evaluation.

PiperOrigin-RevId: 225913130
---
 tensor2tensor/utils/decoding.py    |  7 +++++--
 tensor2tensor/utils/trainer_lib.py | 20 +++++++++++++++-----
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 3ff3893a7..464316900 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -841,13 +841,16 @@ def _decode_input_tensor_to_features_dict(feature_map, hparams):
   return features
 
 
+def get_step_from_ckpt_path(path):
+  return int(os.path.basename(path).split("-")[1])
+
+
 def latest_checkpoint_step(ckpt_dir):
   ckpt = tf.train.get_checkpoint_state(ckpt_dir)
   if not ckpt:
     return None
   path = ckpt.model_checkpoint_path
-  step = int(path.split("-")[-1])
-  return step
+  return get_step_from_ckpt_path(path)
 
 
 class DecodeHookArgs(collections.namedtuple(
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 4e448b5b9..7ac9e8fef 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -524,14 +524,24 @@ def evaluate_on_train_data(self):
 
   def continuous_eval(self):
     """Evaluate until checkpoints stop being produced."""
-    for _ in next_checkpoint(self._hparams.model_dir,
-                             self._hparams.eval_timeout_mins):
+    for ckpt_path in next_checkpoint(self._hparams.model_dir,
+                                     self._hparams.eval_timeout_mins):
+      # Skip zero'th step.
+      train_step = decoding.get_step_from_ckpt_path(ckpt_path)
+      if train_step == 0:
+        tf.logging.info("Skipping evaluation at step 0")
+        continue
       self.evaluate()
 
   def continuous_eval_on_train_data(self):
     """Evaluate on train data until checkpoints stop being produced."""
-    for _ in next_checkpoint(self._hparams.model_dir,
-                             self._hparams.eval_timeout_mins):
+    for ckpt_path in next_checkpoint(self._hparams.model_dir,
+                                     self._hparams.eval_timeout_mins):
+      # Skip zero'th step.
+      train_step = decoding.get_step_from_ckpt_path(ckpt_path)
+      if train_step == 0:
+        tf.logging.info("Skipping evaluation at step 0")
+        continue
       self.evaluate_on_train_data()
 
   def test(self):
@@ -600,7 +610,7 @@ def continuous_decode_on_eval_data(self):
       ckpt_generator = next_checkpoint(self._hparams.model_dir)
 
     for ckpt in ckpt_generator:
-      current_step = int(os.path.basename(ckpt).split("-")[1])
+      current_step = decoding.get_step_from_ckpt_path(ckpt)
       tf.logging.info("Decoding step %d" % current_step)
       # Skip checkpoint 0.
       if current_step == 0:

From b3feed34748e6c711de26a0a03c69c82416acdbc Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <babaeizadeh@gmail.com>
Date: Mon, 17 Dec 2018 19:33:09 -0600
Subject: [PATCH 1389/2720] Fixing decode for extra output frames (#1307)

---
 tensor2tensor/models/video/base_vae.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/video/base_vae.py b/tensor2tensor/models/video/base_vae.py
index d4c25b095..5b1716bf3 100644
--- a/tensor2tensor/models/video/base_vae.py
+++ b/tensor2tensor/models/video/base_vae.py
@@ -105,7 +105,7 @@ def construct_latent_tower(self, images, time_axis):
     latent_num_frames = self.hparams.latent_num_frames
     tf.logging.info("Creating latent tower with %d frames." % latent_num_frames)
     if latent_num_frames > 0:
-      images = images[:latent_num_frames]
+      images = images[:, :latent_num_frames]
 
     return common_video.conv_latent_tower(
         images=images,

From e663a0c232b877ce259f346365919bd0a4316c54 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 18 Dec 2018 11:51:11 -0800
Subject: [PATCH 1390/2720] Fix my typo from an earlier change.

PiperOrigin-RevId: 226033790
---
 tensor2tensor/rl/trainer_model_free.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 049e0d97e..91201d89e 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -131,9 +131,9 @@ def train(hparams, output_dir, report_fn=None):
 def main(_):
   hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
 
-  tf.logging.log("Starting model free training.")
+  tf.logging.info("Starting model free training.")
   train(hparams, FLAGS.output_dir)
-  tf.logging.log("Ended model free training.")
+  tf.logging.info("Ended model free training.")
 
 
 if __name__ == "__main__":

From dbdced05efb86e8a61cfdc57cf781ece1e3b43e1 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 18 Dec 2018 12:38:12 -0800
Subject: [PATCH 1391/2720] Fixes the variable scoping in
 common_video.basic_lstm.

PiperOrigin-RevId: 226041307
---
 tensor2tensor/layers/common_video.py      |  4 +++-
 tensor2tensor/layers/common_video_test.py | 12 ++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 8fc578d1a..cbac48912 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -60,7 +60,9 @@ def decode_to_shape(inputs, shape, scope):
 def basic_lstm(inputs, state, num_units, name=None):
   """Basic LSTM."""
   input_shape = common_layers.shape_list(inputs)
-  cell = tf.contrib.rnn.BasicLSTMCell(num_units, name=name)
+  # reuse parameters across time-steps.
+  cell = tf.nn.rnn_cell.BasicLSTMCell(
+      num_units, name=name, reuse=tf.AUTO_REUSE)
   if state is None:
     state = cell.zero_state(input_shape[0], tf.float32)
   outputs, new_state = cell(inputs, state)
diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index 6dc9751b7..2998879ef 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -132,6 +132,18 @@ def check_if_patch_exists(self, videos, video_patches, num_frames):
         is_present.append(np.allclose(curr_patch, video_patch))
       self.assertTrue(np.any(is_present))
 
+  def testBasicLstm(self):
+    """Tests that the parameters of the LSTM are shared across time."""
+    with tf.Graph().as_default():
+      state = None
+      for _ in range(10):
+        inputs = tf.random_uniform(shape=(32, 16))
+        _, state = common_video.basic_lstm(
+            inputs, state, num_units=100, name="basic")
+      num_params = np.sum([np.prod(v.shape) for v in tf.trainable_variables()])
+      # 4 * ((100 + 16)*100 + 100) => 4 * (W_{xh} + W_{hh} + b)
+      self.assertEqual(num_params, 46800)
+
   @parameterized.named_parameters(
       ("two_frames", 2), ("ten_frames", 10), ("default", -1))
   def testExtractRandomVideoPatch(self, num_frames=2):

From e12c64f96a7fc8794a0b7d5d3129e588ee1bbe9d Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Tue, 18 Dec 2018 13:22:48 -0800
Subject: [PATCH 1392/2720] Make decode_from_file work on TPU. For
 Mesh-TensorFlow, this will only work after the fix to avoid concatentating
 the outfeeds in BROADCAST mode.

PiperOrigin-RevId: 226048381
---
 tensor2tensor/bin/t2t_decoder.py         | 18 ++++---
 tensor2tensor/models/mtf_transformer2.py | 38 ++++++++++++---
 tensor2tensor/utils/decoding.py          | 60 +++++++++++++++---------
 3 files changed, 79 insertions(+), 37 deletions(-)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 8372ed0a7..1656a5df5 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -87,8 +87,6 @@ def decode(estimator, hparams, decode_hp):
     decoding.decode_interactively(estimator, hparams, decode_hp,
                                   checkpoint_path=FLAGS.checkpoint_path)
   elif FLAGS.decode_from_file:
-    if estimator.config.use_tpu:
-      raise ValueError("TPU can only decode from dataset.")
     decoding.decode_from_file(estimator, FLAGS.decode_from_file, hparams,
                               decode_hp, FLAGS.decode_to_file,
                               checkpoint_path=FLAGS.checkpoint_path)
@@ -118,10 +116,10 @@ def score_file(filename):
     batch_inputs = tf.reshape(inputs_ph, [1, -1, 1, 1])  # Make it 4D.
   targets_ph = tf.placeholder(dtype=tf.int32)  # Just length dimension.
   batch_targets = tf.reshape(targets_ph, [1, -1, 1, 1])  # Make it 4D.
-  features = {
-      "inputs": batch_inputs,
-      "targets": batch_targets,
-  } if has_inputs else {"targets": batch_targets}
+  if has_inputs:
+    features = {"inputs": batch_inputs, "targets": batch_targets}
+  else:
+    features = {"targets": batch_targets}
 
   # Prepare the model and the graph when model runs on features.
   model = registry.model(FLAGS.model)(hparams, tf.estimator.ModeKeys.EVAL)
@@ -152,10 +150,10 @@ def score_file(filename):
       if has_inputs:
         inputs_numpy = encoders["inputs"].encode(inputs) + [text_encoder.EOS_ID]
       # Prepare the feed.
-      feed = {
-          inputs_ph: inputs_numpy,
-          targets_ph: targets_numpy
-      } if has_inputs else {targets_ph: targets_numpy}
+      if has_inputs:
+        feed = {inputs_ph: inputs_numpy, targets_ph: targets_numpy}
+      else:
+        feed = {targets_ph: targets_numpy}
       # Get the score.
       np_loss = sess.run(losses["training"], feed)
       results.append(np_loss)
diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index f819121c6..f6a173963 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -540,20 +540,27 @@ def mtr_lm_dense_3():
 
 
 @registry.register_hparams
-def mtr_lm_v1():
+def mtr_lm_v1(num_heads=8, num_memory_heads=0):
   """Model incorporating mixture-of-experts, local and global attention.
 
   ~6B parameters
 
   32 experts in 3 hierarchichal moe layers.
 
+  Args:
+    num_heads: an optional integer
+    num_memory_heads: an optional integer
+
   Returns:
     a hparams
   """
   hparams = mtr_lm_dense(0)
   local_att = transformer_layers.LocalSelfAttention(
-      num_heads=4, key_value_size=128)
-  att = transformer_layers.SelfAttention(num_heads=4, key_value_size=128)
+      num_heads=num_heads, num_memory_heads=num_memory_heads,
+      key_value_size=128)
+  att = transformer_layers.SelfAttention(
+      num_heads=num_heads, num_memory_heads=num_memory_heads,
+      key_value_size=128)
   drd = transformer_layers.DenseReluDense(hidden_size=2048)
   hmoe = moe.MoE2D(expert_x=8, expert_y=4, hidden_size=32768)
   hparams.layer_stack = transformer.LayerStack(
@@ -565,6 +572,12 @@ def mtr_lm_v1():
   return hparams
 
 
+@registry.register_hparams
+def mtr_lm_v1_h1_8():
+  """Version for fast decoding."""
+  return mtr_lm_v1(num_heads=8, num_memory_heads=1)
+
+
 def mtr_tr_dense(sz):
   """Series of machine translation models.
 
@@ -593,7 +606,7 @@ def mtr_tr_dense(sz):
   # one epoch for translate_enfr_wmt32k_packed = 51400 steps
   hparams.learning_rate_decay_steps = 51400
   hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
-  hparams.mesh_shape = "model:4;batch:8"
+  hparams.mesh_shape = "batch:32"
   hparams.label_smoothing = 0.1
   hparams.layer_prepostprocess_dropout = 0.1
   hparams.attention_dropout = 0.1
@@ -613,12 +626,16 @@ def mtr_tr_dense_1():
 
 @registry.register_hparams
 def mtr_tr_dense_2():
-  return mtr_tr_dense(2)
+  hparams = mtr_tr_dense(2)
+  hparams.mesh_shape = "model:4;batch:8"
+  return hparams
 
 
 @registry.register_hparams
 def mtr_tr_dense_3():
-  return mtr_tr_dense(3)
+  hparams = mtr_tr_dense(3)
+  hparams.mesh_shape = "model:4;batch:8"
+  return hparams
 
 
 @registry.register_hparams
@@ -628,6 +645,15 @@ def mtr_tr_dense_3_88():
   return hparams
 
 
+@registry.register_hparams
+def mtr_tr_dense_3_fast():
+  hparams = mtr_tr_dense(3)
+  hparams.decoder_local_attention_radius = 32
+  hparams.decoder_num_heads = 128
+  hparams.decoder_num_memory_heads = 8
+  return hparams
+
+
 def mtr_tr_dense_local(sz):
   """With local self-attention in the decoder."""
   hparams = mtr_tr_dense(sz)
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 464316900..ad33c7a05 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -153,7 +153,7 @@ def fix_and_save_video(vid, prefix):
     if targets is not None and log_results:
       decoded_targets = targets_vocab.decode(_save_until_eos(
           targets, skip_eos_postprocess))
-  if not is_video:
+  if log_results and not is_video:
     tf.logging.info("Inference results OUTPUT: %s" % decoded_outputs)
   if targets is not None and log_results and not is_video:
     tf.logging.info("Inference results TARGET: %s" % decoded_targets)
@@ -373,15 +373,33 @@ def decode_from_file(estimator,
   sorted_inputs, sorted_keys = _get_sorted_inputs(filename, decode_hp.delimiter)
   num_decode_batches = (len(sorted_inputs) - 1) // decode_hp.batch_size + 1
 
-  def input_fn():
-    input_gen = _decode_batch_input_fn(
-        num_decode_batches, sorted_inputs,
-        inputs_vocab, decode_hp.batch_size,
-        decode_hp.max_input_size, task_id=decode_hp.multiproblem_task_id)
-    gen_fn = make_input_fn_from_generator(input_gen)
-    example = gen_fn()
-    return _decode_input_tensor_to_features_dict(example, hparams)
-
+  if estimator.config.use_tpu:
+    length = getattr(hparams, "length", hparams.max_length)
+    batch_ids = []
+    for line in sorted_inputs:
+      ids = inputs_vocab.encode(line.strip()) + [1]
+      if len(ids) < length:
+        ids.extend([0] * (length - len(ids)))
+      else:
+        ids = ids[:length]
+      batch_ids.append(ids)
+    np_ids = np.array(batch_ids, dtype=np.int32)
+    def input_fn(params):
+      batch_size = params["batch_size"]
+      dataset = tf.data.Dataset.from_tensor_slices({"inputs": np_ids})
+      dataset = dataset.map(
+          lambda ex: {"inputs": tf.reshape(ex["inputs"], (length, 1, 1))})
+      dataset = dataset.batch(batch_size)
+      return dataset
+  else:
+    def input_fn():
+      input_gen = _decode_batch_input_fn(
+          num_decode_batches, sorted_inputs,
+          inputs_vocab, decode_hp.batch_size,
+          decode_hp.max_input_size, task_id=decode_hp.multiproblem_task_id)
+      gen_fn = make_input_fn_from_generator(input_gen)
+      example = gen_fn()
+      return _decode_input_tensor_to_features_dict(example, hparams)
   decodes = []
   result_iter = estimator.predict(input_fn, checkpoint_path=checkpoint_path)
 
@@ -448,10 +466,6 @@ def timer(gen):
                   (total_time_per_step / total_cnt,
                    total_time_per_step, total_cnt))
 
-  # Reversing the decoded inputs and outputs because they were reversed in
-  # _decode_batch_input_fn
-  sorted_inputs.reverse()
-  decodes.reverse()
   # If decode_to_file was provided use it as the output filename without change
   # (except for adding shard_id if using more shards for decoding).
   # Otherwise, use the input filename plus model, hp, problem, beam, alpha.
@@ -587,9 +601,6 @@ def _decode_batch_input_fn(num_decode_batches, sorted_inputs, vocabulary,
                            batch_size, max_input_size, task_id=-1):
   """Generator to produce batches of inputs."""
   tf.logging.info(" batch %d" % num_decode_batches)
-  # First reverse all the input sentences so that if you're going to get OOMs,
-  # you'll see it in the first batch
-  sorted_inputs.reverse()
   for b in range(num_decode_batches):
     tf.logging.info("Decoding batch %d" % b)
     batch_length = 0
@@ -604,6 +615,7 @@ def _decode_batch_input_fn(num_decode_batches, sorted_inputs, vocabulary,
       batch_inputs.append(input_ids)
       if len(input_ids) > batch_length:
         batch_length = len(input_ids)
+    batch_length = max_input_size
     final_batch_inputs = []
     for input_ids in batch_inputs:
       assert len(input_ids) <= batch_length
@@ -655,8 +667,8 @@ def _interactive_input_fn(hparams, decode_hp):
               "  dl=<decode_length>  (changes decode length, default: 100)\n"
               "  <%s>                (decode)\n"
               "  q                   (quit)\n"
-              ">" % (num_samples, decode_length, "source_string"
-                     if has_input else "target_prefix"))
+              ">" % (num_samples, decode_length,
+                     "source_string" if has_input else "target_prefix"))
     input_string = input(prompt)
     if input_string == "q":
       return
@@ -728,7 +740,13 @@ def show_and_save_image(img, save_path):
 
 
 def _get_sorted_inputs(filename, delimiter="\n"):
-  """Returning inputs sorted according to length.
+  """Returning inputs sorted according to decreasing length.
+
+  This causes inputs of similar lengths to be processed in the same batch,
+  facilitating early stopping for short sequences.
+
+  Longer sequences are sorted first so that if you're going to get OOMs,
+  you'll see it in the first batch.
 
   Args:
     filename: path to file with inputs, 1 per line.
@@ -746,7 +764,7 @@ def _get_sorted_inputs(filename, delimiter="\n"):
     # Strip the last empty line.
     if not inputs[-1]:
       inputs.pop()
-  input_lens = [(i, len(line.split())) for i, line in enumerate(inputs)]
+  input_lens = [(i, -len(line.split())) for i, line in enumerate(inputs)]
   sorted_input_lens = sorted(input_lens, key=operator.itemgetter(1))
   # We'll need the keys to rearrange the inputs back into their original order
   sorted_keys = {}

From d649d1f867bb4cafe9bad6b4b4c9ef0fb0f4f530 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Wed, 19 Dec 2018 00:15:25 +0100
Subject: [PATCH 1393/2720] Remove policy input standardization (#1312)

---
 tensor2tensor/models/research/rl.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 1ee226659..a32083191 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -447,7 +447,7 @@ class FeedForwardCategoricalPolicy(PolicyBase):
   """Feed-forward categorical."""
 
   def body(self, features):
-    observations = features["inputs"]
+    observations = features["inputs_raw"]
     flat_observations = tf.layers.flatten(observations)
     with tf.variable_scope("policy"):
       x = flat_observations
@@ -471,7 +471,7 @@ class FeedForwardCnnSmallCategoricalPolicy(PolicyBase):
   """Small cnn network with categorical output."""
 
   def body(self, features):
-    observations = features["inputs"]
+    observations = features["inputs_raw"]
     x = tf.transpose(observations, [0, 2, 3, 1, 4])
     x_shape = common_layers.shape_list(x)
     x = tf.reshape(x, x_shape[:-2] + [-1])

From e4d40df04cda1d96b61e1f5459ff9b3f08992394 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 18 Dec 2018 15:15:52 -0800
Subject: [PATCH 1394/2720] internal merge of PR #1312

PiperOrigin-RevId: 226068805
---
 tensor2tensor/models/research/rl.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index a32083191..c959aa712 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -484,7 +484,7 @@ def body(self, features):
                                    activation_fn=tf.nn.relu, padding="SAME")
 
       flat_x = tf.layers.flatten(x)
-      flat_x = tf.nn.dropout(flat_x, rate=dropout)
+      flat_x = tf.layers.dropout(flat_x, rate=dropout)
       x = tf.contrib.layers.fully_connected(flat_x, 128, tf.nn.relu)
 
       logits = tf.layers.dense(

From a65da23d2875fd910299f6157e478834f372f58b Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Tue, 18 Dec 2018 15:31:38 -0800
Subject: [PATCH 1395/2720] Default loss for multiproblems is the LM loss.

PiperOrigin-RevId: 226071083
---
 .../data_generators/multi_problem.py          | 37 ++++++++++++++++---
 tensor2tensor/layers/common_hparams.py        |  2 -
 tensor2tensor/models/transformer.py           |  1 -
 3 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 4fb5f6f71..b6e182566 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -372,6 +372,14 @@ def aggregate_task_losses(hparams,
                           target_modality,
                           feature):
   """Multiproblem loss function."""
+
+  # If no reweighting, we want the default loss to mimic the LM loss.
+  if not hparams.multiproblem_reweight_label_loss:
+    return aggregate_task_lm_losses(hparams=hparams,
+                                    logits=logits,
+                                    target_modality=target_modality,
+                                    feature=feature)
+
   summaries = []
   main_task_id = hparams.problem.task_list[0].task_id
   # Primary task loss
@@ -418,11 +426,7 @@ def aggregate_task_losses(hparams,
         label_loss *= hparams.multiproblem_label_weight
         seq_loss *= (1 - hparams.multiproblem_label_weight)
 
-      if hparams.multiproblem_class_loss_multiplier:
-        label_loss *= hparams.multiproblem_class_loss_multiplier
-        summaries.append([task.name+"_scaled_label_loss", label_loss])
-
-      # This is the training loss for the optimizer after all the scaling.
+      # This is the training loss for the optimizer after scaling.
       task_loss_val = seq_loss + label_loss
 
       loss_den_ = label_loss_den
@@ -458,3 +462,26 @@ def aggregate_task_losses(hparams,
                            loss_den_)
 
   return loss_num, loss_den, summaries
+
+
+def aggregate_task_lm_losses(hparams,
+                             logits,
+                             target_modality,
+                             feature):
+  """LM loss for multiproblems."""
+  summaries = []
+  loss_num = 0.
+  loss_den = 0.
+  for task in hparams.problem.task_list:
+    loss_num_, loss_den_ = target_modality.loss(
+        logits, feature,
+        weights_fn=
+        lambda x: common_layers.weights_multi_problem_all(x, task.task_id))  # pylint: disable=cell-var-from-loop
+
+    loss_num += loss_num_
+    loss_den += loss_den_
+
+    loss_val = loss_num_ / tf.maximum(1.0, loss_den_)
+    summaries.append([task.name+"_loss", loss_val])
+
+  return loss_num, loss_den, summaries
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 8987a4ae9..5b8f5e69e 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -271,8 +271,6 @@ def basic_params1():
       # A list of supported schedules can be found in
       # `data_generators.multi_problem.py`.
       multiproblem_mixing_schedule="constant",
-      # A scalar to upweight the classifier loss in a multiproblem setting.
-      multiproblem_class_loss_multiplier=0.0,
       # A boolean that decides whether input sequence losses and target label
       # losses in classification problems should be reweighted.
       multiproblem_reweight_label_loss=False,
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 7d9531b86..0bbad50ce 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1556,7 +1556,6 @@ def transformer_tall_finetune_textclass():
       "linear_warmup*constant*linear_decay")
   hparams.multiproblem_schedule_max_examples = 0
   hparams.multiproblem_target_eval_only = True
-  hparams.multiproblem_class_loss_multiplier = 4
   hparams.learning_rate_warmup_steps = 50
   # Set train steps to learning_rate_decay_steps or less
   hparams.learning_rate_decay_steps = 25000

From 136297c744f6cb19c74691aebdd544567720a0f2 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Tue, 18 Dec 2018 19:10:28 -0800
Subject: [PATCH 1396/2720] Subsampled summarization data.

PiperOrigin-RevId: 226101272
---
 .../data_generators/cnn_dailymail.py          | 70 +++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index 9380181e6..8da272526 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -22,6 +22,7 @@
 import hashlib
 import io
 import os
+import random
 import tarfile
 import six
 from tensor2tensor.data_generators import generator_utils
@@ -273,3 +274,72 @@ class SummarizeCnnDailymailWikiLMMultiVocab64k(SummarizeCnnDailymail32k):
   @property
   def vocab_filename(self):
     return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
+
+
+@registry.register_problem
+class SummarizeFracCnnDailymailWikiLMSharedVocab64k(SummarizeCnnDailymail32k):
+  """Summarize a fraction of CNN/DM articles using the Wiki 64k vocab."""
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelEnWiki64k().vocab_filename
+
+  def fraction_of_data(self):
+    return 1.
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    del data_dir
+    all_data = []
+    all_files, urls_path = _maybe_download_corpora(tmp_dir, dataset_split)
+    write_raw_text_to_files(all_files, urls_path, dataset_split, tmp_dir)
+    for example in example_generator(all_files, urls_path, sum_token=True):
+      story, summary = _story_summary_split(example)
+      all_data.append((story, summary))
+
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      random.shuffle(all_data)
+      fractional_len = int(self.fraction_of_data() * len(all_data))
+      all_data = all_data[:fractional_len]
+
+    for story, summary in all_data:
+      yield {"inputs": story, "targets": summary}
+
+
+@registry.register_problem
+class SummarizeFrac0p1CnnDailymailWikiLMSharedVocab64k(
+    SummarizeFracCnnDailymailWikiLMSharedVocab64k):
+
+  def fraction_of_data(self):
+    return 0.001
+
+
+@registry.register_problem
+class SummarizeFrac1CnnDailymailWikiLMSharedVocab64k(
+    SummarizeFracCnnDailymailWikiLMSharedVocab64k):
+
+  def fraction_of_data(self):
+    return 0.01
+
+
+@registry.register_problem
+class SummarizeFrac5CnnDailymailWikiLMSharedVocab64k(
+    SummarizeFracCnnDailymailWikiLMSharedVocab64k):
+
+  def fraction_of_data(self):
+    return 0.05
+
+
+@registry.register_problem
+class SummarizeFrac20CnnDailymailWikiLMSharedVocab64k(
+    SummarizeFracCnnDailymailWikiLMSharedVocab64k):
+
+  def fraction_of_data(self):
+    return 0.2
+
+
+@registry.register_problem
+class SummarizeFrac50CnnDailymailWikiLMSharedVocab64k(
+    SummarizeFracCnnDailymailWikiLMSharedVocab64k):
+
+  def fraction_of_data(self):
+    return 0.5

From 0625a4c079b9f1e00ef075402c2789016668e54e Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Tue, 18 Dec 2018 21:06:22 -0800
Subject: [PATCH 1397/2720] Fixes decoding on GPU, currently broken with
 default hparams due to assert on line 606.

PiperOrigin-RevId: 226111073
---
 tensor2tensor/utils/decoding.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index ad33c7a05..bf0dd131d 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -615,7 +615,8 @@ def _decode_batch_input_fn(num_decode_batches, sorted_inputs, vocabulary,
       batch_inputs.append(input_ids)
       if len(input_ids) > batch_length:
         batch_length = len(input_ids)
-    batch_length = max_input_size
+    if max_input_size != -1:
+      batch_length = max_input_size
     final_batch_inputs = []
     for input_ids in batch_inputs:
       assert len(input_ids) <= batch_length

From 629feb9755616fd98f28a036041f4b3d4eb6f084 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 19 Dec 2018 12:27:48 -0800
Subject: [PATCH 1398/2720] Add video-augmentation problem to
 t2t-data_generators.

PiperOrigin-RevId: 226215177
---
 tensor2tensor/data_generators/video_utils.py  | 62 +++++++++++++++++++
 .../data_generators/video_utils_test.py       | 13 ++++
 2 files changed, 75 insertions(+)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 685a8adf9..f242b2dec 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
 import os
 import numpy as np
 import six
@@ -27,6 +28,7 @@
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
@@ -52,6 +54,35 @@ def resize_video_frames(images, size):
   return resized_images
 
 
+def video_augmentation(features, hue=False, saturate=False, contrast=False):
+  """Augments video with optional hue, saturation and constrast.
+
+  Args:
+    features: dict, with keys "inputs", "targets".
+              features["inputs"], 4-D Tensor, shape=(THWC)
+              features["targets"], 4-D Tensor, shape=(THWC)
+    hue: bool, apply hue_transform.
+    saturate: bool, apply saturation transform.
+    contrast: bool, apply constrast transform.
+  Returns:
+    augment_features: dict with transformed "inputs" and "targets".
+  """
+  inputs, targets = features["inputs"], features["targets"]
+  in_steps = common_layers.shape_list(inputs)[0]
+
+  # makes sure that the same augmentation is applied to both input and targets.
+  # if input is 4-D, then tf.image applies the same transform across the batch.
+  video = tf.concat((inputs, targets), axis=0)
+  if hue:
+    video = tf.image.random_hue(video, max_delta=0.2)
+  if saturate:
+    video = tf.image.random_saturation(video, lower=0.5, upper=1.5)
+  if contrast:
+    video = tf.image.random_contrast(video, lower=0.5, upper=1.5)
+  features["inputs"], features["targets"] = video[:in_steps], video[in_steps:]
+  return features
+
+
 def create_border(video, color="blue", border_percent=2):
   """Creates a border around each frame to differentiate input and target.
 
@@ -637,6 +668,37 @@ def eval_metrics(self):
     return eval_metrics
 
 
+class VideoAugmentationProblem(VideoProblem):
+  """Base class for video data-augmentation.
+
+  By default applies a random hue, contrast and saturation transformation
+  to every video. To disable any of these transformations, inherit
+  this class and set the corresponding property to False.
+  """
+
+  @property
+  def hue(self):
+    return True
+
+  @property
+  def contrast(self):
+    return True
+
+  @property
+  def saturate(self):
+    return True
+
+  def preprocess(self, dataset, mode, hparams, interleave=True):
+    dataset = super(VideoAugmentationProblem, self).preprocess(
+        dataset=dataset, mode=mode, hparams=hparams, interleave=interleave)
+    video_augment_func = functools.partial(
+        video_augmentation, hue=self.hue, contrast=self.contrast,
+        saturate=self.saturate)
+    if mode == tf.estimator.ModeKeys.TRAIN:
+      dataset = dataset.map(video_augment_func)
+    return dataset
+
+
 class Video2ClassProblem(VideoProblemOld):
   """Base class for image classification problems."""
 
diff --git a/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index 39d626e46..091b91026 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -45,6 +45,19 @@ def get_predictions(self):
     problem = registry.problem("video_stochastic_shapes10k")
     return predictions, problem
 
+  def testVideoAugmentation(self):
+    # smoke-test, test for shapes.
+    with tf.Graph().as_default():
+      inputs = tf.random_uniform(shape=(3, 64, 64, 3))
+      targets = tf.random_uniform(shape=(10, 64, 64, 3))
+      features = {"inputs": inputs, "targets": targets}
+      augment = video_utils.video_augmentation(
+          features, hue=True, saturate=True, contrast=True)
+      with tf.Session() as sess:
+        augment_dict = sess.run(augment)
+        self.assertEqual(augment_dict["inputs"].shape, (3, 64, 64, 3))
+        self.assertEqual(augment_dict["targets"].shape, (10, 64, 64, 3))
+
   def testDecodeInMemoryTrue(self):
     predictions, problem = self.get_predictions()
     decode_hparams = decoding.decode_hparams()

From f260150491c64c6c0ce31817e8f19538277be779 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Wed, 19 Dec 2018 14:09:51 -0800
Subject: [PATCH 1399/2720] The weights masking may have a bug (reducing over
 axis=1 which is sequence length).

PiperOrigin-RevId: 226231118
---
 tensor2tensor/layers/common_layers_test.py | 34 ++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index a15f37302..7d7c95c22 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -489,6 +489,40 @@ def testDmlLoss(self, batch, height, width, num_mixtures, reduce_sum):
         [actual_loss, expected_loss])
     self.assertAllClose(actual_loss_val, expected_loss_val)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testWeightsMultiProblemAll(self):
+    labels = tf.constant(np.array([[12, 15, 1, 20, 100],
+                                   [67, 1, 34, 45, 124],
+                                   [78, 2, 34, 18, 29],
+                                   [78, 123, 55, 1, 33],
+                                   [1, 18, 22, 36, 59]]), dtype=tf.int32)
+    taskid = 1
+    expected_mask = np.array([[1, 1, 1, 1, 1],
+                              [1, 1, 1, 1, 1],
+                              [0, 0, 0, 0, 0],
+                              [1, 1, 1, 1, 1],
+                              [1, 1, 1, 1, 1]])
+    actual_mask = common_layers.weights_multi_problem_all(labels, taskid)
+    actual_mask_eval = self.evaluate(actual_mask)
+    self.assertAllClose(expected_mask, actual_mask_eval)
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testWeightsMultiProblem(self):
+    labels = tf.constant(np.array([[12, 15, 1, 20, 100],
+                                   [67, 1, 34, 45, 124],
+                                   [78, 2, 34, 18, 29],
+                                   [78, 123, 55, 1, 33],
+                                   [1, 18, 22, 36, 59]]), dtype=tf.int32)
+    taskid = 1
+    expected_mask = np.array([[0, 0, 0, 1, 1],
+                              [0, 0, 1, 1, 1],
+                              [0, 0, 0, 0, 0],
+                              [0, 0, 0, 0, 1],
+                              [0, 1, 1, 1, 1]])
+    actual_mask = common_layers.weights_multi_problem(labels, taskid)
+    actual_mask_eval = self.evaluate(actual_mask)
+    self.assertAllClose(expected_mask, actual_mask_eval)
+
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testDiscretizedMixLogisticLoss(self):
     batch = 2

From 54cbde34fcd1fe5884e5d0f8e31c33b8c3d759bf Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 20 Dec 2018 17:28:48 +0100
Subject: [PATCH 1400/2720] Set is_recurrent_model=True for stochastic discrete
 (#1315)

---
 tensor2tensor/models/video/basic_stochastic.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index aac145600..9f6d50114 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -66,6 +66,10 @@ class NextFrameBasicStochasticDiscrete(
     basic_deterministic.NextFrameBasicDeterministic):
   """Basic next-frame model with a tiny discrete latent."""
 
+  @property
+  def is_recurrent_model(self):
+    return True
+
   def init_internal_states(self):
     if not self.hparams.concat_internal_states:
       return None

From 9452363c6b13400bb45eb8a162f2dc78bd7d2ea7 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 20 Dec 2018 17:29:12 +0100
Subject: [PATCH 1401/2720] Fix simulated env resets (#1314)

---
 tensor2tensor/rl/envs/tf_atari_wrappers.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index de7929489..d72ae756b 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -127,7 +127,8 @@ def _reset_non_empty(self, indices):
     if initial_frames is None:
       inx = [1, self.history, 1, 1, 1]
       initial_frames = tf.tile(tf.expand_dims(new_values, axis=1), inx)
-    assign_op = tf.scatter_update(self._observ, indices, initial_frames)
+    with tf.control_dependencies([new_values]):
+      assign_op = tf.scatter_update(self._observ, indices, initial_frames)
     with tf.control_dependencies([assign_op]):
       return tf.gather(self.observ, indices)
 

From 1de13dbebccb415d89b0658e18a57e9607bafd32 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 20 Dec 2018 08:29:47 -0800
Subject: [PATCH 1402/2720] Allow per-task mixing in multi-problem and add
 Squad to the big MultiProblem; file move in rl.

PiperOrigin-RevId: 226339747
---
 .../data_generators/multi_problem.py          |  23 ++++
 tensor2tensor/data_generators/squad.py        | 120 +++++++++++-------
 .../data_generators/wiki_multi_problems.py    |   2 +
 tensor2tensor/layers/common_hparams.py        |   8 ++
 tensor2tensor/models/transformer.py           |   3 +-
 ...odel_rl_experiment_player.py => player.py} |   0
 .../rl/trainer_model_based_params.py          |   6 +-
 7 files changed, 114 insertions(+), 48 deletions(-)
 rename tensor2tensor/rl/{model_rl_experiment_player.py => player.py} (100%)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index b6e182566..4530aba05 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -275,6 +275,21 @@ def mix_data(example):
         tf.logging.info("Schedule mixing threshold "
                         "%.2f" % hparams.multiproblem_schedule_threshold)
 
+        # If per-task thresholds are specified, use them.
+        thresholds = None
+        if hparams.multiproblem_per_task_threshold:
+          thresholds = hparams.multiproblem_per_task_threshold.split(",")
+          thresholds = [float(t) for t in thresholds]  # Convert to floats.
+          thresholds_sum = sum(thresholds)
+          tf.logging.info("Per-task thresholds: %s." % str(thresholds))
+          thresholds = [t / thresholds_sum for t in thresholds]  # Normalize.
+          thresholds = [sum(thresholds[:i+1]) for i in range(len(thresholds))]
+          tf.logging.info("Per-task threshold sums: %s." % str(thresholds))
+          if len(thresholds) != len(self.task_list):
+            tf.logging.warn("Specified %d thresholds but encountered %d tasks."
+                            % (len(thresholds), len(self.task_list)))
+            thresholds = None
+
         def sample_task(curr_task, num_tasks_left, randnum):
           """A recursive function to sample a task.
 
@@ -294,6 +309,14 @@ def sample_task(curr_task, num_tasks_left, randnum):
           if num_tasks_left == 0:
             return get_next_from_dataset(dataset_iterators[curr_task])
 
+          if thresholds is not None:  # Use per-task thresholds if specified.
+            prob_sum = thresholds[curr_task]
+            return tf.cond(
+                randnum < prob_sum,
+                lambda: get_next_from_dataset(dataset_iterators[curr_task]),
+                lambda: sample_task(curr_task+1, num_tasks_left-1, randnum)
+            )
+
           # When curr_task is 0, the primary task, the new prob is the same as
           # the original probability. `tf.greater` indicates that the primary
           # task receives (1-prob) of the probability mass.
diff --git a/tensor2tensor/data_generators/squad.py b/tensor2tensor/data_generators/squad.py
index 4af5f0d6c..bf9646206 100644
--- a/tensor2tensor/data_generators/squad.py
+++ b/tensor2tensor/data_generators/squad.py
@@ -25,6 +25,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -34,18 +35,18 @@
 class Squad(text_problems.QuestionAndContext2TextProblem):
   """Base class for SquAD question answering problem."""
 
-  _DEV_SET = 'dev-v1.1.json'
-  _URL = 'https://rajpurkar.github.io/SQuAD-explorer/dataset'
-  _TRAINING_SET = 'train-v1.1.json'
+  _DEV_SET = "dev-v1.1.json"
+  _URL = "https://rajpurkar.github.io/SQuAD-explorer/dataset"
+  _TRAINING_SET = "train-v1.1.json"
 
   @property
   def dataset_splits(self):
     return [{
-        'split': problem.DatasetSplit.TRAIN,
-        'shards': 10,
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 10,
     }, {
-        'split': problem.DatasetSplit.EVAL,
-        'shards': 1,
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
     }]
 
   @property
@@ -54,47 +55,48 @@ def is_generate_per_split(self):
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
     url = self._URL
-    file_name = (self._TRAINING_SET if dataset_split ==
-                 problem.DatasetSplit.TRAIN else self._DEV_SET)
+    file_name = self._DEV_SET
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      file_name = self._TRAINING_SET
     squad_file = generator_utils.maybe_download(tmp_dir,
                                                 file_name,
                                                 os.path.join(url, file_name))
-    with tf.gfile.GFile(squad_file, mode='r') as fp:
+    with tf.gfile.GFile(squad_file, mode="r") as fp:
       squad = json.load(fp)
 
-    version = squad['version']
-    for article in squad['data']:
-      if 'title' in article:
-        title = article['title'].strip()
+    version = squad["version"]
+    for article in squad["data"]:
+      if "title" in article:
+        title = article["title"].strip()
       else:
-        title = 'no title'
-      for paragraph in article['paragraphs']:
-        context = paragraph['context'].strip()
-        for qa in paragraph['qas']:
-          question = qa['question'].strip()
-          id_ = qa['id']
+        title = "no title"
+      for paragraph in article["paragraphs"]:
+        context = paragraph["context"].strip()
+        for qa in paragraph["qas"]:
+          question = qa["question"].strip()
+          id_ = qa["id"]
 
-          answer_starts = [answer['answer_start'] for answer in qa['answers']]
-          answers = [answer['text'].strip() for answer in qa['answers']]
+          answer_starts = [answer["answer_start"] for answer in qa["answers"]]
+          answers = [answer["text"].strip() for answer in qa["answers"]]
 
-          # Features currently used are 'context', 'question', and 'answers'.
+          # Features currently used are "context", "question", and "answers".
           # Others are extracted here for the ease of future expansions.
           example = {
-              'version': version,
-              'title': title,
-              'context': context,
-              'question': question,
-              'id': id_,
-              'answer_starts': answer_starts,
-              'answers': answers,
-              'num_answers': len(answers),
-              'is_supervised': True,
+              "version": version,
+              "title": title,
+              "context": context,
+              "question": question,
+              "id": id_,
+              "answer_starts": answer_starts,
+              "answers": answers,
+              "num_answers": len(answers),
+              "is_supervised": True,
           }
           yield {
-              'inputs': example['question'],
+              "inputs": example["question"],
               # TODO(ddohan, wgaj): Figure out a way of extracting all answers.
-              'targets': example['answers'][0],
-              'context': example['context']
+              "targets": example["answers"][0],
+              "context": example["context"]
           }
 
 
@@ -103,21 +105,51 @@ class SquadConcat(Squad):
   """Squad with question and context concatenated together in inputs."""
 
   def dataset_filename(self):
-    return 'squad'
+    return "squad"
 
   def preprocess_example(self, example, unused_mode, unused_model_hparams):
     sep = tf.convert_to_tensor([self.QUESTION_SEPARATOR_ID],
-                               dtype=example['inputs'].dtype)
-    example['inputs'] = tf.concat(
-        [example['inputs'], sep, example['context']], 0)
+                               dtype=example["inputs"].dtype)
+    example["inputs"] = tf.concat(
+        [example["inputs"], sep, example["context"]], 0)
     return example
 
   def hparams(self, defaults, unused_model_hparams):
     (super(SquadConcat, self)
      .hparams(defaults, unused_model_hparams))
     p = defaults
-    del p.modality['context']
-    del p.vocab_size['context']
+    del p.modality["context"]
+    del p.vocab_size["context"]
+
+
+@registry.register_problem
+class SquadConcatMulti64k(SquadConcat):
+  """Squad with question and context concatenated, multi-lingual vocabulary."""
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 100,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  def preprocess_example(self, example, unused_mode, unused_model_hparams):
+    sep = tf.convert_to_tensor([self.QUESTION_SEPARATOR_ID],
+                               dtype=example["inputs"].dtype)
+    example["inputs"] = tf.concat(
+        [example["inputs"], sep, example["context"]], 0)
+    example.pop("context")
+    return example
+
+  def dataset_filename(self):
+    return "squad_multi64k"
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
 
 
 @registry.register_problem
@@ -140,7 +172,7 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
     samples = (super(SquadConcatPositioned, self)
                .generate_encoded_samples(data_dir, tmp_dir, dataset_split))
     for sample in samples:
-      sample['targets'] = self.generate_targets(sample['targets'],
-                                                sample['context'])
-      if sample['targets']:
+      sample["targets"] = self.generate_targets(sample["targets"],
+                                                sample["context"])
+      if sample["targets"]:
         yield sample
diff --git a/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index 7e4c7efe3..c25b580fe 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -22,6 +22,7 @@
 from tensor2tensor.data_generators import cnn_dailymail
 from tensor2tensor.data_generators import multi_problem
 from tensor2tensor.data_generators import multinli
+from tensor2tensor.data_generators import squad
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.data_generators import translate_ende
 from tensor2tensor.data_generators import translate_enfr
@@ -143,6 +144,7 @@ def __init__(self, was_reversed=False, was_copy=False):
     self.task_list.append(
         cnn_dailymail.SummarizeCnnDailymailWikiLMMultiVocab64k())
     self.task_list.append(multinli.MultiNLIWikiLMMultiVocab64k())
+    self.task_list.append(squad.SquadConcatMulti64k())
 
   @property
   def vocab_type(self):
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 5b8f5e69e..096351c4f 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -262,6 +262,14 @@ def basic_params1():
       # mixing should stop (eg: 0.5 means stop at 50-50 mixing, 0.8 means stop
       # at 20-80 mixing for the primary-others mixing case.)
       multiproblem_schedule_threshold=0.5,
+      # For more than 2 tasks, we may want to specify per-task thresholds here.
+      # In that case, this needs to be a string with as many floating point
+      # numbers as the number of tasks in the multi-problem. These numbers
+      # are later normalized to add up to 1 and taken as probabilities for
+      # each task. This enforces a constant mixing schedule and if this is
+      # empty then the threshold from above is used for the first task and
+      # the other tasks get the remaining probability split uniformly.
+      multiproblem_per_task_threshold="",
       # The number of examples at which the proportion of the mixed in datasets
       # is multiproblem_schedule_threshold
       multiproblem_schedule_max_examples=1e7,
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 0bbad50ce..ccec21624 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1604,7 +1604,8 @@ def transformer_tall_pretrain_lm_tpu_adafactor_large():
   hparams.filter_size = 32768
   hparams.batch_size = 4
   hparams.multiproblem_mixing_schedule = "constant"
-  hparams.multiproblem_schedule_threshold = 0.3
+  # Task order: LM, en-de/fr/ro/de-en/fr-en/ro-en/cnndm/mnli/squad.
+  hparams.multiproblem_per_task_threshold = "16,4,8,1,4,8,1,2,1,2"
   return hparams
 
 
diff --git a/tensor2tensor/rl/model_rl_experiment_player.py b/tensor2tensor/rl/player.py
similarity index 100%
rename from tensor2tensor/rl/model_rl_experiment_player.py
rename to tensor2tensor/rl/player.py
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 9a65a0342..610d979a8 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -82,7 +82,7 @@ def _rlmb_base():
       real_sampling_temp=1.0,
 
       # Sampling temperatures to try during eval.
-      eval_sampling_temps=[0.0, 0.2, 0.5, 0.8, 1.0, 2.0],
+      eval_sampling_temps=[0.5, 0.0, 1.0],
       eval_max_num_noops=8,
 
       game="pong",
@@ -94,7 +94,7 @@ def _rlmb_base():
       # Number of batches to run for world model evaluation.
       wm_eval_num_batches=8,
       # Ratios of ppo_epoch_length to report reward_accuracy on.
-      wm_eval_rollout_ratios=[0.25, 0.5, 1, 2],
+      wm_eval_rollout_ratios=[0.25, 1],
       stop_loop_early=False,  # To speed-up tests.
       rl_env_max_episode_steps=-1,  # Use default from gym.make()
       # Number of last observations to feed to the agent and world model.
@@ -135,7 +135,7 @@ def rlmb_ppo_base():
       real_batch_size=1,
       # Number of simulated environments to train on simultaneously.
       simulated_batch_size=16,
-      eval_batch_size=30,
+      eval_batch_size=64,
       wm_policy_param_sharing=False,
 
       # Unused; number of PPO epochs is calculated from the real frame limit.

From 3c1e2bc85e4e6c3ba393f1570521d191ade4da92 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Thu, 20 Dec 2018 08:57:25 -0800
Subject: [PATCH 1403/2720] Ability to make encoder unidirectional when using
 pre-trained

PiperOrigin-RevId: 226343141
---
 tensor2tensor/data_generators/problem.py   |  5 ++++
 tensor2tensor/layers/common_hparams.py     |  6 +++-
 tensor2tensor/layers/transformer_layers.py | 34 +++++++++++++++++-----
 tensor2tensor/models/transformer.py        | 11 ++++++-
 tensor2tensor/utils/t2t_model.py           |  7 +++++
 5 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 8f279f1ec..03ec21ea1 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -1163,6 +1163,11 @@ def _create_modalities(problem_hparams, hparams):
   modality = {}
   for feature_name, modality_cls in six.iteritems(problem_hparams.modality):
     vocab_size = problem_hparams.vocab_size[feature_name]
+    # If needed for using a pre-trained model's vocabulary where extra indices
+    # were allocated for adding new tasks with unique task ids.
+    if (hasattr(hparams, "multiproblem_vocab_size") and
+        hparams.multiproblem_vocab_size > 0):
+      vocab_size = hparams.multiproblem_vocab_size
     modality_cls = modality_overrides.get(feature_name, modality_cls)
     modality[feature_name] = modality_cls(hparams, vocab_size)
   problem_hparams.modality = modality
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 096351c4f..871b61587 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -310,7 +310,11 @@ def basic_params1():
       multiproblem_max_input_length=-1,
       multiproblem_max_target_length=-1,
       # If positive, makes training targets fixed-length in MultiProblem.
-      multiproblem_fixed_train_length=-1
+      multiproblem_fixed_train_length=-1,
+      # Load weights from a second model. For instance, when using
+      # pre-trained weights, you might want to initialize the encoder
+      # and decoder by loading different models.
+      warm_start_from_second=""
   )
 
 
diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 84cc5c557..8da58792a 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -50,17 +50,32 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
     inputs_segmentation = features["inputs_segmentation"]
     inputs_position = features["inputs_position"]
     targets_segmentation = features["targets_segmentation"]
-    encoder_self_attention_bias = common_attention.attention_bias_same_segment(
-        inputs_segmentation, inputs_segmentation)
+    if (hasattr(hparams, "unidirectional_encoder") and
+        hparams.unidirectional_encoder):
+      tf.logging.info("Using unidirectional encoder")
+      encoder_self_attention_bias = (
+          common_attention.attention_bias_lower_triangle(
+              common_layers.shape_list(inputs)[1]))
+    else:
+      encoder_self_attention_bias = (
+          common_attention.attention_bias_same_segment(
+              inputs_segmentation, inputs_segmentation))
     encoder_decoder_attention_bias = (
         common_attention.attention_bias_same_segment(targets_segmentation,
                                                      inputs_segmentation))
   else:
-    # Usual case - not a packed dataset.
     encoder_padding = common_attention.embedding_to_padding(encoder_input)
     ignore_padding = common_attention.attention_bias_ignore_padding(
         encoder_padding)
-    encoder_self_attention_bias = ignore_padding
+    if (hasattr(hparams, "unidirectional_encoder") and
+        hparams.unidirectional_encoder):
+      tf.logging.info("Using unidirectional encoder")
+      encoder_self_attention_bias = (
+          common_attention.attention_bias_lower_triangle(
+              common_layers.shape_list(inputs)[1]))
+    else:
+      # Usual case - not a packed dataset.
+      encoder_self_attention_bias = ignore_padding
     encoder_decoder_attention_bias = ignore_padding
     inputs_position = None
   if hparams.proximity_bias:
@@ -103,7 +118,8 @@ def transformer_encoder(encoder_input,
                         nonpadding=None,
                         save_weights_to=None,
                         make_image_summary=True,
-                        losses=None):
+                        losses=None,
+                        attn_bias_for_padding=None):
   """A stack of transformer layers.
 
   Args:
@@ -123,6 +139,8 @@ def transformer_encoder(encoder_input,
       a string key created from the variable scope (including name).
     make_image_summary: Whether to make an attention image summary.
     losses: optional list onto which to append extra training losses
+    attn_bias_for_padding: Padded attention bias in case a unidirectional
+      encoder is being used where future attention is masked.
 
   Returns:
     y: a Tensors
@@ -149,8 +167,10 @@ def transformer_encoder(encoder_input,
     if nonpadding is not None:
       padding = 1.0 - nonpadding
     else:
-      padding = common_attention.attention_bias_to_padding(
-          encoder_self_attention_bias)
+      attention_bias = encoder_self_attention_bias
+      if attn_bias_for_padding is not None:
+        attention_bias = attn_bias_for_padding
+      padding = common_attention.attention_bias_to_padding(attention_bias)
       nonpadding = 1.0 - padding
     pad_remover = None
     if hparams.use_pad_remover and not common_layers.is_xla_compiled():
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index ccec21624..3ccbe7468 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -92,6 +92,11 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
     encoder_input = tf.nn.dropout(encoder_input,
                                   1.0 - hparams.layer_prepostprocess_dropout)
 
+    attn_bias_for_padding = None
+    # Otherwise the encoder will just use encoder_self_attention_bias.
+    if hparams.unidirectional_encoder:
+      attn_bias_for_padding = encoder_decoder_attention_bias
+
     encoder_output = transformer_encoder(
         encoder_input,
         self_attention_bias,
@@ -99,7 +104,8 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
         nonpadding=features_to_nonpadding(features, "inputs"),
         save_weights_to=self.attention_weights,
         make_image_summary=not common_layers.is_xla_compiled(),
-        losses=losses)
+        losses=losses,
+        attn_bias_for_padding=attn_bias_for_padding)
 
     return encoder_output, encoder_decoder_attention_bias
 
@@ -1375,6 +1381,9 @@ def transformer_base_v1():
   # This is useful for programs that can automatically compare experiments side
   #   by side based on the same metric names.
   hparams.add_hparam("overload_eval_metric_name", "")
+  # For making a transformer encoder unidirectional by using masked
+  # attention.
+  hparams.add_hparam("unidirectional_encoder", False)
   return hparams
 
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 9a0c374ec..393d66cb0 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1456,6 +1456,8 @@ def initialize_from_ckpt(self, ckpt_dir):
     if already_has_ckpt:
       return
 
+    log_info("Checkpoint dir: %s", ckpt_dir)
+
     # TODO(mitchellstern): Add support for partitioned variables?
     reader = tf.contrib.framework.load_checkpoint(ckpt_dir)
     variable_map = {}
@@ -1500,6 +1502,11 @@ def scaffold_fn():
       if self._hparams.warm_start_from:
         self.initialize_from_ckpt(self._hparams.warm_start_from)
 
+      # When loading weights from a pre-trained model, you want to be able to
+      # load separate weights into the encoder and decoder.
+      if self._hparams.warm_start_from_second:
+        self.initialize_from_ckpt(self._hparams.warm_start_from_second)
+
       return tf.estimator.EstimatorSpec(
           tf.estimator.ModeKeys.TRAIN, loss=loss, train_op=train_op)
 

From 0d72b649caaa525ada73f562b1d4e9ce589b0c1b Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Thu, 20 Dec 2018 09:25:33 -0800
Subject: [PATCH 1404/2720] Move out 512 videos from the train-set as holdout
 validation set in BAIR

PiperOrigin-RevId: 226347022
---
 .../data_generators/bair_robot_pushing.py      | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index 0a349d972..c13728207 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -67,9 +67,10 @@ def frame_width(self):
   def is_generate_per_split(self):
     return True
 
+  # num_train_files * num_videos * num_frames
   @property
   def total_number_of_frames(self):
-    return 1305600
+    return 167 * 256 * 30
 
   @property
   def random_skip(self):
@@ -144,12 +145,19 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
     tar.extractall(tmp_dir)
     tar.close()
 
-    if dataset_split == problem.DatasetSplit.TRAIN:
-      base_dir = os.path.join(tmp_dir, "softmotion30_44k/train/*")
-    else:
+    if dataset_split == problem.DatasetSplit.TEST:
       base_dir = os.path.join(tmp_dir, "softmotion30_44k/test/*")
+      filenames = tf.gfile.Glob(base_dir)
+    else:
+      base_dir = os.path.join(tmp_dir, "softmotion30_44k/train/*")
+      filenames = tf.gfile.Glob(base_dir)
+
+      # the test-set contains just 256 videos so this should be sufficient.
+      if dataset_split == problem.DatasetSplit.TRAIN:
+        filenames = filenames[:-2]
+      else:
+        filenames = filenames[-2:]
 
-    filenames = tf.gfile.Glob(base_dir)
     for frame_number, frame, state, action in self.parse_frames(filenames):
       yield {
           "frame_number": [frame_number],

From 0978e8e14ee612bb5a7a222e35894209890a36bc Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Thu, 20 Dec 2018 10:34:12 -0800
Subject: [PATCH 1405/2720] Minor, Allow dilation rates as hparams.

PiperOrigin-RevId: 226357299
---
 tensor2tensor/models/research/glow_ops.py     | 2 +-
 tensor2tensor/models/video/next_frame_glow.py | 1 +
 tensor2tensor/models/video/nfg_conv3d_test.py | 6 +++---
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 71412b5f7..a4957539e 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -679,7 +679,7 @@ def get_dilation_rates(hparams, width):
   # dil_rate=1 means no dilation.
   allowed_dilations = [[1]*5]
   apply_dilations = hparams.get("latent_apply_dilations", False)
-  dilation_rates = [1, 3]   # Number of holes between each filter element.
+  dilation_rates = hparams.get("latent_dilation_rates", [1, 3])
   if apply_dilations:
     for rate in dilation_rates:
       # k + (k - 1) * rate but k is harcoded to be 3 everywhere.
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index d6e0f3d2a..d2892d5be 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -55,6 +55,7 @@ def next_frame_glow_hparams():
   hparams.add_hparam("num_cond_latents", 1)
   hparams.add_hparam("latent_architecture", "glow_resnet")
   hparams.add_hparam("latent_apply_dilations", False)
+  hparams.add_hparam("latent_dilation_rates", [1, 3])
   # Use latent skip connections
   hparams.add_hparam("model_input", False)
   hparams.add_hparam("cond_first_frame", False)
diff --git a/tensor2tensor/models/video/nfg_conv3d_test.py b/tensor2tensor/models/video/nfg_conv3d_test.py
index 6bf5549b4..c539ac20f 100644
--- a/tensor2tensor/models/video/nfg_conv3d_test.py
+++ b/tensor2tensor/models/video/nfg_conv3d_test.py
@@ -24,10 +24,10 @@
 import tensorflow as tf
 
 conv3d_net_hparams = (
-    # ("conv3d_net", 2, 2, "conv3d_net", "conditional", -1, 3),
+    ("conv3d_net", 2, 2, "conv3d_net", "conditional", -1, 3),
     ("conv3d_net_gatu", 2, 2, "conv3d_net", "conditional", -1, 3, False, False,
-     "gatu"),)
-    # ("conv3d_dil", 2, 2, "conv3d_net", "conditional", -1, -1, False, True),)
+     "gatu"),
+    ("conv3d_dil", 2, 2, "conv3d_net", "conditional", -1, -1, False, True),)
 
 
 class NextFrameGlowConv3DTest(nfg_test_utils.NextFrameGlowTest,

From 48061e55e900f9291071fe8586b366715646f3bc Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Thu, 20 Dec 2018 11:51:38 -0800
Subject: [PATCH 1406/2720] Configs for summarization.

PiperOrigin-RevId: 226370199
---
 tensor2tensor/models/transformer.py | 95 +++++++++++++++++++++++------
 1 file changed, 75 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 3ccbe7468..1312a1a80 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1550,9 +1550,83 @@ def transformer_tall():
   hparams.num_hidden_layers = 12
   hparams.num_heads = 12
   hparams.label_smoothing = 0.0
-  hparams.max_length = 512
+  hparams.max_length = 1024
   hparams.eval_drop_long_sequences = True
   hparams.multiproblem_mixing_schedule = "pretrain"
+  hparams.multiproblem_vocab_size = 65536
+  hparams.clip_grad_norm = 1.0
+  return hparams
+
+
+@registry.register_hparams
+def transformer_tall_finetune_tied():
+  """Tied means fine-tune CNN/DM summarization as LM."""
+  hparams = transformer_tall()
+  hparams.multiproblem_max_input_length = 750
+  hparams.multiproblem_max_target_length = 100
+  hparams.multiproblem_schedule_max_examples = 0
+  hparams.learning_rate_schedule = (
+      "linear_warmup*constant*cosdecay")
+  hparams.learning_rate_constant = 5e-5
+  hparams.learning_rate_warmup_steps = 100
+  # Set train steps to learning_rate_decay_steps or less
+  hparams.learning_rate_decay_steps = 80000
+  hparams.multiproblem_target_eval_only = True
+  hparams.multiproblem_reweight_label_loss = True
+  hparams.multiproblem_label_weight = 1.0
+  hparams.optimizer = "TrueAdam"
+  return hparams
+
+
+@registry.register_hparams
+def transformer_tall_train_tied():
+  """Tied means train CNN/DM summarization as LM."""
+  hparams = transformer_tall()
+  hparams.multiproblem_max_input_length = 750
+  hparams.multiproblem_max_target_length = 100
+  hparams.multiproblem_schedule_max_examples = 0
+  hparams.learning_rate_schedule = (
+      "linear_warmup*constant*cosdecay")
+  hparams.learning_rate_constant = 2e-4
+  hparams.learning_rate_warmup_steps = 8000
+  # Set train steps to learning_rate_decay_steps or less
+  hparams.learning_rate_decay_steps = 150000
+  hparams.multiproblem_target_eval_only = True
+  hparams.multiproblem_reweight_label_loss = True
+  hparams.multiproblem_label_weight = 1.0
+  hparams.optimizer = "TrueAdam"
+  return hparams
+
+
+@registry.register_hparams
+def transformer_tall_finetune_uniencdec():
+  """Fine-tune CNN/DM with a unidirectional encoder and decoder."""
+  hparams = transformer_tall()
+  hparams.max_input_seq_length = 750
+  hparams.max_target_seq_length = 100
+  hparams.optimizer = "TrueAdam"
+  hparams.learning_rate_schedule = (
+      "linear_warmup*constant*cosdecay")
+  hparams.learning_rate_decay_steps = 80000
+  hparams.learning_rate_constant = 5e-5
+  hparams.learning_rate_warmup_steps = 100
+  hparams.unidirectional_encoder = True
+  hparams.load_encoder = False
+  return hparams
+
+
+@registry.register_hparams
+def transformer_tall_train_uniencdec():
+  """Train CNN/DM with a unidirectional encoder and decoder."""
+  hparams = transformer_tall()
+  hparams.max_input_seq_length = 750
+  hparams.max_target_seq_length = 100
+  hparams.optimizer = "TrueAdam"
+  hparams.learning_rate_schedule = (
+      "linear_warmup*constant*cosdecay")
+  hparams.learning_rate_decay_steps = 150000
+  hparams.learning_rate_constant = 2e-4
+  hparams.unidirectional_encoder = True
   return hparams
 
 
@@ -1629,25 +1703,6 @@ def transformer_tall_pretrain_lm_tpu():
   return hparams
 
 
-@registry.register_hparams
-def transformer_tall_finetune_cnndm():
-  """Hparams for transformer on LM for finetuning on cnndm summarization."""
-  hparams = transformer_tall()
-  hparams.batch_size = 4096
-  hparams.multiproblem_max_input_length = 412
-  hparams.multiproblem_max_target_length = 100
-  hparams.multiproblem_schedule_max_examples = 0
-  hparams.learning_rate_schedule = (
-      "linear_warmup*constant*cosdecay")
-  hparams.learning_rate_constant = 5e-5
-  hparams.learning_rate_warmup_steps = 100
-  # Set train steps to learning_rate_decay_steps or less
-  hparams.learning_rate_decay_steps = 40000
-  hparams.multiproblem_target_eval_only = True
-  hparams.multiproblem_vocab_size = 2**16
-  return hparams
-
-
 @registry.register_hparams
 def transformer_tall_big():
   """Hparams for transformer on LM+MNLI."""

From b485f4be62f4b6bc45d77aa90b22b82db8de6391 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Thu, 20 Dec 2018 12:55:58 -0800
Subject: [PATCH 1407/2720] Make attention variables 2d (combining heads and
 d_kv dimensions) in order to avoid reshaping the variables - faster
 training/infernce on TPU. Breaks previous mtf_transformer2 checkpoints.
 Optimized beam search code for TPU.

No one uses this code yet, so no worries about breaking anything!

PiperOrigin-RevId: 226379710
---
 tensor2tensor/models/mtf_transformer2.py | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index f6a173963..4ddac0837 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -121,7 +121,9 @@ def model(self):
         output_vocab_size=self._targets_vocab_size,
         autoregressive=self.autoregressive,
         max_length=hparams.max_length,
-        z_loss=hparams.z_loss)
+        z_loss=hparams.z_loss,
+        layout=hparams.layout,
+        mesh_shape=hparams.mesh_shape)
 
   def _mtf_model_fn(self, features, mesh):
     self._original_features = features
@@ -222,7 +224,9 @@ def model(self):
         max_length=hparams.max_length,
         shared_embedding=hparams.shared_embedding,
         label_smoothing=hparams.label_smoothing,
-        z_loss=hparams.z_loss)
+        z_loss=hparams.z_loss,
+        layout=hparams.layout,
+        mesh_shape=hparams.mesh_shape)
 
   def _mtf_model_fn(self, features, mesh):
     self._original_features = features
@@ -647,7 +651,7 @@ def mtr_tr_dense_3_88():
 
 @registry.register_hparams
 def mtr_tr_dense_3_fast():
-  hparams = mtr_tr_dense(3)
+  hparams = mtr_tr_dense_3()
   hparams.decoder_local_attention_radius = 32
   hparams.decoder_num_heads = 128
   hparams.decoder_num_memory_heads = 8
@@ -681,6 +685,15 @@ def mtr_tr_dense_local_0_h1_16():
   return hparams
 
 
+@registry.register_hparams
+def mtr_tr_dense_local_0_h1_8_kv256():
+  hparams = mtr_tr_dense_local_0()
+  hparams.decoder_num_heads = 8
+  hparams.decoder_num_memory_heads = 1
+  hparams.d_kv = 256
+  return hparams
+
+
 @registry.register_hparams
 def mtr_tr_dense_local_0_h1_16_shared_kv():
   hparams = mtr_tr_dense_local_0_h1_16()

From 65f66e70a1b145398e2df3c35f5af5fb135c7467 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Thu, 20 Dec 2018 15:57:21 -0800
Subject: [PATCH 1408/2720] Adds dropout to coupling-layers and distribution
 predictor via hparams.coupling_dropout and hparams.latent_dropout.

PiperOrigin-RevId: 226407884
---
 tensor2tensor/models/research/glow.py         |  1 +
 tensor2tensor/models/research/glow_ops.py     | 45 +++++++++++++------
 .../models/research/glow_ops_test.py          | 42 +++++++++++------
 tensor2tensor/models/video/next_frame_glow.py |  1 +
 4 files changed, 61 insertions(+), 28 deletions(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index ab008dad0..b0ac4e30e 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -54,6 +54,7 @@ def glow_hparams():
   # Coupling layer, additive or affine.
   hparams.add_hparam("coupling", "affine")
   hparams.add_hparam("coupling_width", 512)
+  hparams.add_hparam("coupling_dropout", 0.0)
   hparams.add_hparam("top_prior", "single_conv")
   # init_batch_size denotes the number of examples used for data-dependent
   # initialization. A higher init_batch_size is required for training
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index a4957539e..580a9bd34 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -453,7 +453,8 @@ def conv(name, x, output_channels, filter_size=None, stride=None,
 
 
 @add_arg_scope
-def conv_block(name, x, mid_channels, dilations=None, activation="relu"):
+def conv_block(name, x, mid_channels, dilations=None, activation="relu",
+               dropout=0.0):
   """2 layer conv block used in the affine coupling layer.
 
   Args:
@@ -464,6 +465,7 @@ def conv_block(name, x, mid_channels, dilations=None, activation="relu"):
     activation: relu or gatu.
       If relu, the second layer is relu(W*x)
       If gatu, the second layer is tanh(W1*x) * sigmoid(W2*x)
+    dropout: Dropout probability.
   Returns:
     x: 4-D Tensor: Output activations.
   """
@@ -483,6 +485,8 @@ def conv_block(name, x, mid_channels, dilations=None, activation="relu"):
     x = conv("1_1", x, output_channels=mid_channels, filter_size=first_filter,
              dilations=dilations)
     x = tf.nn.relu(x)
+    if dropout != 0.0:
+      x = tf.nn.dropout(x, keep_prob=1.0 - dropout)
 
     # Padding + conv2d + actnorm + activation.
     # [input, output: 512 channels]
@@ -497,11 +501,15 @@ def conv_block(name, x, mid_channels, dilations=None, activation="relu"):
       x_sigm = conv("1_sigm", x, output_channels=mid_channels,
                     filter_size=second_filter, dilations=dilations)
       x = tf.nn.tanh(x_tanh) * tf.nn.sigmoid(x_sigm)
+
+    if dropout != 0.0:
+      x = tf.nn.dropout(x, keep_prob=1.0 - dropout)
     return x
 
 
 def dilated_conv_stack(name, x, mid_channels, output_channels,
-                       dilation_rates, activation="relu"):
+                       dilation_rates, activation="relu",
+                       dropout=0.0):
   """Dilated convolutional stack.
 
   Features at different rates are computed independently using a 3 layer
@@ -515,6 +523,7 @@ def dilated_conv_stack(name, x, mid_channels, output_channels,
     output_channels: Number of output channels of the last layer.
     dilation_rates: A list of dilation rates.
     activation: Can be either "relu" or "gatu"
+    dropout: dropout.
   Returns:
     output: 5-D Tensor.
   """
@@ -524,14 +533,14 @@ def dilated_conv_stack(name, x, mid_channels, output_channels,
       # TODO(mechcoder) try (concat across channels + 1x1) modulo memory issues.
       curr_out = conv_stack("dil_%d" % dil_ind, x, mid_channels=mid_channels,
                             output_channels=output_channels, dilations=dil_rate,
-                            activation=activation)
+                            activation=activation, dropout=dropout)
       output += curr_out
     return output
 
 
 @add_arg_scope
 def conv_stack(name, x, mid_channels, output_channels, dilations=None,
-               activation="relu"):
+               activation="relu", dropout=0.0):
   """3-layer convolutional stack.
 
   Args:
@@ -544,13 +553,15 @@ def conv_stack(name, x, mid_channels, output_channels, dilations=None,
     activation: relu or gatu.
       If relu, the second layer is relu(W*x)
       If gatu, the second layer is tanh(W1*x) * sigmoid(W2*x)
+    dropout: float, 0.0
   Returns:
     output: output of 3 layer conv network.
   """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
 
     x = conv_block("conv_block", x, mid_channels=mid_channels,
-                   dilations=dilations, activation=activation)
+                   dilations=dilations, activation=activation,
+                   dropout=dropout)
 
     # Final layer.
     x = conv("zeros", x, apply_actnorm=False, conv_init="zeros",
@@ -560,7 +571,7 @@ def conv_stack(name, x, mid_channels, output_channels, dilations=None,
 
 @add_arg_scope
 def additive_coupling(name, x, mid_channels=512, reverse=False,
-                      activation="relu"):
+                      activation="relu", dropout=0.0):
   """Reversible additive coupling layer.
 
   Args:
@@ -569,6 +580,7 @@ def additive_coupling(name, x, mid_channels=512, reverse=False,
     mid_channels: number of channels in the coupling layer.
     reverse: Forward or reverse operation.
     activation: "relu" or "gatu"
+    dropout: default, 0.0
   Returns:
     output:
     objective: 0.0
@@ -579,7 +591,7 @@ def additive_coupling(name, x, mid_channels=512, reverse=False,
 
     z1 = x1
     shift = conv_stack("nn", x1, mid_channels, output_channels=output_channels,
-                       activation=activation)
+                       activation=activation, dropout=dropout)
 
     if not reverse:
       z2 = x2 + shift
@@ -590,7 +602,7 @@ def additive_coupling(name, x, mid_channels=512, reverse=False,
 
 @add_arg_scope
 def affine_coupling(name, x, mid_channels=512, activation="relu",
-                    reverse=False):
+                    reverse=False, dropout=0.0):
   """Reversible affine coupling layer.
 
   Args:
@@ -599,6 +611,7 @@ def affine_coupling(name, x, mid_channels=512, activation="relu",
     mid_channels: number of channels in the coupling layer.
     activation: Can be either "relu" or "gatu".
     reverse: Forward or reverse operation.
+    dropout: default, 0.0
   Returns:
     output: input s
     objective: log-determinant of the jacobian
@@ -614,7 +627,8 @@ def affine_coupling(name, x, mid_channels=512, activation="relu",
     # z2 = (x2 / scale) - shift
     z1 = x1
     log_scale_and_shift = conv_stack(
-        "nn", x1, mid_channels, x_shape[-1], activation=activation)
+        "nn", x1, mid_channels, x_shape[-1], activation=activation,
+        dropout=dropout)
     shift = log_scale_and_shift[:, :, :, 0::2]
     scale = tf.nn.sigmoid(log_scale_and_shift[:, :, :, 1::2] + 2.0)
     if not reverse:
@@ -714,12 +728,14 @@ def temporal_latent_to_dist(name, x, hparams, output_channels=None):
                                 mid_channels=hparams.latent_encoder_width,
                                 output_channels=res_channels,
                                 dilation_rates=dilation_rates,
-                                activation=hparams.latent_activation)
+                                activation=hparams.latent_activation,
+                                dropout=hparams.latent_dropout)
       else:
         h2 = conv_stack("latent_3d_res_%d" % i, h,
                         mid_channels=hparams.latent_encoder_width,
                         output_channels=res_channels,
-                        activation=hparams.latent_activation)
+                        activation=hparams.latent_activation,
+                        dropout=hparams.latent_dropout)
       h += h2
 
     # take last activation that should capture all context since padding is
@@ -798,7 +814,8 @@ def latent_to_dist(name, x, hparams, output_channels=None):
       h = x
       for layer in range(depth):
         h3 = conv_stack("latent_resnet_%d" % layer, h,
-                        mid_channels=width, output_channels=x_shape[-1])
+                        mid_channels=width, output_channels=x_shape[-1],
+                        dropout=hparams.coupling_dropout)
         h += h3
       mean_log_scale = conv("glow_res_final", h, conv_init="zeros",
                             output_channels=2*output_channels,
@@ -1030,12 +1047,12 @@ def revnet_step(name, x, hparams, reverse=True):
       coupling_layer = functools.partial(
           additive_coupling, name="additive", reverse=reverse,
           mid_channels=hparams.coupling_width,
-          activation=hparams.activation)
+          activation=hparams.activation, dropout=hparams.coupling_dropout)
     else:
       coupling_layer = functools.partial(
           affine_coupling, name="affine", reverse=reverse,
           mid_channels=hparams.coupling_width,
-          activation=hparams.activation)
+          activation=hparams.activation, dropout=hparams.coupling_dropout)
     ops = [
         functools.partial(actnorm, name="actnorm", reverse=reverse),
         functools.partial(invertible_1x1_conv, name="invertible",
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index c131d007e..07004294f 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -51,6 +51,7 @@ def get_glow_hparams(self):
     hparams.add_hparam("latent_dist_encoder", "conv_net")
     hparams.add_hparam("latent_time_filter_size", 3)
     hparams.add_hparam("latent_activation", "relu")
+    hparams.add_hparam("latent_dropout", 0.0)
     return hparams
 
   def test_get_variable_ddi(self):
@@ -74,24 +75,29 @@ def test_actnorm(self):
         self.assertTrue(np.allclose(channel_mean, 0.0, atol=1e-3))
         self.assertTrue(np.allclose(channel_var, 1.0, atol=1e-3))
 
-  def check_invertibility(self, op, name):
+  @parameterized.named_parameters(
+      ("inv_1x1", glow_ops.invertible_1x1_conv, "inv_1x1"),
+      ("affine", glow_ops.affine_coupling, "affine_coupling"),
+      ("additive", glow_ops.additive_coupling, "additive_coupling"),
+      ("actnorm", glow_ops.actnorm, "actnorm"),
+      ("affine_drop", glow_ops.affine_coupling, "affine_dropout", 0.5),
+      ("additive_drop", glow_ops.additive_coupling, "additive_dropout", 0.5))
+  def test_invertibility(self, op, name, dropout=0.0):
     with tf.Graph().as_default():
+      tf.set_random_seed(42)
       x = tf.random_uniform(shape=(16, 32, 32, 4))
 
-      x_inv, _ = op(name, x, reverse=False)
-      x_inv_inv, _ = op(name, x_inv, reverse=True)
+      if op in [glow_ops.affine_coupling, glow_ops.additive_coupling]:
+        x_inv, _ = op(name, x, reverse=False, dropout=dropout)
+        x_inv_inv, _ = op(name, x_inv, reverse=True, dropout=dropout)
+      else:
+        x_inv, _ = op(name, x, reverse=False)
+        x_inv_inv, _ = op(name, x_inv, reverse=True)
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())
         diff = session.run(x - x_inv_inv)
         self.assertTrue(np.allclose(diff, 0.0, atol=1e-5))
 
-  def test_invertibility(self):
-    rev_ops = [glow_ops.invertible_1x1_conv, glow_ops.affine_coupling,
-               glow_ops.actnorm, glow_ops.additive_coupling]
-    names = ["inv_1X1_conv", "affine_coupling", "actnorm", "additive_coupling"]
-    for rev_op, name in zip(rev_ops, names):
-      self.check_invertibility(rev_op, name)
-
   def test_add_edge_bias(self):
     with tf.Graph().as_default():
       x = tf.random_uniform(shape=(16, 32, 32, 3))
@@ -323,8 +329,11 @@ def test_split_latent_conditioning(self):
       ("conv_net_skip", "conv_net", True),
       ("conv_net_no_skip", "conv_net", False),
       ("conv3d_skip", "conv3d_net", False),
-      ("conv3d_no_skip", "conv3d_net", True))
-  def test_latent_dist_encoder(self, encoder="conv_lstm", skip=True):
+      ("conv3d_no_skip", "conv3d_net", True),
+      ("conv3d_skip_drop", "conv3d_net", False, 0.1),
+      ("conv3d_no_skip_drop", "conv3d_net", True, 0.1))
+  def test_latent_dist_encoder(self, encoder="conv_lstm", skip=True,
+                               dropout=0.0):
     with tf.Graph().as_default():
       rng = np.random.RandomState(0)
       # Initialize x, latent, state.
@@ -341,6 +350,7 @@ def test_latent_dist_encoder(self, encoder="conv_lstm", skip=True):
       hparams.latent_dist_encoder = encoder
       hparams.latent_skip = skip
       hparams.latent_encoder_width = 256
+      hparams.latent_dropout = dropout
 
       prior_dist, new_state = glow_ops.compute_prior(
           "prior", x_t, latent=latent_t, hparams=hparams, state=init_state,
@@ -411,12 +421,16 @@ def test_actnorm_3d(self):
 
   @parameterized.named_parameters(
       ("dil_relu", True, "relu"), ("no_dil_relu", False, "relu"),
-      ("dil_gatu", True, "gatu"), ("no_dil_gatu", False, "gatu"),)
-  def test_temporal_latent_to_dist(self, apply_dilation, activation):
+      ("dil_gatu", True, "gatu"), ("no_dil_gatu", False, "gatu"),
+      ("dil_relu_drop", True, "relu", 0.1),
+      ("dil_gatu_drop", True, "gatu", 0.1))
+  def test_temporal_latent_to_dist(self, apply_dilation, activation,
+                                   dropout=0.0):
     with tf.Graph().as_default():
       hparams = self.get_glow_hparams()
       hparams.latent_apply_dilations = apply_dilation
       hparams.latent_activation = activation
+      hparams.latent_dropout = dropout
       latent_shape = (16, 5, 32, 32, 48)
       latents = tf.random_normal(latent_shape)
       dist = glow_ops.temporal_latent_to_dist(
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index d2892d5be..910202199 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -62,6 +62,7 @@ def next_frame_glow_hparams():
   hparams.add_hparam("latent_skip", True)
   hparams.add_hparam("latent_encoder_depth", 2)
   hparams.add_hparam("latent_encoder_width", 512)
+  hparams.add_hparam("latent_dropout", 0.0)
   hparams.add_hparam("latent_pre_output_channels", 512)
   hparams.add_hparam("latent_activation", "relu")
   # Pretrains the glow encoder for "pretrain_steps" number of steps.

From 3a23227e832e5847dac8dd4ee6cef2ae2298ac36 Mon Sep 17 00:00:00 2001
From: blazejosinski <804945+blazejosinski@users.noreply.github.com>
Date: Fri, 21 Dec 2018 18:35:10 +0100
Subject: [PATCH 1409/2720] Recording video during training. (#1320)

---
 tensor2tensor/rl/envs/simulated_batch_env.py | 86 ++++++++++++++++++--
 tensor2tensor/rl/trainer_model_based.py      |  7 +-
 2 files changed, 83 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index ba05f8129..0058d70dd 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -23,9 +23,12 @@
 from __future__ import print_function
 
 import copy
+import numpy as np
+import os
 
 from tensor2tensor.data_generators.gym_env import DummyWorldModelProblem
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
 from tensor2tensor.rl.envs import in_graph_batch_env
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
@@ -33,6 +36,18 @@
 import tensorflow as tf
 
 
+# Lazy load PIL.Image
+def PIL_Image():  # pylint: disable=invalid-name
+  from PIL import Image  # pylint: disable=g-import-not-at-top
+  return Image
+
+
+# Lazy load PIL.Image
+def PIL_ImageDraw():  # pylint: disable=invalid-name
+  from PIL import ImageDraw  # pylint: disable=g-import-not-at-top
+  return ImageDraw
+
+
 class HistoryBuffer(object):
   """History Buffer."""
 
@@ -94,7 +109,7 @@ class SimulatedBatchEnv(in_graph_batch_env.InGraphBatchEnv):
   def __init__(
       self, reward_range, observation_space, action_space, frame_stack_size,
       frame_height, frame_width, initial_frame_chooser, batch_size, model_name,
-      model_hparams, model_dir, intrinsic_reward_scale=0.0
+      model_hparams, model_dir, intrinsic_reward_scale=0.0, sim_video_dir=None
   ):
     """Batch of environments inside the TensorFlow graph."""
     super(SimulatedBatchEnv, self).__init__(observation_space, action_space)
@@ -103,6 +118,19 @@ def __init__(
     self._min_reward = reward_range[0]
     self._num_frames = frame_stack_size
     self._intrinsic_reward_scale = intrinsic_reward_scale
+    self._episode_counter = tf.get_variable(
+      "episode_counter", initializer=tf.zeros((), dtype=tf.int32),
+      trainable=False, dtype=tf.int32)
+    if sim_video_dir:
+      self._video_every_epochs = 100
+      self._video_dir = sim_video_dir
+      self._video_writer = None
+      self._video_counter = 0
+      tf.gfile.MakeDirs(self._video_dir)
+      self._video_condition = tf.equal(
+        self._episode_counter.read_value() % self._video_every_epochs, 0)
+    else:
+      self._video_condition = tf.constant(False, dtype=tf.bool, shape=())
 
     model_hparams = copy.copy(model_hparams)
     problem = DummyWorldModelProblem(action_space, reward_range,
@@ -185,9 +213,13 @@ def simulate(self, action):
       done = tf.constant(False, tf.bool, shape=(self.batch_size,))
 
       with tf.control_dependencies([observ]):
+        dump_frame_op = tf.cond(self._video_condition,
+                                lambda: tf.py_func(self._video_dump_frame,
+                                                   [observ, reward], []),
+                                lambda: tf.no_op())
         with tf.control_dependencies(
             [self._observ.assign(observ),
-             self.history_buffer.move_by_one_element(observ)]):
+             self.history_buffer.move_by_one_element(observ), dump_frame_op]):
           clear_reset_model_op = tf.assign(self._reset_model, tf.constant(0.0))
           with tf.control_dependencies([clear_reset_model_op]):
             return tf.identity(reward), tf.identity(done)
@@ -201,12 +233,23 @@ def _reset_non_empty(self, indices):
     Returns:
       Batch tensor of the new observations.
     """
-    with tf.control_dependencies([self.history_buffer.reset(indices)]):
-      with tf.control_dependencies([self._observ.assign(
-          self.history_buffer.get_all_elements()[:, -1, ...])]):
-        reset_model_op = tf.assign(self._reset_model, tf.constant(1.0))
-        with tf.control_dependencies([reset_model_op]):
-          return tf.gather(self._observ.read_value(), indices)
+    reset_video_op = tf.cond(self._video_condition,
+                             lambda: tf.py_func(self._video_reset_writer, [], []),
+                             lambda: tf.no_op())
+    with tf.control_dependencies([reset_video_op]):
+      inc_op = tf.assign_add(self._episode_counter, 1)
+      with tf.control_dependencies([self.history_buffer.reset(indices), inc_op]):
+        initial_frame_dump_op = tf.cond(
+            self._video_condition,
+            lambda: tf.py_func(self._video_dump_frames,
+                               [self.history_buffer.get_all_elements()], []),
+            lambda: tf.no_op())
+        observ_assign_op = self._observ.assign(
+          self.history_buffer.get_all_elements()[:, -1, ...])
+        with tf.control_dependencies([observ_assign_op, initial_frame_dump_op]):
+          reset_model_op = tf.assign(self._reset_model, tf.constant(1.0))
+          with tf.control_dependencies([reset_model_op]):
+            return tf.gather(self._observ.read_value(), indices)
 
   @property
   def observ(self):
@@ -216,3 +259,30 @@ def observ(self):
   @property
   def history_observations(self):
     return self.history_buffer.get_all_elements()
+
+  def _video_dump_frame(self, obs, rews):
+    if self._video_writer is None:
+      self._video_counter += 1
+      self._video_writer = common_video.WholeVideoWriter(
+        fps=10,
+        output_path=os.path.join(self._video_dir,
+                                 "{}.avi".format(self._video_counter)),
+        file_format="avi")
+    img = PIL_Image().new("RGB", (obs.shape[-2], 11), )
+    draw = PIL_ImageDraw().Draw(img)
+    draw.text((0, 0), "r:{:3}".format(int(rews[0])), fill=(255, 0, 0))
+    self._video_writer.write(np.concatenate([np.asarray(img), obs[0]], axis=0))
+
+  def _video_dump_frames(self, obs):
+    zeros = np.zeros(obs.shape[0])
+    for i in range(obs.shape[1]):
+      self._video_dump_frame(obs[:, i, :], zeros)
+
+  def _video_reset_writer(self):
+    if self._video_writer:
+      self._video_writer.finish_to_disk()
+    self._video_writer = None
+
+  def __del__(self):
+    self._video_reset_writer()
+    super(SimulatedBatchEnv, self).__del__()
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 86839f7d0..12c6368e1 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -136,7 +136,8 @@ def choose_subsequence():
 
 
 def make_simulated_env_fn(
-    real_env, hparams, batch_size, initial_frame_chooser, model_dir):
+    real_env, hparams, batch_size, initial_frame_chooser, model_dir,
+    sim_video_dir=None):
   """Creates a simulated env_fn."""
   model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
   if hparams.wm_policy_param_sharing:
@@ -152,6 +153,7 @@ def make_simulated_env_fn(
       model_hparams=trainer_lib.create_hparams(hparams.generative_model_params),
       model_dir=model_dir,
       intrinsic_reward_scale=hparams.intrinsic_reward_scale,
+      sim_video_dir=sim_video_dir,
   )
 
 
@@ -202,7 +204,8 @@ def initial_frame_chooser(batch_size):
     ])
   env_fn = make_simulated_env_fn(
       real_env, hparams, hparams.simulated_batch_size, initial_frame_chooser,
-      world_model_dir
+      world_model_dir, os.path.join(learner.agent_model_dir,
+                                    "sim_videos_{}".format(epoch))
   )
   base_algo_str = hparams.base_algo
   train_hparams = trainer_lib.create_hparams(hparams.base_algo_params)

From f9818cdf7fbc23774ffb0aeab06397bd482da0e3 Mon Sep 17 00:00:00 2001
From: blazejosinski <804945+blazejosinski@users.noreply.github.com>
Date: Fri, 21 Dec 2018 09:35:28 -0800
Subject: [PATCH 1410/2720] internal merge of PR #1320

PiperOrigin-RevId: 226501066
---
 tensor2tensor/rl/envs/simulated_batch_env.py | 39 +++++++++++---------
 1 file changed, 21 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 0058d70dd..dc6848a02 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -23,9 +23,10 @@
 from __future__ import print_function
 
 import copy
-import numpy as np
 import os
 
+import numpy as np
+
 from tensor2tensor.data_generators.gym_env import DummyWorldModelProblem
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
@@ -119,8 +120,8 @@ def __init__(
     self._num_frames = frame_stack_size
     self._intrinsic_reward_scale = intrinsic_reward_scale
     self._episode_counter = tf.get_variable(
-      "episode_counter", initializer=tf.zeros((), dtype=tf.int32),
-      trainable=False, dtype=tf.int32)
+        "episode_counter", initializer=tf.zeros((), dtype=tf.int32),
+        trainable=False, dtype=tf.int32)
     if sim_video_dir:
       self._video_every_epochs = 100
       self._video_dir = sim_video_dir
@@ -128,7 +129,7 @@ def __init__(
       self._video_counter = 0
       tf.gfile.MakeDirs(self._video_dir)
       self._video_condition = tf.equal(
-        self._episode_counter.read_value() % self._video_every_epochs, 0)
+          self._episode_counter.read_value() % self._video_every_epochs, 0)
     else:
       self._video_condition = tf.constant(False, dtype=tf.bool, shape=())
 
@@ -214,9 +215,9 @@ def simulate(self, action):
 
       with tf.control_dependencies([observ]):
         dump_frame_op = tf.cond(self._video_condition,
-                                lambda: tf.py_func(self._video_dump_frame,
+                                lambda: tf.py_func(self._video_dump_frame,  # pylint: disable=g-long-lambda
                                                    [observ, reward], []),
-                                lambda: tf.no_op())
+                                tf.no_op)
         with tf.control_dependencies(
             [self._observ.assign(observ),
              self.history_buffer.move_by_one_element(observ), dump_frame_op]):
@@ -233,19 +234,21 @@ def _reset_non_empty(self, indices):
     Returns:
       Batch tensor of the new observations.
     """
-    reset_video_op = tf.cond(self._video_condition,
-                             lambda: tf.py_func(self._video_reset_writer, [], []),
-                             lambda: tf.no_op())
+    reset_video_op = tf.cond(
+        self._video_condition,
+        lambda: tf.py_func(self._video_reset_writer, [], []),
+        tf.no_op)
     with tf.control_dependencies([reset_video_op]):
       inc_op = tf.assign_add(self._episode_counter, 1)
-      with tf.control_dependencies([self.history_buffer.reset(indices), inc_op]):
+      with tf.control_dependencies([self.history_buffer.reset(indices),
+                                    inc_op]):
         initial_frame_dump_op = tf.cond(
             self._video_condition,
-            lambda: tf.py_func(self._video_dump_frames,
+            lambda: tf.py_func(self._video_dump_frames,  # pylint: disable=g-long-lambda
                                [self.history_buffer.get_all_elements()], []),
-            lambda: tf.no_op())
+            tf.no_op)
         observ_assign_op = self._observ.assign(
-          self.history_buffer.get_all_elements()[:, -1, ...])
+            self.history_buffer.get_all_elements()[:, -1, ...])
         with tf.control_dependencies([observ_assign_op, initial_frame_dump_op]):
           reset_model_op = tf.assign(self._reset_model, tf.constant(1.0))
           with tf.control_dependencies([reset_model_op]):
@@ -264,11 +267,11 @@ def _video_dump_frame(self, obs, rews):
     if self._video_writer is None:
       self._video_counter += 1
       self._video_writer = common_video.WholeVideoWriter(
-        fps=10,
-        output_path=os.path.join(self._video_dir,
-                                 "{}.avi".format(self._video_counter)),
-        file_format="avi")
-    img = PIL_Image().new("RGB", (obs.shape[-2], 11), )
+          fps=10,
+          output_path=os.path.join(self._video_dir,
+                                   "{}.avi".format(self._video_counter)),
+          file_format="avi")
+    img = PIL_Image().new("RGB", (obs.shape[-2], 11),)
     draw = PIL_ImageDraw().Draw(img)
     draw.text((0, 0), "r:{:3}".format(int(rews[0])), fill=(255, 0, 0))
     self._video_writer.write(np.concatenate([np.asarray(img), obs[0]], axis=0))

From 4b8416a26d635b6143e05710a9c8ef45040c1175 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Fri, 21 Dec 2018 09:49:24 -0800
Subject: [PATCH 1411/2720] Move packing to after shuffling.  Previously,
 consecutive sequences from the

PiperOrigin-RevId: 226502722
---
 .../data_generators/generator_utils.py        | 23 ++++++++--
 .../data_generators/text_problems.py          | 43 ++++++++++++++-----
 2 files changed, 52 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index c79babef1..7031b37f3 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -494,22 +494,37 @@ def generate_dataset_and_shuffle(train_gen,
     shuffle_dataset(train_paths + dev_paths)
 
 
-def _shuffle_single(fname):
+def _shuffle_single(fname, extra_fn=None):
+  """Shuffle a single file of records.
+
+  Args:
+    fname: a string
+    extra_fn: an optional function from list of TFRecords to list of TFRecords
+      to be called after shuffling.
+  """
   records = read_records(fname)
   random.shuffle(records)
+  if extra_fn is not None:
+    records = extra_fn(records)
   out_fname = fname.replace(UNSHUFFLED_SUFFIX, "")
   write_records(records, out_fname)
   tf.gfile.Remove(fname)
 
 
-def shuffle_dataset(filenames):
-  """Shuffles the dataset."""
+def shuffle_dataset(filenames, extra_fn=None):
+  """Shuffles the dataset.
+
+  Args:
+    filenames: a list of strings
+    extra_fn: an optional function from list of records to list of records
+      to be called after shuffling a file.
+  """
   if outputs_exist(filenames):
     tf.logging.info("Skipping shuffle because output files exist")
     return
   tf.logging.info("Shuffling data...")
   for filename in filenames:
-    _shuffle_single(filename)
+    _shuffle_single(filename, extra_fn=extra_fn)
   tf.logging.info("Data shuffled.")
 
 
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 118c8a9a1..71561da05 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -242,6 +242,32 @@ def get_or_create_vocab(self, data_dir, tmp_dir, force_get=False):
           "Unrecognized VocabType: %s" % str(self.vocab_type))
     return encoder
 
+  def _pack_fn(self):
+    """For packed datasets, returns a function to pack examples.
+
+    Returns:
+      None or a function from list of TFRecords to list of TFRecords
+    """
+    if not self.packed_length:
+      return None
+    def my_fn(records):
+      """Function from list of TFRecords to list of TFRecords."""
+      examples = []
+      for record in records:
+        x = tf.train.Example()
+        x.ParseFromString(record)
+        example_dict = {}
+        if self.has_inputs:
+          example_dict["inputs"] = [
+              int(i) for i in x.features.feature["inputs"].int64_list.value]
+        example_dict["targets"] = [
+            int(i) for i in x.features.feature["targets"].int64_list.value]
+        examples.append(example_dict)
+      examples = list(self._maybe_pack_examples(examples))
+      return [
+          generator_utils.to_example(x).SerializeToString() for x in examples]
+    return my_fn
+
   def _maybe_pack_examples(self, generator):
     """Wraps generator with packer if self.packed_length."""
     if not self.packed_length:
@@ -302,15 +328,13 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
     if self.is_generate_per_split:
       for split, paths in split_paths:
         generator_utils.generate_files(
-            self._maybe_pack_examples(
-                self.generate_encoded_samples(data_dir, tmp_dir, split)), paths)
+            self.generate_encoded_samples(data_dir, tmp_dir, split), paths)
     else:
       generator_utils.generate_files(
-          self._maybe_pack_examples(
-              self.generate_encoded_samples(
-                  data_dir, tmp_dir, problem.DatasetSplit.TRAIN)), all_paths)
+          self.generate_encoded_samples(
+              data_dir, tmp_dir, problem.DatasetSplit.TRAIN), all_paths)
 
-    generator_utils.shuffle_dataset(all_paths)
+    generator_utils.shuffle_dataset(all_paths, extra_fn=self._pack_fn())
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
@@ -1169,10 +1193,9 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
     # Actually generate examples.
     generator_utils.generate_files(
-        self._maybe_pack_examples(
-            self.generate_encoded_samples(
-                data_dir, tmp_dir, split, input_files)),
+        self.generate_encoded_samples(
+            data_dir, tmp_dir, split, input_files),
         [output_file])
 
     # Shuffle the output.
-    generator_utils.shuffle_dataset([output_file])
+    generator_utils.shuffle_dataset([output_file], extra_fn=self._pack_fn())

From 01a08f4319c7dba27e12c2899cea3149156b09fa Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 21 Dec 2018 11:17:29 -0800
Subject: [PATCH 1412/2720] Format tweaks to glow model

PiperOrigin-RevId: 226515963
---
 tensor2tensor/models/research/glow.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index b0ac4e30e..9f0c427fd 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -61,6 +61,7 @@ def glow_hparams():
   # stability especially when hparams.batch_size is low.
   hparams.add_hparam("init_batch_size", 256)
   hparams.add_hparam("temperature", 1.0)
+
   return hparams
 
 
@@ -192,7 +193,10 @@ def objective_tower(self, features, init=True):
     # through optimisation.
     ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
     with arg_scope(ops, init=init):
-      self.z, encoder_objective, self.eps, _, _ = glow_ops.encoder_decoder(
+      encoder = glow_ops.encoder_decoder
+
+
+      self.z, encoder_objective, self.eps, _, _ = encoder(
           "codec", x, self.hparams, eps=None, reverse=False)
       objective += encoder_objective
 

From 3cbdcc4656f73ef25dc7a7458e2bce02ef851084 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 21 Dec 2018 11:32:01 -0800
Subject: [PATCH 1413/2720] Disable pylint error for MADE not using assertLen.

PiperOrigin-RevId: 226518185
---
 tensor2tensor/layers/reversible_layers_test.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
index f60676952..ec6b67e3f 100644
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from absl.testing import parameterized
 import numpy as np
 
 from tensor2tensor.layers import reversible_layers as reversible
@@ -26,7 +27,7 @@
 import tensorflow as tf
 
 
-class ReversibleLayersTest(tf.test.TestCase):
+class ReversibleLayersTest(parameterized.TestCase, tf.test.TestCase):
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testActNorm(self):
@@ -65,7 +66,9 @@ def testMADELeftToRight(self):
     outputs = network(inputs)
 
     num_weights = sum([np.prod(weight.shape) for weight in network.weights])
+    # Disable lint error for open-source. pylint: disable=g-generic-assert
     self.assertEqual(len(network.weights), 4)
+    # pylint: enable=g-generic-assert
     self.assertEqual(num_weights, (3*1*4 + 4) + (4*3*5 + 3*5))
 
     self.evaluate(tf.global_variables_initializer())
@@ -88,7 +91,9 @@ def testMADERightToLeft(self):
     outputs = network(inputs)
 
     num_weights = sum([np.prod(weight.shape) for weight in network.weights])
+    # Disable lint error for open-source. pylint: disable=g-generic-assert
     self.assertEqual(len(network.weights), 3)
+    # pylint: enable=g-generic-assert
     self.assertEqual(num_weights, 3*5*4 + 4*3 + 3*3*1)
 
     self.evaluate(tf.global_variables_initializer())
@@ -108,7 +113,9 @@ def testMADENoHidden(self):
     outputs = network(inputs)
 
     num_weights = sum([np.prod(weight.shape) for weight in network.weights])
+    # Disable lint error for open-source. pylint: disable=g-generic-assert
     self.assertEqual(len(network.weights), 2)
+    # pylint: enable=g-generic-assert
     self.assertEqual(num_weights, 3*5*3*4 + 3*4)
 
     self.evaluate(tf.global_variables_initializer())

From 9d5a7e90fb50adbd9755281a5c3365d4f33d2238 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Sat, 22 Dec 2018 14:20:07 +0100
Subject: [PATCH 1414/2720] Specify which variables to optimize in PPO training
 (just the policy, not the world model) (#1323)

---
 tensor2tensor/rl/ppo.py         |  3 ++-
 tensor2tensor/utils/optimize.py | 13 ++++++++-----
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index 6002682ab..f152499de 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -61,7 +61,8 @@ def define_ppo_step(data_points, hparams, action_space, lr):
 
   losses = [policy_loss, value_loss, entropy_loss]
   loss = sum(losses)
-  train_op = optimize.optimize(loss, lr, hparams)
+  variables = tf.global_variables(hparams.policy_network + "/.*")
+  train_op = optimize.optimize(loss, lr, hparams, variables=variables)
 
   with tf.control_dependencies([train_op]):
     return [tf.identity(x) for x in losses]
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index b960a2038..d6e8e3e8e 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -30,19 +30,21 @@
 from tensorflow.python.framework import dtypes
 
 
-def optimize(loss, learning_rate, hparams, use_tpu=False):
+def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
   """Minimize loss."""
   loss = weight_decay_and_noise(loss, hparams, learning_rate)
   loss = tf.identity(loss, name="total_loss")
+  if variables is None:
+    variables = tf.trainable_variables()
   # Print trainable variables.
-  log_variable_sizes(verbose=hparams.summarize_vars)
+  log_variable_sizes(variables, verbose=hparams.summarize_vars)
   # Print non-trainable variables.
   non_trainable_variables = list(
-      set(tf.global_variables()) - set(tf.trainable_variables()))
+      set(tf.global_variables()) - set(variables))
   log_variable_sizes(non_trainable_variables, tag="Non-trainable variables",
                      verbose=hparams.summarize_vars)
   if hparams.summarize_vars:
-    summarize_variables()
+    summarize_variables(variables)
     # Summarize non-trainable variables as well
     summarize_variables(non_trainable_variables, tag="Non-trainable variables")
   diet_vars = [
@@ -78,7 +80,8 @@ def optimize(loss, learning_rate, hparams, use_tpu=False):
       gradient_noise_scale=hparams.grad_noise_scale or None,
       optimizer=opt,
       summaries=opt_summaries,
-      colocate_gradients_with_ops=True)
+      colocate_gradients_with_ops=True,
+      variables=variables)
   return train_op
 
 

From 62ba14eeadf5e945691eac52f16088f390c498e7 Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Thu, 27 Dec 2018 01:18:29 -0800
Subject: [PATCH 1415/2720] Adding multiproblems for subsampled cnndm

PiperOrigin-RevId: 227000478
---
 .../data_generators/cnn_dailymail.py          |  16 +++
 .../data_generators/wiki_multi_problems.py    | 102 ++++++++++++++++++
 tensor2tensor/models/transformer.py           |   1 -
 3 files changed, 118 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index 8da272526..5b86c7ad8 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -321,6 +321,14 @@ def fraction_of_data(self):
     return 0.01
 
 
+@registry.register_problem
+class SummarizeFrac2CnnDailymailWikiLMSharedVocab64k(
+    SummarizeFracCnnDailymailWikiLMSharedVocab64k):
+
+  def fraction_of_data(self):
+    return 0.02
+
+
 @registry.register_problem
 class SummarizeFrac5CnnDailymailWikiLMSharedVocab64k(
     SummarizeFracCnnDailymailWikiLMSharedVocab64k):
@@ -329,6 +337,14 @@ def fraction_of_data(self):
     return 0.05
 
 
+@registry.register_problem
+class SummarizeFrac10CnnDailymailWikiLMSharedVocab64k(
+    SummarizeFracCnnDailymailWikiLMSharedVocab64k):
+
+  def fraction_of_data(self):
+    return 0.1
+
+
 @registry.register_problem
 class SummarizeFrac20CnnDailymailWikiLMSharedVocab64k(
     SummarizeFracCnnDailymailWikiLMSharedVocab64k):
diff --git a/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index c25b580fe..5283fec94 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -149,3 +149,105 @@ def __init__(self, was_reversed=False, was_copy=False):
   @property
   def vocab_type(self):
     return text_problems.VocabType.SUBWORD
+
+
+@registry.register_problem
+class LanguagemodelEnWikiLMSummarizeFrac1CnndmSubwords64k(
+    multi_problem.MultiProblem):
+  """Wiki LM and CNN/DM summarization mixed problem class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelEnWikiLMSummarizeFrac1CnndmSubwords64k, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
+    self.task_list.append(
+        cnn_dailymail.SummarizeFrac1CnnDailymailWikiLMSharedVocab64k())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
+
+
+@registry.register_problem
+class LanguagemodelEnWikiLMSummarizeFrac2CnndmSubwords64k(
+    multi_problem.MultiProblem):
+  """Wiki LM and CNN/DM summarization mixed problem class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelEnWikiLMSummarizeFrac2CnndmSubwords64k, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
+    self.task_list.append(
+        cnn_dailymail.SummarizeFrac2CnnDailymailWikiLMSharedVocab64k())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
+
+
+@registry.register_problem
+class LanguagemodelEnWikiLMSummarizeFrac5CnndmSubwords64k(
+    multi_problem.MultiProblem):
+  """Wiki LM and CNN/DM summarization mixed problem class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelEnWikiLMSummarizeFrac5CnndmSubwords64k, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
+    self.task_list.append(
+        cnn_dailymail.SummarizeFrac5CnnDailymailWikiLMSharedVocab64k())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
+
+
+@registry.register_problem
+class LanguagemodelEnWikiLMSummarizeFrac10CnndmSubwords64k(
+    multi_problem.MultiProblem):
+  """Wiki LM and CNN/DM summarization mixed problem class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelEnWikiLMSummarizeFrac10CnndmSubwords64k, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
+    self.task_list.append(
+        cnn_dailymail.SummarizeFrac10CnnDailymailWikiLMSharedVocab64k())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
+
+
+@registry.register_problem
+class LanguagemodelEnWikiLMSummarizeFrac20CnndmSubwords64k(
+    multi_problem.MultiProblem):
+  """Wiki LM and CNN/DM summarization mixed problem class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelEnWikiLMSummarizeFrac20CnndmSubwords64k, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
+    self.task_list.append(
+        cnn_dailymail.SummarizeFrac20CnnDailymailWikiLMSharedVocab64k())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
+
+
+@registry.register_problem
+class LanguagemodelEnWikiLMSummarizeFrac50CnndmSubwords64k(
+    multi_problem.MultiProblem):
+  """Wiki LM and CNN/DM summarization mixed problem class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelEnWikiLMSummarizeFrac50CnndmSubwords64k, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
+    self.task_list.append(
+        cnn_dailymail.SummarizeFrac50CnnDailymailWikiLMSharedVocab64k())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 1312a1a80..5184ef953 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1611,7 +1611,6 @@ def transformer_tall_finetune_uniencdec():
   hparams.learning_rate_constant = 5e-5
   hparams.learning_rate_warmup_steps = 100
   hparams.unidirectional_encoder = True
-  hparams.load_encoder = False
   return hparams
 
 

From 7733fbcb83491c92dda082a76d9673b2382d1060 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Thu, 27 Dec 2018 15:37:51 -0800
Subject: [PATCH 1416/2720] Enable dataset_splits to generate test-data using
 t2t-datagen

PiperOrigin-RevId: 227070522
---
 tensor2tensor/data_generators/bair_robot_pushing.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index c13728207..141844993 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -87,6 +87,14 @@ def only_keep_videos_from_0th_frame(self):
   def use_not_breaking_batching(self):
     return True
 
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each."""
+    return [
+        {"split": problem.DatasetSplit.TRAIN, "shards": 10},
+        {"split": problem.DatasetSplit.EVAL, "shards": 1},
+        {"split": problem.DatasetSplit.TEST, "shards": 1}]
+
   @property
   def extra_reading_spec(self):
     """Additional data fields to store on disk and their decoders."""

From 70f2e3e3a46653a6ade185b826db1757c68ff29d Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Thu, 27 Dec 2018 16:14:52 -0800
Subject: [PATCH 1417/2720] Use test-set in eval job if FLAGS.eval_use_test_set
 is set to True

PiperOrigin-RevId: 227073660
---
 tensor2tensor/bin/t2t_trainer.py   | 1 +
 tensor2tensor/utils/trainer_lib.py | 7 ++++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index e9af961c3..c1a66ba0b 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -188,6 +188,7 @@ def create_experiment_fn():
       eval_early_stopping_metric_minimize=FLAGS
       .eval_early_stopping_metric_minimize,
       eval_timeout_mins=FLAGS.eval_timeout_mins,
+      eval_use_test_set=FLAGS.eval_use_test_set,
       use_tpu=FLAGS.use_tpu,
       use_tpu_estimator=FLAGS.use_tpu_estimator,
       use_xla=FLAGS.xla_compile,
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 7ac9e8fef..847a2ceff 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -662,6 +662,7 @@ def create_experiment(
     eval_early_stopping_metric_delta=None,
     eval_early_stopping_metric_minimize=True,
     eval_timeout_mins=240,
+    eval_use_test_set=False,
     use_tpu=False,
     use_tpu_estimator=False,
     use_xla=False,
@@ -704,8 +705,12 @@ def create_experiment(
   problem = hparams.problem
   train_input_fn = problem.make_estimator_input_fn(tf.estimator.ModeKeys.TRAIN,
                                                    hparams)
+
+  dataset_split = "test" if eval_use_test_set else None
+  dataset_kwargs = {"dataset_split": dataset_split}
   eval_input_fn = problem.make_estimator_input_fn(tf.estimator.ModeKeys.EVAL,
-                                                  hparams)
+                                                  hparams,
+                                                  dataset_kwargs=dataset_kwargs)
 
   # Export
   exporter = None

From ab15710674e8eb94b34af95095604291d60fae6d Mon Sep 17 00:00:00 2001
From: Urvashi Khandelwal <urvashik@google.com>
Date: Sat, 29 Dec 2018 16:04:20 -0800
Subject: [PATCH 1418/2720] Fix decoding.

PiperOrigin-RevId: 227245211
---
 tensor2tensor/utils/decoding.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index bf0dd131d..d33333631 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -615,8 +615,6 @@ def _decode_batch_input_fn(num_decode_batches, sorted_inputs, vocabulary,
       batch_inputs.append(input_ids)
       if len(input_ids) > batch_length:
         batch_length = len(input_ids)
-    if max_input_size != -1:
-      batch_length = max_input_size
     final_batch_inputs = []
     for input_ids in batch_inputs:
       assert len(input_ids) <= batch_length

From 306d71fba5ad018d194a2a426d06dd6397c0a934 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Sat, 29 Dec 2018 18:00:33 -0800
Subject: [PATCH 1419/2720] Add t2t-eval as a wrapper to T2T-Estimator.evaluate
 to evaluate using trained checkpoints on both the validation and the test-set
 as a companion to t2t-trainer and t2t-decode

PiperOrigin-RevId: 227249570
---
 setup.py                      |  1 +
 tensor2tensor/bin/t2t-eval    | 28 +++++++++++++++++
 tensor2tensor/bin/t2t_eval.py | 58 +++++++++++++++++++++++++++++++++++
 3 files changed, 87 insertions(+)
 create mode 100755 tensor2tensor/bin/t2t-eval
 create mode 100644 tensor2tensor/bin/t2t_eval.py

diff --git a/setup.py b/setup.py
index c494ade1c..1b4f1e995 100644
--- a/setup.py
+++ b/setup.py
@@ -24,6 +24,7 @@
         'tensor2tensor/bin/t2t-datagen',
         'tensor2tensor/bin/t2t-decoder',
         'tensor2tensor/bin/t2t-make-tf-configs',
+        'tensor2tensor/bin/t2t-eval',
         'tensor2tensor/bin/t2t-exporter',
         'tensor2tensor/bin/t2t-query-server',
         'tensor2tensor/bin/t2t-insights-server',
diff --git a/tensor2tensor/bin/t2t-eval b/tensor2tensor/bin/t2t-eval
new file mode 100755
index 000000000..fc409359f
--- /dev/null
+++ b/tensor2tensor/bin/t2t-eval
@@ -0,0 +1,28 @@
+#!/usr/bin/env python
+"""Run t2t-eval from a trained checkpoint.
+
+This script is used to run evaluation from a trained checkpoint. Example
+to run evaluation on the test set when trained checkpoint is in /output_dir.
+
+t2t-eval \
+  --problem=image_mnist \
+  --model=imagetransformer \
+  --data_dir=~/t2t
+  --output_dir=/output_dir \
+  --eval_use_test_set=True \
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.bin import t2t_eval
+
+import tensorflow as tf
+
+def main(argv):
+  t2t_eval.main(argv)
+
+
+if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
+  tf.app.run()
diff --git a/tensor2tensor/bin/t2t_eval.py b/tensor2tensor/bin/t2t_eval.py
new file mode 100644
index 000000000..e569fa497
--- /dev/null
+++ b/tensor2tensor/bin/t2t_eval.py
@@ -0,0 +1,58 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+r"""Perform evaluation on trained T2T models using the Estimator API."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.bin import t2t_trainer          # pylint: disable=unused-import
+from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
+from tensor2tensor.utils import trainer_lib
+import tensorflow as tf
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+
+def main(_):
+  tf.logging.set_verbosity(tf.logging.INFO)
+  trainer_lib.set_random_seed(FLAGS.random_seed)
+
+  hparams = trainer_lib.create_hparams(
+      FLAGS.hparams_set, FLAGS.hparams, data_dir=FLAGS.data_dir,
+      problem_name=FLAGS.problem)
+
+  # set appropriate dataset-split, if flags.eval_use_test_set.
+  dataset_split = "test" if FLAGS.eval_use_test_set else None
+  dataset_kwargs = {"dataset_split": dataset_split}
+  eval_input_fn = hparams.problem.make_estimator_input_fn(
+      tf.estimator.ModeKeys.EVAL, hparams, dataset_kwargs=dataset_kwargs)
+  config = t2t_trainer.create_run_config(hparams)
+
+  # summary-hook in tf.estimator.EstimatorSpec requires
+  # hparams.model_dir to be set.
+  hparams.add_hparam("model_dir", config.model_dir)
+
+  estimator = trainer_lib.create_estimator(
+      FLAGS.model, hparams, config, use_tpu=FLAGS.use_tpu)
+  predictions = estimator.evaluate(eval_input_fn, steps=FLAGS.eval_steps)
+  tf.logging.info(predictions)
+
+
+if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
+  tf.app.run()

From 2451614b930c73b2b8dd891b4fc5838d99a151a6 Mon Sep 17 00:00:00 2001
From: Keyon Vafa <vafa@google.com>
Date: Mon, 30 Dec 2019 03:29:33 -0800
Subject: [PATCH 1420/2720] Add SinkhornAutoregressiveFlow to learn permutation
 discrete flow transformation instead of loc-scale.

PiperOrigin-RevId: 227275364
---
 tensor2tensor/layers/reversible_layers.py | 43 +++++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/tensor2tensor/layers/reversible_layers.py b/tensor2tensor/layers/reversible_layers.py
index ba7cf3e73..4a8c02ebc 100644
--- a/tensor2tensor/layers/reversible_layers.py
+++ b/tensor2tensor/layers/reversible_layers.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import numpy as np
+from scipy.optimize import linear_sum_assignment
 import tensorflow as tf
 import tensorflow_probability as tfp
 
@@ -319,3 +320,45 @@ def make_masked_constraint(mask):
   def masked_constraint(x):
     return mask * constraint(x)
   return masked_constraint
+
+
+def sinkhorn(inputs, n_iters=20):
+  """Performs incomplete Sinkhorn normalization to inputs.
+
+  By a theorem by Sinkhorn and Knopp [1], a sufficiently well-behaved  matrix
+  with positive entries can be turned into a doubly-stochastic matrix
+  (i.e. its rows and columns add up to one) via the succesive row and column
+  normalization.
+  -To ensure positivity, the effective input to sinkhorn has to be
+  exp(inputs) (elementwise).
+  -However, for stability, sinkhorn works in the log-space. It is only at
+   return time that entries are exponentiated.
+
+  Code is adapted from Mena et al. [2].
+
+  [1] Richard Sinkhorn and Paul Knopp. Concerning nonnegative matrices and
+  doubly stochastic matrices. Pacific Journal of Mathematics, 1967.
+
+  [2] Gonzalo Mena, David Belanger, Scott Linderman, Jasper Snoek.
+  Learning latent permutations with Gumbel-Sinkhorn networks. International
+  Conference on Learning Representations, 2018.
+
+  Args:
+    inputs: A `Tensor` with shape `[..., vocab_size, vocab_size]`.
+    n_iters: Number of sinkhorn iterations (in practice, as little as 20
+      iterations are needed to achieve decent convergence for `vocab_size` ~100)
+
+  Returns:
+    outputs: A `Tensor` of close-to-doubly-stochastic matrices with shape
+      `[:, vocab_size, vocab_size]`.
+  """
+  vocab_size = tf.shape(inputs)[-1]
+  log_alpha = tf.reshape(inputs, [-1, vocab_size, vocab_size])
+
+  for _ in range(n_iters):
+    log_alpha -= tf.reshape(tf.reduce_logsumexp(log_alpha, axis=2),
+                            [-1, vocab_size, 1])
+    log_alpha -= tf.reshape(tf.reduce_logsumexp(log_alpha, axis=1),
+                            [-1, 1, vocab_size])
+  outputs = tf.exp(log_alpha)
+  return outputs

From e9e40d435df09cb0cb2f68a3026096ec33020c8d Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 2 Jan 2019 11:37:41 -0800
Subject: [PATCH 1421/2720] Minor, change default vgg_ckpt from None to "" to
 avoid tf.contrib.training.HParams parse errors

PiperOrigin-RevId: 227554242
---
 tensor2tensor/utils/decoding.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index d33333631..b89701463 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -83,7 +83,7 @@ def decode_hparams(overrides=""):
       # Used in computation of VGG feature based video metrics.
       # Set this to be the path to a trained VGG ckpt to output
       # useful metrics.
-      vgg_ckpt_path=None,
+      vgg_ckpt_path="",
       # Used for MLPerf compliance logging.
       mlperf_decode_step=0.0,
       mlperf_threshold=25.0,

From 9685f3f4374db64935bdb1ac4c85713e2188c8b3 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 2 Jan 2019 15:47:09 -0800
Subject: [PATCH 1422/2720] Add timeout to continuous decode schedule.

PiperOrigin-RevId: 227595067
---
 tensor2tensor/utils/decoding.py    |  2 ++
 tensor2tensor/utils/trainer_lib.py | 15 ++++++++++-----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index b89701463..b920ca983 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -63,6 +63,8 @@ def decode_hparams(overrides=""):
       delimiter="\n",
       decode_to_file=None,
       decode_in_memory=False,
+      # How much decode should wait for the next checkpoint
+      decode_timeout_mins=240,
       summaries_log_dir="decode",  # Directory to write hook summaries.
       shards=1,    # How many shards of data to decode (treating 1 as None).
       shard_id=0,  # Which shard are we decoding if more than 1 above.
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 847a2ceff..a7d78bd3b 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -594,20 +594,24 @@ def decode(self,
 
   def continuous_decode(self):
     """Decode from dataset on new checkpoint."""
-    for _ in next_checkpoint(self._hparams.model_dir):
+    for _ in next_checkpoint(self._hparams.model_dir,
+                             self._decode_hparams.decode_timeout_mins):
       self.decode()
 
   def continuous_decode_on_train_data(self):
     """Decode from dataset on new checkpoint."""
-    for _ in next_checkpoint(self._hparams.model_dir):
+    for _ in next_checkpoint(self._hparams.model_dir,
+                             self._decode_hparams.decode_timeout_mins):
       self.decode(dataset_split=tf.estimator.ModeKeys.TRAIN)
 
   def continuous_decode_on_eval_data(self):
     """Decode from dataset on new checkpoint."""
     if self._hparams.mlperf_mode:
-      ckpt_generator = next_undecoded_checkpoint(self._hparams.model_dir)
+      ckpt_generator = next_undecoded_checkpoint(
+          self._hparams.model_dir, self._decode_hparams.decode_timeout_mins)
     else:
-      ckpt_generator = next_checkpoint(self._hparams.model_dir)
+      ckpt_generator = next_checkpoint(self._hparams.model_dir,
+                                       self._decode_hparams.decode_timeout_mins)
 
     for ckpt in ckpt_generator:
       current_step = decoding.get_step_from_ckpt_path(ckpt)
@@ -638,7 +642,8 @@ def continuous_decode_on_eval_data(self):
 
   def continuous_decode_from_file(self):
     """Decode from file on new checkpoint."""
-    for _ in next_checkpoint(self._hparams.model_dir):
+    for _ in next_checkpoint(self._hparams.model_dir,
+                             self._decode_hparams.decode_timeout_mins):
       self.decode(decode_from_file=True)
 
 

From 3b9c5fdc977366eb5beee58e66e769be7d9c485e Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 2 Jan 2019 19:00:37 -0800
Subject: [PATCH 1423/2720] Add max_display_decodes as a decode_hparam to limit
 the number of decodes displayed in the decode job.

PiperOrigin-RevId: 227617007
---
 tensor2tensor/data_generators/video_utils.py   |  3 ++-
 .../data_generators/video_utils_test.py        | 18 +++++++++++-------
 tensor2tensor/utils/decoding.py                |  3 ++-
 3 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index f242b2dec..d5f51b5a1 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -166,6 +166,7 @@ def display_video_hooks(hook_args):
   """Hooks to display videos at decode time."""
   predictions = hook_args.predictions
   max_outputs = hook_args.decode_hparams.max_display_outputs
+  max_decodes = hook_args.decode_hparams.max_display_decodes
 
   with tf.Graph().as_default():
     _, best_decodes = video_metrics.compute_video_metrics_from_predictions(
@@ -190,7 +191,7 @@ def display_video_hooks(hook_args):
     all_summaries.extend(summaries)
 
   # Display random decodes for ten conditioning frames.
-  for decode_ind, decode in enumerate(predictions):
+  for decode_ind, decode in enumerate(predictions[: max_decodes]):
     target_videos = video_metrics.stack_data_given_key(decode, "targets")
     output_videos = video_metrics.stack_data_given_key(decode, "outputs")
     input_videos = video_metrics.stack_data_given_key(decode, "inputs")
diff --git a/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index 091b91026..b828fd95e 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -18,6 +18,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+from absl.testing import parameterized
 import numpy as np
 from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
 from tensor2tensor.data_generators import video_utils
@@ -27,9 +28,9 @@
 import tensorflow as tf
 
 
-class VideoUtilsTest(tf.test.TestCase):
+class VideoUtilsTest(parameterized.TestCase, tf.test.TestCase):
 
-  def get_predictions(self):
+  def get_predictions(self, num_decodes=2):
     rng = np.random.RandomState(0)
     # num_samples=4
     inputs = rng.randint(0, 255, (4, 2, 64, 64, 3))
@@ -41,7 +42,7 @@ def get_predictions(self):
       predictions.append(curr_pred)
 
     # num_decodes=2
-    predictions = [predictions] * 2
+    predictions = [predictions] * num_decodes
     problem = registry.problem("video_stochastic_shapes10k")
     return predictions, problem
 
@@ -68,7 +69,9 @@ def testDecodeInMemoryTrue(self):
         predictions=predictions)
     metrics = video_utils.summarize_video_metrics(decode_hooks)
 
-  def testConvertPredictionsToVideoSummaries(self):
+  @parameterized.named_parameters(
+      ("two", 5), ("ten", 10))
+  def testConvertPredictionsToVideoSummaries(self, num_decodes=2):
     # Initialize predictions.
     rng = np.random.RandomState(0)
     inputs = rng.randint(0, 255, (2, 32, 32, 3))
@@ -77,8 +80,9 @@ def testConvertPredictionsToVideoSummaries(self):
 
     # batch it up.
     prediction = [{"outputs": outputs, "inputs": inputs, "targets": targets}]*5
-    predictions = [prediction]
-    decode_hparams = decoding.decode_hparams()
+    predictions = [prediction] * num_decodes
+    decode_hparams = decoding.decode_hparams(
+        overrides="max_display_decodes=5")
 
     decode_hooks = decoding.DecodeHookArgs(
         estimator=None, problem=None, output_dirs=None,
@@ -87,7 +91,7 @@ def testConvertPredictionsToVideoSummaries(self):
     summaries = video_utils.display_video_hooks(decode_hooks)
 
     for summary in summaries:
-      self.assertTrue(isinstance(summary, tf.Summary.Value))
+      self.assertIsInstance(summary, tf.Summary.Value)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index b920ca983..c245b605b 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -80,8 +80,9 @@ def decode_hparams(overrides=""):
       # Creates a blue/red border covering border_percent of the frame.
       border_percent=2,
       # Maximum number of videos displayed.
-      # Total number of videos are max_display_outputs * num_decodes
+      # number of videos displayed = max_display_outputs * max_display_decodes
       max_display_outputs=10,
+      max_display_decodes=5,
       # Used in computation of VGG feature based video metrics.
       # Set this to be the path to a trained VGG ckpt to output
       # useful metrics.

From 75ae65f91085ad1acfa387cd5a4067cb51853ffb Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 2 Jan 2019 21:48:10 -0800
Subject: [PATCH 1424/2720] 1. Do not apply dropout during initialization.

PiperOrigin-RevId: 227629986
---
 tensor2tensor/models/research/glow.py         |  4 +-
 tensor2tensor/models/research/glow_ops.py     | 41 +++++++++++++++++--
 .../models/research/glow_ops_test.py          | 26 ++++++++----
 tensor2tensor/models/video/next_frame_glow.py | 11 +++--
 4 files changed, 64 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 9f0c427fd..7e95579f9 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -118,7 +118,7 @@ def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
     features["targets"] = tf.zeros(shape=(batch_size, 1, 1, 1))
     _, _ = self(features)  # pylint: disable=not-callable
 
-    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
+    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm, glow_ops.get_dropout]
     var_scope = tf.variable_scope("glow/body", reuse=True)
     # If eps=None, images are sampled from the prior.
     with arg_scope(ops, init=False), var_scope:
@@ -191,7 +191,7 @@ def objective_tower(self, features, init=True):
     # the per-channel output activations have zero mean and unit variance
     # ONLY during the first step. After that the parameters are learned
     # through optimisation.
-    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
+    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm, glow_ops.get_dropout]
     with arg_scope(ops, init=init):
       encoder = glow_ops.encoder_decoder
 
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 580a9bd34..a8795bc34 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -111,6 +111,22 @@ def get_variable_ddi(name, shape, initial_value, dtype=tf.float32, init=False,
     return tf.cond(init, lambda: assign(w, initial_value), lambda: w)
 
 
+@add_arg_scope
+def get_dropout(x, rate=0.0, init=True):
+  """Zero dropout during init or prediction time.
+
+  Args:
+    x: 4-D Tensor, shape=(NHWC).
+    rate: Dropout rate.
+    init: Initialization.
+  Returns:
+    x: activations after dropout.
+  """
+  if init or rate == 0:
+    return x
+  return tf.layers.dropout(x, rate=rate, training=True)
+
+
 @add_arg_scope
 def actnorm_3d(name, x, logscale_factor=3.):
   """Applies actnorm to each time-step independently.
@@ -485,8 +501,7 @@ def conv_block(name, x, mid_channels, dilations=None, activation="relu",
     x = conv("1_1", x, output_channels=mid_channels, filter_size=first_filter,
              dilations=dilations)
     x = tf.nn.relu(x)
-    if dropout != 0.0:
-      x = tf.nn.dropout(x, keep_prob=1.0 - dropout)
+    x = get_dropout(x, rate=dropout)
 
     # Padding + conv2d + actnorm + activation.
     # [input, output: 512 channels]
@@ -502,8 +517,7 @@ def conv_block(name, x, mid_channels, dilations=None, activation="relu",
                     filter_size=second_filter, dilations=dilations)
       x = tf.nn.tanh(x_tanh) * tf.nn.sigmoid(x_sigm)
 
-    if dropout != 0.0:
-      x = tf.nn.dropout(x, keep_prob=1.0 - dropout)
+    x = get_dropout(x, rate=dropout)
     return x
 
 
@@ -829,6 +843,22 @@ def latent_to_dist(name, x, hparams, output_channels=None):
     return tfp.distributions.Normal(mean, tf.exp(log_scale))
 
 
+@add_arg_scope
+def noise_op(latents, hparams):
+  """Adds isotropic gaussian-noise to each latent.
+
+  Args:
+    latents: 4-D or 5-D tensor, shape=(NTHWC) or (NHWC).
+    hparams: tf.contrib.training.HParams.
+  Returns:
+    latents: latents with isotropic gaussian noise appended.
+  """
+  if hparams.latent_noise == 0 or hparams.mode != tf.estimator.ModeKeys.TRAIN:
+    return latents
+  latent_shape = common_layers.shape_list(latents)
+  return latents + tf.random_normal(latent_shape, stddev=hparams.latent_noise)
+
+
 @add_arg_scope
 def merge_level_and_latent_dist(level_dist, latent_dist,
                                 merge_std="prev_level"):
@@ -894,6 +924,7 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
     output_channels = common_layers.shape_list(z)[-1]
     last_latent = latent[-1]
     latent_stack = tf.concat([prior_dist.loc] + latent, axis=-1)
+    latent_stack = noise_op(latent_stack, hparams)
     cond_dist = latent_to_dist(
         "latent_stack", latent_stack, hparams=hparams,
         output_channels=output_channels)
@@ -910,6 +941,7 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
     prev_latents = tf.tile(tf.expand_dims(prior_dist.loc, axis=1),
                            [1, num_steps, 1, 1, 1])
     cond_latents = tf.concat((cond_latents, prev_latents), axis=-1)
+    cond_latents = noise_op(cond_latents, hparams)
     cond_dist = temporal_latent_to_dist(
         "latent_stack", cond_latents, hparams, output_channels=output_channels)
 
@@ -917,6 +949,7 @@ def level_cond_prior(prior_dist, z, latent, hparams, state):
     last_latent = latent
     output_channels = common_layers.shape_list(z)[-1]
     latent_stack = tf.concat((prior_dist.loc, latent), axis=-1)
+    latent_stack = noise_op(latent_stack, hparams)
     _, state = common_video.conv_lstm_2d(
         latent_stack, state, hparams.latent_encoder_width, kernel_size=3,
         name="conv_lstm")
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 07004294f..101edb2d3 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -39,6 +39,7 @@ class GlowOpsTest(parameterized.TestCase, tf.test.TestCase):
 
   def get_glow_hparams(self):
     hparams = glow.glow_hparams()
+    hparams.add_hparam("mode", tf.estimator.ModeKeys.TRAIN)
     hparams.add_hparam("num_cond_latents", 1)
     hparams.add_hparam("latent_architecture", "glow_resnet")
     # Use latent skip connections
@@ -52,6 +53,7 @@ def get_glow_hparams(self):
     hparams.add_hparam("latent_time_filter_size", 3)
     hparams.add_hparam("latent_activation", "relu")
     hparams.add_hparam("latent_dropout", 0.0)
+    hparams.add_hparam("latent_noise", 0.0)
     return hparams
 
   def test_get_variable_ddi(self):
@@ -88,8 +90,9 @@ def test_invertibility(self, op, name, dropout=0.0):
       x = tf.random_uniform(shape=(16, 32, 32, 4))
 
       if op in [glow_ops.affine_coupling, glow_ops.additive_coupling]:
-        x_inv, _ = op(name, x, reverse=False, dropout=dropout)
-        x_inv_inv, _ = op(name, x_inv, reverse=True, dropout=dropout)
+        with arg_scope([glow_ops.get_dropout], init=False):
+          x_inv, _ = op(name, x, reverse=False, dropout=dropout)
+          x_inv_inv, _ = op(name, x_inv, reverse=True, dropout=dropout)
       else:
         x_inv, _ = op(name, x, reverse=False)
         x_inv_inv, _ = op(name, x_inv, reverse=True)
@@ -331,9 +334,10 @@ def test_split_latent_conditioning(self):
       ("conv3d_skip", "conv3d_net", False),
       ("conv3d_no_skip", "conv3d_net", True),
       ("conv3d_skip_drop", "conv3d_net", False, 0.1),
-      ("conv3d_no_skip_drop", "conv3d_net", True, 0.1))
+      ("conv3d_no_skip_drop", "conv3d_net", True, 0.1),
+      ("conv3d_no_skip_drop_noise", "conv3d_net", True, 0.1, 0.1),)
   def test_latent_dist_encoder(self, encoder="conv_lstm", skip=True,
-                               dropout=0.0):
+                               dropout=0.0, noise=0.1):
     with tf.Graph().as_default():
       rng = np.random.RandomState(0)
       # Initialize x, latent, state.
@@ -351,10 +355,12 @@ def test_latent_dist_encoder(self, encoder="conv_lstm", skip=True,
       hparams.latent_skip = skip
       hparams.latent_encoder_width = 256
       hparams.latent_dropout = dropout
+      hparams.latent_noise = noise
 
-      prior_dist, new_state = glow_ops.compute_prior(
-          "prior", x_t, latent=latent_t, hparams=hparams, state=init_state,
-          condition=True)
+      with arg_scope([glow_ops.get_dropout], init=False):
+        prior_dist, new_state = glow_ops.compute_prior(
+            "prior", x_t, latent=latent_t, hparams=hparams, state=init_state,
+            condition=True)
       with tf.Session() as sess:
         sess.run(tf.global_variables_initializer())
         # Test initialization:
@@ -423,14 +429,16 @@ def test_actnorm_3d(self):
       ("dil_relu", True, "relu"), ("no_dil_relu", False, "relu"),
       ("dil_gatu", True, "gatu"), ("no_dil_gatu", False, "gatu"),
       ("dil_relu_drop", True, "relu", 0.1),
-      ("dil_gatu_drop", True, "gatu", 0.1))
+      ("dil_gatu_drop", True, "gatu", 0.1),
+      ("dil_gatu_drop_noise", True, "gatu", 0.1, 0.1))
   def test_temporal_latent_to_dist(self, apply_dilation, activation,
-                                   dropout=0.0):
+                                   dropout=0.0, noise=0.1):
     with tf.Graph().as_default():
       hparams = self.get_glow_hparams()
       hparams.latent_apply_dilations = apply_dilation
       hparams.latent_activation = activation
       hparams.latent_dropout = dropout
+      hparams.latent_noise = noise
       latent_shape = (16, 5, 32, 32, 48)
       latents = tf.random_normal(latent_shape)
       dist = glow_ops.temporal_latent_to_dist(
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index 910202199..e18db1867 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -65,6 +65,7 @@ def next_frame_glow_hparams():
   hparams.add_hparam("latent_dropout", 0.0)
   hparams.add_hparam("latent_pre_output_channels", 512)
   hparams.add_hparam("latent_activation", "relu")
+  hparams.add_hparam("latent_noise", 0.0)
   # Pretrains the glow encoder for "pretrain_steps" number of steps.
   # By default, don't pretrain and learn end-to-end
   hparams.add_hparam("pretrain_steps", -1)
@@ -170,7 +171,7 @@ def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
     else:
       num_target_frames = self.hparams.video_num_target_frames
 
-    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
+    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm, glow_ops.get_dropout]
     var_scope = tf.variable_scope("next_frame_glow/body", reuse=True)
     all_frames = []
 
@@ -285,6 +286,7 @@ def top_cond_prior(self, name, cond_top_latents):
         cond_top_latents = tf.concat(cond_top_latents, axis=-1)
 
         # Maps the latent-stack to a distribution.
+        cond_top_latents = glow_ops.noise_op(cond_top_latents, self.hparams)
         top = glow_ops.latent_to_dist(
             name, cond_top_latents, hparams=self.hparams,
             output_channels=output_channels)
@@ -293,6 +295,7 @@ def top_cond_prior(self, name, cond_top_latents):
         output_channels = common_layers.shape_list(cond_top_latents)[-1]
         # (h_t, c_t) = LSTM(z_{t-1}; (h_{t-1}, c_{t-1}))
         # (mu_t, sigma_t) = conv(h_t)
+        cond_top_latents = glow_ops.noise_op(cond_top_latents, self.hparams)
         _, self.top_state = common_video.conv_lstm_2d(
             cond_top_latents, self.top_state, self.hparams.latent_encoder_width,
             kernel_size=3, name="conv_lstm")
@@ -300,8 +303,10 @@ def top_cond_prior(self, name, cond_top_latents):
             name, self.top_state.h, output_channels=output_channels)
       elif self.hparams.latent_dist_encoder == "conv3d_net":
         last_latent = cond_top_latents[-1]
+        cond_top_latents = tf.stack(cond_top_latents, axis=1)
+        cond_top_latents = glow_ops.noise_op(cond_top_latents, self.hparams)
         top = glow_ops.temporal_latent_to_dist(
-            "conv3d", tf.stack(cond_top_latents, axis=1), self.hparams)
+            "conv3d", cond_top_latents, self.hparams)
 
       # mu(z_{t}) = z_{t-1} + latent_encoder(z_{cond})
       if self.hparams.latent_skip:
@@ -477,7 +482,7 @@ def video_objective_tower(self, input_frames, target_frames, init=False):
 
     cond_level_latents, cond_top_latents = None, None
     total_objective = 0.0
-    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
+    ops = [glow_ops.get_variable_ddi, glow_ops.actnorm, glow_ops.get_dropout]
 
     with arg_scope(ops, init=init):
       for frame_ind, frame in enumerate(all_frames):

From 94477898f3a862654f045a2ba90e8a5848037cfd Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 2 Jan 2019 23:12:53 -0800
Subject: [PATCH 1425/2720] Get MTF Transformer language model decoding from
 file working on TPU.

PiperOrigin-RevId: 227636124
---
 tensor2tensor/models/mtf_transformer2.py | 31 +++++++++++-----
 tensor2tensor/utils/decoding.py          | 45 ++++++++++++++++++++++--
 2 files changed, 66 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index 4ddac0837..9f29fd113 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -54,6 +54,12 @@ def batch_dims(self):
           mtf.Dimension("inner_batch",
                         hparams.batch_size // hparams.outer_batch_size)]
 
+  def combine_batch_dims(self, x):
+    if len(self.batch_dims) <= 1:
+      return x
+    return mtf.replace_dimensions(
+        x, self.batch_dims, mtf.combined_dimension(self.batch_dims))
+
   @property
   def autoregressive(self):
     return self._hparams.autoregressive
@@ -151,11 +157,7 @@ def import_feature(key):
   def mtf_model_fn(self, features, mesh):
     logits, loss = self._mtf_model_fn(features, mesh)
     # combine batch dims
-    if len(self.batch_dims) > 1:
-      combined_batch_dim = mtf.Dimension(
-          self.batch_dims[0].name, mtf.Shape(self.batch_dims).size)
-      logits = mtf.reshape(
-          logits, [combined_batch_dim] + logits.shape.dims[-2:])
+    logits = self.combine_batch_dims(logits)
     return logits, loss
 
   @property
@@ -185,16 +187,20 @@ def import_feature(key):
       partial_targets = import_feature("inputs")
       if partial_targets is None:
         partial_targets = import_feature("targets")
-      if partial_targets is None:
+      if partial_targets:
+        partial_targets *= mtf.cast(
+            mtf.not_equal(partial_targets, 1), partial_targets.dtype)
+      else:
         ids_shape = mtf.Shape(self.batch_dims + [self.length_dim])
         partial_targets = mtf.constant(mesh, 0, ids_shape, dtype=tf.int32)
       if hparams.beam_size > 1:
         raise NotImplementedError(
             "Beam search not implemented for unitransformer.")
-      return model.sample_autoregressive(
+      ret = model.sample_autoregressive(
           partial_targets,
           temperature=hparams.sampling_temp,
           variable_dtype=self.variable_dtype)
+      return self.combine_batch_dims(ret)
     else:
       raise ValueError(
           "Don't know how to sample from non-autoregressive unitransformer")
@@ -258,7 +264,7 @@ def sample(self, features, mesh):
     hparams = self._hparams
     model = self.model()
     inputs = self._import_feature(features, mesh, "inputs")
-    return model.decode(
+    ret = model.decode(
         inputs,
         self.variable_dtype,
         beam_size=hparams.beam_size,
@@ -266,6 +272,7 @@ def sample(self, features, mesh):
         temperature=hparams.sampling_temp if hparams.beam_size == 1 else 0,
         decode_length_multiplier=hparams.decode_length_multiplier,
         decode_length_constant=hparams.decode_length_constant)
+    return self.combine_batch_dims(ret)
 
 
 # The following functions construct layers based on hyperparmeters
@@ -524,6 +531,14 @@ def mtr_lm_dense_0():
   return mtr_lm_dense(0)
 
 
+@registry.register_hparams
+def mtr_lm_dense_0_h1_16():
+  hparams = mtr_lm_dense_0()
+  hparams.decoder_num_heads = 16
+  hparams.decoder_num_memory_heads = 1
+  return hparams
+
+
 @registry.register_hparams
 def mtr_lm_dense_1():
   return mtr_lm_dense(1)
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index c245b605b..1045c621f 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -22,6 +22,7 @@
 import operator
 import os
 import re
+import string
 import time
 
 import numpy as np
@@ -373,14 +374,23 @@ def decode_from_file(estimator,
   problem_name = FLAGS.problem
   filename = _add_shard_to_filename(filename, decode_hp)
   tf.logging.info("Performing decoding from file (%s)." % filename)
-  sorted_inputs, sorted_keys = _get_sorted_inputs(filename, decode_hp.delimiter)
+  if has_input:
+    sorted_inputs, sorted_keys = _get_sorted_inputs(
+        filename, decode_hp.delimiter)
+  else:
+    sorted_inputs = _get_language_modeling_inputs(
+        filename, decode_hp.delimiter, repeat=decode_hp.num_decodes)
+    sorted_keys = range(len(sorted_inputs))
   num_decode_batches = (len(sorted_inputs) - 1) // decode_hp.batch_size + 1
 
   if estimator.config.use_tpu:
     length = getattr(hparams, "length", hparams.max_length)
     batch_ids = []
     for line in sorted_inputs:
-      ids = inputs_vocab.encode(line.strip()) + [1]
+      if has_input:
+        ids = inputs_vocab.encode(line.strip()) + [1]
+      else:
+        ids = targets_vocab.encode(line)
       if len(ids) < length:
         ids.extend([0] * (length - len(ids)))
       else:
@@ -741,6 +751,37 @@ def show_and_save_image(img, save_path):
     plt.savefig(sp)
 
 
+def _get_language_modeling_inputs(filename,
+                                  delimiter="\n",
+                                  repeat=1,
+                                  append_space_to_final_punctionation=True):
+  """Read a file of partial texts to continue.
+
+  The purpose of append_space_to_final_punctionation is that SubwordTokenizer
+  groups punctuation and the ensuing space in the same token.  Adding a space
+  causes the token to be completed.
+
+  Args:
+    filename: a string
+    delimiter: a string
+    repeat: an integer - we repeat the entire file that many times.
+    append_space_to_final_punctionation: a boolean
+
+  Returns:
+    a list of strings
+  """
+  with tf.gfile.Open(filename) as f:
+    text = f.read()
+  inputs = text.split(delimiter)
+  if not inputs[-1]:
+    inputs.pop()
+  inputs *= repeat
+  if append_space_to_final_punctionation:
+    inputs = [
+        s + " " if s and s[-1] in string.punctuation else s for s in inputs]
+  return inputs
+
+
 def _get_sorted_inputs(filename, delimiter="\n"):
   """Returning inputs sorted according to decreasing length.
 

From de2b81e2042728ade3399c2950e55d540cfe40bf Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Thu, 3 Jan 2019 22:45:12 -0800
Subject: [PATCH 1426/2720] Minor fixes and configs for training wmt-ende on
 Mesh-TensorFlow.   Good quality requires
 shared_embedding_and_softmax_weights, which wasn't hooked up properly.

PiperOrigin-RevId: 227804947
---
 tensor2tensor/models/mtf_transformer2.py | 39 ++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index 9f29fd113..82aa9f9c1 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -127,6 +127,8 @@ def model(self):
         output_vocab_size=self._targets_vocab_size,
         autoregressive=self.autoregressive,
         max_length=hparams.max_length,
+        shared_embedding_and_softmax_weights=(
+            hparams.shared_embedding_and_softmax_weights),
         z_loss=hparams.z_loss,
         layout=hparams.layout,
         mesh_shape=hparams.mesh_shape)
@@ -229,6 +231,8 @@ def model(self):
         output_vocab_size=self._targets_vocab_size,
         max_length=hparams.max_length,
         shared_embedding=hparams.shared_embedding,
+        shared_embedding_and_softmax_weights=(
+            hparams.shared_embedding_and_softmax_weights),
         label_smoothing=hparams.label_smoothing,
         z_loss=hparams.z_loss,
         layout=hparams.layout,
@@ -700,6 +704,13 @@ def mtr_tr_dense_local_0_h1_16():
   return hparams
 
 
+@registry.register_hparams
+def mtr_tr_dense_local_0_h1_16_shared():
+  hparams = mtr_tr_dense_local_0_h1_16()
+  hparams.shared_embedding_and_softmax_weights = True
+  return hparams
+
+
 @registry.register_hparams
 def mtr_tr_dense_local_0_h1_8_kv256():
   hparams = mtr_tr_dense_local_0()
@@ -772,3 +783,31 @@ def mtr_tr_dense_0_shared_kv():
   hparams = mtr_tr_dense_0()
   hparams.decoder_shared_kv = True
   return hparams
+
+
+@registry.register_hparams
+def mtr_tr_enfr_v0():
+  # good parameters for wmt-en-fr
+  hparams = mtr_tr_dense_local_0_h1_16()
+  return hparams
+
+
+@registry.register_hparams
+def mtr_tr_ende_v0():
+  # good parameters for wmt-en-de
+  hparams = mtr_tr_dense_local_0_h1_16()
+  hparams.learning_rate_decay_steps = 20000
+  hparams.shared_embedding_and_softmax_weights = True
+  hparams.layer_prepostprocess_dropout = 0.2
+  return hparams
+
+
+@registry.register_hparams
+def mtr_tr_ende_deep():
+  hparams = mtr_tr_ende_v0()
+  hparams.decoder_num_heads = 8
+  hparams.encoder_num_heads = 4
+  hparams.d_ff = 2048
+  hparams.encoder_num_layers = 12
+  hparams.decoder_num_layers = 12
+  return hparams

From f780b44c1b5ab42bad6e3ec3470d6d26a3c6dc7c Mon Sep 17 00:00:00 2001
From: hsm207 <hsm207@users.noreply.github.com>
Date: Sat, 5 Jan 2019 03:13:58 +0800
Subject: [PATCH 1427/2720] Fix typo (#1329)

---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 485dc41e1..01ac52b9e 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -423,7 +423,7 @@ def add_timing_signal_1d(x,
   memory inputs to attention.
 
   The use of relative position is possible because sin(x+y) and cos(x+y) can be
-  experessed in terms of y, sin(x) and cos(x).
+  expressed in terms of y, sin(x) and cos(x).
 
   In particular, we use a geometric sequence of timescales starting with
   min_timescale and ending with max_timescale.  The number of different

From 3e483ad36bba231f3e4036d055d22b5d259ba473 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 4 Jan 2019 20:38:00 +0100
Subject: [PATCH 1428/2720] Use tf_inspect for Python 2/3 compatibility (#1324)

---
 tensor2tensor/utils/devices.py   | 2 +-
 tensor2tensor/utils/metrics.py   | 2 +-
 tensor2tensor/utils/registry.py  | 2 +-
 tensor2tensor/utils/t2t_model.py | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index 7017bb03a..2a582de84 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -18,9 +18,9 @@
 from __future__ import division
 from __future__ import print_function
 
-import inspect
 from tensor2tensor.utils import expert_utils as eu
 import tensorflow as tf
+from tensorflow.python.util import tf_inspect as inspect
 
 
 def data_parallelism_from_flags(daisy_chain_variables=True, all_workers=False):
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index aac82e868..461a5e62e 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -18,7 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import inspect
 import numpy as np
 import six
 
@@ -29,6 +28,7 @@
 import tensorflow as tf
 
 from tensorflow.contrib.eager.python import tfe
+from tensorflow.python.util import tf_inspect as inspect
 
 
 class Metrics(object):
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 6b5046382..7c9225f42 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -44,9 +44,9 @@ class MyModel(T2TModel):
 from __future__ import division
 from __future__ import print_function
 
-import inspect
 from tensor2tensor.utils import misc_utils
 import tensorflow as tf
+from tensorflow.python.util import tf_inspect as inspect
 
 _ATTACKS = {}
 _ATTACK_PARAMS = {}
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 393d66cb0..4d99b9f02 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -22,7 +22,6 @@
 import contextlib
 import copy
 import functools
-import inspect
 import math
 import os
 import time
@@ -48,6 +47,7 @@
 from tensorflow.python.layers import base
 from tensorflow.python.ops import inplace_ops
 from tensorflow.python.ops import variable_scope
+from tensorflow.python.util import tf_inspect as inspect
 
 _no_problem_err_str = (
     "The default implementation of %s requires that the "

From e587d4cc5c7c5e63fae712897bde1c4ad3a06858 Mon Sep 17 00:00:00 2001
From: hsm207 <hsm207@users.noreply.github.com>
Date: Fri, 4 Jan 2019 11:14:26 -0800
Subject: [PATCH 1429/2720] internal merge of PR #1329

PiperOrigin-RevId: 227883955
---
 tensor2tensor/utils/devices.py   | 2 +-
 tensor2tensor/utils/metrics.py   | 2 +-
 tensor2tensor/utils/registry.py  | 2 +-
 tensor2tensor/utils/t2t_model.py | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index 2a582de84..7017bb03a 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -18,9 +18,9 @@
 from __future__ import division
 from __future__ import print_function
 
+import inspect
 from tensor2tensor.utils import expert_utils as eu
 import tensorflow as tf
-from tensorflow.python.util import tf_inspect as inspect
 
 
 def data_parallelism_from_flags(daisy_chain_variables=True, all_workers=False):
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 461a5e62e..aac82e868 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import inspect
 import numpy as np
 import six
 
@@ -28,7 +29,6 @@
 import tensorflow as tf
 
 from tensorflow.contrib.eager.python import tfe
-from tensorflow.python.util import tf_inspect as inspect
 
 
 class Metrics(object):
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 7c9225f42..6b5046382 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -44,9 +44,9 @@ class MyModel(T2TModel):
 from __future__ import division
 from __future__ import print_function
 
+import inspect
 from tensor2tensor.utils import misc_utils
 import tensorflow as tf
-from tensorflow.python.util import tf_inspect as inspect
 
 _ATTACKS = {}
 _ATTACK_PARAMS = {}
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 4d99b9f02..393d66cb0 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -22,6 +22,7 @@
 import contextlib
 import copy
 import functools
+import inspect
 import math
 import os
 import time
@@ -47,7 +48,6 @@
 from tensorflow.python.layers import base
 from tensorflow.python.ops import inplace_ops
 from tensorflow.python.ops import variable_scope
-from tensorflow.python.util import tf_inspect as inspect
 
 _no_problem_err_str = (
     "The default implementation of %s requires that the "

From 0fe422035bc0d56f7d728bd1a46f7277f8296907 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 4 Jan 2019 11:38:27 -0800
Subject: [PATCH 1430/2720] internal merge of PR #1324

PiperOrigin-RevId: 227888281
---
 tensor2tensor/utils/devices.py   | 2 +-
 tensor2tensor/utils/metrics.py   | 2 +-
 tensor2tensor/utils/registry.py  | 2 +-
 tensor2tensor/utils/t2t_model.py | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index 7017bb03a..2a582de84 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -18,9 +18,9 @@
 from __future__ import division
 from __future__ import print_function
 
-import inspect
 from tensor2tensor.utils import expert_utils as eu
 import tensorflow as tf
+from tensorflow.python.util import tf_inspect as inspect
 
 
 def data_parallelism_from_flags(daisy_chain_variables=True, all_workers=False):
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index aac82e868..461a5e62e 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -18,7 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import inspect
 import numpy as np
 import six
 
@@ -29,6 +28,7 @@
 import tensorflow as tf
 
 from tensorflow.contrib.eager.python import tfe
+from tensorflow.python.util import tf_inspect as inspect
 
 
 class Metrics(object):
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 6b5046382..7c9225f42 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -44,9 +44,9 @@ class MyModel(T2TModel):
 from __future__ import division
 from __future__ import print_function
 
-import inspect
 from tensor2tensor.utils import misc_utils
 import tensorflow as tf
+from tensorflow.python.util import tf_inspect as inspect
 
 _ATTACKS = {}
 _ATTACK_PARAMS = {}
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 393d66cb0..4d99b9f02 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -22,7 +22,6 @@
 import contextlib
 import copy
 import functools
-import inspect
 import math
 import os
 import time
@@ -48,6 +47,7 @@
 from tensorflow.python.layers import base
 from tensorflow.python.ops import inplace_ops
 from tensorflow.python.ops import variable_scope
+from tensorflow.python.util import tf_inspect as inspect
 
 _no_problem_err_str = (
     "The default implementation of %s requires that the "

From 948b32b747c3755c6a276d1b5cb7a449a67b8d5a Mon Sep 17 00:00:00 2001
From: kngxscn <yincongxian@foxmail.com>
Date: Sat, 5 Jan 2019 06:05:36 +0800
Subject: [PATCH 1431/2720] Fix bug: "is_generate_per_split" should be set to
 property. (#1322)

---
 tensor2tensor/data_generators/translate.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 1a35c9459..82f7eeb8f 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -36,6 +36,7 @@
 class TranslateProblem(text_problems.Text2TextProblem):
   """Base class for translation problems."""
 
+  @property
   def is_generate_per_split(self):
     return True
 

From 0d230010b5ec0bfd6c0e0a3fc494d7e0df0486f4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 4 Jan 2019 14:10:13 -0800
Subject: [PATCH 1432/2720] Add explicit Python 2 kernel metadata to notebooks.

PiperOrigin-RevId: 227911634
---
 tensor2tensor/notebooks/hello_t2t.ipynb | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/notebooks/hello_t2t.ipynb b/tensor2tensor/notebooks/hello_t2t.ipynb
index 5e163d876..f2112adf5 100644
--- a/tensor2tensor/notebooks/hello_t2t.ipynb
+++ b/tensor2tensor/notebooks/hello_t2t.ipynb
@@ -14,6 +14,10 @@
         }
       ],
       "collapsed_sections": []
+    },
+    "kernelspec": {
+      "display_name": "Python 2",
+      "name": "python2"
     }
   },
   "cells": [

From dbab44cf3dce3ea834ff14347b78c5f9f78e1b13 Mon Sep 17 00:00:00 2001
From: Youngwook Kim <youngwook.kim@gmail.com>
Date: Sat, 5 Jan 2019 07:23:23 +0900
Subject: [PATCH 1433/2720] Custom evaluation metrics (#1336)

* Custom evaluation metrics

* Fix Python 2 compatibility issue

* Fix notebook test
---
 tensor2tensor/data_generators/problem.py | 13 ++++++++
 tensor2tensor/utils/metrics.py           | 39 +++++++++++++++---------
 tensor2tensor/utils/t2t_model.py         | 10 +++---
 3 files changed, 43 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 03ec21ea1..b89aac526 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -367,6 +367,19 @@ def eval_metrics(self):
         metrics.Metrics.ACC_PER_SEQ, metrics.Metrics.NEG_LOG_PERPLEXITY
     ]
 
+  def eval_metric_fns(self, model_hparams):
+    metric_names = self.eval_metrics()
+    if not all([m in metrics.METRICS_FNS for m in metric_names]):
+      error_str = ("Unrecognized metric. Problem %s specified metrics "
+                   "%s. Recognized metrics are %s.")
+      raise ValueError(error_str % (self.name,
+                                    metric_names,
+                                    list(metrics.METRICS_FNS.keys())))
+    return {
+        metric_name: metrics.METRICS_FNS[metric_name]
+        for metric_name in metric_names
+    }
+
   def eval_hooks(self, features, logits, hparams):
     del features, logits, hparams
     return []
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 461a5e62e..5608c0565 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -602,15 +602,9 @@ def weights_fn_for_mp(problem_task_id):
     problem_name = problem_instance.name
     if problem_instance.was_reversed:
       problem_name += "_rev"
-    metrics = problem_instance.eval_metrics()
+    metrics = problem_instance.eval_metric_fns(model_hparams)
     if hasattr(model_hparams.problem, "task_list"):
-      metrics = model_hparams.problem.eval_metrics()
-    if not all([m in METRICS_FNS for m in metrics]):
-      error_str = ("Unrecognized metric. Problem %s specified metrics "
-                   "%s. Recognized metrics are %s.")
-      raise ValueError(error_str % (problem_name,
-                                    metrics,
-                                    list(METRICS_FNS.keys())))
+      metrics = model_hparams.problem.eval_metric_fns(model_hparams)
 
     tm = problem_instance.get_hparams(model_hparams).modality["targets"]
     if not isinstance(tm, dict):
@@ -622,8 +616,7 @@ def weights_fn_for_mp(problem_task_id):
         ptid = problem_instance.task_id  # pylint: disable=cell-var-from-loop
         weights_fn = weights_fn_for_mp(ptid)
 
-      for metric in metrics:
-        metric_fn = METRICS_FNS[metric]
+      for metric, metric_fn in six.iteritems(metrics):
         overload_eval_metric_name = getattr(
             model_hparams, "overload_eval_metric_name", None)
         if len(problems) == 1 and overload_eval_metric_name:
@@ -642,9 +635,10 @@ def weights_fn_for_mp(problem_task_id):
 
 def create_eager_metrics_for_problem(problem, model_hparams):
   """See create_eager_metrics."""
-  metric_names = problem.eval_metrics()
+  metric_fns = problem.eval_metric_fns(model_hparams)
   tm = problem.get_hparams(model_hparams).modality["targets"]
-  return create_eager_metrics(metric_names, weights_fn=tm.targets_weights_fn)
+  return create_eager_metrics_internal(
+        metric_fns, weights_fn=tm.targets_weights_fn)
 
 
 def create_eager_metrics(metric_names, weights_fn=common_layers.weights_all):
@@ -662,9 +656,26 @@ def create_eager_metrics(metric_names, weights_fn=common_layers.weights_all):
   """
   metric_fns = dict(
       [(name, METRICS_FNS[name]) for name in metric_names])
+  return create_eager_metrics_internal(metric_fns, weights_fn)
+
+
+def create_eager_metrics_internal(metric_fns,
+                                  weights_fn=common_layers.weights_all):
+  """Create metrics accumulators and averager for Eager mode.
+
+  Args:
+    metric_names: dict<metric name, metric function>
+    weights_fn: function that takes labels and returns a weights mask. Defaults
+      to weights of all 1, i.e. common_layers.weights_all. Use
+      common_layers.weights_nonzero if labels have 0-padding.
+
+  Returns:
+    (accum_fn(predictions, targets) => None,
+     result_fn() => dict<str metric_name, float avg_val>
+  """
   tfe_metrics = dict()
 
-  for name in metric_names:
+  for name in metric_fns:
     tfe_metrics[name] = tfe.metrics.Mean(name=name)
 
   def metric_accum(predictions, targets):
@@ -675,7 +686,7 @@ def metric_accum(predictions, targets):
 
   def metric_means():
     avgs = {}
-    for name in metric_names:
+    for name in metric_fns:
       avgs[name] = tfe_metrics[name].result().numpy()
     return avgs
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 4d99b9f02..f3530b574 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1721,7 +1721,7 @@ def create_tpu_eval_metrics_fn(problem, model_hparams):
   """Create the metrics_fn that TPUEstimatorSpec expects."""
 
   metric_fns = []
-  eval_metrics = problem.eval_metrics()
+  eval_metrics = problem.eval_metric_fns(model_hparams)
 
   tm = _create_target_modality(problem.get_hparams(model_hparams).modality)
   if isinstance(tm, dict):
@@ -1739,12 +1739,12 @@ def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
 
         return wrapped_metric_fn
 
-      for metric in eval_metrics:
+      for metric, metric_fn in six.iteritems(eval_metrics):
         if metric in TPU_METRIC_BLACKLIST:
           log_warn("Skipping eval metric %s in TPU_METRIC_BLACKLIST", metric)
           continue
         name = "%s/metrics-%s/%s" % (k, problem.name, metric)
-        metric_fns.append((name, make_metric_fn(metrics.METRICS_FNS[metric])))
+        metric_fns.append((name, make_metric_fn(metric_fn)))
   else:
     weights_fn = tm.targets_weights_fn
 
@@ -1759,12 +1759,12 @@ def wrapped_metric_fn(logits, labels, features):
 
       return wrapped_metric_fn
 
-    for metric in eval_metrics:
+    for metric, metric_fn in six.iteritems(eval_metrics):
       if metric in TPU_METRIC_BLACKLIST:
         log_warn("Skipping eval metric %s in TPU_METRIC_BLACKLIST", metric)
         continue
       name = "metrics-%s/%s" % (problem.name, metric)
-      metric_fns.append((name, make_metric_fn(metrics.METRICS_FNS[metric])))
+      metric_fns.append((name, make_metric_fn(metric_fn)))
 
   def all_metrics_fn(**kwargs):
     """Construct metrics dictionary."""

From 9f267c38f99e4996c76cdeec0c9f19e9f18d5b18 Mon Sep 17 00:00:00 2001
From: tober <kiyoyamakouhei@gmail.com>
Date: Sat, 5 Jan 2019 07:24:04 +0900
Subject: [PATCH 1434/2720] Feature/new model for scalar regression. In the
 models/transformer.py, a model called Transformer Regressor was created.
 (#1332)

* A model called Transformer Encoder was created in the model.

* solved an inheriting issue.
---
 tensor2tensor/models/transformer.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 5184ef953..19a0618de 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1121,6 +1121,22 @@ def body(self, features):
 
     return encoder_output
 
+@registry.register_model
+class TransformerRegressor(TransformerEncoder):
+  """Transformer inheriting from Encoder, for the regression problem.
+  Final res is a tensor that has a shape of (?, 1, 1, 1)
+  """
+
+  def top(self, body_output, features):
+    """Computes single scalar value from body_output
+    """
+    with tf.variable_scope("reg_top_ffn"):
+      # scalar = common_layers.dense(body_output,hparams)
+      x = body_output
+      x = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
+      res = tf.layers.dense(x, 1, name="model_top")
+      return res
+
 
 def features_to_nonpadding(features, inputs_or_targets="inputs"):
   key = inputs_or_targets + "_segmentation"

From 2fe09661ab837eac05e1b1d012aa4adccc64e99b Mon Sep 17 00:00:00 2001
From: Youngwook Kim <youngwook.kim@gmail.com>
Date: Fri, 4 Jan 2019 14:23:45 -0800
Subject: [PATCH 1435/2720] internal merge of PR #1336

PiperOrigin-RevId: 227913649
---
 tensor2tensor/data_generators/problem.py |  1 +
 tensor2tensor/models/transformer.py      | 16 ----------------
 tensor2tensor/utils/metrics.py           |  4 ++--
 3 files changed, 3 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index b89aac526..219deabdf 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -368,6 +368,7 @@ def eval_metrics(self):
     ]
 
   def eval_metric_fns(self, model_hparams):
+    del model_hparams
     metric_names = self.eval_metrics()
     if not all([m in metrics.METRICS_FNS for m in metric_names]):
       error_str = ("Unrecognized metric. Problem %s specified metrics "
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 19a0618de..5184ef953 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1121,22 +1121,6 @@ def body(self, features):
 
     return encoder_output
 
-@registry.register_model
-class TransformerRegressor(TransformerEncoder):
-  """Transformer inheriting from Encoder, for the regression problem.
-  Final res is a tensor that has a shape of (?, 1, 1, 1)
-  """
-
-  def top(self, body_output, features):
-    """Computes single scalar value from body_output
-    """
-    with tf.variable_scope("reg_top_ffn"):
-      # scalar = common_layers.dense(body_output,hparams)
-      x = body_output
-      x = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
-      res = tf.layers.dense(x, 1, name="model_top")
-      return res
-
 
 def features_to_nonpadding(features, inputs_or_targets="inputs"):
   key = inputs_or_targets + "_segmentation"
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 5608c0565..feea42058 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -638,7 +638,7 @@ def create_eager_metrics_for_problem(problem, model_hparams):
   metric_fns = problem.eval_metric_fns(model_hparams)
   tm = problem.get_hparams(model_hparams).modality["targets"]
   return create_eager_metrics_internal(
-        metric_fns, weights_fn=tm.targets_weights_fn)
+      metric_fns, weights_fn=tm.targets_weights_fn)
 
 
 def create_eager_metrics(metric_names, weights_fn=common_layers.weights_all):
@@ -664,7 +664,7 @@ def create_eager_metrics_internal(metric_fns,
   """Create metrics accumulators and averager for Eager mode.
 
   Args:
-    metric_names: dict<metric name, metric function>
+    metric_fns: dict<metric name, metric function>
     weights_fn: function that takes labels and returns a weights mask. Defaults
       to weights of all 1, i.e. common_layers.weights_all. Use
       common_layers.weights_nonzero if labels have 0-padding.

From da8116220f4cb42fdc1507d2f0f43a3eb0614bfd Mon Sep 17 00:00:00 2001
From: tober <kiyoyamakouhei@gmail.com>
Date: Fri, 4 Jan 2019 14:24:25 -0800
Subject: [PATCH 1436/2720] internal merge of PR #1332

PiperOrigin-RevId: 227913757
---
 tensor2tensor/models/transformer.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 5184ef953..692f88afb 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1122,6 +1122,23 @@ def body(self, features):
     return encoder_output
 
 
+@registry.register_model
+class TransformerRegressor(TransformerEncoder):
+  """Transformer inheriting from Encoder, for the regression problem.
+
+  Final result is a tensor that has a shape of (?, 1, 1, 1).
+  """
+
+  def top(self, body_output, features):
+    """Computes single scalar value from body_output."""
+
+    with tf.variable_scope("reg_top_ffn"):
+      x = body_output
+      x = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
+      res = tf.layers.dense(x, 1, name="model_top")
+      return res
+
+
 def features_to_nonpadding(features, inputs_or_targets="inputs"):
   key = inputs_or_targets + "_segmentation"
   if features and key in features:

From 89052b095b1d5554af75ea8511a42eff4a42069d Mon Sep 17 00:00:00 2001
From: cclauss <cclauss@me.com>
Date: Fri, 4 Jan 2019 23:55:55 +0100
Subject: [PATCH 1437/2720] Add to_unicode_utf8() to text_encoder.py (#1321)

---
 tensor2tensor/data_generators/cnn_dailymail.py       |  6 +-----
 tensor2tensor/data_generators/cola.py                |  6 +-----
 tensor2tensor/data_generators/mrpc.py                |  6 +-----
 tensor2tensor/data_generators/multinli.py            |  6 +-----
 tensor2tensor/data_generators/qnli.py                |  6 +-----
 tensor2tensor/data_generators/quora_qpairs.py        |  6 +-----
 tensor2tensor/data_generators/rte.py                 |  6 +-----
 tensor2tensor/data_generators/scitail.py             |  6 +-----
 tensor2tensor/data_generators/sst_binary.py          |  6 +-----
 tensor2tensor/data_generators/stanford_nli.py        |  6 +-----
 tensor2tensor/data_generators/text_encoder.py        |  4 ++++
 tensor2tensor/data_generators/wiki_revision_utils.py | 10 ++--------
 tensor2tensor/data_generators/wnli.py                |  6 +-----
 13 files changed, 17 insertions(+), 63 deletions(-)

diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index 5b86c7ad8..71599b065 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -24,7 +24,6 @@
 import os
 import random
 import tarfile
-import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -157,10 +156,7 @@ def fix_run_on_sents(line):
     summary = []
     reading_highlights = False
     for line in tf.gfile.Open(story_file, "rb"):
-      if six.PY2:
-        line = unicode(line.strip(), "utf-8")
-      else:
-        line = line.strip().decode("utf-8")
+      line = text_encoder.to_unicode_utf8(line.strip())
       line = fix_run_on_sents(line)
       if not line:
         continue
diff --git a/tensor2tensor/data_generators/cola.py b/tensor2tensor/data_generators/cola.py
index 7a905573c..0f2748fea 100644
--- a/tensor2tensor/data_generators/cola.py
+++ b/tensor2tensor/data_generators/cola.py
@@ -21,7 +21,6 @@
 
 import os
 import zipfile
-import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -83,10 +82,7 @@ def _maybe_download_corpora(self, tmp_dir):
 
   def example_generator(self, filename):
     for line in tf.gfile.Open(filename, "rb"):
-      if six.PY2:
-        line = unicode(line.strip(), "utf-8")
-      else:
-        line = line.strip().decode("utf-8")
+      line = text_encoder.to_unicode_utf8(line.strip())
       _, label, _, sent = line.split("\t")
       yield {
           "inputs": sent,
diff --git a/tensor2tensor/data_generators/mrpc.py b/tensor2tensor/data_generators/mrpc.py
index e8c9e3a39..47c8364d2 100644
--- a/tensor2tensor/data_generators/mrpc.py
+++ b/tensor2tensor/data_generators/mrpc.py
@@ -20,7 +20,6 @@
 from __future__ import print_function
 
 import os
-import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -95,10 +94,7 @@ def download_file(tdir, filepath, url):
   def example_generator(self, filename, dev_ids, dataset_split):
     for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
       if idx == 0: continue  # skip header
-      if six.PY2:
-        line = unicode(line.strip(), "utf-8")
-      else:
-        line = line.strip().decode("utf-8")
+      line = text_encoder.to_unicode_utf8(line.strip())
       l, id1, id2, s1, s2 = line.split("\t")
       is_dev = [id1, id2] in dev_ids
       if dataset_split == problem.DatasetSplit.TRAIN and is_dev:
diff --git a/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index 70ee0107f..e3af79f0e 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -21,7 +21,6 @@
 
 import os
 import zipfile
-import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import lm1b
 from tensor2tensor.data_generators import problem
@@ -87,10 +86,7 @@ def example_generator(self, filename):
     label_list = self.class_labels(data_dir=None)
     for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
       if idx == 0: continue  # skip header
-      if six.PY2:
-        line = unicode(line.strip(), "utf-8")
-      else:
-        line = line.strip().decode("utf-8")
+      line = text_encoder.to_unicode_utf8(line.strip())
       split_line = line.split("\t")
       # Works for both splits even though dev has some extra human labels.
       s1, s2 = split_line[8:10]
diff --git a/tensor2tensor/data_generators/qnli.py b/tensor2tensor/data_generators/qnli.py
index b59db970f..9eeeb2077 100644
--- a/tensor2tensor/data_generators/qnli.py
+++ b/tensor2tensor/data_generators/qnli.py
@@ -21,7 +21,6 @@
 
 import os
 import zipfile
-import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -85,10 +84,7 @@ def example_generator(self, filename):
     label_list = self.class_labels(data_dir=None)
     for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
       if idx == 0: continue  # skip header
-      if six.PY2:
-        line = unicode(line.strip(), "utf-8")
-      else:
-        line = line.strip().decode("utf-8")
+      line = text_encoder.to_unicode_utf8(line.strip())
       _, s1, s2, l = line.split("\t")
       inputs = [s1, s2]
       l = label_list.index(l)
diff --git a/tensor2tensor/data_generators/quora_qpairs.py b/tensor2tensor/data_generators/quora_qpairs.py
index 5960c2488..21f3702ec 100644
--- a/tensor2tensor/data_generators/quora_qpairs.py
+++ b/tensor2tensor/data_generators/quora_qpairs.py
@@ -21,7 +21,6 @@
 
 import os
 import zipfile
-import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -84,10 +83,7 @@ def example_generator(self, filename):
     skipped = 0
     for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
       if idx == 0: continue  # skip header
-      if six.PY2:
-        line = unicode(line.strip(), "utf-8")
-      else:
-        line = line.strip().decode("utf-8")
+      line = text_encoder.to_unicode_utf8(line.strip())
       split_line = line.split("\t")
       if len(split_line) < 6:
         skipped += 1
diff --git a/tensor2tensor/data_generators/rte.py b/tensor2tensor/data_generators/rte.py
index af7fa41e9..2eff16422 100644
--- a/tensor2tensor/data_generators/rte.py
+++ b/tensor2tensor/data_generators/rte.py
@@ -21,7 +21,6 @@
 
 import os
 import zipfile
-import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -85,10 +84,7 @@ def example_generator(self, filename):
     label_list = self.class_labels(data_dir=None)
     for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
       if idx == 0: continue  # skip header
-      if six.PY2:
-        line = unicode(line.strip(), "utf-8")
-      else:
-        line = line.strip().decode("utf-8")
+      line = text_encoder.to_unicode_utf8(line.strip())
       _, s1, s2, l = line.split("\t")
       inputs = [s1, s2]
       l = label_list.index(l)
diff --git a/tensor2tensor/data_generators/scitail.py b/tensor2tensor/data_generators/scitail.py
index 90df97ccb..f600fa560 100644
--- a/tensor2tensor/data_generators/scitail.py
+++ b/tensor2tensor/data_generators/scitail.py
@@ -21,7 +21,6 @@
 
 import os
 import zipfile
-import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import lm1b
 from tensor2tensor.data_generators import problem
@@ -83,10 +82,7 @@ def _maybe_download_corpora(self, tmp_dir):
   def example_generator(self, filename):
     label_list = self.class_labels(data_dir=None)
     for line in tf.gfile.Open(filename, "rb"):
-      if six.PY2:
-        line = unicode(line.strip(), "utf-8")
-      else:
-        line = line.strip().decode("utf-8")
+      line = text_encoder.to_unicode_utf8(line.strip())
       split_line = line.split("\t")
       s1, s2 = split_line[:2]
       l = label_list.index(split_line[2])
diff --git a/tensor2tensor/data_generators/sst_binary.py b/tensor2tensor/data_generators/sst_binary.py
index a8a391c95..9081fc81a 100644
--- a/tensor2tensor/data_generators/sst_binary.py
+++ b/tensor2tensor/data_generators/sst_binary.py
@@ -21,7 +21,6 @@
 
 import os
 import zipfile
-import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -84,10 +83,7 @@ def _maybe_download_corpora(self, tmp_dir):
   def example_generator(self, filename):
     for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
       if idx == 0: continue  # skip header
-      if six.PY2:
-        line = unicode(line.strip(), "utf-8")
-      else:
-        line = line.strip().decode("utf-8")
+      line = text_encoder.to_unicode_utf8(line.strip())
       sent, label = line.split("\t")
       yield {
           "inputs": sent,
diff --git a/tensor2tensor/data_generators/stanford_nli.py b/tensor2tensor/data_generators/stanford_nli.py
index a8aa04602..9c99501ec 100644
--- a/tensor2tensor/data_generators/stanford_nli.py
+++ b/tensor2tensor/data_generators/stanford_nli.py
@@ -21,7 +21,6 @@
 
 import os
 import zipfile
-import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import lm1b
 from tensor2tensor.data_generators import problem
@@ -84,10 +83,7 @@ def example_generator(self, filename):
     label_list = self.class_labels(data_dir=None)
     for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
       if idx == 0: continue  # skip header
-      if six.PY2:
-        line = unicode(line.strip(), "utf-8")
-      else:
-        line = line.strip().decode("utf-8")
+      line = text_encoder.to_unicode_utf8(line.strip())
       split_line = line.split("\t")
       # Works for both splits even though dev has some extra human labels.
       s1, s2 = split_line[5:7]
diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 5580a2f22..3bfa01c9c 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -98,6 +98,10 @@ def to_unicode_ignore_errors(s):
   return to_unicode(s, ignore_errors=True)
 
 
+def to_unicode_utf8(s):
+  return unicode(s, "utf-8") if six.PY2 else s.decode("utf-8")
+
+
 def strip_ids(ids, ids_to_strip):
   """Strip ids_to_strip from the end ids."""
   ids = list(ids)
diff --git a/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index 027bd162b..9704c068c 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -27,18 +27,12 @@
 import re
 import subprocess
 
-import six
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import text_encoder
 
 import tensorflow as tf
 
 
-def to_unicode(s):
-  return unicode(s, "utf-8") if six.PY2 else s.decode("utf-8")
-
-
 def include_revision(revision_num, skip_factor=1.1):
   """Decide whether to include a revision.
 
@@ -118,7 +112,7 @@ def get_title(page):
   assert start_pos != -1
   assert end_pos != -1
   start_pos += len("<title>")
-  return to_unicode(page[start_pos:end_pos])
+  return text_encoder.to_unicode_utf8(page[start_pos:end_pos])
 
 
 def get_id(page):
@@ -257,7 +251,7 @@ def get_text(revision, strip=True):
     ret = revision[end_tag_pos:end_pos]
   if strip:
     ret = strip_text(ret)
-  ret = to_unicode(ret)
+  ret = text_encoder.to_unicode_utf8(ret)
   return ret
 
 
diff --git a/tensor2tensor/data_generators/wnli.py b/tensor2tensor/data_generators/wnli.py
index cd4de046d..9b94b5b43 100644
--- a/tensor2tensor/data_generators/wnli.py
+++ b/tensor2tensor/data_generators/wnli.py
@@ -21,7 +21,6 @@
 
 import os
 import zipfile
-import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -88,10 +87,7 @@ def _maybe_download_corpora(self, tmp_dir):
   def example_generator(self, filename):
     for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
       if idx == 0: continue  # skip header
-      if six.PY2:
-        line = unicode(line.strip(), "utf-8")
-      else:
-        line = line.strip().decode("utf-8")
+      line = text_encoder.to_unicode_utf8(line.strip())
       _, s1, s2, l = line.split("\t")
       inputs = [s1, s2]
       yield {

From ef12bee72270b322165d073c39a650a189de39aa Mon Sep 17 00:00:00 2001
From: cclauss <cclauss@me.com>
Date: Fri, 4 Jan 2019 23:56:15 +0100
Subject: [PATCH 1438/2720] Simplify is_unicode() (#1318)

https://six.readthedocs.io/#six.text_type
---
 tensor2tensor/data_generators/text_encoder.py | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 3bfa01c9c..139fceb82 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -78,13 +78,7 @@ def unicode_to_native(s):
 
 
 def is_unicode(s):
-  if six.PY2:
-    if isinstance(s, unicode):
-      return True
-  else:
-    if isinstance(s, str):
-      return True
-  return False
+  return isinstance(s, six.text_type)
 
 
 def to_unicode(s, ignore_errors=False):

From 3740f2eee4ef3b2670e9d9aeb7a13bb40764a9b9 Mon Sep 17 00:00:00 2001
From: cclauss <cclauss@me.com>
Date: Fri, 4 Jan 2019 14:56:49 -0800
Subject: [PATCH 1439/2720] internal merge of PR #1321

PiperOrigin-RevId: 227918987
---
 tensor2tensor/data_generators/text_encoder.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 139fceb82..3bfa01c9c 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -78,7 +78,13 @@ def unicode_to_native(s):
 
 
 def is_unicode(s):
-  return isinstance(s, six.text_type)
+  if six.PY2:
+    if isinstance(s, unicode):
+      return True
+  else:
+    if isinstance(s, str):
+      return True
+  return False
 
 
 def to_unicode(s, ignore_errors=False):

From 7857ba76be6fe1c56b4d7a0648818eb600bc8731 Mon Sep 17 00:00:00 2001
From: cclauss <cclauss@me.com>
Date: Fri, 4 Jan 2019 15:09:50 -0800
Subject: [PATCH 1440/2720] internal merge of PR #1318

PiperOrigin-RevId: 227921102
---
 tensor2tensor/data_generators/text_encoder.py | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 3bfa01c9c..139fceb82 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -78,13 +78,7 @@ def unicode_to_native(s):
 
 
 def is_unicode(s):
-  if six.PY2:
-    if isinstance(s, unicode):
-      return True
-  else:
-    if isinstance(s, str):
-      return True
-  return False
+  return isinstance(s, six.text_type)
 
 
 def to_unicode(s, ignore_errors=False):

From 5dfb2eb3359200f4c846a3ffd9f43a33fb7fb617 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 4 Jan 2019 15:42:29 -0800
Subject: [PATCH 1441/2720] Fix oss_tests.sh

PiperOrigin-RevId: 227925959
---
 oss_scripts/oss_tests.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index bf221ca74..6297d9f5d 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -74,7 +74,7 @@ then
   # Ignores:
   # * Glow requires the CIFAR-10 dataset to be generated
   pytest tensor2tensor/models/research \
-    --ignore=tensor2tensor/models/research/glow_test.py \
+    --ignore=tensor2tensor/models/research/glow_test.py
   set_status
 fi
 

From 98ec1ee12a9e412188f8d4602c748c8aa4016efe Mon Sep 17 00:00:00 2001
From: Chan Yu <aeloyq@outlook.com>
Date: Sat, 5 Jan 2019 07:56:58 +0800
Subject: [PATCH 1442/2720] Enable Relative Dot Product Visualization (#1303)

* add caching mechanism support for fast decoding with relative_dot_product in transformer model

* fix typo

* enable visualization when use dot_product_relative in self-attention

* clean code
---
 tensor2tensor/layers/common_attention.py     | 10 +++++++++-
 tensor2tensor/visualization/visualization.py | 11 +++++++----
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 01ac52b9e..a36aa9946 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1549,6 +1549,7 @@ def dot_product_attention_relative(q,
                                    max_relative_position,
                                    dropout_rate=0.0,
                                    image_shapes=None,
+                                   save_weights_to=None,
                                    name=None,
                                    make_image_summary=True,
                                    cache=False):
@@ -1566,6 +1567,9 @@ def dot_product_attention_relative(q,
         inputs that unique position embeddings should be learned for.
     dropout_rate: a floating point number.
     image_shapes: optional tuple of integer scalars.
+    save_weights_to: an optional dictionary to capture attention weights
+      for visualization; the weights tensor will be appended there under
+      a string key created from the variable scope (including name).
     name: an optional string.
     make_image_summary: Whether to make an attention image summary.
     cache: whether use cache mode
@@ -1580,7 +1584,7 @@ def dot_product_attention_relative(q,
     raise ValueError("Max relative position (%s) should be > 0 when using "
                      "relative self attention." % (max_relative_position))
   with tf.variable_scope(
-      name, default_name="dot_product_attention_relative", values=[q, k, v]):
+      name, default_name="dot_product_attention_relative", values=[q, k, v]) as scope:
 
     # This calculation only works for self attention.
     # q, k and v must therefore have the same shape.
@@ -1603,6 +1607,9 @@ def dot_product_attention_relative(q,
     if bias is not None:
       logits += bias
     weights = tf.nn.softmax(logits, name="attention_weights")
+    if save_weights_to is not None:
+      save_weights_to[scope.name] = weights
+      save_weights_to[scope.name + "/logits"] = logits
     weights = tf.nn.dropout(weights, 1.0 - dropout_rate)
     if not tf.get_variable_scope().reuse and make_image_summary:
       attention_image_summary(weights, image_shapes)
@@ -3466,6 +3473,7 @@ def multihead_attention(query_antecedent,
           max_relative_position,
           dropout_rate,
           image_shapes,
+          save_weights_to=save_weights_to,
           make_image_summary=make_image_summary,
           cache=cache is not None)
     elif attention_type == "dot_product_unmasked_relative_v2":
diff --git a/tensor2tensor/visualization/visualization.py b/tensor2tensor/visualization/visualization.py
index f21621a25..7b511b42f 100644
--- a/tensor2tensor/visualization/visualization.py
+++ b/tensor2tensor/visualization/visualization.py
@@ -183,15 +183,18 @@ def get_att_mats(translate_model):
   encdec_atts = []
 
   prefix = 'transformer/body/'
-  postfix = '/multihead_attention/dot_product_attention'
+  postfix_self_attention = '/multihead_attention/dot_product_attention'
+  if translate_model.hparams.self_attention_type == "dot_product_relative":
+    postfix_self_attention = '/multihead_attention/dot_product_attention_relative'
+  postfix_encdec = '/multihead_attention/dot_product_attention'
 
   for i in range(translate_model.hparams.num_hidden_layers):
     enc_att = translate_model.attention_weights[
-        '%sencoder/layer_%i/self_attention%s' % (prefix, i, postfix)]
+        '%sencoder/layer_%i/self_attention%s' % (prefix, i, postfix_self_attention)]
     dec_att = translate_model.attention_weights[
-        '%sdecoder/layer_%i/self_attention%s' % (prefix, i, postfix)]
+        '%sdecoder/layer_%i/self_attention%s' % (prefix, i, postfix_self_attention)]
     encdec_att = translate_model.attention_weights[
-        '%sdecoder/layer_%i/encdec_attention%s' % (prefix, i, postfix)]
+        '%sdecoder/layer_%i/encdec_attention%s' % (prefix, i, postfix_encdec)]
     enc_atts.append(enc_att)
     dec_atts.append(dec_att)
     encdec_atts.append(encdec_att)

From 01d76fad4a535bd38803e13a2a8bb002c24bf847 Mon Sep 17 00:00:00 2001
From: Chan Yu <aeloyq@outlook.com>
Date: Fri, 4 Jan 2019 15:57:29 -0800
Subject: [PATCH 1443/2720] internal merge of PR #1303

PiperOrigin-RevId: 227927931
---
 tensor2tensor/layers/common_attention.py     |  3 +-
 tensor2tensor/visualization/visualization.py | 35 +++++++++++---------
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index a36aa9946..dd192e172 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1584,7 +1584,8 @@ def dot_product_attention_relative(q,
     raise ValueError("Max relative position (%s) should be > 0 when using "
                      "relative self attention." % (max_relative_position))
   with tf.variable_scope(
-      name, default_name="dot_product_attention_relative", values=[q, k, v]) as scope:
+      name, default_name="dot_product_attention_relative",
+      values=[q, k, v]) as scope:
 
     # This calculation only works for self attention.
     # q, k and v must therefore have the same shape.
diff --git a/tensor2tensor/visualization/visualization.py b/tensor2tensor/visualization/visualization.py
index 7b511b42f..c70590d98 100644
--- a/tensor2tensor/visualization/visualization.py
+++ b/tensor2tensor/visualization/visualization.py
@@ -51,19 +51,19 @@ def __init__(
 
   def encode(self, input_str):
     """Input str to features dict, ready for inference."""
-    inputs = self.encoders['inputs'].encode(input_str) + [EOS_ID]
+    inputs = self.encoders["inputs"].encode(input_str) + [EOS_ID]
     batch_inputs = np.reshape(inputs, [1, -1, 1, 1])  # Make it 3D.
     return batch_inputs
 
   def decode(self, integers):
     """List of ints to str."""
     integers = list(np.squeeze(integers))
-    return self.encoders['inputs'].decode(integers)
+    return self.encoders["inputs"].decode(integers)
 
   def decode_list(self, integers):
     """List of ints to list of str."""
     integers = list(np.squeeze(integers))
-    return self.encoders['inputs'].decode_list(integers)
+    return self.encoders["inputs"].decode_list(integers)
 
   def get_vis_data_from_string(self, sess, input_string):
     """Constructs the data needed for visualizing attentions.
@@ -135,11 +135,11 @@ def build_model(hparams_set, model_name, data_dir, problem_name, beam_size=1):
   translate_model = registry.model(model_name)(
       hparams, tf.estimator.ModeKeys.EVAL)
 
-  inputs = tf.placeholder(tf.int32, shape=(1, None, 1, 1), name='inputs')
-  targets = tf.placeholder(tf.int32, shape=(1, None, 1, 1), name='targets')
+  inputs = tf.placeholder(tf.int32, shape=(1, None, 1, 1), name="inputs")
+  targets = tf.placeholder(tf.int32, shape=(1, None, 1, 1), name="targets")
   translate_model({
-      'inputs': inputs,
-      'targets': targets,
+      "inputs": inputs,
+      "targets": targets,
   })
 
   # Must be called after building the training graph, so that the dict will
@@ -150,8 +150,8 @@ def build_model(hparams_set, model_name, data_dir, problem_name, beam_size=1):
 
   with tf.variable_scope(tf.get_variable_scope(), reuse=True):
     samples = translate_model.infer({
-        'inputs': inputs,
-    }, beam_size=beam_size)['outputs']
+        "inputs": inputs,
+    }, beam_size=beam_size)["outputs"]
 
   return inputs, targets, samples, att_mats
 
@@ -182,19 +182,22 @@ def get_att_mats(translate_model):
   dec_atts = []
   encdec_atts = []
 
-  prefix = 'transformer/body/'
-  postfix_self_attention = '/multihead_attention/dot_product_attention'
+  prefix = "transformer/body/"
+  postfix_self_attention = "/multihead_attention/dot_product_attention"
   if translate_model.hparams.self_attention_type == "dot_product_relative":
-    postfix_self_attention = '/multihead_attention/dot_product_attention_relative'
-  postfix_encdec = '/multihead_attention/dot_product_attention'
+    postfix_self_attention = ("/multihead_attention/"
+                              "dot_product_attention_relative")
+  postfix_encdec = "/multihead_attention/dot_product_attention"
 
   for i in range(translate_model.hparams.num_hidden_layers):
     enc_att = translate_model.attention_weights[
-        '%sencoder/layer_%i/self_attention%s' % (prefix, i, postfix_self_attention)]
+        "%sencoder/layer_%i/self_attention%s"
+        % (prefix, i, postfix_self_attention)]
     dec_att = translate_model.attention_weights[
-        '%sdecoder/layer_%i/self_attention%s' % (prefix, i, postfix_self_attention)]
+        "%sdecoder/layer_%i/self_attention%s"
+        % (prefix, i, postfix_self_attention)]
     encdec_att = translate_model.attention_weights[
-        '%sdecoder/layer_%i/encdec_attention%s' % (prefix, i, postfix_encdec)]
+        "%sdecoder/layer_%i/encdec_attention%s" % (prefix, i, postfix_encdec)]
     enc_atts.append(enc_att)
     dec_atts.append(dec_att)
     encdec_atts.append(encdec_att)

From 05c972af8e0f8d656cf0fed910839e00a050b2f5 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Mon, 7 Jan 2019 18:54:47 +0000
Subject: [PATCH 1444/2720] Replace deprecated np.random.random_integers with
 np.random.randint (#1346)

---
 tensor2tensor/layers/common_layers_test.py       |  4 ++--
 tensor2tensor/layers/modalities_test.py          | 10 +++++-----
 tensor2tensor/models/basic_test.py               |  4 ++--
 tensor2tensor/models/bytenet_test.py             |  4 ++--
 .../models/image_transformer_2d_test.py          |  8 ++++----
 tensor2tensor/models/image_transformer_test.py   |  4 ++--
 tensor2tensor/models/lstm_test.py                | 16 ++++++++--------
 .../models/mtf_image_transformer_test.py         |  2 +-
 tensor2tensor/models/mtf_transformer_test.py     |  4 ++--
 tensor2tensor/models/neural_gpu_test.py          |  4 ++--
 .../models/research/autoencoders_test.py         |  4 ++--
 .../models/research/gene_expression_test.py      |  4 ++--
 .../models/research/transformer_aux_test.py      |  4 ++--
 .../models/research/transformer_revnet_test.py   |  4 ++--
 .../models/research/transformer_vae_test.py      |  4 ++--
 .../research/universal_transformer_test.py       |  4 ++--
 .../models/research/vqa_attention_test.py        |  8 ++++----
 tensor2tensor/models/resnet_test.py              |  8 ++++----
 tensor2tensor/models/slicenet_test.py            |  4 ++--
 tensor2tensor/models/transformer_test.py         |  4 ++--
 tensor2tensor/models/xception_test.py            |  8 ++++----
 21 files changed, 58 insertions(+), 58 deletions(-)

diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index 7d7c95c22..bc46f0409 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -48,7 +48,7 @@ def testSaturatingSigmoid(self):
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testFlatten4D3D(self):
-    x = np.random.random_integers(1, high=8, size=(3, 5, 2))
+    x = np.random.randint(1, high=9, size=(3, 5, 2))
     y = common_layers.flatten4d3d(common_layers.embedding(x, 10, 7))
     self.evaluate(tf.global_variables_initializer())
     res = self.evaluate(y)
@@ -56,7 +56,7 @@ def testFlatten4D3D(self):
 
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testEmbedding(self):
-    x = np.random.random_integers(1, high=8, size=(3, 5))
+    x = np.random.randint(1, high=9, size=(3, 5))
     y = common_layers.embedding(x, 10, 16)
     self.evaluate(tf.global_variables_initializer())
     res = self.evaluate(y)
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index f1aa4d134..71ba33b86 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -39,7 +39,7 @@ def testSymbolModalityInputs(self):
     model_hparams = common_hparams.basic_params1()
     model_hparams.hidden_size = hidden_size
     model_hparams.mode = tf.estimator.ModeKeys.TRAIN
-    x = -1 + np.random.random_integers(
+    x = np.random.randint(
         vocab_size, size=(batch_size, length, 1, 1))
     m = modalities.SymbolModality(model_hparams, vocab_size)
     data_parallelism = expert_utils.Parallelism(
@@ -62,9 +62,9 @@ def testSymbolModalityTargets(self):
     model_hparams = common_hparams.basic_params1()
     model_hparams.hidden_size = hidden_size
     model_hparams.mode = tf.estimator.ModeKeys.TRAIN
-    body_output = -1 + np.random.random_integers(
+    body_output = np.random.randint(
         100, size=(batch_size, length, height, hidden_size))
-    targets = -1 + np.random.random_integers(
+    targets = np.random.randint(
         vocab_size, size=(batch_size, length, height, 1))
     m = modalities.SymbolModality(model_hparams, vocab_size)
     data_parallelism = expert_utils.Parallelism(
@@ -92,9 +92,9 @@ def testSymbolModalityTargetsFactored(self):
     model_hparams.factored_logits = True
     model_hparams.hidden_size = hidden_size
     model_hparams.mode = tf.estimator.ModeKeys.TRAIN
-    body_output = -1 + np.random.random_integers(
+    body_output = np.random.randint(
         100, size=(batch_size, length, height, hidden_size))
-    targets = -1 + np.random.random_integers(
+    targets = np.random.randint(
         vocab_size, size=(batch_size, length, height, 1))
     m = modalities.SymbolModality(model_hparams, vocab_size)
     data_parallelism = expert_utils.Parallelism(
diff --git a/tensor2tensor/models/basic_test.py b/tensor2tensor/models/basic_test.py
index 829c802bf..eaef58326 100644
--- a/tensor2tensor/models/basic_test.py
+++ b/tensor2tensor/models/basic_test.py
@@ -30,8 +30,8 @@
 class BasicTest(tf.test.TestCase):
 
   def testBasicFcRelu(self):
-    x = np.random.random_integers(0, high=255, size=(1, 28, 28, 1))
-    y = np.random.random_integers(0, high=9, size=(1, 1))
+    x = np.random.randint(256, size=(1, 28, 28, 1))
+    y = np.random.randint(10, size=(1, 1))
     hparams = trainer_lib.create_hparams(
         "basic_fc_small", problem_name="image_mnist", data_dir=".")
     with self.test_session() as session:
diff --git a/tensor2tensor/models/bytenet_test.py b/tensor2tensor/models/bytenet_test.py
index b87d75d5d..e68c10b98 100644
--- a/tensor2tensor/models/bytenet_test.py
+++ b/tensor2tensor/models/bytenet_test.py
@@ -30,8 +30,8 @@ class ByteNetTest(tf.test.TestCase):
 
   def testByteNet(self):
     vocab_size = 9
-    x = np.random.random_integers(1, high=vocab_size - 1, size=(3, 5, 1, 1))
-    y = np.random.random_integers(1, high=vocab_size - 1, size=(3, 6, 1, 1))
+    x = np.random.randint(1, high=vocab_size, size=(3, 5, 1, 1))
+    y = np.random.randint(1, high=vocab_size, size=(3, 6, 1, 1))
     hparams = bytenet.bytenet_base()
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
diff --git a/tensor2tensor/models/image_transformer_2d_test.py b/tensor2tensor/models/image_transformer_2d_test.py
index d5327de78..9beaee900 100644
--- a/tensor2tensor/models/image_transformer_2d_test.py
+++ b/tensor2tensor/models/image_transformer_2d_test.py
@@ -35,8 +35,8 @@ def _test_img2img_transformer(self, net):
     hparams = image_transformer_2d.img2img_transformer2d_tiny()
     hparams.data_dir = ""
     p_hparams = registry.problem("image_celeba").get_hparams(hparams)
-    inputs = np.random.random_integers(0, high=255, size=(3, 4, 4, 3))
-    targets = np.random.random_integers(0, high=255, size=(3, 8, 8, 3))
+    inputs = np.random.randint(256, size=(3, 4, 4, 3))
+    targets = np.random.randint(256, size=(3, 8, 8, 3))
     with self.test_session() as session:
       features = {
           "inputs": tf.constant(inputs, dtype=tf.int32),
@@ -63,9 +63,9 @@ def _test_imagetransformer_2d(self, net):
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
                                                      hparams)
-    inputs = -1 + np.random.random_integers(
+    inputs = np.random.randint(
         vocab_size, size=(batch_size, 1, 1, 1))
-    targets = -1 + np.random.random_integers(
+    targets = np.random.randint(
         vocab_size, size=(batch_size, size, size, 3))
     with self.test_session() as session:
       features = {
diff --git a/tensor2tensor/models/image_transformer_test.py b/tensor2tensor/models/image_transformer_test.py
index cf8417db4..0ccc1f6bd 100644
--- a/tensor2tensor/models/image_transformer_test.py
+++ b/tensor2tensor/models/image_transformer_test.py
@@ -46,9 +46,9 @@ def testImagetransformer(self, net, hparams):
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
                                                      hparams)
-    inputs = -1 + np.random.random_integers(
+    inputs = np.random.randint(
         vocab_size, size=(batch_size, 1, 1, 1))
-    targets = -1 + np.random.random_integers(
+    targets = np.random.randint(
         vocab_size, size=(batch_size, size, size, 3))
     with self.test_session() as session:
       features = {
diff --git a/tensor2tensor/models/lstm_test.py b/tensor2tensor/models/lstm_test.py
index 45c188816..370748cd0 100644
--- a/tensor2tensor/models/lstm_test.py
+++ b/tensor2tensor/models/lstm_test.py
@@ -30,8 +30,8 @@ class LSTMTest(tf.test.TestCase):
 
   def testLSTMSeq2Seq(self):
     vocab_size = 9
-    x = np.random.random_integers(1, high=vocab_size - 1, size=(3, 5, 1, 1))
-    y = np.random.random_integers(1, high=vocab_size - 1, size=(3, 6, 1, 1))
+    x = np.random.randint(1, high=vocab_size, size=(3, 5, 1, 1))
+    y = np.random.randint(1, high=vocab_size, size=(3, 6, 1, 1))
     hparams = lstm.lstm_seq2seq()
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
@@ -50,8 +50,8 @@ def testLSTMSeq2Seq(self):
 
   def testLSTMSeq2SeqAttention(self):
     vocab_size = 9
-    x = np.random.random_integers(1, high=vocab_size - 1, size=(3, 5, 1, 1))
-    y = np.random.random_integers(1, high=vocab_size - 1, size=(3, 6, 1, 1))
+    x = np.random.randint(1, high=vocab_size, size=(3, 5, 1, 1))
+    y = np.random.randint(1, high=vocab_size, size=(3, 6, 1, 1))
     hparams = lstm.lstm_attention()
 
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
@@ -74,8 +74,8 @@ def testLSTMSeq2SeqAttention(self):
 
   def testLSTMSeq2seqBidirectionalEncoder(self):
     vocab_size = 9
-    x = np.random.random_integers(1, high=vocab_size - 1, size=(3, 5, 1, 1))
-    y = np.random.random_integers(1, high=vocab_size - 1, size=(3, 6, 1, 1))
+    x = np.random.randint(1, high=vocab_size, size=(3, 5, 1, 1))
+    y = np.random.randint(1, high=vocab_size, size=(3, 6, 1, 1))
     hparams = lstm.lstm_seq2seq()
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
@@ -94,8 +94,8 @@ def testLSTMSeq2seqBidirectionalEncoder(self):
 
   def testLSTMSeq2seqAttentionBidirectionalEncoder(self):
     vocab_size = 9
-    x = np.random.random_integers(1, high=vocab_size - 1, size=(3, 5, 1, 1))
-    y = np.random.random_integers(1, high=vocab_size - 1, size=(3, 6, 1, 1))
+    x = np.random.randint(1, high=vocab_size, size=(3, 5, 1, 1))
+    y = np.random.randint(1, high=vocab_size, size=(3, 6, 1, 1))
     hparams = lstm.lstm_attention()
 
     p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
diff --git a/tensor2tensor/models/mtf_image_transformer_test.py b/tensor2tensor/models/mtf_image_transformer_test.py
index 27c240ed9..17c72ce96 100644
--- a/tensor2tensor/models/mtf_image_transformer_test.py
+++ b/tensor2tensor/models/mtf_image_transformer_test.py
@@ -49,7 +49,7 @@ def get_model(hparams=None,
   del p_hparams.modality["inputs"]
   hparams.problem_hparams = p_hparams
 
-  targets = -1 + np.random.random_integers(
+  targets = np.random.randint(
       VOCAB_SIZE, size=(BATCH_SIZE, IMG_LENGTH, IMG_LENGTH, 1, 1))
   features = {
       "targets": tf.constant(targets, dtype=tf.int32, name="targets"),
diff --git a/tensor2tensor/models/mtf_transformer_test.py b/tensor2tensor/models/mtf_transformer_test.py
index 58bdee309..672b02d6a 100644
--- a/tensor2tensor/models/mtf_transformer_test.py
+++ b/tensor2tensor/models/mtf_transformer_test.py
@@ -48,9 +48,9 @@ def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
     del p_hparams.modality["inputs"]
   hparams.problem_hparams = p_hparams
 
-  inputs = -1 + np.random.random_integers(
+  inputs = np.random.randint(
       VOCAB_SIZE, size=(BATCH_SIZE, INPUT_LENGTH, 1, 1))
-  targets = -1 + np.random.random_integers(
+  targets = np.random.randint(
       VOCAB_SIZE, size=(BATCH_SIZE, TARGET_LENGTH, 1, 1))
   features = {
       "targets": tf.constant(targets, dtype=tf.int32, name="targets"),
diff --git a/tensor2tensor/models/neural_gpu_test.py b/tensor2tensor/models/neural_gpu_test.py
index cf8731d90..89cdc9a42 100644
--- a/tensor2tensor/models/neural_gpu_test.py
+++ b/tensor2tensor/models/neural_gpu_test.py
@@ -39,9 +39,9 @@ def testNeuralGPU(self):
     p_hparams = problem_hparams.test_problem_hparams(input_vocab_size,
                                                      target_vocab_size,
                                                      hparams)
-    inputs = -1 + np.random.random_integers(
+    inputs = np.random.randint(
         input_vocab_size, size=(batch_size, input_length, 1, 1))
-    targets = -1 + np.random.random_integers(
+    targets = np.random.randint(
         target_vocab_size, size=(batch_size, target_length, 1, 1))
     with self.test_session() as session:
       features = {
diff --git a/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index 3d8832d47..a7e396f8a 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -33,8 +33,8 @@ class AutoencoderTest(tf.test.TestCase):
   def get_mnist_random_output(self, model_name, hparams_set=None,
                               mode=tf.estimator.ModeKeys.TRAIN):
     hparams_set = hparams_set or model_name
-    x = np.random.random_integers(0, high=255, size=(1, 28, 28, 1))
-    y = np.random.random_integers(0, high=9, size=(1, 1))
+    x = np.random.randint(256, size=(1, 28, 28, 1))
+    y = np.random.randint(10, size=(1, 1))
     features = {
         "targets": tf.constant(x, dtype=tf.int32),
         "inputs": tf.constant(y, dtype=tf.int32),
diff --git a/tensor2tensor/models/research/gene_expression_test.py b/tensor2tensor/models/research/gene_expression_test.py
index 60f9019a8..1f0d461c9 100644
--- a/tensor2tensor/models/research/gene_expression_test.py
+++ b/tensor2tensor/models/research/gene_expression_test.py
@@ -42,8 +42,8 @@ def _test_model(self, hparams, model_cls):
     input_length = target_length * 128 // 4  # chunk_size=4
     input_vocab_size = 5
 
-    inputs = np.random.random_integers(
-        input_vocab_size, size=(batch_size, input_length, 1, 1))
+    inputs = np.random.randint(
+        1, input_vocab_size + 1, size=(batch_size, input_length, 1, 1))
     targets = np.random.random_sample((batch_size, target_length, 1,
                                        target_out))
 
diff --git a/tensor2tensor/models/research/transformer_aux_test.py b/tensor2tensor/models/research/transformer_aux_test.py
index b32d7cc56..d9b2f9d6a 100644
--- a/tensor2tensor/models/research/transformer_aux_test.py
+++ b/tensor2tensor/models/research/transformer_aux_test.py
@@ -90,9 +90,9 @@ def test_transformer_aux_body(self):
                                                      vocab_size,
                                                      hparams)
     hparams.problem_hparams = p_hparams
-    inputs = -1 + np.random.random_integers(
+    inputs = np.random.randint(
         vocab_size, size=(batch_size, input_length, 1, 1))
-    targets = -1 + np.random.random_integers(
+    targets = np.random.randint(
         vocab_size, size=(batch_size, target_length, 1, 1))
     features = {
         "inputs": tf.constant(inputs, dtype=tf.int32),
diff --git a/tensor2tensor/models/research/transformer_revnet_test.py b/tensor2tensor/models/research/transformer_revnet_test.py
index f0fabcf8c..ce4a083c7 100644
--- a/tensor2tensor/models/research/transformer_revnet_test.py
+++ b/tensor2tensor/models/research/transformer_revnet_test.py
@@ -47,9 +47,9 @@ def testTransformer(self):
                                                      vocab_size,
                                                      hparams)
     hparams.problem_hparams = p_hparams
-    inputs = -1 + np.random.random_integers(
+    inputs = np.random.randint(
         vocab_size, size=(batch_size, input_length, 1, 1))
-    targets = -1 + np.random.random_integers(
+    targets = np.random.randint(
         vocab_size, size=(batch_size, target_length, 1, 1))
     features = {
         "inputs": tf.constant(inputs, dtype=tf.int32),
diff --git a/tensor2tensor/models/research/transformer_vae_test.py b/tensor2tensor/models/research/transformer_vae_test.py
index c47c485b5..0cee4bc5e 100644
--- a/tensor2tensor/models/research/transformer_vae_test.py
+++ b/tensor2tensor/models/research/transformer_vae_test.py
@@ -37,9 +37,9 @@ def testTransformerAEOnDVQ(self):
                                                      vocab_size,
                                                      hparams)
     hparams.problem_hparams = p_hparams
-    inputs = -1 + np.random.random_integers(
+    inputs = np.random.randint(
         vocab_size, size=(batch_size, input_length, 1, 1))
-    targets = -1 + np.random.random_integers(
+    targets = np.random.randint(
         vocab_size, size=(batch_size, target_length, 1, 1))
     features = {
         "inputs": tf.constant(inputs, dtype=tf.int32),
diff --git a/tensor2tensor/models/research/universal_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
index 184d8945c..8c322a067 100644
--- a/tensor2tensor/models/research/universal_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -48,9 +48,9 @@ def get_model(self,
       del p_hparams.modality["inputs"]
     hparams.problems = [p_hparams]
 
-    inputs = -1 + np.random.random_integers(
+    inputs = np.random.randint(
         VOCAB_SIZE, size=(BATCH_SIZE, INPUT_LENGTH, 1, 1))
-    targets = -1 + np.random.random_integers(
+    targets = np.random.randint(
         VOCAB_SIZE, size=(BATCH_SIZE, TARGET_LENGTH, 1, 1))
     features = {
         "targets": tf.constant(targets, dtype=tf.int32, name="targets"),
diff --git a/tensor2tensor/models/research/vqa_attention_test.py b/tensor2tensor/models/research/vqa_attention_test.py
index d0490ff41..174970dfd 100644
--- a/tensor2tensor/models/research/vqa_attention_test.py
+++ b/tensor2tensor/models/research/vqa_attention_test.py
@@ -39,10 +39,10 @@ def testVqaAttentionBaseline(self):
     question_length = 5
     answer_length = 10
     x = 2 * np.random.rand(batch_size, image_size, image_size, 3) - 1
-    q = np.random.random_integers(
-        1, high=vocab_size - 1, size=(batch_size, question_length, 1, 1))
-    a = np.random.random_integers(
-        0, high=num_classes, size=(batch_size, answer_length, 1, 1))
+    q = np.random.randint(
+        1, high=vocab_size, size=(batch_size, question_length, 1, 1))
+    a = np.random.randint(
+        num_classes + 1, size=(batch_size, answer_length, 1, 1))
     hparams = vqa_attention.vqa_attention_base()
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
diff --git a/tensor2tensor/models/resnet_test.py b/tensor2tensor/models/resnet_test.py
index c9a54c805..776da4f65 100644
--- a/tensor2tensor/models/resnet_test.py
+++ b/tensor2tensor/models/resnet_test.py
@@ -40,10 +40,10 @@ class ResnetTest(tf.test.TestCase):
   def _test_resnet(self, img_size, output_size):
     vocab_size = 9
     batch_size = 2
-    x = np.random.random_integers(
-        0, high=255, size=(batch_size, img_size, img_size, 3))
-    y = np.random.random_integers(
-        1, high=vocab_size - 1, size=(batch_size, 1, 1, 1))
+    x = np.random.randint(
+        256, size=(batch_size, img_size, img_size, 3))
+    y = np.random.randint(
+        1, high=vocab_size, size=(batch_size, 1, 1, 1))
     hparams = resnet_tiny_cpu()
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
diff --git a/tensor2tensor/models/slicenet_test.py b/tensor2tensor/models/slicenet_test.py
index 314d0dae8..b2463e899 100644
--- a/tensor2tensor/models/slicenet_test.py
+++ b/tensor2tensor/models/slicenet_test.py
@@ -31,8 +31,8 @@
 class SliceNetTest(tf.test.TestCase):
 
   def testSliceNet(self):
-    x = np.random.random_integers(0, high=255, size=(3, 5, 5, 3))
-    y = np.random.random_integers(0, high=9, size=(3, 5, 1, 1))
+    x = np.random.randint(256, size=(3, 5, 5, 3))
+    y = np.random.randint(10, size=(3, 5, 1, 1))
     hparams = slicenet.slicenet_params1_tiny()
     hparams.add_hparam("data_dir", "")
     problem = registry.problem("image_cifar10")
diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 4850c817c..e6882a3ed 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -48,9 +48,9 @@ def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
     del p_hparams.modality["inputs"]
   hparams.problem_hparams = p_hparams
 
-  inputs = -1 + np.random.random_integers(
+  inputs = np.random.randint(
       VOCAB_SIZE, size=(BATCH_SIZE, INPUT_LENGTH, 1, 1))
-  targets = -1 + np.random.random_integers(
+  targets = np.random.randint(
       VOCAB_SIZE, size=(BATCH_SIZE, TARGET_LENGTH, 1, 1))
   features = {
       "targets": tf.constant(targets, dtype=tf.int32, name="targets"),
diff --git a/tensor2tensor/models/xception_test.py b/tensor2tensor/models/xception_test.py
index 25dc75898..98ed8e8af 100644
--- a/tensor2tensor/models/xception_test.py
+++ b/tensor2tensor/models/xception_test.py
@@ -33,10 +33,10 @@ class XceptionTest(tf.test.TestCase):
   def _test_xception(self, img_size):
     vocab_size = 9
     batch_size = 3
-    x = np.random.random_integers(
-        0, high=255, size=(batch_size, img_size, img_size, 3))
-    y = np.random.random_integers(
-        1, high=vocab_size - 1, size=(batch_size, 1, 1, 1))
+    x = np.random.randint(
+        256, size=(batch_size, img_size, img_size, 3))
+    y = np.random.randint(
+        1, high=vocab_size, size=(batch_size, 1, 1, 1))
     hparams = xception.xception_tiny()
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,

From 4c7e7ef7582bee09341308e0d94a05a08c307f6f Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Mon, 7 Jan 2019 21:45:48 +0000
Subject: [PATCH 1445/2720] Reduce usage of tf.contrib (#1345)

* tf.contrib.rnn -> tf.nn.rnn_cell

* tf.contrib.framework -> tf.train

tf.contrib.framework.load_checkpoint -> tf.train.load_checkpoint
tf.contrib.framework.list_variables -> tf.train.list_variables
tf.contrib.framework.init_from_checkpoint -> tf.train.init_from_checkpoint

* tf.contrib.learn.ModeKeys -> tf.estimator.ModeKeys

tf.contrib.learn.ModeKeys.TRAIN -> tf.estimator.ModeKeys.TRAIN
tf.contrib.learn.ModeKeys.EVAL -> tf.estimator.ModeKeys.EVAL
tf.contrib.learn.ModeKeys.INFER -> tf.estimator.ModeKeys.PREDICT

* tf.contrib.eager.in_eager_mode -> tf.executing_eagerly

* tf.contrib.data.batch_and_drop_remainder -> .batch(..., drop_remainder=True)

* tf.contrib.data -> tf.data.experimental

* tf.contrib.util.make_tensor_proto -> tf.make_tensor_proto

* Simplify tf.contrib.eager imports

* tf.initialize_all_variables -> tf.global_variables_initializer
---
 tensor2tensor/bin/t2t_attack.py                |  2 +-
 tensor2tensor/bin/t2t_avg_all.py               |  6 +++---
 tensor2tensor/data_generators/image_utils.py   |  2 +-
 tensor2tensor/data_generators/problem.py       |  8 ++++----
 tensor2tensor/layers/common_image_attention.py |  4 ++--
 .../layers/common_image_attention_test.py      |  2 +-
 tensor2tensor/layers/common_layers.py          |  8 ++++----
 tensor2tensor/layers/common_video.py           |  2 +-
 tensor2tensor/layers/discretization.py         |  4 ++--
 tensor2tensor/layers/modalities.py             |  6 +++---
 tensor2tensor/models/image_transformer.py      |  2 +-
 tensor2tensor/models/image_transformer_2d.py   |  6 +++---
 tensor2tensor/models/lstm.py                   | 14 +++++++-------
 tensor2tensor/models/research/glow_ops.py      |  4 ++--
 tensor2tensor/models/research/glow_ops_test.py |  2 +-
 .../models/research/transformer_vae.py         |  2 +-
 tensor2tensor/models/research/vqa_attention.py |  6 +++---
 tensor2tensor/models/shake_shake.py            |  6 +++---
 tensor2tensor/models/video/svg_lp.py           |  8 ++++----
 tensor2tensor/notebooks/asr_transformer.ipynb  |  4 ++--
 tensor2tensor/rl/trainer_model_based.py        |  4 ++--
 tensor2tensor/serving/serving_utils.py         |  2 +-
 tensor2tensor/utils/avg_checkpoints.py         |  6 +++---
 tensor2tensor/utils/expert_utils.py            |  4 ++--
 tensor2tensor/utils/metrics.py                 |  4 ++--
 tensor2tensor/utils/multistep_optimizer.py     |  2 +-
 tensor2tensor/utils/optimize.py                |  2 +-
 tensor2tensor/utils/registry.py                | 14 +++++++-------
 tensor2tensor/utils/t2t_model.py               | 18 +++++++++---------
 tensor2tensor/utils/video/prediction2gif.py    |  2 +-
 tensor2tensor/utils/video/reward_confusion.py  |  2 +-
 tensor2tensor/utils/video_metrics.py           |  2 +-
 32 files changed, 80 insertions(+), 80 deletions(-)

diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index b9e167238..6276d6b4c 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -200,7 +200,7 @@ def main(argv):
     sur_ch_model.get_probs(inputs)
 
     checkpoint_path = os.path.expanduser(FLAGS.surrogate_output_dir)
-    tf.contrib.framework.init_from_checkpoint(
+    tf.train.init_from_checkpoint(
         tf.train.latest_checkpoint(checkpoint_path), {"/": "surrogate/"})
     sess.run(tf.global_variables_initializer())
 
diff --git a/tensor2tensor/bin/t2t_avg_all.py b/tensor2tensor/bin/t2t_avg_all.py
index 080670f35..e3c34be90 100644
--- a/tensor2tensor/bin/t2t_avg_all.py
+++ b/tensor2tensor/bin/t2t_avg_all.py
@@ -60,7 +60,7 @@ def main(_):
   for model in bleu_hook.stepfiles_iterator(model_dir, FLAGS.wait_minutes,
                                             FLAGS.min_steps):
     if models_processed == 0:
-      var_list = tf.contrib.framework.list_variables(model.filename)
+      var_list = tf.train.list_variables(model.filename)
       avg_values = {}
       for (name, shape) in var_list:
         if not (name.startswith("global_step") or
@@ -69,7 +69,7 @@ def main(_):
     models_processed += 1
 
     tf.logging.info("Loading [%d]: %s" % (models_processed, model.filename))
-    reader = tf.contrib.framework.load_checkpoint(model.filename)
+    reader = tf.train.load_checkpoint(model.filename)
     for name in avg_values:
       avg_values[name] += reader.get_tensor(name) / FLAGS.n
     queue.append(model)
@@ -106,7 +106,7 @@ def main(_):
     tf.reset_default_graph()
     first_model = queue.popleft()
 
-    reader = tf.contrib.framework.load_checkpoint(first_model.filename)
+    reader = tf.train.load_checkpoint(first_model.filename)
     for name in avg_values:
       avg_values[name] -= reader.get_tensor(name) / FLAGS.n
 
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 5396d931c..db5ddbb44 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -262,7 +262,7 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
 def encode_images_as_png(images):
   """Yield images encoded as pngs."""
-  if tf.contrib.eager.in_eager_mode():
+  if tf.executing_eagerly():
     for image in images:
       yield tf.image.encode_png(image).numpy()
   else:
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 219deabdf..5406ddf73 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -423,7 +423,7 @@ def _preprocess(example):
 
     if interleave:
       dataset = dataset.apply(
-          tf.contrib.data.parallel_interleave(
+          tf.data.experimental.parallel_interleave(
               _preprocess, sloppy=True, cycle_length=8))
     else:
       dataset = dataset.flat_map(_preprocess)
@@ -674,7 +674,7 @@ def _load_records_and_preprocess(filenames):
     # Create data-set from files by parsing, pre-processing and interleaving.
     if shuffle_files:
       dataset = dataset.apply(
-          tf.contrib.data.parallel_interleave(
+          tf.data.experimental.parallel_interleave(
               _load_records_and_preprocess, sloppy=True, cycle_length=8))
     else:
       dataset = _load_records_and_preprocess(dataset)
@@ -963,7 +963,7 @@ def define_shapes(example):
           batching_scheme["batch_sizes"] = [hparams.batch_size]
           batching_scheme["boundaries"] = []
         dataset = dataset.apply(
-            tf.contrib.data.bucket_by_sequence_length(
+            tf.data.experimental.bucket_by_sequence_length(
                 data_reader.example_length, batching_scheme["boundaries"],
                 batching_scheme["batch_sizes"]))
 
@@ -1040,7 +1040,7 @@ def serving_input_fn(self, hparams):
         tf.shape(serialized_example, out_type=tf.int64)[0],
         dataset.output_shapes)
     dataset = dataset.map(standardize_shapes)
-    features = tf.contrib.data.get_single_element(dataset)
+    features = tf.data.experimental.get_single_element(dataset)
 
     if self.has_inputs:
       features.pop("targets", None)
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 46608d7aa..27191a9dc 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -501,7 +501,7 @@ def postprocess_image(x, rows, cols, hparams):
                               use_bias=True,
                               activation=None,
                               name="output_conv")
-  if (hparams.mode == tf.contrib.learn.ModeKeys.INFER and
+  if (hparams.mode == tf.estimator.ModeKeys.PREDICT and
       hparams.block_raster_scan):
     y = targets
     yshape = common_layers.shape_list(y)
@@ -547,7 +547,7 @@ def prepare_decoder(targets, hparams):
 
   # during training, images are [batch, IMG_LEN, IMG_LEN, 3].
   # At inference, they are [batch, curr_infer_length, 1, 1]
-  if hparams.mode == tf.contrib.learn.ModeKeys.INFER:
+  if hparams.mode == tf.estimator.ModeKeys.PREDICT:
     curr_infer_length = targets_shape[1]
     if hparams.block_raster_scan:
       assert hparams.img_len*channels % hparams.query_shape[1] == 0
diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index e87a2e4d8..088725b18 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -61,7 +61,7 @@ def testPostProcessImageInferMode(self, likelihood, num_mixtures, depth):
         block_raster_scan=True,
         hidden_size=2,
         likelihood=likelihood,
-        mode=tf.contrib.learn.ModeKeys.INFER,
+        mode=tf.estimator.ModeKeys.PREDICT,
         num_mixtures=num_mixtures,
         query_shape=[block_length, block_width],
     )
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index affd409e4..0068d9e86 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -349,7 +349,7 @@ def embedding(x,
     # On the backwards pass, we want to convert the gradient from
     # an indexed-slices to a regular tensor before sending it back to the
     # parameter server. This avoids excess computation on the parameter server.
-    if not tf.contrib.eager.in_eager_mode():
+    if not tf.executing_eagerly():
       embedding_var = convert_gradient_to_tensor(embedding_var)
     x = dropout_no_scaling(x, 1.0 - symbol_dropout_rate)
     emb_x = gather(embedding_var, x, dtype)
@@ -2868,7 +2868,7 @@ def ones_matrix_band_part(rows, cols, num_lower, num_upper, out_shape=None):
 def reshape_like_all_dims(a, b):
   """Reshapes a to match the shape of b."""
   ret = tf.reshape(a, tf.shape(b))
-  if not tf.contrib.eager.in_eager_mode():
+  if not tf.executing_eagerly():
     ret.set_shape(b.get_shape())
   return ret
 
@@ -3193,7 +3193,7 @@ def should_generate_summaries():
 def reshape_like(a, b):
   """Reshapes a to match the shape of b in all but the last dimension."""
   ret = tf.reshape(a, tf.concat([tf.shape(b)[:-1], tf.shape(a)[-1:]], 0))
-  if not tf.contrib.eager.in_eager_mode():
+  if not tf.executing_eagerly():
     ret.set_shape(b.get_shape().as_list()[:-1] + a.get_shape().as_list()[-1:])
   return ret
 
@@ -3205,7 +3205,7 @@ def summarize_video(video, prefix, max_outputs=1):
     raise ValueError("Assuming videos given as tensors in the format "
                      "[batch, time, height, width, channels] but got one "
                      "of shape: %s" % str(video_shape))
-  if tf.contrib.eager.in_eager_mode():
+  if tf.executing_eagerly():
     return
   if video.get_shape().as_list()[1] is None:
     tf.summary.image(
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index cbac48912..ef4568b6e 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -82,7 +82,7 @@ def lstm_cell(inputs,
               name=None):
   """Full LSTM cell."""
   input_shape = common_layers.shape_list(inputs)
-  cell = tf.contrib.rnn.LSTMCell(num_units,
+  cell = tf.nn.rnn_cell.LSTMCell(num_units,
                                  use_peepholes=use_peepholes,
                                  cell_clip=cell_clip,
                                  initializer=initializer,
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 050b3f504..2293cd8b0 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -473,7 +473,7 @@ def gumbel_softmax(x,
     d_dev = -tf.reduce_mean(d_variance)
     ret = s
 
-    if mode != tf.contrib.learn.ModeKeys.TRAIN:
+    if mode != tf.estimator.ModeKeys.TRAIN:
       ret = tf.reshape(maxvhot, common_layers.shape_list(s))  # Just hot @eval.
     return m, ret, d_dev * 5.0 + tf.reduce_mean(kl) * 0.002
 
@@ -822,7 +822,7 @@ def predict_bits_with_lstm(prediction_source, state_size, total_num_bits,
 
   with tf.variable_scope("predict_bits_with_lstm"):
     # Layers and cell state creation.
-    lstm_cell = tf.contrib.rnn.LSTMCell(state_size)
+    lstm_cell = tf.nn.rnn_cell.LSTMCell(state_size)
     discrete_predict = tf.layers.Dense(2**bits_at_once, name="discrete_predict")
     discrete_embed = tf.layers.Dense(state_size, name="discrete_embed")
     batch_size = common_layers.shape_list(prediction_source)[0]
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 0fe4121ad..920124589 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -93,7 +93,7 @@ def _get_weights(self, hidden_dim=None):
     else:
       ret = tf.concat(shards, 0)
     # Convert ret to tensor.
-    if not tf.contrib.eager.in_eager_mode():
+    if not tf.executing_eagerly():
       ret = common_layers.convert_gradient_to_tensor(ret)
     return ret
 
@@ -226,7 +226,7 @@ class ImageModality(modality.Modality):
 
   def bottom(self, x):
     with tf.variable_scope(self.name):
-      if not tf.contrib.eager.in_eager_mode():
+      if not tf.executing_eagerly():
         tf.summary.image(
             "inputs", common_layers.tpu_safe_image_summary(x), max_outputs=2)
       return tf.to_float(x)
@@ -234,7 +234,7 @@ def bottom(self, x):
   def targets_bottom(self, x):
     inputs = x
     with tf.variable_scope(self.name):
-      if not tf.contrib.eager.in_eager_mode():
+      if not tf.executing_eagerly():
         tf.summary.image(
             "targets_bottom",
             common_layers.tpu_safe_image_summary(inputs),
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 4b52f9a42..1f841e28b 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -55,7 +55,7 @@ def body(self, features):
                        "must be ImageChannelBottomIdentityModality and "
                        "num_channels must be 1.")
     if (not tf.get_variable_scope().reuse and
-        hparams.mode != tf.contrib.learn.ModeKeys.INFER and
+        hparams.mode != tf.estimator.ModeKeys.PREDICT and
         hparams.modality["targets"] !=
         modalities.ImageChannelBottomIdentityModality):
       tf.summary.image("targets", tf.to_float(targets), max_outputs=1)
diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 7555150a7..5de2584c2 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -46,7 +46,7 @@ def body(self, features):
     targets = features["targets"]
     targets_shape = common_layers.shape_list(targets)
     if not (tf.get_variable_scope().reuse or
-            hparams.mode == tf.contrib.learn.ModeKeys.INFER):
+            hparams.mode == tf.estimator.ModeKeys.PREDICT):
       tf.summary.image("targets", targets, max_outputs=1)
 
     decoder_input, rows, cols = cia.prepare_decoder(
@@ -76,7 +76,7 @@ def body(self, features):
     targets = features["targets"]
     inputs = features["inputs"]
     if not (tf.get_variable_scope().reuse or
-            hparams.mode == tf.contrib.learn.ModeKeys.INFER):
+            hparams.mode == tf.estimator.ModeKeys.PREDICT):
       tf.summary.image("inputs", inputs, max_outputs=1)
       tf.summary.image("targets", targets, max_outputs=1)
 
@@ -112,7 +112,7 @@ def body(self, features):
     targets = features["targets"]
     inputs = features["inputs"]
     if not (tf.get_variable_scope().reuse or
-            hparams.mode == tf.contrib.learn.ModeKeys.INFER):
+            hparams.mode == tf.estimator.ModeKeys.PREDICT):
       tf.summary.image("inputs", inputs, max_outputs=1)
       tf.summary.image("targets", targets, max_outputs=1)
 
diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 2f657970c..58cadea46 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -29,8 +29,8 @@
 
 
 def _dropout_lstm_cell(hparams, train):
-  return tf.contrib.rnn.DropoutWrapper(
-      tf.contrib.rnn.LSTMCell(hparams.hidden_size),
+  return tf.nn.rnn_cell.DropoutWrapper(
+      tf.nn.rnn_cell.LSTMCell(hparams.hidden_size),
       input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
 
 
@@ -58,7 +58,7 @@ def lstm(inputs, sequence_length, hparams, train, name, initial_state=None):
             for _ in range(hparams.num_hidden_layers)]
   with tf.variable_scope(name):
     return tf.nn.dynamic_rnn(
-        tf.contrib.rnn.MultiRNNCell(layers),
+        tf.nn.rnn_cell.MultiRNNCell(layers),
         inputs,
         sequence_length,
         initial_state=initial_state,
@@ -192,11 +192,11 @@ def lstm_bid_encoder(inputs, sequence_length, hparams, train, name):
   """Bidirectional LSTM for encoding inputs that are [batch x time x size]."""
 
   with tf.variable_scope(name):
-    cell_fw = tf.contrib.rnn.MultiRNNCell(
+    cell_fw = tf.nn.rnn_cell.MultiRNNCell(
         [_dropout_lstm_cell(hparams, train)
          for _ in range(hparams.num_hidden_layers)])
 
-    cell_bw = tf.contrib.rnn.MultiRNNCell(
+    cell_bw = tf.nn.rnn_cell.MultiRNNCell(
         [_dropout_lstm_cell(hparams, train)
          for _ in range(hparams.num_hidden_layers)])
 
@@ -213,7 +213,7 @@ def lstm_bid_encoder(inputs, sequence_length, hparams, train, name):
     encoder_states = []
 
     for i in range(hparams.num_hidden_layers):
-      if isinstance(encoder_fw_state[i], tf.contrib.rnn.LSTMStateTuple):
+      if isinstance(encoder_fw_state[i], tf.nn.rnn_cell.LSTMStateTuple):
         encoder_state_c = tf.concat(
             values=(encoder_fw_state[i].c, encoder_bw_state[i].c),
             axis=1,
@@ -222,7 +222,7 @@ def lstm_bid_encoder(inputs, sequence_length, hparams, train, name):
             values=(encoder_fw_state[i].h, encoder_bw_state[i].h),
             axis=1,
             name="encoder_fw_state_h")
-        encoder_state = tf.contrib.rnn.LSTMStateTuple(
+        encoder_state = tf.nn.rnn_cell.LSTMStateTuple(
             c=encoder_state_c, h=encoder_state_h)
       elif isinstance(encoder_fw_state[i], tf.Tensor):
         encoder_state = tf.concat(
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index a8795bc34..ef806f6e2 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -977,7 +977,7 @@ def compute_prior(name, z, latent, hparams, condition=False, state=None,
             The first-three dimensions of the latent should be the same as z.
     hparams: next_frame_glow_hparams.
     condition: Whether or not to condition the distribution on latent.
-    state: tf.contrib.rnn.LSTMStateTuple.
+    state: tf.nn.rnn_cell.LSTMStateTuple.
            the current state of a LSTM used to model the distribution. Used
            only if hparams.latent_dist_encoder = "conv_lstm".
     temperature: float, temperature with which to sample from the Gaussian.
@@ -1025,7 +1025,7 @@ def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
     eps_std: Sample x2 with the provided eps_std.
     cond_latents: optionally condition x2 on cond_latents.
     hparams: next_frame_glow hparams.
-    state: tf.contrib.rnn.LSTMStateTuple. Current state of the LSTM over z_2.
+    state: tf.nn.rnn_cell.LSTMStateTuple. Current state of the LSTM over z_2.
            Used only when hparams.latent_dist_encoder == "conv_lstm"
     condition: bool, Whether or not to condition the distribution on
                cond_latents.
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 101edb2d3..5266cddbb 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -349,7 +349,7 @@ def test_latent_dist_encoder(self, encoder="conv_lstm", skip=True,
       state_t = tf.convert_to_tensor(state_rand)
       if encoder in ["conv_net", "conv3d_net"]:
         latent_t = [latent_t, latent_t]
-      init_state = tf.contrib.rnn.LSTMStateTuple(state_t, state_t)
+      init_state = tf.nn.rnn_cell.LSTMStateTuple(state_t, state_t)
       hparams = self.get_glow_hparams()
       hparams.latent_dist_encoder = encoder
       hparams.latent_skip = skip
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 58db6dc5b..11446fd64 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -103,7 +103,7 @@ def top_k_softmax(x, k):
 def top_k_experts(x, k, hparams):
   x_shape = common_layers.shape_list(x)
   x_flat = tf.reshape(x, [-1, common_layers.shape_list(x)[-1]])
-  is_training = hparams.mode == tf.contrib.learn.ModeKeys.TRAIN
+  is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
   gates, load = expert_utils.noisy_top_k_gating(
       x_flat, 2 ** hparams.z_size, is_training, k)
   gates_shape = [x_shape[0], x_shape[1], x_shape[2], 2 ** hparams.z_size]
diff --git a/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index aaaeb3a70..774387878 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -234,10 +234,10 @@ def image_encoder(image_feat,
 
 def _get_rnn_cell(hparams):
   if hparams.rnn_type == "lstm":
-    rnn_cell = tf.contrib.rnn.BasicLSTMCell
+    rnn_cell = tf.nn.rnn_cell.BasicLSTMCell
   elif hparams.rnn_type == "lstm_layernorm":
     rnn_cell = tf.contrib.rnn.LayerNormBasicLSTMCell
-  return tf.contrib.rnn.DropoutWrapper(
+  return tf.nn.rnn_cell.DropoutWrapper(
       rnn_cell(hparams.hidden_size),
       output_keep_prob=1.0-hparams.dropout)
 
@@ -269,7 +269,7 @@ def question_encoder(question, hparams, name="encoder"):
 
     # rnn_layers = [_get_rnn_cell(hparams)
     #               for _ in range(hparams.num_rnn_layers)]
-    # rnn_multi_cell = tf.contrib.rnn.MultiRNNCell(rnn_layers)
+    # rnn_multi_cell = tf.nn.rnn_cell.MultiRNNCell(rnn_layers)
     rnn_cell = _get_rnn_cell(hparams)
     # outputs, _ = tf.nn.dynamic_rnn(
     #     rnn_cell, question, length, dtype=tf.float32)
diff --git a/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index 4a83444c0..4af05dc21 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -54,7 +54,7 @@ def shake_shake_skip_connection(x, output_filters, stride, is_training):
 def shake_shake_branch(x, output_filters, stride, rand_forward, rand_backward,
                        hparams):
   """Building a 2 branching convnet."""
-  is_training = hparams.mode == tf.contrib.learn.ModeKeys.TRAIN
+  is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
   x = tf.nn.relu(x)
   x = tf.layers.conv2d(
       x,
@@ -76,7 +76,7 @@ def shake_shake_branch(x, output_filters, stride, rand_forward, rand_backward,
 
 def shake_shake_block(x, output_filters, stride, hparams):
   """Builds a full shake-shake sub layer."""
-  is_training = hparams.mode == tf.contrib.learn.ModeKeys.TRAIN
+  is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
   batch_size = common_layers.shape_list(x)[0]
 
   # Generate random numbers for scaling the branches.
@@ -138,7 +138,7 @@ class ShakeShake(t2t_model.T2TModel):
 
   def body(self, features):
     hparams = self._hparams
-    is_training = hparams.mode == tf.contrib.learn.ModeKeys.TRAIN
+    is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
     inputs = features["inputs"]
     assert (hparams.num_hidden_layers - 2) % 6 == 0
     assert hparams.hidden_size % 16 == 0
diff --git a/tensor2tensor/models/video/svg_lp.py b/tensor2tensor/models/video/svg_lp.py
index 7a49e549b..52e9ec2df 100644
--- a/tensor2tensor/models/video/svg_lp.py
+++ b/tensor2tensor/models/video/svg_lp.py
@@ -57,13 +57,13 @@ def rnn_model(self, hidden_size, nlayers, rnn_type, name):
     """
     layers_units = [hidden_size] * nlayers
     if rnn_type == "lstm":
-      rnn_cell = tf.contrib.rnn.LSTMCell
+      rnn_cell = tf.nn.rnn_cell.LSTMCell
     elif rnn_type == "gru":
-      rnn_cell = tf.contrib.rnn.GRUCell
+      rnn_cell = tf.nn.rnn_cell.GRUCell
     else:
-      rnn_cell = tf.contrib.rnn.RNNCell
+      rnn_cell = tf.nn.rnn_cell.RNNCell
     cells = [rnn_cell(units, name=name) for units in layers_units]
-    stacked_rnn = tf.contrib.rnn.MultiRNNCell(cells)
+    stacked_rnn = tf.nn.rnn_cell.MultiRNNCell(cells)
     return stacked_rnn
 
   def deterministic_rnn(self, cell, inputs, states, output_size, scope):
diff --git a/tensor2tensor/notebooks/asr_transformer.ipynb b/tensor2tensor/notebooks/asr_transformer.ipynb
index d1802d727..82a0728a8 100644
--- a/tensor2tensor/notebooks/asr_transformer.ipynb
+++ b/tensor2tensor/notebooks/asr_transformer.ipynb
@@ -93,8 +93,8 @@
         "from tensor2tensor.utils import metrics\n",
         "\n",
         "# Enable TF Eager execution\n",
-        "from tensorflow.contrib.eager.python import tfe\n",
-        "tfe.enable_eager_execution()\n",
+        "tfe = tf.contrib.eager\n",
+        "tf.enable_eager_execution()\n",
         "\n",
         "# Other setup\n",
         "Modes = tf.estimator.ModeKeys\n",
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 12c6368e1..78abce708 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -177,7 +177,7 @@ def train_agent(real_env, learner, world_model_dir, hparams, epoch):
   """Train the PPO agent in the simulated environment."""
   frame_stack_size = hparams.frame_stack_size
   initial_frame_rollouts = real_env.current_epoch_rollouts(
-      split=tf.contrib.learn.ModeKeys.TRAIN,
+      split=tf.estimator.ModeKeys.TRAIN,
       minimal_rollout_frames=frame_stack_size,
   )
   # TODO(koz4k): Move this to a different module.
@@ -304,7 +304,7 @@ def initial_frame_chooser(batch_size):
       max(hparams.wm_eval_rollout_ratios) * hparams.simulated_rollout_length
   )
   rollouts = real_env.current_epoch_rollouts(
-      split=tf.contrib.learn.ModeKeys.EVAL,
+      split=tf.estimator.ModeKeys.EVAL,
       minimal_rollout_frames=(subsequence_length + frame_stack_size)
   )
 
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 8d50a7756..52695116a 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -108,7 +108,7 @@ def _make_grpc_request(examples):
     request = predict_pb2.PredictRequest()
     request.model_spec.name = servable_name
     request.inputs["input"].CopyFrom(
-        tf.contrib.util.make_tensor_proto(
+        tf.make_tensor_proto(
             [ex.SerializeToString() for ex in examples], shape=[len(examples)]))
     response = stub.Predict(request, timeout_secs)
     outputs = tf.make_ndarray(response.outputs["outputs"])
diff --git a/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index 3d57744e6..01684a394 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -75,13 +75,13 @@ def main(_):
   tf.logging.info("Reading variables and averaging checkpoints:")
   for c in checkpoints:
     tf.logging.info("%s ", c)
-  var_list = tf.contrib.framework.list_variables(checkpoints[0])
+  var_list = tf.train.list_variables(checkpoints[0])
   var_values, var_dtypes = {}, {}
   for (name, shape) in var_list:
     if not name.startswith("global_step"):
       var_values[name] = np.zeros(shape)
   for checkpoint in checkpoints:
-    reader = tf.contrib.framework.load_checkpoint(checkpoint)
+    reader = tf.train.load_checkpoint(checkpoint)
     for name in var_values:
       tensor = reader.get_tensor(name)
       var_dtypes[name] = tensor.dtype
@@ -103,7 +103,7 @@ def main(_):
 
   # Build a model consisting only of variables, set them to the average values.
   with tf.Session() as sess:
-    sess.run(tf.initialize_all_variables())
+    sess.run(tf.global_variables_initializer())
     for p, assign_op, (name, value) in zip(placeholders, assign_ops,
                                            six.iteritems(var_values)):
       sess.run(assign_op, {p: value})
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 01a525ded..829dcd607 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -636,7 +636,7 @@ def remove(self, x):
           x,
           indices=self.nonpad_ids,
       )
-      if not tf.contrib.eager.in_eager_mode():
+      if not tf.executing_eagerly():
         # This is a hack but for some reason, gather_nd return a tensor of
         # undefined shape, so the shape is set up manually
         x.set_shape([None] + x_shape[1:])
@@ -986,7 +986,7 @@ def my_fn(x):
 def flatten_all_but_last(a):
   """Flatten all dimensions of a except the last."""
   ret = tf.reshape(a, [-1, tf.shape(a)[-1]])
-  if not tf.contrib.eager.in_eager_mode():
+  if not tf.executing_eagerly():
     ret.set_shape([None] + a.get_shape().as_list()[-1:])
   return ret
 
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index feea42058..614cefe57 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -26,10 +26,10 @@
 from tensor2tensor.utils import rouge
 
 import tensorflow as tf
-
-from tensorflow.contrib.eager.python import tfe
 from tensorflow.python.util import tf_inspect as inspect
 
+tfe = tf.contrib.eager
+
 
 class Metrics(object):
   """Available evaluation metrics."""
diff --git a/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
index 3f516cd8e..32dd6cdd1 100644
--- a/tensor2tensor/utils/multistep_optimizer.py
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -52,7 +52,7 @@ def _create_slots(self, var_list):
 
   def _get_iter_variable(self):
     graph = (
-        None if tf.contrib.eager.in_eager_mode() else tf.get_default_graph())
+        None if tf.executing_eagerly() else tf.get_default_graph())
     return self._get_non_slot_variable("iter", graph=graph)
 
   def _prepare(self):
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index d6e8e3e8e..ab06a6b2f 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -290,7 +290,7 @@ def get_variable_initializer(hparams):
                                value=hparams.initializer_gain,
                                hparams=hparams)
 
-  if not tf.contrib.eager.in_eager_mode():
+  if not tf.executing_eagerly():
     tf.logging.info("Using variable initializer: %s", hparams.initializer)
   if hparams.initializer == "orthogonal":
     return tf.orthogonal_initializer(gain=hparams.initializer_gain)
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 7c9225f42..b8c3e5991 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -93,7 +93,7 @@ def register_model(name=None):
   def decorator(model_cls, registration_name=None):
     """Registers & returns model_cls with registration_name or default name."""
     model_name = registration_name or default_name(model_cls)
-    if model_name in _MODELS and not tf.contrib.eager.in_eager_mode():
+    if model_name in _MODELS and not tf.executing_eagerly():
       raise LookupError("Model %s already registered." % model_name)
     model_cls.REGISTERED_NAME = model_name
     _MODELS[model_name] = model_cls
@@ -125,7 +125,7 @@ def register_hparams(name=None):
   def decorator(hp_fn, registration_name=None):
     """Registers & returns hp_fn with registration_name or default name."""
     hp_name = registration_name or default_name(hp_fn)
-    if hp_name in _HPARAMS and not tf.contrib.eager.in_eager_mode():
+    if hp_name in _HPARAMS and not tf.executing_eagerly():
       raise LookupError("HParams set %s already registered." % hp_name)
     _HPARAMS[hp_name] = hp_fn
     return hp_fn
@@ -199,7 +199,7 @@ def register_problem(name=None):
   def decorator(p_cls, registration_name=None):
     """Registers & returns p_cls with registration_name or default name."""
     p_name = registration_name or default_name(p_cls)
-    if p_name in _PROBLEMS and not tf.contrib.eager.in_eager_mode():
+    if p_name in _PROBLEMS and not tf.executing_eagerly():
       raise LookupError("Problem %s already registered." % p_name)
 
     _PROBLEMS[p_name] = p_cls
@@ -259,7 +259,7 @@ def register_attack(name=None):
   def decorator(attack_fn, registration_name=None):
     """Registers & returns attack_fn with registration_name or default name."""
     attack_name = registration_name or default_name(attack_fn)
-    if attack_name in _ATTACKS and not tf.contrib.eager.in_eager_mode():
+    if attack_name in _ATTACKS and not tf.executing_eagerly():
       raise LookupError("Attack %s already registered." % attack_name)
     _ATTACKS[attack_name] = attack_fn
     return attack_fn
@@ -299,7 +299,7 @@ def register_attack_params(name=None):
   def decorator(ap_fn, registration_name=None):
     """Registers & returns ap_fn with registration_name or default name."""
     ap_name = registration_name or default_name(ap_fn)
-    if ap_name in _ATTACK_PARAMS and not tf.contrib.eager.in_eager_mode():
+    if ap_name in _ATTACK_PARAMS and not tf.executing_eagerly():
       raise LookupError("Attack HParams set %s already registered." % ap_name)
     _ATTACK_PARAMS[ap_name] = ap_fn
     return ap_fn
@@ -338,7 +338,7 @@ def register_pruning_params(name=None):
   def decorator(pp_fn, registration_name=None):
     """Registers & returns pp_fn with registration_name or default name."""
     pp_name = registration_name or default_name(pp_fn)
-    if pp_name in _PRUNING_PARAMS and not tf.contrib.eager.in_eager_mode():
+    if pp_name in _PRUNING_PARAMS and not tf.executing_eagerly():
       raise LookupError("Pruning HParams set %s already registered." % pp_name)
     _PRUNING_PARAMS[pp_name] = pp_fn
     return pp_fn
@@ -376,7 +376,7 @@ def register_pruning_strategy(name=None):
   def decorator(ps_fn, registration_name=None):
     """Registers & returns ps_fn with registration_name or default name."""
     ps_name = registration_name or default_name(ps_fn)
-    if ps_name in _PRUNING_STRATEGY and not tf.contrib.eager.in_eager_mode():
+    if ps_name in _PRUNING_STRATEGY and not tf.executing_eagerly():
       raise LookupError("Pruning strategy %s already registered." % ps_name)
     _PRUNING_STRATEGY[ps_name] = ps_fn
     return ps_fn
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index f3530b574..d95bd98d7 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -939,7 +939,7 @@ def _slow_greedy_infer_tpu(self, features, decode_length):
 
     def infer_step(i, recent_output, recent_logits, unused_loss):
       """Inference step."""
-      if not tf.contrib.eager.in_eager_mode():
+      if not tf.executing_eagerly():
         recent_output.set_shape([None, None, None, 1])
       padded = tf.pad(recent_output, [[0, 0], [0, 1], [0, 0], [0, 0]])
       features["targets"] = padded
@@ -957,7 +957,7 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
       samples = inplace_ops.alias_inplace_update(samples, i,
                                                  tf.to_int64(cur_sample))
       samples = tf.transpose(samples, perm=[1, 0, 2, 3])
-      if not tf.contrib.eager.in_eager_mode():
+      if not tf.executing_eagerly():
         samples.set_shape([None, None, None, 1])
 
       # Assuming we have one shard for logits.
@@ -1000,7 +1000,7 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
     # tensor padded to [batch_size, decode_length, 1, 1, vocab_size]
     logits = tf.zeros((batch_size, decode_length, 1, 1,
                        target_modality.top_dimensionality))
-    if not tf.contrib.eager.in_eager_mode():
+    if not tf.executing_eagerly():
       logits.set_shape([None, None, None, None, None])
     loss = 0.0
 
@@ -1106,7 +1106,7 @@ def _slow_greedy_infer(self, features, decode_length):
 
     def infer_step(recent_output, recent_logits, unused_loss):
       """Inference step."""
-      if not tf.contrib.eager.in_eager_mode():
+      if not tf.executing_eagerly():
         if self._target_modality_is_real:
           dim = self._problem_hparams.modality["targets"].top_dimensionality
           recent_output.set_shape([None, None, None, dim])
@@ -1130,7 +1130,7 @@ def infer_step(recent_output, recent_logits, unused_loss):
       else:
         cur_sample = tf.to_int64(tf.expand_dims(cur_sample, axis=1))
         samples = tf.concat([recent_output, cur_sample], axis=1)
-        if not tf.contrib.eager.in_eager_mode():
+        if not tf.executing_eagerly():
           samples.set_shape([None, None, None, 1])
 
       # Assuming we have one shard for logits.
@@ -1176,7 +1176,7 @@ def infer_step(recent_output, recent_logits, unused_loss):
       logits = tf.zeros((batch_size, 0, 1, 1,
                          target_modality.top_dimensionality))
       logits_shape_inv = [None, None, None, None, None]
-    if not tf.contrib.eager.in_eager_mode():
+    if not tf.executing_eagerly():
       logits.set_shape(logits_shape_inv)
 
     loss = 0.0
@@ -1459,7 +1459,7 @@ def initialize_from_ckpt(self, ckpt_dir):
     log_info("Checkpoint dir: %s", ckpt_dir)
 
     # TODO(mitchellstern): Add support for partitioned variables?
-    reader = tf.contrib.framework.load_checkpoint(ckpt_dir)
+    reader = tf.train.load_checkpoint(ckpt_dir)
     variable_map = {}
     for var in tf.contrib.framework.get_trainable_variables():
       var_name = var.name.split(":")[0]
@@ -1893,7 +1893,7 @@ def as_default(self):
 
 
 def create_eager_var_store():
-  if tf.contrib.eager.in_eager_mode():
+  if tf.executing_eagerly():
     return variable_scope.EagerVariableStore()
   else:
     return DummyVariableStore()
@@ -1998,7 +1998,7 @@ def summarize_features(features, num_shards=1):
 
 
 def _eager_log(level, *args):
-  if tf.contrib.eager.in_eager_mode() and args in _already_logged:
+  if tf.executing_eagerly() and args in _already_logged:
     return
   _already_logged.add(args)
   getattr(tf.logging, level)(*args)
diff --git a/tensor2tensor/utils/video/prediction2gif.py b/tensor2tensor/utils/video/prediction2gif.py
index 326d8ed06..fc55b7e4c 100644
--- a/tensor2tensor/utils/video/prediction2gif.py
+++ b/tensor2tensor/utils/video/prediction2gif.py
@@ -86,7 +86,7 @@ def main(_):
       data_dir=os.path.expanduser(FLAGS.data_dir),
       hparams=hparams)
 
-  dataset = dataset.apply(tf.contrib.data.batch_and_drop_remainder(num_agents))
+  dataset = dataset.batch(num_agents, drop_remainder=True)
   data = dataset.make_one_shot_iterator().get_next()
   # Setup input placeholders
   input_size = [num_agents, hparams.video_num_input_frames]
diff --git a/tensor2tensor/utils/video/reward_confusion.py b/tensor2tensor/utils/video/reward_confusion.py
index 1cb2c59c4..0e6c7e773 100644
--- a/tensor2tensor/utils/video/reward_confusion.py
+++ b/tensor2tensor/utils/video/reward_confusion.py
@@ -68,7 +68,7 @@ def main(_):
       shuffle_files=False,
       hparams=hparams)
 
-  dataset = dataset.apply(tf.contrib.data.batch_and_drop_remainder(batch_size))
+  dataset = dataset.batch(batch_size, drop_remainder=True)
   data = dataset.make_one_shot_iterator().get_next()
   input_data = dict((k, data[k]) for k in data.keys() if k.startswith("input"))
 
diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index caa84b774..04df03299 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -55,7 +55,7 @@ def load_videos(template, video_length, frame_shape):
   dataset_len = len(filenames)
   filenames = tf.constant(filenames)
   dataset = tf.data.Dataset.from_tensor_slices(filenames)
-  dataset = dataset.apply(tf.contrib.data.map_and_batch(
+  dataset = dataset.apply(tf.data.experimental.map_and_batch(
       lambda filename: load_image_map_function(filename, frame_shape),
       video_length, drop_remainder=True))
   return dataset, dataset_len

From eb6838d6557501d0e5dda6a4ff42e5439fdaf589 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Mon, 7 Jan 2019 10:55:19 -0800
Subject: [PATCH 1446/2720] internal merge of PR #1346

PiperOrigin-RevId: 228195515
---
 tensor2tensor/bin/t2t_attack.py                |  2 +-
 tensor2tensor/bin/t2t_avg_all.py               |  6 +++---
 tensor2tensor/data_generators/image_utils.py   |  2 +-
 tensor2tensor/data_generators/problem.py       |  8 ++++----
 tensor2tensor/layers/common_image_attention.py |  4 ++--
 .../layers/common_image_attention_test.py      |  2 +-
 tensor2tensor/layers/common_layers.py          |  8 ++++----
 tensor2tensor/layers/common_video.py           |  2 +-
 tensor2tensor/layers/discretization.py         |  4 ++--
 tensor2tensor/layers/modalities.py             |  6 +++---
 tensor2tensor/models/image_transformer.py      |  2 +-
 tensor2tensor/models/image_transformer_2d.py   |  6 +++---
 tensor2tensor/models/lstm.py                   | 14 +++++++-------
 tensor2tensor/models/research/glow_ops.py      |  4 ++--
 tensor2tensor/models/research/glow_ops_test.py |  2 +-
 .../models/research/transformer_vae.py         |  2 +-
 tensor2tensor/models/research/vqa_attention.py |  6 +++---
 tensor2tensor/models/shake_shake.py            |  6 +++---
 tensor2tensor/models/video/svg_lp.py           |  8 ++++----
 tensor2tensor/notebooks/asr_transformer.ipynb  |  4 ++--
 tensor2tensor/rl/trainer_model_based.py        |  4 ++--
 tensor2tensor/serving/serving_utils.py         |  2 +-
 tensor2tensor/utils/avg_checkpoints.py         |  6 +++---
 tensor2tensor/utils/expert_utils.py            |  4 ++--
 tensor2tensor/utils/metrics.py                 |  4 ++--
 tensor2tensor/utils/multistep_optimizer.py     |  2 +-
 tensor2tensor/utils/optimize.py                |  2 +-
 tensor2tensor/utils/registry.py                | 14 +++++++-------
 tensor2tensor/utils/t2t_model.py               | 18 +++++++++---------
 tensor2tensor/utils/video/prediction2gif.py    |  2 +-
 tensor2tensor/utils/video/reward_confusion.py  |  2 +-
 tensor2tensor/utils/video_metrics.py           |  2 +-
 32 files changed, 80 insertions(+), 80 deletions(-)

diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 6276d6b4c..b9e167238 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -200,7 +200,7 @@ def main(argv):
     sur_ch_model.get_probs(inputs)
 
     checkpoint_path = os.path.expanduser(FLAGS.surrogate_output_dir)
-    tf.train.init_from_checkpoint(
+    tf.contrib.framework.init_from_checkpoint(
         tf.train.latest_checkpoint(checkpoint_path), {"/": "surrogate/"})
     sess.run(tf.global_variables_initializer())
 
diff --git a/tensor2tensor/bin/t2t_avg_all.py b/tensor2tensor/bin/t2t_avg_all.py
index e3c34be90..080670f35 100644
--- a/tensor2tensor/bin/t2t_avg_all.py
+++ b/tensor2tensor/bin/t2t_avg_all.py
@@ -60,7 +60,7 @@ def main(_):
   for model in bleu_hook.stepfiles_iterator(model_dir, FLAGS.wait_minutes,
                                             FLAGS.min_steps):
     if models_processed == 0:
-      var_list = tf.train.list_variables(model.filename)
+      var_list = tf.contrib.framework.list_variables(model.filename)
       avg_values = {}
       for (name, shape) in var_list:
         if not (name.startswith("global_step") or
@@ -69,7 +69,7 @@ def main(_):
     models_processed += 1
 
     tf.logging.info("Loading [%d]: %s" % (models_processed, model.filename))
-    reader = tf.train.load_checkpoint(model.filename)
+    reader = tf.contrib.framework.load_checkpoint(model.filename)
     for name in avg_values:
       avg_values[name] += reader.get_tensor(name) / FLAGS.n
     queue.append(model)
@@ -106,7 +106,7 @@ def main(_):
     tf.reset_default_graph()
     first_model = queue.popleft()
 
-    reader = tf.train.load_checkpoint(first_model.filename)
+    reader = tf.contrib.framework.load_checkpoint(first_model.filename)
     for name in avg_values:
       avg_values[name] -= reader.get_tensor(name) / FLAGS.n
 
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index db5ddbb44..5396d931c 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -262,7 +262,7 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
 def encode_images_as_png(images):
   """Yield images encoded as pngs."""
-  if tf.executing_eagerly():
+  if tf.contrib.eager.in_eager_mode():
     for image in images:
       yield tf.image.encode_png(image).numpy()
   else:
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 5406ddf73..219deabdf 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -423,7 +423,7 @@ def _preprocess(example):
 
     if interleave:
       dataset = dataset.apply(
-          tf.data.experimental.parallel_interleave(
+          tf.contrib.data.parallel_interleave(
               _preprocess, sloppy=True, cycle_length=8))
     else:
       dataset = dataset.flat_map(_preprocess)
@@ -674,7 +674,7 @@ def _load_records_and_preprocess(filenames):
     # Create data-set from files by parsing, pre-processing and interleaving.
     if shuffle_files:
       dataset = dataset.apply(
-          tf.data.experimental.parallel_interleave(
+          tf.contrib.data.parallel_interleave(
               _load_records_and_preprocess, sloppy=True, cycle_length=8))
     else:
       dataset = _load_records_and_preprocess(dataset)
@@ -963,7 +963,7 @@ def define_shapes(example):
           batching_scheme["batch_sizes"] = [hparams.batch_size]
           batching_scheme["boundaries"] = []
         dataset = dataset.apply(
-            tf.data.experimental.bucket_by_sequence_length(
+            tf.contrib.data.bucket_by_sequence_length(
                 data_reader.example_length, batching_scheme["boundaries"],
                 batching_scheme["batch_sizes"]))
 
@@ -1040,7 +1040,7 @@ def serving_input_fn(self, hparams):
         tf.shape(serialized_example, out_type=tf.int64)[0],
         dataset.output_shapes)
     dataset = dataset.map(standardize_shapes)
-    features = tf.data.experimental.get_single_element(dataset)
+    features = tf.contrib.data.get_single_element(dataset)
 
     if self.has_inputs:
       features.pop("targets", None)
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 27191a9dc..46608d7aa 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -501,7 +501,7 @@ def postprocess_image(x, rows, cols, hparams):
                               use_bias=True,
                               activation=None,
                               name="output_conv")
-  if (hparams.mode == tf.estimator.ModeKeys.PREDICT and
+  if (hparams.mode == tf.contrib.learn.ModeKeys.INFER and
       hparams.block_raster_scan):
     y = targets
     yshape = common_layers.shape_list(y)
@@ -547,7 +547,7 @@ def prepare_decoder(targets, hparams):
 
   # during training, images are [batch, IMG_LEN, IMG_LEN, 3].
   # At inference, they are [batch, curr_infer_length, 1, 1]
-  if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+  if hparams.mode == tf.contrib.learn.ModeKeys.INFER:
     curr_infer_length = targets_shape[1]
     if hparams.block_raster_scan:
       assert hparams.img_len*channels % hparams.query_shape[1] == 0
diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index 088725b18..e87a2e4d8 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -61,7 +61,7 @@ def testPostProcessImageInferMode(self, likelihood, num_mixtures, depth):
         block_raster_scan=True,
         hidden_size=2,
         likelihood=likelihood,
-        mode=tf.estimator.ModeKeys.PREDICT,
+        mode=tf.contrib.learn.ModeKeys.INFER,
         num_mixtures=num_mixtures,
         query_shape=[block_length, block_width],
     )
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 0068d9e86..affd409e4 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -349,7 +349,7 @@ def embedding(x,
     # On the backwards pass, we want to convert the gradient from
     # an indexed-slices to a regular tensor before sending it back to the
     # parameter server. This avoids excess computation on the parameter server.
-    if not tf.executing_eagerly():
+    if not tf.contrib.eager.in_eager_mode():
       embedding_var = convert_gradient_to_tensor(embedding_var)
     x = dropout_no_scaling(x, 1.0 - symbol_dropout_rate)
     emb_x = gather(embedding_var, x, dtype)
@@ -2868,7 +2868,7 @@ def ones_matrix_band_part(rows, cols, num_lower, num_upper, out_shape=None):
 def reshape_like_all_dims(a, b):
   """Reshapes a to match the shape of b."""
   ret = tf.reshape(a, tf.shape(b))
-  if not tf.executing_eagerly():
+  if not tf.contrib.eager.in_eager_mode():
     ret.set_shape(b.get_shape())
   return ret
 
@@ -3193,7 +3193,7 @@ def should_generate_summaries():
 def reshape_like(a, b):
   """Reshapes a to match the shape of b in all but the last dimension."""
   ret = tf.reshape(a, tf.concat([tf.shape(b)[:-1], tf.shape(a)[-1:]], 0))
-  if not tf.executing_eagerly():
+  if not tf.contrib.eager.in_eager_mode():
     ret.set_shape(b.get_shape().as_list()[:-1] + a.get_shape().as_list()[-1:])
   return ret
 
@@ -3205,7 +3205,7 @@ def summarize_video(video, prefix, max_outputs=1):
     raise ValueError("Assuming videos given as tensors in the format "
                      "[batch, time, height, width, channels] but got one "
                      "of shape: %s" % str(video_shape))
-  if tf.executing_eagerly():
+  if tf.contrib.eager.in_eager_mode():
     return
   if video.get_shape().as_list()[1] is None:
     tf.summary.image(
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index ef4568b6e..cbac48912 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -82,7 +82,7 @@ def lstm_cell(inputs,
               name=None):
   """Full LSTM cell."""
   input_shape = common_layers.shape_list(inputs)
-  cell = tf.nn.rnn_cell.LSTMCell(num_units,
+  cell = tf.contrib.rnn.LSTMCell(num_units,
                                  use_peepholes=use_peepholes,
                                  cell_clip=cell_clip,
                                  initializer=initializer,
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 2293cd8b0..050b3f504 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -473,7 +473,7 @@ def gumbel_softmax(x,
     d_dev = -tf.reduce_mean(d_variance)
     ret = s
 
-    if mode != tf.estimator.ModeKeys.TRAIN:
+    if mode != tf.contrib.learn.ModeKeys.TRAIN:
       ret = tf.reshape(maxvhot, common_layers.shape_list(s))  # Just hot @eval.
     return m, ret, d_dev * 5.0 + tf.reduce_mean(kl) * 0.002
 
@@ -822,7 +822,7 @@ def predict_bits_with_lstm(prediction_source, state_size, total_num_bits,
 
   with tf.variable_scope("predict_bits_with_lstm"):
     # Layers and cell state creation.
-    lstm_cell = tf.nn.rnn_cell.LSTMCell(state_size)
+    lstm_cell = tf.contrib.rnn.LSTMCell(state_size)
     discrete_predict = tf.layers.Dense(2**bits_at_once, name="discrete_predict")
     discrete_embed = tf.layers.Dense(state_size, name="discrete_embed")
     batch_size = common_layers.shape_list(prediction_source)[0]
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 920124589..0fe4121ad 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -93,7 +93,7 @@ def _get_weights(self, hidden_dim=None):
     else:
       ret = tf.concat(shards, 0)
     # Convert ret to tensor.
-    if not tf.executing_eagerly():
+    if not tf.contrib.eager.in_eager_mode():
       ret = common_layers.convert_gradient_to_tensor(ret)
     return ret
 
@@ -226,7 +226,7 @@ class ImageModality(modality.Modality):
 
   def bottom(self, x):
     with tf.variable_scope(self.name):
-      if not tf.executing_eagerly():
+      if not tf.contrib.eager.in_eager_mode():
         tf.summary.image(
             "inputs", common_layers.tpu_safe_image_summary(x), max_outputs=2)
       return tf.to_float(x)
@@ -234,7 +234,7 @@ def bottom(self, x):
   def targets_bottom(self, x):
     inputs = x
     with tf.variable_scope(self.name):
-      if not tf.executing_eagerly():
+      if not tf.contrib.eager.in_eager_mode():
         tf.summary.image(
             "targets_bottom",
             common_layers.tpu_safe_image_summary(inputs),
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 1f841e28b..4b52f9a42 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -55,7 +55,7 @@ def body(self, features):
                        "must be ImageChannelBottomIdentityModality and "
                        "num_channels must be 1.")
     if (not tf.get_variable_scope().reuse and
-        hparams.mode != tf.estimator.ModeKeys.PREDICT and
+        hparams.mode != tf.contrib.learn.ModeKeys.INFER and
         hparams.modality["targets"] !=
         modalities.ImageChannelBottomIdentityModality):
       tf.summary.image("targets", tf.to_float(targets), max_outputs=1)
diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 5de2584c2..7555150a7 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -46,7 +46,7 @@ def body(self, features):
     targets = features["targets"]
     targets_shape = common_layers.shape_list(targets)
     if not (tf.get_variable_scope().reuse or
-            hparams.mode == tf.estimator.ModeKeys.PREDICT):
+            hparams.mode == tf.contrib.learn.ModeKeys.INFER):
       tf.summary.image("targets", targets, max_outputs=1)
 
     decoder_input, rows, cols = cia.prepare_decoder(
@@ -76,7 +76,7 @@ def body(self, features):
     targets = features["targets"]
     inputs = features["inputs"]
     if not (tf.get_variable_scope().reuse or
-            hparams.mode == tf.estimator.ModeKeys.PREDICT):
+            hparams.mode == tf.contrib.learn.ModeKeys.INFER):
       tf.summary.image("inputs", inputs, max_outputs=1)
       tf.summary.image("targets", targets, max_outputs=1)
 
@@ -112,7 +112,7 @@ def body(self, features):
     targets = features["targets"]
     inputs = features["inputs"]
     if not (tf.get_variable_scope().reuse or
-            hparams.mode == tf.estimator.ModeKeys.PREDICT):
+            hparams.mode == tf.contrib.learn.ModeKeys.INFER):
       tf.summary.image("inputs", inputs, max_outputs=1)
       tf.summary.image("targets", targets, max_outputs=1)
 
diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 58cadea46..2f657970c 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -29,8 +29,8 @@
 
 
 def _dropout_lstm_cell(hparams, train):
-  return tf.nn.rnn_cell.DropoutWrapper(
-      tf.nn.rnn_cell.LSTMCell(hparams.hidden_size),
+  return tf.contrib.rnn.DropoutWrapper(
+      tf.contrib.rnn.LSTMCell(hparams.hidden_size),
       input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
 
 
@@ -58,7 +58,7 @@ def lstm(inputs, sequence_length, hparams, train, name, initial_state=None):
             for _ in range(hparams.num_hidden_layers)]
   with tf.variable_scope(name):
     return tf.nn.dynamic_rnn(
-        tf.nn.rnn_cell.MultiRNNCell(layers),
+        tf.contrib.rnn.MultiRNNCell(layers),
         inputs,
         sequence_length,
         initial_state=initial_state,
@@ -192,11 +192,11 @@ def lstm_bid_encoder(inputs, sequence_length, hparams, train, name):
   """Bidirectional LSTM for encoding inputs that are [batch x time x size]."""
 
   with tf.variable_scope(name):
-    cell_fw = tf.nn.rnn_cell.MultiRNNCell(
+    cell_fw = tf.contrib.rnn.MultiRNNCell(
         [_dropout_lstm_cell(hparams, train)
          for _ in range(hparams.num_hidden_layers)])
 
-    cell_bw = tf.nn.rnn_cell.MultiRNNCell(
+    cell_bw = tf.contrib.rnn.MultiRNNCell(
         [_dropout_lstm_cell(hparams, train)
          for _ in range(hparams.num_hidden_layers)])
 
@@ -213,7 +213,7 @@ def lstm_bid_encoder(inputs, sequence_length, hparams, train, name):
     encoder_states = []
 
     for i in range(hparams.num_hidden_layers):
-      if isinstance(encoder_fw_state[i], tf.nn.rnn_cell.LSTMStateTuple):
+      if isinstance(encoder_fw_state[i], tf.contrib.rnn.LSTMStateTuple):
         encoder_state_c = tf.concat(
             values=(encoder_fw_state[i].c, encoder_bw_state[i].c),
             axis=1,
@@ -222,7 +222,7 @@ def lstm_bid_encoder(inputs, sequence_length, hparams, train, name):
             values=(encoder_fw_state[i].h, encoder_bw_state[i].h),
             axis=1,
             name="encoder_fw_state_h")
-        encoder_state = tf.nn.rnn_cell.LSTMStateTuple(
+        encoder_state = tf.contrib.rnn.LSTMStateTuple(
             c=encoder_state_c, h=encoder_state_h)
       elif isinstance(encoder_fw_state[i], tf.Tensor):
         encoder_state = tf.concat(
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index ef806f6e2..a8795bc34 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -977,7 +977,7 @@ def compute_prior(name, z, latent, hparams, condition=False, state=None,
             The first-three dimensions of the latent should be the same as z.
     hparams: next_frame_glow_hparams.
     condition: Whether or not to condition the distribution on latent.
-    state: tf.nn.rnn_cell.LSTMStateTuple.
+    state: tf.contrib.rnn.LSTMStateTuple.
            the current state of a LSTM used to model the distribution. Used
            only if hparams.latent_dist_encoder = "conv_lstm".
     temperature: float, temperature with which to sample from the Gaussian.
@@ -1025,7 +1025,7 @@ def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
     eps_std: Sample x2 with the provided eps_std.
     cond_latents: optionally condition x2 on cond_latents.
     hparams: next_frame_glow hparams.
-    state: tf.nn.rnn_cell.LSTMStateTuple. Current state of the LSTM over z_2.
+    state: tf.contrib.rnn.LSTMStateTuple. Current state of the LSTM over z_2.
            Used only when hparams.latent_dist_encoder == "conv_lstm"
     condition: bool, Whether or not to condition the distribution on
                cond_latents.
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 5266cddbb..101edb2d3 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -349,7 +349,7 @@ def test_latent_dist_encoder(self, encoder="conv_lstm", skip=True,
       state_t = tf.convert_to_tensor(state_rand)
       if encoder in ["conv_net", "conv3d_net"]:
         latent_t = [latent_t, latent_t]
-      init_state = tf.nn.rnn_cell.LSTMStateTuple(state_t, state_t)
+      init_state = tf.contrib.rnn.LSTMStateTuple(state_t, state_t)
       hparams = self.get_glow_hparams()
       hparams.latent_dist_encoder = encoder
       hparams.latent_skip = skip
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 11446fd64..58db6dc5b 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -103,7 +103,7 @@ def top_k_softmax(x, k):
 def top_k_experts(x, k, hparams):
   x_shape = common_layers.shape_list(x)
   x_flat = tf.reshape(x, [-1, common_layers.shape_list(x)[-1]])
-  is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+  is_training = hparams.mode == tf.contrib.learn.ModeKeys.TRAIN
   gates, load = expert_utils.noisy_top_k_gating(
       x_flat, 2 ** hparams.z_size, is_training, k)
   gates_shape = [x_shape[0], x_shape[1], x_shape[2], 2 ** hparams.z_size]
diff --git a/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index 774387878..aaaeb3a70 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -234,10 +234,10 @@ def image_encoder(image_feat,
 
 def _get_rnn_cell(hparams):
   if hparams.rnn_type == "lstm":
-    rnn_cell = tf.nn.rnn_cell.BasicLSTMCell
+    rnn_cell = tf.contrib.rnn.BasicLSTMCell
   elif hparams.rnn_type == "lstm_layernorm":
     rnn_cell = tf.contrib.rnn.LayerNormBasicLSTMCell
-  return tf.nn.rnn_cell.DropoutWrapper(
+  return tf.contrib.rnn.DropoutWrapper(
       rnn_cell(hparams.hidden_size),
       output_keep_prob=1.0-hparams.dropout)
 
@@ -269,7 +269,7 @@ def question_encoder(question, hparams, name="encoder"):
 
     # rnn_layers = [_get_rnn_cell(hparams)
     #               for _ in range(hparams.num_rnn_layers)]
-    # rnn_multi_cell = tf.nn.rnn_cell.MultiRNNCell(rnn_layers)
+    # rnn_multi_cell = tf.contrib.rnn.MultiRNNCell(rnn_layers)
     rnn_cell = _get_rnn_cell(hparams)
     # outputs, _ = tf.nn.dynamic_rnn(
     #     rnn_cell, question, length, dtype=tf.float32)
diff --git a/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index 4af05dc21..4a83444c0 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -54,7 +54,7 @@ def shake_shake_skip_connection(x, output_filters, stride, is_training):
 def shake_shake_branch(x, output_filters, stride, rand_forward, rand_backward,
                        hparams):
   """Building a 2 branching convnet."""
-  is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+  is_training = hparams.mode == tf.contrib.learn.ModeKeys.TRAIN
   x = tf.nn.relu(x)
   x = tf.layers.conv2d(
       x,
@@ -76,7 +76,7 @@ def shake_shake_branch(x, output_filters, stride, rand_forward, rand_backward,
 
 def shake_shake_block(x, output_filters, stride, hparams):
   """Builds a full shake-shake sub layer."""
-  is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+  is_training = hparams.mode == tf.contrib.learn.ModeKeys.TRAIN
   batch_size = common_layers.shape_list(x)[0]
 
   # Generate random numbers for scaling the branches.
@@ -138,7 +138,7 @@ class ShakeShake(t2t_model.T2TModel):
 
   def body(self, features):
     hparams = self._hparams
-    is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+    is_training = hparams.mode == tf.contrib.learn.ModeKeys.TRAIN
     inputs = features["inputs"]
     assert (hparams.num_hidden_layers - 2) % 6 == 0
     assert hparams.hidden_size % 16 == 0
diff --git a/tensor2tensor/models/video/svg_lp.py b/tensor2tensor/models/video/svg_lp.py
index 52e9ec2df..7a49e549b 100644
--- a/tensor2tensor/models/video/svg_lp.py
+++ b/tensor2tensor/models/video/svg_lp.py
@@ -57,13 +57,13 @@ def rnn_model(self, hidden_size, nlayers, rnn_type, name):
     """
     layers_units = [hidden_size] * nlayers
     if rnn_type == "lstm":
-      rnn_cell = tf.nn.rnn_cell.LSTMCell
+      rnn_cell = tf.contrib.rnn.LSTMCell
     elif rnn_type == "gru":
-      rnn_cell = tf.nn.rnn_cell.GRUCell
+      rnn_cell = tf.contrib.rnn.GRUCell
     else:
-      rnn_cell = tf.nn.rnn_cell.RNNCell
+      rnn_cell = tf.contrib.rnn.RNNCell
     cells = [rnn_cell(units, name=name) for units in layers_units]
-    stacked_rnn = tf.nn.rnn_cell.MultiRNNCell(cells)
+    stacked_rnn = tf.contrib.rnn.MultiRNNCell(cells)
     return stacked_rnn
 
   def deterministic_rnn(self, cell, inputs, states, output_size, scope):
diff --git a/tensor2tensor/notebooks/asr_transformer.ipynb b/tensor2tensor/notebooks/asr_transformer.ipynb
index 82a0728a8..d1802d727 100644
--- a/tensor2tensor/notebooks/asr_transformer.ipynb
+++ b/tensor2tensor/notebooks/asr_transformer.ipynb
@@ -93,8 +93,8 @@
         "from tensor2tensor.utils import metrics\n",
         "\n",
         "# Enable TF Eager execution\n",
-        "tfe = tf.contrib.eager\n",
-        "tf.enable_eager_execution()\n",
+        "from tensorflow.contrib.eager.python import tfe\n",
+        "tfe.enable_eager_execution()\n",
         "\n",
         "# Other setup\n",
         "Modes = tf.estimator.ModeKeys\n",
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 78abce708..12c6368e1 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -177,7 +177,7 @@ def train_agent(real_env, learner, world_model_dir, hparams, epoch):
   """Train the PPO agent in the simulated environment."""
   frame_stack_size = hparams.frame_stack_size
   initial_frame_rollouts = real_env.current_epoch_rollouts(
-      split=tf.estimator.ModeKeys.TRAIN,
+      split=tf.contrib.learn.ModeKeys.TRAIN,
       minimal_rollout_frames=frame_stack_size,
   )
   # TODO(koz4k): Move this to a different module.
@@ -304,7 +304,7 @@ def initial_frame_chooser(batch_size):
       max(hparams.wm_eval_rollout_ratios) * hparams.simulated_rollout_length
   )
   rollouts = real_env.current_epoch_rollouts(
-      split=tf.estimator.ModeKeys.EVAL,
+      split=tf.contrib.learn.ModeKeys.EVAL,
       minimal_rollout_frames=(subsequence_length + frame_stack_size)
   )
 
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 52695116a..8d50a7756 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -108,7 +108,7 @@ def _make_grpc_request(examples):
     request = predict_pb2.PredictRequest()
     request.model_spec.name = servable_name
     request.inputs["input"].CopyFrom(
-        tf.make_tensor_proto(
+        tf.contrib.util.make_tensor_proto(
             [ex.SerializeToString() for ex in examples], shape=[len(examples)]))
     response = stub.Predict(request, timeout_secs)
     outputs = tf.make_ndarray(response.outputs["outputs"])
diff --git a/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index 01684a394..3d57744e6 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -75,13 +75,13 @@ def main(_):
   tf.logging.info("Reading variables and averaging checkpoints:")
   for c in checkpoints:
     tf.logging.info("%s ", c)
-  var_list = tf.train.list_variables(checkpoints[0])
+  var_list = tf.contrib.framework.list_variables(checkpoints[0])
   var_values, var_dtypes = {}, {}
   for (name, shape) in var_list:
     if not name.startswith("global_step"):
       var_values[name] = np.zeros(shape)
   for checkpoint in checkpoints:
-    reader = tf.train.load_checkpoint(checkpoint)
+    reader = tf.contrib.framework.load_checkpoint(checkpoint)
     for name in var_values:
       tensor = reader.get_tensor(name)
       var_dtypes[name] = tensor.dtype
@@ -103,7 +103,7 @@ def main(_):
 
   # Build a model consisting only of variables, set them to the average values.
   with tf.Session() as sess:
-    sess.run(tf.global_variables_initializer())
+    sess.run(tf.initialize_all_variables())
     for p, assign_op, (name, value) in zip(placeholders, assign_ops,
                                            six.iteritems(var_values)):
       sess.run(assign_op, {p: value})
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 829dcd607..01a525ded 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -636,7 +636,7 @@ def remove(self, x):
           x,
           indices=self.nonpad_ids,
       )
-      if not tf.executing_eagerly():
+      if not tf.contrib.eager.in_eager_mode():
         # This is a hack but for some reason, gather_nd return a tensor of
         # undefined shape, so the shape is set up manually
         x.set_shape([None] + x_shape[1:])
@@ -986,7 +986,7 @@ def my_fn(x):
 def flatten_all_but_last(a):
   """Flatten all dimensions of a except the last."""
   ret = tf.reshape(a, [-1, tf.shape(a)[-1]])
-  if not tf.executing_eagerly():
+  if not tf.contrib.eager.in_eager_mode():
     ret.set_shape([None] + a.get_shape().as_list()[-1:])
   return ret
 
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 614cefe57..feea42058 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -26,9 +26,9 @@
 from tensor2tensor.utils import rouge
 
 import tensorflow as tf
-from tensorflow.python.util import tf_inspect as inspect
 
-tfe = tf.contrib.eager
+from tensorflow.contrib.eager.python import tfe
+from tensorflow.python.util import tf_inspect as inspect
 
 
 class Metrics(object):
diff --git a/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
index 32dd6cdd1..3f516cd8e 100644
--- a/tensor2tensor/utils/multistep_optimizer.py
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -52,7 +52,7 @@ def _create_slots(self, var_list):
 
   def _get_iter_variable(self):
     graph = (
-        None if tf.executing_eagerly() else tf.get_default_graph())
+        None if tf.contrib.eager.in_eager_mode() else tf.get_default_graph())
     return self._get_non_slot_variable("iter", graph=graph)
 
   def _prepare(self):
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index ab06a6b2f..d6e8e3e8e 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -290,7 +290,7 @@ def get_variable_initializer(hparams):
                                value=hparams.initializer_gain,
                                hparams=hparams)
 
-  if not tf.executing_eagerly():
+  if not tf.contrib.eager.in_eager_mode():
     tf.logging.info("Using variable initializer: %s", hparams.initializer)
   if hparams.initializer == "orthogonal":
     return tf.orthogonal_initializer(gain=hparams.initializer_gain)
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index b8c3e5991..7c9225f42 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -93,7 +93,7 @@ def register_model(name=None):
   def decorator(model_cls, registration_name=None):
     """Registers & returns model_cls with registration_name or default name."""
     model_name = registration_name or default_name(model_cls)
-    if model_name in _MODELS and not tf.executing_eagerly():
+    if model_name in _MODELS and not tf.contrib.eager.in_eager_mode():
       raise LookupError("Model %s already registered." % model_name)
     model_cls.REGISTERED_NAME = model_name
     _MODELS[model_name] = model_cls
@@ -125,7 +125,7 @@ def register_hparams(name=None):
   def decorator(hp_fn, registration_name=None):
     """Registers & returns hp_fn with registration_name or default name."""
     hp_name = registration_name or default_name(hp_fn)
-    if hp_name in _HPARAMS and not tf.executing_eagerly():
+    if hp_name in _HPARAMS and not tf.contrib.eager.in_eager_mode():
       raise LookupError("HParams set %s already registered." % hp_name)
     _HPARAMS[hp_name] = hp_fn
     return hp_fn
@@ -199,7 +199,7 @@ def register_problem(name=None):
   def decorator(p_cls, registration_name=None):
     """Registers & returns p_cls with registration_name or default name."""
     p_name = registration_name or default_name(p_cls)
-    if p_name in _PROBLEMS and not tf.executing_eagerly():
+    if p_name in _PROBLEMS and not tf.contrib.eager.in_eager_mode():
       raise LookupError("Problem %s already registered." % p_name)
 
     _PROBLEMS[p_name] = p_cls
@@ -259,7 +259,7 @@ def register_attack(name=None):
   def decorator(attack_fn, registration_name=None):
     """Registers & returns attack_fn with registration_name or default name."""
     attack_name = registration_name or default_name(attack_fn)
-    if attack_name in _ATTACKS and not tf.executing_eagerly():
+    if attack_name in _ATTACKS and not tf.contrib.eager.in_eager_mode():
       raise LookupError("Attack %s already registered." % attack_name)
     _ATTACKS[attack_name] = attack_fn
     return attack_fn
@@ -299,7 +299,7 @@ def register_attack_params(name=None):
   def decorator(ap_fn, registration_name=None):
     """Registers & returns ap_fn with registration_name or default name."""
     ap_name = registration_name or default_name(ap_fn)
-    if ap_name in _ATTACK_PARAMS and not tf.executing_eagerly():
+    if ap_name in _ATTACK_PARAMS and not tf.contrib.eager.in_eager_mode():
       raise LookupError("Attack HParams set %s already registered." % ap_name)
     _ATTACK_PARAMS[ap_name] = ap_fn
     return ap_fn
@@ -338,7 +338,7 @@ def register_pruning_params(name=None):
   def decorator(pp_fn, registration_name=None):
     """Registers & returns pp_fn with registration_name or default name."""
     pp_name = registration_name or default_name(pp_fn)
-    if pp_name in _PRUNING_PARAMS and not tf.executing_eagerly():
+    if pp_name in _PRUNING_PARAMS and not tf.contrib.eager.in_eager_mode():
       raise LookupError("Pruning HParams set %s already registered." % pp_name)
     _PRUNING_PARAMS[pp_name] = pp_fn
     return pp_fn
@@ -376,7 +376,7 @@ def register_pruning_strategy(name=None):
   def decorator(ps_fn, registration_name=None):
     """Registers & returns ps_fn with registration_name or default name."""
     ps_name = registration_name or default_name(ps_fn)
-    if ps_name in _PRUNING_STRATEGY and not tf.executing_eagerly():
+    if ps_name in _PRUNING_STRATEGY and not tf.contrib.eager.in_eager_mode():
       raise LookupError("Pruning strategy %s already registered." % ps_name)
     _PRUNING_STRATEGY[ps_name] = ps_fn
     return ps_fn
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index d95bd98d7..f3530b574 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -939,7 +939,7 @@ def _slow_greedy_infer_tpu(self, features, decode_length):
 
     def infer_step(i, recent_output, recent_logits, unused_loss):
       """Inference step."""
-      if not tf.executing_eagerly():
+      if not tf.contrib.eager.in_eager_mode():
         recent_output.set_shape([None, None, None, 1])
       padded = tf.pad(recent_output, [[0, 0], [0, 1], [0, 0], [0, 0]])
       features["targets"] = padded
@@ -957,7 +957,7 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
       samples = inplace_ops.alias_inplace_update(samples, i,
                                                  tf.to_int64(cur_sample))
       samples = tf.transpose(samples, perm=[1, 0, 2, 3])
-      if not tf.executing_eagerly():
+      if not tf.contrib.eager.in_eager_mode():
         samples.set_shape([None, None, None, 1])
 
       # Assuming we have one shard for logits.
@@ -1000,7 +1000,7 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
     # tensor padded to [batch_size, decode_length, 1, 1, vocab_size]
     logits = tf.zeros((batch_size, decode_length, 1, 1,
                        target_modality.top_dimensionality))
-    if not tf.executing_eagerly():
+    if not tf.contrib.eager.in_eager_mode():
       logits.set_shape([None, None, None, None, None])
     loss = 0.0
 
@@ -1106,7 +1106,7 @@ def _slow_greedy_infer(self, features, decode_length):
 
     def infer_step(recent_output, recent_logits, unused_loss):
       """Inference step."""
-      if not tf.executing_eagerly():
+      if not tf.contrib.eager.in_eager_mode():
         if self._target_modality_is_real:
           dim = self._problem_hparams.modality["targets"].top_dimensionality
           recent_output.set_shape([None, None, None, dim])
@@ -1130,7 +1130,7 @@ def infer_step(recent_output, recent_logits, unused_loss):
       else:
         cur_sample = tf.to_int64(tf.expand_dims(cur_sample, axis=1))
         samples = tf.concat([recent_output, cur_sample], axis=1)
-        if not tf.executing_eagerly():
+        if not tf.contrib.eager.in_eager_mode():
           samples.set_shape([None, None, None, 1])
 
       # Assuming we have one shard for logits.
@@ -1176,7 +1176,7 @@ def infer_step(recent_output, recent_logits, unused_loss):
       logits = tf.zeros((batch_size, 0, 1, 1,
                          target_modality.top_dimensionality))
       logits_shape_inv = [None, None, None, None, None]
-    if not tf.executing_eagerly():
+    if not tf.contrib.eager.in_eager_mode():
       logits.set_shape(logits_shape_inv)
 
     loss = 0.0
@@ -1459,7 +1459,7 @@ def initialize_from_ckpt(self, ckpt_dir):
     log_info("Checkpoint dir: %s", ckpt_dir)
 
     # TODO(mitchellstern): Add support for partitioned variables?
-    reader = tf.train.load_checkpoint(ckpt_dir)
+    reader = tf.contrib.framework.load_checkpoint(ckpt_dir)
     variable_map = {}
     for var in tf.contrib.framework.get_trainable_variables():
       var_name = var.name.split(":")[0]
@@ -1893,7 +1893,7 @@ def as_default(self):
 
 
 def create_eager_var_store():
-  if tf.executing_eagerly():
+  if tf.contrib.eager.in_eager_mode():
     return variable_scope.EagerVariableStore()
   else:
     return DummyVariableStore()
@@ -1998,7 +1998,7 @@ def summarize_features(features, num_shards=1):
 
 
 def _eager_log(level, *args):
-  if tf.executing_eagerly() and args in _already_logged:
+  if tf.contrib.eager.in_eager_mode() and args in _already_logged:
     return
   _already_logged.add(args)
   getattr(tf.logging, level)(*args)
diff --git a/tensor2tensor/utils/video/prediction2gif.py b/tensor2tensor/utils/video/prediction2gif.py
index fc55b7e4c..326d8ed06 100644
--- a/tensor2tensor/utils/video/prediction2gif.py
+++ b/tensor2tensor/utils/video/prediction2gif.py
@@ -86,7 +86,7 @@ def main(_):
       data_dir=os.path.expanduser(FLAGS.data_dir),
       hparams=hparams)
 
-  dataset = dataset.batch(num_agents, drop_remainder=True)
+  dataset = dataset.apply(tf.contrib.data.batch_and_drop_remainder(num_agents))
   data = dataset.make_one_shot_iterator().get_next()
   # Setup input placeholders
   input_size = [num_agents, hparams.video_num_input_frames]
diff --git a/tensor2tensor/utils/video/reward_confusion.py b/tensor2tensor/utils/video/reward_confusion.py
index 0e6c7e773..1cb2c59c4 100644
--- a/tensor2tensor/utils/video/reward_confusion.py
+++ b/tensor2tensor/utils/video/reward_confusion.py
@@ -68,7 +68,7 @@ def main(_):
       shuffle_files=False,
       hparams=hparams)
 
-  dataset = dataset.batch(batch_size, drop_remainder=True)
+  dataset = dataset.apply(tf.contrib.data.batch_and_drop_remainder(batch_size))
   data = dataset.make_one_shot_iterator().get_next()
   input_data = dict((k, data[k]) for k in data.keys() if k.startswith("input"))
 
diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index 04df03299..caa84b774 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -55,7 +55,7 @@ def load_videos(template, video_length, frame_shape):
   dataset_len = len(filenames)
   filenames = tf.constant(filenames)
   dataset = tf.data.Dataset.from_tensor_slices(filenames)
-  dataset = dataset.apply(tf.data.experimental.map_and_batch(
+  dataset = dataset.apply(tf.contrib.data.map_and_batch(
       lambda filename: load_image_map_function(filename, frame_shape),
       video_length, drop_remainder=True))
   return dataset, dataset_len

From 0b1b28c1adee063f3ef197ee17e73cf18211ae68 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Mon, 7 Jan 2019 13:46:11 -0800
Subject: [PATCH 1447/2720] internal merge of PR #1345

PiperOrigin-RevId: 228228161
---
 tensor2tensor/bin/t2t_attack.py                |  2 +-
 tensor2tensor/bin/t2t_avg_all.py               |  6 +++---
 tensor2tensor/data_generators/image_utils.py   |  2 +-
 tensor2tensor/data_generators/problem.py       |  8 ++++----
 tensor2tensor/layers/common_image_attention.py |  4 ++--
 .../layers/common_image_attention_test.py      |  2 +-
 tensor2tensor/layers/common_layers.py          |  8 ++++----
 tensor2tensor/layers/common_video.py           |  2 +-
 tensor2tensor/layers/discretization.py         |  4 ++--
 tensor2tensor/layers/modalities.py             |  6 +++---
 tensor2tensor/models/image_transformer.py      |  2 +-
 tensor2tensor/models/image_transformer_2d.py   |  6 +++---
 tensor2tensor/models/lstm.py                   | 14 +++++++-------
 tensor2tensor/models/research/glow_ops.py      |  4 ++--
 tensor2tensor/models/research/glow_ops_test.py |  2 +-
 .../models/research/transformer_vae.py         |  2 +-
 tensor2tensor/models/research/vqa_attention.py |  6 +++---
 tensor2tensor/models/shake_shake.py            |  6 +++---
 tensor2tensor/models/video/svg_lp.py           |  8 ++++----
 tensor2tensor/notebooks/asr_transformer.ipynb  |  4 ++--
 tensor2tensor/rl/trainer_model_based.py        |  4 ++--
 tensor2tensor/serving/serving_utils.py         |  2 +-
 tensor2tensor/utils/avg_checkpoints.py         |  6 +++---
 tensor2tensor/utils/expert_utils.py            |  4 ++--
 tensor2tensor/utils/multistep_optimizer.py     |  2 +-
 tensor2tensor/utils/optimize.py                |  2 +-
 tensor2tensor/utils/registry.py                | 14 +++++++-------
 tensor2tensor/utils/t2t_model.py               | 18 +++++++++---------
 tensor2tensor/utils/video/prediction2gif.py    |  2 +-
 tensor2tensor/utils/video/reward_confusion.py  |  2 +-
 tensor2tensor/utils/video_metrics.py           |  2 +-
 31 files changed, 78 insertions(+), 78 deletions(-)

diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index b9e167238..6276d6b4c 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -200,7 +200,7 @@ def main(argv):
     sur_ch_model.get_probs(inputs)
 
     checkpoint_path = os.path.expanduser(FLAGS.surrogate_output_dir)
-    tf.contrib.framework.init_from_checkpoint(
+    tf.train.init_from_checkpoint(
         tf.train.latest_checkpoint(checkpoint_path), {"/": "surrogate/"})
     sess.run(tf.global_variables_initializer())
 
diff --git a/tensor2tensor/bin/t2t_avg_all.py b/tensor2tensor/bin/t2t_avg_all.py
index 080670f35..e3c34be90 100644
--- a/tensor2tensor/bin/t2t_avg_all.py
+++ b/tensor2tensor/bin/t2t_avg_all.py
@@ -60,7 +60,7 @@ def main(_):
   for model in bleu_hook.stepfiles_iterator(model_dir, FLAGS.wait_minutes,
                                             FLAGS.min_steps):
     if models_processed == 0:
-      var_list = tf.contrib.framework.list_variables(model.filename)
+      var_list = tf.train.list_variables(model.filename)
       avg_values = {}
       for (name, shape) in var_list:
         if not (name.startswith("global_step") or
@@ -69,7 +69,7 @@ def main(_):
     models_processed += 1
 
     tf.logging.info("Loading [%d]: %s" % (models_processed, model.filename))
-    reader = tf.contrib.framework.load_checkpoint(model.filename)
+    reader = tf.train.load_checkpoint(model.filename)
     for name in avg_values:
       avg_values[name] += reader.get_tensor(name) / FLAGS.n
     queue.append(model)
@@ -106,7 +106,7 @@ def main(_):
     tf.reset_default_graph()
     first_model = queue.popleft()
 
-    reader = tf.contrib.framework.load_checkpoint(first_model.filename)
+    reader = tf.train.load_checkpoint(first_model.filename)
     for name in avg_values:
       avg_values[name] -= reader.get_tensor(name) / FLAGS.n
 
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 5396d931c..db5ddbb44 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -262,7 +262,7 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
 def encode_images_as_png(images):
   """Yield images encoded as pngs."""
-  if tf.contrib.eager.in_eager_mode():
+  if tf.executing_eagerly():
     for image in images:
       yield tf.image.encode_png(image).numpy()
   else:
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 219deabdf..5406ddf73 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -423,7 +423,7 @@ def _preprocess(example):
 
     if interleave:
       dataset = dataset.apply(
-          tf.contrib.data.parallel_interleave(
+          tf.data.experimental.parallel_interleave(
               _preprocess, sloppy=True, cycle_length=8))
     else:
       dataset = dataset.flat_map(_preprocess)
@@ -674,7 +674,7 @@ def _load_records_and_preprocess(filenames):
     # Create data-set from files by parsing, pre-processing and interleaving.
     if shuffle_files:
       dataset = dataset.apply(
-          tf.contrib.data.parallel_interleave(
+          tf.data.experimental.parallel_interleave(
               _load_records_and_preprocess, sloppy=True, cycle_length=8))
     else:
       dataset = _load_records_and_preprocess(dataset)
@@ -963,7 +963,7 @@ def define_shapes(example):
           batching_scheme["batch_sizes"] = [hparams.batch_size]
           batching_scheme["boundaries"] = []
         dataset = dataset.apply(
-            tf.contrib.data.bucket_by_sequence_length(
+            tf.data.experimental.bucket_by_sequence_length(
                 data_reader.example_length, batching_scheme["boundaries"],
                 batching_scheme["batch_sizes"]))
 
@@ -1040,7 +1040,7 @@ def serving_input_fn(self, hparams):
         tf.shape(serialized_example, out_type=tf.int64)[0],
         dataset.output_shapes)
     dataset = dataset.map(standardize_shapes)
-    features = tf.contrib.data.get_single_element(dataset)
+    features = tf.data.experimental.get_single_element(dataset)
 
     if self.has_inputs:
       features.pop("targets", None)
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 46608d7aa..27191a9dc 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -501,7 +501,7 @@ def postprocess_image(x, rows, cols, hparams):
                               use_bias=True,
                               activation=None,
                               name="output_conv")
-  if (hparams.mode == tf.contrib.learn.ModeKeys.INFER and
+  if (hparams.mode == tf.estimator.ModeKeys.PREDICT and
       hparams.block_raster_scan):
     y = targets
     yshape = common_layers.shape_list(y)
@@ -547,7 +547,7 @@ def prepare_decoder(targets, hparams):
 
   # during training, images are [batch, IMG_LEN, IMG_LEN, 3].
   # At inference, they are [batch, curr_infer_length, 1, 1]
-  if hparams.mode == tf.contrib.learn.ModeKeys.INFER:
+  if hparams.mode == tf.estimator.ModeKeys.PREDICT:
     curr_infer_length = targets_shape[1]
     if hparams.block_raster_scan:
       assert hparams.img_len*channels % hparams.query_shape[1] == 0
diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index e87a2e4d8..088725b18 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -61,7 +61,7 @@ def testPostProcessImageInferMode(self, likelihood, num_mixtures, depth):
         block_raster_scan=True,
         hidden_size=2,
         likelihood=likelihood,
-        mode=tf.contrib.learn.ModeKeys.INFER,
+        mode=tf.estimator.ModeKeys.PREDICT,
         num_mixtures=num_mixtures,
         query_shape=[block_length, block_width],
     )
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index affd409e4..0068d9e86 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -349,7 +349,7 @@ def embedding(x,
     # On the backwards pass, we want to convert the gradient from
     # an indexed-slices to a regular tensor before sending it back to the
     # parameter server. This avoids excess computation on the parameter server.
-    if not tf.contrib.eager.in_eager_mode():
+    if not tf.executing_eagerly():
       embedding_var = convert_gradient_to_tensor(embedding_var)
     x = dropout_no_scaling(x, 1.0 - symbol_dropout_rate)
     emb_x = gather(embedding_var, x, dtype)
@@ -2868,7 +2868,7 @@ def ones_matrix_band_part(rows, cols, num_lower, num_upper, out_shape=None):
 def reshape_like_all_dims(a, b):
   """Reshapes a to match the shape of b."""
   ret = tf.reshape(a, tf.shape(b))
-  if not tf.contrib.eager.in_eager_mode():
+  if not tf.executing_eagerly():
     ret.set_shape(b.get_shape())
   return ret
 
@@ -3193,7 +3193,7 @@ def should_generate_summaries():
 def reshape_like(a, b):
   """Reshapes a to match the shape of b in all but the last dimension."""
   ret = tf.reshape(a, tf.concat([tf.shape(b)[:-1], tf.shape(a)[-1:]], 0))
-  if not tf.contrib.eager.in_eager_mode():
+  if not tf.executing_eagerly():
     ret.set_shape(b.get_shape().as_list()[:-1] + a.get_shape().as_list()[-1:])
   return ret
 
@@ -3205,7 +3205,7 @@ def summarize_video(video, prefix, max_outputs=1):
     raise ValueError("Assuming videos given as tensors in the format "
                      "[batch, time, height, width, channels] but got one "
                      "of shape: %s" % str(video_shape))
-  if tf.contrib.eager.in_eager_mode():
+  if tf.executing_eagerly():
     return
   if video.get_shape().as_list()[1] is None:
     tf.summary.image(
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index cbac48912..ef4568b6e 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -82,7 +82,7 @@ def lstm_cell(inputs,
               name=None):
   """Full LSTM cell."""
   input_shape = common_layers.shape_list(inputs)
-  cell = tf.contrib.rnn.LSTMCell(num_units,
+  cell = tf.nn.rnn_cell.LSTMCell(num_units,
                                  use_peepholes=use_peepholes,
                                  cell_clip=cell_clip,
                                  initializer=initializer,
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 050b3f504..2293cd8b0 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -473,7 +473,7 @@ def gumbel_softmax(x,
     d_dev = -tf.reduce_mean(d_variance)
     ret = s
 
-    if mode != tf.contrib.learn.ModeKeys.TRAIN:
+    if mode != tf.estimator.ModeKeys.TRAIN:
       ret = tf.reshape(maxvhot, common_layers.shape_list(s))  # Just hot @eval.
     return m, ret, d_dev * 5.0 + tf.reduce_mean(kl) * 0.002
 
@@ -822,7 +822,7 @@ def predict_bits_with_lstm(prediction_source, state_size, total_num_bits,
 
   with tf.variable_scope("predict_bits_with_lstm"):
     # Layers and cell state creation.
-    lstm_cell = tf.contrib.rnn.LSTMCell(state_size)
+    lstm_cell = tf.nn.rnn_cell.LSTMCell(state_size)
     discrete_predict = tf.layers.Dense(2**bits_at_once, name="discrete_predict")
     discrete_embed = tf.layers.Dense(state_size, name="discrete_embed")
     batch_size = common_layers.shape_list(prediction_source)[0]
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 0fe4121ad..920124589 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -93,7 +93,7 @@ def _get_weights(self, hidden_dim=None):
     else:
       ret = tf.concat(shards, 0)
     # Convert ret to tensor.
-    if not tf.contrib.eager.in_eager_mode():
+    if not tf.executing_eagerly():
       ret = common_layers.convert_gradient_to_tensor(ret)
     return ret
 
@@ -226,7 +226,7 @@ class ImageModality(modality.Modality):
 
   def bottom(self, x):
     with tf.variable_scope(self.name):
-      if not tf.contrib.eager.in_eager_mode():
+      if not tf.executing_eagerly():
         tf.summary.image(
             "inputs", common_layers.tpu_safe_image_summary(x), max_outputs=2)
       return tf.to_float(x)
@@ -234,7 +234,7 @@ def bottom(self, x):
   def targets_bottom(self, x):
     inputs = x
     with tf.variable_scope(self.name):
-      if not tf.contrib.eager.in_eager_mode():
+      if not tf.executing_eagerly():
         tf.summary.image(
             "targets_bottom",
             common_layers.tpu_safe_image_summary(inputs),
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 4b52f9a42..1f841e28b 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -55,7 +55,7 @@ def body(self, features):
                        "must be ImageChannelBottomIdentityModality and "
                        "num_channels must be 1.")
     if (not tf.get_variable_scope().reuse and
-        hparams.mode != tf.contrib.learn.ModeKeys.INFER and
+        hparams.mode != tf.estimator.ModeKeys.PREDICT and
         hparams.modality["targets"] !=
         modalities.ImageChannelBottomIdentityModality):
       tf.summary.image("targets", tf.to_float(targets), max_outputs=1)
diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 7555150a7..5de2584c2 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -46,7 +46,7 @@ def body(self, features):
     targets = features["targets"]
     targets_shape = common_layers.shape_list(targets)
     if not (tf.get_variable_scope().reuse or
-            hparams.mode == tf.contrib.learn.ModeKeys.INFER):
+            hparams.mode == tf.estimator.ModeKeys.PREDICT):
       tf.summary.image("targets", targets, max_outputs=1)
 
     decoder_input, rows, cols = cia.prepare_decoder(
@@ -76,7 +76,7 @@ def body(self, features):
     targets = features["targets"]
     inputs = features["inputs"]
     if not (tf.get_variable_scope().reuse or
-            hparams.mode == tf.contrib.learn.ModeKeys.INFER):
+            hparams.mode == tf.estimator.ModeKeys.PREDICT):
       tf.summary.image("inputs", inputs, max_outputs=1)
       tf.summary.image("targets", targets, max_outputs=1)
 
@@ -112,7 +112,7 @@ def body(self, features):
     targets = features["targets"]
     inputs = features["inputs"]
     if not (tf.get_variable_scope().reuse or
-            hparams.mode == tf.contrib.learn.ModeKeys.INFER):
+            hparams.mode == tf.estimator.ModeKeys.PREDICT):
       tf.summary.image("inputs", inputs, max_outputs=1)
       tf.summary.image("targets", targets, max_outputs=1)
 
diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 2f657970c..58cadea46 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -29,8 +29,8 @@
 
 
 def _dropout_lstm_cell(hparams, train):
-  return tf.contrib.rnn.DropoutWrapper(
-      tf.contrib.rnn.LSTMCell(hparams.hidden_size),
+  return tf.nn.rnn_cell.DropoutWrapper(
+      tf.nn.rnn_cell.LSTMCell(hparams.hidden_size),
       input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
 
 
@@ -58,7 +58,7 @@ def lstm(inputs, sequence_length, hparams, train, name, initial_state=None):
             for _ in range(hparams.num_hidden_layers)]
   with tf.variable_scope(name):
     return tf.nn.dynamic_rnn(
-        tf.contrib.rnn.MultiRNNCell(layers),
+        tf.nn.rnn_cell.MultiRNNCell(layers),
         inputs,
         sequence_length,
         initial_state=initial_state,
@@ -192,11 +192,11 @@ def lstm_bid_encoder(inputs, sequence_length, hparams, train, name):
   """Bidirectional LSTM for encoding inputs that are [batch x time x size]."""
 
   with tf.variable_scope(name):
-    cell_fw = tf.contrib.rnn.MultiRNNCell(
+    cell_fw = tf.nn.rnn_cell.MultiRNNCell(
         [_dropout_lstm_cell(hparams, train)
          for _ in range(hparams.num_hidden_layers)])
 
-    cell_bw = tf.contrib.rnn.MultiRNNCell(
+    cell_bw = tf.nn.rnn_cell.MultiRNNCell(
         [_dropout_lstm_cell(hparams, train)
          for _ in range(hparams.num_hidden_layers)])
 
@@ -213,7 +213,7 @@ def lstm_bid_encoder(inputs, sequence_length, hparams, train, name):
     encoder_states = []
 
     for i in range(hparams.num_hidden_layers):
-      if isinstance(encoder_fw_state[i], tf.contrib.rnn.LSTMStateTuple):
+      if isinstance(encoder_fw_state[i], tf.nn.rnn_cell.LSTMStateTuple):
         encoder_state_c = tf.concat(
             values=(encoder_fw_state[i].c, encoder_bw_state[i].c),
             axis=1,
@@ -222,7 +222,7 @@ def lstm_bid_encoder(inputs, sequence_length, hparams, train, name):
             values=(encoder_fw_state[i].h, encoder_bw_state[i].h),
             axis=1,
             name="encoder_fw_state_h")
-        encoder_state = tf.contrib.rnn.LSTMStateTuple(
+        encoder_state = tf.nn.rnn_cell.LSTMStateTuple(
             c=encoder_state_c, h=encoder_state_h)
       elif isinstance(encoder_fw_state[i], tf.Tensor):
         encoder_state = tf.concat(
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index a8795bc34..ef806f6e2 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -977,7 +977,7 @@ def compute_prior(name, z, latent, hparams, condition=False, state=None,
             The first-three dimensions of the latent should be the same as z.
     hparams: next_frame_glow_hparams.
     condition: Whether or not to condition the distribution on latent.
-    state: tf.contrib.rnn.LSTMStateTuple.
+    state: tf.nn.rnn_cell.LSTMStateTuple.
            the current state of a LSTM used to model the distribution. Used
            only if hparams.latent_dist_encoder = "conv_lstm".
     temperature: float, temperature with which to sample from the Gaussian.
@@ -1025,7 +1025,7 @@ def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
     eps_std: Sample x2 with the provided eps_std.
     cond_latents: optionally condition x2 on cond_latents.
     hparams: next_frame_glow hparams.
-    state: tf.contrib.rnn.LSTMStateTuple. Current state of the LSTM over z_2.
+    state: tf.nn.rnn_cell.LSTMStateTuple. Current state of the LSTM over z_2.
            Used only when hparams.latent_dist_encoder == "conv_lstm"
     condition: bool, Whether or not to condition the distribution on
                cond_latents.
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 101edb2d3..5266cddbb 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -349,7 +349,7 @@ def test_latent_dist_encoder(self, encoder="conv_lstm", skip=True,
       state_t = tf.convert_to_tensor(state_rand)
       if encoder in ["conv_net", "conv3d_net"]:
         latent_t = [latent_t, latent_t]
-      init_state = tf.contrib.rnn.LSTMStateTuple(state_t, state_t)
+      init_state = tf.nn.rnn_cell.LSTMStateTuple(state_t, state_t)
       hparams = self.get_glow_hparams()
       hparams.latent_dist_encoder = encoder
       hparams.latent_skip = skip
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 58db6dc5b..11446fd64 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -103,7 +103,7 @@ def top_k_softmax(x, k):
 def top_k_experts(x, k, hparams):
   x_shape = common_layers.shape_list(x)
   x_flat = tf.reshape(x, [-1, common_layers.shape_list(x)[-1]])
-  is_training = hparams.mode == tf.contrib.learn.ModeKeys.TRAIN
+  is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
   gates, load = expert_utils.noisy_top_k_gating(
       x_flat, 2 ** hparams.z_size, is_training, k)
   gates_shape = [x_shape[0], x_shape[1], x_shape[2], 2 ** hparams.z_size]
diff --git a/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index aaaeb3a70..774387878 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -234,10 +234,10 @@ def image_encoder(image_feat,
 
 def _get_rnn_cell(hparams):
   if hparams.rnn_type == "lstm":
-    rnn_cell = tf.contrib.rnn.BasicLSTMCell
+    rnn_cell = tf.nn.rnn_cell.BasicLSTMCell
   elif hparams.rnn_type == "lstm_layernorm":
     rnn_cell = tf.contrib.rnn.LayerNormBasicLSTMCell
-  return tf.contrib.rnn.DropoutWrapper(
+  return tf.nn.rnn_cell.DropoutWrapper(
       rnn_cell(hparams.hidden_size),
       output_keep_prob=1.0-hparams.dropout)
 
@@ -269,7 +269,7 @@ def question_encoder(question, hparams, name="encoder"):
 
     # rnn_layers = [_get_rnn_cell(hparams)
     #               for _ in range(hparams.num_rnn_layers)]
-    # rnn_multi_cell = tf.contrib.rnn.MultiRNNCell(rnn_layers)
+    # rnn_multi_cell = tf.nn.rnn_cell.MultiRNNCell(rnn_layers)
     rnn_cell = _get_rnn_cell(hparams)
     # outputs, _ = tf.nn.dynamic_rnn(
     #     rnn_cell, question, length, dtype=tf.float32)
diff --git a/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index 4a83444c0..4af05dc21 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -54,7 +54,7 @@ def shake_shake_skip_connection(x, output_filters, stride, is_training):
 def shake_shake_branch(x, output_filters, stride, rand_forward, rand_backward,
                        hparams):
   """Building a 2 branching convnet."""
-  is_training = hparams.mode == tf.contrib.learn.ModeKeys.TRAIN
+  is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
   x = tf.nn.relu(x)
   x = tf.layers.conv2d(
       x,
@@ -76,7 +76,7 @@ def shake_shake_branch(x, output_filters, stride, rand_forward, rand_backward,
 
 def shake_shake_block(x, output_filters, stride, hparams):
   """Builds a full shake-shake sub layer."""
-  is_training = hparams.mode == tf.contrib.learn.ModeKeys.TRAIN
+  is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
   batch_size = common_layers.shape_list(x)[0]
 
   # Generate random numbers for scaling the branches.
@@ -138,7 +138,7 @@ class ShakeShake(t2t_model.T2TModel):
 
   def body(self, features):
     hparams = self._hparams
-    is_training = hparams.mode == tf.contrib.learn.ModeKeys.TRAIN
+    is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
     inputs = features["inputs"]
     assert (hparams.num_hidden_layers - 2) % 6 == 0
     assert hparams.hidden_size % 16 == 0
diff --git a/tensor2tensor/models/video/svg_lp.py b/tensor2tensor/models/video/svg_lp.py
index 7a49e549b..52e9ec2df 100644
--- a/tensor2tensor/models/video/svg_lp.py
+++ b/tensor2tensor/models/video/svg_lp.py
@@ -57,13 +57,13 @@ def rnn_model(self, hidden_size, nlayers, rnn_type, name):
     """
     layers_units = [hidden_size] * nlayers
     if rnn_type == "lstm":
-      rnn_cell = tf.contrib.rnn.LSTMCell
+      rnn_cell = tf.nn.rnn_cell.LSTMCell
     elif rnn_type == "gru":
-      rnn_cell = tf.contrib.rnn.GRUCell
+      rnn_cell = tf.nn.rnn_cell.GRUCell
     else:
-      rnn_cell = tf.contrib.rnn.RNNCell
+      rnn_cell = tf.nn.rnn_cell.RNNCell
     cells = [rnn_cell(units, name=name) for units in layers_units]
-    stacked_rnn = tf.contrib.rnn.MultiRNNCell(cells)
+    stacked_rnn = tf.nn.rnn_cell.MultiRNNCell(cells)
     return stacked_rnn
 
   def deterministic_rnn(self, cell, inputs, states, output_size, scope):
diff --git a/tensor2tensor/notebooks/asr_transformer.ipynb b/tensor2tensor/notebooks/asr_transformer.ipynb
index d1802d727..82a0728a8 100644
--- a/tensor2tensor/notebooks/asr_transformer.ipynb
+++ b/tensor2tensor/notebooks/asr_transformer.ipynb
@@ -93,8 +93,8 @@
         "from tensor2tensor.utils import metrics\n",
         "\n",
         "# Enable TF Eager execution\n",
-        "from tensorflow.contrib.eager.python import tfe\n",
-        "tfe.enable_eager_execution()\n",
+        "tfe = tf.contrib.eager\n",
+        "tf.enable_eager_execution()\n",
         "\n",
         "# Other setup\n",
         "Modes = tf.estimator.ModeKeys\n",
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 12c6368e1..78abce708 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -177,7 +177,7 @@ def train_agent(real_env, learner, world_model_dir, hparams, epoch):
   """Train the PPO agent in the simulated environment."""
   frame_stack_size = hparams.frame_stack_size
   initial_frame_rollouts = real_env.current_epoch_rollouts(
-      split=tf.contrib.learn.ModeKeys.TRAIN,
+      split=tf.estimator.ModeKeys.TRAIN,
       minimal_rollout_frames=frame_stack_size,
   )
   # TODO(koz4k): Move this to a different module.
@@ -304,7 +304,7 @@ def initial_frame_chooser(batch_size):
       max(hparams.wm_eval_rollout_ratios) * hparams.simulated_rollout_length
   )
   rollouts = real_env.current_epoch_rollouts(
-      split=tf.contrib.learn.ModeKeys.EVAL,
+      split=tf.estimator.ModeKeys.EVAL,
       minimal_rollout_frames=(subsequence_length + frame_stack_size)
   )
 
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 8d50a7756..52695116a 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -108,7 +108,7 @@ def _make_grpc_request(examples):
     request = predict_pb2.PredictRequest()
     request.model_spec.name = servable_name
     request.inputs["input"].CopyFrom(
-        tf.contrib.util.make_tensor_proto(
+        tf.make_tensor_proto(
             [ex.SerializeToString() for ex in examples], shape=[len(examples)]))
     response = stub.Predict(request, timeout_secs)
     outputs = tf.make_ndarray(response.outputs["outputs"])
diff --git a/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index 3d57744e6..01684a394 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -75,13 +75,13 @@ def main(_):
   tf.logging.info("Reading variables and averaging checkpoints:")
   for c in checkpoints:
     tf.logging.info("%s ", c)
-  var_list = tf.contrib.framework.list_variables(checkpoints[0])
+  var_list = tf.train.list_variables(checkpoints[0])
   var_values, var_dtypes = {}, {}
   for (name, shape) in var_list:
     if not name.startswith("global_step"):
       var_values[name] = np.zeros(shape)
   for checkpoint in checkpoints:
-    reader = tf.contrib.framework.load_checkpoint(checkpoint)
+    reader = tf.train.load_checkpoint(checkpoint)
     for name in var_values:
       tensor = reader.get_tensor(name)
       var_dtypes[name] = tensor.dtype
@@ -103,7 +103,7 @@ def main(_):
 
   # Build a model consisting only of variables, set them to the average values.
   with tf.Session() as sess:
-    sess.run(tf.initialize_all_variables())
+    sess.run(tf.global_variables_initializer())
     for p, assign_op, (name, value) in zip(placeholders, assign_ops,
                                            six.iteritems(var_values)):
       sess.run(assign_op, {p: value})
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 01a525ded..829dcd607 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -636,7 +636,7 @@ def remove(self, x):
           x,
           indices=self.nonpad_ids,
       )
-      if not tf.contrib.eager.in_eager_mode():
+      if not tf.executing_eagerly():
         # This is a hack but for some reason, gather_nd return a tensor of
         # undefined shape, so the shape is set up manually
         x.set_shape([None] + x_shape[1:])
@@ -986,7 +986,7 @@ def my_fn(x):
 def flatten_all_but_last(a):
   """Flatten all dimensions of a except the last."""
   ret = tf.reshape(a, [-1, tf.shape(a)[-1]])
-  if not tf.contrib.eager.in_eager_mode():
+  if not tf.executing_eagerly():
     ret.set_shape([None] + a.get_shape().as_list()[-1:])
   return ret
 
diff --git a/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
index 3f516cd8e..32dd6cdd1 100644
--- a/tensor2tensor/utils/multistep_optimizer.py
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -52,7 +52,7 @@ def _create_slots(self, var_list):
 
   def _get_iter_variable(self):
     graph = (
-        None if tf.contrib.eager.in_eager_mode() else tf.get_default_graph())
+        None if tf.executing_eagerly() else tf.get_default_graph())
     return self._get_non_slot_variable("iter", graph=graph)
 
   def _prepare(self):
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index d6e8e3e8e..ab06a6b2f 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -290,7 +290,7 @@ def get_variable_initializer(hparams):
                                value=hparams.initializer_gain,
                                hparams=hparams)
 
-  if not tf.contrib.eager.in_eager_mode():
+  if not tf.executing_eagerly():
     tf.logging.info("Using variable initializer: %s", hparams.initializer)
   if hparams.initializer == "orthogonal":
     return tf.orthogonal_initializer(gain=hparams.initializer_gain)
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 7c9225f42..b8c3e5991 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -93,7 +93,7 @@ def register_model(name=None):
   def decorator(model_cls, registration_name=None):
     """Registers & returns model_cls with registration_name or default name."""
     model_name = registration_name or default_name(model_cls)
-    if model_name in _MODELS and not tf.contrib.eager.in_eager_mode():
+    if model_name in _MODELS and not tf.executing_eagerly():
       raise LookupError("Model %s already registered." % model_name)
     model_cls.REGISTERED_NAME = model_name
     _MODELS[model_name] = model_cls
@@ -125,7 +125,7 @@ def register_hparams(name=None):
   def decorator(hp_fn, registration_name=None):
     """Registers & returns hp_fn with registration_name or default name."""
     hp_name = registration_name or default_name(hp_fn)
-    if hp_name in _HPARAMS and not tf.contrib.eager.in_eager_mode():
+    if hp_name in _HPARAMS and not tf.executing_eagerly():
       raise LookupError("HParams set %s already registered." % hp_name)
     _HPARAMS[hp_name] = hp_fn
     return hp_fn
@@ -199,7 +199,7 @@ def register_problem(name=None):
   def decorator(p_cls, registration_name=None):
     """Registers & returns p_cls with registration_name or default name."""
     p_name = registration_name or default_name(p_cls)
-    if p_name in _PROBLEMS and not tf.contrib.eager.in_eager_mode():
+    if p_name in _PROBLEMS and not tf.executing_eagerly():
       raise LookupError("Problem %s already registered." % p_name)
 
     _PROBLEMS[p_name] = p_cls
@@ -259,7 +259,7 @@ def register_attack(name=None):
   def decorator(attack_fn, registration_name=None):
     """Registers & returns attack_fn with registration_name or default name."""
     attack_name = registration_name or default_name(attack_fn)
-    if attack_name in _ATTACKS and not tf.contrib.eager.in_eager_mode():
+    if attack_name in _ATTACKS and not tf.executing_eagerly():
       raise LookupError("Attack %s already registered." % attack_name)
     _ATTACKS[attack_name] = attack_fn
     return attack_fn
@@ -299,7 +299,7 @@ def register_attack_params(name=None):
   def decorator(ap_fn, registration_name=None):
     """Registers & returns ap_fn with registration_name or default name."""
     ap_name = registration_name or default_name(ap_fn)
-    if ap_name in _ATTACK_PARAMS and not tf.contrib.eager.in_eager_mode():
+    if ap_name in _ATTACK_PARAMS and not tf.executing_eagerly():
       raise LookupError("Attack HParams set %s already registered." % ap_name)
     _ATTACK_PARAMS[ap_name] = ap_fn
     return ap_fn
@@ -338,7 +338,7 @@ def register_pruning_params(name=None):
   def decorator(pp_fn, registration_name=None):
     """Registers & returns pp_fn with registration_name or default name."""
     pp_name = registration_name or default_name(pp_fn)
-    if pp_name in _PRUNING_PARAMS and not tf.contrib.eager.in_eager_mode():
+    if pp_name in _PRUNING_PARAMS and not tf.executing_eagerly():
       raise LookupError("Pruning HParams set %s already registered." % pp_name)
     _PRUNING_PARAMS[pp_name] = pp_fn
     return pp_fn
@@ -376,7 +376,7 @@ def register_pruning_strategy(name=None):
   def decorator(ps_fn, registration_name=None):
     """Registers & returns ps_fn with registration_name or default name."""
     ps_name = registration_name or default_name(ps_fn)
-    if ps_name in _PRUNING_STRATEGY and not tf.contrib.eager.in_eager_mode():
+    if ps_name in _PRUNING_STRATEGY and not tf.executing_eagerly():
       raise LookupError("Pruning strategy %s already registered." % ps_name)
     _PRUNING_STRATEGY[ps_name] = ps_fn
     return ps_fn
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index f3530b574..d95bd98d7 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -939,7 +939,7 @@ def _slow_greedy_infer_tpu(self, features, decode_length):
 
     def infer_step(i, recent_output, recent_logits, unused_loss):
       """Inference step."""
-      if not tf.contrib.eager.in_eager_mode():
+      if not tf.executing_eagerly():
         recent_output.set_shape([None, None, None, 1])
       padded = tf.pad(recent_output, [[0, 0], [0, 1], [0, 0], [0, 0]])
       features["targets"] = padded
@@ -957,7 +957,7 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
       samples = inplace_ops.alias_inplace_update(samples, i,
                                                  tf.to_int64(cur_sample))
       samples = tf.transpose(samples, perm=[1, 0, 2, 3])
-      if not tf.contrib.eager.in_eager_mode():
+      if not tf.executing_eagerly():
         samples.set_shape([None, None, None, 1])
 
       # Assuming we have one shard for logits.
@@ -1000,7 +1000,7 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
     # tensor padded to [batch_size, decode_length, 1, 1, vocab_size]
     logits = tf.zeros((batch_size, decode_length, 1, 1,
                        target_modality.top_dimensionality))
-    if not tf.contrib.eager.in_eager_mode():
+    if not tf.executing_eagerly():
       logits.set_shape([None, None, None, None, None])
     loss = 0.0
 
@@ -1106,7 +1106,7 @@ def _slow_greedy_infer(self, features, decode_length):
 
     def infer_step(recent_output, recent_logits, unused_loss):
       """Inference step."""
-      if not tf.contrib.eager.in_eager_mode():
+      if not tf.executing_eagerly():
         if self._target_modality_is_real:
           dim = self._problem_hparams.modality["targets"].top_dimensionality
           recent_output.set_shape([None, None, None, dim])
@@ -1130,7 +1130,7 @@ def infer_step(recent_output, recent_logits, unused_loss):
       else:
         cur_sample = tf.to_int64(tf.expand_dims(cur_sample, axis=1))
         samples = tf.concat([recent_output, cur_sample], axis=1)
-        if not tf.contrib.eager.in_eager_mode():
+        if not tf.executing_eagerly():
           samples.set_shape([None, None, None, 1])
 
       # Assuming we have one shard for logits.
@@ -1176,7 +1176,7 @@ def infer_step(recent_output, recent_logits, unused_loss):
       logits = tf.zeros((batch_size, 0, 1, 1,
                          target_modality.top_dimensionality))
       logits_shape_inv = [None, None, None, None, None]
-    if not tf.contrib.eager.in_eager_mode():
+    if not tf.executing_eagerly():
       logits.set_shape(logits_shape_inv)
 
     loss = 0.0
@@ -1459,7 +1459,7 @@ def initialize_from_ckpt(self, ckpt_dir):
     log_info("Checkpoint dir: %s", ckpt_dir)
 
     # TODO(mitchellstern): Add support for partitioned variables?
-    reader = tf.contrib.framework.load_checkpoint(ckpt_dir)
+    reader = tf.train.load_checkpoint(ckpt_dir)
     variable_map = {}
     for var in tf.contrib.framework.get_trainable_variables():
       var_name = var.name.split(":")[0]
@@ -1893,7 +1893,7 @@ def as_default(self):
 
 
 def create_eager_var_store():
-  if tf.contrib.eager.in_eager_mode():
+  if tf.executing_eagerly():
     return variable_scope.EagerVariableStore()
   else:
     return DummyVariableStore()
@@ -1998,7 +1998,7 @@ def summarize_features(features, num_shards=1):
 
 
 def _eager_log(level, *args):
-  if tf.contrib.eager.in_eager_mode() and args in _already_logged:
+  if tf.executing_eagerly() and args in _already_logged:
     return
   _already_logged.add(args)
   getattr(tf.logging, level)(*args)
diff --git a/tensor2tensor/utils/video/prediction2gif.py b/tensor2tensor/utils/video/prediction2gif.py
index 326d8ed06..fc55b7e4c 100644
--- a/tensor2tensor/utils/video/prediction2gif.py
+++ b/tensor2tensor/utils/video/prediction2gif.py
@@ -86,7 +86,7 @@ def main(_):
       data_dir=os.path.expanduser(FLAGS.data_dir),
       hparams=hparams)
 
-  dataset = dataset.apply(tf.contrib.data.batch_and_drop_remainder(num_agents))
+  dataset = dataset.batch(num_agents, drop_remainder=True)
   data = dataset.make_one_shot_iterator().get_next()
   # Setup input placeholders
   input_size = [num_agents, hparams.video_num_input_frames]
diff --git a/tensor2tensor/utils/video/reward_confusion.py b/tensor2tensor/utils/video/reward_confusion.py
index 1cb2c59c4..0e6c7e773 100644
--- a/tensor2tensor/utils/video/reward_confusion.py
+++ b/tensor2tensor/utils/video/reward_confusion.py
@@ -68,7 +68,7 @@ def main(_):
       shuffle_files=False,
       hparams=hparams)
 
-  dataset = dataset.apply(tf.contrib.data.batch_and_drop_remainder(batch_size))
+  dataset = dataset.batch(batch_size, drop_remainder=True)
   data = dataset.make_one_shot_iterator().get_next()
   input_data = dict((k, data[k]) for k in data.keys() if k.startswith("input"))
 
diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index caa84b774..04df03299 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -55,7 +55,7 @@ def load_videos(template, video_length, frame_shape):
   dataset_len = len(filenames)
   filenames = tf.constant(filenames)
   dataset = tf.data.Dataset.from_tensor_slices(filenames)
-  dataset = dataset.apply(tf.contrib.data.map_and_batch(
+  dataset = dataset.apply(tf.data.experimental.map_and_batch(
       lambda filename: load_image_map_function(filename, frame_shape),
       video_length, drop_remainder=True))
   return dataset, dataset_len

From 7b6ef710cf5cc79cefc2b012e2b396d7c3cd782e Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 7 Jan 2019 15:26:47 -0800
Subject: [PATCH 1448/2720] Dopamine fix, internally the code is at Dopamine 2,
 externally at 1.*

PiperOrigin-RevId: 228247340
---
 tensor2tensor/rl/dopamine_connector.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 7d765172a..89be6dc6b 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -22,7 +22,6 @@
 import copy
 
 from dopamine.agents.dqn import dqn_agent
-from dopamine.discrete_domains import run_experiment
 from dopamine.replay_memory import circular_replay_buffer
 from dopamine.replay_memory.circular_replay_buffer import OutOfGraphReplayBuffer
 from dopamine.replay_memory.circular_replay_buffer import ReplayElement
@@ -40,6 +39,10 @@
   import cv2
 except ImportError:
   cv2 = None
+try:
+  from dopamine.atari import run_experiment
+except ImportError:
+  run_experiment = None
 # pylint: enable=g-import-not-at-top
 
 

From aaabb9672ae0eb2d20d6f85c1ea6a7cdb21cb40d Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Mon, 7 Jan 2019 23:36:19 +0000
Subject: [PATCH 1449/2720] Upgrade Cloud ML version to tf@1.12 (#1344)

tensor2tensor requires [`tensorflow>=1.12.0`](https://github.com/tensorflow/tensor2tensor/blob/master/setup.py#L59). This PR upgrades the version on Cloud ML to 1.12 in order to fix run time errors.

Note: I only tested this change on Cloud ML *without* a TPU.
---
 tensor2tensor/utils/cloud_mlengine.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index c62069d24..4e8b2cfc0 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -34,14 +34,14 @@
 FLAGS = tf.flags.FLAGS
 
 CONSOLE_URL = "https://console.cloud.google.com/mlengine/jobs/"
-RUNTIME_VERSION = "1.9"
+RUNTIME_VERSION = "1.12"
 
 
 class Gcloud(object):
   """gcloud command strings."""
   # Note these can be modified by set_versions
-  VM_VERSION = "tf-1-9"
-  TPU_VERSION = "1.9"
+  VM_VERSION = "tf-1-12"
+  TPU_VERSION = "1.12"
 
   @classmethod
   def set_versions(cls, vm, tpu):

From fb0183551da63d52ba03d9db0150eff8dc808db8 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Mon, 7 Jan 2019 17:00:59 -0800
Subject: [PATCH 1450/2720] Fix bug with relative local 1D attention type key.

PiperOrigin-RevId: 228261235
---
 tensor2tensor/layers/common_image_attention.py | 10 ++++++----
 tensor2tensor/layers/modalities.py             |  2 --
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 27191a9dc..ef5759b6a 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -323,10 +323,12 @@ def transformer_decoder_layers(inputs,
                                attention_type="local_mask_right",
                                q_padding="LEFT", kv_padding="LEFT")
       elif attention_type == AttentionType.RELATIVE_LOCAL_1D:
-        y = local_attention_1d(common_layers.layer_preprocess(x, hparams),
-                               hparams,
-                               attention_type="rel_local_mask_right",
-                               q_padding="LEFT", kv_padding="LEFT")
+        y = local_attention_1d(
+            common_layers.layer_preprocess(x, hparams),
+            hparams,
+            attention_type="local_relative_mask_right",
+            q_padding="LEFT",
+            kv_padding="LEFT")
       elif attention_type == AttentionType.NON_CAUSAL_1D:
         y = local_attention_1d(common_layers.layer_preprocess(x, hparams),
                                hparams,
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 920124589..0915d0aaa 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -825,12 +825,10 @@ class VideoModalityIdentity(VideoModality):
 
   def bottom(self, x):
     common_video.gif_summary("inputs", x, max_outputs=1)
-    x = common_layers.standardize_images(x)
     return x
 
   def targets_bottom(self, x):
     common_video.gif_summary("targets", x, max_outputs=1)
-    x = common_layers.standardize_images(x)
     return x
 
   def top(self, body_output, targets):

From acde95f6cea575c1e5009d7a16d95545a23e0552 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 7 Jan 2019 17:07:32 -0800
Subject: [PATCH 1451/2720] Avoid deepcopy of hparams for json write

PiperOrigin-RevId: 228262066
---
 tensor2tensor/bin/t2t_trainer.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index c1a66ba0b..80efb08a9 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -19,7 +19,6 @@
 from __future__ import print_function
 
 import contextlib
-import copy
 import os
 import sys
 from tensor2tensor import models  # pylint: disable=unused-import
@@ -332,10 +331,11 @@ def save_metadata(hparams):
       f.write(t2t_flags_str)
 
   # Save hparams as hparams.json
-  new_hparams = copy.deepcopy(hparams)
-
+  hp_vals = hparams.values()
   # Modality class is not JSON serializable so remove.
-  new_hparams.del_hparam("modality")
+  del hp_vals["modality"]
+  new_hparams = tf.contrib.training.HParams(**hp_vals)
+
   hparams_fname = os.path.join(output_dir, "hparams.json")
   with tf.gfile.Open(hparams_fname, "w") as f:
     f.write(new_hparams.to_json(indent=0, sort_keys=True))

From 186ffc3f27994172c1cd5968e2ee32f954e369f6 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Tue, 8 Jan 2019 02:57:48 +0100
Subject: [PATCH 1452/2720] Model-Based RL: Player (#1330)

* SimulatedEnv with gym-like interface.

* Initial Player

* Player: Add reward header, keybord reset, CLI, move to player.py

* Player: add WAIT mode, few CLI options.

* Introduce Policy Inferencer

* Recording videos for ppo and player, some refactoring.

* Player refactor. Add real env recording with PPO agent.

* Extend CLI documentation. Remove some imports.

* Pylint

* Extend documentation.

* Correct dopamine import.

* Move gym.utils.play to global imports.

* Remove SimulatedEnv (unnecesarry wrapper for FlatBatchEnv<SimulatedBatchGymEnv>)

* Replace join_and_check with os.path.join.

* Move generation of initial_frame_chooser function to rl_utils.

* Move make_simulated_env_fn from trainer_model_based.py to rl.py

* Remove trainer_model_based imports, clean up player and record_ppo FLAGS.

* Move setup_env and load_t2t_gym_env to T2TGymEnv.

* Correct relative imports.

* Custom policy world_model and data paths for player.

* Enable BatchGymEnv to load directly from checkpoint file.

* Small fix record_ppo.

* Remove unused record_ppo.py.
---
 tensor2tensor/data_generators/gym_env.py      |  55 ++-
 tensor2tensor/models/research/rl.py           |  22 +
 tensor2tensor/rl/envs/simulated_batch_env.py  |   9 +-
 .../rl/envs/simulated_batch_gym_env.py        |   3 +-
 tensor2tensor/rl/player.py                    | 454 ++++++++----------
 tensor2tensor/rl/player_utils.py              | 260 ++++++++++
 tensor2tensor/rl/rl_utils.py                  |  69 ++-
 tensor2tensor/rl/trainer_model_based.py       |  72 +--
 tensor2tensor/rl/trainer_model_free.py        |   6 +-
 9 files changed, 613 insertions(+), 337 deletions(-)
 create mode 100644 tensor2tensor/rl/player_utils.py

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 90f02cdea..7a3a0a943 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -21,7 +21,10 @@
 
 import collections
 import itertools
+import os
 import random
+import re
+
 from gym.spaces import Box
 import numpy as np
 
@@ -185,7 +188,10 @@ def current_epoch_rollouts(self, split=None, minimal_rollout_frames=0):
     if not rollouts_by_split:
       if split is not None:
         raise ValueError(
-            "generate_data() should first be called in the current epoch"
+            "Data is not splitted into train/dev/test. If data created by "
+            "environment interaction (NOT loaded from disk) you should call "
+            "generate_data() first. Note that generate_data() will write to "
+            "disk and can corrupt your experiment data."
         )
       else:
         rollouts = self._current_epoch_rollouts
@@ -636,6 +642,53 @@ def base_env_name(self):
   def num_channels(self):
     return self.observation_space.shape[2]
 
+  @staticmethod
+  def infer_last_epoch_num(data_dir):
+    """Infer highest epoch number from file names in data_dir."""
+    names = os.listdir(data_dir)
+    epochs_str = [re.findall(pattern=r".*\.(-?\d+)$", string=name)
+                  for name in names]
+    epochs_str = sum(epochs_str, [])
+    return max([int(epoch_str) for epoch_str in epochs_str])
+
+  @staticmethod
+  def setup_env_from_hparams(hparams, batch_size, max_num_noops):
+    game_mode = "NoFrameskip-v4"
+    camel_game_name = misc_utils.snakecase_to_camelcase(hparams.game)
+    camel_game_name += game_mode
+    env_name = camel_game_name
+
+    env = T2TGymEnv(base_env_name=env_name,
+                    batch_size=batch_size,
+                    grayscale=hparams.grayscale,
+                    resize_width_factor=hparams.resize_width_factor,
+                    resize_height_factor=hparams.resize_height_factor,
+                    rl_env_max_episode_steps=hparams.rl_env_max_episode_steps,
+                    max_num_noops=max_num_noops, maxskip_envs=True)
+    return env
+
+  @staticmethod
+  def setup_and_load_epoch(hparams, data_dir, which_epoch_data=None):
+    """Load T2TBatchGymEnv with data from one epoch.
+
+    Args:
+        which_epoch_data: data from which epoch to load.
+    """
+    t2t_env = T2TGymEnv.setup_env_from_hparams(
+        hparams, batch_size=hparams.real_batch_size,
+        max_num_noops=hparams.max_num_noops
+    )
+    # Load data.
+    if which_epoch_data is not None:
+      if which_epoch_data == "last":
+        which_epoch_data = T2TGymEnv.infer_last_epoch_num(data_dir)
+      assert isinstance(which_epoch_data, int), \
+        "{}".format(type(which_epoch_data))
+      t2t_env.start_new_epoch(which_epoch_data, data_dir)
+    else:
+      t2t_env.start_new_epoch(-999)
+    return t2t_env
+
   def _derive_observation_space(self, orig_observ_space):
     height, width, channels = orig_observ_space.shape
     if self.grayscale:
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index c959aa712..8d5713684 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -203,6 +203,28 @@ def env_fn(in_graph):
   return env_fn
 
 
+def make_simulated_env_fn_from_hparams(
+    real_env, hparams, batch_size, initial_frame_chooser, model_dir,
+    sim_video_dir=None):
+  """Creates a simulated env_fn."""
+  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
+  if hparams.wm_policy_param_sharing:
+    model_hparams.optimizer_zero_grads = True
+  return make_simulated_env_fn(
+      reward_range=real_env.reward_range,
+      observation_space=real_env.observation_space,
+      action_space=real_env.action_space,
+      frame_stack_size=hparams.frame_stack_size,
+      frame_height=real_env.frame_height, frame_width=real_env.frame_width,
+      initial_frame_chooser=initial_frame_chooser, batch_size=batch_size,
+      model_name=hparams.generative_model,
+      model_hparams=trainer_lib.create_hparams(hparams.generative_model_params),
+      model_dir=model_dir,
+      intrinsic_reward_scale=hparams.intrinsic_reward_scale,
+      sim_video_dir=sim_video_dir,
+  )
+
+
 def get_policy(observations, hparams, action_space):
   """Get a policy network.
 
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index dc6848a02..8062f5bec 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -161,9 +161,12 @@ def initialize(self, sess):
     model_loader = tf.train.Saver(
         var_list=tf.global_variables(scope="next_frame*")  # pylint:disable=unexpected-keyword-arg
     )
-    trainer_lib.restore_checkpoint(
-        self._model_dir, saver=model_loader, sess=sess, must_restore=True
-    )
+    if os.path.isdir(self._model_dir):
+      trainer_lib.restore_checkpoint(
+          self._model_dir, saver=model_loader, sess=sess, must_restore=True
+      )
+    else:
+      model_loader.restore(sess=sess, save_path=self._model_dir)
 
   def __str__(self):
     return "SimulatedEnv"
diff --git a/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index bb815d537..dae336748 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -26,8 +26,7 @@
 
 
 class FlatBatchEnv(Env):
-  """TODO(konradczechowski): Add doc-string."""
-
+  """Gym environment interface for Batched Environments (with batch size = 1)"""
   def __init__(self, batch_env):
     if batch_env.batch_size != 1:
       raise ValueError("Number of environments in batch must be equal to one")
diff --git a/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 1c0bd4617..ea92e9f14 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -13,284 +13,244 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Play with a world model."""
+"""Play with a world model.
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import copy
-import os
-
-from gym.core import Env
-from gym.spaces import Box
-from gym.spaces import Discrete
-from gym.utils import play
-
-import numpy as np
-
-from PIL import Image
-from PIL import ImageDraw
-from PIL import ImageFont
+Controls:
+  WSAD and SPACE to control the agent.
+  R key to reset env.
+  C key to toggle WAIT mode.
+  N to perform NOOP action under WAIT mode.
 
-from tensor2tensor.data_generators import gym_env
-from tensor2tensor.models.research.rl import get_policy
-from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
-from tensor2tensor.rl.trainer_model_based import FLAGS
-from tensor2tensor.rl.trainer_model_based import setup_directories
-from tensor2tensor.rl.trainer_model_based import temporary_flags
-
-from tensor2tensor.utils import registry
-from tensor2tensor.utils import trainer_lib
-import tensorflow as tf
-
-
-_font = None
-FONT_SIZE = 20
+Run this script with the same parameters as trainer_model_based.py. Note that
+values of most of them have no effect on player, so running just
 
+python -m tensor2tensor/rl/player.py \
+    --output_dir=path/to/your/experiment \
+    --loop_hparams_set=rlmb_base
 
-def _get_font():
-  global _font
-  if _font is None:
-    font_paths = []
-    for path in font_paths:
-      try:
-        _font = ImageFont.truetype(path, FONT_SIZE)
-        return _font
-      except:  # pylint: disable=bare-except
-        pass
+might work for you.
 
+More advanced example:
 
-def _assert_image(img):
-  if isinstance(img, np.ndarray):
-    img = Image.fromarray(np.ndarray.astype(img, np.uint8))
-  return img
+python -m tensor2tensor/rl/record_ppo.py \
+    --output_dir=path/to/your/experiment \
+    --loop_hparams_set=rlmb_base \
+    --loop_hparams=game=<right game in case of problems> \
+    --video_dir=my/video/dir \
+    --zoom=6 \
+    --fps=50 \
+    --env=real \
+    --epoch=-1
 
+Check flags definitions under imports for more details.
+"""
 
-def write_on_image(img, text="", position=(0, 0), color=(255, 255, 255)):
-  img = _assert_image(img)
-  if not text:
-    return img
-  draw = ImageDraw.Draw(img)
-  font = _get_font()
-  draw.text(position, text, color, font=font)
-  return img
-
-
-def concatenate_images(imgs, axis=1):
-  imgs = [_assert_image(img) for img in imgs]
-  imgs_np = [np.array(img) for img in imgs]
-  concatenated_im_np = np.concatenate(imgs_np, axis=axis)
-  return _assert_image(concatenated_im_np)
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
+import six
 
-class DebugBatchEnv(Env):
-  """Debugging Environment."""
-  INFO_PANE_WIDTH = 250
+import gym
+from gym.envs.atari.atari_env import ACTION_MEANING
+from gym.utils import play
+import numpy as np
 
-  def __init__(self, hparams, sess=None):
-    self.action_space = Discrete(6)
-    self.observation_space = Box(
-        low=0, high=255, shape=(210, 160+DebugBatchEnv.INFO_PANE_WIDTH, 3),
-        dtype=np.uint8)
-    self._tmp = 1
-    self.res = None
-    self.sess = sess if sess is not None else tf.Session()
-    self._prepare_networks(hparams, self.sess)
+from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image, PIL_ImageDraw
+from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
+from tensor2tensor.rl.player_utils import wrap_with_monitor, \
+  load_data_and_make_simulated_env, infer_paths
+# Import flags from t2t_trainer and trainer_model_based
+from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
+import tensor2tensor.rl.trainer_model_based_params # pylint: disable=unused-import
 
-  def _prepare_networks(self, hparams, sess):
-    self.action = tf.placeholder(shape=(1,), dtype=tf.int32)
-    batch_env = SimulatedBatchEnv(hparams.environment_spec, hparams.num_agents)
-    self.reward, self.done = batch_env.simulate(self.action)
-    self.observation = batch_env.observ
-    self.reset_op = batch_env.reset(tf.constant([0], dtype=tf.int32))
+from tensor2tensor.data_generators.gym_env import T2TGymEnv
+from tensor2tensor.utils import registry
+import tensorflow as tf
 
-    environment_wrappers = hparams.environment_spec.wrappers
-    wrappers = copy.copy(environment_wrappers) if environment_wrappers else []
 
-    to_initialize = [batch_env]
-    for w in wrappers:
-      batch_env = w[0](batch_env, **w[1])
-      to_initialize.append(batch_env)
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("video_dir", "/tmp/gym-results",
+                    "Where to save played trajectories.")
+flags.DEFINE_float("zoom", 4.,
+                   "Resize factor of displayed game.")
+flags.DEFINE_float("fps", 20.,
+                   "Frames per second.")
+flags.DEFINE_string("epoch", "last",
+                    "Data from which epoch to use.")
+flags.DEFINE_boolean("simulated_env", True,
+                     "Either to use 'simulated' or 'real' env.")
+flags.DEFINE_boolean("dry_run", False,
+                     "Dry run - without pygame interaction and display, just "
+                     "some random actions on environment")
+flags.DEFINE_string("model_ckpt", "",
+                    "World model checkpoint path.")
+flags.DEFINE_string("wm_dir", "",
+                    "Directory with world model checkpoints. Inferred from "
+                    "output_dir if empty.")
+flags.DEFINE_string("policy_dir", "",
+                    "Directory with policy. Inferred from output_dir if empty.")
+flags.DEFINE_string("episodes_data_dir", "",
+                    "Path to data for simulated environment initialization. "
+                    "Inferred from output_dir if empty.")
+
+
+class PlayerEnvWrapper(gym.Wrapper):
+  """ Environment Wrapper for gym.utils.play.
+
+  This probably will be highly refactored.
+  """
+
+  RESET_ACTION = 101
+  TOGGLE_WAIT_ACTION = 102
+  WAIT_MODE_NOOP_ACTION = 103
+
+  HEADER_HEIGHT = 12
+
+  def __init__(self, env):
+    super(PlayerEnvWrapper, self).__init__(env)
+
+    # Set observation space
+    orig = self.env.observation_space
+    shape = tuple([orig.shape[0] + self.HEADER_HEIGHT] + list(orig.shape[1:]))
+    self.observation_space = gym.spaces.Box(low=orig.low.min(),
+                                            high=orig.high.max(),
+                                            shape=shape, dtype=orig.dtype)
+
+    # gym play() looks for get_keys_to_action() only on top and bottom level
+    # of env and wrappers stack.
+    self.unwrapped.get_keys_to_action = self.get_keys_to_action
+
+    self._wait = True
+    self.action_meaning = {i: ACTION_MEANING[i]
+                           for i in range(self.action_space.n)}
+    self.name_to_action_num = {v: k for k, v in
+                               six.iteritems(self.action_meaning)}
+
+  def get_action_meanings(self):
+    return [self.action_meaning[i] for i in range(self.action_space.n)]
+
+  def get_keys_to_action(self):
+    # Based on gym atari.py AtariEnv.get_keys_to_action()
+    keyword_to_key = {
+        "UP": ord("w"),
+        "DOWN": ord("s"),
+        "LEFT": ord("a"),
+        "RIGHT": ord("d"),
+        "FIRE": ord(" "),
+    }
+
+    keys_to_action = {}
+
+    for action_id, action_meaning in enumerate(self.get_action_meanings()):
+      keys = []
+      for keyword, key in keyword_to_key.items():
+        if keyword in action_meaning:
+          keys.append(key)
+      keys_tuple = tuple(sorted(keys))
+      del keys
+      assert keys_tuple not in keys_to_action
+      keys_to_action[keys_tuple] = action_id
+
+    # Add utility actions
+    keys_to_action[(ord("r"),)] = self.RESET_ACTION
+    keys_to_action[(ord("c"),)] = self.TOGGLE_WAIT_ACTION
+    keys_to_action[(ord("n"),)] = self.WAIT_MODE_NOOP_ACTION
+
+    return keys_to_action
 
-    def initialization_lambda():
-      for batch_env in to_initialize:
-        batch_env.initialize(sess)
+  def step(self, action):
+    # Special codes
+    if action == self.TOGGLE_WAIT_ACTION:
+      self._wait = not self._wait
+      ob, reward, done, info = self._last_step
+      ob = self.augment_observation(ob, reward, self.total_reward)
+      return ob, reward, done, info
 
-    self.initialize = initialization_lambda
+    if action == self.RESET_ACTION:
+      ob = self.empty_observation()
+      return ob, 0, True, {}
 
-    obs_copy = batch_env.observ + 0
+    if self._wait and action == self.name_to_action_num["NOOP"]:
+      ob, reward, done, info = self._last_step
+      ob = self.augment_observation(ob, reward, self.total_reward)
+      return ob, reward, done, info
 
-    actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
-    self.policy_probs = actor_critic.policy.probs[0, 0, :]
-    self.value = actor_critic.value[0, :]
+    if action == self.WAIT_MODE_NOOP_ACTION:
+      action = self.name_to_action_num["NOOP"]
 
-  def render(self, mode="human"):
-    raise NotImplementedError()
+    ob, reward, done, info = self.env.step(action)
+    self._last_step = ob, reward, done, info
 
-  def _fake_reset(self):
-    self._tmp = 0
-    observ = np.ones(shape=(210, 160, 3), dtype=np.uint8) * 10 * self._tmp
-    observ[0, 0, 0] = 0
-    observ[0, 0, 1] = 255
-    self.res = (observ, 0, False, [0.1, 0.5, 0.5], 1.1)
+    self.total_reward += reward
 
-  def _reset_env(self):
-    observ = self.sess.run(self.reset_op)[0, ...]
-    observ[0, 0, 0] = 0
-    observ[0, 0, 1] = 255
-    # TODO(pmilos): put correct numbers
-    self.res = (observ, 0, False, [0.1, 0.5, 0.5], 1.1)
+    ob = self.augment_observation(ob, reward, self.total_reward)
+    return ob, reward, done, info
 
   def reset(self):
-    self._reset_env()
-    observ = self._augment_observation()
-    return observ
-
-  def _step_fake(self, action):
-    observ = np.ones(shape=(210, 160, 3), dtype=np.uint8)*10*self._tmp
-    observ[0, 0, 0] = 0
-    observ[0, 0, 1] = 255
-
-    self._tmp += 1
-    if self._tmp > 20:
-      self._tmp = 0
-
-    rew = 1
-    done = False
-    probs = np.ones(shape=(6,), dtype=np.float32)/6
-    vf = 0.0
-
-    return observ, rew, done, probs, vf
-
-  def _step_env(self, action):
-    observ, rew, done, probs, vf = self.sess.\
-      run([self.observation, self.reward, self.done, self.policy_probs,
-           self.value],
-          feed_dict={self.action: [action]})
-
-    return observ[0, ...], rew[0, ...], done[0, ...], probs, vf
-
-  def _augment_observation(self):
-    observ, rew, _, probs, vf = self.res
-    info_pane = np.zeros(shape=(210, DebugBatchEnv.INFO_PANE_WIDTH, 3),
-                         dtype=np.uint8)
-    probs_str = ""
-    for p in probs:
-      probs_str += "%.2f" % p + ", "
-
-    probs_str = probs_str[:-2]
-
-    action = np.argmax(probs)
-    info_str = " Policy:{}\n Action:{}\n Value function:{}\n Reward:{}".format(
-        probs_str, action, vf, rew)
-    print("Info str:{}".format(info_str))
-    # info_pane = write_on_image(info_pane, info_str)
-
-    augmented_observ = concatenate_images([observ, info_pane])
-    augmented_observ = np.array(augmented_observ)
-    return augmented_observ
-
-  def step(self, action):
-    # Special codes
-    if action == 100:
-      # skip action
-      _, rew, done, _, _ = self.res
-      observ = self._augment_observation()
-      return observ, rew, done, {}
-
-    if action == 101:
-      # reset
-      self.reset()
-      _, rew, done, _, _ = self.res
-      observ = self._augment_observation()
-      return observ, rew, done, {}
-
-    if action == 102:
-      # play
-      raise NotImplementedError()
-
-    # standard codes
-    observ, rew, done, probs, vf = self._step_env(action)
-    self.res = (observ, rew, done, probs, vf)
-
-    observ = self._augment_observation()
-    return observ, rew, done, {"probs": probs, "vf": vf}
+    ob = self.env.reset()
+    self._last_step = ob, 0, False, {}
+    self.total_reward = 0
+    return self.augment_observation(ob, 0, self.total_reward)
+
+  def empty_observation(self):
+    return np.zeros(self.observation_space.shape)
+
+  def augment_observation(self, ob, reward, total_reward):
+    img = PIL_Image().new("RGB",
+                          (ob.shape[1], PlayerEnvWrapper.HEADER_HEIGHT,))
+    draw = PIL_ImageDraw().Draw(img)
+    draw.text((1, 0), "c:{:3}, r:{:3}".format(int(total_reward), int(reward)),
+              fill=(255, 0, 0))
+    header = np.asarray(img)
+    del img
+    header.setflags(write=1)
+    if self._wait:
+      pixel_fill = (0, 255, 0)
+    else:
+      pixel_fill = (255, 0, 0)
+    header[0, :, :] = pixel_fill
+    return np.concatenate([header, ob], axis=0)
 
 
 def main(_):
+  # gym.logger.set_level(gym.logger.DEBUG)
   hparams = registry.hparams(FLAGS.loop_hparams_set)
   hparams.parse(FLAGS.loop_hparams)
-  output_dir = FLAGS.output_dir
+  # Not important for experiments past 2018
+  if "wm_policy_param_sharing" not in hparams.values().keys():
+    hparams.add_hparam("wm_policy_param_sharing", False)
+  directories = infer_paths(output_dir=FLAGS.output_dir,
+                            world_model=FLAGS.wm_dir,
+                            policy=FLAGS.policy_dir,
+                            data=FLAGS.episodes_data_dir)
+  epoch = FLAGS.epoch if FLAGS.epoch == "last" else int(FLAGS.epoch)
+
+  if FLAGS.simulated_env:
+    env = load_data_and_make_simulated_env(directories["data"],
+                                           directories["world_model"],
+                                           hparams, which_epoch_data=epoch)
+  else:
+    env = T2TGymEnv.setup_and_load_epoch(
+        hparams, data_dir=directories["data"],
+        which_epoch_data=epoch)
+    env = FlatBatchEnv(env)
 
-  subdirectories = ["data", "tmp", "world_model", "ppo"]
-  using_autoencoder = hparams.autoencoder_train_steps > 0
-  if using_autoencoder:
-    subdirectories.append("autoencoder")
-  directories = setup_directories(output_dir, subdirectories)
+  env = PlayerEnvWrapper(env)  # pylint: disable=redefined-variable-type
 
-  if hparams.game in gym_env.ATARI_GAMES:
-    game_with_mode = hparams.game + "_deterministic-v4"
-  else:
-    game_with_mode = hparams.game
+  env = wrap_with_monitor(env, FLAGS.video_dir)
 
-  if using_autoencoder:
-    simulated_problem_name = (
-        "gym_simulated_discrete_problem_with_agent_on_%s_autoencoded"
-        % game_with_mode)
-  else:
-    simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
-                              % game_with_mode)
-    if simulated_problem_name not in registry.list_problems():
-      tf.logging.info("Game Problem %s not found; dynamically registering",
-                      simulated_problem_name)
-      gym_env.register_game(hparams.game, game_mode="Deterministic-v4")
-
-  epoch = hparams.epochs-1
-  epoch_data_dir = os.path.join(directories["data"], str(epoch))
-  ppo_model_dir = directories["ppo"]
-
-  world_model_dir = directories["world_model"]
-
-  gym_problem = registry.problem(simulated_problem_name)
-
-  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
-  environment_spec = copy.copy(gym_problem.environment_spec)
-  environment_spec.simulation_random_starts = hparams.simulation_random_starts
-
-  batch_env_hparams = trainer_lib.create_hparams(hparams.ppo_params)
-  batch_env_hparams.add_hparam("model_hparams", model_hparams)
-  batch_env_hparams.add_hparam("environment_spec", environment_spec)
-  batch_env_hparams.num_agents = 1
-
-  with temporary_flags({
-      "problem": simulated_problem_name,
-      "model": hparams.generative_model,
-      "hparams_set": hparams.generative_model_params,
-      "output_dir": world_model_dir,
-      "data_dir": epoch_data_dir,
-  }):
-    sess = tf.Session()
-    env = DebugBatchEnv(batch_env_hparams, sess)
-    sess.run(tf.global_variables_initializer())
-    env.initialize()
-
-    env_model_loader = tf.train.Saver(
-        tf.global_variables("next_frame*"))
-    trainer_lib.restore_checkpoint(world_model_dir, env_model_loader, sess,
-                                   must_restore=True)
-
-    model_saver = tf.train.Saver(
-        tf.global_variables(".*network_parameters.*"))
-    trainer_lib.restore_checkpoint(ppo_model_dir, model_saver, sess)
-
-    key_mapping = gym_problem.env.env.get_keys_to_action()
-    # map special codes
-    key_mapping[()] = 100
-    key_mapping[(ord("r"),)] = 101
-    key_mapping[(ord("p"),)] = 102
-
-    play.play(env, zoom=2, fps=10, keys_to_action=key_mapping)
+  if FLAGS.dry_run:
+    for _ in range(5):
+      env.reset()
+      for i in range(50):
+        env.step(i % 3)
+      env.step(PlayerEnvWrapper.RESET_ACTION)  # reset
+    return
+
+  play.play(env, zoom=FLAGS.zoom, fps=FLAGS.fps)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
new file mode 100644
index 000000000..3cd118719
--- /dev/null
+++ b/tensor2tensor/rl/player_utils.py
@@ -0,0 +1,260 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Utilities for player.py."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+from copy import deepcopy
+
+import gym
+import six
+from gym import wrappers, spaces
+import numpy as np
+
+from tensor2tensor.rl import rl_utils
+from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
+from tensor2tensor.data_generators.gym_env import T2TGymEnv
+from tensor2tensor.models.research.rl import get_policy
+from tensor2tensor.models.research.rl import make_simulated_env_fn_from_hparams
+from tensor2tensor.utils import trainer_lib
+import tensorflow as tf
+
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+
+def make_simulated_gym_env(real_env, world_model_dir, hparams, random_starts):
+  """Gym environment with world model."""
+  initial_frame_chooser = rl_utils.make_initial_frame_chooser(
+      real_env, hparams.frame_stack_size,
+      simulation_random_starts=random_starts,
+      simulation_flip_first_random_for_beginning=False
+  )
+  env_fn = make_simulated_env_fn_from_hparams(
+      real_env, hparams,
+      batch_size=1,
+      initial_frame_chooser=initial_frame_chooser,
+      model_dir=world_model_dir
+  )
+  env = env_fn(in_graph=False)
+  flat_env = FlatBatchEnv(env)
+  return flat_env
+
+
+def load_data_and_make_simulated_env(
+    data_dir, wm_dir, hparams, which_epoch_data="last", random_starts=True
+):
+  hparams = deepcopy(hparams)
+  t2t_env = T2TGymEnv.setup_and_load_epoch(
+      hparams, data_dir=data_dir,
+      which_epoch_data=which_epoch_data)
+  return make_simulated_gym_env(
+      t2t_env, world_model_dir=wm_dir,
+      hparams=hparams, random_starts=random_starts)
+
+
+class ExtendToEvenDimentions(gym.ObservationWrapper):
+  """ Force even dimentions of both height and width.
+
+  Specifically, it adds single zero row/column to observations if needed.
+  """
+  HW_AXES = (0, 1)
+  def __init__(self, env):
+    gym.ObservationWrapper.__init__(self, env)
+
+    orig_shape = env.observation_space.shape
+    extended_shape = list(orig_shape)
+    for axis in self.HW_AXES:
+      if self.if_odd(orig_shape[axis]):
+        extended_shape[axis] += 1
+
+    assert env.observation_space.dtype == np.uint8
+    self.observation_space = spaces.Box(
+        low=0,
+        high=255,
+        shape=extended_shape,
+        dtype=np.uint8)
+
+  def observation(self, frame):
+    """Add single zero row/column to observation if needed."""
+    if frame.shape == self.observation_space.shape:
+      return frame
+    else:
+      extended_frame = np.zeros(self.observation_space.shape,
+                                self.observation_space.dtype)
+      assert self.HW_AXES == (0, 1)
+      extended_frame[:frame.shape[0], :frame.shape[1]] = frame
+      return extended_frame
+
+  def if_odd(self, n):
+    return n % 2
+
+
+class RenderObservations(gym.Wrapper):
+  """Add observations rendering in 'rgb_array' mode."""
+  def __init__(self, env):
+    super(RenderObservations, self).__init__(env)
+    if "rgb_array" not in self.metadata["render.modes"]:
+      self.metadata["render.modes"].append("rgb_array")
+
+  def step(self, action):
+    ret = self.env.step(action)
+    self.last_observation = ret[0]
+    return ret
+
+  def reset(self, **kwargs):
+    self.last_observation = self.env.reset(**kwargs)
+    return self.last_observation
+
+  def render(self, mode="human", **kwargs):
+    assert mode == "rgb_array"
+    return self.last_observation
+
+
+def wrap_with_monitor(env, video_dir):
+  """Wrap environment with gym.Monitor
+
+  Video recording provided by Monitor requires
+    1) both height and width of observation to be even numbers.
+    2) rendering of environment
+  """
+  env = ExtendToEvenDimentions(env)
+  env = RenderObservations(env)  # pylint: disable=redefined-variable-type
+  env = wrappers.Monitor(env, video_dir, force=True,
+                         video_callable=lambda idx: True,
+                         write_upon_reset=True)
+  return env
+
+
+def create_simulated_env(
+    output_dir, grayscale, resize_width_factor, resize_height_factor,
+    frame_stack_size, generative_model, generative_model_params,
+    random_starts=True, which_epoch_data="last", **other_hparams
+):
+  """"Create SimulatedEnv with minimal subset of hparams."""
+  # We need these, to initialize T2TGymEnv, but these values (hopefully) have
+  # no effect on player.
+  a_bit_risky_defaults = {
+      "game": "pong",  # assumes that T2TGymEnv has always reward_range (-1,1)
+      "real_batch_size": 1,
+      "rl_env_max_episode_steps": -1,
+      "max_num_noops": 0
+  }
+
+  for key in a_bit_risky_defaults:
+    if key not in other_hparams:
+      other_hparams[key] = a_bit_risky_defaults[key]
+
+  hparams = tf.contrib.training.HParams(
+      grayscale=grayscale,
+      resize_width_factor=resize_width_factor,
+      resize_height_factor=resize_height_factor,
+      frame_stack_size=frame_stack_size,
+      generative_model=generative_model,
+      generative_model_params=generative_model_params,
+      **other_hparams
+  )
+  return load_data_and_make_simulated_env(output_dir, hparams,
+                                          which_epoch_data=which_epoch_data,
+                                          random_starts=random_starts)
+
+
+class PPOPolicyInferencer(object):
+  """Non-tensorflow API for infering policy (and value function).
+
+  Example:
+    >>> ppo = PPOPolicyInferencer(...)
+    >>> ppo.reset_frame_stack()
+    >>> ob = env.reset()
+    >>> while not done:
+    >>>   logits, value = ppo.infer(ob)
+    >>>   ob, _, done, _ = env.step(action)
+  """
+  def __init__(self, hparams, action_space, observation_space, policy_dir):
+    assert hparams.base_algo == "ppo"
+    ppo_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
+
+    frame_stack_shape = (1, hparams.frame_stack_size) + observation_space.shape
+    self._frame_stack = np.zeros(frame_stack_shape, dtype=np.uint8)
+
+    with tf.Graph().as_default():
+      self.obs_t = tf.placeholder(shape=self.frame_stack_shape, dtype=np.uint8)
+      self.logits_t, self.value_function_t = get_policy(
+          self.obs_t, ppo_hparams, action_space
+      )
+      model_saver = tf.train.Saver(
+          tf.global_variables(scope=ppo_hparams.policy_network + "/.*")  # pylint: disable=unexpected-keyword-arg
+      )
+      self.sess = tf.Session()
+      self.sess.run(tf.global_variables_initializer())
+      trainer_lib.restore_checkpoint(policy_dir, model_saver,
+                                     self.sess)
+
+  @property
+  def frame_stack_shape(self):
+    return self._frame_stack.shape
+
+  def reset_frame_stack(self, frame_stack=None):
+    if frame_stack is None:
+      self._frame_stack.fill(0)
+    else:
+      assert frame_stack.shape == self.frame_stack_shape, \
+        "{}, {}".format(frame_stack.shape, self.frame_stack_shape)
+      self._frame_stack = frame_stack.copy()
+
+  def _add_to_stack(self, ob):
+    stack = np.roll(self._frame_stack, shift=-1, axis=1)
+    stack[0, -1, ...] = ob
+    self._frame_stack = stack
+
+  def infer(self, ob):
+    """Add new observation to frame stack and infer policy.
+
+    Args:
+      ob: array of shape (height, width, channels)
+    """
+    self._add_to_stack(ob)
+    logits, vf = self.infer_from_frame_stack(self._frame_stack)
+    return logits, vf
+
+  def infer_from_frame_stack(self, ob_stack):
+    """ Infer policy from stack of observations.
+
+    Args:
+      ob_stack: array of shape (1, frame_stack_size, height, width, channels)
+    """
+    logits, vf = self.sess.run([self.logits_t, self.value_function_t],
+                               feed_dict={self.obs_t: ob_stack})
+    return logits, vf
+
+
+def infer_paths(output_dir, **subdirs):
+  """
+
+  Example:
+    >>> infer_paths("/some/output/dir/", policy="", model="custom/path")
+    {"policy": "/some/output/dir/policy", "model": "custom/path",
+    "output_dir":"/some/output/dir/"}
+  """
+  directories = dict()
+  for name, path in six.iteritems(subdirs):
+    directories[name] = path if path else os.path.join(output_dir, name)
+  directories['output_dir'] = output_dir
+  return directories
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 79a5a4dcb..19293e151 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import random
+
 import numpy as np
 import six
 
@@ -26,7 +28,6 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl.dopamine_connector import DQNLearner
 from tensor2tensor.rl.ppo_learner import PPOLearner
-from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -62,7 +63,7 @@ def evaluate_single_config(
 ):
   """Evaluate the PPO agent in the real environment."""
   eval_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
-  env = setup_env(
+  env = T2TGymEnv.setup_env_from_hparams(
       hparams, batch_size=hparams.eval_batch_size, max_num_noops=max_num_noops
   )
   env.start_new_epoch(0)
@@ -103,25 +104,55 @@ def evaluate_all_configs(hparams, agent_model_dir):
 }
 
 
-def setup_env(hparams, batch_size, max_num_noops):
-  """Setup."""
-  game_mode = "NoFrameskip-v4"
-  camel_game_name = misc_utils.snakecase_to_camelcase(hparams.game)
-  camel_game_name += game_mode
-  env_name = camel_game_name
-
-  env = T2TGymEnv(base_env_name=env_name,
-                  batch_size=batch_size,
-                  grayscale=hparams.grayscale,
-                  resize_width_factor=hparams.resize_width_factor,
-                  resize_height_factor=hparams.resize_height_factor,
-                  rl_env_max_episode_steps=hparams.rl_env_max_episode_steps,
-                  max_num_noops=max_num_noops, maxskip_envs=True)
-  return env
-
-
 def update_hparams_from_hparams(target_hparams, source_hparams, prefix):
   """Copy a subset of hparams to target_hparams."""
   for (param_name, param_value) in six.iteritems(source_hparams.values()):
     if param_name.startswith(prefix):
       target_hparams.set_hparam(param_name[len(prefix):], param_value)
+
+
+def random_rollout_subsequences(rollouts, num_subsequences, subsequence_length):
+  """Chooses a random frame sequence of given length from a set of rollouts."""
+  def choose_subsequence():
+    # TODO(koz4k): Weigh rollouts by their lengths so sampling is uniform over
+    # frames and not rollouts.
+    rollout = random.choice(rollouts)
+    try:
+      from_index = random.randrange(len(rollout) - subsequence_length + 1)
+    except ValueError:
+      # Rollout too short; repeat.
+      return choose_subsequence()
+    return rollout[from_index:(from_index + subsequence_length)]
+
+  return [choose_subsequence() for _ in range(num_subsequences)]
+
+
+def make_initial_frame_chooser(real_env, frame_stack_size,
+                               simulation_random_starts,
+                               simulation_flip_first_random_for_beginning):
+  initial_frame_rollouts = real_env.current_epoch_rollouts(
+      split=tf.estimator.ModeKeys.TRAIN,
+      minimal_rollout_frames=frame_stack_size,
+  )
+  def initial_frame_chooser(batch_size):
+    """Frame chooser."""
+
+    deterministic_initial_frames =\
+        initial_frame_rollouts[0][:frame_stack_size]
+    if not simulation_random_starts:
+      # Deterministic starts: repeat first frames from the first rollout.
+      initial_frames = [deterministic_initial_frames] * batch_size
+    else:
+      # Random starts: choose random initial frames from random rollouts.
+      initial_frames = random_rollout_subsequences(
+          initial_frame_rollouts, batch_size, frame_stack_size
+      )
+      if simulation_flip_first_random_for_beginning:
+        # Flip first entry in the batch for deterministic initial frames.
+        initial_frames[0] = deterministic_initial_frames
+
+    return np.stack([
+        [frame.observation.decode() for frame in initial_frame_stack]
+        for initial_frame_stack in initial_frames
+    ])
+  return initial_frame_chooser
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 78abce708..8aeaba516 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -36,9 +36,12 @@
 import numpy as np
 import six
 
+from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw, PIL_Image
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
+from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl
+from tensor2tensor.models.research.rl import make_simulated_env_fn_from_hparams
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl import trainer_model_based_params
 from tensor2tensor.rl.restarter import Restarter
@@ -51,18 +54,6 @@
 FLAGS = flags.FLAGS
 
 
-# Lazy load PIL.Image
-def PIL_Image():  # pylint: disable=invalid-name
-  from PIL import Image  # pylint: disable=g-import-not-at-top
-  return Image
-
-
-# Lazy load PIL.Image
-def PIL_ImageDraw():  # pylint: disable=invalid-name
-  from PIL import ImageDraw  # pylint: disable=g-import-not-at-top
-  return ImageDraw
-
-
 def real_env_step_increment(hparams):
   """Real env step increment."""
   return int(math.ceil(
@@ -135,28 +126,6 @@ def choose_subsequence():
   return [choose_subsequence() for _ in range(num_subsequences)]
 
 
-def make_simulated_env_fn(
-    real_env, hparams, batch_size, initial_frame_chooser, model_dir,
-    sim_video_dir=None):
-  """Creates a simulated env_fn."""
-  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
-  if hparams.wm_policy_param_sharing:
-    model_hparams.optimizer_zero_grads = True
-  return rl.make_simulated_env_fn(
-      reward_range=real_env.reward_range,
-      observation_space=real_env.observation_space,
-      action_space=real_env.action_space,
-      frame_stack_size=hparams.frame_stack_size,
-      frame_height=real_env.frame_height, frame_width=real_env.frame_width,
-      initial_frame_chooser=initial_frame_chooser, batch_size=batch_size,
-      model_name=hparams.generative_model,
-      model_hparams=trainer_lib.create_hparams(hparams.generative_model_params),
-      model_dir=model_dir,
-      intrinsic_reward_scale=hparams.intrinsic_reward_scale,
-      sim_video_dir=sim_video_dir,
-  )
-
-
 def train_supervised(problem, model_name, hparams, data_dir, output_dir,
                      train_steps, eval_steps, local_eval_frequency=None,
                      schedule="continuous_train_and_eval"):
@@ -175,34 +144,11 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
 
 def train_agent(real_env, learner, world_model_dir, hparams, epoch):
   """Train the PPO agent in the simulated environment."""
-  frame_stack_size = hparams.frame_stack_size
-  initial_frame_rollouts = real_env.current_epoch_rollouts(
-      split=tf.estimator.ModeKeys.TRAIN,
-      minimal_rollout_frames=frame_stack_size,
+  initial_frame_chooser = rl_utils.make_initial_frame_chooser(
+      real_env, hparams.frame_stack_size, hparams.simulation_random_starts,
+      hparams.simulation_flip_first_random_for_beginning
   )
-  # TODO(koz4k): Move this to a different module.
-  def initial_frame_chooser(batch_size):
-    """Frame chooser."""
-
-    deterministic_initial_frames =\
-        initial_frame_rollouts[0][:frame_stack_size]
-    if not hparams.simulation_random_starts:
-      # Deterministic starts: repeat first frames from the first rollout.
-      initial_frames = [deterministic_initial_frames] * batch_size
-    else:
-      # Random starts: choose random initial frames from random rollouts.
-      initial_frames = random_rollout_subsequences(
-          initial_frame_rollouts, batch_size, frame_stack_size
-      )
-      if hparams.simulation_flip_first_random_for_beginning:
-        # Flip first entry in the batch for deterministic initial frames.
-        initial_frames[0] = deterministic_initial_frames
-
-    return np.stack([
-        [frame.observation.decode() for frame in initial_frame_stack]
-        for initial_frame_stack in initial_frames
-    ])
-  env_fn = make_simulated_env_fn(
+  env_fn = make_simulated_env_fn_from_hparams(
       real_env, hparams, hparams.simulated_batch_size, initial_frame_chooser,
       world_model_dir, os.path.join(learner.agent_model_dir,
                                     "sim_videos_{}".format(epoch))
@@ -295,7 +241,7 @@ def initial_frame_chooser(batch_size):
         for subsequence in rollout_subsequences
     ])
 
-  env_fn = make_simulated_env_fn(
+  env_fn = make_simulated_env_fn_from_hparams(
       real_env, hparams, hparams.wm_eval_batch_size, initial_frame_chooser,
       world_model_dir
   )
@@ -454,7 +400,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   epoch = -1
   data_dir = directories["data"]
-  env = rl_utils.setup_env(
+  env = T2TGymEnv.setup_env_from_hparams(
       hparams, batch_size=hparams.real_batch_size,
       max_num_noops=hparams.max_num_noops
   )
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 91201d89e..64302dfdc 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -29,6 +29,7 @@
 
 import pprint
 
+from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
@@ -52,8 +53,9 @@
 
 def initialize_env_specs(hparams):
   """Initializes env_specs using T2TGymEnvs."""
-  env = rl_utils.setup_env(hparams, hparams.batch_size,
-                           hparams.eval_max_num_noops)
+  env = T2TGymEnv.setup_env_from_hparams(
+      hparams, hparams.batch_size, hparams.eval_max_num_noops
+  )
   env.start_new_epoch(0)
 
   # TODO(afrozm): Decouple env_fn from hparams and return both, is there

From d36669f84ab8a83d416048d637a45fff93240894 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Mon, 7 Jan 2019 17:58:19 -0800
Subject: [PATCH 1453/2720] internal merge of PR #1330

PiperOrigin-RevId: 228267563
---
 tensor2tensor/data_generators/gym_env.py      |  7 ++-
 tensor2tensor/rl/envs/simulated_batch_env.py  |  1 +
 .../rl/envs/simulated_batch_gym_env.py        |  3 +-
 tensor2tensor/rl/player.py                    | 40 ++++++------
 tensor2tensor/rl/player_utils.py              | 62 ++++++++++++-------
 tensor2tensor/rl/rl_utils.py                  |  1 +
 tensor2tensor/rl/trainer_model_based.py       |  3 +-
 7 files changed, 71 insertions(+), 46 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 7a3a0a943..a777b5aac 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -672,7 +672,12 @@ def setup_and_load_epoch(hparams, data_dir, which_epoch_data=None):
     """Load T2TBatchGymEnv with data from one epoch.
 
     Args:
-        which_epoch_data: data from which epoch to load.
+      hparams: hparams.
+      data_dir: data directory.
+      which_epoch_data: data from which epoch to load.
+
+    Returns:
+      env.
     """
     t2t_env = T2TGymEnv.setup_env_from_hparams(
         hparams, batch_size=hparams.real_batch_size,
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 8062f5bec..f5a46e64b 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -161,6 +161,7 @@ def initialize(self, sess):
     model_loader = tf.train.Saver(
         var_list=tf.global_variables(scope="next_frame*")  # pylint:disable=unexpected-keyword-arg
     )
+    # TODO(afrozm): use TF methods to be on the safe side here.
     if os.path.isdir(self._model_dir):
       trainer_lib.restore_checkpoint(
           self._model_dir, saver=model_loader, sess=sess, must_restore=True
diff --git a/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index dae336748..7bc4c276a 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -26,7 +26,8 @@
 
 
 class FlatBatchEnv(Env):
-  """Gym environment interface for Batched Environments (with batch size = 1)"""
+  """Gym environment interface for Batched Environments (with batch size 1)."""
+
   def __init__(self, batch_env):
     if batch_env.batch_size != 1:
       raise ValueError("Number of environments in batch must be equal to one")
diff --git a/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index ea92e9f14..9ab0a7023 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Play with a world model.
+r"""Play with a world model.
 
 Controls:
   WSAD and SPACE to control the agent.
@@ -49,22 +49,20 @@
 from __future__ import division
 from __future__ import print_function
 
-import six
-
 import gym
 from gym.envs.atari.atari_env import ACTION_MEANING
 from gym.utils import play
 import numpy as np
+import six
 
-from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image, PIL_ImageDraw
-from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
-from tensor2tensor.rl.player_utils import wrap_with_monitor, \
-  load_data_and_make_simulated_env, infer_paths
-# Import flags from t2t_trainer and trainer_model_based
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
-import tensor2tensor.rl.trainer_model_based_params # pylint: disable=unused-import
-
 from tensor2tensor.data_generators.gym_env import T2TGymEnv
+from tensor2tensor.rl import player_utils
+from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image
+from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw
+from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
+# Import flags from t2t_trainer and trainer_model_based
+import tensor2tensor.rl.trainer_model_based_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
 import tensorflow as tf
 
@@ -98,10 +96,7 @@
 
 
 class PlayerEnvWrapper(gym.Wrapper):
-  """ Environment Wrapper for gym.utils.play.
-
-  This probably will be highly refactored.
-  """
+  """Environment Wrapper for gym.utils.play."""
 
   RESET_ACTION = 101
   TOGGLE_WAIT_ACTION = 102
@@ -222,16 +217,17 @@ def main(_):
   # Not important for experiments past 2018
   if "wm_policy_param_sharing" not in hparams.values().keys():
     hparams.add_hparam("wm_policy_param_sharing", False)
-  directories = infer_paths(output_dir=FLAGS.output_dir,
-                            world_model=FLAGS.wm_dir,
-                            policy=FLAGS.policy_dir,
-                            data=FLAGS.episodes_data_dir)
+  directories = player_utils.infer_paths(
+      output_dir=FLAGS.output_dir,
+      world_model=FLAGS.wm_dir,
+      policy=FLAGS.policy_dir,
+      data=FLAGS.episodes_data_dir)
   epoch = FLAGS.epoch if FLAGS.epoch == "last" else int(FLAGS.epoch)
 
   if FLAGS.simulated_env:
-    env = load_data_and_make_simulated_env(directories["data"],
-                                           directories["world_model"],
-                                           hparams, which_epoch_data=epoch)
+    env = player_utils.load_data_and_make_simulated_env(
+        directories["data"], directories["world_model"],
+        hparams, which_epoch_data=epoch)
   else:
     env = T2TGymEnv.setup_and_load_epoch(
         hparams, data_dir=directories["data"],
@@ -240,7 +236,7 @@ def main(_):
 
   env = PlayerEnvWrapper(env)  # pylint: disable=redefined-variable-type
 
-  env = wrap_with_monitor(env, FLAGS.video_dir)
+  env = player_utils.wrap_with_monitor(env, FLAGS.video_dir)
 
   if FLAGS.dry_run:
     for _ in range(5):
diff --git a/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index 3cd118719..7cd2babf8 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -19,19 +19,18 @@
 from __future__ import division
 from __future__ import print_function
 
+import copy
 import os
-from copy import deepcopy
 
 import gym
-import six
-from gym import wrappers, spaces
 import numpy as np
+import six
 
-from tensor2tensor.rl import rl_utils
-from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
 from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.models.research.rl import get_policy
 from tensor2tensor.models.research.rl import make_simulated_env_fn_from_hparams
+from tensor2tensor.rl import rl_utils
+from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
 from tensor2tensor.utils import trainer_lib
 import tensorflow as tf
 
@@ -61,7 +60,7 @@ def make_simulated_gym_env(real_env, world_model_dir, hparams, random_starts):
 def load_data_and_make_simulated_env(
     data_dir, wm_dir, hparams, which_epoch_data="last", random_starts=True
 ):
-  hparams = deepcopy(hparams)
+  hparams = copy.deepcopy(hparams)
   t2t_env = T2TGymEnv.setup_and_load_epoch(
       hparams, data_dir=data_dir,
       which_epoch_data=which_epoch_data)
@@ -71,11 +70,9 @@ def load_data_and_make_simulated_env(
 
 
 class ExtendToEvenDimentions(gym.ObservationWrapper):
-  """ Force even dimentions of both height and width.
-
-  Specifically, it adds single zero row/column to observations if needed.
-  """
+  """Force even dimentions of both height and width by adding zeros."""
   HW_AXES = (0, 1)
+
   def __init__(self, env):
     gym.ObservationWrapper.__init__(self, env)
 
@@ -86,7 +83,7 @@ def __init__(self, env):
         extended_shape[axis] += 1
 
     assert env.observation_space.dtype == np.uint8
-    self.observation_space = spaces.Box(
+    self.observation_space = gym.spaces.Box(
         low=0,
         high=255,
         shape=extended_shape,
@@ -109,6 +106,7 @@ def if_odd(self, n):
 
 class RenderObservations(gym.Wrapper):
   """Add observations rendering in 'rgb_array' mode."""
+
   def __init__(self, env):
     super(RenderObservations, self).__init__(env)
     if "rgb_array" not in self.metadata["render.modes"]:
@@ -129,17 +127,24 @@ def render(self, mode="human", **kwargs):
 
 
 def wrap_with_monitor(env, video_dir):
-  """Wrap environment with gym.Monitor
+  """Wrap environment with gym.Monitor.
 
   Video recording provided by Monitor requires
     1) both height and width of observation to be even numbers.
     2) rendering of environment
+
+  Args:
+    env: environment.
+    video_dir: video directory.
+
+  Returns:
+    wrapped environment.
   """
   env = ExtendToEvenDimentions(env)
   env = RenderObservations(env)  # pylint: disable=redefined-variable-type
-  env = wrappers.Monitor(env, video_dir, force=True,
-                         video_callable=lambda idx: True,
-                         write_upon_reset=True)
+  env = gym.wrappers.Monitor(env, video_dir, force=True,
+                             video_callable=lambda idx: True,
+                             write_upon_reset=True)
   return env
 
 
@@ -171,9 +176,10 @@ def create_simulated_env(
       generative_model_params=generative_model_params,
       **other_hparams
   )
-  return load_data_and_make_simulated_env(output_dir, hparams,
-                                          which_epoch_data=which_epoch_data,
-                                          random_starts=random_starts)
+  return load_data_and_make_simulated_env(
+      output_dir, wm_dir=None, hparams=hparams,
+      which_epoch_data=which_epoch_data,
+      random_starts=random_starts)
 
 
 class PPOPolicyInferencer(object):
@@ -187,6 +193,7 @@ class PPOPolicyInferencer(object):
     >>>   logits, value = ppo.infer(ob)
     >>>   ob, _, done, _ = env.step(action)
   """
+
   def __init__(self, hparams, action_space, observation_space, policy_dir):
     assert hparams.base_algo == "ppo"
     ppo_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
@@ -229,16 +236,22 @@ def infer(self, ob):
 
     Args:
       ob: array of shape (height, width, channels)
+
+    Returns:
+      logits and vf.
     """
     self._add_to_stack(ob)
     logits, vf = self.infer_from_frame_stack(self._frame_stack)
     return logits, vf
 
   def infer_from_frame_stack(self, ob_stack):
-    """ Infer policy from stack of observations.
+    """Infer policy from stack of observations.
 
     Args:
       ob_stack: array of shape (1, frame_stack_size, height, width, channels)
+
+    Returns:
+      logits and vf.
     """
     logits, vf = self.sess.run([self.logits_t, self.value_function_t],
                                feed_dict={self.obs_t: ob_stack})
@@ -246,15 +259,22 @@ def infer_from_frame_stack(self, ob_stack):
 
 
 def infer_paths(output_dir, **subdirs):
-  """
+  """Infers standard paths to policy and model directories.
 
   Example:
     >>> infer_paths("/some/output/dir/", policy="", model="custom/path")
     {"policy": "/some/output/dir/policy", "model": "custom/path",
     "output_dir":"/some/output/dir/"}
+
+  Args:
+    output_dir: output directory.
+    **subdirs: sub-directories.
+
+  Returns:
+    a dictionary with the directories.
   """
   directories = dict()
   for name, path in six.iteritems(subdirs):
     directories[name] = path if path else os.path.join(output_dir, name)
-  directories['output_dir'] = output_dir
+  directories["output_dir"] = output_dir
   return directories
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 19293e151..2a9726eaa 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -130,6 +130,7 @@ def choose_subsequence():
 def make_initial_frame_chooser(real_env, frame_stack_size,
                                simulation_random_starts,
                                simulation_flip_first_random_for_beginning):
+  """Make frame chooser."""
   initial_frame_rollouts = real_env.current_epoch_rollouts(
       split=tf.estimator.ModeKeys.TRAIN,
       minimal_rollout_frames=frame_stack_size,
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 8aeaba516..e5daf6f34 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -36,7 +36,6 @@
 import numpy as np
 import six
 
-from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw, PIL_Image
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
 from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.layers import common_video
@@ -44,6 +43,8 @@
 from tensor2tensor.models.research.rl import make_simulated_env_fn_from_hparams
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl import trainer_model_based_params
+from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image
+from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw
 from tensor2tensor.rl.restarter import Restarter
 from tensor2tensor.utils import trainer_lib
 

From d6667c5c9ac4f2ada01a4f1765aeac090cdddc41 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 7 Jan 2019 18:07:21 -0800
Subject: [PATCH 1454/2720] Add LinearKernel for GPs on top of neural net
 features.

PiperOrigin-RevId: 228268473
---
 tensor2tensor/layers/bayes.py | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 65d11d714..b9be046d5 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -90,6 +90,41 @@ def get_config(self):
     return {'variance': self.variance, 'lengthscale': self.lengthscale}
 
 
+class LinearKernel(object):
+  """Linear kernel, optionally on top of a feature extractor (e.g., encoder)."""
+
+  def __init__(self, variance, bias, encoder=tf.identity):
+    self.variance = variance
+    self.bias = bias
+    self.encoder = encoder
+
+  def __call__(self, x1, x2):
+    """Computes scaled dot product of over all pairs of encoded inputs.
+
+    Args:
+      x1: Tensor of shape [batch_x1] + encoder domain. Slices along the batch
+        axis denote an individual input to be passed to the kernel. It is
+        computed pairwise with each input sliced from x2.
+      x2: Tensor of shape [batch_x2] + encoder domain. Slices along the batch
+        axis denote an individual input to be passed to the kernel. It is
+        computed pairwise with each input sliced from x1.
+
+    Returns:
+      Tensor of shape [batch_x1, batch_x2].
+    """
+    encoded_x1 = self.encoder(x1)
+    encoded_x2 = self.encoder(x2)
+    dot_product = tf.matmul(encoded_x1, encoded_x2, transpose_b=True)
+    return tf.sqrt(self.variance) * dot_product + self.bias
+
+  def get_config(self):
+    return {
+        'variance': self.variance,
+        'bias': self.bias,
+        'encoder': tf.keras.utils.serialize_keras_object(self.encoder),
+    }
+
+
 # TODO(dusenberrymw): Restructure the implementation of a trainable initializer
 # such that callers do not need to have type-conditional logic.
 class TrainableInitializer(tf.keras.initializers.Initializer):

From d2fd2163a704cc2e2cab1fc5b3c232c602c7d9b6 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 7 Jan 2019 19:02:51 -0800
Subject: [PATCH 1455/2720] Add the option to explicitly define a
 "targets_mask" feature to limit loss to select tokens. If targets_mask is
 unspecified computation proceeds as before. "targets_mask" is multiplied by
 target_modality.target_weights_fn().

PiperOrigin-RevId: 228272768
---
 tensor2tensor/utils/t2t_model.py      | 51 ++++++++++++++++++++++++---
 tensor2tensor/utils/t2t_model_test.py | 38 ++++++++++++++++++++
 2 files changed, 84 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index d95bd98d7..0d242fb4e 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -224,7 +224,7 @@ def summarize_hparams(self):
     def create_hparams_summary(hparams, name):
       hparams_strs = [tf.convert_to_tensor([k, str(v)])
                       for k, v in hparams.values().items()]
-      tf.summary.text(name, tf.stack(hparams_strs))
+      tf.summary.text(name, tf.cast(tf.stack(hparams_strs), tf.string))
 
     create_hparams_summary(self._hparams, "%s_hparams" % self.name)
     if self._problem_hparams:
@@ -573,7 +573,7 @@ def top(self, body_output, features):
         target_modality = target_modality["targets"]
       return self._top_single(body_output, target_modality, features)
 
-  def _loss_single(self, logits, target_modality, feature):
+  def _loss_single(self, logits, target_modality, feature, weights=None):
     # The current bfloat16 version still uses float32 for most parts of backward
     # propagation to keep model quality, so cast back before computing the loss
     # value.
@@ -582,11 +582,44 @@ def _loss_single(self, logits, target_modality, feature):
       return (tf.constant(0., dtype=tf.float32),
               tf.constant(1., dtype=tf.float32))
 
-    loss_num, loss_den = target_modality.loss(logits, feature)
+    # Calculate loss contribution.
+    if weights is None:
+      loss_num, loss_den = target_modality.loss(logits, feature)
+    else:
+
+      def weights_fn(labels):
+        """Per-token weights for loss."""
+        # Use target_weights_fn() given by modality as well as explicitly given
+        # weights.
+        modality_weights = target_modality.targets_weights_fn(labels)
+
+        # Broadcast 'weights' along minor dimensions (TF's default is major).
+        explicit_weights = weights
+        if len(explicit_weights.shape) < len(modality_weights.shape):
+          explicit_weights = common_layers.expand_squeeze_to_nd(
+              weights, modality_weights.shape.ndims)
+
+        return explicit_weights * modality_weights
+
+      # Ensure that target.modality_loss() supports "weights_fn" keyword
+      # argument. If it doesn't and "weights" is specified, raise an exception.
+      argument_names = inspect.getargspec(target_modality.loss).args
+      if "weights_fn" not in argument_names:
+        raise ValueError(
+            "Explicit 'weights' given but target_modality.loss doesn't "
+            "support 'weights_fn' keyword argument: %s.loss(%s)." %
+            (type(target_modality), ", ".join(argument_names)))
+
+      loss_num, loss_den = target_modality.loss(
+          logits, feature, weights_fn=weights_fn)
+
     loss_num *= self._problem_hparams.loss_multiplier
 
     if hasattr(self.hparams, "problem") and hasattr(
         self.hparams.problem, "task_list"):
+      if weights_fn is not None:
+        raise NotImplementedError("weights not yet implemented in "
+                                  "multitask setting.")
       loss_num, loss_den, summaries = multi_problem.aggregate_task_losses(
           self.hparams,
           self._problem_hparams,
@@ -613,7 +646,11 @@ def loss(self, logits, features):
             "problem_hparams.modality's dict." % k)
       losses = {}
       for k, v in six.iteritems(logits):
-        losses[k] = self._loss_single(v, target_modality[k], features[k])
+        losses[k] = self._loss_single(
+            v,
+            target_modality[k],
+            features[k],
+            weights=features.get(k + "_mask"))
 
         n, d = losses[k]
         if common_layers.should_generate_summaries():
@@ -637,7 +674,11 @@ def loss(self, logits, features):
             "model_body returned single logits so 'targets' must be a key "
             "since problem_hparams.modality is a dict.")
         target_modality = target_modality["targets"]
-      return self._loss_single(logits, target_modality, features["targets"])
+      return self._loss_single(
+          logits,
+          target_modality,
+          features["targets"],
+          weights=features.get("targets_mask"))
 
   def optimize(self, loss, num_async_replicas=1, use_tpu=False):
     """Return a training op minimizing loss."""
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index bcc0ed44a..9cb71103e 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.utils import modality
 from tensor2tensor.utils import t2t_model
 
 import tensorflow as tf
@@ -38,5 +39,42 @@ def testSummarizeLosses(self):
           len(tf.get_collection(tf.GraphKeys.SUMMARIES, scope="losses")),
           len(losses))
 
+  def testLossSingleWeights(self):
+    """Ensure _loss_single() respects optional 'weights' argument."""
+    with tf.Graph().as_default():
+      with self.test_session() as sess:
+        batch_size = 2
+        sequence_size = 16
+        vocab_size = 3
+
+        model_hparams = tf.contrib.training.HParams()
+        model_hparams.label_smoothing = 0.0
+        model_hparams.shared_embedding_and_softmax_weights = False
+
+        problem_hparams = tf.contrib.training.HParams()
+        problem_hparams.loss_multiplier = 1.0
+        problem_hparams.modality = {}
+
+        model = t2t_model.T2TModel(
+            model_hparams, problem_hparams=problem_hparams)
+        logits = tf.zeros((batch_size, sequence_size, 1, 1, vocab_size))
+        target_modality = modality.Modality(model_hparams)
+        feature = tf.ones((batch_size, sequence_size, 1, 1))
+
+        # all-zero weights == zero loss.
+        weights = tf.zeros((batch_size, sequence_size))
+        loss_num, loss_denom = model._loss_single(
+            logits, target_modality, feature, weights=weights)
+        self.assertAllClose(tf.zeros_like(loss_num), sess.run(loss_num))
+        self.assertAllClose(tf.zeros_like(loss_denom), sess.run(loss_denom))
+
+        # non-zero weights > zero loss.
+        weights = tf.ones((batch_size, sequence_size))
+        loss_num, loss_denom = model._loss_single(
+            logits, target_modality, feature, weights=weights)
+        self.assertAllLess(0.0, sess.run(loss_num))
+        self.assertAllClose(batch_size * sequence_size, sess.run(loss_denom))
+
+
 if __name__ == "__main__":
   tf.test.main()

From 380c9ebc570fdae0c2e2f63eadd331adb5453694 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 7 Jan 2019 20:28:15 -0800
Subject: [PATCH 1456/2720] Adds "max_frames_per_video" to
 video_utils.VideoProblem.

PiperOrigin-RevId: 228277884
---
 .../data_generators/bair_robot_pushing.py     |  3 +++
 tensor2tensor/data_generators/video_utils.py  | 27 ++++++++++++++++++-
 .../data_generators/video_utils_test.py       |  8 +++---
 tensor2tensor/utils/video_metrics.py          |  5 +++-
 4 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index 141844993..a92303013 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -72,6 +72,9 @@ def is_generate_per_split(self):
   def total_number_of_frames(self):
     return 167 * 256 * 30
 
+  def max_frames_per_video(self, hparams):
+    return 30
+
   @property
   def random_skip(self):
     return False
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index d5f51b5a1..3f1c1d54a 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -123,6 +123,7 @@ def convert_videos_to_summaries(input_videos, output_videos, target_videos,
   fps = decode_hparams.frames_per_second
   border_percent = decode_hparams.border_percent
   max_outputs = decode_hparams.max_display_outputs
+  target_steps = target_videos.shape[1]
   all_summaries = []
   input_videos = create_border(
       input_videos, color="blue", border_percent=border_percent)
@@ -146,7 +147,8 @@ def convert_videos_to_summaries(input_videos, output_videos, target_videos,
     all_summaries.extend(input_summ_vals)
 
   # Frame-by-frame summaries
-  iterable = zip(all_input[:max_outputs], all_output[:max_outputs])
+  iterable = zip(output_videos[:max_outputs, :target_steps],
+                 target_videos[:max_outputs])
   for ind, (input_video, output_video) in enumerate(iterable):
     t, h, w, c = input_video.shape
     # Tile vertically
@@ -259,6 +261,25 @@ def __init__(self, *args, **kwargs):
     self.settable_use_not_breaking_batching = True
     self.shuffle = True
 
+  def max_frames_per_video(self, hparams):
+    """Maximum number of frames per video as determined by the dataset.
+
+    This is used only in PREDICT mode and handles the corner case where
+    video_num_input_frames + video_num_target_frames is greater than the
+    maximum number of frames per video in the dataset. For eg, 30 in BAIR.
+
+    For this special case, setting this to return "x" limits the input pipeline
+    to handle "x" (input + target) frames. The corresponding video model can
+    then decode arbitrary number of target frames via
+    hparams.video_num_target_frames.
+
+    Args:
+      hparams: tf.contrib.training.HParams.
+    Returns:
+      num_frames: int.
+    """
+    return hparams.video_num_input_frames + hparams.video_num_target_frames
+
   @property
   def num_channels(self):
     """Number of color channels in each frame."""
@@ -491,8 +512,12 @@ def check_integrity_and_batch(*datasets):
       return dataset
 
     preprocessed_dataset = dataset.map(_preprocess)
+
     num_frames = (
         hparams.video_num_input_frames + hparams.video_num_target_frames)
+    if mode == tf.estimator.ModeKeys.PREDICT:
+      num_frames = min(self.max_frames_per_video(hparams), num_frames)
+
     # We jump by a random position at the beginning to add variety.
     if (self.random_skip and self.settable_random_skip and interleave and
         mode == tf.estimator.ModeKeys.TRAIN):
diff --git a/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index b828fd95e..62304044a 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -70,12 +70,14 @@ def testDecodeInMemoryTrue(self):
     metrics = video_utils.summarize_video_metrics(decode_hooks)
 
   @parameterized.named_parameters(
-      ("two", 5), ("ten", 10))
-  def testConvertPredictionsToVideoSummaries(self, num_decodes=2):
+      ("d5_o6", 5, 6))
+      # ("d5", 5), ("d10", 10), ("d5_o6", 5, 6))
+  def testConvertPredictionsToVideoSummaries(self, num_decodes=5,
+                                             max_output_steps=5):
     # Initialize predictions.
     rng = np.random.RandomState(0)
     inputs = rng.randint(0, 255, (2, 32, 32, 3))
-    outputs = rng.randint(0, 255, (5, 32, 32, 3))
+    outputs = rng.randint(0, 255, (max_output_steps, 32, 32, 3))
     targets = rng.randint(0, 255, (5, 32, 32, 3))
 
     # batch it up.
diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index 04df03299..5a58c2471 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -115,7 +115,10 @@ def get_zipped_dataset_from_predictions(predictions):
   """Creates dataset from in-memory predictions."""
   targets = stack_data_given_key(predictions, "targets")
   outputs = stack_data_given_key(predictions, "outputs")
-  num_videos = len(targets)
+  num_videos, num_steps = targets.shape[:2]
+
+  # Truncate output time-steps to match target time-steps
+  outputs = outputs[:, :num_steps]
 
   targets_placeholder = tf.placeholder(targets.dtype, targets.shape)
   outputs_placeholder = tf.placeholder(outputs.dtype, outputs.shape)

From 5615db2bb0718cc0e94ecc43d5ede5ab68633166 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 8 Jan 2019 07:06:37 +0100
Subject: [PATCH 1457/2720] Disable weight decay for policy (#1349)

---
 tensor2tensor/models/research/rl.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 8d5713684..87acbaa58 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -47,6 +47,7 @@ def ppo_base_v1():
   hparams.learning_rate_schedule = "constant"
   hparams.learning_rate_constant = 1e-4
   hparams.clip_grad_norm = 0.5
+  hparams.weight_decay = 0
   # If set, extends the LR warmup to all epochs except the final one.
   hparams.add_hparam("lr_decay_in_final_epoch", False)
   hparams.add_hparam("init_mean_factor", 0.1)
@@ -141,6 +142,8 @@ def ppo_original_world_model():
       hparams.set_hparam(name, value)
     else:
       hparams.add_hparam(name, value)
+  # Mostly to avoid decaying WM params when training the policy.
+  hparams.weight_decay = 0
   return hparams
 
 
@@ -156,6 +159,7 @@ def ppo_tiny_world_model():
       hparams.set_hparam(name, value)
     else:
       hparams.add_hparam(name, value)
+  hparams.weight_decay = 0
   return hparams
 
 
@@ -173,6 +177,7 @@ def ppo_original_world_model_stochastic_discrete():
       hparams.add_hparam(name, value)
   # To avoid OOM. Probably way to small.
   hparams.optimization_batch_size = 1
+  hparams.weight_decay = 0
   return hparams
 
 

From 41c8754e75570488672f08f4b287d312188e0390 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 8 Jan 2019 10:12:52 -0800
Subject: [PATCH 1458/2720] Add gelu unit to common layers and correct lint.

PiperOrigin-RevId: 228351383
---
 tensor2tensor/layers/common_layers.py | 38 ++++++++++++++++++++-------
 1 file changed, 28 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 0068d9e86..d1930e5d0 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -18,11 +18,9 @@
 from __future__ import division
 from __future__ import print_function
 
-from collections import defaultdict
-
+import collections
 import contextlib
 import functools
-from functools import partial
 import math
 
 import numpy as np
@@ -1253,7 +1251,7 @@ def dense_dropconnect(inputs,
 
   if dropconnect_dropout != 0.0:
     tf.logging.info("Applying dropconnect as the kernel regularization.")
-    kwargs["kernel_regularizer"] = partial(
+    kwargs["kernel_regularizer"] = functools.partial(
         tf.nn.dropout, keep_prob=1.0 - dropconnect_dropout)
 
   return dense(inputs, output_size, use_bias=True, name=name, **kwargs)
@@ -1797,7 +1795,7 @@ def split_to_discretized_mix_logistic_params(inputs):
     shape [batch, height, width, num_mixtures]. Other parameters have shape
     [batch, height, width, num_mixtures, 3].
   """
-  batch, height, width, output_dim = shape_list(inputs)
+  batch, height, width, output_dim = shape_list(inputs)  # pylint: disable=unbalanced-tuple-unpacking
   num_mixtures = output_dim // 10
   logits, locs, log_scales, coeffs = tf.split(
       inputs,
@@ -1852,7 +1850,7 @@ def discretized_mix_logistic_loss(pred, labels):
       pred)
 
   # Tile labels to broadcast compute across the mixture dimension.
-  batch, height, width, num_mixtures = shape_list(logits)
+  batch, height, width, num_mixtures = shape_list(logits)  # pylint: disable=unbalanced-tuple-unpacking
   labels = tf.tile(
       tf.reshape(labels, [batch, height, width, 1, 3]),
       [1, 1, 1, num_mixtures, 1])
@@ -2333,7 +2331,8 @@ def ravanbakhsh_set_layer(layer_size,
 def fn_device_dependency_dict():
   """State container for fn_device_dependency."""
   if not hasattr(tf.get_default_graph(), "dependency_dict"):
-    setattr(tf.get_default_graph(), "dependency_dict", defaultdict(list))
+    setattr(tf.get_default_graph(), "dependency_dict",
+            collections.defaultdict(list))
   return tf.get_default_graph().dependency_dict
 
 
@@ -3064,6 +3063,23 @@ def belu(x):
   return tf.reshape(tf.concat([y1, y2], axis=-1), x_shape)
 
 
+def gelu(x):
+  """Gaussian Error Linear Unit.
+
+  This is a smoother version of the RELU.
+  Original paper: https://arxiv.org/abs/1606.08415
+
+  Args:
+    x: float Tensor to perform activation.
+
+  Returns:
+    x with the GELU activation applied.
+  """
+  cdf = 0.5 * (1.0 + tf.tanh(
+      (np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3)))))
+  return x * cdf
+
+
 def nac(x, depth, name=None, reuse=None):
   """NAC as in https://arxiv.org/abs/1808.00508."""
   with tf.variable_scope(name, default_name="nac", values=[x], reuse=reuse):
@@ -3348,7 +3364,7 @@ def deep_discriminator(x,
   """Discriminator architecture based on InfoGAN."""
   with tf.variable_scope(
       "discriminator", initializer=tf.random_normal_initializer(stddev=0.02)):
-    batch_size, height, width = shape_list(x)[:3]
+    batch_size, height, width = shape_list(x)[:3]  # pylint: disable=unbalanced-tuple-unpacking
     net = tf.layers.conv2d(
         x, filters, filter_size, strides=stride, padding="SAME", name="conv1")
     net = lrelu(net)
@@ -3505,7 +3521,7 @@ def double_discriminator(x, filters1=128, filters2=None,
 
 def upscale(inputs, f, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR):
   """Upscaling the image by a factor of f."""
-  height, width = shape_list(inputs)[1:3]
+  height, width = shape_list(inputs)[1:3]  # pylint: disable=unbalanced-tuple-unpacking
   return tf.image.resize_images(inputs, (height * f, width * f), method)
 
 
@@ -3786,7 +3802,9 @@ def sparse_eye(size):
 # modification from https://github.com/tensorflow/tensorflow/pull/21276
 # without special initialization for g
 class WeightNorm(tf.keras.layers.Wrapper):
-  """ This wrapper reparameterizes a layer by decoupling the weight's
+  """Decouple weight magnitude and direction.
+
+  This wrapper reparameterizes a layer by decoupling the weight's
   magnitude and direction. This speeds up convergence by improving the
   conditioning of the optimization problem.
 

From cf1c451d60099d4a291f24f507d29becdf7e1340 Mon Sep 17 00:00:00 2001
From: cclauss <cclauss@me.com>
Date: Tue, 8 Jan 2019 20:20:06 +0100
Subject: [PATCH 1459/2720] Fix invalid list comprehension (#1343)

Fixes #1002 as the current code will raise a NameError at runtime.
---
 tensor2tensor/utils/t2t_model.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 0d242fb4e..4578c4d69 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -338,9 +338,9 @@ def model_fn_sharded(self, sharded_features):
             sharded_logits[k] = dp(self.top, v, datashard_to_features)
             sharded_losses[k] = dp(self.loss, sharded_logits[k],
                                    datashard_to_features)
-          training_loss_dict = average_sharded_losses([{
+          training_loss_dict = average_sharded_losses([({
               "training": l
-          } for l in loss for loss in sharded_losses.values()])
+          } for l in loss) for loss in sharded_losses.values()])
           losses.update(training_loss_dict)
         else:
           sharded_logits = dp(self.top, body_out, datashard_to_features)

From c54464a3a2f0002dbf2f642dd7b75a9012d96409 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Wed, 9 Jan 2019 00:00:21 +0100
Subject: [PATCH 1460/2720] Add a timelimit for eval and implement a separate
 script for full eval (#1348)

* Add a timelimit for eval

* Implement a separate evaluation script
---
 tensor2tensor/data_generators/gym_env.py      | 53 ------------
 tensor2tensor/models/research/rl.py           |  1 +
 tensor2tensor/rl/evaluator.py                 | 83 +++++++++++++++++++
 tensor2tensor/rl/player.py                    |  3 +-
 tensor2tensor/rl/player_utils.py              | 39 ++++++++-
 tensor2tensor/rl/rl_utils.py                  | 32 ++++++-
 tensor2tensor/rl/trainer_model_based.py       | 17 +---
 .../rl/trainer_model_based_params.py          |  3 +
 tensor2tensor/rl/trainer_model_free.py        |  9 +-
 9 files changed, 164 insertions(+), 76 deletions(-)
 create mode 100644 tensor2tensor/rl/evaluator.py

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index a777b5aac..14aa4ab63 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -23,7 +23,6 @@
 import itertools
 import os
 import random
-import re
 
 from gym.spaces import Box
 import numpy as np
@@ -642,58 +641,6 @@ def base_env_name(self):
   def num_channels(self):
     return self.observation_space.shape[2]
 
-  @staticmethod
-  def infer_last_epoch_num(data_dir):
-    """Infer highest epoch number from file names in data_dir."""
-    names = os.listdir(data_dir)
-    epochs_str = [re.findall(pattern=r".*\.(-?\d+)$", string=name)
-                  for name in names]
-    epochs_str = sum(epochs_str, [])
-    return max([int(epoch_str) for epoch_str in epochs_str])
-
-  @staticmethod
-  def setup_env_from_hparams(hparams, batch_size, max_num_noops):
-    game_mode = "NoFrameskip-v4"
-    camel_game_name = misc_utils.snakecase_to_camelcase(hparams.game)
-    camel_game_name += game_mode
-    env_name = camel_game_name
-
-    env = T2TGymEnv(base_env_name=env_name,
-                    batch_size=batch_size,
-                    grayscale=hparams.grayscale,
-                    resize_width_factor=hparams.resize_width_factor,
-                    resize_height_factor=hparams.resize_height_factor,
-                    rl_env_max_episode_steps=hparams.rl_env_max_episode_steps,
-                    max_num_noops=max_num_noops, maxskip_envs=True)
-    return env
-
-  @staticmethod
-  def setup_and_load_epoch(hparams, data_dir, which_epoch_data=None):
-    """Load T2TBatchGymEnv with data from one epoch.
-
-    Args:
-      hparams: hparams.
-      data_dir: data directory.
-      which_epoch_data: data from which epoch to load.
-
-    Returns:
-      env.
-    """
-    t2t_env = T2TGymEnv.setup_env_from_hparams(
-        hparams, batch_size=hparams.real_batch_size,
-        max_num_noops=hparams.max_num_noops
-    )
-    # Load data.
-    if which_epoch_data is not None:
-      if which_epoch_data == "last":
-        which_epoch_data = T2TGymEnv.infer_last_epoch_num(data_dir)
-      assert isinstance(which_epoch_data, int), \
-        "{}".format(type(which_epoch_data))
-      t2t_env.start_new_epoch(which_epoch_data, data_dir)
-    else:
-      t2t_env.start_new_epoch(-999)
-    return t2t_env
-
   def _derive_observation_space(self, orig_observ_space):
     height, width, channels = orig_observ_space.shape
     if self.grayscale:
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 87acbaa58..ff2b0d72d 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -335,6 +335,7 @@ def rlmf_original():
       frame_stack_size=4,
       eval_sampling_temps=[0.0, 0.2, 0.5, 0.8, 1.0, 2.0],
       eval_max_num_noops=8,
+      eval_rl_env_max_episode_steps=1000,
       resize_height_factor=2,
       resize_width_factor=2,
       grayscale=0,
diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
new file mode 100644
index 000000000..c631b43b2
--- /dev/null
+++ b/tensor2tensor/rl/evaluator.py
@@ -0,0 +1,83 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+r"""Evaluation script for RL agents.
+
+Example invocation:
+
+python -m tensor2tensor.rl.evaluator \
+    --policy_dir=$HOME/t2t/rl_v1/policy \
+    --eval_metrics_dir=$HOME/t2t/rl_v1/full_eval_metrics \
+    --hparams_set=rlmb_base \
+    --hparams='batch_size=64'
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.models.research import rl  # pylint: disable=unused-import
+from tensor2tensor.rl import rl_utils
+from tensor2tensor.rl import trainer_model_based_params  # pylint: disable=unused-import
+from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
+from tensor2tensor.utils import trainer_lib
+
+import tensorflow as tf
+
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+
+flags.DEFINE_string("policy_dir", "", "Directory with policy checkpoints.")
+flags.DEFINE_string(
+    "eval_metrics_dir", "", "Directory to output the eval metrics at."
+)
+flags.DEFINE_bool("full_eval", True, "Whether to ignore the timestep limit.")
+
+
+def evaluate(hparams, policy_dir, eval_metrics_dir, report_fn=None,
+             report_metric=None):
+  """Evaluate."""
+  if report_fn:
+    assert report_metric is not None
+
+  eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
+  eval_metrics = rl_utils.evaluate_all_configs(hparams, policy_dir)
+  rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, 0)
+
+  # Report metrics
+  if report_fn:
+    if report_metric == "mean_reward":
+      metric_name = rl_utils.get_metric_name(
+          sampling_temp=hparams.eval_sampling_temps[0],
+          max_num_noops=hparams.eval_max_num_noops,
+          clipped=False
+      )
+      report_fn(eval_metrics[metric_name], 0)
+    else:
+      report_fn(eval_metrics[report_metric], 0)
+  return eval_metrics
+
+
+def main(_):
+  hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
+  if FLAGS.full_eval:
+    hparams.eval_rl_env_max_episode_steps = -1
+  evaluate(hparams, FLAGS.policy_dir, FLAGS.eval_metrics_dir)
+
+
+if __name__ == "__main__":
+  tf.app.run()
diff --git a/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 9ab0a7023..0dc680093 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -56,7 +56,6 @@
 import six
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
-from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.rl import player_utils
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw
@@ -229,7 +228,7 @@ def main(_):
         directories["data"], directories["world_model"],
         hparams, which_epoch_data=epoch)
   else:
-    env = T2TGymEnv.setup_and_load_epoch(
+    env = player_utils.setup_and_load_epoch(
         hparams, data_dir=directories["data"],
         which_epoch_data=epoch)
     env = FlatBatchEnv(env)
diff --git a/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index 7cd2babf8..5655afff7 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -21,6 +21,7 @@
 
 import copy
 import os
+import re
 
 import gym
 import numpy as np
@@ -39,6 +40,42 @@
 FLAGS = flags.FLAGS
 
 
+def infer_last_epoch_num(data_dir):
+  """Infer highest epoch number from file names in data_dir."""
+  names = os.listdir(data_dir)
+  epochs_str = [re.findall(pattern=r".*\.(-?\d+)$", string=name)
+                for name in names]
+  epochs_str = sum(epochs_str, [])
+  return max([int(epoch_str) for epoch_str in epochs_str])
+
+
+def setup_and_load_epoch(hparams, data_dir, which_epoch_data=None):
+  """Load T2TGymEnv with data from one epoch.
+
+  Args:
+    hparams: hparams.
+    data_dir: data directory.
+    which_epoch_data: data from which epoch to load.
+
+  Returns:
+    env.
+  """
+  t2t_env = rl_utils.setup_env(
+      hparams, batch_size=hparams.real_batch_size,
+      max_num_noops=hparams.max_num_noops
+  )
+  # Load data.
+  if which_epoch_data is not None:
+    if which_epoch_data == "last":
+      which_epoch_data = infer_last_epoch_num(data_dir)
+    assert isinstance(which_epoch_data, int), \
+      "{}".format(type(which_epoch_data))
+    t2t_env.start_new_epoch(which_epoch_data, data_dir)
+  else:
+    t2t_env.start_new_epoch(-999)
+  return t2t_env
+
+
 def make_simulated_gym_env(real_env, world_model_dir, hparams, random_starts):
   """Gym environment with world model."""
   initial_frame_chooser = rl_utils.make_initial_frame_chooser(
@@ -61,7 +98,7 @@ def load_data_and_make_simulated_env(
     data_dir, wm_dir, hparams, which_epoch_data="last", random_starts=True
 ):
   hparams = copy.deepcopy(hparams)
-  t2t_env = T2TGymEnv.setup_and_load_epoch(
+  t2t_env = setup_and_load_epoch(
       hparams, data_dir=data_dir,
       which_epoch_data=which_epoch_data)
   return make_simulated_gym_env(
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 2a9726eaa..876ccb78e 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -28,6 +28,7 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl.dopamine_connector import DQNLearner
 from tensor2tensor.rl.ppo_learner import PPOLearner
+from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -63,8 +64,9 @@ def evaluate_single_config(
 ):
   """Evaluate the PPO agent in the real environment."""
   eval_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
-  env = T2TGymEnv.setup_env_from_hparams(
-      hparams, batch_size=hparams.eval_batch_size, max_num_noops=max_num_noops
+  env = setup_env(
+      hparams, batch_size=hparams.eval_batch_size, max_num_noops=max_num_noops,
+      rl_env_max_episode_steps=hparams.eval_rl_env_max_episode_steps
   )
   env.start_new_epoch(0)
   env_fn = rl.make_real_env_fn(env)
@@ -98,12 +100,38 @@ def evaluate_all_configs(hparams, agent_model_dir):
   return metrics
 
 
+def summarize_metrics(eval_metrics_writer, metrics, epoch):
+  """Write metrics to summary."""
+  for (name, value) in six.iteritems(metrics):
+    summary = tf.Summary()
+    summary.value.add(tag=name, simple_value=value)
+    eval_metrics_writer.add_summary(summary, epoch)
+  eval_metrics_writer.flush()
+
+
 LEARNERS = {
     "ppo": PPOLearner,
     "dqn": DQNLearner,
 }
 
 
+def setup_env(hparams, batch_size, max_num_noops, rl_env_max_episode_steps=-1):
+  """Setup."""
+  game_mode = "NoFrameskip-v4"
+  camel_game_name = misc_utils.snakecase_to_camelcase(hparams.game)
+  camel_game_name += game_mode
+  env_name = camel_game_name
+
+  env = T2TGymEnv(base_env_name=env_name,
+                  batch_size=batch_size,
+                  grayscale=hparams.grayscale,
+                  resize_width_factor=hparams.resize_width_factor,
+                  resize_height_factor=hparams.resize_height_factor,
+                  rl_env_max_episode_steps=rl_env_max_episode_steps,
+                  max_num_noops=max_num_noops, maxskip_envs=True)
+  return env
+
+
 def update_hparams_from_hparams(target_hparams, source_hparams, prefix):
   """Copy a subset of hparams to target_hparams."""
   for (param_name, param_value) in six.iteritems(source_hparams.values()):
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index e5daf6f34..7a55aaad0 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -37,7 +37,6 @@
 import six
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
-from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research.rl import make_simulated_env_fn_from_hparams
@@ -378,15 +377,6 @@ def load_metrics(event_dir, epoch):
   return metrics
 
 
-def summarize_metrics(eval_metrics_writer, metrics, epoch):
-  """Write metrics to summary."""
-  for (name, value) in six.iteritems(metrics):
-    summary = tf.Summary()
-    summary.value.add(tag=name, simple_value=value)
-    eval_metrics_writer.add_summary(summary, epoch)
-  eval_metrics_writer.flush()
-
-
 def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   """Run the main training loop."""
   if report_fn:
@@ -401,9 +391,10 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   epoch = -1
   data_dir = directories["data"]
-  env = T2TGymEnv.setup_env_from_hparams(
+  env = rl_utils.setup_env(
       hparams, batch_size=hparams.real_batch_size,
-      max_num_noops=hparams.max_num_noops
+      max_num_noops=hparams.max_num_noops,
+      rl_env_max_episode_steps=hparams.rl_env_max_episode_steps
   )
   env.start_new_epoch(epoch, data_dir)
 
@@ -493,7 +484,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
         log("World model eval metrics:\n{}".format(pprint.pformat(wm_metrics)))
         metrics.update(wm_metrics)
 
-      summarize_metrics(eval_metrics_writer, metrics, epoch)
+      rl_utils.summarize_metrics(eval_metrics_writer, metrics, epoch)
 
       # Report metrics
       if report_fn:
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 610d979a8..463eac8cf 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -84,6 +84,8 @@ def _rlmb_base():
       # Sampling temperatures to try during eval.
       eval_sampling_temps=[0.5, 0.0, 1.0],
       eval_max_num_noops=8,
+      # To speed up the pipeline. Some games want to run forever.
+      eval_rl_env_max_episode_steps=1000,
 
       game="pong",
       # Whether to evaluate the world model in each iteration of the loop to get
@@ -506,6 +508,7 @@ def _rlmb_tiny_overrides():
       resize_width_factor=2,
       wm_eval_rollout_ratios=[1],
       rl_env_max_episode_steps=7,
+      eval_rl_env_max_episode_steps=7,
       simulated_rollout_length=2,
       eval_sampling_temps=[0.0, 1.0],
   )
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 64302dfdc..1d8d53162 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -20,7 +20,7 @@
 python -m tensor2tensor.rl.trainer_model_free \
     --output_dir=$HOME/t2t/rl_v1 \
     --hparams_set=pong_model_free \
-    --loop_hparams='batch_size=15'
+    --hparams='batch_size=15'
 """
 
 from __future__ import absolute_import
@@ -29,7 +29,6 @@
 
 import pprint
 
-from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
@@ -53,9 +52,9 @@
 
 def initialize_env_specs(hparams):
   """Initializes env_specs using T2TGymEnvs."""
-  env = T2TGymEnv.setup_env_from_hparams(
-      hparams, hparams.batch_size, hparams.eval_max_num_noops
-  )
+  env = rl_utils.setup_env(hparams, hparams.batch_size,
+                           hparams.eval_max_num_noops,
+                           hparams.rl_env_max_episode_steps)
   env.start_new_epoch(0)
 
   # TODO(afrozm): Decouple env_fn from hparams and return both, is there

From 4e620179c026f7f71f49d73582654cdf3a6907a6 Mon Sep 17 00:00:00 2001
From: cclauss <cclauss@me.com>
Date: Tue, 8 Jan 2019 11:20:24 -0800
Subject: [PATCH 1461/2720] internal merge of PR #1343

PiperOrigin-RevId: 228365154
---
 tensor2tensor/data_generators/gym_env.py      | 53 ++++++++++++
 tensor2tensor/models/research/rl.py           |  1 -
 tensor2tensor/rl/evaluator.py                 | 83 -------------------
 tensor2tensor/rl/player.py                    |  3 +-
 tensor2tensor/rl/player_utils.py              | 39 +--------
 tensor2tensor/rl/rl_utils.py                  | 32 +------
 tensor2tensor/rl/trainer_model_based.py       | 17 +++-
 .../rl/trainer_model_based_params.py          |  3 -
 tensor2tensor/rl/trainer_model_free.py        |  9 +-
 9 files changed, 76 insertions(+), 164 deletions(-)
 delete mode 100644 tensor2tensor/rl/evaluator.py

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 14aa4ab63..a777b5aac 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -23,6 +23,7 @@
 import itertools
 import os
 import random
+import re
 
 from gym.spaces import Box
 import numpy as np
@@ -641,6 +642,58 @@ def base_env_name(self):
   def num_channels(self):
     return self.observation_space.shape[2]
 
+  @staticmethod
+  def infer_last_epoch_num(data_dir):
+    """Infer highest epoch number from file names in data_dir."""
+    names = os.listdir(data_dir)
+    epochs_str = [re.findall(pattern=r".*\.(-?\d+)$", string=name)
+                  for name in names]
+    epochs_str = sum(epochs_str, [])
+    return max([int(epoch_str) for epoch_str in epochs_str])
+
+  @staticmethod
+  def setup_env_from_hparams(hparams, batch_size, max_num_noops):
+    game_mode = "NoFrameskip-v4"
+    camel_game_name = misc_utils.snakecase_to_camelcase(hparams.game)
+    camel_game_name += game_mode
+    env_name = camel_game_name
+
+    env = T2TGymEnv(base_env_name=env_name,
+                    batch_size=batch_size,
+                    grayscale=hparams.grayscale,
+                    resize_width_factor=hparams.resize_width_factor,
+                    resize_height_factor=hparams.resize_height_factor,
+                    rl_env_max_episode_steps=hparams.rl_env_max_episode_steps,
+                    max_num_noops=max_num_noops, maxskip_envs=True)
+    return env
+
+  @staticmethod
+  def setup_and_load_epoch(hparams, data_dir, which_epoch_data=None):
+    """Load T2TBatchGymEnv with data from one epoch.
+
+    Args:
+      hparams: hparams.
+      data_dir: data directory.
+      which_epoch_data: data from which epoch to load.
+
+    Returns:
+      env.
+    """
+    t2t_env = T2TGymEnv.setup_env_from_hparams(
+        hparams, batch_size=hparams.real_batch_size,
+        max_num_noops=hparams.max_num_noops
+    )
+    # Load data.
+    if which_epoch_data is not None:
+      if which_epoch_data == "last":
+        which_epoch_data = T2TGymEnv.infer_last_epoch_num(data_dir)
+      assert isinstance(which_epoch_data, int), \
+        "{}".format(type(which_epoch_data))
+      t2t_env.start_new_epoch(which_epoch_data, data_dir)
+    else:
+      t2t_env.start_new_epoch(-999)
+    return t2t_env
+
   def _derive_observation_space(self, orig_observ_space):
     height, width, channels = orig_observ_space.shape
     if self.grayscale:
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index ff2b0d72d..87acbaa58 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -335,7 +335,6 @@ def rlmf_original():
       frame_stack_size=4,
       eval_sampling_temps=[0.0, 0.2, 0.5, 0.8, 1.0, 2.0],
       eval_max_num_noops=8,
-      eval_rl_env_max_episode_steps=1000,
       resize_height_factor=2,
       resize_width_factor=2,
       grayscale=0,
diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
deleted file mode 100644
index c631b43b2..000000000
--- a/tensor2tensor/rl/evaluator.py
+++ /dev/null
@@ -1,83 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
-#
-# 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.
-
-r"""Evaluation script for RL agents.
-
-Example invocation:
-
-python -m tensor2tensor.rl.evaluator \
-    --policy_dir=$HOME/t2t/rl_v1/policy \
-    --eval_metrics_dir=$HOME/t2t/rl_v1/full_eval_metrics \
-    --hparams_set=rlmb_base \
-    --hparams='batch_size=64'
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.models.research import rl  # pylint: disable=unused-import
-from tensor2tensor.rl import rl_utils
-from tensor2tensor.rl import trainer_model_based_params  # pylint: disable=unused-import
-from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
-from tensor2tensor.utils import trainer_lib
-
-import tensorflow as tf
-
-
-flags = tf.flags
-FLAGS = flags.FLAGS
-
-
-flags.DEFINE_string("policy_dir", "", "Directory with policy checkpoints.")
-flags.DEFINE_string(
-    "eval_metrics_dir", "", "Directory to output the eval metrics at."
-)
-flags.DEFINE_bool("full_eval", True, "Whether to ignore the timestep limit.")
-
-
-def evaluate(hparams, policy_dir, eval_metrics_dir, report_fn=None,
-             report_metric=None):
-  """Evaluate."""
-  if report_fn:
-    assert report_metric is not None
-
-  eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
-  eval_metrics = rl_utils.evaluate_all_configs(hparams, policy_dir)
-  rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, 0)
-
-  # Report metrics
-  if report_fn:
-    if report_metric == "mean_reward":
-      metric_name = rl_utils.get_metric_name(
-          sampling_temp=hparams.eval_sampling_temps[0],
-          max_num_noops=hparams.eval_max_num_noops,
-          clipped=False
-      )
-      report_fn(eval_metrics[metric_name], 0)
-    else:
-      report_fn(eval_metrics[report_metric], 0)
-  return eval_metrics
-
-
-def main(_):
-  hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
-  if FLAGS.full_eval:
-    hparams.eval_rl_env_max_episode_steps = -1
-  evaluate(hparams, FLAGS.policy_dir, FLAGS.eval_metrics_dir)
-
-
-if __name__ == "__main__":
-  tf.app.run()
diff --git a/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 0dc680093..9ab0a7023 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -56,6 +56,7 @@
 import six
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
+from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.rl import player_utils
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw
@@ -228,7 +229,7 @@ def main(_):
         directories["data"], directories["world_model"],
         hparams, which_epoch_data=epoch)
   else:
-    env = player_utils.setup_and_load_epoch(
+    env = T2TGymEnv.setup_and_load_epoch(
         hparams, data_dir=directories["data"],
         which_epoch_data=epoch)
     env = FlatBatchEnv(env)
diff --git a/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index 5655afff7..7cd2babf8 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -21,7 +21,6 @@
 
 import copy
 import os
-import re
 
 import gym
 import numpy as np
@@ -40,42 +39,6 @@
 FLAGS = flags.FLAGS
 
 
-def infer_last_epoch_num(data_dir):
-  """Infer highest epoch number from file names in data_dir."""
-  names = os.listdir(data_dir)
-  epochs_str = [re.findall(pattern=r".*\.(-?\d+)$", string=name)
-                for name in names]
-  epochs_str = sum(epochs_str, [])
-  return max([int(epoch_str) for epoch_str in epochs_str])
-
-
-def setup_and_load_epoch(hparams, data_dir, which_epoch_data=None):
-  """Load T2TGymEnv with data from one epoch.
-
-  Args:
-    hparams: hparams.
-    data_dir: data directory.
-    which_epoch_data: data from which epoch to load.
-
-  Returns:
-    env.
-  """
-  t2t_env = rl_utils.setup_env(
-      hparams, batch_size=hparams.real_batch_size,
-      max_num_noops=hparams.max_num_noops
-  )
-  # Load data.
-  if which_epoch_data is not None:
-    if which_epoch_data == "last":
-      which_epoch_data = infer_last_epoch_num(data_dir)
-    assert isinstance(which_epoch_data, int), \
-      "{}".format(type(which_epoch_data))
-    t2t_env.start_new_epoch(which_epoch_data, data_dir)
-  else:
-    t2t_env.start_new_epoch(-999)
-  return t2t_env
-
-
 def make_simulated_gym_env(real_env, world_model_dir, hparams, random_starts):
   """Gym environment with world model."""
   initial_frame_chooser = rl_utils.make_initial_frame_chooser(
@@ -98,7 +61,7 @@ def load_data_and_make_simulated_env(
     data_dir, wm_dir, hparams, which_epoch_data="last", random_starts=True
 ):
   hparams = copy.deepcopy(hparams)
-  t2t_env = setup_and_load_epoch(
+  t2t_env = T2TGymEnv.setup_and_load_epoch(
       hparams, data_dir=data_dir,
       which_epoch_data=which_epoch_data)
   return make_simulated_gym_env(
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 876ccb78e..2a9726eaa 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -28,7 +28,6 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl.dopamine_connector import DQNLearner
 from tensor2tensor.rl.ppo_learner import PPOLearner
-from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -64,9 +63,8 @@ def evaluate_single_config(
 ):
   """Evaluate the PPO agent in the real environment."""
   eval_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
-  env = setup_env(
-      hparams, batch_size=hparams.eval_batch_size, max_num_noops=max_num_noops,
-      rl_env_max_episode_steps=hparams.eval_rl_env_max_episode_steps
+  env = T2TGymEnv.setup_env_from_hparams(
+      hparams, batch_size=hparams.eval_batch_size, max_num_noops=max_num_noops
   )
   env.start_new_epoch(0)
   env_fn = rl.make_real_env_fn(env)
@@ -100,38 +98,12 @@ def evaluate_all_configs(hparams, agent_model_dir):
   return metrics
 
 
-def summarize_metrics(eval_metrics_writer, metrics, epoch):
-  """Write metrics to summary."""
-  for (name, value) in six.iteritems(metrics):
-    summary = tf.Summary()
-    summary.value.add(tag=name, simple_value=value)
-    eval_metrics_writer.add_summary(summary, epoch)
-  eval_metrics_writer.flush()
-
-
 LEARNERS = {
     "ppo": PPOLearner,
     "dqn": DQNLearner,
 }
 
 
-def setup_env(hparams, batch_size, max_num_noops, rl_env_max_episode_steps=-1):
-  """Setup."""
-  game_mode = "NoFrameskip-v4"
-  camel_game_name = misc_utils.snakecase_to_camelcase(hparams.game)
-  camel_game_name += game_mode
-  env_name = camel_game_name
-
-  env = T2TGymEnv(base_env_name=env_name,
-                  batch_size=batch_size,
-                  grayscale=hparams.grayscale,
-                  resize_width_factor=hparams.resize_width_factor,
-                  resize_height_factor=hparams.resize_height_factor,
-                  rl_env_max_episode_steps=rl_env_max_episode_steps,
-                  max_num_noops=max_num_noops, maxskip_envs=True)
-  return env
-
-
 def update_hparams_from_hparams(target_hparams, source_hparams, prefix):
   """Copy a subset of hparams to target_hparams."""
   for (param_name, param_value) in six.iteritems(source_hparams.values()):
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 7a55aaad0..e5daf6f34 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -37,6 +37,7 @@
 import six
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
+from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research.rl import make_simulated_env_fn_from_hparams
@@ -377,6 +378,15 @@ def load_metrics(event_dir, epoch):
   return metrics
 
 
+def summarize_metrics(eval_metrics_writer, metrics, epoch):
+  """Write metrics to summary."""
+  for (name, value) in six.iteritems(metrics):
+    summary = tf.Summary()
+    summary.value.add(tag=name, simple_value=value)
+    eval_metrics_writer.add_summary(summary, epoch)
+  eval_metrics_writer.flush()
+
+
 def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   """Run the main training loop."""
   if report_fn:
@@ -391,10 +401,9 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   epoch = -1
   data_dir = directories["data"]
-  env = rl_utils.setup_env(
+  env = T2TGymEnv.setup_env_from_hparams(
       hparams, batch_size=hparams.real_batch_size,
-      max_num_noops=hparams.max_num_noops,
-      rl_env_max_episode_steps=hparams.rl_env_max_episode_steps
+      max_num_noops=hparams.max_num_noops
   )
   env.start_new_epoch(epoch, data_dir)
 
@@ -484,7 +493,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
         log("World model eval metrics:\n{}".format(pprint.pformat(wm_metrics)))
         metrics.update(wm_metrics)
 
-      rl_utils.summarize_metrics(eval_metrics_writer, metrics, epoch)
+      summarize_metrics(eval_metrics_writer, metrics, epoch)
 
       # Report metrics
       if report_fn:
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 463eac8cf..610d979a8 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -84,8 +84,6 @@ def _rlmb_base():
       # Sampling temperatures to try during eval.
       eval_sampling_temps=[0.5, 0.0, 1.0],
       eval_max_num_noops=8,
-      # To speed up the pipeline. Some games want to run forever.
-      eval_rl_env_max_episode_steps=1000,
 
       game="pong",
       # Whether to evaluate the world model in each iteration of the loop to get
@@ -508,7 +506,6 @@ def _rlmb_tiny_overrides():
       resize_width_factor=2,
       wm_eval_rollout_ratios=[1],
       rl_env_max_episode_steps=7,
-      eval_rl_env_max_episode_steps=7,
       simulated_rollout_length=2,
       eval_sampling_temps=[0.0, 1.0],
   )
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 1d8d53162..64302dfdc 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -20,7 +20,7 @@
 python -m tensor2tensor.rl.trainer_model_free \
     --output_dir=$HOME/t2t/rl_v1 \
     --hparams_set=pong_model_free \
-    --hparams='batch_size=15'
+    --loop_hparams='batch_size=15'
 """
 
 from __future__ import absolute_import
@@ -29,6 +29,7 @@
 
 import pprint
 
+from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
@@ -52,9 +53,9 @@
 
 def initialize_env_specs(hparams):
   """Initializes env_specs using T2TGymEnvs."""
-  env = rl_utils.setup_env(hparams, hparams.batch_size,
-                           hparams.eval_max_num_noops,
-                           hparams.rl_env_max_episode_steps)
+  env = T2TGymEnv.setup_env_from_hparams(
+      hparams, hparams.batch_size, hparams.eval_max_num_noops
+  )
   env.start_new_epoch(0)
 
   # TODO(afrozm): Decouple env_fn from hparams and return both, is there

From 2fdd517f8d05fc4810b3c21db64b8b854f0b7fc9 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 8 Jan 2019 15:00:39 -0800
Subject: [PATCH 1462/2720] internal merge of PR #1348

PiperOrigin-RevId: 228405946
---
 tensor2tensor/data_generators/gym_env.py      | 54 ------------
 tensor2tensor/models/research/rl.py           |  1 +
 tensor2tensor/rl/evaluator.py                 | 83 +++++++++++++++++++
 tensor2tensor/rl/player.py                    |  3 +-
 tensor2tensor/rl/player_utils.py              | 40 ++++++++-
 tensor2tensor/rl/rl_utils.py                  | 32 ++++++-
 tensor2tensor/rl/trainer_model_based.py       | 17 +---
 .../rl/trainer_model_based_params.py          |  3 +
 tensor2tensor/rl/trainer_model_free.py        |  9 +-
 9 files changed, 164 insertions(+), 78 deletions(-)
 create mode 100644 tensor2tensor/rl/evaluator.py

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index a777b5aac..6b2a1c65e 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -21,9 +21,7 @@
 
 import collections
 import itertools
-import os
 import random
-import re
 
 from gym.spaces import Box
 import numpy as np
@@ -642,58 +640,6 @@ def base_env_name(self):
   def num_channels(self):
     return self.observation_space.shape[2]
 
-  @staticmethod
-  def infer_last_epoch_num(data_dir):
-    """Infer highest epoch number from file names in data_dir."""
-    names = os.listdir(data_dir)
-    epochs_str = [re.findall(pattern=r".*\.(-?\d+)$", string=name)
-                  for name in names]
-    epochs_str = sum(epochs_str, [])
-    return max([int(epoch_str) for epoch_str in epochs_str])
-
-  @staticmethod
-  def setup_env_from_hparams(hparams, batch_size, max_num_noops):
-    game_mode = "NoFrameskip-v4"
-    camel_game_name = misc_utils.snakecase_to_camelcase(hparams.game)
-    camel_game_name += game_mode
-    env_name = camel_game_name
-
-    env = T2TGymEnv(base_env_name=env_name,
-                    batch_size=batch_size,
-                    grayscale=hparams.grayscale,
-                    resize_width_factor=hparams.resize_width_factor,
-                    resize_height_factor=hparams.resize_height_factor,
-                    rl_env_max_episode_steps=hparams.rl_env_max_episode_steps,
-                    max_num_noops=max_num_noops, maxskip_envs=True)
-    return env
-
-  @staticmethod
-  def setup_and_load_epoch(hparams, data_dir, which_epoch_data=None):
-    """Load T2TBatchGymEnv with data from one epoch.
-
-    Args:
-      hparams: hparams.
-      data_dir: data directory.
-      which_epoch_data: data from which epoch to load.
-
-    Returns:
-      env.
-    """
-    t2t_env = T2TGymEnv.setup_env_from_hparams(
-        hparams, batch_size=hparams.real_batch_size,
-        max_num_noops=hparams.max_num_noops
-    )
-    # Load data.
-    if which_epoch_data is not None:
-      if which_epoch_data == "last":
-        which_epoch_data = T2TGymEnv.infer_last_epoch_num(data_dir)
-      assert isinstance(which_epoch_data, int), \
-        "{}".format(type(which_epoch_data))
-      t2t_env.start_new_epoch(which_epoch_data, data_dir)
-    else:
-      t2t_env.start_new_epoch(-999)
-    return t2t_env
-
   def _derive_observation_space(self, orig_observ_space):
     height, width, channels = orig_observ_space.shape
     if self.grayscale:
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 87acbaa58..ff2b0d72d 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -335,6 +335,7 @@ def rlmf_original():
       frame_stack_size=4,
       eval_sampling_temps=[0.0, 0.2, 0.5, 0.8, 1.0, 2.0],
       eval_max_num_noops=8,
+      eval_rl_env_max_episode_steps=1000,
       resize_height_factor=2,
       resize_width_factor=2,
       grayscale=0,
diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
new file mode 100644
index 000000000..c631b43b2
--- /dev/null
+++ b/tensor2tensor/rl/evaluator.py
@@ -0,0 +1,83 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+r"""Evaluation script for RL agents.
+
+Example invocation:
+
+python -m tensor2tensor.rl.evaluator \
+    --policy_dir=$HOME/t2t/rl_v1/policy \
+    --eval_metrics_dir=$HOME/t2t/rl_v1/full_eval_metrics \
+    --hparams_set=rlmb_base \
+    --hparams='batch_size=64'
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.models.research import rl  # pylint: disable=unused-import
+from tensor2tensor.rl import rl_utils
+from tensor2tensor.rl import trainer_model_based_params  # pylint: disable=unused-import
+from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
+from tensor2tensor.utils import trainer_lib
+
+import tensorflow as tf
+
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+
+flags.DEFINE_string("policy_dir", "", "Directory with policy checkpoints.")
+flags.DEFINE_string(
+    "eval_metrics_dir", "", "Directory to output the eval metrics at."
+)
+flags.DEFINE_bool("full_eval", True, "Whether to ignore the timestep limit.")
+
+
+def evaluate(hparams, policy_dir, eval_metrics_dir, report_fn=None,
+             report_metric=None):
+  """Evaluate."""
+  if report_fn:
+    assert report_metric is not None
+
+  eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
+  eval_metrics = rl_utils.evaluate_all_configs(hparams, policy_dir)
+  rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, 0)
+
+  # Report metrics
+  if report_fn:
+    if report_metric == "mean_reward":
+      metric_name = rl_utils.get_metric_name(
+          sampling_temp=hparams.eval_sampling_temps[0],
+          max_num_noops=hparams.eval_max_num_noops,
+          clipped=False
+      )
+      report_fn(eval_metrics[metric_name], 0)
+    else:
+      report_fn(eval_metrics[report_metric], 0)
+  return eval_metrics
+
+
+def main(_):
+  hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
+  if FLAGS.full_eval:
+    hparams.eval_rl_env_max_episode_steps = -1
+  evaluate(hparams, FLAGS.policy_dir, FLAGS.eval_metrics_dir)
+
+
+if __name__ == "__main__":
+  tf.app.run()
diff --git a/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 9ab0a7023..0dc680093 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -56,7 +56,6 @@
 import six
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
-from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.rl import player_utils
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw
@@ -229,7 +228,7 @@ def main(_):
         directories["data"], directories["world_model"],
         hparams, which_epoch_data=epoch)
   else:
-    env = T2TGymEnv.setup_and_load_epoch(
+    env = player_utils.setup_and_load_epoch(
         hparams, data_dir=directories["data"],
         which_epoch_data=epoch)
     env = FlatBatchEnv(env)
diff --git a/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index 7cd2babf8..1025816f6 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -21,12 +21,12 @@
 
 import copy
 import os
+import re
 
 import gym
 import numpy as np
 import six
 
-from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.models.research.rl import get_policy
 from tensor2tensor.models.research.rl import make_simulated_env_fn_from_hparams
 from tensor2tensor.rl import rl_utils
@@ -39,6 +39,42 @@
 FLAGS = flags.FLAGS
 
 
+def infer_last_epoch_num(data_dir):
+  """Infer highest epoch number from file names in data_dir."""
+  names = os.listdir(data_dir)
+  epochs_str = [re.findall(pattern=r".*\.(-?\d+)$", string=name)
+                for name in names]
+  epochs_str = sum(epochs_str, [])
+  return max([int(epoch_str) for epoch_str in epochs_str])
+
+
+def setup_and_load_epoch(hparams, data_dir, which_epoch_data=None):
+  """Load T2TGymEnv with data from one epoch.
+
+  Args:
+    hparams: hparams.
+    data_dir: data directory.
+    which_epoch_data: data from which epoch to load.
+
+  Returns:
+    env.
+  """
+  t2t_env = rl_utils.setup_env(
+      hparams, batch_size=hparams.real_batch_size,
+      max_num_noops=hparams.max_num_noops
+  )
+  # Load data.
+  if which_epoch_data is not None:
+    if which_epoch_data == "last":
+      which_epoch_data = infer_last_epoch_num(data_dir)
+    assert isinstance(which_epoch_data, int), \
+      "{}".format(type(which_epoch_data))
+    t2t_env.start_new_epoch(which_epoch_data, data_dir)
+  else:
+    t2t_env.start_new_epoch(-999)
+  return t2t_env
+
+
 def make_simulated_gym_env(real_env, world_model_dir, hparams, random_starts):
   """Gym environment with world model."""
   initial_frame_chooser = rl_utils.make_initial_frame_chooser(
@@ -61,7 +97,7 @@ def load_data_and_make_simulated_env(
     data_dir, wm_dir, hparams, which_epoch_data="last", random_starts=True
 ):
   hparams = copy.deepcopy(hparams)
-  t2t_env = T2TGymEnv.setup_and_load_epoch(
+  t2t_env = setup_and_load_epoch(
       hparams, data_dir=data_dir,
       which_epoch_data=which_epoch_data)
   return make_simulated_gym_env(
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 2a9726eaa..876ccb78e 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -28,6 +28,7 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl.dopamine_connector import DQNLearner
 from tensor2tensor.rl.ppo_learner import PPOLearner
+from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -63,8 +64,9 @@ def evaluate_single_config(
 ):
   """Evaluate the PPO agent in the real environment."""
   eval_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
-  env = T2TGymEnv.setup_env_from_hparams(
-      hparams, batch_size=hparams.eval_batch_size, max_num_noops=max_num_noops
+  env = setup_env(
+      hparams, batch_size=hparams.eval_batch_size, max_num_noops=max_num_noops,
+      rl_env_max_episode_steps=hparams.eval_rl_env_max_episode_steps
   )
   env.start_new_epoch(0)
   env_fn = rl.make_real_env_fn(env)
@@ -98,12 +100,38 @@ def evaluate_all_configs(hparams, agent_model_dir):
   return metrics
 
 
+def summarize_metrics(eval_metrics_writer, metrics, epoch):
+  """Write metrics to summary."""
+  for (name, value) in six.iteritems(metrics):
+    summary = tf.Summary()
+    summary.value.add(tag=name, simple_value=value)
+    eval_metrics_writer.add_summary(summary, epoch)
+  eval_metrics_writer.flush()
+
+
 LEARNERS = {
     "ppo": PPOLearner,
     "dqn": DQNLearner,
 }
 
 
+def setup_env(hparams, batch_size, max_num_noops, rl_env_max_episode_steps=-1):
+  """Setup."""
+  game_mode = "NoFrameskip-v4"
+  camel_game_name = misc_utils.snakecase_to_camelcase(hparams.game)
+  camel_game_name += game_mode
+  env_name = camel_game_name
+
+  env = T2TGymEnv(base_env_name=env_name,
+                  batch_size=batch_size,
+                  grayscale=hparams.grayscale,
+                  resize_width_factor=hparams.resize_width_factor,
+                  resize_height_factor=hparams.resize_height_factor,
+                  rl_env_max_episode_steps=rl_env_max_episode_steps,
+                  max_num_noops=max_num_noops, maxskip_envs=True)
+  return env
+
+
 def update_hparams_from_hparams(target_hparams, source_hparams, prefix):
   """Copy a subset of hparams to target_hparams."""
   for (param_name, param_value) in six.iteritems(source_hparams.values()):
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index e5daf6f34..7a55aaad0 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -37,7 +37,6 @@
 import six
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
-from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research.rl import make_simulated_env_fn_from_hparams
@@ -378,15 +377,6 @@ def load_metrics(event_dir, epoch):
   return metrics
 
 
-def summarize_metrics(eval_metrics_writer, metrics, epoch):
-  """Write metrics to summary."""
-  for (name, value) in six.iteritems(metrics):
-    summary = tf.Summary()
-    summary.value.add(tag=name, simple_value=value)
-    eval_metrics_writer.add_summary(summary, epoch)
-  eval_metrics_writer.flush()
-
-
 def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   """Run the main training loop."""
   if report_fn:
@@ -401,9 +391,10 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   epoch = -1
   data_dir = directories["data"]
-  env = T2TGymEnv.setup_env_from_hparams(
+  env = rl_utils.setup_env(
       hparams, batch_size=hparams.real_batch_size,
-      max_num_noops=hparams.max_num_noops
+      max_num_noops=hparams.max_num_noops,
+      rl_env_max_episode_steps=hparams.rl_env_max_episode_steps
   )
   env.start_new_epoch(epoch, data_dir)
 
@@ -493,7 +484,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
         log("World model eval metrics:\n{}".format(pprint.pformat(wm_metrics)))
         metrics.update(wm_metrics)
 
-      summarize_metrics(eval_metrics_writer, metrics, epoch)
+      rl_utils.summarize_metrics(eval_metrics_writer, metrics, epoch)
 
       # Report metrics
       if report_fn:
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 610d979a8..463eac8cf 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -84,6 +84,8 @@ def _rlmb_base():
       # Sampling temperatures to try during eval.
       eval_sampling_temps=[0.5, 0.0, 1.0],
       eval_max_num_noops=8,
+      # To speed up the pipeline. Some games want to run forever.
+      eval_rl_env_max_episode_steps=1000,
 
       game="pong",
       # Whether to evaluate the world model in each iteration of the loop to get
@@ -506,6 +508,7 @@ def _rlmb_tiny_overrides():
       resize_width_factor=2,
       wm_eval_rollout_ratios=[1],
       rl_env_max_episode_steps=7,
+      eval_rl_env_max_episode_steps=7,
       simulated_rollout_length=2,
       eval_sampling_temps=[0.0, 1.0],
   )
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 64302dfdc..1d8d53162 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -20,7 +20,7 @@
 python -m tensor2tensor.rl.trainer_model_free \
     --output_dir=$HOME/t2t/rl_v1 \
     --hparams_set=pong_model_free \
-    --loop_hparams='batch_size=15'
+    --hparams='batch_size=15'
 """
 
 from __future__ import absolute_import
@@ -29,7 +29,6 @@
 
 import pprint
 
-from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
@@ -53,9 +52,9 @@
 
 def initialize_env_specs(hparams):
   """Initializes env_specs using T2TGymEnvs."""
-  env = T2TGymEnv.setup_env_from_hparams(
-      hparams, hparams.batch_size, hparams.eval_max_num_noops
-  )
+  env = rl_utils.setup_env(hparams, hparams.batch_size,
+                           hparams.eval_max_num_noops,
+                           hparams.rl_env_max_episode_steps)
   env.start_new_epoch(0)
 
   # TODO(afrozm): Decouple env_fn from hparams and return both, is there

From 18c08f321c51f931e9fc6bac41018a3003e35e35 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 8 Jan 2019 17:41:22 -0800
Subject: [PATCH 1463/2720] Correct typo affecting multi-problem training.

PiperOrigin-RevId: 228432054
---
 tensor2tensor/utils/t2t_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 4578c4d69..8688611e9 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -617,7 +617,7 @@ def weights_fn(labels):
 
     if hasattr(self.hparams, "problem") and hasattr(
         self.hparams.problem, "task_list"):
-      if weights_fn is not None:
+      if weights is not None:
         raise NotImplementedError("weights not yet implemented in "
                                   "multitask setting.")
       loss_num, loss_den, summaries = multi_problem.aggregate_task_losses(

From b643525233e5c6c3a4ad93b56358c347b9c87e73 Mon Sep 17 00:00:00 2001
From: Alexander Ku <alexku@google.com>
Date: Tue, 8 Jan 2019 17:41:41 -0800
Subject: [PATCH 1464/2720] SARI score for evaluating paraphrasing and other
 text generation models.

PiperOrigin-RevId: 228432095
---
 tensor2tensor/utils/metrics.py        |   3 +
 tensor2tensor/utils/sari_hook.py      | 252 ++++++++++++++++++++++++++
 tensor2tensor/utils/sari_hook_test.py | 147 +++++++++++++++
 3 files changed, 402 insertions(+)
 create mode 100644 tensor2tensor/utils/sari_hook.py
 create mode 100644 tensor2tensor/utils/sari_hook_test.py

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index feea42058..94d9f64d2 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -24,6 +24,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import bleu_hook
 from tensor2tensor.utils import rouge
+from tensor2tensor.utils import sari_hook
 
 import tensorflow as tf
 
@@ -41,6 +42,7 @@ class Metrics(object):
   NEG_LOG_PERPLEXITY = "neg_log_perplexity"
   MASKED_NEG_LOG_PERPLEXITY = "masked_neg_log_perplexity"
   APPROX_BLEU = "approx_bleu_score"
+  APPROX_SARI = "approx_sari_score"
   RMSE = "rmse"
   LOG_POISSON = "log_poisson"
   PEARSON = "pearson"
@@ -782,6 +784,7 @@ def pearson_correlation_coefficient(predictions, labels, weights_fn=None):
     Metrics.NEG_LOG_PERPLEXITY: padded_neg_log_perplexity,
     Metrics.MASKED_NEG_LOG_PERPLEXITY: padded_neg_log_perplexity_with_masking,
     Metrics.APPROX_BLEU: bleu_hook.bleu_score,
+    Metrics.APPROX_SARI: sari_hook.sari_score,
     Metrics.RMSE: padded_rmse,
     Metrics.LOG_POISSON: padded_log_poisson,
     Metrics.PEARSON: pearson_correlation_coefficient,
diff --git a/tensor2tensor/utils/sari_hook.py b/tensor2tensor/utils/sari_hook.py
new file mode 100644
index 000000000..5e0df0056
--- /dev/null
+++ b/tensor2tensor/utils/sari_hook.py
@@ -0,0 +1,252 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""SARI score for evaluating paraphrasing and other text generation models.
+
+The score is introduced in the following paper:
+
+   Optimizing Statistical Machine Translation for Text Simplification
+   Wei Xu, Courtney Napoles, Ellie Pavlick, Quanze Chen and Chris Callison-Burch
+   In Transactions of the Association for Computational Linguistics (TACL) 2015
+   http://cs.jhu.edu/~napoles/res/tacl2016-optimizing.pdf
+
+This implementation has two differences with the GitHub [1] implementation:
+  (1) Define 0/0=1 instead of 0 to give higher scores for predictions that match
+      a target exactly.
+  (2) Fix an alleged bug [2] in the deletion score computation.
+
+[1] https://github.com/cocoxu/simplification/blob/master/SARI.py
+    (commit 0210f15)
+[2] https://github.com/cocoxu/simplification/issues/6
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+
+import numpy as np
+import tensorflow as tf
+
+# The paper that intoduces the SARI score uses only the precision of the deleted
+# tokens (i.e. beta=0). To give more emphasis on recall, you may set, e.g.,
+# beta=1.
+BETA_FOR_SARI_DELETION_F_MEASURE = 0
+
+
+def _get_ngram_counter(ids, n):
+  """Get a Counter with the ngrams of the given ID list.
+
+  Args:
+    ids: np.array or a list corresponding to a single sentence
+    n: n-gram size
+
+  Returns:
+    collections.Counter with ID tuples as keys and 1s as values.
+  """
+  # Remove zero IDs used to pad the sequence.
+  ids = [token_id for token_id in ids if token_id != 0]
+  ngram_list = [tuple(ids[i:i + n]) for i in range(len(ids) + 1 - n)]
+  ngrams = set(ngram_list)
+  counts = collections.Counter()
+  for ngram in ngrams:
+    counts[ngram] = 1
+  return counts
+
+
+def _get_fbeta_score(true_positives, selected, relevant, beta=1):
+  """Compute Fbeta score.
+
+  Args:
+    true_positives: Number of true positive ngrams.
+    selected: Number of selected ngrams.
+    relevant: Number of relevant ngrams.
+    beta: 0 gives precision only, 1 gives F1 score, and Inf gives recall only.
+
+  Returns:
+    Fbeta score.
+  """
+  precision = 1
+  if selected > 0:
+    precision = true_positives / selected
+  if beta == 0:
+    return precision
+  recall = 1
+  if relevant > 0:
+    recall = true_positives / relevant
+  if precision > 0 and recall > 0:
+    beta2 = beta * beta
+    return (1 + beta2) * precision * recall / (beta2 * precision + recall)
+  else:
+    return 0
+
+
+def get_addition_score(source_counts, prediction_counts, target_counts):
+  """Compute the addition score (Equation 4 in the paper)."""
+  added_to_prediction_counts = prediction_counts - source_counts
+  true_positives = sum((added_to_prediction_counts & target_counts).values())
+  selected = sum(added_to_prediction_counts.values())
+  # Note that in the paper the summation is done over all the ngrams in the
+  # output rather than the ngrams in the following set difference. Since the
+  # former does not make as much sense we compute the latter, which is also done
+  # in the GitHub implementation.
+  relevant = sum((target_counts - source_counts).values())
+  return _get_fbeta_score(true_positives, selected, relevant)
+
+
+def get_keep_score(source_counts, prediction_counts, target_counts):
+  """Compute the keep score (Equation 5 in the paper)."""
+  source_and_prediction_counts = source_counts & prediction_counts
+  source_and_target_counts = source_counts & target_counts
+  true_positives = sum((source_and_prediction_counts &
+                        source_and_target_counts).values())
+  selected = sum(source_and_prediction_counts.values())
+  relevant = sum(source_and_target_counts.values())
+  return _get_fbeta_score(true_positives, selected, relevant)
+
+
+def get_deletion_score(source_counts, prediction_counts, target_counts, beta=0):
+  """Compute the deletion score (Equation 6 in the paper)."""
+  source_not_prediction_counts = source_counts - prediction_counts
+  source_not_target_counts = source_counts - target_counts
+  true_positives = sum((source_not_prediction_counts &
+                        source_not_target_counts).values())
+  selected = sum(source_not_prediction_counts.values())
+  relevant = sum(source_not_target_counts.values())
+  return _get_fbeta_score(true_positives, selected, relevant, beta=beta)
+
+
+def get_sari_score(source_ids, prediction_ids, list_of_targets,
+                   max_gram_size=4, beta_for_deletion=0):
+  """Compute the SARI score for a single prediction and one or more targets.
+
+  Args:
+    source_ids: a list / np.array of SentencePiece IDs
+    prediction_ids: a list / np.array of SentencePiece IDs
+    list_of_targets: a list of target ID lists / np.arrays
+    max_gram_size: int. largest n-gram size we care about (e.g. 3 for unigrams,
+        bigrams, and trigrams)
+    beta_for_deletion: beta for deletion F score.
+
+  Returns:
+    the SARI score and its three components: add, keep, and deletion scores
+  """
+  addition_scores = []
+  keep_scores = []
+  deletion_scores = []
+  for n in range(1, max_gram_size + 1):
+    source_counts = _get_ngram_counter(source_ids, n)
+    prediction_counts = _get_ngram_counter(prediction_ids, n)
+    # All ngrams in the targets with count 1.
+    target_counts = collections.Counter()
+    # All ngrams in the targets with count r/num_targets, where r is the number
+    # of targets where the ngram occurs.
+    weighted_target_counts = collections.Counter()
+    num_nonempty_targets = 0
+    for target_ids_i in list_of_targets:
+      target_counts_i = _get_ngram_counter(target_ids_i, n)
+      if target_counts_i:
+        weighted_target_counts += target_counts_i
+        num_nonempty_targets += 1
+    for gram in weighted_target_counts.keys():
+      weighted_target_counts[gram] /= num_nonempty_targets
+      target_counts[gram] = 1
+    keep_scores.append(get_keep_score(source_counts, prediction_counts,
+                                      weighted_target_counts))
+    deletion_scores.append(get_deletion_score(source_counts, prediction_counts,
+                                              weighted_target_counts,
+                                              beta_for_deletion))
+    addition_scores.append(get_addition_score(source_counts, prediction_counts,
+                                              target_counts))
+
+  avg_keep_score = sum(keep_scores) / max_gram_size
+  avg_addition_score = sum(addition_scores) / max_gram_size
+  avg_deletion_score = sum(deletion_scores) / max_gram_size
+  sari = (avg_keep_score + avg_addition_score + avg_deletion_score) / 3.0
+  return sari, avg_keep_score, avg_addition_score, avg_deletion_score
+
+
+def get_sari(source_ids, prediction_ids, target_ids, max_gram_size=4):
+  """Computes the SARI scores from the given source, prediction and targets.
+
+  Args:
+    source_ids: A 2D tf.Tensor of size (batch_size , sequence_length)
+    prediction_ids: A 2D tf.Tensor of size (batch_size, sequence_length)
+    target_ids: A 3D tf.Tensor of size (batch_size, number_of_targets,
+        sequence_length)
+    max_gram_size: int. largest n-gram size we care about (e.g. 3 for unigrams,
+        bigrams, and trigrams)
+
+  Returns:
+    A 4-tuple of 1D float Tensors of size (batch_size) for the SARI score and
+        the keep, addition and deletion scores.
+  """
+
+  def get_sari_numpy(source_ids, prediction_ids, target_ids):
+    """Iterate over elements in the batch and call the SARI function."""
+    sari_scores = []
+    keep_scores = []
+    add_scores = []
+    deletion_scores = []
+    # Iterate over elements in the batch.
+    for source_ids_i, prediction_ids_i, target_ids_i in zip(
+        source_ids, prediction_ids, target_ids):
+      sari, keep, add, deletion = get_sari_score(
+          source_ids_i, prediction_ids_i, target_ids_i, max_gram_size,
+          BETA_FOR_SARI_DELETION_F_MEASURE)
+      sari_scores.append(sari)
+      keep_scores.append(keep)
+      add_scores.append(add)
+      deletion_scores.append(deletion)
+    return (np.asarray(sari_scores), np.asarray(keep_scores),
+            np.asarray(add_scores), np.asarray(deletion_scores))
+
+  sari, keep, add, deletion = tf.py_function(
+      get_sari_numpy,
+      [source_ids, prediction_ids, target_ids],
+      [tf.float64, tf.float64, tf.float64, tf.float64])
+  return sari, keep, add, deletion
+
+
+def sari_score(predictions, labels, features, **unused_kwargs):
+  """Computes the SARI scores from the given source, prediction and targets.
+
+  An approximate SARI scoring method since we do not glue word pieces or
+  decode the ids and tokenize the output. By default, we use ngram order of 4.
+  Also, this does not have beam search.
+
+  Args:
+    predictions: tensor, model predictions.
+    labels: tensor, gold output.
+    features: dict, containing inputs.
+
+  Returns:
+    sari: int, approx sari score
+  """
+  if "inputs" not in features:
+    raise ValueError("sari_score requires inputs feature")
+
+  # Convert the inputs and outputs to a [batch_size, sequence_length] tensor.
+  inputs = tf.squeeze(features["inputs"], axis=[-1, -2])
+  outputs = tf.to_int32(tf.argmax(predictions, axis=-1))
+  outputs = tf.squeeze(outputs, axis=-1)
+
+  # Convert the labels to a [batch_size, 1, sequence_length] tensor.
+  labels = tf.squeeze(labels, axis=[-1, -2])
+  labels = tf.expand_dims(labels, axis=1)
+
+  score, _, _, _ = get_sari(inputs, outputs, labels)
+  return score, tf.constant(1.0)
diff --git a/tensor2tensor/utils/sari_hook_test.py b/tensor2tensor/utils/sari_hook_test.py
new file mode 100644
index 000000000..1d039847d
--- /dev/null
+++ b/tensor2tensor/utils/sari_hook_test.py
@@ -0,0 +1,147 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.utils.sari_hook."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+
+import numpy as np
+from tensor2tensor.utils import sari_hook
+import tensorflow as tf
+
+
+class SariHookTest(tf.test.TestCase):
+
+  def setUp(self):
+    """Sets up inputs and references from the paper's test cases."""
+    self.input_sentence = (
+        "About 95 species are currently accepted .".split())
+    self.references = [
+        "About 95 species are currently known .".split(),
+        "About 95 species are now accepted .".split(),
+        "95 species are now accepted .".split(),
+    ]
+
+  def testSariSent1(self):
+    """Test case 1 from SARI-paper.
+
+    The score is slightly different from what is reported in the paper (0.2683)
+    since the authors' code seems to contain a bug in the keep recall score
+    computation.
+    """
+    output = "About 95 you now get in ." .split()
+    score, _, _, _ = sari_hook.get_sari_score(self.input_sentence, output,
+                                              self.references)
+    self.assertAlmostEqual(0.2695360, score)
+
+  def testSariSent2(self):
+    """Test case 2 from SARI-paper."""
+    output = "About 95 species are now agreed .".split()
+    score, _, _, _ = sari_hook.get_sari_score(self.input_sentence, output,
+                                              self.references)
+    self.assertAlmostEqual(0.6170966, score)
+
+  def testSariSent3(self):
+    """Test case 3 from SARI-paper."""
+    output = "About 95 species are currently agreed .".split()
+    score, _, _, _ = sari_hook.get_sari_score(self.input_sentence, output,
+                                              self.references)
+    self.assertAlmostEqual(0.5088682, score)
+
+  def testMatchingSentences(self):
+    """If input=output=reference, the score should be 1."""
+    input_sentence = [3, 1, 4, 1, 5, 9, 2, 6, 5]
+    output = input_sentence
+    references = [input_sentence]
+    score, _, _, _ = sari_hook.get_sari_score(input_sentence, output,
+                                              references)
+    self.assertEqual(1, score)
+
+  def testMatchingOutputAndReference(self):
+    """If output=reference, the score should be 1."""
+    input_sentence = [3, 1, 4, 1, 5, 9, 2, 6, 5]
+    output = [3, 1, 4, 1, 80, 70]
+    references = [output]
+    score, _, _, _ = sari_hook.get_sari_score(input_sentence, output,
+                                              references)
+    self.assertEqual(1, score)
+
+  def testMatchingSentencesWithRepetitions(self):
+    """Token frequencies should not matter if we only consider unigrams."""
+    input_sentence = [3, 1, 4]
+    output = [3, 3, 1, 1, 1, 4]
+    references = [[3, 3, 3, 1, 1, 4, 4]]
+    score, _, _, _ = sari_hook.get_sari_score(input_sentence, output,
+                                              references, max_gram_size=1)
+    self.assertEqual(1, score)
+
+  def testKeepScore(self):
+    """Toy example where Input='1 2', Output='2', References=['1 2', 1']."""
+    # Unigram counts.
+    source_counts = collections.Counter({1: 1, 2: 1})
+    prediction_counts = collections.Counter({2: 1})
+    target_counts = collections.Counter({1: 1, 2: 0.5})
+    score = sari_hook.get_keep_score(source_counts, prediction_counts,
+                                     target_counts)
+    self.assertAlmostEqual(6.0/15, score)
+
+  def testDeletionScore(self):
+    """Toy example where Input='1 2', Output='1 2', References=['1']."""
+    # Unigram counts.
+    source_counts = collections.Counter({1: 1, 2: 1})
+    prediction_counts = collections.Counter({1: 1, 2: 1})
+    target_counts = collections.Counter({1: 1})
+    # Output doesn't drop any (incorrect) tokens from the input so precision
+    # should be 1, but since '2' is not dropped, recall should be 0. Thus we
+    # should have F1=0 and F0=precision=1.
+    f1_score = sari_hook.get_deletion_score(source_counts, prediction_counts,
+                                            target_counts, beta=1)
+    self.assertEqual(0, f1_score)
+    f0_score = sari_hook.get_deletion_score(source_counts, prediction_counts,
+                                            target_counts, beta=0)
+    self.assertEqual(1, f0_score)
+
+  def testIdsWithZeros(self):
+    """Zeros should be ignored."""
+    input_sentence = [3, 1, 4, 0, 0, 0]
+    output = [3, 1, 4]
+    references = [[3, 1, 4, 0, 0, 0, 0, 0]]
+    score, _, _, _ = sari_hook.get_sari_score(input_sentence, output,
+                                              references)
+    self.assertEqual(1, score)
+
+  def testSariScoreE2E(self):
+    """Tests the SARI metrics end-to-end."""
+    predictions = np.random.randint(4, size=(12, 12, 1, 12))
+    targets = np.random.randint(4, size=(12, 12, 1, 1))
+    inputs = np.random.randint(4, size=(12, 12, 1, 1))
+    with self.test_session() as session:
+      scores, _ = sari_hook.sari_score(
+          predictions=tf.constant(predictions, dtype=tf.int32),
+          labels=tf.constant(targets, dtype=tf.int32),
+          features={
+              "inputs": tf.constant(inputs, dtype=tf.int32),
+          })
+      a = tf.reduce_mean(scores)
+      session.run(tf.global_variables_initializer())
+      session.run(a)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 05aad3a999829af6c8da14ff4f43d6e728e6f161 Mon Sep 17 00:00:00 2001
From: David Dohan <ddohan@google.com>
Date: Tue, 8 Jan 2019 18:05:14 -0800
Subject: [PATCH 1465/2720] Add warm starting from checkpoint to MTF models

PiperOrigin-RevId: 228434997
---
 tensor2tensor/utils/mtf_model.py | 15 ++++++++++-
 tensor2tensor/utils/t2t_model.py | 43 +++++++++++++++++---------------
 2 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/utils/mtf_model.py b/tensor2tensor/utils/mtf_model.py
index fb829f7fe..a217e3b76 100644
--- a/tensor2tensor/utils/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -160,14 +160,27 @@ def estimator_model_fn(cls,
       else:
         host_call = None
 
+      if hparams.warm_start_from:
+
+        def scaffold_fn():
+          t2t_model.initialize_from_ckpt(
+              ckpt_dir=hparams.warm_start_from, hparams=hparams)
+          return tf.train.Scaffold()
+      else:
+        scaffold_fn = None
+
       t2t_model.remove_summaries()
       return tpu_estimator.TPUEstimatorSpec(
           mode=tf.estimator.ModeKeys.TRAIN,
           loss=tf_loss,
           train_op=train_op,
           host_call=host_call,
-          training_hooks=[restore_hook, saver_hook])
+          training_hooks=[restore_hook, saver_hook],
+          scaffold_fn=scaffold_fn)
     else:
+      if hparams.warm_start_from:
+        t2t_model.initialize_from_ckpt(
+            ckpt_dir=hparams.warm_start_from, hparams=hparams)
       return tf.estimator.EstimatorSpec(
           tf.estimator.ModeKeys.TRAIN, loss=tf_loss, train_op=train_op,
           training_chief_hooks=[restore_hook, saver_hook])
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 8688611e9..8fea747d7 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1491,26 +1491,7 @@ def estimator_model_fn(cls,
         loss, num_async_replicas=num_async_replicas, use_tpu=use_tpu)
 
   def initialize_from_ckpt(self, ckpt_dir):
-    model_dir = self._hparams.get("model_dir", None)
-    already_has_ckpt = (
-        model_dir and tf.train.latest_checkpoint(model_dir) is not None)
-    if already_has_ckpt:
-      return
-
-    log_info("Checkpoint dir: %s", ckpt_dir)
-
-    # TODO(mitchellstern): Add support for partitioned variables?
-    reader = tf.train.load_checkpoint(ckpt_dir)
-    variable_map = {}
-    for var in tf.contrib.framework.get_trainable_variables():
-      var_name = var.name.split(":")[0]
-      if reader.has_tensor(var_name):
-        log_info("Loading variable from checkpoint: %s", var_name)
-        variable_map[var_name] = var
-      else:
-        log_info(
-            "Cannot find variable in checkpoint, skipping: %s", var_name)
-    tf.train.init_from_checkpoint(ckpt_dir, variable_map)
+    return initialize_from_ckpt(ckpt_dir=ckpt_dir, hparams=self._hparams)
 
   def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
     """Constructs `tf.estimator.EstimatorSpec` for TRAIN (training) mode."""
@@ -2105,3 +2086,25 @@ def _create_target_modality(modality_dict):
   # behavior.
   return {k: v for k, v in six.iteritems(modality_dict) if "target" in k
           and k != "targets_segmentation" and k != "targets_position"}
+
+
+def initialize_from_ckpt(ckpt_dir, hparams):
+  """Initialize variables from given directory."""
+  model_dir = hparams.get("model_dir", None)
+  already_has_ckpt = (
+      model_dir and tf.train.latest_checkpoint(model_dir) is not None)
+  if already_has_ckpt:
+    return
+
+  tf.logging.info("Checkpoint dir: %s", ckpt_dir)
+  reader = tf.contrib.framework.load_checkpoint(ckpt_dir)
+  variable_map = {}
+  for var in tf.contrib.framework.get_trainable_variables():
+    var_name = var.name.split(":")[0]
+    if reader.has_tensor(var_name):
+      tf.logging.info("Loading variable from checkpoint: %s", var_name)
+      variable_map[var_name] = var
+    else:
+      tf.logging.info("Cannot find variable in checkpoint, skipping: %s",
+                      var_name)
+  tf.train.init_from_checkpoint(ckpt_dir, variable_map)

From 88f5366ca1dc5c052e0cd39ec55e9302de27513c Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Wed, 9 Jan 2019 19:20:33 +0100
Subject: [PATCH 1466/2720] Quick fix for stochastic discrete param sharing
 (#1354)

---
 tensor2tensor/rl/ppo_learner.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index a68b79425..4077b7cd3 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -134,7 +134,8 @@ def evaluate(self, env_fn, hparams, sampling_temp):
             sampling_temp=sampling_temp,
         )
         model_saver = tf.train.Saver(
-            tf.global_variables(hparams.policy_network + "/.*")
+            tf.global_variables(hparams.policy_network + "/.*") +
+            tf.global_variables("clean_scope.*")  # Needed for the SD model.
         )
 
         with tf.Session() as sess:
@@ -198,6 +199,7 @@ def _run_train(ppo_hparams,
   model_saver = tf.train.Saver(
       tf.global_variables(ppo_hparams.policy_network + "/.*") +
       tf.global_variables("training/" + ppo_hparams.policy_network + "/.*") +
+      tf.global_variables("clean_scope.*") +  # Needed for the SD model.
       tf.global_variables("global_step") +
       tf.global_variables("losses_avg.*") +
       tf.global_variables("train_stats.*")

From e3d82d739fd95c6368a543b4a42653fa8721f463 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 9 Jan 2019 13:06:10 -0800
Subject: [PATCH 1467/2720] Refactor input_fn into a self-standing function in
 data_reader.

PiperOrigin-RevId: 228572907
---
 tensor2tensor/data_generators/problem.py      | 294 +---------------
 .../utils/checkpoint_compatibility_test.py    |   4 +-
 tensor2tensor/utils/data_reader.py            | 318 +++++++++++++++++-
 tensor2tensor/utils/trainer_lib_test.py       |   6 +-
 4 files changed, 337 insertions(+), 285 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 5406ddf73..d510ce8a8 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -20,8 +20,6 @@
 
 import collections
 import copy
-import functools
-import multiprocessing
 import os
 import random
 import six
@@ -163,27 +161,6 @@ def preprocess_example_common(example, hparams, mode):
   return example
 
 
-def _file_num_records_cached(filename):
-  """Return the number of TFRecords in a file."""
-  # Cache the result, as this is expensive to compute
-  if filename in _file_num_records_cache:
-    return _file_num_records_cache[filename]
-  ret = 0
-  for _ in tf.python_io.tf_record_iterator(filename):
-    ret += 1
-  _file_num_records_cache[filename] = ret
-  return ret
-
-
-_file_num_records_cache = {}
-
-
-def cpu_count():
-  """Return the number of available cores."""
-  num_available_cores = multiprocessing.cpu_count()
-  return num_available_cores
-
-
 class Problem(object):
   """Problem base class. Specifies a T2T problem.
 
@@ -846,40 +823,12 @@ def input_fn(self,
       (features_dict<str name, Tensor feature>, Tensor targets)
     """
     partition_id, num_partitions = self._dataset_partition(mode, config)
-
     is_training = mode == tf.estimator.ModeKeys.TRAIN
     if config and config.use_tpu:
       num_threads = 64
     else:
-      num_threads = cpu_count() if is_training else 1
-
-    if config and hasattr(config,
-                          "data_parallelism") and config.data_parallelism:
-      num_shards = config.data_parallelism.n
-    else:
-      num_shards = 1
-
-    max_length = self.max_length(hparams)
-    mlperf_log.transformer_print(
-        key=mlperf_log.INPUT_MAX_LENGTH, value=max_length)
-
-    def tpu_valid_size(example):
-      return data_reader.example_valid_size(example, hparams.min_length,
-                                            max_length)
-
-    def gpu_valid_size(example):
-      drop_long_sequences = is_training or hparams.eval_drop_long_sequences
-      max_validate_length = max_length if drop_long_sequences else 10**9
-      return data_reader.example_valid_size(example, hparams.min_length,
-                                            max_validate_length)
-
-    def define_shapes(example):
-      batch_size = config and config.use_tpu and params["batch_size"]
-      return standardize_shapes(example, batch_size=batch_size)
-
-    # Read and preprocess
+      num_threads = data_reader.cpu_count() if is_training else 1
     data_dir = data_dir or (hasattr(hparams, "data_dir") and hparams.data_dir)
-
     dataset_kwargs = dataset_kwargs or {}
     dataset_kwargs.update({
         "mode": mode,
@@ -889,131 +838,20 @@ def define_shapes(example):
         "partition_id": partition_id,
         "num_partitions": num_partitions,
     })
-
-    dataset = self.dataset(**dataset_kwargs)
-    if (force_repeat or is_training) and not prevent_repeat:
-      # Repeat and skip a random number of records
-      dataset = dataset.repeat()
-
-    if is_training and self.skip_random_fraction_when_training:
-      data_files = tf.contrib.slim.parallel_reader.get_data_files(
-          self.filepattern(data_dir, mode))
-      #  In continuous_train_and_eval when switching between train and
-      #  eval, this input_fn method gets called multiple times and it
-      #  would give you the exact same samples from the last call
-      #  (because the Graph seed is set). So this skip gives you some
-      #  shuffling.
-      dataset = skip_random_fraction(dataset, data_files[0])
-
-    dataset = dataset.map(
-        data_reader.cast_ints_to_int32, num_parallel_calls=num_threads)
-
-    if self.batch_size_means_tokens:
-      batch_size_means_tokens = True
-    else:
-      if _are_shapes_fully_defined(dataset.output_shapes):
-        batch_size_means_tokens = False
-      else:
-        tf.logging.warning(
-            "Shapes are not fully defined. Assuming batch_size means tokens.")
-        batch_size_means_tokens = True
-
-    # Batching
-    if not batch_size_means_tokens:
-      # Batch size means examples per datashard.
-      if config and config.use_tpu:
-        # on TPU, we use params["batch_size"], which specifies the number of
-        # examples across all datashards
-        batch_size = params["batch_size"]
-        dataset = dataset.batch(batch_size, drop_remainder=True)
-      else:
-        batch_size = hparams.batch_size * num_shards
-        dataset = dataset.batch(batch_size)
-    else:
-      # batch_size means tokens per datashard
-      if config and config.use_tpu:
-        dataset = dataset.filter(tpu_valid_size)
-        padded_shapes = self._pad_for_tpu(dataset.output_shapes, hparams)
-        # on TPU, we use params["batch_size"], which specifies the number of
-        # examples across all datashards
-        batch_size = params["batch_size"]
-        if hparams.pad_batch:
-          tf.logging.warn(
-              "Padding the batch to ensure that remainder eval batches are "
-              "processed. This may lead to incorrect metrics for "
-              "non-zero-padded features, e.g. images. Use a smaller batch "
-              "size that has no remainder in that case.")
-          dataset = dataset.padded_batch(
-              batch_size, padded_shapes, drop_remainder=False)
-          dataset = dataset.map(
-              functools.partial(pad_batch, batch_multiple=batch_size),
-              num_parallel_calls=num_threads)
-        else:
-          dataset = dataset.padded_batch(
-              batch_size, padded_shapes, drop_remainder=True)
-      else:
-        # On GPU, bucket by length
-        dataset = dataset.filter(gpu_valid_size)
-        batching_scheme = data_reader.hparams_to_batching_scheme(
-            hparams,
-            shard_multiplier=num_shards,
-            length_multiplier=self.get_hparams().batch_size_multiplier)
-        if hparams.use_fixed_batch_size:
-          # Here  batch_size really means examples per datashard.
-          batching_scheme["batch_sizes"] = [hparams.batch_size]
-          batching_scheme["boundaries"] = []
-        dataset = dataset.apply(
-            tf.data.experimental.bucket_by_sequence_length(
-                data_reader.example_length, batching_scheme["boundaries"],
-                batching_scheme["batch_sizes"]))
-
-        if not is_training:
-          batch_multiple = num_shards
-          if hparams.use_fixed_batch_size:
-            # Make sure the last batch has the same fixed size as the rest.
-            batch_multiple *= hparams.batch_size
-          if batch_multiple > 1:
-            tf.logging.warn(
-                "Padding the batch to ensure that remainder eval batches have "
-                "a batch size divisible by the number of data shards. This may "
-                "lead to incorrect metrics for non-zero-padded features, e.g. "
-                "images. Use a single datashard (i.e. 1 GPU) in that case.")
-            dataset = dataset.map(
-                functools.partial(pad_batch, batch_multiple=batch_multiple),
-                num_parallel_calls=num_threads)
-
-    dataset = dataset.map(define_shapes, num_parallel_calls=num_threads)
-
-    # Add shuffling for training batches. This is necessary along with record
-    # level shuffling in the dataset generation. Record shuffling will shuffle
-    # the examples. However, in some cases, it's possible that the shuffle
-    # buffer size for record shuffling is smaller than the batch size. In such
-    # cases, adding batch shuffling ensures that the data is in random order
-    # during training
-    if (is_training and hasattr(hparams, "batch_shuffle_size") and
-        hparams.batch_shuffle_size):
-      dataset = dataset.shuffle(hparams.batch_shuffle_size)
-
-    def prepare_for_output(example):
-      if not config or not config.use_tpu:
-        _summarize_features(example, num_shards)
-      if mode == tf.estimator.ModeKeys.PREDICT:
-        example["infer_targets"] = example.pop("targets")
-        return example
-      else:
-        return example, example["targets"]
-
-    dataset = dataset.map(prepare_for_output, num_parallel_calls=num_threads)
-    dataset = dataset.prefetch(2)
-
-    if mode == tf.estimator.ModeKeys.PREDICT:
-      # This is because of a bug in the Estimator that short-circuits prediction
-      # if it doesn't see a QueueRunner. DummyQueueRunner implements the
-      # minimal expected interface but does nothing.
-      tf.add_to_collection(tf.GraphKeys.QUEUE_RUNNERS,
-                           data_reader.DummyQueueRunner())
-
-    return dataset
+    return data_reader.input_fn(
+        self.dataset(**dataset_kwargs),
+        self.filepattern(data_dir, mode),
+        self.skip_random_fraction_when_training,
+        self.batch_size_means_tokens,
+        self.get_hparams().batch_size_multiplier,
+        self.max_length(hparams),
+        mode,
+        hparams,
+        data_dir=data_dir,
+        params=params,
+        config=config,
+        force_repeat=force_repeat,
+        prevent_repeat=prevent_repeat)
 
   @property
   def export_assets(self):
@@ -1039,7 +877,7 @@ def serving_input_fn(self, hparams):
     dataset = dataset.padded_batch(
         tf.shape(serialized_example, out_type=tf.int64)[0],
         dataset.output_shapes)
-    dataset = dataset.map(standardize_shapes)
+    dataset = dataset.map(data_reader.standardize_shapes)
     features = tf.data.experimental.get_single_element(dataset)
 
     if self.has_inputs:
@@ -1048,33 +886,6 @@ def serving_input_fn(self, hparams):
     return tf.estimator.export.ServingInputReceiver(
         features=features, receiver_tensors=serialized_example)
 
-  def _pad_for_tpu(self, shapes_dict, hparams):
-    """Pads unknown features' dimensions for TPU."""
-    max_length = self.max_length(hparams)
-    padded_shapes = {}
-
-    def get_filler(specified_max_length):
-      if not specified_max_length:
-        return max_length
-      return min(specified_max_length, max_length)
-
-    inputs_none_filler = get_filler(hparams.max_input_seq_length)
-    targets_none_filler = get_filler(hparams.max_target_seq_length)
-
-    def pad_one_shape(shape, none_filler):
-      return [
-          (dim if dim is not None else none_filler) for dim in shape.as_list()
-      ]
-
-    for key, shape in six.iteritems(shapes_dict):
-      if key == "inputs":
-        padded_shapes[key] = pad_one_shape(shape, inputs_none_filler)
-      elif key == "targets":
-        padded_shapes[key] = pad_one_shape(shape, targets_none_filler)
-      else:
-        padded_shapes[key] = pad_one_shape(shape, max_length)
-    return padded_shapes
-
 
 class FeatureInfo(object):
   """Encapsulates information about a feature."""
@@ -1222,71 +1033,6 @@ def _default_hparams():
       target_space_id=SpaceID.GENERIC)
 
 
-def _are_shapes_fully_defined(shapes_dict):
-  for shape in shapes_dict.values():
-    if not shape.is_fully_defined():
-      return False
-  return True
-
-
-def _summarize_features(features, num_shards=1):
-  with tf.name_scope("input_stats"):
-    for (k, v) in six.iteritems(features):
-      if isinstance(v, tf.Tensor) and v.get_shape().ndims > 1:
-        tf.summary.scalar("%s_batch" % k, tf.shape(v)[0] // num_shards)
-        tf.summary.scalar("%s_length" % k, tf.shape(v)[1])
-        nonpadding = tf.to_float(tf.not_equal(v, 0))
-        nonpadding_tokens = tf.reduce_sum(nonpadding)
-        tf.summary.scalar("%s_nonpadding_tokens" % k, nonpadding_tokens)
-        tf.summary.scalar("%s_nonpadding_fraction" % k,
-                          tf.reduce_mean(nonpadding))
-
-
-def standardize_shapes(features, batch_size=None):
-  """Set the right shapes for the features."""
-
-  for fname in ["inputs", "targets"]:
-    if fname not in features:
-      continue
-
-    f = features[fname]
-    while len(f.get_shape()) < 4:
-      f = tf.expand_dims(f, axis=-1)
-
-    features[fname] = f
-
-  if batch_size:
-    # Ensure batch size is set on all features
-    for _, t in six.iteritems(features):
-      shape = t.get_shape().as_list()
-      shape[0] = batch_size
-      t.set_shape(t.get_shape().merge_with(shape))
-      # Assert shapes are fully known
-      t.get_shape().assert_is_fully_defined()
-
-  return features
-
-
-def pad_batch(features, batch_multiple):
-  """Pad batch dim of features to nearest multiple of batch_multiple."""
-  feature = list(features.items())[0][1]
-  batch_size = tf.shape(feature)[0]
-  mod = batch_size % batch_multiple
-  has_mod = tf.cast(tf.cast(mod, tf.bool), tf.int32)
-  batch_padding = batch_multiple * has_mod - mod
-
-  padded_features = {}
-  for k, feature in features.items():
-    rank = len(feature.shape)
-    paddings = []
-    for _ in range(rank):
-      paddings.append([0, 0])
-    paddings[0][1] = batch_padding
-    padded_feature = tf.pad(feature, paddings)
-    padded_features[k] = padded_feature
-  return padded_features
-
-
 def problem_hparams_to_features(problem_hparams):
   input_space_id, target_space_id = 0, 0
   if problem_hparams:
@@ -1296,11 +1042,3 @@ def problem_hparams_to_features(problem_hparams):
       "input_space_id": input_space_id,
       "target_space_id": target_space_id,
   }
-
-
-def skip_random_fraction(dataset, data_file):
-  # Skip a random fraction at the beginning of the stream.  The skip is
-  # essential for synchronous highly-parallel training to avoid multiple
-  # replicas reading the same data in lock-step.
-  num_skip = random.randint(0, _file_num_records_cached(data_file))
-  return dataset.skip(num_skip)
diff --git a/tensor2tensor/utils/checkpoint_compatibility_test.py b/tensor2tensor/utils/checkpoint_compatibility_test.py
index 1e95ca86e..e99250133 100644
--- a/tensor2tensor/utils/checkpoint_compatibility_test.py
+++ b/tensor2tensor/utils/checkpoint_compatibility_test.py
@@ -34,7 +34,7 @@
 
 from tensor2tensor import models  # pylint: disable=unused-import
 from tensor2tensor import problems  # pylint: disable=unused-import
-from tensor2tensor.data_generators import problem
+from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -72,7 +72,7 @@ def input_fn(self):
     dataset = tf.data.Dataset.from_generator(self.input_generator, types,
                                              shapes)
     dataset = dataset.padded_batch(self.BATCH_SIZE, shapes)
-    dataset = dataset.map(problem.standardize_shapes)
+    dataset = dataset.map(data_reader.standardize_shapes)
     features = dataset.make_one_shot_iterator().get_next()
     return features
 
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index abac8e4a3..1599b6a63 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -18,9 +18,15 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
+import multiprocessing
+import random
+
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
 
+from tensor2tensor.utils import mlperf_log
+
 import tensorflow as tf
 
 
@@ -38,8 +44,9 @@ def example_length(example):
   # Length of the example is the maximum length of the feature lengths
   for _, v in sorted(six.iteritems(example)):
     # For images the sequence length is the size of the spatial dimensions.
-    feature_length = (tf.shape(v)[0] if len(v.get_shape()) < 3 else
-                      tf.shape(v)[0] * tf.shape(v)[1])
+    feature_length = tf.shape(v)[0]
+    if len(v.get_shape()) > 2:
+      feature_length = tf.shape(v)[0] * tf.shape(v)[1]
     length = tf.maximum(length, feature_length)
   return length
 
@@ -182,3 +189,310 @@ def __init__(self):
   def create_threads(self, sess, coord=None, daemon=False, start=False):
     del sess, coord, daemon, start
     return []
+
+
+def _pad_for_tpu(shapes_dict, hparams, max_length):
+  """Pads unknown features' dimensions for TPU."""
+  padded_shapes = {}
+
+  def get_filler(specified_max_length):
+    if not specified_max_length:
+      return max_length
+    return min(specified_max_length, max_length)
+
+  inputs_none_filler = get_filler(hparams.max_input_seq_length)
+  targets_none_filler = get_filler(hparams.max_target_seq_length)
+
+  def pad_one_shape(shape, none_filler):
+    return [
+        (dim if dim is not None else none_filler) for dim in shape.as_list()
+    ]
+
+  for key, shape in six.iteritems(shapes_dict):
+    if key == "inputs":
+      padded_shapes[key] = pad_one_shape(shape, inputs_none_filler)
+    elif key == "targets":
+      padded_shapes[key] = pad_one_shape(shape, targets_none_filler)
+    else:
+      padded_shapes[key] = pad_one_shape(shape, max_length)
+  return padded_shapes
+
+
+def cpu_count():
+  """Return the number of available cores."""
+  num_available_cores = multiprocessing.cpu_count()
+  return num_available_cores
+
+
+def _summarize_features(features, num_shards=1):
+  with tf.name_scope("input_stats"):
+    for (k, v) in six.iteritems(features):
+      if isinstance(v, tf.Tensor) and v.get_shape().ndims > 1:
+        tf.summary.scalar("%s_batch" % k, tf.shape(v)[0] // num_shards)
+        tf.summary.scalar("%s_length" % k, tf.shape(v)[1])
+        nonpadding = tf.to_float(tf.not_equal(v, 0))
+        nonpadding_tokens = tf.reduce_sum(nonpadding)
+        tf.summary.scalar("%s_nonpadding_tokens" % k, nonpadding_tokens)
+        tf.summary.scalar("%s_nonpadding_fraction" % k,
+                          tf.reduce_mean(nonpadding))
+
+
+def standardize_shapes(features, batch_size=None):
+  """Set the right shapes for the features."""
+  for fname in ["inputs", "targets"]:
+    if fname not in features:
+      continue
+    f = features[fname]
+    while len(f.get_shape()) < 4:
+      f = tf.expand_dims(f, axis=-1)
+    features[fname] = f
+
+  if batch_size:
+    # Ensure batch size is set on all features
+    for _, t in six.iteritems(features):
+      shape = t.get_shape().as_list()
+      shape[0] = batch_size
+      t.set_shape(t.get_shape().merge_with(shape))
+      # Assert shapes are fully known
+      t.get_shape().assert_is_fully_defined()
+
+  return features
+
+
+def _are_shapes_fully_defined(shapes_dict):
+  for shape in shapes_dict.values():
+    if not shape.is_fully_defined():
+      return False
+  return True
+
+
+def _file_num_records_cached(filename):
+  """Return the number of TFRecords in a file."""
+  # Cache the result, as this is expensive to compute
+  if filename in _file_num_records_cache:
+    return _file_num_records_cache[filename]
+  ret = 0
+  for _ in tf.python_io.tf_record_iterator(filename):
+    ret += 1
+  _file_num_records_cache[filename] = ret
+  return ret
+
+
+_file_num_records_cache = {}
+
+
+def skip_random_fraction(dataset, data_file):
+  # Skip a random fraction at the beginning of the stream.  The skip is
+  # essential for synchronous highly-parallel training to avoid multiple
+  # replicas reading the same data in lock-step.
+  num_skip = random.randint(0, _file_num_records_cached(data_file))
+  return dataset.skip(num_skip)
+
+
+def _pad_batch(features, batch_multiple):
+  """Pad batch dim of features to nearest multiple of batch_multiple."""
+  feature = list(features.items())[0][1]
+  batch_size = tf.shape(feature)[0]
+  mod = batch_size % batch_multiple
+  has_mod = tf.cast(tf.cast(mod, tf.bool), tf.int32)
+  batch_padding = batch_multiple * has_mod - mod
+
+  padded_features = {}
+  for k, feature in features.items():
+    rank = len(feature.shape)
+    paddings = []
+    for _ in range(rank):
+      paddings.append([0, 0])
+    paddings[0][1] = batch_padding
+    padded_feature = tf.pad(feature, paddings)
+    padded_features[k] = padded_feature
+  return padded_features
+
+
+# TODO(lukaszkaiser): refactor the API to not be just a list of self params
+#   but make sense for other uses too.
+def input_fn(dataset,
+             filepattern,
+             skip_random_fraction_when_training,
+             batch_size_means_tokens_param,
+             batch_size_multiplier,
+             max_length,
+             mode,
+             hparams,
+             data_dir=None,
+             params=None,
+             config=None,
+             force_repeat=False,
+             prevent_repeat=False):
+  """Builds input pipeline for problem.
+
+  Args:
+    dataset: the dataset to make input function from.
+    filepattern: the pattern of files to read from.
+    skip_random_fraction_when_training: whether to skip randomly when training.
+    batch_size_means_tokens_param: whether batch size should mean tokens.
+    batch_size_multiplier: how to multiply batch size when bucketing.
+    max_length: maximum length,
+    mode: tf.estimator.ModeKeys
+    hparams: HParams, model hparams
+    data_dir: str, data directory; if None, will use hparams.data_dir
+    params: dict, may include "batch_size"
+    config: RunConfig; should have the data_parallelism attribute if not using
+      TPU
+    force_repeat: bool, whether to repeat the data even if not training
+    prevent_repeat: bool, whether to not repeat when in training mode.
+      Overrides force_repeat.
+
+  Returns:
+    (features_dict<str name, Tensor feature>, Tensor targets)
+  """
+  is_training = mode == tf.estimator.ModeKeys.TRAIN
+  if config and config.use_tpu:
+    num_threads = 64
+  else:
+    num_threads = cpu_count() if is_training else 1
+
+  if config and hasattr(config,
+                        "data_parallelism") and config.data_parallelism:
+    num_shards = config.data_parallelism.n
+  else:
+    num_shards = 1
+
+  mlperf_log.transformer_print(
+      key=mlperf_log.INPUT_MAX_LENGTH, value=max_length)
+
+  def tpu_valid_size(example):
+    return example_valid_size(example, hparams.min_length, max_length)
+
+  def gpu_valid_size(example):
+    drop_long_sequences = is_training or hparams.eval_drop_long_sequences
+    max_validate_length = max_length if drop_long_sequences else 10**9
+    return example_valid_size(example, hparams.min_length, max_validate_length)
+
+  def define_shapes(example):
+    batch_size = config and config.use_tpu and params["batch_size"]
+    return standardize_shapes(example, batch_size=batch_size)
+
+  # Read and preprocess
+  data_dir = data_dir or (hasattr(hparams, "data_dir") and hparams.data_dir)
+
+  if (force_repeat or is_training) and not prevent_repeat:
+    # Repeat and skip a random number of records
+    dataset = dataset.repeat()
+
+  if is_training and skip_random_fraction_when_training:
+    data_files = tf.contrib.slim.parallel_reader.get_data_files(filepattern)
+    #  In continuous_train_and_eval when switching between train and
+    #  eval, this input_fn method gets called multiple times and it
+    #  would give you the exact same samples from the last call
+    #  (because the Graph seed is set). So this skip gives you some
+    #  shuffling.
+    dataset = skip_random_fraction(dataset, data_files[0])
+
+  dataset = dataset.map(cast_ints_to_int32, num_parallel_calls=num_threads)
+
+  if batch_size_means_tokens_param:
+    batch_size_means_tokens = True
+  else:
+    if _are_shapes_fully_defined(dataset.output_shapes):
+      batch_size_means_tokens = False
+    else:
+      tf.logging.warning(
+          "Shapes are not fully defined. Assuming batch_size means tokens.")
+      batch_size_means_tokens = True
+
+  # Batching
+  if not batch_size_means_tokens:
+    # Batch size means examples per datashard.
+    if config and config.use_tpu:
+      # on TPU, we use params["batch_size"], which specifies the number of
+      # examples across all datashards
+      batch_size = params["batch_size"]
+      dataset = dataset.batch(batch_size, drop_remainder=True)
+    else:
+      batch_size = hparams.batch_size * num_shards
+      dataset = dataset.batch(batch_size)
+  else:
+    # batch_size means tokens per datashard
+    if config and config.use_tpu:
+      dataset = dataset.filter(tpu_valid_size)
+      padded_shapes = _pad_for_tpu(dataset.output_shapes, hparams, max_length)
+      # on TPU, we use params["batch_size"], which specifies the number of
+      # examples across all datashards
+      batch_size = params["batch_size"]
+      if hparams.pad_batch:
+        tf.logging.warn(
+            "Padding the batch to ensure that remainder eval batches are "
+            "processed. This may lead to incorrect metrics for "
+            "non-zero-padded features, e.g. images. Use a smaller batch "
+            "size that has no remainder in that case.")
+        dataset = dataset.padded_batch(
+            batch_size, padded_shapes, drop_remainder=False)
+        dataset = dataset.map(
+            functools.partial(_pad_batch, batch_multiple=batch_size),
+            num_parallel_calls=num_threads)
+      else:
+        dataset = dataset.padded_batch(
+            batch_size, padded_shapes, drop_remainder=True)
+    else:
+      # On GPU, bucket by length
+      dataset = dataset.filter(gpu_valid_size)
+      batching_scheme = hparams_to_batching_scheme(
+          hparams,
+          shard_multiplier=num_shards,
+          length_multiplier=batch_size_multiplier)
+      if hparams.use_fixed_batch_size:
+        # Here  batch_size really means examples per datashard.
+        batching_scheme["batch_sizes"] = [hparams.batch_size]
+        batching_scheme["boundaries"] = []
+      dataset = dataset.apply(
+          tf.data.experimental.bucket_by_sequence_length(
+              example_length, batching_scheme["boundaries"],
+              batching_scheme["batch_sizes"]))
+
+      if not is_training:
+        batch_multiple = num_shards
+        if hparams.use_fixed_batch_size:
+          # Make sure the last batch has the same fixed size as the rest.
+          batch_multiple *= hparams.batch_size
+        if batch_multiple > 1:
+          tf.logging.warn(
+              "Padding the batch to ensure that remainder eval batches have "
+              "a batch size divisible by the number of data shards. This may "
+              "lead to incorrect metrics for non-zero-padded features, e.g. "
+              "images. Use a single datashard (i.e. 1 GPU) in that case.")
+          dataset = dataset.map(
+              functools.partial(_pad_batch, batch_multiple=batch_multiple),
+              num_parallel_calls=num_threads)
+
+  dataset = dataset.map(define_shapes, num_parallel_calls=num_threads)
+
+  # Add shuffling for training batches. This is necessary along with record
+  # level shuffling in the dataset generation. Record shuffling will shuffle
+  # the examples. However, in some cases, it's possible that the shuffle
+  # buffer size for record shuffling is smaller than the batch size. In such
+  # cases, adding batch shuffling ensures that the data is in random order
+  # during training
+  if (is_training and hasattr(hparams, "batch_shuffle_size") and
+      hparams.batch_shuffle_size):
+    dataset = dataset.shuffle(hparams.batch_shuffle_size)
+
+  def prepare_for_output(example):
+    if not config or not config.use_tpu:
+      _summarize_features(example, num_shards)
+    if mode == tf.estimator.ModeKeys.PREDICT:
+      example["infer_targets"] = example.pop("targets")
+      return example
+    else:
+      return example, example["targets"]
+
+  dataset = dataset.map(prepare_for_output, num_parallel_calls=num_threads)
+  dataset = dataset.prefetch(2)
+
+  if mode == tf.estimator.ModeKeys.PREDICT:
+    # This is because of a bug in the Estimator that short-circuits prediction
+    # if it doesn't see a QueueRunner. DummyQueueRunner implements the
+    # minimal expected interface but does nothing.
+    tf.add_to_collection(tf.GraphKeys.QUEUE_RUNNERS, DummyQueueRunner())
+
+  return dataset
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index b91f43dd7..194b771fa 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -21,8 +21,8 @@
 
 import os
 from tensor2tensor.data_generators import algorithmic
-from tensor2tensor.data_generators import problem as problem_lib
 from tensor2tensor.models import transformer  # pylint: disable=unused-import
+from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 import tensorflow as tf
@@ -82,7 +82,7 @@ def testModel(self):
                               algorithmic.TinyAlgo.data_dir)
     dataset = dataset.repeat(None).padded_batch(10, dataset.output_shapes)
     features = dataset.make_one_shot_iterator().get_next()
-    features = problem_lib.standardize_shapes(features)
+    features = data_reader.standardize_shapes(features)
 
     # Model
     model = registry.model("transformer")(hparams, tf.estimator.ModeKeys.TRAIN)
@@ -119,7 +119,7 @@ def testMultipleTargetModalities(self):
                               algorithmic.TinyAlgo.data_dir)
     dataset = dataset.repeat(None).padded_batch(10, dataset.output_shapes)
     features = dataset.make_one_shot_iterator().get_next()
-    features = problem_lib.standardize_shapes(features)
+    features = data_reader.standardize_shapes(features)
     features["targets_A"] = features["targets_B"] = features["targets"]
 
     # Model

From 415585f40d9f21c56df7bda35033bc915d82321e Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Wed, 9 Jan 2019 22:07:29 +0100
Subject: [PATCH 1468/2720] Out-of-graph RL evaluation (#1355)

* Add an option to specify an evaluation function to evaluate_all_configs

* Implement an out-of-graph evaluator

* Add a test for the out-of-graph evaluator
---
 tensor2tensor/rl/evaluator.py      | 61 +++++++++++++++++++++++--
 tensor2tensor/rl/evaluator_test.py | 39 ++++++++++++++++
 tensor2tensor/rl/rl_utils.py       | 73 ++++++++++++++++++++++++++----
 3 files changed, 160 insertions(+), 13 deletions(-)
 create mode 100644 tensor2tensor/rl/evaluator_test.py

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index c631b43b2..e7a69daea 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -34,6 +34,7 @@
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
 
+import numpy as np
 import tensorflow as tf
 
 
@@ -46,16 +47,65 @@
     "eval_metrics_dir", "", "Directory to output the eval metrics at."
 )
 flags.DEFINE_bool("full_eval", True, "Whether to ignore the timestep limit.")
+flags.DEFINE_enum("agent", "policy", ["random", "policy"], "Agent type to use.")
+flags.DEFINE_bool(
+    "eval_with_learner", True,
+    "Whether to use the PolicyLearner.evaluate function instead of an "
+    "out-of-graph one. Works only with --agent=policy."
+)
 
 
-def evaluate(hparams, policy_dir, eval_metrics_dir, report_fn=None,
-             report_metric=None):
+def make_agent(agent_type, action_space):
+  """Factory function for Agents."""
+  return {
+      "random": rl_utils.RandomAgent,
+  }[agent_type](action_space)
+
+
+def make_eval_fn_with_agent(agent_type):
+  """Returns an out-of-graph eval_fn using the Agent API."""
+  def eval_fn(env, hparams, policy_hparams, policy_dir, sampling_temp):
+    agent = make_agent(agent_type, env.action_space)
+    num_dones = 0
+    first_dones = [False] * env.batch_size
+    observations = env.reset()
+    while num_dones < env.batch_size:
+      actions = agent.act(observations)
+      (observations, _, dones) = env.step(actions)
+      observations = list(observations)
+      now_done_indices = []
+      for (i, done) in enumerate(dones):
+        if done and not first_dones[i]:
+          now_done_indices.append(i)
+          first_dones[i] = True
+          num_dones += 1
+      if now_done_indices:
+        # Reset only envs done the first time in this timestep to ensure that
+        # we collect exactly 1 rollout from each env.
+        reset_observations = env.reset(now_done_indices)
+        for (i, observation) in zip(now_done_indices, reset_observations):
+          observations[i] = observation
+      observations = np.array(observations)
+    assert len(env.current_epoch_rollouts()) == env.batch_size
+  return eval_fn
+
+
+def evaluate(
+    hparams, policy_dir, eval_metrics_dir, agent_type, eval_with_learner,
+    report_fn=None, report_metric=None
+):
   """Evaluate."""
+  if eval_with_learner:
+    assert agent_type == "policy"
+
   if report_fn:
     assert report_metric is not None
 
   eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
-  eval_metrics = rl_utils.evaluate_all_configs(hparams, policy_dir)
+  kwargs = {}
+  if not eval_with_learner:
+    kwargs["eval_fn"] = make_eval_fn_with_agent(agent_type)
+  eval_metrics = rl_utils.evaluate_all_configs(hparams, policy_dir, **kwargs)
   rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, 0)
 
   # Report metrics
@@ -76,7 +126,10 @@ def main(_):
   hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
   if FLAGS.full_eval:
     hparams.eval_rl_env_max_episode_steps = -1
-  evaluate(hparams, FLAGS.policy_dir, FLAGS.eval_metrics_dir)
+  evaluate(
+      hparams, FLAGS.policy_dir, FLAGS.eval_metrics_dir, FLAGS.agent,
+      FLAGS.eval_with_learner
+  )
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/rl/evaluator_test.py b/tensor2tensor/rl/evaluator_test.py
new file mode 100644
index 000000000..969232faf
--- /dev/null
+++ b/tensor2tensor/rl/evaluator_test.py
@@ -0,0 +1,39 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 the evaluator."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.rl import evaluator
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+class EvalTest(tf.test.TestCase):
+
+  def test_evaluate_pong_random_agent(self):
+    hparams = registry.hparams("rlmb_tiny")
+    temp_dir = tf.test.get_temp_dir()
+    evaluator.evaluate(
+        hparams, temp_dir, temp_dir, agent_type="random",
+        eval_with_learner=False
+    )
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 876ccb78e..30d43665c 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -59,8 +59,20 @@ def get_metric_name(sampling_temp, max_num_noops, clipped):
   )
 
 
+def _eval_fn_with_learner(
+    env, hparams, policy_hparams, policy_dir, sampling_temp
+):
+  env_fn = rl.make_real_env_fn(env)
+  learner = LEARNERS[hparams.base_algo](
+      hparams.frame_stack_size, base_event_dir=None,
+      agent_model_dir=policy_dir, total_num_epochs=1
+  )
+  learner.evaluate(env_fn, policy_hparams, sampling_temp)
+
+
 def evaluate_single_config(
-    hparams, sampling_temp, max_num_noops, agent_model_dir
+    hparams, sampling_temp, max_num_noops, agent_model_dir,
+    eval_fn=_eval_fn_with_learner
 ):
   """Evaluate the PPO agent in the real environment."""
   eval_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
@@ -69,12 +81,7 @@ def evaluate_single_config(
       rl_env_max_episode_steps=hparams.eval_rl_env_max_episode_steps
   )
   env.start_new_epoch(0)
-  env_fn = rl.make_real_env_fn(env)
-  learner = LEARNERS[hparams.base_algo](
-      hparams.frame_stack_size, base_event_dir=None,
-      agent_model_dir=agent_model_dir, total_num_epochs=1
-  )
-  learner.evaluate(env_fn, eval_hparams, sampling_temp)
+  eval_fn(env, hparams, eval_hparams, agent_model_dir, sampling_temp)
   rollouts = env.current_epoch_rollouts()
   env.close()
 
@@ -83,7 +90,9 @@ def evaluate_single_config(
   )
 
 
-def evaluate_all_configs(hparams, agent_model_dir):
+def evaluate_all_configs(
+    hparams, agent_model_dir, eval_fn=_eval_fn_with_learner
+):
   """Evaluate the agent with multiple eval configurations."""
   metrics = {}
   # Iterate over all combinations of sampling temperatures and whether to do
@@ -91,7 +100,7 @@ def evaluate_all_configs(hparams, agent_model_dir):
   for sampling_temp in hparams.eval_sampling_temps:
     for max_num_noops in (hparams.eval_max_num_noops, 0):
       scores = evaluate_single_config(
-          hparams, sampling_temp, max_num_noops, agent_model_dir
+          hparams, sampling_temp, max_num_noops, agent_model_dir, eval_fn
       )
       for (score, clipped) in zip(scores, (True, False)):
         metric_name = get_metric_name(sampling_temp, max_num_noops, clipped)
@@ -185,3 +194,49 @@ def initial_frame_chooser(batch_size):
         for initial_frame_stack in initial_frames
     ])
   return initial_frame_chooser
+
+
+class BatchAgent(object):
+  """Python API for agents.
+
+  Runs a batch of parallel agents. Operates on Numpy arrays.
+  """
+
+  def __init__(self, action_space):
+    self.action_space = action_space
+
+  def act(self, observations):
+    """Picks actions based on observations.
+
+    Args:
+      observations: A batch of observations.
+
+    Returns:
+      A batch of actions.
+    """
+    raise NotImplementedError
+
+  def estimate_value(self, observations):
+    """Estimates values of states based on observations.
+
+    Used for temporal-difference planning.
+
+    Args:
+      observations: A batch of observations.
+
+    Returns:
+      A batch of values.
+    """
+    raise NotImplementedError
+
+
+class RandomAgent(BatchAgent):
+  """Random agent, sampling actions from the uniform distribution."""
+
+  def act(self, observations):
+    return np.array([
+        self.action_space.sample() for _ in range(observations.shape[0])
+    ])
+
+  def estimate_value(self, observations):
+    return np.zeros(observations.shape[0])

From 1a3ba3345cd40c021a6e91d5a73c94e734130cac Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Wed, 9 Jan 2019 13:07:50 -0800
Subject: [PATCH 1469/2720] internal merge of PR #1355

PiperOrigin-RevId: 228573215
---
 tensor2tensor/models/__init__.py | 3 +++
 tensor2tensor/rl/evaluator.py    | 5 ++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 0ab7b3f9a..f0618faba 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -17,6 +17,9 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+import six
+
 # pylint: disable=unused-import
 
 from tensor2tensor.layers import modalities  # pylint: disable=g-import-not-at-top
diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index e7a69daea..3dea32510 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -28,13 +28,14 @@
 from __future__ import division
 from __future__ import print_function
 
+import numpy as np
+
 from tensor2tensor.models.research import rl  # pylint: disable=unused-import
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl import trainer_model_based_params  # pylint: disable=unused-import
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
 
-import numpy as np
 import tensorflow as tf
 
 
@@ -65,6 +66,8 @@ def make_agent(agent_type, action_space):
 def make_eval_fn_with_agent(agent_type):
   """Returns an out-of-graph eval_fn using the Agent API."""
   def eval_fn(env, hparams, policy_hparams, policy_dir, sampling_temp):
+    """Eval function."""
+    del hparams, policy_hparams, policy_dir, sampling_temp
     agent = make_agent(agent_type, env.action_space)
     num_dones = 0
     first_dones = [False] * env.batch_size

From 7ba9f5a56df4b0bebb65980b0e799696d10ab3f0 Mon Sep 17 00:00:00 2001
From: Chan Yu <aeloyq@outlook.com>
Date: Thu, 10 Jan 2019 05:36:30 +0800
Subject: [PATCH 1470/2720] Avoid error in beam search when "f" is in cache
 (#1302)

* add caching mechanism support for fast decoding with relative_dot_product in transformer model

* fix typo

* remove f in cache if use dense_relu_dense or conv_hidden_relu so that errors won't occur in beamsearch (nest.map function on cache)
---
 tensor2tensor/models/transformer.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 692f88afb..b9d01c018 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -779,10 +779,11 @@ def fast_decode_tpu(encoder_output,
           common_attention.split_heads(
               tf.zeros([batch_size, decode_length, value_channels]),
               hparams.num_heads),
-          "f":
-          tf.zeros([batch_size, decode_length, hparams.hidden_size]),
       } for layer in range(num_layers)
   }
+  if hparams.ffn_layer not in ["dense_relu_dense", "conv_hidden_relu"]:
+    for layer in range(num_layers):
+      cache["layer_%d" % layer]["f"] = tf.zeros([batch_size, 0, hparams.hidden_size])
 
   if encoder_output is not None:
     for layer in range(num_layers):
@@ -957,10 +958,11 @@ def fast_decode(encoder_output,
           "v":
               common_attention.split_heads(
                   tf.zeros([batch_size, 0, value_channels]), hparams.num_heads),
-          "f":
-              tf.zeros([batch_size, 0, hparams.hidden_size]),
       } for layer in range(num_layers)
   }
+  if hparams.ffn_layer not in ["dense_relu_dense", "conv_hidden_relu"]:
+    for layer in range(num_layers):
+      cache["layer_%d" % layer]["f"] = tf.zeros([batch_size, 0, hparams.hidden_size])
 
   if encoder_output is not None:
     for layer in range(num_layers):

From 4bbdc85e3390cf11a4db32ff25da37b9b2fc7155 Mon Sep 17 00:00:00 2001
From: Chan Yu <aeloyq@outlook.com>
Date: Wed, 9 Jan 2019 13:50:20 -0800
Subject: [PATCH 1471/2720] internal merge of PR #1302

PiperOrigin-RevId: 228581110
---
 tensor2tensor/models/transformer.py | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index b9d01c018..a5967879f 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -781,9 +781,16 @@ def fast_decode_tpu(encoder_output,
               hparams.num_heads),
       } for layer in range(num_layers)
   }
+
+  # If `ffn_layer` is in `["dense_relu_dense" or "conv_hidden_relu"]`, then the
+  # cache key "f" won't be used, which means that the` shape of cache["f"]`
+  # won't be changed to
+  # `[beamsize*batch_size, decode_length, hparams.hidden_size]` and may cause
+  # error when applying `nest.map reshape function` on it.
   if hparams.ffn_layer not in ["dense_relu_dense", "conv_hidden_relu"]:
     for layer in range(num_layers):
-      cache["layer_%d" % layer]["f"] = tf.zeros([batch_size, 0, hparams.hidden_size])
+      cache["layer_%d" % layer]["f"] = tf.zeros(
+          [batch_size, 0, hparams.hidden_size])
 
   if encoder_output is not None:
     for layer in range(num_layers):
@@ -960,9 +967,16 @@ def fast_decode(encoder_output,
                   tf.zeros([batch_size, 0, value_channels]), hparams.num_heads),
       } for layer in range(num_layers)
   }
+
+  # If `ffn_layer` is in `["dense_relu_dense" or "conv_hidden_relu"]`, then the
+  # cache key "f" won't be used, which means that the` shape of cache["f"]`
+  # won't be changed to
+  # `[beamsize*batch_size, decode_length, hparams.hidden_size]` and may cause
+  # error when applying `nest.map reshape function` on it.
   if hparams.ffn_layer not in ["dense_relu_dense", "conv_hidden_relu"]:
     for layer in range(num_layers):
-      cache["layer_%d" % layer]["f"] = tf.zeros([batch_size, 0, hparams.hidden_size])
+      cache["layer_%d" % layer]["f"] = tf.zeros(
+          [batch_size, 0, hparams.hidden_size])
 
   if encoder_output is not None:
     for layer in range(num_layers):

From f6a450af64529ce93b6080c0f9b3937b4eecc933 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Wed, 9 Jan 2019 15:57:13 -0800
Subject: [PATCH 1472/2720] Fix sharded decoding name convention.

PiperOrigin-RevId: 228604730
---
 tensor2tensor/utils/decoding.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 1045c621f..36eadb3ed 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -70,6 +70,7 @@ def decode_hparams(overrides=""):
       shards=1,    # How many shards of data to decode (treating 1 as None).
       shard_id=0,  # Which shard are we decoding if more than 1 above.
       shards_start_offset=0,  # Number of the first shard to decode.
+      shard_google_format=False,  # If True use Google shard naming format.
       num_decodes=1,
       force_decode_length=False,
       display_decoded_images=False,
@@ -510,7 +511,11 @@ def timer(gen):
 def _add_shard_to_filename(filename, decode_hp):
   if decode_hp.shards > 1:
     shard_id = decode_hp.shard_id + decode_hp.shards_start_offset
-    filename = filename + ("%.3d" % shard_id)
+    if decode_hp.shard_google_format:
+      filename = filename + "-{0:05d}-of-{1:05d}".format(shard_id,
+                                                         decode_hp.shards)
+    else:
+      filename = filename + ("%.3d" % shard_id)
   return filename
 
 

From cef5491f5bcd9d53e744b4f752eac3117acd18cb Mon Sep 17 00:00:00 2001
From: Giovanni Campagna <scampa.giovanni@gmail.com>
Date: Wed, 9 Jan 2019 17:48:19 -0800
Subject: [PATCH 1473/2720] transformer_fast_decode, beam search: take an
 optional cache and return it (#999)

Some models, eg. semantic parsing models with copying mechanisms,
want to use the output of Transformer for multiple predictions.
One way to do so is to modify the symbols_to_logits_fn to generate
the additional predictions and save it in the cache dictionary.

To do so, though, fast_decode() must allow an externally supplied cache,
and must return it to the caller after the loop.
---
 tensor2tensor/layers/latent_layers.py         |  2 +-
 .../models/research/transformer_nat.py        |  2 +-
 .../models/research/transformer_vae.py        |  2 +-
 tensor2tensor/models/transformer.py           | 15 +++---
 tensor2tensor/utils/beam_search.py            |  4 +-
 tensor2tensor/utils/beam_search_test.py       | 51 ++++++++++++++++---
 tensor2tensor/utils/t2t_model.py              |  2 +-
 7 files changed, 58 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 9a10eaeb0..c32cd691d 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -168,7 +168,7 @@ def symbols_to_logits_fn(ids):
 
   initial_ids = tf.zeros([tf.shape(latents_dense_in)[0]], dtype=tf.int32)
   length = tf.shape(latents_dense_in)[1]
-  ids, _ = beam_search.beam_search(
+  ids, _, _ = beam_search.beam_search(
       symbols_to_logits_fn,
       initial_ids,
       1,
diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 2891d3897..7a775153d 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -228,7 +228,7 @@ def symbols_to_logits_fn(ids):
 
   initial_ids = tf.zeros([tf.shape(latents_dense_in)[0]], dtype=tf.int32)
   length = tf.shape(latents_dense_in)[1]
-  ids, _ = beam_search.beam_search(
+  ids, _, _ = beam_search.beam_search(
       symbols_to_logits_fn,
       initial_ids,
       beam_size=1,
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 11446fd64..88f0eb757 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -290,7 +290,7 @@ def symbols_to_logits_fn(ids):
 
   initial_ids = tf.zeros([tf.shape(latents_dense_in)[0]], dtype=tf.int32)
   length = tf.shape(latents_dense_in)[1]
-  ids, _ = beam_search.beam_search(
+  ids, _, _ = beam_search.beam_search(
       symbols_to_logits_fn, initial_ids, beam_size, length,
       vocab_size, alpha=0.0, eos_id=-1, stop_early=False)
 
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index a5967879f..5e604f95b 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -910,7 +910,8 @@ def fast_decode(encoder_output,
                 eos_id=beam_search.EOS_ID,
                 batch_size=None,
                 force_decode_length=False,
-                scope_prefix="body/"):
+                scope_prefix="body/",
+                cache=None):
   """Given encoder output and a symbols to logits function, does fast decoding.
 
   Implements both greedy and beam search decoding, uses beam search iff
@@ -957,7 +958,9 @@ def fast_decode(encoder_output,
   vars_3d_num_heads = (
       hparams.num_heads if hparams.get("attention_variables_3d") else 0)
 
-  cache = {
+  if cache is None:
+    cache = dict()
+  cache.update({
       "layer_%d" % layer: {
           "k":
               common_attention.split_heads(
@@ -966,7 +969,7 @@ def fast_decode(encoder_output,
               common_attention.split_heads(
                   tf.zeros([batch_size, 0, value_channels]), hparams.num_heads),
       } for layer in range(num_layers)
-  }
+  })
 
   # If `ffn_layer` is in `["dense_relu_dense" or "conv_hidden_relu"]`, then the
   # cache key "f" won't be used, which means that the` shape of cache["f"]`
@@ -1000,7 +1003,7 @@ def fast_decode(encoder_output,
 
   if beam_size > 1:  # Beam Search
     initial_ids = sos_id * tf.ones([batch_size], dtype=tf.int32)
-    decoded_ids, scores = beam_search.beam_search(
+    decoded_ids, scores, cache = beam_search.beam_search(
         symbols_to_logits_fn,
         initial_ids,
         beam_size,
@@ -1047,7 +1050,7 @@ def is_not_finished(i, hit_eos, *_):
     hit_eos = tf.fill([batch_size], False)
     next_id = sos_id * tf.ones([batch_size, 1], dtype=tf.int64)
     initial_log_prob = tf.zeros([batch_size], dtype=tf.float32)
-    _, _, _, decoded_ids, _, log_prob = tf.while_loop(
+    _, _, _, decoded_ids, cache, log_prob = tf.while_loop(
         is_not_finished,
         inner_loop, [
             tf.constant(0), hit_eos, next_id, decoded_ids, cache,
@@ -1063,7 +1066,7 @@ def is_not_finished(i, hit_eos, *_):
         ])
     scores = log_prob
 
-  return {"outputs": decoded_ids, "scores": scores}
+  return {"outputs": decoded_ids, "scores": scores, "cache": cache}
 
 
 @registry.register_model
diff --git a/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index c303d63c1..bddce3fa9 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -752,7 +752,7 @@ def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
         tf.less(i, decode_length), tf.logical_not(bound_is_met))
 
   (_, alive_seq, alive_log_probs, finished_seq, finished_scores,
-   finished_flags, _) = tf.while_loop(
+   finished_flags, states) = tf.while_loop(
        _is_finished,
        inner_loop, [
            tf.constant(0), alive_seq, alive_log_probs, finished_seq,
@@ -786,4 +786,4 @@ def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
       tf.reduce_any(finished_flags, 1), finished_seq, alive_seq)
   finished_scores = tf.where(
       tf.reduce_any(finished_flags, 1), finished_scores, alive_log_probs)
-  return finished_seq, finished_scores
+  return finished_seq, finished_scores, states
diff --git a/tensor2tensor/utils/beam_search_test.py b/tensor2tensor/utils/beam_search_test.py
index 9b6102252..1b11d344c 100644
--- a/tensor2tensor/utils/beam_search_test.py
+++ b/tensor2tensor/utils/beam_search_test.py
@@ -38,7 +38,7 @@ def symbols_to_logits(_):
       # Just return random logits
       return tf.random_uniform((batch_size * beam_size, vocab_size))
 
-    final_ids, final_probs = beam_search.beam_search(
+    final_ids, final_probs, _ = beam_search.beam_search(
         symbols_to_logits, initial_ids, beam_size, decode_length, vocab_size,
         0.)
 
@@ -114,7 +114,7 @@ def symbols_to_logits(ids):
       logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
       return logits
 
-    final_ids, final_probs = beam_search.beam_search(
+    final_ids, final_probs, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
@@ -145,7 +145,7 @@ def symbols_to_logits(ids):
       logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
       return logits
 
-    final_ids, final_probs = beam_search.beam_search(
+    final_ids, final_probs, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
@@ -214,7 +214,7 @@ def symbols_to_logits(ids):
       logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
       return logits
 
-    final_ids, final_probs = beam_search.beam_search(
+    final_ids, final_probs, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
@@ -254,7 +254,7 @@ def symbols_to_logits(ids):
       logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
       return logits
 
-    final_ids, final_scores = beam_search.beam_search(
+    final_ids, final_scores, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
@@ -297,7 +297,7 @@ def symbols_to_logits(ids):
       return logits
 
     # Disable early stopping
-    final_ids, final_scores = beam_search.beam_search(
+    final_ids, final_scores, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
@@ -343,7 +343,7 @@ def symbols_to_logits(ids, _, states):
     states["state"] = tf.placeholder_with_default(
         states["state"], shape=(None, 1))
 
-    final_ids, _ = beam_search.beam_search(
+    final_ids, _, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
@@ -360,6 +360,41 @@ def symbols_to_logits(ids, _, states):
       except tf.errors.InvalidArgumentError as e:
         raise AssertionError(e.message)
 
+  def testStatesAfterLoop(self):
+    batch_size = 1
+    beam_size = 1
+    vocab_size = 2
+    decode_length = 3
+
+    initial_ids = tf.constant([0] * batch_size)  # GO
+    probabilities = tf.constant([[[0.7, 0.3]], [[0.4, 0.6]], [[0.5, 0.5]]])
+
+    def symbols_to_logits(ids, _, states):
+      pos = tf.shape(ids)[1] - 1
+      logits = tf.to_float(tf.log(probabilities[pos, :]))
+      states["state"] += 1
+      return logits, states
+
+    states = {
+        "state": tf.zeros((batch_size, 1)),
+    }
+    states["state"] = tf.placeholder_with_default(
+        states["state"], shape=(None, 1))
+
+    _, _, final_states = beam_search.beam_search(
+        symbols_to_logits,
+        initial_ids,
+        beam_size,
+        decode_length,
+        vocab_size,
+        0.0,
+        eos_id=1,
+        states=states)
+    
+    with self.test_session() as sess:
+      final_states = sess.run(final_states)
+    self.assertAllEqual([[1]], final_states["state"])
+
   def testStateBeamTwo(self):
     batch_size = 1
     beam_size = 2
@@ -393,7 +428,7 @@ def symbols_to_logits(ids, _, states):
     states["state"] = tf.placeholder_with_default(
         states["state"], shape=(None, 1))
 
-    final_ids, _ = beam_search.beam_search(
+    final_ids, _, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 8fea747d7..5443f0c4d 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -891,7 +891,7 @@ def symbols_to_logits_fn(ids, i=None):
       inputs = features["inputs"]
       decode_length = (common_layers.shape_list(inputs)[1] +
                        features.get("decode_length", decode_length))
-    ids, scores = beam_search.beam_search(
+    ids, scores, _ = beam_search.beam_search(
         symbols_to_logits_fn,
         initial_ids,
         beam_size,

From abca2107d0359483bc925dddf13efea6eeb3b2ab Mon Sep 17 00:00:00 2001
From: Giovanni Campagna <scampa.giovanni@gmail.com>
Date: Wed, 9 Jan 2019 17:49:47 -0800
Subject: [PATCH 1474/2720] internal merge of PR #999

PiperOrigin-RevId: 228622817
---
 tensor2tensor/models/transformer.py     |  7 ++++---
 tensor2tensor/utils/beam_search.py      | 17 ++++++++++-------
 tensor2tensor/utils/beam_search_test.py |  8 ++++----
 3 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 5e604f95b..c7ae16d97 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -57,7 +57,7 @@ class Transformer(t2t_model.T2TModel):
 
   def __init__(self, *args, **kwargs):
     super(Transformer, self).__init__(*args, **kwargs)
-    self.attention_weights = dict()  # For visualizing attention heads.
+    self.attention_weights = {}  # For visualizing attention heads.
 
   def encode(self, inputs, target_space, hparams, features=None, losses=None):
     """Encode transformer inputs.
@@ -824,7 +824,7 @@ def fast_decode_tpu(encoder_output,
       hparams=hparams)
   if beam_size > 1:  # Beam Search
     initial_ids = sos_id * tf.ones([batch_size], dtype=tf.int32)
-    decoded_ids, scores = beam_search.beam_search(
+    decoded_ids, scores, _ = beam_search.beam_search(
         symbols_to_logits_fn,
         initial_ids,
         beam_size,
@@ -936,6 +936,7 @@ def fast_decode(encoder_output,
     force_decode_length: bool, whether to force the full decode length, or if
       False, stop when all beams hit eos_id.
     scope_prefix: str, prefix for decoder layer variable scopes.
+    cache: cache dictionary for additional predictions.
 
   Returns:
       A dict of decoding results {
@@ -959,7 +960,7 @@ def fast_decode(encoder_output,
       hparams.num_heads if hparams.get("attention_variables_3d") else 0)
 
   if cache is None:
-    cache = dict()
+    cache = {}
   cache.update({
       "layer_%d" % layer: {
           "k":
diff --git a/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index bddce3fa9..15a4dd266 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -751,6 +751,13 @@ def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
     return tf.logical_and(
         tf.less(i, decode_length), tf.logical_not(bound_is_met))
 
+  inner_shape = tf.TensorShape([None, None, None])
+  if use_tpu:
+    inner_shape = tf.TensorShape([batch_size, beam_size, decode_length + 1])
+  if use_tpu:
+    state_struc = nest.map_structure(lambda state: state.get_shape(), states)
+  else:
+    state_struc = nest.map_structure(get_state_shape_invariants, states)
   (_, alive_seq, alive_log_probs, finished_seq, finished_scores,
    finished_flags, states) = tf.while_loop(
        _is_finished,
@@ -760,16 +767,12 @@ def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
        ],
        shape_invariants=[
            tf.TensorShape([]),
-           (tf.TensorShape([batch_size, beam_size, decode_length + 1])
-            if use_tpu else tf.TensorShape([None, None, None])),
+           inner_shape,
            alive_log_probs.get_shape(),
-           (tf.TensorShape([batch_size, beam_size, decode_length + 1])
-            if use_tpu else tf.TensorShape([None, None, None])),
+           inner_shape,
            finished_scores.get_shape(),
            finished_flags.get_shape(),
-           (nest.map_structure(lambda state: state.get_shape(), states)
-            if use_tpu else
-            nest.map_structure(get_state_shape_invariants, states)),
+           state_struc
        ],
        parallel_iterations=1,
        back_prop=False)
diff --git a/tensor2tensor/utils/beam_search_test.py b/tensor2tensor/utils/beam_search_test.py
index 1b11d344c..7f6aa3595 100644
--- a/tensor2tensor/utils/beam_search_test.py
+++ b/tensor2tensor/utils/beam_search_test.py
@@ -182,7 +182,7 @@ def symbols_to_logits(ids):
       logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
       return logits
 
-    final_ids, final_probs = beam_search.beam_search(
+    final_ids, final_probs, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
@@ -390,10 +390,10 @@ def symbols_to_logits(ids, _, states):
         0.0,
         eos_id=1,
         states=states)
-    
+
     with self.test_session() as sess:
       final_states = sess.run(final_states)
-    self.assertAllEqual([[1]], final_states["state"])
+    self.assertAllEqual([[[2]]], final_states["state"])
 
   def testStateBeamTwo(self):
     batch_size = 1
@@ -476,7 +476,7 @@ def symbols_to_logits(_, i, states):
     states["state"] = tf.placeholder_with_default(
         states["state"], shape=(None, 1))
 
-    final_ids, _ = beam_search.beam_search(
+    final_ids, _, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,

From 04f2aca472a68634cf65f267de515aa25d971f98 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 10 Jan 2019 19:41:13 +0100
Subject: [PATCH 1475/2720] Out-of-graph policy evaluation (#1358)

---
 tensor2tensor/rl/evaluator.py | 23 ++++++---
 tensor2tensor/rl/rl_utils.py  | 97 ++++++++++++++++++++++++++++++++++-
 2 files changed, 113 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 3dea32510..cd415e9ab 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -56,19 +56,30 @@
 )
 
 
-def make_agent(agent_type, action_space):
+def make_agent(
+    agent_type, env, policy_hparams, policy_dir, sampling_temp
+):
   """Factory function for Agents."""
   return {
-      "random": rl_utils.RandomAgent,
-  }[agent_type](action_space)
+      "random": lambda: rl_utils.RandomAgent(
+          env.batch_size, env.observation_space, env.action_space
+      ),
+      "policy": lambda: rl_utils.PolicyAgent(
+          env.batch_size, env.observation_space, env.action_space,
+          policy_hparams, policy_dir, sampling_temp
+      ),
+  }[agent_type]()
 
 
 def make_eval_fn_with_agent(agent_type):
   """Returns an out-of-graph eval_fn using the Agent API."""
   def eval_fn(env, hparams, policy_hparams, policy_dir, sampling_temp):
     """Eval function."""
-    del hparams, policy_hparams, policy_dir, sampling_temp
-    agent = make_agent(agent_type, env.action_space)
+    base_env = env
+    env = rl_utils.BatchStackWrapper(env, hparams.frame_stack_size)
+    agent = make_agent(
+        agent_type, env, policy_hparams, policy_dir, sampling_temp
+    )
     num_dones = 0
     first_dones = [False] * env.batch_size
     observations = env.reset()
@@ -89,7 +100,7 @@ def eval_fn(env, hparams, policy_hparams, policy_dir, sampling_temp):
         for (i, observation) in zip(now_done_indices, reset_observations):
           observations[i] = observation
       observations = np.array(observations)
-    assert len(env.current_epoch_rollouts()) == env.batch_size
+    assert len(base_env.current_epoch_rollouts()) == env.batch_size
   return eval_fn
 
 
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 30d43665c..f7ccb5f92 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -21,10 +21,12 @@
 
 import random
 
+from gym.spaces import Box
 import numpy as np
 import six
 
 from tensor2tensor.data_generators.gym_env import T2TGymEnv
+from tensor2tensor.layers import common_layers
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl.dopamine_connector import DQNLearner
 from tensor2tensor.rl.ppo_learner import PPOLearner
@@ -202,7 +204,9 @@ class BatchAgent(object):
   Runs a batch of parallel agents. Operates on Numpy arrays.
   """
 
-  def __init__(self, action_space):
+  def __init__(self, batch_size, observation_space, action_space):
+    self.batch_size = batch_size
+    self.observation_space = observation_space
     self.action_space = action_space
 
   def act(self, observations):
@@ -240,3 +244,94 @@ def act(self, observations):
 
   def estimate_value(self, observations):
     return np.zeros(observations.shape[0])
+
+
+class PolicyAgent(BatchAgent):
+  """Agent based on a policy network."""
+
+  def __init__(
+      self, batch_size, observation_space, action_space, policy_hparams,
+      policy_dir, sampling_temp
+  ):
+    super(PolicyAgent, self).__init__(
+        batch_size, observation_space, action_space
+    )
+    self._sampling_temp = sampling_temp
+    with tf.Graph().as_default():
+      self._observations_t = tf.placeholder(
+          shape=((batch_size,) + self.observation_space.shape),
+          dtype=self.observation_space.dtype
+      )
+      (logits, self._values_t) = rl.get_policy(
+          self._observations_t, policy_hparams, self.action_space
+      )
+      actions = common_layers.sample_with_temperature(logits, sampling_temp)
+      self._actions_t = tf.cast(actions, tf.int32)
+      model_saver = tf.train.Saver(
+          tf.global_variables(policy_hparams.policy_network + "/.*")  # pylint: disable=unexpected-keyword-arg
+      )
+      self._sess = tf.Session()
+      self._sess.run(tf.global_variables_initializer())
+      trainer_lib.restore_checkpoint(policy_dir, model_saver, self._sess)
+
+  def _run(self, observations):
+    return self._sess.run(
+        [self._actions_t, self._values_t],
+        feed_dict={self._observations_t: observations}
+    )
+
+  def act(self, observations):
+    (actions, _) = self._run(observations)
+    return actions
+
+  def estimate_value(self, observations):
+    (_, values) = self._run(observations)
+    return values
+
+
+# TODO(koz4k): Unify interfaces of batch envs.
+class BatchStackWrapper(object):
+  """Out-of-graph batch stack wrapper."""
+
+  def __init__(self, env, stack_size):
+    self.env = env
+    self.stack_size = stack_size
+    inner_space = env.observation_space
+    self.observation_space = Box(
+        low=np.array([inner_space.low] * self.stack_size),
+        high=np.array([inner_space.high] * self.stack_size),
+        dtype=inner_space.dtype,
+    )
+    self._history_buffer = np.zeros(
+        (self.batch_size,) + self.observation_space.shape
+    )
+
+  @property
+  def batch_size(self):
+    return self.env.batch_size
+
+  @property
+  def action_space(self):
+    return self.env.action_space
+
+  @property
+  def reward_range(self):
+    return self.env.reward_range
+
+  def reset(self, indices=None):
+    if indices is None:
+      indices = range(self.batch_size)
+
+    observations = self.env.reset(indices)
+    for (index, observation) in zip(indices, observations):
+      self._history_buffer[index, ...] = [observation] * self.stack_size
+    return self._history_buffer
+
+  def step(self, actions):
+    (observations, rewards, dones) = self.env.step(actions)
+    self._history_buffer = np.roll(self._history_buffer, shift=-1, axis=1)
+    self._history_buffer[:, -1, ...] = observations
+    return (self._history_buffer, rewards, dones)
+
+  def close(self):
+    self.env.close()

From 681829587860ea20ae7e230a30153604cc834068 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 10 Jan 2019 10:41:30 -0800
Subject: [PATCH 1476/2720] internal merge of PR #1358

PiperOrigin-RevId: 228735764
---
 tensor2tensor/rl/evaluator.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index cd415e9ab..fbbc06e38 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -61,10 +61,10 @@ def make_agent(
 ):
   """Factory function for Agents."""
   return {
-      "random": lambda: rl_utils.RandomAgent(
+      "random": lambda: rl_utils.RandomAgent(  # pylint: disable=g-long-lambda
           env.batch_size, env.observation_space, env.action_space
       ),
-      "policy": lambda: rl_utils.PolicyAgent(
+      "policy": lambda: rl_utils.PolicyAgent(  # pylint: disable=g-long-lambda
           env.batch_size, env.observation_space, env.action_space,
           policy_hparams, policy_dir, sampling_temp
       ),

From 3e8fe2e1f1069af35b46004482f01b2037a521b1 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 10 Jan 2019 11:09:34 -0800
Subject: [PATCH 1477/2720] Corrections and tuning in the model-based RL code.

PiperOrigin-RevId: 228741204
---
 tensor2tensor/data_generators/gym_env.py      |  8 ++++++++
 .../models/video/basic_stochastic.py          |  9 ++++++---
 tensor2tensor/rl/envs/simulated_batch_env.py  |  3 +--
 tensor2tensor/rl/ppo_learner.py               |  6 +++---
 .../rl/trainer_model_based_params.py          | 19 +++++++++++++------
 5 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 6b2a1c65e..6d6d27725 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -746,6 +746,14 @@ def frame_width(self):
     "up_n_down", "video_pinball", "yars_revenge",
 ]
 
+
+# Blacklist a few games where it makes little sense to run on for now.
+ATARI_GAMES_WITH_HUMAN_SCORE_NICE = [
+    g for g in ATARI_GAMES_WITH_HUMAN_SCORE if g not in [
+        "solaris", "pitfall", "montezuma_revenge", "enduro",
+        "video_pinball", "double_dunk"]]
+
+
 ATARI_WHITELIST_GAMES = [
     "amidar",
     "bank_heist",
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 9f6d50114..5d3d7b7ed 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -75,7 +75,8 @@ def init_internal_states(self):
       return None
     # Hardcoded frame shapes.
     max_batch_size = max(64, self.hparams.batch_size)
-    shape = [max_batch_size] + self.hparams.problem.frame_shape[:-1] + [32]
+    shape = [max_batch_size] + self.hparams.problem.frame_shape[:-1] + [
+        self.hparams.recurrent_state_size]
     with tf.variable_scope("clean_scope_for_internal_state"):
       v = tf.get_variable("state", shape, trainable=False,
                           initializer=tf.zeros_initializer())
@@ -106,7 +107,8 @@ def update_internal_states_early(self, internal_states, frames):
     internal_state = internal_states[0][0][:batch_size, :, :, :]
     state_activation = tf.concat([internal_state, frames[0]], axis=-1)
     state_gate_candidate = tf.layers.conv2d(
-        state_activation, 64, (3, 3), padding="SAME", name="state_conv")
+        state_activation, 2 * self.hparams.recurrent_state_size,
+        (3, 3), padding="SAME", name="state_conv")
     state_gate, state_candidate = tf.split(state_gate_candidate, 2, axis=-1)
     state_gate = tf.nn.sigmoid(state_gate)
     state_candidate = tf.tanh(state_candidate)
@@ -267,11 +269,12 @@ def next_frame_basic_stochastic_discrete():
   hparams.add_hparam("discretize_warmup_steps", 40000)
   hparams.add_hparam("latent_rnn_warmup_steps", 40000)
   hparams.add_hparam("latent_rnn_max_sampling", 0.5)
-  hparams.add_hparam("latent_use_max_probability", 0.7)
+  hparams.add_hparam("latent_use_max_probability", 0.8)
   hparams.add_hparam("full_latent_tower", False)
   hparams.add_hparam("latent_predictor_state_size", 128)
   hparams.add_hparam("latent_predictor_temperature", 1.0)
   hparams.add_hparam("complex_addn", True)
+  hparams.add_hparam("recurrent_state_size", 64)
   return hparams
 
 
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index f5a46e64b..c9a7dc61a 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -161,8 +161,7 @@ def initialize(self, sess):
     model_loader = tf.train.Saver(
         var_list=tf.global_variables(scope="next_frame*")  # pylint:disable=unexpected-keyword-arg
     )
-    # TODO(afrozm): use TF methods to be on the safe side here.
-    if os.path.isdir(self._model_dir):
+    if tf.io.gfile.isdir(self._model_dir):
       trainer_lib.restore_checkpoint(
           self._model_dir, saver=model_loader, sess=sess, must_restore=True
       )
diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 4077b7cd3..457e18da1 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -134,8 +134,8 @@ def evaluate(self, env_fn, hparams, sampling_temp):
             sampling_temp=sampling_temp,
         )
         model_saver = tf.train.Saver(
-            tf.global_variables(hparams.policy_network + "/.*") +
-            tf.global_variables("clean_scope.*")  # Needed for the SD model.
+            tf.global_variables(hparams.policy_network + "/.*")
+            # tf.global_variables("clean_scope.*")  # Needed for sharing params.
         )
 
         with tf.Session() as sess:
@@ -199,7 +199,7 @@ def _run_train(ppo_hparams,
   model_saver = tf.train.Saver(
       tf.global_variables(ppo_hparams.policy_network + "/.*") +
       tf.global_variables("training/" + ppo_hparams.policy_network + "/.*") +
-      tf.global_variables("clean_scope.*") +  # Needed for the SD model.
+      # tf.global_variables("clean_scope.*") +  # Needed for sharing params.
       tf.global_variables("global_step") +
       tf.global_variables("losses_avg.*") +
       tf.global_variables("train_stats.*")
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 463eac8cf..3a7815731 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -306,6 +306,8 @@ def rlmb_base_stochastic_discrete():
   hparams.grayscale = False
   hparams.generative_model = "next_frame_basic_stochastic_discrete"
   hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  hparams.ppo_epoch_length = 50
+  hparams.simulated_rollout_length = 50
   return hparams
 
 
@@ -321,12 +323,9 @@ def rlmb_base_stochastic_discrete_param_sharing():
 @registry.register_hparams
 def rlmb_long_stochastic_discrete():
   """Long setting with stochastic discrete model."""
-  hparams = rlmb_base()
-  hparams.learning_rate_bump = 1.0
-  hparams.grayscale = False
-  hparams.generative_model = "next_frame_basic_stochastic_discrete"
+  hparams = rlmb_base_stochastic_discrete()
   hparams.generative_model_params = "next_frame_basic_stochastic_discrete_long"
-  hparams.ppo_epochs_num = 2000
+  hparams.ppo_epochs_num = 1000
   return hparams
 
 
@@ -697,7 +696,15 @@ def rlmb_whitelisted_games(rhp):
 @registry.register_ranged_hparams
 def rlmb_human_score_games(rhp):
   rhp.set_categorical("loop.game",
-                      gym_env.ATARI_GAMES_WITH_HUMAN_SCORE)
+                      gym_env.ATARI_GAMES_WITH_HUMAN_SCORE_NICE)
+  rhp.set_discrete("model.moe_loss_coef", list(range(5)))
+
+
+@registry.register_ranged_hparams
+def rlmb_human_score_games_v100unfriendly(rhp):
+  """Games that for strange reasons often fail on v100s but work on p100s."""
+  rhp.set_categorical("loop.game",
+                      ["chopper_command", "boxing", "asterix", "seaquest"])
   rhp.set_discrete("model.moe_loss_coef", list(range(5)))
 
 

From 0dff89d64c3406d42717280cb9135a5ce7af793c Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 10 Jan 2019 11:36:10 -0800
Subject: [PATCH 1478/2720] Correct sampling in generating sub-sampled
 translation data, use tiny in multi-problem.

PiperOrigin-RevId: 228746388
---
 .../data_generators/translate_enro.py         | 40 ++++++++++++++++++-
 .../data_generators/wiki_multi_problems.py    |  4 +-
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/translate_enro.py b/tensor2tensor/data_generators/translate_enro.py
index 238d79a61..9b8d144a7 100644
--- a/tensor2tensor/data_generators/translate_enro.py
+++ b/tensor2tensor/data_generators/translate_enro.py
@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import random
+
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.data_generators import translate
@@ -99,12 +101,48 @@ def dataset_splits(self):
   def vocab_filename(self):
     return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
 
+  @property
+  def how_many_examples_to_sample(self):
+    return 6000
+
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
     """Generate just the first 6k samples for training."""
+    # If not training, do the same as before.
+    if dataset_split != problem.DatasetSplit.TRAIN:
+      for x in super(TranslateEnroWmtMultiSmall64k, self).generate_samples(
+          data_dir, tmp_dir, dataset_split):
+        yield x
+      raise StopIteration
+    # Now we assume we're training.
     counter = 0
+    # The size of this data-set in total is around 614K, we want to sample so
+    # that in expectation we take the requested number of samples in 1 go.
+    sample_prob = self.how_many_examples_to_sample / float(614000)
+    # Let's sample.
     for x in super(TranslateEnroWmtMultiSmall64k, self).generate_samples(
         data_dir, tmp_dir, dataset_split):
+      if random.random() > sample_prob:
+        continue
       counter += 1
-      if counter > 6000 and dataset_split == problem.DatasetSplit.TRAIN:
+      if counter > self.how_many_examples_to_sample:
         raise StopIteration
       yield x
+    # We do it again if we don't have enough samples.
+    if counter < self.how_many_examples_to_sample:
+      for x in super(TranslateEnroWmtMultiSmall64k, self).generate_samples(
+          data_dir, tmp_dir, dataset_split):
+        if random.random() > sample_prob:
+          continue
+        counter += 1
+        if counter > self.how_many_examples_to_sample:
+          raise StopIteration
+        yield x
+
+
+@registry.register_problem
+class TranslateEnroWmtMultiTiny64k(TranslateEnroWmtMultiSmall64k):
+  """Translation with muli-lingual vocabulary, tiny (6K) training data."""
+
+  @property
+  def how_many_examples_to_sample(self):
+    return 600
diff --git a/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index 5283fec94..87b146c1a 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -134,12 +134,12 @@ def __init__(self, was_reversed=False, was_copy=False):
     self.task_list.append(wiki_lm.LanguagemodelDeEnFrRoWiki64k())
     self.task_list.append(translate_ende.TranslateEndeWmtMulti64k())
     self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k())
-    self.task_list.append(translate_enro.TranslateEnroWmtMultiSmall64k())
+    self.task_list.append(translate_enro.TranslateEnroWmtMultiTiny64k())
     self.task_list.append(translate_ende.TranslateEndeWmtMulti64k(
         was_reversed=True))
     self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k(
         was_reversed=True))
-    self.task_list.append(translate_enro.TranslateEnroWmtMultiSmall64k(
+    self.task_list.append(translate_enro.TranslateEnroWmtMultiTiny64k(
         was_reversed=True))
     self.task_list.append(
         cnn_dailymail.SummarizeCnnDailymailWikiLMMultiVocab64k())

From 6097be464ad98ef27835f020b9030eee1748392a Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Thu, 10 Jan 2019 23:21:55 +0100
Subject: [PATCH 1479/2720] Model-Based RL: Player - run real and simulated env
 side-by-side  (#1359)

* Enable BatchGymEnv to load directly from checkpoint file.

* Player move from Wrapper to Env, refactor.

* Further refactor of PlayerEnv, initializable SimulatedGymEnv.

* Initial Player for simulated and real environment comparison.

* Add restarting simulated env only for Player; refactor main, fix SingleEnvPlayer restart.

* Rename Player methods, add documentation, run pylint.

* Use common function from rl_utils to compute difference between images in trainer_model_based.py.

* Fix epoch choice passing to simulated env in player main.

* Add frame counter to Player display.

* Re-enable infering game name from filenames.

* Correct imports for Python 2 and style guide.
---
 tensor2tensor/rl/player.py              | 441 ++++++++++++++++++++----
 tensor2tensor/rl/player_utils.py        | 116 ++++++-
 tensor2tensor/rl/rl_utils.py            |  15 +
 tensor2tensor/rl/trainer_model_based.py |   5 +-
 4 files changed, 483 insertions(+), 94 deletions(-)

diff --git a/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 0dc680093..77e6e67d1 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -20,6 +20,7 @@
   R key to reset env.
   C key to toggle WAIT mode.
   N to perform NOOP action under WAIT mode.
+  X to reset simulated env only, when running sim-real comparison.
 
 Run this script with the same parameters as trainer_model_based.py. Note that
 values of most of them have no effect on player, so running just
@@ -35,7 +36,9 @@
 python -m tensor2tensor/rl/record_ppo.py \
     --output_dir=path/to/your/experiment \
     --loop_hparams_set=rlmb_base \
-    --loop_hparams=game=<right game in case of problems> \
+    --sim_and_real=False \
+    --simulated_env=False \
+    --loop_hparams=generative_model="next_frame" \
     --video_dir=my/video/dir \
     --zoom=6 \
     --fps=50 \
@@ -60,9 +63,11 @@
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw
 from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
+from tensor2tensor.rl.rl_utils import absolute_hinge_difference
 # Import flags from t2t_trainer and trainer_model_based
 import tensor2tensor.rl.trainer_model_based_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
+
 import tensorflow as tf
 
 
@@ -77,6 +82,8 @@
                    "Frames per second.")
 flags.DEFINE_string("epoch", "last",
                     "Data from which epoch to use.")
+flags.DEFINE_boolean("sim_and_real", True,
+                     "Compare simulated and real environment.")
 flags.DEFINE_boolean("simulated_env", True,
                      "Either to use 'simulated' or 'real' env.")
 flags.DEFINE_boolean("dry_run", False,
@@ -92,42 +99,74 @@
 flags.DEFINE_string("episodes_data_dir", "",
                     "Path to data for simulated environment initialization. "
                     "Inferred from output_dir if empty.")
-
-
-class PlayerEnvWrapper(gym.Wrapper):
-  """Environment Wrapper for gym.utils.play."""
-
-  RESET_ACTION = 101
+flags.DEFINE_boolean("game_from_filenames", True,
+                     "If infer game name from data_dir filenames or from "
+                     "hparams.")
+
+class PlayerEnv(gym.Env):
+  """Base (abstract) environment for interactive human play with gym.utils.play.
+
+  Additionally to normal actions passed to underlying environment(s) it
+  allows to pass special actions by `step` method.
+
+  Special actions:
+    RETURN_DONE_ACTION: Returns done from `step` to force gym.utils.play to
+      call reset.
+    TOGGLE_WAIT_ACTION: Change between real-time-play and wait-for-pressed-key
+      modes.
+    WAIT_MODE_NOOP_ACTION: perform noop action (when wait-for-pressed-key mode
+    is on)
+
+  For keyboard keys related to actions above see `get_keys_to_action` method.
+
+  Naming conventions:
+    envs_step_tuples: Dictionary of tuples similar to these returned by
+      gym.Env.step().
+      {
+        "env_name": (observation, reward, done, info),
+        ...
+      }
+      Keys depend on subclass.
+  """
+
+  # Integers (as taken by step() method) related to special actions.
+  RETURN_DONE_ACTION = 101
   TOGGLE_WAIT_ACTION = 102
   WAIT_MODE_NOOP_ACTION = 103
 
-  HEADER_HEIGHT = 12
-
-  def __init__(self, env):
-    super(PlayerEnvWrapper, self).__init__(env)
-
-    # Set observation space
-    orig = self.env.observation_space
-    shape = tuple([orig.shape[0] + self.HEADER_HEIGHT] + list(orig.shape[1:]))
-    self.observation_space = gym.spaces.Box(low=orig.low.min(),
-                                            high=orig.high.max(),
-                                            shape=shape, dtype=orig.dtype)
-
-    # gym play() looks for get_keys_to_action() only on top and bottom level
-    # of env and wrappers stack.
-    self.unwrapped.get_keys_to_action = self.get_keys_to_action
+  HEADER_HEIGHT = 27
 
+  def __init__(self):
     self._wait = True
+    # If action_space will be needed, one could use e.g. gym.spaces.Dict.
+    self.action_space = None
+    self._last_step_tuples = None
+
+  def _init_action_mappings(self, env):
+    # Atari dependant. In case of problems with keyboard key interpretation
+    # switch to _action_set instead of range(env.action_space.n) (similarly to
+    # how gym AtariEnv does). _action_set can probably be obtain from full
+    # game name.
     self.action_meaning = {i: ACTION_MEANING[i]
-                           for i in range(self.action_space.n)}
+                           for i in range(env.action_space.n)}
     self.name_to_action_num = {v: k for k, v in
                                six.iteritems(self.action_meaning)}
 
-  def get_action_meanings(self):
-    return [self.action_meaning[i] for i in range(self.action_space.n)]
+  def _get_action_meanings(self):
+    return [self.action_meaning[i] for i in range(len(self.action_meaning))]
 
   def get_keys_to_action(self):
-    # Based on gym atari.py AtariEnv.get_keys_to_action()
+    """Get mapping from keyboard keys to actions.
+
+    Required by gym.utils.play in environment or top level wrapper.
+
+    Returns:
+      {
+        Unicode code point for keyboard key: action (formatted for step()),
+        ...
+      }
+    """
+    # Based on gym AtariEnv.get_keys_to_action()
     keyword_to_key = {
         "UP": ord("w"),
         "DOWN": ord("s"),
@@ -138,7 +177,7 @@ def get_keys_to_action(self):
 
     keys_to_action = {}
 
-    for action_id, action_meaning in enumerate(self.get_action_meanings()):
+    for action_id, action_meaning in enumerate(self._get_action_meanings()):
       keys = []
       for keyword, key in keyword_to_key.items():
         if keyword in action_meaning:
@@ -148,59 +187,72 @@ def get_keys_to_action(self):
       assert keys_tuple not in keys_to_action
       keys_to_action[keys_tuple] = action_id
 
-    # Add utility actions
-    keys_to_action[(ord("r"),)] = self.RESET_ACTION
+    # Special actions:
+    keys_to_action[(ord("r"),)] = self.RETURN_DONE_ACTION
     keys_to_action[(ord("c"),)] = self.TOGGLE_WAIT_ACTION
     keys_to_action[(ord("n"),)] = self.WAIT_MODE_NOOP_ACTION
 
     return keys_to_action
 
-  def step(self, action):
-    # Special codes
-    if action == self.TOGGLE_WAIT_ACTION:
-      self._wait = not self._wait
-      ob, reward, done, info = self._last_step
-      ob = self.augment_observation(ob, reward, self.total_reward)
-      return ob, reward, done, info
-
-    if action == self.RESET_ACTION:
-      ob = self.empty_observation()
-      return ob, 0, True, {}
-
-    if self._wait and action == self.name_to_action_num["NOOP"]:
-      ob, reward, done, info = self._last_step
-      ob = self.augment_observation(ob, reward, self.total_reward)
-      return ob, reward, done, info
-
-    if action == self.WAIT_MODE_NOOP_ACTION:
-      action = self.name_to_action_num["NOOP"]
+  def _player_actions(self):
+    return {
+        self.RETURN_DONE_ACTION: self._player_return_done_action,
+        self.TOGGLE_WAIT_ACTION: self._player_toggle_wait_action,
+    }
 
-    ob, reward, done, info = self.env.step(action)
-    self._last_step = ob, reward, done, info
+  def _player_toggle_wait_action(self):
+    self._wait = not self._wait
+    return self._last_step_tuples
 
-    self.total_reward += reward
+  def step(self, action):
+    """Pass action to underlying environment(s) or perform special action.
 
-    ob = self.augment_observation(ob, reward, self.total_reward)
+    For returned tuple explanation see _player_step_tuple() in subclasses.
+    """
+    # Special codes
+    if action in self._player_actions():
+      envs_step_tuples = self._player_actions()[action]()
+    elif self._wait and action == self.name_to_action_num["NOOP"]:
+      # Ignore no-op, do not pass to environment.
+      envs_step_tuples = self._last_step_tuples
+    else:
+      # Run action on environment(s).
+      if action == self.WAIT_MODE_NOOP_ACTION:
+        action = self.name_to_action_num["NOOP"]
+      # Perform action on underlying environment(s).
+      envs_step_tuples = self._step_envs(action)
+      self._update_statistics(envs_step_tuples)
+
+    self._last_step_tuples = envs_step_tuples
+    ob, reward, done, info = self._player_step_tuple(envs_step_tuples)
     return ob, reward, done, info
 
-  def reset(self):
-    ob = self.env.reset()
-    self._last_step = ob, 0, False, {}
-    self.total_reward = 0
-    return self.augment_observation(ob, 0, self.total_reward)
+  def _augment_observation(self, ob, reward, cumulative_reward):
+    """"Expand observation array with additional information header (top rows).
 
-  def empty_observation(self):
-    return np.zeros(self.observation_space.shape)
+    Args:
+      ob: observation
+      reward: reward to be included in header.
+      cumulative_reward: total cumulated reward to be included in header.
 
-  def augment_observation(self, ob, reward, total_reward):
+    Returns:
+      Expanded observation array.
+    """
     img = PIL_Image().new("RGB",
-                          (ob.shape[1], PlayerEnvWrapper.HEADER_HEIGHT,))
+                          (ob.shape[1], self.HEADER_HEIGHT,))
     draw = PIL_ImageDraw().Draw(img)
-    draw.text((1, 0), "c:{:3}, r:{:3}".format(int(total_reward), int(reward)),
-              fill=(255, 0, 0))
+    draw.text(
+        (1, 0), "c:{:3}, r:{:3}".format(int(cumulative_reward), int(reward)),
+        fill=(255, 0, 0)
+    )
+    draw.text(
+        (1, 15), "fc:{:3}".format(int(self._frame_counter)),
+        fill=(255, 0, 0)
+    )
     header = np.asarray(img)
     del img
     header.setflags(write=1)
+    # Top row color indicates if WAIT MODE is on.
     if self._wait:
       pixel_fill = (0, 255, 0)
     else:
@@ -208,6 +260,232 @@ def augment_observation(self, ob, reward, total_reward):
     header[0, :, :] = pixel_fill
     return np.concatenate([header, ob], axis=0)
 
+  def reset(self):
+    raise NotImplementedError
+
+  def _step_envs(self, action):
+    """Perform action on underlying environment(s)."""
+    raise NotImplementedError
+
+  def _update_statistics(self, envs_step_tuples):
+    """Update underlying environment(s) total cumulative rewards."""
+    raise NotImplementedError
+
+  def _player_return_done_action(self):
+    """
+
+    Returns:
+       envs_step_tuples: such that `player_step_tuple(envs_step_tuples)`
+        will return done."""
+    raise NotImplementedError
+
+  def _player_step_tuple(self, envs_step_tuples):
+    """Infer return tuple for step() given underlying environment(s) tuple(s)"""
+    raise NotImplementedError
+
+
+class SimAndRealEnvPlayer(PlayerEnv):
+  """Run simulated and real env side-by-side for comparison.
+
+  Displays three windows - one for real environment, second for simulated
+  and third for their differences.
+
+  Normal actions are passed to both environments.
+
+  Special Actions:
+    RESTART_SIMULATED_ENV_ACTION: restart simulated environment only, using
+      current frames from real environment.
+    See `PlayerEnv` for rest of special actions.
+
+  Naming conventions:
+    envs_step_tuples: dictionary with two keys.
+    {
+      "real_env": (observation, reward, done, info),
+      "sim_env": (observation, reward, done, info)
+    }
+  """
+
+  RESTART_SIMULATED_ENV_ACTION = 110
+
+  def __init__(self, real_env, sim_env):
+    """
+
+    Args:
+      real_env: real environment such as `FlatBatchEnv<T2TGymEnv>`.
+      sim_env: simulation of `real_env` to be compared with. E.g.
+        `SimulatedGymEnv` must allow to update initial frames for next reset
+        with `add_to_initial_stack` method.
+    """
+    super(SimAndRealEnvPlayer, self).__init__()
+    assert real_env.observation_space.shape == sim_env.observation_space.shape
+    self.real_env = real_env
+    self.sim_env = sim_env
+    orig = self.real_env.observation_space
+    # Observation consists three side-to-side images - simulated environment
+    # observation, real environment observation and difference between these
+    # two.
+    shape = (orig.shape[0] + self.HEADER_HEIGHT, orig.shape[1] * 3,
+             orig.shape[2])
+
+    self.observation_space = gym.spaces.Box(low=orig.low.min(),
+                                            high=orig.high.max(),
+                                            shape=shape, dtype=orig.dtype)
+    self._init_action_mappings(sim_env)
+
+  def _player_actions(self):
+    actions = super(SimAndRealEnvPlayer, self)._player_actions()
+    actions.update({
+        self.RESTART_SIMULATED_ENV_ACTION:
+            self.player_restart_simulated_env_action,
+    })
+    return actions
+
+  def get_keys_to_action(self):
+    keys_to_action = super(SimAndRealEnvPlayer, self).get_keys_to_action()
+    keys_to_action[(ord("x"),)] = self.RESTART_SIMULATED_ENV_ACTION
+    return keys_to_action
+
+  def _player_step_tuple(self, envs_step_tuples):
+    """Construct observation, return usual step tuple.
+
+    Returns:
+      Step tuple: ob, reward, done, info
+        ob: concatenated images [simulated observation, real observation,
+          difference], with additional informations in header.
+        reward: real environment reward
+        done: True iff. envs_step_tuples['real_env'][2] is True
+        info: real environment info
+    """
+    ob_real, reward_real, _, _ = envs_step_tuples['real_env']
+    ob_sim, reward_sim, _, _ = envs_step_tuples['sim_env']
+    ob_err = absolute_hinge_difference(ob_sim, ob_real)
+
+    ob_real_aug = self._augment_observation(ob_real, reward_real,
+                                            self.cumulative_real_reward)
+    ob_sim_aug = self._augment_observation(ob_sim, reward_sim,
+                                           self.cumulative_sim_reward)
+    ob_err_aug = self._augment_observation(
+        ob_err, reward_sim - reward_real,
+        self.cumulative_sim_reward - self.cumulative_real_reward
+    )
+    ob = np.concatenate([ob_sim_aug, ob_real_aug, ob_err_aug], axis=1)
+    _, reward, done, info = envs_step_tuples['real_env']
+    return ob, reward, done, info
+
+  def reset(self):
+    """Reset simulated and real environments."""
+    self._frame_counter = 0
+    ob_real = self.real_env.reset()
+    # Initialize simulated environment with frames from real one.
+    self.sim_env.add_to_initial_stack(ob_real)
+    for _ in range(3):
+      ob_real, _, _, _ = self.real_env.step(self.name_to_action_num['NOOP'])
+      self.sim_env.add_to_initial_stack(ob_real)
+    # TODO(konradczechowski): remove when not longer needed.
+    # for i in range(12):
+    #   ob_real, _, _, _ = self.real_env.step(np.random.choice([2,3, 4, 5]))
+    #   self.sim_env.add_to_initial_stack(ob_real)
+    ob_sim = self.sim_env.reset()
+    assert np.all(ob_real == ob_sim)
+    self._last_step_tuples = self._pack_step_tuples((ob_real, 0, False, {}),
+                                                    (ob_sim, 0, False, {}))
+    self.set_zero_cumulative_rewards()
+    ob, _, _, _ = self._player_step_tuple(self._last_step_tuples)
+    return ob
+
+  def _pack_step_tuples(self, real_env_step_tuple, sim_env_step_tuple):
+    return dict(real_env=real_env_step_tuple,
+                sim_env=sim_env_step_tuple)
+
+  def set_zero_cumulative_rewards(self):
+    self.cumulative_real_reward = 0
+    self.cumulative_sim_reward = 0
+
+  def _step_envs(self, action):
+    """Perform `step(action)` on both environments.
+
+    Update initial_frame_stack for simulated environment.
+    """
+    self._frame_counter += 1
+    real_env_step_tuple = self.real_env.step(action)
+    sim_env_step_tuple = self.sim_env.step(action)
+    self.sim_env.add_to_initial_stack(real_env_step_tuple[0])
+    return self._pack_step_tuples(real_env_step_tuple, sim_env_step_tuple)
+
+  def _update_statistics(self, envs_step_tuples):
+    self.cumulative_real_reward += envs_step_tuples['real_env'][1]
+    self.cumulative_sim_reward += envs_step_tuples['sim_env'][1]
+
+  def _player_return_done_action(self):
+    ob = np.zeros(self.real_env.observation_space.shape, dtype=np.uint8)
+    return self._pack_step_tuples((ob, 0, True, {}),
+                                  (ob, 0, True, {}))
+
+  def player_restart_simulated_env_action(self):
+    self._frame_counter = 0
+    ob = self.sim_env.reset()
+    # TODO(konradczechowski): remove when this will be not needed.
+    # new_ob, _, _, _ = self.sim_env.step(2)
+    # print("\n\n\n\ndiff {}\n\n\n\n".format((ob - new_ob).sum()))
+    # ##########
+    assert np.all(self._last_step_tuples['real_env'][0] == ob)
+    self.set_zero_cumulative_rewards()
+    return self._pack_step_tuples(
+        self._last_step_tuples['real_env'], (ob, 0, False, {}))
+
+
+class SingleEnvPlayer(PlayerEnv):
+  """"Play on single (simulated or real) environment.
+
+  See `PlayerEnv` for more details.
+
+  Naming conventions:
+    envs_step_tuples: dictionary with single key.
+      {
+        "env": (observation, reward, done, info),
+      }
+      Plural form used for consistency with `PlayerEnv`.
+  """
+
+  def __init__(self, env):
+    super(SingleEnvPlayer, self).__init__()
+    self.env = env
+    # Set observation space
+    orig = self.env.observation_space
+    shape = tuple([orig.shape[0] + self.HEADER_HEIGHT] + list(orig.shape[1:]))
+    self.observation_space = gym.spaces.Box(low=orig.low.min(),
+                                            high=orig.high.max(),
+                                            shape=shape, dtype=orig.dtype)
+    self._init_action_mappings(env)
+
+  def _player_step_tuple(self, envs_step_tuples):
+    """Augment observation, return usual step tuple."""
+    ob, reward, done, info = envs_step_tuples['env']
+    ob = self._augment_observation(ob, reward, self.cumulative_reward)
+    return ob, reward, done, info
+
+  def _pack_step_tuples(self, env_step_tuple):
+    return dict(env=env_step_tuple)
+
+  def reset(self):
+    self._frame_counter = 0
+    ob = self.env.reset()
+    self._last_step_tuples = self._pack_step_tuples((ob, 0, False, {}))
+    self.cumulative_reward = 0
+    return self._augment_observation(ob, 0, self.cumulative_reward)
+
+  def _step_envs(self, action):
+    self._frame_counter += 1
+    return self._pack_step_tuples(self.env.step(action))
+
+  def _update_statistics(self, envs_step_tuples):
+    _, reward, _, _ = envs_step_tuples['env']
+    self.cumulative_reward += reward
+
+  def _player_return_done_action(self):
+    ob = np.zeros(self.env.observation_space.shape, dtype=np.uint8)
+    return self._pack_step_tuples((ob, 0, True, {}))
+
 
 def main(_):
   # gym.logger.set_level(gym.logger.DEBUG)
@@ -221,19 +499,38 @@ def main(_):
       world_model=FLAGS.wm_dir,
       policy=FLAGS.policy_dir,
       data=FLAGS.episodes_data_dir)
+  if FLAGS.game_from_filenames:
+    hparams.set_hparam(
+        'game', player_utils.infer_game_name_from_filenames(directories['data'])
+    )
   epoch = FLAGS.epoch if FLAGS.epoch == "last" else int(FLAGS.epoch)
 
-  if FLAGS.simulated_env:
-    env = player_utils.load_data_and_make_simulated_env(
-        directories["data"], directories["world_model"],
-        hparams, which_epoch_data=epoch)
-  else:
+  def make_real_env():
     env = player_utils.setup_and_load_epoch(
         hparams, data_dir=directories["data"],
-        which_epoch_data=epoch)
-    env = FlatBatchEnv(env)
+        which_epoch_data=None)
+    env = FlatBatchEnv(env)  # pylint: disable=redefined-variable-type
+    return env
 
-  env = PlayerEnvWrapper(env)  # pylint: disable=redefined-variable-type
+  def make_simulated_env(setable_initial_frames, which_epoch_data):
+    env = player_utils.load_data_and_make_simulated_env(
+        directories["data"], directories["world_model"],
+        hparams, which_epoch_data=which_epoch_data,
+        setable_initial_frames=setable_initial_frames)
+    return env
+
+  if FLAGS.sim_and_real:
+    sim_env = make_simulated_env(
+        which_epoch_data=None, setable_initial_frames=True)
+    real_env = make_real_env()
+    env = SimAndRealEnvPlayer(real_env, sim_env)
+  else:
+    if FLAGS.simulated_env:
+      env = make_simulated_env(  # pylint: disable=redefined-variable-type
+          which_epoch_data=epoch, setable_initial_frames=False)
+    else:
+      env = make_real_env()
+    env = SingleEnvPlayer(env)  # pylint: disable=redefined-variable-type
 
   env = player_utils.wrap_with_monitor(env, FLAGS.video_dir)
 
@@ -242,7 +539,7 @@ def main(_):
       env.reset()
       for i in range(50):
         env.step(i % 3)
-      env.step(PlayerEnvWrapper.RESET_ACTION)  # reset
+      env.step(PlayerEnv.RETURN_DONE_ACTION)  # reset
     return
 
   play.play(env, zoom=FLAGS.zoom, fps=FLAGS.fps)
diff --git a/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index 1025816f6..000c18316 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -32,6 +32,8 @@
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
 from tensor2tensor.utils import trainer_lib
+from tensor2tensor.utils.misc_utils import camelcase_to_snakecase
+
 import tensorflow as tf
 
 
@@ -39,6 +41,84 @@
 FLAGS = flags.FLAGS
 
 
+class SimulatedGymEnv(gym.Env):
+  """Gym environment, running with world model.
+
+  Allows passing custom initial frames.
+
+  Examples:
+    Setup simulated env from some point of real rollout.
+      >>> sim_env = SimulatedGymEnv(setable_initial_frames=True, **kwargs)
+      >>> real_env = FlatBatchEnv(T2TGymEnv(...))
+      >>> while ...:
+      >>>   ob, _, _, _ = real_env.step(action)
+      >>>   sim_env.add_to_initial_stack(ob)
+      >>> sim_env.reset()
+      >>> # Continue sim_env rollout.
+  """
+
+  def __init__(self, real_env, world_model_dir, hparams, random_starts,
+               setable_initial_frames=False):
+    """
+
+    Args:
+       real_env: gym environment.
+       world_model_dir: path to world model checkpoint directory.
+       hparams: hparams for rlmb pipeline.
+       random_starts: if restart world model from random frames, or only
+         from initial ones (from beginning of episodes). Valid only when
+         `setable_initial_fames` set to False.
+       setable_initial_frames: if True, initial_frames for world model should be
+         set by `add_to_initial_stack`.
+    """
+
+    self._setable_initial_frames = setable_initial_frames
+
+    if self._setable_initial_frames:
+      real_obs_shape = real_env.observation_space.shape
+      shape = (1, hparams.frame_stack_size) + real_obs_shape
+      self._initial_frames = np.zeros(shape=shape, dtype=np.uint8)
+      def initial_frame_chooser(batch_size):
+        assert batch_size == 1
+        return self._initial_frames
+
+    else:
+      initial_frame_chooser = rl_utils.make_initial_frame_chooser(
+          real_env, hparams.frame_stack_size,
+          simulation_random_starts=random_starts,
+          simulation_flip_first_random_for_beginning=False
+      )
+    env_fn = make_simulated_env_fn_from_hparams(
+        real_env, hparams,
+        batch_size=1,
+        initial_frame_chooser=initial_frame_chooser,
+        model_dir=world_model_dir,
+    )
+
+    env = env_fn(in_graph=False)
+    self.env = FlatBatchEnv(env)
+
+    self.observation_space = self.env.observation_space
+    self.action_space = self.env.action_space
+
+  def reset(self):
+    return self.env.reset()
+
+  def step(self, action):
+    return self.env.step(action)
+
+  def add_to_initial_stack(self, frame):
+    """Adds new frame to (initial) frame stack, removes last one."""
+    if not self._setable_initial_frames:
+      raise ValueError(
+          "This instace does not allow to manually set initial frame stack.")
+    assert frame.shape == self._initial_frames.shape[2:], \
+        '{}, {}'.format(frame.shape, self._initial_frames.shape[:1])
+    initial_frames = np.roll(self._initial_frames, shift=-1, axis=1)
+    initial_frames[0, -1, ...] = frame
+    self._initial_frames = initial_frames
+
+
 def infer_last_epoch_num(data_dir):
   """Infer highest epoch number from file names in data_dir."""
   names = os.listdir(data_dir)
@@ -75,34 +155,32 @@ def setup_and_load_epoch(hparams, data_dir, which_epoch_data=None):
   return t2t_env
 
 
-def make_simulated_gym_env(real_env, world_model_dir, hparams, random_starts):
-  """Gym environment with world model."""
-  initial_frame_chooser = rl_utils.make_initial_frame_chooser(
-      real_env, hparams.frame_stack_size,
-      simulation_random_starts=random_starts,
-      simulation_flip_first_random_for_beginning=False
-  )
-  env_fn = make_simulated_env_fn_from_hparams(
-      real_env, hparams,
-      batch_size=1,
-      initial_frame_chooser=initial_frame_chooser,
-      model_dir=world_model_dir
-  )
-  env = env_fn(in_graph=False)
-  flat_env = FlatBatchEnv(env)
-  return flat_env
+def infer_game_name_from_filenames(data_dir, snake_case=True):
+  names = os.listdir(data_dir)
+  game_names = [re.findall(pattern=r"^Gym(.*)NoFrameskip", string=name)
+                for name in names]
+  assert game_names, "No data files found in {}".format(data_dir)
+  game_names = sum(game_names, [])
+  game_name = game_names[0]
+  assert all(game_name == other for other in game_names), \
+      "There are multiple different game names in {}".format(data_dir)
+  if snake_case:
+    game_name = camelcase_to_snakecase(game_name)
+  return game_name
 
 
 def load_data_and_make_simulated_env(
-    data_dir, wm_dir, hparams, which_epoch_data="last", random_starts=True
+    data_dir, wm_dir, hparams, which_epoch_data="last", random_starts=True,
+    setable_initial_frames=False
 ):
   hparams = copy.deepcopy(hparams)
   t2t_env = setup_and_load_epoch(
       hparams, data_dir=data_dir,
       which_epoch_data=which_epoch_data)
-  return make_simulated_gym_env(
+  return SimulatedGymEnv(
       t2t_env, world_model_dir=wm_dir,
-      hparams=hparams, random_starts=random_starts)
+      hparams=hparams, random_starts=random_starts,
+      setable_initial_frames=setable_initial_frames)
 
 
 class ExtendToEvenDimentions(gym.ObservationWrapper):
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index f7ccb5f92..56214e73e 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -198,6 +198,21 @@ def initial_frame_chooser(batch_size):
   return initial_frame_chooser
 
 
+def absolute_hinge_difference(arr1, arr2, min_diff=10, dtype=np.uint8):
+  """Point-wise, hinge loss-like, difference between arrays.
+
+  Args:
+    arr1, arr2: integer arrays to compare.
+    min_diff: minimal difference taken into consideration.
+    dtype: dtype of returned array.
+
+  Returns:
+    array
+  """
+  diff = np.abs(arr1.astype(np.int) - arr2, dtype=np.int)
+  return np.maximum(diff - min_diff, 0).astype(dtype)
+
+
 class BatchAgent(object):
   """Python API for agents.
 
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 7a55aaad0..21188bb9e 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -45,6 +45,7 @@
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw
 from tensor2tensor.rl.restarter import Restarter
+from tensor2tensor.rl.rl_utils import absolute_hinge_difference
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -301,9 +302,7 @@ def append_debug_frame_batch(sim_obs, real_obs, sim_cum_rews,
           local_nps.append(np.asarray(img))
         local_nps.append(np.zeros_like(local_nps[0]))
         headers.append(np.concatenate(local_nps, axis=1))
-      errs = np.maximum(
-          np.abs(sim_obs.astype(np.int) - real_obs, dtype=np.int) - 10, 0
-      ).astype(np.uint8)
+      errs = absolute_hinge_difference(sim_obs, real_obs)
       headers = np.stack(headers)
       debug_frame_batches.append(  # pylint: disable=cell-var-from-loop
           np.concatenate([headers,

From a782d04e5d4fcdedd5444b9d2404a744f8717a49 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Thu, 10 Jan 2019 14:39:36 -0800
Subject: [PATCH 1480/2720] internal merge of PR #1359

PiperOrigin-RevId: 228779774
---
 tensor2tensor/rl/player.py       | 46 +++++++++++++++-----------------
 tensor2tensor/rl/player_utils.py |  9 ++++---
 tensor2tensor/rl/rl_utils.py     |  3 ++-
 3 files changed, 29 insertions(+), 29 deletions(-)

diff --git a/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 77e6e67d1..b4da7fdd3 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -103,6 +103,7 @@
                      "If infer game name from data_dir filenames or from "
                      "hparams.")
 
+
 class PlayerEnv(gym.Env):
   """Base (abstract) environment for interactive human play with gym.utils.play.
 
@@ -205,10 +206,7 @@ def _player_toggle_wait_action(self):
     return self._last_step_tuples
 
   def step(self, action):
-    """Pass action to underlying environment(s) or perform special action.
-
-    For returned tuple explanation see _player_step_tuple() in subclasses.
-    """
+    """Pass action to underlying environment(s) or perform special action."""
     # Special codes
     if action in self._player_actions():
       envs_step_tuples = self._player_actions()[action]()
@@ -272,15 +270,16 @@ def _update_statistics(self, envs_step_tuples):
     raise NotImplementedError
 
   def _player_return_done_action(self):
-    """
+    """Function.
 
     Returns:
        envs_step_tuples: such that `player_step_tuple(envs_step_tuples)`
-        will return done."""
+        will return done.
+    """
     raise NotImplementedError
 
   def _player_step_tuple(self, envs_step_tuples):
-    """Infer return tuple for step() given underlying environment(s) tuple(s)"""
+    """Infer return tuple for step() given underlying environment tuple(s)."""
     raise NotImplementedError
 
 
@@ -308,7 +307,7 @@ class SimAndRealEnvPlayer(PlayerEnv):
   RESTART_SIMULATED_ENV_ACTION = 110
 
   def __init__(self, real_env, sim_env):
-    """
+    """Init.
 
     Args:
       real_env: real environment such as `FlatBatchEnv<T2TGymEnv>`.
@@ -348,6 +347,9 @@ def get_keys_to_action(self):
   def _player_step_tuple(self, envs_step_tuples):
     """Construct observation, return usual step tuple.
 
+    Args:
+      envs_step_tuples: tuples.
+
     Returns:
       Step tuple: ob, reward, done, info
         ob: concatenated images [simulated observation, real observation,
@@ -356,8 +358,8 @@ def _player_step_tuple(self, envs_step_tuples):
         done: True iff. envs_step_tuples['real_env'][2] is True
         info: real environment info
     """
-    ob_real, reward_real, _, _ = envs_step_tuples['real_env']
-    ob_sim, reward_sim, _, _ = envs_step_tuples['sim_env']
+    ob_real, reward_real, _, _ = envs_step_tuples["real_env"]
+    ob_sim, reward_sim, _, _ = envs_step_tuples["sim_env"]
     ob_err = absolute_hinge_difference(ob_sim, ob_real)
 
     ob_real_aug = self._augment_observation(ob_real, reward_real,
@@ -369,7 +371,7 @@ def _player_step_tuple(self, envs_step_tuples):
         self.cumulative_sim_reward - self.cumulative_real_reward
     )
     ob = np.concatenate([ob_sim_aug, ob_real_aug, ob_err_aug], axis=1)
-    _, reward, done, info = envs_step_tuples['real_env']
+    _, reward, done, info = envs_step_tuples["real_env"]
     return ob, reward, done, info
 
   def reset(self):
@@ -379,7 +381,7 @@ def reset(self):
     # Initialize simulated environment with frames from real one.
     self.sim_env.add_to_initial_stack(ob_real)
     for _ in range(3):
-      ob_real, _, _, _ = self.real_env.step(self.name_to_action_num['NOOP'])
+      ob_real, _, _, _ = self.real_env.step(self.name_to_action_num["NOOP"])
       self.sim_env.add_to_initial_stack(ob_real)
     # TODO(konradczechowski): remove when not longer needed.
     # for i in range(12):
@@ -402,10 +404,7 @@ def set_zero_cumulative_rewards(self):
     self.cumulative_sim_reward = 0
 
   def _step_envs(self, action):
-    """Perform `step(action)` on both environments.
-
-    Update initial_frame_stack for simulated environment.
-    """
+    """Perform step(action) on environments and update initial_frame_stack."""
     self._frame_counter += 1
     real_env_step_tuple = self.real_env.step(action)
     sim_env_step_tuple = self.sim_env.step(action)
@@ -413,8 +412,8 @@ def _step_envs(self, action):
     return self._pack_step_tuples(real_env_step_tuple, sim_env_step_tuple)
 
   def _update_statistics(self, envs_step_tuples):
-    self.cumulative_real_reward += envs_step_tuples['real_env'][1]
-    self.cumulative_sim_reward += envs_step_tuples['sim_env'][1]
+    self.cumulative_real_reward += envs_step_tuples["real_env"][1]
+    self.cumulative_sim_reward += envs_step_tuples["sim_env"][1]
 
   def _player_return_done_action(self):
     ob = np.zeros(self.real_env.observation_space.shape, dtype=np.uint8)
@@ -427,11 +426,10 @@ def player_restart_simulated_env_action(self):
     # TODO(konradczechowski): remove when this will be not needed.
     # new_ob, _, _, _ = self.sim_env.step(2)
     # print("\n\n\n\ndiff {}\n\n\n\n".format((ob - new_ob).sum()))
-    # ##########
-    assert np.all(self._last_step_tuples['real_env'][0] == ob)
+    assert np.all(self._last_step_tuples["real_env"][0] == ob)
     self.set_zero_cumulative_rewards()
     return self._pack_step_tuples(
-        self._last_step_tuples['real_env'], (ob, 0, False, {}))
+        self._last_step_tuples["real_env"], (ob, 0, False, {}))
 
 
 class SingleEnvPlayer(PlayerEnv):
@@ -460,7 +458,7 @@ def __init__(self, env):
 
   def _player_step_tuple(self, envs_step_tuples):
     """Augment observation, return usual step tuple."""
-    ob, reward, done, info = envs_step_tuples['env']
+    ob, reward, done, info = envs_step_tuples["env"]
     ob = self._augment_observation(ob, reward, self.cumulative_reward)
     return ob, reward, done, info
 
@@ -479,7 +477,7 @@ def _step_envs(self, action):
     return self._pack_step_tuples(self.env.step(action))
 
   def _update_statistics(self, envs_step_tuples):
-    _, reward, _, _ = envs_step_tuples['env']
+    _, reward, _, _ = envs_step_tuples["env"]
     self.cumulative_reward += reward
 
   def _player_return_done_action(self):
@@ -501,7 +499,7 @@ def main(_):
       data=FLAGS.episodes_data_dir)
   if FLAGS.game_from_filenames:
     hparams.set_hparam(
-        'game', player_utils.infer_game_name_from_filenames(directories['data'])
+        "game", player_utils.infer_game_name_from_filenames(directories["data"])
     )
   epoch = FLAGS.epoch if FLAGS.epoch == "last" else int(FLAGS.epoch)
 
diff --git a/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index 000c18316..376ab1f10 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -59,7 +59,7 @@ class SimulatedGymEnv(gym.Env):
 
   def __init__(self, real_env, world_model_dir, hparams, random_starts,
                setable_initial_frames=False):
-    """
+    """Init.
 
     Args:
        real_env: gym environment.
@@ -111,9 +111,9 @@ def add_to_initial_stack(self, frame):
     """Adds new frame to (initial) frame stack, removes last one."""
     if not self._setable_initial_frames:
       raise ValueError(
-          "This instace does not allow to manually set initial frame stack.")
-    assert frame.shape == self._initial_frames.shape[2:], \
-        '{}, {}'.format(frame.shape, self._initial_frames.shape[:1])
+          "This instance does not allow to manually set initial frame stack.")
+    assert_msg = "{}, {}".format(frame.shape, self._initial_frames.shape[:1])
+    assert frame.shape == self._initial_frames.shape[2:], assert_msg
     initial_frames = np.roll(self._initial_frames, shift=-1, axis=1)
     initial_frames[0, -1, ...] = frame
     self._initial_frames = initial_frames
@@ -156,6 +156,7 @@ def setup_and_load_epoch(hparams, data_dir, which_epoch_data=None):
 
 
 def infer_game_name_from_filenames(data_dir, snake_case=True):
+  """Infer name from filenames."""
   names = os.listdir(data_dir)
   game_names = [re.findall(pattern=r"^Gym(.*)NoFrameskip", string=name)
                 for name in names]
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 56214e73e..43a67e111 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -202,7 +202,8 @@ def absolute_hinge_difference(arr1, arr2, min_diff=10, dtype=np.uint8):
   """Point-wise, hinge loss-like, difference between arrays.
 
   Args:
-    arr1, arr2: integer arrays to compare.
+    arr1: integer array to compare.
+    arr2: integer array to compare.
     min_diff: minimal difference taken into consideration.
     dtype: dtype of returned array.
 

From 57a97206f0ddcd1149cebc5b30083e0e3fcfe487 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 11 Jan 2019 02:19:49 +0100
Subject: [PATCH 1481/2720] Reduce usage of tf.contrib.layers (#1350)

* Replace tf.contrib.layers initializers with tf.initializers

* Remove unused imports of tf.contrib.layers

* tf.contrib.layers.conv2d -> tf.layers.conv2d

* tf.contrib.layers.flatten -> tf.layers.flatten

* tf.contrib.layers.fully_connected -> tf.layers.dense

* Fix cyclegan_upsample unit test

See https://github.com/tensorflow/tensor2tensor/pull/1350#discussion_r246194594
---
 tensor2tensor/layers/common_layers.py         | 12 ++---
 tensor2tensor/layers/common_layers_test.py    |  4 +-
 tensor2tensor/models/research/rl.py           | 50 ++++++++-----------
 .../models/research/transformer_vae.py        |  2 +-
 tensor2tensor/models/revnet.py                |  4 +-
 tensor2tensor/models/video/base.py            |  4 --
 .../models/video/basic_deterministic.py       |  4 --
 tensor2tensor/models/video/basic_recurrent.py |  6 ---
 tensor2tensor/models/video/emily.py           |  2 +-
 tensor2tensor/models/video/sv2p.py            |  2 +-
 tensor2tensor/models/video/svg_lp.py          |  4 +-
 tensor2tensor/utils/optimize.py               |  2 +-
 12 files changed, 38 insertions(+), 58 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index d1930e5d0..ed026adb6 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3577,21 +3577,21 @@ def cyclegan_upsample(net, num_outputs, stride, method="conv2d_transpose"):
       net = tf.image.resize_nearest_neighbor(
           net, [stride[0] * height, stride[1] * width])
       net = tf.pad(net, spatial_pad_1, "REFLECT")
-      net = tf.contrib.layers.conv2d(
-          net, num_outputs, kernel_size=[3, 3], padding="valid")
+      net = tf.layers.conv2d(
+          net, num_outputs, (3, 3), activation=tf.nn.relu)
     elif method == "bilinear_upsample_conv":
       net = tf.image.resize_bilinear(net,
                                      [stride[0] * height, stride[1] * width])
       net = tf.pad(net, spatial_pad_1, "REFLECT")
-      net = tf.contrib.layers.conv2d(
-          net, num_outputs, kernel_size=[3, 3], padding="valid")
+      net = tf.layers.conv2d(
+          net, num_outputs, (3, 3), activation=tf.nn.relu)
     elif method == "conv2d_transpose":
       # This corrects 1 pixel offset for images with even width and height.
       # conv2d is left aligned and conv2d_transpose is right aligned for even
       # sized images (while doing "SAME" padding).
       # Note: This doesn"t reflect actual model in paper.
-      net = tf.contrib.layers.conv2d_transpose(
-          net, num_outputs, kernel_size=[3, 3], stride=stride, padding="valid")
+      net = tf.layers.conv2d_transpose(
+          net, num_outputs, (3, 3), strides=stride, activation=tf.nn.relu)
       net = net[:, 1:, 1:, :]
     else:
       raise ValueError("Unknown method: [%s]" % method)
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index bc46f0409..d90b09394 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -685,8 +685,8 @@ def testCycleGANUpsampleConv2dTranspose(self):
     num_channels = 3
     output_filters = 10
     stride = [2, 3]  # we want height to be x2 and width to be x3
-    random_input = np.random.rand(batch, height, width, num_channels).astype(
-        np.float32)
+    random_input = tf.convert_to_tensor(
+        np.random.rand(batch, height, width, num_channels), dtype=tf.float32)
 
     # conv2d_transpose is a little tricky.
     # height_new = (height_old - 1) * stride + kernel - 2*padding - correction
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index ff2b0d72d..cc5761c2f 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -426,8 +426,8 @@ def feed_forward_gaussian_fun(action_space, config, observations):
   if not isinstance(action_space, gym.spaces.box.Box):
     raise ValueError("Expecting continuous action space.")
 
-  mean_weights_initializer = tf.contrib.layers.variance_scaling_initializer(
-      factor=config.init_mean_factor)
+  mean_weights_initializer = tf.initializers.variance_scaling(
+      scale=config.init_mean_factor)
   logstd_initializer = tf.random_normal_initializer(config.init_logstd, 1e-10)
 
   flat_observations = tf.reshape(observations, [
@@ -438,10 +438,10 @@ def feed_forward_gaussian_fun(action_space, config, observations):
     with tf.variable_scope("policy"):
       x = flat_observations
       for size in config.policy_layers:
-        x = tf.contrib.layers.fully_connected(x, size, tf.nn.relu)
-      mean = tf.contrib.layers.fully_connected(
-          x, action_space.shape[0], tf.tanh,
-          weights_initializer=mean_weights_initializer)
+        x = tf.layers.dense(x, size, activation=tf.nn.relu)
+      mean = tf.layers.dense(
+          x, action_space.shape[0], activation=tf.tanh,
+          kernel_initializer=mean_weights_initializer)
       logstd = tf.get_variable(
           "logstd", mean.shape[2:], tf.float32, logstd_initializer)
       logstd = tf.tile(
@@ -450,8 +450,8 @@ def feed_forward_gaussian_fun(action_space, config, observations):
     with tf.variable_scope("value"):
       x = flat_observations
       for size in config.value_layers:
-        x = tf.contrib.layers.fully_connected(x, size, tf.nn.relu)
-      value = tf.contrib.layers.fully_connected(x, 1, None)[..., 0]
+        x = tf.layers.dense(x, size, activation=tf.nn.relu)
+      value = tf.layers.dense(x, 1)[..., 0]
   mean = tf.check_numerics(mean, "mean")
   logstd = tf.check_numerics(logstd, "logstd")
   value = tf.check_numerics(value, "value")
@@ -480,16 +480,14 @@ def body(self, features):
     with tf.variable_scope("policy"):
       x = flat_observations
       for size in self.hparams.policy_layers:
-        x = tf.contrib.layers.fully_connected(x, size, tf.nn.relu)
-      logits = tf.contrib.layers.fully_connected(
-          x, self.hparams.problem.num_actions, activation_fn=None
-      )
+        x = tf.layers.dense(x, size, activation=tf.nn.relu)
+      logits = tf.layers.dense(x, self.hparams.problem.num_actions)
       logits = tf.expand_dims(logits, axis=1)
     with tf.variable_scope("value"):
       x = flat_observations
       for size in self.hparams.value_layers:
-        x = tf.contrib.layers.fully_connected(x, size, tf.nn.relu)
-      value = tf.contrib.layers.fully_connected(x, 1, None)
+        x = tf.layers.dense(x, size, activation=tf.nn.relu)
+      value = tf.layers.dense(x, 1)
     logits = clip_logits(logits, self.hparams)
     return {"target_policy": logits, "target_value": value}
 
@@ -506,14 +504,14 @@ def body(self, features):
     dropout = getattr(self.hparams, "dropout_ppo", 0.0)
     with tf.variable_scope("feed_forward_cnn_small"):
       x = tf.cast(x, tf.float32) / 255.0
-      x = tf.contrib.layers.conv2d(x, 32, [5, 5], [2, 2],
-                                   activation_fn=tf.nn.relu, padding="SAME")
-      x = tf.contrib.layers.conv2d(x, 32, [5, 5], [2, 2],
-                                   activation_fn=tf.nn.relu, padding="SAME")
+      x = tf.layers.conv2d(x, 32, (5, 5), strides=(2, 2),
+                           activation=tf.nn.relu, padding="same")
+      x = tf.layers.conv2d(x, 32, (5, 5), strides=(2, 2),
+                           activation=tf.nn.relu, padding="same")
 
       flat_x = tf.layers.flatten(x)
       flat_x = tf.layers.dropout(flat_x, rate=dropout)
-      x = tf.contrib.layers.fully_connected(flat_x, 128, tf.nn.relu)
+      x = tf.layers.dense(flat_x, 128, activation=tf.nn.relu)
 
       logits = tf.layers.dense(
           x, self.hparams.problem.num_actions, name="dense2"
@@ -521,8 +519,7 @@ def body(self, features):
       logits = clip_logits(logits, self.hparams)
       logits = tf.expand_dims(logits, axis=1)
 
-      value = tf.contrib.layers.fully_connected(
-          x, 1, activation_fn=None)
+      value = tf.layers.dense(x, 1)
     return {"target_policy": logits, "target_value": value}
 
 
@@ -575,15 +572,12 @@ def body(self, features):
     with tf.variable_scope("dense_bitwise"):
       flat_x = discretization.int_to_bit_embed(flat_x, 8, 32)
 
-      x = tf.contrib.layers.fully_connected(flat_x, 256, tf.nn.relu)
-      x = tf.contrib.layers.fully_connected(flat_x, 128, tf.nn.relu)
+      x = tf.layers.dense(flat_x, 256, activation=tf.nn.relu)
+      x = tf.layers.dense(flat_x, 128, activation=tf.nn.relu)
 
-      logits = tf.contrib.layers.fully_connected(
-          x, self.hparams.problem.num_actions, activation_fn=None
-      )
+      logits = tf.layers.dense(x, self.hparams.problem.num_actions)
 
-      value = tf.contrib.layers.fully_connected(
-          x, 1, activation_fn=None)[..., 0]
+      value = tf.layers.dense(x, 1)[..., 0]
 
     return {"target_policy": logits, "target_value": value}
 
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 88f0eb757..c630abc6c 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -587,7 +587,7 @@ def __init__(self, *args, **kwargs):
                 self._hparams.num_residuals, self._hparams.num_blocks,
                 self._hparams.hidden_size, block_dim
             ],
-            initializer=tf.contrib.layers.xavier_initializer(),
+            initializer=tf.initializers.glorot_uniform(),
             trainable=self._hparams.trainable_projections)
 
         self._hparams.bottleneck = functools.partial(
diff --git a/tensor2tensor/models/revnet.py b/tensor2tensor/models/revnet.py
index 036c85840..ebde0dbe8 100644
--- a/tensor2tensor/models/revnet.py
+++ b/tensor2tensor/models/revnet.py
@@ -49,8 +49,8 @@ def wrapped_partial(fn, *args, **kwargs):
   return wrapped
 
 
-conv_initializer = tf.contrib.layers.variance_scaling_initializer(
-    factor=2.0, mode='FAN_OUT')
+conv_initializer = tf.initializers.variance_scaling(
+    scale=2.0, mode='fan_out')
 
 CONFIG = {'2d': {'conv': wrapped_partial(
     tf.layers.conv2d, kernel_initializer=conv_initializer),
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 24905fe0e..2326772c6 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -32,10 +32,6 @@
 import tensorflow as tf
 
 
-tfl = tf.layers
-tfcl = tf.contrib.layers
-
-
 def flat_lists(list_of_lists):
   return [x for l in list_of_lists for x in l]
 
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 54fdb1641..d065042b5 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -29,10 +29,6 @@
 import tensorflow as tf
 
 
-tfl = tf.layers
-tfcl = tf.contrib.layers
-
-
 @registry.register_model
 class NextFrameBasicDeterministic(base.NextFrameBase):
   """Basic next-frame model, may take actions and predict rewards too."""
diff --git a/tensor2tensor/models/video/basic_recurrent.py b/tensor2tensor/models/video/basic_recurrent.py
index c121712eb..42cae07de 100644
--- a/tensor2tensor/models/video/basic_recurrent.py
+++ b/tensor2tensor/models/video/basic_recurrent.py
@@ -23,12 +23,6 @@
 from tensor2tensor.models.video import basic_stochastic
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
-
-
-tfl = tf.layers
-tfcl = tf.contrib.layers
-
 
 @registry.register_model
 class NextFrameBasicRecurrent(
diff --git a/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index 877f3cdb8..221e41eeb 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -276,7 +276,7 @@ def construct_model(self, images, actions, rewards):
     for i, image in enumerate(images):
       with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
         enc, skips = self.encoder(image, rnn_size, has_batchnorm=has_batchnorm)
-        enc = tfcl.flatten(enc)
+        enc = tfl.flatten(enc)
         enc_images.append(enc)
         enc_skips.append(skips)
 
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 1ed42bb46..3c6c75af6 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -314,7 +314,7 @@ def construct_predictive_tower(
 
       if self.hparams.model_options == "CDNA":
         # cdna_input = tf.reshape(hidden5, [int(batch_size), -1])
-        cdna_input = tfcl.flatten(hidden5)
+        cdna_input = tfl.flatten(hidden5)
         transformed += common_video.cdna_transformation(
             input_image, cdna_input, num_masks, int(color_channels),
             self.hparams.dna_kernel_size, self.hparams.relu_shift)
diff --git a/tensor2tensor/models/video/svg_lp.py b/tensor2tensor/models/video/svg_lp.py
index 52e9ec2df..7242ade84 100644
--- a/tensor2tensor/models/video/svg_lp.py
+++ b/tensor2tensor/models/video/svg_lp.py
@@ -180,7 +180,7 @@ def construct_model(self, images, actions, rewards):
     for i, image in enumerate(images):
       with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
         enc, skips = self.encoder(image, g_dim, has_batchnorm=has_batchnorm)
-        enc = tfcl.flatten(enc)
+        enc = tfl.flatten(enc)
         enc_images.append(enc)
         enc_skips.append(skips)
 
@@ -199,7 +199,7 @@ def construct_model(self, images, actions, rewards):
           h_current = enc_images[i-1]
         else:
           h_current, _ = self.encoder(gen_images[-1], g_dim)
-          h_current = tfcl.flatten(h_current)
+          h_current = tfl.flatten(h_current)
 
         # target encoding
         h_target = enc_images[i]
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index ab06a6b2f..dd42168e5 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -304,6 +304,6 @@ def get_variable_initializer(hparams):
     return tf.variance_scaling_initializer(
         hparams.initializer_gain, mode="fan_avg", distribution="uniform")
   elif hparams.initializer == "xavier":
-    return tf.contrib.layers.xavier_initializer()
+    return tf.initializers.glorot_uniform()
   else:
     raise ValueError("Unrecognized initializer: %s" % hparams.initializer)

From 31701a41d752264a3e84aa25c6f7ffaa9cc902ef Mon Sep 17 00:00:00 2001
From: Michael Mezher <mezhermikey@yahoo.com>
Date: Thu, 10 Jan 2019 20:30:02 -0500
Subject: [PATCH 1482/2720] Fixed tensor2tensor language modeling decode
 (#1282)

* Changed reuse val from true to tf.AUTO_REUSE in top to allow for proper weight initialization

* Fixed decoding from file for language modeling problems (when has_input=False)

* nan

* Fixed language modeling decoding from file (allowing variable batch size without padding and EOS insertion)

* Delete misc.xml

* Delete modules.xml

* Delete tensor2tensor.iml

* Delete vcs.xml

* Delete workspace.xml
---
 tensor2tensor/utils/decoding.py | 86 ++++++++++++++++++++++++++++++---
 1 file changed, 79 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 36eadb3ed..139e404e9 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -407,13 +407,18 @@ def input_fn(params):
       return dataset
   else:
     def input_fn():
-      input_gen = _decode_batch_input_fn(
-          num_decode_batches, sorted_inputs,
-          inputs_vocab, decode_hp.batch_size,
-          decode_hp.max_input_size, task_id=decode_hp.multiproblem_task_id)
-      gen_fn = make_input_fn_from_generator(input_gen)
-      example = gen_fn()
-      return _decode_input_tensor_to_features_dict(example, hparams)
+      if has_input:
+        input_gen = _decode_batch_input_fn(
+            num_decode_batches, sorted_inputs,
+            inputs_vocab, decode_hp.batch_size,
+            decode_hp.max_input_size, task_id=decode_hp.multiproblem_task_id)
+      else:
+        input_gen = _decode_batch_input_fn_no_padding(sorted_inputs=sorted_inputs,max_batch_size=decode_hp.batch_size,
+                                                      vocabulary=inputs_vocab,max_input_size=decode_hp.max_input_size,
+                                                      decode_hp=decode_hp)
+    gen_fn = make_input_fn_from_generator(input_gen)
+    example = gen_fn()
+    return _decode_input_tensor_to_features_dict(example, hparams)
   decodes = []
   result_iter = estimator.predict(input_fn, checkpoint_path=checkpoint_path)
 
@@ -643,6 +648,73 @@ def _decode_batch_input_fn(num_decode_batches, sorted_inputs, vocabulary,
         "inputs": np.array(final_batch_inputs).astype(np.int32),
     }
 
+def _decode_batch_input_fn_no_padding(sorted_inputs, max_batch_size, vocabulary, max_input_size, decode_hp):
+    """Generator to produce batches of same length inputs (batch size will be variable)."""
+
+    # First reverse all the input sentences so that if you're going to get OOMs,
+    # you'll see it in the first batch
+    sorted_inputs.reverse()
+
+    #Get variable batch sizes
+    last_batch_length=None
+    batch_lengths, batch_indicies = [],[]
+    for batch_index,elm in enumerate(sorted_inputs):
+        #Exclude whitespace and empty strings from batch length.
+        this_batch_length=len(elm.split(' '))
+        if max_input_size>0:
+            if this_batch_length>max_input_size:
+                this_batch_length=max_input_size
+        if this_batch_length!=last_batch_length:
+            batch_lengths.append(this_batch_length)
+            batch_indicies.append(batch_index)
+            last_batch_length = this_batch_length
+    batch_indicies.append(len(sorted_inputs))
+
+    #Ensure no batches exceed the maximum batch_size
+    batch_sizes = np.diff(batch_indicies)
+    final_batch_sizes = []
+    final_batch_lengths = []
+    for ii,bs in enumerate(batch_sizes):
+      if bs<max_batch_size:
+          final_batch_sizes.append(bs)
+          final_batch_lengths.append(batch_lengths[ii])
+      else:
+          full_batches = bs//max_batch_size
+          partial_batch= bs%max_batch_size
+          for _ in range(full_batches):
+              final_batch_sizes.append(max_batch_size)
+              final_batch_lengths.append(batch_lengths[ii])
+          if partial_batch>0:
+              final_batch_sizes.append(partial_batch)
+              final_batch_lengths.append(batch_lengths[ii])
+
+    #Continue with now variable batch sizes, no need for padding.
+    last_index=0
+    for b,batch_size in enumerate(final_batch_sizes):
+        tf.logging.info("Decoding batch %d" % b)
+        # Batch length should be the same for the entire batch -- Add one additional term for <EOS> token insertion (opt)
+        batch_length = min(max_input_size,final_batch_lengths[b]) + 1
+        batch_inputs = []
+        for inputs in sorted_inputs[last_index:last_index+batch_size]:
+          input_ids = vocabulary.encode(inputs)
+          if max_input_size>0:
+              #For language modeling problems, more recent inputs are often more important.
+              input_ids = input_ids[-max_input_size:]
+          #Padding and <EOS> removed -- for language modeling problems.
+          batch_inputs.append(input_ids)
+        last_index+=batch_size
+
+        final_batch_inputs = []
+        #Ensure consistent batch size
+        for in_ids in batch_inputs:
+          assert len(in_ids) == batch_length
+          x=in_ids
+          final_batch_inputs.append(x)
+
+        yield {
+            "inputs": np.array(final_batch_inputs).astype(np.int32),
+        }
+
 
 def _interactive_input_fn(hparams, decode_hp):
   """Generator that reads from the terminal and yields "interactive inputs".

From 17181cdeefbe561349837846dd42387aa1acddfb Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Thu, 10 Jan 2019 17:20:44 -0800
Subject: [PATCH 1483/2720] internal merge of PR #1350

PiperOrigin-RevId: 228806666
---
 tensor2tensor/utils/decoding.py | 86 +++------------------------------
 1 file changed, 7 insertions(+), 79 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 139e404e9..36eadb3ed 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -407,18 +407,13 @@ def input_fn(params):
       return dataset
   else:
     def input_fn():
-      if has_input:
-        input_gen = _decode_batch_input_fn(
-            num_decode_batches, sorted_inputs,
-            inputs_vocab, decode_hp.batch_size,
-            decode_hp.max_input_size, task_id=decode_hp.multiproblem_task_id)
-      else:
-        input_gen = _decode_batch_input_fn_no_padding(sorted_inputs=sorted_inputs,max_batch_size=decode_hp.batch_size,
-                                                      vocabulary=inputs_vocab,max_input_size=decode_hp.max_input_size,
-                                                      decode_hp=decode_hp)
-    gen_fn = make_input_fn_from_generator(input_gen)
-    example = gen_fn()
-    return _decode_input_tensor_to_features_dict(example, hparams)
+      input_gen = _decode_batch_input_fn(
+          num_decode_batches, sorted_inputs,
+          inputs_vocab, decode_hp.batch_size,
+          decode_hp.max_input_size, task_id=decode_hp.multiproblem_task_id)
+      gen_fn = make_input_fn_from_generator(input_gen)
+      example = gen_fn()
+      return _decode_input_tensor_to_features_dict(example, hparams)
   decodes = []
   result_iter = estimator.predict(input_fn, checkpoint_path=checkpoint_path)
 
@@ -648,73 +643,6 @@ def _decode_batch_input_fn(num_decode_batches, sorted_inputs, vocabulary,
         "inputs": np.array(final_batch_inputs).astype(np.int32),
     }
 
-def _decode_batch_input_fn_no_padding(sorted_inputs, max_batch_size, vocabulary, max_input_size, decode_hp):
-    """Generator to produce batches of same length inputs (batch size will be variable)."""
-
-    # First reverse all the input sentences so that if you're going to get OOMs,
-    # you'll see it in the first batch
-    sorted_inputs.reverse()
-
-    #Get variable batch sizes
-    last_batch_length=None
-    batch_lengths, batch_indicies = [],[]
-    for batch_index,elm in enumerate(sorted_inputs):
-        #Exclude whitespace and empty strings from batch length.
-        this_batch_length=len(elm.split(' '))
-        if max_input_size>0:
-            if this_batch_length>max_input_size:
-                this_batch_length=max_input_size
-        if this_batch_length!=last_batch_length:
-            batch_lengths.append(this_batch_length)
-            batch_indicies.append(batch_index)
-            last_batch_length = this_batch_length
-    batch_indicies.append(len(sorted_inputs))
-
-    #Ensure no batches exceed the maximum batch_size
-    batch_sizes = np.diff(batch_indicies)
-    final_batch_sizes = []
-    final_batch_lengths = []
-    for ii,bs in enumerate(batch_sizes):
-      if bs<max_batch_size:
-          final_batch_sizes.append(bs)
-          final_batch_lengths.append(batch_lengths[ii])
-      else:
-          full_batches = bs//max_batch_size
-          partial_batch= bs%max_batch_size
-          for _ in range(full_batches):
-              final_batch_sizes.append(max_batch_size)
-              final_batch_lengths.append(batch_lengths[ii])
-          if partial_batch>0:
-              final_batch_sizes.append(partial_batch)
-              final_batch_lengths.append(batch_lengths[ii])
-
-    #Continue with now variable batch sizes, no need for padding.
-    last_index=0
-    for b,batch_size in enumerate(final_batch_sizes):
-        tf.logging.info("Decoding batch %d" % b)
-        # Batch length should be the same for the entire batch -- Add one additional term for <EOS> token insertion (opt)
-        batch_length = min(max_input_size,final_batch_lengths[b]) + 1
-        batch_inputs = []
-        for inputs in sorted_inputs[last_index:last_index+batch_size]:
-          input_ids = vocabulary.encode(inputs)
-          if max_input_size>0:
-              #For language modeling problems, more recent inputs are often more important.
-              input_ids = input_ids[-max_input_size:]
-          #Padding and <EOS> removed -- for language modeling problems.
-          batch_inputs.append(input_ids)
-        last_index+=batch_size
-
-        final_batch_inputs = []
-        #Ensure consistent batch size
-        for in_ids in batch_inputs:
-          assert len(in_ids) == batch_length
-          x=in_ids
-          final_batch_inputs.append(x)
-
-        yield {
-            "inputs": np.array(final_batch_inputs).astype(np.int32),
-        }
-
 
 def _interactive_input_fn(hparams, decode_hp):
   """Generator that reads from the terminal and yields "interactive inputs".

From 2794fa7e82cdd1df99409e976ec7dd78c68cae4d Mon Sep 17 00:00:00 2001
From: Michael Mezher <mezhermikey@yahoo.com>
Date: Thu, 10 Jan 2019 17:41:09 -0800
Subject: [PATCH 1484/2720] internal merge of PR #1282

PiperOrigin-RevId: 228809213
---
 tensor2tensor/utils/decoding.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 36eadb3ed..23d22f890 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -410,7 +410,8 @@ def input_fn():
       input_gen = _decode_batch_input_fn(
           num_decode_batches, sorted_inputs,
           inputs_vocab, decode_hp.batch_size,
-          decode_hp.max_input_size, task_id=decode_hp.multiproblem_task_id)
+          decode_hp.max_input_size,
+          task_id=decode_hp.multiproblem_task_id, has_input=has_input)
       gen_fn = make_input_fn_from_generator(input_gen)
       example = gen_fn()
       return _decode_input_tensor_to_features_dict(example, hparams)
@@ -616,7 +617,8 @@ def input_fn():
 
 
 def _decode_batch_input_fn(num_decode_batches, sorted_inputs, vocabulary,
-                           batch_size, max_input_size, task_id=-1):
+                           batch_size, max_input_size,
+                           task_id=-1, has_input=True):
   """Generator to produce batches of inputs."""
   tf.logging.info(" batch %d" % num_decode_batches)
   for b in range(num_decode_batches):
@@ -628,8 +630,9 @@ def _decode_batch_input_fn(num_decode_batches, sorted_inputs, vocabulary,
       if max_input_size > 0:
         # Subtract 1 for the EOS_ID.
         input_ids = input_ids[:max_input_size - 1]
-      final_id = text_encoder.EOS_ID if task_id < 0 else task_id
-      input_ids.append(final_id)
+      if has_input or task_id > -1:  # Do not append EOS for pure LM tasks.
+        final_id = text_encoder.EOS_ID if task_id < 0 else task_id
+        input_ids.append(final_id)
       batch_inputs.append(input_ids)
       if len(input_ids) > batch_length:
         batch_length = len(input_ids)

From bf5743c17de4215fdaee94350d1e694a3995abb9 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 11 Jan 2019 17:09:37 +0100
Subject: [PATCH 1485/2720] Add labels to ml-engine job (#1361)

Cloud ML Engine can use [labels](
https://cloud.google.com/ml-engine/docs/tensorflow/resource-labels) to organize and filter jobs. This PR adds labels for `model`, `problem` and `hparams_set`.
---
 tensor2tensor/utils/cloud_mlengine.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 4e8b2cfc0..9de890f8f 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -202,8 +202,15 @@ def configure_job():
     )
 
   timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-  job_name = "%s_%s_t2t_%s" % (FLAGS.model, FLAGS.problem, timestamp)
-  job_spec = {"jobId": job_name, "trainingInput": training_input}
+  job_spec = {
+    "jobId": "%s_%s_t2t_%s" % (FLAGS.model, FLAGS.problem, timestamp),
+    "labels": {
+      "model": FLAGS.model,
+      "problem": FLAGS.problem,
+      "hparams": FLAGS.hparams_set
+    },
+    "trainingInput": training_input,
+  }
   return job_spec
 
 

From d8735bbc3ae4eba00d7b7148f7d9b4cc4dd9d818 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 11 Jan 2019 08:08:58 -0800
Subject: [PATCH 1486/2720] Try py_func anyway (it may be suboptimal in new TF,
 but if it works for now let's do it and we can change it everywhere later).

PiperOrigin-RevId: 228883864
---
 tensor2tensor/utils/cloud_mlengine.py | 11 ++---------
 tensor2tensor/utils/sari_hook.py      |  2 +-
 2 files changed, 3 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 9de890f8f..4e8b2cfc0 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -202,15 +202,8 @@ def configure_job():
     )
 
   timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-  job_spec = {
-    "jobId": "%s_%s_t2t_%s" % (FLAGS.model, FLAGS.problem, timestamp),
-    "labels": {
-      "model": FLAGS.model,
-      "problem": FLAGS.problem,
-      "hparams": FLAGS.hparams_set
-    },
-    "trainingInput": training_input,
-  }
+  job_name = "%s_%s_t2t_%s" % (FLAGS.model, FLAGS.problem, timestamp)
+  job_spec = {"jobId": job_name, "trainingInput": training_input}
   return job_spec
 
 
diff --git a/tensor2tensor/utils/sari_hook.py b/tensor2tensor/utils/sari_hook.py
index 5e0df0056..d433d5e41 100644
--- a/tensor2tensor/utils/sari_hook.py
+++ b/tensor2tensor/utils/sari_hook.py
@@ -214,7 +214,7 @@ def get_sari_numpy(source_ids, prediction_ids, target_ids):
     return (np.asarray(sari_scores), np.asarray(keep_scores),
             np.asarray(add_scores), np.asarray(deletion_scores))
 
-  sari, keep, add, deletion = tf.py_function(
+  sari, keep, add, deletion = tf.py_func(
       get_sari_numpy,
       [source_ids, prediction_ids, target_ids],
       [tf.float64, tf.float64, tf.float64, tf.float64])

From bd82a924c981f3fba7e6f589c1e8cba3eba69bc7 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 11 Jan 2019 08:09:57 -0800
Subject: [PATCH 1487/2720] internal merge of PR #1361

PiperOrigin-RevId: 228883970
---
 tensor2tensor/utils/cloud_mlengine.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 4e8b2cfc0..4c9c7a701 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -202,8 +202,15 @@ def configure_job():
     )
 
   timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-  job_name = "%s_%s_t2t_%s" % (FLAGS.model, FLAGS.problem, timestamp)
-  job_spec = {"jobId": job_name, "trainingInput": training_input}
+  job_spec = {
+      "jobId": "%s_%s_t2t_%s" % (FLAGS.model, FLAGS.problem, timestamp),
+      "labels": {
+          "model": FLAGS.model,
+          "problem": FLAGS.problem,
+          "hparams": FLAGS.hparams_set
+      },
+      "trainingInput": training_input,
+  }
   return job_spec
 
 

From 7d6a169775e92210c1a7076a0cacd2566cf99930 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Fri, 11 Jan 2019 18:33:37 +0100
Subject: [PATCH 1488/2720] Fix observation reading in SimulatedBatchGymEnv.
 (#1363)

* Fix observation reading in SimulatedBatchGymEnv.

* Remove commments with debug code.
---
 tensor2tensor/rl/envs/simulated_batch_gym_env.py | 3 ++-
 tensor2tensor/rl/player.py                       | 7 -------
 2 files changed, 2 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index 7bc4c276a..9ff2e5ba0 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -53,7 +53,8 @@ def __init__(self, *args, **kwargs):
 
       self._actions_t = tf.placeholder(shape=(self.batch_size,), dtype=tf.int32)
       self._rewards_t, self._dones_t = self._batch_env.simulate(self._actions_t)
-      self._obs_t = self._batch_env.observ
+      with tf.control_dependencies([self._rewards_t]):
+        self._obs_t = self._batch_env.observ
       self._reset_op = self._batch_env.reset(
           tf.range(self.batch_size, dtype=tf.int32)
       )
diff --git a/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index b4da7fdd3..a13351729 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -383,10 +383,6 @@ def reset(self):
     for _ in range(3):
       ob_real, _, _, _ = self.real_env.step(self.name_to_action_num["NOOP"])
       self.sim_env.add_to_initial_stack(ob_real)
-    # TODO(konradczechowski): remove when not longer needed.
-    # for i in range(12):
-    #   ob_real, _, _, _ = self.real_env.step(np.random.choice([2,3, 4, 5]))
-    #   self.sim_env.add_to_initial_stack(ob_real)
     ob_sim = self.sim_env.reset()
     assert np.all(ob_real == ob_sim)
     self._last_step_tuples = self._pack_step_tuples((ob_real, 0, False, {}),
@@ -423,9 +419,6 @@ def _player_return_done_action(self):
   def player_restart_simulated_env_action(self):
     self._frame_counter = 0
     ob = self.sim_env.reset()
-    # TODO(konradczechowski): remove when this will be not needed.
-    # new_ob, _, _, _ = self.sim_env.step(2)
-    # print("\n\n\n\ndiff {}\n\n\n\n".format((ob - new_ob).sum()))
     assert np.all(self._last_step_tuples["real_env"][0] == ob)
     self.set_zero_cumulative_rewards()
     return self._pack_step_tuples(

From 27224cd117fbe4ca45207caaa8251e8fee5591b3 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 11 Jan 2019 10:02:20 -0800
Subject: [PATCH 1489/2720] tf.io.gfile.isdir -> tf.gfile.IsDirectory

PiperOrigin-RevId: 228898979
---
 tensor2tensor/rl/envs/simulated_batch_env.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index c9a7dc61a..f84687890 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -161,7 +161,7 @@ def initialize(self, sess):
     model_loader = tf.train.Saver(
         var_list=tf.global_variables(scope="next_frame*")  # pylint:disable=unexpected-keyword-arg
     )
-    if tf.io.gfile.isdir(self._model_dir):
+    if tf.gfile.IsDirectory(self._model_dir):
       trainer_lib.restore_checkpoint(
           self._model_dir, saver=model_loader, sess=sess, must_restore=True
       )

From 9266e67e9fad84c1f063f823b2cde27f86fab82c Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Fri, 11 Jan 2019 19:13:52 +0100
Subject: [PATCH 1490/2720] Player, correct mapping keyboard keys to actions.
 (#1364)

---
 tensor2tensor/rl/player.py   | 51 ++++++++++++++++++------------------
 tensor2tensor/rl/rl_utils.py | 22 +++++++++++++---
 2 files changed, 43 insertions(+), 30 deletions(-)

diff --git a/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index a13351729..7ab904d46 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -53,17 +53,15 @@
 from __future__ import print_function
 
 import gym
-from gym.envs.atari.atari_env import ACTION_MEANING
 from gym.utils import play
 import numpy as np
-import six
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
 from tensor2tensor.rl import player_utils
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw
 from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
-from tensor2tensor.rl.rl_utils import absolute_hinge_difference
+from tensor2tensor.rl.rl_utils import absolute_hinge_difference, full_game_name
 # Import flags from t2t_trainer and trainer_model_based
 import tensor2tensor.rl.trainer_model_based_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
@@ -137,24 +135,23 @@ class PlayerEnv(gym.Env):
 
   HEADER_HEIGHT = 27
 
-  def __init__(self):
+  def __init__(self, action_meanings):
+    """
+
+    Args:
+      action_meanings: list of strings indicating action names. Can be obtain by
+        >>> env = gym.make("PongNoFrameskip-v4")  # insert your game name
+        >>> env.unwrapped.get_action_meanings()
+        See gym AtariEnv get_action_meanings() for more details.
+    """
+    self.action_meanings = action_meanings
     self._wait = True
     # If action_space will be needed, one could use e.g. gym.spaces.Dict.
     self.action_space = None
     self._last_step_tuples = None
-
-  def _init_action_mappings(self, env):
-    # Atari dependant. In case of problems with keyboard key interpretation
-    # switch to _action_set instead of range(env.action_space.n) (similarly to
-    # how gym AtariEnv does). _action_set can probably be obtain from full
-    # game name.
-    self.action_meaning = {i: ACTION_MEANING[i]
-                           for i in range(env.action_space.n)}
-    self.name_to_action_num = {v: k for k, v in
-                               six.iteritems(self.action_meaning)}
-
-  def _get_action_meanings(self):
-    return [self.action_meaning[i] for i in range(len(self.action_meaning))]
+    self.action_meanings = action_meanings
+    self.name_to_action_num = {name: num for num, name in
+                               enumerate(self.action_meanings)}
 
   def get_keys_to_action(self):
     """Get mapping from keyboard keys to actions.
@@ -178,7 +175,7 @@ def get_keys_to_action(self):
 
     keys_to_action = {}
 
-    for action_id, action_meaning in enumerate(self._get_action_meanings()):
+    for action_id, action_meaning in enumerate(self.action_meanings):
       keys = []
       for keyword, key in keyword_to_key.items():
         if keyword in action_meaning:
@@ -255,6 +252,7 @@ def _augment_observation(self, ob, reward, cumulative_reward):
       pixel_fill = (0, 255, 0)
     else:
       pixel_fill = (255, 0, 0)
+      pixel_fill = (255, 0, 0)
     header[0, :, :] = pixel_fill
     return np.concatenate([header, ob], axis=0)
 
@@ -306,7 +304,7 @@ class SimAndRealEnvPlayer(PlayerEnv):
 
   RESTART_SIMULATED_ENV_ACTION = 110
 
-  def __init__(self, real_env, sim_env):
+  def __init__(self, real_env, sim_env, action_meanings):
     """Init.
 
     Args:
@@ -315,7 +313,7 @@ def __init__(self, real_env, sim_env):
         `SimulatedGymEnv` must allow to update initial frames for next reset
         with `add_to_initial_stack` method.
     """
-    super(SimAndRealEnvPlayer, self).__init__()
+    super(SimAndRealEnvPlayer, self).__init__(action_meanings)
     assert real_env.observation_space.shape == sim_env.observation_space.shape
     self.real_env = real_env
     self.sim_env = sim_env
@@ -329,7 +327,6 @@ def __init__(self, real_env, sim_env):
     self.observation_space = gym.spaces.Box(low=orig.low.min(),
                                             high=orig.high.max(),
                                             shape=shape, dtype=orig.dtype)
-    self._init_action_mappings(sim_env)
 
   def _player_actions(self):
     actions = super(SimAndRealEnvPlayer, self)._player_actions()
@@ -438,8 +435,8 @@ class SingleEnvPlayer(PlayerEnv):
       Plural form used for consistency with `PlayerEnv`.
   """
 
-  def __init__(self, env):
-    super(SingleEnvPlayer, self).__init__()
+  def __init__(self, env, action_meanings):
+    super(SingleEnvPlayer, self).__init__(action_meanings)
     self.env = env
     # Set observation space
     orig = self.env.observation_space
@@ -447,7 +444,6 @@ def __init__(self, env):
     self.observation_space = gym.spaces.Box(low=orig.low.min(),
                                             high=orig.high.max(),
                                             shape=shape, dtype=orig.dtype)
-    self._init_action_mappings(env)
 
   def _player_step_tuple(self, envs_step_tuples):
     """Augment observation, return usual step tuple."""
@@ -494,6 +490,8 @@ def main(_):
     hparams.set_hparam(
         "game", player_utils.infer_game_name_from_filenames(directories["data"])
     )
+  action_meanings = gym.make(full_game_name(hparams.game)).\
+      unwrapped.get_action_meanings()
   epoch = FLAGS.epoch if FLAGS.epoch == "last" else int(FLAGS.epoch)
 
   def make_real_env():
@@ -514,18 +512,19 @@ def make_simulated_env(setable_initial_frames, which_epoch_data):
     sim_env = make_simulated_env(
         which_epoch_data=None, setable_initial_frames=True)
     real_env = make_real_env()
-    env = SimAndRealEnvPlayer(real_env, sim_env)
+    env = SimAndRealEnvPlayer(real_env, sim_env, action_meanings)
   else:
     if FLAGS.simulated_env:
       env = make_simulated_env(  # pylint: disable=redefined-variable-type
           which_epoch_data=epoch, setable_initial_frames=False)
     else:
       env = make_real_env()
-    env = SingleEnvPlayer(env)  # pylint: disable=redefined-variable-type
+    env = SingleEnvPlayer(env, action_meanings)  # pylint: disable=redefined-variable-type
 
   env = player_utils.wrap_with_monitor(env, FLAGS.video_dir)
 
   if FLAGS.dry_run:
+    env.unwrapped.get_keys_to_action()
     for _ in range(5):
       env.reset()
       for i in range(50):
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 43a67e111..638946d6a 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -126,12 +126,26 @@ def summarize_metrics(eval_metrics_writer, metrics, epoch):
 }
 
 
+ATARI_GAME_MODE = "NoFrameskip-v4"
+
+
+def full_game_name(short_name):
+  """CamelCase game name with mode suffix.
+
+  Args:
+    short_name: snake_case name without mode e.g "crazy_climber"
+
+  Returns:
+    full game name e.g. "CrazyClimberNoFrameskip-v4"
+  """
+  camel_game_name = misc_utils.snakecase_to_camelcase(short_name)
+  full_name = camel_game_name + ATARI_GAME_MODE
+  return full_name
+
+
 def setup_env(hparams, batch_size, max_num_noops, rl_env_max_episode_steps=-1):
   """Setup."""
-  game_mode = "NoFrameskip-v4"
-  camel_game_name = misc_utils.snakecase_to_camelcase(hparams.game)
-  camel_game_name += game_mode
-  env_name = camel_game_name
+  env_name = full_game_name(hparams.game)
 
   env = T2TGymEnv(base_env_name=env_name,
                   batch_size=batch_size,

From 29e6950e51023fbcd3af39a3b6b478753625a9c9 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Fri, 11 Jan 2019 10:18:30 -0800
Subject: [PATCH 1491/2720] internal merge of PR #1364

PiperOrigin-RevId: 228901855
---
 tensor2tensor/rl/player.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 7ab904d46..4bc6bd2a0 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -61,7 +61,8 @@
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw
 from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
-from tensor2tensor.rl.rl_utils import absolute_hinge_difference, full_game_name
+from tensor2tensor.rl.rl_utils import absolute_hinge_difference
+from tensor2tensor.rl.rl_utils import full_game_name
 # Import flags from t2t_trainer and trainer_model_based
 import tensor2tensor.rl.trainer_model_based_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
@@ -136,7 +137,7 @@ class PlayerEnv(gym.Env):
   HEADER_HEIGHT = 27
 
   def __init__(self, action_meanings):
-    """
+    """Constructor for PlayerEnv.
 
     Args:
       action_meanings: list of strings indicating action names. Can be obtain by
@@ -252,7 +253,6 @@ def _augment_observation(self, ob, reward, cumulative_reward):
       pixel_fill = (0, 255, 0)
     else:
       pixel_fill = (255, 0, 0)
-      pixel_fill = (255, 0, 0)
     header[0, :, :] = pixel_fill
     return np.concatenate([header, ob], axis=0)
 
@@ -312,6 +312,10 @@ def __init__(self, real_env, sim_env, action_meanings):
       sim_env: simulation of `real_env` to be compared with. E.g.
         `SimulatedGymEnv` must allow to update initial frames for next reset
         with `add_to_initial_stack` method.
+      action_meanings: list of strings indicating action names. Can be obtain by
+        >>> env = gym.make("PongNoFrameskip-v4")  # insert your game name
+        >>> env.unwrapped.get_action_meanings()
+        See gym AtariEnv get_action_meanings() for more details.
     """
     super(SimAndRealEnvPlayer, self).__init__(action_meanings)
     assert real_env.observation_space.shape == sim_env.observation_space.shape

From 77f4437e2c0e635ae2ee74aad3f26fcf7f3a80c2 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 11 Jan 2019 10:57:57 -0800
Subject: [PATCH 1492/2720] Increase tolerances in a gradient check in tests.
 This may not be the correct thing to do, but 2e-6 and 1e-6 seem close enough.

PiperOrigin-RevId: 228908963
---
 tensor2tensor/layers/common_layers_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index d90b09394..60d5961ad 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -631,8 +631,8 @@ def testConvHiddenReluMemoryEfficient(self):
             dx, df1, df2, dnorm_scale, dnorm_bias,
             dx_f, df1_f, df2_f, dnorm_scale_f, dnorm_bias_f])
     self.assertAllClose(y, y_forget)
-    self.assertAllClose(df2, df2_f)
-    self.assertAllClose(df1, df1_f)
+    self.assertAllClose(df2, df2_f, rtol=2e-6, atol=2e-6)
+    self.assertAllClose(df1, df1_f, rtol=2e-6, atol=2e-6)
     self.assertAllClose(dnorm_scale, dnorm_scale_f)
     self.assertAllClose(dnorm_bias, dnorm_bias_f)
     self.assertAllClose(dx, dx_f)

From 7836aa77c3faf54444994f5834084a4b3508a25d Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 11 Jan 2019 11:07:06 -0800
Subject: [PATCH 1493/2720] Check if ffmpeg is installed and guard against it
 in RL env.

PiperOrigin-RevId: 228910916
---
 tensor2tensor/layers/common_video.py         | 16 +++++++++++++---
 tensor2tensor/rl/envs/simulated_batch_env.py |  5 +++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index ef4568b6e..70a5931e3 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -357,8 +357,8 @@ def _encode_gif(images, fps):
   """Encodes numpy images into gif string.
 
   Args:
-    images: A 5-D `uint8` `np.array` (or a list of 4-D images) of shape
-      `[batch_size, time, height, width, channels]` where `channels` is 1 or 3.
+    images: A 4-D `uint8` `np.array` (or a list of 3-D images) of shape
+      `[time, height, width, channels]` where `channels` is 1 or 3.
     fps: frames per second of the animation
 
   Returns:
@@ -372,6 +372,16 @@ def _encode_gif(images, fps):
   return writer.finish()
 
 
+def ffmpeg_works():
+  """Tries to encode images with ffmpeg to check if it works."""
+  images = np.zeros((2, 32, 32, 3), dtype=np.uint8)
+  try:
+    _encode_gif(images, 2)
+    return True
+  except (IOError, OSError):
+    return False
+
+
 def py_gif_summary(tag, images, max_outputs, fps, return_summary_value=False):
   """Outputs a `Summary` protocol buffer with gif animations.
 
@@ -697,7 +707,7 @@ def __init__(self, fps, output_path=None, file_format="gif"):
   def __init_ffmpeg(self, image_shape):
     """Initializes ffmpeg to write frames."""
     import itertools  # pylint: disable=g-import-not-at-top
-    from subprocess import Popen, PIPE  # pylint: disable=g-import-not-at-top,g-multiple-import
+    from subprocess import Popen, PIPE  # pylint: disable=g-import-not-at-top,g-multiple-import,g-importing-member
     ffmpeg = "ffmpeg"
     height, width, channels = image_shape
     self.cmd = [
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index f84687890..874fdccd2 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -115,6 +115,7 @@ def __init__(
     """Batch of environments inside the TensorFlow graph."""
     super(SimulatedBatchEnv, self).__init__(observation_space, action_space)
 
+    self._ffmpeg_works = common_video.ffmpeg_works()
     self.batch_size = batch_size
     self._min_reward = reward_range[0]
     self._num_frames = frame_stack_size
@@ -267,6 +268,8 @@ def history_observations(self):
     return self.history_buffer.get_all_elements()
 
   def _video_dump_frame(self, obs, rews):
+    if not self._ffmpeg_works:
+      return
     if self._video_writer is None:
       self._video_counter += 1
       self._video_writer = common_video.WholeVideoWriter(
@@ -280,6 +283,8 @@ def _video_dump_frame(self, obs, rews):
     self._video_writer.write(np.concatenate([np.asarray(img), obs[0]], axis=0))
 
   def _video_dump_frames(self, obs):
+    if not self._ffmpeg_works:
+      return
     zeros = np.zeros(obs.shape[0])
     for i in range(obs.shape[1]):
       self._video_dump_frame(obs[:, i, :], zeros)

From 152beb0f4248de8d2cda733e0645ac6d22f74b42 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 11 Jan 2019 20:45:39 +0100
Subject: [PATCH 1494/2720] Implement PlannerAgent (#1365)

* Extract a function for running rollouts to rl_utils

* Extract a base class for batch wrappers

* Factorize make_simulated_env_fn_from_hparams by a function returning
just kwargs

* Implement PlannerAgent
---
 tensor2tensor/models/research/rl.py           |  43 +++---
 .../rl/envs/simulated_batch_gym_env.py        |   8 +-
 tensor2tensor/rl/evaluator.py                 |  97 +++++++-----
 tensor2tensor/rl/evaluator_test.py            |   7 +-
 tensor2tensor/rl/rl_utils.py                  | 140 +++++++++++++++---
 tensor2tensor/rl/trainer_model_based.py       |  12 +-
 6 files changed, 224 insertions(+), 83 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index cc5761c2f..ebbcbf506 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -208,25 +208,34 @@ def env_fn(in_graph):
   return env_fn
 
 
-def make_simulated_env_fn_from_hparams(
-    real_env, hparams, batch_size, initial_frame_chooser, model_dir,
-    sim_video_dir=None):
-  """Creates a simulated env_fn."""
-  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
+# TODO(koz4k): Move this and the one below to rl_utils.
+def make_simulated_env_kwargs(real_env, hparams, **extra_kwargs):
+  """Extracts simulated env kwargs from real_env and loop hparams."""
+  objs_and_attrs = [
+      (real_env, [
+          "reward_range", "observation_space", "action_space", "frame_height",
+          "frame_width"
+      ]),
+      (hparams, ["frame_stack_size", "intrinsic_reward_scale"])
+  ]
+  kwargs = {
+      attr: getattr(obj, attr)
+      for (obj, attrs) in objs_and_attrs for attr in attrs
+  }
+  kwargs["model_name"] = hparams.generative_model
+  kwargs["model_hparams"] = trainer_lib.create_hparams(
+      hparams.generative_model_params
+  )
   if hparams.wm_policy_param_sharing:
-    model_hparams.optimizer_zero_grads = True
+    kwargs["model_hparams"].optimizer_zero_grads = True
+  kwargs.update(extra_kwargs)
+  return kwargs
+
+
+def make_simulated_env_fn_from_hparams(real_env, hparams, **extra_kwargs):
+  """Creates a simulated env_fn."""
   return make_simulated_env_fn(
-      reward_range=real_env.reward_range,
-      observation_space=real_env.observation_space,
-      action_space=real_env.action_space,
-      frame_stack_size=hparams.frame_stack_size,
-      frame_height=real_env.frame_height, frame_width=real_env.frame_width,
-      initial_frame_chooser=initial_frame_chooser, batch_size=batch_size,
-      model_name=hparams.generative_model,
-      model_hparams=trainer_lib.create_hparams(hparams.generative_model_params),
-      model_dir=model_dir,
-      intrinsic_reward_scale=hparams.intrinsic_reward_scale,
-      sim_video_dir=sim_video_dir,
+      **make_simulated_env_kwargs(real_env, hparams, **extra_kwargs)
   )
 
 
diff --git a/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index 9ff2e5ba0..def04baaa 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -22,6 +22,7 @@
 from gym import Env
 from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
 
+import numpy as np
 import tensorflow as tf
 
 
@@ -55,6 +56,7 @@ def __init__(self, *args, **kwargs):
       self._rewards_t, self._dones_t = self._batch_env.simulate(self._actions_t)
       with tf.control_dependencies([self._rewards_t]):
         self._obs_t = self._batch_env.observ
+      self._indices_t = tf.placeholder(shape=(self.batch_size,), dtype=tf.int32)
       self._reset_op = self._batch_env.reset(
           tf.range(self.batch_size, dtype=tf.int32)
       )
@@ -79,9 +81,9 @@ def render(self, mode="human"):
     raise NotImplementedError()
 
   def reset(self, indices=None):
-    if indices:
-      raise NotImplementedError()
-    obs = self._sess.run(self._reset_op)
+    if indices is None:
+      indices = np.array(range(self.batch_size))
+    obs = self._sess.run(self._reset_op, feed_dict={self._indices_t: indices})
     # TODO(pmilos): remove if possible
     # obs[:, 0, 0, 0] = 0
     # obs[:, 0, 0, 1] = 255
diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index fbbc06e38..c453a2d03 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -28,13 +28,12 @@
 from __future__ import division
 from __future__ import print_function
 
-import numpy as np
-
 from tensor2tensor.models.research import rl  # pylint: disable=unused-import
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl import trainer_model_based_params  # pylint: disable=unused-import
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
+from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -44,20 +43,38 @@
 
 
 flags.DEFINE_string("policy_dir", "", "Directory with policy checkpoints.")
+flags.DEFINE_string("model_dir", "", "Directory with model checkpoints.")
 flags.DEFINE_string(
     "eval_metrics_dir", "", "Directory to output the eval metrics at."
 )
 flags.DEFINE_bool("full_eval", True, "Whether to ignore the timestep limit.")
-flags.DEFINE_enum("agent", "policy", ["random", "policy"], "Agent type to use.")
+flags.DEFINE_enum(
+    "agent", "policy", ["random", "policy", "planner"], "Agent type to use."
+)
 flags.DEFINE_bool(
     "eval_with_learner", True,
     "Whether to use the PolicyLearner.evaluate function instead of an "
     "out-of-graph one. Works only with --agent=policy."
 )
+flags.DEFINE_string(
+    "planner_hparams_set", "planner_tiny", "Planner hparam set."
+)
+flags.DEFINE_string("planner_hparams", "", "Planner hparam overrides.")
+
+
+@registry.register_hparams
+def planner_tiny():
+  return tf.contrib.training.HParams(
+      num_rollouts=1,
+      planning_horizon=2,
+      rollout_agent_type="random",
+  )
 
 
 def make_agent(
-    agent_type, env, policy_hparams, policy_dir, sampling_temp
+    agent_type, env, policy_hparams, policy_dir, sampling_temp,
+    sim_env_kwargs=None, frame_stack_size=None, planning_horizon=None,
+    rollout_agent_type=None
 ):
   """Factory function for Agents."""
   return {
@@ -68,45 +85,40 @@ def make_agent(
           env.batch_size, env.observation_space, env.action_space,
           policy_hparams, policy_dir, sampling_temp
       ),
+      "planner": lambda: rl_utils.PlannerAgent(  # pylint: disable=g-long-lambda
+          env.batch_size, make_agent(
+              rollout_agent_type, env, policy_hparams, policy_dir, sampling_temp
+          ), rl_utils.SimulatedBatchGymEnvWithFixedInitialFrames(
+              **sim_env_kwargs
+          ), lambda env: rl_utils.BatchStackWrapper(env, frame_stack_size),
+          planning_horizon
+      ),
   }[agent_type]()
 
 
-def make_eval_fn_with_agent(agent_type):
+def make_eval_fn_with_agent(agent_type, planner_hparams, model_dir):
   """Returns an out-of-graph eval_fn using the Agent API."""
-  def eval_fn(env, hparams, policy_hparams, policy_dir, sampling_temp):
+  def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
     """Eval function."""
     base_env = env
-    env = rl_utils.BatchStackWrapper(env, hparams.frame_stack_size)
+    env = rl_utils.BatchStackWrapper(env, loop_hparams.frame_stack_size)
+    sim_env_kwargs = rl.make_simulated_env_kwargs(
+        base_env, loop_hparams, batch_size=planner_hparams.num_rollouts,
+        model_dir=model_dir
+    )
     agent = make_agent(
-        agent_type, env, policy_hparams, policy_dir, sampling_temp
+        agent_type, env, policy_hparams, policy_dir, sampling_temp,
+        sim_env_kwargs, loop_hparams.frame_stack_size,
+        planner_hparams.planning_horizon, planner_hparams.rollout_agent_type
     )
-    num_dones = 0
-    first_dones = [False] * env.batch_size
-    observations = env.reset()
-    while num_dones < env.batch_size:
-      actions = agent.act(observations)
-      (observations, _, dones) = env.step(actions)
-      observations = list(observations)
-      now_done_indices = []
-      for (i, done) in enumerate(dones):
-        if done and not first_dones[i]:
-          now_done_indices.append(i)
-          first_dones[i] = True
-          num_dones += 1
-      if now_done_indices:
-        # Reset only envs done the first time in this timestep to ensure that
-        # we collect exactly 1 rollout from each env.
-        reset_observations = env.reset(now_done_indices)
-        for (i, observation) in zip(now_done_indices, reset_observations):
-          observations[i] = observation
-      observations = np.array(observations)
+    rl_utils.run_rollouts(env, agent, env.reset())
     assert len(base_env.current_epoch_rollouts()) == env.batch_size
   return eval_fn
 
 
 def evaluate(
-    hparams, policy_dir, eval_metrics_dir, agent_type, eval_with_learner,
-    report_fn=None, report_metric=None
+    loop_hparams, planner_hparams, policy_dir, model_dir, eval_metrics_dir,
+    agent_type, eval_with_learner, report_fn=None, report_metric=None
 ):
   """Evaluate."""
   if eval_with_learner:
@@ -118,16 +130,20 @@ def evaluate(
   eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
   kwargs = {}
   if not eval_with_learner:
-    kwargs["eval_fn"] = make_eval_fn_with_agent(agent_type)
-  eval_metrics = rl_utils.evaluate_all_configs(hparams, policy_dir, **kwargs)
+    kwargs["eval_fn"] = make_eval_fn_with_agent(
+        agent_type, planner_hparams, model_dir
+    )
+  eval_metrics = rl_utils.evaluate_all_configs(
+      loop_hparams, policy_dir, **kwargs
+  )
   rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, 0)
 
   # Report metrics
   if report_fn:
     if report_metric == "mean_reward":
       metric_name = rl_utils.get_metric_name(
-          sampling_temp=hparams.eval_sampling_temps[0],
-          max_num_noops=hparams.eval_max_num_noops,
+          sampling_temp=loop_hparams.eval_sampling_temps[0],
+          max_num_noops=loop_hparams.eval_max_num_noops,
           clipped=False
       )
       report_fn(eval_metrics[metric_name], 0)
@@ -137,12 +153,17 @@ def evaluate(
 
 
 def main(_):
-  hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
+  loop_hparams = trainer_lib.create_hparams(
+      FLAGS.loop_hparams_set, FLAGS.loop_hparams
+  )
   if FLAGS.full_eval:
-    hparams.eval_rl_env_max_episode_steps = -1
+    loop_hparams.eval_rl_env_max_episode_steps = -1
+  planner_hparams = trainer_lib.create_hparams(
+      FLAGS.planner_hparams_set, FLAGS.planner_hparams
+  )
   evaluate(
-      hparams, FLAGS.policy_dir, FLAGS.eval_metrics_dir, FLAGS.agent,
-      FLAGS.eval_with_learner
+      loop_hparams, planner_hparams, FLAGS.policy_dir, FLAGS.model_dir,
+      FLAGS.eval_metrics_dir, FLAGS.agent, FLAGS.eval_with_learner
   )
 
 
diff --git a/tensor2tensor/rl/evaluator_test.py b/tensor2tensor/rl/evaluator_test.py
index 969232faf..62acde315 100644
--- a/tensor2tensor/rl/evaluator_test.py
+++ b/tensor2tensor/rl/evaluator_test.py
@@ -27,11 +27,12 @@
 class EvalTest(tf.test.TestCase):
 
   def test_evaluate_pong_random_agent(self):
-    hparams = registry.hparams("rlmb_tiny")
+    loop_hparams = registry.hparams("rlmb_tiny")
+    planner_hparams = registry.hparams("planner_tiny")
     temp_dir = tf.test.get_temp_dir()
     evaluator.evaluate(
-        hparams, temp_dir, temp_dir, agent_type="random",
-        eval_with_learner=False
+        loop_hparams, planner_hparams, temp_dir, temp_dir, temp_dir,
+        agent_type="random", eval_with_learner=False
     )
 
 
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 638946d6a..008c2d4c7 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -29,6 +29,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl.dopamine_connector import DQNLearner
+from tensor2tensor.rl.envs.simulated_batch_gym_env import SimulatedBatchGymEnv
 from tensor2tensor.rl.ppo_learner import PPOLearner
 from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import trainer_lib
@@ -228,6 +229,45 @@ def absolute_hinge_difference(arr1, arr2, min_diff=10, dtype=np.uint8):
   return np.maximum(diff - min_diff, 0).astype(dtype)
 
 
+def run_rollouts(
+    env, agent, initial_observations, step_limit=None, discount_factor=1.0
+):
+  """Runs a batch of rollouts from given initial observations."""
+  num_dones = 0
+  first_dones = [False] * env.batch_size
+  observations = initial_observations
+  step_index = 0
+  cum_rewards = 0
+
+  def proceed():
+    if step_limit is not None:
+      return step_index < step_limit
+    else:
+      return num_dones < env.batch_size
+
+  while proceed():
+    actions = agent.act(observations)
+    (observations, rewards, dones) = env.step(actions)
+    observations = list(observations)
+    now_done_indices = []
+    for (i, done) in enumerate(dones):
+      if done and not first_dones[i]:
+        now_done_indices.append(i)
+        first_dones[i] = True
+        num_dones += 1
+    if now_done_indices:
+      # Reset only envs done the first time in this timestep to ensure that
+      # we collect exactly 1 rollout from each env.
+      reset_observations = env.reset(now_done_indices)
+      for (i, observation) in zip(now_done_indices, reset_observations):
+        observations[i] = observation
+    observations = np.array(observations)
+    cum_rewards = cum_rewards * discount_factor + rewards
+    step_index += 1
+
+  return (observations, cum_rewards)
+
+
 class BatchAgent(object):
   """Python API for agents.
 
@@ -319,12 +359,78 @@ def estimate_value(self, observations):
     return values
 
 
+class PlannerAgent(BatchAgent):
+  """Agent based on temporal difference planning."""
+
+  def __init__(
+      self, batch_size, rollout_agent, sim_env, wrapper_fn, planning_horizon,
+      discount_factor=1.0
+  ):
+    super(PlannerAgent, self).__init__(
+        batch_size, rollout_agent.observation_space, rollout_agent.action_space
+    )
+    self._rollout_agent = rollout_agent
+    self._sim_env = sim_env
+    self._wrapped_env = wrapper_fn(sim_env)
+    self._discount_factor = discount_factor
+    self._planning_horizon = planning_horizon
+
+  def act(self, observations):
+    def run_batch_from(observation, action):
+      self._sim_env.initial_frames = np.array(
+          [observation] * self._sim_env.batch_size
+      )
+      self._wrapped_env.reset()
+      (initial_observations, initial_rewards, _) = self._wrapped_env.step(
+          np.array([action] * self._wrapped_env.batch_size)
+      )
+      (final_observations, cum_rewards) = run_rollouts(
+          self._wrapped_env, self._rollout_agent, initial_observations,
+          discount_factor=self._discount_factor,
+          step_limit=self._planning_horizon
+      )
+      values = self._rollout_agent.estimate_value(final_observations)
+      total_values = (
+          initial_rewards + self._discount_factor * cum_rewards +
+          self._discount_factor ** (self._planning_horizon + 1) * values
+      )
+      return total_values.mean()
+
+    def choose_best_action(observation):
+      return max(
+          range(self.action_space.n),
+          key=(lambda action: run_batch_from(observation, action))
+      )
+
+    return np.array(list(map(choose_best_action, observations)))
+
+
 # TODO(koz4k): Unify interfaces of batch envs.
-class BatchStackWrapper(object):
+class BatchWrapper(object):
+  """Base class for batch env wrappers."""
+
+  def __init__(self, env):
+    self.env = env
+    self.batch_size = env.batch_size
+    self.observation_space = env.observation_space
+    self.action_space = env.action_space
+    self.reward_range = env.reward_range
+
+  def reset(self, indices=None):
+    return self.env.reset(indices)
+
+  def step(self, actions):
+    return self.env.step(actions)
+
+  def close(self):
+    self.env.close()
+
+
+class BatchStackWrapper(BatchWrapper):
   """Out-of-graph batch stack wrapper."""
 
   def __init__(self, env, stack_size):
-    self.env = env
+    super(BatchStackWrapper, self).__init__(env)
     self.stack_size = stack_size
     inner_space = env.observation_space
     self.observation_space = Box(
@@ -333,21 +439,10 @@ def __init__(self, env, stack_size):
         dtype=inner_space.dtype,
     )
     self._history_buffer = np.zeros(
-        (self.batch_size,) + self.observation_space.shape
+        (self.batch_size,) + self.observation_space.shape,
+        dtype=inner_space.dtype
     )
 
-  @property
-  def batch_size(self):
-    return self.env.batch_size
-
-  @property
-  def action_space(self):
-    return self.env.action_space
-
-  @property
-  def reward_range(self):
-    return self.env.reward_range
-
   def reset(self, indices=None):
     if indices is None:
       indices = range(self.batch_size)
@@ -363,5 +458,16 @@ def step(self, actions):
     self._history_buffer[:, -1, ...] = observations
     return (self._history_buffer, rewards, dones)
 
-  def close(self):
-    self.env.close()
+
+class SimulatedBatchGymEnvWithFixedInitialFrames(BatchWrapper):
+  """Wrapper for SimulatedBatchGymEnv that allows to fix initial frames."""
+
+  def __init__(self, *args, **kwargs):
+    self.initial_frames = None
+    def initial_frame_chooser(batch_size):
+      assert batch_size == self.initial_frames.shape[0]
+      return self.initial_frames
+    env = SimulatedBatchGymEnv(
+        *args, initial_frame_chooser=initial_frame_chooser, **kwargs
+    )
+    super(SimulatedBatchGymEnvWithFixedInitialFrames, self).__init__(env)
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 21188bb9e..dcee57a27 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -150,9 +150,11 @@ def train_agent(real_env, learner, world_model_dir, hparams, epoch):
       hparams.simulation_flip_first_random_for_beginning
   )
   env_fn = make_simulated_env_fn_from_hparams(
-      real_env, hparams, hparams.simulated_batch_size, initial_frame_chooser,
-      world_model_dir, os.path.join(learner.agent_model_dir,
-                                    "sim_videos_{}".format(epoch))
+      real_env, hparams, batch_size=hparams.simulated_batch_size,
+      initial_frame_chooser=initial_frame_chooser, model_dir=world_model_dir,
+      sim_video_dir=os.path.join(
+          learner.agent_model_dir, "sim_videos_{}".format(epoch)
+      )
   )
   base_algo_str = hparams.base_algo
   train_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
@@ -243,8 +245,8 @@ def initial_frame_chooser(batch_size):
     ])
 
   env_fn = make_simulated_env_fn_from_hparams(
-      real_env, hparams, hparams.wm_eval_batch_size, initial_frame_chooser,
-      world_model_dir
+      real_env, hparams, batch_size=hparams.wm_eval_batch_size,
+      initial_frame_chooser=initial_frame_chooser, model_dir=world_model_dir
   )
   sim_env = env_fn(in_graph=False)
   subsequence_length = int(

From e9eb66aebba9e626982ab7686028a7c59ef0e113 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 11 Jan 2019 11:46:00 -0800
Subject: [PATCH 1495/2720] internal merge of PR #1365

PiperOrigin-RevId: 228918597
---
 tensor2tensor/rl/envs/simulated_batch_gym_env.py | 4 +++-
 tensor2tensor/rl/evaluator.py                    | 2 +-
 tensor2tensor/rl/rl_utils.py                     | 1 +
 3 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index def04baaa..f5d0da563 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -20,9 +20,11 @@
 from __future__ import print_function
 
 from gym import Env
-from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
 
 import numpy as np
+
+from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
+
 import tensorflow as tf
 
 
diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index c453a2d03..87a8c221e 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -32,8 +32,8 @@
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl import trainer_model_based_params  # pylint: disable=unused-import
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
-from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import registry
+from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
 
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 008c2d4c7..c89e52bea 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -377,6 +377,7 @@ def __init__(
 
   def act(self, observations):
     def run_batch_from(observation, action):
+      """Run a batch of actions."""
       self._sim_env.initial_frames = np.array(
           [observation] * self._sim_env.batch_size
       )

From 3b344700c9fe0259844541ecc161b159dd24d3aa Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 11 Jan 2019 14:02:18 -0800
Subject: [PATCH 1496/2720] Add hparams for RL.

PiperOrigin-RevId: 228941032
---
 tensor2tensor/models/research/rl.py           | 16 +++++++
 .../video/basic_deterministic_params.py       |  9 ++++
 .../rl/trainer_model_based_params.py          | 46 ++++++++++++++++++-
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index ebbcbf506..30a8d0ff2 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -130,6 +130,22 @@ def ppo_original_params():
   return hparams
 
 
+@registry.register_hparams
+def ppo_original_params_gamma95():
+  """Parameters based on the original PPO paper, changed gamma."""
+  hparams = ppo_original_params()
+  hparams.gae_gamma = 0.95
+  return hparams
+
+
+@registry.register_hparams
+def ppo_original_params_gamma90():
+  """Parameters based on the original PPO paper, changed gamma."""
+  hparams = ppo_original_params()
+  hparams.gae_gamma = 0.90
+  return hparams
+
+
 @registry.register_hparams
 def ppo_original_world_model():
   """Atari parameters with world model as policy."""
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 9dee014c3..547c0fcbb 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -59,6 +59,15 @@ def next_frame_pixel_noise():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_pixel_noise_long():
+  """Long scheduled sampling setting."""
+  hparams = next_frame_pixel_noise()
+  hparams.batch_size = 2
+  hparams.video_num_target_frames = 16
+  return hparams
+
+
 @registry.register_hparams
 def next_frame_sampling():
   """Basic conv model with scheduled sampling."""
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 3a7815731..0dd7f5b43 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -306,8 +306,10 @@ def rlmb_base_stochastic_discrete():
   hparams.grayscale = False
   hparams.generative_model = "next_frame_basic_stochastic_discrete"
   hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
+  # The parameters below are the same as base, but repeated for easier reading.
   hparams.ppo_epoch_length = 50
   hparams.simulated_rollout_length = 50
+  hparams.simulated_batch_size = 16
   return hparams
 
 
@@ -320,6 +322,14 @@ def rlmb_base_stochastic_discrete_param_sharing():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_long():
+  """Long setting with base model."""
+  hparams = rlmb_base()
+  hparams.generative_model_params = "next_frame_pixel_noise_long"
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_long_stochastic_discrete():
   """Long setting with stochastic discrete model."""
@@ -330,7 +340,41 @@ def rlmb_long_stochastic_discrete():
 
 
 @registry.register_hparams
-def rlmb_base_stochastic_recurrent():
+def rlmb_long_stochastic_discrete_100steps():
+  """Long setting with stochastic discrete model, changed ppo steps."""
+  hparams = rlmb_long_stochastic_discrete()
+  hparams.ppo_epoch_length = 100
+  hparams.simulated_rollout_length = 100
+  hparams.simulated_batch_size = 8
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_long_stochastic_discrete_25steps():
+  """Long setting with stochastic discrete model, changed ppo steps."""
+  hparams = rlmb_long_stochastic_discrete()
+  hparams.ppo_epoch_length = 25
+  hparams.simulated_rollout_length = 25
+  hparams.simulated_batch_size = 32
+  return hparams
+
+
+def rlmb_long_stochastic_discrete_gamma95():
+  """Long setting with stochastic discrete model, changed gamma."""
+  hparams = rlmb_long_stochastic_discrete()
+  hparams.base_algo_params = "ppo_original_params_gamma95"
+  return hparams
+
+
+def rlmb_long_stochastic_discrete_gamma90():
+  """Long setting with stochastic discrete model, changed gamma."""
+  hparams = rlmb_long_stochastic_discrete()
+  hparams.base_algo_params = "ppo_original_params_gamma90"
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_recurrent():
   """Base setting with recurrent model."""
   hparams = rlmb_base()
   hparams.generative_model = "next_frame_basic_recurrent"

From 4d7fea08ae74140044213e3e3e22124ccb9ba228 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 11 Jan 2019 14:31:18 -0800
Subject: [PATCH 1497/2720] Set filter_size to 1 when number of time-steps
 equal 1 when using 3-D convolutions. This avoids unnecessary padding across
 time for this special case.

PiperOrigin-RevId: 228946248
---
 tensor2tensor/models/research/glow_ops.py      | 16 +++++++++++++---
 tensor2tensor/models/research/glow_ops_test.py |  8 +++++---
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index ef806f6e2..aa9649fb1 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -421,6 +421,7 @@ def conv(name, x, output_channels, filter_size=None, stride=None,
 
   x_shape = common_layers.shape_list(x)
   is_2d = len(x_shape) == 4
+  num_steps = x_shape[1]
 
   # set filter_size, stride and in_channels
   if is_2d:
@@ -435,7 +436,10 @@ def conv(name, x, output_channels, filter_size=None, stride=None,
     conv_filter = tf.nn.conv2d
   else:
     if filter_size is None:
-      filter_size = [2, 3, 3]
+      if num_steps == 1:
+        filter_size = [1, 3, 3]
+      else:
+        filter_size = [2, 3, 3]
     if stride is None:
       stride = [1, 1, 1]
     if dilations is None:
@@ -489,11 +493,17 @@ def conv_block(name, x, mid_channels, dilations=None, activation="relu",
 
     x_shape = common_layers.shape_list(x)
     is_2d = len(x_shape) == 4
+    num_steps = x_shape[1]
     if is_2d:
       first_filter = [3, 3]
       second_filter = [1, 1]
     else:
-      first_filter = [2, 3, 3]
+      # special case when number of steps equal 1 to avoid
+      # padding.
+      if num_steps == 1:
+        first_filter = [1, 3, 3]
+      else:
+        first_filter = [2, 3, 3]
       second_filter = [1, 1, 1]
 
     # Edge Padding + conv2d + actnorm + relu:
@@ -1025,7 +1035,7 @@ def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
     eps_std: Sample x2 with the provided eps_std.
     cond_latents: optionally condition x2 on cond_latents.
     hparams: next_frame_glow hparams.
-    state: tf.nn.rnn_cell.LSTMStateTuple. Current state of the LSTM over z_2.
+    state: tf.nn.rnn_cell.LSTMStateTuple.. Current state of the LSTM over z_2.
            Used only when hparams.latent_dist_encoder == "conv_lstm"
     condition: bool, Whether or not to condition the distribution on
                cond_latents.
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 5266cddbb..19c205157 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -430,16 +430,18 @@ def test_actnorm_3d(self):
       ("dil_gatu", True, "gatu"), ("no_dil_gatu", False, "gatu"),
       ("dil_relu_drop", True, "relu", 0.1),
       ("dil_gatu_drop", True, "gatu", 0.1),
-      ("dil_gatu_drop_noise", True, "gatu", 0.1, 0.1))
+      ("dil_gatu_drop_noise", True, "gatu", 0.1, 0.1),
+      ("gatu_drop_single_step", False, "gatu", 0.1, 0.1, 1),
+      ("dil_gatu_drop_single_step", True, "gatu", 0.1, 0.1, 1),)
   def test_temporal_latent_to_dist(self, apply_dilation, activation,
-                                   dropout=0.0, noise=0.1):
+                                   dropout=0.0, noise=0.1, num_steps=5):
     with tf.Graph().as_default():
       hparams = self.get_glow_hparams()
       hparams.latent_apply_dilations = apply_dilation
       hparams.latent_activation = activation
       hparams.latent_dropout = dropout
       hparams.latent_noise = noise
-      latent_shape = (16, 5, 32, 32, 48)
+      latent_shape = (16, num_steps, 32, 32, 48)
       latents = tf.random_normal(latent_shape)
       dist = glow_ops.temporal_latent_to_dist(
           "tensor_to_dist", latents, hparams)

From 0b099cfd6f987d17149e05f0dc3408057dc9fa24 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 11 Jan 2019 15:04:18 -0800
Subject: [PATCH 1498/2720] Bump setup.py version to 1.12.0

PiperOrigin-RevId: 228952259
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 1b4f1e995..1e83e6bd5 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.11.0',
+    version='1.12.0',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',

From 1f6f621a33755290d528700d3f93b9e5d88d69a8 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 11 Jan 2019 16:00:08 -0800
Subject: [PATCH 1499/2720] Add forgotten registration lines.

PiperOrigin-RevId: 228961118
---
 tensor2tensor/rl/trainer_model_based_params.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 0dd7f5b43..a7df96c30 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -359,6 +359,7 @@ def rlmb_long_stochastic_discrete_25steps():
   return hparams
 
 
+@registry.register_hparams
 def rlmb_long_stochastic_discrete_gamma95():
   """Long setting with stochastic discrete model, changed gamma."""
   hparams = rlmb_long_stochastic_discrete()
@@ -366,6 +367,7 @@ def rlmb_long_stochastic_discrete_gamma95():
   return hparams
 
 
+@registry.register_hparams
 def rlmb_long_stochastic_discrete_gamma90():
   """Long setting with stochastic discrete model, changed gamma."""
   hparams = rlmb_long_stochastic_discrete()

From 6da8b88fee6756f4478a38a5ef87cbc60e30ca6f Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 11 Jan 2019 22:13:55 -0800
Subject: [PATCH 1500/2720] Allow mutliple Transformer layers within UT
 recurrence.

PiperOrigin-RevId: 228990762
---
 .../models/research/universal_transformer.py  | 18 +++++
 .../research/universal_transformer_test.py    |  2 +
 .../research/universal_transformer_util.py    | 75 +++++++++++++------
 tensor2tensor/models/transformer.py           |  6 +-
 4 files changed, 76 insertions(+), 25 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 71b82627f..59403f73a 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -368,6 +368,8 @@ def update_hparams_for_universal_transformer(hparams):
 
   # Number of vanilla transformer layers used to be mixed with u-transofmer.
   hparams.add_hparam("num_mixedin_layers", 2)
+  # Number of transformer layers within the recurrent block (default is 1).
+  hparams.add_hparam("num_inrecurrence_layers", 1)
 
   # Type of recurrency:
   # basic, highway, skip, dwa, act, rnn, gru, lstm.
@@ -521,6 +523,22 @@ def adaptive_universal_transformer_base_tpu():
   return hparams
 
 
+@registry.register_hparams
+def adaptive_universal_transformer_multilayer_tpu():
+  """Multi-layer config for adaptive Transformer on TPU."""
+  hparams = adaptive_universal_transformer_base_tpu()
+  hparams.num_inrecurrence_layers = 2
+  hparams.mix_with_transformer = "before_ut,after_ut"
+  hparams.num_mixedin_layers = 1
+  hparams.transformer_ffn_type = "sepconv"
+  # TODO(lukaszkaiser): the options below don't work on TPU yet, make them work.
+  # hparams.add_step_timing_signal = True
+  # hparams.add_sru = True
+  # hparams.self_attention_type = "dot_product_relative_v2"
+  # hparams.max_relative_position = 256
+  return hparams
+
+
 @registry.register_hparams
 def adaptive_universal_transformer_small():
   hparams = universal_transformer_small()
diff --git a/tensor2tensor/models/research/universal_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
index 8c322a067..15f40ffbd 100644
--- a/tensor2tensor/models/research/universal_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 """Tests for Transformer."""
 
 from __future__ import absolute_import
@@ -40,6 +41,7 @@ def get_model(self,
     hparams.filter_size = 32
     hparams.num_heads = 1
     hparams.layer_prepostprocess_dropout = 0.0
+    hparams.mix_with_transformer = ""
 
     p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE,
                                                      VOCAB_SIZE,
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 3ed3488a9..7ab87b258 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -212,34 +212,36 @@ def universal_transformer_layer(x,
     ValueError: Unknown recurrence type
   """
 
-  def add_vanilla_transformer_layer(x, num_layers):
+  def add_vanilla_transformer_layer(x, num_layers, name):
     """Passes the input through num_layers of vanilla transformer layers.
 
     Args:
      x: input
      num_layers: number of layers
+     name: string, prefix of layer names
 
     Returns:
        output of vanilla_transformer_layer
     """
-
     if hparams.add_position_timing_signal:
       # In case of add_position_timing_signal=true, we set  hparams.pos=None
       # and add position timing signal at the beginning of each step, so for
       # the vanilla transformer, we need to add timing signal here.
       x = common_attention.add_timing_signal_1d(x)
     for layer in range(num_layers):
-      with tf.variable_scope("layer_%d" % layer):
+      with tf.variable_scope(name + "layer_%d" % layer):
         x = ffn_unit(attention_unit(x))
     return x
 
   with tf.variable_scope("universal_transformer_%s" % hparams.recurrence_type):
 
-    if hparams.mix_with_transformer == "before_ut":
-      x = add_vanilla_transformer_layer(x, hparams.num_mixedin_layers)
+    if "before_ut" in hparams.mix_with_transformer:
+      x = add_vanilla_transformer_layer(x, hparams.num_mixedin_layers,
+                                        "before_ut_")
 
     if hparams.recurrence_type == "act":
-      return universal_transformer_act(x, hparams, ffn_unit, attention_unit)
+      output, extra_output = universal_transformer_act(
+          x, hparams, ffn_unit, attention_unit)
 
     else:  # for all the other recurrency types with fixed number of steps
 
@@ -255,8 +257,9 @@ def add_vanilla_transformer_layer(x, num_layers):
           hparams.get("use_memory_as_final_state", False)):
         output = extra_output
 
-    if hparams.mix_with_transformer == "after_ut":
-      output = add_vanilla_transformer_layer(output, hparams.num_mixedin_layers)
+    if "after_ut" in hparams.mix_with_transformer:
+      output = add_vanilla_transformer_layer(output, hparams.num_mixedin_layers,
+                                             "after_ut_")
 
     return output, extra_output
 
@@ -574,9 +577,11 @@ def universal_transformer_basic(layer_inputs,
   """
   state, inputs, memory = tf.unstack(layer_inputs, num=None, axis=0,
                                      name="unstack")
-  state = step_preprocess(state, step, hparams)
+  new_state = step_preprocess(state, step, hparams)
 
-  new_state = ffn_unit(attention_unit(state))
+  for i in range(hparams.num_inrecurrence_layers):
+    with tf.variable_scope("rec_layer_%d" % i):
+      new_state = ffn_unit(attention_unit(new_state))
 
   return new_state, inputs, memory
 
@@ -616,10 +621,13 @@ def universal_transformer_highway(layer_inputs,
   """
 
   state, inputs, memory = layer_inputs
-  state = step_preprocess(state, step, hparams)
+  new_state = step_preprocess(state, step, hparams)
+
+  for i in range(hparams.num_inrecurrence_layers):
+    with tf.variable_scope("rec_layer_%d" % i):
+      new_state = ffn_unit(attention_unit(new_state))
 
-  transformed_state = ffn_unit(attention_unit(state))
-  state.get_shape().assert_is_compatible_with(state.get_shape())
+  transformed_state = new_state
 
   gate_inputs = []
   if "s" in hparams.gates_inputs:
@@ -705,9 +713,13 @@ def universal_transformer_skip(layer_inputs,
   """
 
   state, inputs, memory = layer_inputs
-  state = step_preprocess(state, step, hparams)
+  new_state = step_preprocess(state, step, hparams)
+
+  for i in range(hparams.num_inrecurrence_layers):
+    with tf.variable_scope("rec_layer_%d" % i):
+      new_state = ffn_unit(attention_unit(new_state))
 
-  transformed_state = ffn_unit(attention_unit(state))
+  transformed_state = new_state
 
   inputs.get_shape().assert_is_compatible_with(state.get_shape())
 
@@ -804,10 +816,11 @@ def universal_transformer_depthwise_attention(layer_inputs,
   state_to_be_transformed = tf.reduce_sum(
       (states_so_far * states_so_far_weights), axis=0)
 
-  state_to_be_transformed = step_preprocess(state_to_be_transformed, step,
-                                            hparams)
+  new_state = step_preprocess(state_to_be_transformed, step, hparams)
 
-  new_state = ffn_unit(attention_unit(state_to_be_transformed))
+  for i in range(hparams.num_inrecurrence_layers):
+    with tf.variable_scope("rec_layer_%d" % i):
+      new_state = ffn_unit(attention_unit(new_state))
 
   # add the new state to the memory
   memory = fill_memory_slot(memory, new_state, step + 1)
@@ -1170,7 +1183,10 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
                                     -1)
 
     # apply transformation on the state
-    transformed_state = ffn_unit(attention_unit(state))
+    transformed_state = state
+    for i in range(hparams.num_inrecurrence_layers):
+      with tf.variable_scope("rec_layer_%d" % i):
+        transformed_state = ffn_unit(attention_unit(transformed_state))
 
     # update running part in the weighted state and keep the rest
     new_state = ((transformed_state * update_weights) +
@@ -1323,7 +1339,12 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
                                     -1)
 
     # apply transformation on the state
-    transformed_state = ffn_unit(attention_unit(state))
+    new_state = state
+    for i in range(hparams.num_inrecurrence_layers):
+      with tf.variable_scope("rec_layer_%d" % i):
+        new_state = ffn_unit(attention_unit(new_state))
+
+    transformed_state = new_state
 
     # Add in the weighted state
     accumulated_state = (transformed_state * update_weights) + accumulated_state
@@ -1463,7 +1484,12 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
         tf.expand_dims(p * still_running + new_halted * remainders, -1), -1)
 
     # apply transformation on the state
-    transformed_state = ffn_unit(attention_unit(state))
+    new_state = state
+    for i in range(hparams.num_inrecurrence_layers):
+      with tf.variable_scope("rec_layer_%d" % i):
+        new_state = ffn_unit(attention_unit(new_state))
+
+    transformed_state = new_state
 
     # Add in the weighted state
     new_state = ((transformed_state * update_weights) +
@@ -1608,7 +1634,12 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
                                     -1)
 
     # apply transformation on the state
-    transformed_state = ffn_unit(attention_unit(state))
+    new_state = state
+    for i in range(hparams.num_inrecurrence_layers):
+      with tf.variable_scope("rec_layer_%d" % i):
+        new_state = ffn_unit(attention_unit(new_state))
+
+    transformed_state = new_state
 
     # update running part in the weighted state and keep the rest
     new_state = ((transformed_state * update_weights) +
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index c7ae16d97..027b0d3a8 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1720,11 +1720,11 @@ def transformer_tall_pretrain_lm_tpu_adafactor_large():
   hparams = transformer_tall_pretrain_lm_tpu_adafactor()
   hparams.hidden_size = 1024
   hparams.num_heads = 16
-  hparams.filter_size = 32768
+  hparams.filter_size = 32768  # max fitting in 16G memory is 49152, batch 2
   hparams.batch_size = 4
   hparams.multiproblem_mixing_schedule = "constant"
-  # Task order: LM, en-de/fr/ro/de-en/fr-en/ro-en/cnndm/mnli/squad.
-  hparams.multiproblem_per_task_threshold = "16,4,8,1,4,8,1,2,1,2"
+  # Task order: lm/en-de/en-fr/en-ro/de-en/fr-en/ro-en/cnndm/mnli/squad.
+  hparams.multiproblem_per_task_threshold = "32,8,16,1,8,16,1,2,2,1"
   return hparams
 
 

From 878418865bdc63bc0adc8cb7a1b4242f3becffc5 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 12 Jan 2019 14:40:59 -0800
Subject: [PATCH 1501/2720] Favour an early return in get_hparams(). This way
 we don't do needless extra work when we are returning anyways.

PiperOrigin-RevId: 229035480
---
 tensor2tensor/data_generators/problem.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index d510ce8a8..9b9d00659 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -485,11 +485,12 @@ def get_feature_encoders(self, data_dir=None):
 
   def get_hparams(self, model_hparams=None):
     """Returns problem_hparams."""
-    if model_hparams is None:
-      model_hparams = default_model_hparams()
     if self._hparams is not None:
       return self._hparams
 
+    if model_hparams is None:
+      model_hparams = default_model_hparams()
+
     if self._encoders is None:
       data_dir = (model_hparams and hasattr(model_hparams, "data_dir") and
                   model_hparams.data_dir) or None

From 0dc0a4fd8f1924135ae64123aa35a916ea6903b7 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 12 Jan 2019 14:52:12 -0800
Subject: [PATCH 1502/2720] Rename hparams to model_hparams in
 _create_modalities, since that's what they are.

PiperOrigin-RevId: 229035878
---
 tensor2tensor/data_generators/problem.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 9b9d00659..09c0c2bab 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -972,30 +972,30 @@ def _reverse_problem_hparams(p_hparams):
   p.was_reversed = True
 
 
-def _create_modalities(problem_hparams, hparams):
+def _create_modalities(problem_hparams, model_hparams):
   """Creates modalities and overrides any according to model hparams.
 
   Args:
     problem_hparams: tf.contrib.training.HParams for the Problem. It must have
       modality which is a dict of strings to Modality classes.
-    hparams: tf.contrib.training.HParams for the model. It may have
+    model_hparams: tf.contrib.training.HParams for the model. It may have
       input_modalities and target_modality, which will override
       problem_hparams' modality input and target keys.
 
   Returns:
     None
   """
-  modality_overrides = getattr(hparams, "modality", {})
+  modality_overrides = getattr(model_hparams, "modality", {})
   modality = {}
   for feature_name, modality_cls in six.iteritems(problem_hparams.modality):
     vocab_size = problem_hparams.vocab_size[feature_name]
     # If needed for using a pre-trained model's vocabulary where extra indices
     # were allocated for adding new tasks with unique task ids.
-    if (hasattr(hparams, "multiproblem_vocab_size") and
-        hparams.multiproblem_vocab_size > 0):
-      vocab_size = hparams.multiproblem_vocab_size
+    if (hasattr(model_hparams, "multiproblem_vocab_size") and
+        model_hparams.multiproblem_vocab_size > 0):
+      vocab_size = model_hparams.multiproblem_vocab_size
     modality_cls = modality_overrides.get(feature_name, modality_cls)
-    modality[feature_name] = modality_cls(hparams, vocab_size)
+    modality[feature_name] = modality_cls(model_hparams, vocab_size)
   problem_hparams.modality = modality
 
 

From 54060fff6255904762b6f6a048a541d66f91a109 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 12 Jan 2019 14:53:18 -0800
Subject: [PATCH 1503/2720] Remove extra copy of model_hparams. model_hparams
 doesn't seem to be getting modified in _create_modalities, so I don't
 understand the need for copy.copy -- am I missing something?

PiperOrigin-RevId: 229035908
---
 tensor2tensor/data_generators/problem.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 09c0c2bab..42db40c56 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -19,7 +19,6 @@
 from __future__ import print_function
 
 import collections
-import copy
 import os
 import random
 import six
@@ -511,7 +510,6 @@ def get_hparams(self, model_hparams=None):
     if self._was_copy:
       _copy_problem_hparams(hp)
 
-    model_hparams = copy.copy(model_hparams)
     _create_modalities(hp, model_hparams)
 
     self._hparams = hp

From 6052a30767fb2ff28092a208528ed89700c05f68 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 12 Jan 2019 21:13:36 -0800
Subject: [PATCH 1504/2720] Update documentation for
 Problem.example_reading_spec, copied from convex_hull.py

PiperOrigin-RevId: 229050886
---
 tensor2tensor/data_generators/problem.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 42db40c56..b609f6e54 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -314,6 +314,13 @@ def feature_encoders(self, data_dir):
     }
 
   def example_reading_spec(self):
+    """Define how data is serialized to file and read back.
+
+    Returns:
+      data_fields: A dictionary mapping data names to its feature type.
+      data_items_to_decoders: A dictionary mapping data names to TF Example
+         decoders, to be used when reading back TF examples from disk.
+    """
     data_fields = {
         "inputs": tf.VarLenFeature(tf.int64),
         "targets": tf.VarLenFeature(tf.int64)

From c557db2c8b9d13c8d3360e842f8ddc8d2d73d305 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sun, 13 Jan 2019 09:29:09 -0800
Subject: [PATCH 1505/2720] Make preprocess_example_common and
 preprocess_example have the same argument order.

PiperOrigin-RevId: 229086363
---
 tensor2tensor/data_generators/problem.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index b609f6e54..3a1e32400 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -137,7 +137,7 @@ def default_model_hparams():
       data_dir=None)
 
 
-def preprocess_example_common(example, hparams, mode):
+def preprocess_example_common(example, mode, hparams):
   """Preprocessing steps common to all models."""
   if hparams.max_input_seq_length > 0:
     example["inputs"] = example["inputs"][:hparams.max_input_seq_length]
@@ -197,7 +197,7 @@ class Problem(object):
     * example_reading_spec
         - Specify the names and types of the features on disk.
         - Specify tf.contrib.slim.tfexample_decoder
-    * preprocess_example(example, mode)
+    * preprocess_example(example, mode, hparams)
         - Preprocess the example feature dict from feature name to Tensor or
           SparseTensor.
         - Used in training, eval, and inference (specified by mode).
@@ -342,7 +342,7 @@ def preprocess_example(self, example, mode, hparams):
     Returns:
       dict or Dataset
     """
-    return preprocess_example_common(example, hparams, mode)
+    return preprocess_example_common(example, mode, hparams)
 
   def eval_metrics(self):
     return [

From 48443714ba89dd0e80472cb19b697fa68e63fc18 Mon Sep 17 00:00:00 2001
From: Utku Evci <evcu@google.com>
Date: Mon, 14 Jan 2019 13:50:03 -0800
Subject: [PATCH 1506/2720] Making normalization step at imagenet
 pre-processing optional. So that we can normalize the images according to our
 needs(i.e. vgg preprocessing).

PiperOrigin-RevId: 229247223
---
 tensor2tensor/data_generators/imagenet.py     | 40 ++++++++++++++-----
 .../data_generators/imagenet_test.py          |  6 +++
 2 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 610cab819..71e6807a7 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -97,16 +97,19 @@ def imagenet_pixelrnn_generator(tmp_dir,
       }
 
 
-def imagenet_preprocess_example(example, mode, resize_size=None):
+def imagenet_preprocess_example(example, mode, resize_size=None,
+                                normalize=True):
   """Preprocessing used for Imagenet and similar problems."""
   resize_size = resize_size or [299, 299]
   assert resize_size[0] == resize_size[1]
 
   image = example["inputs"]
   if mode == tf.estimator.ModeKeys.TRAIN:
-    image = preprocess_for_train(image, image_size=resize_size[0])
+    image = preprocess_for_train(image, image_size=resize_size[0],
+                                 normalize=normalize)
   else:
-    image = preprocess_for_eval(image, image_size=resize_size[0])
+    image = preprocess_for_eval(image, image_size=resize_size[0],
+                                normalize=normalize)
 
   example["inputs"] = image
   return example
@@ -142,6 +145,11 @@ def rescale_size(self):
     # return [224, 224]
     raise NotImplementedError()
 
+  @property
+  def normalize_image(self):
+    """Whether the image should be normalized in preprocessing."""
+    return True
+
   def dataset_filename(self):
     return "image_imagenet"  # Reuse Imagenet data.
 
@@ -151,7 +159,8 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
   def preprocess_example(self, example, mode, _):
     return imagenet_preprocess_example(
-        example, mode, resize_size=self.rescale_size)
+        example, mode, resize_size=self.rescale_size,
+        normalize=self.normalize_image)
 
 
 @registry.register_problem
@@ -163,6 +172,16 @@ def rescale_size(self):
     return [224, 224]
 
 
+@registry.register_problem
+class ImageImagenet224NoNormalization(ImageImagenet224):
+  """Imagenet rescaled to 224x224 without normalization."""
+
+  @property
+  def normalize_image(self):
+    """Whether the image should be normalized in preprocessing."""
+    return False
+
+
 @registry.register_problem
 class ImageImagenet256(ImageImagenetRescaled):
   """Imagenet rescaled to 256x256."""
@@ -543,38 +562,37 @@ def _normalize(image):
   return image
 
 
-def preprocess_for_train(image, image_size=224):
+def preprocess_for_train(image, image_size=224, normalize=True):
   """Preprocesses the given image for evaluation.
 
   Args:
     image: `Tensor` representing an image of arbitrary size.
     image_size: int, how large the output image should be.
+    normalize: bool, if True the image is normalized.
 
   Returns:
     A preprocessed image `Tensor`.
   """
   image = _random_crop(image, image_size)
-  image = _normalize(image)
+  if normalize: image = _normalize(image)
   image = _flip(image)
   image = tf.reshape(image, [image_size, image_size, 3])
   return image
 
 
-def preprocess_for_eval(image, image_size=224):
+def preprocess_for_eval(image, image_size=224, normalize=True):
   """Preprocesses the given image for evaluation.
 
   Args:
     image: `Tensor` representing an image of arbitrary size.
     image_size: int, how large the output image should be.
+    normalize: bool, if True the image is normalized.
 
   Returns:
     A preprocessed image `Tensor`.
   """
   image = _do_scale(image, image_size + 32)
-  image = _normalize(image)
+  if normalize: image = _normalize(image)
   image = _center_crop(image, image_size)
   image = tf.reshape(image, [image_size, image_size, 3])
   return image
-
-
-# ==============================================================================
diff --git a/tensor2tensor/data_generators/imagenet_test.py b/tensor2tensor/data_generators/imagenet_test.py
index 3eb671a06..22a731b68 100644
--- a/tensor2tensor/data_generators/imagenet_test.py
+++ b/tensor2tensor/data_generators/imagenet_test.py
@@ -43,6 +43,12 @@ def testImagenetMultiResolutionPreprocessExample(self, resize_method):
     self.assertLen(preprocessed_example, 1)
     self.assertEqual(preprocessed_example["inputs"].shape, (42, 32, 3))
 
+  def testImagenetIsNormalized(self):
+    problem = imagenet.ImageImagenet224()
+    self.assertTrue(problem.normalize_image)
+    problem = imagenet.ImageImagenet224NoNormalization()
+    self.assertFalse(problem.normalize_image)
+
 
 if __name__ == "__main__":
   tf.test.main()

From 21a8b1246b157f31f781372ee51405d1e6b0a0e1 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 15 Jan 2019 01:27:14 +0100
Subject: [PATCH 1507/2720] Planner fixes (#1369)

* Add logging to the evaluator

* Pass the correct num_rollouts to agent

* Add planner_small hparam set

* Pass the discount factor from policy hparams to planner
---
 tensor2tensor/rl/evaluator.py      | 47 ++++++++++++++++++++++--------
 tensor2tensor/rl/evaluator_test.py |  2 +-
 tensor2tensor/rl/rl_utils.py       |  9 +++++-
 3 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 87a8c221e..d654a831b 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -60,6 +60,9 @@
     "planner_hparams_set", "planner_tiny", "Planner hparam set."
 )
 flags.DEFINE_string("planner_hparams", "", "Planner hparam overrides.")
+flags.DEFINE_integer(
+    "log_every_steps", 0, "Log every how many environment steps."
+)
 
 
 @registry.register_hparams
@@ -71,32 +74,46 @@ def planner_tiny():
   )
 
 
+@registry.register_hparams
+def planner_small():
+  return tf.contrib.training.HParams(
+      num_rollouts=16,
+      planning_horizon=16,
+      rollout_agent_type="policy"
+  )
+
+
 def make_agent(
     agent_type, env, policy_hparams, policy_dir, sampling_temp,
     sim_env_kwargs=None, frame_stack_size=None, planning_horizon=None,
-    rollout_agent_type=None
+    rollout_agent_type=None, batch_size=None, num_rollouts=None
 ):
   """Factory function for Agents."""
+  if batch_size is None:
+    batch_size = env.batch_size
   return {
       "random": lambda: rl_utils.RandomAgent(  # pylint: disable=g-long-lambda
-          env.batch_size, env.observation_space, env.action_space
+          batch_size, env.observation_space, env.action_space
       ),
       "policy": lambda: rl_utils.PolicyAgent(  # pylint: disable=g-long-lambda
-          env.batch_size, env.observation_space, env.action_space,
+          batch_size, env.observation_space, env.action_space,
           policy_hparams, policy_dir, sampling_temp
       ),
       "planner": lambda: rl_utils.PlannerAgent(  # pylint: disable=g-long-lambda
-          env.batch_size, make_agent(
-              rollout_agent_type, env, policy_hparams, policy_dir, sampling_temp
+          batch_size, make_agent(
+              rollout_agent_type, env, policy_hparams, policy_dir,
+              sampling_temp, batch_size=num_rollouts
           ), rl_utils.SimulatedBatchGymEnvWithFixedInitialFrames(
               **sim_env_kwargs
           ), lambda env: rl_utils.BatchStackWrapper(env, frame_stack_size),
-          planning_horizon
+          planning_horizon, discount_factor=policy_hparams.gae_gamma
       ),
   }[agent_type]()
 
 
-def make_eval_fn_with_agent(agent_type, planner_hparams, model_dir):
+def make_eval_fn_with_agent(
+    agent_type, planner_hparams, model_dir, log_every_steps=None
+):
   """Returns an out-of-graph eval_fn using the Agent API."""
   def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
     """Eval function."""
@@ -109,16 +126,20 @@ def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
     agent = make_agent(
         agent_type, env, policy_hparams, policy_dir, sampling_temp,
         sim_env_kwargs, loop_hparams.frame_stack_size,
-        planner_hparams.planning_horizon, planner_hparams.rollout_agent_type
+        planner_hparams.planning_horizon, planner_hparams.rollout_agent_type,
+        num_rollouts=planner_hparams.num_rollouts
+    )
+    rl_utils.run_rollouts(
+        env, agent, env.reset(), log_every_steps=log_every_steps
     )
-    rl_utils.run_rollouts(env, agent, env.reset())
     assert len(base_env.current_epoch_rollouts()) == env.batch_size
   return eval_fn
 
 
 def evaluate(
     loop_hparams, planner_hparams, policy_dir, model_dir, eval_metrics_dir,
-    agent_type, eval_with_learner, report_fn=None, report_metric=None
+    agent_type, eval_with_learner, log_every_steps, report_fn=None,
+    report_metric=None
 ):
   """Evaluate."""
   if eval_with_learner:
@@ -131,7 +152,7 @@ def evaluate(
   kwargs = {}
   if not eval_with_learner:
     kwargs["eval_fn"] = make_eval_fn_with_agent(
-        agent_type, planner_hparams, model_dir
+        agent_type, planner_hparams, model_dir, log_every_steps=log_every_steps
     )
   eval_metrics = rl_utils.evaluate_all_configs(
       loop_hparams, policy_dir, **kwargs
@@ -163,9 +184,11 @@ def main(_):
   )
   evaluate(
       loop_hparams, planner_hparams, FLAGS.policy_dir, FLAGS.model_dir,
-      FLAGS.eval_metrics_dir, FLAGS.agent, FLAGS.eval_with_learner
+      FLAGS.eval_metrics_dir, FLAGS.agent, FLAGS.eval_with_learner,
+      FLAGS.log_every_steps if FLAGS.log_every_steps > 0 else None
   )
 
 
 if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/rl/evaluator_test.py b/tensor2tensor/rl/evaluator_test.py
index 62acde315..0de837ea4 100644
--- a/tensor2tensor/rl/evaluator_test.py
+++ b/tensor2tensor/rl/evaluator_test.py
@@ -32,7 +32,7 @@ def test_evaluate_pong_random_agent(self):
     temp_dir = tf.test.get_temp_dir()
     evaluator.evaluate(
         loop_hparams, planner_hparams, temp_dir, temp_dir, temp_dir,
-        agent_type="random", eval_with_learner=False
+        agent_type="random", eval_with_learner=False, log_every_steps=None
     )
 
 
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index c89e52bea..6e751d87b 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -78,6 +78,9 @@ def evaluate_single_config(
     eval_fn=_eval_fn_with_learner
 ):
   """Evaluate the PPO agent in the real environment."""
+  tf.logging.info("Evaluating metric %s", get_metric_name(
+      sampling_temp, max_num_noops, clipped=False
+  ))
   eval_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
   env = setup_env(
       hparams, batch_size=hparams.eval_batch_size, max_num_noops=max_num_noops,
@@ -230,7 +233,8 @@ def absolute_hinge_difference(arr1, arr2, min_diff=10, dtype=np.uint8):
 
 
 def run_rollouts(
-    env, agent, initial_observations, step_limit=None, discount_factor=1.0
+    env, agent, initial_observations, step_limit=None, discount_factor=1.0,
+    log_every_steps=None
 ):
   """Runs a batch of rollouts from given initial observations."""
   num_dones = 0
@@ -265,6 +269,9 @@ def proceed():
     cum_rewards = cum_rewards * discount_factor + rewards
     step_index += 1
 
+    if log_every_steps is not None and step_index % log_every_steps == 0:
+      tf.logging.info("Step %d, mean_score: %f", step_index, cum_rewards.mean())
+
   return (observations, cum_rewards)
 
 

From 6bf4ac1d7e88f5f6a453e4624b2214d4bd3966e0 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Mon, 14 Jan 2019 16:27:34 -0800
Subject: [PATCH 1508/2720] internal merge of PR #1369

PiperOrigin-RevId: 229278674
---
 tensor2tensor/rl/rl_utils.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 6e751d87b..25f733dce 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -269,6 +269,7 @@ def proceed():
     cum_rewards = cum_rewards * discount_factor + rewards
     step_index += 1
 
+    # TODO(afrozm): Clean this up with tf.logging.log_every_n
     if log_every_steps is not None and step_index % log_every_steps == 0:
       tf.logging.info("Step %d, mean_score: %f", step_index, cum_rewards.mean())
 

From f5096a57915abe295b4365121cd0c77e13ae2b22 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 14 Jan 2019 18:57:08 -0800
Subject: [PATCH 1509/2720] Add L2 loss hparams for SAVP.

PiperOrigin-RevId: 229297709
---
 tensor2tensor/models/video/savp_params.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index 4f04fbcdd..834656271 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -50,6 +50,16 @@ def next_frame_savp():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_savp_l2():
+  """SAVP with L2 reconstruction loss."""
+  hparams = next_frame_savp()
+  hparams.modality = {
+      "inputs": modalities.VideoModalityL2Raw,
+      "targets": modalities.VideoModalityL2Raw,
+  }
+
+
 @registry.register_hparams
 def next_frame_savp_vae():
   """SAVP - VAE only model."""

From bfa005773abe31e922941b48b6c0b104c72542ea Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 14 Jan 2019 22:45:11 -0800
Subject: [PATCH 1510/2720] Move static functions out of MultiProblem class and
 document a bit more.

PiperOrigin-RevId: 229315909
---
 .../data_generators/multi_problem.py          | 194 ++++++++++--------
 1 file changed, 111 insertions(+), 83 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 4530aba05..60e969d46 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -35,6 +35,99 @@ class MixingSchedule(object):
   PRETRAIN = "pretrain"
 
 
+def normalize_example_nlp(task, example, is_infer, vocab_type, vocab_offset,
+                          max_input_length, max_target_length,
+                          fixed_train_length):
+  """Normalize the examples from different tasks so they can be merged.
+
+  This function is specific to NLP tasks and normalizes them so that in the
+  end the example only has "targets" and "task_id". For tasks that originally
+  have inputs, this is done by appending task_id to the inputs and prepending
+  targets, so normalized_targets = inputs task_id targets. For classification
+  tasks, targets are constructed by spelling out the class.
+
+  Args:
+    task: the Problem class of the task we are normalizing.
+    example: a dictionary of tensors, the example to normalize.
+    is_infer: bool, whether we are performing inference or not.
+    vocab_type: the type of vocabulary in use.
+    vocab_offset: integer, offset index for subword vocabularies.
+    max_input_length: maximum length to cut inputs to.
+    max_target_length: maximum length to cut targets to.
+    fixed_train_length: set length to this size if > 0.
+
+  Returns:
+    a dictionary of tensors, like example, after normalizing, which in this
+    case means that it only has "targets" and "task_id" as feature.
+  """
+  if task.has_inputs:
+    example["inputs"] = example["inputs"][:-1]  # remove EOS token
+
+  if hasattr(task, "class_labels"):
+    if vocab_type == text_problems.VocabType.CHARACTER:
+      # TODO(urvashik): handle the case where num_labels > 9
+      example["targets"] = tf.cast(discretization.int_to_bit(
+          example["targets"], 1, base=10) + 50, tf.int64)
+      example["targets"] = tf.squeeze(example["targets"], axis=[-1])
+    elif vocab_type == text_problems.VocabType.SUBWORD:
+      example["targets"] = vocab_offset + example["targets"]
+  else:
+    # sequence with inputs and targets eg: summarization
+    if task.has_inputs:
+      if max_input_length > 0:
+        example["inputs"] = example["inputs"][:max_input_length]
+      # Do not truncate targets during inference with beam decoding.
+      if max_target_length > 0 and not is_infer:
+        example["targets"] = example["targets"][:max_target_length]
+
+  def make_constant_shape(x, size):
+    x = x[:size]
+    xlen = tf.shape(x)[0]
+    x = tf.pad(x, [[0, size - xlen]])
+    return tf.reshape(x, [size])
+
+  if task.has_inputs:
+    if is_infer:
+      concat_list = [example["inputs"], [task.task_id]]
+      example["inputs"] = tf.concat(concat_list, axis=0)
+    else:
+      inputs = example.pop("inputs")
+      concat_list = [inputs, [task.task_id], example["targets"]]
+      example["targets"] = tf.concat(concat_list, axis=0)
+      if fixed_train_length > 0:
+        example["targets"] = make_constant_shape(
+            example["targets"], fixed_train_length)
+  else:
+    concat_list = [[task.task_id], example["targets"]]
+    example["targets"] = tf.concat(concat_list, axis=0)
+    if not is_infer and fixed_train_length > 0:
+      example["targets"] = make_constant_shape(
+          example["targets"], fixed_train_length)
+
+  example["task_id"] = tf.constant([task.task_id], dtype=tf.int64)
+  return example
+
+
+def flatten_zip_dataset(*args):
+  """A list of examples to a dataset containing mixed examples.
+
+  Given a list of `n` dataset examples, flatten them by converting
+  each element into a dataset and concatenating them to convert into a
+  single dataset.
+
+  Args:
+    *args: A list containing one example each from `n` different datasets.
+
+  Returns:
+    flattened: A new dataset containing the examples from the list as part
+      of a single dataset.
+  """
+  flattened = tf.data.Dataset.from_tensors(args[0])
+  for ex in args[1:]:
+    flattened = flattened.concatenate(tf.data.Dataset.from_tensors(ex))
+  return flattened
+
+
 class MultiProblem(problem.Problem):
   """MultiProblem base class."""
 
@@ -49,57 +142,16 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
     for task in self.task_list:
       task.generate_data(data_dir, tmp_dir, task_id)
 
-  def add_task_id(self, task, example, encoder, hparams, is_infer):
-    """Convert example to code switching mode by adding a task id."""
-    if task.has_inputs:
-      example["inputs"] = example["inputs"][:-1]  # remove EOS token
-
-    if hasattr(task, "class_labels"):
-      if self.vocab_type == text_problems.VocabType.CHARACTER:
-        # TODO(urvashik): handle the case where num_labels > 9
-        example["targets"] = tf.cast(discretization.int_to_bit(
-            example["targets"], 1, base=10) + 50, tf.int64)
-        example["targets"] = tf.squeeze(example["targets"], axis=[-1])
-      elif self.vocab_type == text_problems.VocabType.SUBWORD:
-        offset = encoder.vocab_size + len(self.task_list)
-        example["targets"] = offset + example["targets"]
-    else:
-      # sequence with inputs and targets eg: summarization
-      if task.has_inputs:
-        if hparams.multiproblem_max_input_length > 0:
-          example["inputs"] = example[
-              "inputs"][:hparams.multiproblem_max_input_length]
-        # Do not truncate targets during inference with beam decoding.
-        if hparams.multiproblem_max_target_length > 0 and not is_infer:
-          example["targets"] = example[
-              "targets"][:hparams.multiproblem_max_target_length]
-
-    def make_constant_shape(x, size):
-      x = x[:size]
-      xlen = tf.shape(x)[0]
-      x = tf.pad(x, [[0, size - xlen]])
-      return tf.reshape(x, [size])
-
-    if task.has_inputs:
-      if is_infer:
-        concat_list = [example["inputs"], [task.task_id]]
-        example["inputs"] = tf.concat(concat_list, axis=0)
-      else:
-        inputs = example.pop("inputs")
-        concat_list = [inputs, [task.task_id], example["targets"]]
-        example["targets"] = tf.concat(concat_list, axis=0)
-        if hparams.multiproblem_fixed_train_length > 0:
-          example["targets"] = make_constant_shape(
-              example["targets"], hparams.multiproblem_fixed_train_length)
-    else:
-      concat_list = [[task.task_id], example["targets"]]
-      example["targets"] = tf.concat(concat_list, axis=0)
-      if not is_infer and hparams.multiproblem_fixed_train_length > 0:
-        example["targets"] = make_constant_shape(
-            example["targets"], hparams.multiproblem_fixed_train_length)
-
-    example["task_id"] = tf.constant([task.task_id], dtype=tf.int64)
-    return example
+  def normalize_example(self, task, example, encoder, hparams, is_infer):
+    """Normalize the examples from different tasks so they can be merged."""
+    # Here we use the default function for NLP tasks that makes everything
+    # a part of "targets" feature. Override in your subclasses for other uses.
+    vocab_offset = encoder.vocab_size + len(self.task_list)
+    return normalize_example_nlp(
+        task, example, is_infer, self.vocab_type, vocab_offset,
+        hparams.multiproblem_max_input_length,
+        hparams.multiproblem_max_target_length,
+        hparams.multiproblem_fixed_train_length)
 
   def filepattern(self, data_dir, mode, shard=None):
     tf.logging.info("Generating multi problem filepattern")
@@ -108,7 +160,6 @@ def filepattern(self, data_dir, mode, shard=None):
   def get_hparams(self, model_hparams=None):
     if self._hparams is not None:
       return self._hparams
-
     self._hparams = self.task_list[0].get_hparams(model_hparams)
     # Increase the vocab size to account for task ids and modify the modality.
     vocab_size_inc = len(self.task_list)
@@ -123,30 +174,8 @@ def get_hparams(self, model_hparams=None):
     self._hparams.vocab_size["targets"] = new_vocab_size
     self._hparams.modality["targets"] = modalities.SymbolModality(
         model_hparams, self._hparams.vocab_size["targets"])
-
     return self._hparams
 
-  def flatten_zip(self, *args):
-    """A list of examples to a dataset containing mixed examples.
-
-    Given a list of `n` dataset examples, flatten them by converting
-    each element into a dataset and concatenating them to convert into a
-    single dataset.
-
-    Args:
-      *args: A list containing one example each from `n` different datasets.
-
-    Returns:
-      flattened: A new dataset containing the examples from the list as part
-        of a single dataset.
-    """
-
-    flattened = tf.data.Dataset.from_tensors(args[0])
-    for ex in args[1:]:
-      flattened = flattened.concatenate(tf.data.Dataset.from_tensors(ex))
-
-    return flattened
-
   def dataset(self,
               mode,
               data_dir=None,
@@ -161,18 +190,12 @@ def dataset(self,
               num_partitions=1,
               shuffle_buffer_size=1024,
               max_records=-1):
-
     # A list of datasets corresponding to the tasks in the task_list object
     # that need to be mixed.
     datasets = []
     is_training = mode == tf.estimator.ModeKeys.TRAIN
     is_infer = mode == tf.estimator.ModeKeys.PREDICT
-
-    primary_task = self.task_list[0]
-    if primary_task.has_inputs:
-      raise ValueError("Only support language models as primary problem which "
-                       "supplies the vocabulary and the hparams.")
-    enc = primary_task.feature_encoders(data_dir=data_dir)["targets"]
+    enc = self.task_list[0].feature_encoders(data_dir=data_dir)["targets"]
     self.update_task_ids(enc.vocab_size)
 
     for task in self.task_list:
@@ -195,8 +218,13 @@ def dataset(self,
 
       # pylint: disable=cell-var-from-loop
       task_dataset = task_dataset.map(
-          lambda x: self.add_task_id(task, x, enc, hparams, is_infer))
+          lambda x: self.normalize_example(task, x, enc, hparams, is_infer))
+      # pylint: enable=cell-var-from-loop
 
+      # To run evaluation, we want to zip datasets from different tasks,
+      # but zipping will cut off at the shortest dataset in tf.Datasets.
+      # For this reason, we add zero padding to the shorter datasets as
+      # it will be ignored in metrics but it provides space for larger data.
       if not is_training and not is_infer:
         zeros = tf.zeros([self._ADDED_EVAL_COUNT, 1], dtype=tf.int64)
         pad_data = tf.data.Dataset.from_tensor_slices({
@@ -340,7 +368,7 @@ def sample_task(curr_task, num_tasks_left, randnum):
         single_mtl_dataset = datasets[1]
       else:
         single_mtl_dataset = tf.data.Dataset.zip(tuple(datasets)).flat_map(
-            self.flatten_zip)
+            flatten_zip_dataset)
 
     return single_mtl_dataset
 

From 5514b059122370f9d196ac253e6680ee8eac732f Mon Sep 17 00:00:00 2001
From: Mike Dusenberry <dusenberrymw@google.com>
Date: Tue, 15 Jan 2019 12:03:06 -0800
Subject: [PATCH 1511/2720] Improve the initialization scheme for the Bayesian
 Layers.

This updates the TrainableNormal initialization to a Normal with mean `0` and
standard deviation `sqrt(1. / fan_in)`, where `fan_in` is the input
dimensionality of the layer. Additionally, a small amount of noise is added to
each to break symmetry.  This also includes the He et al. 2015 variant that
uses the standard deviation `sqrt(2. / fan_in)` for layers that are followed
by ReLU nonlinearities, and the Glorot 2010 variant that uses `sqrt(2. / fan_in
+ fan_out)`, where `fan_out` is the output dimensionality of the layer.

PiperOrigin-RevId: 229414516
---
 tensor2tensor/layers/bayes.py | 163 +++++++++++++++++++++++++++++++++-
 1 file changed, 160 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index b9be046d5..0386b1340 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import math
+
 import tensorflow as tf
 import tensorflow_probability as tfp
 
@@ -125,6 +127,99 @@ def get_config(self):
     }
 
 
+# From `tensorflow/python/ops/init_ops.py`
+def _compute_fans(shape):
+  """Computes the number of input and output units for a weight shape.
+
+  Args:
+    shape: Integer shape tuple or TF tensor shape.
+
+  Returns:
+    A tuple of scalars (fan_in, fan_out).
+  """
+  if len(shape) < 1:  # Just to avoid errors for constants.
+    fan_in = fan_out = 1
+  elif len(shape) == 1:
+    fan_in = fan_out = shape[0]
+  elif len(shape) == 2:
+    fan_in = shape[0]
+    fan_out = shape[1]
+  else:
+    # Assuming convolution kernels (2D, 3D, or more).
+    # kernel shape: (..., input_depth, depth)
+    receptive_field_size = 1.
+    for dim in shape[:-2]:
+      receptive_field_size *= dim
+    fan_in = shape[-2] * receptive_field_size
+    fan_out = shape[-1] * receptive_field_size
+  return fan_in, fan_out
+
+
+class ScaledNormalStdDev(tf.keras.initializers.VarianceScaling):
+  """Initializer capable of adapting its scale to the shape of weights tensors.
+
+  This initializes the standard deviation parameter of a Trainable Normal
+  distribution with a scale based on the shape of the weights tensor.
+  Additionally, A small amount of noise will be added to break weigh symmetry.
+
+  With `distribution="truncated_normal" or "untruncated_normal"`, the standard
+  deviation (after truncation, if used) is `stddev = sqrt(scale / n)`, where n
+  is:
+    - number of input units in the weight tensor, if mode = "fan_in"
+    - number of output units, if mode = "fan_out"
+    - average of the numbers of input and output units, if mode = "fan_avg"
+
+  Args:
+    scale: Scaling factor (positive float).
+    mode: One of "fan_in", "fan_out", "fan_avg".
+    distribution: Random distribution to use. One of "truncated_normal", or
+      "untruncated_normal".
+    seed: A Python integer. Used to create random seeds. See
+      `tf.set_random_seed`
+      for behavior.
+    dtype: The data type. Only floating point types are supported.
+
+  Raises:
+    ValueError: In case of an invalid value for the "scale", mode" or
+      "distribution" arguments.
+  """
+
+  def __init__(self,
+               scale=1.0,
+               mode='fan_in',
+               distribution='untruncated_normal',
+               seed=None,
+               dtype=tf.float32):
+    distribution = distribution.lower()
+    if distribution not in {'truncated_normal', 'untruncated_normal'}:
+      raise ValueError('Invalid `distribution` argument:', distribution)
+    super(ScaledNormalStdDev, self).__init__(scale=scale, mode=mode,
+                                             distribution=distribution,
+                                             seed=seed, dtype=dtype)
+
+  def __call__(self, shape, dtype=None, partition_info=None):
+    if dtype is None:
+      dtype = self.dtype
+    scale = self.scale
+    scale_shape = shape
+    if partition_info is not None:
+      scale_shape = partition_info.full_shape
+    fan_in, fan_out = _compute_fans(scale_shape)
+    if self.mode == 'fan_in':
+      scale /= max(1., fan_in)
+    elif self.mode == 'fan_out':
+      scale /= max(1., fan_out)
+    else:
+      scale /= max(1., (fan_in + fan_out) / 2.)
+    if self.distribution == 'truncated_normal':
+      # constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.)
+      stddev = math.sqrt(scale) / .87962566103423978
+    else:  # self.distribution == 'untruncated_normal':
+      stddev = math.sqrt(scale)
+    return tf.random.truncated_normal(shape, mean=stddev, stddev=stddev*0.1,
+                                      dtype=dtype)
+
+
 # TODO(dusenberrymw): Restructure the implementation of a trainable initializer
 # such that callers do not need to have type-conditional logic.
 class TrainableInitializer(tf.keras.initializers.Initializer):
@@ -146,9 +241,9 @@ class TrainableNormal(TrainableInitializer):
   """Random normal op as an initializer with trainable mean and stddev."""
 
   def __init__(self,
-               mean_initializer=tf.random_normal_initializer(stddev=0.1),
-               stddev_initializer=tf.random_uniform_initializer(
-                   minval=1e-5, maxval=0.1),
+               mean_initializer=tf.keras.initializers.truncated_normal(
+                   stddev=1e-5),
+               stddev_initializer=ScaledNormalStdDev(),
                mean_regularizer=None,
                stddev_regularizer=None,
                mean_constraint=None,
@@ -221,10 +316,72 @@ def get_config(self):
     }
 
 
+class TrainableHeNormal(TrainableNormal):
+  """Trainable normal initialized per He et al. 2015, given a ReLU nonlinearity.
+
+  The distribution is initialized to a Normal scaled by `sqrt(2 / fan_in)`,
+  where `fan_in` is the number of input units. A ReLU nonlinearity is assumed
+  for this initialization scheme.
+
+  References:
+    He K, Zhang X, Ren S, Sun J. Delving deep into rectifiers: Surpassing
+    human-level performance on imagenet classification. In Proceedings of the
+    IEEE international conference on computer vision 2015 (pp. 1026-1034).
+    https://arxiv.org/abs/1502.01852
+  """
+
+  def __init__(self, seed=None, dtype=tf.float32):
+    super(TrainableHeNormal, self).__init__(
+        stddev_initializer=ScaledNormalStdDev(scale=2.0, seed=seed,
+                                              dtype=dtype),
+        seed=seed, dtype=dtype)
+
+  def get_config(self):
+    return {
+        'seed': self.seed,
+        'dtype': self.dtype.name
+    }
+
+
+class TrainableGlorotNormal(TrainableNormal):
+  """Trainable normal initialized per Glorot and Bengio, 2010.
+
+  The distribution is initialized to a Normal scaled by `sqrt(2 / fan_in +
+  fan_out)`, where `fan_in` is the number of input units and `fan_out` is the
+  number of output units.
+
+  References:
+    Glorot X, Bengio Y. Understanding the difficulty of training deep
+    feedforward neural networks. In Proceedings of the thirteenth international
+    conference on artificial intelligence and statistics 2010 Mar 31 (pp.
+    249-256). http://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf
+  """
+
+  def __init__(self, seed=None, dtype=tf.float32):
+    super(TrainableGlorotNormal, self).__init__(
+        stddev_initializer=ScaledNormalStdDev(mode='fan_avg', seed=seed,
+                                              dtype=dtype),
+        seed=seed, dtype=dtype)
+
+  def get_config(self):
+    return {
+        'seed': self.seed,
+        'dtype': self.dtype.name
+    }
+
+
 def trainable_normal():  # alias, following tf.keras.initializers
   return TrainableNormal()
 
 
+def trainable_he_normal():  # alias, following tf.keras.initializers
+  return TrainableHeNormal()
+
+
+def trainable_glorot_normal():  # alias, following tf.keras.initializers
+  return TrainableGlorotNormal()
+
+
 class NormalKLDivergence(tf.keras.regularizers.Regularizer):
   """KL divergence regularizer from one normal distribution to another."""
 

From d9f65a01fde6d3050330394e4240b4bd9ade3e67 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 15 Jan 2019 12:09:13 -0800
Subject: [PATCH 1512/2720] Bug fix for SAVP L2.

PiperOrigin-RevId: 229415828
---
 tensor2tensor/models/video/savp_params.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index 834656271..fd7a3d346 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -58,6 +58,7 @@ def next_frame_savp_l2():
       "inputs": modalities.VideoModalityL2Raw,
       "targets": modalities.VideoModalityL2Raw,
   }
+  return hparams
 
 
 @registry.register_hparams

From 4c6987617716b7dc5ecb2aa9a46d75cbb53d1fd6 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 15 Jan 2019 23:50:51 +0100
Subject: [PATCH 1513/2720] Planner eval (#1371)

* Decouple planner batch size from the number of rollouts

* Allow skipping iterating over eval_max_num_noops
---
 tensor2tensor/rl/evaluator.py | 17 +++++++++++------
 tensor2tensor/rl/rl_utils.py  | 14 ++++++++++----
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index d654a831b..d77272f30 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -71,6 +71,7 @@ def planner_tiny():
       num_rollouts=1,
       planning_horizon=2,
       rollout_agent_type="random",
+      batch_size=1,
   )
 
 
@@ -79,14 +80,16 @@ def planner_small():
   return tf.contrib.training.HParams(
       num_rollouts=16,
       planning_horizon=16,
-      rollout_agent_type="policy"
+      rollout_agent_type="policy",
+      batch_size=16,
   )
 
 
 def make_agent(
     agent_type, env, policy_hparams, policy_dir, sampling_temp,
     sim_env_kwargs=None, frame_stack_size=None, planning_horizon=None,
-    rollout_agent_type=None, batch_size=None, num_rollouts=None
+    rollout_agent_type=None, batch_size=None, num_rollouts=None,
+    inner_batch_size=None
 ):
   """Factory function for Agents."""
   if batch_size is None:
@@ -102,11 +105,12 @@ def make_agent(
       "planner": lambda: rl_utils.PlannerAgent(  # pylint: disable=g-long-lambda
           batch_size, make_agent(
               rollout_agent_type, env, policy_hparams, policy_dir,
-              sampling_temp, batch_size=num_rollouts
+              sampling_temp, batch_size=inner_batch_size
           ), rl_utils.SimulatedBatchGymEnvWithFixedInitialFrames(
               **sim_env_kwargs
           ), lambda env: rl_utils.BatchStackWrapper(env, frame_stack_size),
-          planning_horizon, discount_factor=policy_hparams.gae_gamma
+          num_rollouts, planning_horizon,
+          discount_factor=policy_hparams.gae_gamma
       ),
   }[agent_type]()
 
@@ -120,14 +124,15 @@ def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
     base_env = env
     env = rl_utils.BatchStackWrapper(env, loop_hparams.frame_stack_size)
     sim_env_kwargs = rl.make_simulated_env_kwargs(
-        base_env, loop_hparams, batch_size=planner_hparams.num_rollouts,
+        base_env, loop_hparams, batch_size=planner_hparams.batch_size,
         model_dir=model_dir
     )
     agent = make_agent(
         agent_type, env, policy_hparams, policy_dir, sampling_temp,
         sim_env_kwargs, loop_hparams.frame_stack_size,
         planner_hparams.planning_horizon, planner_hparams.rollout_agent_type,
-        num_rollouts=planner_hparams.num_rollouts
+        num_rollouts=planner_hparams.num_rollouts,
+        inner_batch_size=planner_hparams.batch_size
     )
     rl_utils.run_rollouts(
         env, agent, env.reset(), log_every_steps=log_every_steps
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 25f733dce..a3fb3525a 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -104,7 +104,7 @@ def evaluate_all_configs(
   # Iterate over all combinations of sampling temperatures and whether to do
   # initial no-ops.
   for sampling_temp in hparams.eval_sampling_temps:
-    for max_num_noops in (hparams.eval_max_num_noops, 0):
+    for max_num_noops in {hparams.eval_max_num_noops, 0}:
       scores = evaluate_single_config(
           hparams, sampling_temp, max_num_noops, agent_model_dir, eval_fn
       )
@@ -371,8 +371,8 @@ class PlannerAgent(BatchAgent):
   """Agent based on temporal difference planning."""
 
   def __init__(
-      self, batch_size, rollout_agent, sim_env, wrapper_fn, planning_horizon,
-      discount_factor=1.0
+      self, batch_size, rollout_agent, sim_env, wrapper_fn, num_rollouts,
+      planning_horizon, discount_factor=1.0
   ):
     super(PlannerAgent, self).__init__(
         batch_size, rollout_agent.observation_space, rollout_agent.action_space
@@ -380,6 +380,7 @@ def __init__(
     self._rollout_agent = rollout_agent
     self._sim_env = sim_env
     self._wrapped_env = wrapper_fn(sim_env)
+    self._num_batches = num_rollouts // rollout_agent.batch_size
     self._discount_factor = discount_factor
     self._planning_horizon = planning_horizon
 
@@ -405,10 +406,15 @@ def run_batch_from(observation, action):
       )
       return total_values.mean()
 
+    def run_batches_from(observation, action):
+      return sum(
+          run_batch_from(observation, action) for _ in range(self._num_batches)
+      ) / self._num_batches
+
     def choose_best_action(observation):
       return max(
           range(self.action_space.n),
-          key=(lambda action: run_batch_from(observation, action))
+          key=(lambda action: run_batches_from(observation, action))
       )
 
     return np.array(list(map(choose_best_action, observations)))

From da84863b200ddf59af00d12b55aee878db234662 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 15 Jan 2019 14:51:09 -0800
Subject: [PATCH 1514/2720] internal merge of PR #1371

PiperOrigin-RevId: 229445045
---
 tensor2tensor/rl/rl_utils.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index a3fb3525a..7add43b91 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -37,10 +37,6 @@
 import tensorflow as tf
 
 
-flags = tf.flags
-FLAGS = flags.FLAGS
-
-
 def compute_mean_reward(rollouts, clipped):
   """Calculate mean rewards from given epoch."""
   reward_name = "reward" if clipped else "unclipped_reward"
@@ -104,7 +100,7 @@ def evaluate_all_configs(
   # Iterate over all combinations of sampling temperatures and whether to do
   # initial no-ops.
   for sampling_temp in hparams.eval_sampling_temps:
-    for max_num_noops in {hparams.eval_max_num_noops, 0}:
+    for max_num_noops in [hparams.eval_max_num_noops, 0]:
       scores = evaluate_single_config(
           hparams, sampling_temp, max_num_noops, agent_model_dir, eval_fn
       )

From 3e246128ca66b5ce7222d93cb511b0ab92e18a90 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Wed, 16 Jan 2019 02:06:11 +0100
Subject: [PATCH 1515/2720] Planner videos (#1372)

* Decouple planner batch size from the number of rollouts

* Allow skipping iterating over eval_max_num_noops

* Dump debug videos from planner
---
 tensor2tensor/rl/evaluator.py      | 34 +++++++++----
 tensor2tensor/rl/evaluator_test.py |  3 +-
 tensor2tensor/rl/rl_utils.py       | 79 ++++++++++++++++++++++++++----
 3 files changed, 97 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index d77272f30..726382f71 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -28,6 +28,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl  # pylint: disable=unused-import
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl import trainer_model_based_params  # pylint: disable=unused-import
@@ -63,6 +64,9 @@
 flags.DEFINE_integer(
     "log_every_steps", 0, "Log every how many environment steps."
 )
+flags.DEFINE_string(
+    "debug_video_path", "", "Path to save the planner debug video at."
+)
 
 
 @registry.register_hparams
@@ -89,8 +93,8 @@ def make_agent(
     agent_type, env, policy_hparams, policy_dir, sampling_temp,
     sim_env_kwargs=None, frame_stack_size=None, planning_horizon=None,
     rollout_agent_type=None, batch_size=None, num_rollouts=None,
-    inner_batch_size=None
-):
+    inner_batch_size=None, video_writer=None
+    inner_batch_size=None):
   """Factory function for Agents."""
   if batch_size is None:
     batch_size = env.batch_size
@@ -110,13 +114,14 @@ def make_agent(
               **sim_env_kwargs
           ), lambda env: rl_utils.BatchStackWrapper(env, frame_stack_size),
           num_rollouts, planning_horizon,
-          discount_factor=policy_hparams.gae_gamma
+          discount_factor=policy_hparams.gae_gamma, video_writer=video_writer
       ),
   }[agent_type]()
 
 
 def make_eval_fn_with_agent(
-    agent_type, planner_hparams, model_dir, log_every_steps=None
+    agent_type, planner_hparams, model_dir, log_every_steps=None,
+    video_writer=None
 ):
   """Returns an out-of-graph eval_fn using the Agent API."""
   def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
@@ -132,7 +137,7 @@ def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
         sim_env_kwargs, loop_hparams.frame_stack_size,
         planner_hparams.planning_horizon, planner_hparams.rollout_agent_type,
         num_rollouts=planner_hparams.num_rollouts,
-        inner_batch_size=planner_hparams.batch_size
+        inner_batch_size=planner_hparams.batch_size, video_writer=video_writer
     )
     rl_utils.run_rollouts(
         env, agent, env.reset(), log_every_steps=log_every_steps
@@ -143,8 +148,8 @@ def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
 
 def evaluate(
     loop_hparams, planner_hparams, policy_dir, model_dir, eval_metrics_dir,
-    agent_type, eval_with_learner, log_every_steps, report_fn=None,
-    report_metric=None
+    agent_type, eval_with_learner, log_every_steps, debug_video_path,
+    report_fn=None, report_metric=None
 ):
   """Evaluate."""
   if eval_with_learner:
@@ -156,14 +161,24 @@ def evaluate(
   eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
   kwargs = {}
   if not eval_with_learner:
+    if debug_video_path:
+      video_writer = common_video.WholeVideoWriter(
+          fps=10, output_path=debug_video_path, file_format="avi"
+      )
+    else:
+      video_writer = None
     kwargs["eval_fn"] = make_eval_fn_with_agent(
-        agent_type, planner_hparams, model_dir, log_every_steps=log_every_steps
+        agent_type, planner_hparams, model_dir, log_every_steps=log_every_steps,
+        video_writer=video_writer
     )
   eval_metrics = rl_utils.evaluate_all_configs(
       loop_hparams, policy_dir, **kwargs
   )
   rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, 0)
 
+  if video_writer is not None:
+    video_writer.finish_to_disk()
+
   # Report metrics
   if report_fn:
     if report_metric == "mean_reward":
@@ -190,7 +205,8 @@ def main(_):
   evaluate(
       loop_hparams, planner_hparams, FLAGS.policy_dir, FLAGS.model_dir,
       FLAGS.eval_metrics_dir, FLAGS.agent, FLAGS.eval_with_learner,
-      FLAGS.log_every_steps if FLAGS.log_every_steps > 0 else None
+      FLAGS.log_every_steps if FLAGS.log_every_steps > 0 else None,
+      debug_video_path=FLAGS.debug_video_path
   )
 
 
diff --git a/tensor2tensor/rl/evaluator_test.py b/tensor2tensor/rl/evaluator_test.py
index 0de837ea4..69c203325 100644
--- a/tensor2tensor/rl/evaluator_test.py
+++ b/tensor2tensor/rl/evaluator_test.py
@@ -32,7 +32,8 @@ def test_evaluate_pong_random_agent(self):
     temp_dir = tf.test.get_temp_dir()
     evaluator.evaluate(
         loop_hparams, planner_hparams, temp_dir, temp_dir, temp_dir,
-        agent_type="random", eval_with_learner=False, log_every_steps=None
+        agent_type="random", eval_with_learner=False, log_every_steps=None,
+        debug_video_path=""
     )
 
 
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 7add43b91..4de677907 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -30,6 +30,8 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl.dopamine_connector import DQNLearner
 from tensor2tensor.rl.envs.simulated_batch_gym_env import SimulatedBatchGymEnv
+from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image
+from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw
 from tensor2tensor.rl.ppo_learner import PPOLearner
 from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import trainer_lib
@@ -228,9 +230,35 @@ def absolute_hinge_difference(arr1, arr2, min_diff=10, dtype=np.uint8):
   return np.maximum(diff - min_diff, 0).astype(dtype)
 
 
+# TODO(koz4k): Use this function in player and all debug videos.
+def augment_observation(
+    observation, reward, cum_reward, frame_index, bar_color=None,
+    header_height=27
+):
+  """Augments an observation with debug info."""
+  img = PIL_Image().new(
+      "RGB", (observation.shape[1], header_height,)
+  )
+  draw = PIL_ImageDraw().Draw(img)
+  draw.text(
+      (1, 0), "c:{:3}, r:{:3}".format(int(cum_reward), int(reward)),
+      fill=(255, 0, 0)
+  )
+  draw.text(
+      (1, 15), "f:{:3}".format(int(frame_index)),
+      fill=(255, 0, 0)
+  )
+  header = np.asarray(img)
+  del img
+  header.setflags(write=1)
+  if bar_color is not None:
+    header[0, :, :] = bar_color
+  return np.concatenate([header, observation], axis=0)
+
+
 def run_rollouts(
     env, agent, initial_observations, step_limit=None, discount_factor=1.0,
-    log_every_steps=None
+    log_every_steps=None, video_writer=None
 ):
   """Runs a batch of rollouts from given initial observations."""
   num_dones = 0
@@ -239,6 +267,15 @@ def run_rollouts(
   step_index = 0
   cum_rewards = 0
 
+  if video_writer is not None:
+    obs_stack = initial_observations[0]
+    for (i, ob) in enumerate(obs_stack):
+      debug_frame = augment_observation(
+          ob, reward=0, cum_reward=0, frame_index=(-len(obs_stack) + i + 1),
+          bar_color=(0, 255, 0)
+      )
+      video_writer.write(debug_frame)
+
   def proceed():
     if step_limit is not None:
       return step_index < step_limit
@@ -265,6 +302,14 @@ def proceed():
     cum_rewards = cum_rewards * discount_factor + rewards
     step_index += 1
 
+    if video_writer is not None:
+      ob = observations[0, -1]
+      debug_frame = augment_observation(
+          ob, reward=rewards[0], cum_reward=cum_rewards[0],
+          frame_index=(step_index + 1), bar_color=(255, 0, 0)
+      )
+      video_writer.write(debug_frame)
+
     # TODO(afrozm): Clean this up with tf.logging.log_every_n
     if log_every_steps is not None and step_index % log_every_steps == 0:
       tf.logging.info("Step %d, mean_score: %f", step_index, cum_rewards.mean())
@@ -368,7 +413,7 @@ class PlannerAgent(BatchAgent):
 
   def __init__(
       self, batch_size, rollout_agent, sim_env, wrapper_fn, num_rollouts,
-      planning_horizon, discount_factor=1.0
+      planning_horizon, discount_factor=1.0, video_writer=None
   ):
     super(PlannerAgent, self).__init__(
         batch_size, rollout_agent.observation_space, rollout_agent.action_space
@@ -379,9 +424,13 @@ def __init__(
     self._num_batches = num_rollouts // rollout_agent.batch_size
     self._discount_factor = discount_factor
     self._planning_horizon = planning_horizon
+    self._video_writer = video_writer
 
   def act(self, observations):
-    def run_batch_from(observation, action):
+    # Randomly choose an action to be recorded.
+    recorded_action = self.action_space.sample()
+
+    def run_batch_from(observation, action, planner_index, batch_index):
       """Run a batch of actions."""
       self._sim_env.initial_frames = np.array(
           [observation] * self._sim_env.batch_size
@@ -393,7 +442,13 @@ def run_batch_from(observation, action):
       (final_observations, cum_rewards) = run_rollouts(
           self._wrapped_env, self._rollout_agent, initial_observations,
           discount_factor=self._discount_factor,
-          step_limit=self._planning_horizon
+          step_limit=self._planning_horizon,
+          video_writer=(
+              self._video_writer
+              if planner_index == 0 and batch_index == 0 and
+              action == recorded_action
+              else None
+          )
       )
       values = self._rollout_agent.estimate_value(final_observations)
       total_values = (
@@ -402,18 +457,24 @@ def run_batch_from(observation, action):
       )
       return total_values.mean()
 
-    def run_batches_from(observation, action):
+    def run_batches_from(observation, action, planner_index):
       return sum(
-          run_batch_from(observation, action) for _ in range(self._num_batches)
+          run_batch_from(observation, action, planner_index, i)
+          for i in range(self._num_batches)
       ) / self._num_batches
 
-    def choose_best_action(observation):
+    def choose_best_action(observation, planner_index):
       return max(
           range(self.action_space.n),
-          key=(lambda action: run_batches_from(observation, action))
+          key=(lambda action: run_batches_from(
+              observation, action, planner_index
+          ))
       )
 
-    return np.array(list(map(choose_best_action, observations)))
+    return np.array([
+        choose_best_action(observation, i)
+        for (i, observation) in enumerate(observations)
+    ])
 
 
 # TODO(koz4k): Unify interfaces of batch envs.

From bda17418b8ee2f30b7387f69ceecd4272c49b06d Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 15 Jan 2019 17:06:29 -0800
Subject: [PATCH 1516/2720] internal merge of PR #1372

PiperOrigin-RevId: 229469197
---
 tensor2tensor/rl/evaluator.py |  3 +--
 tensor2tensor/rl/rl_utils.py  | 15 ++++++---------
 2 files changed, 7 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 726382f71..f5afd1e22 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -93,8 +93,7 @@ def make_agent(
     agent_type, env, policy_hparams, policy_dir, sampling_temp,
     sim_env_kwargs=None, frame_stack_size=None, planning_horizon=None,
     rollout_agent_type=None, batch_size=None, num_rollouts=None,
-    inner_batch_size=None, video_writer=None
-    inner_batch_size=None):
+    inner_batch_size=None, video_writer=None):
   """Factory function for Agents."""
   if batch_size is None:
     batch_size = env.batch_size
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 4de677907..a49aa8a90 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -29,9 +29,9 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl.dopamine_connector import DQNLearner
-from tensor2tensor.rl.envs.simulated_batch_gym_env import SimulatedBatchGymEnv
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw
+from tensor2tensor.rl.envs.simulated_batch_gym_env import SimulatedBatchGymEnv
 from tensor2tensor.rl.ppo_learner import PPOLearner
 from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import trainer_lib
@@ -439,17 +439,14 @@ def run_batch_from(observation, action, planner_index, batch_index):
       (initial_observations, initial_rewards, _) = self._wrapped_env.step(
           np.array([action] * self._wrapped_env.batch_size)
       )
+      writer = None
+      if planner_index == 0 and batch_index == 0 and action == recorded_action:
+        writer = self._video_writer
       (final_observations, cum_rewards) = run_rollouts(
           self._wrapped_env, self._rollout_agent, initial_observations,
           discount_factor=self._discount_factor,
           step_limit=self._planning_horizon,
-          video_writer=(
-              self._video_writer
-              if planner_index == 0 and batch_index == 0 and
-              action == recorded_action
-              else None
-          )
-      )
+          video_writer=writer)
       values = self._rollout_agent.estimate_value(final_observations)
       total_values = (
           initial_rewards + self._discount_factor * cum_rewards +
@@ -466,7 +463,7 @@ def run_batches_from(observation, action, planner_index):
     def choose_best_action(observation, planner_index):
       return max(
           range(self.action_space.n),
-          key=(lambda action: run_batches_from(
+          key=(lambda action: run_batches_from(  # pylint: disable=g-long-lambda
               observation, action, planner_index
           ))
       )

From a53aeaebccb1d6a6a4e124caed5bb9297f98d64c Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Wed, 16 Jan 2019 02:55:27 +0100
Subject: [PATCH 1517/2720] Exclude .git/ directory when submitting to Cloud ML
 Engine (#1367)

The `.git` directory isn't used on Cloud ML Engine, therefor we should ignore it in the uploaded archive.

This reduces the size of the `.tar.gz` archive from a fresh clone of `tensor2tensor` from 15 MB to only 3.6 MB.
---
 tensor2tensor/utils/cloud_mlengine.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 4c9c7a701..f7a2043ee 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -232,7 +232,7 @@ def _tar_and_copy(src_dir, target_dir):
   tmp_dir = tempfile.gettempdir().rstrip("/")
   src_base = os.path.basename(src_dir)
   shell_run(
-      "tar -zcf {tmp_dir}/{src_base}.tar.gz -C {src_dir} .",
+      "tar --exclude=.git -zcf {tmp_dir}/{src_base}.tar.gz -C {src_dir} .",
       src_dir=src_dir,
       src_base=src_base,
       tmp_dir=tmp_dir)

From a8cb9f1fc97547aa217dab69e9d5d7a160cfdaf1 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Tue, 15 Jan 2019 17:55:47 -0800
Subject: [PATCH 1518/2720] internal merge of PR #1367

PiperOrigin-RevId: 229475364
---
 tensor2tensor/models/transformer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 027b0d3a8..94c720a3a 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1724,7 +1724,7 @@ def transformer_tall_pretrain_lm_tpu_adafactor_large():
   hparams.batch_size = 4
   hparams.multiproblem_mixing_schedule = "constant"
   # Task order: lm/en-de/en-fr/en-ro/de-en/fr-en/ro-en/cnndm/mnli/squad.
-  hparams.multiproblem_per_task_threshold = "32,8,16,1,8,16,1,2,2,1"
+  hparams.multiproblem_per_task_threshold = "320,80,160,2,80,160,2,20,5,5"
   return hparams
 
 

From a6e418be61e90215d7d0c62b788de837b92b72ff Mon Sep 17 00:00:00 2001
From: Mike Dusenberry <dusenberrymw@google.com>
Date: Tue, 15 Jan 2019 19:57:17 -0800
Subject: [PATCH 1519/2720] Update the bias initializer for the Bayesian LSTM.

This creates the Bayesian equivalent of initializing the bias such that the
forget gate is initialized with a value close to one.

PiperOrigin-RevId: 229487589
---
 tensor2tensor/layers/bayes.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 0386b1340..93a60e684 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -786,6 +786,19 @@ def build(self, input_shape):
 
     if self.use_bias:
       if isinstance(self.bias_initializer, TrainableInitializer):
+        if self.unit_forget_bias:
+          def bias_mean_initializer(_, *args, **kwargs):
+            return tf.concat([
+                tf.keras.initializers.truncated_normal(
+                    stddev=1e-5)((self.units,), *args, **kwargs),
+                tf.keras.initializers.truncated_normal(
+                    mean=1., stddev=1e-5)((self.units,), *args, **kwargs),
+                tf.keras.initializers.truncated_normal(
+                    stddev=1e-5)((self.units * 2,), *args, **kwargs),
+            ], axis=0)
+          self.bias_initializer = TrainableNormal(
+              mean_initializer=bias_mean_initializer)
+
         self.bias_initializer.build(
             [self.units * 4], self.dtype, self.add_weight)
         self.bias = self.bias_initializer()
@@ -797,7 +810,6 @@ def build(self, input_shape):
               'bias', self.bias_initializer, self.bias_regularizer))
       else:
         if self.unit_forget_bias:
-
           def bias_initializer(_, *args, **kwargs):
             return tf.keras.backend.concatenate([
                 self.bias_initializer((self.units,), *args, **kwargs),

From 6a530d35b57fdf5cc76bc744d601b3400dc01188 Mon Sep 17 00:00:00 2001
From: Artit 'Art' Wangperawong <artitw@gmail.com>
Date: Wed, 16 Jan 2019 14:32:31 -0500
Subject: [PATCH 1520/2720] Combine AUT variants into a single function and
 define mix_with_transformer (#1373)

* Combine AUT variants into a single function and define default mix_with_transformer value

* change default mix_with_transformer value back to None
---
 .../research/universal_transformer_util.py    | 583 ++----------------
 1 file changed, 56 insertions(+), 527 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 7ab87b258..c881551f2 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -234,8 +234,7 @@ def add_vanilla_transformer_layer(x, num_layers, name):
     return x
 
   with tf.variable_scope("universal_transformer_%s" % hparams.recurrence_type):
-
-    if "before_ut" in hparams.mix_with_transformer:
+    if hparams.mix_with_transformer and "before_ut" in hparams.mix_with_transformer:
       x = add_vanilla_transformer_layer(x, hparams.num_mixedin_layers,
                                         "before_ut_")
 
@@ -257,7 +256,7 @@ def add_vanilla_transformer_layer(x, num_layers, name):
           hparams.get("use_memory_as_final_state", False)):
         output = extra_output
 
-    if "after_ut" in hparams.mix_with_transformer:
+    if hparams.mix_with_transformer and "after_ut" in hparams.mix_with_transformer:
       output = add_vanilla_transformer_layer(output, hparams.num_mixedin_layers,
                                              "after_ut_")
 
@@ -1038,6 +1037,11 @@ def universal_transformer_act(x, hparams, ffn_unit, attention_unit):
 
   Implementations of all act models are based on craffel@'s cl/160711592.
 
+    (1) Basic AUT based on remainder-distribution ACT (position-wise).
+    (2) AUT with global halting probability (not position-wise).
+    (3) AUT with random halting probability (not position-wise).
+    (4) AUT with final state as accumulation of all states. Similar to the main ACT paper: --> check the issue of differentiability
+    
   Args:
     x: input
     hparams: model hyper-parameters
@@ -1052,65 +1056,30 @@ def universal_transformer_act(x, hparams, ffn_unit, attention_unit):
 
   """
   # TODO(dehghani): Enable pad_remover for the act computations.
-  if hparams.act_type == "basic":
-    return universal_transformer_act_basic(
-        x, hparams, ffn_unit, attention_unit)
-
-  elif hparams.act_type == "accumulated":
-    return universal_transformer_act_accumulated(
-        x, hparams, ffn_unit, attention_unit)
-
-  elif hparams.act_type == "global":
-    return universal_transformer_act_global(
-        x, hparams, ffn_unit, attention_unit)
 
-  elif hparams.act_type == "random":
-    return universal_transformer_act_random(
-        x, hparams, ffn_unit, attention_unit)
-
-  else:
+  if hparams.act_type not in ["basic","global","random","accumulated"]:
     raise ValueError("Unknown act type: %s" % hparams.act_type)
-
-
-def universal_transformer_act_basic(x, hparams, ffn_unit, attention_unit):
-  """Basic universal_transformer with ACT based on remainder-distribution ACT.
-
-  Args:
-    x: input
-    hparams: model hyper-parameters
-    ffn_unit: feed-forward unit
-    attention_unit: multi-head attention unit
-
-  Returns:
-    the output tensor,  (ponder_times, remainders)
-
-  """
-
+    
   state = x
   act_max_steps = hparams.act_max_steps
   threshold = 1.0 - hparams.act_epsilon
-
-  batch_size = tf.shape(state)[0]
-  length = tf.shape(state)[1]
-
+  state_shape_static = state.get_shape()
+
+  state_slice = slice(0,2)
+  if hparams.act_type == "global":
+    state_slice = slice(0,1)
+    
+  # Dynamic shape for update tensors below
+  update_shape = tf.shape(state)[state_slice]
+    
   # Halting probabilities (p_t^n in the paper)
-  halting_probability = tf.zeros(
-      (
-          batch_size,
-          length,
-      ), name="halting_probability")
+  halting_probability = tf.zeros(update_shape, name="halting_probability")
+  
   # Remainders (R(t) in the paper)
-  remainders = tf.zeros(
-      (
-          batch_size,
-          length,
-      ), name="remainder")
+  remainders = tf.zeros(update_shape, name="remainder")
+  
   # Number of updates performed (N(t) in the paper)
-  n_updates = tf.zeros(
-      (
-          batch_size,
-          length,
-      ), name="n_updates")
+  n_updates = tf.zeros(update_shape, name="n_updates")
 
   # Previous cell states (s_t in the paper)
   previous_state = tf.zeros_like(state, name="previous_state")
@@ -1136,19 +1105,30 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
       n_updates: act n_updates
       new_state: new state
     """
-    state_shape = state.get_shape()
-    state = step_preprocess(state, step, hparams)
 
-    with tf.variable_scope("sigmoid_activation_for_pondering"):
-      p = common_layers.dense(
-          state,
-          1,
-          activation=tf.nn.sigmoid,
-          use_bias=True,
-          bias_initializer=tf.constant_initializer(
-              hparams.act_halting_bias_init))
-      p = tf.squeeze(p, axis=-1)
+    state = step_preprocess(state, step, hparams)
 
+    if hparams.act_type == "random":
+        # random as halting probability
+        p = tf.random_uniform(shape=common_layers.shape_list(halting_probability))
+    else:
+        with tf.variable_scope("sigmoid_activation_for_pondering"):
+          p = common_layers.dense(
+              state,
+              1,
+              activation=tf.nn.sigmoid,
+              use_bias=True,
+              bias_initializer=tf.constant_initializer(
+                  hparams.act_halting_bias_init))
+
+          if hparams.act_type == "global":
+            # average over all positions (as a global halting prob)
+            p = tf.reduce_mean(p, axis=1)
+            p = tf.squeeze(p)
+          else:
+            #maintain position-wise probabilities
+            p = tf.squeeze(p, axis=-1)
+            
     # Mask for inputs which have not halted yet
     still_running = tf.cast(tf.less(halting_probability, 1.0), tf.float32)
 
@@ -1181,7 +1161,9 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
     # the remainders when it halted this step
     update_weights = tf.expand_dims(p * still_running + new_halted * remainders,
                                     -1)
-
+    if hparams.act_type == "global":
+        update_weights = tf.expand_dims(update_weights, -1)
+        
     # apply transformation on the state
     transformed_state = state
     for i in range(hparams.num_inrecurrence_layers):
@@ -1190,16 +1172,17 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
 
     # update running part in the weighted state and keep the rest
     new_state = ((transformed_state * update_weights) +
-                 (previous_state * (1 - update_weights)))
+             (previous_state * (1 - update_weights)))
+        
+    if hparams.act_type == "accumulated":
+        # Add in the weighted state
+        new_state = (transformed_state * update_weights) + previous_state
 
     # remind TensorFlow of everything's shape
-    transformed_state.set_shape(state_shape)
+    transformed_state.set_shape(state_shape_static)
     for x in [halting_probability, remainders, n_updates]:
-      x.set_shape([
-          state_shape[0],
-          state_shape[1],
-      ])
-    new_state.set_shape(state_shape)
+      x.set_shape(state_shape_static[state_slice])
+    new_state.set_shape(state_shape_static)
     step += 1
     return (transformed_state, step, halting_probability, remainders, n_updates,
             new_state)
@@ -1226,460 +1209,6 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
 
   return new_state, (ponder_times, remainders)
 
-
-def universal_transformer_act_accumulated(x, hparams, ffn_unit, attention_unit):
-  """The UTAct layer where the final state is the accumulation of all states.
-
-    (similar to the main ACT paper: --> check the issue of differentiability)
-
-  Args:
-    x: input
-    hparams: model hyper-parameters
-    ffn_unit: feed-forward unit
-    attention_unit: multi-head attention unit
-
-  Returns:
-    the output tensor,  (ponder_times, remainders)
-
-  """
-  state = x
-  act_max_steps = hparams.act_max_steps
-  threshold = 1.0 - hparams.act_epsilon
-
-  batch_size = tf.shape(state)[0]
-  length = tf.shape(state)[1]
-
-  # Halting probabilities (p_t^n in the paper)
-  halting_probability = tf.zeros(
-      (
-          batch_size,
-          length,
-      ), name="halting_probability")
-  # Remainders (R(t) in the paper)
-  remainders = tf.zeros(
-      (
-          batch_size,
-          length,
-      ), name="remainder")
-  # Number of updates performed (N(t) in the paper)
-  n_updates = tf.zeros(
-      (
-          batch_size,
-          length,
-      ), name="n_updates")
-
-  # Accumulated cell states (s_t in the paper)
-  accumulated_state = tf.zeros_like(state, name="previous_state")
-  step = tf.constant(0, dtype=tf.int32)
-
-  def ut_function(state, step, halting_probability, remainders, n_updates,
-                  accumulated_state):
-    """Position-wise act.
-
-    Args:
-      state: 3-D Tensor: [batch_size, length, channel]
-      step: indicates number of steps taken so far
-      halting_probability: halting probability
-      remainders: act remainders
-      n_updates: act n_updates
-      accumulated_state: accumulated state
-
-    Returns:
-      transformed_state: transformed state
-      step: step+1
-      halting_probability: halting probability
-      remainders: act remainders
-      n_updates: act n_updates
-      accumulated_state: accumulated state
-    """
-    state_shape = state.get_shape()
-    state = step_preprocess(state, step, hparams)
-
-    with tf.variable_scope("sigmoid_activation_for_pondering"):
-      p = common_layers.dense(
-          state,
-          1,
-          activation=tf.nn.sigmoid,
-          use_bias=True,
-          bias_initializer=tf.constant_initializer(
-              hparams.act_halting_bias_init))
-      p = tf.squeeze(p, axis=-1)
-
-    # Mask for inputs which have not halted yet
-    still_running = tf.cast(tf.less(halting_probability, 1.0), tf.float32)
-
-    # Mask of inputs which halted at this step
-    new_halted = tf.cast(
-        tf.greater(halting_probability + p * still_running, threshold),
-        tf.float32) * still_running
-
-    # Mask of inputs which haven't halted, and didn't halt this step
-    still_running = tf.cast(
-        tf.less_equal(halting_probability + p * still_running, threshold),
-        tf.float32) * still_running
-
-    # Add the halting probability for this step to the halting
-    # probabilities for those input which haven't halted yet
-    halting_probability += p * still_running
-
-    # Compute remainders for the inputs which halted at this step
-    remainders += new_halted * (1 - halting_probability)
-
-    # Add the remainders to those inputs which halted at this step
-    halting_probability += new_halted * remainders
-
-    # Increment n_updates for all inputs which are still running
-    n_updates += still_running + new_halted
-
-    # Compute the weight to be applied to the new state and output
-    # 0 when the input has already halted
-    # p when the input hasn't halted yet
-    # the remainders when it halted this step
-    update_weights = tf.expand_dims(p * still_running + new_halted * remainders,
-                                    -1)
-
-    # apply transformation on the state
-    new_state = state
-    for i in range(hparams.num_inrecurrence_layers):
-      with tf.variable_scope("rec_layer_%d" % i):
-        new_state = ffn_unit(attention_unit(new_state))
-
-    transformed_state = new_state
-
-    # Add in the weighted state
-    accumulated_state = (transformed_state * update_weights) + accumulated_state
-
-    # Remind TensorFlow of everything's shape
-    state.set_shape(state_shape)
-    for x in [halting_probability, remainders, n_updates]:
-      x.set_shape([
-          state_shape[0],
-          state_shape[1],
-      ])
-    accumulated_state.set_shape(state_shape)
-    step += 1
-    return (transformed_state, step, halting_probability, remainders, n_updates,
-            accumulated_state)
-
-  # While loop stops when this predicate is FALSE.
-  # Ie all (probability < 1-eps AND counter < N) are false.
-  def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
-    del u0, u1, u2, u3
-    return tf.reduce_any(
-        tf.logical_and(
-            tf.less(halting_probability, threshold),
-            tf.less(n_updates, act_max_steps)))
-
-  # Do while loop iterations until predicate above is false.
-  (_, _, _, remainder, n_updates, accumulated_state) = tf.while_loop(
-      should_continue, ut_function, (state, step, halting_probability,
-                                     remainders, n_updates, accumulated_state),
-      maximum_iterations=act_max_steps + 1)
-
-  ponder_times = n_updates
-  remainders = remainder
-
-  tf.contrib.summary.scalar("ponder_times", tf.reduce_mean(ponder_times))
-
-  return accumulated_state, (ponder_times, remainders)
-
-
-def universal_transformer_act_global(x, hparams, ffn_unit, attention_unit):
-  """The UTAct  with global halting probability (not position-wise).
-
-  Args:
-    x: input
-    hparams: model hyper-parameters
-    ffn_unit: feed-forward unit
-    attention_unit: multi-head attention unit
-
-  Returns:
-    the output tensor,  (ponder_times, remainders)
-
-  """
-  state = x
-  act_max_steps = hparams.act_max_steps
-  threshold = 1.0 - hparams.act_epsilon
-  act_max_steps = hparams.act_max_steps
-  batch_size = tf.shape(state)[0]
-  state_shape = state.get_shape()
-
-  # Halting probabilities (p_t^n in the paper)
-  halting_probability = tf.zeros((batch_size,), name="halting_probability")
-  # Remainders (R(t) in the paper)
-  remainders = tf.zeros((batch_size,), name="remainder")
-  # Number of updates performed (N(t) in the paper)
-  n_updates = tf.zeros((batch_size,), name="n_updates")
-  # Previous cell states (s_t in the paper)
-  previous_state = tf.zeros_like(state, name="previous_state")
-  step = tf.constant(0, dtype=tf.int32)
-
-  def ut_function(state, step, halting_probability, remainders, n_updates,
-                  previous_state):
-    """implements act (global halting).
-
-    Args:
-      state: 3-D Tensor: [batch_size, length, channel]
-      step: indicates number of steps taken so far
-      halting_probability: halting probability
-      remainders: act remainders
-      n_updates: act n_updates
-      previous_state: previous state
-
-    Returns:
-      transformed_state: transformed state
-      step: step+1
-      halting_probability: halting probability
-      remainders: act remainders
-      n_updates: act n_updates
-      new_state: new state
-
-    """
-
-    state = step_preprocess(state, step, hparams)
-
-    with tf.variable_scope("sigmoid_activation_for_pondering"):
-      p = common_layers.dense(
-          state,
-          1,
-          activation=tf.nn.sigmoid,
-          use_bias=True,
-          bias_initializer=tf.constant_initializer(
-              hparams.act_halting_bias_init))
-      # average over all positions (as a global halting prob)
-      p = tf.reduce_mean(p, axis=1)
-      p = tf.squeeze(p)
-
-    # Mask for inputs which have not halted yet
-    still_running = tf.cast(tf.less(halting_probability, 1.0), tf.float32)
-
-    # Mask of inputs which halted at this step
-    new_halted = tf.cast(
-        tf.greater(halting_probability + p * still_running, threshold),
-        tf.float32) * still_running
-
-    # Mask of inputs which haven't halted, and didn't halt this step
-    still_running = tf.cast(
-        tf.less_equal(halting_probability + p * still_running, threshold),
-        tf.float32) * still_running
-
-    # Add the halting probability for this step to the halting
-    # probabilities for those input which haven't halted yet
-    halting_probability += p * still_running
-
-    # Compute remainders for the inputs which halted at this step
-    remainders += new_halted * (1 - halting_probability)
-
-    # Add the remainders to those inputs which halted at this step
-    halting_probability += new_halted * remainders
-
-    # Increment n_updates for all inputs which are still running
-    n_updates += still_running + new_halted
-
-    # Compute the weight to be applied to the new state and output
-    # 0 when the input has already halted
-    # p when the input hasn't halted yet
-    # the remainders when it halted this step
-    update_weights = tf.expand_dims(
-        tf.expand_dims(p * still_running + new_halted * remainders, -1), -1)
-
-    # apply transformation on the state
-    new_state = state
-    for i in range(hparams.num_inrecurrence_layers):
-      with tf.variable_scope("rec_layer_%d" % i):
-        new_state = ffn_unit(attention_unit(new_state))
-
-    transformed_state = new_state
-
-    # Add in the weighted state
-    new_state = ((transformed_state * update_weights) +
-                 (previous_state * (1 - update_weights)))
-
-    # Remind TensorFlow of everything's shape
-    state.set_shape(state_shape)
-    for x in [halting_probability, remainders, n_updates]:
-      x.set_shape([
-          state_shape[0],
-      ])
-    new_state.set_shape(state_shape)
-
-    step += 1
-    return (transformed_state, step, halting_probability,
-            remainders, n_updates, new_state)
-
-  # While loop stops when this predicate is FALSE.
-  # Ie all (probability < 1-eps AND counter < N) are false.
-  def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
-    del u0, u1, u2, u3
-    return tf.reduce_any(
-        tf.logical_and(
-            tf.less(halting_probability, threshold),
-            tf.less(n_updates, act_max_steps)))
-
-  # Do while loop iterations until predicate above is false.
-  (_, _, _, remainder, n_updates, new_state) = tf.while_loop(
-      should_continue, ut_function,
-      (state, step, halting_probability, remainders, n_updates, previous_state),
-      maximum_iterations=act_max_steps + 1)
-
-  ponder_times = n_updates
-  remainders = remainder
-
-  tf.contrib.summary.scalar("ponder_times", tf.reduce_mean(ponder_times))
-
-  return new_state, (ponder_times, remainders)
-
-
-def universal_transformer_act_random(x, hparams, ffn_unit, attention_unit):
-  """universal_transformer with ACT with random halting probability.
-
-  Args:
-    x: input
-    hparams: model hyper-parameters
-    ffn_unit: feed-forward unit
-    attention_unit: multi-head attention unit
-
-  Returns:
-    the output tensor,  (ponder_times, remainders)
-
-  """
-  state = x
-  act_max_steps = hparams.act_max_steps
-  threshold = 1.0 - hparams.act_epsilon
-
-  batch_size = tf.shape(state)[0]
-  length = tf.shape(state)[1]
-
-  # Halting probabilities (p_t^n in the paper)
-  halting_probability = tf.zeros(
-      (
-          batch_size,
-          length,
-      ), name="halting_probability")
-  # Remainders (R(t) in the paper)
-  remainders = tf.zeros(
-      (
-          batch_size,
-          length,
-      ), name="remainder")
-  # Number of updates performed (N(t) in the paper)
-  n_updates = tf.zeros(
-      (
-          batch_size,
-          length,
-      ), name="n_updates")
-
-  # Previous cell states (s_t in the paper)
-  previous_state = tf.zeros_like(state, name="previous_state")
-  step = tf.constant(0, dtype=tf.int32)
-
-  def ut_function(state, step, halting_probability, remainders, n_updates,
-                  previous_state):
-    """Implements act (position-wise halting).
-
-    Args:
-      state: 3-D Tensor: [batch_size, length, channel]
-      step: indicates number of steps taken so far
-      halting_probability: halting probability
-      remainders: act remainders
-      n_updates: act n_updates
-      previous_state: previous state
-
-    Returns:
-      transformed_state: transformed state
-      step: step+1
-      halting_probability: halting probability
-      remainders: act remainders
-      n_updates: act n_updates
-      new_state: new state
-
-    """
-    state_shape = state.get_shape()
-    state = step_preprocess(state, step, hparams)
-
-    # random as halting probability
-    p = tf.random_uniform(shape=common_layers.shape_list(halting_probability))
-
-    # Mask for inputs which have not halted yet
-    still_running = tf.cast(tf.less(halting_probability, 1.0), tf.float32)
-
-    # Mask of inputs which halted at this step
-    new_halted = tf.cast(
-        tf.greater(halting_probability + p * still_running, threshold),
-        tf.float32) * still_running
-
-    # Mask of inputs which haven't halted, and didn't halt this step
-    still_running = tf.cast(
-        tf.less_equal(halting_probability + p * still_running, threshold),
-        tf.float32) * still_running
-
-    # Add the halting probability for this step to the halting
-    # probabilities for those input which haven't halted yet
-    halting_probability += p * still_running
-
-    # Compute remainders for the inputs which halted at this step
-    remainders += new_halted * (1 - halting_probability)
-
-    # Add the remainders to those inputs which halted at this step
-    halting_probability += new_halted * remainders
-
-    # Increment n_updates for all inputs which are still running
-    n_updates += still_running + new_halted
-
-    # Compute the weight to be applied to the new state and output
-    # 0 when the input has already halted
-    # p when the input hasn't halted yet
-    # the remainders when it halted this step
-    update_weights = tf.expand_dims(p * still_running + new_halted * remainders,
-                                    -1)
-
-    # apply transformation on the state
-    new_state = state
-    for i in range(hparams.num_inrecurrence_layers):
-      with tf.variable_scope("rec_layer_%d" % i):
-        new_state = ffn_unit(attention_unit(new_state))
-
-    transformed_state = new_state
-
-    # update running part in the weighted state and keep the rest
-    new_state = ((transformed_state * update_weights) +
-                 (previous_state * (1 - update_weights)))
-
-    # remind TensorFlow of everything's shape
-    transformed_state.set_shape(state_shape)
-    for x in [halting_probability, remainders, n_updates]:
-      x.set_shape([
-          state_shape[0],
-          state_shape[1],
-      ])
-    new_state.set_shape(state_shape)
-    step += 1
-    return (transformed_state, step,
-            halting_probability, remainders, n_updates, new_state)
-
-  # While loop stops when this predicate is FALSE.
-  # Ie all (probability < 1-eps AND counter < N) are false.
-  def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
-    del u0, u1, u2, u3
-    return tf.reduce_any(
-        tf.logical_and(
-            tf.less(halting_probability, threshold),
-            tf.less(n_updates, act_max_steps)))
-
-  # Do while loop iterations until predicate above is false.
-  (_, _, _, remainder, n_updates, new_state) = tf.while_loop(
-      should_continue, ut_function,
-      (state, step, halting_probability, remainders, n_updates, previous_state),
-      maximum_iterations=act_max_steps + 1)
-
-  ponder_times = n_updates
-  remainders = remainder
-
-  tf.contrib.summary.scalar("ponder_times", tf.reduce_mean(ponder_times))
-
-  return new_state, (ponder_times, remainders)
-
-
 def _ffn_layer_multi_inputs(inputs_list,
                             hparams,
                             ffn_layer_type="dense",

From dd80e2512fe0e688ebdec556560f20888daac1c4 Mon Sep 17 00:00:00 2001
From: Artit 'Art' Wangperawong <artitw@gmail.com>
Date: Wed, 16 Jan 2019 11:32:58 -0800
Subject: [PATCH 1521/2720] internal merge of PR #1373

PiperOrigin-RevId: 229594885
---
 .../research/universal_transformer_util.py    | 92 +++++++++----------
 1 file changed, 46 insertions(+), 46 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index c881551f2..673747e60 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -234,7 +234,8 @@ def add_vanilla_transformer_layer(x, num_layers, name):
     return x
 
   with tf.variable_scope("universal_transformer_%s" % hparams.recurrence_type):
-    if hparams.mix_with_transformer and "before_ut" in hparams.mix_with_transformer:
+    if (hparams.mix_with_transformer and
+        "before_ut" in hparams.mix_with_transformer):
       x = add_vanilla_transformer_layer(x, hparams.num_mixedin_layers,
                                         "before_ut_")
 
@@ -256,7 +257,8 @@ def add_vanilla_transformer_layer(x, num_layers, name):
           hparams.get("use_memory_as_final_state", False)):
         output = extra_output
 
-    if hparams.mix_with_transformer and "after_ut" in hparams.mix_with_transformer:
+    if (hparams.mix_with_transformer and
+        "after_ut" in hparams.mix_with_transformer):
       output = add_vanilla_transformer_layer(output, hparams.num_mixedin_layers,
                                              "after_ut_")
 
@@ -1037,11 +1039,11 @@ def universal_transformer_act(x, hparams, ffn_unit, attention_unit):
 
   Implementations of all act models are based on craffel@'s cl/160711592.
 
-    (1) Basic AUT based on remainder-distribution ACT (position-wise).
-    (2) AUT with global halting probability (not position-wise).
-    (3) AUT with random halting probability (not position-wise).
-    (4) AUT with final state as accumulation of all states. Similar to the main ACT paper: --> check the issue of differentiability
-    
+  (1) Basic AUT based on remainder-distribution ACT (position-wise).
+  (2) AUT with global halting probability (not position-wise).
+  (3) AUT with random halting probability (not position-wise).
+  (4) AUT with final state as accumulation of all states.
+
   Args:
     x: input
     hparams: model hyper-parameters
@@ -1053,31 +1055,28 @@ def universal_transformer_act(x, hparams, ffn_unit, attention_unit):
 
   Raises:
     ValueError: Unknown act type
-
   """
-  # TODO(dehghani): Enable pad_remover for the act computations.
-
-  if hparams.act_type not in ["basic","global","random","accumulated"]:
+  if hparams.act_type not in ["basic", "global", "random", "accumulated"]:
     raise ValueError("Unknown act type: %s" % hparams.act_type)
-    
+
   state = x
   act_max_steps = hparams.act_max_steps
   threshold = 1.0 - hparams.act_epsilon
   state_shape_static = state.get_shape()
 
-  state_slice = slice(0,2)
+  state_slice = slice(0, 2)
   if hparams.act_type == "global":
-    state_slice = slice(0,1)
-    
+    state_slice = slice(0, 1)
+
   # Dynamic shape for update tensors below
   update_shape = tf.shape(state)[state_slice]
-    
+
   # Halting probabilities (p_t^n in the paper)
   halting_probability = tf.zeros(update_shape, name="halting_probability")
-  
+
   # Remainders (R(t) in the paper)
   remainders = tf.zeros(update_shape, name="remainder")
-  
+
   # Number of updates performed (N(t) in the paper)
   n_updates = tf.zeros(update_shape, name="n_updates")
 
@@ -1105,30 +1104,30 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
       n_updates: act n_updates
       new_state: new state
     """
-
     state = step_preprocess(state, step, hparams)
 
     if hparams.act_type == "random":
-        # random as halting probability
-        p = tf.random_uniform(shape=common_layers.shape_list(halting_probability))
+      # random as halting probability
+      p = tf.random_uniform(
+          shape=common_layers.shape_list(halting_probability))
     else:
-        with tf.variable_scope("sigmoid_activation_for_pondering"):
-          p = common_layers.dense(
-              state,
-              1,
-              activation=tf.nn.sigmoid,
-              use_bias=True,
-              bias_initializer=tf.constant_initializer(
-                  hparams.act_halting_bias_init))
-
-          if hparams.act_type == "global":
-            # average over all positions (as a global halting prob)
-            p = tf.reduce_mean(p, axis=1)
-            p = tf.squeeze(p)
-          else:
-            #maintain position-wise probabilities
-            p = tf.squeeze(p, axis=-1)
-            
+      with tf.variable_scope("sigmoid_activation_for_pondering"):
+        p = common_layers.dense(
+            state,
+            1,
+            activation=tf.nn.sigmoid,
+            use_bias=True,
+            bias_initializer=tf.constant_initializer(
+                hparams.act_halting_bias_init))
+
+        if hparams.act_type == "global":
+          # average over all positions (as a global halting prob)
+          p = tf.reduce_mean(p, axis=1)
+          p = tf.squeeze(p)
+        else:
+          # maintain position-wise probabilities
+          p = tf.squeeze(p, axis=-1)
+
     # Mask for inputs which have not halted yet
     still_running = tf.cast(tf.less(halting_probability, 1.0), tf.float32)
 
@@ -1159,11 +1158,11 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
     # 0 when the input has already halted
     # p when the input hasn't halted yet
     # the remainders when it halted this step
-    update_weights = tf.expand_dims(p * still_running + new_halted * remainders,
-                                    -1)
+    update_weights = tf.expand_dims(
+        p * still_running + new_halted * remainders, -1)
     if hparams.act_type == "global":
-        update_weights = tf.expand_dims(update_weights, -1)
-        
+      update_weights = tf.expand_dims(update_weights, -1)
+
     # apply transformation on the state
     transformed_state = state
     for i in range(hparams.num_inrecurrence_layers):
@@ -1172,11 +1171,11 @@ def ut_function(state, step, halting_probability, remainders, n_updates,
 
     # update running part in the weighted state and keep the rest
     new_state = ((transformed_state * update_weights) +
-             (previous_state * (1 - update_weights)))
-        
+                 (previous_state * (1 - update_weights)))
+
     if hparams.act_type == "accumulated":
-        # Add in the weighted state
-        new_state = (transformed_state * update_weights) + previous_state
+      # Add in the weighted state
+      new_state = (transformed_state * update_weights) + previous_state
 
     # remind TensorFlow of everything's shape
     transformed_state.set_shape(state_shape_static)
@@ -1209,6 +1208,7 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
 
   return new_state, (ponder_times, remainders)
 
+
 def _ffn_layer_multi_inputs(inputs_list,
                             hparams,
                             ffn_layer_type="dense",

From 2705dda6ab5e0990974a320dc6e2618adeccff15 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Wed, 16 Jan 2019 21:55:12 +0100
Subject: [PATCH 1522/2720] Planner more fixes (#1377)

* Fix frame numbering in planner debug videos

* Fix BatchStackWrapper resetting

* Allow not iterating over max_num_noops
---
 tensor2tensor/rl/rl_utils.py | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index a49aa8a90..95e392163 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -102,7 +102,9 @@ def evaluate_all_configs(
   # Iterate over all combinations of sampling temperatures and whether to do
   # initial no-ops.
   for sampling_temp in hparams.eval_sampling_temps:
-    for max_num_noops in [hparams.eval_max_num_noops, 0]:
+    # Iterate over a set so that if eval_max_num_noops == 0, there is just one
+    # iteration.
+    for max_num_noops in {hparams.eval_max_num_noops, 0}:
       scores = evaluate_single_config(
           hparams, sampling_temp, max_num_noops, agent_model_dir, eval_fn
       )
@@ -306,7 +308,7 @@ def proceed():
       ob = observations[0, -1]
       debug_frame = augment_observation(
           ob, reward=rewards[0], cum_reward=cum_rewards[0],
-          frame_index=(step_index + 1), bar_color=(255, 0, 0)
+          frame_index=step_index, bar_color=(255, 0, 0)
       )
       video_writer.write(debug_frame)
 
@@ -496,7 +498,10 @@ def close(self):
 
 
 class BatchStackWrapper(BatchWrapper):
-  """Out-of-graph batch stack wrapper."""
+  """Out-of-graph batch stack wrapper.
+
+  Its behavior is consistent with tf_atari_wrappers.StackWrapper.
+  """
 
   def __init__(self, env, stack_size):
     super(BatchStackWrapper, self).__init__(env)
@@ -517,8 +522,14 @@ def reset(self, indices=None):
       indices = range(self.batch_size)
 
     observations = self.env.reset(indices)
-    for (index, observation) in zip(indices, observations):
-      self._history_buffer[index, ...] = [observation] * self.stack_size
+    try:
+      # If we wrap the simulated env, take the initial frames from there.
+      assert self.env.initial_frames.shape[1] == self.stack_size
+      self._history_buffer[...] = self.env.initial_frames
+    except AttributeError:
+      # Otherwise, repeat the first observation stack_size times.
+      for (index, observation) in zip(indices, observations):
+        self._history_buffer[index, ...] = [observation] * self.stack_size
     return self._history_buffer
 
   def step(self, actions):

From a0f9d67013e4d680132b7f5e48685b2104253ddf Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Wed, 16 Jan 2019 12:55:28 -0800
Subject: [PATCH 1523/2720] internal merge of PR #1377

PiperOrigin-RevId: 229609137
---
 tensor2tensor/rl/rl_utils.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 95e392163..2d9d1d046 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -102,9 +102,8 @@ def evaluate_all_configs(
   # Iterate over all combinations of sampling temperatures and whether to do
   # initial no-ops.
   for sampling_temp in hparams.eval_sampling_temps:
-    # Iterate over a set so that if eval_max_num_noops == 0, there is just one
-    # iteration.
-    for max_num_noops in {hparams.eval_max_num_noops, 0}:
+    # Iterate over a set so if eval_max_num_noops == 0 then it's 1 iteration.
+    for max_num_noops in set([hparams.eval_max_num_noops, 0]):
       scores = evaluate_single_config(
           hparams, sampling_temp, max_num_noops, agent_model_dir, eval_fn
       )

From e2c914df28e6d010c4e3e66ada975e07465a3018 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Wed, 16 Jan 2019 22:17:54 +0100
Subject: [PATCH 1524/2720] Pretty print Cloud ML Engine job spec (#1368)

---
 tensor2tensor/utils/cloud_mlengine.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index f7a2043ee..2e2cfb8a4 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -17,6 +17,7 @@
 
 import datetime
 import os
+import pprint
 import shutil
 import subprocess as sp
 import sys
@@ -378,7 +379,7 @@ def launch():
   job_spec = configure_job()
   job_name = job_spec["jobId"]
   tf.logging.info("Launching job %s with ML Engine spec:\n%s", job_name,
-                  job_spec)
+                  pprint.pformat(job_spec))
   assert confirm()
   train_dir = FLAGS.output_dir
   t2t_tar = tar_and_copy_t2t(train_dir)

From d18360c0af81c9211b2b272f76c69a542f64d38c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 16 Jan 2019 13:41:25 -0800
Subject: [PATCH 1525/2720] Fix a typo in discretized_mix_logistic_loss
 docstring.

PiperOrigin-RevId: 229617961
---
 tensor2tensor/layers/common_layers.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index ed026adb6..3b32e1678 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1828,8 +1828,8 @@ def discretized_mix_logistic_loss(pred, labels):
 
   The means tensor is a linear combination of location parameters and previous
   channels. The discretized logistic distribution assigns probability mass to an
-  event P(X=x) via logistic CDFs: P(X <= x + 0.5) - P(X > x - 0.5) for 1 < x <
-  254; P(X <= 0.5) for x = 0; and 1 - P(X > 245.5) for x = 255. Instead of
+  event P(X=x) via logistic CDFs: P(X <= x + 0.5) - P(X < x - 0.5) for 1 < x <
+  254; P(X <= 0.5) for x = 0; and 1 - P(X < 245.5) for x = 255. Instead of
   8-bit inputs, this implementation assumes the events are rescaled to [-1, 1].
 
   Args:

From 659cb7abb40e7ba268450636e7feb797e81b38bf Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 16 Jan 2019 18:10:44 -0800
Subject: [PATCH 1526/2720] Create hparams_lib and add copy_hparams to replace
 deepcopy

PiperOrigin-RevId: 229662503
---
 tensor2tensor/bin/t2t_trainer.py        |  6 +-
 tensor2tensor/utils/hparams_lib.py      | 92 +++++++++++++++++++++++++
 tensor2tensor/utils/hparams_lib_test.py | 44 ++++++++++++
 tensor2tensor/utils/mtf_model.py        |  4 +-
 tensor2tensor/utils/t2t_model.py        | 12 ++--
 tensor2tensor/utils/t2t_model_test.py   |  9 ++-
 tensor2tensor/utils/trainer_lib.py      | 60 ++--------------
 tensor2tensor/utils/trainer_lib_test.py | 11 ---
 8 files changed, 156 insertions(+), 82 deletions(-)
 create mode 100644 tensor2tensor/utils/hparams_lib.py
 create mode 100644 tensor2tensor/utils/hparams_lib_test.py

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 80efb08a9..0dad6b4ee 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -27,6 +27,7 @@
 from tensor2tensor.utils import cloud_mlengine
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
+from tensor2tensor.utils import hparams_lib
 from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
@@ -331,10 +332,9 @@ def save_metadata(hparams):
       f.write(t2t_flags_str)
 
   # Save hparams as hparams.json
-  hp_vals = hparams.values()
+  new_hparams = hparams_lib.copy_hparams(hparams)
   # Modality class is not JSON serializable so remove.
-  del hp_vals["modality"]
-  new_hparams = tf.contrib.training.HParams(**hp_vals)
+  new_hparams.del_hparam("modality")
 
   hparams_fname = os.path.join(output_dir, "hparams.json")
   with tf.gfile.Open(hparams_fname, "w") as f:
diff --git a/tensor2tensor/utils/hparams_lib.py b/tensor2tensor/utils/hparams_lib.py
new file mode 100644
index 000000000..8aee611c9
--- /dev/null
+++ b/tensor2tensor/utils/hparams_lib.py
@@ -0,0 +1,92 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""T2T HParams handling."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import json
+
+from tensor2tensor.data_generators import problem as problem_lib
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+def copy_hparams(hparams):
+  hp_vals = hparams.values()
+  new_hparams = tf.contrib.training.HParams(**hp_vals)
+  other_attrs = ["problem", "problem_hparams"]
+  for attr in other_attrs:
+    attr_val = getattr(hparams, attr, None)
+    if attr_val is not None:
+      setattr(new_hparams, attr, attr_val)
+  return new_hparams
+
+
+def create_hparams(hparams_set,
+                   hparams_overrides_str="",
+                   data_dir=None,
+                   problem_name=None,
+                   hparams_path=None):
+  """Create HParams with data_dir and problem hparams, if kwargs provided."""
+  hparams = registry.hparams(hparams_set)
+  if hparams_path and tf.gfile.Exists(hparams_path):
+    hparams = create_hparams_from_json(hparams_path, hparams)
+  if data_dir:
+    hparams.add_hparam("data_dir", data_dir)
+  if hparams_overrides_str:
+    tf.logging.info("Overriding hparams in %s with %s", hparams_set,
+                    hparams_overrides_str)
+    hparams = hparams.parse(hparams_overrides_str)
+  if problem_name:
+    add_problem_hparams(hparams, problem_name)
+  return hparams
+
+
+def create_hparams_from_json(json_path, hparams=None):
+  """Loading hparams from json; can also start from hparams if specified."""
+  tf.logging.info("Loading hparams from existing json %s" % json_path)
+  with tf.gfile.Open(json_path, "r") as f:
+    hparams_values = json.load(f)
+    new_hparams = tf.contrib.training.HParams(**hparams_values)
+    # Some keys are in new_hparams but not hparams, so we need to be more
+    #   careful than simply using parse_json() from HParams
+    if hparams:  # hparams specified, so update values from json
+      for key in sorted(new_hparams.values().keys()):
+        if hasattr(hparams, key):  # Overlapped keys
+          value = getattr(hparams, key)
+          new_value = getattr(new_hparams, key)
+          if value != new_value:  # Different values
+            tf.logging.info("Overwrite key %s: %s -> %s" % (
+                key, value, new_value))
+            setattr(hparams, key, new_value)
+    else:
+      hparams = new_hparams
+
+  return hparams
+
+
+def add_problem_hparams(hparams, problem_name_or_instance):
+  """Add problem hparams for the problems."""
+  if isinstance(problem_name_or_instance, problem_lib.Problem):
+    problem = problem_name_or_instance
+  else:
+    problem = registry.problem(problem_name_or_instance)
+  p_hparams = problem.get_hparams(hparams)
+  hparams.problem = problem
+  hparams.problem_hparams = p_hparams
diff --git a/tensor2tensor/utils/hparams_lib_test.py b/tensor2tensor/utils/hparams_lib_test.py
new file mode 100644
index 000000000..e0c874ea4
--- /dev/null
+++ b/tensor2tensor/utils/hparams_lib_test.py
@@ -0,0 +1,44 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 trainer_lib."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from tensor2tensor.utils import hparams_lib
+
+import tensorflow as tf
+
+
+class HparamsLibTest(tf.test.TestCase):
+
+  def testCreateHparamsFromJson(self):
+    # Get json_path
+    pkg, _ = os.path.split(__file__)
+    pkg, _ = os.path.split(pkg)
+    json_path = os.path.join(
+        pkg, "test_data", "transformer_test_ckpt", "hparams.json")
+
+    # Create hparams
+    hparams = hparams_lib.create_hparams_from_json(json_path)
+    self.assertEqual(75, len(hparams.values()))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/utils/mtf_model.py b/tensor2tensor/utils/mtf_model.py
index a217e3b76..4a6d88f46 100644
--- a/tensor2tensor/utils/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -20,10 +20,10 @@
 from __future__ import division
 from __future__ import print_function
 
-import copy
 import mesh_tensorflow as mtf
 
 import six
+from tensor2tensor.utils import hparams_lib
 from tensor2tensor.utils import learning_rate
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import t2t_model
@@ -46,7 +46,7 @@ def estimator_model_fn(cls,
                          params=None,
                          decode_hparams=None,
                          use_tpu=False):
-    hparams = copy.deepcopy(hparams)
+    hparams = hparams_lib.copy_hparams(hparams)
     hparams.use_tpu = use_tpu
     # merge decode_hparams into hparams if present
     if mode == tf.estimator.ModeKeys.PREDICT and decode_hparams is not None:
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 5443f0c4d..a4a7c8b16 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -20,7 +20,6 @@
 
 import collections
 import contextlib
-import copy
 import functools
 import math
 import os
@@ -34,6 +33,7 @@
 from tensor2tensor.utils import beam_search
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import expert_utils as eu
+from tensor2tensor.utils import hparams_lib
 from tensor2tensor.utils import learning_rate
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import mlperf_log
@@ -177,7 +177,7 @@ def __init__(self,
     self._problem_hparams = problem_hparams
 
     # Setup hparams
-    hparams = copy.copy(hparams)
+    hparams = hparams_lib.copy_hparams(hparams)
     if self._problem_hparams and hparams.shared_embedding_and_softmax_weights:
       # If vocabularies differ, unset shared_embedding_and_softmax_weights.
       input_modality = self._problem_hparams.modality.get("inputs")
@@ -206,8 +206,8 @@ def __init__(self,
     self._original_hparams = hparams
     self.set_mode(mode)
 
-    self._decode_hparams = copy.copy(decode_hparams or
-                                     decoding.decode_hparams())
+    self._decode_hparams = hparams_lib.copy_hparams(
+        decode_hparams or decoding.decode_hparams())
     self._data_parallelism = data_parallelism or eu.Parallelism([""])
     self._num_datashards = self._data_parallelism.n
     self._ps_devices = self._data_parallelism.ps_devices
@@ -693,7 +693,7 @@ def optimize(self, loss, num_async_replicas=1, use_tpu=False):
   def set_mode(self, mode):
     """Set hparams with the given mode."""
     log_info("Setting T2TModel mode to '%s'", mode)
-    hparams = copy.copy(self._original_hparams)
+    hparams = hparams_lib.copy_hparams(self._original_hparams)
     hparams.add_hparam("mode", mode)
     # When not in training mode, set all forms of dropout to zero.
     if mode != tf.estimator.ModeKeys.TRAIN:
@@ -1396,7 +1396,7 @@ def estimator_model_fn(cls,
     """
     if mode == tf.estimator.ModeKeys.TRAIN:
       create_dummy_vars()
-    hparams = copy.deepcopy(hparams)
+    hparams = hparams_lib.copy_hparams(hparams)
 
     # Instantiate model
     data_parallelism = None
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index 9cb71103e..5692de322 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -47,12 +47,11 @@ def testLossSingleWeights(self):
         sequence_size = 16
         vocab_size = 3
 
-        model_hparams = tf.contrib.training.HParams()
-        model_hparams.label_smoothing = 0.0
-        model_hparams.shared_embedding_and_softmax_weights = False
+        model_hparams = tf.contrib.training.HParams(
+            label_smoothing=0.0,
+            shared_embedding_and_softmax_weights=False)
 
-        problem_hparams = tf.contrib.training.HParams()
-        problem_hparams.loss_multiplier = 1.0
+        problem_hparams = tf.contrib.training.HParams(loss_multiplier=1.0)
         problem_hparams.modality = {}
 
         model = t2t_model.T2TModel(
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index a7d78bd3b..ff56cf8a6 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -25,9 +25,9 @@
 import random
 import numpy as np
 
-from tensor2tensor.data_generators.problem import Problem
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import devices
+from tensor2tensor.utils import hparams_lib
 from tensor2tensor.utils import metrics_hook
 from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import registry
@@ -39,6 +39,10 @@
 from tensorflow.python import debug
 
 
+create_hparams = hparams_lib.create_hparams
+add_problem_hparams = hparams_lib.add_problem_hparams
+
+
 def next_checkpoint(model_dir, timeout_mins=240):
   """Yields successive checkpoints from model_dir.
 
@@ -132,49 +136,6 @@ def create_session_config(log_device_placement=False,
   return config
 
 
-def create_hparams(hparams_set,
-                   hparams_overrides_str="",
-                   data_dir=None,
-                   problem_name=None,
-                   hparams_path=None):
-  """Create HParams with data_dir and problem hparams, if kwargs provided."""
-  hparams = registry.hparams(hparams_set)
-  if hparams_path and tf.gfile.Exists(hparams_path):
-    hparams = _create_hparams_from_json(hparams_path, hparams)
-  if data_dir:
-    hparams.add_hparam("data_dir", data_dir)
-  if hparams_overrides_str:
-    tf.logging.info("Overriding hparams in %s with %s", hparams_set,
-                    hparams_overrides_str)
-    hparams = hparams.parse(hparams_overrides_str)
-  if problem_name:
-    add_problem_hparams(hparams, problem_name)
-  return hparams
-
-
-def _create_hparams_from_json(json_path, hparams=None):
-  """Loading hparams from json; can also start from hparams if specified."""
-  tf.logging.info("Loading hparams from existing json %s" % json_path)
-  with tf.gfile.Open(json_path, "r") as f:
-    hparams_values = json.load(f)
-    new_hparams = tf.contrib.training.HParams(**hparams_values)
-    # Some keys are in new_hparams but not hparams, so we need to be more
-    #   careful than simply using parse_json() from HParams
-    if hparams:  # hparams specified, so update values from json
-      for key in sorted(new_hparams.values().keys()):
-        if hasattr(hparams, key):  # Overlapped keys
-          value = getattr(hparams, key)
-          new_value = getattr(new_hparams, key)
-          if value != new_value:  # Different values
-            tf.logging.info("Overwrite key %s: %s -> %s" % (
-                key, value, new_value))
-            setattr(hparams, key, new_value)
-    else:
-      hparams = new_hparams
-
-  return hparams
-
-
 def is_cloud_async_distributed():
   return ("chief" in
           json.loads(os.environ.get("TF_CONFIG", "{}")).get("cluster", {}))
@@ -808,17 +769,6 @@ def experiment_fn(run_config, hparams):
   return experiment_fn
 
 
-def add_problem_hparams(hparams, problem_name_or_instance):
-  """Add problem hparams for the problems."""
-  if isinstance(problem_name_or_instance, Problem):
-    problem = problem_name_or_instance
-  else:
-    problem = registry.problem(problem_name_or_instance)
-  p_hparams = problem.get_hparams(hparams)
-  hparams.problem = problem
-  hparams.problem_hparams = p_hparams
-
-
 def set_random_seed(seed):
   tf.set_random_seed(seed)
   random.seed(seed)
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index 194b771fa..c37144ca8 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -157,17 +157,6 @@ def testCreateHparams(self):
     base_hparams = trainer_lib.create_hparams("transformer_big")
     self.assertEqual(len(base_hparams.values()), len(hparams.values()))
 
-  def testCreateHparamsFromJson(self):
-    # Get json_path
-    pkg, _ = os.path.split(__file__)
-    pkg, _ = os.path.split(pkg)
-    json_path = os.path.join(
-        pkg, "test_data", "transformer_test_ckpt", "hparams.json")
-
-    # Create hparams
-    hparams = trainer_lib._create_hparams_from_json(json_path)
-    self.assertEqual(75, len(hparams.values()))
-
 
 if __name__ == "__main__":
   tf.test.main()

From 147145a6046b18643640659dcfcb5c931b0818f7 Mon Sep 17 00:00:00 2001
From: etragas-fathom <43351375+etragas-fathom@users.noreply.github.com>
Date: Wed, 16 Jan 2019 21:48:46 -0500
Subject: [PATCH 1527/2720] Mixed Precision Support for Transformer (#1362)

* Add mixed precision to transformer

* remove unnecessary import change
---
 tensor2tensor/layers/common_attention.py   | 53 ++++++++++++++++++----
 tensor2tensor/layers/common_hparams.py     |  5 ++
 tensor2tensor/layers/modalities.py         |  3 +-
 tensor2tensor/layers/transformer_layers.py | 15 +++---
 tensor2tensor/models/transformer.py        | 17 ++++++-
 tensor2tensor/utils/modality.py            |  2 +
 tensor2tensor/utils/optimize.py            | 23 +++++++++-
 tensor2tensor/utils/quantization.py        | 33 ++++++++++++++
 tensor2tensor/utils/t2t_model.py           |  3 ++
 9 files changed, 133 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index dd192e172..6462dffbb 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -37,6 +37,31 @@
 from tensorflow.python.framework import function
 from tensorflow.python.ops import inplace_ops
 
+def large_compatible_negative(tensor):
+  """
+  This function is necessary because the standard value for epsilon
+  in this module (-1e9) cannot be represented using tf.float16
+  """
+
+  if tensor.dtype == tf.float16:
+    return tf.float16.min
+  return -1e9
+
+def mixed_precision_is_enabled(activation_dtype=None, weight_dtype=None, hparams=None):
+  assert not (hparams and (activation_dtype or weight_dtype)), (
+              "Provide only hparams or activation_dtype and weight_dtype")
+
+  if hparams:
+    activation_dtype = hparams.activation_dtype
+    weight_dtype = hparams.weight_dtype
+
+  return activation_dtype == tf.float16 and weight_dtype == tf.float32
+
+def maybe_upcast(logits, activation_dtype=None, weight_dtype=None, hparams=None):
+  if mixed_precision_is_enabled(activation_dtype, weight_dtype, hparams):
+    return tf.cast(logits, tf.float32)
+  return logits
+
 # Struct containing the sequences ids and order on a batch (are send to the
 # expert to allow them to compute the bias mask)
 BatchInfo = collections.namedtuple("BatchInfo", "coordinates, order")
@@ -445,8 +470,7 @@ def add_timing_signal_1d(x,
   channels = common_layers.shape_list(x)[2]
   signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale,
                                 start_index)
-  return x + signal
-
+  return x + common_layers.cast_like(signal, x)
 
 @expert_utils.add_name_scope()
 def get_layer_timing_signal_learned_1d(channels, layer, num_layers):
@@ -881,10 +905,13 @@ def attention_bias_same_segment(query_segment_id, memory_segment_id):
   Returns:
     a `Tensor` with shape [batch, 1, query_length, memory_length].
   """
-  ret = tf.to_float(
+  
+  ret = (tf.to_float(
       tf.not_equal(
           tf.expand_dims(query_segment_id, 2),
-          tf.expand_dims(memory_segment_id, 1))) * -1e9
+          tf.expand_dims(memory_segment_id, 1))) *
+          large_compatible_negative(memory_segment_id))
+
   return tf.expand_dims(ret, axis=1)
 
 
@@ -898,7 +925,9 @@ def attention_bias_ignore_padding(memory_padding):
   Returns:
     a `Tensor` with shape [batch, 1, 1, memory_length].
   """
-  ret = memory_padding * -1e9
+  
+  ret = memory_padding * large_compatible_negative(memory_padding)
+
   return tf.expand_dims(tf.expand_dims(ret, axis=1), axis=1)
 
 
@@ -1435,7 +1464,9 @@ def dot_product_attention(q,
                           name=None,
                           make_image_summary=True,
                           save_weights_to=None,
-                          dropout_broadcast_dims=None):
+                          dropout_broadcast_dims=None,
+                          activation_dtype=None,
+                          weight_dtype=None):
   """Dot-product attention.
 
   Args:
@@ -1455,6 +1486,9 @@ def dot_product_attention(q,
       a string key created from the variable scope (including name).
     dropout_broadcast_dims: an optional list of integers less than rank of q.
       Specifies in which dimensions to broadcast the dropout decisions.
+    activation_dtype: Used to define function activation dtype when using
+      mixed precision.
+    weight_dtype: The dtype weights are stored in when using mixed precision
 
   Returns:
     Tensor with shape [..., length_q, depth_v].
@@ -1465,7 +1499,10 @@ def dot_product_attention(q,
     if bias is not None:
       bias = common_layers.cast_like(bias, logits)
       logits += bias
+    # If logits are fp16, upcast before softmax
+    logits = maybe_upcast(logits, activation_dtype, weight_dtype)
     weights = tf.nn.softmax(logits, name="attention_weights")
+    weights = common_layers.cast_like(weights, q)
     if save_weights_to is not None:
       save_weights_to[scope.name] = weights
       save_weights_to[scope.name + "/logits"] = logits
@@ -3289,7 +3326,6 @@ def compute_qkv(query_antecedent,
       vars_3d_num_heads=vars_3d_num_heads)
   return q, k, v
 
-
 def multihead_attention(query_antecedent,
                         memory_antecedent,
                         bias,
@@ -3464,7 +3500,8 @@ def multihead_attention(query_antecedent,
       x = dot_product_attention(q, k, v, bias, dropout_rate, image_shapes,
                                 save_weights_to=save_weights_to,
                                 make_image_summary=make_image_summary,
-                                dropout_broadcast_dims=dropout_broadcast_dims)
+                                dropout_broadcast_dims=dropout_broadcast_dims,
+                                activation_dtype=kwargs.get('activation_dtype'))
     elif attention_type == "dot_product_relative":
       x = dot_product_attention_relative(
           q,
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 871b61587..712686426 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -70,6 +70,11 @@ def basic_params1():
       optimizer_adafactor_multiply_by_parameter_scale=True,
       # Number of accumulating steps for multi step optimizers.
       optimizer_multistep_accumulate_steps=None,
+      # Loss scaling used.
+      # Generally only necessary with mixed precision training.
+      # Mixed precision training only supports exponential scaling currently
+      # To disable the scaler, see to 0/False
+      mixed_precision_optimizer_loss_scaler='exponential',
       # Whether to zero gradients that were not computed, so that the
       # appropriate slots are created. Useful for sharing checkpoints between
       # models with different sets of heads.
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 0915d0aaa..58aa54a16 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -111,7 +111,7 @@ def bottom_simple(self, x, name, reuse):
       ret = common_layers.gather(var, x)
       if self._model_hparams.multiply_embedding_mode == "sqrt_depth":
         ret *= self._body_input_depth**0.5
-      ret *= tf.expand_dims(tf.to_float(tf.not_equal(x, 0)), -1)
+      ret *= tf.expand_dims(common_layers.cast_like(tf.not_equal(x, 0), ret), -1)
       return ret
 
   def bottom(self, x):
@@ -148,7 +148,6 @@ def top(self, body_output, _):
     else:
       scope_name = "softmax"
       reuse = False
-
     with tf.variable_scope(scope_name, reuse=reuse):
       body_output_shape = common_layers.shape_list(body_output)
       var = self._get_weights(body_output_shape[-1])
diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 8da58792a..6ee0ad5f2 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -88,8 +88,7 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
         32,
         ishape_static[-1],
         name="target_space_embedding",
-        dtype=tf.bfloat16
-        if hparams.activation_dtype == "bfloat16" else tf.float32)
+        dtype=hparams.activation_dtype)
     emb_target_space = tf.reshape(emb_target_space, [1, 1, -1])
     encoder_input += emb_target_space
   if hparams.pos == "timing":
@@ -102,11 +101,9 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
     encoder_input = common_attention.add_positional_embedding(
         encoder_input, hparams.max_length, "inputs_positional_embedding",
         inputs_position)
-  if hparams.activation_dtype == "bfloat16":
-    encoder_self_attention_bias = tf.cast(encoder_self_attention_bias,
-                                          tf.bfloat16)
-    encoder_decoder_attention_bias = tf.cast(encoder_decoder_attention_bias,
-                                             tf.bfloat16)
+  
+  encoder_self_attention_bias = common_layers.cast_like(encoder_self_attention_bias, encoder_input)
+  encoder_decoder_attention_bias = common_layers.cast_like(encoder_decoder_attention_bias, encoder_input)
   return (encoder_input, encoder_self_attention_bias,
           encoder_decoder_attention_bias)
 
@@ -196,7 +193,9 @@ def transformer_encoder(encoder_input,
               make_image_summary=make_image_summary,
               dropout_broadcast_dims=attention_dropout_broadcast_dims,
               max_length=hparams.get("max_length"),
-              vars_3d=hparams.get("attention_variables_3d"))
+              vars_3d=hparams.get("attention_variables_3d"),
+              activation_dtype=hparams.activation_dtype,
+              weight_dtype=hparams.weight_dtype)
           x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 94c720a3a..a5ff38836 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1313,7 +1313,9 @@ def transformer_decoder(decoder_input,
               dropout_broadcast_dims=attention_dropout_broadcast_dims,
               max_length=hparams.get("max_length"),
               decode_loop_step=decode_loop_step,
-              vars_3d=hparams.get("attention_variables_3d"))
+              vars_3d=hparams.get("attention_variables_3d"),
+              activation_dtype=hparams.activation_dtype,
+              weight_dtype=hparams.weight_dtype)
           x = common_layers.layer_postprocess(x, y, hparams)
         if encoder_output is not None:
           with tf.variable_scope("encdec_attention"):
@@ -1335,7 +1337,9 @@ def transformer_decoder(decoder_input,
                 make_image_summary=make_image_summary,
                 dropout_broadcast_dims=attention_dropout_broadcast_dims,
                 max_length=hparams.get("max_length"),
-                vars_3d=hparams.get("attention_variables_3d"))
+                vars_3d=hparams.get("attention_variables_3d"),
+                activation_dtype=hparams.activation_dtype,
+                weight_dtype=hparams.weight_dtype)
             x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(
@@ -2170,6 +2174,15 @@ def transformer_tpu_bf16_activation():
   hparams.activation_dtype = "bfloat16"
   return hparams
 
+@registry.register_hparams
+def transformer_fairseq_fp16_activation_big():
+  """
+  Hparams intended to mirror those used in https://arxiv.org/pdf/1806.00187.pdf
+  """
+  hparams = transformer_big()
+  hparams.activation_dtype = 'float16'
+  hparams.batch_size = 3584
+  return hparams
 
 @registry.register_hparams
 def transformer_packed_tpu():
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index 6041ab310..efb2663d4 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -20,6 +20,7 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import misc_utils
+from tensor2tensor.layers.common_attention import maybe_upcast
 
 import tensorflow as tf
 
@@ -181,6 +182,7 @@ def loss(self, top_out, targets, weights_fn=None):
     logits = top_out
     if weights_fn is None:
       weights_fn = self.targets_weights_fn
+    logits = maybe_upcast(logits,hparams=self._model_hparams)
     return common_layers.padded_cross_entropy(
         logits,
         targets,
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index dd42168e5..26c215cb8 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -27,7 +27,9 @@
 
 import tensorflow as tf
 
+from tensorflow.contrib.mixed_precision import LossScaleOptimizer
 from tensorflow.python.framework import dtypes
+from tensor2tensor.layers.common_attention import mixed_precision_is_enabled
 
 
 def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
@@ -55,7 +57,6 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
   opt = ConditionalOptimizer(hparams.optimizer, learning_rate, hparams, use_tpu)
   if use_tpu:
     opt = tf.contrib.tpu.CrossShardOptimizer(opt)
-
   opt_summaries = []
   if common_layers.should_generate_summaries():
     tf.summary.scalar("learning_rate", learning_rate)
@@ -152,11 +153,30 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
       self._opt = adafactor.adafactor_optimizer_from_hparams(hparams, lr)
     else:
       self._opt = tf.contrib.layers.OPTIMIZER_CLS_NAMES[optimizer_name](lr)
+    if mixed_precision_is_enabled(hparams=hparams):
+      if not hparams.mixed_precision_optimizer_loss_scaler:
+        tf.logging.warning(("Using mixed precision without a loss scaler will ",
+                           "likely cause numerical errors."))
+      elif hparams.mixed_precision_optimizer_loss_scaler != 'exponential':
+        raise ValueError(("Mixed precision training only supports the ",
+                         "exponential loss scaler"))
+      else:
+        tf.logging.info("Using Exponential Update Loss Scaler")
+        loss_scale_manager = tf.contrib.mixed_precision.ExponentialUpdateLossScaleManager(
+            init_loss_scale=2**15,
+            incr_every_n_steps=2000,
+            decr_every_n_nan_or_inf=2,
+            incr_ratio=2,
+            decr_ratio=0.5)
+        self._opt = LossScaleOptimizer(self._opt, loss_scale_manager)
+
 
     self._zero_grads = hparams.optimizer_zero_grads
 
   def compute_gradients(self, loss, var_list=None, **kwargs):  # pylint: disable=arguments-differ
     gradients = self._opt.compute_gradients(loss, var_list, **kwargs)
+    # print("var list:", var_list)
+    # print("Gradients before cast", gradients)
     def cast_grad(g, v):
       if v is not None and g is not None:
         g = common_layers.cast_like(g, v)
@@ -164,6 +184,7 @@ def cast_grad(g, v):
         g = tf.zeros_like(v)
       return (g, v)
     gradients = [cast_grad(g, v) for g, v in gradients]
+    # print("Gradients after cast", gradients)
     return gradients
 
   def apply_gradients(self, grads_and_vars, global_step=None, name=None):
diff --git a/tensor2tensor/utils/quantization.py b/tensor2tensor/utils/quantization.py
index cb51be809..a1fedd622 100644
--- a/tensor2tensor/utils/quantization.py
+++ b/tensor2tensor/utils/quantization.py
@@ -47,6 +47,39 @@ def bfloat16_activations_var_getter(getter, *args, **kwargs):
     var = tf.cast(var, requested_dtype)
   return var
 
+def float16_activations_var_getter(getter, *args, **kwargs):
+  """A custom getter function for float32 parameters and float16 activations.
+   Args:
+    getter: custom getter
+    *args: arguments
+    **kwargs: keyword arguments
+  Returns:
+    variables with the correct dtype.
+  Raises:
+    KeyError: if "dtype" is not provided as a kwarg.
+  This function ensures the following:
+    1. All variables requested with type fp16 are stored as type fp32
+    2. All variables requested with type fp32 are returned as type fp16
+  See https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/#training_tensorflow
+  for more information on this strategy
+  """
+  requested_dtype = kwargs["dtype"]
+
+  if requested_dtype == tf.float16:
+    kwargs["dtype"] = tf.float32
+  
+  if requested_dtype == tf.float32:
+    requested_dtype = tf.float16
+  var = getter(*args, **kwargs)
+  # This if statement is needed to guard the cast, because batch norm
+  # assigns directly to the return value of this custom getter. The cast
+  # makes the return value not a variable so it cannot be assigned. Batch
+  # norm variables are always in fp32 so this if statement is never
+  # triggered for them.
+  if var.dtype.base_dtype != requested_dtype:
+    var = tf.cast(var, requested_dtype)
+  # print("Output var is {}".format(var))
+  return var
 
 def simulated_quantize(x, num_bits, noise):
   """Simulate quantization to num_bits bits, with externally-stored scale.
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index a4a7c8b16..37bef16b4 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -30,6 +30,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators.problem import problem_hparams_to_features
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers.common_attention import mixed_precision_is_enabled
 from tensor2tensor.utils import beam_search
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import expert_utils as eu
@@ -273,6 +274,8 @@ def _custom_getter(self):
           activation_dtype=activation_dtype)
     elif self.hparams.activation_dtype == "bfloat16":
       return quantization.bfloat16_activations_var_getter
+    elif mixed_precision_is_enabled(hparams=self.hparams):
+      return quantization.float16_activations_var_getter
     else:
       return None
 

From 3448fe1937f6eb1e55f66325481da2f28eb67553 Mon Sep 17 00:00:00 2001
From: etragas-fathom <43351375+etragas-fathom@users.noreply.github.com>
Date: Wed, 16 Jan 2019 21:17:46 -0800
Subject: [PATCH 1528/2720] internal merge of PR #1362

PiperOrigin-RevId: 229679153
---
 tensor2tensor/layers/common_attention.py      | 43 +++++++++++--------
 tensor2tensor/layers/common_hparams.py        |  2 +-
 tensor2tensor/layers/modalities.py            | 12 +++---
 tensor2tensor/layers/transformer_layers.py    | 22 +++++-----
 tensor2tensor/models/transformer.py           |  8 ++--
 .../models/video/basic_stochastic.py          |  5 ++-
 tensor2tensor/utils/modality.py               |  4 +-
 tensor2tensor/utils/optimize.py               | 29 +++++++------
 tensor2tensor/utils/quantization.py           | 21 +++++----
 9 files changed, 82 insertions(+), 64 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 6462dffbb..2c6cc6d5a 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -37,31 +37,42 @@
 from tensorflow.python.framework import function
 from tensorflow.python.ops import inplace_ops
 
-def large_compatible_negative(tensor):
-  """
+
+def large_compatible_negative(tensor_type):
+  """Large negative number as Tensor.
+
   This function is necessary because the standard value for epsilon
   in this module (-1e9) cannot be represented using tf.float16
-  """
 
-  if tensor.dtype == tf.float16:
+  Args:
+    tensor_type: a dtype to determine the type.
+
+  Returns:
+    a large negative number.
+  """
+  if tensor_type == tf.float16:
     return tf.float16.min
   return -1e9
 
-def mixed_precision_is_enabled(activation_dtype=None, weight_dtype=None, hparams=None):
-  assert not (hparams and (activation_dtype or weight_dtype)), (
-              "Provide only hparams or activation_dtype and weight_dtype")
 
-  if hparams:
+def mixed_precision_is_enabled(
+    activation_dtype=None, weight_dtype=None, hparams=None):
+  assert not (hparams and (activation_dtype or weight_dtype)), (
+      "Provide only hparams or activation_dtype and weight_dtype")
+  if (hparams and hasattr(hparams, "activation_dtype") and
+      hasattr(hparams, "weight_dtype")):
     activation_dtype = hparams.activation_dtype
     weight_dtype = hparams.weight_dtype
-
   return activation_dtype == tf.float16 and weight_dtype == tf.float32
 
-def maybe_upcast(logits, activation_dtype=None, weight_dtype=None, hparams=None):
+
+def maybe_upcast(logits,
+                 activation_dtype=None, weight_dtype=None, hparams=None):
   if mixed_precision_is_enabled(activation_dtype, weight_dtype, hparams):
     return tf.cast(logits, tf.float32)
   return logits
 
+
 # Struct containing the sequences ids and order on a batch (are send to the
 # expert to allow them to compute the bias mask)
 BatchInfo = collections.namedtuple("BatchInfo", "coordinates, order")
@@ -472,6 +483,7 @@ def add_timing_signal_1d(x,
                                 start_index)
   return x + common_layers.cast_like(signal, x)
 
+
 @expert_utils.add_name_scope()
 def get_layer_timing_signal_learned_1d(channels, layer, num_layers):
   """get n-dimensional embedding as the layer (vertical) timing signal.
@@ -905,13 +917,11 @@ def attention_bias_same_segment(query_segment_id, memory_segment_id):
   Returns:
     a `Tensor` with shape [batch, 1, query_length, memory_length].
   """
-  
   ret = (tf.to_float(
       tf.not_equal(
           tf.expand_dims(query_segment_id, 2),
           tf.expand_dims(memory_segment_id, 1))) *
-          large_compatible_negative(memory_segment_id))
-
+         large_compatible_negative(memory_segment_id.dtype))
   return tf.expand_dims(ret, axis=1)
 
 
@@ -925,9 +935,7 @@ def attention_bias_ignore_padding(memory_padding):
   Returns:
     a `Tensor` with shape [batch, 1, 1, memory_length].
   """
-  
-  ret = memory_padding * large_compatible_negative(memory_padding)
-
+  ret = memory_padding * large_compatible_negative(memory_padding.dtype)
   return tf.expand_dims(tf.expand_dims(ret, axis=1), axis=1)
 
 
@@ -3326,6 +3334,7 @@ def compute_qkv(query_antecedent,
       vars_3d_num_heads=vars_3d_num_heads)
   return q, k, v
 
+
 def multihead_attention(query_antecedent,
                         memory_antecedent,
                         bias,
@@ -3501,7 +3510,7 @@ def multihead_attention(query_antecedent,
                                 save_weights_to=save_weights_to,
                                 make_image_summary=make_image_summary,
                                 dropout_broadcast_dims=dropout_broadcast_dims,
-                                activation_dtype=kwargs.get('activation_dtype'))
+                                activation_dtype=kwargs.get("activation_dtype"))
     elif attention_type == "dot_product_relative":
       x = dot_product_attention_relative(
           q,
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 712686426..2bbfb7618 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -74,7 +74,7 @@ def basic_params1():
       # Generally only necessary with mixed precision training.
       # Mixed precision training only supports exponential scaling currently
       # To disable the scaler, see to 0/False
-      mixed_precision_optimizer_loss_scaler='exponential',
+      mixed_precision_optimizer_loss_scaler="exponential",
       # Whether to zero gradients that were not computed, so that the
       # appropriate slots are created. Useful for sharing checkpoints between
       # models with different sets of heads.
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 58aa54a16..49de33f13 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -111,7 +111,8 @@ def bottom_simple(self, x, name, reuse):
       ret = common_layers.gather(var, x)
       if self._model_hparams.multiply_embedding_mode == "sqrt_depth":
         ret *= self._body_input_depth**0.5
-      ret *= tf.expand_dims(common_layers.cast_like(tf.not_equal(x, 0), ret), -1)
+      ret *= tf.expand_dims(
+          common_layers.cast_like(tf.not_equal(x, 0), ret), -1)
       return ret
 
   def bottom(self, x):
@@ -789,12 +790,11 @@ def name(self):
 
   def bottom(self, x):
     with tf.variable_scope(self.name):
+      multiplier = 1.0
+      if self._model_hparams.multiply_embedding_mode == "sqrt_depth":
+        multiplier = self._body_input_depth**0.5
       return common_layers.embedding(
-          x,
-          self._vocab_size,
-          self._body_input_depth,
-          multiplier=self._body_input_depth**0.5 if
-          self._model_hparams.multiply_embedding_mode == "sqrt_depth" else 1.0)
+          x, self._vocab_size, self._body_input_depth, multiplier=multiplier)
 
   def targets_bottom(self, x):
     with tf.variable_scope(self.name):
diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 6ee0ad5f2..82b7cf19f 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -101,9 +101,11 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
     encoder_input = common_attention.add_positional_embedding(
         encoder_input, hparams.max_length, "inputs_positional_embedding",
         inputs_position)
-  
-  encoder_self_attention_bias = common_layers.cast_like(encoder_self_attention_bias, encoder_input)
-  encoder_decoder_attention_bias = common_layers.cast_like(encoder_decoder_attention_bias, encoder_input)
+
+  encoder_self_attention_bias = common_layers.cast_like(
+      encoder_self_attention_bias, encoder_input)
+  encoder_decoder_attention_bias = common_layers.cast_like(
+      encoder_decoder_attention_bias, encoder_input)
   return (encoder_input, encoder_self_attention_bias,
           encoder_decoder_attention_bias)
 
@@ -323,10 +325,9 @@ def transformer_ffn_layer(x,
   elif ffn_layer == "sru":
     return common_layers.sru(x)
   elif ffn_layer == "local_moe_tpu":
-    overhead = (
-        hparams.moe_overhead_train
-        if hparams.mode == tf.estimator.ModeKeys.TRAIN else
-        hparams.moe_overhead_eval)
+    overhead = hparams.moe_overhead_eval
+    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+      overhead = hparams.moe_overhead_train
     ret, loss = expert_utils.local_moe_tpu(
         x,
         hparams.filter_size // 2,
@@ -335,10 +336,9 @@ def transformer_ffn_layer(x,
         overhead=overhead,
         loss_coef=hparams.moe_loss_coef)
   elif ffn_layer == "local_moe":
-    overhead = (
-        hparams.moe_overhead_train
-        if hparams.mode == tf.estimator.ModeKeys.TRAIN else
-        hparams.moe_overhead_eval)
+    overhead = hparams.moe_overhead_eval
+    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+      overhead = hparams.moe_overhead_train
     ret, loss = expert_utils.local_moe(
         x,
         True,
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index a5ff38836..c8534e751 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -2174,16 +2174,16 @@ def transformer_tpu_bf16_activation():
   hparams.activation_dtype = "bfloat16"
   return hparams
 
+
 @registry.register_hparams
 def transformer_fairseq_fp16_activation_big():
-  """
-  Hparams intended to mirror those used in https://arxiv.org/pdf/1806.00187.pdf
-  """
+  """Hparams intended to mirror those used in arxiv.org/pdf/1806.00187.pdf."""
   hparams = transformer_big()
-  hparams.activation_dtype = 'float16'
+  hparams.activation_dtype = "float16"
   hparams.batch_size = 3584
   return hparams
 
+
 @registry.register_hparams
 def transformer_packed_tpu():
   """Deprecated alias for transformer_tpu()."""
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 5d3d7b7ed..62f158e11 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -33,6 +33,7 @@
 import tensorflow as tf
 
 tfl = tf.layers
+_MAX_BATCH = 128
 
 
 @registry.register_model
@@ -74,7 +75,7 @@ def init_internal_states(self):
     if not self.hparams.concat_internal_states:
       return None
     # Hardcoded frame shapes.
-    max_batch_size = max(64, self.hparams.batch_size)
+    max_batch_size = max(_MAX_BATCH, self.hparams.batch_size)
     shape = [max_batch_size] + self.hparams.problem.frame_shape[:-1] + [
         self.hparams.recurrent_state_size]
     with tf.variable_scope("clean_scope_for_internal_state"):
@@ -114,7 +115,7 @@ def update_internal_states_early(self, internal_states, frames):
     state_candidate = tf.tanh(state_candidate)
     internal_state = internal_state * state_gate
     internal_state += state_candidate * (1.0 - state_gate)
-    max_batch_size = max(64, self.hparams.batch_size)
+    max_batch_size = max(_MAX_BATCH, self.hparams.batch_size)
     diff_batch_size = max_batch_size - batch_size
     internal_state = tf.pad(
         internal_state, [[0, diff_batch_size], [0, 0], [0, 0], [0, 0]])
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index efb2663d4..1bc5bb749 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -18,9 +18,9 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import misc_utils
-from tensor2tensor.layers.common_attention import maybe_upcast
 
 import tensorflow as tf
 
@@ -182,7 +182,7 @@ def loss(self, top_out, targets, weights_fn=None):
     logits = top_out
     if weights_fn is None:
       weights_fn = self.targets_weights_fn
-    logits = maybe_upcast(logits,hparams=self._model_hparams)
+    logits = common_attention.maybe_upcast(logits, hparams=self._model_hparams)
     return common_layers.padded_cross_entropy(
         logits,
         targets,
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 26c215cb8..48ae5d349 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -27,9 +27,16 @@
 
 import tensorflow as tf
 
+
 from tensorflow.contrib.mixed_precision import LossScaleOptimizer
 from tensorflow.python.framework import dtypes
-from tensor2tensor.layers.common_attention import mixed_precision_is_enabled
+
+
+def _mixed_precision_is_enabled(hparams):
+  """Should be the same as in common_attention, avoiding import."""
+  activation_dtype = hparams.activation_dtype
+  weight_dtype = hparams.weight_dtype
+  return activation_dtype == tf.float16 and weight_dtype == tf.float32
 
 
 def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
@@ -153,30 +160,27 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
       self._opt = adafactor.adafactor_optimizer_from_hparams(hparams, lr)
     else:
       self._opt = tf.contrib.layers.OPTIMIZER_CLS_NAMES[optimizer_name](lr)
-    if mixed_precision_is_enabled(hparams=hparams):
+    if _mixed_precision_is_enabled(hparams):
       if not hparams.mixed_precision_optimizer_loss_scaler:
-        tf.logging.warning(("Using mixed precision without a loss scaler will ",
-                           "likely cause numerical errors."))
-      elif hparams.mixed_precision_optimizer_loss_scaler != 'exponential':
-        raise ValueError(("Mixed precision training only supports the ",
-                         "exponential loss scaler"))
+        tf.logging.warning("Using mixed precision without a loss scaler will "
+                           "likely cause numerical errors.")
+      elif hparams.mixed_precision_optimizer_loss_scaler != "exponential":
+        raise ValueError("Mixed precision training only supports the "
+                         "exponential loss scaler")
       else:
         tf.logging.info("Using Exponential Update Loss Scaler")
-        loss_scale_manager = tf.contrib.mixed_precision.ExponentialUpdateLossScaleManager(
+        manager = tf.contrib.mixed_precision.ExponentialUpdateLossScaleManager(
             init_loss_scale=2**15,
             incr_every_n_steps=2000,
             decr_every_n_nan_or_inf=2,
             incr_ratio=2,
             decr_ratio=0.5)
-        self._opt = LossScaleOptimizer(self._opt, loss_scale_manager)
-
+        self._opt = LossScaleOptimizer(self._opt, manager)
 
     self._zero_grads = hparams.optimizer_zero_grads
 
   def compute_gradients(self, loss, var_list=None, **kwargs):  # pylint: disable=arguments-differ
     gradients = self._opt.compute_gradients(loss, var_list, **kwargs)
-    # print("var list:", var_list)
-    # print("Gradients before cast", gradients)
     def cast_grad(g, v):
       if v is not None and g is not None:
         g = common_layers.cast_like(g, v)
@@ -184,7 +188,6 @@ def cast_grad(g, v):
         g = tf.zeros_like(v)
       return (g, v)
     gradients = [cast_grad(g, v) for g, v in gradients]
-    # print("Gradients after cast", gradients)
     return gradients
 
   def apply_gradients(self, grads_and_vars, global_step=None, name=None):
diff --git a/tensor2tensor/utils/quantization.py b/tensor2tensor/utils/quantization.py
index a1fedd622..4ae86309f 100644
--- a/tensor2tensor/utils/quantization.py
+++ b/tensor2tensor/utils/quantization.py
@@ -47,27 +47,32 @@ def bfloat16_activations_var_getter(getter, *args, **kwargs):
     var = tf.cast(var, requested_dtype)
   return var
 
+
 def float16_activations_var_getter(getter, *args, **kwargs):
   """A custom getter function for float32 parameters and float16 activations.
-   Args:
+
+  This function ensures the following:
+    1. All variables requested with type fp16 are stored as type fp32.
+    2. All variables requested with type fp32 are returned as type fp16.
+  See https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/
+  #training_tensorflow for more information on this strategy.
+
+  Args:
     getter: custom getter
     *args: arguments
     **kwargs: keyword arguments
+
   Returns:
     variables with the correct dtype.
+
   Raises:
     KeyError: if "dtype" is not provided as a kwarg.
-  This function ensures the following:
-    1. All variables requested with type fp16 are stored as type fp32
-    2. All variables requested with type fp32 are returned as type fp16
-  See https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/#training_tensorflow
-  for more information on this strategy
   """
   requested_dtype = kwargs["dtype"]
 
   if requested_dtype == tf.float16:
     kwargs["dtype"] = tf.float32
-  
+
   if requested_dtype == tf.float32:
     requested_dtype = tf.float16
   var = getter(*args, **kwargs)
@@ -78,9 +83,9 @@ def float16_activations_var_getter(getter, *args, **kwargs):
   # triggered for them.
   if var.dtype.base_dtype != requested_dtype:
     var = tf.cast(var, requested_dtype)
-  # print("Output var is {}".format(var))
   return var
 
+
 def simulated_quantize(x, num_bits, noise):
   """Simulate quantization to num_bits bits, with externally-stored scale.
 

From f14cf8d3d571b9b811f29b3c9bf5d72b87619880 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 16 Jan 2019 22:01:04 -0800
Subject: [PATCH 1529/2720] Add option in T2TGymEnv to not
 derive/pre-process/encode observations.

PiperOrigin-RevId: 229682563
---
 tensor2tensor/data_generators/gym_env.py      | 63 +++++++++++++------
 tensor2tensor/data_generators/gym_env_test.py | 18 ++++++
 2 files changed, 63 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 6d6d27725..ad24c691c 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -582,11 +582,15 @@ class T2TGymEnv(T2TEnv):
   def __init__(self, base_env_name=None, batch_size=1, grayscale=False,
                resize_height_factor=2, resize_width_factor=2,
                rl_env_max_episode_steps=-1, max_num_noops=0,
-               maxskip_envs=False, **kwargs):
+               maxskip_envs=False,
+               should_derive_observation_space=True,
+               **kwargs):
     if base_env_name is None:
       base_env_name = self.base_env_name
     self._base_env_name = base_env_name
     super(T2TGymEnv, self).__init__(batch_size, **kwargs)
+    # TODO(afrozm): Find a proper way of doing this. Refactor or cleanup.
+    self.should_derive_observation_space = should_derive_observation_space
     self.grayscale = grayscale
     self.resize_height_factor = resize_height_factor
     self.resize_width_factor = resize_width_factor
@@ -612,25 +616,38 @@ def __init__(self, base_env_name=None, batch_size=1, grayscale=False,
                for env in self._envs):
       raise ValueError("All environments must use the same observation space.")
 
-    self.observation_space = self._derive_observation_space(orig_observ_space)
+    self.observation_space = orig_observ_space
+    if self.should_derive_observation_space:
+      self.observation_space = self._derive_observation_space(orig_observ_space)
 
     self.action_space = self._envs[0].action_space
     if not all(env.action_space == self.action_space for env in self._envs):
       raise ValueError("All environments must use the same action space.")
 
-    with self._tf_graph.obj.as_default():
-      self._resize = dict()
-      orig_height, orig_width = orig_observ_space.shape[:2]
-      self._img_batch_t = _Noncopyable(tf.placeholder(
-          dtype=tf.uint8, shape=(None, orig_height, orig_width, 3)))
-      height, width = self.observation_space.shape[:2]
-      resized = tf.image.resize_images(self._img_batch_t.obj,
-                                       [height, width],
-                                       tf.image.ResizeMethod.AREA)
-      resized = tf.cast(resized, tf.as_dtype(self.observation_space.dtype))
-      if self.grayscale:
-        resized = tf.image.rgb_to_grayscale(resized)
-      self._resized_img_batch_t = _Noncopyable(resized)
+    if self.should_derive_observation_space:
+      with self._tf_graph.obj.as_default():
+        self._resize = dict()
+        orig_height, orig_width = orig_observ_space.shape[:2]
+        self._img_batch_t = _Noncopyable(tf.placeholder(
+            dtype=tf.uint8, shape=(None, orig_height, orig_width, 3)))
+        height, width = self.observation_space.shape[:2]
+        resized = tf.image.resize_images(self._img_batch_t.obj,
+                                         [height, width],
+                                         tf.image.ResizeMethod.AREA)
+        resized = tf.cast(resized, tf.as_dtype(self.observation_space.dtype))
+        if self.grayscale:
+          resized = tf.image.rgb_to_grayscale(resized)
+        self._resized_img_batch_t = _Noncopyable(resized)
+
+  # TODO(afrozm): Find a place for this. Till then use self._envs[0]'s hparams.
+  def hparams(self, defaults, unused_model_hparams):
+    if hasattr(self._envs[0], "hparams"):
+      tf.logging.info("Retuning the env's hparams from T2TGymEnv.")
+      return self._envs[0].hparams(defaults, unused_model_hparams)
+
+    # Otherwise just call the super-class' hparams.
+    tf.logging.info("Retuning the T2TGymEnv's superclass' hparams.")
+    super(T2TGymEnv, self).hparams(defaults, unused_model_hparams)
 
   @property
   def base_env_name(self):
@@ -640,6 +657,7 @@ def base_env_name(self):
   def num_channels(self):
     return self.observation_space.shape[2]
 
+  # TODO(afrozm): Why is this separated out from _preprocess_observations?
   def _derive_observation_space(self, orig_observ_space):
     height, width, channels = orig_observ_space.shape
     if self.grayscale:
@@ -654,9 +672,18 @@ def _derive_observation_space(self, orig_observ_space):
   def __str__(self):
     return "T2TGymEnv(%s)" % ", ".join([str(env) for env in self._envs])
 
-  def _preprocess_observations(self, obs):
-    return self._session.obj.run(self._resized_img_batch_t.obj,
-                                 feed_dict={self._img_batch_t.obj: obs})
+  def _encode_observations(self, observations):
+    if not self.should_derive_observation_space:
+      return observations
+    return super(T2TGymEnv, self)._encode_observations(observations)
+
+  def _preprocess_observations(self, observations):
+    # TODO(afrozm): Clean this up.
+    if not self.should_derive_observation_space:
+      return observations
+    return self._session.obj.run(
+        self._resized_img_batch_t.obj,
+        feed_dict={self._img_batch_t.obj: observations})
 
   def _step(self, actions):
     (obs, rewards, dones, _) = zip(*[
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 2d5859b1b..c04a5fab4 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -237,6 +237,24 @@ def test_resize(self):
       self.assertEqual(height, orig_height // resize_height_factor)
       self.assertEqual(width, orig_width // resize_width_factor)
 
+  def test_no_resize_option(self):
+    env_name = TEST_ENV_NAME
+    orig_env = make_gym_env(env_name)
+    resize_height_factor = 2
+    resize_width_factor = 3
+    orig_height, orig_width = orig_env.observation_space.shape[:2]
+    env, obs, _, _ = self.init_batch_and_play(
+        env_name, steps_per_epoch=1,
+        resize_height_factor=resize_height_factor,
+        resize_width_factor=resize_width_factor,
+        should_derive_observation_space=False)
+    for obs_batch in obs:
+      ob = obs_batch[0]
+      self.assertEqual(ob.shape, env.observation_space.shape)
+      height, width = ob.shape[:2]
+      self.assertEqual(height, orig_height)
+      self.assertEqual(width, orig_width)
+
   def assert_channels(self, env, obs, n_channels):
     self.assertEqual(env.observation_space.shape[2], n_channels)
     self.assertEqual(env.num_channels, n_channels)

From 4152ddc0516d134e78f67d1772d7137994d8426c Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 16 Jan 2019 22:18:38 -0800
Subject: [PATCH 1530/2720] Guard some hparams in case they don't exist.

PiperOrigin-RevId: 229684407
---
 tensor2tensor/layers/transformer_layers.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 82b7cf19f..d52d7cb71 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -88,7 +88,7 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
         32,
         ishape_static[-1],
         name="target_space_embedding",
-        dtype=hparams.activation_dtype)
+        dtype=hparams.get("activation_dtype", "float32"))
     emb_target_space = tf.reshape(emb_target_space, [1, 1, -1])
     encoder_input += emb_target_space
   if hparams.pos == "timing":
@@ -196,8 +196,8 @@ def transformer_encoder(encoder_input,
               dropout_broadcast_dims=attention_dropout_broadcast_dims,
               max_length=hparams.get("max_length"),
               vars_3d=hparams.get("attention_variables_3d"),
-              activation_dtype=hparams.activation_dtype,
-              weight_dtype=hparams.weight_dtype)
+              activation_dtype=hparams.get("activation_dtype", "float32"),
+              weight_dtype=hparams.get("weight_dtype", "float32"))
           x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(

From 9448cb010110490c711ac370bf41786d16bbd4cb Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 17 Jan 2019 13:26:45 -0800
Subject: [PATCH 1531/2720] More guards against hparams that may not exist.

PiperOrigin-RevId: 229806871
---
 tensor2tensor/models/transformer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index c8534e751..cc85ab6ec 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1338,8 +1338,8 @@ def transformer_decoder(decoder_input,
                 dropout_broadcast_dims=attention_dropout_broadcast_dims,
                 max_length=hparams.get("max_length"),
                 vars_3d=hparams.get("attention_variables_3d"),
-                activation_dtype=hparams.activation_dtype,
-                weight_dtype=hparams.weight_dtype)
+                activation_dtype=hparams.get("activation_dtype", "float32"),
+                weight_dtype=hparams.get("weight_dtype", "float32"))
             x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(

From 0d9346b2a6df52e0c5ff61f638756e2432e32611 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 18 Jan 2019 01:05:32 +0100
Subject: [PATCH 1532/2720] Allow running planner on real env (#1380)

---
 tensor2tensor/data_generators/gym_env.py | 110 ++++++++++++++++-------
 tensor2tensor/rl/evaluator.py            |  25 ++++--
 tensor2tensor/rl/rl_utils.py             |  62 ++++++++++---
 3 files changed, 152 insertions(+), 45 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index ad24c691c..43d2eca10 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -137,17 +137,22 @@ class T2TEnv(EnvSimulationProblem):
 
   Args:
     batch_size: Number of environments in a batch.
+    store_rollouts: Whether to store collected rollouts in memory and later on
+      disk. Defaults to True.
   """
 
   observation_space = None
   name = None
 
   def __init__(self, batch_size, *args, **kwargs):
+    self._store_rollouts = kwargs.pop("store_rollouts", True)
+
     super(T2TEnv, self).__init__(*args, **kwargs)
 
     self.batch_size = batch_size
     self._rollouts_by_epoch_and_split = collections.OrderedDict()
     self.current_epoch = None
+    self._should_preprocess_on_reset = True
     with tf.Graph().as_default() as tf_graph:
       self._tf_graph = _Noncopyable(tf_graph)
       self._decoded_image_p = _Noncopyable(
@@ -270,7 +275,8 @@ def step(self, actions):
     Raises:
       ValueError: when the data for current epoch has already been loaded.
     """
-    if self._rollouts_by_epoch_and_split[self.current_epoch]:
+    if self._store_rollouts and \
+        self._rollouts_by_epoch_and_split[self.current_epoch]:
       raise ValueError(
           "Data for current epoch has already been loaded from disk."
       )
@@ -278,18 +284,19 @@ def step(self, actions):
     obs = self._preprocess_observations(obs)
     (min_reward, max_reward) = self.reward_range
     rewards = np.around(np.clip(unclipped_rewards, min_reward, max_reward))
-    unclipped_rewards = unclipped_rewards.astype(np.float64)
-    encoded_obs = self._encode_observations(obs)
-    for (rollout, frame, action) in zip(
-        self._current_batch_rollouts, self._current_batch_frames, actions
-    ):
-      rollout.append(frame._replace(action=action))
-
-    # orud = (observation, reward, unclipped_reward, done)
-    self._current_batch_frames = [
-        Frame(*orud, action=None)
-        for orud in zip(encoded_obs, rewards, unclipped_rewards, dones)
-    ]
+    if self._store_rollouts:
+      unclipped_rewards = unclipped_rewards.astype(np.float64)
+      encoded_obs = self._encode_observations(obs)
+      for (rollout, frame, action) in zip(
+          self._current_batch_rollouts, self._current_batch_frames, actions
+      ):
+        rollout.append(frame._replace(action=action))
+
+      # orud = (observation, reward, unclipped_reward, done)
+      self._current_batch_frames = [
+          Frame(*orud, action=None)
+          for orud in zip(encoded_obs, rewards, unclipped_rewards, dones)
+      ]
     return (obs, rewards, dones)
 
   def _reset(self, indices):
@@ -317,7 +324,7 @@ def reset(self, indices=None):
     Raises:
       ValueError: when there's no current epoch.
     """
-    if self.current_epoch is None:
+    if self._store_rollouts and self.current_epoch is None:
       raise ValueError(
           "No current epoch. start_new_epoch() should first be called."
       )
@@ -325,18 +332,21 @@ def reset(self, indices=None):
     if indices is None:
       indices = np.arange(self.batch_size)
     new_obs = self._reset(indices)
-    new_obs = self._preprocess_observations(new_obs)
-    encoded_obs = self._encode_observations(new_obs)
-    for (index, ob) in zip(indices, encoded_obs):
-      frame = self._current_batch_frames[index]
-      if frame is not None:
-        rollout = self._current_batch_rollouts[index]
-        rollout.append(frame._replace(action=0))
-        self._current_epoch_rollouts.append(rollout)
-        self._current_batch_rollouts[index] = []
-      self._current_batch_frames[index] = Frame(
-          observation=ob, reward=0, unclipped_reward=0, done=False, action=None
-      )
+    if self._should_preprocess_on_reset:
+      new_obs = self._preprocess_observations(new_obs)
+    if self._store_rollouts:
+      encoded_obs = self._encode_observations(new_obs)
+      for (index, ob) in zip(indices, encoded_obs):
+        frame = self._current_batch_frames[index]
+        if frame is not None:
+          rollout = self._current_batch_rollouts[index]
+          rollout.append(frame._replace(action=0))
+          self._current_epoch_rollouts.append(rollout)
+          self._current_batch_rollouts[index] = []
+        self._current_batch_frames[index] = Frame(
+            observation=ob, reward=0, unclipped_reward=0, done=False,
+            action=None
+        )
     return new_obs
 
   def close(self):
@@ -594,6 +604,10 @@ def __init__(self, base_env_name=None, batch_size=1, grayscale=False,
     self.grayscale = grayscale
     self.resize_height_factor = resize_height_factor
     self.resize_width_factor = resize_width_factor
+    self.rl_env_max_episode_steps = rl_env_max_episode_steps
+    self.maxskip_envs = maxskip_envs
+    self._initial_state = None
+    self._initial_frames = None
     if not self.name:
       # Set problem name if not registered.
       self.name = "Gym%s" % base_env_name
@@ -649,6 +663,20 @@ def hparams(self, defaults, unused_model_hparams):
     tf.logging.info("Retuning the T2TGymEnv's superclass' hparams.")
     super(T2TGymEnv, self).hparams(defaults, unused_model_hparams)
 
+  def new_like(self, **kwargs):
+    env_kwargs = {
+        "base_env_name": self.base_env_name,
+        "batch_size": self.batch_size,
+        "grayscale": self.grayscale,
+        "resize_height_factor": self.resize_height_factor,
+        "resize_width_factor": self.resize_width_factor,
+        "rl_env_max_episode_steps": self.rl_env_max_episode_steps,
+        "max_num_noops": self.max_num_noops,
+        "maxskip_envs": self.maxskip_envs,
+    }
+    env_kwargs.update(kwargs)
+    return T2TGymEnv(**env_kwargs)
+
   @property
   def base_env_name(self):
     return self._base_env_name
@@ -685,6 +713,17 @@ def _preprocess_observations(self, observations):
         self._resized_img_batch_t.obj,
         feed_dict={self._img_batch_t.obj: observations})
 
+  @property
+  def state(self):
+    """Gets the current state."""
+    return [env.unwrapped.clone_full_state() for env in self._envs]
+
+  def set_initial_state(self, initial_state, initial_frames):
+    """Sets the state that will be used on next reset."""
+    self._initial_state = initial_state
+    self._initial_frames = initial_frames[:, -1, ...]
+    self._should_preprocess_on_reset = False
+
   def _step(self, actions):
     (obs, rewards, dones, _) = zip(*[
         env.step(action) for (env, action) in zip(self._envs, actions)
@@ -692,9 +731,18 @@ def _step(self, actions):
     return tuple(map(np.stack, (obs, rewards, dones)))
 
   def _reset(self, indices):
-    def reset_with_noops(env):
-      """Resets environment and applies random number of NOOP actions on it."""
+    def reset_with_initial_state(env, index):
+      """Resets environment taking self._initial_state into account."""
       obs = env.reset()
+      if self._initial_state is None:
+        return obs
+      else:
+        env.unwrapped.restore_full_state(self._initial_state[index])
+        return self._initial_frames[index, ...]
+
+    def reset_with_noops(env, index):
+      """Resets environment and applies random number of NOOP actions on it."""
+      obs = reset_with_initial_state(env, index)
       try:
         num_noops = random.randint(1, self.max_num_noops)
       except ValueError:
@@ -703,11 +751,13 @@ def reset_with_noops(env):
       for _ in range(num_noops):
         (obs, _, done, _) = env.step(self.noop_action)
         if done:
-          obs = env.reset()
+          obs = reset_with_initial_state(env, index)
 
       return obs
 
-    return np.stack([reset_with_noops(self._envs[index]) for index in indices])
+    return np.stack([
+        reset_with_noops(self._envs[index], index) for index in indices
+    ])
 
   def close(self):
     for env in self._envs:
diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index f5afd1e22..42c95911a 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -76,6 +76,7 @@ def planner_tiny():
       planning_horizon=2,
       rollout_agent_type="random",
       batch_size=1,
+      env_type="simulated",
   )
 
 
@@ -86,14 +87,28 @@ def planner_small():
       planning_horizon=16,
       rollout_agent_type="policy",
       batch_size=16,
+      env_type="simulated",
   )
 
 
+def make_env(env_type, real_env, sim_env_kwargs):
+  """Factory function for envs."""
+  return {
+      "real": lambda: real_env.new_like(
+          batch_size=sim_env_kwargs["batch_size"],
+          store_rollouts=False,
+      ),
+      "simulated": lambda: rl_utils.SimulatedBatchGymEnvWithFixedInitialFrames(
+          **sim_env_kwargs
+      ),
+  }[env_type]()
+
+
 def make_agent(
     agent_type, env, policy_hparams, policy_dir, sampling_temp,
     sim_env_kwargs=None, frame_stack_size=None, planning_horizon=None,
     rollout_agent_type=None, batch_size=None, num_rollouts=None,
-    inner_batch_size=None, video_writer=None):
+    inner_batch_size=None, video_writer=None, env_type=None):
   """Factory function for Agents."""
   if batch_size is None:
     batch_size = env.batch_size
@@ -109,9 +124,8 @@ def make_agent(
           batch_size, make_agent(
               rollout_agent_type, env, policy_hparams, policy_dir,
               sampling_temp, batch_size=inner_batch_size
-          ), rl_utils.SimulatedBatchGymEnvWithFixedInitialFrames(
-              **sim_env_kwargs
-          ), lambda env: rl_utils.BatchStackWrapper(env, frame_stack_size),
+          ), make_env(env_type, env.env, sim_env_kwargs),
+          lambda env: rl_utils.BatchStackWrapper(env, frame_stack_size),
           num_rollouts, planning_horizon,
           discount_factor=policy_hparams.gae_gamma, video_writer=video_writer
       ),
@@ -136,7 +150,8 @@ def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
         sim_env_kwargs, loop_hparams.frame_stack_size,
         planner_hparams.planning_horizon, planner_hparams.rollout_agent_type,
         num_rollouts=planner_hparams.num_rollouts,
-        inner_batch_size=planner_hparams.batch_size, video_writer=video_writer
+        inner_batch_size=planner_hparams.batch_size, video_writer=video_writer,
+        env_type=planner_hparams.env_type
     )
     rl_utils.run_rollouts(
         env, agent, env.reset(), log_every_steps=log_every_steps
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 2d9d1d046..08cc200ef 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from copy import deepcopy
 import random
 
 from gym.spaces import Box
@@ -284,7 +285,10 @@ def proceed():
       return num_dones < env.batch_size
 
   while proceed():
-    actions = agent.act(observations)
+    act_kwargs = {}
+    if agent.needs_env_state:
+      act_kwargs["env_state"] = env.state
+    actions = agent.act(observations, **act_kwargs)
     (observations, rewards, dones) = env.step(actions)
     observations = list(observations)
     now_done_indices = []
@@ -324,12 +328,14 @@ class BatchAgent(object):
   Runs a batch of parallel agents. Operates on Numpy arrays.
   """
 
+  needs_env_state = False
+
   def __init__(self, batch_size, observation_space, action_space):
     self.batch_size = batch_size
     self.observation_space = observation_space
     self.action_space = action_space
 
-  def act(self, observations):
+  def act(self, observations, env_state=None):
     """Picks actions based on observations.
 
     Args:
@@ -357,7 +363,8 @@ def estimate_value(self, observations):
 class RandomAgent(BatchAgent):
   """Random agent, sampling actions from the uniform distribution."""
 
-  def act(self, observations):
+  def act(self, observations, env_state=None):
+    del env_state
     return np.array([
         self.action_space.sample() for _ in range(observations.shape[0])
     ])
@@ -400,7 +407,8 @@ def _run(self, observations):
         feed_dict={self._observations_t: observations}
     )
 
-  def act(self, observations):
+  def act(self, observations, env_state=None):
+    del env_state
     (actions, _) = self._run(observations)
     return actions
 
@@ -412,6 +420,8 @@ def estimate_value(self, observations):
 class PlannerAgent(BatchAgent):
   """Agent based on temporal difference planning."""
 
+  needs_env_state = True
+
   def __init__(
       self, batch_size, rollout_agent, sim_env, wrapper_fn, num_rollouts,
       planning_horizon, discount_factor=1.0, video_writer=None
@@ -427,14 +437,18 @@ def __init__(
     self._planning_horizon = planning_horizon
     self._video_writer = video_writer
 
-  def act(self, observations):
+  def act(self, observations, env_state=None):
     # Randomly choose an action to be recorded.
     recorded_action = self.action_space.sample()
 
     def run_batch_from(observation, action, planner_index, batch_index):
       """Run a batch of actions."""
-      self._sim_env.initial_frames = np.array(
-          [observation] * self._sim_env.batch_size
+      self._wrapped_env.set_initial_state(
+          initial_state=[
+              deepcopy(env_state[planner_index])
+              for _ in range(self._sim_env.batch_size)
+          ],
+          initial_frames=np.array([observation] * self._sim_env.batch_size)
       )
       self._wrapped_env.reset()
       (initial_observations, initial_rewards, _) = self._wrapped_env.step(
@@ -515,6 +529,17 @@ def __init__(self, env, stack_size):
         (self.batch_size,) + self.observation_space.shape,
         dtype=inner_space.dtype
     )
+    self._initial_frames = None
+
+  @property
+  def state(self):
+    """Gets the current state."""
+    return self.env.state
+
+  def set_initial_state(self, initial_state, initial_frames):
+    """Sets the state that will be used on next reset."""
+    self.env.set_initial_state(initial_state, initial_frames)
+    self._initial_frames = initial_frames
 
   def reset(self, indices=None):
     if indices is None:
@@ -526,9 +551,16 @@ def reset(self, indices=None):
       assert self.env.initial_frames.shape[1] == self.stack_size
       self._history_buffer[...] = self.env.initial_frames
     except AttributeError:
-      # Otherwise, repeat the first observation stack_size times.
-      for (index, observation) in zip(indices, observations):
-        self._history_buffer[index, ...] = [observation] * self.stack_size
+      # Otherwise, check if set_initial_state was called and we can take the
+      # frames from there.
+      if self._initial_frames is not None:
+        for (index, observation) in zip(indices, observations):
+          assert (self._initial_frames[index, -1, ...] == observation).all()
+          self._history_buffer[index, ...] = self._initial_frames[index, ...]
+      else:
+        # Otherwise, repeat the first observation stack_size times.
+        for (index, observation) in zip(indices, observations):
+          self._history_buffer[index, ...] = [observation] * self.stack_size
     return self._history_buffer
 
   def step(self, actions):
@@ -550,3 +582,13 @@ def initial_frame_chooser(batch_size):
         *args, initial_frame_chooser=initial_frame_chooser, **kwargs
     )
     super(SimulatedBatchGymEnvWithFixedInitialFrames, self).__init__(env)
+
+  @property
+  def state(self):
+    """Gets the current state."""
+    return [None] * self.batch_size
+
+  def set_initial_state(self, initial_state, initial_frames):
+    """Sets the state that will be used on next reset."""
+    del initial_state
+    self.initial_frames = initial_frames

From e7dc15a0fae5233030a52c6d90234978aba6d3e1 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 17 Jan 2019 16:05:52 -0800
Subject: [PATCH 1533/2720] internal merge of PR #1380

PiperOrigin-RevId: 229835664
---
 tensor2tensor/rl/evaluator.py | 4 ++--
 tensor2tensor/rl/rl_utils.py  | 5 +++--
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 42c95911a..d7a76e4f3 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -94,11 +94,11 @@ def planner_small():
 def make_env(env_type, real_env, sim_env_kwargs):
   """Factory function for envs."""
   return {
-      "real": lambda: real_env.new_like(
+      "real": lambda: real_env.new_like(  # pylint: disable=g-long-lambda
           batch_size=sim_env_kwargs["batch_size"],
           store_rollouts=False,
       ),
-      "simulated": lambda: rl_utils.SimulatedBatchGymEnvWithFixedInitialFrames(
+      "simulated": lambda: rl_utils.SimulatedBatchGymEnvWithFixedInitialFrames(  # pylint: disable=g-long-lambda
           **sim_env_kwargs
       ),
   }[env_type]()
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 08cc200ef..34ef64e3b 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -19,7 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from copy import deepcopy
+import copy
 import random
 
 from gym.spaces import Box
@@ -340,6 +340,7 @@ def act(self, observations, env_state=None):
 
     Args:
       observations: A batch of observations.
+      env_state: State.
 
     Returns:
       A batch of actions.
@@ -445,7 +446,7 @@ def run_batch_from(observation, action, planner_index, batch_index):
       """Run a batch of actions."""
       self._wrapped_env.set_initial_state(
           initial_state=[
-              deepcopy(env_state[planner_index])
+              copy.deepcopy(env_state[planner_index])
               for _ in range(self._sim_env.batch_size)
           ],
           initial_frames=np.array([observation] * self._sim_env.batch_size)

From 08e83030acf3ef13d15ad6eaefaa0a67fb20b59d Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 17 Jan 2019 16:57:02 -0800
Subject: [PATCH 1534/2720] One more forgotten hparam guard.

PiperOrigin-RevId: 229843842
---
 tensor2tensor/models/transformer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index cc85ab6ec..f0d0596e8 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1314,8 +1314,8 @@ def transformer_decoder(decoder_input,
               max_length=hparams.get("max_length"),
               decode_loop_step=decode_loop_step,
               vars_3d=hparams.get("attention_variables_3d"),
-              activation_dtype=hparams.activation_dtype,
-              weight_dtype=hparams.weight_dtype)
+              activation_dtype=hparams.get("activation_dtype", "float32"),
+              weight_dtype=hparams.get("weight_dtype", "float32"))
           x = common_layers.layer_postprocess(x, y, hparams)
         if encoder_output is not None:
           with tf.variable_scope("encdec_attention"):

From 113bf535b3fd8ab32b0559fbc9aab7798e3dfd2e Mon Sep 17 00:00:00 2001
From: Dumitru Erhan <dumitru@google.com>
Date: Thu, 17 Jan 2019 19:24:00 -0800
Subject: [PATCH 1535/2720] New config with deterministic simulation starts.

PiperOrigin-RevId: 229860944
---
 tensor2tensor/rl/trainer_model_based_params.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index a7df96c30..cd18af3b7 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -339,6 +339,16 @@ def rlmb_long_stochastic_discrete():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_long_stochastic_discrete_simulation_deterministic_starts():
+  """Long setting with stochastic discrete model & deterministic sim starts."""
+  hparams = rlmb_base_stochastic_discrete()
+  hparams.generative_model_params = "next_frame_basic_stochastic_discrete_long"
+  hparams.ppo_epochs_num = 1000
+  hparams.simulation_random_starts = False
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_long_stochastic_discrete_100steps():
   """Long setting with stochastic discrete model, changed ppo steps."""

From 83eca3715afdaf8a690f4a41fdce2c7a68e30965 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Sat, 19 Jan 2019 22:27:46 -0800
Subject: [PATCH 1536/2720] Add RL configs with 3 epochs and 1 epoch, remove
 some obsolete ones.

PiperOrigin-RevId: 230103318
---
 .../rl/trainer_model_based_params.py          | 155 +++---------------
 1 file changed, 19 insertions(+), 136 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index cd18af3b7..f12667e81 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -137,7 +137,7 @@ def rlmb_ppo_base():
       real_batch_size=1,
       # Number of simulated environments to train on simultaneously.
       simulated_batch_size=16,
-      eval_batch_size=64,
+      eval_batch_size=32,
       wm_policy_param_sharing=False,
 
       # Unused; number of PPO epochs is calculated from the real frame limit.
@@ -385,6 +385,24 @@ def rlmb_long_stochastic_discrete_gamma90():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_long_stochastic_discrete_3epochs():
+  """Long setting with stochastic discrete model, changed epochs."""
+  hparams = rlmb_long_stochastic_discrete()
+  hparams.epochs = 3
+  hparams.ppo_epochs_num = 2000
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_long_stochastic_discrete_1epoch():
+  """Long setting with stochastic discrete model, changed epochs."""
+  hparams = rlmb_long_stochastic_discrete()
+  hparams.epochs = 1
+  hparams.ppo_epochs_num = 3000
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_base_recurrent():
   """Base setting with recurrent model."""
@@ -440,45 +458,6 @@ def rlmb_base_sv2p_deterministic_softmax():
   return hparams
 
 
-@registry.register_hparams
-def rlmb_base_sv2p_flippy30():
-  """Base setting with sv2p as world model."""
-  hparams = rlmb_base()
-  hparams.epochs = 30
-  hparams.ppo_epochs_num = 1000
-  hparams.model_train_steps = 15000
-  hparams.learning_rate_bump = 1.0
-  hparams.initial_epoch_train_steps_multiplier = 5
-  hparams.generative_model = "next_frame_sv2p"
-  hparams.generative_model_params = "next_frame_sv2p_atari"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_softmax_flippy30():
-  """Base setting with sv2p as world model with softmax."""
-  hparams = rlmb_base_sv2p_flippy30()
-  hparams.generative_model_params = "next_frame_sv2p_atari_softmax"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_deterministic_flippy30():
-  """Base setting with deterministic sv2p as world model."""
-  hparams = rlmb_base_sv2p_flippy30()
-  hparams.generative_model_params = "next_frame_sv2p_atari_deterministic"
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_base_sv2p_deterministic_softmax_flippy30():
-  """Base setting with deterministic sv2p as world model with softmax."""
-  hparams = rlmb_base_sv2p_softmax_flippy30()
-  hparams.generative_model_params = (
-      "next_frame_sv2p_atari_softmax_deterministic")
-  return hparams
-
-
 @registry.register_hparams
 def rlmb_base_sampling():
   """Base setting with a stochastic next-frame model."""
@@ -495,60 +474,6 @@ def rlmb_base_sampling_noresize():
   return hparams
 
 
-@registry.register_hparams
-def rlmb_flippy60():
-  """Schedule with a lot of epochs (slow)."""
-  hparams = rlmb_base_sampling()
-  hparams.epochs = 60
-  hparams.ppo_epochs_num = 500
-  hparams.model_train_steps = 10000
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_flippy30():
-  """Schedule with a lot of epochs (slow)."""
-  hparams = rlmb_base_sampling()
-  hparams.epochs = 30
-  hparams.ppo_epochs_num = 1000
-  hparams.model_train_steps = 15000
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_medium():
-  """Small set for larger testing."""
-  hparams = rlmb_base()
-  hparams.num_real_env_frames //= 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_25k():
-  """Small set for larger testing."""
-  hparams = rlmb_medium()
-  hparams.num_real_env_frames //= 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_short():
-  """Small set for larger testing."""
-  hparams = rlmb_base()
-  hparams.num_real_env_frames //= 5
-  hparams.model_train_steps //= 10
-  hparams.ppo_epochs_num //= 10
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_model_only():
-  hp = rlmb_base()
-  hp.epochs = 1
-  hp.ppo_epochs_num = 0
-  return hp
-
-
 def _rlmb_tiny_overrides():
   """Parameters to override for tiny setting excluding agent-related hparams."""
   return dict(
@@ -637,48 +562,6 @@ def rlmb_tiny_sv2p():
   return hparams
 
 
-@registry.register_hparams
-def rlmb_ae_base():
-  """Parameter set for autoencoders."""
-  hparams = rlmb_base()
-  hparams.ppo_params = "ppo_pong_ae_base"
-  hparams.generative_model_params = "next_frame_ae"
-  hparams.autoencoder_hparams_set = "autoencoder_discrete_pong"
-  hparams.autoencoder_train_steps = 5000
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  hparams.grayscale = False
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_ae_basetest():
-  """Base AE setting but quicker with only 2 epochs."""
-  hparams = rlmb_ae_base()
-  hparams.game = "pong"
-  hparams.epochs = 2
-  hparams.num_real_env_frames = 3200
-  hparams.model_train_steps = 100
-  hparams.autoencoder_train_steps = 10
-  hparams.ppo_epochs_num = 2
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_ae_tiny():
-  """Tiny set for testing autoencoders."""
-  hparams = rlmb_tiny()
-  hparams.ppo_params = "ppo_pong_ae_base"
-  hparams.generative_model_params = "next_frame_ae_tiny"
-  hparams.autoencoder_hparams_set = "autoencoder_discrete_tiny"
-  hparams.resize_height_factor = 1
-  hparams.resize_width_factor = 1
-  hparams.grayscale = False
-  hparams.autoencoder_train_steps = 1
-  hparams.autoencoder_train_steps_initial_multiplier = 0
-  return hparams
-
-
 @registry.register_hparams
 def rlmb_tiny_simulation_deterministic_starts():
   hp = rlmb_tiny()

From 8b88b13dd65bf52b3c27663a128adb7b0a5773fb Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Sat, 19 Jan 2019 22:28:41 -0800
Subject: [PATCH 1537/2720] Allow to run evlauator in distributed mode.

PiperOrigin-RevId: 230103351
---
 tensor2tensor/rl/evaluator.py | 68 +++++++++++++++++++++++++++++------
 1 file changed, 57 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index d7a76e4f3..1359c0fa6 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -28,6 +28,10 @@
 from __future__ import division
 from __future__ import print_function
 
+import datetime
+import os
+
+from tensor2tensor.data_generators import gym_env
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl  # pylint: disable=unused-import
 from tensor2tensor.rl import rl_utils
@@ -42,7 +46,9 @@
 flags = tf.flags
 FLAGS = flags.FLAGS
 
-
+flags.DEFINE_string("output_dir", "", "Main directory for multi-runs.")
+flags.DEFINE_integer("total_num_workers", 1, "How many workers in total.")
+flags.DEFINE_string("worker_to_game_map", "", "How to map workers to games.")
 flags.DEFINE_string("policy_dir", "", "Directory with policy checkpoints.")
 flags.DEFINE_string("model_dir", "", "Directory with model checkpoints.")
 flags.DEFINE_string(
@@ -58,16 +64,23 @@
     "out-of-graph one. Works only with --agent=policy."
 )
 flags.DEFINE_string(
-    "planner_hparams_set", "planner_tiny", "Planner hparam set."
+    "planner_hparams_set", "planner_small", "Planner hparam set."
 )
 flags.DEFINE_string("planner_hparams", "", "Planner hparam overrides.")
 flags.DEFINE_integer(
-    "log_every_steps", 0, "Log every how many environment steps."
+    "log_every_steps", 20, "Log every how many environment steps."
 )
 flags.DEFINE_string(
     "debug_video_path", "", "Path to save the planner debug video at."
 )
 
+# Unused flags needed to pass for multi-run infrastructure.
+flags.DEFINE_bool("autotune", False, "Unused here.")
+flags.DEFINE_string("objective", "", "Unused here.")
+flags.DEFINE_string("client_handle", "client_0", "Unused.")
+flags.DEFINE_bool("maximize_tuner_objective", True, "Unused.")
+flags.DEFINE_integer("vizier_search_algorithm", 0, "Unused.")
+
 
 @registry.register_hparams
 def planner_tiny():
@@ -83,10 +96,10 @@ def planner_tiny():
 @registry.register_hparams
 def planner_small():
   return tf.contrib.training.HParams(
-      num_rollouts=16,
+      num_rollouts=64,
       planning_horizon=16,
       rollout_agent_type="policy",
-      batch_size=16,
+      batch_size=64,
       env_type="simulated",
   )
 
@@ -173,14 +186,12 @@ def evaluate(
     assert report_metric is not None
 
   eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
+  video_writer = None
   kwargs = {}
   if not eval_with_learner:
     if debug_video_path:
       video_writer = common_video.WholeVideoWriter(
-          fps=10, output_path=debug_video_path, file_format="avi"
-      )
-    else:
-      video_writer = None
+          fps=10, output_path=debug_video_path, file_format="avi")
     kwargs["eval_fn"] = make_eval_fn_with_agent(
         agent_type, planner_hparams, model_dir, log_every_steps=log_every_steps,
         video_writer=video_writer
@@ -207,18 +218,53 @@ def evaluate(
   return eval_metrics
 
 
+def get_game_for_worker(map_name, directory_id):
+  """Get game for the given worker (directory) id."""
+  if map_name == "v100unfriendly":
+    games = ["chopper_command", "boxing", "asterix", "seaquest"]
+    worker_per_game = 5
+  elif map_name == "human_nice":
+    games = gym_env.ATARI_GAMES_WITH_HUMAN_SCORE_NICE
+    worker_per_game = 5
+  else:
+    raise ValueError("Unknown worker to game map name: %s" % map_name)
+  games.sort()
+  game_id = (directory_id - 1) // worker_per_game
+  tf.logging.info("Getting game %d from %s." % (game_id, games))
+  return games[game_id]
+
+
 def main(_):
+  now = datetime.datetime.now()
+  now_tag = now.strftime("%Y_%m_%d_%H_%M")
   loop_hparams = trainer_lib.create_hparams(
       FLAGS.loop_hparams_set, FLAGS.loop_hparams
   )
+  if FLAGS.worker_to_game_map and FLAGS.total_num_workers > 1:
+    loop_hparams.game = get_game_for_worker(
+        FLAGS.worker_to_game_map, FLAGS.worker_id + 1)
+    tf.logging.info("Set game to %s." % loop_hparams.game)
   if FLAGS.full_eval:
     loop_hparams.eval_rl_env_max_episode_steps = -1
   planner_hparams = trainer_lib.create_hparams(
       FLAGS.planner_hparams_set, FLAGS.planner_hparams
   )
+  policy_dir = FLAGS.policy_dir
+  model_dir = FLAGS.model_dir
+  eval_metrics_dir = FLAGS.eval_metrics_dir
+  if FLAGS.output_dir:
+    cur_dir = FLAGS.output_dir
+    if FLAGS.total_num_workers > 1:
+      cur_dir = os.path.join(cur_dir, "%d" % (FLAGS.worker_id + 1))
+    policy_dir = os.path.join(cur_dir, "policy")
+    model_dir = os.path.join(cur_dir, "world_model")
+    eval_metrics_dir = os.path.join(cur_dir, "evaluator_" + now_tag)
+    tf.logging.info("Writing metrics to %s." % eval_metrics_dir)
+    if not tf.gfile.Exists(eval_metrics_dir):
+      tf.gfile.MkDir(eval_metrics_dir)
   evaluate(
-      loop_hparams, planner_hparams, FLAGS.policy_dir, FLAGS.model_dir,
-      FLAGS.eval_metrics_dir, FLAGS.agent, FLAGS.eval_with_learner,
+      loop_hparams, planner_hparams, policy_dir, model_dir,
+      eval_metrics_dir, FLAGS.agent, FLAGS.eval_with_learner,
       FLAGS.log_every_steps if FLAGS.log_every_steps > 0 else None,
       debug_video_path=FLAGS.debug_video_path
   )

From 1705fa51ad8192407b8fbf6d11077245532f3b79 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Mon, 21 Jan 2019 01:01:16 +0100
Subject: [PATCH 1538/2720] Sample first action according to policy in planner
 use UCT (#1386)

---
 tensor2tensor/rl/evaluator.py |  15 ++++-
 tensor2tensor/rl/rl_utils.py  | 107 ++++++++++++++++++++++++++--------
 2 files changed, 95 insertions(+), 27 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 1359c0fa6..4ca9c8637 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -90,6 +90,8 @@ def planner_tiny():
       rollout_agent_type="random",
       batch_size=1,
       env_type="simulated",
+      uct_const=0.0,
+      uniform_first_action=True,
   )
 
 
@@ -101,6 +103,8 @@ def planner_small():
       rollout_agent_type="policy",
       batch_size=64,
       env_type="simulated",
+      uct_const=0.0,
+      uniform_first_action=True,
   )
 
 
@@ -121,7 +125,9 @@ def make_agent(
     agent_type, env, policy_hparams, policy_dir, sampling_temp,
     sim_env_kwargs=None, frame_stack_size=None, planning_horizon=None,
     rollout_agent_type=None, batch_size=None, num_rollouts=None,
-    inner_batch_size=None, video_writer=None, env_type=None):
+    inner_batch_size=None, video_writer=None, env_type=None,
+    uct_const=None, uniform_first_action=None
+):
   """Factory function for Agents."""
   if batch_size is None:
     batch_size = env.batch_size
@@ -140,7 +146,9 @@ def make_agent(
           ), make_env(env_type, env.env, sim_env_kwargs),
           lambda env: rl_utils.BatchStackWrapper(env, frame_stack_size),
           num_rollouts, planning_horizon,
-          discount_factor=policy_hparams.gae_gamma, video_writer=video_writer
+          discount_factor=policy_hparams.gae_gamma,
+          uct_const=uct_const, uniform_first_action=uniform_first_action,
+          video_writer=video_writer
       ),
   }[agent_type]()
 
@@ -164,7 +172,8 @@ def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
         planner_hparams.planning_horizon, planner_hparams.rollout_agent_type,
         num_rollouts=planner_hparams.num_rollouts,
         inner_batch_size=planner_hparams.batch_size, video_writer=video_writer,
-        env_type=planner_hparams.env_type
+        env_type=planner_hparams.env_type, uct_const=planner_hparams.uct_const,
+        uniform_first_action=planner_hparams.uniform_first_action
     )
     rl_utils.run_rollouts(
         env, agent, env.reset(), log_every_steps=log_every_steps
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 34ef64e3b..46c9ec809 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import copy
+import math
 import random
 
 from gym.spaces import Box
@@ -360,6 +361,19 @@ def estimate_value(self, observations):
     """
     raise NotImplementedError
 
+  def action_distribution(self, observations):
+    """Calculates action distribution based on observations.
+
+    Used for temporal-difference planning.
+
+    Args:
+      observations: A batch of observations.
+
+    Returns:
+      A batch of action probabilities.
+    """
+    raise NotImplementedError
+
 
 class RandomAgent(BatchAgent):
   """Random agent, sampling actions from the uniform distribution."""
@@ -373,6 +387,11 @@ def act(self, observations, env_state=None):
   def estimate_value(self, observations):
     return np.zeros(observations.shape[0])
 
+  def action_distribution(self, observations):
+    return np.full(
+        (observations.shape[0], self.action_space.n), 1.0 / self.action_space.n
+    )
+
 
 class PolicyAgent(BatchAgent):
   """Agent based on a policy network."""
@@ -394,6 +413,7 @@ def __init__(
           self._observations_t, policy_hparams, self.action_space
       )
       actions = common_layers.sample_with_temperature(logits, sampling_temp)
+      self._probs_t = tf.nn.softmax(logits / sampling_temp)
       self._actions_t = tf.cast(actions, tf.int32)
       model_saver = tf.train.Saver(
           tf.global_variables(policy_hparams.policy_network + "/.*")  # pylint: disable=unexpected-keyword-arg
@@ -404,19 +424,23 @@ def __init__(
 
   def _run(self, observations):
     return self._sess.run(
-        [self._actions_t, self._values_t],
+        [self._actions_t, self._values_t, self._probs_t],
         feed_dict={self._observations_t: observations}
     )
 
   def act(self, observations, env_state=None):
     del env_state
-    (actions, _) = self._run(observations)
+    (actions, _, _) = self._run(observations)
     return actions
 
   def estimate_value(self, observations):
-    (_, values) = self._run(observations)
+    (_, values, _) = self._run(observations)
     return values
 
+  def action_distribution(self, observations):
+    (_, _, probs) = self._run(observations)
+    return probs
+
 
 class PlannerAgent(BatchAgent):
   """Agent based on temporal difference planning."""
@@ -425,7 +449,8 @@ class PlannerAgent(BatchAgent):
 
   def __init__(
       self, batch_size, rollout_agent, sim_env, wrapper_fn, num_rollouts,
-      planning_horizon, discount_factor=1.0, video_writer=None
+      planning_horizon, discount_factor=1.0, uct_const=0,
+      uniform_first_action=True, video_writer=None
   ):
     super(PlannerAgent, self).__init__(
         batch_size, rollout_agent.observation_space, rollout_agent.action_space
@@ -433,30 +458,34 @@ def __init__(
     self._rollout_agent = rollout_agent
     self._sim_env = sim_env
     self._wrapped_env = wrapper_fn(sim_env)
+    self._num_rollouts = num_rollouts
     self._num_batches = num_rollouts // rollout_agent.batch_size
     self._discount_factor = discount_factor
     self._planning_horizon = planning_horizon
+    self._uct_const = uct_const
+    self._uniform_first_action = uniform_first_action
     self._video_writer = video_writer
 
   def act(self, observations, env_state=None):
-    # Randomly choose an action to be recorded.
-    recorded_action = self.action_space.sample()
-
-    def run_batch_from(observation, action, planner_index, batch_index):
+    def run_batch_from(observation, planner_index, batch_index):
       """Run a batch of actions."""
+      repeated_observation = np.array(
+          [observation] * self._wrapped_env.batch_size
+      )
+      actions = self._get_first_actions(repeated_observation)
       self._wrapped_env.set_initial_state(
           initial_state=[
               copy.deepcopy(env_state[planner_index])
               for _ in range(self._sim_env.batch_size)
           ],
-          initial_frames=np.array([observation] * self._sim_env.batch_size)
+          initial_frames=repeated_observation
       )
       self._wrapped_env.reset()
       (initial_observations, initial_rewards, _) = self._wrapped_env.step(
-          np.array([action] * self._wrapped_env.batch_size)
+          actions
       )
       writer = None
-      if planner_index == 0 and batch_index == 0 and action == recorded_action:
+      if planner_index == 0 and batch_index == 0:
         writer = self._video_writer
       (final_observations, cum_rewards) = run_rollouts(
           self._wrapped_env, self._rollout_agent, initial_observations,
@@ -468,27 +497,57 @@ def run_batch_from(observation, action, planner_index, batch_index):
           initial_rewards + self._discount_factor * cum_rewards +
           self._discount_factor ** (self._planning_horizon + 1) * values
       )
-      return total_values.mean()
-
-    def run_batches_from(observation, action, planner_index):
-      return sum(
-          run_batch_from(observation, action, planner_index, i)
-          for i in range(self._num_batches)
-      ) / self._num_batches
+      return list(zip(actions, total_values))
+
+    def run_batches_from(observation, planner_index):
+      sums = {a: 0 for a in range(self.action_space.n)}
+      counts = copy.copy(sums)
+      for i in range(self._num_batches):
+        for (action, total_value) in run_batch_from(
+            observation, planner_index, i
+        ):
+          sums[action] += total_value
+          counts[action] += 1
+      return {a: (sums[a], counts[a]) for a in sums}
 
     def choose_best_action(observation, planner_index):
-      return max(
-          range(self.action_space.n),
-          key=(lambda action: run_batches_from(  # pylint: disable=g-long-lambda
-              observation, action, planner_index
-          ))
-      )
+      action_probs = self._rollout_agent.action_distribution(
+          np.array([observation] * self._rollout_agent.batch_size)
+      )[0, :]
+      sums_and_counts = run_batches_from(observation, planner_index)
+
+      def uct(action):
+        (value_sum, count) = sums_and_counts[action]
+        if count > 0:
+          mean_value = value_sum / count
+        else:
+          mean_value = 0
+        return mean_value + self._uct_bonus(
+            count, action_probs[action]
+        )
+
+      return max(range(self.action_space.n), key=uct)
 
     return np.array([
         choose_best_action(observation, i)
         for (i, observation) in enumerate(observations)
     ])
 
+  def _uct_bonus(self, count, prob):
+    return self._uct_const * prob * math.sqrt(
+        math.log(self._num_rollouts) / (1 + count)
+    )
+
+  def _get_first_actions(self, observations):
+    if self._uniform_first_action:
+      return np.array([
+          int(x) for x in np.linspace(
+              0, self.action_space.n, self._num_rollouts + 1
+          )
+      ])[:self._num_rollouts]
+    else:
+      return list(sorted(self._rollout_agent.act(observations)))
+
 
 # TODO(koz4k): Unify interfaces of batch envs.
 class BatchWrapper(object):

From e98710d50bc7ae0195de0cf18bfab31e952f9179 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Sun, 20 Jan 2019 16:01:34 -0800
Subject: [PATCH 1539/2720] internal merge of PR #1386

PiperOrigin-RevId: 230152220
---
 tensor2tensor/rl/rl_utils.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 46c9ec809..c3e63cb4c 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -511,6 +511,7 @@ def run_batches_from(observation, planner_index):
       return {a: (sums[a], counts[a]) for a in sums}
 
     def choose_best_action(observation, planner_index):
+      """Choose the best action."""
       action_probs = self._rollout_agent.action_distribution(
           np.array([observation] * self._rollout_agent.batch_size)
       )[0, :]

From 271e51cc261022bd0aa2ead8be30070d13469131 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sun, 20 Jan 2019 16:14:57 -0800
Subject: [PATCH 1540/2720] In register_gym_env return the env's name along
 with the env.

PiperOrigin-RevId: 230153145
---
 tensor2tensor/rl/gym_utils.py      | 10 +++++++---
 tensor2tensor/rl/gym_utils_test.py |  4 +++-
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index 7d8a494a0..cc86db34f 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -21,6 +21,7 @@
 
 import gym
 import numpy as np
+import tensorflow as tf
 
 
 class MaxAndSkipEnv(gym.Wrapper):
@@ -102,14 +103,17 @@ def make_gym_env(name, rl_env_max_episode_steps=-1, maxskip_env=False):
 
 
 def register_gym_env(class_entry_point, version="v0"):
-  """Registers the class with its snake case name in Gym and returns it."""
+  """Registers the class in Gym and returns the registered name and the env."""
 
   split_on_colon = class_entry_point.split(":")
   assert len(split_on_colon) == 2
 
   class_name = split_on_colon[1]
   # We have to add the version to conform to gym's API.
-  env_name = "{}-{}".format(class_name, version)
+  env_name = "T2TEnv-{}-{}".format(class_name, version)
   gym.envs.register(id=env_name, entry_point=class_entry_point)
 
-  return gym.make(env_name)
+  tf.logging.info("Entry Point [%s] registered with id [%s]",
+                  class_entry_point, env_name)
+
+  return env_name, gym.make(env_name)
diff --git a/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index 16dbe4f1f..e97762da3 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -66,9 +66,11 @@ def test_unlimited_env(self):
     self.assertTrue(env._max_episode_steps is None)
 
   def test_gym_registration(self):
-    env = gym_utils.register_gym_env(
+    reg_id, env = gym_utils.register_gym_env(
         "tensor2tensor.rl.gym_utils_test:SimpleEnv")
 
+    self.assertEqual("T2TEnv-SimpleEnv-v0", reg_id)
+
     # Most basic check.
     self.assertTrue(isinstance(env, gym.Env))
 

From bcf0b1dcdbd7a05c4f358d2d76380b69df88891f Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sun, 20 Jan 2019 16:15:57 -0800
Subject: [PATCH 1541/2720] Make StackWrapper.simulate handle history = 1 and
 generalize

PiperOrigin-RevId: 230153194
---
 tensor2tensor/rl/envs/tf_atari_wrappers.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index d72ae756b..7d5ac6b20 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -95,6 +95,7 @@ def __init__(self, batch_env, history=4):
     super(StackWrapper, self).__init__(batch_env)
     self.history = history
     self.old_shape = batch_env.observ_shape
+    # TODO(afrozm): Make into tf.get_variable and use_resource=True
     self._observ = tf.Variable(
         tf.zeros((len(self),) + self.observ_shape, self.observ_dtype),
         trainable=False)
@@ -110,6 +111,14 @@ def simulate(self, action):
     reward, done = self._batch_env.simulate(action)
     with tf.control_dependencies([reward, done]):
       new_observ = tf.expand_dims(self._batch_env.observ, axis=1)
+
+      # If we shouldn't stack, i.e. self.history == 1, then just assign
+      # new_observ to self._observ and return from here.
+      if self.history == 1:
+        with tf.control_dependencies([self._observ.assign(new_observ)]):
+          return tf.identity(reward), tf.identity(done)
+
+      # If we should stack, then do the required work.
       old_observ = tf.gather(
           self._observ.read_value(),
           list(range(1, self.history)),
@@ -124,8 +133,11 @@ def _reset_non_empty(self, indices):
     new_values = self._batch_env._reset_non_empty(indices)
     # pylint: enable=protected-access
     initial_frames = getattr(self._batch_env, "history_observations", None)
+
+    num_dimensions_in_env_observation = len(self.old_shape)
+
     if initial_frames is None:
-      inx = [1, self.history, 1, 1, 1]
+      inx = [1, self.history] + ([1] * num_dimensions_in_env_observation)
       initial_frames = tf.tile(tf.expand_dims(new_values, axis=1), inx)
     with tf.control_dependencies([new_values]):
       assign_op = tf.scatter_update(self._observ, indices, initial_frames)

From bb78832a32e58b39a03f28a38f0a535066ea9f32 Mon Sep 17 00:00:00 2001
From: blazejosinski <804945+blazejosinski@users.noreply.github.com>
Date: Mon, 21 Jan 2019 01:17:16 +0100
Subject: [PATCH 1542/2720] Using last env for debug. (#1381)

This is to make debug videos more diverse, as env 0 will start at the beginning of the episode,
when flag simulation_flip_first_random_for_beginning is present.
---
 tensor2tensor/rl/envs/simulated_batch_env.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 874fdccd2..373a1549e 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -279,8 +279,8 @@ def _video_dump_frame(self, obs, rews):
           file_format="avi")
     img = PIL_Image().new("RGB", (obs.shape[-2], 11),)
     draw = PIL_ImageDraw().Draw(img)
-    draw.text((0, 0), "r:{:3}".format(int(rews[0])), fill=(255, 0, 0))
-    self._video_writer.write(np.concatenate([np.asarray(img), obs[0]], axis=0))
+    draw.text((0, 0), "r:{:3}".format(int(rews[-1])), fill=(255, 0, 0))
+    self._video_writer.write(np.concatenate([np.asarray(img), obs[-1]], axis=0))
 
   def _video_dump_frames(self, obs):
     if not self._ffmpeg_works:

From 92ebc7152e0f4f42871251f17dbe6db8409d4fae Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sun, 20 Jan 2019 16:19:05 -0800
Subject: [PATCH 1543/2720] Add a new top-level directory in T2T envs/ with a
 sample tic-tac-toe implementation and tests.

PiperOrigin-RevId: 230153361
---
 tensor2tensor/envs/__init__.py             |  23 ++
 tensor2tensor/envs/tic_tac_toe_env.py      | 249 +++++++++++++++++++++
 tensor2tensor/envs/tic_tac_toe_env_test.py |  82 +++++++
 3 files changed, 354 insertions(+)
 create mode 100644 tensor2tensor/envs/__init__.py
 create mode 100644 tensor2tensor/envs/tic_tac_toe_env.py
 create mode 100644 tensor2tensor/envs/tic_tac_toe_env_test.py

diff --git a/tensor2tensor/envs/__init__.py b/tensor2tensor/envs/__init__.py
new file mode 100644
index 000000000..d5201e5d7
--- /dev/null
+++ b/tensor2tensor/envs/__init__.py
@@ -0,0 +1,23 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Environments defined in T2T. Imports here force registration."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.envs import tic_tac_toe_env
+
diff --git a/tensor2tensor/envs/tic_tac_toe_env.py b/tensor2tensor/envs/tic_tac_toe_env.py
new file mode 100644
index 000000000..e79012e3a
--- /dev/null
+++ b/tensor2tensor/envs/tic_tac_toe_env.py
@@ -0,0 +1,249 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Gym Tic-Tac-Toe environment.
+
+Environment acts like the second player and first player is either environment
+or the agent. The environment follows a random policy for now.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gym
+from gym import spaces
+from gym.utils import seeding
+import numpy as np
+
+from tensor2tensor.data_generators import problem
+from tensor2tensor.layers import modalities
+from tensor2tensor.rl import gym_utils
+
+import tensorflow as tf
+
+
+def encode_pos(i, j):
+  """Encodes a pair (i, j) as a scalar position on the board."""
+  return 3 * i + j
+
+
+def decode_pos(pos):
+  """Decoes a scalar position on the board as a pair (i, j)."""
+  return pos // 3, pos % 3
+
+
+def get_open_spaces(board):
+  """Given a representation of the board, returns a list of open spaces."""
+  open_spaces = []
+  for i in range(3):
+    for j in range(3):
+      if board[i][j] == 0:
+        open_spaces.append(encode_pos(i, j))
+  return open_spaces
+
+
+def get_reward_and_done(board):
+  """Given a representation of the board, returns reward and done."""
+  # Returns (reward, done) where:
+  # reward: -1 means lost, +1 means win, 0 means draw or continuing.
+  # done: True if the game is over, i.e. someone won or it is a draw.
+
+  # Sum all rows ...
+  all_sums = [np.sum(board[i, :]) for i in range(3)]
+  # ... all columns
+  all_sums.extend([np.sum(board[:, i]) for i in range(3)])
+  # and both diagonals.
+  all_sums.append(np.sum([board[i, i] for i in range(3)]))
+  all_sums.append(np.sum([board[i, 2 - i] for i in range(3)]))
+
+  if -3 in all_sums:
+    return -1, True
+
+  if 3 in all_sums:
+    return 1, True
+
+  done = True
+  if get_open_spaces(board):
+    done = False
+
+  return 0, done
+
+
+# TODO(afrozm): This should eventually subclass Problem.
+class TicTacToeEnv(gym.Env):
+  """Simple TicTacToe Env, starts the game randomly half of the time."""
+
+  def __init__(self, strict=False):
+    self.strict = strict
+
+    # What about metadata and spec?
+    self.reward_range = (-1.0, 1.0)
+
+    # Action space -- 9 positions that we can chose to mark.
+    self.action_space = spaces.Discrete(9)
+
+    # Observation space -- this hopefully does what we need.
+    self.observation_space = spaces.Box(
+        low=-1, high=1, shape=(3, 3), dtype=np.int64)
+
+    # Set the seed.
+    self.np_random = None
+    self.seed()
+
+    # Start the game.
+    self.board_state = None
+    self.done = False
+    self.reset()
+
+  def seed(self, seed=None):
+    self.np_random, seed = seeding.np_random(seed)
+    return [seed]
+
+  # TODO(afrozm): Parametrize by some policy so that the env plays in an optimal
+  # way.
+  def play_random_move(self):
+    # Select open spaces.
+    open_spaces = get_open_spaces(self.board_state)
+
+    if not open_spaces:
+      return False
+
+    # Choose a space and mark it.
+    pos = self.np_random.choice(open_spaces)
+    i, j = decode_pos(pos)
+
+    self.board_state[i, j] = -1
+
+  def reset(self):
+    self.board_state = np.zeros((3, 3), dtype=np.int64)
+
+    # We"ll start with a 50% chance.
+    if self.np_random.choice([0, 1]) == 0:
+      self.play_random_move()
+
+    # Return the observation.
+    return self.board_state
+
+  def render(self, mode="human"):
+    # Unused.
+    del mode
+    board_str = ""
+    for i in range(3):
+      for j in range(3):
+        pos = self.board_state[i, j]
+        if pos == -1:
+          board_str += "x"
+        elif pos == 0:
+          board_str += "-"
+        else:
+          board_str += "o"
+      board_str += "\n"
+    return board_str
+
+  def step(self, action):
+    # Are we already done?
+    if self.strict:
+      assert not self.done
+
+    # Action has to belong to the action state.
+    assert self.action_space.contains(action)
+
+    # Is it a legitimate move, i.e. is that position open to play?
+    is_legit_move = action in get_open_spaces(self.board_state)
+
+    # Shouldn"t be an illegal action -- is a noop if not strict.
+    if self.strict:
+      assert is_legit_move
+
+    # If strict mode is off, then let this be a noop and env not play either.
+    if not is_legit_move:
+      return self.board_state, 0, False, {}
+
+    # This is a legit move, perform the action and check if done, etc etc.
+    i, j = decode_pos(action)
+    self.board_state[i, j] = 1
+    reward, done = get_reward_and_done(self.board_state)
+
+    if done:
+      self.done = True
+      return self.board_state, reward, True, {}
+
+    # If not done already, play our move.
+    self.play_random_move()
+    reward, done = get_reward_and_done(self.board_state)
+    self.done = done
+    return self.board_state, reward, self.done, {}
+
+  def hparams(self, defaults, unused_model_hparams):
+    tf.logging.error("@@@ new tictactoe hparams are being called!")
+    p = defaults
+    p.modality = {
+        "inputs": modalities.IdentitySymbolModality,
+        "targets": modalities.IdentitySymbolModality,
+    }
+    p.vocab_size = {
+        "inputs": 3,
+        "targets": 3,
+    }
+    p.input_space_id = 0  # problem.SpaceID.GENERIC
+    p.target_space_id = 0  # problem.SpaceID.GENERIC
+
+    # TODO(afrozm): This doesn't work without returning the hparams object.
+    return p
+
+
+# TODO(afrozm): Figure out how to get rid of this.
+class DummyPolicyProblemTTT(problem.Problem):
+  """Dummy Problem for running the policy."""
+
+  def __init__(self):
+    super(DummyPolicyProblemTTT, self).__init__()
+    self._ttt_env = TicTacToeEnv()
+
+  def hparams(self, defaults, model_hparams):
+    # Update the env's hparams.
+    self._ttt_env.hparams(defaults, model_hparams)
+    # Do these belong here?
+    defaults.modality.update({
+        "input_action": modalities.SymbolModalityWeightsAll,
+        "input_reward": modalities.SymbolModalityWeightsAll,
+        "target_action": modalities.SymbolModalityWeightsAll,
+        "target_reward": modalities.SymbolModalityWeightsAll,
+        "target_policy": modalities.IdentityModality,
+        "target_value": modalities.IdentityModality,
+    })
+    defaults.vocab_size.update({
+        "input_action": self.num_actions,
+        "input_reward": 3,  # -1, 0, +1 ?
+        "target_action": self.num_actions,
+        "target_reward": 3,  # -1, 0, +1 ?
+        "target_policy": None,
+        "target_value": None,
+    })
+
+  @property
+  def num_actions(self):
+    return self._ttt_env.action_space.n
+
+
+def register():
+  # Register this with gym.
+  unused_tictactoe_id, unused_tictactoe_env = gym_utils.register_gym_env(
+      "tensor2tensor.envs.tic_tac_toe_env:TicTacToeEnv", version="v0")
+
+
+# TODO(afrozm): Fix the registration and make it automatic.
+register()
diff --git a/tensor2tensor/envs/tic_tac_toe_env_test.py b/tensor2tensor/envs/tic_tac_toe_env_test.py
new file mode 100644
index 000000000..66fafd669
--- /dev/null
+++ b/tensor2tensor/envs/tic_tac_toe_env_test.py
@@ -0,0 +1,82 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.envs.tic_tac_toe_env."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.envs import tic_tac_toe_env as ttt_env
+import tensorflow as tf
+
+
+class TicTacToeEnvTest(tf.test.TestCase):
+
+  def test_start(self):
+    ttt = ttt_env.TicTacToeEnv(strict=True)
+    self.assertFalse(ttt.done)
+
+    # At max one move may have been played by the env.
+    spaces = ttt_env.get_open_spaces(ttt.board_state)
+    num_open_spaces = len(spaces)
+    # i.e. either 8 or 9
+    self.assertGreater(num_open_spaces, 7)
+
+    # Play a move
+    observation, reward, done, unused_info = ttt.step(spaces[0])
+
+    # The environment should also have played a move.
+    spaces = ttt_env.get_open_spaces(observation)
+    self.assertEqual(num_open_spaces - 2, len(spaces))
+
+    # Since at-max 3 moves have been played, the game can't end.
+    self.assertEqual(reward, 0)
+    self.assertFalse(done)
+
+  def test_env_actions(self):
+    # Environment keeps taking actions and not us, we should eventually lose.
+    ttt = ttt_env.TicTacToeEnv(strict=True)
+    for _ in range(9):
+      ttt.play_random_move()
+      if ttt.done:
+        break
+
+    reward, done = ttt_env.get_reward_and_done(ttt.board_state)
+    self.assertEqual(-1, reward)
+    self.assertTrue(done)
+
+  def test_keep_playing(self):
+    ttt = ttt_env.TicTacToeEnv(strict=False)
+    done = False
+    while not done:
+      # sample an action from the action space.
+      action = ttt.action_space.sample()
+      # play it -- could be a no-op since we don't see if positions are empty.
+      unused_observation, reward, done, unused_info = ttt.step(action)
+
+    # done is True, so either:
+    # we won
+    # env won or
+    # no space left
+
+    we_won = reward == 1
+    env_won = reward == -1
+    no_space = bool(ttt_env.get_open_spaces(ttt.board_state))
+    self.assertTrue(we_won or env_won or no_space)
+
+
+if __name__ == '__main__':
+  tf.test.main()

From d75fe0b97527e4fb5327f6f24ef773d1c78d63a2 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Tue, 22 Jan 2019 05:22:39 +0100
Subject: [PATCH 1544/2720] Model-Based RL: Planner fixes. (#1392)

* Fix uct value when no rollouts.

* Fix num_rollout vs batch_size for Planner with uniform sampling.
---
 tensor2tensor/rl/rl_utils.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index c3e63cb4c..0231578e2 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -522,7 +522,7 @@ def uct(action):
         if count > 0:
           mean_value = value_sum / count
         else:
-          mean_value = 0
+          mean_value = -np.inf
         return mean_value + self._uct_bonus(
             count, action_probs[action]
         )
@@ -543,9 +543,9 @@ def _get_first_actions(self, observations):
     if self._uniform_first_action:
       return np.array([
           int(x) for x in np.linspace(
-              0, self.action_space.n, self._num_rollouts + 1
+              0, self.action_space.n, self._rollout_agent.batch_size + 1
           )
-      ])[:self._num_rollouts]
+      ])[:self._rollout_agent.batch_size]
     else:
       return list(sorted(self._rollout_agent.act(observations)))
 

From 5623a13d50eeceaccd3d83d71c8ac0caa454622c Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 22 Jan 2019 05:40:30 +0100
Subject: [PATCH 1545/2720] Hparam sets for planner (#1393)

---
 tensor2tensor/rl/evaluator.py                  | 13 +++++++++++++
 tensor2tensor/rl/trainer_model_based_params.py | 14 ++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 4ca9c8637..11c08d90d 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -108,6 +108,19 @@ def planner_small():
   )
 
 
+@registry.register_hparams
+def planner_base():
+  return tf.contrib.training.HParams(
+      num_rollouts=96,
+      batch_size=96,
+      planning_horizon=8,
+      rollout_agent_type="policy",
+      env_type="simulated",
+      uct_const=0.0,
+      uniform_first_action=True,
+  )
+
+
 def make_env(env_type, real_env, sim_env_kwargs):
   """Factory function for envs."""
   return {
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index f12667e81..6b18db680 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -339,6 +339,20 @@ def rlmb_long_stochastic_discrete():
   return hparams
 
 
+def _planner_overrides(hparams):
+  hparams.eval_batch_size = 1
+  hparams.eval_sampling_temps = [3.0]
+  hparams.eval_max_num_noops = 0
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_long_stochastic_discrete_planner():
+  hparams = rlmb_long_stochastic_discrete()
+  hparams = _planner_overrides(hparams)
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_long_stochastic_discrete_simulation_deterministic_starts():
   """Long setting with stochastic discrete model & deterministic sim starts."""

From 2cfdaa17ad15b98b01d43ca38ee09185560d1145 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Mon, 21 Jan 2019 20:22:59 -0800
Subject: [PATCH 1546/2720] internal merge of PR #1392

PiperOrigin-RevId: 230276767
---
 tensor2tensor/rl/evaluator.py                  | 13 -------------
 tensor2tensor/rl/trainer_model_based_params.py | 14 --------------
 2 files changed, 27 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 11c08d90d..4ca9c8637 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -108,19 +108,6 @@ def planner_small():
   )
 
 
-@registry.register_hparams
-def planner_base():
-  return tf.contrib.training.HParams(
-      num_rollouts=96,
-      batch_size=96,
-      planning_horizon=8,
-      rollout_agent_type="policy",
-      env_type="simulated",
-      uct_const=0.0,
-      uniform_first_action=True,
-  )
-
-
 def make_env(env_type, real_env, sim_env_kwargs):
   """Factory function for envs."""
   return {
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 6b18db680..f12667e81 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -339,20 +339,6 @@ def rlmb_long_stochastic_discrete():
   return hparams
 
 
-def _planner_overrides(hparams):
-  hparams.eval_batch_size = 1
-  hparams.eval_sampling_temps = [3.0]
-  hparams.eval_max_num_noops = 0
-  return hparams
-
-
-@registry.register_hparams
-def rlmb_long_stochastic_discrete_planner():
-  hparams = rlmb_long_stochastic_discrete()
-  hparams = _planner_overrides(hparams)
-  return hparams
-
-
 @registry.register_hparams
 def rlmb_long_stochastic_discrete_simulation_deterministic_starts():
   """Long setting with stochastic discrete model & deterministic sim starts."""

From 70026a087146784edba15741bae466b1771faff3 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Mon, 21 Jan 2019 20:40:47 -0800
Subject: [PATCH 1547/2720] internal merge of PR #1393

PiperOrigin-RevId: 230277763
---
 tensor2tensor/rl/evaluator.py                  | 13 +++++++++++++
 tensor2tensor/rl/trainer_model_based_params.py |  9 +++++++++
 2 files changed, 22 insertions(+)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 4ca9c8637..11c08d90d 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -108,6 +108,19 @@ def planner_small():
   )
 
 
+@registry.register_hparams
+def planner_base():
+  return tf.contrib.training.HParams(
+      num_rollouts=96,
+      batch_size=96,
+      planning_horizon=8,
+      rollout_agent_type="policy",
+      env_type="simulated",
+      uct_const=0.0,
+      uniform_first_action=True,
+  )
+
+
 def make_env(env_type, real_env, sim_env_kwargs):
   """Factory function for envs."""
   return {
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index f12667e81..2ac8d5e9b 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -339,6 +339,15 @@ def rlmb_long_stochastic_discrete():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_long_stochastic_discrete_planner():
+  hparams = rlmb_long_stochastic_discrete()
+  hparams.eval_batch_size = 1
+  hparams.eval_sampling_temps = [3.0]
+  hparams.eval_max_num_noops = 0
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_long_stochastic_discrete_simulation_deterministic_starts():
   """Long setting with stochastic discrete model & deterministic sim starts."""

From cb0e069887c64e8b006ef32ba2ada99e6870c293 Mon Sep 17 00:00:00 2001
From: blazejosinski <804945+blazejosinski@users.noreply.github.com>
Date: Tue, 22 Jan 2019 06:02:07 +0100
Subject: [PATCH 1548/2720] Fixes and configs for model free training. (#1391)

* Fixes and configs for model free training.

* Hparam range for mf eval.
---
 tensor2tensor/models/research/rl.py    | 23 +++++++++++++++++++++++
 tensor2tensor/rl/trainer_model_free.py | 10 +++++++++-
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 30a8d0ff2..e2fb68ff5 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -21,6 +21,7 @@
 import gym
 import six
 
+from tensor2tensor.data_generators import gym_env
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.layers import common_hparams
@@ -377,6 +378,28 @@ def rlmf_base():
   return hparams
 
 
+@registry.register_hparams
+def rlmf_final_eval():
+  """Base set of hparams for model-free PPO."""
+  hparams = rlmf_original()
+  hparams.batch_size = 8
+  hparams.eval_sampling_temps=[0.0, 1.0]
+  hparams.eval_rl_env_max_episode_steps = -1
+  hparams.add_hparam("ppo_epoch_length", 128)
+  hparams.add_hparam("ppo_optimization_batch_size", 32)
+  hparams.add_hparam("ppo_epochs_num", 10000)
+  hparams.add_hparam("ppo_eval_every_epochs", 500)
+  hparams.add_hparam("attempt", 0)
+  return hparams
+
+
+@registry.register_ranged_hparams
+def rlmf_human_score_games(rhp):
+  rhp.set_categorical("game",
+                      gym_env.ATARI_GAMES_WITH_HUMAN_SCORE_NICE)
+  rhp.set_discrete("attempt", list(range(5)))
+
+
 @registry.register_hparams
 def rlmf_tiny():
   hparams = rlmf_base()
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 1d8d53162..abc08cfe8 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -28,6 +28,7 @@
 from __future__ import print_function
 
 import pprint
+import os
 
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_utils
@@ -108,10 +109,15 @@ def train(hparams, output_dir, report_fn=None):
 
   tf.logging.vlog(1, "metric_name: %s", metric_name)
 
+  eval_metrics_dir = os.path.join(output_dir, "eval_metrics")
+  eval_metrics_dir = os.path.expanduser(eval_metrics_dir)
+  tf.gfile.MakeDirs(eval_metrics_dir)
+  eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
+
   for i, step in enumerate(steps):
     tf.logging.info("Starting training iteration [%d] for [%d] steps.", i, step)
 
-    policy_hparams.epochs_num = step
+    policy_hparams.epochs_num = eval_every_epochs
     learner.train(hparams.env_fn,
                   policy_hparams,
                   simulated=False,
@@ -125,6 +131,8 @@ def train(hparams, output_dir, report_fn=None):
     tf.logging.info(
         "Agent eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
 
+    rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, i)
+
     if report_fn:
       report_fn(eval_metrics[metric_name], step)
 

From ec56d6014ddf2072a439dabaade8622093817409 Mon Sep 17 00:00:00 2001
From: blazejosinski <804945+blazejosinski@users.noreply.github.com>
Date: Mon, 21 Jan 2019 21:02:28 -0800
Subject: [PATCH 1549/2720] internal merge of PR #1391

PiperOrigin-RevId: 230279058
---
 tensor2tensor/models/research/rl.py    | 16 +++++-----------
 tensor2tensor/rl/trainer_model_free.py |  2 +-
 2 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index e2fb68ff5..bace2bc91 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -21,7 +21,6 @@
 import gym
 import six
 
-from tensor2tensor.data_generators import gym_env
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.layers import common_hparams
@@ -379,27 +378,22 @@ def rlmf_base():
 
 
 @registry.register_hparams
-def rlmf_final_eval():
-  """Base set of hparams for model-free PPO."""
+def rlmf_eval():
+  """Eval set of hparams for model-free PPO."""
   hparams = rlmf_original()
   hparams.batch_size = 8
-  hparams.eval_sampling_temps=[0.0, 1.0]
+  hparams.eval_sampling_temps = [0.0, 0.5, 1.0]
   hparams.eval_rl_env_max_episode_steps = -1
   hparams.add_hparam("ppo_epoch_length", 128)
   hparams.add_hparam("ppo_optimization_batch_size", 32)
   hparams.add_hparam("ppo_epochs_num", 10000)
   hparams.add_hparam("ppo_eval_every_epochs", 500)
+  hparams.add_hparam("ppo_eval_every_epochs", 500)
   hparams.add_hparam("attempt", 0)
+  hparams.add_hparam("moe_loss_coef", 0)
   return hparams
 
 
-@registry.register_ranged_hparams
-def rlmf_human_score_games(rhp):
-  rhp.set_categorical("game",
-                      gym_env.ATARI_GAMES_WITH_HUMAN_SCORE_NICE)
-  rhp.set_discrete("attempt", list(range(5)))
-
-
 @registry.register_hparams
 def rlmf_tiny():
   hparams = rlmf_base()
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index abc08cfe8..eb03bd41a 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -27,8 +27,8 @@
 from __future__ import division
 from __future__ import print_function
 
-import pprint
 import os
+import pprint
 
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_utils

From 28e3bee7a858ae87f1fc1a963bdbe25e7714c2d4 Mon Sep 17 00:00:00 2001
From: Dumitru Erhan <dumitru@google.com>
Date: Mon, 21 Jan 2019 22:00:49 -0800
Subject: [PATCH 1550/2720] removing extraneous add_hparam

PiperOrigin-RevId: 230283012
---
 tensor2tensor/models/research/rl.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index bace2bc91..571cdba09 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -388,7 +388,6 @@ def rlmf_eval():
   hparams.add_hparam("ppo_optimization_batch_size", 32)
   hparams.add_hparam("ppo_epochs_num", 10000)
   hparams.add_hparam("ppo_eval_every_epochs", 500)
-  hparams.add_hparam("ppo_eval_every_epochs", 500)
   hparams.add_hparam("attempt", 0)
   hparams.add_hparam("moe_loss_coef", 0)
   return hparams

From 2765aa1b9287f2790983f293b39d1e171d45e416 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 22 Jan 2019 08:36:42 -0800
Subject: [PATCH 1551/2720] Fix a TicTacToe test bug.

PiperOrigin-RevId: 230345644
---
 tensor2tensor/envs/tic_tac_toe_env_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/envs/tic_tac_toe_env_test.py b/tensor2tensor/envs/tic_tac_toe_env_test.py
index 66fafd669..b2da18570 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_test.py
@@ -74,8 +74,8 @@ def test_keep_playing(self):
 
     we_won = reward == 1
     env_won = reward == -1
-    no_space = bool(ttt_env.get_open_spaces(ttt.board_state))
-    self.assertTrue(we_won or env_won or no_space)
+    space = bool(ttt_env.get_open_spaces(ttt.board_state))
+    self.assertTrue(we_won or env_won or not space)
 
 
 if __name__ == '__main__':

From 1f70f10bed6bafb317fc95ae8954cb254d4f7a84 Mon Sep 17 00:00:00 2001
From: Dominic Jack <thedomjack@gmail.com>
Date: Wed, 23 Jan 2019 04:09:26 +1000
Subject: [PATCH 1552/2720] =?UTF-8?q?changed=20=5Fnormalize=5Fbody=5Foutpu?=
 =?UTF-8?q?t=20to=20sum=20the=20mean=20of=20each=20loss=20when=20loss?=
 =?UTF-8?q?=E2=80=A6=20(#1383)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* changed _normalize_body_output to sum the mean of each loss when loss is a list

* generalized for tuples and tensor-like losses
---
 tensor2tensor/utils/t2t_model.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 37bef16b4..cdade2369 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1691,7 +1691,11 @@ def estimator_spec_predict(self, features, use_tpu=False):
   def _normalize_body_output(self, body_out):
     if isinstance(body_out, tuple):
       output, losses = body_out
-      if not isinstance(losses, dict):
+      if isinstance(losses, (list, tuple)):
+        losses = {"extra": tf.add_n([tf.reduce_mean(l) for l in losses])}
+      elif isinstance(losses, dict):
+        pass
+      else:
         losses = {"extra": tf.reduce_mean(losses)}
     else:
       output = body_out

From 32fa4020ea7c3091b8641236391898ac721ad47e Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Tue, 22 Jan 2019 19:15:08 +0100
Subject: [PATCH 1553/2720] Use TensorFlow 1.12 image on FloydHub (#1387)

FloydHub now has [support for TensorFlow 1.12](https://docs.floydhub.com/guides/tensorflow/) which is require for T2T to work correctly.
---
 floyd.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/floyd.yml b/floyd.yml
index 16ca4fd92..2ca96ec31 100644
--- a/floyd.yml
+++ b/floyd.yml
@@ -1,2 +1,2 @@
-env: tensorflow-1.9
+env: tensorflow-1.12
 machine: gpu

From f0fe6315dcdad17d46c47087dd41fcfc103b1556 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Tue, 22 Jan 2019 19:15:49 +0100
Subject: [PATCH 1554/2720] ML Engine: Improve logging (#1390)

This adds usefull commands on how to interact with the submitted job to the CLI.

Example output:
```shell
INFO:tensorflow:Launched shake_shake_image_mnist_t2t_20190121_123909. See console to track: https://console.cloud.google.com/mlengine/jobs/.
INFO:tensorflow:Interact with the training job from the command line:
INFO:tensorflow:Abort job: gcloud ml-engine jobs cancel shake_shake_image_mnist_t2t_20190121_123909
INFO:tensorflow:Stream logs: gcloud ml-engine jobs stream-logs shake_shake_image_mnist_t2t_20190121_123909
INFO:tensorflow:Open tensorboard: tensorboard --logdir gs://lgeiger-test-bucket/training-test
```
---
 tensor2tensor/utils/cloud_mlengine.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 2e2cfb8a4..37a9437c3 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -390,3 +390,7 @@ def launch():
   launch_job(job_spec)
   tf.logging.info("Launched %s. See console to track: %s.", job_name,
                   CONSOLE_URL)
+  tf.logging.info("Interact with the training job from the command line:")
+  tf.logging.info("Abort job: gcloud ml-engine jobs cancel %s", job_name)
+  tf.logging.info("Stream logs: gcloud ml-engine jobs stream-logs %s", job_name)
+  tf.logging.info("Open tensorboard: tensorboard --logdir %s", train_dir)

From 94b6e682496fcb8803a2e6e27e5fbe71f7cf7148 Mon Sep 17 00:00:00 2001
From: Dominic Jack <thedomjack@gmail.com>
Date: Tue, 22 Jan 2019 10:09:50 -0800
Subject: [PATCH 1555/2720] internal merge of PR #1383

PiperOrigin-RevId: 230361728
---
 floyd.yml                             | 2 +-
 tensor2tensor/utils/cloud_mlengine.py | 4 ----
 2 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/floyd.yml b/floyd.yml
index 2ca96ec31..16ca4fd92 100644
--- a/floyd.yml
+++ b/floyd.yml
@@ -1,2 +1,2 @@
-env: tensorflow-1.12
+env: tensorflow-1.9
 machine: gpu
diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 37a9437c3..2e2cfb8a4 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -390,7 +390,3 @@ def launch():
   launch_job(job_spec)
   tf.logging.info("Launched %s. See console to track: %s.", job_name,
                   CONSOLE_URL)
-  tf.logging.info("Interact with the training job from the command line:")
-  tf.logging.info("Abort job: gcloud ml-engine jobs cancel %s", job_name)
-  tf.logging.info("Stream logs: gcloud ml-engine jobs stream-logs %s", job_name)
-  tf.logging.info("Open tensorboard: tensorboard --logdir %s", train_dir)

From 195bcf531de76785ea910b4b24bdcaf9c5b1e937 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Tue, 22 Jan 2019 10:15:39 -0800
Subject: [PATCH 1556/2720] internal merge of PR #1387

PiperOrigin-RevId: 230362924
---
 floyd.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/floyd.yml b/floyd.yml
index 16ca4fd92..2ca96ec31 100644
--- a/floyd.yml
+++ b/floyd.yml
@@ -1,2 +1,2 @@
-env: tensorflow-1.9
+env: tensorflow-1.12
 machine: gpu

From f70d5575d79d6e833ffd5a6bbc0e71b257fda81d Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Tue, 22 Jan 2019 10:16:11 -0800
Subject: [PATCH 1557/2720] internal merge of PR #1390

PiperOrigin-RevId: 230363031
---
 tensor2tensor/utils/cloud_mlengine.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 2e2cfb8a4..37a9437c3 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -390,3 +390,7 @@ def launch():
   launch_job(job_spec)
   tf.logging.info("Launched %s. See console to track: %s.", job_name,
                   CONSOLE_URL)
+  tf.logging.info("Interact with the training job from the command line:")
+  tf.logging.info("Abort job: gcloud ml-engine jobs cancel %s", job_name)
+  tf.logging.info("Stream logs: gcloud ml-engine jobs stream-logs %s", job_name)
+  tf.logging.info("Open tensorboard: tensorboard --logdir %s", train_dir)

From 7aaaea89996aa2cf5ae46d0c33d7d81a45818a62 Mon Sep 17 00:00:00 2001
From: kngxscn <yincongxian@foxmail.com>
Date: Wed, 23 Jan 2019 02:17:57 +0800
Subject: [PATCH 1558/2720] Add a way to set max_subtoken_length in
 translate_enzh problem. (#1385)

---
 .../data_generators/generator_utils.py         |  4 ++--
 tensor2tensor/data_generators/translate.py     | 12 ++++++++++++
 .../data_generators/translate_enzh.py          | 18 ++++++++++++++++--
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 7031b37f3..a07e723c3 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -358,12 +358,12 @@ def get_or_generate_vocab_inner(data_dir, vocab_filename, vocab_size,
 
 
 def get_or_generate_vocab(data_dir, tmp_dir, vocab_filename, vocab_size,
-                          sources, file_byte_budget=1e6):
+                          sources, file_byte_budget=1e6, max_subtoken_length=None):
   """Generate a vocabulary from the datasets in sources."""
 
   vocab_generator = generate_lines_for_vocab(tmp_dir, sources, file_byte_budget)
   return get_or_generate_vocab_inner(data_dir, vocab_filename, vocab_size,
-                                     vocab_generator)
+                                     vocab_generator, max_subtoken_length)
 
 
 def generate_lines_for_vocab(tmp_dir, sources, file_byte_budget=1e6):
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 82f7eeb8f..5feabeed0 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -48,6 +48,18 @@ def source_data_files(self, dataset_split):
     """Files to be passed to compile_data."""
     raise NotImplementedError()
 
+  @property
+  def max_subtoken_length(self):
+    """Maximum subtoken length when generating vocab.
+
+    SubwordTextEncoder vocabulary building is quadratic-time wrt this variable,
+    setting it to None uses the length of the longest token in the corpus.
+
+    Returns:
+      an integer or None
+    """
+    return 200
+
   def vocab_data_files(self):
     """Files to be passed to get_or_generate_vocab."""
     return self.source_data_files(problem.DatasetSplit.TRAIN)
diff --git a/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index 90fd0872a..9122fd45d 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -189,6 +189,18 @@ def source_vocab_name(self):
   def target_vocab_name(self):
     return "%s.zh" % self.vocab_filename
 
+  @property
+  def max_subtoken_length(self):
+    """Maximum subtoken length when generating vocab.
+
+    SubwordTextEncoder vocabulary building is quadratic-time wrt this variable,
+    setting it to None uses the length of the longest token in the corpus.
+
+    Returns:
+      an integer or None
+    """
+    return 200
+
   def get_training_dataset(self, tmp_dir):
     """UN Parallel Corpus and CWMT Corpus need to be downloaded manually.
 
@@ -223,14 +235,16 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
         self.source_vocab_name,
         self.approx_vocab_size,
         source_datasets,
-        file_byte_budget=1e8)
+        file_byte_budget=1e8,
+        max_subtoken_length=self.max_subtoken_length)
     target_vocab = generator_utils.get_or_generate_vocab(
         data_dir,
         tmp_dir,
         self.target_vocab_name,
         self.approx_vocab_size,
         target_datasets,
-        file_byte_budget=1e8)
+        file_byte_budget=1e8,
+        max_subtoken_length=self.max_subtoken_length)
     tag = "train" if train else "dev"
     filename_base = "wmt_enzh_%sk_tok_%s" % (self.approx_vocab_size, tag)
     data_path = translate.compile_data(tmp_dir, datasets, filename_base)

From 1b270351edd8172af8c90d1d1592a824c37ac7e1 Mon Sep 17 00:00:00 2001
From: kngxscn <yincongxian@foxmail.com>
Date: Tue, 22 Jan 2019 10:33:09 -0800
Subject: [PATCH 1559/2720] internal merge of PR #1385

PiperOrigin-RevId: 230366370
---
 tensor2tensor/data_generators/generator_utils.py |  3 ++-
 tensor2tensor/data_generators/translate.py       | 12 ------------
 tensor2tensor/data_generators/translate_enzh.py  | 12 ------------
 3 files changed, 2 insertions(+), 25 deletions(-)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index a07e723c3..9276674eb 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -358,7 +358,8 @@ def get_or_generate_vocab_inner(data_dir, vocab_filename, vocab_size,
 
 
 def get_or_generate_vocab(data_dir, tmp_dir, vocab_filename, vocab_size,
-                          sources, file_byte_budget=1e6, max_subtoken_length=None):
+                          sources, file_byte_budget=1e6,
+                          max_subtoken_length=None):
   """Generate a vocabulary from the datasets in sources."""
 
   vocab_generator = generate_lines_for_vocab(tmp_dir, sources, file_byte_budget)
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 5feabeed0..82f7eeb8f 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -48,18 +48,6 @@ def source_data_files(self, dataset_split):
     """Files to be passed to compile_data."""
     raise NotImplementedError()
 
-  @property
-  def max_subtoken_length(self):
-    """Maximum subtoken length when generating vocab.
-
-    SubwordTextEncoder vocabulary building is quadratic-time wrt this variable,
-    setting it to None uses the length of the longest token in the corpus.
-
-    Returns:
-      an integer or None
-    """
-    return 200
-
   def vocab_data_files(self):
     """Files to be passed to get_or_generate_vocab."""
     return self.source_data_files(problem.DatasetSplit.TRAIN)
diff --git a/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index 9122fd45d..8e4930dbf 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -189,18 +189,6 @@ def source_vocab_name(self):
   def target_vocab_name(self):
     return "%s.zh" % self.vocab_filename
 
-  @property
-  def max_subtoken_length(self):
-    """Maximum subtoken length when generating vocab.
-
-    SubwordTextEncoder vocabulary building is quadratic-time wrt this variable,
-    setting it to None uses the length of the longest token in the corpus.
-
-    Returns:
-      an integer or None
-    """
-    return 200
-
   def get_training_dataset(self, tmp_dir):
     """UN Parallel Corpus and CWMT Corpus need to be downloaded manually.
 

From b67e51f7049ddb266d59dc23df99483b944a95cb Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Tue, 22 Jan 2019 21:10:02 +0100
Subject: [PATCH 1560/2720] MC rollout values normalization for Planner. (when
 ensembled with policy) (#1396)

---
 tensor2tensor/rl/evaluator.py |  9 ++++++---
 tensor2tensor/rl/rl_utils.py  | 32 +++++++++++++++++++++++++-------
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 11c08d90d..5b76dfc67 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -92,6 +92,7 @@ def planner_tiny():
       env_type="simulated",
       uct_const=0.0,
       uniform_first_action=True,
+      uct_std_normalization=False,
   )
 
 
@@ -105,6 +106,7 @@ def planner_small():
       env_type="simulated",
       uct_const=0.0,
       uniform_first_action=True,
+      uct_std_normalization=False,
   )
 
 
@@ -139,7 +141,7 @@ def make_agent(
     sim_env_kwargs=None, frame_stack_size=None, planning_horizon=None,
     rollout_agent_type=None, batch_size=None, num_rollouts=None,
     inner_batch_size=None, video_writer=None, env_type=None,
-    uct_const=None, uniform_first_action=None
+    uct_const=None, uct_std_normalization=None, uniform_first_action=None
 ):
   """Factory function for Agents."""
   if batch_size is None:
@@ -160,8 +162,8 @@ def make_agent(
           lambda env: rl_utils.BatchStackWrapper(env, frame_stack_size),
           num_rollouts, planning_horizon,
           discount_factor=policy_hparams.gae_gamma,
-          uct_const=uct_const, uniform_first_action=uniform_first_action,
-          video_writer=video_writer
+          uct_const=uct_const, uct_std_normalization=uct_std_normalization,
+          uniform_first_action=uniform_first_action, video_writer=video_writer
       ),
   }[agent_type]()
 
@@ -186,6 +188,7 @@ def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
         num_rollouts=planner_hparams.num_rollouts,
         inner_batch_size=planner_hparams.batch_size, video_writer=video_writer,
         env_type=planner_hparams.env_type, uct_const=planner_hparams.uct_const,
+        uct_std_normalization=planner_hparams.uct_std_normalization,
         uniform_first_action=planner_hparams.uniform_first_action
     )
     rl_utils.run_rollouts(
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 0231578e2..a78342a07 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -450,7 +450,7 @@ class PlannerAgent(BatchAgent):
   def __init__(
       self, batch_size, rollout_agent, sim_env, wrapper_fn, num_rollouts,
       planning_horizon, discount_factor=1.0, uct_const=0,
-      uniform_first_action=True, video_writer=None
+      uct_std_normalization=False, uniform_first_action=True, video_writer=None
   ):
     super(PlannerAgent, self).__init__(
         batch_size, rollout_agent.observation_space, rollout_agent.action_space
@@ -463,8 +463,10 @@ def __init__(
     self._discount_factor = discount_factor
     self._planning_horizon = planning_horizon
     self._uct_const = uct_const
+    self._uct_std_normalization = uct_std_normalization
     self._uniform_first_action = uniform_first_action
     self._video_writer = video_writer
+    self._best_mc_values = [[] for _ in range(self.batch_size)]
 
   def act(self, observations, env_state=None):
     def run_batch_from(observation, planner_index, batch_index):
@@ -511,23 +513,39 @@ def run_batches_from(observation, planner_index):
       return {a: (sums[a], counts[a]) for a in sums}
 
     def choose_best_action(observation, planner_index):
-      """Choose the best action."""
+      """Choose the best action, update best Monte Carlo values."""
+      best_mc_values = self._best_mc_values[planner_index]
       action_probs = self._rollout_agent.action_distribution(
           np.array([observation] * self._rollout_agent.batch_size)
       )[0, :]
       sums_and_counts = run_batches_from(observation, planner_index)
 
-      def uct(action):
+      def monte_carlo_value(action):
         (value_sum, count) = sums_and_counts[action]
         if count > 0:
           mean_value = value_sum / count
         else:
           mean_value = -np.inf
-        return mean_value + self._uct_bonus(
-            count, action_probs[action]
-        )
+        return mean_value
 
-      return max(range(self.action_space.n), key=uct)
+      mc_values = np.array(
+          [monte_carlo_value(action) for action in range(self.action_space.n)]
+      )
+      best_mc_values.append(mc_values.max())
+
+      if self._uct_std_normalization:
+        agg = np.std
+      else:
+        agg = lambda x: np.mean(np.std(x))
+      normalizer = max(agg(best_mc_values[-30:]), 0.001)
+      normalized_mc_values = mc_values / normalizer
+
+      uct_bonuses = np.array(
+          [self._uct_bonus(sums_and_counts[action][1], action_probs[action])
+          for action in range(self.action_space.n)]
+      )
+      values = normalized_mc_values + uct_bonuses
+      return np.argmax(values)
 
     return np.array([
         choose_best_action(observation, i)

From 6edd97afff7fd99a2bff5eb067ac288d751e5ac4 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 22 Jan 2019 12:09:33 -0800
Subject: [PATCH 1561/2720] Make MaxAndSkipEnv take dtype from observation
 space rather than hardcoding to uint8.

PiperOrigin-RevId: 230385550
---
 tensor2tensor/rl/evaluator.py |  9 +++------
 tensor2tensor/rl/gym_utils.py |  5 +++--
 tensor2tensor/rl/rl_utils.py  | 32 +++++++-------------------------
 3 files changed, 13 insertions(+), 33 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 5b76dfc67..11c08d90d 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -92,7 +92,6 @@ def planner_tiny():
       env_type="simulated",
       uct_const=0.0,
       uniform_first_action=True,
-      uct_std_normalization=False,
   )
 
 
@@ -106,7 +105,6 @@ def planner_small():
       env_type="simulated",
       uct_const=0.0,
       uniform_first_action=True,
-      uct_std_normalization=False,
   )
 
 
@@ -141,7 +139,7 @@ def make_agent(
     sim_env_kwargs=None, frame_stack_size=None, planning_horizon=None,
     rollout_agent_type=None, batch_size=None, num_rollouts=None,
     inner_batch_size=None, video_writer=None, env_type=None,
-    uct_const=None, uct_std_normalization=None, uniform_first_action=None
+    uct_const=None, uniform_first_action=None
 ):
   """Factory function for Agents."""
   if batch_size is None:
@@ -162,8 +160,8 @@ def make_agent(
           lambda env: rl_utils.BatchStackWrapper(env, frame_stack_size),
           num_rollouts, planning_horizon,
           discount_factor=policy_hparams.gae_gamma,
-          uct_const=uct_const, uct_std_normalization=uct_std_normalization,
-          uniform_first_action=uniform_first_action, video_writer=video_writer
+          uct_const=uct_const, uniform_first_action=uniform_first_action,
+          video_writer=video_writer
       ),
   }[agent_type]()
 
@@ -188,7 +186,6 @@ def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
         num_rollouts=planner_hparams.num_rollouts,
         inner_batch_size=planner_hparams.batch_size, video_writer=video_writer,
         env_type=planner_hparams.env_type, uct_const=planner_hparams.uct_const,
-        uct_std_normalization=planner_hparams.uct_std_normalization,
         uniform_first_action=planner_hparams.uniform_first_action
     )
     rl_utils.run_rollouts(
diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index cc86db34f..ed0a82fc1 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -30,9 +30,10 @@ class MaxAndSkipEnv(gym.Wrapper):
   def __init__(self, env, skip=4):
     """Return only every `skip`-th frame."""
     gym.Wrapper.__init__(self, env)
+    observation_space = env.observation_space
     # Most recent raw observations (for max pooling across time steps).
-    self._obs_buffer = np.zeros((2,) + env.observation_space.shape,
-                                dtype=np.uint8)
+    self._obs_buffer = np.zeros((2,) + observation_space.shape,
+                                dtype=observation_space.dtype)
     self._skip = skip
 
   def __str__(self):
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index a78342a07..0231578e2 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -450,7 +450,7 @@ class PlannerAgent(BatchAgent):
   def __init__(
       self, batch_size, rollout_agent, sim_env, wrapper_fn, num_rollouts,
       planning_horizon, discount_factor=1.0, uct_const=0,
-      uct_std_normalization=False, uniform_first_action=True, video_writer=None
+      uniform_first_action=True, video_writer=None
   ):
     super(PlannerAgent, self).__init__(
         batch_size, rollout_agent.observation_space, rollout_agent.action_space
@@ -463,10 +463,8 @@ def __init__(
     self._discount_factor = discount_factor
     self._planning_horizon = planning_horizon
     self._uct_const = uct_const
-    self._uct_std_normalization = uct_std_normalization
     self._uniform_first_action = uniform_first_action
     self._video_writer = video_writer
-    self._best_mc_values = [[] for _ in range(self.batch_size)]
 
   def act(self, observations, env_state=None):
     def run_batch_from(observation, planner_index, batch_index):
@@ -513,39 +511,23 @@ def run_batches_from(observation, planner_index):
       return {a: (sums[a], counts[a]) for a in sums}
 
     def choose_best_action(observation, planner_index):
-      """Choose the best action, update best Monte Carlo values."""
-      best_mc_values = self._best_mc_values[planner_index]
+      """Choose the best action."""
       action_probs = self._rollout_agent.action_distribution(
           np.array([observation] * self._rollout_agent.batch_size)
       )[0, :]
       sums_and_counts = run_batches_from(observation, planner_index)
 
-      def monte_carlo_value(action):
+      def uct(action):
         (value_sum, count) = sums_and_counts[action]
         if count > 0:
           mean_value = value_sum / count
         else:
           mean_value = -np.inf
-        return mean_value
+        return mean_value + self._uct_bonus(
+            count, action_probs[action]
+        )
 
-      mc_values = np.array(
-          [monte_carlo_value(action) for action in range(self.action_space.n)]
-      )
-      best_mc_values.append(mc_values.max())
-
-      if self._uct_std_normalization:
-        agg = np.std
-      else:
-        agg = lambda x: np.mean(np.std(x))
-      normalizer = max(agg(best_mc_values[-30:]), 0.001)
-      normalized_mc_values = mc_values / normalizer
-
-      uct_bonuses = np.array(
-          [self._uct_bonus(sums_and_counts[action][1], action_probs[action])
-          for action in range(self.action_space.n)]
-      )
-      values = normalized_mc_values + uct_bonuses
-      return np.argmax(values)
+      return max(range(self.action_space.n), key=uct)
 
     return np.array([
         choose_best_action(observation, i)

From 0d9b19d68d58f177779b38c9a7cef538276211e3 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Tue, 22 Jan 2019 12:10:21 -0800
Subject: [PATCH 1562/2720] internal merge of PR #1396

PiperOrigin-RevId: 230385699
---
 tensor2tensor/rl/evaluator.py |  9 ++++++---
 tensor2tensor/rl/rl_utils.py  | 32 +++++++++++++++++++++++++-------
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 11c08d90d..5b76dfc67 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -92,6 +92,7 @@ def planner_tiny():
       env_type="simulated",
       uct_const=0.0,
       uniform_first_action=True,
+      uct_std_normalization=False,
   )
 
 
@@ -105,6 +106,7 @@ def planner_small():
       env_type="simulated",
       uct_const=0.0,
       uniform_first_action=True,
+      uct_std_normalization=False,
   )
 
 
@@ -139,7 +141,7 @@ def make_agent(
     sim_env_kwargs=None, frame_stack_size=None, planning_horizon=None,
     rollout_agent_type=None, batch_size=None, num_rollouts=None,
     inner_batch_size=None, video_writer=None, env_type=None,
-    uct_const=None, uniform_first_action=None
+    uct_const=None, uct_std_normalization=None, uniform_first_action=None
 ):
   """Factory function for Agents."""
   if batch_size is None:
@@ -160,8 +162,8 @@ def make_agent(
           lambda env: rl_utils.BatchStackWrapper(env, frame_stack_size),
           num_rollouts, planning_horizon,
           discount_factor=policy_hparams.gae_gamma,
-          uct_const=uct_const, uniform_first_action=uniform_first_action,
-          video_writer=video_writer
+          uct_const=uct_const, uct_std_normalization=uct_std_normalization,
+          uniform_first_action=uniform_first_action, video_writer=video_writer
       ),
   }[agent_type]()
 
@@ -186,6 +188,7 @@ def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
         num_rollouts=planner_hparams.num_rollouts,
         inner_batch_size=planner_hparams.batch_size, video_writer=video_writer,
         env_type=planner_hparams.env_type, uct_const=planner_hparams.uct_const,
+        uct_std_normalization=planner_hparams.uct_std_normalization,
         uniform_first_action=planner_hparams.uniform_first_action
     )
     rl_utils.run_rollouts(
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 0231578e2..f5b2ecdbb 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -450,7 +450,7 @@ class PlannerAgent(BatchAgent):
   def __init__(
       self, batch_size, rollout_agent, sim_env, wrapper_fn, num_rollouts,
       planning_horizon, discount_factor=1.0, uct_const=0,
-      uniform_first_action=True, video_writer=None
+      uct_std_normalization=False, uniform_first_action=True, video_writer=None
   ):
     super(PlannerAgent, self).__init__(
         batch_size, rollout_agent.observation_space, rollout_agent.action_space
@@ -463,8 +463,10 @@ def __init__(
     self._discount_factor = discount_factor
     self._planning_horizon = planning_horizon
     self._uct_const = uct_const
+    self._uct_std_normalization = uct_std_normalization
     self._uniform_first_action = uniform_first_action
     self._video_writer = video_writer
+    self._best_mc_values = [[] for _ in range(self.batch_size)]
 
   def act(self, observations, env_state=None):
     def run_batch_from(observation, planner_index, batch_index):
@@ -511,23 +513,39 @@ def run_batches_from(observation, planner_index):
       return {a: (sums[a], counts[a]) for a in sums}
 
     def choose_best_action(observation, planner_index):
-      """Choose the best action."""
+      """Choose the best action, update best Monte Carlo values."""
+      best_mc_values = self._best_mc_values[planner_index]
       action_probs = self._rollout_agent.action_distribution(
           np.array([observation] * self._rollout_agent.batch_size)
       )[0, :]
       sums_and_counts = run_batches_from(observation, planner_index)
 
-      def uct(action):
+      def monte_carlo_value(action):
         (value_sum, count) = sums_and_counts[action]
         if count > 0:
           mean_value = value_sum / count
         else:
           mean_value = -np.inf
-        return mean_value + self._uct_bonus(
-            count, action_probs[action]
-        )
+        return mean_value
 
-      return max(range(self.action_space.n), key=uct)
+      mc_values = np.array(
+          [monte_carlo_value(action) for action in range(self.action_space.n)]
+      )
+      best_mc_values.append(mc_values.max())
+
+      if self._uct_std_normalization:
+        agg = np.std
+      else:
+        agg = lambda x: np.mean(np.std(x))
+      normalizer = max(agg(best_mc_values[-30:]), 0.001)
+      normalized_mc_values = mc_values / normalizer
+
+      uct_bonuses = np.array(
+          [self._uct_bonus(sums_and_counts[action][1], action_probs[action])
+           for action in range(self.action_space.n)]
+      )
+      values = normalized_mc_values + uct_bonuses
+      return np.argmax(values)
 
     return np.array([
         choose_best_action(observation, i)

From 00985bd34b1511dc63ae90689341a3d869a85228 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Tue, 22 Jan 2019 23:06:53 +0100
Subject: [PATCH 1563/2720] Model-Based RL: Planner hparams, first guess
 (#1397)

* MC rollout values normalization for Planner. (when ensembled with policy)

* Planner hparams with normalization, first guess.

* Fix paramters, add uct_std_normalization.
---
 tensor2tensor/rl/evaluator.py | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 5b76dfc67..dbbd9f508 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -111,15 +111,42 @@ def planner_small():
 
 
 @registry.register_hparams
-def planner_base():
+def planner_guess1():
   return tf.contrib.training.HParams(
       num_rollouts=96,
       batch_size=96,
       planning_horizon=8,
       rollout_agent_type="policy",
       env_type="simulated",
-      uct_const=0.0,
+      uct_const=0.,
+      uniform_first_action=False,
+      uct_std_normalization=False,
+  )
+
+@registry.register_hparams
+def planner_guess3():
+  return tf.contrib.training.HParams(
+      num_rollouts=96,
+      batch_size=96,
+      planning_horizon=8,
+      rollout_agent_type="policy",
+      env_type="simulated",
+      uct_const=0.2,
       uniform_first_action=True,
+      uct_std_normalization=True,
+  )
+
+@registry.register_hparams
+def planner_guess3():
+  return tf.contrib.training.HParams(
+      num_rollouts=96,
+      batch_size=96,
+      planning_horizon=8,
+      rollout_agent_type="policy",
+      env_type="simulated",
+      uct_const=0.5,
+      uniform_first_action=False,
+      uct_std_normalization=False,
   )
 
 

From 7762a9f1f573d5a904de94d444a80f21b3e820c6 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Tue, 22 Jan 2019 14:07:05 -0800
Subject: [PATCH 1564/2720] internal merge of PR #1397

PiperOrigin-RevId: 230406874
---
 tensor2tensor/rl/evaluator.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index dbbd9f508..6227bdf31 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -123,8 +123,9 @@ def planner_guess1():
       uct_std_normalization=False,
   )
 
+
 @registry.register_hparams
-def planner_guess3():
+def planner_guess2():
   return tf.contrib.training.HParams(
       num_rollouts=96,
       batch_size=96,
@@ -136,6 +137,7 @@ def planner_guess3():
       uct_std_normalization=True,
   )
 
+
 @registry.register_hparams
 def planner_guess3():
   return tf.contrib.training.HParams(

From 890efcc81a4401cc39f4ec641c44f6a5d9f3e9f6 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 22 Jan 2019 14:34:36 -0800
Subject: [PATCH 1565/2720] Add TicTacToe env to our model free setup and its
 hparams.

PiperOrigin-RevId: 230412160
---
 tensor2tensor/models/research/rl.py           | 50 ++++++++++++++++++-
 tensor2tensor/rl/rl_utils.py                  | 40 ++++++++++-----
 .../rl/trainer_model_based_params.py          |  5 ++
 tensor2tensor/rl/trainer_model_free.py        |  3 +-
 .../rl/trainer_model_free_tictactoe_test.py   | 47 +++++++++++++++++
 5 files changed, 131 insertions(+), 14 deletions(-)
 create mode 100644 tensor2tensor/rl/trainer_model_free_tictactoe_test.py

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 571cdba09..2cacc1050 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -23,6 +23,7 @@
 
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
+from tensor2tensor.envs import tic_tac_toe_env
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
@@ -69,6 +70,8 @@ def ppo_base_v1():
   hparams.add_hparam("logits_clip", 0.0)
   hparams.add_hparam("dropout_ppo", 0.1)
   hparams.add_hparam("effective_num_agents", None)
+  # TODO(afrozm): Clean this up, this is used in PPO learner to get modalities.
+  hparams.add_hparam("policy_problem_name", "dummy_policy_problem")
   return hparams
 
 
@@ -130,6 +133,15 @@ def ppo_original_params():
   return hparams
 
 
+@registry.register_hparams
+def ppo_ttt_params():
+  """Parameters based on the original PPO paper."""
+  hparams = ppo_original_params()
+  hparams.policy_network = "feed_forward_categorical_policy"
+  hparams.policy_problem_name = "dummy_policy_problem_ttt"
+  return hparams
+
+
 @registry.register_hparams
 def ppo_original_params_gamma95():
   """Parameters based on the original PPO paper, changed gamma."""
@@ -271,7 +283,16 @@ def get_policy(observations, hparams, action_space):
 
   obs_shape = common_layers.shape_list(observations)
   (frame_height, frame_width) = obs_shape[2:4]
-  policy_problem = DummyPolicyProblem(action_space, frame_height, frame_width)
+
+  # TODO(afrozm): We have these dummy problems mainly for hparams, so cleanup
+  # when possible and do this properly.
+  if hparams.policy_problem_name == "dummy_policy_problem_ttt":
+    tf.logging.info("Using DummyPolicyProblemTTT for the policy.")
+    policy_problem = tic_tac_toe_env.DummyPolicyProblemTTT()
+  else:
+    tf.logging.info("Using DummyPolicyProblem for the policy.")
+    policy_problem = DummyPolicyProblem(action_space, frame_height, frame_width)
+
   trainer_lib.add_problem_hparams(hparams, policy_problem)
   hparams.force_full_predict = True
   model = registry.model(hparams.policy_network)(
@@ -365,9 +386,35 @@ def rlmf_original():
       resize_width_factor=2,
       grayscale=0,
       rl_env_max_episode_steps=-1,
+      # If set, use this as the gym env name, instead of changing game mode etc.
+      rl_env_name="",
+      # Controls whether we should derive observation space, do some
+      # pre-processing etc. See T2TGymEnv._derive_observation_space.
+      rl_should_derive_observation_space=True,
   )
 
 
+@registry.register_hparams
+def rlmf_tictactoe():
+  """Base set of hparams for model-free PPO."""
+  hparams = rlmf_original()
+  hparams.game = "tictactoe"
+  hparams.rl_env_name = "T2TEnv-TicTacToeEnv-v0"
+  # Since we don't have any no-op actions, otherwise we have to have an
+  # attribute called `get_action_meanings`.
+  hparams.eval_max_num_noops = 0
+  hparams.add_hparam("max_num_noops", 0)
+  hparams.rl_should_derive_observation_space = False
+
+  hparams.policy_network = "feed_forward_categorical_policy"
+  hparams.base_algo_params = "ppo_ttt_params"
+
+  # Number of last observations to feed to the agent
+  hparams.frame_stack_size = 1
+
+  return hparams
+
+
 @registry.register_hparams
 def rlmf_base():
   """Base set of hparams for model-free PPO."""
@@ -517,6 +564,7 @@ class FeedForwardCategoricalPolicy(PolicyBase):
 
   def body(self, features):
     observations = features["inputs_raw"]
+    observations = tf.cast(observations, tf.float32)
     flat_observations = tf.layers.flatten(observations)
     with tf.variable_scope("policy"):
       x = flat_observations
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index f5b2ecdbb..d145c4fb5 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -84,8 +84,8 @@ def evaluate_single_config(
   eval_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
   env = setup_env(
       hparams, batch_size=hparams.eval_batch_size, max_num_noops=max_num_noops,
-      rl_env_max_episode_steps=hparams.eval_rl_env_max_episode_steps
-  )
+      rl_env_max_episode_steps=hparams.eval_rl_env_max_episode_steps,
+      env_name=hparams.rl_env_name)
   env.start_new_epoch(0)
   eval_fn(env, hparams, eval_hparams, agent_model_dir, sampling_temp)
   rollouts = env.current_epoch_rollouts()
@@ -148,17 +148,33 @@ def full_game_name(short_name):
   return full_name
 
 
-def setup_env(hparams, batch_size, max_num_noops, rl_env_max_episode_steps=-1):
+def should_apply_max_and_skip_env(hparams):
+  """MaxAndSkipEnv doesn't make sense for some games, so omit it if needed."""
+  return hparams.game != "tictactoe"
+
+
+def setup_env(hparams,
+              batch_size,
+              max_num_noops,
+              rl_env_max_episode_steps=-1,
+              env_name=None):
   """Setup."""
-  env_name = full_game_name(hparams.game)
-
-  env = T2TGymEnv(base_env_name=env_name,
-                  batch_size=batch_size,
-                  grayscale=hparams.grayscale,
-                  resize_width_factor=hparams.resize_width_factor,
-                  resize_height_factor=hparams.resize_height_factor,
-                  rl_env_max_episode_steps=rl_env_max_episode_steps,
-                  max_num_noops=max_num_noops, maxskip_envs=True)
+  if not env_name:
+    env_name = full_game_name(hparams.game)
+
+  maxskip_envs = should_apply_max_and_skip_env(hparams)
+
+  env = T2TGymEnv(
+      base_env_name=env_name,
+      batch_size=batch_size,
+      grayscale=hparams.grayscale,
+      should_derive_observation_space=hparams
+      .rl_should_derive_observation_space,
+      resize_width_factor=hparams.resize_width_factor,
+      resize_height_factor=hparams.resize_height_factor,
+      rl_env_max_episode_steps=rl_env_max_episode_steps,
+      max_num_noops=max_num_noops,
+      maxskip_envs=maxskip_envs)
   return env
 
 
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 2ac8d5e9b..86f60e3a7 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -88,6 +88,11 @@ def _rlmb_base():
       eval_rl_env_max_episode_steps=1000,
 
       game="pong",
+      # If set, use this as the gym env name, instead of changing game mode etc.
+      rl_env_name="",
+      # Controls whether we should derive observation space, do some
+      # pre-processing etc. See T2TGymEnv._derive_observation_space.
+      rl_should_derive_observation_space=True,
       # Whether to evaluate the world model in each iteration of the loop to get
       # the model_reward_accuracy metric.
       eval_world_model=True,
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index eb03bd41a..3607f96f8 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -55,7 +55,8 @@ def initialize_env_specs(hparams):
   """Initializes env_specs using T2TGymEnvs."""
   env = rl_utils.setup_env(hparams, hparams.batch_size,
                            hparams.eval_max_num_noops,
-                           hparams.rl_env_max_episode_steps)
+                           hparams.rl_env_max_episode_steps,
+                           env_name=hparams.rl_env_name)
   env.start_new_epoch(0)
 
   # TODO(afrozm): Decouple env_fn from hparams and return both, is there
diff --git a/tensor2tensor/rl/trainer_model_free_tictactoe_test.py b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
new file mode 100644
index 000000000..24d734734
--- /dev/null
+++ b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
@@ -0,0 +1,47 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 of basic flow of collecting trajectories and training PPO."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.rl import trainer_model_free
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+FLAGS = tf.flags.FLAGS
+
+
+class TrainerModelFreeTicTacToeTest(tf.test.TestCase):
+
+  def test_train_tictactoe(self):
+    hparams = registry.hparams("rlmf_tictactoe")
+    hparams.batch_size = 2
+    hparams.eval_sampling_temps = [0.0, 1.0]
+    hparams.add_hparam("ppo_epochs_num", 2)
+    hparams.add_hparam("ppo_epoch_length", 3)
+
+    hparams.epochs_num = 100
+    hparams.eval_every_epochs = 25
+
+    FLAGS.output_dir = tf.test.get_temp_dir()
+    trainer_model_free.train(hparams, FLAGS.output_dir)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 363a4eb65247d9ff6d020ae7355c28136afbd7a9 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Tue, 22 Jan 2019 23:39:19 +0100
Subject: [PATCH 1566/2720] Increase uct_const for Planner. (#1400)

---
 tensor2tensor/rl/evaluator.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 6227bdf31..9ffe06db7 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -132,7 +132,7 @@ def planner_guess2():
       planning_horizon=8,
       rollout_agent_type="policy",
       env_type="simulated",
-      uct_const=0.2,
+      uct_const=3.,
       uniform_first_action=True,
       uct_std_normalization=True,
   )
@@ -146,7 +146,7 @@ def planner_guess3():
       planning_horizon=8,
       rollout_agent_type="policy",
       env_type="simulated",
-      uct_const=0.5,
+      uct_const=2.,
       uniform_first_action=False,
       uct_std_normalization=False,
   )

From 57b3d59c14073dd3eacc9fc35927be15ad5cda20 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 22 Jan 2019 14:52:51 -0800
Subject: [PATCH 1567/2720] Set default step limit for RL evaluator and allow
 to set eval batch size.

PiperOrigin-RevId: 230415920
---
 tensor2tensor/rl/evaluator.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 9ffe06db7..be6b044de 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -54,12 +54,14 @@
 flags.DEFINE_string(
     "eval_metrics_dir", "", "Directory to output the eval metrics at."
 )
-flags.DEFINE_bool("full_eval", True, "Whether to ignore the timestep limit.")
+flags.DEFINE_integer("eval_batch_size", 64, "Number of games to evaluate.")
+flags.DEFINE_integer("eval_step_limit", 100000,
+                     "Maximum number of time steps, ignored if -1.")
 flags.DEFINE_enum(
     "agent", "policy", ["random", "policy", "planner"], "Agent type to use."
 )
 flags.DEFINE_bool(
-    "eval_with_learner", True,
+    "eval_with_learner", False,
     "Whether to use the PolicyLearner.evaluate function instead of an "
     "out-of-graph one. Works only with --agent=policy."
 )
@@ -298,8 +300,8 @@ def main(_):
     loop_hparams.game = get_game_for_worker(
         FLAGS.worker_to_game_map, FLAGS.worker_id + 1)
     tf.logging.info("Set game to %s." % loop_hparams.game)
-  if FLAGS.full_eval:
-    loop_hparams.eval_rl_env_max_episode_steps = -1
+  loop_hparams.eval_rl_env_max_episode_steps = FLAGS.eval_step_limit
+  loop_hparams.eval_batch_size = FLAGS.eval_batch_size
   planner_hparams = trainer_lib.create_hparams(
       FLAGS.planner_hparams_set, FLAGS.planner_hparams
   )
@@ -312,7 +314,10 @@ def main(_):
       cur_dir = os.path.join(cur_dir, "%d" % (FLAGS.worker_id + 1))
     policy_dir = os.path.join(cur_dir, "policy")
     model_dir = os.path.join(cur_dir, "world_model")
-    eval_metrics_dir = os.path.join(cur_dir, "evaluator_" + now_tag)
+    eval_dir_basename = "evaluator_"
+    if FLAGS.agent == "planner":
+      eval_dir_basename = "planner_"
+    eval_metrics_dir = os.path.join(cur_dir, eval_dir_basename + now_tag)
     tf.logging.info("Writing metrics to %s." % eval_metrics_dir)
     if not tf.gfile.Exists(eval_metrics_dir):
       tf.gfile.MkDir(eval_metrics_dir)

From ee64f6f29884aa66c98130afaf7e4eb182e0ca1f Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 22 Jan 2019 15:27:04 -0800
Subject: [PATCH 1568/2720] Add a generic create_registry method

PiperOrigin-RevId: 230422498
---
 tensor2tensor/utils/registry.py      | 64 ++++++++++++++++++++++++++++
 tensor2tensor/utils/registry_test.py | 29 ++++++++++++-
 2 files changed, 92 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index b8c3e5991..fb0f242b2 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -44,6 +44,8 @@ class MyModel(T2TModel):
 from __future__ import division
 from __future__ import print_function
 
+import collections
+
 from tensor2tensor.utils import misc_utils
 import tensorflow as tf
 from tensorflow.python.util import tf_inspect as inspect
@@ -57,6 +59,68 @@ class MyModel(T2TModel):
 _PRUNING_STRATEGY = {}
 _RANGED_HPARAMS = {}
 
+# Key: registry name, Value: Registry
+_GENERIC_REGISTRIES = {}
+Registry = collections.namedtuple(
+    "_Registry", ["register", "get", "list", "registry"])
+
+
+def registry(registry_name):
+  """Returns `Registry` created by `create_registry`."""
+  if registry_name not in _GENERIC_REGISTRIES:
+    raise KeyError("No registry named %s. Available:\n%s" % (
+        registry_name, sorted(_GENERIC_REGISTRIES)))
+  return _GENERIC_REGISTRIES[registry_name]
+
+
+def create_registry(registry_name):
+  """Create a generic object registry.
+
+  Args:
+    registry_name: str, name of the object registry.
+
+  Returns:
+    `Registry` that contains functions for register (decorator), get, and list.
+
+  Raises:
+    KeyError: if `registry_name` is a pre-existing registry.
+  """
+  if registry_name in _GENERIC_REGISTRIES:
+    raise KeyError(
+        "Registry %s already exists." % registry_name)
+
+  registry_ = {}
+
+  def register(name):
+    """Returns decorator to register an object."""
+
+    def register_dec(obj):
+      if name in registry_:
+        raise KeyError(
+            "Registry %s already contains key %s." % (registry_name, name))
+      registry_[name] = obj
+      return obj
+
+    return register_dec
+
+  def get(name):
+    if name not in registry_:
+      raise KeyError(
+          "Registry %s contains no object named %s" % (registry_name, name))
+    return registry_[name]
+
+  def list_registry():
+    return sorted(registry_)
+
+  registry_obj = Registry(
+      register=register,
+      get=get,
+      list=list_registry,
+      registry=registry_,
+  )
+  _GENERIC_REGISTRIES[registry_name] = registry_obj
+  return registry_obj
+
 
 def _reset():
   for ctr in [_MODELS, _HPARAMS, _RANGED_HPARAMS, _ATTACK_PARAMS]:
diff --git a/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index b56c9c48c..c2515b3ba 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -194,8 +194,35 @@ def rhp_bad2(a, b):  # pylint: disable=unused-argument
         pass
 
 
+class CreateRegistry(tf.test.TestCase):
+  """Test class for `create_registry`."""
+
+  def testCreateRegistry(self):
+    my_registry = registry.create_registry("test_reg1")
+    self.assertIs(my_registry, registry.registry("test_reg1"))
+
+    # Use as decorator on a fn
+    @my_registry.register("foo")
+    def some_fn(num):
+      return num + 2
+
+    # Register a regular object
+    pod_obj = 4
+    my_registry.register("bar")(pod_obj)
+
+    # Register a class
+    @my_registry.register("foobar")
+    class A(object):
+      pass
+
+    self.assertEqual(9, my_registry.get("foo")(7))
+    self.assertEqual(["bar", "foo", "foobar"], my_registry.list())
+    foobar = my_registry.get("foobar")
+    self.assertTrue(isinstance(foobar(), A))
+
+
 class RegistryTest(tf.test.TestCase):
-  """ Test class for common functions."""
+  """Test class for common functions."""
 
   def testRegistryHelp(self):
     help_str = registry.help_string()

From a28a1afda5169aee3f02a7d4d4c4f4ac0641f3b7 Mon Sep 17 00:00:00 2001
From: Keyon Vafa <vafa@google.com>
Date: Tue, 22 Jan 2019 16:00:51 -0800
Subject: [PATCH 1569/2720] Add common_image_transformer test to ensure proper
 decoder masking with global attention.

PiperOrigin-RevId: 230428680
---
 .../layers/common_image_attention_test.py     | 46 +++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index 088725b18..1a3791cc8 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 from absl.testing import parameterized
+from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_image_attention
 
 import tensorflow as tf
@@ -105,5 +106,50 @@ def testCreateOutputTrainMode(self, likelihood, num_mixtures, depth):
     else:
       self.assertEqual(output.shape, (batch, height, width, depth))
 
+  def testTransformerDecoderLayersGlobal(self):
+    one_hot_data = tf.constant([[[0., 1.], [1., 0.]],
+                                [[0., 1.], [1., 0.]],
+                                [[1., 0.], [1., 0.]]])
+
+    hparams = common_hparams.basic_params1()
+    hparams.hidden_size = 4
+    hparams.num_layers = 1
+    hparams.layer_prepostprocess_dropout = 0.
+
+    hparams.add_hparam("attention_key_channels", None)
+    hparams.add_hparam("attention_value_channels", None)
+    hparams.add_hparam("num_heads", 1)
+    hparams.add_hparam("attention_dropout", 0.)
+    hparams.add_hparam("shared_rel", False)
+    hparams.add_hparam("block_width", 1)
+    hparams.add_hparam("block_length", 1)
+    hparams.add_hparam("q_filter_width", 1)
+    hparams.add_hparam("kv_filter_width", 1)
+    hparams.add_hparam("filter_size", 16)
+    hparams.add_hparam("ffn_layer", "conv_hidden_relu")
+    hparams.add_hparam("relu_dropout", 0.)
+
+    conv_1d = tf.keras.layers.Conv1D(filters=hparams.hidden_size,
+                                     kernel_size=1,
+                                     use_bias=False)
+    shifted_data = tf.pad(one_hot_data, [[0, 0], [1, 0], [0, 0]])[..., :-1, :]
+    net = conv_1d(shifted_data)
+    output = common_image_attention.transformer_decoder_layers(
+        inputs=net,
+        encoder_output=None,
+        num_layers=hparams.num_layers,
+        hparams=hparams,
+        self_attention_bias=common_image_attention.get_self_attention_bias(net),
+        attention_type=common_image_attention.AttentionType.GLOBAL)
+    self.evaluate(tf.global_variables_initializer())
+    output_val = self.evaluate(output)
+    # The outputs for the padded dimension should be equal across all data.
+    self.assertAllEqual(output_val[0, 0], output_val[1, 0])
+    self.assertAllEqual(output_val[1, 0], output_val[2, 0])
+    # The first and second elements of the batch are identical, so they should
+    # have the same outputs for the second latent dimension as well.
+    self.assertAllEqual(output_val[0, 1], output_val[1, 1])
+
+
 if __name__ == "__main__":
   tf.test.main()

From 301db5d85c2eceb4978a0e6f272c52fb87aeb927 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 23 Jan 2019 17:45:46 -0800
Subject: [PATCH 1570/2720] Remove unintended log message and a few other
 cleanups.

PiperOrigin-RevId: 230637260
---
 tensor2tensor/envs/tic_tac_toe_env.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/tensor2tensor/envs/tic_tac_toe_env.py b/tensor2tensor/envs/tic_tac_toe_env.py
index e79012e3a..4e0f18c00 100644
--- a/tensor2tensor/envs/tic_tac_toe_env.py
+++ b/tensor2tensor/envs/tic_tac_toe_env.py
@@ -32,8 +32,6 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.rl import gym_utils
 
-import tensorflow as tf
-
 
 def encode_pos(i, j):
   """Encodes a pair (i, j) as a scalar position on the board."""
@@ -188,7 +186,6 @@ def step(self, action):
     return self.board_state, reward, self.done, {}
 
   def hparams(self, defaults, unused_model_hparams):
-    tf.logging.error("@@@ new tictactoe hparams are being called!")
     p = defaults
     p.modality = {
         "inputs": modalities.IdentitySymbolModality,
@@ -201,9 +198,6 @@ def hparams(self, defaults, unused_model_hparams):
     p.input_space_id = 0  # problem.SpaceID.GENERIC
     p.target_space_id = 0  # problem.SpaceID.GENERIC
 
-    # TODO(afrozm): This doesn't work without returning the hparams object.
-    return p
-
 
 # TODO(afrozm): Figure out how to get rid of this.
 class DummyPolicyProblemTTT(problem.Problem):

From 70fd3400ba3c45beb3db96fac67b3945f12f3b77 Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Wed, 23 Jan 2019 20:41:50 -0800
Subject: [PATCH 1571/2720] Unmasked relative attention 2d and supporting
 tests. x and y coordinates are independent.

PiperOrigin-RevId: 230653826
---
 tensor2tensor/layers/common_attention.py      | 151 ++++++++++++++++++
 tensor2tensor/layers/common_attention_test.py | 133 +++++++++++++++
 2 files changed, 284 insertions(+)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 2c6cc6d5a..258d47b8a 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2034,6 +2034,157 @@ def dot_product_unmasked_self_attention_relative_v2(
     return ret
 
 
+def _matmul_with_relative_keys_2d(x, y, heads_share_relative_embedding):
+  """Helper function for dot_product_unmasked_self_attention_relative_2d."""
+  if heads_share_relative_embedding:
+    ret = tf.einsum("bhxyd,md->bhxym", x, y)
+  else:
+    ret = tf.einsum("bhxyd,hmd->bhxym", x, y)
+  return ret
+
+
+def dot_product_unmasked_self_attention_relative_2d(
+    q, k, v, bias, max_relative_position=None, dropout_rate=0.0,
+    image_shapes=None, name=None, make_image_summary=True,
+    dropout_broadcast_dims=None, heads_share_relative_embedding=False,
+    add_relative_to_values=False):
+  """Calculate relative position unmasked dot-product self-attention 2d.
+
+
+  The attention calculation is augmented with learned representations for the
+  relative position between each element in q and each element in k and v in
+  height and width dimensions. for query index (i,j) and key index (l, m),
+  the logit is q_i k_j^T + q_i rh_{l-i}^T + q_i rw_{m-j}^T, where rh and ry are
+  the set of relative embeddings in height and width spatial dimensions,
+  respectively.
+
+  Args:
+    q: a Tensor with shape [batch, heads, height, width, depth].
+    k: a Tensor with shape [batch, heads, height, width, depth].
+    v: a Tensor with shape [batch, heads, height, width, depth].
+    bias: bias Tensor.
+    max_relative_position: an integer the max relative embedding considered.
+      Changing this invalidates checkpoints.
+    dropout_rate: a floating point number.
+    image_shapes: optional tuple of integer scalars.
+    name: an optional string.
+    make_image_summary: Whether to make an attention image summary.
+    dropout_broadcast_dims:  an optional list of integers less than 4
+      specifying in which dimensions to broadcast the dropout decisions.
+      saves memory.
+    heads_share_relative_embedding: a boolean indicating wheather to share
+      relative embeddings between attention heads.
+    add_relative_to_values: a boolean for adding relative embeddings to values.
+
+  Returns:
+    [batch, heads, height, width, depth] tensor, the output of attention.
+    height_key_relative_embeddings: a 3d or 2d tensor, depending on head sharing
+      settings, which are the relative embeddings for height.
+    width_key_relative_embeddings: a 3d or 2d tensor, depending on head sharing
+      settings, which are the relative embeddings for width.
+
+  Raises:
+    ValueError: if max_relative_position is not > 0.
+  """
+  if not max_relative_position:
+    raise ValueError("Max relative position (%s) should be > 0 when using "
+                     "relative self attention." % (max_relative_position))
+
+  if add_relative_to_values:
+    raise ValueError("Adding relative embeddings to values is not implemented")
+
+  with tf.variable_scope(
+      name,
+      default_name="dot_product_self_attention_relative_v2",
+      values=[q, k, v]):
+
+    # This calculation only works for self attention.
+    # q, k and v must therefore have the same shape.
+    q.get_shape().assert_is_compatible_with(k.get_shape())
+    q.get_shape().assert_is_compatible_with(v.get_shape())
+
+    (height, width) = (common_layers.shape_list(q)[2],
+                       common_layers.shape_list(q)[3])
+    k_shape = common_layers.shape_list(k)
+    num_heads = k_shape[1]
+    depth_k = k_shape[-1]
+    depth_v = common_layers.shape_list(v)[-1]
+    # flatten height width
+    flatten_hw = lambda x, d: tf.reshape(x, [-1, num_heads, height*width, d])
+    # [batch, num_heads, query_length, memory_length]
+    logits = tf.matmul(flatten_hw(q, depth_k), flatten_hw(k, depth_k),
+                       transpose_b=True)
+
+    def _compute_2d_relative_logits(
+        query, key_relative_embeddings, height, width,
+        heads_share_relative_embedding, transpose_mask):
+      """compute relative logits."""
+      unmasked_rel_logits = _matmul_with_relative_keys_2d(
+          query, key_relative_embeddings, heads_share_relative_embedding)
+      # collapse height and heads
+      unmasked_rel_logits = tf.reshape(unmasked_rel_logits,
+                                       [-1, num_heads*height, width,
+                                        2*width-1])
+      unmasked_rel_logits = (
+          _relative_position_to_absolute_position_unmasked(
+              unmasked_rel_logits))
+      # shape it back for tiling
+      unmasked_rel_logits = tf.reshape(
+          unmasked_rel_logits, [-1, num_heads, height, width, width])
+      # tiling it height times
+      unmasked_rel_logits = tf.expand_dims(
+          unmasked_rel_logits, axis=3)
+      unmasked_rel_logits = tf.tile(unmasked_rel_logits,
+                                    [1, 1, 1, height, 1, 1])
+      # bringing it to the right shape for adding to the logits.
+      unmasked_rel_logits = tf.transpose(unmasked_rel_logits, transpose_mask)
+      unmasked_rel_logits = tf.reshape(unmasked_rel_logits,
+                                       [-1, num_heads, height*width,
+                                        height*width])
+      return unmasked_rel_logits
+
+    # Relative logits in width dimension first.
+    width_key_relative_embeddings = get_relative_embeddings_left_right(
+        max_relative_position, width, depth_k, num_heads,
+        heads_share_relative_embedding,
+        "width_key_relative_embeddings")
+    # [batch, heads, height, 2*width-1, 2*width-1]
+    width_unmasked_rel_logits = _compute_2d_relative_logits(
+        q, width_key_relative_embeddings, height, width,
+        heads_share_relative_embedding, [0, 1, 2, 4, 3, 5])
+    logits += width_unmasked_rel_logits
+    # Relative logits in height dimension next. For ease, we transpose
+    # height and width and repeat the above steps, and transpose to eventually
+    # put the logits in their right positions.
+    # [batch, heads, height, 2*height-1, 2*width-1]
+    height_key_relative_embeddings = get_relative_embeddings_left_right(
+        max_relative_position, height, depth_k, num_heads,
+        heads_share_relative_embedding,
+        "height_key_relative_embeddings")
+
+    height_unmasked_rel_logits = _compute_2d_relative_logits(
+        tf.transpose(q, [0, 1, 3, 2, 4]),
+        height_key_relative_embeddings,
+        width,
+        height,
+        heads_share_relative_embedding, [0, 1, 4, 2, 5, 3])
+    logits += height_unmasked_rel_logits
+    if bias is not None:
+      logits += bias
+    weights = tf.nn.softmax(logits, name="attention_weights")
+    # dropping out the attention links for each of the heads
+    weights = common_layers.dropout_with_broadcast_dims(
+        weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
+    if common_layers.should_generate_summaries() and make_image_summary:
+      attention_image_summary(weights, image_shapes)
+    ret = tf.matmul(weights, flatten_hw(v, depth_v))
+    # reshape back the same spatial dimensions as q
+    return (
+        tf.reshape(ret, [-1, num_heads, height, width, depth_v]),
+        height_key_relative_embeddings,
+        width_key_relative_embeddings)
+
+
 def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
   """Attention to the source and a neighborhood to the left within a block.
 
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index eb96ea2c9..701e2bb39 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -537,6 +537,139 @@ def testDotProductUnMaskedAttentionRelativeV2(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
+  def python_relative_att(self, q, k, v, batch, num_heads, height, width,
+                          depth, height_key_relative_embeddings,
+                          width_key_relative_embeddings,
+                          heads_share_relative_embedding):
+    """Relative attention computation in numpy.
+
+    For query index (i,j) and key index (l, m) the logit is
+    q_i k_j^T + q_i rh_{l-i}^T + q_i rw_{m-j}^T, where rh and ry are the set of
+    relative embeddings in height and width spatial dimensions, respectively.
+
+    Args:
+      q: [batch, heads, height, width, depth] tensor
+      k: [batch, heads, height, width, depth] tensor
+      v: [batch, heads, height, width, depth] tensor
+      batch: int scalar
+      num_heads: int scalar
+      height: int scalar
+      width: int scalar
+      depth: int scalar
+      height_key_relative_embeddings: a tensor of relative embeddings
+      width_key_relative_embeddings: a tensor of relative embeddings
+      heads_share_relative_embedding: a boolean
+
+    Returns:
+      att_output: A tensor
+    """
+
+    logits = np.zeros((batch, num_heads, height*width, height*width))
+    for b in range(batch):
+      for h in range(num_heads):
+        for i in range(height*width):
+          q_col = i%width
+          q_row = int((i-q_col)/width)
+          for j in range(height*width):
+            k_col = j%width
+            k_row = int((j-k_col)/width)
+            logit = np.dot(q[b][h][q_row][q_col], k[b][h][k_row][k_col])
+            width_rel_dist = k_col - q_col
+            width_rel_index = width-1 + width_rel_dist
+            if heads_share_relative_embedding:
+              width_rel_logit = (
+                  np.dot(q[b][h][q_row][q_col],
+                         width_key_relative_embeddings[width_rel_index]))
+            else:
+              width_rel_logit = (
+                  np.dot(q[b][h][q_row][q_col],
+                         width_key_relative_embeddings[h][width_rel_index]))
+            height_rel_dist = k_row - q_row
+            height_rel_index = height-1 + height_rel_dist
+            if heads_share_relative_embedding:
+              height_rel_logit = (
+                  np.dot(q[b][h][q_row][q_col],
+                         height_key_relative_embeddings[height_rel_index]))
+            else:
+              height_rel_logit = (
+                  np.dot(q[b][h][q_row][q_col],
+                         height_key_relative_embeddings[h][height_rel_index]))
+            logits[b, h, i, j] = logit + width_rel_logit + height_rel_logit
+    # now to do a softmax across the logits
+    att = np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True)
+    # comparing the outputs
+    att_output = np.matmul(att,
+                           np.reshape(v, (
+                               batch, num_heads, height*width, depth)))
+    att_output = np.reshape(att_output,
+                            (batch, num_heads, height, width, depth))
+    return att_output
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testDotProductUnMaskedAttentionRelative2d(self):
+    batch = 1
+    height = 3
+    width = 3
+    num_heads = 2
+    max_relative_position = 6
+    depth = 5
+    heads_share_relative_embedding = False
+    q = np.random.rand(batch, num_heads, height, width, depth)
+    k = np.random.rand(batch, num_heads, height, width, depth)
+    v = np.random.rand(batch, num_heads, height, width, depth)
+    a = common_attention.dot_product_unmasked_self_attention_relative_2d(
+        tf.constant(q, dtype=tf.float32),
+        tf.constant(k, dtype=tf.float32),
+        tf.constant(v, dtype=tf.float32),
+        None,
+        max_relative_position=max_relative_position,
+        heads_share_relative_embedding=heads_share_relative_embedding)
+
+    self.evaluate(tf.global_variables_initializer())
+    res, height_key_relative_embeddings, width_key_relative_embeddings = (
+        self.evaluate(a))
+    att_output = self.python_relative_att(
+        q, k, v, batch, num_heads, height, width, depth,
+        height_key_relative_embeddings, width_key_relative_embeddings,
+        heads_share_relative_embedding)
+    self.assertEqual(res.shape, (batch, num_heads, height, width, depth))
+    self.assertAllClose(res, att_output)
+
+  @parameterized.parameters(
+      (1, 10, 12, 2, 6, 3),
+      (1, 1, 12, 2, 6, 3),
+      (2, 10, 1, 2, 6, 3),
+      (1, 10, 12, 2, 1, 1),
+      (1, 10, 12, 2, 2, 8),
+      (4, 10, 12, 2, 12, 10),
+  )
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testDotProductUnMaskedAttentionRelative2dSharedOneRow(
+      self, batch, height, width, num_heads, max_relative_position, depth):
+    heads_share_relative_embedding = True
+    q = np.random.rand(batch, num_heads, height, width, depth)
+    k = np.random.rand(batch, num_heads, height, width, depth)
+    v = np.random.rand(batch, num_heads, height, width, depth)
+
+    a = common_attention.dot_product_unmasked_self_attention_relative_2d(
+        tf.constant(q, dtype=tf.float32),
+        tf.constant(k, dtype=tf.float32),
+        tf.constant(v, dtype=tf.float32),
+        None,
+        max_relative_position=max_relative_position,
+        heads_share_relative_embedding=heads_share_relative_embedding)
+
+    self.evaluate(tf.global_variables_initializer())
+    (res, height_key_relative_embeddings,
+     width_key_relative_embeddings) = self.evaluate(a)
+    att_output = self.python_relative_att(
+        q, k, v, batch, num_heads, height, width, depth,
+        height_key_relative_embeddings, width_key_relative_embeddings,
+        heads_share_relative_embedding)
+    self.assertEqual(res.shape,
+                     (batch, num_heads, height, width, depth))
+    self.assertAllClose(res, att_output)
+
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testRelativeAttentionV2Unmasked(self):
     # (batch, heads, length, depth)

From 599f9a0fd797aa4a6526bb77a00a4050fa1e1fb4 Mon Sep 17 00:00:00 2001
From: Christopher Beitel <cwbeitel@users.noreply.github.com>
Date: Thu, 24 Jan 2019 09:45:00 -0800
Subject: [PATCH 1572/2720] First draft of multi-problem docs (#1399)

* first draft of multi-problem docs

* simplification of tid lookup docs

* update multi-problem inference from ckpt docs

* minor command fixes; sp.

* polish
---
 docs/multi_problem.md | 188 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 188 insertions(+)
 create mode 100644 docs/multi_problem.md

diff --git a/docs/multi_problem.md b/docs/multi_problem.md
new file mode 100644
index 000000000..0d69394d6
--- /dev/null
+++ b/docs/multi_problem.md
@@ -0,0 +1,188 @@
+# Multi-problem training
+
+Multi-problem training is possible by defining [MultiProblem](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/multi_problem.py) sub-classes that specify a list of [Problem](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/problem.py) objects to include in training. In some cases, multi-problem training can be used to improve performance compared to training on individual problems.
+
+In the following sections we'll discuss MultiProblem from a usage perspective followed by that of someone wishing to build upon it.
+
+Please note the [T2T Walkthrough](https://github.com/tensorflow/tensor2tensor/blob/master/docs/walkthrough.md) documentation is a good place to start to understand the variety of component concepts we'll build on here.
+
+## Usage
+
+### Problem definition and datagen
+
+In this discussion we'll consider the following (large) multi-problem that includes ten different sub-problems. These include:
+
+1. A [language modeling](https://en.wikipedia.org/wiki/Language_model) [problem](https://github.com/tensorflow/tensor2tensor/blob/0dff89d64c3406d42717280cb9135a5ce7af793c/tensor2tensor/data_generators/wiki_lm.py#L223) operating on a corpus of German, English, French, and Romanian language wikipedia articles.
+2. Multiple compatible pairwise language translation problems (En -> De, En -> Fr, En -> Ro, De -> En, Fr -> En, Ro -> En)
+3. A compatible [version](https://github.com/tensorflow/tensor2tensor/blob/ef12bee72270b322165d073c39a650a189de39aa/tensor2tensor/data_generators/cnn_dailymail.py#L267) of the combined CNN/DailyMail news article summarization problem.
+4. A compatible [version](https://github.com/tensorflow/tensor2tensor/blob/ef12bee72270b322165d073c39a650a189de39aa/tensor2tensor/data_generators/multinli.py#L155) of the [MultiNLI](https://www.nyu.edu/projects/bowman/multinli/) textual entailment classification problem.
+5. A compatible [version](https://github.com/tensorflow/tensor2tensor/blob/1de13dbebccb415d89b0658e18a57e9607bafd32/tensor2tensor/data_generators/squad.py#L126) of the [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/) question/answer problem.
+
+```python
+
+@registry.register_problem
+class LanguagemodelMultiWikiTranslate(multi_problem.MultiProblem):
+  """Wiki multi-lingual LM and multiple translations."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelMultiWikiTranslate, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelDeEnFrRoWiki64k())
+    self.task_list.append(translate_ende.TranslateEndeWmtMulti64k())
+    self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k())
+    self.task_list.append(translate_enro.TranslateEnroWmtMultiTiny64k())
+    self.task_list.append(translate_ende.TranslateEndeWmtMulti64k(
+        was_reversed=True))
+    self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k(
+        was_reversed=True))
+    self.task_list.append(translate_enro.TranslateEnroWmtMultiTiny64k(
+        was_reversed=True))
+    self.task_list.append(
+        cnn_dailymail.SummarizeCnnDailymailWikiLMMultiVocab64k())
+    self.task_list.append(multinli.MultiNLIWikiLMMultiVocab64k())
+    self.task_list.append(squad.SquadConcatMulti64k())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
+
+```
+
+The word "compatible" was used a lot above! That's because each of these problems have been modified to use the vocabulary produced by the Wikipedia-based language modeling problem, e.g. the following
+
+```python
+@registry.register_problem
+class SummarizeCnnDailymailWikiLMMultiVocab64k(SummarizeCnnDailymail32k):
+  """Summarize CNN and Daily Mail articles using multi-lingual 64k vocab."""
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
+```
+
+**Important note:** It's easy to miss the key point that, as implemented currently, the first task in the task list must be a language modelling problem and each included task must be modified to use the resulting vocabulary.
+
+With a properly defined and registered multi-problem we can now run datagen as follows:
+
+```bash
+
+t2t-datagen --problem=languagemodel_multi_wiki_translate 
+
+```
+
+This will take approximately the following amount of space (and several hours):
+
+```bash
+(t2t) username@instance-2:~$ du -sh /tmp
+99G     /tmp
+(t2t) username@instance-2:~$ du -sh /tmp/t2t_datagen
+81G     /tmp/t2t_datagen
+```
+
+### Training
+
+Next we're ready to try training a model on this MultiProblem. Note that by not specifying `--data_dir` above TFExample's were by default generated into /tmp so that's what we'll explicitly provide here.
+
+```bash
+
+t2t-trainer --problem=languagemodel_multi_wiki_translate \
+    --model=transformer \
+    --hparams_set=transformer_tall_pretrain_lm_tpu_adafactor_large \
+    --output_dir ~/t2t_train/transformer_multi_2jan19 \
+    --data_dir=/tmp \
+    --train_steps=1 \
+    --eval_steps=1
+
+```
+
+The `hparams_set` parameter we provided above was [transformer_tall_pretrain_lm_tpu_adafactor_large](https://github.com/tensorflow/tensor2tensor/blob/08e83030acf3ef13d15ad6eaefaa0a67fb20b59d/tensor2tensor/models/transformer.py#L1721), also provided below:
+
+```python
+
+@registry.register_hparams
+def transformer_tall_pretrain_lm_tpu_adafactor_large():
+  """Hparams for transformer on LM pretraining on TPU, large model."""
+  hparams = transformer_tall_pretrain_lm_tpu_adafactor()
+  hparams.hidden_size = 1024
+  hparams.num_heads = 16
+  hparams.filter_size = 32768  # max fitting in 16G memory is 49152, batch 2
+  hparams.batch_size = 4
+  hparams.multiproblem_mixing_schedule = "constant"
+  # Task order: lm/en-de/en-fr/en-ro/de-en/fr-en/ro-en/cnndm/mnli/squad.
+  hparams.multiproblem_per_task_threshold = "320,80,160,2,80,160,2,20,5,5"
+  return hparams
+
+```
+
+Here it's worth noting a couple things, one that we have specified a `multi_problem_mixing_schedule` (which is required), consumed by [MultiProblem.mix_data](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/multi_problem.py#L280). When set to "constant" the strategy for sampling examples is not a function of step and is proportional only to the per-task "thresholds" which are by default equal (sample examples from each problem with equal probability).
+
+But notice we have also specified the (non-required) `multiproblem_per_task_threshold` parameter, also consumed by mix_data, and specifically used by [sample_task](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/multi_problem.py#L340) which defines non-uniform thresholds to inform a weighted random sampling. E.g. for two problems with weights 1 and 9 the first would be sampled 1/10 of the time and the other 9/10.
+
+### Inference
+
+You can try translating from English to German using a model previously trained on `LanguagemodelMultiWikiTranslate` (the one shown above) ([gs://tensor2tensor-checkpoints/transformer_multi_2jan19/](https://console.cloud.google.com/storage/browser/tensor2tensor-checkpoints/transformer_multi_2jan19/)). Just copy the checkpoint down to a local directory such as the one given via `--output_dir` below:
+
+```bash
+
+t2t-decoder --problem=languagemodel_multi_wiki_translate \
+    --model=transformer \
+    --hparams_set=transformer_tall_pretrain_lm_tpu_adafactor_large \
+    --decode_hparams='batch_size=1,multiproblem_task_id=64510' \
+    --hparams="" \
+    --output_dir=~/t2t_train/transformer_multi_2jan19 \
+    --decode_from_file ~/newstest2014.en \
+    --data_dir=~/t2t_train/transformer_multi_2jan19
+
+```
+
+Here we'll point `--data_dir` to the checkpoint directory which includes the vocab file `vocab.languagemodel_de_en_fr_ro_wiki64k.64000.subwords`; typically data_dir would point to the directory containing your TFRecord example dataset(s).
+
+The file passed to `--decode_from_file` is simply a file with one sentence to translate on each line (in its original form, not post-vocabulary-encoded).
+
+A key requirement for multi-problem inference is that we specify the ID of the problem for which we want to perform inference. But wait, why is the task ID 64510? We can see from the code for [`MultiProblem.update_task_ids`](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/multi_problem.py#L386) that TID's have a place at the end of the vocabulary.
+
+```python
+
+class MultiProblem(problem.Problem):
+  """MultiProblem base class."""
+
+  ...
+
+  def update_task_ids(self, encoder_vocab_size):
+    """Generate task_ids for each problem.
+    These ids correspond to the index of the task in the task_list.
+    Args:
+      encoder_vocab_size: the size of the vocab which is used to compute
+        the index offset.
+    """
+    for idx, task in enumerate(self.task_list):
+      task.set_task_id(idx + encoder_vocab_size)
+      tf.logging.info("Task %d (%s) has id %d." %
+                      (idx, task.name, task.task_id))
+
+```
+
+We can look up the task_id that is assigned to each task we may want to use for inference by instantiating the MultiProblem subclass and obtaining the value, in this case via the following:
+
+```python
+
+task_index = 1 # The second task in the list is En -> De
+LanguagemodelMultiWikiTranslate().task_list[task_index].task_id
+
+```
+
+For me running the `t2t-decode` command provided above gave the following output:
+
+```bash
+...
+
+INFO:tensorflow:Running local_init_op.
+INFO:tensorflow:Done running local_init_op.
+INFO:tensorflow:Inference results INPUT: hello world was the news of the day
+INFO:tensorflow:Inference results OUTPUT: Hallo Welt war die Nachricht des Tages
+INFO:tensorflow:Elapsed Time: 37.15079
+INFO:tensorflow:Averaged Single Token Generation Time: 3.3009222 (time 36.3101439 count 11)
+
+...
+
+```
\ No newline at end of file

From 51c26b6e7185be1404a410d664746c6951f18649 Mon Sep 17 00:00:00 2001
From: Christopher Beitel <cwbeitel@users.noreply.github.com>
Date: Thu, 24 Jan 2019 09:48:48 -0800
Subject: [PATCH 1573/2720] internal merge of PR #1399

PiperOrigin-RevId: 230738813
---
 docs/multi_problem.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/multi_problem.md b/docs/multi_problem.md
index 0d69394d6..d4e37d09d 100644
--- a/docs/multi_problem.md
+++ b/docs/multi_problem.md
@@ -66,7 +66,7 @@ With a properly defined and registered multi-problem we can now run datagen as f
 
 ```bash
 
-t2t-datagen --problem=languagemodel_multi_wiki_translate 
+t2t-datagen --problem=languagemodel_multi_wiki_translate
 
 ```
 
@@ -185,4 +185,4 @@ INFO:tensorflow:Averaged Single Token Generation Time: 3.3009222 (time 36.310143
 
 ...
 
-```
\ No newline at end of file
+```

From f20ae39221656f721c71d400f352cffa5d6bf044 Mon Sep 17 00:00:00 2001
From: He Jinxin <jxhe@mail.ustc.edu.cn>
Date: Fri, 25 Jan 2019 04:46:44 +0800
Subject: [PATCH 1574/2720] sovle the issue 1219. When use the
 translate_enzh_rev problem module, the decoding resulte is wrong. (#1389)

* Update problem.py

Error Querying Server: Requested more than 0 entries, but params is empty. #1219 After days of work, I have solved this problem. It caused by the function serving_input_fn. The input example no need to be reversed any more. So I deleted the code.

* Update problem.py

Error Querying Server: Requested more than 0 entries, but params is empty. #1219
After days of work, I have solved this problem. It caused by the function serving_input_fn. The input example no need to be reversed any more. So I deleted the code.
---
 tensor2tensor/data_generators/problem.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 3a1e32400..ffb474629 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -877,8 +877,7 @@ def serving_input_fn(self, hparams):
         dtype=tf.string, shape=[None], name="serialized_example")
     dataset = tf.data.Dataset.from_tensor_slices(serialized_example)
     dataset = dataset.map(self.decode_example)
-    dataset = dataset.map(lambda ex: self.preprocess_example(ex, mode, hparams))
-    dataset = dataset.map(self.maybe_reverse_and_copy)
+    dataset = dataset.map(lambda ex: self.preprocess_example(ex, mode, hparams))   
     dataset = dataset.map(data_reader.cast_ints_to_int32)
     dataset = dataset.padded_batch(
         tf.shape(serialized_example, out_type=tf.int64)[0],

From 0d283f3ea07dac5f35aad8fe534a87803f53b1ae Mon Sep 17 00:00:00 2001
From: Mostafa Dehghani <dehghani.mostafa@gmail.com>
Date: Thu, 24 Jan 2019 21:49:31 +0100
Subject: [PATCH 1575/2720] rewrite universal_transformer_base based on
 transformer_base, instead of transformer_big (#1405)

---
 .../models/research/universal_transformer.py        | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 59403f73a..5f70c7487 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -438,14 +438,21 @@ def update_hparams_for_universal_transformer(hparams):
 
 @registry.register_hparams
 def universal_transformer_base():
-  hparams = transformer.transformer_big()
+  hparams = transformer.transformer_base()
+  # To have a similar capacity to the transformer_base with 6 layers,
+  # we need to increase the size of the UT's layer
+  # since, in fact, UT has a single layer repeating multiple times.
+  hparams.hidden_size = 1024
+  hparams.filter_size = 4096
+  hparams.num_heads = 16
+  hparams.layer_prepostprocess_dropout = 0.3
   hparams = update_hparams_for_universal_transformer(hparams)
   return hparams
 
 
 @registry.register_hparams
 def universal_transformer_base_tpu():
-  hparams = transformer.transformer_big()
+  hparams = universal_transformer_base()
   hparams = update_hparams_for_universal_transformer(hparams)
   transformer.update_hparams_for_tpu(hparams)
   hparams.add_step_timing_signal = False
@@ -454,7 +461,7 @@ def universal_transformer_base_tpu():
 
 @registry.register_hparams
 def universal_transformer_big():
-  hparams = transformer.transformer_big()
+  hparams = universal_transformer_base()
   hparams = update_hparams_for_universal_transformer(hparams)
   hparams.hidden_size = 2048
   hparams.filter_size = 8192

From 91ada5747529e5e509c38c7c1f2aa786c8560048 Mon Sep 17 00:00:00 2001
From: He Jinxin <jxhe@mail.ustc.edu.cn>
Date: Thu, 24 Jan 2019 13:01:38 -0800
Subject: [PATCH 1576/2720] internal merge of PR #1389

PiperOrigin-RevId: 230774856
---
 tensor2tensor/data_generators/problem.py            |  2 +-
 .../models/research/universal_transformer.py        | 13 +++----------
 2 files changed, 4 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index ffb474629..858701e4e 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -877,7 +877,7 @@ def serving_input_fn(self, hparams):
         dtype=tf.string, shape=[None], name="serialized_example")
     dataset = tf.data.Dataset.from_tensor_slices(serialized_example)
     dataset = dataset.map(self.decode_example)
-    dataset = dataset.map(lambda ex: self.preprocess_example(ex, mode, hparams))   
+    dataset = dataset.map(lambda ex: self.preprocess_example(ex, mode, hparams))
     dataset = dataset.map(data_reader.cast_ints_to_int32)
     dataset = dataset.padded_batch(
         tf.shape(serialized_example, out_type=tf.int64)[0],
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 5f70c7487..59403f73a 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -438,21 +438,14 @@ def update_hparams_for_universal_transformer(hparams):
 
 @registry.register_hparams
 def universal_transformer_base():
-  hparams = transformer.transformer_base()
-  # To have a similar capacity to the transformer_base with 6 layers,
-  # we need to increase the size of the UT's layer
-  # since, in fact, UT has a single layer repeating multiple times.
-  hparams.hidden_size = 1024
-  hparams.filter_size = 4096
-  hparams.num_heads = 16
-  hparams.layer_prepostprocess_dropout = 0.3
+  hparams = transformer.transformer_big()
   hparams = update_hparams_for_universal_transformer(hparams)
   return hparams
 
 
 @registry.register_hparams
 def universal_transformer_base_tpu():
-  hparams = universal_transformer_base()
+  hparams = transformer.transformer_big()
   hparams = update_hparams_for_universal_transformer(hparams)
   transformer.update_hparams_for_tpu(hparams)
   hparams.add_step_timing_signal = False
@@ -461,7 +454,7 @@ def universal_transformer_base_tpu():
 
 @registry.register_hparams
 def universal_transformer_big():
-  hparams = universal_transformer_base()
+  hparams = transformer.transformer_big()
   hparams = update_hparams_for_universal_transformer(hparams)
   hparams.hidden_size = 2048
   hparams.filter_size = 8192

From af22c2408489ab06f115df0a1be46d6df7bfbb3e Mon Sep 17 00:00:00 2001
From: Mostafa Dehghani <dehghani.mostafa@gmail.com>
Date: Thu, 24 Jan 2019 13:09:56 -0800
Subject: [PATCH 1577/2720] internal merge of PR #1405

PiperOrigin-RevId: 230776557
---
 .../models/research/universal_transformer.py       | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 59403f73a..a160e3627 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -438,14 +438,22 @@ def update_hparams_for_universal_transformer(hparams):
 
 @registry.register_hparams
 def universal_transformer_base():
-  hparams = transformer.transformer_big()
+  """Base parameters for Universal Transformer."""
+  hparams = transformer.transformer_base()
+  # To have a similar capacity to the transformer_base with 6 layers,
+  # we need to increase the size of the UT's layer
+  # since, in fact, UT has a single layer repeating multiple times.
+  hparams.hidden_size = 1024
+  hparams.filter_size = 4096
+  hparams.num_heads = 16
+  hparams.layer_prepostprocess_dropout = 0.3
   hparams = update_hparams_for_universal_transformer(hparams)
   return hparams
 
 
 @registry.register_hparams
 def universal_transformer_base_tpu():
-  hparams = transformer.transformer_big()
+  hparams = universal_transformer_base()
   hparams = update_hparams_for_universal_transformer(hparams)
   transformer.update_hparams_for_tpu(hparams)
   hparams.add_step_timing_signal = False
@@ -454,7 +462,7 @@ def universal_transformer_base_tpu():
 
 @registry.register_hparams
 def universal_transformer_big():
-  hparams = transformer.transformer_big()
+  hparams = universal_transformer_base()
   hparams = update_hparams_for_universal_transformer(hparams)
   hparams.hidden_size = 2048
   hparams.filter_size = 8192

From ba31f44729870ac54445df3c846af5ca1793841a Mon Sep 17 00:00:00 2001
From: Dominic Jack <thedomjack@gmail.com>
Date: Fri, 25 Jan 2019 07:22:31 +1000
Subject: [PATCH 1578/2720] added optimizer registry (#1401)

* added optimizer registry

* fixed adafactor -> Adafactor

* fixed default naming

* improved base optimizer registration implementation
---
 tensor2tensor/utils/optimize.py | 126 ++++++++++++++++++++------------
 tensor2tensor/utils/registry.py |  45 ++++++++++++
 2 files changed, 124 insertions(+), 47 deletions(-)

diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 48ae5d349..89e42c4f9 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -24,6 +24,7 @@
 from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import multistep_optimizer
 from tensor2tensor.utils import yellowfin
+from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -93,6 +94,83 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
   return train_op
 
 
+@registry.register_optimizer
+def adam(learning_rate, hparams):
+  # We change the default epsilon for Adam.
+  # Using LazyAdam as it's much faster for large vocabulary embeddings.
+  return tf.contrib.opt.LazyAdamOptimizer(
+      learning_rate,
+      beta1=hparams.optimizer_adam_beta1,
+      beta2=hparams.optimizer_adam_beta2,
+      epsilon=hparams.optimizer_adam_epsilon)
+
+
+@registry.register_optimizer
+def multistep_adam(learning_rate, hparams):
+  return multistep_optimizer.MultistepAdamOptimizer(
+      learning_rate,
+      beta1=hparams.optimizer_adam_beta1,
+      beta2=hparams.optimizer_adam_beta2,
+      epsilon=hparams.optimizer_adam_epsilon,
+      n=hparams.optimizer_multistep_accumulate_steps)
+
+
+@registry.register_optimizer
+def momentum(learning_rate, hparams):
+  return tf.train.MomentumOptimizer(
+      learning_rate,
+      momentum=hparams.optimizer_momentum_momentum,
+      use_nesterov=hparams.optimizer_momentum_nesterov)
+
+
+@registry.register_optimizer
+def yellow_fin(learning_rate, hparams):
+  return yellowfin.YellowFinOptimizer(
+      learning_rate=learning_rate,
+      momentum=hparams.optimizer_momentum_momentum)
+
+
+@registry.register_optimizer
+def true_adam(learning_rate, hparams):
+  return tf.train.AdamOptimizer(
+      learning_rate,
+      beta1=hparams.optimizer_adam_beta1,
+      beta2=hparams.optimizer_adam_beta2,
+      epsilon=hparams.optimizer_adam_epsilon)
+
+
+@registry.register_optimizer
+def adam_w(learning_rate, hparams):
+  # Openai gpt used weight decay.
+  # Given the internals of AdamW, weight decay dependent on the
+  # learning rate is chosen to match the openai implementation.
+  # The weight decay update to each parameter is applied before the adam
+  # gradients computation, which is different from that described
+  # in the paper and in the openai implementation:
+  # https://arxiv.org/pdf/1711.05101.pdf
+  return tf.contrib.opt.AdamWOptimizer(
+      0.01*learning_rate,
+      learning_rate,
+      beta1=hparams.optimizer_adam_beta1,
+      beta2=hparams.optimizer_adam_beta2,
+      epsilon=hparams.optimizer_adam_epsilon)
+
+
+@registry.register_optimizer("Adafactor")
+def register_adafactor(learning_rate, hparams):
+  return adafactor.adafactor_optimizer_from_hparams(hparams, learning_rate)
+
+
+def _register_base_optimizer(key, fn):
+  registry.register_optimizer(key)(
+      lambda learning_rate, hparams: fn(learning_rate))
+
+
+for k in tf.contrib.layers.OPTIMIZER_CLS_NAMES:
+  if k not in registry._OPTIMIZERS:
+    _register_base_optimizer(k, tf.contrib.layers.OPTIMIZER_CLS_NAMES[k])
+
+
 class ConditionalOptimizer(tf.train.Optimizer):
   """Conditional optimizer."""
 
@@ -113,53 +191,7 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
         value=hparams.optimizer_adam_epsilon,
         hparams=hparams)
 
-    if optimizer_name == "Adam":
-      # We change the default epsilon for Adam.
-      # Using LazyAdam as it's much faster for large vocabulary embeddings.
-      self._opt = tf.contrib.opt.LazyAdamOptimizer(
-          lr,
-          beta1=hparams.optimizer_adam_beta1,
-          beta2=hparams.optimizer_adam_beta2,
-          epsilon=hparams.optimizer_adam_epsilon)
-    elif optimizer_name == "MultistepAdam":
-      self._opt = multistep_optimizer.MultistepAdamOptimizer(
-          lr,
-          beta1=hparams.optimizer_adam_beta1,
-          beta2=hparams.optimizer_adam_beta2,
-          epsilon=hparams.optimizer_adam_epsilon,
-          n=hparams.optimizer_multistep_accumulate_steps)
-    elif optimizer_name == "Momentum":
-      self._opt = tf.train.MomentumOptimizer(
-          lr,
-          momentum=hparams.optimizer_momentum_momentum,
-          use_nesterov=hparams.optimizer_momentum_nesterov)
-    elif optimizer_name == "YellowFin":
-      self._opt = yellowfin.YellowFinOptimizer(
-          learning_rate=lr, momentum=hparams.optimizer_momentum_momentum)
-    elif optimizer_name == "TrueAdam":
-      self._opt = tf.train.AdamOptimizer(
-          lr,
-          beta1=hparams.optimizer_adam_beta1,
-          beta2=hparams.optimizer_adam_beta2,
-          epsilon=hparams.optimizer_adam_epsilon)
-    elif optimizer_name == "AdamW":
-      # Openai gpt used weight decay.
-      # Given the internals of AdamW, weight decay dependent on the
-      # learning rate is chosen to match the openai implementation.
-      # The weight decay update to each parameter is applied before the adam
-      # gradients computation, which is different from that described
-      # in the paper and in the openai implementation:
-      # https://arxiv.org/pdf/1711.05101.pdf
-      self._opt = tf.contrib.opt.AdamWOptimizer(
-          0.01*lr,
-          lr,
-          beta1=hparams.optimizer_adam_beta1,
-          beta2=hparams.optimizer_adam_beta2,
-          epsilon=hparams.optimizer_adam_epsilon)
-    elif optimizer_name == "Adafactor":
-      self._opt = adafactor.adafactor_optimizer_from_hparams(hparams, lr)
-    else:
-      self._opt = tf.contrib.layers.OPTIMIZER_CLS_NAMES[optimizer_name](lr)
+    self._opt = registry.optimizer(optimizer_name)(lr, hparams)
     if _mixed_precision_is_enabled(hparams):
       if not hparams.mixed_precision_optimizer_loss_scaler:
         tf.logging.warning("Using mixed precision without a loss scaler will "
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index fb0f242b2..d8683f234 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -183,6 +183,51 @@ def list_models():
   return list(sorted(_MODELS))
 
 
+_OPTIMIZERS = {}
+
+
+def register_optimizer(name=None):
+  """Register an optimizer. name defaults to upper camel case of fn name."""
+
+  def default_opt_name(opt_fn):
+    return misc_utils.snakecase_to_camelcase(default_name(opt_fn))
+
+  def decorator(opt_fn, registration_name):
+    """Registers and returns optimizer_fn with registration_name or default."""
+    if registration_name is None:
+      registration_name = default_opt_name(opt_fn)
+
+    if registration_name in _OPTIMIZERS and not tf.executing_eagerly():
+      raise LookupError("Optimizer %s already registered." % registration_name)
+    args, varargs, keywords, _ = inspect.getargspec(opt_fn)
+
+    if len(args) != 2 or varargs is not None or keywords is not None:
+      raise ValueError("Optimizer registration function must take two "
+                       "arguments: learning_rate (float) and "
+                       "hparams (HParams).")
+    _OPTIMIZERS[registration_name] = opt_fn
+    return opt_fn
+
+  if callable(name):
+    opt_fn = name
+    registration_name = default_opt_name(opt_fn)
+    return decorator(opt_fn, registration_name=registration_name)
+
+  return lambda opt_fn: decorator(opt_fn, name)
+
+
+def optimizer(name):
+  if name not in _OPTIMIZERS:
+    raise LookupError("Optimizer %s never registered. "
+                      "Available optimizers:\n %s"
+                      % (name, "\n".join(list_optimizers())))
+  return _OPTIMIZERS[name]
+
+
+def list_optimizers():
+  return list(sorted(_OPTIMIZERS))
+
+
 def register_hparams(name=None):
   """Register an HParams set. name defaults to function name snake-cased."""
 

From e5db82e9a5f57e8fabe867f157912a9fab924444 Mon Sep 17 00:00:00 2001
From: Dominic Jack <thedomjack@gmail.com>
Date: Thu, 24 Jan 2019 13:22:52 -0800
Subject: [PATCH 1579/2720] internal merge of PR #1401

PiperOrigin-RevId: 230778721
---
 tensor2tensor/utils/optimize.py | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 89e42c4f9..c8f1a9a50 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -23,8 +23,8 @@
 from tensor2tensor.utils import adafactor
 from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import multistep_optimizer
-from tensor2tensor.utils import yellowfin
 from tensor2tensor.utils import registry
+from tensor2tensor.utils import yellowfin
 
 import tensorflow as tf
 
@@ -94,7 +94,7 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
   return train_op
 
 
-@registry.register_optimizer
+@registry.register_optimizer("adam")
 def adam(learning_rate, hparams):
   # We change the default epsilon for Adam.
   # Using LazyAdam as it's much faster for large vocabulary embeddings.
@@ -105,7 +105,7 @@ def adam(learning_rate, hparams):
       epsilon=hparams.optimizer_adam_epsilon)
 
 
-@registry.register_optimizer
+@registry.register_optimizer("multistep_adam")
 def multistep_adam(learning_rate, hparams):
   return multistep_optimizer.MultistepAdamOptimizer(
       learning_rate,
@@ -115,7 +115,7 @@ def multistep_adam(learning_rate, hparams):
       n=hparams.optimizer_multistep_accumulate_steps)
 
 
-@registry.register_optimizer
+@registry.register_optimizer("momentum")
 def momentum(learning_rate, hparams):
   return tf.train.MomentumOptimizer(
       learning_rate,
@@ -123,14 +123,14 @@ def momentum(learning_rate, hparams):
       use_nesterov=hparams.optimizer_momentum_nesterov)
 
 
-@registry.register_optimizer
+@registry.register_optimizer("yellow_fin")
 def yellow_fin(learning_rate, hparams):
   return yellowfin.YellowFinOptimizer(
       learning_rate=learning_rate,
       momentum=hparams.optimizer_momentum_momentum)
 
 
-@registry.register_optimizer
+@registry.register_optimizer("true_adam")
 def true_adam(learning_rate, hparams):
   return tf.train.AdamOptimizer(
       learning_rate,
@@ -139,7 +139,7 @@ def true_adam(learning_rate, hparams):
       epsilon=hparams.optimizer_adam_epsilon)
 
 
-@registry.register_optimizer
+@registry.register_optimizer("adam_w")
 def adam_w(learning_rate, hparams):
   # Openai gpt used weight decay.
   # Given the internals of AdamW, weight decay dependent on the
@@ -161,13 +161,15 @@ def register_adafactor(learning_rate, hparams):
   return adafactor.adafactor_optimizer_from_hparams(hparams, learning_rate)
 
 
+
+
 def _register_base_optimizer(key, fn):
   registry.register_optimizer(key)(
       lambda learning_rate, hparams: fn(learning_rate))
 
 
 for k in tf.contrib.layers.OPTIMIZER_CLS_NAMES:
-  if k not in registry._OPTIMIZERS:
+  if k not in registry._OPTIMIZERS:  # pylint: disable=protected-access
     _register_base_optimizer(k, tf.contrib.layers.OPTIMIZER_CLS_NAMES[k])
 
 

From 0e982b582397429f2a7b56f59deb53b5684f12a5 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Fri, 25 Jan 2019 09:10:00 -0800
Subject: [PATCH 1580/2720] Use registry to avoid sticking unserializable
 objects in the hparams - gets mtf_transformer2 working again.

PiperOrigin-RevId: 230912618
---
 tensor2tensor/models/mtf_transformer2.py | 199 ++++++++++++-----------
 tensor2tensor/utils/decoding.py          |   2 +-
 2 files changed, 108 insertions(+), 93 deletions(-)

diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index 82aa9f9c1..43dafed38 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -111,11 +111,7 @@ def model(self):
       raise NotImplementedError(
           "Label smoothing not implemented in unitransformer."
           "  Do you really want it?")
-    if isinstance(hparams.layer_stack, transformer.LayerStack):
-      layer_stack = hparams.layer_stack
-    else:
-      # hparams.layer_stack is a function for creating a LayerStack
-      layer_stack = hparams.layer_stack(hparams, "")
+    layer_stack = layer_stack_from_hparams(hparams, "")
     if self.autoregressive:
       input_vocab_size = self._targets_vocab_size
     else:
@@ -214,14 +210,8 @@ class MtfBitransformer(MtfUnitransformer):
 
   def model(self):
     hparams = self._hparams
-    if isinstance(hparams.encoder_layer_stack, transformer.LayerStack):
-      encoder_layer_stack = hparams.encoder_layer_stack
-    else:
-      encoder_layer_stack = hparams.encoder_layer_stack(hparams, "encoder_")
-    if isinstance(hparams.decoder_layer_stack, transformer.LayerStack):
-      decoder_layer_stack = hparams.decoder_layer_stack
-    else:
-      decoder_layer_stack = hparams.decoder_layer_stack(hparams, "decoder_")
+    encoder_layer_stack = layer_stack_from_hparams(hparams, "encoder_")
+    decoder_layer_stack = layer_stack_from_hparams(hparams, "decoder_")
     return transformer.Bitransformer(
         encoder_layer_stack=encoder_layer_stack,
         decoder_layer_stack=decoder_layer_stack,
@@ -279,6 +269,9 @@ def sample(self, features, mesh):
     return self.combine_batch_dims(ret)
 
 
+layers_registry = registry.create_registry("layers")
+
+
 # The following functions construct layers based on hyperparmeters
 def attention_kwargs_from_hparams(hparams):
   return {
@@ -287,65 +280,73 @@ def attention_kwargs_from_hparams(hparams):
   }
 
 
-def self_attention_from_hparams(hparams, prefix):
+@layers_registry.register("self_att")
+def self_attention_layer(hparams, prefix):
   """Create self-attention layer based on hyperparameters."""
-  radius = hparams.get(prefix + "local_attention_radius")
-  if radius:
-    return transformer_layers.LocalSelfAttention(
-        num_heads=hparams.get(prefix + "num_heads"),
-        num_memory_heads=hparams.get(prefix + "num_memory_heads", 0),
-        radius=radius,
-        key_value_size=hparams.d_kv,
-        shared_kv=hparams.get(prefix + "shared_kv", False),
-        attention_kwargs=attention_kwargs_from_hparams(hparams))
-  else:
-    return transformer_layers.SelfAttention(
-        num_heads=hparams.get(prefix + "num_heads"),
-        num_memory_heads=hparams.get(prefix + "num_memory_heads", 0),
-        key_value_size=hparams.d_kv,
-        shared_kv=hparams.get(prefix + "shared_kv", False),
-        attention_kwargs=attention_kwargs_from_hparams(hparams))
-
-
-def enc_dec_attention_from_hparams(hparams, prefix):
+  return transformer_layers.SelfAttention(
+      num_heads=hparams.get(prefix + "num_heads"),
+      num_memory_heads=hparams.get(prefix + "num_memory_heads"),
+      key_value_size=hparams.d_kv,
+      shared_kv=hparams.get(prefix + "shared_kv", False),
+      attention_kwargs=attention_kwargs_from_hparams(hparams))
+
+
+@layers_registry.register("local_self_att")
+def local_self_attention_layer(hparams, prefix):
+  """Create self-attention layer based on hyperparameters."""
+  return transformer_layers.LocalSelfAttention(
+      num_heads=hparams.get(prefix + "num_heads"),
+      num_memory_heads=hparams.get(prefix + "num_memory_heads"),
+      radius=hparams.local_attention_radius,
+      key_value_size=hparams.d_kv,
+      shared_kv=hparams.get(prefix + "shared_kv", False),
+      attention_kwargs=attention_kwargs_from_hparams(hparams))
+
+
+@layers_registry.register("enc_att")
+def enc_dec_attention_layer(hparams, prefix):
   return transformer_layers.EncDecAttention(
       num_heads=hparams.get(prefix + "num_heads"),
-      num_memory_heads=hparams.get(prefix + "num_memory_heads", 0),
+      num_memory_heads=hparams.get(prefix + "num_memory_heads"),
       key_value_size=hparams.d_kv,
       shared_kv=hparams.get(prefix + "shared_kv", False),
       attention_kwargs=attention_kwargs_from_hparams(hparams))
 
 
-def dense_relu_dense_from_hparams(hparams):
+@layers_registry.register("drd")
+def dense_relu_dense_layer(hparams, prefix):
+  del prefix
   return transformer_layers.DenseReluDense(
       hidden_size=hparams.d_ff,
       dropout_rate=hparams.relu_dropout)
 
 
-def layer_stack_from_hparams(hparams, prefix):
-  """Create a layer stack based on the hyperparameter values."""
-  return transformer.LayerStack(
-      [self_attention_from_hparams(hparams, prefix),
-       dense_relu_dense_from_hparams(hparams)
-      ] * hparams.get(prefix + "num_layers"),
-      dropout_rate=hparams.layer_prepostprocess_dropout,
-      norm_epsilon=hparams.norm_epsilon)
+@layers_registry.register("moe_1d")
+def moe_1d_layer(hparams, prefix):
+  del prefix
+  return moe.MoE1D(num_experts=hparams.moe_num_experts,
+                   hidden_size=hparams.moe_hidden_size)
+
 
+@layers_registry.register("moe_2d")
+def moe_2d_layer(hparams, prefix):
+  del prefix
+  return moe.MoE2D(expert_x=hparams.moe_expert_x,
+                   expert_y=hparams.moe_expert_y,
+                   hidden_size=hparams.moe_hidden_size)
 
-def decoder_layer_stack_from_hparams(hparams, prefix):
-  if prefix != "decoder_":
-    raise ValueError("prefix should be 'decoder'")
+
+def layer_stack_from_hparams(hparams, prefix):
+  """Create a layer stack based on the hyperparameter values."""
+  layers = hparams.get(prefix + "layers")
   return transformer.LayerStack(
-      [self_attention_from_hparams(hparams, prefix),
-       enc_dec_attention_from_hparams(hparams, prefix),
-       dense_relu_dense_from_hparams(hparams)
-      ] * hparams.get(prefix + "num_layers"),
+      [layers_registry.get(l)(hparams, prefix) for l in layers],
       dropout_rate=hparams.layer_prepostprocess_dropout,
       norm_epsilon=hparams.norm_epsilon)
 
 
 def mtf_transformer2_base():
-  """Set of hyperparameters."""
+  """Hyperparameters common to both unitransformer and bitransformer."""
   hparams = common_hparams.basic_params1()
 
   hparams.add_hparam("d_model", 1024)
@@ -356,9 +357,7 @@ def mtf_transformer2_base():
   # with bfloat16 activations.
   hparams.add_hparam("z_loss", 1e-4)
 
-  # These hyperparameters are used in layer_stack_from_hparams()
-  # They may not be respected if hparams uses a differet layer stack function.
-  hparams.num_hidden_layers = 6
+  # hparams applying to both encoder and decoder layer stacks.
   hparams.add_hparam("d_ff", 2048)
   hparams.add_hparam("d_kv", 128)
   hparams.add_hparam("attention_dropout", 0.0)
@@ -367,6 +366,12 @@ def mtf_transformer2_base():
   hparams.del_hparam("num_hidden_layers")
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.add_hparam("extra_logit", False)
+  # number of experts for moe_1d
+  hparams.moe_num_experts = 32
+  # number of experts for moe_2d = moe_expert_x * moe_expert_y
+  hparams.add_hparam("moe_expert_x", 8)
+  hparams.add_hparam("moe_expert_y", 4)
+  hparams.add_hparam("moe_hidden_size", 32768)
 
   # round up vocab sizes to be a multiple of this value
   hparams.vocab_divisor = 128
@@ -403,6 +408,7 @@ def mtf_transformer2_base():
       "inputs": modalities.IdentitySymbolModality,
       "targets": modalities.IdentitySymbolModality,
   }
+  hparams.add_hparam("beam_size", 1)
   return hparams
 
 
@@ -411,9 +417,8 @@ def mtf_unitransformer_base():
   """Hyperparameters for single-stack Transformer."""
   hparams = mtf_transformer2_base()
   hparams.add_hparam("autoregressive", True)
-  hparams.layer_stack = layer_stack_from_hparams
   # HYPERPARAMETERS FOR THE SINGLE LAYER STACK
-  hparams.add_hparam("num_layers", 6)
+  hparams.add_hparam("layers", ["self_att", "drd"] * 6)
   # number of heads in multihead attention
   hparams.add_hparam("num_heads", 8)
   # default of 0 for standard transformer behavior
@@ -422,7 +427,7 @@ def mtf_unitransformer_base():
   # share attention keys and values
   hparams.add_hparam("shared_kv", False)
   # if nonzero then use local attention
-  hparams.add_hparam("local_attention_radius", 0)
+  hparams.add_hparam("local_attention_radius", 128)
   return hparams
 
 
@@ -432,14 +437,16 @@ def mtf_bitransformer_base():
   hparams = mtf_transformer2_base()
   hparams.max_length = 256
   hparams.shared_embedding = True
-  hparams.encoder_layer_stack = layer_stack_from_hparams
-  hparams.decoder_layer_stack = decoder_layer_stack_from_hparams
   # HYPERPARAMETERS FOR THE LAYER STACKS
+  hparams.add_hparam("encoder_layers", ["self_att", "drd"] * 6)
+  hparams.add_hparam("decoder_layers", ["self_att", "enc_att", "drd"] * 6)
   hparams.add_hparam("encoder_num_layers", 6)
   hparams.add_hparam("decoder_num_layers", 6)
   # number of heads in multihead attention
   hparams.add_hparam("encoder_num_heads", 8)
   hparams.add_hparam("decoder_num_heads", 8)
+  hparams.add_hparam("local_attention_radius", 128)
+
   # default of 0 for standard transformer behavior
   # 1 means a single set of keys and values that are read by all query heads
   hparams.add_hparam("encoder_num_memory_heads", 0)
@@ -447,9 +454,6 @@ def mtf_bitransformer_base():
   # share attention keys and values
   hparams.add_hparam("encoder_shared_kv", False)
   hparams.add_hparam("decoder_shared_kv", False)
-  # if nonzero then use local attention
-  hparams.add_hparam("encoder_local_attention_radius", 0)
-  hparams.add_hparam("decoder_local_attention_radius", 0)
 
   # Parameters for computing the maximum decode length in beam search.
   # Maximum decode length is:
@@ -467,7 +471,7 @@ def mtf_unitransformer_tiny():
   hparams.batch_size = 2
   hparams.mesh_shape = ""
   hparams.d_model = 128
-  hparams.num_hidden_layers = 2
+  hparams.layers = ["self_att", "drd"] * 2
   hparams.num_heads = 4
   hparams.d_ff = 512
   return hparams
@@ -475,11 +479,13 @@ def mtf_unitransformer_tiny():
 
 @registry.register_hparams
 def mtf_bitransformer_tiny():
+  """Small encoder-decoder model for testing."""
   hparams = mtf_bitransformer_base()
   hparams.batch_size = 2
   hparams.mesh_shape = ""
   hparams.d_model = 128
-  hparams.num_hidden_layers = 2
+  hparams.encoder_layers = ["self_att", "drd"] * 2
+  hparams.decoder_layers = ["self_att", "enc_att", "drd"] * 2
   hparams.num_heads = 4
   hparams.d_ff = 512
   return hparams
@@ -489,12 +495,26 @@ def mtf_bitransformer_tiny():
 def mtf_unitransformer_all_layers_tiny():
   """Test out all the layers on local CPU."""
   hparams = mtf_unitransformer_tiny()
-  hparams.layer_stack = transformer.LayerStack(
-      [transformer_layers.SelfAttention(num_heads=4),
-       transformer_layers.LocalSelfAttention(num_heads=4),
-       moe.MoE1D(num_experts=4, hidden_size=512),
-       moe.MoE2D(expert_x=4, expert_y=4, hidden_size=512),
-       transformer_layers.DenseReluDense(hidden_size=512)])
+  hparams.moe_num_experts = 4
+  hparams.moe_expert_x = 4
+  hparams.moe_expert_y = 4
+  hparams.moe_hidden_size = 512
+  hparams.layers = ["self_att", "local_self_att", "moe_1d", "moe_2d", "drd"]
+  return hparams
+
+
+@registry.register_hparams
+def mtf_bitransformer_all_layers_tiny():
+  """Test out all the layers on local CPU."""
+  hparams = mtf_bitransformer_tiny()
+  hparams.moe_num_experts = 4
+  hparams.moe_expert_x = 4
+  hparams.moe_expert_y = 4
+  hparams.moe_hidden_size = 512
+  hparams.encoder_layers = [
+      "self_att", "local_self_att", "moe_1d", "moe_2d", "drd"]
+  hparams.decoder_layers = [
+      "self_att", "local_self_att", "enc_att", "moe_1d", "moe_2d", "drd"]
   return hparams
 
 
@@ -563,32 +583,26 @@ def mtr_lm_dense_3():
 
 
 @registry.register_hparams
-def mtr_lm_v1(num_heads=8, num_memory_heads=0):
+def mtr_lm_v1():
   """Model incorporating mixture-of-experts, local and global attention.
 
   ~6B parameters
 
   32 experts in 3 hierarchichal moe layers.
 
-  Args:
-    num_heads: an optional integer
-    num_memory_heads: an optional integer
-
   Returns:
     a hparams
   """
   hparams = mtr_lm_dense(0)
-  local_att = transformer_layers.LocalSelfAttention(
-      num_heads=num_heads, num_memory_heads=num_memory_heads,
-      key_value_size=128)
-  att = transformer_layers.SelfAttention(
-      num_heads=num_heads, num_memory_heads=num_memory_heads,
-      key_value_size=128)
-  drd = transformer_layers.DenseReluDense(hidden_size=2048)
-  hmoe = moe.MoE2D(expert_x=8, expert_y=4, hidden_size=32768)
-  hparams.layer_stack = transformer.LayerStack(
-      ([local_att, local_att, drd,
-        att, drd, local_att, local_att, hmoe] * 4)[:-1])
+  hparams.layers = (["local_self_att", "local_self_att", "drd",
+                     "self_att", "drd", "local_self_att",
+                     "local_self_att", "moe_2d"] * 4)[:-1]
+  hparams.d_kv = 128
+  hparams.moe_expert_x = 8
+  hparams.moe_expert_y = 4
+  hparams.moe_hidden_size = 32768
+  hparams.d_ff = 2048
+  hparams.num_memory_heads = 0
   hparams.mesh_shape = "b0:4;b1:8"
   hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
   hparams.outer_batch_size = 4
@@ -598,7 +612,9 @@ def mtr_lm_v1(num_heads=8, num_memory_heads=0):
 @registry.register_hparams
 def mtr_lm_v1_h1_8():
   """Version for fast decoding."""
-  return mtr_lm_v1(num_heads=8, num_memory_heads=1)
+  hparams = mtr_lm_v1()
+  hparams.num_memory_heads = 1
+  return hparams
 
 
 def mtr_tr_dense(sz):
@@ -620,8 +636,6 @@ def mtr_tr_dense(sz):
   hparams.d_model = 1024
   hparams.max_length = 256
   hparams.batch_size = 128
-  # Parameters for my_layer_stack()
-  hparams.num_hidden_layers = 6
   hparams.d_ff = int(4096 * n)
   hparams.d_kv = 128
   hparams.encoder_num_heads = int(8 * n)
@@ -671,7 +685,7 @@ def mtr_tr_dense_3_88():
 @registry.register_hparams
 def mtr_tr_dense_3_fast():
   hparams = mtr_tr_dense_3()
-  hparams.decoder_local_attention_radius = 32
+  hparams.local_attention_radius = 32
   hparams.decoder_num_heads = 128
   hparams.decoder_num_memory_heads = 8
   return hparams
@@ -680,7 +694,8 @@ def mtr_tr_dense_3_fast():
 def mtr_tr_dense_local(sz):
   """With local self-attention in the decoder."""
   hparams = mtr_tr_dense(sz)
-  hparams.decoder_local_attention_radius = 32
+  hparams.decoder_layers = ["local_self_att", "enc_att", "drd"] * 6
+  hparams.local_attention_radius = 32
   return hparams
 
 
@@ -692,7 +707,7 @@ def mtr_tr_dense_local_0():
 @registry.register_hparams
 def mtr_tr_dense_local_0_w8():
   hparams = mtr_tr_dense_local_0()
-  hparams.decoder_local_attention_radius = 8
+  hparams.local_attention_radius = 8
   return hparams
 
 
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 23d22f890..3ee5a88e9 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -385,7 +385,7 @@ def decode_from_file(estimator,
   num_decode_batches = (len(sorted_inputs) - 1) // decode_hp.batch_size + 1
 
   if estimator.config.use_tpu:
-    length = getattr(hparams, "length", hparams.max_length)
+    length = getattr(hparams, "length", 0) or hparams.max_length
     batch_ids = []
     for line in sorted_inputs:
       if has_input:

From e0238fe18b10a3f9099b98c372c27a7c27bf47be Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 25 Jan 2019 22:07:17 +0100
Subject: [PATCH 1581/2720] Planner cleanup and new hparam sets (#1412)

* Cleanup planner parameters

* Record videos in the evaluator
---
 tensor2tensor/rl/evaluator.py | 161 ++++++++++++++++++++++++----------
 tensor2tensor/rl/rl_utils.py  |  57 +++++++-----
 2 files changed, 148 insertions(+), 70 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index be6b044de..7a40223cd 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -75,6 +75,9 @@
 flags.DEFINE_string(
     "debug_video_path", "", "Path to save the planner debug video at."
 )
+flags.DEFINE_integer(
+    "num_debug_videos", 1, "Number of debug videos to generate."
+)
 
 # Unused flags needed to pass for multi-run infrastructure.
 flags.DEFINE_bool("autotune", False, "Unused here.")
@@ -94,7 +97,6 @@ def planner_tiny():
       env_type="simulated",
       uct_const=0.0,
       uniform_first_action=True,
-      uct_std_normalization=False,
   )
 
 
@@ -108,12 +110,11 @@ def planner_small():
       env_type="simulated",
       uct_const=0.0,
       uniform_first_action=True,
-      uct_std_normalization=False,
   )
 
 
 @registry.register_hparams
-def planner_guess1():
+def planner_base():
   return tf.contrib.training.HParams(
       num_rollouts=96,
       batch_size=96,
@@ -121,37 +122,93 @@ def planner_guess1():
       rollout_agent_type="policy",
       env_type="simulated",
       uct_const=0.,
-      uniform_first_action=False,
-      uct_std_normalization=False,
+      uniform_first_action=True,
   )
 
 
+# Tuning of uniform_first_action and uct_const. Default params repeated for
+# clarity.
+
+
+@registry.register_hparams
+def planner_guess1():
+  hparams = planner_base()
+  hparams.uniform_first_action = False
+  hparams.uct_const = 0.
+  return hparams
+
+
 @registry.register_hparams
 def planner_guess2():
-  return tf.contrib.training.HParams(
-      num_rollouts=96,
-      batch_size=96,
-      planning_horizon=8,
-      rollout_agent_type="policy",
-      env_type="simulated",
-      uct_const=3.,
-      uniform_first_action=True,
-      uct_std_normalization=True,
-  )
+  hparams = planner_base()
+  hparams.uniform_first_action = True
+  hparams.uct_const = 3.
+  return hparams
 
 
 @registry.register_hparams
 def planner_guess3():
-  return tf.contrib.training.HParams(
-      num_rollouts=96,
-      batch_size=96,
-      planning_horizon=8,
-      rollout_agent_type="policy",
-      env_type="simulated",
-      uct_const=2.,
-      uniform_first_action=False,
-      uct_std_normalization=False,
-  )
+  hparams = planner_base()
+  hparams.uniform_first_action = False
+  hparams.uct_const = 2.
+  return hparams
+
+
+# Tuning of uct_const, num_collouts and normalizer_window_size.
+
+
+@registry.register_hparams
+def planner_guess4():
+  hparams = planner_base()
+  hparams.uct_const = 2
+  hparams.num_rollouts = 96
+  hparams.normalizer_window_size = 30
+  return hparams
+
+
+@registry.register_hparams
+def planner_guess5():
+  hparams = planner_base()
+  hparams.uct_const = 2
+  hparams.num_rollouts = 3 * 96
+  hparams.normalizer_window_size = 30
+  return hparams
+
+
+@registry.register_hparams
+def planner_guess6():
+  hparams = planner_base()
+  hparams.uct_const = 4
+  hparams.num_rollouts = 96
+  hparams.normalizer_window_size = 30
+  return hparams
+
+
+@registry.register_hparams
+def planner_guess7():
+  hparams = planner_base()
+  hparams.uct_const = 4
+  hparams.num_rollouts = 3 * 96
+  hparams.normalizer_window_size = 30
+  return hparams
+
+
+@registry.register_hparams
+def planner_guess8():
+  hparams = planner_base()
+  hparams.uct_const = 2
+  hparams.num_rollouts = 3 * 96
+  hparams.normalizer_window_size = 300
+  return hparams
+
+
+@registry.register_hparams
+def planner_guess9():
+  hparams = planner_base()
+  hparams.uct_const = 4
+  hparams.num_rollouts = 3 * 96
+  hparams.normalizer_window_size = 300
+  return hparams
 
 
 def make_env(env_type, real_env, sim_env_kwargs):
@@ -169,10 +226,8 @@ def make_env(env_type, real_env, sim_env_kwargs):
 
 def make_agent(
     agent_type, env, policy_hparams, policy_dir, sampling_temp,
-    sim_env_kwargs=None, frame_stack_size=None, planning_horizon=None,
-    rollout_agent_type=None, batch_size=None, num_rollouts=None,
-    inner_batch_size=None, video_writer=None, env_type=None,
-    uct_const=None, uct_std_normalization=None, uniform_first_action=None
+    sim_env_kwargs=None, frame_stack_size=None, rollout_agent_type=None,
+    batch_size=None, inner_batch_size=None, env_type=None, **planner_kwargs
 ):
   """Factory function for Agents."""
   if batch_size is None:
@@ -191,17 +246,14 @@ def make_agent(
               sampling_temp, batch_size=inner_batch_size
           ), make_env(env_type, env.env, sim_env_kwargs),
           lambda env: rl_utils.BatchStackWrapper(env, frame_stack_size),
-          num_rollouts, planning_horizon,
-          discount_factor=policy_hparams.gae_gamma,
-          uct_const=uct_const, uct_std_normalization=uct_std_normalization,
-          uniform_first_action=uniform_first_action, video_writer=video_writer
+          discount_factor=policy_hparams.gae_gamma, **planner_kwargs
       ),
   }[agent_type]()
 
 
 def make_eval_fn_with_agent(
     agent_type, planner_hparams, model_dir, log_every_steps=None,
-    video_writer=None
+    video_writers=()
 ):
   """Returns an out-of-graph eval_fn using the Agent API."""
   def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
@@ -212,18 +264,23 @@ def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
         base_env, loop_hparams, batch_size=planner_hparams.batch_size,
         model_dir=model_dir
     )
+    planner_kwargs = planner_hparams.values()
+    planner_kwargs.pop("batch_size")
+    planner_kwargs.pop("rollout_agent_type")
+    planner_kwargs.pop("env_type")
     agent = make_agent(
         agent_type, env, policy_hparams, policy_dir, sampling_temp,
         sim_env_kwargs, loop_hparams.frame_stack_size,
-        planner_hparams.planning_horizon, planner_hparams.rollout_agent_type,
-        num_rollouts=planner_hparams.num_rollouts,
-        inner_batch_size=planner_hparams.batch_size, video_writer=video_writer,
-        env_type=planner_hparams.env_type, uct_const=planner_hparams.uct_const,
-        uct_std_normalization=planner_hparams.uct_std_normalization,
-        uniform_first_action=planner_hparams.uniform_first_action
+        planner_hparams.rollout_agent_type,
+        inner_batch_size=planner_hparams.batch_size,
+        env_type=planner_hparams.env_type,
+        video_writers=video_writers, **planner_kwargs
     )
+    kwargs = {}
+    if not agent.records_own_videos:
+      kwargs["video_writers"] = video_writers
     rl_utils.run_rollouts(
-        env, agent, env.reset(), log_every_steps=log_every_steps
+        env, agent, env.reset(), log_every_steps=log_every_steps, **kwargs
     )
     assert len(base_env.current_epoch_rollouts()) == env.batch_size
   return eval_fn
@@ -232,7 +289,7 @@ def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
 def evaluate(
     loop_hparams, planner_hparams, policy_dir, model_dir, eval_metrics_dir,
     agent_type, eval_with_learner, log_every_steps, debug_video_path,
-    report_fn=None, report_metric=None
+    num_debug_videos=1, report_fn=None, report_metric=None
 ):
   """Evaluate."""
   if eval_with_learner:
@@ -242,22 +299,29 @@ def evaluate(
     assert report_metric is not None
 
   eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
-  video_writer = None
+  video_writers = ()
   kwargs = {}
   if not eval_with_learner:
     if debug_video_path:
-      video_writer = common_video.WholeVideoWriter(
-          fps=10, output_path=debug_video_path, file_format="avi")
+      tf.gfile.MakeDirs(debug_video_path)
+      video_writers = [
+          common_video.WholeVideoWriter(
+              fps=10,
+              output_path=os.path.join(debug_video_path, "{}.avi".format(i)),
+              file_format="avi",
+          )
+          for i in range(num_debug_videos)
+      ]
     kwargs["eval_fn"] = make_eval_fn_with_agent(
         agent_type, planner_hparams, model_dir, log_every_steps=log_every_steps,
-        video_writer=video_writer
+        video_writers=video_writers
     )
   eval_metrics = rl_utils.evaluate_all_configs(
       loop_hparams, policy_dir, **kwargs
   )
   rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, 0)
 
-  if video_writer is not None:
+  for video_writer in video_writers:
     video_writer.finish_to_disk()
 
   # Report metrics
@@ -325,7 +389,8 @@ def main(_):
       loop_hparams, planner_hparams, policy_dir, model_dir,
       eval_metrics_dir, FLAGS.agent, FLAGS.eval_with_learner,
       FLAGS.log_every_steps if FLAGS.log_every_steps > 0 else None,
-      debug_video_path=FLAGS.debug_video_path
+      debug_video_path=FLAGS.debug_video_path,
+      num_debug_videos=FLAGS.num_debug_videos
   )
 
 
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index d145c4fb5..fc7aeb441 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -277,7 +277,7 @@ def augment_observation(
 
 def run_rollouts(
     env, agent, initial_observations, step_limit=None, discount_factor=1.0,
-    log_every_steps=None, video_writer=None
+    log_every_steps=None, video_writers=(), color_bar=False
 ):
   """Runs a batch of rollouts from given initial observations."""
   num_dones = 0
@@ -286,12 +286,11 @@ def run_rollouts(
   step_index = 0
   cum_rewards = 0
 
-  if video_writer is not None:
-    obs_stack = initial_observations[0]
+  for (video_writer, obs_stack) in zip(video_writers, initial_observations):
     for (i, ob) in enumerate(obs_stack):
       debug_frame = augment_observation(
           ob, reward=0, cum_reward=0, frame_index=(-len(obs_stack) + i + 1),
-          bar_color=(0, 255, 0)
+          bar_color=((0, 255, 0) if color_bar else None)
       )
       video_writer.write(debug_frame)
 
@@ -324,11 +323,13 @@ def proceed():
     cum_rewards = cum_rewards * discount_factor + rewards
     step_index += 1
 
-    if video_writer is not None:
-      ob = observations[0, -1]
+    for (video_writer, obs_stack, reward, cum_reward) in zip(
+        video_writers, observations, rewards, cum_rewards
+    ):
+      ob = obs_stack[-1]
       debug_frame = augment_observation(
-          ob, reward=rewards[0], cum_reward=cum_rewards[0],
-          frame_index=step_index, bar_color=(255, 0, 0)
+          ob, reward=reward, cum_reward=cum_reward,
+          frame_index=step_index, bar_color=((255, 0, 0) if color_bar else None)
       )
       video_writer.write(debug_frame)
 
@@ -346,6 +347,7 @@ class BatchAgent(object):
   """
 
   needs_env_state = False
+  records_own_videos = False
 
   def __init__(self, batch_size, observation_space, action_space):
     self.batch_size = batch_size
@@ -462,11 +464,22 @@ class PlannerAgent(BatchAgent):
   """Agent based on temporal difference planning."""
 
   needs_env_state = True
+  records_own_videos = True
 
   def __init__(
-      self, batch_size, rollout_agent, sim_env, wrapper_fn, num_rollouts,
-      planning_horizon, discount_factor=1.0, uct_const=0,
-      uct_std_normalization=False, uniform_first_action=True, video_writer=None
+      self,
+      batch_size,
+      rollout_agent,
+      sim_env,
+      wrapper_fn,
+      num_rollouts,
+      planning_horizon,
+      discount_factor=1.0,
+      uct_const=0,
+      uniform_first_action=True,
+      normalizer_window_size=30,
+      normalizer_epsilon=0.001,
+      video_writers=(),
   ):
     super(PlannerAgent, self).__init__(
         batch_size, rollout_agent.observation_space, rollout_agent.action_space
@@ -479,9 +492,10 @@ def __init__(
     self._discount_factor = discount_factor
     self._planning_horizon = planning_horizon
     self._uct_const = uct_const
-    self._uct_std_normalization = uct_std_normalization
     self._uniform_first_action = uniform_first_action
-    self._video_writer = video_writer
+    self._normalizer_window_size = normalizer_window_size
+    self._normalizer_epsilon = normalizer_epsilon
+    self._video_writers = video_writers
     self._best_mc_values = [[] for _ in range(self.batch_size)]
 
   def act(self, observations, env_state=None):
@@ -502,14 +516,14 @@ def run_batch_from(observation, planner_index, batch_index):
       (initial_observations, initial_rewards, _) = self._wrapped_env.step(
           actions
       )
-      writer = None
-      if planner_index == 0 and batch_index == 0:
-        writer = self._video_writer
+      video_writers = ()
+      if planner_index < len(self._video_writers) and batch_index == 0:
+        video_writers = (self._video_writers[planner_index],)
       (final_observations, cum_rewards) = run_rollouts(
           self._wrapped_env, self._rollout_agent, initial_observations,
           discount_factor=self._discount_factor,
           step_limit=self._planning_horizon,
-          video_writer=writer)
+          video_writers=video_writers, color_bar=True)
       values = self._rollout_agent.estimate_value(final_observations)
       total_values = (
           initial_rewards + self._discount_factor * cum_rewards +
@@ -549,11 +563,10 @@ def monte_carlo_value(action):
       )
       best_mc_values.append(mc_values.max())
 
-      if self._uct_std_normalization:
-        agg = np.std
-      else:
-        agg = lambda x: np.mean(np.std(x))
-      normalizer = max(agg(best_mc_values[-30:]), 0.001)
+      normalizer = max(
+          np.std(best_mc_values[-self._normalizer_window_size:]),
+          self._normalizer_epsilon
+      )
       normalized_mc_values = mc_values / normalizer
 
       uct_bonuses = np.array(

From e28ba48930129d362bcfe960f824f5e0a7fabea7 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 25 Jan 2019 13:07:37 -0800
Subject: [PATCH 1582/2720] internal merge of PR #1412

PiperOrigin-RevId: 230954839
---
 tensor2tensor/rl/evaluator.py | 2 +-
 tensor2tensor/rl/rl_utils.py  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 7a40223cd..b6e973b42 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -305,7 +305,7 @@ def evaluate(
     if debug_video_path:
       tf.gfile.MakeDirs(debug_video_path)
       video_writers = [
-          common_video.WholeVideoWriter(
+          common_video.WholeVideoWriter(  # pylint: disable=g-complex-comprehension
               fps=10,
               output_path=os.path.join(debug_video_path, "{}.avi".format(i)),
               file_format="avi",
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index fc7aeb441..3686c250b 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -227,7 +227,7 @@ def initial_frame_chooser(batch_size):
         initial_frames[0] = deterministic_initial_frames
 
     return np.stack([
-        [frame.observation.decode() for frame in initial_frame_stack]
+        [frame.observation.decode() for frame in initial_frame_stack]  # pylint: disable=g-complex-comprehension
         for initial_frame_stack in initial_frames
     ])
   return initial_frame_chooser

From de8f44b6dc0460f641af9dd1b5d6e3ed85a22200 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 25 Jan 2019 14:41:23 -0800
Subject: [PATCH 1583/2720] Fix shape mismatch in T2T integration of SARI
 metric.

PiperOrigin-RevId: 230971810
---
 tensor2tensor/utils/sari_hook.py      | 2 +-
 tensor2tensor/utils/sari_hook_test.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/sari_hook.py b/tensor2tensor/utils/sari_hook.py
index d433d5e41..65641b1d9 100644
--- a/tensor2tensor/utils/sari_hook.py
+++ b/tensor2tensor/utils/sari_hook.py
@@ -242,7 +242,7 @@ def sari_score(predictions, labels, features, **unused_kwargs):
   # Convert the inputs and outputs to a [batch_size, sequence_length] tensor.
   inputs = tf.squeeze(features["inputs"], axis=[-1, -2])
   outputs = tf.to_int32(tf.argmax(predictions, axis=-1))
-  outputs = tf.squeeze(outputs, axis=-1)
+  outputs = tf.squeeze(outputs, axis=[-1, -2])
 
   # Convert the labels to a [batch_size, 1, sequence_length] tensor.
   labels = tf.squeeze(labels, axis=[-1, -2])
diff --git a/tensor2tensor/utils/sari_hook_test.py b/tensor2tensor/utils/sari_hook_test.py
index 1d039847d..b168d0143 100644
--- a/tensor2tensor/utils/sari_hook_test.py
+++ b/tensor2tensor/utils/sari_hook_test.py
@@ -128,7 +128,7 @@ def testIdsWithZeros(self):
 
   def testSariScoreE2E(self):
     """Tests the SARI metrics end-to-end."""
-    predictions = np.random.randint(4, size=(12, 12, 1, 12))
+    predictions = np.random.randint(4, size=(12, 12, 1, 1, 12))
     targets = np.random.randint(4, size=(12, 12, 1, 1))
     inputs = np.random.randint(4, size=(12, 12, 1, 1))
     with self.test_session() as session:

From 90f345208de5c5f5544b77d7f1231571c5a95171 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 26 Jan 2019 22:52:45 -0800
Subject: [PATCH 1584/2720] Register optimizers with proper names. I'm opting
 to keep the names explicit still.

PiperOrigin-RevId: 231091053
---
 tensor2tensor/utils/optimize.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index c8f1a9a50..5c3b5b3bf 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -94,7 +94,7 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
   return train_op
 
 
-@registry.register_optimizer("adam")
+@registry.register_optimizer("Adam")
 def adam(learning_rate, hparams):
   # We change the default epsilon for Adam.
   # Using LazyAdam as it's much faster for large vocabulary embeddings.
@@ -105,7 +105,7 @@ def adam(learning_rate, hparams):
       epsilon=hparams.optimizer_adam_epsilon)
 
 
-@registry.register_optimizer("multistep_adam")
+@registry.register_optimizer("MultistepAdam")
 def multistep_adam(learning_rate, hparams):
   return multistep_optimizer.MultistepAdamOptimizer(
       learning_rate,
@@ -115,7 +115,7 @@ def multistep_adam(learning_rate, hparams):
       n=hparams.optimizer_multistep_accumulate_steps)
 
 
-@registry.register_optimizer("momentum")
+@registry.register_optimizer("Momentum")
 def momentum(learning_rate, hparams):
   return tf.train.MomentumOptimizer(
       learning_rate,
@@ -123,14 +123,14 @@ def momentum(learning_rate, hparams):
       use_nesterov=hparams.optimizer_momentum_nesterov)
 
 
-@registry.register_optimizer("yellow_fin")
+@registry.register_optimizer("YellowFin")
 def yellow_fin(learning_rate, hparams):
   return yellowfin.YellowFinOptimizer(
       learning_rate=learning_rate,
       momentum=hparams.optimizer_momentum_momentum)
 
 
-@registry.register_optimizer("true_adam")
+@registry.register_optimizer("TrueAdam")
 def true_adam(learning_rate, hparams):
   return tf.train.AdamOptimizer(
       learning_rate,
@@ -139,7 +139,7 @@ def true_adam(learning_rate, hparams):
       epsilon=hparams.optimizer_adam_epsilon)
 
 
-@registry.register_optimizer("adam_w")
+@registry.register_optimizer("AdamW")
 def adam_w(learning_rate, hparams):
   # Openai gpt used weight decay.
   # Given the internals of AdamW, weight decay dependent on the

From dd604e29287f73894de17ef8bdebc639cfe97800 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 28 Jan 2019 10:26:34 -0800
Subject: [PATCH 1585/2720] Adding a comment for a confusing layer name which
 named dense layers "conv"s.

PiperOrigin-RevId: 231244093
---
 tensor2tensor/layers/common_layers.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 3b32e1678..b48300c26 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1222,6 +1222,8 @@ def dense_relu_dense(inputs,
                      dropout_broadcast_dims=None,
                      name=None):
   """Hidden layer with RELU activation followed by linear projection."""
+  # layer_name is appended with "conv1" or "conv2" in this method only for
+  # historical reasons. These are in fact dense layers.
   layer_name = "%s_{}" % name if name else "{}"
   h = dense(
       inputs,

From 74f37a45f7a0566d4a3dc79086a0598094b575db Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Mon, 28 Jan 2019 14:13:18 -0800
Subject: [PATCH 1586/2720] Code for Multiquery attention paper

PiperOrigin-RevId: 231286854
---
 tensor2tensor/models/__init__.py              |   1 +
 .../models/research/multiquery_paper.py       | 218 ++++++++++++++++++
 2 files changed, 219 insertions(+)
 create mode 100644 tensor2tensor/models/research/multiquery_paper.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index f0618faba..77176a9e0 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -52,6 +52,7 @@
 from tensor2tensor.models.research import glow
 from tensor2tensor.models.research import lm_experiments
 from tensor2tensor.models.research import moe_experiments
+from tensor2tensor.models.research import multiquery_paper
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research import similarity_transformer
 from tensor2tensor.models.research import super_lm
diff --git a/tensor2tensor/models/research/multiquery_paper.py b/tensor2tensor/models/research/multiquery_paper.py
new file mode 100644
index 000000000..5b5491026
--- /dev/null
+++ b/tensor2tensor/models/research/multiquery_paper.py
@@ -0,0 +1,218 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Experiments for Multiquery-Attention Paper.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.models import mtf_transformer2
+from tensor2tensor.utils import registry
+
+
+@registry.register_hparams
+def mqp_ende_base():
+  # params=211M
+  hparams = mtf_transformer2.mtr_tr_dense_0()
+  hparams.learning_rate_decay_steps = 20000
+  hparams.shared_embedding_and_softmax_weights = True
+  hparams.layer_prepostprocess_dropout = 0.2
+  return hparams
+
+
+@registry.register_hparams
+def mqp_ende_local():
+  hparams = mqp_ende_base()
+  hparams.decoder_local_attention_radius = 32
+  return hparams
+
+
+@registry.register_hparams
+def mqp_ende_mq8():
+  # params=178M
+  hparams = mqp_ende_base()
+  hparams.decoder_num_heads = 8
+  hparams.decoder_num_memory_heads = 1
+  hparams.encoder_num_heads = 8
+  hparams.encoder_num_memory_heads = 1
+  return hparams
+
+
+@registry.register_hparams
+def mqp_ende_mq8_ff5440():
+  # params=211M
+  hparams = mqp_ende_mq8()
+  hparams.d_ff = 5440
+  return hparams
+
+
+@registry.register_hparams
+def mqp_ende_mq8_ff5440_local():
+  hparams = mqp_ende_mq8_ff5440()
+  hparams.decoder_local_attention_radius = 32
+  return hparams
+
+
+@registry.register_hparams
+def mqp_ende_h4_kv256():
+  hparams = mqp_ende_base()
+  hparams.decoder_num_heads = 4
+  hparams.encoder_num_heads = 4
+  hparams.d_kv = 256
+  return hparams
+
+
+@registry.register_hparams
+def mqp_ende_h2_kv512():
+  hparams = mqp_ende_base()
+  hparams.decoder_num_heads = 2
+  hparams.encoder_num_heads = 2
+  hparams.d_kv = 512
+  return hparams
+
+
+@registry.register_hparams
+def mqp_ende_h1_kv1024():
+  hparams = mqp_ende_base()
+  hparams.decoder_num_heads = 1
+  hparams.encoder_num_heads = 1
+  hparams.d_kv = 1024
+  return hparams
+
+
+@registry.register_hparams
+def mqp_ende_h4_ff5632():
+  hparams = mqp_ende_base()
+  hparams.decoder_num_heads = 4
+  hparams.encoder_num_heads = 4
+  hparams.d_ff = 5632
+  return hparams
+
+
+@registry.register_hparams
+def mqp_ende_h2_ff6400():
+  hparams = mqp_ende_base()
+  hparams.decoder_num_heads = 2
+  hparams.encoder_num_heads = 2
+  hparams.d_ff = 6400
+  return hparams
+
+
+@registry.register_hparams
+def mqp_ende_h1_ff6784():
+  hparams = mqp_ende_base()
+  hparams.decoder_num_heads = 1
+  hparams.encoder_num_heads = 1
+  hparams.d_ff = 6784
+  return hparams
+
+
+@registry.register_hparams
+def mqp_ende_h2_kv64_ff6784():
+  hparams = mqp_ende_base()
+  hparams.decoder_num_heads = 2
+  hparams.encoder_num_heads = 2
+  hparams.d_kv = 64
+  hparams.d_ff = 6784
+  return hparams
+
+
+@registry.register_hparams
+def mqp_ende_h4_kv32_ff6784():
+  hparams = mqp_ende_base()
+  hparams.decoder_num_heads = 4
+  hparams.encoder_num_heads = 4
+  hparams.d_kv = 32
+  hparams.d_ff = 6784
+  return hparams
+
+
+@registry.register_hparams
+def mqp_ende_h8_kv16_ff6784():
+  hparams = mqp_ende_base()
+  hparams.decoder_num_heads = 8
+  hparams.encoder_num_heads = 8
+  hparams.d_kv = 16
+  return hparams
+
+
+@registry.register_hparams
+def mqp_lm1b_base():
+  """Series of architectures for language modeling."""
+  hparams = mtf_transformer2.mtf_unitransformer_base()
+  hparams.d_model = 1024
+  hparams.max_length = 256
+  hparams.batch_size = 256
+  # Parameters for my_layer_stack()
+  hparams.num_hidden_layers = 6
+  hparams.d_ff = 8192
+  hparams.d_kv = 128
+  hparams.num_heads = 8
+  hparams.learning_rate_decay_steps = 13600
+  hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
+  hparams.mesh_shape = "batch:32"
+  return hparams
+
+
+@registry.register_hparams
+def mqp_lm1b_mq8():
+  hparams = mqp_lm1b_base()
+  hparams.num_heads = 8
+  hparams.num_memory_heads = 1
+  return hparams
+
+
+@registry.register_hparams
+def mqp_lm1b_mq8_ff9088():
+  hparams = mqp_lm1b_mq8()
+  hparams.d_ff = 9088
+  return hparams
+
+
+@registry.register_hparams
+def mqp_lm1b_h1_ff9984():
+  hparams = mqp_lm1b_base()
+  hparams.num_heads = 1
+  hparams.d_ff = 9984
+  return hparams
+
+
+@registry.register_hparams
+def mqp_lm1b_h2_kv64_ff9984():
+  hparams = mqp_lm1b_base()
+  hparams.num_heads = 2
+  hparams.d_kv = 64
+  hparams.d_ff = 9984
+  return hparams
+
+
+@registry.register_hparams
+def mqp_lm1b_h4_kv32_ff9984():
+  hparams = mqp_lm1b_base()
+  hparams.num_heads = 4
+  hparams.d_kv = 32
+  hparams.d_ff = 9984
+  return hparams
+
+
+@registry.register_hparams
+def mqp_lm1b_h8_kv16_ff9984():
+  hparams = mqp_lm1b_base()
+  hparams.num_heads = 8
+  hparams.d_kv = 16
+  hparams.d_ff = 9984
+  return hparams

From ac4cb05e7a1ee3424a1b17f41211f3662f96abd4 Mon Sep 17 00:00:00 2001
From: Dominic Jack <thedomjack@gmail.com>
Date: Tue, 29 Jan 2019 09:24:36 +1000
Subject: [PATCH 1587/2720] Registry refactor (#1410)

* registry refactor and deprecated call-site updates

* added on_problem_set callback, simplified name

* changed optimizer registration names to snake_case, documentation

* removed create_registry
---
 tensor2tensor/bin/t2t_attack.py               |   2 +-
 tensor2tensor/bin/t2t_datagen.py              |   5 +-
 tensor2tensor/bin/t2t_prune.py                |   2 +-
 tensor2tensor/layers/common_hparams.py        |   4 +-
 tensor2tensor/models/mtf_transformer2.py      |   2 +-
 .../models/research/adafactor_experiments.py  |  10 +-
 tensor2tensor/models/research/autoencoders.py |   2 +-
 .../models/research/transformer_nat.py        |   2 +-
 .../models/research/transformer_vae.py        |   4 +-
 .../models/research/vqa_attention.py          |   2 +-
 .../models/research/vqa_self_attention.py     |   2 +-
 tensor2tensor/models/shake_shake.py           |   2 +-
 tensor2tensor/models/transformer.py           |   4 +-
 tensor2tensor/models/vanilla_gan.py           |   2 +-
 tensor2tensor/problems.py                     |   2 +-
 tensor2tensor/rl/datagen_with_agent.py        |   2 +-
 .../transformer_test_ckpt/hparams.json        |   2 +-
 tensor2tensor/utils/adafactor.py              |   2 +-
 tensor2tensor/utils/learning_rate.py          |   2 +-
 tensor2tensor/utils/optimize.py               |  19 +-
 tensor2tensor/utils/registry.py               | 818 +++++++++---------
 tensor2tensor/utils/registry_test.py          | 174 +++-
 22 files changed, 596 insertions(+), 470 deletions(-)

diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 6276d6b4c..f7e0fcc60 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -73,7 +73,7 @@ def create_attack_params():
 
 
 def create_attack(attack):
-  return registry.attacks(attack)
+  return registry.attack(attack)
 
 
 def create_surrogate_hparams():
diff --git a/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index b778bc984..c32a86cf5 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -147,7 +147,7 @@ def main(_):
 
   # Calculate the list of problems to generate.
   problems = sorted(
-      list(_SUPPORTED_PROBLEM_GENERATORS) + registry.list_problems())
+      list(_SUPPORTED_PROBLEM_GENERATORS) + registry.list_base_problems())
   for exclude in FLAGS.exclude_problems.split(","):
     if exclude:
       problems = [p for p in problems if exclude not in p]
@@ -169,7 +169,8 @@ def main(_):
 
   if not problems:
     problems_str = "\n  * ".join(
-        sorted(list(_SUPPORTED_PROBLEM_GENERATORS) + registry.list_problems()))
+        sorted(list(_SUPPORTED_PROBLEM_GENERATORS) +
+               registry.list_base_problems()))
     error_msg = ("You must specify one of the supported problems to "
                  "generate data for:\n  * " + problems_str + "\n")
     error_msg += ("TIMIT and parsing need data_sets specified with "
diff --git a/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
index 008462e63..e7c8c75db 100644
--- a/tensor2tensor/bin/t2t_prune.py
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -54,7 +54,7 @@ def create_pruning_params():
 
 
 def create_pruning_strategy(name):
-  return registry.pruning_strategies(name)
+  return registry.pruning_strategy(name)
 
 
 def main(argv):
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 2bbfb7618..7158c3f4d 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -55,7 +55,7 @@ def basic_params1():
       initializer="orthogonal",
       initializer_gain=1.5,
       label_smoothing=0.1,
-      optimizer="Adam",
+      optimizer="adam",
       optimizer_adam_epsilon=1e-6,
       optimizer_adam_beta1=0.85,
       optimizer_adam_beta2=0.997,
@@ -466,7 +466,7 @@ def basic_range1(ranged_hparams):
   rhp.set_float("optimizer_adam_beta2", 0.995, 0.999)
   rhp.set_categorical(
       "optimizer",
-      ["Adam", "Adagrad", "Momentum", "RMSProp", "SGD", "YellowFin"])
+      ["adam", "adagrad", "momentum", "rms_prop", "sgd", "yellow_fin"])
 
 
 @registry.register_ranged_hparams
diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index 43dafed38..2090228e8 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -269,7 +269,7 @@ def sample(self, features, mesh):
     return self.combine_batch_dims(ret)
 
 
-layers_registry = registry.create_registry("layers")
+layers_registry = registry.Registries.layers
 
 
 # The following functions construct layers based on hyperparmeters
diff --git a/tensor2tensor/models/research/adafactor_experiments.py b/tensor2tensor/models/research/adafactor_experiments.py
index d7d3d4e2c..fbe4dbc2f 100644
--- a/tensor2tensor/models/research/adafactor_experiments.py
+++ b/tensor2tensor/models/research/adafactor_experiments.py
@@ -30,16 +30,16 @@ def mimic_adam_with_adafactor(hparams):
   Some minor things may be different, like epsilon and beta1 correction.
 
   Args:
-    hparams: model hyperparameters where "Adam" in hparams.optimizer
+    hparams: model hyperparameters where "adam" in hparams.optimizer
   """
-  assert "Adam" in hparams.optimizer
-  hparams.optimizer = "Adafactor"
+  assert "adam" in hparams.optimizer
+  hparams.optimizer = "adafactor"
   hparams.optimizer_adafactor_beta1 = hparams.optimizer_adam_beta1
   hparams.optimizer_adafactor_beta2 = hparams.optimizer_adam_beta2
   hparams.optimizer_adafactor_multiply_by_parameter_scale = False
   hparams.optimizer_adafactor_factored = False
   hparams.optimizer_adafactor_clipping_threshold = None
-  hparams.optimizer_adafactor_decay_type = "Adam"
+  hparams.optimizer_adafactor_decay_type = "adam"
 
 
 @registry.register_hparams
@@ -50,7 +50,7 @@ def afx_adam():
   hparams.optimizer_adam_beta2 = 0.999
   hparams.symbol_modality_num_shards = 1
   hparams.batch_size = 2048
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.learning_rate_schedule = (
       "constant*rsqrt_decay*linear_warmup*rsqrt_hidden_size")
   hparams.learning_rate_constant = 2.0
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 3780bd6e6..ecab3faa1 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1020,7 +1020,7 @@ def body(self, features):
 def autoencoder_basic():
   """Basic autoencoder model."""
   hparams = common_hparams.basic_params1()
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.learning_rate_constant = 0.0002
   hparams.learning_rate_warmup_steps = 500
   hparams.learning_rate_schedule = "constant * linear_warmup"
diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 7a775153d..ccea877ed 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -392,7 +392,7 @@ def transformer_nat_small():
   hparams.filter_size = 2048
   hparams.label_smoothing = 0.0
   hparams.force_full_predict = True
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.optimizer_adam_epsilon = 1e-9
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.997
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index c630abc6c..64aecd405 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -767,7 +767,7 @@ def transformer_ae_small():
   hparams.filter_size = 2048
   hparams.add_hparam("compress_filter_size", 2048 * 2)
   hparams.label_smoothing = 0.0
-  hparams.optimizer = "Adam"  # Can be unstable, maybe try Adam.
+  hparams.optimizer = "adam"  # Can be unstable, maybe try Adam.
   hparams.optimizer_adam_epsilon = 1e-9
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.997  # Needs tuning, try 0.98 to 0.999.
@@ -941,7 +941,7 @@ def transformer_ae_a3():
 def transformer_ae_a6():
   """Best hparams for transformer with semhash."""
   hparams = transformer_ae_a3()
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.noise_dev = 0.5
   return hparams
 
diff --git a/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index 774387878..57f1975e6 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -335,7 +335,7 @@ def vqa_attention_base():
   hparams = common_hparams.basic_params1()
   hparams.batch_size = 128
   hparams.use_fixed_batch_size = True,
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.999
   hparams.optimizer_adam_epsilon = 1e-8
diff --git a/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
index 482e18c9e..b8388e606 100644
--- a/tensor2tensor/models/research/vqa_self_attention.py
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -684,7 +684,7 @@ def vqa_self_attention_base():
   hparams = common_hparams.basic_params1()
   hparams.batch_size = 128
   hparams.use_fixed_batch_size = True,
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.997
   hparams.optimizer_adam_epsilon = 1e-9
diff --git a/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index 4af05dc21..22f46f989 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -189,7 +189,7 @@ def shakeshake_small():
 @registry.register_hparams
 def shake_shake_quick():
   hparams = shakeshake_small()
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.learning_rate_cosine_cycle_steps = 1000
   hparams.learning_rate = 0.5
   hparams.batch_size = 100
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index f0d0596e8..5788910d2 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1694,7 +1694,7 @@ def transformer_tall_pretrain_lm():
   hparams.learning_rate_constant = 2e-4
   hparams.learning_rate_schedule = (
       "linear_warmup*constant*cosdecay")
-  hparams.optimizer = "AdamW"
+  hparams.optimizer = "adam_w"
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.999
   hparams.optimizer_adam_epsilon = 1e-8
@@ -1739,7 +1739,7 @@ def transformer_tall_pretrain_lm_tpu():
   # Optimizer gets reset in update_hparams_for_tpu so we set it again here.
   hparams.learning_rate_constant = 2e-4
   hparams.learning_rate_schedule = ("linear_warmup * constant * cosdecay")
-  hparams.optimizer = "AdamW"
+  hparams.optimizer = "adam_w"
   return hparams
 
 
diff --git a/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index cb16a4c0b..00fa89d18 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -199,7 +199,7 @@ def infer(self, *args, **kwargs):  # pylint: disable=arguments-differ
 def sliced_gan():
   """Basic parameters for a vanilla_gan."""
   hparams = common_hparams.basic_params1()
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.learning_rate_constant = 0.0002
   hparams.learning_rate_warmup_steps = 500
   hparams.learning_rate_schedule = "constant * linear_warmup"
diff --git a/tensor2tensor/problems.py b/tensor2tensor/problems.py
index 1337c2a33..2cf8c8762 100644
--- a/tensor2tensor/problems.py
+++ b/tensor2tensor/problems.py
@@ -27,7 +27,7 @@ def problem(name):
 
 
 def available():
-  return sorted(registry.list_problems())
+  return registry.list_base_problems()
 
 
 all_problems.import_modules(all_problems.ALL_MODULES)
diff --git a/tensor2tensor/rl/datagen_with_agent.py b/tensor2tensor/rl/datagen_with_agent.py
index f634e8c87..1e5100537 100644
--- a/tensor2tensor/rl/datagen_with_agent.py
+++ b/tensor2tensor/rl/datagen_with_agent.py
@@ -45,7 +45,7 @@ def main(_):
 
   # Create problem if not already defined
   problem_name = "gym_discrete_problem_with_agent_on_%s" % FLAGS.game
-  if problem_name not in registry.list_problems():
+  if problem_name not in registry.Registries.problems:
     gym_env.register_game(FLAGS.game)
 
   # Generate
diff --git a/tensor2tensor/test_data/transformer_test_ckpt/hparams.json b/tensor2tensor/test_data/transformer_test_ckpt/hparams.json
index df9a654c0..1a6a97223 100644
--- a/tensor2tensor/test_data/transformer_test_ckpt/hparams.json
+++ b/tensor2tensor/test_data/transformer_test_ckpt/hparams.json
@@ -1 +1 @@
-{"daisy_chain_variables": true, "optimizer_adam_beta1": 0.9, "scheduled_sampling_prob": 0.0, "num_hidden_layers": 2, "moe_loss_coef": 0.01, "max_target_seq_length": 0, "clip_grad_norm": 0.0, "pos": "timing", "scheduled_sampling_gold_mixin_prob": 0.5, "initializer": "uniform_unit_scaling", "grad_noise_scale": 0.0, "optimizer_momentum_momentum": 0.9, "nbr_decoder_problems": 1, "attention_key_channels": 0, "eval_drop_long_sequences": false, "learning_rate_cosine_cycle_steps": 250000, "prepend_mode": "none", "weight_decay": 0.0, "symbol_modality_skip_top": false, "weight_noise": 0.0, "target_modality": "default", "attention_dropout": 0.1, "parameter_attention_value_channels": 0, "factored_logits": false, "relu_dropout": 0.1, "no_data_parallelism": false, "layer_preprocess_sequence": "n", "sampling_method": "argmax", "learning_rate": 0.2, "num_heads": 2, "max_length": 256, "summarize_grads": false, "attention_value_channels": 0, "num_encoder_layers": 0, "label_smoothing": 0.1, "use_fixed_batch_size": false, "optimizer": "Adam", "moe_k": 2, "self_attention_type": "dot_product", "learning_rate_decay_scheme": "noam", "sampling_temp": 1.0, "kernel_height": 3, "use_pad_remover": true, "batch_size": 4096, "max_relative_position": 0, "force_full_predict": false, "min_length_bucket": 8, "layer_prepostprocess_dropout": 0.1, "eval_run_autoregressive": false, "shared_embedding_and_softmax_weights": true, "symbol_modality_num_shards": 16, "dropout": 0.2, "compress_steps": 0, "parameter_attention_key_channels": 0, "length_bucket_step": 1.1, "kernel_width": 1, "hidden_size": 16, "num_decoder_layers": 0, "input_modalities": "default", "filter_size": 8, "optimizer_adam_beta2": 0.98, "scheduled_sampling_warmup_steps": 50000, "norm_type": "layer", "min_length": 0, "moe_num_experts": 64, "multiply_embedding_mode": "sqrt_depth", "max_input_seq_length": 0, "learning_rate_warmup_steps": 8000, "proximity_bias": false, "ffn_layer": "dense_relu_dense", "initializer_gain": 1.0, "layer_postprocess_sequence": "da", "moe_hidden_sizes": "2048", "optimizer_adam_epsilon": 1e-09, "norm_epsilon": 1e-06}
+{"daisy_chain_variables": true, "optimizer_adam_beta1": 0.9, "scheduled_sampling_prob": 0.0, "num_hidden_layers": 2, "moe_loss_coef": 0.01, "max_target_seq_length": 0, "clip_grad_norm": 0.0, "pos": "timing", "scheduled_sampling_gold_mixin_prob": 0.5, "initializer": "uniform_unit_scaling", "grad_noise_scale": 0.0, "optimizer_momentum_momentum": 0.9, "nbr_decoder_problems": 1, "attention_key_channels": 0, "eval_drop_long_sequences": false, "learning_rate_cosine_cycle_steps": 250000, "prepend_mode": "none", "weight_decay": 0.0, "symbol_modality_skip_top": false, "weight_noise": 0.0, "target_modality": "default", "attention_dropout": 0.1, "parameter_attention_value_channels": 0, "factored_logits": false, "relu_dropout": 0.1, "no_data_parallelism": false, "layer_preprocess_sequence": "n", "sampling_method": "argmax", "learning_rate": 0.2, "num_heads": 2, "max_length": 256, "summarize_grads": false, "attention_value_channels": 0, "num_encoder_layers": 0, "label_smoothing": 0.1, "use_fixed_batch_size": false, "optimizer": "adam", "moe_k": 2, "self_attention_type": "dot_product", "learning_rate_decay_scheme": "noam", "sampling_temp": 1.0, "kernel_height": 3, "use_pad_remover": true, "batch_size": 4096, "max_relative_position": 0, "force_full_predict": false, "min_length_bucket": 8, "layer_prepostprocess_dropout": 0.1, "eval_run_autoregressive": false, "shared_embedding_and_softmax_weights": true, "symbol_modality_num_shards": 16, "dropout": 0.2, "compress_steps": 0, "parameter_attention_key_channels": 0, "length_bucket_step": 1.1, "kernel_width": 1, "hidden_size": 16, "num_decoder_layers": 0, "input_modalities": "default", "filter_size": 8, "optimizer_adam_beta2": 0.98, "scheduled_sampling_warmup_steps": 50000, "norm_type": "layer", "min_length": 0, "moe_num_experts": 64, "multiply_embedding_mode": "sqrt_depth", "max_input_seq_length": 0, "learning_rate_warmup_steps": 8000, "proximity_bias": false, "ffn_layer": "dense_relu_dense", "initializer_gain": 1.0, "layer_postprocess_sequence": "da", "moe_hidden_sizes": "2048", "optimizer_adam_epsilon": 1e-09, "norm_epsilon": 1e-06}
diff --git a/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index 01b6145ba..9c44d28d9 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -326,7 +326,7 @@ def adafactor_optimizer_from_hparams(hparams, lr):
   Raises:
     ValueError: on illegal values
   """
-  if hparams.optimizer_adafactor_decay_type == "Adam":
+  if hparams.optimizer_adafactor_decay_type == "adam":
     decay_rate = adafactor_decay_rate_adam(
         hparams.optimizer_adafactor_beta2)
   elif hparams.optimizer_adafactor_decay_type == "pow":
diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index bf037a876..709f993e7 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -90,7 +90,7 @@ def legacy_learning_rate_schedule(hparams):
     warmup = _learning_rate_warmup(warmup_steps, hparams=hparams)
     decay = _learning_rate_decay(hparams, warmup_steps)
     ret = tf.where(step_num < warmup_steps, warmup, decay)
-  optimizer_correction = 0.002 if "Adam" in hparams.optimizer else 1.0
+  optimizer_correction = 0.002 if "adam" in hparams.optimizer else 1.0
   tf.logging.info("Base learning rate: %f", hparams.learning_rate)
   return ret * optimizer_correction * hparams.learning_rate
 
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 5c3b5b3bf..641f60601 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -94,7 +94,7 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
   return train_op
 
 
-@registry.register_optimizer("Adam")
+@registry.register_optimizer
 def adam(learning_rate, hparams):
   # We change the default epsilon for Adam.
   # Using LazyAdam as it's much faster for large vocabulary embeddings.
@@ -105,7 +105,7 @@ def adam(learning_rate, hparams):
       epsilon=hparams.optimizer_adam_epsilon)
 
 
-@registry.register_optimizer("MultistepAdam")
+@registry.register_optimizer
 def multistep_adam(learning_rate, hparams):
   return multistep_optimizer.MultistepAdamOptimizer(
       learning_rate,
@@ -115,7 +115,7 @@ def multistep_adam(learning_rate, hparams):
       n=hparams.optimizer_multistep_accumulate_steps)
 
 
-@registry.register_optimizer("Momentum")
+@registry.register_optimizer
 def momentum(learning_rate, hparams):
   return tf.train.MomentumOptimizer(
       learning_rate,
@@ -123,14 +123,14 @@ def momentum(learning_rate, hparams):
       use_nesterov=hparams.optimizer_momentum_nesterov)
 
 
-@registry.register_optimizer("YellowFin")
+@registry.register_optimizer
 def yellow_fin(learning_rate, hparams):
   return yellowfin.YellowFinOptimizer(
       learning_rate=learning_rate,
       momentum=hparams.optimizer_momentum_momentum)
 
 
-@registry.register_optimizer("TrueAdam")
+@registry.register_optimizer
 def true_adam(learning_rate, hparams):
   return tf.train.AdamOptimizer(
       learning_rate,
@@ -139,7 +139,7 @@ def true_adam(learning_rate, hparams):
       epsilon=hparams.optimizer_adam_epsilon)
 
 
-@registry.register_optimizer("AdamW")
+@registry.register_optimizer
 def adam_w(learning_rate, hparams):
   # Openai gpt used weight decay.
   # Given the internals of AdamW, weight decay dependent on the
@@ -156,7 +156,7 @@ def adam_w(learning_rate, hparams):
       epsilon=hparams.optimizer_adam_epsilon)
 
 
-@registry.register_optimizer("Adafactor")
+@registry.register_optimizer("adafactor")
 def register_adafactor(learning_rate, hparams):
   return adafactor.adafactor_optimizer_from_hparams(hparams, learning_rate)
 
@@ -169,8 +169,11 @@ def _register_base_optimizer(key, fn):
 
 
 for k in tf.contrib.layers.OPTIMIZER_CLS_NAMES:
-  if k not in registry._OPTIMIZERS:  # pylint: disable=protected-access
+  if k not in registry.Registries.optimizers and k not in ('SGD', 'RMSProp'):
     _register_base_optimizer(k, tf.contrib.layers.OPTIMIZER_CLS_NAMES[k])
+_register_base_optimizer('sgd', tf.contrib.layers.OPTIMIZER_CLS_NAMES['SGD'])
+_register_base_optimizer(
+    'rms_prop', tf.contrib.layers.OPTIMIZER_CLS_NAMES['RMSProp'])
 
 
 class ConditionalOptimizer(tf.train.Optimizer):
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index d8683f234..56dc11f50 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -13,510 +13,522 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Registry for models, hyperparameter settings, problem types, and datasets.
+"""Registry for
+* models
+* hyperparameter settings
+* ranged hyperparameter settings
+* problem types
+* attacks
+* attack parameters
+* pruning parameters
+* pruning strategies
+* optimizers
+
+`Registries` contains the `Registry` objects used throughout tensor2tensor.
+
+New functions and classes can be registered using `.register`. The can be
+accessed/queried similar to dictionaries, keyed by default by `snake_case`
+equivalents.
+```
+@Registries.models.register
+class MyModel(T2TModel):
+  ...
+
+'my_model' in Registries.models  # True
+for k in Registries.models:
+  print(k)  # prints 'my_model'
+model = Registries.models['my_model'](constructor_arg)
+```
 
+`Registry`s for functions which take no arguments will return the
+result of evaluating those functions (though this is not the default behaviour
+of the `Registry` class in general).
+```
+@Registries.attacks.register
+def my_attack():
+  ...
+
+my_attack_obj = Registries.attacks['my_attack']
+my_attack_obj()  # TypeError: 'Attack' object is not callable
+```
+
+#### Legacy Support
 Define a new model by subclassing T2TModel and register it:
 
 ```
-@registry.register_model
+@register_model
 class MyModel(T2TModel):
   ...
 ```
 
-Access by snake-cased name: `registry.model("my_model")`. If you're using
+Access by snake-cased name: `model("my_model")`. If you're using
 `t2t_trainer.py`, you can pass on the command-line: `--model=my_model`.
 
-See all the models registered: `registry.list_models()`.
+See all the models registered: `list_models()`.
 
 For hyperparameter sets:
-  * Register: `registry.register_hparams`
-  * List: `registry.list_hparams`
-  * Retrieve by name: `registry.hparams`
+  * Register: `register_hparams`
+  * List: `list_hparams`
+  * Retrieve by name: `hparams`
   * Command-line flag in `t2t_trainer.py`: `--hparams_set=name`
 
 For hyperparameter ranges:
-  * Register: `registry.register_ranged_hparams`
-  * List: `registry.list_ranged_hparams`
-  * Retrieve by name: `registry.ranged_hparams`
+  * Register: `register_ranged_hparams`
+  * List: `list_ranged_hparams`
+  * Retrieve by name: `ranged_hparams`
   * Command-line flag in `t2t_trainer.py`: `--hparams_range=name`
 """
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from tensorflow import logging
+import tensorflow.contrib.framework as framework
+from tensorflow.python.util import tf_inspect as inspect
 import collections
 
 from tensor2tensor.utils import misc_utils
-import tensorflow as tf
-from tensorflow.python.util import tf_inspect as inspect
-
-_ATTACKS = {}
-_ATTACK_PARAMS = {}
-_HPARAMS = {}
-_MODELS = {}
-_PROBLEMS = {}
-_PRUNING_PARAMS = {}
-_PRUNING_STRATEGY = {}
-_RANGED_HPARAMS = {}
 
-# Key: registry name, Value: Registry
-_GENERIC_REGISTRIES = {}
-Registry = collections.namedtuple(
-    "_Registry", ["register", "get", "list", "registry"])
 
+def default_name(class_or_fn):
+  """Default name for a class or function.
 
-def registry(registry_name):
-  """Returns `Registry` created by `create_registry`."""
-  if registry_name not in _GENERIC_REGISTRIES:
-    raise KeyError("No registry named %s. Available:\n%s" % (
-        registry_name, sorted(_GENERIC_REGISTRIES)))
-  return _GENERIC_REGISTRIES[registry_name]
-
-
-def create_registry(registry_name):
-  """Create a generic object registry.
+  This is the naming function by default for registries expecting classes or
+  functions.
 
   Args:
-    registry_name: str, name of the object registry.
+    class_or_fn: class or function to be named.
 
   Returns:
-    `Registry` that contains functions for register (decorator), get, and list.
-
-  Raises:
-    KeyError: if `registry_name` is a pre-existing registry.
+    Default name for registration.
   """
-  if registry_name in _GENERIC_REGISTRIES:
-    raise KeyError(
-        "Registry %s already exists." % registry_name)
-
-  registry_ = {}
-
-  def register(name):
-    """Returns decorator to register an object."""
-
-    def register_dec(obj):
-      if name in registry_:
-        raise KeyError(
-            "Registry %s already contains key %s." % (registry_name, name))
-      registry_[name] = obj
-      return obj
+  return misc_utils.camelcase_to_snakecase(class_or_fn.__name__)
 
-    return register_dec
 
-  def get(name):
-    if name not in registry_:
-      raise KeyError(
-          "Registry %s contains no object named %s" % (registry_name, name))
-    return registry_[name]
+class Registry(object):
+  """Dict-like class for managing function registrations.
 
-  def list_registry():
-    return sorted(registry_)
+  ```python
+  my_registry = Registry("custom_name")
 
-  registry_obj = Registry(
-      register=register,
-      get=get,
-      list=list_registry,
-      registry=registry_,
-  )
-  _GENERIC_REGISTRIES[registry_name] = registry_obj
-  return registry_obj
+  @my_registry.register
+  def my_func():
+    pass
 
+  @my_registry.register()
+  def another_func():
+    pass
 
-def _reset():
-  for ctr in [_MODELS, _HPARAMS, _RANGED_HPARAMS, _ATTACK_PARAMS]:
-    ctr.clear()
+  @my_registry.register("non_default_name")
+  def third_func(x, y, z):
+    pass
 
+  def foo():
+    pass
 
-def default_name(obj_class):
-  """Convert a class name to the registry's default name for the class.
+  my_registry.register()(foo)
+  my_registry.register("baz")(lambda (x, y): x + y)
+  my_register.register("bar")
 
-  Args:
-    obj_class: the name of a class
+  print(list(my_registry))
+  # ["my_func", "another_func", "non_default_name", "foo", "baz"]
+  # (order may vary)
+  print(my_registry["non_default_name"] is third_func)  # True
+  print("third_func" in my_registry)                    # False
+  print("bar" in my_registry)                           # False
+  my_registry["non-existent_key"]                       # raises KeyError
+  ```
 
-  Returns:
-    The registry's default name for the class.
+  Optional validation, on_set callback and value transform also supported.
+  See `__init__` doc.
   """
-  return misc_utils.camelcase_to_snakecase(obj_class.__name__)
-
 
-def default_object_name(obj):
-  """Convert an object to the registry's default name for the object class.
+  def __init__(
+      self, registry_name, default_key_fn=default_name, validator=None,
+      on_set=None, value_transformer=(lambda k, v: v)):
+    """Construct a new registry.
 
-  Args:
-    obj: an object instance
+    Args:
+      registry_name: str identifier for the given registry. Used in error msgs.
+      default_key_fn (optional): function mapping value -> key for registration
+        when a key is not provided
+      validator (optional): if given, this is run before setting a given
+        (key, value) pair. Accepts (key, value) and should raise if there is a
+        problem. Overwriting existing keys is not allowed and is checked
+        separately. Values are also checked to be callable separately.
+      on_set (optional): callback function accepting (key, value) pair
+        which is run after an item is successfully set.
+      value_transformer (optional): if run, `__getitem__` will return
+        value_transformer(key, registered_value).
+    """
+    self._registry = {}
+    self._name = registry_name
+    self._default_key_fn = default_key_fn
+    self._validator = validator
+    self._on_set = on_set
+    self._value_transformer = value_transformer
+
+  def default_key(self, value):
+    """Default key used when key not provided. Uses function from __init__."""
+    return self._default_key_fn(value)
+
+  @property
+  def name(self):
+    return self._name
+
+  def validate(self, key, value):
+    """Validation function run before setting. Uses function from __init__."""
+    if self._validator is not None:
+      self._validator(key, value)
+
+  def on_set(self, key, value):
+    """Callback called on successful set. Uses function from __init__."""
+    if self._on_set is not None:
+      self._on_set(key, value)
+
+  def __setitem__(self, key, value):
+    """
+    Validate, set, and (if successfull) call `on_set` for the given item.
 
-  Returns:
-    The registry's default name for the class of the object.
-  """
-  return default_name(obj.__class__)
+    Args:
+      key: key to store value under. If `None`, `self.default_key(value)` is
+        used.
+      value: callable stored under the given key.
 
+    Returns:
+      `None`
+    """
+    if key is None:
+      key = self.default_key(value)
+    if key in self:
+      raise KeyError("key %s already registered in registry %s"
+                     % (key, self._name))
+    if not callable(value):
+      raise ValueError("value must be callable")
+    self.validate(key, value)
+    self._registry[key] = value
+    self.on_set(key, value)
+
+  def register(self, key_or_value=None):
+    """Decorator to register a function, or registration itself.
+
+    This is primarily intended for use as a decorator, either with or without
+    a key/parentheses.
+    ```python
+    @my_registry.register('key1')
+    def value_fn(x, y, z):
+      pass
+
+    @my_registry.register()
+    def another_fn(x, y):
+      pass
+
+    @my_registry.register
+    def third_func():
+      pass
+    ```
+
+    Note if key_or_value is provided as a non-callable, registration only
+    occurs once the returned callback is called with a callable as its only
+    argument.
+    ```python
+    callback = my_registry.register('different_key')
+    'different_key' in my_registry  # False
+    callback(lambda (x, y): x + y)
+    'different_key' in my_registry  # True
+    ```
 
-def register_model(name=None):
-  """Register a model. name defaults to class name snake-cased."""
+    Args:
+      key_or_value (optional): key to access the registered value with, or the
+        function itself. If `None` (default), `self.default_key` will be called
+        on `value` once the returned callback is called with `value` as the
+        only arg. If `key_or_value` is itself callable, it is assumed to be the
+        value and the key is given by `self.default_key(key)`.
 
-  def decorator(model_cls, registration_name=None):
-    """Registers & returns model_cls with registration_name or default name."""
-    model_name = registration_name or default_name(model_cls)
-    if model_name in _MODELS and not tf.executing_eagerly():
-      raise LookupError("Model %s already registered." % model_name)
-    model_cls.REGISTERED_NAME = model_name
-    _MODELS[model_name] = model_cls
-    return model_cls
+    Returns:
+      decorated callback, or callback generated a decorated function.
+    """
+    def decorator(value, key):
+      self[key] = value
+      return value
 
-  # Handle if decorator was used without parens
-  if callable(name):
-    model_cls = name
-    return decorator(model_cls, registration_name=default_name(model_cls))
+    # Handle if decorator was used without parens
+    if callable(key_or_value):
+      return decorator(value=key_or_value, key=None)
+    else:
+      return lambda value: decorator(value, key=key_or_value)
 
-  return lambda model_cls: decorator(model_cls, name)
+  def __getitem__(self, key):
+    if key not in self:
+      raise KeyError("%s never registered with registry %s. Available:\n %s" %
+                     (key, self.name,
+                      display_list_by_prefix(sorted(self), 4)))
+    value = self._registry[key]
+    return self._value_transformer(key, value)
 
+  def __contains__(self, key):
+    return key in self._registry
 
-def model(name):
-  if name not in _MODELS:
-    raise LookupError("Model %s never registered.  Available models:\n %s" %
-                      (name, "\n".join(list_models())))
+  def keys(self):
+    return self._registry.keys()
 
-  return _MODELS[name]
+  def values(self):
+    return (self[k] for k in self)       # complicated because of transformer
 
+  def items(self):
+    return ((k, self[k]) for k in self)  # complicated because of transformer
 
-def list_models():
-  return list(sorted(_MODELS))
+  def __iter__(self):
+    return iter(self._registry)
 
+  def __len__(self):
+    return len(self._registry)
 
-_OPTIMIZERS = {}
+  def _clear(self):
+    self._registry.clear()
 
+  def get(self, key, d=None):
+    return self[key] if key in self else d
 
-def register_optimizer(name=None):
-  """Register an optimizer. name defaults to upper camel case of fn name."""
 
-  def default_opt_name(opt_fn):
-    return misc_utils.snakecase_to_camelcase(default_name(opt_fn))
+def _on_model_set(k, v):
+  v.REGISTERED_NAME = k
 
-  def decorator(opt_fn, registration_name):
-    """Registers and returns optimizer_fn with registration_name or default."""
-    if registration_name is None:
-      registration_name = default_opt_name(opt_fn)
 
-    if registration_name in _OPTIMIZERS and not tf.executing_eagerly():
-      raise LookupError("Optimizer %s already registered." % registration_name)
-    args, varargs, keywords, _ = inspect.getargspec(opt_fn)
+def _nargs_validator(nargs, message):
+  if message is None:
+    message = "Registered function must take exactly %d arguments" % nargs
+  def f(key, value):
+    spec = inspect.getfullargspec(value)
+    if (len(spec.args) != nargs or
+        spec.varargs is not None or
+        spec.varkw is not None):
+      raise ValueError(message)
 
-    if len(args) != 2 or varargs is not None or keywords is not None:
-      raise ValueError("Optimizer registration function must take two "
-                       "arguments: learning_rate (float) and "
-                       "hparams (HParams).")
-    _OPTIMIZERS[registration_name] = opt_fn
-    return opt_fn
+  return f
 
-  if callable(name):
-    opt_fn = name
-    registration_name = default_opt_name(opt_fn)
-    return decorator(opt_fn, registration_name=registration_name)
 
-  return lambda opt_fn: decorator(opt_fn, name)
+ProblemSpec = collections.namedtuple(
+    "ProblemSpec", ["base_name", "was_reversed", "was_copy"])
 
 
-def optimizer(name):
-  if name not in _OPTIMIZERS:
-    raise LookupError("Optimizer %s never registered. "
-                      "Available optimizers:\n %s"
-                      % (name, "\n".join(list_optimizers())))
-  return _OPTIMIZERS[name]
+def parse_problem_name(name):
+  """Determines if problem_name specifies a copy and/or reversal.
 
+  Args:
+    name: str, problem name, possibly with suffixes.
 
-def list_optimizers():
-  return list(sorted(_OPTIMIZERS))
+  Returns:
+    ProblemSpec: namedtuple with ["base_name", "was_reversed", "was_copy"]
 
+  Raises:
+    ValueError if name contains multiple suffixes of the same type
+      ('_rev' or '_copy'). One of each is ok.
+  """
+  # Recursively strip tags until we reach a base name.
+  if name.endswith("_rev"):
+    base, was_rev, was_copy = parse_problem_name(name[:-4])
+    if was_rev:
+      # duplicate rev
+      raise ValueError(
+          "Invalid problem name %s: multiple '_rev' instances" % name)
+    return ProblemSpec(base, True, was_copy)
+  elif name.endswith("_copy"):
+    base, was_reversed, was_copy = parse_problem_name(name[:-5])
+    if was_copy:
+      raise ValueError(
+          "Invalid problem_name %s: multiple '_copy' instances" % name)
+    return ProblemSpec(base, was_reversed, True)
+  else:
+    return ProblemSpec(name, False, False)
+
+
+def get_problem_name(base_name, was_reversed=False, was_copy=False):
+  """
+  Construct a problem name from base and reversed/copy options.
 
-def register_hparams(name=None):
-  """Register an HParams set. name defaults to function name snake-cased."""
+  Inverse of `parse_problem_name`.
 
-  def decorator(hp_fn, registration_name=None):
-    """Registers & returns hp_fn with registration_name or default name."""
-    hp_name = registration_name or default_name(hp_fn)
-    if hp_name in _HPARAMS and not tf.executing_eagerly():
-      raise LookupError("HParams set %s already registered." % hp_name)
-    _HPARAMS[hp_name] = hp_fn
-    return hp_fn
+  Args:
+    base_name: base problem name. Should not end in "_rev" or "_copy"
+    was_reversed: if the problem is to be reversed
+    was_copy: if the problem is to be copied
 
-  # Handle if decorator was used without parens
-  if callable(name):
-    hp_fn = name
-    return decorator(hp_fn, registration_name=default_name(hp_fn))
+  Returns:
+    string name consistent with use with `parse_problem_name`.
 
-  return lambda hp_fn: decorator(hp_fn, name)
+  Raises:
+    ValueError if `base_name` ends with "_rev" or "_copy"
+  """
+  if any(base_name.endswith(suffix) for suffix in ("_rev", "_copy")):
+    raise ValueError("`base_name` cannot end in '_rev' or '_copy'")
+  name = base_name
+  if was_copy:
+    name = "%s_copy" % name
+  if was_reversed:
+    name = "%s_rev" % name
+  return name
+
+
+def _problem_name_validator(k, v):
+  if parse_problem_name(k).base_name != k:
+    raise KeyError(
+        "Invalid problem name: cannot end in %s or %s" % ("_rev", "_copy"))
 
 
-def hparams(name):
-  """Retrieve registered hparams by name."""
-  if name not in _HPARAMS:
-    error_msg = "HParams set %s never registered. Sets registered:\n%s"
-    raise LookupError(
-        error_msg % (name,
-                     display_list_by_prefix(list_hparams(), starting_spaces=4)))
-  hp = _HPARAMS[name]()
-  if hp is None:
-    raise TypeError("HParams %s is None. Make sure the registered function "
-                    "returns the HParams object." % name)
-  return hp
+def _on_problem_set(k, v):
+  v.name = k
 
 
-def list_hparams(prefix=None):
-  if prefix:
-    return [name for name in _HPARAMS if name.startswith(prefix)]
-  return list(_HPARAMS)
+def _call_value(k, v):
+  return v()
 
 
-def register_ranged_hparams(name=None):
-  """Register a RangedHParams set. name defaults to fn name snake-cased."""
+def _hparams_value_transformer(key, value):
+  out = value()
+  if out is None:
+    raise TypeError("HParams %s is None. Make sure the registered function "
+                    "returns the HParams object" % key)
+  return out
 
-  def decorator(rhp_fn, registration_name=None):
-    """Registers & returns hp_fn with registration_name or default name."""
-    rhp_name = registration_name or default_name(rhp_fn)
-    if rhp_name in _RANGED_HPARAMS:
-      raise LookupError("RangedHParams set %s already registered." % rhp_name)
-    # Check that the fn takes a single argument
-    args, varargs, keywords, _ = inspect.getargspec(rhp_fn)
-    if len(args) != 1 or varargs is not None or keywords is not None:
-      raise ValueError("RangedHParams set function must take a single "
-                       "argument, the RangedHParams object.")
 
-    _RANGED_HPARAMS[rhp_name] = rhp_fn
-    return rhp_fn
+class Registries(object):
+  """Object holding `Registry` objects."""
+  def __init__(self):
+    raise RuntimeError("Registries is not intended to be instantiated")
 
-  # Handle if decorator was used without parens
-  if callable(name):
-    rhp_fn = name
-    return decorator(rhp_fn, registration_name=default_name(rhp_fn))
 
-  return lambda rhp_fn: decorator(rhp_fn, name)
+  models = Registry("models", on_set=_on_model_set)
 
+  optimizers = Registry(
+      "optimizers",
+      validator=_nargs_validator(
+          2,
+          "Registered optimizer functions must take exactly two arguments: "
+          "learning_rate (float) and hparams (HParams)."))
 
-def ranged_hparams(name):
-  if name not in _RANGED_HPARAMS:
-    raise LookupError("RangedHParams set %s never registered." % name)
-  return _RANGED_HPARAMS[name]
+  hparams = Registry(
+      "hparams", value_transformer=_hparams_value_transformer)
 
+  ranged_hparams = Registry(
+      "ranged_hparams", validator=_nargs_validator(
+          1,
+          "Registered ranged_hparams functions must take a single argument, "
+          "the RangedHParams object."))
 
-def list_ranged_hparams():
-  return list(_RANGED_HPARAMS)
+  problems = Registry(
+      "problems", validator=_problem_name_validator, on_set=_on_problem_set)
 
+  attacks = Registry(
+      "attacks", value_transformer=_call_value)
 
-def register_problem(name=None):
-  """Register a Problem. name defaults to cls name snake-cased."""
+  attack_params = Registry(
+      "attack_params", value_transformer=_call_value)
 
-  def decorator(p_cls, registration_name=None):
-    """Registers & returns p_cls with registration_name or default name."""
-    p_name = registration_name or default_name(p_cls)
-    if p_name in _PROBLEMS and not tf.executing_eagerly():
-      raise LookupError("Problem %s already registered." % p_name)
+  pruning_params = Registry(
+      "pruning_params", value_transformer=_call_value)
 
-    _PROBLEMS[p_name] = p_cls
-    p_cls.name = p_name
-    return p_cls
+  pruning_strategies = Registry("pruning_strategies")
 
-  # Handle if decorator was used without parens
-  if callable(name):
-    p_cls = name
-    return decorator(p_cls, registration_name=default_name(p_cls))
+  layers = Registry("layers", validator=_nargs_validator(
+      2,
+      "Registered layer functions must take exaction two arguments: "
+      "hparams (HParams) and prefix (str)."
+  ))
 
-  return lambda p_cls: decorator(p_cls, name)
+# consistent version of old API
+model = Registries.models.__getitem__
+list_models = lambda: sorted(Registries.models)
+register_model = Registries.models.register
 
+# optimizer = optimizer_registry.__getitem__
+def optimizer(name):
+  """
+  Get pre-registered optimizer keyed by name.
 
-def problem(name):
-  """Retrieve a problem by name."""
+  `name` should be snake case, though SGD -> sgd, RMSProp -> rms_prop and
+  UpperCamelCase -> snake_case conversions included for legacy support.
 
-  def parse_problem_name(problem_name):
-    """Determines if problem_name specifies a copy and/or reversal.
+  Args:
+    name: name of optimizer used in registration. This should be a snake case
+      identifier, though others supported for legacy reasons.
+  """
+  rest = ("Please update `registry.optimizer` callsite "
+          "(likely due to a `HParams.optimizer` value)")
+  if name == "SGD":
+    name = "sgd"
+    logging.warning("'SGD' optimizer now keyed by 'sgd'. %s" % rest)
+  elif name == 'rms_prop':
+    name = 'rms_prop'
+    logging.warning("'RMSProp' optimizer now keyed by 'rms_prop'. %s" % rest)
+  else:
+    snake_name = misc_utils.camelcase_to_snakecase(name)
+    if name != snake_name:
+      logging.warning(
+          "optimizer names now keyed by snake_case names. %s" % rest)
+      name = snake_name
+  return Registries.optimizers[name]
+
+list_optimizers = lambda: sorted(Registries.optimizers)
+register_optimizer = Registries.optimizers.register
+
+hparams = Registries.hparams.__getitem__
+list_hparams = lambda: sorted(Registries.hparams)
+register_hparams = Registries.hparams.register
+
+ranged_hparams = Registries.ranged_hparams.__getitem__
+list_ranged_hparams = lambda: sorted(Registries.ranged_hparams)
+register_ranged_hparams = Registries.ranged_hparams.register
+
+base_problem = Registries.problems.__getitem__
+list_base_problems = lambda: sorted(Registries.problems)
+register_base_problem = Registries.problems.register
+
+# list_problems won't list all rev/copy combinations,
+# so the name is slightly confusing. Similarly, register_problem will raise an
+# error if attempting to register a value with a non-base key.
+# Keeping for back-compatibility
+list_problems = list_base_problems
+register_problem = register_base_problem
+
+
+def problem(problem_name):
+  """Get possibly copied/reversed problem registered in `base_registry`.
 
-    Args:
-      problem_name: str, problem name, possibly with suffixes.
+  Args:
+    problem_name: string problem name. See `parse_problem_name`.
 
-    Returns:
-      base_name: A string with the base problem name.
-      was_reversed: A boolean.
-      was_copy: A boolean.
-    """
-    # Recursively strip tags until we reach a base name.
-    if problem_name.endswith("_rev"):
-      base, _, was_copy = parse_problem_name(problem_name[:-4])
-      return base, True, was_copy
-    elif problem_name.endswith("_copy"):
-      base, was_reversed, _ = parse_problem_name(problem_name[:-5])
-      return base, was_reversed, True
-    else:
-      return problem_name, False, False
+  Returns:
+    possibly reversed/copied version of base problem registered in the given
+    registry.
+  """
+  spec = parse_problem_name(problem_name)
+  return Registries.problems[spec.base_name](
+      was_copy=spec.was_copy, was_reversed=spec.was_reversed)
 
-  base_name, was_reversed, was_copy = parse_problem_name(name)
-
-  if base_name not in _PROBLEMS:
-    all_problem_names = list_problems()
-    error_lines = ["%s not in the set of supported problems:" % base_name
-                  ] + all_problem_names
-    error_msg = "\n  * ".join(error_lines)
-    raise LookupError(error_msg)
-  return _PROBLEMS[base_name](was_reversed=was_reversed, was_copy=was_copy)
-
-
-def list_problems():
-  return sorted(list(_PROBLEMS))
 
+attack = Registries.attacks.__getitem__
+list_attacks = lambda: sorted(Registries.attacks)
+register_attack = Registries.attacks.register
 
-def register_attack(name=None):
-  """Register an attack HParams set. Same behaviour as register_hparams."""
-
-  def decorator(attack_fn, registration_name=None):
-    """Registers & returns attack_fn with registration_name or default name."""
-    attack_name = registration_name or default_name(attack_fn)
-    if attack_name in _ATTACKS and not tf.executing_eagerly():
-      raise LookupError("Attack %s already registered." % attack_name)
-    _ATTACKS[attack_name] = attack_fn
-    return attack_fn
-
-  # Handle if decorator was used without parens
-  if callable(name):
-    attack_fn = name
-    return decorator(attack_fn, registration_name=default_name(attack_fn))
-
-  return lambda attack_fn: decorator(attack_fn, name)
+attack_params = Registries.attack_params.__getitem__
+list_attack_params = lambda: sorted(Registries.attack_params)
+register_attack_params = Registries.attack_params.register
 
+pruning_params = Registries.pruning_params.__getitem__
+list_pruning_params = lambda: sorted(Registries.pruning_params)
+register_pruning_params = Registries.pruning_params.register
 
-def attacks(name):
-  """Retrieve registered attack by name."""
-  if name not in _ATTACKS:
-    error_msg = "Attack %s never registered. Sets registered:\n%s"
-    raise LookupError(
-        error_msg % (name,
-                     display_list_by_prefix(list_attacks(), starting_spaces=4)))
-  attack = _ATTACKS[name]()
-  if attack is None:
-    raise TypeError(
-        "Attack %s is None. Make sure the registered function returns a "
-        "`cleverhans.attack.Attack` object." % name)
-  return attack
-
-
-def list_attacks(prefix=None):
-  if prefix:
-    return [name for name in _ATTACKS if name.startswith(prefix)]
-  return list(_ATTACKS)
-
-
-def register_attack_params(name=None):
-  """Register an attack HParams set. Same behaviour as register_hparams."""
-
-  def decorator(ap_fn, registration_name=None):
-    """Registers & returns ap_fn with registration_name or default name."""
-    ap_name = registration_name or default_name(ap_fn)
-    if ap_name in _ATTACK_PARAMS and not tf.executing_eagerly():
-      raise LookupError("Attack HParams set %s already registered." % ap_name)
-    _ATTACK_PARAMS[ap_name] = ap_fn
-    return ap_fn
-
-  # Handle if decorator was used without parens
-  if callable(name):
-    ap_fn = name
-    return decorator(ap_fn, registration_name=default_name(ap_fn))
-
-  return lambda ap_fn: decorator(ap_fn, name)
-
-
-def attack_params(name):
-  """Retrieve registered aparams by name."""
-  if name not in _ATTACK_PARAMS:
-    error_msg = "Attack HParams set %s never registered. Sets registered:\n%s"
-    raise LookupError(
-        error_msg %
-        (name, display_list_by_prefix(list_attack_params(), starting_spaces=4)))
-  ap = _ATTACK_PARAMS[name]()
-  if ap is None:
-    raise TypeError("Attack HParams %s is None. Make sure the registered "
-                    "function returns the HParams object." % name)
-  return ap
-
-
-def list_attack_params(prefix=None):
-  if prefix:
-    return [name for name in _ATTACK_PARAMS if name.startswith(prefix)]
-  return list(_ATTACK_PARAMS)
-
-
-def register_pruning_params(name=None):
-  """Register an pruning HParams set. Same behaviour as register_hparams."""
-
-  def decorator(pp_fn, registration_name=None):
-    """Registers & returns pp_fn with registration_name or default name."""
-    pp_name = registration_name or default_name(pp_fn)
-    if pp_name in _PRUNING_PARAMS and not tf.executing_eagerly():
-      raise LookupError("Pruning HParams set %s already registered." % pp_name)
-    _PRUNING_PARAMS[pp_name] = pp_fn
-    return pp_fn
-
-  # Handle if decorator was used without parens
-  if callable(name):
-    pp_fn = name
-    return decorator(pp_fn, registration_name=default_name(pp_fn))
-
-  return lambda pp_fn: decorator(pp_fn, name)
-
-
-def pruning_params(name):
-  """Retrieve registered pruning params by name."""
-  if name not in _PRUNING_PARAMS:
-    error_msg = "Pruning HParams set %s never registered. Sets registered:\n%s"
-    raise LookupError(error_msg % (
-        name, display_list_by_prefix(list_pruning_params(), starting_spaces=4)))
-  pp = _PRUNING_PARAMS[name]()
-  if pp is None:
-    raise TypeError("Pruning HParams %s is None. Make sure the registered "
-                    "function returns the HParams object." % name)
-  return pp
-
-
-def list_pruning_params(prefix=None):
-  if prefix:
-    return [name for name in _PRUNING_PARAMS if name.startswith(prefix)]
-  return list(_PRUNING_PARAMS)
-
-
-def register_pruning_strategy(name=None):
-  """Register an pruning strategy. Same behaviour as register_hparams."""
-
-  def decorator(ps_fn, registration_name=None):
-    """Registers & returns ps_fn with registration_name or default name."""
-    ps_name = registration_name or default_name(ps_fn)
-    if ps_name in _PRUNING_STRATEGY and not tf.executing_eagerly():
-      raise LookupError("Pruning strategy %s already registered." % ps_name)
-    _PRUNING_STRATEGY[ps_name] = ps_fn
-    return ps_fn
-
-  # Handle if decorator was used without parens
-  if callable(name):
-    ps_fn = name
-    return decorator(ps_fn, registration_name=default_name(ps_fn))
-
-  return lambda ps_fn: decorator(ps_fn, name)
-
-
-def pruning_strategies(name):
-  """Retrieve registered pruning strategies by name."""
-  if name not in _PRUNING_STRATEGY:
-    error_msg = "Pruning strategy set %s never registered. Sets registered:\n%s"
-    raise LookupError(
-        error_msg % (name,
-                     display_list_by_prefix(
-                         list_pruning_strategies(), starting_spaces=4)))
-  ps = _PRUNING_STRATEGY[name]
-  if ps is None:
-    raise TypeError("Pruning strategy %s is None. Make sure to register the "
-                    "function." % name)
-  return ps
+pruning_strategy = Registries.pruning_strategies.__getitem__
+list_pruning_strategies = lambda: sorted(Registries.pruning_strategies)
+register_pruning_strategy = Registries.pruning_strategies.register
 
 
-def list_pruning_strategies(prefix=None):
-  if prefix:
-    return [name for name in _PRUNING_STRATEGY if name.startswith(prefix)]
-  return list(_PRUNING_STRATEGY)
+# deprecated functions - plurals inconsistent with rest
+# deprecation decorators added 2019-01-25
+attacks = framework.deprecated(None, "Use registry.attack")(attack)
+pruning_strategies = framework.deprecated(
+    None, "Use registry.pruning_strategy")(pruning_strategy)
 
 
 def display_list_by_prefix(names_list, starting_spaces=0):
@@ -551,6 +563,9 @@ def help_string():
   Problems:
 %s
 
+  Optimizers:
+%s
+
   Attacks:
 %s
 
@@ -563,16 +578,17 @@ def help_string():
   Pruning Strategies:
 %s
 """
-  m, hp, rhp, probs, atks, ap, pp, ps = [
+  lists = tuple(
       display_list_by_prefix(entries, starting_spaces=4) for entries in [
           list_models(),
           list_hparams(),
           list_ranged_hparams(),
-          list_problems(),
+          list_base_problems(),
+          list_optimizers(),
           list_attacks(),
           list_attack_params(),
           list_pruning_params(),
           list_pruning_strategies(),
       ]
-  ]
-  return help_str % (m, hp, rhp, probs, atks, ap, pp, ps)
+  )
+  return help_str % lists
diff --git a/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index c2515b3ba..dc1cdb579 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -26,11 +26,72 @@
 
 # pylint: disable=unused-variable
 
+class RegistryClassTest(tf.test.TestCase):
+  """Test of base registry.Registry class."""
+
+  def testGetterSetter(self):
+    r = registry.Registry("test_registry")
+    r["hello"] = lambda: "world"
+    r["a"] = lambda: "b"
+    self.assertEqual(r["hello"](), "world")
+    self.assertEqual(r["a"](), "b")
+
+  def testDefaultKeyFn(self):
+    r = registry.Registry("test", default_key_fn=lambda x: x().upper())
+    r.register()(lambda: "hello")
+    self.assertEqual(r["HELLO"](), "hello")
+
+  def testNoKeyProvided(self):
+    r = registry.Registry("test")
+    def f():
+      return 3
+    r.register(f)
+    self.assertEqual(r['f'](), 3)
+
+  def testMembership(self):
+    r = registry.Registry("test_registry")
+    r["a"] = lambda: None
+    r["b"] = lambda: 4
+    self.assertTrue("a" in r)
+    self.assertTrue("b" in r)
+
+  def testIteration(self):
+    r = registry.Registry("test_registry")
+    r["a"] = lambda: None
+    r["b"] = lambda: 4
+    self.assertEqual(sorted(r), ["a", "b"])
+
+  def testLen(self):
+    r = registry.Registry("test_registry")
+    self.assertEqual(len(r), 0)
+    r["a"] = lambda: None
+    self.assertEqual(len(r), 1)
+    r["b"] = lambda: 4
+    self.assertEqual(len(r), 2)
+
+  def testTransformer(self):
+    r = registry.Registry(
+        "test_registry", value_transformer=lambda x, y: x + y())
+    r.register(3)(lambda: 5)
+    r.register(10)(lambda: 12)
+    self.assertEqual(r[3], 8)
+    self.assertEqual(r[10], 22)
+    self.assertEqual(set(r.values()), set((8, 22)))
+    self.assertEqual(set(r.items()), set(((3, 8), (10, 22))))
+
+  def testGet(self):
+    r = registry.Registry('test_registry', value_transformer=lambda k, v: v())
+    r["a"] = lambda: "xyz"
+    self.assertEqual(r.get("a"), "xyz")
+    self.assertEqual(r.get("a", 3), "xyz")
+    self.assertIsNone(r.get("b"))
+    self.assertEqual(r.get("b", 3), 3)
+
 
 class ModelRegistryTest(tf.test.TestCase):
 
   def setUp(self):
-    registry._reset()
+    registry.Registries.models._clear()
 
   def testT2TModelRegistration(self):
 
@@ -60,7 +121,7 @@ def model_fn():
     self.assertTrue(model is model_fn)
 
   def testUnknownModel(self):
-    with self.assertRaisesRegexp(LookupError, "never registered"):
+    with self.assertRaisesRegexp(KeyError, "never registered"):
       registry.model("not_registered")
 
   def testDuplicateRegistration(self):
@@ -69,7 +130,7 @@ def testDuplicateRegistration(self):
     def m1():
       pass
 
-    with self.assertRaisesRegexp(LookupError, "already registered"):
+    with self.assertRaisesRegexp(KeyError, "already registered"):
 
       @registry.register_model("m1")
       def m2():
@@ -88,10 +149,82 @@ def m2():
     self.assertSetEqual(set(["m1", "m2"]), set(registry.list_models()))
 
 
+class OptimizerRegistryTest(tf.test.TestCase):
+  def setUp(self):
+    registry.Registries.optimizers._clear()
+
+  def testRegistration(self):
+    @registry.register_optimizer
+    def my_optimizer(learning_rate, hparams):
+      return 3
+
+    @registry.register_optimizer('my_other_optimizer')
+    def another_optimizer(learning_rate, hparams):
+      return 5
+
+    self.assertEqual(registry.optimizer("my_optimizer"), my_optimizer)
+    self.assertEqual(
+        registry.optimizer("my_other_optimizer"), another_optimizer)
+
+  def testMembership(self):
+    @registry.register_optimizer
+    def my_optimizer(learning_rate, hparams):
+      return 3
+
+    @registry.register_optimizer("my_other_optimizer")
+    def another_optimizer(learning_rate, hparams):
+      return 5
+
+    self.assertTrue("my_optimizer" in registry.Registries.optimizers)
+    self.assertTrue("my_other_optimizer" in registry.Registries.optimizers)
+    self.assertFalse("another_optimizer" in registry.Registries.optimizers)
+    self.assertEqual(len(registry.Registries.optimizers), 2)
+
+  def testArgErrorCheck(self):
+    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
+      registry.Registries.optimizers.register('OneArgs')(lambda x: 4)
+    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
+      registry.Registries.optimizers.register('ThreeArgs')(
+          lambda x, y, z: 4)
+    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
+      registry.Registries.optimizers.register('NArgs')(lambda *args: 4)
+    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
+      registry.Registries.optimizers.register("Kwargs")(lambda **kargs: 4)
+    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
+      registry.Registries.optimizers.register("TwoAndKwargs")(
+          lambda a, b, **kargs: 4)
+
+  def testMultipleRegistration(self):
+    @registry.register_optimizer
+    def my_optimizer(learning_rate, hparams):
+      return 3
+
+    with self.assertRaisesRegexp(KeyError, "already registered"):
+
+      @registry.register_optimizer("my_optimizer")
+      def another_fn(learning_rate, hparams):
+        return 5
+
+  def testUnknownOptimizer(self):
+    with self.assertRaisesRegexp(KeyError, "never registered"):
+      registry.optimizer("not_registered_optimizer")
+
+  def testGetterSetterInterface(self):
+    def f(x, y):
+      return 3
+
+    k = 'blah'
+    registry.Registries.optimizers[k] = f
+    self.assertEqual(registry.optimizer(k), f)
+    self.assertEqual(registry.Registries.optimizers[k], f)
+    self.assertEqual(registry.Registries.optimizers[k], registry.optimizer(k))
+
+
 class HParamRegistryTest(tf.test.TestCase):
 
   def setUp(self):
-    registry._reset()
+    registry.Registries.hparams._clear()
+    registry.Registries.ranged_hparams._clear()
 
   def testHParamSet(self):
 
@@ -121,9 +254,9 @@ def my_hparams_range(_):
     self.assertTrue(registry.ranged_hparams("a") is my_hparams_range)
 
   def testUnknownHparams(self):
-    with self.assertRaisesRegexp(LookupError, "never registered"):
+    with self.assertRaisesRegexp(KeyError, "never registered"):
       registry.hparams("not_registered")
-    with self.assertRaisesRegexp(LookupError, "never registered"):
+    with self.assertRaisesRegexp(KeyError, "never registered"):
       registry.ranged_hparams("not_registered")
 
   def testNoneHparams(self):
@@ -194,34 +327,7 @@ def rhp_bad2(a, b):  # pylint: disable=unused-argument
         pass
 
 
-class CreateRegistry(tf.test.TestCase):
-  """Test class for `create_registry`."""
-
-  def testCreateRegistry(self):
-    my_registry = registry.create_registry("test_reg1")
-    self.assertIs(my_registry, registry.registry("test_reg1"))
-
-    # Use as decorator on a fn
-    @my_registry.register("foo")
-    def some_fn(num):
-      return num + 2
-
-    # Register a regular object
-    pod_obj = 4
-    my_registry.register("bar")(pod_obj)
-
-    # Register a class
-    @my_registry.register("foobar")
-    class A(object):
-      pass
-
-    self.assertEqual(9, my_registry.get("foo")(7))
-    self.assertEqual(["bar", "foo", "foobar"], my_registry.list())
-    foobar = my_registry.get("foobar")
-    self.assertTrue(isinstance(foobar(), A))
-
-
-class RegistryTest(tf.test.TestCase):
+class RegistryHelpTest(tf.test.TestCase):
   """Test class for common functions."""
 
   def testRegistryHelp(self):

From ab327c60c49fbd9a783eb6964a2e128ca0ebd15f Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 28 Jan 2019 15:09:03 -0800
Subject: [PATCH 1588/2720] internal

PiperOrigin-RevId: 231297048
---
 tensor2tensor/bin/t2t_attack.py               |   2 +-
 tensor2tensor/bin/t2t_datagen.py              |   5 +-
 tensor2tensor/bin/t2t_prune.py                |   2 +-
 tensor2tensor/layers/common_hparams.py        |   4 +-
 tensor2tensor/models/mtf_transformer2.py      |   2 +-
 .../models/research/adafactor_experiments.py  |  10 +-
 tensor2tensor/models/research/autoencoders.py |   2 +-
 .../models/research/transformer_nat.py        |   2 +-
 .../models/research/transformer_vae.py        |   4 +-
 .../models/research/vqa_attention.py          |   2 +-
 .../models/research/vqa_self_attention.py     |   2 +-
 tensor2tensor/models/shake_shake.py           |   2 +-
 tensor2tensor/models/transformer.py           |   4 +-
 tensor2tensor/models/vanilla_gan.py           |   2 +-
 tensor2tensor/problems.py                     |   2 +-
 tensor2tensor/rl/datagen_with_agent.py        |   2 +-
 .../transformer_test_ckpt/hparams.json        |   2 +-
 tensor2tensor/utils/adafactor.py              |   2 +-
 tensor2tensor/utils/learning_rate.py          |   2 +-
 tensor2tensor/utils/optimize.py               |  19 +-
 tensor2tensor/utils/registry.py               | 818 +++++++++---------
 tensor2tensor/utils/registry_test.py          | 174 +---
 22 files changed, 470 insertions(+), 596 deletions(-)

diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index f7e0fcc60..6276d6b4c 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -73,7 +73,7 @@ def create_attack_params():
 
 
 def create_attack(attack):
-  return registry.attack(attack)
+  return registry.attacks(attack)
 
 
 def create_surrogate_hparams():
diff --git a/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index c32a86cf5..b778bc984 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -147,7 +147,7 @@ def main(_):
 
   # Calculate the list of problems to generate.
   problems = sorted(
-      list(_SUPPORTED_PROBLEM_GENERATORS) + registry.list_base_problems())
+      list(_SUPPORTED_PROBLEM_GENERATORS) + registry.list_problems())
   for exclude in FLAGS.exclude_problems.split(","):
     if exclude:
       problems = [p for p in problems if exclude not in p]
@@ -169,8 +169,7 @@ def main(_):
 
   if not problems:
     problems_str = "\n  * ".join(
-        sorted(list(_SUPPORTED_PROBLEM_GENERATORS) +
-               registry.list_base_problems()))
+        sorted(list(_SUPPORTED_PROBLEM_GENERATORS) + registry.list_problems()))
     error_msg = ("You must specify one of the supported problems to "
                  "generate data for:\n  * " + problems_str + "\n")
     error_msg += ("TIMIT and parsing need data_sets specified with "
diff --git a/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
index e7c8c75db..008462e63 100644
--- a/tensor2tensor/bin/t2t_prune.py
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -54,7 +54,7 @@ def create_pruning_params():
 
 
 def create_pruning_strategy(name):
-  return registry.pruning_strategy(name)
+  return registry.pruning_strategies(name)
 
 
 def main(argv):
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 7158c3f4d..2bbfb7618 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -55,7 +55,7 @@ def basic_params1():
       initializer="orthogonal",
       initializer_gain=1.5,
       label_smoothing=0.1,
-      optimizer="adam",
+      optimizer="Adam",
       optimizer_adam_epsilon=1e-6,
       optimizer_adam_beta1=0.85,
       optimizer_adam_beta2=0.997,
@@ -466,7 +466,7 @@ def basic_range1(ranged_hparams):
   rhp.set_float("optimizer_adam_beta2", 0.995, 0.999)
   rhp.set_categorical(
       "optimizer",
-      ["adam", "adagrad", "momentum", "rms_prop", "sgd", "yellow_fin"])
+      ["Adam", "Adagrad", "Momentum", "RMSProp", "SGD", "YellowFin"])
 
 
 @registry.register_ranged_hparams
diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index 2090228e8..43dafed38 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -269,7 +269,7 @@ def sample(self, features, mesh):
     return self.combine_batch_dims(ret)
 
 
-layers_registry = registry.Registries.layers
+layers_registry = registry.create_registry("layers")
 
 
 # The following functions construct layers based on hyperparmeters
diff --git a/tensor2tensor/models/research/adafactor_experiments.py b/tensor2tensor/models/research/adafactor_experiments.py
index fbe4dbc2f..d7d3d4e2c 100644
--- a/tensor2tensor/models/research/adafactor_experiments.py
+++ b/tensor2tensor/models/research/adafactor_experiments.py
@@ -30,16 +30,16 @@ def mimic_adam_with_adafactor(hparams):
   Some minor things may be different, like epsilon and beta1 correction.
 
   Args:
-    hparams: model hyperparameters where "adam" in hparams.optimizer
+    hparams: model hyperparameters where "Adam" in hparams.optimizer
   """
-  assert "adam" in hparams.optimizer
-  hparams.optimizer = "adafactor"
+  assert "Adam" in hparams.optimizer
+  hparams.optimizer = "Adafactor"
   hparams.optimizer_adafactor_beta1 = hparams.optimizer_adam_beta1
   hparams.optimizer_adafactor_beta2 = hparams.optimizer_adam_beta2
   hparams.optimizer_adafactor_multiply_by_parameter_scale = False
   hparams.optimizer_adafactor_factored = False
   hparams.optimizer_adafactor_clipping_threshold = None
-  hparams.optimizer_adafactor_decay_type = "adam"
+  hparams.optimizer_adafactor_decay_type = "Adam"
 
 
 @registry.register_hparams
@@ -50,7 +50,7 @@ def afx_adam():
   hparams.optimizer_adam_beta2 = 0.999
   hparams.symbol_modality_num_shards = 1
   hparams.batch_size = 2048
-  hparams.optimizer = "adam"
+  hparams.optimizer = "Adam"
   hparams.learning_rate_schedule = (
       "constant*rsqrt_decay*linear_warmup*rsqrt_hidden_size")
   hparams.learning_rate_constant = 2.0
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index ecab3faa1..3780bd6e6 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1020,7 +1020,7 @@ def body(self, features):
 def autoencoder_basic():
   """Basic autoencoder model."""
   hparams = common_hparams.basic_params1()
-  hparams.optimizer = "adam"
+  hparams.optimizer = "Adam"
   hparams.learning_rate_constant = 0.0002
   hparams.learning_rate_warmup_steps = 500
   hparams.learning_rate_schedule = "constant * linear_warmup"
diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index ccea877ed..7a775153d 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -392,7 +392,7 @@ def transformer_nat_small():
   hparams.filter_size = 2048
   hparams.label_smoothing = 0.0
   hparams.force_full_predict = True
-  hparams.optimizer = "adam"
+  hparams.optimizer = "Adam"
   hparams.optimizer_adam_epsilon = 1e-9
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.997
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 64aecd405..c630abc6c 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -767,7 +767,7 @@ def transformer_ae_small():
   hparams.filter_size = 2048
   hparams.add_hparam("compress_filter_size", 2048 * 2)
   hparams.label_smoothing = 0.0
-  hparams.optimizer = "adam"  # Can be unstable, maybe try Adam.
+  hparams.optimizer = "Adam"  # Can be unstable, maybe try Adam.
   hparams.optimizer_adam_epsilon = 1e-9
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.997  # Needs tuning, try 0.98 to 0.999.
@@ -941,7 +941,7 @@ def transformer_ae_a3():
 def transformer_ae_a6():
   """Best hparams for transformer with semhash."""
   hparams = transformer_ae_a3()
-  hparams.optimizer = "adam"
+  hparams.optimizer = "Adam"
   hparams.noise_dev = 0.5
   return hparams
 
diff --git a/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index 57f1975e6..774387878 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -335,7 +335,7 @@ def vqa_attention_base():
   hparams = common_hparams.basic_params1()
   hparams.batch_size = 128
   hparams.use_fixed_batch_size = True,
-  hparams.optimizer = "adam"
+  hparams.optimizer = "Adam"
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.999
   hparams.optimizer_adam_epsilon = 1e-8
diff --git a/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
index b8388e606..482e18c9e 100644
--- a/tensor2tensor/models/research/vqa_self_attention.py
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -684,7 +684,7 @@ def vqa_self_attention_base():
   hparams = common_hparams.basic_params1()
   hparams.batch_size = 128
   hparams.use_fixed_batch_size = True,
-  hparams.optimizer = "adam"
+  hparams.optimizer = "Adam"
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.997
   hparams.optimizer_adam_epsilon = 1e-9
diff --git a/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index 22f46f989..4af05dc21 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -189,7 +189,7 @@ def shakeshake_small():
 @registry.register_hparams
 def shake_shake_quick():
   hparams = shakeshake_small()
-  hparams.optimizer = "adam"
+  hparams.optimizer = "Adam"
   hparams.learning_rate_cosine_cycle_steps = 1000
   hparams.learning_rate = 0.5
   hparams.batch_size = 100
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 5788910d2..f0d0596e8 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1694,7 +1694,7 @@ def transformer_tall_pretrain_lm():
   hparams.learning_rate_constant = 2e-4
   hparams.learning_rate_schedule = (
       "linear_warmup*constant*cosdecay")
-  hparams.optimizer = "adam_w"
+  hparams.optimizer = "AdamW"
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.999
   hparams.optimizer_adam_epsilon = 1e-8
@@ -1739,7 +1739,7 @@ def transformer_tall_pretrain_lm_tpu():
   # Optimizer gets reset in update_hparams_for_tpu so we set it again here.
   hparams.learning_rate_constant = 2e-4
   hparams.learning_rate_schedule = ("linear_warmup * constant * cosdecay")
-  hparams.optimizer = "adam_w"
+  hparams.optimizer = "AdamW"
   return hparams
 
 
diff --git a/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index 00fa89d18..cb16a4c0b 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -199,7 +199,7 @@ def infer(self, *args, **kwargs):  # pylint: disable=arguments-differ
 def sliced_gan():
   """Basic parameters for a vanilla_gan."""
   hparams = common_hparams.basic_params1()
-  hparams.optimizer = "adam"
+  hparams.optimizer = "Adam"
   hparams.learning_rate_constant = 0.0002
   hparams.learning_rate_warmup_steps = 500
   hparams.learning_rate_schedule = "constant * linear_warmup"
diff --git a/tensor2tensor/problems.py b/tensor2tensor/problems.py
index 2cf8c8762..1337c2a33 100644
--- a/tensor2tensor/problems.py
+++ b/tensor2tensor/problems.py
@@ -27,7 +27,7 @@ def problem(name):
 
 
 def available():
-  return registry.list_base_problems()
+  return sorted(registry.list_problems())
 
 
 all_problems.import_modules(all_problems.ALL_MODULES)
diff --git a/tensor2tensor/rl/datagen_with_agent.py b/tensor2tensor/rl/datagen_with_agent.py
index 1e5100537..f634e8c87 100644
--- a/tensor2tensor/rl/datagen_with_agent.py
+++ b/tensor2tensor/rl/datagen_with_agent.py
@@ -45,7 +45,7 @@ def main(_):
 
   # Create problem if not already defined
   problem_name = "gym_discrete_problem_with_agent_on_%s" % FLAGS.game
-  if problem_name not in registry.Registries.problems:
+  if problem_name not in registry.list_problems():
     gym_env.register_game(FLAGS.game)
 
   # Generate
diff --git a/tensor2tensor/test_data/transformer_test_ckpt/hparams.json b/tensor2tensor/test_data/transformer_test_ckpt/hparams.json
index 1a6a97223..df9a654c0 100644
--- a/tensor2tensor/test_data/transformer_test_ckpt/hparams.json
+++ b/tensor2tensor/test_data/transformer_test_ckpt/hparams.json
@@ -1 +1 @@
-{"daisy_chain_variables": true, "optimizer_adam_beta1": 0.9, "scheduled_sampling_prob": 0.0, "num_hidden_layers": 2, "moe_loss_coef": 0.01, "max_target_seq_length": 0, "clip_grad_norm": 0.0, "pos": "timing", "scheduled_sampling_gold_mixin_prob": 0.5, "initializer": "uniform_unit_scaling", "grad_noise_scale": 0.0, "optimizer_momentum_momentum": 0.9, "nbr_decoder_problems": 1, "attention_key_channels": 0, "eval_drop_long_sequences": false, "learning_rate_cosine_cycle_steps": 250000, "prepend_mode": "none", "weight_decay": 0.0, "symbol_modality_skip_top": false, "weight_noise": 0.0, "target_modality": "default", "attention_dropout": 0.1, "parameter_attention_value_channels": 0, "factored_logits": false, "relu_dropout": 0.1, "no_data_parallelism": false, "layer_preprocess_sequence": "n", "sampling_method": "argmax", "learning_rate": 0.2, "num_heads": 2, "max_length": 256, "summarize_grads": false, "attention_value_channels": 0, "num_encoder_layers": 0, "label_smoothing": 0.1, "use_fixed_batch_size": false, "optimizer": "adam", "moe_k": 2, "self_attention_type": "dot_product", "learning_rate_decay_scheme": "noam", "sampling_temp": 1.0, "kernel_height": 3, "use_pad_remover": true, "batch_size": 4096, "max_relative_position": 0, "force_full_predict": false, "min_length_bucket": 8, "layer_prepostprocess_dropout": 0.1, "eval_run_autoregressive": false, "shared_embedding_and_softmax_weights": true, "symbol_modality_num_shards": 16, "dropout": 0.2, "compress_steps": 0, "parameter_attention_key_channels": 0, "length_bucket_step": 1.1, "kernel_width": 1, "hidden_size": 16, "num_decoder_layers": 0, "input_modalities": "default", "filter_size": 8, "optimizer_adam_beta2": 0.98, "scheduled_sampling_warmup_steps": 50000, "norm_type": "layer", "min_length": 0, "moe_num_experts": 64, "multiply_embedding_mode": "sqrt_depth", "max_input_seq_length": 0, "learning_rate_warmup_steps": 8000, "proximity_bias": false, "ffn_layer": "dense_relu_dense", "initializer_gain": 1.0, "layer_postprocess_sequence": "da", "moe_hidden_sizes": "2048", "optimizer_adam_epsilon": 1e-09, "norm_epsilon": 1e-06}
+{"daisy_chain_variables": true, "optimizer_adam_beta1": 0.9, "scheduled_sampling_prob": 0.0, "num_hidden_layers": 2, "moe_loss_coef": 0.01, "max_target_seq_length": 0, "clip_grad_norm": 0.0, "pos": "timing", "scheduled_sampling_gold_mixin_prob": 0.5, "initializer": "uniform_unit_scaling", "grad_noise_scale": 0.0, "optimizer_momentum_momentum": 0.9, "nbr_decoder_problems": 1, "attention_key_channels": 0, "eval_drop_long_sequences": false, "learning_rate_cosine_cycle_steps": 250000, "prepend_mode": "none", "weight_decay": 0.0, "symbol_modality_skip_top": false, "weight_noise": 0.0, "target_modality": "default", "attention_dropout": 0.1, "parameter_attention_value_channels": 0, "factored_logits": false, "relu_dropout": 0.1, "no_data_parallelism": false, "layer_preprocess_sequence": "n", "sampling_method": "argmax", "learning_rate": 0.2, "num_heads": 2, "max_length": 256, "summarize_grads": false, "attention_value_channels": 0, "num_encoder_layers": 0, "label_smoothing": 0.1, "use_fixed_batch_size": false, "optimizer": "Adam", "moe_k": 2, "self_attention_type": "dot_product", "learning_rate_decay_scheme": "noam", "sampling_temp": 1.0, "kernel_height": 3, "use_pad_remover": true, "batch_size": 4096, "max_relative_position": 0, "force_full_predict": false, "min_length_bucket": 8, "layer_prepostprocess_dropout": 0.1, "eval_run_autoregressive": false, "shared_embedding_and_softmax_weights": true, "symbol_modality_num_shards": 16, "dropout": 0.2, "compress_steps": 0, "parameter_attention_key_channels": 0, "length_bucket_step": 1.1, "kernel_width": 1, "hidden_size": 16, "num_decoder_layers": 0, "input_modalities": "default", "filter_size": 8, "optimizer_adam_beta2": 0.98, "scheduled_sampling_warmup_steps": 50000, "norm_type": "layer", "min_length": 0, "moe_num_experts": 64, "multiply_embedding_mode": "sqrt_depth", "max_input_seq_length": 0, "learning_rate_warmup_steps": 8000, "proximity_bias": false, "ffn_layer": "dense_relu_dense", "initializer_gain": 1.0, "layer_postprocess_sequence": "da", "moe_hidden_sizes": "2048", "optimizer_adam_epsilon": 1e-09, "norm_epsilon": 1e-06}
diff --git a/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index 9c44d28d9..01b6145ba 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -326,7 +326,7 @@ def adafactor_optimizer_from_hparams(hparams, lr):
   Raises:
     ValueError: on illegal values
   """
-  if hparams.optimizer_adafactor_decay_type == "adam":
+  if hparams.optimizer_adafactor_decay_type == "Adam":
     decay_rate = adafactor_decay_rate_adam(
         hparams.optimizer_adafactor_beta2)
   elif hparams.optimizer_adafactor_decay_type == "pow":
diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index 709f993e7..bf037a876 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -90,7 +90,7 @@ def legacy_learning_rate_schedule(hparams):
     warmup = _learning_rate_warmup(warmup_steps, hparams=hparams)
     decay = _learning_rate_decay(hparams, warmup_steps)
     ret = tf.where(step_num < warmup_steps, warmup, decay)
-  optimizer_correction = 0.002 if "adam" in hparams.optimizer else 1.0
+  optimizer_correction = 0.002 if "Adam" in hparams.optimizer else 1.0
   tf.logging.info("Base learning rate: %f", hparams.learning_rate)
   return ret * optimizer_correction * hparams.learning_rate
 
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 641f60601..5c3b5b3bf 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -94,7 +94,7 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
   return train_op
 
 
-@registry.register_optimizer
+@registry.register_optimizer("Adam")
 def adam(learning_rate, hparams):
   # We change the default epsilon for Adam.
   # Using LazyAdam as it's much faster for large vocabulary embeddings.
@@ -105,7 +105,7 @@ def adam(learning_rate, hparams):
       epsilon=hparams.optimizer_adam_epsilon)
 
 
-@registry.register_optimizer
+@registry.register_optimizer("MultistepAdam")
 def multistep_adam(learning_rate, hparams):
   return multistep_optimizer.MultistepAdamOptimizer(
       learning_rate,
@@ -115,7 +115,7 @@ def multistep_adam(learning_rate, hparams):
       n=hparams.optimizer_multistep_accumulate_steps)
 
 
-@registry.register_optimizer
+@registry.register_optimizer("Momentum")
 def momentum(learning_rate, hparams):
   return tf.train.MomentumOptimizer(
       learning_rate,
@@ -123,14 +123,14 @@ def momentum(learning_rate, hparams):
       use_nesterov=hparams.optimizer_momentum_nesterov)
 
 
-@registry.register_optimizer
+@registry.register_optimizer("YellowFin")
 def yellow_fin(learning_rate, hparams):
   return yellowfin.YellowFinOptimizer(
       learning_rate=learning_rate,
       momentum=hparams.optimizer_momentum_momentum)
 
 
-@registry.register_optimizer
+@registry.register_optimizer("TrueAdam")
 def true_adam(learning_rate, hparams):
   return tf.train.AdamOptimizer(
       learning_rate,
@@ -139,7 +139,7 @@ def true_adam(learning_rate, hparams):
       epsilon=hparams.optimizer_adam_epsilon)
 
 
-@registry.register_optimizer
+@registry.register_optimizer("AdamW")
 def adam_w(learning_rate, hparams):
   # Openai gpt used weight decay.
   # Given the internals of AdamW, weight decay dependent on the
@@ -156,7 +156,7 @@ def adam_w(learning_rate, hparams):
       epsilon=hparams.optimizer_adam_epsilon)
 
 
-@registry.register_optimizer("adafactor")
+@registry.register_optimizer("Adafactor")
 def register_adafactor(learning_rate, hparams):
   return adafactor.adafactor_optimizer_from_hparams(hparams, learning_rate)
 
@@ -169,11 +169,8 @@ def _register_base_optimizer(key, fn):
 
 
 for k in tf.contrib.layers.OPTIMIZER_CLS_NAMES:
-  if k not in registry.Registries.optimizers and k not in ('SGD', 'RMSProp'):
+  if k not in registry._OPTIMIZERS:  # pylint: disable=protected-access
     _register_base_optimizer(k, tf.contrib.layers.OPTIMIZER_CLS_NAMES[k])
-_register_base_optimizer('sgd', tf.contrib.layers.OPTIMIZER_CLS_NAMES['SGD'])
-_register_base_optimizer(
-    'rms_prop', tf.contrib.layers.OPTIMIZER_CLS_NAMES['RMSProp'])
 
 
 class ConditionalOptimizer(tf.train.Optimizer):
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 56dc11f50..d8683f234 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -13,522 +13,510 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Registry for
-* models
-* hyperparameter settings
-* ranged hyperparameter settings
-* problem types
-* attacks
-* attack parameters
-* pruning parameters
-* pruning strategies
-* optimizers
-
-`Registries` contains the `Registry` objects used throughout tensor2tensor.
-
-New functions and classes can be registered using `.register`. The can be
-accessed/queried similar to dictionaries, keyed by default by `snake_case`
-equivalents.
-```
-@Registries.models.register
-class MyModel(T2TModel):
-  ...
-
-'my_model' in Registries.models  # True
-for k in Registries.models:
-  print(k)  # prints 'my_model'
-model = Registries.models['my_model'](constructor_arg)
-```
+"""Registry for models, hyperparameter settings, problem types, and datasets.
 
-`Registry`s for functions which take no arguments will return the
-result of evaluating those functions (though this is not the default behaviour
-of the `Registry` class in general).
-```
-@Registries.attacks.register
-def my_attack():
-  ...
-
-my_attack_obj = Registries.attacks['my_attack']
-my_attack_obj()  # TypeError: 'Attack' object is not callable
-```
-
-#### Legacy Support
 Define a new model by subclassing T2TModel and register it:
 
 ```
-@register_model
+@registry.register_model
 class MyModel(T2TModel):
   ...
 ```
 
-Access by snake-cased name: `model("my_model")`. If you're using
+Access by snake-cased name: `registry.model("my_model")`. If you're using
 `t2t_trainer.py`, you can pass on the command-line: `--model=my_model`.
 
-See all the models registered: `list_models()`.
+See all the models registered: `registry.list_models()`.
 
 For hyperparameter sets:
-  * Register: `register_hparams`
-  * List: `list_hparams`
-  * Retrieve by name: `hparams`
+  * Register: `registry.register_hparams`
+  * List: `registry.list_hparams`
+  * Retrieve by name: `registry.hparams`
   * Command-line flag in `t2t_trainer.py`: `--hparams_set=name`
 
 For hyperparameter ranges:
-  * Register: `register_ranged_hparams`
-  * List: `list_ranged_hparams`
-  * Retrieve by name: `ranged_hparams`
+  * Register: `registry.register_ranged_hparams`
+  * List: `registry.list_ranged_hparams`
+  * Retrieve by name: `registry.ranged_hparams`
   * Command-line flag in `t2t_trainer.py`: `--hparams_range=name`
 """
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensorflow import logging
-import tensorflow.contrib.framework as framework
-from tensorflow.python.util import tf_inspect as inspect
 import collections
 
 from tensor2tensor.utils import misc_utils
+import tensorflow as tf
+from tensorflow.python.util import tf_inspect as inspect
+
+_ATTACKS = {}
+_ATTACK_PARAMS = {}
+_HPARAMS = {}
+_MODELS = {}
+_PROBLEMS = {}
+_PRUNING_PARAMS = {}
+_PRUNING_STRATEGY = {}
+_RANGED_HPARAMS = {}
 
+# Key: registry name, Value: Registry
+_GENERIC_REGISTRIES = {}
+Registry = collections.namedtuple(
+    "_Registry", ["register", "get", "list", "registry"])
 
-def default_name(class_or_fn):
-  """Default name for a class or function.
 
-  This is the naming function by default for registries expecting classes or
-  functions.
+def registry(registry_name):
+  """Returns `Registry` created by `create_registry`."""
+  if registry_name not in _GENERIC_REGISTRIES:
+    raise KeyError("No registry named %s. Available:\n%s" % (
+        registry_name, sorted(_GENERIC_REGISTRIES)))
+  return _GENERIC_REGISTRIES[registry_name]
+
+
+def create_registry(registry_name):
+  """Create a generic object registry.
 
   Args:
-    class_or_fn: class or function to be named.
+    registry_name: str, name of the object registry.
 
   Returns:
-    Default name for registration.
+    `Registry` that contains functions for register (decorator), get, and list.
+
+  Raises:
+    KeyError: if `registry_name` is a pre-existing registry.
   """
-  return misc_utils.camelcase_to_snakecase(class_or_fn.__name__)
+  if registry_name in _GENERIC_REGISTRIES:
+    raise KeyError(
+        "Registry %s already exists." % registry_name)
 
+  registry_ = {}
 
-class Registry(object):
-  """Dict-like class for managing function registrations.
+  def register(name):
+    """Returns decorator to register an object."""
 
-  ```python
-  my_registry = Registry("custom_name")
+    def register_dec(obj):
+      if name in registry_:
+        raise KeyError(
+            "Registry %s already contains key %s." % (registry_name, name))
+      registry_[name] = obj
+      return obj
 
-  @my_registry.register
-  def my_func():
-    pass
+    return register_dec
 
-  @my_registry.register()
-  def another_func():
-    pass
+  def get(name):
+    if name not in registry_:
+      raise KeyError(
+          "Registry %s contains no object named %s" % (registry_name, name))
+    return registry_[name]
 
-  @my_registry.register("non_default_name")
-  def third_func(x, y, z):
-    pass
+  def list_registry():
+    return sorted(registry_)
 
-  def foo():
-    pass
+  registry_obj = Registry(
+      register=register,
+      get=get,
+      list=list_registry,
+      registry=registry_,
+  )
+  _GENERIC_REGISTRIES[registry_name] = registry_obj
+  return registry_obj
 
-  my_registry.register()(foo)
-  my_registry.register("baz")(lambda (x, y): x + y)
-  my_register.register("bar")
 
-  print(list(my_registry))
-  # ["my_func", "another_func", "non_default_name", "foo", "baz"]
-  # (order may vary)
-  print(my_registry["non_default_name"] is third_func)  # True
-  print("third_func" in my_registry)                    # False
-  print("bar" in my_registry)                           # False
-  my_registry["non-existent_key"]                       # raises KeyError
-  ```
+def _reset():
+  for ctr in [_MODELS, _HPARAMS, _RANGED_HPARAMS, _ATTACK_PARAMS]:
+    ctr.clear()
 
-  Optional validation, on_set callback and value transform also supported.
-  See `__init__` doc.
-  """
 
-  def __init__(
-      self, registry_name, default_key_fn=default_name, validator=None,
-      on_set=None, value_transformer=(lambda k, v: v)):
-    """Construct a new registry.
+def default_name(obj_class):
+  """Convert a class name to the registry's default name for the class.
 
-    Args:
-      registry_name: str identifier for the given registry. Used in error msgs.
-      default_key_fn (optional): function mapping value -> key for registration
-        when a key is not provided
-      validator (optional): if given, this is run before setting a given
-        (key, value) pair. Accepts (key, value) and should raise if there is a
-        problem. Overwriting existing keys is not allowed and is checked
-        separately. Values are also checked to be callable separately.
-      on_set (optional): callback function accepting (key, value) pair
-        which is run after an item is successfully set.
-      value_transformer (optional): if run, `__getitem__` will return
-        value_transformer(key, registered_value).
-    """
-    self._registry = {}
-    self._name = registry_name
-    self._default_key_fn = default_key_fn
-    self._validator = validator
-    self._on_set = on_set
-    self._value_transformer = value_transformer
-
-  def default_key(self, value):
-    """Default key used when key not provided. Uses function from __init__."""
-    return self._default_key_fn(value)
-
-  @property
-  def name(self):
-    return self._name
-
-  def validate(self, key, value):
-    """Validation function run before setting. Uses function from __init__."""
-    if self._validator is not None:
-      self._validator(key, value)
-
-  def on_set(self, key, value):
-    """Callback called on successful set. Uses function from __init__."""
-    if self._on_set is not None:
-      self._on_set(key, value)
-
-  def __setitem__(self, key, value):
-    """
-    Validate, set, and (if successfull) call `on_set` for the given item.
+  Args:
+    obj_class: the name of a class
 
-    Args:
-      key: key to store value under. If `None`, `self.default_key(value)` is
-        used.
-      value: callable stored under the given key.
+  Returns:
+    The registry's default name for the class.
+  """
+  return misc_utils.camelcase_to_snakecase(obj_class.__name__)
 
-    Returns:
-      `None`
-    """
-    if key is None:
-      key = self.default_key(value)
-    if key in self:
-      raise KeyError("key %s already registered in registry %s"
-                     % (key, self._name))
-    if not callable(value):
-      raise ValueError("value must be callable")
-    self.validate(key, value)
-    self._registry[key] = value
-    self.on_set(key, value)
-
-  def register(self, key_or_value=None):
-    """Decorator to register a function, or registration itself.
-
-    This is primarily intended for use as a decorator, either with or without
-    a key/parentheses.
-    ```python
-    @my_registry.register('key1')
-    def value_fn(x, y, z):
-      pass
-
-    @my_registry.register()
-    def another_fn(x, y):
-      pass
-
-    @my_registry.register
-    def third_func():
-      pass
-    ```
-
-    Note if key_or_value is provided as a non-callable, registration only
-    occurs once the returned callback is called with a callable as its only
-    argument.
-    ```python
-    callback = my_registry.register('different_key')
-    'different_key' in my_registry  # False
-    callback(lambda (x, y): x + y)
-    'different_key' in my_registry  # True
-    ```
 
-    Args:
-      key_or_value (optional): key to access the registered value with, or the
-        function itself. If `None` (default), `self.default_key` will be called
-        on `value` once the returned callback is called with `value` as the
-        only arg. If `key_or_value` is itself callable, it is assumed to be the
-        value and the key is given by `self.default_key(key)`.
+def default_object_name(obj):
+  """Convert an object to the registry's default name for the object class.
 
-    Returns:
-      decorated callback, or callback generated a decorated function.
-    """
-    def decorator(value, key):
-      self[key] = value
-      return value
+  Args:
+    obj: an object instance
 
-    # Handle if decorator was used without parens
-    if callable(key_or_value):
-      return decorator(value=key_or_value, key=None)
-    else:
-      return lambda value: decorator(value, key=key_or_value)
+  Returns:
+    The registry's default name for the class of the object.
+  """
+  return default_name(obj.__class__)
 
-  def __getitem__(self, key):
-    if key not in self:
-      raise KeyError("%s never registered with registry %s. Available:\n %s" %
-                     (key, self.name,
-                      display_list_by_prefix(sorted(self), 4)))
-    value = self._registry[key]
-    return self._value_transformer(key, value)
 
-  def __contains__(self, key):
-    return key in self._registry
+def register_model(name=None):
+  """Register a model. name defaults to class name snake-cased."""
 
-  def keys(self):
-    return self._registry.keys()
+  def decorator(model_cls, registration_name=None):
+    """Registers & returns model_cls with registration_name or default name."""
+    model_name = registration_name or default_name(model_cls)
+    if model_name in _MODELS and not tf.executing_eagerly():
+      raise LookupError("Model %s already registered." % model_name)
+    model_cls.REGISTERED_NAME = model_name
+    _MODELS[model_name] = model_cls
+    return model_cls
 
-  def values(self):
-    return (self[k] for k in self)       # complicated because of transformer
+  # Handle if decorator was used without parens
+  if callable(name):
+    model_cls = name
+    return decorator(model_cls, registration_name=default_name(model_cls))
 
-  def items(self):
-    return ((k, self[k]) for k in self)  # complicated because of transformer
+  return lambda model_cls: decorator(model_cls, name)
 
-  def __iter__(self):
-    return iter(self._registry)
 
-  def __len__(self):
-    return len(self._registry)
+def model(name):
+  if name not in _MODELS:
+    raise LookupError("Model %s never registered.  Available models:\n %s" %
+                      (name, "\n".join(list_models())))
 
-  def _clear(self):
-    self._registry.clear()
+  return _MODELS[name]
 
-  def get(self, key, d=None):
-    return self[key] if key in self else d
 
+def list_models():
+  return list(sorted(_MODELS))
 
-def _on_model_set(k, v):
-  v.REGISTERED_NAME = k
 
+_OPTIMIZERS = {}
 
-def _nargs_validator(nargs, message):
-  if message is None:
-    message = "Registered function must take exactly %d arguments" % nargs
-  def f(key, value):
-    spec = inspect.getfullargspec(value)
-    if (len(spec.args) != nargs or
-        spec.varargs is not None or
-        spec.varkw is not None):
-      raise ValueError(message)
 
-  return f
+def register_optimizer(name=None):
+  """Register an optimizer. name defaults to upper camel case of fn name."""
 
+  def default_opt_name(opt_fn):
+    return misc_utils.snakecase_to_camelcase(default_name(opt_fn))
 
-ProblemSpec = collections.namedtuple(
-    "ProblemSpec", ["base_name", "was_reversed", "was_copy"])
+  def decorator(opt_fn, registration_name):
+    """Registers and returns optimizer_fn with registration_name or default."""
+    if registration_name is None:
+      registration_name = default_opt_name(opt_fn)
 
+    if registration_name in _OPTIMIZERS and not tf.executing_eagerly():
+      raise LookupError("Optimizer %s already registered." % registration_name)
+    args, varargs, keywords, _ = inspect.getargspec(opt_fn)
 
-def parse_problem_name(name):
-  """Determines if problem_name specifies a copy and/or reversal.
+    if len(args) != 2 or varargs is not None or keywords is not None:
+      raise ValueError("Optimizer registration function must take two "
+                       "arguments: learning_rate (float) and "
+                       "hparams (HParams).")
+    _OPTIMIZERS[registration_name] = opt_fn
+    return opt_fn
 
-  Args:
-    name: str, problem name, possibly with suffixes.
+  if callable(name):
+    opt_fn = name
+    registration_name = default_opt_name(opt_fn)
+    return decorator(opt_fn, registration_name=registration_name)
 
-  Returns:
-    ProblemSpec: namedtuple with ["base_name", "was_reversed", "was_copy"]
+  return lambda opt_fn: decorator(opt_fn, name)
 
-  Raises:
-    ValueError if name contains multiple suffixes of the same type
-      ('_rev' or '_copy'). One of each is ok.
-  """
-  # Recursively strip tags until we reach a base name.
-  if name.endswith("_rev"):
-    base, was_rev, was_copy = parse_problem_name(name[:-4])
-    if was_rev:
-      # duplicate rev
-      raise ValueError(
-          "Invalid problem name %s: multiple '_rev' instances" % name)
-    return ProblemSpec(base, True, was_copy)
-  elif name.endswith("_copy"):
-    base, was_reversed, was_copy = parse_problem_name(name[:-5])
-    if was_copy:
-      raise ValueError(
-          "Invalid problem_name %s: multiple '_copy' instances" % name)
-    return ProblemSpec(base, was_reversed, True)
-  else:
-    return ProblemSpec(name, False, False)
-
-
-def get_problem_name(base_name, was_reversed=False, was_copy=False):
-  """
-  Construct a problem name from base and reversed/copy options.
 
-  Inverse of `parse_problem_name`.
+def optimizer(name):
+  if name not in _OPTIMIZERS:
+    raise LookupError("Optimizer %s never registered. "
+                      "Available optimizers:\n %s"
+                      % (name, "\n".join(list_optimizers())))
+  return _OPTIMIZERS[name]
 
-  Args:
-    base_name: base problem name. Should not end in "_rev" or "_copy"
-    was_reversed: if the problem is to be reversed
-    was_copy: if the problem is to be copied
 
-  Returns:
-    string name consistent with use with `parse_problem_name`.
+def list_optimizers():
+  return list(sorted(_OPTIMIZERS))
 
-  Raises:
-    ValueError if `base_name` ends with "_rev" or "_copy"
-  """
-  if any(base_name.endswith(suffix) for suffix in ("_rev", "_copy")):
-    raise ValueError("`base_name` cannot end in '_rev' or '_copy'")
-  name = base_name
-  if was_copy:
-    name = "%s_copy" % name
-  if was_reversed:
-    name = "%s_rev" % name
-  return name
-
-
-def _problem_name_validator(k, v):
-  if parse_problem_name(k).base_name != k:
-    raise KeyError(
-        "Invalid problem name: cannot end in %s or %s" % ("_rev", "_copy"))
 
+def register_hparams(name=None):
+  """Register an HParams set. name defaults to function name snake-cased."""
 
-def _on_problem_set(k, v):
-  v.name = k
+  def decorator(hp_fn, registration_name=None):
+    """Registers & returns hp_fn with registration_name or default name."""
+    hp_name = registration_name or default_name(hp_fn)
+    if hp_name in _HPARAMS and not tf.executing_eagerly():
+      raise LookupError("HParams set %s already registered." % hp_name)
+    _HPARAMS[hp_name] = hp_fn
+    return hp_fn
 
+  # Handle if decorator was used without parens
+  if callable(name):
+    hp_fn = name
+    return decorator(hp_fn, registration_name=default_name(hp_fn))
 
-def _call_value(k, v):
-  return v()
+  return lambda hp_fn: decorator(hp_fn, name)
 
 
-def _hparams_value_transformer(key, value):
-  out = value()
-  if out is None:
+def hparams(name):
+  """Retrieve registered hparams by name."""
+  if name not in _HPARAMS:
+    error_msg = "HParams set %s never registered. Sets registered:\n%s"
+    raise LookupError(
+        error_msg % (name,
+                     display_list_by_prefix(list_hparams(), starting_spaces=4)))
+  hp = _HPARAMS[name]()
+  if hp is None:
     raise TypeError("HParams %s is None. Make sure the registered function "
-                    "returns the HParams object" % key)
-  return out
+                    "returns the HParams object." % name)
+  return hp
 
 
-class Registries(object):
-  """Object holding `Registry` objects."""
-  def __init__(self):
-    raise RuntimeError("Registries is not intended to be instantiated")
+def list_hparams(prefix=None):
+  if prefix:
+    return [name for name in _HPARAMS if name.startswith(prefix)]
+  return list(_HPARAMS)
 
 
-  models = Registry("models", on_set=_on_model_set)
+def register_ranged_hparams(name=None):
+  """Register a RangedHParams set. name defaults to fn name snake-cased."""
 
-  optimizers = Registry(
-      "optimizers",
-      validator=_nargs_validator(
-          2,
-          "Registered optimizer functions must take exactly two arguments: "
-          "learning_rate (float) and hparams (HParams)."))
+  def decorator(rhp_fn, registration_name=None):
+    """Registers & returns hp_fn with registration_name or default name."""
+    rhp_name = registration_name or default_name(rhp_fn)
+    if rhp_name in _RANGED_HPARAMS:
+      raise LookupError("RangedHParams set %s already registered." % rhp_name)
+    # Check that the fn takes a single argument
+    args, varargs, keywords, _ = inspect.getargspec(rhp_fn)
+    if len(args) != 1 or varargs is not None or keywords is not None:
+      raise ValueError("RangedHParams set function must take a single "
+                       "argument, the RangedHParams object.")
 
-  hparams = Registry(
-      "hparams", value_transformer=_hparams_value_transformer)
+    _RANGED_HPARAMS[rhp_name] = rhp_fn
+    return rhp_fn
 
-  ranged_hparams = Registry(
-      "ranged_hparams", validator=_nargs_validator(
-          1,
-          "Registered ranged_hparams functions must take a single argument, "
-          "the RangedHParams object."))
+  # Handle if decorator was used without parens
+  if callable(name):
+    rhp_fn = name
+    return decorator(rhp_fn, registration_name=default_name(rhp_fn))
 
-  problems = Registry(
-      "problems", validator=_problem_name_validator, on_set=_on_problem_set)
+  return lambda rhp_fn: decorator(rhp_fn, name)
 
-  attacks = Registry(
-      "attacks", value_transformer=_call_value)
 
-  attack_params = Registry(
-      "attack_params", value_transformer=_call_value)
+def ranged_hparams(name):
+  if name not in _RANGED_HPARAMS:
+    raise LookupError("RangedHParams set %s never registered." % name)
+  return _RANGED_HPARAMS[name]
 
-  pruning_params = Registry(
-      "pruning_params", value_transformer=_call_value)
 
-  pruning_strategies = Registry("pruning_strategies")
+def list_ranged_hparams():
+  return list(_RANGED_HPARAMS)
 
-  layers = Registry("layers", validator=_nargs_validator(
-      2,
-      "Registered layer functions must take exaction two arguments: "
-      "hparams (HParams) and prefix (str)."
-  ))
 
-# consistent version of old API
-model = Registries.models.__getitem__
-list_models = lambda: sorted(Registries.models)
-register_model = Registries.models.register
+def register_problem(name=None):
+  """Register a Problem. name defaults to cls name snake-cased."""
 
-# optimizer = optimizer_registry.__getitem__
-def optimizer(name):
-  """
-  Get pre-registered optimizer keyed by name.
+  def decorator(p_cls, registration_name=None):
+    """Registers & returns p_cls with registration_name or default name."""
+    p_name = registration_name or default_name(p_cls)
+    if p_name in _PROBLEMS and not tf.executing_eagerly():
+      raise LookupError("Problem %s already registered." % p_name)
 
-  `name` should be snake case, though SGD -> sgd, RMSProp -> rms_prop and
-  UpperCamelCase -> snake_case conversions included for legacy support.
+    _PROBLEMS[p_name] = p_cls
+    p_cls.name = p_name
+    return p_cls
 
-  Args:
-    name: name of optimizer used in registration. This should be a snake case
-      identifier, though others supported for legacy reasons.
-  """
-  rest = ("Please update `registry.optimizer` callsite "
-          "(likely due to a `HParams.optimizer` value)")
-  if name == "SGD":
-    name = "sgd"
-    logging.warning("'SGD' optimizer now keyed by 'sgd'. %s" % rest)
-  elif name == 'rms_prop':
-    name = 'rms_prop'
-    logging.warning("'RMSProp' optimizer now keyed by 'rms_prop'. %s" % rest)
-  else:
-    snake_name = misc_utils.camelcase_to_snakecase(name)
-    if name != snake_name:
-      logging.warning(
-          "optimizer names now keyed by snake_case names. %s" % rest)
-      name = snake_name
-  return Registries.optimizers[name]
-
-list_optimizers = lambda: sorted(Registries.optimizers)
-register_optimizer = Registries.optimizers.register
-
-hparams = Registries.hparams.__getitem__
-list_hparams = lambda: sorted(Registries.hparams)
-register_hparams = Registries.hparams.register
-
-ranged_hparams = Registries.ranged_hparams.__getitem__
-list_ranged_hparams = lambda: sorted(Registries.ranged_hparams)
-register_ranged_hparams = Registries.ranged_hparams.register
-
-base_problem = Registries.problems.__getitem__
-list_base_problems = lambda: sorted(Registries.problems)
-register_base_problem = Registries.problems.register
-
-# list_problems won't list all rev/copy combinations,
-# so the name is slightly confusing. Similarly, register_problem will raise an
-# error if attempting to register a value with a non-base key.
-# Keeping for back-compatibility
-list_problems = list_base_problems
-register_problem = register_base_problem
-
-
-def problem(problem_name):
-  """Get possibly copied/reversed problem registered in `base_registry`.
+  # Handle if decorator was used without parens
+  if callable(name):
+    p_cls = name
+    return decorator(p_cls, registration_name=default_name(p_cls))
 
-  Args:
-    problem_name: string problem name. See `parse_problem_name`.
+  return lambda p_cls: decorator(p_cls, name)
 
-  Returns:
-    possibly reversed/copied version of base problem registered in the given
-    registry.
-  """
-  spec = parse_problem_name(problem_name)
-  return Registries.problems[spec.base_name](
-      was_copy=spec.was_copy, was_reversed=spec.was_reversed)
 
+def problem(name):
+  """Retrieve a problem by name."""
+
+  def parse_problem_name(problem_name):
+    """Determines if problem_name specifies a copy and/or reversal.
+
+    Args:
+      problem_name: str, problem name, possibly with suffixes.
+
+    Returns:
+      base_name: A string with the base problem name.
+      was_reversed: A boolean.
+      was_copy: A boolean.
+    """
+    # Recursively strip tags until we reach a base name.
+    if problem_name.endswith("_rev"):
+      base, _, was_copy = parse_problem_name(problem_name[:-4])
+      return base, True, was_copy
+    elif problem_name.endswith("_copy"):
+      base, was_reversed, _ = parse_problem_name(problem_name[:-5])
+      return base, was_reversed, True
+    else:
+      return problem_name, False, False
+
+  base_name, was_reversed, was_copy = parse_problem_name(name)
+
+  if base_name not in _PROBLEMS:
+    all_problem_names = list_problems()
+    error_lines = ["%s not in the set of supported problems:" % base_name
+                  ] + all_problem_names
+    error_msg = "\n  * ".join(error_lines)
+    raise LookupError(error_msg)
+  return _PROBLEMS[base_name](was_reversed=was_reversed, was_copy=was_copy)
+
+
+def list_problems():
+  return sorted(list(_PROBLEMS))
 
-attack = Registries.attacks.__getitem__
-list_attacks = lambda: sorted(Registries.attacks)
-register_attack = Registries.attacks.register
 
-attack_params = Registries.attack_params.__getitem__
-list_attack_params = lambda: sorted(Registries.attack_params)
-register_attack_params = Registries.attack_params.register
+def register_attack(name=None):
+  """Register an attack HParams set. Same behaviour as register_hparams."""
+
+  def decorator(attack_fn, registration_name=None):
+    """Registers & returns attack_fn with registration_name or default name."""
+    attack_name = registration_name or default_name(attack_fn)
+    if attack_name in _ATTACKS and not tf.executing_eagerly():
+      raise LookupError("Attack %s already registered." % attack_name)
+    _ATTACKS[attack_name] = attack_fn
+    return attack_fn
+
+  # Handle if decorator was used without parens
+  if callable(name):
+    attack_fn = name
+    return decorator(attack_fn, registration_name=default_name(attack_fn))
+
+  return lambda attack_fn: decorator(attack_fn, name)
 
-pruning_params = Registries.pruning_params.__getitem__
-list_pruning_params = lambda: sorted(Registries.pruning_params)
-register_pruning_params = Registries.pruning_params.register
 
-pruning_strategy = Registries.pruning_strategies.__getitem__
-list_pruning_strategies = lambda: sorted(Registries.pruning_strategies)
-register_pruning_strategy = Registries.pruning_strategies.register
+def attacks(name):
+  """Retrieve registered attack by name."""
+  if name not in _ATTACKS:
+    error_msg = "Attack %s never registered. Sets registered:\n%s"
+    raise LookupError(
+        error_msg % (name,
+                     display_list_by_prefix(list_attacks(), starting_spaces=4)))
+  attack = _ATTACKS[name]()
+  if attack is None:
+    raise TypeError(
+        "Attack %s is None. Make sure the registered function returns a "
+        "`cleverhans.attack.Attack` object." % name)
+  return attack
+
+
+def list_attacks(prefix=None):
+  if prefix:
+    return [name for name in _ATTACKS if name.startswith(prefix)]
+  return list(_ATTACKS)
+
+
+def register_attack_params(name=None):
+  """Register an attack HParams set. Same behaviour as register_hparams."""
+
+  def decorator(ap_fn, registration_name=None):
+    """Registers & returns ap_fn with registration_name or default name."""
+    ap_name = registration_name or default_name(ap_fn)
+    if ap_name in _ATTACK_PARAMS and not tf.executing_eagerly():
+      raise LookupError("Attack HParams set %s already registered." % ap_name)
+    _ATTACK_PARAMS[ap_name] = ap_fn
+    return ap_fn
+
+  # Handle if decorator was used without parens
+  if callable(name):
+    ap_fn = name
+    return decorator(ap_fn, registration_name=default_name(ap_fn))
+
+  return lambda ap_fn: decorator(ap_fn, name)
+
+
+def attack_params(name):
+  """Retrieve registered aparams by name."""
+  if name not in _ATTACK_PARAMS:
+    error_msg = "Attack HParams set %s never registered. Sets registered:\n%s"
+    raise LookupError(
+        error_msg %
+        (name, display_list_by_prefix(list_attack_params(), starting_spaces=4)))
+  ap = _ATTACK_PARAMS[name]()
+  if ap is None:
+    raise TypeError("Attack HParams %s is None. Make sure the registered "
+                    "function returns the HParams object." % name)
+  return ap
+
+
+def list_attack_params(prefix=None):
+  if prefix:
+    return [name for name in _ATTACK_PARAMS if name.startswith(prefix)]
+  return list(_ATTACK_PARAMS)
+
+
+def register_pruning_params(name=None):
+  """Register an pruning HParams set. Same behaviour as register_hparams."""
+
+  def decorator(pp_fn, registration_name=None):
+    """Registers & returns pp_fn with registration_name or default name."""
+    pp_name = registration_name or default_name(pp_fn)
+    if pp_name in _PRUNING_PARAMS and not tf.executing_eagerly():
+      raise LookupError("Pruning HParams set %s already registered." % pp_name)
+    _PRUNING_PARAMS[pp_name] = pp_fn
+    return pp_fn
+
+  # Handle if decorator was used without parens
+  if callable(name):
+    pp_fn = name
+    return decorator(pp_fn, registration_name=default_name(pp_fn))
+
+  return lambda pp_fn: decorator(pp_fn, name)
+
+
+def pruning_params(name):
+  """Retrieve registered pruning params by name."""
+  if name not in _PRUNING_PARAMS:
+    error_msg = "Pruning HParams set %s never registered. Sets registered:\n%s"
+    raise LookupError(error_msg % (
+        name, display_list_by_prefix(list_pruning_params(), starting_spaces=4)))
+  pp = _PRUNING_PARAMS[name]()
+  if pp is None:
+    raise TypeError("Pruning HParams %s is None. Make sure the registered "
+                    "function returns the HParams object." % name)
+  return pp
+
+
+def list_pruning_params(prefix=None):
+  if prefix:
+    return [name for name in _PRUNING_PARAMS if name.startswith(prefix)]
+  return list(_PRUNING_PARAMS)
+
+
+def register_pruning_strategy(name=None):
+  """Register an pruning strategy. Same behaviour as register_hparams."""
+
+  def decorator(ps_fn, registration_name=None):
+    """Registers & returns ps_fn with registration_name or default name."""
+    ps_name = registration_name or default_name(ps_fn)
+    if ps_name in _PRUNING_STRATEGY and not tf.executing_eagerly():
+      raise LookupError("Pruning strategy %s already registered." % ps_name)
+    _PRUNING_STRATEGY[ps_name] = ps_fn
+    return ps_fn
+
+  # Handle if decorator was used without parens
+  if callable(name):
+    ps_fn = name
+    return decorator(ps_fn, registration_name=default_name(ps_fn))
+
+  return lambda ps_fn: decorator(ps_fn, name)
+
+
+def pruning_strategies(name):
+  """Retrieve registered pruning strategies by name."""
+  if name not in _PRUNING_STRATEGY:
+    error_msg = "Pruning strategy set %s never registered. Sets registered:\n%s"
+    raise LookupError(
+        error_msg % (name,
+                     display_list_by_prefix(
+                         list_pruning_strategies(), starting_spaces=4)))
+  ps = _PRUNING_STRATEGY[name]
+  if ps is None:
+    raise TypeError("Pruning strategy %s is None. Make sure to register the "
+                    "function." % name)
+  return ps
 
 
-# deprecated functions - plurals inconsistent with rest
-# deprecation decorators added 2019-01-25
-attacks = framework.deprecated(None, "Use registry.attack")(attack)
-pruning_strategies = framework.deprecated(
-    None, "Use registry.pruning_strategy")(pruning_strategy)
+def list_pruning_strategies(prefix=None):
+  if prefix:
+    return [name for name in _PRUNING_STRATEGY if name.startswith(prefix)]
+  return list(_PRUNING_STRATEGY)
 
 
 def display_list_by_prefix(names_list, starting_spaces=0):
@@ -563,9 +551,6 @@ def help_string():
   Problems:
 %s
 
-  Optimizers:
-%s
-
   Attacks:
 %s
 
@@ -578,17 +563,16 @@ def help_string():
   Pruning Strategies:
 %s
 """
-  lists = tuple(
+  m, hp, rhp, probs, atks, ap, pp, ps = [
       display_list_by_prefix(entries, starting_spaces=4) for entries in [
           list_models(),
           list_hparams(),
           list_ranged_hparams(),
-          list_base_problems(),
-          list_optimizers(),
+          list_problems(),
           list_attacks(),
           list_attack_params(),
           list_pruning_params(),
           list_pruning_strategies(),
       ]
-  )
-  return help_str % lists
+  ]
+  return help_str % (m, hp, rhp, probs, atks, ap, pp, ps)
diff --git a/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index dc1cdb579..c2515b3ba 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -26,72 +26,11 @@
 
 # pylint: disable=unused-variable
 
-class RegistryClassTest(tf.test.TestCase):
-  """Test of base registry.Registry class."""
-
-  def testGetterSetter(self):
-    r = registry.Registry("test_registry")
-    r["hello"] = lambda: "world"
-    r["a"] = lambda: "b"
-    self.assertEqual(r["hello"](), "world")
-    self.assertEqual(r["a"](), "b")
-
-  def testDefaultKeyFn(self):
-    r = registry.Registry("test", default_key_fn=lambda x: x().upper())
-    r.register()(lambda: "hello")
-    self.assertEqual(r["HELLO"](), "hello")
-
-  def testNoKeyProvided(self):
-    r = registry.Registry("test")
-    def f():
-      return 3
-    r.register(f)
-    self.assertEqual(r['f'](), 3)
-
-  def testMembership(self):
-    r = registry.Registry("test_registry")
-    r["a"] = lambda: None
-    r["b"] = lambda: 4
-    self.assertTrue("a" in r)
-    self.assertTrue("b" in r)
-
-  def testIteration(self):
-    r = registry.Registry("test_registry")
-    r["a"] = lambda: None
-    r["b"] = lambda: 4
-    self.assertEqual(sorted(r), ["a", "b"])
-
-  def testLen(self):
-    r = registry.Registry("test_registry")
-    self.assertEqual(len(r), 0)
-    r["a"] = lambda: None
-    self.assertEqual(len(r), 1)
-    r["b"] = lambda: 4
-    self.assertEqual(len(r), 2)
-
-  def testTransformer(self):
-    r = registry.Registry(
-        "test_registry", value_transformer=lambda x, y: x + y())
-    r.register(3)(lambda: 5)
-    r.register(10)(lambda: 12)
-    self.assertEqual(r[3], 8)
-    self.assertEqual(r[10], 22)
-    self.assertEqual(set(r.values()), set((8, 22)))
-    self.assertEqual(set(r.items()), set(((3, 8), (10, 22))))
-
-  def testGet(self):
-    r = registry.Registry('test_registry', value_transformer=lambda k, v: v())
-    r["a"] = lambda: "xyz"
-    self.assertEqual(r.get("a"), "xyz")
-    self.assertEqual(r.get("a", 3), "xyz")
-    self.assertIsNone(r.get("b"))
-    self.assertEqual(r.get("b", 3), 3)
-
 
 class ModelRegistryTest(tf.test.TestCase):
 
   def setUp(self):
-    registry.Registries.models._clear()
+    registry._reset()
 
   def testT2TModelRegistration(self):
 
@@ -121,7 +60,7 @@ def model_fn():
     self.assertTrue(model is model_fn)
 
   def testUnknownModel(self):
-    with self.assertRaisesRegexp(KeyError, "never registered"):
+    with self.assertRaisesRegexp(LookupError, "never registered"):
       registry.model("not_registered")
 
   def testDuplicateRegistration(self):
@@ -130,7 +69,7 @@ def testDuplicateRegistration(self):
     def m1():
       pass
 
-    with self.assertRaisesRegexp(KeyError, "already registered"):
+    with self.assertRaisesRegexp(LookupError, "already registered"):
 
       @registry.register_model("m1")
       def m2():
@@ -149,82 +88,10 @@ def m2():
     self.assertSetEqual(set(["m1", "m2"]), set(registry.list_models()))
 
 
-class OptimizerRegistryTest(tf.test.TestCase):
-  def setUp(self):
-    registry.Registries.optimizers._clear()
-
-  def testRegistration(self):
-    @registry.register_optimizer
-    def my_optimizer(learning_rate, hparams):
-      return 3
-
-    @registry.register_optimizer('my_other_optimizer')
-    def another_optimizer(learning_rate, hparams):
-      return 5
-
-    self.assertEqual(registry.optimizer("my_optimizer"), my_optimizer)
-    self.assertEqual(
-        registry.optimizer("my_other_optimizer"), another_optimizer)
-
-  def testMembership(self):
-    @registry.register_optimizer
-    def my_optimizer(learning_rate, hparams):
-      return 3
-
-    @registry.register_optimizer("my_other_optimizer")
-    def another_optimizer(learning_rate, hparams):
-      return 5
-
-    self.assertTrue("my_optimizer" in registry.Registries.optimizers)
-    self.assertTrue("my_other_optimizer" in registry.Registries.optimizers)
-    self.assertFalse("another_optimizer" in registry.Registries.optimizers)
-    self.assertEqual(len(registry.Registries.optimizers), 2)
-
-  def testArgErrorCheck(self):
-    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
-      registry.Registries.optimizers.register('OneArgs')(lambda x: 4)
-    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
-      registry.Registries.optimizers.register('ThreeArgs')(
-          lambda x, y, z: 4)
-    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
-      registry.Registries.optimizers.register('NArgs')(lambda *args: 4)
-    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
-      registry.Registries.optimizers.register("Kwargs")(lambda **kargs: 4)
-    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
-      registry.Registries.optimizers.register("TwoAndKwargs")(
-          lambda a, b, **kargs: 4)
-
-  def testMultipleRegistration(self):
-    @registry.register_optimizer
-    def my_optimizer(learning_rate, hparams):
-      return 3
-
-    with self.assertRaisesRegexp(KeyError, "already registered"):
-
-      @registry.register_optimizer("my_optimizer")
-      def another_fn(learning_rate, hparams):
-        return 5
-
-  def testUnknownOptimizer(self):
-    with self.assertRaisesRegexp(KeyError, "never registered"):
-      registry.optimizer("not_registered_optimizer")
-
-  def testGetterSetterInterface(self):
-    def f(x, y):
-      return 3
-
-    k = 'blah'
-    registry.Registries.optimizers[k] = f
-    self.assertEqual(registry.optimizer(k), f)
-    self.assertEqual(registry.Registries.optimizers[k], f)
-    self.assertEqual(registry.Registries.optimizers[k], registry.optimizer(k))
-
-
 class HParamRegistryTest(tf.test.TestCase):
 
   def setUp(self):
-    registry.Registries.hparams._clear()
-    registry.Registries.ranged_hparams._clear()
+    registry._reset()
 
   def testHParamSet(self):
 
@@ -254,9 +121,9 @@ def my_hparams_range(_):
     self.assertTrue(registry.ranged_hparams("a") is my_hparams_range)
 
   def testUnknownHparams(self):
-    with self.assertRaisesRegexp(KeyError, "never registered"):
+    with self.assertRaisesRegexp(LookupError, "never registered"):
       registry.hparams("not_registered")
-    with self.assertRaisesRegexp(KeyError, "never registered"):
+    with self.assertRaisesRegexp(LookupError, "never registered"):
       registry.ranged_hparams("not_registered")
 
   def testNoneHparams(self):
@@ -327,7 +194,34 @@ def rhp_bad2(a, b):  # pylint: disable=unused-argument
         pass
 
 
-class RegistryHelpTest(tf.test.TestCase):
+class CreateRegistry(tf.test.TestCase):
+  """Test class for `create_registry`."""
+
+  def testCreateRegistry(self):
+    my_registry = registry.create_registry("test_reg1")
+    self.assertIs(my_registry, registry.registry("test_reg1"))
+
+    # Use as decorator on a fn
+    @my_registry.register("foo")
+    def some_fn(num):
+      return num + 2
+
+    # Register a regular object
+    pod_obj = 4
+    my_registry.register("bar")(pod_obj)
+
+    # Register a class
+    @my_registry.register("foobar")
+    class A(object):
+      pass
+
+    self.assertEqual(9, my_registry.get("foo")(7))
+    self.assertEqual(["bar", "foo", "foobar"], my_registry.list())
+    foobar = my_registry.get("foobar")
+    self.assertTrue(isinstance(foobar(), A))
+
+
+class RegistryTest(tf.test.TestCase):
   """Test class for common functions."""
 
   def testRegistryHelp(self):

From 5143a056ab26399f38a8d6a4e98850d9208a9f65 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 28 Jan 2019 17:02:30 -0800
Subject: [PATCH 1589/2720] Add ModalityType enum and set Problem modalities to
 the modality type.

For now, Problems with custom modalities and models with overrided modalities still specify the modality as a class.

Future changes:
+ Remove modality classes, splitting the code into functions.
+ Reduce ModalityType's enum size.

PiperOrigin-RevId: 231315839
---
 tensor2tensor/data_generators/algorithmic.py  |  4 +-
 tensor2tensor/data_generators/allen_brain.py  |  4 +-
 tensor2tensor/data_generators/babi_qa.py      |  2 +-
 .../data_generators/bair_robot_pushing.py     |  4 +-
 tensor2tensor/data_generators/celeba.py       |  4 +-
 tensor2tensor/data_generators/celebahq.py     |  2 +-
 tensor2tensor/data_generators/cifar.py        |  4 +-
 tensor2tensor/data_generators/fsns.py         |  4 +-
 .../data_generators/gene_expression.py        |  4 +-
 .../data_generators/google_robot_pushing.py   |  4 +-
 tensor2tensor/data_generators/gym_env.py      | 12 +--
 tensor2tensor/data_generators/ice_parsing.py  |  4 +-
 tensor2tensor/data_generators/image_utils.py  |  4 +-
 tensor2tensor/data_generators/imagenet.py     |  4 +-
 tensor2tensor/data_generators/lambada.py      |  4 +-
 tensor2tensor/data_generators/problem.py      | 16 +++-
 .../data_generators/problem_hparams.py        | 20 ++---
 .../data_generators/speech_recognition.py     |  4 +-
 .../data_generators/text_problems.py          | 18 ++---
 tensor2tensor/data_generators/timeseries.py   |  4 +-
 .../data_generators/video_generated.py        |  4 +-
 tensor2tensor/data_generators/video_utils.py  |  4 +-
 tensor2tensor/data_generators/vqa.py          |  6 +-
 tensor2tensor/layers/common_hparams.py        |  2 +-
 tensor2tensor/layers/modalities.py            | 76 +++++++++++++++++++
 25 files changed, 151 insertions(+), 67 deletions(-)

diff --git a/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index edacffd49..790dcb854 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -84,8 +84,8 @@ def generator_eos(nbr_symbols, max_length, nbr_cases):
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
     vocab_size = self.num_symbols + text_encoder.NUM_RESERVED_TOKENS
-    p.modality = {"inputs": modalities.SymbolModality,
-                  "targets": modalities.SymbolModality}
+    p.modality = {"inputs": modalities.ModalityType.SYMBOL,
+                  "targets": modalities.ModalityType.SYMBOL}
     p.vocab_size = {"inputs": vocab_size,
                     "targets": vocab_size}
     p.input_space_id = problem.SpaceID.DIGIT_0
diff --git a/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index 2fece298c..30a8c9696 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -375,8 +375,8 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.IdentityModality,
-                  "targets": modalities.IdentityModality}
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
     p.vocab_size = {"inputs": 256,
                     "targets": 256}
     p.batch_size_multiplier = 256
diff --git a/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index 56d8d66ae..b6e4c6f67 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -425,7 +425,7 @@ def hparams(self, defaults, unused_model_hparams):
     (super(BabiQa, self).hparams(defaults, unused_model_hparams))
     p = defaults
     num_classes = self._encoders["targets"].vocab_size
-    p.modality = {"targets": modalities.ClassLabelModality}
+    p.modality = {"targets": modalities.ModalityType.CLASS_LABEL}
     p.vocab_size = {"targets": num_classes}
 
   def example_reading_spec(self):
diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index a92303013..63eb6572a 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -112,8 +112,8 @@ def extra_reading_spec(self):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.VideoModality,
-                  "targets": modalities.VideoModality}
+    p.modality = {"inputs": modalities.ModalityType.VIDEO,
+                  "targets": modalities.ModalityType.VIDEO}
     p.vocab_size = {"inputs": 256,
                     "targets": 256}
 
diff --git a/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index 3a80718a1..271c9ecd4 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -56,8 +56,8 @@ class ImageCeleba(image_utils.ImageProblem):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.IdentityModality,
-                  "targets": modalities.IdentityModality}
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
     p.vocab_size = {"inputs": 256,
                     "targets": 256}
     p.batch_size_multiplier = 256
diff --git a/tensor2tensor/data_generators/celebahq.py b/tensor2tensor/data_generators/celebahq.py
index e8a01bfe2..960d49e2a 100644
--- a/tensor2tensor/data_generators/celebahq.py
+++ b/tensor2tensor/data_generators/celebahq.py
@@ -79,7 +79,7 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.batch_size_multiplier = 1
-    p.modality = {"inputs": modalities.IdentityModality}
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY}
     p.vocab_size = {"inputs": 256}
     p.input_space_id = 1
 
diff --git a/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index d5fcafc04..23a106491 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -450,8 +450,8 @@ def preprocess_example(self, example, unused_mode, unused_hparams):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.IdentityModality,
-                  "targets": modalities.IdentityModality}
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
     p.vocab_size = {"inputs": 256,
                     "targets": 256}
     p.batch_size_multiplier = 256
diff --git a/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
index b482c8a70..c7d566b8a 100644
--- a/tensor2tensor/data_generators/fsns.py
+++ b/tensor2tensor/data_generators/fsns.py
@@ -62,8 +62,8 @@ def feature_encoders(self, data_dir):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.ImageModality,
-                  "targets": modalities.SymbolModality}
+    p.modality = {"inputs": modalities.ModalityType.IMAGE,
+                  "targets": modalities.ModalityType.SYMBOL}
     p.vocab_size = {"inputs": 256,
                     "targets": self._encoders["targets"].vocab_size}
     p.batch_size_multiplier = 256
diff --git a/tensor2tensor/data_generators/gene_expression.py b/tensor2tensor/data_generators/gene_expression.py
index 4e82000f2..04b009438 100644
--- a/tensor2tensor/data_generators/gene_expression.py
+++ b/tensor2tensor/data_generators/gene_expression.py
@@ -142,8 +142,8 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.SymbolModality,
-                  "targets": modalities.RealLogPoissonLossModality}
+    p.modality = {"inputs": modalities.ModalityType.SYMBOL,
+                  "targets": modalities.ModalityType.REAL_LOG_POISSON_LOSS}
     p.vocab_size = {"inputs": self._encoders["inputs"].vocab_size,
                     "targets": self.num_output_predictions}
     p.input_space_id = problem.SpaceID.DNA
diff --git a/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index 5b295fccf..b37db990f 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -133,7 +133,7 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.VideoModality,
-                  "targets": modalities.VideoModality}
+    p.modality = {"inputs": modalities.ModalityType.VIDEO,
+                  "targets": modalities.ModalityType.VIDEO}
     p.vocab_size = {"inputs": 256,
                     "targets": 256}
diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 43d2eca10..95fe1c2a2 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -105,12 +105,12 @@ def num_rewards(self):
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.modality = {
-        "inputs": modalities.VideoModality,
-        "input_reward": modalities.SymbolModalityWeightsAll,
-        "input_action": modalities.SymbolModalityWeightsAll,
-        "targets": modalities.VideoModality,
-        "target_reward": modalities.SymbolModalityWeightsAll,
-        "target_action": modalities.SymbolModalityWeightsAll,
+        "inputs": modalities.ModalityType.VIDEO,
+        "input_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "input_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "targets": modalities.ModalityType.VIDEO,
+        "target_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "target_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
     }
     p.vocab_size = {
         "inputs": 256,
diff --git a/tensor2tensor/data_generators/ice_parsing.py b/tensor2tensor/data_generators/ice_parsing.py
index fd07da10e..d23576093 100644
--- a/tensor2tensor/data_generators/ice_parsing.py
+++ b/tensor2tensor/data_generators/ice_parsing.py
@@ -108,8 +108,8 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.SymbolModality,
-                  "targets": modalities.SymbolModality}
+    p.modality = {"inputs": modalities.ModalityType.SYMBOL,
+                  "targets": modalities.ModalityType.SYMBOL}
     p.vocab_size = {"inputs": self._encoders["inputs"].vocab_size,
                     "targets": self.targeted_vocab_size}
     p.input_space_id = self.input_space_id
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index db5ddbb44..30db55bbe 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -355,8 +355,8 @@ def feature_encoders(self, data_dir):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.ImageModality,
-                  "targets": modalities.SymbolModality}
+    p.modality = {"inputs": modalities.ModalityType.IMAGE,
+                  "targets": modalities.ModalityType.SYMBOL}
     p.vocab_size = {"inputs": 256,
                     "targets": self._encoders["targets"].vocab_size}
     p.batch_size_multiplier = 256
diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 71e6807a7..8bd5bb5ff 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -386,8 +386,8 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.IdentityModality,
-                  "targets": modalities.IdentityModality}
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
     p.vocab_size = {"inputs": 256,
                     "targets": 256}
     p.batch_size_multiplier = 256
diff --git a/tensor2tensor/data_generators/lambada.py b/tensor2tensor/data_generators/lambada.py
index 8f3600ac1..9af50cffc 100644
--- a/tensor2tensor/data_generators/lambada.py
+++ b/tensor2tensor/data_generators/lambada.py
@@ -358,8 +358,8 @@ def hparams(self, defaults, unused_model_hparams):
     """
 
     p = defaults
-    p.modality = {"inputs": modalities.SymbolModality,
-                  "targets": modalities.ClassLabelModality}
+    p.modality = {"inputs": modalities.ModalityType.SYMBOL,
+                  "targets": modalities.ModalityType.CLASS_LABEL}
     p.vocab_size = {"inputs": self._encoders["inputs"].vocab_size,
                     "targets": self._encoders["targets"].vocab_size}
 
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 858701e4e..062d26926 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -25,6 +25,7 @@
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import mlperf_log
@@ -981,7 +982,7 @@ def _create_modalities(problem_hparams, model_hparams):
 
   Args:
     problem_hparams: tf.contrib.training.HParams for the Problem. It must have
-      modality which is a dict of strings to Modality classes.
+      modality which is a dict of strings to ModalityTypes or Modality classes.
     model_hparams: tf.contrib.training.HParams for the model. It may have
       input_modalities and target_modality, which will override
       problem_hparams' modality input and target keys.
@@ -991,14 +992,21 @@ def _create_modalities(problem_hparams, model_hparams):
   """
   modality_overrides = getattr(model_hparams, "modality", {})
   modality = {}
-  for feature_name, modality_cls in six.iteritems(problem_hparams.modality):
+  for feature_name, modality_type in six.iteritems(problem_hparams.modality):
     vocab_size = problem_hparams.vocab_size[feature_name]
     # If needed for using a pre-trained model's vocabulary where extra indices
     # were allocated for adding new tasks with unique task ids.
     if (hasattr(model_hparams, "multiproblem_vocab_size") and
         model_hparams.multiproblem_vocab_size > 0):
       vocab_size = model_hparams.multiproblem_vocab_size
-    modality_cls = modality_overrides.get(feature_name, modality_cls)
+    # Override modality using to the associated value in modality_overrides.
+    modality_type = modality_overrides.get(feature_name, modality_type)
+    # Each modality is a ModalityType or class. If ModalityType, get the
+    # corresponding class.
+    if modality_type in modalities.ModalityType.get_choices():
+      modality_cls = getattr(modalities, modality_type)
+    else:
+      modality_cls = modality_type
     modality[feature_name] = modality_cls(model_hparams, vocab_size)
   problem_hparams.modality = modality
 
@@ -1027,7 +1035,7 @@ def _default_hparams():
 
       # Modalities used to map from features to a space compatible with
       # chosen model architecture. It comprises key-value pairs of a feature
-      # name (str) and its modality class.
+      # name (str) and its modality type.
       modality={},
 
       # Identifiers used to tell the model which input/target space will be
diff --git a/tensor2tensor/data_generators/problem_hparams.py b/tensor2tensor/data_generators/problem_hparams.py
index 98ae21052..f67742de4 100644
--- a/tensor2tensor/data_generators/problem_hparams.py
+++ b/tensor2tensor/data_generators/problem_hparams.py
@@ -68,8 +68,8 @@ def feature_encoders(self, _):
 
   def hparams(self, defaults, model_hparams):
     hp = defaults
-    hp.modality = {"inputs": modalities.SpeechRecognitionModality,
-                   "targets": modalities.SymbolModality}
+    hp.modality = {"inputs": modalities.ModalityType.SPEECH_RECOGNITION,
+                   "targets": modalities.ModalityType.SYMBOL}
     hp.vocab_size = {"inputs": None,
                      "targets": 256}
 
@@ -93,8 +93,8 @@ def feature_encoders(self, data_dir):
 
   def hparams(self, defaults, model_hparams):
     hp = defaults
-    hp.modality = {"inputs": modalities.SpeechRecognitionModality,
-                   "targets": modalities.SymbolModality}
+    hp.modality = {"inputs": modalities.ModalityType.SPEECH_RECOGNITION,
+                   "targets": modalities.ModalityType.SYMBOL}
     hp.vocab_size = {
         "inputs": None,
         "targets": self.get_feature_encoders()["targets"].vocab_size,
@@ -130,8 +130,8 @@ def feature_encoders(self, data_dir):
 
   def hparams(self, defaults, model_hparams):
     hp = defaults
-    hp.modality = {"inputs": modalities.SymbolModality,
-                   "targets": modalities.SymbolModality}
+    hp.modality = {"inputs": modalities.ModalityType.SYMBOL,
+                   "targets": modalities.ModalityType.SYMBOL}
     hp.vocab_size = {
         "inputs": self.get_feature_encoders()["inputs"].vocab_size,
         "targets": self.get_feature_encoders()["targets"].vocab_size,
@@ -174,8 +174,8 @@ def feature_encoders(self, data_dir):
 
   def hparams(self, defaults, model_hparams):
     hp = defaults
-    hp.modality = {"inputs": modalities.SymbolModality,
-                   "targets": modalities.SymbolModality}
+    hp.modality = {"inputs": modalities.ModalityType.SYMBOL,
+                   "targets": modalities.ModalityType.SYMBOL}
     hp.vocab_size = {
         "inputs": self.get_feature_encoders()["inputs"].vocab_size,
         "targets": self.get_feature_encoders()["targets"].vocab_size,
@@ -194,8 +194,8 @@ def __init__(self, input_vocab_size, target_vocab_size):
 
   def hparams(self, defaults, model_hparams):
     hp = defaults
-    hp.modality = {"inputs": modalities.SymbolModality,
-                   "targets": modalities.SymbolModality}
+    hp.modality = {"inputs": modalities.ModalityType.SYMBOL,
+                   "targets": modalities.ModalityType.SYMBOL}
     hp.vocab_size = {"inputs": self.input_vocab_size,
                      "targets": self.target_vocab_size}
 
diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 04891d50e..258af10f3 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -60,8 +60,8 @@ def hparams(self, defaults, model_hparams):
     p.add_hparam("num_zeropad_frames", 250)
 
     p = defaults
-    p.modality = {"inputs": modalities.SpeechRecognitionModality,
-                  "targets": modalities.SymbolModality}
+    p.modality = {"inputs": modalities.ModalityType.SPEECH_RECOGNITION,
+                  "targets": modalities.ModalityType.SYMBOL}
     p.vocab_size = {"inputs": None,
                     "targets": 256}
 
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 71561da05..767d6ffa0 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -340,22 +340,22 @@ def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.stop_at_eos = int(True)
 
-    p.modality = {"targets": modalities.SymbolModality}
+    p.modality = {"targets": modalities.ModalityType.SYMBOL}
     p.vocab_size = {"targets": self._encoders["targets"].vocab_size}
     if self.has_inputs:
-      p.modality["inputs"] = modalities.SymbolModality
+      p.modality["inputs"] = modalities.ModalityType.SYMBOL
       p.vocab_size["inputs"] = self._encoders["inputs"].vocab_size
     if self.vocab_type == VocabType.CHARACTER:
       p.loss_multiplier = 2.0
 
     if self.packed_length:
       if self.has_inputs:
-        p.modality["inputs_segmentation"] = modalities.IdentityModality
-        p.modality["inputs_position"] = modalities.IdentityModality
+        p.modality["inputs_segmentation"] = modalities.ModalityType.IDENTITY
+        p.modality["inputs_position"] = modalities.ModalityType.IDENTITY
         p.vocab_size["inputs_segmentation"] = None
         p.vocab_size["inputs_position"] = None
-      p.modality["targets_segmentation"] = modalities.IdentityModality
-      p.modality["targets_position"] = modalities.IdentityModality
+      p.modality["targets_segmentation"] = modalities.ModalityType.IDENTITY
+      p.modality["targets_position"] = modalities.ModalityType.IDENTITY
       p.vocab_size["targets_segmentation"] = None
       p.vocab_size["targets_position"] = None
 
@@ -426,7 +426,7 @@ def hparams(self, defaults, unused_model_hparams):
     (super(QuestionAndContext2TextProblem, self)
      .hparams(defaults, unused_model_hparams))
     p = defaults
-    p.modality["context"] = modalities.SymbolModality
+    p.modality["context"] = modalities.ModalityType.SYMBOL
     p.vocab_size["context"] = self._encoders["context"].vocab_size
     if self.packed_length:
       raise NotImplementedError("QuestionAndContext2Text does not "
@@ -530,8 +530,8 @@ def feature_encoders(self, data_dir):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.SymbolModality,
-                  "targets": modalities.ClassLabelModality}
+    p.modality = {"inputs": modalities.ModalityType.SYMBOL,
+                  "targets": modalities.ModalityType.CLASS_LABEL}
     p.vocab_size = {"inputs": self._encoders["inputs"].vocab_size,
                     "targets": self.num_classes}
 
diff --git a/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index eb7fd7a9f..23154ad73 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -151,8 +151,8 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.RealL2LossModality,
-                  "targets": modalities.RealL2LossModality}
+    p.modality = {"inputs": modalities.ModalityType.REAL_L2_LOSS,
+                  "targets": modalities.ModalityType.REAL_L2_LOSS}
     p.vocab_size = {"inputs": self.num_series,
                     "targets": self.num_series}
     p.input_space_id = problem.SpaceID.REAL
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index d9607f623..419aefa44 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -93,8 +93,8 @@ def extra_reading_spec(self):
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.modality = {
-        "inputs": modalities.VideoModality,
-        "targets": modalities.VideoModality,
+        "inputs": modalities.ModalityType.VIDEO,
+        "targets": modalities.ModalityType.VIDEO,
     }
     p.vocab_size = {
         "inputs": 256,
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 3f1c1d54a..818f6c3bc 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -773,8 +773,8 @@ def example_reading_spec(self):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.ImageModality,
-                  "targets": modalities.ClassLabelModality}
+    p.modality = {"inputs": modalities.ModalityType.IMAGE,
+                  "targets": modalities.ModalityType.CLASS_LABEL}
     p.vocab_size = {"inputs": 256,
                     "targets": self.num_classes}
     p.input_space_id = problem.SpaceID.IMAGE
diff --git a/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
index 5a0d5fc46..2ba91942d 100644
--- a/tensor2tensor/data_generators/vqa.py
+++ b/tensor2tensor/data_generators/vqa.py
@@ -131,9 +131,9 @@ def hparams(self, defaults, unused_model_hparams):
     targets_encoder = self._encoders["targets"]
 
     p.modality = {
-        "inputs": modalities.IdentityModality,
-        "question": modalities.SymbolModality,
-        "targets": modalities.MultiLabelModality,
+        "inputs": modalities.ModalityType.IDENTITY,
+        "question": modalities.ModalityType.SYMBOL,
+        "targets": modalities.ModalityType.MULTI_LABEL,
     }
     p.vocab_size = {
         "inputs": None,
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 2bbfb7618..b422026ce 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -180,7 +180,7 @@ def basic_params1():
       symbol_modality_skip_top=False,
       # Modalities used to map from features to a space compatible with
       # chosen model architecture. It comprises key-value pairs of a feature
-      # name (str) and its modality class.
+      # name (str) and its modality type.
       modality={},
       # The maximum length of "input" sequence.
       # Sequences longer than this value will be truncated. 0 or negative values
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 49de33f13..09e8b9b1c 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -1091,3 +1091,79 @@ def top(self, body_output, _):
       x = body_output
       x = tf.expand_dims(x[:, -1], 1)  # Pick the last timestep
       return tf.layers.dense(x, self._vocab_size)
+
+
+class ModalityType(object):
+  """Types of modalities."""
+
+  SYMBOL = "SymbolModality"
+  SYMBOL_WEIGHTS_ALL = "SymbolModalityWeightsAll"
+  SYMBOL_ONE_HOT = "SymbolModalityOneHot"
+  CTC_SYMBOL = "CTCSymbolModality"
+  IMAGE = "ImageModality"
+  IMAGE_CHANNEL_COMPRESS = "ImageChannelCompressModality"
+  IMAGE_CHANNEL_BOTTOM_IDENTITY = "ImageChannelBottomIdentityModality"
+  IMAGE_CHANNEL_EMBEDDINGS_BOTTOM = "ImageChannelEmbeddingsBottom"
+  AUDIO = "AudioModality"
+  AUDIO_SPECTRAL = "AudioSpectralModality"
+  SPEECH_RECOGNITION = "SpeechRecognitionModality"
+  VIDEO = "VideoModality"
+  VIDEO_BITWISE = "VideoModalityBitwise"
+  VIDEO_PIXEL_NOISE = "VideoModalityPixelNoise"
+  VIDEO_L1 = "VideoModalityL1"
+  VIDEO_L2 = "VideoModalityL2"
+  VIDEO_L2_RAW = "VideoModalityL2Raw"
+  VIDEO_L1_RAW = "VideoModalityL1Raw"
+  CLASS_LABEL = "ClassLabelModality"
+  VIDEO_IDENTITY = "VideoModalityIdentity"
+  MULTI_LABEL = "MultiLabelModality"
+  ONE_HOT_CLASS_LABEL = "OneHotClassLabelModality"
+  IDENTITY = "IdentityModality"
+  GENERIC_L2_LOSS = "GenericL2LossModality"
+  REAL = "RealModality"
+  REAL_L2_LOSS = "RealL2LossModality"
+  REAL_LOG_POISSON_LOSS = "RealLogPoissonLossModality"
+  IDENTITY_SYMBOL = "IdentitySymbolModality"
+  SIGMOID_CLASS_LABEL = "SigmoidClassLabelModality"
+  SIGMOID_MAX_POOLING_CLASS_LABEL = "SigmoidMaxPoolingClassLabelModality"
+  SOFTMAX_MAX_POOLING_CLASS_LABEL = "SoftmaxMaxPoolingClassLabelModality"
+  SOFTMAX_AVERAGE_POOLING_CLASS_LABEL = "SoftmaxAveragePoolingClassLabelModality"
+  SOFTMAX_LAST_TIMESTEP_CLASS_LABEL = "SoftmaxLastTimestepClassLabelModality"
+
+  @staticmethod
+  def get_choices():
+    return [
+        ModalityType.SYMBOL,
+        ModalityType.SYMBOL_WEIGHTS_ALL,
+        ModalityType.SYMBOL_ONE_HOT,
+        ModalityType.CTC_SYMBOL,
+        ModalityType.IMAGE,
+        ModalityType.IMAGE_CHANNEL_COMPRESS,
+        ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
+        ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM,
+        ModalityType.AUDIO,
+        ModalityType.AUDIO_SPECTRAL,
+        ModalityType.SPEECH_RECOGNITION,
+        ModalityType.VIDEO,
+        ModalityType.VIDEO_BITWISE,
+        ModalityType.VIDEO_PIXEL_NOISE,
+        ModalityType.VIDEO_L1,
+        ModalityType.VIDEO_L2,
+        ModalityType.VIDEO_L2_RAW,
+        ModalityType.VIDEO_L1_RAW,
+        ModalityType.CLASS_LABEL,
+        ModalityType.VIDEO_IDENTITY,
+        ModalityType.MULTI_LABEL,
+        ModalityType.ONE_HOT_CLASS_LABEL,
+        ModalityType.IDENTITY,
+        ModalityType.GENERIC_L2_LOSS,
+        ModalityType.REAL,
+        ModalityType.REAL_L2_LOSS,
+        ModalityType.REAL_LOG_POISSON_LOSS,
+        ModalityType.IDENTITY_SYMBOL,
+        ModalityType.SIGMOID_CLASS_LABEL,
+        ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
+        ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL,
+        ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
+        ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
+    ]

From 7ca87dedc7ba27a214441e0b6146e5d7fd24ddbd Mon Sep 17 00:00:00 2001
From: Pablo Samuel Castro <psc@google.com>
Date: Tue, 29 Jan 2019 11:45:00 -0800
Subject: [PATCH 1590/2720] Update Dopamine path to use new discrete_domains
 directory.

PiperOrigin-RevId: 231443642
---
 tensor2tensor/rl/dopamine_connector.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 89be6dc6b..cc59d7dfa 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -40,7 +40,7 @@
 except ImportError:
   cv2 = None
 try:
-  from dopamine.atari import run_experiment
+  from dopamine.discrete_domains import run_experiment
 except ImportError:
   run_experiment = None
 # pylint: enable=g-import-not-at-top
@@ -182,7 +182,7 @@ def get_create_agent(agent_kwargs):
   def create_agent(sess, environment, summary_writer=None):
     """Creates a DQN agent.
 
-    Simplified version of `dopamine.atari.train.create_agent`
+    Simplified version of `dopamine.discrete_domains.train.create_agent`
 
     Args:
       sess: a session

From ba2388a2435960c9e2683e718d0a7a37c035ef3a Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Tue, 29 Jan 2019 12:07:33 -0800
Subject: [PATCH 1591/2720] Replace model-overrided modalities with
 ModalityType.

Note this prepares for the next CL, which removes built-in modality classes. For now, custom modalities remain supported as classes.

PiperOrigin-RevId: 231448049
---
 tensor2tensor/envs/tic_tac_toe_env.py            | 16 ++++++++--------
 tensor2tensor/layers/common_image_attention.py   |  6 ++++--
 tensor2tensor/models/mtf_transformer.py          |  4 ++--
 tensor2tensor/models/mtf_transformer2.py         |  4 ++--
 tensor2tensor/models/research/autoencoders.py    |  4 ++--
 tensor2tensor/models/research/cycle_gan.py       |  4 ++--
 tensor2tensor/models/research/rl.py              | 16 ++++++++--------
 .../models/research/transformer_symshard.py      |  4 ++--
 tensor2tensor/models/video/epva_params.py        |  4 ++--
 tensor2tensor/models/video/next_frame_glow.py    |  4 ++--
 tensor2tensor/models/video/savp_params.py        |  8 ++++----
 tensor2tensor/models/video/sv2p_params.py        |  8 ++++----
 tensor2tensor/utils/data_reader_test.py          |  4 ++--
 13 files changed, 44 insertions(+), 42 deletions(-)

diff --git a/tensor2tensor/envs/tic_tac_toe_env.py b/tensor2tensor/envs/tic_tac_toe_env.py
index 4e0f18c00..43e603a44 100644
--- a/tensor2tensor/envs/tic_tac_toe_env.py
+++ b/tensor2tensor/envs/tic_tac_toe_env.py
@@ -188,8 +188,8 @@ def step(self, action):
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.modality = {
-        "inputs": modalities.IdentitySymbolModality,
-        "targets": modalities.IdentitySymbolModality,
+        "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
+        "targets": modalities.ModalityType.IDENTITY_SYMBOL,
     }
     p.vocab_size = {
         "inputs": 3,
@@ -212,12 +212,12 @@ def hparams(self, defaults, model_hparams):
     self._ttt_env.hparams(defaults, model_hparams)
     # Do these belong here?
     defaults.modality.update({
-        "input_action": modalities.SymbolModalityWeightsAll,
-        "input_reward": modalities.SymbolModalityWeightsAll,
-        "target_action": modalities.SymbolModalityWeightsAll,
-        "target_reward": modalities.SymbolModalityWeightsAll,
-        "target_policy": modalities.IdentityModality,
-        "target_value": modalities.IdentityModality,
+        "input_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "input_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "target_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "target_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "target_policy": modalities.ModalityType.IDENTITY,
+        "target_value": modalities.ModalityType.IDENTITY,
     })
     defaults.vocab_size.update({
         "input_action": self.num_actions,
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index ef5759b6a..9edcda37a 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -610,11 +610,13 @@ def prepare_image(inputs, hparams, name=None):
   channels = hparams.num_channels
 
   hidden_size = hparams.hidden_size
-  # TODO(trandustin): Check via modalities.IdentityModality and not its name.
+  # TODO(trandustin): Check via modalities.ModalityType.IDENTITY and not str.
   # The current implementation is to avoid circular imports, modalities ->
   # discretization -> common_image_attention -> modalities.
   if "targets" in hparams.modality:
-    target_modality_name = hparams.modality["targets"].__name__
+    target_modality_name = hparams.modality["targets"]
+    if not isinstance(target_modality_name, str):
+      target_modality_name = target_modality_name.__name__
   else:
     target_modality_name = None
   if target_modality_name == "IdentityModality":
diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 59de31adb..b70fd6257 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -857,8 +857,8 @@ def mtf_transformer_base():
   # Do not override these, as mtf_transformer does not support other options.
   hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
   hparams.modality = {
-      "inputs": modalities.IdentitySymbolModality,
-      "targets": modalities.IdentitySymbolModality,
+      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
+      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
   }
 
   # Parameters for computing the maximum decode length in beam search.
diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index 43dafed38..977214a49 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -405,8 +405,8 @@ def mtf_transformer2_base():
   hparams.add_hparam("mtf_mode", True)
   hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
   hparams.modality = {
-      "inputs": modalities.IdentitySymbolModality,
-      "targets": modalities.IdentitySymbolModality,
+      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
+      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
   }
   hparams.add_hparam("beam_size", 1)
   return hparams
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 3780bd6e6..4d7bac8b0 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1106,8 +1106,8 @@ def autoencoder_residual_text():
   hparams.max_hidden_size = 512
   hparams.bottleneck_noise = 0.0
   hparams.modality = {
-      "inputs": modalities.IdentitySymbolModality,
-      "targets": modalities.IdentitySymbolModality,
+      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
+      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
   }
   hparams.autoregressive_mode = "none"
   hparams.sample_width = 1
diff --git a/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index 49a2653d1..89693dc27 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -128,8 +128,8 @@ def cycle_gan_small():
   hparams = transformer_vae.transformer_ae_small()
   hparams.batch_size = 2048
   hparams.modality = {
-      "inputs": modalities.IdentitySymbolModality,
-      "targets": modalities.IdentitySymbolModality,
+      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
+      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
   }
   hparams.weight_decay = 3.0
   hparams.learning_rate = 0.05
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 2cacc1050..9f1982636 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -481,14 +481,14 @@ def num_actions(self):
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
     p.modality = {
-        "inputs": modalities.VideoModality,
-        "input_action": modalities.SymbolModalityWeightsAll,
-        "input_reward": modalities.SymbolModalityWeightsAll,
-        "targets": modalities.VideoModality,
-        "target_action": modalities.SymbolModalityWeightsAll,
-        "target_reward": modalities.SymbolModalityWeightsAll,
-        "target_policy": modalities.IdentityModality,
-        "target_value": modalities.IdentityModality,
+        "inputs": modalities.ModalityType.VIDEO,
+        "input_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "input_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "targets": modalities.ModalityType.VIDEO,
+        "target_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "target_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "target_policy": modalities.ModalityType.IDENTITY,
+        "target_value": modalities.ModalityType.IDENTITY,
     }
     p.vocab_size = {
         "inputs": 256,
diff --git a/tensor2tensor/models/research/transformer_symshard.py b/tensor2tensor/models/research/transformer_symshard.py
index 8fa58ab20..8e00e0ca2 100644
--- a/tensor2tensor/models/research/transformer_symshard.py
+++ b/tensor2tensor/models/research/transformer_symshard.py
@@ -368,8 +368,8 @@ def transformer_symshard_base():
   hparams.no_data_parallelism = True
   # bypass the symbol modality so that we can use model parallelism.
   hparams.modality = {
-      "inputs": modalities.IdentitySymbolModality,
-      "targets": modalities.IdentitySymbolModality,
+      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
+      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
   }
   hparams.add_hparam("filter_size", 1280)
   hparams.add_hparam("mix_fraction", 0.5)
diff --git a/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
index 9890f3008..0283f154e 100644
--- a/tensor2tensor/models/video/epva_params.py
+++ b/tensor2tensor/models/video/epva_params.py
@@ -30,8 +30,8 @@ def next_frame_epva():
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 4
   hparams.modality = {
-      "inputs": modalities.VideoModalityL2Raw,
-      "targets": modalities.VideoModalityL2Raw,
+      "inputs": modalities.ModalityType.VIDEO_L2_RAW,
+      "targets": modalities.ModalityType.VIDEO_L2_RAW,
   }
   hparams.learning_rate_schedule = "constant"
   hparams.learning_rate_constant = 1e-05
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index e18db1867..524fbd939 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -70,8 +70,8 @@ def next_frame_glow_hparams():
   # By default, don't pretrain and learn end-to-end
   hparams.add_hparam("pretrain_steps", -1)
   hparams.modality = {
-      "inputs": modalities.VideoModalityL1Raw,
-      "targets": modalities.VideoModalityL1Raw,
+      "inputs": modalities.ModalityType.VIDEO_L1_RAW,
+      "targets": modalities.ModalityType.VIDEO_L1_RAW,
   }
   hparams.init_batch_size = 256
   hparams.batch_size = 32
diff --git a/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index fd7a3d346..ff95e3902 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -37,8 +37,8 @@ def next_frame_savp():
   hparams.add_hparam("gan_vae_loss_multiplier", 0.01)
   hparams.add_hparam("gan_optimization", "joint")
   hparams.modality = {
-      "inputs": modalities.VideoModalityL1Raw,
-      "targets": modalities.VideoModalityL1Raw,
+      "inputs": modalities.ModalityType.VIDEO_L1_RAW,
+      "targets": modalities.ModalityType.VIDEO_L1_RAW,
   }
   hparams.latent_loss_multiplier_schedule = "linear"
   hparams.upsample_method = "bilinear_upsample_conv"
@@ -55,8 +55,8 @@ def next_frame_savp_l2():
   """SAVP with L2 reconstruction loss."""
   hparams = next_frame_savp()
   hparams.modality = {
-      "inputs": modalities.VideoModalityL2Raw,
-      "targets": modalities.VideoModalityL2Raw,
+      "inputs": modalities.ModalityType.VIDEO_L2_RAW,
+      "targets": modalities.ModalityType.VIDEO_L2_RAW,
   }
   return hparams
 
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 98dc9f658..f0d456c29 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -34,8 +34,8 @@ def next_frame_sv2p():
   hparams.video_num_target_frames = 3
   hparams.batch_size = 16
   hparams.modality = {
-      "inputs": modalities.VideoModalityL2Raw,
-      "targets": modalities.VideoModalityL2Raw,
+      "inputs": modalities.ModalityType.VIDEO_L2_RAW,
+      "targets": modalities.ModalityType.VIDEO_L2_RAW,
   }
   hparams.video_modality_loss_cutoff = 0.0
   hparams.scheduled_sampling_mode = "count"
@@ -92,8 +92,8 @@ def next_frame_sv2p_atari_softmax():
   """SV2P model for atari with softmax."""
   hparams = next_frame_sv2p_atari()
   hparams.modality = {
-      "inputs": modalities.VideoModality,
-      "targets": modalities.VideoModality,
+      "inputs": modalities.ModalityType.VIDEO,
+      "targets": modalities.ModalityType.VIDEO,
   }
   hparams.internal_loss = True
   return hparams
diff --git a/tensor2tensor/utils/data_reader_test.py b/tensor2tensor/utils/data_reader_test.py
index a82071d78..9487a8cbd 100644
--- a/tensor2tensor/utils/data_reader_test.py
+++ b/tensor2tensor/utils/data_reader_test.py
@@ -51,8 +51,8 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
   def hparams(self, defaults, model_hparams):
     hp = defaults
-    hp.modality = {"inputs": modalities.SymbolModality,
-                   "targets": modalities.SymbolModality}
+    hp.modality = {"inputs": modalities.ModalityType.SYMBOL,
+                   "targets": modalities.ModalityType.SYMBOL}
     hp.vocab_size = {"inputs": 30,
                      "targets": 30}
 

From fbd58ca430eed4ec303bd7b73c3398321c65c254 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 29 Jan 2019 12:28:38 -0800
Subject: [PATCH 1592/2720] Add support for frame interpolation in the Glow
 latent space as a wrapper around t2t-decode.

PiperOrigin-RevId: 231451719
---
 tensor2tensor/models/research/glow.py         |  11 +-
 tensor2tensor/models/research/glow_ops.py     |  34 +++
 tensor2tensor/models/video/next_frame_glow.py |   2 +-
 tensor2tensor/models/video/nfg_interpolate.py | 221 ++++++++++++++++++
 4 files changed, 257 insertions(+), 11 deletions(-)
 create mode 100644 tensor2tensor/models/video/nfg_interpolate.py

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 7e95579f9..072a04aaa 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -98,15 +98,6 @@ def temperature(self):
       return self.hparams.temperature
     return 1.0
 
-  def scale(self, x):
-    """Scale x from -0.5 - 0.5 to 0 - 255."""
-    x = tf.where(tf.is_nan(x), tf.ones_like(x), x)
-    x = tf.where(tf.is_inf(x), tf.ones_like(x), x)
-    x = tf.clip_by_value(x, -0.5, 0.5)
-    x += 0.5
-    x = x * 2**self.hparams.n_bits_x
-    return tf.cast(tf.clip_by_value(x, 0, 255), dtype=tf.uint8)
-
   @property
   def is_training(self):
     return self.hparams.mode == tf.estimator.ModeKeys.TRAIN
@@ -126,7 +117,7 @@ def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
           "codec", self.z_sample, self.hparams, eps=None, reverse=True,
           temperature=self.temperature)
 
-    return self.scale(predictions)
+    return glow_ops.postprocess(predictions, self.hparams.n_bits_x)
 
   def create_init_batch(self, features):
     """Returns a batch of size "hparams.init_batch_size" for initialization.
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index aa9649fb1..c5b2d0a11 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -31,6 +31,40 @@
 add_arg_scope = tf.contrib.framework.add_arg_scope
 
 
+def linear_interpolate(tensor1, tensor2, coeffs):
+  """Linearly interpolate between two tensors at coeff.
+
+  Args:
+    tensor1: 3-D Tensor, NHWC
+    tensor2: 3-D Tensor, NHWC
+    coeffs: list of floats.
+  Returns:
+    interp_latents: list of interpolated 4-D Tensors, shape=(1HWC)
+  """
+  interp_tensors = []
+  for coeff in coeffs:
+    interp_tensor = tensor1 + coeff * (tensor2 - tensor1)
+    interp_tensors.append(interp_tensor)
+  return tf.concat(interp_tensors, axis=0)
+
+
+def postprocess(x, n_bits_x=8):
+  """Converts x from [-0.5, 0.5], to [0, 255].
+
+  Args:
+    x: 3-D or 4-D Tensor normalized between [-0.5, 0.5]
+    n_bits_x: Number of bits representing each pixel of the output.
+              Defaults to 8, to default to 256 possible values.
+  Returns:
+    x: 3-D or 4-D Tensor representing images or videos.
+  """
+  x = tf.where(tf.is_finite(x), x, tf.ones_like(x))
+  x = tf.clip_by_value(x, -0.5, 0.5)
+  x += 0.5
+  x = x * 2**n_bits_x
+  return tf.cast(tf.clip_by_value(x, 0, 255), dtype=tf.uint8)
+
+
 class TemperedNormal(tfp.distributions.Normal):
   """Normal distribution with temperature T."""
 
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index 524fbd939..9cbcea1a1 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -212,7 +212,7 @@ def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
     if self.hparams.gen_mode == "unconditional":
       predicted_video = tf.tile(
           predicted_video, [1, self.hparams.video_num_target_frames, 1, 1, 1])
-    predicted_video = self.scale(predicted_video)
+    predicted_video = glow_ops.postprocess(predicted_video)
 
     # Output of a single decode / sample.
     output_features = {}
diff --git a/tensor2tensor/models/video/nfg_interpolate.py b/tensor2tensor/models/video/nfg_interpolate.py
new file mode 100644
index 000000000..281536415
--- /dev/null
+++ b/tensor2tensor/models/video/nfg_interpolate.py
@@ -0,0 +1,221 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Utilities for linear interpolation over the next_frame_glow latent space."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import os
+from absl import flags
+import numpy as np
+from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
+from tensor2tensor.data_generators import image_utils
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
+from tensor2tensor.models.research import glow_ops
+from tensor2tensor.utils import decoding
+from tensor2tensor.utils import trainer_lib
+import tensorflow as tf
+
+# Flags placeholders.
+flags.DEFINE_string("checkpoint_path", None,
+                    "Path to the model checkpoint. Overrides output_dir.")
+flags.DEFINE_bool("keep_timestamp", False,
+                  "Set the mtime of the decoded file to the "
+                  "checkpoint_path+'.index' mtime.")
+flags.DEFINE_bool("decode_interactive", False,
+                  "Interactive local inference mode.")
+flags.DEFINE_integer("decode_shards", 1, "Number of decoding replicas.")
+flags.DEFINE_string("score_file", "", "File to score. Each line in the file "
+                    "must be in the format input \t target.")
+flags.DEFINE_bool("decode_in_memory", False, "Decode in memory.")
+# Interpolate between z1 and z2 for alpha = np.linspace(0.0, 1.0, num_interp)
+flags.DEFINE_integer("num_interp", 11, "Number of interpolations")
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+
+arg_scope = tf.contrib.framework.arg_scope
+
+
+def preprocess_frame(frame):
+  """Preprocess frame.
+
+  1. Converts [0, 255] to [-0.5, 0.5]
+  2. Adds uniform noise.
+
+  Args:
+    frame: 3-D Tensor representing pixels.
+  Returns:
+    frame: 3-D Tensor with values in between [-0.5, 0.5]
+  """
+  # Normalize from [0.0, 1.0] -> [-0.5, 0.5]
+  frame = common_layers.convert_rgb_to_real(frame)
+  frame = frame - 0.5
+  frame, _ = glow_ops.uniform_binning_correction(frame)
+  return frame
+
+
+def frame_to_latents(frame, hparams):
+  """Encode frames to latents."""
+  # Preprocess
+  frame = preprocess_frame(frame)
+
+  # Encode [X_t] to [z^1_t, z^2_t .. z^l_t]
+  glow_vals = glow_ops.encoder_decoder(
+      "codec", frame, hparams, eps=None, reverse=False)
+  z_top, _, level_eps, _, _ = glow_vals
+  return z_top, level_eps
+
+
+def latents_to_frames(z_top_interp, level_eps_interp, hparams):
+  """Decodes latents to frames."""
+  # Decode [z^1_t, z^2_t .. z^l_t] to [X_t]
+  images, _, _, _ = glow_ops.encoder_decoder(
+      "codec", z_top_interp, hparams, eps=level_eps_interp, reverse=True)
+  images = glow_ops.postprocess(images)
+  return images
+
+
+def interpolate(features, hparams, num_interp):
+  """Interpolate between the first input frame and last target frame.
+
+  Args:
+    features: dict of tensors
+    hparams: tf.contrib.training.HParams.
+    num_interp: integer.
+  Returns:
+    images: 4-D Tensor, shape=(num_interp, H, W, C)
+  """
+  inputs, targets = features["inputs"], features["targets"]
+  inputs = tf.unstack(inputs, axis=1)
+  targets = tf.unstack(targets, axis=1)
+  coeffs = np.linspace(0.0, 1.0, num_interp)
+
+  # (X_1, X_t) -> (z_1, z_t)
+  first_frame, last_frame = inputs[0], targets[-1]
+  first_top_z, first_level_eps = frame_to_latents(first_frame, hparams)
+  last_top_z, last_level_eps = frame_to_latents(last_frame, hparams)
+
+  # Interpolate top
+  z_top_interp = glow_ops.linear_interpolate(first_top_z, last_top_z, coeffs)
+
+  # Interpolate level.
+  level_eps_interp = []
+  for level in range(hparams.n_levels - 1):
+    level_eps_interp.append(glow_ops.linear_interpolate(
+        first_level_eps[level], last_level_eps[level], coeffs))
+  return latents_to_frames(z_top_interp, level_eps_interp, hparams)
+
+
+def interpolations_to_summary(sample_ind, interpolations, hparams,
+                              decode_hparams):
+  """Converts interpolated frames into tf summaries.
+
+  The summaries consists of:
+    1. Image summary corresponding to the first frame.
+    2. Image summary corresponding to the last frame.
+    3. The interpolated frames as a gif summary.
+
+  Args:
+    sample_ind: int
+    interpolations: Numpy array, shape=(num_interp, 64, 64, 3)
+    hparams: tf.contrib.training.HParams, train hparams
+    decode_hparams: tf.contrib.training.HParams, decode hparams
+  Returns:
+    summaries: list of tf Summary Values.
+  """
+  parent_tag = "sample_%d" % sample_ind
+  frame_shape = hparams.problem.frame_shape
+  interp_shape = [hparams.batch_size, FLAGS.num_interp] + frame_shape
+  interpolations = np.reshape(interpolations, interp_shape)
+  summaries, _ = common_video.py_gif_summary(
+      parent_tag, interpolations, return_summary_value=True,
+      max_outputs=decode_hparams.max_display_outputs,
+      fps=decode_hparams.frames_per_second)
+
+  first_frame, last_frame = interpolations[0, 0], interpolations[0, -1]
+  first_frame_summ = image_utils.image_to_tf_summary_value(
+      first_frame, "%s/first" % parent_tag)
+  last_frame_summ = image_utils.image_to_tf_summary_value(
+      last_frame, "%s/last" % parent_tag)
+  summaries.append(first_frame_summ)
+  summaries.append(last_frame_summ)
+  return summaries
+
+
+def main(_):
+  decode_hparams = decoding.decode_hparams(FLAGS.decode_hparams)
+  trainer_lib.set_random_seed(FLAGS.random_seed)
+  if FLAGS.output_dir is None:
+    raise ValueError("Expected output_dir to be set to a valid path.")
+
+  hparams = trainer_lib.create_hparams(
+      FLAGS.hparams_set, FLAGS.hparams, data_dir=FLAGS.data_dir,
+      problem_name=FLAGS.problem)
+  if hparams.batch_size != 1:
+    raise ValueError("Set batch-size to be equal to 1")
+
+  # prepare dataset using Predict mode.
+  dataset_split = "test" if FLAGS.eval_use_test_set else None
+  dataset = hparams.problem.dataset(
+      tf.estimator.ModeKeys.PREDICT, shuffle_files=False, hparams=hparams,
+      data_dir=FLAGS.data_dir, dataset_split=dataset_split)
+  dataset = dataset.batch(hparams.batch_size)
+  dataset = dataset.make_one_shot_iterator().get_next()
+
+  # Obtain frame interpolations.
+  ops = [glow_ops.get_variable_ddi, glow_ops.actnorm, glow_ops.get_dropout]
+  var_scope = tf.variable_scope("next_frame_glow/body", reuse=tf.AUTO_REUSE)
+  with arg_scope(ops, init=False), var_scope:
+    interpolations = interpolate(dataset, hparams, FLAGS.num_interp)
+
+  var_list = tf.global_variables()
+  saver = tf.train.Saver(var_list)
+
+  # Get latest checkpoints from model_dir.
+  ckpt_path = tf.train.latest_checkpoint(FLAGS.output_dir)
+  child_dir = decode_hparams.summaries_log_dir
+  if dataset_split is not None:
+    child_dir += "_{}".format(dataset_split)
+  final_dir = os.path.join(FLAGS.output_dir, child_dir)
+  summary_writer = tf.summary.FileWriter(final_dir)
+  global_step = decoding.latest_checkpoint_step(FLAGS.output_dir)
+
+  sample_ind = 0
+
+  num_samples = decode_hparams.num_samples
+  all_summaries = []
+
+  with tf.train.MonitoredTrainingSession() as sess:
+    saver.restore(sess, ckpt_path)
+
+    while not sess.should_stop() and sample_ind < num_samples:
+      interp_np = sess.run(interpolations)
+
+      interp_summ = interpolations_to_summary(sample_ind, interp_np, hparams,
+                                              decode_hparams)
+      all_summaries.extend(interp_summ)
+
+      sample_ind += 1
+    all_summaries = tf.Summary(value=list(all_summaries))
+    summary_writer.add_summary(all_summaries, global_step)
+
+
+if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
+  tf.app.run()

From 374a0af7d0f63f997dcb0d8e88973460599006ac Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 29 Jan 2019 12:34:28 -0800
Subject: [PATCH 1593/2720] Adding problem.data_filepaths which internally
 calls {training/dev/test}_filpaths internally.

PiperOrigin-RevId: 231452698
---
 tensor2tensor/data_generators/problem.py      | 10 +++++++
 tensor2tensor/data_generators/problem_test.py | 26 +++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 062d26926..e3a38f9fb 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -435,6 +435,16 @@ def test_filepaths(self, data_dir, num_shards, shuffled):
     return generator_utils.test_data_filenames(file_basename, data_dir,
                                                num_shards)
 
+  def data_filepaths(self, split, output_dir, num_shards, shuffled):
+    if split == DatasetSplit.TRAIN:
+      return self.training_filepaths(output_dir, num_shards, shuffled)
+    elif split == DatasetSplit.EVAL:
+      return self.dev_filepaths(output_dir, num_shards, shuffled)
+    elif split == DatasetSplit.TEST:
+      return self.test_filepaths(output_dir, num_shards, shuffled)
+    else:
+      raise ValueError("Unknown value for split: %s" % split)
+
   def filepattern(self, data_dir, mode, shard=None):
     """Get filepattern for data files for mode.
 
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index ff0084aac..c7b5781be 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -138,5 +138,31 @@ def hparams(self, defaults, model_hparams):
                           modalities.SymbolModality)
     self.assertLen(p_hparams.modality, 1)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testDataFilenames(self):
+    problem = algorithmic.TinyAlgo()
+
+    num_shards = 10
+    shuffled = False
+    data_dir = "/tmp"
+
+    # Test training_filepaths and data_filepaths give the same list on
+    # appropriate arguments.
+    self.assertAllEqual(
+        problem.training_filepaths(data_dir, num_shards, shuffled),
+        problem.data_filepaths(problem_module.DatasetSplit.TRAIN, data_dir,
+                               num_shards, shuffled))
+
+    self.assertAllEqual(
+        problem.dev_filepaths(data_dir, num_shards, shuffled),
+        problem.data_filepaths(problem_module.DatasetSplit.EVAL, data_dir,
+                               num_shards, shuffled))
+
+    self.assertAllEqual(
+        problem.test_filepaths(data_dir, num_shards, shuffled),
+        problem.data_filepaths(problem_module.DatasetSplit.TEST, data_dir,
+                               num_shards, shuffled))
+
+
 if __name__ == "__main__":
   tf.test.main()

From bc850acdcedf3ce4cdba9c02bc0eb6a9a7219033 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 29 Jan 2019 13:00:34 -0800
Subject: [PATCH 1594/2720] Fix glow_ops.test_encoder_decoder.

PiperOrigin-RevId: 231457229
---
 .../models/research/glow_ops_test.py          | 30 +++++++++++--------
 1 file changed, 17 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 19c205157..00f40d3e9 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -206,27 +206,31 @@ def test_encoder_decoder(self):
     with tf.Graph().as_default():
       hparams = glow.glow_hparams()
       hparams.n_levels = 3
-      hparams.depth = 2
-
-      x = tf.random_uniform(shape=(16, 64, 64, 4), seed=0)
-      x_inv, _, eps, z_levels, _ = glow_ops.encoder_decoder(
-          "encoder_decoder", x, hparams, reverse=False)
+      hparams.depth = 6
+      rng = np.random.RandomState(0)
+      x_np = rng.rand(1, 64, 64, 4)
+      x_t = tf.convert_to_tensor(x_np, dtype=tf.float32)
+      init_ops = [glow_ops.get_variable_ddi, glow_ops.actnorm]
+      with arg_scope(init_ops, init=True):
+        x_inv, _, eps, z_levels, _ = glow_ops.encoder_decoder(
+            "encoder_decoder", x_t, hparams, reverse=False)
       x_inv_inv, _, z_inv_levels, _ = glow_ops.encoder_decoder(
           "encoder_decoder", x_inv, hparams, eps=eps, reverse=True)
 
       with tf.Session() as session:
         session.run(tf.global_variables_initializer())
-        diff, x_inv_np, z_levels_np, z_inv_levels_np = session.run(
-            [x - x_inv_inv, x_inv, z_levels, z_inv_levels])
-
+        x_inv_np = session.run(x_inv)
+        z_levels_np, z_inv_levels_np, x_inv_inv_np = session.run(
+            [z_levels, z_inv_levels, x_inv_inv])
+        diff = x_inv_inv_np - x_np
         self.assertLen(z_levels_np, 2)
         self.assertLen(z_inv_levels_np, 2)
         # (h_i, w_i, c_i) = (h_{i-1}/f, w_{i-1}/f, c_{i-1}*(2f)/2) where (f=2)
-        self.assertEqual(z_levels_np[0].shape, (16, 32, 32, 8))
-        self.assertEqual(z_levels_np[1].shape, (16, 16, 16, 16))
-        self.assertEqual(z_inv_levels_np[0].shape, (16, 32, 32, 8))
-        self.assertEqual(z_inv_levels_np[1].shape, (16, 16, 16, 16))
-        self.assertTrue(x_inv_np.shape, (16, 8, 8, 64))
+        self.assertEqual(z_levels_np[0].shape, (1, 32, 32, 8))
+        self.assertEqual(z_levels_np[1].shape, (1, 16, 16, 16))
+        self.assertEqual(z_inv_levels_np[0].shape, (1, 32, 32, 8))
+        self.assertEqual(z_inv_levels_np[1].shape, (1, 16, 16, 16))
+        self.assertTrue(x_inv_np.shape, (1, 8, 8, 64))
         self.assertTrue(np.allclose(diff, 0.0, atol=1e-2))
 
   def test_encoder_decoder_practical_usage(self):

From e02d794244080a9c1fed7cf4692a568e90bc2c42 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 29 Jan 2019 22:32:15 +0100
Subject: [PATCH 1595/2720] Simulation videos (#1419)

* Dump videos from simulated env in evaluator

* Enable model evaluation in evaluator
---
 tensor2tensor/rl/evaluator.py           | 194 +++++++++++++++++++-----
 tensor2tensor/rl/evaluator_test.py      |   4 +-
 tensor2tensor/rl/rl_utils.py            | 174 +++++++++++++++++++--
 tensor2tensor/rl/trainer_model_based.py | 133 +---------------
 4 files changed, 321 insertions(+), 184 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index b6e973b42..03bd87fe5 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -60,6 +60,18 @@
 flags.DEFINE_enum(
     "agent", "policy", ["random", "policy", "planner"], "Agent type to use."
 )
+# Evaluator doesn't report metrics for agent on the simulated env because we
+# don't collect rollouts there. It's just for generating videos.
+# TODO(koz4k): Enable reporting metrics from simulated env by refactoring
+# T2TEnv to a wrapper storing rollouts and providing Problem interface for any
+# batch env.
+flags.DEFINE_enum(
+    "mode", "agent_real", ["agent_real", "agent_simulated", "model"],
+    "Evaluation mode; report agent's score on real or simulated env, or model's"
+    " reward accuracy."
+)
+# TODO(koz4k): Switch to out-of-graph evaluation everywhere and remove this
+# flag.
 flags.DEFINE_bool(
     "eval_with_learner", False,
     "Whether to use the PolicyLearner.evaluate function instead of an "
@@ -73,11 +85,15 @@
     "log_every_steps", 20, "Log every how many environment steps."
 )
 flags.DEFINE_string(
-    "debug_video_path", "", "Path to save the planner debug video at."
+    "debug_video_path", "", "Path to save the debug video at."
 )
 flags.DEFINE_integer(
     "num_debug_videos", 1, "Number of debug videos to generate."
 )
+flags.DEFINE_integer(
+    "random_starts_step_limit", 10000,
+    "Number of frames to choose from for random starts of the simulated env."
+)
 
 # Unused flags needed to pass for multi-run infrastructure.
 flags.DEFINE_bool("autotune", False, "Unused here.")
@@ -251,45 +267,132 @@ def make_agent(
   }[agent_type]()
 
 
+def collect_frames_for_random_starts(
+    storage_env, stacked_env, agent, frame_stack_size, random_starts_step_limit,
+    log_every_steps=None
+):
+  """Collects frames from real env for random starts of simulated env."""
+  storage_env.start_new_epoch(0)
+  tf.logging.info(
+      "Collecting %d frames for random starts.", random_starts_step_limit
+  )
+  rl_utils.run_rollouts(
+      stacked_env, agent, stacked_env.reset(),
+      step_limit=random_starts_step_limit,
+      many_rollouts_from_each_env=True,
+      log_every_steps=log_every_steps,
+  )
+  # Save unfinished rollouts to history.
+  stacked_env.reset()
+
+
+def make_agent_from_hparams(
+    agent_type, base_env, stacked_env, loop_hparams, policy_hparams,
+    planner_hparams, model_dir, policy_dir, sampling_temp, video_writers=()
+):
+  """Creates an Agent from hparams."""
+  sim_env_kwargs = rl.make_simulated_env_kwargs(
+      base_env, loop_hparams, batch_size=planner_hparams.batch_size,
+      model_dir=model_dir
+  )
+  planner_kwargs = planner_hparams.values()
+  planner_kwargs.pop("batch_size")
+  planner_kwargs.pop("rollout_agent_type")
+  planner_kwargs.pop("env_type")
+  return make_agent(
+      agent_type, stacked_env, policy_hparams, policy_dir, sampling_temp,
+      sim_env_kwargs, loop_hparams.frame_stack_size,
+      planner_hparams.rollout_agent_type,
+      inner_batch_size=planner_hparams.batch_size,
+      env_type=planner_hparams.env_type,
+      video_writers=video_writers, **planner_kwargs
+  )
+
+
 def make_eval_fn_with_agent(
-    agent_type, planner_hparams, model_dir, log_every_steps=None,
-    video_writers=()
+    agent_type, eval_mode, planner_hparams, model_dir, log_every_steps=None,
+    video_writers=(), random_starts_step_limit=None
 ):
   """Returns an out-of-graph eval_fn using the Agent API."""
   def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
     """Eval function."""
     base_env = env
     env = rl_utils.BatchStackWrapper(env, loop_hparams.frame_stack_size)
-    sim_env_kwargs = rl.make_simulated_env_kwargs(
-        base_env, loop_hparams, batch_size=planner_hparams.batch_size,
-        model_dir=model_dir
-    )
-    planner_kwargs = planner_hparams.values()
-    planner_kwargs.pop("batch_size")
-    planner_kwargs.pop("rollout_agent_type")
-    planner_kwargs.pop("env_type")
-    agent = make_agent(
-        agent_type, env, policy_hparams, policy_dir, sampling_temp,
-        sim_env_kwargs, loop_hparams.frame_stack_size,
-        planner_hparams.rollout_agent_type,
-        inner_batch_size=planner_hparams.batch_size,
-        env_type=planner_hparams.env_type,
-        video_writers=video_writers, **planner_kwargs
+    agent = make_agent_from_hparams(
+        agent_type, base_env, env, loop_hparams, policy_hparams,
+        planner_hparams, model_dir, policy_dir, sampling_temp, video_writers
     )
+
+    if eval_mode == "agent_simulated":
+      real_env = base_env.new_like(batch_size=1)
+      stacked_env = rl_utils.BatchStackWrapper(
+          real_env, loop_hparams.frame_stack_size
+      )
+      collect_frames_for_random_starts(
+          real_env, stacked_env, agent, loop_hparams.frame_stack_size,
+          random_starts_step_limit, log_every_steps
+      )
+      initial_frame_chooser = rl_utils.make_initial_frame_chooser(
+          real_env, loop_hparams.frame_stack_size,
+          simulation_random_starts=True,
+          simulation_flip_first_random_for_beginning=False,
+          split=None,
+      )
+      env_fn = rl.make_simulated_env_fn_from_hparams(
+          real_env, loop_hparams, batch_size=loop_hparams.eval_batch_size,
+          initial_frame_chooser=initial_frame_chooser, model_dir=model_dir
+      )
+      sim_env = env_fn(in_graph=False)
+      env = rl_utils.BatchStackWrapper(sim_env, loop_hparams.frame_stack_size)
+
     kwargs = {}
     if not agent.records_own_videos:
       kwargs["video_writers"] = video_writers
+    step_limit = base_env.rl_env_max_episode_steps
+    if step_limit == -1:
+      step_limit = None
     rl_utils.run_rollouts(
-        env, agent, env.reset(), log_every_steps=log_every_steps, **kwargs
+        env, agent, env.reset(), log_every_steps=log_every_steps,
+        step_limit=step_limit, **kwargs
     )
-    assert len(base_env.current_epoch_rollouts()) == env.batch_size
+    if eval_mode == "agent_real":
+      assert len(base_env.current_epoch_rollouts()) == env.batch_size
   return eval_fn
 
 
+def evaluate_world_model(
+    agent_type, loop_hparams, planner_hparams, model_dir, policy_dir,
+    random_starts_step_limit, debug_video_path, log_every_steps
+):
+  """Evaluates the world model."""
+  if debug_video_path:
+    debug_video_path = os.path.join(debug_video_path, "0.avi")
+
+  storage_env = rl_utils.setup_env(loop_hparams, batch_size=1, max_num_noops=0)
+  stacked_env = rl_utils.BatchStackWrapper(
+      storage_env, loop_hparams.frame_stack_size
+  )
+  policy_hparams = trainer_lib.create_hparams(loop_hparams.base_algo_params)
+  agent = make_agent_from_hparams(
+      agent_type, storage_env, stacked_env, loop_hparams, policy_hparams,
+      planner_hparams, model_dir, policy_dir,
+      # TODO(koz4k): Loop over eval_sampling_temps?
+      sampling_temp=loop_hparams.eval_sampling_temps[0],
+  )
+  collect_frames_for_random_starts(
+      storage_env, stacked_env, agent, loop_hparams.frame_stack_size,
+      random_starts_step_limit, log_every_steps
+  )
+  return rl_utils.evaluate_world_model(
+      storage_env, loop_hparams, model_dir, debug_video_path, split=None
+  )
+
+
 def evaluate(
     loop_hparams, planner_hparams, policy_dir, model_dir, eval_metrics_dir,
-    agent_type, eval_with_learner, log_every_steps, debug_video_path,
-    num_debug_videos=1, report_fn=None, report_metric=None
+    agent_type, eval_mode, eval_with_learner, log_every_steps, debug_video_path,
+    num_debug_videos=1, random_starts_step_limit=None,
+    report_fn=None, report_metric=None
 ):
   """Evaluate."""
   if eval_with_learner:
@@ -301,24 +404,32 @@ def evaluate(
   eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
   video_writers = ()
   kwargs = {}
-  if not eval_with_learner:
-    if debug_video_path:
-      tf.gfile.MakeDirs(debug_video_path)
-      video_writers = [
-          common_video.WholeVideoWriter(  # pylint: disable=g-complex-comprehension
-              fps=10,
-              output_path=os.path.join(debug_video_path, "{}.avi".format(i)),
-              file_format="avi",
-          )
-          for i in range(num_debug_videos)
-      ]
-    kwargs["eval_fn"] = make_eval_fn_with_agent(
-        agent_type, planner_hparams, model_dir, log_every_steps=log_every_steps,
-        video_writers=video_writers
+  if eval_mode in ["agent_real", "agent_simulated"]:
+    if not eval_with_learner:
+      if debug_video_path:
+        tf.gfile.MakeDirs(debug_video_path)
+        video_writers = [
+            common_video.WholeVideoWriter(  # pylint: disable=g-complex-comprehension
+                fps=10,
+                output_path=os.path.join(debug_video_path, "{}.avi".format(i)),
+                file_format="avi",
+            )
+            for i in range(num_debug_videos)
+        ]
+      kwargs["eval_fn"] = make_eval_fn_with_agent(
+          agent_type, eval_mode, planner_hparams, model_dir,
+          log_every_steps=log_every_steps,
+          video_writers=video_writers,
+          random_starts_step_limit=random_starts_step_limit
+      )
+    eval_metrics = rl_utils.evaluate_all_configs(
+        loop_hparams, policy_dir, **kwargs
+    )
+  else:
+    eval_metrics = evaluate_world_model(
+        agent_type, loop_hparams, planner_hparams, model_dir, policy_dir,
+        random_starts_step_limit, debug_video_path, log_every_steps
     )
-  eval_metrics = rl_utils.evaluate_all_configs(
-      loop_hparams, policy_dir, **kwargs
-  )
   rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, 0)
 
   for video_writer in video_writers:
@@ -387,10 +498,11 @@ def main(_):
       tf.gfile.MkDir(eval_metrics_dir)
   evaluate(
       loop_hparams, planner_hparams, policy_dir, model_dir,
-      eval_metrics_dir, FLAGS.agent, FLAGS.eval_with_learner,
+      eval_metrics_dir, FLAGS.agent, FLAGS.mode, FLAGS.eval_with_learner,
       FLAGS.log_every_steps if FLAGS.log_every_steps > 0 else None,
       debug_video_path=FLAGS.debug_video_path,
-      num_debug_videos=FLAGS.num_debug_videos
+      num_debug_videos=FLAGS.num_debug_videos,
+      random_starts_step_limit=FLAGS.random_starts_step_limit,
   )
 
 
diff --git a/tensor2tensor/rl/evaluator_test.py b/tensor2tensor/rl/evaluator_test.py
index 69c203325..a1f8a2b62 100644
--- a/tensor2tensor/rl/evaluator_test.py
+++ b/tensor2tensor/rl/evaluator_test.py
@@ -32,8 +32,8 @@ def test_evaluate_pong_random_agent(self):
     temp_dir = tf.test.get_temp_dir()
     evaluator.evaluate(
         loop_hparams, planner_hparams, temp_dir, temp_dir, temp_dir,
-        agent_type="random", eval_with_learner=False, log_every_steps=None,
-        debug_video_path=""
+        agent_type="random", eval_mode="agent_real", eval_with_learner=False,
+        log_every_steps=None, debug_video_path=""
     )
 
 
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 3686c250b..6e225883b 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -29,6 +29,7 @@
 
 from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl.dopamine_connector import DQNLearner
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image
@@ -116,6 +117,138 @@ def evaluate_all_configs(
   return metrics
 
 
+def evaluate_world_model(
+    real_env, hparams, world_model_dir, debug_video_path,
+    split=tf.estimator.ModeKeys.EVAL,
+):
+  """Evaluate the world model (reward accuracy)."""
+  frame_stack_size = hparams.frame_stack_size
+  rollout_subsequences = []
+  def initial_frame_chooser(batch_size):
+    assert batch_size == len(rollout_subsequences)
+    return np.stack([
+        [frame.observation.decode() for frame in subsequence[:frame_stack_size]]
+        for subsequence in rollout_subsequences
+    ])
+
+  env_fn = rl.make_simulated_env_fn_from_hparams(
+      real_env, hparams, batch_size=hparams.wm_eval_batch_size,
+      initial_frame_chooser=initial_frame_chooser, model_dir=world_model_dir
+  )
+  sim_env = env_fn(in_graph=False)
+  subsequence_length = int(
+      max(hparams.wm_eval_rollout_ratios) * hparams.simulated_rollout_length
+  )
+  rollouts = real_env.current_epoch_rollouts(
+      split=split,
+      minimal_rollout_frames=(subsequence_length + frame_stack_size)
+  )
+
+  video_writer = common_video.WholeVideoWriter(
+      fps=10, output_path=debug_video_path, file_format="avi"
+  )
+
+  reward_accuracies_by_length = {
+      int(ratio * hparams.simulated_rollout_length): []
+      for ratio in hparams.wm_eval_rollout_ratios
+  }
+  for _ in range(hparams.wm_eval_num_batches):
+    rollout_subsequences[:] = random_rollout_subsequences(
+        rollouts, hparams.wm_eval_batch_size,
+        subsequence_length + frame_stack_size
+    )
+
+    eval_subsequences = [
+        subsequence[(frame_stack_size - 1):]
+        for subsequence in rollout_subsequences
+    ]
+
+    # Check that the initial observation is the same in the real and simulated
+    # rollout.
+    sim_init_obs = sim_env.reset()
+    def decode_real_obs(index):
+      return np.stack([
+          subsequence[index].observation.decode()
+          for subsequence in eval_subsequences  # pylint: disable=cell-var-from-loop
+      ])
+    real_init_obs = decode_real_obs(0)
+    assert np.all(sim_init_obs == real_init_obs)
+
+    debug_frame_batches = []
+    def append_debug_frame_batch(sim_obs, real_obs, sim_cum_rews,
+                                 real_cum_rews, sim_rews, real_rews):
+      """Add a debug frame."""
+      rews = [[sim_cum_rews, sim_rews], [real_cum_rews, real_rews]]
+      headers = []
+      for j in range(len(sim_obs)):
+        local_nps = []
+        for i in range(2):
+          img = PIL_Image().new("RGB", (sim_obs.shape[-2], 11),)
+          draw = PIL_ImageDraw().Draw(img)
+          draw.text((0, 0), "c:{:3}, r:{:3}".format(int(rews[i][0][j]),
+                                                    int(rews[i][1][j])),
+                    fill=(255, 0, 0))
+          local_nps.append(np.asarray(img))
+        local_nps.append(np.zeros_like(local_nps[0]))
+        headers.append(np.concatenate(local_nps, axis=1))
+      errs = absolute_hinge_difference(sim_obs, real_obs)
+      headers = np.stack(headers)
+      debug_frame_batches.append(  # pylint: disable=cell-var-from-loop
+          np.concatenate([headers,
+                          np.concatenate([sim_obs, real_obs, errs], axis=2)],
+                         axis=1)
+      )
+    append_debug_frame_batch(sim_init_obs, real_init_obs,
+                             np.zeros(hparams.wm_eval_batch_size),
+                             np.zeros(hparams.wm_eval_batch_size),
+                             np.zeros(hparams.wm_eval_batch_size),
+                             np.zeros(hparams.wm_eval_batch_size))
+
+    (sim_cum_rewards, real_cum_rewards) = (
+        np.zeros(hparams.wm_eval_batch_size) for _ in range(2)
+    )
+    for i in range(subsequence_length):
+      actions = [subsequence[i].action for subsequence in eval_subsequences]
+      (sim_obs, sim_rewards, _) = sim_env.step(actions)
+      sim_cum_rewards += sim_rewards
+
+      real_rewards = np.array([
+          subsequence[i + 1].reward for subsequence in eval_subsequences
+      ])
+      real_cum_rewards += real_rewards
+      for (length, reward_accuracies) in six.iteritems(
+          reward_accuracies_by_length
+      ):
+        if i + 1 == length:
+          reward_accuracies.append(
+              np.sum(sim_cum_rewards == real_cum_rewards) /
+              len(real_cum_rewards)
+          )
+
+      real_obs = decode_real_obs(i + 1)
+      append_debug_frame_batch(sim_obs, real_obs, sim_cum_rewards,
+                               real_cum_rewards, sim_rewards, real_rewards)
+
+    for debug_frames in np.stack(debug_frame_batches, axis=1):
+      debug_frame = None
+      for debug_frame in debug_frames:
+        video_writer.write(debug_frame)
+
+      if debug_frame is not None:
+        # Append two black frames for aesthetics.
+        for _ in range(2):
+          video_writer.write(np.zeros_like(debug_frame))
+
+  video_writer.finish_to_disk()
+
+  return {
+      "reward_accuracy/at_{}".format(length): np.mean(reward_accuracies)
+      for (length, reward_accuracies) in six.iteritems(
+          reward_accuracies_by_length
+      )
+  }
+
+
 def summarize_metrics(eval_metrics_writer, metrics, epoch):
   """Write metrics to summary."""
   for (name, value) in six.iteritems(metrics):
@@ -201,13 +334,27 @@ def choose_subsequence():
   return [choose_subsequence() for _ in range(num_subsequences)]
 
 
-def make_initial_frame_chooser(real_env, frame_stack_size,
-                               simulation_random_starts,
-                               simulation_flip_first_random_for_beginning):
-  """Make frame chooser."""
+def make_initial_frame_chooser(
+    real_env, frame_stack_size, simulation_random_starts,
+    simulation_flip_first_random_for_beginning,
+    split=tf.estimator.ModeKeys.TRAIN,
+):
+  """Make frame chooser.
+
+  Args:
+    real_env: T2TEnv to take initial frames from.
+    frame_stack_size (int): Number of consecutive frames to extract.
+    simulation_random_starts (bool): Whether to choose frames at random.
+    simulation_flip_first_random_for_beginning (bool): Whether to flip the first
+      frame stack in every batch for the frames at the beginning.
+    split (tf.estimator.ModeKeys or None): Data split to take the frames from,
+      None means use all frames.
+
+  Returns:
+    Function batch_size -> initial_frames.
+  """
   initial_frame_rollouts = real_env.current_epoch_rollouts(
-      split=tf.estimator.ModeKeys.TRAIN,
-      minimal_rollout_frames=frame_stack_size,
+      split=split, minimal_rollout_frames=frame_stack_size,
   )
   def initial_frame_chooser(batch_size):
     """Frame chooser."""
@@ -277,9 +424,15 @@ def augment_observation(
 
 def run_rollouts(
     env, agent, initial_observations, step_limit=None, discount_factor=1.0,
-    log_every_steps=None, video_writers=(), color_bar=False
+    log_every_steps=None, video_writers=(), color_bar=False,
+    many_rollouts_from_each_env=False
 ):
   """Runs a batch of rollouts from given initial observations."""
+  assert step_limit is not None or not many_rollouts_from_each_env, (
+      "When collecting many rollouts from each environment, time limit must "
+      "be set."
+  )
+
   num_dones = 0
   first_dones = [False] * env.batch_size
   observations = initial_observations
@@ -309,13 +462,14 @@ def proceed():
     observations = list(observations)
     now_done_indices = []
     for (i, done) in enumerate(dones):
-      if done and not first_dones[i]:
+      if done and (not first_dones[i] or many_rollouts_from_each_env):
         now_done_indices.append(i)
         first_dones[i] = True
         num_dones += 1
     if now_done_indices:
-      # Reset only envs done the first time in this timestep to ensure that
-      # we collect exactly 1 rollout from each env.
+      # Unless many_rollouts_from_each_env, reset only envs done the first time
+      # in this timestep to ensure that we collect exactly 1 rollout from each
+      # env.
       reset_observations = env.reset(now_done_indices)
       for (i, observation) in zip(now_done_indices, reset_observations):
         observations[i] = observation
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index dcee57a27..2ddb18ae2 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -33,19 +33,13 @@
 import random
 import time
 
-import numpy as np
 import six
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
-from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl
-from tensor2tensor.models.research.rl import make_simulated_env_fn_from_hparams
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl import trainer_model_based_params
-from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image
-from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw
 from tensor2tensor.rl.restarter import Restarter
-from tensor2tensor.rl.rl_utils import absolute_hinge_difference
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -149,7 +143,7 @@ def train_agent(real_env, learner, world_model_dir, hparams, epoch):
       real_env, hparams.frame_stack_size, hparams.simulation_random_starts,
       hparams.simulation_flip_first_random_for_beginning
   )
-  env_fn = make_simulated_env_fn_from_hparams(
+  env_fn = rl.make_simulated_env_fn_from_hparams(
       real_env, hparams, batch_size=hparams.simulated_batch_size,
       initial_frame_chooser=initial_frame_chooser, model_dir=world_model_dir,
       sim_video_dir=os.path.join(
@@ -233,129 +227,6 @@ def train_world_model(
   return world_model_steps_num
 
 
-def evaluate_world_model(real_env, hparams, world_model_dir, debug_video_path):
-  """Evaluate the world model (reward accuracy)."""
-  frame_stack_size = hparams.frame_stack_size
-  rollout_subsequences = []
-  def initial_frame_chooser(batch_size):
-    assert batch_size == len(rollout_subsequences)
-    return np.stack([
-        [frame.observation.decode() for frame in subsequence[:frame_stack_size]]
-        for subsequence in rollout_subsequences
-    ])
-
-  env_fn = make_simulated_env_fn_from_hparams(
-      real_env, hparams, batch_size=hparams.wm_eval_batch_size,
-      initial_frame_chooser=initial_frame_chooser, model_dir=world_model_dir
-  )
-  sim_env = env_fn(in_graph=False)
-  subsequence_length = int(
-      max(hparams.wm_eval_rollout_ratios) * hparams.simulated_rollout_length
-  )
-  rollouts = real_env.current_epoch_rollouts(
-      split=tf.estimator.ModeKeys.EVAL,
-      minimal_rollout_frames=(subsequence_length + frame_stack_size)
-  )
-
-  video_writer = common_video.WholeVideoWriter(
-      fps=10, output_path=debug_video_path, file_format="avi"
-  )
-
-  reward_accuracies_by_length = {
-      int(ratio * hparams.simulated_rollout_length): []
-      for ratio in hparams.wm_eval_rollout_ratios
-  }
-  for _ in range(hparams.wm_eval_num_batches):
-    rollout_subsequences[:] = random_rollout_subsequences(
-        rollouts, hparams.wm_eval_batch_size,
-        subsequence_length + frame_stack_size
-    )
-
-    eval_subsequences = [
-        subsequence[(frame_stack_size - 1):]
-        for subsequence in rollout_subsequences
-    ]
-
-    # Check that the initial observation is the same in the real and simulated
-    # rollout.
-    sim_init_obs = sim_env.reset()
-    def decode_real_obs(index):
-      return np.stack([
-          subsequence[index].observation.decode()
-          for subsequence in eval_subsequences  # pylint: disable=cell-var-from-loop
-      ])
-    real_init_obs = decode_real_obs(0)
-    assert np.all(sim_init_obs == real_init_obs)
-
-    debug_frame_batches = []
-    def append_debug_frame_batch(sim_obs, real_obs, sim_cum_rews,
-                                 real_cum_rews, sim_rews, real_rews):
-      """Add a debug frame."""
-      rews = [[sim_cum_rews, sim_rews], [real_cum_rews, real_rews]]
-      headers = []
-      for j in range(len(sim_obs)):
-        local_nps = []
-        for i in range(2):
-          img = PIL_Image().new("RGB", (sim_obs.shape[-2], 11),)
-          draw = PIL_ImageDraw().Draw(img)
-          draw.text((0, 0), "c:{:3}, r:{:3}".format(int(rews[i][0][j]),
-                                                    int(rews[i][1][j])),
-                    fill=(255, 0, 0))
-          local_nps.append(np.asarray(img))
-        local_nps.append(np.zeros_like(local_nps[0]))
-        headers.append(np.concatenate(local_nps, axis=1))
-      errs = absolute_hinge_difference(sim_obs, real_obs)
-      headers = np.stack(headers)
-      debug_frame_batches.append(  # pylint: disable=cell-var-from-loop
-          np.concatenate([headers,
-                          np.concatenate([sim_obs, real_obs, errs], axis=2)],
-                         axis=1)
-      )
-    append_debug_frame_batch(sim_init_obs, real_init_obs,
-                             np.zeros(hparams.wm_eval_batch_size),
-                             np.zeros(hparams.wm_eval_batch_size),
-                             np.zeros(hparams.wm_eval_batch_size),
-                             np.zeros(hparams.wm_eval_batch_size))
-
-    (sim_cum_rewards, real_cum_rewards) = (
-        np.zeros(hparams.wm_eval_batch_size) for _ in range(2)
-    )
-    for i in range(subsequence_length):
-      actions = [subsequence[i].action for subsequence in eval_subsequences]
-      (sim_obs, sim_rewards, _) = sim_env.step(actions)
-      sim_cum_rewards += sim_rewards
-
-      real_rewards = np.array([
-          subsequence[i + 1].reward for subsequence in eval_subsequences
-      ])
-      real_cum_rewards += real_rewards
-      for (length, reward_accuracies) in six.iteritems(
-          reward_accuracies_by_length
-      ):
-        if i + 1 == length:
-          reward_accuracies.append(
-              np.sum(sim_cum_rewards == real_cum_rewards) /
-              len(real_cum_rewards)
-          )
-
-      real_obs = decode_real_obs(i + 1)
-      append_debug_frame_batch(sim_obs, real_obs, sim_cum_rewards,
-                               real_cum_rewards, sim_rewards, real_rewards)
-
-    for debug_frames in np.stack(debug_frame_batches, axis=1):
-      for debug_frame in debug_frames:
-        video_writer.write(debug_frame)
-
-  video_writer.finish_to_disk()
-
-  return {
-      "reward_accuracy/at_{}".format(length): np.mean(reward_accuracies)
-      for (length, reward_accuracies) in six.iteritems(
-          reward_accuracies_by_length
-      )
-  }
-
-
 def load_metrics(event_dir, epoch):
   """Loads metrics for this epoch if they have already been written.
 
@@ -479,7 +350,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
             directories["world_model", "debug_videos"],
             "{}.avi".format(env.current_epoch)
         )
-        wm_metrics = evaluate_world_model(
+        wm_metrics = rl_utils.evaluate_world_model(
             env, hparams, directories["world_model"], debug_video_path
         )
         log("World model eval metrics:\n{}".format(pprint.pformat(wm_metrics)))

From 7db1cecf933354dcc2cd76c2a291b65679442b7b Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Tue, 29 Jan 2019 13:45:01 -0800
Subject: [PATCH 1596/2720] internal merge of PR #1419

PiperOrigin-RevId: 231465452
---
 tensor2tensor/rl/evaluator.py | 1 +
 tensor2tensor/rl/rl_utils.py  | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 03bd87fe5..a842d609f 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -272,6 +272,7 @@ def collect_frames_for_random_starts(
     log_every_steps=None
 ):
   """Collects frames from real env for random starts of simulated env."""
+  del frame_stack_size
   storage_env.start_new_epoch(0)
   tf.logging.info(
       "Collecting %d frames for random starts.", random_starts_step_limit
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 6e225883b..d385f04b4 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -127,7 +127,7 @@ def evaluate_world_model(
   def initial_frame_chooser(batch_size):
     assert batch_size == len(rollout_subsequences)
     return np.stack([
-        [frame.observation.decode() for frame in subsequence[:frame_stack_size]]
+        [frame.observation.decode() for frame in subsequence[:frame_stack_size]]    # pylint: disable=g-complex-comprehension
         for subsequence in rollout_subsequences
     ])
 

From af17fc746f56bbbd553ba889d2291eb941331dcb Mon Sep 17 00:00:00 2001
From: Dominic Jack <thedomjack@gmail.com>
Date: Tue, 29 Jan 2019 15:34:47 -0800
Subject: [PATCH 1597/2720] internal merge of PR #1410

PiperOrigin-RevId: 231486407
---
 tensor2tensor/bin/t2t_attack.py               |   2 +-
 tensor2tensor/bin/t2t_datagen.py              |   5 +-
 tensor2tensor/bin/t2t_prune.py                |   2 +-
 tensor2tensor/layers/common_hparams.py        |   4 +-
 tensor2tensor/models/mtf_transformer2.py      |   2 +-
 .../models/research/adafactor_experiments.py  |  10 +-
 tensor2tensor/models/research/autoencoders.py |   2 +-
 .../models/research/transformer_nat.py        |   4 +-
 .../models/research/transformer_vae.py        |   4 +-
 .../models/research/vqa_attention.py          |   2 +-
 .../models/research/vqa_self_attention.py     |   2 +-
 tensor2tensor/models/shake_shake.py           |   2 +-
 tensor2tensor/models/transformer.py           |  10 +-
 tensor2tensor/models/vanilla_gan.py           |   2 +-
 tensor2tensor/problems.py                     |   2 +-
 tensor2tensor/rl/datagen_with_agent.py        |   2 +-
 .../transformer_test_ckpt/hparams.json        |   2 +-
 tensor2tensor/utils/adafactor.py              |   2 +-
 tensor2tensor/utils/learning_rate.py          |   2 +-
 tensor2tensor/utils/optimize.py               |  31 +-
 tensor2tensor/utils/registry.py               | 785 +++++++++---------
 tensor2tensor/utils/registry_test.py          | 179 +++-
 22 files changed, 592 insertions(+), 466 deletions(-)

diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 6276d6b4c..f7e0fcc60 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -73,7 +73,7 @@ def create_attack_params():
 
 
 def create_attack(attack):
-  return registry.attacks(attack)
+  return registry.attack(attack)
 
 
 def create_surrogate_hparams():
diff --git a/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index b778bc984..c32a86cf5 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -147,7 +147,7 @@ def main(_):
 
   # Calculate the list of problems to generate.
   problems = sorted(
-      list(_SUPPORTED_PROBLEM_GENERATORS) + registry.list_problems())
+      list(_SUPPORTED_PROBLEM_GENERATORS) + registry.list_base_problems())
   for exclude in FLAGS.exclude_problems.split(","):
     if exclude:
       problems = [p for p in problems if exclude not in p]
@@ -169,7 +169,8 @@ def main(_):
 
   if not problems:
     problems_str = "\n  * ".join(
-        sorted(list(_SUPPORTED_PROBLEM_GENERATORS) + registry.list_problems()))
+        sorted(list(_SUPPORTED_PROBLEM_GENERATORS) +
+               registry.list_base_problems()))
     error_msg = ("You must specify one of the supported problems to "
                  "generate data for:\n  * " + problems_str + "\n")
     error_msg += ("TIMIT and parsing need data_sets specified with "
diff --git a/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
index 008462e63..e7c8c75db 100644
--- a/tensor2tensor/bin/t2t_prune.py
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -54,7 +54,7 @@ def create_pruning_params():
 
 
 def create_pruning_strategy(name):
-  return registry.pruning_strategies(name)
+  return registry.pruning_strategy(name)
 
 
 def main(argv):
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index b422026ce..0d5ce6598 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -55,7 +55,7 @@ def basic_params1():
       initializer="orthogonal",
       initializer_gain=1.5,
       label_smoothing=0.1,
-      optimizer="Adam",
+      optimizer="adam",
       optimizer_adam_epsilon=1e-6,
       optimizer_adam_beta1=0.85,
       optimizer_adam_beta2=0.997,
@@ -466,7 +466,7 @@ def basic_range1(ranged_hparams):
   rhp.set_float("optimizer_adam_beta2", 0.995, 0.999)
   rhp.set_categorical(
       "optimizer",
-      ["Adam", "Adagrad", "Momentum", "RMSProp", "SGD", "YellowFin"])
+      ["adam", "adagrad", "momentum", "rms_prop", "sgd", "yellow_fin"])
 
 
 @registry.register_ranged_hparams
diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index 977214a49..fb5842f87 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -269,7 +269,7 @@ def sample(self, features, mesh):
     return self.combine_batch_dims(ret)
 
 
-layers_registry = registry.create_registry("layers")
+layers_registry = registry.Registries.mtf_layers
 
 
 # The following functions construct layers based on hyperparmeters
diff --git a/tensor2tensor/models/research/adafactor_experiments.py b/tensor2tensor/models/research/adafactor_experiments.py
index d7d3d4e2c..fbe4dbc2f 100644
--- a/tensor2tensor/models/research/adafactor_experiments.py
+++ b/tensor2tensor/models/research/adafactor_experiments.py
@@ -30,16 +30,16 @@ def mimic_adam_with_adafactor(hparams):
   Some minor things may be different, like epsilon and beta1 correction.
 
   Args:
-    hparams: model hyperparameters where "Adam" in hparams.optimizer
+    hparams: model hyperparameters where "adam" in hparams.optimizer
   """
-  assert "Adam" in hparams.optimizer
-  hparams.optimizer = "Adafactor"
+  assert "adam" in hparams.optimizer
+  hparams.optimizer = "adafactor"
   hparams.optimizer_adafactor_beta1 = hparams.optimizer_adam_beta1
   hparams.optimizer_adafactor_beta2 = hparams.optimizer_adam_beta2
   hparams.optimizer_adafactor_multiply_by_parameter_scale = False
   hparams.optimizer_adafactor_factored = False
   hparams.optimizer_adafactor_clipping_threshold = None
-  hparams.optimizer_adafactor_decay_type = "Adam"
+  hparams.optimizer_adafactor_decay_type = "adam"
 
 
 @registry.register_hparams
@@ -50,7 +50,7 @@ def afx_adam():
   hparams.optimizer_adam_beta2 = 0.999
   hparams.symbol_modality_num_shards = 1
   hparams.batch_size = 2048
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.learning_rate_schedule = (
       "constant*rsqrt_decay*linear_warmup*rsqrt_hidden_size")
   hparams.learning_rate_constant = 2.0
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 4d7bac8b0..b9c42e89e 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1020,7 +1020,7 @@ def body(self, features):
 def autoencoder_basic():
   """Basic autoencoder model."""
   hparams = common_hparams.basic_params1()
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.learning_rate_constant = 0.0002
   hparams.learning_rate_warmup_steps = 500
   hparams.learning_rate_schedule = "constant * linear_warmup"
diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 7a775153d..0b11e89cf 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -25,7 +25,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 import tensorflow as tf
-from tensorflow.python.training import moving_averages
+from tensorflow.python.training import moving_averages  # pylint: disable=g-direct-tensorflow-import
 
 
 def init_vq_bottleneck(bottleneck_size, hidden_size):
@@ -392,7 +392,7 @@ def transformer_nat_small():
   hparams.filter_size = 2048
   hparams.label_smoothing = 0.0
   hparams.force_full_predict = True
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.optimizer_adam_epsilon = 1e-9
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.997
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index c630abc6c..64aecd405 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -767,7 +767,7 @@ def transformer_ae_small():
   hparams.filter_size = 2048
   hparams.add_hparam("compress_filter_size", 2048 * 2)
   hparams.label_smoothing = 0.0
-  hparams.optimizer = "Adam"  # Can be unstable, maybe try Adam.
+  hparams.optimizer = "adam"  # Can be unstable, maybe try Adam.
   hparams.optimizer_adam_epsilon = 1e-9
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.997  # Needs tuning, try 0.98 to 0.999.
@@ -941,7 +941,7 @@ def transformer_ae_a3():
 def transformer_ae_a6():
   """Best hparams for transformer with semhash."""
   hparams = transformer_ae_a3()
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.noise_dev = 0.5
   return hparams
 
diff --git a/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index 774387878..57f1975e6 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -335,7 +335,7 @@ def vqa_attention_base():
   hparams = common_hparams.basic_params1()
   hparams.batch_size = 128
   hparams.use_fixed_batch_size = True,
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.999
   hparams.optimizer_adam_epsilon = 1e-8
diff --git a/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
index 482e18c9e..b8388e606 100644
--- a/tensor2tensor/models/research/vqa_self_attention.py
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -684,7 +684,7 @@ def vqa_self_attention_base():
   hparams = common_hparams.basic_params1()
   hparams.batch_size = 128
   hparams.use_fixed_batch_size = True,
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.997
   hparams.optimizer_adam_epsilon = 1e-9
diff --git a/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index 4af05dc21..22f46f989 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -189,7 +189,7 @@ def shakeshake_small():
 @registry.register_hparams
 def shake_shake_quick():
   hparams = shakeshake_small()
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.learning_rate_cosine_cycle_steps = 1000
   hparams.learning_rate = 0.5
   hparams.batch_size = 100
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index f0d0596e8..d7334bb25 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -41,8 +41,10 @@
 
 import tensorflow as tf
 
+# pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.ops import inplace_ops
 from tensorflow.python.util import nest
+# pylint: enable=g-direct-tensorflow-import
 
 
 # Alias some commonly reused layers, here and elsewhere.
@@ -770,7 +772,7 @@ def fast_decode_tpu(encoder_output,
       hparams.num_heads if hparams.get("attention_variables_3d") else 0)
 
   cache = {
-      "layer_%d" % layer: {
+      "layer_%d" % layer: {  # pylint: disable=g-complex-comprehension
           "k":
           common_attention.split_heads(
               tf.zeros([batch_size, decode_length, key_channels]),
@@ -962,7 +964,7 @@ def fast_decode(encoder_output,
   if cache is None:
     cache = {}
   cache.update({
-      "layer_%d" % layer: {
+      "layer_%d" % layer: {  # pylint: disable=g-complex-comprehension
           "k":
               common_attention.split_heads(
                   tf.zeros([batch_size, 0, key_channels]), hparams.num_heads),
@@ -1694,7 +1696,7 @@ def transformer_tall_pretrain_lm():
   hparams.learning_rate_constant = 2e-4
   hparams.learning_rate_schedule = (
       "linear_warmup*constant*cosdecay")
-  hparams.optimizer = "AdamW"
+  hparams.optimizer = "adam_w"
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.999
   hparams.optimizer_adam_epsilon = 1e-8
@@ -1739,7 +1741,7 @@ def transformer_tall_pretrain_lm_tpu():
   # Optimizer gets reset in update_hparams_for_tpu so we set it again here.
   hparams.learning_rate_constant = 2e-4
   hparams.learning_rate_schedule = ("linear_warmup * constant * cosdecay")
-  hparams.optimizer = "AdamW"
+  hparams.optimizer = "adam_w"
   return hparams
 
 
diff --git a/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index cb16a4c0b..00fa89d18 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -199,7 +199,7 @@ def infer(self, *args, **kwargs):  # pylint: disable=arguments-differ
 def sliced_gan():
   """Basic parameters for a vanilla_gan."""
   hparams = common_hparams.basic_params1()
-  hparams.optimizer = "Adam"
+  hparams.optimizer = "adam"
   hparams.learning_rate_constant = 0.0002
   hparams.learning_rate_warmup_steps = 500
   hparams.learning_rate_schedule = "constant * linear_warmup"
diff --git a/tensor2tensor/problems.py b/tensor2tensor/problems.py
index 1337c2a33..2cf8c8762 100644
--- a/tensor2tensor/problems.py
+++ b/tensor2tensor/problems.py
@@ -27,7 +27,7 @@ def problem(name):
 
 
 def available():
-  return sorted(registry.list_problems())
+  return registry.list_base_problems()
 
 
 all_problems.import_modules(all_problems.ALL_MODULES)
diff --git a/tensor2tensor/rl/datagen_with_agent.py b/tensor2tensor/rl/datagen_with_agent.py
index f634e8c87..1e5100537 100644
--- a/tensor2tensor/rl/datagen_with_agent.py
+++ b/tensor2tensor/rl/datagen_with_agent.py
@@ -45,7 +45,7 @@ def main(_):
 
   # Create problem if not already defined
   problem_name = "gym_discrete_problem_with_agent_on_%s" % FLAGS.game
-  if problem_name not in registry.list_problems():
+  if problem_name not in registry.Registries.problems:
     gym_env.register_game(FLAGS.game)
 
   # Generate
diff --git a/tensor2tensor/test_data/transformer_test_ckpt/hparams.json b/tensor2tensor/test_data/transformer_test_ckpt/hparams.json
index df9a654c0..1a6a97223 100644
--- a/tensor2tensor/test_data/transformer_test_ckpt/hparams.json
+++ b/tensor2tensor/test_data/transformer_test_ckpt/hparams.json
@@ -1 +1 @@
-{"daisy_chain_variables": true, "optimizer_adam_beta1": 0.9, "scheduled_sampling_prob": 0.0, "num_hidden_layers": 2, "moe_loss_coef": 0.01, "max_target_seq_length": 0, "clip_grad_norm": 0.0, "pos": "timing", "scheduled_sampling_gold_mixin_prob": 0.5, "initializer": "uniform_unit_scaling", "grad_noise_scale": 0.0, "optimizer_momentum_momentum": 0.9, "nbr_decoder_problems": 1, "attention_key_channels": 0, "eval_drop_long_sequences": false, "learning_rate_cosine_cycle_steps": 250000, "prepend_mode": "none", "weight_decay": 0.0, "symbol_modality_skip_top": false, "weight_noise": 0.0, "target_modality": "default", "attention_dropout": 0.1, "parameter_attention_value_channels": 0, "factored_logits": false, "relu_dropout": 0.1, "no_data_parallelism": false, "layer_preprocess_sequence": "n", "sampling_method": "argmax", "learning_rate": 0.2, "num_heads": 2, "max_length": 256, "summarize_grads": false, "attention_value_channels": 0, "num_encoder_layers": 0, "label_smoothing": 0.1, "use_fixed_batch_size": false, "optimizer": "Adam", "moe_k": 2, "self_attention_type": "dot_product", "learning_rate_decay_scheme": "noam", "sampling_temp": 1.0, "kernel_height": 3, "use_pad_remover": true, "batch_size": 4096, "max_relative_position": 0, "force_full_predict": false, "min_length_bucket": 8, "layer_prepostprocess_dropout": 0.1, "eval_run_autoregressive": false, "shared_embedding_and_softmax_weights": true, "symbol_modality_num_shards": 16, "dropout": 0.2, "compress_steps": 0, "parameter_attention_key_channels": 0, "length_bucket_step": 1.1, "kernel_width": 1, "hidden_size": 16, "num_decoder_layers": 0, "input_modalities": "default", "filter_size": 8, "optimizer_adam_beta2": 0.98, "scheduled_sampling_warmup_steps": 50000, "norm_type": "layer", "min_length": 0, "moe_num_experts": 64, "multiply_embedding_mode": "sqrt_depth", "max_input_seq_length": 0, "learning_rate_warmup_steps": 8000, "proximity_bias": false, "ffn_layer": "dense_relu_dense", "initializer_gain": 1.0, "layer_postprocess_sequence": "da", "moe_hidden_sizes": "2048", "optimizer_adam_epsilon": 1e-09, "norm_epsilon": 1e-06}
+{"daisy_chain_variables": true, "optimizer_adam_beta1": 0.9, "scheduled_sampling_prob": 0.0, "num_hidden_layers": 2, "moe_loss_coef": 0.01, "max_target_seq_length": 0, "clip_grad_norm": 0.0, "pos": "timing", "scheduled_sampling_gold_mixin_prob": 0.5, "initializer": "uniform_unit_scaling", "grad_noise_scale": 0.0, "optimizer_momentum_momentum": 0.9, "nbr_decoder_problems": 1, "attention_key_channels": 0, "eval_drop_long_sequences": false, "learning_rate_cosine_cycle_steps": 250000, "prepend_mode": "none", "weight_decay": 0.0, "symbol_modality_skip_top": false, "weight_noise": 0.0, "target_modality": "default", "attention_dropout": 0.1, "parameter_attention_value_channels": 0, "factored_logits": false, "relu_dropout": 0.1, "no_data_parallelism": false, "layer_preprocess_sequence": "n", "sampling_method": "argmax", "learning_rate": 0.2, "num_heads": 2, "max_length": 256, "summarize_grads": false, "attention_value_channels": 0, "num_encoder_layers": 0, "label_smoothing": 0.1, "use_fixed_batch_size": false, "optimizer": "adam", "moe_k": 2, "self_attention_type": "dot_product", "learning_rate_decay_scheme": "noam", "sampling_temp": 1.0, "kernel_height": 3, "use_pad_remover": true, "batch_size": 4096, "max_relative_position": 0, "force_full_predict": false, "min_length_bucket": 8, "layer_prepostprocess_dropout": 0.1, "eval_run_autoregressive": false, "shared_embedding_and_softmax_weights": true, "symbol_modality_num_shards": 16, "dropout": 0.2, "compress_steps": 0, "parameter_attention_key_channels": 0, "length_bucket_step": 1.1, "kernel_width": 1, "hidden_size": 16, "num_decoder_layers": 0, "input_modalities": "default", "filter_size": 8, "optimizer_adam_beta2": 0.98, "scheduled_sampling_warmup_steps": 50000, "norm_type": "layer", "min_length": 0, "moe_num_experts": 64, "multiply_embedding_mode": "sqrt_depth", "max_input_seq_length": 0, "learning_rate_warmup_steps": 8000, "proximity_bias": false, "ffn_layer": "dense_relu_dense", "initializer_gain": 1.0, "layer_postprocess_sequence": "da", "moe_hidden_sizes": "2048", "optimizer_adam_epsilon": 1e-09, "norm_epsilon": 1e-06}
diff --git a/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index 01b6145ba..9c44d28d9 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -326,7 +326,7 @@ def adafactor_optimizer_from_hparams(hparams, lr):
   Raises:
     ValueError: on illegal values
   """
-  if hparams.optimizer_adafactor_decay_type == "Adam":
+  if hparams.optimizer_adafactor_decay_type == "adam":
     decay_rate = adafactor_decay_rate_adam(
         hparams.optimizer_adafactor_beta2)
   elif hparams.optimizer_adafactor_decay_type == "pow":
diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index bf037a876..709f993e7 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -90,7 +90,7 @@ def legacy_learning_rate_schedule(hparams):
     warmup = _learning_rate_warmup(warmup_steps, hparams=hparams)
     decay = _learning_rate_decay(hparams, warmup_steps)
     ret = tf.where(step_num < warmup_steps, warmup, decay)
-  optimizer_correction = 0.002 if "Adam" in hparams.optimizer else 1.0
+  optimizer_correction = 0.002 if "adam" in hparams.optimizer else 1.0
   tf.logging.info("Base learning rate: %f", hparams.learning_rate)
   return ret * optimizer_correction * hparams.learning_rate
 
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 5c3b5b3bf..ac301189b 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -20,7 +20,7 @@
 import numpy as np
 
 from tensor2tensor.layers import common_layers
-from tensor2tensor.utils import adafactor
+from tensor2tensor.utils import adafactor as adafactor_lib
 from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import multistep_optimizer
 from tensor2tensor.utils import registry
@@ -29,8 +29,7 @@
 import tensorflow as tf
 
 
-from tensorflow.contrib.mixed_precision import LossScaleOptimizer
-from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import dtypes  # pylint: disable=g-direct-tensorflow-import
 
 
 def _mixed_precision_is_enabled(hparams):
@@ -94,7 +93,7 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
   return train_op
 
 
-@registry.register_optimizer("Adam")
+@registry.register_optimizer
 def adam(learning_rate, hparams):
   # We change the default epsilon for Adam.
   # Using LazyAdam as it's much faster for large vocabulary embeddings.
@@ -105,7 +104,7 @@ def adam(learning_rate, hparams):
       epsilon=hparams.optimizer_adam_epsilon)
 
 
-@registry.register_optimizer("MultistepAdam")
+@registry.register_optimizer
 def multistep_adam(learning_rate, hparams):
   return multistep_optimizer.MultistepAdamOptimizer(
       learning_rate,
@@ -115,7 +114,7 @@ def multistep_adam(learning_rate, hparams):
       n=hparams.optimizer_multistep_accumulate_steps)
 
 
-@registry.register_optimizer("Momentum")
+@registry.register_optimizer
 def momentum(learning_rate, hparams):
   return tf.train.MomentumOptimizer(
       learning_rate,
@@ -123,14 +122,14 @@ def momentum(learning_rate, hparams):
       use_nesterov=hparams.optimizer_momentum_nesterov)
 
 
-@registry.register_optimizer("YellowFin")
+@registry.register_optimizer
 def yellow_fin(learning_rate, hparams):
   return yellowfin.YellowFinOptimizer(
       learning_rate=learning_rate,
       momentum=hparams.optimizer_momentum_momentum)
 
 
-@registry.register_optimizer("TrueAdam")
+@registry.register_optimizer
 def true_adam(learning_rate, hparams):
   return tf.train.AdamOptimizer(
       learning_rate,
@@ -139,7 +138,7 @@ def true_adam(learning_rate, hparams):
       epsilon=hparams.optimizer_adam_epsilon)
 
 
-@registry.register_optimizer("AdamW")
+@registry.register_optimizer
 def adam_w(learning_rate, hparams):
   # Openai gpt used weight decay.
   # Given the internals of AdamW, weight decay dependent on the
@@ -156,9 +155,9 @@ def adam_w(learning_rate, hparams):
       epsilon=hparams.optimizer_adam_epsilon)
 
 
-@registry.register_optimizer("Adafactor")
-def register_adafactor(learning_rate, hparams):
-  return adafactor.adafactor_optimizer_from_hparams(hparams, learning_rate)
+@registry.register_optimizer
+def adafactor(learning_rate, hparams):
+  return adafactor_lib.adafactor_optimizer_from_hparams(hparams, learning_rate)
 
 
 
@@ -169,8 +168,11 @@ def _register_base_optimizer(key, fn):
 
 
 for k in tf.contrib.layers.OPTIMIZER_CLS_NAMES:
-  if k not in registry._OPTIMIZERS:  # pylint: disable=protected-access
+  if k not in registry.Registries.optimizers and k not in ("SGD", "RMSProp"):
     _register_base_optimizer(k, tf.contrib.layers.OPTIMIZER_CLS_NAMES[k])
+_register_base_optimizer("sgd", tf.contrib.layers.OPTIMIZER_CLS_NAMES["SGD"])
+_register_base_optimizer(
+    "rms_prop", tf.contrib.layers.OPTIMIZER_CLS_NAMES["RMSProp"])
 
 
 class ConditionalOptimizer(tf.train.Optimizer):
@@ -209,7 +211,8 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
             decr_every_n_nan_or_inf=2,
             incr_ratio=2,
             decr_ratio=0.5)
-        self._opt = LossScaleOptimizer(self._opt, manager)
+        self._opt = tf.contrib.mixed_precision.LossScaleOptimizer(
+            self._opt, manager)
 
     self._zero_grads = hparams.optimizer_zero_grads
 
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index d8683f234..d02a3c99d 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -13,31 +13,53 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Registry for models, hyperparameter settings, problem types, and datasets.
+"""Object registration.
+
+Registries are instances of `Registry`.
+
+See `Registries` for a centralized list of object registries
+(models, problems, hyperparameter sets, etc.).
+
+New functions and classes can be registered using `.register`. The can be
+accessed/queried similar to dictionaries, keyed by default by `snake_case`
+equivalents.
+
+```
+@Registries.models.register
+class MyModel(T2TModel):
+  ...
+
+'my_model' in Registries.models  # True
+for k in Registries.models:
+  print(k)  # prints 'my_model'
+model = Registries.models['my_model'](constructor_arg)
+```
+
+#### Legacy Support
 
 Define a new model by subclassing T2TModel and register it:
 
 ```
-@registry.register_model
+@register_model
 class MyModel(T2TModel):
   ...
 ```
 
-Access by snake-cased name: `registry.model("my_model")`. If you're using
+Access by snake-cased name: `model("my_model")`. If you're using
 `t2t_trainer.py`, you can pass on the command-line: `--model=my_model`.
 
-See all the models registered: `registry.list_models()`.
+See all the models registered: `list_models()`.
 
 For hyperparameter sets:
-  * Register: `registry.register_hparams`
-  * List: `registry.list_hparams`
-  * Retrieve by name: `registry.hparams`
+  * Register: `register_hparams`
+  * List: `list_hparams`
+  * Retrieve by name: `hparams`
   * Command-line flag in `t2t_trainer.py`: `--hparams_set=name`
 
 For hyperparameter ranges:
-  * Register: `registry.register_ranged_hparams`
-  * List: `registry.list_ranged_hparams`
-  * Retrieve by name: `registry.ranged_hparams`
+  * Register: `register_ranged_hparams`
+  * List: `list_ranged_hparams`
+  * Retrieve by name: `ranged_hparams`
   * Command-line flag in `t2t_trainer.py`: `--hparams_range=name`
 """
 from __future__ import absolute_import
@@ -48,475 +70,461 @@ class MyModel(T2TModel):
 
 from tensor2tensor.utils import misc_utils
 import tensorflow as tf
-from tensorflow.python.util import tf_inspect as inspect
-
-_ATTACKS = {}
-_ATTACK_PARAMS = {}
-_HPARAMS = {}
-_MODELS = {}
-_PROBLEMS = {}
-_PRUNING_PARAMS = {}
-_PRUNING_STRATEGY = {}
-_RANGED_HPARAMS = {}
 
-# Key: registry name, Value: Registry
-_GENERIC_REGISTRIES = {}
-Registry = collections.namedtuple(
-    "_Registry", ["register", "get", "list", "registry"])
+from tensorflow.python.util import tf_inspect as inspect  # pylint: disable=g-direct-tensorflow-import
 
 
-def registry(registry_name):
-  """Returns `Registry` created by `create_registry`."""
-  if registry_name not in _GENERIC_REGISTRIES:
-    raise KeyError("No registry named %s. Available:\n%s" % (
-        registry_name, sorted(_GENERIC_REGISTRIES)))
-  return _GENERIC_REGISTRIES[registry_name]
+def default_name(class_or_fn):
+  """Default name for a class or function.
 
-
-def create_registry(registry_name):
-  """Create a generic object registry.
+  This is the naming function by default for registries expecting classes or
+  functions.
 
   Args:
-    registry_name: str, name of the object registry.
+    class_or_fn: class or function to be named.
 
   Returns:
-    `Registry` that contains functions for register (decorator), get, and list.
-
-  Raises:
-    KeyError: if `registry_name` is a pre-existing registry.
+    Default name for registration.
   """
-  if registry_name in _GENERIC_REGISTRIES:
-    raise KeyError(
-        "Registry %s already exists." % registry_name)
+  return misc_utils.camelcase_to_snakecase(class_or_fn.__name__)
 
-  registry_ = {}
 
-  def register(name):
-    """Returns decorator to register an object."""
+default_object_name = lambda obj: default_name(type(obj))
 
-    def register_dec(obj):
-      if name in registry_:
-        raise KeyError(
-            "Registry %s already contains key %s." % (registry_name, name))
-      registry_[name] = obj
-      return obj
 
-    return register_dec
+class Registry(object):
+  """Dict-like class for managing function registrations.
 
-  def get(name):
-    if name not in registry_:
-      raise KeyError(
-          "Registry %s contains no object named %s" % (registry_name, name))
-    return registry_[name]
+  ```python
+  my_registry = Registry("custom_name")
 
-  def list_registry():
-    return sorted(registry_)
+  @my_registry.register
+  def my_func():
+    pass
 
-  registry_obj = Registry(
-      register=register,
-      get=get,
-      list=list_registry,
-      registry=registry_,
-  )
-  _GENERIC_REGISTRIES[registry_name] = registry_obj
-  return registry_obj
+  @my_registry.register()
+  def another_func():
+    pass
 
+  @my_registry.register("non_default_name")
+  def third_func(x, y, z):
+    pass
 
-def _reset():
-  for ctr in [_MODELS, _HPARAMS, _RANGED_HPARAMS, _ATTACK_PARAMS]:
-    ctr.clear()
+  def foo():
+    pass
 
+  my_registry.register()(foo)
+  my_registry.register("baz")(lambda (x, y): x + y)
+  my_register.register("bar")
 
-def default_name(obj_class):
-  """Convert a class name to the registry's default name for the class.
+  print(list(my_registry))
+  # ["my_func", "another_func", "non_default_name", "foo", "baz"]
+  # (order may vary)
+  print(my_registry["non_default_name"] is third_func)  # True
+  print("third_func" in my_registry)                    # False
+  print("bar" in my_registry)                           # False
+  my_registry["non-existent_key"]                       # raises KeyError
+  ```
 
-  Args:
-    obj_class: the name of a class
-
-  Returns:
-    The registry's default name for the class.
+  Optional validation, on_set callback and value transform also supported.
+  See `__init__` doc.
   """
-  return misc_utils.camelcase_to_snakecase(obj_class.__name__)
-
 
-def default_object_name(obj):
-  """Convert an object to the registry's default name for the object class.
-
-  Args:
-    obj: an object instance
+  def __init__(self,
+               registry_name,
+               default_key_fn=default_name,
+               validator=None,
+               on_set=None,
+               value_transformer=(lambda k, v: v)):
+    """Construct a new registry.
 
-  Returns:
-    The registry's default name for the class of the object.
-  """
-  return default_name(obj.__class__)
-
-
-def register_model(name=None):
-  """Register a model. name defaults to class name snake-cased."""
-
-  def decorator(model_cls, registration_name=None):
-    """Registers & returns model_cls with registration_name or default name."""
-    model_name = registration_name or default_name(model_cls)
-    if model_name in _MODELS and not tf.executing_eagerly():
-      raise LookupError("Model %s already registered." % model_name)
-    model_cls.REGISTERED_NAME = model_name
-    _MODELS[model_name] = model_cls
-    return model_cls
-
-  # Handle if decorator was used without parens
-  if callable(name):
-    model_cls = name
-    return decorator(model_cls, registration_name=default_name(model_cls))
-
-  return lambda model_cls: decorator(model_cls, name)
-
-
-def model(name):
-  if name not in _MODELS:
-    raise LookupError("Model %s never registered.  Available models:\n %s" %
-                      (name, "\n".join(list_models())))
-
-  return _MODELS[name]
-
-
-def list_models():
-  return list(sorted(_MODELS))
-
-
-_OPTIMIZERS = {}
-
-
-def register_optimizer(name=None):
-  """Register an optimizer. name defaults to upper camel case of fn name."""
-
-  def default_opt_name(opt_fn):
-    return misc_utils.snakecase_to_camelcase(default_name(opt_fn))
-
-  def decorator(opt_fn, registration_name):
-    """Registers and returns optimizer_fn with registration_name or default."""
-    if registration_name is None:
-      registration_name = default_opt_name(opt_fn)
-
-    if registration_name in _OPTIMIZERS and not tf.executing_eagerly():
-      raise LookupError("Optimizer %s already registered." % registration_name)
-    args, varargs, keywords, _ = inspect.getargspec(opt_fn)
+    Args:
+      registry_name: str identifier for the given registry. Used in error msgs.
+      default_key_fn (optional): function mapping value -> key for registration
+        when a key is not provided
+      validator (optional): if given, this is run before setting a given (key,
+        value) pair. Accepts (key, value) and should raise if there is a
+        problem. Overwriting existing keys is not allowed and is checked
+        separately. Values are also checked to be callable separately.
+      on_set (optional): callback function accepting (key, value) pair which is
+        run after an item is successfully set.
+      value_transformer (optional): if run, `__getitem__` will return
+        value_transformer(key, registered_value).
+    """
+    self._registry = {}
+    self._name = registry_name
+    self._default_key_fn = default_key_fn
+    self._validator = validator
+    self._on_set = on_set
+    self._value_transformer = value_transformer
+
+  def default_key(self, value):
+    """Default key used when key not provided. Uses function from __init__."""
+    return self._default_key_fn(value)
+
+  @property
+  def name(self):
+    return self._name
+
+  def validate(self, key, value):
+    """Validation function run before setting. Uses function from __init__."""
+    if self._validator is not None:
+      self._validator(key, value)
+
+  def on_set(self, key, value):
+    """Callback called on successful set. Uses function from __init__."""
+    if self._on_set is not None:
+      self._on_set(key, value)
+
+  def __setitem__(self, key, value):
+    """Validate, set, and (if successful) call `on_set` for the given item.
 
-    if len(args) != 2 or varargs is not None or keywords is not None:
-      raise ValueError("Optimizer registration function must take two "
-                       "arguments: learning_rate (float) and "
-                       "hparams (HParams).")
-    _OPTIMIZERS[registration_name] = opt_fn
-    return opt_fn
+    Args:
+      key: key to store value under. If `None`, `self.default_key(value)` is
+        used.
+      value: callable stored under the given key.
 
-  if callable(name):
-    opt_fn = name
-    registration_name = default_opt_name(opt_fn)
-    return decorator(opt_fn, registration_name=registration_name)
+    Raises:
+      KeyError: if key is already in registry.
+    """
+    if key is None:
+      key = self.default_key(value)
+    if key in self:
+      raise KeyError(
+          "key %s already registered in registry %s" % (key, self._name))
+    if not callable(value):
+      raise ValueError("value must be callable")
+    self.validate(key, value)
+    self._registry[key] = value
+    self.on_set(key, value)
+
+  def register(self, key_or_value=None):
+    """Decorator to register a function, or registration itself.
+
+    This is primarily intended for use as a decorator, either with or without
+    a key/parentheses.
+    ```python
+    @my_registry.register('key1')
+    def value_fn(x, y, z):
+      pass
+
+    @my_registry.register()
+    def another_fn(x, y):
+      pass
+
+    @my_registry.register
+    def third_func():
+      pass
+    ```
+
+    Note if key_or_value is provided as a non-callable, registration only
+    occurs once the returned callback is called with a callable as its only
+    argument.
+    ```python
+    callback = my_registry.register('different_key')
+    'different_key' in my_registry  # False
+    callback(lambda (x, y): x + y)
+    'different_key' in my_registry  # True
+    ```
 
-  return lambda opt_fn: decorator(opt_fn, name)
+    Args:
+      key_or_value (optional): key to access the registered value with, or the
+        function itself. If `None` (default), `self.default_key` will be called
+        on `value` once the returned callback is called with `value` as the only
+        arg. If `key_or_value` is itself callable, it is assumed to be the value
+        and the key is given by `self.default_key(key)`.
 
+    Returns:
+      decorated callback, or callback generated a decorated function.
+    """
 
-def optimizer(name):
-  if name not in _OPTIMIZERS:
-    raise LookupError("Optimizer %s never registered. "
-                      "Available optimizers:\n %s"
-                      % (name, "\n".join(list_optimizers())))
-  return _OPTIMIZERS[name]
+    def decorator(value, key):
+      self[key] = value
+      return value
 
+    # Handle if decorator was used without parens
+    if callable(key_or_value):
+      return decorator(value=key_or_value, key=None)
+    else:
+      return lambda value: decorator(value, key=key_or_value)
 
-def list_optimizers():
-  return list(sorted(_OPTIMIZERS))
+  def __getitem__(self, key):
+    if key not in self:
+      raise KeyError("%s never registered with registry %s. Available:\n %s" %
+                     (key, self.name, display_list_by_prefix(sorted(self), 4)))
+    value = self._registry[key]
+    return self._value_transformer(key, value)
 
+  def __contains__(self, key):
+    return key in self._registry
 
-def register_hparams(name=None):
-  """Register an HParams set. name defaults to function name snake-cased."""
+  def keys(self):
+    return self._registry.keys()
 
-  def decorator(hp_fn, registration_name=None):
-    """Registers & returns hp_fn with registration_name or default name."""
-    hp_name = registration_name or default_name(hp_fn)
-    if hp_name in _HPARAMS and not tf.executing_eagerly():
-      raise LookupError("HParams set %s already registered." % hp_name)
-    _HPARAMS[hp_name] = hp_fn
-    return hp_fn
+  def values(self):
+    return (self[k] for k in self)  # complicated because of transformer
 
-  # Handle if decorator was used without parens
-  if callable(name):
-    hp_fn = name
-    return decorator(hp_fn, registration_name=default_name(hp_fn))
+  def items(self):
+    return ((k, self[k]) for k in self)  # complicated because of transformer
 
-  return lambda hp_fn: decorator(hp_fn, name)
+  def __iter__(self):
+    return iter(self._registry)
 
+  def __len__(self):
+    return len(self._registry)
 
-def hparams(name):
-  """Retrieve registered hparams by name."""
-  if name not in _HPARAMS:
-    error_msg = "HParams set %s never registered. Sets registered:\n%s"
-    raise LookupError(
-        error_msg % (name,
-                     display_list_by_prefix(list_hparams(), starting_spaces=4)))
-  hp = _HPARAMS[name]()
-  if hp is None:
-    raise TypeError("HParams %s is None. Make sure the registered function "
-                    "returns the HParams object." % name)
-  return hp
+  def _clear(self):
+    self._registry.clear()
 
+  def get(self, key, default=None):
+    return self[key] if key in self else default
 
-def list_hparams(prefix=None):
-  if prefix:
-    return [name for name in _HPARAMS if name.startswith(prefix)]
-  return list(_HPARAMS)
 
+def _on_model_set(k, v):
+  v.REGISTERED_NAME = k
 
-def register_ranged_hparams(name=None):
-  """Register a RangedHParams set. name defaults to fn name snake-cased."""
 
-  def decorator(rhp_fn, registration_name=None):
-    """Registers & returns hp_fn with registration_name or default name."""
-    rhp_name = registration_name or default_name(rhp_fn)
-    if rhp_name in _RANGED_HPARAMS:
-      raise LookupError("RangedHParams set %s already registered." % rhp_name)
-    # Check that the fn takes a single argument
-    args, varargs, keywords, _ = inspect.getargspec(rhp_fn)
-    if len(args) != 1 or varargs is not None or keywords is not None:
-      raise ValueError("RangedHParams set function must take a single "
-                       "argument, the RangedHParams object.")
+def _nargs_validator(nargs, message):
+  """Makes validator for function to ensure it takes nargs args."""
+  if message is None:
+    message = "Registered function must take exactly %d arguments" % nargs
 
-    _RANGED_HPARAMS[rhp_name] = rhp_fn
-    return rhp_fn
+  def f(key, value):
+    del key
+    spec = inspect.getfullargspec(value)
+    if (len(spec.args) != nargs or spec.varargs is not None or
+        spec.varkw is not None):
+      raise ValueError(message)
 
-  # Handle if decorator was used without parens
-  if callable(name):
-    rhp_fn = name
-    return decorator(rhp_fn, registration_name=default_name(rhp_fn))
+  return f
 
-  return lambda rhp_fn: decorator(rhp_fn, name)
 
+ProblemSpec = collections.namedtuple("ProblemSpec",
+                                     ["base_name", "was_reversed", "was_copy"])
 
-def ranged_hparams(name):
-  if name not in _RANGED_HPARAMS:
-    raise LookupError("RangedHParams set %s never registered." % name)
-  return _RANGED_HPARAMS[name]
 
+def parse_problem_name(name):
+  """Determines if problem_name specifies a copy and/or reversal.
 
-def list_ranged_hparams():
-  return list(_RANGED_HPARAMS)
+  Args:
+    name: str, problem name, possibly with suffixes.
 
+  Returns:
+    ProblemSpec: namedtuple with ["base_name", "was_reversed", "was_copy"]
 
-def register_problem(name=None):
-  """Register a Problem. name defaults to cls name snake-cased."""
+  Raises:
+    ValueError if name contains multiple suffixes of the same type
+      ('_rev' or '_copy'). One of each is ok.
+  """
+  # Recursively strip tags until we reach a base name.
+  if name.endswith("_rev"):
+    base, was_reversed, was_copy = parse_problem_name(name[:-4])
+    if was_reversed:
+      # duplicate rev
+      raise ValueError(
+          "Invalid problem name %s: multiple '_rev' instances" % name)
+    return ProblemSpec(base, True, was_copy)
+  elif name.endswith("_copy"):
+    base, was_reversed, was_copy = parse_problem_name(name[:-5])
+    if was_copy:
+      raise ValueError(
+          "Invalid problem_name %s: multiple '_copy' instances" % name)
+    return ProblemSpec(base, was_reversed, True)
+  else:
+    return ProblemSpec(name, False, False)
+
+
+def get_problem_name(base_name, was_reversed=False, was_copy=False):
+  """Construct a problem name from base and reversed/copy options.
+
+  Inverse of `parse_problem_name`.
 
-  def decorator(p_cls, registration_name=None):
-    """Registers & returns p_cls with registration_name or default name."""
-    p_name = registration_name or default_name(p_cls)
-    if p_name in _PROBLEMS and not tf.executing_eagerly():
-      raise LookupError("Problem %s already registered." % p_name)
+  Args:
+    base_name: base problem name. Should not end in "_rev" or "_copy"
+    was_reversed: if the problem is to be reversed
+    was_copy: if the problem is to be copied
 
-    _PROBLEMS[p_name] = p_cls
-    p_cls.name = p_name
-    return p_cls
+  Returns:
+    string name consistent with use with `parse_problem_name`.
 
-  # Handle if decorator was used without parens
-  if callable(name):
-    p_cls = name
-    return decorator(p_cls, registration_name=default_name(p_cls))
+  Raises:
+    ValueError if `base_name` ends with "_rev" or "_copy"
+  """
+  if any(base_name.endswith(suffix) for suffix in ("_rev", "_copy")):
+    raise ValueError("`base_name` cannot end in '_rev' or '_copy'")
+  name = base_name
+  if was_copy:
+    name = "%s_copy" % name
+  if was_reversed:
+    name = "%s_rev" % name
+  return name
+
+
+def _problem_name_validator(k, v):
+  del v
+  if parse_problem_name(k).base_name != k:
+    raise KeyError(
+        "Invalid problem name: cannot end in %s or %s" % ("_rev", "_copy"))
 
-  return lambda p_cls: decorator(p_cls, name)
 
+def _on_problem_set(k, v):
+  v.name = k
 
-def problem(name):
-  """Retrieve a problem by name."""
 
-  def parse_problem_name(problem_name):
-    """Determines if problem_name specifies a copy and/or reversal.
+def _call_value(k, v):
+  del k
+  return v()
 
-    Args:
-      problem_name: str, problem name, possibly with suffixes.
 
-    Returns:
-      base_name: A string with the base problem name.
-      was_reversed: A boolean.
-      was_copy: A boolean.
-    """
-    # Recursively strip tags until we reach a base name.
-    if problem_name.endswith("_rev"):
-      base, _, was_copy = parse_problem_name(problem_name[:-4])
-      return base, True, was_copy
-    elif problem_name.endswith("_copy"):
-      base, was_reversed, _ = parse_problem_name(problem_name[:-5])
-      return base, was_reversed, True
-    else:
-      return problem_name, False, False
+def _hparams_value_transformer(key, value):
+  out = value()
+  if out is None:
+    raise TypeError("HParams %s is None. Make sure the registered function "
+                    "returns the HParams object" % key)
+  return out
 
-  base_name, was_reversed, was_copy = parse_problem_name(name)
 
-  if base_name not in _PROBLEMS:
-    all_problem_names = list_problems()
-    error_lines = ["%s not in the set of supported problems:" % base_name
-                  ] + all_problem_names
-    error_msg = "\n  * ".join(error_lines)
-    raise LookupError(error_msg)
-  return _PROBLEMS[base_name](was_reversed=was_reversed, was_copy=was_copy)
+class Registries(object):
+  """Object holding `Registry` objects."""
 
+  def __init__(self):
+    raise RuntimeError("Registries is not intended to be instantiated")
 
-def list_problems():
-  return sorted(list(_PROBLEMS))
+  models = Registry("models", on_set=_on_model_set)
 
+  optimizers = Registry(
+      "optimizers",
+      validator=_nargs_validator(
+          2, "Registered optimizer functions must take exactly two arguments: "
+          "learning_rate (float) and hparams (HParams)."))
 
-def register_attack(name=None):
-  """Register an attack HParams set. Same behaviour as register_hparams."""
+  hparams = Registry("hparams", value_transformer=_hparams_value_transformer)
 
-  def decorator(attack_fn, registration_name=None):
-    """Registers & returns attack_fn with registration_name or default name."""
-    attack_name = registration_name or default_name(attack_fn)
-    if attack_name in _ATTACKS and not tf.executing_eagerly():
-      raise LookupError("Attack %s already registered." % attack_name)
-    _ATTACKS[attack_name] = attack_fn
-    return attack_fn
+  ranged_hparams = Registry(
+      "ranged_hparams",
+      validator=_nargs_validator(
+          1, "Registered ranged_hparams functions must take a single argument, "
+          "the RangedHParams object."))
 
-  # Handle if decorator was used without parens
-  if callable(name):
-    attack_fn = name
-    return decorator(attack_fn, registration_name=default_name(attack_fn))
+  problems = Registry(
+      "problems", validator=_problem_name_validator, on_set=_on_problem_set)
 
-  return lambda attack_fn: decorator(attack_fn, name)
+  attacks = Registry("attacks", value_transformer=_call_value)
 
+  attack_params = Registry("attack_params", value_transformer=_call_value)
 
-def attacks(name):
-  """Retrieve registered attack by name."""
-  if name not in _ATTACKS:
-    error_msg = "Attack %s never registered. Sets registered:\n%s"
-    raise LookupError(
-        error_msg % (name,
-                     display_list_by_prefix(list_attacks(), starting_spaces=4)))
-  attack = _ATTACKS[name]()
-  if attack is None:
-    raise TypeError(
-        "Attack %s is None. Make sure the registered function returns a "
-        "`cleverhans.attack.Attack` object." % name)
-  return attack
+  pruning_params = Registry("pruning_params", value_transformer=_call_value)
 
+  pruning_strategies = Registry("pruning_strategies")
 
-def list_attacks(prefix=None):
-  if prefix:
-    return [name for name in _ATTACKS if name.startswith(prefix)]
-  return list(_ATTACKS)
+  mtf_layers = Registry(
+      "mtf_layers",
+      validator=_nargs_validator(
+          2, "Registered layer functions must take exaction two arguments: "
+          "hparams (HParams) and prefix (str)."))
 
 
-def register_attack_params(name=None):
-  """Register an attack HParams set. Same behaviour as register_hparams."""
+# consistent version of old API
+model = Registries.models.__getitem__
+list_models = lambda: sorted(Registries.models)
+register_model = Registries.models.register
 
-  def decorator(ap_fn, registration_name=None):
-    """Registers & returns ap_fn with registration_name or default name."""
-    ap_name = registration_name or default_name(ap_fn)
-    if ap_name in _ATTACK_PARAMS and not tf.executing_eagerly():
-      raise LookupError("Attack HParams set %s already registered." % ap_name)
-    _ATTACK_PARAMS[ap_name] = ap_fn
-    return ap_fn
 
-  # Handle if decorator was used without parens
-  if callable(name):
-    ap_fn = name
-    return decorator(ap_fn, registration_name=default_name(ap_fn))
+def optimizer(name):
+  """Get pre-registered optimizer keyed by name.
 
-  return lambda ap_fn: decorator(ap_fn, name)
+  `name` should be snake case, though SGD -> sgd, RMSProp -> rms_prop and
+  UpperCamelCase -> snake_case conversions included for legacy support.
 
+  Args:
+    name: name of optimizer used in registration. This should be a snake case
+      identifier, though others supported for legacy reasons.
 
-def attack_params(name):
-  """Retrieve registered aparams by name."""
-  if name not in _ATTACK_PARAMS:
-    error_msg = "Attack HParams set %s never registered. Sets registered:\n%s"
-    raise LookupError(
-        error_msg %
-        (name, display_list_by_prefix(list_attack_params(), starting_spaces=4)))
-  ap = _ATTACK_PARAMS[name]()
-  if ap is None:
-    raise TypeError("Attack HParams %s is None. Make sure the registered "
-                    "function returns the HParams object." % name)
-  return ap
+  Returns:
+    optimizer
+  """
+  warn_msg = ("Please update `registry.optimizer` callsite "
+              "(likely due to a `HParams.optimizer` value)")
+  if name == "SGD":
+    name = "sgd"
+    tf.logging.warning("'SGD' optimizer now keyed by 'sgd'. %s" % warn_msg)
+  elif name == "RMSProp":
+    name = "rms_prop"
+    tf.logging.warning(
+        "'RMSProp' optimizer now keyed by 'rms_prop'. %s" % warn_msg)
+  else:
+    snake_name = misc_utils.camelcase_to_snakecase(name)
+    if name != snake_name:
+      tf.logging.warning(
+          "optimizer names now keyed by snake_case names. %s" % warn_msg)
+      name = snake_name
+  return Registries.optimizers[name]
+
+
+list_optimizers = lambda: sorted(Registries.optimizers)
+register_optimizer = Registries.optimizers.register
+
+hparams = Registries.hparams.__getitem__
+register_hparams = Registries.hparams.register
 
 
-def list_attack_params(prefix=None):
+def list_hparams(prefix=None):
+  hp_names = sorted(Registries.hparams)
   if prefix:
-    return [name for name in _ATTACK_PARAMS if name.startswith(prefix)]
-  return list(_ATTACK_PARAMS)
-
+    hp_names = [name for name in hp_names if name.startswith(prefix)]
+  return hp_names
 
-def register_pruning_params(name=None):
-  """Register an pruning HParams set. Same behaviour as register_hparams."""
 
-  def decorator(pp_fn, registration_name=None):
-    """Registers & returns pp_fn with registration_name or default name."""
-    pp_name = registration_name or default_name(pp_fn)
-    if pp_name in _PRUNING_PARAMS and not tf.executing_eagerly():
-      raise LookupError("Pruning HParams set %s already registered." % pp_name)
-    _PRUNING_PARAMS[pp_name] = pp_fn
-    return pp_fn
+ranged_hparams = Registries.ranged_hparams.__getitem__
+list_ranged_hparams = lambda: sorted(Registries.ranged_hparams)
+register_ranged_hparams = Registries.ranged_hparams.register
 
-  # Handle if decorator was used without parens
-  if callable(name):
-    pp_fn = name
-    return decorator(pp_fn, registration_name=default_name(pp_fn))
+base_problem = Registries.problems.__getitem__
+list_base_problems = lambda: sorted(Registries.problems)
+register_base_problem = Registries.problems.register
 
-  return lambda pp_fn: decorator(pp_fn, name)
+# Keeping for back-compatibility
+list_problems = list_base_problems
+register_problem = register_base_problem
 
 
-def pruning_params(name):
-  """Retrieve registered pruning params by name."""
-  if name not in _PRUNING_PARAMS:
-    error_msg = "Pruning HParams set %s never registered. Sets registered:\n%s"
-    raise LookupError(error_msg % (
-        name, display_list_by_prefix(list_pruning_params(), starting_spaces=4)))
-  pp = _PRUNING_PARAMS[name]()
-  if pp is None:
-    raise TypeError("Pruning HParams %s is None. Make sure the registered "
-                    "function returns the HParams object." % name)
-  return pp
+def problem(problem_name):
+  """Get possibly copied/reversed problem registered in `base_registry`.
 
+  Args:
+    problem_name: string problem name. See `parse_problem_name`.
 
-def list_pruning_params(prefix=None):
-  if prefix:
-    return [name for name in _PRUNING_PARAMS if name.startswith(prefix)]
-  return list(_PRUNING_PARAMS)
-
-
-def register_pruning_strategy(name=None):
-  """Register an pruning strategy. Same behaviour as register_hparams."""
-
-  def decorator(ps_fn, registration_name=None):
-    """Registers & returns ps_fn with registration_name or default name."""
-    ps_name = registration_name or default_name(ps_fn)
-    if ps_name in _PRUNING_STRATEGY and not tf.executing_eagerly():
-      raise LookupError("Pruning strategy %s already registered." % ps_name)
-    _PRUNING_STRATEGY[ps_name] = ps_fn
-    return ps_fn
+  Returns:
+    possibly reversed/copied version of base problem registered in the given
+    registry.
+  """
+  spec = parse_problem_name(problem_name)
+  return Registries.problems[spec.base_name](
+      was_copy=spec.was_copy, was_reversed=spec.was_reversed)
 
-  # Handle if decorator was used without parens
-  if callable(name):
-    ps_fn = name
-    return decorator(ps_fn, registration_name=default_name(ps_fn))
 
-  return lambda ps_fn: decorator(ps_fn, name)
+attack = Registries.attacks.__getitem__
+list_attacks = lambda: sorted(Registries.attacks)
+register_attack = Registries.attacks.register
 
+attack_params = Registries.attack_params.__getitem__
+list_attack_params = lambda: sorted(Registries.attack_params)
+register_attack_params = Registries.attack_params.register
 
-def pruning_strategies(name):
-  """Retrieve registered pruning strategies by name."""
-  if name not in _PRUNING_STRATEGY:
-    error_msg = "Pruning strategy set %s never registered. Sets registered:\n%s"
-    raise LookupError(
-        error_msg % (name,
-                     display_list_by_prefix(
-                         list_pruning_strategies(), starting_spaces=4)))
-  ps = _PRUNING_STRATEGY[name]
-  if ps is None:
-    raise TypeError("Pruning strategy %s is None. Make sure to register the "
-                    "function." % name)
-  return ps
+pruning_params = Registries.pruning_params.__getitem__
+list_pruning_params = lambda: sorted(Registries.pruning_params)
+register_pruning_params = Registries.pruning_params.register
 
+pruning_strategy = Registries.pruning_strategies.__getitem__
+list_pruning_strategies = lambda: sorted(Registries.pruning_strategies)
+register_pruning_strategy = Registries.pruning_strategies.register
 
-def list_pruning_strategies(prefix=None):
-  if prefix:
-    return [name for name in _PRUNING_STRATEGY if name.startswith(prefix)]
-  return list(_PRUNING_STRATEGY)
+# deprecated functions - plurals inconsistent with rest
+# deprecation decorators added 2019-01-25
+attacks = tf.contrib.framework.deprecated(None, "Use registry.attack")(attack)
+pruning_strategies = tf.contrib.framework.deprecated(
+    None, "Use registry.pruning_strategy")(
+        pruning_strategy)
 
 
 def display_list_by_prefix(names_list, starting_spaces=0):
@@ -551,6 +559,9 @@ def help_string():
   Problems:
 %s
 
+  Optimizers:
+%s
+
   Attacks:
 %s
 
@@ -563,16 +574,16 @@ def help_string():
   Pruning Strategies:
 %s
 """
-  m, hp, rhp, probs, atks, ap, pp, ps = [
-      display_list_by_prefix(entries, starting_spaces=4) for entries in [
+  lists = tuple(
+      display_list_by_prefix(entries, starting_spaces=4) for entries in [  # pylint: disable=g-complex-comprehension
           list_models(),
           list_hparams(),
           list_ranged_hparams(),
-          list_problems(),
+          list_base_problems(),
+          list_optimizers(),
           list_attacks(),
           list_attack_params(),
           list_pruning_params(),
           list_pruning_strategies(),
-      ]
-  ]
-  return help_str % (m, hp, rhp, probs, atks, ap, pp, ps)
+      ])
+  return help_str % lists
diff --git a/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index c2515b3ba..e21536785 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -24,13 +24,76 @@
 
 import tensorflow as tf
 
-# pylint: disable=unused-variable
+
+# pylint: disable=unused-variable,unused-argument
+
+
+class RegistryClassTest(tf.test.TestCase):
+  """Test of base registry.Registry class."""
+
+  def testGetterSetter(self):
+    r = registry.Registry("test_registry")
+    r["hello"] = lambda: "world"
+    r["a"] = lambda: "b"
+    self.assertEqual(r["hello"](), "world")
+    self.assertEqual(r["a"](), "b")
+
+  def testDefaultKeyFn(self):
+    r = registry.Registry("test", default_key_fn=lambda x: x().upper())
+    r.register()(lambda: "hello")
+    self.assertEqual(r["HELLO"](), "hello")
+
+  def testNoKeyProvided(self):
+    r = registry.Registry("test")
+    def f():
+      return 3
+    r.register(f)
+    self.assertEqual(r["f"](), 3)
+
+  def testMembership(self):
+    r = registry.Registry("test_registry")
+    r["a"] = lambda: None
+    r["b"] = lambda: 4
+    self.assertTrue("a" in r)
+    self.assertTrue("b" in r)
+
+  def testIteration(self):
+    r = registry.Registry("test_registry")
+    r["a"] = lambda: None
+    r["b"] = lambda: 4
+    self.assertEqual(sorted(r), ["a", "b"])
+
+  def testLen(self):
+    r = registry.Registry("test_registry")
+    self.assertEqual(len(r), 0)
+    r["a"] = lambda: None
+    self.assertEqual(len(r), 1)
+    r["b"] = lambda: 4
+    self.assertEqual(len(r), 2)
+
+  def testTransformer(self):
+    r = registry.Registry(
+        "test_registry", value_transformer=lambda x, y: x + y())
+    r.register(3)(lambda: 5)
+    r.register(10)(lambda: 12)
+    self.assertEqual(r[3], 8)
+    self.assertEqual(r[10], 22)
+    self.assertEqual(set(r.values()), set((8, 22)))
+    self.assertEqual(set(r.items()), set(((3, 8), (10, 22))))
+
+  def testGet(self):
+    r = registry.Registry("test_registry", value_transformer=lambda k, v: v())
+    r["a"] = lambda: "xyz"
+    self.assertEqual(r.get("a"), "xyz")
+    self.assertEqual(r.get("a", 3), "xyz")
+    self.assertIsNone(r.get("b"))
+    self.assertEqual(r.get("b", 3), 3)
 
 
 class ModelRegistryTest(tf.test.TestCase):
 
   def setUp(self):
-    registry._reset()
+    registry.Registries.models._clear()
 
   def testT2TModelRegistration(self):
 
@@ -60,7 +123,7 @@ def model_fn():
     self.assertTrue(model is model_fn)
 
   def testUnknownModel(self):
-    with self.assertRaisesRegexp(LookupError, "never registered"):
+    with self.assertRaisesRegexp(KeyError, "never registered"):
       registry.model("not_registered")
 
   def testDuplicateRegistration(self):
@@ -69,7 +132,7 @@ def testDuplicateRegistration(self):
     def m1():
       pass
 
-    with self.assertRaisesRegexp(LookupError, "already registered"):
+    with self.assertRaisesRegexp(KeyError, "already registered"):
 
       @registry.register_model("m1")
       def m2():
@@ -88,10 +151,83 @@ def m2():
     self.assertSetEqual(set(["m1", "m2"]), set(registry.list_models()))
 
 
+class OptimizerRegistryTest(tf.test.TestCase):
+
+  def setUp(self):
+    registry.Registries.optimizers._clear()
+
+  def testRegistration(self):
+    @registry.register_optimizer
+    def my_optimizer(learning_rate, hparams):
+      return 3
+
+    @registry.register_optimizer("my_other_optimizer")
+    def another_optimizer(learning_rate, hparams):
+      return 5
+
+    self.assertEqual(registry.optimizer("my_optimizer"), my_optimizer)
+    self.assertEqual(
+        registry.optimizer("my_other_optimizer"), another_optimizer)
+
+  def testMembership(self):
+    @registry.register_optimizer
+    def my_optimizer(learning_rate, hparams):
+      return 3
+
+    @registry.register_optimizer("my_other_optimizer")
+    def another_optimizer(learning_rate, hparams):
+      return 5
+
+    self.assertTrue("my_optimizer" in registry.Registries.optimizers)
+    self.assertTrue("my_other_optimizer" in registry.Registries.optimizers)
+    self.assertFalse("another_optimizer" in registry.Registries.optimizers)
+    self.assertEqual(len(registry.Registries.optimizers), 2)
+
+  def testArgErrorCheck(self):
+    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
+      registry.Registries.optimizers.register("OneArgs")(lambda x: 4)
+    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
+      registry.Registries.optimizers.register("ThreeArgs")(
+          lambda x, y, z: 4)
+    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
+      registry.Registries.optimizers.register("NArgs")(lambda *args: 4)
+    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
+      registry.Registries.optimizers.register("Kwargs")(lambda **kargs: 4)
+    with self.assertRaisesRegexp(ValueError, "must take .* arguments"):
+      registry.Registries.optimizers.register("TwoAndKwargs")(
+          lambda a, b, **kargs: 4)
+
+  def testMultipleRegistration(self):
+    @registry.register_optimizer
+    def my_optimizer(learning_rate, hparams):
+      return 3
+
+    with self.assertRaisesRegexp(KeyError, "already registered"):
+
+      @registry.register_optimizer("my_optimizer")
+      def another_fn(learning_rate, hparams):
+        return 5
+
+  def testUnknownOptimizer(self):
+    with self.assertRaisesRegexp(KeyError, "never registered"):
+      registry.optimizer("not_registered_optimizer")
+
+  def testGetterSetterInterface(self):
+    def f(x, y):
+      return 3
+
+    k = "blah"
+    registry.Registries.optimizers[k] = f
+    self.assertEqual(registry.optimizer(k), f)
+    self.assertEqual(registry.Registries.optimizers[k], f)
+    self.assertEqual(registry.Registries.optimizers[k], registry.optimizer(k))
+
+
 class HParamRegistryTest(tf.test.TestCase):
 
   def setUp(self):
-    registry._reset()
+    registry.Registries.hparams._clear()
+    registry.Registries.ranged_hparams._clear()
 
   def testHParamSet(self):
 
@@ -121,9 +257,9 @@ def my_hparams_range(_):
     self.assertTrue(registry.ranged_hparams("a") is my_hparams_range)
 
   def testUnknownHparams(self):
-    with self.assertRaisesRegexp(LookupError, "never registered"):
+    with self.assertRaisesRegexp(KeyError, "never registered"):
       registry.hparams("not_registered")
-    with self.assertRaisesRegexp(LookupError, "never registered"):
+    with self.assertRaisesRegexp(KeyError, "never registered"):
       registry.ranged_hparams("not_registered")
 
   def testNoneHparams(self):
@@ -194,34 +330,7 @@ def rhp_bad2(a, b):  # pylint: disable=unused-argument
         pass
 
 
-class CreateRegistry(tf.test.TestCase):
-  """Test class for `create_registry`."""
-
-  def testCreateRegistry(self):
-    my_registry = registry.create_registry("test_reg1")
-    self.assertIs(my_registry, registry.registry("test_reg1"))
-
-    # Use as decorator on a fn
-    @my_registry.register("foo")
-    def some_fn(num):
-      return num + 2
-
-    # Register a regular object
-    pod_obj = 4
-    my_registry.register("bar")(pod_obj)
-
-    # Register a class
-    @my_registry.register("foobar")
-    class A(object):
-      pass
-
-    self.assertEqual(9, my_registry.get("foo")(7))
-    self.assertEqual(["bar", "foo", "foobar"], my_registry.list())
-    foobar = my_registry.get("foobar")
-    self.assertTrue(isinstance(foobar(), A))
-
-
-class RegistryTest(tf.test.TestCase):
+class RegistryHelpTest(tf.test.TestCase):
   """Test class for common functions."""
 
   def testRegistryHelp(self):

From 62ca4e8e0dd045c02dc678b69b599e10d59979b4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 29 Jan 2019 16:21:29 -0800
Subject: [PATCH 1598/2720] Adding data generators and problem definitions for
 synthetic transduction tasks: copy, reverse, flip bigrams.

PiperOrigin-RevId: 231495059
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 .../data_generators/transduction_problems.py  | 260 ++++++++++++++++++
 .../transduction_problems_test.py             |  83 ++++++
 3 files changed, 344 insertions(+)
 create mode 100644 tensor2tensor/data_generators/transduction_problems.py
 create mode 100644 tensor2tensor/data_generators/transduction_problems_test.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 688197ab0..c9a8e763f 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -72,6 +72,7 @@
     "tensor2tensor.data_generators.sst_binary",
     "tensor2tensor.data_generators.subject_verb_agreement",
     "tensor2tensor.data_generators.timeseries",
+    "tensor2tensor.data_generators.transduction_problems",
     "tensor2tensor.data_generators.translate_encs",
     "tensor2tensor.data_generators.translate_ende",
     "tensor2tensor.data_generators.translate_enet",
diff --git a/tensor2tensor/data_generators/transduction_problems.py b/tensor2tensor/data_generators/transduction_problems.py
new file mode 100644
index 000000000..01fd5f9a5
--- /dev/null
+++ b/tensor2tensor/data_generators/transduction_problems.py
@@ -0,0 +1,260 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""A suite of sequence transduction problems.
+
+Each problem generates pairs of tokenized input and output sequences which
+represent the effect of the transduction algorithm which must be learned.
+
+These problems are based on the benchmarks outlined in:
+
+Learning to Transduce with Unbounded Memory
+Edward Grefenstette, Karl Moritz Hermann, Mustafa Suleyman, Phil Blunsom
+https://arxiv.org/abs/1506.02516, 2015
+
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import random
+
+from six.moves import xrange  # pylint: disable=redefined-builtin
+
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+MAX_GENERATOR_ATTEMPTS = 100
+
+
+class TransductionProblem(text_problems.Text2TextProblem):
+  """Abstract base clase which all transduction problems inherit from.
+  """
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(TransductionProblem, self).__init__(was_reversed=False,
+                                              was_copy=False)
+    self.vocab = self.build_vocab()
+
+  @property
+  def num_symbols(self):
+    """The number of symbols that can be used as part of a sequence."""
+    return 128
+
+  def min_sequence_length(self, dataset_split):
+    """Determine the minimum sequence length given a dataset_split.
+
+    Args:
+      dataset_split: A problem.DatasetSplit.
+
+    Returns:
+      The minimum length that a sequence can be for this dataset_split.
+    """
+    return {
+        problem.DatasetSplit.TRAIN: 8,
+        problem.DatasetSplit.EVAL: 65,
+        problem.DatasetSplit.TEST: 65
+    }[dataset_split]
+
+  def max_sequence_length(self, dataset_split):
+    """Determine the maximum sequence length given a dataset_split.
+
+    Args:
+      dataset_split: A problem.DatasetSplit.
+
+    Returns:
+      The maximum length that a sequence can be for this dataset_split.
+    """
+    return {
+        problem.DatasetSplit.TRAIN: 64,
+        problem.DatasetSplit.EVAL: 128,
+        problem.DatasetSplit.TEST: 128
+    }[dataset_split]
+
+  def num_samples(self, dataset_split):
+    """Determine the dataset sized given a dataset_split.
+
+    Args:
+      dataset_split: A problem.DatasetSplit.
+
+    Returns:
+      The desired number of samples for this dataset_split.
+    """
+    return {
+        problem.DatasetSplit.TRAIN: 1000000,
+        problem.DatasetSplit.EVAL: 10000,
+        problem.DatasetSplit.TEST: 10000
+    }[dataset_split]
+
+  @property
+  def num_shards(self):
+    """Used to split up datasets into multiple files."""
+    return 10
+
+  @property
+  def is_generate_per_split(self):
+    return False
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.TOKEN
+
+  def sequence_length(self, dataset_split):
+    return random.randint(self.min_sequence_length(dataset_split),
+                          self.max_sequence_length(dataset_split))
+
+  def build_vocab(self):
+    return ["sym_%d" % i for i in xrange(1, self.num_symbols + 1)]
+
+  def get_or_create_vocab(self, data_dir, tmp_dir, force_get=False):
+    vocab_filename = os.path.join(data_dir, self.vocab_filename)
+    if not tf.gfile.Exists(vocab_filename):
+      encoder = text_encoder.TokenTextEncoder(None,
+                                              vocab_list=sorted(self.vocab))
+      encoder.store_to_file(vocab_filename)
+    else:
+      encoder = text_encoder.TokenTextEncoder(vocab_filename,
+                                              replace_oov=self.oov_token)
+    return encoder
+
+  def generate_random_sequence(self, dataset_split):
+    return [random.choice(self.vocab)
+            for _ in range(self.sequence_length(dataset_split))]
+
+  def transpose_sequence(self, input_sequence):
+    raise NotImplementedError()
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    for _ in xrange(self.num_samples(dataset_split)):
+      source = self.generate_random_sequence(dataset_split)
+      target = self.transpose_sequence(source)
+      yield {
+          "inputs": " ".join(source),
+          "targets": " ".join(target),
+      }
+
+
+@registry.register_problem
+class CopySequence(TransductionProblem):
+  """Reproduce a sequence exactly as it was input."""
+
+  def transpose_sequence(self, input_sequence):
+    return input_sequence
+
+
+@registry.register_problem
+class CopySequenceSmall(CopySequence):
+  """Same as CopySequence but with smaller sequences.
+  """
+
+  @property
+  def num_symbols(self):
+    return 64
+
+  def min_sequence_length(self, dataset_split):
+    return {
+        problem.DatasetSplit.TRAIN: 4,
+        problem.DatasetSplit.EVAL: 17,
+        problem.DatasetSplit.TEST: 17
+    }[dataset_split]
+
+  def max_sequence_length(self, dataset_split):
+    return {
+        problem.DatasetSplit.TRAIN: 16,
+        problem.DatasetSplit.EVAL: 32,
+        problem.DatasetSplit.TEST: 32
+    }[dataset_split]
+
+  def num_samples(self, dataset_split):
+    return {
+        problem.DatasetSplit.TRAIN: 100000,
+        problem.DatasetSplit.EVAL: 10000,
+        problem.DatasetSplit.TEST: 10000
+    }[dataset_split]
+
+
+@registry.register_problem
+class ReverseSequence(TransductionProblem):
+  """Reverses the order of the sequence.
+  """
+
+  def transpose_sequence(self, input_sequence):
+    return input_sequence[::-1]
+
+
+@registry.register_problem
+class ReverseSequenceSmall(ReverseSequence):
+  """Same as ReverseSequence but with smaller sequences.
+  """
+
+  @property
+  def num_symbols(self):
+    return 64
+
+  def min_sequence_length(self, dataset_split):
+    return {
+        problem.DatasetSplit.TRAIN: 4,
+        problem.DatasetSplit.EVAL: 17,
+        problem.DatasetSplit.TEST: 17
+    }[dataset_split]
+
+  def max_sequence_length(self, dataset_split):
+    return {
+        problem.DatasetSplit.TRAIN: 16,
+        problem.DatasetSplit.EVAL: 32,
+        problem.DatasetSplit.TEST: 32
+    }[dataset_split]
+
+  def num_samples(self, dataset_split):
+    return {
+        problem.DatasetSplit.TRAIN: 100000,
+        problem.DatasetSplit.EVAL: 10000,
+        problem.DatasetSplit.TEST: 10000
+    }[dataset_split]
+
+
+@registry.register_problem
+class FlipBiGramSequence(TransductionProblem):
+  """Flip every pair of tokens: 1 2 3 4 -> 2 1 4 3.
+  """
+
+  def sequence_length(self, dataset_split):
+    """Only generate sequences with even lengths.
+
+    Args:
+      dataset_split: A problem.DatasetSplit specifying which dataset the
+        sequence is a part of.
+
+    Returns:
+      An even number >= min_sequence_length(dataset_split)
+        and <= max_sequence_length(dataset_split)
+    """
+    min_length = self.min_sequence_length(dataset_split)
+    min_length += min_length % 2
+    max_length = self.max_sequence_length(dataset_split)
+    max_length -= max_length % 2
+    length = random.randint(min_length, max_length)
+    return length - (length % 2)
+
+  def transpose_sequence(self, input_sequence):
+    return [input_sequence[i+1] if i%2 == 0 else input_sequence[i-1]
+            for i in range(len(input_sequence))]
diff --git a/tensor2tensor/data_generators/transduction_problems_test.py b/tensor2tensor/data_generators/transduction_problems_test.py
new file mode 100644
index 000000000..a8dbe5b5e
--- /dev/null
+++ b/tensor2tensor/data_generators/transduction_problems_test.py
@@ -0,0 +1,83 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.data_generators.transduction_problems."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import shutil
+import tempfile
+
+from absl.testing import parameterized
+
+import numpy as np
+
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import transduction_problems
+
+import tensorflow as tf
+
+
+class TransductionProblem(parameterized.TestCase):
+
+  def setUp(self):
+    super(TransductionProblem, self).setUp()
+    # Create a temporary directory
+    self.test_dir = tempfile.mkdtemp()
+
+  def tearDown(self):
+    super(TransductionProblem, self).tearDown()
+    # Remove the directory after the test
+    shutil.rmtree(self.test_dir)
+
+  @parameterized.named_parameters(
+      ('CopySequence',
+       transduction_problems.CopySequence(),
+       lambda x: x),
+      ('CopySequenceSmall',
+       transduction_problems.CopySequenceSmall(),
+       lambda x: x),
+      ('FlipBiGramSequence',
+       transduction_problems.FlipBiGramSequence(),
+       lambda x: [x[i+1] if i%2 == 0 else x[i-1] for i in range(len(x))]),
+      ('ReverseSequence',
+       transduction_problems.ReverseSequence(),
+       lambda x: x[::-1]),
+      ('ReverseSequenceSmall',
+       transduction_problems.ReverseSequenceSmall(),
+       lambda x: x[::-1]),
+  )
+  def testTransduction(self, p, transformation):
+    data_dir = ''
+    dataset_split = problem.DatasetSplit.TEST
+    for sample in p.generate_samples(data_dir, self.test_dir, dataset_split):
+      input_tokens = sample['inputs'].split(' ')
+      target_tokens = sample['targets'].split(' ')
+      self.assertBetween(len(input_tokens),
+                         p.min_sequence_length(dataset_split),
+                         p.max_sequence_length(dataset_split))
+      self.assertBetween(len(target_tokens),
+                         p.min_sequence_length(dataset_split),
+                         p.max_sequence_length(dataset_split))
+
+      transformed_inputs = np.array(transformation(input_tokens))
+
+      np.testing.assert_equal(transformed_inputs, target_tokens)
+
+
+if __name__ == '__main__':
+  tf.test.main()

From be505cf5cf93869236eca64d92af8e52aca912c6 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 29 Jan 2019 22:21:57 -0800
Subject: [PATCH 1599/2720] `EnvProblem` class, that works for an arbitrary env
 and tries to make as few assumptions as possible.

PiperOrigin-RevId: 231534791
---
 tensor2tensor/envs/env_problem.py      | 758 +++++++++++++++++++++++++
 tensor2tensor/envs/env_problem_test.py | 244 ++++++++
 tensor2tensor/envs/gym_spaces_utils.py | 113 ++++
 tensor2tensor/envs/time_step.py        |  67 +++
 tensor2tensor/envs/time_step_test.py   |  57 ++
 tensor2tensor/envs/trajectory.py       | 243 ++++++++
 tensor2tensor/envs/trajectory_test.py  | 212 +++++++
 7 files changed, 1694 insertions(+)
 create mode 100644 tensor2tensor/envs/env_problem.py
 create mode 100644 tensor2tensor/envs/env_problem_test.py
 create mode 100644 tensor2tensor/envs/gym_spaces_utils.py
 create mode 100644 tensor2tensor/envs/time_step.py
 create mode 100644 tensor2tensor/envs/time_step_test.py
 create mode 100644 tensor2tensor/envs/trajectory.py
 create mode 100644 tensor2tensor/envs/trajectory_test.py

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
new file mode 100644
index 000000000..08a9893b2
--- /dev/null
+++ b/tensor2tensor/envs/env_problem.py
@@ -0,0 +1,758 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 class for envs that store their history.
+
+EnvProblem subclasses Problem and also implements the Gym interface (step,
+reset, render, close, seed)
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from gym.core import Env
+import numpy as np
+import six
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.envs import gym_spaces_utils
+from tensor2tensor.envs import trajectory
+from tensor2tensor.layers import modalities
+from tensor2tensor.rl import gym_utils
+import tensorflow as tf
+
+# Names for data fields in stored tf.Examples.
+TIMESTEP_FIELD = "timestep"
+ACTION_FIELD = "action"
+RAW_REWARD_FIELD = "raw_reward"
+PROCESSED_REWARD_FIELD = "reward"
+DONE_FIELD = "done"
+OBSERVATION_FIELD = "observation"
+
+
+class EnvProblem(Env, problem.Problem):
+  """An env which generates data like a problem class.
+
+  EnvProblem is both a gym Env and a Problem, since it subclasses both.
+
+  Conceptually it contains `batch_size` environments on which step (and reset)
+  are called. The data that is generated by the repeated application of step and
+  reset is stored within this class and is persisted on disk when we call
+  `generate_data` on it.
+
+  Subclasses *should* override the following functions, since they are used in
+  the `hparams` function to return modalities and vocab_sizes.
+  - input_modality
+  - input_input_vocab_size
+  - target_modality
+  - target_vocab_size
+
+  NON NATIVELY BATCHED ENVS:
+
+  The default implementations of the other major functions, should work well for
+  cases where the env is not batched by default ex: any gym env. In this case we
+  create `batch_size` number of envs and store them in a list. Any function then
+  that interacts with the envs, like reset, step or close goes over the env list
+  to do the needful, ex: when reset is called with specific indices we reset
+  only those indices, etc.
+
+  The usage of this class will look like the following:
+
+  # 1. Creates and initializes the env_problem.
+  ep = env_problem.EnvProblem(...)
+
+  # 2. One needs to call reset() at the start, this resets all envs.
+  ep.reset()
+
+  # 3. Call step with actions for all envs, i.e. len(action) = batch_size
+  obs, rewards, dones, infos = ep.step(actions)
+
+  # 4. Figure out which envs got done and reset only those.
+  ep.reset(indices=done_indices(dones))
+
+  # 5. Go back to Step #3 to further interact with the env or just dump the
+  # generated data to disk by calling:
+  ep.generate_data(...)
+
+  # 6. If we now need to use this object again to play a few more iterations
+  # perhaps with a different batch size or maybe not recording the data, then
+  # we need to re-initialize environments and do some book-keeping, call:
+  ep.initialize_environments(batch_size)
+
+  # 7. Go back to Step #2, i.e. reset all envs.
+
+  NOTE: Look at `EnvProblemTest.test_interaction_with_env` and/or
+  `EnvProblemTest.test_generate_data`
+
+  NOTE: We rely heavily that the underlying environments expose a gym style
+  interface, i.e. in addition to reset(), step() and close() we have access to
+  the following properties: observation_space, action_space, reward_range.
+
+  NATIVELY BATCHED ENVS:
+
+  If however, our env is a neural network, which can be batched by default, we
+  should:
+
+  # 1 - Give it a gym style interface, by overriding observation_space and
+  action_space.
+
+  # 2 - Override `_reset` and `_step` to do the reset and step in a natively
+  batched manner.
+
+  # 3 - More generally any function that iterates over the self._env list will
+  need to be overridden, ex: `_verify_same_spaces` and `initialize_environments`
+
+  KNOWN LIMITATIONS:
+
+  - observation_space and action_space should be subclasses of gym.spaces
+  - not all subclasses of gym.spaces are supported
+  - no support for continuous action spaces
+
+  """
+
+  def __init__(self,
+               base_env_name=None,
+               base_env_kwargs=None,
+               batch_size=1,
+               reward_range=(-np.inf, np.inf),
+               was_reversed=False,
+               was_copy=False):
+    """Initializes this class by creating the envs and managing trajectories.
+
+    Args:
+      base_env_name: (string) passed to `gym_utils.make_gym_env` to make the
+        underlying environment.
+      base_env_kwargs: (dict) passed to `gym_utils.make_gym_env` to make the
+        underlying environment.
+      batch_size: (int): How many envs to make in the non natively batched mode.
+      reward_range: (tuple(number, number)) the first element is the minimum
+        reward and the second is the maximum reward, used to clip and process
+        the raw reward in `process_rewards`.
+      was_reversed: (bool) should be false, passed to underlying init of
+        Problem.
+      was_copy: (bool) should be false, passed to underlying init of Problem.
+    """
+
+    # Assert on these since they don't make sense.
+    assert not was_reversed and not was_copy
+
+    # Call the super's ctor.
+    problem.Problem.__init__(self, was_reversed=was_reversed, was_copy=was_copy)
+
+    # Name for the base environment, will be used in `gym_utils.make_gym_env` in
+    # the default implementation of `initialize_environments`.
+    self._base_env_name = base_env_name
+
+    # Other arguments for initializing environments, will be used in
+    # `gym_utils.make_gym_env` in the default implementation of
+    # `initialize_environments`.
+    self._base_env_kwargs = base_env_kwargs
+    if not self._base_env_kwargs:
+      self._base_env_kwargs = {}
+
+    # An env generates data when it is given actions by an agent which is either
+    # a policy or a human -- this is supposed to be the `id` of the agent.
+    #
+    # In practice, this is used only to store (and possibly retrieve) history
+    # to an appropriate directory.
+    self._agent_id = "default"
+
+    # We clip rewards to this range before processing them further, as described
+    # in `process_rewards`.
+    self._reward_range = reward_range
+
+    # Initialize the environment(s).
+
+    # This can either be a list of environments of len `batch_size` or this can
+    # be a Neural Network, in which case it will be fed input with first
+    # dimension = `batch_size`.
+    self._envs = None
+
+    self._observation_space = None
+    self._action_space = None
+
+    # A data structure to hold the `batch_size` currently active trajectories
+    # and also the ones that are completed, i.e. done.
+    self._trajectories = None
+
+    self.initialize_environments(batch_size=batch_size)
+
+    # Assert that *all* the above are now set, we should do this since
+    # subclasses can override `initialize_environments`.
+    assert self._envs is not None
+    assert self._observation_space is not None
+    assert self._action_space is not None
+    assert self._reward_range is not None
+    assert self._trajectories is not None
+
+  @property
+  def base_env_name(self):
+    return self._base_env_name
+
+  @property
+  def trajectories(self):
+    return self._trajectories
+
+  def _verify_same_spaces(self):
+    """Verifies that all the envs have the same observation and action space."""
+
+    # Pre-conditions: self._envs is initialized.
+
+    if self._envs is None:
+      raise ValueError("Environments not initialized.")
+
+    if not isinstance(self._envs, list):
+      tf.logging.warning("Not checking observation and action space "
+                         "compatibility across envs, since there is just one.")
+      return
+
+    # NOTE: We compare string representations of observation_space and
+    # action_space because compositional classes like space.Tuple don't return
+    # true on object comparison.
+
+    if not all(
+        str(env.observation_space) == str(self.observation_space)
+        for env in self._envs):
+      err_str = ("All environments should have the same observation space, but "
+                 "don't.")
+      tf.logging.error(err_str)
+      # Log all observation spaces.
+      for i, env in enumerate(self._envs):
+        tf.logging.error("Env[%d] has observation space [%s]", i,
+                         env.observation_space)
+      raise ValueError(err_str)
+
+    if not all(
+        str(env.action_space) == str(self.action_space) for env in self._envs):
+      err_str = "All environments should have the same action space, but don't."
+      tf.logging.error(err_str)
+      # Log all action spaces.
+      for i, env in enumerate(self._envs):
+        tf.logging.error("Env[%d] has action space [%s]", i, env.action_space)
+      raise ValueError(err_str)
+
+  def initialize_environments(self, batch_size=1):
+    """Initializes the environments and trajectories.
+
+    Subclasses can override this if they don't want a default implementation
+    which initializes `batch_size` environments, but must take care to
+    initialize self._trajectories (this is checked in __init__ anyways).
+
+    Args:
+      batch_size: (int) Number of `self.base_env_name` envs to initialize.
+    """
+
+    assert batch_size >= 1
+
+    max_steps = self._base_env_kwargs.get("rl_env_max_episode_steps", -1)
+    maxskip_env = self._base_env_kwargs.get("maxskip_env", False)
+
+    self._envs = []
+    for _ in range(batch_size):
+      self._envs.append(
+          gym_utils.make_gym_env(
+              self.base_env_name,
+              rl_env_max_episode_steps=max_steps,
+              maxskip_env=maxskip_env))
+
+    # If self.observation_space and self.action_space aren't None, then it means
+    # that this is a re-initialization of this class, in that case make sure
+    # that this matches our previous behaviour.
+    if self._observation_space:
+      assert str(self._observation_space) == str(
+          self._envs[0].observation_space)
+    else:
+      # This means that we are initializing this class for the first time.
+      #
+      # We set this equal to the first env's observation space, later on we'll
+      # verify that all envs have the same observation space.
+      self._observation_space = self._envs[0].observation_space
+
+    # Similarly for action_space
+    if self._action_space:
+      assert str(self._action_space) == str(self._envs[0].action_space)
+    else:
+      self._action_space = self._envs[0].action_space
+
+    self._verify_same_spaces()
+
+    # If self.reward_range is None, i.e. this means that we should take the
+    # reward range of the env.
+    if self.reward_range is None:
+      self._reward_range = self._envs[0].reward_range
+
+    # This data structure stores the history of each env.
+    #
+    # NOTE: Even if the env is a NN and can step in all batches concurrently, it
+    # is still valuable to store the trajectories separately.
+    self._trajectories = trajectory.BatchTrajectory(batch_size=batch_size)
+
+  def assert_common_preconditions(self):
+    # Asserts on the common pre-conditions of:
+    #  - self._envs is initialized.
+    #  - self._envs is a list.
+    assert self._envs
+    assert isinstance(self._envs, list)
+
+  @property
+  def observation_space(self):
+    return self._observation_space
+
+  @property
+  def observation_spec(self):
+    """The spec for reading an observation stored in a tf.Example."""
+    return gym_spaces_utils.gym_space_spec(self.observation_space)
+
+  def process_observations(self, observations):
+    """Processes observations prior to saving in the trajectories.
+
+    Args:
+      observations: (np.ndarray) observations to be processed.
+
+    Returns:
+      processed observation
+
+    """
+    return observations
+
+  @property
+  def action_space(self):
+    return self._action_space
+
+  @property
+  def action_spec(self):
+    """The spec for reading an observation stored in a tf.Example."""
+    return gym_spaces_utils.gym_space_spec(self.action_space)
+
+  @property
+  def action_modality(self):
+    raise NotImplementedError
+
+  @property
+  def num_actions(self):
+    """Returns the number of actions in a discrete action space."""
+    return gym_spaces_utils.cardinality(self.action_space)
+
+  @property
+  def reward_range(self):
+    return self._reward_range
+
+  @property
+  def is_reward_range_finite(self):
+    min_reward, max_reward = self.reward_range
+    return (min_reward != -np.inf) and (max_reward != np.inf)
+
+  def process_rewards(self, rewards):
+    """Clips, rounds, adds the min_reward and changes to integer type.
+
+    Args:
+      rewards: numpy array of raw (float) rewards.
+
+    Returns:
+      processed_rewards: numpy array of np.int64
+    """
+
+    min_reward, max_reward = self.reward_range
+
+    # Clips at min and max reward and shift by min (so new min is 0)
+    rewards = np.clip(rewards, min_reward, max_reward) - min_reward
+    # Round to (nearest) int and convert to integral type.
+    rewards = np.around(rewards, decimals=0).astype(np.int64)
+    return rewards
+
+  @property
+  def is_processed_rewards_discrete(self):
+    """Returns true if `self.process_rewards` returns discrete rewards."""
+
+    # Subclasses can override, but it should match their self.process_rewards.
+
+    # This check is a little hackily.
+    return self.process_rewards(0.0).dtype == np.int64
+
+  @property
+  def num_rewards(self):
+    """Returns the number of distinct rewards.
+
+    Returns:
+      Returns None if the reward range is infinite or the processed rewards
+      aren't discrete, otherwise returns the number of distinct rewards.
+    """
+
+    # Pre-conditions: reward range is finite.
+    #               : processed rewards are discrete.
+    if not self.is_reward_range_finite:
+      tf.logging.error("Infinite reward range, `num_rewards returning None`")
+      return None
+    if not self.is_processed_rewards_discrete:
+      tf.logging.error(
+          "Processed rewards are not discrete, `num_rewards` returning None")
+      return None
+
+    min_reward, max_reward = self.reward_range
+    return max_reward - min_reward + 1
+
+  @property
+  def input_modality(self):
+    raise NotImplementedError
+
+  @property
+  def input_vocab_size(self):
+    raise NotImplementedError
+
+  @property
+  def target_modality(self):
+    raise NotImplementedError
+
+  @property
+  def target_vocab_size(self):
+    raise NotImplementedError
+
+  @property
+  def unwrapped(self):
+    return self
+
+  def seed(self, seed=None):
+    if not self._envs:
+      tf.logging.info("`seed` called on non-existent envs, doing nothing.")
+      return
+
+    if not isinstance(self._envs, list):
+      tf.logging.warning("`seed` called on non-list envs, doing nothing.")
+      return
+
+    tf.logging.warning(
+        "Called `seed` on EnvProblem, calling seed on the underlying envs.")
+    for env in self._envs:
+      env.seed(seed)
+
+  def close(self):
+    if not self._envs:
+      tf.logging.info("`close` called on non-existent envs, doing nothing.")
+      return
+
+    if not isinstance(self._envs, list):
+      tf.logging.warning("`close` called on non-list envs, doing nothing.")
+      return
+
+    # Call close on all the envs one by one.
+    for env in self._envs:
+      env.close()
+
+  def _reset(self, indices):
+    """Resets environments at indices shouldn't pre-process or record.
+
+    Subclasses should override this to do the actual reset if something other
+    than the default implementation is desired.
+
+    Args:
+      indices: list of indices of underlying envs to call reset on.
+
+    Returns:
+      np.ndarray of stacked observations from the reset-ed envs.
+    """
+
+    # Pre-conditions: common_preconditions, see `assert_common_preconditions`.
+    self.assert_common_preconditions()
+
+    # This returns a numpy array with first dimension `len(indices)` and the
+    # rest being the dimensionality of the observation.
+    return np.stack([self._envs[index].reset() for index in indices])
+
+  def reset(self, indices=None):
+    """Resets environments at given indices.
+
+    Subclasses should override _reset to do the actual reset if something other
+    than the default implementation is desired.
+
+    Args:
+      indices: Indices of environments to reset. If None all envs are reset.
+
+    Returns:
+      Batch of initial observations of reset environments.
+    """
+
+    if indices is None:
+      indices = np.arange(self.trajectories.batch_size)
+
+    # If this is empty (not None) then don't do anything, no env was done.
+    if indices.size == 0:
+      tf.logging.warning(
+          "`reset` called with empty indices array, this is a no-op.")
+      return None
+
+    observations = self._reset(indices)
+    processed_observations = self.process_observations(observations)
+
+    # Record history.
+    self.trajectories.reset(indices, observations)
+
+    return processed_observations
+
+  def _step(self, actions):
+    """Takes a step in all environments, shouldn't pre-process or record.
+
+    Subclasses should override this to do the actual step if something other
+    than the default implementation is desired.
+
+    Args:
+      actions: (np.ndarray) with first dimension equal to the batch size.
+
+    Returns:
+      a tuple of stacked raw observations, raw rewards, dones and infos.
+    """
+
+    # Pre-conditions: common_preconditions, see `assert_common_preconditions`.
+    #               : len(actions) == len(self._envs)
+    self.assert_common_preconditions()
+    assert len(actions) == len(self._envs)
+
+    observations = []
+    rewards = []
+    dones = []
+    infos = []
+
+    # Take steps in all environments.
+    for env, action in zip(self._envs, actions):
+      observation, reward, done, info = env.step(action)
+
+      observations.append(observation)
+      rewards.append(reward)
+      dones.append(done)
+      infos.append(info)
+
+    # Convert each list (observations, rewards, ...) into np.array and return a
+    # tuple.
+    return tuple(map(np.stack, [observations, rewards, dones, infos]))
+
+  def step(self, actions):
+    """Takes a step in all environments.
+
+    Subclasses should override _step to do the actual reset if something other
+    than the default implementation is desired.
+
+    Args:
+      actions: Batch of actions.
+
+    Returns:
+      (preprocessed_observations, processed_rewards, dones, infos).
+    """
+
+    observations, raw_rewards, dones, infos = self._step(actions)
+
+    # Process rewards.
+    raw_rewards = raw_rewards.astype(np.float32)
+    processed_rewards = self.process_rewards(raw_rewards)
+
+    # Process observations.
+    processed_observations = self.process_observations(observations)
+
+    # Record history.
+    self.trajectories.step(processed_observations, raw_rewards,
+                           processed_rewards, dones, actions)
+
+    return processed_observations, processed_rewards, dones, infos
+
+  @staticmethod
+  def done_indices(dones):
+    """Calculates the indices where dones has True."""
+    return np.argwhere(dones).squeeze(axis=1)
+
+  def example_reading_spec(self):
+    """Data fields to store on disk and their decoders."""
+
+    # Subclasses can override and/or extend.
+
+    processed_reward_type = tf.float32
+    if self.is_processed_rewards_discrete:
+      processed_reward_type = tf.int64
+
+    data_fields = {
+        TIMESTEP_FIELD: tf.FixedLenFeature((1,), tf.int64),
+        RAW_REWARD_FIELD: tf.FixedLenFeature((1,), tf.float32),
+        PROCESSED_REWARD_FIELD: tf.FixedLenFeature((1,), processed_reward_type),
+        DONE_FIELD: tf.FixedLenFeature((1,), tf.int64),  # we wrote this as int.
+
+        # Special treatment because we need to determine type and shape, also
+        # enables classes to override.
+        OBSERVATION_FIELD: self.observation_spec,
+        ACTION_FIELD: self.action_spec,
+    }
+
+    # `data_items_to_decoders` can be None, it will be set to the appropriate
+    # decoder dict in `Problem.decode_example`
+    # TODO(afrozm): Verify that we don't need any special decoder or anything.
+    data_items_to_decoders = None
+
+    return data_fields, data_items_to_decoders
+
+  def hparams(self, defaults, model_hparams):
+    # Usually when using the environment in a supervised setting, given the
+    # observation we are predicting the reward.
+    p = defaults
+
+    # Have to add these the 'proper' way, otherwise __str__ doesn't show them.
+    if "modality" not in p:
+      p.add_hparam("modality", {})
+    if "vocab_size" not in p:
+      p.add_hparam("vocab_size", {})
+
+    # TODO(afrozm): Document what all of these keys are and are supposed to do.
+    p.modality.update({
+        "inputs": self.input_modality,
+        "targets": self.target_modality,
+        "input_reward": modalities.SymbolModalityWeightsAll,
+        "target_reward": modalities.SymbolModalityWeightsAll,
+        "input_action": modalities.SymbolModalityWeightsAll,
+        "target_action": modalities.SymbolModalityWeightsAll,
+        "target_policy": modalities.IdentityModality,
+        "target_value": modalities.IdentityModality,
+    })
+
+    p.vocab_size.update({
+        "inputs": self.input_vocab_size,
+        "targets": self.target_vocab_size,
+        "input_reward": self.num_rewards,
+        "target_reward": self.num_rewards,
+        "input_action": self.num_actions,
+        "target_action": self.num_actions,
+        "target_policy": None,
+        "target_value": None,
+    })
+
+    p.input_space_id = problem.SpaceID.GENERIC
+    p.target_space_id = problem.SpaceID.GENERIC
+
+  @property
+  def agent_id(self):
+    return self._agent_id
+
+  @agent_id.setter
+  def agent_id(self, agent_id):
+    # We use `-` in self.dataset_filename, disallow it here for convenience.
+    if "-" in agent_id:
+      raise ValueError("agent_id shouldn't have - in it.")
+    self._agent_id = agent_id
+
+  def dataset_filename(self):
+    return "{}-{}".format(self.name, self.agent_id)
+
+  @property
+  def num_shards(self):
+    return {
+        problem.DatasetSplit.TRAIN: 10,
+        problem.DatasetSplit.EVAL: 1,
+    }
+
+  def _generate_time_steps(self, trajectory_list):
+    """A generator to yield single time-steps from a list of trajectories."""
+    for single_trajectory in trajectory_list:
+      assert isinstance(single_trajectory, trajectory.Trajectory)
+
+      # Skip writing trajectories that have only a single time-step -- this
+      # could just be a repeated reset.
+
+      if single_trajectory.num_time_steps() <= 1:
+        continue
+
+      for index, time_step in enumerate(single_trajectory.time_steps):
+
+        # The first time-step doesn't have reward/processed_reward, if so, just
+        # setting it to 0.0 / 0 should be OK.
+        raw_reward = time_step.raw_reward
+        if not raw_reward:
+          raw_reward = 0.0
+
+        processed_reward = time_step.processed_reward
+        if not processed_reward:
+          processed_reward = 0
+
+        if time_step.action:
+          action = gym_spaces_utils.gym_space_encode(self.action_space,
+                                                     time_step.action)
+        else:
+          # The last time-step doesn't have action, and this action shouldn't be
+          # used, gym's spaces have a `sample` function, so let's just sample an
+          # action and use that.
+          action = [self.action_space.sample()]
+
+        if six.PY3:
+          # py3 complains that, to_example cannot handle np.int64 !
+
+          action_dtype = self.action_space.dtype
+          if action_dtype in [np.int64, np.int32]:
+            action = list(map(int, action))
+          elif action_dtype in [np.float64, np.float32]:
+            action = list(map(float, action))
+
+          # same with processed_reward.
+          processed_reward = int(processed_reward)
+
+        assert time_step.observation is not None
+
+        yield {
+            TIMESTEP_FIELD: [index],
+            ACTION_FIELD:
+                action,
+            RAW_REWARD_FIELD:
+                [float(raw_reward)],  # to_example errors on np.float32
+            PROCESSED_REWARD_FIELD: [processed_reward],
+            DONE_FIELD: [int(time_step.done)],  # to_example doesn't know bools
+            OBSERVATION_FIELD:
+                gym_spaces_utils.gym_space_encode(self.observation_space,
+                                                  time_step.observation),
+        }
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    # List of files to generate data in.
+    # NOTE: We don't want to shuffle, so we mark the files as shuffled.
+    files_list = []
+    for split, num_shards in self.num_shards.items():
+      files_list.extend(self.data_filepaths(split, data_dir, num_shards, True))
+
+    # At this point some trajectories haven't finished. However we still want to
+    # write those down.
+
+    # A simple way of doing this is to call `self.reset()` here, this will make
+    # all the envs take one (extra) step, but would be a clean way to do it.
+    #
+    # self.reset()
+
+    self.trajectories.complete_all_trajectories()
+
+    # Write the completed data into these files
+
+    num_completed_trajectories = len(self.trajectories.completed_trajectories)
+    num_shards = len(files_list)
+    if num_completed_trajectories < num_shards:
+      tf.logging.warning(
+          "Number of completed trajectories [%d] is less than "
+          "the number of shards [%d], some shards maybe empty.",
+          num_completed_trajectories, num_shards)
+
+    for i, f in enumerate(files_list[:num_completed_trajectories]):
+      # Start at index i of completed trajectories and take every `num_shards`
+      # trajectory. This ensures that the data is approximately a balanced
+      # partition of completed trajectories, also because of the above slicing
+      # of files_list, i will be a valid index into completed_trajectories.
+      trajectories_to_write = self.trajectories.completed_trajectories[
+          i::num_shards]
+
+      # Convert each trajectory from `trajectories_to_write` to a sequence of
+      # time-steps and then send that generator to `generate_files`.
+
+      # `cycle_every_n` isn't needed since file list given to it is a singleton.
+      generator_utils.generate_files(
+          self._generate_time_steps(trajectories_to_write), [f])
diff --git a/tensor2tensor/envs/env_problem_test.py b/tensor2tensor/envs/env_problem_test.py
new file mode 100644
index 000000000..d83b7575b
--- /dev/null
+++ b/tensor2tensor/envs/env_problem_test.py
@@ -0,0 +1,244 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.envs.env_problem."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+from gym.spaces import Box
+from gym.spaces import Discrete
+import numpy as np
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.envs import env_problem
+import tensorflow as tf
+
+
+class EnvProblemTest(tf.test.TestCase):
+
+  def setUp(self):
+    self.tmp_dir = os.path.join(tf.test.get_temp_dir(), "tmp_dir")
+    tf.gfile.MakeDirs(self.tmp_dir)
+
+  def tearDown(self):
+    tf.gfile.DeleteRecursively(self.tmp_dir)
+
+  def test_setup(self):
+    ep = env_problem.EnvProblem(base_env_name="CartPole-v0", batch_size=5)
+    # Checks that environments were created and they are `batch_size` in number.
+    ep.assert_common_preconditions()
+
+    # Expectations on the observation space.
+    observation_space = ep.observation_space
+    self.assertTrue(isinstance(observation_space, Box))
+    self.assertEqual(observation_space.shape, (4,))
+    self.assertEqual(observation_space.dtype, np.float32)
+
+    # Expectations on the action space.
+    action_space = ep.action_space
+    self.assertTrue(isinstance(action_space, Discrete))
+    self.assertEqual(action_space.shape, ())
+    self.assertEqual(action_space.dtype, np.int64)
+    self.assertEqual(ep.num_actions, 2)
+
+    # Reward range is infinite here.
+    self.assertFalse(ep.is_reward_range_finite)
+
+  def test_reward_range(self):
+    # Passing reward_range=None means take the reward range of the underlying
+    # environment as the reward range.
+    ep = env_problem.EnvProblem(
+        base_env_name="FrozenLake-v0", batch_size=5, reward_range=None)
+    ep.assert_common_preconditions()
+
+    # Assert reward range is finite here.
+    self.assertTrue(ep.is_reward_range_finite)
+
+    # Assert that it is as expected of the underlying environment, since reward_
+    self.assertEqual(0, ep.reward_range[0])
+    self.assertEqual(1, ep.reward_range[1])
+
+  def test_default_processed_rewards_discrete(self):
+    # This differs in the above because it has a Tuple observation space.
+    ep = env_problem.EnvProblem(
+        base_env_name="KellyCoinflip-v0", batch_size=5, reward_range=None)
+    ep.assert_common_preconditions()
+
+    # Assert reward range is finite here.
+    self.assertTrue(ep.is_reward_range_finite)
+
+    # Assert that it is as expected of the underlying environment.
+    reward_range = ep.reward_range
+    self.assertEqual(0, reward_range[0])
+    self.assertEqual(ep._envs[0].maxWealth, reward_range[1])
+
+    # Check that the processed rewards are discrete.
+    self.assertTrue(ep.is_processed_rewards_discrete)
+
+    # Assert on the number of rewards.
+    self.assertEqual(ep.num_rewards, reward_range[1] - reward_range[0] + 1)
+
+  def test_interaction_with_env(self):
+    batch_size = 5
+    reward_range = (-1, 1)
+    ep = env_problem.EnvProblem(
+        base_env_name="KellyCoinflip-v0",
+        batch_size=batch_size,
+        reward_range=reward_range)
+    ep.agent_id = "test"
+
+    # Resets all environments.
+    ep.reset()
+
+    # Let's play a few steps.
+    nsteps = 100
+    num_trajectories_completed = 0
+    num_timesteps_completed = 0
+    # If batch_done_at_step[i] = j then it means that i^th env last got done at
+    # step = j.
+    batch_done_at_step = np.full(batch_size, -1)
+    for i in range(nsteps):
+      # Sample batch_size actions from the action space and stack them (since
+      # that is the expected type).
+      actions = np.stack([ep.action_space.sample() for _ in range(batch_size)])
+
+      _, _, dones, _ = ep.step(actions)
+
+      # Do the book-keeping on number of trajectories completed and expect that
+      # it matches ep's completed number.
+
+      num_done = sum(dones)
+      num_trajectories_completed += num_done
+
+      self.assertEqual(num_trajectories_completed,
+                       len(ep.trajectories.completed_trajectories))
+
+      # Get the indices where we are done ...
+      done_indices = env_problem.EnvProblem.done_indices(dones)
+
+      # ... and reset those.
+      ep.reset(indices=done_indices)
+
+      # If nothing got done, go on to the next step.
+      if done_indices.size == 0:
+        # i.e. this is an empty array.
+        continue
+
+      # See when these indices were last done and calculate how many time-steps
+      # each one took to get done.
+      num_timesteps_completed += sum(i + 1 - batch_done_at_step[done_indices])
+      batch_done_at_step[done_indices] = i
+
+      # This should also match the number of time-steps completed given by ep.
+      num_timesteps_completed_ep = sum(
+          ct.num_time_steps() for ct in ep.trajectories.completed_trajectories)
+      self.assertEqual(num_timesteps_completed, num_timesteps_completed_ep)
+
+    # Reset the trajectories.
+    ep.trajectories.reset_batch_trajectories()
+    self.assertEqual(0, len(ep.trajectories.completed_trajectories))
+
+  def read_tfrecord_dataset(self, filenames, ep):
+    # Read the dataset at `filenames` into a tf.data.Dataset and returns the
+    # number of time-steps (just the number of records in the dataset) and the
+    # number of trajectories.
+
+    last_timestep = -1
+    num_time_steps = 0
+    num_trajectories = 0
+    for ex in generator_utils.tfrecord_iterator(
+        filenames, example_spec=ep.example_reading_spec()[0]):
+      num_time_steps += 1
+      this_timestep = ex[env_problem.TIMESTEP_FIELD][0]
+      if 1 + last_timestep != this_timestep:
+        num_trajectories += 1
+        self.assertEqual(0, this_timestep)
+      last_timestep = this_timestep
+    num_trajectories += 1
+
+    return num_trajectories, num_time_steps
+
+  def test_generate_data(self):
+    batch_size = 5
+    reward_range = (-1, 1)
+    ep = env_problem.EnvProblem(
+        base_env_name="CartPole-v0",
+        batch_size=batch_size,
+        reward_range=reward_range)
+    ep.agent_id = "test"
+
+    # Set this in the test to test things, but usually registered subclasses
+    # will set this.
+    ep.name = "CartPoleProblem"
+
+    # Reset all environments.
+    ep.reset()
+
+    # Play for some steps to generate data.
+    nsteps = 100
+    num_dones = 0
+    num_dones_in_last_step = 0
+    for _ in range(nsteps):
+      # Sample actions.
+      actions = np.stack([ep.action_space.sample() for _ in range(batch_size)])
+      # Step through it.
+      _, _, dones, _ = ep.step(actions)
+      # Get the indices where we are done ...
+      done_indices = env_problem.EnvProblem.done_indices(dones)
+      # ... and reset those.
+      ep.reset(indices=done_indices)
+      # count the number of dones we got, in this step and overall.
+      num_dones_in_last_step = sum(dones)
+      num_dones += num_dones_in_last_step
+
+    # This is because every num_dones starts a new trajectory, and a further
+    # batch_size are active at the last step when we call generate_data, but
+    # the ones that got done in the last step (these have only one time-step in
+    # their trajectory) will be skipped.
+    expected_num_trajectories = num_dones + batch_size - num_dones_in_last_step
+
+    # Similar logic as above, nsteps * batch_size overall `step` calls are made.
+    # However, if a `step` call completes a trajectory, one more time-step is
+    # added, but we have to discount the trajectories that got done in the very
+    # last step.
+    expected_num_time_steps = (
+        nsteps * batch_size) + num_dones + batch_size - num_dones_in_last_step
+
+    # Dump the completed data to disk.
+    ep.generate_data(self.tmp_dir, self.tmp_dir)
+
+    # Read the written files and assert on the number of time steps.
+    training_filenames = ep.training_filepaths(
+        self.tmp_dir, ep.num_shards[problem.DatasetSplit.TRAIN], True)
+    dev_filenames = ep.dev_filepaths(
+        self.tmp_dir, ep.num_shards[problem.DatasetSplit.EVAL], True)
+
+    training_trajectories, training_timesteps = self.read_tfrecord_dataset(
+        training_filenames, ep)
+    dev_trajectories, dev_timesteps = self.read_tfrecord_dataset(
+        dev_filenames, ep)
+
+    # This tests what we wrote on disk matches with what we computed.
+    self.assertEqual(expected_num_time_steps,
+                     training_timesteps + dev_timesteps)
+    self.assertEqual(expected_num_trajectories,
+                     training_trajectories + dev_trajectories)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/envs/gym_spaces_utils.py b/tensor2tensor/envs/gym_spaces_utils.py
new file mode 100644
index 000000000..53eef08ba
--- /dev/null
+++ b/tensor2tensor/envs/gym_spaces_utils.py
@@ -0,0 +1,113 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Few utility functions to deal with gym spaces.
+
+gym.spaces.Box and gym.spaces.Discrete are easiest to support.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from gym.spaces import Box
+from gym.spaces import Discrete
+
+import numpy as np
+import tensorflow as tf
+
+
+def box_space_spec(box_space, tf_dtype):
+  return tf.FixedLenFeature(box_space.shape, tf_dtype)
+
+
+def discrete_space_spec(discrete_space, tf_dtype):
+  del discrete_space  # this is not needed.
+  return tf.FixedLenFeature((1,), tf_dtype)
+
+
+def gym_space_spec(gym_space):
+  """Returns a reading spec of a gym space.
+
+  NOTE: Only implemented currently for Box and Discrete.
+
+  Args:
+    gym_space: instance of gym.spaces whose spec we want.
+
+  Returns:
+    Reading spec for that space.
+
+  Raises:
+    NotImplementedError: For spaces whose reading spec we haven't implemented.
+  """
+  # First try to determine the type.
+  try:
+    tf_dtype = tf.as_dtype(gym_space.dtype)
+  except TypeError as e:
+    tf.logging.error("Cannot convert space's type [%s] to tf.dtype",
+                     gym_space.dtype)
+    raise e
+
+  # Now hand it over to the specialized functions.
+  if isinstance(gym_space, Box):
+    return box_space_spec(gym_space, tf_dtype)
+  elif isinstance(gym_space, Discrete):
+    return discrete_space_spec(gym_space, tf_dtype)
+  else:
+    raise NotImplementedError
+
+
+def gym_space_encode(gym_space, observation):
+  # We should return something that generator_utils.to_example can consume.
+  if isinstance(gym_space, Discrete):
+    return [observation]
+
+  if isinstance(gym_space, Box):
+    return list(observation.reshape(-1))
+
+  raise NotImplementedError
+
+
+def cardinality(gym_space):
+  """Number of elements that can be represented by the space.
+
+  Makes the most sense for Discrete or Box type with integral dtype, ex: number
+  of actions in an action space.
+
+  Args:
+    gym_space: The gym space.
+
+  Returns:
+    np.int64 number of observations that can be represented by this space, or
+    returns None when this doesn't make sense, i.e. float boxes etc.
+
+  Raises:
+    NotImplementedError when a space's cardinality makes sense but we haven't
+    implemented it.
+  """
+
+  if (gym_space.dtype == np.float32) or (gym_space.dtype == np.float64):
+    tf.logging.error("Returning None for a float gym space's cardinality: ",
+                     gym_space)
+    return None
+
+  if isinstance(gym_space, Discrete):
+    return gym_space.n
+
+  if isinstance(gym_space, Box):
+    # Construct a box with all possible values in this box and take a product.
+    return np.prod(gym_space.high - gym_space.low + 1)
+
+  raise NotImplementedError
diff --git a/tensor2tensor/envs/time_step.py b/tensor2tensor/envs/time_step.py
new file mode 100644
index 000000000..2e4d8a45c
--- /dev/null
+++ b/tensor2tensor/envs/time_step.py
@@ -0,0 +1,67 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""TimeStep is a simple class that holds the information seen at a time-step.
+
+Let:
+r_t = Reward(s_{t-1}, a_{t-1}, s_t)  - reward for getting into a state.
+d_t = Done(s_t)                      - is this state terminal.
+
+Then the sequence of states, actions and rewards looks like the following:
+
+s0, a0 s1/r1/d1, a1 s2/r2/d2, a2 s3/r3/d3, ...
+
+TimeStep holds (s_t, d_t, r_t, a_t).
+
+NOTE: When we call step on an environment at time-step t, we supply a_t and in
+return the env gives us s_{t+1}, d_{t+1}, r_{t+1}
+
+So, we'd have to add the actions a_t to the current time-step, but add the
+observations, rewards and dones to a new time-step.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+
+
+class TimeStep(
+    collections.namedtuple(
+        "TimeStep",
+        ["observation", "done", "raw_reward", "processed_reward", "action"])):
+  """This class represents the time-step as mentioned above."""
+
+  def replace(self, **kwargs):
+    """Exposes the underlying namedtuple replace."""
+
+    # NOTE: This RETURNS a NEW time-step with the replacements, i.e. doesn't
+    # modify self, since namedtuple is immutable.
+
+    # This allows this to be called like ts.replace(action=a, raw_reward=r) etc.
+
+    return self._replace(**kwargs)
+
+  @classmethod
+  def create_time_step(cls,
+                       observation=None,
+                       done=False,
+                       raw_reward=None,
+                       processed_reward=None,
+                       action=None):
+    """Creates a TimeStep with both rewards and actions as optional."""
+
+    return cls(observation, done, raw_reward, processed_reward, action)
diff --git a/tensor2tensor/envs/time_step_test.py b/tensor2tensor/envs/time_step_test.py
new file mode 100644
index 000000000..410ffb546
--- /dev/null
+++ b/tensor2tensor/envs/time_step_test.py
@@ -0,0 +1,57 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.envs.time_step."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.envs import time_step
+
+import tensorflow as tf
+
+
+class TimeStepTest(tf.test.TestCase):
+
+  def test_create_time_step(self):
+    ts = time_step.TimeStep.create_time_step(
+        observation=1, done=True, raw_reward=1.0, processed_reward=1, action=1)
+
+    self.assertEqual(1, ts.observation)
+    self.assertTrue(ts.done)
+    self.assertNear(1.0, ts.raw_reward, 1e-6)
+    self.assertEqual(1, ts.processed_reward)
+    self.assertEqual(1, ts.action)
+
+  def test_replace(self):
+    ts = time_step.TimeStep.create_time_step(observation=1, action=1)
+    self.assertFalse(ts.done)
+
+    tsr = ts.replace(action=2, done=True)
+
+    # Asert that ts didn't change.
+    self.assertFalse(ts.done)
+    self.assertEqual(1, ts.observation)
+    self.assertEqual(1, ts.action)
+
+    # But tsr is as expected.
+    self.assertTrue(tsr.done)
+    self.assertEqual(1, tsr.observation)  # unchanged
+    self.assertEqual(2, tsr.action)  # changed
+
+
+if __name__ == '__main__':
+  tf.test.main()
diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
new file mode 100644
index 000000000..4a1ff12a2
--- /dev/null
+++ b/tensor2tensor/envs/trajectory.py
@@ -0,0 +1,243 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Trajectory manages a sequence of TimeSteps.
+
+BatchTrajectory manages a batch of trajectories, also keeping account of
+completed trajectories.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.envs import time_step
+
+
+class Trajectory(object):
+  """Basically a list of TimeSteps with convenience methods."""
+
+  def __init__(self):
+    # Contains a list of time steps.
+    self._time_steps = []
+
+  def __str__(self):
+    if not self.time_steps:
+      return "Trajectory[]"
+    return "Trajectory[{}]".format(", ".join(str(ts) for ts in self.time_steps))
+
+  def add_time_step(self, **create_time_step_kwargs):
+    """Creates a time-step and appends it to the list.
+
+    Args:
+      **create_time_step_kwargs: Forwarded to
+                                 time_step.TimeStep.create_time_step.
+    """
+    ts = time_step.TimeStep.create_time_step(**create_time_step_kwargs)
+    assert isinstance(ts, time_step.TimeStep)
+    self._time_steps.append(ts)
+
+  def change_last_time_step(self, **replace_time_step_kwargs):
+    """Replace the last time-steps with the given kwargs."""
+
+    # Pre-conditions: self._time_steps shouldn't be empty.
+    assert self._time_steps
+    self._time_steps[-1] = self._time_steps[-1].replace(
+        **replace_time_step_kwargs)
+
+  def last_time_step(self):
+    # Pre-conditions: self._time_steps shouldn't be empty.
+    assert self._time_steps
+    return self._time_steps[-1]
+
+  # We could have overridden __nonzero__ or __bool__ as well.
+  def num_time_steps(self):
+    return len(self._time_steps)
+
+  def is_active(self):
+    return bool(self.num_time_steps())
+
+  @property
+  def time_steps(self):
+    return self._time_steps
+
+
+class BatchTrajectory(object):
+  """Basically a batch of active trajectories and a list of completed ones."""
+
+  def __init__(self, batch_size=1):
+    self.batch_size = batch_size
+
+    # Stores trajectories that are currently active, i.e. aren't done or reset.
+    self._trajectories = [Trajectory() for _ in range(self.batch_size)]
+
+    # Stores trajectories that are completed.
+    # NOTE: We don't track the index this came from, as it's not needed, right?
+    self._completed_trajectories = []
+
+  def reset_batch_trajectories(self):
+    self.__init__(batch_size=self.batch_size)
+
+  def __str__(self):
+    string = "BatchTrajectory["
+    for i, t in enumerate(self.trajectories):
+      string += "Trajectory {} = {}\n".format(i, str(t))
+    for i, t in enumerate(self.completed_trajectories):
+      string += "Completed Trajectory {} = {}\n".format(i, str(t))
+    return string + "]"
+
+  @property
+  def trajectories(self):
+    return self._trajectories
+
+  @property
+  def completed_trajectories(self):
+    return self._completed_trajectories
+
+  def _complete_trajectory(self, trajectory, index):
+    """Completes the given trajectory at the given index."""
+
+    assert isinstance(trajectory, Trajectory)
+
+    # This *should* be the case.
+    assert trajectory.last_time_step().action is None
+
+    # Add to completed trajectories.
+    self._completed_trajectories.append(trajectory)
+
+    # Make a new one to replace it.
+    self._trajectories[index] = Trajectory()
+
+  def reset(self, indices, observations):
+    """Resets trajectories at given indices and populates observations.
+
+    Reset can either be called right at the beginning, when there are no
+    time-steps, or to reset a currently active trajectory.
+
+    If resetting a currently active trajectory then we save it in
+    self._completed_trajectories.
+
+    Args:
+      indices: 1-D np.ndarray stating the indices to reset.
+      observations: np.ndarray of shape (indices len, obs.shape) of observations
+    """
+
+    # Pre-conditions: indices, observations are np arrays.
+    #               : indices is one-dimensional.
+    #               : their first dimension (batch) is the same.
+    assert isinstance(indices, np.ndarray)
+    assert len(indices.shape) == 1
+    assert isinstance(observations, np.ndarray)
+    assert indices.shape[0] == observations.shape[0]
+
+    for index, observation in zip(indices, observations):
+      trajectory = self._trajectories[index]
+
+      # Are we starting a new trajectory at the given index?
+      if not trajectory.is_active():
+        # Then create a new time-step here with the given observation.
+        trajectory.add_time_step(observation=observation)
+        # That's all we need to do here.
+        continue
+
+      # If however we are resetting a currently active trajectory then we need
+      # to put that in self._completed_trajectories and make a new trajectory
+      # with the current observation.
+
+      # TODO(afrozm): Should we mark these are done? Or is the done=False and
+      # this being the last time-step in the trajectory good enough to recognize
+      # that this was reset?
+
+      # Mark trajectory as completed and move into completed_trajectories.
+      self._complete_trajectory(trajectory, index)
+
+      # Put the observation in the newly created trajectory.
+      # TODO(afrozm): Add 0 reward.
+      self._trajectories[index].add_time_step(observation=observation)
+
+  def complete_all_trajectories(self):
+    """Essentially same as reset, but we don't have observations."""
+    for index in range(self.batch_size):
+      trajectory = self._trajectories[index]
+      assert trajectory.is_active()
+      self._complete_trajectory(trajectory, index)
+
+  def step(self, observations, raw_rewards, processed_rewards, dones, actions):
+    """Record the information obtained from taking a step in all envs.
+
+    Records (observation, rewards, done) in a new time-step and actions in the
+    current time-step.
+
+    If any trajectory gets done, we move that trajectory to
+    completed_trajectories.
+
+    Args:
+      observations: ndarray of first dimension self.batch_size, which has the
+        observations after we've stepped, i.e. s_{t+1} where t is the current
+        state.
+      raw_rewards: ndarray of first dimension self.batch_size containing raw
+        rewards i.e. r_{t+1}.
+      processed_rewards: ndarray of first dimension self.batch_size containing
+        processed rewards. i.e. r_{t+1}
+      dones: ndarray of first dimension self.batch_size, containing true at an
+        index if that env is done, i.e. d_{t+1}
+      actions: ndarray of first dimension self.batch_size, containing actions
+        applied at the current time-step, which leads to the observations
+        rewards and done at the next time-step, i.e. a_t
+    """
+    # Pre-conditions
+    assert isinstance(observations, np.ndarray)
+    assert isinstance(raw_rewards, np.ndarray)
+    assert isinstance(processed_rewards, np.ndarray)
+    assert isinstance(dones, np.ndarray)
+    assert isinstance(actions, np.ndarray)
+
+    # We assume that we step in all envs, i.e. not like reset where we can reset
+    # some envs and not others.
+    assert self.batch_size == observations.shape[0]
+    assert self.batch_size == raw_rewards.shape[0]
+    assert self.batch_size == processed_rewards.shape[0]
+    assert self.batch_size == dones.shape[0]
+    assert self.batch_size == actions.shape[0]
+
+    for index in range(self.batch_size):
+      trajectory = self._trajectories[index]
+
+      # NOTE: If the trajectory isn't active, that means it doesn't have any
+      # time-steps in it, but we are in step, so the assumption is that it has
+      # a prior observation from which we are stepping away from.
+
+      # TODO(afrozm): Let's re-visit this if it becomes too restrictive.
+      assert trajectory.is_active()
+
+      # To this trajectory's last time-step, set actions.
+      trajectory.change_last_time_step(action=actions[index])
+
+      # Create a new time-step to add observation, done & rewards (no actions).
+      trajectory.add_time_step(
+          observation=observations[index],
+          done=dones[index],
+          raw_reward=raw_rewards[index],
+          processed_reward=processed_rewards[index])
+
+      # If the trajectory is completed, i.e. dones[index] == True, then we
+      # account for it right-away.
+      if dones[index]:
+        self._complete_trajectory(trajectory, index)
+
+        # NOTE: The new trajectory at `index` is going to be in-active and
+        # `reset` should be called on it.
+        assert not self._trajectories[index].is_active()
diff --git a/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
new file mode 100644
index 000000000..44b6ca48c
--- /dev/null
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -0,0 +1,212 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 google3.third_party.py.tensor2tensor.envs.trajectory."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.envs import trajectory
+import tensorflow as tf
+
+
+class TrajectoryTest(tf.test.TestCase):
+
+  def test_empty_trajectory(self):
+    t = trajectory.Trajectory()
+    self.assertFalse(t.is_active())
+    self.assertEqual(0, t.num_time_steps())
+
+  def test_add_time_step(self):
+    t = trajectory.Trajectory()
+    t.add_time_step(observation=1, done=True)
+
+    # Test that the trajectory is now active.
+    self.assertTrue(t.is_active())
+
+    added_t = t.last_time_step()
+    self.assertEqual(1, added_t.observation)
+    self.assertTrue(added_t.done)
+    self.assertIsNone(None, added_t.raw_reward)
+    self.assertIsNone(None, added_t.processed_reward)
+    self.assertIsNone(None, added_t.action)
+
+    self.assertEqual(1, t.num_time_steps())
+
+  def test_change_last_time_step(self):
+    t = trajectory.Trajectory()
+    t.add_time_step(observation=1, done=False)
+    t.add_time_step(observation=1, done=True)
+    self.assertTrue(t.is_active())
+
+    num_ts_old = t.num_time_steps()
+    self.assertEqual(2, num_ts_old)
+
+    # Assert on what the last time-step is currently.
+    ts = t.last_time_step()
+    self.assertEqual(1, ts.observation)
+    self.assertTrue(ts.done)
+    self.assertEqual(None, ts.action)
+
+    # Change the last time-step.
+    t.change_last_time_step(done=False, action=5)
+
+    # Assert that it changed.
+    ts = t.last_time_step()
+    self.assertEqual(1, ts.observation)  # unchanged, since we didn't change it.
+    self.assertFalse(ts.done)  # was True earlier
+    self.assertEqual(5, ts.action)  # was None earlier
+
+    # Assert on the number of steps remaining the same as before.
+    self.assertEqual(num_ts_old, t.num_time_steps())
+
+
+class BatchTrajectoryTest(tf.test.TestCase):
+
+  BATCH_SIZE = 10
+  OBSERVATION_SHAPE = (3, 4)
+
+  def test_creation(self):
+    bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
+
+    self.assertEqual(self.BATCH_SIZE, len(bt.trajectories))
+    self.assertEqual(0, len(bt.completed_trajectories))
+
+  def test_reset_all(self):
+    bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
+
+    indices = np.arange(self.BATCH_SIZE)
+    observations = np.random.rand(*(
+        (self.BATCH_SIZE,) + self.OBSERVATION_SHAPE))
+
+    # Call reset.
+    bt.reset(indices, observations)
+
+    # Assert that all trajectories are active and not done (reset never marks
+    # anything as done).
+    self.assertTrue(all(t.is_active() for t in bt.trajectories))
+    self.assertEqual(0, len(bt.completed_trajectories))
+
+  def test_reset_some(self):
+    bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
+
+    indices = np.arange(self.BATCH_SIZE // 2)
+    observations = np.random.rand(*(
+        (self.BATCH_SIZE // 2,) + self.OBSERVATION_SHAPE))
+
+    # Just reset the first half.
+    bt.reset(indices, observations)
+
+    # So first half are active, rest aren't.
+    self.assertTrue(
+        all(t.is_active() for t in bt.trajectories[:self.BATCH_SIZE // 2]))
+    self.assertTrue(
+        all(not t.is_active() for t in bt.trajectories[self.BATCH_SIZE // 2:]))
+
+    # Nothing is done anyways.
+    self.assertEqual(0, len(bt.completed_trajectories))
+
+  def test_step(self):
+    bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
+
+    indices = np.arange(self.BATCH_SIZE)
+    observations = np.random.rand(*(
+        (self.BATCH_SIZE,) + self.OBSERVATION_SHAPE))
+
+    # Have to call reset first.
+    bt.reset(indices, observations)
+
+    # Create some fake data for calling step.
+    new_observations = np.random.rand(*(
+        (self.BATCH_SIZE,) + self.OBSERVATION_SHAPE))
+    raw_rewards = processed_rewards = actions = np.random.randn(self.BATCH_SIZE)
+    processed_rewards = np.int64(processed_rewards)
+    dones = raw_rewards > 0.5
+
+    # Force mark the first one as done anyways, so that there is something to
+    # test.
+    dones[0] = True
+
+    num_done = sum(dones)
+    self.assertLessEqual(1, num_done)  # i.e. num_done is atleast 1.
+
+    num_not_done = len(dones) - num_done
+
+    # Finally call step.
+    bt.step(new_observations, raw_rewards, processed_rewards, dones, actions)
+
+    # Expect to see `num_done` number of completed trajectories.
+    self.assertEqual(num_done, len(bt.completed_trajectories))
+
+    # Expect to see that the rest are marked as active.
+    num_active = sum(t.is_active() for t in bt.trajectories)
+    self.assertEqual(num_not_done, num_active)
+
+  def test_desired_placement_of_rewards_and_actions(self):
+    batch_size = 1
+    bt = trajectory.BatchTrajectory(batch_size=batch_size)
+
+    indices = np.arange(batch_size)
+    observations = np.random.rand(*((batch_size,) + self.OBSERVATION_SHAPE))
+
+    # Have to call reset first.
+    bt.reset(indices, observations)
+
+    # Create some fake data for calling step.
+    new_observations = np.random.rand(*((batch_size,) + self.OBSERVATION_SHAPE))
+    raw_rewards = processed_rewards = actions = np.random.randn(batch_size)
+    processed_rewards = processed_rewards.astype(np.int64)
+    dones = np.full(batch_size, False)
+
+    # Call step.
+    bt.step(new_observations, raw_rewards, processed_rewards, dones, actions)
+
+    # Assert that nothing is done, since dones is False
+    self.assertEqual(0, len(bt.completed_trajectories))
+
+    # The only trajectory is active.
+    self.assertEqual(batch_size, len(bt.trajectories))
+    t = bt.trajectories[0]
+    self.assertTrue(t.is_active())
+    self.assertEqual(2, t.num_time_steps())
+
+    ts = t.time_steps
+
+    # Now assert on placements
+
+    # i.e. the old observation/done is first and the new one comes later.
+    self.assertAllEqual(observations[0], ts[0].observation)
+    self.assertAllEqual(new_observations[0], ts[1].observation)
+
+    self.assertEqual(False, ts[0].done)
+    self.assertEqual(False, ts[1].done)
+
+    # Similarly actions went to the first time-step.
+    self.assertEqual(actions[0], ts[0].action)
+    self.assertIsNone(ts[1].action)
+
+    # However make sure reward went into the second time-step and not the first.
+    self.assertNear(raw_rewards[0], ts[1].raw_reward, 1e-6)
+    self.assertIsNone(ts[0].raw_reward)
+
+    # Similarly with processed_rewards.
+    self.assertEqual(processed_rewards[0], ts[1].processed_reward)
+    self.assertIsNone(ts[0].processed_reward)
+
+
+if __name__ == '__main__':
+  tf.test.main()

From afb1917d2077ee8fc1e27fcee3610cce67a7d37e Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 29 Jan 2019 22:31:42 -0800
Subject: [PATCH 1600/2720] INTERNAL

PiperOrigin-RevId: 231535375
---
 tensor2tensor/envs/trajectory_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index 44b6ca48c..be5ed1de6 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for google3.third_party.py.tensor2tensor.envs.trajectory."""
+"""Tests for tensor2tensor.envs.trajectory."""
 
 from __future__ import absolute_import
 from __future__ import division

From b94bfcbc579af9e08c4ac42466cd0c33a6d8a529 Mon Sep 17 00:00:00 2001
From: Yuwen Yan <ybbaigo@gmail.com>
Date: Thu, 31 Jan 2019 01:26:47 +0800
Subject: [PATCH 1601/2720] add textCNN model to __init__.py (#1421)

---
 tensor2tensor/models/__init__.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 77176a9e0..973c87a81 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -38,6 +38,7 @@
 from tensor2tensor.models import revnet
 from tensor2tensor.models import shake_shake
 from tensor2tensor.models import slicenet
+from tensor2tensor.models import text_cnn
 from tensor2tensor.models import transformer
 from tensor2tensor.models import vanilla_gan
 from tensor2tensor.models import xception

From 847ae4157541053df57fba3851035843827ce8a3 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Wed, 30 Jan 2019 18:27:01 +0100
Subject: [PATCH 1602/2720] Replace deprecated assert calls (#1411)

`self.assertEquals` --> `self.assertEqual`
---
 tensor2tensor/models/revnet_test.py   | 40 +++++++++++++--------------
 tensor2tensor/rl/gym_utils_test.py    |  2 +-
 tensor2tensor/utils/t2t_model_test.py |  2 +-
 3 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/tensor2tensor/models/revnet_test.py b/tensor2tensor/models/revnet_test.py
index 68fec94a2..62b2e9dde 100644
--- a/tensor2tensor/models/revnet_test.py
+++ b/tensor2tensor/models/revnet_test.py
@@ -24,86 +24,86 @@ class RevnetTest(tf.test.TestCase):
   def testH(self):
     rev_block_input = tf.random_uniform([1, 299, 299, 3])
     rev_block_output = revnet.downsample_bottleneck(rev_block_input, 256)
-    self.assertEquals(rev_block_output.get_shape().as_list(),
+    self.assertEqual(rev_block_output.get_shape().as_list(),
                       [1, 299, 299, 256])
 
   def testHStride(self):
     rev_block_input = tf.random_uniform([2, 299, 299, 256])
     rev_block_output = revnet.downsample_bottleneck(
         rev_block_input, 512, stride=2, scope='HStride')
-    self.assertEquals(rev_block_output.get_shape().as_list(),
+    self.assertEqual(rev_block_output.get_shape().as_list(),
                       [2, 150, 150, 512])
 
   def testInit(self):
     images = tf.random_uniform([1, 299, 299, 3])
     x1, x2 = revnet.init(images, 32)
-    self.assertEquals(x1.get_shape().as_list(), [1, 74, 74, 16])
-    self.assertEquals(x2.get_shape().as_list(), [1, 74, 74, 16])
+    self.assertEqual(x1.get_shape().as_list(), [1, 74, 74, 16])
+    self.assertEqual(x2.get_shape().as_list(), [1, 74, 74, 16])
 
   def testInit3D(self):
     images = tf.random_uniform([1, 299, 299, 299, 3])
     x1, x2 = revnet.init(images, 32, dim='3d', scope='init3d')
-    self.assertEquals(x1.get_shape().as_list(), [1, 74, 74, 74, 16])
-    self.assertEquals(x2.get_shape().as_list(), [1, 74, 74, 74, 16])
+    self.assertEqual(x1.get_shape().as_list(), [1, 74, 74, 74, 16])
+    self.assertEqual(x2.get_shape().as_list(), [1, 74, 74, 74, 16])
 
   def testUnit1(self):
     x1 = tf.random_uniform([4, 74, 74, 256])
     x2 = tf.random_uniform([4, 74, 74, 256])
     x1, x2 = revnet.unit(x1, x2, block_num=1, depth=64,
                          first_batch_norm=True, num_layers=1)
-    self.assertEquals(x1.get_shape().as_list(), [4, 74, 74, 256])
-    self.assertEquals(x2.get_shape().as_list(), [4, 74, 74, 256])
+    self.assertEqual(x1.get_shape().as_list(), [4, 74, 74, 256])
+    self.assertEqual(x2.get_shape().as_list(), [4, 74, 74, 256])
 
   def testUnit2(self):
     x1 = tf.random_uniform([4, 74, 74, 256])
     x2 = tf.random_uniform([4, 74, 74, 256])
     x1, x2 = revnet.unit(x1, x2, block_num=2, depth=128,
                          num_layers=1, stride=2)
-    self.assertEquals(x1.get_shape().as_list(), [4, 37, 37, 512])
-    self.assertEquals(x2.get_shape().as_list(), [4, 37, 37, 512])
+    self.assertEqual(x1.get_shape().as_list(), [4, 37, 37, 512])
+    self.assertEqual(x2.get_shape().as_list(), [4, 37, 37, 512])
 
   def testUnit3(self):
     x1 = tf.random_uniform([1, 37, 37, 512])
     x2 = tf.random_uniform([1, 37, 37, 512])
     x1, x2 = revnet.unit(x1, x2, block_num=3, depth=256,
                          num_layers=10, stride=2)
-    self.assertEquals(x1.get_shape().as_list(), [1, 19, 19, 1024])
-    self.assertEquals(x2.get_shape().as_list(), [1, 19, 19, 1024])
+    self.assertEqual(x1.get_shape().as_list(), [1, 19, 19, 1024])
+    self.assertEqual(x2.get_shape().as_list(), [1, 19, 19, 1024])
 
   def testUnit4(self):
     x1 = tf.random_uniform([1, 19, 19, 1024])
     x2 = tf.random_uniform([1, 19, 19, 1024])
     x1, x2 = revnet.unit(x1, x2, block_num=4, depth=416,
                          num_layers=1, stride=2)
-    self.assertEquals(x1.get_shape().as_list(), [1, 10, 10, 1664])
-    self.assertEquals(x2.get_shape().as_list(), [1, 10, 10, 1664])
+    self.assertEqual(x1.get_shape().as_list(), [1, 10, 10, 1664])
+    self.assertEqual(x2.get_shape().as_list(), [1, 10, 10, 1664])
 
   def testUnit3D(self):
     x1 = tf.random_uniform([4, 74, 74, 74, 256])
     x2 = tf.random_uniform([4, 74, 74, 74, 256])
     x1, x2 = revnet.unit(x1, x2, block_num=5, depth=128,
                          num_layers=1, dim='3d', stride=2)
-    self.assertEquals(x1.get_shape().as_list(), [4, 37, 37, 37, 512])
-    self.assertEquals(x2.get_shape().as_list(), [4, 37, 37, 37, 512])
+    self.assertEqual(x1.get_shape().as_list(), [4, 37, 37, 37, 512])
+    self.assertEqual(x2.get_shape().as_list(), [4, 37, 37, 37, 512])
 
   def testFinalBlock(self):
     x1 = tf.random_uniform([5, 10, 10, 1024])
     x2 = tf.random_uniform([5, 10, 10, 1024])
     logits = revnet.final_block(x1, x2)
-    self.assertEquals(logits.shape, [5, 1, 1, 2048])
+    self.assertEqual(logits.shape, [5, 1, 1, 2048])
 
   def testFinalBlock3D(self):
     x1 = tf.random_uniform([5, 10, 10, 10, 1024])
     x2 = tf.random_uniform([5, 10, 10, 10, 1024])
     logits = revnet.final_block(x1, x2, dim='3d', scope='FinalBlock3D')
-    self.assertEquals(logits.shape, [5, 1, 1, 1, 2048])
+    self.assertEqual(logits.shape, [5, 1, 1, 1, 2048])
 
   def testEndToEnd(self):
     images = tf.random_uniform([1, 299, 299, 3])
     hparams = revnet.revnet_base()
     hparams.mode = tf.estimator.ModeKeys.TRAIN
     logits = revnet.revnet(images, hparams)
-    self.assertEquals(logits.shape, [1, 1, 1, 3328])
+    self.assertEqual(logits.shape, [1, 1, 1, 3328])
 
   def testEndToEnd3D(self):
     images = tf.random_uniform([1, 299, 299, 299, 3])
@@ -111,7 +111,7 @@ def testEndToEnd3D(self):
     hparams.dim = '3d'
     hparams.mode = tf.estimator.ModeKeys.TRAIN
     logits = revnet.revnet(images, hparams)
-    self.assertEquals(logits.shape, [1, 1, 1, 1, 3328])
+    self.assertEqual(logits.shape, [1, 1, 1, 1, 3328])
 
 if __name__ == '__main__':
   tf.test.main()
diff --git a/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index e97762da3..ead4563af 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -56,7 +56,7 @@ def test_making_timewrapped_env(self):
     env = gym_utils.make_gym_env("CartPole-v0", rl_env_max_episode_steps=1000)
     self.assertTrue(isinstance(env, gym.Env))
     self.assertTrue(isinstance(env, gym.wrappers.TimeLimit))
-    self.assertEquals(1000, env._max_episode_steps)
+    self.assertEqual(1000, env._max_episode_steps)
 
   # Make a time-wrapped environment with unlimited limit.
   def test_unlimited_env(self):
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index 5692de322..07622fbef 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -35,7 +35,7 @@ def testSummarizeLosses(self):
                 "extra": tf.random_normal([])}
       outputs = model._summarize_losses(losses)
       self.assertIsNone(outputs, None)
-      self.assertEquals(
+      self.assertEqual(
           len(tf.get_collection(tf.GraphKeys.SUMMARIES, scope="losses")),
           len(losses))
 

From 78ae06194bb150f0d756e300864e6beb98862cd6 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Wed, 30 Jan 2019 18:27:33 +0100
Subject: [PATCH 1603/2720] Use efficient squared_difference instead of
 square(diff) (#1413)

* Use efficient squared_difference instead of square(diff)

* Use efficient squared_difference(x, y) instead of (x - y)**2
---
 tensor2tensor/data_generators/speech_recognition.py |  2 +-
 tensor2tensor/layers/common_layers.py               | 10 ++++++----
 tensor2tensor/layers/discretization.py              | 12 +++++++-----
 tensor2tensor/layers/modalities.py                  |  5 +++--
 tensor2tensor/layers/vq_discrete.py                 |  6 ++++--
 tensor2tensor/models/research/autoencoders.py       |  3 ++-
 tensor2tensor/models/research/transformer_nat.py    |  2 +-
 tensor2tensor/models/research/transformer_vae.py    |  3 ++-
 tensor2tensor/models/video/epva.py                  |  3 ++-
 tensor2tensor/utils/expert_utils.py                 |  2 +-
 tensor2tensor/utils/yellowfin.py                    |  2 +-
 11 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 258af10f3..730fe65f0 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -122,7 +122,7 @@ def preprocess_example(self, example, mode, hparams):
       # This replaces CMVN estimation on data
       var_epsilon = 1e-09
       mean = tf.reduce_mean(mel_fbanks, keepdims=True, axis=1)
-      variance = tf.reduce_mean(tf.square(mel_fbanks - mean),
+      variance = tf.reduce_mean(tf.squared_difference(mel_fbanks, mean),
                                 keepdims=True, axis=1)
       mel_fbanks = (mel_fbanks - mean) * tf.rsqrt(variance + var_epsilon)
 
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index b48300c26..b7bb5615b 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -256,7 +256,7 @@ def standardize_images(x):
     x = tf.to_float(tf.reshape(x, [-1] + x_shape[-3:]))
     x_mean = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
     x_variance = tf.reduce_mean(
-        tf.square(x - x_mean), axis=[1, 2], keepdims=True)
+        tf.squared_difference(x, x_mean), axis=[1, 2], keepdims=True)
     num_pixels = tf.to_float(x_shape[-2] * x_shape[-3])
     x = (x - x_mean) / tf.maximum(tf.sqrt(x_variance), tf.rsqrt(num_pixels))
     return tf.reshape(x, x_shape)
@@ -634,7 +634,8 @@ def layer_norm_compute(x, epsilon, scale, bias):
   """Layer norm raw computation."""
   epsilon, scale, bias = [cast_like(t, x) for t in [epsilon, scale, bias]]
   mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
-  variance = tf.reduce_mean(tf.square(x - mean), axis=[-1], keepdims=True)
+  variance = tf.reduce_mean(
+      tf.squared_difference(x, mean), axis=[-1], keepdims=True)
   norm_x = (x - mean) * tf.rsqrt(variance + epsilon)
   return norm_x * scale + bias
 
@@ -690,7 +691,8 @@ def l2_norm(x, filters=None, epsilon=1e-6, name=None, reuse=None):
         "l2_norm_bias", [filters], initializer=tf.zeros_initializer())
     epsilon, scale, bias = [cast_like(t, x) for t in [epsilon, scale, bias]]
     mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
-    l2norm = tf.reduce_sum(tf.square(x - mean), axis=[-1], keepdims=True)
+    l2norm = tf.reduce_sum(
+        tf.squared_difference(x, mean), axis=[-1], keepdims=True)
     norm_x = (x - mean) * tf.rsqrt(l2norm + epsilon)
     return norm_x * scale + bias
 
@@ -3346,7 +3348,7 @@ def get_sorted_projections(x):
 
     proj1 = get_sorted_projections(logits1)
     proj2 = get_sorted_projections(logits2)
-    dist = tf.reduce_mean(tf.square(proj1 - proj2))
+    dist = tf.reduce_mean(tf.squared_difference(proj1, proj2))
     if return_logits:
       return dist, logits1, logits2
     return dist
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 2293cd8b0..020d82d1d 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -217,8 +217,8 @@ def embedding_lookup(x,
 
   # Currently, we use the mean scaling for the commitment loss, as opposed to
   # summing across all non-batch dimensions.
-  q_loss = tf.reduce_mean(tf.square((tf.stop_gradient(x) - x_means)))
-  e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
+  q_loss = tf.reduce_mean(tf.squared_difference(tf.stop_gradient(x), x_means))
+  e_loss = tf.reduce_mean(tf.squared_difference(x, tf.stop_gradient(x_means)))
   return x_means_hot, x_means, q_loss, e_loss, neg_q_entropy
 
 
@@ -469,7 +469,8 @@ def gumbel_softmax(x,
     # Add losses that prevent too few being used.
     distrib = tf.reshape(logsm, [-1, 2**z_size]) * maxvhot
     d_mean = tf.reduce_mean(distrib, axis=[0], keep_dims=True)
-    d_variance = tf.reduce_mean(tf.square(distrib - d_mean), axis=[0])
+    d_variance = tf.reduce_mean(
+        tf.squared_difference(distrib, d_mean), axis=[0])
     d_dev = -tf.reduce_mean(d_variance)
     ret = s
 
@@ -924,7 +925,7 @@ def vq_nearest_neighbor(x, means,
     x_means_hot = tf.one_hot(x_means_idx, bottleneck_size)
   x_means_hot_flat = tf.reshape(x_means_hot, [-1, bottleneck_size])
   x_means = tf.matmul(x_means_hot_flat, means)
-  e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
+  e_loss = tf.reduce_mean(tf.squared_difference(x, tf.stop_gradient(x_means)))
   return x_means_hot, e_loss, dist
 
 
@@ -1333,7 +1334,8 @@ def gumbel_softmax_discrete_bottleneck(x,
   x_means_assignments_flat = tf.reshape(x_means_assignments,
                                         [-1, bottleneck_size])
   x_means = tf.matmul(x_means_assignments_flat, means)
-  commitment_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
+  commitment_loss = tf.reduce_mean(
+      tf.squared_difference(x, tf.stop_gradient(x_means)))
 
   # Update the ema variables.
   updated_ema_count = moving_averages.assign_moving_average(
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 09e8b9b1c..a42a251aa 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -736,7 +736,8 @@ class VideoModalityL2(VideoModalityL1):
   """Modality for videos with L2 loss."""
 
   def internal_loss(self, logits, targets):
-    return tf.nn.relu((logits - targets)**2 - self.cutoff * self.cutoff)
+    return tf.nn.relu(
+        tf.squared_difference(logits, targets) - self.cutoff * self.cutoff)
 
 
 class VideoModalityL2Raw(VideoModalityL2):
@@ -916,7 +917,7 @@ def targets_bottom(self, x):
     return tf.to_float(x)
 
   def loss(self, body_output, targets):
-    loss = tf.square(body_output - tf.to_float(targets))
+    loss = tf.squared_difference(body_output, tf.to_float(targets))
     return tf.reduce_mean(loss), tf.constant(1.0)
 
 
diff --git a/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index f45289df2..344767084 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -138,8 +138,10 @@ def embedding_lookup(self, x, means):
         x_means_hot, [-1, self.hparams.num_blocks, self.hparams.block_v_size])
     x_means = tf.matmul(tf.transpose(x_means_hot_flat, perm=[1, 0, 2]), means)
     x_means = tf.transpose(x_means, [1, 0, 2])
-    q_loss = tf.reduce_mean(tf.square((tf.stop_gradient(x) - x_means)))
-    e_loss = tf.reduce_mean((x - tf.stop_gradient(x_means))**2)
+    q_loss = tf.reduce_mean(
+        tf.squared_difference(tf.stop_gradient(x), x_means))
+    e_loss = tf.reduce_mean(
+        tf.squared_difference(x, tf.stop_gradient(x_means)))
     return x_means_hot, x_means, q_loss, e_loss
 
   def bit_to_int(self, x_bit, num_bits, base=2):
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index b9c42e89e..0da05beba 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -219,7 +219,8 @@ def body(self, features):
         # minimized by just setting x=0 and b=0 -- so we don't want too much
         # of the influence of this, and we stop-gradient to not zero-out x.
         x_stop = tf.stop_gradient(x)
-        xb_loss = tf.reduce_mean(tf.reduce_sum(tf.square(x_stop - b), axis=-1))
+        xb_loss = tf.reduce_mean(tf.reduce_sum(
+            tf.squared_difference(x_stop, b), axis=-1))
         # To prevent this loss from exploding we clip at 1, but anneal clipping.
         clip_max = 1.0 / common_layers.inverse_exp_decay(
             warm_step, min_value=0.001)
diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 0b11e89cf..c6be947fc 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -65,7 +65,7 @@ def vq_nearest_neighbor(x, hparams):
     x_means_idx = tf.argmax(-dist, axis=-1)
     x_means_hot = tf.one_hot(x_means_idx, depth=bottleneck_size)
   x_means = tf.matmul(x_means_hot, means)
-  e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
+  e_loss = tf.reduce_mean(tf.squared_difference(x, tf.stop_gradient(x_means)))
   return x_means_hot, e_loss
 
 
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 64aecd405..44e92c5ce 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -436,7 +436,8 @@ def ae_transformer_internal(inputs,
         losses["neg_q_entropy"] = neg_q_entropy * hparams.entropy_scale
       else:
         inputs_c = decode_transformer(inputs, ed, targets_c, hparams, "dec_c")
-        losses["latent_pred"] = tf.reduce_mean((inputs_c - targets_c)**2) * 20
+        losses["latent_pred"] = tf.reduce_mean(
+            tf.squared_difference(inputs_c, targets_c)) * 20
         def bn_inputs():
           with tf.variable_scope(tf.get_variable_scope(), reuse=True):
             bn, _, _, _, _ = hparams.bottleneck(
diff --git a/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index 95164462a..dbd4e5283 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -590,7 +590,8 @@ def mean_squared_error(true, pred):
   Returns:
     mean squared error between ground truth and predicted image.
   """
-  result = tf.reduce_sum(tf.square(true - pred)) / tf.to_float(tf.size(pred))
+  result = tf.reduce_sum(
+      tf.squared_difference(true, pred)) / tf.to_float(tf.size(pred))
   return result
 
 
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 829dcd607..0630b7ca4 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -364,7 +364,7 @@ def cv_squared(x):
   epsilon = 1e-10
   float_size = tf.to_float(tf.size(x)) + epsilon
   mean = tf.reduce_sum(x) / float_size
-  variance = tf.reduce_sum(tf.square(x - mean)) / float_size
+  variance = tf.reduce_sum(tf.squared_difference(x, mean)) / float_size
   return variance / (tf.square(mean) + epsilon)
 
 
diff --git a/tensor2tensor/utils/yellowfin.py b/tensor2tensor/utils/yellowfin.py
index 4d33c2aad..97acee4fb 100644
--- a/tensor2tensor/utils/yellowfin.py
+++ b/tensor2tensor/utils/yellowfin.py
@@ -392,7 +392,7 @@ def _get_lr_tensor(self):
     Returns:
       The lr_t.
     """
-    lr = (1.0 - tf.sqrt(self._mu))**2 / self._h_min
+    lr = tf.squared_difference(1.0, tf.sqrt(self._mu)) / self._h_min
     return lr
 
   def _get_mu_tensor(self):

From 7548d2974acec9343ea3f482cfc2303fdae48f37 Mon Sep 17 00:00:00 2001
From: etragas-fathom <43351375+etragas-fathom@users.noreply.github.com>
Date: Wed, 30 Jan 2019 12:28:32 -0500
Subject: [PATCH 1604/2720] Partial Mixed Precision for Universal Transformer
 (#1416)

* add hparams

* fix ut

* make small again

* Add init loss scale hparam

* add hparam and docstring

* try out large init loss scale

* try 2**8

* Revert "try 2**8"

This reverts commit 798b3e1fa6b9d7ee38db6713f01c5ec293ccadbe.

* make big

* delete redundant line

* reduce batch size due to oom

* move to base batch size

* make small batches

* fix register

* move to old transformer hparams

* change hparam name

* add bs print

* add hparam

* simplify hparams

* fix space
---
 tensor2tensor/layers/common_hparams.py                    | 2 ++
 tensor2tensor/models/research/universal_transformer.py    | 8 +++++++-
 .../models/research/universal_transformer_util.py         | 4 ++--
 tensor2tensor/utils/optimize.py                           | 6 ++++--
 4 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 0d5ce6598..d08ecbb52 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -75,6 +75,8 @@ def basic_params1():
       # Mixed precision training only supports exponential scaling currently
       # To disable the scaler, see to 0/False
       mixed_precision_optimizer_loss_scaler="exponential",
+      # Determines the initial loss scaling value for mixed precision
+      mixed_precision_optimizer_init_loss_scale=2**15,
       # Whether to zero gradients that were not computed, so that the
       # appropriate slots are created. Useful for sharing checkpoints between
       # models with different sets of heads.
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index a160e3627..38c5de74d 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -459,7 +459,6 @@ def universal_transformer_base_tpu():
   hparams.add_step_timing_signal = False
   return hparams
 
-
 @registry.register_hparams
 def universal_transformer_big():
   hparams = universal_transformer_base()
@@ -469,6 +468,13 @@ def universal_transformer_big():
   return hparams
 
 
+@registry.register_hparams
+def universal_transformer_base_fp16():
+  hparams = transformer.transformer_base()
+  hparams = update_hparams_for_universal_transformer(hparams)
+  hparams.activation_dtype = 'float16'
+  return hparams
+
 @registry.register_hparams
 def universal_transformer_small():
   hparams = transformer.transformer_base()
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 673747e60..54906e5b6 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -1442,7 +1442,7 @@ def add_position_timing_signal(x, step, hparams):
       length, channels, start_index=index)
 
   if hparams.add_or_concat_timing_signal == "add":
-    x_with_timing = x + signal
+    x_with_timing = x + common_layers.cast_like(signal, x)
 
   elif hparams.add_or_concat_timing_signal == "concat":
     batch_size = common_layers.shape_list(x)[0]
@@ -1479,7 +1479,7 @@ def add_step_timing_signal(x, step, hparams):
         channels, step, num_steps)
 
   if hparams.add_or_concat_timing_signal == "add":
-    x_with_timing = x + signal
+    x_with_timing = x + common_layers.cast_like(signal, x)
 
   elif hparams.add_or_concat_timing_signal == "concat":
     batch_size = common_layers.shape_list(x)[0]
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index ac301189b..f0ae96b06 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -204,9 +204,11 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
         raise ValueError("Mixed precision training only supports the "
                          "exponential loss scaler")
       else:
-        tf.logging.info("Using Exponential Update Loss Scaler")
+        tf.logging.info(("Using Exponential Update Loss Scaler with",
+                         "init loss scale of {}".format(
+                           hparams.mixed_precision_optimizer_init_loss_scale)))
         manager = tf.contrib.mixed_precision.ExponentialUpdateLossScaleManager(
-            init_loss_scale=2**15,
+            init_loss_scale=hparams.mixed_precision_optimizer_init_loss_scale,
             incr_every_n_steps=2000,
             decr_every_n_nan_or_inf=2,
             incr_ratio=2,

From e3e425037cf0abf790b6651446e91d3148d17fb0 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Wed, 30 Jan 2019 09:30:32 -0800
Subject: [PATCH 1605/2720] internal merge of PR #1411

PiperOrigin-RevId: 231608988
---
 tensor2tensor/data_generators/speech_recognition.py  |  2 +-
 tensor2tensor/layers/common_hparams.py               |  2 --
 tensor2tensor/layers/common_layers.py                | 10 ++++------
 tensor2tensor/layers/discretization.py               | 12 +++++-------
 tensor2tensor/layers/modalities.py                   |  5 ++---
 tensor2tensor/layers/vq_discrete.py                  |  6 ++----
 tensor2tensor/models/__init__.py                     |  1 -
 tensor2tensor/models/research/autoencoders.py        |  3 +--
 tensor2tensor/models/research/transformer_nat.py     |  2 +-
 tensor2tensor/models/research/transformer_vae.py     |  3 +--
 .../models/research/universal_transformer.py         |  8 +-------
 .../models/research/universal_transformer_util.py    |  4 ++--
 tensor2tensor/models/revnet_test.py                  |  6 ++----
 tensor2tensor/models/video/epva.py                   |  3 +--
 tensor2tensor/utils/expert_utils.py                  |  2 +-
 tensor2tensor/utils/optimize.py                      |  6 ++----
 tensor2tensor/utils/yellowfin.py                     |  2 +-
 17 files changed, 27 insertions(+), 50 deletions(-)

diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 730fe65f0..258af10f3 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -122,7 +122,7 @@ def preprocess_example(self, example, mode, hparams):
       # This replaces CMVN estimation on data
       var_epsilon = 1e-09
       mean = tf.reduce_mean(mel_fbanks, keepdims=True, axis=1)
-      variance = tf.reduce_mean(tf.squared_difference(mel_fbanks, mean),
+      variance = tf.reduce_mean(tf.square(mel_fbanks - mean),
                                 keepdims=True, axis=1)
       mel_fbanks = (mel_fbanks - mean) * tf.rsqrt(variance + var_epsilon)
 
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index d08ecbb52..0d5ce6598 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -75,8 +75,6 @@ def basic_params1():
       # Mixed precision training only supports exponential scaling currently
       # To disable the scaler, see to 0/False
       mixed_precision_optimizer_loss_scaler="exponential",
-      # Determines the initial loss scaling value for mixed precision
-      mixed_precision_optimizer_init_loss_scale=2**15,
       # Whether to zero gradients that were not computed, so that the
       # appropriate slots are created. Useful for sharing checkpoints between
       # models with different sets of heads.
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index b7bb5615b..b48300c26 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -256,7 +256,7 @@ def standardize_images(x):
     x = tf.to_float(tf.reshape(x, [-1] + x_shape[-3:]))
     x_mean = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
     x_variance = tf.reduce_mean(
-        tf.squared_difference(x, x_mean), axis=[1, 2], keepdims=True)
+        tf.square(x - x_mean), axis=[1, 2], keepdims=True)
     num_pixels = tf.to_float(x_shape[-2] * x_shape[-3])
     x = (x - x_mean) / tf.maximum(tf.sqrt(x_variance), tf.rsqrt(num_pixels))
     return tf.reshape(x, x_shape)
@@ -634,8 +634,7 @@ def layer_norm_compute(x, epsilon, scale, bias):
   """Layer norm raw computation."""
   epsilon, scale, bias = [cast_like(t, x) for t in [epsilon, scale, bias]]
   mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
-  variance = tf.reduce_mean(
-      tf.squared_difference(x, mean), axis=[-1], keepdims=True)
+  variance = tf.reduce_mean(tf.square(x - mean), axis=[-1], keepdims=True)
   norm_x = (x - mean) * tf.rsqrt(variance + epsilon)
   return norm_x * scale + bias
 
@@ -691,8 +690,7 @@ def l2_norm(x, filters=None, epsilon=1e-6, name=None, reuse=None):
         "l2_norm_bias", [filters], initializer=tf.zeros_initializer())
     epsilon, scale, bias = [cast_like(t, x) for t in [epsilon, scale, bias]]
     mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
-    l2norm = tf.reduce_sum(
-        tf.squared_difference(x, mean), axis=[-1], keepdims=True)
+    l2norm = tf.reduce_sum(tf.square(x - mean), axis=[-1], keepdims=True)
     norm_x = (x - mean) * tf.rsqrt(l2norm + epsilon)
     return norm_x * scale + bias
 
@@ -3348,7 +3346,7 @@ def get_sorted_projections(x):
 
     proj1 = get_sorted_projections(logits1)
     proj2 = get_sorted_projections(logits2)
-    dist = tf.reduce_mean(tf.squared_difference(proj1, proj2))
+    dist = tf.reduce_mean(tf.square(proj1 - proj2))
     if return_logits:
       return dist, logits1, logits2
     return dist
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 020d82d1d..2293cd8b0 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -217,8 +217,8 @@ def embedding_lookup(x,
 
   # Currently, we use the mean scaling for the commitment loss, as opposed to
   # summing across all non-batch dimensions.
-  q_loss = tf.reduce_mean(tf.squared_difference(tf.stop_gradient(x), x_means))
-  e_loss = tf.reduce_mean(tf.squared_difference(x, tf.stop_gradient(x_means)))
+  q_loss = tf.reduce_mean(tf.square((tf.stop_gradient(x) - x_means)))
+  e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
   return x_means_hot, x_means, q_loss, e_loss, neg_q_entropy
 
 
@@ -469,8 +469,7 @@ def gumbel_softmax(x,
     # Add losses that prevent too few being used.
     distrib = tf.reshape(logsm, [-1, 2**z_size]) * maxvhot
     d_mean = tf.reduce_mean(distrib, axis=[0], keep_dims=True)
-    d_variance = tf.reduce_mean(
-        tf.squared_difference(distrib, d_mean), axis=[0])
+    d_variance = tf.reduce_mean(tf.square(distrib - d_mean), axis=[0])
     d_dev = -tf.reduce_mean(d_variance)
     ret = s
 
@@ -925,7 +924,7 @@ def vq_nearest_neighbor(x, means,
     x_means_hot = tf.one_hot(x_means_idx, bottleneck_size)
   x_means_hot_flat = tf.reshape(x_means_hot, [-1, bottleneck_size])
   x_means = tf.matmul(x_means_hot_flat, means)
-  e_loss = tf.reduce_mean(tf.squared_difference(x, tf.stop_gradient(x_means)))
+  e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
   return x_means_hot, e_loss, dist
 
 
@@ -1334,8 +1333,7 @@ def gumbel_softmax_discrete_bottleneck(x,
   x_means_assignments_flat = tf.reshape(x_means_assignments,
                                         [-1, bottleneck_size])
   x_means = tf.matmul(x_means_assignments_flat, means)
-  commitment_loss = tf.reduce_mean(
-      tf.squared_difference(x, tf.stop_gradient(x_means)))
+  commitment_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
 
   # Update the ema variables.
   updated_ema_count = moving_averages.assign_moving_average(
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index a42a251aa..09e8b9b1c 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -736,8 +736,7 @@ class VideoModalityL2(VideoModalityL1):
   """Modality for videos with L2 loss."""
 
   def internal_loss(self, logits, targets):
-    return tf.nn.relu(
-        tf.squared_difference(logits, targets) - self.cutoff * self.cutoff)
+    return tf.nn.relu((logits - targets)**2 - self.cutoff * self.cutoff)
 
 
 class VideoModalityL2Raw(VideoModalityL2):
@@ -917,7 +916,7 @@ def targets_bottom(self, x):
     return tf.to_float(x)
 
   def loss(self, body_output, targets):
-    loss = tf.squared_difference(body_output, tf.to_float(targets))
+    loss = tf.square(body_output - tf.to_float(targets))
     return tf.reduce_mean(loss), tf.constant(1.0)
 
 
diff --git a/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index 344767084..f45289df2 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -138,10 +138,8 @@ def embedding_lookup(self, x, means):
         x_means_hot, [-1, self.hparams.num_blocks, self.hparams.block_v_size])
     x_means = tf.matmul(tf.transpose(x_means_hot_flat, perm=[1, 0, 2]), means)
     x_means = tf.transpose(x_means, [1, 0, 2])
-    q_loss = tf.reduce_mean(
-        tf.squared_difference(tf.stop_gradient(x), x_means))
-    e_loss = tf.reduce_mean(
-        tf.squared_difference(x, tf.stop_gradient(x_means)))
+    q_loss = tf.reduce_mean(tf.square((tf.stop_gradient(x) - x_means)))
+    e_loss = tf.reduce_mean((x - tf.stop_gradient(x_means))**2)
     return x_means_hot, x_means, q_loss, e_loss
 
   def bit_to_int(self, x_bit, num_bits, base=2):
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 973c87a81..77176a9e0 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -38,7 +38,6 @@
 from tensor2tensor.models import revnet
 from tensor2tensor.models import shake_shake
 from tensor2tensor.models import slicenet
-from tensor2tensor.models import text_cnn
 from tensor2tensor.models import transformer
 from tensor2tensor.models import vanilla_gan
 from tensor2tensor.models import xception
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 0da05beba..b9c42e89e 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -219,8 +219,7 @@ def body(self, features):
         # minimized by just setting x=0 and b=0 -- so we don't want too much
         # of the influence of this, and we stop-gradient to not zero-out x.
         x_stop = tf.stop_gradient(x)
-        xb_loss = tf.reduce_mean(tf.reduce_sum(
-            tf.squared_difference(x_stop, b), axis=-1))
+        xb_loss = tf.reduce_mean(tf.reduce_sum(tf.square(x_stop - b), axis=-1))
         # To prevent this loss from exploding we clip at 1, but anneal clipping.
         clip_max = 1.0 / common_layers.inverse_exp_decay(
             warm_step, min_value=0.001)
diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index c6be947fc..0b11e89cf 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -65,7 +65,7 @@ def vq_nearest_neighbor(x, hparams):
     x_means_idx = tf.argmax(-dist, axis=-1)
     x_means_hot = tf.one_hot(x_means_idx, depth=bottleneck_size)
   x_means = tf.matmul(x_means_hot, means)
-  e_loss = tf.reduce_mean(tf.squared_difference(x, tf.stop_gradient(x_means)))
+  e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
   return x_means_hot, e_loss
 
 
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 44e92c5ce..64aecd405 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -436,8 +436,7 @@ def ae_transformer_internal(inputs,
         losses["neg_q_entropy"] = neg_q_entropy * hparams.entropy_scale
       else:
         inputs_c = decode_transformer(inputs, ed, targets_c, hparams, "dec_c")
-        losses["latent_pred"] = tf.reduce_mean(
-            tf.squared_difference(inputs_c, targets_c)) * 20
+        losses["latent_pred"] = tf.reduce_mean((inputs_c - targets_c)**2) * 20
         def bn_inputs():
           with tf.variable_scope(tf.get_variable_scope(), reuse=True):
             bn, _, _, _, _ = hparams.bottleneck(
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 38c5de74d..a160e3627 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -459,6 +459,7 @@ def universal_transformer_base_tpu():
   hparams.add_step_timing_signal = False
   return hparams
 
+
 @registry.register_hparams
 def universal_transformer_big():
   hparams = universal_transformer_base()
@@ -468,13 +469,6 @@ def universal_transformer_big():
   return hparams
 
 
-@registry.register_hparams
-def universal_transformer_base_fp16():
-  hparams = transformer.transformer_base()
-  hparams = update_hparams_for_universal_transformer(hparams)
-  hparams.activation_dtype = 'float16'
-  return hparams
-
 @registry.register_hparams
 def universal_transformer_small():
   hparams = transformer.transformer_base()
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 54906e5b6..673747e60 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -1442,7 +1442,7 @@ def add_position_timing_signal(x, step, hparams):
       length, channels, start_index=index)
 
   if hparams.add_or_concat_timing_signal == "add":
-    x_with_timing = x + common_layers.cast_like(signal, x)
+    x_with_timing = x + signal
 
   elif hparams.add_or_concat_timing_signal == "concat":
     batch_size = common_layers.shape_list(x)[0]
@@ -1479,7 +1479,7 @@ def add_step_timing_signal(x, step, hparams):
         channels, step, num_steps)
 
   if hparams.add_or_concat_timing_signal == "add":
-    x_with_timing = x + common_layers.cast_like(signal, x)
+    x_with_timing = x + signal
 
   elif hparams.add_or_concat_timing_signal == "concat":
     batch_size = common_layers.shape_list(x)[0]
diff --git a/tensor2tensor/models/revnet_test.py b/tensor2tensor/models/revnet_test.py
index 62b2e9dde..e60bc8420 100644
--- a/tensor2tensor/models/revnet_test.py
+++ b/tensor2tensor/models/revnet_test.py
@@ -24,15 +24,13 @@ class RevnetTest(tf.test.TestCase):
   def testH(self):
     rev_block_input = tf.random_uniform([1, 299, 299, 3])
     rev_block_output = revnet.downsample_bottleneck(rev_block_input, 256)
-    self.assertEqual(rev_block_output.get_shape().as_list(),
-                      [1, 299, 299, 256])
+    self.assertEqual(rev_block_output.get_shape().as_list(), [1, 299, 299, 256])
 
   def testHStride(self):
     rev_block_input = tf.random_uniform([2, 299, 299, 256])
     rev_block_output = revnet.downsample_bottleneck(
         rev_block_input, 512, stride=2, scope='HStride')
-    self.assertEqual(rev_block_output.get_shape().as_list(),
-                      [2, 150, 150, 512])
+    self.assertEqual(rev_block_output.get_shape().as_list(), [2, 150, 150, 512])
 
   def testInit(self):
     images = tf.random_uniform([1, 299, 299, 3])
diff --git a/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index dbd4e5283..95164462a 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -590,8 +590,7 @@ def mean_squared_error(true, pred):
   Returns:
     mean squared error between ground truth and predicted image.
   """
-  result = tf.reduce_sum(
-      tf.squared_difference(true, pred)) / tf.to_float(tf.size(pred))
+  result = tf.reduce_sum(tf.square(true - pred)) / tf.to_float(tf.size(pred))
   return result
 
 
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 0630b7ca4..829dcd607 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -364,7 +364,7 @@ def cv_squared(x):
   epsilon = 1e-10
   float_size = tf.to_float(tf.size(x)) + epsilon
   mean = tf.reduce_sum(x) / float_size
-  variance = tf.reduce_sum(tf.squared_difference(x, mean)) / float_size
+  variance = tf.reduce_sum(tf.square(x - mean)) / float_size
   return variance / (tf.square(mean) + epsilon)
 
 
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index f0ae96b06..ac301189b 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -204,11 +204,9 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
         raise ValueError("Mixed precision training only supports the "
                          "exponential loss scaler")
       else:
-        tf.logging.info(("Using Exponential Update Loss Scaler with",
-                         "init loss scale of {}".format(
-                           hparams.mixed_precision_optimizer_init_loss_scale)))
+        tf.logging.info("Using Exponential Update Loss Scaler")
         manager = tf.contrib.mixed_precision.ExponentialUpdateLossScaleManager(
-            init_loss_scale=hparams.mixed_precision_optimizer_init_loss_scale,
+            init_loss_scale=2**15,
             incr_every_n_steps=2000,
             decr_every_n_nan_or_inf=2,
             incr_ratio=2,
diff --git a/tensor2tensor/utils/yellowfin.py b/tensor2tensor/utils/yellowfin.py
index 97acee4fb..4d33c2aad 100644
--- a/tensor2tensor/utils/yellowfin.py
+++ b/tensor2tensor/utils/yellowfin.py
@@ -392,7 +392,7 @@ def _get_lr_tensor(self):
     Returns:
       The lr_t.
     """
-    lr = tf.squared_difference(1.0, tf.sqrt(self._mu)) / self._h_min
+    lr = (1.0 - tf.sqrt(self._mu))**2 / self._h_min
     return lr
 
   def _get_mu_tensor(self):

From 3af056d88d242a5e4a112d2822de5d28867ba64e Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Wed, 30 Jan 2019 09:44:00 -0800
Subject: [PATCH 1606/2720] internal merge of PR #1413

PiperOrigin-RevId: 231611349
---
 tensor2tensor/data_generators/speech_recognition.py |  2 +-
 tensor2tensor/layers/common_layers.py               | 10 ++++++----
 tensor2tensor/layers/discretization.py              | 12 +++++++-----
 tensor2tensor/layers/modalities.py                  |  5 +++--
 tensor2tensor/layers/vq_discrete.py                 |  6 ++++--
 tensor2tensor/models/research/autoencoders.py       |  3 ++-
 tensor2tensor/models/research/transformer_nat.py    |  2 +-
 tensor2tensor/models/research/transformer_vae.py    |  3 ++-
 tensor2tensor/models/video/epva.py                  |  3 ++-
 tensor2tensor/utils/expert_utils.py                 |  2 +-
 tensor2tensor/utils/yellowfin.py                    |  2 +-
 11 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 258af10f3..730fe65f0 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -122,7 +122,7 @@ def preprocess_example(self, example, mode, hparams):
       # This replaces CMVN estimation on data
       var_epsilon = 1e-09
       mean = tf.reduce_mean(mel_fbanks, keepdims=True, axis=1)
-      variance = tf.reduce_mean(tf.square(mel_fbanks - mean),
+      variance = tf.reduce_mean(tf.squared_difference(mel_fbanks, mean),
                                 keepdims=True, axis=1)
       mel_fbanks = (mel_fbanks - mean) * tf.rsqrt(variance + var_epsilon)
 
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index b48300c26..b7bb5615b 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -256,7 +256,7 @@ def standardize_images(x):
     x = tf.to_float(tf.reshape(x, [-1] + x_shape[-3:]))
     x_mean = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
     x_variance = tf.reduce_mean(
-        tf.square(x - x_mean), axis=[1, 2], keepdims=True)
+        tf.squared_difference(x, x_mean), axis=[1, 2], keepdims=True)
     num_pixels = tf.to_float(x_shape[-2] * x_shape[-3])
     x = (x - x_mean) / tf.maximum(tf.sqrt(x_variance), tf.rsqrt(num_pixels))
     return tf.reshape(x, x_shape)
@@ -634,7 +634,8 @@ def layer_norm_compute(x, epsilon, scale, bias):
   """Layer norm raw computation."""
   epsilon, scale, bias = [cast_like(t, x) for t in [epsilon, scale, bias]]
   mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
-  variance = tf.reduce_mean(tf.square(x - mean), axis=[-1], keepdims=True)
+  variance = tf.reduce_mean(
+      tf.squared_difference(x, mean), axis=[-1], keepdims=True)
   norm_x = (x - mean) * tf.rsqrt(variance + epsilon)
   return norm_x * scale + bias
 
@@ -690,7 +691,8 @@ def l2_norm(x, filters=None, epsilon=1e-6, name=None, reuse=None):
         "l2_norm_bias", [filters], initializer=tf.zeros_initializer())
     epsilon, scale, bias = [cast_like(t, x) for t in [epsilon, scale, bias]]
     mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
-    l2norm = tf.reduce_sum(tf.square(x - mean), axis=[-1], keepdims=True)
+    l2norm = tf.reduce_sum(
+        tf.squared_difference(x, mean), axis=[-1], keepdims=True)
     norm_x = (x - mean) * tf.rsqrt(l2norm + epsilon)
     return norm_x * scale + bias
 
@@ -3346,7 +3348,7 @@ def get_sorted_projections(x):
 
     proj1 = get_sorted_projections(logits1)
     proj2 = get_sorted_projections(logits2)
-    dist = tf.reduce_mean(tf.square(proj1 - proj2))
+    dist = tf.reduce_mean(tf.squared_difference(proj1, proj2))
     if return_logits:
       return dist, logits1, logits2
     return dist
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 2293cd8b0..020d82d1d 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -217,8 +217,8 @@ def embedding_lookup(x,
 
   # Currently, we use the mean scaling for the commitment loss, as opposed to
   # summing across all non-batch dimensions.
-  q_loss = tf.reduce_mean(tf.square((tf.stop_gradient(x) - x_means)))
-  e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
+  q_loss = tf.reduce_mean(tf.squared_difference(tf.stop_gradient(x), x_means))
+  e_loss = tf.reduce_mean(tf.squared_difference(x, tf.stop_gradient(x_means)))
   return x_means_hot, x_means, q_loss, e_loss, neg_q_entropy
 
 
@@ -469,7 +469,8 @@ def gumbel_softmax(x,
     # Add losses that prevent too few being used.
     distrib = tf.reshape(logsm, [-1, 2**z_size]) * maxvhot
     d_mean = tf.reduce_mean(distrib, axis=[0], keep_dims=True)
-    d_variance = tf.reduce_mean(tf.square(distrib - d_mean), axis=[0])
+    d_variance = tf.reduce_mean(
+        tf.squared_difference(distrib, d_mean), axis=[0])
     d_dev = -tf.reduce_mean(d_variance)
     ret = s
 
@@ -924,7 +925,7 @@ def vq_nearest_neighbor(x, means,
     x_means_hot = tf.one_hot(x_means_idx, bottleneck_size)
   x_means_hot_flat = tf.reshape(x_means_hot, [-1, bottleneck_size])
   x_means = tf.matmul(x_means_hot_flat, means)
-  e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
+  e_loss = tf.reduce_mean(tf.squared_difference(x, tf.stop_gradient(x_means)))
   return x_means_hot, e_loss, dist
 
 
@@ -1333,7 +1334,8 @@ def gumbel_softmax_discrete_bottleneck(x,
   x_means_assignments_flat = tf.reshape(x_means_assignments,
                                         [-1, bottleneck_size])
   x_means = tf.matmul(x_means_assignments_flat, means)
-  commitment_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
+  commitment_loss = tf.reduce_mean(
+      tf.squared_difference(x, tf.stop_gradient(x_means)))
 
   # Update the ema variables.
   updated_ema_count = moving_averages.assign_moving_average(
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 09e8b9b1c..a42a251aa 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -736,7 +736,8 @@ class VideoModalityL2(VideoModalityL1):
   """Modality for videos with L2 loss."""
 
   def internal_loss(self, logits, targets):
-    return tf.nn.relu((logits - targets)**2 - self.cutoff * self.cutoff)
+    return tf.nn.relu(
+        tf.squared_difference(logits, targets) - self.cutoff * self.cutoff)
 
 
 class VideoModalityL2Raw(VideoModalityL2):
@@ -916,7 +917,7 @@ def targets_bottom(self, x):
     return tf.to_float(x)
 
   def loss(self, body_output, targets):
-    loss = tf.square(body_output - tf.to_float(targets))
+    loss = tf.squared_difference(body_output, tf.to_float(targets))
     return tf.reduce_mean(loss), tf.constant(1.0)
 
 
diff --git a/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index f45289df2..344767084 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -138,8 +138,10 @@ def embedding_lookup(self, x, means):
         x_means_hot, [-1, self.hparams.num_blocks, self.hparams.block_v_size])
     x_means = tf.matmul(tf.transpose(x_means_hot_flat, perm=[1, 0, 2]), means)
     x_means = tf.transpose(x_means, [1, 0, 2])
-    q_loss = tf.reduce_mean(tf.square((tf.stop_gradient(x) - x_means)))
-    e_loss = tf.reduce_mean((x - tf.stop_gradient(x_means))**2)
+    q_loss = tf.reduce_mean(
+        tf.squared_difference(tf.stop_gradient(x), x_means))
+    e_loss = tf.reduce_mean(
+        tf.squared_difference(x, tf.stop_gradient(x_means)))
     return x_means_hot, x_means, q_loss, e_loss
 
   def bit_to_int(self, x_bit, num_bits, base=2):
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index b9c42e89e..0da05beba 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -219,7 +219,8 @@ def body(self, features):
         # minimized by just setting x=0 and b=0 -- so we don't want too much
         # of the influence of this, and we stop-gradient to not zero-out x.
         x_stop = tf.stop_gradient(x)
-        xb_loss = tf.reduce_mean(tf.reduce_sum(tf.square(x_stop - b), axis=-1))
+        xb_loss = tf.reduce_mean(tf.reduce_sum(
+            tf.squared_difference(x_stop, b), axis=-1))
         # To prevent this loss from exploding we clip at 1, but anneal clipping.
         clip_max = 1.0 / common_layers.inverse_exp_decay(
             warm_step, min_value=0.001)
diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 0b11e89cf..c6be947fc 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -65,7 +65,7 @@ def vq_nearest_neighbor(x, hparams):
     x_means_idx = tf.argmax(-dist, axis=-1)
     x_means_hot = tf.one_hot(x_means_idx, depth=bottleneck_size)
   x_means = tf.matmul(x_means_hot, means)
-  e_loss = tf.reduce_mean(tf.square(x - tf.stop_gradient(x_means)))
+  e_loss = tf.reduce_mean(tf.squared_difference(x, tf.stop_gradient(x_means)))
   return x_means_hot, e_loss
 
 
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 64aecd405..44e92c5ce 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -436,7 +436,8 @@ def ae_transformer_internal(inputs,
         losses["neg_q_entropy"] = neg_q_entropy * hparams.entropy_scale
       else:
         inputs_c = decode_transformer(inputs, ed, targets_c, hparams, "dec_c")
-        losses["latent_pred"] = tf.reduce_mean((inputs_c - targets_c)**2) * 20
+        losses["latent_pred"] = tf.reduce_mean(
+            tf.squared_difference(inputs_c, targets_c)) * 20
         def bn_inputs():
           with tf.variable_scope(tf.get_variable_scope(), reuse=True):
             bn, _, _, _, _ = hparams.bottleneck(
diff --git a/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index 95164462a..dbd4e5283 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -590,7 +590,8 @@ def mean_squared_error(true, pred):
   Returns:
     mean squared error between ground truth and predicted image.
   """
-  result = tf.reduce_sum(tf.square(true - pred)) / tf.to_float(tf.size(pred))
+  result = tf.reduce_sum(
+      tf.squared_difference(true, pred)) / tf.to_float(tf.size(pred))
   return result
 
 
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 829dcd607..0630b7ca4 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -364,7 +364,7 @@ def cv_squared(x):
   epsilon = 1e-10
   float_size = tf.to_float(tf.size(x)) + epsilon
   mean = tf.reduce_sum(x) / float_size
-  variance = tf.reduce_sum(tf.square(x - mean)) / float_size
+  variance = tf.reduce_sum(tf.squared_difference(x, mean)) / float_size
   return variance / (tf.square(mean) + epsilon)
 
 
diff --git a/tensor2tensor/utils/yellowfin.py b/tensor2tensor/utils/yellowfin.py
index 4d33c2aad..97acee4fb 100644
--- a/tensor2tensor/utils/yellowfin.py
+++ b/tensor2tensor/utils/yellowfin.py
@@ -392,7 +392,7 @@ def _get_lr_tensor(self):
     Returns:
       The lr_t.
     """
-    lr = (1.0 - tf.sqrt(self._mu))**2 / self._h_min
+    lr = tf.squared_difference(1.0, tf.sqrt(self._mu)) / self._h_min
     return lr
 
   def _get_mu_tensor(self):

From 80d7983e2464b38535a28174c6cf3d8667772e26 Mon Sep 17 00:00:00 2001
From: Yuwen Yan <ybbaigo@gmail.com>
Date: Wed, 30 Jan 2019 09:44:37 -0800
Subject: [PATCH 1607/2720] internal merge of PR #1421

PiperOrigin-RevId: 231611438
---
 tensor2tensor/models/__init__.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 77176a9e0..973c87a81 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -38,6 +38,7 @@
 from tensor2tensor.models import revnet
 from tensor2tensor.models import shake_shake
 from tensor2tensor.models import slicenet
+from tensor2tensor.models import text_cnn
 from tensor2tensor.models import transformer
 from tensor2tensor.models import vanilla_gan
 from tensor2tensor.models import xception

From da89c04bcd93a36b784abdc25bbd577a062a0b23 Mon Sep 17 00:00:00 2001
From: etragas-fathom <43351375+etragas-fathom@users.noreply.github.com>
Date: Wed, 30 Jan 2019 09:46:30 -0800
Subject: [PATCH 1608/2720] internal merge of PR #1416

PiperOrigin-RevId: 231611731
---
 tensor2tensor/layers/common_hparams.py                    | 2 ++
 tensor2tensor/models/research/universal_transformer.py    | 8 ++++++++
 .../models/research/universal_transformer_util.py         | 4 ++--
 tensor2tensor/utils/optimize.py                           | 7 +++++--
 4 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 0d5ce6598..d08ecbb52 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -75,6 +75,8 @@ def basic_params1():
       # Mixed precision training only supports exponential scaling currently
       # To disable the scaler, see to 0/False
       mixed_precision_optimizer_loss_scaler="exponential",
+      # Determines the initial loss scaling value for mixed precision
+      mixed_precision_optimizer_init_loss_scale=2**15,
       # Whether to zero gradients that were not computed, so that the
       # appropriate slots are created. Useful for sharing checkpoints between
       # models with different sets of heads.
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index a160e3627..6fc230fcb 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -469,6 +469,14 @@ def universal_transformer_big():
   return hparams
 
 
+@registry.register_hparams
+def universal_transformer_base_fp16():
+  hparams = transformer.transformer_base()
+  hparams = update_hparams_for_universal_transformer(hparams)
+  hparams.activation_dtype = "float16"
+  return hparams
+
+
 @registry.register_hparams
 def universal_transformer_small():
   hparams = transformer.transformer_base()
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 673747e60..54906e5b6 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -1442,7 +1442,7 @@ def add_position_timing_signal(x, step, hparams):
       length, channels, start_index=index)
 
   if hparams.add_or_concat_timing_signal == "add":
-    x_with_timing = x + signal
+    x_with_timing = x + common_layers.cast_like(signal, x)
 
   elif hparams.add_or_concat_timing_signal == "concat":
     batch_size = common_layers.shape_list(x)[0]
@@ -1479,7 +1479,7 @@ def add_step_timing_signal(x, step, hparams):
         channels, step, num_steps)
 
   if hparams.add_or_concat_timing_signal == "add":
-    x_with_timing = x + signal
+    x_with_timing = x + common_layers.cast_like(signal, x)
 
   elif hparams.add_or_concat_timing_signal == "concat":
     batch_size = common_layers.shape_list(x)[0]
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index ac301189b..9ff1993ad 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -204,9 +204,12 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
         raise ValueError("Mixed precision training only supports the "
                          "exponential loss scaler")
       else:
-        tf.logging.info("Using Exponential Update Loss Scaler")
+        tf.logging.info(
+            ("Using Exponential Update Loss Scaler with",
+             "init loss scale of {}".format(
+                 hparams.mixed_precision_optimizer_init_loss_scale)))
         manager = tf.contrib.mixed_precision.ExponentialUpdateLossScaleManager(
-            init_loss_scale=2**15,
+            init_loss_scale=hparams.mixed_precision_optimizer_init_loss_scale,
             incr_every_n_steps=2000,
             decr_every_n_nan_or_inf=2,
             incr_ratio=2,

From 96eb592953e56c3d089e7656a5bac702083dd669 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 30 Jan 2019 10:09:47 -0800
Subject: [PATCH 1609/2720] Fix return type of EnvProblem.seed -- don't think
 we'll use it much though.

PiperOrigin-RevId: 231616401
---
 tensor2tensor/envs/env_problem.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 08a9893b2..8f385af14 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -427,17 +427,19 @@ def unwrapped(self):
   def seed(self, seed=None):
     if not self._envs:
       tf.logging.info("`seed` called on non-existent envs, doing nothing.")
-      return
+      return None
 
     if not isinstance(self._envs, list):
       tf.logging.warning("`seed` called on non-list envs, doing nothing.")
-      return
+      return None
 
     tf.logging.warning(
         "Called `seed` on EnvProblem, calling seed on the underlying envs.")
     for env in self._envs:
       env.seed(seed)
 
+    return [seed]
+
   def close(self):
     if not self._envs:
       tf.logging.info("`close` called on non-existent envs, doing nothing.")

From f3b8784f67603ba7289550b09d03843254604672 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 30 Jan 2019 16:30:18 -0800
Subject: [PATCH 1610/2720] Added a test to env_problem_test.py, using
 `Problem.dataset`.

PiperOrigin-RevId: 231691270
---
 tensor2tensor/envs/env_problem_test.py | 139 +++++++++++++++++++++----
 1 file changed, 118 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_test.py b/tensor2tensor/envs/env_problem_test.py
index d83b7575b..92de34a58 100644
--- a/tensor2tensor/envs/env_problem_test.py
+++ b/tensor2tensor/envs/env_problem_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.envs import env_problem
+from tensor2tensor.layers import modalities
 import tensorflow as tf
 
 
@@ -100,7 +101,6 @@ def test_interaction_with_env(self):
         base_env_name="KellyCoinflip-v0",
         batch_size=batch_size,
         reward_range=reward_range)
-    ep.agent_id = "test"
 
     # Resets all environments.
     ep.reset()
@@ -173,39 +173,66 @@ def read_tfrecord_dataset(self, filenames, ep):
 
     return num_trajectories, num_time_steps
 
-  def test_generate_data(self):
-    batch_size = 5
-    reward_range = (-1, 1)
-    ep = env_problem.EnvProblem(
-        base_env_name="CartPole-v0",
-        batch_size=batch_size,
-        reward_range=reward_range)
-    ep.agent_id = "test"
-
-    # Set this in the test to test things, but usually registered subclasses
-    # will set this.
-    ep.name = "CartPoleProblem"
+  def play_env(self,
+               env=None,
+               nsteps=100,
+               base_env_name=None,
+               batch_size=5,
+               reward_range=None):
+    """Creates `EnvProblem` with the given arguments and plays it randomly.
+
+    Args:
+      env: optional env.
+      nsteps: plays the env randomly for nsteps.
+      base_env_name: passed to EnvProblem's init.
+      batch_size: passed to EnvProblem's init.
+      reward_range: passed to EnvProblem's init.
+
+    Returns:
+      tuple of env_problem, number of trajectories done, number of trajectories
+      done in the last step.
+    """
+
+    if env is None:
+      env = env_problem.EnvProblem(
+          base_env_name=base_env_name,
+          batch_size=batch_size,
+          reward_range=reward_range)
+      # Usually done by a registered subclass, we do this manually in the test.
+      env.name = base_env_name
 
     # Reset all environments.
-    ep.reset()
+    env.reset()
 
     # Play for some steps to generate data.
-    nsteps = 100
     num_dones = 0
     num_dones_in_last_step = 0
     for _ in range(nsteps):
       # Sample actions.
-      actions = np.stack([ep.action_space.sample() for _ in range(batch_size)])
+      actions = np.stack([env.action_space.sample() for _ in range(batch_size)])
       # Step through it.
-      _, _, dones, _ = ep.step(actions)
+      _, _, dones, _ = env.step(actions)
       # Get the indices where we are done ...
       done_indices = env_problem.EnvProblem.done_indices(dones)
       # ... and reset those.
-      ep.reset(indices=done_indices)
+      env.reset(indices=done_indices)
       # count the number of dones we got, in this step and overall.
       num_dones_in_last_step = sum(dones)
       num_dones += num_dones_in_last_step
 
+    return env, num_dones, num_dones_in_last_step
+
+  def test_generate_data(self):
+    base_env_name = "CartPole-v0"
+    batch_size = 5
+    reward_range = (-1, 1)
+    nsteps = 100
+    ep, num_dones, num_dones_in_last_step = self.play_env(
+        base_env_name=base_env_name,
+        batch_size=batch_size,
+        reward_range=reward_range,
+        nsteps=nsteps)
+
     # This is because every num_dones starts a new trajectory, and a further
     # batch_size are active at the last step when we call generate_data, but
     # the ones that got done in the last step (these have only one time-step in
@@ -213,9 +240,6 @@ def test_generate_data(self):
     expected_num_trajectories = num_dones + batch_size - num_dones_in_last_step
 
     # Similar logic as above, nsteps * batch_size overall `step` calls are made.
-    # However, if a `step` call completes a trajectory, one more time-step is
-    # added, but we have to discount the trajectories that got done in the very
-    # last step.
     expected_num_time_steps = (
         nsteps * batch_size) + num_dones + batch_size - num_dones_in_last_step
 
@@ -239,6 +263,79 @@ def test_generate_data(self):
     self.assertEqual(expected_num_trajectories,
                      training_trajectories + dev_trajectories)
 
+  def test_problem_dataset_works(self):
+
+    # We need to derive this class to set the required methods.
+    class TestEnv(env_problem.EnvProblem):
+      name = "TestEnv"
+
+      @property
+      def input_modality(self):
+        return modalities.ModalityType.REAL_L2_LOSS
+
+      @property
+      def input_vocab_size(self):
+        return None
+
+      @property
+      def target_modality(self):
+        return modalities.ModalityType.SYMBOL_WEIGHTS_ALL
+
+      @property
+      def target_vocab_size(self):
+        return 2
+
+    base_env_name = "CartPole-v0"
+    batch_size = 5
+    reward_range = (-1, 1)
+
+    env = TestEnv(
+        base_env_name=base_env_name,
+        batch_size=batch_size,
+        reward_range=reward_range)
+
+    nsteps = 100
+    ep, _, _ = self.play_env(env=env, nsteps=nsteps)
+
+    # Dump the completed data to disk.
+    ep.generate_data(self.tmp_dir, self.tmp_dir)
+
+    # Read the actual files and count the trajectories and time-steps.
+    dev_filenames = ep.dev_filepaths(
+        self.tmp_dir, ep.num_shards[problem.DatasetSplit.EVAL], True)
+    dev_trajectories, dev_timesteps = self.read_tfrecord_dataset(
+        dev_filenames, ep)
+
+    # Count them using a tf.data.Dataset.
+    dev_dataset = ep.dataset(tf.estimator.ModeKeys.EVAL, data_dir=self.tmp_dir)
+
+    last_timestep = -1
+    dev_timesteps_ds = 0
+    dev_trajectories_ds = 0
+    iterator = dev_dataset.make_one_shot_iterator()
+    next_element = iterator.get_next()
+    with tf.Session() as session:
+      while True:
+        try:
+          tf_example_dict = session.run(next_element)
+
+          # We have a time-step.
+          dev_timesteps_ds += 1
+
+          this_timestep = tf_example_dict[env_problem.TIMESTEP_FIELD][
+              0]  # [0] since every value in tf_example_dict is an array/list.
+          if 1 + last_timestep != this_timestep:
+            dev_trajectories_ds += 1
+            self.assertEqual(0, this_timestep)
+          last_timestep = this_timestep
+        except tf.errors.OutOfRangeError:
+          dev_trajectories_ds += 1
+          break
+
+    # Make sure that they agree.
+    self.assertEqual(dev_trajectories, dev_trajectories_ds)
+    self.assertEqual(dev_timesteps, dev_timesteps_ds)
+
 
 if __name__ == "__main__":
   tf.test.main()

From 508dc0680a953805b66c44e55ca4d45e81655db2 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 31 Jan 2019 11:18:46 -0800
Subject: [PATCH 1611/2720] Setting to run planner on Atari games.

PiperOrigin-RevId: 231820237
---
 tensor2tensor/rl/evaluator.py | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index a842d609f..5d574aa45 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -55,7 +55,7 @@
     "eval_metrics_dir", "", "Directory to output the eval metrics at."
 )
 flags.DEFINE_integer("eval_batch_size", 64, "Number of games to evaluate.")
-flags.DEFINE_integer("eval_step_limit", 100000,
+flags.DEFINE_integer("eval_step_limit", 50000,
                      "Maximum number of time steps, ignored if -1.")
 flags.DEFINE_enum(
     "agent", "policy", ["random", "policy", "planner"], "Agent type to use."
@@ -82,7 +82,7 @@
 )
 flags.DEFINE_string("planner_hparams", "", "Planner hparam overrides.")
 flags.DEFINE_integer(
-    "log_every_steps", 20, "Log every how many environment steps."
+    "log_every_steps", 5, "Log every how many environment steps."
 )
 flags.DEFINE_string(
     "debug_video_path", "", "Path to save the debug video at."
@@ -227,6 +227,15 @@ def planner_guess9():
   return hparams
 
 
+@registry.register_hparams
+def planner_guess0():
+  hparams = planner_base()
+  hparams.uct_const = 6
+  hparams.num_rollouts = 4 * 96
+  hparams.normalizer_window_size = 30
+  return hparams
+
+
 def make_env(env_type, real_env, sim_env_kwargs):
   """Factory function for envs."""
   return {
@@ -484,6 +493,7 @@ def main(_):
   policy_dir = FLAGS.policy_dir
   model_dir = FLAGS.model_dir
   eval_metrics_dir = FLAGS.eval_metrics_dir
+  debug_video_path = FLAGS.debug_video_path
   if FLAGS.output_dir:
     cur_dir = FLAGS.output_dir
     if FLAGS.total_num_workers > 1:
@@ -492,8 +502,9 @@ def main(_):
     model_dir = os.path.join(cur_dir, "world_model")
     eval_dir_basename = "evaluator_"
     if FLAGS.agent == "planner":
-      eval_dir_basename = "planner_"
+      eval_dir_basename = FLAGS.planner_hparams_set + "_"
     eval_metrics_dir = os.path.join(cur_dir, eval_dir_basename + now_tag)
+    debug_video_path = eval_metrics_dir
     tf.logging.info("Writing metrics to %s." % eval_metrics_dir)
     if not tf.gfile.Exists(eval_metrics_dir):
       tf.gfile.MkDir(eval_metrics_dir)
@@ -501,7 +512,7 @@ def main(_):
       loop_hparams, planner_hparams, policy_dir, model_dir,
       eval_metrics_dir, FLAGS.agent, FLAGS.mode, FLAGS.eval_with_learner,
       FLAGS.log_every_steps if FLAGS.log_every_steps > 0 else None,
-      debug_video_path=FLAGS.debug_video_path,
+      debug_video_path=debug_video_path,
       num_debug_videos=FLAGS.num_debug_videos,
       random_starts_step_limit=FLAGS.random_starts_step_limit,
   )

From b10595aa31c2527c37e343ae1a4c72fad0b3969f Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 31 Jan 2019 13:34:17 -0800
Subject: [PATCH 1612/2720] Modality cleanup in EnvProblem

PiperOrigin-RevId: 231843860
---
 tensor2tensor/envs/env_problem.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 8f385af14..34be5602f 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -615,12 +615,12 @@ def hparams(self, defaults, model_hparams):
     p.modality.update({
         "inputs": self.input_modality,
         "targets": self.target_modality,
-        "input_reward": modalities.SymbolModalityWeightsAll,
-        "target_reward": modalities.SymbolModalityWeightsAll,
-        "input_action": modalities.SymbolModalityWeightsAll,
-        "target_action": modalities.SymbolModalityWeightsAll,
-        "target_policy": modalities.IdentityModality,
-        "target_value": modalities.IdentityModality,
+        "input_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "target_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "input_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "target_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "target_policy": modalities.ModalityType.IDENTITY,
+        "target_value": modalities.ModalityType.IDENTITY,
     })
 
     p.vocab_size.update({

From d52c05915fcd68e13c69a25b924c0c7f8245d240 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 31 Jan 2019 13:37:19 -0800
Subject: [PATCH 1613/2720] Remove the sqrt from the variance of the linear
 kernel.

PiperOrigin-RevId: 231844521
---
 tensor2tensor/layers/bayes.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 93a60e684..fd482e370 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -117,7 +117,7 @@ def __call__(self, x1, x2):
     encoded_x1 = self.encoder(x1)
     encoded_x2 = self.encoder(x2)
     dot_product = tf.matmul(encoded_x1, encoded_x2, transpose_b=True)
-    return tf.sqrt(self.variance) * dot_product + self.bias
+    return self.variance * dot_product + self.bias
 
   def get_config(self):
     return {

From f74837f2a41ab8c06bd4ab3f8eed83bd9bb065b6 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 31 Jan 2019 13:40:15 -0800
Subject: [PATCH 1614/2720] Minor cleanup in trainer_model_free.py

PiperOrigin-RevId: 231845132
---
 tensor2tensor/rl/trainer_model_free.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 3607f96f8..b3d6ea238 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -57,17 +57,15 @@ def initialize_env_specs(hparams):
                            hparams.eval_max_num_noops,
                            hparams.rl_env_max_episode_steps,
                            env_name=hparams.rl_env_name)
+
   env.start_new_epoch(0)
 
-  # TODO(afrozm): Decouple env_fn from hparams and return both, is there
-  # even a need to return hparams? Just return the env_fn?
-  hparams.add_hparam("env_fn", rl.make_real_env_fn(env))
-  return hparams
+  return rl.make_real_env_fn(env)
 
 
 def train(hparams, output_dir, report_fn=None):
   """Train."""
-  hparams = initialize_env_specs(hparams)
+  env_fn = initialize_env_specs(hparams)
 
   tf.logging.vlog(1, "HParams in trainer_model_free.train : %s",
                   misc_utils.pprint_hparams(hparams))
@@ -119,7 +117,7 @@ def train(hparams, output_dir, report_fn=None):
     tf.logging.info("Starting training iteration [%d] for [%d] steps.", i, step)
 
     policy_hparams.epochs_num = eval_every_epochs
-    learner.train(hparams.env_fn,
+    learner.train(env_fn,
                   policy_hparams,
                   simulated=False,
                   save_continuously=True,

From e06634e0e6ca3f9f1610e5c1daa7f70dae748a80 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 31 Jan 2019 19:28:12 -0800
Subject: [PATCH 1615/2720] Add a new registry for `EnvProblem`s, this is
 because some extra initialization is needed to initialize an `EnvProblem`.

PiperOrigin-RevId: 231900357
---
 tensor2tensor/utils/registry.py      | 26 ++++++++++++++++++++++++
 tensor2tensor/utils/registry_test.py | 30 ++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index d02a3c99d..ed20b1c91 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -423,6 +423,8 @@ def __init__(self):
           2, "Registered layer functions must take exaction two arguments: "
           "hparams (HParams) and prefix (str)."))
 
+  env_problems = Registry("env_problems", on_set=_on_problem_set)
+
 
 # consistent version of old API
 model = Registries.models.__getitem__
@@ -467,6 +469,9 @@ def optimizer(name):
 hparams = Registries.hparams.__getitem__
 register_hparams = Registries.hparams.register
 
+list_env_problems = lambda: sorted(Registries.env_problems)
+register_env_problem = Registries.env_problems.register
+
 
 def list_hparams(prefix=None):
   hp_names = sorted(Registries.hparams)
@@ -503,6 +508,23 @@ def problem(problem_name):
       was_copy=spec.was_copy, was_reversed=spec.was_reversed)
 
 
+def env_problem(env_problem_name, batch_size):
+  """Get and initialize the `EnvProblem` with the given name and batch size.
+
+  Args:
+    env_problem_name: string name of the registered env problem.
+    batch_size: batch_size to initialize the env problem with.
+
+  Returns:
+    an initialized EnvProblem with the given batch size.
+  """
+
+  ep_cls = Registries.env_problems[env_problem_name]
+  ep = ep_cls()
+  ep.initialize(batch_size=batch_size)
+  return ep
+
+
 attack = Registries.attacks.__getitem__
 list_attacks = lambda: sorted(Registries.attacks)
 register_attack = Registries.attacks.register
@@ -573,6 +595,9 @@ def help_string():
 
   Pruning Strategies:
 %s
+
+  Env Problems:
+%s
 """
   lists = tuple(
       display_list_by_prefix(entries, starting_spaces=4) for entries in [  # pylint: disable=g-complex-comprehension
@@ -585,5 +610,6 @@ def help_string():
           list_attack_params(),
           list_pruning_params(),
           list_pruning_strategies(),
+          list_env_problems(),
       ])
   return help_str % lists
diff --git a/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index e21536785..1864881e2 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -90,6 +90,36 @@ def testGet(self):
     self.assertEqual(r.get("b", 3), 3)
 
 
+class EnvProblemRegistryTest(tf.test.TestCase):
+
+  def setUp(self):
+    registry.Registries.env_problems._clear()
+
+  def testEnvProblem(self):
+    # Register this class and expect to get it back.
+
+    @registry.register_env_problem
+    class EnvProb(object):
+
+      batch_size = None
+
+      def initialize(self, batch_size):
+        self.batch_size = batch_size
+
+    # Get it with given batch_size.
+    batch_size = 100
+    ep = registry.env_problem("env_prob", batch_size)
+
+    # name property is set.
+    self.assertEqual("env_prob", ep.name)
+
+    # initialize was called and therefore batch_size was set.
+    self.assertEqual(batch_size, ep.batch_size)
+
+    # assert on the type.
+    self.assertTrue(isinstance(ep, EnvProb))
+
+
 class ModelRegistryTest(tf.test.TestCase):
 
   def setUp(self):

From 9e0a894034d8090892c238df1bd9bd3180c2b9a3 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 31 Jan 2019 19:45:17 -0800
Subject: [PATCH 1616/2720] Changes to EnvProblem for better registration and
 initialization.   - __init__ is now light-weight, just setting attributes to
 given values or     None.   - Subclasses just set the appropriate attributes.
   - initialize has just one argument `batch_size`, which then initializes
 envs     and derived attributes and properties.   - convenience function to
 print observation of all the batch.

Added a TicTacToeEnvProblem that is a subclass of EnvProblem and wraps TicTacToe.
  - Tests

PiperOrigin-RevId: 231901996
---
 tensor2tensor/envs/__init__.py                |  2 +
 tensor2tensor/envs/env_problem.py             | 50 +++++++------
 tensor2tensor/envs/tic_tac_toe_env.py         |  5 +-
 tensor2tensor/envs/tic_tac_toe_env_problem.py | 52 +++++++++++++
 .../envs/tic_tac_toe_env_problem_test.py      | 74 +++++++++++++++++++
 5 files changed, 160 insertions(+), 23 deletions(-)
 create mode 100644 tensor2tensor/envs/tic_tac_toe_env_problem.py
 create mode 100644 tensor2tensor/envs/tic_tac_toe_env_problem_test.py

diff --git a/tensor2tensor/envs/__init__.py b/tensor2tensor/envs/__init__.py
index d5201e5d7..0d3981f49 100644
--- a/tensor2tensor/envs/__init__.py
+++ b/tensor2tensor/envs/__init__.py
@@ -19,5 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import tic_tac_toe_env
+from tensor2tensor.envs import tic_tac_toe_env_problem
 
diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 34be5602f..015341546 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -126,10 +126,8 @@ class EnvProblem(Env, problem.Problem):
   def __init__(self,
                base_env_name=None,
                base_env_kwargs=None,
-               batch_size=1,
-               reward_range=(-np.inf, np.inf),
-               was_reversed=False,
-               was_copy=False):
+               batch_size=None,
+               reward_range=(-np.inf, np.inf)):
     """Initializes this class by creating the envs and managing trajectories.
 
     Args:
@@ -137,20 +135,15 @@ def __init__(self,
         underlying environment.
       base_env_kwargs: (dict) passed to `gym_utils.make_gym_env` to make the
         underlying environment.
-      batch_size: (int): How many envs to make in the non natively batched mode.
+      batch_size: (int or None): How many envs to make in the non natively
+        batched mode.
       reward_range: (tuple(number, number)) the first element is the minimum
         reward and the second is the maximum reward, used to clip and process
         the raw reward in `process_rewards`.
-      was_reversed: (bool) should be false, passed to underlying init of
-        Problem.
-      was_copy: (bool) should be false, passed to underlying init of Problem.
     """
 
-    # Assert on these since they don't make sense.
-    assert not was_reversed and not was_copy
-
     # Call the super's ctor.
-    problem.Problem.__init__(self, was_reversed=was_reversed, was_copy=was_copy)
+    problem.Problem.__init__(self, was_reversed=False, was_copy=False)
 
     # Name for the base environment, will be used in `gym_utils.make_gym_env` in
     # the default implementation of `initialize_environments`.
@@ -188,15 +181,8 @@ def __init__(self,
     # and also the ones that are completed, i.e. done.
     self._trajectories = None
 
-    self.initialize_environments(batch_size=batch_size)
-
-    # Assert that *all* the above are now set, we should do this since
-    # subclasses can override `initialize_environments`.
-    assert self._envs is not None
-    assert self._observation_space is not None
-    assert self._action_space is not None
-    assert self._reward_range is not None
-    assert self._trajectories is not None
+    if batch_size is not None:
+      self.initialize(batch_size=batch_size)
 
   @property
   def base_env_name(self):
@@ -244,6 +230,17 @@ def _verify_same_spaces(self):
         tf.logging.error("Env[%d] has action space [%s]", i, env.action_space)
       raise ValueError(err_str)
 
+  def initialize(self, batch_size=1):
+    self.initialize_environments(batch_size=batch_size)
+
+    # Assert that *all* the above are now set, we should do this since
+    # subclasses can override `initialize_environments`.
+    assert self._envs is not None
+    assert self._observation_space is not None
+    assert self._action_space is not None
+    assert self._reward_range is not None
+    assert self._trajectories is not None
+
   def initialize_environments(self, batch_size=1):
     """Initializes the environments and trajectories.
 
@@ -358,6 +355,8 @@ def is_reward_range_finite(self):
   def process_rewards(self, rewards):
     """Clips, rounds, adds the min_reward and changes to integer type.
 
+    The result of the above is that the new minimum is 0.
+
     Args:
       rewards: numpy array of raw (float) rewards.
 
@@ -758,3 +757,12 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
       # `cycle_every_n` isn't needed since file list given to it is a singleton.
       generator_utils.generate_files(
           self._generate_time_steps(trajectories_to_write), [f])
+
+  def print_state(self):
+    for t in self.trajectories.trajectories:
+      print("---------")
+      if not t.is_active():
+        print("trajectory isn't active.")
+        continue
+      last_obs = t.last_time_step().observation
+      print(str(last_obs))
diff --git a/tensor2tensor/envs/tic_tac_toe_env.py b/tensor2tensor/envs/tic_tac_toe_env.py
index 43e603a44..7f54aac48 100644
--- a/tensor2tensor/envs/tic_tac_toe_env.py
+++ b/tensor2tensor/envs/tic_tac_toe_env.py
@@ -192,8 +192,9 @@ def hparams(self, defaults, unused_model_hparams):
         "targets": modalities.ModalityType.IDENTITY_SYMBOL,
     }
     p.vocab_size = {
-        "inputs": 3,
-        "targets": 3,
+        "inputs": 3,  # since at each box, the input is either x, o or -.
+        # nevermind that we have a 3x3 box.
+        "targets": 3,  # -1, 0, 1
     }
     p.input_space_id = 0  # problem.SpaceID.GENERIC
     p.target_space_id = 0  # problem.SpaceID.GENERIC
diff --git a/tensor2tensor/envs/tic_tac_toe_env_problem.py b/tensor2tensor/envs/tic_tac_toe_env_problem.py
new file mode 100644
index 000000000..0d748f82a
--- /dev/null
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem.py
@@ -0,0 +1,52 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""TicTacToeEnvProblem wraps the TicTacToeEnv in an EnvProblem."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.envs import env_problem
+from tensor2tensor.layers import modalities
+from tensor2tensor.utils import registry
+
+
+@registry.register_env_problem
+class TicTacToeEnvProblem(env_problem.EnvProblem):
+  """Plays `batch_size` games of tic-tac-toe."""
+
+  def __init__(self):
+    super(TicTacToeEnvProblem, self).__init__(
+        base_env_name="T2TEnv-TicTacToeEnv-v0",
+        reward_range=(-1, 1))
+
+  @property
+  def input_modality(self):
+    return modalities.ModalityType.IDENTITY_SYMBOL
+
+  @property
+  def input_vocab_size(self):
+    # Since a box can be either x or o or empty.
+    return 3
+
+  @property
+  def target_modality(self):
+    return modalities.ModalityType.IDENTITY_SYMBOL
+
+  @property
+  def target_vocab_size(self):
+    # Since reward is either -1 or 0 or +1.
+    return 3
diff --git a/tensor2tensor/envs/tic_tac_toe_env_problem_test.py b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
new file mode 100644
index 000000000..954be159a
--- /dev/null
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
@@ -0,0 +1,74 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.envs.tic_tac_toe_env_problem."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.envs import tic_tac_toe_env  # pylint: disable=unused-import
+from tensor2tensor.envs import tic_tac_toe_env_problem  # pylint: disable=unused-import
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+
+class TicTacToeEnvProblemTest(tf.test.TestCase):
+
+  def test_registration_and_interaction_with_env_problem(self):
+    batch_size = 5
+    # This ensures that registration has occurred.
+    ep = registry.env_problem("tic_tac_toe_env_problem", batch_size)
+    ep.reset()
+    num_done, num_lost, num_won, num_draw = 0, 0, 0, 0
+    nsteps = 100
+    for _ in range(nsteps):
+      actions = np.stack([ep.action_space.sample() for _ in range(batch_size)])
+      obs, rewards, dones, infos = ep.step(actions)
+
+      # Assert that things are happening batchwise.
+      self.assertEqual(batch_size, len(obs))
+      self.assertEqual(batch_size, len(rewards))
+      self.assertEqual(batch_size, len(dones))
+      self.assertEqual(batch_size, len(infos))
+
+      done_indices = ep.done_indices(dones)
+      ep.reset(done_indices)
+      num_done += sum(dones)
+      for r, d in zip(rewards, dones):
+        if not d:
+          continue
+        # NOTE: r is 0, 1, 2 because the default EnvProblem.process_rewards
+        # shifts the rewards so that min is 0.
+        if r == 0:
+          num_lost += 1
+        elif r == 1:
+          num_draw += 1
+        elif r == 2:
+          num_won += 1
+        else:
+          raise ValueError("reward should be 0, 1, 2 but is {}".format(r))
+
+    # Assert that something got done atleast, without that the next assert is
+    # meaningless.
+    self.assertGreater(num_done, 0)
+
+    # Assert that things are consistent.
+    self.assertEqual(num_done, num_won + num_lost + num_draw)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 211bf893f79330a369fe3819f2c8619e8e2efe18 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 1 Feb 2019 08:52:14 -0800
Subject: [PATCH 1617/2720] Fix convergence issue of t2t models on DF pods.
 Instead of using hardcoded num_shards_per_host of 8 (which is wrong for DF
 pods), directly using num_hosts for num_partitions.

PiperOrigin-RevId: 231978646
---
 tensor2tensor/data_generators/problem.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index e3a38f9fb..7cff2e202 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -778,7 +778,7 @@ def estimator_input_fn(params, config):
 
     return estimator_input_fn
 
-  def _dataset_partition(self, mode, config):
+  def _dataset_partition(self, mode, config, params):
     """Which part of the training data to read.
 
     If there are multiple parallel calls to input_fn (multiple TPU hosts),
@@ -788,6 +788,7 @@ def _dataset_partition(self, mode, config):
     Args:
       mode: tf.estimator.ModeKeys
       config: RunConfig
+      params: A dict that contains parameters.
     Returns:
       partition_id: an integer
       num_partitions: an integer
@@ -802,7 +803,7 @@ def _dataset_partition(self, mode, config):
         phift == tpu_config.InputPipelineConfig.BROADCAST):
       return 0, 1
     if phift:
-      num_partitions = max(config.tpu_config.num_shards // 8, 1)
+      num_partitions = max(params["context"].num_hosts, 1)
     else:
       num_partitions = config.tpu_config.num_shards
     partition_id = getattr(self, "_next_partition_id", 0)
@@ -839,7 +840,7 @@ def input_fn(self,
     Returns:
       (features_dict<str name, Tensor feature>, Tensor targets)
     """
-    partition_id, num_partitions = self._dataset_partition(mode, config)
+    partition_id, num_partitions = self._dataset_partition(mode, config, params)
     is_training = mode == tf.estimator.ModeKeys.TRAIN
     if config and config.use_tpu:
       num_threads = 64

From d45060a27108055e9b0c906da1ee82a64fb36d49 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 1 Feb 2019 10:19:15 -0800
Subject: [PATCH 1618/2720] Fork HParams and put it into t2t/utils. Changes to
 use it come later.

PiperOrigin-RevId: 231993516
---
 tensor2tensor/utils/hparam.py      | 643 +++++++++++++++++++++++++++++
 tensor2tensor/utils/hparam_test.py | 509 +++++++++++++++++++++++
 2 files changed, 1152 insertions(+)
 create mode 100644 tensor2tensor/utils/hparam.py
 create mode 100644 tensor2tensor/utils/hparam_test.py

diff --git a/tensor2tensor/utils/hparam.py b/tensor2tensor/utils/hparam.py
new file mode 100644
index 000000000..ed7f02ded
--- /dev/null
+++ b/tensor2tensor/utils/hparam.py
@@ -0,0 +1,643 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+# Forked with minor changes from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/training/python/training/hparam.py  pylint: disable=line-too-long
+"""Hyperparameter values."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import json
+import numbers
+import re
+import six
+
+# Define the regular expression for parsing a single clause of the input
+# (delimited by commas).  A legal clause looks like:
+#   <variable name>[<index>]? = <rhs>
+# where <rhs> is either a single token or [] enclosed list of tokens.
+# For example:  "var[1] = a" or "x = [1,2,3]"
+PARAM_RE = re.compile(r"""
+  (?P<name>[a-zA-Z][\w\.]*)      # variable name: "var" or "x"
+  (\[\s*(?P<index>\d+)\s*\])?  # (optional) index: "1" or None
+  \s*=\s*
+  ((?P<val>[^,\[]*)            # single value: "a" or None
+   |
+   \[(?P<vals>[^\]]*)\])       # list of values: None or "1,2,3"
+  ($|,\s*)""", re.VERBOSE)
+
+
+def _parse_fail(name, var_type, value, values):
+  """Helper function for raising a value error for bad assignment."""
+  raise ValueError(
+      'Could not parse hparam \'%s\' of type \'%s\' with value \'%s\' in %s' %
+      (name, var_type.__name__, value, values))
+
+
+def _reuse_fail(name, values):
+  """Helper function for raising a value error for reuse of name."""
+  raise ValueError('Multiple assignments to variable \'%s\' in %s' % (name,
+                                                                      values))
+
+
+def _process_scalar_value(name, parse_fn, var_type, m_dict, values,
+                          results_dictionary):
+  """Update results_dictionary with a scalar value.
+
+  Used to update the results_dictionary to be returned by parse_values when
+  encountering a clause with a scalar RHS (e.g.  "s=5" or "arr[0]=5".)
+
+  Mutates results_dictionary.
+
+  Args:
+    name: Name of variable in assignment ("s" or "arr").
+    parse_fn: Function for parsing the actual value.
+    var_type: Type of named variable.
+    m_dict: Dictionary constructed from regex parsing.
+      m_dict['val']: RHS value (scalar)
+      m_dict['index']: List index value (or None)
+    values: Full expression being parsed
+    results_dictionary: The dictionary being updated for return by the parsing
+      function.
+
+  Raises:
+    ValueError: If the name has already been used.
+  """
+  try:
+    parsed_value = parse_fn(m_dict['val'])
+  except ValueError:
+    _parse_fail(name, var_type, m_dict['val'], values)
+
+  # If no index is provided
+  if not m_dict['index']:
+    if name in results_dictionary:
+      _reuse_fail(name, values)
+    results_dictionary[name] = parsed_value
+  else:
+    if name in results_dictionary:
+      # The name has already been used as a scalar, then it
+      # will be in this dictionary and map to a non-dictionary.
+      if not isinstance(results_dictionary.get(name), dict):
+        _reuse_fail(name, values)
+    else:
+      results_dictionary[name] = {}
+
+    index = int(m_dict['index'])
+    # Make sure the index position hasn't already been assigned a value.
+    if index in results_dictionary[name]:
+      _reuse_fail('{}[{}]'.format(name, index), values)
+    results_dictionary[name][index] = parsed_value
+
+
+def _process_list_value(name, parse_fn, var_type, m_dict, values,
+                        results_dictionary):
+  """Update results_dictionary from a list of values.
+
+  Used to update results_dictionary to be returned by parse_values when
+  encountering a clause with a list RHS (e.g.  "arr=[1,2,3]".)
+
+  Mutates results_dictionary.
+
+  Args:
+    name: Name of variable in assignment ("arr").
+    parse_fn: Function for parsing individual values.
+    var_type: Type of named variable.
+    m_dict: Dictionary constructed from regex parsing.
+      m_dict['val']: RHS value (scalar)
+    values: Full expression being parsed
+    results_dictionary: The dictionary being updated for return by the parsing
+      function.
+
+  Raises:
+    ValueError: If the name has an index or the values cannot be parsed.
+  """
+  if m_dict['index'] is not None:
+    raise ValueError('Assignment of a list to a list index.')
+  elements = filter(None, re.split('[ ,]', m_dict['vals']))
+  # Make sure the name hasn't already been assigned a value
+  if name in results_dictionary:
+    raise _reuse_fail(name, values)
+  try:
+    results_dictionary[name] = [parse_fn(e) for e in elements]
+  except ValueError:
+    _parse_fail(name, var_type, m_dict['vals'], values)
+
+
+def _cast_to_type_if_compatible(name, param_type, value):
+  """Cast hparam to the provided type, if compatible.
+
+  Args:
+    name: Name of the hparam to be cast.
+    param_type: The type of the hparam.
+    value: The value to be cast, if compatible.
+
+  Returns:
+    The result of casting `value` to `param_type`.
+
+  Raises:
+    ValueError: If the type of `value` is not compatible with param_type.
+      * If `param_type` is a string type, but `value` is not.
+      * If `param_type` is a boolean, but `value` is not, or vice versa.
+      * If `param_type` is an integer type, but `value` is not.
+      * If `param_type` is a float type, but `value` is not a numeric type.
+  """
+  fail_msg = (
+      "Could not cast hparam '%s' of type '%s' from value %r" %
+      (name, param_type, value))
+
+  # Some callers use None, for which we can't do any casting/checking. :(
+  if issubclass(param_type, type(None)):
+    return value
+
+  # Avoid converting a non-string type to a string.
+  if (issubclass(param_type, (six.string_types, six.binary_type)) and
+      not isinstance(value, (six.string_types, six.binary_type))):
+    raise ValueError(fail_msg)
+
+  # Avoid converting a number or string type to a boolean or vice versa.
+  if issubclass(param_type, bool) != isinstance(value, bool):
+    raise ValueError(fail_msg)
+
+  # Avoid converting float to an integer (the reverse is fine).
+  if (issubclass(param_type, numbers.Integral) and
+      not isinstance(value, numbers.Integral)):
+    raise ValueError(fail_msg)
+
+  # Avoid converting a non-numeric type to a numeric type.
+  if (issubclass(param_type, numbers.Number) and
+      not isinstance(value, numbers.Number)):
+    raise ValueError(fail_msg)
+
+  return param_type(value)
+
+
+def parse_values(values, type_map, ignore_unknown=False):
+  """Parses hyperparameter values from a string into a python map.
+
+  `values` is a string containing comma-separated `name=value` pairs.
+  For each pair, the value of the hyperparameter named `name` is set to
+  `value`.
+
+  If a hyperparameter name appears multiple times in `values`, a ValueError
+  is raised (e.g. 'a=1,a=2', 'a[1]=1,a[1]=2').
+
+  If a hyperparameter name in both an index assignment and scalar assignment,
+  a ValueError is raised.  (e.g. 'a=[1,2,3],a[0] = 1').
+
+  The hyperparameter name may contain '.' symbols, which will result in an
+  attribute name that is only accessible through the getattr and setattr
+  functions.  (And must be first explicit added through add_hparam.)
+
+  WARNING: Use of '.' in your variable names is allowed, but is not well
+  supported and not recommended.
+
+  The `value` in `name=value` must follows the syntax according to the
+  type of the parameter:
+
+  *  Scalar integer: A Python-parsable integer point value.  E.g.: 1,
+     100, -12.
+  *  Scalar float: A Python-parsable floating point value.  E.g.: 1.0,
+     -.54e89.
+  *  Boolean: Either true or false.
+  *  Scalar string: A non-empty sequence of characters, excluding comma,
+     spaces, and square brackets.  E.g.: foo, bar_1.
+  *  List: A comma separated list of scalar values of the parameter type
+     enclosed in square brackets.  E.g.: [1,2,3], [1.0,1e-12], [high,low].
+
+  When index assignment is used, the corresponding type_map key should be the
+  list name.  E.g. for "arr[1]=0" the type_map must have the key "arr" (not
+  "arr[1]").
+
+  Args:
+    values: String.  Comma separated list of `name=value` pairs where
+      'value' must follow the syntax described above.
+    type_map: A dictionary mapping hyperparameter names to types.  Note every
+      parameter name in values must be a key in type_map.  The values must
+      conform to the types indicated, where a value V is said to conform to a
+      type T if either V has type T, or V is a list of elements of type T.
+      Hence, for a multidimensional parameter 'x' taking float values,
+      'x=[0.1,0.2]' will parse successfully if type_map['x'] = float.
+    ignore_unknown: Bool. Whether values that are missing a type in type_map
+      should be ignored. If set to True, a ValueError will not be raised for
+      unknown hyperparameter type.
+
+  Returns:
+    A python map mapping each name to either:
+    * A scalar value.
+    * A list of scalar values.
+    * A dictionary mapping index numbers to scalar values.
+    (e.g. "x=5,L=[1,2],arr[1]=3" results in {'x':5,'L':[1,2],'arr':{1:3}}")
+
+  Raises:
+    ValueError: If there is a problem with input.
+    * If `values` cannot be parsed.
+    * If a list is assigned to a list index (e.g. 'a[1] = [1,2,3]').
+    * If the same rvalue is assigned two different values (e.g. 'a=1,a=2',
+      'a[1]=1,a[1]=2', or 'a=1,a=[1]')
+  """
+  results_dictionary = {}
+  pos = 0
+  while pos < len(values):
+    m = PARAM_RE.match(values, pos)
+    if not m:
+      raise ValueError('Malformed hyperparameter value: %s' % values[pos:])
+    # Check that there is a comma between parameters and move past it.
+    pos = m.end()
+    # Parse the values.
+    m_dict = m.groupdict()
+    name = m_dict['name']
+    if name not in type_map:
+      if ignore_unknown:
+        continue
+      raise ValueError('Unknown hyperparameter type for %s' % name)
+    type_ = type_map[name]
+
+    # Set up correct parsing function (depending on whether type_ is a bool)
+    if type_ == bool:
+
+      def parse_bool(value):
+        if value in ['true', 'True']:
+          return True
+        elif value in ['false', 'False']:
+          return False
+        else:
+          try:
+            return bool(int(value))
+          except ValueError:
+            _parse_fail(name, type_, value, values)
+
+      parse = parse_bool
+    else:
+      parse = type_
+
+    # If a singe value is provided
+    if m_dict['val'] is not None:
+      _process_scalar_value(name, parse, type_, m_dict, values,
+                            results_dictionary)
+
+    # If the assigned value is a list:
+    elif m_dict['vals'] is not None:
+      _process_list_value(name, parse, type_, m_dict, values,
+                          results_dictionary)
+
+    else:  # Not assigned a list or value
+      _parse_fail(name, type_, '', values)
+
+  return results_dictionary
+
+
+class HParams(object):
+  """Class to hold a set of hyperparameters as name-value pairs.
+
+  A `HParams` object holds hyperparameters used to build and train a model,
+  such as the number of hidden units in a neural net layer or the learning rate
+  to use when training.
+
+  You first create a `HParams` object by specifying the names and values of the
+  hyperparameters.
+
+  To make them easily accessible the parameter names are added as direct
+  attributes of the class.  A typical usage is as follows:
+
+  ```python
+  # Create a HParams object specifying names and values of the model
+  # hyperparameters:
+  hparams = HParams(learning_rate=0.1, num_hidden_units=100)
+
+  # The hyperparameter are available as attributes of the HParams object:
+  hparams.learning_rate ==> 0.1
+  hparams.num_hidden_units ==> 100
+  ```
+
+  Hyperparameters have type, which is inferred from the type of their value
+  passed at construction type.   The currently supported types are: integer,
+  float, boolean, string, and list of integer, float, boolean, or string.
+
+  You can override hyperparameter values by calling the
+  [`parse()`](#HParams.parse) method, passing a string of comma separated
+  `name=value` pairs.  This is intended to make it possible to override
+  any hyperparameter values from a single command-line flag to which
+  the user passes 'hyper-param=value' pairs.  It avoids having to define
+  one flag for each hyperparameter.
+
+  The syntax expected for each value depends on the type of the parameter.
+  See `parse()` for a description of the syntax.
+
+  Example:
+
+  ```python
+  # Define a command line flag to pass name=value pairs.
+  # For example using argparse:
+  import argparse
+  parser = argparse.ArgumentParser(description='Train my model.')
+  parser.add_argument('--hparams', type=str,
+                      help='Comma separated list of "name=value" pairs.')
+  args = parser.parse_args()
+  ...
+  def my_program():
+    # Create a HParams object specifying the names and values of the
+    # model hyperparameters:
+    hparams = tf.HParams(learning_rate=0.1, num_hidden_units=100,
+                         activations=['relu', 'tanh'])
+
+    # Override hyperparameters values by parsing the command line
+    hparams.parse(args.hparams)
+
+    # If the user passed `--hparams=learning_rate=0.3` on the command line
+    # then 'hparams' has the following attributes:
+    hparams.learning_rate ==> 0.3
+    hparams.num_hidden_units ==> 100
+    hparams.activations ==> ['relu', 'tanh']
+
+    # If the hyperparameters are in json format use parse_json:
+    hparams.parse_json('{"learning_rate": 0.3, "activations": "relu"}')
+  ```
+  """
+
+  _HAS_DYNAMIC_ATTRIBUTES = True  # Required for pytype checks.
+
+  def __init__(self, model_structure=None, **kwargs):
+    """Create an instance of `HParams` from keyword arguments.
+
+    The keyword arguments specify name-values pairs for the hyperparameters.
+    The parameter types are inferred from the type of the values passed.
+
+    The parameter names are added as attributes of `HParams` object, so they
+    can be accessed directly with the dot notation `hparams._name_`.
+
+    Example:
+
+    ```python
+    # Define 3 hyperparameters: 'learning_rate' is a float parameter,
+    # 'num_hidden_units' an integer parameter, and 'activation' a string
+    # parameter.
+    hparams = tf.HParams(
+        learning_rate=0.1, num_hidden_units=100, activation='relu')
+
+    hparams.activation ==> 'relu'
+    ```
+
+    Note that a few names are reserved and cannot be used as hyperparameter
+    names.  If you use one of the reserved name the constructor raises a
+    `ValueError`.
+
+    Args:
+      model_structure: An instance of ModelStructure, defining the feature
+        crosses to be used in the Trial.
+      **kwargs: Key-value pairs where the key is the hyperparameter name and
+        the value is the value for the parameter.
+
+    Raises:
+      ValueError: If both `hparam_def` and initialization values are provided,
+        or if one of the arguments is invalid.
+
+    """
+    # Register the hyperparameters and their type in _hparam_types.
+    # This simplifies the implementation of parse().
+    # _hparam_types maps the parameter name to a tuple (type, bool).
+    # The type value is the type of the parameter for scalar hyperparameters,
+    # or the type of the list elements for multidimensional hyperparameters.
+    # The bool value is True if the value is a list, False otherwise.
+    self._hparam_types = {}
+    self._model_structure = model_structure
+    for name, value in six.iteritems(kwargs):
+      self.add_hparam(name, value)
+
+  def add_hparam(self, name, value):
+    """Adds {name, value} pair to hyperparameters.
+
+    Args:
+      name: Name of the hyperparameter.
+      value: Value of the hyperparameter. Can be one of the following types:
+        int, float, string, int list, float list, or string list.
+
+    Raises:
+      ValueError: if one of the arguments is invalid.
+    """
+    # Keys in kwargs are unique, but 'name' could the name of a pre-existing
+    # attribute of this object.  In that case we refuse to use it as a
+    # hyperparameter name.
+    if getattr(self, name, None) is not None:
+      raise ValueError('Hyperparameter name is reserved: %s' % name)
+    if isinstance(value, (list, tuple)):
+      if not value:
+        raise ValueError(
+            'Multi-valued hyperparameters cannot be empty: %s' % name)
+      self._hparam_types[name] = (type(value[0]), True)
+    else:
+      self._hparam_types[name] = (type(value), False)
+    setattr(self, name, value)
+
+  def set_hparam(self, name, value):
+    """Set the value of an existing hyperparameter.
+
+    This function verifies that the type of the value matches the type of the
+    existing hyperparameter.
+
+    Args:
+      name: Name of the hyperparameter.
+      value: New value of the hyperparameter.
+
+    Raises:
+      KeyError: If the hyperparameter doesn't exist.
+      ValueError: If there is a type mismatch.
+    """
+    param_type, is_list = self._hparam_types[name]
+    if isinstance(value, list):
+      if not is_list:
+        raise ValueError(
+            'Must not pass a list for single-valued parameter: %s' % name)
+      setattr(self, name, [
+          _cast_to_type_if_compatible(name, param_type, v) for v in value])
+    else:
+      if is_list:
+        raise ValueError(
+            'Must pass a list for multi-valued parameter: %s.' % name)
+      setattr(self, name, _cast_to_type_if_compatible(name, param_type, value))
+
+  def del_hparam(self, name):
+    """Removes the hyperparameter with key 'name'.
+
+    Does nothing if it isn't present.
+
+    Args:
+      name: Name of the hyperparameter.
+    """
+    if hasattr(self, name):
+      delattr(self, name)
+      del self._hparam_types[name]
+
+  def parse(self, values):
+    """Override existing hyperparameter values, parsing new values from a string.
+
+    See parse_values for more detail on the allowed format for values.
+
+    Args:
+      values: String.  Comma separated list of `name=value` pairs where 'value'
+        must follow the syntax described above.
+
+    Returns:
+      The `HParams` instance.
+
+    Raises:
+      ValueError: If `values` cannot be parsed or a hyperparameter in `values`
+      doesn't exist.
+    """
+    type_map = dict()
+    for name, t in self._hparam_types.items():
+      param_type, _ = t
+      type_map[name] = param_type
+
+    values_map = parse_values(values, type_map)
+    return self.override_from_dict(values_map)
+
+  def override_from_dict(self, values_dict):
+    """Override existing hyperparameter values, parsing new values from a dictionary.
+
+    Args:
+      values_dict: Dictionary of name:value pairs.
+
+    Returns:
+      The `HParams` instance.
+
+    Raises:
+      KeyError: If a hyperparameter in `values_dict` doesn't exist.
+      ValueError: If `values_dict` cannot be parsed.
+    """
+    for name, value in values_dict.items():
+      self.set_hparam(name, value)
+    return self
+
+  def set_model_structure(self, model_structure):
+    self._model_structure = model_structure
+
+  def get_model_structure(self):
+    return self._model_structure
+
+  def to_json(self, indent=None, separators=None, sort_keys=False):
+    """Serializes the hyperparameters into JSON.
+
+    Args:
+      indent: If a non-negative integer, JSON array elements and object members
+        will be pretty-printed with that indent level. An indent level of 0, or
+        negative, will only insert newlines. `None` (the default) selects the
+        most compact representation.
+      separators: Optional `(item_separator, key_separator)` tuple. Default is
+        `(', ', ': ')`.
+      sort_keys: If `True`, the output dictionaries will be sorted by key.
+
+    Returns:
+      A JSON string.
+    """
+    return json.dumps(
+        self.values(),
+        indent=indent,
+        separators=separators,
+        sort_keys=sort_keys)
+
+  def parse_json(self, values_json):
+    """Override existing hyperparameter values, parsing new values from a json object.
+
+    Args:
+      values_json: String containing a json object of name:value pairs.
+
+    Returns:
+      The `HParams` instance.
+
+    Raises:
+      KeyError: If a hyperparameter in `values_json` doesn't exist.
+      ValueError: If `values_json` cannot be parsed.
+    """
+    values_map = json.loads(values_json)
+    return self.override_from_dict(values_map)
+
+  def values(self):
+    """Return the hyperparameter values as a Python dictionary.
+
+    Returns:
+      A dictionary with hyperparameter names as keys.  The values are the
+      hyperparameter values.
+    """
+    return {n: getattr(self, n) for n in self._hparam_types.keys()}
+
+  def get(self, key, default=None):
+    """Returns the value of `key` if it exists, else `default`."""
+    if key in self._hparam_types:
+      # Ensure that default is compatible with the parameter type.
+      if default is not None:
+        param_type, is_param_list = self._hparam_types[key]
+        type_str = 'list<%s>' % param_type if is_param_list else str(param_type)
+        fail_msg = ("Hparam '%s' of type '%s' is incompatible with "
+                    'default=%s' % (key, type_str, default))
+
+        is_default_list = isinstance(default, list)
+        if is_param_list != is_default_list:
+          raise ValueError(fail_msg)
+
+        try:
+          if is_default_list:
+            for value in default:
+              _cast_to_type_if_compatible(key, param_type, value)
+          else:
+            _cast_to_type_if_compatible(key, param_type, default)
+        except ValueError as e:
+          raise ValueError('%s. %s' % (fail_msg, e))
+
+      return getattr(self, key)
+
+    return default
+
+  def __contains__(self, key):
+    return key in self._hparam_types
+
+  def __str__(self):
+    return str(sorted(self.values().items()))
+
+  def __repr__(self):
+    return '%s(%s)' % (type(self).__name__, self.__str__())
+
+  @staticmethod
+  def _get_kind_name(param_type, is_list):
+    """Returns the field name given parameter type and is_list.
+
+    Args:
+      param_type: Data type of the hparam.
+      is_list: Whether this is a list.
+
+    Returns:
+      A string representation of the field name.
+
+    Raises:
+      ValueError: If parameter type is not recognized.
+    """
+    if issubclass(param_type, bool):
+      # This check must happen before issubclass(param_type, six.integer_types),
+      # since Python considers bool to be a subclass of int.
+      typename = 'bool'
+    elif issubclass(param_type, six.integer_types):
+      # Setting 'int' and 'long' types to be 'int64' to ensure the type is
+      # compatible with both Python2 and Python3.
+      typename = 'int64'
+    elif issubclass(param_type, (six.string_types, six.binary_type)):
+      # Setting 'string' and 'bytes' types to be 'bytes' to ensure the type is
+      # compatible with both Python2 and Python3.
+      typename = 'bytes'
+    elif issubclass(param_type, float):
+      typename = 'float'
+    else:
+      raise ValueError('Unsupported parameter type: %s' % str(param_type))
+
+    suffix = 'list' if is_list else 'value'
+    return '_'.join([typename, suffix])
diff --git a/tensor2tensor/utils/hparam_test.py b/tensor2tensor/utils/hparam_test.py
new file mode 100644
index 000000000..8c39e3485
--- /dev/null
+++ b/tensor2tensor/utils/hparam_test.py
@@ -0,0 +1,509 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+# Forked with minor changes from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/training/python/training/hparam_test.py pylint: disable=line-too-long
+"""Tests for hparam."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.utils import hparam
+
+import tensorflow as tf
+
+
+class HParamsTest(tf.test.TestCase):
+
+  def testEmpty(self):
+    hparams = hparam.HParams()
+    self.assertDictEqual({}, hparams.values())
+    hparams.parse('')
+    self.assertDictEqual({}, hparams.values())
+    with self.assertRaisesRegexp(ValueError, 'Unknown hyperparameter'):
+      hparams.parse('xyz=123')
+
+  def testContains(self):
+    hparams = hparam.HParams(foo=1)
+    self.assertTrue('foo' in hparams)
+    self.assertFalse('bar' in hparams)
+
+  def testSomeValues(self):
+    hparams = hparam.HParams(aaa=1, b=2.0, c_c='relu6', d='/a/b=c/d')
+    self.assertDictEqual(
+        {'aaa': 1, 'b': 2.0, 'c_c': 'relu6', 'd': '/a/b=c/d'},
+        hparams.values())
+    expected_str = ('[(\'aaa\', 1), (\'b\', 2.0), (\'c_c\', \'relu6\'), '
+                    '(\'d\', \'/a/b=c/d\')]')
+    self.assertEqual(expected_str, str(hparams.__str__()))
+    self.assertEqual(expected_str, str(hparams))
+    self.assertEqual(1, hparams.aaa)
+    self.assertEqual(2.0, hparams.b)
+    self.assertEqual('relu6', hparams.c_c)
+    self.assertEqual('/a/b=c/d', hparams.d)
+    hparams.parse('aaa=12')
+    self.assertDictEqual({
+        'aaa': 12,
+        'b': 2.0,
+        'c_c': 'relu6',
+        'd': '/a/b=c/d'
+    }, hparams.values())
+    self.assertEqual(12, hparams.aaa)
+    self.assertEqual(2.0, hparams.b)
+    self.assertEqual('relu6', hparams.c_c)
+    self.assertEqual('/a/b=c/d', hparams.d)
+    hparams.parse('c_c=relu4, b=-2.0e10')
+    self.assertDictEqual({
+        'aaa': 12,
+        'b': -2.0e10,
+        'c_c': 'relu4',
+        'd': '/a/b=c/d'
+    }, hparams.values())
+    self.assertEqual(12, hparams.aaa)
+    self.assertEqual(-2.0e10, hparams.b)
+    self.assertEqual('relu4', hparams.c_c)
+    self.assertEqual('/a/b=c/d', hparams.d)
+    hparams.parse('c_c=,b=0,')
+    self.assertDictEqual({'aaa': 12, 'b': 0, 'c_c': '', 'd': '/a/b=c/d'},
+                         hparams.values())
+    self.assertEqual(12, hparams.aaa)
+    self.assertEqual(0.0, hparams.b)
+    self.assertEqual('', hparams.c_c)
+    self.assertEqual('/a/b=c/d', hparams.d)
+    hparams.parse('c_c=2.3",b=+2,')
+    self.assertEqual(2.0, hparams.b)
+    self.assertEqual('2.3"', hparams.c_c)
+    hparams.parse('d=/a/b/c/d,aaa=11,')
+    self.assertEqual(11, hparams.aaa)
+    self.assertEqual(2.0, hparams.b)
+    self.assertEqual('2.3"', hparams.c_c)
+    self.assertEqual('/a/b/c/d', hparams.d)
+    hparams.parse('b=1.5,d=/a=b/c/d,aaa=10,')
+    self.assertEqual(10, hparams.aaa)
+    self.assertEqual(1.5, hparams.b)
+    self.assertEqual('2.3"', hparams.c_c)
+    self.assertEqual('/a=b/c/d', hparams.d)
+    with self.assertRaisesRegexp(ValueError, 'Unknown hyperparameter'):
+      hparams.parse('x=123')
+    with self.assertRaisesRegexp(ValueError, 'Could not parse'):
+      hparams.parse('aaa=poipoi')
+    with self.assertRaisesRegexp(ValueError, 'Could not parse'):
+      hparams.parse('aaa=1.0')
+    with self.assertRaisesRegexp(ValueError, 'Could not parse'):
+      hparams.parse('b=12x')
+    with self.assertRaisesRegexp(ValueError, 'Could not parse'):
+      hparams.parse('b=relu')
+    with self.assertRaisesRegexp(ValueError, 'Must not pass a list'):
+      hparams.parse('aaa=[123]')
+    self.assertEqual(10, hparams.aaa)
+    self.assertEqual(1.5, hparams.b)
+    self.assertEqual('2.3"', hparams.c_c)
+    self.assertEqual('/a=b/c/d', hparams.d)
+
+  def testWithPeriodInVariableName(self):
+    hparams = hparam.HParams()
+    hparams.add_hparam(name='a.b', value=0.0)
+    hparams.parse('a.b=1.0')
+    self.assertEqual(1.0, getattr(hparams, 'a.b'))
+    hparams.add_hparam(name='c.d', value=0.0)
+    with self.assertRaisesRegexp(ValueError, 'Could not parse'):
+      hparams.parse('c.d=abc')
+    hparams.add_hparam(name='e.f', value='')
+    hparams.parse('e.f=abc')
+    self.assertEqual('abc', getattr(hparams, 'e.f'))
+    hparams.add_hparam(name='d..', value=0.0)
+    hparams.parse('d..=10.0')
+    self.assertEqual(10.0, getattr(hparams, 'd..'))
+
+  def testSetFromMap(self):
+    hparams = hparam.HParams(a=1, b=2.0, c='tanh')
+    hparams.override_from_dict({'a': -2, 'c': 'identity'})
+    self.assertDictEqual({'a': -2, 'c': 'identity', 'b': 2.0}, hparams.values())
+
+    hparams = hparam.HParams(x=1, b=2.0, d=[0.5])
+    hparams.override_from_dict({'d': [0.1, 0.2, 0.3]})
+    self.assertDictEqual({'d': [0.1, 0.2, 0.3], 'x': 1, 'b': 2.0},
+                         hparams.values())
+
+  def testBoolParsing(self):
+    for value in 'true', 'false', 'True', 'False', '1', '0':
+      for initial in False, True:
+        hparams = hparam.HParams(use_gpu=initial)
+        hparams.parse('use_gpu=' + value)
+        self.assertEqual(hparams.use_gpu, value in ['True', 'true', '1'])
+
+  def testBoolParsingFail(self):
+    hparams = hparam.HParams(use_gpu=True)
+    with self.assertRaisesRegexp(ValueError, r'Could not parse.*use_gpu'):
+      hparams.parse('use_gpu=yep')
+
+  def testLists(self):
+    hparams = hparam.HParams(aaa=[1], b=[2.0, 3.0], c_c=['relu6'])
+    self.assertDictEqual({
+        'aaa': [1],
+        'b': [2.0, 3.0],
+        'c_c': ['relu6']
+    }, hparams.values())
+    self.assertEqual([1], hparams.aaa)
+    self.assertEqual([2.0, 3.0], hparams.b)
+    self.assertEqual(['relu6'], hparams.c_c)
+    hparams.parse('aaa=[12]')
+    self.assertEqual([12], hparams.aaa)
+    hparams.parse('aaa=[12,34,56]')
+    self.assertEqual([12, 34, 56], hparams.aaa)
+    hparams.parse('c_c=[relu4,relu12],b=[1.0]')
+    self.assertEqual(['relu4', 'relu12'], hparams.c_c)
+    self.assertEqual([1.0], hparams.b)
+    hparams.parse('c_c=[],aaa=[-34]')
+    self.assertEqual([-34], hparams.aaa)
+    self.assertEqual([], hparams.c_c)
+    hparams.parse('c_c=[_12,3\'4"],aaa=[+3]')
+    self.assertEqual([3], hparams.aaa)
+    self.assertEqual(['_12', '3\'4"'], hparams.c_c)
+    with self.assertRaisesRegexp(ValueError, 'Unknown hyperparameter'):
+      hparams.parse('x=[123]')
+    with self.assertRaisesRegexp(ValueError, 'Could not parse'):
+      hparams.parse('aaa=[poipoi]')
+    with self.assertRaisesRegexp(ValueError, 'Could not parse'):
+      hparams.parse('aaa=[1.0]')
+    with self.assertRaisesRegexp(ValueError, 'Could not parse'):
+      hparams.parse('b=[12x]')
+    with self.assertRaisesRegexp(ValueError, 'Could not parse'):
+      hparams.parse('b=[relu]')
+    with self.assertRaisesRegexp(ValueError, 'Must pass a list'):
+      hparams.parse('aaa=123')
+
+  def testParseValuesWithIndexAssigment1(self):
+    """Assignment to an index position."""
+    parse_dict = hparam.parse_values('arr[1]=10', {'arr': int})
+    self.assertEqual(len(parse_dict), 1)
+    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertDictEqual(parse_dict['arr'], {1: 10})
+
+  def testParseValuesWithIndexAssigment1_IgnoreUnknown(self):
+    """Assignment to an index position."""
+    parse_dict = hparam.parse_values(
+        'arr[1]=10,b=5', {'arr': int}, ignore_unknown=True)
+    self.assertEqual(len(parse_dict), 1)
+    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertDictEqual(parse_dict['arr'], {1: 10})
+
+  def testParseValuesWithIndexAssigment2(self):
+    """Assignment to multiple index positions."""
+    parse_dict = hparam.parse_values('arr[0]=10,arr[5]=20', {'arr': int})
+    self.assertEqual(len(parse_dict), 1)
+    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertDictEqual(parse_dict['arr'], {0: 10, 5: 20})
+
+  def testParseValuesWithIndexAssigment2_IgnoreUnknown(self):
+    """Assignment to multiple index positions."""
+    parse_dict = hparam.parse_values(
+        'arr[0]=10,arr[5]=20,foo=bar', {'arr': int}, ignore_unknown=True)
+    self.assertEqual(len(parse_dict), 1)
+    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertDictEqual(parse_dict['arr'], {0: 10, 5: 20})
+
+  def testParseValuesWithIndexAssigment3(self):
+    """Assignment to index positions in multiple names."""
+    parse_dict = hparam.parse_values('arr[0]=10,arr[1]=20,L[5]=100,L[10]=200',
+                                     {'arr': int,
+                                      'L': int})
+    self.assertEqual(len(parse_dict), 2)
+    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertDictEqual(parse_dict['arr'], {0: 10, 1: 20})
+    self.assertTrue(isinstance(parse_dict['L'], dict))
+    self.assertDictEqual(parse_dict['L'], {5: 100, 10: 200})
+
+  def testParseValuesWithIndexAssigment3_IgnoreUnknown(self):
+    """Assignment to index positions in multiple names."""
+    parse_dict = hparam.parse_values(
+        'arr[0]=10,C=5,arr[1]=20,B[0]=kkk,L[5]=100,L[10]=200',
+        {'arr': int, 'L': int}, ignore_unknown=True)
+    self.assertEqual(len(parse_dict), 2)
+    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertDictEqual(parse_dict['arr'], {0: 10, 1: 20})
+    self.assertTrue(isinstance(parse_dict['L'], dict))
+    self.assertDictEqual(parse_dict['L'], {5: 100, 10: 200})
+
+  def testParseValuesWithIndexAssigment4(self):
+    """Assignment of index positions and scalars."""
+    parse_dict = hparam.parse_values('x=10,arr[1]=20,y=30',
+                                     {'x': int,
+                                      'y': int,
+                                      'arr': int})
+    self.assertEqual(len(parse_dict), 3)
+    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertDictEqual(parse_dict['arr'], {1: 20})
+    self.assertEqual(parse_dict['x'], 10)
+    self.assertEqual(parse_dict['y'], 30)
+
+  def testParseValuesWithIndexAssigment4_IgnoreUnknown(self):
+    """Assignment of index positions and scalars."""
+    parse_dict = hparam.parse_values(
+        'x=10,foo[0]=bar,arr[1]=20,zzz=78,y=30',
+        {'x': int, 'y': int, 'arr': int}, ignore_unknown=True)
+    self.assertEqual(len(parse_dict), 3)
+    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertDictEqual(parse_dict['arr'], {1: 20})
+    self.assertEqual(parse_dict['x'], 10)
+    self.assertEqual(parse_dict['y'], 30)
+
+  def testParseValuesWithIndexAssigment5(self):
+    """Different variable types."""
+    parse_dict = hparam.parse_values('a[0]=5,b[1]=true,c[2]=abc,d[3]=3.14', {
+        'a': int,
+        'b': bool,
+        'c': str,
+        'd': float
+    })
+    self.assertEqual(set(parse_dict.keys()), {'a', 'b', 'c', 'd'})
+    self.assertTrue(isinstance(parse_dict['a'], dict))
+    self.assertDictEqual(parse_dict['a'], {0: 5})
+    self.assertTrue(isinstance(parse_dict['b'], dict))
+    self.assertDictEqual(parse_dict['b'], {1: True})
+    self.assertTrue(isinstance(parse_dict['c'], dict))
+    self.assertDictEqual(parse_dict['c'], {2: 'abc'})
+    self.assertTrue(isinstance(parse_dict['d'], dict))
+    self.assertDictEqual(parse_dict['d'], {3: 3.14})
+
+  def testParseValuesWithIndexAssigment5_IgnoreUnknown(self):
+    """Different variable types."""
+    parse_dict = hparam.parse_values(
+        'a[0]=5,cc=4,b[1]=true,c[2]=abc,mm=2,d[3]=3.14',
+        {'a': int, 'b': bool, 'c': str, 'd': float},
+        ignore_unknown=True)
+    self.assertEqual(set(parse_dict.keys()), {'a', 'b', 'c', 'd'})
+    self.assertTrue(isinstance(parse_dict['a'], dict))
+    self.assertDictEqual(parse_dict['a'], {0: 5})
+    self.assertTrue(isinstance(parse_dict['b'], dict))
+    self.assertDictEqual(parse_dict['b'], {1: True})
+    self.assertTrue(isinstance(parse_dict['c'], dict))
+    self.assertDictEqual(parse_dict['c'], {2: 'abc'})
+    self.assertTrue(isinstance(parse_dict['d'], dict))
+    self.assertDictEqual(parse_dict['d'], {3: 3.14})
+
+  def testParseValuesWithBadIndexAssigment1(self):
+    """Reject assignment of list to variable type."""
+    with self.assertRaisesRegexp(ValueError,
+                                 r'Assignment of a list to a list index.'):
+      hparam.parse_values('arr[1]=[1,2,3]', {'arr': int})
+
+  def testParseValuesWithBadIndexAssigment1_IgnoreUnknown(self):
+    """Reject assignment of list to variable type."""
+    with self.assertRaisesRegexp(ValueError,
+                                 r'Assignment of a list to a list index.'):
+      hparam.parse_values(
+          'arr[1]=[1,2,3],c=8', {'arr': int}, ignore_unknown=True)
+
+  def testParseValuesWithBadIndexAssigment2(self):
+    """Reject if type missing."""
+    with self.assertRaisesRegexp(ValueError,
+                                 r'Unknown hyperparameter type for arr'):
+      hparam.parse_values('arr[1]=5', {})
+
+  def testParseValuesWithBadIndexAssigment2_IgnoreUnknown(self):
+    """Ignore missing type."""
+    hparam.parse_values('arr[1]=5', {}, ignore_unknown=True)
+
+  def testParseValuesWithBadIndexAssigment3(self):
+    """Reject type of the form name[index]."""
+    with self.assertRaisesRegexp(ValueError,
+                                 'Unknown hyperparameter type for arr'):
+      hparam.parse_values('arr[1]=1', {'arr[1]': int})
+
+  def testParseValuesWithBadIndexAssigment3_IgnoreUnknown(self):
+    """Ignore type of the form name[index]."""
+    hparam.parse_values('arr[1]=1', {'arr[1]': int}, ignore_unknown=True)
+
+  def testWithReusedVariables(self):
+    with self.assertRaisesRegexp(ValueError,
+                                 'Multiple assignments to variable \'x\''):
+      hparam.parse_values('x=1,x=1', {'x': int})
+
+    with self.assertRaisesRegexp(ValueError,
+                                 'Multiple assignments to variable \'arr\''):
+      hparam.parse_values('arr=[100,200],arr[0]=10', {'arr': int})
+
+    with self.assertRaisesRegexp(
+        ValueError, r'Multiple assignments to variable \'arr\[0\]\''):
+      hparam.parse_values('arr[0]=10,arr[0]=20', {'arr': int})
+
+    with self.assertRaisesRegexp(ValueError,
+                                 'Multiple assignments to variable \'arr\''):
+      hparam.parse_values('arr[0]=10,arr=[100]', {'arr': int})
+
+  def testJson(self):
+    hparams = hparam.HParams(aaa=1, b=2.0, c_c='relu6', d=True)
+    self.assertDictEqual({
+        'aaa': 1,
+        'b': 2.0,
+        'c_c': 'relu6',
+        'd': True
+    }, hparams.values())
+    self.assertEqual(1, hparams.aaa)
+    self.assertEqual(2.0, hparams.b)
+    self.assertEqual('relu6', hparams.c_c)
+    hparams.parse_json('{"aaa": 12, "b": 3.0, "c_c": "relu4", "d": false}')
+    self.assertDictEqual({
+        'aaa': 12,
+        'b': 3.0,
+        'c_c': 'relu4',
+        'd': False
+    }, hparams.values())
+    self.assertEqual(12, hparams.aaa)
+    self.assertEqual(3.0, hparams.b)
+    self.assertEqual('relu4', hparams.c_c)
+
+    json_str = hparams.to_json()
+    hparams2 = hparam.HParams(aaa=10, b=20.0, c_c='hello', d=False)
+    hparams2.parse_json(json_str)
+    self.assertEqual(12, hparams2.aaa)
+    self.assertEqual(3.0, hparams2.b)
+    self.assertEqual('relu4', hparams2.c_c)
+    self.assertEqual(False, hparams2.d)
+
+    hparams3 = hparam.HParams(aaa=123)
+    self.assertEqual('{"aaa": 123}', hparams3.to_json())
+    self.assertEqual('{\n  "aaa": 123\n}', hparams3.to_json(indent=2))
+    self.assertEqual('{"aaa"=123}', hparams3.to_json(separators=(';', '=')))
+
+    hparams4 = hparam.HParams(aaa=123, b='hello', c_c=False)
+    self.assertEqual(
+        '{"aaa": 123, "b": "hello", "c_c": false}',
+        hparams4.to_json(sort_keys=True))
+
+  def testSetHParam(self):
+    hparams = hparam.HParams(aaa=1, b=2.0, c_c='relu6', d=True)
+    self.assertDictEqual({
+        'aaa': 1,
+        'b': 2.0,
+        'c_c': 'relu6',
+        'd': True
+    }, hparams.values())
+    self.assertEqual(1, hparams.aaa)
+    self.assertEqual(2.0, hparams.b)
+    self.assertEqual('relu6', hparams.c_c)
+
+    hparams.set_hparam('aaa', 12)
+    hparams.set_hparam('b', 3.0)
+    hparams.set_hparam('c_c', 'relu4')
+    hparams.set_hparam('d', False)
+    self.assertDictEqual({
+        'aaa': 12,
+        'b': 3.0,
+        'c_c': 'relu4',
+        'd': False
+    }, hparams.values())
+    self.assertEqual(12, hparams.aaa)
+    self.assertEqual(3.0, hparams.b)
+    self.assertEqual('relu4', hparams.c_c)
+
+  def testSetHParamListNonListMismatch(self):
+    hparams = hparam.HParams(a=1, b=[2.0, 3.0])
+    with self.assertRaisesRegexp(ValueError, r'Must not pass a list'):
+      hparams.set_hparam('a', [1.0])
+    with self.assertRaisesRegexp(ValueError, r'Must pass a list'):
+      hparams.set_hparam('b', 1.0)
+
+  def testSetHParamTypeMismatch(self):
+    hparams = hparam.HParams(
+        int_=1, str_='str', bool_=True, float_=1.1, list_int=[1, 2], none=None)
+
+    with self.assertRaises(ValueError):
+      hparams.set_hparam('str_', 2.2)
+
+    with self.assertRaises(ValueError):
+      hparams.set_hparam('int_', False)
+
+    with self.assertRaises(ValueError):
+      hparams.set_hparam('bool_', 1)
+
+    with self.assertRaises(ValueError):
+      hparams.set_hparam('int_', 2.2)
+
+    with self.assertRaises(ValueError):
+      hparams.set_hparam('list_int', [2, 3.3])
+
+    with self.assertRaises(ValueError):
+      hparams.set_hparam('int_', '2')
+
+    # Casting int to float is OK
+    hparams.set_hparam('float_', 1)
+
+    # Getting stuck with NoneType :(
+    hparams.set_hparam('none', '1')
+    self.assertEqual('1', hparams.none)
+
+  def testGet(self):
+    hparams = hparam.HParams(aaa=1, b=2.0, c_c='relu6', d=True, e=[5.0, 6.0])
+
+    # Existing parameters with default=None.
+    self.assertEqual(1, hparams.get('aaa'))
+    self.assertEqual(2.0, hparams.get('b'))
+    self.assertEqual('relu6', hparams.get('c_c'))
+    self.assertEqual(True, hparams.get('d'))
+    self.assertEqual([5.0, 6.0], hparams.get('e', None))
+
+    # Existing parameters with compatible defaults.
+    self.assertEqual(1, hparams.get('aaa', 2))
+    self.assertEqual(2.0, hparams.get('b', 3.0))
+    self.assertEqual(2.0, hparams.get('b', 3))
+    self.assertEqual('relu6', hparams.get('c_c', 'default'))
+    self.assertEqual(True, hparams.get('d', True))
+    self.assertEqual([5.0, 6.0], hparams.get('e', [1.0, 2.0, 3.0]))
+    self.assertEqual([5.0, 6.0], hparams.get('e', [1, 2, 3]))
+
+    # Existing parameters with incompatible defaults.
+    with self.assertRaises(ValueError):
+      hparams.get('aaa', 2.0)
+
+    with self.assertRaises(ValueError):
+      hparams.get('b', False)
+
+    with self.assertRaises(ValueError):
+      hparams.get('c_c', [1, 2, 3])
+
+    with self.assertRaises(ValueError):
+      hparams.get('d', 'relu')
+
+    with self.assertRaises(ValueError):
+      hparams.get('e', 123.0)
+
+    with self.assertRaises(ValueError):
+      hparams.get('e', ['a', 'b', 'c'])
+
+    # Nonexistent parameters.
+    self.assertEqual(None, hparams.get('unknown'))
+    self.assertEqual(123, hparams.get('unknown', 123))
+    self.assertEqual([1, 2, 3], hparams.get('unknown', [1, 2, 3]))
+
+  def testDel(self):
+    hparams = hparam.HParams(aaa=1, b=2.0)
+
+    with self.assertRaises(ValueError):
+      hparams.set_hparam('aaa', 'will fail')
+
+    with self.assertRaises(ValueError):
+      hparams.add_hparam('aaa', 'will fail')
+
+    hparams.del_hparam('aaa')
+    hparams.add_hparam('aaa', 'will work')
+    self.assertEqual('will work', hparams.get('aaa'))
+
+    hparams.set_hparam('aaa', 'still works')
+    self.assertEqual('still works', hparams.get('aaa'))
+
+
+if __name__ == '__main__':
+  tf.test.main()

From 512c9a2d48b3d3cb59806bce1a2395b419b384d8 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 1 Feb 2019 10:43:10 -0800
Subject: [PATCH 1619/2720] Migrate away from contrib.HParams to the newly
 forked over one.

PiperOrigin-RevId: 231998023
---
 README.md                                     |  2 +-
 docs/walkthrough.md                           |  2 +-
 tensor2tensor/data_generators/celeba_test.py  |  3 +-
 .../data_generators/imagenet_test.py          |  3 +-
 tensor2tensor/data_generators/mscoco_test.py  |  3 +-
 tensor2tensor/data_generators/problem.py      | 11 ++++----
 tensor2tensor/data_generators/video_utils.py  |  4 +--
 tensor2tensor/layers/common_hparams.py        |  3 +-
 .../layers/common_image_attention.py          |  4 +--
 .../layers/common_image_attention_test.py     |  7 +++--
 tensor2tensor/layers/latent_layers.py         | 28 +++++++++----------
 tensor2tensor/models/lstm.py                  |  4 +--
 tensor2tensor/models/research/glow_ops.py     |  6 ++--
 .../models/research/glow_ops_test.py          |  3 +-
 tensor2tensor/models/research/rl.py           |  7 +++--
 tensor2tensor/models/resnet.py                |  5 ++--
 tensor2tensor/models/shake_shake.py           |  3 +-
 tensor2tensor/models/video/nfg_interpolate.py |  6 ++--
 tensor2tensor/rl/evaluator.py                 |  7 +++--
 tensor2tensor/rl/player_utils.py              |  3 +-
 .../rl/trainer_model_based_params.py          |  9 +++---
 tensor2tensor/serving/query.py                |  3 +-
 tensor2tensor/utils/decoding.py               |  3 +-
 tensor2tensor/utils/diet.py                   |  3 +-
 tensor2tensor/utils/hparams_lib.py            |  5 ++--
 tensor2tensor/utils/misc_utils_test.py        |  3 +-
 tensor2tensor/utils/t2t_model.py              |  4 +--
 tensor2tensor/utils/t2t_model_test.py         |  7 +++--
 tensor2tensor/utils/video_metrics.py          |  4 ++-
 29 files changed, 88 insertions(+), 67 deletions(-)

diff --git a/README.md b/README.md
index 7e72ce13b..9248cd9d9 100644
--- a/README.md
+++ b/README.md
@@ -383,7 +383,7 @@ and are registered with
 **Hyperparameter sets** are defined and registered in code with
 [`@registry.register_hparams`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/registry.py)
 and are encoded in
-[`tf.contrib.training.HParams`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/training/python/training/hparam.py)
+[`HParams`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/hparam.py)
 objects. The `HParams` are available to both the problem specification and the
 model. A basic set of hyperparameters are defined in
 [`common_hparams.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/layers/common_hparams.py)
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 7e72ce13b..9248cd9d9 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -383,7 +383,7 @@ and are registered with
 **Hyperparameter sets** are defined and registered in code with
 [`@registry.register_hparams`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/registry.py)
 and are encoded in
-[`tf.contrib.training.HParams`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/training/python/training/hparam.py)
+[`HParams`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/hparam.py)
 objects. The `HParams` are available to both the problem specification and the
 model. A basic set of hyperparameters are defined in
 [`common_hparams.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/layers/common_hparams.py)
diff --git a/tensor2tensor/data_generators/celeba_test.py b/tensor2tensor/data_generators/celeba_test.py
index a347b3734..9f9247a65 100644
--- a/tensor2tensor/data_generators/celeba_test.py
+++ b/tensor2tensor/data_generators/celeba_test.py
@@ -21,6 +21,7 @@
 
 from absl.testing import parameterized
 from tensor2tensor.data_generators import celeba
+from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -34,7 +35,7 @@ class CelebaTest(parameterized.TestCase, tf.test.TestCase):
   def testCelebaMultiResolutionPreprocessExample(self, resize_method):
     example = {"inputs": tf.random_uniform([218, 178, 3], minval=-1.)}
     mode = tf.estimator.ModeKeys.TRAIN
-    hparams = tf.contrib.training.HParams(resolutions=[8, 16, 32])
+    hparams = HParams(resolutions=[8, 16, 32])
     if resize_method is not None:
       hparams.resize_method = resize_method
 
diff --git a/tensor2tensor/data_generators/imagenet_test.py b/tensor2tensor/data_generators/imagenet_test.py
index 22a731b68..36f2954b0 100644
--- a/tensor2tensor/data_generators/imagenet_test.py
+++ b/tensor2tensor/data_generators/imagenet_test.py
@@ -21,6 +21,7 @@
 
 from absl.testing import parameterized
 from tensor2tensor.data_generators import imagenet
+from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -34,7 +35,7 @@ class ImagenetTest(parameterized.TestCase, tf.test.TestCase):
   def testImagenetMultiResolutionPreprocessExample(self, resize_method):
     example = {"inputs": tf.random_uniform([64, 64, 3], minval=-1.)}
     mode = tf.estimator.ModeKeys.TRAIN
-    hparams = tf.contrib.training.HParams(resolutions=[8, 16, 32])
+    hparams = HParams(resolutions=[8, 16, 32])
     if resize_method is not None:
       hparams.resize_method = resize_method
 
diff --git a/tensor2tensor/data_generators/mscoco_test.py b/tensor2tensor/data_generators/mscoco_test.py
index 9c5564a71..6958180bf 100644
--- a/tensor2tensor/data_generators/mscoco_test.py
+++ b/tensor2tensor/data_generators/mscoco_test.py
@@ -21,6 +21,7 @@
 
 from absl.testing import parameterized
 from tensor2tensor.data_generators import mscoco
+from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -34,7 +35,7 @@ class MscocoTest(parameterized.TestCase, tf.test.TestCase):
   def testMsCocoMultiResolutionPreprocessExample(self, resize_method):
     example = {"inputs": tf.random_uniform([400, 400, 3], minval=-1.)}
     mode = tf.estimator.ModeKeys.TRAIN
-    hparams = tf.contrib.training.HParams(resolutions=[8, 16, 32])
+    hparams = HParams(resolutions=[8, 16, 32])
     if resize_method is not None:
       hparams.resize_method = resize_method
 
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 7cff2e202..e3c8c6c1e 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -29,6 +29,7 @@
 from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import mlperf_log
+from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 from tensorflow.contrib.tpu.python.tpu import tpu_config
@@ -130,7 +131,7 @@ class TaskID(object):
 
 
 def default_model_hparams():
-  return tf.contrib.training.HParams(
+  return HParams(
       max_input_seq_length=0,
       max_target_seq_length=0,
       prepend_mode="none",
@@ -596,7 +597,7 @@ def dataset(self,
       output_buffer_size: int, how many elements to prefetch at end of pipeline.
       shuffle_files: whether to shuffle input files. Default behavior (i.e. when
         shuffle_files=None) is to shuffle if mode == TRAIN.
-      hparams: tf.contrib.training.HParams; hparams to be passed to
+      hparams: HParams; hparams to be passed to
         Problem.preprocess_example and Problem.hparams. If None, will use a
         default set that is a no-op.
       preprocess: bool, whether to map the Dataset through
@@ -992,9 +993,9 @@ def _create_modalities(problem_hparams, model_hparams):
   """Creates modalities and overrides any according to model hparams.
 
   Args:
-    problem_hparams: tf.contrib.training.HParams for the Problem. It must have
+    problem_hparams: HParams for the Problem. It must have
       modality which is a dict of strings to ModalityTypes or Modality classes.
-    model_hparams: tf.contrib.training.HParams for the model. It may have
+    model_hparams: HParams for the model. It may have
       input_modalities and target_modality, which will override
       problem_hparams' modality input and target keys.
 
@@ -1024,7 +1025,7 @@ def _create_modalities(problem_hparams, model_hparams):
 
 def _default_hparams():
   """A set of basic model hyperparameters."""
-  return tf.contrib.training.HParams(
+  return HParams(
       # Use this parameter to get comparable perplexity numbers with different
       # tokenizations.  This value should be set to the ratio of the number of
       # tokens in the test set according to the tokenization used to the number
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 818f6c3bc..28f1f73d3 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -115,7 +115,7 @@ def convert_videos_to_summaries(input_videos, output_videos, target_videos,
     output_videos: 5-D NumPy array, (NTHWC) model predictions.
     target_videos: 5-D NumPy array, (NTHWC) target frames.
     tag: tf summary tag.
-    decode_hparams: tf.contrib.training.HParams.
+    decode_hparams: HParams.
     display_ground_truth: Whether or not to display ground truth videos.
   Returns:
     summaries: a list of tf frame-by-frame and video summaries.
@@ -274,7 +274,7 @@ def max_frames_per_video(self, hparams):
     hparams.video_num_target_frames.
 
     Args:
-      hparams: tf.contrib.training.HParams.
+      hparams: HParams.
     Returns:
       num_frames: int.
     """
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index d08ecbb52..7602e6b84 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 from six.moves import zip  # pylint: disable=redefined-builtin
 from tensor2tensor.utils import registry
+from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -27,7 +28,7 @@
 @registry.register_hparams("basic_1")
 def basic_params1():
   """A set of basic hyperparameters."""
-  return tf.contrib.training.HParams(
+  return HParams(
       # If the problem consists of variable-length sequences
       # (see problem.batch_size_means_tokens()), then this is the number
       # of tokens per batch per GPU or per TPU core.  Otherwise, this is
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 9edcda37a..1580ae376 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -477,7 +477,7 @@ def postprocess_image(x, rows, cols, hparams):
       number of elements in x is batch * rows * cols * hparams.hidden_size.
     rows: Integer representing number of rows in a 2-D data point.
     cols: Integer representing number of columns in a 2-D data point.
-    hparams: tf.contrib.training.HParams set.
+    hparams: HParams set.
 
   Returns:
     Tensor of shape [batch, rows, cols, depth], where depth is
@@ -639,7 +639,7 @@ def create_output(decoder_output, rows, cols, targets, hparams):
     cols: Integer representing number of columns in a 2-D data point.
     targets: Tensor of shape [batch, hparams.img_len, hparams.img_len,
       hparams.num_channels].
-    hparams: tf.contrib.training.HParams set.
+    hparams: HParams set.
 
   Returns:
     Tensor of shape [batch, hparams.img_len, hparams.img_len,
diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index 1a3791cc8..bf6ce85e6 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -22,6 +22,7 @@
 from absl.testing import parameterized
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_image_attention
+from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -36,7 +37,7 @@ def testPostProcessImageTrainMode(self, likelihood, num_mixtures, depth):
     batch = 1
     rows = 8
     cols = 24
-    hparams = tf.contrib.training.HParams(
+    hparams = HParams(
         hidden_size=2,
         likelihood=likelihood,
         mode=tf.estimator.ModeKeys.TRAIN,
@@ -58,7 +59,7 @@ def testPostProcessImageInferMode(self, likelihood, num_mixtures, depth):
     cols = 24
     block_length = 4
     block_width = 2
-    hparams = tf.contrib.training.HParams(
+    hparams = HParams(
         block_raster_scan=True,
         hidden_size=2,
         likelihood=likelihood,
@@ -90,7 +91,7 @@ def testCreateOutputTrainMode(self, likelihood, num_mixtures, depth):
       cols = channels * width
     else:
       cols = width
-    hparams = tf.contrib.training.HParams(
+    hparams = HParams(
         hidden_size=2,
         likelihood=likelihood,
         mode=tf.estimator.ModeKeys.TRAIN,
diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index c32cd691d..b493f5d04 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -106,7 +106,7 @@ def ae_latent_softmax(latents_pred, latents_discrete_hot, vocab_size, hparams):
     latents_pred: Tensor of shape [..., depth].
     latents_discrete_hot: Tensor of shape [..., vocab_size].
     vocab_size: an int representing the vocab size.
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
 
   Returns:
     sample: Tensor of shape [...], a sample from a multinomial distribution.
@@ -143,7 +143,7 @@ def ae_latent_sample_beam(latents_dense_in, inputs, ed, embed, hparams):
       length_kv]. Encoder-decoder attention bias.
     embed: Callable which embeds discrete latent hot-vectors and a hidden size
       and returns dense vectors.
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
 
   Returns:
     Tensor of shape [batch, length].
@@ -192,7 +192,7 @@ def residual_block_layer(inputs, hparams):
 
   Args:
     inputs: Tensor of shape [batch, height, width, hparams.hidden_size].
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
 
   Returns:
     Tensor of shape [batch, height, width, hparams.hidden_size].
@@ -228,7 +228,7 @@ def compress_encoder(inputs,
 
   Args:
     inputs: Tensor of shape [batch, height, width, channels].
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     strides: Tuple, strides for conv block.
     kernel_size: Tuple, kernel window size for conv block.
     name: string, variable scope.
@@ -275,7 +275,7 @@ def compress_encoder_2d(x, hparams, name=None):
 
   Args:
     x: Tensor of shape [batch, height, width, channels].
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     name: string, variable scope.
 
   Returns:
@@ -296,7 +296,7 @@ def compress_encoder_1d(x, hparams, name=None):
 
   Args:
     x: Tensor of shape [batch, length, channels].
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     name: string, variable scope.
 
   Returns:
@@ -321,7 +321,7 @@ def decompress_decoder(inputs,
 
   Args:
     inputs: Tensor of shape [batch, compress_height, compress_width, channels].
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     strides: Tuple, strides for conv block.
     kernel: Tuple, kernel window size for conv block.
     name: string, variable scope.
@@ -357,7 +357,7 @@ def decompress_decoder_2d(x, hparams, name=None):
 
   Args:
     x: Tensor of shape [batch, compress_height, compress_width, channels].
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     name: string, variable scope.
 
   Returns:
@@ -374,7 +374,7 @@ def decompress_decoder_1d(x, hparams, name=None):
 
   Args:
     x: Tensor of shape [batch, compress_length, channels].
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     name: string, variable scope.
 
   Returns:
@@ -397,7 +397,7 @@ def transformer_text_encoder(inputs,
   Args:
     inputs: Tensor of shape [batch, length, 1, hparams.hidden_size].
     target_space: int. Used for encoding inputs under a target space id.
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     name: string, variable scope.
 
   Returns:
@@ -432,7 +432,7 @@ def transformer_image_decoder(targets,
     encoder_output: Tensor of shape [batch, length_kv, hparams.hidden_size].
     ed_attention_bias: Tensor which broadcasts with shape [batch,
       hparams.num_heads, length_q, length_kv]. Encoder-decoder attention bias.
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     name: string, variable scope.
 
   Returns:
@@ -476,7 +476,7 @@ def transformer_latent_decoder(x,
     encoder_output: Tensor of shape [batch, length_kv, hparams.hidden_size].
     ed_attention_bias: Tensor which broadcasts with shape [batch,
       hparams.num_heads, length_q, length_kv]. Encoder-decoder attention bias.
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     name: string, variable scope.
 
   Returns:
@@ -547,7 +547,7 @@ def latent_prediction_model(inputs,
     latents_dense: Tensor of shape [batch, length_q, hparams.hidden_size].
       length_q is the latent length, which is
       height * width * hparams.num_latents / (2**hparams.num_compress_steps).
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     vocab_size: int or None. If None, it is 2**hparams.bottleneck_bits.
     name: string, variable scope.
 
@@ -586,7 +586,7 @@ def transformer_autoencoder(inputs,
     targets: Tensor of shape [batch, ..., channels]. Ellipses may be 1 or 2
       dimensions denoting sequence length.
     target_space: int. Used for encoding inputs under a target space id.
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     cache: Tensor of shape [batch, length] or None.
     predict_mask: Tensor masking whether to use gold targets or predictions.
 
diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 58cadea46..9841c6019 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -41,7 +41,7 @@ def lstm(inputs, sequence_length, hparams, train, name, initial_state=None):
     inputs: The input `Tensor`, shaped `[batch_size, time_steps, hidden_size]`.
     sequence_length: Lengths of the actual input sequence, excluding padding; a
         `Tensor` shaped `[batch_size]`.
-    hparams: tf.contrib.training.HParams; hyperparameters.
+    hparams: HParams; hyperparameters.
     train: bool; `True` when constructing training graph to enable dropout.
     name: string; Create variable names under this scope.
     initial_state: tuple of `LSTMStateTuple`s; the initial state of each layer.
@@ -74,7 +74,7 @@ def lstm_attention_decoder(inputs, hparams, train, name, initial_state,
   Args:
     inputs: The decoder input `Tensor`, shaped `[batch_size, decoder_steps,
         hidden_size]`.
-    hparams: tf.contrib.training.HParams; hyperparameters.
+    hparams: HParams; hyperparameters.
     train: bool; `True` when constructing training graph to enable dropout.
     name: string; Create variable names under this scope.
     initial_state: Tuple of `LSTMStateTuple`s; the initial state of each layer.
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index c5b2d0a11..3a9ba8a76 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -742,7 +742,7 @@ def get_dilation_rates(hparams, width):
   """Get a list of valid dilation rates.
 
   Args:
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     width: spatial dimension. Ensures that the effective filter size is
            not larger than the spatial dimension.
   Returns:
@@ -832,7 +832,7 @@ def latent_to_dist(name, x, hparams, output_channels=None):
   Args:
     name: variable scope.
     x: 4-D Tensor of shape (NHWC)
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
       latent_architecture - can be "single_conv", "glow_nn" or "glow_resnet",
                             default = single_conv
       latent_encoder_depth - int, depth of architecture, valid if
@@ -893,7 +893,7 @@ def noise_op(latents, hparams):
 
   Args:
     latents: 4-D or 5-D tensor, shape=(NTHWC) or (NHWC).
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
   Returns:
     latents: latents with isotropic gaussian noise appended.
   """
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 00f40d3e9..45b6ae8d4 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -25,6 +25,7 @@
 import numpy as np
 from tensor2tensor.models.research import glow
 from tensor2tensor.models.research import glow_ops
+from tensor2tensor.utils.hparam import HParams
 import tensorflow as tf
 
 arg_scope = tf.contrib.framework.arg_scope
@@ -156,7 +157,7 @@ def test_conv_stack(self, activation="relu"):
   def check_latent_to_dist(self, architecture):
     with tf.Graph().as_default():
       x = tf.random_uniform(shape=(16, 5, 5, 32))
-      hparams = tf.contrib.training.HParams(architecture=architecture)
+      hparams = HParams(architecture=architecture)
       x_prior = glow_ops.latent_to_dist("split_prior", x, hparams=hparams,
                                         output_channels=64)
       mean_t, scale_t = x_prior.loc, x_prior.scale
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 9f1982636..0f8e1aa46 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -36,6 +36,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 from tensor2tensor.utils import trainer_lib
+from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 import tensorflow_probability as tfp
@@ -78,7 +79,7 @@ def ppo_base_v1():
 @registry.register_hparams
 def basic_policy_parameters():
   wrappers = None
-  return tf.contrib.training.HParams(wrappers=wrappers)
+  return HParams(wrappers=wrappers)
 
 
 @registry.register_hparams
@@ -335,7 +336,7 @@ def ppo_pong_ae_base():
 def dqn_atari_base():
   # These params are based on agents/dqn/configs/dqn.gin
   # with some modifications taking into account our code
-  return tf.contrib.training.HParams(
+  return HParams(
       agent_gamma=0.99,
       agent_update_horizon=1,
       agent_min_replay_history=20000,  # agent steps
@@ -372,7 +373,7 @@ def dqn_original_params():
 
 @registry.register_hparams
 def rlmf_original():
-  return tf.contrib.training.HParams(
+  return HParams(
       game="pong",
       base_algo="ppo",
       base_algo_params="ppo_original_params",
diff --git a/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index 8dfabefed..39277a3ab 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -24,6 +24,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
+from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -787,7 +788,7 @@ def resnet_200():
 # Pruning parameters
 @registry.register_pruning_params
 def resnet_weight():
-  hp = tf.contrib.training.HParams()
+  hp = HParams()
   hp.add_hparam("strategy", "weight")
   hp.add_hparam("black_list", ["logits", "bias"])
   hp.add_hparam("white_list", ["td_conv"])
@@ -805,7 +806,7 @@ def resnet_unit():
 # Adversarial attack parameters
 @registry.register_attack_params
 def resnet_fgsm():
-  aparams = tf.contrib.training.HParams()
+  aparams = HParams()
   aparams.attack = "fgsm"
   aparams.epsilon_name = "eps"
   aparams.attack_epsilons = [i * 0.8 for i in range(20)]
diff --git a/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index 22f46f989..01681bb05 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -22,6 +22,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
+from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -214,7 +215,7 @@ def shakeshake_tpu():
 
 @registry.register_attack_params
 def shake_shake_fgsm():
-  aparams = tf.contrib.training.HParams()
+  aparams = HParams()
   aparams.attack = "fgsm"
   aparams.attack_epsilons = [(i+1) * 0.1 for i in range(12)]
   aparams.add_hparam("clip_min", 0.0)
diff --git a/tensor2tensor/models/video/nfg_interpolate.py b/tensor2tensor/models/video/nfg_interpolate.py
index 281536415..71a550a2c 100644
--- a/tensor2tensor/models/video/nfg_interpolate.py
+++ b/tensor2tensor/models/video/nfg_interpolate.py
@@ -96,7 +96,7 @@ def interpolate(features, hparams, num_interp):
 
   Args:
     features: dict of tensors
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     num_interp: integer.
   Returns:
     images: 4-D Tensor, shape=(num_interp, H, W, C)
@@ -134,8 +134,8 @@ def interpolations_to_summary(sample_ind, interpolations, hparams,
   Args:
     sample_ind: int
     interpolations: Numpy array, shape=(num_interp, 64, 64, 3)
-    hparams: tf.contrib.training.HParams, train hparams
-    decode_hparams: tf.contrib.training.HParams, decode hparams
+    hparams: HParams, train hparams
+    decode_hparams: HParams, decode hparams
   Returns:
     summaries: list of tf Summary Values.
   """
diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 5d574aa45..c34e75c9b 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -39,6 +39,7 @@
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
+from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -105,7 +106,7 @@
 
 @registry.register_hparams
 def planner_tiny():
-  return tf.contrib.training.HParams(
+  return HParams(
       num_rollouts=1,
       planning_horizon=2,
       rollout_agent_type="random",
@@ -118,7 +119,7 @@ def planner_tiny():
 
 @registry.register_hparams
 def planner_small():
-  return tf.contrib.training.HParams(
+  return HParams(
       num_rollouts=64,
       planning_horizon=16,
       rollout_agent_type="policy",
@@ -131,7 +132,7 @@ def planner_small():
 
 @registry.register_hparams
 def planner_base():
-  return tf.contrib.training.HParams(
+  return HParams(
       num_rollouts=96,
       batch_size=96,
       planning_horizon=8,
diff --git a/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index 376ab1f10..63695c094 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -32,6 +32,7 @@
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
 from tensor2tensor.utils import trainer_lib
+from tensor2tensor.utils.hparam import HParams
 from tensor2tensor.utils.misc_utils import camelcase_to_snakecase
 
 import tensorflow as tf
@@ -282,7 +283,7 @@ def create_simulated_env(
     if key not in other_hparams:
       other_hparams[key] = a_bit_risky_defaults[key]
 
-  hparams = tf.contrib.training.HParams(
+  hparams = HParams(
       grayscale=grayscale,
       resize_width_factor=resize_width_factor,
       resize_height_factor=resize_height_factor,
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 86f60e3a7..169eff73c 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -23,6 +23,7 @@
 
 from tensor2tensor.data_generators import gym_env
 from tensor2tensor.utils import registry
+from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -44,7 +45,7 @@
 
 
 def _rlmb_base():
-  return tf.contrib.training.HParams(
+  return HParams(
       epochs=15,
       # Total frames used for training. This will be distributed evenly across
       # hparams.epochs.
@@ -814,7 +815,7 @@ def merge_unscoped_hparams(scopes_and_hparams):
       scoped_key = "%s.%s" % (scope, key)
       merged_values[scoped_key] = value
 
-  return tf.contrib.training.HParams(**merged_values)
+  return HParams(**merged_values)
 
 
 def split_scoped_hparams(scopes, merged_hparams):
@@ -827,7 +828,7 @@ def split_scoped_hparams(scopes, merged_hparams):
     split_values[scope][key] = value
 
   return [
-      tf.contrib.training.HParams(**split_values[scope]) for scope in scopes
+      HParams(**split_values[scope]) for scope in scopes
   ]
 
 
@@ -881,7 +882,7 @@ def dynamic_register_hparams(name, hparams):
 
   @registry.register_hparams(name)
   def new_hparams_set():
-    return tf.contrib.training.HParams(**hparams.values())
+    return HParams(**hparams.values())
 
   return new_hparams_set
 
diff --git a/tensor2tensor/serving/query.py b/tensor2tensor/serving/query.py
index 1af0e9f2d..a8c46ce80 100644
--- a/tensor2tensor/serving/query.py
+++ b/tensor2tensor/serving/query.py
@@ -27,6 +27,7 @@
 from tensor2tensor.serving import serving_utils
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import usr_dir
+from tensor2tensor.utils.hparam import HParams
 import tensorflow as tf
 
 flags = tf.flags
@@ -80,7 +81,7 @@ def main(_):
   validate_flags()
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
   problem = registry.problem(FLAGS.problem)
-  hparams = tf.contrib.training.HParams(
+  hparams = HParams(
       data_dir=os.path.expanduser(FLAGS.data_dir))
   problem.get_hparams(hparams)
   request_fn = make_request_fn()
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 3ee5a88e9..7f0d5bbc1 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -35,6 +35,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import registry
+from tensor2tensor.utils.hparam import HParams
 import tensorflow as tf
 
 FLAGS = tf.flags.FLAGS
@@ -45,7 +46,7 @@
 
 def decode_hparams(overrides=""):
   """Hyperparameters for decoding."""
-  hp = tf.contrib.training.HParams(
+  hp = HParams(
       save_images=False,
       log_results=True,
       extra_length=100,
diff --git a/tensor2tensor/utils/diet.py b/tensor2tensor/utils/diet.py
index 700957280..a8d5d7c2f 100644
--- a/tensor2tensor/utils/diet.py
+++ b/tensor2tensor/utils/diet.py
@@ -27,6 +27,7 @@
 import math
 
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils.hparam import HParams
 import tensorflow as tf
 
 
@@ -36,7 +37,7 @@ def diet_adam_optimizer_params():
   Returns:
     a hyperparameters object.
   """
-  return tf.contrib.training.HParams(
+  return HParams(
       quantize=True,  # use 16-bit fixed-point
       quantization_scale=10.0 / tf.int16.max,
       optimizer="DietAdam",
diff --git a/tensor2tensor/utils/hparams_lib.py b/tensor2tensor/utils/hparams_lib.py
index 8aee611c9..e72c4f3f8 100644
--- a/tensor2tensor/utils/hparams_lib.py
+++ b/tensor2tensor/utils/hparams_lib.py
@@ -23,13 +23,14 @@
 
 from tensor2tensor.data_generators import problem as problem_lib
 from tensor2tensor.utils import registry
+from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
 
 def copy_hparams(hparams):
   hp_vals = hparams.values()
-  new_hparams = tf.contrib.training.HParams(**hp_vals)
+  new_hparams = HParams(**hp_vals)
   other_attrs = ["problem", "problem_hparams"]
   for attr in other_attrs:
     attr_val = getattr(hparams, attr, None)
@@ -63,7 +64,7 @@ def create_hparams_from_json(json_path, hparams=None):
   tf.logging.info("Loading hparams from existing json %s" % json_path)
   with tf.gfile.Open(json_path, "r") as f:
     hparams_values = json.load(f)
-    new_hparams = tf.contrib.training.HParams(**hparams_values)
+    new_hparams = HParams(**hparams_values)
     # Some keys are in new_hparams but not hparams, so we need to be more
     #   careful than simply using parse_json() from HParams
     if hparams:  # hparams specified, so update values from json
diff --git a/tensor2tensor/utils/misc_utils_test.py b/tensor2tensor/utils/misc_utils_test.py
index ab8da2bd1..11228de68 100644
--- a/tensor2tensor/utils/misc_utils_test.py
+++ b/tensor2tensor/utils/misc_utils_test.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 from tensor2tensor.utils import misc_utils
+from tensor2tensor.utils.hparam import HParams
 import tensorflow as tf
 
 
@@ -56,7 +57,7 @@ def test_snakecase_to_camelcase(self):
                      misc_utils.snakecase_to_camelcase("lstm_seq2_seq"))
 
   def test_pprint_hparams(self):
-    hparams = tf.contrib.training.HParams(
+    hparams = HParams(
         int_=1, str_="str", bool_=True, float_=1.1, list_int=[1, 2], none=None)
 
     # pylint: disable=g-inconsistent-quotes
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index cdade2369..417427fd1 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -155,9 +155,9 @@ def __init__(self,
     """Creates a T2TModel.
 
     Args:
-      hparams: tf.contrib.training.HParams, model hyperparameters.
+      hparams: HParams, model hyperparameters.
       mode: tf.estimator.ModeKeys, the execution mode.
-      problem_hparams: tf.contrib.training.HParams, hyperparameters for the
+      problem_hparams: HParams, hyperparameters for the
         Problem. If provided here or in hparams.problem_hparams, the model will
         automatically determine bottom, top, and loss methods. If not provided,
         calling the model will only invoke body.
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index 07622fbef..7e8cef93a 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -21,6 +21,7 @@
 
 from tensor2tensor.utils import modality
 from tensor2tensor.utils import t2t_model
+from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -30,7 +31,7 @@ class T2TModelTest(tf.test.TestCase):
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testSummarizeLosses(self):
     with tf.Graph().as_default():
-      model = t2t_model.T2TModel(tf.contrib.training.HParams())
+      model = t2t_model.T2TModel(HParams())
       losses = {"training": tf.random_normal([]),
                 "extra": tf.random_normal([])}
       outputs = model._summarize_losses(losses)
@@ -47,11 +48,11 @@ def testLossSingleWeights(self):
         sequence_size = 16
         vocab_size = 3
 
-        model_hparams = tf.contrib.training.HParams(
+        model_hparams = HParams(
             label_smoothing=0.0,
             shared_embedding_and_softmax_weights=False)
 
-        problem_hparams = tf.contrib.training.HParams(loss_multiplier=1.0)
+        problem_hparams = HParams(loss_multiplier=1.0)
         problem_hparams.modality = {}
 
         model = t2t_model.T2TModel(
diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index 5a58c2471..177916021 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -22,6 +22,8 @@
 import os
 import numpy as np
 import six
+
+
 import tensorflow as tf
 
 
@@ -224,7 +226,7 @@ def compute_video_metrics_from_predictions(predictions, decode_hparams):
   Args:
     predictions: list of list of dicts.
                  outer length: num_decodes, inner_length: num_samples
-    decode_hparams: Decode hparams. instance of tf.contrib.training.HParams.
+    decode_hparams: Decode hparams. instance of HParams.
   Returns:
     statistics: dict of Tensors, key being the metric with each Tensor
                 having the shape (num_samples, num_frames).

From f7c64216cb3013331e45dbc35b51f79010370b3e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 1 Feb 2019 11:39:14 -0800
Subject: [PATCH 1620/2720] Evolved Transformer encoder.

PiperOrigin-RevId: 232009541
---
 tensor2tensor/layers/transformer_layers.py | 164 +++++++++++++++++++++
 1 file changed, 164 insertions(+)

diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index d52d7cb71..c63151286 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -217,6 +217,170 @@ def transformer_encoder(encoder_input,
     return common_layers.layer_preprocess(x, hparams)
 
 
+def evolved_transformer_encoder(encoder_input,
+                                encoder_self_attention_bias,
+                                hparams,
+                                name="encoder",
+                                nonpadding=None,
+                                save_weights_to=None,
+                                make_image_summary=True,
+                                losses=None,
+                                attn_bias_for_padding=None):
+  """Evolved Transformer encoder. See arxiv.org/abs/1901.11117 for more details.
+
+  Note: Pad remover is not supported.
+
+  Args:
+    encoder_input: a Tensor.
+    encoder_self_attention_bias: bias Tensor for self-attention (see
+      common_attention.attention_bias()).
+    hparams: hyperparameters for model.
+    name: a string.
+    nonpadding: optional Tensor with shape [batch_size, encoder_length]
+      indicating what positions are not padding.  This must either be passed in,
+      which we do for "packed" datasets, or inferred from
+      encoder_self_attention_bias.  The knowledge about padding is used for
+      pad_remover(efficiency) and to mask out padding in convolutional layers.
+    save_weights_to: an optional dictionary to capture attention weights for
+      visualization; the weights tensor will be appended there under a string
+      key created from the variable scope (including name).
+    make_image_summary: Whether to make an attention image summary.
+    losses: Not used.
+    attn_bias_for_padding: Padded attention bias in case a unidirectional
+      encoder is being used where future attention is masked.
+
+  Returns:
+    Tensor encoder output.
+  """
+  del losses
+
+  hidden_state = encoder_input
+  attention_dropout_broadcast_dims = (
+      common_layers.comma_separated_string_to_integer_list(
+          getattr(hparams, "attention_dropout_broadcast_dims", "")))
+
+  with tf.variable_scope(name):
+    if nonpadding is not None:
+      padding = 1.0 - nonpadding
+    else:
+      attention_bias = encoder_self_attention_bias
+      if attn_bias_for_padding is not None:
+        attention_bias = attn_bias_for_padding
+      padding = common_attention.attention_bias_to_padding(attention_bias)
+      nonpadding = 1.0 - padding
+
+    for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):
+      with tf.variable_scope("layer_%d" % layer):
+
+        with tf.variable_scope("gated_linear_unit"):
+
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          values = tf.layers.dense(hidden_state, hparams.hidden_size)
+          gates = tf.layers.dense(
+              hidden_state, hparams.hidden_size, activation=tf.nn.sigmoid)
+          hidden_state = values * gates
+
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, hidden_state, hparams)
+
+        with tf.variable_scope("conv_branches"):
+
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+          # Mask padding from conv layers.
+          mask = tf.tile(
+              tf.expand_dims(nonpadding, 2), [1, 1, hparams.hidden_size])
+          hidden_state *= mask
+
+          left_output_dim = int(hparams.hidden_size * 4)
+          left_state = tf.layers.dense(
+              hidden_state, left_output_dim, activation=tf.nn.relu)
+          left_state = tf.nn.dropout(left_state,
+                                     1 - hparams.layer_prepostprocess_dropout)
+
+          right_output_dim = int(hparams.hidden_size / 2)
+          right_state = tf.layers.conv1d(
+              hidden_state,
+              right_output_dim,
+              3,
+              padding="SAME",
+              name="standard_conv_3x1",
+              activation=tf.nn.relu)
+          right_state = tf.nn.dropout(right_state,
+                                      1 - hparams.layer_prepostprocess_dropout)
+
+          right_state = tf.pad(
+              right_state,
+              [[0, 0], [0, 0], [0, left_output_dim - right_output_dim]],
+              constant_values=0)
+          hidden_state = left_state + right_state
+
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+          # Mask padding from conv layer.
+          mask = tf.tile(tf.expand_dims(nonpadding, 2), [1, 1, left_output_dim])
+          hidden_state *= mask
+
+          separable_conv_9x1 = tf.layers.SeparableConv1D(
+              right_output_dim, 9, padding="SAME", name="separable_conv_9x1")
+          hidden_state = separable_conv_9x1.apply(hidden_state)
+          hidden_state = tf.pad(
+              hidden_state,
+              [[0, 0], [0, 0], [0, hparams.hidden_size - right_output_dim]],
+              constant_values=0)
+
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, hidden_state, hparams)
+
+        with tf.variable_scope("self_attention"):
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          hidden_state = common_attention.multihead_attention(
+              hidden_state,
+              None,
+              encoder_self_attention_bias,
+              hparams.attention_key_channels or hparams.hidden_size,
+              hparams.attention_value_channels or hparams.hidden_size,
+              hparams.hidden_size,
+              hparams.num_heads,
+              hparams.attention_dropout,
+              attention_type=hparams.self_attention_type,
+              max_relative_position=hparams.max_relative_position,
+              heads_share_relative_embedding=(
+                  hparams.heads_share_relative_embedding),
+              add_relative_to_values=hparams.add_relative_to_values,
+              save_weights_to=save_weights_to,
+              make_image_summary=make_image_summary,
+              dropout_broadcast_dims=attention_dropout_broadcast_dims,
+              max_length=hparams.get("max_length"),
+              vars_3d=hparams.get("attention_variables_3d"),
+              activation_dtype=hparams.get("activation_dtype", "float32"),
+              weight_dtype=hparams.get("weight_dtype", "float32"))
+
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, hidden_state, hparams)
+
+        with tf.variable_scope("dense_layers"):
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          hidden_state = tf.layers.dense(
+              hidden_state, int(hparams.hidden_size * 4), activation=tf.nn.relu)
+          hidden_state = tf.nn.dropout(hidden_state,
+                                       1 - hparams.layer_prepostprocess_dropout)
+
+          hidden_state = tf.layers.dense(hidden_state, hparams.hidden_size)
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, hidden_state, hparams)
+
+    # If normalization is done in layer_preprocess, then it should also be done
+    # on the output, since the output can grow very large, being the sum of
+    # a whole stack of unnormalized layer outputs.
+    return common_layers.layer_preprocess(hidden_state, hparams)
+
+
 def transformer_ffn_layer(x,
                           hparams,
                           pad_remover=None,

From e6ef7b25fd5034f833e4ef50ac5e196457421a76 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 1 Feb 2019 11:57:10 -0800
Subject: [PATCH 1621/2720] Add English-German translation data with Paracrawl
 and remove an obsolete one.

PiperOrigin-RevId: 232012614
---
 tensor2tensor/data_generators/translate.py    | 29 ++++++-
 .../data_generators/translate_ende.py         | 77 ++++++-------------
 tensor2tensor/models/transformer.py           |  2 +-
 3 files changed, 52 insertions(+), 56 deletions(-)

diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 82f7eeb8f..e9f8356fa 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -144,7 +144,33 @@ def compile_data(tmp_dir, datasets, filename):
         if url.startswith("http"):
           generator_utils.maybe_download(tmp_dir, compressed_filename, url)
 
-        if dataset[1][0] == "tsv":
+        if dataset[1][0] == "tmx":
+          tmx_filename = os.path.join(tmp_dir, dataset[1][1])
+          if tmx_filename.endswith(".gz"):
+            new_filename = tmx_filename.strip(".gz")
+            if not tf.gfile.Exists(new_filename):
+              generator_utils.gunzip_file(tmx_filename, new_filename)
+            tmx_filename = new_filename
+          source, target = None, None
+          with tf.gfile.Open(tmx_filename) as tmx_file:
+            for line in tmx_file:
+              text = line.strip()
+              if text.startswith("<seg>"):
+                if text.endswith("</seg>"):
+                  sentence = text[5:-6]  # Strip <seg> and </seg>.
+                  if source is None:
+                    source = sentence
+                  else:
+                    target = sentence
+              if source is not None and target is not None:
+                if source and target:  # Prevent empty string examples.
+                  lang1_resfile.write(source)
+                  lang1_resfile.write("\n")
+                  lang2_resfile.write(target)
+                  lang2_resfile.write("\n")
+                source, target = None, None
+
+        elif dataset[1][0] == "tsv":
           _, src_column, trg_column, glob_pattern = dataset[1]
           filenames = tf.gfile.Glob(os.path.join(tmp_dir, glob_pattern))
           if not filenames:
@@ -169,6 +195,7 @@ def compile_data(tmp_dir, datasets, filename):
                     lang1_resfile.write("\n")
                     lang2_resfile.write(target)
                     lang2_resfile.write("\n")
+
         else:
           lang1_filename, lang2_filename = dataset[1]
           lang1_filepath = os.path.join(tmp_dir, lang1_filename)
diff --git a/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index 259236967..80a34c28c 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -19,17 +19,12 @@
 from __future__ import division
 from __future__ import print_function
 
-import os
-import tarfile
-from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
-from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.data_generators import translate
 from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
 
 _ENDE_TRAIN_DATASETS = [
     [
@@ -46,7 +41,7 @@
         ("training/europarl-v7.de-en.en", "training/europarl-v7.de-en.de")
     ],
 ]
-_ENDE_TEST_DATASETS = [
+_ENDE_EVAL_DATASETS = [
     [
         "http://data.statmt.org/wmt17/translation-task/dev.tgz",
         ("dev/newstest2013.en", "dev/newstest2013.de")
@@ -54,53 +49,6 @@
 ]
 
 
-def _get_wmt_ende_bpe_dataset(directory, filename):
-  """Extract the WMT en-de corpus `filename` to directory unless it's there."""
-  train_path = os.path.join(directory, filename)
-  if not (tf.gfile.Exists(train_path + ".de") and
-          tf.gfile.Exists(train_path + ".en")):
-    url = ("https://drive.google.com/uc?export=download&id="
-           "0B_bZck-ksdkpM25jRUN2X2UxMm8")
-    corpus_file = generator_utils.maybe_download_from_drive(
-        directory, "wmt16_en_de.tar.gz", url)
-    with tarfile.open(corpus_file, "r:gz") as corpus_tar:
-      corpus_tar.extractall(directory)
-  return train_path
-
-
-@registry.register_problem
-class TranslateEndeWmtBpe32k(translate.TranslateProblem):
-  """Problem spec for WMT En-De translation, BPE version."""
-
-  @property
-  def vocab_type(self):
-    return text_problems.VocabType.TOKEN
-
-  @property
-  def oov_token(self):
-    return "UNK"
-
-  def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    """Instance of token generator for the WMT en->de task, training set."""
-    train = dataset_split == problem.DatasetSplit.TRAIN
-    dataset_path = ("train.tok.clean.bpe.32000"
-                    if train else "newstest2013.tok.bpe.32000")
-    train_path = _get_wmt_ende_bpe_dataset(tmp_dir, dataset_path)
-
-    # Vocab
-    vocab_path = os.path.join(data_dir, self.vocab_filename)
-    if not tf.gfile.Exists(vocab_path):
-      bpe_vocab = os.path.join(tmp_dir, "vocab.bpe.32000")
-      with tf.gfile.Open(bpe_vocab) as f:
-        vocab_list = f.read().split("\n")
-      vocab_list.append(self.oov_token)
-      text_encoder.TokenTextEncoder(
-          None, vocab_list=vocab_list).store_to_file(vocab_path)
-
-    return text_problems.text2text_txt_iterator(train_path + ".en",
-                                                train_path + ".de")
-
-
 @registry.register_problem
 class TranslateEndeWmt8k(translate.TranslateProblem):
   """Problem spec for WMT En-De translation."""
@@ -109,9 +57,15 @@ class TranslateEndeWmt8k(translate.TranslateProblem):
   def approx_vocab_size(self):
     return 2**13  # 8192
 
+  @property
+  def additional_training_datasets(self):
+    """Allow subclasses to add training datasets."""
+    return []
+
   def source_data_files(self, dataset_split):
     train = dataset_split == problem.DatasetSplit.TRAIN
-    return _ENDE_TRAIN_DATASETS if train else _ENDE_TEST_DATASETS
+    train_datasets = _ENDE_TRAIN_DATASETS + self.additional_training_datasets
+    return train_datasets if train else _ENDE_EVAL_DATASETS
 
 
 @registry.register_problem
@@ -122,6 +76,21 @@ def approx_vocab_size(self):
     return 2**15  # 32768
 
 
+@registry.register_problem
+class TranslateEndeWmtParacrawlBicleaner32k(TranslateEndeWmt32k):
+  """WMT en-de corpus with extra data from Paracrawl, cleaned with Bicleaner."""
+
+  @property
+  def vocab_filename(self):
+    return TranslateEndeWmt32k().vocab_filename
+
+  @property
+  def additional_training_datasets(self):
+    paracrawl = "https://s3.amazonaws.com/web-language-models/paracrawl/"
+    return [(paracrawl + "release3/en-de.bicleaner07.tmx.gz",
+             ("tmx", "en-de.bicleaner07.tmx.gz"))]
+
+
 @registry.register_problem
 class TranslateEndeWmt32kPacked(TranslateEndeWmt32k):
 
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index d7334bb25..e76fe8f9f 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1730,7 +1730,7 @@ def transformer_tall_pretrain_lm_tpu_adafactor_large():
   hparams.batch_size = 4
   hparams.multiproblem_mixing_schedule = "constant"
   # Task order: lm/en-de/en-fr/en-ro/de-en/fr-en/ro-en/cnndm/mnli/squad.
-  hparams.multiproblem_per_task_threshold = "320,80,160,2,80,160,2,20,5,5"
+  hparams.multiproblem_per_task_threshold = "320,80,160,1,80,160,2,20,10,5"
   return hparams
 
 

From 9626aff2f5b9ce9ca03aa576811af550e3c5ccff Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 1 Feb 2019 12:14:08 -0800
Subject: [PATCH 1622/2720] Evolved Transformer decoder.

PiperOrigin-RevId: 232015730
---
 tensor2tensor/models/transformer.py | 400 +++++++++++++++++++++++-----
 1 file changed, 326 insertions(+), 74 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index e76fe8f9f..42f9fd7dd 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -46,7 +46,6 @@
 from tensorflow.python.util import nest
 # pylint: enable=g-direct-tensorflow-import
 
-
 # Alias some commonly reused layers, here and elsewhere.
 transformer_prepare_encoder = transformer_layers.transformer_prepare_encoder
 transformer_encoder = transformer_layers.transformer_encoder
@@ -69,8 +68,8 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
         will be flattened along the two spatial dimensions.
       target_space: scalar, target space ID.
       hparams: hyperparameters for model.
-      features: optionally pass the entire features dictionary as well.
-        This is needed now for "packed" datasets.
+      features: optionally pass the entire features dictionary as well. This is
+        needed now for "packed" datasets.
       losses: optional list onto which to append extra training losses
 
     Returns:
@@ -124,19 +123,19 @@ def decode(self,
     """Decode Transformer outputs from encoder representation.
 
     Args:
-      decoder_input: inputs to bottom of the model.
-          [batch_size, decoder_length, hidden_dim]
-      encoder_output: Encoder representation.
-          [batch_size, input_length, hidden_dim]
-      encoder_decoder_attention_bias: Bias and mask weights for
-          encoder-decoder attention. [batch_size, input_length]
+      decoder_input: inputs to bottom of the model. [batch_size, decoder_length,
+        hidden_dim]
+      encoder_output: Encoder representation. [batch_size, input_length,
+        hidden_dim]
+      encoder_decoder_attention_bias: Bias and mask weights for encoder-decoder
+        attention. [batch_size, input_length]
       decoder_self_attention_bias: Bias and mask weights for decoder
-          self-attention. [batch_size, decoder_length]
+        self-attention. [batch_size, decoder_length]
       hparams: hyperparameters for model.
       cache: dict, containing tensors which are the results of previous
-          attentions, used for fast decoding.
-      decode_loop_step: An integer, step number of the decoding loop.
-          Only used for inference on TPU.
+        attentions, used for fast decoding.
+      decode_loop_step: An integer, step number of the decoding loop. Only used
+        for inference on TPU.
       nonpadding: optional Tensor with shape [batch_size, decoder_length]
       losses: optional list onto which to append extra training losses
 
@@ -176,10 +175,10 @@ def body(self, features):
 
     Args:
       features: Map of features to the model. Should contain the following:
-          "inputs": Transformer inputs.
-              [batch_size, input_length, 1, hidden_dim].
-          "targets": Target decoder outputs.
-              [batch_size, decoder_length, 1, hidden_dim]
+          "inputs": Transformer inputs. [batch_size, input_length, 1,
+            hidden_dim].
+          "targets": Target decoder outputs. [batch_size, decoder_length, 1,
+            hidden_dim]
           "target_space_id": A scalar int from data_generators.problem.SpaceID.
 
     Returns:
@@ -248,7 +247,7 @@ def _greedy_infer(self, features, decode_length, use_tpu=False):
     # For real-valued modalities use the slow decode path for now.
     if (self._target_modality_is_real or
         self._hparams.self_attention_type != "dot_product"):
-      return  super(Transformer, self)._greedy_infer(features, decode_length)
+      return super(Transformer, self)._greedy_infer(features, decode_length)
     with tf.variable_scope(self.name):
       if use_tpu:
         return self._fast_decode_tpu(features, decode_length)
@@ -281,8 +280,9 @@ def _beam_decode(self,
               None if using greedy decoding (beam_size=1)
       }
     """
-    if (self._hparams.self_attention_type not in ["dot_product",
-                                                  "dot_product_relative"]):
+    if (self._hparams.self_attention_type not in [
+        "dot_product", "dot_product_relative"
+    ]):
       # Caching is not guaranteed to work with attention types other than
       # dot_product.
       # TODO(petershaw): Support fast decoding when using relative
@@ -291,10 +291,10 @@ def _beam_decode(self,
                                     top_beams, alpha, use_tpu)
     with tf.variable_scope(self.name):
       if use_tpu:
-        return self._fast_decode_tpu(
-            features, decode_length, beam_size, top_beams, alpha)
-      return self._fast_decode(
-          features, decode_length, beam_size, top_beams, alpha)
+        return self._fast_decode_tpu(features, decode_length, beam_size,
+                                     top_beams, alpha)
+      return self._fast_decode(features, decode_length, beam_size, top_beams,
+                               alpha)
 
   def _fast_decode_tpu(self,
                        features,
@@ -443,9 +443,9 @@ def symbols_to_logits_tpu_fn(ids, i, cache):
       Args:
         ids: A tensor, symbol IDs.
         i: An integer, step number of the decoding loop. Only used for inference
-            on TPU.
+          on TPU.
         cache: A dict, containing tensors which are the results of previous
-            attentions, used for fast decoding.
+          attentions, used for fast decoding.
 
       Returns:
         ret: A tensor, computed logits.
@@ -611,8 +611,8 @@ def _fast_decode(self,
           decode_length + 1, hparams.hidden_size)
     elif hparams.pos == "emb":
       positional_encoding = common_attention.add_positional_embedding(
-          tf.zeros([1, decode_length, hparams.hidden_size]),
-          hparams.max_length, "body/targets_positional_embedding", None)
+          tf.zeros([1, decode_length, hparams.hidden_size]), hparams.max_length,
+          "body/targets_positional_embedding", None)
     else:
       positional_encoding = None
 
@@ -733,9 +733,9 @@ def fast_decode_tpu(encoder_output,
   Args:
     encoder_output: A tensor, output from encoder.
     encoder_decoder_attention_bias: A tensor, bias for use in encoder-decoder
-        attention.
-    symbols_to_logits_fn: Incremental decoding, function mapping triple
-        `(ids, step, cache)` to symbol logits.
+      attention.
+    symbols_to_logits_fn: Incremental decoding, function mapping triple `(ids,
+      step, cache)` to symbol logits.
     hparams: Run hyperparameters.
     decode_length: An integer, how many additional timesteps to decode.
     vocab_size: Output vocabulary size.
@@ -747,7 +747,7 @@ def fast_decode_tpu(encoder_output,
     eos_id: End-of-sequence symbol.
     batch_size: An integer, must be passed if there is no input.
     force_decode_length: A bool, whether to force the full decode length, or if
-        False, stop when all beams hit eos_id.
+      False, stop when all beams hit eos_id.
     scope_prefix: str, prefix for decoder layer variable scopes.
 
   Returns:
@@ -797,15 +797,18 @@ def fast_decode_tpu(encoder_output,
   if encoder_output is not None:
     for layer in range(num_layers):
       layer_name = "layer_%d" % layer
-      with tf.variable_scope(
-          "%sdecoder/%s/encdec_attention/multihead_attention" % (scope_prefix,
-                                                                 layer_name)):
+      with tf.variable_scope("%sdecoder/%s/encdec_attention/multihead_attention"
+                             % (scope_prefix, layer_name)):
         k_encdec = common_attention.compute_attention_component(
-            encoder_output, key_channels, name="k",
+            encoder_output,
+            key_channels,
+            name="k",
             vars_3d_num_heads=vars_3d_num_heads)
         k_encdec = common_attention.split_heads(k_encdec, hparams.num_heads)
         v_encdec = common_attention.compute_attention_component(
-            encoder_output, value_channels, name="v",
+            encoder_output,
+            value_channels,
+            name="v",
             vars_3d_num_heads=vars_3d_num_heads)
         v_encdec = common_attention.split_heads(v_encdec, hparams.num_heads)
       cache[layer_name]["k_encdec"] = k_encdec
@@ -845,6 +848,7 @@ def fast_decode_tpu(encoder_output,
       decoded_ids = decoded_ids[:, :top_beams, 1:]
       scores = scores[:, :top_beams]
   else:  # Greedy
+
     def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       """One step of greedy decoding."""
       logits, cache = symbols_to_logits_fn(next_id, i, cache)
@@ -855,8 +859,8 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       next_id = common_layers.sample_with_temperature(logits, temperature)
       hit_eos |= tf.equal(next_id, eos_id)
 
-      log_prob_indices = tf.stack(
-          [tf.range(tf.to_int64(batch_size)), next_id], axis=1)
+      log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],
+                                  axis=1)
       log_prob += tf.gather_nd(log_probs, log_prob_indices)
 
       next_id = tf.expand_dims(next_id, axis=1)
@@ -923,8 +927,8 @@ def fast_decode(encoder_output,
     encoder_output: Output from encoder.
     encoder_decoder_attention_bias: a bias tensor for use in encoder-decoder
       attention
-    symbols_to_logits_fn: Incremental decoding; function mapping triple
-      `(ids, step, cache)` to symbol logits.
+    symbols_to_logits_fn: Incremental decoding; function mapping triple `(ids,
+      step, cache)` to symbol logits.
     hparams: run hyperparameters
     decode_length: an integer.  How many additional timesteps to decode.
     vocab_size: Output vocabulary size.
@@ -987,15 +991,18 @@ def fast_decode(encoder_output,
   if encoder_output is not None:
     for layer in range(num_layers):
       layer_name = "layer_%d" % layer
-      with tf.variable_scope(
-          "%sdecoder/%s/encdec_attention/multihead_attention" % (scope_prefix,
-                                                                 layer_name)):
+      with tf.variable_scope("%sdecoder/%s/encdec_attention/multihead_attention"
+                             % (scope_prefix, layer_name)):
         k_encdec = common_attention.compute_attention_component(
-            encoder_output, key_channels, name="k",
+            encoder_output,
+            key_channels,
+            name="k",
             vars_3d_num_heads=vars_3d_num_heads)
         k_encdec = common_attention.split_heads(k_encdec, hparams.num_heads)
         v_encdec = common_attention.compute_attention_component(
-            encoder_output, value_channels, name="v",
+            encoder_output,
+            value_channels,
+            name="v",
             vars_3d_num_heads=vars_3d_num_heads)
         v_encdec = common_attention.split_heads(v_encdec, hparams.num_heads)
       cache[layer_name]["k_encdec"] = k_encdec
@@ -1035,8 +1042,8 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       next_id = common_layers.sample_with_temperature(logits, temperature)
       hit_eos |= tf.equal(next_id, eos_id)
 
-      log_prob_indices = tf.stack(
-          [tf.range(tf.to_int64(batch_size)), next_id], axis=1)
+      log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],
+                                  axis=1)
       log_prob += tf.gather_nd(log_probs, log_prob_indices)
 
       next_id = tf.expand_dims(next_id, axis=1)
@@ -1174,8 +1181,8 @@ def transformer_prepare_decoder(targets, hparams, features=None):
   Args:
     targets: a Tensor.
     hparams: run hyperparameters
-    features: optionally pass the entire features dictionary as well.
-      This is needed now for "packed" datasets.
+    features: optionally pass the entire features dictionary as well. This is
+      needed now for "packed" datasets.
 
   Returns:
     decoder_input: a Tensor, bottom of decoder stack
@@ -1243,24 +1250,24 @@ def transformer_decoder(decoder_input,
   Args:
     decoder_input: a Tensor
     encoder_output: a Tensor
-    decoder_self_attention_bias: bias Tensor for self-attention
-      (see common_attention.attention_bias())
+    decoder_self_attention_bias: bias Tensor for self-attention (see
+      common_attention.attention_bias())
     encoder_decoder_attention_bias: bias Tensor for encoder-decoder attention
       (see common_attention.attention_bias())
     hparams: hyperparameters for model
     cache: dict, containing tensors which are the results of previous
-        attentions, used for fast decoding.
-    decode_loop_step: An integer, step number of the decoding loop.
-        Only used for inference on TPU.
+      attentions, used for fast decoding.
+    decode_loop_step: An integer, step number of the decoding loop. Only used
+      for inference on TPU.
     name: a string
     nonpadding: optional Tensor with shape [batch_size, encoder_length]
-      indicating what positions are not padding.  This is used
-      to mask out padding in convolutional layers.  We generally only
-      need this mask for "packed" datasets, because for ordinary datasets,
-      no padding is ever followed by nonpadding.
-    save_weights_to: an optional dictionary to capture attention weights
-      for visualization; the weights tensor will be appended there under
-      a string key created from the variable scope (including name).
+      indicating what positions are not padding.  This is used to mask out
+      padding in convolutional layers.  We generally only need this mask for
+      "packed" datasets, because for ordinary datasets, no padding is ever
+      followed by nonpadding.
+    save_weights_to: an optional dictionary to capture attention weights for
+      visualization; the weights tensor will be appended there under a string
+      key created from the variable scope (including name).
     make_image_summary: Whether to make an attention image summary.
     losses: optional list onto which to append extra training losses
 
@@ -1363,6 +1370,257 @@ def transformer_decoder(decoder_input,
     return common_layers.layer_preprocess(x, hparams)
 
 
+def evolved_transformer_decoder(decoder_input,
+                                encoder_output,
+                                decoder_self_attention_bias,
+                                encoder_decoder_attention_bias,
+                                hparams,
+                                cache=None,
+                                decode_loop_step=None,
+                                name="decoder",
+                                nonpadding=None,
+                                save_weights_to=None,
+                                make_image_summary=True,
+                                losses=None):
+  """Evolved Transformer decoder. See arxiv.org/abs/1901.11117 for more details.
+
+  Args:
+    decoder_input: a Tensor.
+    encoder_output: a Tensor.
+    decoder_self_attention_bias: bias Tensor for self-attention (see
+      common_attention.attention_bias()).
+    encoder_decoder_attention_bias: bias Tensor for encoder-decoder attention
+      (see common_attention.attention_bias()).
+    hparams: hyperparameters for model.
+    cache: Not supported.
+    decode_loop_step: An integer, step number of the decoding loop. Only used
+      for inference on TPU.
+    name: a string.
+    nonpadding: optional Tensor with shape [batch_size, encoder_length]
+      indicating what positions are not padding.  This is used to mask out
+      padding in convolutional layers.  We generally only need this mask for
+      "packed" datasets, because for ordinary datasets, no padding is ever
+      followed by nonpadding.
+    save_weights_to: an optional dictionary to capture attention weights for
+      visualization; the weights tensor will be appended there under a string
+      key created from the variable scope (including name).
+    make_image_summary: Whether to make an attention image summary.
+    losses: Not supported.
+
+  Returns:
+    Decoder output tensor.
+  """
+  del cache, losses
+
+  attention_dropout_broadcast_dims = (
+      common_layers.comma_separated_string_to_integer_list(
+          getattr(hparams, "attention_dropout_broadcast_dims", "")))
+
+  with tf.variable_scope(name):
+    hidden_state = decoder_input
+    layer_cache = None
+
+    for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
+      with tf.variable_scope("layer_%d" % layer):
+
+        with tf.variable_scope("16_head_self_attention"):
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          # 16 head attention. Hard coding number of heads.
+          left_state = common_attention.multihead_attention(
+              hidden_state,
+              None,
+              decoder_self_attention_bias,
+              hparams.attention_key_channels or hparams.hidden_size,
+              hparams.attention_value_channels or hparams.hidden_size,
+              hparams.hidden_size,
+              16,  # Heads are hard coded to replicate paper.
+              hparams.attention_dropout,
+              attention_type=hparams.self_attention_type,
+              max_relative_position=hparams.max_relative_position,
+              heads_share_relative_embedding=(
+                  hparams.heads_share_relative_embedding),
+              add_relative_to_values=hparams.add_relative_to_values,
+              save_weights_to=save_weights_to,
+              cache=layer_cache,
+              make_image_summary=make_image_summary,
+              dropout_broadcast_dims=attention_dropout_broadcast_dims,
+              max_length=hparams.get("max_length"),
+              decode_loop_step=decode_loop_step,
+              vars_3d=hparams.get("attention_variables_3d"),
+              activation_dtype=hparams.get("activation_dtype", "float32"),
+              weight_dtype=hparams.get("weight_dtype", "float32"))
+
+        if encoder_output is not None:
+          with tf.variable_scope("first_attend_to_encoder"):
+            right_state = common_attention.multihead_attention(
+                hidden_state,
+                encoder_output,
+                encoder_decoder_attention_bias,
+                hparams.attention_key_channels or hparams.hidden_size,
+                hparams.attention_value_channels or hparams.hidden_size,
+                hparams.hidden_size,
+                hparams.num_heads,
+                hparams.attention_dropout,
+                max_relative_position=hparams.max_relative_position,
+                heads_share_relative_embedding=(
+                    hparams.heads_share_relative_embedding),
+                add_relative_to_values=hparams.add_relative_to_values,
+                save_weights_to=save_weights_to,
+                cache=layer_cache,
+                make_image_summary=make_image_summary,
+                dropout_broadcast_dims=attention_dropout_broadcast_dims,
+                max_length=hparams.get("max_length"),
+                vars_3d=hparams.get("attention_variables_3d"),
+                activation_dtype=hparams.get("activation_dtype", "float32"),
+                weight_dtype=hparams.get("weight_dtype", "float32"))
+
+            left_state = tf.nn.dropout(left_state,
+                                       1 - hparams.layer_prepostprocess_dropout)
+            right_state = tf.nn.dropout(
+                right_state, 1 - hparams.layer_prepostprocess_dropout)
+
+            hidden_state = residual_state + left_state + right_state
+
+        else:
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, left_state, hparams)
+
+        with tf.variable_scope("conv_branches"):
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          if nonpadding:
+            # Mask padding from conv layers.
+            mask = tf.tile(
+                tf.expand_dims(nonpadding, 2), [1, 1, hparams.hidden_size])
+            hidden_state *= mask
+
+          # Shift inputs so that future tokens cannot be seen.
+          left_state = tf.pad(hidden_state, paddings=[[0, 0], [10, 0], [0, 0]])
+          left_output_dim = int(hparams.hidden_size * 2)
+          separable_conv_11x1 = tf.layers.SeparableConv1D(
+              left_output_dim,
+              11,
+              padding="VALID",
+              name="separable_conv11x1",
+              activation=tf.nn.relu)
+          left_state = separable_conv_11x1.apply(left_state)
+          left_state = tf.nn.dropout(left_state,
+                                     1 - hparams.layer_prepostprocess_dropout)
+
+          right_state = tf.pad(hidden_state, paddings=[[0, 0], [6, 0], [0, 0]])
+          right_output_dim = int(hparams.hidden_size / 2)
+          separable_conv_7x1_1 = tf.layers.SeparableConv1D(
+              right_output_dim, 7, padding="VALID", name="separable_conv_7x1_1")
+          right_state = separable_conv_7x1_1.apply(right_state)
+          right_state = tf.nn.dropout(right_state,
+                                      1 - hparams.layer_prepostprocess_dropout)
+          right_state = tf.pad(
+              right_state,
+              [[0, 0], [0, 0], [0, left_output_dim - right_output_dim]],
+              constant_values=0)
+
+          hidden_state = left_state + right_state
+
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+          if nonpadding:
+            # Mask padding from conv layers.
+            mask = tf.tile(
+                tf.expand_dims(nonpadding, 2), [1, 1, hparams.hidden_size])
+            hidden_state *= mask
+
+          hidden_state = tf.pad(hidden_state, paddings=[[0, 0], [6, 0], [0, 0]])
+          separable_conv_7x1_2 = tf.layers.SeparableConv1D(
+              hparams.hidden_size,
+              7,
+              padding="VALID",
+              name="separable_conv_7x1_2")
+          hidden_state = separable_conv_7x1_2.apply(hidden_state)
+
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, hidden_state, hparams)
+
+        with tf.variable_scope("self_attention"):
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          hidden_state = common_attention.multihead_attention(
+              hidden_state,
+              None,
+              decoder_self_attention_bias,
+              hparams.attention_key_channels or hparams.hidden_size,
+              hparams.attention_value_channels or hparams.hidden_size,
+              hparams.hidden_size,
+              hparams.num_heads,
+              hparams.attention_dropout,
+              attention_type=hparams.self_attention_type,
+              max_relative_position=hparams.max_relative_position,
+              heads_share_relative_embedding=(
+                  hparams.heads_share_relative_embedding),
+              add_relative_to_values=hparams.add_relative_to_values,
+              save_weights_to=save_weights_to,
+              cache=layer_cache,
+              make_image_summary=make_image_summary,
+              dropout_broadcast_dims=attention_dropout_broadcast_dims,
+              max_length=hparams.get("max_length"),
+              decode_loop_step=decode_loop_step,
+              vars_3d=hparams.get("attention_variables_3d"),
+              activation_dtype=hparams.get("activation_dtype", "float32"),
+              weight_dtype=hparams.get("weight_dtype", "float32"))
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, hidden_state, hparams)
+
+        if encoder_output is not None:
+          with tf.variable_scope("second_attend_to_encoder"):
+            residual_state = hidden_state
+            hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+            hidden_state = common_attention.multihead_attention(
+                hidden_state,
+                encoder_output,
+                encoder_decoder_attention_bias,
+                hparams.attention_key_channels or hparams.hidden_size,
+                hparams.attention_value_channels or hparams.hidden_size,
+                hparams.hidden_size,
+                hparams.num_heads,
+                hparams.attention_dropout,
+                max_relative_position=hparams.max_relative_position,
+                heads_share_relative_embedding=(
+                    hparams.heads_share_relative_embedding),
+                add_relative_to_values=hparams.add_relative_to_values,
+                save_weights_to=save_weights_to,
+                cache=layer_cache,
+                make_image_summary=make_image_summary,
+                dropout_broadcast_dims=attention_dropout_broadcast_dims,
+                max_length=hparams.get("max_length"),
+                vars_3d=hparams.get("attention_variables_3d"),
+                activation_dtype=hparams.get("activation_dtype", "float32"),
+                weight_dtype=hparams.get("weight_dtype", "float32"))
+            hidden_state = common_layers.layer_postprocess(
+                residual_state, hidden_state, hparams)
+
+        with tf.variable_scope("dense_layers"):
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          hidden_state = tf.layers.dense(
+              hidden_state,
+              int(hparams.hidden_size * 4),
+              activation=tf.nn.swish)
+          hidden_state = tf.nn.dropout(hidden_state,
+                                       1 - hparams.layer_prepostprocess_dropout)
+
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          hidden_state = tf.layers.dense(hidden_state, hparams.hidden_size)
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, hidden_state, hparams)
+
+    return common_layers.layer_preprocess(hidden_state, hparams)
+
+
 @registry.register_hparams
 def transformer_base_v1():
   """Set of hyperparameters."""
@@ -1608,8 +1866,7 @@ def transformer_tall_finetune_tied():
   hparams.multiproblem_max_input_length = 750
   hparams.multiproblem_max_target_length = 100
   hparams.multiproblem_schedule_max_examples = 0
-  hparams.learning_rate_schedule = (
-      "linear_warmup*constant*cosdecay")
+  hparams.learning_rate_schedule = ("linear_warmup*constant*cosdecay")
   hparams.learning_rate_constant = 5e-5
   hparams.learning_rate_warmup_steps = 100
   # Set train steps to learning_rate_decay_steps or less
@@ -1628,8 +1885,7 @@ def transformer_tall_train_tied():
   hparams.multiproblem_max_input_length = 750
   hparams.multiproblem_max_target_length = 100
   hparams.multiproblem_schedule_max_examples = 0
-  hparams.learning_rate_schedule = (
-      "linear_warmup*constant*cosdecay")
+  hparams.learning_rate_schedule = ("linear_warmup*constant*cosdecay")
   hparams.learning_rate_constant = 2e-4
   hparams.learning_rate_warmup_steps = 8000
   # Set train steps to learning_rate_decay_steps or less
@@ -1648,8 +1904,7 @@ def transformer_tall_finetune_uniencdec():
   hparams.max_input_seq_length = 750
   hparams.max_target_seq_length = 100
   hparams.optimizer = "TrueAdam"
-  hparams.learning_rate_schedule = (
-      "linear_warmup*constant*cosdecay")
+  hparams.learning_rate_schedule = ("linear_warmup*constant*cosdecay")
   hparams.learning_rate_decay_steps = 80000
   hparams.learning_rate_constant = 5e-5
   hparams.learning_rate_warmup_steps = 100
@@ -1664,8 +1919,7 @@ def transformer_tall_train_uniencdec():
   hparams.max_input_seq_length = 750
   hparams.max_target_seq_length = 100
   hparams.optimizer = "TrueAdam"
-  hparams.learning_rate_schedule = (
-      "linear_warmup*constant*cosdecay")
+  hparams.learning_rate_schedule = ("linear_warmup*constant*cosdecay")
   hparams.learning_rate_decay_steps = 150000
   hparams.learning_rate_constant = 2e-4
   hparams.unidirectional_encoder = True
@@ -1677,8 +1931,7 @@ def transformer_tall_finetune_textclass():
   """Hparams for transformer on LM for finetuning on text class problems."""
   hparams = transformer_tall()
   hparams.learning_rate_constant = 6.25e-5
-  hparams.learning_rate_schedule = (
-      "linear_warmup*constant*linear_decay")
+  hparams.learning_rate_schedule = ("linear_warmup*constant*linear_decay")
   hparams.multiproblem_schedule_max_examples = 0
   hparams.multiproblem_target_eval_only = True
   hparams.learning_rate_warmup_steps = 50
@@ -1694,8 +1947,7 @@ def transformer_tall_pretrain_lm():
   """Hparams for transformer on LM pretraining (with 64k vocab)."""
   hparams = transformer_tall()
   hparams.learning_rate_constant = 2e-4
-  hparams.learning_rate_schedule = (
-      "linear_warmup*constant*cosdecay")
+  hparams.learning_rate_schedule = ("linear_warmup*constant*cosdecay")
   hparams.optimizer = "adam_w"
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.999

From 44323a19c0233babbd7b25a0e9b0adea7917219d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 1 Feb 2019 14:53:05 -0800
Subject: [PATCH 1623/2720] Required changes needed to run model free training
 with EnvProblem.

PiperOrigin-RevId: 232042371
---
 tensor2tensor/envs/env_problem.py          | 9 +++++++++
 tensor2tensor/rl/envs/py_func_batch_env.py | 9 ++++++++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 015341546..c57bfe287 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -181,9 +181,17 @@ def __init__(self,
     # and also the ones that are completed, i.e. done.
     self._trajectories = None
 
+    self._batch_size = None
+
     if batch_size is not None:
       self.initialize(batch_size=batch_size)
 
+  @property
+  def batch_size(self):
+    # TODO(afrozm): I've added this here since it is being used in a lot of
+    # places in ppo_learner.py -- re-evaluate if needed.
+    return self._batch_size
+
   @property
   def base_env_name(self):
     return self._base_env_name
@@ -253,6 +261,7 @@ def initialize_environments(self, batch_size=1):
     """
 
     assert batch_size >= 1
+    self._batch_size = batch_size
 
     max_steps = self._base_env_kwargs.get("rl_env_max_episode_steps", -1)
     maxskip_env = self._base_env_kwargs.get("maxskip_env", False)
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index 5f11fe95d..b83cf5377 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -91,7 +91,14 @@ def simulate(self, action):
       if action.dtype in (tf.float16, tf.float32, tf.float64):
         action = tf.check_numerics(action, "action")
       def step(action):
-        (observ, reward, done) = self._batch_env.step(action)
+        step_response = self._batch_env.step(action)
+        # Current env doesn't return `info`, but EnvProblem does.
+        # TODO(afrozm): The proper way to do this is to make T2TGymEnv return
+        # an empty info return value.
+        if len(step_response) == 3:
+          (observ, reward, done) = step_response
+        else:
+          (observ, reward, done, _) = step_response
         return (observ, reward.astype(np.float32), done)
       observ, reward, done = tf.py_func(
           step, [action],

From 8e888982218d708150f0c084b4c91e726d9d2ef7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 1 Feb 2019 23:59:05 -0800
Subject: [PATCH 1624/2720] Fix build breakage by cl/231978646.

PiperOrigin-RevId: 232093621
---
 tensor2tensor/data_generators/problem.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index e3c8c6c1e..7f0ac8104 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -804,7 +804,9 @@ def _dataset_partition(self, mode, config, params):
         phift == tpu_config.InputPipelineConfig.BROADCAST):
       return 0, 1
     if phift:
-      num_partitions = max(params["context"].num_hosts, 1)
+      num_hosts = (params["context"].num_hosts if "context" in params
+                   else config.tpu_config.num_shards // 8)
+      num_partitions = max(num_hosts, 1)
     else:
       num_partitions = config.tpu_config.num_shards
     partition_id = getattr(self, "_next_partition_id", 0)

From 21723411f0f0abda45e2b4cd49ac282a16797d71 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 4 Feb 2019 11:06:35 -0800
Subject: [PATCH 1625/2720] Remove contrib run_test_in_graph_and_eager and rely
 on a fork of TFDS' run_in_graph_and_eager_modes.

Also add a run_in_graph_mode_only decorator, since we have to switch on eager in tests for the first decorator to work.

PiperOrigin-RevId: 232329353
---
 tensor2tensor/data_generators/problem_test.py |  14 +-
 tensor2tensor/utils/t2t_model_test.py         |   4 +-
 tensor2tensor/utils/test_utils.py             | 124 ++++++++++++++++++
 tensor2tensor/utils/test_utils_test.py        |  76 +++++++++++
 4 files changed, 212 insertions(+), 6 deletions(-)
 create mode 100644 tensor2tensor/utils/test_utils.py
 create mode 100644 tensor2tensor/utils/test_utils_test.py

diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index c7b5781be..4cab788e6 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -26,8 +26,10 @@
 from tensor2tensor.data_generators import problem as problem_module
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.layers import modalities
+from tensor2tensor.utils import test_utils
 
 import tensorflow as tf
+tf.compat.v1.enable_eager_execution()
 
 
 def assert_tensors_equal(sess, t1, t2, n):
@@ -52,6 +54,7 @@ class ProblemTest(parameterized.TestCase, tf.test.TestCase):
   def setUpClass(cls):
     algorithmic.TinyAlgo.setup_for_test()
 
+  @test_utils.run_in_graph_mode_only()
   def testNoShuffleDeterministic(self):
     problem = algorithmic.TinyAlgo()
     dataset = problem.dataset(mode=tf.estimator.ModeKeys.TRAIN,
@@ -64,6 +67,7 @@ def testNoShuffleDeterministic(self):
     with tf.Session() as sess:
       self.assertTrue(assert_tensors_equal(sess, tensor1, tensor2, 20))
 
+  @test_utils.run_in_graph_mode_only()
   def testNoShufflePreprocess(self):
 
     problem = algorithmic.TinyAlgo()
@@ -80,7 +84,7 @@ def testNoShufflePreprocess(self):
     with tf.Session() as sess:
       self.assertTrue(assert_tensors_equal(sess, tensor1, tensor2, 20))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testProblemHparamsModality(self):
     problem = problem_hparams.TestProblem(input_vocab_size=2,
                                           target_vocab_size=3)
@@ -90,7 +94,7 @@ def testProblemHparamsModality(self):
     self.assertIsInstance(p_hparams.modality["targets"],
                           modalities.SymbolModality)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testProblemHparamsModalityObj(self):
     class ModalityObjProblem(problem_module.Problem):
 
@@ -108,7 +112,7 @@ def hparams(self, defaults, model_hparams):
     self.assertIsInstance(p_hparams.modality["targets"],
                           modalities.SymbolModality)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testProblemHparamsInputOnlyModality(self):
     class InputOnlyProblem(problem_module.Problem):
 
@@ -123,7 +127,7 @@ def hparams(self, defaults, model_hparams):
                           modalities.SymbolModality)
     self.assertLen(p_hparams.modality, 1)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testProblemHparamsTargetOnlyModality(self):
     class TargetOnlyProblem(problem_module.Problem):
 
@@ -138,7 +142,7 @@ def hparams(self, defaults, model_hparams):
                           modalities.SymbolModality)
     self.assertLen(p_hparams.modality, 1)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testDataFilenames(self):
     problem = algorithmic.TinyAlgo()
 
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index 7e8cef93a..76847439b 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -21,14 +21,16 @@
 
 from tensor2tensor.utils import modality
 from tensor2tensor.utils import t2t_model
+from tensor2tensor.utils import test_utils
 from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
+tf.compat.v1.enable_eager_execution()
 
 
 class T2TModelTest(tf.test.TestCase):
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testSummarizeLosses(self):
     with tf.Graph().as_default():
       model = t2t_model.T2TModel(HParams())
diff --git a/tensor2tensor/utils/test_utils.py b/tensor2tensor/utils/test_utils.py
new file mode 100644
index 000000000..0d48dd3bf
--- /dev/null
+++ b/tensor2tensor/utils/test_utils.py
@@ -0,0 +1,124 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Test utilities."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+
+def run_in_graph_and_eager_modes(func=None,
+                                 config=None,
+                                 use_gpu=True):
+  """Execute the decorated test with and without enabling eager execution.
+
+  This function returns a decorator intended to be applied to test methods in
+  a `tf.test.TestCase` class. Doing so will cause the contents of the test
+  method to be executed twice - once in graph mode, and once with eager
+  execution enabled. This allows unittests to confirm the equivalence between
+  eager and graph execution.
+
+  NOTE: This decorator can only be used when executing eagerly in the
+  outer scope.
+
+  For example, consider the following unittest:
+
+  ```python
+  tf.compat.v1.enable_eager_execution()
+
+  class SomeTest(tf.test.TestCase):
+
+    @test_utils.run_in_graph_and_eager_modes
+    def test_foo(self):
+      x = tf.constant([1, 2])
+      y = tf.constant([3, 4])
+      z = tf.add(x, y)
+      self.assertAllEqual([4, 6], self.evaluate(z))
+
+  if __name__ == "__main__":
+    tf.test.main()
+  ```
+
+  This test validates that `tf.add()` has the same behavior when computed with
+  eager execution enabled as it does when constructing a TensorFlow graph and
+  executing the `z` tensor with a session.
+
+  Args:
+    func: function to be annotated. If `func` is None, this method returns a
+      decorator the can be applied to a function. If `func` is not None this
+      returns the decorator applied to `func`.
+    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.
+
+  Returns:
+    Returns a decorator that will run the decorated test method twice:
+    once by constructing and executing a graph in a session and once with
+    eager execution enabled.
+  """
+
+  def decorator(f):
+    """Decorator for a method."""
+    def decorated(self, *args, **kwargs):
+      """Run the decorated test method."""
+      if not tf.executing_eagerly():
+        raise ValueError("Must be executing eagerly when using the "
+                         "run_in_graph_and_eager_modes decorator.")
+
+      # Run eager block
+      f(self, *args, **kwargs)
+      self.tearDown()
+
+      # Run in graph mode block
+      with tf.Graph().as_default():
+        self.setUp()
+        with self.test_session(use_gpu=use_gpu, config=config):
+          f(self, *args, **kwargs)
+
+    return decorated
+
+  if func is not None:
+    return decorator(func)
+
+  return decorator
+
+
+def run_in_graph_mode_only(func=None, config=None, use_gpu=True):
+  """Runs a test in graph mode only, when eager is enabled by default."""
+  def decorator(f):
+    """Decorator for a method."""
+    def decorated(self, *args, **kwargs):
+      """Run the decorated test method."""
+      self.tearDown()
+      # Run in graph mode block
+      with tf.Graph().as_default():
+        self.setUp()
+        with self.test_session(use_gpu=use_gpu, config=config):
+          f(self, *args, **kwargs)
+
+    return decorated
+
+  if func is not None:
+    return decorator(func)
+
+  return decorator
+
+
+def test_main():
+  tf.compat.v1.enable_eager_execution()
+  tf.test.main()
diff --git a/tensor2tensor/utils/test_utils_test.py b/tensor2tensor/utils/test_utils_test.py
new file mode 100644
index 000000000..d7bb25ab6
--- /dev/null
+++ b/tensor2tensor/utils/test_utils_test.py
@@ -0,0 +1,76 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.utils.test_utils."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.utils import test_utils
+
+import tensorflow as tf
+
+tf.compat.v1.enable_eager_execution()
+
+
+class RunInGraphAndEagerTest(tf.test.TestCase):
+
+  def test_run_in_graph_and_eager_modes(self):
+    l = []
+    def inc(self, with_brackets):
+      del self  # self argument is required by run_in_graph_and_eager_modes.
+      mode = "eager" if tf.executing_eagerly() else "graph"
+      with_brackets = "with_brackets" if with_brackets else "without_brackets"
+      l.append((with_brackets, mode))
+
+    f = test_utils.run_in_graph_and_eager_modes(inc)
+    f(self, with_brackets=False)
+    f = test_utils.run_in_graph_and_eager_modes()(inc)
+    f(self, with_brackets=True)
+
+    self.assertEqual(len(l), 4)
+    self.assertEqual(set(l), {
+        ("with_brackets", "graph"),
+        ("with_brackets", "eager"),
+        ("without_brackets", "graph"),
+        ("without_brackets", "eager"),
+    })
+
+  def test_run_in_graph_and_eager_modes_setup_in_same_mode(self):
+    modes = []
+    mode_name = lambda: "eager" if tf.executing_eagerly() else "graph"
+
+    class ExampleTest(tf.test.TestCase):
+
+      def runTest(self):
+        pass
+
+      def setUp(self):
+        modes.append("setup_" + mode_name())
+
+      @test_utils.run_in_graph_and_eager_modes
+      def testBody(self):
+        modes.append("run_" + mode_name())
+
+    e = ExampleTest()
+    e.setUp()
+    e.testBody()
+
+    self.assertEqual(modes[0:2], ["setup_eager", "run_eager"])
+    self.assertEqual(modes[2:], ["setup_graph", "run_graph"])
+
+if __name__ == "__main__":
+  tf.test.main()

From b08dfeb4477095093fe2449b84ffef17d1ab1afb Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 4 Feb 2019 11:28:10 -0800
Subject: [PATCH 1626/2720] Move the remaining uses of
 tf.contrib.eager.run_test_in_graph_and_eager_modes decorator to the added
 test_utils.run_in_graph_and_eager_modes and
 test_utils.run_in_graph_mode_only.

A lot of tests that are annotated with run_in_graph_mode_only seem to be easily
convertible to be run in both modes, so that should be done as well.

PiperOrigin-RevId: 232333669
---
 tensor2tensor/layers/bayes_test.py            | 24 ++---
 tensor2tensor/layers/common_attention_test.py | 49 +++++-----
 tensor2tensor/layers/common_layers_test.py    | 91 +++++++++++--------
 tensor2tensor/layers/common_video_test.py     | 18 ++--
 tensor2tensor/layers/discretization_test.py   | 24 +++--
 tensor2tensor/layers/latent_layers_test.py    |  4 +-
 tensor2tensor/layers/modalities_test.py       |  7 +-
 tensor2tensor/layers/ngram_test.py            |  7 +-
 .../layers/reversible_layers_test.py          | 12 ++-
 9 files changed, 137 insertions(+), 99 deletions(-)

diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index e9ec46370..15eb0564b 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -23,13 +23,15 @@
 import numpy as np
 
 from tensor2tensor.layers import bayes
+from tensor2tensor.utils import test_utils
 
 import tensorflow as tf
+tf.compat.v1.enable_eager_execution()
 
 
 class BayesTest(parameterized.TestCase, tf.test.TestCase):
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes
+  @test_utils.run_in_graph_and_eager_modes
   def testTrainableNormalStddevConstraint(self):
     layer = bayes.DenseReparameterization(
         100, kernel_initializer=bayes.TrainableNormal())
@@ -48,7 +50,7 @@ def testTrainableNormalStddevConstraint(self):
       {"testcase_name": "_bias_uncertainty", "kernel_initializer": "zeros",
        "bias_initializer": None, "all_close": False},
   )
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes
+  @test_utils.run_in_graph_and_eager_modes
   def testDenseReparameterizationKernel(
       self, kernel_initializer, bias_initializer, all_close):
     inputs = tf.to_float(np.random.rand(5, 3, 12))
@@ -67,7 +69,7 @@ def testDenseReparameterizationKernel(
       self.assertNotAllClose(res1, res2)
     layer.get_config()
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testDenseReparameterizationKL(self):
     inputs = tf.to_float(np.random.rand(5, 12))
     layer = bayes.DenseReparameterization(10)
@@ -104,7 +106,7 @@ def testDenseReparameterizationKL(self):
     for grad in grads:
       self.assertIsNotNone(grad)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testDenseReparameterizationModel(self):
     inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
     model = tf.keras.Sequential([
@@ -121,7 +123,7 @@ def testDenseReparameterizationModel(self):
     self.assertEqual(res.shape, (3, 2))
     self.assertLen(model.losses, 1)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testGaussianProcessPosterior(self):
     train_batch_size = 3
     test_batch_size = 2
@@ -142,7 +144,7 @@ def testGaussianProcessPosterior(self):
     self.assertGreaterEqual(test_nats_val, 0.)
     self.assertEqual(outputs_val.shape, (test_batch_size, output_dim))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testGaussianProcessPrior(self):
     batch_size = 3
     input_dim = 4
@@ -175,7 +177,7 @@ def testGaussianProcessPrior(self):
        "recurrent_initializer": "orthogonal", "bias_initializer": None,
        "all_close": False},
   )
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes
+  @test_utils.run_in_graph_and_eager_modes
   def testLSTMCellReparameterization(
       self, kernel_initializer, recurrent_initializer, bias_initializer,
       all_close):
@@ -203,7 +205,7 @@ def testLSTMCellReparameterization(
       self.assertNotAllClose(res1, res3)
     cell.get_config()
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testLSTMCellReparameterizationKL(self):
     inputs = tf.to_float(np.random.rand(5, 1, 12))
     cell = bayes.LSTMCellReparameterization(10)
@@ -249,7 +251,7 @@ def testLSTMCellReparameterizationKL(self):
     for grad in grads:
       self.assertIsNotNone(grad)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testLSTMCellReparameterizationModel(self):
     batch_size, timesteps, dim = 5, 3, 12
     hidden_size = 10
@@ -279,7 +281,7 @@ def testLSTMCellReparameterizationModel(self):
     self.assertAllClose(res2, res3)
     self.assertLen(model.losses, 2)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testBayesianLinearModel(self):
     """Tests that model makes reasonable predictions."""
     np.random.seed(42)
@@ -310,7 +312,7 @@ def testBayesianLinearModel(self):
     self.assertAllClose(test_predictions_val, test_labels_val, atol=0.1)
     self.assertAllLessEqual(test_predictions_variance_val, noise_variance)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testMixtureLogistic(self):
     batch_size = 3
     features = tf.to_float(np.random.rand(batch_size, 4))
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 701e2bb39..ce9f4417d 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -24,13 +24,15 @@
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import test_utils
 
 import tensorflow as tf
+tf.compat.v1.enable_eager_execution()
 
 
 class CommonAttentionTest(parameterized.TestCase, tf.test.TestCase):
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testAddPositionalEmbedding(self):
     x = np.random.rand(5, 3, 12)
     y = common_attention.add_positional_embedding(
@@ -46,7 +48,7 @@ def testAddPositionalEmbedding(self):
       {"input_shape": (5, 5, 5, 12)},
       {"input_shape": (5, 3, 3, 3, 12)},
   )
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testAddPositionalEmbeddingNd(self, input_shape):
     x = np.random.rand(*input_shape)
     y = common_attention.add_positional_embedding_nd(
@@ -57,7 +59,7 @@ def testAddPositionalEmbeddingNd(self, input_shape):
     res = self.evaluate(y)
     self.assertEqual(res.shape, input_shape)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testDotProductAttention(self):
     x = np.random.rand(5, 7, 12, 32)
     y = np.random.rand(5, 7, 12, 32)
@@ -191,6 +193,7 @@ def testLocalUnmaskedAttention2D(self, batch, heads, length,
 
     self.assertEqual(res.shape, (batch, heads, length, length, depth_v))
 
+  @test_utils.run_in_graph_mode_only()
   def testMultiheadSelfAttentionMemoryEfficient(self):
     num_heads = 4
     io_size = 16
@@ -235,7 +238,7 @@ def testMultiheadSelfAttentionMemoryEfficient(self):
     self.assertAllClose(dnorm_bias, dnorm_bias_f)
     self.assertAllClose(dx, dx_f)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def test2dGatherAndScatterInvertibility(self):
     """2d gather and scatter invertibility test."""
     batch_size = 2
@@ -254,7 +257,7 @@ def test2dGatherAndScatterInvertibility(self):
     res = self.evaluate(scattered_x)
     self.assertAllClose(x, res)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def test2dBlockRasterScanMask(self):
     """Testing the 2d block raster scan mask."""
     query_shape = (2, 3)
@@ -277,7 +280,7 @@ def test2dBlockRasterScanMask(self):
           1.0, 0.0, 0.0, 0.0, 0.0, 1.0]])
     self.assertAllClose(correct_mask, res)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def test2dGather(self):
     """Testing 2d index gather and block gather functions."""
     batch_size = 2
@@ -316,7 +319,7 @@ def test2dGather(self):
     self.assertAllEqual(correct_indices, x_indices)
     self.assertAllClose(correct_gathered_x, gathered_x)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testGetMemoryRegion(self):
     """Testing the function that gathers the flanged memory region."""
     np.set_printoptions(threshold=np.inf)
@@ -393,7 +396,7 @@ def testGetMemoryRegion(self):
     self.assertAllClose(correct_x_flange, x_flange)
     self.assertAllClose(correct_x_center, x_center)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testGetShiftedCenterBlocks(self):
     """Testing the function that gathers the flanged memory region."""
     np.set_printoptions(threshold=np.inf)
@@ -458,7 +461,7 @@ def testGetShiftedCenterBlocks(self):
     x_indices, gathered_x = self.evaluate([x_indices, gathered_x])
     self.assertAllClose(correct_gathered_x, gathered_x)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testDotProductAttentionRelative(self):
     x = np.random.rand(5, 7, 12, 32)
     y = np.random.rand(5, 7, 12, 32)
@@ -472,7 +475,7 @@ def testDotProductAttentionRelative(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testRelativeAttentionV2(self):
     # (batch, heads, length, depth)
     x = np.random.rand(5, 4, 16, 7)
@@ -489,7 +492,7 @@ def testRelativeAttentionV2(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 4, 16, 7))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testRelativeAttentionV2SharedRel(self):
     # (batch, heads, length, depth)
     x = np.random.rand(5, 4, 16, 7)
@@ -506,7 +509,7 @@ def testRelativeAttentionV2SharedRel(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 4, 16, 7))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testRelativeAttentionV2MaxRelativeLargerThanLength(self):
     # (batch, heads, length, depth)
     x = np.random.rand(5, 4, 3, 7)
@@ -523,7 +526,7 @@ def testRelativeAttentionV2MaxRelativeLargerThanLength(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 4, 3, 7))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testDotProductUnMaskedAttentionRelativeV2(self):
     x = np.random.rand(5, 7, 12, 32)
     y = np.random.rand(5, 7, 12, 32)
@@ -605,7 +608,7 @@ def python_relative_att(self, q, k, v, batch, num_heads, height, width,
                             (batch, num_heads, height, width, depth))
     return att_output
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testDotProductUnMaskedAttentionRelative2d(self):
     batch = 1
     height = 3
@@ -643,7 +646,7 @@ def testDotProductUnMaskedAttentionRelative2d(self):
       (1, 10, 12, 2, 2, 8),
       (4, 10, 12, 2, 12, 10),
   )
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testDotProductUnMaskedAttentionRelative2dSharedOneRow(
       self, batch, height, width, num_heads, max_relative_position, depth):
     heads_share_relative_embedding = True
@@ -670,7 +673,7 @@ def testDotProductUnMaskedAttentionRelative2dSharedOneRow(
                      (batch, num_heads, height, width, depth))
     self.assertAllClose(res, att_output)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testRelativeAttentionV2Unmasked(self):
     # (batch, heads, length, depth)
     x = np.random.rand(5, 4, 16, 7)
@@ -687,7 +690,7 @@ def testRelativeAttentionV2Unmasked(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 4, 16, 7))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testRelativeAttentionV2UnmaskedSharedRel(self):
     # (batch, heads, length, depth)
     x = np.random.rand(5, 4, 16, 7)
@@ -704,7 +707,7 @@ def testRelativeAttentionV2UnmaskedSharedRel(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 4, 16, 7))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testRelativeAttentionV2UnmaskedRelativeLargerThanLength(self):
     # (batch, heads, length, depth)
     x = np.random.rand(5, 4, 3, 7)
@@ -721,7 +724,7 @@ def testRelativeAttentionV2UnmaskedRelativeLargerThanLength(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 4, 3, 7))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testMaskedRelativeLocalAttentionV2(self):
     # (batch, heads, length, depth)
     x = np.random.rand(5, 4, 16, 7)
@@ -739,7 +742,7 @@ def testMaskedRelativeLocalAttentionV2(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 4, 16, 7))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testMaskedRelativeLocalAttentionV2AddRelativeValues(self):
     # (batch, heads, length, depth)
     x = np.random.rand(5, 4, 16, 7)
@@ -757,7 +760,7 @@ def testMaskedRelativeLocalAttentionV2AddRelativeValues(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 4, 16, 7))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testMaskedRelativeLocalAttentionV2SeqShorterThanBlockLength(self):
     # (batch, heads, length, depth)
     x = np.random.rand(5, 7, 2, 7)
@@ -774,7 +777,7 @@ def testMaskedRelativeLocalAttentionV2SeqShorterThanBlockLength(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 2, 7))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testMaskedRelativeLocalAttentionV2SeqShorterThanTwiceBlockLength(self):
     # (batch, heads, length, depth)
     x = np.random.rand(5, 7, 5, 7)
@@ -814,7 +817,7 @@ def testBiasBatchCoordinates(self):
     bias = common_attention.attention_bias_coordinates(q, k)
     self.assertAllClose(self.evaluate(bias), ground_truth)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testBiasFuture(self):
     """Testing the sequence order mask."""
     q = tf.constant([0, 1, 2, 3, 0, 1, 2, 0, 1], dtype=tf.int32)
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index 60d5961ad..50fc19950 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -23,13 +23,16 @@
 import numpy as np
 
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import test_utils
 
 import tensorflow as tf
 
+tf.compat.v1.enable_eager_execution()
+
 
 class CommonLayersTest(parameterized.TestCase, tf.test.TestCase):
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testIndexLastDimWithIndices(self):
     x = np.array([[2., 3., 4., 5.],
                   [6., 7., 8., 9.]])
@@ -39,14 +42,14 @@ def testIndexLastDimWithIndices(self):
     expected = np.array([4., 6.])
     self.assertAllEqual(expected, self.evaluate(x_idx))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testSaturatingSigmoid(self):
     x = np.array([-120.0, -100.0, 0.0, 100.0, 120.0], dtype=np.float32)
     y = common_layers.saturating_sigmoid(tf.constant(x))
     res = self.evaluate(y)
     self.assertAllClose(res, [0.0, 0.0, 0.5, 1.0, 1.0])
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testFlatten4D3D(self):
     x = np.random.randint(1, high=9, size=(3, 5, 2))
     y = common_layers.flatten4d3d(common_layers.embedding(x, 10, 7))
@@ -54,7 +57,7 @@ def testFlatten4D3D(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (3, 5 * 2, 7))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testEmbedding(self):
     x = np.random.randint(1, high=9, size=(3, 5))
     y = common_layers.embedding(x, 10, 16)
@@ -62,6 +65,7 @@ def testEmbedding(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (3, 5, 16))
 
+  @test_utils.run_in_graph_mode_only()
   def testShakeShake(self):
     x = np.random.rand(5, 7)
     with self.test_session() as session:
@@ -70,7 +74,7 @@ def testShakeShake(self):
       inp, res = session.run([x, y])
     self.assertAllClose(res, inp)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testConv(self):
     x = np.random.rand(5, 7, 1, 11)
     y = common_layers.conv(tf.constant(x, dtype=tf.float32), 13, (3, 1))
@@ -78,7 +82,7 @@ def testConv(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 5, 1, 13))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testConv1d(self):
     x = np.random.rand(5, 7, 11)
     y = common_layers.conv1d(tf.constant(x, dtype=tf.float32), 13, 1)
@@ -86,7 +90,7 @@ def testConv1d(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 13))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testSeparableConv(self):
     x = np.random.rand(5, 7, 1, 11)
     y = common_layers.separable_conv(
@@ -95,7 +99,7 @@ def testSeparableConv(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 5, 1, 13))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testSubSeparableConv(self):
     for sep in [0, 1, 2, 4]:
       x = np.random.rand(5, 7, 1, 12)
@@ -106,7 +110,7 @@ def testSubSeparableConv(self):
       res = self.evaluate(y)
       self.assertEqual(res.shape, (5, 5, 1, 16))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testConvBlock(self):
     x = np.random.rand(5, 7, 1, 11)
     y = common_layers.conv_block(
@@ -118,7 +122,7 @@ def testConvBlock(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 1, 13))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testSeparableConvBlock(self):
     x = np.random.rand(5, 7, 1, 11)
     y = common_layers.separable_conv_block(
@@ -129,7 +133,7 @@ def testSeparableConvBlock(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 1, 13))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testSubSeparableConvBlock(self):
     for sep in [0, 1, 2, 4]:
       x = np.random.rand(5, 7, 1, 12)
@@ -143,7 +147,7 @@ def testSubSeparableConvBlock(self):
       res = self.evaluate(y)
       self.assertEqual(res.shape, (5, 7, 1, 16))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testPool(self):
     x = np.random.rand(5, 8, 1, 11)
     y = common_layers.pool(
@@ -152,7 +156,7 @@ def testPool(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 8, 1, 11))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testConvBlockDownsample(self):
     x = np.random.rand(5, 7, 1, 11)
     y = common_layers.conv_block_downsample(
@@ -161,7 +165,7 @@ def testConvBlockDownsample(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 4, 1, 27))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testGetTimingSignal(self):
     length = 7
     num_timescales = 10
@@ -169,7 +173,7 @@ def testGetTimingSignal(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (length, 2 * num_timescales))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testAddTimingSignal(self):
     batch = 5
     length = 7
@@ -180,7 +184,7 @@ def testAddTimingSignal(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (batch, length, height, depth))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testConvGRU(self):
     x = np.random.rand(5, 7, 3, 11)
     y = common_layers.conv_gru(tf.constant(x, dtype=tf.float32), (1, 3), 11)
@@ -192,6 +196,7 @@ def testConvGRU(self):
     self.assertEqual(res1.shape, (5, 7, 3, 11))
     self.assertEqual(res2.shape, (5, 7, 3, 11))
 
+  @test_utils.run_in_graph_mode_only
   def testSRU(self):
     x = np.random.rand(5, 7, 3, 11)
     with self.test_session() as session:
@@ -200,7 +205,7 @@ def testSRU(self):
       res = session.run(y)
     self.assertEqual(res.shape, (5, 7, 3, 11))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testLayerNorm(self):
     x = np.random.rand(5, 7, 11)
     y = common_layers.layer_norm(tf.constant(x, dtype=tf.float32), 11)
@@ -208,7 +213,7 @@ def testLayerNorm(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 11))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testGroupNorm(self):
     x = np.random.rand(5, 7, 3, 16)
     y = common_layers.group_norm(tf.constant(x, dtype=tf.float32))
@@ -216,7 +221,7 @@ def testGroupNorm(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 3, 16))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testConvLSTM(self):
     x = np.random.rand(5, 7, 11, 13)
     y = common_layers.conv_lstm(tf.constant(x, dtype=tf.float32), (1, 3), 13)
@@ -224,7 +229,7 @@ def testConvLSTM(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 11, 13))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testPadToSameLength(self):
     x1 = np.random.rand(5, 7, 11)
     x2 = np.random.rand(5, 9, 11)
@@ -241,7 +246,7 @@ def testPadToSameLength(self):
     self.assertEqual(res1a.shape, (5, 12, 11))
     self.assertEqual(res2a.shape, (5, 12, 11))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testShiftLeft(self):
     x1 = np.zeros((5, 7, 1, 11))
     x1[:, 0, :] = np.ones_like(x1[:, 0, :])
@@ -251,7 +256,7 @@ def testShiftLeft(self):
     actual = self.evaluate(a)
     self.assertAllEqual(actual, expected)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testConvStride2MultiStep(self):
     x1 = np.random.rand(5, 32, 16, 11)
     a = common_layers.conv_stride2_multistep(
@@ -260,7 +265,7 @@ def testConvStride2MultiStep(self):
     actual = self.evaluate(a[0])
     self.assertEqual(actual.shape, (5, 2, 1, 16))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testDeconvStride2MultiStep(self):
     x1 = np.random.rand(5, 2, 1, 11)
     a = common_layers.deconv_stride2_multistep(
@@ -269,7 +274,7 @@ def testDeconvStride2MultiStep(self):
     actual = self.evaluate(a)
     self.assertEqual(actual.shape, (5, 32, 1, 16))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testApplyNormLayer(self):
     x1 = np.random.rand(5, 2, 1, 11)
     x2 = common_layers.apply_norm(
@@ -278,7 +283,7 @@ def testApplyNormLayer(self):
     actual = self.evaluate(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testApplyNormNoam(self):
     x1 = np.random.rand(5, 2, 1, 11)
     x2 = common_layers.apply_norm(
@@ -287,7 +292,7 @@ def testApplyNormNoam(self):
     actual = self.evaluate(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testApplyNormBatch(self):
     x1 = np.random.rand(5, 2, 1, 11)
     x2 = common_layers.apply_norm(
@@ -296,7 +301,7 @@ def testApplyNormBatch(self):
     actual = self.evaluate(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testApplyNormNone(self):
     x1 = np.random.rand(5, 2, 1, 11)
     x2 = common_layers.apply_norm(
@@ -355,21 +360,21 @@ def testRavanbakhshSetLayer(self):
     actual = self.evaluate(layer)
     self.assertEqual(actual.shape, (5, 4, 32))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testBReLU(self):
     x = np.random.rand(5, 2, 1, 12)
     y = common_layers.brelu(tf.constant(x, dtype=tf.float32))
     actual = self.evaluate(y)
     self.assertEqual(actual.shape, (5, 2, 1, 12))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testBELU(self):
     x = np.random.rand(5, 2, 1, 12)
     y = common_layers.belu(tf.constant(x, dtype=tf.float32))
     actual = self.evaluate(y)
     self.assertEqual(actual.shape, (5, 2, 1, 12))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testNAC(self):
     x = np.random.rand(5, 2, 1, 12)
     y = common_layers.nac(tf.constant(x, dtype=tf.float32), 14)
@@ -377,7 +382,7 @@ def testNAC(self):
     actual = self.evaluate(y)
     self.assertEqual(actual.shape, (5, 2, 1, 14))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testNALU(self):
     x = np.random.rand(5, 2, 1, 12)
     y = common_layers.nalu(tf.constant(x, dtype=tf.float32), 14)
@@ -385,7 +390,7 @@ def testNALU(self):
     actual = self.evaluate(y)
     self.assertEqual(actual.shape, (5, 2, 1, 14))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testNALUzeros(self):
     x = np.random.rand(5, 2, 1, 12)
     y = common_layers.nalu(tf.zeros_like(x, dtype=tf.float32), 14)
@@ -394,6 +399,7 @@ def testNALUzeros(self):
     self.assertTrue(np.all(np.isfinite(actual)))
     self.assertEqual(actual.shape, (5, 2, 1, 14))
 
+  @test_utils.run_in_graph_mode_only
   def testPaddingCrossEntropyFactored(self):
     vocab_size = 19
     rows = 5
@@ -427,6 +433,7 @@ def testPaddingCrossEntropyFactored(self):
     self.assertAllClose(num, num_f)
     self.assertAllClose(den, den_f)
 
+  @test_utils.run_in_graph_mode_only
   def testPaddingCrossEntropyFactoredGrad(self):
     vocab_size = 19
     rows = 5
@@ -489,7 +496,7 @@ def testDmlLoss(self, batch, height, width, num_mixtures, reduce_sum):
         [actual_loss, expected_loss])
     self.assertAllClose(actual_loss_val, expected_loss_val)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testWeightsMultiProblemAll(self):
     labels = tf.constant(np.array([[12, 15, 1, 20, 100],
                                    [67, 1, 34, 45, 124],
@@ -506,7 +513,7 @@ def testWeightsMultiProblemAll(self):
     actual_mask_eval = self.evaluate(actual_mask)
     self.assertAllClose(expected_mask, actual_mask_eval)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testWeightsMultiProblem(self):
     labels = tf.constant(np.array([[12, 15, 1, 20, 100],
                                    [67, 1, 34, 45, 124],
@@ -523,7 +530,7 @@ def testWeightsMultiProblem(self):
     actual_mask_eval = self.evaluate(actual_mask)
     self.assertAllClose(expected_mask, actual_mask_eval)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testDiscretizedMixLogisticLoss(self):
     batch = 2
     height = 4
@@ -560,7 +567,7 @@ def testDiscretizedMixLogisticLoss(self):
         [actual_loss, expected_loss])
     self.assertAllClose(actual_loss_val, expected_loss_val, rtol=1e-5)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testSampleFromDiscretizedMixLogistic(self):
     batch = 2
     height = 4
@@ -588,7 +595,7 @@ def testSampleFromDiscretizedMixLogistic(self):
     # implementation clips log-scales so they always contribute to sampling.
     self.assertAllClose(actual_sample_val, expected_sample_val, atol=1e-2)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testFactoredTensorImplicitConversion(self):
     a = np.random.rand(3, 4, 5)
     b = np.random.rand(6, 5)
@@ -600,6 +607,7 @@ def testFactoredTensorImplicitConversion(self):
     out = self.evaluate(d)
     self.assertEqual(out.shape, (3, 4, 6))
 
+  @test_utils.run_in_graph_mode_only()
   def testConvHiddenReluMemoryEfficient(self):
     batch = 3
     length = 23
@@ -637,7 +645,7 @@ def testConvHiddenReluMemoryEfficient(self):
     self.assertAllClose(dnorm_bias, dnorm_bias_f)
     self.assertAllClose(dx, dx_f)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testCycleGANUpsampleNnUpsampleConv(self):
     batch = 8
     height = 32
@@ -657,7 +665,7 @@ def testCycleGANUpsampleNnUpsampleConv(self):
         [batch, height * stride[0], width * stride[1], output_filters],
         self.evaluate(upsampled_output_shape))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testCycleGANUpsampleBilinearUpsampleConv(self):
     batch = 8
     height = 32
@@ -677,7 +685,7 @@ def testCycleGANUpsampleBilinearUpsampleConv(self):
         [batch, height * stride[0], width * stride[1], output_filters],
         self.evaluate(upsampled_output_shape))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testCycleGANUpsampleConv2dTranspose(self):
     batch = 8
     height = 32
@@ -724,6 +732,7 @@ def testSpectralNorm(self):
 
 class FnWithCustomGradTest(tf.test.TestCase):
 
+  @test_utils.run_in_graph_mode_only()
   def testCorrectness(self):
 
     w = tf.random_uniform([6, 10])
@@ -769,6 +778,7 @@ def grad_fn(inputs, variables, outputs, grad_outputs):
       for g1, g2 in zip(grads_val, custom_grads_val):
         self.assertAllClose(g1, g2)
 
+  @test_utils.run_in_graph_mode_only()
   def testCustomGrad(self):
 
     def fn(a, b, c):
@@ -801,6 +811,7 @@ def grad_fn(inputs, variables, unused_outputs, unused_grad_outputs):
 
 class RecomputeTest(tf.test.TestCase):
 
+  @test_utils.run_in_graph_mode_only()
   def testRecompute(self):
 
     def layer(x, name=None):
diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index 2998879ef..a58bda345 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -23,7 +23,10 @@
 import numpy as np
 
 from tensor2tensor.layers import common_video
+from tensor2tensor.utils import test_utils
+
 import tensorflow as tf
+tf.compat.v1.enable_eager_execution()
 
 
 class CommonVideoTest(parameterized.TestCase, tf.test.TestCase):
@@ -37,45 +40,45 @@ def _run_scheduled_sample_func(self, func, var, batch_size):
     output = self.evaluate([ground_truth_x, generated_x, ss_out])
     return output
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testScheduledSampleProbStart(self):
     ground_truth_x, _, ss_out = self._run_scheduled_sample_func(
         common_video.scheduled_sample_prob, 1.0, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testScheduledSampleProbMid(self):
     _, _, ss_out = self._run_scheduled_sample_func(
         common_video.scheduled_sample_prob, 0.5, 1000)
     positive_count = np.sum(ss_out > 0)
     self.assertAlmostEqual(positive_count / 1000.0, 0.5, places=1)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testScheduledSampleProbEnd(self):
     _, generated_x, ss_out = self._run_scheduled_sample_func(
         common_video.scheduled_sample_prob, 0.0, 10)
     self.assertAllEqual(generated_x, ss_out)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testScheduledSampleCountStart(self):
     ground_truth_x, _, ss_out = self._run_scheduled_sample_func(
         common_video.scheduled_sample_count, 10, 10)
     self.assertAllEqual(ground_truth_x, ss_out)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testScheduledSampleCountMid(self):
     _, _, ss_out = self._run_scheduled_sample_func(
         common_video.scheduled_sample_count, 5, 10)
     positive_count = np.sum(ss_out > 0)
     self.assertEqual(positive_count, 5)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testScheduledSampleCountEnd(self):
     _, generated_x, ss_out = self._run_scheduled_sample_func(
         common_video.scheduled_sample_count, 0, 10)
     self.assertAllEqual(generated_x, ss_out)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testDynamicTileAndConcat(self):
     # image = (1 X 4 X 4 X 1)
     image = [[1, 2, 3, 4],
@@ -103,6 +106,7 @@ def testDynamicTileAndConcat(self):
          [90, 90, 90, 90],
          [100, 100, 100, 100]])
 
+  @test_utils.run_in_graph_mode_only()
   def testGifSummary(self):
     for c in (1, 3):
       images_shape = (1, 12, 48, 64, c)  # batch, time, height, width, channels
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index 0b29a04f6..effe43c4b 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -21,7 +21,10 @@
 
 import numpy as np
 from tensor2tensor.layers import discretization
+from tensor2tensor.utils import test_utils
+
 import tensorflow as tf
+tf.compat.v1.enable_eager_execution()
 
 
 class DiscretizationTest(tf.test.TestCase):
@@ -31,7 +34,7 @@ def setUp(self):
     tf.set_random_seed(1234)
     np.random.seed(123)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testBitToIntZeros(self):
     x_bit = tf.zeros(shape=[1, 10], dtype=tf.float32)
     x_int = tf.zeros(shape=[1], dtype=tf.int32)
@@ -39,7 +42,7 @@ def testBitToIntZeros(self):
     d = self.evaluate(diff)
     self.assertEqual(d, 0)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testBitToIntOnes(self):
     x_bit = tf.ones(shape=[1, 3], dtype=tf.float32)
     x_int = 7 * tf.ones(shape=[1], dtype=tf.int32)
@@ -47,7 +50,7 @@ def testBitToIntOnes(self):
     d = self.evaluate(diff)
     self.assertEqual(d, 0)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testIntToBitZeros(self):
     x_bit = tf.zeros(shape=[1, 10], dtype=tf.float32)
     x_int = tf.zeros(shape=[1], dtype=tf.int32)
@@ -55,7 +58,7 @@ def testIntToBitZeros(self):
     d = self.evaluate(diff)
     self.assertTrue(np.all(d == 0))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testIntToBitOnes(self):
     x_bit = tf.ones(shape=[1, 3], dtype=tf.float32)
     x_int = 7 * tf.ones(shape=[1], dtype=tf.int32)
@@ -63,7 +66,7 @@ def testIntToBitOnes(self):
     d = self.evaluate(diff)
     self.assertTrue(np.all(d == 0))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testProjectHidden(self):
     hidden_size = 60
     block_dim = 20
@@ -77,7 +80,7 @@ def testProjectHidden(self):
     self.assertEqual(np.shape(x_projected_eval), (1, 1, num_blocks, block_dim))
     self.assertTrue(np.all(x_projected_eval == 0))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testSliceHiddenZeros(self):
     hidden_size = 60
     block_dim = 20
@@ -88,7 +91,7 @@ def testSliceHiddenZeros(self):
     self.assertEqual(np.shape(x_sliced_eval), (1, 1, num_blocks, block_dim))
     self.assertTrue(np.all(x_sliced_eval == 0))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testSliceHiddenOnes(self):
     hidden_size = 60
     block_dim = 20
@@ -99,7 +102,7 @@ def testSliceHiddenOnes(self):
     self.assertEqual(np.shape(x_sliced_eval), (1, 1, num_blocks, block_dim))
     self.assertTrue(np.all(x_sliced_eval == 1))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testNearestNeighbors(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
     x = tf.reshape(x, [1, 1, 2, 3])
@@ -114,6 +117,7 @@ def testNearestNeighbors(self):
     self.assertEqual(np.shape(x_means_hot_eval), (1, 2, 4))
     self.assertTrue(np.all(x_means_hot_eval == x_means_hot_test))
 
+  @test_utils.run_in_graph_mode_only()
   def testGetVQBottleneck(self):
     bottleneck_bits = 2
     bottleneck_size = 2**bottleneck_bits
@@ -129,7 +133,7 @@ def testGetVQBottleneck(self):
       self.assertTrue(np.all(sess.run(means_new) == 0))
       self.assertTrue(np.all(sess.run(ema_count) == 0))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testVQNearestNeighbors(self):
     x = tf.constant([[0, 0.9, 0], [0.8, 0., 0.]], dtype=tf.float32)
     means = tf.constant(
@@ -163,6 +167,7 @@ def testGumbelSoftmaxDiscreteBottleneck(self):
     x_means_hot_eval = self.evaluate(x_means_hot)
     self.assertEqual(np.shape(x_means_hot_eval), (2, 4))
 
+  @test_utils.run_in_graph_mode_only()
   def testDiscreteBottleneckVQ(self):
     hidden_size = 60
     z_size = 4
@@ -198,6 +203,7 @@ def testDiscreteBottleneckVQ(self):
       self.assertTrue(np.all(means_eval == np.zeros(
           (1, 1, 2**z_size, hidden_size))))
 
+  @test_utils.run_in_graph_mode_only()
   def testDiscreteBottleneckVQCond(self):
     hidden_size = 60
     z_size = 4
diff --git a/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
index 884db383c..0e5b2ac73 100644
--- a/tensor2tensor/layers/latent_layers_test.py
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -26,8 +26,10 @@
 from tensor2tensor.layers import discretization
 from tensor2tensor.layers import latent_layers
 from tensor2tensor.models import transformer
+from tensor2tensor.utils import test_utils
 
 import tensorflow as tf
+tf.compat.v1.enable_eager_execution()
 
 
 def imagetransformer_latent_tiny():
@@ -90,7 +92,7 @@ def imagetransformer_latent_tiny():
 
 class LatentLayersTest(tf.test.TestCase):
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testTransformerAutoencoder(self):
     hparams = imagetransformer_latent_tiny()
     hparams.mode = tf.estimator.ModeKeys.TRAIN
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 71ba33b86..7ae536207 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -23,13 +23,15 @@
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import expert_utils
+from tensor2tensor.utils import test_utils
 
 import tensorflow as tf
+tf.compat.v1.enable_eager_execution()
 
 
 class ModalityTest(tf.test.TestCase):
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testSymbolModalityInputs(self):
     batch_size = 10
     num_datashards = 5
@@ -51,7 +53,7 @@ def testSymbolModalityInputs(self):
     res = self.evaluate(output)
     self.assertEqual(res.shape, (batch_size, length, 1, hidden_size))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testSymbolModalityTargets(self):
     batch_size = 10
     num_datashards = 5
@@ -81,6 +83,7 @@ def testSymbolModalityTargets(self):
     self.assertEqual(res1.shape, (batch_size, length, height, 1, vocab_size))
     self.assertEqual(res2.shape, ())
 
+  @test_utils.run_in_graph_mode_only()
   def testSymbolModalityTargetsFactored(self):
     batch_size = 10
     num_datashards = 5
diff --git a/tensor2tensor/layers/ngram_test.py b/tensor2tensor/layers/ngram_test.py
index 0eac013f1..475d5f4cb 100644
--- a/tensor2tensor/layers/ngram_test.py
+++ b/tensor2tensor/layers/ngram_test.py
@@ -21,12 +21,15 @@
 
 from tensor2tensor.layers import ngram
 
+from tensor2tensor.utils import test_utils
+
 import tensorflow as tf
+tf.compat.v1.enable_eager_execution()
 
 
 class NGramTest(tf.test.TestCase):
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testNGramLayerShape(self):
     batch_size = 2
     length = 8
@@ -41,7 +44,7 @@ def testNGramLayerShape(self):
     num_ngrams = sum([vocab_size**n for n in range(minval, maxval)])
     self.assertEqual(outputs_val.shape, (batch_size, num_ngrams))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testNGramLayerOutput(self):
     inputs = tf.constant(
         [[0, 0, 0, 0, 1],
diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
index ec6b67e3f..309e543f6 100644
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -24,12 +24,16 @@
 
 from tensor2tensor.layers import reversible_layers as reversible
 
+from tensor2tensor.utils import test_utils
+
 import tensorflow as tf
+tf.compat.v1.enable_eager_execution()
+
 
 
 class ReversibleLayersTest(parameterized.TestCase, tf.test.TestCase):
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testActNorm(self):
     np.random.seed(83243)
     batch_size = 25
@@ -54,7 +58,7 @@ def testActNorm(self):
     self.assertAllClose(mean_val, np.zeros(channels), atol=0.25)
     self.assertAllClose(variance_val, np.ones(channels), atol=0.25)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testMADELeftToRight(self):
     np.random.seed(83243)
     batch_size = 2
@@ -76,7 +80,7 @@ def testMADELeftToRight(self):
     self.assertAllEqual(outputs_val[:, 0, :], np.zeros((batch_size, units)))
     self.assertEqual(outputs_val.shape, (batch_size, length, units))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testMADERightToLeft(self):
     np.random.seed(1328)
     batch_size = 2
@@ -101,7 +105,7 @@ def testMADERightToLeft(self):
     self.assertAllEqual(outputs_val[:, -1, :], np.zeros((batch_size, units)))
     self.assertEqual(outputs_val.shape, (batch_size, length, units))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @test_utils.run_in_graph_and_eager_modes()
   def testMADENoHidden(self):
     np.random.seed(532)
     batch_size = 2

From c8f1afc032a9d183dcfbaddd3e848bd82a344d7e Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 4 Feb 2019 11:43:20 -0800
Subject: [PATCH 1627/2720] Fix contrib.layers optimizer registration names

PiperOrigin-RevId: 232336767
---
 tensor2tensor/utils/optimize.py      | 16 +++++-----
 tensor2tensor/utils/optimize_test.py | 48 ++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 8 deletions(-)
 create mode 100644 tensor2tensor/utils/optimize_test.py

diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 9ff1993ad..dda809797 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -21,6 +21,7 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import adafactor as adafactor_lib
+from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import multistep_optimizer
 from tensor2tensor.utils import registry
@@ -162,17 +163,16 @@ def adafactor(learning_rate, hparams):
 
 
 
-def _register_base_optimizer(key, fn):
+def _register_base_optimizer(name, opt):
+  key = misc_utils.camelcase_to_snakecase(name)
+  if key in registry.Registries.optimizers:
+    return
   registry.register_optimizer(key)(
-      lambda learning_rate, hparams: fn(learning_rate))
+      lambda learning_rate, hparams: opt(learning_rate))
 
 
-for k in tf.contrib.layers.OPTIMIZER_CLS_NAMES:
-  if k not in registry.Registries.optimizers and k not in ("SGD", "RMSProp"):
-    _register_base_optimizer(k, tf.contrib.layers.OPTIMIZER_CLS_NAMES[k])
-_register_base_optimizer("sgd", tf.contrib.layers.OPTIMIZER_CLS_NAMES["SGD"])
-_register_base_optimizer(
-    "rms_prop", tf.contrib.layers.OPTIMIZER_CLS_NAMES["RMSProp"])
+for _name, _opt in tf.contrib.layers.OPTIMIZER_CLS_NAMES.items():
+  _register_base_optimizer(_name, _opt)
 
 
 class ConditionalOptimizer(tf.train.Optimizer):
diff --git a/tensor2tensor/utils/optimize_test.py b/tensor2tensor/utils/optimize_test.py
new file mode 100644
index 000000000..ea1478647
--- /dev/null
+++ b/tensor2tensor/utils/optimize_test.py
@@ -0,0 +1,48 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.utils.optimize."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+from tensor2tensor.utils import hparams_lib
+from tensor2tensor.utils import optimize
+import tensorflow as tf
+
+
+class OptimizeTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(
+      "sgd",
+      "SGD",
+      "rms_prop",
+      "RMSProp",
+      "adagrad",
+      "Adagrad",
+      "adam",
+      "Adam",
+      "adam_w",
+      "AdamW",
+  )
+  def test_names(self, opt_name):
+    hparams = hparams_lib.create_hparams("basic_1")
+    optimize.ConditionalOptimizer(opt_name, 0.1, hparams)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 57a59e1a284125294fb14cdf0a89b7ead88c28ce Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Mon, 4 Feb 2019 13:45:43 -0800
Subject: [PATCH 1628/2720] Latent-space interpolations, Part 2

PiperOrigin-RevId: 232358762
---
 tensor2tensor/models/research/glow_ops.py     |  32 +++++
 .../models/research/glow_ops_test.py          |  22 +++
 tensor2tensor/models/video/nfg_interpolate.py | 136 ++++++++++++------
 3 files changed, 145 insertions(+), 45 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 3a9ba8a76..ef58517c0 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -48,6 +48,38 @@ def linear_interpolate(tensor1, tensor2, coeffs):
   return tf.concat(interp_tensors, axis=0)
 
 
+def linear_interpolate_rank(tensor1, tensor2, coeffs, rank=1):
+  """Linearly interpolate channel at "rank" between two tensors.
+
+  The channels are ranked according to their L2 norm between tensor1[channel]
+  and tensor2[channel].
+
+  Args:
+    tensor1: 4-D Tensor, NHWC
+    tensor2: 4-D Tensor, NHWC
+    coeffs: list of floats.
+    rank: integer.
+  Returns:
+    interp_latents: list of interpolated 4-D Tensors, shape=(NHWC)
+  """
+  # sum across space, max across channels.
+  _, _, _, num_channels = common_layers.shape_list(tensor1)
+  diff_sq_sum = tf.reduce_sum((tensor1 - tensor2)**2, axis=(0, 1, 2))
+  _, feature_ranks = tf.math.top_k(diff_sq_sum, k=rank)
+  feature_rank = feature_ranks[-1]
+  channel_inds = tf.range(num_channels, dtype=tf.int32)
+  channel_mask = tf.equal(channel_inds, feature_rank)
+  ones_t = tf.ones(num_channels, dtype=tf.float32)
+  zeros_t = tf.zeros(num_channels, dtype=tf.float32)
+
+  interp_tensors = []
+  for coeff in coeffs:
+    curr_coeff = tf.where(channel_mask, coeff * ones_t, zeros_t)
+    interp_tensor = tensor1 + curr_coeff * (tensor2 - tensor1)
+    interp_tensors.append(interp_tensor)
+  return tf.concat(interp_tensors, axis=0)
+
+
 def postprocess(x, n_bits_x=8):
   """Converts x from [-0.5, 0.5], to [0, 255].
 
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 45b6ae8d4..dfef7b727 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -484,6 +484,28 @@ def test_temperature_normal(self, temperature):
         self.assertTrue(np.allclose(loc_exp, loc_act, atol=1e-2))
         self.assertTrue(np.allclose(scale_exp, scale_act, atol=1e-2))
 
+  def linear_interpolate_rank(self):
+    with tf.Graph().as_default():
+      # Since rank is 1, the first channel should remain 1.0.
+      # and the second channel should be interpolated between 1.0 and 6.0
+      z1 = np.ones(shape=(4, 4, 2))
+      z2 = np.copy(z1)
+      z2[:, :, 0] += 0.01
+      z2[:, :, 1] += 5.0
+      coeffs = np.linspace(0.0, 1.0, 11)
+      z1 = np.expand_dims(z1, axis=0)
+      z2 = np.expand_dims(z2, axis=0)
+      tensor1 = tf.convert_to_tensor(z1, dtype=tf.float32)
+      tensor2 = tf.convert_to_tensor(z2, dtype=tf.float32)
+      lin_interp_max = glow_ops.linear_interpolate_rank(
+          tensor1, tensor2, coeffs)
+      with tf.Session() as sess:
+        lin_interp_np_max = sess.run(lin_interp_max)
+        for lin_interp_np, coeff in zip(lin_interp_np_max, coeffs):
+          exp_val = 1.0 + coeff * (6.0 - 1.0)
+          self.assertTrue(np.allclose(lin_interp_np[:, :, 0], 1.0))
+          self.assertTrue(np.allclose(lin_interp_np[:, :, 1], exp_val))
+
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/models/video/nfg_interpolate.py b/tensor2tensor/models/video/nfg_interpolate.py
index 71a550a2c..d9d746ea6 100644
--- a/tensor2tensor/models/video/nfg_interpolate.py
+++ b/tensor2tensor/models/video/nfg_interpolate.py
@@ -42,8 +42,6 @@
 flags.DEFINE_string("score_file", "", "File to score. Each line in the file "
                     "must be in the format input \t target.")
 flags.DEFINE_bool("decode_in_memory", False, "Decode in memory.")
-# Interpolate between z1 and z2 for alpha = np.linspace(0.0, 1.0, num_interp)
-flags.DEFINE_integer("num_interp", 11, "Number of interpolations")
 
 flags = tf.flags
 FLAGS = flags.FLAGS
@@ -52,6 +50,23 @@
 arg_scope = tf.contrib.framework.arg_scope
 
 
+def decode_hparams(overrides=""):
+  """Hparams for decoding."""
+  hparams = decoding.decode_hparams()
+  # Number of interpolations between [0.0, 1.0].
+  hparams.add_hparam("num_interp", 11)
+  # Which level(s) to interpolate.
+  hparams.add_hparam("level_interp", [0, 1, 2])
+  # "all" or "ranked", interpolate all channels or a "ranked".
+  hparams.add_hparam("channel_interp", "all")
+  # interpolate channels ranked according to squared L2 norm.
+  hparams.add_hparam("rank_interp", 1)
+  # Whether on not to save frames as summaries
+  hparams.add_hparam("save_frames", True)
+  hparams.parse(overrides)
+  return hparams
+
+
 def preprocess_frame(frame):
   """Preprocess frame.
 
@@ -91,39 +106,66 @@ def latents_to_frames(z_top_interp, level_eps_interp, hparams):
   return images
 
 
-def interpolate(features, hparams, num_interp):
+def interpolate(features, hparams, decode_hp):
   """Interpolate between the first input frame and last target frame.
 
   Args:
     features: dict of tensors
-    hparams: HParams.
-    num_interp: integer.
+    hparams: tf.contrib.training.HParams, training hparams.
+    decode_hp: tf.contrib.training.HParams, decode hparams.
   Returns:
-    images: 4-D Tensor, shape=(num_interp, H, W, C)
+    images: interpolated images, 4-D Tensor, shape=(num_interp, H, W, C)
+    first_frame: image, 3-D Tensor, shape=(1, H, W, C)
+    last_frame: image, 3-D Tensor, shape=(1, H, W, C)
   """
   inputs, targets = features["inputs"], features["targets"]
   inputs = tf.unstack(inputs, axis=1)
   targets = tf.unstack(targets, axis=1)
-  coeffs = np.linspace(0.0, 1.0, num_interp)
+  coeffs = np.linspace(0.0, 1.0, decode_hp.num_interp)
 
   # (X_1, X_t) -> (z_1, z_t)
   first_frame, last_frame = inputs[0], targets[-1]
   first_top_z, first_level_eps = frame_to_latents(first_frame, hparams)
   last_top_z, last_level_eps = frame_to_latents(last_frame, hparams)
 
-  # Interpolate top
-  z_top_interp = glow_ops.linear_interpolate(first_top_z, last_top_z, coeffs)
-
-  # Interpolate level.
-  level_eps_interp = []
-  for level in range(hparams.n_levels - 1):
-    level_eps_interp.append(glow_ops.linear_interpolate(
-        first_level_eps[level], last_level_eps[level], coeffs))
-  return latents_to_frames(z_top_interp, level_eps_interp, hparams)
+  # Interpolate latents at all levels.
+  first_lats = first_level_eps + [first_top_z]
+  last_lats = last_level_eps + [last_top_z]
+  interp_lats = []
+  lat_iterator = enumerate(zip(first_lats, last_lats))
+  for level_ind, (first_lat, last_lat) in lat_iterator:
+    if level_ind in decode_hp.level_interp:
+      if decode_hp.channel_interp == "all":
+        interp_lat = glow_ops.linear_interpolate(first_lat, last_lat, coeffs)
+      else:
+        interp_lat = glow_ops.linear_interpolate_rank(
+            first_lat, last_lat, coeffs, decode_hp.rank_interp)
+    else:
+      interp_lat = tf.tile(first_lat, [decode_hp.num_interp, 1, 1, 1])
+    interp_lats.append(interp_lat)
+
+  level_eps_interp = interp_lats[:hparams.n_levels-1]
+  z_top_interp = interp_lats[-1]
+  images = latents_to_frames(z_top_interp, level_eps_interp, hparams)
+  return images, first_frame, last_frame
+
+
+def get_summaries_log_dir(decode_hp, output_dir, dataset_split):
+  """Get nested summaries_log_dir based on decode_hp."""
+  child_dir = decode_hp.summaries_log_dir
+  level_dir = "".join([str(level) for level in decode_hp.level_interp])
+  if decode_hp.channel_interp == "all":
+    rank_dir = "all"
+  else:
+    rank_dir = "rank_%d" % decode_hp.rank_interp
+  child_dir = "%s/%s_%s" % (child_dir, level_dir, rank_dir)
+  if dataset_split is not None:
+    child_dir += "_{}".format(dataset_split)
+  return os.path.join(output_dir, child_dir)
 
 
-def interpolations_to_summary(sample_ind, interpolations, hparams,
-                              decode_hparams):
+def interpolations_to_summary(sample_ind, interpolations, first_frame,
+                              last_frame, hparams, decode_hp):
   """Converts interpolated frames into tf summaries.
 
   The summaries consists of:
@@ -133,33 +175,38 @@ def interpolations_to_summary(sample_ind, interpolations, hparams,
 
   Args:
     sample_ind: int
-    interpolations: Numpy array, shape=(num_interp, 64, 64, 3)
-    hparams: HParams, train hparams
-    decode_hparams: HParams, decode hparams
+    interpolations: Numpy array, shape=(num_interp, H, W, 3)
+    first_frame: Numpy array, shape=(HWC)
+    last_frame: Numpy array, shape=(HWC)
+    hparams: tf.contrib.training.HParams, train hparams
+    decode_hp: tf.contrib.training.HParams, decode hparams
   Returns:
     summaries: list of tf Summary Values.
   """
   parent_tag = "sample_%d" % sample_ind
   frame_shape = hparams.problem.frame_shape
-  interp_shape = [hparams.batch_size, FLAGS.num_interp] + frame_shape
+  interp_shape = [hparams.batch_size, decode_hp.num_interp] + frame_shape
   interpolations = np.reshape(interpolations, interp_shape)
+  interp_tag = "%s/interp/%s" % (parent_tag, decode_hp.channel_interp)
+  if decode_hp.channel_interp == "ranked":
+    interp_tag = "%s/rank_%d" % (interp_tag, decode_hp.rank_interp)
   summaries, _ = common_video.py_gif_summary(
-      parent_tag, interpolations, return_summary_value=True,
-      max_outputs=decode_hparams.max_display_outputs,
-      fps=decode_hparams.frames_per_second)
-
-  first_frame, last_frame = interpolations[0, 0], interpolations[0, -1]
-  first_frame_summ = image_utils.image_to_tf_summary_value(
-      first_frame, "%s/first" % parent_tag)
-  last_frame_summ = image_utils.image_to_tf_summary_value(
-      last_frame, "%s/last" % parent_tag)
-  summaries.append(first_frame_summ)
-  summaries.append(last_frame_summ)
+      interp_tag, interpolations, return_summary_value=True,
+      max_outputs=decode_hp.max_display_outputs,
+      fps=decode_hp.frames_per_second)
+
+  if decode_hp.save_frames:
+    first_frame_summ = image_utils.image_to_tf_summary_value(
+        first_frame, "%s/first" % parent_tag)
+    last_frame_summ = image_utils.image_to_tf_summary_value(
+        last_frame, "%s/last" % parent_tag)
+    summaries.append(first_frame_summ)
+    summaries.append(last_frame_summ)
   return summaries
 
 
 def main(_):
-  decode_hparams = decoding.decode_hparams(FLAGS.decode_hparams)
+  decode_hp = decode_hparams(FLAGS.decode_hparams)
   trainer_lib.set_random_seed(FLAGS.random_seed)
   if FLAGS.output_dir is None:
     raise ValueError("Expected output_dir to be set to a valid path.")
@@ -182,35 +229,34 @@ def main(_):
   ops = [glow_ops.get_variable_ddi, glow_ops.actnorm, glow_ops.get_dropout]
   var_scope = tf.variable_scope("next_frame_glow/body", reuse=tf.AUTO_REUSE)
   with arg_scope(ops, init=False), var_scope:
-    interpolations = interpolate(dataset, hparams, FLAGS.num_interp)
+    interpolations, first_frame, last_frame = interpolate(
+        dataset, hparams, decode_hp)
 
   var_list = tf.global_variables()
   saver = tf.train.Saver(var_list)
 
   # Get latest checkpoints from model_dir.
   ckpt_path = tf.train.latest_checkpoint(FLAGS.output_dir)
-  child_dir = decode_hparams.summaries_log_dir
-  if dataset_split is not None:
-    child_dir += "_{}".format(dataset_split)
-  final_dir = os.path.join(FLAGS.output_dir, child_dir)
+  final_dir = get_summaries_log_dir(decode_hp, FLAGS.output_dir, dataset_split)
   summary_writer = tf.summary.FileWriter(final_dir)
   global_step = decoding.latest_checkpoint_step(FLAGS.output_dir)
 
   sample_ind = 0
-
-  num_samples = decode_hparams.num_samples
+  num_samples = decode_hp.num_samples
   all_summaries = []
 
   with tf.train.MonitoredTrainingSession() as sess:
     saver.restore(sess, ckpt_path)
 
     while not sess.should_stop() and sample_ind < num_samples:
-      interp_np = sess.run(interpolations)
+      interp_np, first_frame_np, last_frame_np = sess.run(
+          [interpolations, first_frame, last_frame])
 
-      interp_summ = interpolations_to_summary(sample_ind, interp_np, hparams,
-                                              decode_hparams)
+      interp_summ = interpolations_to_summary(sample_ind, interp_np,
+                                              first_frame_np[0],
+                                              last_frame_np[0],
+                                              hparams, decode_hp)
       all_summaries.extend(interp_summ)
-
       sample_ind += 1
     all_summaries = tf.Summary(value=list(all_summaries))
     summary_writer.add_summary(all_summaries, global_step)

From 9dfa71b61c861a9df84ffc13a0be753a59d68c1f Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Mon, 4 Feb 2019 15:15:30 -0800
Subject: [PATCH 1629/2720] Add multiproblems for Squad in multi problem

PiperOrigin-RevId: 232376213
---
 tensor2tensor/data_generators/squad.py            | 12 ++++++++++++
 .../data_generators/wiki_multi_problems.py        | 15 +++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/tensor2tensor/data_generators/squad.py b/tensor2tensor/data_generators/squad.py
index bf9646206..1cf541984 100644
--- a/tensor2tensor/data_generators/squad.py
+++ b/tensor2tensor/data_generators/squad.py
@@ -152,6 +152,18 @@ def vocab_filename(self):
     return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
 
 
+@registry.register_problem
+class SquadConcatSharedVocab(SquadConcatMulti64k):
+  """Squad with question and context concatenated, multi-lingual vocabulary."""
+
+  def dataset_filename(self):
+    return "squad"
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelEnWiki32k().vocab_filename
+
+
 @registry.register_problem
 class SquadConcatPositioned(SquadConcat):
   """SquadConcat with targets in format of answer position + answer length."""
diff --git a/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index 87b146c1a..faef16eea 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -251,3 +251,18 @@ def __init__(self, was_reversed=False, was_copy=False):
   @property
   def vocab_type(self):
     return text_problems.VocabType.SUBWORD
+
+
+@registry.register_problem
+class LanguagemodelEnWikiLMSquadConcatSubwords(multi_problem.MultiProblem):
+  """Wiki LM and MNLI mixed problem class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelEnWikiLMSquadConcatSubwords, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelEnWiki32k())
+    self.task_list.append(multinli.SquadConcatSharedVocab())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD

From 5997c193654072802242a31f2fb2e2f4a503ba0e Mon Sep 17 00:00:00 2001
From: Alexander Ku <alexku@google.com>
Date: Mon, 4 Feb 2019 20:38:55 -0800
Subject: [PATCH 1630/2720] Generalizing the MultiProblem class and adding
 tests.

PiperOrigin-RevId: 232416765
---
 .../data_generators/multi_problem_v2.py       | 356 ++++++++++++++++++
 .../data_generators/multi_problem_v2_test.py  | 207 ++++++++++
 .../data_generators/wiki_multi_problems.py    |  28 ++
 3 files changed, 591 insertions(+)
 create mode 100644 tensor2tensor/data_generators/multi_problem_v2.py
 create mode 100644 tensor2tensor/data_generators/multi_problem_v2_test.py

diff --git a/tensor2tensor/data_generators/multi_problem_v2.py b/tensor2tensor/data_generators/multi_problem_v2.py
new file mode 100644
index 000000000..7f0717bcf
--- /dev/null
+++ b/tensor2tensor/data_generators/multi_problem_v2.py
@@ -0,0 +1,356 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Multi-problem scheduling in T2T.
+
+Data sampling schedules are specified by an interpolation method i and a
+sequence of tuples (t, pmf), where i can either be 'linear' or 'step',
+t is the global_step at training, and pmf is the distribution from which
+training examples from each problem are sampled.
+
+Linear interpolation constructs a piecewise linear training schedule, connecting
+pmfs with linear segments. Step interpolation abruptly shifts the sampling
+distribution to pmf at global_step t. Both interpolation methods can approximate
+any continuous sampling process with sufficient points of interpolation.
+
+Continuation of the interpolant is constant outside the domain specified by
+the schedule. That is, we sample from pmfs[0] for global_step < ts[0] and
+pmfs[-1] for global_step > ts[-1].
+
+Examples of schedule strings include:
+
+(1) 'step @0 0.7, 0.3': Sample from problem 0 w.p. 0.7 and problem 1 w.p. 0.3
+    for the entirety of training. Since there is only one point, the choice of
+    interpolation method and global_step does not matter.
+
+(2) 'step @0 1.0 0.0 @100 0.0 1.0': Train on problem 0 for the first 100 steps
+    then train on problem 1 for the rest of training.
+
+(3) 'step @0 0.5 0.5 0.0 @100 1.0 0.0 0.0': Pretrain on problems 0 and 1 for the
+    first 100 steps then fine tune on problem 2 for the rest of training.
+
+(4) 'linear @0 1.0 0.0 @100 0.0 1.0' Linear transition from training on problem
+    0 to problem 1 over 100 steps, then train on problem 1 for the rest of
+    training.
+
+(5) 'linear @0 1.0 0.0 @100 0.9 0.1  @200 0.4 0.6  @300 0.0 1.0': Approximate
+    inverse exponential decay from problem 0 to problem 1 over 300 steps, then
+    train on problem 1 for the rest of training.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+import os
+
+import numpy as np
+
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+import tensorflow as tf
+
+
+class MultiProblemV2(problem.Problem):
+  """Dataset scheduling for multiple problems."""
+
+  def __init__(self, problems, schedule, **kwargs):
+    """Creates a MultiProblem object.
+
+    Args:
+      problems: A list of problem.Problem objects.
+      schedule: A schedule tuple, see encode_schedule for details.
+      **kwargs: Keywords for problem.Problem.__init__.
+    """
+    super(MultiProblemV2, self).__init__(**kwargs)
+    self.problems = problems
+    self.schedule = schedule
+
+  def filepattern(self, *args, **kwargs):
+    """Returns a list of filepatterns, one for each problem."""
+    return [p.filepattern(*args, **kwargs) for p in self.problems]
+
+  def generate_data(self, *args, **kwargs):
+    """Generates data for each problem."""
+    for p in self.problems:
+      p.generate_data(*args, **kwargs)
+
+  @property
+  def only_eval_first_problem(self):
+    """Only run validation on examples from the first problem."""
+    return False
+
+  def normalize_example(self, example, hparams):
+    """Preprocesses examples from different problems before mixing."""
+    del hparams  # Unused.
+    return example
+
+  def dataset(self, mode, hparams=None, global_step=None, **kwargs):
+    """Returns a dataset containing examples from multiple problems.
+
+    Args:
+      mode: A member of problem.DatasetSplit.
+      hparams: A tf.HParams object, the model hparams.
+      global_step: A scalar tensor used to compute the sampling distribution.
+        If global_step is None, we call tf.train.get_or_create_global_step by
+        default.
+      **kwargs: Keywords for problem.Problem.Dataset.
+
+    Returns:
+      A dataset containing examples from multiple problems.
+    """
+    datasets = [p.dataset(mode, **kwargs) for p in self.problems]
+    datasets = [
+        d.map(lambda x, i=j: self.normalize_example(  # pylint: disable=g-long-lambda
+            dict(x, problem_id=tf.constant([i])), hparams))
+        for j, d in enumerate(datasets)  # Tag examples with a problem_id.
+    ]
+    if mode is problem.DatasetSplit.TRAIN:
+      if global_step is None:
+        global_step = tf.train.get_or_create_global_step()
+      pmf = get_schedule_distribution(self.schedule, global_step)
+      return get_multi_dataset(datasets, pmf)
+    elif self.only_eval_first_problem:
+      return datasets[0]
+    else:
+      datasets = [d.repeat() for d in datasets]
+      return tf.data.Dataset.zip(tuple(datasets)).flat_map(
+          lambda *x: functools.reduce(  # pylint: disable=g-long-lambda
+              tf.data.Dataset.concatenate,
+              map(tf.data.Dataset.from_tensors, x)))
+
+
+class MultiText2TextProblem(MultiProblemV2, text_problems.Text2TextProblem):
+  """Dataset scheduling for multiple text-to-text problems."""
+
+  def normalize_example(self, example, hparams):
+    """Assumes that example contains both inputs and targets."""
+
+    def _to_constant_shape(tensor):
+      max_length = self.max_length(hparams)
+      tensor = tensor[:max_length]
+      tensor = tf.pad(tensor, [(0, max_length - tf.shape(tensor)[0])])
+      return tf.reshape(tensor, [max_length])
+
+    if self.has_inputs:
+      example['inputs'] = _to_constant_shape(example['inputs'])
+      example['targets'] = _to_constant_shape(example['targets'])
+    elif 'inputs' in example:
+      inputs = example.pop('inputs')[:-1]  # Remove EOS token.
+      targets = tf.concat([inputs, example['targets']], 0)
+      example['targets'] = _to_constant_shape(targets)
+    else:
+      example['targets'] = _to_constant_shape(example['targets'])
+    return example
+
+  def generate_data_with_shared_vocab(self, data_dir, tmp_dir, task_id=-1):
+    """Generates TF-Records for problems using a global vocabulary file."""
+    global_vocab_filename = os.path.join(data_dir, self.vocab_filename)
+    if not tf.gfile.Exists(global_vocab_filename):
+      raise ValueError(
+          'Global vocabulary file: %s does not exist, '
+          'please create one using build_vocab.py' % global_vocab_filename)
+    # Before generating data, we copy the global vocabulary file to the children
+    # locations. Although this is not the most disk efficient strategy, it
+    # imposes the fewest changes to the text-to-text API.
+    for p in self.problems:
+      local_vocab_filename = os.path.join(data_dir, p.vocab_filename)
+      if not tf.gfile.Exists(local_vocab_filename):
+        tf.gfile.Copy(global_vocab_filename, local_vocab_filename)
+      p.generate_data(data_dir, tmp_dir, task_id)
+
+
+def get_multi_dataset(datasets, pmf=None):
+  """Returns a Dataset that samples records from one or more Datasets.
+
+  Args:
+    datasets: A list of one or more Dataset objects to sample from.
+    pmf: A tensor of shape [len(datasets)], the probabilities to sample each
+      dataset with. This tensor is often constructed with the global_step. If
+      this is None, we sample from the datasets uniformly at random.
+
+  Returns:
+    A Dataset object containing records from multiple datasets. Note that
+    because this dataset iterates through other datasets it is stateful, thus
+    you will need to call make_initializable_iterator instead of
+    make_one_shot_iterator.
+  """
+  pmf = tf.fill([len(datasets)], 1.0 / len(datasets)) if pmf is None else pmf
+  samplers = [d.repeat().make_one_shot_iterator().get_next for d in datasets]
+  sample = lambda _: categorical_case(pmf, samplers)
+  return tf.data.Dataset.from_tensors([]).repeat().map(sample)
+
+
+def get_schedule_distribution(schedule, global_step=None):
+  """Computes the pmf of a schedule given the global_step.
+
+  Args:
+    schedule: A schedule tuple, see encode_schedule for details.
+    global_step: A scalar tensor, the step to query the schedule.
+
+  Returns:
+    A 1-D tensor of probs, the sampling distribution of the global_step.
+  """
+  interpolation, steps, pmfs = schedule
+  if global_step is None:
+    global_step = tf.train.get_or_create_global_step()
+  if interpolation == 'step':
+    interpolation_fn = step_interpolation
+  elif interpolation == 'linear':
+    interpolation_fn = linear_interpolation
+  else:
+    raise ValueError('Invalid interpolation strategy: %s' % interpolation)
+  return tf.reshape(
+      tf.py_func(
+          func=lambda x: interpolation_fn(x, np.array(steps), np.array(pmfs)),
+          inp=[global_step], Tout=tf.float32), [len(pmfs[0])])
+
+
+def categorical_case(pmf, fns, rand=None):
+  """Returns the outputs of fns[i] with probability pmf[i].
+
+  Args:
+    pmf: A 1-D tensor of probabilities, the probability mass function.
+    fns: A list of callables that return tensors, same length as pmf.
+    rand: An optional scalar between 0.0 and 1.0, the output of an RNG.
+
+  Returns:
+    A tensor, the output of fns[i] with probability pmf[i].
+  """
+  rand = tf.random_uniform([]) if rand is None else rand
+  cmf = tf.pad(tf.cumsum(pmf), [(1, 0)])
+  cmf = [cmf[i] for i in range(len(fns) + 1)]
+  preds = [(rand >= a) & (rand < b) for a, b in zip(cmf[:-1], cmf[1:])]
+  return tf.case(list(zip(preds, fns)), exclusive=True)
+
+
+def linear_interpolation(x, xp, fp, **kwargs):
+  """Multi-dimensional linear interpolation.
+
+  Returns the multi-dimensional piecewise linear interpolant to a function with
+  given discrete data points (xp, fp), evaluated at x.
+
+  Note that *N and *M indicate zero or more dimensions.
+
+  Args:
+    x: An array of shape [*N], the x-coordinates of the interpolated values.
+    xp: An np.array of shape [D], the x-coordinates of the data points, must be
+      increasing.
+    fp: An np.array of shape [D, *M], the y-coordinates of the data points.
+    **kwargs: Keywords for np.interp.
+
+  Returns:
+    An array of shape [*N, *M], the interpolated values.
+  """
+  yp = fp.reshape([fp.shape[0], -1]).transpose()
+  y = np.stack([np.interp(x, xp, zp, **kwargs) for zp in yp]).transpose()
+  return y.reshape(x.shape[:1] + fp.shape[1:]).astype(np.float32)
+
+
+def step_interpolation(x, xp, fp, **kwargs):
+  """Multi-dimensional step interpolation.
+
+  Returns the multi-dimensional step interpolant to a function with
+  given discrete data points (xp, fp), evaluated at x.
+
+  Note that *N and *M indicate zero or more dimensions.
+
+  Args:
+    x: An array of shape [*N], the x-coordinates of the interpolated values.
+    xp: An np.array of shape [D], the x-coordinates of the data points, must be
+      increasing.
+    fp: An np.array of shape [D, *M], the y-coordinates of the data points.
+    **kwargs: Unused.
+
+  Returns:
+    An array of shape [*N, *M], the interpolated values.
+  """
+  del kwargs  # Unused.
+  xp = np.expand_dims(xp, -1)
+  lower, upper = xp[:-1], xp[1:]
+  conditions = (x >= lower) & (x < upper)
+  # Underflow and overflow conditions and values. Values default to fp[0] and
+  # fp[-1] respectively.
+  conditions = np.concatenate([[x < xp[0]], conditions, [x >= xp[-1]]])
+  values = np.concatenate([[fp[0]], fp])
+  assert np.all(np.sum(conditions, 0) == 1), 'xp must be increasing.'
+  indices = np.argmax(conditions, 0)
+  return values[indices].astype(np.float32)
+
+
+def constant_schedule(pmf):
+  """Returns a schedule tuple for constant sampling distribution.
+
+  Args:
+    pmf: An array of shape [N] of probabilities. The sampling distribution to
+      use throughout training. Probabilities must sum to one.
+
+  Returns:
+    A schedule tuple, see encode_schedule for details.
+  """
+  return ('step', (0,), (tuplize(pmf),))
+
+
+def encode_schedule(schedule):
+  """Encodes a schedule tuple into a string.
+
+  Args:
+    schedule: A tuple containing (interpolation, steps, pmfs), where
+      interpolation is a string specifying the interpolation strategy, steps
+      is an int array_like of shape [N] specifying the global steps, and pmfs is
+      an array_like of shape [N, M] where pmf[i] is the sampling distribution
+      at global step steps[i]. N is the number of schedule requirements to
+      interpolate and M is the size of the probability space.
+
+  Returns:
+    The string encoding of the schedule tuple.
+  """
+  interpolation, steps, pmfs = schedule
+  return interpolation + ' ' + ' '.join(
+      '@' + str(s) + ' ' + ' '.join(map(str, p)) for s, p in zip(steps, pmfs))
+
+
+def decode_schedule(string):
+  """Decodes a string into a schedule tuple.
+
+  Args:
+    string: The string encoding of a schedule tuple.
+
+  Returns:
+    A schedule tuple, see encode_schedule for details.
+  """
+  splits = string.split()
+  steps = [int(x[1:]) for x in splits[1:] if x[0] == '@']
+  pmfs = np.reshape(
+      [float(x) for x in splits[1:] if x[0] != '@'], [len(steps), -1])
+  return splits[0], tuplize(steps), tuplize(pmfs)
+
+
+def tuplize(nested):
+  """Recursively converts iterables into tuples.
+
+  Args:
+    nested: A nested structure of items and iterables.
+
+  Returns:
+    A nested structure of items and tuples.
+  """
+  if isinstance(nested, str):
+    return nested
+  try:
+    return tuple(map(tuplize, nested))
+  except TypeError:
+    return nested
diff --git a/tensor2tensor/data_generators/multi_problem_v2_test.py b/tensor2tensor/data_generators/multi_problem_v2_test.py
new file mode 100644
index 000000000..c189856de
--- /dev/null
+++ b/tensor2tensor/data_generators/multi_problem_v2_test.py
@@ -0,0 +1,207 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.data_generators.multi_problem."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensor2tensor.data_generators import multi_problem_v2
+from tensor2tensor.data_generators import problem
+import tensorflow as tf
+
+
+class MultiProblemV2Test(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(
+      {
+          'inputs': [(0.0, ['string', 12]), np.array([12, 10])],
+          'targets': ((0.0, ('string', 12)), (12, 10)),
+      },
+      {
+          'inputs': [1.0, np.ones([2, 3])],
+          'targets': (1.0, ((1.0, 1.0, 1.0), (1.0, 1.0, 1.0))),
+      },
+  )
+  def test_tuplize(self, inputs, targets):
+    self.assertEqual(multi_problem_v2.tuplize(inputs), targets)
+
+  @parameterized.parameters(
+      {
+          'schedule': ('step', (100,), ((0.25, 0.75),)),
+          'string': 'step @100 0.25 0.75',
+      },
+      {
+          'schedule': ('step', (100, 200), ((0.25, 0.75), (0.62, 0.38))),
+          'string': 'step @100 0.25 0.75 @200 0.62 0.38',
+      },
+      {
+          'schedule': ('linear', (100, 200), ((0.25, 0.75), (0.62, 0.38))),
+          'string': 'linear @100 0.25 0.75 @200 0.62 0.38',
+      },
+  )
+  def test_encode_decode_schedule(self, schedule, string):
+    self.assertEqual(multi_problem_v2.encode_schedule(schedule), string)
+    self.assertEqual(multi_problem_v2.decode_schedule(string), schedule)
+
+  @parameterized.parameters(
+      {
+          'x': np.array([-1.0, 0.0, 0.25, 0.5, 0.75, 1.0, 2.0]),
+          'xp': np.array([0.0, 1.0]),
+          'fp': np.array([0.2, 0.4]),
+          'y': np.array([0.2, 0.2, 0.25, 0.3, 0.35, 0.4, 0.4]),
+      },
+      {
+          'x': np.array([-1.0, 0.0, 0.5, 1.0, 2.0]),
+          'xp': np.array([0.0, 1.0]),
+          'fp': np.array([[0.2, 0.4], [0.4, 0.2]]),
+          'y': np.array(
+              [[0.2, 0.4], [0.2, 0.4], [0.3, 0.3], [0.4, 0.2], [0.4, 0.2]]),
+      },
+  )
+  def test_linear_interpolation(self, x, xp, fp, y):
+    self.assertAllClose(multi_problem_v2.linear_interpolation(x, xp, fp), y)
+
+  @parameterized.parameters(
+      {
+          'x': np.array([-1.0, 0.0, 0.25, 0.5, 0.75, 1.0, 2.0]),
+          'xp': np.array([0.0, 0.6, 0.9]),
+          'fp': np.array([0.1, 0.9, 0.6]),
+          'y': np.array([0.1, 0.1, 0.1, 0.1, 0.9, 0.6, 0.6]),
+      },
+      {
+          'x': np.array([-1.0, 0.0, 0.5, 1.0, 2.0]),
+          'xp': np.array([0.0, 0.6, 0.9]),
+          'fp': np.array([[0.1, 0.4], [0.9, 0.2], [0.6, 0.9]]),
+          'y': np.array(
+              [[0.1, 0.4], [0.1, 0.4], [0.1, 0.4], [0.6, 0.9], [0.6, 0.9]]),
+      },
+  )
+  def test_step_interpolation(self, x, xp, fp, y):
+    self.assertAllClose(multi_problem_v2.step_interpolation(x, xp, fp), y)
+
+  @parameterized.parameters(
+      {
+          'schedule': ('linear', (100, 200), ((0.25, 0.75), (0.62, 0.38))),
+          'steps': np.array([50, 100, 150, 200, 250]),
+          'pmfs': np.array(
+              [[0.25, 0.75], [0.25, 0.75], [0.435, 0.565], [0.62, 0.38],
+               [0.62, 0.38]]),
+      },
+      {
+          'schedule': ('step', (100, 200), ((0.25, 0.75), (0.62, 0.38))),
+          'steps': np.array([50, 100, 150, 200, 250]),
+          'pmfs': np.array(
+              [[0.25, 0.75], [0.25, 0.75], [0.25, 0.75], [0.62, 0.38],
+               [0.62, 0.38]]),
+      },
+  )
+  def test_get_schedule_distribution(self, schedule, steps, pmfs):
+    with self.test_session() as sess:
+      global_step = tf.train.get_or_create_global_step()
+      output = multi_problem_v2.get_schedule_distribution(schedule, global_step)
+      sess.run(global_step.initializer)
+      for step, pmf in zip(steps, pmfs):
+        sess.run(global_step.assign(step))
+        self.assertAllClose(sess.run(output), pmf)
+
+  @parameterized.parameters(
+      {
+          'pmf': np.array([1.0, 0.0], np.float32),
+          'fns': [lambda: 0, lambda: 1],
+          'rands': np.array([0.1, 0.4, 0.6, 0.9], np.float32),
+          'targets': np.array([0, 0, 0, 0], np.float32),
+      },
+      {
+          'pmf': np.array([0.2, 0.6, 0.2], np.float32),
+          'fns': [lambda: 0, lambda: 1, lambda: 2],
+          'rands': np.array([0.1, 0.4, 0.6, 0.9], np.float32),
+          'targets': np.array([0, 1, 1, 2], np.float32),
+      },
+  )
+  def test_categorical_case(self, pmf, fns, rands, targets):
+    with self.test_session() as sess:
+      for rand, target in zip(rands, targets):
+        output = multi_problem_v2.categorical_case(pmf, fns, rand)
+        self.assertEqual(sess.run(output), target)
+
+  @parameterized.parameters(
+      {
+          'pmf': np.array([1.0, 0.0], np.float32),
+          'num_datasets': 2,
+          'sample_size': 10,
+      },
+      {
+          'pmf': np.array([0.3, 0.7], np.float32),
+          'num_datasets': 2,
+          'sample_size': 400,
+      },
+      {
+          'pmf': None,
+          'num_datasets': 2,
+          'sample_size': 400,
+      },
+  )
+  def test_get_multi_dataset(self, pmf, num_datasets, sample_size):
+    with self.test_session() as sess:
+      datasets = [tf.data.Dataset.from_tensors(i) for i in range(num_datasets)]
+      multi_dataset = multi_problem_v2.get_multi_dataset(datasets, pmf)
+      multi_dataset = multi_dataset.batch(sample_size)
+      iterator = multi_dataset.make_initializable_iterator()
+      sess.run(iterator.initializer)
+      sample_pmf = tf.reduce_mean(
+          tf.one_hot(iterator.get_next(), num_datasets), 0)
+      if pmf is None:
+        pmf = np.array([1.0 / num_datasets] * num_datasets, np.float32)
+      self.assertAllClose(sess.run(sample_pmf), pmf, rtol=0.1, atol=0.1)
+
+  @parameterized.parameters(
+      {
+          'schedule': ('step', (100, 200), ((1.0, 0.0), (0.0, 1.0))),
+          'num_datasets': 2,
+          'sample_size': 20,
+      },
+      {
+          'schedule': ('linear', (100, 200), ((0.6, 0.4), (0.1, 0.9))),
+          'num_datasets': 2,
+          'sample_size': 400,
+      },
+  )
+  def test_multi_problem_v2(self, schedule, num_datasets, sample_size):
+
+    class DummyProblem(problem.Problem):
+
+      def dataset(self, *args, **kwargs):
+        return tf.data.Dataset.from_tensors({'targets': 0.0})
+
+    with self.test_session() as sess:
+      for mode in [problem.DatasetSplit.TRAIN, problem.DatasetSplit.EVAL]:
+        p = multi_problem_v2.MultiProblemV2(
+            [DummyProblem() for _ in range(num_datasets)], schedule)
+        global_step = tf.train.get_or_create_global_step()
+        dataset = p.dataset(mode, global_step).batch(sample_size)
+        iterator = dataset.make_initializable_iterator()
+        features = iterator.get_next()
+        sess.run(global_step.initializer)
+        sess.run(iterator.initializer)
+        sess.run(features)
+
+
+if __name__ == '__main__':
+  tf.test.main()
diff --git a/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index faef16eea..b86cb1d8e 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -21,6 +21,7 @@
 
 from tensor2tensor.data_generators import cnn_dailymail
 from tensor2tensor.data_generators import multi_problem
+from tensor2tensor.data_generators import multi_problem_v2
 from tensor2tensor.data_generators import multinli
 from tensor2tensor.data_generators import squad
 from tensor2tensor.data_generators import text_problems
@@ -46,6 +47,33 @@ def vocab_type(self):
     return text_problems.VocabType.SUBWORD
 
 
+@registry.register_problem
+class LanguagemodelEnWikiLMMultiNLISubwordsV2(
+    multi_problem_v2.MultiText2TextProblem):
+  """Wiki LM and MNLI mixed problem class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    problems = [
+        wiki_lm.LanguagemodelEnWiki32k(),
+        multinli.MultiNLIWikiLMSharedVocab(),
+    ]
+    schedule = multi_problem_v2.constant_schedule([0.5, 0.5])
+    super(LanguagemodelEnWikiLMMultiNLISubwordsV2, self).__init__(
+        problems, schedule, was_reversed=was_reversed, was_copy=was_copy)
+
+  @property
+  def has_inputs(self):
+    return False
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelEnWiki32k().vocab_filename
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
+
+
 @registry.register_problem
 class LanguagemodelEnWikiLMMultiNLISubwords64k(multi_problem.MultiProblem):
   """Wiki LM and MNLI mixed problem class."""

From e413606376595c67501a5a13716b77f0a644d5b0 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 5 Feb 2019 10:50:34 -0800
Subject: [PATCH 1631/2720] Evolved Transformer model and hparams.

PiperOrigin-RevId: 232516655
---
 tensor2tensor/models/__init__.py              |   1 +
 tensor2tensor/models/evolved_transformer.py   | 107 ++++++++++++++++++
 .../models/evolved_transformer_test.py        |  68 +++++++++++
 tensor2tensor/models/transformer.py           |   6 +-
 tensor2tensor/utils/learning_rate.py          |   7 ++
 5 files changed, 187 insertions(+), 2 deletions(-)
 create mode 100644 tensor2tensor/models/evolved_transformer.py
 create mode 100644 tensor2tensor/models/evolved_transformer_test.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 973c87a81..ca3f7f3f2 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -26,6 +26,7 @@
 from tensor2tensor.models import basic
 from tensor2tensor.models import bytenet
 from tensor2tensor.models import distillation
+from tensor2tensor.models import evolved_transformer
 from tensor2tensor.models import image_transformer
 from tensor2tensor.models import image_transformer_2d
 from tensor2tensor.models import lstm
diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
new file mode 100644
index 000000000..0395fbcbd
--- /dev/null
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -0,0 +1,107 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Evolved Transformer model.
+
+This implements the model described in arxiv.org/abs/1901.11117 .
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import transformer_layers
+from tensor2tensor.models import transformer
+from tensor2tensor.utils import registry
+
+
+@registry.register_model
+class EvolvedTransformer(transformer.Transformer):
+  """The Evolved Transformer from arxiv.org/abs/1901.11117 ."""
+
+  def __init__(self, *args, **kwargs):
+    super(EvolvedTransformer, self).__init__(*args, **kwargs)
+    self._encoder_function = transformer_layers.evolved_transformer_encoder
+    self._decoder_function = transformer.evolved_transformer_decoder
+
+  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha,
+                   use_tpu):
+    """Forced slow beam decode because cache is not supported.
+
+    Args:
+      features: an map of string to `Tensor`.
+      decode_length: an integer.  How many additional timesteps to decode.
+      beam_size: number of beams.
+      top_beams: an integer. How many of the beams to return.
+      alpha: Float that controls the length penalty. larger the alpha, stronger
+        the preference for longer translations.
+      use_tpu: Whether or not TPU is being used.
+
+    Returns:
+      A dict of decoding results {
+          "outputs": integer `Tensor` of decoded ids of shape
+              [batch_size, <= decode_length] if beam_size == 1 or
+              [batch_size, top_beams, <= decode_length].
+          "scores": decoding log probs from the beam search,
+              None if using greedy decoding (beam_size=1).
+      }
+    """
+    return self._beam_decode_slow(features, decode_length, beam_size, top_beams,
+                                  alpha, use_tpu)
+
+
+# TODO(davidso): Update optimizer, learning rate, and decay to match paper.
+def add_evolved_transformer_hparams(hparams):
+  """Add Evolved Transformer hparams.
+
+  Note: These are for the Adam optimizer, not the Adafactor optimizer used in
+  the paper.
+
+  Args:
+    hparams: Current hparams.
+
+  Returns:
+    hparams updated with Evolved Transformer values.
+  """
+  # Evolved Transformer "layers" are twice as deep as Transformer, so roughly
+  # halve the number that we use. These numbers are taken from
+  # arxiv.org/abs/1901.11117 .
+  hparams.num_encoder_layers = 3
+  hparams.num_decoder_layers = 4
+
+  # Learning rate and decay scheme that mimics the transformer Adam config,
+  # but with cosine decay instead of rsqrt.
+  hparams.learning_rate_constant /= hparams.learning_rate_warmup_steps ** 0.5
+  hparams.learning_rate_schedule = (
+      "constant*linear_warmup*single_cycle_cos_decay*rsqrt_hidden_size")
+  # The current infrastructure does not support exposing
+  # `train_steps` to the decay functions, and so we are hard coding the decay
+  # steps here to match the default number of train steps used in `t2t_trainer`.
+  # TODO(davidso): Thread `train_steps` through to decay functions so we do not
+  # have to worry about a `learning_rate_decay_steps` mismatch.
+  hparams.learning_rate_decay_steps = 250000
+  return hparams
+
+
+@registry.register_hparams
+def evolved_transformer_base():
+  """Base parameters for Evolved Transformer model."""
+  return add_evolved_transformer_hparams(transformer.transformer_base())
+
+
+@registry.register_hparams
+def evolved_transformer_big():
+  """Big parameters for Evolved Transformer model on WMT."""
+  return add_evolved_transformer_hparams(transformer.transformer_big())
diff --git a/tensor2tensor/models/evolved_transformer_test.py b/tensor2tensor/models/evolved_transformer_test.py
new file mode 100644
index 000000000..0056ed2f0
--- /dev/null
+++ b/tensor2tensor/models/evolved_transformer_test.py
@@ -0,0 +1,68 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the Evolved Transformer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import numpy as np
+
+from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.models import evolved_transformer
+from tensor2tensor.models import transformer
+
+import tensorflow as tf
+
+BATCH_SIZE = 3
+INPUT_LENGTH = 5
+TARGET_LENGTH = 7
+VOCAB_SIZE = 10
+
+
+def get_model():
+  hparams = transformer.transformer_tiny()
+  hparams.layer_prepostprocess_dropout = 0.0
+
+  p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE, VOCAB_SIZE,
+                                                   hparams)
+  hparams.problem_hparams = p_hparams
+
+  inputs = np.random.randint(VOCAB_SIZE, size=(BATCH_SIZE, INPUT_LENGTH, 1, 1))
+  targets = np.random.randint(
+      VOCAB_SIZE, size=(BATCH_SIZE, TARGET_LENGTH, 1, 1))
+  features = {
+      "targets": tf.constant(targets, dtype=tf.int32, name="targets"),
+      "target_space_id": tf.constant(1, dtype=tf.int32),
+      "inputs": tf.constant(inputs, dtype=tf.int32, name="inputs"),
+  }
+
+  return (evolved_transformer.EvolvedTransformer(
+      hparams, tf.estimator.ModeKeys.TRAIN, p_hparams), features)
+
+
+class EvolvedTransformerTest(tf.test.TestCase):
+
+  def testEvolvedTransformer(self):
+    model, features = get_model()
+    logits, _ = model(features)
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      res = session.run(logits)
+    self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, 1, 1, VOCAB_SIZE))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 42f9fd7dd..29005b2a2 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -59,6 +59,8 @@ class Transformer(t2t_model.T2TModel):
   def __init__(self, *args, **kwargs):
     super(Transformer, self).__init__(*args, **kwargs)
     self.attention_weights = {}  # For visualizing attention heads.
+    self._encoder_function = transformer_encoder
+    self._decoder_function = transformer_decoder
 
   def encode(self, inputs, target_space, hparams, features=None, losses=None):
     """Encode transformer inputs.
@@ -98,7 +100,7 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
     if hparams.unidirectional_encoder:
       attn_bias_for_padding = encoder_decoder_attention_bias
 
-    encoder_output = transformer_encoder(
+    encoder_output = self._encoder_function(
         encoder_input,
         self_attention_bias,
         hparams,
@@ -149,7 +151,7 @@ def decode(self,
     decoder_input = tf.nn.dropout(decoder_input,
                                   1.0 - hparams.layer_prepostprocess_dropout)
 
-    decoder_output = transformer_decoder(
+    decoder_output = self._decoder_function(
         decoder_input,
         encoder_output,
         decoder_self_attention_bias,
diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index 709f993e7..b78402760 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -40,6 +40,13 @@ def learning_rate_factor(name, step_num, hparams):
         np.pi * step_num / hparams.learning_rate_decay_steps))
     # if in warmup stage return 1 else return the decayed value
     return in_warmup * 1 + (1 - in_warmup) * ret
+  elif name == "single_cycle_cos_decay":
+    # Cosine decay to zero with a single cycle. This is different from
+    # "cosdecay" because it starts at 1 when the warmup steps end.
+    x = tf.maximum(step_num, hparams.learning_rate_warmup_steps)
+    step = x - hparams.learning_rate_warmup_steps
+    return tf.math.cos(
+        step * np.pi / hparams.learning_rate_decay_steps) / 2.0 + 0.5
   elif name == "rsqrt_decay":
     return tf.rsqrt(tf.maximum(step_num, hparams.learning_rate_warmup_steps))
   elif name == "rsqrt_normalized_decay":

From 28adf2690c551ef0f570d41bef2019d9c502ec7e Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 5 Feb 2019 20:30:45 -0800
Subject: [PATCH 1632/2720] Starting a simplfied TFv2 trainer and the most
 basic model.

PiperOrigin-RevId: 232608829
---
 tensor2tensor/bin/t2t_trainer.py  |  22 ++++
 tensor2tensor/v2/keras_utils.py   |  39 ++++++
 tensor2tensor/v2/models/basic.py  |  97 +++++++++++++++
 tensor2tensor/v2/models/resnet.py |  64 ++++++++++
 tensor2tensor/v2/t2t.py           | 192 ++++++++++++++++++++++++++++++
 tensor2tensor/v2/t2t_trainer.py   |  64 ++++++++++
 6 files changed, 478 insertions(+)
 create mode 100644 tensor2tensor/v2/keras_utils.py
 create mode 100644 tensor2tensor/v2/models/basic.py
 create mode 100644 tensor2tensor/v2/models/resnet.py
 create mode 100644 tensor2tensor/v2/t2t.py
 create mode 100644 tensor2tensor/v2/t2t_trainer.py

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 0dad6b4ee..43432a082 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -32,6 +32,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
+from tensor2tensor.v2 import t2t as t2t_v2
 import tensorflow as tf
 
 from tensorflow.contrib.tpu.python.tpu import tpu_config
@@ -69,6 +70,7 @@
 flags.DEFINE_integer("intra_op_parallelism_threads", 0,
                      "Number of intra_op_parallelism_threads to use for CPU. "
                      "See TensorFlow config.proto for details.")
+flags.DEFINE_bool("v2", False, "Whether to use T2T v2.")
 # TODO(lukaszkaiser): resolve memory and variable assign issues and set to True.
 flags.DEFINE_bool(
     "optionally_use_dist_strat", False,
@@ -357,6 +359,26 @@ def run_std_server():
 def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
 
+  if FLAGS.v2:
+    tf.enable_v2_behavior()
+    # Hacking main v1 flags to work with v2.
+    config_strs = []
+    config_strs.append(
+        "train_fn.train_steps=" + str(FLAGS.train_steps))
+    config_strs.append(
+        "train_fn.eval_steps=" + str(FLAGS.eval_steps))
+    config_strs.append(
+        "train_fn.eval_frequency=" + str(FLAGS.local_eval_frequency))
+    if FLAGS.hparams:
+      config_strs.extend(str(FLAGS.hparams).split(","))
+    config_str = "\n".join(config_strs)
+    data_dir = os.path.expanduser(FLAGS.data_dir)
+    output_dir = os.path.expanduser(FLAGS.output_dir)
+    t2t_v2.t2t_train(FLAGS.model, FLAGS.problem,
+                     data_dir=data_dir, output_dir=output_dir,
+                     config_file=FLAGS.hparams_set, config=config_str)
+    return
+
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
 
   # If we just have to print the registry, do that and exit early.
diff --git a/tensor2tensor/v2/keras_utils.py b/tensor2tensor/v2/keras_utils.py
new file mode 100644
index 000000000..a9eeebb42
--- /dev/null
+++ b/tensor2tensor/v2/keras_utils.py
@@ -0,0 +1,39 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Utilities to use TF v1 layers with Keras and TF v2 easily."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+
+class FunctionLayer(tf.compat.v2.keras.layers.Layer):
+  """Layer made of a function. Stores all variables."""
+
+  def __init__(self, function, name=None):
+    if name is None:
+      name = function.__name__
+    super(FunctionLayer, self).__init__(name=name)
+    self._template = tf.compat.v1.make_template(name, function)
+
+  @property
+  def losses(self):
+    return []
+
+  def call(self, *args, **kwargs):
+    return self._template(*args, **kwargs)
diff --git a/tensor2tensor/v2/models/basic.py b/tensor2tensor/v2/models/basic.py
new file mode 100644
index 000000000..4d2cbe054
--- /dev/null
+++ b/tensor2tensor/v2/models/basic.py
@@ -0,0 +1,97 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Basic models for testing simple tasks."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+import gin.tf
+
+
+@gin.configurable(whitelist=["num_hidden_layers", "hidden_size", "dropout"])
+class BasicFcRelu(tf.keras.Model):
+  """Basic fully-connected + ReLU model."""
+
+  def __init__(self, features_info=None, supervised_keys=None,
+               num_hidden_layers=2, hidden_size=64, dropout=0.1):
+    super(BasicFcRelu, self).__init__()
+    assert features_info is not None
+    assert supervised_keys is not None
+    self._input_key = supervised_keys[0]
+    label_key = supervised_keys[1]
+    input_shape = features_info[self._input_key].shape
+    num_output_classes = features_info[label_key].num_classes
+    self._num_hidden_layers = num_hidden_layers
+    self._dense_layers = []
+    self._dropout_layers = []
+
+    # Now the model.
+    self._flatten_layer = tf.keras.layers.Flatten(input_shape=input_shape)
+    for i in range(num_hidden_layers):
+      self._dense_layers.append(tf.keras.layers.Dense(
+          hidden_size, activation="relu", name="layer_%d" % i))
+      self._dropout_layers.append(tf.keras.layers.Dropout(
+          rate=dropout))
+    self._logits = tf.keras.layers.Dense(
+        num_output_classes, activation="softmax")
+
+  def call(self, inputs, training=False):
+    x = tf.cast(inputs[self._input_key], tf.float32) / 255.0
+    x = self._flatten_layer(x)
+    for i in range(self._num_hidden_layers):
+      x = self._dense_layers[i](x)
+      x = self._dropout_layers[i](x, training=training)
+    return self._logits(x)
+
+
+def basic_fc_large():
+  """Large set of parameters for this model."""
+  gin.bind_parameter("BasicFcRelu.num_hidden_layers", 3)
+  gin.bind_parameter("BasicFcRelu.hidden_size", 128)
+  gin.bind_parameter("BasicFcRelu.dropout", 0.3)
+  return BasicFcRelu
+
+
+# TODO(lukaszkaiser): could we allow coding like this? it's much easier!
+# This will run fine, but not train as new layers are made in each step!
+@gin.configurable(whitelist=["num_hidden_layers", "hidden_size", "dropout"])
+class BasicFcReluV2(tf.keras.Model):
+  """Basic fully-connected + ReLU model, nicer code version."""
+
+  def __init__(self, features_info=None, supervised_keys=None,
+               num_hidden_layers=2, hidden_size=64, dropout=0.1):
+    super(BasicFcReluV2, self).__init__()
+    assert features_info is not None
+    assert supervised_keys is not None
+    self._input_key = supervised_keys[0]
+    self._input_shape = features_info[self._input_key].shape
+    self._num_output_classes = features_info[supervised_keys[1]].num_classes
+    self._num_hidden_layers = num_hidden_layers
+    self._dropout = dropout
+    self._hidden_size = hidden_size
+
+  def call(self, inputs, training=False):
+    x = tf.cast(inputs[self._input_key], tf.float32) / 255.0
+    x = tf.keras.layers.Flatten(
+        input_shape=self._input_shape)(x)
+    for i in range(self._num_hidden_layers):
+      x = tf.keras.layers.Dense(
+          self._hidden_size, activation="relu", name="layer_%d" % i)(x)
+      x = tf.keras.layers.Dropout(rate=self._dropout)(x, training=training)
+    return tf.keras.layers.Dense(
+        self._num_output_classes, activation="softmax")(x)
diff --git a/tensor2tensor/v2/models/resnet.py b/tensor2tensor/v2/models/resnet.py
new file mode 100644
index 000000000..d4f6b3d89
--- /dev/null
+++ b/tensor2tensor/v2/models/resnet.py
@@ -0,0 +1,64 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Basic models for testing simple tasks."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.models import resnet
+from tensor2tensor.v2 import keras_utils
+import tensorflow as tf
+import gin.tf
+
+
+@gin.configurable(whitelist=["layer_sizes", "filter_sizes"])
+class Resnet(tf.keras.Model):
+  """Resnet."""
+
+  def __init__(self, features_info=None, supervised_keys=None,
+               layer_sizes=None, filter_sizes=None):
+    # Base config for resnet-50.
+    if layer_sizes is None:
+      layer_sizes = [3, 4, 6, 3]
+    if filter_sizes is None:
+      filter_sizes = [64, 64, 128, 256, 512]
+    assert features_info is not None
+    assert supervised_keys is not None
+    super(Resnet, self).__init__()
+    self._input_key = supervised_keys[0]
+    label_key = supervised_keys[1]
+    num_output_classes = features_info[label_key].num_classes
+
+    # Now the model.
+    def resnet_model(inputs, training):
+      return resnet.resnet_v2(
+          inputs,
+          resnet.bottleneck_block,
+          layer_sizes,
+          filter_sizes,
+          is_training=training,
+          is_cifar=True)
+
+    self._resnet = keras_utils.FunctionLayer(resnet_model)
+    self._logits = tf.keras.layers.Dense(
+        num_output_classes, activation="softmax")
+
+  def call(self, inputs, training=False):
+    x = tf.cast(inputs[self._input_key], tf.float32) / 255.0
+    x = self._resnet(x, training)
+    x = tf.reduce_mean(x, axis=[1, 2])
+    return self._logits(x)
diff --git a/tensor2tensor/v2/t2t.py b/tensor2tensor/v2/t2t.py
new file mode 100644
index 000000000..422527eb3
--- /dev/null
+++ b/tensor2tensor/v2/t2t.py
@@ -0,0 +1,192 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""T2T models, configs and main training functions."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+from tensor2tensor import problems
+from tensor2tensor.v2.models import basic
+from tensor2tensor.v2.models import resnet
+
+import tensorflow as tf
+import tensorflow_datasets as tfds
+
+import gin.tf
+
+
+# Since there are few models and configs for now, we use this simple registry.
+# TODO(lukaszkaiser): find a better way to do this or remove altogether.
+_MODEL_REGISTRY = {
+    "basic_fc_relu": lambda: basic.BasicFcRelu,
+    "basic_fc_large": basic.basic_fc_large,
+    "basic_fc_relu_v2": lambda: basic.BasicFcReluV2,
+    "resnet": lambda: resnet.Resnet,
+}
+
+
+def train_and_eval_dataset(dataset_name, data_dir):
+  """Return train and evaluation datasets, feature info and supervised keys.
+
+  Args:
+    dataset_name: a string, the name of the dataset; if it starts with "v1_"
+      then we'll search T2T Problem registry for it, otherwise we assume it
+      is a dataset from TFDS and load it from there.
+    data_dir: directory where the data is located.
+
+  Returns:
+    a 4-tuple consisting of:
+     * the train tf.Daataset
+     * the eval tf.Daataset
+     * information about features: a python dictionary with feature names
+         as keys and an object as value that provides .shape and .num_classes.
+     * supervised_keys: information what's the input and what's the target.
+  """
+  if dataset_name.startswith("v1_"):
+    return _train_and_eval_dataset_v1(dataset_name[3:], data_dir)
+  dataset_builder = tfds.builder(dataset_name, data_dir=data_dir)
+  info = dataset_builder.info
+  splits = dataset_builder.info.splits
+  if tfds.Split.TRAIN not in splits:
+    raise ValueError("To train we require a train split in the dataset.")
+  if tfds.Split.VALIDATION not in splits and "test" not in splits:
+    raise ValueError("We require a validation or test split in the dataset.")
+  eval_split = tfds.Split.VALIDATION
+  if tfds.Split.VALIDATION not in splits:
+    eval_split = tfds.Split.TEST
+  train, valid = tfds.load(
+      name=dataset_name, split=[tfds.Split.TRAIN, eval_split])
+  return train, valid, info.features, info.supervised_keys
+
+
+def _make_info(shape_list, num_classes):
+  """Create an info-like tuple for feature given some shapes and vocab size."""
+  feature_info = collections.namedtuple("FeatureInfo", ["shape", "num_classes"])
+  cur_shape = list(shape_list[0])
+  # We need to merge the provided shapes, put None where they disagree.
+  for shape in shape_list:
+    if len(shape) != len(cur_shape):
+      raise ValueError("Shapes need to have the same number of dimensions.")
+    for i in range(len(shape)):
+      if cur_shape[i] is not None:
+        if shape[i] != cur_shape[i]:
+          cur_shape[i] = None
+  return feature_info(cur_shape, num_classes)
+
+
+def _train_and_eval_dataset_v1(problem_name, data_dir):
+  """Return train and evaluation datasets, feature info and supervised keys."""
+  problem = problems.problem(problem_name)
+  train_dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, data_dir)
+  eval_dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, data_dir)
+  supervised_keys = ("inputs", "targets")
+  hparams = problem.get_hparams()
+  # We take a few training examples to guess the shapes.
+  input_shapes, target_shapes = [], []
+  for example in train_dataset.take(3):
+    input_shapes.append(example["inputs"].shape.as_list())
+    target_shapes.append(example["targets"].shape.as_list())
+  input_info = _make_info(
+      input_shapes, hparams.modality["inputs"].top_dimensionality)
+  target_info = _make_info(
+      target_shapes, hparams.modality["targets"].top_dimensionality)
+  info = {"inputs": input_info, "targets": target_info}
+  return train_dataset, eval_dataset, info, supervised_keys
+
+
+def shuffle_and_batch_data(dataset, batch_size, target_key):
+  """Shuffle and batch the given dataset."""
+  shuffled = dataset.shuffle(128).batch(batch_size).prefetch(8)
+  return shuffled.map(lambda ex: (ex, ex[target_key]))
+
+
+@gin.configurable(blacklist=["model"])
+def model_compile(model,
+                  optimizer="adam",
+                  loss="sparse_categorical_crossentropy",
+                  metrics=None):
+  """Compile the model in Keras."""
+  metrics = ["accuracy"] if metrics is None else metrics
+  return model.compile(optimizer=optimizer,
+                       loss=loss,
+                       metrics=metrics)
+
+
+# We include in gin config everything that could be useful to share between
+# users, so when it gets saved in a .gin file it can be re-ran with few flags.
+@gin.configurable(blacklist=["data_dir", "output_dir"])
+def train_fn(data_dir=None, output_dir=None,
+             model_class=gin.REQUIRED, dataset=gin.REQUIRED,
+             batch_size=32, train_steps=1000, eval_steps=1, eval_frequency=100):
+  """Train the given model on the given dataset.
+
+  Args:
+    data_dir: Directory where the data is located.
+    output_dir: Directory where to put the logs and checkpoints.
+    model_class: The model class to train.
+    dataset: The name of the dataset to train on.
+    batch_size: integer, how many examples per batch.
+    train_steps: for how many steps to train.
+    eval_steps: for how many steps to do evaluation.
+    eval_frequency: how often (every this many steps) to run evaluation.
+  """
+  train_data, eval_data, features_info, keys = train_and_eval_dataset(
+      dataset, data_dir)
+  model = model_class(features_info=features_info, supervised_keys=keys)
+  model_compile(model)
+  train_batches = shuffle_and_batch_data(train_data, batch_size, keys[1])
+  eval_batches = shuffle_and_batch_data(eval_data, batch_size, keys[1])
+
+  # Training loop.
+  callbacks = []
+  callbacks.append(tf.keras.callbacks.History())
+  callbacks.append(tf.keras.callbacks.BaseLogger())
+  if output_dir is not None:
+    callbacks.append(tf.keras.callbacks.TensorBoard(log_dir=output_dir))
+    # TODO(lukaszkaiser): the one below doesn't seem to work, why?
+    # callbacks.append(tf.keras.callbacks.ModelCheckpoint(
+    #     filepath=output_dir))
+  model.fit(train_batches.repeat(),
+            epochs=train_steps // eval_frequency,
+            steps_per_epoch=eval_frequency,
+            validation_data=eval_batches,
+            validation_steps=eval_steps,
+            callbacks=callbacks)
+
+
+def t2t_train(model_name, dataset_name,
+              data_dir=None, output_dir=None, config_file=None, config=None):
+  """Main function to train the given model on the given dataset.
+
+  Args:
+    model_name: The name of the model to train.
+    dataset_name: The name of the dataset to train on.
+    data_dir: Directory where the data is located.
+    output_dir: Directory where to put the logs and checkpoints.
+    config_file: the gin configuration file to use.
+    config: string (in gin format) to override gin parameters.
+  """
+  if model_name not in _MODEL_REGISTRY:
+    raise ValueError("Model %s not in registry. Available models:\n * %s." %
+                     (model_name, "\n * ".join(_MODEL_REGISTRY.keys())))
+  model_class = _MODEL_REGISTRY[model_name]()
+  gin.bind_parameter("train_fn.model_class", model_class)
+  gin.bind_parameter("train_fn.dataset", dataset_name)
+  gin.parse_config_files_and_bindings(config_file, config)
+  # TODO(lukaszkaiser): save gin config in output_dir if provided?
+  train_fn(data_dir, output_dir=output_dir)
diff --git a/tensor2tensor/v2/t2t_trainer.py b/tensor2tensor/v2/t2t_trainer.py
new file mode 100644
index 000000000..aa9abe8f1
--- /dev/null
+++ b/tensor2tensor/v2/t2t_trainer.py
@@ -0,0 +1,64 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# 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.
+
+r"""T2T trainer for TF 2.0.
+
+This trainer only supports a subset of models and features for now.
+
+Examples:
+
+- train a basic model on mnist:
+    v2/t2t_trainer.py --dataset=mnist --model=basic_fc_relu
+      --config="train_fn.train_steps=4000"
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+from absl import app
+from absl import flags
+from tensor2tensor.v2 import t2t
+import tensorflow as tf
+
+tf.enable_v2_behavior()
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("dataset", None, "Which dataset to use.")
+flags.DEFINE_string("model", None, "Which model to train.")
+flags.DEFINE_string("data_dir", None, "Path to the directory with data.")
+flags.DEFINE_string("output_dir", None,
+                    "Path to the directory to save logs and checkpoints.")
+flags.DEFINE_multi_string("config_file", None,
+                          "Configuration file with parameters (.gin).")
+flags.DEFINE_multi_string("config", None,
+                          "Configuration parameters (gin string).")
+
+
+def main(argv):
+  del argv
+  data_dir, output_dir = FLAGS.data_dir, FLAGS.output_dir
+  if data_dir is not None:
+    data_dir = os.path.expanduser(data_dir)
+  if output_dir is not None:
+    output_dir = os.path.expanduser(output_dir)
+  t2t.t2t_train(FLAGS.model, FLAGS.dataset,
+                data_dir=data_dir, output_dir=output_dir,
+                config_file=FLAGS.config_file, config=FLAGS.config)
+
+
+if __name__ == "__main__":
+  app.run(main)

From 2ea1518eeac80fd365ee7bf0c8c76e2e008b001c Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 6 Feb 2019 09:20:13 -0800
Subject: [PATCH 1633/2720] Remove
 Modality.{bottom,targets_bottom,top,loss}_sharded.

This is an incremental change toward replacing Modality classes with a simpler ModalityType enum.

PiperOrigin-RevId: 232687333
---
 tensor2tensor/layers/modalities_test.py | 26 +++++++++-----
 tensor2tensor/models/transformer.py     | 12 +++----
 tensor2tensor/utils/modality.py         | 48 -------------------------
 tensor2tensor/utils/t2t_model.py        | 16 +++++----
 4 files changed, 33 insertions(+), 69 deletions(-)

diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 7ae536207..d14de2927 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -47,7 +47,7 @@ def testSymbolModalityInputs(self):
     data_parallelism = expert_utils.Parallelism(
         ["/device:CPU:0"] * num_datashards)
     xs = tf.split(x, num_datashards)
-    sharded_output = m.bottom_sharded(xs, data_parallelism)
+    sharded_output = data_parallelism(m.bottom, xs)
     output = tf.concat(sharded_output, 0)
     self.evaluate(tf.global_variables_initializer())
     res = self.evaluate(output)
@@ -73,10 +73,14 @@ def testSymbolModalityTargets(self):
         ["/device:CPU:0"] * num_datashards)
     sharded_body_output = tf.split(tf.to_float(body_output), num_datashards)
     sharded_targets = tf.split(targets, num_datashards)
-    sharded_logits = m.top_sharded(sharded_body_output, sharded_targets,
-                                   data_parallelism)
-    train_loss = m.loss_sharded(sharded_logits, sharded_targets,
-                                data_parallelism)
+    sharded_logits = data_parallelism(m.top,
+                                      sharded_body_output,
+                                      sharded_targets)
+    sharded_loss_num, sharded_loss_den = data_parallelism(m.loss,
+                                                          sharded_logits,
+                                                          sharded_targets)
+    train_loss = (tf.add_n(sharded_loss_num) /
+                  tf.maximum(1.0, tf.add_n(sharded_loss_den)))
     logits = tf.concat(sharded_logits, 0)
     self.evaluate(tf.global_variables_initializer())
     res1, res2 = self.evaluate((logits, train_loss))
@@ -105,10 +109,14 @@ def testSymbolModalityTargetsFactored(self):
     with self.test_session() as session:
       sharded_body_output = tf.split(tf.to_float(body_output), num_datashards)
       sharded_targets = tf.split(targets, num_datashards)
-      sharded_logits = m.top_sharded(sharded_body_output, sharded_targets,
-                                     data_parallelism)
-      train_loss = m.loss_sharded(sharded_logits, sharded_targets,
-                                  data_parallelism)
+      sharded_logits = data_parallelism(m.top,
+                                        sharded_body_output,
+                                        sharded_targets)
+      sharded_loss_num, sharded_loss_den = data_parallelism(m.loss,
+                                                            sharded_logits,
+                                                            sharded_targets)
+      train_loss = (tf.add_n(sharded_loss_num) /
+                    tf.maximum(1.0, tf.add_n(sharded_loss_den)))
       logits = tf.concat(sharded_logits, 0)
       session.run(tf.global_variables_initializer())
       res1, res2 = session.run((logits, train_loss))
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 29005b2a2..5429208d1 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -360,7 +360,7 @@ def _fast_decode_tpu(self,
       inputs = self._shard_features({"inputs": inputs})["inputs"]
       input_modality = self._problem_hparams.modality["inputs"]
       with tf.variable_scope(input_modality.name):
-        inputs = input_modality.bottom_sharded(inputs, dp)
+        inputs = dp(input_modality.bottom, inputs)
       with tf.variable_scope("body"):
         encoder_output, encoder_decoder_attention_bias = dp(
             self.encode,
@@ -419,7 +419,7 @@ def preprocess_targets(targets, i):
       # _shard_features called to ensure that the variable names match
       targets = self._shard_features({"targets": targets})["targets"]
       with tf.variable_scope(target_modality.name):
-        targets = target_modality.targets_bottom_sharded(targets, dp)[0]
+        targets = dp(target_modality.targets_bottom, targets)[0]
       targets = common_layers.flatten4d3d(targets)
 
       # TODO(llion): Explain! Is this even needed?
@@ -475,7 +475,7 @@ def symbols_to_logits_tpu_fn(ids, i, cache):
             nonpadding=features_to_nonpadding(features, "targets"))
 
       with tf.variable_scope(target_modality.name):
-        logits = target_modality.top_sharded(body_outputs, None, dp)[0]
+        logits = dp(target_modality.top, body_outputs, None)[0]
 
       ret = tf.squeeze(logits, axis=[1, 2, 3])
       if partial_targets is not None:
@@ -577,7 +577,7 @@ def _fast_decode(self,
       inputs = self._shard_features({"inputs": inputs})["inputs"]
       input_modality = self._problem_hparams.modality["inputs"]
       with tf.variable_scope(input_modality.name):
-        inputs = input_modality.bottom_sharded(inputs, dp)
+        inputs = dp(input_modality.bottom, inputs)
       with tf.variable_scope("body"):
         encoder_output, encoder_decoder_attention_bias = dp(
             self.encode,
@@ -636,7 +636,7 @@ def preprocess_targets(targets, i):
       # _shard_features called to ensure that the variable names match
       targets = self._shard_features({"targets": targets})["targets"]
       with tf.variable_scope(target_modality.name):
-        targets = target_modality.targets_bottom_sharded(targets, dp)[0]
+        targets = dp(target_modality.targets_bottom, targets)[0]
       targets = common_layers.flatten4d3d(targets)
 
       # TODO(llion): Explain! Is this even needed?
@@ -673,7 +673,7 @@ def symbols_to_logits_fn(ids, i, cache):
             nonpadding=features_to_nonpadding(features, "targets"))
 
       with tf.variable_scope(target_modality.name):
-        logits = target_modality.top_sharded(body_outputs, None, dp)[0]
+        logits = dp(target_modality.top, body_outputs, None)[0]
 
       ret = tf.squeeze(logits, axis=[1, 2, 3])
       if partial_targets is not None:
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index 1bc5bb749..55c69bd4d 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -112,19 +112,6 @@ def bottom(self, x):
     """
     raise NotImplementedError("Abstract Method")
 
-  def bottom_sharded(self, xs, data_parallelism):
-    """Transform the inputs.
-
-    Args:
-      xs: A list of num_datashards Tensors (one per shard)
-        each with shape [batch, p0, p1, depth]
-      data_parallelism: a expert_utils.Parallelism object
-    Returns:
-      shaded_body_input: A list of num_datashards Tensors, each with shape
-        [batch, p0, p1, body_input_depth].
-    """
-    return data_parallelism(self.bottom, xs)
-
   def targets_bottom(self, x):
     """Transform one shard of targets.
 
@@ -136,19 +123,6 @@ def targets_bottom(self, x):
     with tf.variable_scope("targets_bottom"):
       return self.bottom(x)
 
-  def targets_bottom_sharded(self, xs, data_parallelism):
-    """Transform the targets.
-
-    Args:
-      xs: A list of num_datashards Tensors (one per shard)
-        each with shape [batch, p0, p1, target_channels]
-      data_parallelism: a expert_utils.Parallelism object
-    Returns:
-      shaded_body_input: A list of num_datashards Tensors, each with shape
-        [batch, p0, p1, body_input_depth].
-    """
-    return data_parallelism(self.targets_bottom, xs)
-
   def top(self, body_output, targets):
     """Generate predictions/logits for one shard of output.
 
@@ -163,20 +137,6 @@ def top(self, body_output, targets):
     """
     raise NotImplementedError("Abstract Method")
 
-  def top_sharded(self, sharded_body_output, sharded_targets, data_parallelism):
-    """Generate predictions/logits for all shards.
-
-    Classes with cross-shard interaction will override this function.
-
-    Args:
-      sharded_body_output: A list of Tensors.
-      sharded_targets: A list of Tensors.
-      data_parallelism: a expert_utils.Parallelism object.
-    Returns:
-      sharded_logits: A list of Tensors.
-    """
-    return data_parallelism(self.top, sharded_body_output, sharded_targets)
-
   def loss(self, top_out, targets, weights_fn=None):
     """Compute loss numerator and denominator for one shard of output."""
     logits = top_out
@@ -189,14 +149,6 @@ def loss(self, top_out, targets, weights_fn=None):
         self._model_hparams.label_smoothing,
         weights_fn=weights_fn)
 
-  def loss_sharded(self, sharded_top_out, sharded_targets, data_parallelism):
-    """Compute loss for all shards."""
-    sharded_loss_num, sharded_loss_den = data_parallelism(
-        self.loss, sharded_top_out, sharded_targets)
-    loss = tf.add_n(sharded_loss_num) / tf.maximum(1.0,
-                                                   tf.add_n(sharded_loss_den))
-    return loss
-
   @property
   def is_class_modality(self):
     return self.name.startswith("class_label")
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 417427fd1..cd96d6350 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1955,18 +1955,22 @@ def sampled_results():
     new_features = transformed_features
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
       with tf.variable_scope(target_modality.name):
-        new_features["targets"] = target_modality.targets_bottom_sharded(
-            new_targets, dp)
+        new_features["targets"] = dp(target_modality.targets_bottom,
+                                     new_targets)
       with tf.variable_scope("body"):
         body_outputs, losses = model.model_fn_sharded(new_features)
         if not isinstance(losses, dict):  # If it's a single extra loss.
           losses = {"extra": losses}
       with tf.variable_scope(target_modality.name):
-        new_sharded_logits = target_modality.top_sharded(
-            body_outputs, sharded_features["targets"], dp)
+        new_sharded_logits = dp(target_modality.top,
+                                body_outputs,
+                                sharded_features["targets"])
         if "training" not in losses:
-          training_loss = target_modality.loss_sharded(
-              sharded_logits, sharded_features["targets"], dp)
+          sharded_loss_num, sharded_loss_den = dp(target_modality.loss,
+                                                  sharded_logits,
+                                                  sharded_features["targets"])
+          training_loss = (tf.add_n(sharded_loss_num) /
+                           tf.maximum(1.0, tf.add_n(sharded_loss_den)))
           training_loss *= problem_hparams.loss_multiplier
           losses["training"] = training_loss
     return new_sharded_logits, losses

From 5a867d031bd493eeb7d2776e1118d1594ff0a623 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 6 Feb 2019 12:00:50 -0800
Subject: [PATCH 1634/2720] Represent explicitly when one problem uses the
 vocabulary from another

PiperOrigin-RevId: 232718349
---
 .../data_generators/text_problems.py          | 18 +++++++++++++
 .../data_generators/translate_ende.py         | 25 +++++++++++++------
 .../data_generators/translate_enfr.py         | 21 +++++++++++-----
 .../data_generators/translate_enro.py         | 17 ++++++++++---
 tensor2tensor/data_generators/wiki_lm.py      |  4 +--
 5 files changed, 65 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 767d6ffa0..e137791c1 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -212,6 +212,9 @@ def generate_text_for_vocab(self, data_dir, tmp_dir):
 
   @property
   def vocab_filename(self):
+    other_problem = self.use_vocab_from_other_problem
+    if other_problem:
+      return other_problem.vocab_filename
     if self.vocab_type == VocabType.SUBWORD:
       return "vocab.%s.%d.%s" % (self.dataset_filename(),
                                  self.approx_vocab_size,
@@ -219,6 +222,18 @@ def vocab_filename(self):
     else:
       return "vocab.%s.%s" % (self.dataset_filename(), VocabType.TOKEN)
 
+  @property
+  def use_vocab_from_other_problem(self):
+    """Optional - use the vocabulary from a different problem.
+
+    TODO(noam): problems should override this method instead of overriding
+    vocab_filename(), so as to generate the correct vocabulary. Fix everywhere.
+
+    Returns:
+       a Text2TextProblem instance or None
+    """
+    return None
+
   def get_or_create_vocab(self, data_dir, tmp_dir, force_get=False):
     if self.vocab_type == VocabType.CHARACTER:
       encoder = text_encoder.ByteTextEncoder()
@@ -227,6 +242,9 @@ def get_or_create_vocab(self, data_dir, tmp_dir, force_get=False):
         vocab_filepath = os.path.join(data_dir, self.vocab_filename)
         encoder = text_encoder.SubwordTextEncoder(vocab_filepath)
       else:
+        other_problem = self.use_vocab_from_other_problem
+        if other_problem:
+          return other_problem.get_or_create_vocab(data_dir, tmp_dir, force_get)
         encoder = generator_utils.get_or_generate_vocab_inner(
             data_dir, self.vocab_filename, self.approx_vocab_size,
             self.generate_text_for_vocab(data_dir, tmp_dir),
diff --git a/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index 80a34c28c..45481711b 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -81,8 +81,8 @@ class TranslateEndeWmtParacrawlBicleaner32k(TranslateEndeWmt32k):
   """WMT en-de corpus with extra data from Paracrawl, cleaned with Bicleaner."""
 
   @property
-  def vocab_filename(self):
-    return TranslateEndeWmt32k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return TranslateEndeWmt32k()
 
   @property
   def additional_training_datasets(self):
@@ -99,8 +99,8 @@ def packed_length(self):
     return 256
 
   @property
-  def vocab_filename(self):
-    return TranslateEndeWmt32k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return TranslateEndeWmt32k()
 
 
 @registry.register_problem
@@ -111,8 +111,8 @@ def packed_length(self):
     return 256
 
   @property
-  def vocab_filename(self):
-    return TranslateEndeWmt8k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return TranslateEndeWmt8k()
 
 
 @registry.register_problem
@@ -129,5 +129,14 @@ class TranslateEndeWmtMulti64k(TranslateEndeWmt8k):
   """Translation with muli-lingual vocabulary."""
 
   @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
+
+
+@registry.register_problem
+class TranslateEndeWmtMulti64kPacked1k(TranslateEndeWmtMulti64k):
+  """Translation with muli-lingual vocabulary."""
+
+  @property
+  def packed_length(self):
+    return 1024
diff --git a/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index 95aa15d20..7d12e67ce 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -138,8 +138,8 @@ def packed_length(self):
     return 256
 
   @property
-  def vocab_filename(self):
-    return TranslateEnfrWmt32k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return TranslateEnfrWmt32k()
 
 
 @registry.register_problem
@@ -147,8 +147,8 @@ class TranslateEnfrWmt32kWithBacktranslateFr(TranslateEnfrWmt32k):
   """En-Fr translation with added French data, back-translated."""
 
   @property
-  def vocab_filename(self):
-    return TranslateEnfrWmt32k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return TranslateEnfrWmt32k()
 
   @property
   def already_shuffled(self):
@@ -248,5 +248,14 @@ def use_small_dataset(self):
     return False
 
   @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
+
+
+@registry.register_problem
+class TranslateEnfrWmtMulti64kPacked1k(TranslateEnfrWmtMulti64k):
+  """Translation with muli-lingual vocabulary."""
+
+  @property
+  def packed_length(self):
+    return 1024
diff --git a/tensor2tensor/data_generators/translate_enro.py b/tensor2tensor/data_generators/translate_enro.py
index 9b8d144a7..96421f714 100644
--- a/tensor2tensor/data_generators/translate_enro.py
+++ b/tensor2tensor/data_generators/translate_enro.py
@@ -78,8 +78,8 @@ class TranslateEnroWmtMulti64k(TranslateEnroWmt8k):
   """Translation with muli-lingual vocabulary."""
 
   @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
 
 
 @registry.register_problem
@@ -98,8 +98,8 @@ def dataset_splits(self):
     }]
 
   @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
 
   @property
   def how_many_examples_to_sample(self):
@@ -146,3 +146,12 @@ class TranslateEnroWmtMultiTiny64k(TranslateEnroWmtMultiSmall64k):
   @property
   def how_many_examples_to_sample(self):
     return 600
+
+
+@registry.register_problem
+class TranslateEnroWmtMultiTiny64kPacked1k(TranslateEnroWmtMultiTiny64k):
+  """Translation with muli-lingual vocabulary."""
+
+  @property
+  def packed_length(self):
+    return 1024
diff --git a/tensor2tensor/data_generators/wiki_lm.py b/tensor2tensor/data_generators/wiki_lm.py
index 6db24bb04..891fd1b01 100644
--- a/tensor2tensor/data_generators/wiki_lm.py
+++ b/tensor2tensor/data_generators/wiki_lm.py
@@ -162,8 +162,8 @@ def combine_characters_threshold(self):
     return 384*8
 
   @property
-  def vocab_filename(self):
-    return LanguagemodelEnWiki64k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return LanguagemodelEnWiki64k()
 
 
 @registry.register_problem

From f9f63c3f3c5aab8ed8070f10662ad8012c09c17b Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 7 Feb 2019 11:18:01 -0800
Subject: [PATCH 1635/2720] [V2] Make checkpointing work, support multiple
 targets.

PiperOrigin-RevId: 232910297
---
 tensor2tensor/__init__.py                     |  2 +-
 tensor2tensor/bin/__init__.py                 |  2 +-
 tensor2tensor/bin/build_vocab.py              |  2 +-
 tensor2tensor/bin/make_tf_configs.py          |  2 +-
 tensor2tensor/bin/t2t_attack.py               |  2 +-
 tensor2tensor/bin/t2t_avg_all.py              |  2 +-
 tensor2tensor/bin/t2t_bleu.py                 |  2 +-
 tensor2tensor/bin/t2t_datagen.py              |  2 +-
 tensor2tensor/bin/t2t_decoder.py              |  2 +-
 tensor2tensor/bin/t2t_distill.py              |  2 +-
 tensor2tensor/bin/t2t_eval.py                 |  2 +-
 tensor2tensor/bin/t2t_prune.py                |  2 +-
 tensor2tensor/bin/t2t_trainer.py              |  2 +-
 tensor2tensor/bin/t2t_trainer_test.py         |  2 +-
 tensor2tensor/bin/t2t_translate_all.py        |  2 +-
 tensor2tensor/data_generators/__init__.py     |  2 +-
 tensor2tensor/data_generators/algorithmic.py  |  2 +-
 .../data_generators/algorithmic_math.py       |  2 +-
 .../data_generators/algorithmic_math_test.py  |  2 +-
 .../data_generators/algorithmic_test.py       |  2 +-
 tensor2tensor/data_generators/all_problems.py |  2 +-
 tensor2tensor/data_generators/allen_brain.py  |  2 +-
 .../data_generators/allen_brain_test.py       |  2 +-
 tensor2tensor/data_generators/audio.py        |  2 +-
 .../data_generators/audio_encoder.py          |  2 +-
 tensor2tensor/data_generators/audio_test.py   |  2 +-
 tensor2tensor/data_generators/babi_qa.py      |  2 +-
 .../data_generators/bair_robot_pushing.py     |  2 +-
 tensor2tensor/data_generators/celeba.py       |  2 +-
 tensor2tensor/data_generators/celeba_test.py  |  2 +-
 tensor2tensor/data_generators/celebahq.py     |  2 +-
 tensor2tensor/data_generators/cifar.py        |  2 +-
 tensor2tensor/data_generators/cipher.py       |  2 +-
 .../data_generators/cnn_dailymail.py          |  2 +-
 tensor2tensor/data_generators/cola.py         |  2 +-
 tensor2tensor/data_generators/common_voice.py |  2 +-
 .../data_generators/common_voice_test.py      |  2 +-
 tensor2tensor/data_generators/conll_ner.py    |  2 +-
 tensor2tensor/data_generators/desc2code.py    |  2 +-
 .../data_generators/desc2code_test.py         |  2 +-
 tensor2tensor/data_generators/dna_encoder.py  |  2 +-
 .../data_generators/dna_encoder_test.py       |  2 +-
 tensor2tensor/data_generators/fsns.py         |  2 +-
 .../data_generators/function_docstring.py     |  2 +-
 .../data_generators/gene_expression.py        |  2 +-
 .../data_generators/gene_expression_test.py   |  2 +-
 .../data_generators/generator_utils.py        |  2 +-
 .../data_generators/generator_utils_test.py   |  2 +-
 .../data_generators/google_robot_pushing.py   |  2 +-
 tensor2tensor/data_generators/gym_env.py      |  2 +-
 tensor2tensor/data_generators/gym_env_test.py |  2 +-
 tensor2tensor/data_generators/ice_parsing.py  |  2 +-
 tensor2tensor/data_generators/image_lsun.py   |  2 +-
 tensor2tensor/data_generators/image_utils.py  |  2 +-
 .../data_generators/image_utils_test.py       |  2 +-
 tensor2tensor/data_generators/imagenet.py     |  2 +-
 .../data_generators/imagenet_test.py          |  2 +-
 tensor2tensor/data_generators/imdb.py         |  2 +-
 .../data_generators/inspect_tfrecord.py       |  2 +-
 tensor2tensor/data_generators/lambada.py      |  2 +-
 tensor2tensor/data_generators/librispeech.py  |  2 +-
 tensor2tensor/data_generators/lm1b.py         |  2 +-
 tensor2tensor/data_generators/lm1b_imdb.py    |  2 +-
 tensor2tensor/data_generators/lm1b_mnli.py    |  2 +-
 .../mathematical_language_understanding.py    |  2 +-
 tensor2tensor/data_generators/mnist.py        |  2 +-
 tensor2tensor/data_generators/mrpc.py         |  2 +-
 tensor2tensor/data_generators/mscoco.py       |  2 +-
 tensor2tensor/data_generators/mscoco_test.py  |  2 +-
 .../data_generators/multi_problem.py          |  2 +-
 .../data_generators/multi_problem_v2.py       |  2 +-
 .../data_generators/multi_problem_v2_test.py  |  2 +-
 tensor2tensor/data_generators/multinli.py     |  2 +-
 tensor2tensor/data_generators/ocr.py          |  2 +-
 .../data_generators/paraphrase_ms_coco.py     |  2 +-
 .../paraphrase_ms_coco_test.py                |  2 +-
 .../data_generators/pointer_generator_word.py |  2 +-
 tensor2tensor/data_generators/problem.py      |  2 +-
 .../data_generators/problem_hparams.py        |  2 +-
 tensor2tensor/data_generators/problem_test.py |  2 +-
 .../data_generators/program_search.py         |  2 +-
 .../data_generators/program_search_test.py    |  2 +-
 tensor2tensor/data_generators/ptb.py          |  2 +-
 tensor2tensor/data_generators/qnli.py         |  2 +-
 tensor2tensor/data_generators/quora_qpairs.py |  2 +-
 tensor2tensor/data_generators/rte.py          |  2 +-
 tensor2tensor/data_generators/scitail.py      |  2 +-
 tensor2tensor/data_generators/snli.py         |  2 +-
 .../data_generators/speech_recognition.py     |  2 +-
 tensor2tensor/data_generators/squad.py        |  2 +-
 tensor2tensor/data_generators/sst_binary.py   |  2 +-
 tensor2tensor/data_generators/stanford_nli.py |  2 +-
 .../data_generators/style_transfer.py         |  2 +-
 .../data_generators/style_transfer_test.py    |  2 +-
 .../data_generators/subject_verb_agreement.py |  2 +-
 tensor2tensor/data_generators/text_encoder.py |  2 +-
 .../text_encoder_build_subword.py             |  2 +-
 .../data_generators/text_encoder_test.py      |  2 +-
 .../data_generators/text_problems.py          |  2 +-
 .../data_generators/text_problems_test.py     |  2 +-
 tensor2tensor/data_generators/timeseries.py   |  2 +-
 .../timeseries_data_generator.py              |  2 +-
 .../timeseries_data_generator_test.py         |  2 +-
 .../data_generators/timeseries_test.py        |  2 +-
 tensor2tensor/data_generators/tokenizer.py    |  2 +-
 .../data_generators/tokenizer_test.py         |  2 +-
 .../data_generators/transduction_problems.py  |  2 +-
 .../transduction_problems_test.py             |  2 +-
 tensor2tensor/data_generators/translate.py    |  2 +-
 .../data_generators/translate_encs.py         |  2 +-
 .../data_generators/translate_ende.py         |  2 +-
 .../data_generators/translate_enet.py         |  2 +-
 .../data_generators/translate_enfr.py         |  2 +-
 .../data_generators/translate_enid.py         |  2 +-
 .../data_generators/translate_enmk.py         |  2 +-
 .../data_generators/translate_enro.py         |  2 +-
 .../data_generators/translate_entn.py         |  2 +-
 .../data_generators/translate_envi.py         |  2 +-
 .../data_generators/translate_enzh.py         |  2 +-
 .../data_generators/translate_test.py         |  2 +-
 .../data_generators/video_generated.py        |  2 +-
 tensor2tensor/data_generators/video_utils.py  |  2 +-
 .../data_generators/video_utils_test.py       |  2 +-
 tensor2tensor/data_generators/vqa.py          |  2 +-
 tensor2tensor/data_generators/vqa_utils.py    |  2 +-
 tensor2tensor/data_generators/wiki.py         |  2 +-
 tensor2tensor/data_generators/wiki_lm.py      |  2 +-
 .../data_generators/wiki_multi_problems.py    |  2 +-
 .../data_generators/wiki_revision.py          |  2 +-
 .../data_generators/wiki_revision_utils.py    |  2 +-
 .../data_generators/wikisum/__init__.py       |  2 +-
 .../data_generators/wikisum/generate_vocab.py |  2 +-
 .../wikisum/get_references_commoncrawl.py     |  2 +-
 .../wikisum/get_references_web.py             |  2 +-
 .../get_references_web_single_group.py        |  2 +-
 tensor2tensor/data_generators/wikisum/html.py |  2 +-
 .../wikisum/parallel_launch.py                |  2 +-
 .../wikisum/produce_examples.py               |  2 +-
 .../data_generators/wikisum/utils.py          |  2 +-
 .../data_generators/wikisum/utils_test.py     |  2 +-
 .../data_generators/wikisum/validate_data.py  |  2 +-
 .../data_generators/wikisum/wikisum.py        |  2 +-
 tensor2tensor/data_generators/wikitext103.py  |  2 +-
 tensor2tensor/data_generators/wnli.py         |  2 +-
 tensor2tensor/data_generators/wsj_parsing.py  |  2 +-
 tensor2tensor/envs/__init__.py                |  2 +-
 tensor2tensor/envs/env_problem.py             |  2 +-
 tensor2tensor/envs/env_problem_test.py        |  2 +-
 tensor2tensor/envs/gym_spaces_utils.py        |  2 +-
 tensor2tensor/envs/tic_tac_toe_env.py         |  2 +-
 tensor2tensor/envs/tic_tac_toe_env_problem.py |  2 +-
 .../envs/tic_tac_toe_env_problem_test.py      |  2 +-
 tensor2tensor/envs/tic_tac_toe_env_test.py    |  2 +-
 tensor2tensor/envs/time_step.py               |  2 +-
 tensor2tensor/envs/time_step_test.py          |  2 +-
 tensor2tensor/envs/trajectory.py              |  2 +-
 tensor2tensor/envs/trajectory_test.py         |  2 +-
 tensor2tensor/insights/__init__.py            |  2 +-
 tensor2tensor/insights/graph.py               |  2 +-
 tensor2tensor/insights/query_processor.py     |  2 +-
 tensor2tensor/insights/server.py              |  2 +-
 tensor2tensor/insights/transformer_model.py   |  2 +-
 tensor2tensor/layers/__init__.py              |  2 +-
 tensor2tensor/layers/bayes.py                 |  2 +-
 tensor2tensor/layers/bayes_test.py            |  2 +-
 tensor2tensor/layers/common_attention.py      |  2 +-
 tensor2tensor/layers/common_attention_test.py |  2 +-
 tensor2tensor/layers/common_audio.py          |  2 +-
 tensor2tensor/layers/common_hparams.py        |  2 +-
 .../layers/common_image_attention.py          |  2 +-
 .../layers/common_image_attention_test.py     |  2 +-
 tensor2tensor/layers/common_layers.py         |  2 +-
 tensor2tensor/layers/common_layers_test.py    |  2 +-
 .../common_message_passing_attention.py       |  2 +-
 tensor2tensor/layers/common_video.py          |  2 +-
 tensor2tensor/layers/common_video_test.py     |  2 +-
 tensor2tensor/layers/discretization.py        |  2 +-
 tensor2tensor/layers/discretization_test.py   |  2 +-
 tensor2tensor/layers/latent_layers.py         |  2 +-
 tensor2tensor/layers/latent_layers_test.py    |  2 +-
 tensor2tensor/layers/modalities.py            |  2 +-
 tensor2tensor/layers/modalities_test.py       |  2 +-
 tensor2tensor/layers/ngram.py                 |  2 +-
 tensor2tensor/layers/ngram_test.py            |  2 +-
 tensor2tensor/layers/reversible_layers.py     |  2 +-
 .../layers/reversible_layers_test.py          |  2 +-
 tensor2tensor/layers/transformer_layers.py    |  2 +-
 tensor2tensor/layers/vq_discrete.py           |  2 +-
 tensor2tensor/layers/vqa_layers.py            |  2 +-
 tensor2tensor/metrics/__init__.py             |  2 +-
 .../metrics/video_conditional_fvd.py          |  2 +-
 .../metrics/video_conditional_fvd_test.py     |  2 +-
 tensor2tensor/models/__init__.py              |  2 +-
 tensor2tensor/models/basic.py                 |  2 +-
 tensor2tensor/models/basic_test.py            |  2 +-
 tensor2tensor/models/bytenet.py               |  2 +-
 tensor2tensor/models/bytenet_test.py          |  2 +-
 tensor2tensor/models/distillation.py          |  2 +-
 tensor2tensor/models/evolved_transformer.py   |  2 +-
 .../models/evolved_transformer_test.py        |  2 +-
 tensor2tensor/models/image_transformer.py     |  2 +-
 tensor2tensor/models/image_transformer_2d.py  |  2 +-
 .../models/image_transformer_2d_test.py       |  2 +-
 .../models/image_transformer_test.py          |  2 +-
 tensor2tensor/models/lstm.py                  |  2 +-
 tensor2tensor/models/lstm_test.py             |  2 +-
 tensor2tensor/models/mtf_image_transformer.py |  2 +-
 .../models/mtf_image_transformer_test.py      |  2 +-
 tensor2tensor/models/mtf_resnet.py            |  2 +-
 tensor2tensor/models/mtf_transformer.py       |  2 +-
 tensor2tensor/models/mtf_transformer2.py      |  2 +-
 tensor2tensor/models/mtf_transformer_test.py  |  2 +-
 tensor2tensor/models/neural_gpu.py            |  2 +-
 tensor2tensor/models/neural_gpu_test.py       |  2 +-
 tensor2tensor/models/research/__init__.py     |  2 +-
 .../models/research/adafactor_experiments.py  |  2 +-
 tensor2tensor/models/research/aligned.py      |  2 +-
 tensor2tensor/models/research/attention_lm.py |  2 +-
 .../models/research/attention_lm_moe.py       |  2 +-
 tensor2tensor/models/research/autoencoders.py |  2 +-
 .../models/research/autoencoders_test.py      |  2 +-
 tensor2tensor/models/research/cycle_gan.py    |  2 +-
 .../models/research/gene_expression.py        |  2 +-
 .../models/research/gene_expression_test.py   |  2 +-
 tensor2tensor/models/research/glow.py         |  2 +-
 .../models/research/glow_init_hook.py         |  2 +-
 tensor2tensor/models/research/glow_ops.py     |  2 +-
 .../models/research/glow_ops_test.py          |  2 +-
 tensor2tensor/models/research/glow_test.py    |  2 +-
 .../models/research/lm_experiments.py         |  2 +-
 tensor2tensor/models/research/moe.py          |  2 +-
 .../models/research/moe_experiments.py        |  2 +-
 .../models/research/multiquery_paper.py       |  2 +-
 tensor2tensor/models/research/rl.py           |  2 +-
 .../models/research/similarity_transformer.py |  2 +-
 tensor2tensor/models/research/super_lm.py     |  2 +-
 .../models/research/transformer_aux.py        |  2 +-
 .../models/research/transformer_aux_test.py   |  2 +-
 .../models/research/transformer_moe.py        |  2 +-
 .../models/research/transformer_nat.py        |  2 +-
 .../models/research/transformer_parallel.py   |  2 +-
 .../models/research/transformer_revnet.py     |  2 +-
 .../research/transformer_revnet_test.py       |  2 +-
 .../models/research/transformer_sketch.py     |  2 +-
 .../models/research/transformer_symshard.py   |  2 +-
 .../models/research/transformer_vae.py        |  2 +-
 .../models/research/transformer_vae_test.py   |  2 +-
 .../models/research/universal_transformer.py  |  2 +-
 .../research/universal_transformer_test.py    |  2 +-
 .../research/universal_transformer_util.py    |  2 +-
 .../models/research/vqa_attention.py          |  2 +-
 .../models/research/vqa_attention_test.py     |  2 +-
 .../research/vqa_recurrent_self_attention.py  |  2 +-
 .../models/research/vqa_self_attention.py     |  2 +-
 tensor2tensor/models/resnet.py                |  2 +-
 tensor2tensor/models/resnet_test.py           |  2 +-
 tensor2tensor/models/revnet.py                |  2 +-
 tensor2tensor/models/revnet_test.py           |  2 +-
 tensor2tensor/models/shake_shake.py           |  2 +-
 tensor2tensor/models/slicenet.py              |  2 +-
 tensor2tensor/models/slicenet_test.py         |  2 +-
 tensor2tensor/models/text_cnn.py              |  2 +-
 tensor2tensor/models/transformer.py           |  2 +-
 tensor2tensor/models/transformer_test.py      |  2 +-
 tensor2tensor/models/vanilla_gan.py           |  2 +-
 tensor2tensor/models/video/__init__.py        |  2 +-
 tensor2tensor/models/video/base.py            |  2 +-
 tensor2tensor/models/video/base_vae.py        |  2 +-
 .../models/video/basic_deterministic.py       |  2 +-
 .../video/basic_deterministic_params.py       |  2 +-
 .../models/video/basic_deterministic_test.py  |  2 +-
 tensor2tensor/models/video/basic_recurrent.py |  2 +-
 .../models/video/basic_recurrent_test.py      |  2 +-
 .../models/video/basic_stochastic.py          |  2 +-
 .../models/video/basic_stochastic_test.py     |  2 +-
 tensor2tensor/models/video/emily.py           |  2 +-
 tensor2tensor/models/video/emily_test.py      |  2 +-
 tensor2tensor/models/video/epva.py            |  2 +-
 tensor2tensor/models/video/epva_params.py     |  2 +-
 tensor2tensor/models/video/next_frame_glow.py |  2 +-
 tensor2tensor/models/video/nfg_conv3d_test.py |  2 +-
 .../models/video/nfg_conv_lstm_test.py        |  2 +-
 tensor2tensor/models/video/nfg_conv_test.py   |  2 +-
 tensor2tensor/models/video/nfg_interpolate.py |  2 +-
 tensor2tensor/models/video/nfg_test_utils.py  |  2 +-
 tensor2tensor/models/video/nfg_uncond_test.py |  2 +-
 tensor2tensor/models/video/savp.py            |  2 +-
 tensor2tensor/models/video/savp_params.py     |  2 +-
 tensor2tensor/models/video/savp_test.py       |  2 +-
 tensor2tensor/models/video/sv2p.py            |  2 +-
 tensor2tensor/models/video/sv2p_params.py     |  2 +-
 tensor2tensor/models/video/sv2p_test.py       |  2 +-
 tensor2tensor/models/video/svg_lp.py          |  2 +-
 tensor2tensor/models/video/tests_utils.py     |  2 +-
 tensor2tensor/models/xception.py              |  2 +-
 tensor2tensor/models/xception_test.py         |  2 +-
 tensor2tensor/problems.py                     |  2 +-
 tensor2tensor/problems_test.py                |  2 +-
 tensor2tensor/rl/__init__.py                  |  2 +-
 tensor2tensor/rl/datagen_with_agent.py        |  2 +-
 tensor2tensor/rl/dopamine_connector.py        |  2 +-
 tensor2tensor/rl/envs/__init__.py             |  2 +-
 tensor2tensor/rl/envs/in_graph_batch_env.py   |  2 +-
 tensor2tensor/rl/envs/py_func_batch_env.py    |  2 +-
 tensor2tensor/rl/envs/simulated_batch_env.py  |  2 +-
 .../rl/envs/simulated_batch_gym_env.py        |  2 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py    |  2 +-
 tensor2tensor/rl/evaluator.py                 |  2 +-
 tensor2tensor/rl/evaluator_test.py            |  2 +-
 tensor2tensor/rl/gym_utils.py                 |  2 +-
 tensor2tensor/rl/gym_utils_test.py            |  2 +-
 tensor2tensor/rl/player.py                    |  2 +-
 tensor2tensor/rl/player_utils.py              |  2 +-
 tensor2tensor/rl/policy_learner.py            |  2 +-
 tensor2tensor/rl/ppo.py                       |  2 +-
 tensor2tensor/rl/ppo_learner.py               |  2 +-
 tensor2tensor/rl/restarter.py                 |  2 +-
 tensor2tensor/rl/restarter_test.py            |  2 +-
 tensor2tensor/rl/rl_utils.py                  |  2 +-
 tensor2tensor/rl/trainer_model_based.py       |  2 +-
 .../rl/trainer_model_based_agent_only.py      |  2 +-
 .../rl/trainer_model_based_dqn_test.py        |  2 +-
 .../rl/trainer_model_based_params.py          |  2 +-
 .../rl/trainer_model_based_recurrent_test.py  |  2 +-
 .../rl/trainer_model_based_stochastic_test.py |  2 +-
 .../rl/trainer_model_based_sv2p_test.py       |  2 +-
 tensor2tensor/rl/trainer_model_based_test.py  |  2 +-
 tensor2tensor/rl/trainer_model_free.py        |  2 +-
 tensor2tensor/rl/trainer_model_free_test.py   |  2 +-
 .../rl/trainer_model_free_tictactoe_test.py   |  2 +-
 tensor2tensor/serving/__init__.py             |  2 +-
 tensor2tensor/serving/export.py               |  2 +-
 tensor2tensor/serving/query.py                |  2 +-
 tensor2tensor/serving/serving_utils.py        |  2 +-
 .../test_data/example_usr_dir/__init__.py     |  2 +-
 .../test_data/example_usr_dir/my_submodule.py |  2 +-
 tensor2tensor/utils/__init__.py               |  2 +-
 tensor2tensor/utils/adafactor.py              |  2 +-
 tensor2tensor/utils/adv_attack_utils.py       |  2 +-
 tensor2tensor/utils/avg_checkpoints.py        |  2 +-
 tensor2tensor/utils/beam_search.py            |  2 +-
 tensor2tensor/utils/beam_search_test.py       |  2 +-
 tensor2tensor/utils/bleu_hook.py              |  2 +-
 tensor2tensor/utils/bleu_hook_test.py         |  2 +-
 .../utils/checkpoint_compatibility_test.py    |  2 +-
 tensor2tensor/utils/cloud_mlengine.py         |  2 +-
 tensor2tensor/utils/compute_video_metrics.py  |  2 +-
 tensor2tensor/utils/data_reader.py            |  2 +-
 tensor2tensor/utils/data_reader_test.py       |  2 +-
 tensor2tensor/utils/decoding.py               |  2 +-
 tensor2tensor/utils/devices.py                |  2 +-
 tensor2tensor/utils/diet.py                   |  2 +-
 tensor2tensor/utils/diet_test.py              |  2 +-
 tensor2tensor/utils/expert_utils.py           |  2 +-
 tensor2tensor/utils/expert_utils_test.py      |  2 +-
 tensor2tensor/utils/flags.py                  |  2 +-
 tensor2tensor/utils/get_rouge.py              |  2 +-
 tensor2tensor/utils/hparam.py                 |  2 +-
 tensor2tensor/utils/hparam_test.py            |  2 +-
 tensor2tensor/utils/hparams_lib.py            |  2 +-
 tensor2tensor/utils/hparams_lib_test.py       |  2 +-
 tensor2tensor/utils/learning_rate.py          |  2 +-
 tensor2tensor/utils/metrics.py                |  2 +-
 tensor2tensor/utils/metrics_hook.py           |  2 +-
 tensor2tensor/utils/metrics_hook_test.py      |  2 +-
 tensor2tensor/utils/metrics_test.py           |  2 +-
 tensor2tensor/utils/misc_utils.py             |  2 +-
 tensor2tensor/utils/misc_utils_test.py        |  2 +-
 tensor2tensor/utils/mlperf_log.py             |  2 +-
 tensor2tensor/utils/mlperf_tags.py            |  2 +-
 tensor2tensor/utils/modality.py               |  2 +-
 tensor2tensor/utils/mtf_model.py              |  2 +-
 tensor2tensor/utils/multistep_optimizer.py    |  2 +-
 .../utils/multistep_optimizer_test.py         |  2 +-
 tensor2tensor/utils/optimize.py               |  2 +-
 tensor2tensor/utils/optimize_test.py          |  2 +-
 tensor2tensor/utils/pruning_utils.py          |  2 +-
 tensor2tensor/utils/quantization.py           |  2 +-
 tensor2tensor/utils/registry.py               |  2 +-
 tensor2tensor/utils/registry_test.py          |  2 +-
 tensor2tensor/utils/restore_hook.py           |  2 +-
 tensor2tensor/utils/rouge.py                  |  2 +-
 tensor2tensor/utils/rouge_test.py             |  2 +-
 tensor2tensor/utils/sari_hook.py              |  2 +-
 tensor2tensor/utils/sari_hook_test.py         |  2 +-
 tensor2tensor/utils/t2t_model.py              |  2 +-
 tensor2tensor/utils/t2t_model_test.py         |  2 +-
 tensor2tensor/utils/test_utils.py             |  2 +-
 tensor2tensor/utils/test_utils_test.py        |  2 +-
 tensor2tensor/utils/trainer_lib.py            |  2 +-
 tensor2tensor/utils/trainer_lib_test.py       |  2 +-
 tensor2tensor/utils/update_ops_hook.py        |  2 +-
 tensor2tensor/utils/usr_dir.py                |  2 +-
 tensor2tensor/utils/video/prediction2gif.py   |  2 +-
 tensor2tensor/utils/video/reward_confusion.py |  2 +-
 tensor2tensor/utils/video2gif.py              |  2 +-
 tensor2tensor/utils/video_metrics.py          |  2 +-
 tensor2tensor/utils/video_metrics_test.py     |  2 +-
 tensor2tensor/utils/yellowfin.py              |  2 +-
 tensor2tensor/utils/yellowfin_test.py         |  2 +-
 tensor2tensor/v2/keras_utils.py               |  2 +-
 tensor2tensor/v2/models/basic.py              | 27 ++++-----
 tensor2tensor/v2/models/resnet.py             | 15 ++---
 tensor2tensor/v2/t2t.py                       | 55 ++++++++++++++-----
 tensor2tensor/v2/t2t_trainer.py               |  2 +-
 tensor2tensor/visualization/__init__.py       |  2 +-
 tensor2tensor/visualization/attention.py      |  2 +-
 tensor2tensor/visualization/visualization.py  |  2 +-
 .../visualization/visualization_test.py       |  2 +-
 409 files changed, 465 insertions(+), 444 deletions(-)

diff --git a/tensor2tensor/__init__.py b/tensor2tensor/__init__.py
index 4bd418a74..4872e5d5d 100644
--- a/tensor2tensor/__init__.py
+++ b/tensor2tensor/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/__init__.py b/tensor2tensor/bin/__init__.py
index 4bd418a74..4872e5d5d 100644
--- a/tensor2tensor/bin/__init__.py
+++ b/tensor2tensor/bin/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/build_vocab.py b/tensor2tensor/bin/build_vocab.py
index b93ec94c2..7b5be823d 100644
--- a/tensor2tensor/bin/build_vocab.py
+++ b/tensor2tensor/bin/build_vocab.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/make_tf_configs.py b/tensor2tensor/bin/make_tf_configs.py
index ef847bc3e..a27e3e376 100644
--- a/tensor2tensor/bin/make_tf_configs.py
+++ b/tensor2tensor/bin/make_tf_configs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index f7e0fcc60..6dbcc627e 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_avg_all.py b/tensor2tensor/bin/t2t_avg_all.py
index e3c34be90..4bb9f9949 100644
--- a/tensor2tensor/bin/t2t_avg_all.py
+++ b/tensor2tensor/bin/t2t_avg_all.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_bleu.py b/tensor2tensor/bin/t2t_bleu.py
index ab4a5f014..61fcb2438 100644
--- a/tensor2tensor/bin/t2t_bleu.py
+++ b/tensor2tensor/bin/t2t_bleu.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index c32a86cf5..a985eaedb 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 1656a5df5..61df0d4b3 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_distill.py b/tensor2tensor/bin/t2t_distill.py
index f970a56c2..2c05c9409 100644
--- a/tensor2tensor/bin/t2t_distill.py
+++ b/tensor2tensor/bin/t2t_distill.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_eval.py b/tensor2tensor/bin/t2t_eval.py
index e569fa497..2df551810 100644
--- a/tensor2tensor/bin/t2t_eval.py
+++ b/tensor2tensor/bin/t2t_eval.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
index e7c8c75db..1d2c668fd 100644
--- a/tensor2tensor/bin/t2t_prune.py
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 43432a082..c922e74dc 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_trainer_test.py b/tensor2tensor/bin/t2t_trainer_test.py
index 63c206ab0..9ed90dc63 100644
--- a/tensor2tensor/bin/t2t_trainer_test.py
+++ b/tensor2tensor/bin/t2t_trainer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_translate_all.py b/tensor2tensor/bin/t2t_translate_all.py
index 2d12e8246..401004fa0 100644
--- a/tensor2tensor/bin/t2t_translate_all.py
+++ b/tensor2tensor/bin/t2t_translate_all.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/__init__.py b/tensor2tensor/data_generators/__init__.py
index 4bd418a74..4872e5d5d 100644
--- a/tensor2tensor/data_generators/__init__.py
+++ b/tensor2tensor/data_generators/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index 790dcb854..ee59b486f 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math.py b/tensor2tensor/data_generators/algorithmic_math.py
index fcf038c7a..42767d421 100644
--- a/tensor2tensor/data_generators/algorithmic_math.py
+++ b/tensor2tensor/data_generators/algorithmic_math.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math_test.py b/tensor2tensor/data_generators/algorithmic_math_test.py
index bb4a3fae0..035ec16ce 100644
--- a/tensor2tensor/data_generators/algorithmic_math_test.py
+++ b/tensor2tensor/data_generators/algorithmic_math_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_test.py b/tensor2tensor/data_generators/algorithmic_test.py
index 1034409a8..c15bc76d3 100644
--- a/tensor2tensor/data_generators/algorithmic_test.py
+++ b/tensor2tensor/data_generators/algorithmic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index c9a8e763f..199a4bfa8 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index 30a8c9696..2f58de3e9 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/allen_brain_test.py b/tensor2tensor/data_generators/allen_brain_test.py
index 7871c4657..e2c875789 100644
--- a/tensor2tensor/data_generators/allen_brain_test.py
+++ b/tensor2tensor/data_generators/allen_brain_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio.py b/tensor2tensor/data_generators/audio.py
index ea6863d7f..2ad7ce29e 100644
--- a/tensor2tensor/data_generators/audio.py
+++ b/tensor2tensor/data_generators/audio.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio_encoder.py b/tensor2tensor/data_generators/audio_encoder.py
index da224e7fb..26fc01889 100644
--- a/tensor2tensor/data_generators/audio_encoder.py
+++ b/tensor2tensor/data_generators/audio_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio_test.py b/tensor2tensor/data_generators/audio_test.py
index 4cc6eb461..1af2ffd74 100644
--- a/tensor2tensor/data_generators/audio_test.py
+++ b/tensor2tensor/data_generators/audio_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index b6e4c6f67..d312a8272 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index 63eb6572a..c9bcba37f 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index 271c9ecd4..abfd0b784 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celeba_test.py b/tensor2tensor/data_generators/celeba_test.py
index 9f9247a65..598a081a6 100644
--- a/tensor2tensor/data_generators/celeba_test.py
+++ b/tensor2tensor/data_generators/celeba_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celebahq.py b/tensor2tensor/data_generators/celebahq.py
index 960d49e2a..259de8407 100644
--- a/tensor2tensor/data_generators/celebahq.py
+++ b/tensor2tensor/data_generators/celebahq.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index 23a106491..8bb7f5d4e 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cipher.py b/tensor2tensor/data_generators/cipher.py
index db07f5a7b..ddafeddb1 100644
--- a/tensor2tensor/data_generators/cipher.py
+++ b/tensor2tensor/data_generators/cipher.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index 71599b065..cbd9d84cf 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cola.py b/tensor2tensor/data_generators/cola.py
index 0f2748fea..be5fa6da5 100644
--- a/tensor2tensor/data_generators/cola.py
+++ b/tensor2tensor/data_generators/cola.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/common_voice.py b/tensor2tensor/data_generators/common_voice.py
index 1bc8fc126..2dda2ef99 100644
--- a/tensor2tensor/data_generators/common_voice.py
+++ b/tensor2tensor/data_generators/common_voice.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/common_voice_test.py b/tensor2tensor/data_generators/common_voice_test.py
index d7caefd32..337c4228c 100644
--- a/tensor2tensor/data_generators/common_voice_test.py
+++ b/tensor2tensor/data_generators/common_voice_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/conll_ner.py b/tensor2tensor/data_generators/conll_ner.py
index 0ec57bdc1..ff742b2c6 100644
--- a/tensor2tensor/data_generators/conll_ner.py
+++ b/tensor2tensor/data_generators/conll_ner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/desc2code.py b/tensor2tensor/data_generators/desc2code.py
index 383425f02..11216a06b 100644
--- a/tensor2tensor/data_generators/desc2code.py
+++ b/tensor2tensor/data_generators/desc2code.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/desc2code_test.py b/tensor2tensor/data_generators/desc2code_test.py
index 2b47f8271..b8a2bffac 100644
--- a/tensor2tensor/data_generators/desc2code_test.py
+++ b/tensor2tensor/data_generators/desc2code_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dna_encoder.py b/tensor2tensor/data_generators/dna_encoder.py
index e9c7afc93..c155dc9e2 100644
--- a/tensor2tensor/data_generators/dna_encoder.py
+++ b/tensor2tensor/data_generators/dna_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dna_encoder_test.py b/tensor2tensor/data_generators/dna_encoder_test.py
index b35c28824..2ceb91c75 100644
--- a/tensor2tensor/data_generators/dna_encoder_test.py
+++ b/tensor2tensor/data_generators/dna_encoder_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
index c7d566b8a..8dff69547 100644
--- a/tensor2tensor/data_generators/fsns.py
+++ b/tensor2tensor/data_generators/fsns.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
index 3d030a3d9..d2ce005ea 100644
--- a/tensor2tensor/data_generators/function_docstring.py
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gene_expression.py b/tensor2tensor/data_generators/gene_expression.py
index 04b009438..ee51db912 100644
--- a/tensor2tensor/data_generators/gene_expression.py
+++ b/tensor2tensor/data_generators/gene_expression.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gene_expression_test.py b/tensor2tensor/data_generators/gene_expression_test.py
index ab98da1c7..2461bbed9 100644
--- a/tensor2tensor/data_generators/gene_expression_test.py
+++ b/tensor2tensor/data_generators/gene_expression_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 9276674eb..30ffa9f5d 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/generator_utils_test.py b/tensor2tensor/data_generators/generator_utils_test.py
index 4249d60aa..2b4fc1fe4 100644
--- a/tensor2tensor/data_generators/generator_utils_test.py
+++ b/tensor2tensor/data_generators/generator_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index b37db990f..f49bd46e4 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 95fe1c2a2..8161102e4 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index c04a5fab4..011b6a273 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ice_parsing.py b/tensor2tensor/data_generators/ice_parsing.py
index d23576093..d5c42903a 100644
--- a/tensor2tensor/data_generators/ice_parsing.py
+++ b/tensor2tensor/data_generators/ice_parsing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_lsun.py b/tensor2tensor/data_generators/image_lsun.py
index 8f41dd1ce..c26e312cc 100644
--- a/tensor2tensor/data_generators/image_lsun.py
+++ b/tensor2tensor/data_generators/image_lsun.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 30db55bbe..2c0f8c584 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_utils_test.py b/tensor2tensor/data_generators/image_utils_test.py
index 326947e46..dfdf843a0 100644
--- a/tensor2tensor/data_generators/image_utils_test.py
+++ b/tensor2tensor/data_generators/image_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 8bd5bb5ff..b414ab4e3 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imagenet_test.py b/tensor2tensor/data_generators/imagenet_test.py
index 36f2954b0..66569cc0b 100644
--- a/tensor2tensor/data_generators/imagenet_test.py
+++ b/tensor2tensor/data_generators/imagenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imdb.py b/tensor2tensor/data_generators/imdb.py
index 0e8b6e28a..801b7e865 100644
--- a/tensor2tensor/data_generators/imdb.py
+++ b/tensor2tensor/data_generators/imdb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/inspect_tfrecord.py b/tensor2tensor/data_generators/inspect_tfrecord.py
index a245edb9e..ce2025ca2 100644
--- a/tensor2tensor/data_generators/inspect_tfrecord.py
+++ b/tensor2tensor/data_generators/inspect_tfrecord.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lambada.py b/tensor2tensor/data_generators/lambada.py
index 9af50cffc..753767c94 100644
--- a/tensor2tensor/data_generators/lambada.py
+++ b/tensor2tensor/data_generators/lambada.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index c19191513..2161edfd6 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b.py b/tensor2tensor/data_generators/lm1b.py
index dc120df83..2c19bbb42 100644
--- a/tensor2tensor/data_generators/lm1b.py
+++ b/tensor2tensor/data_generators/lm1b.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b_imdb.py b/tensor2tensor/data_generators/lm1b_imdb.py
index 05309c6f3..49aeac9b9 100644
--- a/tensor2tensor/data_generators/lm1b_imdb.py
+++ b/tensor2tensor/data_generators/lm1b_imdb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b_mnli.py b/tensor2tensor/data_generators/lm1b_mnli.py
index 66530dc36..8ab1af68c 100644
--- a/tensor2tensor/data_generators/lm1b_mnli.py
+++ b/tensor2tensor/data_generators/lm1b_mnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mathematical_language_understanding.py b/tensor2tensor/data_generators/mathematical_language_understanding.py
index 5d0f431ba..f01295d7b 100644
--- a/tensor2tensor/data_generators/mathematical_language_understanding.py
+++ b/tensor2tensor/data_generators/mathematical_language_understanding.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mnist.py b/tensor2tensor/data_generators/mnist.py
index 322aecb60..40f03977b 100644
--- a/tensor2tensor/data_generators/mnist.py
+++ b/tensor2tensor/data_generators/mnist.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mrpc.py b/tensor2tensor/data_generators/mrpc.py
index 47c8364d2..7848a98a4 100644
--- a/tensor2tensor/data_generators/mrpc.py
+++ b/tensor2tensor/data_generators/mrpc.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index d57bd616a..93c51d3ee 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mscoco_test.py b/tensor2tensor/data_generators/mscoco_test.py
index 6958180bf..ce87332a2 100644
--- a/tensor2tensor/data_generators/mscoco_test.py
+++ b/tensor2tensor/data_generators/mscoco_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 60e969d46..27ae5c118 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem_v2.py b/tensor2tensor/data_generators/multi_problem_v2.py
index 7f0717bcf..004aa1cb8 100644
--- a/tensor2tensor/data_generators/multi_problem_v2.py
+++ b/tensor2tensor/data_generators/multi_problem_v2.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem_v2_test.py b/tensor2tensor/data_generators/multi_problem_v2_test.py
index c189856de..e714850c9 100644
--- a/tensor2tensor/data_generators/multi_problem_v2_test.py
+++ b/tensor2tensor/data_generators/multi_problem_v2_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index e3af79f0e..f3f6d3d92 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ocr.py b/tensor2tensor/data_generators/ocr.py
index 39cfa56a8..8f68d9885 100644
--- a/tensor2tensor/data_generators/ocr.py
+++ b/tensor2tensor/data_generators/ocr.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/paraphrase_ms_coco.py b/tensor2tensor/data_generators/paraphrase_ms_coco.py
index 4c459a324..df7bd30a8 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/paraphrase_ms_coco_test.py b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
index 763ad7af5..aa57bce3b 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/pointer_generator_word.py b/tensor2tensor/data_generators/pointer_generator_word.py
index 0180a9ad4..ea14ba1c3 100644
--- a/tensor2tensor/data_generators/pointer_generator_word.py
+++ b/tensor2tensor/data_generators/pointer_generator_word.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 7f0ac8104..4d8a972f0 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem_hparams.py b/tensor2tensor/data_generators/problem_hparams.py
index f67742de4..ad0b76d27 100644
--- a/tensor2tensor/data_generators/problem_hparams.py
+++ b/tensor2tensor/data_generators/problem_hparams.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 4cab788e6..9c0344767 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/program_search.py b/tensor2tensor/data_generators/program_search.py
index b0fe553ca..160949857 100644
--- a/tensor2tensor/data_generators/program_search.py
+++ b/tensor2tensor/data_generators/program_search.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/program_search_test.py b/tensor2tensor/data_generators/program_search_test.py
index e73a59a6b..41fcaa37d 100644
--- a/tensor2tensor/data_generators/program_search_test.py
+++ b/tensor2tensor/data_generators/program_search_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ptb.py b/tensor2tensor/data_generators/ptb.py
index 5684b21eb..6f63077cc 100644
--- a/tensor2tensor/data_generators/ptb.py
+++ b/tensor2tensor/data_generators/ptb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/qnli.py b/tensor2tensor/data_generators/qnli.py
index 9eeeb2077..732a46c18 100644
--- a/tensor2tensor/data_generators/qnli.py
+++ b/tensor2tensor/data_generators/qnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/quora_qpairs.py b/tensor2tensor/data_generators/quora_qpairs.py
index 21f3702ec..003f34e28 100644
--- a/tensor2tensor/data_generators/quora_qpairs.py
+++ b/tensor2tensor/data_generators/quora_qpairs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/rte.py b/tensor2tensor/data_generators/rte.py
index 2eff16422..892c40f32 100644
--- a/tensor2tensor/data_generators/rte.py
+++ b/tensor2tensor/data_generators/rte.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/scitail.py b/tensor2tensor/data_generators/scitail.py
index f600fa560..a9d023a55 100644
--- a/tensor2tensor/data_generators/scitail.py
+++ b/tensor2tensor/data_generators/scitail.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/snli.py b/tensor2tensor/data_generators/snli.py
index 446562265..8bad0f699 100644
--- a/tensor2tensor/data_generators/snli.py
+++ b/tensor2tensor/data_generators/snli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 730fe65f0..5426988b3 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/squad.py b/tensor2tensor/data_generators/squad.py
index 1cf541984..82105e013 100644
--- a/tensor2tensor/data_generators/squad.py
+++ b/tensor2tensor/data_generators/squad.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/sst_binary.py b/tensor2tensor/data_generators/sst_binary.py
index 9081fc81a..c5682404c 100644
--- a/tensor2tensor/data_generators/sst_binary.py
+++ b/tensor2tensor/data_generators/sst_binary.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/stanford_nli.py b/tensor2tensor/data_generators/stanford_nli.py
index 9c99501ec..2ff90d0d6 100644
--- a/tensor2tensor/data_generators/stanford_nli.py
+++ b/tensor2tensor/data_generators/stanford_nli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/style_transfer.py b/tensor2tensor/data_generators/style_transfer.py
index d7a046564..cdd9ceacc 100644
--- a/tensor2tensor/data_generators/style_transfer.py
+++ b/tensor2tensor/data_generators/style_transfer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/style_transfer_test.py b/tensor2tensor/data_generators/style_transfer_test.py
index fe9618ae2..534f6335a 100644
--- a/tensor2tensor/data_generators/style_transfer_test.py
+++ b/tensor2tensor/data_generators/style_transfer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/subject_verb_agreement.py b/tensor2tensor/data_generators/subject_verb_agreement.py
index 18b802d8c..57f0992d4 100644
--- a/tensor2tensor/data_generators/subject_verb_agreement.py
+++ b/tensor2tensor/data_generators/subject_verb_agreement.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 139fceb82..4fe071093 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_encoder_build_subword.py b/tensor2tensor/data_generators/text_encoder_build_subword.py
index 89c6b9516..a7c9c37b3 100644
--- a/tensor2tensor/data_generators/text_encoder_build_subword.py
+++ b/tensor2tensor/data_generators/text_encoder_build_subword.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_encoder_test.py b/tensor2tensor/data_generators/text_encoder_test.py
index ba5bbf998..29f7c808f 100644
--- a/tensor2tensor/data_generators/text_encoder_test.py
+++ b/tensor2tensor/data_generators/text_encoder_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index e137791c1..9eeba462e 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_problems_test.py b/tensor2tensor/data_generators/text_problems_test.py
index 31d0992f8..51f948fcf 100644
--- a/tensor2tensor/data_generators/text_problems_test.py
+++ b/tensor2tensor/data_generators/text_problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index 23154ad73..e8f01e511 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_data_generator.py b/tensor2tensor/data_generators/timeseries_data_generator.py
index 13fe70af0..94c2c17c3 100644
--- a/tensor2tensor/data_generators/timeseries_data_generator.py
+++ b/tensor2tensor/data_generators/timeseries_data_generator.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_data_generator_test.py b/tensor2tensor/data_generators/timeseries_data_generator_test.py
index 6d4c208c9..1050ab45b 100644
--- a/tensor2tensor/data_generators/timeseries_data_generator_test.py
+++ b/tensor2tensor/data_generators/timeseries_data_generator_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index a203b5306..e0947108a 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/tokenizer.py b/tensor2tensor/data_generators/tokenizer.py
index d5449748e..0bce48fd3 100644
--- a/tensor2tensor/data_generators/tokenizer.py
+++ b/tensor2tensor/data_generators/tokenizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/tokenizer_test.py b/tensor2tensor/data_generators/tokenizer_test.py
index ca9c3b8e5..932529d90 100644
--- a/tensor2tensor/data_generators/tokenizer_test.py
+++ b/tensor2tensor/data_generators/tokenizer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/transduction_problems.py b/tensor2tensor/data_generators/transduction_problems.py
index 01fd5f9a5..6e1cff743 100644
--- a/tensor2tensor/data_generators/transduction_problems.py
+++ b/tensor2tensor/data_generators/transduction_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/transduction_problems_test.py b/tensor2tensor/data_generators/transduction_problems_test.py
index a8dbe5b5e..3678b24e4 100644
--- a/tensor2tensor/data_generators/transduction_problems_test.py
+++ b/tensor2tensor/data_generators/transduction_problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index e9f8356fa..dceb9d96b 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_encs.py b/tensor2tensor/data_generators/translate_encs.py
index d5d6a3f99..5fd7a321b 100644
--- a/tensor2tensor/data_generators/translate_encs.py
+++ b/tensor2tensor/data_generators/translate_encs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index 45481711b..d49a4f0b0 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enet.py b/tensor2tensor/data_generators/translate_enet.py
index 4c3b4e929..77ad70f07 100644
--- a/tensor2tensor/data_generators/translate_enet.py
+++ b/tensor2tensor/data_generators/translate_enet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index 7d12e67ce..fe4e31e80 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enid.py b/tensor2tensor/data_generators/translate_enid.py
index aa070637e..bf598cae2 100644
--- a/tensor2tensor/data_generators/translate_enid.py
+++ b/tensor2tensor/data_generators/translate_enid.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enmk.py b/tensor2tensor/data_generators/translate_enmk.py
index 6caa1ee9d..07947bf04 100644
--- a/tensor2tensor/data_generators/translate_enmk.py
+++ b/tensor2tensor/data_generators/translate_enmk.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enro.py b/tensor2tensor/data_generators/translate_enro.py
index 96421f714..e34fe7566 100644
--- a/tensor2tensor/data_generators/translate_enro.py
+++ b/tensor2tensor/data_generators/translate_enro.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_entn.py b/tensor2tensor/data_generators/translate_entn.py
index 2ae6e6647..ab78df4d1 100644
--- a/tensor2tensor/data_generators/translate_entn.py
+++ b/tensor2tensor/data_generators/translate_entn.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_envi.py b/tensor2tensor/data_generators/translate_envi.py
index b3a1a7885..5e333fd55 100644
--- a/tensor2tensor/data_generators/translate_envi.py
+++ b/tensor2tensor/data_generators/translate_envi.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index 8e4930dbf..e5ed9cddb 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_test.py b/tensor2tensor/data_generators/translate_test.py
index a489b10d3..f0bbdd6dc 100644
--- a/tensor2tensor/data_generators/translate_test.py
+++ b/tensor2tensor/data_generators/translate_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 419aefa44..aa0e5bcf3 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 28f1f73d3..0072b68d0 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index 62304044a..7db91dade 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
index 2ba91942d..8a39f85f6 100644
--- a/tensor2tensor/data_generators/vqa.py
+++ b/tensor2tensor/data_generators/vqa.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/vqa_utils.py b/tensor2tensor/data_generators/vqa_utils.py
index a1a6bc6e2..660d3c943 100644
--- a/tensor2tensor/data_generators/vqa_utils.py
+++ b/tensor2tensor/data_generators/vqa_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki.py b/tensor2tensor/data_generators/wiki.py
index ee3dcbc17..3ea3d9e39 100644
--- a/tensor2tensor/data_generators/wiki.py
+++ b/tensor2tensor/data_generators/wiki.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_lm.py b/tensor2tensor/data_generators/wiki_lm.py
index 891fd1b01..37cb3b79b 100644
--- a/tensor2tensor/data_generators/wiki_lm.py
+++ b/tensor2tensor/data_generators/wiki_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index b86cb1d8e..37e0a4204 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_revision.py b/tensor2tensor/data_generators/wiki_revision.py
index 62bc90c63..089c941c1 100644
--- a/tensor2tensor/data_generators/wiki_revision.py
+++ b/tensor2tensor/data_generators/wiki_revision.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index 9704c068c..cd6fa57f6 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/__init__.py b/tensor2tensor/data_generators/wikisum/__init__.py
index 4bd418a74..4872e5d5d 100644
--- a/tensor2tensor/data_generators/wikisum/__init__.py
+++ b/tensor2tensor/data_generators/wikisum/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/generate_vocab.py b/tensor2tensor/data_generators/wikisum/generate_vocab.py
index 0231e0e80..431b2373f 100644
--- a/tensor2tensor/data_generators/wikisum/generate_vocab.py
+++ b/tensor2tensor/data_generators/wikisum/generate_vocab.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py b/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
index 52cec62a6..d8f346662 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_web.py b/tensor2tensor/data_generators/wikisum/get_references_web.py
index 2a8d9699d..9c6524b1a 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py b/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
index d27686a4b..49d04a70a 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/html.py b/tensor2tensor/data_generators/wikisum/html.py
index e8483ae8b..5a70dcb2d 100644
--- a/tensor2tensor/data_generators/wikisum/html.py
+++ b/tensor2tensor/data_generators/wikisum/html.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/parallel_launch.py b/tensor2tensor/data_generators/wikisum/parallel_launch.py
index 3332a8978..194ee3dcf 100644
--- a/tensor2tensor/data_generators/wikisum/parallel_launch.py
+++ b/tensor2tensor/data_generators/wikisum/parallel_launch.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/produce_examples.py b/tensor2tensor/data_generators/wikisum/produce_examples.py
index 95f736978..16b9a6c98 100644
--- a/tensor2tensor/data_generators/wikisum/produce_examples.py
+++ b/tensor2tensor/data_generators/wikisum/produce_examples.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/utils.py b/tensor2tensor/data_generators/wikisum/utils.py
index fcaf0399a..09e60e7a1 100644
--- a/tensor2tensor/data_generators/wikisum/utils.py
+++ b/tensor2tensor/data_generators/wikisum/utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/utils_test.py b/tensor2tensor/data_generators/wikisum/utils_test.py
index a20afbf4c..36397bae6 100644
--- a/tensor2tensor/data_generators/wikisum/utils_test.py
+++ b/tensor2tensor/data_generators/wikisum/utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/validate_data.py b/tensor2tensor/data_generators/wikisum/validate_data.py
index 618759a74..8ebf0c4ae 100644
--- a/tensor2tensor/data_generators/wikisum/validate_data.py
+++ b/tensor2tensor/data_generators/wikisum/validate_data.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/wikisum.py b/tensor2tensor/data_generators/wikisum/wikisum.py
index 19e1816c1..691490156 100644
--- a/tensor2tensor/data_generators/wikisum/wikisum.py
+++ b/tensor2tensor/data_generators/wikisum/wikisum.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikitext103.py b/tensor2tensor/data_generators/wikitext103.py
index 24ebc2763..1a25b2984 100644
--- a/tensor2tensor/data_generators/wikitext103.py
+++ b/tensor2tensor/data_generators/wikitext103.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wnli.py b/tensor2tensor/data_generators/wnli.py
index 9b94b5b43..413ffbf88 100644
--- a/tensor2tensor/data_generators/wnli.py
+++ b/tensor2tensor/data_generators/wnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wsj_parsing.py b/tensor2tensor/data_generators/wsj_parsing.py
index b9adf2a7e..ae42e73d2 100644
--- a/tensor2tensor/data_generators/wsj_parsing.py
+++ b/tensor2tensor/data_generators/wsj_parsing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/__init__.py b/tensor2tensor/envs/__init__.py
index 0d3981f49..675172fb9 100644
--- a/tensor2tensor/envs/__init__.py
+++ b/tensor2tensor/envs/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index c57bfe287..2732dafe3 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/env_problem_test.py b/tensor2tensor/envs/env_problem_test.py
index 92de34a58..eedb6c171 100644
--- a/tensor2tensor/envs/env_problem_test.py
+++ b/tensor2tensor/envs/env_problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_spaces_utils.py b/tensor2tensor/envs/gym_spaces_utils.py
index 53eef08ba..1176db0ff 100644
--- a/tensor2tensor/envs/gym_spaces_utils.py
+++ b/tensor2tensor/envs/gym_spaces_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env.py b/tensor2tensor/envs/tic_tac_toe_env.py
index 7f54aac48..4a6b1a1e7 100644
--- a/tensor2tensor/envs/tic_tac_toe_env.py
+++ b/tensor2tensor/envs/tic_tac_toe_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_problem.py b/tensor2tensor/envs/tic_tac_toe_env_problem.py
index 0d748f82a..3a57a11ef 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_problem_test.py b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
index 954be159a..214ab7196 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_test.py b/tensor2tensor/envs/tic_tac_toe_env_test.py
index b2da18570..8925d1dbf 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/time_step.py b/tensor2tensor/envs/time_step.py
index 2e4d8a45c..f02567d9e 100644
--- a/tensor2tensor/envs/time_step.py
+++ b/tensor2tensor/envs/time_step.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/time_step_test.py b/tensor2tensor/envs/time_step_test.py
index 410ffb546..02eeb9612 100644
--- a/tensor2tensor/envs/time_step_test.py
+++ b/tensor2tensor/envs/time_step_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index 4a1ff12a2..0aee091df 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index be5ed1de6..ec577716d 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/__init__.py b/tensor2tensor/insights/__init__.py
index 4bd418a74..4872e5d5d 100644
--- a/tensor2tensor/insights/__init__.py
+++ b/tensor2tensor/insights/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/graph.py b/tensor2tensor/insights/graph.py
index 17e18ea3c..2ca6bd91c 100644
--- a/tensor2tensor/insights/graph.py
+++ b/tensor2tensor/insights/graph.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/query_processor.py b/tensor2tensor/insights/query_processor.py
index d6703af44..3456b0725 100644
--- a/tensor2tensor/insights/query_processor.py
+++ b/tensor2tensor/insights/query_processor.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/server.py b/tensor2tensor/insights/server.py
index 783b865bd..ed2b15aa8 100644
--- a/tensor2tensor/insights/server.py
+++ b/tensor2tensor/insights/server.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/transformer_model.py b/tensor2tensor/insights/transformer_model.py
index da8cf5fe3..8cad91251 100644
--- a/tensor2tensor/insights/transformer_model.py
+++ b/tensor2tensor/insights/transformer_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/__init__.py b/tensor2tensor/layers/__init__.py
index 4bd418a74..4872e5d5d 100644
--- a/tensor2tensor/layers/__init__.py
+++ b/tensor2tensor/layers/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index fd482e370..2ef1a0cc9 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 15eb0564b..66a0f107a 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 258d47b8a..f0f6674ea 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index ce9f4417d..8d1575df3 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_audio.py b/tensor2tensor/layers/common_audio.py
index 14432576f..8c388728e 100644
--- a/tensor2tensor/layers/common_audio.py
+++ b/tensor2tensor/layers/common_audio.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 7602e6b84..390dfff9b 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 1580ae376..4f5ea7b67 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index bf6ce85e6..1d22927b6 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index b7bb5615b..3b60583dc 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index 50fc19950..b3cdf6be2 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_message_passing_attention.py b/tensor2tensor/layers/common_message_passing_attention.py
index b744ffd07..949527fe2 100644
--- a/tensor2tensor/layers/common_message_passing_attention.py
+++ b/tensor2tensor/layers/common_message_passing_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 70a5931e3..81f7f0568 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index a58bda345..50493069d 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 020d82d1d..bb6c89b6f 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index effe43c4b..1ba6fde3e 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index b493f5d04..2788f5e8f 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
index 0e5b2ac73..2e1ddc77b 100644
--- a/tensor2tensor/layers/latent_layers_test.py
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index a42a251aa..79a453ac8 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index d14de2927..603ef154f 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/ngram.py b/tensor2tensor/layers/ngram.py
index a865d489d..675a20448 100644
--- a/tensor2tensor/layers/ngram.py
+++ b/tensor2tensor/layers/ngram.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/ngram_test.py b/tensor2tensor/layers/ngram_test.py
index 475d5f4cb..02ef6ba06 100644
--- a/tensor2tensor/layers/ngram_test.py
+++ b/tensor2tensor/layers/ngram_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/reversible_layers.py b/tensor2tensor/layers/reversible_layers.py
index 4a8c02ebc..4c7003c15 100644
--- a/tensor2tensor/layers/reversible_layers.py
+++ b/tensor2tensor/layers/reversible_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
index 309e543f6..bc0e73870 100644
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index c63151286..d019eef71 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index 344767084..ff6c8ddd1 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/vqa_layers.py b/tensor2tensor/layers/vqa_layers.py
index 537e7ac8f..f81c66709 100644
--- a/tensor2tensor/layers/vqa_layers.py
+++ b/tensor2tensor/layers/vqa_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/__init__.py b/tensor2tensor/metrics/__init__.py
index 4bd418a74..4872e5d5d 100644
--- a/tensor2tensor/metrics/__init__.py
+++ b/tensor2tensor/metrics/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/video_conditional_fvd.py b/tensor2tensor/metrics/video_conditional_fvd.py
index b9965e1cb..d97b946d2 100644
--- a/tensor2tensor/metrics/video_conditional_fvd.py
+++ b/tensor2tensor/metrics/video_conditional_fvd.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/video_conditional_fvd_test.py b/tensor2tensor/metrics/video_conditional_fvd_test.py
index ec1712eec..84446325e 100644
--- a/tensor2tensor/metrics/video_conditional_fvd_test.py
+++ b/tensor2tensor/metrics/video_conditional_fvd_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index ca3f7f3f2..475da8231 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/basic.py b/tensor2tensor/models/basic.py
index a0a46ff58..155a73ec5 100644
--- a/tensor2tensor/models/basic.py
+++ b/tensor2tensor/models/basic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/basic_test.py b/tensor2tensor/models/basic_test.py
index eaef58326..a89550da9 100644
--- a/tensor2tensor/models/basic_test.py
+++ b/tensor2tensor/models/basic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/bytenet.py b/tensor2tensor/models/bytenet.py
index 3eec2afda..de5e6c550 100644
--- a/tensor2tensor/models/bytenet.py
+++ b/tensor2tensor/models/bytenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/bytenet_test.py b/tensor2tensor/models/bytenet_test.py
index e68c10b98..934b6ead4 100644
--- a/tensor2tensor/models/bytenet_test.py
+++ b/tensor2tensor/models/bytenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/distillation.py b/tensor2tensor/models/distillation.py
index ed08e076c..1dbe46bc7 100644
--- a/tensor2tensor/models/distillation.py
+++ b/tensor2tensor/models/distillation.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 0395fbcbd..4c97231cb 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/evolved_transformer_test.py b/tensor2tensor/models/evolved_transformer_test.py
index 0056ed2f0..ae29f1c83 100644
--- a/tensor2tensor/models/evolved_transformer_test.py
+++ b/tensor2tensor/models/evolved_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 1f841e28b..0b3975193 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 5de2584c2..fe5d31444 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_2d_test.py b/tensor2tensor/models/image_transformer_2d_test.py
index 9beaee900..7f903fb15 100644
--- a/tensor2tensor/models/image_transformer_2d_test.py
+++ b/tensor2tensor/models/image_transformer_2d_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_test.py b/tensor2tensor/models/image_transformer_test.py
index 0ccc1f6bd..3b61cfb40 100644
--- a/tensor2tensor/models/image_transformer_test.py
+++ b/tensor2tensor/models/image_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 9841c6019..9d343cdb9 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/lstm_test.py b/tensor2tensor/models/lstm_test.py
index 370748cd0..95b6fe4f0 100644
--- a/tensor2tensor/models/lstm_test.py
+++ b/tensor2tensor/models/lstm_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_image_transformer.py b/tensor2tensor/models/mtf_image_transformer.py
index 1564af7f7..e169e221c 100644
--- a/tensor2tensor/models/mtf_image_transformer.py
+++ b/tensor2tensor/models/mtf_image_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_image_transformer_test.py b/tensor2tensor/models/mtf_image_transformer_test.py
index 17c72ce96..8bcef7a45 100644
--- a/tensor2tensor/models/mtf_image_transformer_test.py
+++ b/tensor2tensor/models/mtf_image_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_resnet.py b/tensor2tensor/models/mtf_resnet.py
index 430c46406..3ffac2d19 100644
--- a/tensor2tensor/models/mtf_resnet.py
+++ b/tensor2tensor/models/mtf_resnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index b70fd6257..167ab7394 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index fb5842f87..6ea437901 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer_test.py b/tensor2tensor/models/mtf_transformer_test.py
index 672b02d6a..b6b8b5334 100644
--- a/tensor2tensor/models/mtf_transformer_test.py
+++ b/tensor2tensor/models/mtf_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_gpu.py b/tensor2tensor/models/neural_gpu.py
index d7369d247..1b824d349 100644
--- a/tensor2tensor/models/neural_gpu.py
+++ b/tensor2tensor/models/neural_gpu.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_gpu_test.py b/tensor2tensor/models/neural_gpu_test.py
index 89cdc9a42..8d3e7068b 100644
--- a/tensor2tensor/models/neural_gpu_test.py
+++ b/tensor2tensor/models/neural_gpu_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/__init__.py b/tensor2tensor/models/research/__init__.py
index 4bd418a74..4872e5d5d 100644
--- a/tensor2tensor/models/research/__init__.py
+++ b/tensor2tensor/models/research/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/adafactor_experiments.py b/tensor2tensor/models/research/adafactor_experiments.py
index fbe4dbc2f..9fe1c38e9 100644
--- a/tensor2tensor/models/research/adafactor_experiments.py
+++ b/tensor2tensor/models/research/adafactor_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/aligned.py b/tensor2tensor/models/research/aligned.py
index 7a36ec991..26a857647 100644
--- a/tensor2tensor/models/research/aligned.py
+++ b/tensor2tensor/models/research/aligned.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/attention_lm.py b/tensor2tensor/models/research/attention_lm.py
index 3a03efafb..93d6f5e1a 100644
--- a/tensor2tensor/models/research/attention_lm.py
+++ b/tensor2tensor/models/research/attention_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/attention_lm_moe.py b/tensor2tensor/models/research/attention_lm_moe.py
index 4bb0c71f8..7cfadb1ed 100644
--- a/tensor2tensor/models/research/attention_lm_moe.py
+++ b/tensor2tensor/models/research/attention_lm_moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 0da05beba..8105ffd32 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index a7e396f8a..6b7ff62d8 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index 89693dc27..5ba2a01f8 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/gene_expression.py b/tensor2tensor/models/research/gene_expression.py
index e15d2fd8d..955980298 100644
--- a/tensor2tensor/models/research/gene_expression.py
+++ b/tensor2tensor/models/research/gene_expression.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/gene_expression_test.py b/tensor2tensor/models/research/gene_expression_test.py
index 1f0d461c9..142b4315b 100644
--- a/tensor2tensor/models/research/gene_expression_test.py
+++ b/tensor2tensor/models/research/gene_expression_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 072a04aaa..9fa39d742 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_init_hook.py b/tensor2tensor/models/research/glow_init_hook.py
index 0a88c973b..86009f3e4 100644
--- a/tensor2tensor/models/research/glow_init_hook.py
+++ b/tensor2tensor/models/research/glow_init_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index ef58517c0..daab40340 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index dfef7b727..ef01ff9b7 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
index 12414b857..653b23769 100644
--- a/tensor2tensor/models/research/glow_test.py
+++ b/tensor2tensor/models/research/glow_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/lm_experiments.py b/tensor2tensor/models/research/lm_experiments.py
index 97fef03a2..cf9f2a75f 100644
--- a/tensor2tensor/models/research/lm_experiments.py
+++ b/tensor2tensor/models/research/lm_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/moe.py b/tensor2tensor/models/research/moe.py
index 11c31a1b6..f70b0eece 100644
--- a/tensor2tensor/models/research/moe.py
+++ b/tensor2tensor/models/research/moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/moe_experiments.py b/tensor2tensor/models/research/moe_experiments.py
index 04a09cfd1..ee21605f3 100644
--- a/tensor2tensor/models/research/moe_experiments.py
+++ b/tensor2tensor/models/research/moe_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/multiquery_paper.py b/tensor2tensor/models/research/multiquery_paper.py
index 5b5491026..4fdc07b6f 100644
--- a/tensor2tensor/models/research/multiquery_paper.py
+++ b/tensor2tensor/models/research/multiquery_paper.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 0f8e1aa46..833425664 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/similarity_transformer.py b/tensor2tensor/models/research/similarity_transformer.py
index 43b0b31ed..4d316a961 100644
--- a/tensor2tensor/models/research/similarity_transformer.py
+++ b/tensor2tensor/models/research/similarity_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/super_lm.py b/tensor2tensor/models/research/super_lm.py
index 6b4b137ec..fc0b403c5 100644
--- a/tensor2tensor/models/research/super_lm.py
+++ b/tensor2tensor/models/research/super_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_aux.py b/tensor2tensor/models/research/transformer_aux.py
index fb486757f..ff5aafacf 100644
--- a/tensor2tensor/models/research/transformer_aux.py
+++ b/tensor2tensor/models/research/transformer_aux.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_aux_test.py b/tensor2tensor/models/research/transformer_aux_test.py
index d9b2f9d6a..11ec146f0 100644
--- a/tensor2tensor/models/research/transformer_aux_test.py
+++ b/tensor2tensor/models/research/transformer_aux_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index 26129cd4f..2a4899f04 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index c6be947fc..1eb01c5cd 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_parallel.py b/tensor2tensor/models/research/transformer_parallel.py
index 671db7eb7..23b04f99f 100644
--- a/tensor2tensor/models/research/transformer_parallel.py
+++ b/tensor2tensor/models/research/transformer_parallel.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_revnet.py b/tensor2tensor/models/research/transformer_revnet.py
index b21338a33..21a17788b 100644
--- a/tensor2tensor/models/research/transformer_revnet.py
+++ b/tensor2tensor/models/research/transformer_revnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_revnet_test.py b/tensor2tensor/models/research/transformer_revnet_test.py
index ce4a083c7..37c3648ba 100644
--- a/tensor2tensor/models/research/transformer_revnet_test.py
+++ b/tensor2tensor/models/research/transformer_revnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_sketch.py b/tensor2tensor/models/research/transformer_sketch.py
index a64d959b6..9b3a0cf94 100644
--- a/tensor2tensor/models/research/transformer_sketch.py
+++ b/tensor2tensor/models/research/transformer_sketch.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_symshard.py b/tensor2tensor/models/research/transformer_symshard.py
index 8e00e0ca2..98c9cf0ee 100644
--- a/tensor2tensor/models/research/transformer_symshard.py
+++ b/tensor2tensor/models/research/transformer_symshard.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 44e92c5ce..e8f769bb6 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae_test.py b/tensor2tensor/models/research/transformer_vae_test.py
index 0cee4bc5e..bb04f44bc 100644
--- a/tensor2tensor/models/research/transformer_vae_test.py
+++ b/tensor2tensor/models/research/transformer_vae_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 6fc230fcb..9945bb140 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
index 15f40ffbd..947cecc6f 100644
--- a/tensor2tensor/models/research/universal_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 54906e5b6..2937e7dea 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index 57f1975e6..c73f25ab3 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_attention_test.py b/tensor2tensor/models/research/vqa_attention_test.py
index 174970dfd..6d8904a2a 100644
--- a/tensor2tensor/models/research/vqa_attention_test.py
+++ b/tensor2tensor/models/research/vqa_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_recurrent_self_attention.py b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
index 1be385c1f..6765bce16 100644
--- a/tensor2tensor/models/research/vqa_recurrent_self_attention.py
+++ b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
index b8388e606..4702b598f 100644
--- a/tensor2tensor/models/research/vqa_self_attention.py
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index 39277a3ab..eba093135 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/resnet_test.py b/tensor2tensor/models/resnet_test.py
index 776da4f65..0e15dbbf4 100644
--- a/tensor2tensor/models/resnet_test.py
+++ b/tensor2tensor/models/resnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/revnet.py b/tensor2tensor/models/revnet.py
index ebde0dbe8..8a4f24542 100644
--- a/tensor2tensor/models/revnet.py
+++ b/tensor2tensor/models/revnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/revnet_test.py b/tensor2tensor/models/revnet_test.py
index e60bc8420..8344f4090 100644
--- a/tensor2tensor/models/revnet_test.py
+++ b/tensor2tensor/models/revnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index 01681bb05..51b52af2e 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/slicenet.py b/tensor2tensor/models/slicenet.py
index bbaaeddaa..f9d300e98 100644
--- a/tensor2tensor/models/slicenet.py
+++ b/tensor2tensor/models/slicenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/slicenet_test.py b/tensor2tensor/models/slicenet_test.py
index b2463e899..cf38c2b04 100644
--- a/tensor2tensor/models/slicenet_test.py
+++ b/tensor2tensor/models/slicenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/text_cnn.py b/tensor2tensor/models/text_cnn.py
index 532506853..0a6957f2f 100644
--- a/tensor2tensor/models/text_cnn.py
+++ b/tensor2tensor/models/text_cnn.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 5429208d1..a0da5b1cc 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index e6882a3ed..20f1ee726 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index 00fa89d18..bf922df50 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/__init__.py b/tensor2tensor/models/video/__init__.py
index 4bd418a74..4872e5d5d 100644
--- a/tensor2tensor/models/video/__init__.py
+++ b/tensor2tensor/models/video/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 2326772c6..b54d6356f 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/base_vae.py b/tensor2tensor/models/video/base_vae.py
index 5b1716bf3..ca9620dae 100644
--- a/tensor2tensor/models/video/base_vae.py
+++ b/tensor2tensor/models/video/base_vae.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index d065042b5..882fc01b8 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 547c0fcbb..05566bb7a 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic_test.py b/tensor2tensor/models/video/basic_deterministic_test.py
index 20d007478..0c16666a7 100644
--- a/tensor2tensor/models/video/basic_deterministic_test.py
+++ b/tensor2tensor/models/video/basic_deterministic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_recurrent.py b/tensor2tensor/models/video/basic_recurrent.py
index 42cae07de..e2af2747a 100644
--- a/tensor2tensor/models/video/basic_recurrent.py
+++ b/tensor2tensor/models/video/basic_recurrent.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_recurrent_test.py b/tensor2tensor/models/video/basic_recurrent_test.py
index d9deb5c98..dfcef5bd7 100644
--- a/tensor2tensor/models/video/basic_recurrent_test.py
+++ b/tensor2tensor/models/video/basic_recurrent_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 62f158e11..db08ae864 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_stochastic_test.py b/tensor2tensor/models/video/basic_stochastic_test.py
index d74f42117..4eb339d64 100644
--- a/tensor2tensor/models/video/basic_stochastic_test.py
+++ b/tensor2tensor/models/video/basic_stochastic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index 221e41eeb..8d0ab7953 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/emily_test.py b/tensor2tensor/models/video/emily_test.py
index b670c5de8..9ccf34c79 100644
--- a/tensor2tensor/models/video/emily_test.py
+++ b/tensor2tensor/models/video/emily_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index dbd4e5283..a1e4941b7 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
index 0283f154e..a8b8f98c9 100644
--- a/tensor2tensor/models/video/epva_params.py
+++ b/tensor2tensor/models/video/epva_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index 9cbcea1a1..56a9da10b 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv3d_test.py b/tensor2tensor/models/video/nfg_conv3d_test.py
index c539ac20f..9c7f0fab8 100644
--- a/tensor2tensor/models/video/nfg_conv3d_test.py
+++ b/tensor2tensor/models/video/nfg_conv3d_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv_lstm_test.py b/tensor2tensor/models/video/nfg_conv_lstm_test.py
index 965ea8ff4..046cec72d 100644
--- a/tensor2tensor/models/video/nfg_conv_lstm_test.py
+++ b/tensor2tensor/models/video/nfg_conv_lstm_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv_test.py b/tensor2tensor/models/video/nfg_conv_test.py
index 2ccc54e53..51e00ed1c 100644
--- a/tensor2tensor/models/video/nfg_conv_test.py
+++ b/tensor2tensor/models/video/nfg_conv_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_interpolate.py b/tensor2tensor/models/video/nfg_interpolate.py
index d9d746ea6..e7284dde7 100644
--- a/tensor2tensor/models/video/nfg_interpolate.py
+++ b/tensor2tensor/models/video/nfg_interpolate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_test_utils.py b/tensor2tensor/models/video/nfg_test_utils.py
index 01744e9de..ec48cf818 100644
--- a/tensor2tensor/models/video/nfg_test_utils.py
+++ b/tensor2tensor/models/video/nfg_test_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_uncond_test.py b/tensor2tensor/models/video/nfg_uncond_test.py
index e2ceb50d3..97fa984d1 100644
--- a/tensor2tensor/models/video/nfg_uncond_test.py
+++ b/tensor2tensor/models/video/nfg_uncond_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index f73a36ef2..bfda04240 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index ff95e3902..708c888f3 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp_test.py b/tensor2tensor/models/video/savp_test.py
index 2061eb06b..5ec184b4e 100644
--- a/tensor2tensor/models/video/savp_test.py
+++ b/tensor2tensor/models/video/savp_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 3c6c75af6..551dd6171 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index f0d456c29..674b5b211 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p_test.py b/tensor2tensor/models/video/sv2p_test.py
index e261f2cfa..3694af807 100644
--- a/tensor2tensor/models/video/sv2p_test.py
+++ b/tensor2tensor/models/video/sv2p_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/svg_lp.py b/tensor2tensor/models/video/svg_lp.py
index 7242ade84..f495fb9ff 100644
--- a/tensor2tensor/models/video/svg_lp.py
+++ b/tensor2tensor/models/video/svg_lp.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/tests_utils.py b/tensor2tensor/models/video/tests_utils.py
index 7bda4b093..b8677396c 100644
--- a/tensor2tensor/models/video/tests_utils.py
+++ b/tensor2tensor/models/video/tests_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/xception.py b/tensor2tensor/models/xception.py
index 3f59690b0..7c6c6ca40 100644
--- a/tensor2tensor/models/xception.py
+++ b/tensor2tensor/models/xception.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/xception_test.py b/tensor2tensor/models/xception_test.py
index 98ed8e8af..487e5aafe 100644
--- a/tensor2tensor/models/xception_test.py
+++ b/tensor2tensor/models/xception_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/problems.py b/tensor2tensor/problems.py
index 2cf8c8762..b6f6f8bdd 100644
--- a/tensor2tensor/problems.py
+++ b/tensor2tensor/problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/problems_test.py b/tensor2tensor/problems_test.py
index 07387fd53..c49611ed6 100644
--- a/tensor2tensor/problems_test.py
+++ b/tensor2tensor/problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/__init__.py b/tensor2tensor/rl/__init__.py
index 4bd418a74..4872e5d5d 100644
--- a/tensor2tensor/rl/__init__.py
+++ b/tensor2tensor/rl/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/datagen_with_agent.py b/tensor2tensor/rl/datagen_with_agent.py
index 1e5100537..43c52186f 100644
--- a/tensor2tensor/rl/datagen_with_agent.py
+++ b/tensor2tensor/rl/datagen_with_agent.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index cc59d7dfa..9e1ad3765 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/__init__.py b/tensor2tensor/rl/envs/__init__.py
index 4bd418a74..4872e5d5d 100644
--- a/tensor2tensor/rl/envs/__init__.py
+++ b/tensor2tensor/rl/envs/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index b20d18f10..1164f7872 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index b83cf5377..50d3ee3dc 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 373a1549e..eac13e012 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index f5d0da563..e0ffd9805 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 7d5ac6b20..5c801c9ba 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index c34e75c9b..0459da703 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/evaluator_test.py b/tensor2tensor/rl/evaluator_test.py
index a1f8a2b62..dea1a6666 100644
--- a/tensor2tensor/rl/evaluator_test.py
+++ b/tensor2tensor/rl/evaluator_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index ed0a82fc1..72bff9f2f 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index ead4563af..ce4f86e5c 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 4bc6bd2a0..4ec8f0e2d 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index 63695c094..41789c868 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
index db6bc5854..82051f159 100644
--- a/tensor2tensor/rl/policy_learner.py
+++ b/tensor2tensor/rl/policy_learner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index f152499de..a40a46f59 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 457e18da1..e5ba6dfb7 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/restarter.py b/tensor2tensor/rl/restarter.py
index c9a2b1110..e996984b9 100644
--- a/tensor2tensor/rl/restarter.py
+++ b/tensor2tensor/rl/restarter.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/restarter_test.py b/tensor2tensor/rl/restarter_test.py
index 9db1148d1..ff0455ecd 100644
--- a/tensor2tensor/rl/restarter_test.py
+++ b/tensor2tensor/rl/restarter_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index d385f04b4..d6630fd6a 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 2ddb18ae2..31b90726a 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_agent_only.py b/tensor2tensor/rl/trainer_model_based_agent_only.py
index 801a54df2..e221106a7 100644
--- a/tensor2tensor/rl/trainer_model_based_agent_only.py
+++ b/tensor2tensor/rl/trainer_model_based_agent_only.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_dqn_test.py b/tensor2tensor/rl/trainer_model_based_dqn_test.py
index d8af9de9d..50c056649 100644
--- a/tensor2tensor/rl/trainer_model_based_dqn_test.py
+++ b/tensor2tensor/rl/trainer_model_based_dqn_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 169eff73c..00f384507 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_recurrent_test.py b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
index 42d1089f6..beda282a6 100644
--- a/tensor2tensor/rl/trainer_model_based_recurrent_test.py
+++ b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_stochastic_test.py b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
index 018ab2675..b0598a320 100644
--- a/tensor2tensor/rl/trainer_model_based_stochastic_test.py
+++ b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_sv2p_test.py b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
index 8be291555..905b2cfc8 100644
--- a/tensor2tensor/rl/trainer_model_based_sv2p_test.py
+++ b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_test.py b/tensor2tensor/rl/trainer_model_based_test.py
index 83cab7d60..0c836162c 100644
--- a/tensor2tensor/rl/trainer_model_based_test.py
+++ b/tensor2tensor/rl/trainer_model_based_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index b3d6ea238..73cae5684 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index aa75aac9c..90c728ff0 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free_tictactoe_test.py b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
index 24d734734..76fb37339 100644
--- a/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
+++ b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/__init__.py b/tensor2tensor/serving/__init__.py
index 4bd418a74..4872e5d5d 100644
--- a/tensor2tensor/serving/__init__.py
+++ b/tensor2tensor/serving/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 8a8f85a63..1e17c56f4 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/query.py b/tensor2tensor/serving/query.py
index a8c46ce80..a0ea895c7 100644
--- a/tensor2tensor/serving/query.py
+++ b/tensor2tensor/serving/query.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 52695116a..5d873aa8d 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/test_data/example_usr_dir/__init__.py b/tensor2tensor/test_data/example_usr_dir/__init__.py
index 61a511e17..ed4c6aca1 100644
--- a/tensor2tensor/test_data/example_usr_dir/__init__.py
+++ b/tensor2tensor/test_data/example_usr_dir/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/test_data/example_usr_dir/my_submodule.py b/tensor2tensor/test_data/example_usr_dir/my_submodule.py
index e3ffd962c..1ba2b439a 100644
--- a/tensor2tensor/test_data/example_usr_dir/my_submodule.py
+++ b/tensor2tensor/test_data/example_usr_dir/my_submodule.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/__init__.py b/tensor2tensor/utils/__init__.py
index 4bd418a74..4872e5d5d 100644
--- a/tensor2tensor/utils/__init__.py
+++ b/tensor2tensor/utils/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index 9c44d28d9..e578e6b19 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/adv_attack_utils.py b/tensor2tensor/utils/adv_attack_utils.py
index 070f9cfd4..b1ec90d9d 100644
--- a/tensor2tensor/utils/adv_attack_utils.py
+++ b/tensor2tensor/utils/adv_attack_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index 01684a394..10e58bd31 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index 15a4dd266..0c57f0c6d 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/beam_search_test.py b/tensor2tensor/utils/beam_search_test.py
index 7f6aa3595..47991405d 100644
--- a/tensor2tensor/utils/beam_search_test.py
+++ b/tensor2tensor/utils/beam_search_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/bleu_hook.py b/tensor2tensor/utils/bleu_hook.py
index 2fa4ce0a4..7346b2b74 100644
--- a/tensor2tensor/utils/bleu_hook.py
+++ b/tensor2tensor/utils/bleu_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/bleu_hook_test.py b/tensor2tensor/utils/bleu_hook_test.py
index 999a240f8..ae086342d 100644
--- a/tensor2tensor/utils/bleu_hook_test.py
+++ b/tensor2tensor/utils/bleu_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/checkpoint_compatibility_test.py b/tensor2tensor/utils/checkpoint_compatibility_test.py
index e99250133..6657adb9c 100644
--- a/tensor2tensor/utils/checkpoint_compatibility_test.py
+++ b/tensor2tensor/utils/checkpoint_compatibility_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 37a9437c3..8d2de835f 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/compute_video_metrics.py b/tensor2tensor/utils/compute_video_metrics.py
index 9db41038d..f41888fdd 100644
--- a/tensor2tensor/utils/compute_video_metrics.py
+++ b/tensor2tensor/utils/compute_video_metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index 1599b6a63..4aa21438d 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/data_reader_test.py b/tensor2tensor/utils/data_reader_test.py
index 9487a8cbd..32544da36 100644
--- a/tensor2tensor/utils/data_reader_test.py
+++ b/tensor2tensor/utils/data_reader_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 7f0d5bbc1..b20901180 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index 2a582de84..bed966566 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/diet.py b/tensor2tensor/utils/diet.py
index a8d5d7c2f..73d4ffeda 100644
--- a/tensor2tensor/utils/diet.py
+++ b/tensor2tensor/utils/diet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/diet_test.py b/tensor2tensor/utils/diet_test.py
index 667ad1f7f..0a41e58f1 100644
--- a/tensor2tensor/utils/diet_test.py
+++ b/tensor2tensor/utils/diet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 0630b7ca4..8a97c764c 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/expert_utils_test.py b/tensor2tensor/utils/expert_utils_test.py
index 5c9cc714e..3c6feec2f 100644
--- a/tensor2tensor/utils/expert_utils_test.py
+++ b/tensor2tensor/utils/expert_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index 05d51ec12..5525a1796 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/get_rouge.py b/tensor2tensor/utils/get_rouge.py
index e374811aa..ac8ea0761 100644
--- a/tensor2tensor/utils/get_rouge.py
+++ b/tensor2tensor/utils/get_rouge.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparam.py b/tensor2tensor/utils/hparam.py
index ed7f02ded..aeb53cc2e 100644
--- a/tensor2tensor/utils/hparam.py
+++ b/tensor2tensor/utils/hparam.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparam_test.py b/tensor2tensor/utils/hparam_test.py
index 8c39e3485..600349ffd 100644
--- a/tensor2tensor/utils/hparam_test.py
+++ b/tensor2tensor/utils/hparam_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparams_lib.py b/tensor2tensor/utils/hparams_lib.py
index e72c4f3f8..0d02a291c 100644
--- a/tensor2tensor/utils/hparams_lib.py
+++ b/tensor2tensor/utils/hparams_lib.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparams_lib_test.py b/tensor2tensor/utils/hparams_lib_test.py
index e0c874ea4..651743806 100644
--- a/tensor2tensor/utils/hparams_lib_test.py
+++ b/tensor2tensor/utils/hparams_lib_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index b78402760..377280fa1 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 94d9f64d2..f3b05a00b 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_hook.py b/tensor2tensor/utils/metrics_hook.py
index 2fed06c18..6186b1df6 100644
--- a/tensor2tensor/utils/metrics_hook.py
+++ b/tensor2tensor/utils/metrics_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_hook_test.py b/tensor2tensor/utils/metrics_hook_test.py
index f707f24af..3069a8f8d 100644
--- a/tensor2tensor/utils/metrics_hook_test.py
+++ b/tensor2tensor/utils/metrics_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index a8dbde9ab..ede480d8f 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/misc_utils.py b/tensor2tensor/utils/misc_utils.py
index bef5fc0c6..90d640221 100644
--- a/tensor2tensor/utils/misc_utils.py
+++ b/tensor2tensor/utils/misc_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/misc_utils_test.py b/tensor2tensor/utils/misc_utils_test.py
index 11228de68..5d5ae805d 100644
--- a/tensor2tensor/utils/misc_utils_test.py
+++ b/tensor2tensor/utils/misc_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mlperf_log.py b/tensor2tensor/utils/mlperf_log.py
index 81e11cf93..304d60b6d 100644
--- a/tensor2tensor/utils/mlperf_log.py
+++ b/tensor2tensor/utils/mlperf_log.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mlperf_tags.py b/tensor2tensor/utils/mlperf_tags.py
index 764c4399e..6de1fd71d 100644
--- a/tensor2tensor/utils/mlperf_tags.py
+++ b/tensor2tensor/utils/mlperf_tags.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index 55c69bd4d..a2a6fddb1 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mtf_model.py b/tensor2tensor/utils/mtf_model.py
index 4a6d88f46..d39fe9902 100644
--- a/tensor2tensor/utils/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
index 32dd6cdd1..2e997e2bc 100644
--- a/tensor2tensor/utils/multistep_optimizer.py
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_optimizer_test.py b/tensor2tensor/utils/multistep_optimizer_test.py
index 19d248b03..05b9fb059 100644
--- a/tensor2tensor/utils/multistep_optimizer_test.py
+++ b/tensor2tensor/utils/multistep_optimizer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index dda809797..ae962c32f 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/optimize_test.py b/tensor2tensor/utils/optimize_test.py
index ea1478647..1a24a9d34 100644
--- a/tensor2tensor/utils/optimize_test.py
+++ b/tensor2tensor/utils/optimize_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/pruning_utils.py b/tensor2tensor/utils/pruning_utils.py
index 2cc4cb35b..1023ce904 100644
--- a/tensor2tensor/utils/pruning_utils.py
+++ b/tensor2tensor/utils/pruning_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/quantization.py b/tensor2tensor/utils/quantization.py
index 4ae86309f..06e84ec5c 100644
--- a/tensor2tensor/utils/quantization.py
+++ b/tensor2tensor/utils/quantization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index ed20b1c91..a89dd6739 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index 1864881e2..55e1359fd 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/restore_hook.py b/tensor2tensor/utils/restore_hook.py
index c649231c7..6d649c209 100644
--- a/tensor2tensor/utils/restore_hook.py
+++ b/tensor2tensor/utils/restore_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/rouge.py b/tensor2tensor/utils/rouge.py
index cb3c9af4b..e6de0df2c 100644
--- a/tensor2tensor/utils/rouge.py
+++ b/tensor2tensor/utils/rouge.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/rouge_test.py b/tensor2tensor/utils/rouge_test.py
index eab9eaeeb..5dc6a42cf 100644
--- a/tensor2tensor/utils/rouge_test.py
+++ b/tensor2tensor/utils/rouge_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/sari_hook.py b/tensor2tensor/utils/sari_hook.py
index 65641b1d9..fba2899ea 100644
--- a/tensor2tensor/utils/sari_hook.py
+++ b/tensor2tensor/utils/sari_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/sari_hook_test.py b/tensor2tensor/utils/sari_hook_test.py
index b168d0143..e3c9a3aa8 100644
--- a/tensor2tensor/utils/sari_hook_test.py
+++ b/tensor2tensor/utils/sari_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index cd96d6350..66b748ebb 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index 76847439b..0c9ac408e 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/test_utils.py b/tensor2tensor/utils/test_utils.py
index 0d48dd3bf..292949f6e 100644
--- a/tensor2tensor/utils/test_utils.py
+++ b/tensor2tensor/utils/test_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/test_utils_test.py b/tensor2tensor/utils/test_utils_test.py
index d7bb25ab6..f9f07502e 100644
--- a/tensor2tensor/utils/test_utils_test.py
+++ b/tensor2tensor/utils/test_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index ff56cf8a6..89fbb6604 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index c37144ca8..2995a4e84 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/update_ops_hook.py b/tensor2tensor/utils/update_ops_hook.py
index c2615c661..f3252060b 100644
--- a/tensor2tensor/utils/update_ops_hook.py
+++ b/tensor2tensor/utils/update_ops_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/usr_dir.py b/tensor2tensor/utils/usr_dir.py
index 06b5955ea..6817e9831 100644
--- a/tensor2tensor/utils/usr_dir.py
+++ b/tensor2tensor/utils/usr_dir.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video/prediction2gif.py b/tensor2tensor/utils/video/prediction2gif.py
index fc55b7e4c..cd2d15f53 100644
--- a/tensor2tensor/utils/video/prediction2gif.py
+++ b/tensor2tensor/utils/video/prediction2gif.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video/reward_confusion.py b/tensor2tensor/utils/video/reward_confusion.py
index 0e6c7e773..115303934 100644
--- a/tensor2tensor/utils/video/reward_confusion.py
+++ b/tensor2tensor/utils/video/reward_confusion.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video2gif.py b/tensor2tensor/utils/video2gif.py
index a66e75402..2e09df0df 100644
--- a/tensor2tensor/utils/video2gif.py
+++ b/tensor2tensor/utils/video2gif.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index 177916021..4515bb808 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video_metrics_test.py b/tensor2tensor/utils/video_metrics_test.py
index 18ddcec5c..b3f7b53fa 100644
--- a/tensor2tensor/utils/video_metrics_test.py
+++ b/tensor2tensor/utils/video_metrics_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/yellowfin.py b/tensor2tensor/utils/yellowfin.py
index 97acee4fb..e21025f84 100644
--- a/tensor2tensor/utils/yellowfin.py
+++ b/tensor2tensor/utils/yellowfin.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/yellowfin_test.py b/tensor2tensor/utils/yellowfin_test.py
index 67999707c..0a59c3917 100644
--- a/tensor2tensor/utils/yellowfin_test.py
+++ b/tensor2tensor/utils/yellowfin_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/v2/keras_utils.py b/tensor2tensor/v2/keras_utils.py
index a9eeebb42..566da85f1 100644
--- a/tensor2tensor/v2/keras_utils.py
+++ b/tensor2tensor/v2/keras_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/v2/models/basic.py b/tensor2tensor/v2/models/basic.py
index 4d2cbe054..8ed3d758e 100644
--- a/tensor2tensor/v2/models/basic.py
+++ b/tensor2tensor/v2/models/basic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -27,15 +27,12 @@
 class BasicFcRelu(tf.keras.Model):
   """Basic fully-connected + ReLU model."""
 
-  def __init__(self, features_info=None, supervised_keys=None,
+  def __init__(self, features_info=None, input_names=None, target_names=None,
                num_hidden_layers=2, hidden_size=64, dropout=0.1):
     super(BasicFcRelu, self).__init__()
-    assert features_info is not None
-    assert supervised_keys is not None
-    self._input_key = supervised_keys[0]
-    label_key = supervised_keys[1]
-    input_shape = features_info[self._input_key].shape
-    num_output_classes = features_info[label_key].num_classes
+    self._input_name = input_names[0]
+    input_shape = features_info[self._input_name].shape
+    num_output_classes = features_info[target_names[0]].num_classes
     self._num_hidden_layers = num_hidden_layers
     self._dense_layers = []
     self._dropout_layers = []
@@ -51,7 +48,7 @@ def __init__(self, features_info=None, supervised_keys=None,
         num_output_classes, activation="softmax")
 
   def call(self, inputs, training=False):
-    x = tf.cast(inputs[self._input_key], tf.float32) / 255.0
+    x = tf.cast(inputs[self._input_name], tf.float32) / 255.0
     x = self._flatten_layer(x)
     for i in range(self._num_hidden_layers):
       x = self._dense_layers[i](x)
@@ -73,20 +70,18 @@ def basic_fc_large():
 class BasicFcReluV2(tf.keras.Model):
   """Basic fully-connected + ReLU model, nicer code version."""
 
-  def __init__(self, features_info=None, supervised_keys=None,
+  def __init__(self, features_info=None, input_names=None, target_names=None,
                num_hidden_layers=2, hidden_size=64, dropout=0.1):
     super(BasicFcReluV2, self).__init__()
-    assert features_info is not None
-    assert supervised_keys is not None
-    self._input_key = supervised_keys[0]
-    self._input_shape = features_info[self._input_key].shape
-    self._num_output_classes = features_info[supervised_keys[1]].num_classes
+    self._input_name = input_names[0]
+    self._input_shape = features_info[self._input_name].shape
+    self._num_output_classes = features_info[target_names[0]].num_classes
     self._num_hidden_layers = num_hidden_layers
     self._dropout = dropout
     self._hidden_size = hidden_size
 
   def call(self, inputs, training=False):
-    x = tf.cast(inputs[self._input_key], tf.float32) / 255.0
+    x = tf.cast(inputs[self._input_name], tf.float32) / 255.0
     x = tf.keras.layers.Flatten(
         input_shape=self._input_shape)(x)
     for i in range(self._num_hidden_layers):
diff --git a/tensor2tensor/v2/models/resnet.py b/tensor2tensor/v2/models/resnet.py
index d4f6b3d89..33406ba73 100644
--- a/tensor2tensor/v2/models/resnet.py
+++ b/tensor2tensor/v2/models/resnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -29,19 +29,16 @@
 class Resnet(tf.keras.Model):
   """Resnet."""
 
-  def __init__(self, features_info=None, supervised_keys=None,
+  def __init__(self, features_info=None, input_names=None, target_names=None,
                layer_sizes=None, filter_sizes=None):
+    super(Resnet, self).__init__()
     # Base config for resnet-50.
     if layer_sizes is None:
       layer_sizes = [3, 4, 6, 3]
     if filter_sizes is None:
       filter_sizes = [64, 64, 128, 256, 512]
-    assert features_info is not None
-    assert supervised_keys is not None
-    super(Resnet, self).__init__()
-    self._input_key = supervised_keys[0]
-    label_key = supervised_keys[1]
-    num_output_classes = features_info[label_key].num_classes
+    self._input_name = input_names[0]
+    num_output_classes = features_info[target_names[0]].num_classes
 
     # Now the model.
     def resnet_model(inputs, training):
@@ -58,7 +55,7 @@ def resnet_model(inputs, training):
         num_output_classes, activation="softmax")
 
   def call(self, inputs, training=False):
-    x = tf.cast(inputs[self._input_key], tf.float32) / 255.0
+    x = tf.cast(inputs[self._input_name], tf.float32) / 255.0
     x = self._resnet(x, training)
     x = tf.reduce_mean(x, axis=[1, 2])
     return self._logits(x)
diff --git a/tensor2tensor/v2/t2t.py b/tensor2tensor/v2/t2t.py
index 422527eb3..739170384 100644
--- a/tensor2tensor/v2/t2t.py
+++ b/tensor2tensor/v2/t2t.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -55,7 +55,8 @@ def train_and_eval_dataset(dataset_name, data_dir):
      * the eval tf.Daataset
      * information about features: a python dictionary with feature names
          as keys and an object as value that provides .shape and .num_classes.
-     * supervised_keys: information what's the input and what's the target.
+     * supervised_keys: information what's the input and what's the target,
+         ie., a pair of lists with input and target feature names.
   """
   if dataset_name.startswith("v1_"):
     return _train_and_eval_dataset_v1(dataset_name[3:], data_dir)
@@ -71,7 +72,10 @@ def train_and_eval_dataset(dataset_name, data_dir):
     eval_split = tfds.Split.TEST
   train, valid = tfds.load(
       name=dataset_name, split=[tfds.Split.TRAIN, eval_split])
-  return train, valid, info.features, info.supervised_keys
+  keys = None
+  if info.supervised_keys:
+    keys = ([info.supervised_keys[0]], [info.supervised_keys[1]])
+  return train, valid, info.features, keys
 
 
 def _make_info(shape_list, num_classes):
@@ -94,7 +98,7 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
   problem = problems.problem(problem_name)
   train_dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, data_dir)
   eval_dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, data_dir)
-  supervised_keys = ("inputs", "targets")
+  supervised_keys = (["inputs"], ["targets"])
   hparams = problem.get_hparams()
   # We take a few training examples to guess the shapes.
   input_shapes, target_shapes = [], []
@@ -109,10 +113,21 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
   return train_dataset, eval_dataset, info, supervised_keys
 
 
-def shuffle_and_batch_data(dataset, batch_size, target_key):
+def shuffle_and_batch_data(dataset, batch_size, target_names, repeat=False):
   """Shuffle and batch the given dataset."""
+  def append_targets(example):
+    """Append targets to the example dictionary. Needed for Keras."""
+    if len(target_names) == 1:
+      return (example, example[target_names[0]])
+    targets = {}
+    for name in target_names:
+      targets[name] = example[name]
+    return (example, targets)
+  dataset = dataset.map(append_targets)
+  if repeat:
+    dataset = dataset.repeat()
   shuffled = dataset.shuffle(128).batch(batch_size).prefetch(8)
-  return shuffled.map(lambda ex: (ex, ex[target_key]))
+  return shuffled
 
 
 @gin.configurable(blacklist=["model"])
@@ -132,6 +147,7 @@ def model_compile(model,
 @gin.configurable(blacklist=["data_dir", "output_dir"])
 def train_fn(data_dir=None, output_dir=None,
              model_class=gin.REQUIRED, dataset=gin.REQUIRED,
+             input_names=None, target_names=None,
              batch_size=32, train_steps=1000, eval_steps=1, eval_frequency=100):
   """Train the given model on the given dataset.
 
@@ -140,6 +156,8 @@ def train_fn(data_dir=None, output_dir=None,
     output_dir: Directory where to put the logs and checkpoints.
     model_class: The model class to train.
     dataset: The name of the dataset to train on.
+    input_names: List of strings with the names of the features on input.
+    target_names: List of strings with the names of the target features.
     batch_size: integer, how many examples per batch.
     train_steps: for how many steps to train.
     eval_steps: for how many steps to do evaluation.
@@ -147,10 +165,22 @@ def train_fn(data_dir=None, output_dir=None,
   """
   train_data, eval_data, features_info, keys = train_and_eval_dataset(
       dataset, data_dir)
-  model = model_class(features_info=features_info, supervised_keys=keys)
+  if input_names is None:
+    input_names = keys[0]
+  if target_names is None:
+    target_names = keys[1]
+  # TODO(lukaszkaiser): The use of distribution strategy below fails like this:
+  #   .../keras/models.py", line 93, in _clone_functional_model
+  #      for layer in model._input_layers:
+  #   AttributeError: 'BasicFcRelu' object has no attribute '_input_layers'
+  # strategy = tf.distribute.MirroredStrategy()
+  # with strategy.scope():
+  model = model_class(features_info=features_info,
+                      input_names=input_names, target_names=target_names)
   model_compile(model)
-  train_batches = shuffle_and_batch_data(train_data, batch_size, keys[1])
-  eval_batches = shuffle_and_batch_data(eval_data, batch_size, keys[1])
+  train_batches = shuffle_and_batch_data(
+      train_data, batch_size, target_names, repeat=True)
+  eval_batches = shuffle_and_batch_data(eval_data, batch_size, target_names)
 
   # Training loop.
   callbacks = []
@@ -158,10 +188,9 @@ def train_fn(data_dir=None, output_dir=None,
   callbacks.append(tf.keras.callbacks.BaseLogger())
   if output_dir is not None:
     callbacks.append(tf.keras.callbacks.TensorBoard(log_dir=output_dir))
-    # TODO(lukaszkaiser): the one below doesn't seem to work, why?
-    # callbacks.append(tf.keras.callbacks.ModelCheckpoint(
-    #     filepath=output_dir))
-  model.fit(train_batches.repeat(),
+    callbacks.append(tf.keras.callbacks.ModelCheckpoint(
+        filepath=output_dir, save_weights_only=True))
+  model.fit(train_batches,
             epochs=train_steps // eval_frequency,
             steps_per_epoch=eval_frequency,
             validation_data=eval_batches,
diff --git a/tensor2tensor/v2/t2t_trainer.py b/tensor2tensor/v2/t2t_trainer.py
index aa9abe8f1..9153b0978 100644
--- a/tensor2tensor/v2/t2t_trainer.py
+++ b/tensor2tensor/v2/t2t_trainer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/__init__.py b/tensor2tensor/visualization/__init__.py
index c2a9550b0..b775a72bd 100644
--- a/tensor2tensor/visualization/__init__.py
+++ b/tensor2tensor/visualization/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/attention.py b/tensor2tensor/visualization/attention.py
index f4f78f3f6..739edf79e 100644
--- a/tensor2tensor/visualization/attention.py
+++ b/tensor2tensor/visualization/attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/visualization.py b/tensor2tensor/visualization/visualization.py
index c70590d98..0fae8722e 100644
--- a/tensor2tensor/visualization/visualization.py
+++ b/tensor2tensor/visualization/visualization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/visualization_test.py b/tensor2tensor/visualization/visualization_test.py
index 9c1371a7d..8ad86d107 100644
--- a/tensor2tensor/visualization/visualization_test.py
+++ b/tensor2tensor/visualization/visualization_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The Tensor2Tensor Authors.
+# Copyright 2019 The Tensor2Tensor Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

From f29aa09292efe1dc5f70ee35d6a43e7626136f1d Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Thu, 7 Feb 2019 14:48:12 -0800
Subject: [PATCH 1636/2720] Work on datasets for one giant model. - Add
 LanguagemodelMultiWikiTranslatePacked1k:    a multi-text2text problem with
 packed examples suitable for training on TPU. - Added support for appending
 specifying string prefixes to the input and output of   a Text2TextProblem. 
 We can specify the problem in the input instead of using a   separate id
 (simplifies the interface, supports compositional problem descriptions). -
 Sequence-to-sequence fill-in-the-blank version of wiki_lm problem. - support
 packed examples in multi_problem_v2 - Avoid using pyfunc to compute a
 constant mixing schedule (pyfunc doesn't work on TPU) - Added a convenience
 function to specify mixing fraction in terms of epoch rates   (ratios get
 multiplied by sizes of training sets) - text2text versions of MNLI and SQUAD

Mesh-Tensorflow Transformer:
   Use correct positions for packed examples.
   Use input id 0 instead of EOS for first position of later sequences
     in a packed example in the autoregressive decoder.

PiperOrigin-RevId: 232949422
---
 .../data_generators/cnn_dailymail.py          |  41 ++++-
 .../data_generators/generator_utils.py        |  39 +++++
 .../data_generators/multi_problem_v2.py       |  80 ++++++++-
 tensor2tensor/data_generators/multinli.py     | 147 ++++++++++++-----
 tensor2tensor/data_generators/problem.py      |   5 +
 tensor2tensor/data_generators/squad.py        | 152 ++++++++++++------
 .../data_generators/text_problems.py          |  26 ++-
 .../data_generators/translate_ende.py         |  14 ++
 .../data_generators/translate_enfr.py         |  12 ++
 .../data_generators/translate_enro.py         |  14 +-
 tensor2tensor/data_generators/wiki_lm.py      |  39 ++++-
 .../data_generators/wiki_multi_problems.py    |  47 +++++-
 tensor2tensor/models/mtf_transformer2.py      |  26 ++-
 13 files changed, 524 insertions(+), 118 deletions(-)

diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index cbd9d84cf..5e77b5599 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -250,8 +250,8 @@ class SummarizeCnnDailymailWikiLMSharedVocab(SummarizeCnnDailymail32k):
   """Summarize CNN and Daily Mail articles using the Wiki 32k vocab."""
 
   @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelEnWiki32k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelEnWiki32k()
 
 
 @registry.register_problem
@@ -259,8 +259,8 @@ class SummarizeCnnDailymailWikiLMSharedVocab64k(SummarizeCnnDailymail32k):
   """Summarize CNN and Daily Mail articles using the Wiki 64k vocab."""
 
   @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelEnWiki64k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelEnWiki64k()
 
 
 @registry.register_problem
@@ -268,8 +268,33 @@ class SummarizeCnnDailymailWikiLMMultiVocab64k(SummarizeCnnDailymail32k):
   """Summarize CNN and Daily Mail articles using multi-lingual 64k vocab."""
 
   @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
+
+
+@registry.register_problem
+class SummarizeCnnDailymailMulti64kPacked1k(SummarizeCnnDailymail32k):
+  """Summarize CNN and Daily Mail articles using multi-lingual 64k vocab."""
+
+  @property
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
+
+  @property
+  def packed_length(self):
+    return 1024
+
+  @property
+  def num_training_examples(self):
+    return 252600
+
+  @property
+  def inputs_prefix(self):
+    return "CNN Daily Mail article to summary "
+
+  @property
+  def targets_prefix(self):
+    return "CNN Daily Mail summary to article "
 
 
 @registry.register_problem
@@ -277,8 +302,8 @@ class SummarizeFracCnnDailymailWikiLMSharedVocab64k(SummarizeCnnDailymail32k):
   """Summarize a fraction of CNN/DM articles using the Wiki 64k vocab."""
 
   @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelEnWiki64k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelEnWiki64k()
 
   def fraction_of_data(self):
     return 1.
diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 30ffa9f5d..1fd616c19 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import gzip
+import math
 import os
 import random
 import stat
@@ -720,3 +721,41 @@ def _parse_example(ex_ser):
           yield ex
         except tf.errors.OutOfRangeError:
           break
+
+
+def random_deinterleave(text, separator_symbol="X"):
+  """Create a fill-in-the-blanks training example from text.
+
+  Split on spaces, then cut into segments at random points.  Alternate segments
+  are assigned to the two output strings. separator_symbol separates segments
+  within each of the outputs.
+
+  example:
+    text="The quick brown fox jumps over the lazy dog."
+    returns: ("X quick brown X the lazy X", "The X fox jumps over X dog.")
+
+  The two outputs can also be reversed to yield an instance of the same problem.
+
+  Args:
+    text: a string
+    separator_symbol: a string
+  Returns:
+    a pair of strings
+  """
+  words = text.strip().split(" ")
+  n = len(words)
+  if n <= 1:
+    return text, ""
+  cut = [False] * n
+  cut[0] = True
+  num_cuts = int(math.exp(random.uniform(0, math.log(n))))
+  for _ in xrange(num_cuts):
+    cut[random.randint(1, n -1)] = True
+  out = [[], []]
+  part = random.randint(0, 1)
+  for i in xrange(n):
+    if cut[i]:
+      out[part].append(separator_symbol)
+      part = 1 - part
+    out[part].append(words[i])
+  return " ".join(out[0]), " ".join(out[1])
diff --git a/tensor2tensor/data_generators/multi_problem_v2.py b/tensor2tensor/data_generators/multi_problem_v2.py
index 004aa1cb8..63a6cddf0 100644
--- a/tensor2tensor/data_generators/multi_problem_v2.py
+++ b/tensor2tensor/data_generators/multi_problem_v2.py
@@ -139,21 +139,45 @@ class MultiText2TextProblem(MultiProblemV2, text_problems.Text2TextProblem):
   def normalize_example(self, example, hparams):
     """Assumes that example contains both inputs and targets."""
 
+    length = self.max_length(hparams)
     def _to_constant_shape(tensor):
-      max_length = self.max_length(hparams)
-      tensor = tensor[:max_length]
-      tensor = tf.pad(tensor, [(0, max_length - tf.shape(tensor)[0])])
-      return tf.reshape(tensor, [max_length])
+      tensor = tensor[:length]
+      tensor = tf.pad(tensor, [(0, length - tf.shape(tensor)[0])])
+      return tf.reshape(tensor, [length])
 
     if self.has_inputs:
       example['inputs'] = _to_constant_shape(example['inputs'])
       example['targets'] = _to_constant_shape(example['targets'])
     elif 'inputs' in example:
+      if self.packed_length:
+        raise ValueError('cannot concatenate packed examples on the fly.')
       inputs = example.pop('inputs')[:-1]  # Remove EOS token.
       targets = tf.concat([inputs, example['targets']], 0)
       example['targets'] = _to_constant_shape(targets)
     else:
       example['targets'] = _to_constant_shape(example['targets'])
+    if self.packed_length:
+      if self.has_inputs:
+        if 'inputs_segmentation' in example:
+          example['inputs_segmentation'] = _to_constant_shape(
+              example['inputs_segmentation'])
+          example['inputs_position'] = _to_constant_shape(
+              example['inputs_position'])
+        else:
+          example['inputs_segmentation'] = tf.to_int64(
+              tf.not_equal(example['inputs'], 0))
+          example['inputs_position'] = (
+              example['inputs_segmentation'] * tf.range(length, dtype=tf.int64))
+      if 'targets_segmentation' in example:
+        example['targets_segmentation'] = _to_constant_shape(
+            example['targets_segmentation'])
+        example['targets_position'] = _to_constant_shape(
+            example['targets_position'])
+      else:
+        example['targets_segmentation'] = tf.to_int64(
+            tf.not_equal(example['targets'], 0))
+        example['targets_position'] = (
+            example['targets_segmentation'] * tf.range(length, dtype=tf.int64))
     return example
 
   def generate_data_with_shared_vocab(self, data_dir, tmp_dir, task_id=-1):
@@ -172,6 +196,11 @@ def generate_data_with_shared_vocab(self, data_dir, tmp_dir, task_id=-1):
         tf.gfile.Copy(global_vocab_filename, local_vocab_filename)
       p.generate_data(data_dir, tmp_dir, task_id)
 
+  @property
+  def packed_length(self):
+    """Set this to a positive integer if some of the problems are packed."""
+    return None
+
 
 def get_multi_dataset(datasets, pmf=None):
   """Returns a Dataset that samples records from one or more Datasets.
@@ -205,6 +234,11 @@ def get_schedule_distribution(schedule, global_step=None):
     A 1-D tensor of probs, the sampling distribution of the global_step.
   """
   interpolation, steps, pmfs = schedule
+  if len(pmfs) == 1:
+    # py_func doesn't seem to work on TPU - at least get the constant case to
+    # run.
+    # TODO(noam): get the general case working.
+    return pmfs[0]
   if global_step is None:
     global_step = tf.train.get_or_create_global_step()
   if interpolation == 'step':
@@ -304,6 +338,44 @@ def constant_schedule(pmf):
   return ('step', (0,), (tuplize(pmf),))
 
 
+def example_rates_to_pmf(example_rates):
+  """Creates a probability-mass-function based on relative example rates.
+
+  Args:
+    example_rates: a list or tuple
+  Returns:
+    a list of floats
+  """
+  total = sum(example_rates)
+  return [r / total for r in example_rates]
+
+
+def epoch_rates_to_pmf(problems, epoch_rates=None):
+  """Create a probability-mass-function based on relative epoch rates.
+
+  if epoch_rates=None, then we use uniform epoch rates [1.0] * len(problems)
+  i.e. it takes each problem the same time to go through one epoch.
+
+  If epoch_rates is given, then these are the relative numbers of epochs
+  of each problem to go through in a given amount of time.
+
+  Each must have problem.num_training_examples implemented.
+
+  Args:
+    problems: a list of Problem instances.
+    epoch_rates: an optional list of float
+
+  Returns:
+    a list of floating point values.
+  """
+  if epoch_rates is None:
+    epoch_rates = [1.0] * len(problems)
+  example_rates = []
+  for p, epoch_rate in zip(problems, epoch_rates):
+    example_rates.append(epoch_rate * p.num_training_examples)
+  return example_rates_to_pmf(example_rates)
+
+
 def encode_schedule(schedule):
   """Encodes a schedule tuple into a string.
 
diff --git a/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index f3f6d3d92..29f236e44 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -32,17 +32,57 @@
 
 EOS = text_encoder.EOS
 
+# Link to data from GLUE: https://gluebenchmark.com/tasks
+_MNLI_URL = ("https://firebasestorage.googleapis.com/v0/b/"
+             "mtl-sentence-representations.appspot.com/o/"
+             "data%2FMNLI.zip?alt=media&token=50329ea1-e339-"
+             "40e2-809c-10c40afff3ce")
+
+
+def _maybe_download_corpora(tmp_dir):
+  """Download corpora for multinli.
+
+  Args:
+    tmp_dir: a string
+  Returns:
+    a string
+  """
+  mnli_filename = "MNLI.zip"
+  mnli_finalpath = os.path.join(tmp_dir, "MNLI")
+  if not tf.gfile.Exists(mnli_finalpath):
+    zip_filepath = generator_utils.maybe_download(
+        tmp_dir, mnli_filename, _MNLI_URL)
+    zip_ref = zipfile.ZipFile(zip_filepath, "r")
+    zip_ref.extractall(tmp_dir)
+    zip_ref.close()
+
+  return mnli_finalpath
+
+
+def _example_generator(filename):
+  """Generate mnli examples.
+
+  Args:
+    filename: a string
+  Yields:
+    dictionaries containing "premise", "hypothesis" and "label" strings
+  """
+  for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
+    if idx == 0: continue  # skip header
+    line = text_encoder.to_unicode_utf8(line.strip())
+    split_line = line.split("\t")
+    # Works for both splits even though dev has some extra human labels.
+    yield {
+        "premise": split_line[8],
+        "hypothesis": split_line[9],
+        "label": split_line[-1]
+    }
+
 
 @registry.register_problem
 class MultiNLI(text_problems.TextConcat2ClassProblem):
   """MultiNLI classification problems."""
 
-  # Link to data from GLUE: https://gluebenchmark.com/tasks
-  _MNLI_URL = ("https://firebasestorage.googleapis.com/v0/b/"
-               "mtl-sentence-representations.appspot.com/o/"
-               "data%2FMNLI.zip?alt=media&token=50329ea1-e339-"
-               "40e2-809c-10c40afff3ce")
-
   @property
   def is_generate_per_split(self):
     return True
@@ -70,46 +110,69 @@ def class_labels(self, data_dir):
     # Note this binary classification is different from usual MNLI.
     return ["contradiction", "entailment", "neutral"]
 
-  def _maybe_download_corpora(self, tmp_dir):
-    mnli_filename = "MNLI.zip"
-    mnli_finalpath = os.path.join(tmp_dir, "MNLI")
-    if not tf.gfile.Exists(mnli_finalpath):
-      zip_filepath = generator_utils.maybe_download(
-          tmp_dir, mnli_filename, self._MNLI_URL)
-      zip_ref = zipfile.ZipFile(zip_filepath, "r")
-      zip_ref.extractall(tmp_dir)
-      zip_ref.close()
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    mnli_dir = _maybe_download_corpora(tmp_dir)
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      filesplit = ["train.tsv"]
+    else:
+      # Using dev matched as the default for eval. Can also switch this to
+      # dev_mismatched.tsv
+      filesplit = ["dev_matched.tsv"]
+    label_list = self.class_labels(data_dir=None)
+    for fs in filesplit:
+      filename = os.path.join(mnli_dir, fs)
+      for example in _example_generator(filename):
+        yield {
+            "inputs": [example["premise"], example["hypothesis"]],
+            "label": label_list.index(example["label"])
+        }
 
-    return mnli_finalpath
 
-  def example_generator(self, filename):
-    label_list = self.class_labels(data_dir=None)
-    for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
-      if idx == 0: continue  # skip header
-      line = text_encoder.to_unicode_utf8(line.strip())
-      split_line = line.split("\t")
-      # Works for both splits even though dev has some extra human labels.
-      s1, s2 = split_line[8:10]
-      l = label_list.index(split_line[-1])
-      inputs = [s1, s2]
-      yield {
-          "inputs": inputs,
-          "label": l
-      }
+@registry.register_problem
+class MultiNLIText2text(text_problems.Text2TextProblem):
+  """MultiNLI classification problems."""
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def approx_vocab_size(self):
+    return 2**15
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    mnli_dir = self._maybe_download_corpora(tmp_dir)
+    mnli_dir = _maybe_download_corpora(tmp_dir)
     if dataset_split == problem.DatasetSplit.TRAIN:
       filesplit = ["train.tsv"]
     else:
       # Using dev matched as the default for eval. Can also switch this to
       # dev_mismatched.tsv
       filesplit = ["dev_matched.tsv"]
-
     for fs in filesplit:
       filename = os.path.join(mnli_dir, fs)
-      for example in self.example_generator(filename):
-        yield example
+      for example in _example_generator(filename):
+        yield {
+            "inputs": "multinli premise: %s hypothesis: %s" % (
+                example["premise"], example["hypothesis"]),
+            "targets": example["label"]
+        }
+
+
+@registry.register_problem
+class MultiNLIText2textMulti64kPacked1k(MultiNLIText2text):
+  """MultiNLI classification problems with the multi-lingual vocabulary."""
+
+  @property
+  def packed_length(self):
+    return 1024
+
+  @property
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
+
+  @property
+  def num_training_examples(self):
+    return 18300
 
 
 @registry.register_problem
@@ -129,8 +192,8 @@ class MultiNLISharedVocab(MultiNLI):
   """MultiNLI classification problems with the LM1b vocabulary."""
 
   @property
-  def vocab_filename(self):
-    return lm1b.LanguagemodelLm1b32k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return lm1b.LanguagemodelLm1b32k()
 
 
 @registry.register_problem
@@ -138,8 +201,8 @@ class MultiNLIWikiLMSharedVocab(MultiNLI):
   """MultiNLI classification problems with the Wiki vocabulary."""
 
   @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelEnWiki32k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelEnWiki32k()
 
 
 @registry.register_problem
@@ -147,8 +210,8 @@ class MultiNLIWikiLMSharedVocab64k(MultiNLIWikiLMSharedVocab):
   """MultiNLI classification problems with the Wiki vocabulary."""
 
   @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelEnWiki64k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelEnWiki64k()
 
 
 @registry.register_problem
@@ -156,5 +219,5 @@ class MultiNLIWikiLMMultiVocab64k(MultiNLIWikiLMSharedVocab):
   """MultiNLI classification problems with the multi-lingual vocabulary."""
 
   @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 4d8a972f0..6c651ba74 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -234,6 +234,11 @@ def num_generate_tasks(self):
     """Needed if multiprocess_generate is True."""
     raise NotImplementedError()
 
+  @property
+  def num_training_examples(self):
+    """Used when mixing problems - how many examples are in the dataset."""
+    raise NotImplementedError()
+
   def prepare_to_generate(self, data_dir, tmp_dir):
     """Prepare to generate data in parallel on different processes.
 
diff --git a/tensor2tensor/data_generators/squad.py b/tensor2tensor/data_generators/squad.py
index 82105e013..dc9117f6b 100644
--- a/tensor2tensor/data_generators/squad.py
+++ b/tensor2tensor/data_generators/squad.py
@@ -31,14 +31,99 @@
 import tensorflow as tf
 
 
+_DEV_SET = "dev-v1.1.json"
+_URL = "https://rajpurkar.github.io/SQuAD-explorer/dataset"
+_TRAINING_SET = "train-v1.1.json"
+
+
+def _generate_examples(tmp_dir, dataset_split):
+  """Generate squad examples.
+
+  Args:
+    tmp_dir: a string
+    dataset_split: problem.DatasetSplit.TRAIN or problem.DatasetSplit.EVAL
+  Yields:
+    dictionaries representing examples
+  """
+  if dataset_split == problem.DatasetSplit.TRAIN:
+    file_name = _TRAINING_SET
+  else:
+    file_name = _DEV_SET
+  squad_file = generator_utils.maybe_download(tmp_dir,
+                                              file_name,
+                                              os.path.join(_URL, file_name))
+  with tf.gfile.GFile(squad_file, mode="r") as fp:
+    squad = json.load(fp)
+
+  version = squad["version"]
+  for article in squad["data"]:
+    if "title" in article:
+      title = article["title"].strip()
+    else:
+      title = "no title"
+    for paragraph in article["paragraphs"]:
+      context = paragraph["context"].strip()
+      for qa in paragraph["qas"]:
+        question = qa["question"].strip()
+        id_ = qa["id"]
+        answer_starts = [answer["answer_start"] for answer in qa["answers"]]
+        answers = [answer["text"].strip() for answer in qa["answers"]]
+
+        # Features currently used are "context", "question", and "answers".
+        # Others are extracted here for the ease of future expansions.
+        example = {
+            "version": version,
+            "title": title,
+            "context": context,
+            "question": question,
+            "id": id_,
+            "answer_starts": answer_starts,
+            "answers": answers,
+            "num_answers": len(answers),
+            "is_supervised": True,
+        }
+        yield example
+
+
+@registry.register_problem
+class SquadText2text(text_problems.Text2TextProblem):
+  """Squad as a Text2TextProblem."""
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    for example in _generate_examples(tmp_dir, dataset_split):
+      yield {
+          "inputs": "squad context: %s question: %s" % (
+              example["context"], example["question"]),
+          # TODO(ddohan, wgaj): Figure out a way of extracting all answers.
+          "targets": example["answers"][0],
+      }
+
+
+@registry.register_problem
+class SquadText2textMulti64kPacked1k(SquadText2text):
+  """Squad with multi-lingual vocabulary."""
+
+  @property
+  def packed_length(self):
+    return 1024
+
+  @property
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
+
+  @property
+  def num_training_examples(self):
+    return 16300
+
+
 @registry.register_problem
 class Squad(text_problems.QuestionAndContext2TextProblem):
   """Base class for SquAD question answering problem."""
 
-  _DEV_SET = "dev-v1.1.json"
-  _URL = "https://rajpurkar.github.io/SQuAD-explorer/dataset"
-  _TRAINING_SET = "train-v1.1.json"
-
   @property
   def dataset_splits(self):
     return [{
@@ -54,50 +139,13 @@ def is_generate_per_split(self):
     return True
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    url = self._URL
-    file_name = self._DEV_SET
-    if dataset_split == problem.DatasetSplit.TRAIN:
-      file_name = self._TRAINING_SET
-    squad_file = generator_utils.maybe_download(tmp_dir,
-                                                file_name,
-                                                os.path.join(url, file_name))
-    with tf.gfile.GFile(squad_file, mode="r") as fp:
-      squad = json.load(fp)
-
-    version = squad["version"]
-    for article in squad["data"]:
-      if "title" in article:
-        title = article["title"].strip()
-      else:
-        title = "no title"
-      for paragraph in article["paragraphs"]:
-        context = paragraph["context"].strip()
-        for qa in paragraph["qas"]:
-          question = qa["question"].strip()
-          id_ = qa["id"]
-
-          answer_starts = [answer["answer_start"] for answer in qa["answers"]]
-          answers = [answer["text"].strip() for answer in qa["answers"]]
-
-          # Features currently used are "context", "question", and "answers".
-          # Others are extracted here for the ease of future expansions.
-          example = {
-              "version": version,
-              "title": title,
-              "context": context,
-              "question": question,
-              "id": id_,
-              "answer_starts": answer_starts,
-              "answers": answers,
-              "num_answers": len(answers),
-              "is_supervised": True,
-          }
-          yield {
-              "inputs": example["question"],
-              # TODO(ddohan, wgaj): Figure out a way of extracting all answers.
-              "targets": example["answers"][0],
-              "context": example["context"]
-          }
+    for example in _generate_examples(tmp_dir, dataset_split):
+      yield {
+          "inputs": example["question"],
+          # TODO(ddohan, wgaj): Figure out a way of extracting all answers.
+          "targets": example["answers"][0],
+          "context": example["context"]
+      }
 
 
 @registry.register_problem
@@ -148,8 +196,8 @@ def dataset_filename(self):
     return "squad_multi64k"
 
   @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
 
 
 @registry.register_problem
@@ -160,8 +208,8 @@ def dataset_filename(self):
     return "squad"
 
   @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelEnWiki32k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelEnWiki32k()
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 9eeba462e..0a5fd66ac 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -306,7 +306,9 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
     generator = self.generate_samples(data_dir, tmp_dir, dataset_split)
     encoder = self.get_or_create_vocab(data_dir, tmp_dir)
     return text2text_generate_encoded(generator, encoder,
-                                      has_inputs=self.has_inputs)
+                                      has_inputs=self.has_inputs,
+                                      inputs_prefix=self.inputs_prefix,
+                                      targets_prefix=self.targets_prefix)
 
   @property
   def max_subtoken_length(self):
@@ -328,6 +330,16 @@ def batch_size_means_tokens(self):
   def already_shuffled(self):
     return False
 
+  @property
+  def inputs_prefix(self):
+    """String to prepend to inputs before tokenization."""
+    return ""
+
+  @property
+  def targets_prefix(self):
+    """String to prepend to targets before tokenization."""
+    return ""
+
   def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
     filepath_fns = {
@@ -665,14 +677,16 @@ def text2text_txt_tab_iterator(txt_path):
 def text2text_generate_encoded(sample_generator,
                                vocab,
                                targets_vocab=None,
-                               has_inputs=True):
+                               has_inputs=True,
+                               inputs_prefix="",
+                               targets_prefix=""):
   """Encode Text2Text samples from the generator with the vocab."""
   targets_vocab = targets_vocab or vocab
   for sample in sample_generator:
     if has_inputs:
-      sample["inputs"] = vocab.encode(sample["inputs"])
+      sample["inputs"] = vocab.encode(inputs_prefix + sample["inputs"])
       sample["inputs"].append(text_encoder.EOS_ID)
-    sample["targets"] = targets_vocab.encode(sample["targets"])
+    sample["targets"] = targets_vocab.encode(targets_prefix + sample["targets"])
     sample["targets"].append(text_encoder.EOS_ID)
     yield sample
 
@@ -1195,7 +1209,9 @@ def generate_encoded_samples(self,
     generator = self.generate_samples(data_dir, tmp_dir, dataset_split,
                                       input_files)
     return text2text_generate_encoded(
-        generator, encoder, has_inputs=self.has_inputs)
+        generator, encoder, has_inputs=self.has_inputs,
+        inputs_prefix=self.inputs_prefix,
+        targets_prefix=self.targets_prefix)
 
   def generate_data(self, data_dir, tmp_dir, task_id=-1):
     # task_id should be in [0, self.num_output_shards)
diff --git a/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index d49a4f0b0..ac7718726 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -140,3 +140,17 @@ class TranslateEndeWmtMulti64kPacked1k(TranslateEndeWmtMulti64k):
   @property
   def packed_length(self):
     return 1024
+
+  @property
+  def num_training_examples(self):
+    return 173800
+
+  @property
+  def inputs_prefix(self):
+    return "translate English German "
+
+  @property
+  def targets_prefix(self):
+    return "translate German English "
+
+
diff --git a/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index fe4e31e80..4484b5c4e 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -259,3 +259,15 @@ class TranslateEnfrWmtMulti64kPacked1k(TranslateEnfrWmtMulti64k):
   @property
   def packed_length(self):
     return 1024
+
+  @property
+  def num_training_examples(self):
+    return 1760600
+
+  @property
+  def inputs_prefix(self):
+    return "translate English French "
+
+  @property
+  def targets_prefix(self):
+    return "translate French English "
diff --git a/tensor2tensor/data_generators/translate_enro.py b/tensor2tensor/data_generators/translate_enro.py
index e34fe7566..f022f638d 100644
--- a/tensor2tensor/data_generators/translate_enro.py
+++ b/tensor2tensor/data_generators/translate_enro.py
@@ -141,7 +141,7 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
 
 @registry.register_problem
 class TranslateEnroWmtMultiTiny64k(TranslateEnroWmtMultiSmall64k):
-  """Translation with muli-lingual vocabulary, tiny (6K) training data."""
+  """Translation with muli-lingual vocabulary, tiny (600) training data."""
 
   @property
   def how_many_examples_to_sample(self):
@@ -155,3 +155,15 @@ class TranslateEnroWmtMultiTiny64kPacked1k(TranslateEnroWmtMultiTiny64k):
   @property
   def packed_length(self):
     return 1024
+
+  @property
+  def num_training_examples(self):
+    return 32
+
+  @property
+  def inputs_prefix(self):
+    return "translate English Romanian "
+
+  @property
+  def targets_prefix(self):
+    return "translate Romanian English "
diff --git a/tensor2tensor/data_generators/wiki_lm.py b/tensor2tensor/data_generators/wiki_lm.py
index 37cb3b79b..1e6dd3ca9 100644
--- a/tensor2tensor/data_generators/wiki_lm.py
+++ b/tensor2tensor/data_generators/wiki_lm.py
@@ -35,7 +35,7 @@ def concat_generator(filename, up_threshold, low_threshold=10):
   txt = ""
   for line in tf.gfile.Open(filename):
     line = line.strip()
-    if len(txt) + len(line) > up_threshold:
+    if len(txt) + len(line) + 1 >= up_threshold:
       ret = txt
       txt = ""
       # We don't yield very short long parts to prevent noisy examples.
@@ -238,3 +238,40 @@ def approx_vocab_size(self):
   @property
   def max_samples_for_vocab(self):
     return 256000  # Samples are intertwined, take more to cover 4 languages.
+
+
+@registry.register_problem
+class LanguagemodelDeEnFrRoWiki64kFitbPacked1k(
+    LanguagemodelDeEnFrRoWiki64k):
+  """4 languages fill-in-the-blanks text-to-text problem."""
+
+  @property
+  def use_vocab_from_other_problem(self):
+    return LanguagemodelDeEnFrRoWiki64k()
+
+  @property
+  def has_inputs(self):
+    return True
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    for example in super(
+        LanguagemodelDeEnFrRoWiki64kFitbPacked1k, self).generate_samples(
+            data_dir, tmp_dir, dataset_split):
+      a, b = generator_utils.random_deinterleave(example["targets"])
+      yield {"inputs": a, "targets": b}
+
+  @property
+  def num_training_examples(self):
+    return 3597800
+
+  @property
+  def packed_length(self):
+    return 1024
+
+  @property
+  def inputs_prefix(self):
+    return "wiki fill "
+
+  @property
+  def targets_prefix(self):
+    return "wiki fill "
diff --git a/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index 37e0a4204..49d248749 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -66,14 +66,57 @@ def has_inputs(self):
     return False
 
   @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelEnWiki32k().vocab_filename
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelEnWiki32k()
 
   @property
   def vocab_type(self):
     return text_problems.VocabType.SUBWORD
 
 
+@registry.register_problem
+class LanguagemodelMultiWikiTranslatePacked1k(
+    multi_problem_v2.MultiText2TextProblem):
+  """Wiki-LM, Translation, MNLI, SQUAD mixed problem class."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    problems = [
+        # TODO(noam): uncommonet once data is generated
+        wiki_lm.LanguagemodelDeEnFrRoWiki64kFitbPacked1k(),
+        wiki_lm.LanguagemodelDeEnFrRoWiki64kFitbPacked1k(was_reversed=True),
+        translate_ende.TranslateEndeWmtMulti64kPacked1k(),
+        translate_ende.TranslateEndeWmtMulti64kPacked1k(was_reversed=True),
+        translate_enfr.TranslateEnfrWmtMulti64kPacked1k(),
+        translate_enfr.TranslateEnfrWmtMulti64kPacked1k(was_reversed=True),
+        translate_enro.TranslateEnroWmtMultiTiny64kPacked1k(),
+        translate_enro.TranslateEnroWmtMultiTiny64kPacked1k(was_reversed=True),
+        cnn_dailymail.SummarizeCnnDailymailMulti64kPacked1k(),
+        cnn_dailymail.SummarizeCnnDailymailMulti64kPacked1k(was_reversed=True),
+        multinli.MultiNLIText2textMulti64kPacked1k(),
+        squad.SquadText2textMulti64kPacked1k(),
+    ]
+    schedule = multi_problem_v2.constant_schedule(
+        multi_problem_v2.epoch_rates_to_pmf(problems))
+    super(LanguagemodelMultiWikiTranslatePacked1k, self).__init__(
+        problems, schedule, was_reversed=was_reversed, was_copy=was_copy)
+
+  @property
+  def has_inputs(self):
+    return True
+
+  @property
+  def use_vocab_from_other_problem(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
+
+  @property
+  def packed_length(self):
+    return 1024
+
+
 @registry.register_problem
 class LanguagemodelEnWikiLMMultiNLISubwords64k(multi_problem.MultiProblem):
   """Wiki LM and MNLI mixed problem class."""
diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index 6ea437901..c279ee034 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -135,13 +135,17 @@ def _mtf_model_fn(self, features, mesh):
     def import_feature(key):
       return self._import_feature(features, mesh, key)
     targets = import_feature("targets")
+    sequence_id = import_feature("targets_segmentation")
+    position = import_feature("targets_position")
     if self.autoregressive:
       inputs = mtf.shift(
           targets, offset=1, dim=self.length_dim, wrap=False)
+      if position is not None:
+        # first input in later sequences should be 0
+        inputs *= mtf.to_int32(mtf.not_equal(position, 0))
     else:
       inputs = import_feature("inputs")
       # TODO(noam): options for bert-style masking here?
-    sequence_id = import_feature("targets_segmentation")
     model = self.model()
     logits, loss = model.call_simple(
         inputs=inputs,
@@ -149,7 +153,8 @@ def import_feature(key):
         compute_loss=True,
         mode=hparams.mode,
         variable_dtype=self.variable_dtype,
-        sequence_id=sequence_id)
+        sequence_id=sequence_id,
+        position=position)
     return logits, loss
 
   def mtf_model_fn(self, features, mesh):
@@ -243,6 +248,8 @@ def import_feature(key):
     decoder_sequence_id = import_feature("targets_segmentation")
     if decoder_sequence_id is None:
       decoder_sequence_id = mtf.to_int32(mtf.not_equal(targets, 0))
+    encoder_position = import_feature("inputs_position")
+    decoder_position = import_feature("targets_position")
     model = self.model()
     logits, loss = model.call_simple(
         inputs=inputs,
@@ -251,7 +258,9 @@ def import_feature(key):
         mode=hparams.mode,
         variable_dtype=self.variable_dtype,
         encoder_sequence_id=encoder_sequence_id,
-        decoder_sequence_id=decoder_sequence_id)
+        decoder_sequence_id=decoder_sequence_id,
+        encoder_position=encoder_position,
+        decoder_position=decoder_position)
     return logits, loss
 
   def sample(self, features, mesh):
@@ -461,6 +470,8 @@ def mtf_bitransformer_base():
   #        decode_length_multiplier * input_length + decode_length_constant)
   hparams.add_hparam("decode_length_multiplier", 1.5)
   hparams.add_hparam("decode_length_constant", 10.0)
+  # used during decoding
+  hparams.add_hparam("alpha", 0.6)
   hparams.sampling_temp = 0.0
   return hparams
 
@@ -826,3 +837,12 @@ def mtr_tr_ende_deep():
   hparams.encoder_num_layers = 12
   hparams.decoder_num_layers = 12
   return hparams
+
+
+@registry.register_hparams
+def ogm_dense_0():
+  hparams = mtr_tr_dense(0)
+  hparams.max_length = 1024
+  hparams.batch_size = 128
+  hparams.shared_embedding_and_softmax_weights = True
+  return hparams

From 8190c499acb7ecf7917e978b13cd39e82929c331 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 7 Feb 2019 16:37:01 -0800
Subject: [PATCH 1637/2720] Evolved Transformer TPU hparams.

PiperOrigin-RevId: 232969826
---
 tensor2tensor/models/evolved_transformer.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 4c97231cb..0ac840b16 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -105,3 +105,23 @@ def evolved_transformer_base():
 def evolved_transformer_big():
   """Big parameters for Evolved Transformer model on WMT."""
   return add_evolved_transformer_hparams(transformer.transformer_big())
+
+
+@registry.register_hparams
+def evolved_transformer_base_tpu():
+  """Base parameters for Evolved Transformer model on TPU."""
+  hparams = add_evolved_transformer_hparams(transformer.transformer_tpu())
+  hparams.learning_rate_constant = 1 / hparams.learning_rate_warmup_steps ** 0.5
+  hparams.learning_rate_schedule = (
+      "constant*single_cycle_cos_decay")
+  return hparams
+
+
+@registry.register_hparams
+def evolved_transformer_big_tpu():
+  """Big parameters for Evolved Transformer model on TPU."""
+  hparams = add_evolved_transformer_hparams(transformer.transformer_big_tpu())
+  hparams.learning_rate_constant = 1 / hparams.learning_rate_warmup_steps ** 0.5
+  hparams.learning_rate_schedule = (
+      "constant*single_cycle_cos_decay")
+  return hparams

From 63e6638835576eb209e9f3f574617592345f2e4c Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 7 Feb 2019 19:56:57 -0800
Subject: [PATCH 1638/2720] Add tf.keras getter functionality to specify args
 as strings.

PiperOrigin-RevId: 232993050
---
 tensor2tensor/layers/bayes.py      | 137 ++++++++++++++++-------------
 tensor2tensor/layers/bayes_test.py |  42 ++++++---
 2 files changed, 104 insertions(+), 75 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 2ef1a0cc9..64a0b4e1c 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -21,6 +21,7 @@
 
 import math
 
+import six
 import tensorflow as tf
 import tensorflow_probability as tfp
 
@@ -40,10 +41,6 @@ def get_config(self):
     return {'epsilon': self.epsilon}
 
 
-def positive():  # alias, following tf.keras.constraints
-  return Positive()
-
-
 class Zeros(object):
   """Function returning zeros tensor of same shape excluding the last dim."""
 
@@ -243,21 +240,21 @@ class TrainableNormal(TrainableInitializer):
   def __init__(self,
                mean_initializer=tf.keras.initializers.truncated_normal(
                    stddev=1e-5),
-               stddev_initializer=ScaledNormalStdDev(),
+               stddev_initializer='scaled_normal_std_dev',
                mean_regularizer=None,
                stddev_regularizer=None,
                mean_constraint=None,
-               stddev_constraint=positive(),
+               stddev_constraint='positive',
                seed=None,
                dtype=tf.float32):
     """Constructs the initializer."""
     super(TrainableNormal, self).__init__()
-    self.mean_initializer = mean_initializer
-    self.stddev_initializer = stddev_initializer
-    self.mean_regularizer = mean_regularizer
-    self.stddev_regularizer = stddev_regularizer
-    self.mean_constraint = mean_constraint
-    self.stddev_constraint = stddev_constraint
+    self.mean_initializer = get(mean_initializer)
+    self.stddev_initializer = get(stddev_initializer)
+    self.mean_regularizer = get(mean_regularizer)
+    self.stddev_regularizer = get(stddev_regularizer)
+    self.mean_constraint = get(mean_constraint)
+    self.stddev_constraint = get(stddev_constraint)
     self.seed = seed
     self.dtype = tf.as_dtype(dtype)
 
@@ -370,18 +367,6 @@ def get_config(self):
     }
 
 
-def trainable_normal():  # alias, following tf.keras.initializers
-  return TrainableNormal()
-
-
-def trainable_he_normal():  # alias, following tf.keras.initializers
-  return TrainableHeNormal()
-
-
-def trainable_glorot_normal():  # alias, following tf.keras.initializers
-  return TrainableGlorotNormal()
-
-
 class NormalKLDivergence(tf.keras.regularizers.Regularizer):
   """KL divergence regularizer from one normal distribution to another."""
 
@@ -409,10 +394,6 @@ def get_config(self):
     }
 
 
-def normal_kl_divergence():  # alias, following tf.keras.regularizers
-  return NormalKLDivergence()
-
-
 class DenseReparameterization(tf.keras.layers.Dense):
   """Bayesian densely-connected layer estimated via reparameterization.
 
@@ -436,25 +417,21 @@ def __init__(self,
                units,
                activation=None,
                use_bias=True,
-               kernel_initializer=None,
+               kernel_initializer='trainable_normal',
                bias_initializer='zero',
-               kernel_regularizer=normal_kl_divergence(),
+               kernel_regularizer='normal_kl_divergence',
                bias_regularizer=None,
                activity_regularizer=None,
                **kwargs):
-    if not kernel_initializer:
-      kernel_initializer = trainable_normal()
-    if not bias_initializer:
-      bias_initializer = trainable_normal()
     super(DenseReparameterization, self).__init__(
         units=units,
-        activation=activation,
+        activation=get(activation),
         use_bias=use_bias,
-        kernel_initializer=kernel_initializer,
-        bias_initializer=bias_initializer,
-        kernel_regularizer=kernel_regularizer,
-        bias_regularizer=bias_regularizer,
-        activity_regularizer=activity_regularizer,
+        kernel_initializer=get(kernel_initializer),
+        bias_initializer=get(bias_initializer),
+        kernel_regularizer=get(kernel_regularizer),
+        bias_regularizer=get(bias_regularizer),
+        activity_regularizer=get(activity_regularizer),
         **kwargs)
 
   @property
@@ -699,12 +676,12 @@ def __init__(self,
                activation='tanh',
                recurrent_activation='hard_sigmoid',
                use_bias=True,
-               kernel_initializer=None,
-               recurrent_initializer=None,
+               kernel_initializer='trainable_normal',
+               recurrent_initializer='trainable_normal',
                bias_initializer='zeros',
                unit_forget_bias=True,
-               kernel_regularizer=normal_kl_divergence(),
-               recurrent_regularizer=normal_kl_divergence(),
+               kernel_regularizer='normal_kl_divergence',
+               recurrent_regularizer='normal_kl_divergence',
                bias_regularizer=None,
                kernel_constraint=None,
                recurrent_constraint=None,
@@ -713,27 +690,21 @@ def __init__(self,
                recurrent_dropout=0.,
                implementation=1,
                **kwargs):
-    if not kernel_initializer:
-      kernel_initializer = trainable_normal()
-    if not recurrent_initializer:
-      recurrent_initializer = trainable_normal()
-    if not bias_initializer:
-      bias_initializer = trainable_normal()
     super(LSTMCellReparameterization, self).__init__(
         units=units,
-        activation=activation,
-        recurrent_activation=recurrent_activation,
+        activation=get(activation),
+        recurrent_activation=get(recurrent_activation),
         use_bias=use_bias,
-        kernel_initializer=kernel_initializer,
-        recurrent_initializer=recurrent_initializer,
-        bias_initializer=bias_initializer,
+        kernel_initializer=get(kernel_initializer),
+        recurrent_initializer=get(recurrent_initializer),
+        bias_initializer=get(bias_initializer),
         unit_forget_bias=unit_forget_bias,
-        kernel_regularizer=kernel_regularizer,
-        recurrent_regularizer=recurrent_regularizer,
-        bias_regularizer=bias_regularizer,
-        kernel_constraint=kernel_constraint,
-        recurrent_constraint=recurrent_constraint,
-        bias_constraint=bias_constraint,
+        kernel_regularizer=get(kernel_regularizer),
+        recurrent_regularizer=get(recurrent_regularizer),
+        bias_regularizer=get(bias_regularizer),
+        kernel_constraint=get(kernel_constraint),
+        recurrent_constraint=get(recurrent_constraint),
+        bias_constraint=get(bias_constraint),
         dropout=dropout,
         recurrent_dropout=recurrent_dropout,
         implementation=implementation,
@@ -954,3 +925,47 @@ def get_config(self):
     config = {'num_components': self.num_components}
     base_config = super(MixtureLogistic, self).get_config()
     return dict(list(base_config.items()) + list(config.items()))
+
+
+# Compatibility aliases, following tf.keras
+
+# pylint: disable=invalid-name
+positive = Positive
+scaled_normal_std_dev = ScaledNormalStdDev
+trainable_normal = TrainableNormal
+trainable_he_normal = TrainableHeNormal
+trainable_glorot_normal = TrainableGlorotNormal
+normal_kl_divergence = NormalKLDivergence
+# pylint: enable=invalid-name
+
+# Utility functions, following tf.keras
+
+
+def deserialize(config, custom_objects=None):
+  return tf.keras.utils.deserialize_keras_object(
+      config,
+      module_objects=globals(),
+      custom_objects=custom_objects,
+      printable_module_name='bayes')
+
+
+def get(identifier, value=None):
+  """Getter for loading from strings; returns value if can't load."""
+  if value is None:
+    value = identifier
+  if identifier is None:
+    return None
+  elif isinstance(identifier, dict):
+    try:
+      return deserialize(identifier)
+    except ValueError:
+      return value
+  elif isinstance(identifier, six.string_types):
+    config = {'class_name': str(identifier), 'config': {}}
+    try:
+      return deserialize(config)
+    except ValueError:
+      return value
+  elif callable(identifier):
+    return identifier
+  return value
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 66a0f107a..c94d0d0cf 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -43,12 +43,18 @@ def testTrainableNormalStddevConstraint(self):
     self.assertAllGreater(res, 0.)
 
   @parameterized.named_parameters(
-      {"testcase_name": "_no_uncertainty", "kernel_initializer": "zeros",
-       "bias_initializer": "zeros", "all_close": True},
-      {"testcase_name": "_kernel_uncertainty", "kernel_initializer": None,
-       "bias_initializer": "zeros", "all_close": False},
-      {"testcase_name": "_bias_uncertainty", "kernel_initializer": "zeros",
-       "bias_initializer": None, "all_close": False},
+      {"testcase_name": "_no_uncertainty",
+       "kernel_initializer": "zeros",
+       "bias_initializer": "zeros",
+       "all_close": True},
+      {"testcase_name": "_kernel_uncertainty",
+       "kernel_initializer": "trainable_normal",
+       "bias_initializer": "zeros",
+       "all_close": False},
+      {"testcase_name": "_bias_uncertainty",
+       "kernel_initializer": "zeros",
+       "bias_initializer": "trainable_normal",
+       "all_close": False},
   )
   @test_utils.run_in_graph_and_eager_modes
   def testDenseReparameterizationKernel(
@@ -164,17 +170,25 @@ def testGaussianProcessPrior(self):
     self.assertEqual(outputs_val.shape, (batch_size, output_dim))
 
   @parameterized.named_parameters(
-      {"testcase_name": "_no_uncertainty", "kernel_initializer": "zeros",
-       "recurrent_initializer": "orthogonal", "bias_initializer": "zeros",
+      {"testcase_name": "_no_uncertainty",
+       "kernel_initializer": "zeros",
+       "recurrent_initializer": "orthogonal",
+       "bias_initializer": "zeros",
        "all_close": True},
-      {"testcase_name": "_kernel_uncertainty", "kernel_initializer": None,
-       "recurrent_initializer": "orthogonal", "bias_initializer": "zeros",
+      {"testcase_name": "_kernel_uncertainty",
+       "kernel_initializer": "trainable_normal",
+       "recurrent_initializer": "orthogonal",
+       "bias_initializer": "zeros",
        "all_close": False},
-      {"testcase_name": "_recurrent_uncertainty", "kernel_initializer": "zeros",
-       "recurrent_initializer": None, "bias_initializer": "zeros",
+      {"testcase_name": "_recurrent_uncertainty",
+       "kernel_initializer": "zeros",
+       "recurrent_initializer": "trainable_normal",
+       "bias_initializer": "zeros",
        "all_close": False},
-      {"testcase_name": "_bias_uncertainty", "kernel_initializer": "zeros",
-       "recurrent_initializer": "orthogonal", "bias_initializer": None,
+      {"testcase_name": "_bias_uncertainty",
+       "kernel_initializer": "zeros",
+       "recurrent_initializer": "orthogonal",
+       "bias_initializer": "trainable_normal",
        "all_close": False},
   )
   @test_utils.run_in_graph_and_eager_modes

From 90ce10cee4978558caed6dfec49cab99958d9ab7 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 7 Feb 2019 20:17:00 -0800
Subject: [PATCH 1639/2720] Add Conv2DReparameterization.

I mirrored the DenseReparameterization implementation.

PiperOrigin-RevId: 232994978
---
 tensor2tensor/layers/bayes.py      | 139 +++++++++++++++++++++++++++++
 tensor2tensor/layers/bayes_test.py |  53 +++++++++++
 2 files changed, 192 insertions(+)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 64a0b4e1c..4e3ef4ab3 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -26,6 +26,8 @@
 import tensorflow_probability as tfp
 
 from tensorflow_probability import edward2 as ed
+from tensorflow.python.keras.utils import conv_utils  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.ops import nn_ops  # pylint: disable=g-direct-tensorflow-import
 
 
 class Positive(tf.keras.constraints.Constraint):
@@ -497,6 +499,143 @@ def build(self, input_shape):
     self.built = True
 
 
+class Conv2DReparameterization(tf.keras.layers.Conv2D):
+  """2D convolution layer (e.g. spatial convolution over images).
+
+  The layer computes a variational Bayesian approximation to the distribution
+  over convolutional layers,
+
+  ```
+  p(outputs | inputs) = int conv2d(inputs; weights, bias) p(weights, bias)
+    dweights dbias.
+  ```
+
+  It does this with a stochastic forward pass, sampling from learnable
+  distributions on the kernel and bias. Gradients with respect to the
+  distributions' learnable parameters backpropagate via reparameterization.
+  Minimizing cross-entropy plus the layer's losses performs variational
+  minimum description length, i.e., it minimizes an upper bound to the negative
+  marginal likelihood.
+  """
+
+  def __init__(self,
+               filters,
+               kernel_size,
+               strides=(1, 1),
+               padding='valid',
+               data_format=None,
+               dilation_rate=(1, 1),
+               activation=None,
+               use_bias=True,
+               kernel_initializer='trainable_normal',
+               bias_initializer='zeros',
+               kernel_regularizer='normal_kl_divergence',
+               bias_regularizer=None,
+               activity_regularizer=None,
+               kernel_constraint=None,
+               bias_constraint=None,
+               **kwargs):
+    super(Conv2DReparameterization, self).__init__(
+        filters=filters,
+        kernel_size=kernel_size,
+        strides=strides,
+        padding=padding,
+        data_format=data_format,
+        dilation_rate=dilation_rate,
+        activation=get(activation),
+        use_bias=use_bias,
+        kernel_initializer=get(kernel_initializer),
+        bias_initializer=get(bias_initializer),
+        kernel_regularizer=get(kernel_regularizer),
+        bias_regularizer=get(bias_regularizer),
+        activity_regularizer=get(activity_regularizer),
+        kernel_constraint=get(kernel_constraint),
+        bias_constraint=get(bias_constraint),
+        **kwargs)
+
+  @property
+  def kernel(self):
+    if isinstance(self.kernel_initializer, TrainableInitializer):
+      return self.kernel_initializer()
+    else:
+      return self._kernel
+
+  @property
+  def bias(self):
+    if isinstance(self.bias_initializer, TrainableInitializer):
+      return self.bias_initializer()
+    else:
+      return self._bias
+
+  def build(self, input_shape):
+    input_shape = tf.TensorShape(input_shape)
+    if self.data_format == 'channels_first':
+      channel_axis = 1
+    else:
+      channel_axis = -1
+    if input_shape.dims[channel_axis].value is None:
+      raise ValueError('The channel dimension of the inputs '
+                       'should be defined. Found `None`.')
+    input_dim = int(input_shape[channel_axis])
+    kernel_shape = self.kernel_size + (input_dim, self.filters)
+
+    if isinstance(self.kernel_initializer, TrainableInitializer):
+      self.kernel_initializer.build(kernel_shape,
+                                    self.dtype,
+                                    self.add_weight)
+      if self.kernel_regularizer is not None:
+        self.add_loss(create_regularization_loss_fn(
+            'kernel', lambda: self.kernel, self.kernel_regularizer))
+
+    else:
+      self._kernel = self.add_weight(
+          name='kernel',
+          shape=kernel_shape,
+          initializer=self.kernel_initializer,
+          regularizer=self.kernel_regularizer,
+          constraint=self.kernel_constraint,
+          trainable=True,
+          dtype=self.dtype)
+
+    if self.use_bias:
+      if isinstance(self.bias_initializer, TrainableInitializer):
+        self.bias_initializer.build((self.filters,),
+                                    self.dtype,
+                                    self.add_weight)
+        if self.bias_regularizer is not None:
+          self.add_loss(create_regularization_loss_fn(
+              'bias', lambda: self.bias, self.bias_regularizer))
+      else:
+        self._bias = self.add_weight(
+            name='bias',
+            shape=(self.filters,),
+            initializer=self.bias_initializer,
+            regularizer=self.bias_regularizer,
+            constraint=self.bias_constraint,
+            trainable=True,
+            dtype=self.dtype)
+    else:
+      self._bias = None
+
+    self.input_spec = tf.layers.InputSpec(ndim=self.rank + 2,
+                                          axes={channel_axis: input_dim})
+    if self.padding == 'causal':
+      op_padding = 'valid'
+    else:
+      op_padding = self.padding
+    if not isinstance(op_padding, (list, tuple)):
+      op_padding = op_padding.upper()
+    self._convolution_op = nn_ops.Convolution(
+        input_shape,
+        filter_shape=self.kernel.get_shape(),
+        dilation_rate=self.dilation_rate,
+        strides=self.strides,
+        padding=op_padding,
+        data_format=conv_utils.convert_data_format(self.data_format,
+                                                   self.rank + 2))
+    self.built = True
+
+
 class GaussianProcess(tf.keras.layers.Layer):
   r"""Gaussian process layer.
 
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index c94d0d0cf..4a2797e8a 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -129,6 +129,59 @@ def testDenseReparameterizationModel(self):
     self.assertEqual(res.shape, (3, 2))
     self.assertLen(model.losses, 1)
 
+  @parameterized.named_parameters(
+      {"testcase_name": "_no_uncertainty",
+       "kernel_initializer": "zeros",
+       "bias_initializer": "zeros",
+       "all_close": True},
+      {"testcase_name": "_kernel_uncertainty",
+       "kernel_initializer": "trainable_normal",
+       "bias_initializer": "zeros",
+       "all_close": False},
+      {"testcase_name": "_bias_uncertainty",
+       "kernel_initializer": "zeros",
+       "bias_initializer": "trainable_normal",
+       "all_close": False},
+  )
+  @test_utils.run_in_graph_and_eager_modes
+  def testConv2DReparameterizationKernel(
+      self, kernel_initializer, bias_initializer, all_close):
+    inputs = tf.to_float(np.random.rand(5, 4, 4, 12))
+    layer = bayes.Conv2DReparameterization(
+        4,
+        kernel_size=2,
+        kernel_initializer=kernel_initializer,
+        bias_initializer=bias_initializer,
+        activation=tf.nn.relu)
+    outputs1 = layer(inputs)
+    outputs2 = layer(inputs)
+    self.evaluate(tf.global_variables_initializer())
+    res1, res2 = self.evaluate([outputs1, outputs2])
+    self.assertEqual(res1.shape, (5, 3, 3, 4))
+    self.assertAllGreaterEqual(res1, 0.)
+    if all_close:
+      self.assertAllClose(res1, res2)
+    else:
+      self.assertNotAllClose(res1, res2)
+    layer.get_config()
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testConv2DReparameterizationModel(self):
+    inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
+    model = tf.keras.Sequential([
+        bayes.Conv2DReparameterization(3,
+                                       kernel_size=2,
+                                       padding="SAME",
+                                       activation=tf.nn.relu),
+        tf.keras.layers.Flatten(),
+        tf.keras.layers.Dense(2, activation=None),
+    ])
+    outputs = model(inputs)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(outputs)
+    self.assertEqual(res.shape, (3, 2))
+    self.assertLen(model.losses, 1)
+
   @test_utils.run_in_graph_and_eager_modes()
   def testGaussianProcessPosterior(self):
     train_batch_size = 3

From df290d61565139ec9e45fa92ad2be640761d8fe8 Mon Sep 17 00:00:00 2001
From: blazejosinski <804945+blazejosinski@users.noreply.github.com>
Date: Fri, 8 Feb 2019 19:35:05 +0100
Subject: [PATCH 1640/2720] Improve and simplify mf eval. (#1443)

---
 tensor2tensor/rl/ppo_learner.py        | 11 ++++++--
 tensor2tensor/rl/trainer_model_free.py | 39 ++++++++++----------------
 2 files changed, 23 insertions(+), 27 deletions(-)

diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index e5ba6dfb7..f17c523f0 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -53,7 +53,8 @@ def train(self,
             num_env_steps=None,
             env_step_multiplier=1,
             eval_env_fn=None,
-            report_fn=None):
+            report_fn=None,
+            model_save_fn=None):
     assert sampling_temp == 1.0 or hparams.learning_rate == 0.0, \
         "Sampling with non-1 temperature does not make sense during training."
 
@@ -118,7 +119,8 @@ def train(self,
             train_summary_op,
             eval_summary_op,
             initializers,
-            report_fn=report_fn)
+            report_fn=report_fn,
+            model_save_fn=model_save_fn)
 
   def evaluate(self, env_fn, hparams, sampling_temp):
     with tf.Graph().as_default():
@@ -191,7 +193,8 @@ def _run_train(ppo_hparams,
                train_summary_op,
                eval_summary_op,
                initializers,
-               report_fn=None):
+               report_fn=None,
+               model_save_fn=None):
   """Train."""
   summary_writer = tf.summary.FileWriter(
       event_dir, graph=tf.get_default_graph(), flush_secs=60)
@@ -244,6 +247,8 @@ def _run_train(ppo_hparams,
               "model.ckpt-{}".format(tf.train.global_step(sess, global_step))
           )
           model_saver.save(sess, ckpt_path)
+          if model_save_fn:
+            model_save_fn(model_dir)
 
 
 def _rollout_metadata(batch_env):
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 73cae5684..39bc13bad 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -62,6 +62,7 @@ def initialize_env_specs(hparams):
 
   return rl.make_real_env_fn(env)
 
+step = 0
 
 def train(hparams, output_dir, report_fn=None):
   """Train."""
@@ -94,12 +95,6 @@ def train(hparams, output_dir, report_fn=None):
     eval_every_epochs = total_steps
   policy_hparams.eval_every_epochs = 0
 
-  steps = list(range(eval_every_epochs, total_steps+1, eval_every_epochs))
-  if not steps or steps[-1] < eval_every_epochs:
-    steps.append(eval_every_epochs)
-
-  tf.logging.vlog(1, "steps: [%s]", ",".join([str(s) for s in steps]))
-
   metric_name = rl_utils.get_metric_name(
       sampling_temp=hparams.eval_sampling_temps[0],
       max_num_noops=hparams.eval_max_num_noops,
@@ -113,28 +108,24 @@ def train(hparams, output_dir, report_fn=None):
   tf.gfile.MakeDirs(eval_metrics_dir)
   eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
 
-  for i, step in enumerate(steps):
-    tf.logging.info("Starting training iteration [%d] for [%d] steps.", i, step)
-
-    policy_hparams.epochs_num = eval_every_epochs
-    learner.train(env_fn,
-                  policy_hparams,
-                  simulated=False,
-                  save_continuously=True,
-                  epoch=0)
-
-    tf.logging.info("Ended training iteration [%d] for [%d] steps.", i, step)
-
-    eval_metrics = rl_utils.evaluate_all_configs(hparams, output_dir)
-
+  def evaluate_on_new_model(model_dir_path):
+    global step
+    eval_metrics = rl_utils.evaluate_all_configs(hparams, model_dir_path)
     tf.logging.info(
         "Agent eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
-
-    rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, i)
-
+    rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, step)
     if report_fn:
       report_fn(eval_metrics[metric_name], step)
-
+    step += 1
+
+  policy_hparams.epochs_num = total_steps
+  policy_hparams.save_models_every_epochs = eval_every_epochs
+  learner.train(env_fn,
+                policy_hparams,
+                simulated=False,
+                save_continuously=True,
+                epoch=0,
+                model_save_fn=evaluate_on_new_model)
 
 def main(_):
   hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)

From 8e409f9eca381272331b2bf9a457d4195a9aa895 Mon Sep 17 00:00:00 2001
From: blazejosinski <804945+blazejosinski@users.noreply.github.com>
Date: Fri, 8 Feb 2019 10:49:22 -0800
Subject: [PATCH 1641/2720] Merge of PR #1443

PiperOrigin-RevId: 233087778
---
 tensor2tensor/rl/ppo_learner.py        | 2 +-
 tensor2tensor/rl/trainer_model_free.py | 3 +++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index f17c523f0..86b78ea9a 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -351,7 +351,7 @@ def initialization_lambda(sess):
         batch_env.initialize(sess)
 
     memory = [
-        tf.get_variable(
+        tf.get_variable(  # pylint: disable=g-complex-comprehension
             "collect_memory_%d_%s" % (epoch_length, name),
             shape=[epoch_length] + shape,
             dtype=dtype,
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 39bc13bad..4a6ce15ba 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -62,8 +62,10 @@ def initialize_env_specs(hparams):
 
   return rl.make_real_env_fn(env)
 
+
 step = 0
 
+
 def train(hparams, output_dir, report_fn=None):
   """Train."""
   env_fn = initialize_env_specs(hparams)
@@ -127,6 +129,7 @@ def evaluate_on_new_model(model_dir_path):
                 epoch=0,
                 model_save_fn=evaluate_on_new_model)
 
+
 def main(_):
   hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
 

From 7c6428d156139dd56ebcbbdf69a4de2e9b3f19db Mon Sep 17 00:00:00 2001
From: Stefan Schweter <stefan@schweter.it>
Date: Fri, 8 Feb 2019 20:24:21 +0100
Subject: [PATCH 1642/2720] doc: add evolved transformer paper (#1426)

---
 README.md           | 1 +
 docs/walkthrough.md | 1 +
 2 files changed, 2 insertions(+)

diff --git a/README.md b/README.md
index 9248cd9d9..6c5885e2b 100644
--- a/README.md
+++ b/README.md
@@ -484,5 +484,6 @@ T2T](https://research.googleblog.com/2017/06/accelerating-deep-learning-research
 * [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://arxiv.org/abs/1804.04235)
 * [Universal Transformers](https://arxiv.org/abs/1807.03819)
 * [Attending to Mathematical Language with Transformers](https://arxiv.org/abs/1812.02825)
+* [The Evolved Transformer](https://arxiv.org/abs/1901.11117)
 
 *Note: This is not an official Google product.*
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 9248cd9d9..6c5885e2b 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -484,5 +484,6 @@ T2T](https://research.googleblog.com/2017/06/accelerating-deep-learning-research
 * [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://arxiv.org/abs/1804.04235)
 * [Universal Transformers](https://arxiv.org/abs/1807.03819)
 * [Attending to Mathematical Language with Transformers](https://arxiv.org/abs/1812.02825)
+* [The Evolved Transformer](https://arxiv.org/abs/1901.11117)
 
 *Note: This is not an official Google product.*

From c81f56c15a05a52dabb53f12cab93a1deb26c395 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 8 Feb 2019 20:38:06 +0100
Subject: [PATCH 1643/2720] Replace numerically more stable log1p and expm1
 (#1424)

This PR replaces `log(1 + x)` with `log1p(x)` and `exp(x) - 1` with `expm1(x)`. These functions are more precise if x is close to zero.
---
 tensor2tensor/data_generators/wiki_revision_utils.py | 2 +-
 tensor2tensor/layers/common_attention.py             | 2 +-
 tensor2tensor/layers/common_layers.py                | 2 +-
 tensor2tensor/layers/discretization.py               | 2 +-
 tensor2tensor/models/research/autoencoders.py        | 4 ++--
 tensor2tensor/models/research/cycle_gan.py           | 2 +-
 6 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index cd6fa57f6..4864f25c3 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -51,7 +51,7 @@ def include_revision(revision_num, skip_factor=1.1):
   """
   if skip_factor <= 1.0:
     return True
-  return (int(math.log(revision_num + 1.0) / math.log(skip_factor)) != int(
+  return (int(math.log1p(revision_num) / math.log(skip_factor)) != int(
       math.log(revision_num + 2.0) / math.log(skip_factor)))
 
 
diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index f0f6674ea..3f5dcafbe 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1002,7 +1002,7 @@ def attention_bias_proximal(length):
   """
   r = tf.to_float(tf.range(length))
   diff = tf.expand_dims(r, 0) - tf.expand_dims(r, 1)
-  return tf.expand_dims(tf.expand_dims(-tf.log(1 + tf.abs(diff)), 0), 0)
+  return tf.expand_dims(tf.expand_dims(-tf.log1p(tf.abs(diff)), 0), 0)
 
 
 @expert_utils.add_name_scope()
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 3b60583dc..158495733 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1939,7 +1939,7 @@ def sample_from_discretized_mix_logistic(pred, seed=None):
   # nearest 8-bit value when sampling.
   uniform_noise = tf.random_uniform(
       tf.shape(locs), minval=1e-5, maxval=1. - 1e-5, seed=seed)
-  logistic_noise = tf.log(uniform_noise) - tf.log(1. - uniform_noise)
+  logistic_noise = tf.log(uniform_noise) - tf.log1p(-uniform_noise)
   x = locs + tf.exp(log_scales) * logistic_noise
   x0 = x[..., 0]
   x1 = x[..., 1] + coeffs[..., 0] * x0
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index bb6c89b6f..30daeecf6 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -379,7 +379,7 @@ def vae(x, z_size, name=None):
     epsilon = tf.random_normal([shape[0], shape[1], 1, z_size])
     z = mu + tf.exp(log_sigma / 2) * epsilon
     kl = 0.5 * tf.reduce_mean(
-        tf.exp(log_sigma) + tf.square(mu) - 1. - log_sigma, axis=-1)
+        tf.expm1(log_sigma) + tf.square(mu) - log_sigma, axis=-1)
     free_bits = z_size // 4
     kl_loss = tf.reduce_mean(tf.maximum(kl - free_bits, 0.0))
     return z, kl_loss, mu, log_sigma
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 8105ffd32..6e5aecdb4 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -719,7 +719,7 @@ def bottleneck(self, x):
       epsilon = tf.random_normal(x_shape[:-1] + [z_size])
       z = mu + tf.exp(log_sigma / 2) * epsilon
       kl = 0.5 * tf.reduce_mean(
-          tf.exp(log_sigma) + tf.square(mu) - 1. - log_sigma, axis=-1)
+          tf.expm1(log_sigma) + tf.square(mu) - log_sigma, axis=-1)
       free_bits = z_size // 4
       kl_loss = tf.reduce_mean(tf.maximum(kl - free_bits, 0.0))
     return z, kl_loss * hparams.kl_beta
@@ -825,7 +825,7 @@ def bottleneck(self, x):  # pylint: disable=arguments-differ
     if hparams.mode == tf.estimator.ModeKeys.TRAIN:
       # We want a number p such that p^bottleneck_bits = 1 - noise.
       # So log(p) * bottleneck_bits = log(noise)
-      log_p = tf.log(1 - float(noise) / 2) / float(hparams.bottleneck_bits)
+      log_p = tf.log1p(-float(noise) / 2) / float(hparams.bottleneck_bits)
       # Probabilities of flipping are p, p^2, p^3, ..., p^bottleneck_bits.
       noise_mask = 1.0 - tf.exp(tf.cumsum(tf.zeros_like(x) + log_p, axis=-1))
       # Having the no-noise mask, we can make noise just uniformly at random.
diff --git a/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index 5ba2a01f8..94842b2db 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -57,7 +57,7 @@ def lossfn(real_input, fake_input, compress, hparams, lsgan, name):
       loss = (dloss + gloss)/2
     else:  # cross_entropy
       dloss = -tf.reduce_mean(
-          tf.log(d1 + eps)) - tf.reduce_mean(tf.log(1 - d2 + eps))
+          tf.log(d1 + eps)) - tf.reduce_mean(tf.log1p(eps - d2))
       gloss = -tf.reduce_mean(tf.log(d2 + eps))
       loss = (dloss + gloss)/2
     return loss

From 2d78a7b1fc080550ba132c4067c55ca0727dc72a Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 8 Feb 2019 20:39:17 +0100
Subject: [PATCH 1644/2720] Remove unnecessary use of six.iterkeys (#1444)

---
 tensor2tensor/data_generators/problem.py   | 4 ++--
 tensor2tensor/layers/latent_layers_test.py | 5 ++---
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 6c651ba74..7c1327504 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -950,7 +950,7 @@ def _reverse_problem_hparams(p_hparams):
   # 'target', and each intended feature to swap has feature name 'input'.
   # In the future, remove need for this behavior.
   reversed_modality = {}
-  for feature_name in six.iterkeys(p.modality):
+  for feature_name in p.modality:
     reversed_feature_name = feature_name.replace("target", "input")
     if "target" in feature_name and reversed_feature_name in p.modality:
       reversed_modality[feature_name] = p.modality[reversed_feature_name]
@@ -962,7 +962,7 @@ def _reverse_problem_hparams(p_hparams):
 
   # Swap vocab sizes.
   reversed_vocab_size = {}
-  for feature_name in six.iterkeys(p.vocab_size):
+  for feature_name in p.vocab_size:
     reversed_feature_name = feature_name.replace("target", "input")
     if "target" in feature_name and reversed_feature_name in p.vocab_size:
       reversed_vocab_size[feature_name] = p.vocab_size[reversed_feature_name]
diff --git a/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
index 2e1ddc77b..78a83d008 100644
--- a/tensor2tensor/layers/latent_layers_test.py
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -20,7 +20,6 @@
 from __future__ import print_function
 
 import functools
-import six
 
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import discretization
@@ -138,8 +137,7 @@ def testTransformerAutoencoder(self):
     decoder_output, losses, cache = latent_layers.transformer_autoencoder(
         inputs, targets, target_space_id, hparams)
 
-    self.assertEqual(set(six.iterkeys(losses)),
-                     {"extra", "extra_loss", "latent_pred"})
+    self.assertEqual(set(losses), {"extra", "extra_loss", "latent_pred"})
 
     self.evaluate(tf.global_variables_initializer())
     decoder_output_, extra_loss_, latent_pred_ = self.evaluate(
@@ -154,5 +152,6 @@ def testTransformerAutoencoder(self):
     self.assertAllGreaterEqual(latent_pred_, 0.)
     self.assertEqual(cache, None)
 
+
 if __name__ == "__main__":
   tf.test.main()

From 25b675951b302bb7eaa9f778ed35d6e14b1b1f33 Mon Sep 17 00:00:00 2001
From: Stefan Schweter <stefan@schweter.it>
Date: Fri, 8 Feb 2019 11:39:15 -0800
Subject: [PATCH 1645/2720] Merge of PR #1426

PiperOrigin-RevId: 233097594
---
 tensor2tensor/data_generators/problem.py             | 4 ++--
 tensor2tensor/data_generators/wiki_revision_utils.py | 2 +-
 tensor2tensor/layers/common_attention.py             | 2 +-
 tensor2tensor/layers/common_layers.py                | 2 +-
 tensor2tensor/layers/discretization.py               | 2 +-
 tensor2tensor/layers/latent_layers_test.py           | 5 +++--
 tensor2tensor/models/research/autoencoders.py        | 4 ++--
 tensor2tensor/models/research/cycle_gan.py           | 2 +-
 8 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 7c1327504..6c651ba74 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -950,7 +950,7 @@ def _reverse_problem_hparams(p_hparams):
   # 'target', and each intended feature to swap has feature name 'input'.
   # In the future, remove need for this behavior.
   reversed_modality = {}
-  for feature_name in p.modality:
+  for feature_name in six.iterkeys(p.modality):
     reversed_feature_name = feature_name.replace("target", "input")
     if "target" in feature_name and reversed_feature_name in p.modality:
       reversed_modality[feature_name] = p.modality[reversed_feature_name]
@@ -962,7 +962,7 @@ def _reverse_problem_hparams(p_hparams):
 
   # Swap vocab sizes.
   reversed_vocab_size = {}
-  for feature_name in p.vocab_size:
+  for feature_name in six.iterkeys(p.vocab_size):
     reversed_feature_name = feature_name.replace("target", "input")
     if "target" in feature_name and reversed_feature_name in p.vocab_size:
       reversed_vocab_size[feature_name] = p.vocab_size[reversed_feature_name]
diff --git a/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index 4864f25c3..cd6fa57f6 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -51,7 +51,7 @@ def include_revision(revision_num, skip_factor=1.1):
   """
   if skip_factor <= 1.0:
     return True
-  return (int(math.log1p(revision_num) / math.log(skip_factor)) != int(
+  return (int(math.log(revision_num + 1.0) / math.log(skip_factor)) != int(
       math.log(revision_num + 2.0) / math.log(skip_factor)))
 
 
diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 3f5dcafbe..f0f6674ea 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1002,7 +1002,7 @@ def attention_bias_proximal(length):
   """
   r = tf.to_float(tf.range(length))
   diff = tf.expand_dims(r, 0) - tf.expand_dims(r, 1)
-  return tf.expand_dims(tf.expand_dims(-tf.log1p(tf.abs(diff)), 0), 0)
+  return tf.expand_dims(tf.expand_dims(-tf.log(1 + tf.abs(diff)), 0), 0)
 
 
 @expert_utils.add_name_scope()
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 158495733..3b60583dc 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1939,7 +1939,7 @@ def sample_from_discretized_mix_logistic(pred, seed=None):
   # nearest 8-bit value when sampling.
   uniform_noise = tf.random_uniform(
       tf.shape(locs), minval=1e-5, maxval=1. - 1e-5, seed=seed)
-  logistic_noise = tf.log(uniform_noise) - tf.log1p(-uniform_noise)
+  logistic_noise = tf.log(uniform_noise) - tf.log(1. - uniform_noise)
   x = locs + tf.exp(log_scales) * logistic_noise
   x0 = x[..., 0]
   x1 = x[..., 1] + coeffs[..., 0] * x0
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 30daeecf6..bb6c89b6f 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -379,7 +379,7 @@ def vae(x, z_size, name=None):
     epsilon = tf.random_normal([shape[0], shape[1], 1, z_size])
     z = mu + tf.exp(log_sigma / 2) * epsilon
     kl = 0.5 * tf.reduce_mean(
-        tf.expm1(log_sigma) + tf.square(mu) - log_sigma, axis=-1)
+        tf.exp(log_sigma) + tf.square(mu) - 1. - log_sigma, axis=-1)
     free_bits = z_size // 4
     kl_loss = tf.reduce_mean(tf.maximum(kl - free_bits, 0.0))
     return z, kl_loss, mu, log_sigma
diff --git a/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
index 78a83d008..2e1ddc77b 100644
--- a/tensor2tensor/layers/latent_layers_test.py
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import functools
+import six
 
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import discretization
@@ -137,7 +138,8 @@ def testTransformerAutoencoder(self):
     decoder_output, losses, cache = latent_layers.transformer_autoencoder(
         inputs, targets, target_space_id, hparams)
 
-    self.assertEqual(set(losses), {"extra", "extra_loss", "latent_pred"})
+    self.assertEqual(set(six.iterkeys(losses)),
+                     {"extra", "extra_loss", "latent_pred"})
 
     self.evaluate(tf.global_variables_initializer())
     decoder_output_, extra_loss_, latent_pred_ = self.evaluate(
@@ -152,6 +154,5 @@ def testTransformerAutoencoder(self):
     self.assertAllGreaterEqual(latent_pred_, 0.)
     self.assertEqual(cache, None)
 
-
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 6e5aecdb4..8105ffd32 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -719,7 +719,7 @@ def bottleneck(self, x):
       epsilon = tf.random_normal(x_shape[:-1] + [z_size])
       z = mu + tf.exp(log_sigma / 2) * epsilon
       kl = 0.5 * tf.reduce_mean(
-          tf.expm1(log_sigma) + tf.square(mu) - log_sigma, axis=-1)
+          tf.exp(log_sigma) + tf.square(mu) - 1. - log_sigma, axis=-1)
       free_bits = z_size // 4
       kl_loss = tf.reduce_mean(tf.maximum(kl - free_bits, 0.0))
     return z, kl_loss * hparams.kl_beta
@@ -825,7 +825,7 @@ def bottleneck(self, x):  # pylint: disable=arguments-differ
     if hparams.mode == tf.estimator.ModeKeys.TRAIN:
       # We want a number p such that p^bottleneck_bits = 1 - noise.
       # So log(p) * bottleneck_bits = log(noise)
-      log_p = tf.log1p(-float(noise) / 2) / float(hparams.bottleneck_bits)
+      log_p = tf.log(1 - float(noise) / 2) / float(hparams.bottleneck_bits)
       # Probabilities of flipping are p, p^2, p^3, ..., p^bottleneck_bits.
       noise_mask = 1.0 - tf.exp(tf.cumsum(tf.zeros_like(x) + log_p, axis=-1))
       # Having the no-noise mask, we can make noise just uniformly at random.
diff --git a/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index 94842b2db..5ba2a01f8 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -57,7 +57,7 @@ def lossfn(real_input, fake_input, compress, hparams, lsgan, name):
       loss = (dloss + gloss)/2
     else:  # cross_entropy
       dloss = -tf.reduce_mean(
-          tf.log(d1 + eps)) - tf.reduce_mean(tf.log1p(eps - d2))
+          tf.log(d1 + eps)) - tf.reduce_mean(tf.log(1 - d2 + eps))
       gloss = -tf.reduce_mean(tf.log(d2 + eps))
       loss = (dloss + gloss)/2
     return loss

From bded2def654a5e85d26d12bf1831454200865c97 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 8 Feb 2019 11:54:06 -0800
Subject: [PATCH 1646/2720] Update LICENSEs following OSS change; remove
 Modality._body_input_depth.

This is an incremental change toward replacing Modality classes with a simpler ModalityType enum.

PiperOrigin-RevId: 233100202
---
 tensor2tensor/layers/modalities.py | 70 ++++++++++++++++++------------
 tensor2tensor/utils/modality.py    |  7 ---
 2 files changed, 43 insertions(+), 34 deletions(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 79a453ac8..5a0411679 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -42,7 +42,8 @@ class SymbolModality(modality.Modality):
 
   @property
   def name(self):
-    return "symbol_modality_%d_%d" % (self._vocab_size, self._body_input_depth)
+    return "symbol_modality_%d_%d" % (self._vocab_size,
+                                      self._model_hparams.hidden_size)
 
   @property
   def top_is_pointwise(self):
@@ -71,13 +72,13 @@ def _get_weights(self, hidden_dim=None):
     """Create or get concatenated embedding or softmax variable.
 
     Args:
-      hidden_dim: dim of the variable. Defaults to self._body_input_depth
+      hidden_dim: dim of the variable. Defaults to _model_hparams' hidden_size
 
     Returns:
        a list of self._num_shards Tensors.
     """
     if hidden_dim is None:
-      hidden_dim = self._body_input_depth
+      hidden_dim = self._model_hparams.hidden_size
     num_shards = self._model_hparams.symbol_modality_num_shards
     shards = []
     for i in range(num_shards):
@@ -110,7 +111,7 @@ def bottom_simple(self, x, name, reuse):
           x, 1.0 - self._model_hparams.symbol_dropout)
       ret = common_layers.gather(var, x)
       if self._model_hparams.multiply_embedding_mode == "sqrt_depth":
-        ret *= self._body_input_depth**0.5
+        ret *= self._model_hparams.hidden_size**0.5
       ret *= tf.expand_dims(
           common_layers.cast_like(tf.not_equal(x, 0), ret), -1)
       return ret
@@ -136,7 +137,9 @@ def top(self, body_output, _):
     """Generate logits.
 
     Args:
-      body_output: A Tensor with shape [batch, p0, p1, body_input_depth]
+      body_output: A Tensor with shape
+        [batch, p0, p1, self._model_hparams.hidden_size].
+
     Returns:
       logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].
     """
@@ -255,7 +258,7 @@ def targets_bottom(self, x):
       embedded = tf.reshape(embedded, inputs_shape[:3] + [merged_size])
       merged = tf.layers.dense(
           embedded,
-          self._body_input_depth,
+          self._model_hparams.hidden_size,
           name="merge_pixel_embedded_channels")
       return merged
 
@@ -308,7 +311,8 @@ def bottom_compress(self, inputs, name="bottom"):
       name: string, scope.
 
     Returns:
-      body_input: Tensor of shape [batch, img_len, img_len, body_input_depth].
+      body_input: Tensor of shape
+        [batch, img_len, img_len, self._model_hparams.hidden_size].
     """
     with tf.variable_scope(name):
       inputs = tf.to_float(inputs)
@@ -328,7 +332,7 @@ def bottom_compress(self, inputs, name="bottom"):
       # Compress RGB intensities for each pixel using a convolution.
       outputs = tf.layers.conv2d(
           inputs,
-          self._body_input_depth,
+          self._model_hparams.hidden_size,
           kernel_size=(1, self.num_channels),
           padding="VALID",
           strides=(1, self.num_channels),
@@ -435,8 +439,10 @@ def bottom(self, x):
 
     Args:
       x: A Tensor with shape [batch, ...]
+
     Returns:
-      body_input: A Tensor with shape [batch, ?, ?, body_input_depth].
+      body_input: A Tensor with shape [batch, ?, ?,
+        self._model_hparams.hidden_size].
     """
     inputs = x
     with tf.variable_scope(self.name):
@@ -467,7 +473,9 @@ def xnet_resblock(x, filters, res_relu, name):
       x.set_shape([None, None, None, 1])
       for i in range(self._model_hparams.audio_compression):
         x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
-      return xnet_resblock(x, self._body_input_depth, False,
+      return xnet_resblock(x,
+                           self._model_hparams.hidden_size,
+                           False,
                            "compress_block_final")
 
 
@@ -479,8 +487,10 @@ def bottom(self, x):
 
     Args:
       x: A Tensor with shape [batch, ...]
+
     Returns:
-      body_input: A Tensor with shape [batch, ?, ?, body_input_depth].
+      body_input: A Tensor with shape [batch, ?, ?,
+        self._model_hparams.hidden_size].
     """
     inputs = x
     with tf.variable_scope(self.name):
@@ -512,7 +522,9 @@ def xnet_resblock(x, filters, res_relu, name):
       x.set_shape([None, None, None, 1])
       for i in range(self._model_hparams.audio_compression):
         x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
-      return xnet_resblock(x, self._body_input_depth, False,
+      return xnet_resblock(x,
+                           self._model_hparams.hidden_size,
+                           False,
                            "compress_block_final")
 
 
@@ -651,7 +663,7 @@ def bottom(self, x):
       # Project.
       return tf.layers.dense(
           embedded,
-          self._body_input_depth,
+          self._model_hparams.hidden_size,
           name="merge_pixel_embedded_frames")
 
   def targets_bottom(self, x):  # pylint: disable=arguments-differ
@@ -666,7 +678,7 @@ def targets_bottom(self, x):  # pylint: disable=arguments-differ
       transposed = common_layers.time_to_channels(embedded)
       return tf.layers.dense(
           transposed,
-          self._body_input_depth,
+          self._model_hparams.hidden_size,
           name="merge_pixel_embedded_frames")
 
 
@@ -787,20 +799,24 @@ class ClassLabelModality(modality.Modality):
   @property
   def name(self):
     return "class_label_modality_%d_%d" % (self._vocab_size,
-                                           self._body_input_depth)
+                                           self._model_hparams.hidden_size)
 
   def bottom(self, x):
     with tf.variable_scope(self.name):
       multiplier = 1.0
       if self._model_hparams.multiply_embedding_mode == "sqrt_depth":
-        multiplier = self._body_input_depth**0.5
-      return common_layers.embedding(
-          x, self._vocab_size, self._body_input_depth, multiplier=multiplier)
+        multiplier = self._model_hparams.hidden_size**0.5
+      return common_layers.embedding(x,
+                                     self._vocab_size,
+                                     self._model_hparams.hidden_size,
+                                     multiplier=multiplier)
 
   def targets_bottom(self, x):
     with tf.variable_scope(self.name):
-      return tf.zeros(
-          [common_layers.shape_list(x)[0], 1, 1, self._body_input_depth])
+      return tf.zeros([common_layers.shape_list(x)[0],
+                       1,
+                       1,
+                       self._model_hparams.hidden_size])
 
   def top(self, body_output, _):
     """Transform inputs from model space to target space.
@@ -935,7 +951,7 @@ def top_is_pointwise(self):
   def bottom(self, x):
     with tf.variable_scope("real"):
       return tf.layers.dense(
-          tf.to_float(x), self._body_input_depth, name="bottom")
+          tf.to_float(x), self._model_hparams.hidden_size, name="bottom")
 
   def top(self, body_output, _):
     with tf.variable_scope("real"):
@@ -1000,8 +1016,8 @@ class SigmoidClassLabelModality(ClassLabelModality):
 
   @property
   def name(self):
-    return "sigmoid_class_symbol_modality_%d_%d" % (self._vocab_size,
-                                                    self._body_input_depth)
+    return "sigmoid_class_symbol_modality_%d_%d" % (
+        self._vocab_size, self._model_hparams.hidden_size)
 
   def loss(self, top_out, targets):
     # Expect inputs of size [batch-size, timesteps, 1, num-classes], where the
@@ -1020,7 +1036,7 @@ class SigmoidMaxPoolingClassLabelModality(ClassLabelModality):
   @property
   def name(self):
     return "sigmoid_max_pooling_class_symbol_modality_%d_%d" % (
-        self._vocab_size, self._body_input_depth)
+        self._vocab_size, self._model_hparams.hidden_size)
 
   def top(self, body_output, _):
     """Transform inputs from model space to target space.
@@ -1055,7 +1071,7 @@ class SoftmaxMaxPoolingClassLabelModality(OneHotClassLabelModality):
   @property
   def name(self):
     return "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
-        self._vocab_size, self._body_input_depth)
+        self._vocab_size, self._model_hparams.hidden_size)
 
   def top(self, body_output, _):
     with tf.variable_scope(self.name):
@@ -1070,7 +1086,7 @@ class SoftmaxAveragePoolingClassLabelModality(OneHotClassLabelModality):
   @property
   def name(self):
     return "softmax_average_pooling_onehot_class_label_modality_%d_%d" % (
-        self._vocab_size, self._body_input_depth)
+        self._vocab_size, self._model_hparams.hidden_size)
 
   def top(self, body_output, _):
     with tf.variable_scope(self.name):
@@ -1085,7 +1101,7 @@ class SoftmaxLastTimestepClassLabelModality(OneHotClassLabelModality):
   @property
   def name(self):
     return "softmax_last_timestep_onehot_class_label_modality_%d_%d" % (
-        self._vocab_size, self._body_input_depth)
+        self._vocab_size, self._model_hparams.hidden_size)
 
   def top(self, body_output, _):
     with tf.variable_scope(self.name):
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index a2a6fddb1..f3ab31e18 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -46,9 +46,6 @@ class Modality(object):
     same as the `bottom` function, and that is the default we use. But, e.g.,
     for images, a different function might be needed to regress properly.
   * `loss` would compare the generated image to the target image and score it.
-
-  All the functions have simple and sharded versions. A sub-class only needs to
-  implement the simple version, the default sharding will be used then.
   """
 
   def __init__(self, model_hparams, vocab_size=None):
@@ -66,10 +63,6 @@ def top_dimensionality(self):
     """Integer, the last dimension of the predictions (vocab size)."""
     return self._vocab_size
 
-  @property
-  def _body_input_depth(self):
-    return self._model_hparams.hidden_size
-
   @property
   def top_is_pointwise(self):
     """Whether the top mapping of the modality is pointwise.

From eac8c5f10fbeac8d3a1ed93aed1f6b9fbe1eac54 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 8 Feb 2019 11:55:05 -0800
Subject: [PATCH 1647/2720] Merge of PR #1424

PiperOrigin-RevId: 233100331
---
 tensor2tensor/data_generators/wiki_revision_utils.py | 2 +-
 tensor2tensor/layers/common_attention.py             | 2 +-
 tensor2tensor/layers/common_layers.py                | 2 +-
 tensor2tensor/layers/discretization.py               | 2 +-
 tensor2tensor/models/research/autoencoders.py        | 4 ++--
 tensor2tensor/models/research/cycle_gan.py           | 2 +-
 6 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index cd6fa57f6..4864f25c3 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -51,7 +51,7 @@ def include_revision(revision_num, skip_factor=1.1):
   """
   if skip_factor <= 1.0:
     return True
-  return (int(math.log(revision_num + 1.0) / math.log(skip_factor)) != int(
+  return (int(math.log1p(revision_num) / math.log(skip_factor)) != int(
       math.log(revision_num + 2.0) / math.log(skip_factor)))
 
 
diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index f0f6674ea..3f5dcafbe 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1002,7 +1002,7 @@ def attention_bias_proximal(length):
   """
   r = tf.to_float(tf.range(length))
   diff = tf.expand_dims(r, 0) - tf.expand_dims(r, 1)
-  return tf.expand_dims(tf.expand_dims(-tf.log(1 + tf.abs(diff)), 0), 0)
+  return tf.expand_dims(tf.expand_dims(-tf.log1p(tf.abs(diff)), 0), 0)
 
 
 @expert_utils.add_name_scope()
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 3b60583dc..158495733 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1939,7 +1939,7 @@ def sample_from_discretized_mix_logistic(pred, seed=None):
   # nearest 8-bit value when sampling.
   uniform_noise = tf.random_uniform(
       tf.shape(locs), minval=1e-5, maxval=1. - 1e-5, seed=seed)
-  logistic_noise = tf.log(uniform_noise) - tf.log(1. - uniform_noise)
+  logistic_noise = tf.log(uniform_noise) - tf.log1p(-uniform_noise)
   x = locs + tf.exp(log_scales) * logistic_noise
   x0 = x[..., 0]
   x1 = x[..., 1] + coeffs[..., 0] * x0
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index bb6c89b6f..30daeecf6 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -379,7 +379,7 @@ def vae(x, z_size, name=None):
     epsilon = tf.random_normal([shape[0], shape[1], 1, z_size])
     z = mu + tf.exp(log_sigma / 2) * epsilon
     kl = 0.5 * tf.reduce_mean(
-        tf.exp(log_sigma) + tf.square(mu) - 1. - log_sigma, axis=-1)
+        tf.expm1(log_sigma) + tf.square(mu) - log_sigma, axis=-1)
     free_bits = z_size // 4
     kl_loss = tf.reduce_mean(tf.maximum(kl - free_bits, 0.0))
     return z, kl_loss, mu, log_sigma
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 8105ffd32..6e5aecdb4 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -719,7 +719,7 @@ def bottleneck(self, x):
       epsilon = tf.random_normal(x_shape[:-1] + [z_size])
       z = mu + tf.exp(log_sigma / 2) * epsilon
       kl = 0.5 * tf.reduce_mean(
-          tf.exp(log_sigma) + tf.square(mu) - 1. - log_sigma, axis=-1)
+          tf.expm1(log_sigma) + tf.square(mu) - log_sigma, axis=-1)
       free_bits = z_size // 4
       kl_loss = tf.reduce_mean(tf.maximum(kl - free_bits, 0.0))
     return z, kl_loss * hparams.kl_beta
@@ -825,7 +825,7 @@ def bottleneck(self, x):  # pylint: disable=arguments-differ
     if hparams.mode == tf.estimator.ModeKeys.TRAIN:
       # We want a number p such that p^bottleneck_bits = 1 - noise.
       # So log(p) * bottleneck_bits = log(noise)
-      log_p = tf.log(1 - float(noise) / 2) / float(hparams.bottleneck_bits)
+      log_p = tf.log1p(-float(noise) / 2) / float(hparams.bottleneck_bits)
       # Probabilities of flipping are p, p^2, p^3, ..., p^bottleneck_bits.
       noise_mask = 1.0 - tf.exp(tf.cumsum(tf.zeros_like(x) + log_p, axis=-1))
       # Having the no-noise mask, we can make noise just uniformly at random.
diff --git a/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index 5ba2a01f8..94842b2db 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -57,7 +57,7 @@ def lossfn(real_input, fake_input, compress, hparams, lsgan, name):
       loss = (dloss + gloss)/2
     else:  # cross_entropy
       dloss = -tf.reduce_mean(
-          tf.log(d1 + eps)) - tf.reduce_mean(tf.log(1 - d2 + eps))
+          tf.log(d1 + eps)) - tf.reduce_mean(tf.log1p(eps - d2))
       gloss = -tf.reduce_mean(tf.log(d2 + eps))
       loss = (dloss + gloss)/2
     return loss

From 215be538067123fcb09d19c158a4516b7824e7fa Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 8 Feb 2019 12:14:18 -0800
Subject: [PATCH 1648/2720] Allow to chunk examples with long targets into
 multiple batches.

PiperOrigin-RevId: 233103880
---
 tensor2tensor/layers/common_hparams.py |  3 +++
 tensor2tensor/utils/data_reader.py     | 34 ++++++++++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 390dfff9b..4bf961050 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -151,6 +151,9 @@ def basic_params1():
       # than max_length.
       # If max_length==0, we use hparams.batch_size instead.
       max_length=0,
+      # Split targets on the first axis into chunks of this length.
+      split_targets_chunk_length=0,
+      split_targets_max_chunks=100,
       # Maximum length in the smallest length bucket.  Setting this
       # flag too high will result in wasteful padding of short
       # sequences.  Due to some (hopefully) temporary hacks in the
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index 4aa21438d..ee063ebdf 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -477,6 +477,40 @@ def define_shapes(example):
       hparams.batch_shuffle_size):
     dataset = dataset.shuffle(hparams.batch_shuffle_size)
 
+  # Split batches into chunks if targets are too long.
+  # The new "chunk_number" feature is 0 for the first chunk and goes up then.
+  # Chunks are reversed so the 0th chunk comes first, then the 1st and so on,
+  # so models can attend to them in the order they arrive. The last chunk is
+  # usually the one containing the end of the target sentence (EOS).
+  chunk_length = hparams.get("split_targets_chunk_length", 0)
+  max_chunks = hparams.get("split_targets_max_chunks", 100)
+  if chunk_length > 0:
+    def is_nonzero_chunk(example):
+      """A chunk is zero if all targets are 0s."""
+      return tf.less(0, tf.reduce_sum(tf.abs(example["targets"])))
+
+    def split_on_length(example):
+      """Split a batch of ditcs on length."""
+      x = example["targets"]
+      length_diff = chunk_length * max_chunks - tf.shape(x)[1]
+      padded_x = tf.pad(x, [(0, 0), (0, length_diff), (0, 0), (0, 0)])
+      chunks = [padded_x[:, i*chunk_length:(i+1)*chunk_length, :, :]
+                for i in range(max_chunks - 1)]
+      chunks.append(padded_x[:, (max_chunks - 1)*chunk_length:, :, :])
+      new_example = {}
+      new_example["chunk_number"] = tf.range(max_chunks)
+      new_example["targets"] = tf.concat(
+          [tf.expand_dims(c, axis=0) for c in chunks], axis=0)
+      for k in example:
+        if k != "targets":
+          new_example[k] = tf.concat(
+              [tf.expand_dims(example[k], axis=0) for _ in range(max_chunks)],
+              axis=0)
+      return tf.data.Dataset.from_tensor_slices(new_example)
+
+    dataset = dataset.flat_map(split_on_length)
+    dataset = dataset.filter(is_nonzero_chunk)
+
   def prepare_for_output(example):
     if not config or not config.use_tpu:
       _summarize_features(example, num_shards)

From 69046f61b0b0002636202d3a773efc3ddb71a411 Mon Sep 17 00:00:00 2001
From: Artit 'Art' Wangperawong <artitw@gmail.com>
Date: Fri, 8 Feb 2019 16:02:56 -0500
Subject: [PATCH 1649/2720] Fix Algorithmic Math Two Variables (#1442)

* Fix Algorithmic Math Two Variables

* Fix Algorithmic Math Two Variables data download

* Fix Algorithmic Math Two Variables data file paths
---
 README.md                                     |  2 +-
 ...g.py => algorithmic_math_two_variables.py} | 41 +++++++++++++++----
 tensor2tensor/data_generators/all_problems.py |  2 +-
 3 files changed, 34 insertions(+), 11 deletions(-)
 rename tensor2tensor/data_generators/{mathematical_language_understanding.py => algorithmic_math_two_variables.py} (76%)

diff --git a/README.md b/README.md
index 6c5885e2b..e0bd75a1d 100644
--- a/README.md
+++ b/README.md
@@ -85,7 +85,7 @@ to modify the hyperparameters if you run on a different setup.
 For evaluating mathematical expressions at the character level involving addition, subtraction and multiplication of both positive and negative decimal numbers with variable digits assigned to symbolic variables, use
 
 * the [MLU](https://art.wangperawong.com/mathematical_language_understanding_train.tar.gz) data-set:
- `--problem=mathematical_language_understanding`
+ `--problem=algorithmic_math_two_variables`
 
 You can try solving the problem with different transformer models and hyperparameters as described in the [paper](https://arxiv.org/abs/1812.02825):
 * Standard transformer:
diff --git a/tensor2tensor/data_generators/mathematical_language_understanding.py b/tensor2tensor/data_generators/algorithmic_math_two_variables.py
similarity index 76%
rename from tensor2tensor/data_generators/mathematical_language_understanding.py
rename to tensor2tensor/data_generators/algorithmic_math_two_variables.py
index f01295d7b..d5b230ac3 100644
--- a/tensor2tensor/data_generators/mathematical_language_understanding.py
+++ b/tensor2tensor/data_generators/algorithmic_math_two_variables.py
@@ -53,13 +53,40 @@
 import tensorflow as tf
 
 
+_URL = ("https://art.wangperawong.com/mathematical_language_understanding"
+         "_train.tar.gz")
+
+def _download_mlu_data(tmp_dir, data_dir):
+  """Downloads and extracts the dataset.
+
+  Args:
+    tmp_dir: temp directory to download and extract the dataset
+    data_dir: The base directory where data and vocab files are stored.
+
+  Returns:
+    tmp_dir: temp directory containing the raw data.
+  """
+  if not tf.gfile.Exists(data_dir):
+    tf.gfile.MakeDirs(data_dir)
+
+  filename = os.path.basename(_URL)
+  file_path = os.path.join(tmp_dir, filename)
+  headers = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_1) "
+                           "AppleWebKit/537.36 (KHTML, like Gecko) "
+                           "Chrome/63.0.3239.132 Safari/537.36"}
+  resp = requests.get(_URL, headers=headers)
+  with open(file_path, "wb") as f:
+    f.write(resp.content)
+
+  with tarfile.open(file_path, "r:gz") as tar:
+    tar.extractall(tmp_dir)
+
+  return tmp_dir
+
 @registry.register_problem
 class AlgorithmicMathTwoVariables(text_problems.Text2TextProblem):
   """Mathematical language understanding, see arxiv.org/abs/1812.02825."""
 
-  URL = ("https://art.wangperawong.com/mathematical_language_understanding"
-         "_train.tar.gz")
-
   @property
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
@@ -96,12 +123,8 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
       tf.gfile.MakeDirs(data_dir)
 
     # Download and extract.
-    compressed_filename = os.path.basename(self.URL)
-    download_path = generator_utils.maybe_download(
-        tmp_dir, compressed_filename, self.URL)
-    with tarfile.open(download_path, "r:gz") as tar:
-      tar.extractall(tmp_dir)
-    filepath = os.path.join(tmp_dir, "symbolic_math_train.txt")
+    download_path = download_mlu_data(tmp_dir, data_dir)
+    filepath = os.path.join(download_path, "symbolic_math_train.txt")
     with open(filepath, "r") as fp:
       for l in fp:
         prob, ans = l.strip().split(":")
diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 199a4bfa8..ece1d1b48 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -50,7 +50,7 @@
     "tensor2tensor.data_generators.lm1b",
     "tensor2tensor.data_generators.lm1b_imdb",
     "tensor2tensor.data_generators.lm1b_mnli",
-    "tensor2tensor.data_generators.mathematical_language_understanding",
+    "tensor2tensor.data_generators.algorithmic_math_two_variables",
     "tensor2tensor.data_generators.mnist",
     "tensor2tensor.data_generators.mrpc",
     "tensor2tensor.data_generators.mscoco",

From 8bf3c3321c01b280f25bcd7e938ce884b4cb1ac1 Mon Sep 17 00:00:00 2001
From: Artit 'Art' Wangperawong <artitw@gmail.com>
Date: Fri, 8 Feb 2019 13:10:19 -0800
Subject: [PATCH 1650/2720] Merge of PR #1442

PiperOrigin-RevId: 233113072
---
 docs/walkthrough.md                                       | 2 +-
 .../data_generators/algorithmic_math_two_variables.py     | 8 +++++---
 tensor2tensor/data_generators/all_problems.py             | 2 +-
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 6c5885e2b..e0bd75a1d 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -85,7 +85,7 @@ to modify the hyperparameters if you run on a different setup.
 For evaluating mathematical expressions at the character level involving addition, subtraction and multiplication of both positive and negative decimal numbers with variable digits assigned to symbolic variables, use
 
 * the [MLU](https://art.wangperawong.com/mathematical_language_understanding_train.tar.gz) data-set:
- `--problem=mathematical_language_understanding`
+ `--problem=algorithmic_math_two_variables`
 
 You can try solving the problem with different transformer models and hyperparameters as described in the [paper](https://arxiv.org/abs/1812.02825):
 * Standard transformer:
diff --git a/tensor2tensor/data_generators/algorithmic_math_two_variables.py b/tensor2tensor/data_generators/algorithmic_math_two_variables.py
index d5b230ac3..2d40a437b 100644
--- a/tensor2tensor/data_generators/algorithmic_math_two_variables.py
+++ b/tensor2tensor/data_generators/algorithmic_math_two_variables.py
@@ -44,8 +44,8 @@
 
 import os
 import tarfile
+import requests
 
-from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
@@ -54,7 +54,8 @@
 
 
 _URL = ("https://art.wangperawong.com/mathematical_language_understanding"
-         "_train.tar.gz")
+        "_train.tar.gz")
+
 
 def _download_mlu_data(tmp_dir, data_dir):
   """Downloads and extracts the dataset.
@@ -83,6 +84,7 @@ def _download_mlu_data(tmp_dir, data_dir):
 
   return tmp_dir
 
+
 @registry.register_problem
 class AlgorithmicMathTwoVariables(text_problems.Text2TextProblem):
   """Mathematical language understanding, see arxiv.org/abs/1812.02825."""
@@ -123,7 +125,7 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
       tf.gfile.MakeDirs(data_dir)
 
     # Download and extract.
-    download_path = download_mlu_data(tmp_dir, data_dir)
+    download_path = _download_mlu_data(tmp_dir, data_dir)
     filepath = os.path.join(download_path, "symbolic_math_train.txt")
     with open(filepath, "r") as fp:
       for l in fp:
diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index ece1d1b48..157970685 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -24,6 +24,7 @@
 MODULES = [
     "tensor2tensor.data_generators.algorithmic",
     "tensor2tensor.data_generators.algorithmic_math",
+    "tensor2tensor.data_generators.algorithmic_math_two_variables",
     "tensor2tensor.data_generators.allen_brain",
     "tensor2tensor.data_generators.audio",
     "tensor2tensor.data_generators.babi_qa",
@@ -50,7 +51,6 @@
     "tensor2tensor.data_generators.lm1b",
     "tensor2tensor.data_generators.lm1b_imdb",
     "tensor2tensor.data_generators.lm1b_mnli",
-    "tensor2tensor.data_generators.algorithmic_math_two_variables",
     "tensor2tensor.data_generators.mnist",
     "tensor2tensor.data_generators.mrpc",
     "tensor2tensor.data_generators.mscoco",

From a5f26c1963ab262bc6d43b26e8d49612c422caa0 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 8 Feb 2019 13:18:18 -0800
Subject: [PATCH 1651/2720] Restructure trainable initializers as
 tf.keras.layers.Layers.

With trainable initializers as Layers, and overloading self.add_weight(), I was able to remove build() boilerplate in Layers. I kept it for LSTMCellReparameterization however as it's needed for the initializer logic on the forget gate.

PiperOrigin-RevId: 233114333
---
 tensor2tensor/layers/bayes.py      | 392 +++++++++++------------------
 tensor2tensor/layers/bayes_test.py |  67 +++--
 2 files changed, 188 insertions(+), 271 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 4e3ef4ab3..558d960a1 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -26,8 +26,6 @@
 import tensorflow_probability as tfp
 
 from tensorflow_probability import edward2 as ed
-from tensorflow.python.keras.utils import conv_utils  # pylint: disable=g-direct-tensorflow-import
-from tensorflow.python.ops import nn_ops  # pylint: disable=g-direct-tensorflow-import
 
 
 class Positive(tf.keras.constraints.Constraint):
@@ -151,6 +149,10 @@ def _compute_fans(shape):
       receptive_field_size *= dim
     fan_in = shape[-2] * receptive_field_size
     fan_out = shape[-1] * receptive_field_size
+  if isinstance(fan_in, tf.Dimension):
+    fan_in = fan_in.value
+  if isinstance(fan_out, tf.Dimension):
+    fan_out = fan_out.value
   return fan_in, fan_out
 
 
@@ -219,24 +221,7 @@ def __call__(self, shape, dtype=None, partition_info=None):
                                       dtype=dtype)
 
 
-# TODO(dusenberrymw): Restructure the implementation of a trainable initializer
-# such that callers do not need to have type-conditional logic.
-class TrainableInitializer(tf.keras.initializers.Initializer):
-  """An initializer with trainable variables.
-
-  In this implementation, a layer must call `build` before usage in order to
-  capture the variables.
-  """
-
-  def __init__(self):
-    self.built = False
-
-  def build(self, shape, dtype=None, add_variable_fn=None):
-    """Builds the initializer, with the variables captured by the caller."""
-    raise NotImplementedError
-
-
-class TrainableNormal(TrainableInitializer):
+class TrainableNormal(tf.keras.layers.Layer):
   """Random normal op as an initializer with trainable mean and stddev."""
 
   def __init__(self,
@@ -248,9 +233,10 @@ def __init__(self,
                mean_constraint=None,
                stddev_constraint='positive',
                seed=None,
-               dtype=tf.float32):
+               dtype=tf.float32,
+               **kwargs):
     """Constructs the initializer."""
-    super(TrainableNormal, self).__init__()
+    super(TrainableNormal, self).__init__(dtype=dtype, **kwargs)
     self.mean_initializer = get(mean_initializer)
     self.stddev_initializer = get(stddev_initializer)
     self.mean_regularizer = get(mean_regularizer)
@@ -258,16 +244,12 @@ def __init__(self,
     self.mean_constraint = get(mean_constraint)
     self.stddev_constraint = get(stddev_constraint)
     self.seed = seed
-    self.dtype = tf.as_dtype(dtype)
 
-  def build(self, shape, dtype=None, add_variable_fn=None):
-    """Builds the initializer, with the variables captured by the caller."""
+  def build(self, shape, dtype=None):
     if dtype is None:
       dtype = self.dtype
-    self.shape = shape
-    self.dtype = tf.as_dtype(dtype)
 
-    self.mean = add_variable_fn(
+    self.mean = self.add_weight(
         'mean',
         shape=shape,
         initializer=self.mean_initializer,
@@ -275,7 +257,7 @@ def build(self, shape, dtype=None, add_variable_fn=None):
         constraint=self.mean_constraint,
         dtype=dtype,
         trainable=True)
-    self.stddev = add_variable_fn(
+    self.stddev = self.add_weight(
         'stddev',
         shape=shape,
         initializer=self.stddev_initializer,
@@ -285,16 +267,13 @@ def build(self, shape, dtype=None, add_variable_fn=None):
         trainable=True)
     self.built = True
 
-  def __call__(self, shape=None, dtype=None, partition_info=None):
-    del shape, dtype, partition_info  # Unused in TrainableInitializers.
-    # TODO(dusenberrymw): Restructure so that we can build as needed.
+  def __call__(self, shape, dtype=None, partition_info=None):
+    del partition_info  # unused arg
     if not self.built:
-      raise ValueError('A TrainableInitializer must be built by a layer before '
-                       'usage, and is currently only compatible with Bayesian '
-                       'layers.')
+      self.build(shape, dtype)
     return ed.Independent(
         ed.Normal(loc=self.mean, scale=self.stddev).distribution,
-        reinterpreted_batch_ndims=len(self.shape))
+        reinterpreted_batch_ndims=len(shape))
 
   def get_config(self):
     return {
@@ -311,7 +290,7 @@ def get_config(self):
         'stddev_constraint':
             tf.keras.constraints.serialize(self.stddev_constraint),
         'seed': self.seed,
-        'dtype': self.dtype.name,
+        'dtype': self.dtype,
     }
 
 
@@ -436,67 +415,42 @@ def __init__(self,
         activity_regularizer=get(activity_regularizer),
         **kwargs)
 
-  @property
-  def kernel(self):
-    if isinstance(self.kernel_initializer, TrainableInitializer):
-      return self.kernel_initializer()
-    else:
-      return self._kernel
-
-  @property
-  def bias(self):
-    if isinstance(self.bias_initializer, TrainableInitializer):
-      return self.bias_initializer()
-    else:
-      return self._bias
-
-  def build(self, input_shape):
-    input_shape = tf.TensorShape(input_shape)
-    last_dim = input_shape[-1]
-    if isinstance(last_dim, tf.Dimension):
-      last_dim = last_dim.value
-    if last_dim is None:
-      raise ValueError('The last dimension of the inputs to `Dense` '
-                       'should be defined. Found `None`.')
-    self.input_spec = tf.layers.InputSpec(min_ndim=2, axes={-1: last_dim})
-
-    if isinstance(self.kernel_initializer, TrainableInitializer):
-      self.kernel_initializer.build([last_dim, self.units],
-                                    self.dtype,
-                                    self.add_weight)
-      if self.kernel_regularizer is not None:
-        self.add_loss(create_regularization_loss_fn(
-            'kernel', lambda: self.kernel, self.kernel_regularizer))
-
-    else:
-      self._kernel = self.add_weight(
-          'kernel',
-          shape=[last_dim, self.units],
-          initializer=self.kernel_initializer,
-          regularizer=self.kernel_regularizer,
-          constraint=self.kernel_constraint,
-          dtype=self.dtype,
-          trainable=True)
-
-    if self.use_bias:
-      if isinstance(self.bias_initializer, TrainableInitializer):
-        self.bias_initializer.build([self.units], self.dtype, self.add_weight)
-        if self.bias_regularizer is not None:
-          self.add_loss(create_regularization_loss_fn(
-              'bias', lambda: self.bias, self.bias_regularizer))
-      else:
-        self._bias = self.add_weight(
-            'bias',
-            shape=[self.units],
-            initializer=self.bias_initializer,
-            regularizer=self.bias_regularizer,
-            constraint=self.bias_constraint,
-            dtype=self.dtype,
-            trainable=True)
-
-    else:
-      self._bias = None
-    self.built = True
+  # TODO(trandustin): This name is not accurate. Rename or move functionality
+  # into random variables to resample/recreate their init ops.
+  def sample_weights(self):
+    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
+      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
+    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
+      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
+
+  def call(self, *args, **kwargs):
+    self.sample_weights()
+    return super(DenseReparameterization, self).call(*args, **kwargs)
+
+  def add_weight(self,
+                 name=None,
+                 shape=None,
+                 dtype=None,
+                 initializer=None,
+                 regularizer=None,
+                 **kwargs):
+    if isinstance(initializer, tf.keras.layers.Layer):
+      weight = initializer(shape, dtype)
+      self._trainable_weights.extend(initializer.trainable_weights)
+      self._non_trainable_weights.extend(initializer.non_trainable_weights)
+      if regularizer is not None:
+        self.add_loss(
+            create_regularization_loss_fn(name,
+                                          lambda: initializer(shape, dtype),
+                                          regularizer))
+      return weight
+    return super(DenseReparameterization, self).add_weight(
+        name=name,
+        shape=shape,
+        dtype=dtype,
+        initializer=initializer,
+        regularizer=regularizer,
+        **kwargs)
 
 
 class Conv2DReparameterization(tf.keras.layers.Conv2D):
@@ -553,87 +507,40 @@ def __init__(self,
         bias_constraint=get(bias_constraint),
         **kwargs)
 
-  @property
-  def kernel(self):
-    if isinstance(self.kernel_initializer, TrainableInitializer):
-      return self.kernel_initializer()
-    else:
-      return self._kernel
-
-  @property
-  def bias(self):
-    if isinstance(self.bias_initializer, TrainableInitializer):
-      return self.bias_initializer()
-    else:
-      return self._bias
-
-  def build(self, input_shape):
-    input_shape = tf.TensorShape(input_shape)
-    if self.data_format == 'channels_first':
-      channel_axis = 1
-    else:
-      channel_axis = -1
-    if input_shape.dims[channel_axis].value is None:
-      raise ValueError('The channel dimension of the inputs '
-                       'should be defined. Found `None`.')
-    input_dim = int(input_shape[channel_axis])
-    kernel_shape = self.kernel_size + (input_dim, self.filters)
-
-    if isinstance(self.kernel_initializer, TrainableInitializer):
-      self.kernel_initializer.build(kernel_shape,
-                                    self.dtype,
-                                    self.add_weight)
-      if self.kernel_regularizer is not None:
-        self.add_loss(create_regularization_loss_fn(
-            'kernel', lambda: self.kernel, self.kernel_regularizer))
-
-    else:
-      self._kernel = self.add_weight(
-          name='kernel',
-          shape=kernel_shape,
-          initializer=self.kernel_initializer,
-          regularizer=self.kernel_regularizer,
-          constraint=self.kernel_constraint,
-          trainable=True,
-          dtype=self.dtype)
-
-    if self.use_bias:
-      if isinstance(self.bias_initializer, TrainableInitializer):
-        self.bias_initializer.build((self.filters,),
-                                    self.dtype,
-                                    self.add_weight)
-        if self.bias_regularizer is not None:
-          self.add_loss(create_regularization_loss_fn(
-              'bias', lambda: self.bias, self.bias_regularizer))
-      else:
-        self._bias = self.add_weight(
-            name='bias',
-            shape=(self.filters,),
-            initializer=self.bias_initializer,
-            regularizer=self.bias_regularizer,
-            constraint=self.bias_constraint,
-            trainable=True,
-            dtype=self.dtype)
-    else:
-      self._bias = None
-
-    self.input_spec = tf.layers.InputSpec(ndim=self.rank + 2,
-                                          axes={channel_axis: input_dim})
-    if self.padding == 'causal':
-      op_padding = 'valid'
-    else:
-      op_padding = self.padding
-    if not isinstance(op_padding, (list, tuple)):
-      op_padding = op_padding.upper()
-    self._convolution_op = nn_ops.Convolution(
-        input_shape,
-        filter_shape=self.kernel.get_shape(),
-        dilation_rate=self.dilation_rate,
-        strides=self.strides,
-        padding=op_padding,
-        data_format=conv_utils.convert_data_format(self.data_format,
-                                                   self.rank + 2))
-    self.built = True
+  def sample_weights(self):
+    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
+      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
+    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
+      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
+
+  def call(self, *args, **kwargs):
+    self.sample_weights()
+    return super(Conv2DReparameterization, self).call(*args, **kwargs)
+
+  def add_weight(self,
+                 name=None,
+                 shape=None,
+                 dtype=None,
+                 initializer=None,
+                 regularizer=None,
+                 **kwargs):
+    if isinstance(initializer, tf.keras.layers.Layer):
+      weight = initializer(shape, dtype)
+      self._trainable_weights.extend(initializer.trainable_weights)
+      self._non_trainable_weights.extend(initializer.non_trainable_weights)
+      if regularizer is not None:
+        self.add_loss(
+            create_regularization_loss_fn(name,
+                                          lambda: initializer(shape, dtype),
+                                          regularizer))
+      return weight
+    return super(Conv2DReparameterization, self).add_weight(
+        name=name,
+        shape=shape,
+        dtype=dtype,
+        initializer=initializer,
+        regularizer=regularizer,
+        **kwargs)
 
 
 class GaussianProcess(tf.keras.layers.Layer):
@@ -854,49 +761,22 @@ def build(self, input_shape):
     input_dim = input_shape[-1]
     if isinstance(input_dim, tf.Dimension):
       input_dim = input_dim.value
-
-    if isinstance(self.kernel_initializer, TrainableInitializer):
-      self.kernel_initializer.build(
-          [input_dim, self.units * 4], self.dtype, self.add_weight)
-      self.kernel = self.kernel_initializer()
-      if self.kernel_regularizer is not None:
-        self.add_loss(create_regularization_loss_fn(
-            # Can't use the kernel directly because we actually need to create a
-            # new Edward RV.  The Dense layer already does this.
-            # Also note that the initializer is a callable.
-            'kernel', self.kernel_initializer, self.kernel_regularizer))
-
-    else:
-      self.kernel = self.add_weight(
-          shape=(input_dim, self.units * 4),
-          name='kernel',
-          initializer=self.kernel_initializer,
-          regularizer=self.kernel_regularizer,
-          constraint=self.kernel_constraint)
-
-    if isinstance(self.recurrent_initializer, TrainableInitializer):
-      self.recurrent_initializer.build(
-          [self.units, self.units * 4], self.dtype, self.add_weight)
-      self.recurrent_kernel = self.recurrent_initializer()
-      if self.recurrent_regularizer is not None:
-        self.add_loss(create_regularization_loss_fn(
-            # Can't use the kernel directly because we actually need to create a
-            # new Edward RV.  The Dense layer already does this.
-            # Also note that the initializer is a callable.
-            'recurrent_kernel', self.recurrent_initializer,
-            self.recurrent_regularizer))
-
-    else:
-      self.recurrent_kernel = self.add_weight(
-          shape=(self.units, self.units * 4),
-          name='recurrent_kernel',
-          initializer=self.recurrent_initializer,
-          regularizer=self.recurrent_regularizer,
-          constraint=self.recurrent_constraint)
+    self.kernel = self.add_weight(
+        shape=(input_dim, self.units * 4),
+        name='kernel',
+        initializer=self.kernel_initializer,
+        regularizer=self.kernel_regularizer,
+        constraint=self.kernel_constraint)
+    self.recurrent_kernel = self.add_weight(
+        shape=(self.units, self.units * 4),
+        name='recurrent_kernel',
+        initializer=self.recurrent_initializer,
+        regularizer=self.recurrent_regularizer,
+        constraint=self.recurrent_constraint)
 
     if self.use_bias:
-      if isinstance(self.bias_initializer, TrainableInitializer):
-        if self.unit_forget_bias:
+      if self.unit_forget_bias:
+        if isinstance(self.bias_initializer, tf.keras.layers.Layer):
           def bias_mean_initializer(_, *args, **kwargs):
             return tf.concat([
                 tf.keras.initializers.truncated_normal(
@@ -906,45 +786,35 @@ def bias_mean_initializer(_, *args, **kwargs):
                 tf.keras.initializers.truncated_normal(
                     stddev=1e-5)((self.units * 2,), *args, **kwargs),
             ], axis=0)
-          self.bias_initializer = TrainableNormal(
+          bias_initializer = TrainableNormal(
               mean_initializer=bias_mean_initializer)
-
-        self.bias_initializer.build(
-            [self.units * 4], self.dtype, self.add_weight)
-        self.bias = self.bias_initializer()
-        if self.bias_regularizer is not None:
-          self.add_loss(create_regularization_loss_fn(
-              # Can't use the bias directly because we actually need to create a
-              # new Edward RV.  The Dense layer already does this.
-              # Also note that the initializer is a callable.
-              'bias', self.bias_initializer, self.bias_regularizer))
-      else:
-        if self.unit_forget_bias:
+        else:
           def bias_initializer(_, *args, **kwargs):
             return tf.keras.backend.concatenate([
                 self.bias_initializer((self.units,), *args, **kwargs),
                 tf.keras.initializers.Ones()((self.units,), *args, **kwargs),
                 self.bias_initializer((self.units * 2,), *args, **kwargs),
             ])
-        else:
-          bias_initializer = self.bias_initializer
-        self.bias = self.add_weight(
-            shape=(self.units * 4,),
-            name='bias',
-            initializer=bias_initializer,
-            regularizer=self.bias_regularizer,
-            constraint=self.bias_constraint)
+      else:
+        bias_initializer = self.bias_initializer
+      self.bias = self.add_weight(
+          shape=(self.units * 4,),
+          name='bias',
+          initializer=bias_initializer,
+          regularizer=self.bias_regularizer,
+          constraint=self.bias_constraint)
     else:
       self.bias = None
     self.built = True
 
   def sample_weights(self):
-    if isinstance(self.kernel_initializer, TrainableInitializer):
-      self.kernel = self.kernel_initializer()
-    if isinstance(self.recurrent_initializer, TrainableInitializer):
-      self.recurrent_kernel = self.recurrent_initializer()
-    if isinstance(self.bias_initializer, TrainableInitializer):
-      self.bias = self.bias_initializer()
+    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
+      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
+    if isinstance(self.recurrent_initializer, tf.keras.layers.Layer):
+      self.recurrent_kernel = self.recurrent_initializer(
+          self.recurrent_kernel.shape, self.dtype)
+    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
+      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
 
   # NOTE: This will not be called in TF < 1.11.
   def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
@@ -954,7 +824,35 @@ def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
     return super(LSTMCellReparameterization, self).get_initial_state(
         inputs=inputs, batch_size=batch_size, dtype=dtype)
 
+  def add_weight(self,
+                 name=None,
+                 shape=None,
+                 dtype=None,
+                 initializer=None,
+                 regularizer=None,
+                 **kwargs):
+    if isinstance(initializer, tf.keras.layers.Layer):
+      weight = initializer(shape, dtype)
+      self._trainable_weights.extend(initializer.trainable_weights)
+      self._non_trainable_weights.extend(initializer.non_trainable_weights)
+      if regularizer is not None:
+        self.add_loss(
+            create_regularization_loss_fn(name,
+                                          lambda: initializer(shape, dtype),
+                                          regularizer))
+      return weight
+    return super(LSTMCellReparameterization, self).add_weight(
+        name=name,
+        shape=shape,
+        dtype=dtype,
+        initializer=initializer,
+        regularizer=regularizer,
+        **kwargs)
+
 
+# TODO(trandustin): Replace need for this function with
+# Layer._handle_weight_regularization. For Eager compatibility, random variable
+# __init__s cannot apply TF ops (cl/220898007).
 def create_regularization_loss_fn(name, variable_fn, regularizer_fn):
   """Create a regularization loss function.
 
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 4a2797e8a..4563cfa33 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -76,15 +76,17 @@ def testDenseReparameterizationKernel(
     layer.get_config()
 
   @test_utils.run_in_graph_and_eager_modes()
-  def testDenseReparameterizationKL(self):
-    inputs = tf.to_float(np.random.rand(5, 12))
+  def testDenseReparameterizationLoss(self):
+    features = tf.to_float(np.random.rand(5, 12))
+    labels = tf.to_float(np.random.rand(5, 10))
     layer = bayes.DenseReparameterization(10)
 
     # Imagine this is the 1st epoch.
-    with tf.GradientTape() as tape:
-      layer(inputs)  # first call forces a build, here inside this tape
-      layer(inputs)  # ensure robustness after multiple calls
-      loss = sum(layer.losses)
+    with tf.GradientTape(persistent=True) as tape:
+      predictions = layer(features)  # first call forces build
+      layer(features)  # ensure robustness after multiple calls
+      nll = tf.losses.mean_squared_error(labels, predictions)
+      kl = sum(layer.losses)
 
     variables = [layer.kernel_initializer.mean, layer.kernel_initializer.stddev]
     for v in variables:
@@ -92,14 +94,18 @@ def testDenseReparameterizationKL(self):
 
     # This will be fine, since the layer was built inside this tape, and thus
     # the distribution init ops were inside this tape.
-    grads = tape.gradient(loss, variables)
+    grads = tape.gradient(nll, variables)
+    for grad in grads:
+      self.assertIsNotNone(grad)
+    grads = tape.gradient(kl, variables)
     for grad in grads:
       self.assertIsNotNone(grad)
 
     # Imagine this is the 2nd epoch.
-    with tf.GradientTape() as tape:
-      layer(inputs)  # build won't be called again
-      loss = sum(layer.losses)
+    with tf.GradientTape(persistent=True) as tape:
+      predictions = layer(features)  # build is not called
+      nll = tf.losses.mean_squared_error(labels, predictions)
+      kl = sum(layer.losses)
 
     variables = [layer.kernel_initializer.mean, layer.kernel_initializer.stddev]
     for v in variables:
@@ -108,7 +114,10 @@ def testDenseReparameterizationKL(self):
     # This would fail, since the layer was built inside the tape from the 1st
     # epoch, and thus the distribution init ops were inside that tape instead of
     # this tape. By using a callable for the variable, this will no longer fail.
-    grads = tape.gradient(loss, variables)
+    grads = tape.gradient(nll, variables)
+    for grad in grads:
+      self.assertIsNotNone(grad)
+    grads = tape.gradient(kl, variables)
     for grad in grads:
       self.assertIsNotNone(grad)
 
@@ -273,18 +282,20 @@ def testLSTMCellReparameterization(
     cell.get_config()
 
   @test_utils.run_in_graph_and_eager_modes()
-  def testLSTMCellReparameterizationKL(self):
-    inputs = tf.to_float(np.random.rand(5, 1, 12))
+  def testLSTMCellReparameterizationLoss(self):
+    features = tf.to_float(np.random.rand(5, 1, 12))
+    labels = tf.to_float(np.random.rand(5, 10))
     cell = bayes.LSTMCellReparameterization(10)
     state = (tf.zeros([1, 10]), tf.zeros([1, 10]))
 
     # Imagine this is the 1st epoch.
-    with tf.GradientTape() as tape:
-      cell(inputs[:, 0, :], state)  # first call forces a build, inside the tape
-      cell(inputs[:, 0, :], state)  # ensure robustness after multiple calls
-      cell.get_initial_state(inputs[:, 0, :])
-      cell(inputs[:, 0, :], state)  # ensure robustness after multiple calls
-      loss = sum(cell.losses)
+    with tf.GradientTape(persistent=True) as tape:
+      predictions, _ = cell(features[:, 0, :], state)  # first call forces build
+      cell(features[:, 0, :], state)  # ensure robustness after multiple calls
+      cell.get_initial_state(features[:, 0, :])
+      cell(features[:, 0, :], state)  # ensure robustness after multiple calls
+      nll = tf.losses.mean_squared_error(labels, predictions)
+      kl = sum(cell.losses)
 
     variables = [
         cell.kernel_initializer.mean, cell.kernel_initializer.stddev,
@@ -295,14 +306,19 @@ def testLSTMCellReparameterizationKL(self):
 
     # This will be fine, since the layer was built inside this tape, and thus
     # the distribution init ops were inside this tape.
-    grads = tape.gradient(loss, variables)
+    grads = tape.gradient(nll, variables)
+    for grad in grads:
+      self.assertIsNotNone(grad)
+    grads = tape.gradient(kl, variables)
     for grad in grads:
       self.assertIsNotNone(grad)
 
     # Imagine this is the 2nd epoch.
-    with tf.GradientTape() as tape:
-      cell(inputs[:, 0, :], state)  # build won't be called again
-      loss = sum(cell.losses)
+    with tf.GradientTape(persistent=True) as tape:
+      cell.get_initial_state(features[:, 0, :])
+      predictions, _ = cell(features[:, 0, :], state)  # build is not called
+      nll = tf.losses.mean_squared_error(labels, predictions)
+      kl = sum(cell.losses)
 
     variables = [
         cell.kernel_initializer.mean, cell.kernel_initializer.stddev,
@@ -314,7 +330,10 @@ def testLSTMCellReparameterizationKL(self):
     # This would fail, since the layer was built inside the tape from the 1st
     # epoch, and thus the distribution init ops were inside that tape instead of
     # this tape. By using a callable for the variable, this will no longer fail.
-    grads = tape.gradient(loss, variables)
+    grads = tape.gradient(nll, variables)
+    for grad in grads:
+      self.assertIsNotNone(grad)
+    grads = tape.gradient(kl, variables)
     for grad in grads:
       self.assertIsNotNone(grad)
 

From 9d6ff58345ffc9b8575c43adc35690cba25f47f7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 8 Feb 2019 13:32:29 -0800
Subject: [PATCH 1652/2720] Implemented functions for area attention.

PiperOrigin-RevId: 233116679
---
 tensor2tensor/layers/area_attention.py      | 428 ++++++++++++++++++++
 tensor2tensor/layers/area_attention_test.py | 276 +++++++++++++
 2 files changed, 704 insertions(+)
 create mode 100644 tensor2tensor/layers/area_attention.py
 create mode 100644 tensor2tensor/layers/area_attention_test.py

diff --git a/tensor2tensor/layers/area_attention.py b/tensor2tensor/layers/area_attention.py
new file mode 100644
index 000000000..8372718f2
--- /dev/null
+++ b/tensor2tensor/layers/area_attention.py
@@ -0,0 +1,428 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Utilities for area attention."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from six.moves import range  # pylint: disable=redefined-builtin
+from tensor2tensor.layers import common_layers
+import tensorflow as tf
+
+
+def lengths_to_area_mask(feature_length, length, max_area_size):
+  """Generates a non-padding mask for areas based on lengths.
+
+  Args:
+    feature_length: a tensor of [batch_size]
+    length: the length of the batch
+    max_area_size: the maximum area size considered
+  Returns:
+    mask: a tensor in shape of [batch_size, num_areas]
+  """
+
+  paddings = tf.cast(tf.expand_dims(
+      tf.logical_not(
+          tf.sequence_mask(feature_length, maxlen=length)), 2), tf.float32)
+  _, _, area_sum, _, _ = compute_area_features(paddings,
+                                               max_area_width=max_area_size)
+  mask = tf.squeeze(tf.logical_not(tf.cast(area_sum, tf.bool)), [2])
+  return mask
+
+
+def _max_pool_one_shape(features_2d, area_width, area_height, batch_size,
+                        width, height, depth, name=None):
+  """Computes area max for features_2d.
+
+  Args:
+    features_2d: a Tensor in a shape of [batch_size, height * width, depth].
+    area_width: the max width allowed for an area.
+    area_height: the max height allowed for an area.
+    batch_size: the batch size.
+    width: the width of the memory.
+    height: the height of the memory.
+    depth: the depth of the features.
+    name: the op name.
+  Returns:
+    max_tensor: A Tensor of shape [batch_size, num_areas, depth]
+  """
+  with tf.name_scope(name, default_name="max_pool_one_shape"):
+    images = []
+    for y_shift in range(area_height):
+      image_height = tf.maximum(height - area_height + 1 + y_shift, 0)
+      for x_shift in range(area_width):
+        image_width = tf.maximum(width - area_width + 1 + x_shift, 0)
+        area = features_2d[:, y_shift:image_height, x_shift:image_width, :]
+        flatten_area = tf.reshape(area, [batch_size, -1, depth, 1])
+        images.append(flatten_area)
+    image_tensor = tf.concat(images, axis=3)
+    max_tensor = tf.reduce_max(image_tensor, axis=3)
+  return max_tensor
+
+
+def max_pool(features, max_area_width, max_area_height=1, height=1,
+             name=None):
+  """Computes area max for features.
+
+  Args:
+    features: a Tensor in a shape of [batch_size, height * width, depth].
+    max_area_width: the max width allowed for an area.
+    max_area_height: the max height allowed for an area.
+    height: the height of the image.
+    name: the namescope.
+  Returns:
+    max_results: A Tensor of shape [batch_size, num_areas, depth]
+    area_heights: A Tensor of shape [batch_size, num_areas, 1]
+    area_widths: A Tensor of shape [batch_size, num_areas, 1]
+  """
+  with tf.name_scope(name, default_name="max_pool"):
+    feature_shape = common_layers.shape_list(features)
+    batch_size = feature_shape[0]
+    length = feature_shape[-2]
+    depth = feature_shape[-1]
+    width = length // height
+    features_2d = tf.reshape(features, [batch_size, height, width, depth])
+    height_list = []
+    width_list = []
+    max_list = []
+    size_tensor = tf.ones_like(features_2d[:, :, :, 0], dtype=tf.int32)
+    for area_height in range(max_area_height):
+      for area_width in range(max_area_width):
+        max_tensor = _max_pool_one_shape(features_2d,
+                                         area_width=area_width + 1,
+                                         area_height=area_height + 1,
+                                         batch_size=batch_size,
+                                         width=width,
+                                         height=height,
+                                         depth=depth)
+        max_list.append(
+            tf.reshape(max_tensor, [batch_size, -1, depth]))
+        height_list.append(
+            tf.reshape(
+                size_tensor[:, area_height:, area_width:] *\
+                (area_height + 1), [batch_size, -1]))
+        width_list.append(
+            tf.reshape(
+                size_tensor[:, area_height:, area_width:] *\
+                (area_width + 1), [batch_size, -1]))
+    max_results = tf.concat(max_list, axis=1)
+    area_heights = tf.expand_dims(tf.concat(height_list, axis=1), 2)
+    area_widths = tf.expand_dims(tf.concat(width_list, axis=1), 2)
+  return max_results, area_heights, area_widths
+
+
+def _compute_sum_image(features, max_area_width, max_area_height=1, height=1,
+                       name=None):
+  """Computes area sums for features.
+
+  Args:
+    features: a Tensor in a shape of [batch_size, height * width, depth].
+    max_area_width: the max width allowed for an area.
+    max_area_height: the max height allowed for an area.
+    height: the height of the image.
+    name: the namescope.
+  Returns:
+    sum_image: A Tensor of shape [batch_size, num_areas, depth]
+    area_heights: A Tensor of shape [batch_size, num_areas, 1]
+    area_widths: A Tensor of shape [batch_size, num_areas, 1]
+  """
+  with tf.name_scope(name, default_name="compute_sum_image"):
+    feature_shape = common_layers.shape_list(features)
+    batch_size = feature_shape[0]
+    length = feature_shape[-2]
+    depth = feature_shape[-1]
+    width = length // height
+    features_2d = tf.reshape(features, [batch_size, height, width, depth])
+    width_cum = tf.cumsum(features_2d, axis=-2, name="compute_integral_h")
+    integral_image = tf.cumsum(width_cum, axis=-3, name="compute_integral_v")
+    padded_image = tf.pad(
+        integral_image, [[0, 0], [1, 0], [1, 0], [0, 0]], constant_values=0)
+    height_list = []
+    width_list = []
+    dst_images = []
+    src_images_diag = []
+    src_images_h = []
+    src_images_v = []
+    size_tensor = tf.ones_like(padded_image[:, :, :, 0],
+                               dtype=tf.int32)
+    for area_height in range(max_area_height):
+      for area_width in range(max_area_width):
+        dst_images.append(
+            tf.reshape(
+                padded_image[:, area_height + 1:, area_width + 1:, :],
+                [batch_size, -1, depth]))
+        src_images_diag.append(
+            tf.reshape(
+                padded_image[:, :-area_height - 1, :-area_width - 1, :],
+                [batch_size, -1, depth]))
+        src_images_h.append(
+            tf.reshape(
+                padded_image[:, area_height + 1:, :-area_width - 1, :],
+                [batch_size, -1, depth]))
+        src_images_v.append(
+            tf.reshape(
+                padded_image[:, :-area_height - 1, area_width + 1:, :],
+                [batch_size, -1, depth]))
+        height_list.append(
+            tf.reshape(
+                size_tensor[:, area_height + 1:, area_width + 1:] *\
+                (area_height + 1), [batch_size, -1]))
+        width_list.append(
+            tf.reshape(
+                size_tensor[:, area_height + 1:, area_width + 1:] *\
+                (area_width + 1), [batch_size, -1]))
+    sum_image = tf.subtract(
+        tf.concat(dst_images, axis=1) + tf.concat(src_images_diag, axis=1),
+        tf.concat(src_images_v, axis=1) + tf.concat(src_images_h, axis=1))
+    area_heights = tf.expand_dims(tf.concat(height_list, axis=1), 2)
+    area_widths = tf.expand_dims(tf.concat(width_list, axis=1), 2)
+  return sum_image, area_heights, area_widths
+
+
+def compute_area_features(features, max_area_width, max_area_height=1, height=1,
+                          epsilon=1e-6):
+  """Computes features for each area.
+
+  Args:
+    features: a Tensor in a shape of [batch_size, height * width, depth].
+    max_area_width: the max width allowed for an area.
+    max_area_height: the max height allowed for an area.
+    height: the height of the image.
+    epsilon: the epsilon added to the variance for computing standard deviation.
+  Returns:
+    area_mean: A Tensor of shape [batch_size, num_areas, depth]
+    area_std: A Tensor of shape [batch_size, num_areas, depth]
+    area_sum: A Tensor of shape [batch_size, num_areas, depth]
+    area_heights: A Tensor of shape [batch_size, num_areas, 1]
+    area_widths: A Tensor of shape [batch_size, num_areas, 1]
+  """
+  with tf.name_scope("compute_area_features"):
+    tf.logging.info("area_attention compute_area_features: %d x %d",
+                    max_area_height, max_area_width)
+    area_sum, area_heights, area_widths = _compute_sum_image(
+        features, max_area_width=max_area_width,
+        max_area_height=max_area_height, height=height)
+    area_squared_sum, _, _ = _compute_sum_image(
+        tf.pow(features, 2), max_area_width=max_area_width,
+        max_area_height=max_area_height, height=height)
+    sizes = tf.multiply(area_heights, area_widths)
+    float_area_sizes = tf.to_float(sizes)
+    area_mean = tf.div(area_sum, float_area_sizes)
+    s2_n = tf.div(area_squared_sum, float_area_sizes)
+    area_variance = tf.subtract(s2_n, tf.pow(area_mean, 2))
+    area_std = tf.sqrt(tf.abs(area_variance) + epsilon)
+    return area_mean, area_std, area_sum, area_heights, area_widths
+
+
+def compute_area_key(features, max_area_width, max_area_height=1, height=1,
+                     mode="mean", training=True, name=None):
+  """Computes the key for each area.
+
+  Args:
+    features: a Tensor in a shape of [batch_size, height * width, depth].
+    max_area_width: the max width allowed for an area.
+    max_area_height: the max height allowed for an area.
+    height: the height of the image.
+    mode: whether to combine different area features or only use
+        the vector mean of each area, which can be "mean", "concat", "sum",
+        "sample_concat", and "sample_sum".
+    training: indicating if it is in the training mode.
+    name: the name for setting the variable scope.
+  Returns:
+    area_key: a Tensor in the shape of [batch_size, num_areas, depth]
+  """
+
+  tf.logging.info("area_attention mode=%s", mode)
+  area_mean, area_std, _, area_heights, area_widths =\
+      compute_area_features(features, max_area_width=max_area_width,
+                            max_area_height=max_area_height, height=height)
+  if mode == "mean":
+    return area_mean
+  elif mode == "max":
+    area_max, _, _ = max_pool(features, max_area_width=max_area_width,
+                              max_area_height=max_area_height, height=height)
+    return area_max
+  elif mode == "sample":
+    if training:
+      area_mean += (area_std * tf.random_normal(tf.shape(area_std)))
+    return area_mean
+  with tf.variable_scope(
+      name, default_name="combine_area_features",
+      values=[area_mean, area_std, area_heights, area_widths]):
+    depth = common_layers.shape_list(area_mean)[-1]
+    height_embed = tf.nn.embedding_lookup(
+        params=tf.get_variable("area_height_emb",
+                               [max_area_height, depth // 2]),
+        ids=area_heights[:, :, 0] - 1)
+    width_embed = tf.nn.embedding_lookup(
+        params=tf.get_variable("area_width_emb",
+                               [max_area_width, depth // 2]),
+        ids=area_widths[:, :, 0] - 1)
+    size_embed = tf.concat([height_embed, width_embed], -1)
+    if mode == "concat":
+      feature_concat = tf.concat([area_mean, area_std, size_embed], -1)
+    elif mode == "max_concat":
+      area_max, _, _ = max_pool(features, max_area_width=max_area_width,
+                                max_area_height=max_area_height, height=height)
+      feature_concat = tf.concat([area_max, size_embed], -1)
+    elif mode == "sum":
+      feature_concat = size_embed + area_mean + area_std
+    elif mode == "sample_concat":
+      if training:
+        area_mean += (area_std * tf.random_normal(tf.shape(area_std)))
+      feature_concat = tf.concat([area_mean, size_embed], -1)
+    elif mode == "sample_sum":
+      if training:
+        area_mean += (area_std * tf.random_normal(tf.shape(area_std)))
+      feature_concat = area_mean + size_embed
+    else:
+      raise ValueError("Unsupported area key mode=%s" % mode)
+    feature_hidden = tf.layers.dense(inputs=feature_concat,
+                                     units=depth,
+                                     activation=tf.nn.relu)
+    area_key = tf.layers.dense(feature_hidden, units=depth)
+    return area_key
+
+
+def dot_product_area_attention(q,
+                               k,
+                               v,
+                               bias,
+                               dropout_rate=0.0,
+                               image_shapes=None,
+                               name=None,
+                               attention_image_summary=None,
+                               save_weights_to=None,
+                               dropout_broadcast_dims=None,
+                               max_area_width=1,
+                               max_area_height=1,
+                               memory_height=1,
+                               area_key_mode="mean",
+                               area_value_mode="sum",
+                               top_k_areas=0,
+                               area_temperature=1.0,
+                               training=True):
+  """Dot-product area attention.
+
+  Args:
+    q: Tensor with shape [..., length_q, depth_k].
+    k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
+      match with q.
+    v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
+      match with q.
+    bias: bias Tensor (see attention_bias())
+    dropout_rate: a float.
+    image_shapes: optional tuple of integer scalars.
+      see comments for attention_image_summary()
+    name: an optional string
+    attention_image_summary: the callback for making image summary of attention.
+    save_weights_to: an optional dictionary to capture attention weights
+      for visualization; the weights tensor will be appended there under
+      a string key created from the variable scope (including name).
+    dropout_broadcast_dims: an optional list of integers less than rank of q.
+      Specifies in which dimensions to broadcast the dropout decisions.
+    max_area_width: the max width allowed for an area.
+    max_area_height: the max height allowed for an area.
+    memory_height: the height of the memory.
+    area_key_mode: the mode for computing area keys, which can be "mean",
+      "concat", "sum", "sample_concat", and "sample_sum".
+    area_value_mode: the mode for computing area values, which can be either
+      "mean", or "sum".
+    top_k_areas: Use the top key areas for attention.
+    area_temperature: the temperature for attention softmax.
+    training: indicating if it is in the training mode.
+  Returns:
+    Tensor with shape [..., length_q, depth_v].
+  """
+
+  tf.logging.info("dot_product_area_attention: "
+                  "area_h=%d, area_w=%d, mem_h=%d, "
+                  "area_key_mode=%s, area_value_mode=%s, "
+                  "area_temperature=%f, top_k_areas=%d",
+                  max_area_height, max_area_width, memory_height,
+                  area_key_mode, area_value_mode,
+                  area_temperature, top_k_areas)
+  with tf.variable_scope(
+      name, default_name="dot_product_area_attention",
+      values=[q, k, v]) as scope:
+    mem_shape = common_layers.shape_list(k)
+    batch_size = mem_shape[0]
+    head_size = mem_shape[1]
+    length = mem_shape[2]
+    depth = mem_shape[3]
+    k_area = compute_area_key(
+        tf.reshape(k, [-1, length, depth]),
+        max_area_width=max_area_width,
+        max_area_height=max_area_height,
+        height=memory_height,
+        mode=area_key_mode,
+        training=training)
+    if area_value_mode == "mean":
+      v_area, _, _, _, _ = compute_area_features(
+          tf.reshape(v, [-1, length, depth]), max_area_width=max_area_width,
+          max_area_height=max_area_height, height=memory_height)
+    elif area_value_mode == "max":
+      v_area, _, _ = max_pool(tf.reshape(v, [-1, length, depth]),
+                              max_area_width=max_area_width,
+                              max_area_height=max_area_height,
+                              height=memory_height)
+    elif area_value_mode == "sum":
+      _, _, v_area, _, _ = compute_area_features(
+          tf.reshape(v, [-1, length, depth]), max_area_width=max_area_width,
+          max_area_height=max_area_height, height=memory_height)
+    else:
+      raise ValueError("Unsupported area value mode=%s" % area_value_mode)
+    k = tf.reshape(k_area, [batch_size, head_size, -1, depth])
+    v = tf.reshape(v_area, [batch_size, head_size, -1, depth])
+    logits = tf.matmul(q, k, transpose_b=True)  # [..., length_q, length_kv]
+    if bias is not None:
+      bias = common_layers.cast_like(bias, logits)
+      with tf.name_scope("compute_area_att_bias", values=[bias]):
+        bias_shape = common_layers.shape_list(bias)
+        mem_length = bias_shape[-1]
+        bias_values = tf.reshape(
+            tf.to_float(tf.less(bias, -1)), [-1, mem_length, 1])
+        _, _, padding_sum, _, _ = compute_area_features(
+            bias_values, max_area_width=max_area_width,
+            max_area_height=max_area_height, height=memory_height)
+        bias = tf.where(
+            tf.cast(tf.to_int32(padding_sum), tf.bool),
+            tf.fill(tf.shape(padding_sum), -np.inf),
+            tf.zeros_like(padding_sum, dtype=tf.float32))
+        bias = tf.reshape(bias,
+                          [bias_shape[0], bias_shape[1],
+                           bias_shape[2], -1])
+      logits += bias
+    logits = logits / area_temperature
+    weights = tf.nn.softmax(logits, name="attention_weights")
+    if top_k_areas > 0:
+      tf.logging.info("area_attention top_k_areas=%d", top_k_areas)
+      top_k = tf.minimum(common_layers.shape_list(weights)[-1], top_k_areas)
+      top_weights, _ = tf.nn.top_k(weights, k=top_k)
+      min_values = tf.reduce_min(top_weights, -1, keepdims=True)
+      weights = tf.where(tf.greater_equal(weights, min_values),
+                         weights, tf.zeros_like(weights))
+      weights = tf.div(weights, tf.reduce_sum(weights, -1, keepdims=True))
+    if save_weights_to is not None:
+      save_weights_to[scope.name] = weights
+      save_weights_to[scope.name + "/logits"] = logits
+    # Drop out attention links for each head.
+    weights = common_layers.dropout_with_broadcast_dims(
+        weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
+    if common_layers.should_generate_summaries() and attention_image_summary:
+      attention_image_summary(weights, image_shapes)
+    return tf.matmul(weights, v)
diff --git a/tensor2tensor/layers/area_attention_test.py b/tensor2tensor/layers/area_attention_test.py
new file mode 100644
index 000000000..1b464a6a0
--- /dev/null
+++ b/tensor2tensor/layers/area_attention_test.py
@@ -0,0 +1,276 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 area attention."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+import numpy as np
+from tensor2tensor.layers import area_attention
+import tensorflow as tf
+
+
+class AreaAttentionTest(parameterized.TestCase, tf.test.TestCase):
+
+  def testComputeAreaFeatures1D(self):
+    features = tf.constant([[[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]],
+                            [[1.1, 2.1], [3.1, 4.1], [5.1, 6.1], [7.1, 8.1],
+                             [9.1, 10.1]]],
+                           dtype=tf.float32)
+    area_mean, area_std, area_sum, area_height, area_widths = (
+        area_attention.compute_area_features(features, max_area_width=3,
+                                             epsilon=0.))
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      res1, res2, res3, res4, res5 = session.run([area_mean, area_std, area_sum,
+                                                  area_height, area_widths])
+    self.assertAllClose(((((1, 2), (3, 4), (5, 6), (7, 8), (9, 10),
+                           (2, 3), (4, 5), (6, 7), (8, 9),
+                           (3, 4), (5, 6), (7, 8)),
+                          ((1.1, 2.1), (3.1, 4.1), (5.1, 6.1), (7.1, 8.1),
+                           (9.1, 10.1),
+                           (2.1, 3.1), (4.1, 5.1), (6.1, 7.1), (8.1, 9.1),
+                           (3.1, 4.1), (5.1, 6.1), (7.1, 8.1)))),
+                        res1,
+                        msg="mean_1d")
+    expected_std = np.array([[[0, 0], [0, 0], [0, 0], [0, 0], [0, 0],
+                              [1, 1], [1, 1], [1, 1], [1, 1],
+                              [1.63299, 1.63299], [1.63299, 1.63299],
+                              [1.63299, 1.63299]],
+                             [[0, 0], [0, 0], [0, 0], [0, 0], [0, 0],
+                              [1, 1], [1, 1], [1, 1], [1, 1],
+                              [1.63299, 1.63299], [1.63299, 1.63299],
+                              [1.63299, 1.63299]]])
+    self.assertAllClose(expected_std, res2, atol=1e-2, msg="std_1d")
+    self.assertAllClose([[[1, 2], [3, 4], [5, 6], [7, 8], [9, 10],
+                          [4, 6], [8, 10], [12, 14], [16, 18],
+                          [9, 12], [15, 18], [21, 24]],
+                         [[1.1, 2.1], [3.1, 4.1], [5.1, 6.1], [7.1, 8.1],
+                          [9.1, 10.1],
+                          [4.2, 6.2], [8.2, 10.2], [12.2, 14.2], [16.2, 18.2],
+                          [9.3, 12.3], [15.3, 18.3], [21.3, 24.3]]],
+                        res3,
+                        msg="sum_1d")
+    self.assertAllEqual([[[1], [1], [1], [1], [1],
+                          [1], [1], [1], [1],
+                          [1], [1], [1]],
+                         [[1], [1], [1], [1], [1],
+                          [1], [1], [1], [1],
+                          [1], [1], [1]]],
+                        res4,
+                        msg="height_1d")
+    self.assertAllEqual([[[1], [1], [1], [1], [1],
+                          [2], [2], [2], [2],
+                          [3], [3], [3]],
+                         [[1], [1], [1], [1], [1],
+                          [2], [2], [2], [2],
+                          [3], [3], [3]]],
+                        res5,
+                        msg="width_1d")
+
+  def testComputeAreaFeatures2D(self):
+    features = tf.constant([[[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]],
+                            [[1.1, 2.1], [3.1, 4.1], [5.1, 6.1], [7.1, 8.1],
+                             [9.1, 10.1], [11.1, 12.1]]],
+                           dtype=tf.float32)
+    area_mean, area_std, area_sum, area_height, area_widths = (
+        area_attention.compute_area_features(features, max_area_width=3,
+                                             max_area_height=2,
+                                             height=2, epsilon=0.))
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      res1, _, res3, res4, res5 = session.run([area_mean, area_std, area_sum,
+                                               area_height, area_widths])
+    expected_means = [[[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12],
+                       [2, 3], [4, 5], [8, 9], [10, 11],
+                       [3, 4], [9, 10],
+                       [4, 5], [6, 7], [8, 9],
+                       [5, 6], [7, 8],
+                       [6, 7]],
+                      [[1.1, 2.1], [3.1, 4.1], [5.1, 6.1], [7.1, 8.1],
+                       [9.1, 10.1], [11.1, 12.1],
+                       [2.1, 3.1], [4.1, 5.1], [8.1, 9.1], [10.1, 11.1],
+                       [3.1, 4.1], [9.1, 10.1],
+                       [4.1, 5.1], [6.1, 7.1], [8.1, 9.1],
+                       [5.1, 6.1], [7.1, 8.1],
+                       [6.1, 7.1]]]
+    self.assertAllClose(expected_means, res1, msg="mean_1d")
+    expected_heights = [[[1], [1], [1], [1], [1], [1],
+                         # 1x2
+                         [1], [1], [1], [1],
+                         # 1x3
+                         [1], [1],
+                         # 2x1
+                         [2], [2], [2],
+                         # 2x2
+                         [2], [2],
+                         # 2x3
+                         [2]],
+                        [[1], [1], [1], [1], [1], [1],
+                         # 1x2
+                         [1], [1], [1], [1],
+                         # 1x3
+                         [1], [1],
+                         # 2x1
+                         [2], [2], [2],
+                         # 2x2
+                         [2], [2],
+                         # 2x3
+                         [2]]]
+    self.assertAllEqual(expected_heights, res4, msg="height_1d")
+    expected_widths = [[[1], [1], [1], [1], [1], [1],
+                        # 1x2
+                        [2], [2], [2], [2],
+                        # 1x3
+                        [3], [3],
+                        # 2x1
+                        [1], [1], [1],
+                        # 2x2
+                        [2], [2],
+                        # 2x3
+                        [3]],
+                       [[1], [1], [1], [1], [1], [1],
+                        # 1x2
+                        [2], [2], [2], [2],
+                        # 1x3
+                        [3], [3],
+                        # 2x1
+                        [1], [1], [1],
+                        # 2x2
+                        [2], [2],
+                        # 2x3
+                        [3]]]
+    self.assertAllEqual(expected_widths, res5, msg="width_1d")
+    sizes = np.multiply(np.array(expected_heights), np.array(expected_widths))
+    expected_sums = np.multiply(np.array(expected_means), sizes)
+    self.assertAllClose(expected_sums, res3, msg="sum_1d")
+
+  def testAreaMean(self):
+    batch_size = 256
+    feature_len = 100
+    memory_height = 10
+    heads = 2
+    key_len = 2
+    depth = 128
+    max_area_height = 3
+    max_area_width = 3
+    queries = tf.random_uniform([batch_size, heads, key_len, depth],
+                                minval=-10.0, maxval=10.0)
+    features = tf.random_uniform([batch_size, heads, feature_len, depth],
+                                 minval=-10.0, maxval=10.0)
+    target_values = tf.random_uniform([batch_size, heads, key_len, depth],
+                                      minval=-0.2, maxval=0.2)
+    keys = tf.layers.dense(features, units=depth)
+    values = tf.layers.dense(features, units=depth)
+    mean_attention = area_attention.dot_product_area_attention(
+        queries, keys, values,
+        bias=None,
+        area_key_mode="mean",
+        name="mean_key",
+        max_area_width=max_area_width,
+        max_area_height=max_area_height,
+        memory_height=memory_height)
+    mean_gradients = tf.gradients(
+        tf.reduce_mean(
+            tf.pow(target_values - mean_attention, 2)), features)
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      result = session.run([mean_gradients])
+    self.assertFalse(np.any(np.logical_not(np.isfinite(result))))
+
+  def test2DAreaMax(self):
+    batch_size = 256
+    feature_len = 100
+    memory_height = 10
+    heads = 2
+    key_len = 6
+    depth = 128
+    max_area_height = 3
+    max_area_width = 3
+    queries = tf.random_uniform([batch_size, heads, key_len, depth],
+                                minval=-10.0, maxval=10.0)
+    features = tf.random_uniform([batch_size, heads, feature_len, depth],
+                                 minval=-10.0, maxval=10.0)
+    target_values = tf.random_uniform([batch_size, heads, key_len, depth],
+                                      minval=-0.2, maxval=0.2)
+    keys = tf.layers.dense(features, units=depth)
+    values = tf.layers.dense(features, units=depth)
+    max_attention = area_attention.dot_product_area_attention(
+        queries, keys, values,
+        bias=None,
+        area_key_mode="max",
+        area_value_mode="max",
+        name="max_key",
+        max_area_width=max_area_width,
+        max_area_height=max_area_height,
+        memory_height=memory_height)
+    max_gradients = tf.gradients(tf.reduce_mean(
+        tf.pow(target_values - max_attention, 2)), features)
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      result1, result2 = session.run([max_gradients, max_attention])
+    self.assertFalse(np.any(np.logical_not(np.isfinite(result1))))
+    self.assertFalse(np.any(np.logical_not(np.isfinite(result2))))
+
+  def test1DAreaMax(self):
+    batch_size = 256
+    feature_len = 100
+    heads = 2
+    key_len = 15
+    depth = 128
+    max_area_width = 3
+    queries = tf.random_uniform([batch_size, heads, key_len, depth],
+                                minval=-10.0, maxval=10.0)
+    features = tf.random_uniform([batch_size, heads, feature_len, depth],
+                                 minval=-10.0, maxval=10.0)
+    feature_length = tf.constant(
+        np.concatenate(
+            (np.random.randint(max_area_width, feature_len, [batch_size - 1]),
+             np.array([feature_len])), axis=0), tf.int32)
+    base_mask = tf.expand_dims(tf.sequence_mask(feature_length), 1)
+    mask = tf.expand_dims(base_mask, 3)
+    mask = tf.tile(mask, [1, heads, 1, depth])
+    features = tf.where(mask, features, tf.zeros_like(features))
+    # [batch, 1, 1, memory_length]
+    bias_mask = tf.expand_dims(base_mask, 1)
+    bias = tf.where(
+        bias_mask,
+        tf.zeros_like(bias_mask, tf.float32),
+        tf.ones_like(bias_mask, tf.float32) * -1e9)
+    target_values = tf.random_uniform([batch_size, heads, key_len, depth],
+                                      minval=-0.2, maxval=0.2)
+    keys = tf.layers.dense(features, units=depth)
+    values = tf.layers.dense(features, units=depth)
+    max_attention = area_attention.dot_product_area_attention(
+        queries, keys, values,
+        bias=bias,
+        area_key_mode="max",
+        area_value_mode="max",
+        name="max_key",
+        max_area_width=max_area_width)
+    max_gradients = tf.gradients(
+        tf.reduce_mean(
+            tf.pow(target_values - max_attention, 2)), features)
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      result1, result2 = session.run([max_gradients, max_attention])
+    self.assertFalse(np.any(np.logical_not(np.isfinite(result1))))
+    self.assertFalse(np.any(np.logical_not(np.isfinite(result2))))
+
+if __name__ == "__main__":
+  tf.test.main()

From d54ed849043baebe52870b417a0e6a7cbe68bc03 Mon Sep 17 00:00:00 2001
From: Dumitru Erhan <dumitru@google.com>
Date: Fri, 8 Feb 2019 14:22:36 -0800
Subject: [PATCH 1653/2720] s/TrueAdam/true_adam

PiperOrigin-RevId: 233126026
---
 tensor2tensor/models/image_transformer_2d.py | 2 +-
 tensor2tensor/models/transformer.py          | 8 ++++----
 tensor2tensor/models/video/sv2p_params.py    | 2 +-
 tensor2tensor/models/xception.py             | 2 +-
 4 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index fe5d31444..467febe16 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -783,7 +783,7 @@ def imagetransformer2d_tiny():
 
 def update_hparams_for_tpu(hparams):
   hparams.use_pad_remover = False  # where op not supported
-  hparams.optimizer = "TrueAdam"
+  hparams.optimizer = "true_adam"
   hparams.batch_size = 4
 
 
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index a0da5b1cc..d7644db93 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1876,7 +1876,7 @@ def transformer_tall_finetune_tied():
   hparams.multiproblem_target_eval_only = True
   hparams.multiproblem_reweight_label_loss = True
   hparams.multiproblem_label_weight = 1.0
-  hparams.optimizer = "TrueAdam"
+  hparams.optimizer = "true_adam"
   return hparams
 
 
@@ -1895,7 +1895,7 @@ def transformer_tall_train_tied():
   hparams.multiproblem_target_eval_only = True
   hparams.multiproblem_reweight_label_loss = True
   hparams.multiproblem_label_weight = 1.0
-  hparams.optimizer = "TrueAdam"
+  hparams.optimizer = "true_adam"
   return hparams
 
 
@@ -1905,7 +1905,7 @@ def transformer_tall_finetune_uniencdec():
   hparams = transformer_tall()
   hparams.max_input_seq_length = 750
   hparams.max_target_seq_length = 100
-  hparams.optimizer = "TrueAdam"
+  hparams.optimizer = "true_adam"
   hparams.learning_rate_schedule = ("linear_warmup*constant*cosdecay")
   hparams.learning_rate_decay_steps = 80000
   hparams.learning_rate_constant = 5e-5
@@ -1920,7 +1920,7 @@ def transformer_tall_train_uniencdec():
   hparams = transformer_tall()
   hparams.max_input_seq_length = 750
   hparams.max_target_seq_length = 100
-  hparams.optimizer = "TrueAdam"
+  hparams.optimizer = "true_adam"
   hparams.learning_rate_schedule = ("linear_warmup*constant*cosdecay")
   hparams.learning_rate_decay_steps = 150000
   hparams.learning_rate_constant = 2e-4
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 674b5b211..f5687865b 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -27,7 +27,7 @@
 def next_frame_sv2p():
   """SV2P model hparams."""
   hparams = basic_stochastic.next_frame_basic_stochastic()
-  hparams.optimizer = "TrueAdam"
+  hparams.optimizer = "true_adam"
   hparams.learning_rate_schedule = "constant"
   hparams.learning_rate_constant = 1e-3
   hparams.video_num_input_frames = 1
diff --git a/tensor2tensor/models/xception.py b/tensor2tensor/models/xception.py
index 7c6c6ca40..9361745e0 100644
--- a/tensor2tensor/models/xception.py
+++ b/tensor2tensor/models/xception.py
@@ -182,5 +182,5 @@ def xception_tiny_tpu():
   hparams.batch_size = 2
   hparams.num_hidden_layers = 2
   hparams.hidden_size = 128
-  hparams.optimizer = "TrueAdam"
+  hparams.optimizer = "true_adam"
   return hparams

From d86eac3dfd02c8dea8f562087837b01d4326426f Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 8 Feb 2019 18:57:24 -0800
Subject: [PATCH 1654/2720] Merge of PR #1444

PiperOrigin-RevId: 233163674
---
 tensor2tensor/data_generators/problem.py   | 4 ++--
 tensor2tensor/layers/latent_layers_test.py | 5 ++---
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 6c651ba74..7c1327504 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -950,7 +950,7 @@ def _reverse_problem_hparams(p_hparams):
   # 'target', and each intended feature to swap has feature name 'input'.
   # In the future, remove need for this behavior.
   reversed_modality = {}
-  for feature_name in six.iterkeys(p.modality):
+  for feature_name in p.modality:
     reversed_feature_name = feature_name.replace("target", "input")
     if "target" in feature_name and reversed_feature_name in p.modality:
       reversed_modality[feature_name] = p.modality[reversed_feature_name]
@@ -962,7 +962,7 @@ def _reverse_problem_hparams(p_hparams):
 
   # Swap vocab sizes.
   reversed_vocab_size = {}
-  for feature_name in six.iterkeys(p.vocab_size):
+  for feature_name in p.vocab_size:
     reversed_feature_name = feature_name.replace("target", "input")
     if "target" in feature_name and reversed_feature_name in p.vocab_size:
       reversed_vocab_size[feature_name] = p.vocab_size[reversed_feature_name]
diff --git a/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
index 2e1ddc77b..78a83d008 100644
--- a/tensor2tensor/layers/latent_layers_test.py
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -20,7 +20,6 @@
 from __future__ import print_function
 
 import functools
-import six
 
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import discretization
@@ -138,8 +137,7 @@ def testTransformerAutoencoder(self):
     decoder_output, losses, cache = latent_layers.transformer_autoencoder(
         inputs, targets, target_space_id, hparams)
 
-    self.assertEqual(set(six.iterkeys(losses)),
-                     {"extra", "extra_loss", "latent_pred"})
+    self.assertEqual(set(losses), {"extra", "extra_loss", "latent_pred"})
 
     self.evaluate(tf.global_variables_initializer())
     decoder_output_, extra_loss_, latent_pred_ = self.evaluate(
@@ -154,5 +152,6 @@ def testTransformerAutoencoder(self):
     self.assertAllGreaterEqual(latent_pred_, 0.)
     self.assertEqual(cache, None)
 
+
 if __name__ == "__main__":
   tf.test.main()

From 28a6305850674055226e3080aab0a4d145e9e2ee Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Mon, 11 Feb 2019 09:57:25 -0800
Subject: [PATCH 1655/2720] Add an hparam
 use_global_position_in_packed_sequence in mtf_transformer2. If True
 (default), then we use the global position in the packed example as the input
 to the positional embedding.  If False, then we use the position in the
 individual sequence. It is counterintuitive why we want to make True the
 default, since False seems to make more sense. However, the previous
 submitted CL had the effect of changing from True to False, which caused some
 models to diverge.  This CL restores the previous working state. TODO(noam):
 investigate why the models diverge with False.

PiperOrigin-RevId: 233427027
---
 .../data_generators/wiki_multi_problems.py    | 57 +++++++++++++------
 tensor2tensor/models/mtf_transformer2.py      | 38 +++++++------
 tensor2tensor/utils/trainer_lib.py            |  3 +-
 3 files changed, 64 insertions(+), 34 deletions(-)

diff --git a/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index 49d248749..4c140f139 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -80,26 +80,30 @@ class LanguagemodelMultiWikiTranslatePacked1k(
   """Wiki-LM, Translation, MNLI, SQUAD mixed problem class."""
 
   def __init__(self, was_reversed=False, was_copy=False):
-    problems = [
-        # TODO(noam): uncommonet once data is generated
-        wiki_lm.LanguagemodelDeEnFrRoWiki64kFitbPacked1k(),
-        wiki_lm.LanguagemodelDeEnFrRoWiki64kFitbPacked1k(was_reversed=True),
-        translate_ende.TranslateEndeWmtMulti64kPacked1k(),
-        translate_ende.TranslateEndeWmtMulti64kPacked1k(was_reversed=True),
-        translate_enfr.TranslateEnfrWmtMulti64kPacked1k(),
-        translate_enfr.TranslateEnfrWmtMulti64kPacked1k(was_reversed=True),
-        translate_enro.TranslateEnroWmtMultiTiny64kPacked1k(),
-        translate_enro.TranslateEnroWmtMultiTiny64kPacked1k(was_reversed=True),
-        cnn_dailymail.SummarizeCnnDailymailMulti64kPacked1k(),
-        cnn_dailymail.SummarizeCnnDailymailMulti64kPacked1k(was_reversed=True),
-        multinli.MultiNLIText2textMulti64kPacked1k(),
-        squad.SquadText2textMulti64kPacked1k(),
-    ]
-    schedule = multi_problem_v2.constant_schedule(
-        multi_problem_v2.epoch_rates_to_pmf(problems))
+    problems = []
+    rates = []
+    for rate, also_reverse, cls in self.problems_and_rates:
+      for r in [False, True] if also_reverse else [False]:
+        problems.append(cls(was_reversed=r))
+        rates.append(rate)
+    pmf = multi_problem_v2.epoch_rates_to_pmf(problems, epoch_rates=rates)
+    schedule = multi_problem_v2.constant_schedule(pmf)
     super(LanguagemodelMultiWikiTranslatePacked1k, self).__init__(
         problems, schedule, was_reversed=was_reversed, was_copy=was_copy)
 
+  @property
+  def problems_and_rates(self):
+    """Returns a list of (weight, also_reverse, problem_class) triples."""
+    return [
+        (1.0, True, wiki_lm.LanguagemodelDeEnFrRoWiki64kFitbPacked1k),
+        (1.0, True, translate_ende.TranslateEndeWmtMulti64kPacked1k),
+        (1.0, True, translate_enfr.TranslateEnfrWmtMulti64kPacked1k),
+        (1.0, True, translate_enro.TranslateEnroWmtMultiTiny64kPacked1k),
+        (1.0, True, cnn_dailymail.SummarizeCnnDailymailMulti64kPacked1k),
+        (1.0, False, multinli.MultiNLIText2textMulti64kPacked1k),
+        (1.0, False, squad.SquadText2textMulti64kPacked1k),
+    ]
+
   @property
   def has_inputs(self):
     return True
@@ -117,6 +121,25 @@ def packed_length(self):
     return 1024
 
 
+@registry.register_problem
+class LanguagemodelMultiWikiTranslatePacked1kV2(
+    LanguagemodelMultiWikiTranslatePacked1k):
+  """Higher rates for rarer problems."""
+
+  @property
+  def problems_and_rates(self):
+    """Returns a list of (weight, also_reverse, problem_class) triples."""
+    return [
+        (1.0, True, wiki_lm.LanguagemodelDeEnFrRoWiki64kFitbPacked1k),
+        (3.0, True, translate_ende.TranslateEndeWmtMulti64kPacked1k),
+        (1.0, True, translate_enfr.TranslateEnfrWmtMulti64kPacked1k),
+        (100.0, True, translate_enro.TranslateEnroWmtMultiTiny64kPacked1k),
+        (1.0, True, cnn_dailymail.SummarizeCnnDailymailMulti64kPacked1k),
+        (10.0, False, multinli.MultiNLIText2textMulti64kPacked1k),
+        (10.0, False, squad.SquadText2textMulti64kPacked1k),
+    ]
+
+
 @registry.register_problem
 class LanguagemodelEnWikiLMMultiNLISubwords64k(multi_problem.MultiProblem):
   """Wiki LM and MNLI mixed problem class."""
diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index c279ee034..e91505cfd 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -136,13 +136,16 @@ def import_feature(key):
       return self._import_feature(features, mesh, key)
     targets = import_feature("targets")
     sequence_id = import_feature("targets_segmentation")
-    position = import_feature("targets_position")
+    if hparams.use_global_position_in_packed_sequence:
+      position = None
+    else:
+      position = import_feature("targets_position")
     if self.autoregressive:
       inputs = mtf.shift(
           targets, offset=1, dim=self.length_dim, wrap=False)
-      if position is not None:
-        # first input in later sequences should be 0
-        inputs *= mtf.to_int32(mtf.not_equal(position, 0))
+      # We should have a 0 at the beginning of each sequence rather than the
+      # shifted EOS (1) from the previous sequence.
+      inputs -= mtf.to_int32(mtf.equal(inputs, 1))
     else:
       inputs = import_feature("inputs")
       # TODO(noam): options for bert-style masking here?
@@ -248,8 +251,12 @@ def import_feature(key):
     decoder_sequence_id = import_feature("targets_segmentation")
     if decoder_sequence_id is None:
       decoder_sequence_id = mtf.to_int32(mtf.not_equal(targets, 0))
-    encoder_position = import_feature("inputs_position")
-    decoder_position = import_feature("targets_position")
+    if hparams.use_global_position_in_packed_sequence:
+      encoder_position = None
+      decoder_position = None
+    else:
+      encoder_position = import_feature("inputs_position")
+      decoder_position = import_feature("targets_position")
     model = self.model()
     logits, loss = model.call_simple(
         inputs=inputs,
@@ -349,7 +356,7 @@ def layer_stack_from_hparams(hparams, prefix):
   """Create a layer stack based on the hyperparameter values."""
   layers = hparams.get(prefix + "layers")
   return transformer.LayerStack(
-      [layers_registry.get(l)(hparams, prefix) for l in layers],
+      [layers_registry[l](hparams, prefix) for l in layers],
       dropout_rate=hparams.layer_prepostprocess_dropout,
       norm_epsilon=hparams.norm_epsilon)
 
@@ -418,6 +425,14 @@ def mtf_transformer2_base():
       "targets": modalities.ModalityType.IDENTITY_SYMBOL,
   }
   hparams.add_hparam("beam_size", 1)
+
+  # If this is True, then in a packed dataset (where exaples are concatenated
+  # to form longer examples) we use the global position (within the concatenated
+  # sequence) to compute the positional embedding, instead of the position
+  # within the individual sequence.  This is counterintuitive, but for some
+  # reason, it keeps the model from diverging.
+  hparams.add_hparam("use_global_position_in_packed_sequence", True)
+
   return hparams
 
 
@@ -837,12 +852,3 @@ def mtr_tr_ende_deep():
   hparams.encoder_num_layers = 12
   hparams.decoder_num_layers = 12
   return hparams
-
-
-@registry.register_hparams
-def ogm_dense_0():
-  hparams = mtr_tr_dense(0)
-  hparams.max_length = 1024
-  hparams.batch_size = 128
-  hparams.shared_embedding_and_softmax_weights = True
-  return hparams
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 89fbb6604..e8af60126 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -132,7 +132,8 @@ def create_session_config(log_device_placement=False,
       gpu_options=gpu_options,
       log_device_placement=log_device_placement,
       inter_op_parallelism_threads=inter_op_parallelism_threads,
-      intra_op_parallelism_threads=intra_op_parallelism_threads)
+      intra_op_parallelism_threads=intra_op_parallelism_threads,
+      isolate_session_state=True)
   return config
 
 

From c41d20208b6fb2290ec4818ac86ad2f5eb1ed544 Mon Sep 17 00:00:00 2001
From: Eric Purdy <epurdy@uchicago.edu>
Date: Mon, 11 Feb 2019 10:42:20 -0800
Subject: [PATCH 1656/2720] Fix for string features (#1440)

---
 tensor2tensor/utils/t2t_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 66b748ebb..6c8a29696 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -2017,7 +2017,7 @@ def summarize_features(features, num_shards=1):
 
   with tf.name_scope("input_stats"):
     for (k, v) in sorted(six.iteritems(features)):
-      if isinstance(v, tf.Tensor) and v.get_shape().ndims > 1:
+      if isinstance(v, tf.Tensor) and v.get_shape().ndims > 1 and v.dtype != tf.string:
         tf.summary.scalar("%s_batch" % k, tf.shape(v)[0] // num_shards)
         tf.summary.scalar("%s_length" % k, tf.shape(v)[1])
         nonpadding = tf.to_float(tf.not_equal(v, 0))

From a8962dffe4ee3732c1394b330af9c6d2c3dffb3c Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Mon, 11 Feb 2019 19:42:53 +0100
Subject: [PATCH 1657/2720] CI: Test tf-nighly on Python 3 and cleanup config
 to speedup test (#1445)

Closes #1431
---
 .travis.yml                    | 12 ++----------
 oss_scripts/oss_pip_install.sh |  8 ++++----
 2 files changed, 6 insertions(+), 14 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 367601195..fc22a2bd0 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -1,7 +1,8 @@
 sudo: required
 language: python
+cache: pip
 git:
-  depth: 10
+  depth: 1
   quiet: true
 services:
   - docker
@@ -21,15 +22,6 @@ env:
     # If updating, also update TF_LATEST above
     - TF_VERSION="1.12.*"
     - TF_VERSION="tf-nightly"
-matrix:
-  exclude:
-    # We test against all versions in Python 2 but only the latest stable
-    # version in Python 3
-    - python: "3.6"
-      env: TF_VERSION="tf-nightly"
-before_install:
-  - sudo apt-get update -qq
-  - sudo apt-get install -qq libhdf5-dev
 install:
   - ./oss_scripts/oss_pip_install.sh
 script:
diff --git a/oss_scripts/oss_pip_install.sh b/oss_scripts/oss_pip_install.sh
index 12ffeb2c1..333f49b41 100755
--- a/oss_scripts/oss_pip_install.sh
+++ b/oss_scripts/oss_pip_install.sh
@@ -5,6 +5,10 @@ set -e  # fail and exit on any command erroring
 
 : "${TF_VERSION:?}"
 
+# Make sure we have the latest version of numpy - avoid problems we were
+# seeing with Python 3
+pip install -q -U numpy
+
 if [[ "$TF_VERSION" == "tf-nightly"  ]]
 then
   pip install tf-nightly;
@@ -21,7 +25,3 @@ t2t-datagen 2>&1 | grep translate_ende 2>&1 >/dev/null && echo passed
 pip install -q -e .[tests,allen]
 # Make sure to install the atari extras for gym
 pip install "gym[atari]"
-
-# Make sure we have the latest version of numpy - avoid problems we were
-# seeing with Python 3
-pip install -q -U numpy

From 928c62f5819f9e4af96a6236d69ce46e7753cd37 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Mon, 11 Feb 2019 19:43:29 +0100
Subject: [PATCH 1658/2720] Slightly simplify python loops and attr acces
 (#1446)

Just some minor readability imrpovements that accumulated during debugging bugs in my custom model code.
---
 tensor2tensor/data_generators/cipher.py              |  6 +++---
 tensor2tensor/data_generators/text_encoder.py        |  4 ++--
 tensor2tensor/data_generators/wiki_revision_utils.py |  8 +++-----
 tensor2tensor/layers/common_layers.py                | 11 +++++------
 tensor2tensor/rl/trainer_model_based.py              |  2 +-
 tensor2tensor/serving/serving_utils.py               |  6 +++---
 6 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/data_generators/cipher.py b/tensor2tensor/data_generators/cipher.py
index ddafeddb1..fc8a59022 100644
--- a/tensor2tensor/data_generators/cipher.py
+++ b/tensor2tensor/data_generators/cipher.py
@@ -213,9 +213,9 @@ def encipher_vigenere(plaintext, plain_vocab, key):
   """
   ciphertext = []
   # generate Vigenere table
-  layers = []
-  for i in range(len(plain_vocab)):
-    layers.append(ShiftEncryptionLayer(plain_vocab, i))
+  layers = [
+      ShiftEncryptionLayer(plain_vocab, i) for i in range(len(plain_vocab))
+  ]
 
   for i, sentence in enumerate(plaintext):
     cipher_sentence = []
diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 4fe071093..54a30c4dd 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -393,8 +393,8 @@ def store_to_file(self, filename):
       filename: Full path of the file to store the vocab to.
     """
     with tf.gfile.Open(filename, "w") as f:
-      for i in range(len(self._id_to_token)):
-        f.write(self._id_to_token[i] + "\n")
+      for token in self._id_to_token:
+        f.write(token + "\n")
 
 
 def _escape_token(token, alphabet):
diff --git a/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index 4864f25c3..22f86b966 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -491,10 +491,8 @@ def edit_distance_filter(source_target_input, max_equal_to_diff_ratio=0):
   if not max_equal_to_diff_ratio:
     return source_target_input, thrown_out_count
 
-  for i in range(len(source_target_input)):
-    src = source_target_input[i][0]
-    tgt = source_target_input[i][1]
-    opcodes = fast_match_sequences(src, tgt)
+  for src_tgt in source_target_input:
+    opcodes = fast_match_sequences(*src_tgt)
     diff_char_count = 0
     equal_char_count = 0
     for tag, i1, i2, j1, j2 in opcodes:
@@ -504,7 +502,7 @@ def edit_distance_filter(source_target_input, max_equal_to_diff_ratio=0):
       else:
         equal_char_count += i2 - i1
     if diff_char_count <= max_equal_to_diff_ratio * equal_char_count:
-      source_target_output.append(source_target_input[i])
+      source_target_output.append(src_tgt)
     else:
       thrown_out_count += 1
   return source_target_output, thrown_out_count
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 158495733..9072b2820 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2334,10 +2334,10 @@ def ravanbakhsh_set_layer(layer_size,
 
 def fn_device_dependency_dict():
   """State container for fn_device_dependency."""
-  if not hasattr(tf.get_default_graph(), "dependency_dict"):
-    setattr(tf.get_default_graph(), "dependency_dict",
-            collections.defaultdict(list))
-  return tf.get_default_graph().dependency_dict
+  default_graph = tf.get_default_graph()
+  if not hasattr(default_graph, "dependency_dict"):
+    default_graph.dependency_dict = collections.defaultdict(list)
+  return default_graph.dependency_dict
 
 
 @contextlib.contextmanager
@@ -2791,8 +2791,7 @@ def shape_list(x):
   shape = tf.shape(x)
 
   ret = []
-  for i in range(len(static)):
-    dim = static[i]
+  for i, dim in enumerate(static):
     if dim is None:
       dim = shape[i]
     ret.append(dim)
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 31b90726a..f34303d92 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -126,7 +126,7 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
                      schedule="continuous_train_and_eval"):
   """Train supervised."""
   if local_eval_frequency is None:
-    local_eval_frequency = getattr(FLAGS, "local_eval_frequency")
+    local_eval_frequency = FLAGS.local_eval_frequency
 
   exp_fn = trainer_lib.create_experiment_fn(
       model_name, problem, data_dir, train_steps, eval_steps,
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 5d873aa8d..487a312b1 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -115,9 +115,9 @@ def _make_grpc_request(examples):
     scores = tf.make_ndarray(response.outputs["scores"])
     assert len(outputs) == len(scores)
     return [{
-        "outputs": outputs[i],
-        "scores": scores[i]
-    } for i in range(len(outputs))]
+        "outputs": output,
+        "scores": score
+    } for output, score in zip(outputs, scores)]
 
   return _make_grpc_request
 

From 8919e8248d965cc9a3172d5ac5dc64817626a541 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Mon, 11 Feb 2019 10:45:29 -0800
Subject: [PATCH 1659/2720] Merge of PR #1445

PiperOrigin-RevId: 233435966
---
 tensor2tensor/data_generators/cipher.py              |  6 +++---
 tensor2tensor/data_generators/text_encoder.py        |  4 ++--
 tensor2tensor/data_generators/wiki_revision_utils.py |  8 +++++---
 tensor2tensor/layers/common_layers.py                | 11 ++++++-----
 tensor2tensor/rl/trainer_model_based.py              |  2 +-
 tensor2tensor/serving/serving_utils.py               |  6 +++---
 tensor2tensor/utils/t2t_model.py                     |  2 +-
 7 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/data_generators/cipher.py b/tensor2tensor/data_generators/cipher.py
index fc8a59022..ddafeddb1 100644
--- a/tensor2tensor/data_generators/cipher.py
+++ b/tensor2tensor/data_generators/cipher.py
@@ -213,9 +213,9 @@ def encipher_vigenere(plaintext, plain_vocab, key):
   """
   ciphertext = []
   # generate Vigenere table
-  layers = [
-      ShiftEncryptionLayer(plain_vocab, i) for i in range(len(plain_vocab))
-  ]
+  layers = []
+  for i in range(len(plain_vocab)):
+    layers.append(ShiftEncryptionLayer(plain_vocab, i))
 
   for i, sentence in enumerate(plaintext):
     cipher_sentence = []
diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 54a30c4dd..4fe071093 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -393,8 +393,8 @@ def store_to_file(self, filename):
       filename: Full path of the file to store the vocab to.
     """
     with tf.gfile.Open(filename, "w") as f:
-      for token in self._id_to_token:
-        f.write(token + "\n")
+      for i in range(len(self._id_to_token)):
+        f.write(self._id_to_token[i] + "\n")
 
 
 def _escape_token(token, alphabet):
diff --git a/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index 22f86b966..4864f25c3 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -491,8 +491,10 @@ def edit_distance_filter(source_target_input, max_equal_to_diff_ratio=0):
   if not max_equal_to_diff_ratio:
     return source_target_input, thrown_out_count
 
-  for src_tgt in source_target_input:
-    opcodes = fast_match_sequences(*src_tgt)
+  for i in range(len(source_target_input)):
+    src = source_target_input[i][0]
+    tgt = source_target_input[i][1]
+    opcodes = fast_match_sequences(src, tgt)
     diff_char_count = 0
     equal_char_count = 0
     for tag, i1, i2, j1, j2 in opcodes:
@@ -502,7 +504,7 @@ def edit_distance_filter(source_target_input, max_equal_to_diff_ratio=0):
       else:
         equal_char_count += i2 - i1
     if diff_char_count <= max_equal_to_diff_ratio * equal_char_count:
-      source_target_output.append(src_tgt)
+      source_target_output.append(source_target_input[i])
     else:
       thrown_out_count += 1
   return source_target_output, thrown_out_count
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 9072b2820..158495733 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2334,10 +2334,10 @@ def ravanbakhsh_set_layer(layer_size,
 
 def fn_device_dependency_dict():
   """State container for fn_device_dependency."""
-  default_graph = tf.get_default_graph()
-  if not hasattr(default_graph, "dependency_dict"):
-    default_graph.dependency_dict = collections.defaultdict(list)
-  return default_graph.dependency_dict
+  if not hasattr(tf.get_default_graph(), "dependency_dict"):
+    setattr(tf.get_default_graph(), "dependency_dict",
+            collections.defaultdict(list))
+  return tf.get_default_graph().dependency_dict
 
 
 @contextlib.contextmanager
@@ -2791,7 +2791,8 @@ def shape_list(x):
   shape = tf.shape(x)
 
   ret = []
-  for i, dim in enumerate(static):
+  for i in range(len(static)):
+    dim = static[i]
     if dim is None:
       dim = shape[i]
     ret.append(dim)
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index f34303d92..31b90726a 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -126,7 +126,7 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
                      schedule="continuous_train_and_eval"):
   """Train supervised."""
   if local_eval_frequency is None:
-    local_eval_frequency = FLAGS.local_eval_frequency
+    local_eval_frequency = getattr(FLAGS, "local_eval_frequency")
 
   exp_fn = trainer_lib.create_experiment_fn(
       model_name, problem, data_dir, train_steps, eval_steps,
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 487a312b1..5d873aa8d 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -115,9 +115,9 @@ def _make_grpc_request(examples):
     scores = tf.make_ndarray(response.outputs["scores"])
     assert len(outputs) == len(scores)
     return [{
-        "outputs": output,
-        "scores": score
-    } for output, score in zip(outputs, scores)]
+        "outputs": outputs[i],
+        "scores": scores[i]
+    } for i in range(len(outputs))]
 
   return _make_grpc_request
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 6c8a29696..66b748ebb 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -2017,7 +2017,7 @@ def summarize_features(features, num_shards=1):
 
   with tf.name_scope("input_stats"):
     for (k, v) in sorted(six.iteritems(features)):
-      if isinstance(v, tf.Tensor) and v.get_shape().ndims > 1 and v.dtype != tf.string:
+      if isinstance(v, tf.Tensor) and v.get_shape().ndims > 1:
         tf.summary.scalar("%s_batch" % k, tf.shape(v)[0] // num_shards)
         tf.summary.scalar("%s_length" % k, tf.shape(v)[1])
         nonpadding = tf.to_float(tf.not_equal(v, 0))

From f1076df84f2453821113ee59a56f45f1860147da Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Mon, 11 Feb 2019 10:55:59 -0800
Subject: [PATCH 1660/2720] Merge of PR #1446

PiperOrigin-RevId: 233438134
---
 tensor2tensor/data_generators/cipher.py              |  6 +++---
 tensor2tensor/data_generators/wiki_revision_utils.py |  8 +++-----
 tensor2tensor/layers/common_layers.py                | 11 +++++------
 tensor2tensor/rl/trainer_model_based.py              |  2 +-
 tensor2tensor/serving/serving_utils.py               |  6 +++---
 5 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/data_generators/cipher.py b/tensor2tensor/data_generators/cipher.py
index ddafeddb1..fc8a59022 100644
--- a/tensor2tensor/data_generators/cipher.py
+++ b/tensor2tensor/data_generators/cipher.py
@@ -213,9 +213,9 @@ def encipher_vigenere(plaintext, plain_vocab, key):
   """
   ciphertext = []
   # generate Vigenere table
-  layers = []
-  for i in range(len(plain_vocab)):
-    layers.append(ShiftEncryptionLayer(plain_vocab, i))
+  layers = [
+      ShiftEncryptionLayer(plain_vocab, i) for i in range(len(plain_vocab))
+  ]
 
   for i, sentence in enumerate(plaintext):
     cipher_sentence = []
diff --git a/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index 4864f25c3..22f86b966 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -491,10 +491,8 @@ def edit_distance_filter(source_target_input, max_equal_to_diff_ratio=0):
   if not max_equal_to_diff_ratio:
     return source_target_input, thrown_out_count
 
-  for i in range(len(source_target_input)):
-    src = source_target_input[i][0]
-    tgt = source_target_input[i][1]
-    opcodes = fast_match_sequences(src, tgt)
+  for src_tgt in source_target_input:
+    opcodes = fast_match_sequences(*src_tgt)
     diff_char_count = 0
     equal_char_count = 0
     for tag, i1, i2, j1, j2 in opcodes:
@@ -504,7 +502,7 @@ def edit_distance_filter(source_target_input, max_equal_to_diff_ratio=0):
       else:
         equal_char_count += i2 - i1
     if diff_char_count <= max_equal_to_diff_ratio * equal_char_count:
-      source_target_output.append(source_target_input[i])
+      source_target_output.append(src_tgt)
     else:
       thrown_out_count += 1
   return source_target_output, thrown_out_count
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 158495733..9072b2820 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2334,10 +2334,10 @@ def ravanbakhsh_set_layer(layer_size,
 
 def fn_device_dependency_dict():
   """State container for fn_device_dependency."""
-  if not hasattr(tf.get_default_graph(), "dependency_dict"):
-    setattr(tf.get_default_graph(), "dependency_dict",
-            collections.defaultdict(list))
-  return tf.get_default_graph().dependency_dict
+  default_graph = tf.get_default_graph()
+  if not hasattr(default_graph, "dependency_dict"):
+    default_graph.dependency_dict = collections.defaultdict(list)
+  return default_graph.dependency_dict
 
 
 @contextlib.contextmanager
@@ -2791,8 +2791,7 @@ def shape_list(x):
   shape = tf.shape(x)
 
   ret = []
-  for i in range(len(static)):
-    dim = static[i]
+  for i, dim in enumerate(static):
     if dim is None:
       dim = shape[i]
     ret.append(dim)
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 31b90726a..f34303d92 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -126,7 +126,7 @@ def train_supervised(problem, model_name, hparams, data_dir, output_dir,
                      schedule="continuous_train_and_eval"):
   """Train supervised."""
   if local_eval_frequency is None:
-    local_eval_frequency = getattr(FLAGS, "local_eval_frequency")
+    local_eval_frequency = FLAGS.local_eval_frequency
 
   exp_fn = trainer_lib.create_experiment_fn(
       model_name, problem, data_dir, train_steps, eval_steps,
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 5d873aa8d..487a312b1 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -115,9 +115,9 @@ def _make_grpc_request(examples):
     scores = tf.make_ndarray(response.outputs["scores"])
     assert len(outputs) == len(scores)
     return [{
-        "outputs": outputs[i],
-        "scores": scores[i]
-    } for i in range(len(outputs))]
+        "outputs": output,
+        "scores": score
+    } for output, score in zip(outputs, scores)]
 
   return _make_grpc_request
 

From 5e32d3137cc1e9ead5cc44801e295a70ebfc2c1f Mon Sep 17 00:00:00 2001
From: Eric Purdy <epurdy@uchicago.edu>
Date: Mon, 11 Feb 2019 10:56:25 -0800
Subject: [PATCH 1661/2720] Merge of PR #1440

PiperOrigin-RevId: 233438227
---
 tensor2tensor/utils/t2t_model.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 66b748ebb..3014a4e69 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -2017,7 +2017,8 @@ def summarize_features(features, num_shards=1):
 
   with tf.name_scope("input_stats"):
     for (k, v) in sorted(six.iteritems(features)):
-      if isinstance(v, tf.Tensor) and v.get_shape().ndims > 1:
+      if (isinstance(v, tf.Tensor) and (v.get_shape().ndims > 1) and
+          (v.dtype != tf.string)):
         tf.summary.scalar("%s_batch" % k, tf.shape(v)[0] // num_shards)
         tf.summary.scalar("%s_length" % k, tf.shape(v)[1])
         nonpadding = tf.to_float(tf.not_equal(v, 0))

From 474463e23d4df33f1c87062ffb28b7a8db49c5e1 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 11 Feb 2019 11:33:03 -0800
Subject: [PATCH 1662/2720] Remove Modality.top_dimensionality.

This is an incremental change toward replacing Modality classes with a simpler ModalityType enum.

PiperOrigin-RevId: 233445977
---
 tensor2tensor/data_generators/problem.py      |   2 +-
 tensor2tensor/layers/modalities.py            |  28 ++---
 tensor2tensor/models/mtf_image_transformer.py |   6 +-
 tensor2tensor/models/mtf_transformer.py       |   6 +-
 tensor2tensor/models/mtf_transformer2.py      |   6 +-
 tensor2tensor/models/research/autoencoders.py |  12 +-
 tensor2tensor/models/resnet.py                |   4 +-
 tensor2tensor/models/transformer.py           |  12 +-
 tensor2tensor/utils/modality.py               |   5 -
 tensor2tensor/utils/t2t_model.py              | 107 ++++++++++--------
 tensor2tensor/v2/t2t.py                       |   8 +-
 11 files changed, 108 insertions(+), 88 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 7c1327504..70ecbcc48 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -749,7 +749,7 @@ def feature_info(self):
     for feature_name, modality_cls in six.iteritems(hp.modality):
       finfo = features[feature_name]
       finfo.modality = modality_cls
-      finfo.vocab_size = modality_cls.top_dimensionality
+      finfo.vocab_size = hp.vocab_size[feature_name]
 
     vocabs = hp.vocabulary
     for name, encoder in six.iteritems(vocabs):
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 5a0411679..4b725f569 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -246,12 +246,12 @@ def targets_bottom(self, x):
       if len(inputs_shape) != 4:
         raise ValueError("Assuming images given as int tensors in the format "
                          "[batch, height, width, channels] (256 values).")
-      # We embed each of 256=self.top_dimensionality possible pixel values.
+      # We embed each of 256=self._vocab_size possible pixel values.
       embedding_var = tf.get_variable(
           "pixel_embedding",
-          [self.top_dimensionality, self.PIXEL_EMBEDDING_SIZE])
-      hot_inputs = tf.one_hot(tf.to_int32(inputs), self.top_dimensionality)
-      hot_inputs = tf.reshape(hot_inputs, [-1, self.top_dimensionality])
+          [self._vocab_size, self.PIXEL_EMBEDDING_SIZE])
+      hot_inputs = tf.one_hot(tf.to_int32(inputs), self._vocab_size)
+      hot_inputs = tf.reshape(hot_inputs, [-1, self._vocab_size])
       embedded = tf.matmul(hot_inputs, embedding_var)
       # Let's now merge all channels that were embedded into a single vector.
       merged_size = self.PIXEL_EMBEDDING_SIZE * inputs_shape[3]
@@ -268,8 +268,8 @@ def top(self, body_output, _):
     with tf.variable_scope("rgb_softmax"):
       body_output_shape = common_layers.shape_list(body_output)
       reshape_shape = body_output_shape[:3]
-      reshape_shape.extend([num_channels, self.top_dimensionality])
-      res = tf.layers.dense(body_output, self.top_dimensionality * num_channels)
+      reshape_shape.extend([num_channels, self._vocab_size])
+      res = tf.layers.dense(body_output, self._vocab_size * num_channels)
       res = tf.reshape(res, reshape_shape)
       if not tf.get_variable_scope().reuse:
         res_argmax = tf.argmax(res, axis=-1)
@@ -353,7 +353,7 @@ def top(self, body_output, _):
       body_output: Tensor of shape [batch, img_len, img_len, depth].
 
     Returns:
-      Tensor of shape [batch, img_len, img_len, channels, top_dimensionality].
+      Tensor of shape [batch, img_len, img_len, channels, vocab_size].
     """
     with tf.variable_scope(self.name):
       hidden_size = self._model_hparams.hidden_size
@@ -371,12 +371,12 @@ def top(self, body_output, _):
       x = tf.reshape(x, [batch, img_len, img_len * channels, hidden_size])
       x = common_layers.layer_preprocess(x, self._model_hparams)
       x = tf.layers.dense(x,
-                          self.top_dimensionality,
+                          self._vocab_size,
                           use_bias=True,
                           activation=None,
                           name="output_conv")
       x = tf.reshape(
-          x, [batch, img_len, img_len, channels, self.top_dimensionality])
+          x, [batch, img_len, img_len, channels, self._vocab_size])
       return x
 
 
@@ -427,7 +427,7 @@ def top(self, body_output, _):
       x = tf.layers.dense(
           body_output, 256, use_bias=True, activation=None, name="output_conv")
       x = tf.reshape(x,
-                     [-1, img_len, img_len, channels, self.top_dimensionality])
+                     [-1, img_len, img_len, channels, self._vocab_size])
       return x
 
 
@@ -626,10 +626,10 @@ def targets_bottom(self, x):
   def top(self, body_output, targets):
     num_channels = self._model_hparams.problem.num_channels
     shape = common_layers.shape_list(body_output)
-    reshape_shape = shape[:-1] + [num_channels, self.top_dimensionality]
+    reshape_shape = shape[:-1] + [num_channels, self._vocab_size]
     res = tf.reshape(body_output, reshape_shape)
     # Calculate argmax so as to have a summary with the produced images.
-    x = tf.argmax(tf.reshape(res, [-1, self.top_dimensionality]), axis=-1)
+    x = tf.argmax(tf.reshape(res, [-1, self._vocab_size]), axis=-1)
     x = tf.reshape(x, shape[:-1] + [num_channels])
     common_video.gif_summary("results", x, max_outputs=1)
     return res
@@ -657,7 +657,7 @@ def bottom(self, x):
     with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
       common_layers.summarize_video(inputs, "bottom")
       # Embed bitwise.
-      assert self.top_dimensionality == 256
+      assert self._vocab_size == 256
       embedded = discretization.int_to_bit_embed(inputs, 8,
                                                  self.PIXEL_EMBEDDING_SIZE)
       # Project.
@@ -671,7 +671,7 @@ def targets_bottom(self, x):  # pylint: disable=arguments-differ
     with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
       common_layers.summarize_video(inputs, "targets_bottom")
       # Embed bitwise.
-      assert self.top_dimensionality == 256
+      assert self._vocab_size == 256
       embedded = discretization.int_to_bit_embed(inputs, 8,
                                                  self.PIXEL_EMBEDDING_SIZE)
       # Transpose and project.
diff --git a/tensor2tensor/models/mtf_image_transformer.py b/tensor2tensor/models/mtf_image_transformer.py
index e169e221c..67915f656 100644
--- a/tensor2tensor/models/mtf_image_transformer.py
+++ b/tensor2tensor/models/mtf_image_transformer.py
@@ -46,8 +46,10 @@ def inputs_vocab_dim(self):
 
   @property
   def targets_vocab_dim(self):
-    return mtf.Dimension(
-        "vocab", self._problem_hparams.modality["targets"].top_dimensionality)
+    vocab_size = self._problem_hparams.vocab_size["targets"]
+    if hasattr(self._hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
+    return mtf.Dimension("vocab", vocab_size)
 
   @property
   def outputs_vocab_dim(self):
diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 167ab7394..3e32e9c03 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -403,15 +403,13 @@ def mtf_model_fn(self, features, mesh):
 
   @property
   def _targets_vocab_size(self):
-    targets_vocab_size = self._problem_hparams.modality[
-        "targets"].top_dimensionality
+    targets_vocab_size = self._problem_hparams.vocab_size["targets"]
     targets_vocab_size += (-targets_vocab_size) % self._hparams.vocab_divisor
     return targets_vocab_size
 
   @property
   def _inputs_vocab_size(self):
-    inputs_vocab_size = self._problem_hparams.modality[
-        "inputs"].top_dimensionality
+    inputs_vocab_size = self._problem_hparams.vocab_size["inputs"]
     inputs_vocab_size += (-inputs_vocab_size) % self._hparams.vocab_divisor
     return inputs_vocab_size
 
diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index e91505cfd..1196f7075 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -168,15 +168,13 @@ def mtf_model_fn(self, features, mesh):
 
   @property
   def _targets_vocab_size(self):
-    targets_vocab_size = self._problem_hparams.modality[
-        "targets"].top_dimensionality
+    targets_vocab_size = self._problem_hparams.vocab_size["targets"]
     targets_vocab_size += (-targets_vocab_size) % self._hparams.vocab_divisor
     return targets_vocab_size
 
   @property
   def _inputs_vocab_size(self):
-    inputs_vocab_size = self._problem_hparams.modality[
-        "inputs"].top_dimensionality
+    inputs_vocab_size = self._problem_hparams.vocab_size["inputs"]
     inputs_vocab_size += (-inputs_vocab_size) % self._hparams.vocab_divisor
     return inputs_vocab_size
 
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 6e5aecdb4..7486fb61d 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -159,7 +159,9 @@ def decoder(self, x, encoder_layers):
   def gumbel_sample(self, reconstr_gan):
     hparams = self.hparams
     is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
-    vocab_size = self._problem_hparams.modality["targets"].top_dimensionality
+    vocab_size = self._problem_hparams.vocab_size["targets"]
+    if hasattr(self._hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
     reconstr_gan = tf.nn.log_softmax(reconstr_gan)
     if is_training and hparams.gumbel_temperature > 0.0:
       gumbel_samples = discretization.gumbel_sample(
@@ -180,7 +182,9 @@ def gumbel_sample(self, reconstr_gan):
   def body(self, features):
     hparams = self.hparams
     is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
-    vocab_size = self._problem_hparams.modality["targets"].top_dimensionality
+    vocab_size = self._problem_hparams.vocab_size["targets"]
+    if hasattr(self._hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
     encoder_layers = None
     self.is1d = hparams.sample_width == 1
     if (hparams.mode != tf.estimator.ModeKeys.PREDICT
@@ -460,7 +464,9 @@ def body(self, features):
       plain_training_loss = losses.pop("training")
       losses["plain"] = plain_training_loss
     res_shape = common_layers.shape_list(basic_result)
-    vocab_size = self._problem_hparams.modality["targets"].top_dimensionality
+    vocab_size = self._problem_hparams.vocab_size["targets"]
+    if hasattr(self._hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
     targets = tf.one_hot(features["targets_raw"], vocab_size)
     # Prepare inputs for autoregressive modes.
     if common_layers.shape_list(features["targets"])[1] == 1:
diff --git a/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index eba093135..e1d358048 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -569,7 +569,9 @@ def body(self, features):
       return out
 
     out = tf.reduce_mean(out, [1, 2])
-    num_classes = self._problem_hparams.modality["targets"].top_dimensionality
+    num_classes = self._problem_hparams.vocab_size["targets"]
+    if hasattr(self._hparams, "vocab_divisor"):
+      num_classes += (-num_classes) % self._hparams.vocab_divisor
     logits = tf.layers.dense(out, num_classes, name="logits")
 
     losses = {"training": 0.0}
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index d7644db93..43b879cd2 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -497,13 +497,17 @@ def forced_logits():
             tf.less(i, partial_targets_length), forced_logits, lambda: ret)
       return ret, cache
 
+    vocab_size = self._problem_hparams.vocab_size["targets"]
+    if hasattr(self._hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
+
     ret = fast_decode_tpu(
         encoder_output=encoder_output,
         encoder_decoder_attention_bias=encoder_decoder_attention_bias,
         symbols_to_logits_fn=symbols_to_logits_tpu_fn,
         hparams=hparams,
         decode_length=decode_length,
-        vocab_size=target_modality.top_dimensionality,
+        vocab_size=vocab_size,
         beam_size=beam_size,
         top_beams=top_beams,
         alpha=alpha,
@@ -693,13 +697,17 @@ def forced_logits():
             tf.less(i, partial_targets_length), forced_logits, lambda: ret)
       return ret, cache
 
+    vocab_size = self._problem_hparams.vocab_size["targets"]
+    if hasattr(self._hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
+
     ret = fast_decode(
         encoder_output=encoder_output,
         encoder_decoder_attention_bias=encoder_decoder_attention_bias,
         symbols_to_logits_fn=symbols_to_logits_fn,
         hparams=hparams,
         decode_length=decode_length,
-        vocab_size=target_modality.top_dimensionality,
+        vocab_size=vocab_size,
         beam_size=beam_size,
         top_beams=top_beams,
         alpha=alpha,
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index f3ab31e18..399eed4a5 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -58,11 +58,6 @@ def __init__(self, model_hparams, vocab_size=None):
   def name(self):
     return misc_utils.camelcase_to_snakecase(type(self).__name__)
 
-  @property
-  def top_dimensionality(self):
-    """Integer, the last dimension of the predictions (vocab size)."""
-    return self._vocab_size
-
   @property
   def top_is_pointwise(self):
     """Whether the top mapping of the modality is pointwise.
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 3014a4e69..957829d9e 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -38,7 +38,6 @@
 from tensor2tensor.utils import learning_rate
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import mlperf_log
-from tensor2tensor.utils import modality
 from tensor2tensor.utils import optimize
 from tensor2tensor.utils import quantization
 from tensor2tensor.utils import registry
@@ -181,28 +180,27 @@ def __init__(self,
     hparams = hparams_lib.copy_hparams(hparams)
     if self._problem_hparams and hparams.shared_embedding_and_softmax_weights:
       # If vocabularies differ, unset shared_embedding_and_softmax_weights.
-      input_modality = self._problem_hparams.modality.get("inputs")
-      target_modality = self._problem_hparams.modality.get("targets")
-      if (isinstance(input_modality, modality.Modality) and
-          isinstance(target_modality, modality.Modality) and
-          input_modality.top_dimensionality !=
-          target_modality.top_dimensionality):
+      input_vocab_size = self._problem_hparams.vocab_size["inputs"]
+      target_vocab_size = self._problem_hparams.vocab_size["targets"]
+      if input_vocab_size is not None and hasattr(hparams, "vocab_divisor"):
+        input_vocab_size += (-input_vocab_size) % hparams.vocab_divisor
+      if target_vocab_size is not None and hasattr(hparams, "vocab_divisor"):
+        target_vocab_size += (-target_vocab_size) % hparams.vocab_divisor
+      if input_vocab_size != target_vocab_size:
         log_info("Unsetting shared_embedding_and_softmax_weights.")
         hparams.shared_embedding_and_softmax_weights = 0
 
-      if isinstance(target_modality, modality.Modality):
-        if hparams.hidden_size:
-          hidden_size = hparams.hidden_size
-        else:
-          hidden_size = 1024
-
-        mlperf_log.transformer_print(
-            key=mlperf_log.MODEL_HP_EMBEDDING_SHARED_WEIGHTS,
-            value={
-                "vocab_size": target_modality.top_dimensionality,
-                "hidden_size": hidden_size
-            },
-            hparams=hparams)
+      if hparams.hidden_size:
+        hidden_size = hparams.hidden_size
+      else:
+        hidden_size = 1024
+      mlperf_log.transformer_print(
+          key=mlperf_log.MODEL_HP_EMBEDDING_SHARED_WEIGHTS,
+          value={
+              "vocab_size": target_vocab_size,
+              "hidden_size": hidden_size
+          },
+          hparams=hparams)
 
     self._original_hparams = hparams
     self.set_mode(mode)
@@ -428,7 +426,7 @@ def bottom(self, features):
     target_modality = _create_target_modality(self._problem_hparams.modality)
 
     # Transform features via its corresponding modality.
-    for feature_name, modality_obj in sorted(
+    for feature_name, modality in sorted(
         six.iteritems(self._problem_hparams.modality)):
       if feature_name not in features:
         tf.logging.warning("Missing feature %s - ignoring." % feature_name)
@@ -438,27 +436,27 @@ def bottom(self, features):
       # target modality; and to reuse variable scopes for only input modalities.
       if feature_name in target_modality:
         if len(target_modality) > 1:
-          variable_scope_name = "%s/%s" % (modality_obj.name, feature_name)
+          variable_scope_name = "%s/%s" % (modality.name, feature_name)
         else:
-          variable_scope_name = modality_obj.name
+          variable_scope_name = modality.name
         # TODO(aidangomez): share variables?
         with tf.variable_scope(variable_scope_name) as vs:
           self._add_variable_scope(variable_scope_name, vs)
           log_info("Transforming feature '%s' with %s.targets_bottom",
                    feature_name,
-                   modality_obj.name)
-          transformed_features[feature_name] = modality_obj.targets_bottom(
+                   modality.name)
+          transformed_features[feature_name] = modality.targets_bottom(
               features[feature_name])
       else:
-        do_reuse = modality_obj.name in all_previous_modalities
-        with tf.variable_scope(modality_obj.name, reuse=do_reuse) as vs:
-          self._add_variable_scope(modality_obj.name, vs)
+        do_reuse = modality.name in all_previous_modalities
+        with tf.variable_scope(modality.name, reuse=do_reuse) as vs:
+          self._add_variable_scope(modality.name, vs)
           log_info("Transforming feature '%s' with %s.bottom",
                    feature_name,
-                   modality_obj.name)
-          transformed_features[feature_name] = modality_obj.bottom(
+                   modality.name)
+          transformed_features[feature_name] = modality.bottom(
               features[feature_name])
-        all_previous_modalities.append(modality_obj.name)
+        all_previous_modalities.append(modality.name)
 
     for key in features:
       if key not in transformed_features:
@@ -708,9 +706,9 @@ def set_mode(self, mode):
 
     if self._problem_hparams:
       # Set model hparams in problem_hparams' modalities, which also store them.
-      for modality_obj in six.itervalues(self._problem_hparams.modality):
-        if modality_obj is not None:
-          modality_obj._model_hparams = self._hparams  # pylint: disable=protected-access
+      for modality in six.itervalues(self._problem_hparams.modality):
+        if modality is not None:
+          modality._model_hparams = self._hparams  # pylint: disable=protected-access
 
   def prepare_features_for_infer(self, features):
     """Called before inference to allow adding infer-specific features."""
@@ -887,8 +885,10 @@ def symbols_to_logits_fn(ids, i=None):
       features["inputs"] = tf.reshape(features["inputs"],
                                       [s[0] * s[1], s[2], s[3], s[4]])
 
-    target_modality = self._problem_hparams.modality["targets"]
-    vocab_size = target_modality.top_dimensionality
+    vocab_size = self._problem_hparams.vocab_size["targets"]
+    if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
+
     # Setting decode length to input length + decode_length
     if "partial_targets" not in features:
       inputs = features["inputs"]
@@ -1042,8 +1042,10 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
          tf.zeros([batch_size, decode_length, 1, 1], tf.int64)],
         axis=1)
     # tensor padded to [batch_size, decode_length, 1, 1, vocab_size]
-    logits = tf.zeros((batch_size, decode_length, 1, 1,
-                       target_modality.top_dimensionality))
+    vocab_size = self._problem_hparams.vocab_size["targets"]
+    if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
+    logits = tf.zeros((batch_size, decode_length, 1, 1, vocab_size))
     if not tf.executing_eagerly():
       logits.set_shape([None, None, None, None, None])
     loss = 0.0
@@ -1076,16 +1078,17 @@ def fn_not_eos():
             lambda: not_overflow)
       return not_overflow
 
+    vocab_size = self._problem_hparams.vocab_size["targets"]
+    if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
+
     _, result, logits, loss = tf.while_loop(
         while_exit_cond,
         infer_step, [tf.constant(0), result, logits, loss],
         shape_invariants=[
             tf.TensorShape([]),
             tf.TensorShape([batch_size, decode_length, 1, 1]),
-            tf.TensorShape([
-                batch_size, decode_length, 1, 1,
-                target_modality.top_dimensionality
-            ]),
+            tf.TensorShape([batch_size, decode_length, 1, 1, vocab_size]),
             tf.TensorShape([]),
         ],
         back_prop=False,
@@ -1152,7 +1155,9 @@ def infer_step(recent_output, recent_logits, unused_loss):
       """Inference step."""
       if not tf.executing_eagerly():
         if self._target_modality_is_real:
-          dim = self._problem_hparams.modality["targets"].top_dimensionality
+          dim = self._problem_hparams.vocab_size["targets"]
+          if dim is not None and hasattr(self._hparams, "vocab_divisor"):
+            dim += (-dim) % self._hparams.vocab_divisor
           recent_output.set_shape([None, None, None, dim])
         else:
           recent_output.set_shape([None, None, None, 1])
@@ -1192,7 +1197,9 @@ def infer_step(recent_output, recent_logits, unused_loss):
     else:
       batch_size = common_layers.shape_list(features["inputs"])[0]
       if self._target_modality_is_real:
-        dim = self._problem_hparams.modality["targets"].top_dimensionality
+        dim = self._problem_hparams.vocab_size["targets"]
+        if dim is not None and hasattr(self._hparams, "vocab_divisor"):
+          dim += (-dim) % self._hparams.vocab_divisor
         initial_output = tf.zeros((batch_size, 0, 1, dim), dtype=tf.float32)
       else:
         initial_output = tf.zeros((batch_size, 0, 1, 1), dtype=tf.int64)
@@ -1212,13 +1219,15 @@ def infer_step(recent_output, recent_logits, unused_loss):
 
     # Initial values of result, logits and loss.
     result = initial_output
+    vocab_size = self._problem_hparams.vocab_size["targets"]
+    if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
     if self._target_modality_is_real:
-      logits = tf.zeros((batch_size, 0, 1, target_modality.top_dimensionality))
+      logits = tf.zeros((batch_size, 0, 1, vocab_size))
       logits_shape_inv = [None, None, None, None]
     else:
       # tensor of shape [batch_size, time, 1, 1, vocab_size]
-      logits = tf.zeros((batch_size, 0, 1, 1,
-                         target_modality.top_dimensionality))
+      logits = tf.zeros((batch_size, 0, 1, 1, vocab_size))
       logits_shape_inv = [None, None, None, None, None]
     if not tf.executing_eagerly():
       logits.set_shape(logits_shape_inv)
@@ -1935,7 +1944,9 @@ def scheduled_sampling(hparams, problem_hparams, dp, sharded_logits, losses,
 
   def sample(x):
     """Multinomial sampling from a n-dimensional tensor."""
-    vocab_size = target_modality.top_dimensionality
+    vocab_size = problem_hparams.vocab_size["targets"]
+    if vocab_size is not None and hasattr(hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % hparams.vocab_divisor
     samples = tf.multinomial(tf.reshape(x, [-1, vocab_size]), 1)
     reshaped_samples = tf.reshape(samples, common_layers.shape_list(x)[:-1])
     return tf.to_int32(reshaped_samples)
diff --git a/tensor2tensor/v2/t2t.py b/tensor2tensor/v2/t2t.py
index 739170384..a4bf77214 100644
--- a/tensor2tensor/v2/t2t.py
+++ b/tensor2tensor/v2/t2t.py
@@ -105,10 +105,10 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
   for example in train_dataset.take(3):
     input_shapes.append(example["inputs"].shape.as_list())
     target_shapes.append(example["targets"].shape.as_list())
-  input_info = _make_info(
-      input_shapes, hparams.modality["inputs"].top_dimensionality)
-  target_info = _make_info(
-      target_shapes, hparams.modality["targets"].top_dimensionality)
+  input_vocab_size = hparams.vocab_size["inputs"]
+  target_vocab_size = hparams.vocab_size["targets"]
+  input_info = _make_info(input_shapes, input_vocab_size)
+  target_info = _make_info(target_shapes, target_vocab_size)
   info = {"inputs": input_info, "targets": target_info}
   return train_dataset, eval_dataset, info, supervised_keys
 

From 7a0584cbe372906ddc55cf83e934a3502e6ac6c9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 11 Feb 2019 14:43:28 -0800
Subject: [PATCH 1663/2720] trainer_model_free.py changes needed for using the
 new EnvProblem.

PiperOrigin-RevId: 233482446
---
 tensor2tensor/rl/trainer_model_free.py        | 38 +++++++++++++------
 .../rl/trainer_model_free_tictactoe_test.py   |  3 +-
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 4a6ce15ba..40fdb9c2a 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -21,6 +21,16 @@
     --output_dir=$HOME/t2t/rl_v1 \
     --hparams_set=pong_model_free \
     --hparams='batch_size=15'
+
+Example invocation with EnvProblem interface:
+
+python -m tensor2tensor.rl.trainer_model_free \
+  --env_problem_name=tic_tac_toe_env_problem \
+  --hparams_set=rlmf_tictactoe \
+  --output_dir=${OUTPUTDIR} \
+  --log_dir=${LOGDIR} \
+  --alsologtostderr \
+  --vmodule=*/tensor2tensor/*=2 \
 """
 
 from __future__ import absolute_import
@@ -34,6 +44,7 @@
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
 from tensor2tensor.utils import misc_utils
+from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -43,6 +54,9 @@
 FLAGS = flags.FLAGS
 
 
+flags.DEFINE_string("env_problem_name", "",
+                    "Which registered env_problem do we want?")
+
 # To maintain compatibility with some internal libs, we guard against these flag
 # definitions possibly erring. Apologies for the ugliness.
 try:
@@ -51,14 +65,16 @@
   pass
 
 
-def initialize_env_specs(hparams):
-  """Initializes env_specs using T2TGymEnvs."""
-  env = rl_utils.setup_env(hparams, hparams.batch_size,
-                           hparams.eval_max_num_noops,
-                           hparams.rl_env_max_episode_steps,
-                           env_name=hparams.rl_env_name)
-
-  env.start_new_epoch(0)
+def initialize_env_specs(hparams, env_problem_name):
+  """Initializes env_specs using the appropriate env."""
+  if env_problem_name:
+    env = registry.env_problem(env_problem_name, hparams.batch_size)
+  else:
+    env = rl_utils.setup_env(hparams, hparams.batch_size,
+                             hparams.eval_max_num_noops,
+                             hparams.rl_env_max_episode_steps,
+                             env_name=hparams.rl_env_name)
+    env.start_new_epoch(0)
 
   return rl.make_real_env_fn(env)
 
@@ -66,9 +82,9 @@ def initialize_env_specs(hparams):
 step = 0
 
 
-def train(hparams, output_dir, report_fn=None):
+def train(hparams, output_dir, env_problem_name, report_fn=None):
   """Train."""
-  env_fn = initialize_env_specs(hparams)
+  env_fn = initialize_env_specs(hparams, env_problem_name)
 
   tf.logging.vlog(1, "HParams in trainer_model_free.train : %s",
                   misc_utils.pprint_hparams(hparams))
@@ -134,7 +150,7 @@ def main(_):
   hparams = trainer_lib.create_hparams(FLAGS.hparams_set, FLAGS.hparams)
 
   tf.logging.info("Starting model free training.")
-  train(hparams, FLAGS.output_dir)
+  train(hparams, FLAGS.output_dir, FLAGS.env_problem_name)
   tf.logging.info("Ended model free training.")
 
 
diff --git a/tensor2tensor/rl/trainer_model_free_tictactoe_test.py b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
index 76fb37339..f57e46546 100644
--- a/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
+++ b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
@@ -40,7 +40,8 @@ def test_train_tictactoe(self):
     hparams.eval_every_epochs = 25
 
     FLAGS.output_dir = tf.test.get_temp_dir()
-    trainer_model_free.train(hparams, FLAGS.output_dir)
+    FLAGS.env_problem_name = "tic_tac_toe_env_problem"
+    trainer_model_free.train(hparams, FLAGS.output_dir, FLAGS.env_problem_name)
 
 
 if __name__ == "__main__":

From d17df09faef4a1042041d4938f2de6974cf891f3 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 11 Feb 2019 15:49:58 -0800
Subject: [PATCH 1664/2720] Replace add_weight boilerplate with decorator.

This makes it even simpler to add a new Bayesian Layer.

Speaking self-critically, this feels like a bit of overoptimization. However, I implemented it now because I was curious if it was even possible. :-)

PiperOrigin-RevId: 233494984
---
 tensor2tensor/layers/bayes.py | 137 +++++++++-------------------------
 1 file changed, 36 insertions(+), 101 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 558d960a1..581438625 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -375,6 +375,36 @@ def get_config(self):
     }
 
 
+def _add_weight(layer,
+                name=None,
+                shape=None,
+                dtype=None,
+                initializer=None,
+                regularizer=None,
+                **kwargs):
+  """Adds weight."""
+  if isinstance(initializer, tf.keras.layers.Layer):
+    weight = initializer(shape, dtype)
+    layer._trainable_weights.extend(initializer.trainable_weights)  # pylint: disable=protected-access
+    layer._non_trainable_weights.extend(initializer.non_trainable_weights)  # pylint: disable=protected-access
+    if regularizer is not None:
+      # TODO(trandustin): Replace need for this with
+      # Layer._handle_weight_regularization. For Eager compatibility, random
+      # variable __init__s cannot apply TF ops (cl/220898007).
+      def loss_fn():
+        """Creates a regularization loss `Tensor`."""
+        with tf.name_scope(name + '/Regularizer'):
+          return regularizer(initializer(shape, dtype))
+      layer.add_loss(loss_fn)
+    return weight
+  return super(layer.__class__, layer).add_weight(name=name,
+                                                  shape=shape,
+                                                  dtype=dtype,
+                                                  initializer=initializer,
+                                                  regularizer=regularizer,
+                                                  **kwargs)
+
+
 class DenseReparameterization(tf.keras.layers.Dense):
   """Bayesian densely-connected layer estimated via reparameterization.
 
@@ -415,6 +445,8 @@ def __init__(self,
         activity_regularizer=get(activity_regularizer),
         **kwargs)
 
+  add_weight = _add_weight
+
   # TODO(trandustin): This name is not accurate. Rename or move functionality
   # into random variables to resample/recreate their init ops.
   def sample_weights(self):
@@ -427,31 +459,6 @@ def call(self, *args, **kwargs):
     self.sample_weights()
     return super(DenseReparameterization, self).call(*args, **kwargs)
 
-  def add_weight(self,
-                 name=None,
-                 shape=None,
-                 dtype=None,
-                 initializer=None,
-                 regularizer=None,
-                 **kwargs):
-    if isinstance(initializer, tf.keras.layers.Layer):
-      weight = initializer(shape, dtype)
-      self._trainable_weights.extend(initializer.trainable_weights)
-      self._non_trainable_weights.extend(initializer.non_trainable_weights)
-      if regularizer is not None:
-        self.add_loss(
-            create_regularization_loss_fn(name,
-                                          lambda: initializer(shape, dtype),
-                                          regularizer))
-      return weight
-    return super(DenseReparameterization, self).add_weight(
-        name=name,
-        shape=shape,
-        dtype=dtype,
-        initializer=initializer,
-        regularizer=regularizer,
-        **kwargs)
-
 
 class Conv2DReparameterization(tf.keras.layers.Conv2D):
   """2D convolution layer (e.g. spatial convolution over images).
@@ -507,6 +514,8 @@ def __init__(self,
         bias_constraint=get(bias_constraint),
         **kwargs)
 
+  add_weight = _add_weight
+
   def sample_weights(self):
     if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
       self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
@@ -517,31 +526,6 @@ def call(self, *args, **kwargs):
     self.sample_weights()
     return super(Conv2DReparameterization, self).call(*args, **kwargs)
 
-  def add_weight(self,
-                 name=None,
-                 shape=None,
-                 dtype=None,
-                 initializer=None,
-                 regularizer=None,
-                 **kwargs):
-    if isinstance(initializer, tf.keras.layers.Layer):
-      weight = initializer(shape, dtype)
-      self._trainable_weights.extend(initializer.trainable_weights)
-      self._non_trainable_weights.extend(initializer.non_trainable_weights)
-      if regularizer is not None:
-        self.add_loss(
-            create_regularization_loss_fn(name,
-                                          lambda: initializer(shape, dtype),
-                                          regularizer))
-      return weight
-    return super(Conv2DReparameterization, self).add_weight(
-        name=name,
-        shape=shape,
-        dtype=dtype,
-        initializer=initializer,
-        regularizer=regularizer,
-        **kwargs)
-
 
 class GaussianProcess(tf.keras.layers.Layer):
   r"""Gaussian process layer.
@@ -756,6 +740,8 @@ def __init__(self,
         implementation=implementation,
         **kwargs)
 
+  add_weight = _add_weight
+
   def build(self, input_shape):
     input_shape = tf.TensorShape(input_shape)
     input_dim = input_shape[-1]
@@ -824,57 +810,6 @@ def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
     return super(LSTMCellReparameterization, self).get_initial_state(
         inputs=inputs, batch_size=batch_size, dtype=dtype)
 
-  def add_weight(self,
-                 name=None,
-                 shape=None,
-                 dtype=None,
-                 initializer=None,
-                 regularizer=None,
-                 **kwargs):
-    if isinstance(initializer, tf.keras.layers.Layer):
-      weight = initializer(shape, dtype)
-      self._trainable_weights.extend(initializer.trainable_weights)
-      self._non_trainable_weights.extend(initializer.non_trainable_weights)
-      if regularizer is not None:
-        self.add_loss(
-            create_regularization_loss_fn(name,
-                                          lambda: initializer(shape, dtype),
-                                          regularizer))
-      return weight
-    return super(LSTMCellReparameterization, self).add_weight(
-        name=name,
-        shape=shape,
-        dtype=dtype,
-        initializer=initializer,
-        regularizer=regularizer,
-        **kwargs)
-
-
-# TODO(trandustin): Replace need for this function with
-# Layer._handle_weight_regularization. For Eager compatibility, random variable
-# __init__s cannot apply TF ops (cl/220898007).
-def create_regularization_loss_fn(name, variable_fn, regularizer_fn):
-  """Create a regularization loss function.
-
-  The callable representing the variable allows for use with Bayesian Layers.
-
-  Args:
-    name: String name scope prefix.
-    variable_fn: Callable that returns a TF Variable or ed.RandomVariable.
-    regularizer_fn: Callable that returns a loss tensor when called with a TF
-      Variable or ed.RandomVariable.
-
-  Returns:
-    A callable that returns a regularization loss tensor when called.
-  """
-  def loss_fn():
-    """Creates a regularization loss `Tensor`."""
-    with tf.name_scope(name + '/Regularizer'):
-      regularization = regularizer_fn(variable_fn())
-    return regularization
-
-  return loss_fn
-
 
 class BayesianLinearModel(tf.keras.Model):
   r"""Bayesian linear model with standard normal prior over its coefficients.

From 68aa47a4bbc1e83874740a82a36d6b48af25acff Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 11 Feb 2019 20:13:51 -0800
Subject: [PATCH 1665/2720] [V2] Porting Transformer, first step.

PiperOrigin-RevId: 233530607
---
 tensor2tensor/layers/transformer_layers.py |   3 +-
 tensor2tensor/models/transformer.py        | 230 ++++++++++++---------
 tensor2tensor/utils/data_reader.py         |  28 +--
 tensor2tensor/utils/data_reader_test.py    |  14 +-
 tensor2tensor/v2/keras_utils.py            |  13 +-
 tensor2tensor/v2/models/basic.py           |   4 +-
 tensor2tensor/v2/models/resnet.py          |   2 +-
 tensor2tensor/v2/models/transformer.py     | 130 ++++++++++++
 tensor2tensor/v2/t2t.py                    | 156 ++++++++++++--
 9 files changed, 441 insertions(+), 139 deletions(-)
 create mode 100644 tensor2tensor/v2/models/transformer.py

diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index d019eef71..8288e9761 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -81,7 +81,8 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
   if hparams.proximity_bias:
     encoder_self_attention_bias += common_attention.attention_bias_proximal(
         common_layers.shape_list(inputs)[1])
-  if hparams.get("use_target_space_embedding", True):
+  if target_space is not None and hparams.get("use_target_space_embedding",
+                                              True):
     # Append target_space_id embedding to inputs.
     emb_target_space = common_layers.embedding(
         target_space,
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 43b879cd2..479dfd45e 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -52,6 +52,124 @@
 transformer_ffn_layer = transformer_layers.transformer_ffn_layer
 
 
+def transformer_encode(encoder_function, inputs, target_space, hparams,
+                       attention_weights=None, features=None, losses=None):
+  """Encode transformer inputs.
+
+  Args:
+    encoder_function: the encoder function
+    inputs: Transformer inputs [batch_size, input_length, 1, hidden_dim] which
+      will be flattened along the two spatial dimensions.
+    target_space: scalar, target space ID.
+    hparams: hyperparameters for model.
+    attention_weights: weight to store attention to.
+    features: optionally pass the entire features dictionary as well. This is
+      needed now for "packed" datasets.
+    losses: optional list onto which to append extra training losses
+
+  Returns:
+    Tuple of:
+        encoder_output: Encoder representation.
+            [batch_size, input_length, hidden_dim]
+        encoder_decoder_attention_bias: Bias and mask weights for
+            encoder-decoder attention. [batch_size, input_length]
+  """
+  inputs = common_layers.flatten4d3d(inputs)
+
+  encoder_input, self_attention_bias, encoder_decoder_attention_bias = (
+      transformer_prepare_encoder(
+          inputs, target_space, hparams, features=features))
+
+  mlperf_log.transformer_print(
+      key=mlperf_log.MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
+      value=hparams.layer_prepostprocess_dropout,
+      hparams=hparams)
+
+  encoder_input = tf.nn.dropout(encoder_input,
+                                1.0 - hparams.layer_prepostprocess_dropout)
+
+  attn_bias_for_padding = None
+  # Otherwise the encoder will just use encoder_self_attention_bias.
+  if hparams.unidirectional_encoder:
+    attn_bias_for_padding = encoder_decoder_attention_bias
+
+  encoder_output = encoder_function(
+      encoder_input,
+      self_attention_bias,
+      hparams,
+      nonpadding=features_to_nonpadding(features, "inputs"),
+      save_weights_to=attention_weights,
+      make_image_summary=not common_layers.is_xla_compiled(),
+      losses=losses,
+      attn_bias_for_padding=attn_bias_for_padding)
+
+  return encoder_output, encoder_decoder_attention_bias
+
+
+def transformer_decode(decoder_function,
+                       decoder_input,
+                       encoder_output,
+                       encoder_decoder_attention_bias,
+                       decoder_self_attention_bias,
+                       hparams,
+                       attention_weights=None,
+                       cache=None,
+                       decode_loop_step=None,
+                       nonpadding=None,
+                       losses=None):
+  """Decode Transformer outputs from encoder representation.
+
+  Args:
+    decoder_function: the decoder function
+    decoder_input: inputs to bottom of the model. [batch_size, decoder_length,
+      hidden_dim]
+    encoder_output: Encoder representation. [batch_size, input_length,
+      hidden_dim]
+    encoder_decoder_attention_bias: Bias and mask weights for encoder-decoder
+      attention. [batch_size, input_length]
+    decoder_self_attention_bias: Bias and mask weights for decoder
+      self-attention. [batch_size, decoder_length]
+    hparams: hyperparameters for model.
+    attention_weights: weight to store attention to.
+    cache: dict, containing tensors which are the results of previous
+      attentions, used for fast decoding.
+    decode_loop_step: An integer, step number of the decoding loop. Only used
+      for inference on TPU.
+    nonpadding: optional Tensor with shape [batch_size, decoder_length]
+    losses: optional list onto which to append extra training losses
+
+  Returns:
+    Final decoder representation. [batch_size, decoder_length, hidden_dim]
+  """
+  mlperf_log.transformer_print(
+      key=mlperf_log.MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
+      value=hparams.layer_prepostprocess_dropout,
+      hparams=hparams)
+  decoder_input = tf.nn.dropout(decoder_input,
+                                1.0 - hparams.layer_prepostprocess_dropout)
+
+  decoder_output = decoder_function(
+      decoder_input,
+      encoder_output,
+      decoder_self_attention_bias,
+      encoder_decoder_attention_bias,
+      hparams,
+      cache=cache,
+      decode_loop_step=decode_loop_step,
+      nonpadding=nonpadding,
+      save_weights_to=attention_weights,
+      losses=losses)
+
+  if (common_layers.is_xla_compiled() and
+      hparams.mode == tf.estimator.ModeKeys.TRAIN):
+    # TPU does not react kindly to extra dimensions.
+    # TODO(noam): remove this once TPU is more forgiving of extra dims.
+    return decoder_output
+  else:
+    # Expand since t2t expects 4d tensors.
+    return tf.expand_dims(decoder_output, axis=2)
+
+
 @registry.register_model
 class Transformer(t2t_model.T2TModel):
   """Attention net.  See file docstring."""
@@ -63,54 +181,11 @@ def __init__(self, *args, **kwargs):
     self._decoder_function = transformer_decoder
 
   def encode(self, inputs, target_space, hparams, features=None, losses=None):
-    """Encode transformer inputs.
-
-    Args:
-      inputs: Transformer inputs [batch_size, input_length, 1, hidden_dim] which
-        will be flattened along the two spatial dimensions.
-      target_space: scalar, target space ID.
-      hparams: hyperparameters for model.
-      features: optionally pass the entire features dictionary as well. This is
-        needed now for "packed" datasets.
-      losses: optional list onto which to append extra training losses
-
-    Returns:
-      Tuple of:
-          encoder_output: Encoder representation.
-              [batch_size, input_length, hidden_dim]
-          encoder_decoder_attention_bias: Bias and mask weights for
-              encoder-decoder attention. [batch_size, input_length]
-    """
-    inputs = common_layers.flatten4d3d(inputs)
-
-    encoder_input, self_attention_bias, encoder_decoder_attention_bias = (
-        transformer_prepare_encoder(
-            inputs, target_space, hparams, features=features))
-
-    mlperf_log.transformer_print(
-        key=mlperf_log.MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
-        value=hparams.layer_prepostprocess_dropout,
-        hparams=hparams)
-
-    encoder_input = tf.nn.dropout(encoder_input,
-                                  1.0 - hparams.layer_prepostprocess_dropout)
-
-    attn_bias_for_padding = None
-    # Otherwise the encoder will just use encoder_self_attention_bias.
-    if hparams.unidirectional_encoder:
-      attn_bias_for_padding = encoder_decoder_attention_bias
-
-    encoder_output = self._encoder_function(
-        encoder_input,
-        self_attention_bias,
-        hparams,
-        nonpadding=features_to_nonpadding(features, "inputs"),
-        save_weights_to=self.attention_weights,
-        make_image_summary=not common_layers.is_xla_compiled(),
-        losses=losses,
-        attn_bias_for_padding=attn_bias_for_padding)
-
-    return encoder_output, encoder_decoder_attention_bias
+    """Encode transformer inputs, see transformer_encode."""
+    return transformer_encode(
+        self._encoder_function, inputs, target_space, hparams,
+        attention_weights=self.attention_weights,
+        features=features, losses=losses)
 
   def decode(self,
              decoder_input,
@@ -122,55 +197,12 @@ def decode(self,
              decode_loop_step=None,
              nonpadding=None,
              losses=None):
-    """Decode Transformer outputs from encoder representation.
-
-    Args:
-      decoder_input: inputs to bottom of the model. [batch_size, decoder_length,
-        hidden_dim]
-      encoder_output: Encoder representation. [batch_size, input_length,
-        hidden_dim]
-      encoder_decoder_attention_bias: Bias and mask weights for encoder-decoder
-        attention. [batch_size, input_length]
-      decoder_self_attention_bias: Bias and mask weights for decoder
-        self-attention. [batch_size, decoder_length]
-      hparams: hyperparameters for model.
-      cache: dict, containing tensors which are the results of previous
-        attentions, used for fast decoding.
-      decode_loop_step: An integer, step number of the decoding loop. Only used
-        for inference on TPU.
-      nonpadding: optional Tensor with shape [batch_size, decoder_length]
-      losses: optional list onto which to append extra training losses
-
-    Returns:
-      Final decoder representation. [batch_size, decoder_length, hidden_dim]
-    """
-    mlperf_log.transformer_print(
-        key=mlperf_log.MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
-        value=hparams.layer_prepostprocess_dropout,
-        hparams=hparams)
-    decoder_input = tf.nn.dropout(decoder_input,
-                                  1.0 - hparams.layer_prepostprocess_dropout)
-
-    decoder_output = self._decoder_function(
-        decoder_input,
-        encoder_output,
-        decoder_self_attention_bias,
-        encoder_decoder_attention_bias,
-        hparams,
-        cache=cache,
-        decode_loop_step=decode_loop_step,
-        nonpadding=nonpadding,
-        save_weights_to=self.attention_weights,
-        losses=losses)
-
-    if (common_layers.is_xla_compiled() and
-        hparams.mode == tf.estimator.ModeKeys.TRAIN):
-      # TPU does not react kindly to extra dimensions.
-      # TODO(noam): remove this once TPU is more forgiving of extra dims.
-      return decoder_output
-    else:
-      # Expand since t2t expects 4d tensors.
-      return tf.expand_dims(decoder_output, axis=2)
+    """Decode Transformer outputs, see transformer_decode."""
+    return transformer_decode(
+        self._decoder_function, decoder_input, encoder_output,
+        encoder_decoder_attention_bias, decoder_self_attention_bias,
+        hparams, attention_weights=self.attention_weights, cache=cache,
+        decode_loop_step=decode_loop_step, nonpadding=nonpadding, losses=losses)
 
   def body(self, features):
     """Transformer main model_fn.
@@ -2028,7 +2060,9 @@ def transformer_big_single_gpu():
 def transformer_base_single_gpu():
   """HParams for transformer base model for single GPU."""
   hparams = transformer_base()
-  hparams.batch_size = 2048
+  hparams.batch_size = 1024
+  hparams.learning_rate_schedule = "constant*linear_warmup*rsqrt_decay"
+  hparams.learning_rate_constant = 0.1
   hparams.learning_rate_warmup_steps = 16000
   return hparams
 
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index ee063ebdf..3456c7fe0 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -77,14 +77,14 @@ def _bucket_boundaries(max_length, min_length=8, length_bucket_step=1.1):
   return boundaries
 
 
-def _batching_scheme(batch_size,
-                     max_length,
-                     min_length_bucket,
-                     length_bucket_step,
-                     drop_long_sequences=False,
-                     shard_multiplier=1,
-                     length_multiplier=1,
-                     min_length=0):
+def batching_scheme(batch_size,
+                    max_length,
+                    min_length_bucket,
+                    length_bucket_step,
+                    drop_long_sequences=False,
+                    shard_multiplier=1,
+                    length_multiplier=1,
+                    min_length=0):
   """A batching scheme based on model hyperparameters.
 
   Every batch contains a number of sequences divisible by `shard_multiplier`.
@@ -169,7 +169,7 @@ def hparams_to_batching_scheme(hparams,
                                shard_multiplier=1,
                                length_multiplier=1):
   """Wrapper around _batching_scheme with hparams."""
-  return _batching_scheme(
+  return batching_scheme(
       batch_size=hparams.batch_size,
       min_length=hparams.min_length,
       max_length=hparams.max_length,
@@ -437,18 +437,18 @@ def define_shapes(example):
     else:
       # On GPU, bucket by length
       dataset = dataset.filter(gpu_valid_size)
-      batching_scheme = hparams_to_batching_scheme(
+      cur_batching_scheme = hparams_to_batching_scheme(
           hparams,
           shard_multiplier=num_shards,
           length_multiplier=batch_size_multiplier)
       if hparams.use_fixed_batch_size:
         # Here  batch_size really means examples per datashard.
-        batching_scheme["batch_sizes"] = [hparams.batch_size]
-        batching_scheme["boundaries"] = []
+        cur_batching_scheme["batch_sizes"] = [hparams.batch_size]
+        cur_batching_scheme["boundaries"] = []
       dataset = dataset.apply(
           tf.data.experimental.bucket_by_sequence_length(
-              example_length, batching_scheme["boundaries"],
-              batching_scheme["batch_sizes"]))
+              example_length, cur_batching_scheme["boundaries"],
+              cur_batching_scheme["batch_sizes"]))
 
       if not is_training:
         batch_multiple = num_shards
diff --git a/tensor2tensor/utils/data_reader_test.py b/tensor2tensor/utils/data_reader_test.py
index 32544da36..ef4c1f314 100644
--- a/tensor2tensor/utils/data_reader_test.py
+++ b/tensor2tensor/utils/data_reader_test.py
@@ -139,7 +139,7 @@ def testLengthFilter(self):
     self.assertAllEqual(list(range(1, max_len + 1)), sorted(ex_lens))
 
   def testBatchingSchemeMaxLength(self):
-    scheme = data_reader._batching_scheme(
+    scheme = data_reader.batching_scheme(
         batch_size=20,
         max_length=None,
         min_length_bucket=8,
@@ -147,7 +147,7 @@ def testBatchingSchemeMaxLength(self):
         drop_long_sequences=False)
     self.assertGreater(scheme["max_length"], 10000)
 
-    scheme = data_reader._batching_scheme(
+    scheme = data_reader.batching_scheme(
         batch_size=20,
         max_length=None,
         min_length_bucket=8,
@@ -155,7 +155,7 @@ def testBatchingSchemeMaxLength(self):
         drop_long_sequences=True)
     self.assertEqual(scheme["max_length"], 20)
 
-    scheme = data_reader._batching_scheme(
+    scheme = data_reader.batching_scheme(
         batch_size=20,
         max_length=15,
         min_length_bucket=8,
@@ -163,7 +163,7 @@ def testBatchingSchemeMaxLength(self):
         drop_long_sequences=True)
     self.assertEqual(scheme["max_length"], 15)
 
-    scheme = data_reader._batching_scheme(
+    scheme = data_reader.batching_scheme(
         batch_size=20,
         max_length=15,
         min_length_bucket=8,
@@ -172,7 +172,7 @@ def testBatchingSchemeMaxLength(self):
     self.assertGreater(scheme["max_length"], 10000)
 
   def testBatchingSchemeBuckets(self):
-    scheme = data_reader._batching_scheme(
+    scheme = data_reader.batching_scheme(
         batch_size=128,
         max_length=0,
         min_length_bucket=8,
@@ -190,7 +190,7 @@ def testBatchingSchemeBuckets(self):
     ]
     self.assertEqual(expected_batch_sizes, batch_sizes)
 
-    scheme = data_reader._batching_scheme(
+    scheme = data_reader.batching_scheme(
         batch_size=128,
         max_length=0,
         min_length_bucket=8,
@@ -200,7 +200,7 @@ def testBatchingSchemeBuckets(self):
     self.assertAllEqual([bs * 2 for bs in expected_batch_sizes], batch_sizes)
     self.assertEqual(expected_boundaries, boundaries)
 
-    scheme = data_reader._batching_scheme(
+    scheme = data_reader.batching_scheme(
         batch_size=128,
         max_length=0,
         min_length_bucket=8,
diff --git a/tensor2tensor/v2/keras_utils.py b/tensor2tensor/v2/keras_utils.py
index 566da85f1..251b03796 100644
--- a/tensor2tensor/v2/keras_utils.py
+++ b/tensor2tensor/v2/keras_utils.py
@@ -30,10 +30,21 @@ def __init__(self, function, name=None):
       name = function.__name__
     super(FunctionLayer, self).__init__(name=name)
     self._template = tf.compat.v1.make_template(name, function)
+    self._was_called = False
 
   @property
   def losses(self):
     return []
 
-  def call(self, *args, **kwargs):
+  def compute_mask(self, inputs, previous_mask):
+    return previous_mask
+
+  @tf.function
+  def _template_call(self, *args, **kwargs):
+    """Call to template but made in graph mode for better speed."""
     return self._template(*args, **kwargs)
+
+  def call(self, *args, **kwargs):
+    if not self._was_called:  # Create variables at first call.
+      return self._template(*args, **kwargs)
+    return self._template_call(*args, **kwargs)
diff --git a/tensor2tensor/v2/models/basic.py b/tensor2tensor/v2/models/basic.py
index 8ed3d758e..113dcafc2 100644
--- a/tensor2tensor/v2/models/basic.py
+++ b/tensor2tensor/v2/models/basic.py
@@ -45,7 +45,7 @@ def __init__(self, features_info=None, input_names=None, target_names=None,
       self._dropout_layers.append(tf.keras.layers.Dropout(
           rate=dropout))
     self._logits = tf.keras.layers.Dense(
-        num_output_classes, activation="softmax")
+        num_output_classes, activation=None)
 
   def call(self, inputs, training=False):
     x = tf.cast(inputs[self._input_name], tf.float32) / 255.0
@@ -89,4 +89,4 @@ def call(self, inputs, training=False):
           self._hidden_size, activation="relu", name="layer_%d" % i)(x)
       x = tf.keras.layers.Dropout(rate=self._dropout)(x, training=training)
     return tf.keras.layers.Dense(
-        self._num_output_classes, activation="softmax")(x)
+        self._num_output_classes, activation=None)(x)
diff --git a/tensor2tensor/v2/models/resnet.py b/tensor2tensor/v2/models/resnet.py
index 33406ba73..a471522b0 100644
--- a/tensor2tensor/v2/models/resnet.py
+++ b/tensor2tensor/v2/models/resnet.py
@@ -52,7 +52,7 @@ def resnet_model(inputs, training):
 
     self._resnet = keras_utils.FunctionLayer(resnet_model)
     self._logits = tf.keras.layers.Dense(
-        num_output_classes, activation="softmax")
+        num_output_classes, activation=None)
 
   def call(self, inputs, training=False):
     x = tf.cast(inputs[self._input_name], tf.float32) / 255.0
diff --git a/tensor2tensor/v2/models/transformer.py b/tensor2tensor/v2/models/transformer.py
new file mode 100644
index 000000000..edb27dcf2
--- /dev/null
+++ b/tensor2tensor/v2/models/transformer.py
@@ -0,0 +1,130 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Transformer model from "Attention Is All You Need".
+
+The Transformer model consists of an encoder and a decoder. Both are stacks
+of self-attention layers followed by feed-forward layers. This model yields
+good results on a number of problems, especially in NLP and machine translation.
+
+See "Attention Is All You Need" (https://arxiv.org/abs/1706.03762) for the full
+description of the model and the results obtained with its early version.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import transformer_layers
+from tensor2tensor.models import transformer
+from tensor2tensor.v2 import keras_utils
+import tensorflow as tf
+import gin.tf
+
+
+@gin.configurable(whitelist=["hidden_size", "filter_size"])
+class Transformer(tf.keras.Model):
+  """Transformer."""
+
+  def __init__(self, features_info=None, input_names=None, target_names=None,
+               hidden_size=512, filter_size=2048):
+    super(Transformer, self).__init__()
+    # TODO(lukaszkaiser): gin'ify and split into encoder/decoder classes.
+    self._has_input = True if input_names else False
+    self._input_name = input_names[0]
+    self._target_name = target_names[0]
+    try:
+      target_vocab_size = features_info[self._target_name].num_classes
+    except AttributeError:
+      target_vocab_size = features_info[self._target_name].encoder.vocab_size
+    hparams = transformer.transformer_base()
+    hparams.hidden_size = hidden_size
+    hparams.filter_size = filter_size
+
+    # Now the model.
+    self._embedding = tf.keras.layers.Embedding(
+        target_vocab_size, hidden_size, mask_zero=True)
+    def transformer_encoder(inputs, features):
+      return transformer.transformer_encode(
+          transformer_layers.transformer_encoder, inputs, None,
+          hparams, features=features)
+
+    def transformer_prepare_decoder(targets, features):
+      return transformer.transformer_prepare_decoder(targets, hparams, features)
+
+    def transformer_decoder(decoder_input, encoder_output,
+                            encoder_decoder_attention_bias,
+                            decoder_self_attention_bias,
+                            features):
+      return transformer.transformer_decode(
+          transformer.transformer_decoder,
+          decoder_input,
+          encoder_output,
+          encoder_decoder_attention_bias,
+          decoder_self_attention_bias,
+          hparams,
+          nonpadding=transformer.features_to_nonpadding(features, "targets"))
+
+    if self._has_input:
+      self._encoder = keras_utils.FunctionLayer(transformer_encoder)
+    self._prepare_decoder = keras_utils.FunctionLayer(
+        transformer_prepare_decoder)
+    self._decoder = keras_utils.FunctionLayer(transformer_decoder)
+    self._logits = tf.keras.layers.Dense(
+        target_vocab_size, activation=None)
+
+  def call(self, features, training=False):
+    """Transformer main model_fn.
+
+    Args:
+      features: Map of features to the model. Should contain the following:
+          "inputs": Transformer inputs. [batch_size, input_length, 1,
+            hidden_dim].
+          "targets": Target decoder outputs. [batch_size, decoder_length, 1,
+            hidden_dim]
+          "target_space_id": A scalar int from data_generators.problem.SpaceID.
+      training: Whether we are training or not.
+
+    Returns:
+      Final decoder representation. [batch_size, decoder_length, hidden_dim]
+    """
+    if self._has_input:
+      inputs = features[self._input_name]
+      inputs = tf.expand_dims(self._embedding(inputs), 2)
+      encoder_output, encoder_decoder_attention_bias = self._encoder(
+          inputs, features)
+    else:
+      encoder_output, encoder_decoder_attention_bias = (None, None)
+
+    targets = features[self._target_name]
+    targets = self._embedding(targets)
+    decoder_input, decoder_self_attention_bias = self._prepare_decoder(
+        targets, features)
+    decoder_output = self._decoder(
+        decoder_input,
+        encoder_output,
+        encoder_decoder_attention_bias,
+        decoder_self_attention_bias,
+        features)
+
+    return self._logits(tf.squeeze(decoder_output, axis=2))
+
+
+def transformer_base_single_gpu():
+  """Single-gpu set of parameters for Transformer."""
+  gin.bind_parameter("T2TLearningRateSchedule.warmup_steps", 16000)
+  gin.bind_parameter("preprocess_fn.max_target_length", 256)
+  gin.bind_parameter("batch_fn.eval_batch_size", 8)
+  return Transformer
diff --git a/tensor2tensor/v2/t2t.py b/tensor2tensor/v2/t2t.py
index a4bf77214..ce2e906ee 100644
--- a/tensor2tensor/v2/t2t.py
+++ b/tensor2tensor/v2/t2t.py
@@ -20,9 +20,13 @@
 from __future__ import print_function
 
 import collections
+import os
+
 from tensor2tensor import problems
+from tensor2tensor.utils import data_reader
 from tensor2tensor.v2.models import basic
 from tensor2tensor.v2.models import resnet
+from tensor2tensor.v2.models import transformer
 
 import tensorflow as tf
 import tensorflow_datasets as tfds
@@ -37,6 +41,7 @@
     "basic_fc_large": basic.basic_fc_large,
     "basic_fc_relu_v2": lambda: basic.BasicFcReluV2,
     "resnet": lambda: resnet.Resnet,
+    "transformer": transformer.transformer_base_single_gpu,
 }
 
 
@@ -93,11 +98,19 @@ def _make_info(shape_list, num_classes):
   return feature_info(cur_shape, num_classes)
 
 
+def _select_features(example, feature_list=None):
+  """Select a subset of features from the example dict."""
+  feature_list = feature_list or ["inputs", "targets"]
+  return {f: example[f] for f in feature_list}
+
+
 def _train_and_eval_dataset_v1(problem_name, data_dir):
   """Return train and evaluation datasets, feature info and supervised keys."""
   problem = problems.problem(problem_name)
   train_dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, data_dir)
+  train_dataset = train_dataset.map(_select_features)
   eval_dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, data_dir)
+  eval_dataset = eval_dataset.map(_select_features)
   supervised_keys = (["inputs"], ["targets"])
   hparams = problem.get_hparams()
   # We take a few training examples to guess the shapes.
@@ -113,7 +126,55 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
   return train_dataset, eval_dataset, info, supervised_keys
 
 
-def shuffle_and_batch_data(dataset, batch_size, target_names, repeat=False):
+@gin.configurable(blacklist=["dataset", "training"])
+def preprocess_fn(dataset, training, max_target_length=-1):
+  def target_right_length(_, target):
+    if max_target_length < 1 or not training:
+      return tf.constant(True)
+    return tf.less(tf.shape(target)[0], max_target_length + 1)
+  dataset = dataset.filter(target_right_length)
+  return dataset
+
+
+@gin.configurable(blacklist=["dataset", "training", "shapes", "target_names"])
+def batch_fn(dataset, training, shapes, target_names,
+             batch_size=32, eval_batch_size=32, bucket_batch_length=32,
+             bucket_max_length=256, bucket_min_length=8,
+             bucket_length_step=1.1, buckets=None):
+  """Batching function."""
+  del target_names
+  # If bucketing is not specified, check if target shapes are variable.
+  cur_batch_size = batch_size if training else eval_batch_size
+  if buckets is None:
+    variable_target_shapes = False
+    target_shape = shapes[1]
+    for dim in target_shape:
+      if dim is None:
+        variable_target_shapes = True
+    tf.logging.info("Heuristically setting bucketing to %s based on shapes "
+                    "of target tensors." % variable_target_shapes)
+    if variable_target_shapes:
+      batch_size_per_token = cur_batch_size * bucket_batch_length
+      scheme = data_reader.batching_scheme(batch_size_per_token,
+                                           bucket_max_length,
+                                           bucket_min_length,
+                                           bucket_length_step,
+                                           drop_long_sequences=training)
+      buckets = (scheme["boundaries"], scheme["batch_sizes"])
+
+  if buckets:
+    tf.logging.info("Bucketing with buckets %s." % str(buckets))
+    def example_length(_, target):
+      return tf.shape(target)[0]
+    boundaries, batch_sizes = buckets
+    dataset = dataset.apply(tf.data.experimental.bucket_by_sequence_length(
+        example_length, boundaries, batch_sizes))
+  else:
+    dataset = dataset.padded_batch(cur_batch_size, shapes)
+  return dataset
+
+
+def shuffle_and_batch_data(dataset, target_names, features_info, training):
   """Shuffle and batch the given dataset."""
   def append_targets(example):
     """Append targets to the example dictionary. Needed for Keras."""
@@ -124,19 +185,69 @@ def append_targets(example):
       targets[name] = example[name]
     return (example, targets)
   dataset = dataset.map(append_targets)
-  if repeat:
+  if training:
     dataset = dataset.repeat()
-  shuffled = dataset.shuffle(128).batch(batch_size).prefetch(8)
-  return shuffled
+  shapes = {k: features_info[k].shape for k in features_info}
+  shapes = (shapes, shapes[target_names[0]])
+  dataset = dataset.shuffle(128)
+  dataset = preprocess_fn(dataset, training)
+  dataset = batch_fn(dataset, training, shapes, target_names)
+  return dataset.prefetch(8)
+
+
+@gin.configurable()
+class T2TLearningRateSchedule(
+    tf.keras.optimizers.schedules.LearningRateSchedule):
+  """A LearningRateSchedule that uses a T2T config."""
+
+  def __init__(self, schedule=None, constant=0.1, warmup_steps=200):
+    """Applies the give T2T schedule string with the given parameters."""
+    super(T2TLearningRateSchedule, self).__init__()
+    self.schedule = schedule or "constant * linear_warmup * rsqrt_decay"
+    self.constant = constant
+    self.warmup_steps = warmup_steps
+
+  def __call__(self, step):
+    ret = tf.constant(1.0)
+    for name in [n.strip() for n in self.schedule.split("*")]:
+      if name == "constant":
+        ret *= self.constant
+      elif name == "linear_warmup":
+        ret *= tf.minimum(1.0, step / self.warmup_steps)
+      elif name == "rsqrt_decay":
+        ret *= tf.rsqrt(tf.maximum(step, self.warmup_steps))
+      else:
+        raise ValueError("Unknown factor %s." % name)
+    tf.contrib.summary.scalar("learning_rate", ret)
+    return ret
+
+  def get_config(self):
+    return {
+        "schedule": self.schedule,
+        "constant": self.constant,
+        "warmup_steps": self.warmup_steps,
+    }
 
 
 @gin.configurable(blacklist=["model"])
-def model_compile(model,
-                  optimizer="adam",
-                  loss="sparse_categorical_crossentropy",
-                  metrics=None):
+def optimize_fn(model,
+                optimizer=None,
+                learning_rate_schedule=None,
+                loss=None,
+                metrics=None):
   """Compile the model in Keras."""
-  metrics = ["accuracy"] if metrics is None else metrics
+  learning_rate_schedule = learning_rate_schedule or T2TLearningRateSchedule()
+  if optimizer:
+    optimizer = optimizer(learning_rate=learning_rate_schedule)
+  else:  # We use Adam by default with adjusted parameters.
+    optimizer = tf.keras.optimizers.Adam(
+        learning_rate=learning_rate_schedule,
+        beta_1=0.9, beta_2=0.997, epsilon=1e-9)
+  metrics = metrics or [tf.keras.metrics.sparse_categorical_accuracy]
+  def xent_loss(y, x):
+    return tf.keras.backend.sparse_categorical_crossentropy(
+        y, x, from_logits=True)
+  loss = loss or xent_loss
   return model.compile(optimizer=optimizer,
                        loss=loss,
                        metrics=metrics)
@@ -148,7 +259,7 @@ def model_compile(model,
 def train_fn(data_dir=None, output_dir=None,
              model_class=gin.REQUIRED, dataset=gin.REQUIRED,
              input_names=None, target_names=None,
-             batch_size=32, train_steps=1000, eval_steps=1, eval_frequency=100):
+             train_steps=1000, eval_steps=1, eval_frequency=100):
   """Train the given model on the given dataset.
 
   Args:
@@ -158,7 +269,6 @@ def train_fn(data_dir=None, output_dir=None,
     dataset: The name of the dataset to train on.
     input_names: List of strings with the names of the features on input.
     target_names: List of strings with the names of the target features.
-    batch_size: integer, how many examples per batch.
     train_steps: for how many steps to train.
     eval_steps: for how many steps to do evaluation.
     eval_frequency: how often (every this many steps) to run evaluation.
@@ -177,24 +287,40 @@ def train_fn(data_dir=None, output_dir=None,
   # with strategy.scope():
   model = model_class(features_info=features_info,
                       input_names=input_names, target_names=target_names)
-  model_compile(model)
+  optimize_fn(model)
   train_batches = shuffle_and_batch_data(
-      train_data, batch_size, target_names, repeat=True)
-  eval_batches = shuffle_and_batch_data(eval_data, batch_size, target_names)
+      train_data, target_names, features_info, training=True)
+  eval_batches = shuffle_and_batch_data(
+      eval_data, target_names, features_info, training=False)
+  # Need to run one training step just to get optimizer variables to load.
+  model.fit(train_batches, epochs=1, steps_per_epoch=1)
 
   # Training loop.
   callbacks = []
   callbacks.append(tf.keras.callbacks.History())
   callbacks.append(tf.keras.callbacks.BaseLogger())
+  last_epoch = 0
   if output_dir is not None:
     callbacks.append(tf.keras.callbacks.TensorBoard(log_dir=output_dir))
+    output_format = os.path.join(output_dir, "model-{epoch:05d}")
     callbacks.append(tf.keras.callbacks.ModelCheckpoint(
-        filepath=output_dir, save_weights_only=True))
+        filepath=output_format, save_weights_only=True))
+    checkpoints = tf.gfile.Glob(os.path.join(output_dir, "model-*"))
+    # Take basenames and strip the "model-" prefix.
+    checkpoints = [os.path.basename(ckpt)[6:] for ckpt in checkpoints]
+    # Get epoch numbers from the filenames and sort to obtain last epoch.
+    epoch_numbers = [int(ckpt[:5]) for ckpt in checkpoints if len(ckpt) > 4]
+    epoch_numbers.sort()
+    if epoch_numbers:
+      last_epoch = epoch_numbers[-1]
+      saved_path = os.path.join(output_dir, "model-%05d" % last_epoch)
+      model.load_weights(saved_path)
   model.fit(train_batches,
             epochs=train_steps // eval_frequency,
             steps_per_epoch=eval_frequency,
             validation_data=eval_batches,
             validation_steps=eval_steps,
+            initial_epoch=last_epoch,
             callbacks=callbacks)
 
 

From df6cf315200683b8e346e23652c0dd6df6949133 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 12 Feb 2019 11:57:23 -0800
Subject: [PATCH 1666/2720] Add unsophisticated but effective cleaning for
 ParaCrawl data.

PiperOrigin-RevId: 233651175
---
 .../data_generators/cleaner_en_xx.py          | 172 ++++++++++++++++++
 tensor2tensor/data_generators/translate.py    |  22 +--
 2 files changed, 178 insertions(+), 16 deletions(-)
 create mode 100644 tensor2tensor/data_generators/cleaner_en_xx.py

diff --git a/tensor2tensor/data_generators/cleaner_en_xx.py b/tensor2tensor/data_generators/cleaner_en_xx.py
new file mode 100644
index 000000000..669666d1a
--- /dev/null
+++ b/tensor2tensor/data_generators/cleaner_en_xx.py
@@ -0,0 +1,172 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+# encoding=UTF-8
+"""An unsophisticated data cleaner for en-.. sentence translation pairs.
+
+This pattern-based English-... cleaner aims fairly aggressively for clean
+sentence-like pairs. It discards pairs if the English member has signs of
+non-sentence noise or origin, e.g., lacks expected punctuation or has suspicious
+character sequences. It also simplistically detects and corrects some missing
+sentence breaks. It makes minimal assumptions about the other language, mainly
+that its sentences can end in one of '.!?' and that its sentences can start
+with an ASCII capital letter.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+
+import itertools
+import re
+
+import tensorflow as tf
+
+
+_RE_GOOD_S_START = re.compile(r'^["“”]?[A-Z]')
+_RE_GOOD_S_END = re.compile(r'\w[.?!]["”]?$', re.UNICODE)
+
+_RE_LABEL_COLON = re.compile(r'^\w+\.?( \w+)?: ', re.UNICODE)
+_RE_DIGIT_SPACE_DIGIT = re.compile(r'\d +\d', re.UNICODE)
+_RE_ALL_CAP_WORDS = re.compile(r'^[A-Z]\S*(\s+[A-Z]\S+)+\s*$')
+
+_RE_DQ_ONE = re.compile(r'^[^"“”]*["“”][^"“”]*$')
+_RE_DQ_INITIAL = re.compile(r'^["“”]([^"“”]+)$')
+_RE_DQ_FINAL = re.compile(r'^[^"“”]+["“”]$')
+_RE_DQ_LINE = re.compile(r'^["“”].*["“”]$')
+
+_RE_DQ_MANY = re.compile(r'(["“”].*){3,}')
+_RE_SQ_MANY = re.compile(r'''(['‘’][^st].*){3,}''')
+_RE_CHARS_QQ = re.compile(r'''["“”'‘’]\s*["“”'‘’]''')
+_RE_SPACE_PUNCT_SPACE = re.compile(r'''\s["“”'‘’,:;]\s''')
+
+_RE_COPYRIGHT = re.compile(r'©|^Copyright|^\(C\)')
+_RE_UNMATCHED_PAREN_LEFT = re.compile(r'[(][^)]*$')
+_RE_UNMATCHED_PAREN_RIGHT = re.compile(r'^[^(]*[)]')
+_RE_TAGLINE_CITY = re.compile(r'^[A-Z]{2,}(\s+[A-Z]+)*\s+-')
+_RE_CHARS_UPPER_UNDERSCORE = re.compile(r'^[A-Z]+[a-z]*_')
+
+
+def paracrawl_v3_pairs(paracrawl_file):
+  """Generates raw (English, other) pairs from a ParaCrawl V3.0 data file.
+
+  Args:
+    paracrawl_file: A ParaCrawl V3.0 en-.. data file.
+  Yields:
+    Pairs of (sentence_en, sentence_xx), as Unicode strings.
+  Raises:
+    StopIteration: If the file ends while this method is in the middle of
+        creating a translation pair.
+  """
+  raw_sentences = _raw_sentences(paracrawl_file)
+  for s_en in raw_sentences:
+    try:
+      s_xx = next(raw_sentences)
+      if s_en and s_xx:  # Prevent empty string examples.
+        yield s_en, s_xx
+    except StopIteration:
+      tf.logging.error(
+          'Unmatched final sentence while reading in sentence pairs: [%s]',
+          s_en)
+
+
+def _raw_sentences(paracrawl_file):
+  """Generates Unicode strings, one for each <seg> in a ParaCrawl data file.
+
+  Also decodes some of the most common HTML entities found in ParaCrawl data.
+
+  Args:
+    paracrawl_file: A ParaCrawl V3.0 en-.. data file.
+  Yields:
+    One Unicode string for each <seg> element in the ParaCrawl data file.
+  """
+  for line_utf8 in paracrawl_file:
+    line_uni = line_utf8.decode('UTF-8')
+    text_match = re.match(r' +<seg>(.*)</seg>$', line_uni)
+    if text_match:
+      txt = text_match.group(1)
+      txt = re.sub(r'&amp;', r'&', txt)
+      txt = re.sub(r'& ?amp;', r'&', txt)
+      txt = re.sub(r'& ?apos;', r"'", txt)
+      txt = re.sub(r'& ?quot;', r'"', txt)
+      txt = re.sub(r'& ?lt;', r'<', txt)
+      txt = re.sub(r'& ?gt;', r'>', txt)
+      yield txt
+
+
+def clean_en_xx_pairs(en_xx_pairs):
+  """Generates a cleaned-up stream of (English, other) translation pairs.
+
+  Cleaning includes both filtering and simplistic sentence splitting, with
+  minimal assumptions on the non-English pair member: (1) All filtering is
+  done based on the English member of the pair, and (2) sentence splitting
+  assumes only that sentences can end with one of '.!?' and begin with an
+  ASCII uppercase letter. Input pairs that would get split into different
+  numbers of sentences (e.g., three English sentences vs. two German ones) are
+  discarded.
+
+  Args:
+    en_xx_pairs: A stream (iterable) of Unicode string pairs. Each item in the
+        stream should be a (sentence_en, sentence_xx) pair.
+  Yields:
+    Cleaned-up (sentence_en, sentence_xx) pairs.
+  """
+  for s1, s2 in en_xx_pairs:
+    if _regex_filter(s1):
+      continue
+    s1_list, s2_list = _split_sentences(s1, s2)
+    if len(s1_list) != len(s2_list):
+      pass  # discard this pair
+    elif len(s1_list) == 1:
+      yield s1, s2
+    else:
+      for s1_subsentence, s2_subsentence in itertools.izip(s1_list, s2_list):
+        if _regex_filter(s1_subsentence):
+          continue
+        yield s1_subsentence, s2_subsentence
+
+
+def _regex_filter(sentence):
+  return (not _is_match(sentence, _RE_GOOD_S_START)
+          or not _is_match(sentence, _RE_GOOD_S_END)
+          or _is_match(sentence, _RE_LABEL_COLON)
+          or _is_match(sentence, _RE_DIGIT_SPACE_DIGIT)
+          or _is_match(sentence, _RE_DQ_ONE)
+          or _is_match(sentence, _RE_DQ_INITIAL)
+          or _is_match(sentence, _RE_DQ_FINAL)
+          or _is_match(sentence, _RE_DQ_LINE)
+          or _is_match(sentence, _RE_DQ_MANY)
+          or _is_match(sentence, _RE_SQ_MANY)
+          or _is_match(sentence, _RE_CHARS_QQ)
+          or _is_match(sentence, _RE_SPACE_PUNCT_SPACE)
+          or _is_match(sentence, _RE_COPYRIGHT)
+          or _is_match(sentence, _RE_UNMATCHED_PAREN_LEFT)
+          or _is_match(sentence, _RE_UNMATCHED_PAREN_RIGHT)
+          or _is_match(sentence, _RE_TAGLINE_CITY)
+          or _is_match(sentence, _RE_CHARS_UPPER_UNDERSCORE))
+
+
+def _is_match(sentence, regex):
+  return regex.search(sentence)
+
+
+def _split_sentences(s1, s2):
+  s1 = re.sub(r'(\w[A-Z]|[0-9a-z])([.!?]) ([A-Z])', r'\1\2__|__\3', s1)
+  s2 = re.sub(r'([^0-9][.!?]) ([A-Z])', r'\1__|__\2', s2)
+  s1_subsentences = s1.split('__|__')
+  s2_subsentences = s2.split('__|__')
+  return s1_subsentences, s2_subsentences
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index dceb9d96b..ce55cf2df 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -21,6 +21,7 @@
 
 import os
 import tarfile
+from tensor2tensor.data_generators import cleaner_en_xx
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -153,22 +154,11 @@ def compile_data(tmp_dir, datasets, filename):
             tmx_filename = new_filename
           source, target = None, None
           with tf.gfile.Open(tmx_filename) as tmx_file:
-            for line in tmx_file:
-              text = line.strip()
-              if text.startswith("<seg>"):
-                if text.endswith("</seg>"):
-                  sentence = text[5:-6]  # Strip <seg> and </seg>.
-                  if source is None:
-                    source = sentence
-                  else:
-                    target = sentence
-              if source is not None and target is not None:
-                if source and target:  # Prevent empty string examples.
-                  lang1_resfile.write(source)
-                  lang1_resfile.write("\n")
-                  lang2_resfile.write(target)
-                  lang2_resfile.write("\n")
-                source, target = None, None
+            for source, target in cleaner_en_xx.paracrawl_v3_pairs(tmx_file):
+              lang1_resfile.write(source)
+              lang1_resfile.write("\n")
+              lang2_resfile.write(target)
+              lang2_resfile.write("\n")
 
         elif dataset[1][0] == "tsv":
           _, src_column, trg_column, glob_pattern = dataset[1]

From ba763e558534331c93355234441820c71b698d1d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 12 Feb 2019 13:05:04 -0800
Subject: [PATCH 1667/2720] Changes for using kfac in the transformer model in
 T2T.

PiperOrigin-RevId: 233662431
---
 setup.py                                      |  1 +
 tensor2tensor/layers/common_attention.py      | 44 +++++++++--
 tensor2tensor/layers/common_attention_test.py | 10 +++
 tensor2tensor/layers/common_layers.py         | 74 +++++++++++++++++--
 tensor2tensor/layers/common_layers_test.py    | 29 ++++++++
 tensor2tensor/layers/transformer_layers.py    |  8 +-
 tensor2tensor/models/transformer.py           | 29 +++++---
 7 files changed, 167 insertions(+), 28 deletions(-)

diff --git a/setup.py b/setup.py
index 1e83e6bd5..e6b66dcac 100644
--- a/setup.py
+++ b/setup.py
@@ -42,6 +42,7 @@
         'gunicorn',
         'gym',
         'h5py',
+        'kfac',
         'mesh-tensorflow',
         'numpy',
         'oauth2client',
diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 3f5dcafbe..b0c255aff 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -3396,7 +3396,8 @@ def compute_attention_component(antecedent,
                                 filter_width=1,
                                 padding="VALID",
                                 name="c",
-                                vars_3d_num_heads=0):
+                                vars_3d_num_heads=0,
+                                layer_collection=None):
   """Computes attention compoenent (query, key or value).
 
   Args:
@@ -3407,10 +3408,18 @@ def compute_attention_component(antecedent,
     padding: One of "VALID", "SAME" or "LEFT". Default is VALID: No padding.
     name: a string specifying scope name.
     vars_3d_num_heads: an optional integer (if we want to use 3d variables)
+    layer_collection: A tensorflow_kfac.LayerCollection. Only used by the
+      KFAC optimizer. Default is None.
 
   Returns:
     c : [batch, length, depth] tensor
   """
+  if layer_collection is not None:
+    if filter_width != 1 or vars_3d_num_heads != 0:
+      raise ValueError(
+          "KFAC implementation only supports filter_width=1 (actual: {}) and "
+          "vars_3d_num_heads=0 (actual: {}).".format(
+              filter_width, vars_3d_num_heads))
   if vars_3d_num_heads > 0:
     assert filter_width == 1
     input_depth = antecedent.get_shape().as_list()[-1]
@@ -3428,7 +3437,8 @@ def compute_attention_component(antecedent,
     return tf.tensordot(antecedent, var, axes=1)
   if filter_width == 1:
     return common_layers.dense(
-        antecedent, total_depth, use_bias=False, name=name)
+        antecedent, total_depth, use_bias=False, name=name,
+        layer_collection=layer_collection)
   else:
     return common_layers.conv1d(
         antecedent, total_depth, filter_width, padding=padding, name=name)
@@ -3442,7 +3452,8 @@ def compute_qkv(query_antecedent,
                 kv_filter_width=1,
                 q_padding="VALID",
                 kv_padding="VALID",
-                vars_3d_num_heads=0):
+                vars_3d_num_heads=0,
+                layer_collection=None):
   """Computes query, key and value.
 
   Args:
@@ -3456,6 +3467,8 @@ def compute_qkv(query_antecedent,
     q_padding: One of "VALID", "SAME" or "LEFT". Default is VALID: No padding.
     kv_padding: One of "VALID", "SAME" or "LEFT". Default is VALID: No padding.
     vars_3d_num_heads: an optional (if we want to use 3d variables)
+    layer_collection: A tensorflow_kfac.LayerCollection. Only used by the
+      KFAC optimizer. Default is None.
 
   Returns:
     q, k, v : [batch, length, depth] tensors
@@ -3468,21 +3481,24 @@ def compute_qkv(query_antecedent,
       q_filter_width,
       q_padding,
       "q",
-      vars_3d_num_heads=vars_3d_num_heads)
+      vars_3d_num_heads=vars_3d_num_heads,
+      layer_collection=layer_collection)
   k = compute_attention_component(
       memory_antecedent,
       total_key_depth,
       kv_filter_width,
       kv_padding,
       "k",
-      vars_3d_num_heads=vars_3d_num_heads)
+      vars_3d_num_heads=vars_3d_num_heads,
+      layer_collection=layer_collection)
   v = compute_attention_component(
       memory_antecedent,
       total_value_depth,
       kv_filter_width,
       kv_padding,
       "v",
-      vars_3d_num_heads=vars_3d_num_heads)
+      vars_3d_num_heads=vars_3d_num_heads,
+      layer_collection=layer_collection)
   return q, k, v
 
 
@@ -3513,6 +3529,7 @@ def multihead_attention(query_antecedent,
                         make_image_summary=True,
                         dropout_broadcast_dims=None,
                         vars_3d=False,
+                        layer_collection=None,
                         **kwargs):
   """Multihead scaled-dot-product attention with input/output transformations.
 
@@ -3564,6 +3581,8 @@ def multihead_attention(query_antecedent,
       specifying in which dimensions to broadcast the dropout decisions.
       saves memory.
     vars_3d: use 3-dimensional variables for input/output transformations
+    layer_collection: A tensorflow_kfac.LayerCollection. Only used by the
+      KFAC optimizer. Default is None.
     **kwargs (dict): Parameters for the attention function
 
   Caching:
@@ -3594,6 +3613,13 @@ def multihead_attention(query_antecedent,
     raise ValueError("Value depth (%d) must be divisible by the number of "
                      "attention heads (%d)." % (total_value_depth, num_heads))
   vars_3d_num_heads = num_heads if vars_3d else 0
+
+  if layer_collection is not None:
+    if cache is not None:
+      raise ValueError("KFAC implementation only supports cache is None.")
+    if vars_3d:
+      raise ValueError("KFAC implementation does not support 3d vars.")
+
   with tf.variable_scope(name, default_name="multihead_attention",
                          values=[query_antecedent, memory_antecedent]):
 
@@ -3601,7 +3627,8 @@ def multihead_attention(query_antecedent,
       q, k, v = compute_qkv(query_antecedent, memory_antecedent,
                             total_key_depth, total_value_depth, q_filter_width,
                             kv_filter_width, q_padding, kv_padding,
-                            vars_3d_num_heads=vars_3d_num_heads)
+                            vars_3d_num_heads=vars_3d_num_heads,
+                            layer_collection=layer_collection)
     if cache is not None:
       if attention_type not in ["dot_product", "dot_product_relative"]:
         # TODO(petershaw): Support caching when using relative position
@@ -3744,7 +3771,8 @@ def multihead_attention(query_antecedent,
       x = tf.tensordot(x, o_var, axes=1)
     else:
       x = common_layers.dense(
-          x, output_depth, use_bias=False, name="output_transform")
+          x, output_depth, use_bias=False, name="output_transform",
+          layer_collection=layer_collection)
     if additional_returned_value is not None:
       return x, additional_returned_value
     return x
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 8d1575df3..ec1b6f2e9 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 from absl.testing import parameterized
+import kfac
 import numpy as np
 
 from tensor2tensor.layers import common_attention
@@ -841,6 +842,15 @@ def testBiasFuture(self):
     bias = common_attention.attention_bias_future(q, k)
     self.assertAllClose(self.evaluate(bias), ground_truth)
 
+  @test_utils.run_in_graph_mode_only()
+  def testMultiheadAttentionWithLayerCollection(self):
+    """Testing multihead attention with layer collection for kfac."""
+    x = tf.zeros([3, 4, 5], tf.float32)
+    layer_collection = kfac.LayerCollection()
+    common_attention.multihead_attention(
+        x, None, None, 10, 10, 10, 2, 0.2,
+        layer_collection=layer_collection)
+    self.assertLen(layer_collection.get_blocks(), 4)
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 9072b2820..6b23d7afa 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -640,13 +640,22 @@ def layer_norm_compute(x, epsilon, scale, bias):
   return norm_x * scale + bias
 
 
-def layer_norm(x, filters=None, epsilon=1e-6, name=None, reuse=None):
+def layer_norm(x,
+               filters=None,
+               epsilon=1e-6,
+               name=None,
+               reuse=None,
+               layer_collection=None):
   """Layer normalize the tensor x, averaging over the last dimension."""
   if filters is None:
     filters = shape_list(x)[-1]
   with tf.variable_scope(
       name, default_name="layer_norm", values=[x], reuse=reuse):
     scale, bias = layer_norm_vars(filters)
+    if layer_collection:
+      tf.logging.info("Registering layer norm to collection with (scale, bias):"
+                      " ({}, {})".format(scale, bias))
+      layer_collection.register_generic((scale, bias), shape_list(x)[0])
     return layer_norm_compute(x, epsilon, scale, bias)
 
 
@@ -737,10 +746,13 @@ def apply_spectral_norm(x):
   return tf.divide(x, spectral_norm), assign_op
 
 
-def apply_norm(x, norm_type, depth, epsilon):
+def apply_norm(x, norm_type, depth, epsilon, layer_collection=None):
   """Apply Normalization."""
+  if layer_collection is not None:
+    assert norm_type == "layer"
   if norm_type == "layer":
-    return layer_norm(x, filters=depth, epsilon=epsilon)
+    return layer_norm(
+        x, filters=depth, epsilon=epsilon, layer_collection=layer_collection)
   if norm_type == "group":
     return group_norm(x, filters=depth, epsilon=epsilon)
   if norm_type == "batch":
@@ -786,7 +798,8 @@ def layer_prepostprocess(previous_value,
                          epsilon,
                          default_name,
                          name=None,
-                         dropout_broadcast_dims=None):
+                         dropout_broadcast_dims=None,
+                         layer_collection=None):
   """Apply a sequence of functions to the input or output of a layer.
 
   The sequence is specified as a string which may contain the following
@@ -812,6 +825,8 @@ def layer_prepostprocess(previous_value,
     dropout_broadcast_dims:  an optional list of integers less than 3
       specifying in which dimensions to broadcast the dropout decisions.
       saves memory.
+    layer_collection: A tensorflow_kfac.LayerCollection. Only used by the
+      KFAC optimizer. Default is None.
 
   Returns:
     a Tensor
@@ -825,7 +840,8 @@ def layer_prepostprocess(previous_value,
       elif c == "z":
         x = zero_add(previous_value, x)
       elif c == "n":
-        x = apply_norm(x, norm_type, depth, epsilon)
+        x = apply_norm(
+            x, norm_type, depth, epsilon, layer_collection=layer_collection)
       else:
         assert c == "d", ("Unknown sequence step %s" % c)
         x = dropout_with_broadcast_dims(
@@ -833,7 +849,7 @@ def layer_prepostprocess(previous_value,
     return x
 
 
-def layer_preprocess(layer_input, hparams):
+def layer_preprocess(layer_input, hparams, layer_collection=None):
   """Apply layer preprocessing.
 
   See layer_prepostprocess() for details.
@@ -850,6 +866,8 @@ def layer_preprocess(layer_input, hparams):
   Args:
     layer_input: a Tensor
     hparams: a hyperparameters object.
+    layer_collection: A tensorflow_kfac.LayerCollection. Only used by the
+      KFAC optimizer. Default is None.
 
   Returns:
     a Tensor
@@ -868,7 +886,8 @@ def layer_preprocess(layer_input, hparams):
       epsilon=hparams.norm_epsilon,
       dropout_broadcast_dims=comma_separated_string_to_integer_list(
           getattr(hparams, "layer_prepostprocess_dropout_broadcast_dims", "")),
-      default_name="layer_prepostprocess")
+      default_name="layer_prepostprocess",
+      layer_collection=layer_collection)
 
 
 def layer_postprocess(layer_input, layer_output, hparams):
@@ -1222,6 +1241,7 @@ def dense_relu_dense(inputs,
                      output_activation=None,
                      dropout=0.0,
                      dropout_broadcast_dims=None,
+                     layer_collection=None,
                      name=None):
   """Hidden layer with RELU activation followed by linear projection."""
   # layer_name is appended with "conv1" or "conv2" in this method only for
@@ -1232,6 +1252,7 @@ def dense_relu_dense(inputs,
       filter_size,
       use_bias=True,
       activation=tf.nn.relu,
+      layer_collection=layer_collection,
       name=layer_name.format("conv1"))
 
   if dropout != 0.0:
@@ -1242,6 +1263,7 @@ def dense_relu_dense(inputs,
       output_size,
       activation=output_activation,
       use_bias=True,
+      layer_collection=layer_collection,
       name=layer_name.format("conv2"))
   return o
 
@@ -2935,7 +2957,43 @@ def fn_with_recompute(*args):
 
 def dense(x, units, **kwargs):
   """Identical to tf.layers.dense."""
-  return tf.layers.dense(x, units, **kwargs)
+  layer_collection = kwargs.pop("layer_collection", None)
+  activations = tf.layers.dense(x, units, **kwargs)
+  if layer_collection:
+    # We need to find the layer parameters using scope name for the layer, so
+    # check that the layer is named. Otherwise parameters for different layers
+    # may get mixed up.
+    layer_name = tf.get_variable_scope().name
+    if (not layer_name) or ("name" not in kwargs):
+      raise ValueError(
+          "Variable scope and layer name cannot be empty. Actual: "
+          "variable_scope={}, layer name={}".format(
+              layer_name, kwargs.get("name", None)))
+
+    layer_name += "/" + kwargs["name"]
+    layer_params = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
+                                     scope=layer_name)
+    assert layer_params
+    if len(layer_params) == 1:
+      layer_params = layer_params[0]
+
+    tf.logging.info(
+        "Registering dense layer to collection for tensor: {}".format(
+            layer_params))
+
+    x_shape = x.shape.as_list()
+    if len(x_shape) == 3:
+      # Handle [batch, time, depth] inputs by folding batch and time into
+      # one dimension: reshaping inputs to [batchxtime, depth].
+      x_2d = tf.reshape(x, [-1, x_shape[2]])
+      activations_shape = activations.shape.as_list()
+      activations_2d = tf.reshape(activations, [-1, activations_shape[2]])
+      layer_collection.register_fully_connected_multi(
+          layer_params, x_2d, activations_2d, num_uses=x_shape[1])
+      activations = tf.reshape(activations_2d, activations_shape)
+    else:
+      layer_collection.register_fully_connected(layer_params, x, activations)
+  return activations
 
 
 def batch_dense(inputs,
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index b3cdf6be2..a11aff86e 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 from absl.testing import parameterized
+import kfac
 import numpy as np
 
 from tensor2tensor.layers import common_layers
@@ -213,6 +214,11 @@ def testLayerNorm(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 11))
 
+    # Testing layer collection.
+    layer_collection = kfac.LayerCollection()
+    common_layers.layer_norm(x, layer_collection=layer_collection)
+    self.assertLen(layer_collection.get_blocks(), 1)
+
   @test_utils.run_in_graph_and_eager_modes()
   def testGroupNorm(self):
     x = np.random.rand(5, 7, 3, 16)
@@ -311,6 +317,29 @@ def testApplyNormNone(self):
     self.assertEqual(actual.shape, (5, 2, 1, 11))
     self.assertAllClose(actual, x1, atol=1e-03)
 
+  @test_utils.run_in_graph_mode_only()
+  def testApplyNormWithLayerCollection(self):
+    x = np.random.rand(5, 2, 1, 11)
+    layer_collection = kfac.LayerCollection()
+    common_layers.apply_norm(x, "layer", depth=11, epsilon=1e-6,
+                             layer_collection=layer_collection)
+    self.assertLen(layer_collection.get_blocks(), 1)
+
+  @test_utils.run_in_graph_mode_only()
+  def testDenseWithLayerCollection(self):
+    with tf.variable_scope("test_layer_collection"):
+      x1 = tf.zeros([3, 4], tf.float32)
+      layer_collection = kfac.LayerCollection()
+      common_layers.dense(
+          x1, units=10, layer_collection=layer_collection, name="y1")
+      self.assertLen(layer_collection.get_blocks(), 1)
+
+      # 3D inputs.
+      x2 = tf.zeros([3, 4, 5], tf.float32)
+      common_layers.dense(
+          x2, units=10, layer_collection=layer_collection, name="y2")
+      self.assertLen(layer_collection.get_blocks(), 2)
+
   def testGlobalPool1d(self):
     x1 = np.random.rand(5, 4, 11)
     no_mask = np.ones((5, 4))
diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 8288e9761..636f6e4a8 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -390,7 +390,8 @@ def transformer_ffn_layer(x,
                           losses=None,
                           cache=None,
                           decode_loop_step=None,
-                          readout_filter_size=0):
+                          readout_filter_size=0,
+                          layer_collection=None):
   """Feed-forward layer in the transformer.
 
   Args:
@@ -411,6 +412,8 @@ def transformer_ffn_layer(x,
         Only used for inference on TPU.
     readout_filter_size: if it's greater than 0, then it will be used instead of
       filter_size
+    layer_collection: A tensorflow_kfac.LayerCollection. Only used by the
+      KFAC optimizer. Default is None.
 
 
   Returns:
@@ -453,7 +456,8 @@ def transformer_ffn_layer(x,
         hparams.filter_size,
         hparams.hidden_size,
         dropout=hparams.relu_dropout,
-        dropout_broadcast_dims=relu_dropout_broadcast_dims)
+        dropout_broadcast_dims=relu_dropout_broadcast_dims,
+        layer_collection=layer_collection)
     if pad_remover:
       # Restore `conv_output` to the original shape of `x`, including padding.
       conv_output = tf.reshape(
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 479dfd45e..ddd17294d 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1286,7 +1286,8 @@ def transformer_decoder(decoder_input,
                         nonpadding=None,
                         save_weights_to=None,
                         make_image_summary=True,
-                        losses=None):
+                        losses=None,
+                        layer_collection=None):
   """A stack of transformer layers.
 
   Args:
@@ -1312,6 +1313,8 @@ def transformer_decoder(decoder_input,
       key created from the variable scope (including name).
     make_image_summary: Whether to make an attention image summary.
     losses: optional list onto which to append extra training losses
+    layer_collection: A tensorflow_kfac.LayerCollection. Only used by the
+      KFAC optimizer. Default is None.
 
   Returns:
     y: a Tensors
@@ -1345,7 +1348,8 @@ def transformer_decoder(decoder_input,
       with tf.variable_scope(layer_name):
         with tf.variable_scope("self_attention"):
           y = common_attention.multihead_attention(
-              common_layers.layer_preprocess(x, hparams),
+              common_layers.layer_preprocess(
+                  x, hparams, layer_collection=layer_collection),
               None,
               decoder_self_attention_bias,
               hparams.attention_key_channels or hparams.hidden_size,
@@ -1366,12 +1370,14 @@ def transformer_decoder(decoder_input,
               decode_loop_step=decode_loop_step,
               vars_3d=hparams.get("attention_variables_3d"),
               activation_dtype=hparams.get("activation_dtype", "float32"),
-              weight_dtype=hparams.get("weight_dtype", "float32"))
+              weight_dtype=hparams.get("weight_dtype", "float32"),
+              layer_collection=layer_collection)
           x = common_layers.layer_postprocess(x, y, hparams)
         if encoder_output is not None:
           with tf.variable_scope("encdec_attention"):
             y = common_attention.multihead_attention(
-                common_layers.layer_preprocess(x, hparams),
+                common_layers.layer_preprocess(
+                    x, hparams, layer_collection=layer_collection),
                 encoder_output,
                 encoder_decoder_attention_bias,
                 hparams.attention_key_channels or hparams.hidden_size,
@@ -1390,26 +1396,29 @@ def transformer_decoder(decoder_input,
                 max_length=hparams.get("max_length"),
                 vars_3d=hparams.get("attention_variables_3d"),
                 activation_dtype=hparams.get("activation_dtype", "float32"),
-                weight_dtype=hparams.get("weight_dtype", "float32"))
+                weight_dtype=hparams.get("weight_dtype", "float32"),
+                layer_collection=layer_collection)
             x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(
-              common_layers.layer_preprocess(x, hparams),
+              common_layers.layer_preprocess(
+                  x, hparams, layer_collection=layer_collection),
               hparams,
               conv_padding="LEFT",
               nonpadding_mask=nonpadding,
               losses=losses,
               cache=layer_cache,
-              decode_loop_step=decode_loop_step)
+              decode_loop_step=decode_loop_step,
+              layer_collection=layer_collection)
           x = common_layers.layer_postprocess(x, y, hparams)
     # if normalization is done in layer_preprocess, then it should also be done
     # on the output, since the output can grow very large, being the sum of
     # a whole stack of unnormalized layer outputs.
     mlperf_log.transformer_print(
         key=mlperf_log.MODEL_HP_NORM,
-        value={"hidden_size": hparams.hidden_size},
-        hparams=hparams)
-    return common_layers.layer_preprocess(x, hparams)
+        value={"hidden_size": hparams.hidden_size})
+    return common_layers.layer_preprocess(
+        x, hparams, layer_collection=layer_collection)
 
 
 def evolved_transformer_decoder(decoder_input,

From 0ae37698435720b2109bb3ca320c4c49bc7986c2 Mon Sep 17 00:00:00 2001
From: Mike Dusenberry <dusenberrymw@google.com>
Date: Tue, 12 Feb 2019 13:33:11 -0800
Subject: [PATCH 1668/2720] Revert to the full decorator approach for the
 add_weight method.

For classes that subclass our Bayesian layers without overriding `add_weight`,
`return super(layer.__class__, layer).add_weight(...)` ends up recursively
calling the same function. This reverts the manual decorator equivalent back to
a full decorator.

PiperOrigin-RevId: 233664423
---
 tensor2tensor/layers/bayes.py      | 75 ++++++++++++++++--------------
 tensor2tensor/layers/bayes_test.py | 20 ++++++++
 2 files changed, 59 insertions(+), 36 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 581438625..18c799ff1 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
 import math
 
 import six
@@ -375,36 +376,42 @@ def get_config(self):
     }
 
 
-def _add_weight(layer,
-                name=None,
-                shape=None,
-                dtype=None,
-                initializer=None,
-                regularizer=None,
-                **kwargs):
-  """Adds weight."""
-  if isinstance(initializer, tf.keras.layers.Layer):
-    weight = initializer(shape, dtype)
-    layer._trainable_weights.extend(initializer.trainable_weights)  # pylint: disable=protected-access
-    layer._non_trainable_weights.extend(initializer.non_trainable_weights)  # pylint: disable=protected-access
-    if regularizer is not None:
-      # TODO(trandustin): Replace need for this with
-      # Layer._handle_weight_regularization. For Eager compatibility, random
-      # variable __init__s cannot apply TF ops (cl/220898007).
-      def loss_fn():
-        """Creates a regularization loss `Tensor`."""
-        with tf.name_scope(name + '/Regularizer'):
-          return regularizer(initializer(shape, dtype))
-      layer.add_loss(loss_fn)
-    return weight
-  return super(layer.__class__, layer).add_weight(name=name,
-                                                  shape=shape,
-                                                  dtype=dtype,
-                                                  initializer=initializer,
-                                                  regularizer=regularizer,
-                                                  **kwargs)
-
-
+def add_weight(cls):
+  """Decorator for Layers, overriding add_weight for trainable initializers."""
+  @functools.wraps(cls.add_weight)
+  def _add_weight(self,
+                  name=None,
+                  shape=None,
+                  dtype=None,
+                  initializer=None,
+                  regularizer=None,
+                  **kwargs):
+    """Adds weight."""
+    if isinstance(initializer, tf.keras.layers.Layer):
+      weight = initializer(shape, dtype)
+      self._trainable_weights.extend(initializer.trainable_weights)  # pylint: disable=protected-access
+      self._non_trainable_weights.extend(initializer.non_trainable_weights)  # pylint: disable=protected-access
+      if regularizer is not None:
+        # TODO(trandustin): Replace need for this with
+        # Layer._handle_weight_regularization. For Eager compatibility, random
+        # variable __init__s cannot apply TF ops (cl/220898007).
+        def loss_fn():
+          """Creates a regularization loss `Tensor`."""
+          with tf.name_scope(name + '/Regularizer'):
+            return regularizer(initializer(shape, dtype))
+        self.add_loss(loss_fn)
+      return weight
+    return super(cls, self).add_weight(name=name,
+                                       shape=shape,
+                                       dtype=dtype,
+                                       initializer=initializer,
+                                       regularizer=regularizer,
+                                       **kwargs)
+  cls.add_weight = _add_weight
+  return cls
+
+
+@add_weight
 class DenseReparameterization(tf.keras.layers.Dense):
   """Bayesian densely-connected layer estimated via reparameterization.
 
@@ -445,8 +452,6 @@ def __init__(self,
         activity_regularizer=get(activity_regularizer),
         **kwargs)
 
-  add_weight = _add_weight
-
   # TODO(trandustin): This name is not accurate. Rename or move functionality
   # into random variables to resample/recreate their init ops.
   def sample_weights(self):
@@ -460,6 +465,7 @@ def call(self, *args, **kwargs):
     return super(DenseReparameterization, self).call(*args, **kwargs)
 
 
+@add_weight
 class Conv2DReparameterization(tf.keras.layers.Conv2D):
   """2D convolution layer (e.g. spatial convolution over images).
 
@@ -514,8 +520,6 @@ def __init__(self,
         bias_constraint=get(bias_constraint),
         **kwargs)
 
-  add_weight = _add_weight
-
   def sample_weights(self):
     if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
       self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
@@ -680,6 +684,7 @@ def get_config(self):
     return dict(list(base_config.items()) + list(config.items()))
 
 
+@add_weight
 class LSTMCellReparameterization(tf.keras.layers.LSTMCell):
   """Bayesian LSTM cell class estimated via reparameterization.
 
@@ -740,8 +745,6 @@ def __init__(self,
         implementation=implementation,
         **kwargs)
 
-  add_weight = _add_weight
-
   def build(self, input_shape):
     input_shape = tf.TensorShape(input_shape)
     input_dim = input_shape[-1]
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 4563cfa33..bcde4ecfa 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -138,6 +138,26 @@ def testDenseReparameterizationModel(self):
     self.assertEqual(res.shape, (3, 2))
     self.assertLen(model.losses, 1)
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testDenseReparameterizationSubclass(self):
+    class DenseReparameterizationSubclass(bayes.DenseReparameterization):
+      pass
+
+    inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
+    model = tf.keras.Sequential([
+        tf.keras.layers.Conv2D(3,
+                               kernel_size=2,
+                               padding="SAME",
+                               activation=tf.nn.relu),
+        tf.keras.layers.Flatten(),
+        DenseReparameterizationSubclass(2, activation=None),
+    ])
+    outputs = model(inputs)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(outputs)
+    self.assertEqual(res.shape, (3, 2))
+    self.assertLen(model.losses, 1)
+
   @parameterized.named_parameters(
       {"testcase_name": "_no_uncertainty",
        "kernel_initializer": "zeros",

From 0b156ac533ab53f65f44966381f6e147c7371eee Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Tue, 12 Feb 2019 21:07:35 -0800
Subject: [PATCH 1669/2720] Fixes to mixture-of-experts code for
 mtf-transformer2 - feed in a tensor that indicates nonpadding positions.

PiperOrigin-RevId: 233684322
---
 .../models/research/moe_experiments.py        | 75 +++++++++++++++++++
 1 file changed, 75 insertions(+)

diff --git a/tensor2tensor/models/research/moe_experiments.py b/tensor2tensor/models/research/moe_experiments.py
index ee21605f3..a817a6314 100644
--- a/tensor2tensor/models/research/moe_experiments.py
+++ b/tensor2tensor/models/research/moe_experiments.py
@@ -21,10 +21,85 @@
 from __future__ import print_function
 
 from tensor2tensor.models import mtf_transformer
+from tensor2tensor.models import mtf_transformer2
 from tensor2tensor.models.research import moe
 from tensor2tensor.utils import registry
 
 
+@registry.register_hparams
+def xmoe_tr_dense_2k():
+  """Series of architectural experiments on Translation.
+
+  # run on 8-core setup
+
+  119M params, einsum=0.95e13
+
+  Returns:
+    a hparams
+  """
+  hparams = mtf_transformer2.mtf_bitransformer_base()
+  hparams.encoder_layers = ["self_att", "drd"] * 4
+  hparams.decoder_layers = ["self_att", "enc_att", "drd"] * 4
+  hparams.batch_size = 64
+  hparams.shared_embedding_and_softmax_weights = True
+  hparams.mesh_shape = "batch:8"
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_tr_dense_32k():
+  """Bigger d_ff.
+
+  623M params, einsum=3.42e13
+
+  Returns:
+    a hparams
+  """
+  hparams = xmoe_tr_dense_2k()
+  hparams.d_ff = 32768
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_tr_1d():
+  """Mixture of experts (16 experts).
+
+
+  623M Params, einsum=1.09e13
+
+  Returns:
+    a hparams
+  """
+  hparams = xmoe_tr_dense_2k()
+  hparams.encoder_layers = ["self_att", "moe_1d"] * 4
+  hparams.decoder_layers = ["self_att", "enc_att", "moe_1d"] * 4
+  hparams.layout = "batch:batch;experts:batch"
+  hparams.moe_hidden_size = 2048
+  hparams.moe_num_experts = 16
+  return hparams
+
+
+@registry.register_hparams
+def xmoe_tr_2d():
+  """Mixture of experts (16 experts).
+
+  623M Params, einsum=1.09e13
+
+  Returns:
+    a hparams
+  """
+  hparams = xmoe_tr_dense_2k()
+  hparams.mesh_shape = "b0:2;b1:4"
+  hparams.outer_batch_size = 4
+  hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
+  hparams.encoder_layers = ["self_att", "moe_2d"] * 4
+  hparams.decoder_layers = ["self_att", "enc_att", "moe_2d"] * 4
+  hparams.moe_hidden_size = 2048
+  hparams.moe_experts_x = 4
+  hparams.moe_experts_y = 4
+  return hparams
+
+
 @registry.register_hparams
 def xmoe_dense_4k():
   """Series of architectural experiments on cheap language models.

From 2d10f58303634f5f77c73c10fb8b17c0e84ce2f0 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 13 Feb 2019 17:48:01 -0800
Subject: [PATCH 1670/2720] Modify `gym_utils.register_gym_env` to be able to
 pass arguments to the class.

PiperOrigin-RevId: 233866405
---
 tensor2tensor/rl/gym_utils.py      |  4 +--
 tensor2tensor/rl/gym_utils_test.py | 57 ++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index 72bff9f2f..cf8c0675e 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -103,7 +103,7 @@ def make_gym_env(name, rl_env_max_episode_steps=-1, maxskip_env=False):
   return gym.wrappers.TimeLimit(env, max_episode_steps=rl_env_max_episode_steps)
 
 
-def register_gym_env(class_entry_point, version="v0"):
+def register_gym_env(class_entry_point, version="v0", kwargs=None):
   """Registers the class in Gym and returns the registered name and the env."""
 
   split_on_colon = class_entry_point.split(":")
@@ -112,7 +112,7 @@ def register_gym_env(class_entry_point, version="v0"):
   class_name = split_on_colon[1]
   # We have to add the version to conform to gym's API.
   env_name = "T2TEnv-{}-{}".format(class_name, version)
-  gym.envs.register(id=env_name, entry_point=class_entry_point)
+  gym.envs.register(id=env_name, entry_point=class_entry_point, kwargs=kwargs)
 
   tf.logging.info("Entry Point [%s] registered with id [%s]",
                   class_entry_point, env_name)
diff --git a/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index ce4f86e5c..95d92faf7 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -44,6 +44,20 @@ def step(self, action):
       return self.observation_space.high, +1.0, True, {}
 
 
+class EnvWithOptions(SimpleEnv):
+  """A simple env that takes arguments on init."""
+
+  def __init__(self, done_action=0):
+    super(EnvWithOptions, self).__init__()
+    self.action_space = spaces.Discrete(3)
+    self._done_action = done_action
+
+  def step(self, action):
+    if action == self._done_action:
+      return self.observation_space.high, +1.0, True, {}
+    return self.reset(), -1.0, False, {}
+
+
 class GymUtilsTest(tf.test.TestCase):
 
   # Just make an environment and expect to get one.
@@ -81,6 +95,49 @@ def test_gym_registration(self):
     _, _, done, _ = env.step(1)
     self.assertTrue(done)
 
+  def test_gym_registration_with_kwargs(self):
+    reg_id, env = gym_utils.register_gym_env(
+        "tensor2tensor.rl.gym_utils_test:EnvWithOptions",
+        kwargs={"done_action": 2}
+    )
+
+    self.assertEqual("T2TEnv-EnvWithOptions-v0", reg_id)
+
+    # Obligatory reset.
+    env.reset()
+
+    # Make sure that on action = 0, 1 we are not done, but on 2 we are.
+    _, _, done, _ = env.step(0)
+    self.assertFalse(done)
+
+    _, _, done, _ = env.step(1)
+    self.assertFalse(done)
+
+    _, _, done, _ = env.step(2)
+    self.assertTrue(done)
+
+    # Now lets try to change the env -- note we have to change the version.
+    reg_id, env = gym_utils.register_gym_env(
+        "tensor2tensor.rl.gym_utils_test:EnvWithOptions",
+        version="v1",
+        kwargs={"done_action": 1}
+    )
+
+    self.assertEqual("T2TEnv-EnvWithOptions-v1", reg_id)
+
+    # Obligatory reset.
+    env.reset()
+
+    # Make sure that on action = 0, 2 we are not done, but on 1 we are.
+    _, _, done, _ = env.step(0)
+    self.assertFalse(done)
+
+    _, _, done, _ = env.step(2)
+    self.assertFalse(done)
+
+    _, _, done, _ = env.step(1)
+    self.assertTrue(done)
+
 
 if __name__ == "__main__":
   tf.test.main()

From 6375c91587b300e459947ef5818298114af22379 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 14 Feb 2019 12:42:18 -0800
Subject: [PATCH 1671/2720] Enable beam search for examples with "knowledge"
 feature.

PiperOrigin-RevId: 234009368
---
 tensor2tensor/utils/t2t_model.py | 37 +++++++++++++++++++++-----------
 1 file changed, 24 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 957829d9e..5721d7417 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -872,18 +872,29 @@ def symbols_to_logits_fn(ids, i=None):
       logits = logits[:, current_output_position, :, :]
       return tf.squeeze(logits, axis=[1, 2])
 
-    initial_ids = tf.zeros([batch_size], dtype=tf.int32)
+    def _clone_examples_for_beam(old_feature, n):
+      """Clone each example n times."""
+      old_shape = common_layers.shape_list(old_feature)
+      assert len(old_shape) >= 1
 
-    if self.has_input:
-      inputs_old = features["inputs"]
-      features["inputs"] = tf.expand_dims(features["inputs"], 1)
-      if len(features["inputs"].shape) < 5:
-        features["inputs"] = tf.expand_dims(features["inputs"], 4)
       # Expand the inputs in to the beam size.
-      features["inputs"] = tf.tile(features["inputs"], [1, beam_size, 1, 1, 1])
-      s = common_layers.shape_list(features["inputs"])
-      features["inputs"] = tf.reshape(features["inputs"],
-                                      [s[0] * s[1], s[2], s[3], s[4]])
+      feature = tf.expand_dims(old_feature, 1)
+      feature = tf.tile(feature, [1, n] + [1] * (len(old_shape) - 1))
+      new_shape = common_layers.shape_list(feature)
+      feature = tf.reshape(feature,
+                           [new_shape[0] * new_shape[1]] + new_shape[2:])
+      return feature
+
+    initial_ids = tf.zeros([batch_size], dtype=tf.int32)
+
+    # Clone select features multiple times to account for beam size.
+    old_features = {}
+    for feature_name in ["inputs", "knowledge"]:
+      if feature_name not in features:
+        continue
+      old_features[feature_name] = features[feature_name]
+      features[feature_name] = _clone_examples_for_beam(
+          features[feature_name], beam_size)
 
     vocab_size = self._problem_hparams.vocab_size["targets"]
     if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
@@ -904,9 +915,9 @@ def symbols_to_logits_fn(ids, i=None):
         stop_early=(top_beams == 1),
         use_tpu=use_tpu)
 
-    # Set inputs back to the unexpanded inputs to not to confuse the Estimator!
-    if self.has_input:
-      features["inputs"] = inputs_old
+    # Set features back to the unexpanded form to not to confuse the
+    # Estimator!
+    features.update(old_features)
 
     # Return `top_beams` decodings (also remove initial id from the beam search)
     # TODO(lukaszkaiser): make it work multi-problem.

From b8ed3bd5c87d08de2adbf4f68d33eb760117d540 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 14 Feb 2019 14:39:18 -0800
Subject: [PATCH 1672/2720] Keras'ify ResNet so it works with pure V2 (needs a
 layers hack for now).

PiperOrigin-RevId: 234032345
---
 tensor2tensor/layers/common_layers.py | 20 ++++++++++++++
 tensor2tensor/models/resnet.py        | 40 ++++++++++++++-------------
 tensor2tensor/v2/models/resnet.py     |  7 +++--
 3 files changed, 45 insertions(+), 22 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 6b23d7afa..b2d56d1e5 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -35,6 +35,26 @@
 from tensorflow.python.ops import inplace_ops
 
 
+_cached_layers = None
+
+
+def layers():
+  """Get the layers module good for TF 1 and TF 2 work for now."""
+  global _cached_layers
+  if _cached_layers is not None:
+    return _cached_layers
+  layers_module = tf.layers
+  try:
+    from tensorflow.python import tf2  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+    if tf2.enabled():
+      tf.logging.info("Running in V2 mode, using Keras layers.")
+      layers_module = tf.keras.layers
+  except ImportError:
+    pass
+  _cached_layers = layers_module
+  return layers_module
+
+
 @function.Defun(
     python_grad_func=lambda x, dy: tf.convert_to_tensor(dy),
     shape_func=lambda op: [op.inputs[0].get_shape()])
diff --git a/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index e1d358048..19ace57af 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -28,10 +28,16 @@
 
 import tensorflow as tf
 
+
 BATCH_NORM_DECAY = 0.9
 BATCH_NORM_EPSILON = 1e-5
 
 
+# TODO(lukaszkaiser): remove or simplify after V2 work is done.
+def layers():
+  return common_layers.layers()
+
+
 def batch_norm_relu(inputs,
                     is_training,
                     relu=True,
@@ -61,16 +67,14 @@ def batch_norm_relu(inputs,
   else:
     axis = 3
 
-  inputs = tf.layers.batch_normalization(
-      inputs=inputs,
+  inputs = layers().BatchNormalization(
       axis=axis,
       momentum=BATCH_NORM_DECAY,
       epsilon=BATCH_NORM_EPSILON,
       center=True,
       scale=True,
-      training=is_training,
       fused=True,
-      gamma_initializer=gamma_initializer)
+      gamma_initializer=gamma_initializer)(inputs, training=is_training)
 
   if relu:
     inputs = tf.nn.relu(inputs)
@@ -172,15 +176,14 @@ def conv2d_fixed_padding(inputs,
         use_bias=False,
         kernel_initializer=tf.variance_scaling_initializer())
   else:
-    y = tf.layers.conv2d(
-        inputs=inputs,
+    y = layers().Conv2D(
         filters=filters,
         kernel_size=kernel_size,
         strides=strides,
         padding=("SAME" if strides == 1 else "VALID"),
         use_bias=False,
         kernel_initializer=tf.variance_scaling_initializer(),
-        data_format=data_format)
+        data_format=data_format)(inputs)
 
   return y
 
@@ -423,7 +426,7 @@ def projection_shortcut(inputs):
 
 def resnet_v2(inputs,
               block_fn,
-              layers,
+              layer_blocks,
               filters,
               data_format="channels_first",
               is_training=False,
@@ -437,9 +440,9 @@ def resnet_v2(inputs,
     inputs: `Tensor` images.
     block_fn: `function` for the block to use within the model. Either
         `residual_block` or `bottleneck_block`.
-    layers: list of 3 or 4 `int`s denoting the number of blocks to include in
-      each of the 3 or 4 block groups. Each group consists of blocks that take
-      inputs of the same resolution.
+    layer_blocks: list of 3 or 4 `int`s denoting the number of blocks to include
+      in each of the 3 or 4 block groups. Each group consists of blocks that
+      take inputs of the same resolution.
     filters: list of 4 or 5 `int`s denoting the number of filter to include in
       block.
     data_format: `str`, "channels_first" `[batch, channels, height,
@@ -459,7 +462,7 @@ def resnet_v2(inputs,
       inputs=inputs,
       filters=filters[1],
       block_fn=block_fn,
-      blocks=layers[0],
+      blocks=layer_blocks[0],
       strides=1,
       is_training=is_training,
       name="block_layer1",
@@ -471,7 +474,7 @@ def resnet_v2(inputs,
       inputs=inputs,
       filters=filters[2],
       block_fn=block_fn,
-      blocks=layers[1],
+      blocks=layer_blocks[1],
       strides=2,
       is_training=is_training,
       name="block_layer2",
@@ -483,7 +486,7 @@ def resnet_v2(inputs,
       inputs=inputs,
       filters=filters[3],
       block_fn=block_fn,
-      blocks=layers[2],
+      blocks=layer_blocks[2],
       strides=2,
       is_training=is_training,
       name="block_layer3",
@@ -496,7 +499,7 @@ def resnet_v2(inputs,
         inputs=inputs,
         filters=filters[4],
         block_fn=block_fn,
-        blocks=layers[3],
+        blocks=layer_blocks[3],
         strides=2,
         is_training=is_training,
         name="block_layer4",
@@ -542,12 +545,11 @@ def body(self, features):
     inputs = batch_norm_relu(inputs, is_training, data_format=data_format)
 
     if not hp.is_cifar:
-      inputs = tf.layers.max_pooling2d(
-          inputs=inputs,
+      inputs = layers().MaxPooling2D(
           pool_size=3,
           strides=2,
           padding="SAME",
-          data_format=data_format)
+          data_format=data_format)(inputs)
       inputs = tf.identity(inputs, "initial_max_pool")
 
     out = resnet_v2(
@@ -572,7 +574,7 @@ def body(self, features):
     num_classes = self._problem_hparams.vocab_size["targets"]
     if hasattr(self._hparams, "vocab_divisor"):
       num_classes += (-num_classes) % self._hparams.vocab_divisor
-    logits = tf.layers.dense(out, num_classes, name="logits")
+    logits = layers().Dense(num_classes, name="logits")(out)
 
     losses = {"training": 0.0}
     if is_training:
diff --git a/tensor2tensor/v2/models/resnet.py b/tensor2tensor/v2/models/resnet.py
index a471522b0..db584e3fc 100644
--- a/tensor2tensor/v2/models/resnet.py
+++ b/tensor2tensor/v2/models/resnet.py
@@ -20,7 +20,6 @@
 from __future__ import print_function
 
 from tensor2tensor.models import resnet
-from tensor2tensor.v2 import keras_utils
 import tensorflow as tf
 import gin.tf
 
@@ -41,7 +40,7 @@ def __init__(self, features_info=None, input_names=None, target_names=None,
     num_output_classes = features_info[target_names[0]].num_classes
 
     # Now the model.
-    def resnet_model(inputs, training):
+    def resnet_model(inputs, training=None):
       return resnet.resnet_v2(
           inputs,
           resnet.bottleneck_block,
@@ -50,7 +49,9 @@ def resnet_model(inputs, training):
           is_training=training,
           is_cifar=True)
 
-    self._resnet = keras_utils.FunctionLayer(resnet_model)
+    inputs = tf.keras.Input(shape=(32, 32, 3))
+    outputs = resnet_model(inputs)
+    self._resnet = tf.keras.Model(inputs=inputs, outputs=outputs)
     self._logits = tf.keras.layers.Dense(
         num_output_classes, activation=None)
 

From 164d26baf74aa0ac5ab585d02f84b224ab86c129 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 14 Feb 2019 17:48:54 -0800
Subject: [PATCH 1673/2720] Layer cleaning and abstraction to help with V2
 porting.

PiperOrigin-RevId: 234062474
---
 tensor2tensor/layers/common_attention.py      |  19 +--
 tensor2tensor/layers/common_layers.py         | 109 +++++++++---------
 ...ention.py => message_passing_attention.py} |   0
 tensor2tensor/layers/transformer_layers.py    |  30 ++---
 4 files changed, 85 insertions(+), 73 deletions(-)
 rename tensor2tensor/layers/{common_message_passing_attention.py => message_passing_attention.py} (100%)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index b0c255aff..6cdc900a1 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -38,6 +38,11 @@
 from tensorflow.python.ops import inplace_ops
 
 
+# TODO(lukaszkaiser): remove this function when not needed any more.
+def layers():
+  return common_layers.layers()
+
+
 def large_compatible_negative(tensor_type):
   """Large negative number as Tensor.
 
@@ -275,7 +280,7 @@ def memeff_attention_fn(*args, **kwargs):
 
   # Define all available layers
 
-  layers = dict(
+  cur_layers = dict(
       # Attention layers:
       a=multihead_attention_fn,  # Multihead full attention
       loc=local_attention_fn,  # Local attention
@@ -288,7 +293,7 @@ def memeff_attention_fn(*args, **kwargs):
       sep=sep_conv_relu,  # Separable convolution (unmasked)
       sepm=sep_conv_relu_masked,  # Separable convolution (masked)
   )
-  return layers
+  return cur_layers
 
 
 def add_standard_attention_hparams(hparams):
@@ -4604,14 +4609,13 @@ def deconv_elems_1d(x, factor, out_depth=None):
   """
   out_depth = out_depth or x.get_shape().as_list()[-1]
   x = tf.expand_dims(x, 1)  # [batch_size, 1, length, depth]
-  x = tf.layers.conv2d_transpose(
-      inputs=x,
+  x = layers().Conv2DTranspose(
       filters=out_depth,
       kernel_size=(1, factor),
       strides=(1, factor),
       padding="valid",
       data_format="channels_last",
-  )  # [batch_size, 1, length*factor, out_depth]
+  )(x)  # [batch_size, 1, length*factor, out_depth]
   x = tf.squeeze(x, 1)  # [batch_size, length*factor, depth]
   return x
 
@@ -4637,14 +4641,13 @@ def conv_elems_1d(x, factor, out_depth=None):
   # with tf.control_dependencies(  # Dynamic assertion
   #     [tf.assert_equal(tf.shape(x)[1] % factor, 0)]):
   x = tf.expand_dims(x, 1)  # [batch_size, 1, length, depth]
-  x = tf.layers.conv2d(
-      inputs=x,
+  x = layers().Conv2D(
       filters=out_depth,
       kernel_size=(1, factor),
       strides=(1, factor),
       padding="valid",
       data_format="channels_last",
-  )  # [batch_size, 1, length//factor, out_depth]
+  )(x)  # [batch_size, 1, length//factor, out_depth]
   x = tf.squeeze(x, 1)  # [batch_size, length//factor, depth]
   return x
 
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index b2d56d1e5..d498f801e 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -38,6 +38,7 @@
 _cached_layers = None
 
 
+# TODO(lukaszkaiser): remove this function when not needed any more.
 def layers():
   """Get the layers module good for TF 1 and TF 2 work for now."""
   global _cached_layers
@@ -546,8 +547,10 @@ def conv2d_kernel(kernel_size_arg, name_suffix):
 
 
 def conv(inputs, filters, kernel_size, dilation_rate=(1, 1), **kwargs):
+  def _conv2d(x, *args, **kwargs):
+    return layers().Conv2D(*args, **kwargs)(x)
   return conv_internal(
-      tf.layers.conv2d,
+      _conv2d,
       inputs,
       filters,
       kernel_size,
@@ -557,16 +560,15 @@ def conv(inputs, filters, kernel_size, dilation_rate=(1, 1), **kwargs):
 
 def conv1d(inputs, filters, kernel_size, dilation_rate=1, **kwargs):
   return tf.squeeze(
-      conv(
-          tf.expand_dims(inputs, 2),
-          filters, (kernel_size, 1),
-          dilation_rate=(dilation_rate, 1),
-          **kwargs), 2)
+      conv(tf.expand_dims(inputs, 2), filters, (kernel_size, 1),
+           dilation_rate=(dilation_rate, 1), **kwargs),
+      2)
 
 
 def separable_conv(inputs, filters, kernel_size, **kwargs):
-  return conv_internal(tf.layers.separable_conv2d, inputs, filters, kernel_size,
-                       **kwargs)
+  def _sep_conv2d(x, *args, **kwargs):
+    return layers().SeparableConv2D(*args, **kwargs)(x)
+  return conv_internal(_sep_conv2d, inputs, filters, kernel_size, **kwargs)
 
 
 def subseparable_conv(inputs, filters, kernel_size, **kwargs):
@@ -584,22 +586,22 @@ def conv_fn(inputs, filters, kernel_size, **kwargs):
         with tf.variable_scope("part_%d" % split_idx):
           if separability > 0:
             parts.append(
-                tf.layers.conv2d(split, filters // separability, kernel_size,
-                                 **kwargs))
+                layers().Conv2D(filters // separability, kernel_size,
+                                **kwargs)(split))
           else:
             parts.append(
-                tf.layers.separable_conv2d(split, filters // abs_sep,
-                                           kernel_size, **kwargs))
+                layers().SeparableConv2D(filters // abs_sep,
+                                         kernel_size, **kwargs)(split))
       if separability > 1:
-        result = tf.layers.conv2d(tf.concat(parts, axis=3), filters, (1, 1))
+        result = layers().Conv2D(filters, (1, 1))(tf.concat(parts, axis=3))
       elif abs_sep == 1:  # If we have just one block, return it.
         assert len(parts) == 1
         result = parts[0]
       else:
         result = tf.concat(parts, axis=3)
     else:
-      result = tf.layers.separable_conv2d(inputs, filters, kernel_size,
-                                          **kwargs)
+      result = layers().SeparableConv2D(filters, kernel_size,
+                                        **kwargs)(inputs)
     if separability is not None:
       kwargs["separability"] = separability
     return result
@@ -776,7 +778,7 @@ def apply_norm(x, norm_type, depth, epsilon, layer_collection=None):
   if norm_type == "group":
     return group_norm(x, filters=depth, epsilon=epsilon)
   if norm_type == "batch":
-    return tf.layers.batch_normalization(x, epsilon=epsilon)
+    return layers().BatchNormalization(epsilon=epsilon)(x)
   if norm_type == "noam":
     return noam_norm(x, epsilon)
   if norm_type == "l2":
@@ -2121,7 +2123,7 @@ def gated_linear_unit_layer(x, name=None):
   """
   with tf.variable_scope(name, default_name="glu_layer", values=[x]):
     depth = shape_list(x)[-1]
-    x = tf.layers.dense(x, depth * 2, activation=None)
+    x = layers().Dense(depth * 2, activation=None)(x)
     x, gating_x = tf.split(x, 2, axis=-1)
     return x * tf.nn.sigmoid(gating_x)
 
@@ -2170,7 +2172,7 @@ def next_state(cur_state, args_tup):
       # The parallel part of the SRU.
       x_orig = x
       x, f, r = tf.split(
-          tf.layers.dense(x, 3 * x_shape[-1], name="kernel_%d" % i), 3, axis=-1)
+          layers().Dense(3 * x_shape[-1], name="kernel_%d" % i)(x), 3, axis=-1)
       f, r = tf.sigmoid(f), tf.sigmoid(r)
       x_times_one_minus_f = x * (1.0 - f)  # Compute in parallel for speed.
       # Calculate states.
@@ -2262,7 +2264,7 @@ def sru(x,
       # The parallel part of the SRU.
       x_orig = x
       x, f, r = tf.split(
-          tf.layers.dense(x, 3 * x_shape[-1], name="kernel_%d" % i), 3, axis=-1)
+          layers().Dense(3 * x_shape[-1], name="kernel_%d" % i)(x), 3, axis=-1)
       f, r = tf.sigmoid(f), tf.sigmoid(r)
       x_times_one_minus_f = x * (1.0 - f)  # Compute in parallel for speed.
       # Calculate states.
@@ -2976,9 +2978,9 @@ def fn_with_recompute(*args):
 
 
 def dense(x, units, **kwargs):
-  """Identical to tf.layers.dense."""
+  """Identical to layers.dense."""
   layer_collection = kwargs.pop("layer_collection", None)
-  activations = tf.layers.dense(x, units, **kwargs)
+  activations = layers().Dense(units, **kwargs)(x)
   if layer_collection:
     # We need to find the layer parameters using scope name for the layer, so
     # check that the layer is named. Otherwise parameters for different layers
@@ -3327,7 +3329,12 @@ def cast_like(x, y):
 
   cast_x = tf.cast(x, y.dtype)
   if cast_x.device != x.device:
-    tf.logging.warning("Cast for %s may induce copy from '%s' to '%s'", x.name,
+    x_name = "(eager Tensor)"
+    try:
+      x_name = x.name
+    except AttributeError:
+      pass
+    tf.logging.warning("Cast for %s may induce copy from '%s' to '%s'", x_name,
                        x.device, cast_x.device)
   return cast_x
 
@@ -3446,20 +3453,19 @@ def deep_discriminator(x,
   with tf.variable_scope(
       "discriminator", initializer=tf.random_normal_initializer(stddev=0.02)):
     batch_size, height, width = shape_list(x)[:3]  # pylint: disable=unbalanced-tuple-unpacking
-    net = tf.layers.conv2d(
-        x, filters, filter_size, strides=stride, padding="SAME", name="conv1")
+    net = layers().Conv2D(
+        filters, filter_size, strides=stride, padding="SAME", name="conv1")(x)
     net = lrelu(net)
-    net = tf.layers.conv2d(
-        net,
+    net = layers().Conv2D(
         2 * filters,
         filter_size,
         strides=stride,
         padding="SAME",
-        name="conv2")
+        name="conv2")(net)
     # [bs, h/4, w/4, 128]
     if batch_norm:
-      net = tf.layers.batch_normalization(
-          net, training=is_training, momentum=0.999, name="d_bn2")
+      net = layers().BatchNormalization(
+          training=is_training, momentum=0.999, name="d_bn2")(net)
     net = lrelu(net)
     size = height * width
     x_shape = x.get_shape().as_list()
@@ -3467,10 +3473,10 @@ def deep_discriminator(x,
       net = tf.reduce_mean(net, axis=[1, 2])
     else:
       net = tf.reshape(net, [batch_size, size * 8])
-    net = tf.layers.dense(net, output_size, name="d_fc3")
+    net = layers().Dense(output_size, name="d_fc3")(net)
     if batch_norm:
-      net = tf.layers.batch_normalization(
-          net, training=is_training, momentum=0.999, name="d_bn3")
+      net = layers().BatchNormalization(
+          training=is_training, momentum=0.999, name="d_bn3")(net)
     net = lrelu(net)
     return net
 
@@ -3502,17 +3508,16 @@ def general_conv(x,
                  relufactor=0):
   """Generalized convolution layer."""
   with tf.variable_scope(name):
-    x = tf.layers.conv2d(
-        x,
+    x = layers().Conv2D(
         num_filters,
         filter_size,
         stride,
         padding,
         activation=None,
         kernel_initializer=tf.truncated_normal_initializer(stddev=stddev),
-        bias_initializer=tf.constant_initializer(0.0))
+        bias_initializer=tf.constant_initializer(0.0))(x)
     if do_norm == "layer":
-      x = tf.contrib.layers.layer_norm(x)
+      x = layer_norm(x)
     elif do_norm == "instance":
       x = instance_norm(x)
 
@@ -3553,22 +3558,22 @@ def mean_with_attention(x, name, num_heads=4):
   with tf.variable_scope(name):
     shape = shape_list(x)
     m = tf.reduce_mean(x, [1, 2])
-    a = tf.layers.dense(x, num_heads, name="mean_attn")
+    a = layers().Dense(num_heads, name="mean_attn")(x)
     s = tf.reshape(a, [shape[0], -1, num_heads])
     s = tf.nn.softmax(s, axis=1)
     s = tf.reshape(s, shape[:-1] + [1, num_heads])
     am = tf.reduce_mean(tf.expand_dims(x, axis=-1) * s, [1, 2])
     l = tf.concat([am, tf.expand_dims(m, axis=-1)], axis=-1)
-    return tf.layers.dense(tf.reshape(l, [shape[0], (num_heads+1) * shape[-1]]),
-                           2 * shape[-1], name="mean_attn_final")
+    return layers().Dense(2 * shape[-1], name="mean_attn_final")(
+        tf.reshape(l, [shape[0], (num_heads+1) * shape[-1]]))
 
 
 def single_discriminator(x, filters=128, kernel_size=8,
                          strides=4, pure_mean=False):
   """A simple single-layer convolutional discriminator."""
   with tf.variable_scope("discriminator"):
-    net = tf.layers.conv2d(
-        x, filters, kernel_size, strides=strides, padding="SAME", name="conv1")
+    net = layers().Conv2D(
+        filters, kernel_size, strides=strides, padding="SAME", name="conv1")(x)
     if pure_mean:
       net = tf.reduce_mean(net, [1, 2])
     else:
@@ -3583,16 +3588,16 @@ def double_discriminator(x, filters1=128, filters2=None,
     filters2 = 4 * filters1
   with tf.variable_scope("discriminator"):
     batch_size = shape_list(x)[0]
-    net = tf.layers.conv2d(
-        x, filters1, kernel_size, strides=strides, padding="SAME", name="conv1")
+    net = layers().Conv2D(
+        filters1, kernel_size, strides=strides, padding="SAME", name="conv1")(x)
     if pure_mean:
       net1 = tf.reduce_mean(net, [1, 2])
     else:
       net1 = mean_with_attention(net, "mean_with_attention1")
       tf.reshape(net, [batch_size, -1])
     net = tf.nn.relu(net)
-    net = tf.layers.conv2d(
-        x, filters2, kernel_size, strides=strides, padding="SAME", name="conv2")
+    net = layers().Conv2D(
+        filters2, kernel_size, strides=strides, padding="SAME", name="conv2")(x)
     if pure_mean:
       net2 = tf.reduce_mean(net, [1, 2])
     else:
@@ -3658,21 +3663,21 @@ def cyclegan_upsample(net, num_outputs, stride, method="conv2d_transpose"):
       net = tf.image.resize_nearest_neighbor(
           net, [stride[0] * height, stride[1] * width])
       net = tf.pad(net, spatial_pad_1, "REFLECT")
-      net = tf.layers.conv2d(
-          net, num_outputs, (3, 3), activation=tf.nn.relu)
+      net = layers().Conv2D(
+          num_outputs, (3, 3), activation=tf.nn.relu)(net)
     elif method == "bilinear_upsample_conv":
       net = tf.image.resize_bilinear(net,
                                      [stride[0] * height, stride[1] * width])
       net = tf.pad(net, spatial_pad_1, "REFLECT")
-      net = tf.layers.conv2d(
-          net, num_outputs, (3, 3), activation=tf.nn.relu)
+      net = layers().Conv2D(
+          num_outputs, (3, 3), activation=tf.nn.relu)(net)
     elif method == "conv2d_transpose":
       # This corrects 1 pixel offset for images with even width and height.
       # conv2d is left aligned and conv2d_transpose is right aligned for even
       # sized images (while doing "SAME" padding).
       # Note: This doesn"t reflect actual model in paper.
-      net = tf.layers.conv2d_transpose(
-          net, num_outputs, (3, 3), strides=stride, activation=tf.nn.relu)
+      net = layers().Conv2DTranspose(
+          num_outputs, (3, 3), strides=stride, activation=tf.nn.relu)(net)
       net = net[:, 1:, 1:, :]
     else:
       raise ValueError("Unknown method: [%s]" % method)
@@ -3957,7 +3962,7 @@ def _data_dep_init(self, inputs):
   def build(self, input_shape=None):
     """Build `Layer`."""
     input_shape = tf.TensorShape(input_shape).as_list()
-    self.input_spec = tf.layers.InputSpec(shape=input_shape)
+    self.input_spec = layers().InputSpec(shape=input_shape)
 
     if not self.layer.built:
       self.layer.build(input_shape)
diff --git a/tensor2tensor/layers/common_message_passing_attention.py b/tensor2tensor/layers/message_passing_attention.py
similarity index 100%
rename from tensor2tensor/layers/common_message_passing_attention.py
rename to tensor2tensor/layers/message_passing_attention.py
diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 636f6e4a8..85c76823b 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -27,6 +27,11 @@
 import tensorflow as tf
 
 
+# TODO(lukaszkaiser): remove this function when not needed any more.
+def layers():
+  return common_layers.layers()
+
+
 def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
   """Prepare one shard of the model for the encoder.
 
@@ -278,9 +283,9 @@ def evolved_transformer_encoder(encoder_input,
           residual_state = hidden_state
           hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
 
-          values = tf.layers.dense(hidden_state, hparams.hidden_size)
-          gates = tf.layers.dense(
-              hidden_state, hparams.hidden_size, activation=tf.nn.sigmoid)
+          values = layers().Dense(hparams.hidden_size)(hidden_state)
+          gates = layers().Dense(
+              hparams.hidden_size, activation=tf.nn.sigmoid)(hidden_state)
           hidden_state = values * gates
 
           hidden_state = common_layers.layer_postprocess(
@@ -296,19 +301,18 @@ def evolved_transformer_encoder(encoder_input,
           hidden_state *= mask
 
           left_output_dim = int(hparams.hidden_size * 4)
-          left_state = tf.layers.dense(
-              hidden_state, left_output_dim, activation=tf.nn.relu)
+          left_state = layers().Dense(
+              left_output_dim, activation=tf.nn.relu)(hidden_state)
           left_state = tf.nn.dropout(left_state,
                                      1 - hparams.layer_prepostprocess_dropout)
 
           right_output_dim = int(hparams.hidden_size / 2)
-          right_state = tf.layers.conv1d(
-              hidden_state,
+          right_state = layers().Conv1D(
               right_output_dim,
               3,
               padding="SAME",
               name="standard_conv_3x1",
-              activation=tf.nn.relu)
+              activation=tf.nn.relu)(hidden_state)
           right_state = tf.nn.dropout(right_state,
                                       1 - hparams.layer_prepostprocess_dropout)
 
@@ -323,9 +327,9 @@ def evolved_transformer_encoder(encoder_input,
           mask = tf.tile(tf.expand_dims(nonpadding, 2), [1, 1, left_output_dim])
           hidden_state *= mask
 
-          separable_conv_9x1 = tf.layers.SeparableConv1D(
+          separable_conv_9x1 = layers().SeparableConv1D(
               right_output_dim, 9, padding="SAME", name="separable_conv_9x1")
-          hidden_state = separable_conv_9x1.apply(hidden_state)
+          hidden_state = separable_conv_9x1(hidden_state)
           hidden_state = tf.pad(
               hidden_state,
               [[0, 0], [0, 0], [0, hparams.hidden_size - right_output_dim]],
@@ -367,12 +371,12 @@ def evolved_transformer_encoder(encoder_input,
           residual_state = hidden_state
           hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
 
-          hidden_state = tf.layers.dense(
-              hidden_state, int(hparams.hidden_size * 4), activation=tf.nn.relu)
+          hidden_state = layers().Dense(
+              int(hparams.hidden_size * 4), activation=tf.nn.relu)(hidden_state)
           hidden_state = tf.nn.dropout(hidden_state,
                                        1 - hparams.layer_prepostprocess_dropout)
 
-          hidden_state = tf.layers.dense(hidden_state, hparams.hidden_size)
+          hidden_state = layers().Dense(hparams.hidden_size)(hidden_state)
           hidden_state = common_layers.layer_postprocess(
               residual_state, hidden_state, hparams)
 

From 058a16f303604837080f0f2dac70523b71f6f6fa Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 15 Feb 2019 09:28:07 -0800
Subject: [PATCH 1674/2720] Internal

PiperOrigin-RevId: 234157661
---
 tensor2tensor/utils/decoding.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index b20901180..02d880d90 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -50,6 +50,7 @@ def decode_hparams(overrides=""):
       save_images=False,
       log_results=True,
       extra_length=100,
+      min_length_ratio=0.0,
       batch_size=0,
       beam_size=4,
       alpha=0.6,

From 1eae825756a293d7806ebeeddfd06a5566c270aa Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 15 Feb 2019 10:07:48 -0800
Subject: [PATCH 1675/2720] Internal

PiperOrigin-RevId: 234164775
---
 tensor2tensor/utils/decoding.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 02d880d90..1b69cbf77 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -58,6 +58,7 @@ def decode_hparams(overrides=""):
       block_size=0,
       guess_and_check_top_k=0,
       guess_and_check_epsilon=-1,
+      insertion_parallel=False,
       return_beams=False,
       write_beam_scores=False,
       max_input_size=-1,

From 50fc5c9b4c040f988b478b04b87162268bae2123 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 15 Feb 2019 10:37:05 -0800
Subject: [PATCH 1676/2720] Update to depend on stable tensorflow-datasets

PiperOrigin-RevId: 234170650
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index e6b66dcac..9221ee573 100644
--- a/setup.py
+++ b/setup.py
@@ -53,7 +53,7 @@
         'sympy',
         'tensorflow-probability',
         'tf-agents',
-        'tfds-nightly',
+        'tensorflow-datasets',
         'tqdm',
     ],
     extras_require={

From 97f18214e01abcc8d1f1f1c37480c25a1c1d83af Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 15 Feb 2019 13:11:51 -0800
Subject: [PATCH 1677/2720] Replace more modality settings with ModalityType.

PiperOrigin-RevId: 234199556
---
 tensor2tensor/data_generators/cifar.py            |  4 ++--
 tensor2tensor/data_generators/image_utils.py      |  4 ++--
 tensor2tensor/data_generators/problem_test.py     |  8 ++++----
 tensor2tensor/data_generators/wikisum/wikisum.py  |  4 ++--
 tensor2tensor/models/image_transformer.py         | 15 +++++++++------
 tensor2tensor/models/image_transformer_2d.py      |  2 +-
 tensor2tensor/models/research/autoencoders.py     |  4 ++--
 tensor2tensor/models/research/super_lm.py         |  2 +-
 .../models/video/basic_deterministic_params.py    | 10 +++++-----
 9 files changed, 28 insertions(+), 25 deletions(-)

diff --git a/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index 8bb7f5d4e..92bcb97da 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -240,8 +240,8 @@ def preprocess_example(self, example, unused_mode, unused_hparams):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.IdentityModality,
-                  "targets": modalities.IdentityModality}
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
     p.vocab_size = {"inputs": 256,
                     "targets": 256}
     p.batch_size_multiplier = 256
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 2c0f8c584..8313f0560 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -241,8 +241,8 @@ def example_reading_spec(self):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.modality = {"inputs": modalities.ImageModality,
-                  "targets": modalities.ClassLabelModality}
+    p.modality = {"inputs": modalities.ModalityType.IMAGE,
+                  "targets": modalities.ModalityType.CLASS_LABEL}
     p.vocab_size = {"inputs": 256,
                     "targets": self.num_classes}
     p.batch_size_multiplier = 4 if self.is_small else 256
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 9c0344767..097ef1586 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -100,8 +100,8 @@ class ModalityObjProblem(problem_module.Problem):
 
       def hparams(self, defaults, model_hparams):
         hp = defaults
-        hp.modality = {"inputs": modalities.SymbolModality,
-                       "targets": modalities.SymbolModality}
+        hp.modality = {"inputs": modalities.ModalityType.SYMBOL,
+                       "targets": modalities.ModalityType.SYMBOL}
         hp.vocab_size = {"inputs": 2,
                          "targets": 3}
 
@@ -118,7 +118,7 @@ class InputOnlyProblem(problem_module.Problem):
 
       def hparams(self, defaults, model_hparams):
         hp = defaults
-        hp.modality = {"inputs": modalities.SymbolModality}
+        hp.modality = {"inputs": modalities.ModalityType.SYMBOL}
         hp.vocab_size = {"inputs": 2}
 
     problem = InputOnlyProblem(False, False)
@@ -133,7 +133,7 @@ class TargetOnlyProblem(problem_module.Problem):
 
       def hparams(self, defaults, model_hparams):
         hp = defaults
-        hp.modality = {"targets": modalities.SymbolModality}
+        hp.modality = {"targets": modalities.ModalityType.SYMBOL}
         hp.vocab_size = {"targets": 3}
 
     problem = TargetOnlyProblem(False, False)
diff --git a/tensor2tensor/data_generators/wikisum/wikisum.py b/tensor2tensor/data_generators/wikisum/wikisum.py
index 691490156..8e994ea9d 100644
--- a/tensor2tensor/data_generators/wikisum/wikisum.py
+++ b/tensor2tensor/data_generators/wikisum/wikisum.py
@@ -90,8 +90,8 @@ def hparams(self, defaults, unused_model_hparams):
         "targets": self._encoders["targets"].vocab_size,
     }
     p.modality = {
-        "inputs": modalities.SymbolModality,
-        "targets": modalities.SymbolModality,
+        "inputs": modalities.ModalityType.SYMBOL,
+        "targets": modalities.ModalityType.SYMBOL,
     }
 
   def eval_metrics(self):
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 0b3975193..81d8de434 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -49,7 +49,7 @@ def body(self, features):
     targets = features["targets"]
     if (hparams.likelihood == cia.DistributionType.DMOL and
         (hparams.modality["targets"] !=
-         modalities.ImageChannelBottomIdentityModality or
+         modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY or
          hparams.num_channels != 1)):
       raise ValueError("When using DMOL for the likelihood,modality['targets'] "
                        "must be ImageChannelBottomIdentityModality and "
@@ -57,7 +57,7 @@ def body(self, features):
     if (not tf.get_variable_scope().reuse and
         hparams.mode != tf.estimator.ModeKeys.PREDICT and
         hparams.modality["targets"] !=
-        modalities.ImageChannelBottomIdentityModality):
+        modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY):
       tf.summary.image("targets", tf.to_float(targets), max_outputs=1)
 
     # Extra losses list if we want to use moe.
@@ -193,7 +193,7 @@ def image_transformer_base():
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.98
   hparams.label_smoothing = 0.0
-  hparams.modality["targets"] = modalities.IdentityModality
+  hparams.modality["targets"] = modalities.ModalityType.IDENTITY
   hparams.norm_type = "layer"
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.add_hparam("filter_size", 512)  # Add new ones like this.
@@ -280,7 +280,8 @@ def imagetransformer_cifar10_base_dmol():
   hparams = image_transformer_base()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.modality["targets"] = modalities.ImageChannelBottomIdentityModality
+  hparams.modality["targets"] = (
+      modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY)
   hparams.num_heads = 8
   hparams.batch_size = 8
   hparams.sampling_method = "random"
@@ -421,7 +422,8 @@ def imagetransformerpp_sep_channels_8l_8h():
   hparams = imagetransformer_base()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.modality["targets"] = modalities.ImageChannelBottomIdentityModality
+  hparams.modality["targets"] = (
+      modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY)
   hparams.num_heads = 8
   hparams.batch_size = 4
   hparams.attention_key_channels = hparams.attention_value_channels = 0
@@ -884,7 +886,8 @@ def imagetransformerpp_tiny():
   hparams = imagetransformer_tiny()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.modality["targets"] = modalities.ImageChannelBottomIdentityModality
+  hparams.modality["targets"] = (
+      modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY)
   return hparams
 
 
diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 467febe16..acf793440 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -382,7 +382,7 @@ def image_transformer2d_base():
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.98
   hparams.label_smoothing = 0.0
-  hparams.modality["targets"] = modalities.IdentityModality
+  hparams.modality["targets"] = modalities.ModalityType.IDENTITY
   hparams.norm_type = "layer"
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.add_hparam("filter_size", 512)  # Add new ones like this.
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 7486fb61d..02dde9e03 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1220,8 +1220,8 @@ def autoencoder_ordered_text():
   hparams.autoregressive_mode = "conv5"
   hparams.max_hidden_size = 1024
   hparams.modality = {
-      "inputs": modalities.IdentitySymbolModality,
-      "targets": modalities.IdentitySymbolModality,
+      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
+      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
   }
   hparams.sample_height = 128
   hparams.sample_width = 1
diff --git a/tensor2tensor/models/research/super_lm.py b/tensor2tensor/models/research/super_lm.py
index fc0b403c5..5a02ccf94 100644
--- a/tensor2tensor/models/research/super_lm.py
+++ b/tensor2tensor/models/research/super_lm.py
@@ -265,7 +265,7 @@ def super_lm_base():
   # we only want one data shard.
   hparams.no_data_parallelism = True
   # bypass the symbol modality so that we can use model parallelism.
-  hparams.modality["targets"] = modalities.IdentitySymbolModality
+  hparams.modality["targets"] = modalities.ModalityType.IDENTITY_SYMBOL
   hparams.add_hparam("filter_size", 512)
   hparams.add_hparam("mix_fraction", 0.5)
   # attention-related flags
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 05566bb7a..7616d3549 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -55,7 +55,7 @@ def next_frame_pixel_noise():
   """Basic 2-frame conv model with pixel noise."""
   hparams = next_frame_basic_deterministic()
   hparams.add_hparam("video_modality_input_noise", 0.05)
-  hparams.modality["inputs"] = modalities.VideoModalityPixelNoise
+  hparams.modality["inputs"] = modalities.ModalityType.VIDEO_PIXEL_NOISE
   return hparams
 
 
@@ -89,7 +89,7 @@ def next_frame_tpu():
 def next_frame_ae():
   """Conv autoencoder."""
   hparams = next_frame_basic_deterministic()
-  hparams.modality["inputs"] = modalities.VideoModalityBitwise
+  hparams.modality["inputs"] = modalities.ModalityType.VIDEO_BITWISE
   hparams.hidden_size = 256
   hparams.batch_size = 8
   hparams.num_hidden_layers = 4
@@ -102,7 +102,7 @@ def next_frame_ae():
 def next_frame_ae_tiny():
   """Conv autoencoder, tiny set for testing."""
   hparams = next_frame_tiny()
-  hparams.modality["inputs"] = modalities.VideoModalityBitwise
+  hparams.modality["inputs"] = modalities.ModalityType.VIDEO_BITWISE
   hparams.batch_size = 8
   hparams.dropout = 0.4
   return hparams
@@ -131,7 +131,7 @@ def next_frame_tiny():
 def next_frame_l1():
   """Basic conv model with L1 modality."""
   hparams = next_frame_basic_deterministic()
-  hparams.modality["targets"] = modalities.VideoModalityL1
+  hparams.modality["targets"] = modalities.ModalityType.VIDEO_L1
   hparams.video_modality_loss_cutoff = 2.4
   return hparams
 
@@ -140,7 +140,7 @@ def next_frame_l1():
 def next_frame_l2():
   """Basic conv model with L2 modality."""
   hparams = next_frame_basic_deterministic()
-  hparams.modality["targets"] = modalities.VideoModalityL2
+  hparams.modality["targets"] = modalities.ModalityType.VIDEO_L2
   hparams.video_modality_loss_cutoff = 2.4
   return hparams
 

From 8a8661eb722f64c161ab5ca594afc711ec841837 Mon Sep 17 00:00:00 2001
From: Allen Lavoie <allenl@google.com>
Date: Sat, 16 Feb 2019 10:19:48 -0800
Subject: [PATCH 1678/2720] Cleanup after checkpointable -> trackable rename

Sorry about that

PiperOrigin-RevId: 234306111
---
 tensor2tensor/layers/common_layers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index d498f801e..6b583549f 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3928,7 +3928,7 @@ def __init__(self, layer, data_init=False, **kwargs):
           "`Layer` instance. You passed: {input}".format(input=layer))
 
     super(WeightNorm, self).__init__(layer, **kwargs)
-    self._track_checkpointable(layer, name="layer")
+    self._track_trackable(layer, name="layer")
 
   def _compute_weights(self):
     """Generate weights with normalization."""

From f7342698de2bf7187f6f1899fec247a35b1f8fca Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Sun, 17 Feb 2019 20:44:54 -0800
Subject: [PATCH 1679/2720] Code for "packing" text-examples on the fly.

PiperOrigin-RevId: 234416767
---
 .../data_generators/text_problems.py          | 142 ++++++++++++++++++
 tensor2tensor/models/mtf_transformer2.py      |   2 +-
 2 files changed, 143 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 0a5fd66ac..30d1176a0 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -1233,3 +1233,145 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
     # Shuffle the output.
     generator_utils.shuffle_dataset([output_file], extra_fn=self._pack_fn())
+
+
+def pack_dataset(dataset, length):
+  """Creates a 'packed' version of a dataset on-the-fly.
+
+  This is meant to replace the irritation of having to create a separate
+  "packed" version of a dataset to train efficiently on TPU.
+
+  Each example in the output dataset represents several examples in the
+  input dataset.
+
+  For each key in the input dataset, two additional keys are created:
+  <key>_segmentation: an int32 tensor identifying the parts
+     representing the original example.
+  <key>_position: an int32 tensor identifying the position within the original
+     example.
+
+  Example:
+  Two input examples get combined to form an output example.
+  The input examples are:
+  {"inputs": [8, 7, 1, 0], "targets":[4, 1, 0]}
+  {"inputs": [2, 3, 4, 1], "targets":[5, 6, 1]}
+  The output example is:
+  {
+                 "inputs": [8, 7, 1, 2, 3, 4, 1, 0, 0, 0]
+    "inputs_segmentation": [1, 1, 1, 2, 2, 2, 2, 0, 0, 0]
+        "inputs_position": [0, 1, 2, 0, 1, 2, 3, 0, 0, 0]
+                "targets": [4, 1, 5, 6, 1, 0, 0, 0, 0, 0]
+   "targets_segmentation": [1, 1, 2, 2, 2, 0, 0, 0, 0, 0]
+       "targets_position": [0, 1, 0, 1, 2, 0, 0, 0, 0, 0]
+  }
+
+  For now, the fields in the input sequences must end in 1 (EOS).
+  TODO(noam): remove the need for the input sequences to end in 1.
+
+  0 represents padding in both the inputs and the outputs.
+
+  Sequences in the incoming examples are truncated to length "length", and the
+  sequences in the output examples all have fixed (padded) length "length".
+
+  TODO(noam): Implement a more intelligent algorithm to achieve a more
+  dense packing (keep multiple active combined examples).
+
+  Args:
+    dataset: a tf.data.Dataset
+    length: an integer
+  Returns:
+    a tf.data.Dataset
+  """
+  shapes = dataset.output_shapes
+  keys = shapes.keys()
+  tf.logging.info("pack_dataset: shapes=%s" % (dataset.output_shapes,))
+  batch_size = length
+  dataset = dataset.padded_batch(
+      batch_size, padded_shapes={k: [-1] for k in keys})
+  empty_example = {k: tf.zeros([0], dtype=tf.int64) for k in keys}
+
+  def write_packed_example(partial, outputs):
+    new_partial = empty_example
+    new_outputs = {}
+    for k in keys:
+      new_outputs[k] = outputs[k].write(
+          outputs[k].size(),
+          tf.pad(partial[k], [[0, length - tf.size(partial[k])]]))
+    return new_partial, new_outputs
+
+  def map_fn(x):
+    """Internal function to flat_map over.
+
+    Consumes a batch of input examples and produces a variable number of output
+    examples.
+
+    Args:
+      x: a single example
+    Returns:
+      a tf.data.Dataset
+    """
+    partial = empty_example
+    i = tf.zeros([], dtype=tf.int32)
+    dynamic_batch_size = tf.shape(x[keys[0]])[0]
+    outputs = {}
+    for k in keys:
+      outputs[k] = tf.TensorArray(
+          tf.int64, size=0, dynamic_size=True, element_shape=[length])
+    def cond_fn(i, partial, outputs):
+      del partial, outputs
+      return i < dynamic_batch_size
+    def body_fn(i, partial, outputs):
+      """Body function for while_loop.
+
+      Args:
+        i: integer scalar
+        partial: dictionary of Tensor (partially-constructed example)
+        outputs: dictionary of TensorArray
+      Returns:
+        A triple containing the new values of the inputs.
+      """
+      can_append = True
+      one_example = {}
+      for k in keys:
+        val = x[k][i]
+        val = val[:tf.reduce_sum(tf.to_int32(tf.not_equal(val, 0)))]
+        one_example[k] = val
+      for k in keys:
+        can_append = tf.logical_and(
+            can_append,
+            tf.less_equal(
+                tf.size(partial[k]) + tf.size(one_example[k]), length))
+      def false_fn():
+        return write_packed_example(partial, outputs)
+      def true_fn():
+        return partial, outputs
+      partial, outputs = tf.cond(can_append, true_fn, false_fn)
+      partial = {
+          k: tf.concat([partial[k], one_example[k][:length]], 0) for k in keys}
+      return i+1, partial, outputs
+
+    i, partial, outputs = tf.while_loop(
+        cond_fn, body_fn, (i, partial, outputs),
+        back_prop=False,
+        shape_invariants=(
+            tf.TensorShape([]),
+            {k: tf.TensorShape([None]) for k in keys},
+            {k: tf.TensorShape(None) for k in keys}))
+    partial, outputs = write_packed_example(partial, outputs)
+    packed = {}
+    for k in keys:
+      ids = outputs[k].stack()
+      packed[k] = ids
+      eoss = tf.to_int32(tf.equal(ids, 1))
+      eos_positions = tf.to_int32(tf.reshape(tf.where(eoss), [-1]))
+      nonpadding = tf.to_int32(tf.not_equal(ids, 0))
+      segment_start = tf.concat([[0, 0], eos_positions + 1], axis=0)
+      segmentation = (tf.cumsum(eoss, axis=1) + 1) * nonpadding
+      position = nonpadding * (
+          tf.range(length, dtype=tf.int32)
+          - tf.gather(segment_start, segmentation))
+      packed[k + "_segmentation"] = segmentation
+      packed[k + "_position"] = position
+    return tf.data.Dataset.from_tensor_slices(packed)
+  dataset = dataset.flat_map(map_fn)
+  return dataset
diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index 1196f7075..c74d6bcbb 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -97,7 +97,7 @@ def _import_feature(self, features, mesh, key):
     x = tf.to_int32(features[key])
     x = common_layers.expand_squeeze_to_nd(x, 2)
     batch_size = mtf.Shape(self.batch_dims).size
-    # pad to length
+    x = x[:][:self.length_dim.size]
     extra_length = self.length_dim.size - tf.shape(x)[1]
     extra_batch = batch_size - tf.shape(x)[0]
     x = tf.pad(x, [[0, extra_batch], [0, extra_length]])

From 0561a1c454103c7b0012a3922beab647233e6f92 Mon Sep 17 00:00:00 2001
From: Ben Goodrich <bgoodrich@google.com>
Date: Tue, 19 Feb 2019 12:29:59 -0800
Subject: [PATCH 1680/2720] Support disabling top_k_with_unique in TPU
 beam_search

PiperOrigin-RevId: 234651573
---
 tensor2tensor/models/transformer.py |  8 ++++--
 tensor2tensor/utils/beam_search.py  | 39 ++++++++++++++++++++++-------
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index ddd17294d..4d161adce 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -766,7 +766,8 @@ def fast_decode_tpu(encoder_output,
                     eos_id=beam_search.EOS_ID,
                     batch_size=None,
                     force_decode_length=False,
-                    scope_prefix="body/"):
+                    scope_prefix="body/",
+                    use_top_k_with_unique=True):
   """Given encoder output and a symbols to logits function, does fast decoding.
 
   Implements both greedy and beam search decoding for TPU, uses beam search iff
@@ -791,6 +792,8 @@ def fast_decode_tpu(encoder_output,
     force_decode_length: A bool, whether to force the full decode length, or if
       False, stop when all beams hit eos_id.
     scope_prefix: str, prefix for decoder layer variable scopes.
+    use_top_k_with_unique: bool, whether to use a fast (but decreased precision)
+      top_k during beam search.
 
   Returns:
     A dict of decoding results {
@@ -881,7 +884,8 @@ def fast_decode_tpu(encoder_output,
         states=cache,
         eos_id=eos_id,
         stop_early=(top_beams == 1),
-        use_tpu=True)
+        use_tpu=True,
+        use_top_k_with_unique=use_top_k_with_unique)
 
     if top_beams == 1:
       decoded_ids = decoded_ids[:, 0, 1:]
diff --git a/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index 0c57f0c6d..ba2cd999f 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -295,9 +295,16 @@ def top_k_with_unique(inputs, k):
   return top_values, indices
 
 
-def compute_topk_scores_and_seq(sequences, scores, scores_to_gather, flags,
-                                beam_size, batch_size, prefix="default",
-                                states_to_gather=None, use_tpu=False):
+def compute_topk_scores_and_seq(sequences,
+                                scores,
+                                scores_to_gather,
+                                flags,
+                                beam_size,
+                                batch_size,
+                                prefix="default",
+                                states_to_gather=None,
+                                use_tpu=False,
+                                use_top_k_with_unique=True):
   """Given sequences and scores, will gather the top k=beam size sequences.
 
   This function is used to grow alive, and finished. It takes sequences,
@@ -327,6 +334,8 @@ def compute_topk_scores_and_seq(sequences, scores, scores_to_gather, flags,
     prefix: string that will prefix unique names for the ops run.
     states_to_gather: dict (possibly nested) of decoding states.
     use_tpu: A bool, whether to compute topk scores and sequences on TPU.
+    use_top_k_with_unique: bool, whether to use a fast (but decreased precision)
+      top_k during TPU beam search.
 
   Returns:
     Tuple of
@@ -362,7 +371,10 @@ def gather(tensor, name):
     else:
       topk_gathered_states = states_to_gather
   else:
-    _, topk_indexes = top_k_with_unique(scores, k=beam_size)
+    if use_top_k_with_unique:
+      _, topk_indexes = top_k_with_unique(scores, k=beam_size)
+    else:
+      _, topk_indexes = tf.nn.top_k(scores, k=beam_size)
     # Gather up the highest scoring sequences.  For each operation added, give
     # it a concrete name to simplify observing these operations with tfdbg.
     # Clients can capture these tensors by watching these node names.
@@ -390,7 +402,8 @@ def beam_search(symbols_to_logits_fn,
                 states=None,
                 eos_id=EOS_ID,
                 stop_early=True,
-                use_tpu=False):
+                use_tpu=False,
+                use_top_k_with_unique=True):
   """Beam search with length penalties.
 
   Requires a function that can take the currently decoded symbols and return
@@ -432,6 +445,8 @@ def beam_search(symbols_to_logits_fn,
     eos_id: ID for end of sentence.
     stop_early: a boolean - stop once best sequence is provably determined.
     use_tpu: A bool, whether to do beam search on TPU.
+    use_top_k_with_unique: bool, whether to use a fast (but decreased precision)
+      top_k during TPU beam search.
 
   Returns:
     Tuple of
@@ -501,9 +516,15 @@ def grow_finished(finished_seq, finished_scores, finished_flags, curr_seq,
     curr_finished_scores = tf.concat([finished_scores, curr_scores], axis=1)
     curr_finished_flags = tf.concat([finished_flags, curr_finished], axis=1)
     return compute_topk_scores_and_seq(
-        curr_finished_seq, curr_finished_scores, curr_finished_scores,
-        curr_finished_flags, beam_size, batch_size, "grow_finished",
-        use_tpu=use_tpu)
+        curr_finished_seq,
+        curr_finished_scores,
+        curr_finished_scores,
+        curr_finished_flags,
+        beam_size,
+        batch_size,
+        "grow_finished",
+        use_tpu=use_tpu,
+        use_top_k_with_unique=use_top_k_with_unique)
 
   def grow_alive(curr_seq, curr_scores, curr_log_probs, curr_finished, states):
     """Given sequences and scores, will gather the top k=beam size sequences.
@@ -590,7 +611,7 @@ def grow_topk(i, alive_seq, alive_log_probs, states):
     # Flatten out (beam_size, vocab_size) probs in to a list of possibilities
     flat_curr_scores = tf.reshape(curr_scores, [-1, beam_size * vocab_size])
 
-    if use_tpu:
+    if use_tpu and use_top_k_with_unique:
       topk_scores, topk_ids = top_k_with_unique(
           flat_curr_scores, k=beam_size * 2)
     else:

From 93d34d69092f86b203f0f0a8230fcd9ecbe9086f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 19 Feb 2019 12:50:52 -0800
Subject: [PATCH 1681/2720] Fix training for language modeling problems

The code for verifying shared input/target vocabularies must handle the case where there are only targets but not inputs (as is the case for language modeling problems)

PiperOrigin-RevId: 234655644
---
 tensor2tensor/utils/t2t_model.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 5721d7417..4af6e5dc7 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -180,13 +180,14 @@ def __init__(self,
     hparams = hparams_lib.copy_hparams(hparams)
     if self._problem_hparams and hparams.shared_embedding_and_softmax_weights:
       # If vocabularies differ, unset shared_embedding_and_softmax_weights.
-      input_vocab_size = self._problem_hparams.vocab_size["inputs"]
-      target_vocab_size = self._problem_hparams.vocab_size["targets"]
+      input_vocab_size = self._problem_hparams.vocab_size.get("inputs")
+      target_vocab_size = self._problem_hparams.vocab_size.get("targets")
       if input_vocab_size is not None and hasattr(hparams, "vocab_divisor"):
         input_vocab_size += (-input_vocab_size) % hparams.vocab_divisor
       if target_vocab_size is not None and hasattr(hparams, "vocab_divisor"):
         target_vocab_size += (-target_vocab_size) % hparams.vocab_divisor
-      if input_vocab_size != target_vocab_size:
+      if (input_vocab_size is not None and target_vocab_size is not None and
+          input_vocab_size != target_vocab_size):
         log_info("Unsetting shared_embedding_and_softmax_weights.")
         hparams.shared_embedding_and_softmax_weights = 0
 

From 02a24cabfb40d1571e4160bd210df882cc463505 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 19 Feb 2019 17:53:12 -0800
Subject: [PATCH 1682/2720] Implements the mixture model clustering for
 compressed representation

PiperOrigin-RevId: 234710761
---
 tensor2tensor/layers/transformer_memory.py    | 126 ++++++++++++++++++
 .../layers/transformer_memory_test.py         | 111 +++++++++++++++
 2 files changed, 237 insertions(+)
 create mode 100644 tensor2tensor/layers/transformer_memory.py
 create mode 100644 tensor2tensor/layers/transformer_memory_test.py

diff --git a/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
new file mode 100644
index 000000000..0042173da
--- /dev/null
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -0,0 +1,126 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""The memory unit for remembering a sequence as a collection of clusters."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_layers
+import tensorflow as tf
+
+
+class TransformerMemory(object):
+  """Implements the Memory module.
+
+  It compresses a sequence by storing items into appropriate clusters.
+  A single item can be allocated into multiple clusters like a mixture model.
+  Each vector in the memory represents the centroid of the cluster that is
+  updated in an online fashion. The memory also keeps the total amount of
+  probability mass that is used for updating each item that indicates the amount
+  of change that has been made to each cluster.
+  """
+
+  def __init__(self, batch_size, feature_dim, memory_size):
+    """Initialize the memory object.
+
+    Args:
+      batch_size: the batch size.
+      feature_dim: the depth of the feature.
+      memory_size: the number of clusters to maintain in the memory, which does
+          not have to be the same as the segment length.
+    """
+    self.feature_dim = feature_dim
+    self.batch_size = batch_size
+    self.memory_size = memory_size
+    self.mem_vals = tf.get_variable(
+        "memvals", [self.batch_size, self.memory_size, self.feature_dim],
+        dtype=tf.float32, trainable=False,
+        initializer=tf.constant_initializer(.0))
+    self.mem_times = tf.get_variable(
+        "memtimes", [self.batch_size, self.memory_size], dtype=tf.float32,
+        trainable=False, initializer=tf.constant_initializer(.0))
+    self.seq_length_so_far = tf.get_variable(
+        "seqlensofar", [self.batch_size], dtype=tf.int32,
+        trainable=False, initializer=tf.constant_initializer(0))
+
+  def set(self, mem_vals, mem_times, seq_length_so_far):
+    set_op = tf.group([
+        self.mem_vals.assign(mem_vals),
+        self.mem_times.assign(mem_times),
+        self.seq_length_so_far.assign(seq_length_so_far)])
+    return set_op
+
+  def get(self):
+    return self.mem_vals, self.mem_times, self.seq_length_so_far
+
+  def update(self, segment):
+    """Update the memory given the segment of events.
+
+    It might be useful to consider adding a decay to each cluster to favor
+    recent events.
+
+    Args:
+      segment: a tensor of shape [batch_size, segment_length, depth].
+    Returns:
+      the update op.
+    """
+    attention_logits = tf.matmul(segment, tf.transpose(
+        self.mem_vals, [0, 2, 1]))
+    alloc_probs = tf.nn.softmax(attention_logits)
+    aggregated_alloc_probs = tf.reduce_sum(alloc_probs, axis=1)
+    time_increment = tf.where(
+        tf.equal(self.seq_length_so_far, 0),
+        tf.ones_like(self.mem_times),
+        aggregated_alloc_probs)
+    update_times = self.mem_times.assign_add(time_increment)
+    with tf.control_dependencies([update_times]):
+      allocations = tf.multiply(
+          tf.expand_dims(alloc_probs, 3), tf.expand_dims(segment, 2))
+      allocations = tf.reduce_sum(allocations, axis=1)
+      add_to_vals = tf.where(
+          tf.equal(self.seq_length_so_far, 0),
+          segment,
+          tf.div(allocations - self.mem_vals,
+                 tf.expand_dims(self.mem_times, 2)))
+      update_vals = self.mem_vals.assign_add(add_to_vals)
+      with tf.control_dependencies([update_vals]):
+        segment_length = common_layers.shape_list(segment)[1]
+        update_seq_length = self.seq_length_so_far.assign_add(
+            tf.tile(tf.expand_dims(segment_length, 0), [self.batch_size]))
+    return update_seq_length
+
+  def reset(self, entries_to_reset):
+    """Reset the entries in the memory.
+
+    Args:
+      entries_to_reset: a 1D tensor.
+    Returns:
+      the reset op.
+    """
+    num_updates = tf.size(entries_to_reset)
+    update_vals = tf.scatter_update(
+        self.mem_vals, entries_to_reset,
+        tf.tile(tf.expand_dims(
+            tf.fill([self.memory_size, self.feature_dim], .0), 0),
+                [num_updates, 1, 1]))
+    update_times = tf.scatter_update(
+        self.mem_times, entries_to_reset,
+        tf.tile(tf.expand_dims(
+            tf.fill([self.memory_size], .0), 0), [num_updates, 1]))
+    update_segs = tf.scatter_update(
+        self.seq_length_so_far, entries_to_reset, tf.fill([num_updates], 0))
+    reset_op = tf.group([update_vals, update_times, update_segs])
+    return reset_op
diff --git a/tensor2tensor/layers/transformer_memory_test.py b/tensor2tensor/layers/transformer_memory_test.py
new file mode 100644
index 000000000..4df01abae
--- /dev/null
+++ b/tensor2tensor/layers/transformer_memory_test.py
@@ -0,0 +1,111 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.layers.transformer_memory."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+from tensor2tensor.layers import transformer_memory
+import tensorflow as tf
+
+
+class TransformerMemoryTest(parameterized.TestCase, tf.test.TestCase):
+
+  def testInitialize(self):
+    batch_size = 2
+    feature_dim = 3
+    memory_size = 4
+    memory = transformer_memory.TransformerMemory(
+        batch_size, feature_dim, memory_size)
+    segment = tf.constant([[[1., 2., 3.], [1., 1., 1.],
+                            [3., 2., 1.], [2., 2., 2.]],
+                           [[3., 3., 3.], [1., 2., 3.],
+                            [3., 2., 1.], [2., 2., 2.]]])
+    update_op = memory.update(segment)
+    mem_vals, mem_times, mem_len_so_far = memory.get()
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      session.run(update_op)
+      vals, times, length_so_far = session.run([
+          mem_vals, mem_times, mem_len_so_far])
+    self.assertAllEqual([[[1., 2., 3.], [1., 1., 1.],
+                          [3., 2., 1.], [2., 2., 2.]],
+                         [[3., 3., 3.], [1., 2., 3.],
+                          [3., 2., 1.], [2., 2., 2.]]], vals)
+    self.assertAllEqual([[1., 1., 1., 1.], [1., 1., 1., 1.]], times)
+    self.assertAllEqual([4, 4], length_so_far)
+
+  def testUpdate(self):
+    batch_size = 2
+    feature_dim = 3
+    memory_size = 4
+    memory = transformer_memory.TransformerMemory(
+        batch_size, feature_dim, memory_size)
+    segment = tf.constant([[[1., 2., 3.], [2., 2., 2.],
+                            [3., 2., 1.], [2., 2., 2.]],
+                           [[2., 2., 2.], [1., 2., 3.],
+                            [3., 2., 1.], [2., 2., 2.]]])
+    init_op = memory.set(segment, [[1., 2., 3., 4.], [2., 1., 5., 1.]],
+                         [10, 9])
+    new_segment = tf.constant(
+        [[[1., 2., 3.], [3., 2., 1.],
+          [2., 2., 2.], [2., 2., 2.]],
+         [[2., 2., 2.], [1., 2., 3.],
+          [3., 2., 1.], [2., 2., 2.]]])
+    update_op = memory.update(new_segment)
+    mem_vals, mem_times, mem_len_so_far = memory.get()
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      session.run(init_op)
+      session.run(update_op)
+      vals, times, length_so_far = session.run([
+          mem_vals, mem_times, mem_len_so_far])
+      print(vals, times, length_so_far)
+    self.assertAllEqual([2, 4, 3], vals.shape)
+    self.assertAllEqual([2, 4], times.shape)
+    self.assertAllEqual([14, 13], length_so_far)
+
+  def testReset(self):
+    batch_size = 2
+    feature_dim = 3
+    memory_size = 4
+    memory = transformer_memory.TransformerMemory(
+        batch_size, feature_dim, memory_size)
+    segment = tf.constant([[[1., 2., 3.], [1., 1., 1.],
+                            [3., 2., 1.], [2., 2., 2.]],
+                           [[3., 3., 3.], [1., 2., 3.],
+                            [3., 2., 1.], [2., 2., 2.]]])
+    update_op = memory.set(segment, [[1., 2., 3., 4.], [2., 1., 5., 1.]],
+                           [10, 9])
+    reset_op = memory.reset([1])
+    mem_vals, mem_times, mem_len_so_far = memory.get()
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      session.run(update_op)
+      session.run(reset_op)
+      vals, times, length_so_far = session.run([
+          mem_vals, mem_times, mem_len_so_far])
+    self.assertAllEqual([[[1., 2., 3.], [1., 1., 1.],
+                          [3., 2., 1.], [2., 2., 2.]],
+                         [[0., 0., 0.], [0., 0., 0.],
+                          [0., 0., 0.], [0., 0., 0.]]], vals)
+    self.assertAllEqual([[1., 2., 3., 4.], [0., 0., 0., 0.]], times)
+    self.assertAllEqual([10, 0], length_so_far)
+
+if __name__ == "__main__":
+  tf.test.main()

From 74dc1e1cdcd139dbe8aaf64cb8042c92c5598cb8 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 20 Feb 2019 08:30:48 -0800
Subject: [PATCH 1683/2720] fix incorrect slice notation.

PiperOrigin-RevId: 234801687
---
 tensor2tensor/models/mtf_transformer2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index c74d6bcbb..1d1e2a71d 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -97,7 +97,7 @@ def _import_feature(self, features, mesh, key):
     x = tf.to_int32(features[key])
     x = common_layers.expand_squeeze_to_nd(x, 2)
     batch_size = mtf.Shape(self.batch_dims).size
-    x = x[:][:self.length_dim.size]
+    x = x[:, :self.length_dim.size]
     extra_length = self.length_dim.size - tf.shape(x)[1]
     extra_batch = batch_size - tf.shape(x)[0]
     x = tf.pad(x, [[0, extra_batch], [0, extra_length]])

From c9cf76eff2acf60552c8015bef2853d89459bbcc Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Wed, 20 Feb 2019 14:17:02 -0800
Subject: [PATCH 1684/2720] Add capability for on-the-fly sequence packing
 using hparams.pack_dataset.

PiperOrigin-RevId: 234867141
---
 .../data_generators/generator_utils.py        | 172 ++++++++++++++++++
 .../data_generators/ops/pack_sequences_ops.cc | 155 ++++++++++++++++
 .../ops/pack_sequences_ops_test.py            |  77 ++++++++
 .../ops/subword_text_encoder.cc               | 128 +++++++++++++
 .../ops/subword_text_encoder.h                |  44 +++++
 .../ops/subword_text_encoder_ops.cc           |  65 +++++++
 .../ops/subword_text_encoder_ops_test.py      |  40 ++++
 .../ops/subword_text_encoder_test.cc          |  50 +++++
 .../data_generators/ops/testdata/subwords     |  31 ++++
 tensor2tensor/data_generators/problem.py      |   4 +
 .../data_generators/text_problems.py          | 142 ---------------
 tensor2tensor/layers/common_hparams.py        |   4 +
 12 files changed, 770 insertions(+), 142 deletions(-)
 create mode 100644 tensor2tensor/data_generators/ops/pack_sequences_ops.cc
 create mode 100644 tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
 create mode 100644 tensor2tensor/data_generators/ops/subword_text_encoder.cc
 create mode 100644 tensor2tensor/data_generators/ops/subword_text_encoder.h
 create mode 100644 tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
 create mode 100644 tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
 create mode 100644 tensor2tensor/data_generators/ops/subword_text_encoder_test.cc
 create mode 100644 tensor2tensor/data_generators/ops/testdata/subwords

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 1fd616c19..4e6bb57ea 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -33,6 +33,7 @@
 import six.moves.urllib_request as urllib
 
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators.ops import pack_sequences_ops
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
@@ -660,6 +661,177 @@ def pack_examples(examples,
     yield c.to_dict()
 
 
+def pack_dataset(dataset, length, keys=None, use_custom_ops=False):
+  """Creates a 'packed' version of a dataset on-the-fly.
+
+  This is meant to replace the irritation of having to create a separate
+  "packed" version of a dataset to train efficiently on TPU.
+
+  Each example in the output dataset represents several examples in the
+  input dataset.
+
+  For each key in the input dataset, two additional keys are created:
+  <key>_segmentation: an int32 tensor identifying the parts
+     representing the original example.
+  <key>_position: an int32 tensor identifying the position within the original
+     example.
+
+  Example:
+  Two input examples get combined to form an output example.
+  The input examples are:
+  {"inputs": [8, 7, 1, 0], "targets":[4, 1, 0]}
+  {"inputs": [2, 3, 4, 1], "targets":[5, 6, 1]}
+  The output example is:
+  {
+                 "inputs": [8, 7, 1, 2, 3, 4, 1, 0, 0, 0]
+    "inputs_segmentation": [1, 1, 1, 2, 2, 2, 2, 0, 0, 0]
+        "inputs_position": [0, 1, 2, 0, 1, 2, 3, 0, 0, 0]
+                "targets": [4, 1, 5, 6, 1, 0, 0, 0, 0, 0]
+   "targets_segmentation": [1, 1, 2, 2, 2, 0, 0, 0, 0, 0]
+       "targets_position": [0, 1, 0, 1, 2, 0, 0, 0, 0, 0]
+  }
+
+  0 represents padding in both the inputs and the outputs.
+
+  Sequences in the incoming examples are truncated to length "length", and the
+  sequences in the output examples all have fixed (padded) length "length".
+
+  TODO(noam): This code is slow - the use_custom_ops option is faster, but
+  requiers a custom-built binary.  Resolve this so that it is easy to get
+  good perfomrance.
+
+  Args:
+    dataset: a tf.data.Dataset
+    length: an integer
+    keys: a list of strings (e.g. ["inputs", "targets"])
+    use_custom_ops: use a custom c++ op not included in standard tf (faster)
+
+  Returns:
+    a tf.data.Dataset
+  """
+  if keys is None:
+    keys = dataset.output_shapes.keys
+  # trim to length
+  dataset = dataset.map(lambda x: {k: x[k][:length] for k in keys})
+
+  batch_size = length
+  dataset = dataset.padded_batch(
+      batch_size, padded_shapes={k: [-1] for k in keys})
+  if use_custom_ops and len(keys) == 2:
+    # faster and better packing but requires custom-built binary.
+    k1, k2 = keys
+    def map_fn_custom(x):
+      """Map-function."""
+      (k1_packed, k1_segmengation, k1_position,
+       k2_packed, k2_segmentation, k2_position) = (
+           pack_sequences_ops.pack_sequences2(x[k1], x[k2], length))
+      packed = {
+          k1: k1_packed,
+          k1 + "_inputs": k1_segmengation,
+          k1 + "_position": k1_position,
+          k2: k2_packed,
+          k2 + "_inputs": k2_segmentation,
+          k2 + "_position": k2_position,
+      }
+      return tf.data.Dataset.from_tensor_slices(packed)
+    dataset = dataset.flat_map(map_fn_custom)
+    return dataset
+
+  empty_example = {}
+  for k in keys:
+    empty_example[k] = tf.zeros([0], dtype=tf.int64)
+    empty_example[k + "_position"] = tf.zeros([0], dtype=tf.int32)
+  keys_etc = empty_example.keys()
+
+  def write_packed_example(partial, outputs):
+    new_partial = empty_example.copy()
+    new_outputs = {}
+    for k in keys_etc:
+      new_outputs[k] = outputs[k].write(
+          outputs[k].size(),
+          tf.pad(partial[k], [[0, length - tf.size(partial[k])]]))
+    return new_partial, new_outputs
+
+  def map_fn(x):
+    """Internal function to flat_map over.
+
+    Consumes a batch of input examples and produces a variable number of output
+    examples.
+
+    Args:
+      x: a single example
+    Returns:
+      a tf.data.Dataset
+    """
+    partial = empty_example.copy()
+    i = tf.zeros([], dtype=tf.int32)
+    dynamic_batch_size = tf.shape(x[keys[0]])[0]
+    outputs = {}
+    for k in keys:
+      outputs[k] = tf.TensorArray(
+          tf.int64, size=0, dynamic_size=True, element_shape=[length])
+      outputs[k + "_position"] = tf.TensorArray(
+          tf.int32, size=0, dynamic_size=True, element_shape=[length])
+    def cond_fn(i, partial, outputs):
+      del partial, outputs
+      return i < dynamic_batch_size
+    def body_fn(i, partial, outputs):
+      """Body function for while_loop.
+
+      Args:
+        i: integer scalar
+        partial: dictionary of Tensor (partially-constructed example)
+        outputs: dictionary of TensorArray
+      Returns:
+        A triple containing the new values of the inputs.
+      """
+      can_append = True
+      one_example = {}
+      for k in keys:
+        val = x[k][i]
+        val = val[:tf.reduce_sum(tf.to_int32(tf.not_equal(val, 0)))]
+        one_example[k] = val
+      for k in keys:
+        can_append = tf.logical_and(
+            can_append,
+            tf.less_equal(
+                tf.size(partial[k]) + tf.size(one_example[k]), length))
+      def false_fn():
+        return write_packed_example(partial, outputs)
+      def true_fn():
+        return partial, outputs
+      partial, outputs = tf.cond(can_append, true_fn, false_fn)
+      new_partial = {}
+      for k in keys:
+        new_seq = one_example[k][:length]
+        new_seq_len = tf.size(new_seq)
+        new_partial[k] = tf.concat([partial[k], new_seq], 0)
+        new_partial[k + "_position"] = tf.concat(
+            [partial[k + "_position"],
+             tf.range(new_seq_len, dtype=tf.int32)], 0)
+      partial = new_partial
+      return i+1, partial, outputs
+
+    i, partial, outputs = tf.while_loop(
+        cond_fn, body_fn, (i, partial, outputs),
+        back_prop=False,
+        shape_invariants=(
+            tf.TensorShape([]),
+            {k: tf.TensorShape([None]) for k in keys_etc},
+            {k: tf.TensorShape(None) for k in keys_etc},
+            ))
+    partial, outputs = write_packed_example(partial, outputs)
+    packed = {k: outputs[k].stack() for k in keys_etc}
+    for k in keys:
+      packed[k + "_segmentation"] = (
+          tf.cumsum(tf.to_int32(tf.equal(packed[k + "_position"], 0)), axis=1) *
+          tf.to_int32(tf.not_equal(packed[k], 0)))
+
+    return tf.data.Dataset.from_tensor_slices(packed)
+  dataset = dataset.flat_map(map_fn)
+  return dataset
+
+
 def make_tmp_dir(suffix="", prefix="tmp", dir=None):  # pylint: disable=redefined-builtin
   """Make a temporary directory."""
   if dir is None:
diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
new file mode 100644
index 000000000..b04d94356
--- /dev/null
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
@@ -0,0 +1,155 @@
+#include "third_party/tensorflow/core/framework/op_kernel.h"
+#include "third_party/tensorflow/core/framework/shape_inference.h"
+#include "third_party/tensorflow/core/framework/tensor.h"
+#include "third_party/tensorflow/core/framework/types.h"
+
+namespace tensor2tensor {
+namespace {
+
+using ::tensorflow::DEVICE_CPU;
+using ::tensorflow::OpKernel;
+using ::tensorflow::OpKernelConstruction;
+using ::tensorflow::OpKernelContext;
+using ::tensorflow::Status;
+using ::tensorflow::Tensor;
+using ::tensorflow::TensorShape;
+using ::tensorflow::shape_inference::InferenceContext;
+
+REGISTER_OP("PackSequences2")
+    .Input("inputs: int64")
+    .Input("targets: int64")
+    .Input("max_length: int32")
+    .Output("inputs_packed: int64")
+    .Output("inputs_segmentation: int32")
+    .Output("inputs_position: int32")
+    .Output("targets_packed: int64")
+    .Output("targets_segmentation: int32")
+    .Output("targets_position: int32")
+    .SetShapeFn([](InferenceContext* ctx) {
+                  for (int i=0; i < ctx->num_outputs(); i++) {
+                    ctx->set_output(i, ctx->Matrix(ctx->UnknownDim(),
+                                                   ctx->UnknownDim()));
+                  }
+                  return Status::OK();
+                });
+
+class PackSequences2Op : public OpKernel {
+ public:
+  explicit PackSequences2Op(
+      OpKernelConstruction* ctx) : OpKernel(ctx) {
+  }
+
+  void Compute(OpKernelContext* ctx) override {
+    auto inputs = ctx->input(0).matrix<int64>();
+    auto targets = ctx->input(1).matrix<int64>();
+    int max_length = ctx->input(2).scalar<int32>()();
+    int n = inputs.dimension(0);
+    std::vector<int> inputs_lengths(n);
+    std::vector<int> targets_lengths(n);
+    int padded_inputs_length = min(static_cast<int>(inputs.dimension(1)),
+                                   max_length);
+    for (int i = 0; i < n; i++) {
+      for (int j = 0; j < padded_inputs_length; j++) {
+          if (inputs(i, j) != 0)
+            inputs_lengths[i]++;
+      }
+    }
+    int padded_targets_length = min(static_cast<int>(targets.dimension(1)),
+                                    max_length);
+    for (int i = 0; i < n; i++) {
+      for (int j = 0; j < padded_targets_length; j++) {
+          if (targets(i, j) != 0)
+            targets_lengths[i]++;
+      }
+    }
+    int num_combined = 0;
+    std::vector<int> combined_inputs_length;
+    std::vector<int> combined_targets_length;
+    std::vector<std::vector<int> > combined_sequence_ids;
+    for (int seq_id = 0; seq_id < n; seq_id++) {
+      int inputs_length = inputs_lengths[seq_id];
+      int targets_length = targets_lengths[seq_id];
+      for (int combined_id = max(0, num_combined - 10); true; combined_id++) {
+        if (combined_id == num_combined) {
+          combined_inputs_length.push_back(inputs_length);
+          combined_targets_length.push_back(targets_length);
+          combined_sequence_ids.push_back(std::vector<int>(1, seq_id));
+          num_combined++;
+          break;
+        } else if (
+            (combined_inputs_length[combined_id] + inputs_length
+             <= max_length) &&
+            (combined_targets_length[combined_id] + targets_length
+             <= max_length)) {
+          combined_inputs_length[combined_id] += inputs_length;
+          combined_targets_length[combined_id] += targets_length;
+          combined_sequence_ids[combined_id].push_back(seq_id);
+          break;
+        }
+      }
+    }
+
+    auto output_shape = TensorShape(
+        {static_cast<int64>(num_combined), static_cast<int64>(max_length)});
+
+    Tensor* inputs_packed;
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, output_shape, &inputs_packed));
+    auto inputs_packed_m = inputs_packed->matrix<int64>();
+    inputs_packed_m.setZero();
+
+    Tensor* inputs_segmentation;
+    OP_REQUIRES_OK(
+        ctx, ctx->allocate_output(1, output_shape, &inputs_segmentation));
+    auto inputs_segmentation_m = inputs_segmentation->matrix<int32>();
+    inputs_segmentation_m.setZero();
+
+    Tensor* inputs_position;
+    OP_REQUIRES_OK(
+        ctx, ctx->allocate_output(2, output_shape, &inputs_position));
+    auto inputs_position_m = inputs_position->matrix<int32>();
+    inputs_position_m.setZero();
+
+    Tensor* targets_packed;
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(3, output_shape, &targets_packed));
+    auto targets_packed_m = targets_packed->matrix<int64>();
+    targets_packed_m.setZero();
+
+    Tensor* targets_segmentation;
+    OP_REQUIRES_OK(
+        ctx, ctx->allocate_output(4, output_shape, &targets_segmentation));
+    auto targets_segmentation_m = targets_segmentation->matrix<int32>();
+    targets_segmentation_m.setZero();
+
+    Tensor* targets_position;
+    OP_REQUIRES_OK(
+        ctx, ctx->allocate_output(5, output_shape, &targets_position));
+    auto targets_position_m = targets_position->matrix<int32>();
+    targets_position_m.setZero();
+
+    for (int combined_id = 0; combined_id < num_combined; combined_id++) {
+      int inputs_pos = 0;
+      int targets_pos = 0;
+      for (int i=0; i < combined_sequence_ids[combined_id].size(); i++) {
+        int seq_id = combined_sequence_ids[combined_id][i];
+        for (int j=0; j < inputs_lengths[seq_id]; j++) {
+          inputs_packed_m(combined_id, inputs_pos) = inputs(seq_id, j);
+          inputs_segmentation_m(combined_id, inputs_pos) = i + 1;
+          inputs_position_m(combined_id, inputs_pos) = j;
+          inputs_pos++;
+        }
+        for (int j=0; j < targets_lengths[seq_id]; j++) {
+          targets_packed_m(combined_id, targets_pos) = targets(seq_id, j);
+          targets_segmentation_m(combined_id, targets_pos) = i + 1;
+          targets_position_m(combined_id, targets_pos) = j;
+          targets_pos++;
+        }
+      }
+    }
+  }
+};
+
+REGISTER_KERNEL_BUILDER(Name("PackSequences2").Device(DEVICE_CPU),
+                        PackSequences2Op);
+
+}  // namespace
+}  // namespace tensor2tensor
diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
new file mode 100644
index 000000000..77b42e7da
--- /dev/null
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
@@ -0,0 +1,77 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 pack_sequences_ops."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators.ops import pack_sequences_ops
+import tensorflow as tf
+
+
+class PackSequencesOpsTest(tf.test.TestCase):
+
+  def test_pack_sequences(self):
+    inputs = [
+        [1, 2, 3],
+        [4, 5, 0],
+        [6, 0, 0],
+    ]
+    targets = [
+        [10, 0, 0],
+        [20, 30, 40],
+        [50, 60, 0],
+    ]
+    max_length = 5
+    (inputs_packed, inputs_segmentation, inputs_position,
+     targets_packed, targets_segmentation, targets_position) = (
+         pack_sequences_ops.pack_sequences2(inputs, targets, max_length))
+    self.assertAllEqual(
+        inputs_packed, [
+            [1, 2, 3, 4, 5],
+            [6, 0, 0, 0, 0],
+        ])
+    self.assertAllEqual(
+        inputs_segmentation, [
+            [1, 1, 1, 2, 2],
+            [1, 0, 0, 0, 0],
+        ])
+    self.assertAllEqual(
+        inputs_position, [
+            [0, 1, 2, 0, 1],
+            [0, 0, 0, 0, 0],
+        ])
+    self.assertAllEqual(
+        targets_packed, [
+            [10, 20, 30, 40, 0],
+            [50, 60, 0, 0, 0],
+        ])
+    self.assertAllEqual(
+        targets_segmentation, [
+            [1, 2, 2, 2, 0],
+            [1, 1, 0, 0, 0],
+        ])
+    self.assertAllEqual(
+        targets_position, [
+            [0, 0, 1, 2, 0],
+            [0, 1, 0, 0, 0],
+        ])
+
+
+if __name__ == "__main__":
+  tf.enable_eager_execution()
+  tf.test.main()
diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder.cc b/tensor2tensor/data_generators/ops/subword_text_encoder.cc
new file mode 100644
index 000000000..9199e2a83
--- /dev/null
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder.cc
@@ -0,0 +1,128 @@
+#include "third_party/py/tensor2tensor/data_generators/ops/subword_text_encoder.h"
+
+#include "third_party/absl/strings/str_cat.h"
+#include "third_party/absl/strings/str_split.h"
+#include "third_party/absl/strings/string_view.h"
+#include "third_party/icu/include/unicode/uchar.h"
+#include "third_party/icu/include/unicode/utf8.h"
+#include "third_party/tensorflow/core/framework/tensor.h"
+#include "third_party/tensorflow/core/platform/env.h"
+
+namespace tensor2tensor {
+namespace {
+
+using ::tensorflow::Env;
+using ::tensorflow::Tensor;
+
+// End of Sequence token ID to insert at end of encoded text.
+constexpr int64 kEosTokenId = 1;
+
+}  // namespace
+
+SubwordTextEncoder::SubwordTextEncoder(const string& vocab_filename) {
+  // TODO(ormandi): Add a unified vocabulary reader function.
+  string vocab_contents;
+  TF_CHECK_OK(
+      ReadFileToString(Env::Default(), vocab_filename, &vocab_contents));
+  std::vector<absl::string_view> vocab_list =
+      absl::StrSplit(vocab_contents, '\n');
+  // Strip trailing newline by skipping last element, then strip the first and
+  // last chars to remove enclosing quotes.
+  auto vocab_size = vocab_list.size() - vocab_list.back().empty();
+  for (auto i = 0; i < vocab_size; ++i) {
+    absl::string_view token =
+        vocab_list[i].substr(1, vocab_list[i].length() - 2);
+    int char_index = 0;
+    do {
+      // Note throughout that these strings are unicode so we iterate over utf-8
+      // code points, which may be between 8-32 bits long, using U8_NEXT. It is
+      // important never to iterate directly over ascii characters or models
+      // will fail to handle non-ascii alphabets properly.
+      UChar32 c;
+      U8_NEXT(token, char_index, token.length(), c);
+      CHECK_GE(c, 0);
+      alphabet_.insert(c);
+    } while (char_index < token.length());
+    vocab_.insert({string(token), i});
+  }
+}
+
+void SubwordTextEncoder::Encode(absl::string_view text, std::vector<int>* ids) {
+  ids->clear();
+  int token_start = 0;
+  int token_end = 0;
+  UChar32 c;
+  UChar32 next_c;
+  U8_NEXT(text, token_end, text.length(), c);
+  CHECK_GE(c, 0);
+  while (token_end <= text.length()) {
+    int next_end = token_end;
+    U8_NEXT(text, next_end, text.length(), next_c);
+    CHECK_GE(next_c, 0);
+    // Subtoken break when switching from non-alphanum to alphanum, or when
+    // reaching the end of the original token.
+    if (u_isalnum(next_c) != u_isalnum(c) || token_end >= text.length()) {
+      absl::string_view next_token =
+          text.substr(token_start, token_end - token_start);
+      if (next_token != " ") {
+        EncodeSubtokens(next_token, ids);
+      }
+      token_start = token_end;
+    }
+    token_end = next_end;
+    c = next_c;
+  }
+  ids->push_back(kEosTokenId);
+}
+
+void SubwordTextEncoder::EncodeSubtokens(
+    absl::string_view token, std::vector<int> *ids) {
+  string token_s = EscapeToken(token);
+  token = token_s;
+  int subtoken_start = 0;
+  // TODO(noam): this algorithm is quadratic in the length of the token.
+  //   We should instead start with a length equal to the maximum subtoken
+  //   length in the vocabulary.
+  int subtoken_end = token.length();
+  while (subtoken_start < token.length()) {
+    absl::string_view subtoken =
+        token.substr(subtoken_start, subtoken_end - subtoken_start);
+    auto iter = vocab_.find(subtoken);
+    if (iter != vocab_.end()) {
+      ids->push_back(iter->second);
+      subtoken_start = subtoken_end;
+      // TODO(noam): again, set subtoken_end forward only enough to catch
+      // the longest subtoken in the vocabulary.
+      subtoken_end = token.length();
+    } else {
+      U8_BACK_1((const uint8_t*)token_s.data(), 0, subtoken_end);
+      if (subtoken_end <= subtoken_start) {
+        LOG(FATAL) << "Unencodable tokens found.";
+      }
+    }
+  }
+}
+
+string SubwordTextEncoder::EscapeToken(absl::string_view token) {
+  string token_s;
+  int i = 0;
+  do {
+    int prev = i;
+    UChar32 c;
+    U8_NEXT(token, i, token.length(), c);
+    CHECK_GE(c, 0);
+    if (c == '_') {
+      absl::StrAppend(&token_s, "\\u");
+    } else if (c == '\\') {
+      absl::StrAppend(&token_s, "\\\\");
+    } else if (c == '\n' || alphabet_.find(c) == alphabet_.end()) {
+      absl::StrAppend(&token_s, "\\", c, ";");
+    } else {
+      absl::StrAppend(&token_s, token.substr(prev, i - prev));
+    }
+  } while (i < token.length());
+  absl::StrAppend(&token_s, "_");
+  return token_s;
+}
+
+}  // namespace tensor2tensor
diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder.h b/tensor2tensor/data_generators/ops/subword_text_encoder.h
new file mode 100644
index 000000000..a06cdc2d2
--- /dev/null
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder.h
@@ -0,0 +1,44 @@
+#ifndef TENSOR2TESNOR_DATA_GENERATORS_OPS_SUBWORD_TEXT_ENCODER_H_
+#define TENSOR2TESNOR_DATA_GENERATORS_OPS_SUBWORD_TEXT_ENCODER_H_
+
+#include "third_party/absl/container/flat_hash_map.h"
+#include "third_party/absl/container/flat_hash_set.h"
+#include "third_party/absl/strings/string_view.h"
+#include "third_party/icu/include/unicode/uchar.h"
+#include "third_party/tensorflow/core/framework/tensor.h"
+
+namespace tensor2tensor {
+
+// A subword text encoder with built in tokenizer.
+//
+// Equivalent to tensor2tensor's subword text
+// https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/text_encoder.py,
+// This code (or a suitable replacement) should eventually move into tfds
+//   and should be deleted from tensor2tensor.
+
+class SubwordTextEncoder {
+ public:
+  explicit SubwordTextEncoder(const string& vocab_filename);
+  virtual ~SubwordTextEncoder() {}
+
+  // Breaks up input text into subtokens.
+  void Encode(absl::string_view text, std::vector<int>* ids);
+
+ private:
+  // Given a full token as input, breaks the token up into subtokens and appends
+  // corresponding IDs to the ids vector.
+  void EncodeSubtokens(absl::string_view token, std::vector<int>* ids);
+
+  // Escapes a token so unencodable characters are replaced by escape sequences.
+  string EscapeToken(absl::string_view token);
+
+  // Maps subword tokens to IDs.
+  absl::flat_hash_map<string, int64> vocab_;
+  // A set containing all valid unicode code points that can be encoded without
+  // being escaped.
+  absl::flat_hash_set<UChar32> alphabet_;
+};
+
+}  // namespace tensor2tensor
+
+#endif  // TENSOR2TESNOR_DATA_GENERATORS_OPS_SUBWORD_TEXT_ENCODER_H_
diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc b/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
new file mode 100644
index 000000000..d0ba6ec8f
--- /dev/null
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
@@ -0,0 +1,65 @@
+#include "third_party/py/tensor2tensor/data_generators/ops/subword_text_encoder.h"
+#include "third_party/tensorflow/core/framework/op_kernel.h"
+#include "third_party/tensorflow/core/framework/shape_inference.h"
+#include "third_party/tensorflow/core/framework/tensor.h"
+#include "third_party/tensorflow/core/framework/types.h"
+
+namespace tensor2tensor {
+namespace {
+
+using ::tensorflow::DEVICE_CPU;
+using ::tensorflow::OpKernel;
+using ::tensorflow::OpKernelConstruction;
+using ::tensorflow::OpKernelContext;
+using ::tensorflow::Status;
+using ::tensorflow::Tensor;
+using ::tensorflow::TensorShape;
+using ::tensorflow::shape_inference::InferenceContext;
+
+REGISTER_OP("SubwordTextEncoderEncode")
+    .Input("s: string")
+    .Output("encoded: int64")
+    .Attr("vocab_filename: string")
+    .SetShapeFn([](InferenceContext* ctx) {
+      ctx->set_output(0, ctx->Vector(ctx->UnknownDim()));
+      return Status::OK();
+    });
+
+class SubwordTextEncoderEncodeOp : public OpKernel {
+ public:
+  explicit SubwordTextEncoderEncodeOp(
+      OpKernelConstruction* ctx) : OpKernel(ctx) {
+    string vocab_filename;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("vocab_filename", &vocab_filename));
+    encoder_ = absl::make_unique<SubwordTextEncoder>(vocab_filename);
+  }
+
+  void Compute(OpKernelContext* ctx) override {
+    // Get input string and deserialize into ArticleExample proto.
+    const string& s = ctx->input(0).scalar<string>()();
+
+    // Construct encoded output tensors.
+    std::vector<int> encoded_ids;
+    encoder_->Encode(s, &encoded_ids);
+    Tensor* encoded;
+    OP_REQUIRES_OK(
+        ctx,
+        ctx->allocate_output(0, TensorShape(
+            {static_cast<int64>(encoded_ids.size())}), &encoded));
+    auto encoded_vec = encoded->vec<int64>();
+    // TODO(noam): find someone who remembers c++ eigen and ask the proper way
+    // to copy a std::Vector to an Eigen whatever-this-is
+    for (int i = 0; i < encoded_ids.size(); i++) {
+      encoded_vec(i) = encoded_ids[i];
+    }
+  }
+
+ private:
+  std::unique_ptr<SubwordTextEncoder> encoder_;
+};
+
+REGISTER_KERNEL_BUILDER(Name("SubwordTextEncoderEncode").Device(DEVICE_CPU),
+                        SubwordTextEncoderEncodeOp);
+
+}  // namespace
+}  // namespace tensor2tensor
diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py b/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
new file mode 100644
index 000000000..677f95de1
--- /dev/null
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
@@ -0,0 +1,40 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 subword_text_encoder_ops."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators.ops import subword_text_encoder_ops
+import tensorflow as tf
+
+vocab_file = (
+    "third_party/py/tensor2tensor/data_generators/ops/testdata/subwords")
+
+
+class SubwordTextEncoderOpsTest(tf.test.TestCase):
+
+  def test_subword_text_encoder_encode(self):
+    s = "the quick brown fox jumps over the lazy dog"
+    encoded = subword_text_encoder_ops.subword_text_encoder_encode(
+        s, vocab_file)
+    self.assertAllEqual(encoded, [2, 3, 4, 5, 6, 7, 8, 9, 2, 11, 12, 1])
+
+
+if __name__ == "__main__":
+  tf.enable_eager_execution()
+  tf.test.main()
diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder_test.cc b/tensor2tensor/data_generators/ops/subword_text_encoder_test.cc
new file mode 100644
index 000000000..baef07c5f
--- /dev/null
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder_test.cc
@@ -0,0 +1,50 @@
+#include "third_party/py/tensor2tensor/data_generators/ops/subword_text_encoder.h"
+
+#include "testing/base/public/gunit.h"
+#include "third_party/tensorflow/core/framework/tensor.h"
+#include "third_party/tensorflow/core/framework/tensor_testutil.h"
+
+namespace tensor2tensor {
+namespace {
+
+using ::tensorflow::DT_INT64;
+using ::tensorflow::Tensor;
+using ::tensorflow::TensorShape;
+using ::tensorflow::test::AsTensor;
+using ::tensorflow::test::ExpectTensorEqual;
+
+TEST(SubwordTextEncoderTest, EncodesSubTokens) {
+  SubwordTextEncoder encoder("third_party/py/tensor2tensor/"
+                             "data_generators/ops/testdata/subwords");
+  std::vector<int> t;
+  encoder.Encode("the quick brown fox jumps over the lazy dog", &t);
+  EXPECT_EQ(t, std::vector<int>({2, 3, 4, 5, 6, 7, 8, 9, 2, 11, 12, 1}));
+}
+
+TEST(SubwordTextEncoderTest, EncodesUnicodeSubTokens) {
+  SubwordTextEncoder encoder("third_party/py/tensor2tensor/"
+                             "data_generators/ops/testdata/subwords");
+  std::vector<int> t;
+  encoder.Encode("ɧęĻĽÒ", &t);
+  EXPECT_EQ(t, std::vector<int>({13, 14, 1}));
+}
+
+TEST(SubwordTextEncoderTest, EncodesUnicodeCodePoints) {
+  SubwordTextEncoder encoder("third_party/py/tensor2tensor/"
+                             "data_generators/ops/testdata/subwords");
+  std::vector<int> t;
+  encoder.Encode("⻦ ⻭", &t);
+  EXPECT_EQ(t, std::vector<int>({15, 18, 16, 17, 1}));
+}
+
+TEST(SubwordTextEncoderTest, EncodesCharactersNotInAlphabet) {
+  SubwordTextEncoder encoder("third_party/py/tensor2tensor/"
+                             "data_generators/ops/testdata/subwords");
+  std::vector<int> t;
+  encoder.Encode("!", &t);
+  // Subtokens: '\', '3', '3', ';', '_', '<eos>', '<pad>'.
+  EXPECT_EQ(t, std::vector<int>({19, 23, 23, 30, 17, 1}));
+}
+
+}  // namespace
+}  // namespace tensor2tensor
diff --git a/tensor2tensor/data_generators/ops/testdata/subwords b/tensor2tensor/data_generators/ops/testdata/subwords
new file mode 100644
index 000000000..2591acac1
--- /dev/null
+++ b/tensor2tensor/data_generators/ops/testdata/subwords
@@ -0,0 +1,31 @@
+'<pad>'
+'<eos>'
+'the_'
+'quick_'
+'brow'
+'n_'
+'fox_'
+'jump'
+'s_'
+'over_'
+'the_'
+'lazy_'
+'dog_'
+'ɧę'
+'ĻĽÒ_'
+'⻦'
+'⻭'
+'_'
+' '
+'\'
+'0'
+'1'
+'2'
+'3'
+'4'
+'5'
+'6'
+'7'
+'8'
+'9'
+';'
\ No newline at end of file
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 70ecbcc48..ad9a127dc 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -686,6 +686,10 @@ def _load_records_and_preprocess(filenames):
     ## Shuffle records only for training examples.
     if shuffle_files and is_training:
       dataset = dataset.shuffle(shuffle_buffer_size)
+    if hparams.get("pack_dataset", False):
+      dataset = generator_utils.pack_dataset(
+          dataset, hparams.max_length, keys=["inputs", "targets"],
+          use_custom_ops=hparams.get("use_custom_ops", False))
     if output_buffer_size:
       dataset = dataset.prefetch(output_buffer_size)
 
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 30d1176a0..0a5fd66ac 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -1233,145 +1233,3 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
     # Shuffle the output.
     generator_utils.shuffle_dataset([output_file], extra_fn=self._pack_fn())
-
-
-def pack_dataset(dataset, length):
-  """Creates a 'packed' version of a dataset on-the-fly.
-
-  This is meant to replace the irritation of having to create a separate
-  "packed" version of a dataset to train efficiently on TPU.
-
-  Each example in the output dataset represents several examples in the
-  input dataset.
-
-  For each key in the input dataset, two additional keys are created:
-  <key>_segmentation: an int32 tensor identifying the parts
-     representing the original example.
-  <key>_position: an int32 tensor identifying the position within the original
-     example.
-
-  Example:
-  Two input examples get combined to form an output example.
-  The input examples are:
-  {"inputs": [8, 7, 1, 0], "targets":[4, 1, 0]}
-  {"inputs": [2, 3, 4, 1], "targets":[5, 6, 1]}
-  The output example is:
-  {
-                 "inputs": [8, 7, 1, 2, 3, 4, 1, 0, 0, 0]
-    "inputs_segmentation": [1, 1, 1, 2, 2, 2, 2, 0, 0, 0]
-        "inputs_position": [0, 1, 2, 0, 1, 2, 3, 0, 0, 0]
-                "targets": [4, 1, 5, 6, 1, 0, 0, 0, 0, 0]
-   "targets_segmentation": [1, 1, 2, 2, 2, 0, 0, 0, 0, 0]
-       "targets_position": [0, 1, 0, 1, 2, 0, 0, 0, 0, 0]
-  }
-
-  For now, the fields in the input sequences must end in 1 (EOS).
-  TODO(noam): remove the need for the input sequences to end in 1.
-
-  0 represents padding in both the inputs and the outputs.
-
-  Sequences in the incoming examples are truncated to length "length", and the
-  sequences in the output examples all have fixed (padded) length "length".
-
-  TODO(noam): Implement a more intelligent algorithm to achieve a more
-  dense packing (keep multiple active combined examples).
-
-  Args:
-    dataset: a tf.data.Dataset
-    length: an integer
-  Returns:
-    a tf.data.Dataset
-  """
-  shapes = dataset.output_shapes
-  keys = shapes.keys()
-  tf.logging.info("pack_dataset: shapes=%s" % (dataset.output_shapes,))
-  batch_size = length
-  dataset = dataset.padded_batch(
-      batch_size, padded_shapes={k: [-1] for k in keys})
-  empty_example = {k: tf.zeros([0], dtype=tf.int64) for k in keys}
-
-  def write_packed_example(partial, outputs):
-    new_partial = empty_example
-    new_outputs = {}
-    for k in keys:
-      new_outputs[k] = outputs[k].write(
-          outputs[k].size(),
-          tf.pad(partial[k], [[0, length - tf.size(partial[k])]]))
-    return new_partial, new_outputs
-
-  def map_fn(x):
-    """Internal function to flat_map over.
-
-    Consumes a batch of input examples and produces a variable number of output
-    examples.
-
-    Args:
-      x: a single example
-    Returns:
-      a tf.data.Dataset
-    """
-    partial = empty_example
-    i = tf.zeros([], dtype=tf.int32)
-    dynamic_batch_size = tf.shape(x[keys[0]])[0]
-    outputs = {}
-    for k in keys:
-      outputs[k] = tf.TensorArray(
-          tf.int64, size=0, dynamic_size=True, element_shape=[length])
-    def cond_fn(i, partial, outputs):
-      del partial, outputs
-      return i < dynamic_batch_size
-    def body_fn(i, partial, outputs):
-      """Body function for while_loop.
-
-      Args:
-        i: integer scalar
-        partial: dictionary of Tensor (partially-constructed example)
-        outputs: dictionary of TensorArray
-      Returns:
-        A triple containing the new values of the inputs.
-      """
-      can_append = True
-      one_example = {}
-      for k in keys:
-        val = x[k][i]
-        val = val[:tf.reduce_sum(tf.to_int32(tf.not_equal(val, 0)))]
-        one_example[k] = val
-      for k in keys:
-        can_append = tf.logical_and(
-            can_append,
-            tf.less_equal(
-                tf.size(partial[k]) + tf.size(one_example[k]), length))
-      def false_fn():
-        return write_packed_example(partial, outputs)
-      def true_fn():
-        return partial, outputs
-      partial, outputs = tf.cond(can_append, true_fn, false_fn)
-      partial = {
-          k: tf.concat([partial[k], one_example[k][:length]], 0) for k in keys}
-      return i+1, partial, outputs
-
-    i, partial, outputs = tf.while_loop(
-        cond_fn, body_fn, (i, partial, outputs),
-        back_prop=False,
-        shape_invariants=(
-            tf.TensorShape([]),
-            {k: tf.TensorShape([None]) for k in keys},
-            {k: tf.TensorShape(None) for k in keys}))
-    partial, outputs = write_packed_example(partial, outputs)
-    packed = {}
-    for k in keys:
-      ids = outputs[k].stack()
-      packed[k] = ids
-      eoss = tf.to_int32(tf.equal(ids, 1))
-      eos_positions = tf.to_int32(tf.reshape(tf.where(eoss), [-1]))
-      nonpadding = tf.to_int32(tf.not_equal(ids, 0))
-      segment_start = tf.concat([[0, 0], eos_positions + 1], axis=0)
-      segmentation = (tf.cumsum(eoss, axis=1) + 1) * nonpadding
-      position = nonpadding * (
-          tf.range(length, dtype=tf.int32)
-          - tf.gather(segment_start, segmentation))
-      packed[k + "_segmentation"] = segmentation
-      packed[k + "_position"] = position
-    return tf.data.Dataset.from_tensor_slices(packed)
-  dataset = dataset.flat_map(map_fn)
-  return dataset
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 4bf961050..d67bdea9c 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -151,6 +151,10 @@ def basic_params1():
       # than max_length.
       # If max_length==0, we use hparams.batch_size instead.
       max_length=0,
+      # Pack examples on the fly.
+      pack_dataset=False,
+      # Use custom ops not included in standard tensorflow.
+      use_custom_ops=True,
       # Split targets on the first axis into chunks of this length.
       split_targets_chunk_length=0,
       split_targets_max_chunks=100,

From f6404dcc5492a8d74730e698ee69047f717458a7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 20 Feb 2019 14:17:49 -0800
Subject: [PATCH 1685/2720] Optimize local unmasked 1d attention.

PiperOrigin-RevId: 234867283
---
 tensor2tensor/layers/common_attention.py | 97 ++++++++++++++----------
 1 file changed, 58 insertions(+), 39 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 6cdc900a1..2484422f2 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2629,14 +2629,15 @@ def local_attention_1d(q, k, v, block_length=128, filter_width=100, name=None):
 
   The sequence is divided into blocks of length block_length. Attention for a
   given query position can see all memory positions in the corresponding block
-  and filter_width many positions to the left of the block.
+  and filter_width many positions to the left and right of the block.
 
   Args:
     q: a Tensor with shape [batch, heads, length, depth_k]
     k: a Tensor with shape [batch, heads, length, depth_k]
     v: a Tensor with shape [batch, heads, length, depth_v]
     block_length: an integer
-    filter_width: an integer indicating how much to look left.
+    filter_width: an integer indicating how much to look left and right of the
+      block.
     name: an optional string
 
   Returns:
@@ -2644,8 +2645,11 @@ def local_attention_1d(q, k, v, block_length=128, filter_width=100, name=None):
   """
   with tf.variable_scope(
       name, default_name="local_self_attention_1d", values=[q, k, v]):
+    # Check that q, k, v have the same shape except in their depth dimension.
+    q.get_shape()[:-1].assert_is_compatible_with(k.get_shape()[:-1])
+    q.get_shape()[:-1].assert_is_compatible_with(v.get_shape()[:-1])
+
     batch_size, num_heads, original_length, _ = common_layers.shape_list(q)
-    depth_v = common_layers.shape_list(v)[-1]
 
     # Pad query, key, value to ensure multiple of corresponding lengths.
     def pad_to_multiple(x, pad_length):
@@ -2655,48 +2659,63 @@ def pad_to_multiple(x, pad_length):
     def pad_l_and_r(x, pad_length):
       return tf.pad(x, [[0, 0], [0, 0], [pad_length, pad_length], [0, 0]])
 
-    q = pad_to_multiple(q, block_length)
-    k = pad_to_multiple(k, block_length)
-    v = pad_to_multiple(v, block_length)
-
     # Set up query blocks.
-    new_q_shape = common_layers.shape_list(q)
-    q = reshape_by_blocks(q, new_q_shape, block_length)
+    # [batch, heads, blocks_q, block_length, depth_k]
+    q = pad_to_multiple(q, block_length)
+    q = reshape_by_blocks(q, common_layers.shape_list(q), block_length)
+    total_query_blocks = common_layers.shape_list(q)[2]
 
     # Set up key and value blocks.
-    # Get gather indices.
-    k = pad_l_and_r(k, filter_width)
-    v = pad_l_and_r(v, filter_width)
-    length = common_layers.shape_list(k)[2]
-    full_filter_width = block_length + 2 * filter_width
-    indices = tf.range(0, length, delta=1, name="index_range")
-    indices = tf.reshape(indices, [1, -1, 1])  # [1, length, 1] for convs
-    kernel = tf.expand_dims(tf.eye(full_filter_width), axis=1)
-    gather_indices = tf.nn.conv1d(
-        tf.cast(indices, tf.float32),
-        kernel,
-        block_length,
-        padding="VALID",
-        name="gather_conv")
-
-    gather_indices = tf.squeeze(tf.cast(gather_indices, tf.int32), axis=0)
-
-    # Reshape keys and values to [length, batch, heads, dim] for gather. Then
-    # reshape to [batch, heads, blocks, block_length + filter_width, dim].
-    k_t = tf.transpose(k, [2, 0, 1, 3])
-    k_new = tf.gather(k_t, gather_indices)
-    k_new = tf.transpose(k_new, [2, 3, 0, 1, 4])
-
-    attention_bias = tf.expand_dims(embedding_to_padding(k_new) * -1e9, axis=-2)
-
-    v_t = tf.transpose(v, [2, 0, 1, 3])
-    v_new = tf.gather(v_t, gather_indices)
-    v_new = tf.transpose(v_new, [2, 3, 0, 1, 4])
+    # [batch, heads, blocks_k, block_length, depth_k]
+    blocks_per_filter_width = filter_width // block_length
+    remaining_items = filter_width % block_length
+    k = pad_to_multiple(k, block_length)
+    v = pad_to_multiple(v, block_length)
+    k = pad_l_and_r(k, filter_width + block_length - remaining_items)
+    v = pad_l_and_r(v, filter_width + block_length - remaining_items)
+    k = reshape_by_blocks(k, common_layers.shape_list(k), block_length)
+    v = reshape_by_blocks(v, common_layers.shape_list(v), block_length)
+
+    total_kv_blocks = common_layers.shape_list(k)[2]
+
+    slices = []
+    # prepare the left-most and right-most partial blocks if needed
+    if remaining_items:
+      first_partial_block_k = tf.slice(
+          k, [0, 0, 0, block_length - remaining_items, 0],
+          [-1, -1, total_query_blocks, -1, -1])
+      first_partial_block_v = tf.slice(
+          v, [0, 0, 0, block_length - remaining_items, 0],
+          [-1, -1, total_query_blocks, -1, -1])
+      last_partial_block_k = tf.slice(
+          k, [0, 0, total_kv_blocks - total_query_blocks, 0, 0],
+          [-1, -1, -1, remaining_items, -1])
+      last_partial_block_v = tf.slice(
+          v, [0, 0, total_kv_blocks - total_query_blocks, 0, 0],
+          [-1, -1, -1, remaining_items, -1])
+      slices.append((first_partial_block_k, first_partial_block_v))
+      slices.append((last_partial_block_k, last_partial_block_v))
+
+    # Prepare the rest of the blocks
+    first_block_index = 1 if remaining_items else 0
+    attention_blocks = 2 * blocks_per_filter_width + 1
+    for i in range(first_block_index, attention_blocks + first_block_index):
+      block_k = tf.slice(k, [0, 0, i, 0, 0],
+                         [-1, -1, total_query_blocks, -1, -1])
+      block_v = tf.slice(v, [0, 0, i, 0, 0],
+                         [-1, -1, total_query_blocks, -1, -1])
+      slices.append((block_k, block_v))
+    # [batch, heads, blocks_q, block_length + 2 * filter_width, depth_k]
+    k = tf.concat([s[0] for s in slices], axis=3)
+    v = tf.concat([s[1] for s in slices], axis=3)
+
+    attention_bias = tf.expand_dims(embedding_to_padding(k) * -1e9, axis=-2)
+    depth_v = common_layers.shape_list(v)[-1]
 
     output = dot_product_attention(
         q,
-        k_new,
-        v_new,
+        k,
+        v,
         attention_bias,
         dropout_rate=0.,
         name="local_1d",

From b5e4f53d3fdce6bb3b0048795e1e56080ebaeb8a Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 20 Feb 2019 14:19:49 -0800
Subject: [PATCH 1686/2720] Dont subtract min reward in
 `EnvProblem.process_reward` this leads to bad ppo behaviour, this makes it
 similar to what T2TEnv does (just clip and integerize).

PiperOrigin-RevId: 234867629
---
 tensor2tensor/envs/env_problem.py              | 18 ++++++++----------
 .../envs/tic_tac_toe_env_problem_test.py       | 10 ++++------
 2 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 2732dafe3..a0d2524d0 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -362,9 +362,7 @@ def is_reward_range_finite(self):
     return (min_reward != -np.inf) and (max_reward != np.inf)
 
   def process_rewards(self, rewards):
-    """Clips, rounds, adds the min_reward and changes to integer type.
-
-    The result of the above is that the new minimum is 0.
+    """Clips, rounds, and changes to integer type.
 
     Args:
       rewards: numpy array of raw (float) rewards.
@@ -375,8 +373,8 @@ def process_rewards(self, rewards):
 
     min_reward, max_reward = self.reward_range
 
-    # Clips at min and max reward and shift by min (so new min is 0)
-    rewards = np.clip(rewards, min_reward, max_reward) - min_reward
+    # Clips at min and max reward.
+    rewards = np.clip(rewards, min_reward, max_reward)
     # Round to (nearest) int and convert to integral type.
     rewards = np.around(rewards, decimals=0).astype(np.int64)
     return rewards
@@ -714,12 +712,12 @@ def _generate_time_steps(self, trajectory_list):
 
         yield {
             TIMESTEP_FIELD: [index],
-            ACTION_FIELD:
-                action,
-            RAW_REWARD_FIELD:
-                [float(raw_reward)],  # to_example errors on np.float32
+            ACTION_FIELD: action,
+            # to_example errors on np.float32
+            RAW_REWARD_FIELD: [float(raw_reward)],
             PROCESSED_REWARD_FIELD: [processed_reward],
-            DONE_FIELD: [int(time_step.done)],  # to_example doesn't know bools
+            # to_example doesn't know bools
+            DONE_FIELD: [int(time_step.done)],
             OBSERVATION_FIELD:
                 gym_spaces_utils.gym_space_encode(self.observation_space,
                                                   time_step.observation),
diff --git a/tensor2tensor/envs/tic_tac_toe_env_problem_test.py b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
index 214ab7196..bf69fe6f3 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
@@ -51,16 +51,14 @@ def test_registration_and_interaction_with_env_problem(self):
       for r, d in zip(rewards, dones):
         if not d:
           continue
-        # NOTE: r is 0, 1, 2 because the default EnvProblem.process_rewards
-        # shifts the rewards so that min is 0.
-        if r == 0:
+        if r == -1:
           num_lost += 1
-        elif r == 1:
+        elif r == 0:
           num_draw += 1
-        elif r == 2:
+        elif r == 1:
           num_won += 1
         else:
-          raise ValueError("reward should be 0, 1, 2 but is {}".format(r))
+          raise ValueError("reward should be -1, 0, 1 but is {}".format(r))
 
     # Assert that something got done atleast, without that the next assert is
     # meaningless.

From bac0da4f17c9d72c9f1ead192035888f17582c90 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 20 Feb 2019 16:44:00 -0800
Subject: [PATCH 1687/2720] Replace queries to self.hparams.problem_hparams
 with just self.problem_hparams.

PiperOrigin-RevId: 234895520
---
 tensor2tensor/data_generators/problem.py      |  1 +
 .../models/research/transformer_parallel.py   |  2 +-
 tensor2tensor/models/video/base.py            | 23 +++++++++++--------
 tensor2tensor/utils/t2t_model.py              |  4 ++++
 4 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index ad9a127dc..153f5adcc 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -1060,6 +1060,7 @@ def _default_hparams():
       # chosen model architecture. It comprises key-value pairs of a feature
       # name (str) and its modality type.
       modality={},
+      vocab_size={},
 
       # Identifiers used to tell the model which input/target space will be
       # expected. For example, it can tell that we expect French as characters
diff --git a/tensor2tensor/models/research/transformer_parallel.py b/tensor2tensor/models/research/transformer_parallel.py
index 23b04f99f..166c512e0 100644
--- a/tensor2tensor/models/research/transformer_parallel.py
+++ b/tensor2tensor/models/research/transformer_parallel.py
@@ -154,7 +154,7 @@ def infer_step(result, length):
       """Inference step."""
 
       def print_info(result, length, new_length):
-        vocab = self._hparams.problem_hparams.vocabulary["targets"]
+        vocab = self.problem_hparams.vocabulary["targets"]
         tf.logging.info(
             "length=%s new_length=%s length_diff=%s new_suffix=%s",
             length,
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index b54d6356f..c5502569a 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -26,6 +26,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 from tensor2tensor.layers import discretization
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -194,13 +195,16 @@ def __init__(self, *args, **kwargs):
 
   @property
   def _target_modality(self):
-    # TODO(mbz): get rid of this somehow.
-    modality = self.hparams.problem_hparams.modality["targets"]
-    return modality.__class__.__name__
+    target_modality = self.hparams.modality.get(
+        "targets",
+        self.problem_hparams.modality["targets"])
+    if target_modality not in modalities.ModalityType.get_choices():
+      target_modality = target_modality.__class__.__name__
+    return target_modality
 
   @property
   def is_per_pixel_softmax(self):
-    return self._target_modality == "VideoModality"
+    return self._target_modality == modalities.ModalityType.VIDEO
 
   def get_iteration_num(self):
     step_num = tf.train.get_global_step()
@@ -336,9 +340,9 @@ def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
     Raises:
       ValueError: in case of unknown modality.
     """
-    if self._target_modality == "VideoModalityL2Raw":
+    if self._target_modality == modalities.ModalityType.VIDEO_L2_RAW:
       recon_loss = tf.losses.mean_squared_error(extra_gts, extra_pds)
-    elif self._target_modality == "VideoModality":
+    elif self._target_modality == modalities.ModalityType.VIDEO:
       shape = common_layers.shape_list(extra_pds)
       updated_shape = shape[:-1] + [3, 256]
       extra_pds = tf.reshape(extra_pds, updated_shape)
@@ -347,7 +351,8 @@ def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
       targets = extra_raw_gts
       targets_shape = common_layers.shape_list(targets)
       targets = tf.reshape(targets, [-1] + targets_shape[2:])
-      mod = self.hparams.problem_hparams.modality["targets"]
+      mod = self.hparams.modality.get("targets",
+                                      self.problem_hparams.modality["targets"])
       numerator, denominator = common_layers.padded_cross_entropy(
           logits,
           targets,
@@ -467,8 +472,8 @@ def logits_to_samples(logits, key):
                        hparams.video_num_target_frames, 1, 1, num_channels]
 
     features["targets"] = tf.zeros(targets_shape, dtype=tf.int32)
-    reward_in_mod = "target_reward" in hparams.problem_hparams.modality
-    action_in_mod = "target_action" in hparams.problem_hparams.modality
+    reward_in_mod = "target_reward" in self.problem_hparams.modality
+    action_in_mod = "target_action" in self.problem_hparams.modality
     if reward_in_mod:
       # TODO(lukaszkaiser): this is a hack. get the actual reward history.
       if "input_reward" not in features:
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 4af6e5dc7..749788784 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -245,6 +245,10 @@ def eval_hooks(hook_context):
   def hparams(self):
     return self._hparams
 
+  @property
+  def problem_hparams(self):
+    return self._problem_hparams
+
   @property
   def is_training(self):
     return self._hparams.mode == tf.estimator.ModeKeys.TRAIN

From ac74489b6aa1e9c5a9abd12d919ab66d9a99c1e8 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 20 Feb 2019 18:57:58 -0800
Subject: [PATCH 1688/2720] Replace Modality class methods with functions.

In detail, this CL does the following:

+ Rewrites following Modality methods as functions: bottom, loss, name, targets_weights_fn, top, top_is_pointwise.
+ Implements get identifiers so T2TModel applies default transformations. They can be overriden by hparams.{bottom,loss,name,etc.}.

This change touches many files:

+ Major changes in modality.py, modalities.py, and t2t_model.py.
+ All model files calling modality methods. They're changed to use the get identifier.
+ All custom modality files. They're changed from class methods to functions.

Sorry for the big CL. This was the smallest incremental change I found. CLs for non-T2T directories will be submitted after this one.

PiperOrigin-RevId: 234913884
---
 .../data_generators/multi_problem.py          |  51 +-
 tensor2tensor/data_generators/problem.py      |  37 -
 tensor2tensor/data_generators/problem_test.py |  34 +-
 tensor2tensor/layers/modalities.py            | 662 +++++++++++-------
 tensor2tensor/layers/modalities_test.py       |  45 +-
 .../models/research/vqa_attention_test.py     |  11 +-
 tensor2tensor/models/resnet_test.py           |   5 +-
 tensor2tensor/models/slicenet.py              |   5 +-
 tensor2tensor/models/transformer.py           |  83 ++-
 tensor2tensor/models/video/base.py            |   8 +-
 tensor2tensor/models/video/tests_utils.py     |  35 +-
 tensor2tensor/models/xception_test.py         |   5 +-
 tensor2tensor/utils/metrics.py                |  14 +-
 tensor2tensor/utils/modality.py               |  72 +-
 tensor2tensor/utils/t2t_model.py              | 259 ++++---
 tensor2tensor/utils/t2t_model_test.py         |  15 +-
 tensor2tensor/utils/trainer_lib_test.py       |   5 +
 17 files changed, 796 insertions(+), 550 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 27ae5c118..5f165a2b5 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -172,8 +172,7 @@ def get_hparams(self, model_hparams=None):
     self.update_task_ids(vocab_size)
     tf.logging.info("New vocabulary size: %d" % new_vocab_size)
     self._hparams.vocab_size["targets"] = new_vocab_size
-    self._hparams.modality["targets"] = modalities.SymbolModality(
-        model_hparams, self._hparams.vocab_size["targets"])
+    self._hparams.modality["targets"] = modalities.ModalityType.SYMBOL
     return self._hparams
 
   def dataset(self,
@@ -420,24 +419,30 @@ def get_max_num_classes(self):
 def aggregate_task_losses(hparams,
                           problem_hparams,
                           logits,
-                          target_modality,
+                          feature_name,
                           feature):
   """Multiproblem loss function."""
 
   # If no reweighting, we want the default loss to mimic the LM loss.
   if not hparams.multiproblem_reweight_label_loss:
     return aggregate_task_lm_losses(hparams=hparams,
+                                    problem_hparams=problem_hparams,
                                     logits=logits,
-                                    target_modality=target_modality,
+                                    feature_name=feature_name,
                                     feature=feature)
 
   summaries = []
   main_task_id = hparams.problem.task_list[0].task_id
+  vocab_size = problem_hparams.vocab_size[feature_name]
+  if vocab_size is not None and hasattr(hparams, "vocab_divisor"):
+    vocab_size += (-vocab_size) % hparams.vocab_divisor
+  modality = problem_hparams.modality[feature_name]
+  loss = hparams.loss.get(feature_name, modalities.get_loss(modality))
   # Primary task loss
-  loss_num, loss_den = target_modality.loss(
+  loss_num, loss_den = loss(
       logits, feature,
-      weights_fn=
-      lambda x: common_layers.weights_multi_problem_all(x, main_task_id))
+      lambda x: common_layers.weights_multi_problem_all(x, main_task_id),
+      hparams, vocab_size)
 
   loss_val = loss_num / tf.maximum(1.0, loss_den)
   summaries.append([hparams.problem.task_list[0].name+"_loss", loss_val])
@@ -450,10 +455,10 @@ def aggregate_task_losses(hparams,
 
   for task in hparams.problem.task_list[1:]:
     # Loss only from the input sequence -- the auxiliary LM loss.
-    seq_loss_num, seq_loss_den = target_modality.loss(
+    seq_loss_num, seq_loss_den = loss(
         logits, feature,
-        weights_fn=
-        lambda x: common_layers.weights_multi_problem_input(x, task.task_id))  # pylint: disable=cell-var-from-loop
+        lambda x: common_layers.weights_multi_problem_input(x, task.task_id),  # pylint: disable=cell-var-from-loop
+        hparams, vocab_size)
     seq_loss_num *= problem_hparams.loss_multiplier
 
     # Unscaled sequence loss.
@@ -462,10 +467,10 @@ def aggregate_task_losses(hparams,
 
     if hasattr(task, "num_classes"):
       # Loss only from the classification label.
-      label_loss_num, label_loss_den = target_modality.loss(
+      label_loss_num, label_loss_den = loss(
           logits, feature,
-          weights_fn=
-          lambda x: common_layers.weights_multi_problem(x, task.task_id))  # pylint: disable=cell-var-from-loop
+          lambda x: common_layers.weights_multi_problem(x, task.task_id),  # pylint: disable=cell-var-from-loop
+          hparams, vocab_size)
       label_loss_num *= problem_hparams.loss_multiplier
 
       # Unscaled classification label loss.
@@ -484,10 +489,10 @@ def aggregate_task_losses(hparams,
 
     else:
       # Loss only from the target sequence.
-      target_loss_num, target_loss_den = target_modality.loss(
+      target_loss_num, target_loss_den = loss(
           logits, feature,
-          weights_fn=
-          lambda x: common_layers.weights_multi_problem(x, task.task_id))  # pylint: disable=cell-var-from-loop
+          lambda x: common_layers.weights_multi_problem(x, task.task_id),  # pylint: disable=cell-var-from-loop
+          hparams, vocab_size)
       target_loss_num *= problem_hparams.loss_multiplier
 
       # Unscaled target sequence loss.
@@ -516,18 +521,24 @@ def aggregate_task_losses(hparams,
 
 
 def aggregate_task_lm_losses(hparams,
+                             problem_hparams,
                              logits,
-                             target_modality,
+                             feature_name,
                              feature):
   """LM loss for multiproblems."""
   summaries = []
+  vocab_size = problem_hparams.vocab_size[feature_name]
+  if vocab_size is not None and hasattr(hparams, "vocab_divisor"):
+    vocab_size += (-vocab_size) % hparams.vocab_divisor
+  modality = problem_hparams.modality[feature_name]
+  loss = hparams.loss.get(feature_name, modalities.get_loss(modality))
   loss_num = 0.
   loss_den = 0.
   for task in hparams.problem.task_list:
-    loss_num_, loss_den_ = target_modality.loss(
+    loss_num_, loss_den_ = loss(
         logits, feature,
-        weights_fn=
-        lambda x: common_layers.weights_multi_problem_all(x, task.task_id))  # pylint: disable=cell-var-from-loop
+        lambda x: common_layers.weights_multi_problem_all(x, task.task_id),  # pylint: disable=cell-var-from-loop
+        hparams, vocab_size)
 
     loss_num += loss_num_
     loss_den += loss_den_
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 153f5adcc..8329dc69c 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -25,7 +25,6 @@
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import text_encoder
-from tensor2tensor.layers import modalities
 from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import mlperf_log
@@ -534,8 +533,6 @@ def get_hparams(self, model_hparams=None):
     if self._was_copy:
       _copy_problem_hparams(hp)
 
-    _create_modalities(hp, model_hparams)
-
     self._hparams = hp
     return self._hparams
 
@@ -1000,40 +997,6 @@ def _reverse_problem_hparams(p_hparams):
   p.was_reversed = True
 
 
-def _create_modalities(problem_hparams, model_hparams):
-  """Creates modalities and overrides any according to model hparams.
-
-  Args:
-    problem_hparams: HParams for the Problem. It must have
-      modality which is a dict of strings to ModalityTypes or Modality classes.
-    model_hparams: HParams for the model. It may have
-      input_modalities and target_modality, which will override
-      problem_hparams' modality input and target keys.
-
-  Returns:
-    None
-  """
-  modality_overrides = getattr(model_hparams, "modality", {})
-  modality = {}
-  for feature_name, modality_type in six.iteritems(problem_hparams.modality):
-    vocab_size = problem_hparams.vocab_size[feature_name]
-    # If needed for using a pre-trained model's vocabulary where extra indices
-    # were allocated for adding new tasks with unique task ids.
-    if (hasattr(model_hparams, "multiproblem_vocab_size") and
-        model_hparams.multiproblem_vocab_size > 0):
-      vocab_size = model_hparams.multiproblem_vocab_size
-    # Override modality using to the associated value in modality_overrides.
-    modality_type = modality_overrides.get(feature_name, modality_type)
-    # Each modality is a ModalityType or class. If ModalityType, get the
-    # corresponding class.
-    if modality_type in modalities.ModalityType.get_choices():
-      modality_cls = getattr(modalities, modality_type)
-    else:
-      modality_cls = modality_type
-    modality[feature_name] = modality_cls(model_hparams, vocab_size)
-  problem_hparams.modality = modality
-
-
 def _default_hparams():
   """A set of basic model hyperparameters."""
   return HParams(
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 097ef1586..03e42a193 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -89,28 +89,10 @@ def testProblemHparamsModality(self):
     problem = problem_hparams.TestProblem(input_vocab_size=2,
                                           target_vocab_size=3)
     p_hparams = problem.get_hparams()
-    self.assertIsInstance(p_hparams.modality["inputs"],
-                          modalities.SymbolModality)
-    self.assertIsInstance(p_hparams.modality["targets"],
-                          modalities.SymbolModality)
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testProblemHparamsModalityObj(self):
-    class ModalityObjProblem(problem_module.Problem):
-
-      def hparams(self, defaults, model_hparams):
-        hp = defaults
-        hp.modality = {"inputs": modalities.ModalityType.SYMBOL,
-                       "targets": modalities.ModalityType.SYMBOL}
-        hp.vocab_size = {"inputs": 2,
-                         "targets": 3}
-
-    problem = ModalityObjProblem(False, False)
-    p_hparams = problem.get_hparams()
-    self.assertIsInstance(p_hparams.modality["inputs"],
-                          modalities.SymbolModality)
-    self.assertIsInstance(p_hparams.modality["targets"],
-                          modalities.SymbolModality)
+    self.assertEqual(p_hparams.modality["inputs"],
+                     modalities.ModalityType.SYMBOL)
+    self.assertEqual(p_hparams.modality["targets"],
+                     modalities.ModalityType.SYMBOL)
 
   @test_utils.run_in_graph_and_eager_modes()
   def testProblemHparamsInputOnlyModality(self):
@@ -123,8 +105,8 @@ def hparams(self, defaults, model_hparams):
 
     problem = InputOnlyProblem(False, False)
     p_hparams = problem.get_hparams()
-    self.assertIsInstance(p_hparams.modality["inputs"],
-                          modalities.SymbolModality)
+    self.assertEqual(p_hparams.modality["inputs"],
+                     modalities.ModalityType.SYMBOL)
     self.assertLen(p_hparams.modality, 1)
 
   @test_utils.run_in_graph_and_eager_modes()
@@ -138,8 +120,8 @@ def hparams(self, defaults, model_hparams):
 
     problem = TargetOnlyProblem(False, False)
     p_hparams = problem.get_hparams()
-    self.assertIsInstance(p_hparams.modality["targets"],
-                          modalities.SymbolModality)
+    self.assertEqual(p_hparams.modality["targets"],
+                     modalities.ModalityType.SYMBOL)
     self.assertLen(p_hparams.modality, 1)
 
   @test_utils.run_in_graph_and_eager_modes()
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 4b725f569..b481419f0 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -17,6 +17,8 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+import sys
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_attention
@@ -40,20 +42,19 @@ class SymbolModality(modality.Modality):
     Linear transformation + softmax.
   """
 
-  @property
-  def name(self):
-    return "symbol_modality_%d_%d" % (self._vocab_size,
-                                      self._model_hparams.hidden_size)
+  @staticmethod
+  def name(model_hparams, vocab_size):
+    return "symbol_modality_%d_%d" % (vocab_size, model_hparams.hidden_size)
 
-  @property
-  def top_is_pointwise(self):
+  @staticmethod
+  def top_is_pointwise():
     return True
 
-  @property
-  def targets_weights_fn(self):
+  @staticmethod
+  def targets_weights_fn(model_hparams):
     weights_fn = common_layers.weights_nonzero
 
-    hp = self._model_hparams
+    hp = model_hparams
     if hp and hp.prepend_mode != "none":
       assert (hp.prepend_mode == "prepend_inputs_masked_attention" or
               hp.prepend_mode == "prepend_inputs_full_attention")
@@ -68,22 +69,25 @@ def targets_weights_fn(self):
 
     return weights_fn
 
-  def _get_weights(self, hidden_dim=None):
+  @staticmethod
+  def _get_weights(model_hparams, vocab_size, hidden_dim=None):
     """Create or get concatenated embedding or softmax variable.
 
     Args:
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
       hidden_dim: dim of the variable. Defaults to _model_hparams' hidden_size
 
     Returns:
-       a list of self._num_shards Tensors.
+       a list of num_shards Tensors.
     """
     if hidden_dim is None:
-      hidden_dim = self._model_hparams.hidden_size
-    num_shards = self._model_hparams.symbol_modality_num_shards
+      hidden_dim = model_hparams.hidden_size
+    num_shards = model_hparams.symbol_modality_num_shards
     shards = []
     for i in range(num_shards):
-      shard_size = (self._vocab_size // num_shards) + (
-          1 if i < self._vocab_size % num_shards else 0)
+      shard_size = (vocab_size // num_shards) + (
+          1 if i < vocab_size % num_shards else 0)
       var_name = "weights_%d" % i
       shards.append(
           tf.get_variable(
@@ -98,7 +102,8 @@ def _get_weights(self, hidden_dim=None):
       ret = common_layers.convert_gradient_to_tensor(ret)
     return ret
 
-  def bottom_simple(self, x, name, reuse):
+  @classmethod
+  def bottom_simple(cls, x, model_hparams, vocab_size, name, reuse):
     with tf.variable_scope(name, reuse=reuse):
       # Ensure the inputs are 3-D
       if len(x.get_shape()) == 4:
@@ -106,47 +111,59 @@ def bottom_simple(self, x, name, reuse):
       while len(x.get_shape()) < 3:
         x = tf.expand_dims(x, axis=-1)
 
-      var = self._get_weights()
+      var = cls._get_weights(model_hparams, vocab_size)
       x = common_layers.dropout_no_scaling(
-          x, 1.0 - self._model_hparams.symbol_dropout)
+          x, 1.0 - model_hparams.symbol_dropout)
       ret = common_layers.gather(var, x)
-      if self._model_hparams.multiply_embedding_mode == "sqrt_depth":
-        ret *= self._model_hparams.hidden_size**0.5
+      if model_hparams.multiply_embedding_mode == "sqrt_depth":
+        ret *= model_hparams.hidden_size**0.5
       ret *= tf.expand_dims(
           common_layers.cast_like(tf.not_equal(x, 0), ret), -1)
       return ret
 
-  def bottom(self, x):
-    if (self._model_hparams.shared_embedding_and_softmax_weights or
-        self._model_hparams.get("shared_embedding")):
-      return self.bottom_simple(x, "shared", reuse=None)
-    return self.bottom_simple(x, "input_emb", reuse=None)
-
-  def targets_bottom(self, x):
-    if (self._model_hparams.shared_embedding_and_softmax_weights or
-        self._model_hparams.get("shared_embedding")):
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    if (model_hparams.shared_embedding_and_softmax_weights or
+        model_hparams.get("shared_embedding")):
+      return cls.bottom_simple(
+          x, model_hparams, vocab_size, "shared", reuse=None)
+    return cls.bottom_simple(
+        x, model_hparams, vocab_size, "input_emb", reuse=None)
+
+  @classmethod
+  def targets_bottom(cls, x, model_hparams, vocab_size):
+    if (model_hparams.shared_embedding_and_softmax_weights or
+        model_hparams.get("shared_embedding")):
       try:
-        return self.bottom_simple(x, "shared", reuse=True)
+        return cls.bottom_simple(
+            x, model_hparams, vocab_size, "shared", reuse=True)
       except ValueError:
         # perhaps there were no inputs, and this is a new variable.
-        return self.bottom_simple(x, "shared", reuse=None)
+        return cls.bottom_simple(
+            x, model_hparams, vocab_size, "shared", reuse=None)
     else:
-      return self.bottom_simple(x, "target_emb", reuse=None)
+      return cls.bottom_simple(
+          x, model_hparams, vocab_size, "target_emb", reuse=None)
 
-  def top(self, body_output, _):
+  @classmethod
+  def top(cls, body_output, targets, model_hparams, vocab_size):
     """Generate logits.
 
     Args:
       body_output: A Tensor with shape
-        [batch, p0, p1, self._model_hparams.hidden_size].
+        [batch, p0, p1, model_hparams.hidden_size].
+      targets: Unused.
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
 
     Returns:
       logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].
     """
-    if self._model_hparams.symbol_modality_skip_top:
+    del targets  # unused arg
+    if model_hparams.symbol_modality_skip_top:
       return tf.expand_dims(body_output, 3)
 
-    if self._model_hparams.shared_embedding_and_softmax_weights:
+    if model_hparams.shared_embedding_and_softmax_weights:
       scope_name = "shared"
       reuse = tf.AUTO_REUSE
     else:
@@ -154,9 +171,9 @@ def top(self, body_output, _):
       reuse = False
     with tf.variable_scope(scope_name, reuse=reuse):
       body_output_shape = common_layers.shape_list(body_output)
-      var = self._get_weights(body_output_shape[-1])
-      if (self._model_hparams.factored_logits and
-          self._model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
+      var = cls._get_weights(model_hparams, vocab_size, body_output_shape[-1])
+      if (model_hparams.factored_logits and
+          model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
         # insert channels dimension
         body_output = tf.expand_dims(body_output, 3)
         return common_layers.FactoredTensor(body_output, var)
@@ -164,31 +181,35 @@ def top(self, body_output, _):
         body_output = tf.reshape(body_output, [-1, body_output_shape[-1]])
         logits = tf.matmul(body_output, var, transpose_b=True)
         return tf.reshape(logits,
-                          body_output_shape[:-1] + [1, self._vocab_size])
+                          body_output_shape[:-1] + [1, vocab_size])
 
 
 class SymbolModalityWeightsAll(SymbolModality):
   """SymbolModality for features that do not have 0-padding."""
 
-  @property
-  def targets_weights_fn(self):
+  @staticmethod
+  def targets_weights_fn(model_hparams):
     return common_layers.weights_all
 
 
 class SymbolModalityOneHot(SymbolModality):
   """Simple SymbolModality with one hot as embeddings."""
 
-  def bottom(self, x):
-    return tf.one_hot(x, self._vocab_size)
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size):
+    return tf.one_hot(x, vocab_size)
 
-  def targets_bottom(self, x):
-    return tf.one_hot(x, self._vocab_size)
+  @staticmethod
+  def targets_bottom(x, model_hparams, vocab_size):
+    return tf.one_hot(x, vocab_size)
 
-  def top(self, body_output, _):
+  @staticmethod
+  def top(body_output, _, model_hparams, vocab_size):
     return body_output
 
-  def loss(self, top_out, targets):
-    labels = tf.one_hot(targets, self._vocab_size)
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size):
+    labels = tf.one_hot(targets, vocab_size)
     loss = tf.nn.softmax_cross_entropy_with_logits(
         logits=top_out, labels=labels)
     return tf.reduce_mean(loss), tf.constant(1.0)
@@ -197,7 +218,8 @@ def loss(self, top_out, targets):
 class CTCSymbolModality(SymbolModality):
   """SymbolModality that uses CTC loss."""
 
-  def loss(self, top_out, targets):
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size):
     """Compute the CTC loss."""
     logits = top_out
     with tf.name_scope("ctc_loss", values=[logits, targets]):
@@ -219,24 +241,26 @@ def loss(self, top_out, targets):
           time_major=False,
           preprocess_collapse_repeated=False,
           ctc_merge_repeated=False)
-      weights = self.targets_weights_fn(targets)  # pylint: disable=not-callable
+      weights = cls.targets_weights_fn(targets)  # pylint: disable=not-callable
       return tf.reduce_sum(xent), tf.reduce_sum(weights)
 
 
 class ImageModality(modality.Modality):
   """Modality for images."""
-  PIXEL_EMBEDDING_SIZE = 64
 
-  def bottom(self, x):
-    with tf.variable_scope(self.name):
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
       if not tf.executing_eagerly():
         tf.summary.image(
             "inputs", common_layers.tpu_safe_image_summary(x), max_outputs=2)
       return tf.to_float(x)
 
-  def targets_bottom(self, x):
+  @classmethod
+  def targets_bottom(cls, x, model_hparams, vocab_size):
+    pixel_embedding_size = 64
     inputs = x
-    with tf.variable_scope(self.name):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
       if not tf.executing_eagerly():
         tf.summary.image(
             "targets_bottom",
@@ -246,30 +270,31 @@ def targets_bottom(self, x):
       if len(inputs_shape) != 4:
         raise ValueError("Assuming images given as int tensors in the format "
                          "[batch, height, width, channels] (256 values).")
-      # We embed each of 256=self._vocab_size possible pixel values.
+      # We embed each of 256=vocab_size possible pixel values.
       embedding_var = tf.get_variable(
           "pixel_embedding",
-          [self._vocab_size, self.PIXEL_EMBEDDING_SIZE])
-      hot_inputs = tf.one_hot(tf.to_int32(inputs), self._vocab_size)
-      hot_inputs = tf.reshape(hot_inputs, [-1, self._vocab_size])
+          [vocab_size, pixel_embedding_size])
+      hot_inputs = tf.one_hot(tf.to_int32(inputs), vocab_size)
+      hot_inputs = tf.reshape(hot_inputs, [-1, vocab_size])
       embedded = tf.matmul(hot_inputs, embedding_var)
       # Let's now merge all channels that were embedded into a single vector.
-      merged_size = self.PIXEL_EMBEDDING_SIZE * inputs_shape[3]
+      merged_size = pixel_embedding_size * inputs_shape[3]
       embedded = tf.reshape(embedded, inputs_shape[:3] + [merged_size])
       merged = tf.layers.dense(
           embedded,
-          self._model_hparams.hidden_size,
+          model_hparams.hidden_size,
           name="merge_pixel_embedded_channels")
       return merged
 
-  def top(self, body_output, _):
+  @staticmethod
+  def top(body_output, _, model_hparams, vocab_size):
     # TODO(lukaszkaiser): is this a universal enough way to get channels?
-    num_channels = self._model_hparams.problem.num_channels
+    num_channels = model_hparams.problem.num_channels
     with tf.variable_scope("rgb_softmax"):
       body_output_shape = common_layers.shape_list(body_output)
       reshape_shape = body_output_shape[:3]
-      reshape_shape.extend([num_channels, self._vocab_size])
-      res = tf.layers.dense(body_output, self._vocab_size * num_channels)
+      reshape_shape.extend([num_channels, vocab_size])
+      res = tf.layers.dense(body_output, vocab_size * num_channels)
       res = tf.reshape(res, reshape_shape)
       if not tf.get_variable_scope().reuse:
         res_argmax = tf.argmax(res, axis=-1)
@@ -279,26 +304,24 @@ def top(self, body_output, _):
             max_outputs=1)
       return res
 
-  def loss(self, top_out, targets):
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size):
     """Compute loss numerator and denominator for one shard of output."""
     logits = top_out
-    cutoff = getattr(self._model_hparams, "video_modality_loss_cutoff", 0.0)
+    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.0)
     return common_layers.padded_cross_entropy(
         logits,
         targets,
-        self._model_hparams.label_smoothing,
+        model_hparams.label_smoothing,
         cutoff=cutoff,
-        weights_fn=self.targets_weights_fn)
+        weights_fn=cls.targets_weights_fn(model_hparams))
 
 
 class ImageChannelCompressModality(modality.Modality):
   """Modality for images using channel compression for generation."""
 
-  @property
-  def num_channels(self):
-    return 3
-
-  def bottom_compress(self, inputs, name="bottom"):
+  @staticmethod
+  def bottom_compress(inputs, model_hparams, name="bottom"):
     """Compresses channel-wise input pixels into whole pixel representions.
 
     Perform conversion of RGB pixel values to a real number in the range -1 to
@@ -308,15 +331,17 @@ def bottom_compress(self, inputs, name="bottom"):
     Args:
       inputs: Tensor representing RGB pixel intensities as integers, of shape
         [batch, img_len, img_len, channels].
+      model_hparams: tf.HParams, model hyperparmeters.
       name: string, scope.
 
     Returns:
       body_input: Tensor of shape
-        [batch, img_len, img_len, self._model_hparams.hidden_size].
+        [batch, img_len, img_len, model_hparams.hidden_size].
     """
+    num_channels = 3
     with tf.variable_scope(name):
       inputs = tf.to_float(inputs)
-      hp = self._model_hparams
+      hp = model_hparams
       if hp.mode != tf.estimator.ModeKeys.PREDICT:
         tf.summary.image(
             "inputs",
@@ -332,33 +357,38 @@ def bottom_compress(self, inputs, name="bottom"):
       # Compress RGB intensities for each pixel using a convolution.
       outputs = tf.layers.conv2d(
           inputs,
-          self._model_hparams.hidden_size,
-          kernel_size=(1, self.num_channels),
+          model_hparams.hidden_size,
+          kernel_size=(1, num_channels),
           padding="VALID",
-          strides=(1, self.num_channels),
+          strides=(1, num_channels),
           activation=tf.nn.relu,
           name="conv_input")
       return outputs
 
-  def bottom(self, x):
-    return self.bottom_compress(x, "input_bottom")
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    return cls.bottom_compress(x, model_hparams, "input_bottom")
 
-  def targets_bottom(self, x):
-    return self.bottom_compress(x, "output_bottom")
+  @classmethod
+  def targets_bottom(cls, x, model_hparams, vocab_size):
+    return cls.bottom_compress(x, model_hparams, "output_bottom")
 
-  def top(self, body_output, _):
+  @classmethod
+  def top(cls, body_output, _, model_hparams, vocab_size):
     """Transforms body output to return logits.
 
     Args:
       body_output: Tensor of shape [batch, img_len, img_len, depth].
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
 
     Returns:
       Tensor of shape [batch, img_len, img_len, channels, vocab_size].
     """
-    with tf.variable_scope(self.name):
-      hidden_size = self._model_hparams.hidden_size
-      img_len = self._model_hparams.img_len
-      channels = self.num_channels  # RGB
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      hidden_size = model_hparams.hidden_size
+      img_len = model_hparams.img_len
+      channels = 3  # RGB
       batch = common_layers.shape_list(body_output)[0]
       x = tf.layers.conv2d(
           body_output,
@@ -369,28 +399,29 @@ def top(self, body_output, _):
           activation=tf.nn.relu,
           name="decompress_conv")
       x = tf.reshape(x, [batch, img_len, img_len * channels, hidden_size])
-      x = common_layers.layer_preprocess(x, self._model_hparams)
+      x = common_layers.layer_preprocess(x, model_hparams)
       x = tf.layers.dense(x,
-                          self._vocab_size,
+                          vocab_size,
                           use_bias=True,
                           activation=None,
                           name="output_conv")
       x = tf.reshape(
-          x, [batch, img_len, img_len, channels, self._vocab_size])
+          x, [batch, img_len, img_len, channels, vocab_size])
       return x
 
 
 class ImageChannelBottomIdentityModality(ImageChannelCompressModality):
 
-  def top(self, body_output, _):
+  @staticmethod
+  def top(body_output, _, model_hparams, vocab_size):
     return body_output
 
 
 class ImageChannelEmbeddingsBottom(modality.Modality):
   """Modality for images using channel compression for generation."""
 
-  def get_channel_embeddings(self,
-                             io_depth,
+  @staticmethod
+  def get_channel_embeddings(io_depth,
                              targets,
                              hidden_size,
                              name="channel"):
@@ -410,42 +441,47 @@ def get_channel_embeddings(self,
 
     return tf.concat(channel_target_embs, axis=-1)
 
-  def targets_bottom(self, x):
+  @staticmethod
+  def targets_bottom(x, model_hparams, vocab_size):
     inputs = x
-    io_depth = self._model_hparams.num_channels
+    io_depth = model_hparams.num_channels
     tshape = common_layers.shape_list(inputs)
-    hidden_size = self._model_hparams.hidden_size
-    target_embeddings = self.get_channel_embeddings(io_depth, inputs,
-                                                    hidden_size, "input_bottom")
+    hidden_size = model_hparams.hidden_size
+    target_embeddings = ImageChannelEmbeddingsBottom.get_channel_embeddings(
+        io_depth, inputs, hidden_size, "input_bottom")
     return tf.reshape(target_embeddings,
                       [tshape[0], tshape[1], tshape[2] * io_depth, hidden_size])
 
-  def top(self, body_output, _):
-    with tf.variable_scope(self.name):
-      img_len = self._model_hparams.img_len
-      channels = self._model_hparams.num_channels
+  @classmethod
+  def top(cls, body_output, _, model_hparams, vocab_size):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      img_len = model_hparams.img_len
+      channels = model_hparams.num_channels
       x = tf.layers.dense(
           body_output, 256, use_bias=True, activation=None, name="output_conv")
       x = tf.reshape(x,
-                     [-1, img_len, img_len, channels, self._vocab_size])
+                     [-1, img_len, img_len, channels, vocab_size])
       return x
 
 
 class AudioModality(modality.Modality):
   """Performs strided conv compressions for audio data."""
 
-  def bottom(self, x):
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
     """Transform input from data space to model space.
 
     Args:
       x: A Tensor with shape [batch, ...]
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
 
     Returns:
       body_input: A Tensor with shape [batch, ?, ?,
-        self._model_hparams.hidden_size].
+        model_hparams.hidden_size].
     """
     inputs = x
-    with tf.variable_scope(self.name):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
       # TODO(aidangomez): Will need to sort out a better audio pipeline
       def xnet_resblock(x, filters, res_relu, name):
         """Xception block."""
@@ -471,10 +507,10 @@ def xnet_resblock(x, filters, res_relu, name):
 
       x = tf.to_float(inputs) / 255.
       x.set_shape([None, None, None, 1])
-      for i in range(self._model_hparams.audio_compression):
+      for i in range(model_hparams.audio_compression):
         x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
       return xnet_resblock(x,
-                           self._model_hparams.hidden_size,
+                           model_hparams.hidden_size,
                            False,
                            "compress_block_final")
 
@@ -482,18 +518,21 @@ def xnet_resblock(x, filters, res_relu, name):
 class AudioSpectralModality(modality.Modality):
   """Performs strided conv compressions for audio spectral data."""
 
-  def bottom(self, x):
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
     """Transform input from data space to model space.
 
     Args:
       x: A Tensor with shape [batch, ...]
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
 
     Returns:
       body_input: A Tensor with shape [batch, ?, ?,
-        self._model_hparams.hidden_size].
+        model_hparams.hidden_size].
     """
     inputs = x
-    with tf.variable_scope(self.name):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
       # TODO(aidangomez): Will need to sort out a better audio pipeline
       def xnet_resblock(x, filters, res_relu, name):
         """Xception-like block."""
@@ -520,10 +559,10 @@ def xnet_resblock(x, filters, res_relu, name):
       # Bitcast back from int32
       x = tf.bitcast(inputs, tf.float32)
       x.set_shape([None, None, None, 1])
-      for i in range(self._model_hparams.audio_compression):
+      for i in range(model_hparams.audio_compression):
         x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
       return xnet_resblock(x,
-                           self._model_hparams.hidden_size,
+                           model_hparams.hidden_size,
                            False,
                            "compress_block_final")
 
@@ -531,22 +570,25 @@ def xnet_resblock(x, filters, res_relu, name):
 class SpeechRecognitionModality(modality.Modality):
   """Common ASR filterbank processing."""
 
-  def bottom(self, x):
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
     """Use batchnorm instead of CMVN and shorten the stft with strided convs.
 
     Args:
       x: float32 tensor with shape [batch_size, len, 1, freqs * channels]
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
 
     Returns:
       float32 tensor with shape [batch_size, shorter_len, 1, hidden_size]
     """
     inputs = x
-    p = self._model_hparams
+    p = model_hparams
 
     num_mel_bins = p.audio_num_mel_bins
     num_channels = 3 if p.audio_add_delta_deltas else 1
 
-    with tf.variable_scope(self.name):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
       if p.audio_preproc_in_bottom:
         # Compute filterbanks
         with tf.variable_scope("fbanks"):
@@ -613,102 +655,110 @@ def bottom(self, x):
 class VideoModality(modality.Modality):
   """Modality for videos, i.e., time-sequences of frames."""
 
-  def bottom(self, x):
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size):
     common_video.gif_summary("inputs", x, max_outputs=1)
     x = common_layers.standardize_images(x)
     return x
 
-  def targets_bottom(self, x):
+  @staticmethod
+  def targets_bottom(x, model_hparams, vocab_size):
     common_video.gif_summary("targets", x, max_outputs=1)
     x = common_layers.standardize_images(x)
     return x
 
-  def top(self, body_output, targets):
-    num_channels = self._model_hparams.problem.num_channels
+  @staticmethod
+  def top(body_output, targets, model_hparams, vocab_size):
+    num_channels = model_hparams.problem.num_channels
     shape = common_layers.shape_list(body_output)
-    reshape_shape = shape[:-1] + [num_channels, self._vocab_size]
+    reshape_shape = shape[:-1] + [num_channels, vocab_size]
     res = tf.reshape(body_output, reshape_shape)
     # Calculate argmax so as to have a summary with the produced images.
-    x = tf.argmax(tf.reshape(res, [-1, self._vocab_size]), axis=-1)
+    x = tf.argmax(tf.reshape(res, [-1, vocab_size]), axis=-1)
     x = tf.reshape(x, shape[:-1] + [num_channels])
     common_video.gif_summary("results", x, max_outputs=1)
     return res
 
-  def loss(self, top_out, targets):
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size):
     """Compute loss numerator and denominator for one shard of output."""
     logits = top_out
     logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
     targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-    cutoff = getattr(self._model_hparams, "video_modality_loss_cutoff", 0.01)
+    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
     return common_layers.padded_cross_entropy(
         logits,
         targets,
-        self._model_hparams.label_smoothing,
+        model_hparams.label_smoothing,
         cutoff=cutoff,
-        weights_fn=self.targets_weights_fn)
+        weights_fn=cls.targets_weights_fn(model_hparams))
 
 
 class VideoModalityBitwise(VideoModality):
   """Video Modality where bottom embeds pixels bitwise."""
-  PIXEL_EMBEDDING_SIZE = 64
 
-  def bottom(self, x):
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    pixel_embedding_size = 64
     inputs = x
-    with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size),
+                           reuse=tf.AUTO_REUSE):
       common_layers.summarize_video(inputs, "bottom")
       # Embed bitwise.
-      assert self._vocab_size == 256
+      assert vocab_size == 256
       embedded = discretization.int_to_bit_embed(inputs, 8,
-                                                 self.PIXEL_EMBEDDING_SIZE)
+                                                 pixel_embedding_size)
       # Project.
       return tf.layers.dense(
           embedded,
-          self._model_hparams.hidden_size,
+          model_hparams.hidden_size,
           name="merge_pixel_embedded_frames")
 
-  def targets_bottom(self, x):  # pylint: disable=arguments-differ
+  @classmethod
+  def targets_bottom(cls, x, model_hparams, vocab_size):  # pylint: disable=arguments-differ
+    pixel_embedding_size = 64
     inputs = x
-    with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size),
+                           reuse=tf.AUTO_REUSE):
       common_layers.summarize_video(inputs, "targets_bottom")
       # Embed bitwise.
-      assert self._vocab_size == 256
+      assert vocab_size == 256
       embedded = discretization.int_to_bit_embed(inputs, 8,
-                                                 self.PIXEL_EMBEDDING_SIZE)
+                                                 pixel_embedding_size)
       # Transpose and project.
       transposed = common_layers.time_to_channels(embedded)
       return tf.layers.dense(
           transposed,
-          self._model_hparams.hidden_size,
+          model_hparams.hidden_size,
           name="merge_pixel_embedded_frames")
 
 
 class VideoModalityPixelNoise(VideoModality):
   """Video modality that introduces pixel noise on input during training."""
 
-  def bottom(self, x):
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    input_noise = getattr(model_hparams, "video_modality_input_noise", 0.25)
     inputs = x
-    if self._model_hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    if model_hparams.mode == tf.estimator.ModeKeys.TRAIN:
       background = tfp.distributions.percentile(inputs, 50., axis=[0, 1, 2, 3])
       input_shape = common_layers.shape_list(inputs)
       input_size = tf.reduce_prod(input_shape[:-1])
       input_mask = tf.multinomial(
-          tf.log([[self.input_noise, 1.-self.input_noise]]), input_size)
+          tf.log([[input_noise, 1.-input_noise]]), input_size)
       input_mask = tf.reshape(tf.cast(input_mask, tf.int32),
                               input_shape[:-1]+[1])
       inputs = inputs * input_mask + background * (1 - input_mask)
-    return super(VideoModalityPixelNoise, self).bottom(inputs)
-
-  @property
-  def input_noise(self):
-    return getattr(self._model_hparams, "video_modality_input_noise", 0.25)
+    return cls.bottom(inputs, model_hparams, vocab_size)
 
 
 class VideoModalityL1(VideoModality):
   """Video modality that predicts a scalar per channel with an L1 loss."""
 
-  def top(self, body_output, _):
-    num_channels = self._model_hparams.problem.num_channels
-    num_frames = self._model_hparams.video_num_target_frames
+  @staticmethod
+  def top(body_output, _, model_hparams, vocab_size):
+    num_channels = model_hparams.problem.num_channels
+    num_frames = model_hparams.video_num_target_frames
     with tf.variable_scope("rgb"):
       body_output_shape = common_layers.shape_list(body_output)
       res = tf.layers.dense(body_output, num_channels * num_frames, name="cast")
@@ -722,55 +772,61 @@ def top(self, body_output, _):
             max_outputs=1)
       return tf.expand_dims(res, axis=-1)  # Add an axis like in perplexity.
 
-  @property
-  def cutoff(self):
-    return getattr(self._model_hparams, "video_modality_loss_cutoff", 0.2)
-
-  def internal_loss(self, logits, targets):
-    return tf.nn.relu(tf.abs(logits - targets) - self.cutoff)
+  @staticmethod
+  def internal_loss(logits, targets, model_hparams):
+    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
+    return tf.nn.relu(tf.abs(logits - targets) - cutoff)
 
-  def loss(self, top_out, targets):
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size):
     """Compute loss numerator and denominator for one shard of output."""
     logits = top_out
     logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:-1])
     targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-    weights = self.targets_weights_fn(targets)
+    targets_weights_fn = VideoModalityL1.targets_weights_fn(model_hparams)
+    weights = targets_weights_fn(targets)
     # Shift targets by 0.5 so later just casting to int gives the prediction.
     # So for int targets, say 0 and 7, we actually train to predict 0.5 and 7.5.
     # Later (in merics or infer) this is cast to int anyway. Also, we have no
-    # loss beyond self.cutoff = 0.2 as these are already correct predictions.
+    # loss beyond cutoff = 0.2 as these are already correct predictions.
     targets = tf.to_float(targets) + 0.5
-    loss = self.internal_loss(logits, targets)
+    loss = VideoModalityL1.internal_loss(logits, targets, model_hparams)
     return tf.reduce_sum(loss * weights), tf.reduce_sum(weights)
 
 
 class VideoModalityL2(VideoModalityL1):
   """Modality for videos with L2 loss."""
 
-  def internal_loss(self, logits, targets):
+  @staticmethod
+  def internal_loss(logits, targets, model_hparams):
+    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
     return tf.nn.relu(
-        tf.squared_difference(logits, targets) - self.cutoff * self.cutoff)
+        tf.squared_difference(logits, targets) - cutoff * cutoff)
 
 
 class VideoModalityL2Raw(VideoModalityL2):
   """Modality with L2 loss and raw input (sequences of frames)."""
 
-  def convert_rgb_to_real(self, prediction, targets):
+  @staticmethod
+  def convert_rgb_to_real(prediction, targets):
     """Convert prediction and target from rgb to real."""
     prediction = tf.squeeze(prediction, axis=-1)
     prediction = common_layers.convert_rgb_to_real(prediction)
     targets = common_layers.convert_rgb_to_real(targets)
     return prediction, targets
 
-  def bottom(self, x):
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size):
     common_video.gif_summary("inputs", x)
     return common_layers.convert_rgb_to_real(x)
 
-  def targets_bottom(self, x):  # pylint: disable=arguments-differ
+  @staticmethod
+  def targets_bottom(x, model_hparams, vocab_size):  # pylint: disable=arguments-differ
     common_video.gif_summary("targets_bottom", x)
     return common_layers.convert_rgb_to_real(x)
 
-  def top(self, body_output, _):
+  @staticmethod
+  def top(body_output, _, model_hparams, vocab_size):
     frames = body_output
     if isinstance(body_output, list):
       frames = tf.stack(body_output, axis=1)
@@ -778,8 +834,9 @@ def top(self, body_output, _):
     common_video.gif_summary("body_output", rgb_frames)
     return tf.expand_dims(rgb_frames, axis=-1)
 
-  def loss(self, top_out, targets):
-    prediction, groundtruth = self.convert_rgb_to_real(top_out, targets)
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size):
+    prediction, groundtruth = cls.convert_rgb_to_real(top_out, targets)
     loss = tf.losses.mean_squared_error(prediction, groundtruth)
     return loss, tf.constant(1.0)
 
@@ -787,8 +844,9 @@ def loss(self, top_out, targets):
 class VideoModalityL1Raw(VideoModalityL2Raw):
   """Modality with L1 loss and raw input (sequences of frames)."""
 
-  def loss(self, top_out, targets):
-    prediction, groundtruth = self.convert_rgb_to_real(top_out, targets)
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size):
+    prediction, groundtruth = cls.convert_rgb_to_real(top_out, targets)
     loss = tf.losses.absolute_difference(prediction, groundtruth)
     return loss, tf.constant(1.0)
 
@@ -796,84 +854,94 @@ def loss(self, top_out, targets):
 class ClassLabelModality(modality.Modality):
   """Used for label data."""
 
-  @property
-  def name(self):
-    return "class_label_modality_%d_%d" % (self._vocab_size,
-                                           self._model_hparams.hidden_size)
+  @staticmethod
+  def name(model_hparams, vocab_size):
+    return "class_label_modality_%d_%d" % (vocab_size,
+                                           model_hparams.hidden_size)
 
-  def bottom(self, x):
-    with tf.variable_scope(self.name):
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
       multiplier = 1.0
-      if self._model_hparams.multiply_embedding_mode == "sqrt_depth":
-        multiplier = self._model_hparams.hidden_size**0.5
+      if model_hparams.multiply_embedding_mode == "sqrt_depth":
+        multiplier = model_hparams.hidden_size**0.5
       return common_layers.embedding(x,
-                                     self._vocab_size,
-                                     self._model_hparams.hidden_size,
+                                     vocab_size,
+                                     model_hparams.hidden_size,
                                      multiplier=multiplier)
 
-  def targets_bottom(self, x):
-    with tf.variable_scope(self.name):
+  @classmethod
+  def targets_bottom(cls, x, model_hparams, vocab_size):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
       return tf.zeros([common_layers.shape_list(x)[0],
                        1,
                        1,
-                       self._model_hparams.hidden_size])
+                       model_hparams.hidden_size])
 
-  def top(self, body_output, _):
+  @classmethod
+  def top(cls, body_output, _, model_hparams, vocab_size):
     """Transform inputs from model space to target space.
 
     Average over inner dims and a linear layer to logits.
 
     Args:
       body_output: A Tensor with shape [batch, ?, ?, body_output_size].
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
 
     Returns:
       a Tensors, each with shape [batch_size, 1, 1, 1, vocab_size]
     """
-    with tf.variable_scope(self.name):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
       x = body_output
       x = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
-      res = tf.layers.dense(x, self._vocab_size)
+      res = tf.layers.dense(x, vocab_size)
       return tf.expand_dims(res, 3)
 
 
 class VideoModalityIdentity(VideoModality):
   """Video Modality where top and bottom is an identity function."""
 
-  def bottom(self, x):
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size):
     common_video.gif_summary("inputs", x, max_outputs=1)
     return x
 
-  def targets_bottom(self, x):
+  @staticmethod
+  def targets_bottom(x, model_hparams, vocab_size):
     common_video.gif_summary("targets", x, max_outputs=1)
     return x
 
-  def top(self, body_output, targets):
+  @staticmethod
+  def top(body_output, targets, model_hparams, vocab_size):
     return body_output
 
-  def loss(self, top_out, targets):
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size):
     """Compute loss numerator and denominator for one shard of output."""
     # TODO(nikip): Try L2 loss
     logits = top_out
     logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
     targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-    cutoff = getattr(self._model_hparams, "video_modality_loss_cutoff", 0.01)
+    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
     return common_layers.padded_cross_entropy(
         logits,
         targets,
-        self._model_hparams.label_smoothing,
+        model_hparams.label_smoothing,
         cutoff=cutoff,
-        weights_fn=self.targets_weights_fn)
+        weights_fn=VideoModalityIdentity.targets_weights_fn(model_hparams))
 
 
 class MultiLabelModality(ClassLabelModality):
   """Used for multi label task."""
 
-  @property
-  def targets_weights_fn(self):
+  @staticmethod
+  def targets_weights_fn(model_hparams):
     """Target weight function for multi label, defaults to nonzero labels."""
     return common_layers.weights_nonzero
 
-  def loss(self, top_out, targets):
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size):
     """Average loss over the labels."""
     logits = top_out
     num_labels = tf.shape(targets)[1]
@@ -882,8 +950,8 @@ def loss(self, top_out, targets):
     xent, weights = common_layers.padded_cross_entropy(
         logits,
         targets,
-        self._model_hparams.label_smoothing,
-        weights_fn=self.targets_weights_fn,
+        model_hparams.label_smoothing,
+        weights_fn=cls.targets_weights_fn(model_hparams),
         reduce_sum=False,
     )
     xent = tf.squeeze(xent, [2, 3])
@@ -900,18 +968,23 @@ def loss(self, top_out, targets):
 class OneHotClassLabelModality(ClassLabelModality):
   """Used for one-hot encoded class labels."""
 
-  def loss(self, top_out, targets):
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size):
     """Apply softmax cross-entropy between outputs and targets.
 
     Args:
       top_out: logits Tensor with shape [batch, ?, ?, num_classes]
       targets: one-hot encoding Tensor with shape [batch, ?, ?, num_classes]
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+
     Returns:
       loss_scale (cross-entropy), loss_denom
     """
     loss_scale = tf.losses.softmax_cross_entropy(
         onehot_labels=targets, logits=top_out)
-    weights = self.targets_weights_fn(targets)
+    targets_weights_fn = cls.targets_weights_fn(model_hparams)
+    weights = targets_weights_fn(targets)
     loss_denom = tf.reduce_sum(weights)
     return loss_scale, loss_denom
 
@@ -919,20 +992,24 @@ def loss(self, top_out, targets):
 class IdentityModality(modality.Modality):
   """Does nothing."""
 
-  def bottom(self, x):
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size):
     return tf.to_float(x)
 
-  def top(self, body_output, _):
+  @staticmethod
+  def top(body_output, _, model_hparams, vocab_size):
     return body_output
 
 
 class GenericL2LossModality(IdentityModality):
   """Generic modality with L2 as Loss."""
 
-  def targets_bottom(self, x):
+  @staticmethod
+  def targets_bottom(x, model_hparams, vocab_size):
     return tf.to_float(x)
 
-  def loss(self, body_output, targets):
+  @staticmethod
+  def loss(body_output, targets, model_hparams, vocab_size):
     loss = tf.squared_difference(body_output, tf.to_float(targets))
     return tf.reduce_mean(loss), tf.constant(1.0)
 
@@ -944,33 +1021,38 @@ class RealModality(modality.Modality):
   * Top is a linear projection layer to vocab_size.
   """
 
-  @property
-  def top_is_pointwise(self):
+  @staticmethod
+  def top_is_pointwise():
     return True
 
-  def bottom(self, x):
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size):
     with tf.variable_scope("real"):
       return tf.layers.dense(
-          tf.to_float(x), self._model_hparams.hidden_size, name="bottom")
+          tf.to_float(x), model_hparams.hidden_size, name="bottom")
 
-  def top(self, body_output, _):
+  @staticmethod
+  def top(body_output, _, model_hparams, vocab_size):
     with tf.variable_scope("real"):
-      return tf.layers.dense(body_output, self._vocab_size, name="top")
+      return tf.layers.dense(body_output, vocab_size, name="top")
 
-  def loss(self, top_out, targets):
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size):
     raise NotImplementedError()
 
 
 class RealL2LossModality(RealModality):
   """Modality for real (i.e. float) vectors with L2 (Gaussian) loss."""
 
-  def loss(self, top_out, targets):
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size):
     predictions = top_out
     if (len(common_layers.shape_list(top_out)) != len(
         common_layers.shape_list(targets))):
       predictions = tf.squeeze(top_out, axis=[-1])
     with tf.name_scope("l2"):
-      weights = self.targets_weights_fn(targets)
+      targets_weights_fn = cls.targets_weights_fn(model_hparams)
+      weights = targets_weights_fn(targets)
       l2 = tf.pow(predictions - targets, 2)
       return tf.reduce_sum(l2 * weights), tf.reduce_sum(weights)
 
@@ -978,13 +1060,15 @@ def loss(self, top_out, targets):
 class RealLogPoissonLossModality(RealModality):
   """Modality for real (i.e. float) vectors with log Poisson regression loss."""
 
-  def loss(self, top_out, targets):
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size):
     predictions = top_out
     if (len(common_layers.shape_list(top_out)) != len(
         common_layers.shape_list(targets))):
       predictions = tf.squeeze(top_out, axis=[-1])
     with tf.name_scope("log_possion"):
-      weights = self.targets_weights_fn(targets)
+      targets_weights_fn = cls.targets_weights_fn(model_hparams)
+      weights = targets_weights_fn(targets)
       lp_loss = tf.nn.log_poisson_loss(targets, predictions)
       return tf.reduce_sum(lp_loss * weights), tf.reduce_sum(weights)
 
@@ -995,18 +1079,21 @@ class IdentitySymbolModality(SymbolModality):
   Uses the weights_fn from SymbolModality so that loss/metrics ignore padding.
   """
 
-  def bottom(self, x):
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size):
     return tf.to_float(x)
 
-  def top(self, body_output, _):
+  @staticmethod
+  def top(body_output, _, model_hparams, vocab_size):
     return body_output
 
-  def targets_bottom(self, x):
+  @classmethod
+  def targets_bottom(cls, x, model_hparams, vocab_size):
     """SymbolModality overrides targets_bottom, so need to override here too."""
-    return self.bottom(x)
+    return cls.bottom(x, model_hparams, vocab_size)
 
-  @property
-  def top_is_pointwise(self):
+  @staticmethod
+  def top_is_pointwise():
     # pointwise mode manipulates body output, not logits, so it fails here.
     return False
 
@@ -1014,18 +1101,20 @@ def top_is_pointwise(self):
 class SigmoidClassLabelModality(ClassLabelModality):
   """Sigmoid cross-entropy for independent class labels."""
 
-  @property
-  def name(self):
-    return "sigmoid_class_symbol_modality_%d_%d" % (
-        self._vocab_size, self._model_hparams.hidden_size)
+  @staticmethod
+  def name(model_hparams, vocab_size):
+    return "sigmoid_class_symbol_modality_%d_%d" % (vocab_size,
+                                                    model_hparams.hidden_size)
 
-  def loss(self, top_out, targets):
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size):
     # Expect inputs of size [batch-size, timesteps, 1, num-classes], where the
     # last dimension of num-classes represents logits for binary labels
     loss_scale = tf.losses.sigmoid_cross_entropy(
         multi_class_labels=targets, logits=top_out)
     # Weigh all classes equally
-    weights = self.targets_weights_fn(targets)
+    targets_weights_fn = cls.targets_weights_fn(model_hparams)
+    weights = targets_weights_fn(targets)
     loss_denom = tf.reduce_sum(weights)
     return loss_scale, loss_denom
 
@@ -1033,34 +1122,38 @@ def loss(self, top_out, targets):
 class SigmoidMaxPoolingClassLabelModality(ClassLabelModality):
   """Sigmoid cross-entropy applied on max-pooling over timesteps."""
 
-  @property
-  def name(self):
+  @staticmethod
+  def name(model_hparams, vocab_size):
     return "sigmoid_max_pooling_class_symbol_modality_%d_%d" % (
-        self._vocab_size, self._model_hparams.hidden_size)
+        vocab_size, model_hparams.hidden_size)
 
-  def top(self, body_output, _):
+  @classmethod
+  def top(cls, body_output, _, model_hparams, vocab_size):
     """Transform inputs from model space to target space.
 
     Average over inner dims and a linear layer to logits.
 
     Args:
       body_output: A Tensor with shape [batch, timesteps, 1, body_output_size].
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
 
     Returns:
       a Tensors, each with shape [batch_size, 1, 1, vocab_size]
     """
-    with tf.variable_scope(self.name):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
       x = body_output
       x = tf.reduce_max(x, axis=1, keepdims=True)
-      return tf.layers.dense(x, self._vocab_size)
+      return tf.layers.dense(x, vocab_size)
 
-  def loss(self, top_out, targets):
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size):
     # Expect inputs of size [batch-size, 1, 1, num-classes], where the
     # last dimension of num-classes represents logits for binary labels
     loss_scale = tf.losses.sigmoid_cross_entropy(
         multi_class_labels=targets, logits=top_out)
     # Weigh all classes equally
-    weights = self.targets_weights_fn(targets)
+    weights = cls.targets_weights_fn(model_hparams)(targets)
     loss_denom = tf.reduce_sum(weights)
     return loss_scale, loss_denom
 
@@ -1068,46 +1161,49 @@ def loss(self, top_out, targets):
 class SoftmaxMaxPoolingClassLabelModality(OneHotClassLabelModality):
   """Softmax cross-entropy applied on max-pooling over timesteps."""
 
-  @property
-  def name(self):
+  @staticmethod
+  def name(model_hparams, vocab_size):
     return "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
-        self._vocab_size, self._model_hparams.hidden_size)
+        vocab_size, model_hparams.hidden_size)
 
-  def top(self, body_output, _):
-    with tf.variable_scope(self.name):
+  @classmethod
+  def top(cls, body_output, _, model_hparams, vocab_size):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
       x = body_output
       x = tf.reduce_max(x, axis=1, keepdims=True)
-      return tf.layers.dense(x, self._vocab_size)
+      return tf.layers.dense(x, vocab_size)
 
 
 class SoftmaxAveragePoolingClassLabelModality(OneHotClassLabelModality):
   """Softmax cross-entropy applied on average-pooling over timesteps."""
 
-  @property
-  def name(self):
+  @staticmethod
+  def name(model_hparams, vocab_size):
     return "softmax_average_pooling_onehot_class_label_modality_%d_%d" % (
-        self._vocab_size, self._model_hparams.hidden_size)
+        vocab_size, model_hparams.hidden_size)
 
-  def top(self, body_output, _):
-    with tf.variable_scope(self.name):
+  @classmethod
+  def top(cls, body_output, _, model_hparams, vocab_size):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
       x = body_output
       x = tf.reduce_mean(x, axis=1, keepdims=True)
-      return tf.layers.dense(x, self._vocab_size)
+      return tf.layers.dense(x, vocab_size)
 
 
 class SoftmaxLastTimestepClassLabelModality(OneHotClassLabelModality):
   """Softmax cross-entropy applied on last-timestep encoding."""
 
-  @property
-  def name(self):
+  @staticmethod
+  def name(model_hparams, vocab_size):
     return "softmax_last_timestep_onehot_class_label_modality_%d_%d" % (
-        self._vocab_size, self._model_hparams.hidden_size)
+        vocab_size, model_hparams.hidden_size)
 
-  def top(self, body_output, _):
-    with tf.variable_scope(self.name):
+  @classmethod
+  def top(cls, body_output, _, model_hparams, vocab_size):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
       x = body_output
       x = tf.expand_dims(x[:, -1], 1)  # Pick the last timestep
-      return tf.layers.dense(x, self._vocab_size)
+      return tf.layers.dense(x, vocab_size)
 
 
 class ModalityType(object):
@@ -1184,3 +1280,63 @@ def get_choices():
         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
     ]
+
+
+# Utility functions, similar to tf.keras
+current_module = sys.modules[__name__]
+
+
+def get_bottom(modality_type, value=None):
+  """Gets default bottom transformation; if none available, return value."""
+  if modality_type in ModalityType.get_choices():
+    modality_cls = getattr(current_module, modality_type)
+    return modality_cls.bottom
+  return value
+
+
+def get_loss(modality_type, value=None):
+  """Gets default loss transformation; if none available, return value."""
+  if modality_type in ModalityType.get_choices():
+    modality_cls = getattr(current_module, modality_type)
+    return modality_cls.loss
+  return value
+
+
+def get_name(modality_type, value=None):
+  """Gets default name for transformations; if none available, return value."""
+  if modality_type in ModalityType.get_choices():
+    modality_cls = getattr(current_module, modality_type)
+    return modality_cls.name
+  return value
+
+
+def get_targets_bottom(modality_type, value=None):
+  """Gets default bottom transformation for targets; if none, return value."""
+  if modality_type in ModalityType.get_choices():
+    modality_cls = getattr(current_module, modality_type)
+    return modality_cls.targets_bottom
+  return value
+
+
+def get_targets_weights_fn(modality_type, value=None):
+  """Gets default weights function; if none available, return value."""
+  if modality_type in ModalityType.get_choices():
+    modality_cls = getattr(current_module, modality_type)
+    return modality_cls.targets_weights_fn
+  return value
+
+
+def get_top(modality_type, value=None):
+  """Gets default top transformation; if none available, return value."""
+  if modality_type in ModalityType.get_choices():
+    modality_cls = getattr(current_module, modality_type)
+    return modality_cls.top
+  return value
+
+
+def get_top_is_pointwise(modality_type, value=None):
+  """Gets whether default top is pointwise; if none available, return value."""
+  if modality_type in ModalityType.get_choices():
+    modality_cls = getattr(current_module, modality_type)
+    return modality_cls(None, None).top_is_pointwise
+  return value
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 603ef154f..bc9f8d0c2 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -43,11 +43,14 @@ def testSymbolModalityInputs(self):
     model_hparams.mode = tf.estimator.ModeKeys.TRAIN
     x = np.random.randint(
         vocab_size, size=(batch_size, length, 1, 1))
-    m = modalities.SymbolModality(model_hparams, vocab_size)
     data_parallelism = expert_utils.Parallelism(
         ["/device:CPU:0"] * num_datashards)
     xs = tf.split(x, num_datashards)
-    sharded_output = data_parallelism(m.bottom, xs)
+    sharded_output = data_parallelism(
+        modalities.get_bottom(modalities.ModalityType.SYMBOL),
+        xs,
+        model_hparams,
+        vocab_size)
     output = tf.concat(sharded_output, 0)
     self.evaluate(tf.global_variables_initializer())
     res = self.evaluate(output)
@@ -68,17 +71,22 @@ def testSymbolModalityTargets(self):
         100, size=(batch_size, length, height, hidden_size))
     targets = np.random.randint(
         vocab_size, size=(batch_size, length, height, 1))
-    m = modalities.SymbolModality(model_hparams, vocab_size)
     data_parallelism = expert_utils.Parallelism(
         ["/device:CPU:0"] * num_datashards)
     sharded_body_output = tf.split(tf.to_float(body_output), num_datashards)
     sharded_targets = tf.split(targets, num_datashards)
-    sharded_logits = data_parallelism(m.top,
-                                      sharded_body_output,
-                                      sharded_targets)
-    sharded_loss_num, sharded_loss_den = data_parallelism(m.loss,
-                                                          sharded_logits,
-                                                          sharded_targets)
+    sharded_logits = data_parallelism(
+        modalities.get_top(modalities.ModalityType.SYMBOL),
+        sharded_body_output,
+        sharded_targets,
+        model_hparams,
+        vocab_size)
+    sharded_loss_num, sharded_loss_den = data_parallelism(
+        modalities.get_loss(modalities.ModalityType.SYMBOL),
+        sharded_logits,
+        sharded_targets,
+        model_hparams,
+        vocab_size)
     train_loss = (tf.add_n(sharded_loss_num) /
                   tf.maximum(1.0, tf.add_n(sharded_loss_den)))
     logits = tf.concat(sharded_logits, 0)
@@ -103,18 +111,23 @@ def testSymbolModalityTargetsFactored(self):
         100, size=(batch_size, length, height, hidden_size))
     targets = np.random.randint(
         vocab_size, size=(batch_size, length, height, 1))
-    m = modalities.SymbolModality(model_hparams, vocab_size)
     data_parallelism = expert_utils.Parallelism(
         ["/device:CPU:0"] * num_datashards)
     with self.test_session() as session:
       sharded_body_output = tf.split(tf.to_float(body_output), num_datashards)
       sharded_targets = tf.split(targets, num_datashards)
-      sharded_logits = data_parallelism(m.top,
-                                        sharded_body_output,
-                                        sharded_targets)
-      sharded_loss_num, sharded_loss_den = data_parallelism(m.loss,
-                                                            sharded_logits,
-                                                            sharded_targets)
+      sharded_logits = data_parallelism(
+          modalities.get_top(modalities.ModalityType.SYMBOL),
+          sharded_body_output,
+          sharded_targets,
+          model_hparams,
+          vocab_size)
+      sharded_loss_num, sharded_loss_den = data_parallelism(
+          modalities.SymbolModality.loss,
+          sharded_logits,
+          sharded_targets,
+          model_hparams,
+          vocab_size)
       train_loss = (tf.add_n(sharded_loss_num) /
                     tf.maximum(1.0, tf.add_n(sharded_loss_den)))
       logits = tf.concat(sharded_logits, 0)
diff --git a/tensor2tensor/models/research/vqa_attention_test.py b/tensor2tensor/models/research/vqa_attention_test.py
index 6d8904a2a..da6fc0bc3 100644
--- a/tensor2tensor/models/research/vqa_attention_test.py
+++ b/tensor2tensor/models/research/vqa_attention_test.py
@@ -45,13 +45,12 @@ def testVqaAttentionBaseline(self):
         num_classes + 1, size=(batch_size, answer_length, 1, 1))
     hparams = vqa_attention.vqa_attention_base()
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
-                                                     vocab_size,
+                                                     num_classes + 1,
                                                      hparams)
-    p_hparams.modality["inputs"] = modalities.ImageModality(hparams)
-    p_hparams.modality["question"] = modalities.SymbolModality(
-        hparams, vocab_size)
-    p_hparams.modality["targets"] = modalities.MultiLabelModality(
-        hparams, num_classes + 1)
+    p_hparams.modality["inputs"] = modalities.ModalityType.IMAGE
+    p_hparams.modality["targets"] = modalities.ModalityType.MULTI_LABEL
+    p_hparams.modality["question"] = modalities.ModalityType.SYMBOL
+    p_hparams.vocab_size["question"] = vocab_size
     with self.test_session() as session:
       features = {
           "inputs": tf.constant(x, dtype=tf.float32),
diff --git a/tensor2tensor/models/resnet_test.py b/tensor2tensor/models/resnet_test.py
index 0e15dbbf4..17e930a79 100644
--- a/tensor2tensor/models/resnet_test.py
+++ b/tensor2tensor/models/resnet_test.py
@@ -48,9 +48,8 @@ def _test_resnet(self, img_size, output_size):
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
                                                      hparams)
-    p_hparams.modality["inputs"] = modalities.ImageModality(hparams)
-    p_hparams.modality["targets"] = modalities.ClassLabelModality(
-        hparams, vocab_size)
+    p_hparams.modality["inputs"] = modalities.ModalityType.IMAGE
+    p_hparams.modality["targets"] = modalities.ModalityType.CLASS_LABEL
     with self.test_session() as session:
       features = {
           "inputs": tf.constant(x, dtype=tf.int32),
diff --git a/tensor2tensor/models/slicenet.py b/tensor2tensor/models/slicenet.py
index f9d300e98..8006974b5 100644
--- a/tensor2tensor/models/slicenet.py
+++ b/tensor2tensor/models/slicenet.py
@@ -23,6 +23,7 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -264,9 +265,9 @@ def slicenet_internal(inputs, targets, target_space, hparams, run_decoder=True):
 class SliceNet(t2t_model.T2TModel):
 
   def body(self, features):
-    target_modality_name = self._problem_hparams.modality["targets"].name
+    target_modality = self._problem_hparams.modality["targets"]
     # If we're just predicting a class, there is no use for a decoder.
-    run_decoder = "class_label_modality" not in target_modality_name
+    run_decoder = target_modality != modalities.ModalityType.CLASS_LABEL
     return slicenet_internal(
         features["inputs"],
         features["targets"],
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 4d161adce..5361c4016 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -32,6 +32,7 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.layers import transformer_layers
 from tensor2tensor.utils import beam_search
 from tensor2tensor.utils import expert_utils
@@ -371,10 +372,13 @@ def _fast_decode_tpu(self,
     dp = self._data_parallelism
     hparams = self._hparams
     target_modality = self._problem_hparams.modality["targets"]
+    target_vocab_size = self._problem_hparams.vocab_size["targets"]
+    if target_vocab_size is not None and hasattr(hparams, "vocab_divisor"):
+      target_vocab_size += (-target_vocab_size) % hparams.vocab_divisor
 
     if self.has_input:
       inputs = features["inputs"]
-      if target_modality.is_class_modality:
+      if target_modality == modalities.ModalityType.CLASS_LABEL:
         decode_length = 1
       else:
         decode_length = (
@@ -391,8 +395,16 @@ def _fast_decode_tpu(self,
       # _shard_features called to ensure that the variable names match
       inputs = self._shard_features({"inputs": inputs})["inputs"]
       input_modality = self._problem_hparams.modality["inputs"]
-      with tf.variable_scope(input_modality.name):
-        inputs = dp(input_modality.bottom, inputs)
+      input_vocab_size = self._problem_hparams.vocab_size["inputs"]
+      if input_vocab_size is not None and hasattr(hparams, "vocab_divisor"):
+        input_vocab_size += (-input_vocab_size) % hparams.vocab_divisor
+      modality_name = hparams.name.get(
+          "inputs",
+          modalities.get_name(input_modality))(hparams, input_vocab_size)
+      with tf.variable_scope(modality_name):
+        bottom = hparams.bottom.get("inputs",
+                                    modalities.get_bottom(input_modality))
+        inputs = dp(bottom, inputs, hparams, input_vocab_size)
       with tf.variable_scope("body"):
         encoder_output, encoder_decoder_attention_bias = dp(
             self.encode,
@@ -450,8 +462,13 @@ def preprocess_targets(targets, i):
       """
       # _shard_features called to ensure that the variable names match
       targets = self._shard_features({"targets": targets})["targets"]
-      with tf.variable_scope(target_modality.name):
-        targets = dp(target_modality.targets_bottom, targets)[0]
+      modality_name = hparams.name.get(
+          "targets",
+          modalities.get_name(target_modality))(hparams, target_vocab_size)
+      with tf.variable_scope(modality_name):
+        bottom = hparams.bottom.get(
+            "targets", modalities.get_targets_bottom(target_modality))
+        targets = dp(bottom, targets, hparams, target_vocab_size)[0]
       targets = common_layers.flatten4d3d(targets)
 
       # TODO(llion): Explain! Is this even needed?
@@ -505,9 +522,13 @@ def symbols_to_logits_tpu_fn(ids, i, cache):
             cache,
             i,
             nonpadding=features_to_nonpadding(features, "targets"))
-
-      with tf.variable_scope(target_modality.name):
-        logits = dp(target_modality.top, body_outputs, None)[0]
+      modality_name = hparams.name.get(
+          "targets",
+          modalities.get_name(target_modality))(hparams, target_vocab_size)
+      with tf.variable_scope(modality_name):
+        top = hparams.top.get("targets",
+                              modalities.get_top(target_modality))
+        logits = dp(top, body_outputs, None, hparams, target_vocab_size)[0]
 
       ret = tf.squeeze(logits, axis=[1, 2, 3])
       if partial_targets is not None:
@@ -529,17 +550,13 @@ def forced_logits():
             tf.less(i, partial_targets_length), forced_logits, lambda: ret)
       return ret, cache
 
-    vocab_size = self._problem_hparams.vocab_size["targets"]
-    if hasattr(self._hparams, "vocab_divisor"):
-      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
-
     ret = fast_decode_tpu(
         encoder_output=encoder_output,
         encoder_decoder_attention_bias=encoder_decoder_attention_bias,
         symbols_to_logits_fn=symbols_to_logits_tpu_fn,
         hparams=hparams,
         decode_length=decode_length,
-        vocab_size=vocab_size,
+        vocab_size=target_vocab_size,
         beam_size=beam_size,
         top_beams=top_beams,
         alpha=alpha,
@@ -588,6 +605,9 @@ def _fast_decode(self,
     dp = self._data_parallelism
     hparams = self._hparams
     target_modality = self._problem_hparams.modality["targets"]
+    target_vocab_size = self._problem_hparams.vocab_size["targets"]
+    if target_vocab_size is not None and hasattr(hparams, "vocab_divisor"):
+      target_vocab_size += (-target_vocab_size) % hparams.vocab_divisor
     if "targets_segmentation" in features:
       raise NotImplementedError(
           "Decoding not supported on packed datasets "
@@ -595,7 +615,7 @@ def _fast_decode(self,
           " of the dataset when decoding.")
     if self.has_input:
       inputs = features["inputs"]
-      if target_modality.is_class_modality:
+      if target_modality == modalities.ModalityType.CLASS_LABEL:
         decode_length = 1
       else:
         decode_length = (
@@ -612,8 +632,16 @@ def _fast_decode(self,
       # _shard_features called to ensure that the variable names match
       inputs = self._shard_features({"inputs": inputs})["inputs"]
       input_modality = self._problem_hparams.modality["inputs"]
-      with tf.variable_scope(input_modality.name):
-        inputs = dp(input_modality.bottom, inputs)
+      input_vocab_size = self._problem_hparams.vocab_size["inputs"]
+      if input_vocab_size is not None and hasattr(hparams, "vocab_divisor"):
+        input_vocab_size += (-input_vocab_size) % hparams.vocab_divisor
+      modality_name = hparams.name.get(
+          "inputs",
+          modalities.get_name(input_modality))(hparams, input_vocab_size)
+      with tf.variable_scope(modality_name):
+        bottom = hparams.bottom.get("inputs",
+                                    modalities.get_bottom(input_modality))
+        inputs = dp(bottom, inputs, hparams, input_vocab_size)
       with tf.variable_scope("body"):
         encoder_output, encoder_decoder_attention_bias = dp(
             self.encode,
@@ -671,8 +699,13 @@ def preprocess_targets(targets, i):
       """
       # _shard_features called to ensure that the variable names match
       targets = self._shard_features({"targets": targets})["targets"]
-      with tf.variable_scope(target_modality.name):
-        targets = dp(target_modality.targets_bottom, targets)[0]
+      modality_name = hparams.name.get(
+          "targets",
+          modalities.get_name(target_modality))(hparams, target_vocab_size)
+      with tf.variable_scope(modality_name):
+        bottom = hparams.bottom.get(
+            "targets", modalities.get_targets_bottom(target_modality))
+        targets = dp(bottom, targets, hparams, target_vocab_size)[0]
       targets = common_layers.flatten4d3d(targets)
 
       # TODO(llion): Explain! Is this even needed?
@@ -708,8 +741,12 @@ def symbols_to_logits_fn(ids, i, cache):
             cache,
             nonpadding=features_to_nonpadding(features, "targets"))
 
-      with tf.variable_scope(target_modality.name):
-        logits = dp(target_modality.top, body_outputs, None)[0]
+      modality_name = hparams.name.get(
+          "targets",
+          modalities.get_name(target_modality))(hparams, target_vocab_size)
+      with tf.variable_scope(modality_name):
+        top = hparams.top.get("targets", modalities.get_top(target_modality))
+        logits = dp(top, body_outputs, None, hparams, target_vocab_size)[0]
 
       ret = tf.squeeze(logits, axis=[1, 2, 3])
       if partial_targets is not None:
@@ -729,17 +766,13 @@ def forced_logits():
             tf.less(i, partial_targets_length), forced_logits, lambda: ret)
       return ret, cache
 
-    vocab_size = self._problem_hparams.vocab_size["targets"]
-    if hasattr(self._hparams, "vocab_divisor"):
-      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
-
     ret = fast_decode(
         encoder_output=encoder_output,
         encoder_decoder_attention_bias=encoder_decoder_attention_bias,
         symbols_to_logits_fn=symbols_to_logits_fn,
         hparams=hparams,
         decode_length=decode_length,
-        vocab_size=vocab_size,
+        vocab_size=target_vocab_size,
         beam_size=beam_size,
         top_beams=top_beams,
         alpha=alpha,
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index c5502569a..f3516b8b1 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -351,14 +351,16 @@ def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
       targets = extra_raw_gts
       targets_shape = common_layers.shape_list(targets)
       targets = tf.reshape(targets, [-1] + targets_shape[2:])
-      mod = self.hparams.modality.get("targets",
-                                      self.problem_hparams.modality["targets"])
+      modality = self.hparams.problem_hparams.modality["targets"]
+      targets_weights_fn = self.hparams.targets_weights_fn.get(
+          "targets",
+          modalities.get_targets_weights_fn(modality))(self.hparams)
       numerator, denominator = common_layers.padded_cross_entropy(
           logits,
           targets,
           self.hparams.label_smoothing,
           cutoff=getattr(self.hparams, "video_modality_loss_cutoff", 0.01),
-          weights_fn=mod.targets_weights_fn)
+          weights_fn=targets_weights_fn)
       recon_loss = numerator / denominator
     else:
       raise ValueError("internal loss only supports specific modalities.")
diff --git a/tensor2tensor/models/video/tests_utils.py b/tensor2tensor/models/video/tests_utils.py
index b8677396c..e28936993 100644
--- a/tensor2tensor/models/video/tests_utils.py
+++ b/tensor2tensor/models/video/tests_utils.py
@@ -41,11 +41,18 @@ def fill_hparams(hparams, in_frames, out_frames):
 
 
 def action_modalities(hparams):
+  """Modalities with actions."""
   hparams.problem_hparams.modality = {
-      "inputs": modalities.VideoModalityL2Raw(hparams, 256),
-      "input_action": modalities.SymbolModality(hparams, 5),
-      "targets": modalities.VideoModalityL2Raw(hparams, 256),
-      "target_action": modalities.SymbolModality(hparams, 5),
+      "inputs": modalities.ModalityType.VIDEO_L2_RAW,
+      "input_action": modalities.ModalityType.SYMBOL,
+      "targets": modalities.ModalityType.VIDEO_L2_RAW,
+      "target_action": modalities.ModalityType.SYMBOL,
+  }
+  hparams.problem_hparams.vocab_size = {
+      "inputs": 256,
+      "input_action": 5,
+      "targets": 256,
+      "target_action": 5,
   }
   return hparams
 
@@ -53,12 +60,20 @@ def action_modalities(hparams):
 def full_modalities(hparams):
   """Full modalities with actions and rewards."""
   hparams.problem_hparams.modality = {
-      "inputs": modalities.VideoModalityL2Raw(hparams, 256),
-      "input_reward": modalities.SymbolModality(hparams, 3),
-      "input_action": modalities.SymbolModality(hparams, 5),
-      "targets": modalities.VideoModalityL2Raw(hparams, 256),
-      "target_reward": modalities.SymbolModality(hparams, 3),
-      "target_action": modalities.SymbolModality(hparams, 5),
+      "inputs": modalities.ModalityType.VIDEO_L2_RAW,
+      "input_action": modalities.ModalityType.SYMBOL,
+      "input_reward": modalities.ModalityType.SYMBOL,
+      "targets": modalities.ModalityType.VIDEO_L2_RAW,
+      "target_action": modalities.ModalityType.SYMBOL,
+      "target_reward": modalities.ModalityType.SYMBOL,
+  }
+  hparams.problem_hparams.vocab_size = {
+      "inputs": 256,
+      "input_action": 5,
+      "input_reward": 3,
+      "targets": 256,
+      "target_action": 5,
+      "target_reward": 3,
   }
   hparams.force_full_predict = True
   return hparams
diff --git a/tensor2tensor/models/xception_test.py b/tensor2tensor/models/xception_test.py
index 487e5aafe..041081889 100644
--- a/tensor2tensor/models/xception_test.py
+++ b/tensor2tensor/models/xception_test.py
@@ -41,9 +41,8 @@ def _test_xception(self, img_size):
     p_hparams = problem_hparams.test_problem_hparams(vocab_size,
                                                      vocab_size,
                                                      hparams)
-    p_hparams.modality["inputs"] = modalities.ImageModality(hparams)
-    p_hparams.modality["targets"] = modalities.ClassLabelModality(
-        hparams, vocab_size)
+    p_hparams.modality["inputs"] = modalities.ModalityType.IMAGE
+    p_hparams.modality["targets"] = modalities.ModalityType.CLASS_LABEL
     with self.test_session() as session:
       features = {
           "inputs": tf.constant(x, dtype=tf.int32),
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index f3b05a00b..151551713 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -22,6 +22,7 @@
 import six
 
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import bleu_hook
 from tensor2tensor.utils import rouge
 from tensor2tensor.utils import sari_hook
@@ -613,7 +614,9 @@ def weights_fn_for_mp(problem_task_id):
       tm = {"targets": tm}
 
     for target_name, modality in six.iteritems(tm):
-      weights_fn = modality.targets_weights_fn
+      weights_fn = model_hparams.targets_weights_fn.get(
+          "targets",
+          modalities.get_targets_weights_fn(modality))(model_hparams)
       if hasattr(model_hparams.problem, "task_list"):
         ptid = problem_instance.task_id  # pylint: disable=cell-var-from-loop
         weights_fn = weights_fn_for_mp(ptid)
@@ -638,9 +641,12 @@ def weights_fn_for_mp(problem_task_id):
 def create_eager_metrics_for_problem(problem, model_hparams):
   """See create_eager_metrics."""
   metric_fns = problem.eval_metric_fns(model_hparams)
-  tm = problem.get_hparams(model_hparams).modality["targets"]
-  return create_eager_metrics_internal(
-      metric_fns, weights_fn=tm.targets_weights_fn)
+  problem_hparams = problem.get_hparams(model_hparams)
+  target_modality = problem_hparams.modality["targets"]
+  weights_fn = model_hparams.targets_weights_fn.get(
+      "targets",
+      modalities.get_targets_weights_fn(target_modality))(model_hparams)
+  return create_eager_metrics_internal(metric_fns, weights_fn=weights_fn)
 
 
 def create_eager_metrics(metric_names, weights_fn=common_layers.weights_all):
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index 399eed4a5..4518a729c 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -49,17 +49,18 @@ class Modality(object):
   """
 
   def __init__(self, model_hparams, vocab_size=None):
-    self._model_hparams = model_hparams
-    if vocab_size is not None and hasattr(model_hparams, "vocab_divisor"):
-      vocab_size += (0 - vocab_size) % model_hparams.vocab_divisor
-    self._vocab_size = vocab_size
-
-  @property
-  def name(self):
-    return misc_utils.camelcase_to_snakecase(type(self).__name__)
-
-  @property
-  def top_is_pointwise(self):
+    # __init__ args are unused in any methods. They're maintained for
+    # backwards compatibility for now. In the future, Modality classes will be
+    # removed altogether.
+    del model_hparams, vocab_size
+
+  @classmethod
+  def name(cls, model_hparams, vocab_size=None):
+    del model_hparams, vocab_size  # unused arg
+    return misc_utils.camelcase_to_snakecase(type(cls).__name__)
+
+  @staticmethod
+  def top_is_pointwise():
     """Whether the top mapping of the modality is pointwise.
 
     An example of a pointwise top mapping is a linear layer followed by
@@ -74,8 +75,8 @@ def top_is_pointwise(self):
     """
     return False
 
-  @property
-  def targets_weights_fn(self):
+  @staticmethod
+  def targets_weights_fn(model_hparams):
     """The weights function to use for loss and eval metrics.
 
     A weights function takes labels and returns a Tensor that assigns weights
@@ -85,33 +86,46 @@ def targets_weights_fn(self):
       * weights_all: 1. for all labels
       * weights_nonzero: 1. for all non-zero labels (e.g. to deal with padding)
 
+    Args:
+      model_hparams: tf.HParams, model hyperparmeters.
+
     Returns:
       Callable: (targets) -> weights Tensor
     """
+    del model_hparams  # unused arg
     return common_layers.weights_all
 
-  def bottom(self, x):
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size=None):
     """Transform one shard of input.
 
     Args:
       x: An int32 Tensor with shape [batch, p0, p1, input_channels]
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+
     Returns:
       A float32 Tensor with shape [batch, p0, p1, body_input_depth]
     """
     raise NotImplementedError("Abstract Method")
 
-  def targets_bottom(self, x):
+  @classmethod
+  def targets_bottom(cls, x, model_hparams, vocab_size=None):
     """Transform one shard of targets.
 
     Args:
       x: An int32 Tensor with shape [batch, p0, p1, target_channels]
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+
     Returns:
       A float32 Tensor with shape [batch, p0, p1, body_input_depth]
     """
     with tf.variable_scope("targets_bottom"):
-      return self.bottom(x)
+      return cls.bottom(x, model_hparams, vocab_size)
 
-  def top(self, body_output, targets):
+  @staticmethod
+  def top(body_output, targets, model_hparams, vocab_size=None):
     """Generate predictions/logits for one shard of output.
 
     Most classes will override this function.
@@ -120,23 +134,29 @@ def top(self, body_output, targets):
       body_output: A Tensor with shape [batch, p0, p1, body_output_depth]
       targets: A Tensor with shape [batch, p0, p1, targets_channels,
         top_dimensionality]
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+
     Returns:
       A Tensor of class logits.
     """
     raise NotImplementedError("Abstract Method")
 
-  def loss(self, top_out, targets, weights_fn=None):
+  @classmethod
+  def loss(cls,
+           top_out,
+           targets,
+           model_hparams,
+           vocab_size=None,
+           weights_fn=None):
     """Compute loss numerator and denominator for one shard of output."""
-    logits = top_out
+    del vocab_size  # unused arg
     if weights_fn is None:
-      weights_fn = self.targets_weights_fn
-    logits = common_attention.maybe_upcast(logits, hparams=self._model_hparams)
+      weights_fn = cls.targets_weights_fn(model_hparams)
+    logits = top_out
+    logits = common_attention.maybe_upcast(logits, hparams=model_hparams)
     return common_layers.padded_cross_entropy(
         logits,
         targets,
-        self._model_hparams.label_smoothing,
+        model_hparams.label_smoothing,
         weights_fn=weights_fn)
-
-  @property
-  def is_class_modality(self):
-    return self.name.startswith("class_label")
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 749788784..dfbbc2ab8 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -30,6 +30,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators.problem import problem_hparams_to_features
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.layers.common_attention import mixed_precision_is_enabled
 from tensor2tensor.utils import beam_search
 from tensor2tensor.utils import decoding
@@ -203,6 +204,37 @@ def __init__(self,
           },
           hparams=hparams)
 
+    # TODO(trandustin): For now, we get custom feature transformations via
+    # hparams.modality. Once modality classes are removed, let users
+    # individually specify custom transformations for bottom, loss, etc.
+    if not hasattr(hparams, "modality"):
+      hparams.add_hparam("modality", {})
+    if not hasattr(hparams, "bottom"):
+      hparams.add_hparam("bottom", {})
+    if not hasattr(hparams, "loss"):
+      hparams.add_hparam("loss", {})
+    if not hasattr(hparams, "name"):
+      hparams.add_hparam("name", {})
+    if not hasattr(hparams, "targets_weights_fn"):
+      hparams.add_hparam("targets_weights_fn", {})
+    if not hasattr(hparams, "top"):
+      hparams.add_hparam("top", {})
+    if not hasattr(hparams, "top_is_pointwise"):
+      hparams.add_hparam("top_is_pointwise", {})
+    target_modalities = _create_target_modality(hparams.modality)
+    for feature_name, modality in six.iteritems(hparams.modality):
+      if modality in modalities.ModalityType.get_choices():
+        modality = getattr(modalities, modality)
+      if feature_name in target_modalities:
+        hparams.bottom[feature_name] = modality.targets_bottom
+      else:
+        hparams.bottom[feature_name] = modality.bottom
+      hparams.loss[feature_name] = modality.loss
+      hparams.name[feature_name] = modality.name
+      hparams.targets_weights_fn[feature_name] = modality.targets_weights_fn
+      hparams.top[feature_name] = modality.top
+      hparams.top_is_pointwise[feature_name] = modality.top_is_pointwise
+
     self._original_hparams = hparams
     self.set_mode(mode)
 
@@ -285,8 +317,14 @@ def _custom_getter(self):
   @property
   def _target_modality_is_real(self):
     """Whether the target modality is real-valued."""
-    target_modality = self._problem_hparams.modality["targets"]
-    return target_modality.name.startswith("real_")
+    vocab_size = self._problem_hparams.vocab_size["targets"]
+    if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
+    modality = self._problem_hparams.modality["targets"]
+    modality_name = self._hparams.name.get(
+        "targets",
+        modalities.get_name(modality))(self._hparams, vocab_size)
+    return modality_name.startswith("Real")
 
   def call(self, inputs, **kwargs):
     del kwargs
@@ -436,32 +474,45 @@ def bottom(self, features):
       if feature_name not in features:
         tf.logging.warning("Missing feature %s - ignoring." % feature_name)
         continue
+      vocab_size = self._problem_hparams.vocab_size[feature_name]
+      if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
+        vocab_size += (-vocab_size) % self._hparams.vocab_divisor
+      modality_name = self._hparams.name.get(
+          feature_name,
+          modalities.get_name(modality))(self._hparams, vocab_size)
       # Use if-else clauses to preserve behavior of previous changes: namely,
       # the variable scope name for the targets feature if there is only one
       # target modality; and to reuse variable scopes for only input modalities.
       if feature_name in target_modality:
         if len(target_modality) > 1:
-          variable_scope_name = "%s/%s" % (modality.name, feature_name)
+          variable_scope_name = "%s/%s" % (modality_name, feature_name)
         else:
-          variable_scope_name = modality.name
+          variable_scope_name = modality_name
+        bottom = self._hparams.bottom.get(
+            feature_name,
+            modalities.get_targets_bottom(modality))
         # TODO(aidangomez): share variables?
         with tf.variable_scope(variable_scope_name) as vs:
           self._add_variable_scope(variable_scope_name, vs)
           log_info("Transforming feature '%s' with %s.targets_bottom",
                    feature_name,
-                   modality.name)
-          transformed_features[feature_name] = modality.targets_bottom(
-              features[feature_name])
+                   modality_name)
+          transformed_features[feature_name] = bottom(features[feature_name],
+                                                      self._hparams,
+                                                      vocab_size)
       else:
-        do_reuse = modality.name in all_previous_modalities
-        with tf.variable_scope(modality.name, reuse=do_reuse) as vs:
-          self._add_variable_scope(modality.name, vs)
+        bottom = self._hparams.bottom.get(feature_name,
+                                          modalities.get_bottom(modality))
+        do_reuse = modality_name in all_previous_modalities
+        with tf.variable_scope(modality_name, reuse=do_reuse) as vs:
+          self._add_variable_scope(modality_name, vs)
           log_info("Transforming feature '%s' with %s.bottom",
                    feature_name,
-                   modality.name)
-          transformed_features[feature_name] = modality.bottom(
-              features[feature_name])
-        all_previous_modalities.append(modality.name)
+                   modality_name)
+          transformed_features[feature_name] = bottom(features[feature_name],
+                                                      self._hparams,
+                                                      vocab_size)
+        all_previous_modalities.append(modality_name)
 
     for key in features:
       if key not in transformed_features:
@@ -493,20 +544,31 @@ def body(self, features):
     """
     raise NotImplementedError("Abstract Method")
 
-  def _top_single(self, body_output, target_modality, features):
-    if not target_modality:
+  def _top_single(self, body_output, feature_name, features):
+    if not self._problem_hparams:
       log_warn("Without a Problem, T2TModel.top is a passthrough.")
       return body_output
 
-    with tf.variable_scope(target_modality.name) as tm_vs:
+    modality = self._problem_hparams.modality[feature_name]
+    vocab_size = self._problem_hparams.vocab_size[feature_name]
+    if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
+    name = self._hparams.name.get(
+        feature_name,
+        modalities.get_name(modality))(self._hparams, vocab_size)
+    with tf.variable_scope(name) as tm_vs:
       self._add_variable_scope(tm_vs.name, tm_vs)
-      log_info("Transforming body output with %s.top", target_modality.name)
-      last_only = (
-          target_modality.top_is_pointwise and
-          self.hparams.mode == tf.estimator.ModeKeys.PREDICT and
-          not self.hparams.force_full_predict)
+      log_info("Transforming body output with %s.top", name)
+      top_is_pointwise = self._hparams.top_is_pointwise.get(
+          feature_name,
+          modalities.get_top_is_pointwise(modality))()
+      last_only = (top_is_pointwise and
+                   self.hparams.mode == tf.estimator.ModeKeys.PREDICT and
+                   not self.hparams.force_full_predict)
       if not last_only:
-        logits = target_modality.top(body_output, features.get("targets"))
+        top = self._hparams.top.get(feature_name, modalities.get_top(modality))
+        logits = top(body_output, features.get("targets"),
+                     self._hparams, vocab_size)
       else:
         # Take body outputs for the last position only, and targets too.
         if "decode_loop_step" not in features:
@@ -525,8 +587,9 @@ def _top_single(self, body_output, target_modality, features):
           last_position_targets = tf.slice(
               features["targets"], [0, features["decode_loop_step"][0], 0, 0],
               [target_shape[0], 1, target_shape[2], target_shape[3]])
-        logits = target_modality.top(last_position_body_output,
-                                     last_position_targets)
+        top = self._hparams.top.get(feature_name, modalities.get_top(modality))
+        logits = top(last_position_body_output, last_position_targets,
+                     self._hparams, vocab_size)
     return logits
 
   def top(self, body_output, features):
@@ -550,54 +613,43 @@ def top(self, body_output, features):
           }
     """
     if isinstance(body_output, dict):
-      if self._problem_hparams:
-        target_modality = _create_target_modality(
-            self._problem_hparams.modality)
-      else:
-        target_modality = {k: None for k in body_output.keys()}
-      for k in body_output.keys():
-        assert k in target_modality.keys(), (
-            "The key %s of model_body's returned logits dict must be in "
-            "problem_hparams.modality's dict." % k)
       logits = {}
       for k, v in six.iteritems(body_output):
         # TODO(aidangomez): share variables here?
         with tf.variable_scope(k) as top_vs:
           self._add_variable_scope("top_%s" % k, top_vs)
-          logits[k] = self._top_single(v, target_modality[k], features)
+          logits[k] = self._top_single(v, k, features)
       return logits
     else:
-      if self._problem_hparams:
-        target_modality = _create_target_modality(
-            self._problem_hparams.modality)
-      else:
-        target_modality = None
-      if isinstance(target_modality, dict):
-        assert "targets" in target_modality, (
-            "model_body returned single logits so 'targets' must be a key "
-            "since problem_hparams.modality is a dict.")
-        target_modality = target_modality["targets"]
-      return self._top_single(body_output, target_modality, features)
-
-  def _loss_single(self, logits, target_modality, feature, weights=None):
+      return self._top_single(body_output, "targets", features)
+
+  def _loss_single(self, logits, feature_name, feature, weights=None):
     # The current bfloat16 version still uses float32 for most parts of backward
     # propagation to keep model quality, so cast back before computing the loss
     # value.
-    if not target_modality:
+    if not self._problem_hparams:
       log_warn(_no_problem_err("loss"))
       return (tf.constant(0., dtype=tf.float32),
               tf.constant(1., dtype=tf.float32))
 
     # Calculate loss contribution.
+    modality = self._problem_hparams.modality[feature_name]
+    vocab_size = self._problem_hparams.vocab_size[feature_name]
+    if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
+    loss = self._hparams.loss.get(feature_name, modalities.get_loss(modality))
     if weights is None:
-      loss_num, loss_den = target_modality.loss(logits, feature)
+      loss_num, loss_den = loss(logits, feature, self._hparams, vocab_size)
     else:
 
       def weights_fn(labels):
         """Per-token weights for loss."""
         # Use target_weights_fn() given by modality as well as explicitly given
         # weights.
-        modality_weights = target_modality.targets_weights_fn(labels)
+        targets_weights_fn = self._hparams.targets_weights_fn.get(
+            feature_name,
+            modalities.get_targets_weights_fn(modality))(self._hparams)
+        modality_weights = targets_weights_fn(labels)
 
         # Broadcast 'weights' along minor dimensions (TF's default is major).
         explicit_weights = weights
@@ -609,15 +661,15 @@ def weights_fn(labels):
 
       # Ensure that target.modality_loss() supports "weights_fn" keyword
       # argument. If it doesn't and "weights" is specified, raise an exception.
-      argument_names = inspect.getargspec(target_modality.loss).args
+      argument_names = inspect.getargspec(loss).args
       if "weights_fn" not in argument_names:
         raise ValueError(
-            "Explicit 'weights' given but target_modality.loss doesn't "
+            "Explicit 'weights' given but default loss for modality doesn't "
             "support 'weights_fn' keyword argument: %s.loss(%s)." %
-            (type(target_modality), ", ".join(argument_names)))
+            (modality, ", ".join(argument_names)))
 
-      loss_num, loss_den = target_modality.loss(
-          logits, feature, weights_fn=weights_fn)
+      loss_num, loss_den = loss(
+          logits, feature, self._hparams, vocab_size, weights_fn=weights_fn)
 
     loss_num *= self._problem_hparams.loss_multiplier
 
@@ -630,7 +682,7 @@ def weights_fn(labels):
           self.hparams,
           self._problem_hparams,
           logits,
-          target_modality,
+          feature_name,
           feature
       )
 
@@ -641,20 +693,11 @@ def weights_fn(labels):
 
   def loss(self, logits, features):
     if isinstance(logits, dict):
-      if self._problem_hparams:
-        target_modality = _create_target_modality(
-            self._problem_hparams.modality)
-      else:
-        target_modality = {k: None for k in logits.keys()}
-      for k in logits.keys():
-        assert k in target_modality.keys(), (
-            "The key %s of model_body's returned logits dict must be in "
-            "problem_hparams.modality's dict." % k)
       losses = {}
       for k, v in six.iteritems(logits):
         losses[k] = self._loss_single(
             v,
-            target_modality[k],
+            k,
             features[k],
             weights=features.get(k + "_mask"))
 
@@ -670,19 +713,9 @@ def loss(self, logits, features):
 
       return tf.add_n([n / d for n, d in losses.values()])
     else:
-      if self._problem_hparams:
-        target_modality = _create_target_modality(
-            self._problem_hparams.modality)
-      else:
-        target_modality = None
-      if isinstance(target_modality, dict):
-        assert "targets" in target_modality, (
-            "model_body returned single logits so 'targets' must be a key "
-            "since problem_hparams.modality is a dict.")
-        target_modality = target_modality["targets"]
       return self._loss_single(
           logits,
-          target_modality,
+          "targets",
           features["targets"],
           weights=features.get("targets_mask"))
 
@@ -709,12 +742,6 @@ def set_mode(self, mode):
           setattr(hparams, key, 0.0)
     self._hparams = hparams
 
-    if self._problem_hparams:
-      # Set model hparams in problem_hparams' modalities, which also store them.
-      for modality in six.itervalues(self._problem_hparams.modality):
-        if modality is not None:
-          modality._model_hparams = self._hparams  # pylint: disable=protected-access
-
   def prepare_features_for_infer(self, features):
     """Called before inference to allow adding infer-specific features."""
     pass
@@ -789,7 +816,7 @@ def infer(self,
 
       if self._problem_hparams:
         target_modality = self._problem_hparams.modality["targets"]
-        if target_modality.is_class_modality:
+        if target_modality == modalities.ModalityType.CLASS_LABEL:
           beam_size = 1  # No use to run beam-search for a single class.
       if beam_size == 1:
         log_info("Greedy Decoding")
@@ -870,7 +897,11 @@ def symbols_to_logits_fn(ids, i=None):
       # it has shape [batch_size] and contains floats between 0 and
       # source_length.
       if self._problem_hparams:
-        if self._problem_hparams.modality["targets"].top_is_pointwise:
+        target_modality = self._problem_hparams.modality["targets"]
+        top_is_pointwise = self._hparams.top_is_pointwise.get(
+            "targets",
+            modalities.get_top_is_pointwise(target_modality))()
+        if top_is_pointwise:
           return tf.squeeze(logits, axis=[1, 2, 3])
       # -1 due to the pad above.
       current_output_position = common_layers.shape_list(ids)[1] - 1
@@ -1009,7 +1040,10 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
       samples, logits, losses = self.sample(features)
       # Concatenate the already-generated recent_output with last timestep
       # of the newly-generated samples.
-      if target_modality.top_is_pointwise:
+      top_is_pointwise = self._hparams.top_is_pointwise.get(
+          "targets",
+          modalities.get_top_is_pointwise(target_modality))()
+      if top_is_pointwise:
         cur_sample = samples[:, -1, :, :]
       else:
         cur_sample = samples[:, i, :, :]
@@ -1043,7 +1077,7 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
     initial_output = tf.slice(initial_output, [0, 0, 0, 0],
                               common_layers.shape_list(initial_output))
     target_modality = self._problem_hparams.modality["targets"]
-    if target_modality.is_class_modality:
+    if target_modality == modalities.ModalityType.CLASS_LABEL:
       decode_length = 1
     else:
       if "partial_targets" in features:
@@ -1094,10 +1128,6 @@ def fn_not_eos():
             lambda: not_overflow)
       return not_overflow
 
-    vocab_size = self._problem_hparams.vocab_size["targets"]
-    if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
-      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
-
     _, result, logits, loss = tf.while_loop(
         while_exit_cond,
         infer_step, [tf.constant(0), result, logits, loss],
@@ -1184,7 +1214,10 @@ def infer_step(recent_output, recent_logits, unused_loss):
       samples, logits, losses = self.sample(features)
       # Concatenate the already-generated recent_output with last timestep
       # of the newly-generated samples.
-      if target_modality.top_is_pointwise:
+      top_is_pointwise = self._hparams.top_is_pointwise.get(
+          "targets",
+          modalities.get_top_is_pointwise(target_modality))()
+      if top_is_pointwise:
         cur_sample = samples[:, -1, :, :]
       else:
         cur_sample = samples[:,
@@ -1224,7 +1257,7 @@ def infer_step(recent_output, recent_logits, unused_loss):
     initial_output = tf.slice(initial_output, [0, 0, 0, 0],
                               common_layers.shape_list(initial_output))
     target_modality = self._problem_hparams.modality["targets"]
-    if target_modality.is_class_modality:
+    if target_modality == modalities.ModalityType.CLASS_LABEL:
       decode_length = 1
     else:
       if "partial_targets" in features:
@@ -1780,7 +1813,7 @@ def create_tpu_eval_metrics_fn(problem, model_hparams):
   tm = _create_target_modality(problem.get_hparams(model_hparams).modality)
   if isinstance(tm, dict):
     for k, v in six.iteritems(tm):
-      weights_fn = v.targets_weights_fn
+      weights_fn = v.targets_weights_fn(model_hparams)
 
       def make_metric_fn(metric_fn):
         def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
@@ -1800,7 +1833,7 @@ def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
         name = "%s/metrics-%s/%s" % (k, problem.name, metric)
         metric_fns.append((name, make_metric_fn(metric_fn)))
   else:
-    weights_fn = tm.targets_weights_fn
+    weights_fn = tm.targets_weights_fn(model_hparams)
 
     def make_metric_fn(metric_fn):
       def wrapped_metric_fn(logits, labels, features):
@@ -1956,13 +1989,13 @@ def create_eager_var_store():
 def scheduled_sampling(hparams, problem_hparams, dp, sharded_logits, losses,
                        sharded_features, transformed_features, model):
   """Scheduled sampling."""
-  target_modality = problem_hparams.modality["targets"]
+  modality = problem_hparams.modality["targets"]
+  vocab_size = problem_hparams.vocab_size["targets"]
+  if vocab_size is not None and hasattr(hparams, "vocab_divisor"):
+    vocab_size += (-vocab_size) % hparams.vocab_divisor
 
   def sample(x):
     """Multinomial sampling from a n-dimensional tensor."""
-    vocab_size = problem_hparams.vocab_size["targets"]
-    if vocab_size is not None and hasattr(hparams, "vocab_divisor"):
-      vocab_size += (-vocab_size) % hparams.vocab_divisor
     samples = tf.multinomial(tf.reshape(x, [-1, vocab_size]), 1)
     reshaped_samples = tf.reshape(samples, common_layers.shape_list(x)[:-1])
     return tf.to_int32(reshaped_samples)
@@ -1981,21 +2014,31 @@ def sampled_results():
                      sampled_targets)
     new_features = transformed_features
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
-      with tf.variable_scope(target_modality.name):
-        new_features["targets"] = dp(target_modality.targets_bottom,
-                                     new_targets)
+      modality_name = hparams.name.get(
+          "targets",
+          modalities.get_name(modality))(hparams, vocab_size)
+      with tf.variable_scope(modality_name):
+        bottom = hparams.bottom.get(
+            "targets", modalities.get_targets_bottom(modality))
+        new_features["targets"] = dp(bottom, new_targets, hparams, vocab_size)
       with tf.variable_scope("body"):
         body_outputs, losses = model.model_fn_sharded(new_features)
         if not isinstance(losses, dict):  # If it's a single extra loss.
           losses = {"extra": losses}
-      with tf.variable_scope(target_modality.name):
-        new_sharded_logits = dp(target_modality.top,
+      with tf.variable_scope(modality_name):
+        top = hparams.top.get("targets", modalities.get_top(modality))
+        new_sharded_logits = dp(top,
                                 body_outputs,
-                                sharded_features["targets"])
+                                sharded_features["targets"],
+                                hparams,
+                                vocab_size)
         if "training" not in losses:
-          sharded_loss_num, sharded_loss_den = dp(target_modality.loss,
+          loss = hparams.loss.get("targets", modalities.get_loss(modality))
+          sharded_loss_num, sharded_loss_den = dp(loss,
                                                   sharded_logits,
-                                                  sharded_features["targets"])
+                                                  sharded_features["targets"],
+                                                  hparams,
+                                                  vocab_size)
           training_loss = (tf.add_n(sharded_loss_num) /
                            tf.maximum(1.0, tf.add_n(sharded_loss_den)))
           training_loss *= problem_hparams.loss_multiplier
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index 0c9ac408e..c9b680302 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -19,7 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.utils import modality
+from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.utils import t2t_model
 from tensor2tensor.utils import test_utils
 from tensor2tensor.utils.hparam import HParams
@@ -51,29 +51,28 @@ def testLossSingleWeights(self):
         vocab_size = 3
 
         model_hparams = HParams(
+            prepend_mode="none",
             label_smoothing=0.0,
             shared_embedding_and_softmax_weights=False)
 
-        problem_hparams = HParams(loss_multiplier=1.0)
-        problem_hparams.modality = {}
+        ph = problem_hparams.TestProblem(
+            vocab_size, vocab_size).get_hparams(model_hparams)
 
-        model = t2t_model.T2TModel(
-            model_hparams, problem_hparams=problem_hparams)
+        model = t2t_model.T2TModel(model_hparams, problem_hparams=ph)
         logits = tf.zeros((batch_size, sequence_size, 1, 1, vocab_size))
-        target_modality = modality.Modality(model_hparams)
         feature = tf.ones((batch_size, sequence_size, 1, 1))
 
         # all-zero weights == zero loss.
         weights = tf.zeros((batch_size, sequence_size))
         loss_num, loss_denom = model._loss_single(
-            logits, target_modality, feature, weights=weights)
+            logits, "targets", feature, weights=weights)
         self.assertAllClose(tf.zeros_like(loss_num), sess.run(loss_num))
         self.assertAllClose(tf.zeros_like(loss_denom), sess.run(loss_denom))
 
         # non-zero weights > zero loss.
         weights = tf.ones((batch_size, sequence_size))
         loss_num, loss_denom = model._loss_single(
-            logits, target_modality, feature, weights=weights)
+            logits, "targets", feature, weights=weights)
         self.assertAllLess(0.0, sess.run(loss_num))
         self.assertAllClose(batch_size * sequence_size, sess.run(loss_denom))
 
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index 2995a4e84..d45550fcf 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -111,6 +111,11 @@ def testMultipleTargetModalities(self):
         "targets_A": hparams.problem_hparams.modality["targets"],
         "targets_B": hparams.problem_hparams.modality["targets"],
     }
+    hparams.problem_hparams.vocab_size = {
+        "targets": hparams.problem_hparams.vocab_size["targets"],
+        "targets_A": hparams.problem_hparams.vocab_size["targets"],
+        "targets_B": hparams.problem_hparams.vocab_size["targets"],
+    }
     hparams.problem._hparams = hparams.problem_hparams
 
     # Dataset

From 616421a12d842ae2d689d248249d2b5db0f2e0d3 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 21 Feb 2019 09:34:24 -0800
Subject: [PATCH 1689/2720] Use JAX in a T2T way.

What works:
* jax/j2j_trainer --dataset=mnist --model=mlp
* t2t_trainer --jax --problem=mnist --model=mlp --output_dir ~/t2t_train/testjx/ --train_steps=1000 --local_eval_frequency=100

PiperOrigin-RevId: 235010175
---
 setup.py                            |   3 +
 tensor2tensor/bin/t2t_trainer.py    |  31 +++--
 tensor2tensor/jax/input_pipeline.py | 189 ++++++++++++++++++++++++++
 tensor2tensor/jax/j2j.py            | 201 ++++++++++++++++++++++++++++
 tensor2tensor/jax/j2j_trainer.py    |  94 +++++++++++++
 tensor2tensor/jax/models.py         |  35 +++++
 6 files changed, 544 insertions(+), 9 deletions(-)
 create mode 100644 tensor2tensor/jax/input_pipeline.py
 create mode 100644 tensor2tensor/jax/j2j.py
 create mode 100644 tensor2tensor/jax/j2j_trainer.py
 create mode 100644 tensor2tensor/jax/models.py

diff --git a/setup.py b/setup.py
index 9221ee573..6c18258e6 100644
--- a/setup.py
+++ b/setup.py
@@ -38,10 +38,13 @@
         'flask',
         'future',
         'gevent',
+        'gin-config',
         'google-api-python-client',
         'gunicorn',
         'gym',
         'h5py',
+        'jax',
+        'jaxlib',
         'kfac',
         'mesh-tensorflow',
         'numpy',
diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index c922e74dc..ead630db7 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -21,9 +21,11 @@
 import contextlib
 import os
 import sys
+import gin
 from tensor2tensor import models  # pylint: disable=unused-import
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
 from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
+from tensor2tensor.jax import j2j
 from tensor2tensor.utils import cloud_mlengine
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
@@ -71,6 +73,7 @@
                      "Number of intra_op_parallelism_threads to use for CPU. "
                      "See TensorFlow config.proto for details.")
 flags.DEFINE_bool("v2", False, "Whether to use T2T v2.")
+flags.DEFINE_bool("jax", False, "Whether to use J2J.")
 # TODO(lukaszkaiser): resolve memory and variable assign issues and set to True.
 flags.DEFINE_bool(
     "optionally_use_dist_strat", False,
@@ -359,25 +362,35 @@ def run_std_server():
 def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
 
-  if FLAGS.v2:
+  if FLAGS.v2 or FLAGS.jax:
     tf.enable_v2_behavior()
-    # Hacking main v1 flags to work with v2.
+    # Hacking main v1 flags to work with v2 and jax.
+    prefix = "t2t." if FLAGS.v2 else "j2j."
     config_strs = []
     config_strs.append(
-        "train_fn.train_steps=" + str(FLAGS.train_steps))
+        prefix + "train_fn.train_steps=" + str(FLAGS.train_steps))
     config_strs.append(
-        "train_fn.eval_steps=" + str(FLAGS.eval_steps))
+        prefix + "train_fn.eval_steps=" + str(FLAGS.eval_steps))
     config_strs.append(
-        "train_fn.eval_frequency=" + str(FLAGS.local_eval_frequency))
+        prefix + "train_fn.eval_frequency=" + str(FLAGS.local_eval_frequency))
     if FLAGS.hparams:
       config_strs.extend(str(FLAGS.hparams).split(","))
     config_str = "\n".join(config_strs)
     data_dir = os.path.expanduser(FLAGS.data_dir)
     output_dir = os.path.expanduser(FLAGS.output_dir)
-    t2t_v2.t2t_train(FLAGS.model, FLAGS.problem,
-                     data_dir=data_dir, output_dir=output_dir,
-                     config_file=FLAGS.hparams_set, config=config_str)
-    return
+
+    if FLAGS.v2:
+      t2t_v2.t2t_train(FLAGS.model, FLAGS.problem,
+                       data_dir=data_dir, output_dir=output_dir,
+                       config_file=FLAGS.hparams_set, config=config_str)
+      return
+
+    if FLAGS.jax:
+      gin.bind_parameter("j2j.train_fn.dataset", FLAGS.problem)
+      config_strs += ["j2j.train_fn.model=@models." + FLAGS.model]
+      gin.parse_config_files_and_bindings(FLAGS.hparams_set, config_strs)
+      j2j.train_fn(data_dir, output_dir=output_dir)
+      return
 
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
 
diff --git a/tensor2tensor/jax/input_pipeline.py b/tensor2tensor/jax/input_pipeline.py
new file mode 100644
index 000000000..64c7cac1b
--- /dev/null
+++ b/tensor2tensor/jax/input_pipeline.py
@@ -0,0 +1,189 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""T2T models, configs and main training functions."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+import gin
+
+from tensor2tensor import problems
+
+import tensorflow as tf
+import tensorflow_datasets as tfds
+
+
+def train_and_eval_dataset(dataset_name, data_dir):
+  """Return train and evaluation datasets, feature info and supervised keys.
+
+  Args:
+    dataset_name: a string, the name of the dataset; if it starts with "v1_"
+      then we'll search T2T Problem registry for it, otherwise we assume it
+      is a dataset from TFDS and load it from there.
+    data_dir: directory where the data is located.
+
+  Returns:
+    a 4-tuple consisting of:
+     * the train tf.Daataset
+     * the eval tf.Daataset
+     * information about features: a python dictionary with feature names
+         as keys and an object as value that provides .shape and .num_classes.
+     * supervised_keys: information what's the input and what's the target,
+         ie., a pair of lists with input and target feature names.
+  """
+  if dataset_name.startswith("v1_"):
+    return _train_and_eval_dataset_v1(dataset_name[3:], data_dir)
+  dataset_builder = tfds.builder(dataset_name, data_dir=data_dir)
+  info = dataset_builder.info
+  splits = dataset_builder.info.splits
+  if tfds.Split.TRAIN not in splits:
+    raise ValueError("To train we require a train split in the dataset.")
+  if tfds.Split.VALIDATION not in splits and "test" not in splits:
+    raise ValueError("We require a validation or test split in the dataset.")
+  eval_split = tfds.Split.VALIDATION
+  if tfds.Split.VALIDATION not in splits:
+    eval_split = tfds.Split.TEST
+  train, valid = tfds.load(
+      name=dataset_name, split=[tfds.Split.TRAIN, eval_split])
+  keys = None
+  if info.supervised_keys:
+    keys = ([info.supervised_keys[0]], [info.supervised_keys[1]])
+  return train, valid, info.features, keys
+
+
+def _make_info(shape_list, num_classes):
+  """Create an info-like tuple for feature given some shapes and vocab size."""
+  feature_info = collections.namedtuple("FeatureInfo", ["shape", "num_classes"])
+  cur_shape = list(shape_list[0])
+  # We need to merge the provided shapes, put None where they disagree.
+  for shape in shape_list:
+    if len(shape) != len(cur_shape):
+      raise ValueError("Shapes need to have the same number of dimensions.")
+    for i in range(len(shape)):
+      if cur_shape[i] is not None:
+        if shape[i] != cur_shape[i]:
+          cur_shape[i] = None
+  return feature_info(cur_shape, num_classes)
+
+
+def _select_features(example, feature_list=None):
+  """Select a subset of features from the example dict."""
+  feature_list = feature_list or ["inputs", "targets"]
+  return {f: example[f] for f in feature_list}
+
+
+def _train_and_eval_dataset_v1(problem_name, data_dir):
+  """Return train and evaluation datasets, feature info and supervised keys."""
+  problem = problems.problem(problem_name)
+  train_dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, data_dir)
+  train_dataset = train_dataset.map(_select_features)
+  eval_dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, data_dir)
+  eval_dataset = eval_dataset.map(_select_features)
+  supervised_keys = (["inputs"], ["targets"])
+  hparams = problem.get_hparams()
+  # We take a few training examples to guess the shapes.
+  input_shapes, target_shapes = [], []
+  for example in train_dataset.take(3):
+    input_shapes.append(example["inputs"].shape.as_list())
+    target_shapes.append(example["targets"].shape.as_list())
+  input_vocab_size = hparams.vocab_size["inputs"]
+  target_vocab_size = hparams.vocab_size["targets"]
+  input_info = _make_info(input_shapes, input_vocab_size)
+  target_info = _make_info(target_shapes, target_vocab_size)
+  info = {"inputs": input_info, "targets": target_info}
+  return train_dataset, eval_dataset, info, supervised_keys
+
+
+@gin.configurable(whitelist=["max_target_length"])
+def preprocess_fn(dataset, training, max_target_length=-1):
+  def target_right_length(_, target):
+    if max_target_length < 1 or not training:
+      return tf.constant(True)
+    return tf.less(tf.shape(target)[0], max_target_length + 1)
+  dataset = dataset.filter(target_right_length)
+  return dataset
+
+
+@gin.configurable(blacklist=["dataset", "training", "shapes", "target_names"])
+def batch_fn(dataset, training, shapes, target_names,
+             batch_size=32, eval_batch_size=32, bucket_length=32, buckets=None):
+  """Batching function."""
+  del target_names
+  # If bucketing is not specified, check if target shapes are variable.
+  cur_batch_size = batch_size if training else eval_batch_size
+  if buckets is None:
+    variable_target_shapes = False
+    target_shape = shapes[1]
+    for dim in target_shape:
+      if dim is None:
+        variable_target_shapes = True
+    tf.logging.info("Heuristically setting bucketing to %s based on shapes "
+                    "of target tensors." % variable_target_shapes)
+    if variable_target_shapes:
+      bucket_boundaries = [bucket_length // 4, bucket_length // 2,
+                           bucket_length, bucket_length * 2,
+                           bucket_length * 4, bucket_length * 8]
+      bucket_batch_sizes = [cur_batch_size * 4, cur_batch_size * 2,
+                            cur_batch_size, cur_batch_size // 2,
+                            cur_batch_size // 4, cur_batch_size // 8]
+      buckets = (bucket_boundaries, bucket_batch_sizes)
+
+  if buckets:
+    tf.logging.info("Bucketing with buckets %s." % str(buckets))
+    def example_length(_, target):
+      return tf.shape(target)[0]
+    boundaries, batch_sizes = buckets
+    dataset = dataset.apply(tf.data.experimental.bucket_by_sequence_length(
+        example_length, boundaries, batch_sizes, pad_to_bucket_boundary=True))
+  else:
+    dataset = dataset.padded_batch(cur_batch_size, shapes)
+  return dataset
+
+
+def shuffle_and_batch_data(dataset, target_names, features_info, training):
+  """Shuffle and batch the given dataset."""
+  def append_targets(example):
+    """Append targets to the example dictionary. Needed for Keras."""
+    if len(target_names) == 1:
+      return (example, example[target_names[0]])
+    targets = {}
+    for name in target_names:
+      targets[name] = example[name]
+    return (example, targets)
+  dataset = dataset.map(append_targets)
+  if training:
+    dataset = dataset.repeat()
+  shapes = {k: features_info[k].shape for k in features_info}
+  shapes = (shapes, shapes[target_names[0]])
+  dataset = dataset.shuffle(1024)
+  dataset = preprocess_fn(dataset, training)
+  dataset = batch_fn(dataset, training, shapes, target_names)
+  return dataset.prefetch(32)
+
+
+def train_and_eval_batches(dataset, data_dir):
+  """Return train and eval batches with input name and shape."""
+  (train_data, eval_data, features_info, keys) = train_and_eval_dataset(
+      dataset, data_dir)
+  input_names, target_names = keys[0], keys[1]
+  train_batches = shuffle_and_batch_data(
+      train_data, target_names, features_info, training=True)
+  eval_batches = shuffle_and_batch_data(
+      eval_data, target_names, features_info, training=False)
+  input_shape = features_info[input_names[0]].shape
+  return train_batches, eval_batches, input_names[0], list(input_shape)
diff --git a/tensor2tensor/jax/j2j.py b/tensor2tensor/jax/j2j.py
new file mode 100644
index 000000000..fcab6fc91
--- /dev/null
+++ b/tensor2tensor/jax/j2j.py
@@ -0,0 +1,201 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""J2J main training functions."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import pickle
+import time
+
+from absl import logging
+import gin
+
+import jax
+from jax.experimental import optimizers
+import jax.numpy as np
+
+from tensor2tensor.jax import input_pipeline
+# Import for gin configurable models
+from tensor2tensor.jax import models  # pylint: disable=unused-import
+
+from tensorflow import gfile
+
+import tensorflow_datasets as tfds
+
+
+@gin.configurable(blacklist=["step"])
+def learning_rate(step, schedule=None, constant=0.001, warmup_steps=200):
+  """Learning rate."""
+  schedule = schedule or "constant * linear_warmup * rsqrt_decay"
+  ret = 1.0
+  for name in [n.strip() for n in schedule.split("*")]:
+    if name == "constant":
+      ret *= constant
+    elif name == "linear_warmup":
+      ret *= np.minimum(1.0, step / warmup_steps)
+    elif name == "rsqrt_decay":
+      ret /= np.sqrt(np.maximum(step, warmup_steps))
+    else:
+      raise ValueError("Unknown factor %s." % name)
+  return ret
+
+
+@gin.configurable()
+def optimizer(name="adam",
+              momentum_mass=0.9, rmsprop_gamma=0.9, rmsprop_eps=1e-8,
+              adam_b1=0.9, adam_b2=0.997, adam_eps=1e-8):
+  """Return the optimizer, by name."""
+  if name == "sgd":
+    return optimizers.sgd(learning_rate)
+  if name == "momentum":
+    return optimizers.momentum(learning_rate, mass=momentum_mass)
+  if name == "rmsprop":
+    return optimizers.rmsprop(
+        learning_rate, gamma=rmsprop_gamma, eps=rmsprop_eps)
+  if name == "adam":
+    return optimizers.adam(learning_rate, b1=adam_b1, b2=adam_b2, eps=adam_eps)
+  raise ValueError("Unknown optimizer %s" % str(name))
+
+
+def one_hot(x, k, dtype=np.float32):
+  """Create a one-hot encoding of x of size k."""
+  return np.array(x[:, None] == np.arange(k), dtype)
+
+
+def accuracy(params, batch, model_predict):
+  """Calculate accuracy."""
+  inputs, targets = batch
+  predicted_class = np.argmax(model_predict(params, inputs), axis=1)
+  return np.mean(predicted_class == targets)
+
+
+def loss(params, batch, model_predict):
+  """Calculate loss."""
+  inputs, targets = batch
+  preds = model_predict(params, inputs)
+  return -np.mean(preds * one_hot(targets, preds.shape[-1]))
+
+
+def dataset_to_stream(batches, input_name):
+  """Takes a tf.Dataset and creates a numpy stream of ready batches."""
+  for example in tfds.as_numpy(batches):
+    inp, out = example[0][input_name], example[1]
+    yield inp, out
+
+
+def log(s, stdout=True):
+  logging.info(s)
+  if stdout:
+    print(s)
+
+
+def save_params_and_step(params, step, output_dir):
+  """Save params and step in output dir."""
+  if output_dir is not None:
+    if not gfile.Exists(output_dir):
+      log("Creating directory %s" % output_dir)
+      gfile.MkDir(output_dir)
+    params_file = os.path.join(output_dir, "model.pkl")
+    with gfile.Open(params_file, "wb") as f:
+      pickle.dump((params, step), f)
+    log("Model saved to %s" % params_file, stdout=False)
+
+
+def load_params_and_step(output_dir):
+  """Save params and step in output dir."""
+  if output_dir is None:
+    return None, None
+  if not gfile.Exists(output_dir):
+    return None, None
+  params_file = os.path.join(output_dir, "model.pkl")
+  if not gfile.Exists(params_file):
+    return None, None
+  with gfile.Open(params_file, "r") as f:
+    (params, step) = pickle.load(f)
+  log("Model loaded from %s" % params_file)
+  return params, step
+
+
+# We include in gin config everything that could be useful to share between
+# users, so when it gets saved in a .gin file it can be re-ran with few flags.
+@gin.configurable(blacklist=["data_dir", "output_dir"])
+def train_fn(data_dir=None, output_dir=None,
+             model=gin.REQUIRED,
+             dataset=gin.REQUIRED,
+             train_steps=1000, eval_steps=10, eval_frequency=100):
+  """Train the given model on the given dataset.
+
+  Args:
+    data_dir: Directory where the data is located.
+    output_dir: Directory where to put the logs and checkpoints.
+    model: The model to train (a function).
+    dataset: The name of the dataset to train on.
+    train_steps: for how many steps to train.
+    eval_steps: for how many steps to do evaluation.
+    eval_frequency: how often (every this many steps) to run evaluation.
+  """
+  (train_batches, eval_batches,
+   input_name, input_shape) = input_pipeline.train_and_eval_batches(
+       dataset, data_dir)
+  train_stream = dataset_to_stream(train_batches, input_name)
+
+  # Training loop.
+  opt_init, opt_update = optimizer()
+  model_init, model_predict = model()
+
+  @jax.jit
+  def update(i, opt_state, batch):
+    params = optimizers.get_params(opt_state)
+    return opt_update(i, jax.grad(loss)(
+        params, batch, model_predict), opt_state)
+
+  _, init_params = model_init([-1] + input_shape)
+  step = 0
+  if output_dir is not None:
+    loaded_params, loaded_step = load_params_and_step(output_dir)
+    if loaded_params is not None:
+      init_params = loaded_params
+    if loaded_step is not None:
+      step = loaded_step
+  opt_state = opt_init(init_params)
+
+  log("Starting training.")
+  while step < train_steps:
+    # Training.
+    start_time = time.time()
+    for _ in range(eval_frequency):
+      opt_state = update(step, opt_state, next(train_stream))
+      step += 1
+    epoch_time = time.time() - start_time
+    log("Step {}, last {} steps in {:0.2f} sec".format(
+        step, eval_frequency, epoch_time))
+
+    # Save the model.
+    params = optimizers.get_params(opt_state)
+    save_params_and_step(params, step, output_dir)
+
+    # Evaluation.
+    eval_stream = dataset_to_stream(eval_batches, input_name)
+    eval_train_stream = dataset_to_stream(train_batches, input_name)
+    train_acc, eval_acc = 0, 0
+    for _ in range(eval_steps):
+      train_acc += accuracy(params, next(eval_train_stream), model_predict)
+      eval_acc += accuracy(params, next(eval_stream), model_predict)
+    log("Train set accuracy {:0.4f}".format(train_acc / eval_steps))
+    log("Eval  set accuracy {:0.4f}".format(eval_acc / eval_steps))
diff --git a/tensor2tensor/jax/j2j_trainer.py b/tensor2tensor/jax/j2j_trainer.py
new file mode 100644
index 000000000..5cb081022
--- /dev/null
+++ b/tensor2tensor/jax/j2j_trainer.py
@@ -0,0 +1,94 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+r"""J2J trainer.
+
+Examples:
+
+- train a basic model on mnist:
+    jax/j2j_trainer.py --dataset=mnist --model=mlp
+      --config="train_fn.train_steps=4000" --output_dir ~/j2j/test1
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from absl import app
+from absl import flags
+from absl import logging
+
+import gin
+from tensor2tensor.jax import j2j
+
+import tensorflow as tf
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("dataset", None, "Which dataset to use.")
+flags.DEFINE_string("model", None, "Which model to train.")
+flags.DEFINE_string("data_dir", None, "Path to the directory with data.")
+flags.DEFINE_string("output_dir", None,
+                    "Path to the directory to save logs and checkpoints.")
+flags.DEFINE_multi_string("config_file", None,
+                          "Configuration file with parameters (.gin).")
+flags.DEFINE_multi_string("config", None,
+                          "Configuration parameters (gin string).")
+
+# For iterators over datasets so we can do "for example in dataset".
+tf.enable_v2_behavior()
+
+
+def j2j_train(model_name, dataset_name,
+              data_dir=None, output_dir=None, config_file=None, config=None):
+  """Main function to train the given model on the given dataset.
+
+  Args:
+    model_name: The name of the model to train.
+    dataset_name: The name of the dataset to train on.
+    data_dir: Directory where the data is located.
+    output_dir: Directory where to put the logs and checkpoints.
+    config_file: the gin configuration file to use.
+    config: string (in gin format) to override gin parameters.
+  """
+  gin.bind_parameter("train_fn.dataset", dataset_name)
+  if FLAGS.model:
+    config = []  if config is None else config
+    config += ["train_fn.model=@models." + model_name]
+  gin.parse_config_files_and_bindings(config_file, config)
+  if output_dir:
+    if not tf.gfile.Exists(output_dir):
+      tf.gfile.MkDir(output_dir)
+    config_path = os.path.join(output_dir, "gin.config")
+    # TODO(lukaszkaiser): why is the file empty if there's no provided config?
+    with tf.gfile.Open(config_path, "w") as f:
+      f.write(gin.operative_config_str())
+  j2j.train_fn(data_dir, output_dir=output_dir)
+
+
+def main(argv):
+  del argv
+  logging.set_verbosity(logging.INFO)
+  data_dir, output_dir = FLAGS.data_dir, FLAGS.output_dir
+  data_dir = data_dir and os.path.expanduser(data_dir)
+  output_dir = output_dir and os.path.expanduser(output_dir)
+  j2j_train(FLAGS.model, FLAGS.dataset,
+            data_dir=data_dir, output_dir=output_dir,
+            config_file=FLAGS.config_file, config=FLAGS.config)
+
+
+if __name__ == "__main__":
+  app.run(main)
diff --git a/tensor2tensor/jax/models.py b/tensor2tensor/jax/models.py
new file mode 100644
index 000000000..86d392316
--- /dev/null
+++ b/tensor2tensor/jax/models.py
@@ -0,0 +1,35 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""J2J models."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gin
+
+from jax.experimental import stax
+
+
+@gin.configurable()
+def mlp(num_hidden_layers=2,
+        hidden_size=512,
+        activation_fn=stax.Relu,
+        num_output_classes=10):
+  layers = [stax.Flatten]
+  layers += [stax.Dense(hidden_size), activation_fn] * num_hidden_layers
+  layers += [stax.Dense(num_output_classes), stax.LogSoftmax]
+  return stax.serial(*layers)

From 28b1c770d5330afd02aa30557c26e1b0bda6c585 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Thu, 21 Feb 2019 10:57:18 -0800
Subject: [PATCH 1690/2720] fixes proposed by rsepassi

PiperOrigin-RevId: 235026395
---
 .../data_generators/generator_utils.py        | 90 ++++++++++++++-----
 .../data_generators/ops/pack_sequences_ops.cc |  2 +
 2 files changed, 71 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 4e6bb57ea..8a722101a 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -33,7 +33,6 @@
 import six.moves.urllib_request as urllib
 
 from tensor2tensor.data_generators import text_encoder
-from tensor2tensor.data_generators.ops import pack_sequences_ops
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
@@ -709,34 +708,83 @@ def pack_dataset(dataset, length, keys=None, use_custom_ops=False):
   Returns:
     a tf.data.Dataset
   """
+  shapes = dataset.output_shapes
   if keys is None:
-    keys = dataset.output_shapes.keys
+    keys = shapes.keys()
+  for k in keys:
+    if k not in shapes:
+      raise ValueError("Key %s not found in dataset.  Available keys are %s"
+                       % (k, shapes.keys()))
+    if not shapes[k].is_compatible_with(tf.TensorShape([None])):
+      raise ValueError("Tensors to be packed must be one-dimensional.")
+
   # trim to length
   dataset = dataset.map(lambda x: {k: x[k][:length] for k in keys})
-
+  # Setting batch_size=length ensures that the concatenated sequences (if they
+  # have length >=1) are sufficient to fill at least one packed example.
   batch_size = length
   dataset = dataset.padded_batch(
       batch_size, padded_shapes={k: [-1] for k in keys})
   if use_custom_ops and len(keys) == 2:
-    # faster and better packing but requires custom-built binary.
-    k1, k2 = keys
-    def map_fn_custom(x):
-      """Map-function."""
-      (k1_packed, k1_segmengation, k1_position,
-       k2_packed, k2_segmentation, k2_position) = (
-           pack_sequences_ops.pack_sequences2(x[k1], x[k2], length))
-      packed = {
-          k1: k1_packed,
-          k1 + "_inputs": k1_segmengation,
-          k1 + "_position": k1_position,
-          k2: k2_packed,
-          k2 + "_inputs": k2_segmentation,
-          k2 + "_position": k2_position,
-      }
-      return tf.data.Dataset.from_tensor_slices(packed)
-    dataset = dataset.flat_map(map_fn_custom)
-    return dataset
+    # custom op only handles 2 keys.
+    # TODO(noam): support other numbers of keys.
+    return _pack_with_custom_ops(dataset, keys, length)
+  else:
+    return _pack_with_tf_ops(dataset, keys, length)
+
+
+def _pack_with_custom_ops(dataset, keys, length):
+  """Helper-function for packing a dataset which has already been batched.
+
+  See pack_dataset()
 
+  Relies on custom ops which require a custom compiled binary.
+  Faster than _pack_with_tf_ops(), and denser packing.
+
+  Args:
+    dataset: a dataset containing padded batches of examples.
+    keys: a list of strings (must have length 2)
+    length: an integer
+
+  Returns:
+    a dataset.
+  """
+  from tensor2tensor.data_generators.ops import pack_sequences_ops  # pylint: disable=g-import-not-at-top
+  # faster and better packing but requires custom-built binary.
+  k1, k2 = keys
+  def map_fn_custom(x):
+    """Map-function."""
+    (k1_packed, k1_segmengation, k1_position,
+     k2_packed, k2_segmentation, k2_position) = (
+         pack_sequences_ops.pack_sequences2(x[k1], x[k2], length))
+    packed = {
+        k1: k1_packed,
+        k1 + "_segmentation": k1_segmengation,
+        k1 + "_position": k1_position,
+        k2: k2_packed,
+        k2 + "_segmentation": k2_segmentation,
+        k2 + "_position": k2_position,
+    }
+    return tf.data.Dataset.from_tensor_slices(packed)
+  dataset = dataset.flat_map(map_fn_custom)
+  return dataset
+
+
+def _pack_with_tf_ops(dataset, keys, length):
+  """Helper-function for packing a dataset which has already been batched.
+
+  See pack_dataset()
+
+  Uses tf.while_loop.  Slow.
+
+  Args:
+    dataset: a dataset containing padded batches of examples.
+    keys: a list of strings
+    length: an integer
+
+  Returns:
+    a dataset.
+  """
   empty_example = {}
   for k in keys:
     empty_example[k] = tf.zeros([0], dtype=tf.int64)
diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
index b04d94356..4c07c7bdd 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
@@ -15,6 +15,8 @@ using ::tensorflow::Tensor;
 using ::tensorflow::TensorShape;
 using ::tensorflow::shape_inference::InferenceContext;
 
+// TODO(noam): this op packs a dataset of pairs of sequaneces (inputs, targets)
+// Generalize later to an arbitrary number of sequences.
 REGISTER_OP("PackSequences2")
     .Input("inputs: int64")
     .Input("targets: int64")

From 1d2d336d0f6b553448c3dd86035c8678b9d4c1b6 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 21 Feb 2019 13:06:49 -0800
Subject: [PATCH 1691/2720] Remove Modality.top_is_pointwise.

PiperOrigin-RevId: 235053495
---
 tensor2tensor/layers/modalities.py | 44 ++++++++++++++++--------------
 tensor2tensor/utils/modality.py    | 16 -----------
 tensor2tensor/utils/t2t_model.py   | 34 ++++++++---------------
 3 files changed, 35 insertions(+), 59 deletions(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index b481419f0..c93ec0109 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -32,6 +32,26 @@
 import tensorflow_probability as tfp
 
 
+def is_pointwise(func):
+  """Decorator for whether the function is pointwise.
+
+  An example of a pointwise function is a linear layer followed by
+  a softmax. Given a tensor [batch, length, height, depth] it operates
+  only on the last axis, on every point in [batch, length, height] fully
+  independently. In contrast, a classifier that first averages over length
+  and height is not pointwise, as it depends on the whole field. It is useful
+  to know if top functions are pointwise to speed up decoding in certain models.
+
+  Args:
+    func: Function to decorate.
+
+  Returns:
+    Original function with an attribute pointwise set to True.
+  """
+  func.pointwise = True
+  return func
+
+
 class SymbolModality(modality.Modality):
   """Modality for sets of discrete symbols.
 
@@ -46,10 +66,6 @@ class SymbolModality(modality.Modality):
   def name(model_hparams, vocab_size):
     return "symbol_modality_%d_%d" % (vocab_size, model_hparams.hidden_size)
 
-  @staticmethod
-  def top_is_pointwise():
-    return True
-
   @staticmethod
   def targets_weights_fn(model_hparams):
     weights_fn = common_layers.weights_nonzero
@@ -146,6 +162,7 @@ def targets_bottom(cls, x, model_hparams, vocab_size):
           x, model_hparams, vocab_size, "target_emb", reuse=None)
 
   @classmethod
+  @is_pointwise
   def top(cls, body_output, targets, model_hparams, vocab_size):
     """Generate logits.
 
@@ -204,6 +221,7 @@ def targets_bottom(x, model_hparams, vocab_size):
     return tf.one_hot(x, vocab_size)
 
   @staticmethod
+  @is_pointwise
   def top(body_output, _, model_hparams, vocab_size):
     return body_output
 
@@ -1021,10 +1039,6 @@ class RealModality(modality.Modality):
   * Top is a linear projection layer to vocab_size.
   """
 
-  @staticmethod
-  def top_is_pointwise():
-    return True
-
   @staticmethod
   def bottom(x, model_hparams, vocab_size):
     with tf.variable_scope("real"):
@@ -1032,6 +1046,7 @@ def bottom(x, model_hparams, vocab_size):
           tf.to_float(x), model_hparams.hidden_size, name="bottom")
 
   @staticmethod
+  @is_pointwise
   def top(body_output, _, model_hparams, vocab_size):
     with tf.variable_scope("real"):
       return tf.layers.dense(body_output, vocab_size, name="top")
@@ -1092,11 +1107,6 @@ def targets_bottom(cls, x, model_hparams, vocab_size):
     """SymbolModality overrides targets_bottom, so need to override here too."""
     return cls.bottom(x, model_hparams, vocab_size)
 
-  @staticmethod
-  def top_is_pointwise():
-    # pointwise mode manipulates body output, not logits, so it fails here.
-    return False
-
 
 class SigmoidClassLabelModality(ClassLabelModality):
   """Sigmoid cross-entropy for independent class labels."""
@@ -1332,11 +1342,3 @@ def get_top(modality_type, value=None):
     modality_cls = getattr(current_module, modality_type)
     return modality_cls.top
   return value
-
-
-def get_top_is_pointwise(modality_type, value=None):
-  """Gets whether default top is pointwise; if none available, return value."""
-  if modality_type in ModalityType.get_choices():
-    modality_cls = getattr(current_module, modality_type)
-    return modality_cls(None, None).top_is_pointwise
-  return value
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index 4518a729c..0a8925950 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -59,22 +59,6 @@ def name(cls, model_hparams, vocab_size=None):
     del model_hparams, vocab_size  # unused arg
     return misc_utils.camelcase_to_snakecase(type(cls).__name__)
 
-  @staticmethod
-  def top_is_pointwise():
-    """Whether the top mapping of the modality is pointwise.
-
-    An example of a pointwise top mapping is a linear layer followed by
-    a softmax. Given a tensor [batch, length, height, depth] it operates
-    only on the last axis, on every point in [batch, length, height] fully
-    independently. In contrast, a classifier that first averages over length
-    and height is not pointwise, as it depends on the whole field. It is useful
-    to know if a top is pointwise to speed up decoding in certain models.
-
-    Returns:
-      A Boolean, True if the modality is pointwise, False otherwise (default).
-    """
-    return False
-
   @staticmethod
   def targets_weights_fn(model_hparams):
     """The weights function to use for loss and eval metrics.
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index dfbbc2ab8..3f22eb7db 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -219,8 +219,6 @@ def __init__(self,
       hparams.add_hparam("targets_weights_fn", {})
     if not hasattr(hparams, "top"):
       hparams.add_hparam("top", {})
-    if not hasattr(hparams, "top_is_pointwise"):
-      hparams.add_hparam("top_is_pointwise", {})
     target_modalities = _create_target_modality(hparams.modality)
     for feature_name, modality in six.iteritems(hparams.modality):
       if modality in modalities.ModalityType.get_choices():
@@ -233,7 +231,6 @@ def __init__(self,
       hparams.name[feature_name] = modality.name
       hparams.targets_weights_fn[feature_name] = modality.targets_weights_fn
       hparams.top[feature_name] = modality.top
-      hparams.top_is_pointwise[feature_name] = modality.top_is_pointwise
 
     self._original_hparams = hparams
     self.set_mode(mode)
@@ -559,14 +556,12 @@ def _top_single(self, body_output, feature_name, features):
     with tf.variable_scope(name) as tm_vs:
       self._add_variable_scope(tm_vs.name, tm_vs)
       log_info("Transforming body output with %s.top", name)
-      top_is_pointwise = self._hparams.top_is_pointwise.get(
-          feature_name,
-          modalities.get_top_is_pointwise(modality))()
+      top = self._hparams.top.get(feature_name, modalities.get_top(modality))
+      top_is_pointwise = getattr(top, "pointwise", False)
       last_only = (top_is_pointwise and
                    self.hparams.mode == tf.estimator.ModeKeys.PREDICT and
                    not self.hparams.force_full_predict)
       if not last_only:
-        top = self._hparams.top.get(feature_name, modalities.get_top(modality))
         logits = top(body_output, features.get("targets"),
                      self._hparams, vocab_size)
       else:
@@ -587,7 +582,6 @@ def _top_single(self, body_output, feature_name, features):
           last_position_targets = tf.slice(
               features["targets"], [0, features["decode_loop_step"][0], 0, 0],
               [target_shape[0], 1, target_shape[2], target_shape[3]])
-        top = self._hparams.top.get(feature_name, modalities.get_top(modality))
         logits = top(last_position_body_output, last_position_targets,
                      self._hparams, vocab_size)
     return logits
@@ -897,11 +891,9 @@ def symbols_to_logits_fn(ids, i=None):
       # it has shape [batch_size] and contains floats between 0 and
       # source_length.
       if self._problem_hparams:
-        target_modality = self._problem_hparams.modality["targets"]
-        top_is_pointwise = self._hparams.top_is_pointwise.get(
-            "targets",
-            modalities.get_top_is_pointwise(target_modality))()
-        if top_is_pointwise:
+        modality = self._problem_hparams.modality["targets"]
+        top = self._hparams.top.get("targets", modalities.get_top(modality))
+        if getattr(top, "pointwise", False):
           return tf.squeeze(logits, axis=[1, 2, 3])
       # -1 due to the pad above.
       current_output_position = common_layers.shape_list(ids)[1] - 1
@@ -1039,11 +1031,10 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
       features["decode_loop_step"] = i
       samples, logits, losses = self.sample(features)
       # Concatenate the already-generated recent_output with last timestep
-      # of the newly-generated samples.
-      top_is_pointwise = self._hparams.top_is_pointwise.get(
-          "targets",
-          modalities.get_top_is_pointwise(target_modality))()
-      if top_is_pointwise:
+      # of the newly-generated samples.z
+      top = self._hparams.top.get("targets",
+                                  modalities.get_top(target_modality))
+      if getattr(top, "pointwise", False):
         cur_sample = samples[:, -1, :, :]
       else:
         cur_sample = samples[:, i, :, :]
@@ -1214,10 +1205,9 @@ def infer_step(recent_output, recent_logits, unused_loss):
       samples, logits, losses = self.sample(features)
       # Concatenate the already-generated recent_output with last timestep
       # of the newly-generated samples.
-      top_is_pointwise = self._hparams.top_is_pointwise.get(
-          "targets",
-          modalities.get_top_is_pointwise(target_modality))()
-      if top_is_pointwise:
+      top = self._hparams.top.get("targets",
+                                  modalities.get_top(target_modality))
+      if getattr(top, "pointwise", False):
         cur_sample = samples[:, -1, :, :]
       else:
         cur_sample = samples[:,

From 2dba035b4870020de004747271b4661c161da760 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 21 Feb 2019 14:17:42 -0800
Subject: [PATCH 1692/2720] EnvProblem.agent_id accepts int agent_id now. This
 will let us call code like env_problem.agent_id = counter; counter += 1

PiperOrigin-RevId: 235067205
---
 tensor2tensor/envs/env_problem.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index a0d2524d0..13d8d7baa 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -649,6 +649,8 @@ def agent_id(self):
 
   @agent_id.setter
   def agent_id(self, agent_id):
+    # Lets us call agent_id with integers that we increment.
+    agent_id = str(agent_id)
     # We use `-` in self.dataset_filename, disallow it here for convenience.
     if "-" in agent_id:
       raise ValueError("agent_id shouldn't have - in it.")

From cd977f58a6292caa2947bfa95d37da6b237220d8 Mon Sep 17 00:00:00 2001
From: Ben Goodrich <bgoodrich@google.com>
Date: Thu, 21 Feb 2019 15:31:17 -0800
Subject: [PATCH 1693/2720] Check for dtype.is_integer instead of int32

PiperOrigin-RevId: 235082354
---
 tensor2tensor/utils/beam_search.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index ba2cd999f..0632d3b12 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -153,11 +153,11 @@ def _gather(params, indices):
         gather_result = tf.cast(gather_result, dtype)
       return gather_result
 
-    # If the dtype is int32, use the gather instead of one_hot matmul to avoid
+    # If the dtype is int, use the gather instead of one_hot matmul to avoid
     # precision loss. The max int value can be represented by bfloat16 in MXU is
     # 256, which is smaller than the possible id values. Encoding/decoding can
     # potentially used to make it work, but the benenfit is small right now.
-    if dtype == tf.int32:
+    if dtype.is_integer:
       gather_result = tf.batch_gather(params, indices)
     else:
       gather_result = _gather(params, indices)

From e472b556710cb3b29eff69f2c766ce38fae91a36 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 21 Feb 2019 16:20:01 -0800
Subject: [PATCH 1694/2720] Return decoded strings from decode_once() if
 decode_hp.decode_in_memory is True.

PiperOrigin-RevId: 235091906
---
 tensor2tensor/utils/decoding.py | 46 ++++++++++++++++++++++++++++-----
 1 file changed, 39 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 1b69cbf77..2a171f63b 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -63,9 +63,9 @@ def decode_hparams(overrides=""):
       write_beam_scores=False,
       max_input_size=-1,
       identity_output=False,
-      num_samples=-1,
+      num_samples=-1,  # Number of examples to decode.
       delimiter="\n",
-      decode_to_file=None,
+      decode_to_file=None,  # str. Prefix for filename to write decodings to.
       decode_in_memory=False,
       # How much decode should wait for the next checkpoint
       decode_timeout_mins=240,
@@ -74,7 +74,7 @@ def decode_hparams(overrides=""):
       shard_id=0,  # Which shard are we decoding if more than 1 above.
       shards_start_offset=0,  # Number of the first shard to decode.
       shard_google_format=False,  # If True use Google shard naming format.
-      num_decodes=1,
+      num_decodes=1,  # Number of times to go over the dataset.
       force_decode_length=False,
       display_decoded_images=False,
       # Multi-problem decoding task id.
@@ -181,7 +181,7 @@ def decode_from_dataset(estimator,
   # We assume that worker_id corresponds to shard number.
   shard = decode_hp.shard_id if decode_hp.shards > 1 else None
 
-  # Setup decode output directory for any artifacts that may be written out
+  # Setup output directory for any artifacts that may be written out.
   output_dir = os.path.join(estimator.model_dir, "decode")
   tf.gfile.MakeDirs(output_dir)
 
@@ -218,7 +218,7 @@ def decode_from_dataset(estimator,
                          decode_hp,
                          decode_to_file,
                          output_dir,
-                         log_results=not decode_hp.decode_in_memory,
+                         log_results=True,
                          checkpoint_path=checkpoint_path)
 
     if decode_hp.decode_in_memory:
@@ -249,7 +249,30 @@ def decode_once(estimator,
                 output_dir,
                 log_results=True,
                 checkpoint_path=None):
-  """Decodes once."""
+  """Decodes once.
+
+  Args:
+    estimator: tf.estimator.Estimator instance. Used to generate encoded
+      predictions.
+    problem_name: str. Name of problem.
+    hparams: tf.HParams instance. HParams for model training.
+    infer_input_fn: zero-arg function. Input function for estimator.
+    decode_hp: tf.HParams instance. See decode_hparams() above.
+    decode_to_file: str. Prefix for filenames. Used to generated filenames to
+      which decoded predictions are written.
+    output_dir: str. Output directory. Only used for writing images.
+    log_results: bool. If False, return encoded predictions without any
+      further processing.
+    checkpoint_path: str. Path to load model checkpoint from. If unspecified,
+      Estimator's default is used.
+
+  Returns:
+    If decode_hp.decode_in_memory is True:
+      List of dicts, one per example. Values are either numpy arrays or decoded
+      strings.
+    If decode_hp.decode_in_memory is False:
+      An empty list.
+  """
 
   # Get the predictions as an iterable
   predictions = estimator.predict(infer_input_fn,
@@ -281,6 +304,10 @@ def decode_once(estimator,
   targets_vocab = problem_hparams.vocabulary["targets"]
 
   num_eval_samples = 0
+
+  # all_outputs[i][j] = (input: str, output: str, target: str). Input,
+  # decoded output, and target strings for example i, beam rank j.
+  all_outputs = []
   for num_predictions, prediction in enumerate(predictions):
     num_eval_samples += 1
     num_predictions += 1
@@ -289,8 +316,11 @@ def decode_once(estimator,
     outputs = prediction.get("outputs")
 
     # Log predictions
-    decoded_outputs = []
+    decoded_outputs = []  # [(str, str, str)]. See all_outputs above.
+    if decode_hp.decode_in_memory:
+      all_outputs.append(decoded_outputs)
     decoded_scores = []
+
     if decode_hp.return_beams:
       output_beams = np.split(outputs, decode_hp.beam_size, axis=0)
       scores = None
@@ -355,6 +385,8 @@ def decode_once(estimator,
     target_file.close()
     input_file.close()
 
+  return all_outputs
+
 
 def decode_from_file(estimator,
                      filename,

From fe9f2cf7ab82ce15e90e6dcacb5d4405dd30b695 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 21 Feb 2019 16:57:07 -0800
Subject: [PATCH 1695/2720] Remove v2 and guard against jax in t2t bin to make
 OSS work.

PiperOrigin-RevId: 235098195
---
 tensor2tensor/bin/t2t_trainer.py    | 40 ++++++++++++-----------------
 tensor2tensor/jax/input_pipeline.py | 26 +++++++++++--------
 tensor2tensor/jax/j2j_trainer.py    |  3 ---
 3 files changed, 32 insertions(+), 37 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index ead630db7..bb0ed809a 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -25,7 +25,7 @@
 from tensor2tensor import models  # pylint: disable=unused-import
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
 from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
-from tensor2tensor.jax import j2j
+
 from tensor2tensor.utils import cloud_mlengine
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
@@ -34,11 +34,16 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
-from tensor2tensor.v2 import t2t as t2t_v2
 import tensorflow as tf
 
 from tensorflow.contrib.tpu.python.tpu import tpu_config
 
+try:
+  from tensor2tensor.jax import j2j  # pylint: disable=g-import-not-at-top
+except TypeError:
+  pass
+
+
 flags = tf.flags
 FLAGS = flags.FLAGS
 
@@ -72,7 +77,6 @@
 flags.DEFINE_integer("intra_op_parallelism_threads", 0,
                      "Number of intra_op_parallelism_threads to use for CPU. "
                      "See TensorFlow config.proto for details.")
-flags.DEFINE_bool("v2", False, "Whether to use T2T v2.")
 flags.DEFINE_bool("jax", False, "Whether to use J2J.")
 # TODO(lukaszkaiser): resolve memory and variable assign issues and set to True.
 flags.DEFINE_bool(
@@ -362,35 +366,25 @@ def run_std_server():
 def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
 
-  if FLAGS.v2 or FLAGS.jax:
-    tf.enable_v2_behavior()
-    # Hacking main v1 flags to work with v2 and jax.
-    prefix = "t2t." if FLAGS.v2 else "j2j."
+  if FLAGS.jax:
+    # Hacking main v1 flags to work with jax.
     config_strs = []
     config_strs.append(
-        prefix + "train_fn.train_steps=" + str(FLAGS.train_steps))
+        "train_fn.train_steps=" + str(FLAGS.train_steps))
     config_strs.append(
-        prefix + "train_fn.eval_steps=" + str(FLAGS.eval_steps))
+        "train_fn.eval_steps=" + str(FLAGS.eval_steps))
     config_strs.append(
-        prefix + "train_fn.eval_frequency=" + str(FLAGS.local_eval_frequency))
+        "train_fn.eval_frequency=" + str(FLAGS.local_eval_frequency))
     if FLAGS.hparams:
       config_strs.extend(str(FLAGS.hparams).split(","))
-    config_str = "\n".join(config_strs)
     data_dir = os.path.expanduser(FLAGS.data_dir)
     output_dir = os.path.expanduser(FLAGS.output_dir)
 
-    if FLAGS.v2:
-      t2t_v2.t2t_train(FLAGS.model, FLAGS.problem,
-                       data_dir=data_dir, output_dir=output_dir,
-                       config_file=FLAGS.hparams_set, config=config_str)
-      return
-
-    if FLAGS.jax:
-      gin.bind_parameter("j2j.train_fn.dataset", FLAGS.problem)
-      config_strs += ["j2j.train_fn.model=@models." + FLAGS.model]
-      gin.parse_config_files_and_bindings(FLAGS.hparams_set, config_strs)
-      j2j.train_fn(data_dir, output_dir=output_dir)
-      return
+    gin.bind_parameter("train_fn.dataset", FLAGS.problem)
+    config_strs += ["train_fn.model=@models." + FLAGS.model]
+    gin.parse_config_files_and_bindings(FLAGS.hparams_set, config_strs)
+    j2j.train_fn(data_dir, output_dir=output_dir)
+    return
 
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
 
diff --git a/tensor2tensor/jax/input_pipeline.py b/tensor2tensor/jax/input_pipeline.py
index 64c7cac1b..18aa984a8 100644
--- a/tensor2tensor/jax/input_pipeline.py
+++ b/tensor2tensor/jax/input_pipeline.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""T2T models, configs and main training functions."""
+"""J2J input pipeline."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -32,7 +32,7 @@ def train_and_eval_dataset(dataset_name, data_dir):
   """Return train and evaluation datasets, feature info and supervised keys.
 
   Args:
-    dataset_name: a string, the name of the dataset; if it starts with "v1_"
+    dataset_name: a string, the name of the dataset; if it starts with "t2t_"
       then we'll search T2T Problem registry for it, otherwise we assume it
       is a dataset from TFDS and load it from there.
     data_dir: directory where the data is located.
@@ -46,8 +46,8 @@ def train_and_eval_dataset(dataset_name, data_dir):
      * supervised_keys: information what's the input and what's the target,
          ie., a pair of lists with input and target feature names.
   """
-  if dataset_name.startswith("v1_"):
-    return _train_and_eval_dataset_v1(dataset_name[3:], data_dir)
+  if dataset_name.startswith("t2t_"):
+    return _train_and_eval_dataset_v1(dataset_name[4:], data_dir)
   dataset_builder = tfds.builder(dataset_name, data_dir=data_dir)
   info = dataset_builder.info
   splits = dataset_builder.info.splits
@@ -98,9 +98,14 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
   hparams = problem.get_hparams()
   # We take a few training examples to guess the shapes.
   input_shapes, target_shapes = [], []
-  for example in train_dataset.take(3):
-    input_shapes.append(example["inputs"].shape.as_list())
-    target_shapes.append(example["targets"].shape.as_list())
+  example_tensor = train_dataset.make_one_shot_iterator().get_next()
+  sess = tf.Session()
+  example1 = sess.run(example_tensor)
+  example2 = sess.run(example_tensor)
+  example3 = sess.run(example_tensor)
+  for example in [example1, example2, example3]:
+    input_shapes.append(list(example["inputs"].shape))
+    target_shapes.append(list(example["targets"].shape))
   input_vocab_size = hparams.vocab_size["inputs"]
   target_vocab_size = hparams.vocab_size["targets"]
   input_info = _make_info(input_shapes, input_vocab_size)
@@ -109,13 +114,12 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
   return train_dataset, eval_dataset, info, supervised_keys
 
 
-@gin.configurable(whitelist=["max_target_length"])
+@gin.configurable(blacklist=["dataset", "training"])
 def preprocess_fn(dataset, training, max_target_length=-1):
   def target_right_length(_, target):
-    if max_target_length < 1 or not training:
-      return tf.constant(True)
     return tf.less(tf.shape(target)[0], max_target_length + 1)
-  dataset = dataset.filter(target_right_length)
+  if max_target_length > 0 and training:
+    dataset = dataset.filter(target_right_length)
   return dataset
 
 
diff --git a/tensor2tensor/jax/j2j_trainer.py b/tensor2tensor/jax/j2j_trainer.py
index 5cb081022..e6e7803d4 100644
--- a/tensor2tensor/jax/j2j_trainer.py
+++ b/tensor2tensor/jax/j2j_trainer.py
@@ -48,9 +48,6 @@
 flags.DEFINE_multi_string("config", None,
                           "Configuration parameters (gin string).")
 
-# For iterators over datasets so we can do "for example in dataset".
-tf.enable_v2_behavior()
-
 
 def j2j_train(model_name, dataset_name,
               data_dir=None, output_dir=None, config_file=None, config=None):

From 0fb50f3ed6052bf76af341770fdd7fa698064ba9 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 21 Feb 2019 19:12:38 -0800
Subject: [PATCH 1696/2720] SVG FP and LP implementations within emily.py

PiperOrigin-RevId: 235116515
---
 tensor2tensor/models/__init__.py     |   1 -
 tensor2tensor/models/video/emily.py  | 182 ++++++++++---
 tensor2tensor/models/video/svg_lp.py | 366 ---------------------------
 3 files changed, 149 insertions(+), 400 deletions(-)
 delete mode 100644 tensor2tensor/models/video/svg_lp.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 475da8231..b9fd5bf47 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -78,7 +78,6 @@
 from tensor2tensor.models.video import next_frame_glow
 from tensor2tensor.models.video import savp
 from tensor2tensor.models.video import sv2p
-from tensor2tensor.models.video import svg_lp
 
 from tensor2tensor.utils import registry
 
diff --git a/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index 8d0ab7953..58afe4890 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -205,7 +205,8 @@ def stacked_lstm(self, inputs, states, hidden_size, output_size, nlayers):
         net, output_size, activation=tf.nn.tanh, name="af2")
     return net, states
 
-  def lstm_gaussian(self, inputs, states, hidden_size, output_size, nlayers):
+  def lstm_gaussian(self, inputs, states, hidden_size, output_size, nlayers,
+                    name):
     """Stacked LSTM layers with FC layer as input and gaussian as output.
 
     Args:
@@ -214,18 +215,19 @@ def lstm_gaussian(self, inputs, states, hidden_size, output_size, nlayers):
       hidden_size: number of lstm units
       output_size: size of the output
       nlayers: number of lstm layers
+      name: the lstm name for scope definition
     Returns:
       mu: mean of the predicted gaussian
       logvar: log(var) of the predicted gaussian
       skips: a list of updated lstm states for each layer
     """
     net = inputs
-    net = tfl.dense(net, hidden_size, activation=None, name="bf1")
+    net = tfl.dense(net, hidden_size, activation=None, name="%sf1"%name)
     for i in range(nlayers):
       net, states[i] = common_video.basic_lstm(
-          net, states[i], hidden_size, name="blstm%d"%i)
-    mu = tfl.dense(net, output_size, activation=None, name="bf2mu")
-    logvar = tfl.dense(net, output_size, activation=None, name="bf2log")
+          net, states[i], hidden_size, name="%slstm%d"%(name, i))
+    mu = tfl.dense(net, output_size, activation=None, name="%sf2mu"%name)
+    logvar = tfl.dense(net, output_size, activation=None, name="%sf2log"%name)
     return mu, logvar, states
 
   def construct_model(self, images, actions, rewards):
@@ -258,6 +260,7 @@ def construct_model(self, images, actions, rewards):
     z_dim = self.hparams.z_dim
     g_dim = self.hparams.g_dim
     rnn_size = self.hparams.rnn_size
+    prior_rnn_layers = self.hparams.prior_rnn_layers
     posterior_rnn_layers = self.hparams.posterior_rnn_layers
     predictor_rnn_layers = self.hparams.predictor_rnn_layers
     context_frames = self.hparams.video_num_input_frames
@@ -266,8 +269,9 @@ def construct_model(self, images, actions, rewards):
     seq_len, batch_size, _, _, color_channels = common_layers.shape_list(images)
 
     # LSTM initial sizesstates.
-    predictor_states = [None] * predictor_rnn_layers
+    prior_states = [None] * prior_rnn_layers
     posterior_states = [None] * posterior_rnn_layers
+    predictor_states = [None] * predictor_rnn_layers
 
     tf.logging.info(">>>> Encoding")
     # Encoding:
@@ -275,58 +279,168 @@ def construct_model(self, images, actions, rewards):
     images = tf.unstack(images, axis=0)
     for i, image in enumerate(images):
       with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
-        enc, skips = self.encoder(image, rnn_size, has_batchnorm=has_batchnorm)
+        enc, skips = self.encoder(image, g_dim, has_batchnorm=has_batchnorm)
         enc = tfl.flatten(enc)
         enc_images.append(enc)
         enc_skips.append(skips)
 
     tf.logging.info(">>>> Prediction")
     # Prediction
-    pred_enc, pred_mu, pred_logvar = [], [], []
+    pred_mu_pos = []
+    pred_logvar_pos = []
+    pred_mu_prior = []
+    pred_logvar_prior = []
+    gen_images = []
     for i in range(1, seq_len):
-      with tf.variable_scope("prediction", reuse=tf.AUTO_REUSE):
+      with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
         # current encoding
-        h_current = enc_images[i-1]
+        if self.is_training or len(gen_images) < context_frames:
+          h_current = enc_images[i - 1]
+        else:
+          h_current, _ = self.encoder(gen_images[-1], g_dim)
+          h_current = tfl.flatten(h_current)
+
         # target encoding
         h_target = enc_images[i]
 
-        z = tf.random_normal([batch_size, z_dim], 0, 1, dtype=tf.float32)
-        mu, logvar = tf.zeros_like(z), tf.zeros_like(z)
+      with tf.variable_scope("prediction", reuse=tf.AUTO_REUSE):
+        # Prior parameters
+        if self.hparams.learned_prior:
+          mu_prior, logvar_prior, prior_states = self.lstm_gaussian(
+              h_current, prior_states, rnn_size, z_dim, prior_rnn_layers,
+              "prior")
+        else:
+          mu_prior = tf.zeros((batch_size, z_dim))
+          logvar_prior = tf.zeros((batch_size, z_dim))
 
         # Only use Posterior if it's training time
-        if self.hparams.mode == tf.estimator.ModeKeys.TRAIN:
-          mu, logvar, posterior_states = self.lstm_gaussian(
-              h_target, posterior_states, rnn_size, z_dim, posterior_rnn_layers)
-
+        if self.is_training or len(gen_images) < context_frames:
+          mu_pos, logvar_pos, posterior_states = self.lstm_gaussian(
+              h_target, posterior_states, rnn_size, z_dim, posterior_rnn_layers,
+              "posterior")
           # Sample z from posterior distribution
-          z = z * tf.exp(tf.multiply(0.5, logvar)) + mu
+          z = common_video.get_gaussian_tensor(mu_pos, logvar_pos)
+        else:
+          mu_pos = tf.zeros_like(mu_prior)
+          logvar_pos = tf.zeros_like(logvar_prior)
+          z = common_video.get_gaussian_tensor(mu_prior, logvar_prior)
 
         # Predict output encoding
         h_pred, predictor_states = self.stacked_lstm(
             tf.concat([h_current, z], axis=1),
             predictor_states, rnn_size, g_dim, predictor_rnn_layers)
 
-        pred_enc.append(h_pred)
-        pred_mu.append(mu)
-        pred_logvar.append(logvar)
+        pred_mu_pos.append(tf.identity(mu_pos, "mu_pos"))
+        pred_logvar_pos.append(tf.identity(logvar_pos, "logvar_pos"))
+        pred_mu_prior.append(tf.identity(mu_prior, "mu_prior"))
+        pred_logvar_prior.append(tf.identity(logvar_prior, "logvar_prior"))
 
-    tf.logging.info(">>>> Decoding")
-    # Decoding
-    gen_images = []
-    for i in range(seq_len-1):
       with tf.variable_scope("decoding", reuse=tf.AUTO_REUSE):
-        # use skip values of last available frame
-        skip_index = min(context_frames-1, i)
-
-        h_pred = tf.reshape(pred_enc[i], [batch_size, 1, 1, g_dim])
-        x_pred = self.decoder(
-            h_pred, color_channels, enc_skips[skip_index],
-            has_batchnorm=has_batchnorm)
+        skip_index = min(context_frames-1, i-1)
+        h_pred = tf.reshape(h_pred, [batch_size, 1, 1, g_dim])
+        if self.hparams.has_skips:
+          x_pred = self.decoder(
+              h_pred, color_channels,
+              skips=enc_skips[skip_index], has_batchnorm=has_batchnorm)
+        else:
+          x_pred = self.decoder(
+              h_pred, color_channels, has_batchnorm=has_batchnorm)
         gen_images.append(x_pred)
 
     tf.logging.info(">>>> Done")
     gen_images = tf.stack(gen_images, axis=0)
-    return gen_images, fake_reward_prediction, pred_mu, pred_logvar
+    return {"gen_images": gen_images,
+            "fake_reward_prediction": fake_reward_prediction,
+            "pred_mu_pos": pred_mu_pos,
+            "pred_logvar_pos": pred_logvar_pos,
+            "pred_mu_prior": pred_mu_prior,
+            "pred_logvar_prior": pred_logvar_prior}
+
+  def get_extra_loss(self,
+                     latent_means_pos, latent_logvars_pos,
+                     latent_means_prior, latent_logvars_prior):
+    """Losses in addition to the default modality losses."""
+    return self.get_kl_loss(
+        latent_means_pos, latent_logvars_pos,
+        latent_means_prior, latent_logvars_prior)
+
+  def body(self, features):
+    hparams = self.hparams
+    batch_size = common_layers.shape_list(features["inputs"])[0]
+
+    # Swap time and batch axes.
+    input_frames = common_video.swap_time_and_batch_axes(features["inputs"])
+    target_frames = common_video.swap_time_and_batch_axes(features["targets"])
+
+    # Get actions if exist otherwise use zeros
+    input_actions = self.get_input_if_exists(
+        features, "input_action", batch_size, hparams.video_num_input_frames)
+    target_actions = self.get_input_if_exists(
+        features, "target_action", batch_size, hparams.video_num_target_frames)
+
+    # Get rewards if exist otherwise use zeros
+    input_rewards = self.get_input_if_exists(
+        features, "input_reward", batch_size, hparams.video_num_input_frames)
+    target_rewards = self.get_input_if_exists(
+        features, "target_reward", batch_size, hparams.video_num_target_frames)
+
+    all_actions = tf.concat([input_actions, target_actions], axis=0)
+    all_rewards = tf.concat([input_rewards, target_rewards], axis=0)
+    all_frames = tf.concat([input_frames, target_frames], axis=0)
+
+    # Each image is being used twice, in latent tower and main tower.
+    # This is to make sure we are using the *same* image for both, ...
+    # ... given how TF queues work.
+    # NOT sure if this is required at all. Doesn"t hurt though! :)
+    all_frames = tf.identity(all_frames)
+
+    retvals = self.construct_model(
+        images=all_frames, actions=all_actions, rewards=all_rewards)
+
+    # retrieve tensors returned by the model contructor
+    gen_images = retvals["gen_images"]
+    gen_rewards = retvals["fake_reward_prediction"]
+    latent_means_pos = retvals["pred_mu_pos"]
+    latent_logvars_pos = retvals["pred_logvar_pos"]
+    latent_means_prior = retvals["pred_mu_prior"]
+    latent_logvars_prior = retvals["pred_logvar_prior"]
+
+    extra_loss = self.get_extra_loss(
+        latent_means_pos=latent_means_pos,
+        latent_logvars_pos=latent_logvars_pos,
+        latent_means_prior=latent_means_prior,
+        latent_logvars_prior=latent_logvars_prior)
+
+    # Visualize predictions in Tensorboard
+    if self.is_training:
+      self.visualize_predictions(all_frames[1:], gen_images)
+
+    # Ignore the predictions from the input frames.
+    # This is NOT the same as original paper/implementation.
+    predictions = gen_images[hparams.video_num_input_frames-1:]
+    reward_pred = gen_rewards[hparams.video_num_input_frames-1:]
+    reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove extra dimension.
+
+    # Swap back time and batch axes.
+    predictions = common_video.swap_time_and_batch_axes(predictions)
+    reward_pred = common_video.swap_time_and_batch_axes(reward_pred)
+
+    if self.is_training and hparams.internal_loss:
+      # add the loss for input frames as well.
+      extra_gts = all_frames[1:hparams.video_num_input_frames]
+      extra_gts = common_video.swap_time_and_batch_axes(extra_gts)
+      extra_pds = gen_images[:hparams.video_num_input_frames-1]
+      extra_pds = common_video.swap_time_and_batch_axes(extra_pds)
+      extra_raw_gts = features["inputs_raw"][:, 1:]
+      recon_loss = self.get_extra_internal_loss(
+          extra_raw_gts, extra_gts, extra_pds)
+      extra_loss += recon_loss
+
+    return_targets = predictions
+    if hparams.reward_prediction:
+      return_targets = {"targets": predictions, "target_reward": reward_pred}
+
+    return return_targets, extra_loss
 
 
 @registry.register_hparams
@@ -349,9 +463,11 @@ def next_frame_emily():
   hparams.optimizer_adam_epsilon = 1e-08
   hparams.anneal_end = -1
   hparams.clip_grad_norm = 5.0
-  hparams.add_hparam("z_dim", 10)
+  hparams.add_hparam("learned_prior", True)
+  hparams.add_hparam("z_dim", 64)
   hparams.add_hparam("g_dim", 128)
   hparams.add_hparam("rnn_size", 256)
+  hparams.add_hparam("prior_rnn_layers", 1)
   hparams.add_hparam("posterior_rnn_layers", 1)
   hparams.add_hparam("predictor_rnn_layers", 2)
   hparams.add_hparam("has_skips", True)
diff --git a/tensor2tensor/models/video/svg_lp.py b/tensor2tensor/models/video/svg_lp.py
deleted file mode 100644
index f495fb9ff..000000000
--- a/tensor2tensor/models/video/svg_lp.py
+++ /dev/null
@@ -1,366 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Model architecture for video prediction model.
-
-   based on following paper:
-   "Stochastic Video Generation with a Learned Prior"
-   https://arxiv.org/pdf/1802.07687.pdf
-   by Emily Denton and Rob Fergus.
-
-   This code is a translation of the original code from PyTorch:
-   https://github.com/edenton/svg
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.layers import common_layers
-from tensor2tensor.layers import common_video
-from tensor2tensor.models.video import emily
-from tensor2tensor.models.video import sv2p_params
-from tensor2tensor.utils import registry
-
-import tensorflow as tf
-
-tfl = tf.layers
-tfcl = tf.contrib.layers
-
-
-@registry.register_model
-class NextFrameSVGLP(emily.NextFrameEmily):
-  """Stochastic Variational Video Prediction With Learned Prior."""
-
-  def rnn_model(self, hidden_size, nlayers, rnn_type, name):
-    """Stacked RNN cell constructor.
-
-    Args:
-      hidden_size: number of lstm units
-      nlayers: number of lstm layers
-      rnn_type: type of RNN cell to use
-      name: RNN name
-    Returns:
-      stacked_rnn: stacked RNN cell
-    """
-    layers_units = [hidden_size] * nlayers
-    if rnn_type == "lstm":
-      rnn_cell = tf.nn.rnn_cell.LSTMCell
-    elif rnn_type == "gru":
-      rnn_cell = tf.nn.rnn_cell.GRUCell
-    else:
-      rnn_cell = tf.nn.rnn_cell.RNNCell
-    cells = [rnn_cell(units, name=name) for units in layers_units]
-    stacked_rnn = tf.nn.rnn_cell.MultiRNNCell(cells)
-    return stacked_rnn
-
-  def deterministic_rnn(self, cell, inputs, states, output_size, scope):
-    """Deterministic RNN step function.
-
-    Args:
-      cell: RNN cell to forward through
-      inputs: input to RNN cell
-      states: previous RNN state
-      output_size: size of the output
-      scope: scope of the current RNN forward computation parameters
-    Returns:
-      outputs: deterministic RNN output vector
-      states: updated RNN states
-    """
-    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
-      embedded = tfl.dense(
-          inputs, cell.output_size, activation=tf.nn.relu, name="embed")
-      hidden, states = cell(embedded, states)
-      outputs = tfl.dense(
-          hidden, output_size, activation=tf.nn.relu, name="output")
-
-    return outputs, states
-
-  def gaussian_rnn(self, cell, inputs, states, output_size, scope):
-    """Deterministic RNN step function.
-
-    Args:
-      cell: RNN cell to forward through
-      inputs: input to RNN cell
-      states: previous RNN state
-      output_size: size of the output
-      scope: scope of the current RNN forward computation parameters
-    Returns:
-      mu: mean of the predicted gaussian
-      logvar: log(var) of the predicted gaussian
-      states: updated RNN states
-    """
-    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
-      embedded = tfl.dense(
-          inputs, cell.output_size, activation=tf.nn.relu, name="embed")
-      hidden, states = cell(embedded, states)
-      mu = tfl.dense(
-          hidden, output_size, activation=None, name="mu")
-      logvar = tfl.dense(
-          hidden, output_size, activation=None, name="logvar")
-
-    return mu, logvar, states
-
-  def sample(self, mu, logvar):
-    eps = tf.random_normal([self.hparams.batch_size, self.hparams.z_dim], 0, 1)
-    sigma = tf.exp(tf.multiply(0.5, logvar))
-    z = tf.add(mu, tf.multiply(sigma, eps))
-
-    return z
-
-  def construct_model(self, images, actions, rewards):
-    """Builds the stochastic model.
-
-    The model first encodes all the images (x_t) in the sequence
-    using the encoder. Let"s call the output e_t. Then it predicts the
-    latent state of the next frame using a recurrent posterior network
-    z ~ q(z|e_{0:t}) = N(mu(e_{0:t}), sigma(e_{0:t})).
-    Another recurrent network predicts the embedding of the next frame
-    using the approximated posterior e_{t+1} = p(e_{t+1}|e_{0:t}, z)
-    Finally, the decoder decodes e_{t+1} into x_{t+1}.
-    Skip connections from encoder to decoder help with reconstruction.
-
-    Args:
-      images: tensor of ground truth image sequences
-      actions: NOT used list of action tensors
-      rewards: NOT used list of reward tensors
-
-    Returns:
-      gen_images: generated images
-      fakr_rewards: input rewards as reward prediction!
-      pred_mu: predited means of posterior
-      pred_logvar: predicted log(var) of posterior
-    """
-    # model does not support action conditioned and reward prediction
-    fake_reward_prediction = rewards
-    del actions, rewards
-
-    mode = self.hparams.mode
-    z_dim = self.hparams.z_dim
-    g_dim = self.hparams.g_dim
-    rnn_size = self.hparams.rnn_size
-    rnn_type = self.hparams.rnn_type
-    prior_rnn_layers = self.hparams.prior_rnn_layers
-    posterior_rnn_layers = self.hparams.posterior_rnn_layers
-    predictor_rnn_layers = self.hparams.predictor_rnn_layers
-    context_frames = self.hparams.video_num_input_frames
-    has_batchnorm = self.hparams.has_batchnorm
-
-    # Create RNN cells
-    predictor_cell = self.rnn_model(
-        rnn_size, predictor_rnn_layers, rnn_type, "frame_predictor")
-    prior_cell = self.rnn_model(
-        rnn_size, prior_rnn_layers, rnn_type, "prior")
-    posterior_cell = self.rnn_model(
-        rnn_size, posterior_rnn_layers, rnn_type, "posterior")
-
-    seq_len, batch_size, _, _, color_channels = common_layers.shape_list(images)
-
-    # RNN initialize states.
-    prior_states = prior_cell.zero_state(batch_size, tf.float32)
-    predictor_states = predictor_cell.zero_state(batch_size, tf.float32)
-    posterior_states = posterior_cell.zero_state(batch_size, tf.float32)
-
-    tf.logging.info(">>>> Encoding")
-    # Encoding:
-    enc_images, enc_skips = [], []
-    images = tf.unstack(images, axis=0)
-    for i, image in enumerate(images):
-      with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
-        enc, skips = self.encoder(image, g_dim, has_batchnorm=has_batchnorm)
-        enc = tfl.flatten(enc)
-        enc_images.append(enc)
-        enc_skips.append(skips)
-
-    tf.logging.info(">>>> Prediction")
-    # Prediction
-    pred_mu = []
-    pred_logvar = []
-    pred_mu_p = []
-    pred_logvar_p = []
-    gen_images = []
-    for i in range(1, seq_len):
-      with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
-        # current encoding
-        if (mode == tf.estimator.ModeKeys.TRAIN or
-            len(gen_images) < context_frames):
-          h_current = enc_images[i-1]
-        else:
-          h_current, _ = self.encoder(gen_images[-1], g_dim)
-          h_current = tfl.flatten(h_current)
-
-        # target encoding
-        h_target = enc_images[i]
-
-      with tf.variable_scope("prediction", reuse=tf.AUTO_REUSE):
-        # Prior parameters
-        mu_p, logvar_p, prior_states = self.gaussian_rnn(
-            prior_cell, h_current, prior_states, z_dim, "prior")
-
-        # Only use Posterior if it's training time
-        if mode == tf.estimator.ModeKeys.TRAIN:
-          mu, logvar, posterior_states = self.gaussian_rnn(
-              posterior_cell, h_target, posterior_states, z_dim, "posterior")
-          z = self.sample(mu, logvar)
-        else:
-          mu = tf.zeros_like(mu_p)
-          logvar = tf.zeros_like(logvar_p)
-          z = self.sample(mu_p, logvar_p)
-
-        # Predict output images
-        h_pred, predictor_states = self.deterministic_rnn(
-            predictor_cell, tf.concat([h_current, z], axis=1),
-            predictor_states, g_dim, "predictor")
-
-        pred_mu.append(tf.identity(mu, "mu"))
-        pred_logvar.append(tf.identity(logvar, "logvar"))
-        pred_mu_p.append(tf.identity(mu_p, "mu_p"))
-        pred_logvar_p.append(tf.identity(logvar_p, "log_var_p"))
-
-      with tf.variable_scope("decoding", reuse=tf.AUTO_REUSE):
-        skip_index = min(context_frames-1, i-1)
-        h_pred = tf.reshape(h_pred, [batch_size, 1, 1, g_dim])
-        if self.hparams.has_skips:
-          x_pred = self.decoder(
-              h_pred, color_channels,
-              skips=enc_skips[skip_index], has_batchnorm=has_batchnorm)
-        else:
-          x_pred = self.decoder(
-              h_pred, color_channels, has_batchnorm=has_batchnorm)
-        gen_images.append(x_pred)
-
-    tf.logging.info(">>>> Done")
-    gen_images = tf.stack(gen_images, axis=0)
-    return (gen_images, fake_reward_prediction,
-            pred_mu, pred_logvar, pred_mu_p, pred_logvar_p)
-
-  def get_extra_loss(self,
-                     latent_means, latent_logvars,
-                     latent_means_p, latent_logvars_p):
-    """Losses in addition to the default modality losses."""
-    return self.get_kl_loss(
-        latent_means, latent_logvars, latent_means_p, latent_logvars_p)
-
-  def body(self, features):
-    hparams = self.hparams
-    batch_size = common_layers.shape_list(features["inputs"])[0]
-
-    # Swap time and batch axes.
-    input_frames = common_video.swap_time_and_batch_axes(features["inputs"])
-    target_frames = common_video.swap_time_and_batch_axes(features["targets"])
-
-    # Get actions if exist otherwise use zeros
-    input_actions = self.get_input_if_exists(
-        features, "input_action", batch_size, hparams.video_num_input_frames)
-    target_actions = self.get_input_if_exists(
-        features, "target_action", batch_size, hparams.video_num_target_frames)
-
-    # Get rewards if exist otherwise use zeros
-    input_rewards = self.get_input_if_exists(
-        features, "input_reward", batch_size, hparams.video_num_input_frames)
-    target_rewards = self.get_input_if_exists(
-        features, "target_reward", batch_size, hparams.video_num_target_frames)
-
-    all_actions = tf.concat([input_actions, target_actions], axis=0)
-    all_rewards = tf.concat([input_rewards, target_rewards], axis=0)
-    all_frames = tf.concat([input_frames, target_frames], axis=0)
-
-    # Each image is being used twice, in latent tower and main tower.
-    # This is to make sure we are using the *same* image for both, ...
-    # ... given how TF queues work.
-    # NOT sure if this is required at all. Doesn"t hurt though! :)
-    all_frames = tf.identity(all_frames)
-
-    retvals = self.construct_model(
-        images=all_frames, actions=all_actions, rewards=all_rewards)
-
-    # retrieve tensors returned by the model contructor
-    gen_images = retvals[0]
-    gen_rewards = retvals[1]
-    latent_means = retvals[2]
-    latent_logvars = retvals[3]
-    latent_means_p = retvals[4]
-    latent_logvars_p = retvals[5]
-
-    extra_loss = self.get_extra_loss(
-        latent_means=latent_means,
-        latent_logvars=latent_logvars,
-        latent_means_p=latent_means_p,
-        latent_logvars_p=latent_logvars_p)
-
-    # Visualize predictions in Tensorboard
-    if self.is_training:
-      self.visualize_predictions(all_frames[1:], gen_images)
-
-    # Ignore the predictions from the input frames.
-    # This is NOT the same as original paper/implementation.
-    predictions = gen_images[hparams.video_num_input_frames-1:]
-    reward_pred = gen_rewards[hparams.video_num_input_frames-1:]
-    reward_pred = tf.squeeze(reward_pred, axis=2)  # Remove extra dimension.
-
-    # Swap back time and batch axes.
-    predictions = common_video.swap_time_and_batch_axes(predictions)
-    reward_pred = common_video.swap_time_and_batch_axes(reward_pred)
-
-    if self.is_training and hparams.internal_loss:
-      # add the loss for input frames as well.
-      extra_gts = all_frames[1:hparams.video_num_input_frames]
-      extra_gts = common_video.swap_time_and_batch_axes(extra_gts)
-      extra_pds = gen_images[:hparams.video_num_input_frames-1]
-      extra_pds = common_video.swap_time_and_batch_axes(extra_pds)
-      extra_raw_gts = features["inputs_raw"][:, 1:]
-      recon_loss = self.get_extra_internal_loss(
-          extra_raw_gts, extra_gts, extra_pds)
-      extra_loss += recon_loss
-
-    return_targets = predictions
-    if hparams.reward_prediction:
-      return_targets = {"targets": predictions, "target_reward": reward_pred}
-
-    return return_targets, extra_loss
-
-
-@registry.register_hparams
-def next_frame_svglp():
-  """SVG with learned prior model hparams."""
-  hparams = sv2p_params.next_frame_sv2p()
-  hparams.video_num_input_frames = 2
-  hparams.video_num_target_frames = 10
-  hparams.learning_rate_constant = 1e-4
-  seq_length = hparams.video_num_input_frames + hparams.video_num_target_frames
-  # The latent_loss_multiplier is divided by the number of frames because
-  # the image sequence loss in t2t is averaged instead of added through
-  # time as they do in the SVG-LP paper
-  hparams.latent_loss_multiplier = 1e-4 / seq_length
-  hparams.reward_prediction = False
-  hparams.num_iterations_1st_stage = -1
-  hparams.num_iterations_2nd_stage = -1
-  hparams.optimizer_adam_beta1 = 0.9
-  hparams.optimizer_adam_beta2 = 0.999
-  hparams.optimizer_adam_epsilon = 1e-08
-  hparams.anneal_end = -1
-  hparams.clip_grad_norm = 5.0
-  hparams.add_hparam("learned_prior", True)
-  hparams.add_hparam("z_dim", 64)
-  hparams.add_hparam("g_dim", 128)
-  hparams.add_hparam("rnn_size", 256)
-  hparams.add_hparam("rnn_type", "lstm")
-  hparams.add_hparam("prior_rnn_layers", 1)
-  hparams.add_hparam("posterior_rnn_layers", 1)
-  hparams.add_hparam("predictor_rnn_layers", 2)
-  hparams.add_hparam("has_skips", True)
-  hparams.add_hparam("has_batchnorm", True)
-  return hparams
-

From 4829943e1368921c44c54479d8a435098bf814e8 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 21 Feb 2019 22:22:58 -0800
Subject: [PATCH 1697/2720] Add a reward and done property to a trajectory,
 this comes in use when we want to calculate total reward over a trajectory in
 evaluation.

PiperOrigin-RevId: 235133956
---
 tensor2tensor/envs/trajectory.py      | 17 +++++++++++++++++
 tensor2tensor/envs/trajectory_test.py | 16 ++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index 0aee091df..59fc3e1c7 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -74,6 +74,23 @@ def is_active(self):
   def time_steps(self):
     return self._time_steps
 
+  @property
+  def done(self):
+    return self.is_active() and self.last_time_step().done
+
+  # TODO(afrozm): Add discounting and rewards-to-go when it makes sense.
+  @property
+  def reward(self):
+    """Returns a tuple of sum of raw and processed rewards."""
+    raw_rewards, processed_rewards = 0, 0
+    for ts in self.time_steps:
+      # NOTE: raw_reward and processed_reward are None for the first time-step.
+      if ts.raw_reward is not None:
+        raw_rewards += ts.raw_reward
+      if ts.processed_reward is not None:
+        processed_rewards += ts.processed_reward
+    return raw_rewards, processed_rewards
+
 
 class BatchTrajectory(object):
   """Basically a batch of active trajectories and a list of completed ones."""
diff --git a/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index ec577716d..36ab17cac 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -30,6 +30,7 @@ def test_empty_trajectory(self):
     t = trajectory.Trajectory()
     self.assertFalse(t.is_active())
     self.assertEqual(0, t.num_time_steps())
+    self.assertFalse(t.done)
 
   def test_add_time_step(self):
     t = trajectory.Trajectory()
@@ -74,6 +75,21 @@ def test_change_last_time_step(self):
     # Assert on the number of steps remaining the same as before.
     self.assertEqual(num_ts_old, t.num_time_steps())
 
+  def test_reward(self):
+    t = trajectory.Trajectory()
+    # first time-step doesn't have rewards, since they are on entering a state.
+    t.add_time_step(
+        observation=1, raw_reward=None, processed_reward=None, done=False)
+    t.add_time_step(
+        observation=2, raw_reward=2, processed_reward=200, done=False)
+    t.add_time_step(
+        observation=3, raw_reward=3, processed_reward=300, done=True)
+
+    raw_reward, processed_reward = t.reward
+
+    self.assertEqual(5, raw_reward)
+    self.assertEqual(500, processed_reward)
+
 
 class BatchTrajectoryTest(tf.test.TestCase):
 

From 189d04e2391ef4991bd034f9c998323e9f0df924 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 21 Feb 2019 22:53:14 -0800
Subject: [PATCH 1698/2720] Simplify Modality.targets_weights_fn(hparams) to
 Modality.targets_weights_fn.

To do so, I cleaned up Modality.loss to always require weights_fn as an argument. I also moved the only logic requiring hparams into T2TModel.

PiperOrigin-RevId: 235136229
---
 .../data_generators/multi_problem.py          |   8 +-
 tensor2tensor/layers/modalities.py            | 107 +++++++-----------
 tensor2tensor/layers/modalities_test.py       |   6 +-
 tensor2tensor/models/video/base.py            |   2 +-
 tensor2tensor/utils/metrics.py                |   4 +-
 tensor2tensor/utils/modality.py               |  31 +----
 tensor2tensor/utils/t2t_model.py              |  31 +++--
 7 files changed, 81 insertions(+), 108 deletions(-)

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 5f165a2b5..13ec711ae 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -438,11 +438,13 @@ def aggregate_task_losses(hparams,
     vocab_size += (-vocab_size) % hparams.vocab_divisor
   modality = problem_hparams.modality[feature_name]
   loss = hparams.loss.get(feature_name, modalities.get_loss(modality))
+  weights_fn = hparams.targets_weights_fn.get(
+      feature_name, modalities.get_targets_weights_fn(modality))
   # Primary task loss
   loss_num, loss_den = loss(
       logits, feature,
       lambda x: common_layers.weights_multi_problem_all(x, main_task_id),
-      hparams, vocab_size)
+      hparams, vocab_size, weights_fn)
 
   loss_val = loss_num / tf.maximum(1.0, loss_den)
   summaries.append([hparams.problem.task_list[0].name+"_loss", loss_val])
@@ -532,13 +534,15 @@ def aggregate_task_lm_losses(hparams,
     vocab_size += (-vocab_size) % hparams.vocab_divisor
   modality = problem_hparams.modality[feature_name]
   loss = hparams.loss.get(feature_name, modalities.get_loss(modality))
+  weights_fn = hparams.targets_weights_fn.get(
+      feature_name, modalities.get_targets_weights_fn(modality))
   loss_num = 0.
   loss_den = 0.
   for task in hparams.problem.task_list:
     loss_num_, loss_den_ = loss(
         logits, feature,
         lambda x: common_layers.weights_multi_problem_all(x, task.task_id),  # pylint: disable=cell-var-from-loop
-        hparams, vocab_size)
+        hparams, vocab_size, weights_fn)
 
     loss_num += loss_num_
     loss_den += loss_den_
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index c93ec0109..cbcd62c7c 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -66,24 +66,7 @@ class SymbolModality(modality.Modality):
   def name(model_hparams, vocab_size):
     return "symbol_modality_%d_%d" % (vocab_size, model_hparams.hidden_size)
 
-  @staticmethod
-  def targets_weights_fn(model_hparams):
-    weights_fn = common_layers.weights_nonzero
-
-    hp = model_hparams
-    if hp and hp.prepend_mode != "none":
-      assert (hp.prepend_mode == "prepend_inputs_masked_attention" or
-              hp.prepend_mode == "prepend_inputs_full_attention")
-
-      if (
-          # In masked attention mode, during training, the network try to
-          # autoregressively predicting the inputs portion, while the
-          # evaluation is only done on the output
-          hp.prepend_mode != "prepend_inputs_masked_attention" or
-          hp.mode != tf.estimator.ModeKeys.TRAIN):
-        weights_fn = common_layers.weights_prepend_inputs_to_targets
-
-    return weights_fn
+  targets_weights_fn = staticmethod(common_layers.weights_nonzero)
 
   @staticmethod
   def _get_weights(model_hparams, vocab_size, hidden_dim=None):
@@ -204,9 +187,7 @@ def top(cls, body_output, targets, model_hparams, vocab_size):
 class SymbolModalityWeightsAll(SymbolModality):
   """SymbolModality for features that do not have 0-padding."""
 
-  @staticmethod
-  def targets_weights_fn(model_hparams):
-    return common_layers.weights_all
+  targets_weights_fn = staticmethod(common_layers.weights_all)
 
 
 class SymbolModalityOneHot(SymbolModality):
@@ -226,7 +207,8 @@ def top(body_output, _, model_hparams, vocab_size):
     return body_output
 
   @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size):
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+    del weights_fn  # unused arg
     labels = tf.one_hot(targets, vocab_size)
     loss = tf.nn.softmax_cross_entropy_with_logits(
         logits=top_out, labels=labels)
@@ -237,7 +219,7 @@ class CTCSymbolModality(SymbolModality):
   """SymbolModality that uses CTC loss."""
 
   @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size):
+  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
     """Compute the CTC loss."""
     logits = top_out
     with tf.name_scope("ctc_loss", values=[logits, targets]):
@@ -259,7 +241,7 @@ def loss(cls, top_out, targets, model_hparams, vocab_size):
           time_major=False,
           preprocess_collapse_repeated=False,
           ctc_merge_repeated=False)
-      weights = cls.targets_weights_fn(targets)  # pylint: disable=not-callable
+      weights = weights_fn(targets)
       return tf.reduce_sum(xent), tf.reduce_sum(weights)
 
 
@@ -323,7 +305,7 @@ def top(body_output, _, model_hparams, vocab_size):
       return res
 
   @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size):
+  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
     """Compute loss numerator and denominator for one shard of output."""
     logits = top_out
     cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.0)
@@ -332,7 +314,7 @@ def loss(cls, top_out, targets, model_hparams, vocab_size):
         targets,
         model_hparams.label_smoothing,
         cutoff=cutoff,
-        weights_fn=cls.targets_weights_fn(model_hparams))
+        weights_fn=weights_fn)
 
 
 class ImageChannelCompressModality(modality.Modality):
@@ -697,8 +679,8 @@ def top(body_output, targets, model_hparams, vocab_size):
     common_video.gif_summary("results", x, max_outputs=1)
     return res
 
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size):
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
     """Compute loss numerator and denominator for one shard of output."""
     logits = top_out
     logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
@@ -709,7 +691,7 @@ def loss(cls, top_out, targets, model_hparams, vocab_size):
         targets,
         model_hparams.label_smoothing,
         cutoff=cutoff,
-        weights_fn=cls.targets_weights_fn(model_hparams))
+        weights_fn=weights_fn)
 
 
 class VideoModalityBitwise(VideoModality):
@@ -796,13 +778,12 @@ def internal_loss(logits, targets, model_hparams):
     return tf.nn.relu(tf.abs(logits - targets) - cutoff)
 
   @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size):
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
     """Compute loss numerator and denominator for one shard of output."""
     logits = top_out
     logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:-1])
     targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-    targets_weights_fn = VideoModalityL1.targets_weights_fn(model_hparams)
-    weights = targets_weights_fn(targets)
+    weights = weights_fn(targets)
     # Shift targets by 0.5 so later just casting to int gives the prediction.
     # So for int targets, say 0 and 7, we actually train to predict 0.5 and 7.5.
     # Later (in merics or infer) this is cast to int anyway. Also, we have no
@@ -853,7 +834,8 @@ def top(body_output, _, model_hparams, vocab_size):
     return tf.expand_dims(rgb_frames, axis=-1)
 
   @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size):
+  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
+    del weights_fn  # unused arg
     prediction, groundtruth = cls.convert_rgb_to_real(top_out, targets)
     loss = tf.losses.mean_squared_error(prediction, groundtruth)
     return loss, tf.constant(1.0)
@@ -863,7 +845,7 @@ class VideoModalityL1Raw(VideoModalityL2Raw):
   """Modality with L1 loss and raw input (sequences of frames)."""
 
   @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size):
+  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
     prediction, groundtruth = cls.convert_rgb_to_real(top_out, targets)
     loss = tf.losses.absolute_difference(prediction, groundtruth)
     return loss, tf.constant(1.0)
@@ -935,7 +917,7 @@ def top(body_output, targets, model_hparams, vocab_size):
     return body_output
 
   @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size):
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
     """Compute loss numerator and denominator for one shard of output."""
     # TODO(nikip): Try L2 loss
     logits = top_out
@@ -947,19 +929,16 @@ def loss(top_out, targets, model_hparams, vocab_size):
         targets,
         model_hparams.label_smoothing,
         cutoff=cutoff,
-        weights_fn=VideoModalityIdentity.targets_weights_fn(model_hparams))
+        weights_fn=weights_fn)
 
 
 class MultiLabelModality(ClassLabelModality):
   """Used for multi label task."""
 
-  @staticmethod
-  def targets_weights_fn(model_hparams):
-    """Target weight function for multi label, defaults to nonzero labels."""
-    return common_layers.weights_nonzero
+  targets_weights_fn = staticmethod(common_layers.weights_nonzero)
 
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size):
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
     """Average loss over the labels."""
     logits = top_out
     num_labels = tf.shape(targets)[1]
@@ -969,7 +948,7 @@ def loss(cls, top_out, targets, model_hparams, vocab_size):
         logits,
         targets,
         model_hparams.label_smoothing,
-        weights_fn=cls.targets_weights_fn(model_hparams),
+        weights_fn=weights_fn,
         reduce_sum=False,
     )
     xent = tf.squeeze(xent, [2, 3])
@@ -986,8 +965,8 @@ def loss(cls, top_out, targets, model_hparams, vocab_size):
 class OneHotClassLabelModality(ClassLabelModality):
   """Used for one-hot encoded class labels."""
 
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size):
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
     """Apply softmax cross-entropy between outputs and targets.
 
     Args:
@@ -995,14 +974,14 @@ def loss(cls, top_out, targets, model_hparams, vocab_size):
       targets: one-hot encoding Tensor with shape [batch, ?, ?, num_classes]
       model_hparams: tf.HParams, model hyperparmeters.
       vocab_size: int, vocabulary size.
+      weights_fn: Function mapping targets to weights.
 
     Returns:
       loss_scale (cross-entropy), loss_denom
     """
     loss_scale = tf.losses.softmax_cross_entropy(
         onehot_labels=targets, logits=top_out)
-    targets_weights_fn = cls.targets_weights_fn(model_hparams)
-    weights = targets_weights_fn(targets)
+    weights = weights_fn(targets)
     loss_denom = tf.reduce_sum(weights)
     return loss_scale, loss_denom
 
@@ -1027,7 +1006,8 @@ def targets_bottom(x, model_hparams, vocab_size):
     return tf.to_float(x)
 
   @staticmethod
-  def loss(body_output, targets, model_hparams, vocab_size):
+  def loss(body_output, targets, model_hparams, vocab_size, weights_fn):
+    del weights_fn  # unused
     loss = tf.squared_difference(body_output, tf.to_float(targets))
     return tf.reduce_mean(loss), tf.constant(1.0)
 
@@ -1052,22 +1032,21 @@ def top(body_output, _, model_hparams, vocab_size):
       return tf.layers.dense(body_output, vocab_size, name="top")
 
   @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size):
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
     raise NotImplementedError()
 
 
 class RealL2LossModality(RealModality):
   """Modality for real (i.e. float) vectors with L2 (Gaussian) loss."""
 
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size):
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
     predictions = top_out
     if (len(common_layers.shape_list(top_out)) != len(
         common_layers.shape_list(targets))):
       predictions = tf.squeeze(top_out, axis=[-1])
     with tf.name_scope("l2"):
-      targets_weights_fn = cls.targets_weights_fn(model_hparams)
-      weights = targets_weights_fn(targets)
+      weights = weights_fn(targets)
       l2 = tf.pow(predictions - targets, 2)
       return tf.reduce_sum(l2 * weights), tf.reduce_sum(weights)
 
@@ -1075,15 +1054,14 @@ def loss(cls, top_out, targets, model_hparams, vocab_size):
 class RealLogPoissonLossModality(RealModality):
   """Modality for real (i.e. float) vectors with log Poisson regression loss."""
 
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size):
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
     predictions = top_out
     if (len(common_layers.shape_list(top_out)) != len(
         common_layers.shape_list(targets))):
       predictions = tf.squeeze(top_out, axis=[-1])
     with tf.name_scope("log_possion"):
-      targets_weights_fn = cls.targets_weights_fn(model_hparams)
-      weights = targets_weights_fn(targets)
+      weights = weights_fn(targets)
       lp_loss = tf.nn.log_poisson_loss(targets, predictions)
       return tf.reduce_sum(lp_loss * weights), tf.reduce_sum(weights)
 
@@ -1116,15 +1094,13 @@ def name(model_hparams, vocab_size):
     return "sigmoid_class_symbol_modality_%d_%d" % (vocab_size,
                                                     model_hparams.hidden_size)
 
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size):
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
     # Expect inputs of size [batch-size, timesteps, 1, num-classes], where the
     # last dimension of num-classes represents logits for binary labels
     loss_scale = tf.losses.sigmoid_cross_entropy(
         multi_class_labels=targets, logits=top_out)
-    # Weigh all classes equally
-    targets_weights_fn = cls.targets_weights_fn(model_hparams)
-    weights = targets_weights_fn(targets)
+    weights = weights_fn(targets)
     loss_denom = tf.reduce_sum(weights)
     return loss_scale, loss_denom
 
@@ -1156,14 +1132,13 @@ def top(cls, body_output, _, model_hparams, vocab_size):
       x = tf.reduce_max(x, axis=1, keepdims=True)
       return tf.layers.dense(x, vocab_size)
 
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size):
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
     # Expect inputs of size [batch-size, 1, 1, num-classes], where the
     # last dimension of num-classes represents logits for binary labels
     loss_scale = tf.losses.sigmoid_cross_entropy(
         multi_class_labels=targets, logits=top_out)
-    # Weigh all classes equally
-    weights = cls.targets_weights_fn(model_hparams)(targets)
+    weights = weights_fn(targets)
     loss_denom = tf.reduce_sum(weights)
     return loss_scale, loss_denom
 
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index bc9f8d0c2..6cb6efc98 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -86,7 +86,8 @@ def testSymbolModalityTargets(self):
         sharded_logits,
         sharded_targets,
         model_hparams,
-        vocab_size)
+        vocab_size,
+        modalities.get_targets_weights_fn(modalities.ModalityType.SYMBOL))
     train_loss = (tf.add_n(sharded_loss_num) /
                   tf.maximum(1.0, tf.add_n(sharded_loss_den)))
     logits = tf.concat(sharded_logits, 0)
@@ -127,7 +128,8 @@ def testSymbolModalityTargetsFactored(self):
           sharded_logits,
           sharded_targets,
           model_hparams,
-          vocab_size)
+          vocab_size,
+          modalities.get_targets_weights_fn(modalities.ModalityType.SYMBOL))
       train_loss = (tf.add_n(sharded_loss_num) /
                     tf.maximum(1.0, tf.add_n(sharded_loss_den)))
       logits = tf.concat(sharded_logits, 0)
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index f3516b8b1..c3dfb53ea 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -354,7 +354,7 @@ def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
       modality = self.hparams.problem_hparams.modality["targets"]
       targets_weights_fn = self.hparams.targets_weights_fn.get(
           "targets",
-          modalities.get_targets_weights_fn(modality))(self.hparams)
+          modalities.get_targets_weights_fn(modality))
       numerator, denominator = common_layers.padded_cross_entropy(
           logits,
           targets,
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 151551713..d3852c648 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -616,7 +616,7 @@ def weights_fn_for_mp(problem_task_id):
     for target_name, modality in six.iteritems(tm):
       weights_fn = model_hparams.targets_weights_fn.get(
           "targets",
-          modalities.get_targets_weights_fn(modality))(model_hparams)
+          modalities.get_targets_weights_fn(modality))
       if hasattr(model_hparams.problem, "task_list"):
         ptid = problem_instance.task_id  # pylint: disable=cell-var-from-loop
         weights_fn = weights_fn_for_mp(ptid)
@@ -645,7 +645,7 @@ def create_eager_metrics_for_problem(problem, model_hparams):
   target_modality = problem_hparams.modality["targets"]
   weights_fn = model_hparams.targets_weights_fn.get(
       "targets",
-      modalities.get_targets_weights_fn(target_modality))(model_hparams)
+      modalities.get_targets_weights_fn(target_modality))
   return create_eager_metrics_internal(metric_fns, weights_fn=weights_fn)
 
 
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
index 0a8925950..09a83d050 100644
--- a/tensor2tensor/utils/modality.py
+++ b/tensor2tensor/utils/modality.py
@@ -59,25 +59,7 @@ def name(cls, model_hparams, vocab_size=None):
     del model_hparams, vocab_size  # unused arg
     return misc_utils.camelcase_to_snakecase(type(cls).__name__)
 
-  @staticmethod
-  def targets_weights_fn(model_hparams):
-    """The weights function to use for loss and eval metrics.
-
-    A weights function takes labels and returns a Tensor that assigns weights
-    (usually either 1. or 0.) to each one.
-
-    Common weights functions are:
-      * weights_all: 1. for all labels
-      * weights_nonzero: 1. for all non-zero labels (e.g. to deal with padding)
-
-    Args:
-      model_hparams: tf.HParams, model hyperparmeters.
-
-    Returns:
-      Callable: (targets) -> weights Tensor
-    """
-    del model_hparams  # unused arg
-    return common_layers.weights_all
+  targets_weights_fn = staticmethod(common_layers.weights_all)
 
   @staticmethod
   def bottom(x, model_hparams, vocab_size=None):
@@ -126,17 +108,10 @@ def top(body_output, targets, model_hparams, vocab_size=None):
     """
     raise NotImplementedError("Abstract Method")
 
-  @classmethod
-  def loss(cls,
-           top_out,
-           targets,
-           model_hparams,
-           vocab_size=None,
-           weights_fn=None):
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
     """Compute loss numerator and denominator for one shard of output."""
     del vocab_size  # unused arg
-    if weights_fn is None:
-      weights_fn = cls.targets_weights_fn(model_hparams)
     logits = top_out
     logits = common_attention.maybe_upcast(logits, hparams=model_hparams)
     return common_layers.padded_cross_entropy(
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 3f22eb7db..f3b5c8fec 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -232,6 +232,20 @@ def __init__(self,
       hparams.targets_weights_fn[feature_name] = modality.targets_weights_fn
       hparams.top[feature_name] = modality.top
 
+    if self._problem_hparams:
+      for feature_name, modality in six.iteritems(
+          self._problem_hparams.modality):
+        # If prepend mode, set targets_weights_fn to appropriately handle it.
+        if (modality in (modalities.ModalityType.SYMBOL,
+                         modalities.ModalityType.SYMBOL_ONE_HOT,
+                         modalities.ModalityType.CTC_SYMBOL,
+                         modalities.ModalityType.IDENTITY_SYMBOL)):
+          if (hparams.prepend_mode == "prepend_inputs_full_attention" or
+              (hparams.prepend_mode == "prepend_inputs_masked_attention" and
+               mode != tf.estimator.ModeKeys.TRAIN)):
+            weights_fn = common_layers.weights_prepend_inputs_to_targets
+            hparams.targets_weights_fn[feature_name] = weights_fn
+
     self._original_hparams = hparams
     self.set_mode(mode)
 
@@ -632,17 +646,17 @@ def _loss_single(self, logits, feature_name, feature, weights=None):
     if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
       vocab_size += (-vocab_size) % self._hparams.vocab_divisor
     loss = self._hparams.loss.get(feature_name, modalities.get_loss(modality))
+    targets_weights_fn = self._hparams.targets_weights_fn.get(
+        "targets", modalities.get_targets_weights_fn(modality))
     if weights is None:
-      loss_num, loss_den = loss(logits, feature, self._hparams, vocab_size)
+      loss_num, loss_den = loss(logits, feature, self._hparams, vocab_size,
+                                weights_fn=targets_weights_fn)
     else:
 
       def weights_fn(labels):
         """Per-token weights for loss."""
         # Use target_weights_fn() given by modality as well as explicitly given
         # weights.
-        targets_weights_fn = self._hparams.targets_weights_fn.get(
-            feature_name,
-            modalities.get_targets_weights_fn(modality))(self._hparams)
         modality_weights = targets_weights_fn(labels)
 
         # Broadcast 'weights' along minor dimensions (TF's default is major).
@@ -1803,7 +1817,7 @@ def create_tpu_eval_metrics_fn(problem, model_hparams):
   tm = _create_target_modality(problem.get_hparams(model_hparams).modality)
   if isinstance(tm, dict):
     for k, v in six.iteritems(tm):
-      weights_fn = v.targets_weights_fn(model_hparams)
+      weights_fn = v.targets_weights_fn
 
       def make_metric_fn(metric_fn):
         def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
@@ -1823,7 +1837,7 @@ def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
         name = "%s/metrics-%s/%s" % (k, problem.name, metric)
         metric_fns.append((name, make_metric_fn(metric_fn)))
   else:
-    weights_fn = tm.targets_weights_fn(model_hparams)
+    weights_fn = tm.targets_weights_fn
 
     def make_metric_fn(metric_fn):
       def wrapped_metric_fn(logits, labels, features):
@@ -2024,11 +2038,14 @@ def sampled_results():
                                 vocab_size)
         if "training" not in losses:
           loss = hparams.loss.get("targets", modalities.get_loss(modality))
+          weights_fn = hparams.targets_weights_fn.get(
+              "targets", modalities.get_targets_weights_fn(modality))
           sharded_loss_num, sharded_loss_den = dp(loss,
                                                   sharded_logits,
                                                   sharded_features["targets"],
                                                   hparams,
-                                                  vocab_size)
+                                                  vocab_size,
+                                                  weights_fn=weights_fn)
           training_loss = (tf.add_n(sharded_loss_num) /
                            tf.maximum(1.0, tf.add_n(sharded_loss_den)))
           training_loss *= problem_hparams.loss_multiplier

From 728078d67e9d79647177b83a9e1df656894661c9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 22 Feb 2019 09:32:34 -0800
Subject: [PATCH 1699/2720] Move EnvProblem configuation options from
 constructor to initialize. This gives us more flexibility.

PiperOrigin-RevId: 235204944
---
 tensor2tensor/envs/env_problem.py             | 26 ++++++-------------
 .../envs/tic_tac_toe_env_problem_test.py      |  2 +-
 tensor2tensor/rl/trainer_model_free.py        |  2 +-
 tensor2tensor/utils/registry.py               |  6 ++---
 tensor2tensor/utils/registry_test.py          |  2 +-
 5 files changed, 14 insertions(+), 24 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 13d8d7baa..1b2ffde24 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -125,7 +125,6 @@ class EnvProblem(Env, problem.Problem):
 
   def __init__(self,
                base_env_name=None,
-               base_env_kwargs=None,
                batch_size=None,
                reward_range=(-np.inf, np.inf)):
     """Initializes this class by creating the envs and managing trajectories.
@@ -133,8 +132,6 @@ def __init__(self,
     Args:
       base_env_name: (string) passed to `gym_utils.make_gym_env` to make the
         underlying environment.
-      base_env_kwargs: (dict) passed to `gym_utils.make_gym_env` to make the
-        underlying environment.
       batch_size: (int or None): How many envs to make in the non natively
         batched mode.
       reward_range: (tuple(number, number)) the first element is the minimum
@@ -149,13 +146,6 @@ def __init__(self,
     # the default implementation of `initialize_environments`.
     self._base_env_name = base_env_name
 
-    # Other arguments for initializing environments, will be used in
-    # `gym_utils.make_gym_env` in the default implementation of
-    # `initialize_environments`.
-    self._base_env_kwargs = base_env_kwargs
-    if not self._base_env_kwargs:
-      self._base_env_kwargs = {}
-
     # An env generates data when it is given actions by an agent which is either
     # a policy or a human -- this is supposed to be the `id` of the agent.
     #
@@ -238,8 +228,8 @@ def _verify_same_spaces(self):
         tf.logging.error("Env[%d] has action space [%s]", i, env.action_space)
       raise ValueError(err_str)
 
-  def initialize(self, batch_size=1):
-    self.initialize_environments(batch_size=batch_size)
+  def initialize(self, **kwargs):
+    self.initialize_environments(**kwargs)
 
     # Assert that *all* the above are now set, we should do this since
     # subclasses can override `initialize_environments`.
@@ -249,7 +239,8 @@ def initialize(self, batch_size=1):
     assert self._reward_range is not None
     assert self._trajectories is not None
 
-  def initialize_environments(self, batch_size=1):
+  def initialize_environments(self, batch_size=1, max_episode_steps=-1,
+                              max_and_skip_env=False):
     """Initializes the environments and trajectories.
 
     Subclasses can override this if they don't want a default implementation
@@ -258,21 +249,20 @@ def initialize_environments(self, batch_size=1):
 
     Args:
       batch_size: (int) Number of `self.base_env_name` envs to initialize.
+      max_episode_steps: (int) Passed on to `gym_utils.make_gym_env`.
+      max_and_skip_env: (boolean) Passed on to `gym_utils.make_gym_env`.
     """
 
     assert batch_size >= 1
     self._batch_size = batch_size
 
-    max_steps = self._base_env_kwargs.get("rl_env_max_episode_steps", -1)
-    maxskip_env = self._base_env_kwargs.get("maxskip_env", False)
-
     self._envs = []
     for _ in range(batch_size):
       self._envs.append(
           gym_utils.make_gym_env(
               self.base_env_name,
-              rl_env_max_episode_steps=max_steps,
-              maxskip_env=maxskip_env))
+              rl_env_max_episode_steps=max_episode_steps,
+              maxskip_env=max_and_skip_env))
 
     # If self.observation_space and self.action_space aren't None, then it means
     # that this is a re-initialization of this class, in that case make sure
diff --git a/tensor2tensor/envs/tic_tac_toe_env_problem_test.py b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
index bf69fe6f3..f9d4e6a03 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
@@ -31,7 +31,7 @@ class TicTacToeEnvProblemTest(tf.test.TestCase):
   def test_registration_and_interaction_with_env_problem(self):
     batch_size = 5
     # This ensures that registration has occurred.
-    ep = registry.env_problem("tic_tac_toe_env_problem", batch_size)
+    ep = registry.env_problem("tic_tac_toe_env_problem", batch_size=batch_size)
     ep.reset()
     num_done, num_lost, num_won, num_draw = 0, 0, 0, 0
     nsteps = 100
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 40fdb9c2a..5d7a0974a 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -68,7 +68,7 @@
 def initialize_env_specs(hparams, env_problem_name):
   """Initializes env_specs using the appropriate env."""
   if env_problem_name:
-    env = registry.env_problem(env_problem_name, hparams.batch_size)
+    env = registry.env_problem(env_problem_name, batch_size=hparams.batch_size)
   else:
     env = rl_utils.setup_env(hparams, hparams.batch_size,
                              hparams.eval_max_num_noops,
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index a89dd6739..cbde96852 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -508,12 +508,12 @@ def problem(problem_name):
       was_copy=spec.was_copy, was_reversed=spec.was_reversed)
 
 
-def env_problem(env_problem_name, batch_size):
+def env_problem(env_problem_name, **kwargs):
   """Get and initialize the `EnvProblem` with the given name and batch size.
 
   Args:
     env_problem_name: string name of the registered env problem.
-    batch_size: batch_size to initialize the env problem with.
+    **kwargs: forwarded to env problem's initialize method.
 
   Returns:
     an initialized EnvProblem with the given batch size.
@@ -521,7 +521,7 @@ def env_problem(env_problem_name, batch_size):
 
   ep_cls = Registries.env_problems[env_problem_name]
   ep = ep_cls()
-  ep.initialize(batch_size=batch_size)
+  ep.initialize(**kwargs)
   return ep
 
 
diff --git a/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index 55e1359fd..e50d1ce6a 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -108,7 +108,7 @@ def initialize(self, batch_size):
 
     # Get it with given batch_size.
     batch_size = 100
-    ep = registry.env_problem("env_prob", batch_size)
+    ep = registry.env_problem("env_prob", batch_size=batch_size)
 
     # name property is set.
     self.assertEqual("env_prob", ep.name)

From 7fdbcc1b2b6dee44a951eb781ddaf2086bd2d691 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 22 Feb 2019 10:03:16 -0800
Subject: [PATCH 1700/2720] OSS T2T currently breaks on import jax when
 ImportError is thrown. Switch TypeError to ImportError to catch.

PiperOrigin-RevId: 235210246
---
 tensor2tensor/bin/t2t_trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index bb0ed809a..f2285b11a 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -40,7 +40,7 @@
 
 try:
   from tensor2tensor.jax import j2j  # pylint: disable=g-import-not-at-top
-except TypeError:
+except (TypeError, ImportError):
   pass
 
 

From 83a3e21db2dd077da5521f85ce09611788607aea Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 22 Feb 2019 12:01:00 -0800
Subject: [PATCH 1701/2720] tensorboard summary exporter for jax

PiperOrigin-RevId: 235232910
---
 tensor2tensor/jax/jaxboard.py      | 309 +++++++++++++++++++++++++++++
 tensor2tensor/jax/jaxboard_demo.py | 122 ++++++++++++
 2 files changed, 431 insertions(+)
 create mode 100644 tensor2tensor/jax/jaxboard.py
 create mode 100644 tensor2tensor/jax/jaxboard_demo.py

diff --git a/tensor2tensor/jax/jaxboard.py b/tensor2tensor/jax/jaxboard.py
new file mode 100644
index 000000000..278d3abba
--- /dev/null
+++ b/tensor2tensor/jax/jaxboard.py
@@ -0,0 +1,309 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Write Summaries from JAX for use with Tensorboard.
+
+See jaxboard_demo.py for example usage.
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import io
+import struct
+import warnings
+import wave
+import matplotlib as mpl
+# Necessary to prevent attempted Tk import:
+with warnings.catch_warnings():
+  warnings.simplefilter('ignore')
+  mpl.use('Agg')
+# pylint: disable=g-import-not-at-top
+import matplotlib.pyplot as plt
+import numpy as onp
+import tensorflow as tf
+from tensorflow import gfile
+from tensorflow import HistogramProto
+from tensorflow import Summary
+from tensorflow import SummaryMetadata
+
+
+def _pack_images(images, rows, cols):
+  """Helper utility to make a tiled field of images from numpy arrays.
+
+  Args:
+    images: Image tensor in shape [N, W, H, C].
+    rows: Number of images per row in tiled image.
+    cols: Number of images per column in tiled image.
+
+  Returns:
+    A tiled image of shape [W * rows, H * cols, C].
+    Truncates incomplete rows.
+  """
+  shape = onp.shape(images)
+  width, height, depth = shape[-3:]
+  images = onp.reshape(images, (-1, width, height, depth))
+  batch = onp.shape(images)[0]
+  rows = onp.minimum(rows, batch)
+  cols = onp.minimum(batch // rows, cols)
+  images = images[:rows * cols]
+  images = onp.reshape(images, (rows, cols, width, height, depth))
+  images = onp.transpose(images, [0, 2, 1, 3, 4])
+  images = onp.reshape(images, [rows * width, cols * height, depth])
+  return images
+
+
+class SummaryWriter(object):
+  """Saves data in event and summary protos for tensorboard."""
+
+  def __init__(self, log_dir):
+    """Create a new SummaryWriter.
+
+    Args:
+      log_dir: path to record tfevents files in.
+    """
+    # If needed, create log_dir directory as well as missing parent directories.
+    if not gfile.IsDirectory(log_dir):
+      gfile.MakeDirs(log_dir)
+
+    self.writer = tf.summary.FileWriter(log_dir, graph=None)
+    self.end_summaries = []
+    self.step = 0
+    self.closed = False
+
+  def close(self):
+    """Close SummaryWriter. Final!"""
+    if not self.closed:
+      for summary in self.end_summaries:
+        self.writer.add_summary(summary, self.step)
+      self.writer.close()
+      self.closed = True
+      del self.writer
+
+  def __del__(self):  # safe?
+    self.close()
+
+  def scalar(self, tag, value, step=None):
+    """Saves scalar value.
+
+    Args:
+      tag: str: label for this data
+      value: int/float: number to log
+      step: int: training step
+    """
+    value = float(onp.array(value))
+    if step is None:
+      step = self.step
+    else:
+      self.step = step
+    summary = Summary(value=[Summary.Value(tag=tag, simple_value=value)])
+    self.writer.add_summary(summary, step)
+
+  def image(self, tag, image, step=None):
+    """Saves RGB image summary from onp.ndarray [H,W], [H,W,1], or [H,W,3].
+
+    Args:
+      tag: str: label for this data
+      image: ndarray: [H,W], [H,W,1], [H,W,3] save image in greyscale or colors/
+      step: int: training step
+    """
+    image = onp.array(image)
+    if step is None:
+      step = self.step
+    else:
+      self.step = step
+    if len(onp.shape(image)) == 2:
+      image = image[:, :, onp.newaxis]
+    if onp.shape(image)[-1] == 1:
+      image = onp.repeat(image, 3, axis=-1)
+    image_strio = io.BytesIO()
+    plt.imsave(image_strio, image, format='png')
+    image_summary = Summary.Image(
+        encoded_image_string=image_strio.getvalue(),
+        colorspace=3,
+        height=image.shape[0],
+        width=image.shape[1])
+    summary = Summary(value=[Summary.Value(tag=tag, image=image_summary)])
+    self.writer.add_summary(summary, step)
+
+  def images(self, tag, images, step=None, rows=None, cols=None):
+    """Saves (rows, cols) tiled images from onp.ndarray.
+
+    If either rows or cols aren't given, they are determined automatically
+    from the size of the image batch, if neither are given a long column
+    of images is produced. This truncates the image batch rather than padding
+    if it doesn't fill the final row.
+
+    Args:
+      tag: str: label for this data
+      images: ndarray: [N,H,W,1] or [N,H,W,3] to tile in 2d
+      step: int: training step
+      rows: int: number of rows in tile
+      cols: int: number of columns in tile
+    """
+    images = onp.array(images)
+    if step is None:
+      step = self.step
+    else:
+      self.step = step
+    n_images = onp.shape(images)[0]
+    if rows is None and cols is None:
+      rows = 1
+      cols = n_images
+    elif rows is None:
+      rows = n_images // cols
+    elif cols is None:
+      cols = n_images // rows
+    tiled_images = _pack_images(images, rows, cols)
+    self.image(tag, tiled_images, step=step)
+
+  def plot(self, tag, mpl_plt, step=None, close_plot=True):
+    """Saves matplotlib plot output to summary image.
+
+    Args:
+      tag: str: label for this data
+      mpl_plt: matplotlib stateful pyplot object with prepared plotting state
+      step: int: training step
+      close_plot: bool: automatically closes plot
+    """
+    if step is None:
+      step = self.step
+    else:
+      self.step = step
+    fig = mpl_plt.get_current_fig_manager()
+    img_w, img_h = fig.canvas.get_width_height()
+    image_buf = io.BytesIO()
+    mpl_plt.savefig(image_buf, format='png')
+    image_summary = Summary.Image(
+        encoded_image_string=image_buf.getvalue(),
+        colorspace=4,  # RGBA
+        height=img_h,
+        width=img_w)
+    summary = Summary(value=[Summary.Value(tag=tag, image=image_summary)])
+    self.writer.add_summary(summary, step)
+    if close_plot:
+      mpl_plt.close()
+
+  def audio(self, tag, audiodata, step=None, sample_rate=44100):
+    """Saves audio.
+
+    NB: single channel only right now.
+
+    Args:
+      tag: str: label for this data
+      audiodata: ndarray [Nsamples,]: data between (-1.0,1.0) to save as wave
+      step: int: training step
+      sample_rate: sample rate of passed in audio buffer
+    """
+    audiodata = onp.array(audiodata)
+    if step is None:
+      step = self.step
+    else:
+      self.step = step
+    audiodata = onp.clip(onp.squeeze(audiodata), -1, 1)
+    if audiodata.ndim != 1:
+      raise ValueError('Audio data must be 1D.')
+    sample_list = (32767.0 * audiodata).astype(int).tolist()
+    wio = io.BytesIO()
+    wav_buf = wave.open(wio, 'wb')
+    wav_buf.setnchannels(1)
+    wav_buf.setsampwidth(2)
+    wav_buf.setframerate(sample_rate)
+    enc = b''.join([struct.pack('<h', v) for v in sample_list])
+    wav_buf.writeframes(enc)
+    wav_buf.close()
+    encoded_audio_bytes = wio.getvalue()
+    wio.close()
+    audio = Summary.Audio(
+        sample_rate=sample_rate,
+        num_channels=1,
+        length_frames=len(sample_list),
+        encoded_audio_string=encoded_audio_bytes,
+        content_type='audio/wav')
+    summary = Summary(value=[Summary.Value(tag=tag, audio=audio)])
+    self.writer.add_summary(summary, step)
+
+  def histogram(self, tag, values, bins, step=None):
+    """Saves histogram of values.
+
+    Args:
+      tag: str: label for this data
+      values: ndarray: will be flattened by this routine
+      bins: number of bins in histogram, or array of bins for onp.histogram
+      step: int: training step
+    """
+    if step is None:
+      step = self.step
+    else:
+      self.step = step
+    values = onp.array(values)
+    bins = onp.array(bins)
+    values = onp.reshape(values, -1)
+    counts, limits = onp.histogram(values, bins=bins)
+    # boundary logic
+    cum_counts = onp.cumsum(onp.greater(counts, 0, dtype=onp.int32))
+    start, end = onp.searchsorted(
+        cum_counts, [0, cum_counts[-1] - 1], side='right')
+    start, end = int(start), int(end) + 1
+    counts = (
+        counts[start -
+               1:end] if start > 0 else onp.concatenate([[0], counts[:end]]))
+    limits = limits[start:end + 1]
+    sum_sq = values.dot(values)
+    histo = HistogramProto(
+        min=values.min(),
+        max=values.max(),
+        num=len(values),
+        sum=values.sum(),
+        sum_squares=sum_sq,
+        bucket_limit=limits.tolist(),
+        bucket=counts.tolist())
+    summary = Summary(value=[Summary.Value(tag=tag, histo=histo)])
+    self.writer.add_summary(summary, step)
+
+  def text(self, tag, textdata, step=None):
+    """Saves a text summary.
+
+    Args:
+      tag: str: label for this data
+      textdata: string, or 1D/2D list/numpy array of strings
+      step: int: training step
+    Note: markdown formatting is rendered by tensorboard.
+    """
+    if step is None:
+      step = self.step
+    else:
+      self.step = step
+    smd = SummaryMetadata(
+        plugin_data=SummaryMetadata.PluginData(plugin_name='text'))
+    if isinstance(textdata, (str, bytes)):
+      tensor = tf.make_tensor_proto(
+          values=[textdata.encode(encoding='utf_8')], shape=(1,))
+    else:
+      textdata = onp.array(textdata)  # convert lists, jax arrays, etc.
+      datashape = onp.shape(textdata)
+      if len(datashape) == 1:
+        tensor = tf.make_tensor_proto(
+            values=[td.encode(encoding='utf_8') for td in textdata],
+            shape=(datashape[0],))
+      elif len(datashape) == 2:
+        tensor = tf.make_tensor_proto(
+            values=[
+                td.encode(encoding='utf_8') for td in onp.reshape(textdata, -1)
+            ],
+            shape=(datashape[0], datashape[1]))
+    summary = Summary(
+        value=[Summary.Value(tag=tag, metadata=smd, tensor=tensor)])
+    self.writer.add_summary(summary, step)
diff --git a/tensor2tensor/jax/jaxboard_demo.py b/tensor2tensor/jax/jaxboard_demo.py
new file mode 100644
index 000000000..53b740415
--- /dev/null
+++ b/tensor2tensor/jax/jaxboard_demo.py
@@ -0,0 +1,122 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+r"""Jaxboard Summary Types Demo."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+from absl import app
+from absl import flags
+
+import warnings  # pylint: disable=g-bad-import-order
+import matplotlib as mpl
+with warnings.catch_warnings():
+  warnings.simplefilter('ignore')
+  mpl.use('Agg')
+# pylint: disable=g-import-not-at-top
+from matplotlib import pyplot as plt
+import numpy as onp
+
+from tensor2tensor.jax import jaxboard
+
+flags.DEFINE_string('tb_log_dir', '/tmp/tb_logs',
+                    'Path where we store summaries.')
+FLAGS = flags.FLAGS
+
+
+def demo():
+  """Run Summary Types Demo."""
+  sw = jaxboard.SummaryWriter(
+      os.path.join(FLAGS.tb_log_dir, 'demo', 'summarydemo'))
+
+  # Scalars.  We pass in step explicitly.
+  for i, v in enumerate(onp.sin(onp.linspace(0.0, 1.0, 100))):
+    sw.scalar('summarydemo_loss', v + 0.1 * onp.random.random(), step=i)
+
+  # SummaryWriter stores last step variable passed in, we can also set it
+  # explicitly for a set of exports to avoid providing the kwarg.
+  sw.step = 2
+
+  # Images. [H,W] or [H,W,C] with C = 1 or 3
+  sw.image('pic_c0', onp.random.random((100, 100)))
+  sw.image('pic_c1', onp.random.random((100, 100, 1)))
+  sw.image('pic_c3', onp.random.random((100, 100, 3)))
+
+  # Tiled sets of images. Must be [N,H,W,C] with C = 1 or 3
+  bw_tiles = onp.stack([
+      0.1 * onp.random.random((100, 100, 1)), 0.2 * onp.random.random(
+          (100, 100, 1)), 0.4 * onp.random.random((100, 100, 1)),
+      0.8 * onp.random.random((100, 100, 1))
+  ])
+  sw.images('pics_tiled_c1', bw_tiles, rows=2, cols=2)
+  clr_tiles = onp.stack([
+      0.1 * onp.random.random((100, 100, 3)), 0.2 * onp.random.random(
+          (100, 100, 3)), 0.4 * onp.random.random((100, 100, 3)),
+      0.8 * onp.random.random((100, 100, 3))
+  ])
+  sw.images('pics_tiled_c3', clr_tiles, rows=2, cols=2)
+
+  # Matplotlib plots. Just pass in prepared stateful pyplot object.
+  # -- scatter
+  plt.figure(figsize=(4, 4))
+  plt.scatter(
+      onp.random.randint(size=(10,), low=0, high=10),
+      onp.random.randint(size=(10,), low=0, high=10))
+  sw.plot('plot1', plt)
+
+  # -- imshow
+  plt.figure(figsize=(4, 4))
+  plt.imshow(
+      onp.random.randint(size=(50, 50, 3), low=0, high=255),
+      cmap='viridis',
+      interpolation='nearest')
+  sw.plot('plot2', plt)
+
+  # Audio.
+  t = onp.linspace(0, 1.0, 44100)
+  sinwave = (
+      0.1 * onp.sin(440. * onp.pi * t) *
+      # slow ramp-up to prevent 'pop'
+      onp.where(t > 0.2, 1.0, t / 0.2))
+  sw.audio('audio', sinwave)
+
+  # Text.
+  # -- tensorboard text plugin supports some markdown formatting!
+  sw.text('text', 'Colorless _green_ __ideas__ sleep furiously.')
+
+  # -- 1d/2d arrays of strings rendered as tables by plugin:
+  sw.text('text1d', ['Colorless', 'green', 'ideas', 'sleep', 'furiously.'])
+  sw.text('text2d', onp.array([['foo', 'bar'], ['baz', 'qup']]))
+
+  # Histograms / Distributions.
+  # (bins can be int or array - passed into onp.histogram bins arg)
+  sw.histogram('histo', onp.random.normal(size=(1000,)), 25, step=3)
+  sw.histogram('histo', onp.random.normal(size=(1000,)), 25, step=4)
+  sw.histogram('histo', onp.random.normal(size=(1000,)), 25, step=5)
+
+  # Fin.
+  sw.close()
+
+
+def main(argv):
+  del argv
+  demo()
+
+
+if __name__ == '__main__':
+  app.run(main)

From 99c1dbbc44f0a46b7e2fb5474f25197f77d6cf5e Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 22 Feb 2019 21:12:41 +0100
Subject: [PATCH 1702/2720] Readme (#1462)

* Fix VideoModalityPixelNoise infinite recursion

* Fix cumulative reward calculation for batches of rollouts

* New README for the RL module

* Describe the player controls in README
---
 tensor2tensor/layers/modalities.py |   3 +-
 tensor2tensor/rl/README.md         | 225 ++++++++++++++++++++++-------
 tensor2tensor/rl/rl_utils.py       |  20 ++-
 3 files changed, 187 insertions(+), 61 deletions(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index cbcd62c7c..fe190651e 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -749,7 +749,8 @@ def bottom(cls, x, model_hparams, vocab_size):
       input_mask = tf.reshape(tf.cast(input_mask, tf.int32),
                               input_shape[:-1]+[1])
       inputs = inputs * input_mask + background * (1 - input_mask)
-    return cls.bottom(inputs, model_hparams, vocab_size)
+    return super(VideoModalityPixelNoise, cls).bottom(
+        inputs, model_hparams, vocab_size)
 
 
 class VideoModalityL1(VideoModality):
diff --git a/tensor2tensor/rl/README.md b/tensor2tensor/rl/README.md
index 195222a60..c1fbacb19 100644
--- a/tensor2tensor/rl/README.md
+++ b/tensor2tensor/rl/README.md
@@ -1,6 +1,4 @@
-# Tensor2Tensor experimental Model-Based Reinforcement Learning.
-
-**Note**: Experimental and under development.
+# Tensor2Tensor Model-Based Reinforcement Learning.
 
 The `rl` package provides the ability to run model-based reinforcement learning
 algorithms using models trained with Tensor2Tensor.
@@ -11,87 +9,210 @@ Proximal Policy Optimization (PPO). See `trainer_model_based.py`.
 As a baseline, you can also run PPO without the model using
 `trainer_model_free.py`.
 
-## Model-based training
+To use this package, you need to install the Atari dependencies for OpenAI Gym:
+
+```
+pip install gym[atari]
+```
+
+## Evaluating a pretrained policy
+
+We provide a set of pretrained policies and models you can use. To evaluate and
+generate videos for a pretrained policy on Pong:
+
+```
+OUTPUT_DIR=~/t2t_train/pong_pretrained
+python -m tensor2tensor.rl.evaluator \
+  --loop_hparams_set=rlmb_long_stochastic_discrete \
+  --loop_hparams=game=pong \
+  --policy_dir=gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/policy \
+  --eval_metrics_dir=$OUTPUT_DIR \
+  --debug_video_path=$OUTPUT_DIR \
+  --num_debug_videos=4
+```
+
+By default, it will run a grid of different evaluation settings (sampling
+temperatures and whether to do initial rollouts). You can override those
+settings:
+
+```
+  --loop_hparams=game=pong,eval_max_num_noops=0,eval_sampling_temps=[0.0]
+```
+
+TensorBoard metrics are exported to the `eval_metrics_dir`. To view them, run:
 
-Alternate training a world model and a PPO agent within that model using the
-base hyperparameters on Freeway:
+```
+tensorboard --logdir=~/t2t_train/pong_pretrained
+```
+
+Description of player controls and flags can be found in `tensor2tensor/rl/player.py`.
+
+## Model-based training with pretrained world models
+
+To train a policy with a pretrained world model (requires Google Cloud SDK):
 
 ```
+OUTPUT_DIR=~/t2t_train/mb_sd_pong_pretrained
+mkdir -p $OUTPUT_DIR
+gsutil -m cp -r \
+  gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/world_model \
+  $OUTPUT_DIR/
 python -m tensor2tensor.rl.trainer_model_based \
-  --output_dir=$OUT_DIR \
-  --loop_hparams_set=rl_modelrl_base \
-  --loop_hparams='game=freeway'
+  --loop_hparams_set=rlmb_long_stochastic_discrete \
+  --loop_hparams=game=pong,epochs=1,model_train_steps=0 \
+  --eval_world_model=False \
+  --output_dir=$OUTPUT_DIR
 ```
 
-All hyperparameter sets are defined in `trainer_model_based.py` and are derived
-from `rl_modelrl_base`.
+Note that this command will collect some frames from the real environment for
+random starts.
+
+The same command can be used to resume interrupted training - checkpoints are
+saved in `output_dir`.
+
+We use `NoFrameskip-v4` game mode with our own frame skip (4 by default).
 
-The hyperparameters for the environment model and agent are nested within the
-`loop_hparams` by name. For example:
+The training script runs periodic evaluation, but with timestep limit 1000 to
+make it faster. To do full evaluation after training, run:
 
 ```
-  generative_model="next_frame_basic",
-  generative_model_params="next_frame_pixel_noise",
-  ppo_params="ppo_original_params",
+python -m tensor2tensor.rl.evaluator \
+  --loop_hparams_set=rlmb_long_stochastic_discrete \
+  --hparams=game=pong \
+  --policy_dir=$OUTPUT_DIR \
+  --eval_metrics_dir=$OUTPUT_DIR/full_eval_metrics
 ```
 
-## Model-free training
+## Model training with random trajectories
 
-**TODO(piotrmilos): Update**
+The simplest way to train your own model is to use random trajectories. Then you
+can train a policy on it as described in the previous section.
 
-Training an agent in `Pendulum-v0`:
+To train a deterministic model:
 
 ```
-python -m tensor2tensor.rl.trainer_model_free \
-  --problem=Pendulum-v0 \
-  --hparams_set ppo_continuous_action_base \
-  --output_dir $OUT_DIR
+python -m tensor2tensor.rl.trainer_model_based \
+  --loop_hparams_set=rlmb_base \
+  --loop_hparams=game=pong,epochs=1,ppo_epochs_num=0 \
+  --output_dir=~/t2t_train/mb_det_pong_random
 ```
 
-Training an agent in `PongNoFrameskip-v0`:
+To train a stochastic discrete model (it will require more time and memory):
 
 ```
-python -m tensor2tensor.rl.trainer_model_free \
-  --problem stacked_pong \
-  --hparams_set ppo_original_params \
-  --hparams num_agents=5 \
-  --output_dir dir_location
+python -m tensor2tensor.rl.trainer_model_based \
+  --loop_hparams_set=rlmb_long_stochastic_discrete \
+  --loop_hparams=game=pong,epochs=1,ppo_epochs_num=0 \
+  --output_dir=~/t2t_train/mb_sd_pong_random
+```
+
+## Full model-based training
+
+Our full training pipeline involves alternating between collecting data using
+policy, training the world model and training the policy inside the model. It
+requires significantly more time (several days to a week, depending on your
+hardware and the model you use).
+
+To train a deterministic model:
+
+```
+python -m tensor2tensor.rl.trainer_model_based \
+  --loop_hparams_set=rlmb_base \
+  --loop_hparams=game=pong \
+  --output_dir ~/t2t_train/mb_det_pong
+```
+
+To train a stochastic discrete model:
+
+```
+python -m tensor2tensor.rl.trainer_model_based \
+  --loop_hparams_set=rlmb_long_stochastic_discrete \
+  --loop_hparams=game=pong \
+  --output_dir ~/t2t_train/mb_sd_pong
+```
+
+Hyperparameter sets are defined in
+`tensor2tensor/rl/trainer_model_based_params.py`. Hyperparameter sets for the
+world model and agent are nested within `loop_hparams` by name. You can change
+them with:
+
+```
+  --loop_hparams=game=freeway,generative_model=next_frame_basic_deterministic,base_algo_params=ppo_original_params
 ```
 
-## Model training on random trajectories
+Game names should be provided in `snake_case`.
 
-Generate trajectories with a random policy:
+## Playing in the model
 
+To assess world model quality you can play in it, as in an Atari emulator
+(you need a machine with GPU for this). First install `pygame`:
+
+```
+pip install pygame
 ```
-python -m tensor2tensor.rl.datagen_with_agent \
-  --data_dir=$HOME/t2t/data \
-  --tmp_dir=$HOME/t2t/tmp \
-  --game=pong \
-  --num_env_steps=30000
+
+Then you can run the player, specifying a path to world model checkpoints:
+
 ```
+OUTPUT_DIR=~/t2t_train/mb_sd_pong_pretrained
+mkdir -p $OUTPUT_DIR
+gsutil -m cp -r \
+  gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/world_model \
+  $OUTPUT_DIR/
+python -m tensor2tensor.rl.player \
+  --wm_dir=$OUTPUT_DIR/world_model \
+  --loop_hparams_set=rlmb_long_stochastic_discrete \
+  --loop_hparams=game=pong \
+  --game_from_filenames=False \
+  --zoom=3 \
+  --fps=5
+```
+
+The screen is split into 3 columns: frame from the world model, corresponding
+frame from the real environment and the difference between the two. Use WSAD
+and space to control the agent. The model will likely diverge quickly, press X
+to reset it using the current state of the real environment. Note that frames
+fed to the model were likely never seen by it during training, so the model's
+performance will be worse than during the policy training.
+
+For more details on controls and flags see `tensor2tensor/rl/player.py`.
 
-Train model on trajectories:
+## Model-free training
+
+Training model-free on Pong:
 
 ```
-python -m tensor2tensor.bin.t2t_trainer \
-  --data_dir=$HOME/t2t/data \
-  --output_dir=$HOME/t2t/train/pong_model \
-  --problem=gym_pong_random \
-  --model=next_frame_basic \
-  --hparams_set=next_frame
+python -m tensor2tensor.rl.trainer_model_free \
+  --hparams_set=rlmf_base \
+  --hparams=game=pong \
+  --output_dir=~/t2t_train/mf_pong
 ```
 
+Hyperparameter sets are defined in `tensor2tensor/models/research/rl.py`. You
+can override them using the `hparams` flag, e.g.
 
-## Collect trajectories using a trained agent
+```
+  --hparams=game=kung_fu_master,frame_stack_size=5
+```
+
+As in model-based training, the periodic evaluation runs with timestep limit
+of 1000. To do full evaluation after training, run:
 
 ```
-python -m tensor2tensor.rl.datagen_with_agent \
-  --data_dir=$HOME/t2t/data \
-  --tmp_dir=$HOME/t2t/tmp \
-  --game=pong \
-  --num_env_steps=30000 \
-  --agent_policy_path=$AGENT_CKPT_PATH
+OUTPUT_DIR=~/t2t_train/mf_pong
+python -m tensor2tensor.rl.evaluator \
+  --loop_hparams_set=rlmf_base \
+  --hparams=game=pong \
+  --policy_dir=$OUTPUT_DIR \
+  --eval_metrics_dir=$OUTPUT_DIR/full_eval_metrics
 ```
 
-Add `--eval` if you want to evaluate the agent against the environment instead
-of generating trajectories for training the world model.
+## Using checkpoints for other games
+
+We provide pretrained policies and stochastic discrete models for most of the
+Atari games in OpenAI Gym. They are available in Google Cloud Storage at
+`gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/N`, where `N` is
+a run number in range 1 - 180. Games with checkpoints are defined in
+`tensor2tensor.data_generators.gym_env.ATARI_GAMES_WITH_HUMAN_SCORE_NICE` and
+are numbered according to this order, with 5 runs per game. For example, runs
+for Amidar have numbers 6 - 10.
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index d6630fd6a..9eaac4f48 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -434,10 +434,10 @@ def run_rollouts(
   )
 
   num_dones = 0
-  first_dones = [False] * env.batch_size
+  first_dones = np.array([False] * env.batch_size)
   observations = initial_observations
   step_index = 0
-  cum_rewards = 0
+  cum_rewards = np.zeros(env.batch_size)
 
   for (video_writer, obs_stack) in zip(video_writers, initial_observations):
     for (i, ob) in enumerate(obs_stack):
@@ -448,10 +448,10 @@ def run_rollouts(
       video_writer.write(debug_frame)
 
   def proceed():
-    if step_limit is not None:
-      return step_index < step_limit
+    if step_index < step_limit:
+      return num_dones < env.batch_size or many_rollouts_from_each_env
     else:
-      return num_dones < env.batch_size
+      return False
 
   while proceed():
     act_kwargs = {}
@@ -474,12 +474,16 @@ def proceed():
       for (i, observation) in zip(now_done_indices, reset_observations):
         observations[i] = observation
     observations = np.array(observations)
-    cum_rewards = cum_rewards * discount_factor + rewards
+    cum_rewards[~first_dones] = (
+        cum_rewards[~first_dones] * discount_factor + rewards[~first_dones]
+    )
     step_index += 1
 
-    for (video_writer, obs_stack, reward, cum_reward) in zip(
-        video_writers, observations, rewards, cum_rewards
+    for (video_writer, obs_stack, reward, cum_reward, done) in zip(
+        video_writers, observations, rewards, cum_rewards, first_dones
     ):
+      if done:
+        continue
       ob = obs_stack[-1]
       debug_frame = augment_observation(
           ob, reward=reward, cum_reward=cum_reward,

From 065f470b0f9987b99ed07d0ce7e2d82488cbb708 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 22 Feb 2019 12:13:57 -0800
Subject: [PATCH 1703/2720] Correct gin imports.

PiperOrigin-RevId: 235235419
---
 tensor2tensor/v2/t2t.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/v2/t2t.py b/tensor2tensor/v2/t2t.py
index ce2e906ee..5921187b3 100644
--- a/tensor2tensor/v2/t2t.py
+++ b/tensor2tensor/v2/t2t.py
@@ -22,6 +22,8 @@
 import collections
 import os
 
+import gin
+
 from tensor2tensor import problems
 from tensor2tensor.utils import data_reader
 from tensor2tensor.v2.models import basic
@@ -31,8 +33,6 @@
 import tensorflow as tf
 import tensorflow_datasets as tfds
 
-import gin.tf
-
 
 # Since there are few models and configs for now, we use this simple registry.
 # TODO(lukaszkaiser): find a better way to do this or remove altogether.

From 41d17dcfcd1dca76b58f5e87a082ec8a058a224f Mon Sep 17 00:00:00 2001
From: qixiuai <qixiuai@163.com>
Date: Sat, 23 Feb 2019 04:15:35 +0800
Subject: [PATCH 1704/2720] Fix typo (#1459)

Fix typo: replace tf.Daataset to tf.data.Dataset
---
 tensor2tensor/v2/t2t.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/v2/t2t.py b/tensor2tensor/v2/t2t.py
index 5921187b3..ffa4db73e 100644
--- a/tensor2tensor/v2/t2t.py
+++ b/tensor2tensor/v2/t2t.py
@@ -56,8 +56,8 @@ def train_and_eval_dataset(dataset_name, data_dir):
 
   Returns:
     a 4-tuple consisting of:
-     * the train tf.Daataset
-     * the eval tf.Daataset
+     * the train tf.data.Dataset
+     * the eval tf.data.Dataset
      * information about features: a python dictionary with feature names
          as keys and an object as value that provides .shape and .num_classes.
      * supervised_keys: information what's the input and what's the target,

From d84c4e96605d30e9bdda4f3634ae6880d7966cf9 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 22 Feb 2019 12:39:09 -0800
Subject: [PATCH 1705/2720] Remove symbol_modality_skip_top in common_hparams.

No models seem to use this anymore, so I dropped it.

PiperOrigin-RevId: 235239623
---
 tensor2tensor/layers/common_hparams.py | 2 --
 tensor2tensor/layers/modalities.py     | 3 ---
 2 files changed, 5 deletions(-)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index d67bdea9c..654ee4829 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -186,8 +186,6 @@ def basic_params1():
       # by using a problem_hparams that uses the same modality object for
       # the input modality and target modality.
       shared_embedding=False,
-      # In SymbolModality, skip the top layer, assume we're providing logits.
-      symbol_modality_skip_top=False,
       # Modalities used to map from features to a space compatible with
       # chosen model architecture. It comprises key-value pairs of a feature
       # name (str) and its modality type.
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index fe190651e..9a87aee1b 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -160,9 +160,6 @@ def top(cls, body_output, targets, model_hparams, vocab_size):
       logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].
     """
     del targets  # unused arg
-    if model_hparams.symbol_modality_skip_top:
-      return tf.expand_dims(body_output, 3)
-
     if model_hparams.shared_embedding_and_softmax_weights:
       scope_name = "shared"
       reuse = tf.AUTO_REUSE

From 3efdec0a241490fcfd0af6dbea46a670da08340f Mon Sep 17 00:00:00 2001
From: "Xiaoming (Jason) Cui" <xiaoming.cui@intel.com>
Date: Fri, 22 Feb 2019 13:20:06 -0800
Subject: [PATCH 1706/2720] =?UTF-8?q?Added=20log=20information=20of=20perf?=
 =?UTF-8?q?ormance=20of=20inference=20of=20decoding,=20if=20bat=E2=80=A6?=
 =?UTF-8?q?=20(#1447)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Added log information of performance of inference of decoding, if batch_size=1 log average latency performance, if batch_size>1, log throughput performance

* Changed the format of the python file so that it follow google python coding convension
---
 tensor2tensor/utils/decoding.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 2a171f63b..b928d482a 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -417,7 +417,8 @@ def decode_from_file(estimator,
     sorted_inputs = _get_language_modeling_inputs(
         filename, decode_hp.delimiter, repeat=decode_hp.num_decodes)
     sorted_keys = range(len(sorted_inputs))
-  num_decode_batches = (len(sorted_inputs) - 1) // decode_hp.batch_size + 1
+  num_sentences = len(sorted_inputs)
+  num_decode_batches = (num_sentences - 1) // decode_hp.batch_size + 1
 
   if estimator.config.use_tpu:
     length = getattr(hparams, "length", 0) or hparams.max_length
@@ -510,11 +511,20 @@ def timer(gen):
       decodes.append(decoded_outputs)
     total_time_per_step += elapsed_time
     total_cnt += result["outputs"].shape[-1]
-  tf.logging.info("Elapsed Time: %5.5f" % (time.time() - start_time))
+  duration = time.time() - start_time
+  tf.logging.info("Elapsed Time: %5.5f" % duration)
   tf.logging.info("Averaged Single Token Generation Time: %5.7f "
                   "(time %5.7f count %d)" %
                   (total_time_per_step / total_cnt,
                    total_time_per_step, total_cnt))
+  if decode_hp.batch_size is 1:
+    tf.logging.info("Inference time %.4f seconds "
+                    "(Latency = %.4f ms/setences)" %
+                    (duration, 1000.0*duration/num_sentences))
+  else:
+    tf.logging.info("Inference time %.4f seconds "
+                    "(Throughput = %.4f sentences/second)" %
+                    (duration, num_sentences/duration))
 
   # If decode_to_file was provided use it as the output filename without change
   # (except for adding shard_id if using more shards for decoding).

From d2ee6989373ee58b8bae37706cfba89939c850a6 Mon Sep 17 00:00:00 2001
From: "Xiaoming (Jason) Cui" <xiaoming.cui@intel.com>
Date: Fri, 22 Feb 2019 13:20:25 -0800
Subject: [PATCH 1707/2720] Merge of PR #1447

PiperOrigin-RevId: 235246544
---
 tensor2tensor/utils/decoding.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index b928d482a..df61b5f08 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -517,7 +517,7 @@ def timer(gen):
                   "(time %5.7f count %d)" %
                   (total_time_per_step / total_cnt,
                    total_time_per_step, total_cnt))
-  if decode_hp.batch_size is 1:
+  if decode_hp.batch_size == 1:
     tf.logging.info("Inference time %.4f seconds "
                     "(Latency = %.4f ms/setences)" %
                     (duration, 1000.0*duration/num_sentences))

From e45e03cc6c7f308b9aefbf592cfc26c555411547 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 22 Feb 2019 14:07:10 -0800
Subject: [PATCH 1708/2720] enable video summaries for non RGB videos.

PiperOrigin-RevId: 235255355
---
 tensor2tensor/data_generators/video_utils.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 0072b68d0..4dc5560ff 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -93,6 +93,9 @@ def create_border(video, color="blue", border_percent=2):
   Returns:
     video: 5-D NumPy array.
   """
+  # Do not create border if the video is not in RGB format
+  if video.shape[-1] != 3:
+    return video
   color_to_axis = {"blue": 2, "red": 0, "green": 1}
   axis = color_to_axis[color]
   _, _, height, width, _ = video.shape

From 6b366ce6c4a4f0180a112698b241580b9b84a19c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 22 Feb 2019 17:22:07 -0800
Subject: [PATCH 1709/2720] Add monochrome support for image summary
 generation.

PiperOrigin-RevId: 235288012
---
 tensor2tensor/data_generators/image_utils.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 8313f0560..61766a89e 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -51,6 +51,9 @@ def image_to_tf_summary_value(image, tag):
   """
   curr_image = np.asarray(image, dtype=np.uint8)
   height, width, n_channels = curr_image.shape
+  # If monochrome image, then reshape to [height, width]
+  if n_channels == 1:
+    curr_image = np.reshape(curr_image, [height, width])
   s = io.BytesIO()
   matplotlib_pyplot().imsave(s, curr_image, format="png")
   img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(),

From 9df56f233794129e043569220e1948ec1118ee5e Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 22 Feb 2019 18:29:11 -0800
Subject: [PATCH 1710/2720] Add Resnet50 and structure models and configs.

What works:
jax/j2j_trainer --config_file jax/configs/resnet50_imagenet_8gb.gin

PiperOrigin-RevId: 235295575
---
 tensor2tensor/bin/t2t_trainer.py              |  7 +-
 .../jax/configs/resnet50_imagenet_8gb.gin     | 39 ++++++++
 tensor2tensor/jax/j2j.py                      | 71 +++++++++++---
 tensor2tensor/jax/j2j_trainer.py              | 46 +++------
 tensor2tensor/jax/models/__init__.py          | 24 +++++
 .../jax/{models.py => models/mlp.py}          |  2 +-
 tensor2tensor/jax/models/resnet.py            | 97 +++++++++++++++++++
 7 files changed, 236 insertions(+), 50 deletions(-)
 create mode 100644 tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin
 create mode 100644 tensor2tensor/jax/models/__init__.py
 rename tensor2tensor/jax/{models.py => models/mlp.py} (97%)
 create mode 100644 tensor2tensor/jax/models/resnet.py

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index f2285b11a..f63c22ecb 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -381,8 +381,11 @@ def main(argv):
     output_dir = os.path.expanduser(FLAGS.output_dir)
 
     gin.bind_parameter("train_fn.dataset", FLAGS.problem)
-    config_strs += ["train_fn.model=@models." + FLAGS.model]
-    gin.parse_config_files_and_bindings(FLAGS.hparams_set, config_strs)
+    config_strs += ["train_fn.model=@" + FLAGS.model]
+    config_files = []
+    if FLAGS.hparams_set:
+      config_files = [os.path.expanduser(FLAGS.hparams_set)]
+    gin.parse_config_files_and_bindings(config_files, config_strs)
     j2j.train_fn(data_dir, output_dir=output_dir)
     return
 
diff --git a/tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin
new file mode 100644
index 000000000..c75bef109
--- /dev/null
+++ b/tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin
@@ -0,0 +1,39 @@
+# Parameters for batch_fn:
+# ==============================================================================
+batch_fn.batch_size = 16
+batch_fn.bucket_length = 32
+batch_fn.buckets = None
+batch_fn.eval_batch_size = 16
+
+# Parameters for learning_rate:
+# ==============================================================================
+learning_rate.constant = 0.1
+learning_rate.schedule = 'constant * linear_warmup * rsqrt_decay'
+learning_rate.warmup_steps = 10000
+
+# Parameters for optimizer:
+# ==============================================================================
+optimizer.adam_b1 = 0.9
+optimizer.adam_b2 = 0.997
+optimizer.adam_eps = 1e-08
+optimizer.momentum_mass = 0.9
+optimizer.name = 'momentum'
+optimizer.rmsprop_eps = 1e-08
+optimizer.rmsprop_gamma = 0.9
+
+# Parameters for preprocess_fn:
+# ==============================================================================
+preprocess_fn.max_target_length = -1
+
+# Parameters for Resnet50:
+# ==============================================================================
+Resnet50.hidden_size = 64
+Resnet50.num_output_classes = 1001
+
+# Parameters for train_fn:
+# ==============================================================================
+train_fn.dataset = 't2t_image_imagenet'
+train_fn.eval_frequency = 100
+train_fn.eval_steps = 2
+train_fn.model = @Resnet50
+train_fn.train_steps = 100000
diff --git a/tensor2tensor/jax/j2j.py b/tensor2tensor/jax/j2j.py
index fcab6fc91..e3351bc72 100644
--- a/tensor2tensor/jax/j2j.py
+++ b/tensor2tensor/jax/j2j.py
@@ -31,18 +31,21 @@
 import jax.numpy as np
 
 from tensor2tensor.jax import input_pipeline
+from tensor2tensor.jax import jaxboard
 # Import for gin configurable models
 from tensor2tensor.jax import models  # pylint: disable=unused-import
 
-from tensorflow import gfile
+from tensorflow.io import gfile
 
 import tensorflow_datasets as tfds
 
 
 @gin.configurable(blacklist=["step"])
-def learning_rate(step, schedule=None, constant=0.001, warmup_steps=200):
+def learning_rate(step,
+                  schedule="constant * linear_warmup * rsqrt_decay",
+                  constant=0.001,
+                  warmup_steps=100):
   """Learning rate."""
-  schedule = schedule or "constant * linear_warmup * rsqrt_decay"
   ret = 1.0
   for name in [n.strip() for n in schedule.split("*")]:
     if name == "constant":
@@ -105,14 +108,19 @@ def log(s, stdout=True):
     print(s)
 
 
+def _make_directory(path):
+  """Helper function: create directory if it doesn't exist yet."""
+  if not gfile.exists(path):
+    log("Creating directory %s" % path)
+    gfile.mkdir(path)
+
+
 def save_params_and_step(params, step, output_dir):
   """Save params and step in output dir."""
   if output_dir is not None:
-    if not gfile.Exists(output_dir):
-      log("Creating directory %s" % output_dir)
-      gfile.MkDir(output_dir)
+    _make_directory(output_dir)
     params_file = os.path.join(output_dir, "model.pkl")
-    with gfile.Open(params_file, "wb") as f:
+    with gfile.GFile(params_file, "wb") as f:
       pickle.dump((params, step), f)
     log("Model saved to %s" % params_file, stdout=False)
 
@@ -121,17 +129,22 @@ def load_params_and_step(output_dir):
   """Save params and step in output dir."""
   if output_dir is None:
     return None, None
-  if not gfile.Exists(output_dir):
+  if not gfile.exists(output_dir):
     return None, None
   params_file = os.path.join(output_dir, "model.pkl")
-  if not gfile.Exists(params_file):
+  if not gfile.exists(params_file):
     return None, None
-  with gfile.Open(params_file, "r") as f:
+  with gfile.GFile(params_file, "r") as f:
     (params, step) = pickle.load(f)
   log("Model loaded from %s" % params_file)
   return params, step
 
 
+def _make_summary_writer(output_path):
+  _make_directory(output_path)
+  return jaxboard.SummaryWriter(output_path)
+
+
 # We include in gin config everything that could be useful to share between
 # users, so when it gets saved in a .gin file it can be re-ran with few flags.
 @gin.configurable(blacklist=["data_dir", "output_dir"])
@@ -166,16 +179,23 @@ def update(i, opt_state, batch):
         params, batch, model_predict), opt_state)
 
   _, init_params = model_init([-1] + input_shape)
-  step = 0
+  step, train_sw, eval_sw = 0, None, None
   if output_dir is not None:
+    _make_directory(output_dir)
+    # Load parameters.
     loaded_params, loaded_step = load_params_and_step(output_dir)
     if loaded_params is not None:
       init_params = loaded_params
     if loaded_step is not None:
       step = loaded_step
-  opt_state = opt_init(init_params)
+
+    # Create summary writers.
+    eval_sw = _make_summary_writer(os.path.join(output_dir, "eval_log"))
+    train_sw = _make_summary_writer(os.path.join(output_dir, "train_log"))
 
   log("Starting training.")
+  opt_state = opt_init(init_params)
+  gin_config_saved = False
   while step < train_steps:
     # Training.
     start_time = time.time()
@@ -190,12 +210,33 @@ def update(i, opt_state, batch):
     params = optimizers.get_params(opt_state)
     save_params_and_step(params, step, output_dir)
 
+    # Save the config if not saved yet.
+    # Gin file only includes used parameters, so we save it at this point.
+    if output_dir and not gin_config_saved:
+      gin_config_saved = True
+      config_path = os.path.join(output_dir, "config.gin")
+      with gfile.GFile(config_path, "w") as f:
+        f.write(gin.operative_config_str())
+
     # Evaluation.
     eval_stream = dataset_to_stream(eval_batches, input_name)
     eval_train_stream = dataset_to_stream(train_batches, input_name)
-    train_acc, eval_acc = 0, 0
+    train_acc, eval_acc, train_loss, eval_loss = 0.0, 0.0, 0.0, 0.0
     for _ in range(eval_steps):
       train_acc += accuracy(params, next(eval_train_stream), model_predict)
       eval_acc += accuracy(params, next(eval_stream), model_predict)
-    log("Train set accuracy {:0.4f}".format(train_acc / eval_steps))
-    log("Eval  set accuracy {:0.4f}".format(eval_acc / eval_steps))
+      train_loss += loss(params, next(eval_train_stream), model_predict)
+      eval_loss += loss(params, next(eval_stream), model_predict)
+    train_acc /= eval_steps
+    eval_acc /= eval_steps
+    train_loss /= eval_steps
+    eval_loss /= eval_steps
+    log("Train accuracy {:0.4f} loss {:0.8f}".format(train_acc, train_loss))
+    if train_sw:
+      train_sw.scalar("steps/s", epoch_time / eval_frequency, step=step)
+      train_sw.scalar("accuracy", train_acc, step=step)
+      train_sw.scalar("loss", train_loss, step=step)
+    log("Eval  accuracy {:0.4f} loss {:0.8f}".format(eval_acc, eval_loss))
+    if eval_sw:
+      eval_sw.scalar("accuracy", eval_acc, step=step)
+      train_sw.scalar("loss", eval_loss, step=step)
diff --git a/tensor2tensor/jax/j2j_trainer.py b/tensor2tensor/jax/j2j_trainer.py
index e6e7803d4..7ae480421 100644
--- a/tensor2tensor/jax/j2j_trainer.py
+++ b/tensor2tensor/jax/j2j_trainer.py
@@ -34,8 +34,6 @@
 import gin
 from tensor2tensor.jax import j2j
 
-import tensorflow as tf
-
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string("dataset", None, "Which dataset to use.")
@@ -49,43 +47,27 @@
                           "Configuration parameters (gin string).")
 
 
-def j2j_train(model_name, dataset_name,
-              data_dir=None, output_dir=None, config_file=None, config=None):
-  """Main function to train the given model on the given dataset.
-
-  Args:
-    model_name: The name of the model to train.
-    dataset_name: The name of the dataset to train on.
-    data_dir: Directory where the data is located.
-    output_dir: Directory where to put the logs and checkpoints.
-    config_file: the gin configuration file to use.
-    config: string (in gin format) to override gin parameters.
-  """
-  gin.bind_parameter("train_fn.dataset", dataset_name)
+def _setup_gin():
+  configs = FLAGS.config or []
+  # Override with --dataset and --model
+  if FLAGS.dataset:
+    configs.append("train_fn.dataset='%s'" % FLAGS.dataset)
   if FLAGS.model:
-    config = []  if config is None else config
-    config += ["train_fn.model=@models." + model_name]
-  gin.parse_config_files_and_bindings(config_file, config)
-  if output_dir:
-    if not tf.gfile.Exists(output_dir):
-      tf.gfile.MkDir(output_dir)
-    config_path = os.path.join(output_dir, "gin.config")
-    # TODO(lukaszkaiser): why is the file empty if there's no provided config?
-    with tf.gfile.Open(config_path, "w") as f:
-      f.write(gin.operative_config_str())
-  j2j.train_fn(data_dir, output_dir=output_dir)
+    configs.append("train_fn.model=@" + FLAGS.model)
+  gin.parse_config_files_and_bindings(FLAGS.config_file, configs)
 
 
-def main(argv):
-  del argv
-  logging.set_verbosity(logging.INFO)
+def main(_):
+  _setup_gin()
+
+  # Setup directories
   data_dir, output_dir = FLAGS.data_dir, FLAGS.output_dir
   data_dir = data_dir and os.path.expanduser(data_dir)
   output_dir = output_dir and os.path.expanduser(output_dir)
-  j2j_train(FLAGS.model, FLAGS.dataset,
-            data_dir=data_dir, output_dir=output_dir,
-            config_file=FLAGS.config_file, config=FLAGS.config)
+
+  j2j.train_fn(data_dir, output_dir=output_dir)
 
 
 if __name__ == "__main__":
+  logging.set_verbosity(logging.INFO)
   app.run(main)
diff --git a/tensor2tensor/jax/models/__init__.py b/tensor2tensor/jax/models/__init__.py
new file mode 100644
index 000000000..1769e36ac
--- /dev/null
+++ b/tensor2tensor/jax/models/__init__.py
@@ -0,0 +1,24 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Models defined in J2J."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# pylint: disable=unused-import
+from tensor2tensor.jax.models import mlp
+from tensor2tensor.jax.models import resnet
+# pylint: enable=unused-import
diff --git a/tensor2tensor/jax/models.py b/tensor2tensor/jax/models/mlp.py
similarity index 97%
rename from tensor2tensor/jax/models.py
rename to tensor2tensor/jax/models/mlp.py
index 86d392316..90aeda0d4 100644
--- a/tensor2tensor/jax/models.py
+++ b/tensor2tensor/jax/models/mlp.py
@@ -25,7 +25,7 @@
 
 
 @gin.configurable()
-def mlp(num_hidden_layers=2,
+def MLP(num_hidden_layers=2,
         hidden_size=512,
         activation_fn=stax.Relu,
         num_output_classes=10):
diff --git a/tensor2tensor/jax/models/resnet.py b/tensor2tensor/jax/models/resnet.py
new file mode 100644
index 000000000..a054edcb1
--- /dev/null
+++ b/tensor2tensor/jax/models/resnet.py
@@ -0,0 +1,97 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""J2J models."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gin
+
+from jax.experimental import stax
+
+
+def ConvBlock(kernel_size, filters, strides):
+  """ResNet convolutional striding block."""
+  ks = kernel_size
+  filters1, filters2, filters3 = filters
+  main = stax.serial(
+      stax.Conv(filters1, (1, 1), strides),
+      stax.BatchNorm(), stax.Relu,
+      stax.Conv(filters2, (ks, ks), padding='SAME'),
+      stax.BatchNorm(), stax.Relu,
+      stax.Conv(filters3, (1, 1)), stax.BatchNorm())
+  shortcut = stax.serial(
+      stax.Conv(filters3, (1, 1), strides),
+      stax.BatchNorm())
+  return stax.serial(
+      stax.FanOut(2),
+      stax.parallel(main, shortcut),
+      stax.FanInSum, stax.Relu)
+
+
+def IdentityBlock(kernel_size, filters):
+  """ResNet identical size block."""
+  ks = kernel_size
+  filters1, filters2 = filters
+  def MakeMain(input_shape):
+    # the number of output channels depends on the number of input channels
+    return stax.serial(
+        stax.Conv(filters1, (1, 1)),
+        stax.BatchNorm(), stax.Relu,
+        stax.Conv(filters2, (ks, ks), padding='SAME'),
+        stax.BatchNorm(), stax.Relu,
+        stax.Conv(input_shape[3], (1, 1)), stax.BatchNorm())
+  main = stax.shape_dependent(MakeMain)
+  return stax.serial(
+      stax.FanOut(2),
+      stax.parallel(main, stax.Identity),
+      stax.FanInSum, stax.Relu)
+
+
+@gin.configurable()
+def Resnet50(hidden_size=64, num_output_classes=1001):
+  """ResNet.
+
+  Args:
+    hidden_size: the size of the first hidden layer (multiplied later).
+    num_output_classes: how many classes to distinguish.
+
+  Returns:
+    The ResNet model with the given layer and output sizes.
+  """
+  return stax.serial(
+      stax.Conv(hidden_size, (7, 7), (2, 2), 'SAME'),
+      stax.BatchNorm(), stax.Relu,
+      stax.MaxPool((3, 3), strides=(2, 2)),
+      ConvBlock(3, [hidden_size, hidden_size, 4 * hidden_size], (1, 1)),
+      IdentityBlock(3, [hidden_size, hidden_size]),
+      IdentityBlock(3, [hidden_size, hidden_size]),
+      ConvBlock(3, [2 * hidden_size, 2 * hidden_size, 8 * hidden_size], (2, 2)),
+      IdentityBlock(3, [2 * hidden_size, 2 * hidden_size]),
+      IdentityBlock(3, [2 * hidden_size, 2 * hidden_size]),
+      IdentityBlock(3, [2 * hidden_size, 2 * hidden_size]),
+      ConvBlock(3, [4 * hidden_size, 4 * hidden_size, 16*hidden_size], (2, 2)),
+      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size]),
+      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size]),
+      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size]),
+      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size]),
+      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size]),
+      ConvBlock(3, [8 * hidden_size, 8 * hidden_size, 32*hidden_size], (2, 2)),
+      IdentityBlock(3, [8 * hidden_size, 8 * hidden_size]),
+      IdentityBlock(3, [8 * hidden_size, 8 * hidden_size]),
+      stax.AvgPool((7, 7)), stax.Flatten,
+      stax.Dense(num_output_classes), stax.LogSoftmax)

From fcf9819430ebfb4839ae5faa91ce0ff767254439 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 22 Feb 2019 22:20:13 -0800
Subject: [PATCH 1711/2720] use decode.log_results.

PiperOrigin-RevId: 235310788
---
 tensor2tensor/utils/decoding.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index df61b5f08..6d6a15624 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -218,7 +218,7 @@ def decode_from_dataset(estimator,
                          decode_hp,
                          decode_to_file,
                          output_dir,
-                         log_results=True,
+                         log_results=decode_hp.log_results,
                          checkpoint_path=checkpoint_path)
 
     if decode_hp.decode_in_memory:
@@ -340,7 +340,7 @@ def decode_once(estimator,
             output_dir=output_dir,
             identity_output=decode_hp.identity_output,
             targets=targets,
-            log_results=decode_hp.log_results)
+            log_results=log_results)
         decoded_outputs.append(decoded)
         if decode_hp.write_beam_scores:
           decoded_scores.append(score)
@@ -356,7 +356,7 @@ def decode_once(estimator,
           output_dir=output_dir,
           identity_output=decode_hp.identity_output,
           targets=targets,
-          log_results=decode_hp.log_results)
+          log_results=log_results)
       decoded_outputs.append(decoded)
 
     # Write out predictions if decode_to_file passed

From 576b7a83b0d755491ce5845a72b884cc103c3b90 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Sat, 23 Feb 2019 22:35:55 -0800
Subject: [PATCH 1712/2720] make exception catcher more narrow.  Previous code
 was catching import errors it should not have.

PiperOrigin-RevId: 235387243
---
 tensor2tensor/data_generators/all_problems.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 157970685..c1f54f18f 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -19,7 +19,6 @@
 from __future__ import print_function
 
 import importlib
-import re
 
 MODULES = [
     "tensor2tensor.data_generators.algorithmic",
@@ -97,8 +96,11 @@
 
 
 def _is_import_err_msg(err_str, module):
-  module_pattern = "(.)?".join(["(%s)?" % m for m in module.split(".")])
-  return re.match("^No module named (')?%s(')?$" % module_pattern, err_str)
+  parts = module.split(".")
+  suffixes = [".".join(parts[i:]) for i in xrange(len(parts))]
+  return err_str in (
+      ["No module named %s" % suffix for suffix in suffixes] +
+      ["No module named '%s'" % suffix for suffix in suffixes])
 
 
 def _handle_errors(errors):

From 2ec4e737785af36ef52e42b793e55d697b2a6f1a Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Sun, 24 Feb 2019 18:43:54 -0800
Subject: [PATCH 1713/2720] Update modality description in README.md.

PiperOrigin-RevId: 235452836
---
 README.md           | 22 +++++++++-------------
 docs/walkthrough.md | 22 +++++++++-------------
 2 files changed, 18 insertions(+), 26 deletions(-)

diff --git a/README.md b/README.md
index e0bd75a1d..fea78ad2e 100644
--- a/README.md
+++ b/README.md
@@ -354,24 +354,20 @@ and many common sequence datasets are already available for generation and use.
 
 ### Problems and Modalities
 
-**Problems** define training-time hyperparameters for the dataset and task,
-mainly by setting input and output **modalities** (e.g. symbol, image, audio,
-label) and vocabularies, if applicable. All problems are defined either in
-[`problem_hparams.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/problem_hparams.py)
+**Problems** consist of features such as inputs and targets, and metadata such
+as each feature's modality (e.g. symbol, image, audio) and vocabularies. All
+problems are imported in
+[`all_problems.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/all_problems.py)
 or are registered with `@registry.register_problem` (run `t2t-datagen` to see
 the list of all available problems).
-**Modalities**, defined in
-[`modality.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/modality.py),
-abstract away the input and output data types so that **models** may deal with
-modality-independent tensors.
 
 ### Models
 
-**`T2TModel`s** define the core tensor-to-tensor transformation, independent of
-input/output modality or task. Models take dense tensors in and produce dense
-tensors that may then be transformed in a final step by a **modality** depending
-on the task (e.g. fed through a final linear transform to produce logits for a
-softmax over classes). All models are imported in the
+**`T2TModel`s** define the core tensor-to-tensor computation. They apply a
+default transformation to each input and output so that models may deal with
+modality-independent tensors (e.g. embeddings at the input; and a linear
+transform at the output to produce logits for a softmax over classes). All
+models are imported in the
 [`models` subpackage](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/models/__init__.py),
 inherit from `T2TModel` - defined in
 [`t2t_model.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/t2t_model.py) -
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index e0bd75a1d..fea78ad2e 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -354,24 +354,20 @@ and many common sequence datasets are already available for generation and use.
 
 ### Problems and Modalities
 
-**Problems** define training-time hyperparameters for the dataset and task,
-mainly by setting input and output **modalities** (e.g. symbol, image, audio,
-label) and vocabularies, if applicable. All problems are defined either in
-[`problem_hparams.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/problem_hparams.py)
+**Problems** consist of features such as inputs and targets, and metadata such
+as each feature's modality (e.g. symbol, image, audio) and vocabularies. All
+problems are imported in
+[`all_problems.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/all_problems.py)
 or are registered with `@registry.register_problem` (run `t2t-datagen` to see
 the list of all available problems).
-**Modalities**, defined in
-[`modality.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/modality.py),
-abstract away the input and output data types so that **models** may deal with
-modality-independent tensors.
 
 ### Models
 
-**`T2TModel`s** define the core tensor-to-tensor transformation, independent of
-input/output modality or task. Models take dense tensors in and produce dense
-tensors that may then be transformed in a final step by a **modality** depending
-on the task (e.g. fed through a final linear transform to produce logits for a
-softmax over classes). All models are imported in the
+**`T2TModel`s** define the core tensor-to-tensor computation. They apply a
+default transformation to each input and output so that models may deal with
+modality-independent tensors (e.g. embeddings at the input; and a linear
+transform at the output to produce logits for a softmax over classes). All
+models are imported in the
 [`models` subpackage](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/models/__init__.py),
 inherit from `T2TModel` - defined in
 [`t2t_model.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/t2t_model.py) -

From f1655a28e630c883b17179a16399665b50add1af Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 25 Feb 2019 12:33:44 -0800
Subject: [PATCH 1714/2720] Refine T2T README.md.

I cleaned up some descriptions. Happy to make any additional changes.

PiperOrigin-RevId: 235581990
---
 README.md           | 42 +++++++++++++++++-------------------------
 docs/walkthrough.md | 42 +++++++++++++++++-------------------------
 2 files changed, 34 insertions(+), 50 deletions(-)

diff --git a/README.md b/README.md
index fea78ad2e..0f2f1ff01 100644
--- a/README.md
+++ b/README.md
@@ -333,8 +333,8 @@ python -c "from tensor2tensor.models.transformer import Transformer"
   request for public datasets!).
 * Models can be used with any dataset and input mode (or even multiple); all
   modality-specific processing (e.g. embedding lookups for text tokens) is done
-  with `Modality` objects, which are specified per-feature in the dataset/task
-  specification.
+  with `bottom` and `top` transformations, which are specified per-feature in the
+  model.
 * Support for multi-GPU machines and synchronous (1 master, many workers) and
   asynchronous (independent workers synchronizing through a parameter server)
   [distributed training](https://tensorflow.github.io/tensor2tensor/distributed_training.html).
@@ -344,22 +344,17 @@ python -c "from tensor2tensor.models.transformer import Transformer"
 
 ## T2T overview
 
-### Datasets
-
-**Datasets** are all standardized on `TFRecord` files with `tensorflow.Example`
-protocol buffers. All datasets are registered and generated with the
-[data
-generator](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/bin/t2t-datagen)
-and many common sequence datasets are already available for generation and use.
-
-### Problems and Modalities
+### Problems
 
 **Problems** consist of features such as inputs and targets, and metadata such
-as each feature's modality (e.g. symbol, image, audio) and vocabularies. All
+as each feature's modality (e.g. symbol, image, audio) and vocabularies. Problem
+features are given by a dataset, which is stored as a `TFRecord` file with
+`tensorflow.Example` protocol buffers. All
 problems are imported in
 [`all_problems.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/all_problems.py)
-or are registered with `@registry.register_problem` (run `t2t-datagen` to see
-the list of all available problems).
+or are registered with `@registry.register_problem`. Run
+[`t2t-datagen`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/bin/t2t-datagen)
+to see the list of available problems and download them.
 
 ### Models
 
@@ -369,25 +364,24 @@ modality-independent tensors (e.g. embeddings at the input; and a linear
 transform at the output to produce logits for a softmax over classes). All
 models are imported in the
 [`models` subpackage](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/models/__init__.py),
-inherit from `T2TModel` - defined in
-[`t2t_model.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/t2t_model.py) -
+inherit from [`T2TModel`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/t2t_model.py),
 and are registered with
 [`@registry.register_model`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/registry.py).
 
 ### Hyperparameter Sets
 
-**Hyperparameter sets** are defined and registered in code with
-[`@registry.register_hparams`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/registry.py)
-and are encoded in
+**Hyperparameter sets** are encoded in
 [`HParams`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/hparam.py)
-objects. The `HParams` are available to both the problem specification and the
-model. A basic set of hyperparameters are defined in
+objects, and are registered with
+[`@registry.register_hparams`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/registry.py).
+Every model and problem has a `HParams`. A basic set of hyperparameters are
+defined in
 [`common_hparams.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/layers/common_hparams.py)
 and hyperparameter set functions can compose other hyperparameter set functions.
 
 ### Trainer
 
-The **trainer** binary is the main entrypoint for training, evaluation, and
+The **trainer** binary is the entrypoint for training, evaluation, and
 inference. Users can easily switch between problems, models, and hyperparameter
 sets by using the `--model`, `--problem`, and `--hparams_set` flags. Specific
 hyperparameters can be overridden with the `--hparams` flag. `--schedule` and
@@ -413,9 +407,7 @@ To add a new dataset, subclass
 [`Problem`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/problem.py)
 and register it with `@registry.register_problem`. See
 [`TranslateEndeWmt8k`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/translate_ende.py)
-for an example.
-
-Also see the [data generators
+for an example. Also see the [data generators
 README](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/README.md).
 
 ## Run on FloydHub
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index fea78ad2e..0f2f1ff01 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -333,8 +333,8 @@ python -c "from tensor2tensor.models.transformer import Transformer"
   request for public datasets!).
 * Models can be used with any dataset and input mode (or even multiple); all
   modality-specific processing (e.g. embedding lookups for text tokens) is done
-  with `Modality` objects, which are specified per-feature in the dataset/task
-  specification.
+  with `bottom` and `top` transformations, which are specified per-feature in the
+  model.
 * Support for multi-GPU machines and synchronous (1 master, many workers) and
   asynchronous (independent workers synchronizing through a parameter server)
   [distributed training](https://tensorflow.github.io/tensor2tensor/distributed_training.html).
@@ -344,22 +344,17 @@ python -c "from tensor2tensor.models.transformer import Transformer"
 
 ## T2T overview
 
-### Datasets
-
-**Datasets** are all standardized on `TFRecord` files with `tensorflow.Example`
-protocol buffers. All datasets are registered and generated with the
-[data
-generator](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/bin/t2t-datagen)
-and many common sequence datasets are already available for generation and use.
-
-### Problems and Modalities
+### Problems
 
 **Problems** consist of features such as inputs and targets, and metadata such
-as each feature's modality (e.g. symbol, image, audio) and vocabularies. All
+as each feature's modality (e.g. symbol, image, audio) and vocabularies. Problem
+features are given by a dataset, which is stored as a `TFRecord` file with
+`tensorflow.Example` protocol buffers. All
 problems are imported in
 [`all_problems.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/all_problems.py)
-or are registered with `@registry.register_problem` (run `t2t-datagen` to see
-the list of all available problems).
+or are registered with `@registry.register_problem`. Run
+[`t2t-datagen`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/bin/t2t-datagen)
+to see the list of available problems and download them.
 
 ### Models
 
@@ -369,25 +364,24 @@ modality-independent tensors (e.g. embeddings at the input; and a linear
 transform at the output to produce logits for a softmax over classes). All
 models are imported in the
 [`models` subpackage](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/models/__init__.py),
-inherit from `T2TModel` - defined in
-[`t2t_model.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/t2t_model.py) -
+inherit from [`T2TModel`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/t2t_model.py),
 and are registered with
 [`@registry.register_model`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/registry.py).
 
 ### Hyperparameter Sets
 
-**Hyperparameter sets** are defined and registered in code with
-[`@registry.register_hparams`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/registry.py)
-and are encoded in
+**Hyperparameter sets** are encoded in
 [`HParams`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/hparam.py)
-objects. The `HParams` are available to both the problem specification and the
-model. A basic set of hyperparameters are defined in
+objects, and are registered with
+[`@registry.register_hparams`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/registry.py).
+Every model and problem has a `HParams`. A basic set of hyperparameters are
+defined in
 [`common_hparams.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/layers/common_hparams.py)
 and hyperparameter set functions can compose other hyperparameter set functions.
 
 ### Trainer
 
-The **trainer** binary is the main entrypoint for training, evaluation, and
+The **trainer** binary is the entrypoint for training, evaluation, and
 inference. Users can easily switch between problems, models, and hyperparameter
 sets by using the `--model`, `--problem`, and `--hparams_set` flags. Specific
 hyperparameters can be overridden with the `--hparams` flag. `--schedule` and
@@ -413,9 +407,7 @@ To add a new dataset, subclass
 [`Problem`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/problem.py)
 and register it with `@registry.register_problem`. See
 [`TranslateEndeWmt8k`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/translate_ende.py)
-for an example.
-
-Also see the [data generators
+for an example. Also see the [data generators
 README](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/README.md).
 
 ## Run on FloydHub

From c564d0322cc0c8edbb9f5f2c434026e54f3bff54 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 25 Feb 2019 13:08:24 -0800
Subject: [PATCH 1715/2720] Add a function `play_env_problem_randomly` that
 plays an env_problem randomly. Change Trajectory.{last_time_step,
 num_time_steps, is_active} to be properties. Add two functions
 BatchTrajectory.{num_completed_time_steps, num_time_steps}.

PiperOrigin-RevId: 235588714
---
 tensor2tensor/envs/env_problem.py            |  6 +-
 tensor2tensor/envs/env_problem_test.py       |  2 +-
 tensor2tensor/envs/env_problem_utils.py      | 44 +++++++++++
 tensor2tensor/envs/env_problem_utils_test.py | 48 ++++++++++++
 tensor2tensor/envs/trajectory.py             | 33 ++++++---
 tensor2tensor/envs/trajectory_test.py        | 77 ++++++++++++--------
 6 files changed, 166 insertions(+), 44 deletions(-)
 create mode 100644 tensor2tensor/envs/env_problem_utils.py
 create mode 100644 tensor2tensor/envs/env_problem_utils_test.py

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 1b2ffde24..04b11d0b7 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -664,7 +664,7 @@ def _generate_time_steps(self, trajectory_list):
       # Skip writing trajectories that have only a single time-step -- this
       # could just be a repeated reset.
 
-      if single_trajectory.num_time_steps() <= 1:
+      if single_trajectory.num_time_steps <= 1:
         continue
 
       for index, time_step in enumerate(single_trajectory.time_steps):
@@ -760,8 +760,8 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
   def print_state(self):
     for t in self.trajectories.trajectories:
       print("---------")
-      if not t.is_active():
+      if not t.is_active:
         print("trajectory isn't active.")
         continue
-      last_obs = t.last_time_step().observation
+      last_obs = t.last_time_step.observation
       print(str(last_obs))
diff --git a/tensor2tensor/envs/env_problem_test.py b/tensor2tensor/envs/env_problem_test.py
index eedb6c171..f0d9db00a 100644
--- a/tensor2tensor/envs/env_problem_test.py
+++ b/tensor2tensor/envs/env_problem_test.py
@@ -146,7 +146,7 @@ def test_interaction_with_env(self):
 
       # This should also match the number of time-steps completed given by ep.
       num_timesteps_completed_ep = sum(
-          ct.num_time_steps() for ct in ep.trajectories.completed_trajectories)
+          ct.num_time_steps for ct in ep.trajectories.completed_trajectories)
       self.assertEqual(num_timesteps_completed, num_timesteps_completed_ep)
 
     # Reset the trajectories.
diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
new file mode 100644
index 000000000..0391ba347
--- /dev/null
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -0,0 +1,44 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Utilities to deal with EnvProblem."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensor2tensor.envs import env_problem as env_problem_lib
+
+
+def play_env_problem_randomly(env_problem,
+                              num_steps):
+  """Plays the env problem by randomly sampling actions for `num_steps`."""
+  # Reset all environments.
+  env_problem.reset()
+
+  # Play all environments, sampling random actions each time.
+  for _ in range(num_steps):
+    # Sample batch_size actions from the action space and stack them.
+    actions = np.stack([env_problem.action_space.sample() for _ in range(
+        env_problem.batch_size)])
+
+    # Execute actions, observations are stored in `env_problem`.
+    _, _, dones, _ = env_problem.step(actions)
+
+    # Get the indices where we are done and reset those.
+    done_indices = env_problem_lib.EnvProblem.done_indices(dones)
+    env_problem.reset(indices=done_indices)
diff --git a/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
new file mode 100644
index 000000000..8f1e0db51
--- /dev/null
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -0,0 +1,48 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 env_problem_utils."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.envs import env_problem_utils
+from tensor2tensor.envs import tic_tac_toe_env  # pylint: disable=unused-import
+from tensor2tensor.envs import tic_tac_toe_env_problem
+
+import tensorflow as tf
+
+
+class EnvProblemUtilsTest(tf.test.TestCase):
+
+  def test_play_env_problem_randomly(self):
+    batch_size = 5
+    num_steps = 100
+
+    ep = tic_tac_toe_env_problem.TicTacToeEnvProblem()
+    ep.initialize(batch_size=batch_size)
+
+    env_problem_utils.play_env_problem_randomly(ep, num_steps)
+
+    # We've played num_steps * batch_size steps + everytime we get 'done' we
+    # create another step + batch_size number of pending steps.
+    self.assertEqual(
+        num_steps * batch_size + len(ep.trajectories.completed_trajectories) +
+        batch_size, ep.trajectories.num_time_steps)
+
+
+if __name__ == '__main__':
+  tf.test.main()
diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index 59fc3e1c7..557f2db3d 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -44,7 +44,7 @@ def add_time_step(self, **create_time_step_kwargs):
 
     Args:
       **create_time_step_kwargs: Forwarded to
-                                 time_step.TimeStep.create_time_step.
+        time_step.TimeStep.create_time_step.
     """
     ts = time_step.TimeStep.create_time_step(**create_time_step_kwargs)
     assert isinstance(ts, time_step.TimeStep)
@@ -58,17 +58,19 @@ def change_last_time_step(self, **replace_time_step_kwargs):
     self._time_steps[-1] = self._time_steps[-1].replace(
         **replace_time_step_kwargs)
 
+  @property
   def last_time_step(self):
     # Pre-conditions: self._time_steps shouldn't be empty.
     assert self._time_steps
     return self._time_steps[-1]
 
-  # We could have overridden __nonzero__ or __bool__ as well.
+  @property
   def num_time_steps(self):
     return len(self._time_steps)
 
+  @property
   def is_active(self):
-    return bool(self.num_time_steps())
+    return bool(self.num_time_steps)
 
   @property
   def time_steps(self):
@@ -76,7 +78,7 @@ def time_steps(self):
 
   @property
   def done(self):
-    return self.is_active() and self.last_time_step().done
+    return self.is_active and self.last_time_step.done
 
   # TODO(afrozm): Add discounting and rewards-to-go when it makes sense.
   @property
@@ -130,7 +132,7 @@ def _complete_trajectory(self, trajectory, index):
     assert isinstance(trajectory, Trajectory)
 
     # This *should* be the case.
-    assert trajectory.last_time_step().action is None
+    assert trajectory.last_time_step.action is None
 
     # Add to completed trajectories.
     self._completed_trajectories.append(trajectory)
@@ -164,7 +166,7 @@ def reset(self, indices, observations):
       trajectory = self._trajectories[index]
 
       # Are we starting a new trajectory at the given index?
-      if not trajectory.is_active():
+      if not trajectory.is_active:
         # Then create a new time-step here with the given observation.
         trajectory.add_time_step(observation=observation)
         # That's all we need to do here.
@@ -189,7 +191,7 @@ def complete_all_trajectories(self):
     """Essentially same as reset, but we don't have observations."""
     for index in range(self.batch_size):
       trajectory = self._trajectories[index]
-      assert trajectory.is_active()
+      assert trajectory.is_active
       self._complete_trajectory(trajectory, index)
 
   def step(self, observations, raw_rewards, processed_rewards, dones, actions):
@@ -238,7 +240,7 @@ def step(self, observations, raw_rewards, processed_rewards, dones, actions):
       # a prior observation from which we are stepping away from.
 
       # TODO(afrozm): Let's re-visit this if it becomes too restrictive.
-      assert trajectory.is_active()
+      assert trajectory.is_active
 
       # To this trajectory's last time-step, set actions.
       trajectory.change_last_time_step(action=actions[index])
@@ -257,4 +259,17 @@ def step(self, observations, raw_rewards, processed_rewards, dones, actions):
 
         # NOTE: The new trajectory at `index` is going to be in-active and
         # `reset` should be called on it.
-        assert not self._trajectories[index].is_active()
+        assert not self._trajectories[index].is_active
+
+  @property
+  def num_completed_time_steps(self):
+    """Returns the number of time-steps in completed trajectories."""
+
+    return sum(t.num_time_steps for t in self.completed_trajectories)
+
+  @property
+  def num_time_steps(self):
+    """Returns the number of time-steps in completed and incomplete trajectories."""
+
+    num_time_steps = sum(t.num_time_steps for t in self.trajectories)
+    return num_time_steps + self.num_completed_time_steps
diff --git a/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index 36ab17cac..92105c220 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -28,8 +28,8 @@ class TrajectoryTest(tf.test.TestCase):
 
   def test_empty_trajectory(self):
     t = trajectory.Trajectory()
-    self.assertFalse(t.is_active())
-    self.assertEqual(0, t.num_time_steps())
+    self.assertFalse(t.is_active)
+    self.assertEqual(0, t.num_time_steps)
     self.assertFalse(t.done)
 
   def test_add_time_step(self):
@@ -37,28 +37,28 @@ def test_add_time_step(self):
     t.add_time_step(observation=1, done=True)
 
     # Test that the trajectory is now active.
-    self.assertTrue(t.is_active())
+    self.assertTrue(t.is_active)
 
-    added_t = t.last_time_step()
+    added_t = t.last_time_step
     self.assertEqual(1, added_t.observation)
     self.assertTrue(added_t.done)
     self.assertIsNone(None, added_t.raw_reward)
     self.assertIsNone(None, added_t.processed_reward)
     self.assertIsNone(None, added_t.action)
 
-    self.assertEqual(1, t.num_time_steps())
+    self.assertEqual(1, t.num_time_steps)
 
   def test_change_last_time_step(self):
     t = trajectory.Trajectory()
     t.add_time_step(observation=1, done=False)
     t.add_time_step(observation=1, done=True)
-    self.assertTrue(t.is_active())
+    self.assertTrue(t.is_active)
 
-    num_ts_old = t.num_time_steps()
+    num_ts_old = t.num_time_steps
     self.assertEqual(2, num_ts_old)
 
     # Assert on what the last time-step is currently.
-    ts = t.last_time_step()
+    ts = t.last_time_step
     self.assertEqual(1, ts.observation)
     self.assertTrue(ts.done)
     self.assertEqual(None, ts.action)
@@ -67,13 +67,13 @@ def test_change_last_time_step(self):
     t.change_last_time_step(done=False, action=5)
 
     # Assert that it changed.
-    ts = t.last_time_step()
+    ts = t.last_time_step
     self.assertEqual(1, ts.observation)  # unchanged, since we didn't change it.
     self.assertFalse(ts.done)  # was True earlier
     self.assertEqual(5, ts.action)  # was None earlier
 
     # Assert on the number of steps remaining the same as before.
-    self.assertEqual(num_ts_old, t.num_time_steps())
+    self.assertEqual(num_ts_old, t.num_time_steps)
 
   def test_reward(self):
     t = trajectory.Trajectory()
@@ -96,6 +96,17 @@ class BatchTrajectoryTest(tf.test.TestCase):
   BATCH_SIZE = 10
   OBSERVATION_SHAPE = (3, 4)
 
+  def get_random_observations_rewards_actions_dones(self, batch_size=None):
+    batch_size = batch_size or self.BATCH_SIZE
+    # Random observations, rewards, actions, done of the expected shape.
+    observations = np.random.rand(*((batch_size,) + self.OBSERVATION_SHAPE))
+    raw_rewards = np.random.randn(batch_size)
+    actions = np.random.randn(batch_size)
+    # 40% change of being done.
+    dones = np.random.random((batch_size,)) > 0.6
+
+    return observations, raw_rewards, actions, dones
+
   def test_creation(self):
     bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
 
@@ -106,32 +117,37 @@ def test_reset_all(self):
     bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
 
     indices = np.arange(self.BATCH_SIZE)
-    observations = np.random.rand(*(
-        (self.BATCH_SIZE,) + self.OBSERVATION_SHAPE))
+    observations, _, _, _ = self.get_random_observations_rewards_actions_dones()
 
     # Call reset.
     bt.reset(indices, observations)
 
     # Assert that all trajectories are active and not done (reset never marks
     # anything as done).
-    self.assertTrue(all(t.is_active() for t in bt.trajectories))
+    self.assertTrue(all(t.is_active for t in bt.trajectories))
     self.assertEqual(0, len(bt.completed_trajectories))
 
+  def test_num_time_steps(self):
+    bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
+
+    self.assertEqual(0, bt.num_completed_time_steps)
+    self.assertEqual(0, bt.num_time_steps)
+
   def test_reset_some(self):
     bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
 
     indices = np.arange(self.BATCH_SIZE // 2)
-    observations = np.random.rand(*(
-        (self.BATCH_SIZE // 2,) + self.OBSERVATION_SHAPE))
+    observations, _, _, _ = self.get_random_observations_rewards_actions_dones(
+        batch_size=self.BATCH_SIZE // 2)
 
     # Just reset the first half.
     bt.reset(indices, observations)
 
     # So first half are active, rest aren't.
     self.assertTrue(
-        all(t.is_active() for t in bt.trajectories[:self.BATCH_SIZE // 2]))
+        all(t.is_active for t in bt.trajectories[:self.BATCH_SIZE // 2]))
     self.assertTrue(
-        all(not t.is_active() for t in bt.trajectories[self.BATCH_SIZE // 2:]))
+        all(not t.is_active for t in bt.trajectories[self.BATCH_SIZE // 2:]))
 
     # Nothing is done anyways.
     self.assertEqual(0, len(bt.completed_trajectories))
@@ -140,18 +156,15 @@ def test_step(self):
     bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
 
     indices = np.arange(self.BATCH_SIZE)
-    observations = np.random.rand(*(
-        (self.BATCH_SIZE,) + self.OBSERVATION_SHAPE))
+    observations, _, _, _ = self.get_random_observations_rewards_actions_dones()
 
     # Have to call reset first.
     bt.reset(indices, observations)
 
     # Create some fake data for calling step.
-    new_observations = np.random.rand(*(
-        (self.BATCH_SIZE,) + self.OBSERVATION_SHAPE))
-    raw_rewards = processed_rewards = actions = np.random.randn(self.BATCH_SIZE)
-    processed_rewards = np.int64(processed_rewards)
-    dones = raw_rewards > 0.5
+    new_observations, raw_rewards, actions, dones = (
+        self.get_random_observations_rewards_actions_dones())
+    processed_rewards = raw_rewards.astype(np.int64)
 
     # Force mark the first one as done anyways, so that there is something to
     # test.
@@ -169,7 +182,7 @@ def test_step(self):
     self.assertEqual(num_done, len(bt.completed_trajectories))
 
     # Expect to see that the rest are marked as active.
-    num_active = sum(t.is_active() for t in bt.trajectories)
+    num_active = sum(t.is_active for t in bt.trajectories)
     self.assertEqual(num_not_done, num_active)
 
   def test_desired_placement_of_rewards_and_actions(self):
@@ -177,15 +190,17 @@ def test_desired_placement_of_rewards_and_actions(self):
     bt = trajectory.BatchTrajectory(batch_size=batch_size)
 
     indices = np.arange(batch_size)
-    observations = np.random.rand(*((batch_size,) + self.OBSERVATION_SHAPE))
+    observations, _, _, _ = self.get_random_observations_rewards_actions_dones(
+        batch_size=batch_size)
 
     # Have to call reset first.
     bt.reset(indices, observations)
 
     # Create some fake data for calling step.
-    new_observations = np.random.rand(*((batch_size,) + self.OBSERVATION_SHAPE))
-    raw_rewards = processed_rewards = actions = np.random.randn(batch_size)
-    processed_rewards = processed_rewards.astype(np.int64)
+    new_observations, raw_rewards, actions, _ = (
+        self.get_random_observations_rewards_actions_dones(
+            batch_size=batch_size))
+    processed_rewards = raw_rewards.astype(np.int64)
     dones = np.full(batch_size, False)
 
     # Call step.
@@ -197,8 +212,8 @@ def test_desired_placement_of_rewards_and_actions(self):
     # The only trajectory is active.
     self.assertEqual(batch_size, len(bt.trajectories))
     t = bt.trajectories[0]
-    self.assertTrue(t.is_active())
-    self.assertEqual(2, t.num_time_steps())
+    self.assertTrue(t.is_active)
+    self.assertEqual(2, t.num_time_steps)
 
     ts = t.time_steps
 

From c46548e3b1e9c3aef8f6010b1df6bf0356a38619 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 25 Feb 2019 13:56:07 -0800
Subject: [PATCH 1716/2720] Move EnvProblem.done_indices to env_problem_utils,
 it seemed like a more appropriate place.

PiperOrigin-RevId: 235597040
---
 tensor2tensor/envs/env_problem.py                  | 7 +------
 tensor2tensor/envs/env_problem_test.py             | 5 +++--
 tensor2tensor/envs/env_problem_utils.py            | 8 +++++---
 tensor2tensor/envs/tic_tac_toe_env_problem_test.py | 3 ++-
 4 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 04b11d0b7..ee8e62878 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -81,7 +81,7 @@ class EnvProblem(Env, problem.Problem):
   obs, rewards, dones, infos = ep.step(actions)
 
   # 4. Figure out which envs got done and reset only those.
-  ep.reset(indices=done_indices(dones))
+  ep.reset(indices=env_problem_utils.done_indices(dones))
 
   # 5. Go back to Step #3 to further interact with the env or just dump the
   # generated data to disk by calling:
@@ -563,11 +563,6 @@ def step(self, actions):
 
     return processed_observations, processed_rewards, dones, infos
 
-  @staticmethod
-  def done_indices(dones):
-    """Calculates the indices where dones has True."""
-    return np.argwhere(dones).squeeze(axis=1)
-
   def example_reading_spec(self):
     """Data fields to store on disk and their decoders."""
 
diff --git a/tensor2tensor/envs/env_problem_test.py b/tensor2tensor/envs/env_problem_test.py
index f0d9db00a..1f902fb19 100644
--- a/tensor2tensor/envs/env_problem_test.py
+++ b/tensor2tensor/envs/env_problem_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.envs import env_problem
+from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.layers import modalities
 import tensorflow as tf
 
@@ -129,7 +130,7 @@ def test_interaction_with_env(self):
                        len(ep.trajectories.completed_trajectories))
 
       # Get the indices where we are done ...
-      done_indices = env_problem.EnvProblem.done_indices(dones)
+      done_indices = env_problem_utils.done_indices(dones)
 
       # ... and reset those.
       ep.reset(indices=done_indices)
@@ -213,7 +214,7 @@ def play_env(self,
       # Step through it.
       _, _, dones, _ = env.step(actions)
       # Get the indices where we are done ...
-      done_indices = env_problem.EnvProblem.done_indices(dones)
+      done_indices = env_problem_utils.done_indices(dones)
       # ... and reset those.
       env.reset(indices=done_indices)
       # count the number of dones we got, in this step and overall.
diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 0391ba347..24e845cd4 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -21,7 +21,10 @@
 
 import numpy as np
 
-from tensor2tensor.envs import env_problem as env_problem_lib
+
+def done_indices(dones):
+  """Calculates the indices where dones has True."""
+  return np.argwhere(dones).squeeze(axis=1)
 
 
 def play_env_problem_randomly(env_problem,
@@ -40,5 +43,4 @@ def play_env_problem_randomly(env_problem,
     _, _, dones, _ = env_problem.step(actions)
 
     # Get the indices where we are done and reset those.
-    done_indices = env_problem_lib.EnvProblem.done_indices(dones)
-    env_problem.reset(indices=done_indices)
+    env_problem.reset(indices=done_indices(dones))
diff --git a/tensor2tensor/envs/tic_tac_toe_env_problem_test.py b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
index f9d4e6a03..ac5a549cf 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import numpy as np
+from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.envs import tic_tac_toe_env  # pylint: disable=unused-import
 from tensor2tensor.envs import tic_tac_toe_env_problem  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
@@ -45,7 +46,7 @@ def test_registration_and_interaction_with_env_problem(self):
       self.assertEqual(batch_size, len(dones))
       self.assertEqual(batch_size, len(infos))
 
-      done_indices = ep.done_indices(dones)
+      done_indices = env_problem_utils.done_indices(dones)
       ep.reset(done_indices)
       num_done += sum(dones)
       for r, d in zip(rewards, dones):

From 6f3b35063c3a638e9f0d5ceb11881160e8765ad9 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 25 Feb 2019 19:14:07 -0800
Subject: [PATCH 1717/2720] Improve logs in jax training.

PiperOrigin-RevId: 235645051
---
 tensor2tensor/data_generators/imagenet.py     |  2 +
 .../jax/configs/resnet50_imagenet_8gb.gin     |  8 +-
 tensor2tensor/jax/j2j.py                      | 86 +++++++++++++------
 3 files changed, 68 insertions(+), 28 deletions(-)

diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index b414ab4e3..9d724fe6e 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -573,6 +573,7 @@ def preprocess_for_train(image, image_size=224, normalize=True):
   Returns:
     A preprocessed image `Tensor`.
   """
+  if normalize: image = tf.to_float(image) / 255.0
   image = _random_crop(image, image_size)
   if normalize: image = _normalize(image)
   image = _flip(image)
@@ -591,6 +592,7 @@ def preprocess_for_eval(image, image_size=224, normalize=True):
   Returns:
     A preprocessed image `Tensor`.
   """
+  if normalize: image = tf.to_float(image) / 255.0
   image = _do_scale(image, image_size + 32)
   if normalize: image = _normalize(image)
   image = _center_crop(image, image_size)
diff --git a/tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin
index c75bef109..c825e0cbc 100644
--- a/tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin
@@ -1,15 +1,15 @@
 # Parameters for batch_fn:
 # ==============================================================================
-batch_fn.batch_size = 16
+batch_fn.batch_size = 32
 batch_fn.bucket_length = 32
 batch_fn.buckets = None
 batch_fn.eval_batch_size = 16
 
 # Parameters for learning_rate:
 # ==============================================================================
-learning_rate.constant = 0.1
+learning_rate.constant = 10.0
 learning_rate.schedule = 'constant * linear_warmup * rsqrt_decay'
-learning_rate.warmup_steps = 10000
+learning_rate.warmup_steps = 400
 
 # Parameters for optimizer:
 # ==============================================================================
@@ -32,7 +32,7 @@ Resnet50.num_output_classes = 1001
 
 # Parameters for train_fn:
 # ==============================================================================
-train_fn.dataset = 't2t_image_imagenet'
+train_fn.dataset = 't2t_image_imagenet224'
 train_fn.eval_frequency = 100
 train_fn.eval_steps = 2
 train_fn.model = @Resnet50
diff --git a/tensor2tensor/jax/j2j.py b/tensor2tensor/jax/j2j.py
index e3351bc72..1768bbb29 100644
--- a/tensor2tensor/jax/j2j.py
+++ b/tensor2tensor/jax/j2j.py
@@ -81,24 +81,33 @@ def one_hot(x, k, dtype=np.float32):
   return np.array(x[:, None] == np.arange(k), dtype)
 
 
-def accuracy(params, batch, model_predict):
+def accuracy(batch, model_predictions):
   """Calculate accuracy."""
-  inputs, targets = batch
-  predicted_class = np.argmax(model_predict(params, inputs), axis=1)
+  _, targets = batch
+  predicted_class = np.argmax(model_predictions, axis=-1)
   return np.mean(predicted_class == targets)
 
 
+def neg_log_perplexity(batch, model_predictions):
+  """Calculate negative log perplexity."""
+  _, targets = batch
+  hot_targets = one_hot(targets, model_predictions.shape[-1])
+  return np.mean(np.sum(model_predictions * hot_targets, axis=-1))
+
+
 def loss(params, batch, model_predict):
   """Calculate loss."""
   inputs, targets = batch
   preds = model_predict(params, inputs)
-  return -np.mean(preds * one_hot(targets, preds.shape[-1]))
+  return - np.mean(preds * one_hot(targets, preds.shape[-1]))
 
 
 def dataset_to_stream(batches, input_name):
   """Takes a tf.Dataset and creates a numpy stream of ready batches."""
   for example in tfds.as_numpy(batches):
     inp, out = example[0][input_name], example[1]
+    if len(out.shape) > 1 and out.shape[-1] == 1:
+      out = np.squeeze(out, axis=-1)
     yield inp, out
 
 
@@ -145,6 +154,14 @@ def _make_summary_writer(output_path):
   return jaxboard.SummaryWriter(output_path)
 
 
+# Metrics to calculate and report.
+_metrics = {
+    "accuracy": accuracy,
+    "neg_log_perplexity": neg_log_perplexity,
+    "loss": lambda x, y: - neg_log_perplexity(x, y),
+}
+
+
 # We include in gin config everything that could be useful to share between
 # users, so when it gets saved in a .gin file it can be re-ran with few flags.
 @gin.configurable(blacklist=["data_dir", "output_dir"])
@@ -190,21 +207,25 @@ def update(i, opt_state, batch):
       step = loaded_step
 
     # Create summary writers.
-    eval_sw = _make_summary_writer(os.path.join(output_dir, "eval_log"))
-    train_sw = _make_summary_writer(os.path.join(output_dir, "train_log"))
+    eval_sw = _make_summary_writer(os.path.join(output_dir, "eval"))
+    train_sw = _make_summary_writer(os.path.join(output_dir, "train"))
 
   log("Starting training.")
   opt_state = opt_init(init_params)
   gin_config_saved = False
+  cur_eval_frequency = 1  # First evaluation after the first training step.
   while step < train_steps:
     # Training.
     start_time = time.time()
-    for _ in range(eval_frequency):
+    for _ in range(cur_eval_frequency):
       opt_state = update(step, opt_state, next(train_stream))
+      if train_sw and step % 10 == 0:  # Log learning rate curve each 10 steps.
+        train_sw.scalar("training/learning rate",
+                        learning_rate(step), step=step)
       step += 1
     epoch_time = time.time() - start_time
     log("Step {}, last {} steps in {:0.2f} sec".format(
-        step, eval_frequency, epoch_time))
+        step, cur_eval_frequency, epoch_time))
 
     # Save the model.
     params = optimizers.get_params(opt_state)
@@ -221,22 +242,39 @@ def update(i, opt_state, batch):
     # Evaluation.
     eval_stream = dataset_to_stream(eval_batches, input_name)
     eval_train_stream = dataset_to_stream(train_batches, input_name)
-    train_acc, eval_acc, train_loss, eval_loss = 0.0, 0.0, 0.0, 0.0
+    train_metrics = {key: 0.0 for key in _metrics}
+    eval_metrics = {key: 0.0 for key in _metrics}
     for _ in range(eval_steps):
-      train_acc += accuracy(params, next(eval_train_stream), model_predict)
-      eval_acc += accuracy(params, next(eval_stream), model_predict)
-      train_loss += loss(params, next(eval_train_stream), model_predict)
-      eval_loss += loss(params, next(eval_stream), model_predict)
-    train_acc /= eval_steps
-    eval_acc /= eval_steps
-    train_loss /= eval_steps
-    eval_loss /= eval_steps
-    log("Train accuracy {:0.4f} loss {:0.8f}".format(train_acc, train_loss))
+      train_batch = next(eval_train_stream)
+      train_predictions = model_predict(params, train_batch[0])
+      eval_batch = next(eval_stream)
+      eval_predictions = model_predict(params, eval_batch[0])
+      for m in _metrics:
+        train_metrics[m] += _metrics[m](
+            train_batch, train_predictions) / float(eval_steps)
+        eval_metrics[m] += _metrics[m](
+            eval_batch, eval_predictions) / float(eval_steps)
+
+    for m in _metrics:
+      log("Step %d train %s %.8f" % (step, m, train_metrics[m]))
+      prefix = "metrics/"
+      if train_sw:
+        train_sw.scalar(prefix + m, train_metrics[m], step=step)
+      log("Step %d eval  %s %.8f" % (step, m, eval_metrics[m]))
+      if eval_sw:
+        eval_sw.scalar(prefix + m, eval_metrics[m], step=step)
+
+    # Log non-metric reports and flush.
     if train_sw:
-      train_sw.scalar("steps/s", epoch_time / eval_frequency, step=step)
-      train_sw.scalar("accuracy", train_acc, step=step)
-      train_sw.scalar("loss", train_loss, step=step)
-    log("Eval  accuracy {:0.4f} loss {:0.8f}".format(eval_acc, eval_loss))
+      if step > 1:  # Don't log performance of the first step.
+        train_sw.scalar("training/steps per second",
+                        cur_eval_frequency / epoch_time, step=step)
+      train_sw.writer.flush()
     if eval_sw:
-      eval_sw.scalar("accuracy", eval_acc, step=step)
-      train_sw.scalar("loss", eval_loss, step=step)
+      eval_sw.writer.flush()
+
+    # After the first step, Evaluate every eval_frequency steps.
+    if cur_eval_frequency == 1 and eval_frequency != 1:
+      cur_eval_frequency = eval_frequency - 1
+    else:
+      cur_eval_frequency = eval_frequency

From 1852cca03271345ecd795d37a8d64608034cc63b Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 25 Feb 2019 21:51:29 -0800
Subject: [PATCH 1718/2720] Remove Modality classes.

Originally, users override default feature transformations by setting hparams.modality. Now, users override transformations by setting hparams.bottom, hparams.loss, hparams.top, and hparams.weights_fn.

See changes to common_hparams.py, modalities.py, modality.py. Other changes are maintenance.

PiperOrigin-RevId: 235659008
---
 .../data_generators/multi_problem.py          |    8 +-
 tensor2tensor/layers/common_hparams.py        |   31 +-
 .../layers/common_image_attention.py          |   35 +-
 .../layers/common_image_attention_test.py     |    1 +
 tensor2tensor/layers/modalities.py            | 2476 +++++++++--------
 tensor2tensor/layers/modalities_test.py       |   26 +-
 tensor2tensor/models/image_transformer.py     |   28 +-
 tensor2tensor/models/image_transformer_2d.py  |    5 +-
 tensor2tensor/models/mtf_transformer.py       |    9 +-
 tensor2tensor/models/mtf_transformer2.py      |    9 +-
 tensor2tensor/models/research/autoencoders.py |   18 +-
 tensor2tensor/models/research/cycle_gan.py    |    9 +-
 tensor2tensor/models/research/super_lm.py     |    8 +-
 .../models/research/transformer_symshard.py   |    9 +-
 .../models/research/transformer_vae.py        |    3 +-
 tensor2tensor/models/video/base.py            |   24 +-
 .../video/basic_deterministic_params.py       |   15 +-
 tensor2tensor/models/video/epva_params.py     |   12 +-
 tensor2tensor/models/video/next_frame_glow.py |   12 +-
 tensor2tensor/models/video/savp_params.py     |   17 +-
 tensor2tensor/models/video/sv2p_params.py     |   19 +-
 tensor2tensor/utils/metrics.py                |    8 +-
 tensor2tensor/utils/modality.py               |  121 -
 tensor2tensor/utils/t2t_model.py              |   52 +-
 tensor2tensor/utils/t2t_model_test.py         |    2 +
 25 files changed, 1521 insertions(+), 1436 deletions(-)
 delete mode 100644 tensor2tensor/utils/modality.py

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 13ec711ae..fe534f050 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -438,8 +438,8 @@ def aggregate_task_losses(hparams,
     vocab_size += (-vocab_size) % hparams.vocab_divisor
   modality = problem_hparams.modality[feature_name]
   loss = hparams.loss.get(feature_name, modalities.get_loss(modality))
-  weights_fn = hparams.targets_weights_fn.get(
-      feature_name, modalities.get_targets_weights_fn(modality))
+  weights_fn = hparams.weights_fn.get(
+      feature_name, modalities.get_weights_fn(modality))
   # Primary task loss
   loss_num, loss_den = loss(
       logits, feature,
@@ -534,8 +534,8 @@ def aggregate_task_lm_losses(hparams,
     vocab_size += (-vocab_size) % hparams.vocab_divisor
   modality = problem_hparams.modality[feature_name]
   loss = hparams.loss.get(feature_name, modalities.get_loss(modality))
-  weights_fn = hparams.targets_weights_fn.get(
-      feature_name, modalities.get_targets_weights_fn(modality))
+  weights_fn = hparams.weights_fn.get(
+      feature_name, modalities.get_weights_fn(modality))
   loss_num = 0.
   loss_den = 0.
   for task in hparams.problem.task_list:
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 654ee4829..8d3f23eaf 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -141,7 +141,6 @@ def basic_params1():
       norm_type="layer",  # "batch", layer", "noam", "none".
       # epsilon parameter to normalization function
       norm_epsilon=1e-6,
-      symbol_modality_num_shards=1,
       # pad vocabularies so that this value divides the vocabulary size.
       vocab_divisor=1,
       # During training, we drop sequences whose inputs and targets are shorter
@@ -176,20 +175,26 @@ def basic_params1():
       # If True, run the model autoregressively instead of teacher-forcing
       # during eval
       eval_run_autoregressive=False,
-      # TODO(lukaszkaiser): these parameters should probably be set elsewhere.
-      # (SymbolModality) - If this flag is on, we try to share all of the input
-      # embeddings, the target embeddings and the softmax weights.
+      # (For features with symbol modality) If True, share all of the
+      # input embeddings, target embeddings, and softmax weights.
       shared_embedding_and_softmax_weights=False,
-      # (SymbolModality) - If this flag is on, we try to share the input
-      # embeddings and the target embeddings.
-      # You can also share the input embeddings with the target embeddings
-      # by using a problem_hparams that uses the same modality object for
-      # the input modality and target modality.
+      # (For features with symbol modality) If True, share the input embeddings
+      # and target embeddings.
       shared_embedding=False,
-      # Modalities used to map from features to a space compatible with
-      # chosen model architecture. It comprises key-value pairs of a feature
-      # name (str) and its modality type.
-      modality={},
+      # (For features with symbol modality) Number to shard embeddings by.
+      symbol_modality_num_shards=1,
+      # Feature transformations are optional dictionaries comprising key-value
+      # pairs of a feature name (str) and its transformation (function). If not
+      # specified, T2TModel applies a default transformation according to the
+      # feature's modality. Bottom is applicable to all features; loss, top, and
+      # weights_fn are only applicable to target features.
+      # TODO(trandustin): `name` is an optional hparam for legacy reasons,
+      # defining variable scope names. Remove this hparam in the future.
+      bottom={},
+      loss={},
+      name={},
+      top={},
+      weights_fn={},
       # The maximum length of "input" sequence.
       # Sequences longer than this value will be truncated. 0 or negative values
       # mean there is no maximum or truncation.
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 4f5ea7b67..ddc0ecd9d 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -603,30 +603,9 @@ def prepare_decoder(targets, hparams):
 
 def prepare_image(inputs, hparams, name=None):
   """Prepare image."""
-  inputs_shape = common_layers.shape_list(inputs)
-  batch = inputs_shape[0]
-  orig_rows = inputs_shape[1]
-  orig_cols = inputs_shape[2]
-  channels = hparams.num_channels
-
-  hidden_size = hparams.hidden_size
-  # TODO(trandustin): Check via modalities.ModalityType.IDENTITY and not str.
-  # The current implementation is to avoid circular imports, modalities ->
-  # discretization -> common_image_attention -> modalities.
-  if "targets" in hparams.modality:
-    target_modality_name = hparams.modality["targets"]
-    if not isinstance(target_modality_name, str):
-      target_modality_name = target_modality_name.__name__
-  else:
-    target_modality_name = None
-  if target_modality_name == "IdentityModality":
-    inputs = tf.to_int32(inputs)
-    x = get_channel_embeddings(channels, inputs, hidden_size, name=name)
-  else:
-    x = inputs
-  x = tf.reshape(x, [batch, orig_rows, orig_cols * channels, hidden_size])
-
-  return x
+  # TODO(trandustin): This is a legacy function. Remove its usage.
+  del hparams, name  # unused arg
+  return inputs
 
 
 def create_output(decoder_output, rows, cols, targets, hparams):
@@ -647,17 +626,19 @@ def create_output(decoder_output, rows, cols, targets, hparams):
     [batch, hparams.img_len, hparams.img_len, hparams.num_channels, 256].
     In the special case of predict mode, it is a Tensor of rank 5.
   """
+  del targets  # unused arg
   decoded_image = postprocess_image(decoder_output, rows, cols, hparams)
+  batch = common_layers.shape_list(decoded_image)[0]
   depth = common_layers.shape_list(decoded_image)[-1]
-  batch, height, width, channels = common_layers.shape_list(targets)
   likelihood = getattr(hparams, "likelihood", DistributionType.CAT)
   if hparams.mode == tf.estimator.ModeKeys.PREDICT:
     y = tf.reshape(decoded_image, [batch, -1, 1, 1, depth])
-    output = y[:, :height, :, :, :]
+    output = y[:, :rows, :, :, :]
   elif likelihood == DistributionType.CAT:
     # Unpack the cols dimension of the Categorical.
+    channels = hparams.num_channels
     output = tf.reshape(decoded_image,
-                        [batch, height, width, channels, depth])
+                        [batch, rows, cols // channels, channels, depth])
   else:
     output = decoded_image
   return output
diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index 1d22927b6..fae6806fa 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -94,6 +94,7 @@ def testCreateOutputTrainMode(self, likelihood, num_mixtures, depth):
     hparams = HParams(
         hidden_size=2,
         likelihood=likelihood,
+        num_channels=channels,
         mode=tf.estimator.ModeKeys.TRAIN,
         num_mixtures=num_mixtures,
     )
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 9a87aee1b..a8fbb29f6 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -13,20 +13,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Modalities define the bottom and top of the model (not the body)."""
+"""Modalities, which specify a feature's domain.
+
+T2TModel applies a default transformation to each feature according to its
+modality. Override them by specifying a model's
+hparams.{bottom,loss,name,top,weights_fn}.
+"""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import sys
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_audio
+from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 from tensor2tensor.layers import discretization
-from tensor2tensor.utils import modality
 
 import tensorflow as tf
 import tensorflow_probability as tfp
@@ -52,1266 +56,1440 @@ def is_pointwise(func):
   return func
 
 
-class SymbolModality(modality.Modality):
-  """Modality for sets of discrete symbols.
-
-  Input:
-    Embedding.
-
-  Output:
-    Linear transformation + softmax.
-  """
-
-  @staticmethod
-  def name(model_hparams, vocab_size):
-    return "symbol_modality_%d_%d" % (vocab_size, model_hparams.hidden_size)
-
-  targets_weights_fn = staticmethod(common_layers.weights_nonzero)
+def generic_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  """Compute loss numerator and denominator for one shard of output."""
+  del vocab_size  # unused arg
+  logits = top_out
+  logits = common_attention.maybe_upcast(logits, hparams=model_hparams)
+  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.0)
+  return common_layers.padded_cross_entropy(
+      logits,
+      targets,
+      model_hparams.label_smoothing,
+      cutoff=cutoff,
+      weights_fn=weights_fn)
 
-  @staticmethod
-  def _get_weights(model_hparams, vocab_size, hidden_dim=None):
-    """Create or get concatenated embedding or softmax variable.
-
-    Args:
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-      hidden_dim: dim of the variable. Defaults to _model_hparams' hidden_size
-
-    Returns:
-       a list of num_shards Tensors.
-    """
-    if hidden_dim is None:
-      hidden_dim = model_hparams.hidden_size
-    num_shards = model_hparams.symbol_modality_num_shards
-    shards = []
-    for i in range(num_shards):
-      shard_size = (vocab_size // num_shards) + (
-          1 if i < vocab_size % num_shards else 0)
-      var_name = "weights_%d" % i
-      shards.append(
-          tf.get_variable(
-              var_name, [shard_size, hidden_dim],
-              initializer=tf.random_normal_initializer(0.0, hidden_dim**-0.5)))
-    if num_shards == 1:
-      ret = shards[0]
-    else:
-      ret = tf.concat(shards, 0)
-    # Convert ret to tensor.
-    if not tf.executing_eagerly():
-      ret = common_layers.convert_gradient_to_tensor(ret)
-    return ret
-
-  @classmethod
-  def bottom_simple(cls, x, model_hparams, vocab_size, name, reuse):
-    with tf.variable_scope(name, reuse=reuse):
-      # Ensure the inputs are 3-D
-      if len(x.get_shape()) == 4:
-        x = tf.squeeze(x, axis=3)
-      while len(x.get_shape()) < 3:
-        x = tf.expand_dims(x, axis=-1)
-
-      var = cls._get_weights(model_hparams, vocab_size)
-      x = common_layers.dropout_no_scaling(
-          x, 1.0 - model_hparams.symbol_dropout)
-      ret = common_layers.gather(var, x)
-      if model_hparams.multiply_embedding_mode == "sqrt_depth":
-        ret *= model_hparams.hidden_size**0.5
-      ret *= tf.expand_dims(
-          common_layers.cast_like(tf.not_equal(x, 0), ret), -1)
-      return ret
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    if (model_hparams.shared_embedding_and_softmax_weights or
-        model_hparams.get("shared_embedding")):
-      return cls.bottom_simple(
-          x, model_hparams, vocab_size, "shared", reuse=None)
-    return cls.bottom_simple(
-        x, model_hparams, vocab_size, "input_emb", reuse=None)
-
-  @classmethod
-  def targets_bottom(cls, x, model_hparams, vocab_size):
-    if (model_hparams.shared_embedding_and_softmax_weights or
-        model_hparams.get("shared_embedding")):
-      try:
-        return cls.bottom_simple(
-            x, model_hparams, vocab_size, "shared", reuse=True)
-      except ValueError:
-        # perhaps there were no inputs, and this is a new variable.
-        return cls.bottom_simple(
-            x, model_hparams, vocab_size, "shared", reuse=None)
-    else:
-      return cls.bottom_simple(
-          x, model_hparams, vocab_size, "target_emb", reuse=None)
-
-  @classmethod
-  @is_pointwise
-  def top(cls, body_output, targets, model_hparams, vocab_size):
-    """Generate logits.
-
-    Args:
-      body_output: A Tensor with shape
-        [batch, p0, p1, model_hparams.hidden_size].
-      targets: Unused.
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].
-    """
-    del targets  # unused arg
-    if model_hparams.shared_embedding_and_softmax_weights:
-      scope_name = "shared"
-      reuse = tf.AUTO_REUSE
-    else:
-      scope_name = "softmax"
-      reuse = False
-    with tf.variable_scope(scope_name, reuse=reuse):
-      body_output_shape = common_layers.shape_list(body_output)
-      var = cls._get_weights(model_hparams, vocab_size, body_output_shape[-1])
-      if (model_hparams.factored_logits and
-          model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
-        # insert channels dimension
-        body_output = tf.expand_dims(body_output, 3)
-        return common_layers.FactoredTensor(body_output, var)
-      else:
-        body_output = tf.reshape(body_output, [-1, body_output_shape[-1]])
-        logits = tf.matmul(body_output, var, transpose_b=True)
-        return tf.reshape(logits,
-                          body_output_shape[:-1] + [1, vocab_size])
 
-
-class SymbolModalityWeightsAll(SymbolModality):
-  """SymbolModality for features that do not have 0-padding."""
-
-  targets_weights_fn = staticmethod(common_layers.weights_all)
-
-
-class SymbolModalityOneHot(SymbolModality):
-  """Simple SymbolModality with one hot as embeddings."""
-
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size):
-    return tf.one_hot(x, vocab_size)
-
-  @staticmethod
+def make_targets_bottom(bottom):
   def targets_bottom(x, model_hparams, vocab_size):
-    return tf.one_hot(x, vocab_size)
-
-  @staticmethod
-  @is_pointwise
-  def top(body_output, _, model_hparams, vocab_size):
-    return body_output
+    with tf.variable_scope("targets_bottom"):
+      return bottom(x, model_hparams, vocab_size)
+  return targets_bottom
 
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    del weights_fn  # unused arg
-    labels = tf.one_hot(targets, vocab_size)
-    loss = tf.nn.softmax_cross_entropy_with_logits(
-        logits=top_out, labels=labels)
-    return tf.reduce_mean(loss), tf.constant(1.0)
-
-
-class CTCSymbolModality(SymbolModality):
-  """SymbolModality that uses CTC loss."""
-
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Compute the CTC loss."""
-    logits = top_out
-    with tf.name_scope("ctc_loss", values=[logits, targets]):
-      # For CTC we assume targets are 1d, [batch, length, 1, 1] here.
-      targets_shape = targets.get_shape().as_list()
-      assert len(targets_shape) == 4
-      assert targets_shape[2] == 1
-      assert targets_shape[3] == 1
-      targets = tf.squeeze(targets, axis=[2, 3])
-      logits = tf.squeeze(logits, axis=[2, 3])
-      targets_mask = 1 - tf.to_int32(tf.equal(targets, 0))
-      targets_lengths = tf.reduce_sum(targets_mask, axis=1)
-      sparse_targets = tf.keras.backend.ctc_label_dense_to_sparse(
-          targets, targets_lengths)
-      xent = tf.nn.ctc_loss(
-          sparse_targets,
-          logits,
-          targets_lengths,
-          time_major=False,
-          preprocess_collapse_repeated=False,
-          ctc_merge_repeated=False)
-      weights = weights_fn(targets)
-      return tf.reduce_sum(xent), tf.reduce_sum(weights)
-
-
-class ImageModality(modality.Modality):
-  """Modality for images."""
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      if not tf.executing_eagerly():
-        tf.summary.image(
-            "inputs", common_layers.tpu_safe_image_summary(x), max_outputs=2)
-      return tf.to_float(x)
-
-  @classmethod
-  def targets_bottom(cls, x, model_hparams, vocab_size):
-    pixel_embedding_size = 64
-    inputs = x
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      if not tf.executing_eagerly():
-        tf.summary.image(
-            "targets_bottom",
-            common_layers.tpu_safe_image_summary(inputs),
-            max_outputs=1)
-      inputs_shape = common_layers.shape_list(inputs)
-      if len(inputs_shape) != 4:
-        raise ValueError("Assuming images given as int tensors in the format "
-                         "[batch, height, width, channels] (256 values).")
-      # We embed each of 256=vocab_size possible pixel values.
-      embedding_var = tf.get_variable(
-          "pixel_embedding",
-          [vocab_size, pixel_embedding_size])
-      hot_inputs = tf.one_hot(tf.to_int32(inputs), vocab_size)
-      hot_inputs = tf.reshape(hot_inputs, [-1, vocab_size])
-      embedded = tf.matmul(hot_inputs, embedding_var)
-      # Let's now merge all channels that were embedded into a single vector.
-      merged_size = pixel_embedding_size * inputs_shape[3]
-      embedded = tf.reshape(embedded, inputs_shape[:3] + [merged_size])
-      merged = tf.layers.dense(
-          embedded,
-          model_hparams.hidden_size,
-          name="merge_pixel_embedded_channels")
-      return merged
-
-  @staticmethod
-  def top(body_output, _, model_hparams, vocab_size):
-    # TODO(lukaszkaiser): is this a universal enough way to get channels?
-    num_channels = model_hparams.problem.num_channels
-    with tf.variable_scope("rgb_softmax"):
-      body_output_shape = common_layers.shape_list(body_output)
-      reshape_shape = body_output_shape[:3]
-      reshape_shape.extend([num_channels, vocab_size])
-      res = tf.layers.dense(body_output, vocab_size * num_channels)
-      res = tf.reshape(res, reshape_shape)
-      if not tf.get_variable_scope().reuse:
-        res_argmax = tf.argmax(res, axis=-1)
-        tf.summary.image(
-            "result",
-            common_layers.tpu_safe_image_summary(res_argmax),
-            max_outputs=1)
-      return res
-
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Compute loss numerator and denominator for one shard of output."""
-    logits = top_out
-    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.0)
-    return common_layers.padded_cross_entropy(
-        logits,
-        targets,
-        model_hparams.label_smoothing,
-        cutoff=cutoff,
-        weights_fn=weights_fn)
 
+def get_weights(model_hparams, vocab_size, hidden_dim=None):
+  """Create or get concatenated embedding or softmax variable.
 
-class ImageChannelCompressModality(modality.Modality):
-  """Modality for images using channel compression for generation."""
-
-  @staticmethod
-  def bottom_compress(inputs, model_hparams, name="bottom"):
-    """Compresses channel-wise input pixels into whole pixel representions.
-
-    Perform conversion of RGB pixel values to a real number in the range -1 to
-    1. This combines pixel channels to form a representation of shape
-    [img_len, img_len].
-
-    Args:
-      inputs: Tensor representing RGB pixel intensities as integers, of shape
-        [batch, img_len, img_len, channels].
-      model_hparams: tf.HParams, model hyperparmeters.
-      name: string, scope.
-
-    Returns:
-      body_input: Tensor of shape
-        [batch, img_len, img_len, model_hparams.hidden_size].
-    """
-    num_channels = 3
-    with tf.variable_scope(name):
-      inputs = tf.to_float(inputs)
-      hp = model_hparams
-      if hp.mode != tf.estimator.ModeKeys.PREDICT:
-        tf.summary.image(
-            "inputs",
-            common_layers.tpu_safe_image_summary(inputs),
-            max_outputs=2)
-      inputs = common_layers.convert_rgb_to_symmetric_real(inputs)
-
-      # Reshape inputs to apply convolutions across [img_len, img_len*channels].
-      inputs_shape = common_layers.shape_list(inputs)
-      inputs = tf.reshape(
-          inputs, [-1, inputs_shape[1], inputs_shape[2] * inputs_shape[3], 1])
-
-      # Compress RGB intensities for each pixel using a convolution.
-      outputs = tf.layers.conv2d(
-          inputs,
-          model_hparams.hidden_size,
-          kernel_size=(1, num_channels),
-          padding="VALID",
-          strides=(1, num_channels),
-          activation=tf.nn.relu,
-          name="conv_input")
-      return outputs
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    return cls.bottom_compress(x, model_hparams, "input_bottom")
-
-  @classmethod
-  def targets_bottom(cls, x, model_hparams, vocab_size):
-    return cls.bottom_compress(x, model_hparams, "output_bottom")
-
-  @classmethod
-  def top(cls, body_output, _, model_hparams, vocab_size):
-    """Transforms body output to return logits.
-
-    Args:
-      body_output: Tensor of shape [batch, img_len, img_len, depth].
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      Tensor of shape [batch, img_len, img_len, channels, vocab_size].
-    """
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      hidden_size = model_hparams.hidden_size
-      img_len = model_hparams.img_len
-      channels = 3  # RGB
-      batch = common_layers.shape_list(body_output)[0]
-      x = tf.layers.conv2d(
-          body_output,
-          hidden_size * channels,
-          kernel_size=(1, 1),
-          strides=(1, 1),
-          padding="VALID",
-          activation=tf.nn.relu,
-          name="decompress_conv")
-      x = tf.reshape(x, [batch, img_len, img_len * channels, hidden_size])
-      x = common_layers.layer_preprocess(x, model_hparams)
-      x = tf.layers.dense(x,
-                          vocab_size,
-                          use_bias=True,
-                          activation=None,
-                          name="output_conv")
-      x = tf.reshape(
-          x, [batch, img_len, img_len, channels, vocab_size])
-      return x
-
-
-class ImageChannelBottomIdentityModality(ImageChannelCompressModality):
-
-  @staticmethod
-  def top(body_output, _, model_hparams, vocab_size):
-    return body_output
-
-
-class ImageChannelEmbeddingsBottom(modality.Modality):
-  """Modality for images using channel compression for generation."""
+  Args:
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
+    hidden_dim: dim of the variable. Defaults to _model_hparams' hidden_size
 
-  @staticmethod
-  def get_channel_embeddings(io_depth,
-                             targets,
-                             hidden_size,
-                             name="channel"):
-    """Get separate embedding for each of the channels."""
-    targets_split = tf.split(targets, io_depth, axis=3)
-    rgb_embedding_var = tf.get_variable("rgb_target_emb_%s" % name,
-                                        [256 * io_depth, hidden_size])
-    rgb_embedding_var = tf.identity(rgb_embedding_var)
-    rgb_embedding_var *= float(hidden_size)**0.5
-    channel_target_embs = []
-    for i in range(io_depth):
-      # Adding the channel offsets to get the right embedding since the
-      # embedding tensor has shape 256 * io_depth, hidden_size
-      target_ids = tf.squeeze(targets_split[i], axis=3) + i * 256
-      target_embs = common_layers.gather(rgb_embedding_var, target_ids)
-      channel_target_embs.append(target_embs)
-
-    return tf.concat(channel_target_embs, axis=-1)
+  Returns:
+     a list of num_shards Tensors.
+  """
+  if hidden_dim is None:
+    hidden_dim = model_hparams.hidden_size
+  num_shards = model_hparams.symbol_modality_num_shards
+  shards = []
+  for i in range(num_shards):
+    shard_size = (vocab_size // num_shards) + (
+        1 if i < vocab_size % num_shards else 0)
+    var_name = "weights_%d" % i
+    shards.append(
+        tf.get_variable(
+            var_name, [shard_size, hidden_dim],
+            initializer=tf.random_normal_initializer(0.0, hidden_dim**-0.5)))
+  if num_shards == 1:
+    ret = shards[0]
+  else:
+    ret = tf.concat(shards, 0)
+  # Convert ret to tensor.
+  if not tf.executing_eagerly():
+    ret = common_layers.convert_gradient_to_tensor(ret)
+  return ret
+
+
+def _symbol_bottom_simple(x, model_hparams, vocab_size, name, reuse):
+  """Bottom transformation for symbols."""
+  with tf.variable_scope(name, reuse=reuse):
+    # Ensure the inputs are 3-D
+    if len(x.get_shape()) == 4:
+      x = tf.squeeze(x, axis=3)
+    while len(x.get_shape()) < 3:
+      x = tf.expand_dims(x, axis=-1)
+
+    var = get_weights(model_hparams, vocab_size)
+    x = common_layers.dropout_no_scaling(
+        x, 1.0 - model_hparams.symbol_dropout)
+    ret = common_layers.gather(var, x)
+    if model_hparams.multiply_embedding_mode == "sqrt_depth":
+      ret *= model_hparams.hidden_size**0.5
+    ret *= tf.expand_dims(
+        common_layers.cast_like(tf.not_equal(x, 0), ret), -1)
+    return ret
 
-  @staticmethod
-  def targets_bottom(x, model_hparams, vocab_size):
-    inputs = x
-    io_depth = model_hparams.num_channels
-    tshape = common_layers.shape_list(inputs)
-    hidden_size = model_hparams.hidden_size
-    target_embeddings = ImageChannelEmbeddingsBottom.get_channel_embeddings(
-        io_depth, inputs, hidden_size, "input_bottom")
-    return tf.reshape(target_embeddings,
-                      [tshape[0], tshape[1], tshape[2] * io_depth, hidden_size])
-
-  @classmethod
-  def top(cls, body_output, _, model_hparams, vocab_size):
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      img_len = model_hparams.img_len
-      channels = model_hparams.num_channels
-      x = tf.layers.dense(
-          body_output, 256, use_bias=True, activation=None, name="output_conv")
-      x = tf.reshape(x,
-                     [-1, img_len, img_len, channels, vocab_size])
-      return x
-
-
-class AudioModality(modality.Modality):
-  """Performs strided conv compressions for audio data."""
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    """Transform input from data space to model space.
-
-    Args:
-      x: A Tensor with shape [batch, ...]
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      body_input: A Tensor with shape [batch, ?, ?,
-        model_hparams.hidden_size].
-    """
-    inputs = x
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      # TODO(aidangomez): Will need to sort out a better audio pipeline
-      def xnet_resblock(x, filters, res_relu, name):
-        """Xception block."""
-        with tf.variable_scope(name):
-          # Typically audio samples are >100k samples in length and have a width
-          # of 2 or 4. Mono audio has a single channel while stereo has 2.
-          y = common_layers.separable_conv_block(
-              x,
-              filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
-              first_relu=True,
-              padding="SAME",
-              force2d=True,
-              name="sep_conv_block")
-          y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 2))
-          return y + common_layers.conv_block(
-              x,
-              filters, [((1, 1), (1, 1))],
-              padding="SAME",
-              strides=(2, 2),
-              first_relu=res_relu,
-              force2d=True,
-              name="res_conv0")
-
-      x = tf.to_float(inputs) / 255.
-      x.set_shape([None, None, None, 1])
-      for i in range(model_hparams.audio_compression):
-        x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
-      return xnet_resblock(x,
-                           model_hparams.hidden_size,
-                           False,
-                           "compress_block_final")
-
-
-class AudioSpectralModality(modality.Modality):
-  """Performs strided conv compressions for audio spectral data."""
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    """Transform input from data space to model space.
-
-    Args:
-      x: A Tensor with shape [batch, ...]
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      body_input: A Tensor with shape [batch, ?, ?,
-        model_hparams.hidden_size].
-    """
-    inputs = x
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      # TODO(aidangomez): Will need to sort out a better audio pipeline
-      def xnet_resblock(x, filters, res_relu, name):
-        """Xception-like block."""
-        with tf.variable_scope(name):
-          # We only stride along the length dimension to preserve the spectral
-          # bins (which are tiny in dimensionality relative to length)
-          y = common_layers.separable_conv_block(
-              x,
-              filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
-              first_relu=True,
-              padding="SAME",
-              force2d=True,
-              name="sep_conv_block")
-          y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 1))
-          return y + common_layers.conv_block(
-              x,
-              filters, [((1, 1), (1, 1))],
-              padding="SAME",
-              strides=(2, 1),
-              first_relu=res_relu,
-              force2d=True,
-              name="res_conv0")
-
-      # Bitcast back from int32
-      x = tf.bitcast(inputs, tf.float32)
-      x.set_shape([None, None, None, 1])
-      for i in range(model_hparams.audio_compression):
-        x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
-      return xnet_resblock(x,
-                           model_hparams.hidden_size,
-                           False,
-                           "compress_block_final")
-
-
-class SpeechRecognitionModality(modality.Modality):
-  """Common ASR filterbank processing."""
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    """Use batchnorm instead of CMVN and shorten the stft with strided convs.
-
-    Args:
-      x: float32 tensor with shape [batch_size, len, 1, freqs * channels]
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      float32 tensor with shape [batch_size, shorter_len, 1, hidden_size]
-    """
-    inputs = x
-    p = model_hparams
-
-    num_mel_bins = p.audio_num_mel_bins
-    num_channels = 3 if p.audio_add_delta_deltas else 1
-
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      if p.audio_preproc_in_bottom:
-        # Compute filterbanks
-        with tf.variable_scope("fbanks"):
-          waveforms = tf.squeeze(inputs, [2, 3])
-          mel_fbanks = common_audio.compute_mel_filterbank_features(
-              waveforms,
-              sample_rate=p.audio_sample_rate,
-              dither=p.audio_dither,
-              preemphasis=p.audio_preemphasis,
-              frame_length=p.audio_frame_length,
-              frame_step=p.audio_frame_step,
-              lower_edge_hertz=p.audio_lower_edge_hertz,
-              upper_edge_hertz=p.audio_upper_edge_hertz,
-              num_mel_bins=p.audio_num_mel_bins,
-              apply_mask=True)
-          if p.audio_add_delta_deltas:
-            mel_fbanks = common_audio.add_delta_deltas(mel_fbanks)
-          x = tf.reshape(mel_fbanks,
-                         common_layers.shape_list(mel_fbanks)[:2] +
-                         [num_mel_bins, num_channels])
-
-          nonpadding_mask = 1. - common_attention.embedding_to_padding(x)
-          num_of_nonpadding_elements = tf.reduce_sum(
-              nonpadding_mask) * num_mel_bins * num_channels
-
-          # This replaces CMVN estimation on data
-          var_epsilon = 1e-09
-          mean = tf.reduce_sum(
-              x, axis=[1], keepdims=True) / num_of_nonpadding_elements
-          variance = (num_of_nonpadding_elements * mean**2. -
-                      2. * mean * tf.reduce_sum(x, axis=[1], keepdims=True) +
-                      tf.reduce_sum(x**2, axis=[1], keepdims=True)
-                     ) / num_of_nonpadding_elements
-          x = (x - mean) * tf.rsqrt(variance + var_epsilon) * tf.expand_dims(
-              nonpadding_mask, -1)
-      else:
-        x = inputs
-
-      # The convention is that the models are flattened along the spatial,
-      # dimensions, thus the speech preprocessor treats frequencies and
-      # channels as image colors (last axis)
-      x.set_shape([None, None, num_mel_bins, num_channels])
-
-      # TODO(chorowski): how to specify bottom's hparams and avoid hardcoding?
-      x = tf.pad(x, [[0, 0], [0, 8], [0, 0], [0, 0]])
-      for _ in range(2):
-        x = tf.layers.conv2d(
-            x, 128, (3, 3), (2, 2), use_bias=False)
-        x = common_layers.layer_norm(x)
-        x = tf.nn.relu(x)
-
-      xshape = common_layers.shape_list(x)
-      # apply a conv that will remove all frequencies and at the same time
-      # project the output into desired hidden_size
-      x = tf.pad(x, [[0, 0], [0, 2], [0, 0], [0, 0]])
-      x = tf.layers.conv2d(x, p.hidden_size, (3, xshape[2]), use_bias=False)
-
-      assert common_layers.shape_list(x)[2] == 1
-      x = common_layers.layer_norm(x)
-      x = tf.nn.relu(x)
-    return x
 
+def symbol_bottom(x, model_hparams, vocab_size):
+  if (model_hparams.shared_embedding_and_softmax_weights or
+      model_hparams.get("shared_embedding")):
+    return _symbol_bottom_simple(
+        x, model_hparams, vocab_size, "shared", reuse=None)
+  return _symbol_bottom_simple(
+      x, model_hparams, vocab_size, "input_emb", reuse=None)
+
+
+def symbol_targets_bottom(x, model_hparams, vocab_size):
+  """Bottom transformation for target symbols."""
+  if (model_hparams.shared_embedding_and_softmax_weights or
+      model_hparams.get("shared_embedding")):
+    try:
+      return _symbol_bottom_simple(
+          x, model_hparams, vocab_size, "shared", reuse=True)
+    except ValueError:
+      # perhaps there were no inputs, and this is a new variable.
+      return _symbol_bottom_simple(
+          x, model_hparams, vocab_size, "shared", reuse=None)
+  else:
+    return _symbol_bottom_simple(
+        x, model_hparams, vocab_size, "target_emb", reuse=None)
 
-class VideoModality(modality.Modality):
-  """Modality for videos, i.e., time-sequences of frames."""
 
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size):
-    common_video.gif_summary("inputs", x, max_outputs=1)
-    x = common_layers.standardize_images(x)
-    return x
+@is_pointwise
+def symbol_top(body_output, targets, model_hparams, vocab_size):
+  """Generate logits.
 
-  @staticmethod
-  def targets_bottom(x, model_hparams, vocab_size):
-    common_video.gif_summary("targets", x, max_outputs=1)
-    x = common_layers.standardize_images(x)
-    return x
-
-  @staticmethod
-  def top(body_output, targets, model_hparams, vocab_size):
-    num_channels = model_hparams.problem.num_channels
-    shape = common_layers.shape_list(body_output)
-    reshape_shape = shape[:-1] + [num_channels, vocab_size]
-    res = tf.reshape(body_output, reshape_shape)
-    # Calculate argmax so as to have a summary with the produced images.
-    x = tf.argmax(tf.reshape(res, [-1, vocab_size]), axis=-1)
-    x = tf.reshape(x, shape[:-1] + [num_channels])
-    common_video.gif_summary("results", x, max_outputs=1)
-    return res
+  Args:
+    body_output: A Tensor with shape
+      [batch, p0, p1, model_hparams.hidden_size].
+    targets: Unused.
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
 
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Compute loss numerator and denominator for one shard of output."""
-    logits = top_out
-    logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
-    targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
-    return common_layers.padded_cross_entropy(
+  Returns:
+    logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].
+  """
+  del targets  # unused arg
+  if model_hparams.shared_embedding_and_softmax_weights:
+    scope_name = "shared"
+    reuse = tf.AUTO_REUSE
+  else:
+    scope_name = "softmax"
+    reuse = False
+  with tf.variable_scope(scope_name, reuse=reuse):
+    body_output_shape = common_layers.shape_list(body_output)
+    var = get_weights(model_hparams, vocab_size, body_output_shape[-1])
+    if (model_hparams.factored_logits and
+        model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
+      # insert channels dimension
+      body_output = tf.expand_dims(body_output, 3)
+      return common_layers.FactoredTensor(body_output, var)
+    else:
+      body_output = tf.reshape(body_output, [-1, body_output_shape[-1]])
+      logits = tf.matmul(body_output, var, transpose_b=True)
+      return tf.reshape(logits,
+                        body_output_shape[:-1] + [1, vocab_size])
+
+
+def symbol_one_hot_bottom(x, model_hparams, vocab_size):
+  del model_hparams  # unused arg
+  return tf.one_hot(x, vocab_size)
+
+
+@is_pointwise
+def symbol_one_hot_top(body_output, targets, model_hparams, vocab_size):
+  del targets, model_hparams, vocab_size  # unused arg
+  return body_output
+
+
+def symbol_one_hot_loss(top_out,
+                        targets,
+                        model_hparams,
+                        vocab_size,
+                        weights_fn):
+  del model_hparams, weights_fn  # unused arg
+  labels = tf.one_hot(targets, vocab_size)
+  loss = tf.nn.softmax_cross_entropy_with_logits(
+      logits=top_out, labels=labels)
+  return tf.reduce_mean(loss), tf.constant(1.0)
+
+
+def ctc_symbol_loss(top_out, targets, model_hparams, vocab_size, weight_fn):
+  """Compute the CTC loss."""
+  del model_hparams, vocab_size  # unused arg
+  logits = top_out
+  with tf.name_scope("ctc_loss", values=[logits, targets]):
+    # For CTC we assume targets are 1d, [batch, length, 1, 1] here.
+    targets_shape = targets.get_shape().as_list()
+    assert len(targets_shape) == 4
+    assert targets_shape[2] == 1
+    assert targets_shape[3] == 1
+    targets = tf.squeeze(targets, axis=[2, 3])
+    logits = tf.squeeze(logits, axis=[2, 3])
+    targets_mask = 1 - tf.to_int32(tf.equal(targets, 0))
+    targets_lengths = tf.reduce_sum(targets_mask, axis=1)
+    sparse_targets = tf.keras.backend.ctc_label_dense_to_sparse(
+        targets, targets_lengths)
+    xent = tf.nn.ctc_loss(
+        sparse_targets,
         logits,
-        targets,
-        model_hparams.label_smoothing,
-        cutoff=cutoff,
-        weights_fn=weights_fn)
-
-
-class VideoModalityBitwise(VideoModality):
-  """Video Modality where bottom embeds pixels bitwise."""
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    pixel_embedding_size = 64
-    inputs = x
-    with tf.variable_scope(cls.name(model_hparams, vocab_size),
-                           reuse=tf.AUTO_REUSE):
-      common_layers.summarize_video(inputs, "bottom")
-      # Embed bitwise.
-      assert vocab_size == 256
-      embedded = discretization.int_to_bit_embed(inputs, 8,
-                                                 pixel_embedding_size)
-      # Project.
-      return tf.layers.dense(
-          embedded,
-          model_hparams.hidden_size,
-          name="merge_pixel_embedded_frames")
-
-  @classmethod
-  def targets_bottom(cls, x, model_hparams, vocab_size):  # pylint: disable=arguments-differ
-    pixel_embedding_size = 64
-    inputs = x
-    with tf.variable_scope(cls.name(model_hparams, vocab_size),
-                           reuse=tf.AUTO_REUSE):
-      common_layers.summarize_video(inputs, "targets_bottom")
-      # Embed bitwise.
-      assert vocab_size == 256
-      embedded = discretization.int_to_bit_embed(inputs, 8,
-                                                 pixel_embedding_size)
-      # Transpose and project.
-      transposed = common_layers.time_to_channels(embedded)
-      return tf.layers.dense(
-          transposed,
-          model_hparams.hidden_size,
-          name="merge_pixel_embedded_frames")
-
-
-class VideoModalityPixelNoise(VideoModality):
-  """Video modality that introduces pixel noise on input during training."""
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    input_noise = getattr(model_hparams, "video_modality_input_noise", 0.25)
-    inputs = x
-    if model_hparams.mode == tf.estimator.ModeKeys.TRAIN:
-      background = tfp.distributions.percentile(inputs, 50., axis=[0, 1, 2, 3])
-      input_shape = common_layers.shape_list(inputs)
-      input_size = tf.reduce_prod(input_shape[:-1])
-      input_mask = tf.multinomial(
-          tf.log([[input_noise, 1.-input_noise]]), input_size)
-      input_mask = tf.reshape(tf.cast(input_mask, tf.int32),
-                              input_shape[:-1]+[1])
-      inputs = inputs * input_mask + background * (1 - input_mask)
-    return super(VideoModalityPixelNoise, cls).bottom(
-        inputs, model_hparams, vocab_size)
-
-
-class VideoModalityL1(VideoModality):
-  """Video modality that predicts a scalar per channel with an L1 loss."""
-
-  @staticmethod
-  def top(body_output, _, model_hparams, vocab_size):
-    num_channels = model_hparams.problem.num_channels
-    num_frames = model_hparams.video_num_target_frames
-    with tf.variable_scope("rgb"):
-      body_output_shape = common_layers.shape_list(body_output)
-      res = tf.layers.dense(body_output, num_channels * num_frames, name="cast")
-      res = tf.reshape(res, body_output_shape[:3] + [num_channels, num_frames])
-      res = tf.transpose(res, [0, 4, 1, 2, 3])  # Move frames next to batch.
-      if not tf.get_variable_scope().reuse:
-        res_argmax = res[:, -1, :, :, :]
-        tf.summary.image(
-            "result",
-            common_layers.tpu_safe_image_summary(res_argmax),
-            max_outputs=1)
-      return tf.expand_dims(res, axis=-1)  # Add an axis like in perplexity.
-
-  @staticmethod
-  def internal_loss(logits, targets, model_hparams):
-    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
-    return tf.nn.relu(tf.abs(logits - targets) - cutoff)
-
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Compute loss numerator and denominator for one shard of output."""
-    logits = top_out
-    logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:-1])
-    targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-    weights = weights_fn(targets)
-    # Shift targets by 0.5 so later just casting to int gives the prediction.
-    # So for int targets, say 0 and 7, we actually train to predict 0.5 and 7.5.
-    # Later (in merics or infer) this is cast to int anyway. Also, we have no
-    # loss beyond cutoff = 0.2 as these are already correct predictions.
-    targets = tf.to_float(targets) + 0.5
-    loss = VideoModalityL1.internal_loss(logits, targets, model_hparams)
-    return tf.reduce_sum(loss * weights), tf.reduce_sum(weights)
-
-
-class VideoModalityL2(VideoModalityL1):
-  """Modality for videos with L2 loss."""
-
-  @staticmethod
-  def internal_loss(logits, targets, model_hparams):
-    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
-    return tf.nn.relu(
-        tf.squared_difference(logits, targets) - cutoff * cutoff)
-
-
-class VideoModalityL2Raw(VideoModalityL2):
-  """Modality with L2 loss and raw input (sequences of frames)."""
+        targets_lengths,
+        time_major=False,
+        preprocess_collapse_repeated=False,
+        ctc_merge_repeated=False)
+    weights = weight_fn(targets)
+    return tf.reduce_sum(xent), tf.reduce_sum(weights)
 
-  @staticmethod
-  def convert_rgb_to_real(prediction, targets):
-    """Convert prediction and target from rgb to real."""
-    prediction = tf.squeeze(prediction, axis=-1)
-    prediction = common_layers.convert_rgb_to_real(prediction)
-    targets = common_layers.convert_rgb_to_real(targets)
-    return prediction, targets
-
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size):
-    common_video.gif_summary("inputs", x)
-    return common_layers.convert_rgb_to_real(x)
 
-  @staticmethod
-  def targets_bottom(x, model_hparams, vocab_size):  # pylint: disable=arguments-differ
-    common_video.gif_summary("targets_bottom", x)
-    return common_layers.convert_rgb_to_real(x)
-
-  @staticmethod
-  def top(body_output, _, model_hparams, vocab_size):
-    frames = body_output
-    if isinstance(body_output, list):
-      frames = tf.stack(body_output, axis=1)
-    rgb_frames = common_layers.convert_real_to_rgb(frames)
-    common_video.gif_summary("body_output", rgb_frames)
-    return tf.expand_dims(rgb_frames, axis=-1)
+def image_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  with tf.variable_scope("image_modality"):
+    if not tf.executing_eagerly():
+      tf.summary.image(
+          "inputs", common_layers.tpu_safe_image_summary(x), max_outputs=2)
+    return tf.to_float(x)
 
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
-    del weights_fn  # unused arg
-    prediction, groundtruth = cls.convert_rgb_to_real(top_out, targets)
-    loss = tf.losses.mean_squared_error(prediction, groundtruth)
-    return loss, tf.constant(1.0)
 
+def image_targets_bottom(x, model_hparams, vocab_size):
+  """Bottom transformation for target images."""
+  pixel_embedding_size = 64
+  inputs = x
+  with tf.variable_scope("image_modality"):
+    if not tf.executing_eagerly():
+      tf.summary.image(
+          "targets_bottom",
+          common_layers.tpu_safe_image_summary(inputs),
+          max_outputs=1)
+    inputs_shape = common_layers.shape_list(inputs)
+    if len(inputs_shape) != 4:
+      raise ValueError("Assuming images given as int tensors in the format "
+                       "[batch, height, width, channels] (256 values).")
+    # We embed each of 256=vocab_size possible pixel values.
+    embedding_var = tf.get_variable(
+        "pixel_embedding",
+        [vocab_size, pixel_embedding_size])
+    hot_inputs = tf.one_hot(tf.to_int32(inputs), vocab_size)
+    hot_inputs = tf.reshape(hot_inputs, [-1, vocab_size])
+    embedded = tf.matmul(hot_inputs, embedding_var)
+    # Let's now merge all channels that were embedded into a single vector.
+    merged_size = pixel_embedding_size * inputs_shape[3]
+    embedded = tf.reshape(embedded, inputs_shape[:3] + [merged_size])
+    merged = tf.layers.dense(
+        embedded,
+        model_hparams.hidden_size,
+        name="merge_pixel_embedded_channels")
+    return merged
+
+
+def image_top(body_output, targets, model_hparams, vocab_size):
+  """Top transformation for images."""
+  del targets  # unused arg
+  # TODO(lukaszkaiser): is this a universal enough way to get channels?
+  num_channels = model_hparams.problem.num_channels
+  with tf.variable_scope("rgb_softmax"):
+    body_output_shape = common_layers.shape_list(body_output)
+    reshape_shape = body_output_shape[:3]
+    reshape_shape.extend([num_channels, vocab_size])
+    res = tf.layers.dense(body_output, vocab_size * num_channels)
+    res = tf.reshape(res, reshape_shape)
+    if not tf.get_variable_scope().reuse:
+      res_argmax = tf.argmax(res, axis=-1)
+      tf.summary.image(
+          "result",
+          common_layers.tpu_safe_image_summary(res_argmax),
+          max_outputs=1)
+    return res
 
-class VideoModalityL1Raw(VideoModalityL2Raw):
-  """Modality with L1 loss and raw input (sequences of frames)."""
 
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
-    prediction, groundtruth = cls.convert_rgb_to_real(top_out, targets)
-    loss = tf.losses.absolute_difference(prediction, groundtruth)
-    return loss, tf.constant(1.0)
+def _image_channel_compress_bottom(inputs, model_hparams, name="bottom"):
+  """Compresses channel-wise input pixels into whole pixel representions.
 
+  Perform conversion of RGB pixel values to a real number in the range -1 to
+  1. This combines pixel channels to form a representation of shape
+  [img_len, img_len].
 
-class ClassLabelModality(modality.Modality):
-  """Used for label data."""
+  Args:
+    inputs: Tensor representing RGB pixel intensities as integers, of shape
+      [batch, img_len, img_len, channels].
+    model_hparams: tf.HParams, model hyperparmeters.
+    name: string, scope.
 
-  @staticmethod
-  def name(model_hparams, vocab_size):
-    return "class_label_modality_%d_%d" % (vocab_size,
-                                           model_hparams.hidden_size)
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      multiplier = 1.0
-      if model_hparams.multiply_embedding_mode == "sqrt_depth":
-        multiplier = model_hparams.hidden_size**0.5
-      return common_layers.embedding(x,
-                                     vocab_size,
-                                     model_hparams.hidden_size,
-                                     multiplier=multiplier)
-
-  @classmethod
-  def targets_bottom(cls, x, model_hparams, vocab_size):
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      return tf.zeros([common_layers.shape_list(x)[0],
-                       1,
-                       1,
-                       model_hparams.hidden_size])
-
-  @classmethod
-  def top(cls, body_output, _, model_hparams, vocab_size):
-    """Transform inputs from model space to target space.
-
-    Average over inner dims and a linear layer to logits.
-
-    Args:
-      body_output: A Tensor with shape [batch, ?, ?, body_output_size].
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      a Tensors, each with shape [batch_size, 1, 1, 1, vocab_size]
-    """
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      x = body_output
-      x = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
-      res = tf.layers.dense(x, vocab_size)
-      return tf.expand_dims(res, 3)
-
-
-class VideoModalityIdentity(VideoModality):
-  """Video Modality where top and bottom is an identity function."""
+  Returns:
+    body_input: Tensor of shape
+      [batch, img_len, img_len, model_hparams.hidden_size].
+  """
+  num_channels = 3
+  with tf.variable_scope(name):
+    inputs = tf.to_float(inputs)
+    hp = model_hparams
+    if hp.mode != tf.estimator.ModeKeys.PREDICT:
+      tf.summary.image(
+          "inputs",
+          common_layers.tpu_safe_image_summary(inputs),
+          max_outputs=2)
+    inputs = common_layers.convert_rgb_to_symmetric_real(inputs)
+
+    # Reshape inputs to apply convolutions across [img_len, img_len*channels].
+    inputs_shape = common_layers.shape_list(inputs)
+    inputs = tf.reshape(
+        inputs, [-1, inputs_shape[1], inputs_shape[2] * inputs_shape[3], 1])
+
+    # Compress RGB intensities for each pixel using a convolution.
+    outputs = tf.layers.conv2d(
+        inputs,
+        model_hparams.hidden_size,
+        kernel_size=(1, num_channels),
+        padding="VALID",
+        strides=(1, num_channels),
+        activation=tf.nn.relu,
+        name="conv_input")
+    return outputs
+
+
+def image_channel_compress_bottom(x, model_hparams, vocab_size):
+  del vocab_size  # unused arg
+  return _image_channel_compress_bottom(x, model_hparams, "input_bottom")
+
+
+def image_channel_compress_targets_bottom(x, model_hparams, vocab_size):
+  del vocab_size  # unused arg
+  return _image_channel_compress_bottom(x, model_hparams, "output_bottom")
+
+
+def image_channel_compress_top(body_output, targets, model_hparams, vocab_size):
+  """Transforms body output to return logits.
 
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size):
-    common_video.gif_summary("inputs", x, max_outputs=1)
-    return x
+  Args:
+    body_output: Tensor of shape [batch, img_len, img_len, depth].
+    targets:
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
 
-  @staticmethod
-  def targets_bottom(x, model_hparams, vocab_size):
-    common_video.gif_summary("targets", x, max_outputs=1)
+  Returns:
+    Tensor of shape [batch, img_len, img_len, channels, vocab_size].
+  """
+  del targets  # unused arg
+  with tf.variable_scope("image_channel_compress_modality"):
+    hidden_size = model_hparams.hidden_size
+    img_len = model_hparams.img_len
+    channels = 3  # RGB
+    batch = common_layers.shape_list(body_output)[0]
+    x = tf.layers.conv2d(
+        body_output,
+        hidden_size * channels,
+        kernel_size=(1, 1),
+        strides=(1, 1),
+        padding="VALID",
+        activation=tf.nn.relu,
+        name="decompress_conv")
+    x = tf.reshape(x, [batch, img_len, img_len * channels, hidden_size])
+    x = common_layers.layer_preprocess(x, model_hparams)
+    x = tf.layers.dense(x,
+                        vocab_size,
+                        use_bias=True,
+                        activation=None,
+                        name="output_conv")
+    x = tf.reshape(
+        x, [batch, img_len, img_len, channels, vocab_size])
     return x
 
-  @staticmethod
-  def top(body_output, targets, model_hparams, vocab_size):
-    return body_output
-
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Compute loss numerator and denominator for one shard of output."""
-    # TODO(nikip): Try L2 loss
-    logits = top_out
-    logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
-    targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
-    return common_layers.padded_cross_entropy(
-        logits,
-        targets,
-        model_hparams.label_smoothing,
-        cutoff=cutoff,
-        weights_fn=weights_fn)
-
-
-class MultiLabelModality(ClassLabelModality):
-  """Used for multi label task."""
-
-  targets_weights_fn = staticmethod(common_layers.weights_nonzero)
-
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Average loss over the labels."""
-    logits = top_out
-    num_labels = tf.shape(targets)[1]
-    logits = tf.tile(logits, [1, num_labels, 1, 1, 1])
-
-    xent, weights = common_layers.padded_cross_entropy(
-        logits,
-        targets,
-        model_hparams.label_smoothing,
-        weights_fn=weights_fn,
-        reduce_sum=False,
-    )
-    xent = tf.squeeze(xent, [2, 3])
-    weights = tf.squeeze(weights, [2, 3])
-    # average loss over all labels
-    loss = tf.reduce_sum(xent, axis=1)
-    weights = tf.reduce_sum(weights, axis=1)
-    loss /= (weights + 1e-8)
-    weights = tf.to_float(tf.greater(weights, 0.))
-
-    return tf.reduce_sum(loss*weights), tf.reduce_sum(weights)
-
 
-class OneHotClassLabelModality(ClassLabelModality):
-  """Used for one-hot encoded class labels."""
-
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Apply softmax cross-entropy between outputs and targets.
-
-    Args:
-      top_out: logits Tensor with shape [batch, ?, ?, num_classes]
-      targets: one-hot encoding Tensor with shape [batch, ?, ?, num_classes]
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-      weights_fn: Function mapping targets to weights.
-
-    Returns:
-      loss_scale (cross-entropy), loss_denom
-    """
-    loss_scale = tf.losses.softmax_cross_entropy(
-        onehot_labels=targets, logits=top_out)
-    weights = weights_fn(targets)
-    loss_denom = tf.reduce_sum(weights)
-    return loss_scale, loss_denom
-
-
-class IdentityModality(modality.Modality):
-  """Does nothing."""
-
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size):
-    return tf.to_float(x)
-
-  @staticmethod
-  def top(body_output, _, model_hparams, vocab_size):
-    return body_output
-
-
-class GenericL2LossModality(IdentityModality):
-  """Generic modality with L2 as Loss."""
-
-  @staticmethod
-  def targets_bottom(x, model_hparams, vocab_size):
-    return tf.to_float(x)
+def image_channel_embeddings_bottom(x, model_hparams, vocab_size):
+  """Bottom transformation for image targets."""
+  del vocab_size  # unused arg
+  inputs = tf.to_int32(x)
+  io_depth = model_hparams.num_channels
+  tshape = common_layers.shape_list(inputs)
+  hidden_size = model_hparams.hidden_size
+  target_embeddings = cia.get_channel_embeddings(
+      io_depth, inputs, hidden_size, "input_bottom")
+  return tf.reshape(target_embeddings,
+                    [tshape[0], tshape[1], tshape[2] * io_depth, hidden_size])
+
+
+def image_channel_embeddings_top(body_output,
+                                 targets,
+                                 model_hparams,
+                                 vocab_size):
+  """Top transformation for images."""
+  del targets  # unused arg
+  with tf.variable_scope("image_channel_embeddings_bottom"):
+    img_len = model_hparams.img_len
+    channels = model_hparams.num_channels
+    x = tf.layers.dense(
+        body_output, 256, use_bias=True, activation=None, name="output_conv")
+    x = tf.reshape(x,
+                   [-1, img_len, img_len, channels, vocab_size])
+    return x
 
-  @staticmethod
-  def loss(body_output, targets, model_hparams, vocab_size, weights_fn):
-    del weights_fn  # unused
-    loss = tf.squared_difference(body_output, tf.to_float(targets))
-    return tf.reduce_mean(loss), tf.constant(1.0)
 
+def audio_bottom(x, model_hparams, vocab_size):
+  """Transform input from data space to model space.
 
-class RealModality(modality.Modality):
-  """Base class for real (i.e. float) vectors.
+  Args:
+    x: A Tensor with shape [batch, ...]
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
 
-  * Bottom is a linear projection layer to hparams.hidden_size.
-  * Top is a linear projection layer to vocab_size.
+  Returns:
+    body_input: A Tensor with shape [batch, ?, ?,
+      model_hparams.hidden_size].
   """
+  del vocab_size  # unused arg
+  inputs = x
+  with tf.variable_scope("audio_modality"):
+    # TODO(aidangomez): Will need to sort out a better audio pipeline
+    def xnet_resblock(x, filters, res_relu, name):
+      """Xception block."""
+      with tf.variable_scope(name):
+        # Typically audio samples are >100k samples in length and have a width
+        # of 2 or 4. Mono audio has a single channel while stereo has 2.
+        y = common_layers.separable_conv_block(
+            x,
+            filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
+            first_relu=True,
+            padding="SAME",
+            force2d=True,
+            name="sep_conv_block")
+        y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 2))
+        return y + common_layers.conv_block(
+            x,
+            filters, [((1, 1), (1, 1))],
+            padding="SAME",
+            strides=(2, 2),
+            first_relu=res_relu,
+            force2d=True,
+            name="res_conv0")
+
+    x = tf.to_float(inputs) / 255.
+    x.set_shape([None, None, None, 1])
+    for i in range(model_hparams.audio_compression):
+      x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
+    return xnet_resblock(x,
+                         model_hparams.hidden_size,
+                         False,
+                         "compress_block_final")
+
+
+def audio_spectral_bottom(x, model_hparams, vocab_size):
+  """Transform input from data space to model space.
 
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size):
-    with tf.variable_scope("real"):
-      return tf.layers.dense(
-          tf.to_float(x), model_hparams.hidden_size, name="bottom")
-
-  @staticmethod
-  @is_pointwise
-  def top(body_output, _, model_hparams, vocab_size):
-    with tf.variable_scope("real"):
-      return tf.layers.dense(body_output, vocab_size, name="top")
-
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    raise NotImplementedError()
+  Args:
+    x: A Tensor with shape [batch, ...]
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
 
+  Returns:
+    body_input: A Tensor with shape [batch, ?, ?,
+      model_hparams.hidden_size].
+  """
+  del vocab_size  # unused arg
+  inputs = x
+  with tf.variable_scope("audio_spectral_modality"):
+    # TODO(aidangomez): Will need to sort out a better audio pipeline
+    def xnet_resblock(x, filters, res_relu, name):
+      """Xception-like block."""
+      with tf.variable_scope(name):
+        # We only stride along the length dimension to preserve the spectral
+        # bins (which are tiny in dimensionality relative to length)
+        y = common_layers.separable_conv_block(
+            x,
+            filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
+            first_relu=True,
+            padding="SAME",
+            force2d=True,
+            name="sep_conv_block")
+        y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 1))
+        return y + common_layers.conv_block(
+            x,
+            filters, [((1, 1), (1, 1))],
+            padding="SAME",
+            strides=(2, 1),
+            first_relu=res_relu,
+            force2d=True,
+            name="res_conv0")
+
+    # Bitcast back from int32
+    x = tf.bitcast(inputs, tf.float32)
+    x.set_shape([None, None, None, 1])
+    for i in range(model_hparams.audio_compression):
+      x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
+    return xnet_resblock(x,
+                         model_hparams.hidden_size,
+                         False,
+                         "compress_block_final")
+
+
+def speech_recognition_bottom(x, model_hparams, vocab_size):
+  """Use batchnorm instead of CMVN and shorten the stft with strided convs.
 
-class RealL2LossModality(RealModality):
-  """Modality for real (i.e. float) vectors with L2 (Gaussian) loss."""
+  Args:
+    x: float32 tensor with shape [batch_size, len, 1, freqs * channels]
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
 
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    predictions = top_out
-    if (len(common_layers.shape_list(top_out)) != len(
-        common_layers.shape_list(targets))):
-      predictions = tf.squeeze(top_out, axis=[-1])
-    with tf.name_scope("l2"):
-      weights = weights_fn(targets)
-      l2 = tf.pow(predictions - targets, 2)
-      return tf.reduce_sum(l2 * weights), tf.reduce_sum(weights)
+  Returns:
+    float32 tensor with shape [batch_size, shorter_len, 1, hidden_size]
+  """
+  del vocab_size  # unused arg
+  inputs = x
+  p = model_hparams
+
+  num_mel_bins = p.audio_num_mel_bins
+  num_channels = 3 if p.audio_add_delta_deltas else 1
+
+  with tf.variable_scope("speech_recognition_modality"):
+    if p.audio_preproc_in_bottom:
+      # Compute filterbanks
+      with tf.variable_scope("fbanks"):
+        waveforms = tf.squeeze(inputs, [2, 3])
+        mel_fbanks = common_audio.compute_mel_filterbank_features(
+            waveforms,
+            sample_rate=p.audio_sample_rate,
+            dither=p.audio_dither,
+            preemphasis=p.audio_preemphasis,
+            frame_length=p.audio_frame_length,
+            frame_step=p.audio_frame_step,
+            lower_edge_hertz=p.audio_lower_edge_hertz,
+            upper_edge_hertz=p.audio_upper_edge_hertz,
+            num_mel_bins=p.audio_num_mel_bins,
+            apply_mask=True)
+        if p.audio_add_delta_deltas:
+          mel_fbanks = common_audio.add_delta_deltas(mel_fbanks)
+        x = tf.reshape(mel_fbanks,
+                       common_layers.shape_list(mel_fbanks)[:2] +
+                       [num_mel_bins, num_channels])
+
+        nonpadding_mask = 1. - common_attention.embedding_to_padding(x)
+        num_of_nonpadding_elements = tf.reduce_sum(
+            nonpadding_mask) * num_mel_bins * num_channels
+
+        # This replaces CMVN estimation on data
+        var_epsilon = 1e-09
+        mean = tf.reduce_sum(
+            x, axis=[1], keepdims=True) / num_of_nonpadding_elements
+        variance = (num_of_nonpadding_elements * mean**2. -
+                    2. * mean * tf.reduce_sum(x, axis=[1], keepdims=True) +
+                    tf.reduce_sum(x**2, axis=[1], keepdims=True)
+                   ) / num_of_nonpadding_elements
+        x = (x - mean) * tf.rsqrt(variance + var_epsilon) * tf.expand_dims(
+            nonpadding_mask, -1)
+    else:
+      x = inputs
 
+    # The convention is that the models are flattened along the spatial,
+    # dimensions, thus the speech preprocessor treats frequencies and
+    # channels as image colors (last axis)
+    x.set_shape([None, None, num_mel_bins, num_channels])
 
-class RealLogPoissonLossModality(RealModality):
-  """Modality for real (i.e. float) vectors with log Poisson regression loss."""
+    # TODO(chorowski): how to specify bottom's hparams and avoid hardcoding?
+    x = tf.pad(x, [[0, 0], [0, 8], [0, 0], [0, 0]])
+    for _ in range(2):
+      x = tf.layers.conv2d(
+          x, 128, (3, 3), (2, 2), use_bias=False)
+      x = common_layers.layer_norm(x)
+      x = tf.nn.relu(x)
+
+    xshape = common_layers.shape_list(x)
+    # apply a conv that will remove all frequencies and at the same time
+    # project the output into desired hidden_size
+    x = tf.pad(x, [[0, 0], [0, 2], [0, 0], [0, 0]])
+    x = tf.layers.conv2d(x, p.hidden_size, (3, xshape[2]), use_bias=False)
+
+    assert common_layers.shape_list(x)[2] == 1
+    x = common_layers.layer_norm(x)
+    x = tf.nn.relu(x)
+  return x
+
+
+def video_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("inputs", x, max_outputs=1)
+  x = common_layers.standardize_images(x)
+  return x
+
+
+def video_targets_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("targets", x, max_outputs=1)
+  x = common_layers.standardize_images(x)
+  return x
+
+
+def video_top(body_output, targets, model_hparams, vocab_size):
+  """Top transformation for video."""
+  del targets  # unused arg
+  num_channels = model_hparams.problem.num_channels
+  shape = common_layers.shape_list(body_output)
+  reshape_shape = shape[:-1] + [num_channels, vocab_size]
+  res = tf.reshape(body_output, reshape_shape)
+  # Calculate argmax so as to have a summary with the produced images.
+  x = tf.argmax(tf.reshape(res, [-1, vocab_size]), axis=-1)
+  x = tf.reshape(x, shape[:-1] + [num_channels])
+  common_video.gif_summary("results", x, max_outputs=1)
+  return res
+
+
+def video_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  """Compute loss numerator and denominator for one shard of output."""
+  del vocab_size  # unused arg
+  logits = top_out
+  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
+  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
+  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
+  return common_layers.padded_cross_entropy(
+      logits,
+      targets,
+      model_hparams.label_smoothing,
+      cutoff=cutoff,
+      weights_fn=weights_fn)
+
+
+def video_bitwise_bottom(x, model_hparams, vocab_size):
+  """Bottom transformation for embedding video bitwise."""
+  pixel_embedding_size = 64
+  inputs = x
+  with tf.variable_scope("video_modality_bitwise", reuse=tf.AUTO_REUSE):
+    common_layers.summarize_video(inputs, "bottom")
+    # Embed bitwise.
+    assert vocab_size == 256
+    embedded = discretization.int_to_bit_embed(inputs, 8,
+                                               pixel_embedding_size)
+    # Project.
+    return tf.layers.dense(
+        embedded,
+        model_hparams.hidden_size,
+        name="merge_pixel_embedded_frames")
+
+
+def video_bitwise_targets_bottom(x, model_hparams, vocab_size):
+  """Bottom transformation for embedding target video bitwise."""
+  pixel_embedding_size = 64
+  inputs = x
+  with tf.variable_scope("video_modality_bitwise", reuse=tf.AUTO_REUSE):
+    common_layers.summarize_video(inputs, "targets_bottom")
+    # Embed bitwise.
+    assert vocab_size == 256
+    embedded = discretization.int_to_bit_embed(inputs, 8,
+                                               pixel_embedding_size)
+    # Transpose and project.
+    transposed = common_layers.time_to_channels(embedded)
+    return tf.layers.dense(
+        transposed,
+        model_hparams.hidden_size,
+        name="merge_pixel_embedded_frames")
+
+
+def video_pixel_noise_bottom(x, model_hparams, vocab_size):
+  """Bottom transformation for video."""
+  input_noise = getattr(model_hparams, "video_modality_input_noise", 0.25)
+  inputs = x
+  if model_hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    background = tfp.distributions.percentile(inputs, 50., axis=[0, 1, 2, 3])
+    input_shape = common_layers.shape_list(inputs)
+    input_size = tf.reduce_prod(input_shape[:-1])
+    input_mask = tf.multinomial(
+        tf.log([[input_noise, 1.-input_noise]]), input_size)
+    input_mask = tf.reshape(tf.cast(input_mask, tf.int32),
+                            input_shape[:-1]+[1])
+    inputs = inputs * input_mask + background * (1 - input_mask)
+  return video_bottom(inputs, model_hparams, vocab_size)
+
+
+def video_l1_top(body_output, targets, model_hparams, vocab_size):
+  """Top transformation for video."""
+  del targets, vocab_size  # unused arg
+  num_channels = model_hparams.problem.num_channels
+  num_frames = model_hparams.video_num_target_frames
+  with tf.variable_scope("rgb"):
+    body_output_shape = common_layers.shape_list(body_output)
+    res = tf.layers.dense(body_output, num_channels * num_frames, name="cast")
+    res = tf.reshape(res, body_output_shape[:3] + [num_channels, num_frames])
+    res = tf.transpose(res, [0, 4, 1, 2, 3])  # Move frames next to batch.
+    if not tf.get_variable_scope().reuse:
+      res_argmax = res[:, -1, :, :, :]
+      tf.summary.image(
+          "result",
+          common_layers.tpu_safe_image_summary(res_argmax),
+          max_outputs=1)
+    return tf.expand_dims(res, axis=-1)  # Add an axis like in perplexity.
+
+
+def video_l1_internal_loss(logits, targets, model_hparams):
+  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
+  return tf.nn.relu(tf.abs(logits - targets) - cutoff)
+
+
+def video_l1_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  """Compute loss numerator and denominator for one shard of output."""
+  del vocab_size  # unused arg
+  logits = top_out
+  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:-1])
+  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
+  weights = weights_fn(targets)
+  # Shift targets by 0.5 so later just casting to int gives the prediction.
+  # So for int targets, say 0 and 7, we actually train to predict 0.5 and 7.5.
+  # Later (in merics or infer) this is cast to int anyway. Also, we have no
+  # loss beyond cutoff = 0.2 as these are already correct predictions.
+  targets = tf.to_float(targets) + 0.5
+  loss = video_l1_internal_loss(logits, targets, model_hparams)
+  return tf.reduce_sum(loss * weights), tf.reduce_sum(weights)
+
+
+def video_l2_internal_loss(logits, targets, model_hparams):
+  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
+  return tf.nn.relu(
+      tf.squared_difference(logits, targets) - cutoff * cutoff)
+
+
+def video_l2_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  """Compute loss numerator and denominator for one shard of output."""
+  del vocab_size  # unused arg
+  logits = top_out
+  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:-1])
+  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
+  weights = weights_fn(targets)
+  # Shift targets by 0.5 so later just casting to int gives the prediction.
+  # So for int targets, say 0 and 7, we actually train to predict 0.5 and 7.5.
+  # Later (in merics or infer) this is cast to int anyway. Also, we have no
+  # loss beyond cutoff = 0.2 as these are already correct predictions.
+  targets = tf.to_float(targets) + 0.5
+  loss = video_l2_internal_loss(logits, targets, model_hparams)
+  return tf.reduce_sum(loss * weights), tf.reduce_sum(weights)
+
+
+def convert_rgb_to_real(prediction, targets):
+  """Convert prediction and target from rgb to real."""
+  prediction = tf.squeeze(prediction, axis=-1)
+  prediction = common_layers.convert_rgb_to_real(prediction)
+  targets = common_layers.convert_rgb_to_real(targets)
+  return prediction, targets
+
+
+def video_raw_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("inputs", x)
+  return common_layers.convert_rgb_to_real(x)
+
+
+def video_raw_targets_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("targets_bottom", x)
+  return common_layers.convert_rgb_to_real(x)
+
+
+def video_raw_top(body_output, targets, model_hparams, vocab_size):
+  del targets, model_hparams, vocab_size  # unused arg
+  frames = body_output
+  if isinstance(body_output, list):
+    frames = tf.stack(body_output, axis=1)
+  rgb_frames = common_layers.convert_real_to_rgb(frames)
+  common_video.gif_summary("body_output", rgb_frames)
+  return tf.expand_dims(rgb_frames, axis=-1)
+
+
+def video_l2_raw_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  del model_hparams, vocab_size, weights_fn  # unused arg
+  prediction, groundtruth = convert_rgb_to_real(top_out, targets)
+  loss = tf.losses.mean_squared_error(prediction, groundtruth)
+  return loss, tf.constant(1.0)
+
+
+def video_l1_raw_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  del model_hparams, vocab_size, weights_fn  # unused arg
+  prediction, groundtruth = convert_rgb_to_real(top_out, targets)
+  loss = tf.losses.absolute_difference(prediction, groundtruth)
+  return loss, tf.constant(1.0)
+
+
+def class_label_bottom(x, model_hparams, vocab_size):
+  with tf.variable_scope("class_label_modality_%d_%d" % (
+      vocab_size, model_hparams.hidden_size)):
+    multiplier = 1.0
+    if model_hparams.multiply_embedding_mode == "sqrt_depth":
+      multiplier = model_hparams.hidden_size**0.5
+    return common_layers.embedding(x,
+                                   vocab_size,
+                                   model_hparams.hidden_size,
+                                   multiplier=multiplier)
 
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    predictions = top_out
-    if (len(common_layers.shape_list(top_out)) != len(
-        common_layers.shape_list(targets))):
-      predictions = tf.squeeze(top_out, axis=[-1])
-    with tf.name_scope("log_possion"):
-      weights = weights_fn(targets)
-      lp_loss = tf.nn.log_poisson_loss(targets, predictions)
-      return tf.reduce_sum(lp_loss * weights), tf.reduce_sum(weights)
 
+def class_label_targets_bottom(x, model_hparams, vocab_size):
+  with tf.variable_scope("class_label_modality_%d_%d" % (
+      vocab_size, model_hparams.hidden_size)):
+    return tf.zeros([common_layers.shape_list(x)[0],
+                     1,
+                     1,
+                     model_hparams.hidden_size])
+
+
+def class_label_top(body_output, targets, model_hparams, vocab_size):
+  """Transform inputs from model space to target space.
+
+  Average over inner dims and a linear layer to logits.
 
-class IdentitySymbolModality(SymbolModality):
-  """Symbol modality with identity top and bottom transformations.
+  Args:
+    body_output: A Tensor with shape [batch, ?, ?, body_output_size].
+    targets:
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
 
-  Uses the weights_fn from SymbolModality so that loss/metrics ignore padding.
+  Returns:
+    a Tensors, each with shape [batch_size, 1, 1, 1, vocab_size]
   """
+  del targets  # unused arg
+  with tf.variable_scope("class_label_modality_%d_%d" % (
+      vocab_size, model_hparams.hidden_size)):
+    x = body_output
+    x = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
+    res = tf.layers.dense(x, vocab_size)
+    return tf.expand_dims(res, 3)
+
+
+def video_identity_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("inputs", x, max_outputs=1)
+  return x
+
+
+def video_identity_targets_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("targets", x, max_outputs=1)
+  return x
+
+
+def video_identity_loss(top_out,
+                        targets,
+                        model_hparams,
+                        vocab_size,
+                        weights_fn):
+  """Compute loss numerator and denominator for one shard of output."""
+  del vocab_size  # unused arg
+  # TODO(nikip): Try L2 loss
+  logits = top_out
+  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
+  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
+  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
+  return common_layers.padded_cross_entropy(
+      logits,
+      targets,
+      model_hparams.label_smoothing,
+      cutoff=cutoff,
+      weights_fn=weights_fn)
+
+
+def multi_label_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  """Average loss over the labels."""
+  del vocab_size  # unused arg
+  logits = top_out
+  num_labels = tf.shape(targets)[1]
+  logits = tf.tile(logits, [1, num_labels, 1, 1, 1])
+
+  xent, weights = common_layers.padded_cross_entropy(
+      logits,
+      targets,
+      model_hparams.label_smoothing,
+      weights_fn=weights_fn,
+      reduce_sum=False,
+  )
+  xent = tf.squeeze(xent, [2, 3])
+  weights = tf.squeeze(weights, [2, 3])
+  # average loss over all labels
+  loss = tf.reduce_sum(xent, axis=1)
+  weights = tf.reduce_sum(weights, axis=1)
+  loss /= (weights + 1e-8)
+  weights = tf.to_float(tf.greater(weights, 0.))
+
+  return tf.reduce_sum(loss*weights), tf.reduce_sum(weights)
+
+
+def one_hot_class_label_loss(top_out,
+                             targets,
+                             model_hparams,
+                             vocab_size,
+                             weights_fn):
+  """Apply softmax cross-entropy between outputs and targets.
 
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size):
-    return tf.to_float(x)
-
-  @staticmethod
-  def top(body_output, _, model_hparams, vocab_size):
-    return body_output
+  Args:
+    top_out: logits Tensor with shape [batch, ?, ?, num_classes]
+    targets: one-hot encoding Tensor with shape [batch, ?, ?, num_classes]
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
+    weights_fn:
 
-  @classmethod
-  def targets_bottom(cls, x, model_hparams, vocab_size):
-    """SymbolModality overrides targets_bottom, so need to override here too."""
-    return cls.bottom(x, model_hparams, vocab_size)
+  Returns:
+    loss_scale (cross-entropy), loss_denom
+  """
+  del model_hparams, vocab_size  # unused arg
+  loss_scale = tf.losses.softmax_cross_entropy(
+      onehot_labels=targets, logits=top_out)
+  weights = weights_fn(targets)
+  loss_denom = tf.reduce_sum(weights)
+  return loss_scale, loss_denom
 
 
-class SigmoidClassLabelModality(ClassLabelModality):
-  """Sigmoid cross-entropy for independent class labels."""
+def identity_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  return tf.to_float(x)
 
-  @staticmethod
-  def name(model_hparams, vocab_size):
-    return "sigmoid_class_symbol_modality_%d_%d" % (vocab_size,
-                                                    model_hparams.hidden_size)
 
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    # Expect inputs of size [batch-size, timesteps, 1, num-classes], where the
-    # last dimension of num-classes represents logits for binary labels
-    loss_scale = tf.losses.sigmoid_cross_entropy(
-        multi_class_labels=targets, logits=top_out)
-    weights = weights_fn(targets)
-    loss_denom = tf.reduce_sum(weights)
-    return loss_scale, loss_denom
+def identity_top(body_output, targets, model_hparams, vocab_size):
+  del targets, model_hparams, vocab_size  # unused arg
+  return body_output
 
 
-class SigmoidMaxPoolingClassLabelModality(ClassLabelModality):
-  """Sigmoid cross-entropy applied on max-pooling over timesteps."""
+def generic_l2_loss(body_output,
+                    targets,
+                    model_hparams,
+                    vocab_size,
+                    weights_fn):
+  del model_hparams, vocab_size, weights_fn  # unused arg
+  loss = tf.squared_difference(body_output, tf.to_float(targets))
+  return tf.reduce_mean(loss), tf.constant(1.0)
 
-  @staticmethod
-  def name(model_hparams, vocab_size):
-    return "sigmoid_max_pooling_class_symbol_modality_%d_%d" % (
-        vocab_size, model_hparams.hidden_size)
 
-  @classmethod
-  def top(cls, body_output, _, model_hparams, vocab_size):
-    """Transform inputs from model space to target space.
+def real_bottom(x, model_hparams, vocab_size):
+  del vocab_size  # unused arg
+  with tf.variable_scope("real"):
+    return tf.layers.dense(
+        tf.to_float(x), model_hparams.hidden_size, name="bottom")
 
-    Average over inner dims and a linear layer to logits.
 
-    Args:
-      body_output: A Tensor with shape [batch, timesteps, 1, body_output_size].
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
+@is_pointwise
+def real_top(body_output, targets, model_hparams, vocab_size):
+  del targets, model_hparams  # unused arg
+  with tf.variable_scope("real"):
+    return tf.layers.dense(body_output, vocab_size, name="top")
 
-    Returns:
-      a Tensors, each with shape [batch_size, 1, 1, vocab_size]
-    """
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      x = body_output
-      x = tf.reduce_max(x, axis=1, keepdims=True)
-      return tf.layers.dense(x, vocab_size)
 
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    # Expect inputs of size [batch-size, 1, 1, num-classes], where the
-    # last dimension of num-classes represents logits for binary labels
-    loss_scale = tf.losses.sigmoid_cross_entropy(
-        multi_class_labels=targets, logits=top_out)
+def real_l2_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  del model_hparams, vocab_size  # unused arg
+  predictions = top_out
+  if (len(common_layers.shape_list(top_out)) != len(
+      common_layers.shape_list(targets))):
+    predictions = tf.squeeze(top_out, axis=[-1])
+  with tf.name_scope("l2"):
     weights = weights_fn(targets)
-    loss_denom = tf.reduce_sum(weights)
-    return loss_scale, loss_denom
-
-
-class SoftmaxMaxPoolingClassLabelModality(OneHotClassLabelModality):
-  """Softmax cross-entropy applied on max-pooling over timesteps."""
-
-  @staticmethod
-  def name(model_hparams, vocab_size):
-    return "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
-        vocab_size, model_hparams.hidden_size)
-
-  @classmethod
-  def top(cls, body_output, _, model_hparams, vocab_size):
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      x = body_output
-      x = tf.reduce_max(x, axis=1, keepdims=True)
-      return tf.layers.dense(x, vocab_size)
+    l2 = tf.pow(predictions - targets, 2)
+    return tf.reduce_sum(l2 * weights), tf.reduce_sum(weights)
 
 
-class SoftmaxAveragePoolingClassLabelModality(OneHotClassLabelModality):
-  """Softmax cross-entropy applied on average-pooling over timesteps."""
-
-  @staticmethod
-  def name(model_hparams, vocab_size):
-    return "softmax_average_pooling_onehot_class_label_modality_%d_%d" % (
-        vocab_size, model_hparams.hidden_size)
-
-  @classmethod
-  def top(cls, body_output, _, model_hparams, vocab_size):
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      x = body_output
-      x = tf.reduce_mean(x, axis=1, keepdims=True)
-      return tf.layers.dense(x, vocab_size)
+def real_log_poisson_loss(top_out,
+                          targets,
+                          model_hparams,
+                          vocab_size,
+                          weights_fn):
+  """Poisson loss for real."""
+  del model_hparams, vocab_size  # unused arg
+  predictions = top_out
+  if (len(common_layers.shape_list(top_out)) != len(
+      common_layers.shape_list(targets))):
+    predictions = tf.squeeze(top_out, axis=[-1])
+  with tf.name_scope("log_possion"):
+    weights = weights_fn(targets)
+    lp_loss = tf.nn.log_poisson_loss(targets, predictions)
+    return tf.reduce_sum(lp_loss * weights), tf.reduce_sum(weights)
 
 
-class SoftmaxLastTimestepClassLabelModality(OneHotClassLabelModality):
-  """Softmax cross-entropy applied on last-timestep encoding."""
+def sigmoid_class_label_loss(top_out,
+                             targets,
+                             model_hparams,
+                             vocab_size,
+                             weights_fn):
+  """Loss for class label."""
+  # Expect inputs of size [batch-size, timesteps, 1, num-classes], where the
+  # last dimension of num-classes represents logits for binary labels
+  del model_hparams, vocab_size  # unused arg
+  loss_scale = tf.losses.sigmoid_cross_entropy(
+      multi_class_labels=targets, logits=top_out)
+  weights = weights_fn(targets)
+  loss_denom = tf.reduce_sum(weights)
+  return loss_scale, loss_denom
+
+
+def sigmoid_max_pooling_class_label_top(body_output,
+                                        targets,
+                                        model_hparams,
+                                        vocab_size):
+  """Transform inputs from model space to target space.
+
+  Average over inner dims and a linear layer to logits.
 
-  @staticmethod
-  def name(model_hparams, vocab_size):
-    return "softmax_last_timestep_onehot_class_label_modality_%d_%d" % (
-        vocab_size, model_hparams.hidden_size)
+  Args:
+    body_output: A Tensor with shape [batch, timesteps, 1, body_output_size].
+    targets:
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
 
-  @classmethod
-  def top(cls, body_output, _, model_hparams, vocab_size):
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      x = body_output
-      x = tf.expand_dims(x[:, -1], 1)  # Pick the last timestep
-      return tf.layers.dense(x, vocab_size)
+  Returns:
+    a Tensors, each with shape [batch_size, 1, 1, vocab_size]
+  """
+  del targets  # unused arg
+  with tf.variable_scope(
+      "sigmoid_max_pooling_class_symbol_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)):
+    x = body_output
+    x = tf.reduce_max(x, axis=1, keepdims=True)
+    return tf.layers.dense(x, vocab_size)
+
+
+def sigmoid_max_pooling_class_label_loss(top_out,
+                                         targets,
+                                         model_hparams,
+                                         vocab_size,
+                                         weights_fn):
+  """Loss for class label."""
+  # Expect inputs of size [batch-size, 1, 1, num-classes], where the
+  # last dimension of num-classes represents logits for binary labels
+  del model_hparams, vocab_size  # unused arg
+  loss_scale = tf.losses.sigmoid_cross_entropy(
+      multi_class_labels=targets, logits=top_out)
+  weights = weights_fn(targets)
+  loss_denom = tf.reduce_sum(weights)
+  return loss_scale, loss_denom
+
+
+def softmax_max_pooling_class_label_top(body_output,
+                                        targets,
+                                        model_hparams,
+                                        vocab_size):
+  """Loss for class label."""
+  del targets  # unused arg
+  with tf.variable_scope(
+      "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)):
+    x = body_output
+    x = tf.reduce_max(x, axis=1, keepdims=True)
+    return tf.layers.dense(x, vocab_size)
+
+
+def softmax_average_pooling_class_label_top(body_output,
+                                            targets,
+                                            model_hparams,
+                                            vocab_size):
+  """Loss for class label."""
+  del targets  # unused arg
+  with tf.variable_scope(
+      "softmax_average_pooling_onehot_class_label_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)):
+    x = body_output
+    x = tf.reduce_mean(x, axis=1, keepdims=True)
+    return tf.layers.dense(x, vocab_size)
+
+
+def softmax_last_timestep_class_label_top(body_output,
+                                          targets,
+                                          model_hparams,
+                                          vocab_size):
+  """Loss for class label."""
+  del targets  # unused arg
+  with tf.variable_scope(
+      "softmax_last_timestep_onehot_class_label_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)):
+    x = body_output
+    x = tf.expand_dims(x[:, -1], 1)  # Pick the last timestep
+    return tf.layers.dense(x, vocab_size)
 
 
 class ModalityType(object):
   """Types of modalities."""
 
-  SYMBOL = "SymbolModality"
-  SYMBOL_WEIGHTS_ALL = "SymbolModalityWeightsAll"
-  SYMBOL_ONE_HOT = "SymbolModalityOneHot"
-  CTC_SYMBOL = "CTCSymbolModality"
-  IMAGE = "ImageModality"
-  IMAGE_CHANNEL_COMPRESS = "ImageChannelCompressModality"
-  IMAGE_CHANNEL_BOTTOM_IDENTITY = "ImageChannelBottomIdentityModality"
-  IMAGE_CHANNEL_EMBEDDINGS_BOTTOM = "ImageChannelEmbeddingsBottom"
-  AUDIO = "AudioModality"
-  AUDIO_SPECTRAL = "AudioSpectralModality"
-  SPEECH_RECOGNITION = "SpeechRecognitionModality"
-  VIDEO = "VideoModality"
-  VIDEO_BITWISE = "VideoModalityBitwise"
-  VIDEO_PIXEL_NOISE = "VideoModalityPixelNoise"
-  VIDEO_L1 = "VideoModalityL1"
-  VIDEO_L2 = "VideoModalityL2"
-  VIDEO_L2_RAW = "VideoModalityL2Raw"
-  VIDEO_L1_RAW = "VideoModalityL1Raw"
-  CLASS_LABEL = "ClassLabelModality"
-  VIDEO_IDENTITY = "VideoModalityIdentity"
-  MULTI_LABEL = "MultiLabelModality"
-  ONE_HOT_CLASS_LABEL = "OneHotClassLabelModality"
-  IDENTITY = "IdentityModality"
-  GENERIC_L2_LOSS = "GenericL2LossModality"
-  REAL = "RealModality"
-  REAL_L2_LOSS = "RealL2LossModality"
-  REAL_LOG_POISSON_LOSS = "RealLogPoissonLossModality"
-  IDENTITY_SYMBOL = "IdentitySymbolModality"
-  SIGMOID_CLASS_LABEL = "SigmoidClassLabelModality"
-  SIGMOID_MAX_POOLING_CLASS_LABEL = "SigmoidMaxPoolingClassLabelModality"
-  SOFTMAX_MAX_POOLING_CLASS_LABEL = "SoftmaxMaxPoolingClassLabelModality"
-  SOFTMAX_AVERAGE_POOLING_CLASS_LABEL = "SoftmaxAveragePoolingClassLabelModality"
-  SOFTMAX_LAST_TIMESTEP_CLASS_LABEL = "SoftmaxLastTimestepClassLabelModality"
+  AUDIO = "audio"
+  AUDIO_SPECTRAL = "audio_spectral"
+  CLASS_LABEL = "class_label"
+  CTC_SYMBOL = "ctc_symbol"  # symbol with CTC loss
+  GENERIC_L2_LOSS = "generic_l2"  # identity modality with L2 loss
+  IDENTITY = "identity"  # identity top and bottom
+  IDENTITY_SYMBOL = "identity_symbol"  # symbol with identity top and bottom
+  IMAGE = "image"
+  # images using channel compression for generation
+  IMAGE_CHANNEL_BOTTOM_IDENTITY = "image_channel_bottom_identity"
+  # images using channel compression for generation
+  IMAGE_CHANNEL_COMPRESS = "image_channel_compress"
+  IMAGE_CHANNEL_EMBEDDINGS_BOTTOM = "image_channel_embeddings_bottom"
+  MULTI_LABEL = "multi_label"
+  ONE_HOT_CLASS_LABEL = "one_hot_class_label"
+  REAL = "real"  # real vectors
+  REAL_L2_LOSS = "real_l2"  # real vectors with L2 as loss
+  # real vectors with log Poisson regression loss
+  REAL_LOG_POISSON_LOSS = "real_log_poisson"
+  SIGMOID_CLASS_LABEL = "sigmoid_class_label"  # sigmoid cross-entropy loss
+  # sigmoid cross-entropy applied on max-pooling over timesteps
+  SIGMOID_MAX_POOLING_CLASS_LABEL = "sigmoid_max_pooling_class_label"
+  # softmax cross-entropy applied on average-pooling over timesteps
+  SOFTMAX_AVERAGE_POOLING_CLASS_LABEL = "softmax_average_pooling_class_label"
+  # softmax cross-entropy applied on last-timestep encoding
+  SOFTMAX_LAST_TIMESTEP_CLASS_LABEL = "softmax_last_timestep_class_label"
+  # softmax cross-entropy applied on max-pooling over timesteps
+  SOFTMAX_MAX_POOLING_CLASS_LABEL = "softmax_max_pooling_class_label"
+  SPEECH_RECOGNITION = "speech_recognition"
+  SYMBOL = "symbol"
+  SYMBOL_WEIGHTS_ALL = "symbol_weights_all"  # symbol for features w/o 0-padding
+  SYMBOL_ONE_HOT = "symbol_one_hot"  # symbol with one hot as embeddings
+  VIDEO = "video"
+  VIDEO_BITWISE = "video_bitwise"  # video where bottom embeds pixels bitwise
+  VIDEO_IDENTITY = "video_identity"  # video with identity top and bottom
+  VIDEO_L1 = "video_l1"  # video with L2 loss
+  VIDEO_L2 = "video_l2"  # video with L1 loss
+  # video with L1 loss and raw input (sequences of frames)
+  VIDEO_L1_RAW = "video_l1_raw"
+  # video with L2 loss and raw input (sequences of frames)
+  VIDEO_L2_RAW = "video_l2_raw"
+  # video with pixel noise on input during training
+  VIDEO_PIXEL_NOISE = "video_pixel_noise"
 
   @staticmethod
   def get_choices():
     return [
-        ModalityType.SYMBOL,
-        ModalityType.SYMBOL_WEIGHTS_ALL,
-        ModalityType.SYMBOL_ONE_HOT,
+        ModalityType.AUDIO,
+        ModalityType.AUDIO_SPECTRAL,
+        ModalityType.CLASS_LABEL,
         ModalityType.CTC_SYMBOL,
+        ModalityType.GENERIC_L2_LOSS,
+        ModalityType.IDENTITY,
+        ModalityType.IDENTITY_SYMBOL,
         ModalityType.IMAGE,
-        ModalityType.IMAGE_CHANNEL_COMPRESS,
         ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
+        ModalityType.IMAGE_CHANNEL_COMPRESS,
         ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM,
-        ModalityType.AUDIO,
-        ModalityType.AUDIO_SPECTRAL,
-        ModalityType.SPEECH_RECOGNITION,
-        ModalityType.VIDEO,
-        ModalityType.VIDEO_BITWISE,
-        ModalityType.VIDEO_PIXEL_NOISE,
-        ModalityType.VIDEO_L1,
-        ModalityType.VIDEO_L2,
-        ModalityType.VIDEO_L2_RAW,
-        ModalityType.VIDEO_L1_RAW,
-        ModalityType.CLASS_LABEL,
-        ModalityType.VIDEO_IDENTITY,
         ModalityType.MULTI_LABEL,
         ModalityType.ONE_HOT_CLASS_LABEL,
-        ModalityType.IDENTITY,
-        ModalityType.GENERIC_L2_LOSS,
         ModalityType.REAL,
         ModalityType.REAL_L2_LOSS,
         ModalityType.REAL_LOG_POISSON_LOSS,
-        ModalityType.IDENTITY_SYMBOL,
         ModalityType.SIGMOID_CLASS_LABEL,
         ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
-        ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL,
         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
+        ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL,
+        ModalityType.SPEECH_RECOGNITION,
+        ModalityType.SYMBOL,
+        ModalityType.SYMBOL_ONE_HOT,
+        ModalityType.SYMBOL_WEIGHTS_ALL,
+        ModalityType.VIDEO,
+        ModalityType.VIDEO_BITWISE,
+        ModalityType.VIDEO_IDENTITY,
+        ModalityType.VIDEO_L1,
+        ModalityType.VIDEO_L2,
+        ModalityType.VIDEO_L1_RAW,
+        ModalityType.VIDEO_L2_RAW,
+        ModalityType.VIDEO_PIXEL_NOISE,
     ]
 
 
 # Utility functions, similar to tf.keras
-current_module = sys.modules[__name__]
 
 
 def get_bottom(modality_type, value=None):
   """Gets default bottom transformation; if none available, return value."""
-  if modality_type in ModalityType.get_choices():
-    modality_cls = getattr(current_module, modality_type)
-    return modality_cls.bottom
+  if modality_type == ModalityType.AUDIO:
+    return audio_bottom
+  elif modality_type == ModalityType.AUDIO_SPECTRAL:
+    return audio_spectral_bottom
+  elif modality_type in (ModalityType.CLASS_LABEL,
+                         ModalityType.MULTI_LABEL,
+                         ModalityType.ONE_HOT_CLASS_LABEL,
+                         ModalityType.SIGMOID_CLASS_LABEL,
+                         ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
+                         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
+                         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
+                         ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL):
+    return class_label_bottom
+  elif modality_type in (ModalityType.CTC_SYMBOL,
+                         ModalityType.SYMBOL,
+                         ModalityType.SYMBOL_WEIGHTS_ALL):
+    return symbol_bottom
+  elif modality_type in (ModalityType.GENERIC_L2_LOSS,
+                         ModalityType.IDENTITY,
+                         ModalityType.IDENTITY_SYMBOL,
+                         ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM):
+    return identity_bottom
+  elif modality_type == ModalityType.IMAGE:
+    return image_bottom
+  elif modality_type in (ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
+                         ModalityType.IMAGE_CHANNEL_COMPRESS):
+    return image_channel_compress_bottom
+  elif modality_type in (ModalityType.REAL,
+                         ModalityType.REAL_L2_LOSS,
+                         ModalityType.REAL_LOG_POISSON_LOSS):
+    return real_bottom
+  elif modality_type == ModalityType.SPEECH_RECOGNITION:
+    return speech_recognition_bottom
+  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
+    return symbol_one_hot_bottom
+  elif modality_type in (ModalityType.VIDEO,
+                         ModalityType.VIDEO_L1,
+                         ModalityType.VIDEO_L2):
+    return video_bottom
+  elif modality_type == ModalityType.VIDEO_BITWISE:
+    return video_bitwise_bottom
+  elif modality_type == ModalityType.VIDEO_IDENTITY:
+    return video_identity_bottom
+  elif modality_type in (ModalityType.VIDEO_L1_RAW,
+                         ModalityType.VIDEO_L2_RAW):
+    return video_raw_bottom
+  elif modality_type == ModalityType.VIDEO_PIXEL_NOISE:
+    return video_pixel_noise_bottom
   return value
 
 
 def get_loss(modality_type, value=None):
   """Gets default loss transformation; if none available, return value."""
-  if modality_type in ModalityType.get_choices():
-    modality_cls = getattr(current_module, modality_type)
-    return modality_cls.loss
+  if modality_type in (ModalityType.AUDIO,
+                       ModalityType.AUDIO_SPECTRAL,
+                       ModalityType.CLASS_LABEL,
+                       ModalityType.IDENTITY,
+                       ModalityType.IDENTITY_SYMBOL,
+                       ModalityType.IMAGE,
+                       ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
+                       ModalityType.IMAGE_CHANNEL_COMPRESS,
+                       ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM,
+                       ModalityType.REAL,
+                       ModalityType.SPEECH_RECOGNITION,
+                       ModalityType.SYMBOL,
+                       ModalityType.SYMBOL_WEIGHTS_ALL):
+    return generic_loss
+  elif modality_type == ModalityType.CTC_SYMBOL:
+    return ctc_symbol_loss
+  elif modality_type == ModalityType.GENERIC_L2_LOSS:
+    return generic_l2_loss
+  elif modality_type == ModalityType.MULTI_LABEL:
+    return multi_label_loss
+  elif modality_type in (ModalityType.ONE_HOT_CLASS_LABEL,
+                         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
+                         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
+                         ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL):
+    return one_hot_class_label_loss
+  elif modality_type == ModalityType.REAL_L2_LOSS:
+    return real_l2_loss
+  elif modality_type == ModalityType.REAL_LOG_POISSON_LOSS:
+    return real_log_poisson_loss
+  elif modality_type == ModalityType.SIGMOID_CLASS_LABEL:
+    return sigmoid_class_label_loss
+  elif modality_type == ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL:
+    return sigmoid_max_pooling_class_label_loss
+  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
+    return symbol_one_hot_loss
+  elif modality_type in (ModalityType.VIDEO,
+                         ModalityType.VIDEO_BITWISE,
+                         ModalityType.VIDEO_PIXEL_NOISE):
+    return video_loss
+  elif modality_type == ModalityType.VIDEO_IDENTITY:
+    return video_identity_loss
+  elif modality_type == ModalityType.VIDEO_L1:
+    return video_l1_loss
+  elif modality_type == ModalityType.VIDEO_L1_RAW:
+    return video_l1_raw_loss
+  elif modality_type == ModalityType.VIDEO_L2:
+    return video_l2_loss
+  elif modality_type == ModalityType.VIDEO_L2_RAW:
+    return video_l2_raw_loss
   return value
 
 
 def get_name(modality_type, value=None):
   """Gets default name for transformations; if none available, return value."""
-  if modality_type in ModalityType.get_choices():
-    modality_cls = getattr(current_module, modality_type)
-    return modality_cls.name
+  # For legacy reasons, modalities vary in their naming scheme.
+  if modality_type == ModalityType.AUDIO:
+    return lambda model_hparams, vocab_size: "audio_modality"
+  elif modality_type == ModalityType.AUDIO_SPECTRAL:
+    return lambda model_hparams, vocab_size: "audio_spectral_modality"
+  elif modality_type == ModalityType.GENERIC_L2_LOSS:
+    return lambda model_hparams, vocab_size: "generic_l2_loss_modality"
+  elif modality_type == ModalityType.IDENTITY:
+    return lambda model_hparams, vocab_size: "identity_modality"
+  elif modality_type == ModalityType.IMAGE:
+    return lambda model_hparams, vocab_size: "image_modality"
+  elif modality_type == ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY:
+    return (lambda model_hparams, vocab_size:  # pylint: disable=g-long-lambda
+            "image_channel_bottom_identity_modality")
+  elif modality_type == ModalityType.IMAGE_CHANNEL_COMPRESS:
+    return lambda model_hparams, vocab_size: "image_channel_compress_modality"
+  elif modality_type == ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM:
+    return lambda model_hparams, vocab_size: "image_channel_embeddings_bottom"
+  elif modality_type == ModalityType.REAL:
+    return lambda model_hparams, vocab_size: "real_modality"
+  elif modality_type == ModalityType.REAL_L2_LOSS:
+    return lambda model_hparams, vocab_size: "real_l2_loss_modality"
+  elif modality_type == ModalityType.REAL_LOG_POISSON_LOSS:
+    return lambda model_hparams, vocab_size: "real_log_poisson_loss_modality"
+  elif modality_type == ModalityType.SPEECH_RECOGNITION:
+    return lambda model_hparams, vocab_size: "speech_recognition_modality"
+  elif modality_type == ModalityType.VIDEO:
+    return lambda model_hparams, vocab_size: "video_modality"
+  elif modality_type == ModalityType.VIDEO_BITWISE:
+    return lambda model_hparams, vocab_size: "video_modality_bitwise"
+  elif modality_type == ModalityType.VIDEO_IDENTITY:
+    return lambda model_hparams, vocab_size: "video_modality_identity"
+  elif modality_type == ModalityType.VIDEO_L1:
+    return lambda model_hparams, vocab_size: "video_modality_l1"
+  elif modality_type == ModalityType.VIDEO_L1_RAW:
+    return lambda model_hparams, vocab_size: "video_modality_l1_raw"
+  elif modality_type == ModalityType.VIDEO_L2:
+    return lambda model_hparams, vocab_size: "video_modality_l2"
+  elif modality_type == ModalityType.VIDEO_L2_RAW:
+    return lambda model_hparams, vocab_size: "video_modality_l2_raw"
+  elif modality_type == ModalityType.VIDEO_PIXEL_NOISE:
+    return lambda model_hparams, vocab_size: "video_modality_pixel_noise"
+  elif modality_type in (ModalityType.CLASS_LABEL,
+                         ModalityType.MULTI_LABEL,
+                         ModalityType.ONE_HOT_CLASS_LABEL):
+    def name(model_hparams, vocab_size):
+      return "class_label_modality_%d_%d" % (vocab_size,
+                                             model_hparams.hidden_size)
+    return name
+  elif modality_type in (ModalityType.CTC_SYMBOL,
+                         ModalityType.IDENTITY_SYMBOL,
+                         ModalityType.SYMBOL,
+                         ModalityType.SYMBOL_WEIGHTS_ALL,
+                         ModalityType.SYMBOL_ONE_HOT):
+    def name(model_hparams, vocab_size):
+      return "symbol_modality_%d_%d" % (vocab_size, model_hparams.hidden_size)
+    return name
+  elif modality_type == ModalityType.SIGMOID_CLASS_LABEL:
+    def name(model_hparams, vocab_size):
+      return "sigmoid_class_symbol_modality_%d_%d" % (vocab_size,
+                                                      model_hparams.hidden_size)
+    return name
+  elif modality_type == ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL:
+    def name(model_hparams, vocab_size):
+      return "sigmoid_max_pooling_class_symbol_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)
+    return name
+  elif modality_type == ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL:
+    def name(model_hparams, vocab_size):
+      return "softmax_average_pooling_onehot_class_label_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)
+    return name
+  elif modality_type == ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL:
+    def name(model_hparams, vocab_size):
+      return "softmax_last_timestep_onehot_class_label_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)
+    return name
+  elif modality_type == ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL:
+    def name(model_hparams, vocab_size):
+      return "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)
+    return name
   return value
 
 
 def get_targets_bottom(modality_type, value=None):
   """Gets default bottom transformation for targets; if none, return value."""
-  if modality_type in ModalityType.get_choices():
-    modality_cls = getattr(current_module, modality_type)
-    return modality_cls.targets_bottom
+  if modality_type == ModalityType.AUDIO:
+    return make_targets_bottom(audio_bottom)
+  elif modality_type == ModalityType.AUDIO_SPECTRAL:
+    return make_targets_bottom(audio_spectral_bottom)
+  elif modality_type in (ModalityType.CLASS_LABEL,
+                         ModalityType.MULTI_LABEL,
+                         ModalityType.ONE_HOT_CLASS_LABEL,
+                         ModalityType.SIGMOID_CLASS_LABEL,
+                         ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
+                         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
+                         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
+                         ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL):
+    return class_label_targets_bottom
+  elif modality_type in (ModalityType.CTC_SYMBOL,
+                         ModalityType.SYMBOL,
+                         ModalityType.SYMBOL_WEIGHTS_ALL):
+    return symbol_targets_bottom
+  elif modality_type in (ModalityType.GENERIC_L2_LOSS,
+                         ModalityType.IDENTITY_SYMBOL):
+    return identity_bottom
+  elif modality_type == ModalityType.IDENTITY:
+    return make_targets_bottom(identity_bottom)
+  elif modality_type == ModalityType.IMAGE:
+    return image_targets_bottom
+  elif modality_type in (ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
+                         ModalityType.IMAGE_CHANNEL_COMPRESS):
+    return image_channel_compress_targets_bottom
+  elif modality_type == ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM:
+    return image_channel_embeddings_bottom
+  elif modality_type in (ModalityType.REAL,
+                         ModalityType.REAL_L2_LOSS,
+                         ModalityType.REAL_LOG_POISSON_LOSS):
+    return make_targets_bottom(real_bottom)
+  elif modality_type == ModalityType.SPEECH_RECOGNITION:
+    return make_targets_bottom(speech_recognition_bottom)
+  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
+    return symbol_one_hot_bottom
+  elif modality_type in (ModalityType.VIDEO,
+                         ModalityType.VIDEO_L1,
+                         ModalityType.VIDEO_L2):
+    return video_targets_bottom
+  elif modality_type == ModalityType.VIDEO_BITWISE:
+    return video_bitwise_targets_bottom
+  elif modality_type == ModalityType.VIDEO_IDENTITY:
+    return video_identity_targets_bottom
+  elif modality_type in (ModalityType.VIDEO_L1_RAW,
+                         ModalityType.VIDEO_L2_RAW):
+    return video_raw_targets_bottom
+  elif modality_type == ModalityType.VIDEO_PIXEL_NOISE:
+    return make_targets_bottom(video_pixel_noise_bottom)
   return value
 
 
-def get_targets_weights_fn(modality_type, value=None):
-  """Gets default weights function; if none available, return value."""
-  if modality_type in ModalityType.get_choices():
-    modality_cls = getattr(current_module, modality_type)
-    return modality_cls.targets_weights_fn
+def get_top(modality_type, value=None):
+  """Gets default top transformation; if none available, return value."""
+  if modality_type in (ModalityType.AUDIO,
+                       ModalityType.AUDIO_SPECTRAL,
+                       ModalityType.GENERIC_L2_LOSS,
+                       ModalityType.IDENTITY,
+                       ModalityType.IDENTITY_SYMBOL,
+                       ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
+                       ModalityType.SPEECH_RECOGNITION,
+                       ModalityType.VIDEO_IDENTITY):
+    return identity_top
+  elif modality_type in (ModalityType.CLASS_LABEL,
+                         ModalityType.MULTI_LABEL,
+                         ModalityType.ONE_HOT_CLASS_LABEL,
+                         ModalityType.SIGMOID_CLASS_LABEL):
+    return class_label_top
+  elif modality_type in (ModalityType.CTC_SYMBOL,
+                         ModalityType.SYMBOL,
+                         ModalityType.SYMBOL_WEIGHTS_ALL):
+    return symbol_top
+  elif modality_type == ModalityType.IMAGE:
+    return image_top
+  elif modality_type == ModalityType.IMAGE_CHANNEL_COMPRESS:
+    return image_channel_compress_top
+  elif modality_type == ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM:
+    return image_channel_embeddings_top
+  elif modality_type in (ModalityType.REAL,
+                         ModalityType.REAL_L2_LOSS,
+                         ModalityType.REAL_LOG_POISSON_LOSS):
+    return real_top
+  elif modality_type == ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL:
+    return sigmoid_max_pooling_class_label_top
+  elif modality_type == ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL:
+    return softmax_average_pooling_class_label_top
+  elif modality_type == ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL:
+    return softmax_last_timestep_class_label_top
+  elif modality_type == ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL:
+    return softmax_max_pooling_class_label_top
+  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
+    return symbol_one_hot_top
+  elif modality_type in (ModalityType.VIDEO,
+                         ModalityType.VIDEO_BITWISE,
+                         ModalityType.VIDEO_PIXEL_NOISE):
+    return video_top
+  elif modality_type in (ModalityType.VIDEO_L1,
+                         ModalityType.VIDEO_L2):
+    return video_l1_top
+  elif modality_type in (ModalityType.VIDEO_L1_RAW,
+                         ModalityType.VIDEO_L2_RAW):
+    return video_raw_top
   return value
 
 
-def get_top(modality_type, value=None):
-  """Gets default top transformation; if none available, return value."""
-  if modality_type in ModalityType.get_choices():
-    modality_cls = getattr(current_module, modality_type)
-    return modality_cls.top
+def get_weights_fn(modality_type, value=None):
+  """Gets default weights function; if none available, return value."""
+  if modality_type in (ModalityType.CTC_SYMBOL,
+                       ModalityType.IDENTITY_SYMBOL,
+                       ModalityType.MULTI_LABEL,
+                       ModalityType.SYMBOL,
+                       ModalityType.SYMBOL_ONE_HOT):
+    return common_layers.weights_nonzero
+  elif modality_type in ModalityType.get_choices():
+    return common_layers.weights_all
   return value
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 6cb6efc98..69504bb8e 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -31,6 +31,26 @@
 
 class ModalityTest(tf.test.TestCase):
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testGetForAllModalities(self):
+    for modality in modalities.ModalityType.get_choices():
+      bottom = modalities.get_bottom(modality)
+      loss = modalities.get_loss(modality)
+      name = modalities.get_name(modality)
+      targets_bottom = modalities.get_targets_bottom(modality)
+      top = modalities.get_top(modality)
+      weights_fn = modalities.get_weights_fn(modality)
+      self.assertIsNotNone(bottom,
+                           msg="{} has no default bottom".format(modality))
+      self.assertIsNotNone(loss, msg="{} has no default loss".format(modality))
+      self.assertIsNotNone(name, msg="{} has no default name".format(modality))
+      self.assertIsNotNone(
+          targets_bottom,
+          msg="{} has no default targets_bottom".format(modality))
+      self.assertIsNotNone(top, msg="{} has no default top".format(modality))
+      self.assertIsNotNone(weights_fn,
+                           msg="{} has no default weights_fn".format(modality))
+
   @test_utils.run_in_graph_and_eager_modes()
   def testSymbolModalityInputs(self):
     batch_size = 10
@@ -87,7 +107,7 @@ def testSymbolModalityTargets(self):
         sharded_targets,
         model_hparams,
         vocab_size,
-        modalities.get_targets_weights_fn(modalities.ModalityType.SYMBOL))
+        modalities.get_weights_fn(modalities.ModalityType.SYMBOL))
     train_loss = (tf.add_n(sharded_loss_num) /
                   tf.maximum(1.0, tf.add_n(sharded_loss_den)))
     logits = tf.concat(sharded_logits, 0)
@@ -124,12 +144,12 @@ def testSymbolModalityTargetsFactored(self):
           model_hparams,
           vocab_size)
       sharded_loss_num, sharded_loss_den = data_parallelism(
-          modalities.SymbolModality.loss,
+          modalities.get_loss(modalities.ModalityType.SYMBOL),
           sharded_logits,
           sharded_targets,
           model_hparams,
           vocab_size,
-          modalities.get_targets_weights_fn(modalities.ModalityType.SYMBOL))
+          modalities.get_weights_fn(modalities.ModalityType.SYMBOL))
       train_loss = (tf.add_n(sharded_loss_num) /
                     tf.maximum(1.0, tf.add_n(sharded_loss_den)))
       logits = tf.concat(sharded_logits, 0)
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 81d8de434..d43f96c15 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -48,16 +48,11 @@ def body(self, features):
     hparams = copy.copy(self._hparams)
     targets = features["targets"]
     if (hparams.likelihood == cia.DistributionType.DMOL and
-        (hparams.modality["targets"] !=
-         modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY or
-         hparams.num_channels != 1)):
-      raise ValueError("When using DMOL for the likelihood,modality['targets'] "
-                       "must be ImageChannelBottomIdentityModality and "
-                       "num_channels must be 1.")
+        hparams.num_channels != 1):
+      raise ValueError("When using DMOL for the likelihood, bottom function "
+                       " must be identity and num_channels must be 1.")
     if (not tf.get_variable_scope().reuse and
-        hparams.mode != tf.estimator.ModeKeys.PREDICT and
-        hparams.modality["targets"] !=
-        modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY):
+        hparams.mode != tf.estimator.ModeKeys.PREDICT):
       tf.summary.image("targets", tf.to_float(targets), max_outputs=1)
 
     # Extra losses list if we want to use moe.
@@ -193,7 +188,8 @@ def image_transformer_base():
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.98
   hparams.label_smoothing = 0.0
-  hparams.modality["targets"] = modalities.ModalityType.IDENTITY
+  hparams.bottom["targets"] = modalities.image_channel_embeddings_bottom
+  hparams.top["targets"] = modalities.identity_top
   hparams.norm_type = "layer"
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.add_hparam("filter_size", 512)  # Add new ones like this.
@@ -280,8 +276,8 @@ def imagetransformer_cifar10_base_dmol():
   hparams = image_transformer_base()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.modality["targets"] = (
-      modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY)
+  hparams.bottom["targets"] = modalities.image_channel_compress_targets_bottom
+  hparams.top["targets"] = modalities.identity_top
   hparams.num_heads = 8
   hparams.batch_size = 8
   hparams.sampling_method = "random"
@@ -422,8 +418,8 @@ def imagetransformerpp_sep_channels_8l_8h():
   hparams = imagetransformer_base()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.modality["targets"] = (
-      modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY)
+  hparams.bottom["targets"] = modalities.image_channel_compress_targets_bottom
+  hparams.top["targets"] = modalities.identity_top
   hparams.num_heads = 8
   hparams.batch_size = 4
   hparams.attention_key_channels = hparams.attention_value_channels = 0
@@ -886,8 +882,8 @@ def imagetransformerpp_tiny():
   hparams = imagetransformer_tiny()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.modality["targets"] = (
-      modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY)
+  hparams.bottom["targets"] = modalities.image_channel_compress_targets_bottom
+  hparams.top["targets"] = modalities.identity_top
   return hparams
 
 
diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index acf793440..80ce13d53 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -382,7 +382,9 @@ def image_transformer2d_base():
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.98
   hparams.label_smoothing = 0.0
-  hparams.modality["targets"] = modalities.ModalityType.IDENTITY
+  hparams.bottom["targets"] = modalities.make_targets_bottom(
+      modalities.image_channel_embeddings_bottom)
+  hparams.top["targets"] = modalities.identity_top
   hparams.norm_type = "layer"
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.add_hparam("filter_size", 512)  # Add new ones like this.
@@ -593,6 +595,7 @@ def img2img_transformer2d_base():
   hparams.filter_size = 2048
   hparams.num_encoder_layers = 4
   hparams.num_decoder_layers = 8
+  hparams.bottom["inputs"] = modalities.image_channel_embeddings_bottom
   hparams.dec_attention_type = cia.AttentionType.LOCAL_2D
   hparams.block_raster_scan = True
   return hparams
diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 3e32e9c03..b8d4c56d8 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -854,9 +854,12 @@ def mtf_transformer_base():
   # These parameters make Transformer model compatible with MtfTransformer
   # Do not override these, as mtf_transformer does not support other options.
   hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
-  hparams.modality = {
-      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
-      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
+  hparams.bottom = {
+      "inputs": modalities.identity_bottom,
+      "targets": modalities.identity_bottom,
+  }
+  hparams.top = {
+      "targets": modalities.identity_top,
   }
 
   # Parameters for computing the maximum decode length in beam search.
diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index 1d1e2a71d..c8cb5354d 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -418,9 +418,12 @@ def mtf_transformer2_base():
   hparams.use_fixed_batch_size = True
   hparams.add_hparam("mtf_mode", True)
   hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
-  hparams.modality = {
-      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
-      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
+  hparams.bottom = {
+      "inputs": modalities.identity_bottom,
+      "targets": modalities.identity_bottom,
+  }
+  hparams.top = {
+      "targets": modalities.identity_top,
   }
   hparams.add_hparam("beam_size", 1)
 
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 02dde9e03..cf4305a30 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1112,9 +1112,12 @@ def autoencoder_residual_text():
   hparams.hidden_size = 64
   hparams.max_hidden_size = 512
   hparams.bottleneck_noise = 0.0
-  hparams.modality = {
-      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
-      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
+  hparams.bottom = {
+      "inputs": modalities.identity_bottom,
+      "targets": modalities.identity_bottom,
+  }
+  hparams.top = {
+      "targets": modalities.identity_top,
   }
   hparams.autoregressive_mode = "none"
   hparams.sample_width = 1
@@ -1219,9 +1222,12 @@ def autoencoder_ordered_text():
   hparams.batch_size = 1024
   hparams.autoregressive_mode = "conv5"
   hparams.max_hidden_size = 1024
-  hparams.modality = {
-      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
-      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
+  hparams.bottom = {
+      "inputs": modalities.identity_bottom,
+      "targets": modalities.identity_bottom,
+  }
+  hparams.top = {
+      "targets": modalities.identity_top,
   }
   hparams.sample_height = 128
   hparams.sample_width = 1
diff --git a/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index 94842b2db..3ba2e64ab 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -127,9 +127,12 @@ def cycle_gan_small():
   """Set of hyperparameters."""
   hparams = transformer_vae.transformer_ae_small()
   hparams.batch_size = 2048
-  hparams.modality = {
-      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
-      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
+  hparams.bottom = {
+      "inputs": modalities.identity_bottom,
+      "targets": modalities.identity_bottom,
+  }
+  hparams.top = {
+      "targets": modalities.identity_top,
   }
   hparams.weight_decay = 3.0
   hparams.learning_rate = 0.05
diff --git a/tensor2tensor/models/research/super_lm.py b/tensor2tensor/models/research/super_lm.py
index 5a02ccf94..0e7dd58c8 100644
--- a/tensor2tensor/models/research/super_lm.py
+++ b/tensor2tensor/models/research/super_lm.py
@@ -265,7 +265,13 @@ def super_lm_base():
   # we only want one data shard.
   hparams.no_data_parallelism = True
   # bypass the symbol modality so that we can use model parallelism.
-  hparams.modality["targets"] = modalities.ModalityType.IDENTITY_SYMBOL
+  hparams.bottom = {
+      "inputs": modalities.identity_bottom,
+      "targets": modalities.identity_bottom,
+  }
+  hparams.top = {
+      "targets": modalities.identity_top,
+  }
   hparams.add_hparam("filter_size", 512)
   hparams.add_hparam("mix_fraction", 0.5)
   # attention-related flags
diff --git a/tensor2tensor/models/research/transformer_symshard.py b/tensor2tensor/models/research/transformer_symshard.py
index 98c9cf0ee..7f5bf39dd 100644
--- a/tensor2tensor/models/research/transformer_symshard.py
+++ b/tensor2tensor/models/research/transformer_symshard.py
@@ -367,9 +367,12 @@ def transformer_symshard_base():
   # we only want one data shard.
   hparams.no_data_parallelism = True
   # bypass the symbol modality so that we can use model parallelism.
-  hparams.modality = {
-      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
-      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
+  hparams.bottom = {
+      "inputs": modalities.identity_bottom,
+      "targets": modalities.identity_bottom,
+  }
+  hparams.top = {
+      "targets": modalities.identity_top,
   }
   hparams.add_hparam("filter_size", 1280)
   hparams.add_hparam("mix_fraction", 0.5)
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index e8f769bb6..9cfe80933 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -896,7 +896,8 @@ def imagetransformer_ae_cifar():
 
   hparams.add_hparam("unconditional", False)  # unconditional generation
 
-  hparams.modality["targets"] = modalities.ImageChannelEmbeddingsBottom
+  hparams.bottom["targets"] = modalities.image_channel_embeddings_bottom
+  hparams.top["targets"] = modalities.image_channel_embeddings_top
   hparams.drop_inputs = True
   hparams.do_attend_compress = False
   hparams.do_attend_decompress = False
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index c3dfb53ea..ec068e452 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -195,16 +195,12 @@ def __init__(self, *args, **kwargs):
 
   @property
   def _target_modality(self):
-    target_modality = self.hparams.modality.get(
-        "targets",
-        self.problem_hparams.modality["targets"])
-    if target_modality not in modalities.ModalityType.get_choices():
-      target_modality = target_modality.__class__.__name__
-    return target_modality
+    return self.problem_hparams.modality["targets"]
 
   @property
   def is_per_pixel_softmax(self):
-    return self._target_modality == modalities.ModalityType.VIDEO
+    # TODO(trandustin): This is a hack.
+    return "targets" not in self.hparams.get("loss")
 
   def get_iteration_num(self):
     step_num = tf.train.get_global_step()
@@ -338,11 +334,12 @@ def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
       Additional reconstruction loss.
 
     Raises:
-      ValueError: in case of unknown modality.
+      ValueError: in case of unknown loss transformation.
     """
-    if self._target_modality == modalities.ModalityType.VIDEO_L2_RAW:
+    # TODO(trandustin): This logic should be moved elsewhere.
+    if self.hparams.loss.get("targets") == modalities.video_l2_raw_loss:
       recon_loss = tf.losses.mean_squared_error(extra_gts, extra_pds)
-    elif self._target_modality == modalities.ModalityType.VIDEO:
+    elif "targets" not in self.hparams.loss:
       shape = common_layers.shape_list(extra_pds)
       updated_shape = shape[:-1] + [3, 256]
       extra_pds = tf.reshape(extra_pds, updated_shape)
@@ -351,10 +348,9 @@ def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
       targets = extra_raw_gts
       targets_shape = common_layers.shape_list(targets)
       targets = tf.reshape(targets, [-1] + targets_shape[2:])
-      modality = self.hparams.problem_hparams.modality["targets"]
-      targets_weights_fn = self.hparams.targets_weights_fn.get(
+      targets_weights_fn = self.hparams.weights_fn.get(
           "targets",
-          modalities.get_targets_weights_fn(modality))
+          modalities.get_weights_fn(self._target_modality))
       numerator, denominator = common_layers.padded_cross_entropy(
           logits,
           targets,
@@ -363,7 +359,7 @@ def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
           weights_fn=targets_weights_fn)
       recon_loss = numerator / denominator
     else:
-      raise ValueError("internal loss only supports specific modalities.")
+      raise ValueError("internal loss only supports specific hparams.loss.")
     tf.summary.scalar("recon_extra", recon_loss)
     return recon_loss
 
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 7616d3549..8d8a2c6ff 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -55,7 +55,8 @@ def next_frame_pixel_noise():
   """Basic 2-frame conv model with pixel noise."""
   hparams = next_frame_basic_deterministic()
   hparams.add_hparam("video_modality_input_noise", 0.05)
-  hparams.modality["inputs"] = modalities.ModalityType.VIDEO_PIXEL_NOISE
+  hparams.bottom["inputs"] = modalities.video_pixel_noise_bottom
+  hparams.top["inputs"] = modalities.video_top
   return hparams
 
 
@@ -89,7 +90,8 @@ def next_frame_tpu():
 def next_frame_ae():
   """Conv autoencoder."""
   hparams = next_frame_basic_deterministic()
-  hparams.modality["inputs"] = modalities.ModalityType.VIDEO_BITWISE
+  hparams.bottom["inputs"] = modalities.video_bitwise_bottom
+  hparams.top["inputs"] = modalities.video_top
   hparams.hidden_size = 256
   hparams.batch_size = 8
   hparams.num_hidden_layers = 4
@@ -102,7 +104,8 @@ def next_frame_ae():
 def next_frame_ae_tiny():
   """Conv autoencoder, tiny set for testing."""
   hparams = next_frame_tiny()
-  hparams.modality["inputs"] = modalities.ModalityType.VIDEO_BITWISE
+  hparams.bottom["inputs"] = modalities.video_bitwise_bottom
+  hparams.top["inputs"] = modalities.video_top
   hparams.batch_size = 8
   hparams.dropout = 0.4
   return hparams
@@ -131,7 +134,8 @@ def next_frame_tiny():
 def next_frame_l1():
   """Basic conv model with L1 modality."""
   hparams = next_frame_basic_deterministic()
-  hparams.modality["targets"] = modalities.ModalityType.VIDEO_L1
+  hparams.loss["targets"] = modalities.video_l1_loss
+  hparams.top["targets"] = modalities.video_l1_top
   hparams.video_modality_loss_cutoff = 2.4
   return hparams
 
@@ -140,7 +144,8 @@ def next_frame_l1():
 def next_frame_l2():
   """Basic conv model with L2 modality."""
   hparams = next_frame_basic_deterministic()
-  hparams.modality["targets"] = modalities.ModalityType.VIDEO_L2
+  hparams.loss["targets"] = modalities.video_l2_loss
+  hparams.top["targets"] = modalities.video_l1_top
   hparams.video_modality_loss_cutoff = 2.4
   return hparams
 
diff --git a/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
index a8b8f98c9..e42c94d9d 100644
--- a/tensor2tensor/models/video/epva_params.py
+++ b/tensor2tensor/models/video/epva_params.py
@@ -29,9 +29,15 @@ def next_frame_epva():
   hparams = basic_deterministic_params.next_frame_basic_deterministic()
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 4
-  hparams.modality = {
-      "inputs": modalities.ModalityType.VIDEO_L2_RAW,
-      "targets": modalities.ModalityType.VIDEO_L2_RAW,
+  hparams.bottom = {
+      "inputs": modalities.video_raw_bottom,
+      "targets": modalities.video_raw_targets_bottom,
+  }
+  hparams.loss = {
+      "targets": modalities.video_l2_raw_loss,
+  }
+  hparams.top = {
+      "targets": modalities.video_raw_top,
   }
   hparams.learning_rate_schedule = "constant"
   hparams.learning_rate_constant = 1e-05
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index 56a9da10b..c323ae088 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -69,9 +69,15 @@ def next_frame_glow_hparams():
   # Pretrains the glow encoder for "pretrain_steps" number of steps.
   # By default, don't pretrain and learn end-to-end
   hparams.add_hparam("pretrain_steps", -1)
-  hparams.modality = {
-      "inputs": modalities.ModalityType.VIDEO_L1_RAW,
-      "targets": modalities.ModalityType.VIDEO_L1_RAW,
+  hparams.bottom = {
+      "inputs": modalities.video_raw_bottom,
+      "targets": modalities.video_raw_targets_bottom,
+  }
+  hparams.loss = {
+      "targets": modalities.video_l1_raw_loss,
+  }
+  hparams.top = {
+      "targets": modalities.video_raw_top,
   }
   hparams.init_batch_size = 256
   hparams.batch_size = 32
diff --git a/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index 708c888f3..c28c4378a 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -36,9 +36,15 @@ def next_frame_savp():
   hparams.add_hparam("gan_loss_multiplier", 0.01)
   hparams.add_hparam("gan_vae_loss_multiplier", 0.01)
   hparams.add_hparam("gan_optimization", "joint")
-  hparams.modality = {
-      "inputs": modalities.ModalityType.VIDEO_L1_RAW,
-      "targets": modalities.ModalityType.VIDEO_L1_RAW,
+  hparams.bottom = {
+      "inputs": modalities.video_raw_bottom,
+      "targets": modalities.video_raw_targets_bottom,
+  }
+  hparams.loss = {
+      "targets": modalities.video_l1_raw_loss,
+  }
+  hparams.top = {
+      "targets": modalities.video_raw_top,
   }
   hparams.latent_loss_multiplier_schedule = "linear"
   hparams.upsample_method = "bilinear_upsample_conv"
@@ -54,9 +60,8 @@ def next_frame_savp():
 def next_frame_savp_l2():
   """SAVP with L2 reconstruction loss."""
   hparams = next_frame_savp()
-  hparams.modality = {
-      "inputs": modalities.ModalityType.VIDEO_L2_RAW,
-      "targets": modalities.ModalityType.VIDEO_L2_RAW,
+  hparams.loss = {
+      "targets": modalities.video_l2_raw_loss,
   }
   return hparams
 
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index f5687865b..f7352c714 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -33,9 +33,15 @@ def next_frame_sv2p():
   hparams.video_num_input_frames = 1
   hparams.video_num_target_frames = 3
   hparams.batch_size = 16
-  hparams.modality = {
-      "inputs": modalities.ModalityType.VIDEO_L2_RAW,
-      "targets": modalities.ModalityType.VIDEO_L2_RAW,
+  hparams.bottom = {
+      "inputs": modalities.video_raw_bottom,
+      "targets": modalities.video_raw_targets_bottom,
+  }
+  hparams.loss = {
+      "targets": modalities.video_l2_raw_loss,
+  }
+  hparams.top = {
+      "targets": modalities.video_raw_top,
   }
   hparams.video_modality_loss_cutoff = 0.0
   hparams.scheduled_sampling_mode = "count"
@@ -91,10 +97,9 @@ def next_frame_sv2p_atari():
 def next_frame_sv2p_atari_softmax():
   """SV2P model for atari with softmax."""
   hparams = next_frame_sv2p_atari()
-  hparams.modality = {
-      "inputs": modalities.ModalityType.VIDEO,
-      "targets": modalities.ModalityType.VIDEO,
-  }
+  hparams.bottom = {}
+  hparams.loss = {}
+  hparams.top = {}
   hparams.internal_loss = True
   return hparams
 
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index d3852c648..951cbc3d5 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -614,9 +614,9 @@ def weights_fn_for_mp(problem_task_id):
       tm = {"targets": tm}
 
     for target_name, modality in six.iteritems(tm):
-      weights_fn = model_hparams.targets_weights_fn.get(
+      weights_fn = model_hparams.weights_fn.get(
           "targets",
-          modalities.get_targets_weights_fn(modality))
+          modalities.get_weights_fn(modality))
       if hasattr(model_hparams.problem, "task_list"):
         ptid = problem_instance.task_id  # pylint: disable=cell-var-from-loop
         weights_fn = weights_fn_for_mp(ptid)
@@ -643,9 +643,9 @@ def create_eager_metrics_for_problem(problem, model_hparams):
   metric_fns = problem.eval_metric_fns(model_hparams)
   problem_hparams = problem.get_hparams(model_hparams)
   target_modality = problem_hparams.modality["targets"]
-  weights_fn = model_hparams.targets_weights_fn.get(
+  weights_fn = model_hparams.weights_fn.get(
       "targets",
-      modalities.get_targets_weights_fn(target_modality))
+      modalities.get_weights_fn(target_modality))
   return create_eager_metrics_internal(metric_fns, weights_fn=weights_fn)
 
 
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
deleted file mode 100644
index 09a83d050..000000000
--- a/tensor2tensor/utils/modality.py
+++ /dev/null
@@ -1,121 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Modality base class - defines the bottom and top of the model."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.layers import common_attention
-from tensor2tensor.layers import common_layers
-from tensor2tensor.utils import misc_utils
-
-import tensorflow as tf
-
-
-class Modality(object):
-  """Abstract Modality class for data transformations.
-
-  An abstract class representing modalities for transforming data to a space
-  interpretable by T2T models. It has 4 functions:
-  * bottom: called on inputs entering the model.
-  * targets_bottom: called on targets entering the model (e.g., the decoder).
-  * top: called on model outputs to generate predictions (e.g., logits).
-  * loss: called on predictions (outputs of top) and targets.
-
-  For example, think about a modality for images:
-  * `bottom` represents the part of the model applied to an incoming image,
-    e.g., an entry flow of a convolutional network.
-  * `top` represents the top part of a model that is generating images, e.g., a
-    PixelCNN network.
-  * `targets_bottom` represents the auto-regressive part of the network.  It is
-    applied to the already-generated part of an image, which is given to the
-    decoder to generate the next part. In some cases, e.g., for text, it is the
-    same as the `bottom` function, and that is the default we use. But, e.g.,
-    for images, a different function might be needed to regress properly.
-  * `loss` would compare the generated image to the target image and score it.
-  """
-
-  def __init__(self, model_hparams, vocab_size=None):
-    # __init__ args are unused in any methods. They're maintained for
-    # backwards compatibility for now. In the future, Modality classes will be
-    # removed altogether.
-    del model_hparams, vocab_size
-
-  @classmethod
-  def name(cls, model_hparams, vocab_size=None):
-    del model_hparams, vocab_size  # unused arg
-    return misc_utils.camelcase_to_snakecase(type(cls).__name__)
-
-  targets_weights_fn = staticmethod(common_layers.weights_all)
-
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size=None):
-    """Transform one shard of input.
-
-    Args:
-      x: An int32 Tensor with shape [batch, p0, p1, input_channels]
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      A float32 Tensor with shape [batch, p0, p1, body_input_depth]
-    """
-    raise NotImplementedError("Abstract Method")
-
-  @classmethod
-  def targets_bottom(cls, x, model_hparams, vocab_size=None):
-    """Transform one shard of targets.
-
-    Args:
-      x: An int32 Tensor with shape [batch, p0, p1, target_channels]
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      A float32 Tensor with shape [batch, p0, p1, body_input_depth]
-    """
-    with tf.variable_scope("targets_bottom"):
-      return cls.bottom(x, model_hparams, vocab_size)
-
-  @staticmethod
-  def top(body_output, targets, model_hparams, vocab_size=None):
-    """Generate predictions/logits for one shard of output.
-
-    Most classes will override this function.
-
-    Args:
-      body_output: A Tensor with shape [batch, p0, p1, body_output_depth]
-      targets: A Tensor with shape [batch, p0, p1, targets_channels,
-        top_dimensionality]
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      A Tensor of class logits.
-    """
-    raise NotImplementedError("Abstract Method")
-
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Compute loss numerator and denominator for one shard of output."""
-    del vocab_size  # unused arg
-    logits = top_out
-    logits = common_attention.maybe_upcast(logits, hparams=model_hparams)
-    return common_layers.padded_cross_entropy(
-        logits,
-        targets,
-        model_hparams.label_smoothing,
-        weights_fn=weights_fn)
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index f3b5c8fec..dc3451cbf 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -204,47 +204,19 @@ def __init__(self,
           },
           hparams=hparams)
 
-    # TODO(trandustin): For now, we get custom feature transformations via
-    # hparams.modality. Once modality classes are removed, let users
-    # individually specify custom transformations for bottom, loss, etc.
-    if not hasattr(hparams, "modality"):
-      hparams.add_hparam("modality", {})
-    if not hasattr(hparams, "bottom"):
-      hparams.add_hparam("bottom", {})
-    if not hasattr(hparams, "loss"):
-      hparams.add_hparam("loss", {})
-    if not hasattr(hparams, "name"):
-      hparams.add_hparam("name", {})
-    if not hasattr(hparams, "targets_weights_fn"):
-      hparams.add_hparam("targets_weights_fn", {})
-    if not hasattr(hparams, "top"):
-      hparams.add_hparam("top", {})
-    target_modalities = _create_target_modality(hparams.modality)
-    for feature_name, modality in six.iteritems(hparams.modality):
-      if modality in modalities.ModalityType.get_choices():
-        modality = getattr(modalities, modality)
-      if feature_name in target_modalities:
-        hparams.bottom[feature_name] = modality.targets_bottom
-      else:
-        hparams.bottom[feature_name] = modality.bottom
-      hparams.loss[feature_name] = modality.loss
-      hparams.name[feature_name] = modality.name
-      hparams.targets_weights_fn[feature_name] = modality.targets_weights_fn
-      hparams.top[feature_name] = modality.top
-
     if self._problem_hparams:
       for feature_name, modality in six.iteritems(
           self._problem_hparams.modality):
-        # If prepend mode, set targets_weights_fn to appropriately handle it.
-        if (modality in (modalities.ModalityType.SYMBOL,
-                         modalities.ModalityType.SYMBOL_ONE_HOT,
-                         modalities.ModalityType.CTC_SYMBOL,
-                         modalities.ModalityType.IDENTITY_SYMBOL)):
+        # If prepend mode, set weights_fn to appropriately handle it.
+        if (modality in (modalities.ModalityType.CTC_SYMBOL,
+                         modalities.ModalityType.IDENTITY_SYMBOL,
+                         modalities.ModalityType.SYMBOL,
+                         modalities.ModalityType.SYMBOL_ONE_HOT)):
           if (hparams.prepend_mode == "prepend_inputs_full_attention" or
               (hparams.prepend_mode == "prepend_inputs_masked_attention" and
                mode != tf.estimator.ModeKeys.TRAIN)):
             weights_fn = common_layers.weights_prepend_inputs_to_targets
-            hparams.targets_weights_fn[feature_name] = weights_fn
+            hparams.weights_fn[feature_name] = weights_fn
 
     self._original_hparams = hparams
     self.set_mode(mode)
@@ -646,8 +618,8 @@ def _loss_single(self, logits, feature_name, feature, weights=None):
     if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
       vocab_size += (-vocab_size) % self._hparams.vocab_divisor
     loss = self._hparams.loss.get(feature_name, modalities.get_loss(modality))
-    targets_weights_fn = self._hparams.targets_weights_fn.get(
-        "targets", modalities.get_targets_weights_fn(modality))
+    targets_weights_fn = self._hparams.weights_fn.get(
+        "targets", modalities.get_weights_fn(modality))
     if weights is None:
       loss_num, loss_den = loss(logits, feature, self._hparams, vocab_size,
                                 weights_fn=targets_weights_fn)
@@ -1817,7 +1789,7 @@ def create_tpu_eval_metrics_fn(problem, model_hparams):
   tm = _create_target_modality(problem.get_hparams(model_hparams).modality)
   if isinstance(tm, dict):
     for k, v in six.iteritems(tm):
-      weights_fn = v.targets_weights_fn
+      weights_fn = v.weights_fn
 
       def make_metric_fn(metric_fn):
         def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
@@ -1837,7 +1809,7 @@ def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
         name = "%s/metrics-%s/%s" % (k, problem.name, metric)
         metric_fns.append((name, make_metric_fn(metric_fn)))
   else:
-    weights_fn = tm.targets_weights_fn
+    weights_fn = tm.weights_fn
 
     def make_metric_fn(metric_fn):
       def wrapped_metric_fn(logits, labels, features):
@@ -2038,8 +2010,8 @@ def sampled_results():
                                 vocab_size)
         if "training" not in losses:
           loss = hparams.loss.get("targets", modalities.get_loss(modality))
-          weights_fn = hparams.targets_weights_fn.get(
-              "targets", modalities.get_targets_weights_fn(modality))
+          weights_fn = hparams.weights_fn.get(
+              "targets", modalities.get_weights_fn(modality))
           sharded_loss_num, sharded_loss_den = dp(loss,
                                                   sharded_logits,
                                                   sharded_features["targets"],
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index c9b680302..229846832 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -52,6 +52,8 @@ def testLossSingleWeights(self):
 
         model_hparams = HParams(
             prepend_mode="none",
+            loss={},
+            weights_fn={},
             label_smoothing=0.0,
             shared_embedding_and_softmax_weights=False)
 

From 5923c0d89a5c15c4e6d9293bcf93f31becb43fc5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 25 Feb 2019 22:24:50 -0800
Subject: [PATCH 1719/2720] Remove Modality classes.

Originally, users override default feature transformations by setting hparams.modality. Now, users override transformations by setting hparams.bottom, hparams.loss, hparams.top, and hparams.weights_fn.

See changes to common_hparams.py, modalities.py, modality.py. Other changes are maintenance.

PiperOrigin-RevId: 235661775
---
 .../data_generators/multi_problem.py          |    8 +-
 tensor2tensor/layers/common_hparams.py        |   31 +-
 .../layers/common_image_attention.py          |   35 +-
 .../layers/common_image_attention_test.py     |    1 -
 tensor2tensor/layers/modalities.py            | 2476 ++++++++---------
 tensor2tensor/layers/modalities_test.py       |   26 +-
 tensor2tensor/models/image_transformer.py     |   28 +-
 tensor2tensor/models/image_transformer_2d.py  |    5 +-
 tensor2tensor/models/mtf_transformer.py       |    9 +-
 tensor2tensor/models/mtf_transformer2.py      |    9 +-
 tensor2tensor/models/research/autoencoders.py |   18 +-
 tensor2tensor/models/research/cycle_gan.py    |    9 +-
 tensor2tensor/models/research/super_lm.py     |    8 +-
 .../models/research/transformer_symshard.py   |    9 +-
 .../models/research/transformer_vae.py        |    3 +-
 tensor2tensor/models/video/base.py            |   24 +-
 .../video/basic_deterministic_params.py       |   15 +-
 tensor2tensor/models/video/epva_params.py     |   12 +-
 tensor2tensor/models/video/next_frame_glow.py |   12 +-
 tensor2tensor/models/video/savp_params.py     |   17 +-
 tensor2tensor/models/video/sv2p_params.py     |   19 +-
 tensor2tensor/utils/metrics.py                |    8 +-
 tensor2tensor/utils/modality.py               |  121 +
 tensor2tensor/utils/t2t_model.py              |   52 +-
 tensor2tensor/utils/t2t_model_test.py         |    2 -
 25 files changed, 1436 insertions(+), 1521 deletions(-)
 create mode 100644 tensor2tensor/utils/modality.py

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index fe534f050..13ec711ae 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -438,8 +438,8 @@ def aggregate_task_losses(hparams,
     vocab_size += (-vocab_size) % hparams.vocab_divisor
   modality = problem_hparams.modality[feature_name]
   loss = hparams.loss.get(feature_name, modalities.get_loss(modality))
-  weights_fn = hparams.weights_fn.get(
-      feature_name, modalities.get_weights_fn(modality))
+  weights_fn = hparams.targets_weights_fn.get(
+      feature_name, modalities.get_targets_weights_fn(modality))
   # Primary task loss
   loss_num, loss_den = loss(
       logits, feature,
@@ -534,8 +534,8 @@ def aggregate_task_lm_losses(hparams,
     vocab_size += (-vocab_size) % hparams.vocab_divisor
   modality = problem_hparams.modality[feature_name]
   loss = hparams.loss.get(feature_name, modalities.get_loss(modality))
-  weights_fn = hparams.weights_fn.get(
-      feature_name, modalities.get_weights_fn(modality))
+  weights_fn = hparams.targets_weights_fn.get(
+      feature_name, modalities.get_targets_weights_fn(modality))
   loss_num = 0.
   loss_den = 0.
   for task in hparams.problem.task_list:
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 8d3f23eaf..654ee4829 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -141,6 +141,7 @@ def basic_params1():
       norm_type="layer",  # "batch", layer", "noam", "none".
       # epsilon parameter to normalization function
       norm_epsilon=1e-6,
+      symbol_modality_num_shards=1,
       # pad vocabularies so that this value divides the vocabulary size.
       vocab_divisor=1,
       # During training, we drop sequences whose inputs and targets are shorter
@@ -175,26 +176,20 @@ def basic_params1():
       # If True, run the model autoregressively instead of teacher-forcing
       # during eval
       eval_run_autoregressive=False,
-      # (For features with symbol modality) If True, share all of the
-      # input embeddings, target embeddings, and softmax weights.
+      # TODO(lukaszkaiser): these parameters should probably be set elsewhere.
+      # (SymbolModality) - If this flag is on, we try to share all of the input
+      # embeddings, the target embeddings and the softmax weights.
       shared_embedding_and_softmax_weights=False,
-      # (For features with symbol modality) If True, share the input embeddings
-      # and target embeddings.
+      # (SymbolModality) - If this flag is on, we try to share the input
+      # embeddings and the target embeddings.
+      # You can also share the input embeddings with the target embeddings
+      # by using a problem_hparams that uses the same modality object for
+      # the input modality and target modality.
       shared_embedding=False,
-      # (For features with symbol modality) Number to shard embeddings by.
-      symbol_modality_num_shards=1,
-      # Feature transformations are optional dictionaries comprising key-value
-      # pairs of a feature name (str) and its transformation (function). If not
-      # specified, T2TModel applies a default transformation according to the
-      # feature's modality. Bottom is applicable to all features; loss, top, and
-      # weights_fn are only applicable to target features.
-      # TODO(trandustin): `name` is an optional hparam for legacy reasons,
-      # defining variable scope names. Remove this hparam in the future.
-      bottom={},
-      loss={},
-      name={},
-      top={},
-      weights_fn={},
+      # Modalities used to map from features to a space compatible with
+      # chosen model architecture. It comprises key-value pairs of a feature
+      # name (str) and its modality type.
+      modality={},
       # The maximum length of "input" sequence.
       # Sequences longer than this value will be truncated. 0 or negative values
       # mean there is no maximum or truncation.
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index ddc0ecd9d..4f5ea7b67 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -603,9 +603,30 @@ def prepare_decoder(targets, hparams):
 
 def prepare_image(inputs, hparams, name=None):
   """Prepare image."""
-  # TODO(trandustin): This is a legacy function. Remove its usage.
-  del hparams, name  # unused arg
-  return inputs
+  inputs_shape = common_layers.shape_list(inputs)
+  batch = inputs_shape[0]
+  orig_rows = inputs_shape[1]
+  orig_cols = inputs_shape[2]
+  channels = hparams.num_channels
+
+  hidden_size = hparams.hidden_size
+  # TODO(trandustin): Check via modalities.ModalityType.IDENTITY and not str.
+  # The current implementation is to avoid circular imports, modalities ->
+  # discretization -> common_image_attention -> modalities.
+  if "targets" in hparams.modality:
+    target_modality_name = hparams.modality["targets"]
+    if not isinstance(target_modality_name, str):
+      target_modality_name = target_modality_name.__name__
+  else:
+    target_modality_name = None
+  if target_modality_name == "IdentityModality":
+    inputs = tf.to_int32(inputs)
+    x = get_channel_embeddings(channels, inputs, hidden_size, name=name)
+  else:
+    x = inputs
+  x = tf.reshape(x, [batch, orig_rows, orig_cols * channels, hidden_size])
+
+  return x
 
 
 def create_output(decoder_output, rows, cols, targets, hparams):
@@ -626,19 +647,17 @@ def create_output(decoder_output, rows, cols, targets, hparams):
     [batch, hparams.img_len, hparams.img_len, hparams.num_channels, 256].
     In the special case of predict mode, it is a Tensor of rank 5.
   """
-  del targets  # unused arg
   decoded_image = postprocess_image(decoder_output, rows, cols, hparams)
-  batch = common_layers.shape_list(decoded_image)[0]
   depth = common_layers.shape_list(decoded_image)[-1]
+  batch, height, width, channels = common_layers.shape_list(targets)
   likelihood = getattr(hparams, "likelihood", DistributionType.CAT)
   if hparams.mode == tf.estimator.ModeKeys.PREDICT:
     y = tf.reshape(decoded_image, [batch, -1, 1, 1, depth])
-    output = y[:, :rows, :, :, :]
+    output = y[:, :height, :, :, :]
   elif likelihood == DistributionType.CAT:
     # Unpack the cols dimension of the Categorical.
-    channels = hparams.num_channels
     output = tf.reshape(decoded_image,
-                        [batch, rows, cols // channels, channels, depth])
+                        [batch, height, width, channels, depth])
   else:
     output = decoded_image
   return output
diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index fae6806fa..1d22927b6 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -94,7 +94,6 @@ def testCreateOutputTrainMode(self, likelihood, num_mixtures, depth):
     hparams = HParams(
         hidden_size=2,
         likelihood=likelihood,
-        num_channels=channels,
         mode=tf.estimator.ModeKeys.TRAIN,
         num_mixtures=num_mixtures,
     )
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index a8fbb29f6..9a87aee1b 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -13,24 +13,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Modalities, which specify a feature's domain.
-
-T2TModel applies a default transformation to each feature according to its
-modality. Override them by specifying a model's
-hparams.{bottom,loss,name,top,weights_fn}.
-"""
+"""Modalities define the bottom and top of the model (not the body)."""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import sys
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_audio
-from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 from tensor2tensor.layers import discretization
+from tensor2tensor.utils import modality
 
 import tensorflow as tf
 import tensorflow_probability as tfp
@@ -56,1440 +52,1266 @@ def is_pointwise(func):
   return func
 
 
-def generic_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-  """Compute loss numerator and denominator for one shard of output."""
-  del vocab_size  # unused arg
-  logits = top_out
-  logits = common_attention.maybe_upcast(logits, hparams=model_hparams)
-  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.0)
-  return common_layers.padded_cross_entropy(
-      logits,
-      targets,
-      model_hparams.label_smoothing,
-      cutoff=cutoff,
-      weights_fn=weights_fn)
-
+class SymbolModality(modality.Modality):
+  """Modality for sets of discrete symbols.
 
-def make_targets_bottom(bottom):
-  def targets_bottom(x, model_hparams, vocab_size):
-    with tf.variable_scope("targets_bottom"):
-      return bottom(x, model_hparams, vocab_size)
-  return targets_bottom
+  Input:
+    Embedding.
 
+  Output:
+    Linear transformation + softmax.
+  """
 
-def get_weights(model_hparams, vocab_size, hidden_dim=None):
-  """Create or get concatenated embedding or softmax variable.
+  @staticmethod
+  def name(model_hparams, vocab_size):
+    return "symbol_modality_%d_%d" % (vocab_size, model_hparams.hidden_size)
 
-  Args:
-    model_hparams: tf.HParams, model hyperparmeters.
-    vocab_size: int, vocabulary size.
-    hidden_dim: dim of the variable. Defaults to _model_hparams' hidden_size
+  targets_weights_fn = staticmethod(common_layers.weights_nonzero)
 
-  Returns:
-     a list of num_shards Tensors.
-  """
-  if hidden_dim is None:
-    hidden_dim = model_hparams.hidden_size
-  num_shards = model_hparams.symbol_modality_num_shards
-  shards = []
-  for i in range(num_shards):
-    shard_size = (vocab_size // num_shards) + (
-        1 if i < vocab_size % num_shards else 0)
-    var_name = "weights_%d" % i
-    shards.append(
-        tf.get_variable(
-            var_name, [shard_size, hidden_dim],
-            initializer=tf.random_normal_initializer(0.0, hidden_dim**-0.5)))
-  if num_shards == 1:
-    ret = shards[0]
-  else:
-    ret = tf.concat(shards, 0)
-  # Convert ret to tensor.
-  if not tf.executing_eagerly():
-    ret = common_layers.convert_gradient_to_tensor(ret)
-  return ret
-
-
-def _symbol_bottom_simple(x, model_hparams, vocab_size, name, reuse):
-  """Bottom transformation for symbols."""
-  with tf.variable_scope(name, reuse=reuse):
-    # Ensure the inputs are 3-D
-    if len(x.get_shape()) == 4:
-      x = tf.squeeze(x, axis=3)
-    while len(x.get_shape()) < 3:
-      x = tf.expand_dims(x, axis=-1)
-
-    var = get_weights(model_hparams, vocab_size)
-    x = common_layers.dropout_no_scaling(
-        x, 1.0 - model_hparams.symbol_dropout)
-    ret = common_layers.gather(var, x)
-    if model_hparams.multiply_embedding_mode == "sqrt_depth":
-      ret *= model_hparams.hidden_size**0.5
-    ret *= tf.expand_dims(
-        common_layers.cast_like(tf.not_equal(x, 0), ret), -1)
+  @staticmethod
+  def _get_weights(model_hparams, vocab_size, hidden_dim=None):
+    """Create or get concatenated embedding or softmax variable.
+
+    Args:
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+      hidden_dim: dim of the variable. Defaults to _model_hparams' hidden_size
+
+    Returns:
+       a list of num_shards Tensors.
+    """
+    if hidden_dim is None:
+      hidden_dim = model_hparams.hidden_size
+    num_shards = model_hparams.symbol_modality_num_shards
+    shards = []
+    for i in range(num_shards):
+      shard_size = (vocab_size // num_shards) + (
+          1 if i < vocab_size % num_shards else 0)
+      var_name = "weights_%d" % i
+      shards.append(
+          tf.get_variable(
+              var_name, [shard_size, hidden_dim],
+              initializer=tf.random_normal_initializer(0.0, hidden_dim**-0.5)))
+    if num_shards == 1:
+      ret = shards[0]
+    else:
+      ret = tf.concat(shards, 0)
+    # Convert ret to tensor.
+    if not tf.executing_eagerly():
+      ret = common_layers.convert_gradient_to_tensor(ret)
     return ret
 
-
-def symbol_bottom(x, model_hparams, vocab_size):
-  if (model_hparams.shared_embedding_and_softmax_weights or
-      model_hparams.get("shared_embedding")):
-    return _symbol_bottom_simple(
-        x, model_hparams, vocab_size, "shared", reuse=None)
-  return _symbol_bottom_simple(
-      x, model_hparams, vocab_size, "input_emb", reuse=None)
-
-
-def symbol_targets_bottom(x, model_hparams, vocab_size):
-  """Bottom transformation for target symbols."""
-  if (model_hparams.shared_embedding_and_softmax_weights or
-      model_hparams.get("shared_embedding")):
-    try:
-      return _symbol_bottom_simple(
-          x, model_hparams, vocab_size, "shared", reuse=True)
-    except ValueError:
-      # perhaps there were no inputs, and this is a new variable.
-      return _symbol_bottom_simple(
+  @classmethod
+  def bottom_simple(cls, x, model_hparams, vocab_size, name, reuse):
+    with tf.variable_scope(name, reuse=reuse):
+      # Ensure the inputs are 3-D
+      if len(x.get_shape()) == 4:
+        x = tf.squeeze(x, axis=3)
+      while len(x.get_shape()) < 3:
+        x = tf.expand_dims(x, axis=-1)
+
+      var = cls._get_weights(model_hparams, vocab_size)
+      x = common_layers.dropout_no_scaling(
+          x, 1.0 - model_hparams.symbol_dropout)
+      ret = common_layers.gather(var, x)
+      if model_hparams.multiply_embedding_mode == "sqrt_depth":
+        ret *= model_hparams.hidden_size**0.5
+      ret *= tf.expand_dims(
+          common_layers.cast_like(tf.not_equal(x, 0), ret), -1)
+      return ret
+
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    if (model_hparams.shared_embedding_and_softmax_weights or
+        model_hparams.get("shared_embedding")):
+      return cls.bottom_simple(
           x, model_hparams, vocab_size, "shared", reuse=None)
-  else:
-    return _symbol_bottom_simple(
-        x, model_hparams, vocab_size, "target_emb", reuse=None)
+    return cls.bottom_simple(
+        x, model_hparams, vocab_size, "input_emb", reuse=None)
+
+  @classmethod
+  def targets_bottom(cls, x, model_hparams, vocab_size):
+    if (model_hparams.shared_embedding_and_softmax_weights or
+        model_hparams.get("shared_embedding")):
+      try:
+        return cls.bottom_simple(
+            x, model_hparams, vocab_size, "shared", reuse=True)
+      except ValueError:
+        # perhaps there were no inputs, and this is a new variable.
+        return cls.bottom_simple(
+            x, model_hparams, vocab_size, "shared", reuse=None)
+    else:
+      return cls.bottom_simple(
+          x, model_hparams, vocab_size, "target_emb", reuse=None)
+
+  @classmethod
+  @is_pointwise
+  def top(cls, body_output, targets, model_hparams, vocab_size):
+    """Generate logits.
+
+    Args:
+      body_output: A Tensor with shape
+        [batch, p0, p1, model_hparams.hidden_size].
+      targets: Unused.
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+
+    Returns:
+      logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].
+    """
+    del targets  # unused arg
+    if model_hparams.shared_embedding_and_softmax_weights:
+      scope_name = "shared"
+      reuse = tf.AUTO_REUSE
+    else:
+      scope_name = "softmax"
+      reuse = False
+    with tf.variable_scope(scope_name, reuse=reuse):
+      body_output_shape = common_layers.shape_list(body_output)
+      var = cls._get_weights(model_hparams, vocab_size, body_output_shape[-1])
+      if (model_hparams.factored_logits and
+          model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
+        # insert channels dimension
+        body_output = tf.expand_dims(body_output, 3)
+        return common_layers.FactoredTensor(body_output, var)
+      else:
+        body_output = tf.reshape(body_output, [-1, body_output_shape[-1]])
+        logits = tf.matmul(body_output, var, transpose_b=True)
+        return tf.reshape(logits,
+                          body_output_shape[:-1] + [1, vocab_size])
 
 
-@is_pointwise
-def symbol_top(body_output, targets, model_hparams, vocab_size):
-  """Generate logits.
+class SymbolModalityWeightsAll(SymbolModality):
+  """SymbolModality for features that do not have 0-padding."""
 
-  Args:
-    body_output: A Tensor with shape
-      [batch, p0, p1, model_hparams.hidden_size].
-    targets: Unused.
-    model_hparams: tf.HParams, model hyperparmeters.
-    vocab_size: int, vocabulary size.
+  targets_weights_fn = staticmethod(common_layers.weights_all)
 
-  Returns:
-    logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].
-  """
-  del targets  # unused arg
-  if model_hparams.shared_embedding_and_softmax_weights:
-    scope_name = "shared"
-    reuse = tf.AUTO_REUSE
-  else:
-    scope_name = "softmax"
-    reuse = False
-  with tf.variable_scope(scope_name, reuse=reuse):
-    body_output_shape = common_layers.shape_list(body_output)
-    var = get_weights(model_hparams, vocab_size, body_output_shape[-1])
-    if (model_hparams.factored_logits and
-        model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
-      # insert channels dimension
-      body_output = tf.expand_dims(body_output, 3)
-      return common_layers.FactoredTensor(body_output, var)
-    else:
-      body_output = tf.reshape(body_output, [-1, body_output_shape[-1]])
-      logits = tf.matmul(body_output, var, transpose_b=True)
-      return tf.reshape(logits,
-                        body_output_shape[:-1] + [1, vocab_size])
-
-
-def symbol_one_hot_bottom(x, model_hparams, vocab_size):
-  del model_hparams  # unused arg
-  return tf.one_hot(x, vocab_size)
-
-
-@is_pointwise
-def symbol_one_hot_top(body_output, targets, model_hparams, vocab_size):
-  del targets, model_hparams, vocab_size  # unused arg
-  return body_output
-
-
-def symbol_one_hot_loss(top_out,
-                        targets,
-                        model_hparams,
-                        vocab_size,
-                        weights_fn):
-  del model_hparams, weights_fn  # unused arg
-  labels = tf.one_hot(targets, vocab_size)
-  loss = tf.nn.softmax_cross_entropy_with_logits(
-      logits=top_out, labels=labels)
-  return tf.reduce_mean(loss), tf.constant(1.0)
-
-
-def ctc_symbol_loss(top_out, targets, model_hparams, vocab_size, weight_fn):
-  """Compute the CTC loss."""
-  del model_hparams, vocab_size  # unused arg
-  logits = top_out
-  with tf.name_scope("ctc_loss", values=[logits, targets]):
-    # For CTC we assume targets are 1d, [batch, length, 1, 1] here.
-    targets_shape = targets.get_shape().as_list()
-    assert len(targets_shape) == 4
-    assert targets_shape[2] == 1
-    assert targets_shape[3] == 1
-    targets = tf.squeeze(targets, axis=[2, 3])
-    logits = tf.squeeze(logits, axis=[2, 3])
-    targets_mask = 1 - tf.to_int32(tf.equal(targets, 0))
-    targets_lengths = tf.reduce_sum(targets_mask, axis=1)
-    sparse_targets = tf.keras.backend.ctc_label_dense_to_sparse(
-        targets, targets_lengths)
-    xent = tf.nn.ctc_loss(
-        sparse_targets,
-        logits,
-        targets_lengths,
-        time_major=False,
-        preprocess_collapse_repeated=False,
-        ctc_merge_repeated=False)
-    weights = weight_fn(targets)
-    return tf.reduce_sum(xent), tf.reduce_sum(weights)
 
+class SymbolModalityOneHot(SymbolModality):
+  """Simple SymbolModality with one hot as embeddings."""
 
-def image_bottom(x, model_hparams, vocab_size):
-  del model_hparams, vocab_size  # unused arg
-  with tf.variable_scope("image_modality"):
-    if not tf.executing_eagerly():
-      tf.summary.image(
-          "inputs", common_layers.tpu_safe_image_summary(x), max_outputs=2)
-    return tf.to_float(x)
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size):
+    return tf.one_hot(x, vocab_size)
 
+  @staticmethod
+  def targets_bottom(x, model_hparams, vocab_size):
+    return tf.one_hot(x, vocab_size)
 
-def image_targets_bottom(x, model_hparams, vocab_size):
-  """Bottom transformation for target images."""
-  pixel_embedding_size = 64
-  inputs = x
-  with tf.variable_scope("image_modality"):
-    if not tf.executing_eagerly():
-      tf.summary.image(
-          "targets_bottom",
-          common_layers.tpu_safe_image_summary(inputs),
-          max_outputs=1)
-    inputs_shape = common_layers.shape_list(inputs)
-    if len(inputs_shape) != 4:
-      raise ValueError("Assuming images given as int tensors in the format "
-                       "[batch, height, width, channels] (256 values).")
-    # We embed each of 256=vocab_size possible pixel values.
-    embedding_var = tf.get_variable(
-        "pixel_embedding",
-        [vocab_size, pixel_embedding_size])
-    hot_inputs = tf.one_hot(tf.to_int32(inputs), vocab_size)
-    hot_inputs = tf.reshape(hot_inputs, [-1, vocab_size])
-    embedded = tf.matmul(hot_inputs, embedding_var)
-    # Let's now merge all channels that were embedded into a single vector.
-    merged_size = pixel_embedding_size * inputs_shape[3]
-    embedded = tf.reshape(embedded, inputs_shape[:3] + [merged_size])
-    merged = tf.layers.dense(
-        embedded,
-        model_hparams.hidden_size,
-        name="merge_pixel_embedded_channels")
-    return merged
-
-
-def image_top(body_output, targets, model_hparams, vocab_size):
-  """Top transformation for images."""
-  del targets  # unused arg
-  # TODO(lukaszkaiser): is this a universal enough way to get channels?
-  num_channels = model_hparams.problem.num_channels
-  with tf.variable_scope("rgb_softmax"):
-    body_output_shape = common_layers.shape_list(body_output)
-    reshape_shape = body_output_shape[:3]
-    reshape_shape.extend([num_channels, vocab_size])
-    res = tf.layers.dense(body_output, vocab_size * num_channels)
-    res = tf.reshape(res, reshape_shape)
-    if not tf.get_variable_scope().reuse:
-      res_argmax = tf.argmax(res, axis=-1)
-      tf.summary.image(
-          "result",
-          common_layers.tpu_safe_image_summary(res_argmax),
-          max_outputs=1)
-    return res
+  @staticmethod
+  @is_pointwise
+  def top(body_output, _, model_hparams, vocab_size):
+    return body_output
 
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+    del weights_fn  # unused arg
+    labels = tf.one_hot(targets, vocab_size)
+    loss = tf.nn.softmax_cross_entropy_with_logits(
+        logits=top_out, labels=labels)
+    return tf.reduce_mean(loss), tf.constant(1.0)
+
+
+class CTCSymbolModality(SymbolModality):
+  """SymbolModality that uses CTC loss."""
+
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
+    """Compute the CTC loss."""
+    logits = top_out
+    with tf.name_scope("ctc_loss", values=[logits, targets]):
+      # For CTC we assume targets are 1d, [batch, length, 1, 1] here.
+      targets_shape = targets.get_shape().as_list()
+      assert len(targets_shape) == 4
+      assert targets_shape[2] == 1
+      assert targets_shape[3] == 1
+      targets = tf.squeeze(targets, axis=[2, 3])
+      logits = tf.squeeze(logits, axis=[2, 3])
+      targets_mask = 1 - tf.to_int32(tf.equal(targets, 0))
+      targets_lengths = tf.reduce_sum(targets_mask, axis=1)
+      sparse_targets = tf.keras.backend.ctc_label_dense_to_sparse(
+          targets, targets_lengths)
+      xent = tf.nn.ctc_loss(
+          sparse_targets,
+          logits,
+          targets_lengths,
+          time_major=False,
+          preprocess_collapse_repeated=False,
+          ctc_merge_repeated=False)
+      weights = weights_fn(targets)
+      return tf.reduce_sum(xent), tf.reduce_sum(weights)
+
+
+class ImageModality(modality.Modality):
+  """Modality for images."""
+
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      if not tf.executing_eagerly():
+        tf.summary.image(
+            "inputs", common_layers.tpu_safe_image_summary(x), max_outputs=2)
+      return tf.to_float(x)
+
+  @classmethod
+  def targets_bottom(cls, x, model_hparams, vocab_size):
+    pixel_embedding_size = 64
+    inputs = x
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      if not tf.executing_eagerly():
+        tf.summary.image(
+            "targets_bottom",
+            common_layers.tpu_safe_image_summary(inputs),
+            max_outputs=1)
+      inputs_shape = common_layers.shape_list(inputs)
+      if len(inputs_shape) != 4:
+        raise ValueError("Assuming images given as int tensors in the format "
+                         "[batch, height, width, channels] (256 values).")
+      # We embed each of 256=vocab_size possible pixel values.
+      embedding_var = tf.get_variable(
+          "pixel_embedding",
+          [vocab_size, pixel_embedding_size])
+      hot_inputs = tf.one_hot(tf.to_int32(inputs), vocab_size)
+      hot_inputs = tf.reshape(hot_inputs, [-1, vocab_size])
+      embedded = tf.matmul(hot_inputs, embedding_var)
+      # Let's now merge all channels that were embedded into a single vector.
+      merged_size = pixel_embedding_size * inputs_shape[3]
+      embedded = tf.reshape(embedded, inputs_shape[:3] + [merged_size])
+      merged = tf.layers.dense(
+          embedded,
+          model_hparams.hidden_size,
+          name="merge_pixel_embedded_channels")
+      return merged
 
-def _image_channel_compress_bottom(inputs, model_hparams, name="bottom"):
-  """Compresses channel-wise input pixels into whole pixel representions.
+  @staticmethod
+  def top(body_output, _, model_hparams, vocab_size):
+    # TODO(lukaszkaiser): is this a universal enough way to get channels?
+    num_channels = model_hparams.problem.num_channels
+    with tf.variable_scope("rgb_softmax"):
+      body_output_shape = common_layers.shape_list(body_output)
+      reshape_shape = body_output_shape[:3]
+      reshape_shape.extend([num_channels, vocab_size])
+      res = tf.layers.dense(body_output, vocab_size * num_channels)
+      res = tf.reshape(res, reshape_shape)
+      if not tf.get_variable_scope().reuse:
+        res_argmax = tf.argmax(res, axis=-1)
+        tf.summary.image(
+            "result",
+            common_layers.tpu_safe_image_summary(res_argmax),
+            max_outputs=1)
+      return res
+
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
+    """Compute loss numerator and denominator for one shard of output."""
+    logits = top_out
+    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.0)
+    return common_layers.padded_cross_entropy(
+        logits,
+        targets,
+        model_hparams.label_smoothing,
+        cutoff=cutoff,
+        weights_fn=weights_fn)
 
-  Perform conversion of RGB pixel values to a real number in the range -1 to
-  1. This combines pixel channels to form a representation of shape
-  [img_len, img_len].
 
-  Args:
-    inputs: Tensor representing RGB pixel intensities as integers, of shape
-      [batch, img_len, img_len, channels].
-    model_hparams: tf.HParams, model hyperparmeters.
-    name: string, scope.
+class ImageChannelCompressModality(modality.Modality):
+  """Modality for images using channel compression for generation."""
 
-  Returns:
-    body_input: Tensor of shape
-      [batch, img_len, img_len, model_hparams.hidden_size].
-  """
-  num_channels = 3
-  with tf.variable_scope(name):
-    inputs = tf.to_float(inputs)
-    hp = model_hparams
-    if hp.mode != tf.estimator.ModeKeys.PREDICT:
-      tf.summary.image(
-          "inputs",
-          common_layers.tpu_safe_image_summary(inputs),
-          max_outputs=2)
-    inputs = common_layers.convert_rgb_to_symmetric_real(inputs)
-
-    # Reshape inputs to apply convolutions across [img_len, img_len*channels].
-    inputs_shape = common_layers.shape_list(inputs)
-    inputs = tf.reshape(
-        inputs, [-1, inputs_shape[1], inputs_shape[2] * inputs_shape[3], 1])
-
-    # Compress RGB intensities for each pixel using a convolution.
-    outputs = tf.layers.conv2d(
-        inputs,
-        model_hparams.hidden_size,
-        kernel_size=(1, num_channels),
-        padding="VALID",
-        strides=(1, num_channels),
-        activation=tf.nn.relu,
-        name="conv_input")
-    return outputs
-
-
-def image_channel_compress_bottom(x, model_hparams, vocab_size):
-  del vocab_size  # unused arg
-  return _image_channel_compress_bottom(x, model_hparams, "input_bottom")
-
-
-def image_channel_compress_targets_bottom(x, model_hparams, vocab_size):
-  del vocab_size  # unused arg
-  return _image_channel_compress_bottom(x, model_hparams, "output_bottom")
-
-
-def image_channel_compress_top(body_output, targets, model_hparams, vocab_size):
-  """Transforms body output to return logits.
+  @staticmethod
+  def bottom_compress(inputs, model_hparams, name="bottom"):
+    """Compresses channel-wise input pixels into whole pixel representions.
+
+    Perform conversion of RGB pixel values to a real number in the range -1 to
+    1. This combines pixel channels to form a representation of shape
+    [img_len, img_len].
+
+    Args:
+      inputs: Tensor representing RGB pixel intensities as integers, of shape
+        [batch, img_len, img_len, channels].
+      model_hparams: tf.HParams, model hyperparmeters.
+      name: string, scope.
+
+    Returns:
+      body_input: Tensor of shape
+        [batch, img_len, img_len, model_hparams.hidden_size].
+    """
+    num_channels = 3
+    with tf.variable_scope(name):
+      inputs = tf.to_float(inputs)
+      hp = model_hparams
+      if hp.mode != tf.estimator.ModeKeys.PREDICT:
+        tf.summary.image(
+            "inputs",
+            common_layers.tpu_safe_image_summary(inputs),
+            max_outputs=2)
+      inputs = common_layers.convert_rgb_to_symmetric_real(inputs)
+
+      # Reshape inputs to apply convolutions across [img_len, img_len*channels].
+      inputs_shape = common_layers.shape_list(inputs)
+      inputs = tf.reshape(
+          inputs, [-1, inputs_shape[1], inputs_shape[2] * inputs_shape[3], 1])
+
+      # Compress RGB intensities for each pixel using a convolution.
+      outputs = tf.layers.conv2d(
+          inputs,
+          model_hparams.hidden_size,
+          kernel_size=(1, num_channels),
+          padding="VALID",
+          strides=(1, num_channels),
+          activation=tf.nn.relu,
+          name="conv_input")
+      return outputs
+
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    return cls.bottom_compress(x, model_hparams, "input_bottom")
+
+  @classmethod
+  def targets_bottom(cls, x, model_hparams, vocab_size):
+    return cls.bottom_compress(x, model_hparams, "output_bottom")
+
+  @classmethod
+  def top(cls, body_output, _, model_hparams, vocab_size):
+    """Transforms body output to return logits.
+
+    Args:
+      body_output: Tensor of shape [batch, img_len, img_len, depth].
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+
+    Returns:
+      Tensor of shape [batch, img_len, img_len, channels, vocab_size].
+    """
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      hidden_size = model_hparams.hidden_size
+      img_len = model_hparams.img_len
+      channels = 3  # RGB
+      batch = common_layers.shape_list(body_output)[0]
+      x = tf.layers.conv2d(
+          body_output,
+          hidden_size * channels,
+          kernel_size=(1, 1),
+          strides=(1, 1),
+          padding="VALID",
+          activation=tf.nn.relu,
+          name="decompress_conv")
+      x = tf.reshape(x, [batch, img_len, img_len * channels, hidden_size])
+      x = common_layers.layer_preprocess(x, model_hparams)
+      x = tf.layers.dense(x,
+                          vocab_size,
+                          use_bias=True,
+                          activation=None,
+                          name="output_conv")
+      x = tf.reshape(
+          x, [batch, img_len, img_len, channels, vocab_size])
+      return x
 
-  Args:
-    body_output: Tensor of shape [batch, img_len, img_len, depth].
-    targets:
-    model_hparams: tf.HParams, model hyperparmeters.
-    vocab_size: int, vocabulary size.
 
-  Returns:
-    Tensor of shape [batch, img_len, img_len, channels, vocab_size].
-  """
-  del targets  # unused arg
-  with tf.variable_scope("image_channel_compress_modality"):
+class ImageChannelBottomIdentityModality(ImageChannelCompressModality):
+
+  @staticmethod
+  def top(body_output, _, model_hparams, vocab_size):
+    return body_output
+
+
+class ImageChannelEmbeddingsBottom(modality.Modality):
+  """Modality for images using channel compression for generation."""
+
+  @staticmethod
+  def get_channel_embeddings(io_depth,
+                             targets,
+                             hidden_size,
+                             name="channel"):
+    """Get separate embedding for each of the channels."""
+    targets_split = tf.split(targets, io_depth, axis=3)
+    rgb_embedding_var = tf.get_variable("rgb_target_emb_%s" % name,
+                                        [256 * io_depth, hidden_size])
+    rgb_embedding_var = tf.identity(rgb_embedding_var)
+    rgb_embedding_var *= float(hidden_size)**0.5
+    channel_target_embs = []
+    for i in range(io_depth):
+      # Adding the channel offsets to get the right embedding since the
+      # embedding tensor has shape 256 * io_depth, hidden_size
+      target_ids = tf.squeeze(targets_split[i], axis=3) + i * 256
+      target_embs = common_layers.gather(rgb_embedding_var, target_ids)
+      channel_target_embs.append(target_embs)
+
+    return tf.concat(channel_target_embs, axis=-1)
+
+  @staticmethod
+  def targets_bottom(x, model_hparams, vocab_size):
+    inputs = x
+    io_depth = model_hparams.num_channels
+    tshape = common_layers.shape_list(inputs)
     hidden_size = model_hparams.hidden_size
-    img_len = model_hparams.img_len
-    channels = 3  # RGB
-    batch = common_layers.shape_list(body_output)[0]
-    x = tf.layers.conv2d(
-        body_output,
-        hidden_size * channels,
-        kernel_size=(1, 1),
-        strides=(1, 1),
-        padding="VALID",
-        activation=tf.nn.relu,
-        name="decompress_conv")
-    x = tf.reshape(x, [batch, img_len, img_len * channels, hidden_size])
-    x = common_layers.layer_preprocess(x, model_hparams)
-    x = tf.layers.dense(x,
-                        vocab_size,
-                        use_bias=True,
-                        activation=None,
-                        name="output_conv")
-    x = tf.reshape(
-        x, [batch, img_len, img_len, channels, vocab_size])
+    target_embeddings = ImageChannelEmbeddingsBottom.get_channel_embeddings(
+        io_depth, inputs, hidden_size, "input_bottom")
+    return tf.reshape(target_embeddings,
+                      [tshape[0], tshape[1], tshape[2] * io_depth, hidden_size])
+
+  @classmethod
+  def top(cls, body_output, _, model_hparams, vocab_size):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      img_len = model_hparams.img_len
+      channels = model_hparams.num_channels
+      x = tf.layers.dense(
+          body_output, 256, use_bias=True, activation=None, name="output_conv")
+      x = tf.reshape(x,
+                     [-1, img_len, img_len, channels, vocab_size])
+      return x
+
+
+class AudioModality(modality.Modality):
+  """Performs strided conv compressions for audio data."""
+
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    """Transform input from data space to model space.
+
+    Args:
+      x: A Tensor with shape [batch, ...]
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+
+    Returns:
+      body_input: A Tensor with shape [batch, ?, ?,
+        model_hparams.hidden_size].
+    """
+    inputs = x
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      # TODO(aidangomez): Will need to sort out a better audio pipeline
+      def xnet_resblock(x, filters, res_relu, name):
+        """Xception block."""
+        with tf.variable_scope(name):
+          # Typically audio samples are >100k samples in length and have a width
+          # of 2 or 4. Mono audio has a single channel while stereo has 2.
+          y = common_layers.separable_conv_block(
+              x,
+              filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
+              first_relu=True,
+              padding="SAME",
+              force2d=True,
+              name="sep_conv_block")
+          y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 2))
+          return y + common_layers.conv_block(
+              x,
+              filters, [((1, 1), (1, 1))],
+              padding="SAME",
+              strides=(2, 2),
+              first_relu=res_relu,
+              force2d=True,
+              name="res_conv0")
+
+      x = tf.to_float(inputs) / 255.
+      x.set_shape([None, None, None, 1])
+      for i in range(model_hparams.audio_compression):
+        x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
+      return xnet_resblock(x,
+                           model_hparams.hidden_size,
+                           False,
+                           "compress_block_final")
+
+
+class AudioSpectralModality(modality.Modality):
+  """Performs strided conv compressions for audio spectral data."""
+
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    """Transform input from data space to model space.
+
+    Args:
+      x: A Tensor with shape [batch, ...]
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+
+    Returns:
+      body_input: A Tensor with shape [batch, ?, ?,
+        model_hparams.hidden_size].
+    """
+    inputs = x
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      # TODO(aidangomez): Will need to sort out a better audio pipeline
+      def xnet_resblock(x, filters, res_relu, name):
+        """Xception-like block."""
+        with tf.variable_scope(name):
+          # We only stride along the length dimension to preserve the spectral
+          # bins (which are tiny in dimensionality relative to length)
+          y = common_layers.separable_conv_block(
+              x,
+              filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
+              first_relu=True,
+              padding="SAME",
+              force2d=True,
+              name="sep_conv_block")
+          y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 1))
+          return y + common_layers.conv_block(
+              x,
+              filters, [((1, 1), (1, 1))],
+              padding="SAME",
+              strides=(2, 1),
+              first_relu=res_relu,
+              force2d=True,
+              name="res_conv0")
+
+      # Bitcast back from int32
+      x = tf.bitcast(inputs, tf.float32)
+      x.set_shape([None, None, None, 1])
+      for i in range(model_hparams.audio_compression):
+        x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
+      return xnet_resblock(x,
+                           model_hparams.hidden_size,
+                           False,
+                           "compress_block_final")
+
+
+class SpeechRecognitionModality(modality.Modality):
+  """Common ASR filterbank processing."""
+
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    """Use batchnorm instead of CMVN and shorten the stft with strided convs.
+
+    Args:
+      x: float32 tensor with shape [batch_size, len, 1, freqs * channels]
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+
+    Returns:
+      float32 tensor with shape [batch_size, shorter_len, 1, hidden_size]
+    """
+    inputs = x
+    p = model_hparams
+
+    num_mel_bins = p.audio_num_mel_bins
+    num_channels = 3 if p.audio_add_delta_deltas else 1
+
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      if p.audio_preproc_in_bottom:
+        # Compute filterbanks
+        with tf.variable_scope("fbanks"):
+          waveforms = tf.squeeze(inputs, [2, 3])
+          mel_fbanks = common_audio.compute_mel_filterbank_features(
+              waveforms,
+              sample_rate=p.audio_sample_rate,
+              dither=p.audio_dither,
+              preemphasis=p.audio_preemphasis,
+              frame_length=p.audio_frame_length,
+              frame_step=p.audio_frame_step,
+              lower_edge_hertz=p.audio_lower_edge_hertz,
+              upper_edge_hertz=p.audio_upper_edge_hertz,
+              num_mel_bins=p.audio_num_mel_bins,
+              apply_mask=True)
+          if p.audio_add_delta_deltas:
+            mel_fbanks = common_audio.add_delta_deltas(mel_fbanks)
+          x = tf.reshape(mel_fbanks,
+                         common_layers.shape_list(mel_fbanks)[:2] +
+                         [num_mel_bins, num_channels])
+
+          nonpadding_mask = 1. - common_attention.embedding_to_padding(x)
+          num_of_nonpadding_elements = tf.reduce_sum(
+              nonpadding_mask) * num_mel_bins * num_channels
+
+          # This replaces CMVN estimation on data
+          var_epsilon = 1e-09
+          mean = tf.reduce_sum(
+              x, axis=[1], keepdims=True) / num_of_nonpadding_elements
+          variance = (num_of_nonpadding_elements * mean**2. -
+                      2. * mean * tf.reduce_sum(x, axis=[1], keepdims=True) +
+                      tf.reduce_sum(x**2, axis=[1], keepdims=True)
+                     ) / num_of_nonpadding_elements
+          x = (x - mean) * tf.rsqrt(variance + var_epsilon) * tf.expand_dims(
+              nonpadding_mask, -1)
+      else:
+        x = inputs
+
+      # The convention is that the models are flattened along the spatial,
+      # dimensions, thus the speech preprocessor treats frequencies and
+      # channels as image colors (last axis)
+      x.set_shape([None, None, num_mel_bins, num_channels])
+
+      # TODO(chorowski): how to specify bottom's hparams and avoid hardcoding?
+      x = tf.pad(x, [[0, 0], [0, 8], [0, 0], [0, 0]])
+      for _ in range(2):
+        x = tf.layers.conv2d(
+            x, 128, (3, 3), (2, 2), use_bias=False)
+        x = common_layers.layer_norm(x)
+        x = tf.nn.relu(x)
+
+      xshape = common_layers.shape_list(x)
+      # apply a conv that will remove all frequencies and at the same time
+      # project the output into desired hidden_size
+      x = tf.pad(x, [[0, 0], [0, 2], [0, 0], [0, 0]])
+      x = tf.layers.conv2d(x, p.hidden_size, (3, xshape[2]), use_bias=False)
+
+      assert common_layers.shape_list(x)[2] == 1
+      x = common_layers.layer_norm(x)
+      x = tf.nn.relu(x)
     return x
 
 
-def image_channel_embeddings_bottom(x, model_hparams, vocab_size):
-  """Bottom transformation for image targets."""
-  del vocab_size  # unused arg
-  inputs = tf.to_int32(x)
-  io_depth = model_hparams.num_channels
-  tshape = common_layers.shape_list(inputs)
-  hidden_size = model_hparams.hidden_size
-  target_embeddings = cia.get_channel_embeddings(
-      io_depth, inputs, hidden_size, "input_bottom")
-  return tf.reshape(target_embeddings,
-                    [tshape[0], tshape[1], tshape[2] * io_depth, hidden_size])
-
-
-def image_channel_embeddings_top(body_output,
-                                 targets,
-                                 model_hparams,
-                                 vocab_size):
-  """Top transformation for images."""
-  del targets  # unused arg
-  with tf.variable_scope("image_channel_embeddings_bottom"):
-    img_len = model_hparams.img_len
-    channels = model_hparams.num_channels
-    x = tf.layers.dense(
-        body_output, 256, use_bias=True, activation=None, name="output_conv")
-    x = tf.reshape(x,
-                   [-1, img_len, img_len, channels, vocab_size])
+class VideoModality(modality.Modality):
+  """Modality for videos, i.e., time-sequences of frames."""
+
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size):
+    common_video.gif_summary("inputs", x, max_outputs=1)
+    x = common_layers.standardize_images(x)
     return x
 
+  @staticmethod
+  def targets_bottom(x, model_hparams, vocab_size):
+    common_video.gif_summary("targets", x, max_outputs=1)
+    x = common_layers.standardize_images(x)
+    return x
 
-def audio_bottom(x, model_hparams, vocab_size):
-  """Transform input from data space to model space.
+  @staticmethod
+  def top(body_output, targets, model_hparams, vocab_size):
+    num_channels = model_hparams.problem.num_channels
+    shape = common_layers.shape_list(body_output)
+    reshape_shape = shape[:-1] + [num_channels, vocab_size]
+    res = tf.reshape(body_output, reshape_shape)
+    # Calculate argmax so as to have a summary with the produced images.
+    x = tf.argmax(tf.reshape(res, [-1, vocab_size]), axis=-1)
+    x = tf.reshape(x, shape[:-1] + [num_channels])
+    common_video.gif_summary("results", x, max_outputs=1)
+    return res
 
-  Args:
-    x: A Tensor with shape [batch, ...]
-    model_hparams: tf.HParams, model hyperparmeters.
-    vocab_size: int, vocabulary size.
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+    """Compute loss numerator and denominator for one shard of output."""
+    logits = top_out
+    logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
+    targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
+    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
+    return common_layers.padded_cross_entropy(
+        logits,
+        targets,
+        model_hparams.label_smoothing,
+        cutoff=cutoff,
+        weights_fn=weights_fn)
+
+
+class VideoModalityBitwise(VideoModality):
+  """Video Modality where bottom embeds pixels bitwise."""
+
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    pixel_embedding_size = 64
+    inputs = x
+    with tf.variable_scope(cls.name(model_hparams, vocab_size),
+                           reuse=tf.AUTO_REUSE):
+      common_layers.summarize_video(inputs, "bottom")
+      # Embed bitwise.
+      assert vocab_size == 256
+      embedded = discretization.int_to_bit_embed(inputs, 8,
+                                                 pixel_embedding_size)
+      # Project.
+      return tf.layers.dense(
+          embedded,
+          model_hparams.hidden_size,
+          name="merge_pixel_embedded_frames")
+
+  @classmethod
+  def targets_bottom(cls, x, model_hparams, vocab_size):  # pylint: disable=arguments-differ
+    pixel_embedding_size = 64
+    inputs = x
+    with tf.variable_scope(cls.name(model_hparams, vocab_size),
+                           reuse=tf.AUTO_REUSE):
+      common_layers.summarize_video(inputs, "targets_bottom")
+      # Embed bitwise.
+      assert vocab_size == 256
+      embedded = discretization.int_to_bit_embed(inputs, 8,
+                                                 pixel_embedding_size)
+      # Transpose and project.
+      transposed = common_layers.time_to_channels(embedded)
+      return tf.layers.dense(
+          transposed,
+          model_hparams.hidden_size,
+          name="merge_pixel_embedded_frames")
+
+
+class VideoModalityPixelNoise(VideoModality):
+  """Video modality that introduces pixel noise on input during training."""
+
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    input_noise = getattr(model_hparams, "video_modality_input_noise", 0.25)
+    inputs = x
+    if model_hparams.mode == tf.estimator.ModeKeys.TRAIN:
+      background = tfp.distributions.percentile(inputs, 50., axis=[0, 1, 2, 3])
+      input_shape = common_layers.shape_list(inputs)
+      input_size = tf.reduce_prod(input_shape[:-1])
+      input_mask = tf.multinomial(
+          tf.log([[input_noise, 1.-input_noise]]), input_size)
+      input_mask = tf.reshape(tf.cast(input_mask, tf.int32),
+                              input_shape[:-1]+[1])
+      inputs = inputs * input_mask + background * (1 - input_mask)
+    return super(VideoModalityPixelNoise, cls).bottom(
+        inputs, model_hparams, vocab_size)
+
+
+class VideoModalityL1(VideoModality):
+  """Video modality that predicts a scalar per channel with an L1 loss."""
 
-  Returns:
-    body_input: A Tensor with shape [batch, ?, ?,
-      model_hparams.hidden_size].
-  """
-  del vocab_size  # unused arg
-  inputs = x
-  with tf.variable_scope("audio_modality"):
-    # TODO(aidangomez): Will need to sort out a better audio pipeline
-    def xnet_resblock(x, filters, res_relu, name):
-      """Xception block."""
-      with tf.variable_scope(name):
-        # Typically audio samples are >100k samples in length and have a width
-        # of 2 or 4. Mono audio has a single channel while stereo has 2.
-        y = common_layers.separable_conv_block(
-            x,
-            filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
-            first_relu=True,
-            padding="SAME",
-            force2d=True,
-            name="sep_conv_block")
-        y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 2))
-        return y + common_layers.conv_block(
-            x,
-            filters, [((1, 1), (1, 1))],
-            padding="SAME",
-            strides=(2, 2),
-            first_relu=res_relu,
-            force2d=True,
-            name="res_conv0")
-
-    x = tf.to_float(inputs) / 255.
-    x.set_shape([None, None, None, 1])
-    for i in range(model_hparams.audio_compression):
-      x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
-    return xnet_resblock(x,
-                         model_hparams.hidden_size,
-                         False,
-                         "compress_block_final")
-
-
-def audio_spectral_bottom(x, model_hparams, vocab_size):
-  """Transform input from data space to model space.
+  @staticmethod
+  def top(body_output, _, model_hparams, vocab_size):
+    num_channels = model_hparams.problem.num_channels
+    num_frames = model_hparams.video_num_target_frames
+    with tf.variable_scope("rgb"):
+      body_output_shape = common_layers.shape_list(body_output)
+      res = tf.layers.dense(body_output, num_channels * num_frames, name="cast")
+      res = tf.reshape(res, body_output_shape[:3] + [num_channels, num_frames])
+      res = tf.transpose(res, [0, 4, 1, 2, 3])  # Move frames next to batch.
+      if not tf.get_variable_scope().reuse:
+        res_argmax = res[:, -1, :, :, :]
+        tf.summary.image(
+            "result",
+            common_layers.tpu_safe_image_summary(res_argmax),
+            max_outputs=1)
+      return tf.expand_dims(res, axis=-1)  # Add an axis like in perplexity.
 
-  Args:
-    x: A Tensor with shape [batch, ...]
-    model_hparams: tf.HParams, model hyperparmeters.
-    vocab_size: int, vocabulary size.
+  @staticmethod
+  def internal_loss(logits, targets, model_hparams):
+    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
+    return tf.nn.relu(tf.abs(logits - targets) - cutoff)
 
-  Returns:
-    body_input: A Tensor with shape [batch, ?, ?,
-      model_hparams.hidden_size].
-  """
-  del vocab_size  # unused arg
-  inputs = x
-  with tf.variable_scope("audio_spectral_modality"):
-    # TODO(aidangomez): Will need to sort out a better audio pipeline
-    def xnet_resblock(x, filters, res_relu, name):
-      """Xception-like block."""
-      with tf.variable_scope(name):
-        # We only stride along the length dimension to preserve the spectral
-        # bins (which are tiny in dimensionality relative to length)
-        y = common_layers.separable_conv_block(
-            x,
-            filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
-            first_relu=True,
-            padding="SAME",
-            force2d=True,
-            name="sep_conv_block")
-        y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 1))
-        return y + common_layers.conv_block(
-            x,
-            filters, [((1, 1), (1, 1))],
-            padding="SAME",
-            strides=(2, 1),
-            first_relu=res_relu,
-            force2d=True,
-            name="res_conv0")
-
-    # Bitcast back from int32
-    x = tf.bitcast(inputs, tf.float32)
-    x.set_shape([None, None, None, 1])
-    for i in range(model_hparams.audio_compression):
-      x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
-    return xnet_resblock(x,
-                         model_hparams.hidden_size,
-                         False,
-                         "compress_block_final")
-
-
-def speech_recognition_bottom(x, model_hparams, vocab_size):
-  """Use batchnorm instead of CMVN and shorten the stft with strided convs.
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+    """Compute loss numerator and denominator for one shard of output."""
+    logits = top_out
+    logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:-1])
+    targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
+    weights = weights_fn(targets)
+    # Shift targets by 0.5 so later just casting to int gives the prediction.
+    # So for int targets, say 0 and 7, we actually train to predict 0.5 and 7.5.
+    # Later (in merics or infer) this is cast to int anyway. Also, we have no
+    # loss beyond cutoff = 0.2 as these are already correct predictions.
+    targets = tf.to_float(targets) + 0.5
+    loss = VideoModalityL1.internal_loss(logits, targets, model_hparams)
+    return tf.reduce_sum(loss * weights), tf.reduce_sum(weights)
 
-  Args:
-    x: float32 tensor with shape [batch_size, len, 1, freqs * channels]
-    model_hparams: tf.HParams, model hyperparmeters.
-    vocab_size: int, vocabulary size.
 
-  Returns:
-    float32 tensor with shape [batch_size, shorter_len, 1, hidden_size]
-  """
-  del vocab_size  # unused arg
-  inputs = x
-  p = model_hparams
-
-  num_mel_bins = p.audio_num_mel_bins
-  num_channels = 3 if p.audio_add_delta_deltas else 1
-
-  with tf.variable_scope("speech_recognition_modality"):
-    if p.audio_preproc_in_bottom:
-      # Compute filterbanks
-      with tf.variable_scope("fbanks"):
-        waveforms = tf.squeeze(inputs, [2, 3])
-        mel_fbanks = common_audio.compute_mel_filterbank_features(
-            waveforms,
-            sample_rate=p.audio_sample_rate,
-            dither=p.audio_dither,
-            preemphasis=p.audio_preemphasis,
-            frame_length=p.audio_frame_length,
-            frame_step=p.audio_frame_step,
-            lower_edge_hertz=p.audio_lower_edge_hertz,
-            upper_edge_hertz=p.audio_upper_edge_hertz,
-            num_mel_bins=p.audio_num_mel_bins,
-            apply_mask=True)
-        if p.audio_add_delta_deltas:
-          mel_fbanks = common_audio.add_delta_deltas(mel_fbanks)
-        x = tf.reshape(mel_fbanks,
-                       common_layers.shape_list(mel_fbanks)[:2] +
-                       [num_mel_bins, num_channels])
-
-        nonpadding_mask = 1. - common_attention.embedding_to_padding(x)
-        num_of_nonpadding_elements = tf.reduce_sum(
-            nonpadding_mask) * num_mel_bins * num_channels
-
-        # This replaces CMVN estimation on data
-        var_epsilon = 1e-09
-        mean = tf.reduce_sum(
-            x, axis=[1], keepdims=True) / num_of_nonpadding_elements
-        variance = (num_of_nonpadding_elements * mean**2. -
-                    2. * mean * tf.reduce_sum(x, axis=[1], keepdims=True) +
-                    tf.reduce_sum(x**2, axis=[1], keepdims=True)
-                   ) / num_of_nonpadding_elements
-        x = (x - mean) * tf.rsqrt(variance + var_epsilon) * tf.expand_dims(
-            nonpadding_mask, -1)
-    else:
-      x = inputs
+class VideoModalityL2(VideoModalityL1):
+  """Modality for videos with L2 loss."""
 
-    # The convention is that the models are flattened along the spatial,
-    # dimensions, thus the speech preprocessor treats frequencies and
-    # channels as image colors (last axis)
-    x.set_shape([None, None, num_mel_bins, num_channels])
+  @staticmethod
+  def internal_loss(logits, targets, model_hparams):
+    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
+    return tf.nn.relu(
+        tf.squared_difference(logits, targets) - cutoff * cutoff)
 
-    # TODO(chorowski): how to specify bottom's hparams and avoid hardcoding?
-    x = tf.pad(x, [[0, 0], [0, 8], [0, 0], [0, 0]])
-    for _ in range(2):
-      x = tf.layers.conv2d(
-          x, 128, (3, 3), (2, 2), use_bias=False)
-      x = common_layers.layer_norm(x)
-      x = tf.nn.relu(x)
-
-    xshape = common_layers.shape_list(x)
-    # apply a conv that will remove all frequencies and at the same time
-    # project the output into desired hidden_size
-    x = tf.pad(x, [[0, 0], [0, 2], [0, 0], [0, 0]])
-    x = tf.layers.conv2d(x, p.hidden_size, (3, xshape[2]), use_bias=False)
-
-    assert common_layers.shape_list(x)[2] == 1
-    x = common_layers.layer_norm(x)
-    x = tf.nn.relu(x)
-  return x
-
-
-def video_bottom(x, model_hparams, vocab_size):
-  del model_hparams, vocab_size  # unused arg
-  common_video.gif_summary("inputs", x, max_outputs=1)
-  x = common_layers.standardize_images(x)
-  return x
-
-
-def video_targets_bottom(x, model_hparams, vocab_size):
-  del model_hparams, vocab_size  # unused arg
-  common_video.gif_summary("targets", x, max_outputs=1)
-  x = common_layers.standardize_images(x)
-  return x
-
-
-def video_top(body_output, targets, model_hparams, vocab_size):
-  """Top transformation for video."""
-  del targets  # unused arg
-  num_channels = model_hparams.problem.num_channels
-  shape = common_layers.shape_list(body_output)
-  reshape_shape = shape[:-1] + [num_channels, vocab_size]
-  res = tf.reshape(body_output, reshape_shape)
-  # Calculate argmax so as to have a summary with the produced images.
-  x = tf.argmax(tf.reshape(res, [-1, vocab_size]), axis=-1)
-  x = tf.reshape(x, shape[:-1] + [num_channels])
-  common_video.gif_summary("results", x, max_outputs=1)
-  return res
-
-
-def video_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-  """Compute loss numerator and denominator for one shard of output."""
-  del vocab_size  # unused arg
-  logits = top_out
-  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
-  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
-  return common_layers.padded_cross_entropy(
-      logits,
-      targets,
-      model_hparams.label_smoothing,
-      cutoff=cutoff,
-      weights_fn=weights_fn)
-
-
-def video_bitwise_bottom(x, model_hparams, vocab_size):
-  """Bottom transformation for embedding video bitwise."""
-  pixel_embedding_size = 64
-  inputs = x
-  with tf.variable_scope("video_modality_bitwise", reuse=tf.AUTO_REUSE):
-    common_layers.summarize_video(inputs, "bottom")
-    # Embed bitwise.
-    assert vocab_size == 256
-    embedded = discretization.int_to_bit_embed(inputs, 8,
-                                               pixel_embedding_size)
-    # Project.
-    return tf.layers.dense(
-        embedded,
-        model_hparams.hidden_size,
-        name="merge_pixel_embedded_frames")
-
-
-def video_bitwise_targets_bottom(x, model_hparams, vocab_size):
-  """Bottom transformation for embedding target video bitwise."""
-  pixel_embedding_size = 64
-  inputs = x
-  with tf.variable_scope("video_modality_bitwise", reuse=tf.AUTO_REUSE):
-    common_layers.summarize_video(inputs, "targets_bottom")
-    # Embed bitwise.
-    assert vocab_size == 256
-    embedded = discretization.int_to_bit_embed(inputs, 8,
-                                               pixel_embedding_size)
-    # Transpose and project.
-    transposed = common_layers.time_to_channels(embedded)
-    return tf.layers.dense(
-        transposed,
-        model_hparams.hidden_size,
-        name="merge_pixel_embedded_frames")
-
-
-def video_pixel_noise_bottom(x, model_hparams, vocab_size):
-  """Bottom transformation for video."""
-  input_noise = getattr(model_hparams, "video_modality_input_noise", 0.25)
-  inputs = x
-  if model_hparams.mode == tf.estimator.ModeKeys.TRAIN:
-    background = tfp.distributions.percentile(inputs, 50., axis=[0, 1, 2, 3])
-    input_shape = common_layers.shape_list(inputs)
-    input_size = tf.reduce_prod(input_shape[:-1])
-    input_mask = tf.multinomial(
-        tf.log([[input_noise, 1.-input_noise]]), input_size)
-    input_mask = tf.reshape(tf.cast(input_mask, tf.int32),
-                            input_shape[:-1]+[1])
-    inputs = inputs * input_mask + background * (1 - input_mask)
-  return video_bottom(inputs, model_hparams, vocab_size)
-
-
-def video_l1_top(body_output, targets, model_hparams, vocab_size):
-  """Top transformation for video."""
-  del targets, vocab_size  # unused arg
-  num_channels = model_hparams.problem.num_channels
-  num_frames = model_hparams.video_num_target_frames
-  with tf.variable_scope("rgb"):
-    body_output_shape = common_layers.shape_list(body_output)
-    res = tf.layers.dense(body_output, num_channels * num_frames, name="cast")
-    res = tf.reshape(res, body_output_shape[:3] + [num_channels, num_frames])
-    res = tf.transpose(res, [0, 4, 1, 2, 3])  # Move frames next to batch.
-    if not tf.get_variable_scope().reuse:
-      res_argmax = res[:, -1, :, :, :]
-      tf.summary.image(
-          "result",
-          common_layers.tpu_safe_image_summary(res_argmax),
-          max_outputs=1)
-    return tf.expand_dims(res, axis=-1)  # Add an axis like in perplexity.
-
-
-def video_l1_internal_loss(logits, targets, model_hparams):
-  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
-  return tf.nn.relu(tf.abs(logits - targets) - cutoff)
-
-
-def video_l1_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-  """Compute loss numerator and denominator for one shard of output."""
-  del vocab_size  # unused arg
-  logits = top_out
-  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:-1])
-  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-  weights = weights_fn(targets)
-  # Shift targets by 0.5 so later just casting to int gives the prediction.
-  # So for int targets, say 0 and 7, we actually train to predict 0.5 and 7.5.
-  # Later (in merics or infer) this is cast to int anyway. Also, we have no
-  # loss beyond cutoff = 0.2 as these are already correct predictions.
-  targets = tf.to_float(targets) + 0.5
-  loss = video_l1_internal_loss(logits, targets, model_hparams)
-  return tf.reduce_sum(loss * weights), tf.reduce_sum(weights)
-
-
-def video_l2_internal_loss(logits, targets, model_hparams):
-  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
-  return tf.nn.relu(
-      tf.squared_difference(logits, targets) - cutoff * cutoff)
-
-
-def video_l2_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-  """Compute loss numerator and denominator for one shard of output."""
-  del vocab_size  # unused arg
-  logits = top_out
-  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:-1])
-  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-  weights = weights_fn(targets)
-  # Shift targets by 0.5 so later just casting to int gives the prediction.
-  # So for int targets, say 0 and 7, we actually train to predict 0.5 and 7.5.
-  # Later (in merics or infer) this is cast to int anyway. Also, we have no
-  # loss beyond cutoff = 0.2 as these are already correct predictions.
-  targets = tf.to_float(targets) + 0.5
-  loss = video_l2_internal_loss(logits, targets, model_hparams)
-  return tf.reduce_sum(loss * weights), tf.reduce_sum(weights)
-
-
-def convert_rgb_to_real(prediction, targets):
-  """Convert prediction and target from rgb to real."""
-  prediction = tf.squeeze(prediction, axis=-1)
-  prediction = common_layers.convert_rgb_to_real(prediction)
-  targets = common_layers.convert_rgb_to_real(targets)
-  return prediction, targets
-
-
-def video_raw_bottom(x, model_hparams, vocab_size):
-  del model_hparams, vocab_size  # unused arg
-  common_video.gif_summary("inputs", x)
-  return common_layers.convert_rgb_to_real(x)
-
-
-def video_raw_targets_bottom(x, model_hparams, vocab_size):
-  del model_hparams, vocab_size  # unused arg
-  common_video.gif_summary("targets_bottom", x)
-  return common_layers.convert_rgb_to_real(x)
-
-
-def video_raw_top(body_output, targets, model_hparams, vocab_size):
-  del targets, model_hparams, vocab_size  # unused arg
-  frames = body_output
-  if isinstance(body_output, list):
-    frames = tf.stack(body_output, axis=1)
-  rgb_frames = common_layers.convert_real_to_rgb(frames)
-  common_video.gif_summary("body_output", rgb_frames)
-  return tf.expand_dims(rgb_frames, axis=-1)
-
-
-def video_l2_raw_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-  del model_hparams, vocab_size, weights_fn  # unused arg
-  prediction, groundtruth = convert_rgb_to_real(top_out, targets)
-  loss = tf.losses.mean_squared_error(prediction, groundtruth)
-  return loss, tf.constant(1.0)
-
-
-def video_l1_raw_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-  del model_hparams, vocab_size, weights_fn  # unused arg
-  prediction, groundtruth = convert_rgb_to_real(top_out, targets)
-  loss = tf.losses.absolute_difference(prediction, groundtruth)
-  return loss, tf.constant(1.0)
-
-
-def class_label_bottom(x, model_hparams, vocab_size):
-  with tf.variable_scope("class_label_modality_%d_%d" % (
-      vocab_size, model_hparams.hidden_size)):
-    multiplier = 1.0
-    if model_hparams.multiply_embedding_mode == "sqrt_depth":
-      multiplier = model_hparams.hidden_size**0.5
-    return common_layers.embedding(x,
-                                   vocab_size,
-                                   model_hparams.hidden_size,
-                                   multiplier=multiplier)
 
+class VideoModalityL2Raw(VideoModalityL2):
+  """Modality with L2 loss and raw input (sequences of frames)."""
 
-def class_label_targets_bottom(x, model_hparams, vocab_size):
-  with tf.variable_scope("class_label_modality_%d_%d" % (
-      vocab_size, model_hparams.hidden_size)):
-    return tf.zeros([common_layers.shape_list(x)[0],
-                     1,
-                     1,
-                     model_hparams.hidden_size])
-
-
-def class_label_top(body_output, targets, model_hparams, vocab_size):
-  """Transform inputs from model space to target space.
-
-  Average over inner dims and a linear layer to logits.
+  @staticmethod
+  def convert_rgb_to_real(prediction, targets):
+    """Convert prediction and target from rgb to real."""
+    prediction = tf.squeeze(prediction, axis=-1)
+    prediction = common_layers.convert_rgb_to_real(prediction)
+    targets = common_layers.convert_rgb_to_real(targets)
+    return prediction, targets
 
-  Args:
-    body_output: A Tensor with shape [batch, ?, ?, body_output_size].
-    targets:
-    model_hparams: tf.HParams, model hyperparmeters.
-    vocab_size: int, vocabulary size.
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size):
+    common_video.gif_summary("inputs", x)
+    return common_layers.convert_rgb_to_real(x)
 
-  Returns:
-    a Tensors, each with shape [batch_size, 1, 1, 1, vocab_size]
-  """
-  del targets  # unused arg
-  with tf.variable_scope("class_label_modality_%d_%d" % (
-      vocab_size, model_hparams.hidden_size)):
-    x = body_output
-    x = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
-    res = tf.layers.dense(x, vocab_size)
-    return tf.expand_dims(res, 3)
-
-
-def video_identity_bottom(x, model_hparams, vocab_size):
-  del model_hparams, vocab_size  # unused arg
-  common_video.gif_summary("inputs", x, max_outputs=1)
-  return x
-
-
-def video_identity_targets_bottom(x, model_hparams, vocab_size):
-  del model_hparams, vocab_size  # unused arg
-  common_video.gif_summary("targets", x, max_outputs=1)
-  return x
-
-
-def video_identity_loss(top_out,
-                        targets,
-                        model_hparams,
-                        vocab_size,
-                        weights_fn):
-  """Compute loss numerator and denominator for one shard of output."""
-  del vocab_size  # unused arg
-  # TODO(nikip): Try L2 loss
-  logits = top_out
-  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
-  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
-  return common_layers.padded_cross_entropy(
-      logits,
-      targets,
-      model_hparams.label_smoothing,
-      cutoff=cutoff,
-      weights_fn=weights_fn)
-
-
-def multi_label_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-  """Average loss over the labels."""
-  del vocab_size  # unused arg
-  logits = top_out
-  num_labels = tf.shape(targets)[1]
-  logits = tf.tile(logits, [1, num_labels, 1, 1, 1])
-
-  xent, weights = common_layers.padded_cross_entropy(
-      logits,
-      targets,
-      model_hparams.label_smoothing,
-      weights_fn=weights_fn,
-      reduce_sum=False,
-  )
-  xent = tf.squeeze(xent, [2, 3])
-  weights = tf.squeeze(weights, [2, 3])
-  # average loss over all labels
-  loss = tf.reduce_sum(xent, axis=1)
-  weights = tf.reduce_sum(weights, axis=1)
-  loss /= (weights + 1e-8)
-  weights = tf.to_float(tf.greater(weights, 0.))
-
-  return tf.reduce_sum(loss*weights), tf.reduce_sum(weights)
-
-
-def one_hot_class_label_loss(top_out,
-                             targets,
-                             model_hparams,
-                             vocab_size,
-                             weights_fn):
-  """Apply softmax cross-entropy between outputs and targets.
+  @staticmethod
+  def targets_bottom(x, model_hparams, vocab_size):  # pylint: disable=arguments-differ
+    common_video.gif_summary("targets_bottom", x)
+    return common_layers.convert_rgb_to_real(x)
 
-  Args:
-    top_out: logits Tensor with shape [batch, ?, ?, num_classes]
-    targets: one-hot encoding Tensor with shape [batch, ?, ?, num_classes]
-    model_hparams: tf.HParams, model hyperparmeters.
-    vocab_size: int, vocabulary size.
-    weights_fn:
+  @staticmethod
+  def top(body_output, _, model_hparams, vocab_size):
+    frames = body_output
+    if isinstance(body_output, list):
+      frames = tf.stack(body_output, axis=1)
+    rgb_frames = common_layers.convert_real_to_rgb(frames)
+    common_video.gif_summary("body_output", rgb_frames)
+    return tf.expand_dims(rgb_frames, axis=-1)
 
-  Returns:
-    loss_scale (cross-entropy), loss_denom
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
+    del weights_fn  # unused arg
+    prediction, groundtruth = cls.convert_rgb_to_real(top_out, targets)
+    loss = tf.losses.mean_squared_error(prediction, groundtruth)
+    return loss, tf.constant(1.0)
+
+
+class VideoModalityL1Raw(VideoModalityL2Raw):
+  """Modality with L1 loss and raw input (sequences of frames)."""
+
+  @classmethod
+  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
+    prediction, groundtruth = cls.convert_rgb_to_real(top_out, targets)
+    loss = tf.losses.absolute_difference(prediction, groundtruth)
+    return loss, tf.constant(1.0)
+
+
+class ClassLabelModality(modality.Modality):
+  """Used for label data."""
+
+  @staticmethod
+  def name(model_hparams, vocab_size):
+    return "class_label_modality_%d_%d" % (vocab_size,
+                                           model_hparams.hidden_size)
+
+  @classmethod
+  def bottom(cls, x, model_hparams, vocab_size):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      multiplier = 1.0
+      if model_hparams.multiply_embedding_mode == "sqrt_depth":
+        multiplier = model_hparams.hidden_size**0.5
+      return common_layers.embedding(x,
+                                     vocab_size,
+                                     model_hparams.hidden_size,
+                                     multiplier=multiplier)
+
+  @classmethod
+  def targets_bottom(cls, x, model_hparams, vocab_size):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      return tf.zeros([common_layers.shape_list(x)[0],
+                       1,
+                       1,
+                       model_hparams.hidden_size])
+
+  @classmethod
+  def top(cls, body_output, _, model_hparams, vocab_size):
+    """Transform inputs from model space to target space.
+
+    Average over inner dims and a linear layer to logits.
+
+    Args:
+      body_output: A Tensor with shape [batch, ?, ?, body_output_size].
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+
+    Returns:
+      a Tensors, each with shape [batch_size, 1, 1, 1, vocab_size]
+    """
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      x = body_output
+      x = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
+      res = tf.layers.dense(x, vocab_size)
+      return tf.expand_dims(res, 3)
+
+
+class VideoModalityIdentity(VideoModality):
+  """Video Modality where top and bottom is an identity function."""
+
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size):
+    common_video.gif_summary("inputs", x, max_outputs=1)
+    return x
+
+  @staticmethod
+  def targets_bottom(x, model_hparams, vocab_size):
+    common_video.gif_summary("targets", x, max_outputs=1)
+    return x
+
+  @staticmethod
+  def top(body_output, targets, model_hparams, vocab_size):
+    return body_output
+
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+    """Compute loss numerator and denominator for one shard of output."""
+    # TODO(nikip): Try L2 loss
+    logits = top_out
+    logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
+    targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
+    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
+    return common_layers.padded_cross_entropy(
+        logits,
+        targets,
+        model_hparams.label_smoothing,
+        cutoff=cutoff,
+        weights_fn=weights_fn)
+
+
+class MultiLabelModality(ClassLabelModality):
+  """Used for multi label task."""
+
+  targets_weights_fn = staticmethod(common_layers.weights_nonzero)
+
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+    """Average loss over the labels."""
+    logits = top_out
+    num_labels = tf.shape(targets)[1]
+    logits = tf.tile(logits, [1, num_labels, 1, 1, 1])
+
+    xent, weights = common_layers.padded_cross_entropy(
+        logits,
+        targets,
+        model_hparams.label_smoothing,
+        weights_fn=weights_fn,
+        reduce_sum=False,
+    )
+    xent = tf.squeeze(xent, [2, 3])
+    weights = tf.squeeze(weights, [2, 3])
+    # average loss over all labels
+    loss = tf.reduce_sum(xent, axis=1)
+    weights = tf.reduce_sum(weights, axis=1)
+    loss /= (weights + 1e-8)
+    weights = tf.to_float(tf.greater(weights, 0.))
+
+    return tf.reduce_sum(loss*weights), tf.reduce_sum(weights)
+
+
+class OneHotClassLabelModality(ClassLabelModality):
+  """Used for one-hot encoded class labels."""
+
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+    """Apply softmax cross-entropy between outputs and targets.
+
+    Args:
+      top_out: logits Tensor with shape [batch, ?, ?, num_classes]
+      targets: one-hot encoding Tensor with shape [batch, ?, ?, num_classes]
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+      weights_fn: Function mapping targets to weights.
+
+    Returns:
+      loss_scale (cross-entropy), loss_denom
+    """
+    loss_scale = tf.losses.softmax_cross_entropy(
+        onehot_labels=targets, logits=top_out)
+    weights = weights_fn(targets)
+    loss_denom = tf.reduce_sum(weights)
+    return loss_scale, loss_denom
+
+
+class IdentityModality(modality.Modality):
+  """Does nothing."""
+
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size):
+    return tf.to_float(x)
+
+  @staticmethod
+  def top(body_output, _, model_hparams, vocab_size):
+    return body_output
+
+
+class GenericL2LossModality(IdentityModality):
+  """Generic modality with L2 as Loss."""
+
+  @staticmethod
+  def targets_bottom(x, model_hparams, vocab_size):
+    return tf.to_float(x)
+
+  @staticmethod
+  def loss(body_output, targets, model_hparams, vocab_size, weights_fn):
+    del weights_fn  # unused
+    loss = tf.squared_difference(body_output, tf.to_float(targets))
+    return tf.reduce_mean(loss), tf.constant(1.0)
+
+
+class RealModality(modality.Modality):
+  """Base class for real (i.e. float) vectors.
+
+  * Bottom is a linear projection layer to hparams.hidden_size.
+  * Top is a linear projection layer to vocab_size.
   """
-  del model_hparams, vocab_size  # unused arg
-  loss_scale = tf.losses.softmax_cross_entropy(
-      onehot_labels=targets, logits=top_out)
-  weights = weights_fn(targets)
-  loss_denom = tf.reduce_sum(weights)
-  return loss_scale, loss_denom
 
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size):
+    with tf.variable_scope("real"):
+      return tf.layers.dense(
+          tf.to_float(x), model_hparams.hidden_size, name="bottom")
 
-def identity_bottom(x, model_hparams, vocab_size):
-  del model_hparams, vocab_size  # unused arg
-  return tf.to_float(x)
+  @staticmethod
+  @is_pointwise
+  def top(body_output, _, model_hparams, vocab_size):
+    with tf.variable_scope("real"):
+      return tf.layers.dense(body_output, vocab_size, name="top")
 
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+    raise NotImplementedError()
 
-def identity_top(body_output, targets, model_hparams, vocab_size):
-  del targets, model_hparams, vocab_size  # unused arg
-  return body_output
 
+class RealL2LossModality(RealModality):
+  """Modality for real (i.e. float) vectors with L2 (Gaussian) loss."""
 
-def generic_l2_loss(body_output,
-                    targets,
-                    model_hparams,
-                    vocab_size,
-                    weights_fn):
-  del model_hparams, vocab_size, weights_fn  # unused arg
-  loss = tf.squared_difference(body_output, tf.to_float(targets))
-  return tf.reduce_mean(loss), tf.constant(1.0)
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+    predictions = top_out
+    if (len(common_layers.shape_list(top_out)) != len(
+        common_layers.shape_list(targets))):
+      predictions = tf.squeeze(top_out, axis=[-1])
+    with tf.name_scope("l2"):
+      weights = weights_fn(targets)
+      l2 = tf.pow(predictions - targets, 2)
+      return tf.reduce_sum(l2 * weights), tf.reduce_sum(weights)
 
 
-def real_bottom(x, model_hparams, vocab_size):
-  del vocab_size  # unused arg
-  with tf.variable_scope("real"):
-    return tf.layers.dense(
-        tf.to_float(x), model_hparams.hidden_size, name="bottom")
+class RealLogPoissonLossModality(RealModality):
+  """Modality for real (i.e. float) vectors with log Poisson regression loss."""
 
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+    predictions = top_out
+    if (len(common_layers.shape_list(top_out)) != len(
+        common_layers.shape_list(targets))):
+      predictions = tf.squeeze(top_out, axis=[-1])
+    with tf.name_scope("log_possion"):
+      weights = weights_fn(targets)
+      lp_loss = tf.nn.log_poisson_loss(targets, predictions)
+      return tf.reduce_sum(lp_loss * weights), tf.reduce_sum(weights)
+
+
+class IdentitySymbolModality(SymbolModality):
+  """Symbol modality with identity top and bottom transformations.
+
+  Uses the weights_fn from SymbolModality so that loss/metrics ignore padding.
+  """
+
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size):
+    return tf.to_float(x)
+
+  @staticmethod
+  def top(body_output, _, model_hparams, vocab_size):
+    return body_output
+
+  @classmethod
+  def targets_bottom(cls, x, model_hparams, vocab_size):
+    """SymbolModality overrides targets_bottom, so need to override here too."""
+    return cls.bottom(x, model_hparams, vocab_size)
 
-@is_pointwise
-def real_top(body_output, targets, model_hparams, vocab_size):
-  del targets, model_hparams  # unused arg
-  with tf.variable_scope("real"):
-    return tf.layers.dense(body_output, vocab_size, name="top")
 
+class SigmoidClassLabelModality(ClassLabelModality):
+  """Sigmoid cross-entropy for independent class labels."""
 
-def real_l2_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-  del model_hparams, vocab_size  # unused arg
-  predictions = top_out
-  if (len(common_layers.shape_list(top_out)) != len(
-      common_layers.shape_list(targets))):
-    predictions = tf.squeeze(top_out, axis=[-1])
-  with tf.name_scope("l2"):
+  @staticmethod
+  def name(model_hparams, vocab_size):
+    return "sigmoid_class_symbol_modality_%d_%d" % (vocab_size,
+                                                    model_hparams.hidden_size)
+
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+    # Expect inputs of size [batch-size, timesteps, 1, num-classes], where the
+    # last dimension of num-classes represents logits for binary labels
+    loss_scale = tf.losses.sigmoid_cross_entropy(
+        multi_class_labels=targets, logits=top_out)
     weights = weights_fn(targets)
-    l2 = tf.pow(predictions - targets, 2)
-    return tf.reduce_sum(l2 * weights), tf.reduce_sum(weights)
+    loss_denom = tf.reduce_sum(weights)
+    return loss_scale, loss_denom
 
 
-def real_log_poisson_loss(top_out,
-                          targets,
-                          model_hparams,
-                          vocab_size,
-                          weights_fn):
-  """Poisson loss for real."""
-  del model_hparams, vocab_size  # unused arg
-  predictions = top_out
-  if (len(common_layers.shape_list(top_out)) != len(
-      common_layers.shape_list(targets))):
-    predictions = tf.squeeze(top_out, axis=[-1])
-  with tf.name_scope("log_possion"):
+class SigmoidMaxPoolingClassLabelModality(ClassLabelModality):
+  """Sigmoid cross-entropy applied on max-pooling over timesteps."""
+
+  @staticmethod
+  def name(model_hparams, vocab_size):
+    return "sigmoid_max_pooling_class_symbol_modality_%d_%d" % (
+        vocab_size, model_hparams.hidden_size)
+
+  @classmethod
+  def top(cls, body_output, _, model_hparams, vocab_size):
+    """Transform inputs from model space to target space.
+
+    Average over inner dims and a linear layer to logits.
+
+    Args:
+      body_output: A Tensor with shape [batch, timesteps, 1, body_output_size].
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+
+    Returns:
+      a Tensors, each with shape [batch_size, 1, 1, vocab_size]
+    """
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      x = body_output
+      x = tf.reduce_max(x, axis=1, keepdims=True)
+      return tf.layers.dense(x, vocab_size)
+
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+    # Expect inputs of size [batch-size, 1, 1, num-classes], where the
+    # last dimension of num-classes represents logits for binary labels
+    loss_scale = tf.losses.sigmoid_cross_entropy(
+        multi_class_labels=targets, logits=top_out)
     weights = weights_fn(targets)
-    lp_loss = tf.nn.log_poisson_loss(targets, predictions)
-    return tf.reduce_sum(lp_loss * weights), tf.reduce_sum(weights)
+    loss_denom = tf.reduce_sum(weights)
+    return loss_scale, loss_denom
 
 
-def sigmoid_class_label_loss(top_out,
-                             targets,
-                             model_hparams,
-                             vocab_size,
-                             weights_fn):
-  """Loss for class label."""
-  # Expect inputs of size [batch-size, timesteps, 1, num-classes], where the
-  # last dimension of num-classes represents logits for binary labels
-  del model_hparams, vocab_size  # unused arg
-  loss_scale = tf.losses.sigmoid_cross_entropy(
-      multi_class_labels=targets, logits=top_out)
-  weights = weights_fn(targets)
-  loss_denom = tf.reduce_sum(weights)
-  return loss_scale, loss_denom
-
-
-def sigmoid_max_pooling_class_label_top(body_output,
-                                        targets,
-                                        model_hparams,
-                                        vocab_size):
-  """Transform inputs from model space to target space.
-
-  Average over inner dims and a linear layer to logits.
+class SoftmaxMaxPoolingClassLabelModality(OneHotClassLabelModality):
+  """Softmax cross-entropy applied on max-pooling over timesteps."""
 
-  Args:
-    body_output: A Tensor with shape [batch, timesteps, 1, body_output_size].
-    targets:
-    model_hparams: tf.HParams, model hyperparmeters.
-    vocab_size: int, vocabulary size.
+  @staticmethod
+  def name(model_hparams, vocab_size):
+    return "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
+        vocab_size, model_hparams.hidden_size)
 
-  Returns:
-    a Tensors, each with shape [batch_size, 1, 1, vocab_size]
-  """
-  del targets  # unused arg
-  with tf.variable_scope(
-      "sigmoid_max_pooling_class_symbol_modality_%d_%d" % (
-          vocab_size, model_hparams.hidden_size)):
-    x = body_output
-    x = tf.reduce_max(x, axis=1, keepdims=True)
-    return tf.layers.dense(x, vocab_size)
-
-
-def sigmoid_max_pooling_class_label_loss(top_out,
-                                         targets,
-                                         model_hparams,
-                                         vocab_size,
-                                         weights_fn):
-  """Loss for class label."""
-  # Expect inputs of size [batch-size, 1, 1, num-classes], where the
-  # last dimension of num-classes represents logits for binary labels
-  del model_hparams, vocab_size  # unused arg
-  loss_scale = tf.losses.sigmoid_cross_entropy(
-      multi_class_labels=targets, logits=top_out)
-  weights = weights_fn(targets)
-  loss_denom = tf.reduce_sum(weights)
-  return loss_scale, loss_denom
-
-
-def softmax_max_pooling_class_label_top(body_output,
-                                        targets,
-                                        model_hparams,
-                                        vocab_size):
-  """Loss for class label."""
-  del targets  # unused arg
-  with tf.variable_scope(
-      "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
-          vocab_size, model_hparams.hidden_size)):
-    x = body_output
-    x = tf.reduce_max(x, axis=1, keepdims=True)
-    return tf.layers.dense(x, vocab_size)
-
-
-def softmax_average_pooling_class_label_top(body_output,
-                                            targets,
-                                            model_hparams,
-                                            vocab_size):
-  """Loss for class label."""
-  del targets  # unused arg
-  with tf.variable_scope(
-      "softmax_average_pooling_onehot_class_label_modality_%d_%d" % (
-          vocab_size, model_hparams.hidden_size)):
-    x = body_output
-    x = tf.reduce_mean(x, axis=1, keepdims=True)
-    return tf.layers.dense(x, vocab_size)
-
-
-def softmax_last_timestep_class_label_top(body_output,
-                                          targets,
-                                          model_hparams,
-                                          vocab_size):
-  """Loss for class label."""
-  del targets  # unused arg
-  with tf.variable_scope(
-      "softmax_last_timestep_onehot_class_label_modality_%d_%d" % (
-          vocab_size, model_hparams.hidden_size)):
-    x = body_output
-    x = tf.expand_dims(x[:, -1], 1)  # Pick the last timestep
-    return tf.layers.dense(x, vocab_size)
+  @classmethod
+  def top(cls, body_output, _, model_hparams, vocab_size):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      x = body_output
+      x = tf.reduce_max(x, axis=1, keepdims=True)
+      return tf.layers.dense(x, vocab_size)
+
+
+class SoftmaxAveragePoolingClassLabelModality(OneHotClassLabelModality):
+  """Softmax cross-entropy applied on average-pooling over timesteps."""
+
+  @staticmethod
+  def name(model_hparams, vocab_size):
+    return "softmax_average_pooling_onehot_class_label_modality_%d_%d" % (
+        vocab_size, model_hparams.hidden_size)
+
+  @classmethod
+  def top(cls, body_output, _, model_hparams, vocab_size):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      x = body_output
+      x = tf.reduce_mean(x, axis=1, keepdims=True)
+      return tf.layers.dense(x, vocab_size)
+
+
+class SoftmaxLastTimestepClassLabelModality(OneHotClassLabelModality):
+  """Softmax cross-entropy applied on last-timestep encoding."""
+
+  @staticmethod
+  def name(model_hparams, vocab_size):
+    return "softmax_last_timestep_onehot_class_label_modality_%d_%d" % (
+        vocab_size, model_hparams.hidden_size)
+
+  @classmethod
+  def top(cls, body_output, _, model_hparams, vocab_size):
+    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
+      x = body_output
+      x = tf.expand_dims(x[:, -1], 1)  # Pick the last timestep
+      return tf.layers.dense(x, vocab_size)
 
 
 class ModalityType(object):
   """Types of modalities."""
 
-  AUDIO = "audio"
-  AUDIO_SPECTRAL = "audio_spectral"
-  CLASS_LABEL = "class_label"
-  CTC_SYMBOL = "ctc_symbol"  # symbol with CTC loss
-  GENERIC_L2_LOSS = "generic_l2"  # identity modality with L2 loss
-  IDENTITY = "identity"  # identity top and bottom
-  IDENTITY_SYMBOL = "identity_symbol"  # symbol with identity top and bottom
-  IMAGE = "image"
-  # images using channel compression for generation
-  IMAGE_CHANNEL_BOTTOM_IDENTITY = "image_channel_bottom_identity"
-  # images using channel compression for generation
-  IMAGE_CHANNEL_COMPRESS = "image_channel_compress"
-  IMAGE_CHANNEL_EMBEDDINGS_BOTTOM = "image_channel_embeddings_bottom"
-  MULTI_LABEL = "multi_label"
-  ONE_HOT_CLASS_LABEL = "one_hot_class_label"
-  REAL = "real"  # real vectors
-  REAL_L2_LOSS = "real_l2"  # real vectors with L2 as loss
-  # real vectors with log Poisson regression loss
-  REAL_LOG_POISSON_LOSS = "real_log_poisson"
-  SIGMOID_CLASS_LABEL = "sigmoid_class_label"  # sigmoid cross-entropy loss
-  # sigmoid cross-entropy applied on max-pooling over timesteps
-  SIGMOID_MAX_POOLING_CLASS_LABEL = "sigmoid_max_pooling_class_label"
-  # softmax cross-entropy applied on average-pooling over timesteps
-  SOFTMAX_AVERAGE_POOLING_CLASS_LABEL = "softmax_average_pooling_class_label"
-  # softmax cross-entropy applied on last-timestep encoding
-  SOFTMAX_LAST_TIMESTEP_CLASS_LABEL = "softmax_last_timestep_class_label"
-  # softmax cross-entropy applied on max-pooling over timesteps
-  SOFTMAX_MAX_POOLING_CLASS_LABEL = "softmax_max_pooling_class_label"
-  SPEECH_RECOGNITION = "speech_recognition"
-  SYMBOL = "symbol"
-  SYMBOL_WEIGHTS_ALL = "symbol_weights_all"  # symbol for features w/o 0-padding
-  SYMBOL_ONE_HOT = "symbol_one_hot"  # symbol with one hot as embeddings
-  VIDEO = "video"
-  VIDEO_BITWISE = "video_bitwise"  # video where bottom embeds pixels bitwise
-  VIDEO_IDENTITY = "video_identity"  # video with identity top and bottom
-  VIDEO_L1 = "video_l1"  # video with L2 loss
-  VIDEO_L2 = "video_l2"  # video with L1 loss
-  # video with L1 loss and raw input (sequences of frames)
-  VIDEO_L1_RAW = "video_l1_raw"
-  # video with L2 loss and raw input (sequences of frames)
-  VIDEO_L2_RAW = "video_l2_raw"
-  # video with pixel noise on input during training
-  VIDEO_PIXEL_NOISE = "video_pixel_noise"
+  SYMBOL = "SymbolModality"
+  SYMBOL_WEIGHTS_ALL = "SymbolModalityWeightsAll"
+  SYMBOL_ONE_HOT = "SymbolModalityOneHot"
+  CTC_SYMBOL = "CTCSymbolModality"
+  IMAGE = "ImageModality"
+  IMAGE_CHANNEL_COMPRESS = "ImageChannelCompressModality"
+  IMAGE_CHANNEL_BOTTOM_IDENTITY = "ImageChannelBottomIdentityModality"
+  IMAGE_CHANNEL_EMBEDDINGS_BOTTOM = "ImageChannelEmbeddingsBottom"
+  AUDIO = "AudioModality"
+  AUDIO_SPECTRAL = "AudioSpectralModality"
+  SPEECH_RECOGNITION = "SpeechRecognitionModality"
+  VIDEO = "VideoModality"
+  VIDEO_BITWISE = "VideoModalityBitwise"
+  VIDEO_PIXEL_NOISE = "VideoModalityPixelNoise"
+  VIDEO_L1 = "VideoModalityL1"
+  VIDEO_L2 = "VideoModalityL2"
+  VIDEO_L2_RAW = "VideoModalityL2Raw"
+  VIDEO_L1_RAW = "VideoModalityL1Raw"
+  CLASS_LABEL = "ClassLabelModality"
+  VIDEO_IDENTITY = "VideoModalityIdentity"
+  MULTI_LABEL = "MultiLabelModality"
+  ONE_HOT_CLASS_LABEL = "OneHotClassLabelModality"
+  IDENTITY = "IdentityModality"
+  GENERIC_L2_LOSS = "GenericL2LossModality"
+  REAL = "RealModality"
+  REAL_L2_LOSS = "RealL2LossModality"
+  REAL_LOG_POISSON_LOSS = "RealLogPoissonLossModality"
+  IDENTITY_SYMBOL = "IdentitySymbolModality"
+  SIGMOID_CLASS_LABEL = "SigmoidClassLabelModality"
+  SIGMOID_MAX_POOLING_CLASS_LABEL = "SigmoidMaxPoolingClassLabelModality"
+  SOFTMAX_MAX_POOLING_CLASS_LABEL = "SoftmaxMaxPoolingClassLabelModality"
+  SOFTMAX_AVERAGE_POOLING_CLASS_LABEL = "SoftmaxAveragePoolingClassLabelModality"
+  SOFTMAX_LAST_TIMESTEP_CLASS_LABEL = "SoftmaxLastTimestepClassLabelModality"
 
   @staticmethod
   def get_choices():
     return [
-        ModalityType.AUDIO,
-        ModalityType.AUDIO_SPECTRAL,
-        ModalityType.CLASS_LABEL,
+        ModalityType.SYMBOL,
+        ModalityType.SYMBOL_WEIGHTS_ALL,
+        ModalityType.SYMBOL_ONE_HOT,
         ModalityType.CTC_SYMBOL,
-        ModalityType.GENERIC_L2_LOSS,
-        ModalityType.IDENTITY,
-        ModalityType.IDENTITY_SYMBOL,
         ModalityType.IMAGE,
-        ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
         ModalityType.IMAGE_CHANNEL_COMPRESS,
+        ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
         ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM,
+        ModalityType.AUDIO,
+        ModalityType.AUDIO_SPECTRAL,
+        ModalityType.SPEECH_RECOGNITION,
+        ModalityType.VIDEO,
+        ModalityType.VIDEO_BITWISE,
+        ModalityType.VIDEO_PIXEL_NOISE,
+        ModalityType.VIDEO_L1,
+        ModalityType.VIDEO_L2,
+        ModalityType.VIDEO_L2_RAW,
+        ModalityType.VIDEO_L1_RAW,
+        ModalityType.CLASS_LABEL,
+        ModalityType.VIDEO_IDENTITY,
         ModalityType.MULTI_LABEL,
         ModalityType.ONE_HOT_CLASS_LABEL,
+        ModalityType.IDENTITY,
+        ModalityType.GENERIC_L2_LOSS,
         ModalityType.REAL,
         ModalityType.REAL_L2_LOSS,
         ModalityType.REAL_LOG_POISSON_LOSS,
+        ModalityType.IDENTITY_SYMBOL,
         ModalityType.SIGMOID_CLASS_LABEL,
         ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
+        ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL,
         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
-        ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL,
-        ModalityType.SPEECH_RECOGNITION,
-        ModalityType.SYMBOL,
-        ModalityType.SYMBOL_ONE_HOT,
-        ModalityType.SYMBOL_WEIGHTS_ALL,
-        ModalityType.VIDEO,
-        ModalityType.VIDEO_BITWISE,
-        ModalityType.VIDEO_IDENTITY,
-        ModalityType.VIDEO_L1,
-        ModalityType.VIDEO_L2,
-        ModalityType.VIDEO_L1_RAW,
-        ModalityType.VIDEO_L2_RAW,
-        ModalityType.VIDEO_PIXEL_NOISE,
     ]
 
 
 # Utility functions, similar to tf.keras
+current_module = sys.modules[__name__]
 
 
 def get_bottom(modality_type, value=None):
   """Gets default bottom transformation; if none available, return value."""
-  if modality_type == ModalityType.AUDIO:
-    return audio_bottom
-  elif modality_type == ModalityType.AUDIO_SPECTRAL:
-    return audio_spectral_bottom
-  elif modality_type in (ModalityType.CLASS_LABEL,
-                         ModalityType.MULTI_LABEL,
-                         ModalityType.ONE_HOT_CLASS_LABEL,
-                         ModalityType.SIGMOID_CLASS_LABEL,
-                         ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
-                         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
-                         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
-                         ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL):
-    return class_label_bottom
-  elif modality_type in (ModalityType.CTC_SYMBOL,
-                         ModalityType.SYMBOL,
-                         ModalityType.SYMBOL_WEIGHTS_ALL):
-    return symbol_bottom
-  elif modality_type in (ModalityType.GENERIC_L2_LOSS,
-                         ModalityType.IDENTITY,
-                         ModalityType.IDENTITY_SYMBOL,
-                         ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM):
-    return identity_bottom
-  elif modality_type == ModalityType.IMAGE:
-    return image_bottom
-  elif modality_type in (ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
-                         ModalityType.IMAGE_CHANNEL_COMPRESS):
-    return image_channel_compress_bottom
-  elif modality_type in (ModalityType.REAL,
-                         ModalityType.REAL_L2_LOSS,
-                         ModalityType.REAL_LOG_POISSON_LOSS):
-    return real_bottom
-  elif modality_type == ModalityType.SPEECH_RECOGNITION:
-    return speech_recognition_bottom
-  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
-    return symbol_one_hot_bottom
-  elif modality_type in (ModalityType.VIDEO,
-                         ModalityType.VIDEO_L1,
-                         ModalityType.VIDEO_L2):
-    return video_bottom
-  elif modality_type == ModalityType.VIDEO_BITWISE:
-    return video_bitwise_bottom
-  elif modality_type == ModalityType.VIDEO_IDENTITY:
-    return video_identity_bottom
-  elif modality_type in (ModalityType.VIDEO_L1_RAW,
-                         ModalityType.VIDEO_L2_RAW):
-    return video_raw_bottom
-  elif modality_type == ModalityType.VIDEO_PIXEL_NOISE:
-    return video_pixel_noise_bottom
+  if modality_type in ModalityType.get_choices():
+    modality_cls = getattr(current_module, modality_type)
+    return modality_cls.bottom
   return value
 
 
 def get_loss(modality_type, value=None):
   """Gets default loss transformation; if none available, return value."""
-  if modality_type in (ModalityType.AUDIO,
-                       ModalityType.AUDIO_SPECTRAL,
-                       ModalityType.CLASS_LABEL,
-                       ModalityType.IDENTITY,
-                       ModalityType.IDENTITY_SYMBOL,
-                       ModalityType.IMAGE,
-                       ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
-                       ModalityType.IMAGE_CHANNEL_COMPRESS,
-                       ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM,
-                       ModalityType.REAL,
-                       ModalityType.SPEECH_RECOGNITION,
-                       ModalityType.SYMBOL,
-                       ModalityType.SYMBOL_WEIGHTS_ALL):
-    return generic_loss
-  elif modality_type == ModalityType.CTC_SYMBOL:
-    return ctc_symbol_loss
-  elif modality_type == ModalityType.GENERIC_L2_LOSS:
-    return generic_l2_loss
-  elif modality_type == ModalityType.MULTI_LABEL:
-    return multi_label_loss
-  elif modality_type in (ModalityType.ONE_HOT_CLASS_LABEL,
-                         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
-                         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
-                         ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL):
-    return one_hot_class_label_loss
-  elif modality_type == ModalityType.REAL_L2_LOSS:
-    return real_l2_loss
-  elif modality_type == ModalityType.REAL_LOG_POISSON_LOSS:
-    return real_log_poisson_loss
-  elif modality_type == ModalityType.SIGMOID_CLASS_LABEL:
-    return sigmoid_class_label_loss
-  elif modality_type == ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL:
-    return sigmoid_max_pooling_class_label_loss
-  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
-    return symbol_one_hot_loss
-  elif modality_type in (ModalityType.VIDEO,
-                         ModalityType.VIDEO_BITWISE,
-                         ModalityType.VIDEO_PIXEL_NOISE):
-    return video_loss
-  elif modality_type == ModalityType.VIDEO_IDENTITY:
-    return video_identity_loss
-  elif modality_type == ModalityType.VIDEO_L1:
-    return video_l1_loss
-  elif modality_type == ModalityType.VIDEO_L1_RAW:
-    return video_l1_raw_loss
-  elif modality_type == ModalityType.VIDEO_L2:
-    return video_l2_loss
-  elif modality_type == ModalityType.VIDEO_L2_RAW:
-    return video_l2_raw_loss
+  if modality_type in ModalityType.get_choices():
+    modality_cls = getattr(current_module, modality_type)
+    return modality_cls.loss
   return value
 
 
 def get_name(modality_type, value=None):
   """Gets default name for transformations; if none available, return value."""
-  # For legacy reasons, modalities vary in their naming scheme.
-  if modality_type == ModalityType.AUDIO:
-    return lambda model_hparams, vocab_size: "audio_modality"
-  elif modality_type == ModalityType.AUDIO_SPECTRAL:
-    return lambda model_hparams, vocab_size: "audio_spectral_modality"
-  elif modality_type == ModalityType.GENERIC_L2_LOSS:
-    return lambda model_hparams, vocab_size: "generic_l2_loss_modality"
-  elif modality_type == ModalityType.IDENTITY:
-    return lambda model_hparams, vocab_size: "identity_modality"
-  elif modality_type == ModalityType.IMAGE:
-    return lambda model_hparams, vocab_size: "image_modality"
-  elif modality_type == ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY:
-    return (lambda model_hparams, vocab_size:  # pylint: disable=g-long-lambda
-            "image_channel_bottom_identity_modality")
-  elif modality_type == ModalityType.IMAGE_CHANNEL_COMPRESS:
-    return lambda model_hparams, vocab_size: "image_channel_compress_modality"
-  elif modality_type == ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM:
-    return lambda model_hparams, vocab_size: "image_channel_embeddings_bottom"
-  elif modality_type == ModalityType.REAL:
-    return lambda model_hparams, vocab_size: "real_modality"
-  elif modality_type == ModalityType.REAL_L2_LOSS:
-    return lambda model_hparams, vocab_size: "real_l2_loss_modality"
-  elif modality_type == ModalityType.REAL_LOG_POISSON_LOSS:
-    return lambda model_hparams, vocab_size: "real_log_poisson_loss_modality"
-  elif modality_type == ModalityType.SPEECH_RECOGNITION:
-    return lambda model_hparams, vocab_size: "speech_recognition_modality"
-  elif modality_type == ModalityType.VIDEO:
-    return lambda model_hparams, vocab_size: "video_modality"
-  elif modality_type == ModalityType.VIDEO_BITWISE:
-    return lambda model_hparams, vocab_size: "video_modality_bitwise"
-  elif modality_type == ModalityType.VIDEO_IDENTITY:
-    return lambda model_hparams, vocab_size: "video_modality_identity"
-  elif modality_type == ModalityType.VIDEO_L1:
-    return lambda model_hparams, vocab_size: "video_modality_l1"
-  elif modality_type == ModalityType.VIDEO_L1_RAW:
-    return lambda model_hparams, vocab_size: "video_modality_l1_raw"
-  elif modality_type == ModalityType.VIDEO_L2:
-    return lambda model_hparams, vocab_size: "video_modality_l2"
-  elif modality_type == ModalityType.VIDEO_L2_RAW:
-    return lambda model_hparams, vocab_size: "video_modality_l2_raw"
-  elif modality_type == ModalityType.VIDEO_PIXEL_NOISE:
-    return lambda model_hparams, vocab_size: "video_modality_pixel_noise"
-  elif modality_type in (ModalityType.CLASS_LABEL,
-                         ModalityType.MULTI_LABEL,
-                         ModalityType.ONE_HOT_CLASS_LABEL):
-    def name(model_hparams, vocab_size):
-      return "class_label_modality_%d_%d" % (vocab_size,
-                                             model_hparams.hidden_size)
-    return name
-  elif modality_type in (ModalityType.CTC_SYMBOL,
-                         ModalityType.IDENTITY_SYMBOL,
-                         ModalityType.SYMBOL,
-                         ModalityType.SYMBOL_WEIGHTS_ALL,
-                         ModalityType.SYMBOL_ONE_HOT):
-    def name(model_hparams, vocab_size):
-      return "symbol_modality_%d_%d" % (vocab_size, model_hparams.hidden_size)
-    return name
-  elif modality_type == ModalityType.SIGMOID_CLASS_LABEL:
-    def name(model_hparams, vocab_size):
-      return "sigmoid_class_symbol_modality_%d_%d" % (vocab_size,
-                                                      model_hparams.hidden_size)
-    return name
-  elif modality_type == ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL:
-    def name(model_hparams, vocab_size):
-      return "sigmoid_max_pooling_class_symbol_modality_%d_%d" % (
-          vocab_size, model_hparams.hidden_size)
-    return name
-  elif modality_type == ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL:
-    def name(model_hparams, vocab_size):
-      return "softmax_average_pooling_onehot_class_label_modality_%d_%d" % (
-          vocab_size, model_hparams.hidden_size)
-    return name
-  elif modality_type == ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL:
-    def name(model_hparams, vocab_size):
-      return "softmax_last_timestep_onehot_class_label_modality_%d_%d" % (
-          vocab_size, model_hparams.hidden_size)
-    return name
-  elif modality_type == ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL:
-    def name(model_hparams, vocab_size):
-      return "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
-          vocab_size, model_hparams.hidden_size)
-    return name
+  if modality_type in ModalityType.get_choices():
+    modality_cls = getattr(current_module, modality_type)
+    return modality_cls.name
   return value
 
 
 def get_targets_bottom(modality_type, value=None):
   """Gets default bottom transformation for targets; if none, return value."""
-  if modality_type == ModalityType.AUDIO:
-    return make_targets_bottom(audio_bottom)
-  elif modality_type == ModalityType.AUDIO_SPECTRAL:
-    return make_targets_bottom(audio_spectral_bottom)
-  elif modality_type in (ModalityType.CLASS_LABEL,
-                         ModalityType.MULTI_LABEL,
-                         ModalityType.ONE_HOT_CLASS_LABEL,
-                         ModalityType.SIGMOID_CLASS_LABEL,
-                         ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
-                         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
-                         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
-                         ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL):
-    return class_label_targets_bottom
-  elif modality_type in (ModalityType.CTC_SYMBOL,
-                         ModalityType.SYMBOL,
-                         ModalityType.SYMBOL_WEIGHTS_ALL):
-    return symbol_targets_bottom
-  elif modality_type in (ModalityType.GENERIC_L2_LOSS,
-                         ModalityType.IDENTITY_SYMBOL):
-    return identity_bottom
-  elif modality_type == ModalityType.IDENTITY:
-    return make_targets_bottom(identity_bottom)
-  elif modality_type == ModalityType.IMAGE:
-    return image_targets_bottom
-  elif modality_type in (ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
-                         ModalityType.IMAGE_CHANNEL_COMPRESS):
-    return image_channel_compress_targets_bottom
-  elif modality_type == ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM:
-    return image_channel_embeddings_bottom
-  elif modality_type in (ModalityType.REAL,
-                         ModalityType.REAL_L2_LOSS,
-                         ModalityType.REAL_LOG_POISSON_LOSS):
-    return make_targets_bottom(real_bottom)
-  elif modality_type == ModalityType.SPEECH_RECOGNITION:
-    return make_targets_bottom(speech_recognition_bottom)
-  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
-    return symbol_one_hot_bottom
-  elif modality_type in (ModalityType.VIDEO,
-                         ModalityType.VIDEO_L1,
-                         ModalityType.VIDEO_L2):
-    return video_targets_bottom
-  elif modality_type == ModalityType.VIDEO_BITWISE:
-    return video_bitwise_targets_bottom
-  elif modality_type == ModalityType.VIDEO_IDENTITY:
-    return video_identity_targets_bottom
-  elif modality_type in (ModalityType.VIDEO_L1_RAW,
-                         ModalityType.VIDEO_L2_RAW):
-    return video_raw_targets_bottom
-  elif modality_type == ModalityType.VIDEO_PIXEL_NOISE:
-    return make_targets_bottom(video_pixel_noise_bottom)
+  if modality_type in ModalityType.get_choices():
+    modality_cls = getattr(current_module, modality_type)
+    return modality_cls.targets_bottom
   return value
 
 
-def get_top(modality_type, value=None):
-  """Gets default top transformation; if none available, return value."""
-  if modality_type in (ModalityType.AUDIO,
-                       ModalityType.AUDIO_SPECTRAL,
-                       ModalityType.GENERIC_L2_LOSS,
-                       ModalityType.IDENTITY,
-                       ModalityType.IDENTITY_SYMBOL,
-                       ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
-                       ModalityType.SPEECH_RECOGNITION,
-                       ModalityType.VIDEO_IDENTITY):
-    return identity_top
-  elif modality_type in (ModalityType.CLASS_LABEL,
-                         ModalityType.MULTI_LABEL,
-                         ModalityType.ONE_HOT_CLASS_LABEL,
-                         ModalityType.SIGMOID_CLASS_LABEL):
-    return class_label_top
-  elif modality_type in (ModalityType.CTC_SYMBOL,
-                         ModalityType.SYMBOL,
-                         ModalityType.SYMBOL_WEIGHTS_ALL):
-    return symbol_top
-  elif modality_type == ModalityType.IMAGE:
-    return image_top
-  elif modality_type == ModalityType.IMAGE_CHANNEL_COMPRESS:
-    return image_channel_compress_top
-  elif modality_type == ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM:
-    return image_channel_embeddings_top
-  elif modality_type in (ModalityType.REAL,
-                         ModalityType.REAL_L2_LOSS,
-                         ModalityType.REAL_LOG_POISSON_LOSS):
-    return real_top
-  elif modality_type == ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL:
-    return sigmoid_max_pooling_class_label_top
-  elif modality_type == ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL:
-    return softmax_average_pooling_class_label_top
-  elif modality_type == ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL:
-    return softmax_last_timestep_class_label_top
-  elif modality_type == ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL:
-    return softmax_max_pooling_class_label_top
-  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
-    return symbol_one_hot_top
-  elif modality_type in (ModalityType.VIDEO,
-                         ModalityType.VIDEO_BITWISE,
-                         ModalityType.VIDEO_PIXEL_NOISE):
-    return video_top
-  elif modality_type in (ModalityType.VIDEO_L1,
-                         ModalityType.VIDEO_L2):
-    return video_l1_top
-  elif modality_type in (ModalityType.VIDEO_L1_RAW,
-                         ModalityType.VIDEO_L2_RAW):
-    return video_raw_top
+def get_targets_weights_fn(modality_type, value=None):
+  """Gets default weights function; if none available, return value."""
+  if modality_type in ModalityType.get_choices():
+    modality_cls = getattr(current_module, modality_type)
+    return modality_cls.targets_weights_fn
   return value
 
 
-def get_weights_fn(modality_type, value=None):
-  """Gets default weights function; if none available, return value."""
-  if modality_type in (ModalityType.CTC_SYMBOL,
-                       ModalityType.IDENTITY_SYMBOL,
-                       ModalityType.MULTI_LABEL,
-                       ModalityType.SYMBOL,
-                       ModalityType.SYMBOL_ONE_HOT):
-    return common_layers.weights_nonzero
-  elif modality_type in ModalityType.get_choices():
-    return common_layers.weights_all
+def get_top(modality_type, value=None):
+  """Gets default top transformation; if none available, return value."""
+  if modality_type in ModalityType.get_choices():
+    modality_cls = getattr(current_module, modality_type)
+    return modality_cls.top
   return value
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 69504bb8e..6cb6efc98 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -31,26 +31,6 @@
 
 class ModalityTest(tf.test.TestCase):
 
-  @test_utils.run_in_graph_and_eager_modes()
-  def testGetForAllModalities(self):
-    for modality in modalities.ModalityType.get_choices():
-      bottom = modalities.get_bottom(modality)
-      loss = modalities.get_loss(modality)
-      name = modalities.get_name(modality)
-      targets_bottom = modalities.get_targets_bottom(modality)
-      top = modalities.get_top(modality)
-      weights_fn = modalities.get_weights_fn(modality)
-      self.assertIsNotNone(bottom,
-                           msg="{} has no default bottom".format(modality))
-      self.assertIsNotNone(loss, msg="{} has no default loss".format(modality))
-      self.assertIsNotNone(name, msg="{} has no default name".format(modality))
-      self.assertIsNotNone(
-          targets_bottom,
-          msg="{} has no default targets_bottom".format(modality))
-      self.assertIsNotNone(top, msg="{} has no default top".format(modality))
-      self.assertIsNotNone(weights_fn,
-                           msg="{} has no default weights_fn".format(modality))
-
   @test_utils.run_in_graph_and_eager_modes()
   def testSymbolModalityInputs(self):
     batch_size = 10
@@ -107,7 +87,7 @@ def testSymbolModalityTargets(self):
         sharded_targets,
         model_hparams,
         vocab_size,
-        modalities.get_weights_fn(modalities.ModalityType.SYMBOL))
+        modalities.get_targets_weights_fn(modalities.ModalityType.SYMBOL))
     train_loss = (tf.add_n(sharded_loss_num) /
                   tf.maximum(1.0, tf.add_n(sharded_loss_den)))
     logits = tf.concat(sharded_logits, 0)
@@ -144,12 +124,12 @@ def testSymbolModalityTargetsFactored(self):
           model_hparams,
           vocab_size)
       sharded_loss_num, sharded_loss_den = data_parallelism(
-          modalities.get_loss(modalities.ModalityType.SYMBOL),
+          modalities.SymbolModality.loss,
           sharded_logits,
           sharded_targets,
           model_hparams,
           vocab_size,
-          modalities.get_weights_fn(modalities.ModalityType.SYMBOL))
+          modalities.get_targets_weights_fn(modalities.ModalityType.SYMBOL))
       train_loss = (tf.add_n(sharded_loss_num) /
                     tf.maximum(1.0, tf.add_n(sharded_loss_den)))
       logits = tf.concat(sharded_logits, 0)
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index d43f96c15..81d8de434 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -48,11 +48,16 @@ def body(self, features):
     hparams = copy.copy(self._hparams)
     targets = features["targets"]
     if (hparams.likelihood == cia.DistributionType.DMOL and
-        hparams.num_channels != 1):
-      raise ValueError("When using DMOL for the likelihood, bottom function "
-                       " must be identity and num_channels must be 1.")
+        (hparams.modality["targets"] !=
+         modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY or
+         hparams.num_channels != 1)):
+      raise ValueError("When using DMOL for the likelihood,modality['targets'] "
+                       "must be ImageChannelBottomIdentityModality and "
+                       "num_channels must be 1.")
     if (not tf.get_variable_scope().reuse and
-        hparams.mode != tf.estimator.ModeKeys.PREDICT):
+        hparams.mode != tf.estimator.ModeKeys.PREDICT and
+        hparams.modality["targets"] !=
+        modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY):
       tf.summary.image("targets", tf.to_float(targets), max_outputs=1)
 
     # Extra losses list if we want to use moe.
@@ -188,8 +193,7 @@ def image_transformer_base():
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.98
   hparams.label_smoothing = 0.0
-  hparams.bottom["targets"] = modalities.image_channel_embeddings_bottom
-  hparams.top["targets"] = modalities.identity_top
+  hparams.modality["targets"] = modalities.ModalityType.IDENTITY
   hparams.norm_type = "layer"
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.add_hparam("filter_size", 512)  # Add new ones like this.
@@ -276,8 +280,8 @@ def imagetransformer_cifar10_base_dmol():
   hparams = image_transformer_base()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.bottom["targets"] = modalities.image_channel_compress_targets_bottom
-  hparams.top["targets"] = modalities.identity_top
+  hparams.modality["targets"] = (
+      modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY)
   hparams.num_heads = 8
   hparams.batch_size = 8
   hparams.sampling_method = "random"
@@ -418,8 +422,8 @@ def imagetransformerpp_sep_channels_8l_8h():
   hparams = imagetransformer_base()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.bottom["targets"] = modalities.image_channel_compress_targets_bottom
-  hparams.top["targets"] = modalities.identity_top
+  hparams.modality["targets"] = (
+      modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY)
   hparams.num_heads = 8
   hparams.batch_size = 4
   hparams.attention_key_channels = hparams.attention_value_channels = 0
@@ -882,8 +886,8 @@ def imagetransformerpp_tiny():
   hparams = imagetransformer_tiny()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.bottom["targets"] = modalities.image_channel_compress_targets_bottom
-  hparams.top["targets"] = modalities.identity_top
+  hparams.modality["targets"] = (
+      modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY)
   return hparams
 
 
diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 80ce13d53..acf793440 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -382,9 +382,7 @@ def image_transformer2d_base():
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.98
   hparams.label_smoothing = 0.0
-  hparams.bottom["targets"] = modalities.make_targets_bottom(
-      modalities.image_channel_embeddings_bottom)
-  hparams.top["targets"] = modalities.identity_top
+  hparams.modality["targets"] = modalities.ModalityType.IDENTITY
   hparams.norm_type = "layer"
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.add_hparam("filter_size", 512)  # Add new ones like this.
@@ -595,7 +593,6 @@ def img2img_transformer2d_base():
   hparams.filter_size = 2048
   hparams.num_encoder_layers = 4
   hparams.num_decoder_layers = 8
-  hparams.bottom["inputs"] = modalities.image_channel_embeddings_bottom
   hparams.dec_attention_type = cia.AttentionType.LOCAL_2D
   hparams.block_raster_scan = True
   return hparams
diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index b8d4c56d8..3e32e9c03 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -854,12 +854,9 @@ def mtf_transformer_base():
   # These parameters make Transformer model compatible with MtfTransformer
   # Do not override these, as mtf_transformer does not support other options.
   hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
-  hparams.bottom = {
-      "inputs": modalities.identity_bottom,
-      "targets": modalities.identity_bottom,
-  }
-  hparams.top = {
-      "targets": modalities.identity_top,
+  hparams.modality = {
+      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
+      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
   }
 
   # Parameters for computing the maximum decode length in beam search.
diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index c8cb5354d..1d1e2a71d 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -418,12 +418,9 @@ def mtf_transformer2_base():
   hparams.use_fixed_batch_size = True
   hparams.add_hparam("mtf_mode", True)
   hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
-  hparams.bottom = {
-      "inputs": modalities.identity_bottom,
-      "targets": modalities.identity_bottom,
-  }
-  hparams.top = {
-      "targets": modalities.identity_top,
+  hparams.modality = {
+      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
+      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
   }
   hparams.add_hparam("beam_size", 1)
 
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index cf4305a30..02dde9e03 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1112,12 +1112,9 @@ def autoencoder_residual_text():
   hparams.hidden_size = 64
   hparams.max_hidden_size = 512
   hparams.bottleneck_noise = 0.0
-  hparams.bottom = {
-      "inputs": modalities.identity_bottom,
-      "targets": modalities.identity_bottom,
-  }
-  hparams.top = {
-      "targets": modalities.identity_top,
+  hparams.modality = {
+      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
+      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
   }
   hparams.autoregressive_mode = "none"
   hparams.sample_width = 1
@@ -1222,12 +1219,9 @@ def autoencoder_ordered_text():
   hparams.batch_size = 1024
   hparams.autoregressive_mode = "conv5"
   hparams.max_hidden_size = 1024
-  hparams.bottom = {
-      "inputs": modalities.identity_bottom,
-      "targets": modalities.identity_bottom,
-  }
-  hparams.top = {
-      "targets": modalities.identity_top,
+  hparams.modality = {
+      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
+      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
   }
   hparams.sample_height = 128
   hparams.sample_width = 1
diff --git a/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index 3ba2e64ab..94842b2db 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -127,12 +127,9 @@ def cycle_gan_small():
   """Set of hyperparameters."""
   hparams = transformer_vae.transformer_ae_small()
   hparams.batch_size = 2048
-  hparams.bottom = {
-      "inputs": modalities.identity_bottom,
-      "targets": modalities.identity_bottom,
-  }
-  hparams.top = {
-      "targets": modalities.identity_top,
+  hparams.modality = {
+      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
+      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
   }
   hparams.weight_decay = 3.0
   hparams.learning_rate = 0.05
diff --git a/tensor2tensor/models/research/super_lm.py b/tensor2tensor/models/research/super_lm.py
index 0e7dd58c8..5a02ccf94 100644
--- a/tensor2tensor/models/research/super_lm.py
+++ b/tensor2tensor/models/research/super_lm.py
@@ -265,13 +265,7 @@ def super_lm_base():
   # we only want one data shard.
   hparams.no_data_parallelism = True
   # bypass the symbol modality so that we can use model parallelism.
-  hparams.bottom = {
-      "inputs": modalities.identity_bottom,
-      "targets": modalities.identity_bottom,
-  }
-  hparams.top = {
-      "targets": modalities.identity_top,
-  }
+  hparams.modality["targets"] = modalities.ModalityType.IDENTITY_SYMBOL
   hparams.add_hparam("filter_size", 512)
   hparams.add_hparam("mix_fraction", 0.5)
   # attention-related flags
diff --git a/tensor2tensor/models/research/transformer_symshard.py b/tensor2tensor/models/research/transformer_symshard.py
index 7f5bf39dd..98c9cf0ee 100644
--- a/tensor2tensor/models/research/transformer_symshard.py
+++ b/tensor2tensor/models/research/transformer_symshard.py
@@ -367,12 +367,9 @@ def transformer_symshard_base():
   # we only want one data shard.
   hparams.no_data_parallelism = True
   # bypass the symbol modality so that we can use model parallelism.
-  hparams.bottom = {
-      "inputs": modalities.identity_bottom,
-      "targets": modalities.identity_bottom,
-  }
-  hparams.top = {
-      "targets": modalities.identity_top,
+  hparams.modality = {
+      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
+      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
   }
   hparams.add_hparam("filter_size", 1280)
   hparams.add_hparam("mix_fraction", 0.5)
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 9cfe80933..e8f769bb6 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -896,8 +896,7 @@ def imagetransformer_ae_cifar():
 
   hparams.add_hparam("unconditional", False)  # unconditional generation
 
-  hparams.bottom["targets"] = modalities.image_channel_embeddings_bottom
-  hparams.top["targets"] = modalities.image_channel_embeddings_top
+  hparams.modality["targets"] = modalities.ImageChannelEmbeddingsBottom
   hparams.drop_inputs = True
   hparams.do_attend_compress = False
   hparams.do_attend_decompress = False
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index ec068e452..c3dfb53ea 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -195,12 +195,16 @@ def __init__(self, *args, **kwargs):
 
   @property
   def _target_modality(self):
-    return self.problem_hparams.modality["targets"]
+    target_modality = self.hparams.modality.get(
+        "targets",
+        self.problem_hparams.modality["targets"])
+    if target_modality not in modalities.ModalityType.get_choices():
+      target_modality = target_modality.__class__.__name__
+    return target_modality
 
   @property
   def is_per_pixel_softmax(self):
-    # TODO(trandustin): This is a hack.
-    return "targets" not in self.hparams.get("loss")
+    return self._target_modality == modalities.ModalityType.VIDEO
 
   def get_iteration_num(self):
     step_num = tf.train.get_global_step()
@@ -334,12 +338,11 @@ def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
       Additional reconstruction loss.
 
     Raises:
-      ValueError: in case of unknown loss transformation.
+      ValueError: in case of unknown modality.
     """
-    # TODO(trandustin): This logic should be moved elsewhere.
-    if self.hparams.loss.get("targets") == modalities.video_l2_raw_loss:
+    if self._target_modality == modalities.ModalityType.VIDEO_L2_RAW:
       recon_loss = tf.losses.mean_squared_error(extra_gts, extra_pds)
-    elif "targets" not in self.hparams.loss:
+    elif self._target_modality == modalities.ModalityType.VIDEO:
       shape = common_layers.shape_list(extra_pds)
       updated_shape = shape[:-1] + [3, 256]
       extra_pds = tf.reshape(extra_pds, updated_shape)
@@ -348,9 +351,10 @@ def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
       targets = extra_raw_gts
       targets_shape = common_layers.shape_list(targets)
       targets = tf.reshape(targets, [-1] + targets_shape[2:])
-      targets_weights_fn = self.hparams.weights_fn.get(
+      modality = self.hparams.problem_hparams.modality["targets"]
+      targets_weights_fn = self.hparams.targets_weights_fn.get(
           "targets",
-          modalities.get_weights_fn(self._target_modality))
+          modalities.get_targets_weights_fn(modality))
       numerator, denominator = common_layers.padded_cross_entropy(
           logits,
           targets,
@@ -359,7 +363,7 @@ def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
           weights_fn=targets_weights_fn)
       recon_loss = numerator / denominator
     else:
-      raise ValueError("internal loss only supports specific hparams.loss.")
+      raise ValueError("internal loss only supports specific modalities.")
     tf.summary.scalar("recon_extra", recon_loss)
     return recon_loss
 
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 8d8a2c6ff..7616d3549 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -55,8 +55,7 @@ def next_frame_pixel_noise():
   """Basic 2-frame conv model with pixel noise."""
   hparams = next_frame_basic_deterministic()
   hparams.add_hparam("video_modality_input_noise", 0.05)
-  hparams.bottom["inputs"] = modalities.video_pixel_noise_bottom
-  hparams.top["inputs"] = modalities.video_top
+  hparams.modality["inputs"] = modalities.ModalityType.VIDEO_PIXEL_NOISE
   return hparams
 
 
@@ -90,8 +89,7 @@ def next_frame_tpu():
 def next_frame_ae():
   """Conv autoencoder."""
   hparams = next_frame_basic_deterministic()
-  hparams.bottom["inputs"] = modalities.video_bitwise_bottom
-  hparams.top["inputs"] = modalities.video_top
+  hparams.modality["inputs"] = modalities.ModalityType.VIDEO_BITWISE
   hparams.hidden_size = 256
   hparams.batch_size = 8
   hparams.num_hidden_layers = 4
@@ -104,8 +102,7 @@ def next_frame_ae():
 def next_frame_ae_tiny():
   """Conv autoencoder, tiny set for testing."""
   hparams = next_frame_tiny()
-  hparams.bottom["inputs"] = modalities.video_bitwise_bottom
-  hparams.top["inputs"] = modalities.video_top
+  hparams.modality["inputs"] = modalities.ModalityType.VIDEO_BITWISE
   hparams.batch_size = 8
   hparams.dropout = 0.4
   return hparams
@@ -134,8 +131,7 @@ def next_frame_tiny():
 def next_frame_l1():
   """Basic conv model with L1 modality."""
   hparams = next_frame_basic_deterministic()
-  hparams.loss["targets"] = modalities.video_l1_loss
-  hparams.top["targets"] = modalities.video_l1_top
+  hparams.modality["targets"] = modalities.ModalityType.VIDEO_L1
   hparams.video_modality_loss_cutoff = 2.4
   return hparams
 
@@ -144,8 +140,7 @@ def next_frame_l1():
 def next_frame_l2():
   """Basic conv model with L2 modality."""
   hparams = next_frame_basic_deterministic()
-  hparams.loss["targets"] = modalities.video_l2_loss
-  hparams.top["targets"] = modalities.video_l1_top
+  hparams.modality["targets"] = modalities.ModalityType.VIDEO_L2
   hparams.video_modality_loss_cutoff = 2.4
   return hparams
 
diff --git a/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
index e42c94d9d..a8b8f98c9 100644
--- a/tensor2tensor/models/video/epva_params.py
+++ b/tensor2tensor/models/video/epva_params.py
@@ -29,15 +29,9 @@ def next_frame_epva():
   hparams = basic_deterministic_params.next_frame_basic_deterministic()
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 4
-  hparams.bottom = {
-      "inputs": modalities.video_raw_bottom,
-      "targets": modalities.video_raw_targets_bottom,
-  }
-  hparams.loss = {
-      "targets": modalities.video_l2_raw_loss,
-  }
-  hparams.top = {
-      "targets": modalities.video_raw_top,
+  hparams.modality = {
+      "inputs": modalities.ModalityType.VIDEO_L2_RAW,
+      "targets": modalities.ModalityType.VIDEO_L2_RAW,
   }
   hparams.learning_rate_schedule = "constant"
   hparams.learning_rate_constant = 1e-05
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index c323ae088..56a9da10b 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -69,15 +69,9 @@ def next_frame_glow_hparams():
   # Pretrains the glow encoder for "pretrain_steps" number of steps.
   # By default, don't pretrain and learn end-to-end
   hparams.add_hparam("pretrain_steps", -1)
-  hparams.bottom = {
-      "inputs": modalities.video_raw_bottom,
-      "targets": modalities.video_raw_targets_bottom,
-  }
-  hparams.loss = {
-      "targets": modalities.video_l1_raw_loss,
-  }
-  hparams.top = {
-      "targets": modalities.video_raw_top,
+  hparams.modality = {
+      "inputs": modalities.ModalityType.VIDEO_L1_RAW,
+      "targets": modalities.ModalityType.VIDEO_L1_RAW,
   }
   hparams.init_batch_size = 256
   hparams.batch_size = 32
diff --git a/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index c28c4378a..708c888f3 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -36,15 +36,9 @@ def next_frame_savp():
   hparams.add_hparam("gan_loss_multiplier", 0.01)
   hparams.add_hparam("gan_vae_loss_multiplier", 0.01)
   hparams.add_hparam("gan_optimization", "joint")
-  hparams.bottom = {
-      "inputs": modalities.video_raw_bottom,
-      "targets": modalities.video_raw_targets_bottom,
-  }
-  hparams.loss = {
-      "targets": modalities.video_l1_raw_loss,
-  }
-  hparams.top = {
-      "targets": modalities.video_raw_top,
+  hparams.modality = {
+      "inputs": modalities.ModalityType.VIDEO_L1_RAW,
+      "targets": modalities.ModalityType.VIDEO_L1_RAW,
   }
   hparams.latent_loss_multiplier_schedule = "linear"
   hparams.upsample_method = "bilinear_upsample_conv"
@@ -60,8 +54,9 @@ def next_frame_savp():
 def next_frame_savp_l2():
   """SAVP with L2 reconstruction loss."""
   hparams = next_frame_savp()
-  hparams.loss = {
-      "targets": modalities.video_l2_raw_loss,
+  hparams.modality = {
+      "inputs": modalities.ModalityType.VIDEO_L2_RAW,
+      "targets": modalities.ModalityType.VIDEO_L2_RAW,
   }
   return hparams
 
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index f7352c714..f5687865b 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -33,15 +33,9 @@ def next_frame_sv2p():
   hparams.video_num_input_frames = 1
   hparams.video_num_target_frames = 3
   hparams.batch_size = 16
-  hparams.bottom = {
-      "inputs": modalities.video_raw_bottom,
-      "targets": modalities.video_raw_targets_bottom,
-  }
-  hparams.loss = {
-      "targets": modalities.video_l2_raw_loss,
-  }
-  hparams.top = {
-      "targets": modalities.video_raw_top,
+  hparams.modality = {
+      "inputs": modalities.ModalityType.VIDEO_L2_RAW,
+      "targets": modalities.ModalityType.VIDEO_L2_RAW,
   }
   hparams.video_modality_loss_cutoff = 0.0
   hparams.scheduled_sampling_mode = "count"
@@ -97,9 +91,10 @@ def next_frame_sv2p_atari():
 def next_frame_sv2p_atari_softmax():
   """SV2P model for atari with softmax."""
   hparams = next_frame_sv2p_atari()
-  hparams.bottom = {}
-  hparams.loss = {}
-  hparams.top = {}
+  hparams.modality = {
+      "inputs": modalities.ModalityType.VIDEO,
+      "targets": modalities.ModalityType.VIDEO,
+  }
   hparams.internal_loss = True
   return hparams
 
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 951cbc3d5..d3852c648 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -614,9 +614,9 @@ def weights_fn_for_mp(problem_task_id):
       tm = {"targets": tm}
 
     for target_name, modality in six.iteritems(tm):
-      weights_fn = model_hparams.weights_fn.get(
+      weights_fn = model_hparams.targets_weights_fn.get(
           "targets",
-          modalities.get_weights_fn(modality))
+          modalities.get_targets_weights_fn(modality))
       if hasattr(model_hparams.problem, "task_list"):
         ptid = problem_instance.task_id  # pylint: disable=cell-var-from-loop
         weights_fn = weights_fn_for_mp(ptid)
@@ -643,9 +643,9 @@ def create_eager_metrics_for_problem(problem, model_hparams):
   metric_fns = problem.eval_metric_fns(model_hparams)
   problem_hparams = problem.get_hparams(model_hparams)
   target_modality = problem_hparams.modality["targets"]
-  weights_fn = model_hparams.weights_fn.get(
+  weights_fn = model_hparams.targets_weights_fn.get(
       "targets",
-      modalities.get_weights_fn(target_modality))
+      modalities.get_targets_weights_fn(target_modality))
   return create_eager_metrics_internal(metric_fns, weights_fn=weights_fn)
 
 
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
new file mode 100644
index 000000000..09a83d050
--- /dev/null
+++ b/tensor2tensor/utils/modality.py
@@ -0,0 +1,121 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Modality base class - defines the bottom and top of the model."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import misc_utils
+
+import tensorflow as tf
+
+
+class Modality(object):
+  """Abstract Modality class for data transformations.
+
+  An abstract class representing modalities for transforming data to a space
+  interpretable by T2T models. It has 4 functions:
+  * bottom: called on inputs entering the model.
+  * targets_bottom: called on targets entering the model (e.g., the decoder).
+  * top: called on model outputs to generate predictions (e.g., logits).
+  * loss: called on predictions (outputs of top) and targets.
+
+  For example, think about a modality for images:
+  * `bottom` represents the part of the model applied to an incoming image,
+    e.g., an entry flow of a convolutional network.
+  * `top` represents the top part of a model that is generating images, e.g., a
+    PixelCNN network.
+  * `targets_bottom` represents the auto-regressive part of the network.  It is
+    applied to the already-generated part of an image, which is given to the
+    decoder to generate the next part. In some cases, e.g., for text, it is the
+    same as the `bottom` function, and that is the default we use. But, e.g.,
+    for images, a different function might be needed to regress properly.
+  * `loss` would compare the generated image to the target image and score it.
+  """
+
+  def __init__(self, model_hparams, vocab_size=None):
+    # __init__ args are unused in any methods. They're maintained for
+    # backwards compatibility for now. In the future, Modality classes will be
+    # removed altogether.
+    del model_hparams, vocab_size
+
+  @classmethod
+  def name(cls, model_hparams, vocab_size=None):
+    del model_hparams, vocab_size  # unused arg
+    return misc_utils.camelcase_to_snakecase(type(cls).__name__)
+
+  targets_weights_fn = staticmethod(common_layers.weights_all)
+
+  @staticmethod
+  def bottom(x, model_hparams, vocab_size=None):
+    """Transform one shard of input.
+
+    Args:
+      x: An int32 Tensor with shape [batch, p0, p1, input_channels]
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+
+    Returns:
+      A float32 Tensor with shape [batch, p0, p1, body_input_depth]
+    """
+    raise NotImplementedError("Abstract Method")
+
+  @classmethod
+  def targets_bottom(cls, x, model_hparams, vocab_size=None):
+    """Transform one shard of targets.
+
+    Args:
+      x: An int32 Tensor with shape [batch, p0, p1, target_channels]
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+
+    Returns:
+      A float32 Tensor with shape [batch, p0, p1, body_input_depth]
+    """
+    with tf.variable_scope("targets_bottom"):
+      return cls.bottom(x, model_hparams, vocab_size)
+
+  @staticmethod
+  def top(body_output, targets, model_hparams, vocab_size=None):
+    """Generate predictions/logits for one shard of output.
+
+    Most classes will override this function.
+
+    Args:
+      body_output: A Tensor with shape [batch, p0, p1, body_output_depth]
+      targets: A Tensor with shape [batch, p0, p1, targets_channels,
+        top_dimensionality]
+      model_hparams: tf.HParams, model hyperparmeters.
+      vocab_size: int, vocabulary size.
+
+    Returns:
+      A Tensor of class logits.
+    """
+    raise NotImplementedError("Abstract Method")
+
+  @staticmethod
+  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+    """Compute loss numerator and denominator for one shard of output."""
+    del vocab_size  # unused arg
+    logits = top_out
+    logits = common_attention.maybe_upcast(logits, hparams=model_hparams)
+    return common_layers.padded_cross_entropy(
+        logits,
+        targets,
+        model_hparams.label_smoothing,
+        weights_fn=weights_fn)
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index dc3451cbf..f3b5c8fec 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -204,19 +204,47 @@ def __init__(self,
           },
           hparams=hparams)
 
+    # TODO(trandustin): For now, we get custom feature transformations via
+    # hparams.modality. Once modality classes are removed, let users
+    # individually specify custom transformations for bottom, loss, etc.
+    if not hasattr(hparams, "modality"):
+      hparams.add_hparam("modality", {})
+    if not hasattr(hparams, "bottom"):
+      hparams.add_hparam("bottom", {})
+    if not hasattr(hparams, "loss"):
+      hparams.add_hparam("loss", {})
+    if not hasattr(hparams, "name"):
+      hparams.add_hparam("name", {})
+    if not hasattr(hparams, "targets_weights_fn"):
+      hparams.add_hparam("targets_weights_fn", {})
+    if not hasattr(hparams, "top"):
+      hparams.add_hparam("top", {})
+    target_modalities = _create_target_modality(hparams.modality)
+    for feature_name, modality in six.iteritems(hparams.modality):
+      if modality in modalities.ModalityType.get_choices():
+        modality = getattr(modalities, modality)
+      if feature_name in target_modalities:
+        hparams.bottom[feature_name] = modality.targets_bottom
+      else:
+        hparams.bottom[feature_name] = modality.bottom
+      hparams.loss[feature_name] = modality.loss
+      hparams.name[feature_name] = modality.name
+      hparams.targets_weights_fn[feature_name] = modality.targets_weights_fn
+      hparams.top[feature_name] = modality.top
+
     if self._problem_hparams:
       for feature_name, modality in six.iteritems(
           self._problem_hparams.modality):
-        # If prepend mode, set weights_fn to appropriately handle it.
-        if (modality in (modalities.ModalityType.CTC_SYMBOL,
-                         modalities.ModalityType.IDENTITY_SYMBOL,
-                         modalities.ModalityType.SYMBOL,
-                         modalities.ModalityType.SYMBOL_ONE_HOT)):
+        # If prepend mode, set targets_weights_fn to appropriately handle it.
+        if (modality in (modalities.ModalityType.SYMBOL,
+                         modalities.ModalityType.SYMBOL_ONE_HOT,
+                         modalities.ModalityType.CTC_SYMBOL,
+                         modalities.ModalityType.IDENTITY_SYMBOL)):
           if (hparams.prepend_mode == "prepend_inputs_full_attention" or
               (hparams.prepend_mode == "prepend_inputs_masked_attention" and
                mode != tf.estimator.ModeKeys.TRAIN)):
             weights_fn = common_layers.weights_prepend_inputs_to_targets
-            hparams.weights_fn[feature_name] = weights_fn
+            hparams.targets_weights_fn[feature_name] = weights_fn
 
     self._original_hparams = hparams
     self.set_mode(mode)
@@ -618,8 +646,8 @@ def _loss_single(self, logits, feature_name, feature, weights=None):
     if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
       vocab_size += (-vocab_size) % self._hparams.vocab_divisor
     loss = self._hparams.loss.get(feature_name, modalities.get_loss(modality))
-    targets_weights_fn = self._hparams.weights_fn.get(
-        "targets", modalities.get_weights_fn(modality))
+    targets_weights_fn = self._hparams.targets_weights_fn.get(
+        "targets", modalities.get_targets_weights_fn(modality))
     if weights is None:
       loss_num, loss_den = loss(logits, feature, self._hparams, vocab_size,
                                 weights_fn=targets_weights_fn)
@@ -1789,7 +1817,7 @@ def create_tpu_eval_metrics_fn(problem, model_hparams):
   tm = _create_target_modality(problem.get_hparams(model_hparams).modality)
   if isinstance(tm, dict):
     for k, v in six.iteritems(tm):
-      weights_fn = v.weights_fn
+      weights_fn = v.targets_weights_fn
 
       def make_metric_fn(metric_fn):
         def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
@@ -1809,7 +1837,7 @@ def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
         name = "%s/metrics-%s/%s" % (k, problem.name, metric)
         metric_fns.append((name, make_metric_fn(metric_fn)))
   else:
-    weights_fn = tm.weights_fn
+    weights_fn = tm.targets_weights_fn
 
     def make_metric_fn(metric_fn):
       def wrapped_metric_fn(logits, labels, features):
@@ -2010,8 +2038,8 @@ def sampled_results():
                                 vocab_size)
         if "training" not in losses:
           loss = hparams.loss.get("targets", modalities.get_loss(modality))
-          weights_fn = hparams.weights_fn.get(
-              "targets", modalities.get_weights_fn(modality))
+          weights_fn = hparams.targets_weights_fn.get(
+              "targets", modalities.get_targets_weights_fn(modality))
           sharded_loss_num, sharded_loss_den = dp(loss,
                                                   sharded_logits,
                                                   sharded_features["targets"],
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index 229846832..c9b680302 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -52,8 +52,6 @@ def testLossSingleWeights(self):
 
         model_hparams = HParams(
             prepend_mode="none",
-            loss={},
-            weights_fn={},
             label_smoothing=0.0,
             shared_embedding_and_softmax_weights=False)
 

From 150b2615b0ca6e2baf7dabb22dc7b0d67281ebd2 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 26 Feb 2019 08:50:27 -0800
Subject: [PATCH 1720/2720] Refactoring: state save/restore, evaluation,
 logging, inputs

PiperOrigin-RevId: 235729093
---
 tensor2tensor/jax/__init__.py                 |  15 ++
 .../jax/{input_pipeline.py => inputs.py}      |  31 +++
 tensor2tensor/jax/j2j.py                      | 229 +++++++++---------
 tensor2tensor/jax/j2j_trainer.py              |  10 +-
 tensor2tensor/jax/models/__init__.py          |   2 -
 5 files changed, 172 insertions(+), 115 deletions(-)
 create mode 100644 tensor2tensor/jax/__init__.py
 rename tensor2tensor/jax/{input_pipeline.py => inputs.py} (90%)

diff --git a/tensor2tensor/jax/__init__.py b/tensor2tensor/jax/__init__.py
new file mode 100644
index 000000000..4872e5d5d
--- /dev/null
+++ b/tensor2tensor/jax/__init__.py
@@ -0,0 +1,15 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
diff --git a/tensor2tensor/jax/input_pipeline.py b/tensor2tensor/jax/inputs.py
similarity index 90%
rename from tensor2tensor/jax/input_pipeline.py
rename to tensor2tensor/jax/inputs.py
index 18aa984a8..457534b6c 100644
--- a/tensor2tensor/jax/input_pipeline.py
+++ b/tensor2tensor/jax/inputs.py
@@ -22,12 +22,43 @@
 import collections
 import gin
 
+import jax.numpy as np
+
 from tensor2tensor import problems
 
 import tensorflow as tf
 import tensorflow_datasets as tfds
 
 
+Inputs = collections.namedtuple(
+    "_Inputs", ["train_fn", "eval_fn", "input_shape"])
+
+
+def make_inputs(dataset_name, data_dir):
+  """Make Inputs."""
+  (train_batches, eval_batches,
+   input_name, input_shape) = train_and_eval_batches(
+       dataset_name, data_dir)
+
+  def train_input_fn():
+    return dataset_to_stream(train_batches, input_name)
+
+  def eval_input_fn():
+    return dataset_to_stream(eval_batches, input_name)
+
+  return Inputs(train_fn=train_input_fn, eval_fn=eval_input_fn,
+                input_shape=input_shape)
+
+
+def dataset_to_stream(dataset, input_name):
+  """Takes a tf.Dataset and creates a numpy stream of ready batches."""
+  for example in tfds.as_numpy(dataset):
+    inp, out = example[0][input_name], example[1]
+    if len(out.shape) > 1 and out.shape[-1] == 1:
+      out = np.squeeze(out, axis=-1)
+    yield inp, out
+
+
 def train_and_eval_dataset(dataset_name, data_dir):
   """Return train and evaluation datasets, feature info and supervised keys.
 
diff --git a/tensor2tensor/jax/j2j.py b/tensor2tensor/jax/j2j.py
index 1768bbb29..0ca5accc5 100644
--- a/tensor2tensor/jax/j2j.py
+++ b/tensor2tensor/jax/j2j.py
@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
+import functools
 import os
 import pickle
 import time
@@ -30,15 +32,13 @@
 from jax.experimental import optimizers
 import jax.numpy as np
 
-from tensor2tensor.jax import input_pipeline
+from tensor2tensor.jax import inputs as inputs_lib
 from tensor2tensor.jax import jaxboard
 # Import for gin configurable models
 from tensor2tensor.jax import models  # pylint: disable=unused-import
 
 from tensorflow.io import gfile
 
-import tensorflow_datasets as tfds
-
 
 @gin.configurable(blacklist=["step"])
 def learning_rate(step,
@@ -102,60 +102,54 @@ def loss(params, batch, model_predict):
   return - np.mean(preds * one_hot(targets, preds.shape[-1]))
 
 
-def dataset_to_stream(batches, input_name):
-  """Takes a tf.Dataset and creates a numpy stream of ready batches."""
-  for example in tfds.as_numpy(batches):
-    inp, out = example[0][input_name], example[1]
-    if len(out.shape) > 1 and out.shape[-1] == 1:
-      out = np.squeeze(out, axis=-1)
-    yield inp, out
-
-
 def log(s, stdout=True):
   logging.info(s)
   if stdout:
     print(s)
 
 
-def _make_directory(path):
-  """Helper function: create directory if it doesn't exist yet."""
-  if not gfile.exists(path):
-    log("Creating directory %s" % path)
-    gfile.mkdir(path)
+def step_log(step, s):
+  log("Step % 6d: %s" % (step, s))
 
 
-def save_params_and_step(params, step, output_dir):
-  """Save params and step in output dir."""
-  if output_dir is not None:
-    _make_directory(output_dir)
-    params_file = os.path.join(output_dir, "model.pkl")
-    with gfile.GFile(params_file, "wb") as f:
-      pickle.dump((params, step), f)
-    log("Model saved to %s" % params_file, stdout=False)
+State = collections.namedtuple("_State", ["step", "params"])
 
 
-def load_params_and_step(output_dir):
-  """Save params and step in output dir."""
-  if output_dir is None:
-    return None, None
-  if not gfile.exists(output_dir):
-    return None, None
+def restore_state(output_dir):
+  """Restore State."""
+  empty_state = State(step=None, params=None)
+  if output_dir is None or not gfile.exists(output_dir):
+    return empty_state
   params_file = os.path.join(output_dir, "model.pkl")
   if not gfile.exists(params_file):
-    return None, None
-  with gfile.GFile(params_file, "r") as f:
+    return empty_state
+
+  with gfile.GFile(params_file, "rb") as f:
     (params, step) = pickle.load(f)
   log("Model loaded from %s" % params_file)
-  return params, step
+  return State(step=step, params=params)
 
 
-def _make_summary_writer(output_path):
-  _make_directory(output_path)
-  return jaxboard.SummaryWriter(output_path)
+def save_state(state, output_dir, save_gin=True):
+  """Save State and optionally gin config."""
+  if not output_dir:
+    return
+
+  gfile.makedirs(output_dir)
+  params_file = os.path.join(output_dir, "model.pkl")
+  with gfile.GFile(params_file, "wb") as f:
+    pickle.dump((state.params, state.step), f)
+  log("Model saved to %s" % params_file, stdout=False)
+
+  # Gin file only includes used parameters, so we save it at this point.
+  if save_gin:
+    config_path = os.path.join(output_dir, "config.gin")
+    with gfile.GFile(config_path, "w") as f:
+      f.write(gin.operative_config_str())
 
 
 # Metrics to calculate and report.
-_metrics = {
+_METRICS = {
     "accuracy": accuracy,
     "neg_log_perplexity": neg_log_perplexity,
     "loss": lambda x, y: - neg_log_perplexity(x, y),
@@ -163,12 +157,16 @@ def _make_summary_writer(output_path):
 
 
 # We include in gin config everything that could be useful to share between
-# users, so when it gets saved in a .gin file it can be re-ran with few flags.
+# users, so when it gets saved in a .gin file it can be re-run with minimal
+# flags.
 @gin.configurable(blacklist=["data_dir", "output_dir"])
-def train_fn(data_dir=None, output_dir=None,
+def train_fn(data_dir=None,
+             output_dir=None,
              model=gin.REQUIRED,
              dataset=gin.REQUIRED,
-             train_steps=1000, eval_steps=10, eval_frequency=100):
+             train_steps=1000,
+             eval_steps=10,
+             eval_frequency=100):
   """Train the given model on the given dataset.
 
   Args:
@@ -180,101 +178,112 @@ def train_fn(data_dir=None, output_dir=None,
     eval_steps: for how many steps to do evaluation.
     eval_frequency: how often (every this many steps) to run evaluation.
   """
-  (train_batches, eval_batches,
-   input_name, input_shape) = input_pipeline.train_and_eval_batches(
-       dataset, data_dir)
-  train_stream = dataset_to_stream(train_batches, input_name)
+  # Make Inputs
+  inputs = inputs_lib.make_inputs(dataset, data_dir)
 
-  # Training loop.
+  # Setup optimizer and model
   opt_init, opt_update = optimizer()
   model_init, model_predict = model()
-
+  jit_predict = jax.jit(model_predict)
+
+  # Setup state
+  state = restore_state(output_dir)
+  step = state.step or 0
+  _, init_params = model_init([-1] + inputs.input_shape)
+  opt_state = opt_init(state.params or init_params)
+
+  # Create summary writers.
+  train_sw, eval_sw = None, None
+  if output_dir:
+    gfile.makedirs(output_dir)
+    train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
+    eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
+
+  # Make fast update function
   @jax.jit
   def update(i, opt_state, batch):
     params = optimizers.get_params(opt_state)
     return opt_update(i, jax.grad(loss)(
         params, batch, model_predict), opt_state)
 
-  _, init_params = model_init([-1] + input_shape)
-  step, train_sw, eval_sw = 0, None, None
-  if output_dir is not None:
-    _make_directory(output_dir)
-    # Load parameters.
-    loaded_params, loaded_step = load_params_and_step(output_dir)
-    if loaded_params is not None:
-      init_params = loaded_params
-    if loaded_step is not None:
-      step = loaded_step
-
-    # Create summary writers.
-    eval_sw = _make_summary_writer(os.path.join(output_dir, "eval"))
-    train_sw = _make_summary_writer(os.path.join(output_dir, "train"))
-
-  log("Starting training.")
-  opt_state = opt_init(init_params)
-  gin_config_saved = False
-  cur_eval_frequency = 1  # First evaluation after the first training step.
+  step_log(step, "starting training")
+  train_gen = inputs.train_fn()
+  is_first_step = True
+  epoch_steps = 1  # First evaluation after the first training step.
   while step < train_steps:
-    # Training.
+    # Train
     start_time = time.time()
-    for _ in range(cur_eval_frequency):
-      opt_state = update(step, opt_state, next(train_stream))
+    for _ in range(epoch_steps):
+      opt_state = update(step, opt_state, next(train_gen))
       if train_sw and step % 10 == 0:  # Log learning rate curve each 10 steps.
         train_sw.scalar("training/learning rate",
                         learning_rate(step), step=step)
       step += 1
     epoch_time = time.time() - start_time
-    log("Step {}, last {} steps in {:0.2f} sec".format(
-        step, cur_eval_frequency, epoch_time))
+    print()
+    step_log(step, "%d train steps in %0.2f secs" % (epoch_steps, epoch_time))
 
-    # Save the model.
+    # Save state
     params = optimizers.get_params(opt_state)
-    save_params_and_step(params, step, output_dir)
-
-    # Save the config if not saved yet.
-    # Gin file only includes used parameters, so we save it at this point.
-    if output_dir and not gin_config_saved:
-      gin_config_saved = True
-      config_path = os.path.join(output_dir, "config.gin")
-      with gfile.GFile(config_path, "w") as f:
-        f.write(gin.operative_config_str())
-
-    # Evaluation.
-    eval_stream = dataset_to_stream(eval_batches, input_name)
-    eval_train_stream = dataset_to_stream(train_batches, input_name)
-    train_metrics = {key: 0.0 for key in _metrics}
-    eval_metrics = {key: 0.0 for key in _metrics}
-    for _ in range(eval_steps):
-      train_batch = next(eval_train_stream)
-      train_predictions = model_predict(params, train_batch[0])
-      eval_batch = next(eval_stream)
-      eval_predictions = model_predict(params, eval_batch[0])
-      for m in _metrics:
-        train_metrics[m] += _metrics[m](
-            train_batch, train_predictions) / float(eval_steps)
-        eval_metrics[m] += _metrics[m](
-            eval_batch, eval_predictions) / float(eval_steps)
-
-    for m in _metrics:
-      log("Step %d train %s %.8f" % (step, m, train_metrics[m]))
+    save_state(State(params=params, step=step), output_dir,
+               save_gin=is_first_step)
+
+    # Evaluate
+    print()
+    step_log(step, "starting evaluation")
+    train_metrics, eval_metrics = evaluate(
+        inputs, functools.partial(jit_predict, params), eval_steps)
+
+    for m in _METRICS:
+      step_log(step, "train %s %.8f" % (m, train_metrics[m]))
       prefix = "metrics/"
       if train_sw:
         train_sw.scalar(prefix + m, train_metrics[m], step=step)
-      log("Step %d eval  %s %.8f" % (step, m, eval_metrics[m]))
+      step_log(step, "eval  %s %.8f" % (m, eval_metrics[m]))
       if eval_sw:
         eval_sw.scalar(prefix + m, eval_metrics[m], step=step)
 
     # Log non-metric reports and flush.
-    if train_sw:
-      if step > 1:  # Don't log performance of the first step.
-        train_sw.scalar("training/steps per second",
-                        cur_eval_frequency / epoch_time, step=step)
+    if train_sw and not is_first_step:
+      train_sw.scalar("training/steps per second",
+                      epoch_steps / epoch_time, step=step)
       train_sw.writer.flush()
     if eval_sw:
       eval_sw.writer.flush()
 
-    # After the first step, Evaluate every eval_frequency steps.
-    if cur_eval_frequency == 1 and eval_frequency != 1:
-      cur_eval_frequency = eval_frequency - 1
-    else:
-      cur_eval_frequency = eval_frequency
+    # After the first step, train for eval_frequency steps before evaluating
+    epoch_steps = (eval_frequency - 1) if is_first_step else eval_frequency
+    is_first_step = False
+
+  step_log(step, "finished training.")
+
+
+def evaluate(inputs, predict_fn, eval_steps):
+  """Evaluate.
+
+  Args:
+    inputs: Inputs namedtuple.
+    predict_fn: function from inputs to predictions. params should already be
+      partially applied.
+    eval_steps: int, number of evaluation steps.
+
+  Returns:
+    train_metrics: dict
+    eval_metrics: dict
+  """
+  eval_stream = inputs.eval_fn()
+  eval_train_stream = inputs.train_fn()
+  train_metrics = {key: 0.0 for key in _METRICS}
+  eval_metrics = {key: 0.0 for key in _METRICS}
+  for _ in range(eval_steps):
+    train_batch = next(eval_train_stream)
+    train_predictions = predict_fn(train_batch[0])
+    eval_batch = next(eval_stream)
+    eval_predictions = predict_fn(eval_batch[0])
+    for m in _METRICS:
+      train_metrics[m] += _METRICS[m](
+          train_batch, train_predictions) / float(eval_steps)
+      eval_metrics[m] += _METRICS[m](
+          eval_batch, eval_predictions) / float(eval_steps)
+
+  return train_metrics, eval_metrics
diff --git a/tensor2tensor/jax/j2j_trainer.py b/tensor2tensor/jax/j2j_trainer.py
index 7ae480421..4b9b4f761 100644
--- a/tensor2tensor/jax/j2j_trainer.py
+++ b/tensor2tensor/jax/j2j_trainer.py
@@ -17,9 +17,13 @@
 
 Examples:
 
-- train a basic model on mnist:
-    jax/j2j_trainer.py --dataset=mnist --model=mlp
-      --config="train_fn.train_steps=4000" --output_dir ~/j2j/test1
+
+MLP on mnist:
+  python -m tensor2tensor.jax.j2j_trainer \
+    --dataset=mnist \
+    --model=MLP \
+    --config="train_fn.train_steps=4000" \
+    --output_dir=~/j2j/test1
 """
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/jax/models/__init__.py b/tensor2tensor/jax/models/__init__.py
index 1769e36ac..c663e829e 100644
--- a/tensor2tensor/jax/models/__init__.py
+++ b/tensor2tensor/jax/models/__init__.py
@@ -18,7 +18,5 @@
 from __future__ import division
 from __future__ import print_function
 
-# pylint: disable=unused-import
 from tensor2tensor.jax.models import mlp
 from tensor2tensor.jax.models import resnet
-# pylint: enable=unused-import

From 2bae6e870afe6fe44fc3a34c26a1f66c5fc402e4 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 26 Feb 2019 13:06:26 -0800
Subject: [PATCH 1721/2720] Add VideoFlow hparams for the Stochastic Movement
 Dataset and BAIR qualitative adn quantitative results.

PiperOrigin-RevId: 235777401
---
 tensor2tensor/models/video/next_frame_glow.py | 51 +++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index 56a9da10b..e734ed37f 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -80,6 +80,57 @@ def next_frame_glow_hparams():
   return hparams
 
 
+@registry.register_hparams
+def next_frame_glow_bair_quant():
+  """Hparams to reproduce bits-per-pixel results on BAIR action-free dataset."""
+  hparams = next_frame_glow_hparams()
+  hparams.video_num_input_frames = 3
+  hparams.video_num_target_frames = 10
+  hparams.num_train_frames = 4
+  hparams.num_cond_latents = 3
+  hparams.depth = 24
+  hparams.latent_dist_encoder = "conv3d_net"
+  hparams.latent_encoder_width = 256
+  hparams.latent_architecture = "glow_resnet"
+  hparams.latent_encoder_depth = 5
+  hparams.latent_apply_dilations = True
+  hparams.latent_activation = "gatu"
+  hparams.activation = "gatu"
+  hparams.learning_rate_constant = 3e-4
+  hparams.learning_rate_schedule = "constant*linear_warmup"
+  hparams.learning_rate_warmup_steps = 10000
+  hparams.init_batch_size = 128
+  hparams.batch_size = 5
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_glow_bair_qual():
+  """Hparams for qualitative video generation results."""
+  hparams = next_frame_glow_bair_quant()
+  hparams.coupling = "additive"
+  hparams.temperature = 0.5
+  hparams.coupling_width = 392
+  return hparams
+
+
+@registry.register_hparams
+def next_frame_glow_shapes():
+  """Hparams for qualitative and quantitative results on shapes dataset."""
+  hparams = next_frame_glow_bair_quant()
+  hparams.video_num_input_frames = 1
+  hparams.video_num_target_frames = 2
+  hparams.num_train_frames = 2
+  hparams.num_cond_latents = 1
+  hparams.coupling = "additive"
+  hparams.coupling_width = 512
+  hparams.latent_encoder_depth = 10
+  hparams.latent_skip = False
+  hparams.learning_rate_constant = 1e-4
+  hparams.batch_size = 10
+  return hparams
+
+
 @registry.register_hparams
 def frame_glow_hparams():
   """Unconditional generation on video-frames."""

From bc40f3de6ee2228f610256832a2e8fb810f9bb27 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Tue, 26 Feb 2019 13:37:52 -0800
Subject: [PATCH 1722/2720] Require data_dir and output_dir; make default
 output_dir Improve metrics logging

PiperOrigin-RevId: 235783250
---
 tensor2tensor/bin/t2t_trainer.py |  2 +-
 tensor2tensor/jax/j2j.py         | 87 ++++++++++++++++----------------
 tensor2tensor/jax/j2j_trainer.py | 27 ++++++++--
 3 files changed, 68 insertions(+), 48 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index f63c22ecb..89cefa2fc 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -386,7 +386,7 @@ def main(argv):
     if FLAGS.hparams_set:
       config_files = [os.path.expanduser(FLAGS.hparams_set)]
     gin.parse_config_files_and_bindings(config_files, config_strs)
-    j2j.train_fn(data_dir, output_dir=output_dir)
+    j2j.train_fn(data_dir=data_dir, output_dir=output_dir)
     return
 
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
diff --git a/tensor2tensor/jax/j2j.py b/tensor2tensor/jax/j2j.py
index 0ca5accc5..f7c7b4927 100644
--- a/tensor2tensor/jax/j2j.py
+++ b/tensor2tensor/jax/j2j.py
@@ -32,6 +32,8 @@
 from jax.experimental import optimizers
 import jax.numpy as np
 
+import six
+
 from tensor2tensor.jax import inputs as inputs_lib
 from tensor2tensor.jax import jaxboard
 # Import for gin configurable models
@@ -117,12 +119,9 @@ def step_log(step, s):
 
 def restore_state(output_dir):
   """Restore State."""
-  empty_state = State(step=None, params=None)
-  if output_dir is None or not gfile.exists(output_dir):
-    return empty_state
   params_file = os.path.join(output_dir, "model.pkl")
   if not gfile.exists(params_file):
-    return empty_state
+    return State(step=None, params=None)
 
   with gfile.GFile(params_file, "rb") as f:
     (params, step) = pickle.load(f)
@@ -132,10 +131,6 @@ def restore_state(output_dir):
 
 def save_state(state, output_dir, save_gin=True):
   """Save State and optionally gin config."""
-  if not output_dir:
-    return
-
-  gfile.makedirs(output_dir)
   params_file = os.path.join(output_dir, "model.pkl")
   with gfile.GFile(params_file, "wb") as f:
     pickle.dump((state.params, state.step), f)
@@ -160,8 +155,8 @@ def save_state(state, output_dir, save_gin=True):
 # users, so when it gets saved in a .gin file it can be re-run with minimal
 # flags.
 @gin.configurable(blacklist=["data_dir", "output_dir"])
-def train_fn(data_dir=None,
-             output_dir=None,
+def train_fn(output_dir,
+             data_dir,
              model=gin.REQUIRED,
              dataset=gin.REQUIRED,
              train_steps=1000,
@@ -170,58 +165,64 @@ def train_fn(data_dir=None,
   """Train the given model on the given dataset.
 
   Args:
-    data_dir: Directory where the data is located.
     output_dir: Directory where to put the logs and checkpoints.
-    model: The model to train (a function).
-    dataset: The name of the dataset to train on.
-    train_steps: for how many steps to train.
-    eval_steps: for how many steps to do evaluation.
-    eval_frequency: how often (every this many steps) to run evaluation.
+    data_dir: Directory where the data is located.
+    model: The model to train as a callable returning 2 callables, an init_fun
+      and apply_fun.
+    dataset: The name of the TFDS dataset to train on. To train on a T2T
+      dataset, prefix the name with "t2t_".
+    train_steps: int, total number of training steps.
+    eval_steps: int, num of steps per evaluation.
+    eval_frequency: int, how often to run evaluation (every eval_frequency
+      steps).
   """
+  gfile.makedirs(output_dir)
+
   # Make Inputs
   inputs = inputs_lib.make_inputs(dataset, data_dir)
 
   # Setup optimizer and model
   opt_init, opt_update = optimizer()
   model_init, model_predict = model()
-  jit_predict = jax.jit(model_predict)
 
   # Setup state
   state = restore_state(output_dir)
   step = state.step or 0
-  _, init_params = model_init([-1] + inputs.input_shape)
-  opt_state = opt_init(state.params or init_params)
+  params_initializer = lambda: model_init([-1] + inputs.input_shape)[1]
+  opt_state = opt_init(state.params or params_initializer())
 
   # Create summary writers.
-  train_sw, eval_sw = None, None
-  if output_dir:
-    gfile.makedirs(output_dir)
-    train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
-    eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
+  train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
+  eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
+
+  # jit model_predict and update so they're fast
+  jit_predict = jax.jit(model_predict)  # for evaluation
 
-  # Make fast update function
   @jax.jit
   def update(i, opt_state, batch):
     params = optimizers.get_params(opt_state)
     return opt_update(i, jax.grad(loss)(
         params, batch, model_predict), opt_state)
 
+  print()
   step_log(step, "starting training")
   train_gen = inputs.train_fn()
   is_first_step = True
   epoch_steps = 1  # First evaluation after the first training step.
   while step < train_steps:
+    print()
+
     # Train
     start_time = time.time()
     for _ in range(epoch_steps):
       opt_state = update(step, opt_state, next(train_gen))
-      if train_sw and step % 10 == 0:  # Log learning rate curve each 10 steps.
+      if step % 10 == 0:  # Log learning rate curve each 10 steps.
         train_sw.scalar("training/learning rate",
                         learning_rate(step), step=step)
       step += 1
     epoch_time = time.time() - start_time
-    print()
-    step_log(step, "%d train steps in %0.2f secs" % (epoch_steps, epoch_time))
+    step_log(step, "ran %d train steps in %0.2f secs" %
+             (epoch_steps, epoch_time))
 
     # Save state
     params = optimizers.get_params(opt_state)
@@ -229,33 +230,25 @@ def update(i, opt_state, batch):
                save_gin=is_first_step)
 
     # Evaluate
-    print()
     step_log(step, "starting evaluation")
     train_metrics, eval_metrics = evaluate(
         inputs, functools.partial(jit_predict, params), eval_steps)
-
-    for m in _METRICS:
-      step_log(step, "train %s %.8f" % (m, train_metrics[m]))
-      prefix = "metrics/"
-      if train_sw:
-        train_sw.scalar(prefix + m, train_metrics[m], step=step)
-      step_log(step, "eval  %s %.8f" % (m, eval_metrics[m]))
-      if eval_sw:
-        eval_sw.scalar(prefix + m, eval_metrics[m], step=step)
+    log_metrics(train_metrics, train_sw, "train", step)
+    log_metrics(eval_metrics, eval_sw, "eval ", step)
 
     # Log non-metric reports and flush.
-    if train_sw and not is_first_step:
+    if not is_first_step:
       train_sw.scalar("training/steps per second",
                       epoch_steps / epoch_time, step=step)
-      train_sw.writer.flush()
-    if eval_sw:
-      eval_sw.writer.flush()
+    train_sw.writer.flush()
+    eval_sw.writer.flush()
 
     # After the first step, train for eval_frequency steps before evaluating
     epoch_steps = (eval_frequency - 1) if is_first_step else eval_frequency
     is_first_step = False
 
-  step_log(step, "finished training.")
+  print()
+  step_log(step, "finished training")
 
 
 def evaluate(inputs, predict_fn, eval_steps):
@@ -287,3 +280,11 @@ def evaluate(inputs, predict_fn, eval_steps):
           eval_batch, eval_predictions) / float(eval_steps)
 
   return train_metrics, eval_metrics
+
+
+def log_metrics(metrics, summ_writer, log_prefix, step):
+  rjust_len = max([len(name) for name in metrics])
+  for name, value in six.iteritems(metrics):
+    step_log(step, "%s %s | % .8f" % (log_prefix, name.rjust(rjust_len), value))
+    if summ_writer:
+      summ_writer.scalar("metrics/" + name, value, step)
diff --git a/tensor2tensor/jax/j2j_trainer.py b/tensor2tensor/jax/j2j_trainer.py
index 4b9b4f761..0e09a8d7b 100644
--- a/tensor2tensor/jax/j2j_trainer.py
+++ b/tensor2tensor/jax/j2j_trainer.py
@@ -29,6 +29,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import datetime
 import os
 
 from absl import app
@@ -61,15 +62,33 @@ def _setup_gin():
   gin.parse_config_files_and_bindings(FLAGS.config_file, configs)
 
 
+def _default_output_dir():
+  """Default output directory."""
+  dir_name = "{model_name}_{dataset_name}_{timestamp}".format(
+      model_name=gin.query_parameter("train_fn.model").configurable.name,
+      dataset_name=gin.query_parameter("train_fn.dataset"),
+      timestamp=datetime.datetime.now().strftime("%Y%m%d_%H%M"),
+  )
+  dir_path = os.path.join("~", "j2j", dir_name)
+  print()
+  j2j.log("No output_dir specified")
+  return dir_path
+
+
 def main(_):
   _setup_gin()
 
   # Setup directories
-  data_dir, output_dir = FLAGS.data_dir, FLAGS.output_dir
-  data_dir = data_dir and os.path.expanduser(data_dir)
-  output_dir = output_dir and os.path.expanduser(output_dir)
+  data_dir = FLAGS.data_dir
+  output_dir = FLAGS.output_dir or _default_output_dir()
+  assert data_dir, "Must specify a data directory"
+  assert output_dir, "Must specify an output directory"
+  j2j.log("Using output_dir %s" % output_dir)
+
+  data_dir = os.path.expanduser(data_dir)
+  output_dir = os.path.expanduser(output_dir)
 
-  j2j.train_fn(data_dir, output_dir=output_dir)
+  j2j.train_fn(data_dir=data_dir, output_dir=output_dir)
 
 
 if __name__ == "__main__":

From 92d5dd0503a3ce98fc5ce0f88fb64c96a44995c0 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 26 Feb 2019 15:48:26 -0800
Subject: [PATCH 1723/2720] Add a version of wikitext103 with longer context
 (4k tokens)

PiperOrigin-RevId: 235809389
---
 tensor2tensor/data_generators/wikitext103.py | 30 ++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/tensor2tensor/data_generators/wikitext103.py b/tensor2tensor/data_generators/wikitext103.py
index 1a25b2984..6c6f9eb1f 100644
--- a/tensor2tensor/data_generators/wikitext103.py
+++ b/tensor2tensor/data_generators/wikitext103.py
@@ -161,3 +161,33 @@ class LanguagemodelWikitext103Characters(LanguagemodelWikitext103):
   @property
   def vocab_type(self):
     return text_problems.VocabType.CHARACTER
+
+
+@registry.register_problem
+class LanguagemodelWikitext103L4k(LanguagemodelWikitext103):
+  """Wikitext-103, token-level, with examples up to 4,096 tokens long."""
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    samples_by_line = super(LanguagemodelWikitext103L4k,
+                            self).generate_samples(data_dir, tmp_dir,
+                                                   dataset_split)
+
+    def _generate_samples():
+      tokens = []
+      for sample in samples_by_line:
+        sample_tokens = sample["targets"].split()
+        if len(tokens) + len(sample_tokens) < self.sequence_length:
+          tokens.extend(sample_tokens)
+        else:
+          yield {"targets": " ".join(tokens)}
+          tokens = sample_tokens
+
+    return _generate_samples()
+
+  def max_length(self, model_hparams):
+    return model_hparams.split_to_length or self.sequence_length
+
+  @property
+  def sequence_length(self):
+    """Length of each example (in tokens)."""
+    return 4096

From d06ce2581a82623ef6fda5996618b63967a4d5a1 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 26 Feb 2019 22:08:24 -0800
Subject: [PATCH 1724/2720] Correct bug in loss computation in J2J (don't
 average over vocab size), adjust learning rate accordingly.

PiperOrigin-RevId: 235854726
---
 tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin | 8 ++++----
 tensor2tensor/jax/j2j.py                            | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin
index c825e0cbc..9483f47c8 100644
--- a/tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin
@@ -3,11 +3,11 @@
 batch_fn.batch_size = 32
 batch_fn.bucket_length = 32
 batch_fn.buckets = None
-batch_fn.eval_batch_size = 16
+batch_fn.eval_batch_size = 32
 
 # Parameters for learning_rate:
 # ==============================================================================
-learning_rate.constant = 10.0
+learning_rate.constant = 5.0
 learning_rate.schedule = 'constant * linear_warmup * rsqrt_decay'
 learning_rate.warmup_steps = 400
 
@@ -34,6 +34,6 @@ Resnet50.num_output_classes = 1001
 # ==============================================================================
 train_fn.dataset = 't2t_image_imagenet224'
 train_fn.eval_frequency = 100
-train_fn.eval_steps = 2
+train_fn.eval_steps = 3
 train_fn.model = @Resnet50
-train_fn.train_steps = 100000
+train_fn.train_steps = 200000
diff --git a/tensor2tensor/jax/j2j.py b/tensor2tensor/jax/j2j.py
index f7c7b4927..d67da5fdc 100644
--- a/tensor2tensor/jax/j2j.py
+++ b/tensor2tensor/jax/j2j.py
@@ -101,7 +101,7 @@ def loss(params, batch, model_predict):
   """Calculate loss."""
   inputs, targets = batch
   preds = model_predict(params, inputs)
-  return - np.mean(preds * one_hot(targets, preds.shape[-1]))
+  return - np.mean(np.sum(preds * one_hot(targets, preds.shape[-1]), axis=-1))
 
 
 def log(s, stdout=True):

From 236bf310c0b80490052c3c828bc24d1be5994fb6 Mon Sep 17 00:00:00 2001
From: Mike Dusenberry <dusenberrymw@google.com>
Date: Wed, 27 Feb 2019 12:42:05 -0800
Subject: [PATCH 1725/2720] Fixed a serialization bug in the Bayesian
 initializers.

PiperOrigin-RevId: 235972599
---
 tensor2tensor/layers/bayes.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 18c799ff1..d6e2162f1 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -318,7 +318,7 @@ def __init__(self, seed=None, dtype=tf.float32):
   def get_config(self):
     return {
         'seed': self.seed,
-        'dtype': self.dtype.name
+        'dtype': self.dtype,
     }
 
 
@@ -345,7 +345,7 @@ def __init__(self, seed=None, dtype=tf.float32):
   def get_config(self):
     return {
         'seed': self.seed,
-        'dtype': self.dtype.name
+        'dtype': self.dtype
     }
 
 

From cbe5c82e28aebd9449693f7e6ea03b21ad4f10fd Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 27 Feb 2019 16:01:43 -0800
Subject: [PATCH 1726/2720] j2j is now trax Make models and optimizers external
 configurables

PiperOrigin-RevId: 236010981
---
 tensor2tensor/bin/t2t_trainer.py              | 16 ++---
 tensor2tensor/trax/README.md                  | 37 +++++++++++
 tensor2tensor/{jax => trax}/__init__.py       |  0
 .../configs/resnet50_imagenet_8gb.gin         | 23 +++----
 tensor2tensor/{jax => trax}/inputs.py         |  2 +-
 tensor2tensor/{jax => trax}/jaxboard.py       |  0
 tensor2tensor/{jax => trax}/jaxboard_demo.py  |  0
 .../{jax => trax}/models/__init__.py          | 12 +++-
 tensor2tensor/{jax => trax}/models/mlp.py     |  5 +-
 tensor2tensor/{jax => trax}/models/resnet.py  |  5 +-
 tensor2tensor/trax/optimizers.py              | 41 ++++++++++++
 .../{jax/j2j_trainer.py => trax/trainer.py}   | 30 +++------
 tensor2tensor/{jax/j2j.py => trax/trax.py}    | 64 +++++++++----------
 13 files changed, 147 insertions(+), 88 deletions(-)
 create mode 100644 tensor2tensor/trax/README.md
 rename tensor2tensor/{jax => trax}/__init__.py (100%)
 rename tensor2tensor/{jax => trax}/configs/resnet50_imagenet_8gb.gin (69%)
 rename tensor2tensor/{jax => trax}/inputs.py (99%)
 rename tensor2tensor/{jax => trax}/jaxboard.py (100%)
 rename tensor2tensor/{jax => trax}/jaxboard_demo.py (100%)
 rename tensor2tensor/{jax => trax}/models/__init__.py (74%)
 rename tensor2tensor/{jax => trax}/models/mlp.py (95%)
 rename tensor2tensor/{jax => trax}/models/resnet.py (98%)
 create mode 100644 tensor2tensor/trax/optimizers.py
 rename tensor2tensor/{jax/j2j_trainer.py => trax/trainer.py} (76%)
 rename tensor2tensor/{jax/j2j.py => trax/trax.py} (85%)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 89cefa2fc..db223d899 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -39,7 +39,7 @@
 from tensorflow.contrib.tpu.python.tpu import tpu_config
 
 try:
-  from tensor2tensor.jax import j2j  # pylint: disable=g-import-not-at-top
+  from tensor2tensor.trax import trax  # pylint: disable=g-import-not-at-top
 except (TypeError, ImportError):
   pass
 
@@ -77,7 +77,7 @@
 flags.DEFINE_integer("intra_op_parallelism_threads", 0,
                      "Number of intra_op_parallelism_threads to use for CPU. "
                      "See TensorFlow config.proto for details.")
-flags.DEFINE_bool("jax", False, "Whether to use J2J.")
+flags.DEFINE_bool("jax", False, "Whether to use trax.")
 # TODO(lukaszkaiser): resolve memory and variable assign issues and set to True.
 flags.DEFINE_bool(
     "optionally_use_dist_strat", False,
@@ -370,23 +370,23 @@ def main(argv):
     # Hacking main v1 flags to work with jax.
     config_strs = []
     config_strs.append(
-        "train_fn.train_steps=" + str(FLAGS.train_steps))
+        "train.train_steps=" + str(FLAGS.train_steps))
     config_strs.append(
-        "train_fn.eval_steps=" + str(FLAGS.eval_steps))
+        "train.eval_steps=" + str(FLAGS.eval_steps))
     config_strs.append(
-        "train_fn.eval_frequency=" + str(FLAGS.local_eval_frequency))
+        "train.eval_frequency=" + str(FLAGS.local_eval_frequency))
     if FLAGS.hparams:
       config_strs.extend(str(FLAGS.hparams).split(","))
     data_dir = os.path.expanduser(FLAGS.data_dir)
     output_dir = os.path.expanduser(FLAGS.output_dir)
 
-    gin.bind_parameter("train_fn.dataset", FLAGS.problem)
-    config_strs += ["train_fn.model=@" + FLAGS.model]
+    gin.bind_parameter("train.dataset", FLAGS.problem)
+    config_strs += ["train.model=@" + FLAGS.model]
     config_files = []
     if FLAGS.hparams_set:
       config_files = [os.path.expanduser(FLAGS.hparams_set)]
     gin.parse_config_files_and_bindings(config_files, config_strs)
-    j2j.train_fn(data_dir=data_dir, output_dir=output_dir)
+    trax.train(data_dir=data_dir, output_dir=output_dir)
     return
 
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
diff --git a/tensor2tensor/trax/README.md b/tensor2tensor/trax/README.md
new file mode 100644
index 000000000..d5c269f73
--- /dev/null
+++ b/tensor2tensor/trax/README.md
@@ -0,0 +1,37 @@
+# `trax`: Train Neural Nets with JAX
+
+![train tracks](https://images.pexels.com/photos/461772/pexels-photo-461772.jpeg?dl&fit=crop&crop=entropy&w=640&h=426)
+
+* Configuration is done with [`gin`](https://github.com/google/gin-config).
+  `trainer.py` takes `--config_file` as well as `--config` for file overrides.
+* Models are defined with [`stax`](https://github.com/google/jax/blob/master/jax/experimental/stax.py) in
+  `models/`. They are made gin-configurable in `models/__init__.py`.
+* Datasets are simple iterators over batches. Datasets from
+  [`tensorflow/datasets`](https://github.com/tensorflow/datasets)
+  and [`tensor2tensor`](https://github.com/tensorflow/tensor2tensor)
+  are built-in and can be addressed by name.
+
+Entrypoints:
+
+* Script: `trainer.py`
+* Main library entrypoint: `trax.train`
+
+### Examples
+
+#### MLP on MNIST
+
+
+```
+python -m tensor2tensor.trax.trainer \
+  --dataset=mnist \
+  --model=MLP \
+  --config="train.train_steps=1000"
+```
+
+#### Resnet50 on Imagenet
+
+
+```
+python -m tensor2tensor.trax.trainer \
+  --config_file=$PWD/trax/configs/resnet50_imagenet_8gb.gin
+```
diff --git a/tensor2tensor/jax/__init__.py b/tensor2tensor/trax/__init__.py
similarity index 100%
rename from tensor2tensor/jax/__init__.py
rename to tensor2tensor/trax/__init__.py
diff --git a/tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
similarity index 69%
rename from tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin
rename to tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
index 9483f47c8..049a5e169 100644
--- a/tensor2tensor/jax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
@@ -11,15 +11,9 @@ learning_rate.constant = 5.0
 learning_rate.schedule = 'constant * linear_warmup * rsqrt_decay'
 learning_rate.warmup_steps = 400
 
-# Parameters for optimizer:
+# Parameters for momentum:
 # ==============================================================================
-optimizer.adam_b1 = 0.9
-optimizer.adam_b2 = 0.997
-optimizer.adam_eps = 1e-08
-optimizer.momentum_mass = 0.9
-optimizer.name = 'momentum'
-optimizer.rmsprop_eps = 1e-08
-optimizer.rmsprop_gamma = 0.9
+momentum.mass = 0.9
 
 # Parameters for preprocess_fn:
 # ==============================================================================
@@ -30,10 +24,11 @@ preprocess_fn.max_target_length = -1
 Resnet50.hidden_size = 64
 Resnet50.num_output_classes = 1001
 
-# Parameters for train_fn:
+# Parameters for train:
 # ==============================================================================
-train_fn.dataset = 't2t_image_imagenet224'
-train_fn.eval_frequency = 100
-train_fn.eval_steps = 3
-train_fn.model = @Resnet50
-train_fn.train_steps = 200000
+train.dataset = 't2t_image_imagenet224'
+train.eval_frequency = 100
+train.eval_steps = 3
+train.model = @models.Resnet50
+train.optimizer = @optimizers.momentum
+train.train_steps = 200000
diff --git a/tensor2tensor/jax/inputs.py b/tensor2tensor/trax/inputs.py
similarity index 99%
rename from tensor2tensor/jax/inputs.py
rename to tensor2tensor/trax/inputs.py
index 457534b6c..c82007f95 100644
--- a/tensor2tensor/jax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""J2J input pipeline."""
+"""trax input pipeline."""
 
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/jax/jaxboard.py b/tensor2tensor/trax/jaxboard.py
similarity index 100%
rename from tensor2tensor/jax/jaxboard.py
rename to tensor2tensor/trax/jaxboard.py
diff --git a/tensor2tensor/jax/jaxboard_demo.py b/tensor2tensor/trax/jaxboard_demo.py
similarity index 100%
rename from tensor2tensor/jax/jaxboard_demo.py
rename to tensor2tensor/trax/jaxboard_demo.py
diff --git a/tensor2tensor/jax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
similarity index 74%
rename from tensor2tensor/jax/models/__init__.py
rename to tensor2tensor/trax/models/__init__.py
index c663e829e..21c555f5e 100644
--- a/tensor2tensor/jax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -13,10 +13,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Models defined in J2J."""
+"""Models defined in trax."""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.jax.models import mlp
-from tensor2tensor.jax.models import resnet
+import gin
+
+from tensor2tensor.trax.models import mlp
+from tensor2tensor.trax.models import resnet
+
+# Ginify
+gin.external_configurable(mlp.MLP, module="models")
+gin.external_configurable(resnet.Resnet50, module="models")
diff --git a/tensor2tensor/jax/models/mlp.py b/tensor2tensor/trax/models/mlp.py
similarity index 95%
rename from tensor2tensor/jax/models/mlp.py
rename to tensor2tensor/trax/models/mlp.py
index 90aeda0d4..62b679199 100644
--- a/tensor2tensor/jax/models/mlp.py
+++ b/tensor2tensor/trax/models/mlp.py
@@ -13,18 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""J2J models."""
+"""MLP."""
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import gin
-
 from jax.experimental import stax
 
 
-@gin.configurable()
 def MLP(num_hidden_layers=2,
         hidden_size=512,
         activation_fn=stax.Relu,
diff --git a/tensor2tensor/jax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
similarity index 98%
rename from tensor2tensor/jax/models/resnet.py
rename to tensor2tensor/trax/models/resnet.py
index a054edcb1..9de5c4232 100644
--- a/tensor2tensor/jax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -13,14 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""J2J models."""
+"""ResNet."""
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import gin
-
 from jax.experimental import stax
 
 
@@ -62,7 +60,6 @@ def MakeMain(input_shape):
       stax.FanInSum, stax.Relu)
 
 
-@gin.configurable()
 def Resnet50(hidden_size=64, num_output_classes=1001):
   """ResNet.
 
diff --git a/tensor2tensor/trax/optimizers.py b/tensor2tensor/trax/optimizers.py
new file mode 100644
index 000000000..de20a0c61
--- /dev/null
+++ b/tensor2tensor/trax/optimizers.py
@@ -0,0 +1,41 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""gin-configurable optimizers and learning rate functions."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gin
+
+from jax.experimental import optimizers as opt
+
+
+def opt_configure(*args, **kwargs):
+  kwargs["module"] = "optimizers"
+  return gin.external_configurable(*args, **kwargs)
+
+# Optimizers
+sgd = opt_configure(opt.sgd)
+adam = opt_configure(opt.adam)
+momentum = opt_configure(opt.momentum)
+rmsprop = opt_configure(opt.rmsprop)
+
+# Learning rates
+constant = opt_configure(opt.constant)
+exponential_decay = opt_configure(opt.exponential_decay)
+inverse_time_decay = opt_configure(opt.inverse_time_decay)
+piecewise_constant = opt_configure(opt.piecewise_constant)
diff --git a/tensor2tensor/jax/j2j_trainer.py b/tensor2tensor/trax/trainer.py
similarity index 76%
rename from tensor2tensor/jax/j2j_trainer.py
rename to tensor2tensor/trax/trainer.py
index 0e09a8d7b..f29cf8883 100644
--- a/tensor2tensor/jax/j2j_trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -13,18 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-r"""J2J trainer.
+"""trax trainer."""
 
-Examples:
-
-
-MLP on mnist:
-  python -m tensor2tensor.jax.j2j_trainer \
-    --dataset=mnist \
-    --model=MLP \
-    --config="train_fn.train_steps=4000" \
-    --output_dir=~/j2j/test1
-"""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
@@ -37,7 +27,7 @@
 from absl import logging
 
 import gin
-from tensor2tensor.jax import j2j
+from tensor2tensor.trax import trax
 
 FLAGS = flags.FLAGS
 
@@ -56,22 +46,22 @@ def _setup_gin():
   configs = FLAGS.config or []
   # Override with --dataset and --model
   if FLAGS.dataset:
-    configs.append("train_fn.dataset='%s'" % FLAGS.dataset)
+    configs.append("train.dataset='%s'" % FLAGS.dataset)
   if FLAGS.model:
-    configs.append("train_fn.model=@" + FLAGS.model)
+    configs.append("train.model=@" + FLAGS.model)
   gin.parse_config_files_and_bindings(FLAGS.config_file, configs)
 
 
 def _default_output_dir():
   """Default output directory."""
   dir_name = "{model_name}_{dataset_name}_{timestamp}".format(
-      model_name=gin.query_parameter("train_fn.model").configurable.name,
-      dataset_name=gin.query_parameter("train_fn.dataset"),
+      model_name=gin.query_parameter("train.model").configurable.name,
+      dataset_name=gin.query_parameter("train.dataset"),
       timestamp=datetime.datetime.now().strftime("%Y%m%d_%H%M"),
   )
-  dir_path = os.path.join("~", "j2j", dir_name)
+  dir_path = os.path.join("~", "trax", dir_name)
   print()
-  j2j.log("No output_dir specified")
+  trax.log("No --output_dir specified")
   return dir_path
 
 
@@ -83,12 +73,12 @@ def main(_):
   output_dir = FLAGS.output_dir or _default_output_dir()
   assert data_dir, "Must specify a data directory"
   assert output_dir, "Must specify an output directory"
-  j2j.log("Using output_dir %s" % output_dir)
+  trax.log("Using --output_dir %s" % output_dir)
 
   data_dir = os.path.expanduser(data_dir)
   output_dir = os.path.expanduser(output_dir)
 
-  j2j.train_fn(data_dir=data_dir, output_dir=output_dir)
+  trax.train(data_dir=data_dir, output_dir=output_dir)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/jax/j2j.py b/tensor2tensor/trax/trax.py
similarity index 85%
rename from tensor2tensor/jax/j2j.py
rename to tensor2tensor/trax/trax.py
index d67da5fdc..fe9a7294d 100644
--- a/tensor2tensor/jax/j2j.py
+++ b/tensor2tensor/trax/trax.py
@@ -26,18 +26,21 @@
 import time
 
 from absl import logging
+
 import gin
 
 import jax
-from jax.experimental import optimizers
+from jax.experimental import optimizers as jax_opt
 import jax.numpy as np
 
 import six
 
-from tensor2tensor.jax import inputs as inputs_lib
-from tensor2tensor.jax import jaxboard
-# Import for gin configurable models
-from tensor2tensor.jax import models  # pylint: disable=unused-import
+from tensor2tensor.trax import inputs as inputs_lib
+from tensor2tensor.trax import jaxboard
+
+# Imports for gin configurables
+from tensor2tensor.trax import models as _trax_models  # pylint: disable=unused-import
+from tensor2tensor.trax import optimizers as trax_opt
 
 from tensorflow.io import gfile
 
@@ -61,23 +64,6 @@ def learning_rate(step,
   return ret
 
 
-@gin.configurable()
-def optimizer(name="adam",
-              momentum_mass=0.9, rmsprop_gamma=0.9, rmsprop_eps=1e-8,
-              adam_b1=0.9, adam_b2=0.997, adam_eps=1e-8):
-  """Return the optimizer, by name."""
-  if name == "sgd":
-    return optimizers.sgd(learning_rate)
-  if name == "momentum":
-    return optimizers.momentum(learning_rate, mass=momentum_mass)
-  if name == "rmsprop":
-    return optimizers.rmsprop(
-        learning_rate, gamma=rmsprop_gamma, eps=rmsprop_eps)
-  if name == "adam":
-    return optimizers.adam(learning_rate, b1=adam_b1, b2=adam_b2, eps=adam_eps)
-  raise ValueError("Unknown optimizer %s" % str(name))
-
-
 def one_hot(x, k, dtype=np.float32):
   """Create a one-hot encoding of x of size k."""
   return np.array(x[:, None] == np.arange(k), dtype)
@@ -150,18 +136,26 @@ def save_state(state, output_dir, save_gin=True):
     "loss": lambda x, y: - neg_log_perplexity(x, y),
 }
 
+# TODO(trax):
+# * Make Inputs an argument to train
+# * If eval_steps=None/0 or eval_frequency=None/0, disable evaluation
+# * Make learning rate configurable; possibly combine with optimizer
+# * Make loss configurable
+# * Make eval metrics configurable
+
 
 # We include in gin config everything that could be useful to share between
 # users, so when it gets saved in a .gin file it can be re-run with minimal
 # flags.
 @gin.configurable(blacklist=["data_dir", "output_dir"])
-def train_fn(output_dir,
-             data_dir,
-             model=gin.REQUIRED,
-             dataset=gin.REQUIRED,
-             train_steps=1000,
-             eval_steps=10,
-             eval_frequency=100):
+def train(output_dir,
+          data_dir,
+          model=gin.REQUIRED,
+          dataset=gin.REQUIRED,
+          optimizer=trax_opt.adam,
+          train_steps=1000,
+          eval_steps=10,
+          eval_frequency=100):
   """Train the given model on the given dataset.
 
   Args:
@@ -171,6 +165,8 @@ def train_fn(output_dir,
       and apply_fun.
     dataset: The name of the TFDS dataset to train on. To train on a T2T
       dataset, prefix the name with "t2t_".
+    optimizer: The optimizer as a callable taking a learning_rate callable and
+      returning 2 callables, opt_init and opt_update.
     train_steps: int, total number of training steps.
     eval_steps: int, num of steps per evaluation.
     eval_frequency: int, how often to run evaluation (every eval_frequency
@@ -182,7 +178,7 @@ def train_fn(output_dir,
   inputs = inputs_lib.make_inputs(dataset, data_dir)
 
   # Setup optimizer and model
-  opt_init, opt_update = optimizer()
+  opt_init, opt_update = optimizer(learning_rate)
   model_init, model_predict = model()
 
   # Setup state
@@ -200,13 +196,13 @@ def train_fn(output_dir,
 
   @jax.jit
   def update(i, opt_state, batch):
-    params = optimizers.get_params(opt_state)
+    params = jax_opt.get_params(opt_state)
     return opt_update(i, jax.grad(loss)(
         params, batch, model_predict), opt_state)
 
   print()
   step_log(step, "starting training")
-  train_gen = inputs.train_fn()
+  inputs_stream = inputs.train_fn()
   is_first_step = True
   epoch_steps = 1  # First evaluation after the first training step.
   while step < train_steps:
@@ -215,7 +211,7 @@ def update(i, opt_state, batch):
     # Train
     start_time = time.time()
     for _ in range(epoch_steps):
-      opt_state = update(step, opt_state, next(train_gen))
+      opt_state = update(step, opt_state, next(inputs_stream))
       if step % 10 == 0:  # Log learning rate curve each 10 steps.
         train_sw.scalar("training/learning rate",
                         learning_rate(step), step=step)
@@ -225,7 +221,7 @@ def update(i, opt_state, batch):
              (epoch_steps, epoch_time))
 
     # Save state
-    params = optimizers.get_params(opt_state)
+    params = jax_opt.get_params(opt_state)
     save_state(State(params=params, step=step), output_dir,
                save_gin=is_first_step)
 

From 77673a0cd2cee3a1568a69d7ff0108e6501d9ffb Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Wed, 27 Feb 2019 18:39:22 -0800
Subject: [PATCH 1727/2720] Make trax inputs configurable

PiperOrigin-RevId: 236035398
---
 tensor2tensor/bin/t2t_trainer.py              |  48 +++---
 tensor2tensor/data_generators/all_problems.py |   2 +-
 .../trax/configs/resnet50_imagenet_8gb.gin    |  11 +-
 tensor2tensor/trax/inputs.py                  |  19 ++-
 tensor2tensor/trax/models/__init__.py         |  10 +-
 tensor2tensor/trax/optimizers.py              |   2 +-
 tensor2tensor/trax/trainer.py                 |  34 ++--
 tensor2tensor/trax/trax.py                    | 149 +++++++++---------
 8 files changed, 154 insertions(+), 121 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index db223d899..122c87680 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -367,27 +367,33 @@ def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
 
   if FLAGS.jax:
-    # Hacking main v1 flags to work with jax.
-    config_strs = []
-    config_strs.append(
-        "train.train_steps=" + str(FLAGS.train_steps))
-    config_strs.append(
-        "train.eval_steps=" + str(FLAGS.eval_steps))
-    config_strs.append(
-        "train.eval_frequency=" + str(FLAGS.local_eval_frequency))
-    if FLAGS.hparams:
-      config_strs.extend(str(FLAGS.hparams).split(","))
-    data_dir = os.path.expanduser(FLAGS.data_dir)
-    output_dir = os.path.expanduser(FLAGS.output_dir)
-
-    gin.bind_parameter("train.dataset", FLAGS.problem)
-    config_strs += ["train.model=@" + FLAGS.model]
-    config_files = []
-    if FLAGS.hparams_set:
-      config_files = [os.path.expanduser(FLAGS.hparams_set)]
-    gin.parse_config_files_and_bindings(config_files, config_strs)
-    trax.train(data_dir=data_dir, output_dir=output_dir)
-    return
+    # Setup trax FLAGS
+    dataset = FLAGS.problem
+    model = FLAGS.model
+    data_dir = FLAGS.data_dir
+    output_dir = FLAGS.output_dir
+    config_file = [FLAGS.hparams_set]
+    config = [
+        "train.train_steps=%d" % FLAGS.train_steps,
+        "train.eval_steps=%d" % FLAGS.eval_steps,
+        "train.eval_frequency=%d" % FLAGS.local_eval_frequency,
+    ] + str(FLAGS.hparams).split(",")
+
+    # Copied _setup_gin exactly from trax/trainer.py and removed "FLAGS."
+
+    def _setup_gin():
+      configs = config or []
+      # Override with --dataset and --model
+      if dataset:
+        configs.append("inputs.dataset_name='%s'" % dataset)
+        configs.append("inputs.data_dir='%s'" % data_dir)
+        configs.append("train.inputs=@trax.inputs.inputs")
+      if model:
+        configs.append("train.model=@trax.models.%s" % model)
+      gin.parse_config_files_and_bindings(config_file, configs)
+
+    _setup_gin()
+    trax.train(output_dir=output_dir)
 
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
 
diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index c1f54f18f..5e0a12ca2 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -108,7 +108,7 @@ def _handle_errors(errors):
   if not errors:
     return
   log_all = True  # pylint: disable=unused-variable
-  err_msg = "Skipped importing {num_missing} data_generators modules."
+  err_msg = "T2T: skipped importing {num_missing} data_generators modules."
   print(err_msg.format(num_missing=len(errors)))
   for module, err in errors:
     err_str = str(err)
diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
index 049a5e169..9f1a6a9ba 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
@@ -5,6 +5,11 @@ batch_fn.bucket_length = 32
 batch_fn.buckets = None
 batch_fn.eval_batch_size = 32
 
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_image_imagenet224'
+
 # Parameters for learning_rate:
 # ==============================================================================
 learning_rate.constant = 5.0
@@ -26,9 +31,9 @@ Resnet50.num_output_classes = 1001
 
 # Parameters for train:
 # ==============================================================================
-train.dataset = 't2t_image_imagenet224'
 train.eval_frequency = 100
 train.eval_steps = 3
-train.model = @models.Resnet50
-train.optimizer = @optimizers.momentum
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.Resnet50
+train.optimizer = @trax.optimizers.momentum
 train.train_steps = 200000
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index c82007f95..b8765ecca 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -20,6 +20,8 @@
 from __future__ import print_function
 
 import collections
+import os
+
 import gin
 
 import jax.numpy as np
@@ -34,8 +36,21 @@
     "_Inputs", ["train_fn", "eval_fn", "input_shape"])
 
 
-def make_inputs(dataset_name, data_dir):
-  """Make Inputs."""
+@gin.configurable()
+def inputs(dataset_name, data_dir):
+  """Make Inputs for built-in datasets.
+
+  Args:
+    dataset_name: a TFDS or T2T dataset name. If it's a T2T dataset name, prefix
+      with "t2t_".
+    data_dir: data directory.
+
+  Returns:
+    trax.inputs.Inputs
+  """
+  assert data_dir, "Must provide a data directory"
+  data_dir = os.path.expanduser(data_dir)
+
   (train_batches, eval_batches,
    input_name, input_shape) = train_and_eval_batches(
        dataset_name, data_dir)
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index 21c555f5e..c885faed2 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -23,6 +23,12 @@
 from tensor2tensor.trax.models import mlp
 from tensor2tensor.trax.models import resnet
 
+
 # Ginify
-gin.external_configurable(mlp.MLP, module="models")
-gin.external_configurable(resnet.Resnet50, module="models")
+def model_configure(*args, **kwargs):
+  kwargs["module"] = "trax.models"
+  return gin.external_configurable(*args, **kwargs)
+
+
+model_configure(mlp.MLP)
+model_configure(resnet.Resnet50)
diff --git a/tensor2tensor/trax/optimizers.py b/tensor2tensor/trax/optimizers.py
index de20a0c61..f835d1d49 100644
--- a/tensor2tensor/trax/optimizers.py
+++ b/tensor2tensor/trax/optimizers.py
@@ -25,7 +25,7 @@
 
 
 def opt_configure(*args, **kwargs):
-  kwargs["module"] = "optimizers"
+  kwargs["module"] = "trax.optimizers"
   return gin.external_configurable(*args, **kwargs)
 
 # Optimizers
diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index f29cf8883..9c0d2c64c 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -27,6 +27,7 @@
 from absl import logging
 
 import gin
+
 from tensor2tensor.trax import trax
 
 FLAGS = flags.FLAGS
@@ -42,21 +43,11 @@
                           "Configuration parameters (gin string).")
 
 
-def _setup_gin():
-  configs = FLAGS.config or []
-  # Override with --dataset and --model
-  if FLAGS.dataset:
-    configs.append("train.dataset='%s'" % FLAGS.dataset)
-  if FLAGS.model:
-    configs.append("train.model=@" + FLAGS.model)
-  gin.parse_config_files_and_bindings(FLAGS.config_file, configs)
-
-
 def _default_output_dir():
   """Default output directory."""
   dir_name = "{model_name}_{dataset_name}_{timestamp}".format(
       model_name=gin.query_parameter("train.model").configurable.name,
-      dataset_name=gin.query_parameter("train.dataset"),
+      dataset_name=gin.query_parameter("inputs.dataset_name"),
       timestamp=datetime.datetime.now().strftime("%Y%m%d_%H%M"),
   )
   dir_path = os.path.join("~", "trax", dir_name)
@@ -65,20 +56,27 @@ def _default_output_dir():
   return dir_path
 
 
+def _setup_gin():
+  configs = FLAGS.config or []
+  # Override with --dataset and --model
+  if FLAGS.dataset:
+    configs.append("inputs.dataset_name='%s'" % FLAGS.dataset)
+    configs.append("inputs.data_dir='%s'" % FLAGS.data_dir)
+    configs.append("train.inputs=@trax.inputs.inputs")
+  if FLAGS.model:
+    configs.append("train.model=@trax.models.%s" % FLAGS.model)
+  gin.parse_config_files_and_bindings(FLAGS.config_file, configs)
+
+
 def main(_):
   _setup_gin()
 
-  # Setup directories
-  data_dir = FLAGS.data_dir
+  # Setup output directory
   output_dir = FLAGS.output_dir or _default_output_dir()
-  assert data_dir, "Must specify a data directory"
-  assert output_dir, "Must specify an output directory"
   trax.log("Using --output_dir %s" % output_dir)
-
-  data_dir = os.path.expanduser(data_dir)
   output_dir = os.path.expanduser(output_dir)
 
-  trax.train(data_dir=data_dir, output_dir=output_dir)
+  trax.train(output_dir=output_dir)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index fe9a7294d..67cca22cc 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""J2J main training functions."""
+"""trax main training functions."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -35,15 +35,19 @@
 
 import six
 
-from tensor2tensor.trax import inputs as inputs_lib
 from tensor2tensor.trax import jaxboard
-
-# Imports for gin configurables
-from tensor2tensor.trax import models as _trax_models  # pylint: disable=unused-import
 from tensor2tensor.trax import optimizers as trax_opt
 
 from tensorflow.io import gfile
 
+# Imports for gin configurables
+# TODO(trax): Move to trainer.py. Only here because of t2t_trainer usage.
+# pylint: disable=unused-import,g-bad-import-order,reimported
+from tensor2tensor.trax import inputs as _trax_inputs
+from tensor2tensor.trax import models as _trax_models
+from tensor2tensor.trax import optimizers as _trax_opt
+# pylint: disable=unused-import,g-bad-import-order,reimported
+
 
 @gin.configurable(blacklist=["step"])
 def learning_rate(step,
@@ -136,46 +140,79 @@ def save_state(state, output_dir, save_gin=True):
     "loss": lambda x, y: - neg_log_perplexity(x, y),
 }
 
+
+def evaluate(inputs, predict_fn, eval_steps):
+  """Evaluate.
+
+  Args:
+    inputs: Inputs namedtuple.
+    predict_fn: function from inputs to predictions. params should already be
+      partially applied.
+    eval_steps: int, number of evaluation steps.
+
+  Returns:
+    train_metrics: dict
+    eval_metrics: dict
+  """
+  eval_stream = inputs.eval_fn()
+  eval_train_stream = inputs.train_fn()
+  train_metrics = {key: 0.0 for key in _METRICS}
+  eval_metrics = {key: 0.0 for key in _METRICS}
+  for _ in range(eval_steps):
+    train_batch = next(eval_train_stream)
+    train_predictions = predict_fn(train_batch[0])
+    eval_batch = next(eval_stream)
+    eval_predictions = predict_fn(eval_batch[0])
+    for m in _METRICS:
+      train_metrics[m] += (_METRICS[m](train_batch, train_predictions)
+                           / float(eval_steps))
+      eval_metrics[m] += (_METRICS[m](eval_batch, eval_predictions)
+                          / float(eval_steps))
+
+  return train_metrics, eval_metrics
+
+
+def log_metrics(metrics, summ_writer, log_prefix, step):
+  rjust_len = max([len(name) for name in metrics])
+  for name, value in six.iteritems(metrics):
+    step_log(step, "%s %s | % .8f" % (log_prefix, name.rjust(rjust_len), value))
+    if summ_writer:
+      summ_writer.scalar("metrics/" + name, value, step)
+
+
 # TODO(trax):
-# * Make Inputs an argument to train
-# * If eval_steps=None/0 or eval_frequency=None/0, disable evaluation
-# * Make learning rate configurable; possibly combine with optimizer
-# * Make loss configurable
-# * Make eval metrics configurable
+# * Make configurable:
+#   * loss
+#   * metrics
+#   * learning rate
+# * Save/restore: pickle unsafe. Use np.array.savez + MessagePack?
 
 
-# We include in gin config everything that could be useful to share between
-# users, so when it gets saved in a .gin file it can be re-run with minimal
-# flags.
-@gin.configurable(blacklist=["data_dir", "output_dir"])
+@gin.configurable(blacklist=["output_dir"])
 def train(output_dir,
-          data_dir,
           model=gin.REQUIRED,
-          dataset=gin.REQUIRED,
+          inputs=gin.REQUIRED,
           optimizer=trax_opt.adam,
           train_steps=1000,
           eval_steps=10,
           eval_frequency=100):
-  """Train the given model on the given dataset.
+  """Train the model on the inputs.
 
   Args:
     output_dir: Directory where to put the logs and checkpoints.
-    data_dir: Directory where the data is located.
     model: The model to train as a callable returning 2 callables, an init_fun
       and apply_fun.
-    dataset: The name of the TFDS dataset to train on. To train on a T2T
-      dataset, prefix the name with "t2t_".
+    inputs: callable returning trax.inputs.Inputs.
     optimizer: The optimizer as a callable taking a learning_rate callable and
       returning 2 callables, opt_init and opt_update.
     train_steps: int, total number of training steps.
-    eval_steps: int, num of steps per evaluation.
+    eval_steps: int, num of steps per evaluation. If None or 0, eval disabled.
     eval_frequency: int, how often to run evaluation (every eval_frequency
-      steps).
+      steps). If None or 0, eval disabled.
   """
   gfile.makedirs(output_dir)
 
-  # Make Inputs
-  inputs = inputs_lib.make_inputs(dataset, data_dir)
+  inputs = inputs()
 
   # Setup optimizer and model
   opt_init, opt_update = optimizer(learning_rate)
@@ -203,10 +240,13 @@ def update(i, opt_state, batch):
   print()
   step_log(step, "starting training")
   inputs_stream = inputs.train_fn()
+  eval_enabled = eval_steps and eval_frequency
   is_first_step = True
-  epoch_steps = 1  # First evaluation after the first training step.
+  # Evaluate after the first training step, then reset to normal_epoch_steps
+  normal_epoch_steps = (eval_enabled and eval_frequency) or train_steps
+  epoch_steps = 1
   while step < train_steps:
-    print()
+    print()  # separate logging for each loop iteration
 
     # Train
     start_time = time.time()
@@ -226,61 +266,24 @@ def update(i, opt_state, batch):
                save_gin=is_first_step)
 
     # Evaluate
-    step_log(step, "starting evaluation")
-    train_metrics, eval_metrics = evaluate(
-        inputs, functools.partial(jit_predict, params), eval_steps)
-    log_metrics(train_metrics, train_sw, "train", step)
-    log_metrics(eval_metrics, eval_sw, "eval ", step)
+    if eval_enabled:
+      step_log(step, "starting evaluation")
+      train_metrics, eval_metrics = evaluate(
+          inputs, functools.partial(jit_predict, params), eval_steps)
+      log_metrics(train_metrics, train_sw, "train", step)
+      log_metrics(eval_metrics, eval_sw, "eval ", step)
+      eval_sw.writer.flush()
 
     # Log non-metric reports and flush.
     if not is_first_step:
       train_sw.scalar("training/steps per second",
                       epoch_steps / epoch_time, step=step)
     train_sw.writer.flush()
-    eval_sw.writer.flush()
 
-    # After the first step, train for eval_frequency steps before evaluating
-    epoch_steps = (eval_frequency - 1) if is_first_step else eval_frequency
+    # After the first step, train for normal_epoch_steps steps before evaluating
+    epoch_steps = (
+        (normal_epoch_steps - 1) if is_first_step else normal_epoch_steps)
     is_first_step = False
 
   print()
   step_log(step, "finished training")
-
-
-def evaluate(inputs, predict_fn, eval_steps):
-  """Evaluate.
-
-  Args:
-    inputs: Inputs namedtuple.
-    predict_fn: function from inputs to predictions. params should already be
-      partially applied.
-    eval_steps: int, number of evaluation steps.
-
-  Returns:
-    train_metrics: dict
-    eval_metrics: dict
-  """
-  eval_stream = inputs.eval_fn()
-  eval_train_stream = inputs.train_fn()
-  train_metrics = {key: 0.0 for key in _METRICS}
-  eval_metrics = {key: 0.0 for key in _METRICS}
-  for _ in range(eval_steps):
-    train_batch = next(eval_train_stream)
-    train_predictions = predict_fn(train_batch[0])
-    eval_batch = next(eval_stream)
-    eval_predictions = predict_fn(eval_batch[0])
-    for m in _METRICS:
-      train_metrics[m] += _METRICS[m](
-          train_batch, train_predictions) / float(eval_steps)
-      eval_metrics[m] += _METRICS[m](
-          eval_batch, eval_predictions) / float(eval_steps)
-
-  return train_metrics, eval_metrics
-
-
-def log_metrics(metrics, summ_writer, log_prefix, step):
-  rjust_len = max([len(name) for name in metrics])
-  for name, value in six.iteritems(metrics):
-    step_log(step, "%s %s | % .8f" % (log_prefix, name.rjust(rjust_len), value))
-    if summ_writer:
-      summ_writer.scalar("metrics/" + name, value, step)

From 428b8c09453ab1de047279184f283d02644ab581 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 28 Feb 2019 09:55:43 -0800
Subject: [PATCH 1728/2720] Add en-de data-sets with cleaning.

PiperOrigin-RevId: 236141252
---
 .../data_generators/cleaner_en_xx.py          |  6 +-
 tensor2tensor/data_generators/translate.py    | 60 ++++++++++++++-----
 .../data_generators/translate_ende.py         | 30 +++++++++-
 3 files changed, 77 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/data_generators/cleaner_en_xx.py b/tensor2tensor/data_generators/cleaner_en_xx.py
index 669666d1a..15243d17b 100644
--- a/tensor2tensor/data_generators/cleaner_en_xx.py
+++ b/tensor2tensor/data_generators/cleaner_en_xx.py
@@ -34,6 +34,8 @@
 import itertools
 import re
 
+from tensor2tensor.data_generators import text_encoder
+
 import tensorflow as tf
 
 
@@ -130,7 +132,7 @@ def clean_en_xx_pairs(en_xx_pairs):
       continue
     s1_list, s2_list = _split_sentences(s1, s2)
     if len(s1_list) != len(s2_list):
-      pass  # discard this pair
+      continue  # discard this pair
     elif len(s1_list) == 1:
       yield s1, s2
     else:
@@ -165,6 +167,8 @@ def _is_match(sentence, regex):
 
 
 def _split_sentences(s1, s2):
+  s1 = text_encoder.native_to_unicode(s1)
+  s2 = text_encoder.native_to_unicode(s2)
   s1 = re.sub(r'(\w[A-Z]|[0-9a-z])([.!?]) ([A-Z])', r'\1\2__|__\3', s1)
   s2 = re.sub(r'([^0-9][.!?]) ([A-Z])', r'\1__|__\2', s2)
   s1_subsentences = s1.split('__|__')
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index ce55cf2df..769911b92 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -45,6 +45,10 @@ def is_generate_per_split(self):
   def approx_vocab_size(self):
     return 2**15
 
+  @property
+  def datatypes_to_clean(self):
+    return None
+
   def source_data_files(self, dataset_split):
     """Files to be passed to compile_data."""
     raise NotImplementedError()
@@ -55,9 +59,14 @@ def vocab_data_files(self):
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
     datasets = self.source_data_files(dataset_split)
-    tag = "train" if dataset_split == problem.DatasetSplit.TRAIN else "dev"
-    data_path = compile_data(tmp_dir, datasets, "%s-compiled-%s" % (self.name,
-                                                                    tag))
+    tag = "dev"
+    datatypes_to_clean = None
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      tag = "train"
+      datatypes_to_clean = self.datatypes_to_clean
+    data_path = compile_data(
+        tmp_dir, datasets, "%s-compiled-%s" % (self.name, tag),
+        datatypes_to_clean=datatypes_to_clean)
     return text_problems.text2text_txt_iterator(data_path + ".lang1",
                                                 data_path + ".lang2")
 
@@ -127,8 +136,16 @@ def _preprocess_sgm(line, is_sgm):
     return line[i + 1:-6]  # Strip first <seg ...> and last </seg>.
 
 
-def compile_data(tmp_dir, datasets, filename):
-  """Concatenate all `datasets` and save to `filename`."""
+def _clean_sentences(sentence_pairs):
+  res_pairs = []
+  for cleaned in cleaner_en_xx.clean_en_xx_pairs(sentence_pairs):
+    res_pairs.append(cleaned)
+  return res_pairs
+
+
+def compile_data(tmp_dir, datasets, filename, datatypes_to_clean=None):
+  """Concatenates all `datasets` and saves to `filename`."""
+  datatypes_to_clean = datatypes_to_clean or []
   filename = os.path.join(tmp_dir, filename)
   lang1_fname = filename + ".lang1"
   lang2_fname = filename + ".lang2"
@@ -154,7 +171,10 @@ def compile_data(tmp_dir, datasets, filename):
             tmx_filename = new_filename
           source, target = None, None
           with tf.gfile.Open(tmx_filename) as tmx_file:
-            for source, target in cleaner_en_xx.paracrawl_v3_pairs(tmx_file):
+            stream = cleaner_en_xx.paracrawl_v3_pairs(tmx_file)
+            if "tmx" in datatypes_to_clean:
+              stream = cleaner_en_xx.clean_en_xx_pairs(stream)
+            for source, target in stream:
               lang1_resfile.write(source)
               lang1_resfile.write("\n")
               lang2_resfile.write(target)
@@ -180,11 +200,15 @@ def compile_data(tmp_dir, datasets, filename):
                   parts = line.split("\t")
                   source, target = parts[src_column], parts[trg_column]
                   source, target = source.strip(), target.strip()
-                  if source and target:
-                    lang1_resfile.write(source)
-                    lang1_resfile.write("\n")
-                    lang2_resfile.write(target)
-                    lang2_resfile.write("\n")
+                  clean_pairs = [(source, target)]
+                  if "tsv" in datatypes_to_clean:
+                    clean_pairs = cleaner_en_xx.clean_en_xx_pairs(clean_pairs)
+                  for source, target in clean_pairs:
+                    if source and target:
+                      lang1_resfile.write(source)
+                      lang1_resfile.write("\n")
+                      lang2_resfile.write(target)
+                      lang2_resfile.write("\n")
 
         else:
           lang1_filename, lang2_filename = dataset[1]
@@ -212,11 +236,15 @@ def compile_data(tmp_dir, datasets, filename):
               lang1_filepath, lang2_filepath):
             line1res = _preprocess_sgm(example["inputs"], is_sgm)
             line2res = _preprocess_sgm(example["targets"], is_sgm)
-            if line1res and line2res:
-              lang1_resfile.write(line1res)
-              lang1_resfile.write("\n")
-              lang2_resfile.write(line2res)
-              lang2_resfile.write("\n")
+            clean_pairs = [(line1res, line2res)]
+            if "txt" in datatypes_to_clean:
+              clean_pairs = cleaner_en_xx.clean_en_xx_pairs(clean_pairs)
+            for line1res, line2res in clean_pairs:
+              if line1res and line2res:
+                lang1_resfile.write(line1res)
+                lang1_resfile.write("\n")
+                lang2_resfile.write(line2res)
+                lang2_resfile.write("\n")
 
   return filename
 
diff --git a/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index ac7718726..b788553b8 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -76,6 +76,18 @@ def approx_vocab_size(self):
     return 2**15  # 32768
 
 
+@registry.register_problem
+class TranslateEndeWmtClean32k(TranslateEndeWmt32k):
+
+  @property
+  def use_vocab_from_other_problem(self):
+    return TranslateEndeWmt32k()
+
+  @property
+  def datatypes_to_clean(self):
+    return ["txt"]
+
+
 @registry.register_problem
 class TranslateEndeWmtParacrawlBicleaner32k(TranslateEndeWmt32k):
   """WMT en-de corpus with extra data from Paracrawl, cleaned with Bicleaner."""
@@ -91,6 +103,22 @@ def additional_training_datasets(self):
              ("tmx", "en-de.bicleaner07.tmx.gz"))]
 
 
+@registry.register_problem
+class TranslateEndeWmtParacrawlClean32k(TranslateEndeWmtParacrawlBicleaner32k):
+
+  @property
+  def datatypes_to_clean(self):
+    return ["tmx"]
+
+
+@registry.register_problem
+class TranslateEndeWmtParacrawlAllClean32k(TranslateEndeWmtParacrawlClean32k):
+
+  @property
+  def datatypes_to_clean(self):
+    return ["tmx", "txt"]
+
+
 @registry.register_problem
 class TranslateEndeWmt32kPacked(TranslateEndeWmt32k):
 
@@ -152,5 +180,3 @@ def inputs_prefix(self):
   @property
   def targets_prefix(self):
     return "translate German English "
-
-

From 85791b5218aa186ddd2fb52120a0227601dd21fc Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 28 Feb 2019 11:17:44 -0800
Subject: [PATCH 1729/2720] Gin config tensorboard summary; return State from
 train

PiperOrigin-RevId: 236159218
---
 tensor2tensor/trax/jaxboard.py | 28 ++++++++++++++++++++++++++++
 tensor2tensor/trax/trax.py     | 32 ++++++++++++++++++++++----------
 2 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/trax/jaxboard.py b/tensor2tensor/trax/jaxboard.py
index 278d3abba..92cbc0c25 100644
--- a/tensor2tensor/trax/jaxboard.py
+++ b/tensor2tensor/trax/jaxboard.py
@@ -307,3 +307,31 @@ def text(self, tag, textdata, step=None):
     summary = Summary(
         value=[Summary.Value(tag=tag, metadata=smd, tensor=tensor)])
     self.writer.add_summary(summary, step)
+
+
+# Copied from gin/tf/utils.py:GinConfigSaverHook
+def markdownify_operative_config_str(string):
+  """Convert an operative config string to markdown format."""
+
+  # TODO(b/37527917): Total hack below. Implement more principled formatting.
+  def process(line):
+    """Convert a single line to markdown format."""
+    if not line.startswith('#'):
+      return '    ' + line
+
+    line = line[2:]
+    if line.startswith('===='):
+      return ''
+    if line.startswith('None'):
+      return '    # None.'
+    if line.endswith(':'):
+      return '#### ' + line
+    return line
+
+  output_lines = []
+  for line in string.splitlines():
+    procd_line = process(line)
+    if procd_line is not None:
+      output_lines.append(procd_line)
+
+  return '\n'.join(output_lines)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 67cca22cc..908cbd6bf 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -119,19 +119,23 @@ def restore_state(output_dir):
   return State(step=step, params=params)
 
 
-def save_state(state, output_dir, save_gin=True):
+def save_gin(output_dir, sw=None):
+  config_path = os.path.join(output_dir, "config.gin")
+  config_str = gin.operative_config_str()
+  with gfile.GFile(config_path, "w") as f:
+    f.write(config_str)
+  if sw:
+    sw.text("gin_config",
+            jaxboard.markdownify_operative_config_str(config_str))
+
+
+def save_state(state, output_dir):
   """Save State and optionally gin config."""
   params_file = os.path.join(output_dir, "model.pkl")
   with gfile.GFile(params_file, "wb") as f:
     pickle.dump((state.params, state.step), f)
   log("Model saved to %s" % params_file, stdout=False)
 
-  # Gin file only includes used parameters, so we save it at this point.
-  if save_gin:
-    config_path = os.path.join(output_dir, "config.gin")
-    with gfile.GFile(config_path, "w") as f:
-      f.write(gin.operative_config_str())
-
 
 # Metrics to calculate and report.
 _METRICS = {
@@ -186,6 +190,7 @@ def log_metrics(metrics, summ_writer, log_prefix, step):
 #   * metrics
 #   * learning rate
 # * Save/restore: pickle unsafe. Use np.array.savez + MessagePack?
+# * Move metrics to metrics.py
 
 
 @gin.configurable(blacklist=["output_dir"])
@@ -209,6 +214,9 @@ def train(output_dir,
     eval_steps: int, num of steps per evaluation. If None or 0, eval disabled.
     eval_frequency: int, how often to run evaluation (every eval_frequency
       steps). If None or 0, eval disabled.
+
+  Returns:
+    trax.State
   """
   gfile.makedirs(output_dir)
 
@@ -262,8 +270,7 @@ def update(i, opt_state, batch):
 
     # Save state
     params = jax_opt.get_params(opt_state)
-    save_state(State(params=params, step=step), output_dir,
-               save_gin=is_first_step)
+    save_state(State(params=params, step=step), output_dir)
 
     # Evaluate
     if eval_enabled:
@@ -274,7 +281,11 @@ def update(i, opt_state, batch):
       log_metrics(eval_metrics, eval_sw, "eval ", step)
       eval_sw.writer.flush()
 
-    # Log non-metric reports and flush.
+    # Gin only tracks the used parameters, so we save it after the first step.
+    if is_first_step:
+      save_gin(output_dir, train_sw)
+
+    # Log non-metric reports.
     if not is_first_step:
       train_sw.scalar("training/steps per second",
                       epoch_steps / epoch_time, step=step)
@@ -287,3 +298,4 @@ def update(i, opt_state, batch):
 
   print()
   step_log(step, "finished training")
+  return State(params=params, step=step)

From 2930afbcca393b3cb0b4af5d00f8d80489d2a4d0 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 28 Feb 2019 14:03:52 -0800
Subject: [PATCH 1730/2720] Update TF_LATEST in .travis.yml to 1.13.*

PiperOrigin-RevId: 236193341
---
 .travis.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index fc22a2bd0..6484ecb17 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -14,13 +14,13 @@ env:
     - T2T_PROBLEM=algorithmic_reverse_binary40_test
     - T2T_DATA_DIR=/tmp/t2t-data
     - T2T_TRAIN_DIR=/tmp/t2t-train
-    - TF_LATEST="1.12.*"
+    - TF_LATEST="1.13.*"
     # This is necessary to have gsutil work with Python 2.7
     - BOTO_CONFIG=/dev/null
   matrix:
     # We test against the latest stable TensorFlow and tf-nightly.
     # If updating, also update TF_LATEST above
-    - TF_VERSION="1.12.*"
+    - TF_VERSION="1.13.*"
     - TF_VERSION="tf-nightly"
 install:
   - ./oss_scripts/oss_pip_install.sh

From c317d6214f2995f5d14c332095ff8324624b96a9 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 1 Mar 2019 01:09:24 -0800
Subject: [PATCH 1731/2720] Move trax learning rate schedule into a separate
 library and add a schedule adjustable by eval performance.

PiperOrigin-RevId: 236276518
---
 tensor2tensor/trax/history.py       |  59 +++++++++++++++
 tensor2tensor/trax/jaxboard.py      |   6 +-
 tensor2tensor/trax/learning_rate.py | 110 ++++++++++++++++++++++++++++
 tensor2tensor/trax/trax.py          |  71 ++++++++----------
 4 files changed, 204 insertions(+), 42 deletions(-)
 create mode 100644 tensor2tensor/trax/history.py
 create mode 100644 tensor2tensor/trax/learning_rate.py

diff --git a/tensor2tensor/trax/history.py b/tensor2tensor/trax/history.py
new file mode 100644
index 000000000..a7421e745
--- /dev/null
+++ b/tensor2tensor/trax/history.py
@@ -0,0 +1,59 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""trax history."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+
+
+class History(object):
+  """History of metrics.
+
+  History contains the metrics recorded during training and evaluation.
+  Save data with history.append(metric, value, step, mode) and get a sequence
+  of data by calling history.get(metric, mode). For example:
+
+  history.append("metrics/accuracy", 0.04, 1, "train")
+  history.append("metrics/accuracy", 0.31, 1000, "train")
+  history.get("metrics/accuracy", "train")
+  # returns [(1, 0.04), (1000, 0.31)]
+  """
+
+  def __init__(self):
+    # Structure is
+    # values = {
+    #   "mode1": {
+    #     "metric1": [val1, val2],
+    #     ...
+    #   },
+    #   "mode2": ...
+    # }
+    self._values = {}
+
+  def append(self, metric, value, step, mode):
+    """Append (step, value) pair to history for the given mode and metric."""
+    if mode not in self._values:
+      self._values[mode] = collections.defaultdict(list)
+    self._values[mode][metric].append((step, value))
+
+  def get(self, metric, mode):
+    """Get the history for the given metric and mode."""
+    if mode not in self._values:
+      return []
+    return list(self._values[mode][metric])
diff --git a/tensor2tensor/trax/jaxboard.py b/tensor2tensor/trax/jaxboard.py
index 92cbc0c25..69ac5afce 100644
--- a/tensor2tensor/trax/jaxboard.py
+++ b/tensor2tensor/trax/jaxboard.py
@@ -34,10 +34,10 @@
 import matplotlib.pyplot as plt
 import numpy as onp
 import tensorflow as tf
-from tensorflow import gfile
 from tensorflow import HistogramProto
 from tensorflow import Summary
 from tensorflow import SummaryMetadata
+from tensorflow.io import gfile
 
 
 def _pack_images(images, rows, cols):
@@ -75,8 +75,8 @@ def __init__(self, log_dir):
       log_dir: path to record tfevents files in.
     """
     # If needed, create log_dir directory as well as missing parent directories.
-    if not gfile.IsDirectory(log_dir):
-      gfile.MakeDirs(log_dir)
+    if not gfile.isdir(log_dir):
+      gfile.makedirs(log_dir)
 
     self.writer = tf.summary.FileWriter(log_dir, graph=None)
     self.end_summaries = []
diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
new file mode 100644
index 000000000..cb81f84af
--- /dev/null
+++ b/tensor2tensor/trax/learning_rate.py
@@ -0,0 +1,110 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""trax learning rate schedules."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gin
+import jax.numpy as np
+
+
+@gin.configurable(blacklist=["history"])
+def make_default_schedule(history=None,
+                          schedule="constant * linear_warmup * rsqrt_decay",
+                          constant=0.001,
+                          warmup_steps=100):
+  """Default learning rate  schedule.
+
+  Note: the learning rate schedule takes arguments and return a function,
+  learning_rate: step -> lr, that only takes a step and return the rate.
+  The reason is that learning_rate(step) is called at every training step,
+  so should be efficient, while the schedule is re-computed only when
+  evaluating the model, so usually only every 100 or 1000 steps.
+
+  Interprets factors in the schedule string which can consist of:
+  * constant: interpreted as the constant value,
+  * linear_warmup: interpreted as linear warmup until warmup_steps,
+  * rsqrt_decay: divide by square root of max(step, warmup_steps)
+
+  Args:
+    history: the history of training and evaluation (History object).
+    schedule: a string with factors separated by "*" that defines the schedule.
+    constant: float, the starting constant for the learning rate schedule.
+    warmup_steps: how many steps to warm up for in the warmup schedule.
+
+  Returns:
+    a function learning_rate(step): float -> float, the step-dependent lr.
+  """
+  del history
+  factors = [n.strip() for n in schedule.split("*")]
+
+  def learning_rate(step):
+    """Step to learning rate function."""
+    ret = 1.0
+    for name in factors:
+      if name == "constant":
+        ret *= constant
+      elif name == "linear_warmup":
+        ret *= np.minimum(1.0, step / warmup_steps)
+      elif name == "rsqrt_decay":
+        ret /= np.sqrt(np.maximum(step, warmup_steps))
+      else:
+        raise ValueError("Unknown factor %s." % name)
+    return ret
+
+  return learning_rate
+
+
+@gin.configurable(blacklist=["history"])
+def make_eval_adjusting_schedule(history,
+                                 constant=0.001,
+                                 steps_to_decrease=10,
+                                 improvement_margin=0.01,
+                                 decrease_rate=2.0,
+                                 metric="metrics/accuracy"):
+  """Learning rate that decreases when eval metric stalls.
+
+  If the chosen metric does not improve by improvement_margin for as many as
+  steps_to_decrease steps, then the constant gets decreased by decrease rate.
+  Finally, the default schedule gets called with the adjusted constant.
+
+  Args:
+    history: the history of training and evaluation (History object).
+    constant: float, the starting constant for the learning rate schedule.
+    steps_to_decrease: int, after how many steps without improvement
+      should we decrease the constant.
+    improvement_margin: how much we need to improve to count it.
+    decrease_rate: by how much to decrease.
+    metric: which evaluation metric to use for adjustments.
+
+  Returns:
+    a function learning_rate(step): float -> float, the step-dependent lr.
+  """
+  metric = history.get(metric, "eval")
+  adjusted = constant
+  steps_without_improvement = 0
+  while len(metric) > 1:
+    last = metric.pop()
+    if last[1] < metric[-1][1] * (1 + improvement_margin):
+      steps_without_improvement += 1
+    else:
+      steps_without_improvement = 0
+    if steps_without_improvement >= steps_to_decrease:
+      adjusted /= decrease_rate
+      steps_without_improvement = 0
+  return make_default_schedule(history, constant=adjusted)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 908cbd6bf..bd22888e9 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -35,7 +35,9 @@
 
 import six
 
+from tensor2tensor.trax import history as trax_history
 from tensor2tensor.trax import jaxboard
+from tensor2tensor.trax import learning_rate as lr
 from tensor2tensor.trax import optimizers as trax_opt
 
 from tensorflow.io import gfile
@@ -49,25 +51,6 @@
 # pylint: disable=unused-import,g-bad-import-order,reimported
 
 
-@gin.configurable(blacklist=["step"])
-def learning_rate(step,
-                  schedule="constant * linear_warmup * rsqrt_decay",
-                  constant=0.001,
-                  warmup_steps=100):
-  """Learning rate."""
-  ret = 1.0
-  for name in [n.strip() for n in schedule.split("*")]:
-    if name == "constant":
-      ret *= constant
-    elif name == "linear_warmup":
-      ret *= np.minimum(1.0, step / warmup_steps)
-    elif name == "rsqrt_decay":
-      ret /= np.sqrt(np.maximum(step, warmup_steps))
-    else:
-      raise ValueError("Unknown factor %s." % name)
-  return ret
-
-
 def one_hot(x, k, dtype=np.float32):
   """Create a one-hot encoding of x of size k."""
   return np.array(x[:, None] == np.arange(k), dtype)
@@ -104,19 +87,19 @@ def step_log(step, s):
   log("Step % 6d: %s" % (step, s))
 
 
-State = collections.namedtuple("_State", ["step", "params"])
+State = collections.namedtuple("_State", ["step", "params", "history"])
 
 
 def restore_state(output_dir):
   """Restore State."""
   params_file = os.path.join(output_dir, "model.pkl")
   if not gfile.exists(params_file):
-    return State(step=None, params=None)
+    return State(step=None, params=None, history=trax_history.History())
 
   with gfile.GFile(params_file, "rb") as f:
-    (params, step) = pickle.load(f)
+    (params, step, history) = pickle.load(f)
   log("Model loaded from %s" % params_file)
-  return State(step=step, params=params)
+  return State(step=step, params=params, history=history)
 
 
 def save_gin(output_dir, sw=None):
@@ -133,7 +116,7 @@ def save_state(state, output_dir):
   """Save State and optionally gin config."""
   params_file = os.path.join(output_dir, "model.pkl")
   with gfile.GFile(params_file, "wb") as f:
-    pickle.dump((state.params, state.step), f)
+    pickle.dump((state.params, state.step, state.history), f)
   log("Model saved to %s" % params_file, stdout=False)
 
 
@@ -176,19 +159,22 @@ def evaluate(inputs, predict_fn, eval_steps):
   return train_metrics, eval_metrics
 
 
-def log_metrics(metrics, summ_writer, log_prefix, step):
+def log_metrics(metrics, summ_writer, log_prefix, step, history=None):
+  """Log metrics to summary writer and history."""
   rjust_len = max([len(name) for name in metrics])
   for name, value in six.iteritems(metrics):
     step_log(step, "%s %s | % .8f" % (log_prefix, name.rjust(rjust_len), value))
+    full_name = "metrics/" + name
+    if history:
+      history.append(full_name, value, step, log_prefix)
     if summ_writer:
-      summ_writer.scalar("metrics/" + name, value, step)
+      summ_writer.scalar(full_name, value, step)
 
 
 # TODO(trax):
 # * Make configurable:
 #   * loss
 #   * metrics
-#   * learning rate
 # * Save/restore: pickle unsafe. Use np.array.savez + MessagePack?
 # * Move metrics to metrics.py
 
@@ -198,6 +184,7 @@ def train(output_dir,
           model=gin.REQUIRED,
           inputs=gin.REQUIRED,
           optimizer=trax_opt.adam,
+          learning_rate_fn=lr.make_default_schedule,
           train_steps=1000,
           eval_steps=10,
           eval_frequency=100):
@@ -210,6 +197,8 @@ def train(output_dir,
     inputs: callable returning trax.inputs.Inputs.
     optimizer: The optimizer as a callable taking a learning_rate callable and
       returning 2 callables, opt_init and opt_update.
+    learning_rate_fn: The learning rate callable that takes history and returns
+      a function from step to learning rate (a float).
     train_steps: int, total number of training steps.
     eval_steps: int, num of steps per evaluation. If None or 0, eval disabled.
     eval_frequency: int, how often to run evaluation (every eval_frequency
@@ -219,23 +208,24 @@ def train(output_dir,
     trax.State
   """
   gfile.makedirs(output_dir)
+  # Create summary writers and history.
+  train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
+  eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
 
   inputs = inputs()
 
   # Setup optimizer and model
+  state = restore_state(output_dir)
+  history = state.history
+  learning_rate = learning_rate_fn(history)
   opt_init, opt_update = optimizer(learning_rate)
   model_init, model_predict = model()
 
   # Setup state
-  state = restore_state(output_dir)
   step = state.step or 0
   params_initializer = lambda: model_init([-1] + inputs.input_shape)[1]
   opt_state = opt_init(state.params or params_initializer())
 
-  # Create summary writers.
-  train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
-  eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
-
   # jit model_predict and update so they're fast
   jit_predict = jax.jit(model_predict)  # for evaluation
 
@@ -268,19 +258,19 @@ def update(i, opt_state, batch):
     step_log(step, "ran %d train steps in %0.2f secs" %
              (epoch_steps, epoch_time))
 
-    # Save state
-    params = jax_opt.get_params(opt_state)
-    save_state(State(params=params, step=step), output_dir)
-
     # Evaluate
+    params = jax_opt.get_params(opt_state)
     if eval_enabled:
       step_log(step, "starting evaluation")
       train_metrics, eval_metrics = evaluate(
           inputs, functools.partial(jit_predict, params), eval_steps)
-      log_metrics(train_metrics, train_sw, "train", step)
-      log_metrics(eval_metrics, eval_sw, "eval ", step)
+      log_metrics(train_metrics, train_sw, "train", step, history=history)
+      log_metrics(eval_metrics, eval_sw, "eval ", step, history=history)
       eval_sw.writer.flush()
 
+    # Save state
+    save_state(State(params=params, step=step, history=history), output_dir)
+
     # Gin only tracks the used parameters, so we save it after the first step.
     if is_first_step:
       save_gin(output_dir, train_sw)
@@ -291,6 +281,9 @@ def update(i, opt_state, batch):
                       epoch_steps / epoch_time, step=step)
     train_sw.writer.flush()
 
+    # Update learning rate with new history.
+    learning_rate = learning_rate_fn(history)
+
     # After the first step, train for normal_epoch_steps steps before evaluating
     epoch_steps = (
         (normal_epoch_steps - 1) if is_first_step else normal_epoch_steps)
@@ -298,4 +291,4 @@ def update(i, opt_state, batch):
 
   print()
   step_log(step, "finished training")
-  return State(params=params, step=step)
+  return State(params=params, step=step, history=history)

From 50b2455293a5ec7306fd122a922d987e384d0d03 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Fri, 1 Mar 2019 12:19:03 -0800
Subject: [PATCH 1732/2720] Mv gin configurable imports from trax to trainer

PiperOrigin-RevId: 236355882
---
 tensor2tensor/bin/t2t_trainer.py                     | 8 ++++++++
 tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin | 5 +++++
 tensor2tensor/trax/trainer.py                        | 8 ++++++++
 tensor2tensor/trax/trax.py                           | 8 --------
 4 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 122c87680..304d49e0b 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -382,6 +382,14 @@ def main(argv):
     # Copied _setup_gin exactly from trax/trainer.py and removed "FLAGS."
 
     def _setup_gin():
+      """Setup gin configuration."""
+      # Imports for configurables
+      # pylint: disable=g-import-not-at-top,unused-import,g-bad-import-order,reimported,unused-variable
+      from tensor2tensor.trax import inputs as _trax_inputs
+      from tensor2tensor.trax import models as _trax_models
+      from tensor2tensor.trax import optimizers as _trax_opt
+      # pylint: disable=g-import-not-at-top,unused-import,g-bad-import-order,reimported,unused-variable
+
       configs = config or []
       # Override with --dataset and --model
       if dataset:
diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
index 9f1a6a9ba..cae941553 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
@@ -1,3 +1,8 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
 # Parameters for batch_fn:
 # ==============================================================================
 batch_fn.batch_size = 32
diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index 9c0d2c64c..fb56ec83a 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -57,6 +57,14 @@ def _default_output_dir():
 
 
 def _setup_gin():
+  """Setup gin configuration."""
+  # Imports for configurables
+  # pylint: disable=g-import-not-at-top,unused-import,g-bad-import-order,reimported,unused-variable
+  from tensor2tensor.trax import inputs as _trax_inputs
+  from tensor2tensor.trax import models as _trax_models
+  from tensor2tensor.trax import optimizers as _trax_opt
+  # pylint: disable=g-import-not-at-top,unused-import,g-bad-import-order,reimported,unused-variable
+
   configs = FLAGS.config or []
   # Override with --dataset and --model
   if FLAGS.dataset:
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index bd22888e9..4466057fe 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -42,14 +42,6 @@
 
 from tensorflow.io import gfile
 
-# Imports for gin configurables
-# TODO(trax): Move to trainer.py. Only here because of t2t_trainer usage.
-# pylint: disable=unused-import,g-bad-import-order,reimported
-from tensor2tensor.trax import inputs as _trax_inputs
-from tensor2tensor.trax import models as _trax_models
-from tensor2tensor.trax import optimizers as _trax_opt
-# pylint: disable=unused-import,g-bad-import-order,reimported
-
 
 def one_hot(x, k, dtype=np.float32):
   """Create a one-hot encoding of x of size k."""

From 663adec0880b65334921cabc15e19a87e29c33c8 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Fri, 1 Mar 2019 22:36:18 +0100
Subject: [PATCH 1733/2720] Correct model-free PPO test. (#1474)

---
 tensor2tensor/rl/trainer_model_free_test.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index 90c728ff0..98c28feca 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -35,7 +35,8 @@ def test_train_pong(self):
     hparams.add_hparam("ppo_epochs_num", 2)
     hparams.add_hparam("ppo_epoch_length", 3)
     FLAGS.output_dir = tf.test.get_temp_dir()
-    trainer_model_free.train(hparams, FLAGS.output_dir)
+    trainer_model_free.train(hparams, FLAGS.output_dir,
+                             env_problem_name=None)
 
 
 if __name__ == "__main__":

From 3f4a99700383b7e3e04a7dcec89f88144bd2be5b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 1 Mar 2019 13:42:27 -0800
Subject: [PATCH 1734/2720] Bugfix for evolved_transformer_decode and move
 evolved transformer related code into one file (evolved_transformer.py).

PiperOrigin-RevId: 236370260
---
 tensor2tensor/layers/transformer_layers.py  | 163 --------
 tensor2tensor/models/evolved_transformer.py | 425 +++++++++++++++++++-
 tensor2tensor/models/transformer.py         | 251 ------------
 3 files changed, 422 insertions(+), 417 deletions(-)

diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 85c76823b..2202fdefb 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -223,169 +223,6 @@ def transformer_encoder(encoder_input,
     return common_layers.layer_preprocess(x, hparams)
 
 
-def evolved_transformer_encoder(encoder_input,
-                                encoder_self_attention_bias,
-                                hparams,
-                                name="encoder",
-                                nonpadding=None,
-                                save_weights_to=None,
-                                make_image_summary=True,
-                                losses=None,
-                                attn_bias_for_padding=None):
-  """Evolved Transformer encoder. See arxiv.org/abs/1901.11117 for more details.
-
-  Note: Pad remover is not supported.
-
-  Args:
-    encoder_input: a Tensor.
-    encoder_self_attention_bias: bias Tensor for self-attention (see
-      common_attention.attention_bias()).
-    hparams: hyperparameters for model.
-    name: a string.
-    nonpadding: optional Tensor with shape [batch_size, encoder_length]
-      indicating what positions are not padding.  This must either be passed in,
-      which we do for "packed" datasets, or inferred from
-      encoder_self_attention_bias.  The knowledge about padding is used for
-      pad_remover(efficiency) and to mask out padding in convolutional layers.
-    save_weights_to: an optional dictionary to capture attention weights for
-      visualization; the weights tensor will be appended there under a string
-      key created from the variable scope (including name).
-    make_image_summary: Whether to make an attention image summary.
-    losses: Not used.
-    attn_bias_for_padding: Padded attention bias in case a unidirectional
-      encoder is being used where future attention is masked.
-
-  Returns:
-    Tensor encoder output.
-  """
-  del losses
-
-  hidden_state = encoder_input
-  attention_dropout_broadcast_dims = (
-      common_layers.comma_separated_string_to_integer_list(
-          getattr(hparams, "attention_dropout_broadcast_dims", "")))
-
-  with tf.variable_scope(name):
-    if nonpadding is not None:
-      padding = 1.0 - nonpadding
-    else:
-      attention_bias = encoder_self_attention_bias
-      if attn_bias_for_padding is not None:
-        attention_bias = attn_bias_for_padding
-      padding = common_attention.attention_bias_to_padding(attention_bias)
-      nonpadding = 1.0 - padding
-
-    for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):
-      with tf.variable_scope("layer_%d" % layer):
-
-        with tf.variable_scope("gated_linear_unit"):
-
-          residual_state = hidden_state
-          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
-
-          values = layers().Dense(hparams.hidden_size)(hidden_state)
-          gates = layers().Dense(
-              hparams.hidden_size, activation=tf.nn.sigmoid)(hidden_state)
-          hidden_state = values * gates
-
-          hidden_state = common_layers.layer_postprocess(
-              residual_state, hidden_state, hparams)
-
-        with tf.variable_scope("conv_branches"):
-
-          residual_state = hidden_state
-          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
-          # Mask padding from conv layers.
-          mask = tf.tile(
-              tf.expand_dims(nonpadding, 2), [1, 1, hparams.hidden_size])
-          hidden_state *= mask
-
-          left_output_dim = int(hparams.hidden_size * 4)
-          left_state = layers().Dense(
-              left_output_dim, activation=tf.nn.relu)(hidden_state)
-          left_state = tf.nn.dropout(left_state,
-                                     1 - hparams.layer_prepostprocess_dropout)
-
-          right_output_dim = int(hparams.hidden_size / 2)
-          right_state = layers().Conv1D(
-              right_output_dim,
-              3,
-              padding="SAME",
-              name="standard_conv_3x1",
-              activation=tf.nn.relu)(hidden_state)
-          right_state = tf.nn.dropout(right_state,
-                                      1 - hparams.layer_prepostprocess_dropout)
-
-          right_state = tf.pad(
-              right_state,
-              [[0, 0], [0, 0], [0, left_output_dim - right_output_dim]],
-              constant_values=0)
-          hidden_state = left_state + right_state
-
-          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
-          # Mask padding from conv layer.
-          mask = tf.tile(tf.expand_dims(nonpadding, 2), [1, 1, left_output_dim])
-          hidden_state *= mask
-
-          separable_conv_9x1 = layers().SeparableConv1D(
-              right_output_dim, 9, padding="SAME", name="separable_conv_9x1")
-          hidden_state = separable_conv_9x1(hidden_state)
-          hidden_state = tf.pad(
-              hidden_state,
-              [[0, 0], [0, 0], [0, hparams.hidden_size - right_output_dim]],
-              constant_values=0)
-
-          hidden_state = common_layers.layer_postprocess(
-              residual_state, hidden_state, hparams)
-
-        with tf.variable_scope("self_attention"):
-          residual_state = hidden_state
-          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
-
-          hidden_state = common_attention.multihead_attention(
-              hidden_state,
-              None,
-              encoder_self_attention_bias,
-              hparams.attention_key_channels or hparams.hidden_size,
-              hparams.attention_value_channels or hparams.hidden_size,
-              hparams.hidden_size,
-              hparams.num_heads,
-              hparams.attention_dropout,
-              attention_type=hparams.self_attention_type,
-              max_relative_position=hparams.max_relative_position,
-              heads_share_relative_embedding=(
-                  hparams.heads_share_relative_embedding),
-              add_relative_to_values=hparams.add_relative_to_values,
-              save_weights_to=save_weights_to,
-              make_image_summary=make_image_summary,
-              dropout_broadcast_dims=attention_dropout_broadcast_dims,
-              max_length=hparams.get("max_length"),
-              vars_3d=hparams.get("attention_variables_3d"),
-              activation_dtype=hparams.get("activation_dtype", "float32"),
-              weight_dtype=hparams.get("weight_dtype", "float32"))
-
-          hidden_state = common_layers.layer_postprocess(
-              residual_state, hidden_state, hparams)
-
-        with tf.variable_scope("dense_layers"):
-          residual_state = hidden_state
-          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
-
-          hidden_state = layers().Dense(
-              int(hparams.hidden_size * 4), activation=tf.nn.relu)(hidden_state)
-          hidden_state = tf.nn.dropout(hidden_state,
-                                       1 - hparams.layer_prepostprocess_dropout)
-
-          hidden_state = layers().Dense(hparams.hidden_size)(hidden_state)
-          hidden_state = common_layers.layer_postprocess(
-              residual_state, hidden_state, hparams)
-
-    # If normalization is done in layer_preprocess, then it should also be done
-    # on the output, since the output can grow very large, being the sum of
-    # a whole stack of unnormalized layer outputs.
-    return common_layers.layer_preprocess(hidden_state, hparams)
-
-
 def transformer_ffn_layer(x,
                           hparams,
                           pad_remover=None,
diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 0ac840b16..a99a01e81 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -22,10 +22,13 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.layers import transformer_layers
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
 from tensor2tensor.utils import registry
 
+import tensorflow as tf
+
 
 @registry.register_model
 class EvolvedTransformer(transformer.Transformer):
@@ -33,8 +36,8 @@ class EvolvedTransformer(transformer.Transformer):
 
   def __init__(self, *args, **kwargs):
     super(EvolvedTransformer, self).__init__(*args, **kwargs)
-    self._encoder_function = transformer_layers.evolved_transformer_encoder
-    self._decoder_function = transformer.evolved_transformer_decoder
+    self._encoder_function = evolved_transformer_encoder
+    self._decoder_function = evolved_transformer_decoder
 
   def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha,
                    use_tpu):
@@ -62,6 +65,422 @@ def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha,
                                   alpha, use_tpu)
 
 
+def evolved_transformer_encoder(encoder_input,
+                                encoder_self_attention_bias,
+                                hparams,
+                                name="encoder",
+                                nonpadding=None,
+                                save_weights_to=None,
+                                make_image_summary=True,
+                                losses=None,
+                                attn_bias_for_padding=None):
+  """Evolved Transformer encoder. See arxiv.org/abs/1901.11117 for more details.
+
+  Note: Pad remover is not supported.
+
+  Args:
+    encoder_input: a Tensor.
+    encoder_self_attention_bias: bias Tensor for self-attention (see
+      common_attention.attention_bias()).
+    hparams: hyperparameters for model.
+    name: a string.
+    nonpadding: optional Tensor with shape [batch_size, encoder_length]
+      indicating what positions are not padding.  This must either be passed in,
+      which we do for "packed" datasets, or inferred from
+      encoder_self_attention_bias.  The knowledge about padding is used for
+      pad_remover(efficiency) and to mask out padding in convolutional layers.
+    save_weights_to: an optional dictionary to capture attention weights for
+      visualization; the weights tensor will be appended there under a string
+      key created from the variable scope (including name).
+    make_image_summary: Whether to make an attention image summary.
+    losses: Not used.
+    attn_bias_for_padding: Padded attention bias in case a unidirectional
+      encoder is being used where future attention is masked.
+
+  Returns:
+    Tensor encoder output.
+  """
+  del losses
+
+  hidden_state = encoder_input
+  attention_dropout_broadcast_dims = (
+      common_layers.comma_separated_string_to_integer_list(
+          getattr(hparams, "attention_dropout_broadcast_dims", "")))
+
+  with tf.variable_scope(name):
+    if nonpadding is not None:
+      padding = 1.0 - nonpadding
+    else:
+      attention_bias = encoder_self_attention_bias
+      if attn_bias_for_padding is not None:
+        attention_bias = attn_bias_for_padding
+      padding = common_attention.attention_bias_to_padding(attention_bias)
+      nonpadding = 1.0 - padding
+
+    for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):
+      with tf.variable_scope("layer_%d" % layer):
+
+        with tf.variable_scope("gated_linear_unit"):
+
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          values = common_layers.layers().Dense(
+              hparams.hidden_size)(hidden_state)
+          gates = common_layers.layers().Dense(
+              hparams.hidden_size, activation=tf.nn.sigmoid)(hidden_state)
+          hidden_state = values * gates
+
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, hidden_state, hparams)
+
+        with tf.variable_scope("conv_branches"):
+
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+          # Mask padding from conv layers.
+          mask = tf.tile(
+              tf.expand_dims(nonpadding, 2), [1, 1, hparams.hidden_size])
+          hidden_state *= mask
+
+          left_output_dim = int(hparams.hidden_size * 4)
+          left_state = common_layers.layers().Dense(
+              left_output_dim, activation=tf.nn.relu)(hidden_state)
+          left_state = tf.nn.dropout(left_state,
+                                     1 - hparams.layer_prepostprocess_dropout)
+
+          right_output_dim = int(hparams.hidden_size / 2)
+          right_state = common_layers.layers().Conv1D(
+              right_output_dim,
+              3,
+              padding="SAME",
+              name="standard_conv_3x1",
+              activation=tf.nn.relu)(hidden_state)
+          right_state = tf.nn.dropout(right_state,
+                                      1 - hparams.layer_prepostprocess_dropout)
+
+          right_state = tf.pad(
+              right_state,
+              [[0, 0], [0, 0], [0, left_output_dim - right_output_dim]],
+              constant_values=0)
+          hidden_state = left_state + right_state
+
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+          # Mask padding from conv layer.
+          mask = tf.tile(tf.expand_dims(nonpadding, 2), [1, 1, left_output_dim])
+          hidden_state *= mask
+
+          separable_conv_9x1 = common_layers.layers().SeparableConv1D(
+              right_output_dim, 9, padding="SAME", name="separable_conv_9x1")
+          hidden_state = separable_conv_9x1(hidden_state)
+          hidden_state = tf.pad(
+              hidden_state,
+              [[0, 0], [0, 0], [0, hparams.hidden_size - right_output_dim]],
+              constant_values=0)
+
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, hidden_state, hparams)
+
+        with tf.variable_scope("self_attention"):
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          hidden_state = common_attention.multihead_attention(
+              hidden_state,
+              None,
+              encoder_self_attention_bias,
+              hparams.attention_key_channels or hparams.hidden_size,
+              hparams.attention_value_channels or hparams.hidden_size,
+              hparams.hidden_size,
+              hparams.num_heads,
+              hparams.attention_dropout,
+              attention_type=hparams.self_attention_type,
+              max_relative_position=hparams.max_relative_position,
+              heads_share_relative_embedding=(
+                  hparams.heads_share_relative_embedding),
+              add_relative_to_values=hparams.add_relative_to_values,
+              save_weights_to=save_weights_to,
+              make_image_summary=make_image_summary,
+              dropout_broadcast_dims=attention_dropout_broadcast_dims,
+              max_length=hparams.get("max_length"),
+              vars_3d=hparams.get("attention_variables_3d"),
+              activation_dtype=hparams.get("activation_dtype", "float32"),
+              weight_dtype=hparams.get("weight_dtype", "float32"))
+
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, hidden_state, hparams)
+
+        with tf.variable_scope("dense_layers"):
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          hidden_state = common_layers.layers().Dense(
+              int(hparams.hidden_size * 4), activation=tf.nn.relu)(hidden_state)
+          hidden_state = tf.nn.dropout(hidden_state,
+                                       1 - hparams.layer_prepostprocess_dropout)
+
+          hidden_state = common_layers.layers().Dense(
+              hparams.hidden_size)(hidden_state)
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, hidden_state, hparams)
+
+    # If normalization is done in layer_preprocess, then it should also be done
+    # on the output, since the output can grow very large, being the sum of
+    # a whole stack of unnormalized layer outputs.
+    return common_layers.layer_preprocess(hidden_state, hparams)
+
+
+def evolved_transformer_decoder(decoder_input,
+                                encoder_output,
+                                decoder_self_attention_bias,
+                                encoder_decoder_attention_bias,
+                                hparams,
+                                cache=None,
+                                decode_loop_step=None,
+                                name="decoder",
+                                nonpadding=None,
+                                save_weights_to=None,
+                                make_image_summary=True,
+                                losses=None):
+  """Evolved Transformer decoder. See arxiv.org/abs/1901.11117 for more details.
+
+  Args:
+    decoder_input: a Tensor.
+    encoder_output: a Tensor.
+    decoder_self_attention_bias: bias Tensor for self-attention (see
+      common_attention.attention_bias()).
+    encoder_decoder_attention_bias: bias Tensor for encoder-decoder attention
+      (see common_attention.attention_bias()).
+    hparams: hyperparameters for model.
+    cache: Not supported.
+    decode_loop_step: An integer, step number of the decoding loop. Only used
+      for inference on TPU.
+    name: a string.
+    nonpadding: optional Tensor with shape [batch_size, encoder_length]
+      indicating what positions are not padding.  This is used to mask out
+      padding in convolutional layers.  We generally only need this mask for
+      "packed" datasets, because for ordinary datasets, no padding is ever
+      followed by nonpadding.
+    save_weights_to: an optional dictionary to capture attention weights for
+      visualization; the weights tensor will be appended there under a string
+      key created from the variable scope (including name).
+    make_image_summary: Whether to make an attention image summary.
+    losses: Not supported.
+
+  Returns:
+    Decoder output tensor.
+  """
+  del cache, losses
+
+  attention_dropout_broadcast_dims = (
+      common_layers.comma_separated_string_to_integer_list(
+          getattr(hparams, "attention_dropout_broadcast_dims", "")))
+
+  with tf.variable_scope(name):
+    hidden_state = decoder_input
+    layer_cache = None
+
+    for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
+      with tf.variable_scope("layer_%d" % layer):
+
+        with tf.variable_scope("16_head_self_attention"):
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          # 16 head attention. Hard coding number of heads.
+          left_state = common_attention.multihead_attention(
+              hidden_state,
+              None,
+              decoder_self_attention_bias,
+              hparams.attention_key_channels or hparams.hidden_size,
+              hparams.attention_value_channels or hparams.hidden_size,
+              hparams.hidden_size,
+              16,  # Heads are hard coded to replicate paper.
+              hparams.attention_dropout,
+              attention_type=hparams.self_attention_type,
+              max_relative_position=hparams.max_relative_position,
+              heads_share_relative_embedding=(
+                  hparams.heads_share_relative_embedding),
+              add_relative_to_values=hparams.add_relative_to_values,
+              save_weights_to=save_weights_to,
+              cache=layer_cache,
+              make_image_summary=make_image_summary,
+              dropout_broadcast_dims=attention_dropout_broadcast_dims,
+              max_length=hparams.get("max_length"),
+              decode_loop_step=decode_loop_step,
+              vars_3d=hparams.get("attention_variables_3d"),
+              activation_dtype=hparams.get("activation_dtype", "float32"),
+              weight_dtype=hparams.get("weight_dtype", "float32"))
+
+        if encoder_output is not None:
+          with tf.variable_scope("first_attend_to_encoder"):
+            right_state = common_attention.multihead_attention(
+                hidden_state,
+                encoder_output,
+                encoder_decoder_attention_bias,
+                hparams.attention_key_channels or hparams.hidden_size,
+                hparams.attention_value_channels or hparams.hidden_size,
+                hparams.hidden_size,
+                hparams.num_heads,
+                hparams.attention_dropout,
+                max_relative_position=hparams.max_relative_position,
+                heads_share_relative_embedding=(
+                    hparams.heads_share_relative_embedding),
+                add_relative_to_values=hparams.add_relative_to_values,
+                save_weights_to=save_weights_to,
+                cache=layer_cache,
+                make_image_summary=make_image_summary,
+                dropout_broadcast_dims=attention_dropout_broadcast_dims,
+                max_length=hparams.get("max_length"),
+                vars_3d=hparams.get("attention_variables_3d"),
+                activation_dtype=hparams.get("activation_dtype", "float32"),
+                weight_dtype=hparams.get("weight_dtype", "float32"))
+
+            left_state = tf.nn.dropout(left_state,
+                                       1 - hparams.layer_prepostprocess_dropout)
+            right_state = tf.nn.dropout(
+                right_state, 1 - hparams.layer_prepostprocess_dropout)
+
+            hidden_state = residual_state + left_state + right_state
+
+        else:
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, left_state, hparams)
+
+        with tf.variable_scope("conv_branches"):
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          if nonpadding is not None:
+            # Mask padding from conv layers.
+            mask = tf.tile(
+                tf.expand_dims(nonpadding, 2), [1, 1, hparams.hidden_size])
+            hidden_state *= mask
+
+          # Shift inputs so that future tokens cannot be seen.
+          left_state = tf.pad(hidden_state, paddings=[[0, 0], [10, 0], [0, 0]])
+          left_output_dim = int(hparams.hidden_size * 2)
+          separable_conv_11x1 = tf.layers.SeparableConv1D(
+              left_output_dim,
+              11,
+              padding="VALID",
+              name="separable_conv11x1",
+              activation=tf.nn.relu)
+          left_state = separable_conv_11x1.apply(left_state)
+          left_state = tf.nn.dropout(left_state,
+                                     1 - hparams.layer_prepostprocess_dropout)
+
+          right_state = tf.pad(hidden_state, paddings=[[0, 0], [6, 0], [0, 0]])
+          right_output_dim = int(hparams.hidden_size / 2)
+          separable_conv_7x1_1 = tf.layers.SeparableConv1D(
+              right_output_dim, 7, padding="VALID", name="separable_conv_7x1_1")
+          right_state = separable_conv_7x1_1.apply(right_state)
+          right_state = tf.nn.dropout(right_state,
+                                      1 - hparams.layer_prepostprocess_dropout)
+          right_state = tf.pad(
+              right_state,
+              [[0, 0], [0, 0], [0, left_output_dim - right_output_dim]],
+              constant_values=0)
+
+          hidden_state = left_state + right_state
+
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+          if nonpadding is not None:
+            # Mask padding from conv layers.
+            mask = tf.tile(
+                tf.expand_dims(nonpadding, 2), [1, 1, hparams.hidden_size * 2])
+            hidden_state *= mask
+
+          hidden_state = tf.pad(hidden_state, paddings=[[0, 0], [6, 0], [0, 0]])
+          separable_conv_7x1_2 = tf.layers.SeparableConv1D(
+              hparams.hidden_size,
+              7,
+              padding="VALID",
+              name="separable_conv_7x1_2")
+          hidden_state = separable_conv_7x1_2.apply(hidden_state)
+
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, hidden_state, hparams)
+
+        with tf.variable_scope("self_attention"):
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          hidden_state = common_attention.multihead_attention(
+              hidden_state,
+              None,
+              decoder_self_attention_bias,
+              hparams.attention_key_channels or hparams.hidden_size,
+              hparams.attention_value_channels or hparams.hidden_size,
+              hparams.hidden_size,
+              hparams.num_heads,
+              hparams.attention_dropout,
+              attention_type=hparams.self_attention_type,
+              max_relative_position=hparams.max_relative_position,
+              heads_share_relative_embedding=(
+                  hparams.heads_share_relative_embedding),
+              add_relative_to_values=hparams.add_relative_to_values,
+              save_weights_to=save_weights_to,
+              cache=layer_cache,
+              make_image_summary=make_image_summary,
+              dropout_broadcast_dims=attention_dropout_broadcast_dims,
+              max_length=hparams.get("max_length"),
+              decode_loop_step=decode_loop_step,
+              vars_3d=hparams.get("attention_variables_3d"),
+              activation_dtype=hparams.get("activation_dtype", "float32"),
+              weight_dtype=hparams.get("weight_dtype", "float32"))
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, hidden_state, hparams)
+
+        if encoder_output is not None:
+          with tf.variable_scope("second_attend_to_encoder"):
+            residual_state = hidden_state
+            hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+            hidden_state = common_attention.multihead_attention(
+                hidden_state,
+                encoder_output,
+                encoder_decoder_attention_bias,
+                hparams.attention_key_channels or hparams.hidden_size,
+                hparams.attention_value_channels or hparams.hidden_size,
+                hparams.hidden_size,
+                hparams.num_heads,
+                hparams.attention_dropout,
+                max_relative_position=hparams.max_relative_position,
+                heads_share_relative_embedding=(
+                    hparams.heads_share_relative_embedding),
+                add_relative_to_values=hparams.add_relative_to_values,
+                save_weights_to=save_weights_to,
+                cache=layer_cache,
+                make_image_summary=make_image_summary,
+                dropout_broadcast_dims=attention_dropout_broadcast_dims,
+                max_length=hparams.get("max_length"),
+                vars_3d=hparams.get("attention_variables_3d"),
+                activation_dtype=hparams.get("activation_dtype", "float32"),
+                weight_dtype=hparams.get("weight_dtype", "float32"))
+            hidden_state = common_layers.layer_postprocess(
+                residual_state, hidden_state, hparams)
+
+        with tf.variable_scope("dense_layers"):
+          residual_state = hidden_state
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          hidden_state = tf.layers.dense(
+              hidden_state,
+              int(hparams.hidden_size * 4),
+              activation=tf.nn.swish)
+          hidden_state = tf.nn.dropout(hidden_state,
+                                       1 - hparams.layer_prepostprocess_dropout)
+
+          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+
+          hidden_state = tf.layers.dense(hidden_state, hparams.hidden_size)
+          hidden_state = common_layers.layer_postprocess(
+              residual_state, hidden_state, hparams)
+
+    return common_layers.layer_preprocess(hidden_state, hparams)
+
+
 # TODO(davidso): Update optimizer, learning rate, and decay to match paper.
 def add_evolved_transformer_hparams(hparams):
   """Add Evolved Transformer hparams.
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 5361c4016..4fead25d2 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1458,257 +1458,6 @@ def transformer_decoder(decoder_input,
         x, hparams, layer_collection=layer_collection)
 
 
-def evolved_transformer_decoder(decoder_input,
-                                encoder_output,
-                                decoder_self_attention_bias,
-                                encoder_decoder_attention_bias,
-                                hparams,
-                                cache=None,
-                                decode_loop_step=None,
-                                name="decoder",
-                                nonpadding=None,
-                                save_weights_to=None,
-                                make_image_summary=True,
-                                losses=None):
-  """Evolved Transformer decoder. See arxiv.org/abs/1901.11117 for more details.
-
-  Args:
-    decoder_input: a Tensor.
-    encoder_output: a Tensor.
-    decoder_self_attention_bias: bias Tensor for self-attention (see
-      common_attention.attention_bias()).
-    encoder_decoder_attention_bias: bias Tensor for encoder-decoder attention
-      (see common_attention.attention_bias()).
-    hparams: hyperparameters for model.
-    cache: Not supported.
-    decode_loop_step: An integer, step number of the decoding loop. Only used
-      for inference on TPU.
-    name: a string.
-    nonpadding: optional Tensor with shape [batch_size, encoder_length]
-      indicating what positions are not padding.  This is used to mask out
-      padding in convolutional layers.  We generally only need this mask for
-      "packed" datasets, because for ordinary datasets, no padding is ever
-      followed by nonpadding.
-    save_weights_to: an optional dictionary to capture attention weights for
-      visualization; the weights tensor will be appended there under a string
-      key created from the variable scope (including name).
-    make_image_summary: Whether to make an attention image summary.
-    losses: Not supported.
-
-  Returns:
-    Decoder output tensor.
-  """
-  del cache, losses
-
-  attention_dropout_broadcast_dims = (
-      common_layers.comma_separated_string_to_integer_list(
-          getattr(hparams, "attention_dropout_broadcast_dims", "")))
-
-  with tf.variable_scope(name):
-    hidden_state = decoder_input
-    layer_cache = None
-
-    for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
-      with tf.variable_scope("layer_%d" % layer):
-
-        with tf.variable_scope("16_head_self_attention"):
-          residual_state = hidden_state
-          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
-
-          # 16 head attention. Hard coding number of heads.
-          left_state = common_attention.multihead_attention(
-              hidden_state,
-              None,
-              decoder_self_attention_bias,
-              hparams.attention_key_channels or hparams.hidden_size,
-              hparams.attention_value_channels or hparams.hidden_size,
-              hparams.hidden_size,
-              16,  # Heads are hard coded to replicate paper.
-              hparams.attention_dropout,
-              attention_type=hparams.self_attention_type,
-              max_relative_position=hparams.max_relative_position,
-              heads_share_relative_embedding=(
-                  hparams.heads_share_relative_embedding),
-              add_relative_to_values=hparams.add_relative_to_values,
-              save_weights_to=save_weights_to,
-              cache=layer_cache,
-              make_image_summary=make_image_summary,
-              dropout_broadcast_dims=attention_dropout_broadcast_dims,
-              max_length=hparams.get("max_length"),
-              decode_loop_step=decode_loop_step,
-              vars_3d=hparams.get("attention_variables_3d"),
-              activation_dtype=hparams.get("activation_dtype", "float32"),
-              weight_dtype=hparams.get("weight_dtype", "float32"))
-
-        if encoder_output is not None:
-          with tf.variable_scope("first_attend_to_encoder"):
-            right_state = common_attention.multihead_attention(
-                hidden_state,
-                encoder_output,
-                encoder_decoder_attention_bias,
-                hparams.attention_key_channels or hparams.hidden_size,
-                hparams.attention_value_channels or hparams.hidden_size,
-                hparams.hidden_size,
-                hparams.num_heads,
-                hparams.attention_dropout,
-                max_relative_position=hparams.max_relative_position,
-                heads_share_relative_embedding=(
-                    hparams.heads_share_relative_embedding),
-                add_relative_to_values=hparams.add_relative_to_values,
-                save_weights_to=save_weights_to,
-                cache=layer_cache,
-                make_image_summary=make_image_summary,
-                dropout_broadcast_dims=attention_dropout_broadcast_dims,
-                max_length=hparams.get("max_length"),
-                vars_3d=hparams.get("attention_variables_3d"),
-                activation_dtype=hparams.get("activation_dtype", "float32"),
-                weight_dtype=hparams.get("weight_dtype", "float32"))
-
-            left_state = tf.nn.dropout(left_state,
-                                       1 - hparams.layer_prepostprocess_dropout)
-            right_state = tf.nn.dropout(
-                right_state, 1 - hparams.layer_prepostprocess_dropout)
-
-            hidden_state = residual_state + left_state + right_state
-
-        else:
-          hidden_state = common_layers.layer_postprocess(
-              residual_state, left_state, hparams)
-
-        with tf.variable_scope("conv_branches"):
-          residual_state = hidden_state
-          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
-
-          if nonpadding:
-            # Mask padding from conv layers.
-            mask = tf.tile(
-                tf.expand_dims(nonpadding, 2), [1, 1, hparams.hidden_size])
-            hidden_state *= mask
-
-          # Shift inputs so that future tokens cannot be seen.
-          left_state = tf.pad(hidden_state, paddings=[[0, 0], [10, 0], [0, 0]])
-          left_output_dim = int(hparams.hidden_size * 2)
-          separable_conv_11x1 = tf.layers.SeparableConv1D(
-              left_output_dim,
-              11,
-              padding="VALID",
-              name="separable_conv11x1",
-              activation=tf.nn.relu)
-          left_state = separable_conv_11x1.apply(left_state)
-          left_state = tf.nn.dropout(left_state,
-                                     1 - hparams.layer_prepostprocess_dropout)
-
-          right_state = tf.pad(hidden_state, paddings=[[0, 0], [6, 0], [0, 0]])
-          right_output_dim = int(hparams.hidden_size / 2)
-          separable_conv_7x1_1 = tf.layers.SeparableConv1D(
-              right_output_dim, 7, padding="VALID", name="separable_conv_7x1_1")
-          right_state = separable_conv_7x1_1.apply(right_state)
-          right_state = tf.nn.dropout(right_state,
-                                      1 - hparams.layer_prepostprocess_dropout)
-          right_state = tf.pad(
-              right_state,
-              [[0, 0], [0, 0], [0, left_output_dim - right_output_dim]],
-              constant_values=0)
-
-          hidden_state = left_state + right_state
-
-          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
-          if nonpadding:
-            # Mask padding from conv layers.
-            mask = tf.tile(
-                tf.expand_dims(nonpadding, 2), [1, 1, hparams.hidden_size])
-            hidden_state *= mask
-
-          hidden_state = tf.pad(hidden_state, paddings=[[0, 0], [6, 0], [0, 0]])
-          separable_conv_7x1_2 = tf.layers.SeparableConv1D(
-              hparams.hidden_size,
-              7,
-              padding="VALID",
-              name="separable_conv_7x1_2")
-          hidden_state = separable_conv_7x1_2.apply(hidden_state)
-
-          hidden_state = common_layers.layer_postprocess(
-              residual_state, hidden_state, hparams)
-
-        with tf.variable_scope("self_attention"):
-          residual_state = hidden_state
-          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
-
-          hidden_state = common_attention.multihead_attention(
-              hidden_state,
-              None,
-              decoder_self_attention_bias,
-              hparams.attention_key_channels or hparams.hidden_size,
-              hparams.attention_value_channels or hparams.hidden_size,
-              hparams.hidden_size,
-              hparams.num_heads,
-              hparams.attention_dropout,
-              attention_type=hparams.self_attention_type,
-              max_relative_position=hparams.max_relative_position,
-              heads_share_relative_embedding=(
-                  hparams.heads_share_relative_embedding),
-              add_relative_to_values=hparams.add_relative_to_values,
-              save_weights_to=save_weights_to,
-              cache=layer_cache,
-              make_image_summary=make_image_summary,
-              dropout_broadcast_dims=attention_dropout_broadcast_dims,
-              max_length=hparams.get("max_length"),
-              decode_loop_step=decode_loop_step,
-              vars_3d=hparams.get("attention_variables_3d"),
-              activation_dtype=hparams.get("activation_dtype", "float32"),
-              weight_dtype=hparams.get("weight_dtype", "float32"))
-          hidden_state = common_layers.layer_postprocess(
-              residual_state, hidden_state, hparams)
-
-        if encoder_output is not None:
-          with tf.variable_scope("second_attend_to_encoder"):
-            residual_state = hidden_state
-            hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
-
-            hidden_state = common_attention.multihead_attention(
-                hidden_state,
-                encoder_output,
-                encoder_decoder_attention_bias,
-                hparams.attention_key_channels or hparams.hidden_size,
-                hparams.attention_value_channels or hparams.hidden_size,
-                hparams.hidden_size,
-                hparams.num_heads,
-                hparams.attention_dropout,
-                max_relative_position=hparams.max_relative_position,
-                heads_share_relative_embedding=(
-                    hparams.heads_share_relative_embedding),
-                add_relative_to_values=hparams.add_relative_to_values,
-                save_weights_to=save_weights_to,
-                cache=layer_cache,
-                make_image_summary=make_image_summary,
-                dropout_broadcast_dims=attention_dropout_broadcast_dims,
-                max_length=hparams.get("max_length"),
-                vars_3d=hparams.get("attention_variables_3d"),
-                activation_dtype=hparams.get("activation_dtype", "float32"),
-                weight_dtype=hparams.get("weight_dtype", "float32"))
-            hidden_state = common_layers.layer_postprocess(
-                residual_state, hidden_state, hparams)
-
-        with tf.variable_scope("dense_layers"):
-          residual_state = hidden_state
-          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
-
-          hidden_state = tf.layers.dense(
-              hidden_state,
-              int(hparams.hidden_size * 4),
-              activation=tf.nn.swish)
-          hidden_state = tf.nn.dropout(hidden_state,
-                                       1 - hparams.layer_prepostprocess_dropout)
-
-          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
-
-          hidden_state = tf.layers.dense(hidden_state, hparams.hidden_size)
-          hidden_state = common_layers.layer_postprocess(
-              residual_state, hidden_state, hparams)
-
-    return common_layers.layer_preprocess(hidden_state, hparams)
-
-
 @registry.register_hparams
 def transformer_base_v1():
   """Set of hyperparameters."""

From 33b8c5e1d79a9dbf77fb13cfb093673bc9674d29 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 1 Mar 2019 21:51:33 +0000
Subject: [PATCH 1735/2720] Prefer literals over list and dict functions
 (#1454)

* d = dict() --> d = {}

* l = list() --> []
---
 tensor2tensor/data_generators/audio.py            | 2 +-
 tensor2tensor/data_generators/gym_env.py          | 2 +-
 tensor2tensor/data_generators/gym_env_test.py     | 8 ++++----
 tensor2tensor/data_generators/librispeech.py      | 2 +-
 tensor2tensor/data_generators/mscoco.py           | 2 +-
 tensor2tensor/models/research/attention_lm_moe.py | 2 +-
 tensor2tensor/rl/player_utils.py                  | 2 +-
 tensor2tensor/utils/hparam.py                     | 2 +-
 tensor2tensor/utils/metrics.py                    | 4 ++--
 tensor2tensor/utils/rouge.py                      | 2 +-
 tensor2tensor/utils/t2t_model.py                  | 2 +-
 11 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/data_generators/audio.py b/tensor2tensor/data_generators/audio.py
index 2ad7ce29e..7a3480f6f 100644
--- a/tensor2tensor/data_generators/audio.py
+++ b/tensor2tensor/data_generators/audio.py
@@ -60,7 +60,7 @@ def _collect_data(directory, input_ext, target_ext):
   #   if the datafile was "/path/to/datafile.wav" then the key would be
   #   "/path/to/datafile"
   # value: a pair of strings (input_filepath, target_filepath)
-  data_files = dict()
+  data_files = {}
   for root, _, filenames in os.walk(directory):
     input_files = [filename for filename in filenames if input_ext in filename]
     for input_filename in input_files:
diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 8161102e4..06989c417 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -640,7 +640,7 @@ def __init__(self, base_env_name=None, batch_size=1, grayscale=False,
 
     if self.should_derive_observation_space:
       with self._tf_graph.obj.as_default():
-        self._resize = dict()
+        self._resize = {}
         orig_height, orig_width = orig_observ_space.shape[:2]
         self._img_batch_t = _Noncopyable(tf.placeholder(
             dtype=tf.uint8, shape=(None, orig_height, orig_width, 3)))
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 011b6a273..3efd54715 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -84,8 +84,8 @@ def setUp(self):
   def init_batch_and_play(self, env_name, steps_per_epoch=1, epochs=(0,),
                           generate_data=False, batch_size=2, **kwargs):
     env = gym_env.T2TGymEnv(env_name, batch_size=batch_size, **kwargs)
-    obs = list()
-    rewards = list()
+    obs = []
+    rewards = []
     num_dones = 0
     for epoch in epochs:
       env.start_new_epoch(epoch, self.out_dir)
@@ -100,8 +100,8 @@ def init_batch_and_play(self, env_name, steps_per_epoch=1, epochs=(0,),
     return env, obs, rewards, num_dones
 
   def play(self, env, n_steps):
-    obs = list()
-    rewards = list()
+    obs = []
+    rewards = []
     obs.append(env.reset())
     num_dones = 0
     for _ in range(n_steps):
diff --git a/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index 2161edfd6..1bab1d79a 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -67,7 +67,7 @@ def _collect_data(directory, input_ext, transcription_ext):
   #   if the datafile was "/path/to/datafile.wav" then the key would be
   #   "/path/to/datafile"
   # value: a pair of strings (media_filepath, label)
-  data_files = dict()
+  data_files = {}
   for root, _, filenames in os.walk(directory):
     transcripts = [filename for filename in filenames
                    if transcription_ext in filename]
diff --git a/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index 93c51d3ee..93d5f68f7 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -107,7 +107,7 @@ def get_vocab():
   caption_file = io.open(caption_filepath)
   caption_json = json.load(caption_file)
   # Dictionary from image_id to ((filename, height, width), captions).
-  image_dict = dict()
+  image_dict = {}
   for image in caption_json["images"]:
     image_dict[image["id"]] = [(image["file_name"], image["height"],
                                 image["width"]), []]
diff --git a/tensor2tensor/models/research/attention_lm_moe.py b/tensor2tensor/models/research/attention_lm_moe.py
index 7cfadb1ed..6aa7112e0 100644
--- a/tensor2tensor/models/research/attention_lm_moe.py
+++ b/tensor2tensor/models/research/attention_lm_moe.py
@@ -308,7 +308,7 @@ def print_shape(x, suffix, debug=False):
                 x,
                 hparams.filter_size)
           else:
-            additional_conv_params = dict()
+            additional_conv_params = {}
             if hparams.use_sepconv:
               additional_conv_params = dict(
                   padding="LEFT",
diff --git a/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index 41789c868..59317f80f 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -389,7 +389,7 @@ def infer_paths(output_dir, **subdirs):
   Returns:
     a dictionary with the directories.
   """
-  directories = dict()
+  directories = {}
   for name, path in six.iteritems(subdirs):
     directories[name] = path if path else os.path.join(output_dir, name)
   directories["output_dir"] = output_dir
diff --git a/tensor2tensor/utils/hparam.py b/tensor2tensor/utils/hparam.py
index aeb53cc2e..6d056438a 100644
--- a/tensor2tensor/utils/hparam.py
+++ b/tensor2tensor/utils/hparam.py
@@ -495,7 +495,7 @@ def parse(self, values):
       ValueError: If `values` cannot be parsed or a hyperparameter in `values`
       doesn't exist.
     """
-    type_map = dict()
+    type_map = {}
     for name, t in self._hparam_types.items():
       param_type, _ = t
       type_map[name] = param_type
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index d3852c648..52f9cdebf 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -600,7 +600,7 @@ def image_wrapped_metric_fn(predictions,
   def weights_fn_for_mp(problem_task_id):
     return lambda x: common_layers.weights_multi_problem(x, problem_task_id)
 
-  eval_metrics = dict()
+  eval_metrics = {}
   for problem_instance in problems:
     problem_name = problem_instance.name
     if problem_instance.was_reversed:
@@ -681,7 +681,7 @@ def create_eager_metrics_internal(metric_fns,
     (accum_fn(predictions, targets) => None,
      result_fn() => dict<str metric_name, float avg_val>
   """
-  tfe_metrics = dict()
+  tfe_metrics = {}
 
   for name in metric_fns:
     tfe_metrics[name] = tfe.metrics.Mean(name=name)
diff --git a/tensor2tensor/utils/rouge.py b/tensor2tensor/utils/rouge.py
index e6de0df2c..c3e6c7ebc 100644
--- a/tensor2tensor/utils/rouge.py
+++ b/tensor2tensor/utils/rouge.py
@@ -62,7 +62,7 @@ def _lcs(x, y):
     Table of dictionary of coord and len lcs
   """
   n, m = len(x), len(y)
-  table = dict()
+  table = {}
   for i in range(n + 1):
     for j in range(m + 1):
       if i == 0 or j == 0:
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index f3b5c8fec..bb429a37b 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1375,7 +1375,7 @@ def multinomial_squeeze(logits, temperature=1.0):
     return samples, logits, losses
 
   def _shard_features(self, features):  # pylint: disable=missing-docstring
-    sharded_features = dict()
+    sharded_features = {}
     for k, v in sorted(six.iteritems(features)):
       v = tf.convert_to_tensor(v)
       v_shape = common_layers.shape_list(v)

From f0f09484a4af76a917aa31c35985d32716928318 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Hubert=20Bry=C5=82kowski?= <hubert+github@brylkowski.com>
Date: Fri, 1 Mar 2019 22:52:33 +0100
Subject: [PATCH 1736/2720] Update mscoco.py (#1466)

---
 tensor2tensor/data_generators/mscoco.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index 93d5f68f7..dcb029723 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -125,7 +125,7 @@ def get_vocab():
   for image_info, labels in data:
     image_filename = image_info[0]
     image_filepath = os.path.join(tmp_dir, prefix, image_filename)
-    with tf.gfile.Open(image_filepath, "r") as f:
+    with tf.gfile.Open(image_filepath, "rb") as f:
       encoded_image_data = f.read()
       height, width = image_info[1], image_info[2]
       for label in labels:

From e87e85f34b12e877237c462aa2073d76dc04c457 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 1 Mar 2019 13:52:01 -0800
Subject: [PATCH 1737/2720] Merge of PR #1454

PiperOrigin-RevId: 236372024
---
 tensor2tensor/data_generators/mscoco.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index dcb029723..93d5f68f7 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -125,7 +125,7 @@ def get_vocab():
   for image_info, labels in data:
     image_filename = image_info[0]
     image_filepath = os.path.join(tmp_dir, prefix, image_filename)
-    with tf.gfile.Open(image_filepath, "rb") as f:
+    with tf.gfile.Open(image_filepath, "r") as f:
       encoded_image_data = f.read()
       height, width = image_info[1], image_info[2]
       for label in labels:

From d88a7f79c0744fed1b64572f2cfb64a4e255612f Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 1 Mar 2019 13:54:17 -0800
Subject: [PATCH 1738/2720] Merge of PR #1468

PiperOrigin-RevId: 236372428
---
 tensor2tensor/data_generators/all_problems.py          | 3 ++-
 tensor2tensor/data_generators/generator_utils.py       | 4 ++--
 tensor2tensor/data_generators/transduction_problems.py | 6 +++---
 tensor2tensor/models/research/vqa_self_attention.py    | 4 ++--
 4 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 5e0a12ca2..83bec50d6 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -19,6 +19,7 @@
 from __future__ import print_function
 
 import importlib
+from six.moves import range  # pylint: disable=redefined-builtin
 
 MODULES = [
     "tensor2tensor.data_generators.algorithmic",
@@ -97,7 +98,7 @@
 
 def _is_import_err_msg(err_str, module):
   parts = module.split(".")
-  suffixes = [".".join(parts[i:]) for i in xrange(len(parts))]
+  suffixes = [".".join(parts[i:]) for i in range(len(parts))]
   return err_str in (
       ["No module named %s" % suffix for suffix in suffixes] +
       ["No module named '%s'" % suffix for suffix in suffixes])
diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 8a722101a..aae3e5572 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -969,11 +969,11 @@ def random_deinterleave(text, separator_symbol="X"):
   cut = [False] * n
   cut[0] = True
   num_cuts = int(math.exp(random.uniform(0, math.log(n))))
-  for _ in xrange(num_cuts):
+  for _ in range(num_cuts):
     cut[random.randint(1, n -1)] = True
   out = [[], []]
   part = random.randint(0, 1)
-  for i in xrange(n):
+  for i in range(n):
     if cut[i]:
       out[part].append(separator_symbol)
       part = 1 - part
diff --git a/tensor2tensor/data_generators/transduction_problems.py b/tensor2tensor/data_generators/transduction_problems.py
index 6e1cff743..37082dcbf 100644
--- a/tensor2tensor/data_generators/transduction_problems.py
+++ b/tensor2tensor/data_generators/transduction_problems.py
@@ -33,7 +33,7 @@
 import os
 import random
 
-from six.moves import xrange  # pylint: disable=redefined-builtin
+from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
@@ -123,7 +123,7 @@ def sequence_length(self, dataset_split):
                           self.max_sequence_length(dataset_split))
 
   def build_vocab(self):
-    return ["sym_%d" % i for i in xrange(1, self.num_symbols + 1)]
+    return ["sym_%d" % i for i in range(1, self.num_symbols + 1)]
 
   def get_or_create_vocab(self, data_dir, tmp_dir, force_get=False):
     vocab_filename = os.path.join(data_dir, self.vocab_filename)
@@ -144,7 +144,7 @@ def transpose_sequence(self, input_sequence):
     raise NotImplementedError()
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    for _ in xrange(self.num_samples(dataset_split)):
+    for _ in range(self.num_samples(dataset_split)):
       source = self.generate_random_sequence(dataset_split)
       target = self.transpose_sequence(source)
       yield {
diff --git a/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
index 4702b598f..c9ae04680 100644
--- a/tensor2tensor/models/research/vqa_self_attention.py
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -19,7 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from six.moves import xrange
+from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
@@ -657,7 +657,7 @@ def iterative_encoder_decoder(encoder_input,
                               query,
                               hparams):
   """Iterative encoder decoder."""
-  for _ in xrange(hparams.num_rec_steps):
+  for _ in range(hparams.num_rec_steps):
     with tf.variable_scope("step", reuse=tf.AUTO_REUSE):
       encoder_output = image_question_encoder(
           encoder_input,

From 55225459e9e17b85c04328348e0d1d1a8c1bfe30 Mon Sep 17 00:00:00 2001
From: Hubert Bry?kowski <hubert+github@brylkowski.com>
Date: Fri, 1 Mar 2019 14:16:00 -0800
Subject: [PATCH 1739/2720] Merge of PR #1466

PiperOrigin-RevId: 236376505
---
 tensor2tensor/data_generators/mscoco.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index 93d5f68f7..dcb029723 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -125,7 +125,7 @@ def get_vocab():
   for image_info, labels in data:
     image_filename = image_info[0]
     image_filepath = os.path.join(tmp_dir, prefix, image_filename)
-    with tf.gfile.Open(image_filepath, "r") as f:
+    with tf.gfile.Open(image_filepath, "rb") as f:
       encoded_image_data = f.read()
       height, width = image_info[1], image_info[2]
       for label in labels:

From a8c246c83bc44d1b4132330659d9de4e807447be Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 1 Mar 2019 14:40:26 -0800
Subject: [PATCH 1740/2720] Correct modalities for TPU eval after recent
 changes.

PiperOrigin-RevId: 236380935
---
 tensor2tensor/utils/t2t_model.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index bb429a37b..f81caf5e2 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1817,7 +1817,7 @@ def create_tpu_eval_metrics_fn(problem, model_hparams):
   tm = _create_target_modality(problem.get_hparams(model_hparams).modality)
   if isinstance(tm, dict):
     for k, v in six.iteritems(tm):
-      weights_fn = v.targets_weights_fn
+      weights_fn = modalities.get_targets_weights_fn(v)
 
       def make_metric_fn(metric_fn):
         def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
@@ -1837,7 +1837,7 @@ def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
         name = "%s/metrics-%s/%s" % (k, problem.name, metric)
         metric_fns.append((name, make_metric_fn(metric_fn)))
   else:
-    weights_fn = tm.targets_weights_fn
+    weights_fn = modalities.get_targets_weights_fn(tm)
 
     def make_metric_fn(metric_fn):
       def wrapped_metric_fn(logits, labels, features):

From abd5ea21dd8b9c561fdb5c26ef5f720b0a454b4b Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 1 Mar 2019 23:02:15 +0000
Subject: [PATCH 1741/2720] Replace .append loop with list comprehension
 (#1451)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This PR replaces loops that only append to a list with list comprehensions. This makes it more concise and arguably more readably.
List comprehensions also are quite a bit faster than appending to a list.
Toy example in Python 3.6:
```python
In [1]: %%timeit
   ...: l = []
   ...: for i in range(5000):
   ...:     l.append(i)
375 µs ± 1.73 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [2]: %%timeit
   ...: l = [i for i in range(5000)]
168 µs ± 1.08 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
```
---
 tensor2tensor/data_generators/cnn_dailymail.py   |  4 +---
 .../data_generators/multi_problem_v2.py          |  5 ++---
 tensor2tensor/data_generators/video_utils.py     |  9 ++-------
 tensor2tensor/envs/env_problem.py                | 12 +++++-------
 tensor2tensor/insights/server.py                 | 12 +++++-------
 tensor2tensor/layers/discretization.py           | 14 +++++---------
 tensor2tensor/layers/vq_discrete.py              | 16 +++++++---------
 tensor2tensor/models/video/base.py               |  7 +++----
 tensor2tensor/rl/player.py                       |  9 +++------
 tensor2tensor/utils/data_reader.py               |  4 +---
 tensor2tensor/visualization/attention.py         |  6 +-----
 11 files changed, 35 insertions(+), 63 deletions(-)

diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index 5e77b5599..af11c928b 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -118,9 +118,7 @@ def generate_hash(inp):
 
   all_files_map = {f.split("/")[-1]: f for f in all_files}
 
-  urls = []
-  for line in tf.gfile.Open(url_file):
-    urls.append(line.strip().encode("utf-8"))
+  urls = [line.strip().encode("utf-8") for line in tf.gfile.Open(url_file)]
 
   filelist = []
   for url in urls:
diff --git a/tensor2tensor/data_generators/multi_problem_v2.py b/tensor2tensor/data_generators/multi_problem_v2.py
index 63a6cddf0..62bd8c4ed 100644
--- a/tensor2tensor/data_generators/multi_problem_v2.py
+++ b/tensor2tensor/data_generators/multi_problem_v2.py
@@ -370,9 +370,8 @@ def epoch_rates_to_pmf(problems, epoch_rates=None):
   """
   if epoch_rates is None:
     epoch_rates = [1.0] * len(problems)
-  example_rates = []
-  for p, epoch_rate in zip(problems, epoch_rates):
-    example_rates.append(epoch_rate * p.num_training_examples)
+  example_rates = [epoch_rate * p.num_training_examples
+                   for p, epoch_rate in zip(problems, epoch_rates)]
   return example_rates_to_pmf(example_rates)
 
 
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 4dc5560ff..4939efcc3 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -45,13 +45,8 @@
 
 
 def resize_video_frames(images, size):
-  resized_images = []
-  for image in images:
-    resized_images.append(
-        tf.to_int64(
-            tf.image.resize_images(image, [size, size],
-                                   tf.image.ResizeMethod.BILINEAR)))
-  return resized_images
+  return [tf.to_int64(tf.image.resize_images(
+      image, [size, size], tf.image.ResizeMethod.BILINEAR)) for image in images]
 
 
 def video_augmentation(features, hue=False, saturate=False, contrast=False):
diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index ee8e62878..450a94488 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -256,13 +256,11 @@ def initialize_environments(self, batch_size=1, max_episode_steps=-1,
     assert batch_size >= 1
     self._batch_size = batch_size
 
-    self._envs = []
-    for _ in range(batch_size):
-      self._envs.append(
-          gym_utils.make_gym_env(
-              self.base_env_name,
-              rl_env_max_episode_steps=max_episode_steps,
-              maxskip_env=max_and_skip_env))
+    self._envs = [
+        gym_utils.make_gym_env(
+            self.base_env_name,
+            rl_env_max_episode_steps=max_episode_steps,
+            maxskip_env=max_and_skip_env) for _ in range(batch_size)]
 
     # If self.observation_space and self.action_space aren't None, then it means
     # that this is a re-initialization of this class, in that case make sure
diff --git a/tensor2tensor/insights/server.py b/tensor2tensor/insights/server.py
index ed2b15aa8..e5a1c16ca 100644
--- a/tensor2tensor/insights/server.py
+++ b/tensor2tensor/insights/server.py
@@ -145,13 +145,11 @@ def list_models():  # pylint: disable=unused-variable
     Returns:
       JSON for the supported models.
     """
-    configuration_list = []
-    for source_code, target_code, label in processors:
-      configuration_list.append({
-          "id": label,
-          "source_language": languages[source_code],
-          "target_language": languages[target_code],
-      })
+    configuration_list = [{
+        "id": label,
+        "source_language": languages[source_code],
+        "target_language": languages[target_code],
+        } for source_code, target_code, label in processors]
     return jsonify({
         "configuration": configuration_list
     })
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 30daeecf6..669f6ae6e 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -235,9 +235,8 @@ def bit_to_int(x_bit, num_bits, base=2):
     Integer representation of this number.
   """
   x_l = tf.stop_gradient(tf.to_int32(tf.reshape(x_bit, [-1, num_bits])))
-  x_labels = []
-  for i in range(num_bits):
-    x_labels.append(x_l[:, i] * tf.to_int32(base)**tf.to_int32(i))
+  x_labels = [
+      x_l[:, i] * tf.to_int32(base)**tf.to_int32(i) for i in range(num_bits)]
   res = sum(x_labels)
   return tf.to_int32(tf.reshape(res, common_layers.shape_list(x_bit)[:-1]))
 
@@ -254,12 +253,9 @@ def int_to_bit(x_int, num_bits, base=2):
     Corresponding number expressed in base.
   """
   x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
-  x_labels = []
-  for i in range(num_bits):
-    x_labels.append(
-        tf.floormod(
-            tf.floordiv(tf.to_int32(x_l),
-                        tf.to_int32(base)**i), tf.to_int32(base)))
+  x_labels = [tf.floormod(
+      tf.floordiv(tf.to_int32(x_l), tf.to_int32(base)**i), tf.to_int32(base))
+      for i in range(num_bits)]
   res = tf.concat(x_labels, axis=-1)
   return tf.to_float(res)
 
diff --git a/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index ff6c8ddd1..eb4937cd3 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -158,9 +158,8 @@ def bit_to_int(self, x_bit, num_bits, base=2):
         Integer representation of this number.
     """
     x_l = tf.stop_gradient(tf.to_int32(tf.reshape(x_bit, [-1, num_bits])))
-    x_labels = []
-    for i in range(num_bits):
-      x_labels.append(x_l[:, i] * tf.to_int32(base)**tf.to_int32(i))
+    x_labels = [
+        x_l[:, i] * tf.to_int32(base)**tf.to_int32(i) for i in range(num_bits)]
     res = sum(x_labels)
     return tf.to_int32(tf.reshape(res, common_layers.shape_list(x_bit)[:-1]))
 
@@ -177,12 +176,11 @@ def int_to_bit(self, x_int, num_bits, base=2):
         Corresponding number expressed in base.
     """
     x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
-    x_labels = []
-    for i in range(num_bits):
-      x_labels.append(
-          tf.floormod(
-              tf.floordiv(tf.to_int32(x_l),
-                          tf.to_int32(base)**i), tf.to_int32(base)))
+    x_labels = [
+        tf.floormod(
+            tf.floordiv(tf.to_int32(x_l),
+                        tf.to_int32(base)**i), tf.to_int32(base))
+        for i in range(num_bits)]
     res = tf.concat(x_labels, axis=-1)
     return tf.to_float(res)
 
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index c3dfb53ea..637cb7ccc 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -303,10 +303,9 @@ def get_scheduled_sample_inputs(self,
     def sample():
       """Calculate the scheduled sampling params based on iteration number."""
       with tf.variable_scope("scheduled_sampling", reuse=tf.AUTO_REUSE):
-        output_items = []
-        for item_gt, item_gen in zip(groundtruth_items, generated_items):
-          output_items.append(scheduled_sampling_func(item_gt, item_gen))
-        return output_items
+        return [
+            scheduled_sampling_func(item_gt, item_gen)
+            for item_gt, item_gen in zip(groundtruth_items, generated_items)]
 
     cases = [
         (tf.logical_not(done_warm_start), lambda: groundtruth_items),
diff --git a/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 4ec8f0e2d..5a66d3d35 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -177,12 +177,9 @@ def get_keys_to_action(self):
     keys_to_action = {}
 
     for action_id, action_meaning in enumerate(self.action_meanings):
-      keys = []
-      for keyword, key in keyword_to_key.items():
-        if keyword in action_meaning:
-          keys.append(key)
-      keys_tuple = tuple(sorted(keys))
-      del keys
+      keys_tuple = tuple(sorted([
+          key for keyword, key in keyword_to_key.items()
+          if keyword in action_meaning]))
       assert keys_tuple not in keys_to_action
       keys_to_action[keys_tuple] = action_id
 
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index 3456c7fe0..995c47a31 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -300,9 +300,7 @@ def _pad_batch(features, batch_multiple):
   padded_features = {}
   for k, feature in features.items():
     rank = len(feature.shape)
-    paddings = []
-    for _ in range(rank):
-      paddings.append([0, 0])
+    paddings = [[0, 0] for _ in range(rank)]
     paddings[0][1] = batch_padding
     padded_feature = tf.pad(feature, paddings)
     padded_features[k] = padded_feature
diff --git a/tensor2tensor/visualization/attention.py b/tensor2tensor/visualization/attention.py
index 739edf79e..2b3b2e894 100644
--- a/tensor2tensor/visualization/attention.py
+++ b/tensor2tensor/visualization/attention.py
@@ -135,11 +135,7 @@ def get_out_out_attention(layer):
 
   def get_attentions(get_attention_fn):
     num_layers = len(enc_atts)
-    attentions = []
-    for i in range(num_layers):
-      attentions.append(get_attention_fn(i))
-
-    return attentions
+    return [get_attention_fn(i) for i in range(num_layers)]
 
   attentions = {
       'all': {

From 7a9d6659f298e32f73f8d770e88af09e491cddfb Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 1 Mar 2019 23:02:48 +0000
Subject: [PATCH 1742/2720] Use xrange from six to fix Python 3 support (#1468)

T2T uses a mix of `xrange` and `range` leading to Python 3 errors when `xrange` isn't imported from `six.moves`.
This PR switches all usage of `xrange` to `six.moves.range` which is equivalent and supports both Python 3 and 2.

From a5191bbaeec3b4d4090d0f135243dea148bf652c Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 1 Mar 2019 15:13:14 -0800
Subject: [PATCH 1743/2720] Define the name `params` outside the training loop,
 since if we don't go in it the name is not initialized at the `return
 State(params=params ...` line leading to a crash.

Discovered it while trying to restore and run model from an already run output
dir.

PiperOrigin-RevId: 236386819
---
 tensor2tensor/data_generators/cnn_dailymail.py   |  4 +++-
 .../data_generators/multi_problem_v2.py          |  5 +++--
 tensor2tensor/data_generators/video_utils.py     |  9 +++++++--
 tensor2tensor/envs/env_problem.py                | 12 +++++++-----
 tensor2tensor/insights/server.py                 | 12 +++++++-----
 tensor2tensor/layers/discretization.py           | 14 +++++++++-----
 tensor2tensor/layers/vq_discrete.py              | 16 +++++++++-------
 tensor2tensor/models/video/base.py               |  7 ++++---
 tensor2tensor/rl/player.py                       |  9 ++++++---
 tensor2tensor/trax/trax.py                       |  3 ++-
 tensor2tensor/utils/data_reader.py               |  4 +++-
 tensor2tensor/visualization/attention.py         |  6 +++++-
 12 files changed, 65 insertions(+), 36 deletions(-)

diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index af11c928b..5e77b5599 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -118,7 +118,9 @@ def generate_hash(inp):
 
   all_files_map = {f.split("/")[-1]: f for f in all_files}
 
-  urls = [line.strip().encode("utf-8") for line in tf.gfile.Open(url_file)]
+  urls = []
+  for line in tf.gfile.Open(url_file):
+    urls.append(line.strip().encode("utf-8"))
 
   filelist = []
   for url in urls:
diff --git a/tensor2tensor/data_generators/multi_problem_v2.py b/tensor2tensor/data_generators/multi_problem_v2.py
index 62bd8c4ed..63a6cddf0 100644
--- a/tensor2tensor/data_generators/multi_problem_v2.py
+++ b/tensor2tensor/data_generators/multi_problem_v2.py
@@ -370,8 +370,9 @@ def epoch_rates_to_pmf(problems, epoch_rates=None):
   """
   if epoch_rates is None:
     epoch_rates = [1.0] * len(problems)
-  example_rates = [epoch_rate * p.num_training_examples
-                   for p, epoch_rate in zip(problems, epoch_rates)]
+  example_rates = []
+  for p, epoch_rate in zip(problems, epoch_rates):
+    example_rates.append(epoch_rate * p.num_training_examples)
   return example_rates_to_pmf(example_rates)
 
 
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 4939efcc3..4dc5560ff 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -45,8 +45,13 @@
 
 
 def resize_video_frames(images, size):
-  return [tf.to_int64(tf.image.resize_images(
-      image, [size, size], tf.image.ResizeMethod.BILINEAR)) for image in images]
+  resized_images = []
+  for image in images:
+    resized_images.append(
+        tf.to_int64(
+            tf.image.resize_images(image, [size, size],
+                                   tf.image.ResizeMethod.BILINEAR)))
+  return resized_images
 
 
 def video_augmentation(features, hue=False, saturate=False, contrast=False):
diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 450a94488..ee8e62878 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -256,11 +256,13 @@ def initialize_environments(self, batch_size=1, max_episode_steps=-1,
     assert batch_size >= 1
     self._batch_size = batch_size
 
-    self._envs = [
-        gym_utils.make_gym_env(
-            self.base_env_name,
-            rl_env_max_episode_steps=max_episode_steps,
-            maxskip_env=max_and_skip_env) for _ in range(batch_size)]
+    self._envs = []
+    for _ in range(batch_size):
+      self._envs.append(
+          gym_utils.make_gym_env(
+              self.base_env_name,
+              rl_env_max_episode_steps=max_episode_steps,
+              maxskip_env=max_and_skip_env))
 
     # If self.observation_space and self.action_space aren't None, then it means
     # that this is a re-initialization of this class, in that case make sure
diff --git a/tensor2tensor/insights/server.py b/tensor2tensor/insights/server.py
index e5a1c16ca..ed2b15aa8 100644
--- a/tensor2tensor/insights/server.py
+++ b/tensor2tensor/insights/server.py
@@ -145,11 +145,13 @@ def list_models():  # pylint: disable=unused-variable
     Returns:
       JSON for the supported models.
     """
-    configuration_list = [{
-        "id": label,
-        "source_language": languages[source_code],
-        "target_language": languages[target_code],
-        } for source_code, target_code, label in processors]
+    configuration_list = []
+    for source_code, target_code, label in processors:
+      configuration_list.append({
+          "id": label,
+          "source_language": languages[source_code],
+          "target_language": languages[target_code],
+      })
     return jsonify({
         "configuration": configuration_list
     })
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 669f6ae6e..30daeecf6 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -235,8 +235,9 @@ def bit_to_int(x_bit, num_bits, base=2):
     Integer representation of this number.
   """
   x_l = tf.stop_gradient(tf.to_int32(tf.reshape(x_bit, [-1, num_bits])))
-  x_labels = [
-      x_l[:, i] * tf.to_int32(base)**tf.to_int32(i) for i in range(num_bits)]
+  x_labels = []
+  for i in range(num_bits):
+    x_labels.append(x_l[:, i] * tf.to_int32(base)**tf.to_int32(i))
   res = sum(x_labels)
   return tf.to_int32(tf.reshape(res, common_layers.shape_list(x_bit)[:-1]))
 
@@ -253,9 +254,12 @@ def int_to_bit(x_int, num_bits, base=2):
     Corresponding number expressed in base.
   """
   x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
-  x_labels = [tf.floormod(
-      tf.floordiv(tf.to_int32(x_l), tf.to_int32(base)**i), tf.to_int32(base))
-      for i in range(num_bits)]
+  x_labels = []
+  for i in range(num_bits):
+    x_labels.append(
+        tf.floormod(
+            tf.floordiv(tf.to_int32(x_l),
+                        tf.to_int32(base)**i), tf.to_int32(base)))
   res = tf.concat(x_labels, axis=-1)
   return tf.to_float(res)
 
diff --git a/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index eb4937cd3..ff6c8ddd1 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -158,8 +158,9 @@ def bit_to_int(self, x_bit, num_bits, base=2):
         Integer representation of this number.
     """
     x_l = tf.stop_gradient(tf.to_int32(tf.reshape(x_bit, [-1, num_bits])))
-    x_labels = [
-        x_l[:, i] * tf.to_int32(base)**tf.to_int32(i) for i in range(num_bits)]
+    x_labels = []
+    for i in range(num_bits):
+      x_labels.append(x_l[:, i] * tf.to_int32(base)**tf.to_int32(i))
     res = sum(x_labels)
     return tf.to_int32(tf.reshape(res, common_layers.shape_list(x_bit)[:-1]))
 
@@ -176,11 +177,12 @@ def int_to_bit(self, x_int, num_bits, base=2):
         Corresponding number expressed in base.
     """
     x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
-    x_labels = [
-        tf.floormod(
-            tf.floordiv(tf.to_int32(x_l),
-                        tf.to_int32(base)**i), tf.to_int32(base))
-        for i in range(num_bits)]
+    x_labels = []
+    for i in range(num_bits):
+      x_labels.append(
+          tf.floormod(
+              tf.floordiv(tf.to_int32(x_l),
+                          tf.to_int32(base)**i), tf.to_int32(base)))
     res = tf.concat(x_labels, axis=-1)
     return tf.to_float(res)
 
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 637cb7ccc..c3dfb53ea 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -303,9 +303,10 @@ def get_scheduled_sample_inputs(self,
     def sample():
       """Calculate the scheduled sampling params based on iteration number."""
       with tf.variable_scope("scheduled_sampling", reuse=tf.AUTO_REUSE):
-        return [
-            scheduled_sampling_func(item_gt, item_gen)
-            for item_gt, item_gen in zip(groundtruth_items, generated_items)]
+        output_items = []
+        for item_gt, item_gen in zip(groundtruth_items, generated_items):
+          output_items.append(scheduled_sampling_func(item_gt, item_gen))
+        return output_items
 
     cases = [
         (tf.logical_not(done_warm_start), lambda: groundtruth_items),
diff --git a/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 5a66d3d35..4ec8f0e2d 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -177,9 +177,12 @@ def get_keys_to_action(self):
     keys_to_action = {}
 
     for action_id, action_meaning in enumerate(self.action_meanings):
-      keys_tuple = tuple(sorted([
-          key for keyword, key in keyword_to_key.items()
-          if keyword in action_meaning]))
+      keys = []
+      for keyword, key in keyword_to_key.items():
+        if keyword in action_meaning:
+          keys.append(key)
+      keys_tuple = tuple(sorted(keys))
+      del keys
       assert keys_tuple not in keys_to_action
       keys_to_action[keys_tuple] = action_id
 
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 4466057fe..98b3d0568 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -216,7 +216,8 @@ def train(output_dir,
   # Setup state
   step = state.step or 0
   params_initializer = lambda: model_init([-1] + inputs.input_shape)[1]
-  opt_state = opt_init(state.params or params_initializer())
+  params = state.params or params_initializer()
+  opt_state = opt_init(params)
 
   # jit model_predict and update so they're fast
   jit_predict = jax.jit(model_predict)  # for evaluation
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index 995c47a31..3456c7fe0 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -300,7 +300,9 @@ def _pad_batch(features, batch_multiple):
   padded_features = {}
   for k, feature in features.items():
     rank = len(feature.shape)
-    paddings = [[0, 0] for _ in range(rank)]
+    paddings = []
+    for _ in range(rank):
+      paddings.append([0, 0])
     paddings[0][1] = batch_padding
     padded_feature = tf.pad(feature, paddings)
     padded_features[k] = padded_feature
diff --git a/tensor2tensor/visualization/attention.py b/tensor2tensor/visualization/attention.py
index 2b3b2e894..739edf79e 100644
--- a/tensor2tensor/visualization/attention.py
+++ b/tensor2tensor/visualization/attention.py
@@ -135,7 +135,11 @@ def get_out_out_attention(layer):
 
   def get_attentions(get_attention_fn):
     num_layers = len(enc_atts)
-    return [get_attention_fn(i) for i in range(num_layers)]
+    attentions = []
+    for i in range(num_layers):
+      attentions.append(get_attention_fn(i))
+
+    return attentions
 
   attentions = {
       'all': {

From 3b38635f12348036ea1e0166857f43a6b971ab07 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Fri, 1 Mar 2019 15:23:28 -0800
Subject: [PATCH 1744/2720] Merge of PR #1451

PiperOrigin-RevId: 236388504
---
 tensor2tensor/data_generators/cnn_dailymail.py |  4 +---
 .../data_generators/multi_problem_v2.py        |  5 ++---
 tensor2tensor/data_generators/video_utils.py   |  9 ++-------
 tensor2tensor/envs/env_problem.py              | 13 ++++++-------
 tensor2tensor/insights/server.py               | 13 ++++++-------
 tensor2tensor/layers/discretization.py         | 14 +++++---------
 tensor2tensor/layers/vq_discrete.py            | 18 +++++++++---------
 tensor2tensor/models/video/base.py             |  7 +++----
 tensor2tensor/rl/player.py                     |  9 +++------
 tensor2tensor/utils/data_reader.py             |  4 +---
 tensor2tensor/visualization/attention.py       |  6 +-----
 11 files changed, 39 insertions(+), 63 deletions(-)

diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index 5e77b5599..af11c928b 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -118,9 +118,7 @@ def generate_hash(inp):
 
   all_files_map = {f.split("/")[-1]: f for f in all_files}
 
-  urls = []
-  for line in tf.gfile.Open(url_file):
-    urls.append(line.strip().encode("utf-8"))
+  urls = [line.strip().encode("utf-8") for line in tf.gfile.Open(url_file)]
 
   filelist = []
   for url in urls:
diff --git a/tensor2tensor/data_generators/multi_problem_v2.py b/tensor2tensor/data_generators/multi_problem_v2.py
index 63a6cddf0..62bd8c4ed 100644
--- a/tensor2tensor/data_generators/multi_problem_v2.py
+++ b/tensor2tensor/data_generators/multi_problem_v2.py
@@ -370,9 +370,8 @@ def epoch_rates_to_pmf(problems, epoch_rates=None):
   """
   if epoch_rates is None:
     epoch_rates = [1.0] * len(problems)
-  example_rates = []
-  for p, epoch_rate in zip(problems, epoch_rates):
-    example_rates.append(epoch_rate * p.num_training_examples)
+  example_rates = [epoch_rate * p.num_training_examples
+                   for p, epoch_rate in zip(problems, epoch_rates)]
   return example_rates_to_pmf(example_rates)
 
 
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 4dc5560ff..4939efcc3 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -45,13 +45,8 @@
 
 
 def resize_video_frames(images, size):
-  resized_images = []
-  for image in images:
-    resized_images.append(
-        tf.to_int64(
-            tf.image.resize_images(image, [size, size],
-                                   tf.image.ResizeMethod.BILINEAR)))
-  return resized_images
+  return [tf.to_int64(tf.image.resize_images(
+      image, [size, size], tf.image.ResizeMethod.BILINEAR)) for image in images]
 
 
 def video_augmentation(features, hue=False, saturate=False, contrast=False):
diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index ee8e62878..4d7312001 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -256,13 +256,12 @@ def initialize_environments(self, batch_size=1, max_episode_steps=-1,
     assert batch_size >= 1
     self._batch_size = batch_size
 
-    self._envs = []
-    for _ in range(batch_size):
-      self._envs.append(
-          gym_utils.make_gym_env(
-              self.base_env_name,
-              rl_env_max_episode_steps=max_episode_steps,
-              maxskip_env=max_and_skip_env))
+    # pylint: disable=g-complex-comprehension
+    self._envs = [
+        gym_utils.make_gym_env(
+            self.base_env_name,
+            rl_env_max_episode_steps=max_episode_steps,
+            maxskip_env=max_and_skip_env) for _ in range(batch_size)]
 
     # If self.observation_space and self.action_space aren't None, then it means
     # that this is a re-initialization of this class, in that case make sure
diff --git a/tensor2tensor/insights/server.py b/tensor2tensor/insights/server.py
index ed2b15aa8..942e87504 100644
--- a/tensor2tensor/insights/server.py
+++ b/tensor2tensor/insights/server.py
@@ -145,13 +145,12 @@ def list_models():  # pylint: disable=unused-variable
     Returns:
       JSON for the supported models.
     """
-    configuration_list = []
-    for source_code, target_code, label in processors:
-      configuration_list.append({
-          "id": label,
-          "source_language": languages[source_code],
-          "target_language": languages[target_code],
-      })
+    # pylint: disable=g-complex-comprehension
+    configuration_list = [{
+        "id": label,
+        "source_language": languages[source_code],
+        "target_language": languages[target_code],
+        } for source_code, target_code, label in processors]
     return jsonify({
         "configuration": configuration_list
     })
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 30daeecf6..506fd6864 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -235,9 +235,8 @@ def bit_to_int(x_bit, num_bits, base=2):
     Integer representation of this number.
   """
   x_l = tf.stop_gradient(tf.to_int32(tf.reshape(x_bit, [-1, num_bits])))
-  x_labels = []
-  for i in range(num_bits):
-    x_labels.append(x_l[:, i] * tf.to_int32(base)**tf.to_int32(i))
+  x_labels = [
+      x_l[:, i] * tf.to_int32(base)**tf.to_int32(i) for i in range(num_bits)]
   res = sum(x_labels)
   return tf.to_int32(tf.reshape(res, common_layers.shape_list(x_bit)[:-1]))
 
@@ -254,12 +253,9 @@ def int_to_bit(x_int, num_bits, base=2):
     Corresponding number expressed in base.
   """
   x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
-  x_labels = []
-  for i in range(num_bits):
-    x_labels.append(
-        tf.floormod(
-            tf.floordiv(tf.to_int32(x_l),
-                        tf.to_int32(base)**i), tf.to_int32(base)))
+  x_labels = [tf.floormod(
+      tf.floordiv(tf.to_int32(x_l), tf.to_int32(base)**i), tf.to_int32(base))
+              for i in range(num_bits)]
   res = tf.concat(x_labels, axis=-1)
   return tf.to_float(res)
 
diff --git a/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index ff6c8ddd1..b593d7a87 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -158,9 +158,9 @@ def bit_to_int(self, x_bit, num_bits, base=2):
         Integer representation of this number.
     """
     x_l = tf.stop_gradient(tf.to_int32(tf.reshape(x_bit, [-1, num_bits])))
-    x_labels = []
-    for i in range(num_bits):
-      x_labels.append(x_l[:, i] * tf.to_int32(base)**tf.to_int32(i))
+    # pylint: disable=g-complex-comprehension
+    x_labels = [
+        x_l[:, i] * tf.to_int32(base)**tf.to_int32(i) for i in range(num_bits)]
     res = sum(x_labels)
     return tf.to_int32(tf.reshape(res, common_layers.shape_list(x_bit)[:-1]))
 
@@ -177,12 +177,12 @@ def int_to_bit(self, x_int, num_bits, base=2):
         Corresponding number expressed in base.
     """
     x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
-    x_labels = []
-    for i in range(num_bits):
-      x_labels.append(
-          tf.floormod(
-              tf.floordiv(tf.to_int32(x_l),
-                          tf.to_int32(base)**i), tf.to_int32(base)))
+    # pylint: disable=g-complex-comprehension
+    x_labels = [
+        tf.floormod(
+            tf.floordiv(tf.to_int32(x_l),
+                        tf.to_int32(base)**i), tf.to_int32(base))
+        for i in range(num_bits)]
     res = tf.concat(x_labels, axis=-1)
     return tf.to_float(res)
 
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index c3dfb53ea..637cb7ccc 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -303,10 +303,9 @@ def get_scheduled_sample_inputs(self,
     def sample():
       """Calculate the scheduled sampling params based on iteration number."""
       with tf.variable_scope("scheduled_sampling", reuse=tf.AUTO_REUSE):
-        output_items = []
-        for item_gt, item_gen in zip(groundtruth_items, generated_items):
-          output_items.append(scheduled_sampling_func(item_gt, item_gen))
-        return output_items
+        return [
+            scheduled_sampling_func(item_gt, item_gen)
+            for item_gt, item_gen in zip(groundtruth_items, generated_items)]
 
     cases = [
         (tf.logical_not(done_warm_start), lambda: groundtruth_items),
diff --git a/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 4ec8f0e2d..5a66d3d35 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -177,12 +177,9 @@ def get_keys_to_action(self):
     keys_to_action = {}
 
     for action_id, action_meaning in enumerate(self.action_meanings):
-      keys = []
-      for keyword, key in keyword_to_key.items():
-        if keyword in action_meaning:
-          keys.append(key)
-      keys_tuple = tuple(sorted(keys))
-      del keys
+      keys_tuple = tuple(sorted([
+          key for keyword, key in keyword_to_key.items()
+          if keyword in action_meaning]))
       assert keys_tuple not in keys_to_action
       keys_to_action[keys_tuple] = action_id
 
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index 3456c7fe0..995c47a31 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -300,9 +300,7 @@ def _pad_batch(features, batch_multiple):
   padded_features = {}
   for k, feature in features.items():
     rank = len(feature.shape)
-    paddings = []
-    for _ in range(rank):
-      paddings.append([0, 0])
+    paddings = [[0, 0] for _ in range(rank)]
     paddings[0][1] = batch_padding
     padded_feature = tf.pad(feature, paddings)
     padded_features[k] = padded_feature
diff --git a/tensor2tensor/visualization/attention.py b/tensor2tensor/visualization/attention.py
index 739edf79e..2b3b2e894 100644
--- a/tensor2tensor/visualization/attention.py
+++ b/tensor2tensor/visualization/attention.py
@@ -135,11 +135,7 @@ def get_out_out_attention(layer):
 
   def get_attentions(get_attention_fn):
     num_layers = len(enc_atts)
-    attentions = []
-    for i in range(num_layers):
-      attentions.append(get_attention_fn(i))
-
-    return attentions
+    return [get_attention_fn(i) for i in range(num_layers)]
 
   attentions = {
       'all': {

From c05874acf2af1397ca3f0a7f7e1e351b77f938ba Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 1 Mar 2019 17:23:20 -0800
Subject: [PATCH 1745/2720] Remove Modality classes.

Originally, users override default feature transformations by setting hparams.modality. Now, users override transformations by setting hparams.bottom, hparams.loss, hparams.top, and hparams.weights_fn.

See changes to common_hparams.py, modalities.py, modality.py. Other changes are maintenance.

PiperOrigin-RevId: 236407013
---
 .../data_generators/multi_problem.py          |    8 +-
 tensor2tensor/layers/common_hparams.py        |   31 +-
 .../layers/common_image_attention.py          |   35 +-
 .../layers/common_image_attention_test.py     |    1 +
 tensor2tensor/layers/modalities.py            | 2476 +++++++++--------
 tensor2tensor/layers/modalities_test.py       |   26 +-
 tensor2tensor/models/image_transformer.py     |   28 +-
 tensor2tensor/models/image_transformer_2d.py  |    5 +-
 tensor2tensor/models/mtf_transformer.py       |    9 +-
 tensor2tensor/models/mtf_transformer2.py      |    9 +-
 tensor2tensor/models/research/autoencoders.py |   18 +-
 tensor2tensor/models/research/cycle_gan.py    |    9 +-
 tensor2tensor/models/research/super_lm.py     |    8 +-
 .../models/research/transformer_symshard.py   |    9 +-
 .../models/research/transformer_vae.py        |    3 +-
 tensor2tensor/models/video/base.py            |   24 +-
 .../video/basic_deterministic_params.py       |   15 +-
 tensor2tensor/models/video/epva_params.py     |   12 +-
 tensor2tensor/models/video/next_frame_glow.py |   12 +-
 tensor2tensor/models/video/savp_params.py     |   17 +-
 tensor2tensor/models/video/sv2p_params.py     |   19 +-
 tensor2tensor/utils/hparam.py                 |    3 +-
 tensor2tensor/utils/metrics.py                |    8 +-
 tensor2tensor/utils/modality.py               |  121 -
 tensor2tensor/utils/t2t_model.py              |   52 +-
 tensor2tensor/utils/t2t_model_test.py         |    2 +
 26 files changed, 1523 insertions(+), 1437 deletions(-)
 delete mode 100644 tensor2tensor/utils/modality.py

diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 13ec711ae..fe534f050 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -438,8 +438,8 @@ def aggregate_task_losses(hparams,
     vocab_size += (-vocab_size) % hparams.vocab_divisor
   modality = problem_hparams.modality[feature_name]
   loss = hparams.loss.get(feature_name, modalities.get_loss(modality))
-  weights_fn = hparams.targets_weights_fn.get(
-      feature_name, modalities.get_targets_weights_fn(modality))
+  weights_fn = hparams.weights_fn.get(
+      feature_name, modalities.get_weights_fn(modality))
   # Primary task loss
   loss_num, loss_den = loss(
       logits, feature,
@@ -534,8 +534,8 @@ def aggregate_task_lm_losses(hparams,
     vocab_size += (-vocab_size) % hparams.vocab_divisor
   modality = problem_hparams.modality[feature_name]
   loss = hparams.loss.get(feature_name, modalities.get_loss(modality))
-  weights_fn = hparams.targets_weights_fn.get(
-      feature_name, modalities.get_targets_weights_fn(modality))
+  weights_fn = hparams.weights_fn.get(
+      feature_name, modalities.get_weights_fn(modality))
   loss_num = 0.
   loss_den = 0.
   for task in hparams.problem.task_list:
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 654ee4829..8d3f23eaf 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -141,7 +141,6 @@ def basic_params1():
       norm_type="layer",  # "batch", layer", "noam", "none".
       # epsilon parameter to normalization function
       norm_epsilon=1e-6,
-      symbol_modality_num_shards=1,
       # pad vocabularies so that this value divides the vocabulary size.
       vocab_divisor=1,
       # During training, we drop sequences whose inputs and targets are shorter
@@ -176,20 +175,26 @@ def basic_params1():
       # If True, run the model autoregressively instead of teacher-forcing
       # during eval
       eval_run_autoregressive=False,
-      # TODO(lukaszkaiser): these parameters should probably be set elsewhere.
-      # (SymbolModality) - If this flag is on, we try to share all of the input
-      # embeddings, the target embeddings and the softmax weights.
+      # (For features with symbol modality) If True, share all of the
+      # input embeddings, target embeddings, and softmax weights.
       shared_embedding_and_softmax_weights=False,
-      # (SymbolModality) - If this flag is on, we try to share the input
-      # embeddings and the target embeddings.
-      # You can also share the input embeddings with the target embeddings
-      # by using a problem_hparams that uses the same modality object for
-      # the input modality and target modality.
+      # (For features with symbol modality) If True, share the input embeddings
+      # and target embeddings.
       shared_embedding=False,
-      # Modalities used to map from features to a space compatible with
-      # chosen model architecture. It comprises key-value pairs of a feature
-      # name (str) and its modality type.
-      modality={},
+      # (For features with symbol modality) Number to shard embeddings by.
+      symbol_modality_num_shards=1,
+      # Feature transformations are optional dictionaries comprising key-value
+      # pairs of a feature name (str) and its transformation (function). If not
+      # specified, T2TModel applies a default transformation according to the
+      # feature's modality. Bottom is applicable to all features; loss, top, and
+      # weights_fn are only applicable to target features.
+      # TODO(trandustin): `name` is an optional hparam for legacy reasons,
+      # defining variable scope names. Remove this hparam in the future.
+      bottom={},
+      loss={},
+      name={},
+      top={},
+      weights_fn={},
       # The maximum length of "input" sequence.
       # Sequences longer than this value will be truncated. 0 or negative values
       # mean there is no maximum or truncation.
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 4f5ea7b67..ddc0ecd9d 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -603,30 +603,9 @@ def prepare_decoder(targets, hparams):
 
 def prepare_image(inputs, hparams, name=None):
   """Prepare image."""
-  inputs_shape = common_layers.shape_list(inputs)
-  batch = inputs_shape[0]
-  orig_rows = inputs_shape[1]
-  orig_cols = inputs_shape[2]
-  channels = hparams.num_channels
-
-  hidden_size = hparams.hidden_size
-  # TODO(trandustin): Check via modalities.ModalityType.IDENTITY and not str.
-  # The current implementation is to avoid circular imports, modalities ->
-  # discretization -> common_image_attention -> modalities.
-  if "targets" in hparams.modality:
-    target_modality_name = hparams.modality["targets"]
-    if not isinstance(target_modality_name, str):
-      target_modality_name = target_modality_name.__name__
-  else:
-    target_modality_name = None
-  if target_modality_name == "IdentityModality":
-    inputs = tf.to_int32(inputs)
-    x = get_channel_embeddings(channels, inputs, hidden_size, name=name)
-  else:
-    x = inputs
-  x = tf.reshape(x, [batch, orig_rows, orig_cols * channels, hidden_size])
-
-  return x
+  # TODO(trandustin): This is a legacy function. Remove its usage.
+  del hparams, name  # unused arg
+  return inputs
 
 
 def create_output(decoder_output, rows, cols, targets, hparams):
@@ -647,17 +626,19 @@ def create_output(decoder_output, rows, cols, targets, hparams):
     [batch, hparams.img_len, hparams.img_len, hparams.num_channels, 256].
     In the special case of predict mode, it is a Tensor of rank 5.
   """
+  del targets  # unused arg
   decoded_image = postprocess_image(decoder_output, rows, cols, hparams)
+  batch = common_layers.shape_list(decoded_image)[0]
   depth = common_layers.shape_list(decoded_image)[-1]
-  batch, height, width, channels = common_layers.shape_list(targets)
   likelihood = getattr(hparams, "likelihood", DistributionType.CAT)
   if hparams.mode == tf.estimator.ModeKeys.PREDICT:
     y = tf.reshape(decoded_image, [batch, -1, 1, 1, depth])
-    output = y[:, :height, :, :, :]
+    output = y[:, :rows, :, :, :]
   elif likelihood == DistributionType.CAT:
     # Unpack the cols dimension of the Categorical.
+    channels = hparams.num_channels
     output = tf.reshape(decoded_image,
-                        [batch, height, width, channels, depth])
+                        [batch, rows, cols // channels, channels, depth])
   else:
     output = decoded_image
   return output
diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index 1d22927b6..fae6806fa 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -94,6 +94,7 @@ def testCreateOutputTrainMode(self, likelihood, num_mixtures, depth):
     hparams = HParams(
         hidden_size=2,
         likelihood=likelihood,
+        num_channels=channels,
         mode=tf.estimator.ModeKeys.TRAIN,
         num_mixtures=num_mixtures,
     )
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 9a87aee1b..a8fbb29f6 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -13,20 +13,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Modalities define the bottom and top of the model (not the body)."""
+"""Modalities, which specify a feature's domain.
+
+T2TModel applies a default transformation to each feature according to its
+modality. Override them by specifying a model's
+hparams.{bottom,loss,name,top,weights_fn}.
+"""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import sys
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_audio
+from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 from tensor2tensor.layers import discretization
-from tensor2tensor.utils import modality
 
 import tensorflow as tf
 import tensorflow_probability as tfp
@@ -52,1266 +56,1440 @@ def is_pointwise(func):
   return func
 
 
-class SymbolModality(modality.Modality):
-  """Modality for sets of discrete symbols.
-
-  Input:
-    Embedding.
-
-  Output:
-    Linear transformation + softmax.
-  """
-
-  @staticmethod
-  def name(model_hparams, vocab_size):
-    return "symbol_modality_%d_%d" % (vocab_size, model_hparams.hidden_size)
-
-  targets_weights_fn = staticmethod(common_layers.weights_nonzero)
+def generic_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  """Compute loss numerator and denominator for one shard of output."""
+  del vocab_size  # unused arg
+  logits = top_out
+  logits = common_attention.maybe_upcast(logits, hparams=model_hparams)
+  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.0)
+  return common_layers.padded_cross_entropy(
+      logits,
+      targets,
+      model_hparams.label_smoothing,
+      cutoff=cutoff,
+      weights_fn=weights_fn)
 
-  @staticmethod
-  def _get_weights(model_hparams, vocab_size, hidden_dim=None):
-    """Create or get concatenated embedding or softmax variable.
-
-    Args:
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-      hidden_dim: dim of the variable. Defaults to _model_hparams' hidden_size
-
-    Returns:
-       a list of num_shards Tensors.
-    """
-    if hidden_dim is None:
-      hidden_dim = model_hparams.hidden_size
-    num_shards = model_hparams.symbol_modality_num_shards
-    shards = []
-    for i in range(num_shards):
-      shard_size = (vocab_size // num_shards) + (
-          1 if i < vocab_size % num_shards else 0)
-      var_name = "weights_%d" % i
-      shards.append(
-          tf.get_variable(
-              var_name, [shard_size, hidden_dim],
-              initializer=tf.random_normal_initializer(0.0, hidden_dim**-0.5)))
-    if num_shards == 1:
-      ret = shards[0]
-    else:
-      ret = tf.concat(shards, 0)
-    # Convert ret to tensor.
-    if not tf.executing_eagerly():
-      ret = common_layers.convert_gradient_to_tensor(ret)
-    return ret
-
-  @classmethod
-  def bottom_simple(cls, x, model_hparams, vocab_size, name, reuse):
-    with tf.variable_scope(name, reuse=reuse):
-      # Ensure the inputs are 3-D
-      if len(x.get_shape()) == 4:
-        x = tf.squeeze(x, axis=3)
-      while len(x.get_shape()) < 3:
-        x = tf.expand_dims(x, axis=-1)
-
-      var = cls._get_weights(model_hparams, vocab_size)
-      x = common_layers.dropout_no_scaling(
-          x, 1.0 - model_hparams.symbol_dropout)
-      ret = common_layers.gather(var, x)
-      if model_hparams.multiply_embedding_mode == "sqrt_depth":
-        ret *= model_hparams.hidden_size**0.5
-      ret *= tf.expand_dims(
-          common_layers.cast_like(tf.not_equal(x, 0), ret), -1)
-      return ret
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    if (model_hparams.shared_embedding_and_softmax_weights or
-        model_hparams.get("shared_embedding")):
-      return cls.bottom_simple(
-          x, model_hparams, vocab_size, "shared", reuse=None)
-    return cls.bottom_simple(
-        x, model_hparams, vocab_size, "input_emb", reuse=None)
-
-  @classmethod
-  def targets_bottom(cls, x, model_hparams, vocab_size):
-    if (model_hparams.shared_embedding_and_softmax_weights or
-        model_hparams.get("shared_embedding")):
-      try:
-        return cls.bottom_simple(
-            x, model_hparams, vocab_size, "shared", reuse=True)
-      except ValueError:
-        # perhaps there were no inputs, and this is a new variable.
-        return cls.bottom_simple(
-            x, model_hparams, vocab_size, "shared", reuse=None)
-    else:
-      return cls.bottom_simple(
-          x, model_hparams, vocab_size, "target_emb", reuse=None)
-
-  @classmethod
-  @is_pointwise
-  def top(cls, body_output, targets, model_hparams, vocab_size):
-    """Generate logits.
-
-    Args:
-      body_output: A Tensor with shape
-        [batch, p0, p1, model_hparams.hidden_size].
-      targets: Unused.
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].
-    """
-    del targets  # unused arg
-    if model_hparams.shared_embedding_and_softmax_weights:
-      scope_name = "shared"
-      reuse = tf.AUTO_REUSE
-    else:
-      scope_name = "softmax"
-      reuse = False
-    with tf.variable_scope(scope_name, reuse=reuse):
-      body_output_shape = common_layers.shape_list(body_output)
-      var = cls._get_weights(model_hparams, vocab_size, body_output_shape[-1])
-      if (model_hparams.factored_logits and
-          model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
-        # insert channels dimension
-        body_output = tf.expand_dims(body_output, 3)
-        return common_layers.FactoredTensor(body_output, var)
-      else:
-        body_output = tf.reshape(body_output, [-1, body_output_shape[-1]])
-        logits = tf.matmul(body_output, var, transpose_b=True)
-        return tf.reshape(logits,
-                          body_output_shape[:-1] + [1, vocab_size])
 
-
-class SymbolModalityWeightsAll(SymbolModality):
-  """SymbolModality for features that do not have 0-padding."""
-
-  targets_weights_fn = staticmethod(common_layers.weights_all)
-
-
-class SymbolModalityOneHot(SymbolModality):
-  """Simple SymbolModality with one hot as embeddings."""
-
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size):
-    return tf.one_hot(x, vocab_size)
-
-  @staticmethod
+def make_targets_bottom(bottom):
   def targets_bottom(x, model_hparams, vocab_size):
-    return tf.one_hot(x, vocab_size)
-
-  @staticmethod
-  @is_pointwise
-  def top(body_output, _, model_hparams, vocab_size):
-    return body_output
+    with tf.variable_scope("targets_bottom"):
+      return bottom(x, model_hparams, vocab_size)
+  return targets_bottom
 
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    del weights_fn  # unused arg
-    labels = tf.one_hot(targets, vocab_size)
-    loss = tf.nn.softmax_cross_entropy_with_logits(
-        logits=top_out, labels=labels)
-    return tf.reduce_mean(loss), tf.constant(1.0)
-
-
-class CTCSymbolModality(SymbolModality):
-  """SymbolModality that uses CTC loss."""
-
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Compute the CTC loss."""
-    logits = top_out
-    with tf.name_scope("ctc_loss", values=[logits, targets]):
-      # For CTC we assume targets are 1d, [batch, length, 1, 1] here.
-      targets_shape = targets.get_shape().as_list()
-      assert len(targets_shape) == 4
-      assert targets_shape[2] == 1
-      assert targets_shape[3] == 1
-      targets = tf.squeeze(targets, axis=[2, 3])
-      logits = tf.squeeze(logits, axis=[2, 3])
-      targets_mask = 1 - tf.to_int32(tf.equal(targets, 0))
-      targets_lengths = tf.reduce_sum(targets_mask, axis=1)
-      sparse_targets = tf.keras.backend.ctc_label_dense_to_sparse(
-          targets, targets_lengths)
-      xent = tf.nn.ctc_loss(
-          sparse_targets,
-          logits,
-          targets_lengths,
-          time_major=False,
-          preprocess_collapse_repeated=False,
-          ctc_merge_repeated=False)
-      weights = weights_fn(targets)
-      return tf.reduce_sum(xent), tf.reduce_sum(weights)
-
-
-class ImageModality(modality.Modality):
-  """Modality for images."""
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      if not tf.executing_eagerly():
-        tf.summary.image(
-            "inputs", common_layers.tpu_safe_image_summary(x), max_outputs=2)
-      return tf.to_float(x)
-
-  @classmethod
-  def targets_bottom(cls, x, model_hparams, vocab_size):
-    pixel_embedding_size = 64
-    inputs = x
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      if not tf.executing_eagerly():
-        tf.summary.image(
-            "targets_bottom",
-            common_layers.tpu_safe_image_summary(inputs),
-            max_outputs=1)
-      inputs_shape = common_layers.shape_list(inputs)
-      if len(inputs_shape) != 4:
-        raise ValueError("Assuming images given as int tensors in the format "
-                         "[batch, height, width, channels] (256 values).")
-      # We embed each of 256=vocab_size possible pixel values.
-      embedding_var = tf.get_variable(
-          "pixel_embedding",
-          [vocab_size, pixel_embedding_size])
-      hot_inputs = tf.one_hot(tf.to_int32(inputs), vocab_size)
-      hot_inputs = tf.reshape(hot_inputs, [-1, vocab_size])
-      embedded = tf.matmul(hot_inputs, embedding_var)
-      # Let's now merge all channels that were embedded into a single vector.
-      merged_size = pixel_embedding_size * inputs_shape[3]
-      embedded = tf.reshape(embedded, inputs_shape[:3] + [merged_size])
-      merged = tf.layers.dense(
-          embedded,
-          model_hparams.hidden_size,
-          name="merge_pixel_embedded_channels")
-      return merged
-
-  @staticmethod
-  def top(body_output, _, model_hparams, vocab_size):
-    # TODO(lukaszkaiser): is this a universal enough way to get channels?
-    num_channels = model_hparams.problem.num_channels
-    with tf.variable_scope("rgb_softmax"):
-      body_output_shape = common_layers.shape_list(body_output)
-      reshape_shape = body_output_shape[:3]
-      reshape_shape.extend([num_channels, vocab_size])
-      res = tf.layers.dense(body_output, vocab_size * num_channels)
-      res = tf.reshape(res, reshape_shape)
-      if not tf.get_variable_scope().reuse:
-        res_argmax = tf.argmax(res, axis=-1)
-        tf.summary.image(
-            "result",
-            common_layers.tpu_safe_image_summary(res_argmax),
-            max_outputs=1)
-      return res
-
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Compute loss numerator and denominator for one shard of output."""
-    logits = top_out
-    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.0)
-    return common_layers.padded_cross_entropy(
-        logits,
-        targets,
-        model_hparams.label_smoothing,
-        cutoff=cutoff,
-        weights_fn=weights_fn)
 
+def get_weights(model_hparams, vocab_size, hidden_dim=None):
+  """Create or get concatenated embedding or softmax variable.
 
-class ImageChannelCompressModality(modality.Modality):
-  """Modality for images using channel compression for generation."""
-
-  @staticmethod
-  def bottom_compress(inputs, model_hparams, name="bottom"):
-    """Compresses channel-wise input pixels into whole pixel representions.
-
-    Perform conversion of RGB pixel values to a real number in the range -1 to
-    1. This combines pixel channels to form a representation of shape
-    [img_len, img_len].
-
-    Args:
-      inputs: Tensor representing RGB pixel intensities as integers, of shape
-        [batch, img_len, img_len, channels].
-      model_hparams: tf.HParams, model hyperparmeters.
-      name: string, scope.
-
-    Returns:
-      body_input: Tensor of shape
-        [batch, img_len, img_len, model_hparams.hidden_size].
-    """
-    num_channels = 3
-    with tf.variable_scope(name):
-      inputs = tf.to_float(inputs)
-      hp = model_hparams
-      if hp.mode != tf.estimator.ModeKeys.PREDICT:
-        tf.summary.image(
-            "inputs",
-            common_layers.tpu_safe_image_summary(inputs),
-            max_outputs=2)
-      inputs = common_layers.convert_rgb_to_symmetric_real(inputs)
-
-      # Reshape inputs to apply convolutions across [img_len, img_len*channels].
-      inputs_shape = common_layers.shape_list(inputs)
-      inputs = tf.reshape(
-          inputs, [-1, inputs_shape[1], inputs_shape[2] * inputs_shape[3], 1])
-
-      # Compress RGB intensities for each pixel using a convolution.
-      outputs = tf.layers.conv2d(
-          inputs,
-          model_hparams.hidden_size,
-          kernel_size=(1, num_channels),
-          padding="VALID",
-          strides=(1, num_channels),
-          activation=tf.nn.relu,
-          name="conv_input")
-      return outputs
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    return cls.bottom_compress(x, model_hparams, "input_bottom")
-
-  @classmethod
-  def targets_bottom(cls, x, model_hparams, vocab_size):
-    return cls.bottom_compress(x, model_hparams, "output_bottom")
-
-  @classmethod
-  def top(cls, body_output, _, model_hparams, vocab_size):
-    """Transforms body output to return logits.
-
-    Args:
-      body_output: Tensor of shape [batch, img_len, img_len, depth].
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      Tensor of shape [batch, img_len, img_len, channels, vocab_size].
-    """
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      hidden_size = model_hparams.hidden_size
-      img_len = model_hparams.img_len
-      channels = 3  # RGB
-      batch = common_layers.shape_list(body_output)[0]
-      x = tf.layers.conv2d(
-          body_output,
-          hidden_size * channels,
-          kernel_size=(1, 1),
-          strides=(1, 1),
-          padding="VALID",
-          activation=tf.nn.relu,
-          name="decompress_conv")
-      x = tf.reshape(x, [batch, img_len, img_len * channels, hidden_size])
-      x = common_layers.layer_preprocess(x, model_hparams)
-      x = tf.layers.dense(x,
-                          vocab_size,
-                          use_bias=True,
-                          activation=None,
-                          name="output_conv")
-      x = tf.reshape(
-          x, [batch, img_len, img_len, channels, vocab_size])
-      return x
-
-
-class ImageChannelBottomIdentityModality(ImageChannelCompressModality):
-
-  @staticmethod
-  def top(body_output, _, model_hparams, vocab_size):
-    return body_output
-
-
-class ImageChannelEmbeddingsBottom(modality.Modality):
-  """Modality for images using channel compression for generation."""
+  Args:
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
+    hidden_dim: dim of the variable. Defaults to _model_hparams' hidden_size
 
-  @staticmethod
-  def get_channel_embeddings(io_depth,
-                             targets,
-                             hidden_size,
-                             name="channel"):
-    """Get separate embedding for each of the channels."""
-    targets_split = tf.split(targets, io_depth, axis=3)
-    rgb_embedding_var = tf.get_variable("rgb_target_emb_%s" % name,
-                                        [256 * io_depth, hidden_size])
-    rgb_embedding_var = tf.identity(rgb_embedding_var)
-    rgb_embedding_var *= float(hidden_size)**0.5
-    channel_target_embs = []
-    for i in range(io_depth):
-      # Adding the channel offsets to get the right embedding since the
-      # embedding tensor has shape 256 * io_depth, hidden_size
-      target_ids = tf.squeeze(targets_split[i], axis=3) + i * 256
-      target_embs = common_layers.gather(rgb_embedding_var, target_ids)
-      channel_target_embs.append(target_embs)
-
-    return tf.concat(channel_target_embs, axis=-1)
+  Returns:
+     a list of num_shards Tensors.
+  """
+  if hidden_dim is None:
+    hidden_dim = model_hparams.hidden_size
+  num_shards = model_hparams.symbol_modality_num_shards
+  shards = []
+  for i in range(num_shards):
+    shard_size = (vocab_size // num_shards) + (
+        1 if i < vocab_size % num_shards else 0)
+    var_name = "weights_%d" % i
+    shards.append(
+        tf.get_variable(
+            var_name, [shard_size, hidden_dim],
+            initializer=tf.random_normal_initializer(0.0, hidden_dim**-0.5)))
+  if num_shards == 1:
+    ret = shards[0]
+  else:
+    ret = tf.concat(shards, 0)
+  # Convert ret to tensor.
+  if not tf.executing_eagerly():
+    ret = common_layers.convert_gradient_to_tensor(ret)
+  return ret
+
+
+def _symbol_bottom_simple(x, model_hparams, vocab_size, name, reuse):
+  """Bottom transformation for symbols."""
+  with tf.variable_scope(name, reuse=reuse):
+    # Ensure the inputs are 3-D
+    if len(x.get_shape()) == 4:
+      x = tf.squeeze(x, axis=3)
+    while len(x.get_shape()) < 3:
+      x = tf.expand_dims(x, axis=-1)
+
+    var = get_weights(model_hparams, vocab_size)
+    x = common_layers.dropout_no_scaling(
+        x, 1.0 - model_hparams.symbol_dropout)
+    ret = common_layers.gather(var, x)
+    if model_hparams.multiply_embedding_mode == "sqrt_depth":
+      ret *= model_hparams.hidden_size**0.5
+    ret *= tf.expand_dims(
+        common_layers.cast_like(tf.not_equal(x, 0), ret), -1)
+    return ret
 
-  @staticmethod
-  def targets_bottom(x, model_hparams, vocab_size):
-    inputs = x
-    io_depth = model_hparams.num_channels
-    tshape = common_layers.shape_list(inputs)
-    hidden_size = model_hparams.hidden_size
-    target_embeddings = ImageChannelEmbeddingsBottom.get_channel_embeddings(
-        io_depth, inputs, hidden_size, "input_bottom")
-    return tf.reshape(target_embeddings,
-                      [tshape[0], tshape[1], tshape[2] * io_depth, hidden_size])
-
-  @classmethod
-  def top(cls, body_output, _, model_hparams, vocab_size):
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      img_len = model_hparams.img_len
-      channels = model_hparams.num_channels
-      x = tf.layers.dense(
-          body_output, 256, use_bias=True, activation=None, name="output_conv")
-      x = tf.reshape(x,
-                     [-1, img_len, img_len, channels, vocab_size])
-      return x
-
-
-class AudioModality(modality.Modality):
-  """Performs strided conv compressions for audio data."""
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    """Transform input from data space to model space.
-
-    Args:
-      x: A Tensor with shape [batch, ...]
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      body_input: A Tensor with shape [batch, ?, ?,
-        model_hparams.hidden_size].
-    """
-    inputs = x
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      # TODO(aidangomez): Will need to sort out a better audio pipeline
-      def xnet_resblock(x, filters, res_relu, name):
-        """Xception block."""
-        with tf.variable_scope(name):
-          # Typically audio samples are >100k samples in length and have a width
-          # of 2 or 4. Mono audio has a single channel while stereo has 2.
-          y = common_layers.separable_conv_block(
-              x,
-              filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
-              first_relu=True,
-              padding="SAME",
-              force2d=True,
-              name="sep_conv_block")
-          y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 2))
-          return y + common_layers.conv_block(
-              x,
-              filters, [((1, 1), (1, 1))],
-              padding="SAME",
-              strides=(2, 2),
-              first_relu=res_relu,
-              force2d=True,
-              name="res_conv0")
-
-      x = tf.to_float(inputs) / 255.
-      x.set_shape([None, None, None, 1])
-      for i in range(model_hparams.audio_compression):
-        x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
-      return xnet_resblock(x,
-                           model_hparams.hidden_size,
-                           False,
-                           "compress_block_final")
-
-
-class AudioSpectralModality(modality.Modality):
-  """Performs strided conv compressions for audio spectral data."""
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    """Transform input from data space to model space.
-
-    Args:
-      x: A Tensor with shape [batch, ...]
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      body_input: A Tensor with shape [batch, ?, ?,
-        model_hparams.hidden_size].
-    """
-    inputs = x
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      # TODO(aidangomez): Will need to sort out a better audio pipeline
-      def xnet_resblock(x, filters, res_relu, name):
-        """Xception-like block."""
-        with tf.variable_scope(name):
-          # We only stride along the length dimension to preserve the spectral
-          # bins (which are tiny in dimensionality relative to length)
-          y = common_layers.separable_conv_block(
-              x,
-              filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
-              first_relu=True,
-              padding="SAME",
-              force2d=True,
-              name="sep_conv_block")
-          y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 1))
-          return y + common_layers.conv_block(
-              x,
-              filters, [((1, 1), (1, 1))],
-              padding="SAME",
-              strides=(2, 1),
-              first_relu=res_relu,
-              force2d=True,
-              name="res_conv0")
-
-      # Bitcast back from int32
-      x = tf.bitcast(inputs, tf.float32)
-      x.set_shape([None, None, None, 1])
-      for i in range(model_hparams.audio_compression):
-        x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
-      return xnet_resblock(x,
-                           model_hparams.hidden_size,
-                           False,
-                           "compress_block_final")
-
-
-class SpeechRecognitionModality(modality.Modality):
-  """Common ASR filterbank processing."""
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    """Use batchnorm instead of CMVN and shorten the stft with strided convs.
-
-    Args:
-      x: float32 tensor with shape [batch_size, len, 1, freqs * channels]
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      float32 tensor with shape [batch_size, shorter_len, 1, hidden_size]
-    """
-    inputs = x
-    p = model_hparams
-
-    num_mel_bins = p.audio_num_mel_bins
-    num_channels = 3 if p.audio_add_delta_deltas else 1
-
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      if p.audio_preproc_in_bottom:
-        # Compute filterbanks
-        with tf.variable_scope("fbanks"):
-          waveforms = tf.squeeze(inputs, [2, 3])
-          mel_fbanks = common_audio.compute_mel_filterbank_features(
-              waveforms,
-              sample_rate=p.audio_sample_rate,
-              dither=p.audio_dither,
-              preemphasis=p.audio_preemphasis,
-              frame_length=p.audio_frame_length,
-              frame_step=p.audio_frame_step,
-              lower_edge_hertz=p.audio_lower_edge_hertz,
-              upper_edge_hertz=p.audio_upper_edge_hertz,
-              num_mel_bins=p.audio_num_mel_bins,
-              apply_mask=True)
-          if p.audio_add_delta_deltas:
-            mel_fbanks = common_audio.add_delta_deltas(mel_fbanks)
-          x = tf.reshape(mel_fbanks,
-                         common_layers.shape_list(mel_fbanks)[:2] +
-                         [num_mel_bins, num_channels])
-
-          nonpadding_mask = 1. - common_attention.embedding_to_padding(x)
-          num_of_nonpadding_elements = tf.reduce_sum(
-              nonpadding_mask) * num_mel_bins * num_channels
-
-          # This replaces CMVN estimation on data
-          var_epsilon = 1e-09
-          mean = tf.reduce_sum(
-              x, axis=[1], keepdims=True) / num_of_nonpadding_elements
-          variance = (num_of_nonpadding_elements * mean**2. -
-                      2. * mean * tf.reduce_sum(x, axis=[1], keepdims=True) +
-                      tf.reduce_sum(x**2, axis=[1], keepdims=True)
-                     ) / num_of_nonpadding_elements
-          x = (x - mean) * tf.rsqrt(variance + var_epsilon) * tf.expand_dims(
-              nonpadding_mask, -1)
-      else:
-        x = inputs
-
-      # The convention is that the models are flattened along the spatial,
-      # dimensions, thus the speech preprocessor treats frequencies and
-      # channels as image colors (last axis)
-      x.set_shape([None, None, num_mel_bins, num_channels])
-
-      # TODO(chorowski): how to specify bottom's hparams and avoid hardcoding?
-      x = tf.pad(x, [[0, 0], [0, 8], [0, 0], [0, 0]])
-      for _ in range(2):
-        x = tf.layers.conv2d(
-            x, 128, (3, 3), (2, 2), use_bias=False)
-        x = common_layers.layer_norm(x)
-        x = tf.nn.relu(x)
-
-      xshape = common_layers.shape_list(x)
-      # apply a conv that will remove all frequencies and at the same time
-      # project the output into desired hidden_size
-      x = tf.pad(x, [[0, 0], [0, 2], [0, 0], [0, 0]])
-      x = tf.layers.conv2d(x, p.hidden_size, (3, xshape[2]), use_bias=False)
-
-      assert common_layers.shape_list(x)[2] == 1
-      x = common_layers.layer_norm(x)
-      x = tf.nn.relu(x)
-    return x
 
+def symbol_bottom(x, model_hparams, vocab_size):
+  if (model_hparams.shared_embedding_and_softmax_weights or
+      model_hparams.get("shared_embedding")):
+    return _symbol_bottom_simple(
+        x, model_hparams, vocab_size, "shared", reuse=None)
+  return _symbol_bottom_simple(
+      x, model_hparams, vocab_size, "input_emb", reuse=None)
+
+
+def symbol_targets_bottom(x, model_hparams, vocab_size):
+  """Bottom transformation for target symbols."""
+  if (model_hparams.shared_embedding_and_softmax_weights or
+      model_hparams.get("shared_embedding")):
+    try:
+      return _symbol_bottom_simple(
+          x, model_hparams, vocab_size, "shared", reuse=True)
+    except ValueError:
+      # perhaps there were no inputs, and this is a new variable.
+      return _symbol_bottom_simple(
+          x, model_hparams, vocab_size, "shared", reuse=None)
+  else:
+    return _symbol_bottom_simple(
+        x, model_hparams, vocab_size, "target_emb", reuse=None)
 
-class VideoModality(modality.Modality):
-  """Modality for videos, i.e., time-sequences of frames."""
 
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size):
-    common_video.gif_summary("inputs", x, max_outputs=1)
-    x = common_layers.standardize_images(x)
-    return x
+@is_pointwise
+def symbol_top(body_output, targets, model_hparams, vocab_size):
+  """Generate logits.
 
-  @staticmethod
-  def targets_bottom(x, model_hparams, vocab_size):
-    common_video.gif_summary("targets", x, max_outputs=1)
-    x = common_layers.standardize_images(x)
-    return x
-
-  @staticmethod
-  def top(body_output, targets, model_hparams, vocab_size):
-    num_channels = model_hparams.problem.num_channels
-    shape = common_layers.shape_list(body_output)
-    reshape_shape = shape[:-1] + [num_channels, vocab_size]
-    res = tf.reshape(body_output, reshape_shape)
-    # Calculate argmax so as to have a summary with the produced images.
-    x = tf.argmax(tf.reshape(res, [-1, vocab_size]), axis=-1)
-    x = tf.reshape(x, shape[:-1] + [num_channels])
-    common_video.gif_summary("results", x, max_outputs=1)
-    return res
+  Args:
+    body_output: A Tensor with shape
+      [batch, p0, p1, model_hparams.hidden_size].
+    targets: Unused.
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
 
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Compute loss numerator and denominator for one shard of output."""
-    logits = top_out
-    logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
-    targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
-    return common_layers.padded_cross_entropy(
+  Returns:
+    logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].
+  """
+  del targets  # unused arg
+  if model_hparams.shared_embedding_and_softmax_weights:
+    scope_name = "shared"
+    reuse = tf.AUTO_REUSE
+  else:
+    scope_name = "softmax"
+    reuse = False
+  with tf.variable_scope(scope_name, reuse=reuse):
+    body_output_shape = common_layers.shape_list(body_output)
+    var = get_weights(model_hparams, vocab_size, body_output_shape[-1])
+    if (model_hparams.factored_logits and
+        model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
+      # insert channels dimension
+      body_output = tf.expand_dims(body_output, 3)
+      return common_layers.FactoredTensor(body_output, var)
+    else:
+      body_output = tf.reshape(body_output, [-1, body_output_shape[-1]])
+      logits = tf.matmul(body_output, var, transpose_b=True)
+      return tf.reshape(logits,
+                        body_output_shape[:-1] + [1, vocab_size])
+
+
+def symbol_one_hot_bottom(x, model_hparams, vocab_size):
+  del model_hparams  # unused arg
+  return tf.one_hot(x, vocab_size)
+
+
+@is_pointwise
+def symbol_one_hot_top(body_output, targets, model_hparams, vocab_size):
+  del targets, model_hparams, vocab_size  # unused arg
+  return body_output
+
+
+def symbol_one_hot_loss(top_out,
+                        targets,
+                        model_hparams,
+                        vocab_size,
+                        weights_fn):
+  del model_hparams, weights_fn  # unused arg
+  labels = tf.one_hot(targets, vocab_size)
+  loss = tf.nn.softmax_cross_entropy_with_logits(
+      logits=top_out, labels=labels)
+  return tf.reduce_mean(loss), tf.constant(1.0)
+
+
+def ctc_symbol_loss(top_out, targets, model_hparams, vocab_size, weight_fn):
+  """Compute the CTC loss."""
+  del model_hparams, vocab_size  # unused arg
+  logits = top_out
+  with tf.name_scope("ctc_loss", values=[logits, targets]):
+    # For CTC we assume targets are 1d, [batch, length, 1, 1] here.
+    targets_shape = targets.get_shape().as_list()
+    assert len(targets_shape) == 4
+    assert targets_shape[2] == 1
+    assert targets_shape[3] == 1
+    targets = tf.squeeze(targets, axis=[2, 3])
+    logits = tf.squeeze(logits, axis=[2, 3])
+    targets_mask = 1 - tf.to_int32(tf.equal(targets, 0))
+    targets_lengths = tf.reduce_sum(targets_mask, axis=1)
+    sparse_targets = tf.keras.backend.ctc_label_dense_to_sparse(
+        targets, targets_lengths)
+    xent = tf.nn.ctc_loss(
+        sparse_targets,
         logits,
-        targets,
-        model_hparams.label_smoothing,
-        cutoff=cutoff,
-        weights_fn=weights_fn)
-
-
-class VideoModalityBitwise(VideoModality):
-  """Video Modality where bottom embeds pixels bitwise."""
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    pixel_embedding_size = 64
-    inputs = x
-    with tf.variable_scope(cls.name(model_hparams, vocab_size),
-                           reuse=tf.AUTO_REUSE):
-      common_layers.summarize_video(inputs, "bottom")
-      # Embed bitwise.
-      assert vocab_size == 256
-      embedded = discretization.int_to_bit_embed(inputs, 8,
-                                                 pixel_embedding_size)
-      # Project.
-      return tf.layers.dense(
-          embedded,
-          model_hparams.hidden_size,
-          name="merge_pixel_embedded_frames")
-
-  @classmethod
-  def targets_bottom(cls, x, model_hparams, vocab_size):  # pylint: disable=arguments-differ
-    pixel_embedding_size = 64
-    inputs = x
-    with tf.variable_scope(cls.name(model_hparams, vocab_size),
-                           reuse=tf.AUTO_REUSE):
-      common_layers.summarize_video(inputs, "targets_bottom")
-      # Embed bitwise.
-      assert vocab_size == 256
-      embedded = discretization.int_to_bit_embed(inputs, 8,
-                                                 pixel_embedding_size)
-      # Transpose and project.
-      transposed = common_layers.time_to_channels(embedded)
-      return tf.layers.dense(
-          transposed,
-          model_hparams.hidden_size,
-          name="merge_pixel_embedded_frames")
-
-
-class VideoModalityPixelNoise(VideoModality):
-  """Video modality that introduces pixel noise on input during training."""
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    input_noise = getattr(model_hparams, "video_modality_input_noise", 0.25)
-    inputs = x
-    if model_hparams.mode == tf.estimator.ModeKeys.TRAIN:
-      background = tfp.distributions.percentile(inputs, 50., axis=[0, 1, 2, 3])
-      input_shape = common_layers.shape_list(inputs)
-      input_size = tf.reduce_prod(input_shape[:-1])
-      input_mask = tf.multinomial(
-          tf.log([[input_noise, 1.-input_noise]]), input_size)
-      input_mask = tf.reshape(tf.cast(input_mask, tf.int32),
-                              input_shape[:-1]+[1])
-      inputs = inputs * input_mask + background * (1 - input_mask)
-    return super(VideoModalityPixelNoise, cls).bottom(
-        inputs, model_hparams, vocab_size)
-
-
-class VideoModalityL1(VideoModality):
-  """Video modality that predicts a scalar per channel with an L1 loss."""
-
-  @staticmethod
-  def top(body_output, _, model_hparams, vocab_size):
-    num_channels = model_hparams.problem.num_channels
-    num_frames = model_hparams.video_num_target_frames
-    with tf.variable_scope("rgb"):
-      body_output_shape = common_layers.shape_list(body_output)
-      res = tf.layers.dense(body_output, num_channels * num_frames, name="cast")
-      res = tf.reshape(res, body_output_shape[:3] + [num_channels, num_frames])
-      res = tf.transpose(res, [0, 4, 1, 2, 3])  # Move frames next to batch.
-      if not tf.get_variable_scope().reuse:
-        res_argmax = res[:, -1, :, :, :]
-        tf.summary.image(
-            "result",
-            common_layers.tpu_safe_image_summary(res_argmax),
-            max_outputs=1)
-      return tf.expand_dims(res, axis=-1)  # Add an axis like in perplexity.
-
-  @staticmethod
-  def internal_loss(logits, targets, model_hparams):
-    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
-    return tf.nn.relu(tf.abs(logits - targets) - cutoff)
-
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Compute loss numerator and denominator for one shard of output."""
-    logits = top_out
-    logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:-1])
-    targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-    weights = weights_fn(targets)
-    # Shift targets by 0.5 so later just casting to int gives the prediction.
-    # So for int targets, say 0 and 7, we actually train to predict 0.5 and 7.5.
-    # Later (in merics or infer) this is cast to int anyway. Also, we have no
-    # loss beyond cutoff = 0.2 as these are already correct predictions.
-    targets = tf.to_float(targets) + 0.5
-    loss = VideoModalityL1.internal_loss(logits, targets, model_hparams)
-    return tf.reduce_sum(loss * weights), tf.reduce_sum(weights)
-
-
-class VideoModalityL2(VideoModalityL1):
-  """Modality for videos with L2 loss."""
-
-  @staticmethod
-  def internal_loss(logits, targets, model_hparams):
-    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
-    return tf.nn.relu(
-        tf.squared_difference(logits, targets) - cutoff * cutoff)
-
-
-class VideoModalityL2Raw(VideoModalityL2):
-  """Modality with L2 loss and raw input (sequences of frames)."""
+        targets_lengths,
+        time_major=False,
+        preprocess_collapse_repeated=False,
+        ctc_merge_repeated=False)
+    weights = weight_fn(targets)
+    return tf.reduce_sum(xent), tf.reduce_sum(weights)
 
-  @staticmethod
-  def convert_rgb_to_real(prediction, targets):
-    """Convert prediction and target from rgb to real."""
-    prediction = tf.squeeze(prediction, axis=-1)
-    prediction = common_layers.convert_rgb_to_real(prediction)
-    targets = common_layers.convert_rgb_to_real(targets)
-    return prediction, targets
-
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size):
-    common_video.gif_summary("inputs", x)
-    return common_layers.convert_rgb_to_real(x)
 
-  @staticmethod
-  def targets_bottom(x, model_hparams, vocab_size):  # pylint: disable=arguments-differ
-    common_video.gif_summary("targets_bottom", x)
-    return common_layers.convert_rgb_to_real(x)
-
-  @staticmethod
-  def top(body_output, _, model_hparams, vocab_size):
-    frames = body_output
-    if isinstance(body_output, list):
-      frames = tf.stack(body_output, axis=1)
-    rgb_frames = common_layers.convert_real_to_rgb(frames)
-    common_video.gif_summary("body_output", rgb_frames)
-    return tf.expand_dims(rgb_frames, axis=-1)
+def image_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  with tf.variable_scope("image_modality"):
+    if not tf.executing_eagerly():
+      tf.summary.image(
+          "inputs", common_layers.tpu_safe_image_summary(x), max_outputs=2)
+    return tf.to_float(x)
 
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
-    del weights_fn  # unused arg
-    prediction, groundtruth = cls.convert_rgb_to_real(top_out, targets)
-    loss = tf.losses.mean_squared_error(prediction, groundtruth)
-    return loss, tf.constant(1.0)
 
+def image_targets_bottom(x, model_hparams, vocab_size):
+  """Bottom transformation for target images."""
+  pixel_embedding_size = 64
+  inputs = x
+  with tf.variable_scope("image_modality"):
+    if not tf.executing_eagerly():
+      tf.summary.image(
+          "targets_bottom",
+          common_layers.tpu_safe_image_summary(inputs),
+          max_outputs=1)
+    inputs_shape = common_layers.shape_list(inputs)
+    if len(inputs_shape) != 4:
+      raise ValueError("Assuming images given as int tensors in the format "
+                       "[batch, height, width, channels] (256 values).")
+    # We embed each of 256=vocab_size possible pixel values.
+    embedding_var = tf.get_variable(
+        "pixel_embedding",
+        [vocab_size, pixel_embedding_size])
+    hot_inputs = tf.one_hot(tf.to_int32(inputs), vocab_size)
+    hot_inputs = tf.reshape(hot_inputs, [-1, vocab_size])
+    embedded = tf.matmul(hot_inputs, embedding_var)
+    # Let's now merge all channels that were embedded into a single vector.
+    merged_size = pixel_embedding_size * inputs_shape[3]
+    embedded = tf.reshape(embedded, inputs_shape[:3] + [merged_size])
+    merged = tf.layers.dense(
+        embedded,
+        model_hparams.hidden_size,
+        name="merge_pixel_embedded_channels")
+    return merged
+
+
+def image_top(body_output, targets, model_hparams, vocab_size):
+  """Top transformation for images."""
+  del targets  # unused arg
+  # TODO(lukaszkaiser): is this a universal enough way to get channels?
+  num_channels = model_hparams.problem.num_channels
+  with tf.variable_scope("rgb_softmax"):
+    body_output_shape = common_layers.shape_list(body_output)
+    reshape_shape = body_output_shape[:3]
+    reshape_shape.extend([num_channels, vocab_size])
+    res = tf.layers.dense(body_output, vocab_size * num_channels)
+    res = tf.reshape(res, reshape_shape)
+    if not tf.get_variable_scope().reuse:
+      res_argmax = tf.argmax(res, axis=-1)
+      tf.summary.image(
+          "result",
+          common_layers.tpu_safe_image_summary(res_argmax),
+          max_outputs=1)
+    return res
 
-class VideoModalityL1Raw(VideoModalityL2Raw):
-  """Modality with L1 loss and raw input (sequences of frames)."""
 
-  @classmethod
-  def loss(cls, top_out, targets, model_hparams, vocab_size, weights_fn):
-    prediction, groundtruth = cls.convert_rgb_to_real(top_out, targets)
-    loss = tf.losses.absolute_difference(prediction, groundtruth)
-    return loss, tf.constant(1.0)
+def _image_channel_compress_bottom(inputs, model_hparams, name="bottom"):
+  """Compresses channel-wise input pixels into whole pixel representions.
 
+  Perform conversion of RGB pixel values to a real number in the range -1 to
+  1. This combines pixel channels to form a representation of shape
+  [img_len, img_len].
 
-class ClassLabelModality(modality.Modality):
-  """Used for label data."""
+  Args:
+    inputs: Tensor representing RGB pixel intensities as integers, of shape
+      [batch, img_len, img_len, channels].
+    model_hparams: tf.HParams, model hyperparmeters.
+    name: string, scope.
 
-  @staticmethod
-  def name(model_hparams, vocab_size):
-    return "class_label_modality_%d_%d" % (vocab_size,
-                                           model_hparams.hidden_size)
-
-  @classmethod
-  def bottom(cls, x, model_hparams, vocab_size):
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      multiplier = 1.0
-      if model_hparams.multiply_embedding_mode == "sqrt_depth":
-        multiplier = model_hparams.hidden_size**0.5
-      return common_layers.embedding(x,
-                                     vocab_size,
-                                     model_hparams.hidden_size,
-                                     multiplier=multiplier)
-
-  @classmethod
-  def targets_bottom(cls, x, model_hparams, vocab_size):
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      return tf.zeros([common_layers.shape_list(x)[0],
-                       1,
-                       1,
-                       model_hparams.hidden_size])
-
-  @classmethod
-  def top(cls, body_output, _, model_hparams, vocab_size):
-    """Transform inputs from model space to target space.
-
-    Average over inner dims and a linear layer to logits.
-
-    Args:
-      body_output: A Tensor with shape [batch, ?, ?, body_output_size].
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      a Tensors, each with shape [batch_size, 1, 1, 1, vocab_size]
-    """
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      x = body_output
-      x = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
-      res = tf.layers.dense(x, vocab_size)
-      return tf.expand_dims(res, 3)
-
-
-class VideoModalityIdentity(VideoModality):
-  """Video Modality where top and bottom is an identity function."""
+  Returns:
+    body_input: Tensor of shape
+      [batch, img_len, img_len, model_hparams.hidden_size].
+  """
+  num_channels = 3
+  with tf.variable_scope(name):
+    inputs = tf.to_float(inputs)
+    hp = model_hparams
+    if hp.mode != tf.estimator.ModeKeys.PREDICT:
+      tf.summary.image(
+          "inputs",
+          common_layers.tpu_safe_image_summary(inputs),
+          max_outputs=2)
+    inputs = common_layers.convert_rgb_to_symmetric_real(inputs)
+
+    # Reshape inputs to apply convolutions across [img_len, img_len*channels].
+    inputs_shape = common_layers.shape_list(inputs)
+    inputs = tf.reshape(
+        inputs, [-1, inputs_shape[1], inputs_shape[2] * inputs_shape[3], 1])
+
+    # Compress RGB intensities for each pixel using a convolution.
+    outputs = tf.layers.conv2d(
+        inputs,
+        model_hparams.hidden_size,
+        kernel_size=(1, num_channels),
+        padding="VALID",
+        strides=(1, num_channels),
+        activation=tf.nn.relu,
+        name="conv_input")
+    return outputs
+
+
+def image_channel_compress_bottom(x, model_hparams, vocab_size):
+  del vocab_size  # unused arg
+  return _image_channel_compress_bottom(x, model_hparams, "input_bottom")
+
+
+def image_channel_compress_targets_bottom(x, model_hparams, vocab_size):
+  del vocab_size  # unused arg
+  return _image_channel_compress_bottom(x, model_hparams, "output_bottom")
+
+
+def image_channel_compress_top(body_output, targets, model_hparams, vocab_size):
+  """Transforms body output to return logits.
 
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size):
-    common_video.gif_summary("inputs", x, max_outputs=1)
-    return x
+  Args:
+    body_output: Tensor of shape [batch, img_len, img_len, depth].
+    targets:
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
 
-  @staticmethod
-  def targets_bottom(x, model_hparams, vocab_size):
-    common_video.gif_summary("targets", x, max_outputs=1)
+  Returns:
+    Tensor of shape [batch, img_len, img_len, channels, vocab_size].
+  """
+  del targets  # unused arg
+  with tf.variable_scope("image_channel_compress_modality"):
+    hidden_size = model_hparams.hidden_size
+    img_len = model_hparams.img_len
+    channels = 3  # RGB
+    batch = common_layers.shape_list(body_output)[0]
+    x = tf.layers.conv2d(
+        body_output,
+        hidden_size * channels,
+        kernel_size=(1, 1),
+        strides=(1, 1),
+        padding="VALID",
+        activation=tf.nn.relu,
+        name="decompress_conv")
+    x = tf.reshape(x, [batch, img_len, img_len * channels, hidden_size])
+    x = common_layers.layer_preprocess(x, model_hparams)
+    x = tf.layers.dense(x,
+                        vocab_size,
+                        use_bias=True,
+                        activation=None,
+                        name="output_conv")
+    x = tf.reshape(
+        x, [batch, img_len, img_len, channels, vocab_size])
     return x
 
-  @staticmethod
-  def top(body_output, targets, model_hparams, vocab_size):
-    return body_output
-
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Compute loss numerator and denominator for one shard of output."""
-    # TODO(nikip): Try L2 loss
-    logits = top_out
-    logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
-    targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-    cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
-    return common_layers.padded_cross_entropy(
-        logits,
-        targets,
-        model_hparams.label_smoothing,
-        cutoff=cutoff,
-        weights_fn=weights_fn)
-
-
-class MultiLabelModality(ClassLabelModality):
-  """Used for multi label task."""
-
-  targets_weights_fn = staticmethod(common_layers.weights_nonzero)
-
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Average loss over the labels."""
-    logits = top_out
-    num_labels = tf.shape(targets)[1]
-    logits = tf.tile(logits, [1, num_labels, 1, 1, 1])
-
-    xent, weights = common_layers.padded_cross_entropy(
-        logits,
-        targets,
-        model_hparams.label_smoothing,
-        weights_fn=weights_fn,
-        reduce_sum=False,
-    )
-    xent = tf.squeeze(xent, [2, 3])
-    weights = tf.squeeze(weights, [2, 3])
-    # average loss over all labels
-    loss = tf.reduce_sum(xent, axis=1)
-    weights = tf.reduce_sum(weights, axis=1)
-    loss /= (weights + 1e-8)
-    weights = tf.to_float(tf.greater(weights, 0.))
-
-    return tf.reduce_sum(loss*weights), tf.reduce_sum(weights)
-
 
-class OneHotClassLabelModality(ClassLabelModality):
-  """Used for one-hot encoded class labels."""
-
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Apply softmax cross-entropy between outputs and targets.
-
-    Args:
-      top_out: logits Tensor with shape [batch, ?, ?, num_classes]
-      targets: one-hot encoding Tensor with shape [batch, ?, ?, num_classes]
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-      weights_fn: Function mapping targets to weights.
-
-    Returns:
-      loss_scale (cross-entropy), loss_denom
-    """
-    loss_scale = tf.losses.softmax_cross_entropy(
-        onehot_labels=targets, logits=top_out)
-    weights = weights_fn(targets)
-    loss_denom = tf.reduce_sum(weights)
-    return loss_scale, loss_denom
-
-
-class IdentityModality(modality.Modality):
-  """Does nothing."""
-
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size):
-    return tf.to_float(x)
-
-  @staticmethod
-  def top(body_output, _, model_hparams, vocab_size):
-    return body_output
-
-
-class GenericL2LossModality(IdentityModality):
-  """Generic modality with L2 as Loss."""
-
-  @staticmethod
-  def targets_bottom(x, model_hparams, vocab_size):
-    return tf.to_float(x)
+def image_channel_embeddings_bottom(x, model_hparams, vocab_size):
+  """Bottom transformation for image targets."""
+  del vocab_size  # unused arg
+  inputs = tf.to_int32(x)
+  io_depth = model_hparams.num_channels
+  tshape = common_layers.shape_list(inputs)
+  hidden_size = model_hparams.hidden_size
+  target_embeddings = cia.get_channel_embeddings(
+      io_depth, inputs, hidden_size, "input_bottom")
+  return tf.reshape(target_embeddings,
+                    [tshape[0], tshape[1], tshape[2] * io_depth, hidden_size])
+
+
+def image_channel_embeddings_top(body_output,
+                                 targets,
+                                 model_hparams,
+                                 vocab_size):
+  """Top transformation for images."""
+  del targets  # unused arg
+  with tf.variable_scope("image_channel_embeddings_bottom"):
+    img_len = model_hparams.img_len
+    channels = model_hparams.num_channels
+    x = tf.layers.dense(
+        body_output, 256, use_bias=True, activation=None, name="output_conv")
+    x = tf.reshape(x,
+                   [-1, img_len, img_len, channels, vocab_size])
+    return x
 
-  @staticmethod
-  def loss(body_output, targets, model_hparams, vocab_size, weights_fn):
-    del weights_fn  # unused
-    loss = tf.squared_difference(body_output, tf.to_float(targets))
-    return tf.reduce_mean(loss), tf.constant(1.0)
 
+def audio_bottom(x, model_hparams, vocab_size):
+  """Transform input from data space to model space.
 
-class RealModality(modality.Modality):
-  """Base class for real (i.e. float) vectors.
+  Args:
+    x: A Tensor with shape [batch, ...]
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
 
-  * Bottom is a linear projection layer to hparams.hidden_size.
-  * Top is a linear projection layer to vocab_size.
+  Returns:
+    body_input: A Tensor with shape [batch, ?, ?,
+      model_hparams.hidden_size].
   """
+  del vocab_size  # unused arg
+  inputs = x
+  with tf.variable_scope("audio_modality"):
+    # TODO(aidangomez): Will need to sort out a better audio pipeline
+    def xnet_resblock(x, filters, res_relu, name):
+      """Xception block."""
+      with tf.variable_scope(name):
+        # Typically audio samples are >100k samples in length and have a width
+        # of 2 or 4. Mono audio has a single channel while stereo has 2.
+        y = common_layers.separable_conv_block(
+            x,
+            filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
+            first_relu=True,
+            padding="SAME",
+            force2d=True,
+            name="sep_conv_block")
+        y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 2))
+        return y + common_layers.conv_block(
+            x,
+            filters, [((1, 1), (1, 1))],
+            padding="SAME",
+            strides=(2, 2),
+            first_relu=res_relu,
+            force2d=True,
+            name="res_conv0")
+
+    x = tf.to_float(inputs) / 255.
+    x.set_shape([None, None, None, 1])
+    for i in range(model_hparams.audio_compression):
+      x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
+    return xnet_resblock(x,
+                         model_hparams.hidden_size,
+                         False,
+                         "compress_block_final")
+
+
+def audio_spectral_bottom(x, model_hparams, vocab_size):
+  """Transform input from data space to model space.
 
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size):
-    with tf.variable_scope("real"):
-      return tf.layers.dense(
-          tf.to_float(x), model_hparams.hidden_size, name="bottom")
-
-  @staticmethod
-  @is_pointwise
-  def top(body_output, _, model_hparams, vocab_size):
-    with tf.variable_scope("real"):
-      return tf.layers.dense(body_output, vocab_size, name="top")
-
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    raise NotImplementedError()
+  Args:
+    x: A Tensor with shape [batch, ...]
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
 
+  Returns:
+    body_input: A Tensor with shape [batch, ?, ?,
+      model_hparams.hidden_size].
+  """
+  del vocab_size  # unused arg
+  inputs = x
+  with tf.variable_scope("audio_spectral_modality"):
+    # TODO(aidangomez): Will need to sort out a better audio pipeline
+    def xnet_resblock(x, filters, res_relu, name):
+      """Xception-like block."""
+      with tf.variable_scope(name):
+        # We only stride along the length dimension to preserve the spectral
+        # bins (which are tiny in dimensionality relative to length)
+        y = common_layers.separable_conv_block(
+            x,
+            filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
+            first_relu=True,
+            padding="SAME",
+            force2d=True,
+            name="sep_conv_block")
+        y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 1))
+        return y + common_layers.conv_block(
+            x,
+            filters, [((1, 1), (1, 1))],
+            padding="SAME",
+            strides=(2, 1),
+            first_relu=res_relu,
+            force2d=True,
+            name="res_conv0")
+
+    # Bitcast back from int32
+    x = tf.bitcast(inputs, tf.float32)
+    x.set_shape([None, None, None, 1])
+    for i in range(model_hparams.audio_compression):
+      x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
+    return xnet_resblock(x,
+                         model_hparams.hidden_size,
+                         False,
+                         "compress_block_final")
+
+
+def speech_recognition_bottom(x, model_hparams, vocab_size):
+  """Use batchnorm instead of CMVN and shorten the stft with strided convs.
 
-class RealL2LossModality(RealModality):
-  """Modality for real (i.e. float) vectors with L2 (Gaussian) loss."""
+  Args:
+    x: float32 tensor with shape [batch_size, len, 1, freqs * channels]
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
 
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    predictions = top_out
-    if (len(common_layers.shape_list(top_out)) != len(
-        common_layers.shape_list(targets))):
-      predictions = tf.squeeze(top_out, axis=[-1])
-    with tf.name_scope("l2"):
-      weights = weights_fn(targets)
-      l2 = tf.pow(predictions - targets, 2)
-      return tf.reduce_sum(l2 * weights), tf.reduce_sum(weights)
+  Returns:
+    float32 tensor with shape [batch_size, shorter_len, 1, hidden_size]
+  """
+  del vocab_size  # unused arg
+  inputs = x
+  p = model_hparams
+
+  num_mel_bins = p.audio_num_mel_bins
+  num_channels = 3 if p.audio_add_delta_deltas else 1
+
+  with tf.variable_scope("speech_recognition_modality"):
+    if p.audio_preproc_in_bottom:
+      # Compute filterbanks
+      with tf.variable_scope("fbanks"):
+        waveforms = tf.squeeze(inputs, [2, 3])
+        mel_fbanks = common_audio.compute_mel_filterbank_features(
+            waveforms,
+            sample_rate=p.audio_sample_rate,
+            dither=p.audio_dither,
+            preemphasis=p.audio_preemphasis,
+            frame_length=p.audio_frame_length,
+            frame_step=p.audio_frame_step,
+            lower_edge_hertz=p.audio_lower_edge_hertz,
+            upper_edge_hertz=p.audio_upper_edge_hertz,
+            num_mel_bins=p.audio_num_mel_bins,
+            apply_mask=True)
+        if p.audio_add_delta_deltas:
+          mel_fbanks = common_audio.add_delta_deltas(mel_fbanks)
+        x = tf.reshape(mel_fbanks,
+                       common_layers.shape_list(mel_fbanks)[:2] +
+                       [num_mel_bins, num_channels])
+
+        nonpadding_mask = 1. - common_attention.embedding_to_padding(x)
+        num_of_nonpadding_elements = tf.reduce_sum(
+            nonpadding_mask) * num_mel_bins * num_channels
+
+        # This replaces CMVN estimation on data
+        var_epsilon = 1e-09
+        mean = tf.reduce_sum(
+            x, axis=[1], keepdims=True) / num_of_nonpadding_elements
+        variance = (num_of_nonpadding_elements * mean**2. -
+                    2. * mean * tf.reduce_sum(x, axis=[1], keepdims=True) +
+                    tf.reduce_sum(x**2, axis=[1], keepdims=True)
+                   ) / num_of_nonpadding_elements
+        x = (x - mean) * tf.rsqrt(variance + var_epsilon) * tf.expand_dims(
+            nonpadding_mask, -1)
+    else:
+      x = inputs
 
+    # The convention is that the models are flattened along the spatial,
+    # dimensions, thus the speech preprocessor treats frequencies and
+    # channels as image colors (last axis)
+    x.set_shape([None, None, num_mel_bins, num_channels])
 
-class RealLogPoissonLossModality(RealModality):
-  """Modality for real (i.e. float) vectors with log Poisson regression loss."""
+    # TODO(chorowski): how to specify bottom's hparams and avoid hardcoding?
+    x = tf.pad(x, [[0, 0], [0, 8], [0, 0], [0, 0]])
+    for _ in range(2):
+      x = tf.layers.conv2d(
+          x, 128, (3, 3), (2, 2), use_bias=False)
+      x = common_layers.layer_norm(x)
+      x = tf.nn.relu(x)
+
+    xshape = common_layers.shape_list(x)
+    # apply a conv that will remove all frequencies and at the same time
+    # project the output into desired hidden_size
+    x = tf.pad(x, [[0, 0], [0, 2], [0, 0], [0, 0]])
+    x = tf.layers.conv2d(x, p.hidden_size, (3, xshape[2]), use_bias=False)
+
+    assert common_layers.shape_list(x)[2] == 1
+    x = common_layers.layer_norm(x)
+    x = tf.nn.relu(x)
+  return x
+
+
+def video_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("inputs", x, max_outputs=1)
+  x = common_layers.standardize_images(x)
+  return x
+
+
+def video_targets_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("targets", x, max_outputs=1)
+  x = common_layers.standardize_images(x)
+  return x
+
+
+def video_top(body_output, targets, model_hparams, vocab_size):
+  """Top transformation for video."""
+  del targets  # unused arg
+  num_channels = model_hparams.problem.num_channels
+  shape = common_layers.shape_list(body_output)
+  reshape_shape = shape[:-1] + [num_channels, vocab_size]
+  res = tf.reshape(body_output, reshape_shape)
+  # Calculate argmax so as to have a summary with the produced images.
+  x = tf.argmax(tf.reshape(res, [-1, vocab_size]), axis=-1)
+  x = tf.reshape(x, shape[:-1] + [num_channels])
+  common_video.gif_summary("results", x, max_outputs=1)
+  return res
+
+
+def video_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  """Compute loss numerator and denominator for one shard of output."""
+  del vocab_size  # unused arg
+  logits = top_out
+  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
+  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
+  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
+  return common_layers.padded_cross_entropy(
+      logits,
+      targets,
+      model_hparams.label_smoothing,
+      cutoff=cutoff,
+      weights_fn=weights_fn)
+
+
+def video_bitwise_bottom(x, model_hparams, vocab_size):
+  """Bottom transformation for embedding video bitwise."""
+  pixel_embedding_size = 64
+  inputs = x
+  with tf.variable_scope("video_modality_bitwise", reuse=tf.AUTO_REUSE):
+    common_layers.summarize_video(inputs, "bottom")
+    # Embed bitwise.
+    assert vocab_size == 256
+    embedded = discretization.int_to_bit_embed(inputs, 8,
+                                               pixel_embedding_size)
+    # Project.
+    return tf.layers.dense(
+        embedded,
+        model_hparams.hidden_size,
+        name="merge_pixel_embedded_frames")
+
+
+def video_bitwise_targets_bottom(x, model_hparams, vocab_size):
+  """Bottom transformation for embedding target video bitwise."""
+  pixel_embedding_size = 64
+  inputs = x
+  with tf.variable_scope("video_modality_bitwise", reuse=tf.AUTO_REUSE):
+    common_layers.summarize_video(inputs, "targets_bottom")
+    # Embed bitwise.
+    assert vocab_size == 256
+    embedded = discretization.int_to_bit_embed(inputs, 8,
+                                               pixel_embedding_size)
+    # Transpose and project.
+    transposed = common_layers.time_to_channels(embedded)
+    return tf.layers.dense(
+        transposed,
+        model_hparams.hidden_size,
+        name="merge_pixel_embedded_frames")
+
+
+def video_pixel_noise_bottom(x, model_hparams, vocab_size):
+  """Bottom transformation for video."""
+  input_noise = getattr(model_hparams, "video_modality_input_noise", 0.25)
+  inputs = x
+  if model_hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    background = tfp.distributions.percentile(inputs, 50., axis=[0, 1, 2, 3])
+    input_shape = common_layers.shape_list(inputs)
+    input_size = tf.reduce_prod(input_shape[:-1])
+    input_mask = tf.multinomial(
+        tf.log([[input_noise, 1.-input_noise]]), input_size)
+    input_mask = tf.reshape(tf.cast(input_mask, tf.int32),
+                            input_shape[:-1]+[1])
+    inputs = inputs * input_mask + background * (1 - input_mask)
+  return video_bottom(inputs, model_hparams, vocab_size)
+
+
+def video_l1_top(body_output, targets, model_hparams, vocab_size):
+  """Top transformation for video."""
+  del targets, vocab_size  # unused arg
+  num_channels = model_hparams.problem.num_channels
+  num_frames = model_hparams.video_num_target_frames
+  with tf.variable_scope("rgb"):
+    body_output_shape = common_layers.shape_list(body_output)
+    res = tf.layers.dense(body_output, num_channels * num_frames, name="cast")
+    res = tf.reshape(res, body_output_shape[:3] + [num_channels, num_frames])
+    res = tf.transpose(res, [0, 4, 1, 2, 3])  # Move frames next to batch.
+    if not tf.get_variable_scope().reuse:
+      res_argmax = res[:, -1, :, :, :]
+      tf.summary.image(
+          "result",
+          common_layers.tpu_safe_image_summary(res_argmax),
+          max_outputs=1)
+    return tf.expand_dims(res, axis=-1)  # Add an axis like in perplexity.
+
+
+def video_l1_internal_loss(logits, targets, model_hparams):
+  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
+  return tf.nn.relu(tf.abs(logits - targets) - cutoff)
+
+
+def video_l1_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  """Compute loss numerator and denominator for one shard of output."""
+  del vocab_size  # unused arg
+  logits = top_out
+  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:-1])
+  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
+  weights = weights_fn(targets)
+  # Shift targets by 0.5 so later just casting to int gives the prediction.
+  # So for int targets, say 0 and 7, we actually train to predict 0.5 and 7.5.
+  # Later (in merics or infer) this is cast to int anyway. Also, we have no
+  # loss beyond cutoff = 0.2 as these are already correct predictions.
+  targets = tf.to_float(targets) + 0.5
+  loss = video_l1_internal_loss(logits, targets, model_hparams)
+  return tf.reduce_sum(loss * weights), tf.reduce_sum(weights)
+
+
+def video_l2_internal_loss(logits, targets, model_hparams):
+  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
+  return tf.nn.relu(
+      tf.squared_difference(logits, targets) - cutoff * cutoff)
+
+
+def video_l2_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  """Compute loss numerator and denominator for one shard of output."""
+  del vocab_size  # unused arg
+  logits = top_out
+  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:-1])
+  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
+  weights = weights_fn(targets)
+  # Shift targets by 0.5 so later just casting to int gives the prediction.
+  # So for int targets, say 0 and 7, we actually train to predict 0.5 and 7.5.
+  # Later (in merics or infer) this is cast to int anyway. Also, we have no
+  # loss beyond cutoff = 0.2 as these are already correct predictions.
+  targets = tf.to_float(targets) + 0.5
+  loss = video_l2_internal_loss(logits, targets, model_hparams)
+  return tf.reduce_sum(loss * weights), tf.reduce_sum(weights)
+
+
+def convert_rgb_to_real(prediction, targets):
+  """Convert prediction and target from rgb to real."""
+  prediction = tf.squeeze(prediction, axis=-1)
+  prediction = common_layers.convert_rgb_to_real(prediction)
+  targets = common_layers.convert_rgb_to_real(targets)
+  return prediction, targets
+
+
+def video_raw_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("inputs", x)
+  return common_layers.convert_rgb_to_real(x)
+
+
+def video_raw_targets_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("targets_bottom", x)
+  return common_layers.convert_rgb_to_real(x)
+
+
+def video_raw_top(body_output, targets, model_hparams, vocab_size):
+  del targets, model_hparams, vocab_size  # unused arg
+  frames = body_output
+  if isinstance(body_output, list):
+    frames = tf.stack(body_output, axis=1)
+  rgb_frames = common_layers.convert_real_to_rgb(frames)
+  common_video.gif_summary("body_output", rgb_frames)
+  return tf.expand_dims(rgb_frames, axis=-1)
+
+
+def video_l2_raw_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  del model_hparams, vocab_size, weights_fn  # unused arg
+  prediction, groundtruth = convert_rgb_to_real(top_out, targets)
+  loss = tf.losses.mean_squared_error(prediction, groundtruth)
+  return loss, tf.constant(1.0)
+
+
+def video_l1_raw_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  del model_hparams, vocab_size, weights_fn  # unused arg
+  prediction, groundtruth = convert_rgb_to_real(top_out, targets)
+  loss = tf.losses.absolute_difference(prediction, groundtruth)
+  return loss, tf.constant(1.0)
+
+
+def class_label_bottom(x, model_hparams, vocab_size):
+  with tf.variable_scope("class_label_modality_%d_%d" % (
+      vocab_size, model_hparams.hidden_size)):
+    multiplier = 1.0
+    if model_hparams.multiply_embedding_mode == "sqrt_depth":
+      multiplier = model_hparams.hidden_size**0.5
+    return common_layers.embedding(x,
+                                   vocab_size,
+                                   model_hparams.hidden_size,
+                                   multiplier=multiplier)
 
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    predictions = top_out
-    if (len(common_layers.shape_list(top_out)) != len(
-        common_layers.shape_list(targets))):
-      predictions = tf.squeeze(top_out, axis=[-1])
-    with tf.name_scope("log_possion"):
-      weights = weights_fn(targets)
-      lp_loss = tf.nn.log_poisson_loss(targets, predictions)
-      return tf.reduce_sum(lp_loss * weights), tf.reduce_sum(weights)
 
+def class_label_targets_bottom(x, model_hparams, vocab_size):
+  with tf.variable_scope("class_label_modality_%d_%d" % (
+      vocab_size, model_hparams.hidden_size)):
+    return tf.zeros([common_layers.shape_list(x)[0],
+                     1,
+                     1,
+                     model_hparams.hidden_size])
+
+
+def class_label_top(body_output, targets, model_hparams, vocab_size):
+  """Transform inputs from model space to target space.
+
+  Average over inner dims and a linear layer to logits.
 
-class IdentitySymbolModality(SymbolModality):
-  """Symbol modality with identity top and bottom transformations.
+  Args:
+    body_output: A Tensor with shape [batch, ?, ?, body_output_size].
+    targets:
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
 
-  Uses the weights_fn from SymbolModality so that loss/metrics ignore padding.
+  Returns:
+    a Tensors, each with shape [batch_size, 1, 1, 1, vocab_size]
   """
+  del targets  # unused arg
+  with tf.variable_scope("class_label_modality_%d_%d" % (
+      vocab_size, model_hparams.hidden_size)):
+    x = body_output
+    x = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
+    res = tf.layers.dense(x, vocab_size)
+    return tf.expand_dims(res, 3)
+
+
+def video_identity_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("inputs", x, max_outputs=1)
+  return x
+
+
+def video_identity_targets_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("targets", x, max_outputs=1)
+  return x
+
+
+def video_identity_loss(top_out,
+                        targets,
+                        model_hparams,
+                        vocab_size,
+                        weights_fn):
+  """Compute loss numerator and denominator for one shard of output."""
+  del vocab_size  # unused arg
+  # TODO(nikip): Try L2 loss
+  logits = top_out
+  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
+  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
+  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
+  return common_layers.padded_cross_entropy(
+      logits,
+      targets,
+      model_hparams.label_smoothing,
+      cutoff=cutoff,
+      weights_fn=weights_fn)
+
+
+def multi_label_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  """Average loss over the labels."""
+  del vocab_size  # unused arg
+  logits = top_out
+  num_labels = tf.shape(targets)[1]
+  logits = tf.tile(logits, [1, num_labels, 1, 1, 1])
+
+  xent, weights = common_layers.padded_cross_entropy(
+      logits,
+      targets,
+      model_hparams.label_smoothing,
+      weights_fn=weights_fn,
+      reduce_sum=False,
+  )
+  xent = tf.squeeze(xent, [2, 3])
+  weights = tf.squeeze(weights, [2, 3])
+  # average loss over all labels
+  loss = tf.reduce_sum(xent, axis=1)
+  weights = tf.reduce_sum(weights, axis=1)
+  loss /= (weights + 1e-8)
+  weights = tf.to_float(tf.greater(weights, 0.))
+
+  return tf.reduce_sum(loss*weights), tf.reduce_sum(weights)
+
+
+def one_hot_class_label_loss(top_out,
+                             targets,
+                             model_hparams,
+                             vocab_size,
+                             weights_fn):
+  """Apply softmax cross-entropy between outputs and targets.
 
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size):
-    return tf.to_float(x)
-
-  @staticmethod
-  def top(body_output, _, model_hparams, vocab_size):
-    return body_output
+  Args:
+    top_out: logits Tensor with shape [batch, ?, ?, num_classes]
+    targets: one-hot encoding Tensor with shape [batch, ?, ?, num_classes]
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
+    weights_fn:
 
-  @classmethod
-  def targets_bottom(cls, x, model_hparams, vocab_size):
-    """SymbolModality overrides targets_bottom, so need to override here too."""
-    return cls.bottom(x, model_hparams, vocab_size)
+  Returns:
+    loss_scale (cross-entropy), loss_denom
+  """
+  del model_hparams, vocab_size  # unused arg
+  loss_scale = tf.losses.softmax_cross_entropy(
+      onehot_labels=targets, logits=top_out)
+  weights = weights_fn(targets)
+  loss_denom = tf.reduce_sum(weights)
+  return loss_scale, loss_denom
 
 
-class SigmoidClassLabelModality(ClassLabelModality):
-  """Sigmoid cross-entropy for independent class labels."""
+def identity_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  return tf.to_float(x)
 
-  @staticmethod
-  def name(model_hparams, vocab_size):
-    return "sigmoid_class_symbol_modality_%d_%d" % (vocab_size,
-                                                    model_hparams.hidden_size)
 
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    # Expect inputs of size [batch-size, timesteps, 1, num-classes], where the
-    # last dimension of num-classes represents logits for binary labels
-    loss_scale = tf.losses.sigmoid_cross_entropy(
-        multi_class_labels=targets, logits=top_out)
-    weights = weights_fn(targets)
-    loss_denom = tf.reduce_sum(weights)
-    return loss_scale, loss_denom
+def identity_top(body_output, targets, model_hparams, vocab_size):
+  del targets, model_hparams, vocab_size  # unused arg
+  return body_output
 
 
-class SigmoidMaxPoolingClassLabelModality(ClassLabelModality):
-  """Sigmoid cross-entropy applied on max-pooling over timesteps."""
+def generic_l2_loss(body_output,
+                    targets,
+                    model_hparams,
+                    vocab_size,
+                    weights_fn):
+  del model_hparams, vocab_size, weights_fn  # unused arg
+  loss = tf.squared_difference(body_output, tf.to_float(targets))
+  return tf.reduce_mean(loss), tf.constant(1.0)
 
-  @staticmethod
-  def name(model_hparams, vocab_size):
-    return "sigmoid_max_pooling_class_symbol_modality_%d_%d" % (
-        vocab_size, model_hparams.hidden_size)
 
-  @classmethod
-  def top(cls, body_output, _, model_hparams, vocab_size):
-    """Transform inputs from model space to target space.
+def real_bottom(x, model_hparams, vocab_size):
+  del vocab_size  # unused arg
+  with tf.variable_scope("real"):
+    return tf.layers.dense(
+        tf.to_float(x), model_hparams.hidden_size, name="bottom")
 
-    Average over inner dims and a linear layer to logits.
 
-    Args:
-      body_output: A Tensor with shape [batch, timesteps, 1, body_output_size].
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
+@is_pointwise
+def real_top(body_output, targets, model_hparams, vocab_size):
+  del targets, model_hparams  # unused arg
+  with tf.variable_scope("real"):
+    return tf.layers.dense(body_output, vocab_size, name="top")
 
-    Returns:
-      a Tensors, each with shape [batch_size, 1, 1, vocab_size]
-    """
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      x = body_output
-      x = tf.reduce_max(x, axis=1, keepdims=True)
-      return tf.layers.dense(x, vocab_size)
 
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    # Expect inputs of size [batch-size, 1, 1, num-classes], where the
-    # last dimension of num-classes represents logits for binary labels
-    loss_scale = tf.losses.sigmoid_cross_entropy(
-        multi_class_labels=targets, logits=top_out)
+def real_l2_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  del model_hparams, vocab_size  # unused arg
+  predictions = top_out
+  if (len(common_layers.shape_list(top_out)) != len(
+      common_layers.shape_list(targets))):
+    predictions = tf.squeeze(top_out, axis=[-1])
+  with tf.name_scope("l2"):
     weights = weights_fn(targets)
-    loss_denom = tf.reduce_sum(weights)
-    return loss_scale, loss_denom
-
-
-class SoftmaxMaxPoolingClassLabelModality(OneHotClassLabelModality):
-  """Softmax cross-entropy applied on max-pooling over timesteps."""
-
-  @staticmethod
-  def name(model_hparams, vocab_size):
-    return "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
-        vocab_size, model_hparams.hidden_size)
-
-  @classmethod
-  def top(cls, body_output, _, model_hparams, vocab_size):
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      x = body_output
-      x = tf.reduce_max(x, axis=1, keepdims=True)
-      return tf.layers.dense(x, vocab_size)
+    l2 = tf.pow(predictions - targets, 2)
+    return tf.reduce_sum(l2 * weights), tf.reduce_sum(weights)
 
 
-class SoftmaxAveragePoolingClassLabelModality(OneHotClassLabelModality):
-  """Softmax cross-entropy applied on average-pooling over timesteps."""
-
-  @staticmethod
-  def name(model_hparams, vocab_size):
-    return "softmax_average_pooling_onehot_class_label_modality_%d_%d" % (
-        vocab_size, model_hparams.hidden_size)
-
-  @classmethod
-  def top(cls, body_output, _, model_hparams, vocab_size):
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      x = body_output
-      x = tf.reduce_mean(x, axis=1, keepdims=True)
-      return tf.layers.dense(x, vocab_size)
+def real_log_poisson_loss(top_out,
+                          targets,
+                          model_hparams,
+                          vocab_size,
+                          weights_fn):
+  """Poisson loss for real."""
+  del model_hparams, vocab_size  # unused arg
+  predictions = top_out
+  if (len(common_layers.shape_list(top_out)) != len(
+      common_layers.shape_list(targets))):
+    predictions = tf.squeeze(top_out, axis=[-1])
+  with tf.name_scope("log_possion"):
+    weights = weights_fn(targets)
+    lp_loss = tf.nn.log_poisson_loss(targets, predictions)
+    return tf.reduce_sum(lp_loss * weights), tf.reduce_sum(weights)
 
 
-class SoftmaxLastTimestepClassLabelModality(OneHotClassLabelModality):
-  """Softmax cross-entropy applied on last-timestep encoding."""
+def sigmoid_class_label_loss(top_out,
+                             targets,
+                             model_hparams,
+                             vocab_size,
+                             weights_fn):
+  """Loss for class label."""
+  # Expect inputs of size [batch-size, timesteps, 1, num-classes], where the
+  # last dimension of num-classes represents logits for binary labels
+  del model_hparams, vocab_size  # unused arg
+  loss_scale = tf.losses.sigmoid_cross_entropy(
+      multi_class_labels=targets, logits=top_out)
+  weights = weights_fn(targets)
+  loss_denom = tf.reduce_sum(weights)
+  return loss_scale, loss_denom
+
+
+def sigmoid_max_pooling_class_label_top(body_output,
+                                        targets,
+                                        model_hparams,
+                                        vocab_size):
+  """Transform inputs from model space to target space.
+
+  Average over inner dims and a linear layer to logits.
 
-  @staticmethod
-  def name(model_hparams, vocab_size):
-    return "softmax_last_timestep_onehot_class_label_modality_%d_%d" % (
-        vocab_size, model_hparams.hidden_size)
+  Args:
+    body_output: A Tensor with shape [batch, timesteps, 1, body_output_size].
+    targets:
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
 
-  @classmethod
-  def top(cls, body_output, _, model_hparams, vocab_size):
-    with tf.variable_scope(cls.name(model_hparams, vocab_size)):
-      x = body_output
-      x = tf.expand_dims(x[:, -1], 1)  # Pick the last timestep
-      return tf.layers.dense(x, vocab_size)
+  Returns:
+    a Tensors, each with shape [batch_size, 1, 1, vocab_size]
+  """
+  del targets  # unused arg
+  with tf.variable_scope(
+      "sigmoid_max_pooling_class_symbol_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)):
+    x = body_output
+    x = tf.reduce_max(x, axis=1, keepdims=True)
+    return tf.layers.dense(x, vocab_size)
+
+
+def sigmoid_max_pooling_class_label_loss(top_out,
+                                         targets,
+                                         model_hparams,
+                                         vocab_size,
+                                         weights_fn):
+  """Loss for class label."""
+  # Expect inputs of size [batch-size, 1, 1, num-classes], where the
+  # last dimension of num-classes represents logits for binary labels
+  del model_hparams, vocab_size  # unused arg
+  loss_scale = tf.losses.sigmoid_cross_entropy(
+      multi_class_labels=targets, logits=top_out)
+  weights = weights_fn(targets)
+  loss_denom = tf.reduce_sum(weights)
+  return loss_scale, loss_denom
+
+
+def softmax_max_pooling_class_label_top(body_output,
+                                        targets,
+                                        model_hparams,
+                                        vocab_size):
+  """Loss for class label."""
+  del targets  # unused arg
+  with tf.variable_scope(
+      "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)):
+    x = body_output
+    x = tf.reduce_max(x, axis=1, keepdims=True)
+    return tf.layers.dense(x, vocab_size)
+
+
+def softmax_average_pooling_class_label_top(body_output,
+                                            targets,
+                                            model_hparams,
+                                            vocab_size):
+  """Loss for class label."""
+  del targets  # unused arg
+  with tf.variable_scope(
+      "softmax_average_pooling_onehot_class_label_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)):
+    x = body_output
+    x = tf.reduce_mean(x, axis=1, keepdims=True)
+    return tf.layers.dense(x, vocab_size)
+
+
+def softmax_last_timestep_class_label_top(body_output,
+                                          targets,
+                                          model_hparams,
+                                          vocab_size):
+  """Loss for class label."""
+  del targets  # unused arg
+  with tf.variable_scope(
+      "softmax_last_timestep_onehot_class_label_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)):
+    x = body_output
+    x = tf.expand_dims(x[:, -1], 1)  # Pick the last timestep
+    return tf.layers.dense(x, vocab_size)
 
 
 class ModalityType(object):
   """Types of modalities."""
 
-  SYMBOL = "SymbolModality"
-  SYMBOL_WEIGHTS_ALL = "SymbolModalityWeightsAll"
-  SYMBOL_ONE_HOT = "SymbolModalityOneHot"
-  CTC_SYMBOL = "CTCSymbolModality"
-  IMAGE = "ImageModality"
-  IMAGE_CHANNEL_COMPRESS = "ImageChannelCompressModality"
-  IMAGE_CHANNEL_BOTTOM_IDENTITY = "ImageChannelBottomIdentityModality"
-  IMAGE_CHANNEL_EMBEDDINGS_BOTTOM = "ImageChannelEmbeddingsBottom"
-  AUDIO = "AudioModality"
-  AUDIO_SPECTRAL = "AudioSpectralModality"
-  SPEECH_RECOGNITION = "SpeechRecognitionModality"
-  VIDEO = "VideoModality"
-  VIDEO_BITWISE = "VideoModalityBitwise"
-  VIDEO_PIXEL_NOISE = "VideoModalityPixelNoise"
-  VIDEO_L1 = "VideoModalityL1"
-  VIDEO_L2 = "VideoModalityL2"
-  VIDEO_L2_RAW = "VideoModalityL2Raw"
-  VIDEO_L1_RAW = "VideoModalityL1Raw"
-  CLASS_LABEL = "ClassLabelModality"
-  VIDEO_IDENTITY = "VideoModalityIdentity"
-  MULTI_LABEL = "MultiLabelModality"
-  ONE_HOT_CLASS_LABEL = "OneHotClassLabelModality"
-  IDENTITY = "IdentityModality"
-  GENERIC_L2_LOSS = "GenericL2LossModality"
-  REAL = "RealModality"
-  REAL_L2_LOSS = "RealL2LossModality"
-  REAL_LOG_POISSON_LOSS = "RealLogPoissonLossModality"
-  IDENTITY_SYMBOL = "IdentitySymbolModality"
-  SIGMOID_CLASS_LABEL = "SigmoidClassLabelModality"
-  SIGMOID_MAX_POOLING_CLASS_LABEL = "SigmoidMaxPoolingClassLabelModality"
-  SOFTMAX_MAX_POOLING_CLASS_LABEL = "SoftmaxMaxPoolingClassLabelModality"
-  SOFTMAX_AVERAGE_POOLING_CLASS_LABEL = "SoftmaxAveragePoolingClassLabelModality"
-  SOFTMAX_LAST_TIMESTEP_CLASS_LABEL = "SoftmaxLastTimestepClassLabelModality"
+  AUDIO = "audio"
+  AUDIO_SPECTRAL = "audio_spectral"
+  CLASS_LABEL = "class_label"
+  CTC_SYMBOL = "ctc_symbol"  # symbol with CTC loss
+  GENERIC_L2_LOSS = "generic_l2"  # identity modality with L2 loss
+  IDENTITY = "identity"  # identity top and bottom
+  IDENTITY_SYMBOL = "identity_symbol"  # symbol with identity top and bottom
+  IMAGE = "image"
+  # images using channel compression for generation
+  IMAGE_CHANNEL_BOTTOM_IDENTITY = "image_channel_bottom_identity"
+  # images using channel compression for generation
+  IMAGE_CHANNEL_COMPRESS = "image_channel_compress"
+  IMAGE_CHANNEL_EMBEDDINGS_BOTTOM = "image_channel_embeddings_bottom"
+  MULTI_LABEL = "multi_label"
+  ONE_HOT_CLASS_LABEL = "one_hot_class_label"
+  REAL = "real"  # real vectors
+  REAL_L2_LOSS = "real_l2"  # real vectors with L2 as loss
+  # real vectors with log Poisson regression loss
+  REAL_LOG_POISSON_LOSS = "real_log_poisson"
+  SIGMOID_CLASS_LABEL = "sigmoid_class_label"  # sigmoid cross-entropy loss
+  # sigmoid cross-entropy applied on max-pooling over timesteps
+  SIGMOID_MAX_POOLING_CLASS_LABEL = "sigmoid_max_pooling_class_label"
+  # softmax cross-entropy applied on average-pooling over timesteps
+  SOFTMAX_AVERAGE_POOLING_CLASS_LABEL = "softmax_average_pooling_class_label"
+  # softmax cross-entropy applied on last-timestep encoding
+  SOFTMAX_LAST_TIMESTEP_CLASS_LABEL = "softmax_last_timestep_class_label"
+  # softmax cross-entropy applied on max-pooling over timesteps
+  SOFTMAX_MAX_POOLING_CLASS_LABEL = "softmax_max_pooling_class_label"
+  SPEECH_RECOGNITION = "speech_recognition"
+  SYMBOL = "symbol"
+  SYMBOL_WEIGHTS_ALL = "symbol_weights_all"  # symbol for features w/o 0-padding
+  SYMBOL_ONE_HOT = "symbol_one_hot"  # symbol with one hot as embeddings
+  VIDEO = "video"
+  VIDEO_BITWISE = "video_bitwise"  # video where bottom embeds pixels bitwise
+  VIDEO_IDENTITY = "video_identity"  # video with identity top and bottom
+  VIDEO_L1 = "video_l1"  # video with L2 loss
+  VIDEO_L2 = "video_l2"  # video with L1 loss
+  # video with L1 loss and raw input (sequences of frames)
+  VIDEO_L1_RAW = "video_l1_raw"
+  # video with L2 loss and raw input (sequences of frames)
+  VIDEO_L2_RAW = "video_l2_raw"
+  # video with pixel noise on input during training
+  VIDEO_PIXEL_NOISE = "video_pixel_noise"
 
   @staticmethod
   def get_choices():
     return [
-        ModalityType.SYMBOL,
-        ModalityType.SYMBOL_WEIGHTS_ALL,
-        ModalityType.SYMBOL_ONE_HOT,
+        ModalityType.AUDIO,
+        ModalityType.AUDIO_SPECTRAL,
+        ModalityType.CLASS_LABEL,
         ModalityType.CTC_SYMBOL,
+        ModalityType.GENERIC_L2_LOSS,
+        ModalityType.IDENTITY,
+        ModalityType.IDENTITY_SYMBOL,
         ModalityType.IMAGE,
-        ModalityType.IMAGE_CHANNEL_COMPRESS,
         ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
+        ModalityType.IMAGE_CHANNEL_COMPRESS,
         ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM,
-        ModalityType.AUDIO,
-        ModalityType.AUDIO_SPECTRAL,
-        ModalityType.SPEECH_RECOGNITION,
-        ModalityType.VIDEO,
-        ModalityType.VIDEO_BITWISE,
-        ModalityType.VIDEO_PIXEL_NOISE,
-        ModalityType.VIDEO_L1,
-        ModalityType.VIDEO_L2,
-        ModalityType.VIDEO_L2_RAW,
-        ModalityType.VIDEO_L1_RAW,
-        ModalityType.CLASS_LABEL,
-        ModalityType.VIDEO_IDENTITY,
         ModalityType.MULTI_LABEL,
         ModalityType.ONE_HOT_CLASS_LABEL,
-        ModalityType.IDENTITY,
-        ModalityType.GENERIC_L2_LOSS,
         ModalityType.REAL,
         ModalityType.REAL_L2_LOSS,
         ModalityType.REAL_LOG_POISSON_LOSS,
-        ModalityType.IDENTITY_SYMBOL,
         ModalityType.SIGMOID_CLASS_LABEL,
         ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
-        ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL,
         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
+        ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL,
+        ModalityType.SPEECH_RECOGNITION,
+        ModalityType.SYMBOL,
+        ModalityType.SYMBOL_ONE_HOT,
+        ModalityType.SYMBOL_WEIGHTS_ALL,
+        ModalityType.VIDEO,
+        ModalityType.VIDEO_BITWISE,
+        ModalityType.VIDEO_IDENTITY,
+        ModalityType.VIDEO_L1,
+        ModalityType.VIDEO_L2,
+        ModalityType.VIDEO_L1_RAW,
+        ModalityType.VIDEO_L2_RAW,
+        ModalityType.VIDEO_PIXEL_NOISE,
     ]
 
 
 # Utility functions, similar to tf.keras
-current_module = sys.modules[__name__]
 
 
 def get_bottom(modality_type, value=None):
   """Gets default bottom transformation; if none available, return value."""
-  if modality_type in ModalityType.get_choices():
-    modality_cls = getattr(current_module, modality_type)
-    return modality_cls.bottom
+  if modality_type == ModalityType.AUDIO:
+    return audio_bottom
+  elif modality_type == ModalityType.AUDIO_SPECTRAL:
+    return audio_spectral_bottom
+  elif modality_type in (ModalityType.CLASS_LABEL,
+                         ModalityType.MULTI_LABEL,
+                         ModalityType.ONE_HOT_CLASS_LABEL,
+                         ModalityType.SIGMOID_CLASS_LABEL,
+                         ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
+                         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
+                         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
+                         ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL):
+    return class_label_bottom
+  elif modality_type in (ModalityType.CTC_SYMBOL,
+                         ModalityType.SYMBOL,
+                         ModalityType.SYMBOL_WEIGHTS_ALL):
+    return symbol_bottom
+  elif modality_type in (ModalityType.GENERIC_L2_LOSS,
+                         ModalityType.IDENTITY,
+                         ModalityType.IDENTITY_SYMBOL,
+                         ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM):
+    return identity_bottom
+  elif modality_type == ModalityType.IMAGE:
+    return image_bottom
+  elif modality_type in (ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
+                         ModalityType.IMAGE_CHANNEL_COMPRESS):
+    return image_channel_compress_bottom
+  elif modality_type in (ModalityType.REAL,
+                         ModalityType.REAL_L2_LOSS,
+                         ModalityType.REAL_LOG_POISSON_LOSS):
+    return real_bottom
+  elif modality_type == ModalityType.SPEECH_RECOGNITION:
+    return speech_recognition_bottom
+  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
+    return symbol_one_hot_bottom
+  elif modality_type in (ModalityType.VIDEO,
+                         ModalityType.VIDEO_L1,
+                         ModalityType.VIDEO_L2):
+    return video_bottom
+  elif modality_type == ModalityType.VIDEO_BITWISE:
+    return video_bitwise_bottom
+  elif modality_type == ModalityType.VIDEO_IDENTITY:
+    return video_identity_bottom
+  elif modality_type in (ModalityType.VIDEO_L1_RAW,
+                         ModalityType.VIDEO_L2_RAW):
+    return video_raw_bottom
+  elif modality_type == ModalityType.VIDEO_PIXEL_NOISE:
+    return video_pixel_noise_bottom
   return value
 
 
 def get_loss(modality_type, value=None):
   """Gets default loss transformation; if none available, return value."""
-  if modality_type in ModalityType.get_choices():
-    modality_cls = getattr(current_module, modality_type)
-    return modality_cls.loss
+  if modality_type in (ModalityType.AUDIO,
+                       ModalityType.AUDIO_SPECTRAL,
+                       ModalityType.CLASS_LABEL,
+                       ModalityType.IDENTITY,
+                       ModalityType.IDENTITY_SYMBOL,
+                       ModalityType.IMAGE,
+                       ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
+                       ModalityType.IMAGE_CHANNEL_COMPRESS,
+                       ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM,
+                       ModalityType.REAL,
+                       ModalityType.SPEECH_RECOGNITION,
+                       ModalityType.SYMBOL,
+                       ModalityType.SYMBOL_WEIGHTS_ALL):
+    return generic_loss
+  elif modality_type == ModalityType.CTC_SYMBOL:
+    return ctc_symbol_loss
+  elif modality_type == ModalityType.GENERIC_L2_LOSS:
+    return generic_l2_loss
+  elif modality_type == ModalityType.MULTI_LABEL:
+    return multi_label_loss
+  elif modality_type in (ModalityType.ONE_HOT_CLASS_LABEL,
+                         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
+                         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
+                         ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL):
+    return one_hot_class_label_loss
+  elif modality_type == ModalityType.REAL_L2_LOSS:
+    return real_l2_loss
+  elif modality_type == ModalityType.REAL_LOG_POISSON_LOSS:
+    return real_log_poisson_loss
+  elif modality_type == ModalityType.SIGMOID_CLASS_LABEL:
+    return sigmoid_class_label_loss
+  elif modality_type == ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL:
+    return sigmoid_max_pooling_class_label_loss
+  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
+    return symbol_one_hot_loss
+  elif modality_type in (ModalityType.VIDEO,
+                         ModalityType.VIDEO_BITWISE,
+                         ModalityType.VIDEO_PIXEL_NOISE):
+    return video_loss
+  elif modality_type == ModalityType.VIDEO_IDENTITY:
+    return video_identity_loss
+  elif modality_type == ModalityType.VIDEO_L1:
+    return video_l1_loss
+  elif modality_type == ModalityType.VIDEO_L1_RAW:
+    return video_l1_raw_loss
+  elif modality_type == ModalityType.VIDEO_L2:
+    return video_l2_loss
+  elif modality_type == ModalityType.VIDEO_L2_RAW:
+    return video_l2_raw_loss
   return value
 
 
 def get_name(modality_type, value=None):
   """Gets default name for transformations; if none available, return value."""
-  if modality_type in ModalityType.get_choices():
-    modality_cls = getattr(current_module, modality_type)
-    return modality_cls.name
+  # For legacy reasons, modalities vary in their naming scheme.
+  if modality_type == ModalityType.AUDIO:
+    return lambda model_hparams, vocab_size: "audio_modality"
+  elif modality_type == ModalityType.AUDIO_SPECTRAL:
+    return lambda model_hparams, vocab_size: "audio_spectral_modality"
+  elif modality_type == ModalityType.GENERIC_L2_LOSS:
+    return lambda model_hparams, vocab_size: "generic_l2_loss_modality"
+  elif modality_type == ModalityType.IDENTITY:
+    return lambda model_hparams, vocab_size: "identity_modality"
+  elif modality_type == ModalityType.IMAGE:
+    return lambda model_hparams, vocab_size: "image_modality"
+  elif modality_type == ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY:
+    return (lambda model_hparams, vocab_size:  # pylint: disable=g-long-lambda
+            "image_channel_bottom_identity_modality")
+  elif modality_type == ModalityType.IMAGE_CHANNEL_COMPRESS:
+    return lambda model_hparams, vocab_size: "image_channel_compress_modality"
+  elif modality_type == ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM:
+    return lambda model_hparams, vocab_size: "image_channel_embeddings_bottom"
+  elif modality_type == ModalityType.REAL:
+    return lambda model_hparams, vocab_size: "real_modality"
+  elif modality_type == ModalityType.REAL_L2_LOSS:
+    return lambda model_hparams, vocab_size: "real_l2_loss_modality"
+  elif modality_type == ModalityType.REAL_LOG_POISSON_LOSS:
+    return lambda model_hparams, vocab_size: "real_log_poisson_loss_modality"
+  elif modality_type == ModalityType.SPEECH_RECOGNITION:
+    return lambda model_hparams, vocab_size: "speech_recognition_modality"
+  elif modality_type == ModalityType.VIDEO:
+    return lambda model_hparams, vocab_size: "video_modality"
+  elif modality_type == ModalityType.VIDEO_BITWISE:
+    return lambda model_hparams, vocab_size: "video_modality_bitwise"
+  elif modality_type == ModalityType.VIDEO_IDENTITY:
+    return lambda model_hparams, vocab_size: "video_modality_identity"
+  elif modality_type == ModalityType.VIDEO_L1:
+    return lambda model_hparams, vocab_size: "video_modality_l1"
+  elif modality_type == ModalityType.VIDEO_L1_RAW:
+    return lambda model_hparams, vocab_size: "video_modality_l1_raw"
+  elif modality_type == ModalityType.VIDEO_L2:
+    return lambda model_hparams, vocab_size: "video_modality_l2"
+  elif modality_type == ModalityType.VIDEO_L2_RAW:
+    return lambda model_hparams, vocab_size: "video_modality_l2_raw"
+  elif modality_type == ModalityType.VIDEO_PIXEL_NOISE:
+    return lambda model_hparams, vocab_size: "video_modality_pixel_noise"
+  elif modality_type in (ModalityType.CLASS_LABEL,
+                         ModalityType.MULTI_LABEL,
+                         ModalityType.ONE_HOT_CLASS_LABEL):
+    def name(model_hparams, vocab_size):
+      return "class_label_modality_%d_%d" % (vocab_size,
+                                             model_hparams.hidden_size)
+    return name
+  elif modality_type in (ModalityType.CTC_SYMBOL,
+                         ModalityType.IDENTITY_SYMBOL,
+                         ModalityType.SYMBOL,
+                         ModalityType.SYMBOL_WEIGHTS_ALL,
+                         ModalityType.SYMBOL_ONE_HOT):
+    def name(model_hparams, vocab_size):
+      return "symbol_modality_%d_%d" % (vocab_size, model_hparams.hidden_size)
+    return name
+  elif modality_type == ModalityType.SIGMOID_CLASS_LABEL:
+    def name(model_hparams, vocab_size):
+      return "sigmoid_class_symbol_modality_%d_%d" % (vocab_size,
+                                                      model_hparams.hidden_size)
+    return name
+  elif modality_type == ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL:
+    def name(model_hparams, vocab_size):
+      return "sigmoid_max_pooling_class_symbol_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)
+    return name
+  elif modality_type == ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL:
+    def name(model_hparams, vocab_size):
+      return "softmax_average_pooling_onehot_class_label_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)
+    return name
+  elif modality_type == ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL:
+    def name(model_hparams, vocab_size):
+      return "softmax_last_timestep_onehot_class_label_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)
+    return name
+  elif modality_type == ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL:
+    def name(model_hparams, vocab_size):
+      return "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)
+    return name
   return value
 
 
 def get_targets_bottom(modality_type, value=None):
   """Gets default bottom transformation for targets; if none, return value."""
-  if modality_type in ModalityType.get_choices():
-    modality_cls = getattr(current_module, modality_type)
-    return modality_cls.targets_bottom
+  if modality_type == ModalityType.AUDIO:
+    return make_targets_bottom(audio_bottom)
+  elif modality_type == ModalityType.AUDIO_SPECTRAL:
+    return make_targets_bottom(audio_spectral_bottom)
+  elif modality_type in (ModalityType.CLASS_LABEL,
+                         ModalityType.MULTI_LABEL,
+                         ModalityType.ONE_HOT_CLASS_LABEL,
+                         ModalityType.SIGMOID_CLASS_LABEL,
+                         ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
+                         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
+                         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
+                         ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL):
+    return class_label_targets_bottom
+  elif modality_type in (ModalityType.CTC_SYMBOL,
+                         ModalityType.SYMBOL,
+                         ModalityType.SYMBOL_WEIGHTS_ALL):
+    return symbol_targets_bottom
+  elif modality_type in (ModalityType.GENERIC_L2_LOSS,
+                         ModalityType.IDENTITY_SYMBOL):
+    return identity_bottom
+  elif modality_type == ModalityType.IDENTITY:
+    return make_targets_bottom(identity_bottom)
+  elif modality_type == ModalityType.IMAGE:
+    return image_targets_bottom
+  elif modality_type in (ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
+                         ModalityType.IMAGE_CHANNEL_COMPRESS):
+    return image_channel_compress_targets_bottom
+  elif modality_type == ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM:
+    return image_channel_embeddings_bottom
+  elif modality_type in (ModalityType.REAL,
+                         ModalityType.REAL_L2_LOSS,
+                         ModalityType.REAL_LOG_POISSON_LOSS):
+    return make_targets_bottom(real_bottom)
+  elif modality_type == ModalityType.SPEECH_RECOGNITION:
+    return make_targets_bottom(speech_recognition_bottom)
+  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
+    return symbol_one_hot_bottom
+  elif modality_type in (ModalityType.VIDEO,
+                         ModalityType.VIDEO_L1,
+                         ModalityType.VIDEO_L2):
+    return video_targets_bottom
+  elif modality_type == ModalityType.VIDEO_BITWISE:
+    return video_bitwise_targets_bottom
+  elif modality_type == ModalityType.VIDEO_IDENTITY:
+    return video_identity_targets_bottom
+  elif modality_type in (ModalityType.VIDEO_L1_RAW,
+                         ModalityType.VIDEO_L2_RAW):
+    return video_raw_targets_bottom
+  elif modality_type == ModalityType.VIDEO_PIXEL_NOISE:
+    return make_targets_bottom(video_pixel_noise_bottom)
   return value
 
 
-def get_targets_weights_fn(modality_type, value=None):
-  """Gets default weights function; if none available, return value."""
-  if modality_type in ModalityType.get_choices():
-    modality_cls = getattr(current_module, modality_type)
-    return modality_cls.targets_weights_fn
+def get_top(modality_type, value=None):
+  """Gets default top transformation; if none available, return value."""
+  if modality_type in (ModalityType.AUDIO,
+                       ModalityType.AUDIO_SPECTRAL,
+                       ModalityType.GENERIC_L2_LOSS,
+                       ModalityType.IDENTITY,
+                       ModalityType.IDENTITY_SYMBOL,
+                       ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
+                       ModalityType.SPEECH_RECOGNITION,
+                       ModalityType.VIDEO_IDENTITY):
+    return identity_top
+  elif modality_type in (ModalityType.CLASS_LABEL,
+                         ModalityType.MULTI_LABEL,
+                         ModalityType.ONE_HOT_CLASS_LABEL,
+                         ModalityType.SIGMOID_CLASS_LABEL):
+    return class_label_top
+  elif modality_type in (ModalityType.CTC_SYMBOL,
+                         ModalityType.SYMBOL,
+                         ModalityType.SYMBOL_WEIGHTS_ALL):
+    return symbol_top
+  elif modality_type == ModalityType.IMAGE:
+    return image_top
+  elif modality_type == ModalityType.IMAGE_CHANNEL_COMPRESS:
+    return image_channel_compress_top
+  elif modality_type == ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM:
+    return image_channel_embeddings_top
+  elif modality_type in (ModalityType.REAL,
+                         ModalityType.REAL_L2_LOSS,
+                         ModalityType.REAL_LOG_POISSON_LOSS):
+    return real_top
+  elif modality_type == ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL:
+    return sigmoid_max_pooling_class_label_top
+  elif modality_type == ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL:
+    return softmax_average_pooling_class_label_top
+  elif modality_type == ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL:
+    return softmax_last_timestep_class_label_top
+  elif modality_type == ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL:
+    return softmax_max_pooling_class_label_top
+  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
+    return symbol_one_hot_top
+  elif modality_type in (ModalityType.VIDEO,
+                         ModalityType.VIDEO_BITWISE,
+                         ModalityType.VIDEO_PIXEL_NOISE):
+    return video_top
+  elif modality_type in (ModalityType.VIDEO_L1,
+                         ModalityType.VIDEO_L2):
+    return video_l1_top
+  elif modality_type in (ModalityType.VIDEO_L1_RAW,
+                         ModalityType.VIDEO_L2_RAW):
+    return video_raw_top
   return value
 
 
-def get_top(modality_type, value=None):
-  """Gets default top transformation; if none available, return value."""
-  if modality_type in ModalityType.get_choices():
-    modality_cls = getattr(current_module, modality_type)
-    return modality_cls.top
+def get_weights_fn(modality_type, value=None):
+  """Gets default weights function; if none available, return value."""
+  if modality_type in (ModalityType.CTC_SYMBOL,
+                       ModalityType.IDENTITY_SYMBOL,
+                       ModalityType.MULTI_LABEL,
+                       ModalityType.SYMBOL,
+                       ModalityType.SYMBOL_ONE_HOT):
+    return common_layers.weights_nonzero
+  elif modality_type in ModalityType.get_choices():
+    return common_layers.weights_all
   return value
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 6cb6efc98..69504bb8e 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -31,6 +31,26 @@
 
 class ModalityTest(tf.test.TestCase):
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testGetForAllModalities(self):
+    for modality in modalities.ModalityType.get_choices():
+      bottom = modalities.get_bottom(modality)
+      loss = modalities.get_loss(modality)
+      name = modalities.get_name(modality)
+      targets_bottom = modalities.get_targets_bottom(modality)
+      top = modalities.get_top(modality)
+      weights_fn = modalities.get_weights_fn(modality)
+      self.assertIsNotNone(bottom,
+                           msg="{} has no default bottom".format(modality))
+      self.assertIsNotNone(loss, msg="{} has no default loss".format(modality))
+      self.assertIsNotNone(name, msg="{} has no default name".format(modality))
+      self.assertIsNotNone(
+          targets_bottom,
+          msg="{} has no default targets_bottom".format(modality))
+      self.assertIsNotNone(top, msg="{} has no default top".format(modality))
+      self.assertIsNotNone(weights_fn,
+                           msg="{} has no default weights_fn".format(modality))
+
   @test_utils.run_in_graph_and_eager_modes()
   def testSymbolModalityInputs(self):
     batch_size = 10
@@ -87,7 +107,7 @@ def testSymbolModalityTargets(self):
         sharded_targets,
         model_hparams,
         vocab_size,
-        modalities.get_targets_weights_fn(modalities.ModalityType.SYMBOL))
+        modalities.get_weights_fn(modalities.ModalityType.SYMBOL))
     train_loss = (tf.add_n(sharded_loss_num) /
                   tf.maximum(1.0, tf.add_n(sharded_loss_den)))
     logits = tf.concat(sharded_logits, 0)
@@ -124,12 +144,12 @@ def testSymbolModalityTargetsFactored(self):
           model_hparams,
           vocab_size)
       sharded_loss_num, sharded_loss_den = data_parallelism(
-          modalities.SymbolModality.loss,
+          modalities.get_loss(modalities.ModalityType.SYMBOL),
           sharded_logits,
           sharded_targets,
           model_hparams,
           vocab_size,
-          modalities.get_targets_weights_fn(modalities.ModalityType.SYMBOL))
+          modalities.get_weights_fn(modalities.ModalityType.SYMBOL))
       train_loss = (tf.add_n(sharded_loss_num) /
                     tf.maximum(1.0, tf.add_n(sharded_loss_den)))
       logits = tf.concat(sharded_logits, 0)
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 81d8de434..d43f96c15 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -48,16 +48,11 @@ def body(self, features):
     hparams = copy.copy(self._hparams)
     targets = features["targets"]
     if (hparams.likelihood == cia.DistributionType.DMOL and
-        (hparams.modality["targets"] !=
-         modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY or
-         hparams.num_channels != 1)):
-      raise ValueError("When using DMOL for the likelihood,modality['targets'] "
-                       "must be ImageChannelBottomIdentityModality and "
-                       "num_channels must be 1.")
+        hparams.num_channels != 1):
+      raise ValueError("When using DMOL for the likelihood, bottom function "
+                       " must be identity and num_channels must be 1.")
     if (not tf.get_variable_scope().reuse and
-        hparams.mode != tf.estimator.ModeKeys.PREDICT and
-        hparams.modality["targets"] !=
-        modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY):
+        hparams.mode != tf.estimator.ModeKeys.PREDICT):
       tf.summary.image("targets", tf.to_float(targets), max_outputs=1)
 
     # Extra losses list if we want to use moe.
@@ -193,7 +188,8 @@ def image_transformer_base():
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.98
   hparams.label_smoothing = 0.0
-  hparams.modality["targets"] = modalities.ModalityType.IDENTITY
+  hparams.bottom["targets"] = modalities.image_channel_embeddings_bottom
+  hparams.top["targets"] = modalities.identity_top
   hparams.norm_type = "layer"
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.add_hparam("filter_size", 512)  # Add new ones like this.
@@ -280,8 +276,8 @@ def imagetransformer_cifar10_base_dmol():
   hparams = image_transformer_base()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.modality["targets"] = (
-      modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY)
+  hparams.bottom["targets"] = modalities.image_channel_compress_targets_bottom
+  hparams.top["targets"] = modalities.identity_top
   hparams.num_heads = 8
   hparams.batch_size = 8
   hparams.sampling_method = "random"
@@ -422,8 +418,8 @@ def imagetransformerpp_sep_channels_8l_8h():
   hparams = imagetransformer_base()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.modality["targets"] = (
-      modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY)
+  hparams.bottom["targets"] = modalities.image_channel_compress_targets_bottom
+  hparams.top["targets"] = modalities.identity_top
   hparams.num_heads = 8
   hparams.batch_size = 4
   hparams.attention_key_channels = hparams.attention_value_channels = 0
@@ -886,8 +882,8 @@ def imagetransformerpp_tiny():
   hparams = imagetransformer_tiny()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.modality["targets"] = (
-      modalities.ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY)
+  hparams.bottom["targets"] = modalities.image_channel_compress_targets_bottom
+  hparams.top["targets"] = modalities.identity_top
   return hparams
 
 
diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index acf793440..80ce13d53 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -382,7 +382,9 @@ def image_transformer2d_base():
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.98
   hparams.label_smoothing = 0.0
-  hparams.modality["targets"] = modalities.ModalityType.IDENTITY
+  hparams.bottom["targets"] = modalities.make_targets_bottom(
+      modalities.image_channel_embeddings_bottom)
+  hparams.top["targets"] = modalities.identity_top
   hparams.norm_type = "layer"
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.add_hparam("filter_size", 512)  # Add new ones like this.
@@ -593,6 +595,7 @@ def img2img_transformer2d_base():
   hparams.filter_size = 2048
   hparams.num_encoder_layers = 4
   hparams.num_decoder_layers = 8
+  hparams.bottom["inputs"] = modalities.image_channel_embeddings_bottom
   hparams.dec_attention_type = cia.AttentionType.LOCAL_2D
   hparams.block_raster_scan = True
   return hparams
diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 3e32e9c03..b8d4c56d8 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -854,9 +854,12 @@ def mtf_transformer_base():
   # These parameters make Transformer model compatible with MtfTransformer
   # Do not override these, as mtf_transformer does not support other options.
   hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
-  hparams.modality = {
-      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
-      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
+  hparams.bottom = {
+      "inputs": modalities.identity_bottom,
+      "targets": modalities.identity_bottom,
+  }
+  hparams.top = {
+      "targets": modalities.identity_top,
   }
 
   # Parameters for computing the maximum decode length in beam search.
diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index 1d1e2a71d..c8cb5354d 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -418,9 +418,12 @@ def mtf_transformer2_base():
   hparams.use_fixed_batch_size = True
   hparams.add_hparam("mtf_mode", True)
   hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
-  hparams.modality = {
-      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
-      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
+  hparams.bottom = {
+      "inputs": modalities.identity_bottom,
+      "targets": modalities.identity_bottom,
+  }
+  hparams.top = {
+      "targets": modalities.identity_top,
   }
   hparams.add_hparam("beam_size", 1)
 
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 02dde9e03..cf4305a30 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1112,9 +1112,12 @@ def autoencoder_residual_text():
   hparams.hidden_size = 64
   hparams.max_hidden_size = 512
   hparams.bottleneck_noise = 0.0
-  hparams.modality = {
-      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
-      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
+  hparams.bottom = {
+      "inputs": modalities.identity_bottom,
+      "targets": modalities.identity_bottom,
+  }
+  hparams.top = {
+      "targets": modalities.identity_top,
   }
   hparams.autoregressive_mode = "none"
   hparams.sample_width = 1
@@ -1219,9 +1222,12 @@ def autoencoder_ordered_text():
   hparams.batch_size = 1024
   hparams.autoregressive_mode = "conv5"
   hparams.max_hidden_size = 1024
-  hparams.modality = {
-      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
-      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
+  hparams.bottom = {
+      "inputs": modalities.identity_bottom,
+      "targets": modalities.identity_bottom,
+  }
+  hparams.top = {
+      "targets": modalities.identity_top,
   }
   hparams.sample_height = 128
   hparams.sample_width = 1
diff --git a/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index 94842b2db..3ba2e64ab 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -127,9 +127,12 @@ def cycle_gan_small():
   """Set of hyperparameters."""
   hparams = transformer_vae.transformer_ae_small()
   hparams.batch_size = 2048
-  hparams.modality = {
-      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
-      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
+  hparams.bottom = {
+      "inputs": modalities.identity_bottom,
+      "targets": modalities.identity_bottom,
+  }
+  hparams.top = {
+      "targets": modalities.identity_top,
   }
   hparams.weight_decay = 3.0
   hparams.learning_rate = 0.05
diff --git a/tensor2tensor/models/research/super_lm.py b/tensor2tensor/models/research/super_lm.py
index 5a02ccf94..0e7dd58c8 100644
--- a/tensor2tensor/models/research/super_lm.py
+++ b/tensor2tensor/models/research/super_lm.py
@@ -265,7 +265,13 @@ def super_lm_base():
   # we only want one data shard.
   hparams.no_data_parallelism = True
   # bypass the symbol modality so that we can use model parallelism.
-  hparams.modality["targets"] = modalities.ModalityType.IDENTITY_SYMBOL
+  hparams.bottom = {
+      "inputs": modalities.identity_bottom,
+      "targets": modalities.identity_bottom,
+  }
+  hparams.top = {
+      "targets": modalities.identity_top,
+  }
   hparams.add_hparam("filter_size", 512)
   hparams.add_hparam("mix_fraction", 0.5)
   # attention-related flags
diff --git a/tensor2tensor/models/research/transformer_symshard.py b/tensor2tensor/models/research/transformer_symshard.py
index 98c9cf0ee..7f5bf39dd 100644
--- a/tensor2tensor/models/research/transformer_symshard.py
+++ b/tensor2tensor/models/research/transformer_symshard.py
@@ -367,9 +367,12 @@ def transformer_symshard_base():
   # we only want one data shard.
   hparams.no_data_parallelism = True
   # bypass the symbol modality so that we can use model parallelism.
-  hparams.modality = {
-      "inputs": modalities.ModalityType.IDENTITY_SYMBOL,
-      "targets": modalities.ModalityType.IDENTITY_SYMBOL,
+  hparams.bottom = {
+      "inputs": modalities.identity_bottom,
+      "targets": modalities.identity_bottom,
+  }
+  hparams.top = {
+      "targets": modalities.identity_top,
   }
   hparams.add_hparam("filter_size", 1280)
   hparams.add_hparam("mix_fraction", 0.5)
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index e8f769bb6..9cfe80933 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -896,7 +896,8 @@ def imagetransformer_ae_cifar():
 
   hparams.add_hparam("unconditional", False)  # unconditional generation
 
-  hparams.modality["targets"] = modalities.ImageChannelEmbeddingsBottom
+  hparams.bottom["targets"] = modalities.image_channel_embeddings_bottom
+  hparams.top["targets"] = modalities.image_channel_embeddings_top
   hparams.drop_inputs = True
   hparams.do_attend_compress = False
   hparams.do_attend_decompress = False
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 637cb7ccc..dbae5eb1a 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -195,16 +195,12 @@ def __init__(self, *args, **kwargs):
 
   @property
   def _target_modality(self):
-    target_modality = self.hparams.modality.get(
-        "targets",
-        self.problem_hparams.modality["targets"])
-    if target_modality not in modalities.ModalityType.get_choices():
-      target_modality = target_modality.__class__.__name__
-    return target_modality
+    return self.problem_hparams.modality["targets"]
 
   @property
   def is_per_pixel_softmax(self):
-    return self._target_modality == modalities.ModalityType.VIDEO
+    # TODO(trandustin): This is a hack.
+    return "targets" not in self.hparams.get("loss")
 
   def get_iteration_num(self):
     step_num = tf.train.get_global_step()
@@ -337,11 +333,12 @@ def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
       Additional reconstruction loss.
 
     Raises:
-      ValueError: in case of unknown modality.
+      ValueError: in case of unknown loss transformation.
     """
-    if self._target_modality == modalities.ModalityType.VIDEO_L2_RAW:
+    # TODO(trandustin): This logic should be moved elsewhere.
+    if self.hparams.loss.get("targets") == modalities.video_l2_raw_loss:
       recon_loss = tf.losses.mean_squared_error(extra_gts, extra_pds)
-    elif self._target_modality == modalities.ModalityType.VIDEO:
+    elif "targets" not in self.hparams.loss:
       shape = common_layers.shape_list(extra_pds)
       updated_shape = shape[:-1] + [3, 256]
       extra_pds = tf.reshape(extra_pds, updated_shape)
@@ -350,10 +347,9 @@ def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
       targets = extra_raw_gts
       targets_shape = common_layers.shape_list(targets)
       targets = tf.reshape(targets, [-1] + targets_shape[2:])
-      modality = self.hparams.problem_hparams.modality["targets"]
-      targets_weights_fn = self.hparams.targets_weights_fn.get(
+      targets_weights_fn = self.hparams.weights_fn.get(
           "targets",
-          modalities.get_targets_weights_fn(modality))
+          modalities.get_weights_fn(self._target_modality))
       numerator, denominator = common_layers.padded_cross_entropy(
           logits,
           targets,
@@ -362,7 +358,7 @@ def get_extra_internal_loss(self, extra_raw_gts, extra_gts, extra_pds):
           weights_fn=targets_weights_fn)
       recon_loss = numerator / denominator
     else:
-      raise ValueError("internal loss only supports specific modalities.")
+      raise ValueError("internal loss only supports specific hparams.loss.")
     tf.summary.scalar("recon_extra", recon_loss)
     return recon_loss
 
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 7616d3549..8d8a2c6ff 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -55,7 +55,8 @@ def next_frame_pixel_noise():
   """Basic 2-frame conv model with pixel noise."""
   hparams = next_frame_basic_deterministic()
   hparams.add_hparam("video_modality_input_noise", 0.05)
-  hparams.modality["inputs"] = modalities.ModalityType.VIDEO_PIXEL_NOISE
+  hparams.bottom["inputs"] = modalities.video_pixel_noise_bottom
+  hparams.top["inputs"] = modalities.video_top
   return hparams
 
 
@@ -89,7 +90,8 @@ def next_frame_tpu():
 def next_frame_ae():
   """Conv autoencoder."""
   hparams = next_frame_basic_deterministic()
-  hparams.modality["inputs"] = modalities.ModalityType.VIDEO_BITWISE
+  hparams.bottom["inputs"] = modalities.video_bitwise_bottom
+  hparams.top["inputs"] = modalities.video_top
   hparams.hidden_size = 256
   hparams.batch_size = 8
   hparams.num_hidden_layers = 4
@@ -102,7 +104,8 @@ def next_frame_ae():
 def next_frame_ae_tiny():
   """Conv autoencoder, tiny set for testing."""
   hparams = next_frame_tiny()
-  hparams.modality["inputs"] = modalities.ModalityType.VIDEO_BITWISE
+  hparams.bottom["inputs"] = modalities.video_bitwise_bottom
+  hparams.top["inputs"] = modalities.video_top
   hparams.batch_size = 8
   hparams.dropout = 0.4
   return hparams
@@ -131,7 +134,8 @@ def next_frame_tiny():
 def next_frame_l1():
   """Basic conv model with L1 modality."""
   hparams = next_frame_basic_deterministic()
-  hparams.modality["targets"] = modalities.ModalityType.VIDEO_L1
+  hparams.loss["targets"] = modalities.video_l1_loss
+  hparams.top["targets"] = modalities.video_l1_top
   hparams.video_modality_loss_cutoff = 2.4
   return hparams
 
@@ -140,7 +144,8 @@ def next_frame_l1():
 def next_frame_l2():
   """Basic conv model with L2 modality."""
   hparams = next_frame_basic_deterministic()
-  hparams.modality["targets"] = modalities.ModalityType.VIDEO_L2
+  hparams.loss["targets"] = modalities.video_l2_loss
+  hparams.top["targets"] = modalities.video_l1_top
   hparams.video_modality_loss_cutoff = 2.4
   return hparams
 
diff --git a/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
index a8b8f98c9..e42c94d9d 100644
--- a/tensor2tensor/models/video/epva_params.py
+++ b/tensor2tensor/models/video/epva_params.py
@@ -29,9 +29,15 @@ def next_frame_epva():
   hparams = basic_deterministic_params.next_frame_basic_deterministic()
   hparams.video_num_input_frames = 4
   hparams.video_num_target_frames = 4
-  hparams.modality = {
-      "inputs": modalities.ModalityType.VIDEO_L2_RAW,
-      "targets": modalities.ModalityType.VIDEO_L2_RAW,
+  hparams.bottom = {
+      "inputs": modalities.video_raw_bottom,
+      "targets": modalities.video_raw_targets_bottom,
+  }
+  hparams.loss = {
+      "targets": modalities.video_l2_raw_loss,
+  }
+  hparams.top = {
+      "targets": modalities.video_raw_top,
   }
   hparams.learning_rate_schedule = "constant"
   hparams.learning_rate_constant = 1e-05
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index e734ed37f..cc29c6b02 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -69,9 +69,15 @@ def next_frame_glow_hparams():
   # Pretrains the glow encoder for "pretrain_steps" number of steps.
   # By default, don't pretrain and learn end-to-end
   hparams.add_hparam("pretrain_steps", -1)
-  hparams.modality = {
-      "inputs": modalities.ModalityType.VIDEO_L1_RAW,
-      "targets": modalities.ModalityType.VIDEO_L1_RAW,
+  hparams.bottom = {
+      "inputs": modalities.video_raw_bottom,
+      "targets": modalities.video_raw_targets_bottom,
+  }
+  hparams.loss = {
+      "targets": modalities.video_l1_raw_loss,
+  }
+  hparams.top = {
+      "targets": modalities.video_raw_top,
   }
   hparams.init_batch_size = 256
   hparams.batch_size = 32
diff --git a/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index 708c888f3..c28c4378a 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -36,9 +36,15 @@ def next_frame_savp():
   hparams.add_hparam("gan_loss_multiplier", 0.01)
   hparams.add_hparam("gan_vae_loss_multiplier", 0.01)
   hparams.add_hparam("gan_optimization", "joint")
-  hparams.modality = {
-      "inputs": modalities.ModalityType.VIDEO_L1_RAW,
-      "targets": modalities.ModalityType.VIDEO_L1_RAW,
+  hparams.bottom = {
+      "inputs": modalities.video_raw_bottom,
+      "targets": modalities.video_raw_targets_bottom,
+  }
+  hparams.loss = {
+      "targets": modalities.video_l1_raw_loss,
+  }
+  hparams.top = {
+      "targets": modalities.video_raw_top,
   }
   hparams.latent_loss_multiplier_schedule = "linear"
   hparams.upsample_method = "bilinear_upsample_conv"
@@ -54,9 +60,8 @@ def next_frame_savp():
 def next_frame_savp_l2():
   """SAVP with L2 reconstruction loss."""
   hparams = next_frame_savp()
-  hparams.modality = {
-      "inputs": modalities.ModalityType.VIDEO_L2_RAW,
-      "targets": modalities.ModalityType.VIDEO_L2_RAW,
+  hparams.loss = {
+      "targets": modalities.video_l2_raw_loss,
   }
   return hparams
 
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index f5687865b..f7352c714 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -33,9 +33,15 @@ def next_frame_sv2p():
   hparams.video_num_input_frames = 1
   hparams.video_num_target_frames = 3
   hparams.batch_size = 16
-  hparams.modality = {
-      "inputs": modalities.ModalityType.VIDEO_L2_RAW,
-      "targets": modalities.ModalityType.VIDEO_L2_RAW,
+  hparams.bottom = {
+      "inputs": modalities.video_raw_bottom,
+      "targets": modalities.video_raw_targets_bottom,
+  }
+  hparams.loss = {
+      "targets": modalities.video_l2_raw_loss,
+  }
+  hparams.top = {
+      "targets": modalities.video_raw_top,
   }
   hparams.video_modality_loss_cutoff = 0.0
   hparams.scheduled_sampling_mode = "count"
@@ -91,10 +97,9 @@ def next_frame_sv2p_atari():
 def next_frame_sv2p_atari_softmax():
   """SV2P model for atari with softmax."""
   hparams = next_frame_sv2p_atari()
-  hparams.modality = {
-      "inputs": modalities.ModalityType.VIDEO,
-      "targets": modalities.ModalityType.VIDEO,
-  }
+  hparams.bottom = {}
+  hparams.loss = {}
+  hparams.top = {}
   hparams.internal_loss = True
   return hparams
 
diff --git a/tensor2tensor/utils/hparam.py b/tensor2tensor/utils/hparam.py
index 6d056438a..130dbf685 100644
--- a/tensor2tensor/utils/hparam.py
+++ b/tensor2tensor/utils/hparam.py
@@ -542,7 +542,8 @@ def to_json(self, indent=None, separators=None, sort_keys=False):
       A JSON string.
     """
     return json.dumps(
-        self.values(),
+        {k: v.__name__ if callable(v) else v
+         for k, v in six.iteritems(self.values())},
         indent=indent,
         separators=separators,
         sort_keys=sort_keys)
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 52f9cdebf..91aaad345 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -614,9 +614,9 @@ def weights_fn_for_mp(problem_task_id):
       tm = {"targets": tm}
 
     for target_name, modality in six.iteritems(tm):
-      weights_fn = model_hparams.targets_weights_fn.get(
+      weights_fn = model_hparams.weights_fn.get(
           "targets",
-          modalities.get_targets_weights_fn(modality))
+          modalities.get_weights_fn(modality))
       if hasattr(model_hparams.problem, "task_list"):
         ptid = problem_instance.task_id  # pylint: disable=cell-var-from-loop
         weights_fn = weights_fn_for_mp(ptid)
@@ -643,9 +643,9 @@ def create_eager_metrics_for_problem(problem, model_hparams):
   metric_fns = problem.eval_metric_fns(model_hparams)
   problem_hparams = problem.get_hparams(model_hparams)
   target_modality = problem_hparams.modality["targets"]
-  weights_fn = model_hparams.targets_weights_fn.get(
+  weights_fn = model_hparams.weights_fn.get(
       "targets",
-      modalities.get_targets_weights_fn(target_modality))
+      modalities.get_weights_fn(target_modality))
   return create_eager_metrics_internal(metric_fns, weights_fn=weights_fn)
 
 
diff --git a/tensor2tensor/utils/modality.py b/tensor2tensor/utils/modality.py
deleted file mode 100644
index 09a83d050..000000000
--- a/tensor2tensor/utils/modality.py
+++ /dev/null
@@ -1,121 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Modality base class - defines the bottom and top of the model."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.layers import common_attention
-from tensor2tensor.layers import common_layers
-from tensor2tensor.utils import misc_utils
-
-import tensorflow as tf
-
-
-class Modality(object):
-  """Abstract Modality class for data transformations.
-
-  An abstract class representing modalities for transforming data to a space
-  interpretable by T2T models. It has 4 functions:
-  * bottom: called on inputs entering the model.
-  * targets_bottom: called on targets entering the model (e.g., the decoder).
-  * top: called on model outputs to generate predictions (e.g., logits).
-  * loss: called on predictions (outputs of top) and targets.
-
-  For example, think about a modality for images:
-  * `bottom` represents the part of the model applied to an incoming image,
-    e.g., an entry flow of a convolutional network.
-  * `top` represents the top part of a model that is generating images, e.g., a
-    PixelCNN network.
-  * `targets_bottom` represents the auto-regressive part of the network.  It is
-    applied to the already-generated part of an image, which is given to the
-    decoder to generate the next part. In some cases, e.g., for text, it is the
-    same as the `bottom` function, and that is the default we use. But, e.g.,
-    for images, a different function might be needed to regress properly.
-  * `loss` would compare the generated image to the target image and score it.
-  """
-
-  def __init__(self, model_hparams, vocab_size=None):
-    # __init__ args are unused in any methods. They're maintained for
-    # backwards compatibility for now. In the future, Modality classes will be
-    # removed altogether.
-    del model_hparams, vocab_size
-
-  @classmethod
-  def name(cls, model_hparams, vocab_size=None):
-    del model_hparams, vocab_size  # unused arg
-    return misc_utils.camelcase_to_snakecase(type(cls).__name__)
-
-  targets_weights_fn = staticmethod(common_layers.weights_all)
-
-  @staticmethod
-  def bottom(x, model_hparams, vocab_size=None):
-    """Transform one shard of input.
-
-    Args:
-      x: An int32 Tensor with shape [batch, p0, p1, input_channels]
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      A float32 Tensor with shape [batch, p0, p1, body_input_depth]
-    """
-    raise NotImplementedError("Abstract Method")
-
-  @classmethod
-  def targets_bottom(cls, x, model_hparams, vocab_size=None):
-    """Transform one shard of targets.
-
-    Args:
-      x: An int32 Tensor with shape [batch, p0, p1, target_channels]
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      A float32 Tensor with shape [batch, p0, p1, body_input_depth]
-    """
-    with tf.variable_scope("targets_bottom"):
-      return cls.bottom(x, model_hparams, vocab_size)
-
-  @staticmethod
-  def top(body_output, targets, model_hparams, vocab_size=None):
-    """Generate predictions/logits for one shard of output.
-
-    Most classes will override this function.
-
-    Args:
-      body_output: A Tensor with shape [batch, p0, p1, body_output_depth]
-      targets: A Tensor with shape [batch, p0, p1, targets_channels,
-        top_dimensionality]
-      model_hparams: tf.HParams, model hyperparmeters.
-      vocab_size: int, vocabulary size.
-
-    Returns:
-      A Tensor of class logits.
-    """
-    raise NotImplementedError("Abstract Method")
-
-  @staticmethod
-  def loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-    """Compute loss numerator and denominator for one shard of output."""
-    del vocab_size  # unused arg
-    logits = top_out
-    logits = common_attention.maybe_upcast(logits, hparams=model_hparams)
-    return common_layers.padded_cross_entropy(
-        logits,
-        targets,
-        model_hparams.label_smoothing,
-        weights_fn=weights_fn)
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index f81caf5e2..5a141f8ae 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -204,47 +204,19 @@ def __init__(self,
           },
           hparams=hparams)
 
-    # TODO(trandustin): For now, we get custom feature transformations via
-    # hparams.modality. Once modality classes are removed, let users
-    # individually specify custom transformations for bottom, loss, etc.
-    if not hasattr(hparams, "modality"):
-      hparams.add_hparam("modality", {})
-    if not hasattr(hparams, "bottom"):
-      hparams.add_hparam("bottom", {})
-    if not hasattr(hparams, "loss"):
-      hparams.add_hparam("loss", {})
-    if not hasattr(hparams, "name"):
-      hparams.add_hparam("name", {})
-    if not hasattr(hparams, "targets_weights_fn"):
-      hparams.add_hparam("targets_weights_fn", {})
-    if not hasattr(hparams, "top"):
-      hparams.add_hparam("top", {})
-    target_modalities = _create_target_modality(hparams.modality)
-    for feature_name, modality in six.iteritems(hparams.modality):
-      if modality in modalities.ModalityType.get_choices():
-        modality = getattr(modalities, modality)
-      if feature_name in target_modalities:
-        hparams.bottom[feature_name] = modality.targets_bottom
-      else:
-        hparams.bottom[feature_name] = modality.bottom
-      hparams.loss[feature_name] = modality.loss
-      hparams.name[feature_name] = modality.name
-      hparams.targets_weights_fn[feature_name] = modality.targets_weights_fn
-      hparams.top[feature_name] = modality.top
-
     if self._problem_hparams:
       for feature_name, modality in six.iteritems(
           self._problem_hparams.modality):
-        # If prepend mode, set targets_weights_fn to appropriately handle it.
-        if (modality in (modalities.ModalityType.SYMBOL,
-                         modalities.ModalityType.SYMBOL_ONE_HOT,
-                         modalities.ModalityType.CTC_SYMBOL,
-                         modalities.ModalityType.IDENTITY_SYMBOL)):
+        # If prepend mode, set weights_fn to appropriately handle it.
+        if (modality in (modalities.ModalityType.CTC_SYMBOL,
+                         modalities.ModalityType.IDENTITY_SYMBOL,
+                         modalities.ModalityType.SYMBOL,
+                         modalities.ModalityType.SYMBOL_ONE_HOT)):
           if (hparams.prepend_mode == "prepend_inputs_full_attention" or
               (hparams.prepend_mode == "prepend_inputs_masked_attention" and
                mode != tf.estimator.ModeKeys.TRAIN)):
             weights_fn = common_layers.weights_prepend_inputs_to_targets
-            hparams.targets_weights_fn[feature_name] = weights_fn
+            hparams.weights_fn[feature_name] = weights_fn
 
     self._original_hparams = hparams
     self.set_mode(mode)
@@ -646,8 +618,8 @@ def _loss_single(self, logits, feature_name, feature, weights=None):
     if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
       vocab_size += (-vocab_size) % self._hparams.vocab_divisor
     loss = self._hparams.loss.get(feature_name, modalities.get_loss(modality))
-    targets_weights_fn = self._hparams.targets_weights_fn.get(
-        "targets", modalities.get_targets_weights_fn(modality))
+    targets_weights_fn = self._hparams.weights_fn.get(
+        "targets", modalities.get_weights_fn(modality))
     if weights is None:
       loss_num, loss_den = loss(logits, feature, self._hparams, vocab_size,
                                 weights_fn=targets_weights_fn)
@@ -1817,7 +1789,7 @@ def create_tpu_eval_metrics_fn(problem, model_hparams):
   tm = _create_target_modality(problem.get_hparams(model_hparams).modality)
   if isinstance(tm, dict):
     for k, v in six.iteritems(tm):
-      weights_fn = modalities.get_targets_weights_fn(v)
+      weights_fn = modalities.get_weights_fn(v)
 
       def make_metric_fn(metric_fn):
         def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
@@ -1837,7 +1809,7 @@ def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
         name = "%s/metrics-%s/%s" % (k, problem.name, metric)
         metric_fns.append((name, make_metric_fn(metric_fn)))
   else:
-    weights_fn = modalities.get_targets_weights_fn(tm)
+    weights_fn = modalities.get_weights_fn(tm)
 
     def make_metric_fn(metric_fn):
       def wrapped_metric_fn(logits, labels, features):
@@ -2038,8 +2010,8 @@ def sampled_results():
                                 vocab_size)
         if "training" not in losses:
           loss = hparams.loss.get("targets", modalities.get_loss(modality))
-          weights_fn = hparams.targets_weights_fn.get(
-              "targets", modalities.get_targets_weights_fn(modality))
+          weights_fn = hparams.weights_fn.get(
+              "targets", modalities.get_weights_fn(modality))
           sharded_loss_num, sharded_loss_den = dp(loss,
                                                   sharded_logits,
                                                   sharded_features["targets"],
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index c9b680302..229846832 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -52,6 +52,8 @@ def testLossSingleWeights(self):
 
         model_hparams = HParams(
             prepend_mode="none",
+            loss={},
+            weights_fn={},
             label_smoothing=0.0,
             shared_embedding_and_softmax_weights=False)
 

From 90ce2fec7f8fa8fda0a272afd2a048cfbf3024c1 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Sat, 2 Mar 2019 14:55:27 -0800
Subject: [PATCH 1746/2720] Lazy load t2t problems (cuts import time down
 significantly)

PiperOrigin-RevId: 236483583
---
 tensor2tensor/trax/inputs.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index b8765ecca..ab2b24478 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -26,8 +26,6 @@
 
 import jax.numpy as np
 
-from tensor2tensor import problems
-
 import tensorflow as tf
 import tensorflow_datasets as tfds
 
@@ -135,6 +133,7 @@ def _select_features(example, feature_list=None):
 
 def _train_and_eval_dataset_v1(problem_name, data_dir):
   """Return train and evaluation datasets, feature info and supervised keys."""
+  from tensor2tensor import problems  # pylint: disable=g-import-not-at-top
   problem = problems.problem(problem_name)
   train_dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, data_dir)
   train_dataset = train_dataset.map(_select_features)

From 334dd05874a4a83d372f2964e35ba304a69b4c1a Mon Sep 17 00:00:00 2001
From: Wanqi Zhu <1213.ghs@gmail.com>
Date: Sun, 3 Mar 2019 06:56:42 +0800
Subject: [PATCH 1747/2720] t2t_decoder checkpoint_path fix (#1471)

Update t2t_decoder.py to pass in checkpoint_path FLAG correctly in decode_from_dataset
---
 tensor2tensor/bin/t2t_decoder.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 61df0d4b3..768fdac7b 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -100,7 +100,8 @@ def decode(estimator, hparams, decode_hp):
         hparams,
         decode_hp,
         decode_to_file=FLAGS.decode_to_file,
-        dataset_split="test" if FLAGS.eval_use_test_set else None)
+        dataset_split="test" if FLAGS.eval_use_test_set else None,
+        checkpoint_path=FLAGS.checkpoint_path)
 
 
 def score_file(filename):

From 228f2cf21f2e883afcc1bda0d2c5e9c6f8815fe7 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 4 Mar 2019 09:18:05 -0800
Subject: [PATCH 1748/2720] * Clean up training loop to make clear what the
 different pieces are. We'll want to think about how to make this extensible.
 fast.ai uses a callback system. Estimator has hooks. Maybe we can do better,
 or maybe we use the same, but having a cleaner loop will make it easier to
 see what can/should be factored out and how. * fn->fun to follow the jax
 convention, because it's more fun :) * Use the stax/Keras-like convention of
 making "constructors" capitalized for the learning rates.

PiperOrigin-RevId: 236665210
---
 .../trax/configs/resnet50_imagenet_8gb.gin    |  22 +--
 tensor2tensor/trax/inputs.py                  |  20 +-
 tensor2tensor/trax/learning_rate.py           |  24 +--
 tensor2tensor/trax/trax.py                    | 176 +++++++++++-------
 4 files changed, 144 insertions(+), 98 deletions(-)

diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
index cae941553..953310085 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
@@ -3,31 +3,31 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fn:
+# Parameters for batch_fun:
 # ==============================================================================
-batch_fn.batch_size = 32
-batch_fn.bucket_length = 32
-batch_fn.buckets = None
-batch_fn.eval_batch_size = 32
+batch_fun.batch_size = 32
+batch_fun.bucket_length = 32
+batch_fun.buckets = None
+batch_fun.eval_batch_size = 32
 
 # Parameters for inputs:
 # ==============================================================================
 inputs.data_dir = None
 inputs.dataset_name = 't2t_image_imagenet224'
 
-# Parameters for learning_rate:
+# Parameters for DefaultSchedule:
 # ==============================================================================
-learning_rate.constant = 5.0
-learning_rate.schedule = 'constant * linear_warmup * rsqrt_decay'
-learning_rate.warmup_steps = 400
+DefaultSchedule.constant = 5.0
+DefaultSchedule.schedule = 'constant * linear_warmup * rsqrt_decay'
+DefaultSchedule.warmup_steps = 400
 
 # Parameters for momentum:
 # ==============================================================================
 momentum.mass = 0.9
 
-# Parameters for preprocess_fn:
+# Parameters for preprocess_fun:
 # ==============================================================================
-preprocess_fn.max_target_length = -1
+preprocess_fun.max_target_length = -1
 
 # Parameters for Resnet50:
 # ==============================================================================
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index ab2b24478..992e0cabc 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -31,7 +31,7 @@
 
 
 Inputs = collections.namedtuple(
-    "_Inputs", ["train_fn", "eval_fn", "input_shape"])
+    "_Inputs", ["train_stream", "eval_stream", "input_shape"])
 
 
 @gin.configurable()
@@ -53,13 +53,14 @@ def inputs(dataset_name, data_dir):
    input_name, input_shape) = train_and_eval_batches(
        dataset_name, data_dir)
 
-  def train_input_fn():
+  def train_input_fun():
     return dataset_to_stream(train_batches, input_name)
 
-  def eval_input_fn():
+  def eval_input_fun():
     return dataset_to_stream(eval_batches, input_name)
 
-  return Inputs(train_fn=train_input_fn, eval_fn=eval_input_fn,
+  return Inputs(train_stream=train_input_fun,
+                eval_stream=eval_input_fun,
                 input_shape=input_shape)
 
 
@@ -160,7 +161,7 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
 
 
 @gin.configurable(blacklist=["dataset", "training"])
-def preprocess_fn(dataset, training, max_target_length=-1):
+def preprocess_fun(dataset, training, max_target_length=-1):
   def target_right_length(_, target):
     return tf.less(tf.shape(target)[0], max_target_length + 1)
   if max_target_length > 0 and training:
@@ -169,8 +170,9 @@ def target_right_length(_, target):
 
 
 @gin.configurable(blacklist=["dataset", "training", "shapes", "target_names"])
-def batch_fn(dataset, training, shapes, target_names,
-             batch_size=32, eval_batch_size=32, bucket_length=32, buckets=None):
+def batch_fun(dataset, training, shapes, target_names,
+              batch_size=32, eval_batch_size=32,
+              bucket_length=32, buckets=None):
   """Batching function."""
   del target_names
   # If bucketing is not specified, check if target shapes are variable.
@@ -220,8 +222,8 @@ def append_targets(example):
   shapes = {k: features_info[k].shape for k in features_info}
   shapes = (shapes, shapes[target_names[0]])
   dataset = dataset.shuffle(1024)
-  dataset = preprocess_fn(dataset, training)
-  dataset = batch_fn(dataset, training, shapes, target_names)
+  dataset = preprocess_fun(dataset, training)
+  dataset = batch_fun(dataset, training, shapes, target_names)
   return dataset.prefetch(32)
 
 
diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index cb81f84af..0bd526362 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -24,10 +24,10 @@
 
 
 @gin.configurable(blacklist=["history"])
-def make_default_schedule(history=None,
-                          schedule="constant * linear_warmup * rsqrt_decay",
-                          constant=0.001,
-                          warmup_steps=100):
+def DefaultSchedule(history=None,
+                    schedule="constant * linear_warmup * rsqrt_decay",
+                    constant=0.001,
+                    warmup_steps=100):
   """Default learning rate  schedule.
 
   Note: the learning rate schedule takes arguments and return a function,
@@ -53,7 +53,7 @@ def make_default_schedule(history=None,
   del history
   factors = [n.strip() for n in schedule.split("*")]
 
-  def learning_rate(step):
+  def learning_rate(step):  # pylint: disable=invalid-name
     """Step to learning rate function."""
     ret = 1.0
     for name in factors:
@@ -71,12 +71,12 @@ def learning_rate(step):
 
 
 @gin.configurable(blacklist=["history"])
-def make_eval_adjusting_schedule(history,
-                                 constant=0.001,
-                                 steps_to_decrease=10,
-                                 improvement_margin=0.01,
-                                 decrease_rate=2.0,
-                                 metric="metrics/accuracy"):
+def EvalAdjustingSchedule(history,
+                          constant=0.001,
+                          steps_to_decrease=10,
+                          improvement_margin=0.01,
+                          decrease_rate=2.0,
+                          metric="metrics/accuracy"):
   """Learning rate that decreases when eval metric stalls.
 
   If the chosen metric does not improve by improvement_margin for as many as
@@ -107,4 +107,4 @@ def make_eval_adjusting_schedule(history,
     if steps_without_improvement >= steps_to_decrease:
       adjusted /= decrease_rate
       steps_without_improvement = 0
-  return make_default_schedule(history, constant=adjusted)
+  return DefaultSchedule(history, constant=adjusted)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 98b3d0568..a15febfe9 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -21,6 +21,7 @@
 
 import collections
 import functools
+import itertools
 import os
 import pickle
 import time
@@ -120,35 +121,46 @@ def save_state(state, output_dir):
 }
 
 
-def evaluate(inputs, predict_fn, eval_steps):
+def evaluate_train_and_eval(step, inputs, predict_fun, eval_steps,
+                            train_sw=None, eval_sw=None, history=None):
+  """Evalaute on train and eval data, and log metrics."""
+  step_log(step, "Evaluation")
+  train_metrics, eval_metrics = [
+      evaluate(  # pylint: disable=g-complex-comprehension
+          itertools.islice(input_stream(), eval_steps),
+          predict_fun,
+          _METRICS)
+      for input_stream in
+      [inputs.train_stream, inputs.eval_stream]]
+  if train_sw:
+    log_metrics(train_metrics, train_sw, "train", step, history=history)
+  if eval_sw:
+    log_metrics(eval_metrics, eval_sw, "eval ", step, history=history)
+  return train_metrics, eval_metrics
+
+
+def evaluate(inputs_stream, predict_fun, metric_funs):
   """Evaluate.
 
   Args:
-    inputs: Inputs namedtuple.
-    predict_fn: function from inputs to predictions. params should already be
+    inputs_stream: iterable of inputs to evaluate on.
+    predict_fun: function from inputs to predictions. params should already be
       partially applied.
-    eval_steps: int, number of evaluation steps.
+    metric_funs: dict from metric name to metric function, which takes inputs
+      and predictions and returns a scalar metric value.
 
   Returns:
-    train_metrics: dict
-    eval_metrics: dict
+    metrics: dict from metric name to metric value averaged over the number of
+      inputs.
   """
-  eval_stream = inputs.eval_fn()
-  eval_train_stream = inputs.train_fn()
-  train_metrics = {key: 0.0 for key in _METRICS}
-  eval_metrics = {key: 0.0 for key in _METRICS}
-  for _ in range(eval_steps):
-    train_batch = next(eval_train_stream)
-    train_predictions = predict_fn(train_batch[0])
-    eval_batch = next(eval_stream)
-    eval_predictions = predict_fn(eval_batch[0])
-    for m in _METRICS:
-      train_metrics[m] += (_METRICS[m](train_batch, train_predictions)
-                           / float(eval_steps))
-      eval_metrics[m] += (_METRICS[m](eval_batch, eval_predictions)
-                          / float(eval_steps))
-
-  return train_metrics, eval_metrics
+  metrics = collections.defaultdict(float)
+  count = 0
+  for inp in inputs_stream:
+    count += 1
+    preds = predict_fun(inp[0])
+    for m, f in six.iteritems(metric_funs):
+      metrics[m] += f(inp, preds)
+  return {m: v / count for (m, v) in six.iteritems(metrics)}
 
 
 def log_metrics(metrics, summ_writer, log_prefix, step, history=None):
@@ -167,8 +179,37 @@ def log_metrics(metrics, summ_writer, log_prefix, step, history=None):
 # * Make configurable:
 #   * loss
 #   * metrics
+# * Training loop callbacks/hooks/...
 # * Save/restore: pickle unsafe. Use np.array.savez + MessagePack?
 # * Move metrics to metrics.py
+# * Setup namedtuples for interfaces (e.g. lr fun constructors can take a
+#   LearningRateInit, metric funs, etc.).
+# * Allow disabling eval
+
+
+def epochs(steps=None, epoch_steps=1):
+  """Iterator over epochs until steps is reached. 1-indexed.
+
+  Args:
+    steps: int, total number of steps. Infinite if None.
+    epoch_steps: int, number of steps per epoch. Can also be an iterable<int> to
+      enable variable length epochs.
+
+  Yields:
+    (epoch: int, epoch id, epoch_steps: int, number of steps in this epoch)
+  """
+  try:
+    iter(epoch_steps)
+  except TypeError:
+    epoch_steps = itertools.repeat(epoch_steps)
+
+  step = 0
+  for epoch, epoch_steps in enumerate(epoch_steps):
+    epoch_steps = min(epoch_steps, steps - step)
+    yield (epoch + 1, epoch_steps)
+    step += epoch_steps
+    if steps and step >= steps:
+      break
 
 
 @gin.configurable(blacklist=["output_dir"])
@@ -176,7 +217,7 @@ def train(output_dir,
           model=gin.REQUIRED,
           inputs=gin.REQUIRED,
           optimizer=trax_opt.adam,
-          learning_rate_fn=lr.make_default_schedule,
+          lr_schedule=lr.DefaultSchedule,
           train_steps=1000,
           eval_steps=10,
           eval_frequency=100):
@@ -189,8 +230,8 @@ def train(output_dir,
     inputs: callable returning trax.inputs.Inputs.
     optimizer: The optimizer as a callable taking a learning_rate callable and
       returning 2 callables, opt_init and opt_update.
-    learning_rate_fn: The learning rate callable that takes history and returns
-      a function from step to learning rate (a float).
+    lr_schedule: A learning rate schedule as a function that takes history and
+      returns a function from step to learning rate (a float).
     train_steps: int, total number of training steps.
     eval_steps: int, num of steps per evaluation. If None or 0, eval disabled.
     eval_frequency: int, how often to run evaluation (every eval_frequency
@@ -209,8 +250,8 @@ def train(output_dir,
   # Setup optimizer and model
   state = restore_state(output_dir)
   history = state.history
-  learning_rate = learning_rate_fn(history)
-  opt_init, opt_update = optimizer(learning_rate)
+  lr_fun = lr_schedule(history)
+  opt_init, opt_update = optimizer(lr_fun)
   model_init, model_predict = model()
 
   # Setup state
@@ -229,59 +270,62 @@ def update(i, opt_state, batch):
         params, batch, model_predict), opt_state)
 
   print()
-  step_log(step, "starting training")
-  inputs_stream = inputs.train_fn()
-  eval_enabled = eval_steps and eval_frequency
-  is_first_step = True
-  # Evaluate after the first training step, then reset to normal_epoch_steps
-  normal_epoch_steps = (eval_enabled and eval_frequency) or train_steps
-  epoch_steps = 1
-  while step < train_steps:
-    print()  # separate logging for each loop iteration
-
-    # Train
+  train_stream = inputs.train_stream()
+  epoch_steps = itertools.chain([1,  # first epoch only 1 step
+                                 eval_frequency - 1],
+                                itertools.repeat(eval_frequency))
+  step_log(step, "Starting training")
+
+  for epoch, epoch_steps in epochs(train_steps, epoch_steps):
+    # Log separator
+    print()
+
+    # Timer
     start_time = time.time()
+
     for _ in range(epoch_steps):
-      opt_state = update(step, opt_state, next(inputs_stream))
-      if step % 10 == 0:  # Log learning rate curve each 10 steps.
-        train_sw.scalar("training/learning rate",
-                        learning_rate(step), step=step)
+      # Train
+      opt_state = update(step, opt_state, next(train_stream))
       step += 1
+
+      # LR log
+      if step == 1 or step % 10 == 0:
+        train_sw.scalar("training/learning rate",
+                        lr_fun(step), step=step)
+
+    # Timer
     epoch_time = time.time() - start_time
-    step_log(step, "ran %d train steps in %0.2f secs" %
+    step_log(step, "Ran %d train steps in %0.2f secs" %
              (epoch_steps, epoch_time))
+    if epoch_steps > 1:
+      train_sw.scalar("training/steps per second",
+                      epoch_steps / epoch_time, step=step)
 
     # Evaluate
     params = jax_opt.get_params(opt_state)
-    if eval_enabled:
-      step_log(step, "starting evaluation")
-      train_metrics, eval_metrics = evaluate(
-          inputs, functools.partial(jit_predict, params), eval_steps)
-      log_metrics(train_metrics, train_sw, "train", step, history=history)
-      log_metrics(eval_metrics, eval_sw, "eval ", step, history=history)
-      eval_sw.writer.flush()
+    evaluate_train_and_eval(
+        step=step,
+        inputs=inputs,
+        predict_fun=functools.partial(jit_predict, params),
+        eval_steps=eval_steps,
+        train_sw=train_sw,
+        eval_sw=eval_sw,
+        history=history)
 
     # Save state
     save_state(State(params=params, step=step, history=history), output_dir)
 
-    # Gin only tracks the used parameters, so we save it after the first step.
-    if is_first_step:
+    # Save Gin config
+    # Gin only tracks the used parameters, so we save it after the first epoch.
+    if epoch == 1:
       save_gin(output_dir, train_sw)
 
-    # Log non-metric reports.
-    if not is_first_step:
-      train_sw.scalar("training/steps per second",
-                      epoch_steps / epoch_time, step=step)
-    train_sw.writer.flush()
+    # Update learning rate with new history
+    lr_fun = lr_schedule(history)
 
-    # Update learning rate with new history.
-    learning_rate = learning_rate_fn(history)
-
-    # After the first step, train for normal_epoch_steps steps before evaluating
-    epoch_steps = (
-        (normal_epoch_steps - 1) if is_first_step else normal_epoch_steps)
-    is_first_step = False
+    # Flush summary writers
+    train_sw.writer.flush()
+    eval_sw.writer.flush()
 
-  print()
-  step_log(step, "finished training")
+  step_log(step, "Training done")
   return State(params=params, step=step, history=history)

From cfc816bcb7969116fe5b6837e16420b4b0b51018 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 4 Mar 2019 11:21:27 -0800
Subject: [PATCH 1749/2720] Added a method to insert one event to the memory.

PiperOrigin-RevId: 236690193
---
 tensor2tensor/layers/transformer_memory.py | 27 ++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index 0042173da..c16c638bd 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -66,6 +66,33 @@ def set(self, mem_vals, mem_times, seq_length_so_far):
   def get(self):
     return self.mem_vals, self.mem_times, self.seq_length_so_far
 
+  def incremental_update(self, event):
+    """Add a new event to the memory and also advance the time.
+
+    Args:
+      event: a tensor in the shape of [batch_size, depth].
+    Returns:
+      the update op.
+    """
+    event = tf.expand_dims(event, 1)
+    similarity_logits = tf.matmul(event, tf.transpose(
+        self.mem_vals, [0, 2, 1]))
+    similarity_logits = tf.squeeze(similarity_logits, [1])
+    max_logits = tf.reduce_max(similarity_logits, -1, keep_dims=True)
+    similarity_logits = tf.where(
+        tf.less(self.mem_times, 0.5),
+        tf.tile(max_logits, [1, self.memory_size]) + 1.0,
+        similarity_logits)
+    _, indices = tf.nn.top_k(similarity_logits)
+    update_mask = tf.cast(tf.one_hot(indices, self.memory_size), tf.float32)
+    update_times = self.mem_times.assign_add(update_mask)
+    with tf.control_dependencies([update_times]):
+      add_to_vals = tf.where(
+          tf.cast(update_mask, tf.bool),
+          tf.zeros_like(self.mem_vals),
+          tf.div(event - self.mem_vals, tf.expand_dims(self.mem_times, 2)))
+      return self.mem_vals.assign_add(add_to_vals)
+
   def update(self, segment):
     """Update the memory given the segment of events.
 

From 4dfdfaa712c38938b06b339facfc9e7235b270df Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 4 Mar 2019 13:59:53 -0800
Subject: [PATCH 1750/2720] Make JSON dumps resilient to nested functions.

PiperOrigin-RevId: 236720218
---
 tensor2tensor/utils/hparam.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/hparam.py b/tensor2tensor/utils/hparam.py
index 130dbf685..15ce59677 100644
--- a/tensor2tensor/utils/hparam.py
+++ b/tensor2tensor/utils/hparam.py
@@ -541,9 +541,16 @@ def to_json(self, indent=None, separators=None, sort_keys=False):
     Returns:
       A JSON string.
     """
+    def remove_callables(x):
+      if callable(x):
+        return x.__name__
+      if isinstance(x, dict):
+        return {k: remove_callables(v) for k, v in six.iteritems(x)}
+      if isinstance(x, list):
+        return [remove_callables(i) for i in x]
+      return x
     return json.dumps(
-        {k: v.__name__ if callable(v) else v
-         for k, v in six.iteritems(self.values())},
+        {k: remove_callables(v) for k, v in six.iteritems(self.values())},
         indent=indent,
         separators=separators,
         sort_keys=sort_keys)

From 767b2fa373dae5eb2fa7bb880b0dc85d1cfc6d13 Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Mon, 4 Mar 2019 15:08:45 -0800
Subject: [PATCH 1751/2720] * Add trax_test.py * Renamed
 learning_rate.DefaultSchedule -> FactorBasedSchedule * Rm need to reset
 learning rates (they keep the history themselves and manage   whether they
 need to recompute something) * Change arg order in History (mode, metric,
 step, val) * Disable EvalAdjustingSchedule as it's not currently working

PiperOrigin-RevId: 236732800
---
 .../trax/configs/resnet50_imagenet_8gb.gin    |   8 +-
 tensor2tensor/trax/history.py                 |  30 +++-
 tensor2tensor/trax/learning_rate.py           | 137 +++++++++++-------
 tensor2tensor/trax/models/__init__.py         |   5 +-
 tensor2tensor/trax/trax.py                    |  14 +-
 tensor2tensor/trax/trax_test.py               |  89 ++++++++++++
 6 files changed, 209 insertions(+), 74 deletions(-)
 create mode 100644 tensor2tensor/trax/trax_test.py

diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
index 953310085..811117dfc 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
@@ -15,11 +15,11 @@ batch_fun.eval_batch_size = 32
 inputs.data_dir = None
 inputs.dataset_name = 't2t_image_imagenet224'
 
-# Parameters for DefaultSchedule:
+# Parameters for MultifactorSchedule:
 # ==============================================================================
-DefaultSchedule.constant = 5.0
-DefaultSchedule.schedule = 'constant * linear_warmup * rsqrt_decay'
-DefaultSchedule.warmup_steps = 400
+MultifactorSchedule.constant = 5.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 400
 
 # Parameters for momentum:
 # ==============================================================================
diff --git a/tensor2tensor/trax/history.py b/tensor2tensor/trax/history.py
index a7421e745..e99150b85 100644
--- a/tensor2tensor/trax/history.py
+++ b/tensor2tensor/trax/history.py
@@ -21,17 +21,20 @@
 
 import collections
 
+from absl import logging
+
 
 class History(object):
   """History of metrics.
 
   History contains the metrics recorded during training and evaluation.
-  Save data with history.append(metric, value, step, mode) and get a sequence
-  of data by calling history.get(metric, mode). For example:
+  Save data with history.append and get a sequence of data by calling
+  history.get.
 
-  history.append("metrics/accuracy", 0.04, 1, "train")
-  history.append("metrics/accuracy", 0.31, 1000, "train")
-  history.get("metrics/accuracy", "train")
+  For example:
+  history.append("train", "metrics/accuracy", 1, 0.04)
+  history.append("train", "metrics/accuracy", 1000, 0.31)
+  history.get("train", "metrics/accuracy")
   # returns [(1, 0.04), (1000, 0.31)]
   """
 
@@ -46,14 +49,27 @@ def __init__(self):
     # }
     self._values = {}
 
-  def append(self, metric, value, step, mode):
+  def append(self, mode, metric, step, value):
     """Append (step, value) pair to history for the given mode and metric."""
     if mode not in self._values:
       self._values[mode] = collections.defaultdict(list)
     self._values[mode][metric].append((step, value))
 
-  def get(self, metric, mode):
+  def get(self, mode, metric):
     """Get the history for the given metric and mode."""
     if mode not in self._values:
+      logging.info("Metric %s not found for mode %s", metric, mode)
       return []
     return list(self._values[mode][metric])
+
+  @property
+  def modes(self):
+    """Current tracked modes."""
+    return sorted(list(self._values.keys()))
+
+  def metrics_for_mode(self, mode):
+    """Metrics available for a given mode."""
+    if mode not in self._values:
+      logging.info("Mode %s not found", mode)
+      return []
+    return sorted(list(self._values[mode].keys()))
diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index 0bd526362..4f4f191c9 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -13,37 +13,39 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""trax learning rate schedules."""
+"""trax learning rate schedules.
+
+The learning rate schedules here all have the signature:
+  lr: history -> (step -> lr)
+
+That is, they are functions that take a trax.history.History and return a
+function that takes a step and returns a learning rate.
+"""
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
 import gin
+
 import jax.numpy as np
 
 
 @gin.configurable(blacklist=["history"])
-def DefaultSchedule(history=None,
-                    schedule="constant * linear_warmup * rsqrt_decay",
-                    constant=0.001,
-                    warmup_steps=100):
-  """Default learning rate  schedule.
+def MultifactorSchedule(history=None,
+                        factors="constant * linear_warmup * rsqrt_decay",
+                        constant=0.001,
+                        warmup_steps=100):
+  """Factor-based learning rate schedule.
 
-  Note: the learning rate schedule takes arguments and return a function,
-  learning_rate: step -> lr, that only takes a step and return the rate.
-  The reason is that learning_rate(step) is called at every training step,
-  so should be efficient, while the schedule is re-computed only when
-  evaluating the model, so usually only every 100 or 1000 steps.
-
-  Interprets factors in the schedule string which can consist of:
+  Interprets factors in the factors string which can consist of:
   * constant: interpreted as the constant value,
   * linear_warmup: interpreted as linear warmup until warmup_steps,
   * rsqrt_decay: divide by square root of max(step, warmup_steps)
 
   Args:
     history: the history of training and evaluation (History object).
-    schedule: a string with factors separated by "*" that defines the schedule.
+    factors: a string with factors separated by "*" that defines the schedule.
     constant: float, the starting constant for the learning rate schedule.
     warmup_steps: how many steps to warm up for in the warmup schedule.
 
@@ -51,7 +53,7 @@ def DefaultSchedule(history=None,
     a function learning_rate(step): float -> float, the step-dependent lr.
   """
   del history
-  factors = [n.strip() for n in schedule.split("*")]
+  factors = [n.strip() for n in factors.split("*")]
 
   def learning_rate(step):  # pylint: disable=invalid-name
     """Step to learning rate function."""
@@ -70,41 +72,70 @@ def learning_rate(step):  # pylint: disable=invalid-name
   return learning_rate
 
 
-@gin.configurable(blacklist=["history"])
-def EvalAdjustingSchedule(history,
-                          constant=0.001,
-                          steps_to_decrease=10,
-                          improvement_margin=0.01,
-                          decrease_rate=2.0,
-                          metric="metrics/accuracy"):
-  """Learning rate that decreases when eval metric stalls.
-
-  If the chosen metric does not improve by improvement_margin for as many as
-  steps_to_decrease steps, then the constant gets decreased by decrease rate.
-  Finally, the default schedule gets called with the adjusted constant.
-
-  Args:
-    history: the history of training and evaluation (History object).
-    constant: float, the starting constant for the learning rate schedule.
-    steps_to_decrease: int, after how many steps without improvement
-      should we decrease the constant.
-    improvement_margin: how much we need to improve to count it.
-    decrease_rate: by how much to decrease.
-    metric: which evaluation metric to use for adjustments.
-
-  Returns:
-    a function learning_rate(step): float -> float, the step-dependent lr.
-  """
-  metric = history.get(metric, "eval")
-  adjusted = constant
-  steps_without_improvement = 0
-  while len(metric) > 1:
-    last = metric.pop()
-    if last[1] < metric[-1][1] * (1 + improvement_margin):
-      steps_without_improvement += 1
-    else:
-      steps_without_improvement = 0
-    if steps_without_improvement >= steps_to_decrease:
-      adjusted /= decrease_rate
-      steps_without_improvement = 0
-  return DefaultSchedule(history, constant=adjusted)
+# TODO(trax): Find a way to enable this with @jit.
+# Currently disabled because it does not work with @jit. To use properly, would
+# need to re-initialize this learning rate schedule function, the optimizer, and
+# the update jit.
+# @gin.configurable(blacklist=["history"])
+# def EvalAdjustingSchedule(history,
+#                           constant=0.001,
+#                           steps_to_decrease=10,
+#                           improvement_margin=0.01,
+#                           decrease_rate=2.0,
+#                           adjustment_frequency=100,
+#                           history_mode="eval",
+#                           metric="metrics/accuracy"):
+#   """Learning rate that decreases when eval metric stalls.
+#
+#   If the chosen metric does not improve by improvement_margin for as many as
+#   steps_to_decrease steps, then the constant gets decreased by decrease rate.
+#   Finally, the MultifactorSchedule gets called with the adjusted constant.
+#
+#   Args:
+#     history: trax.history.History, the history of training and evaluation.
+#     constant: float, the starting constant for the learning rate schedule.
+#     steps_to_decrease: int, after how many steps without improvement
+#       should we decrease the constant.
+#     improvement_margin: how much we need to improve to consider the metric
+#       improved.
+#     decrease_rate: by what fraction to decrease (i.e. lr /= decrease_rate).
+#     adjustment_frequency: int, how often to reset the learning rate based on
+#       the latest history.
+#     history_mode: str, which mode of the history to use.
+#     metric: which evaluation metric to use for adjustments.
+#
+#   Returns:
+#     a function learning_rate(step): float -> float, the step-dependent lr.
+#   """
+#
+#   def get_constant_from_history():
+#     metrics = history.get(history_mode, metric)
+#     adjusted = constant
+#     steps_without_improvement = 0
+#     while len(metrics) > 1:
+#       last = metrics.pop()
+#       if last[1] < metrics[-1][1] * (1 + improvement_margin):
+#         steps_without_improvement += 1
+#       else:
+#         steps_without_improvement = 0
+#       if steps_without_improvement >= steps_to_decrease:
+#         adjusted /= decrease_rate
+#         steps_without_improvement = 0
+#     return adjusted
+#
+#   state = {
+#       "schedule": None,
+#   }
+#
+#   def reset_schedule():
+#     state["schedule"] = MultifactorSchedule(
+#         history, constant=get_constant_from_history())
+#
+#   reset_schedule()
+#
+#   def lr_step(step):
+#     if step % adjustment_frequency == 0:
+#       reset_schedule()
+#     return state["schedule"](step)
+#
+#   return lr_step
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index c885faed2..ac8237c66 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -30,5 +30,6 @@ def model_configure(*args, **kwargs):
   return gin.external_configurable(*args, **kwargs)
 
 
-model_configure(mlp.MLP)
-model_configure(resnet.Resnet50)
+# pylint: disable=invalid-name
+MLP = model_configure(mlp.MLP)
+Resnet50 = model_configure(resnet.Resnet50)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index a15febfe9..373e89be7 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -135,7 +135,7 @@ def evaluate_train_and_eval(step, inputs, predict_fun, eval_steps,
   if train_sw:
     log_metrics(train_metrics, train_sw, "train", step, history=history)
   if eval_sw:
-    log_metrics(eval_metrics, eval_sw, "eval ", step, history=history)
+    log_metrics(eval_metrics, eval_sw, "eval", step, history=history)
   return train_metrics, eval_metrics
 
 
@@ -167,10 +167,11 @@ def log_metrics(metrics, summ_writer, log_prefix, step, history=None):
   """Log metrics to summary writer and history."""
   rjust_len = max([len(name) for name in metrics])
   for name, value in six.iteritems(metrics):
-    step_log(step, "%s %s | % .8f" % (log_prefix, name.rjust(rjust_len), value))
+    step_log(step, "%s %s | % .8f" % (
+        log_prefix.ljust(5), name.rjust(rjust_len), value))
     full_name = "metrics/" + name
     if history:
-      history.append(full_name, value, step, log_prefix)
+      history.append(log_prefix, full_name, step, value)
     if summ_writer:
       summ_writer.scalar(full_name, value, step)
 
@@ -217,7 +218,7 @@ def train(output_dir,
           model=gin.REQUIRED,
           inputs=gin.REQUIRED,
           optimizer=trax_opt.adam,
-          lr_schedule=lr.DefaultSchedule,
+          lr_schedule=lr.MultifactorSchedule,
           train_steps=1000,
           eval_steps=10,
           eval_frequency=100):
@@ -256,7 +257,7 @@ def train(output_dir,
 
   # Setup state
   step = state.step or 0
-  params_initializer = lambda: model_init([-1] + inputs.input_shape)[1]
+  params_initializer = lambda: model_init([-1] + list(inputs.input_shape))[1]
   params = state.params or params_initializer()
   opt_state = opt_init(params)
 
@@ -320,9 +321,6 @@ def update(i, opt_state, batch):
     if epoch == 1:
       save_gin(output_dir, train_sw)
 
-    # Update learning rate with new history
-    lr_fun = lr_schedule(history)
-
     # Flush summary writers
     train_sw.writer.flush()
     eval_sw.writer.flush()
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
new file mode 100644
index 000000000..7a2905e83
--- /dev/null
+++ b/tensor2tensor/trax/trax_test.py
@@ -0,0 +1,89 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""trax test."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import contextlib
+import functools
+import tempfile
+
+import numpy as np
+
+from tensor2tensor.trax import inputs
+from tensor2tensor.trax import models
+from tensor2tensor.trax import trax
+
+from tensorflow import test
+from tensorflow.io import gfile
+
+
+def test_inputs(num_classes):
+  """Make trax.inputs.Inputs."""
+  batch_size = 2
+  input_shape = (6, 6, 3)
+
+  def input_stream():
+    while True:
+      yield (np.random.rand(*([batch_size] + list(input_shape))),
+             np.random.randint(num_classes, size=batch_size))
+
+  return inputs.Inputs(
+      train_stream=input_stream,
+      eval_stream=input_stream,
+      input_shape=input_shape)
+
+
+class TraxTest(test.TestCase):
+
+  @contextlib.contextmanager
+  def tmp_dir(self):
+    tmp = tempfile.mkdtemp(dir=self.get_temp_dir())
+    yield tmp
+    gfile.rmtree(tmp)
+
+  @property
+  def train_args(self):
+    num_classes = 4
+    return dict(
+        model=functools.partial(models.MLP,
+                                hidden_size=16,
+                                num_output_classes=num_classes),
+        inputs=lambda: test_inputs(num_classes),
+        train_steps=3,
+        eval_steps=2)
+
+  def _test_train(self, train_args):
+    with self.tmp_dir() as output_dir:
+      state = trax.train(output_dir, **train_args)
+
+      # Assert total train steps
+      self.assertEqual(train_args["train_steps"], state.step)
+
+      # Assert 2 epochs ran
+      train_acc = state.history.get("train", "metrics/accuracy")
+      eval_acc = state.history.get("eval", "metrics/accuracy")
+      self.assertEqual(len(train_acc), len(eval_acc))
+      self.assertEqual(2, len(eval_acc))
+
+  def test_train(self):
+    self._test_train(self.train_args)
+
+
+if __name__ == "__main__":
+  test.main()

From e166c8b62e98cefd1ff932580e7c61e903727874 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 5 Mar 2019 09:41:20 -0800
Subject: [PATCH 1752/2720] Make adjustable learning rate work by re-jit'ing.

PiperOrigin-RevId: 236864640
---
 tensor2tensor/trax/README.md        |  35 +++++---
 tensor2tensor/trax/learning_rate.py | 123 +++++++++++++---------------
 tensor2tensor/trax/trax.py          |  28 +++++--
 3 files changed, 101 insertions(+), 85 deletions(-)

diff --git a/tensor2tensor/trax/README.md b/tensor2tensor/trax/README.md
index d5c269f73..c13f296a3 100644
--- a/tensor2tensor/trax/README.md
+++ b/tensor2tensor/trax/README.md
@@ -1,17 +1,21 @@
-# `trax`: Train Neural Nets with JAX
+## `trax`: Train Neural Nets with JAX
 
 ![train tracks](https://images.pexels.com/photos/461772/pexels-photo-461772.jpeg?dl&fit=crop&crop=entropy&w=640&h=426)
 
-* Configuration is done with [`gin`](https://github.com/google/gin-config).
-  `trainer.py` takes `--config_file` as well as `--config` for file overrides.
-* Models are defined with [`stax`](https://github.com/google/jax/blob/master/jax/experimental/stax.py) in
-  `models/`. They are made gin-configurable in `models/__init__.py`.
-* Datasets are simple iterators over batches. Datasets from
-  [`tensorflow/datasets`](https://github.com/tensorflow/datasets)
-  and [`tensor2tensor`](https://github.com/tensorflow/tensor2tensor)
-  are built-in and can be addressed by name.
+### `trax`: T2T Radically Simpler with JAX
+
+*Why?* Because T2T has gotten too complex. We are simplifying the main code too,
+but we wanted to try a more radical step. So you can write code as in pure
+NumPy and debug directly. So you can easily pinpoint each line where things
+happen and understand each function. But we also want it to run fast on
+accelerators, and that's possible with [JAX](https://github.com/google/jax).
+
+*Status:* preview; things work: models train, checkpoints are saved, TensorBoard
+has summaries, you can decode. But we are changing a lot every day for now.
+Please let us know what we should add, delete, keep, change. We plan to move
+the best parts into core JAX.
 
-Entrypoints:
+*Entrypoints:*
 
 * Script: `trainer.py`
 * Main library entrypoint: `trax.train`
@@ -35,3 +39,14 @@ python -m tensor2tensor.trax.trainer \
 python -m tensor2tensor.trax.trainer \
   --config_file=$PWD/trax/configs/resnet50_imagenet_8gb.gin
 ```
+
+### How `trax` differs from T2T
+
+* Configuration is done with [`gin`](https://github.com/google/gin-config).
+  `trainer.py` takes `--config_file` as well as `--config` for file overrides.
+* Models are defined with [`stax`](https://github.com/google/jax/blob/master/jax/experimental/stax.py) in
+  `models/`. They are made gin-configurable in `models/__init__.py`.
+* Datasets are simple iterators over batches. Datasets from
+  [`tensorflow/datasets`](https://github.com/tensorflow/datasets)
+  and [`tensor2tensor`](https://github.com/tensorflow/tensor2tensor)
+  are built-in and can be addressed by name.
diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index 4f4f191c9..e767a57a5 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -31,6 +31,14 @@
 import jax.numpy as np
 
 
+# A dictionary to memoize results of the MultifactorSchedule below.
+# We memoize because returning exactly the same function objects allows
+# later (in the training loop) to optimize re-compiling the function (for
+# running on an accelerator) only when it changes. Note that this does not
+# affect correctness, it is done purely for performance reasons.
+_memoized_multifactor_schedules = {}
+
+
 @gin.configurable(blacklist=["history"])
 def MultifactorSchedule(history=None,
                         factors="constant * linear_warmup * rsqrt_decay",
@@ -53,6 +61,11 @@ def MultifactorSchedule(history=None,
     a function learning_rate(step): float -> float, the step-dependent lr.
   """
   del history
+
+  cache_args = (factors, constant, warmup_steps)
+  if cache_args in _memoized_multifactor_schedules:
+    return _memoized_multifactor_schedules[cache_args]
+
   factors = [n.strip() for n in factors.split("*")]
 
   def learning_rate(step):  # pylint: disable=invalid-name
@@ -69,73 +82,49 @@ def learning_rate(step):  # pylint: disable=invalid-name
         raise ValueError("Unknown factor %s." % name)
     return ret
 
+  _memoized_multifactor_schedules[cache_args] = learning_rate
   return learning_rate
 
 
-# TODO(trax): Find a way to enable this with @jit.
-# Currently disabled because it does not work with @jit. To use properly, would
-# need to re-initialize this learning rate schedule function, the optimizer, and
-# the update jit.
-# @gin.configurable(blacklist=["history"])
-# def EvalAdjustingSchedule(history,
-#                           constant=0.001,
-#                           steps_to_decrease=10,
-#                           improvement_margin=0.01,
-#                           decrease_rate=2.0,
-#                           adjustment_frequency=100,
-#                           history_mode="eval",
-#                           metric="metrics/accuracy"):
-#   """Learning rate that decreases when eval metric stalls.
-#
-#   If the chosen metric does not improve by improvement_margin for as many as
-#   steps_to_decrease steps, then the constant gets decreased by decrease rate.
-#   Finally, the MultifactorSchedule gets called with the adjusted constant.
-#
-#   Args:
-#     history: trax.history.History, the history of training and evaluation.
-#     constant: float, the starting constant for the learning rate schedule.
-#     steps_to_decrease: int, after how many steps without improvement
-#       should we decrease the constant.
-#     improvement_margin: how much we need to improve to consider the metric
-#       improved.
-#     decrease_rate: by what fraction to decrease (i.e. lr /= decrease_rate).
-#     adjustment_frequency: int, how often to reset the learning rate based on
-#       the latest history.
-#     history_mode: str, which mode of the history to use.
-#     metric: which evaluation metric to use for adjustments.
-#
-#   Returns:
-#     a function learning_rate(step): float -> float, the step-dependent lr.
-#   """
-#
-#   def get_constant_from_history():
-#     metrics = history.get(history_mode, metric)
-#     adjusted = constant
-#     steps_without_improvement = 0
-#     while len(metrics) > 1:
-#       last = metrics.pop()
-#       if last[1] < metrics[-1][1] * (1 + improvement_margin):
-#         steps_without_improvement += 1
-#       else:
-#         steps_without_improvement = 0
-#       if steps_without_improvement >= steps_to_decrease:
-#         adjusted /= decrease_rate
-#         steps_without_improvement = 0
-#     return adjusted
-#
-#   state = {
-#       "schedule": None,
-#   }
-#
-#   def reset_schedule():
-#     state["schedule"] = MultifactorSchedule(
-#         history, constant=get_constant_from_history())
-#
-#   reset_schedule()
-#
-#   def lr_step(step):
-#     if step % adjustment_frequency == 0:
-#       reset_schedule()
-#     return state["schedule"](step)
-#
-#   return lr_step
+@gin.configurable(blacklist=["history"])
+def EvalAdjustingSchedule(history,
+                          constant=0.001,
+                          steps_to_decrease=10,
+                          improvement_margin=0.01,
+                          decrease_rate=2.0,
+                          history_mode="eval",
+                          metric="metrics/accuracy"):
+  """Learning rate that decreases when eval metric stalls.
+
+  If the chosen metric does not improve by improvement_margin for as many as
+  steps_to_decrease steps, then the constant gets decreased by decrease rate.
+  Finally, the MultifactorSchedule gets called with the adjusted constant.
+
+  Args:
+    history: trax.history.History, the history of training and evaluation.
+    constant: float, the starting constant for the learning rate schedule.
+    steps_to_decrease: int, after how many steps without improvement
+      should we decrease the constant.
+    improvement_margin: how much we need to improve to consider the metric
+      improved.
+    decrease_rate: by what fraction to decrease (i.e. lr /= decrease_rate).
+    history_mode: str, which mode of the history to use.
+    metric: which evaluation metric to use for adjustments.
+
+  Returns:
+    a function learning_rate(step): float -> float, the step-dependent lr.
+  """
+  metrics = history.get(history_mode, metric)
+  adjusted = constant
+  steps_without_improvement = 0
+  while len(metrics) > 1:
+    last = metrics.pop()
+    if last[1] < metrics[-1][1] * (1 + improvement_margin):
+      steps_without_improvement += 1
+    else:
+      steps_without_improvement = 0
+    if steps_without_improvement >= steps_to_decrease:
+      adjusted /= decrease_rate
+      steps_without_improvement = 0
+
+  return MultifactorSchedule(history, constant=adjusted)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 373e89be7..fd1b3d58e 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -213,6 +213,17 @@ def epochs(steps=None, epoch_steps=1):
       break
 
 
+def _jit_update_fun(predict_fun, loss_fun, optimizer, lr_fun):
+  """Get jit-ed update function for loss, optimizer, learning rate function."""
+  @jax.jit
+  def update(i, opt_state, batch):
+    _, opt_update = optimizer(lr_fun)
+    params = jax_opt.get_params(opt_state)
+    return opt_update(i, jax.grad(loss_fun)(
+        params, batch, predict_fun), opt_state)
+  return update
+
+
 @gin.configurable(blacklist=["output_dir"])
 def train(output_dir,
           model=gin.REQUIRED,
@@ -252,7 +263,7 @@ def train(output_dir,
   state = restore_state(output_dir)
   history = state.history
   lr_fun = lr_schedule(history)
-  opt_init, opt_update = optimizer(lr_fun)
+  opt_init, _ = optimizer(lr_fun)
   model_init, model_predict = model()
 
   # Setup state
@@ -263,12 +274,7 @@ def train(output_dir,
 
   # jit model_predict and update so they're fast
   jit_predict = jax.jit(model_predict)  # for evaluation
-
-  @jax.jit
-  def update(i, opt_state, batch):
-    params = jax_opt.get_params(opt_state)
-    return opt_update(i, jax.grad(loss)(
-        params, batch, model_predict), opt_state)
+  update_fun = _jit_update_fun(model_predict, loss, optimizer, lr_fun)
 
   print()
   train_stream = inputs.train_stream()
@@ -286,7 +292,7 @@ def update(i, opt_state, batch):
 
     for _ in range(epoch_steps):
       # Train
-      opt_state = update(step, opt_state, next(train_stream))
+      opt_state = update_fun(step, opt_state, next(train_stream))
       step += 1
 
       # LR log
@@ -321,6 +327,12 @@ def update(i, opt_state, batch):
     if epoch == 1:
       save_gin(output_dir, train_sw)
 
+    # Update learning rate with new history
+    old_lr_fun = lr_fun
+    lr_fun = lr_schedule(history)
+    if lr_fun != old_lr_fun:  # For performance, only jit if there is a change.
+      update_fun = _jit_update_fun(model_predict, loss, optimizer, lr_fun)
+
     # Flush summary writers
     train_sw.writer.flush()
     eval_sw.writer.flush()

From d27a628eefaf7945a5b7adb924b016b5a2c50803 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 5 Mar 2019 11:52:26 -0800
Subject: [PATCH 1753/2720] initial attention layers, transformer, and stax
 layers extensions

PiperOrigin-RevId: 236892081
---
 tensor2tensor/trax/models/__init__.py    |   2 +
 tensor2tensor/trax/models/transformer.py | 276 ++++++++++++++++
 tensor2tensor/trax/stax/README.md        | 124 ++++++++
 tensor2tensor/trax/stax/__init__.py      |  27 ++
 tensor2tensor/trax/stax/attention.py     | 208 ++++++++++++
 tensor2tensor/trax/stax/losses.py        |  49 +++
 tensor2tensor/trax/stax/slax.py          | 386 +++++++++++++++++++++++
 tensor2tensor/trax/stax/slax_test.py     | 210 ++++++++++++
 8 files changed, 1282 insertions(+)
 create mode 100644 tensor2tensor/trax/models/transformer.py
 create mode 100644 tensor2tensor/trax/stax/README.md
 create mode 100644 tensor2tensor/trax/stax/__init__.py
 create mode 100644 tensor2tensor/trax/stax/attention.py
 create mode 100644 tensor2tensor/trax/stax/losses.py
 create mode 100644 tensor2tensor/trax/stax/slax.py
 create mode 100644 tensor2tensor/trax/stax/slax_test.py

diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index ac8237c66..07fce9b0f 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -22,6 +22,7 @@
 
 from tensor2tensor.trax.models import mlp
 from tensor2tensor.trax.models import resnet
+from tensor2tensor.trax.models import transformer
 
 
 # Ginify
@@ -33,3 +34,4 @@ def model_configure(*args, **kwargs):
 # pylint: disable=invalid-name
 MLP = model_configure(mlp.MLP)
 Resnet50 = model_configure(resnet.Resnet50)
+Transformer = model_configure(transformer.Transformer)
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
new file mode 100644
index 000000000..0d79bf14c
--- /dev/null
+++ b/tensor2tensor/trax/models/transformer.py
@@ -0,0 +1,276 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Transformer Model."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+from jax import random
+import jax.numpy as np
+import tensor2tensor.trax.stax as stax
+
+
+def TransformerEncoder(mode='train',  # pylint: disable=invalid-name
+                       num_layers=6,
+                       feature_depth=512,
+                       feedforward_depth=2048,
+                       num_heads=8,
+                       dropout=0.9):
+  """Transformer Encoder Stack.
+
+  Args:
+    mode: str: 'train' or 'eval'
+    num_layers: int: number of encoder/decoder layers
+    feature_depth: int:  depth of embedding
+    feedforward_depth: int: depth of feed-forward layer
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate - Stax follows TF's KEEP probability convention
+
+  Returns:
+    A staxlayer for implementing a raw Transformer encoder stack.  No embedding
+    or positional signals are added by this layer.
+  """
+
+  # Multi-headed Attention and Feed-forward layers
+  multi_attention = stax.MultiHeadedAttention(
+      feature_depth, num_heads=num_heads, dropout=dropout, mode=mode)
+
+  feed_forward = stax.serial(
+      stax.Dense(feedforward_depth, W_init=stax.xavier_uniform()),
+      stax.Relu,
+      stax.Dropout(dropout, mode=mode),
+      stax.Dense(feature_depth, W_init=stax.xavier_uniform())
+  )
+
+  @stax.Lambda
+  def encoder(embedded_source, source_mask):
+    """Transformer encoder stack.
+
+    Args:
+      embedded_source: staxlayer variable: embedded source sequences
+      source_mask: staxlayer variable: self-attention mask
+
+    Returns:
+      Staxlayer variable that outputs encoded source.
+    """
+    encoder_layer = stax.serial(
+        # input attends to self
+        stax.residual(stax.LayerNorm(feature_depth),
+                      stax.multiplex(stax.Identity,  # query
+                                     stax.Identity,  # key
+                                     stax.Identity,  # value
+                                     source_mask),  # attention mask
+                      multi_attention,
+                      stax.Dropout(dropout, mode=mode)),
+        # feed-forward
+        stax.residual(stax.LayerNorm(feature_depth),
+                      feed_forward,
+                      stax.Dropout(dropout, mode=mode))
+    )
+    return stax.serial(
+        embedded_source,
+        stax.repeat(encoder_layer, num_layers),
+        stax.LayerNorm(feature_depth),
+    )
+
+  return encoder
+
+
+def Transformer(source_vocab_size,  # pylint: disable=invalid-name
+                target_vocab_size,
+                mode='train',
+                num_layers=6,
+                feature_depth=512,
+                feedforward_depth=2048,
+                num_heads=8,
+                dropout=0.9,
+                shared_embedding=True,
+                max_len=200,
+                return_evals=False):
+  """Transformer model.
+
+  Args:
+    source_vocab_size: int: source vocab size
+    target_vocab_size: int: target vocab size
+    mode: str: 'train' or 'eval'
+    num_layers: int: number of encoder/decoder layers
+    feature_depth: int:  depth of embedding
+    feedforward_depth: int: depth of feed-forward layer
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate - Stax follows TF's KEEP probability convention
+    shared_embedding: bool: specify whether source/target embeddings are tied.
+    max_len: int: maximum symbol length for positional encoding
+    return_evals: bool: whether to generate decode-time evaluation functions
+
+  Returns:
+    A namedtuple containing model 'init' and 'apply' functions for training and
+  the 'evals' functions that itself returns a namedtuple containing evaluation
+  functions for the trained encoder, decoder, and generator substax.
+  """
+
+  # Input embedding and positional encoding
+  inject_position = stax.serial(
+      stax.PositionalEncoding(feature_depth, max_len=max_len),
+      stax.Dropout(dropout, mode=mode)
+  )
+  if shared_embedding:
+    assert source_vocab_size == target_vocab_size
+    # Weight-shared Embedding
+    embedding = stax.Share(stax.Embedding(feature_depth, source_vocab_size))
+    source_embedding_layer = stax.serial(embedding, inject_position)
+    target_embedding_layer = source_embedding_layer
+  else:
+    source_embedding = stax.Embedding(feature_depth, source_vocab_size)
+    target_embedding = stax.Embedding(feature_depth, target_vocab_size)
+    source_embedding_layer = stax.serial(source_embedding, inject_position)
+    target_embedding_layer = stax.serial(target_embedding, inject_position)
+
+  # Multi-headed Attention and Feed-forward layers
+  multi_attention = stax.MultiHeadedAttention(
+      feature_depth, num_heads=num_heads, dropout=dropout, mode=mode)
+
+  feed_forward = stax.serial(
+      stax.Dense(feedforward_depth, W_init=stax.xavier_uniform()),
+      stax.Relu,
+      stax.Dropout(dropout, mode=mode),
+      stax.Dense(feature_depth, W_init=stax.xavier_uniform())
+  )
+
+  # Encoder
+  @stax.Lambda
+  def encoder(source, source_mask):
+    """Transformer encoder stack.
+
+    Args:
+      source: staxlayer variable: raw source sequences
+      source_mask: staxlayer variable: self-attention mask
+
+    Returns:
+      Staxlayer variable that outputs encoded source.
+    """
+    encoder_layer = stax.serial(
+        # input attends to self
+        stax.residual(stax.LayerNorm(feature_depth),
+                      stax.multiplex(stax.Identity,  # query
+                                     stax.Identity,  # key
+                                     stax.Identity,  # value
+                                     source_mask),  # attention mask
+                      multi_attention,
+                      stax.Dropout(dropout, mode=mode)),
+        # feed-forward
+        stax.residual(stax.LayerNorm(feature_depth),
+                      feed_forward,
+                      stax.Dropout(dropout, mode=mode))
+    )
+    return stax.serial(
+        source,
+        source_embedding_layer,
+        stax.repeat(encoder_layer, num_layers),
+        stax.LayerNorm(feature_depth),
+    )
+
+  # Decoder
+  @stax.Lambda
+  def decoder(memory, target, target_mask, memory_mask):
+    """Transformer decoder stack.
+
+    Args:
+      memory: staxlayer variable: encoded source sequences
+      target: staxlayer variable: raw target sequences
+      target_mask: staxlayer variable: self-attention mask
+      memory_mask: staxlayer variable: memory attention mask
+
+    Returns:
+      Staxlayer variable that outputs encoded source.
+    """
+    decoder_layer = stax.serial(
+        # target attends to self
+        stax.residual(stax.LayerNorm(feature_depth),
+                      stax.multiplex(stax.Identity,  # query
+                                     stax.Identity,  # key
+                                     stax.Identity,  # value
+                                     target_mask),  # attention mask
+                      multi_attention,
+                      stax.Dropout(dropout, mode=mode)),
+        # target attends to encoded source
+        stax.residual(stax.LayerNorm(feature_depth),
+                      stax.multiplex(stax.Identity,  # query
+                                     memory,  # key
+                                     memory,  # value
+                                     memory_mask),  # attention mask
+                      multi_attention,
+                      stax.Dropout(dropout, mode=mode)),
+        # feed-forward
+        stax.residual(stax.LayerNorm(feature_depth),
+                      feed_forward,
+                      stax.Dropout(dropout, mode=mode))
+    )
+    return stax.serial(
+        target,
+        target_embedding_layer,
+        stax.repeat(decoder_layer, num_layers),
+        stax.LayerNorm(feature_depth),
+    )
+
+  # The Transformer
+  @stax.Lambda
+  def transformer(source, target, source_mask, target_mask, memory_mask):
+    encoded_source = encoder(source, source_mask)
+    return decoder(encoded_source, target, target_mask, memory_mask)
+
+  # Finally, bind the generator transform to use later for inference.
+  @stax.Lambda
+  def generator(encoded_target):
+    return stax.serial(
+        encoded_target,
+        stax.Dense(target_vocab_size, W_init=stax.xavier_uniform()),
+        stax.LogSoftmax
+    )
+
+  # Model-Building and Evaluation Functions
+  # Get entire model's init and apply pair
+  top_init, top_apply = generator(transformer)
+
+  # By default act as a normal Stax constructor and emit an (init, apply) pair.
+  if not return_evals:
+    return (top_init, top_apply)
+  else:
+    # Inference-time function for binding trained params to model and returning
+    # the python-bound sub-expressions for evaluation and sequence generation.
+    def make_namedtuple(**kwargs):
+      return collections.namedtuple('Model', kwargs.keys())(**kwargs)
+
+    def get_evals(params):
+      # We need to feed _concrete_ trained parameters through the network once.
+      # Otherwise the bound parameters point to abstract tracer values.
+      # The inputs don't matter.
+      fake_inputs = 5 * (np.ones((1), dtype=np.int32),)
+      fake_key = random.PRNGKey(1)
+      top_apply(params, fake_inputs, rng=fake_key)
+      # We can now return eval functions from the bound pieces of the model.
+      return make_namedtuple(
+          encoder=stax.make_apply_fun(encoder),
+          generator=stax.make_apply_fun(generator),
+          decoder=stax.make_apply_fun(decoder),
+      )
+
+    # We return the functions needed to train and evaluate the Transformer.
+    return make_namedtuple(
+        init=top_init,
+        apply=top_apply,
+        evals=get_evals,
+    )
diff --git a/tensor2tensor/trax/stax/README.md b/tensor2tensor/trax/stax/README.md
new file mode 100644
index 000000000..d92f8520a
--- /dev/null
+++ b/tensor2tensor/trax/stax/README.md
@@ -0,0 +1,124 @@
+# Stax - Layer Extensions
+
+# Convenience layers and combinators
+
+SLAX implements repeat, residual, and multiplex combinators, parallel
+input-tuple sub-selection with Take layer, and graph inputs shape-logging with
+LogInputs layer for debugging.
+
+# Name Binding
+
+SLAX implements Share, Bind, Var, Vars, Lambda, and make_apply_fun.
+These operators augment the point-free Stax API with name-binding operations.
+This provides a concise way of pointifying Stax notation when needed for
+complicated neural net models while retaining a very functional style overall.
+
+### Bind
+
+Layer name-binding. Caches the results of the layer on its first application
+in the computation DAG so that it can be referred to elsewhere in a model
+definition and used as though it were a pointer to the cached variable.
+
+We use the name-bound layer inside the main model just like a normal stax layer:
+
+```python
+# bind a layer with Bind:
+encoder = Bind(serial(Dense(10), Relu))
+
+# elsewhere in stax definition:
+model = serial(
+    # ...
+    encoder, # evaluated and cached here
+    # ...
+    encoder, # this always returns the same value
+    #...
+)
+
+# after training, we can access its params:
+encoder.params
+
+# or its last activations:
+encoder.value
+
+# or we can re-evaluate it with its trained set of params:
+eval_time_result = make_apply_fun(encoder)(inputs, **kwargs)
+```
+
+Also note the convenience functions __Var__ and __Vars__(_N_), which are just
+bound Identity layers. This is convenient for capturing input values to be used
+elsewhere in the model. These can be used with __parallel__ and the helper
+__multiplex__ combinators to easily route data around inside a stax model.
+
+### Share
+
+Parameter name-binding, for shared parameters. Just like __Bind__, but __Share__
+doesn't bind the cached _results_ of a layer, but only it's _parameters_. This
+allows us to create a weight-sharing layer by name. This works transparently
+with jax.grad and optimizers as they only ever see one set of real parameters
+from the state tree in the traced computations, so there's no inefficiency
+introduced.
+
+```python
+# bind a layer with Share:
+shared_layer = Share(serial(Dense(10), Relu))
+
+# elsewhere in stax definition:
+tower_A = serial(..., shared_layer, ...)
+tower_B = serial(..., shared_layer, ...)
+
+# after training, we can access its params:
+shared_layer.params
+
+# or we can re-evaluate it with its trained set of params:
+eval_time_result = make_apply_fun(shared_layer)(inputs, **kwargs)
+```
+
+### Lambda
+
+A function wrapper to allow concise function definitions of model layers.
+This uses __Bind__ behind the scenes to fill in the values of the named
+arguments with an input layer that captures the tuple of inputs, finally it
+wraps the output of the function with a special form of __Bind__ that overloads
+the `__call__` operator to make it easy to couple this subgraph to inputs as if
+it were a normal function call.
+
+```python
+# we wrap a normal python function (*args only no **kwargs supported, but they
+# always can be fed in from an outer scope.) e.g.:
+
+some_layer_outside = serial(...)
+@Lambda
+def fun(x, y):
+    tmp = serial(x, serial(Dense(10), Relu)))
+    return serial(parallel(tmp, y), FanInSum, some_layer_outside)
+
+# Later we can simply call the function with staxlayer arguments, even within
+# another Lambda wrapped function:
+result = fun(input1, input2)
+
+# or chain them:
+result = fun(input3, fun(input1, input2))
+
+# Lambda is doing the "spiritual equivalent" to following:
+x, y = Var(), Var()
+Bind(
+  serial(
+    parallel(x, y),
+    fun(x, y)
+  )
+)
+# But Lambda also takes care of some annoying technical issues with
+# combinators behind the scenes to make this work as well using a
+# special Bind that overloads __call__ to make the result act like a
+# function.
+
+# after training, we can access its params:
+fun.params
+
+# or its last activations:
+fun.value
+
+# or we can re-evaluate it with its trained set of params:
+eval_time_result = make_apply_fun(fun)(inputs, **kwargs)
+
+```
diff --git a/tensor2tensor/trax/stax/__init__.py b/tensor2tensor/trax/stax/__init__.py
new file mode 100644
index 000000000..aca079496
--- /dev/null
+++ b/tensor2tensor/trax/stax/__init__.py
@@ -0,0 +1,27 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Layers defined in trax."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# We create a flat stax.* namespace for uniform calling conventions as we
+# upstream changes.
+# pylint: disable=wildcard-import
+from jax.experimental.stax import *
+from tensor2tensor.trax.stax.attention import *
+from tensor2tensor.trax.stax.losses import *
+from tensor2tensor.trax.stax.slax import *
diff --git a/tensor2tensor/trax/stax/attention.py b/tensor2tensor/trax/stax/attention.py
new file mode 100644
index 000000000..a954e2d27
--- /dev/null
+++ b/tensor2tensor/trax/stax/attention.py
@@ -0,0 +1,208 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Attention Layers."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from jax import random
+import jax.experimental.stax as stax
+import jax.numpy as np
+import numpy as onp
+import numpy.random as npr
+
+
+def causal_mask(size, dtype=np.uint8):
+  """Causal attention mask."""
+  return onp.tril(onp.ones((1, size, size), dtype=dtype), k=0)
+
+
+def xavier_uniform(out_dim=0, in_dim=1, rng=npr):
+  """An initializer function for random uniform xavier-scaled coefficients."""
+  def init(shape):
+    fan_in, fan_out = shape[in_dim], shape[out_dim]
+    std = np.sqrt(2.0 / (fan_in + fan_out))
+    a = onp.sqrt(3.0) * std
+    return rng.uniform(low=-a, high=a, size=shape).astype('float32')
+  return init
+
+
+def LayerNorm(features, epsilon=1e-5):  # pylint: disable=invalid-name
+  """Layer construction function for Layer Normalization layer.."""
+  def init_fun(input_shape):
+    a_2 = np.ones(features)
+    b_2 = np.zeros(features)
+    return input_shape, (a_2, b_2)
+
+  def apply_fun(params, inputs, **kwargs):
+    del kwargs
+    (a_2, b_2) = params
+    mean = np.mean(inputs, axis=-1, keepdims=True)
+    std = np.std(inputs, axis=-1, keepdims=True)
+    return a_2 * (inputs - mean) / (std + epsilon) + b_2
+
+  return init_fun, apply_fun
+
+
+def Embedding(feature_depth, vocab_size):  # pylint: disable=invalid-name
+  """Layer constructor function for a dense embedding layer."""
+  def init_fun(input_shape):
+    output_shape = input_shape + (feature_depth,)
+    dense_embedding = xavier_uniform()((vocab_size, feature_depth))
+    return output_shape, dense_embedding
+  def apply_fun(params, inputs, **kwargs):
+    del kwargs
+    dense_embedding = params
+    return np.take(dense_embedding, inputs, axis=0)
+  return init_fun, apply_fun
+
+
+def PositionalEncoding(feature_depth, max_len):  # pylint: disable=invalid-name
+  """Implements bare positional encoding."""
+  def init_fun(input_shape):
+    # Compute the positional encodings once in log space.
+    pe = onp.zeros((max_len, feature_depth), dtype=onp.float32)
+    position = onp.arange(0, max_len)[:, onp.newaxis]
+    div_term = onp.exp(
+        onp.arange(0, feature_depth, 2) * -(onp.log(10000.0) / feature_depth))
+    pe[:, 0::2] = onp.sin(position * div_term)
+    pe[:, 1::2] = onp.cos(position * div_term)
+    pe = np.array(pe[onp.newaxis, :])  # send to device
+    return input_shape, pe
+
+  def apply_fun(params, inputs, **kwargs):
+    del kwargs
+    pe = params
+    symbol_size = np.shape(inputs)[1]
+    return inputs + pe[:, :symbol_size]
+
+  return init_fun, apply_fun
+
+
+def dot_product_attention(query, key, value, mask, dropout, mode, rng):
+  """Core dot product self-attention.
+
+  Args:
+    query: array of representations
+    key: array of representations
+    value: array of representations
+    mask: attention-mask, gates attention
+    dropout: float: dropout rate - keep probability
+    mode: 'eval' or 'train': whether to use dropout
+    rng: JAX PRNGKey: subkey for disposable use
+
+  Returns:
+    Self attention for q, k, v arrays.
+  """
+  depth = np.shape(query)[-1]
+  dots = np.matmul(query, np.swapaxes(key, -1, -2)) / np.sqrt(depth)
+  if mask is not None:
+    dots = np.where(mask, dots, -1e9)
+  dots = stax.softmax(dots, axis=-1)
+  if dropout is not None and mode == 'train':
+    keep = random.bernoulli(rng, dropout, dots.shape)
+    dots = np.where(keep, dots / dropout, 0)
+  out = np.matmul(dots, value)
+  return out
+
+
+def PureDotProductAttention(dropout=1.0, mode='train'):  # pylint: disable=invalid-name
+  """Pure single-headed self-attention.
+
+  Args:
+    dropout: float: dropout rate - keep probability
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    Pure single-headed attention layer. (No Dense transforms on input.)
+  """
+  def init_fun(input_shapes):
+    q_shape, _, v_shape, _ = input_shapes
+    output_shape = q_shape[:-1] + (v_shape[-1],)
+    return output_shape, ()
+  def apply_fun(params, inputs, **kwargs):
+    del params
+    q, k, v, mask = inputs
+    rng = kwargs.get('rng', None)
+    return dot_product_attention(q, k, v, mask,
+                                 dropout=dropout, mode=mode, rng=rng)
+  return init_fun, apply_fun
+
+
+def PureMultiHeadedAttention(  # pylint: disable=invalid-name
+    feature_depth, num_heads=8, dropout=1.0, mode='train'):
+  """Pure transformer-style multi-headed attention.
+
+  Args:
+    feature_depth: int:  depth of embedding
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate - keep probability
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    Pure Multi-headed attention layer. (No Dense transforms on input.)
+  """
+  def init_fun(input_shapes):
+    input_shape = input_shapes[0]
+    output_shape = input_shape[:-1] + (feature_depth,)
+    return output_shape, ()
+  def apply_fun(params, inputs, **kwargs):  # pylint: disable=missing-docstring
+    del params
+    rng = kwargs.get('rng', None)
+    q, k, v, mask = inputs
+    assert feature_depth % num_heads == 0
+    head_depth = feature_depth // num_heads
+    nbatch = np.shape(q)[0]
+    # nbatch, seqlen, feature_depth --> nbatch, num_heads, seqlen, head_depth
+    def split_heads(x):
+      return np.transpose(
+          np.reshape(x, (nbatch, -1, num_heads, head_depth)), (0, 2, 1, 3))
+    # nbatch, num_heads, seqlen, head_depth --> nbatch, seqlen, feature_depth
+    def join_heads(x):
+      return np.reshape(
+          np.transpose(x, (0, 2, 1, 3)), (nbatch, -1, num_heads*head_depth))
+    # Split heads, dot-product attention, rejoin heads.
+    return join_heads(
+        dot_product_attention(
+            split_heads(q), split_heads(k), split_heads(v), mask,
+            dropout=dropout, mode=mode, rng=rng))
+  return init_fun, apply_fun
+
+
+def MultiHeadedAttention(  # pylint: disable=invalid-name
+    feature_depth, num_heads=8, dropout=1.0, mode='train'):
+  """Transformer-style multi-headed attention.
+
+  Args:
+    feature_depth: int:  depth of embedding
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate - keep probability
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    Multi-headed self-attention layer.
+  """
+  return stax.serial(
+      stax.parallel(
+          stax.Dense(feature_depth, W_init=xavier_uniform()),
+          stax.Dense(feature_depth, W_init=xavier_uniform()),
+          stax.Dense(feature_depth, W_init=xavier_uniform()),
+          stax.Identity
+      ),
+      PureMultiHeadedAttention(
+          feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),
+      stax.Dense(feature_depth, W_init=xavier_uniform()),
+  )
diff --git a/tensor2tensor/trax/stax/losses.py b/tensor2tensor/trax/stax/losses.py
new file mode 100644
index 000000000..68c3112c3
--- /dev/null
+++ b/tensor2tensor/trax/stax/losses.py
@@ -0,0 +1,49 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Loss functions and layers."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gin
+import jax.numpy as np
+from tensor2tensor.trax.stax import slax
+
+
+@gin.configurable(blacklist=['logpred', 'target'])
+def kl_div(logpred, target, eps=np.finfo(np.float32).eps):
+  """Calculate KL-divergence."""
+  return np.sum(target * (np.log(target + eps) - logpred))
+
+
+def crossentropy_loss(logpred, target):
+  """Calculate crossentropy loss."""
+  return - np.mean(
+      np.sum(logpred * slax.one_hot(target, logpred.shape[-1]), axis=-1))
+
+
+@gin.configurable(blacklist=['logpred', 'target', 'size'])
+def label_smoothed_loss(logpred, target, size, padding_idx=0, smoothing=0.0):
+  """Returns a label-smoothing loss-criterion function."""
+  confidence = 1.0 - smoothing
+  zerosmoothed = smoothing / (size - 2)
+  delta = confidence - zerosmoothed
+  assert logpred.shape[1] == size
+  truedist = (np.full_like(logpred, zerosmoothed) +
+              delta * slax.one_hot(target, size))
+  truedist *= (1 - (np.arange(size) == padding_idx))
+  truedist *= (1 - (target == padding_idx))[:, np.newaxis]
+  return kl_div(logpred, truedist, eps=1e-6)
diff --git a/tensor2tensor/trax/stax/slax.py b/tensor2tensor/trax/stax/slax.py
new file mode 100644
index 000000000..e97a8c2fc
--- /dev/null
+++ b/tensor2tensor/trax/stax/slax.py
@@ -0,0 +1,386 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""SLAX - Layer eXtensions to Stax."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import inspect
+from absl import logging
+import jax.experimental.stax as stax
+import jax.numpy as np
+from jax.tree_util import register_pytree_node as _register_pytree_node
+
+
+# Utility functions
+# ------------------------------------------------------------------------------
+def one_hot(x, size, dtype=np.float32):
+  """Make a n+1 dim one-hot array from n dim int-categorical array."""
+  return np.array(x[..., np.newaxis] == np.arange(size), dtype)
+
+
+# Utility Combinators
+# ------------------------------------------------------------------------------
+def repeat(layer, num_repeats):
+  """Repeats layers serially num_repeats times."""
+  if num_repeats < 1:
+    raise ValueError('Repeat combinator num_repeats must be >= 1.')
+  layers = num_repeats * (layer,)
+  return stax.serial(*layers)
+
+
+def residual(*layers, **kwargs):
+  """Constructs a residual version of layers, summing input to layers output."""
+  res = kwargs.get('res', stax.Identity)
+  if len(layers) > 1:
+    return stax.serial(
+        stax.FanOut(2),
+        stax.parallel(stax.serial(*layers), res),
+        stax.FanInSum
+    )
+  elif len(layers) == 1:
+    return stax.serial(
+        stax.FanOut(2),
+        stax.parallel(layers[0], res),
+        stax.FanInSum
+    )
+  else:
+    raise ValueError('Empty residual combinator.')
+
+
+def multiplex(*args):
+  """Helper to form input argument lists of bound variables.
+
+  Args:
+    *args: list of bound layers or raw stax Identity layers.
+
+  Returns:
+    A layer returning in parallel the bound variables as well as
+  (multiple) copies of this layer's input wherever Identity has been specified.
+  """
+  return stax.serial(
+      stax.FanOut(len(args)),
+      stax.parallel(*args)
+  )
+
+
+# Utility Layers
+# ------------------------------------------------------------------------------
+def Take(*args):  # pylint: disable=invalid-name
+  """Layer to pick subset of inputs from parallel input stream.
+
+  Args:
+    *args: a sequence of ints
+
+  Returns:
+    A new layer that selects inputs from an incoming parallel stream.
+    In numpy notation: outputs = parallel_inputs[args]
+    If the resulting output list has only one member, it is automatically
+    unwrapped and the contents are passed on directly.
+  """
+  def init_fun(input_shape):
+    output_shape = []
+    for arg in args:
+      output_shape.append(input_shape[arg])
+    if len(output_shape) == 1:
+      output_shape = output_shape[0]
+    return (output_shape, ())
+  def apply_fun(params, inputs, **kwargs):
+    del params, kwargs
+    outputs = []
+    for arg in args:
+      outputs.append(inputs[arg])
+    if len(outputs) == 1:
+      outputs = outputs[0]
+    return outputs
+  return init_fun, apply_fun
+
+
+def LogInputs(prefix='', debug=True):  # pylint: disable=invalid-name
+  """Logging side-effects layer, equivalent to Identity.
+
+  Args:
+    prefix: string: logging prefix
+    debug: bool: if True this will print logs, otherwise not.
+
+  Returns:
+    An Identity layer with log-printing side-effects. This
+  prints the types and shapes of the inputs.  NB: at the moment
+  this doesn't handle printing nested tuple/list shapes!
+  """
+  def return_shapes(inputs):
+    """Return shape information of inputs."""
+    if isinstance(inputs, _PlaceholderTree):
+      return []
+    if isinstance(inputs, (list, tuple)):
+      return [x.shape for x in inputs]
+    elif isinstance(inputs, dict):
+      return [inputs[k].shape for k in inputs.keys()]
+    else:
+      return inputs.shape
+  def init_fun(input_shape):
+    if debug:
+      logging.info('%s [init]: %s', prefix, input_shape)
+    return input_shape, ()
+  def apply_fun(params, inputs, **kwargs):
+    del params, kwargs
+    if debug:
+      logging.info('%s: %s %s', prefix, type(inputs), return_shapes(inputs))
+    return inputs
+  return init_fun, apply_fun
+
+
+# Staxlayer binding to python variables
+# ------------------------------------------------------------------------------
+# Stax params-tree leaf type to mark bound subtrees references.
+class _TreeMarker(dict):
+  pass
+# Add this leaf-type to JAX's tree-walker.
+_register_pytree_node(_TreeMarker,
+                      lambda xs: (tuple(), None),
+                      lambda _, xs: _TreeMarker())
+
+
+# TODO(levskaya, rsepassi): abstract away tuple-subclassing to StaxLayer?
+class Share(tuple):
+  """Layer parameter caching function to allow weight sharing.
+
+  Args:
+    A staxlayer: an (init_fun, apply_fun) pair.
+
+  Returns:
+    A 'parameter-bound' staxlayer that can be assigned to a python variable.
+  Wherever this value is needed elsewhere in the stax tree, call this bound
+  variable and all occurrences will share parameters that will automatically
+  be updated by Stax optimizers.
+  """
+
+  def __init__(self, staxlayer):  # pylint: disable=super-init-not-called
+    self._orig_init_fun, self._orig_apply_fun = staxlayer
+    self._first_init = True
+    self.params = None  # cached staxlayer params
+
+  def _init_fun(self, input_shape):  # pylint: disable=missing-docstring
+    if self._first_init:
+      # point of first subgraph initialization call: sets params, output_shape
+      self._first_init = False
+      out_shape, self.params = self._orig_init_fun(input_shape)
+      return out_shape, self.params
+    else:
+      # point of subgraph reuse:
+      # params are just a marker to apply_funs signalling subgraph params reuse
+      out_shape, _ = self._orig_init_fun(input_shape)
+      return out_shape, _TreeMarker()
+
+  def _apply_fun(self, params, inputs, **kwargs):
+    if isinstance(params, _TreeMarker):
+      # point of subgraph reuse: calculate new value with cached params
+      return self._orig_apply_fun(self.params, inputs, **kwargs)
+    else:
+      # point of first subgraph application to params: cache params
+      self.params = params
+      return self._orig_apply_fun(params, inputs, **kwargs)
+
+  # when unpacking this (init, apply) pair we return the wrapped funs
+  def __iter__(self):
+    return iter((self._init_fun, self._apply_fun))
+
+
+class Bind(tuple):
+  """Layer/variable caching function to allow name binding.
+
+  Args:
+    A staxlayer: an (init_fun, apply_fun) pair.
+
+  Returns:
+    A 'bound' staxlayer that can be assigned to a python variable.
+  Wherever this value is needed elsewhere in the stax tree, call this bound
+  variable and all occurrences will share output values.
+  """
+
+  def __init__(self, staxlayer):  # pylint: disable=super-init-not-called
+    self._orig_init_fun, self._orig_apply_fun = staxlayer
+    self._first_init = True
+    self._out_shape = None  # cached staxlayer output shape
+    self.params = None  # cached staxlayer params
+    self.value = None  # cached staxlayer output value
+
+  def _init_fun(self, input_shape):
+    if self._first_init:
+      # point of first subgraph initialization call: sets params, output_shape
+      self._first_init = False
+      self._out_shape, self.params = self._orig_init_fun(input_shape)
+      return self._out_shape, self.params
+    else:
+      # point of subgraph reuse:
+      # params are just a marker to apply_funs signalling subgraph value reuse
+      return self._out_shape, _TreeMarker()
+
+  def _apply_fun(self, params, inputs, **kwargs):
+    if isinstance(params, _TreeMarker):
+      # point of subgraph reuse: return cached value
+      return self.value
+    else:
+      # point of first subgraph application to params: cache value
+      self.params = params
+      self.value = self._orig_apply_fun(params, inputs, **kwargs)
+      return self.value
+
+  # when unpacking this (init, apply) pair we return the wrapped funs
+  def __iter__(self):
+    return iter((self._init_fun, self._apply_fun))
+
+
+# Convenience methods for common use-case of input variable capture and reuse.
+Var = lambda: Bind(stax.Identity)  # pylint: disable=invalid-name,
+Vars = lambda num_vars: tuple(Bind(stax.Identity) for _ in range(num_vars))  # pylint: disable=invalid-name,
+
+
+def make_apply_fun(bound_layer):
+  """Returns an apply function partially applied to bound params.
+
+  Requires that the top-level model apply_fun be fed params with
+  concrete values for these bound params to be numerically meaningful!
+  (e.g. not JaxprTrace arrays from a JAX JIT pass!)
+
+  Args:
+    bound_layer: Share/Bind/Lambda-bound staxlayer
+
+  Returns:
+    An apply function for this subgraph.
+  """
+  if not isinstance(bound_layer, (Share, Bind)):
+    raise ValueError('Can only create apply function from bound layer.')
+  def partial_apply_fun(inputs, **kwargs):
+    return bound_layer._orig_apply_fun(  # pylint: disable=protected-access
+        bound_layer.params, inputs, **kwargs)
+  return partial_apply_fun
+
+
+# Lambda
+# ------------------------------------------------------------------------------
+# The below provide a nicer syntax for 'pointy' function definition than using
+# raw bound variables.
+class LambdaBind(Bind):
+  """Layer/variable caching function to allow name binding for Lambda layers.
+
+  Args:
+    A staxlayer: an (init_fun, apply_fun) pair.
+
+  Returns:
+    A 'bound' staxlayer that can be assigned to a python variable.
+  Wherever this value is needed elsewhere in the stax tree, call this bound
+  variable and all occurrences will share output values.  Overloads __call__
+  to provide syntactic sugar for Lambda-like invocation.
+  """
+
+  # Syntactic sugar for applying this Lambda to other staxlayers
+  # NB: we do not bind the result by default here!
+  def __call__(self, *args):
+    if len(args) > 1:
+      return stax.serial(stax.parallel(*args), self)
+    elif len(args) == 1:
+      return stax.serial(args[0], self)
+    else:
+      return self
+
+
+class _PlaceholderTree(tuple):
+  """Placeholder tree object for 'initializing' combinators inside Lambdas.
+
+  When we create a Lambda, we're cutting off normal Stax data flow into
+  the subgraph that Lambda wraps with its bound inputs.  This is a
+  problem for any (potentially nested) parallel/serial combinators that
+  are input-facing, as they'll try to unpack the input_shape, inputs, and
+  params trees to feed their sub-layers.  We can't easily know what series
+  of nested access patterns are in a function, so we instead provide
+  recursive placeholder trees to placate the combinators. These placeholders
+  should feed into Lambda input nodes that completely ignore their inputs
+  anyway, but they'll break immediately if the user tries to use unbound
+  inputs from the Stax chain, which is a useful way to force the semantics
+  of Lambda.  This is aggressively tested for correctness in our unit tests.
+  """
+
+  def __init__(self):  # pylint: disable=super-init-not-called
+    self.shape = 0
+    # set generous safety limits for placeholder tree recursion and traversal
+    self.iterator_limit = 1000
+    self.recursion_limit = 30
+
+  def __getitem__(self, _):
+    if self.recursion_limit > 0:
+      self.recursion_limit -= 1
+      return self
+    else:
+      raise IndexError('_PlaceholderTree reached maximum depth')
+
+  def __iter__(self):
+    return self
+
+  def __next__(self):  # PY3
+    return self.next()
+
+  def next(self):  # PY2
+    if self.iterator_limit > 0:
+      self.iterator_limit -= 1
+      return self
+    else:
+      raise StopIteration
+# Register this class with tree-walker to be ignored by optimizers' init fns.
+_register_pytree_node(_PlaceholderTree,
+                      lambda xs: (tuple(), None),
+                      lambda _, xs: _PlaceholderTree())
+
+
+def _PlaceholderInputs():  # pylint: disable=invalid-name
+  """Feeds placeholders into input combinators of a Lambda-bound staxlayer."""
+  init_fun = lambda input_shape: iter((_PlaceholderTree(), _PlaceholderTree()))
+  apply_fun = lambda params, inputs, **kwargs: _PlaceholderTree()
+  return init_fun, apply_fun
+_PlaceholderInputs = _PlaceholderInputs()  # pylint: disable=invalid-name
+
+
+def Lambda(fn):  # pylint: disable=invalid-name
+  """Turn a normal function into a bound, callable Stax layer.
+
+  Args:
+    fn: a python function with _named_ args (i.e. no *args) and no kwargs.
+
+  Returns:
+    A callable, 'bound' staxlayer that can be assigned to a python variable and
+    called like a function with other staxlayers as arguments.  Like Bind,
+    wherever this value is placed in the stax tree, it will always output the
+    same cached value.
+  """
+  # fn's args are just symbolic names that we fill with Vars.
+  num_args = len(inspect.getargspec(fn).args)
+  if num_args > 1:
+    bound_args = Vars(num_args)
+    return LambdaBind(stax.serial(
+        stax.parallel(*bound_args),  # capture inputs
+        _PlaceholderInputs,  # placeholders for input combinators inside fn
+        fn(*bound_args)  # feed captured inputs into fn's args
+    ))
+  elif num_args == 1:
+    bound_arg = Var()
+    return LambdaBind(stax.serial(
+        bound_arg,  # capture input
+        _PlaceholderInputs,  # placeholders for input combinators inside fn
+        fn(bound_arg)  # feed captured inputs into fn's args
+    ))
+  # LambdaBind when no args are given:
+  else:
+    return LambdaBind(fn())
diff --git a/tensor2tensor/trax/stax/slax_test.py b/tensor2tensor/trax/stax/slax_test.py
new file mode 100644
index 000000000..49cca2890
--- /dev/null
+++ b/tensor2tensor/trax/stax/slax_test.py
@@ -0,0 +1,210 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 Stax Extensions."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import itertools
+from absl.testing import absltest
+from jax import random
+import numpy as onp
+import tensor2tensor.trax.stax as stax
+
+
+def random_inputs(rng, input_shape):
+  if isinstance(input_shape, tuple):
+    return rng.randn(*input_shape).astype(onp.float32)
+  elif isinstance(input_shape, list):
+    return [random_inputs(rng, shape) for shape in input_shape]
+  else:
+    raise TypeError(type(input_shape))
+
+
+def check_shape_agreement(test_case, init_fun, apply_fun, input_shape):
+  result_shape, params = init_fun(input_shape)
+  inputs = random_inputs(onp.random.RandomState(0), input_shape)
+  rng_key = random.PRNGKey(0)
+  result = apply_fun(params, inputs, rng=rng_key)
+  test_case.assertEqual(result.shape, result_shape)
+
+
+def check_staxlayer(test_case, staxlayer, input_shape):
+  init_fun, apply_fun = staxlayer
+  check_shape_agreement(test_case, init_fun, apply_fun, input_shape)
+
+
+# Helper functions for testing Lambda wrapper against functions involving
+# complicated input trees:
+def _enumerate_trees_w_leaves(n_leaves):
+  """Construct all rooted trees with n leaves."""
+  def enumtree(*args):
+    n_args = len(args)
+    # trivial cases:
+    if n_args == 0:
+      return []
+    if n_args == 1:
+      return args
+    # general case of 2 or more args:
+    # build index array
+    idxs = range(0, n_args)
+    trees = []
+    # we consider all possible subsets of size n_set to gather
+    for n_set in range(2, n_args+1):
+      idxsets = list(itertools.combinations(idxs, n_set))
+      for idxset in idxsets:
+        # recurse by joining all subtrees with
+        # n_set leaves and (n_args - n_set) leaves
+        arg_set = tuple(args[i] for i in idxs if i in idxset)
+        arg_coset = tuple(args[i] for i in idxs if i not in idxset)
+        if arg_coset:
+          trees.extend(tuple(itertools.product(enumtree(*arg_set),
+                                               enumtree(*arg_coset))))
+        else:
+          # trivial case where arg_set is entire set
+          trees.append(arg_set)
+    return trees
+  # return enumerated trees with integers as leaves
+  return enumtree(*range(n_leaves))
+
+
+def _build_combinator_tree(input_treespec, in_vars):
+  """Build a trivial Staxlayer that takes a complicated tree of inputs."""
+  parallel_args = []
+  for e in input_treespec:
+    if isinstance(e, int):
+      parallel_args.append(in_vars[e])
+    elif isinstance(e, tuple):
+      parallel_args.append(_build_combinator_tree(e, in_vars))
+  return stax.serial(stax.parallel(*parallel_args), stax.FanInSum)
+
+
+class SlaxTest(absltest.TestCase):
+
+  # Lambdas replace the staxlayer input stream with a placeholder that
+  # _should_ break any use of unbound variables in the input stream.
+  def testLambda_forbidden_access(self):
+    with self.assertRaises(ValueError):
+      for tree_spec in _enumerate_trees_w_leaves(2):
+        @stax.Lambda
+        def lambda_fun(x, y):  # pylint: disable=unused-argument
+          return _build_combinator_tree(tree_spec,  # pylint: disable=cell-var-from-loop
+                                        # try to read from input stream
+                                        # rather than bound vars
+                                        (x, stax.Identity))
+        check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*2)
+
+  # Exhaustively test the tricky part of Lambda - input combinator
+  # "initialization" for all 2412 trees of stax serial and parallel
+  # combinators of up to six variables.  This probably covers most
+  # practical use patterns!
+
+  # The variables in for loops below are used immediately, disable lint warning
+  # for this section:
+  # pylint: disable=cell-var-from-loop
+  def testLambda_1_arg(self):
+    @stax.Lambda
+    def lambda_fun(x):
+      return _build_combinator_tree((0,), (x,))
+    check_staxlayer(self, lambda_fun, (1, 5, 7, 11))
+
+  def testLambda_2_args(self):
+    for tree_spec in _enumerate_trees_w_leaves(2):
+      @stax.Lambda
+      def lambda_fun(x, y):
+        return _build_combinator_tree(tree_spec, (x, y))
+      check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*2)
+
+  def testLambda_3_args(self):
+    for tree_spec in _enumerate_trees_w_leaves(3):
+      @stax.Lambda
+      def lambda_fun(x, y, z):
+        return _build_combinator_tree(tree_spec, (x, y, z))
+      check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*3)
+
+  def testLambda_4_args(self):
+    for tree_spec in _enumerate_trees_w_leaves(4):
+      @stax.Lambda
+      def lambda_fun(x, y, z, w):
+        return _build_combinator_tree(tree_spec, (x, y, z, w))
+      check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*4)
+
+  def testLambda_5_args(self):
+    for tree_spec in _enumerate_trees_w_leaves(5):
+      @stax.Lambda
+      def lambda_fun(x, y, z, w, v):
+        return _build_combinator_tree(tree_spec, (x, y, z, w, v))
+      check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*5)
+
+  def testLambda_6_args(self):
+    for tree_spec in _enumerate_trees_w_leaves(6):
+      @stax.Lambda
+      def lambda_fun(x, y, z, w, v, u):
+        return _build_combinator_tree(tree_spec, (x, y, z, w, v, u))
+      check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*6)
+
+  # Test a few other cases, unused variables, non-input-tree use of
+  # bound Lambda input variables.
+  def testLambda_4_args_only_3_used(self):
+    for tree_spec in _enumerate_trees_w_leaves(3):
+      @stax.Lambda
+      def lambda_fun(x, y, z, w):  # pylint: disable=unused-argument
+        return _build_combinator_tree(tree_spec, (x, y, z))
+      check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*4)
+
+  def testLambda_4_args_only_2_used(self):
+    for tree_spec in _enumerate_trees_w_leaves(2):
+      @stax.Lambda
+      def lambda_fun(x, y, z, w):  # pylint: disable=unused-argument
+        return _build_combinator_tree(tree_spec, (x, y))
+      check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*4)
+
+  def testLambda_4_args_only_1_used(self):
+    @stax.Lambda
+    def lambda_fun(x, y, z, w):  # pylint: disable=unused-argument
+      return _build_combinator_tree((0,), (x,))
+    check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*4)
+
+  def testLambda_5_args_2_post_input_tree(self):
+    for tree_spec in _enumerate_trees_w_leaves(3):
+      @stax.Lambda
+      def lambda_fun1(x, y, z, w, v):
+        input_tree = _build_combinator_tree(tree_spec, (x, y, z))
+        return stax.serial(input_tree,
+                           stax.multiplex(stax.Identity, w, v),
+                           stax.FanInSum)
+      check_staxlayer(self, lambda_fun1, [(1, 5, 7, 11),]*5)
+
+      @stax.Lambda
+      def lambda_fun2(x, y, z, w, v):
+        input_tree = _build_combinator_tree(tree_spec, (x, y, z))
+        return stax.serial(input_tree,
+                           stax.multiplex(w, stax.Identity, v),
+                           stax.FanInSum)
+      check_staxlayer(self, lambda_fun2, [(1, 5, 7, 11),]*5)
+
+      @stax.Lambda
+      def lambda_fun3(x, y, z, w, v):
+        input_tree = _build_combinator_tree(tree_spec, (x, y, z))
+        return stax.serial(input_tree,
+                           stax.multiplex(w, v, stax.Identity),
+                           stax.FanInSum)
+      check_staxlayer(self, lambda_fun3, [(1, 5, 7, 11),]*5)
+  # pylint: enable=cell-var-from-loop
+
+
+if __name__ == "__main__":
+  absltest.main()

From 7217e0fb8aabf256f7db3399758ce94966724d0d Mon Sep 17 00:00:00 2001
From: Brian Patton <bjp@google.com>
Date: Tue, 5 Mar 2019 12:22:27 -0800
Subject: [PATCH 1754/2720] Internal change

PiperOrigin-RevId: 236897455
---
 tensor2tensor/layers/modalities.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index a8fbb29f6..c8ee107f7 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -685,7 +685,7 @@ def video_pixel_noise_bottom(x, model_hparams, vocab_size):
   input_noise = getattr(model_hparams, "video_modality_input_noise", 0.25)
   inputs = x
   if model_hparams.mode == tf.estimator.ModeKeys.TRAIN:
-    background = tfp.distributions.percentile(inputs, 50., axis=[0, 1, 2, 3])
+    background = tfp.stats.percentile(inputs, 50., axis=[0, 1, 2, 3])
     input_shape = common_layers.shape_list(inputs)
     input_size = tf.reduce_prod(input_shape[:-1])
     input_mask = tf.multinomial(

From 6c869e209eb9ba8cece047f65c2960f563ca5179 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 5 Mar 2019 14:45:31 -0800
Subject: [PATCH 1755/2720] Add debug logging in some places.

PiperOrigin-RevId: 236924859
---
 tensor2tensor/trax/history.py | 3 +++
 tensor2tensor/trax/trainer.py | 4 +++-
 tensor2tensor/trax/trax.py    | 3 ++-
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/history.py b/tensor2tensor/trax/history.py
index e99150b85..4ffc9af3f 100644
--- a/tensor2tensor/trax/history.py
+++ b/tensor2tensor/trax/history.py
@@ -73,3 +73,6 @@ def metrics_for_mode(self, mode):
       logging.info("Mode %s not found", mode)
       return []
     return sorted(list(self._values[mode].keys()))
+
+  def __str__(self):
+    return str(self._values)
diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index fb56ec83a..32ce88107 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -41,6 +41,7 @@
                           "Configuration file with parameters (.gin).")
 flags.DEFINE_multi_string("config", None,
                           "Configuration parameters (gin string).")
+flags.DEFINE_integer("log_level", logging.INFO, "Log level.")
 
 
 def _default_output_dir():
@@ -77,6 +78,8 @@ def _setup_gin():
 
 
 def main(_):
+  logging.set_verbosity(FLAGS.log_level)
+
   _setup_gin()
 
   # Setup output directory
@@ -88,5 +91,4 @@ def main(_):
 
 
 if __name__ == "__main__":
-  logging.set_verbosity(logging.INFO)
   app.run(main)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index fd1b3d58e..67730603a 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -91,7 +91,8 @@ def restore_state(output_dir):
 
   with gfile.GFile(params_file, "rb") as f:
     (params, step, history) = pickle.load(f)
-  log("Model loaded from %s" % params_file)
+  log("Model loaded from %s at step %d" % (params_file, step))
+  logging.debug("From loaded model : history = %s", history)
   return State(step=step, params=params, history=history)
 
 

From 749927fee201d02d2f125a1e7b6f1152e9e6a9e4 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 5 Mar 2019 16:28:19 -0800
Subject: [PATCH 1756/2720]  - don't test subword_text_encoder_ops_test and
 pack_sequences_ops_test  - tests that enable eager are tested separately
 (since pytest runs everything in the same binary the graph only tests that
 don't have the decorator fail)  - ignore list is just sorted now.

PiperOrigin-RevId: 236943115
---
 oss_scripts/oss_tests.sh | 39 ++++++++++++++++++++++++++++++++++-----
 1 file changed, 34 insertions(+), 5 deletions(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 6297d9f5d..17895f586 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -38,23 +38,52 @@ set_status
 #   * allen_brain_test
 #   * models/research
 # algorithmic_math_test: flaky
+# subword_text_encoder_ops_test, pack_sequences_ops_test: interface with C++ ops
+# others (see below) enable eager, so can't be tested along with the others in
+# pytest
 pytest \
-  --ignore=tensor2tensor/utils/registry_test.py \
-  --ignore=tensor2tensor/utils/trainer_lib_test.py \
-  --ignore=tensor2tensor/visualization/visualization_test.py \
   --ignore=tensor2tensor/bin/t2t_trainer_test.py \
   --ignore=tensor2tensor/data_generators/algorithmic_math_test.py \
   --ignore=tensor2tensor/data_generators/allen_brain_test.py \
-  --ignore=tensor2tensor/rl \
+  --ignore=tensor2tensor/data_generators/ops/pack_sequences_ops_test.py \
+  --ignore=tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py \
+  --ignore=tensor2tensor/data_generators/problem_test.py \
+  --ignore=tensor2tensor/layers/bayes_test.py \
+  --ignore=tensor2tensor/layers/common_attention_test.py \
+  --ignore=tensor2tensor/layers/common_layers_test.py \
+  --ignore=tensor2tensor/layers/common_video_test.py \
+  --ignore=tensor2tensor/layers/discretization_test.py \
+  --ignore=tensor2tensor/layers/latent_layers_test.py \
+  --ignore=tensor2tensor/layers/modalities_test.py \
+  --ignore=tensor2tensor/layers/ngram_test.py \
+  --ignore=tensor2tensor/layers/reversible_layers_test.py \
   --ignore=tensor2tensor/models/research \
-  --ignore=tensor2tensor/models/video/nfg_conv_test.py \
   --ignore=tensor2tensor/models/video/nfg_conv3d_test.py \
   --ignore=tensor2tensor/models/video/nfg_conv_lstm_test.py \
+  --ignore=tensor2tensor/models/video/nfg_conv_test.py \
   --ignore=tensor2tensor/models/video/nfg_uncond_test.py \
+  --ignore=tensor2tensor/rl \
+  --ignore=tensor2tensor/utils/t2t_model_test.py \
+  --ignore=tensor2tensor/utils/test_utils_test.py \
+  --ignore=tensor2tensor/visualization/visualization_test.py \
   --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary \
   --deselect=tensor2tensor/utils/beam_search_test.py::BeamSearchTest::testTPUBeam
 set_status
 
+# These tests enable eager, so are tested separately.
+pytest tensor2tensor/layers/bayes_test.py \
+  tensor2tensor/layers/common_layers_test.py \
+  tensor2tensor/layers/common_attention_test.py \
+  tensor2tensor/layers/reversible_layers_test.py \
+  tensor2tensor/data_generators/problem_test.py \
+  tensor2tensor/layers/discretization_test.py \
+  tensor2tensor/layers/common_video_test.py \
+  tensor2tensor/utils/t2t_model_test.py \
+  tensor2tensor/layers/latent_layers_test.py \
+  tensor2tensor/layers/ngram_test.py \
+  tensor2tensor/layers/modalities_test.py \
+  tensor2tensor/utils/test_utils_test.py
+
 pytest tensor2tensor/utils/registry_test.py
 set_status
 

From 97ca3bfd70d72365de8ed307c994d06de11d8cfc Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 5 Mar 2019 17:29:45 -0800
Subject: [PATCH 1757/2720] Ignore test_utils.py from being tested, pytest
 thinks it is a test.

PiperOrigin-RevId: 236953133
---
 oss_scripts/oss_tests.sh | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 17895f586..f78b02492 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -38,6 +38,7 @@ set_status
 #   * allen_brain_test
 #   * models/research
 # algorithmic_math_test: flaky
+# test_utils.py is not a test, but pytest thinks it is.
 # subword_text_encoder_ops_test, pack_sequences_ops_test: interface with C++ ops
 # others (see below) enable eager, so can't be tested along with the others in
 # pytest
@@ -64,6 +65,7 @@ pytest \
   --ignore=tensor2tensor/models/video/nfg_uncond_test.py \
   --ignore=tensor2tensor/rl \
   --ignore=tensor2tensor/utils/t2t_model_test.py \
+  --ignore=tensor2tensor/utils/test_utils.py \
   --ignore=tensor2tensor/utils/test_utils_test.py \
   --ignore=tensor2tensor/visualization/visualization_test.py \
   --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary \

From 49bf55d17dfab694de54ec66fed2f143e71df6bc Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 5 Mar 2019 17:49:11 -0800
Subject: [PATCH 1758/2720] Work around differences in internal and external
 Gym.

PiperOrigin-RevId: 236956022
---
 tensor2tensor/envs/env_problem_test.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/envs/env_problem_test.py b/tensor2tensor/envs/env_problem_test.py
index 1f902fb19..03aa7ef74 100644
--- a/tensor2tensor/envs/env_problem_test.py
+++ b/tensor2tensor/envs/env_problem_test.py
@@ -87,7 +87,12 @@ def test_default_processed_rewards_discrete(self):
     # Assert that it is as expected of the underlying environment.
     reward_range = ep.reward_range
     self.assertEqual(0, reward_range[0])
-    self.assertEqual(ep._envs[0].maxWealth, reward_range[1])
+
+    # Google's version of Gym has maxWealth, vs max_wealth externally.
+    max_wealth = getattr(ep._envs[0], "maxWealth",
+                         getattr(ep._envs[0], "max_wealth", None))
+    self.assertIsNotNone(max_wealth)
+    self.assertEqual(max_wealth, reward_range[1])
 
     # Check that the processed rewards are discrete.
     self.assertTrue(ep.is_processed_rewards_discrete)

From 18e916abe0357d6f4fb1d3ab47ad0161997b60d4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 5 Mar 2019 18:14:45 -0800
Subject: [PATCH 1759/2720] added masking and shifted helper functions for
 transformer

PiperOrigin-RevId: 236959333
---
 tensor2tensor/trax/stax/attention.py | 33 ++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/tensor2tensor/trax/stax/attention.py b/tensor2tensor/trax/stax/attention.py
index a954e2d27..238f8b9ec 100644
--- a/tensor2tensor/trax/stax/attention.py
+++ b/tensor2tensor/trax/stax/attention.py
@@ -30,6 +30,39 @@ def causal_mask(size, dtype=np.uint8):
   return onp.tril(onp.ones((1, size, size), dtype=dtype), k=0)
 
 
+def make_target_mask(target, pad=0):
+  """Create an attention mask to hide padding and future words."""
+  target_mask = (target != pad)[ :, np.newaxis, :]
+  target_dtype = target_mask.dtype
+  target_mask = (
+      (target_mask & stax.causal_mask(target.shape[-1])).astype(target_dtype))
+  return np.expand_dims(target_mask, axis=1)
+
+
+def prepare_paired_sequence_batch(source, target_in, pad=0):
+  """Build masks for this batch.
+
+  Args:
+    source: (batch, source_len) array of integer-coded symbols for inputs
+    target_in: (batch, batch_len) array of integer-coded symbols for targets
+    pad: int: the padding symbol used to pad the above
+
+  Returns:
+    Prepared batch of tuple of arrays: source, input-target, shifted-target,
+    source mask, target mask, source-target "memory" mask, minibatch token count
+  """
+  target = target_in[:, :-1]
+  target_y = target_in[:, 1:]
+  source_mask = np.reshape(source != pad,
+                           (source.shape[0], 1, 1, source.shape[-1]))
+  target_mask = make_target_mask(target, pad)
+  memory_mask = (
+      np.reshape(np.arange(target.shape[-1]) < source.shape[-1], [-1, 1]))
+  ntokens = np.sum(target_y != pad)
+  return (source, target, target_y,
+          source_mask, target_mask, memory_mask, ntokens)
+
+
 def xavier_uniform(out_dim=0, in_dim=1, rng=npr):
   """An initializer function for random uniform xavier-scaled coefficients."""
   def init(shape):

From 41af9000449e3e69382df8046db2058c6a67492d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 5 Mar 2019 18:25:57 -0800
Subject: [PATCH 1760/2720]  - trainer_lib_test/registry_test are tested
 separately, this got accidentally reverted in an earlier change.  - Not
 testing TRAX for now, this needs "GLIBCXX_3.4.21" so we'll figure this out
 separately.  - Move the testGifSummary to the new stanza, since
 common_video_test.py was moved to the new stanza.

PiperOrigin-RevId: 236960609
---
 oss_scripts/oss_tests.sh | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index f78b02492..b5084c18d 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -40,6 +40,7 @@ set_status
 # algorithmic_math_test: flaky
 # test_utils.py is not a test, but pytest thinks it is.
 # subword_text_encoder_ops_test, pack_sequences_ops_test: interface with C++ ops
+# trax tests need C++
 # others (see below) enable eager, so can't be tested along with the others in
 # pytest
 pytest \
@@ -64,16 +65,23 @@ pytest \
   --ignore=tensor2tensor/models/video/nfg_conv_test.py \
   --ignore=tensor2tensor/models/video/nfg_uncond_test.py \
   --ignore=tensor2tensor/rl \
+  --ignore=tensor2tensor/trax \
   --ignore=tensor2tensor/utils/t2t_model_test.py \
   --ignore=tensor2tensor/utils/test_utils.py \
   --ignore=tensor2tensor/utils/test_utils_test.py \
+  --ignore=tensor2tensor/utils/registry_test.py \
+  --ignore=tensor2tensor/utils/trainer_lib_test.py \
   --ignore=tensor2tensor/visualization/visualization_test.py \
-  --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary \
   --deselect=tensor2tensor/utils/beam_search_test.py::BeamSearchTest::testTPUBeam
 set_status
 
+# TODO(afrozm): Enable trax tests they currently need GLIBCXX_3.4.21
+# Travis Error:
+# ImportError: /usr/lib/x86_64-linux-gnu/libstdc++.so.6: version `GLIBCXX_3.4.21' not found (required by /home/travis/virtualenv/python3.6.3/lib/python3.6/site-packages/jaxlib/_pywrap_xla.so)
+
 # These tests enable eager, so are tested separately.
-pytest tensor2tensor/layers/bayes_test.py \
+pytest \
+  tensor2tensor/layers/bayes_test.py \
   tensor2tensor/layers/common_layers_test.py \
   tensor2tensor/layers/common_attention_test.py \
   tensor2tensor/layers/reversible_layers_test.py \
@@ -84,7 +92,8 @@ pytest tensor2tensor/layers/bayes_test.py \
   tensor2tensor/layers/latent_layers_test.py \
   tensor2tensor/layers/ngram_test.py \
   tensor2tensor/layers/modalities_test.py \
-  tensor2tensor/utils/test_utils_test.py
+  tensor2tensor/utils/test_utils_test.py \
+  --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary
 
 pytest tensor2tensor/utils/registry_test.py
 set_status

From 05d3c356959548df53dc596b7cc0b5e698349aea Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 5 Mar 2019 18:52:02 -0800
Subject: [PATCH 1761/2720] No-OP - Simply sorting test names in a pytest
 stanza.

PiperOrigin-RevId: 236963466
---
 oss_scripts/oss_tests.sh | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index b5084c18d..263a11630 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -81,17 +81,17 @@ set_status
 
 # These tests enable eager, so are tested separately.
 pytest \
+  tensor2tensor/data_generators/problem_test.py \
   tensor2tensor/layers/bayes_test.py \
-  tensor2tensor/layers/common_layers_test.py \
   tensor2tensor/layers/common_attention_test.py \
-  tensor2tensor/layers/reversible_layers_test.py \
-  tensor2tensor/data_generators/problem_test.py \
-  tensor2tensor/layers/discretization_test.py \
+  tensor2tensor/layers/common_layers_test.py \
   tensor2tensor/layers/common_video_test.py \
-  tensor2tensor/utils/t2t_model_test.py \
+  tensor2tensor/layers/discretization_test.py \
   tensor2tensor/layers/latent_layers_test.py \
-  tensor2tensor/layers/ngram_test.py \
   tensor2tensor/layers/modalities_test.py \
+  tensor2tensor/layers/ngram_test.py \
+  tensor2tensor/layers/reversible_layers_test.py \
+  tensor2tensor/utils/t2t_model_test.py \
   tensor2tensor/utils/test_utils_test.py \
   --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary
 

From d0dd8638a8cf5e0329110f12d4d4b47ce8163979 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 5 Mar 2019 18:52:50 -0800
Subject: [PATCH 1762/2720] Bump version to 1.13 in preparation for a release.

PiperOrigin-RevId: 236963547
---
 setup.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/setup.py b/setup.py
index 6c18258e6..3fda512dd 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.12.0',
+    version='1.13.0',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',
@@ -60,8 +60,8 @@
         'tqdm',
     ],
     extras_require={
-        'tensorflow': ['tensorflow>=1.12.0'],
-        'tensorflow_gpu': ['tensorflow-gpu>=1.12.0'],
+        'tensorflow': ['tensorflow>=1.13.1'],
+        'tensorflow_gpu': ['tensorflow-gpu>=1.13.1'],
         'tensorflow-hub': ['tensorflow-hub>=0.1.1'],
         'tests': [
             'absl-py',

From 9bc05874118884a1643c08c68a0cef1ecd175746 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 5 Mar 2019 21:24:00 -0800
Subject: [PATCH 1763/2720] Add Gym Mujoco envs to T2T.

PiperOrigin-RevId: 236977901
---
 tensor2tensor/envs/env_problem.py             | 38 ++++-----
 tensor2tensor/envs/env_problem_test.py        |  4 +
 tensor2tensor/envs/mujoco_problems.py         | 68 +++++++++++++++
 tensor2tensor/envs/mujoco_problems_test.py    | 72 ++++++++++++++++
 tensor2tensor/envs/tic_tac_toe_env_problem.py |  4 +
 tensor2tensor/rl/gym_utils.py                 | 85 ++++++++++++++-----
 6 files changed, 228 insertions(+), 43 deletions(-)
 create mode 100644 tensor2tensor/envs/mujoco_problems.py
 create mode 100644 tensor2tensor/envs/mujoco_problems_test.py

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 4d7312001..840eab902 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -23,6 +23,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import gym
 from gym.core import Env
 import numpy as np
 import six
@@ -31,7 +32,6 @@
 from tensor2tensor.envs import gym_spaces_utils
 from tensor2tensor.envs import trajectory
 from tensor2tensor.layers import modalities
-from tensor2tensor.rl import gym_utils
 import tensorflow as tf
 
 # Names for data fields in stored tf.Examples.
@@ -56,9 +56,10 @@ class EnvProblem(Env, problem.Problem):
   Subclasses *should* override the following functions, since they are used in
   the `hparams` function to return modalities and vocab_sizes.
   - input_modality
-  - input_input_vocab_size
+  - input_vocab_size
   - target_modality
   - target_vocab_size
+  - action_modality
 
   NON NATIVELY BATCHED ENVS:
 
@@ -119,21 +120,23 @@ class EnvProblem(Env, problem.Problem):
 
   - observation_space and action_space should be subclasses of gym.spaces
   - not all subclasses of gym.spaces are supported
-  - no support for continuous action spaces
 
   """
 
   def __init__(self,
                base_env_name=None,
                batch_size=None,
+               env_wrapper_fn=None,
                reward_range=(-np.inf, np.inf)):
     """Initializes this class by creating the envs and managing trajectories.
 
     Args:
-      base_env_name: (string) passed to `gym_utils.make_gym_env` to make the
-        underlying environment.
-      batch_size: (int or None): How many envs to make in the non natively
+      base_env_name: (string) passed to `gym.make` to make the underlying
+        environment.
+      batch_size: (int or None) How many envs to make in the non natively
         batched mode.
+      env_wrapper_fn: (callable(env): env) Applies gym wrappers to the base
+        environment.
       reward_range: (tuple(number, number)) the first element is the minimum
         reward and the second is the maximum reward, used to clip and process
         the raw reward in `process_rewards`.
@@ -142,7 +145,7 @@ def __init__(self,
     # Call the super's ctor.
     problem.Problem.__init__(self, was_reversed=False, was_copy=False)
 
-    # Name for the base environment, will be used in `gym_utils.make_gym_env` in
+    # Name for the base environment, will be used in `gym.make` in
     # the default implementation of `initialize_environments`.
     self._base_env_name = base_env_name
 
@@ -173,6 +176,8 @@ def __init__(self,
 
     self._batch_size = None
 
+    self._env_wrapper_fn = env_wrapper_fn
+
     if batch_size is not None:
       self.initialize(batch_size=batch_size)
 
@@ -239,8 +244,7 @@ def initialize(self, **kwargs):
     assert self._reward_range is not None
     assert self._trajectories is not None
 
-  def initialize_environments(self, batch_size=1, max_episode_steps=-1,
-                              max_and_skip_env=False):
+  def initialize_environments(self, batch_size=1):
     """Initializes the environments and trajectories.
 
     Subclasses can override this if they don't want a default implementation
@@ -249,19 +253,13 @@ def initialize_environments(self, batch_size=1, max_episode_steps=-1,
 
     Args:
       batch_size: (int) Number of `self.base_env_name` envs to initialize.
-      max_episode_steps: (int) Passed on to `gym_utils.make_gym_env`.
-      max_and_skip_env: (boolean) Passed on to `gym_utils.make_gym_env`.
     """
-
     assert batch_size >= 1
     self._batch_size = batch_size
 
-    # pylint: disable=g-complex-comprehension
-    self._envs = [
-        gym_utils.make_gym_env(
-            self.base_env_name,
-            rl_env_max_episode_steps=max_episode_steps,
-            maxskip_env=max_and_skip_env) for _ in range(batch_size)]
+    self._envs = [gym.make(self.base_env_name) for _ in range(batch_size)]
+    if self._env_wrapper_fn is not None:
+      self._envs = list(map(self._env_wrapper_fn, self._envs))
 
     # If self.observation_space and self.action_space aren't None, then it means
     # that this is a re-initialization of this class, in that case make sure
@@ -607,8 +605,8 @@ def hparams(self, defaults, model_hparams):
         "targets": self.target_modality,
         "input_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
         "target_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
-        "input_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
-        "target_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "input_action": self.action_modality,
+        "target_action": self.action_modality,
         "target_policy": modalities.ModalityType.IDENTITY,
         "target_value": modalities.ModalityType.IDENTITY,
     })
diff --git a/tensor2tensor/envs/env_problem_test.py b/tensor2tensor/envs/env_problem_test.py
index 03aa7ef74..e1b8f3e89 100644
--- a/tensor2tensor/envs/env_problem_test.py
+++ b/tensor2tensor/envs/env_problem_test.py
@@ -291,6 +291,10 @@ def target_modality(self):
       def target_vocab_size(self):
         return 2
 
+      @property
+      def action_modality(self):
+        return modalities.ModalityType.SYMBOL_WEIGHTS_ALL
+
     base_env_name = "CartPole-v0"
     batch_size = 5
     reward_range = (-1, 1)
diff --git a/tensor2tensor/envs/mujoco_problems.py b/tensor2tensor/envs/mujoco_problems.py
new file mode 100644
index 000000000..559c58bcc
--- /dev/null
+++ b/tensor2tensor/envs/mujoco_problems.py
@@ -0,0 +1,68 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Mujoco Gym environments."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+from tensor2tensor.envs import env_problem
+from tensor2tensor.layers import modalities
+from tensor2tensor.rl import gym_utils
+from tensor2tensor.utils import registry
+
+# BEGIN GOOGLE_INTERNAL
+import google3.robotics.reinforcement_learning.environments.gym_mujoco  # pylint: disable=unused-import
+# END GOOGLE_INTERNAL
+
+
+@registry.register_env_problem
+class ReacherEnvProblem(env_problem.EnvProblem):
+  """Mujoco's reacher environment."""
+
+  def __init__(self):
+    base_env_name = "Reacher-v2"
+    if base_env_name is None:
+      base_env_name = "Reacher-v2"
+    wrapper_fn = functools.partial(
+        gym_utils.gym_env_wrapper, **{
+            "rl_env_max_episode_steps": -1,
+            "maxskip_env": False,
+            "rendered_env": True
+        })
+    super(ReacherEnvProblem, self).__init__(
+        base_env_name=base_env_name, env_wrapper_fn=wrapper_fn)
+
+  @property
+  def input_modality(self):
+    return modalities.ModalityType.VIDEO
+
+  @property
+  def target_modality(self):
+    return modalities.ModalityType.VIDEO
+
+  @property
+  def action_modality(self):
+    return modalities.ModalityType.IDENTITY
+
+  @property
+  def input_vocab_size(self):
+    return 256
+
+  @property
+  def target_vocab_size(self):
+    return 256
diff --git a/tensor2tensor/envs/mujoco_problems_test.py b/tensor2tensor/envs/mujoco_problems_test.py
new file mode 100644
index 000000000..8d90f616a
--- /dev/null
+++ b/tensor2tensor/envs/mujoco_problems_test.py
@@ -0,0 +1,72 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.envs.mujoco_problems."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.envs import env_problem_utils
+from tensor2tensor.envs import mujoco_problems  # pylint: disable=unused-import
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+
+class ReacherEnvProblemTest(tf.test.TestCase):
+
+  def test_registration_and_interaction_with_env_problem(self):
+    batch_size = 5
+    # This ensures that registration has occurred.
+    ep = registry.env_problem("reacher_env_problem", batch_size=batch_size)
+    ep.reset()
+    num_done, num_lost, num_won, num_draw = 0, 0, 0, 0
+    nsteps = 100
+    for _ in range(nsteps):
+      actions = np.stack([ep.action_space.sample() for _ in range(batch_size)])
+      obs, rewards, dones, infos = ep.step(actions)
+
+      # Assert that things are happening batchwise.
+      self.assertEqual(batch_size, len(obs))
+      self.assertEqual(batch_size, len(rewards))
+      self.assertEqual(batch_size, len(dones))
+      self.assertEqual(batch_size, len(infos))
+
+      done_indices = env_problem_utils.done_indices(dones)
+      ep.reset(done_indices)
+      num_done += sum(dones)
+      for r, d in zip(rewards, dones):
+        if not d:
+          continue
+        if r == -1:
+          num_lost += 1
+        elif r == 0:
+          num_draw += 1
+        elif r == 1:
+          num_won += 1
+        else:
+          raise ValueError("reward should be -1, 0, 1 but is {}".format(r))
+
+    # Assert that something got done atleast, without that the next assert is
+    # meaningless.
+    self.assertGreater(num_done, 0)
+
+    # Assert that things are consistent.
+    self.assertEqual(num_done, num_won + num_lost + num_draw)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/envs/tic_tac_toe_env_problem.py b/tensor2tensor/envs/tic_tac_toe_env_problem.py
index 3a57a11ef..ae5b877d8 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem.py
@@ -50,3 +50,7 @@ def target_modality(self):
   def target_vocab_size(self):
     # Since reward is either -1 or 0 or +1.
     return 3
+
+  @property
+  def action_modality(self):
+    return modalities.ModalityType.SYMBOL_WEIGHTS_ALL
diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index cf8c0675e..9ba522d8f 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -32,8 +32,8 @@ def __init__(self, env, skip=4):
     gym.Wrapper.__init__(self, env)
     observation_space = env.observation_space
     # Most recent raw observations (for max pooling across time steps).
-    self._obs_buffer = np.zeros((2,) + observation_space.shape,
-                                dtype=observation_space.dtype)
+    self._obs_buffer = np.zeros(
+        (2,) + observation_space.shape, dtype=observation_space.dtype)
     self._skip = skip
 
   def __str__(self):
@@ -45,8 +45,10 @@ def step(self, action):
     done = None
     for i in range(self._skip):
       obs, reward, done, info = self.env.step(action)
-      if i == self._skip - 2: self._obs_buffer[0] = obs
-      if i == self._skip - 1: self._obs_buffer[1] = obs
+      if i == self._skip - 2:
+        self._obs_buffer[0] = obs
+      if i == self._skip - 1:
+        self._obs_buffer[1] = obs
       total_reward += reward
       if done:
         break
@@ -58,29 +60,34 @@ def reset(self, **kwargs):
     return self.env.reset(**kwargs)
 
 
-def make_gym_env(name, rl_env_max_episode_steps=-1, maxskip_env=False):
-  """Create a gym env optionally with a time limit and maxskip wrapper.
+class RenderedEnv(gym.Wrapper):
+  """Simple Env wrapper to override observations with rendered rgb values."""
 
-  NOTE: The returned env may already be wrapped with TimeLimit!
+  def __init__(self, env, mode="rgb_array", low=0, high=255):
+    gym.Wrapper.__init__(self, env)
+    # Get a sample frame to correctly set observation space
+    self.mode = mode
+    sample_frame = self.render(mode=self.mode)
+    assert sample_frame is not None
+    self.observation_space = gym.spaces.Box(
+        low=low, high=high, shape=sample_frame.shape, dtype=sample_frame.dtype)
 
-  Args:
-    name: `str` - base name of the gym env to make.
-    rl_env_max_episode_steps: `int` or None - Using any value < 0 returns the
-      env as-in, otherwise we impose the requested timelimit. Setting this to
-      None returns a wrapped env that doesn't have a step limit.
-    maxskip_env: whether to also use MaxAndSkip wrapper before time limit.
+  def step(self, action):
+    _, reward, done, info = self.env.step(action)
+    obs = self.env.render(mode=self.mode)
+    return obs, reward, done, info
 
-  Returns:
-    An instance of `gym.Env` or `gym.wrappers.TimeLimit` with the requested
-    step limit.
-  """
+  def reset(self, **kwargs):
+    self.env.reset(**kwargs)
+    return self.env.render(mode=self.mode)
 
+
+def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env):
+  """Wraps a gym environment. see make_gym_environment for details."""
   # rl_env_max_episode_steps is None or int.
   assert ((not rl_env_max_episode_steps) or
           isinstance(rl_env_max_episode_steps, int))
 
-  env = gym.make(name)
-
   # If nothing to do, then return the env.
   if rl_env_max_episode_steps and rl_env_max_episode_steps < 0:
     if maxskip_env:
@@ -88,8 +95,11 @@ def make_gym_env(name, rl_env_max_episode_steps=-1, maxskip_env=False):
         # Unwrap time limit and put it above MaxAndSkip for consistency.
         max_episode_steps = env._max_episode_steps  # pylint: disable=protected-access
         env = MaxAndSkipEnv(env.env)
-        return gym.wrappers.TimeLimit(env, max_episode_steps=max_episode_steps)
-      return MaxAndSkipEnv(env)
+        env = gym.wrappers.TimeLimit(env, max_episode_steps=max_episode_steps)
+      else:
+        env = MaxAndSkipEnv(env)
+    if rendered_env:
+      env = RenderedEnv(env)
     return env
 
   # Sometimes (mostly?) the env is already wrapped in a TimeLimit wrapper, in
@@ -100,9 +110,38 @@ def make_gym_env(name, rl_env_max_episode_steps=-1, maxskip_env=False):
   if maxskip_env:
     env = MaxAndSkipEnv(env)
 
+  if rendered_env:
+    env = RenderedEnv(env)
+
   return gym.wrappers.TimeLimit(env, max_episode_steps=rl_env_max_episode_steps)
 
 
+def make_gym_env(name,
+                 rl_env_max_episode_steps=-1,
+                 maxskip_env=False,
+                 rendered_env=False):
+  """Create a gym env optionally with a time limit and maxskip wrapper.
+
+  NOTE: The returned env may already be wrapped with TimeLimit!
+
+  Args:
+    name: `str` - base name of the gym env to make.
+    rl_env_max_episode_steps: `int` or None - Using any value < 0 returns the
+      env as-in, otherwise we impose the requested timelimit. Setting this to
+      None returns a wrapped env that doesn't have a step limit.
+    maxskip_env: whether to also use MaxAndSkip wrapper before time limit.
+    rendered_env: whether to force render for observations. Use this for
+      environments that are not natively rendering the scene for observations.
+
+  Returns:
+    An instance of `gym.Env` or `gym.wrappers.TimeLimit` with the requested
+    step limit.
+  """
+  env = gym.make(name)
+  return gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env,
+                         rendered_env)
+
+
 def register_gym_env(class_entry_point, version="v0", kwargs=None):
   """Registers the class in Gym and returns the registered name and the env."""
 
@@ -114,7 +153,7 @@ def register_gym_env(class_entry_point, version="v0", kwargs=None):
   env_name = "T2TEnv-{}-{}".format(class_name, version)
   gym.envs.register(id=env_name, entry_point=class_entry_point, kwargs=kwargs)
 
-  tf.logging.info("Entry Point [%s] registered with id [%s]",
-                  class_entry_point, env_name)
+  tf.logging.info("Entry Point [%s] registered with id [%s]", class_entry_point,
+                  env_name)
 
   return env_name, gym.make(env_name)

From 500668d3b67236978d9eb02a4f9323ee9d554e38 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 5 Mar 2019 21:46:17 -0800
Subject: [PATCH 1764/2720] Add Gym Mujoco envs to T2T.

PiperOrigin-RevId: 236980030
---
 tensor2tensor/envs/env_problem.py             | 38 +++++----
 tensor2tensor/envs/env_problem_test.py        |  4 -
 tensor2tensor/envs/mujoco_problems.py         | 68 ---------------
 tensor2tensor/envs/mujoco_problems_test.py    | 72 ----------------
 tensor2tensor/envs/tic_tac_toe_env_problem.py |  4 -
 tensor2tensor/rl/gym_utils.py                 | 85 +++++--------------
 6 files changed, 43 insertions(+), 228 deletions(-)
 delete mode 100644 tensor2tensor/envs/mujoco_problems.py
 delete mode 100644 tensor2tensor/envs/mujoco_problems_test.py

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 840eab902..4d7312001 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -23,7 +23,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import gym
 from gym.core import Env
 import numpy as np
 import six
@@ -32,6 +31,7 @@
 from tensor2tensor.envs import gym_spaces_utils
 from tensor2tensor.envs import trajectory
 from tensor2tensor.layers import modalities
+from tensor2tensor.rl import gym_utils
 import tensorflow as tf
 
 # Names for data fields in stored tf.Examples.
@@ -56,10 +56,9 @@ class EnvProblem(Env, problem.Problem):
   Subclasses *should* override the following functions, since they are used in
   the `hparams` function to return modalities and vocab_sizes.
   - input_modality
-  - input_vocab_size
+  - input_input_vocab_size
   - target_modality
   - target_vocab_size
-  - action_modality
 
   NON NATIVELY BATCHED ENVS:
 
@@ -120,23 +119,21 @@ class EnvProblem(Env, problem.Problem):
 
   - observation_space and action_space should be subclasses of gym.spaces
   - not all subclasses of gym.spaces are supported
+  - no support for continuous action spaces
 
   """
 
   def __init__(self,
                base_env_name=None,
                batch_size=None,
-               env_wrapper_fn=None,
                reward_range=(-np.inf, np.inf)):
     """Initializes this class by creating the envs and managing trajectories.
 
     Args:
-      base_env_name: (string) passed to `gym.make` to make the underlying
-        environment.
-      batch_size: (int or None) How many envs to make in the non natively
+      base_env_name: (string) passed to `gym_utils.make_gym_env` to make the
+        underlying environment.
+      batch_size: (int or None): How many envs to make in the non natively
         batched mode.
-      env_wrapper_fn: (callable(env): env) Applies gym wrappers to the base
-        environment.
       reward_range: (tuple(number, number)) the first element is the minimum
         reward and the second is the maximum reward, used to clip and process
         the raw reward in `process_rewards`.
@@ -145,7 +142,7 @@ def __init__(self,
     # Call the super's ctor.
     problem.Problem.__init__(self, was_reversed=False, was_copy=False)
 
-    # Name for the base environment, will be used in `gym.make` in
+    # Name for the base environment, will be used in `gym_utils.make_gym_env` in
     # the default implementation of `initialize_environments`.
     self._base_env_name = base_env_name
 
@@ -176,8 +173,6 @@ def __init__(self,
 
     self._batch_size = None
 
-    self._env_wrapper_fn = env_wrapper_fn
-
     if batch_size is not None:
       self.initialize(batch_size=batch_size)
 
@@ -244,7 +239,8 @@ def initialize(self, **kwargs):
     assert self._reward_range is not None
     assert self._trajectories is not None
 
-  def initialize_environments(self, batch_size=1):
+  def initialize_environments(self, batch_size=1, max_episode_steps=-1,
+                              max_and_skip_env=False):
     """Initializes the environments and trajectories.
 
     Subclasses can override this if they don't want a default implementation
@@ -253,13 +249,19 @@ def initialize_environments(self, batch_size=1):
 
     Args:
       batch_size: (int) Number of `self.base_env_name` envs to initialize.
+      max_episode_steps: (int) Passed on to `gym_utils.make_gym_env`.
+      max_and_skip_env: (boolean) Passed on to `gym_utils.make_gym_env`.
     """
+
     assert batch_size >= 1
     self._batch_size = batch_size
 
-    self._envs = [gym.make(self.base_env_name) for _ in range(batch_size)]
-    if self._env_wrapper_fn is not None:
-      self._envs = list(map(self._env_wrapper_fn, self._envs))
+    # pylint: disable=g-complex-comprehension
+    self._envs = [
+        gym_utils.make_gym_env(
+            self.base_env_name,
+            rl_env_max_episode_steps=max_episode_steps,
+            maxskip_env=max_and_skip_env) for _ in range(batch_size)]
 
     # If self.observation_space and self.action_space aren't None, then it means
     # that this is a re-initialization of this class, in that case make sure
@@ -605,8 +607,8 @@ def hparams(self, defaults, model_hparams):
         "targets": self.target_modality,
         "input_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
         "target_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
-        "input_action": self.action_modality,
-        "target_action": self.action_modality,
+        "input_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "target_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
         "target_policy": modalities.ModalityType.IDENTITY,
         "target_value": modalities.ModalityType.IDENTITY,
     })
diff --git a/tensor2tensor/envs/env_problem_test.py b/tensor2tensor/envs/env_problem_test.py
index e1b8f3e89..03aa7ef74 100644
--- a/tensor2tensor/envs/env_problem_test.py
+++ b/tensor2tensor/envs/env_problem_test.py
@@ -291,10 +291,6 @@ def target_modality(self):
       def target_vocab_size(self):
         return 2
 
-      @property
-      def action_modality(self):
-        return modalities.ModalityType.SYMBOL_WEIGHTS_ALL
-
     base_env_name = "CartPole-v0"
     batch_size = 5
     reward_range = (-1, 1)
diff --git a/tensor2tensor/envs/mujoco_problems.py b/tensor2tensor/envs/mujoco_problems.py
deleted file mode 100644
index 559c58bcc..000000000
--- a/tensor2tensor/envs/mujoco_problems.py
+++ /dev/null
@@ -1,68 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Mujoco Gym environments."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import functools
-from tensor2tensor.envs import env_problem
-from tensor2tensor.layers import modalities
-from tensor2tensor.rl import gym_utils
-from tensor2tensor.utils import registry
-
-# BEGIN GOOGLE_INTERNAL
-import google3.robotics.reinforcement_learning.environments.gym_mujoco  # pylint: disable=unused-import
-# END GOOGLE_INTERNAL
-
-
-@registry.register_env_problem
-class ReacherEnvProblem(env_problem.EnvProblem):
-  """Mujoco's reacher environment."""
-
-  def __init__(self):
-    base_env_name = "Reacher-v2"
-    if base_env_name is None:
-      base_env_name = "Reacher-v2"
-    wrapper_fn = functools.partial(
-        gym_utils.gym_env_wrapper, **{
-            "rl_env_max_episode_steps": -1,
-            "maxskip_env": False,
-            "rendered_env": True
-        })
-    super(ReacherEnvProblem, self).__init__(
-        base_env_name=base_env_name, env_wrapper_fn=wrapper_fn)
-
-  @property
-  def input_modality(self):
-    return modalities.ModalityType.VIDEO
-
-  @property
-  def target_modality(self):
-    return modalities.ModalityType.VIDEO
-
-  @property
-  def action_modality(self):
-    return modalities.ModalityType.IDENTITY
-
-  @property
-  def input_vocab_size(self):
-    return 256
-
-  @property
-  def target_vocab_size(self):
-    return 256
diff --git a/tensor2tensor/envs/mujoco_problems_test.py b/tensor2tensor/envs/mujoco_problems_test.py
deleted file mode 100644
index 8d90f616a..000000000
--- a/tensor2tensor/envs/mujoco_problems_test.py
+++ /dev/null
@@ -1,72 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.envs.mujoco_problems."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-from tensor2tensor.envs import env_problem_utils
-from tensor2tensor.envs import mujoco_problems  # pylint: disable=unused-import
-from tensor2tensor.utils import registry
-import tensorflow as tf
-
-
-class ReacherEnvProblemTest(tf.test.TestCase):
-
-  def test_registration_and_interaction_with_env_problem(self):
-    batch_size = 5
-    # This ensures that registration has occurred.
-    ep = registry.env_problem("reacher_env_problem", batch_size=batch_size)
-    ep.reset()
-    num_done, num_lost, num_won, num_draw = 0, 0, 0, 0
-    nsteps = 100
-    for _ in range(nsteps):
-      actions = np.stack([ep.action_space.sample() for _ in range(batch_size)])
-      obs, rewards, dones, infos = ep.step(actions)
-
-      # Assert that things are happening batchwise.
-      self.assertEqual(batch_size, len(obs))
-      self.assertEqual(batch_size, len(rewards))
-      self.assertEqual(batch_size, len(dones))
-      self.assertEqual(batch_size, len(infos))
-
-      done_indices = env_problem_utils.done_indices(dones)
-      ep.reset(done_indices)
-      num_done += sum(dones)
-      for r, d in zip(rewards, dones):
-        if not d:
-          continue
-        if r == -1:
-          num_lost += 1
-        elif r == 0:
-          num_draw += 1
-        elif r == 1:
-          num_won += 1
-        else:
-          raise ValueError("reward should be -1, 0, 1 but is {}".format(r))
-
-    # Assert that something got done atleast, without that the next assert is
-    # meaningless.
-    self.assertGreater(num_done, 0)
-
-    # Assert that things are consistent.
-    self.assertEqual(num_done, num_won + num_lost + num_draw)
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/envs/tic_tac_toe_env_problem.py b/tensor2tensor/envs/tic_tac_toe_env_problem.py
index ae5b877d8..3a57a11ef 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem.py
@@ -50,7 +50,3 @@ def target_modality(self):
   def target_vocab_size(self):
     # Since reward is either -1 or 0 or +1.
     return 3
-
-  @property
-  def action_modality(self):
-    return modalities.ModalityType.SYMBOL_WEIGHTS_ALL
diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index 9ba522d8f..cf8c0675e 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -32,8 +32,8 @@ def __init__(self, env, skip=4):
     gym.Wrapper.__init__(self, env)
     observation_space = env.observation_space
     # Most recent raw observations (for max pooling across time steps).
-    self._obs_buffer = np.zeros(
-        (2,) + observation_space.shape, dtype=observation_space.dtype)
+    self._obs_buffer = np.zeros((2,) + observation_space.shape,
+                                dtype=observation_space.dtype)
     self._skip = skip
 
   def __str__(self):
@@ -45,10 +45,8 @@ def step(self, action):
     done = None
     for i in range(self._skip):
       obs, reward, done, info = self.env.step(action)
-      if i == self._skip - 2:
-        self._obs_buffer[0] = obs
-      if i == self._skip - 1:
-        self._obs_buffer[1] = obs
+      if i == self._skip - 2: self._obs_buffer[0] = obs
+      if i == self._skip - 1: self._obs_buffer[1] = obs
       total_reward += reward
       if done:
         break
@@ -60,34 +58,29 @@ def reset(self, **kwargs):
     return self.env.reset(**kwargs)
 
 
-class RenderedEnv(gym.Wrapper):
-  """Simple Env wrapper to override observations with rendered rgb values."""
-
-  def __init__(self, env, mode="rgb_array", low=0, high=255):
-    gym.Wrapper.__init__(self, env)
-    # Get a sample frame to correctly set observation space
-    self.mode = mode
-    sample_frame = self.render(mode=self.mode)
-    assert sample_frame is not None
-    self.observation_space = gym.spaces.Box(
-        low=low, high=high, shape=sample_frame.shape, dtype=sample_frame.dtype)
+def make_gym_env(name, rl_env_max_episode_steps=-1, maxskip_env=False):
+  """Create a gym env optionally with a time limit and maxskip wrapper.
 
-  def step(self, action):
-    _, reward, done, info = self.env.step(action)
-    obs = self.env.render(mode=self.mode)
-    return obs, reward, done, info
+  NOTE: The returned env may already be wrapped with TimeLimit!
 
-  def reset(self, **kwargs):
-    self.env.reset(**kwargs)
-    return self.env.render(mode=self.mode)
+  Args:
+    name: `str` - base name of the gym env to make.
+    rl_env_max_episode_steps: `int` or None - Using any value < 0 returns the
+      env as-in, otherwise we impose the requested timelimit. Setting this to
+      None returns a wrapped env that doesn't have a step limit.
+    maxskip_env: whether to also use MaxAndSkip wrapper before time limit.
 
+  Returns:
+    An instance of `gym.Env` or `gym.wrappers.TimeLimit` with the requested
+    step limit.
+  """
 
-def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env):
-  """Wraps a gym environment. see make_gym_environment for details."""
   # rl_env_max_episode_steps is None or int.
   assert ((not rl_env_max_episode_steps) or
           isinstance(rl_env_max_episode_steps, int))
 
+  env = gym.make(name)
+
   # If nothing to do, then return the env.
   if rl_env_max_episode_steps and rl_env_max_episode_steps < 0:
     if maxskip_env:
@@ -95,11 +88,8 @@ def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env):
         # Unwrap time limit and put it above MaxAndSkip for consistency.
         max_episode_steps = env._max_episode_steps  # pylint: disable=protected-access
         env = MaxAndSkipEnv(env.env)
-        env = gym.wrappers.TimeLimit(env, max_episode_steps=max_episode_steps)
-      else:
-        env = MaxAndSkipEnv(env)
-    if rendered_env:
-      env = RenderedEnv(env)
+        return gym.wrappers.TimeLimit(env, max_episode_steps=max_episode_steps)
+      return MaxAndSkipEnv(env)
     return env
 
   # Sometimes (mostly?) the env is already wrapped in a TimeLimit wrapper, in
@@ -110,38 +100,9 @@ def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env):
   if maxskip_env:
     env = MaxAndSkipEnv(env)
 
-  if rendered_env:
-    env = RenderedEnv(env)
-
   return gym.wrappers.TimeLimit(env, max_episode_steps=rl_env_max_episode_steps)
 
 
-def make_gym_env(name,
-                 rl_env_max_episode_steps=-1,
-                 maxskip_env=False,
-                 rendered_env=False):
-  """Create a gym env optionally with a time limit and maxskip wrapper.
-
-  NOTE: The returned env may already be wrapped with TimeLimit!
-
-  Args:
-    name: `str` - base name of the gym env to make.
-    rl_env_max_episode_steps: `int` or None - Using any value < 0 returns the
-      env as-in, otherwise we impose the requested timelimit. Setting this to
-      None returns a wrapped env that doesn't have a step limit.
-    maxskip_env: whether to also use MaxAndSkip wrapper before time limit.
-    rendered_env: whether to force render for observations. Use this for
-      environments that are not natively rendering the scene for observations.
-
-  Returns:
-    An instance of `gym.Env` or `gym.wrappers.TimeLimit` with the requested
-    step limit.
-  """
-  env = gym.make(name)
-  return gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env,
-                         rendered_env)
-
-
 def register_gym_env(class_entry_point, version="v0", kwargs=None):
   """Registers the class in Gym and returns the registered name and the env."""
 
@@ -153,7 +114,7 @@ def register_gym_env(class_entry_point, version="v0", kwargs=None):
   env_name = "T2TEnv-{}-{}".format(class_name, version)
   gym.envs.register(id=env_name, entry_point=class_entry_point, kwargs=kwargs)
 
-  tf.logging.info("Entry Point [%s] registered with id [%s]", class_entry_point,
-                  env_name)
+  tf.logging.info("Entry Point [%s] registered with id [%s]",
+                  class_entry_point, env_name)
 
   return env_name, gym.make(env_name)

From ef7fb5539da9f649ac2d6bbd9020473bf4a1a89d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 6 Mar 2019 10:37:13 -0800
Subject: [PATCH 1765/2720] Follow convention to call the jit-ted version of
 function `f` as `jit_f`.  - s/jit_predict/jit_model_predict  -
 s/update_fun/jit_update_fun

PiperOrigin-RevId: 237075442
---
 tensor2tensor/trax/trax.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 67730603a..98c5330f9 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -274,8 +274,8 @@ def train(output_dir,
   opt_state = opt_init(params)
 
   # jit model_predict and update so they're fast
-  jit_predict = jax.jit(model_predict)  # for evaluation
-  update_fun = _jit_update_fun(model_predict, loss, optimizer, lr_fun)
+  jit_model_predict = jax.jit(model_predict)  # for evaluation
+  jit_update_fun = _jit_update_fun(model_predict, loss, optimizer, lr_fun)
 
   print()
   train_stream = inputs.train_stream()
@@ -293,7 +293,7 @@ def train(output_dir,
 
     for _ in range(epoch_steps):
       # Train
-      opt_state = update_fun(step, opt_state, next(train_stream))
+      opt_state = jit_update_fun(step, opt_state, next(train_stream))
       step += 1
 
       # LR log
@@ -314,7 +314,7 @@ def train(output_dir,
     evaluate_train_and_eval(
         step=step,
         inputs=inputs,
-        predict_fun=functools.partial(jit_predict, params),
+        predict_fun=functools.partial(jit_model_predict, params),
         eval_steps=eval_steps,
         train_sw=train_sw,
         eval_sw=eval_sw,
@@ -332,7 +332,7 @@ def train(output_dir,
     old_lr_fun = lr_fun
     lr_fun = lr_schedule(history)
     if lr_fun != old_lr_fun:  # For performance, only jit if there is a change.
-      update_fun = _jit_update_fun(model_predict, loss, optimizer, lr_fun)
+      jit_update_fun = _jit_update_fun(model_predict, loss, optimizer, lr_fun)
 
     # Flush summary writers
     train_sw.writer.flush()

From 1c5522f72cfcb0a83185df1f38043cc68bf96e83 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 6 Mar 2019 12:08:39 -0800
Subject: [PATCH 1766/2720] Add Gym Mujoco envs to T2T.

This introduces the following changes:
1- Correctly sets action modalities.
2- remove dependency to gym_make_env and use dependency injection for wrappers.
3- Implements RenderedEnv Gym wrapper for environments that are not natively doing it for their observations.
4- Add reacher Gym environment problem.

PiperOrigin-RevId: 237094163
---
 tensor2tensor/envs/env_problem.py             | 38 ++++-----
 tensor2tensor/envs/env_problem_test.py        |  4 +
 tensor2tensor/envs/mujoco_problems.py         | 63 ++++++++++++++
 tensor2tensor/envs/mujoco_problems_test.py    | 72 ++++++++++++++++
 tensor2tensor/envs/tic_tac_toe_env_problem.py |  4 +
 tensor2tensor/rl/gym_utils.py                 | 85 ++++++++++++++-----
 6 files changed, 223 insertions(+), 43 deletions(-)
 create mode 100644 tensor2tensor/envs/mujoco_problems.py
 create mode 100644 tensor2tensor/envs/mujoco_problems_test.py

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 4d7312001..840eab902 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -23,6 +23,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import gym
 from gym.core import Env
 import numpy as np
 import six
@@ -31,7 +32,6 @@
 from tensor2tensor.envs import gym_spaces_utils
 from tensor2tensor.envs import trajectory
 from tensor2tensor.layers import modalities
-from tensor2tensor.rl import gym_utils
 import tensorflow as tf
 
 # Names for data fields in stored tf.Examples.
@@ -56,9 +56,10 @@ class EnvProblem(Env, problem.Problem):
   Subclasses *should* override the following functions, since they are used in
   the `hparams` function to return modalities and vocab_sizes.
   - input_modality
-  - input_input_vocab_size
+  - input_vocab_size
   - target_modality
   - target_vocab_size
+  - action_modality
 
   NON NATIVELY BATCHED ENVS:
 
@@ -119,21 +120,23 @@ class EnvProblem(Env, problem.Problem):
 
   - observation_space and action_space should be subclasses of gym.spaces
   - not all subclasses of gym.spaces are supported
-  - no support for continuous action spaces
 
   """
 
   def __init__(self,
                base_env_name=None,
                batch_size=None,
+               env_wrapper_fn=None,
                reward_range=(-np.inf, np.inf)):
     """Initializes this class by creating the envs and managing trajectories.
 
     Args:
-      base_env_name: (string) passed to `gym_utils.make_gym_env` to make the
-        underlying environment.
-      batch_size: (int or None): How many envs to make in the non natively
+      base_env_name: (string) passed to `gym.make` to make the underlying
+        environment.
+      batch_size: (int or None) How many envs to make in the non natively
         batched mode.
+      env_wrapper_fn: (callable(env): env) Applies gym wrappers to the base
+        environment.
       reward_range: (tuple(number, number)) the first element is the minimum
         reward and the second is the maximum reward, used to clip and process
         the raw reward in `process_rewards`.
@@ -142,7 +145,7 @@ def __init__(self,
     # Call the super's ctor.
     problem.Problem.__init__(self, was_reversed=False, was_copy=False)
 
-    # Name for the base environment, will be used in `gym_utils.make_gym_env` in
+    # Name for the base environment, will be used in `gym.make` in
     # the default implementation of `initialize_environments`.
     self._base_env_name = base_env_name
 
@@ -173,6 +176,8 @@ def __init__(self,
 
     self._batch_size = None
 
+    self._env_wrapper_fn = env_wrapper_fn
+
     if batch_size is not None:
       self.initialize(batch_size=batch_size)
 
@@ -239,8 +244,7 @@ def initialize(self, **kwargs):
     assert self._reward_range is not None
     assert self._trajectories is not None
 
-  def initialize_environments(self, batch_size=1, max_episode_steps=-1,
-                              max_and_skip_env=False):
+  def initialize_environments(self, batch_size=1):
     """Initializes the environments and trajectories.
 
     Subclasses can override this if they don't want a default implementation
@@ -249,19 +253,13 @@ def initialize_environments(self, batch_size=1, max_episode_steps=-1,
 
     Args:
       batch_size: (int) Number of `self.base_env_name` envs to initialize.
-      max_episode_steps: (int) Passed on to `gym_utils.make_gym_env`.
-      max_and_skip_env: (boolean) Passed on to `gym_utils.make_gym_env`.
     """
-
     assert batch_size >= 1
     self._batch_size = batch_size
 
-    # pylint: disable=g-complex-comprehension
-    self._envs = [
-        gym_utils.make_gym_env(
-            self.base_env_name,
-            rl_env_max_episode_steps=max_episode_steps,
-            maxskip_env=max_and_skip_env) for _ in range(batch_size)]
+    self._envs = [gym.make(self.base_env_name) for _ in range(batch_size)]
+    if self._env_wrapper_fn is not None:
+      self._envs = list(map(self._env_wrapper_fn, self._envs))
 
     # If self.observation_space and self.action_space aren't None, then it means
     # that this is a re-initialization of this class, in that case make sure
@@ -607,8 +605,8 @@ def hparams(self, defaults, model_hparams):
         "targets": self.target_modality,
         "input_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
         "target_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
-        "input_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
-        "target_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "input_action": self.action_modality,
+        "target_action": self.action_modality,
         "target_policy": modalities.ModalityType.IDENTITY,
         "target_value": modalities.ModalityType.IDENTITY,
     })
diff --git a/tensor2tensor/envs/env_problem_test.py b/tensor2tensor/envs/env_problem_test.py
index 03aa7ef74..e1b8f3e89 100644
--- a/tensor2tensor/envs/env_problem_test.py
+++ b/tensor2tensor/envs/env_problem_test.py
@@ -291,6 +291,10 @@ def target_modality(self):
       def target_vocab_size(self):
         return 2
 
+      @property
+      def action_modality(self):
+        return modalities.ModalityType.SYMBOL_WEIGHTS_ALL
+
     base_env_name = "CartPole-v0"
     batch_size = 5
     reward_range = (-1, 1)
diff --git a/tensor2tensor/envs/mujoco_problems.py b/tensor2tensor/envs/mujoco_problems.py
new file mode 100644
index 000000000..a172aa69f
--- /dev/null
+++ b/tensor2tensor/envs/mujoco_problems.py
@@ -0,0 +1,63 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Mujoco Gym environments."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+from tensor2tensor.envs import env_problem
+from tensor2tensor.layers import modalities
+from tensor2tensor.rl import gym_utils
+from tensor2tensor.utils import registry
+
+
+
+@registry.register_env_problem
+class ReacherEnvProblem(env_problem.EnvProblem):
+  """Mujoco's reacher environment."""
+
+  def __init__(self):
+    base_env_name = "Reacher-v2"
+    wrapper_fn = functools.partial(
+        gym_utils.gym_env_wrapper, **{
+            "rl_env_max_episode_steps": -1,
+            "maxskip_env": False,
+            "rendered_env": True
+        })
+    super(ReacherEnvProblem, self).__init__(
+        base_env_name=base_env_name, env_wrapper_fn=wrapper_fn)
+
+  @property
+  def input_modality(self):
+    return modalities.ModalityType.VIDEO
+
+  @property
+  def target_modality(self):
+    return modalities.ModalityType.VIDEO
+
+  @property
+  def action_modality(self):
+    return modalities.ModalityType.IDENTITY
+
+  @property
+  def input_vocab_size(self):
+    return 256
+
+  @property
+  def target_vocab_size(self):
+    return 256
diff --git a/tensor2tensor/envs/mujoco_problems_test.py b/tensor2tensor/envs/mujoco_problems_test.py
new file mode 100644
index 000000000..8d90f616a
--- /dev/null
+++ b/tensor2tensor/envs/mujoco_problems_test.py
@@ -0,0 +1,72 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.envs.mujoco_problems."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.envs import env_problem_utils
+from tensor2tensor.envs import mujoco_problems  # pylint: disable=unused-import
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+
+class ReacherEnvProblemTest(tf.test.TestCase):
+
+  def test_registration_and_interaction_with_env_problem(self):
+    batch_size = 5
+    # This ensures that registration has occurred.
+    ep = registry.env_problem("reacher_env_problem", batch_size=batch_size)
+    ep.reset()
+    num_done, num_lost, num_won, num_draw = 0, 0, 0, 0
+    nsteps = 100
+    for _ in range(nsteps):
+      actions = np.stack([ep.action_space.sample() for _ in range(batch_size)])
+      obs, rewards, dones, infos = ep.step(actions)
+
+      # Assert that things are happening batchwise.
+      self.assertEqual(batch_size, len(obs))
+      self.assertEqual(batch_size, len(rewards))
+      self.assertEqual(batch_size, len(dones))
+      self.assertEqual(batch_size, len(infos))
+
+      done_indices = env_problem_utils.done_indices(dones)
+      ep.reset(done_indices)
+      num_done += sum(dones)
+      for r, d in zip(rewards, dones):
+        if not d:
+          continue
+        if r == -1:
+          num_lost += 1
+        elif r == 0:
+          num_draw += 1
+        elif r == 1:
+          num_won += 1
+        else:
+          raise ValueError("reward should be -1, 0, 1 but is {}".format(r))
+
+    # Assert that something got done atleast, without that the next assert is
+    # meaningless.
+    self.assertGreater(num_done, 0)
+
+    # Assert that things are consistent.
+    self.assertEqual(num_done, num_won + num_lost + num_draw)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/envs/tic_tac_toe_env_problem.py b/tensor2tensor/envs/tic_tac_toe_env_problem.py
index 3a57a11ef..ae5b877d8 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem.py
@@ -50,3 +50,7 @@ def target_modality(self):
   def target_vocab_size(self):
     # Since reward is either -1 or 0 or +1.
     return 3
+
+  @property
+  def action_modality(self):
+    return modalities.ModalityType.SYMBOL_WEIGHTS_ALL
diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index cf8c0675e..9ba522d8f 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -32,8 +32,8 @@ def __init__(self, env, skip=4):
     gym.Wrapper.__init__(self, env)
     observation_space = env.observation_space
     # Most recent raw observations (for max pooling across time steps).
-    self._obs_buffer = np.zeros((2,) + observation_space.shape,
-                                dtype=observation_space.dtype)
+    self._obs_buffer = np.zeros(
+        (2,) + observation_space.shape, dtype=observation_space.dtype)
     self._skip = skip
 
   def __str__(self):
@@ -45,8 +45,10 @@ def step(self, action):
     done = None
     for i in range(self._skip):
       obs, reward, done, info = self.env.step(action)
-      if i == self._skip - 2: self._obs_buffer[0] = obs
-      if i == self._skip - 1: self._obs_buffer[1] = obs
+      if i == self._skip - 2:
+        self._obs_buffer[0] = obs
+      if i == self._skip - 1:
+        self._obs_buffer[1] = obs
       total_reward += reward
       if done:
         break
@@ -58,29 +60,34 @@ def reset(self, **kwargs):
     return self.env.reset(**kwargs)
 
 
-def make_gym_env(name, rl_env_max_episode_steps=-1, maxskip_env=False):
-  """Create a gym env optionally with a time limit and maxskip wrapper.
+class RenderedEnv(gym.Wrapper):
+  """Simple Env wrapper to override observations with rendered rgb values."""
 
-  NOTE: The returned env may already be wrapped with TimeLimit!
+  def __init__(self, env, mode="rgb_array", low=0, high=255):
+    gym.Wrapper.__init__(self, env)
+    # Get a sample frame to correctly set observation space
+    self.mode = mode
+    sample_frame = self.render(mode=self.mode)
+    assert sample_frame is not None
+    self.observation_space = gym.spaces.Box(
+        low=low, high=high, shape=sample_frame.shape, dtype=sample_frame.dtype)
 
-  Args:
-    name: `str` - base name of the gym env to make.
-    rl_env_max_episode_steps: `int` or None - Using any value < 0 returns the
-      env as-in, otherwise we impose the requested timelimit. Setting this to
-      None returns a wrapped env that doesn't have a step limit.
-    maxskip_env: whether to also use MaxAndSkip wrapper before time limit.
+  def step(self, action):
+    _, reward, done, info = self.env.step(action)
+    obs = self.env.render(mode=self.mode)
+    return obs, reward, done, info
 
-  Returns:
-    An instance of `gym.Env` or `gym.wrappers.TimeLimit` with the requested
-    step limit.
-  """
+  def reset(self, **kwargs):
+    self.env.reset(**kwargs)
+    return self.env.render(mode=self.mode)
 
+
+def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env):
+  """Wraps a gym environment. see make_gym_environment for details."""
   # rl_env_max_episode_steps is None or int.
   assert ((not rl_env_max_episode_steps) or
           isinstance(rl_env_max_episode_steps, int))
 
-  env = gym.make(name)
-
   # If nothing to do, then return the env.
   if rl_env_max_episode_steps and rl_env_max_episode_steps < 0:
     if maxskip_env:
@@ -88,8 +95,11 @@ def make_gym_env(name, rl_env_max_episode_steps=-1, maxskip_env=False):
         # Unwrap time limit and put it above MaxAndSkip for consistency.
         max_episode_steps = env._max_episode_steps  # pylint: disable=protected-access
         env = MaxAndSkipEnv(env.env)
-        return gym.wrappers.TimeLimit(env, max_episode_steps=max_episode_steps)
-      return MaxAndSkipEnv(env)
+        env = gym.wrappers.TimeLimit(env, max_episode_steps=max_episode_steps)
+      else:
+        env = MaxAndSkipEnv(env)
+    if rendered_env:
+      env = RenderedEnv(env)
     return env
 
   # Sometimes (mostly?) the env is already wrapped in a TimeLimit wrapper, in
@@ -100,9 +110,38 @@ def make_gym_env(name, rl_env_max_episode_steps=-1, maxskip_env=False):
   if maxskip_env:
     env = MaxAndSkipEnv(env)
 
+  if rendered_env:
+    env = RenderedEnv(env)
+
   return gym.wrappers.TimeLimit(env, max_episode_steps=rl_env_max_episode_steps)
 
 
+def make_gym_env(name,
+                 rl_env_max_episode_steps=-1,
+                 maxskip_env=False,
+                 rendered_env=False):
+  """Create a gym env optionally with a time limit and maxskip wrapper.
+
+  NOTE: The returned env may already be wrapped with TimeLimit!
+
+  Args:
+    name: `str` - base name of the gym env to make.
+    rl_env_max_episode_steps: `int` or None - Using any value < 0 returns the
+      env as-in, otherwise we impose the requested timelimit. Setting this to
+      None returns a wrapped env that doesn't have a step limit.
+    maxskip_env: whether to also use MaxAndSkip wrapper before time limit.
+    rendered_env: whether to force render for observations. Use this for
+      environments that are not natively rendering the scene for observations.
+
+  Returns:
+    An instance of `gym.Env` or `gym.wrappers.TimeLimit` with the requested
+    step limit.
+  """
+  env = gym.make(name)
+  return gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env,
+                         rendered_env)
+
+
 def register_gym_env(class_entry_point, version="v0", kwargs=None):
   """Registers the class in Gym and returns the registered name and the env."""
 
@@ -114,7 +153,7 @@ def register_gym_env(class_entry_point, version="v0", kwargs=None):
   env_name = "T2TEnv-{}-{}".format(class_name, version)
   gym.envs.register(id=env_name, entry_point=class_entry_point, kwargs=kwargs)
 
-  tf.logging.info("Entry Point [%s] registered with id [%s]",
-                  class_entry_point, env_name)
+  tf.logging.info("Entry Point [%s] registered with id [%s]", class_entry_point,
+                  env_name)
 
   return env_name, gym.make(env_name)

From 19cd29e68e82f176ec6a493314385c12728ca3cb Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 6 Mar 2019 13:09:34 -0800
Subject: [PATCH 1767/2720] Organize feature transformations for modalities.

They're grouped by bottom, loss, top followed by alphabetical order.

PiperOrigin-RevId: 237105482
---
 tensor2tensor/layers/modalities.py | 1384 ++++++++++++++--------------
 1 file changed, 697 insertions(+), 687 deletions(-)

diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index c8ee107f7..21409ad73 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -17,7 +17,7 @@
 
 T2TModel applies a default transformation to each feature according to its
 modality. Override them by specifying a model's
-hparams.{bottom,loss,name,top,weights_fn}.
+hparams.{bottom,loss,top,weights_fn}.
 """
 from __future__ import absolute_import
 from __future__ import division
@@ -36,209 +36,216 @@
 import tensorflow_probability as tfp
 
 
-def is_pointwise(func):
-  """Decorator for whether the function is pointwise.
-
-  An example of a pointwise function is a linear layer followed by
-  a softmax. Given a tensor [batch, length, height, depth] it operates
-  only on the last axis, on every point in [batch, length, height] fully
-  independently. In contrast, a classifier that first averages over length
-  and height is not pointwise, as it depends on the whole field. It is useful
-  to know if top functions are pointwise to speed up decoding in certain models.
-
-  Args:
-    func: Function to decorate.
-
-  Returns:
-    Original function with an attribute pointwise set to True.
-  """
-  func.pointwise = True
-  return func
+class ModalityType(object):
+  """Types of modalities."""
 
+  AUDIO = "audio"
+  AUDIO_SPECTRAL = "audio_spectral"
+  CLASS_LABEL = "class_label"
+  CTC_SYMBOL = "ctc_symbol"  # symbol with CTC loss
+  GENERIC_L2_LOSS = "generic_l2"  # identity modality with L2 loss
+  IDENTITY = "identity"  # identity top and bottom
+  IDENTITY_SYMBOL = "identity_symbol"  # symbol with identity top and bottom
+  IMAGE = "image"
+  # images using channel compression for generation
+  IMAGE_CHANNEL_BOTTOM_IDENTITY = "image_channel_bottom_identity"
+  # images using channel compression for generation
+  IMAGE_CHANNEL_COMPRESS = "image_channel_compress"
+  IMAGE_CHANNEL_EMBEDDINGS_BOTTOM = "image_channel_embeddings_bottom"
+  MULTI_LABEL = "multi_label"
+  ONE_HOT_CLASS_LABEL = "one_hot_class_label"
+  REAL = "real"  # real vectors
+  REAL_L2_LOSS = "real_l2"  # real vectors with L2 as loss
+  # real vectors with log Poisson regression loss
+  REAL_LOG_POISSON_LOSS = "real_log_poisson"
+  SIGMOID_CLASS_LABEL = "sigmoid_class_label"  # sigmoid cross-entropy loss
+  # sigmoid cross-entropy applied on max-pooling over timesteps
+  SIGMOID_MAX_POOLING_CLASS_LABEL = "sigmoid_max_pooling_class_label"
+  # softmax cross-entropy applied on average-pooling over timesteps
+  SOFTMAX_AVERAGE_POOLING_CLASS_LABEL = "softmax_average_pooling_class_label"
+  # softmax cross-entropy applied on last-timestep encoding
+  SOFTMAX_LAST_TIMESTEP_CLASS_LABEL = "softmax_last_timestep_class_label"
+  # softmax cross-entropy applied on max-pooling over timesteps
+  SOFTMAX_MAX_POOLING_CLASS_LABEL = "softmax_max_pooling_class_label"
+  SPEECH_RECOGNITION = "speech_recognition"
+  SYMBOL = "symbol"
+  SYMBOL_WEIGHTS_ALL = "symbol_weights_all"  # symbol for features w/o 0-padding
+  SYMBOL_ONE_HOT = "symbol_one_hot"  # symbol with one hot as embeddings
+  VIDEO = "video"
+  VIDEO_BITWISE = "video_bitwise"  # video where bottom embeds pixels bitwise
+  VIDEO_IDENTITY = "video_identity"  # video with identity top and bottom
+  VIDEO_L1 = "video_l1"  # video with L2 loss
+  VIDEO_L2 = "video_l2"  # video with L1 loss
+  # video with L1 loss and raw input (sequences of frames)
+  VIDEO_L1_RAW = "video_l1_raw"
+  # video with L2 loss and raw input (sequences of frames)
+  VIDEO_L2_RAW = "video_l2_raw"
+  # video with pixel noise on input during training
+  VIDEO_PIXEL_NOISE = "video_pixel_noise"
 
-def generic_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-  """Compute loss numerator and denominator for one shard of output."""
-  del vocab_size  # unused arg
-  logits = top_out
-  logits = common_attention.maybe_upcast(logits, hparams=model_hparams)
-  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.0)
-  return common_layers.padded_cross_entropy(
-      logits,
-      targets,
-      model_hparams.label_smoothing,
-      cutoff=cutoff,
-      weights_fn=weights_fn)
+  @staticmethod
+  def get_choices():
+    return [
+        ModalityType.AUDIO,
+        ModalityType.AUDIO_SPECTRAL,
+        ModalityType.CLASS_LABEL,
+        ModalityType.CTC_SYMBOL,
+        ModalityType.GENERIC_L2_LOSS,
+        ModalityType.IDENTITY,
+        ModalityType.IDENTITY_SYMBOL,
+        ModalityType.IMAGE,
+        ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
+        ModalityType.IMAGE_CHANNEL_COMPRESS,
+        ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM,
+        ModalityType.MULTI_LABEL,
+        ModalityType.ONE_HOT_CLASS_LABEL,
+        ModalityType.REAL,
+        ModalityType.REAL_L2_LOSS,
+        ModalityType.REAL_LOG_POISSON_LOSS,
+        ModalityType.SIGMOID_CLASS_LABEL,
+        ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
+        ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
+        ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
+        ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL,
+        ModalityType.SPEECH_RECOGNITION,
+        ModalityType.SYMBOL,
+        ModalityType.SYMBOL_ONE_HOT,
+        ModalityType.SYMBOL_WEIGHTS_ALL,
+        ModalityType.VIDEO,
+        ModalityType.VIDEO_BITWISE,
+        ModalityType.VIDEO_IDENTITY,
+        ModalityType.VIDEO_L1,
+        ModalityType.VIDEO_L2,
+        ModalityType.VIDEO_L1_RAW,
+        ModalityType.VIDEO_L2_RAW,
+        ModalityType.VIDEO_PIXEL_NOISE,
+    ]
 
 
-def make_targets_bottom(bottom):
-  def targets_bottom(x, model_hparams, vocab_size):
-    with tf.variable_scope("targets_bottom"):
-      return bottom(x, model_hparams, vocab_size)
-  return targets_bottom
+# Bottom transformations, applied to all features
 
 
-def get_weights(model_hparams, vocab_size, hidden_dim=None):
-  """Create or get concatenated embedding or softmax variable.
+def audio_bottom(x, model_hparams, vocab_size):
+  """Transform input from data space to model space.
 
   Args:
+    x: A Tensor with shape [batch, ...]
     model_hparams: tf.HParams, model hyperparmeters.
     vocab_size: int, vocabulary size.
-    hidden_dim: dim of the variable. Defaults to _model_hparams' hidden_size
 
   Returns:
-     a list of num_shards Tensors.
+    body_input: A Tensor with shape [batch, ?, ?,
+      model_hparams.hidden_size].
   """
-  if hidden_dim is None:
-    hidden_dim = model_hparams.hidden_size
-  num_shards = model_hparams.symbol_modality_num_shards
-  shards = []
-  for i in range(num_shards):
-    shard_size = (vocab_size // num_shards) + (
-        1 if i < vocab_size % num_shards else 0)
-    var_name = "weights_%d" % i
-    shards.append(
-        tf.get_variable(
-            var_name, [shard_size, hidden_dim],
-            initializer=tf.random_normal_initializer(0.0, hidden_dim**-0.5)))
-  if num_shards == 1:
-    ret = shards[0]
-  else:
-    ret = tf.concat(shards, 0)
-  # Convert ret to tensor.
-  if not tf.executing_eagerly():
-    ret = common_layers.convert_gradient_to_tensor(ret)
-  return ret
-
-
-def _symbol_bottom_simple(x, model_hparams, vocab_size, name, reuse):
-  """Bottom transformation for symbols."""
-  with tf.variable_scope(name, reuse=reuse):
-    # Ensure the inputs are 3-D
-    if len(x.get_shape()) == 4:
-      x = tf.squeeze(x, axis=3)
-    while len(x.get_shape()) < 3:
-      x = tf.expand_dims(x, axis=-1)
-
-    var = get_weights(model_hparams, vocab_size)
-    x = common_layers.dropout_no_scaling(
-        x, 1.0 - model_hparams.symbol_dropout)
-    ret = common_layers.gather(var, x)
-    if model_hparams.multiply_embedding_mode == "sqrt_depth":
-      ret *= model_hparams.hidden_size**0.5
-    ret *= tf.expand_dims(
-        common_layers.cast_like(tf.not_equal(x, 0), ret), -1)
-    return ret
-
-
-def symbol_bottom(x, model_hparams, vocab_size):
-  if (model_hparams.shared_embedding_and_softmax_weights or
-      model_hparams.get("shared_embedding")):
-    return _symbol_bottom_simple(
-        x, model_hparams, vocab_size, "shared", reuse=None)
-  return _symbol_bottom_simple(
-      x, model_hparams, vocab_size, "input_emb", reuse=None)
-
+  del vocab_size  # unused arg
+  inputs = x
+  with tf.variable_scope("audio_modality"):
+    # TODO(aidangomez): Will need to sort out a better audio pipeline
+    def xnet_resblock(x, filters, res_relu, name):
+      """Xception block."""
+      with tf.variable_scope(name):
+        # Typically audio samples are >100k samples in length and have a width
+        # of 2 or 4. Mono audio has a single channel while stereo has 2.
+        y = common_layers.separable_conv_block(
+            x,
+            filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
+            first_relu=True,
+            padding="SAME",
+            force2d=True,
+            name="sep_conv_block")
+        y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 2))
+        return y + common_layers.conv_block(
+            x,
+            filters, [((1, 1), (1, 1))],
+            padding="SAME",
+            strides=(2, 2),
+            first_relu=res_relu,
+            force2d=True,
+            name="res_conv0")
 
-def symbol_targets_bottom(x, model_hparams, vocab_size):
-  """Bottom transformation for target symbols."""
-  if (model_hparams.shared_embedding_and_softmax_weights or
-      model_hparams.get("shared_embedding")):
-    try:
-      return _symbol_bottom_simple(
-          x, model_hparams, vocab_size, "shared", reuse=True)
-    except ValueError:
-      # perhaps there were no inputs, and this is a new variable.
-      return _symbol_bottom_simple(
-          x, model_hparams, vocab_size, "shared", reuse=None)
-  else:
-    return _symbol_bottom_simple(
-        x, model_hparams, vocab_size, "target_emb", reuse=None)
+    x = tf.to_float(inputs) / 255.
+    x.set_shape([None, None, None, 1])
+    for i in range(model_hparams.audio_compression):
+      x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
+    return xnet_resblock(x,
+                         model_hparams.hidden_size,
+                         False,
+                         "compress_block_final")
 
 
-@is_pointwise
-def symbol_top(body_output, targets, model_hparams, vocab_size):
-  """Generate logits.
+def audio_spectral_bottom(x, model_hparams, vocab_size):
+  """Transform input from data space to model space.
 
   Args:
-    body_output: A Tensor with shape
-      [batch, p0, p1, model_hparams.hidden_size].
-    targets: Unused.
+    x: A Tensor with shape [batch, ...]
     model_hparams: tf.HParams, model hyperparmeters.
     vocab_size: int, vocabulary size.
 
   Returns:
-    logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].
+    body_input: A Tensor with shape [batch, ?, ?,
+      model_hparams.hidden_size].
   """
-  del targets  # unused arg
-  if model_hparams.shared_embedding_and_softmax_weights:
-    scope_name = "shared"
-    reuse = tf.AUTO_REUSE
-  else:
-    scope_name = "softmax"
-    reuse = False
-  with tf.variable_scope(scope_name, reuse=reuse):
-    body_output_shape = common_layers.shape_list(body_output)
-    var = get_weights(model_hparams, vocab_size, body_output_shape[-1])
-    if (model_hparams.factored_logits and
-        model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
-      # insert channels dimension
-      body_output = tf.expand_dims(body_output, 3)
-      return common_layers.FactoredTensor(body_output, var)
-    else:
-      body_output = tf.reshape(body_output, [-1, body_output_shape[-1]])
-      logits = tf.matmul(body_output, var, transpose_b=True)
-      return tf.reshape(logits,
-                        body_output_shape[:-1] + [1, vocab_size])
-
-
-def symbol_one_hot_bottom(x, model_hparams, vocab_size):
-  del model_hparams  # unused arg
-  return tf.one_hot(x, vocab_size)
+  del vocab_size  # unused arg
+  inputs = x
+  with tf.variable_scope("audio_spectral_modality"):
+    # TODO(aidangomez): Will need to sort out a better audio pipeline
+    def xnet_resblock(x, filters, res_relu, name):
+      """Xception-like block."""
+      with tf.variable_scope(name):
+        # We only stride along the length dimension to preserve the spectral
+        # bins (which are tiny in dimensionality relative to length)
+        y = common_layers.separable_conv_block(
+            x,
+            filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
+            first_relu=True,
+            padding="SAME",
+            force2d=True,
+            name="sep_conv_block")
+        y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 1))
+        return y + common_layers.conv_block(
+            x,
+            filters, [((1, 1), (1, 1))],
+            padding="SAME",
+            strides=(2, 1),
+            first_relu=res_relu,
+            force2d=True,
+            name="res_conv0")
 
+    # Bitcast back from int32
+    x = tf.bitcast(inputs, tf.float32)
+    x.set_shape([None, None, None, 1])
+    for i in range(model_hparams.audio_compression):
+      x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
+    return xnet_resblock(x,
+                         model_hparams.hidden_size,
+                         False,
+                         "compress_block_final")
 
-@is_pointwise
-def symbol_one_hot_top(body_output, targets, model_hparams, vocab_size):
-  del targets, model_hparams, vocab_size  # unused arg
-  return body_output
 
+def class_label_bottom(x, model_hparams, vocab_size):
+  with tf.variable_scope("class_label_modality_%d_%d" % (
+      vocab_size, model_hparams.hidden_size)):
+    multiplier = 1.0
+    if model_hparams.multiply_embedding_mode == "sqrt_depth":
+      multiplier = model_hparams.hidden_size**0.5
+    return common_layers.embedding(x,
+                                   vocab_size,
+                                   model_hparams.hidden_size,
+                                   multiplier=multiplier)
 
-def symbol_one_hot_loss(top_out,
-                        targets,
-                        model_hparams,
-                        vocab_size,
-                        weights_fn):
-  del model_hparams, weights_fn  # unused arg
-  labels = tf.one_hot(targets, vocab_size)
-  loss = tf.nn.softmax_cross_entropy_with_logits(
-      logits=top_out, labels=labels)
-  return tf.reduce_mean(loss), tf.constant(1.0)
 
+def class_label_targets_bottom(x, model_hparams, vocab_size):
+  with tf.variable_scope("class_label_modality_%d_%d" % (
+      vocab_size, model_hparams.hidden_size)):
+    return tf.zeros([common_layers.shape_list(x)[0],
+                     1,
+                     1,
+                     model_hparams.hidden_size])
 
-def ctc_symbol_loss(top_out, targets, model_hparams, vocab_size, weight_fn):
-  """Compute the CTC loss."""
-  del model_hparams, vocab_size  # unused arg
-  logits = top_out
-  with tf.name_scope("ctc_loss", values=[logits, targets]):
-    # For CTC we assume targets are 1d, [batch, length, 1, 1] here.
-    targets_shape = targets.get_shape().as_list()
-    assert len(targets_shape) == 4
-    assert targets_shape[2] == 1
-    assert targets_shape[3] == 1
-    targets = tf.squeeze(targets, axis=[2, 3])
-    logits = tf.squeeze(logits, axis=[2, 3])
-    targets_mask = 1 - tf.to_int32(tf.equal(targets, 0))
-    targets_lengths = tf.reduce_sum(targets_mask, axis=1)
-    sparse_targets = tf.keras.backend.ctc_label_dense_to_sparse(
-        targets, targets_lengths)
-    xent = tf.nn.ctc_loss(
-        sparse_targets,
-        logits,
-        targets_lengths,
-        time_major=False,
-        preprocess_collapse_repeated=False,
-        ctc_merge_repeated=False)
-    weights = weight_fn(targets)
-    return tf.reduce_sum(xent), tf.reduce_sum(weights)
+
+def identity_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  return tf.to_float(x)
 
 
 def image_bottom(x, model_hparams, vocab_size):
@@ -281,26 +288,6 @@ def image_targets_bottom(x, model_hparams, vocab_size):
     return merged
 
 
-def image_top(body_output, targets, model_hparams, vocab_size):
-  """Top transformation for images."""
-  del targets  # unused arg
-  # TODO(lukaszkaiser): is this a universal enough way to get channels?
-  num_channels = model_hparams.problem.num_channels
-  with tf.variable_scope("rgb_softmax"):
-    body_output_shape = common_layers.shape_list(body_output)
-    reshape_shape = body_output_shape[:3]
-    reshape_shape.extend([num_channels, vocab_size])
-    res = tf.layers.dense(body_output, vocab_size * num_channels)
-    res = tf.reshape(res, reshape_shape)
-    if not tf.get_variable_scope().reuse:
-      res_argmax = tf.argmax(res, axis=-1)
-      tf.summary.image(
-          "result",
-          common_layers.tpu_safe_image_summary(res_argmax),
-          max_outputs=1)
-    return res
-
-
 def _image_channel_compress_bottom(inputs, model_hparams, name="bottom"):
   """Compresses channel-wise input pixels into whole pixel representions.
 
@@ -356,44 +343,6 @@ def image_channel_compress_targets_bottom(x, model_hparams, vocab_size):
   return _image_channel_compress_bottom(x, model_hparams, "output_bottom")
 
 
-def image_channel_compress_top(body_output, targets, model_hparams, vocab_size):
-  """Transforms body output to return logits.
-
-  Args:
-    body_output: Tensor of shape [batch, img_len, img_len, depth].
-    targets:
-    model_hparams: tf.HParams, model hyperparmeters.
-    vocab_size: int, vocabulary size.
-
-  Returns:
-    Tensor of shape [batch, img_len, img_len, channels, vocab_size].
-  """
-  del targets  # unused arg
-  with tf.variable_scope("image_channel_compress_modality"):
-    hidden_size = model_hparams.hidden_size
-    img_len = model_hparams.img_len
-    channels = 3  # RGB
-    batch = common_layers.shape_list(body_output)[0]
-    x = tf.layers.conv2d(
-        body_output,
-        hidden_size * channels,
-        kernel_size=(1, 1),
-        strides=(1, 1),
-        padding="VALID",
-        activation=tf.nn.relu,
-        name="decompress_conv")
-    x = tf.reshape(x, [batch, img_len, img_len * channels, hidden_size])
-    x = common_layers.layer_preprocess(x, model_hparams)
-    x = tf.layers.dense(x,
-                        vocab_size,
-                        use_bias=True,
-                        activation=None,
-                        name="output_conv")
-    x = tf.reshape(
-        x, [batch, img_len, img_len, channels, vocab_size])
-    return x
-
-
 def image_channel_embeddings_bottom(x, model_hparams, vocab_size):
   """Bottom transformation for image targets."""
   del vocab_size  # unused arg
@@ -407,117 +356,18 @@ def image_channel_embeddings_bottom(x, model_hparams, vocab_size):
                     [tshape[0], tshape[1], tshape[2] * io_depth, hidden_size])
 
 
-def image_channel_embeddings_top(body_output,
-                                 targets,
-                                 model_hparams,
-                                 vocab_size):
-  """Top transformation for images."""
-  del targets  # unused arg
-  with tf.variable_scope("image_channel_embeddings_bottom"):
-    img_len = model_hparams.img_len
-    channels = model_hparams.num_channels
-    x = tf.layers.dense(
-        body_output, 256, use_bias=True, activation=None, name="output_conv")
-    x = tf.reshape(x,
-                   [-1, img_len, img_len, channels, vocab_size])
-    return x
-
-
-def audio_bottom(x, model_hparams, vocab_size):
-  """Transform input from data space to model space.
-
-  Args:
-    x: A Tensor with shape [batch, ...]
-    model_hparams: tf.HParams, model hyperparmeters.
-    vocab_size: int, vocabulary size.
-
-  Returns:
-    body_input: A Tensor with shape [batch, ?, ?,
-      model_hparams.hidden_size].
-  """
-  del vocab_size  # unused arg
-  inputs = x
-  with tf.variable_scope("audio_modality"):
-    # TODO(aidangomez): Will need to sort out a better audio pipeline
-    def xnet_resblock(x, filters, res_relu, name):
-      """Xception block."""
-      with tf.variable_scope(name):
-        # Typically audio samples are >100k samples in length and have a width
-        # of 2 or 4. Mono audio has a single channel while stereo has 2.
-        y = common_layers.separable_conv_block(
-            x,
-            filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
-            first_relu=True,
-            padding="SAME",
-            force2d=True,
-            name="sep_conv_block")
-        y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 2))
-        return y + common_layers.conv_block(
-            x,
-            filters, [((1, 1), (1, 1))],
-            padding="SAME",
-            strides=(2, 2),
-            first_relu=res_relu,
-            force2d=True,
-            name="res_conv0")
-
-    x = tf.to_float(inputs) / 255.
-    x.set_shape([None, None, None, 1])
-    for i in range(model_hparams.audio_compression):
-      x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
-    return xnet_resblock(x,
-                         model_hparams.hidden_size,
-                         False,
-                         "compress_block_final")
-
-
-def audio_spectral_bottom(x, model_hparams, vocab_size):
-  """Transform input from data space to model space.
+def make_targets_bottom(bottom):
+  def targets_bottom(x, model_hparams, vocab_size):
+    with tf.variable_scope("targets_bottom"):
+      return bottom(x, model_hparams, vocab_size)
+  return targets_bottom
 
-  Args:
-    x: A Tensor with shape [batch, ...]
-    model_hparams: tf.HParams, model hyperparmeters.
-    vocab_size: int, vocabulary size.
 
-  Returns:
-    body_input: A Tensor with shape [batch, ?, ?,
-      model_hparams.hidden_size].
-  """
+def real_bottom(x, model_hparams, vocab_size):
   del vocab_size  # unused arg
-  inputs = x
-  with tf.variable_scope("audio_spectral_modality"):
-    # TODO(aidangomez): Will need to sort out a better audio pipeline
-    def xnet_resblock(x, filters, res_relu, name):
-      """Xception-like block."""
-      with tf.variable_scope(name):
-        # We only stride along the length dimension to preserve the spectral
-        # bins (which are tiny in dimensionality relative to length)
-        y = common_layers.separable_conv_block(
-            x,
-            filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
-            first_relu=True,
-            padding="SAME",
-            force2d=True,
-            name="sep_conv_block")
-        y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 1))
-        return y + common_layers.conv_block(
-            x,
-            filters, [((1, 1), (1, 1))],
-            padding="SAME",
-            strides=(2, 1),
-            first_relu=res_relu,
-            force2d=True,
-            name="res_conv0")
-
-    # Bitcast back from int32
-    x = tf.bitcast(inputs, tf.float32)
-    x.set_shape([None, None, None, 1])
-    for i in range(model_hparams.audio_compression):
-      x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
-    return xnet_resblock(x,
-                         model_hparams.hidden_size,
-                         False,
-                         "compress_block_final")
+  with tf.variable_scope("real"):
+    return tf.layers.dense(
+        tf.to_float(x), model_hparams.hidden_size, name="bottom")
 
 
 def speech_recognition_bottom(x, model_hparams, vocab_size):
@@ -602,6 +452,89 @@ def speech_recognition_bottom(x, model_hparams, vocab_size):
   return x
 
 
+def get_weights(model_hparams, vocab_size, hidden_dim=None):
+  """Create or get concatenated embedding or softmax variable.
+
+  Args:
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
+    hidden_dim: dim of the variable. Defaults to _model_hparams' hidden_size
+
+  Returns:
+     a list of num_shards Tensors.
+  """
+  if hidden_dim is None:
+    hidden_dim = model_hparams.hidden_size
+  num_shards = model_hparams.symbol_modality_num_shards
+  shards = []
+  for i in range(num_shards):
+    shard_size = (vocab_size // num_shards) + (
+        1 if i < vocab_size % num_shards else 0)
+    var_name = "weights_%d" % i
+    shards.append(
+        tf.get_variable(
+            var_name, [shard_size, hidden_dim],
+            initializer=tf.random_normal_initializer(0.0, hidden_dim**-0.5)))
+  if num_shards == 1:
+    ret = shards[0]
+  else:
+    ret = tf.concat(shards, 0)
+  # Convert ret to tensor.
+  if not tf.executing_eagerly():
+    ret = common_layers.convert_gradient_to_tensor(ret)
+  return ret
+
+
+def _symbol_bottom_simple(x, model_hparams, vocab_size, name, reuse):
+  """Bottom transformation for symbols."""
+  with tf.variable_scope(name, reuse=reuse):
+    # Ensure the inputs are 3-D
+    if len(x.get_shape()) == 4:
+      x = tf.squeeze(x, axis=3)
+    while len(x.get_shape()) < 3:
+      x = tf.expand_dims(x, axis=-1)
+
+    var = get_weights(model_hparams, vocab_size)
+    x = common_layers.dropout_no_scaling(
+        x, 1.0 - model_hparams.symbol_dropout)
+    ret = common_layers.gather(var, x)
+    if model_hparams.multiply_embedding_mode == "sqrt_depth":
+      ret *= model_hparams.hidden_size**0.5
+    ret *= tf.expand_dims(
+        common_layers.cast_like(tf.not_equal(x, 0), ret), -1)
+    return ret
+
+
+def symbol_bottom(x, model_hparams, vocab_size):
+  if (model_hparams.shared_embedding_and_softmax_weights or
+      model_hparams.get("shared_embedding")):
+    return _symbol_bottom_simple(
+        x, model_hparams, vocab_size, "shared", reuse=None)
+  return _symbol_bottom_simple(
+      x, model_hparams, vocab_size, "input_emb", reuse=None)
+
+
+def symbol_targets_bottom(x, model_hparams, vocab_size):
+  """Bottom transformation for target symbols."""
+  if (model_hparams.shared_embedding_and_softmax_weights or
+      model_hparams.get("shared_embedding")):
+    try:
+      return _symbol_bottom_simple(
+          x, model_hparams, vocab_size, "shared", reuse=True)
+    except ValueError:
+      # perhaps there were no inputs, and this is a new variable.
+      return _symbol_bottom_simple(
+          x, model_hparams, vocab_size, "shared", reuse=None)
+  else:
+    return _symbol_bottom_simple(
+        x, model_hparams, vocab_size, "target_emb", reuse=None)
+
+
+def symbol_one_hot_bottom(x, model_hparams, vocab_size):
+  del model_hparams  # unused arg
+  return tf.one_hot(x, vocab_size)
+
+
 def video_bottom(x, model_hparams, vocab_size):
   del model_hparams, vocab_size  # unused arg
   common_video.gif_summary("inputs", x, max_outputs=1)
@@ -616,35 +549,6 @@ def video_targets_bottom(x, model_hparams, vocab_size):
   return x
 
 
-def video_top(body_output, targets, model_hparams, vocab_size):
-  """Top transformation for video."""
-  del targets  # unused arg
-  num_channels = model_hparams.problem.num_channels
-  shape = common_layers.shape_list(body_output)
-  reshape_shape = shape[:-1] + [num_channels, vocab_size]
-  res = tf.reshape(body_output, reshape_shape)
-  # Calculate argmax so as to have a summary with the produced images.
-  x = tf.argmax(tf.reshape(res, [-1, vocab_size]), axis=-1)
-  x = tf.reshape(x, shape[:-1] + [num_channels])
-  common_video.gif_summary("results", x, max_outputs=1)
-  return res
-
-
-def video_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-  """Compute loss numerator and denominator for one shard of output."""
-  del vocab_size  # unused arg
-  logits = top_out
-  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
-  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
-  return common_layers.padded_cross_entropy(
-      logits,
-      targets,
-      model_hparams.label_smoothing,
-      cutoff=cutoff,
-      weights_fn=weights_fn)
-
-
 def video_bitwise_bottom(x, model_hparams, vocab_size):
   """Bottom transformation for embedding video bitwise."""
   pixel_embedding_size = 64
@@ -680,6 +584,18 @@ def video_bitwise_targets_bottom(x, model_hparams, vocab_size):
         name="merge_pixel_embedded_frames")
 
 
+def video_identity_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("inputs", x, max_outputs=1)
+  return x
+
+
+def video_identity_targets_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("targets", x, max_outputs=1)
+  return x
+
+
 def video_pixel_noise_bottom(x, model_hparams, vocab_size):
   """Bottom transformation for video."""
   input_noise = getattr(model_hparams, "video_modality_input_noise", 0.25)
@@ -696,23 +612,236 @@ def video_pixel_noise_bottom(x, model_hparams, vocab_size):
   return video_bottom(inputs, model_hparams, vocab_size)
 
 
-def video_l1_top(body_output, targets, model_hparams, vocab_size):
-  """Top transformation for video."""
-  del targets, vocab_size  # unused arg
-  num_channels = model_hparams.problem.num_channels
-  num_frames = model_hparams.video_num_target_frames
-  with tf.variable_scope("rgb"):
-    body_output_shape = common_layers.shape_list(body_output)
-    res = tf.layers.dense(body_output, num_channels * num_frames, name="cast")
-    res = tf.reshape(res, body_output_shape[:3] + [num_channels, num_frames])
-    res = tf.transpose(res, [0, 4, 1, 2, 3])  # Move frames next to batch.
-    if not tf.get_variable_scope().reuse:
-      res_argmax = res[:, -1, :, :, :]
-      tf.summary.image(
-          "result",
-          common_layers.tpu_safe_image_summary(res_argmax),
-          max_outputs=1)
-    return tf.expand_dims(res, axis=-1)  # Add an axis like in perplexity.
+def convert_rgb_to_real(prediction, targets):
+  """Convert prediction and target from rgb to real."""
+  prediction = tf.squeeze(prediction, axis=-1)
+  prediction = common_layers.convert_rgb_to_real(prediction)
+  targets = common_layers.convert_rgb_to_real(targets)
+  return prediction, targets
+
+
+def video_raw_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("inputs", x)
+  return common_layers.convert_rgb_to_real(x)
+
+
+def video_raw_targets_bottom(x, model_hparams, vocab_size):
+  del model_hparams, vocab_size  # unused arg
+  common_video.gif_summary("targets_bottom", x)
+  return common_layers.convert_rgb_to_real(x)
+
+
+# Loss transformations, applied to target features
+
+
+def ctc_symbol_loss(top_out, targets, model_hparams, vocab_size, weight_fn):
+  """Compute the CTC loss."""
+  del model_hparams, vocab_size  # unused arg
+  logits = top_out
+  with tf.name_scope("ctc_loss", values=[logits, targets]):
+    # For CTC we assume targets are 1d, [batch, length, 1, 1] here.
+    targets_shape = targets.get_shape().as_list()
+    assert len(targets_shape) == 4
+    assert targets_shape[2] == 1
+    assert targets_shape[3] == 1
+    targets = tf.squeeze(targets, axis=[2, 3])
+    logits = tf.squeeze(logits, axis=[2, 3])
+    targets_mask = 1 - tf.to_int32(tf.equal(targets, 0))
+    targets_lengths = tf.reduce_sum(targets_mask, axis=1)
+    sparse_targets = tf.keras.backend.ctc_label_dense_to_sparse(
+        targets, targets_lengths)
+    xent = tf.nn.ctc_loss(
+        sparse_targets,
+        logits,
+        targets_lengths,
+        time_major=False,
+        preprocess_collapse_repeated=False,
+        ctc_merge_repeated=False)
+    weights = weight_fn(targets)
+    return tf.reduce_sum(xent), tf.reduce_sum(weights)
+
+
+def generic_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  """Compute loss numerator and denominator for one shard of output."""
+  del vocab_size  # unused arg
+  logits = top_out
+  logits = common_attention.maybe_upcast(logits, hparams=model_hparams)
+  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.0)
+  return common_layers.padded_cross_entropy(
+      logits,
+      targets,
+      model_hparams.label_smoothing,
+      cutoff=cutoff,
+      weights_fn=weights_fn)
+
+
+def generic_l2_loss(body_output,
+                    targets,
+                    model_hparams,
+                    vocab_size,
+                    weights_fn):
+  del model_hparams, vocab_size, weights_fn  # unused arg
+  loss = tf.squared_difference(body_output, tf.to_float(targets))
+  return tf.reduce_mean(loss), tf.constant(1.0)
+
+
+def multi_label_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  """Average loss over the labels."""
+  del vocab_size  # unused arg
+  logits = top_out
+  num_labels = tf.shape(targets)[1]
+  logits = tf.tile(logits, [1, num_labels, 1, 1, 1])
+
+  xent, weights = common_layers.padded_cross_entropy(
+      logits,
+      targets,
+      model_hparams.label_smoothing,
+      weights_fn=weights_fn,
+      reduce_sum=False,
+  )
+  xent = tf.squeeze(xent, [2, 3])
+  weights = tf.squeeze(weights, [2, 3])
+  # average loss over all labels
+  loss = tf.reduce_sum(xent, axis=1)
+  weights = tf.reduce_sum(weights, axis=1)
+  loss /= (weights + 1e-8)
+  weights = tf.to_float(tf.greater(weights, 0.))
+
+  return tf.reduce_sum(loss*weights), tf.reduce_sum(weights)
+
+
+def one_hot_class_label_loss(top_out,
+                             targets,
+                             model_hparams,
+                             vocab_size,
+                             weights_fn):
+  """Apply softmax cross-entropy between outputs and targets.
+
+  Args:
+    top_out: logits Tensor with shape [batch, ?, ?, num_classes]
+    targets: one-hot encoding Tensor with shape [batch, ?, ?, num_classes]
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
+    weights_fn:
+
+  Returns:
+    loss_scale (cross-entropy), loss_denom
+  """
+  del model_hparams, vocab_size  # unused arg
+  loss_scale = tf.losses.softmax_cross_entropy(
+      onehot_labels=targets, logits=top_out)
+  weights = weights_fn(targets)
+  loss_denom = tf.reduce_sum(weights)
+  return loss_scale, loss_denom
+
+
+def real_l2_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  del model_hparams, vocab_size  # unused arg
+  predictions = top_out
+  if (len(common_layers.shape_list(top_out)) != len(
+      common_layers.shape_list(targets))):
+    predictions = tf.squeeze(top_out, axis=[-1])
+  with tf.name_scope("l2"):
+    weights = weights_fn(targets)
+    l2 = tf.pow(predictions - targets, 2)
+    return tf.reduce_sum(l2 * weights), tf.reduce_sum(weights)
+
+
+def real_log_poisson_loss(top_out,
+                          targets,
+                          model_hparams,
+                          vocab_size,
+                          weights_fn):
+  """Poisson loss for real."""
+  del model_hparams, vocab_size  # unused arg
+  predictions = top_out
+  if (len(common_layers.shape_list(top_out)) != len(
+      common_layers.shape_list(targets))):
+    predictions = tf.squeeze(top_out, axis=[-1])
+  with tf.name_scope("log_possion"):
+    weights = weights_fn(targets)
+    lp_loss = tf.nn.log_poisson_loss(targets, predictions)
+    return tf.reduce_sum(lp_loss * weights), tf.reduce_sum(weights)
+
+
+def sigmoid_class_label_loss(top_out,
+                             targets,
+                             model_hparams,
+                             vocab_size,
+                             weights_fn):
+  """Loss for class label."""
+  # Expect inputs of size [batch-size, timesteps, 1, num-classes], where the
+  # last dimension of num-classes represents logits for binary labels
+  del model_hparams, vocab_size  # unused arg
+  loss_scale = tf.losses.sigmoid_cross_entropy(
+      multi_class_labels=targets, logits=top_out)
+  weights = weights_fn(targets)
+  loss_denom = tf.reduce_sum(weights)
+  return loss_scale, loss_denom
+
+
+def sigmoid_max_pooling_class_label_loss(top_out,
+                                         targets,
+                                         model_hparams,
+                                         vocab_size,
+                                         weights_fn):
+  """Loss for class label."""
+  # Expect inputs of size [batch-size, 1, 1, num-classes], where the
+  # last dimension of num-classes represents logits for binary labels
+  del model_hparams, vocab_size  # unused arg
+  loss_scale = tf.losses.sigmoid_cross_entropy(
+      multi_class_labels=targets, logits=top_out)
+  weights = weights_fn(targets)
+  loss_denom = tf.reduce_sum(weights)
+  return loss_scale, loss_denom
+
+
+def symbol_one_hot_loss(top_out,
+                        targets,
+                        model_hparams,
+                        vocab_size,
+                        weights_fn):
+  del model_hparams, weights_fn  # unused arg
+  labels = tf.one_hot(targets, vocab_size)
+  loss = tf.nn.softmax_cross_entropy_with_logits(
+      logits=top_out, labels=labels)
+  return tf.reduce_mean(loss), tf.constant(1.0)
+
+
+def video_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  """Compute loss numerator and denominator for one shard of output."""
+  del vocab_size  # unused arg
+  logits = top_out
+  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
+  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
+  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
+  return common_layers.padded_cross_entropy(
+      logits,
+      targets,
+      model_hparams.label_smoothing,
+      cutoff=cutoff,
+      weights_fn=weights_fn)
+
+
+def video_identity_loss(top_out,
+                        targets,
+                        model_hparams,
+                        vocab_size,
+                        weights_fn):
+  """Compute loss numerator and denominator for one shard of output."""
+  del vocab_size  # unused arg
+  # TODO(nikip): Try L2 loss
+  logits = top_out
+  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
+  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
+  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
+  return common_layers.padded_cross_entropy(
+      logits,
+      targets,
+      model_hparams.label_smoothing,
+      cutoff=cutoff,
+      weights_fn=weights_fn)
 
 
 def video_l1_internal_loss(logits, targets, model_hparams):
@@ -758,36 +887,6 @@ def video_l2_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
   return tf.reduce_sum(loss * weights), tf.reduce_sum(weights)
 
 
-def convert_rgb_to_real(prediction, targets):
-  """Convert prediction and target from rgb to real."""
-  prediction = tf.squeeze(prediction, axis=-1)
-  prediction = common_layers.convert_rgb_to_real(prediction)
-  targets = common_layers.convert_rgb_to_real(targets)
-  return prediction, targets
-
-
-def video_raw_bottom(x, model_hparams, vocab_size):
-  del model_hparams, vocab_size  # unused arg
-  common_video.gif_summary("inputs", x)
-  return common_layers.convert_rgb_to_real(x)
-
-
-def video_raw_targets_bottom(x, model_hparams, vocab_size):
-  del model_hparams, vocab_size  # unused arg
-  common_video.gif_summary("targets_bottom", x)
-  return common_layers.convert_rgb_to_real(x)
-
-
-def video_raw_top(body_output, targets, model_hparams, vocab_size):
-  del targets, model_hparams, vocab_size  # unused arg
-  frames = body_output
-  if isinstance(body_output, list):
-    frames = tf.stack(body_output, axis=1)
-  rgb_frames = common_layers.convert_real_to_rgb(frames)
-  common_video.gif_summary("body_output", rgb_frames)
-  return tf.expand_dims(rgb_frames, axis=-1)
-
-
 def video_l2_raw_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
   del model_hparams, vocab_size, weights_fn  # unused arg
   prediction, groundtruth = convert_rgb_to_real(top_out, targets)
@@ -802,25 +901,27 @@ def video_l1_raw_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
   return loss, tf.constant(1.0)
 
 
-def class_label_bottom(x, model_hparams, vocab_size):
-  with tf.variable_scope("class_label_modality_%d_%d" % (
-      vocab_size, model_hparams.hidden_size)):
-    multiplier = 1.0
-    if model_hparams.multiply_embedding_mode == "sqrt_depth":
-      multiplier = model_hparams.hidden_size**0.5
-    return common_layers.embedding(x,
-                                   vocab_size,
-                                   model_hparams.hidden_size,
-                                   multiplier=multiplier)
+# Top transformations, applied to target features
 
 
-def class_label_targets_bottom(x, model_hparams, vocab_size):
-  with tf.variable_scope("class_label_modality_%d_%d" % (
-      vocab_size, model_hparams.hidden_size)):
-    return tf.zeros([common_layers.shape_list(x)[0],
-                     1,
-                     1,
-                     model_hparams.hidden_size])
+def is_pointwise(func):
+  """Decorator for whether the function is pointwise.
+
+  An example of a pointwise function is a linear layer followed by
+  a softmax. Given a tensor [batch, length, height, depth] it operates
+  only on the last axis, on every point in [batch, length, height] fully
+  independently. In contrast, a classifier that first averages over length
+  and height is not pointwise, as it depends on the whole field. It is useful
+  to know if top functions are pointwise to speed up decoding in certain models.
+
+  Args:
+    func: Function to decorate.
+
+  Returns:
+    Original function with an attribute pointwise set to True.
+  """
+  func.pointwise = True
+  return func
 
 
 def class_label_top(body_output, targets, model_hparams, vocab_size):
@@ -846,113 +947,83 @@ def class_label_top(body_output, targets, model_hparams, vocab_size):
     return tf.expand_dims(res, 3)
 
 
-def video_identity_bottom(x, model_hparams, vocab_size):
-  del model_hparams, vocab_size  # unused arg
-  common_video.gif_summary("inputs", x, max_outputs=1)
-  return x
-
-
-def video_identity_targets_bottom(x, model_hparams, vocab_size):
-  del model_hparams, vocab_size  # unused arg
-  common_video.gif_summary("targets", x, max_outputs=1)
-  return x
-
-
-def video_identity_loss(top_out,
-                        targets,
-                        model_hparams,
-                        vocab_size,
-                        weights_fn):
-  """Compute loss numerator and denominator for one shard of output."""
-  del vocab_size  # unused arg
-  # TODO(nikip): Try L2 loss
-  logits = top_out
-  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
-  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
-  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
-  return common_layers.padded_cross_entropy(
-      logits,
-      targets,
-      model_hparams.label_smoothing,
-      cutoff=cutoff,
-      weights_fn=weights_fn)
-
-
-def multi_label_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-  """Average loss over the labels."""
-  del vocab_size  # unused arg
-  logits = top_out
-  num_labels = tf.shape(targets)[1]
-  logits = tf.tile(logits, [1, num_labels, 1, 1, 1])
+def identity_top(body_output, targets, model_hparams, vocab_size):
+  del targets, model_hparams, vocab_size  # unused arg
+  return body_output
 
-  xent, weights = common_layers.padded_cross_entropy(
-      logits,
-      targets,
-      model_hparams.label_smoothing,
-      weights_fn=weights_fn,
-      reduce_sum=False,
-  )
-  xent = tf.squeeze(xent, [2, 3])
-  weights = tf.squeeze(weights, [2, 3])
-  # average loss over all labels
-  loss = tf.reduce_sum(xent, axis=1)
-  weights = tf.reduce_sum(weights, axis=1)
-  loss /= (weights + 1e-8)
-  weights = tf.to_float(tf.greater(weights, 0.))
 
-  return tf.reduce_sum(loss*weights), tf.reduce_sum(weights)
+def image_top(body_output, targets, model_hparams, vocab_size):
+  """Top transformation for images."""
+  del targets  # unused arg
+  # TODO(lukaszkaiser): is this a universal enough way to get channels?
+  num_channels = model_hparams.problem.num_channels
+  with tf.variable_scope("rgb_softmax"):
+    body_output_shape = common_layers.shape_list(body_output)
+    reshape_shape = body_output_shape[:3]
+    reshape_shape.extend([num_channels, vocab_size])
+    res = tf.layers.dense(body_output, vocab_size * num_channels)
+    res = tf.reshape(res, reshape_shape)
+    if not tf.get_variable_scope().reuse:
+      res_argmax = tf.argmax(res, axis=-1)
+      tf.summary.image(
+          "result",
+          common_layers.tpu_safe_image_summary(res_argmax),
+          max_outputs=1)
+    return res
 
 
-def one_hot_class_label_loss(top_out,
-                             targets,
-                             model_hparams,
-                             vocab_size,
-                             weights_fn):
-  """Apply softmax cross-entropy between outputs and targets.
+def image_channel_compress_top(body_output, targets, model_hparams, vocab_size):
+  """Transforms body output to return logits.
 
   Args:
-    top_out: logits Tensor with shape [batch, ?, ?, num_classes]
-    targets: one-hot encoding Tensor with shape [batch, ?, ?, num_classes]
+    body_output: Tensor of shape [batch, img_len, img_len, depth].
+    targets:
     model_hparams: tf.HParams, model hyperparmeters.
     vocab_size: int, vocabulary size.
-    weights_fn:
 
   Returns:
-    loss_scale (cross-entropy), loss_denom
+    Tensor of shape [batch, img_len, img_len, channels, vocab_size].
   """
-  del model_hparams, vocab_size  # unused arg
-  loss_scale = tf.losses.softmax_cross_entropy(
-      onehot_labels=targets, logits=top_out)
-  weights = weights_fn(targets)
-  loss_denom = tf.reduce_sum(weights)
-  return loss_scale, loss_denom
-
-
-def identity_bottom(x, model_hparams, vocab_size):
-  del model_hparams, vocab_size  # unused arg
-  return tf.to_float(x)
-
-
-def identity_top(body_output, targets, model_hparams, vocab_size):
-  del targets, model_hparams, vocab_size  # unused arg
-  return body_output
-
-
-def generic_l2_loss(body_output,
-                    targets,
-                    model_hparams,
-                    vocab_size,
-                    weights_fn):
-  del model_hparams, vocab_size, weights_fn  # unused arg
-  loss = tf.squared_difference(body_output, tf.to_float(targets))
-  return tf.reduce_mean(loss), tf.constant(1.0)
+  del targets  # unused arg
+  with tf.variable_scope("image_channel_compress_modality"):
+    hidden_size = model_hparams.hidden_size
+    img_len = model_hparams.img_len
+    channels = 3  # RGB
+    batch = common_layers.shape_list(body_output)[0]
+    x = tf.layers.conv2d(
+        body_output,
+        hidden_size * channels,
+        kernel_size=(1, 1),
+        strides=(1, 1),
+        padding="VALID",
+        activation=tf.nn.relu,
+        name="decompress_conv")
+    x = tf.reshape(x, [batch, img_len, img_len * channels, hidden_size])
+    x = common_layers.layer_preprocess(x, model_hparams)
+    x = tf.layers.dense(x,
+                        vocab_size,
+                        use_bias=True,
+                        activation=None,
+                        name="output_conv")
+    x = tf.reshape(
+        x, [batch, img_len, img_len, channels, vocab_size])
+    return x
 
 
-def real_bottom(x, model_hparams, vocab_size):
-  del vocab_size  # unused arg
-  with tf.variable_scope("real"):
-    return tf.layers.dense(
-        tf.to_float(x), model_hparams.hidden_size, name="bottom")
+def image_channel_embeddings_top(body_output,
+                                 targets,
+                                 model_hparams,
+                                 vocab_size):
+  """Top transformation for images."""
+  del targets  # unused arg
+  with tf.variable_scope("image_channel_embeddings_bottom"):
+    img_len = model_hparams.img_len
+    channels = model_hparams.num_channels
+    x = tf.layers.dense(
+        body_output, 256, use_bias=True, activation=None, name="output_conv")
+    x = tf.reshape(x,
+                   [-1, img_len, img_len, channels, vocab_size])
+    return x
 
 
 @is_pointwise
@@ -962,51 +1033,6 @@ def real_top(body_output, targets, model_hparams, vocab_size):
     return tf.layers.dense(body_output, vocab_size, name="top")
 
 
-def real_l2_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
-  del model_hparams, vocab_size  # unused arg
-  predictions = top_out
-  if (len(common_layers.shape_list(top_out)) != len(
-      common_layers.shape_list(targets))):
-    predictions = tf.squeeze(top_out, axis=[-1])
-  with tf.name_scope("l2"):
-    weights = weights_fn(targets)
-    l2 = tf.pow(predictions - targets, 2)
-    return tf.reduce_sum(l2 * weights), tf.reduce_sum(weights)
-
-
-def real_log_poisson_loss(top_out,
-                          targets,
-                          model_hparams,
-                          vocab_size,
-                          weights_fn):
-  """Poisson loss for real."""
-  del model_hparams, vocab_size  # unused arg
-  predictions = top_out
-  if (len(common_layers.shape_list(top_out)) != len(
-      common_layers.shape_list(targets))):
-    predictions = tf.squeeze(top_out, axis=[-1])
-  with tf.name_scope("log_possion"):
-    weights = weights_fn(targets)
-    lp_loss = tf.nn.log_poisson_loss(targets, predictions)
-    return tf.reduce_sum(lp_loss * weights), tf.reduce_sum(weights)
-
-
-def sigmoid_class_label_loss(top_out,
-                             targets,
-                             model_hparams,
-                             vocab_size,
-                             weights_fn):
-  """Loss for class label."""
-  # Expect inputs of size [batch-size, timesteps, 1, num-classes], where the
-  # last dimension of num-classes represents logits for binary labels
-  del model_hparams, vocab_size  # unused arg
-  loss_scale = tf.losses.sigmoid_cross_entropy(
-      multi_class_labels=targets, logits=top_out)
-  weights = weights_fn(targets)
-  loss_denom = tf.reduce_sum(weights)
-  return loss_scale, loss_denom
-
-
 def sigmoid_max_pooling_class_label_top(body_output,
                                         targets,
                                         model_hparams,
@@ -1033,36 +1059,6 @@ def sigmoid_max_pooling_class_label_top(body_output,
     return tf.layers.dense(x, vocab_size)
 
 
-def sigmoid_max_pooling_class_label_loss(top_out,
-                                         targets,
-                                         model_hparams,
-                                         vocab_size,
-                                         weights_fn):
-  """Loss for class label."""
-  # Expect inputs of size [batch-size, 1, 1, num-classes], where the
-  # last dimension of num-classes represents logits for binary labels
-  del model_hparams, vocab_size  # unused arg
-  loss_scale = tf.losses.sigmoid_cross_entropy(
-      multi_class_labels=targets, logits=top_out)
-  weights = weights_fn(targets)
-  loss_denom = tf.reduce_sum(weights)
-  return loss_scale, loss_denom
-
-
-def softmax_max_pooling_class_label_top(body_output,
-                                        targets,
-                                        model_hparams,
-                                        vocab_size):
-  """Loss for class label."""
-  del targets  # unused arg
-  with tf.variable_scope(
-      "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
-          vocab_size, model_hparams.hidden_size)):
-    x = body_output
-    x = tf.reduce_max(x, axis=1, keepdims=True)
-    return tf.layers.dense(x, vocab_size)
-
-
 def softmax_average_pooling_class_label_top(body_output,
                                             targets,
                                             model_hparams,
@@ -1091,93 +1087,106 @@ def softmax_last_timestep_class_label_top(body_output,
     return tf.layers.dense(x, vocab_size)
 
 
-class ModalityType(object):
-  """Types of modalities."""
+def softmax_max_pooling_class_label_top(body_output,
+                                        targets,
+                                        model_hparams,
+                                        vocab_size):
+  """Loss for class label."""
+  del targets  # unused arg
+  with tf.variable_scope(
+      "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
+          vocab_size, model_hparams.hidden_size)):
+    x = body_output
+    x = tf.reduce_max(x, axis=1, keepdims=True)
+    return tf.layers.dense(x, vocab_size)
 
-  AUDIO = "audio"
-  AUDIO_SPECTRAL = "audio_spectral"
-  CLASS_LABEL = "class_label"
-  CTC_SYMBOL = "ctc_symbol"  # symbol with CTC loss
-  GENERIC_L2_LOSS = "generic_l2"  # identity modality with L2 loss
-  IDENTITY = "identity"  # identity top and bottom
-  IDENTITY_SYMBOL = "identity_symbol"  # symbol with identity top and bottom
-  IMAGE = "image"
-  # images using channel compression for generation
-  IMAGE_CHANNEL_BOTTOM_IDENTITY = "image_channel_bottom_identity"
-  # images using channel compression for generation
-  IMAGE_CHANNEL_COMPRESS = "image_channel_compress"
-  IMAGE_CHANNEL_EMBEDDINGS_BOTTOM = "image_channel_embeddings_bottom"
-  MULTI_LABEL = "multi_label"
-  ONE_HOT_CLASS_LABEL = "one_hot_class_label"
-  REAL = "real"  # real vectors
-  REAL_L2_LOSS = "real_l2"  # real vectors with L2 as loss
-  # real vectors with log Poisson regression loss
-  REAL_LOG_POISSON_LOSS = "real_log_poisson"
-  SIGMOID_CLASS_LABEL = "sigmoid_class_label"  # sigmoid cross-entropy loss
-  # sigmoid cross-entropy applied on max-pooling over timesteps
-  SIGMOID_MAX_POOLING_CLASS_LABEL = "sigmoid_max_pooling_class_label"
-  # softmax cross-entropy applied on average-pooling over timesteps
-  SOFTMAX_AVERAGE_POOLING_CLASS_LABEL = "softmax_average_pooling_class_label"
-  # softmax cross-entropy applied on last-timestep encoding
-  SOFTMAX_LAST_TIMESTEP_CLASS_LABEL = "softmax_last_timestep_class_label"
-  # softmax cross-entropy applied on max-pooling over timesteps
-  SOFTMAX_MAX_POOLING_CLASS_LABEL = "softmax_max_pooling_class_label"
-  SPEECH_RECOGNITION = "speech_recognition"
-  SYMBOL = "symbol"
-  SYMBOL_WEIGHTS_ALL = "symbol_weights_all"  # symbol for features w/o 0-padding
-  SYMBOL_ONE_HOT = "symbol_one_hot"  # symbol with one hot as embeddings
-  VIDEO = "video"
-  VIDEO_BITWISE = "video_bitwise"  # video where bottom embeds pixels bitwise
-  VIDEO_IDENTITY = "video_identity"  # video with identity top and bottom
-  VIDEO_L1 = "video_l1"  # video with L2 loss
-  VIDEO_L2 = "video_l2"  # video with L1 loss
-  # video with L1 loss and raw input (sequences of frames)
-  VIDEO_L1_RAW = "video_l1_raw"
-  # video with L2 loss and raw input (sequences of frames)
-  VIDEO_L2_RAW = "video_l2_raw"
-  # video with pixel noise on input during training
-  VIDEO_PIXEL_NOISE = "video_pixel_noise"
 
-  @staticmethod
-  def get_choices():
-    return [
-        ModalityType.AUDIO,
-        ModalityType.AUDIO_SPECTRAL,
-        ModalityType.CLASS_LABEL,
-        ModalityType.CTC_SYMBOL,
-        ModalityType.GENERIC_L2_LOSS,
-        ModalityType.IDENTITY,
-        ModalityType.IDENTITY_SYMBOL,
-        ModalityType.IMAGE,
-        ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
-        ModalityType.IMAGE_CHANNEL_COMPRESS,
-        ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM,
-        ModalityType.MULTI_LABEL,
-        ModalityType.ONE_HOT_CLASS_LABEL,
-        ModalityType.REAL,
-        ModalityType.REAL_L2_LOSS,
-        ModalityType.REAL_LOG_POISSON_LOSS,
-        ModalityType.SIGMOID_CLASS_LABEL,
-        ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
-        ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
-        ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
-        ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL,
-        ModalityType.SPEECH_RECOGNITION,
-        ModalityType.SYMBOL,
-        ModalityType.SYMBOL_ONE_HOT,
-        ModalityType.SYMBOL_WEIGHTS_ALL,
-        ModalityType.VIDEO,
-        ModalityType.VIDEO_BITWISE,
-        ModalityType.VIDEO_IDENTITY,
-        ModalityType.VIDEO_L1,
-        ModalityType.VIDEO_L2,
-        ModalityType.VIDEO_L1_RAW,
-        ModalityType.VIDEO_L2_RAW,
-        ModalityType.VIDEO_PIXEL_NOISE,
-    ]
+@is_pointwise
+def symbol_top(body_output, targets, model_hparams, vocab_size):
+  """Generate logits.
+
+  Args:
+    body_output: A Tensor with shape
+      [batch, p0, p1, model_hparams.hidden_size].
+    targets: Unused.
+    model_hparams: tf.HParams, model hyperparmeters.
+    vocab_size: int, vocabulary size.
+
+  Returns:
+    logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].
+  """
+  del targets  # unused arg
+  if model_hparams.shared_embedding_and_softmax_weights:
+    scope_name = "shared"
+    reuse = tf.AUTO_REUSE
+  else:
+    scope_name = "softmax"
+    reuse = False
+  with tf.variable_scope(scope_name, reuse=reuse):
+    body_output_shape = common_layers.shape_list(body_output)
+    var = get_weights(model_hparams, vocab_size, body_output_shape[-1])
+    if (model_hparams.factored_logits and
+        model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
+      # insert channels dimension
+      body_output = tf.expand_dims(body_output, 3)
+      return common_layers.FactoredTensor(body_output, var)
+    else:
+      body_output = tf.reshape(body_output, [-1, body_output_shape[-1]])
+      logits = tf.matmul(body_output, var, transpose_b=True)
+      return tf.reshape(logits,
+                        body_output_shape[:-1] + [1, vocab_size])
+
+
+@is_pointwise
+def symbol_one_hot_top(body_output, targets, model_hparams, vocab_size):
+  del targets, model_hparams, vocab_size  # unused arg
+  return body_output
+
+
+def video_top(body_output, targets, model_hparams, vocab_size):
+  """Top transformation for video."""
+  del targets  # unused arg
+  num_channels = model_hparams.problem.num_channels
+  shape = common_layers.shape_list(body_output)
+  reshape_shape = shape[:-1] + [num_channels, vocab_size]
+  res = tf.reshape(body_output, reshape_shape)
+  # Calculate argmax so as to have a summary with the produced images.
+  x = tf.argmax(tf.reshape(res, [-1, vocab_size]), axis=-1)
+  x = tf.reshape(x, shape[:-1] + [num_channels])
+  common_video.gif_summary("results", x, max_outputs=1)
+  return res
+
+
+def video_l1_top(body_output, targets, model_hparams, vocab_size):
+  """Top transformation for video."""
+  del targets, vocab_size  # unused arg
+  num_channels = model_hparams.problem.num_channels
+  num_frames = model_hparams.video_num_target_frames
+  with tf.variable_scope("rgb"):
+    body_output_shape = common_layers.shape_list(body_output)
+    res = tf.layers.dense(body_output, num_channels * num_frames, name="cast")
+    res = tf.reshape(res, body_output_shape[:3] + [num_channels, num_frames])
+    res = tf.transpose(res, [0, 4, 1, 2, 3])  # Move frames next to batch.
+    if not tf.get_variable_scope().reuse:
+      res_argmax = res[:, -1, :, :, :]
+      tf.summary.image(
+          "result",
+          common_layers.tpu_safe_image_summary(res_argmax),
+          max_outputs=1)
+    return tf.expand_dims(res, axis=-1)  # Add an axis like in perplexity.
+
+
+def video_raw_top(body_output, targets, model_hparams, vocab_size):
+  del targets, model_hparams, vocab_size  # unused arg
+  frames = body_output
+  if isinstance(body_output, list):
+    frames = tf.stack(body_output, axis=1)
+  rgb_frames = common_layers.convert_real_to_rgb(frames)
+  common_video.gif_summary("body_output", rgb_frames)
+  return tf.expand_dims(rgb_frames, axis=-1)
 
 
-# Utility functions, similar to tf.keras
+# Utility functions similar to tf.keras for default transformations
 
 
 def get_bottom(modality_type, value=None):
@@ -1289,7 +1298,8 @@ def get_loss(modality_type, value=None):
 
 def get_name(modality_type, value=None):
   """Gets default name for transformations; if none available, return value."""
-  # For legacy reasons, modalities vary in their naming scheme.
+  # For legacy reasons, modalities vary in their naming scheme. Future plans are
+  # to remove any need for get_name. We do not recommend using it.
   if modality_type == ModalityType.AUDIO:
     return lambda model_hparams, vocab_size: "audio_modality"
   elif modality_type == ModalityType.AUDIO_SPECTRAL:

From 613483a63b61ea720a5cb1e5ad4e97138e08405f Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 6 Mar 2019 16:46:59 -0800
Subject: [PATCH 1768/2720] Re-order RL readme for a more natural sequence.

PiperOrigin-RevId: 237146400
---
 tensor2tensor/rl/README.md | 158 +++++++++++++++++++------------------
 1 file changed, 81 insertions(+), 77 deletions(-)

diff --git a/tensor2tensor/rl/README.md b/tensor2tensor/rl/README.md
index c1fbacb19..99c4a866f 100644
--- a/tensor2tensor/rl/README.md
+++ b/tensor2tensor/rl/README.md
@@ -15,7 +15,7 @@ To use this package, you need to install the Atari dependencies for OpenAI Gym:
 pip install gym[atari]
 ```
 
-## Evaluating a pretrained policy
+## Play using a pre-trained policy
 
 We provide a set of pretrained policies and models you can use. To evaluate and
 generate videos for a pretrained policy on Pong:
@@ -47,46 +47,41 @@ tensorboard --logdir=~/t2t_train/pong_pretrained
 
 Description of player controls and flags can be found in `tensor2tensor/rl/player.py`.
 
-## Model-based training with pretrained world models
 
-To train a policy with a pretrained world model (requires Google Cloud SDK):
+## Train your policy (model-free training)
+
+Training model-free on Pong:
 
 ```
-OUTPUT_DIR=~/t2t_train/mb_sd_pong_pretrained
-mkdir -p $OUTPUT_DIR
-gsutil -m cp -r \
-  gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/world_model \
-  $OUTPUT_DIR/
-python -m tensor2tensor.rl.trainer_model_based \
-  --loop_hparams_set=rlmb_long_stochastic_discrete \
-  --loop_hparams=game=pong,epochs=1,model_train_steps=0 \
-  --eval_world_model=False \
-  --output_dir=$OUTPUT_DIR
+python -m tensor2tensor.rl.trainer_model_free \
+  --hparams_set=rlmf_base \
+  --hparams=game=pong \
+  --output_dir=~/t2t_train/mf_pong
 ```
 
-Note that this command will collect some frames from the real environment for
-random starts.
-
-The same command can be used to resume interrupted training - checkpoints are
-saved in `output_dir`.
+Hyperparameter sets are defined in `tensor2tensor/models/research/rl.py`. You
+can override them using the `hparams` flag, e.g.
 
-We use `NoFrameskip-v4` game mode with our own frame skip (4 by default).
+```
+  --hparams=game=kung_fu_master,frame_stack_size=5
+```
 
-The training script runs periodic evaluation, but with timestep limit 1000 to
-make it faster. To do full evaluation after training, run:
+As in model-based training, the periodic evaluation runs with timestep limit
+of 1000. To do full evaluation after training, run:
 
 ```
+OUTPUT_DIR=~/t2t_train/mf_pong
 python -m tensor2tensor.rl.evaluator \
-  --loop_hparams_set=rlmb_long_stochastic_discrete \
+  --loop_hparams_set=rlmf_base \
   --hparams=game=pong \
   --policy_dir=$OUTPUT_DIR \
   --eval_metrics_dir=$OUTPUT_DIR/full_eval_metrics
 ```
 
-## Model training with random trajectories
+## World Model training (with random trajectories)
 
-The simplest way to train your own model is to use random trajectories. Then you
-can train a policy on it as described in the previous section.
+The simplest way to train your own world model is to use random trajectories.
+Then you can train a policy on it as described next.
 
 To train a deterministic model:
 
@@ -106,43 +101,7 @@ python -m tensor2tensor.rl.trainer_model_based \
   --output_dir=~/t2t_train/mb_sd_pong_random
 ```
 
-## Full model-based training
-
-Our full training pipeline involves alternating between collecting data using
-policy, training the world model and training the policy inside the model. It
-requires significantly more time (several days to a week, depending on your
-hardware and the model you use).
-
-To train a deterministic model:
-
-```
-python -m tensor2tensor.rl.trainer_model_based \
-  --loop_hparams_set=rlmb_base \
-  --loop_hparams=game=pong \
-  --output_dir ~/t2t_train/mb_det_pong
-```
-
-To train a stochastic discrete model:
-
-```
-python -m tensor2tensor.rl.trainer_model_based \
-  --loop_hparams_set=rlmb_long_stochastic_discrete \
-  --loop_hparams=game=pong \
-  --output_dir ~/t2t_train/mb_sd_pong
-```
-
-Hyperparameter sets are defined in
-`tensor2tensor/rl/trainer_model_based_params.py`. Hyperparameter sets for the
-world model and agent are nested within `loop_hparams` by name. You can change
-them with:
-
-```
-  --loop_hparams=game=freeway,generative_model=next_frame_basic_deterministic,base_algo_params=ppo_original_params
-```
-
-Game names should be provided in `snake_case`.
-
-## Playing in the model
+## Playing in the world model
 
 To assess world model quality you can play in it, as in an Atari emulator
 (you need a machine with GPU for this). First install `pygame`:
@@ -177,36 +136,81 @@ performance will be worse than during the policy training.
 
 For more details on controls and flags see `tensor2tensor/rl/player.py`.
 
-## Model-free training
 
-Training model-free on Pong:
+## Model-based training with pre-trained world models
+
+To train a policy with a pretrained world model (requires Google Cloud SDK):
 
 ```
-python -m tensor2tensor.rl.trainer_model_free \
-  --hparams_set=rlmf_base \
-  --hparams=game=pong \
-  --output_dir=~/t2t_train/mf_pong
+OUTPUT_DIR=~/t2t_train/mb_sd_pong_pretrained
+mkdir -p $OUTPUT_DIR
+gsutil -m cp -r \
+  gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/world_model \
+  $OUTPUT_DIR/
+python -m tensor2tensor.rl.trainer_model_based \
+  --loop_hparams_set=rlmb_long_stochastic_discrete \
+  --loop_hparams=game=pong,epochs=1,model_train_steps=0 \
+  --eval_world_model=False \
+  --output_dir=$OUTPUT_DIR
 ```
 
-Hyperparameter sets are defined in `tensor2tensor/models/research/rl.py`. You
-can override them using the `hparams` flag, e.g.
+Note that this command will collect some frames from the real environment for
+random starts.
 
-```
-  --hparams=game=kung_fu_master,frame_stack_size=5
-```
+The same command can be used to resume interrupted training - checkpoints are
+saved in `output_dir`.
 
-As in model-based training, the periodic evaluation runs with timestep limit
-of 1000. To do full evaluation after training, run:
+We use `NoFrameskip-v4` game mode with our own frame skip (4 by default).
+
+The training script runs periodic evaluation, but with timestep limit 1000 to
+make it faster. To do full evaluation after training, run:
 
 ```
-OUTPUT_DIR=~/t2t_train/mf_pong
 python -m tensor2tensor.rl.evaluator \
-  --loop_hparams_set=rlmf_base \
+  --loop_hparams_set=rlmb_long_stochastic_discrete \
   --hparams=game=pong \
   --policy_dir=$OUTPUT_DIR \
   --eval_metrics_dir=$OUTPUT_DIR/full_eval_metrics
 ```
 
+
+## Full model-based training
+
+Our full training pipeline involves alternating between collecting data using
+policy, training the world model and training the policy inside the model. It
+requires significantly more time (several days to a week, depending on your
+hardware and the model you use).
+
+To train a deterministic model:
+
+```
+python -m tensor2tensor.rl.trainer_model_based \
+  --loop_hparams_set=rlmb_base \
+  --loop_hparams=game=pong \
+  --output_dir ~/t2t_train/mb_det_pong
+```
+
+To train a stochastic discrete model:
+
+```
+python -m tensor2tensor.rl.trainer_model_based \
+  --loop_hparams_set=rlmb_long_stochastic_discrete \
+  --loop_hparams=game=pong \
+  --output_dir ~/t2t_train/mb_sd_pong
+```
+
+Hyperparameter sets are defined in
+`tensor2tensor/rl/trainer_model_based_params.py`. Hyperparameter sets for the
+world model and agent are nested within `loop_hparams` by name. You can change
+them with:
+
+```
+  --loop_hparams=game=freeway,generative_model=next_frame_basic_deterministic,base_algo_params=ppo_original_params
+```
+
+Game names should be provided in `snake_case`.
+
+
 ## Using checkpoints for other games
 
 We provide pretrained policies and stochastic discrete models for most of the

From 209d07f21cbe2c30c5a85a17dba636c4bfbe0ebd Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 6 Mar 2019 17:21:02 -0800
Subject: [PATCH 1769/2720] Add three new en-de problem definitions.

PiperOrigin-RevId: 237151609
---
 tensor2tensor/data_generators/translate.py    | 35 +++++++------
 .../data_generators/translate_ende.py         | 49 +++++++++++++++++--
 2 files changed, 66 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 769911b92..448d970f1 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import gzip
 import os
 import tarfile
 from tensor2tensor.data_generators import cleaner_en_xx
@@ -143,6 +144,18 @@ def _clean_sentences(sentence_pairs):
   return res_pairs
 
 
+def _tmx_to_source_target(tmx_file, source_resfile, target_resfile,
+                          do_cleaning=False):
+  source_target_pairs = cleaner_en_xx.paracrawl_v3_pairs(tmx_file)
+  if do_cleaning:
+    source_target_pairs = cleaner_en_xx.clean_en_xx_pairs(source_target_pairs)
+  for source, target in source_target_pairs:
+    source_resfile.write(source)
+    source_resfile.write("\n")
+    target_resfile.write(target)
+    target_resfile.write("\n")
+
+
 def compile_data(tmp_dir, datasets, filename, datatypes_to_clean=None):
   """Concatenates all `datasets` and saves to `filename`."""
   datatypes_to_clean = datatypes_to_clean or []
@@ -163,22 +176,16 @@ def compile_data(tmp_dir, datasets, filename, datatypes_to_clean=None):
           generator_utils.maybe_download(tmp_dir, compressed_filename, url)
 
         if dataset[1][0] == "tmx":
+          cleaning_requested = "tmx" in datatypes_to_clean
           tmx_filename = os.path.join(tmp_dir, dataset[1][1])
           if tmx_filename.endswith(".gz"):
-            new_filename = tmx_filename.strip(".gz")
-            if not tf.gfile.Exists(new_filename):
-              generator_utils.gunzip_file(tmx_filename, new_filename)
-            tmx_filename = new_filename
-          source, target = None, None
-          with tf.gfile.Open(tmx_filename) as tmx_file:
-            stream = cleaner_en_xx.paracrawl_v3_pairs(tmx_file)
-            if "tmx" in datatypes_to_clean:
-              stream = cleaner_en_xx.clean_en_xx_pairs(stream)
-            for source, target in stream:
-              lang1_resfile.write(source)
-              lang1_resfile.write("\n")
-              lang2_resfile.write(target)
-              lang2_resfile.write("\n")
+            with gzip.open(tmx_filename, "rb") as tmx_file:
+              _tmx_to_source_target(tmx_file, lang1_resfile, lang2_resfile,
+                                    do_cleaning=cleaning_requested)
+          else:
+            with tf.gfile.Open(tmx_filename) as tmx_file:
+              _tmx_to_source_target(tmx_file, lang1_resfile, lang2_resfile,
+                                    do_cleaning=cleaning_requested)
 
         elif dataset[1][0] == "tsv":
           _, src_column, trg_column, glob_pattern = dataset[1]
diff --git a/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index b788553b8..6c62dbbce 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -47,6 +47,12 @@
         ("dev/newstest2013.en", "dev/newstest2013.de")
     ],
 ]
+_ENDE_PARACRAWL_DATASETS = [
+    [
+        "https://s3.amazonaws.com/web-language-models/paracrawl/release3/en-de.bicleaner07.tmx.gz",  # pylint: disable=line-too-long
+        ("tmx", "en-de.bicleaner07.tmx.gz")
+    ]
+]
 
 
 @registry.register_problem
@@ -88,6 +94,35 @@ def datatypes_to_clean(self):
     return ["txt"]
 
 
+@registry.register_problem
+class TranslateEndeParacrawl32k(translate.TranslateProblem):
+  """Problem spec for Paracrawl en-de translation."""
+
+  @property
+  def use_vocab_from_other_problem(self):
+    return TranslateEndeWmt32k()
+
+  @property
+  def additional_training_datasets(self):
+    """Allow subclasses to add training datasets."""
+    return []
+
+  def source_data_files(self, dataset_split):
+    train = dataset_split == problem.DatasetSplit.TRAIN
+    train_datasets = (
+        _ENDE_PARACRAWL_DATASETS + self.additional_training_datasets)
+    return train_datasets if train else _ENDE_EVAL_DATASETS
+
+
+@registry.register_problem
+class TranslateEndeParacrawlClean32k(TranslateEndeParacrawl32k):
+  """Paracrawl en-de Bicleaner corpus, with additional cleaning."""
+
+  @property
+  def datatypes_to_clean(self):
+    return ["tmx"]
+
+
 @registry.register_problem
 class TranslateEndeWmtParacrawlBicleaner32k(TranslateEndeWmt32k):
   """WMT en-de corpus with extra data from Paracrawl, cleaned with Bicleaner."""
@@ -98,9 +133,15 @@ def use_vocab_from_other_problem(self):
 
   @property
   def additional_training_datasets(self):
-    paracrawl = "https://s3.amazonaws.com/web-language-models/paracrawl/"
-    return [(paracrawl + "release3/en-de.bicleaner07.tmx.gz",
-             ("tmx", "en-de.bicleaner07.tmx.gz"))]
+    return _ENDE_PARACRAWL_DATASETS
+
+
+@registry.register_problem
+class TranslateEndeWmtCleanParacrawl32k(TranslateEndeWmtParacrawlBicleaner32k):
+
+  @property
+  def datatypes_to_clean(self):
+    return ["txt"]
 
 
 @registry.register_problem
@@ -116,7 +157,7 @@ class TranslateEndeWmtParacrawlAllClean32k(TranslateEndeWmtParacrawlClean32k):
 
   @property
   def datatypes_to_clean(self):
-    return ["tmx", "txt"]
+    return ["txt", "tmx"]
 
 
 @registry.register_problem

From 1b08988f0d691223d4534b85f426f29dc5a39b1c Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 6 Mar 2019 17:26:48 -0800
Subject: [PATCH 1770/2720] Transformer LM runs in TRAX.

PiperOrigin-RevId: 237152409
---
 tensor2tensor/trax/README.md                  |  8 +++
 .../trax/configs/transformer_lm1b_8gb.gin     | 48 ++++++++++++++
 tensor2tensor/trax/inputs.py                  | 26 +++++---
 tensor2tensor/trax/models/__init__.py         |  2 +-
 tensor2tensor/trax/models/transformer.py      | 63 ++++++++++++++++++-
 tensor2tensor/trax/stax/attention.py          | 11 +++-
 tensor2tensor/trax/stax/slax.py               | 12 ++++
 tensor2tensor/trax/trainer.py                 |  5 +-
 tensor2tensor/trax/trax.py                    | 31 ++++++---
 9 files changed, 183 insertions(+), 23 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/transformer_lm1b_8gb.gin

diff --git a/tensor2tensor/trax/README.md b/tensor2tensor/trax/README.md
index c13f296a3..71715aaf2 100644
--- a/tensor2tensor/trax/README.md
+++ b/tensor2tensor/trax/README.md
@@ -40,6 +40,14 @@ python -m tensor2tensor.trax.trainer \
   --config_file=$PWD/trax/configs/resnet50_imagenet_8gb.gin
 ```
 
+#### TransformerDecoder on LM1B
+
+
+```
+python -m tensor2tensor.trax.trainer \
+  --config_file=$PWD/trax/configs/transformer_lm1b_8gb.gin
+```
+
 ### How `trax` differs from T2T
 
 * Configuration is done with [`gin`](https://github.com/google/gin-config).
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
new file mode 100644
index 000000000..7a3f28c5a
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -0,0 +1,48 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size = 32
+batch_fun.eval_batch_size = 32
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_languagemodel_lm1b32k'
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 0.05
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+preprocess_fun.max_target_length = 256
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 1000
+train.eval_steps = 1
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerLM
+train.run_debug_step = False
+train.train_steps = 50000
+
+# Parameters for train_and_eval_batches:
+# ==============================================================================
+train_and_eval_batches.input_name = 'targets'
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.dropout = 0.1
+TransformerLM.feature_depth = 512
+TransformerLM.feedforward_depth = 2048
+TransformerLM.max_len = 256
+TransformerLM.mode = 'train'
+TransformerLM.num_heads = 8
+TransformerLM.num_layers = 6
+TransformerLM.vocab_size = 32000
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 992e0cabc..a24cbb6ca 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -129,7 +129,7 @@ def _make_info(shape_list, num_classes):
 def _select_features(example, feature_list=None):
   """Select a subset of features from the example dict."""
   feature_list = feature_list or ["inputs", "targets"]
-  return {f: example[f] for f in feature_list}
+  return {f: example[f] for f in feature_list if f in example}
 
 
 def _train_and_eval_dataset_v1(problem_name, data_dir):
@@ -140,7 +140,6 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
   train_dataset = train_dataset.map(_select_features)
   eval_dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, data_dir)
   eval_dataset = eval_dataset.map(_select_features)
-  supervised_keys = (["inputs"], ["targets"])
   hparams = problem.get_hparams()
   # We take a few training examples to guess the shapes.
   input_shapes, target_shapes = [], []
@@ -149,14 +148,18 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
   example1 = sess.run(example_tensor)
   example2 = sess.run(example_tensor)
   example3 = sess.run(example_tensor)
+  # We use "inputs" as input except for purely auto-regressive tasks like
+  # language models where "targets" are used as input_key.
+  input_key = "inputs" if "inputs" in example1 else "targets"
+  supervised_keys = ([input_key], ["targets"])
   for example in [example1, example2, example3]:
-    input_shapes.append(list(example["inputs"].shape))
+    input_shapes.append(list(example[input_key].shape))
     target_shapes.append(list(example["targets"].shape))
-  input_vocab_size = hparams.vocab_size["inputs"]
+  input_vocab_size = hparams.vocab_size[input_key]
   target_vocab_size = hparams.vocab_size["targets"]
   input_info = _make_info(input_shapes, input_vocab_size)
   target_info = _make_info(target_shapes, target_vocab_size)
-  info = {"inputs": input_info, "targets": target_info}
+  info = {input_key: input_info, "targets": target_info}
   return train_dataset, eval_dataset, info, supervised_keys
 
 
@@ -191,7 +194,7 @@ def batch_fun(dataset, training, shapes, target_names,
                            bucket_length * 4, bucket_length * 8]
       bucket_batch_sizes = [cur_batch_size * 4, cur_batch_size * 2,
                             cur_batch_size, cur_batch_size // 2,
-                            cur_batch_size // 4, cur_batch_size // 8]
+                            cur_batch_size // 4, cur_batch_size // 8, 1]
       buckets = (bucket_boundaries, bucket_batch_sizes)
 
   if buckets:
@@ -200,7 +203,8 @@ def example_length(_, target):
       return tf.shape(target)[0]
     boundaries, batch_sizes = buckets
     dataset = dataset.apply(tf.data.experimental.bucket_by_sequence_length(
-        example_length, boundaries, batch_sizes, pad_to_bucket_boundary=True))
+        example_length, boundaries, batch_sizes,
+        pad_to_bucket_boundary=training))
   else:
     dataset = dataset.padded_batch(cur_batch_size, shapes)
   return dataset
@@ -227,7 +231,8 @@ def append_targets(example):
   return dataset.prefetch(32)
 
 
-def train_and_eval_batches(dataset, data_dir):
+@gin.configurable(whitelist=["input_name"])
+def train_and_eval_batches(dataset, data_dir, input_name=None):
   """Return train and eval batches with input name and shape."""
   (train_data, eval_data, features_info, keys) = train_and_eval_dataset(
       dataset, data_dir)
@@ -236,5 +241,6 @@ def train_and_eval_batches(dataset, data_dir):
       train_data, target_names, features_info, training=True)
   eval_batches = shuffle_and_batch_data(
       eval_data, target_names, features_info, training=False)
-  input_shape = features_info[input_names[0]].shape
-  return train_batches, eval_batches, input_names[0], list(input_shape)
+  input_name = input_name or input_names[0]
+  input_shape = features_info[input_name].shape
+  return train_batches, eval_batches, input_name, list(input_shape)
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index 07fce9b0f..ceaf732bf 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -34,4 +34,4 @@ def model_configure(*args, **kwargs):
 # pylint: disable=invalid-name
 MLP = model_configure(mlp.MLP)
 Resnet50 = model_configure(resnet.Resnet50)
-Transformer = model_configure(transformer.Transformer)
+TransformerLM = model_configure(transformer.TransformerLM)
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 0d79bf14c..34d73e744 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -44,7 +44,6 @@ def TransformerEncoder(mode='train',  # pylint: disable=invalid-name
     A staxlayer for implementing a raw Transformer encoder stack.  No embedding
     or positional signals are added by this layer.
   """
-
   # Multi-headed Attention and Feed-forward layers
   multi_attention = stax.MultiHeadedAttention(
       feature_depth, num_heads=num_heads, dropout=dropout, mode=mode)
@@ -90,6 +89,68 @@ def encoder(embedded_source, source_mask):
   return encoder
 
 
+def TransformerLM(vocab_size,  # pylint: disable=invalid-name
+                  mode='train',
+                  num_layers=6,
+                  feature_depth=512,
+                  feedforward_depth=2048,
+                  num_heads=8,
+                  dropout=0.9,
+                  max_len=256):
+  """Transformer language model (only uses the decoder part of Transformer).
+
+  Args:
+    vocab_size: int: vocab size
+    mode: str: 'train' or 'eval'
+    num_layers: int: number of encoder/decoder layers
+    feature_depth: int:  depth of embedding
+    feedforward_depth: int: depth of feed-forward layer
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate - Stax follows TF's KEEP probability convention
+    max_len: int: maximum symbol length for positional encoding
+
+  Returns:
+    init and apply.
+  """
+  # Multi-headed Attention and Feed-forward layers
+  multi_attention = stax.MultiHeadedAttention(
+      feature_depth, num_heads=num_heads, dropout=dropout, mode=mode)
+
+  feed_forward = stax.serial(
+      stax.Dense(feedforward_depth, W_init=stax.xavier_uniform()),
+      stax.Relu,
+      stax.Dropout(dropout, mode=mode),
+      stax.Dense(feature_depth, W_init=stax.xavier_uniform())
+  )
+
+  # Single decoder layer
+  decoder_layer = stax.serial(
+      # target attends to self
+      stax.residual(stax.LayerNorm(feature_depth),
+                    stax.multiplex(stax.Identity,  # query
+                                   stax.Identity,  # key
+                                   stax.Identity,  # value
+                                   stax.CausalMask(axis=-2)),  # attention mask
+                    multi_attention,
+                    stax.Dropout(dropout, mode=mode)),
+      # feed-forward
+      stax.residual(stax.LayerNorm(feature_depth),
+                    feed_forward,
+                    stax.Dropout(dropout, mode=mode))
+  )
+
+  return stax.serial(
+      stax.ShiftRight(),
+      stax.Embedding(feature_depth, vocab_size),
+      stax.PositionalEncoding(feature_depth, max_len=max_len),
+      stax.Dropout(dropout, mode=mode),
+      stax.repeat(decoder_layer, num_layers),
+      stax.LayerNorm(feature_depth),
+      stax.Dense(vocab_size, W_init=stax.xavier_uniform()),
+      stax.LogSoftmax
+  )
+
+
 def Transformer(source_vocab_size,  # pylint: disable=invalid-name
                 target_vocab_size,
                 mode='train',
diff --git a/tensor2tensor/trax/stax/attention.py b/tensor2tensor/trax/stax/attention.py
index 238f8b9ec..d9f2fd52c 100644
--- a/tensor2tensor/trax/stax/attention.py
+++ b/tensor2tensor/trax/stax/attention.py
@@ -30,6 +30,15 @@ def causal_mask(size, dtype=np.uint8):
   return onp.tril(onp.ones((1, size, size), dtype=dtype), k=0)
 
 
+def CausalMask(axis=-1):  # pylint: disable=invalid-name
+  """Layer to create a causal mask for its inputs."""
+  init_fun = lambda input_shape: (input_shape, ())
+  def apply_fun(params, inputs, **kwargs):
+    del params, kwargs
+    return causal_mask(inputs.shape[axis], dtype=inputs.dtype)
+  return init_fun, apply_fun
+
+
 def make_target_mask(target, pad=0):
   """Create an attention mask to hide padding and future words."""
   target_mask = (target != pad)[ :, np.newaxis, :]
@@ -93,7 +102,7 @@ def apply_fun(params, inputs, **kwargs):
 def Embedding(feature_depth, vocab_size):  # pylint: disable=invalid-name
   """Layer constructor function for a dense embedding layer."""
   def init_fun(input_shape):
-    output_shape = input_shape + (feature_depth,)
+    output_shape = tuple(input_shape) + (feature_depth,)
     dense_embedding = xavier_uniform()((vocab_size, feature_depth))
     return output_shape, dense_embedding
   def apply_fun(params, inputs, **kwargs):
diff --git a/tensor2tensor/trax/stax/slax.py b/tensor2tensor/trax/stax/slax.py
index e97a8c2fc..d092f4387 100644
--- a/tensor2tensor/trax/stax/slax.py
+++ b/tensor2tensor/trax/stax/slax.py
@@ -32,6 +32,18 @@ def one_hot(x, size, dtype=np.float32):
   return np.array(x[..., np.newaxis] == np.arange(size), dtype)
 
 
+def ShiftRight():  # pylint: disable=invalid-name
+  """Layer to shift the tensor to the right by padding on axis 1."""
+  init_fun = lambda input_shape: (input_shape, ())
+  def apply_fun(params, inputs, **kwargs):
+    del params, kwargs
+    pad_widths = [(0, 0), (1, 0)]
+    pad_widths += [(0, 0) for _ in range(len(inputs.shape) - 2)]
+    padded = np.pad(inputs, pad_widths, mode='constant')
+    return padded[:, :-1, ...]
+  return init_fun, apply_fun
+
+
 # Utility Combinators
 # ------------------------------------------------------------------------------
 def repeat(layer, num_repeats):
diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index 32ce88107..9f4cfb6db 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -70,7 +70,10 @@ def _setup_gin():
   # Override with --dataset and --model
   if FLAGS.dataset:
     configs.append("inputs.dataset_name='%s'" % FLAGS.dataset)
-    configs.append("inputs.data_dir='%s'" % FLAGS.data_dir)
+    if FLAGS.data_dir:
+      configs.append("inputs.data_dir='%s'" % FLAGS.data_dir)
+    else:
+      configs.append("inputs.data_dir=None")
     configs.append("train.inputs=@trax.inputs.inputs")
   if FLAGS.model:
     configs.append("train.model=@trax.models.%s" % FLAGS.model)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 98c5330f9..876ddc95d 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -33,6 +33,7 @@
 import jax
 from jax.experimental import optimizers as jax_opt
 import jax.numpy as np
+import jax.random as random
 
 import six
 
@@ -40,15 +41,11 @@
 from tensor2tensor.trax import jaxboard
 from tensor2tensor.trax import learning_rate as lr
 from tensor2tensor.trax import optimizers as trax_opt
+import tensor2tensor.trax.stax as stax
 
 from tensorflow.io import gfile
 
 
-def one_hot(x, k, dtype=np.float32):
-  """Create a one-hot encoding of x of size k."""
-  return np.array(x[:, None] == np.arange(k), dtype)
-
-
 def accuracy(batch, model_predictions):
   """Calculate accuracy."""
   _, targets = batch
@@ -59,7 +56,7 @@ def accuracy(batch, model_predictions):
 def neg_log_perplexity(batch, model_predictions):
   """Calculate negative log perplexity."""
   _, targets = batch
-  hot_targets = one_hot(targets, model_predictions.shape[-1])
+  hot_targets = stax.one_hot(targets, model_predictions.shape[-1])
   return np.mean(np.sum(model_predictions * hot_targets, axis=-1))
 
 
@@ -67,7 +64,8 @@ def loss(params, batch, model_predict):
   """Calculate loss."""
   inputs, targets = batch
   preds = model_predict(params, inputs)
-  return - np.mean(np.sum(preds * one_hot(targets, preds.shape[-1]), axis=-1))
+  return - np.mean(np.sum(preds * stax.one_hot(targets, preds.shape[-1]),
+                          axis=-1))
 
 
 def log(s, stdout=True):
@@ -233,7 +231,8 @@ def train(output_dir,
           lr_schedule=lr.MultifactorSchedule,
           train_steps=1000,
           eval_steps=10,
-          eval_frequency=100):
+          eval_frequency=100,
+          run_debug_step=False):
   """Train the model on the inputs.
 
   Args:
@@ -249,10 +248,13 @@ def train(output_dir,
     eval_steps: int, num of steps per evaluation. If None or 0, eval disabled.
     eval_frequency: int, how often to run evaluation (every eval_frequency
       steps). If None or 0, eval disabled.
+    run_debug_step: bool, if True, will run the model and loss without @jit for
+      one step.
 
   Returns:
     trax.State
   """
+  rng = random.PRNGKey(0)
   gfile.makedirs(output_dir)
   # Create summary writers and history.
   train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
@@ -265,7 +267,13 @@ def train(output_dir,
   history = state.history
   lr_fun = lr_schedule(history)
   opt_init, _ = optimizer(lr_fun)
-  model_init, model_predict = model()
+  model_init, model_predict_original = model()
+  # We need a model_predict that fills in the random generator if needed.
+  def model_predict(x, y, **kwargs):
+    """Same as model_predict_original but fill in rng if it isn't passed."""
+    if "rng" in kwargs:
+      return model_predict_original(x, y, **kwargs)
+    return model_predict_original(x, y, rng=rng, **kwargs)
 
   # Setup state
   step = state.step or 0
@@ -284,6 +292,11 @@ def train(output_dir,
                                 itertools.repeat(eval_frequency))
   step_log(step, "Starting training")
 
+  # Non-compiled debug step helps find problems in models easier.
+  if run_debug_step:
+    debug_loss = loss(params, next(train_stream), model_predict)
+    step_log(step, "Debug step loss %.8f" % debug_loss)
+
   for epoch, epoch_steps in epochs(train_steps, epoch_steps):
     # Log separator
     print()

From b44804149ef543385c584f756bf2d2ad90b68299 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 6 Mar 2019 21:01:50 -0800
Subject: [PATCH 1771/2720] Don't export functions when saving hparams to json.

PiperOrigin-RevId: 237176017
---
 tensor2tensor/utils/hparam.py      | 12 ++++++------
 tensor2tensor/utils/hparam_test.py |  9 +++++++++
 tensor2tensor/utils/hparams_lib.py |  8 ++++++++
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/utils/hparam.py b/tensor2tensor/utils/hparam.py
index 15ce59677..5bd58035b 100644
--- a/tensor2tensor/utils/hparam.py
+++ b/tensor2tensor/utils/hparam.py
@@ -542,15 +542,15 @@ def to_json(self, indent=None, separators=None, sort_keys=False):
       A JSON string.
     """
     def remove_callables(x):
-      if callable(x):
-        return x.__name__
+      """Omit callable elements from input with arbitrary nesting."""
       if isinstance(x, dict):
-        return {k: remove_callables(v) for k, v in six.iteritems(x)}
-      if isinstance(x, list):
-        return [remove_callables(i) for i in x]
+        return {k: remove_callables(v) for k, v in six.iteritems(x)
+                if not callable(v)}
+      elif isinstance(x, list):
+        return [remove_callables(i) for i in x if not callable(i)]
       return x
     return json.dumps(
-        {k: remove_callables(v) for k, v in six.iteritems(self.values())},
+        remove_callables(self.values()),
         indent=indent,
         separators=separators,
         sort_keys=sort_keys)
diff --git a/tensor2tensor/utils/hparam_test.py b/tensor2tensor/utils/hparam_test.py
index 600349ffd..85790ded2 100644
--- a/tensor2tensor/utils/hparam_test.py
+++ b/tensor2tensor/utils/hparam_test.py
@@ -136,6 +136,15 @@ def testSetFromMap(self):
     self.assertDictEqual({'d': [0.1, 0.2, 0.3], 'x': 1, 'b': 2.0},
                          hparams.values())
 
+  def testFunction(self):
+    def f(x):
+      return x
+    hparams = hparam.HParams(function=f)
+    self.assertEqual(hparams.function, f)
+
+    json_str = hparams.to_json()
+    self.assertEqual(json_str, '{}')
+
   def testBoolParsing(self):
     for value in 'true', 'false', 'True', 'False', '1', '0':
       for initial in False, True:
diff --git a/tensor2tensor/utils/hparams_lib.py b/tensor2tensor/utils/hparams_lib.py
index 0d02a291c..3db112e96 100644
--- a/tensor2tensor/utils/hparams_lib.py
+++ b/tensor2tensor/utils/hparams_lib.py
@@ -64,6 +64,14 @@ def create_hparams_from_json(json_path, hparams=None):
   tf.logging.info("Loading hparams from existing json %s" % json_path)
   with tf.gfile.Open(json_path, "r") as f:
     hparams_values = json.load(f)
+    # Prevent certain keys from overwriting the passed-in hparams.
+    # TODO(trandustin): Remove this hack after registries are available to avoid
+    # saving them as functions.
+    hparams_values.pop("bottom", None)
+    hparams_values.pop("loss", None)
+    hparams_values.pop("name", None)
+    hparams_values.pop("top", None)
+    hparams_values.pop("weights_fn", None)
     new_hparams = HParams(**hparams_values)
     # Some keys are in new_hparams but not hparams, so we need to be more
     #   careful than simply using parse_json() from HParams

From 8de05840752a44ea472adf027aed88696c69216d Mon Sep 17 00:00:00 2001
From: Ryan Sepassi <rsepassi@google.com>
Date: Thu, 7 Mar 2019 10:29:13 -0800
Subject: [PATCH 1772/2720] Update trax_test to test_train_eval_predict

PiperOrigin-RevId: 237274100
---
 tensor2tensor/trax/trax_test.py | 41 ++++++++++++++++++---------------
 1 file changed, 23 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index 7a2905e83..8eadee4fe 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -25,7 +25,7 @@
 
 import numpy as np
 
-from tensor2tensor.trax import inputs
+from tensor2tensor.trax import inputs as inputs_lib
 from tensor2tensor.trax import models
 from tensor2tensor.trax import trax
 
@@ -43,7 +43,7 @@ def input_stream():
       yield (np.random.rand(*([batch_size] + list(input_shape))),
              np.random.randint(num_classes, size=batch_size))
 
-  return inputs.Inputs(
+  return inputs_lib.Inputs(
       train_stream=input_stream,
       eval_stream=input_stream,
       input_shape=input_shape)
@@ -57,32 +57,37 @@ def tmp_dir(self):
     yield tmp
     gfile.rmtree(tmp)
 
-  @property
-  def train_args(self):
-    num_classes = 4
-    return dict(
-        model=functools.partial(models.MLP,
+  def test_train_eval_predict(self):
+    with self.tmp_dir() as output_dir:
+      # Prepare model and inputs
+      num_classes = 4
+      train_steps = 2
+      eval_steps = 2
+      model = functools.partial(models.MLP,
                                 hidden_size=16,
-                                num_output_classes=num_classes),
-        inputs=lambda: test_inputs(num_classes),
-        train_steps=3,
-        eval_steps=2)
+                                num_output_classes=num_classes)
+      inputs = lambda: test_inputs(num_classes)
 
-  def _test_train(self, train_args):
-    with self.tmp_dir() as output_dir:
-      state = trax.train(output_dir, **train_args)
+      # Train and evaluate
+      state = trax.train(output_dir,
+                         model=model,
+                         inputs=inputs,
+                         train_steps=train_steps,
+                         eval_steps=eval_steps)
 
       # Assert total train steps
-      self.assertEqual(train_args["train_steps"], state.step)
+      self.assertEqual(train_steps, state.step)
 
-      # Assert 2 epochs ran
+      # Assert 2 evaluations ran
       train_acc = state.history.get("train", "metrics/accuracy")
       eval_acc = state.history.get("eval", "metrics/accuracy")
       self.assertEqual(len(train_acc), len(eval_acc))
       self.assertEqual(2, len(eval_acc))
 
-  def test_train(self):
-    self._test_train(self.train_args)
+      # Predict with final params
+      _, predict_fun = model()
+      inputs = inputs().train_stream()
+      predict_fun(state.params, next(inputs)[0])
 
 
 if __name__ == "__main__":

From 19c88717bc52bc9a212c1ab1edbdb3f65b125c39 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 7 Mar 2019 14:03:36 -0800
Subject: [PATCH 1773/2720]  - Disable the MuJoCo test from Travis since it
 needs installation, license and setup.  - Added a missing set_status from a
 previous change.  - Disable overly verbose warnings summary from pytest.

PiperOrigin-RevId: 237319913
---
 oss_scripts/oss_tests.sh | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 263a11630..ea31b21ce 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -43,13 +43,14 @@ set_status
 # trax tests need C++
 # others (see below) enable eager, so can't be tested along with the others in
 # pytest
-pytest \
+pytest --disable-warnings \
   --ignore=tensor2tensor/bin/t2t_trainer_test.py \
   --ignore=tensor2tensor/data_generators/algorithmic_math_test.py \
   --ignore=tensor2tensor/data_generators/allen_brain_test.py \
   --ignore=tensor2tensor/data_generators/ops/pack_sequences_ops_test.py \
   --ignore=tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py \
   --ignore=tensor2tensor/data_generators/problem_test.py \
+  --ignore=tensor2tensor/envs/mujoco_problems_test.py \
   --ignore=tensor2tensor/layers/bayes_test.py \
   --ignore=tensor2tensor/layers/common_attention_test.py \
   --ignore=tensor2tensor/layers/common_layers_test.py \
@@ -66,10 +67,10 @@ pytest \
   --ignore=tensor2tensor/models/video/nfg_uncond_test.py \
   --ignore=tensor2tensor/rl \
   --ignore=tensor2tensor/trax \
+  --ignore=tensor2tensor/utils/registry_test.py \
   --ignore=tensor2tensor/utils/t2t_model_test.py \
   --ignore=tensor2tensor/utils/test_utils.py \
   --ignore=tensor2tensor/utils/test_utils_test.py \
-  --ignore=tensor2tensor/utils/registry_test.py \
   --ignore=tensor2tensor/utils/trainer_lib_test.py \
   --ignore=tensor2tensor/visualization/visualization_test.py \
   --deselect=tensor2tensor/utils/beam_search_test.py::BeamSearchTest::testTPUBeam
@@ -80,7 +81,7 @@ set_status
 # ImportError: /usr/lib/x86_64-linux-gnu/libstdc++.so.6: version `GLIBCXX_3.4.21' not found (required by /home/travis/virtualenv/python3.6.3/lib/python3.6/site-packages/jaxlib/_pywrap_xla.so)
 
 # These tests enable eager, so are tested separately.
-pytest \
+pytest --disable-warnings \
   tensor2tensor/data_generators/problem_test.py \
   tensor2tensor/layers/bayes_test.py \
   tensor2tensor/layers/common_attention_test.py \
@@ -94,17 +95,18 @@ pytest \
   tensor2tensor/utils/t2t_model_test.py \
   tensor2tensor/utils/test_utils_test.py \
   --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary
+set_status
 
-pytest tensor2tensor/utils/registry_test.py
+pytest --disable-warnings tensor2tensor/utils/registry_test.py
 set_status
 
-pytest tensor2tensor/utils/trainer_lib_test.py
+pytest --disable-warnings tensor2tensor/utils/trainer_lib_test.py
 set_status
 
-pytest tensor2tensor/visualization/visualization_test.py
+pytest --disable-warnings tensor2tensor/visualization/visualization_test.py
 set_status
 
-pytest tensor2tensor/data_generators/allen_brain_test.py
+pytest --disable-warnings tensor2tensor/data_generators/allen_brain_test.py
 set_status
 
 
@@ -113,7 +115,7 @@ if [[ "$TRAVIS_PYTHON_VERSION" == "2.7"  ]] && [[ "$TF_VERSION" == "tf-nightly"
 then
   # Ignores:
   # * Glow requires the CIFAR-10 dataset to be generated
-  pytest tensor2tensor/models/research \
+  pytest --disable-warnings tensor2tensor/models/research \
     --ignore=tensor2tensor/models/research/glow_test.py
   set_status
 fi
@@ -121,9 +123,11 @@ fi
 if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]
 then
     # TODO(afrozm): Once we drop support for 1.10 we can get rid of this.
-    pytest tensor2tensor/utils/beam_search_test.py::BeamSearchTest::testTPUBeam
+    pytest --disable-warnings \
+      tensor2tensor/utils/beam_search_test.py::BeamSearchTest::testTPUBeam
     set_status
     # TODO(afrozm): Enable other tests in the RL directory.
+    # Can't add disable warning here since it parses flags.
     pytest tensor2tensor/rl/trainer_model_based_test.py
     set_status
     jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/hello_t2t.ipynb

From 6671139107636a62ca37ffc8b2063572a19116b6 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 7 Mar 2019 14:25:33 -0800
Subject: [PATCH 1774/2720] Update datagen to work with environment problems.
 This will run the environment with randomly sampled actions and records the
 environment state + actions.

Flags:
 - env_problem_batch_size: Controls how many examples we generate
 - env_problem_max_env_steps: Controls how many steps to run in each example.
PiperOrigin-RevId: 237324759
---
 tensor2tensor/bin/t2t_datagen.py              | 159 +++++++++++-------
 tensor2tensor/data_generators/all_problems.py |   2 +
 tensor2tensor/envs/env_problem.py             |   9 +-
 3 files changed, 106 insertions(+), 64 deletions(-)

diff --git a/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index a985eaedb..9f5acb2e5 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -38,6 +38,7 @@
 
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
 from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import usr_dir
 
@@ -54,7 +55,6 @@
 # Improrting here to prevent pylint from ungrouped-imports warning.
 import tensorflow as tf  # pylint: disable=g-import-not-at-top
 
-
 flags = tf.flags
 FLAGS = flags.FLAGS
 
@@ -65,10 +65,18 @@
                     "The name of the problem to generate data for.")
 flags.DEFINE_string("exclude_problems", "",
                     "Comma-separates list of problems to exclude.")
-flags.DEFINE_integer("num_shards", 0, "How many shards to use. Ignored for "
-                     "registered Problems.")
+flags.DEFINE_integer(
+    "num_shards", 0, "How many shards to use. Ignored for "
+    "registered Problems.")
 flags.DEFINE_integer("max_cases", 0,
                      "Maximum number of cases to generate (unbounded if 0).")
+flags.DEFINE_integer(
+    "env_problem_max_env_steps", 0,
+    "Maximum number of steps to take for environment-based problems. "
+    "Actions are chosen randomly")
+flags.DEFINE_integer(
+    "env_problem_batch_size", 0,
+    "Number of environments to simulate for environment-based problems.")
 flags.DEFINE_bool("only_list", False,
                   "If true, we only list the problems that will be generated.")
 flags.DEFINE_integer("random_seed", 429459, "Random seed to use.")
@@ -78,58 +86,66 @@
 flags.DEFINE_integer(
     "num_concurrent_processes", None,
     "Applies only to problems for which multiprocess_generate=True.")
-flags.DEFINE_string("t2t_usr_dir", "",
-                    "Path to a Python module that will be imported. The "
-                    "__init__.py file should include the necessary imports. "
-                    "The imported files should contain registrations, "
-                    "e.g. @registry.register_problem calls, that will then be "
-                    "available to t2t-datagen.")
+flags.DEFINE_string(
+    "t2t_usr_dir", "", "Path to a Python module that will be imported. The "
+    "__init__.py file should include the necessary imports. "
+    "The imported files should contain registrations, "
+    "e.g. @registry.register_problem calls, that will then be "
+    "available to t2t-datagen.")
 
 # Mapping from problems that we can generate data for to their generators.
 # pylint: disable=g-long-lambda
 _SUPPORTED_PROBLEM_GENERATORS = {
-    "algorithmic_algebra_inverse": (
-        lambda: algorithmic_math.algebra_inverse(26, 0, 2, 100000),
-        lambda: algorithmic_math.algebra_inverse(26, 3, 3, 10000),
-        lambda: None),  # test set
-    "parsing_english_ptb8k": (
-        lambda: wsj_parsing.parsing_token_generator(
+    "algorithmic_algebra_inverse":
+        (lambda: algorithmic_math.algebra_inverse(26, 0, 2, 100000),
+         lambda: algorithmic_math.algebra_inverse(26, 3, 3, 10000),
+         lambda: None),  # test set
+    "parsing_english_ptb8k":
+        (lambda: wsj_parsing.parsing_token_generator(
             FLAGS.data_dir, FLAGS.tmp_dir, True, 2**13, 2**9),
-        lambda: wsj_parsing.parsing_token_generator(
-            FLAGS.data_dir, FLAGS.tmp_dir, False, 2**13, 2**9),
-        lambda: None),  # test set
-    "parsing_english_ptb16k": (
-        lambda: wsj_parsing.parsing_token_generator(
+         lambda: wsj_parsing.parsing_token_generator(
+             FLAGS.data_dir, FLAGS.tmp_dir, False, 2**13, 2**9),
+         lambda: None),  # test set
+    "parsing_english_ptb16k":
+        (lambda: wsj_parsing.parsing_token_generator(
             FLAGS.data_dir, FLAGS.tmp_dir, True, 2**14, 2**9),
-        lambda: wsj_parsing.parsing_token_generator(
-            FLAGS.data_dir, FLAGS.tmp_dir, False, 2**14, 2**9),
-        lambda: None),  # test set
-    "inference_snli32k": (
-        lambda: snli.snli_token_generator(FLAGS.tmp_dir, True, 2**15),
-        lambda: snli.snli_token_generator(FLAGS.tmp_dir, False, 2**15),
-        lambda: None),  # test set
-    "audio_timit_characters_test": (
-        lambda: audio.timit_generator(
-            FLAGS.data_dir, FLAGS.tmp_dir, True, 1718),
-        lambda: audio.timit_generator(
-            FLAGS.data_dir, FLAGS.tmp_dir, False, 626),
-        lambda: None),  # test set
-    "audio_timit_tokens_8k_test": (
-        lambda: audio.timit_generator(
-            FLAGS.data_dir, FLAGS.tmp_dir, True, 1718,
-            vocab_filename="vocab.endefr.%d" % 2**13, vocab_size=2**13),
-        lambda: audio.timit_generator(
-            FLAGS.data_dir, FLAGS.tmp_dir, False, 626,
-            vocab_filename="vocab.endefr.%d" % 2**13, vocab_size=2**13),
-        lambda: None),  # test set
-    "audio_timit_tokens_32k_test": (
-        lambda: audio.timit_generator(
-            FLAGS.data_dir, FLAGS.tmp_dir, True, 1718,
-            vocab_filename="vocab.endefr.%d" % 2**15, vocab_size=2**15),
-        lambda: audio.timit_generator(
-            FLAGS.data_dir, FLAGS.tmp_dir, False, 626,
-            vocab_filename="vocab.endefr.%d" % 2**15, vocab_size=2**15),
-        lambda: None),  # test set
+         lambda: wsj_parsing.parsing_token_generator(
+             FLAGS.data_dir, FLAGS.tmp_dir, False, 2**14, 2**9),
+         lambda: None),  # test set
+    "inference_snli32k":
+        (lambda: snli.snli_token_generator(FLAGS.tmp_dir, True, 2**15),
+         lambda: snli.snli_token_generator(FLAGS.tmp_dir, False, 2**15),
+         lambda: None),  # test set
+    "audio_timit_characters_test": (lambda: audio.timit_generator(
+        FLAGS.data_dir, FLAGS.tmp_dir, True, 1718
+    ), lambda: audio.timit_generator(FLAGS.data_dir, FLAGS.tmp_dir, False, 626),
+                                    lambda: None),  # test set
+    "audio_timit_tokens_8k_test": (lambda: audio.timit_generator(
+        FLAGS.data_dir,
+        FLAGS.tmp_dir,
+        True,
+        1718,
+        vocab_filename="vocab.endefr.%d" % 2**13,
+        vocab_size=2**13), lambda: audio.timit_generator(
+            FLAGS.data_dir,
+            FLAGS.tmp_dir,
+            False,
+            626,
+            vocab_filename="vocab.endefr.%d" % 2**13,
+            vocab_size=2**13), lambda: None),  # test set
+    "audio_timit_tokens_32k_test": (lambda: audio.timit_generator(
+        FLAGS.data_dir,
+        FLAGS.tmp_dir,
+        True,
+        1718,
+        vocab_filename="vocab.endefr.%d" % 2**15,
+        vocab_size=2**15), lambda: audio.timit_generator(
+            FLAGS.data_dir,
+            FLAGS.tmp_dir,
+            False,
+            626,
+            vocab_filename="vocab.endefr.%d" % 2**15,
+            vocab_size=2**15), lambda: None),  # test set
 }
 
 # pylint: enable=g-long-lambda
@@ -147,7 +163,8 @@ def main(_):
 
   # Calculate the list of problems to generate.
   problems = sorted(
-      list(_SUPPORTED_PROBLEM_GENERATORS) + registry.list_base_problems())
+      list(_SUPPORTED_PROBLEM_GENERATORS) + registry.list_base_problems() +
+      registry.list_env_problems())
   for exclude in FLAGS.exclude_problems.split(","):
     if exclude:
       problems = [p for p in problems if exclude not in p]
@@ -169,8 +186,9 @@ def main(_):
 
   if not problems:
     problems_str = "\n  * ".join(
-        sorted(list(_SUPPORTED_PROBLEM_GENERATORS) +
-               registry.list_base_problems()))
+        sorted(
+            list(_SUPPORTED_PROBLEM_GENERATORS) +
+            registry.list_base_problems() + registry.list_env_problems()))
     error_msg = ("You must specify one of the supported problems to "
                  "generate data for:\n  * " + problems_str + "\n")
     error_msg += ("TIMIT and parsing need data_sets specified with "
@@ -179,15 +197,14 @@ def main(_):
 
   if not FLAGS.data_dir:
     FLAGS.data_dir = tempfile.gettempdir()
-    tf.logging.warning("It is strongly recommended to specify --data_dir. "
-                       "Data will be written to default data_dir=%s.",
-                       FLAGS.data_dir)
+    tf.logging.warning(
+        "It is strongly recommended to specify --data_dir. "
+        "Data will be written to default data_dir=%s.", FLAGS.data_dir)
   FLAGS.data_dir = os.path.expanduser(FLAGS.data_dir)
   tf.gfile.MakeDirs(FLAGS.data_dir)
 
-  tf.logging.info("Generating problems:\n%s"
-                  % registry.display_list_by_prefix(problems,
-                                                    starting_spaces=4))
+  tf.logging.info("Generating problems:\n%s" %
+                  registry.display_list_by_prefix(problems, starting_spaces=4))
   if FLAGS.only_list:
     return
   for problem in problems:
@@ -195,8 +212,13 @@ def main(_):
 
     if problem in _SUPPORTED_PROBLEM_GENERATORS:
       generate_data_for_problem(problem)
-    else:
+    elif problem in registry.list_base_problems():
       generate_data_for_registered_problem(problem)
+    elif problem in registry.list_env_problems():
+      generate_data_for_env_problem(problem)
+    else:
+      tf.logging.error("Problem %s is not a supported problem for datagen.",
+                       problem)
 
 
 def generate_data_for_problem(problem):
@@ -235,6 +257,24 @@ def generate_data_in_process(arg):
   problem.generate_data(data_dir, tmp_dir, task_id)
 
 
+def generate_data_for_env_problem(problem_name):
+  """Generate data for `EnvProblem`s."""
+  assert FLAGS.env_problem_max_env_steps > 0, ("--env_problem_max_env_steps "
+                                               "should be greater than zero")
+  assert FLAGS.env_problem_batch_size > 0, ("--env_problem_batch_size should be"
+                                            " greather than zero")
+  problem = registry.env_problem(problem_name)
+  task_id = None if FLAGS.task_id < 0 else FLAGS.task_id
+  data_dir = os.path.expanduser(FLAGS.data_dir)
+  tmp_dir = os.path.expanduser(FLAGS.tmp_dir)
+  # TODO(msaffar): Handle large values for env_problem_batch_size where we
+  #  cannot create that many environments within the same process.
+  problem.initialize(batch_size=FLAGS.env_problem_batch_size)
+  env_problem_utils.play_env_problem_randomly(
+      problem, num_steps=FLAGS.env_problem_max_env_steps)
+  problem.generate_data(data_dir=data_dir, tmp_dir=tmp_dir, task_id=task_id)
+
+
 def generate_data_for_registered_problem(problem_name):
   """Generate data for a registered problem."""
   tf.logging.info("Generating data for %s.", problem_name)
@@ -260,6 +300,7 @@ def generate_data_for_registered_problem(problem_name):
   else:
     problem.generate_data(data_dir, tmp_dir, task_id)
 
+
 if __name__ == "__main__":
   tf.logging.set_verbosity(tf.logging.INFO)
   tf.app.run()
diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 83bec50d6..006a881be 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -91,6 +91,8 @@
     "tensor2tensor.data_generators.wikitext103",
     "tensor2tensor.data_generators.wsj_parsing",
     "tensor2tensor.data_generators.wnli",
+    "tensor2tensor.envs.mujoco_problems",
+    "tensor2tensor.envs.tic_tac_toe_env_problem",
 ]
 ALL_MODULES = list(MODULES)
 
diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 840eab902..4acb5a343 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -671,14 +671,13 @@ def _generate_time_steps(self, trajectory_list):
         if not processed_reward:
           processed_reward = 0
 
-        if time_step.action:
-          action = gym_spaces_utils.gym_space_encode(self.action_space,
-                                                     time_step.action)
-        else:
+        action = time_step.action
+        if action is None:
           # The last time-step doesn't have action, and this action shouldn't be
           # used, gym's spaces have a `sample` function, so let's just sample an
           # action and use that.
-          action = [self.action_space.sample()]
+          action = self.action_space.sample()
+        action = gym_spaces_utils.gym_space_encode(self.action_space, action)
 
         if six.PY3:
           # py3 complains that, to_example cannot handle np.int64 !

From 83d98cd359049a368d7b83442d12d6e4e2d04d48 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 7 Mar 2019 14:33:37 -0800
Subject: [PATCH 1775/2720] TPU compatibility for data chunking

PiperOrigin-RevId: 237326592
---
 tensor2tensor/utils/data_reader.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index 995c47a31..956d69c93 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -490,17 +490,27 @@ def is_nonzero_chunk(example):
     def split_on_length(example):
       """Split a batch of ditcs on length."""
       x = example["targets"]
+      # TODO(kitaev): This code breaks if chunk_length * max_chunks < batch_size
       length_diff = chunk_length * max_chunks - tf.shape(x)[1]
       padded_x = tf.pad(x, [(0, 0), (0, length_diff), (0, 0), (0, 0)])
       chunks = [padded_x[:, i*chunk_length:(i+1)*chunk_length, :, :]
                 for i in range(max_chunks - 1)]
       chunks.append(padded_x[:, (max_chunks - 1)*chunk_length:, :, :])
       new_example = {}
-      new_example["chunk_number"] = tf.range(max_chunks)
+      # Setting chunk_number to be tf.range(max_chunks) is incompatible with TPU
+      new_example["chunk_number"] = tf.concat([
+          tf.expand_dims(tf.ones_like(c) * n, axis=0)
+          for n, c in enumerate(chunks)
+      ],
+                                              axis=0)
       new_example["targets"] = tf.concat(
           [tf.expand_dims(c, axis=0) for c in chunks], axis=0)
       for k in example:
         if k != "targets":
+          assert k != "chunk_number", (
+              "Chunking code expects the chunk_number feature name to be "
+              "available"
+          )
           new_example[k] = tf.concat(
               [tf.expand_dims(example[k], axis=0) for _ in range(max_chunks)],
               axis=0)

From b6bc169e6ccce0076f523a6aac95acedaa9956a7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 8 Mar 2019 10:52:37 -0800
Subject: [PATCH 1776/2720] Fix gym space encoding to support nested numpy
 structures.

PiperOrigin-RevId: 237479355
---
 tensor2tensor/envs/gym_spaces_utils.py      |  2 +-
 tensor2tensor/envs/gym_spaces_utils_test.py | 59 +++++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)
 create mode 100644 tensor2tensor/envs/gym_spaces_utils_test.py

diff --git a/tensor2tensor/envs/gym_spaces_utils.py b/tensor2tensor/envs/gym_spaces_utils.py
index 1176db0ff..a4ca516ff 100644
--- a/tensor2tensor/envs/gym_spaces_utils.py
+++ b/tensor2tensor/envs/gym_spaces_utils.py
@@ -75,7 +75,7 @@ def gym_space_encode(gym_space, observation):
     return [observation]
 
   if isinstance(gym_space, Box):
-    return list(observation.reshape(-1))
+    return observation.reshape(-1).tolist()
 
   raise NotImplementedError
 
diff --git a/tensor2tensor/envs/gym_spaces_utils_test.py b/tensor2tensor/envs/gym_spaces_utils_test.py
new file mode 100644
index 000000000..af7b76112
--- /dev/null
+++ b/tensor2tensor/envs/gym_spaces_utils_test.py
@@ -0,0 +1,59 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 gym_spaces_utils.py."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from gym.spaces import Box
+from gym.spaces import Discrete
+import numpy as np
+from tensor2tensor.envs import gym_spaces_utils
+import tensorflow as tf
+
+
+class GymSpacesUtilsTest(tf.test.TestCase):
+
+  def test_discrete_space_spec(self):
+    discrete_space = Discrete(100)
+    spec = gym_spaces_utils.gym_space_spec(discrete_space)
+    self.assertIsInstance(spec, tf.FixedLenFeature)
+    self.assertEqual(spec.dtype, tf.int64)
+    self.assertListEqual(list(spec.shape), [1])
+
+  def test_box_space_spec(self):
+    box_space = Box(low=0, high=10, shape=[5, 6], dtype=np.float32)
+    spec = gym_spaces_utils.gym_space_spec(box_space)
+    self.assertIsInstance(spec, tf.FixedLenFeature)
+    self.assertEqual(spec.dtype, tf.float32)
+    self.assertListEqual(list(spec.shape), [5, 6])
+
+  def test_discrete_space_encode(self):
+    discrete_space = Discrete(100)
+    value = discrete_space.sample()
+    encoded_value = gym_spaces_utils.gym_space_encode(discrete_space, value)
+    self.assertListEqual([value], encoded_value)
+
+  def test_box_space_encode(self):
+    box_space = Box(low=0, high=10, shape=[2], dtype=np.int64)
+    value = np.array([2, 3])
+    encoded_value = gym_spaces_utils.gym_space_encode(box_space, value)
+    self.assertListEqual([2, 3], encoded_value)
+
+
+if __name__ == '__main__':
+  tf.test.main()

From 3de03e72df6babd8b272963004a3acf9f44fd7ab Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 8 Mar 2019 11:52:53 -0800
Subject: [PATCH 1777/2720] Fixed "KeyError: 'inputs'" by redirecting 'inputs'
 to 'targets' in estimator_model_fn when use_tpu is used.

PiperOrigin-RevId: 237492137
---
 tensor2tensor/utils/t2t_model.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 5a141f8ae..010f619d1 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1450,7 +1450,9 @@ def estimator_model_fn(cls,
     # PREDICT mode
     if mode == tf.estimator.ModeKeys.PREDICT:
       if use_tpu:
-        inputs = features["inputs"]
+        inputs = features.get("inputs")
+        if inputs is None:
+          inputs = features["targets"]
         shape = inputs.get_shape().as_list()
         if shape[0] is None:
           shape[0] = decode_hparams.batch_size or hparams.batch_size

From 2569b77f11e6867b43adfac9370e5df57c0cb3ce Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 11 Mar 2019 14:02:38 -0700
Subject: [PATCH 1778/2720] Added hparam for top-k random sampling in
 Transformer.

PiperOrigin-RevId: 237877301
---
 tensor2tensor/layers/common_hparams.py |  1 +
 tensor2tensor/layers/common_layers.py  | 20 ++++++++++++++++++--
 tensor2tensor/models/transformer.py    | 10 ++++++++--
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 8d3f23eaf..55b147987 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -105,6 +105,7 @@ def basic_params1():
       learning_rate=0.1,
       sampling_method="argmax",  # "argmax" or "random"
       sampling_temp=1.0,  # temperature for sampling
+      sampling_keep_top_k=-1,  # If >0, ignore all but the top k logits
       # expand the logits a piece at a time - saves memory.
       factored_logits=False,
       multiply_embedding_mode="sqrt_depth",
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 6b583549f..805dfc02a 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2849,13 +2849,13 @@ def list_product(els):
   return prod
 
 
-def sample_with_temperature(logits, temperature):
+def sample_with_temperature(logits, temperature, sampling_keep_top_k=-1):
   """Either argmax or random sampling.
 
   Args:
     logits: a Tensor.
     temperature: a float  0.0=argmax 1.0=random
-
+    sampling_keep_top_k: If not -1, only sample from the top k logits.
   Returns:
     a Tensor with one fewer dimension than logits.
   """
@@ -2866,6 +2866,22 @@ def sample_with_temperature(logits, temperature):
     return tf.reshape(argmax, logits_shape[:-1])
   else:
     assert temperature > 0.0
+
+    if sampling_keep_top_k != -1:
+      if sampling_keep_top_k <= 0:
+        raise ValueError("sampling_keep_top_k must either be -1 or positive.")
+
+      vocab_size = shape_list(logits)[1]
+
+      k_largest = tf.contrib.nn.nth_element(
+          logits, n=sampling_keep_top_k, reverse=True)
+      k_largest = tf.tile(tf.reshape(k_largest, [-1, 1]), [1, vocab_size])
+
+      # Force every position that is not in the top k to have probability near
+      # 0 by setting the logit to be very negative.
+      logits = tf.where(tf.less_equal(logits, k_largest),
+                        tf.ones_like(logits)*-1e3, logits)
+
     reshaped_logits = (
         tf.reshape(logits, [-1, shape_list(logits)[-1]]) / temperature)
     choices = tf.multinomial(reshaped_logits, 1)
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 4fead25d2..f2e4e53fd 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -935,7 +935,10 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       temperature = getattr(hparams, "sampling_temp", 0.0)
       if hparams.sampling_method == "argmax":
         temperature = 0.0
-      next_id = common_layers.sample_with_temperature(logits, temperature)
+      keep_top = hparams.sampling_keep_top_k
+      next_id = common_layers.sample_with_temperature(
+          logits, temperature, keep_top)
+
       hit_eos |= tf.equal(next_id, eos_id)
 
       log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],
@@ -1116,9 +1119,12 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       logits, cache = symbols_to_logits_fn(next_id, i, cache)
       log_probs = common_layers.log_prob_from_logits(logits)
       temperature = getattr(hparams, "sampling_temp", 0.0)
+
       if hparams.sampling_method == "argmax":
         temperature = 0.0
-      next_id = common_layers.sample_with_temperature(logits, temperature)
+      keep_top = hparams.sampling_keep_top_k
+      next_id = common_layers.sample_with_temperature(
+          logits, temperature, keep_top)
       hit_eos |= tf.equal(next_id, eos_id)
 
       log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],

From a1a6404c829f31d46a0402772d9e2ee4b125d430 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 11 Mar 2019 15:42:50 -0700
Subject: [PATCH 1779/2720] Added hparam for top-k random sampling in
 Transformer.

PiperOrigin-RevId: 237899161
---
 tensor2tensor/layers/common_hparams.py |  1 -
 tensor2tensor/layers/common_layers.py  | 20 ++------------------
 tensor2tensor/models/transformer.py    | 10 ++--------
 3 files changed, 4 insertions(+), 27 deletions(-)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 55b147987..8d3f23eaf 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -105,7 +105,6 @@ def basic_params1():
       learning_rate=0.1,
       sampling_method="argmax",  # "argmax" or "random"
       sampling_temp=1.0,  # temperature for sampling
-      sampling_keep_top_k=-1,  # If >0, ignore all but the top k logits
       # expand the logits a piece at a time - saves memory.
       factored_logits=False,
       multiply_embedding_mode="sqrt_depth",
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 805dfc02a..6b583549f 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2849,13 +2849,13 @@ def list_product(els):
   return prod
 
 
-def sample_with_temperature(logits, temperature, sampling_keep_top_k=-1):
+def sample_with_temperature(logits, temperature):
   """Either argmax or random sampling.
 
   Args:
     logits: a Tensor.
     temperature: a float  0.0=argmax 1.0=random
-    sampling_keep_top_k: If not -1, only sample from the top k logits.
+
   Returns:
     a Tensor with one fewer dimension than logits.
   """
@@ -2866,22 +2866,6 @@ def sample_with_temperature(logits, temperature, sampling_keep_top_k=-1):
     return tf.reshape(argmax, logits_shape[:-1])
   else:
     assert temperature > 0.0
-
-    if sampling_keep_top_k != -1:
-      if sampling_keep_top_k <= 0:
-        raise ValueError("sampling_keep_top_k must either be -1 or positive.")
-
-      vocab_size = shape_list(logits)[1]
-
-      k_largest = tf.contrib.nn.nth_element(
-          logits, n=sampling_keep_top_k, reverse=True)
-      k_largest = tf.tile(tf.reshape(k_largest, [-1, 1]), [1, vocab_size])
-
-      # Force every position that is not in the top k to have probability near
-      # 0 by setting the logit to be very negative.
-      logits = tf.where(tf.less_equal(logits, k_largest),
-                        tf.ones_like(logits)*-1e3, logits)
-
     reshaped_logits = (
         tf.reshape(logits, [-1, shape_list(logits)[-1]]) / temperature)
     choices = tf.multinomial(reshaped_logits, 1)
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index f2e4e53fd..4fead25d2 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -935,10 +935,7 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       temperature = getattr(hparams, "sampling_temp", 0.0)
       if hparams.sampling_method == "argmax":
         temperature = 0.0
-      keep_top = hparams.sampling_keep_top_k
-      next_id = common_layers.sample_with_temperature(
-          logits, temperature, keep_top)
-
+      next_id = common_layers.sample_with_temperature(logits, temperature)
       hit_eos |= tf.equal(next_id, eos_id)
 
       log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],
@@ -1119,12 +1116,9 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       logits, cache = symbols_to_logits_fn(next_id, i, cache)
       log_probs = common_layers.log_prob_from_logits(logits)
       temperature = getattr(hparams, "sampling_temp", 0.0)
-
       if hparams.sampling_method == "argmax":
         temperature = 0.0
-      keep_top = hparams.sampling_keep_top_k
-      next_id = common_layers.sample_with_temperature(
-          logits, temperature, keep_top)
+      next_id = common_layers.sample_with_temperature(logits, temperature)
       hit_eos |= tf.equal(next_id, eos_id)
 
       log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],

From b2346e1f1dbbe7d8e1df876222591a6fb32bfe71 Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Mon, 11 Mar 2019 15:53:37 -0700
Subject: [PATCH 1780/2720] local 2d attention that works on TPU.

PiperOrigin-RevId: 237901376
---
 tensor2tensor/layers/common_attention.py      | 289 +++++++++++++++++-
 tensor2tensor/layers/common_attention_test.py | 254 +++++++++++++++
 2 files changed, 536 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 2484422f2..48254978c 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2190,6 +2190,279 @@ def _compute_2d_relative_logits(
         width_key_relative_embeddings)
 
 
+def _split_along_width(x_left_right_blocks):
+  """Helper function for local 2d attention.
+
+  Takes a tensor of [batch, heads, num_h_blocks, num_w_blocks,
+  height, width, depth] and returns two tensors which contain every alternate
+  position along the width
+
+
+  Args:
+    x_left_right_blocks: A [batch, heads, num_h_blocks, num_w_blocks,
+                            height, width, depth] tensor
+
+  Returns:
+    x_left_blocks, x_right_blocks: two [batch, heads, num_h_blocks,
+                                        (num_w_blocks-2)/2, height, width,
+                                        depth] tensors
+
+  """
+  (_, num_heads, x_num_h_blocks, x_num_outer_w_blocks, x_memory_flange_h,
+   x_memory_flange_w, depth) = common_layers.shape_list(x_left_right_blocks)
+  x_num_w_blocks = (x_num_outer_w_blocks-1)//2
+  # get it ready for splitting the left and right memory blocks
+  x_left_right_blocks = tf.reshape(x_left_right_blocks,
+                                   [-1, num_heads,
+                                    x_num_h_blocks,
+                                    x_num_outer_w_blocks//2, 2,
+                                    x_memory_flange_h,
+                                    x_memory_flange_w, depth])
+
+  x_left_blocks, x_right_blocks = tf.split(x_left_right_blocks,
+                                           num_or_size_splits=2, axis=4)
+  x_left_blocks = tf.squeeze(x_left_blocks, axis=4)
+  x_right_blocks = tf.squeeze(x_right_blocks, axis=4)
+  x_left_blocks = tf.slice(x_left_blocks, [0, 0, 0, 0, 0, 0, 0],
+                           [-1, -1, -1, x_num_w_blocks, -1, -1, -1])
+  x_right_blocks = tf.slice(x_right_blocks, [0, 0, 0, 1, 0, 0, 0],
+                            [-1, -1, -1, x_num_w_blocks, -1, -1, -1])
+  return x_left_blocks, x_right_blocks
+
+
+def _get_left_right_blocks(x):
+  """Helper function. Assumes that memory_flange is half of query sizes.
+
+  This function splits the tensor of width 'n' into two halves, where the
+  first half gets the width indices 0, 2, 4.. and the second half gets the
+  width indices 3, 5, ... We also fuse two blocks along the h dimension.
+
+  Args:
+    x: a 7-d tensor.
+
+  Returns:
+    x_left_blocks, x_right_blocks: Two 7-d tensors
+  """
+  (_, num_heads, x_num_outer_h_blocks, x_num_outer_w_blocks, x_memory_flange_h,
+   x_memory_flange_w, depth) = common_layers.shape_list(x)
+  x_left_right_blocks = tf.slice(x,
+                                 [0, 0, 1, 0, 0, 0, 0],
+                                 [-1, -1, x_num_outer_h_blocks-2, -1, -1,
+                                  -1, -1])
+  num_blocks_h = (x_num_outer_h_blocks-2)//2
+  x_left_right_blocks = tf.reshape(x_left_right_blocks,
+                                   [-1, num_heads,
+                                    num_blocks_h,
+                                    2, x_num_outer_w_blocks,
+                                    x_memory_flange_h,
+                                    x_memory_flange_w, depth])
+  x_left_right_blocks = tf.transpose(x_left_right_blocks,
+                                     [0, 1, 2, 4, 3, 5, 6, 7])
+  x_left_right_blocks = tf.reshape(x_left_right_blocks,
+                                   [-1, num_heads, num_blocks_h,
+                                    x_num_outer_w_blocks, 2*x_memory_flange_h,
+                                    x_memory_flange_w, depth])
+  # get it ready for splitting the left and right memory blocks
+  x_left_blocks, x_right_blocks = _split_along_width(x_left_right_blocks)
+
+  return x_left_blocks, x_right_blocks
+  # return x_left_right_blocks
+
+
+def _extract_blocks(x, block_h, block_w):
+  """Helper function for local 2d attention.
+
+  Args:
+    x: a [batch, num_heads, height, width, depth] tensor
+    block_h: An integer. block height
+    block_w: An inteter. block width
+
+  returns:
+    a [batch, num_heads, height/block_h, width/block_w, depth] tensor
+  """
+  (_, num_heads, height, width, depth) = common_layers.shape_list(x)
+  assert height % block_h == 0
+  assert width % block_w == 0
+  x = tf.reshape(x, [-1, num_heads, height//block_h, block_h,
+                     width//block_w, block_w, depth])
+  return tf.transpose(x, [0, 1, 2, 4, 3, 5, 6])
+
+
+def get_2d_local_memory(x, query_shape, memory_flange):
+  """Stitches together the local 2d memory blocks.
+
+  Args:
+    x: a [batch, heads, height, width, depth tensor]
+    query_shape: 2-d integer list of query shape
+    memory_flange: 2-d integer list of memory flanges
+
+  Returns:
+    x: A [batch, heads, num_h_blocks, num_w_blocks,
+          query_shape[0]+2*memory_flange[0],query_shape[1]+2*memory_flange[1]]
+          tensor.
+  """
+  (_, num_heads, height, width, depth_x) = common_layers.shape_list(x)
+  x_center_blocks = _extract_blocks(x, query_shape[0], query_shape[1])
+  # add extra padding to x so that we can extract the memory region
+  # around the center
+  paddings = [[0, 0], [0, 0], [memory_flange[0], memory_flange[0]],
+              [memory_flange[1], memory_flange[1]], [0, 0]]
+  padded_x = tf.pad(x, paddings)
+  padded_x.set_shape([None, num_heads, height+2*memory_flange[0],
+                      width+2*memory_flange[1], depth_x])
+  x_outer_memory_blocks = _extract_blocks(padded_x,
+                                          memory_flange[0], memory_flange[1])
+  # We'll extract left and right memory blocks, top and bottom memory blocks,
+  # and then the corner memory blocks
+
+  # Each of these after  will have shape
+  # [batch, num_heads, num_h_blocks, num_w_blocks, query_shape[0],
+  # memory_flange[1], depth]
+  x_left_blocks, x_right_blocks = _get_left_right_blocks(
+      x_outer_memory_blocks)
+  t_hw_block = lambda x: tf.transpose(x, [0, 1, 3, 2, 5, 4, 6])
+  # now to get top and bottom blocks, we should just transpose the outer
+  # blocks, call the same function and transpose back to get shape
+  # [batch, num_heads, num_h_blocks, num_w_blocks, memory_flange[0],
+  # query_shape[1], depth]
+  x_top_center_blocks, x_bottom_center_blocks = (
+      map(t_hw_block, _get_left_right_blocks(
+          t_hw_block(x_outer_memory_blocks))))
+
+  # now to get the corner blocks
+  x_left_corner_blocks, x_right_corner_blocks = _split_along_width(
+      x_outer_memory_blocks)
+  # now to extract top and bottom for both k and v
+  # we need to transpose because _split_along_width separates along
+  # the width
+  # each of these should have shape [batch, num_heads, num_h_blocks,
+  # num_w_blocks, memory_flange[0], memory_flange[1], depth]
+
+  t_hw = lambda x: tf.transpose(x, [0, 1, 3, 2, 4, 5, 6])
+  x_top_left_corner_blocks, x_bottom_left_corner_blocks = (
+      map(t_hw, _split_along_width(t_hw(x_left_corner_blocks))))
+  x_top_right_corner_blocks, x_bottom_right_corner_blocks = (
+      map(t_hw, _split_along_width(t_hw(x_right_corner_blocks))))
+
+  # The memory is top_left     top_center    top_right
+  #               left_center  middle        right_center
+  #               bottom_left  bottom_center bottom_right
+  # Assembling the above row by row
+  # first [x_top_left, x_top, x_top_right]
+  # to get [batch, num_heads, num_h_blocks, num_w_blocks, memory_flange[0],
+  # query_shape[1]+2*memory_flange[1], depth]
+  # then [x_left, x_center, x_right]
+  # then [x_bottom_left, x_bottom, x_bottom_right]
+  x_top_memory = tf.concat(
+      [x_top_left_corner_blocks,
+       x_top_center_blocks,
+       x_top_right_corner_blocks], axis=5)
+  x_middle_memory = tf.concat(
+      [x_left_blocks, x_center_blocks, x_right_blocks], axis=5)
+  x_bottom_memory = tf.concat(
+      [x_bottom_left_corner_blocks,
+       x_bottom_center_blocks,
+       x_bottom_right_corner_blocks], axis=5)
+
+  # concat along height
+  x = tf.concat([x_top_memory, x_middle_memory, x_bottom_memory], axis=4)
+  return x
+
+
+def dot_product_unmasked_attention_local_2d_tpu(
+    q, k, v, bias, max_relative_position=None, query_shape=(8, 8),
+    dropout_rate=0.0, image_shapes=None, name=None, make_image_summary=False,
+    dropout_broadcast_dims=None):
+  """Calculate unmasked dot-product local self-attention 2d on tpu.
+
+  Args:
+    q: a Tensor with shape [batch, heads, height, width, depth].
+    k: a Tensor with shape [batch, heads, height, width, depth].
+    v: a Tensor with shape [batch, heads, height, width, depth].
+    bias: bias Tensor.
+    max_relative_position: an integer the max relative embedding considered.
+      Changing this invalidates checkpoints.
+    query_shape: a two tuple indicating query shape
+    dropout_rate: a floating point number.
+    image_shapes: optional tuple of integer scalars.
+    name: an optional string.
+    make_image_summary: Whether to make an attention image summary.
+    dropout_broadcast_dims:  an optional list of integers less than 4
+      specifying in which dimensions to broadcast the dropout decisions.
+      saves memory.
+
+  Returns:
+    [batch, heads, height, width, depth] tensor, the output of attention.
+    height_key_relative_embeddings: a 3d or 2d tensor, depending on head sharing
+      settings, which are the relative embeddings for height.
+    width_key_relative_embeddings: a 3d or 2d tensor, depending on head sharing
+      settings, which are the relative embeddings for width.
+
+  """
+  if max_relative_position:
+    raise ValueError("Relative local 2d attention not implemented")
+
+  with tf.variable_scope(
+      name,
+      default_name="dot_product_self_attention_relative_v2",
+      values=[q, k, v]):
+
+    # This calculation only works for self attention.
+    # q, k and v must therefore have the same shape.
+    q.get_shape().assert_is_compatible_with(k.get_shape())
+    q.get_shape().assert_is_compatible_with(v.get_shape())
+    q_shape = common_layers.shape_list(q)
+    (height, width) = (q_shape[2],
+                       q_shape[3])
+    _, num_heads, height, width, depth_k = common_layers.shape_list(k)
+    depth_v = common_layers.shape_list(v)[-1]
+    num_h_blocks = height//query_shape[0]
+    num_w_blocks = width//query_shape[1]
+    # Pad query, key, value to ensure multiple of corresponding lengths.
+    memory_flange = [int(query_shape[0]//2), int(query_shape[1]//2)]
+    q = pad_to_multiple_2d(q, query_shape)
+    k = pad_to_multiple_2d(k, query_shape)
+    v = pad_to_multiple_2d(v, query_shape)
+
+    # Extract center queries, keys, and values
+
+    queries = _extract_blocks(
+        q, query_shape[0], query_shape[1])
+    keys = get_2d_local_memory(
+        k, query_shape, memory_flange)
+    values = get_2d_local_memory(
+        v, query_shape, memory_flange)
+    memory_h = query_shape[0] + 2*memory_flange[0]
+    memory_w = query_shape[1] + 2*memory_flange[1]
+    queries = tf.reshape(queries, [-1, num_heads, num_h_blocks, num_w_blocks,
+                                   query_shape[0]*query_shape[1], depth_k])
+    keys = tf.reshape(keys, [-1, num_heads, num_h_blocks, num_w_blocks,
+                             memory_h*memory_w, depth_k])
+    values = tf.reshape(values, [-1, num_heads, num_h_blocks, num_w_blocks,
+                                 memory_h*memory_w, depth_v])
+    logits = tf.matmul(queries, keys, transpose_b=True)
+    if bias is not None:
+      logits += bias
+
+    weights = tf.nn.softmax(logits, name="attention_weights")
+    # Dropping out the attention links for each of the heads
+    weights = common_layers.dropout_with_broadcast_dims(
+        weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
+    if common_layers.should_generate_summaries() and make_image_summary:
+      attention_image_summary(weights, image_shapes)
+    ret = tf.matmul(weights, values)
+    # we need to get it back to shape [batch, heads, height, width]
+    ret = tf.reshape(ret, [-1, num_heads, num_h_blocks, num_w_blocks,
+                           query_shape[0], query_shape[1], depth_v])
+    ret = tf.transpose(ret, [0, 1, 2, 4, 3, 5, 6])
+    ret = tf.reshape(ret, [-1, num_heads, num_h_blocks*query_shape[0],
+                           num_w_blocks*query_shape[1], depth_v])
+    # slice if padding was introduced
+    ret = tf.slice(ret, [0, 0, 0, 0, 0], [-1, -1, q_shape[2], q_shape[3],
+                                          -1])
+    return ret
+
+
 def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
   """Attention to the source and a neighborhood to the left within a block.
 
@@ -3189,12 +3462,10 @@ def make_2d_block_raster_mask(query_shape, memory_flange):
   split_query_masks = tf.split(query_triangle, query_shape[0], axis=1)
   # adding mask for left and right
   mask_pieces = [
-      tf.concat(
-          [
-              tf.ones([np.prod(query_shape), memory_flange[1]]),
-              split_query_masks[i],
-              tf.zeros([np.prod(query_shape), memory_flange[1]])
-          ],
+      tf.concat(  # pylint: disable=g-complex-comprehension
+          [tf.ones([np.prod(query_shape), memory_flange[1]]),
+           split_query_masks[i],
+           tf.zeros([np.prod(query_shape), memory_flange[1]])],
           axis=1) for i in range(query_shape[0])
   ]
   # adding mask for top
@@ -3854,10 +4125,14 @@ def multihead_attention_2d(query_antecedent,
     if attention_type == "local_attention_2d":
       x = local_attention_2d(
           q, k, v, query_shape=query_shape, memory_flange=memory_flange)
-    else:
+    elif attention_type == "masked_local_attention_2d":
       assert attention_type == "masked_local_attention_2d"
       x = masked_local_attention_2d(
           q, k, v, query_shape=query_shape, memory_flange=memory_flange)
+    else:
+      assert attention_type == "unmasked_local_attention_2d_tpu"
+      x = dot_product_unmasked_attention_local_2d_tpu(
+          q, k, v, None, max_relative_position=None, query_shape=query_shape)
     x = combine_heads_2d(x)
     x = common_layers.dense(
         x, output_depth, use_bias=False, name="output_transform")
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index ec1b6f2e9..17780faaa 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -541,6 +541,260 @@ def testDotProductUnMaskedAttentionRelativeV2(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testExtractblocks(self):
+
+    batch_size = 1
+    num_heads = 3
+    height = 6
+    width = 10
+    depth = 15
+    block_h = 3
+    block_w = 2
+    t = np.random.rand(batch_size, num_heads, height, width, depth)
+    a = common_attention._extract_blocks(t, block_h, block_w)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
+    self.assertEqual(res.shape, (batch_size, num_heads, height//block_h,
+                                 width//block_w, block_h, block_w, depth))
+    # also check if the content is right
+    out = np.zeros((batch_size, num_heads, height//block_h,
+                    width//block_w, block_h, block_w, depth))
+    for b in range(batch_size):
+      for h in range(num_heads):
+        for x in range(height//block_h):
+          for y in range(width//block_w):
+            for v in range(block_h):
+              for w in range(block_w):
+                out[b, h, x, y, v, w] = t[b, h, block_h*x+v, block_w*y+w]
+    self.assertAllClose(res, out)
+
+  def python_get_2d_local_memory(self, t, batch_size, num_heads, height, width,
+                                 num_h_blocks, num_w_blocks, query_shape,
+                                 memory_flange, depth):
+    # also check if the content is right
+    out = np.zeros((batch_size, num_heads, height//query_shape[0],
+                    width//query_shape[1], query_shape[0]+2*memory_flange[0],
+                    query_shape[1]+2*memory_flange[1], depth))
+    memory_height = query_shape[0]+2*memory_flange[0]
+    memory_width = query_shape[1]+2*memory_flange[1]
+    t_padded = np.pad(t, ((0, 0), (0, 0), (memory_flange[0], memory_flange[0]),
+                          (memory_flange[1], memory_flange[1]), (0, 0)),
+                      "constant",
+                      constant_values=((0, 0), (0, 0), (0, 0), (0, 0), (0, 0)))
+    for b in range(batch_size):
+      for h in range(num_heads):
+        for x in range(num_h_blocks):
+          for y in range(num_w_blocks):
+            for v in range(memory_height):
+              for w in range(memory_width):
+                memory_h_start = x*query_shape[0]
+                memory_w_start = y*query_shape[1]
+                memory_h_index = memory_h_start + v
+                memory_w_index = memory_w_start + w
+                out[b, h, x, y, v, w] = t_padded[b, h, memory_h_index,
+                                                 memory_w_index]
+    return out
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testGet2dLocalMemory(self):
+    batch_size = 3
+    num_heads = 3
+    height = 12
+    width = 6
+    depth = 15
+    num_h_blocks = 3
+    num_w_blocks = 3
+    memory_flange = [2, 1]
+    query_shape = [4, 2]
+    t = np.random.rand(batch_size, num_heads, height, width, depth)
+    a = common_attention.get_2d_local_memory(
+        t, query_shape, memory_flange)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
+    self.assertEqual(res.shape, (batch_size, num_heads,
+                                 num_h_blocks,
+                                 num_w_blocks,
+                                 query_shape[0]+2*memory_flange[0],
+                                 query_shape[1]+2*memory_flange[1], depth))
+    out = self.python_get_2d_local_memory(t, batch_size, num_heads,
+                                          height, width, num_h_blocks,
+                                          num_w_blocks, query_shape,
+                                          memory_flange, depth)
+
+    self.assertAllClose(res, out)
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testSplitAlongWidth(self):
+    batch_size = 1
+    num_heads = 3
+    num_outer_h_blocks = 4
+    num_outer_w_blocks = 8
+    memory_flange = [2, 2]
+    num_w_blocks = 3
+    depth = 15
+    t = np.random.rand(batch_size, num_heads, num_outer_h_blocks,
+                       num_outer_w_blocks, memory_flange[0], memory_flange[1],
+                       depth)
+    a = common_attention._split_along_width(t)
+    # self.evaluate(tf.global_variables_initializer())
+    res_l, res_r = self.evaluate(a)
+    # res = self.evaluate(a)
+    self.assertEqual(res_l.shape, (batch_size, num_heads, num_outer_h_blocks,
+                                   num_w_blocks, memory_flange[0],
+                                   memory_flange[1], depth))
+    self.assertEqual(res_r.shape, (batch_size, num_heads, num_outer_h_blocks,
+                                   num_w_blocks, memory_flange[0],
+                                   memory_flange[1], depth))
+    # also check if the content is right
+    out_l = np.zeros((batch_size, num_heads, num_outer_h_blocks, num_w_blocks,
+                      memory_flange[0], memory_flange[1], depth))
+    out_r = np.zeros((batch_size, num_heads, num_outer_h_blocks, num_w_blocks,
+                      memory_flange[0], memory_flange[1], depth))
+    block_h = memory_flange[0]
+    block_w = memory_flange[1]
+    for b in range(batch_size):
+      for h in range(num_heads):
+        for x in range(num_outer_h_blocks):
+          for y in range(num_w_blocks):
+            for v in range(block_h):
+              for w in range(block_w):
+                # we should compute the index of the position in the
+                out_l[b, h, x, y, v, w] = (
+                    t[b, h, x, 2*y, v, w]
+                    )
+                out_r[b, h, x, y, v, w] = (
+                    t[b, h, x, 2*y+3, v, w]
+                    )
+    self.assertAllClose(res_l, out_l)
+    self.assertAllClose(res_r, out_r)
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testGetLeftRightBlocks(self):
+    batch_size = 1
+    num_heads = 3
+    num_outer_h_blocks = 6
+    num_outer_w_blocks = 6
+    memory_flange = [2, 2]
+    num_h_blocks = 2
+    num_w_blocks = 2
+    depth = 15
+    t = np.random.rand(batch_size, num_heads, num_outer_h_blocks,
+                       num_outer_w_blocks, memory_flange[0], memory_flange[1],
+                       depth)
+    a = common_attention._get_left_right_blocks(t)
+    self.evaluate(tf.global_variables_initializer())
+    res_l, res_r = self.evaluate(a)
+    self.assertEqual(res_l.shape, (batch_size, num_heads, num_h_blocks,
+                                   num_w_blocks, memory_flange[0]*2,
+                                   memory_flange[1], depth))
+    self.assertEqual(res_r.shape, (batch_size, num_heads, num_h_blocks,
+                                   num_w_blocks, memory_flange[0]*2,
+                                   memory_flange[1], depth))
+    # also check if the content is right
+    block_h = memory_flange[0]*2
+    block_w = memory_flange[1]
+    out_l = np.zeros((batch_size, num_heads, num_h_blocks,
+                      num_w_blocks, memory_flange[0]*2, memory_flange[1],
+                      depth))
+    out_r = np.zeros((batch_size, num_heads, num_h_blocks,
+                      num_w_blocks, memory_flange[0]*2, memory_flange[1],
+                      depth))
+    block_h = memory_flange[0]*2
+    block_w = memory_flange[1]
+    for b in range(batch_size):
+      for h in range(num_heads):
+        for x in range(num_h_blocks):
+          for y in range(num_w_blocks):
+            for v in range(block_h):
+              for w in range(block_w):
+                # we should compute the index of the position in the
+                outer_block_h_index = (
+                    1 + block_h//memory_flange[0]*x + v//2)
+                h_index = v%memory_flange[0]
+                left_outer_w_index = 2*y
+                right_outer_w_index = 2*y + 3
+                out_l[b, h, x, y, v, w] = (
+                    t[b, h, outer_block_h_index, left_outer_w_index, h_index,
+                      w]
+                    )
+                out_r[b, h, x, y, v, w] = (
+                    t[b, h, outer_block_h_index, right_outer_w_index, h_index,
+                      w]
+                    )
+    self.assertAllClose(res_l, out_l)
+    self.assertAllClose(res_r, out_r)
+
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testDotProductUnmaskedAttentionLocal2dTpu(self):
+    batch_size = 1
+    num_heads = 3
+    height = 4
+    width = 12
+    depth = 15
+    num_h_blocks = 2
+    num_w_blocks = 2
+    memory_flange = [1, 3]
+    query_shape = [2, 6]
+    memory_h = query_shape[0] + 2*memory_flange[0]
+    memory_w = query_shape[1] + 2*memory_flange[1]
+
+    q = np.random.rand(batch_size, num_heads, height, width, depth)
+    k = np.random.rand(batch_size, num_heads, height, width, depth)
+    v = np.random.rand(batch_size, num_heads, height, width, depth)
+    a = common_attention.dot_product_unmasked_attention_local_2d_tpu(
+        tf.constant(q, dtype=tf.float32),
+        tf.constant(k, dtype=tf.float32),
+        tf.constant(v, dtype=tf.float32), None, max_relative_position=None,
+        query_shape=query_shape, dropout_rate=0.0, image_shapes=None,
+        name=None, make_image_summary=False, dropout_broadcast_dims=None)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(a)
+    self.assertEqual(res.shape, (batch_size, num_heads,
+                                 height, width, depth))
+    # now to check the content too
+    queries = self.python_get_2d_local_memory(q, batch_size, num_heads,
+                                              height, width, num_h_blocks,
+                                              num_w_blocks, query_shape, [0, 0],
+                                              depth)
+    keys = self.python_get_2d_local_memory(k, batch_size, num_heads,
+                                           height, width, num_h_blocks,
+                                           num_w_blocks, query_shape,
+                                           memory_flange, depth)
+    values = self.python_get_2d_local_memory(v, batch_size, num_heads,
+                                             height, width, num_h_blocks,
+                                             num_w_blocks, query_shape,
+                                             memory_flange, depth)
+    logits = np.matmul(
+        np.reshape(queries, (batch_size, num_heads,
+                             num_h_blocks, num_w_blocks,
+                             query_shape[0]*query_shape[1], depth)),
+        np.transpose(
+            np.reshape(keys, (batch_size, num_heads, num_h_blocks, num_w_blocks,
+                              memory_h*memory_w, depth)), (0, 1, 2, 3, 5, 4)))
+    # now to do a softmax across the logits
+    att = np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True)
+    att_output = np.matmul(att, np.reshape(
+        values, (batch_size, num_heads, num_h_blocks, num_w_blocks,
+                 memory_h*memory_w, depth)))
+    att_output = np.reshape(att_output,
+                            (batch_size, num_heads, num_h_blocks, num_w_blocks,
+                             query_shape[0], query_shape[1], depth))
+    # putting the attention results back into the right place
+    out = np.zeros((batch_size, num_heads, height, width, depth))
+    for b in range(batch_size):
+      for h in range(num_heads):
+        for x in range(height):
+          for y in range(width):
+            h_block_index = x//query_shape[0]
+            w_block_index = y//query_shape[1]
+            inside_h_index = x%query_shape[0]
+            inside_w_index = y%query_shape[1]
+            out[b, h, x, y] = (
+                att_output[b, h, h_block_index, w_block_index, inside_h_index,
+                           inside_w_index])
+    self.assertAllClose(res, out)
+
   def python_relative_att(self, q, k, v, batch, num_heads, height, width,
                           depth, height_key_relative_embeddings,
                           width_key_relative_embeddings,

From c9b8032ed40b693fd5b8a81771ca875b83c55f3e Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Mon, 11 Mar 2019 17:32:54 -0700
Subject: [PATCH 1781/2720] Fixed bug in 2d local TPU.

PiperOrigin-RevId: 237920758
---
 tensor2tensor/layers/common_attention.py      | 16 +++----
 tensor2tensor/layers/common_attention_test.py | 43 +++++++++++++------
 2 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 48254978c..ab51af221 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2411,6 +2411,12 @@ def dot_product_unmasked_attention_local_2d_tpu(
     # q, k and v must therefore have the same shape.
     q.get_shape().assert_is_compatible_with(k.get_shape())
     q.get_shape().assert_is_compatible_with(v.get_shape())
+    orig_q_shape = common_layers.shape_list(q)
+    # Pad query, key, value to ensure multiple of corresponding lengths.
+    memory_flange = [int(query_shape[0]//2), int(query_shape[1]//2)]
+    q = pad_to_multiple_2d(q, query_shape)
+    k = pad_to_multiple_2d(k, query_shape)
+    v = pad_to_multiple_2d(v, query_shape)
     q_shape = common_layers.shape_list(q)
     (height, width) = (q_shape[2],
                        q_shape[3])
@@ -2418,12 +2424,6 @@ def dot_product_unmasked_attention_local_2d_tpu(
     depth_v = common_layers.shape_list(v)[-1]
     num_h_blocks = height//query_shape[0]
     num_w_blocks = width//query_shape[1]
-    # Pad query, key, value to ensure multiple of corresponding lengths.
-    memory_flange = [int(query_shape[0]//2), int(query_shape[1]//2)]
-    q = pad_to_multiple_2d(q, query_shape)
-    k = pad_to_multiple_2d(k, query_shape)
-    v = pad_to_multiple_2d(v, query_shape)
-
     # Extract center queries, keys, and values
 
     queries = _extract_blocks(
@@ -2458,8 +2458,8 @@ def dot_product_unmasked_attention_local_2d_tpu(
     ret = tf.reshape(ret, [-1, num_heads, num_h_blocks*query_shape[0],
                            num_w_blocks*query_shape[1], depth_v])
     # slice if padding was introduced
-    ret = tf.slice(ret, [0, 0, 0, 0, 0], [-1, -1, q_shape[2], q_shape[3],
-                                          -1])
+    ret = tf.slice(ret, [0, 0, 0, 0, 0], [-1, -1, orig_q_shape[2],
+                                          orig_q_shape[3], -1])
     return ret
 
 
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 17780faaa..210909035 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -729,13 +729,13 @@ def testGetLeftRightBlocks(self):
   def testDotProductUnmaskedAttentionLocal2dTpu(self):
     batch_size = 1
     num_heads = 3
-    height = 4
+    height = 7
     width = 12
     depth = 15
-    num_h_blocks = 2
-    num_w_blocks = 2
-    memory_flange = [1, 3]
-    query_shape = [2, 6]
+    num_h_blocks = 4
+    num_w_blocks = 6
+    memory_flange = [1, 1]
+    query_shape = [2, 2]
     memory_h = query_shape[0] + 2*memory_flange[0]
     memory_w = query_shape[1] + 2*memory_flange[1]
 
@@ -753,17 +753,33 @@ def testDotProductUnmaskedAttentionLocal2dTpu(self):
     self.assertEqual(res.shape, (batch_size, num_heads,
                                  height, width, depth))
     # now to check the content too
+    # first pad q, k, ad v
+    height_padding = -height % query_shape[0]
+    width_padding = -width % query_shape[1]
+    new_height = height + -height % query_shape[0]
+    new_width = width + -width % query_shape[1]
+    q = np.pad(q, ((0, 0), (0, 0), (0, height_padding),
+                   (0, width_padding), (0, 0)), "constant",
+               constant_values=((0, 0), (0, 0), (0, 0), (0, 0), (0, 0)))
+    k = np.pad(k, ((0, 0), (0, 0), (0, height_padding),
+                   (0, width_padding), (0, 0)), "constant",
+               constant_values=((0, 0), (0, 0), (0, 0), (0, 0), (0, 0)))
+    v = np.pad(v, ((0, 0), (0, 0), (0, height_padding),
+                   (0, width_padding), (0, 0)), "constant",
+               constant_values=((0, 0), (0, 0), (0, 0), (0, 0), (0, 0)))
     queries = self.python_get_2d_local_memory(q, batch_size, num_heads,
-                                              height, width, num_h_blocks,
-                                              num_w_blocks, query_shape, [0, 0],
+                                              new_height, new_width,
+                                              num_h_blocks, num_w_blocks,
+                                              query_shape, [0, 0],
                                               depth)
     keys = self.python_get_2d_local_memory(k, batch_size, num_heads,
-                                           height, width, num_h_blocks,
+                                           new_height, new_width, num_h_blocks,
                                            num_w_blocks, query_shape,
                                            memory_flange, depth)
     values = self.python_get_2d_local_memory(v, batch_size, num_heads,
-                                             height, width, num_h_blocks,
-                                             num_w_blocks, query_shape,
+                                             new_height, new_width,
+                                             num_h_blocks, num_w_blocks,
+                                             query_shape,
                                              memory_flange, depth)
     logits = np.matmul(
         np.reshape(queries, (batch_size, num_heads,
@@ -781,11 +797,11 @@ def testDotProductUnmaskedAttentionLocal2dTpu(self):
                             (batch_size, num_heads, num_h_blocks, num_w_blocks,
                              query_shape[0], query_shape[1], depth))
     # putting the attention results back into the right place
-    out = np.zeros((batch_size, num_heads, height, width, depth))
+    out = np.zeros((batch_size, num_heads, new_height, new_width, depth))
     for b in range(batch_size):
       for h in range(num_heads):
-        for x in range(height):
-          for y in range(width):
+        for x in range(new_height):
+          for y in range(new_width):
             h_block_index = x//query_shape[0]
             w_block_index = y//query_shape[1]
             inside_h_index = x%query_shape[0]
@@ -793,6 +809,7 @@ def testDotProductUnmaskedAttentionLocal2dTpu(self):
             out[b, h, x, y] = (
                 att_output[b, h, h_block_index, w_block_index, inside_h_index,
                            inside_w_index])
+    out = out[:, :, :height, :width, :]
     self.assertAllClose(res, out)
 
   def python_relative_att(self, q, k, v, batch, num_heads, height, width,

From 42a41209aa0e1001bc37a70c3d78d0daa4e47370 Mon Sep 17 00:00:00 2001
From: Zi Yang <ziy@google.com>
Date: Mon, 11 Mar 2019 21:47:10 -0700
Subject: [PATCH 1782/2720] Updated Tensor2Tensor's serving_input_fn to allow
 correctly pad features and prepare the batch size at exporting time for use
 at the serving time, if TPU is used.

PiperOrigin-RevId: 237947564
---
 tensor2tensor/data_generators/problem.py      | 21 ++++++++--
 tensor2tensor/data_generators/problem_test.py | 40 +++++++++++++++++++
 tensor2tensor/serving/export.py               |  9 +++--
 tensor2tensor/utils/data_reader.py            | 10 ++---
 tensor2tensor/utils/trainer_lib.py            |  5 ++-
 5 files changed, 72 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 8329dc69c..d83ef5bbe 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -19,6 +19,7 @@
 from __future__ import print_function
 
 import collections
+import functools
 import os
 import random
 import six
@@ -891,7 +892,7 @@ def export_assets(self):
 
     return None
 
-  def serving_input_fn(self, hparams):
+  def serving_input_fn(self, hparams, decode_hparams=None, use_tpu=False):
     """Input fn for serving export, starting from serialized example."""
     mode = tf.estimator.ModeKeys.PREDICT
     serialized_example = tf.placeholder(
@@ -900,9 +901,21 @@ def serving_input_fn(self, hparams):
     dataset = dataset.map(self.decode_example)
     dataset = dataset.map(lambda ex: self.preprocess_example(ex, mode, hparams))
     dataset = dataset.map(data_reader.cast_ints_to_int32)
-    dataset = dataset.padded_batch(
-        tf.shape(serialized_example, out_type=tf.int64)[0],
-        dataset.output_shapes)
+
+    if use_tpu:
+      padded_shapes = data_reader.pad_for_tpu(dataset.output_shapes, hparams,
+                                              hparams.max_length)
+      batch_size = 1 if not decode_hparams else getattr(decode_hparams,
+                                                        "batch_size", 1)
+      dataset = dataset.padded_batch(
+          batch_size, padded_shapes, drop_remainder=False)
+      dataset = dataset.map(
+          functools.partial(data_reader.pad_batch, batch_multiple=batch_size))
+    else:
+      dataset = dataset.padded_batch(
+          tf.shape(serialized_example, out_type=tf.int64)[0],
+          dataset.output_shapes)
+
     dataset = dataset.map(data_reader.standardize_shapes)
     features = tf.data.experimental.get_single_element(dataset)
 
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 03e42a193..7bf0137a9 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.data_generators import problem as problem_module
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.layers import modalities
+from tensor2tensor.utils import hparam
 from tensor2tensor.utils import test_utils
 
 import tensorflow as tf
@@ -149,6 +150,45 @@ def testDataFilenames(self):
         problem.data_filepaths(problem_module.DatasetSplit.TEST, data_dir,
                                num_shards, shuffled))
 
+  @test_utils.run_in_graph_mode_only()
+  def testServingInputFnUseTpu(self):
+    problem = problem_module.Problem()
+    max_length = 128
+    batch_size = 10
+    hparams = hparam.HParams(
+        max_length=max_length,
+        max_input_seq_length=max_length,
+        max_target_seq_length=max_length,
+        prepend_mode="none",
+        split_to_length=0)
+    decode_hparams = hparam.HParams(batch_size=batch_size)
+    serving_input_receiver = problem.serving_input_fn(
+        hparams=hparams, decode_hparams=decode_hparams, use_tpu=True)
+    serving_input_fn_input = getattr(serving_input_receiver,
+                                     "receiver_tensors")["input"]
+    serving_input_fn_output = getattr(serving_input_receiver,
+                                      "features")["inputs"]
+    example_1 = tf.train.Example(
+        features=tf.train.Features(feature={
+            "inputs": tf.train.Feature(
+                int64_list=tf.train.Int64List(value=[0]))
+        }))
+    example_2 = tf.train.Example(
+        features=tf.train.Features(feature={
+            "inputs": tf.train.Feature(
+                int64_list=tf.train.Int64List(value=[1]))
+        }))
+    serialized_examples = [
+        example_1.SerializeToString(),
+        example_2.SerializeToString()
+    ]
+    with self.test_session() as sess:
+      output_shape = sess.run(
+          tf.shape(serving_input_fn_output),
+          feed_dict={serving_input_fn_input: serialized_examples})
+      self.assertEqual(output_shape[0], batch_size)
+      self.assertEqual(output_shape[1], max_length)
+
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 1e17c56f4..ff582db38 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -128,7 +128,8 @@ def hub_module_fn():
         hparams,
         decode_hparams=decode_hparams,
         use_tpu=FLAGS.use_tpu)
-    features = problem.serving_input_fn(hparams).features
+    features = problem.serving_input_fn(
+        hparams, decode_hparams, use_tpu=FLAGS.use_tpu).features
 
     # we must do a copy of the features, as the model_fn can add additional
     # entries there (like hyperparameter settings etc).
@@ -168,12 +169,12 @@ def main(_):
   hparams = create_hparams()
   hparams.no_data_parallelism = True  # To clear the devices
   problem = hparams.problem
+  decode_hparams = decoding.decode_hparams(FLAGS.decode_hparams)
 
   export_dir = FLAGS.export_dir or os.path.join(ckpt_dir, "export")
 
   if FLAGS.export_as_tfhub:
     checkpoint_path = tf.train.latest_checkpoint(ckpt_dir)
-    decode_hparams = decoding.decode_hparams(FLAGS.decode_hparams)
     export_as_tfhub_module(FLAGS.model, hparams, decode_hparams, problem,
                            checkpoint_path, export_dir)
     return
@@ -183,7 +184,9 @@ def main(_):
   estimator = create_estimator(run_config, hparams)
 
   exporter = tf.estimator.FinalExporter(
-      "exporter", lambda: problem.serving_input_fn(hparams), as_text=True)
+      "exporter",
+      lambda: problem.serving_input_fn(hparams, decode_hparams, FLAGS.use_tpu),
+      as_text=True)
 
   exporter.export(
       estimator,
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index 956d69c93..c623ea47b 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -191,7 +191,7 @@ def create_threads(self, sess, coord=None, daemon=False, start=False):
     return []
 
 
-def _pad_for_tpu(shapes_dict, hparams, max_length):
+def pad_for_tpu(shapes_dict, hparams, max_length):
   """Pads unknown features' dimensions for TPU."""
   padded_shapes = {}
 
@@ -289,7 +289,7 @@ def skip_random_fraction(dataset, data_file):
   return dataset.skip(num_skip)
 
 
-def _pad_batch(features, batch_multiple):
+def pad_batch(features, batch_multiple):
   """Pad batch dim of features to nearest multiple of batch_multiple."""
   feature = list(features.items())[0][1]
   batch_size = tf.shape(feature)[0]
@@ -414,7 +414,7 @@ def define_shapes(example):
     # batch_size means tokens per datashard
     if config and config.use_tpu:
       dataset = dataset.filter(tpu_valid_size)
-      padded_shapes = _pad_for_tpu(dataset.output_shapes, hparams, max_length)
+      padded_shapes = pad_for_tpu(dataset.output_shapes, hparams, max_length)
       # on TPU, we use params["batch_size"], which specifies the number of
       # examples across all datashards
       batch_size = params["batch_size"]
@@ -427,7 +427,7 @@ def define_shapes(example):
         dataset = dataset.padded_batch(
             batch_size, padded_shapes, drop_remainder=False)
         dataset = dataset.map(
-            functools.partial(_pad_batch, batch_multiple=batch_size),
+            functools.partial(pad_batch, batch_multiple=batch_size),
             num_parallel_calls=num_threads)
       else:
         dataset = dataset.padded_batch(
@@ -460,7 +460,7 @@ def define_shapes(example):
               "lead to incorrect metrics for non-zero-padded features, e.g. "
               "images. Use a single datashard (i.e. 1 GPU) in that case.")
           dataset = dataset.map(
-              functools.partial(_pad_batch, batch_multiple=batch_multiple),
+              functools.partial(pad_batch, batch_multiple=batch_multiple),
               num_parallel_calls=num_threads)
 
   dataset = dataset.map(define_shapes, num_parallel_calls=num_threads)
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index e8af60126..ec8606c82 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -686,9 +686,12 @@ def compare_fn(best_eval_result, current_eval_result):
       metric = eval_early_stopping_metric or "loss"
       return current_eval_result[metric] < best_eval_result[metric]
 
+    def serving_input_receiver_fn(hparams, decode_hparams, use_tpu):
+      return problem.serving_input_fn(hparams, decode_hparams, use_tpu)
+
     exporter = tf.estimator.BestExporter(
         name="best",
-        serving_input_receiver_fn=lambda: problem.serving_input_fn(hparams),
+        serving_input_receiver_fn=serving_input_receiver_fn,
         compare_fn=compare_fn,
         assets_extra=problem.export_assets)
 

From d0a8ed02aedbcbc81cb20ab85582699942cca688 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 12 Mar 2019 11:25:41 -0700
Subject: [PATCH 1783/2720] Add RenderedEnvProblem.

This is the base class for any Gym environment problem with rgb array as observations to behave as a VideoProblem.

PiperOrigin-RevId: 238058750
---
 oss_scripts/oss_tests.sh                      |   1 +
 tensor2tensor/envs/env_problem.py             |  20 ++-
 tensor2tensor/envs/env_problem_test.py        |   4 +
 tensor2tensor/envs/mujoco_problems.py         |   8 +-
 tensor2tensor/envs/rendered_env_problem.py    | 132 ++++++++++++++++++
 .../envs/rendered_env_problem_test.py         |  70 ++++++++++
 6 files changed, 226 insertions(+), 9 deletions(-)
 create mode 100644 tensor2tensor/envs/rendered_env_problem.py
 create mode 100644 tensor2tensor/envs/rendered_env_problem_test.py

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index ea31b21ce..4d104a643 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -51,6 +51,7 @@ pytest --disable-warnings \
   --ignore=tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py \
   --ignore=tensor2tensor/data_generators/problem_test.py \
   --ignore=tensor2tensor/envs/mujoco_problems_test.py \
+  --ignore=tensor2tensor/envs/rendered_env_problem_test.py \
   --ignore=tensor2tensor/layers/bayes_test.py \
   --ignore=tensor2tensor/layers/common_attention_test.py \
   --ignore=tensor2tensor/layers/common_layers_test.py \
diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 4acb5a343..0cddfa74e 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -60,6 +60,7 @@ class EnvProblem(Env, problem.Problem):
   - target_modality
   - target_vocab_size
   - action_modality
+  - reward_modality
 
   NON NATIVELY BATCHED ENVS:
 
@@ -401,6 +402,10 @@ def num_rewards(self):
   def input_modality(self):
     raise NotImplementedError
 
+  @property
+  def reward_modality(self):
+    raise NotImplementedError
+
   @property
   def input_vocab_size(self):
     raise NotImplementedError
@@ -581,10 +586,10 @@ def example_reading_spec(self):
         ACTION_FIELD: self.action_spec,
     }
 
-    # `data_items_to_decoders` can be None, it will be set to the appropriate
-    # decoder dict in `Problem.decode_example`
-    # TODO(afrozm): Verify that we don't need any special decoder or anything.
-    data_items_to_decoders = None
+    data_items_to_decoders = {
+        field: tf.contrib.slim.tfexample_decoder.Tensor(field)
+        for field in data_fields
+    }
 
     return data_fields, data_items_to_decoders
 
@@ -603,8 +608,8 @@ def hparams(self, defaults, model_hparams):
     p.modality.update({
         "inputs": self.input_modality,
         "targets": self.target_modality,
-        "input_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
-        "target_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        "input_reward": self.reward_modality,
+        "target_reward": self.reward_modality,
         "input_action": self.action_modality,
         "target_action": self.action_modality,
         "target_policy": modalities.ModalityType.IDENTITY,
@@ -695,7 +700,8 @@ def _generate_time_steps(self, trajectory_list):
 
         yield {
             TIMESTEP_FIELD: [index],
-            ACTION_FIELD: action,
+            ACTION_FIELD:
+                action,
             # to_example errors on np.float32
             RAW_REWARD_FIELD: [float(raw_reward)],
             PROCESSED_REWARD_FIELD: [processed_reward],
diff --git a/tensor2tensor/envs/env_problem_test.py b/tensor2tensor/envs/env_problem_test.py
index e1b8f3e89..ba2fc3678 100644
--- a/tensor2tensor/envs/env_problem_test.py
+++ b/tensor2tensor/envs/env_problem_test.py
@@ -295,6 +295,10 @@ def target_vocab_size(self):
       def action_modality(self):
         return modalities.ModalityType.SYMBOL_WEIGHTS_ALL
 
+      @property
+      def reward_modality(self):
+        return modalities.ModalityType.SYMBOL_WEIGHTS_ALL
+
     base_env_name = "CartPole-v0"
     batch_size = 5
     reward_range = (-1, 1)
diff --git a/tensor2tensor/envs/mujoco_problems.py b/tensor2tensor/envs/mujoco_problems.py
index a172aa69f..60f36009e 100644
--- a/tensor2tensor/envs/mujoco_problems.py
+++ b/tensor2tensor/envs/mujoco_problems.py
@@ -20,7 +20,7 @@
 from __future__ import print_function
 
 import functools
-from tensor2tensor.envs import env_problem
+from tensor2tensor.envs import rendered_env_problem
 from tensor2tensor.layers import modalities
 from tensor2tensor.rl import gym_utils
 from tensor2tensor.utils import registry
@@ -28,7 +28,7 @@
 
 
 @registry.register_env_problem
-class ReacherEnvProblem(env_problem.EnvProblem):
+class ReacherEnvProblem(rendered_env_problem.RenderedEnvProblem):
   """Mujoco's reacher environment."""
 
   def __init__(self):
@@ -54,6 +54,10 @@ def target_modality(self):
   def action_modality(self):
     return modalities.ModalityType.IDENTITY
 
+  @property
+  def reward_modality(self):
+    return modalities.ModalityType.IDENTITY
+
   @property
   def input_vocab_size(self):
     return 256
diff --git a/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
new file mode 100644
index 000000000..5760ae47a
--- /dev/null
+++ b/tensor2tensor/envs/rendered_env_problem.py
@@ -0,0 +1,132 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 class for env problems with RGB array as observation space."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import png
+import six
+from tensor2tensor.data_generators import video_utils
+from tensor2tensor.envs import env_problem
+import tensorflow as tf
+
+_IMAGE_ENCODED_FIELD = "image/encoded"
+_IMAGE_FORMAT_FIELD = "image/format"
+_IMAGE_HEIGHT_FIELD = "image/height"
+_IMAGE_WIDTH_FIELD = "image/width"
+_FRAME_NUMBER_FIELD = "frame_number"
+
+_FORMAT = "png"
+
+
+class RenderedEnvProblem(env_problem.EnvProblem, video_utils.VideoProblem):
+  """An `EnvProblem` when observations are RGB arrays.
+
+  This takes care of wrapping a rendered gym environment to behave like a
+  `VideoProblem`. This class assumes the underlying gym environment is either a
+  `gym_utils.RenderedEnv` or it natively returns rendered scene for
+  observations. i.e. the underlying gym environment should have a
+  `Box` observation space with the following shape: [frame_height, frame_width,
+  channels]
+
+  Note: The method resolution order for this class is:
+  `RenderedEnvProblem`, `EnvProblem`, `Env`, `VideoProblem`, `Problem`
+  """
+
+  def __init__(self,
+               base_env_name=None,
+               batch_size=None,
+               env_wrapper_fn=None,
+               reward_range=(-np.inf, np.inf)):
+    """Initialize by calling both parents' constructors."""
+    env_problem.EnvProblem.__init__(self, base_env_name, batch_size,
+                                    env_wrapper_fn, reward_range)
+    video_utils.VideoProblem.__init__(self)
+
+  def initialize_environments(self, batch_size=1):
+    env_problem.EnvProblem.initialize_environments(self, batch_size)
+    # Assert the underlying gym environment has correct observation space
+    assert len(self.observation_spec.shape) == 3
+
+  def example_reading_spec(self):
+    """Return a mix of env and video data fields and decoders."""
+    video_fields, video_decoders = (
+        video_utils.VideoProblem.example_reading_spec(self))
+    env_fields, env_decoders = env_problem.EnvProblem.example_reading_spec(self)
+
+    # Remove raw observations field since we want to capture them as videos.
+    env_fields.pop(env_problem.OBSERVATION_FIELD)
+    env_decoders.pop(env_problem.OBSERVATION_FIELD)
+
+    # Add frame number spec and decoder.
+    env_fields[_FRAME_NUMBER_FIELD] = tf.FixedLenFeature((1,), tf.int64)
+    env_decoders[
+        _FRAME_NUMBER_FIELD] = tf.contrib.slim.tfexample_decoder.Tensor(
+            _FRAME_NUMBER_FIELD)
+
+    # Add video fields and decoders
+    env_fields.update(video_fields)
+    env_decoders.update(video_decoders)
+    return env_fields, env_decoders
+
+  def _generate_time_steps(self, trajectory_list):
+    """Transforms time step observations to frames of a video."""
+    for time_step in env_problem.EnvProblem._generate_time_steps(
+        self, trajectory_list):
+      # Convert the rendered observations from numpy to png format.
+      frame_np = np.array(time_step.pop(env_problem.OBSERVATION_FIELD))
+      frame_np = frame_np.reshape(
+          [self.frame_height, self.frame_width, self.num_channels])
+      # TODO(msaffar) Add support for non RGB rendered environments
+      frame = png.from_array(frame_np, "RGB", info={"bitdepth": 8})
+      frame_buffer = six.BytesIO()
+      frame.save(frame_buffer)
+
+      # Put the encoded frame back.
+      time_step[_IMAGE_ENCODED_FIELD] = [frame_buffer.getvalue()]
+      time_step[_IMAGE_FORMAT_FIELD] = [_FORMAT]
+      time_step[_IMAGE_HEIGHT_FIELD] = [self.frame_height]
+      time_step[_IMAGE_WIDTH_FIELD] = [self.frame_width]
+
+      # Add the frame number
+      time_step[_FRAME_NUMBER_FIELD] = time_step[env_problem.TIMESTEP_FIELD]
+      yield time_step
+
+  @property
+  def num_channels(self):
+    return self.observation_spec.shape[2]
+
+  @property
+  def frame_height(self):
+    return self.observation_spec.shape[0]
+
+  @property
+  def frame_width(self):
+    return self.observation_spec.shape[1]
+
+  @property
+  def total_number_of_frames(self):
+    """Upper bound on the total number of frames across all environments.
+
+    This is used to decide sharding. See `VideoProblem.total_number_of_frames`
+    for more details.
+
+    Returns:
+      number of frames among all examples in the dataset.
+    """
+    return self.trajectories.num_time_steps
diff --git a/tensor2tensor/envs/rendered_env_problem_test.py b/tensor2tensor/envs/rendered_env_problem_test.py
new file mode 100644
index 000000000..d52dce618
--- /dev/null
+++ b/tensor2tensor/envs/rendered_env_problem_test.py
@@ -0,0 +1,70 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.envs.rendered_env_problem."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.envs import env_problem
+from tensor2tensor.envs import env_problem_utils
+from tensor2tensor.envs import rendered_env_problem
+from tensor2tensor.envs.mujoco_problems import ReacherEnvProblem
+import tensorflow as tf
+
+
+class RenderedEnvProblemTest(tf.test.TestCase):
+
+  def test_generate_timesteps(self):
+    env = ReacherEnvProblem()
+    env.initialize(batch_size=2)
+    env_problem_utils.play_env_problem_randomly(env, num_steps=5)
+    env.trajectories.complete_all_trajectories()
+
+    frame_number = 0
+    for time_step in env._generate_time_steps(
+        env.trajectories.completed_trajectories):
+      # original observation should not be in time_step
+      self.assertNotIn(env_problem.OBSERVATION_FIELD, time_step)
+      # validate frame
+      self.assertIn(rendered_env_problem._IMAGE_ENCODED_FIELD, time_step)
+      self.assertIn(rendered_env_problem._IMAGE_HEIGHT_FIELD, time_step)
+      self.assertIn(rendered_env_problem._IMAGE_WIDTH_FIELD, time_step)
+      self.assertIn(rendered_env_problem._IMAGE_FORMAT_FIELD, time_step)
+      self.assertIn(rendered_env_problem._FRAME_NUMBER_FIELD, time_step)
+
+      decoded_frame = tf.image.decode_png(
+          time_step[rendered_env_problem._IMAGE_ENCODED_FIELD][0])
+
+      decoded_frame = self.evaluate(decoded_frame)
+
+      self.assertListEqual(
+          [env.frame_height, env.frame_width, env.num_channels],
+          list(decoded_frame.shape))
+      self.assertListEqual([rendered_env_problem._FORMAT],
+                           time_step[rendered_env_problem._IMAGE_FORMAT_FIELD])
+      self.assertListEqual([frame_number],
+                           time_step[rendered_env_problem._FRAME_NUMBER_FIELD])
+      self.assertListEqual([env.frame_width],
+                           time_step[rendered_env_problem._IMAGE_WIDTH_FIELD])
+      self.assertListEqual([env.frame_height],
+                           time_step[rendered_env_problem._IMAGE_HEIGHT_FIELD])
+      frame_number += 1
+      frame_number %= 6
+
+
+if __name__ == "__main__":
+  tf.test.main()

From b4394423063029836522744fdf3ccc59014cb207 Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Tue, 12 Mar 2019 15:05:59 -0700
Subject: [PATCH 1784/2720] Fold head and batch because tf does not like > 6
 dimensions for gradients. Adjusted tests accordingly.

PiperOrigin-RevId: 238105879
---
 tensor2tensor/layers/common_attention.py      |  82 ++++++-------
 tensor2tensor/layers/common_attention_test.py | 113 +++++++++---------
 2 files changed, 100 insertions(+), 95 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index ab51af221..a90306c87 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2199,34 +2199,34 @@ def _split_along_width(x_left_right_blocks):
 
 
   Args:
-    x_left_right_blocks: A [batch, heads, num_h_blocks, num_w_blocks,
+    x_left_right_blocks: A [batch, num_h_blocks, num_w_blocks,
                             height, width, depth] tensor
 
   Returns:
-    x_left_blocks, x_right_blocks: two [batch, heads, num_h_blocks,
+    x_left_blocks, x_right_blocks: two [batch, num_h_blocks,
                                         (num_w_blocks-2)/2, height, width,
                                         depth] tensors
 
   """
-  (_, num_heads, x_num_h_blocks, x_num_outer_w_blocks, x_memory_flange_h,
+  (_, x_num_h_blocks, x_num_outer_w_blocks, x_memory_flange_h,
    x_memory_flange_w, depth) = common_layers.shape_list(x_left_right_blocks)
   x_num_w_blocks = (x_num_outer_w_blocks-1)//2
   # get it ready for splitting the left and right memory blocks
   x_left_right_blocks = tf.reshape(x_left_right_blocks,
-                                   [-1, num_heads,
+                                   [-1,
                                     x_num_h_blocks,
                                     x_num_outer_w_blocks//2, 2,
                                     x_memory_flange_h,
                                     x_memory_flange_w, depth])
 
   x_left_blocks, x_right_blocks = tf.split(x_left_right_blocks,
-                                           num_or_size_splits=2, axis=4)
-  x_left_blocks = tf.squeeze(x_left_blocks, axis=4)
-  x_right_blocks = tf.squeeze(x_right_blocks, axis=4)
-  x_left_blocks = tf.slice(x_left_blocks, [0, 0, 0, 0, 0, 0, 0],
-                           [-1, -1, -1, x_num_w_blocks, -1, -1, -1])
-  x_right_blocks = tf.slice(x_right_blocks, [0, 0, 0, 1, 0, 0, 0],
-                            [-1, -1, -1, x_num_w_blocks, -1, -1, -1])
+                                           num_or_size_splits=2, axis=3)
+  x_left_blocks = tf.squeeze(x_left_blocks, axis=3)
+  x_right_blocks = tf.squeeze(x_right_blocks, axis=3)
+  x_left_blocks = tf.slice(x_left_blocks, [0, 0, 0, 0, 0, 0],
+                           [-1, -1, x_num_w_blocks, -1, -1, -1])
+  x_right_blocks = tf.slice(x_right_blocks, [0, 0, 1, 0, 0, 0],
+                            [-1, -1, x_num_w_blocks, -1, -1, -1])
   return x_left_blocks, x_right_blocks
 
 
@@ -2238,28 +2238,28 @@ def _get_left_right_blocks(x):
   width indices 3, 5, ... We also fuse two blocks along the h dimension.
 
   Args:
-    x: a 7-d tensor.
+    x: a 6-d tensor.
 
   Returns:
-    x_left_blocks, x_right_blocks: Two 7-d tensors
+    x_left_blocks, x_right_blocks: Two 6-d tensors
   """
-  (_, num_heads, x_num_outer_h_blocks, x_num_outer_w_blocks, x_memory_flange_h,
+  (_, x_num_outer_h_blocks, x_num_outer_w_blocks, x_memory_flange_h,
    x_memory_flange_w, depth) = common_layers.shape_list(x)
   x_left_right_blocks = tf.slice(x,
-                                 [0, 0, 1, 0, 0, 0, 0],
-                                 [-1, -1, x_num_outer_h_blocks-2, -1, -1,
+                                 [0, 1, 0, 0, 0, 0],
+                                 [-1, x_num_outer_h_blocks-2, -1, -1,
                                   -1, -1])
   num_blocks_h = (x_num_outer_h_blocks-2)//2
   x_left_right_blocks = tf.reshape(x_left_right_blocks,
-                                   [-1, num_heads,
+                                   [-1,
                                     num_blocks_h,
                                     2, x_num_outer_w_blocks,
                                     x_memory_flange_h,
                                     x_memory_flange_w, depth])
   x_left_right_blocks = tf.transpose(x_left_right_blocks,
-                                     [0, 1, 2, 4, 3, 5, 6, 7])
+                                     [0, 1, 3, 2, 4, 5, 6])
   x_left_right_blocks = tf.reshape(x_left_right_blocks,
-                                   [-1, num_heads, num_blocks_h,
+                                   [-1, num_blocks_h,
                                     x_num_outer_w_blocks, 2*x_memory_flange_h,
                                     x_memory_flange_w, depth])
   # get it ready for splitting the left and right memory blocks
@@ -2273,42 +2273,42 @@ def _extract_blocks(x, block_h, block_w):
   """Helper function for local 2d attention.
 
   Args:
-    x: a [batch, num_heads, height, width, depth] tensor
+    x: a [batch, height, width, depth] tensor
     block_h: An integer. block height
     block_w: An inteter. block width
 
   returns:
     a [batch, num_heads, height/block_h, width/block_w, depth] tensor
   """
-  (_, num_heads, height, width, depth) = common_layers.shape_list(x)
+  (_, height, width, depth) = common_layers.shape_list(x)
   assert height % block_h == 0
   assert width % block_w == 0
-  x = tf.reshape(x, [-1, num_heads, height//block_h, block_h,
+  x = tf.reshape(x, [-1, height//block_h, block_h,
                      width//block_w, block_w, depth])
-  return tf.transpose(x, [0, 1, 2, 4, 3, 5, 6])
+  return tf.transpose(x, [0, 1, 3, 2, 4, 5])
 
 
 def get_2d_local_memory(x, query_shape, memory_flange):
   """Stitches together the local 2d memory blocks.
 
   Args:
-    x: a [batch, heads, height, width, depth tensor]
+    x: a [batch, height, width, depth tensor]
     query_shape: 2-d integer list of query shape
     memory_flange: 2-d integer list of memory flanges
 
   Returns:
-    x: A [batch, heads, num_h_blocks, num_w_blocks,
+    x: A [batch, num_h_blocks, num_w_blocks,
           query_shape[0]+2*memory_flange[0],query_shape[1]+2*memory_flange[1]]
           tensor.
   """
-  (_, num_heads, height, width, depth_x) = common_layers.shape_list(x)
+  (_, height, width, depth_x) = common_layers.shape_list(x)
   x_center_blocks = _extract_blocks(x, query_shape[0], query_shape[1])
   # add extra padding to x so that we can extract the memory region
   # around the center
-  paddings = [[0, 0], [0, 0], [memory_flange[0], memory_flange[0]],
+  paddings = [[0, 0], [memory_flange[0], memory_flange[0]],
               [memory_flange[1], memory_flange[1]], [0, 0]]
   padded_x = tf.pad(x, paddings)
-  padded_x.set_shape([None, num_heads, height+2*memory_flange[0],
+  padded_x.set_shape([None, height+2*memory_flange[0],
                       width+2*memory_flange[1], depth_x])
   x_outer_memory_blocks = _extract_blocks(padded_x,
                                           memory_flange[0], memory_flange[1])
@@ -2316,14 +2316,14 @@ def get_2d_local_memory(x, query_shape, memory_flange):
   # and then the corner memory blocks
 
   # Each of these after  will have shape
-  # [batch, num_heads, num_h_blocks, num_w_blocks, query_shape[0],
+  # [batch, num_h_blocks, num_w_blocks, query_shape[0],
   # memory_flange[1], depth]
   x_left_blocks, x_right_blocks = _get_left_right_blocks(
       x_outer_memory_blocks)
-  t_hw_block = lambda x: tf.transpose(x, [0, 1, 3, 2, 5, 4, 6])
+  t_hw_block = lambda x: tf.transpose(x, [0, 2, 1, 4, 3, 5])
   # now to get top and bottom blocks, we should just transpose the outer
   # blocks, call the same function and transpose back to get shape
-  # [batch, num_heads, num_h_blocks, num_w_blocks, memory_flange[0],
+  # [batch, num_h_blocks, num_w_blocks, memory_flange[0],
   # query_shape[1], depth]
   x_top_center_blocks, x_bottom_center_blocks = (
       map(t_hw_block, _get_left_right_blocks(
@@ -2335,10 +2335,10 @@ def get_2d_local_memory(x, query_shape, memory_flange):
   # now to extract top and bottom for both k and v
   # we need to transpose because _split_along_width separates along
   # the width
-  # each of these should have shape [batch, num_heads, num_h_blocks,
+  # each of these should have shape [batch, num_h_blocks,
   # num_w_blocks, memory_flange[0], memory_flange[1], depth]
 
-  t_hw = lambda x: tf.transpose(x, [0, 1, 3, 2, 4, 5, 6])
+  t_hw = lambda x: tf.transpose(x, [0, 2, 1, 3, 4, 5])
   x_top_left_corner_blocks, x_bottom_left_corner_blocks = (
       map(t_hw, _split_along_width(t_hw(x_left_corner_blocks))))
   x_top_right_corner_blocks, x_bottom_right_corner_blocks = (
@@ -2349,23 +2349,23 @@ def get_2d_local_memory(x, query_shape, memory_flange):
   #               bottom_left  bottom_center bottom_right
   # Assembling the above row by row
   # first [x_top_left, x_top, x_top_right]
-  # to get [batch, num_heads, num_h_blocks, num_w_blocks, memory_flange[0],
+  # to get [batch, num_h_blocks, num_w_blocks, memory_flange[0],
   # query_shape[1]+2*memory_flange[1], depth]
   # then [x_left, x_center, x_right]
   # then [x_bottom_left, x_bottom, x_bottom_right]
   x_top_memory = tf.concat(
       [x_top_left_corner_blocks,
        x_top_center_blocks,
-       x_top_right_corner_blocks], axis=5)
+       x_top_right_corner_blocks], axis=4)
   x_middle_memory = tf.concat(
-      [x_left_blocks, x_center_blocks, x_right_blocks], axis=5)
+      [x_left_blocks, x_center_blocks, x_right_blocks], axis=4)
   x_bottom_memory = tf.concat(
       [x_bottom_left_corner_blocks,
        x_bottom_center_blocks,
-       x_bottom_right_corner_blocks], axis=5)
+       x_bottom_right_corner_blocks], axis=4)
 
   # concat along height
-  x = tf.concat([x_top_memory, x_middle_memory, x_bottom_memory], axis=4)
+  x = tf.concat([x_top_memory, x_middle_memory, x_bottom_memory], axis=3)
   return x
 
 
@@ -2404,7 +2404,7 @@ def dot_product_unmasked_attention_local_2d_tpu(
 
   with tf.variable_scope(
       name,
-      default_name="dot_product_self_attention_relative_v2",
+      default_name="dot_product_unmasked_attention_local_2d_tpu",
       values=[q, k, v]):
 
     # This calculation only works for self attention.
@@ -2425,11 +2425,13 @@ def dot_product_unmasked_attention_local_2d_tpu(
     num_h_blocks = height//query_shape[0]
     num_w_blocks = width//query_shape[1]
     # Extract center queries, keys, and values
-
+    q = tf.reshape(q, [-1, height, width, depth_k])
     queries = _extract_blocks(
         q, query_shape[0], query_shape[1])
+    k = tf.reshape(k, [-1, height, width, depth_k])
     keys = get_2d_local_memory(
         k, query_shape, memory_flange)
+    v = tf.reshape(v, [-1, height, width, depth_v])
     values = get_2d_local_memory(
         v, query_shape, memory_flange)
     memory_h = query_shape[0] + 2*memory_flange[0]
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 210909035..e54a941c1 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -551,22 +551,21 @@ def testExtractblocks(self):
     depth = 15
     block_h = 3
     block_w = 2
-    t = np.random.rand(batch_size, num_heads, height, width, depth)
+    t = np.random.rand(batch_size * num_heads, height, width, depth)
     a = common_attention._extract_blocks(t, block_h, block_w)
     self.evaluate(tf.global_variables_initializer())
     res = self.evaluate(a)
-    self.assertEqual(res.shape, (batch_size, num_heads, height//block_h,
+    self.assertEqual(res.shape, (batch_size * num_heads, height//block_h,
                                  width//block_w, block_h, block_w, depth))
     # also check if the content is right
-    out = np.zeros((batch_size, num_heads, height//block_h,
+    out = np.zeros((batch_size*num_heads, height//block_h,
                     width//block_w, block_h, block_w, depth))
-    for b in range(batch_size):
-      for h in range(num_heads):
-        for x in range(height//block_h):
-          for y in range(width//block_w):
-            for v in range(block_h):
-              for w in range(block_w):
-                out[b, h, x, y, v, w] = t[b, h, block_h*x+v, block_w*y+w]
+    for b in range(batch_size*num_heads):
+      for x in range(height//block_h):
+        for y in range(width//block_w):
+          for v in range(block_h):
+            for w in range(block_w):
+              out[b, x, y, v, w] = t[b, block_h*x+v, block_w*y+w]
     self.assertAllClose(res, out)
 
   def python_get_2d_local_memory(self, t, batch_size, num_heads, height, width,
@@ -609,10 +608,11 @@ def testGet2dLocalMemory(self):
     query_shape = [4, 2]
     t = np.random.rand(batch_size, num_heads, height, width, depth)
     a = common_attention.get_2d_local_memory(
-        t, query_shape, memory_flange)
+        np.reshape(t, (batch_size*num_heads, height, width, depth)),
+        query_shape, memory_flange)
     self.evaluate(tf.global_variables_initializer())
     res = self.evaluate(a)
-    self.assertEqual(res.shape, (batch_size, num_heads,
+    self.assertEqual(res.shape, (batch_size*num_heads,
                                  num_h_blocks,
                                  num_w_blocks,
                                  query_shape[0]+2*memory_flange[0],
@@ -621,6 +621,11 @@ def testGet2dLocalMemory(self):
                                           height, width, num_h_blocks,
                                           num_w_blocks, query_shape,
                                           memory_flange, depth)
+    out = np.reshape(out, (batch_size*num_heads,
+                           num_h_blocks,
+                           num_w_blocks,
+                           query_shape[0]+2*memory_flange[0],
+                           query_shape[1]+2*memory_flange[1], depth))
 
     self.assertAllClose(res, out)
 
@@ -633,39 +638,38 @@ def testSplitAlongWidth(self):
     memory_flange = [2, 2]
     num_w_blocks = 3
     depth = 15
-    t = np.random.rand(batch_size, num_heads, num_outer_h_blocks,
+    t = np.random.rand(batch_size*num_heads, num_outer_h_blocks,
                        num_outer_w_blocks, memory_flange[0], memory_flange[1],
                        depth)
     a = common_attention._split_along_width(t)
     # self.evaluate(tf.global_variables_initializer())
     res_l, res_r = self.evaluate(a)
     # res = self.evaluate(a)
-    self.assertEqual(res_l.shape, (batch_size, num_heads, num_outer_h_blocks,
+    self.assertEqual(res_l.shape, (batch_size*num_heads, num_outer_h_blocks,
                                    num_w_blocks, memory_flange[0],
                                    memory_flange[1], depth))
-    self.assertEqual(res_r.shape, (batch_size, num_heads, num_outer_h_blocks,
+    self.assertEqual(res_r.shape, (batch_size*num_heads, num_outer_h_blocks,
                                    num_w_blocks, memory_flange[0],
                                    memory_flange[1], depth))
     # also check if the content is right
-    out_l = np.zeros((batch_size, num_heads, num_outer_h_blocks, num_w_blocks,
+    out_l = np.zeros((batch_size*num_heads, num_outer_h_blocks, num_w_blocks,
                       memory_flange[0], memory_flange[1], depth))
-    out_r = np.zeros((batch_size, num_heads, num_outer_h_blocks, num_w_blocks,
+    out_r = np.zeros((batch_size*num_heads, num_outer_h_blocks, num_w_blocks,
                       memory_flange[0], memory_flange[1], depth))
     block_h = memory_flange[0]
     block_w = memory_flange[1]
-    for b in range(batch_size):
-      for h in range(num_heads):
-        for x in range(num_outer_h_blocks):
-          for y in range(num_w_blocks):
-            for v in range(block_h):
-              for w in range(block_w):
-                # we should compute the index of the position in the
-                out_l[b, h, x, y, v, w] = (
-                    t[b, h, x, 2*y, v, w]
-                    )
-                out_r[b, h, x, y, v, w] = (
-                    t[b, h, x, 2*y+3, v, w]
-                    )
+    for b in range(batch_size*num_heads):
+      for x in range(num_outer_h_blocks):
+        for y in range(num_w_blocks):
+          for v in range(block_h):
+            for w in range(block_w):
+              # we should compute the index of the position in the
+              out_l[b, x, y, v, w] = (
+                  t[b, x, 2*y, v, w]
+                  )
+              out_r[b, x, y, v, w] = (
+                  t[b, x, 2*y+3, v, w]
+                  )
     self.assertAllClose(res_l, out_l)
     self.assertAllClose(res_r, out_r)
 
@@ -679,49 +683,48 @@ def testGetLeftRightBlocks(self):
     num_h_blocks = 2
     num_w_blocks = 2
     depth = 15
-    t = np.random.rand(batch_size, num_heads, num_outer_h_blocks,
+    t = np.random.rand(batch_size*num_heads, num_outer_h_blocks,
                        num_outer_w_blocks, memory_flange[0], memory_flange[1],
                        depth)
     a = common_attention._get_left_right_blocks(t)
     self.evaluate(tf.global_variables_initializer())
     res_l, res_r = self.evaluate(a)
-    self.assertEqual(res_l.shape, (batch_size, num_heads, num_h_blocks,
+    self.assertEqual(res_l.shape, (batch_size*num_heads, num_h_blocks,
                                    num_w_blocks, memory_flange[0]*2,
                                    memory_flange[1], depth))
-    self.assertEqual(res_r.shape, (batch_size, num_heads, num_h_blocks,
+    self.assertEqual(res_r.shape, (batch_size*num_heads, num_h_blocks,
                                    num_w_blocks, memory_flange[0]*2,
                                    memory_flange[1], depth))
     # also check if the content is right
     block_h = memory_flange[0]*2
     block_w = memory_flange[1]
-    out_l = np.zeros((batch_size, num_heads, num_h_blocks,
+    out_l = np.zeros((batch_size*num_heads, num_h_blocks,
                       num_w_blocks, memory_flange[0]*2, memory_flange[1],
                       depth))
-    out_r = np.zeros((batch_size, num_heads, num_h_blocks,
+    out_r = np.zeros((batch_size*num_heads, num_h_blocks,
                       num_w_blocks, memory_flange[0]*2, memory_flange[1],
                       depth))
     block_h = memory_flange[0]*2
     block_w = memory_flange[1]
-    for b in range(batch_size):
-      for h in range(num_heads):
-        for x in range(num_h_blocks):
-          for y in range(num_w_blocks):
-            for v in range(block_h):
-              for w in range(block_w):
-                # we should compute the index of the position in the
-                outer_block_h_index = (
-                    1 + block_h//memory_flange[0]*x + v//2)
-                h_index = v%memory_flange[0]
-                left_outer_w_index = 2*y
-                right_outer_w_index = 2*y + 3
-                out_l[b, h, x, y, v, w] = (
-                    t[b, h, outer_block_h_index, left_outer_w_index, h_index,
-                      w]
-                    )
-                out_r[b, h, x, y, v, w] = (
-                    t[b, h, outer_block_h_index, right_outer_w_index, h_index,
-                      w]
-                    )
+    for b in range(batch_size*num_heads):
+      for x in range(num_h_blocks):
+        for y in range(num_w_blocks):
+          for v in range(block_h):
+            for w in range(block_w):
+              # we should compute the index of the position in the
+              outer_block_h_index = (
+                  1 + block_h//memory_flange[0]*x + v//2)
+              h_index = v%memory_flange[0]
+              left_outer_w_index = 2*y
+              right_outer_w_index = 2*y + 3
+              out_l[b, x, y, v, w] = (
+                  t[b, outer_block_h_index, left_outer_w_index, h_index,
+                    w]
+                  )
+              out_r[b, x, y, v, w] = (
+                  t[b, outer_block_h_index, right_outer_w_index, h_index,
+                    w]
+                  )
     self.assertAllClose(res_l, out_l)
     self.assertAllClose(res_r, out_r)
 

From cb1f6a95d210dbc26903cf2f7a7b786c86a40084 Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Wed, 13 Mar 2019 10:27:56 -0700
Subject: [PATCH 1785/2720] Simpler local 2d tpu attention. Algorithm developed
 by avaswani@ and nikip@. Will only work if memory flange is half of query
 size.

PiperOrigin-RevId: 238252818
---
 tensor2tensor/layers/common_attention.py      | 46 ++++++++++++++++++-
 tensor2tensor/layers/common_attention_test.py |  8 ++--
 2 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index a90306c87..c2992fc47 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2369,6 +2369,48 @@ def get_2d_local_memory(x, query_shape, memory_flange):
   return x
 
 
+def get_2d_local_memory_v2(x, query_shape, memory_flange):
+  """Transposeless local 2d memory block construction.
+
+    Only works if memory flanges are half of query sizes.
+
+  Args:
+    x: a [batch, height, width, depth tensor]
+    query_shape: 2-d integer list of query shape
+    memory_flange: 2-d integer list of memory flanges
+
+  Returns:
+    x: A [batch, num_h_blocks, num_w_blocks,
+          query_shape[0]+2*memory_flange[0],query_shape[1]+2*memory_flange[1]]
+          tensor.
+  """
+  (_, height, width, depth_x) = common_layers.shape_list(x)
+  # add extra padding to x so that we can extract the memory region
+  # around the center
+  paddings = [[0, 0], [memory_flange[0], memory_flange[0]],
+              [memory_flange[1], memory_flange[1]], [0, 0]]
+  padded_x = tf.pad(x, paddings)
+  padded_x.set_shape([None, height+2*memory_flange[0],
+                      width+2*memory_flange[1], depth_x])
+  num_h_memory_blocks = height//query_shape[0] + 1
+  num_w_memory_blocks = width//query_shape[1] + 1
+  x_memory_blocks = _extract_blocks(padded_x,
+                                    query_shape[0], query_shape[1])
+  x_left_width = tf.slice(x_memory_blocks, [0, 0, 0, 0, 0, 0],
+                          [-1, -1, num_w_memory_blocks - 1, -1, -1, -1])
+  x_right_width = tf.slice(x_memory_blocks, [0, 0, 1, 0, 0, 0],
+                           [-1, -1, -1, - 1, -1, -1])
+  x_memory_blocks = tf.concat([x_left_width, x_right_width], axis=4)
+
+  x_top_height = tf.slice(x_memory_blocks, [0, 0, 0, 0, 0, 0],
+                          [-1, num_h_memory_blocks-1, -1, - 1, -1, -1])
+  x_bottom_height = tf.slice(x_memory_blocks, [0, 1, 0, 0, 0, 0],
+                             [-1, -1, -1, - 1, -1, -1])
+  x = tf.concat([x_top_height, x_bottom_height], axis=3)
+
+  return x
+
+
 def dot_product_unmasked_attention_local_2d_tpu(
     q, k, v, bias, max_relative_position=None, query_shape=(8, 8),
     dropout_rate=0.0, image_shapes=None, name=None, make_image_summary=False,
@@ -2429,10 +2471,10 @@ def dot_product_unmasked_attention_local_2d_tpu(
     queries = _extract_blocks(
         q, query_shape[0], query_shape[1])
     k = tf.reshape(k, [-1, height, width, depth_k])
-    keys = get_2d_local_memory(
+    keys = get_2d_local_memory_v2(
         k, query_shape, memory_flange)
     v = tf.reshape(v, [-1, height, width, depth_v])
-    values = get_2d_local_memory(
+    values = get_2d_local_memory_v2(
         v, query_shape, memory_flange)
     memory_h = query_shape[0] + 2*memory_flange[0]
     memory_w = query_shape[1] + 2*memory_flange[1]
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index e54a941c1..ae688576c 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -599,15 +599,15 @@ def python_get_2d_local_memory(self, t, batch_size, num_heads, height, width,
   def testGet2dLocalMemory(self):
     batch_size = 3
     num_heads = 3
-    height = 12
+    height = 6
     width = 6
     depth = 15
     num_h_blocks = 3
     num_w_blocks = 3
-    memory_flange = [2, 1]
-    query_shape = [4, 2]
+    memory_flange = [1, 1]
+    query_shape = [2, 2]
     t = np.random.rand(batch_size, num_heads, height, width, depth)
-    a = common_attention.get_2d_local_memory(
+    a = common_attention.get_2d_local_memory_v2(
         np.reshape(t, (batch_size*num_heads, height, width, depth)),
         query_shape, memory_flange)
     self.evaluate(tf.global_variables_initializer())

From b79c310aace8ebffc9b309bd0d04797b2489a483 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 13 Mar 2019 13:57:25 -0700
Subject: [PATCH 1786/2720] Search in base and env registries to create a
 problem.

PiperOrigin-RevId: 238301625
---
 tensor2tensor/utils/registry.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index cbde96852..42f673517 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -493,19 +493,24 @@ def list_hparams(prefix=None):
 register_problem = register_base_problem
 
 
-def problem(problem_name):
-  """Get possibly copied/reversed problem registered in `base_registry`.
+def problem(problem_name, **kwargs):
+  """Get possibly copied/reversed problem in `base_registry` or `env_registry`.
 
   Args:
     problem_name: string problem name. See `parse_problem_name`.
+    **kwargs: forwarded to env problem's initialize method.
 
   Returns:
     possibly reversed/copied version of base problem registered in the given
     registry.
   """
   spec = parse_problem_name(problem_name)
-  return Registries.problems[spec.base_name](
-      was_copy=spec.was_copy, was_reversed=spec.was_reversed)
+  try:
+    return Registries.problems[spec.base_name](
+        was_copy=spec.was_copy, was_reversed=spec.was_reversed)
+  except KeyError:
+    # If name is not found in base problems then try creating an env problem
+    return env_problem(problem_name, **kwargs)
 
 
 def env_problem(env_problem_name, **kwargs):

From 6838765ccd9b469cde3c2686f695fb0bb03e258e Mon Sep 17 00:00:00 2001
From: Trevor Gale <tgale@google.com>
Date: Thu, 14 Mar 2019 16:32:47 -0700
Subject: [PATCH 1787/2720] Fixing dev token count for lm1b32k.

PiperOrigin-RevId: 238544184
---
 tensor2tensor/data_generators/lm1b.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/lm1b.py b/tensor2tensor/data_generators/lm1b.py
index 2c19bbb42..c44d82d35 100644
--- a/tensor2tensor/data_generators/lm1b.py
+++ b/tensor2tensor/data_generators/lm1b.py
@@ -108,7 +108,7 @@ class LanguagemodelLm1b32k(text_problems.Text2SelfProblem):
   """A language model on the 1B words corpus.
 
   Ratio of dev tokens (including eos) to dev words (including eos)
-  176884 / 159658 = 1.107893; multiply log_ppl by this to compare results.
+  176923 / 159658 = 1.108137; multiply log_ppl by this to compare results.
   """
 
   @property

From c3b602493be1a9b910a8b79668657a78f20b46dc Mon Sep 17 00:00:00 2001
From: Shivani Agrawal <shivaniagrawal@google.com>
Date: Thu, 14 Mar 2019 22:30:26 -0700
Subject: [PATCH 1788/2720] Separated out skip_summary() method from
 summary_op_util, fixing accordingly.

PiperOrigin-RevId: 238585287
---
 tensor2tensor/layers/common_video.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 81f7f0568..015bc3d97 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -24,7 +24,8 @@
 from tensor2tensor.layers import common_layers
 import tensorflow as tf
 
-from tensorflow.python.ops import summary_op_util
+from tensorflow.python.distribute import summary_op_util as distribute_summary_op_util  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.ops import summary_op_util  # pylint: disable=g-direct-tensorflow-import
 
 tfl = tf.layers
 tfcl = tf.contrib.layers
@@ -475,7 +476,7 @@ def gif_summary(name, tensor, max_outputs=3, fps=10, collections=None,
                      "[batch, time, height, width, channels] but got one "
                      "of shape: %s" % str(tensor.get_shape()))
   tensor = tf.cast(tensor, tf.uint8)
-  if summary_op_util.skip_summary():
+  if distribute_summary_op_util.skip_summary():
     return tf.constant("")
   with summary_op_util.summary_scope(
       name, family, values=[tensor]) as (tag, scope):

From e84c425df0a6db9a9b9c9e2cdd9c782c5e2505ab Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Fri, 15 Mar 2019 08:33:25 -0700
Subject: [PATCH 1789/2720] Replaced gather local 2d with splits instead of
 slices. Slightly better MXU utilization.

PiperOrigin-RevId: 238646469
---
 tensor2tensor/layers/common_attention.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index c2992fc47..dc8344611 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2370,7 +2370,7 @@ def get_2d_local_memory(x, query_shape, memory_flange):
 
 
 def get_2d_local_memory_v2(x, query_shape, memory_flange):
-  """Transposeless local 2d memory block construction.
+  """Gathering memory blocks around query blocks. flange is half of query .
 
     Only works if memory flanges are half of query sizes.
 
@@ -2396,16 +2396,15 @@ def get_2d_local_memory_v2(x, query_shape, memory_flange):
   num_w_memory_blocks = width//query_shape[1] + 1
   x_memory_blocks = _extract_blocks(padded_x,
                                     query_shape[0], query_shape[1])
-  x_left_width = tf.slice(x_memory_blocks, [0, 0, 0, 0, 0, 0],
-                          [-1, -1, num_w_memory_blocks - 1, -1, -1, -1])
-  x_right_width = tf.slice(x_memory_blocks, [0, 0, 1, 0, 0, 0],
-                           [-1, -1, -1, - 1, -1, -1])
+  x_width_blocks = tf.split(x_memory_blocks, num_w_memory_blocks,
+                            2)
+  x_left_width = tf.concat(x_width_blocks[:num_w_memory_blocks - 1], axis=2)
+  x_right_width = tf.concat(x_width_blocks[1:], axis=2)
   x_memory_blocks = tf.concat([x_left_width, x_right_width], axis=4)
 
-  x_top_height = tf.slice(x_memory_blocks, [0, 0, 0, 0, 0, 0],
-                          [-1, num_h_memory_blocks-1, -1, - 1, -1, -1])
-  x_bottom_height = tf.slice(x_memory_blocks, [0, 1, 0, 0, 0, 0],
-                             [-1, -1, -1, - 1, -1, -1])
+  x_height_blocks = tf.split(x_memory_blocks, num_h_memory_blocks, 1)
+  x_top_height = tf.concat(x_height_blocks[:num_h_memory_blocks - 1], axis=1)
+  x_bottom_height = tf.concat(x_height_blocks[1:], axis=1)
   x = tf.concat([x_top_height, x_bottom_height], axis=3)
 
   return x

From fc6037d3ea9a4878ae85a0d8761a8d66a292fbdf Mon Sep 17 00:00:00 2001
From: cbockman <c.bockman@gmail.com>
Date: Fri, 15 Mar 2019 08:52:06 -0700
Subject: [PATCH 1790/2720] (minor) spelling deault -> default (#1485)

---
 tensor2tensor/models/research/universal_transformer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 9945bb140..9f13e4125 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -350,7 +350,7 @@ def body(self, features):
 
 
 def update_hparams_for_universal_transformer(hparams):
-  """Adds deault hparams for all of the variants of the Universal Transformer.
+  """Adds default hparams for all of the variants of the Universal Transformer.
 
   Args:
     hparams: default hparams (usually one of the standard hparams from

From f23f535d19644ce3693045c4674d1044cb454751 Mon Sep 17 00:00:00 2001
From: Le Zhang <lzhang10@users.noreply.github.com>
Date: Fri, 15 Mar 2019 11:53:25 -0400
Subject: [PATCH 1791/2720] Fix step size extraction for checkpoint name with -
 in it such as avg-model.ckpt-1234 (#1487)

---
 tensor2tensor/utils/decoding.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 6d6a15624..b54390cf6 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -956,7 +956,7 @@ def _decode_input_tensor_to_features_dict(feature_map, hparams):
 
 
 def get_step_from_ckpt_path(path):
-  return int(os.path.basename(path).split("-")[1])
+  return int(os.path.basename(path).split("-")[-1])
 
 
 def latest_checkpoint_step(ckpt_dir):

From 3edecf91e7420a2dce19bc4e2229351e7475964f Mon Sep 17 00:00:00 2001
From: cbockman <c.bockman@gmail.com>
Date: Fri, 15 Mar 2019 08:52:26 -0700
Subject: [PATCH 1792/2720] Merge of PR #1485

PiperOrigin-RevId: 238649310
---
 tensor2tensor/utils/decoding.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index b54390cf6..6d6a15624 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -956,7 +956,7 @@ def _decode_input_tensor_to_features_dict(feature_map, hparams):
 
 
 def get_step_from_ckpt_path(path):
-  return int(os.path.basename(path).split("-")[-1])
+  return int(os.path.basename(path).split("-")[1])
 
 
 def latest_checkpoint_step(ckpt_dir):

From a187657b3f5f81ee1efc8a47a69837c03c80b27b Mon Sep 17 00:00:00 2001
From: Le Zhang <lzhang10@users.noreply.github.com>
Date: Fri, 15 Mar 2019 08:53:47 -0700
Subject: [PATCH 1793/2720] Merge of PR #1487

PiperOrigin-RevId: 238649533
---
 tensor2tensor/utils/decoding.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 6d6a15624..b54390cf6 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -956,7 +956,7 @@ def _decode_input_tensor_to_features_dict(feature_map, hparams):
 
 
 def get_step_from_ckpt_path(path):
-  return int(os.path.basename(path).split("-")[1])
+  return int(os.path.basename(path).split("-")[-1])
 
 
 def latest_checkpoint_step(ckpt_dir):

From 0542447c4a7172336f7b91b9dc388ece89f0d801 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 15 Mar 2019 16:18:50 -0700
Subject: [PATCH 1794/2720] Implemented Neural Turing Machine.

PiperOrigin-RevId: 238730108
---
 tensor2tensor/layers/transformer_memory.py    | 180 +++++++++---------
 .../layers/transformer_memory_test.py         | 109 +++++------
 2 files changed, 141 insertions(+), 148 deletions(-)

diff --git a/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index c16c638bd..383df5013 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -13,121 +13,124 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""The memory unit for remembering a sequence as a collection of clusters."""
+"""The memory unit for Transformer."""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.layers import common_layers
 import tensorflow as tf
 
 
 class TransformerMemory(object):
   """Implements the Memory module.
 
-  It compresses a sequence by storing items into appropriate clusters.
-  A single item can be allocated into multiple clusters like a mixture model.
-  Each vector in the memory represents the centroid of the cluster that is
-  updated in an online fashion. The memory also keeps the total amount of
-  probability mass that is used for updating each item that indicates the amount
-  of change that has been made to each cluster.
+  Based on Neural Turing Machines: arXiv:1410.5401 [cs.NE]
   """
 
-  def __init__(self, batch_size, feature_dim, memory_size):
+  def __init__(self, batch_size, key_depth, val_depth, memory_size,
+               sharpen_factor=1.):
     """Initialize the memory object.
 
     Args:
       batch_size: the batch size.
-      feature_dim: the depth of the feature.
-      memory_size: the number of clusters to maintain in the memory, which does
-          not have to be the same as the segment length.
+      key_depth: the depth of the memory keys.
+      val_depth: the depth of the memory values.
+      memory_size: the number of items in the memory.
+      sharpen_factor: the sharpen_factor for addressing the memory.
     """
-    self.feature_dim = feature_dim
     self.batch_size = batch_size
+    self.key_depth = key_depth
+    self.val_depth = val_depth
     self.memory_size = memory_size
+    self.sharpen_factor = sharpen_factor
     self.mem_vals = tf.get_variable(
-        "memvals", [self.batch_size, self.memory_size, self.feature_dim],
+        "memvals", [self.batch_size, self.memory_size, self.val_depth],
+        dtype=tf.float32, trainable=False,
+        initializer=tf.constant_initializer(.0))
+    self.mean_logits = tf.get_variable(
+        "meanlogits", [self.batch_size, self.memory_size],
         dtype=tf.float32, trainable=False,
         initializer=tf.constant_initializer(.0))
-    self.mem_times = tf.get_variable(
-        "memtimes", [self.batch_size, self.memory_size], dtype=tf.float32,
-        trainable=False, initializer=tf.constant_initializer(.0))
-    self.seq_length_so_far = tf.get_variable(
-        "seqlensofar", [self.batch_size], dtype=tf.int32,
-        trainable=False, initializer=tf.constant_initializer(0))
-
-  def set(self, mem_vals, mem_times, seq_length_so_far):
-    set_op = tf.group([
-        self.mem_vals.assign(mem_vals),
-        self.mem_times.assign(mem_times),
-        self.seq_length_so_far.assign(seq_length_so_far)])
-    return set_op
 
-  def get(self):
-    return self.mem_vals, self.mem_times, self.seq_length_so_far
+  def _address_content(self, x):
+    """Address the memory based on content similarity.
 
-  def incremental_update(self, event):
-    """Add a new event to the memory and also advance the time.
+    Args:
+      x: a tensor in the shape of [batch_size, length, depth].
+    Returns:
+      the logits for each memory entry [batch_size, length, memory_size].
+    """
+    mem_keys = tf.layers.dense(self.mem_vals, self.key_depth, name="mem_key")
+    mem_query = tf.layers.dense(x, self.key_depth, name="mem_query")
+    norm = tf.matmul(
+        tf.norm(mem_query, axis=-1, keepdims=True),
+        tf.norm(mem_keys, axis=-1, keepdims=True), transpose_b=True)
+    cos_dist = tf.div(
+        tf.matmul(mem_query, mem_keys, transpose_b=True), norm,
+        name="cos_dist")
+    access_logits = self.sharpen_factor * cos_dist
+    return access_logits
+
+  def read(self, x):
+    """Read from the memory.
+
+    An external component can use the results via a simple MLP,
+    e.g., fn(x W_x + retrieved_mem W_m).
 
     Args:
-      event: a tensor in the shape of [batch_size, depth].
+      x: a tensor in the shape of [batch_size, length, depth].
     Returns:
-      the update op.
+      access_logits: the logits for accessing the memory in shape of
+          [batch_size, length, memory_size].
+      retrieved_mem: the retrieved results in the shape of
+          [batch_size, length, val_depth].
     """
-    event = tf.expand_dims(event, 1)
-    similarity_logits = tf.matmul(event, tf.transpose(
-        self.mem_vals, [0, 2, 1]))
-    similarity_logits = tf.squeeze(similarity_logits, [1])
-    max_logits = tf.reduce_max(similarity_logits, -1, keep_dims=True)
-    similarity_logits = tf.where(
-        tf.less(self.mem_times, 0.5),
-        tf.tile(max_logits, [1, self.memory_size]) + 1.0,
-        similarity_logits)
-    _, indices = tf.nn.top_k(similarity_logits)
-    update_mask = tf.cast(tf.one_hot(indices, self.memory_size), tf.float32)
-    update_times = self.mem_times.assign_add(update_mask)
-    with tf.control_dependencies([update_times]):
-      add_to_vals = tf.where(
-          tf.cast(update_mask, tf.bool),
-          tf.zeros_like(self.mem_vals),
-          tf.div(event - self.mem_vals, tf.expand_dims(self.mem_times, 2)))
-      return self.mem_vals.assign_add(add_to_vals)
-
-  def update(self, segment):
-    """Update the memory given the segment of events.
-
-    It might be useful to consider adding a decay to each cluster to favor
-    recent events.
+    access_logits = self._address_content(x)
+    weights = tf.nn.softmax(access_logits)
+    retrieved_mem = tf.reduce_sum(
+        tf.multiply(tf.expand_dims(weights, 3),
+                    tf.expand_dims(self.mem_vals, axis=1)), axis=2)
+    return access_logits, retrieved_mem
+
+  def write(self, x, access_logits):
+    """Write to the memory based on a combination of similarity and least used.
+
+    Based on arXiv:1607.00036v2 [cs.LG].
 
     Args:
-      segment: a tensor of shape [batch_size, segment_length, depth].
+      x: a tensor in the shape of [batch_size, length, depth].
+      access_logits: the logits for accessing the memory.
     Returns:
       the update op.
     """
-    attention_logits = tf.matmul(segment, tf.transpose(
-        self.mem_vals, [0, 2, 1]))
-    alloc_probs = tf.nn.softmax(attention_logits)
-    aggregated_alloc_probs = tf.reduce_sum(alloc_probs, axis=1)
-    time_increment = tf.where(
-        tf.equal(self.seq_length_so_far, 0),
-        tf.ones_like(self.mem_times),
-        aggregated_alloc_probs)
-    update_times = self.mem_times.assign_add(time_increment)
-    with tf.control_dependencies([update_times]):
-      allocations = tf.multiply(
-          tf.expand_dims(alloc_probs, 3), tf.expand_dims(segment, 2))
-      allocations = tf.reduce_sum(allocations, axis=1)
-      add_to_vals = tf.where(
-          tf.equal(self.seq_length_so_far, 0),
-          segment,
-          tf.div(allocations - self.mem_vals,
-                 tf.expand_dims(self.mem_times, 2)))
-      update_vals = self.mem_vals.assign_add(add_to_vals)
-      with tf.control_dependencies([update_vals]):
-        segment_length = common_layers.shape_list(segment)[1]
-        update_seq_length = self.seq_length_so_far.assign_add(
-            tf.tile(tf.expand_dims(segment_length, 0), [self.batch_size]))
-    return update_seq_length
+    gamma = tf.layers.dense(x, 1, activation=tf.sigmoid, name="gamma")
+    write_logits = access_logits - gamma * tf.expand_dims(self.mean_logits, 1)
+    candidate_value = tf.layers.dense(x, self.val_depth,
+                                      activation=tf.nn.relu,
+                                      name="candidate_value")
+    erase_gates = tf.layers.dense(x, self.memory_size,
+                                  activation=tf.nn.sigmoid,
+                                  name="erase")
+    write_weights = tf.nn.softmax(write_logits)
+    erase = tf.multiply(tf.expand_dims(1 - erase_gates * write_weights, 3),
+                        tf.expand_dims(self.mem_vals, 1))
+    addition = tf.multiply(
+        tf.expand_dims(write_weights, 3), tf.expand_dims(candidate_value, 2))
+    update_value_op = self.mem_vals.assign(
+        tf.reduce_sum(erase + addition, axis=1))
+    with tf.control_dependencies([update_value_op]):
+      write_op = self.mean_logits.assign(
+          self.mean_logits * 0.1 + tf.reduce_sum(write_logits * 0.9, axis=1))
+      return write_op
+
+  def set(self, mem_vals, mean_logits):
+    set_op = tf.group([
+        self.mem_vals.assign(mem_vals),
+        self.mean_logits.assign(mean_logits)])
+    return set_op
+
+  def get(self):
+    return self.mem_vals, self.mean_logits
 
   def reset(self, entries_to_reset):
     """Reset the entries in the memory.
@@ -141,13 +144,12 @@ def reset(self, entries_to_reset):
     update_vals = tf.scatter_update(
         self.mem_vals, entries_to_reset,
         tf.tile(tf.expand_dims(
-            tf.fill([self.memory_size, self.feature_dim], .0), 0),
+            tf.fill([self.memory_size, self.val_depth], .0), 0),
                 [num_updates, 1, 1]))
-    update_times = tf.scatter_update(
-        self.mem_times, entries_to_reset,
+    update_logits = tf.scatter_update(
+        self.mean_logits, entries_to_reset,
         tf.tile(tf.expand_dims(
-            tf.fill([self.memory_size], .0), 0), [num_updates, 1]))
-    update_segs = tf.scatter_update(
-        self.seq_length_so_far, entries_to_reset, tf.fill([num_updates], 0))
-    reset_op = tf.group([update_vals, update_times, update_segs])
+            tf.fill([self.memory_size], .0), 0),
+                [num_updates, 1]))
+    reset_op = tf.group([update_vals, update_logits])
     return reset_op
diff --git a/tensor2tensor/layers/transformer_memory_test.py b/tensor2tensor/layers/transformer_memory_test.py
index 4df01abae..b5c365c41 100644
--- a/tensor2tensor/layers/transformer_memory_test.py
+++ b/tensor2tensor/layers/transformer_memory_test.py
@@ -26,86 +26,77 @@
 
 class TransformerMemoryTest(parameterized.TestCase, tf.test.TestCase):
 
-  def testInitialize(self):
+  def testRead(self):
     batch_size = 2
-    feature_dim = 3
+    key_depth = 3
+    val_depth = 5
     memory_size = 4
+    window_size = 6
+    x_depth = 10
     memory = transformer_memory.TransformerMemory(
-        batch_size, feature_dim, memory_size)
-    segment = tf.constant([[[1., 2., 3.], [1., 1., 1.],
-                            [3., 2., 1.], [2., 2., 2.]],
-                           [[3., 3., 3.], [1., 2., 3.],
-                            [3., 2., 1.], [2., 2., 2.]]])
-    update_op = memory.update(segment)
-    mem_vals, mem_times, mem_len_so_far = memory.get()
+        batch_size, key_depth, val_depth, memory_size)
+    x = tf.random_uniform([batch_size, window_size, x_depth], minval=1.0)
+    vals = tf.random_uniform([batch_size, memory_size, val_depth], minval=1.0)
+    logits = tf.random_uniform([batch_size, memory_size], minval=1.0)
+    update_op = memory.set(vals, logits)
+    with tf.control_dependencies([update_op]):
+      logits, retrieved_values = memory.read(x)
     with self.test_session() as session:
       session.run(tf.global_variables_initializer())
-      session.run(update_op)
-      vals, times, length_so_far = session.run([
-          mem_vals, mem_times, mem_len_so_far])
-    self.assertAllEqual([[[1., 2., 3.], [1., 1., 1.],
-                          [3., 2., 1.], [2., 2., 2.]],
-                         [[3., 3., 3.], [1., 2., 3.],
-                          [3., 2., 1.], [2., 2., 2.]]], vals)
-    self.assertAllEqual([[1., 1., 1., 1.], [1., 1., 1., 1.]], times)
-    self.assertAllEqual([4, 4], length_so_far)
+      logits_values, values = session.run([logits, retrieved_values])
+    self.assertAllEqual([batch_size, window_size, memory_size],
+                        logits_values.shape)
+    self.assertAllEqual([batch_size, window_size, val_depth], values.shape)
 
-  def testUpdate(self):
+  def testWrite(self):
     batch_size = 2
-    feature_dim = 3
+    key_depth = 3
+    val_depth = 5
     memory_size = 4
+    window_size = 6
+    x_depth = 10
     memory = transformer_memory.TransformerMemory(
-        batch_size, feature_dim, memory_size)
-    segment = tf.constant([[[1., 2., 3.], [2., 2., 2.],
-                            [3., 2., 1.], [2., 2., 2.]],
-                           [[2., 2., 2.], [1., 2., 3.],
-                            [3., 2., 1.], [2., 2., 2.]]])
-    init_op = memory.set(segment, [[1., 2., 3., 4.], [2., 1., 5., 1.]],
-                         [10, 9])
-    new_segment = tf.constant(
-        [[[1., 2., 3.], [3., 2., 1.],
-          [2., 2., 2.], [2., 2., 2.]],
-         [[2., 2., 2.], [1., 2., 3.],
-          [3., 2., 1.], [2., 2., 2.]]])
-    update_op = memory.update(new_segment)
-    mem_vals, mem_times, mem_len_so_far = memory.get()
+        batch_size, key_depth, val_depth, memory_size)
+    x = tf.random_uniform([batch_size, window_size, x_depth], minval=1.0)
+    vals = tf.random_uniform([batch_size, memory_size, val_depth], minval=1.0)
+    logits = tf.random_uniform([batch_size, memory_size], minval=1.0)
+    update_op = memory.set(vals, logits)
+    with tf.control_dependencies([update_op]):
+      logits, _ = memory.read(x)
+      write_op = memory.write(x, logits)
+    mem_vals, mem_logits = memory.get()
     with self.test_session() as session:
       session.run(tf.global_variables_initializer())
-      session.run(init_op)
-      session.run(update_op)
-      vals, times, length_so_far = session.run([
-          mem_vals, mem_times, mem_len_so_far])
-      print(vals, times, length_so_far)
-    self.assertAllEqual([2, 4, 3], vals.shape)
-    self.assertAllEqual([2, 4], times.shape)
-    self.assertAllEqual([14, 13], length_so_far)
+      session.run(write_op)
+      updated_vals, updated_logits = session.run([mem_vals, mem_logits])
+    self.assertAllEqual([batch_size, memory_size, val_depth],
+                        updated_vals.shape)
+    self.assertAllEqual([batch_size, memory_size], updated_logits.shape)
 
   def testReset(self):
     batch_size = 2
-    feature_dim = 3
+    key_depth = 3
+    val_depth = 5
     memory_size = 4
     memory = transformer_memory.TransformerMemory(
-        batch_size, feature_dim, memory_size)
-    segment = tf.constant([[[1., 2., 3.], [1., 1., 1.],
-                            [3., 2., 1.], [2., 2., 2.]],
-                           [[3., 3., 3.], [1., 2., 3.],
-                            [3., 2., 1.], [2., 2., 2.]]])
-    update_op = memory.set(segment, [[1., 2., 3., 4.], [2., 1., 5., 1.]],
-                           [10, 9])
+        batch_size, key_depth, val_depth, memory_size)
+    vals = tf.random_uniform([batch_size, memory_size, val_depth], minval=1.0)
+    logits = tf.random_uniform([batch_size, memory_size], minval=1.0)
+    update_op = memory.set(vals, logits)
     reset_op = memory.reset([1])
-    mem_vals, mem_times, mem_len_so_far = memory.get()
+    mem_vals, mem_logits = memory.get()
+    assert_op1 = tf.assert_equal(mem_vals[0], vals[0])
+    assert_op2 = tf.assert_equal(mem_logits[0], logits[0])
+    with tf.control_dependencies([assert_op1, assert_op2]):
+      all_zero1 = tf.reduce_sum(tf.abs(mem_vals[1]))
+      all_zero2 = tf.reduce_sum(tf.abs(mem_logits[1]))
     with self.test_session() as session:
       session.run(tf.global_variables_initializer())
       session.run(update_op)
       session.run(reset_op)
-      vals, times, length_so_far = session.run([
-          mem_vals, mem_times, mem_len_so_far])
-    self.assertAllEqual([[[1., 2., 3.], [1., 1., 1.],
-                          [3., 2., 1.], [2., 2., 2.]],
-                         [[0., 0., 0.], [0., 0., 0.],
-                          [0., 0., 0.], [0., 0., 0.]]], vals)
-    self.assertAllEqual([[1., 2., 3., 4.], [0., 0., 0., 0.]], times)
-    self.assertAllEqual([10, 0], length_so_far)
+      zero1, zero2 = session.run([all_zero1, all_zero2])
+    self.assertAllEqual(0, zero1)
+    self.assertAllEqual(0, zero2)
 
 if __name__ == "__main__":
   tf.test.main()

From 8d93b2eb65d30a36a1b14ba861fd2a3c59c56d0f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 15 Mar 2019 17:33:43 -0700
Subject: [PATCH 1795/2720] Base hparam sets for languagemodel_wikitext103_l4k

PiperOrigin-RevId: 238741124
---
 tensor2tensor/models/transformer.py | 34 +++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 4fead25d2..5888df41c 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -2501,3 +2501,37 @@ def transformer_tpu_1b():
   # maximize number of parameters relative to computation by not sharing.
   hparams.shared_embedding_and_softmax_weights = False
   return hparams
+
+
+@registry.register_hparams
+def transformer_wikitext103_l4k_v0():
+  """HParams for training languagemodel_wikitext103_l4k."""
+  hparams = transformer_big()
+
+  # Adafactor uses less memory than Adam.
+  # switch to Adafactor with its recommended learning rate scheme.
+  hparams.optimizer = "Adafactor"
+  hparams.learning_rate_schedule = "rsqrt_decay"
+  hparams.learning_rate_warmup_steps = 10000
+
+  hparams.num_heads = 4
+  hparams.max_length = 4096
+  hparams.batch_size = 4096
+  hparams.shared_embedding_and_softmax_weights = False
+
+  hparams.num_hidden_layers = 8
+  hparams.attention_dropout = 0.1
+  hparams.layer_prepostprocess_dropout = 0.2
+  hparams.relu_dropout = 0.1
+  hparams.label_smoothing = 0.0
+
+  # Using noise broadcast in the dropout layers saves memory during training.
+  hparams.attention_dropout_broadcast_dims = "0,1"  # batch, heads
+  hparams.relu_dropout_broadcast_dims = "1"  # length
+  hparams.layer_prepostprocess_dropout_broadcast_dims = "1"  # length
+
+  # Avoid an expensive concat on TPU.
+  # >1 shards helps with faster parameter distribution on multi-GPU machines
+  hparams.symbol_modality_num_shards = 1
+
+  return hparams

From 64eeb9de755cc367cf23a9b2d487a3580f2b6bcc Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 18 Mar 2019 12:05:58 -0700
Subject: [PATCH 1796/2720] Correct eval adjusting schedule and use for resnet
 in trax (improves training).

PiperOrigin-RevId: 239033536
---
 tensor2tensor/bin/t2t_trainer.py              |  2 +-
 .../trax/configs/resnet50_imagenet_8gb.gin    | 12 ++++++-----
 tensor2tensor/trax/learning_rate.py           | 21 ++++++++++++-------
 3 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 304d49e0b..c82043b3a 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -388,7 +388,7 @@ def _setup_gin():
       from tensor2tensor.trax import inputs as _trax_inputs
       from tensor2tensor.trax import models as _trax_models
       from tensor2tensor.trax import optimizers as _trax_opt
-      # pylint: disable=g-import-not-at-top,unused-import,g-bad-import-order,reimported,unused-variable
+      # pylint: enable=g-import-not-at-top,unused-import,g-bad-import-order,reimported,unused-variable
 
       configs = config or []
       # Override with --dataset and --model
diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
index 811117dfc..5c9bb365b 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
@@ -1,4 +1,5 @@
 import tensor2tensor.trax.inputs
+import tensor2tensor.trax.learning_rate
 import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
@@ -17,8 +18,8 @@ inputs.dataset_name = 't2t_image_imagenet224'
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-MultifactorSchedule.constant = 5.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+EvalAdjustingSchedule.constant = 1.0
+MultifactorSchedule.factors = 'constant * linear_warmup'
 MultifactorSchedule.warmup_steps = 400
 
 # Parameters for momentum:
@@ -36,9 +37,10 @@ Resnet50.num_output_classes = 1001
 
 # Parameters for train:
 # ==============================================================================
-train.eval_frequency = 100
-train.eval_steps = 3
+train.eval_frequency = 2000
+train.eval_steps = 20
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.Resnet50
 train.optimizer = @trax.optimizers.momentum
-train.train_steps = 200000
+train.train_steps = 500000
+train.lr_schedule = @learning_rate.EvalAdjustingSchedule
diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index e767a57a5..e14f43efb 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -42,7 +42,7 @@
 @gin.configurable(blacklist=["history"])
 def MultifactorSchedule(history=None,
                         factors="constant * linear_warmup * rsqrt_decay",
-                        constant=0.001,
+                        constant=0.1,
                         warmup_steps=100):
   """Factor-based learning rate schedule.
 
@@ -88,10 +88,10 @@ def learning_rate(step):  # pylint: disable=invalid-name
 
 @gin.configurable(blacklist=["history"])
 def EvalAdjustingSchedule(history,
-                          constant=0.001,
-                          steps_to_decrease=10,
-                          improvement_margin=0.01,
-                          decrease_rate=2.0,
+                          constant=0.1,
+                          steps_to_decrease=20,
+                          improvement_margin=0.001,
+                          decrease_rate=1.5,
                           history_mode="eval",
                           metric="metrics/accuracy"):
   """Learning rate that decreases when eval metric stalls.
@@ -116,15 +116,22 @@ def EvalAdjustingSchedule(history,
   """
   metrics = history.get(history_mode, metric)
   adjusted = constant
+  if len(metrics) < 2:
+    return MultifactorSchedule(history, constant=adjusted)
+
   steps_without_improvement = 0
+  cur = metrics.pop()[1]  # The most-recent value of the metric.
   while len(metrics) > 1:
-    last = metrics.pop()
-    if last[1] < metrics[-1][1] * (1 + improvement_margin):
+    # The one-before value of metrics as .pop() removes one element each time.
+    prev = metrics.pop()[1]
+    if cur < prev * (1 + improvement_margin):
       steps_without_improvement += 1
     else:
+      cur = prev
       steps_without_improvement = 0
     if steps_without_improvement >= steps_to_decrease:
       adjusted /= decrease_rate
+      cur = prev
       steps_without_improvement = 0
 
   return MultifactorSchedule(history, constant=adjusted)

From 191d9ad56af3209ceaa92adbf368893751f32ed5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 18 Mar 2019 14:13:59 -0700
Subject: [PATCH 1797/2720] Resolve TODO.

PiperOrigin-RevId: 239058553
---
 tensor2tensor/models/transformer.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 5888df41c..261927e96 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -471,7 +471,10 @@ def preprocess_targets(targets, i):
         targets = dp(bottom, targets, hparams, target_vocab_size)[0]
       targets = common_layers.flatten4d3d(targets)
 
-      # TODO(llion): Explain! Is this even needed?
+      # GO embeddings are all zero, this is because transformer_prepare_decoder
+      # Shifts the targets along by one for the input which pads with zeros.
+      # If the modality already maps GO to the zero embeddings this is not
+      # needed.
       targets = tf.cond(
           tf.equal(i, 0), lambda: tf.zeros_like(targets), lambda: targets)
 
@@ -708,7 +711,10 @@ def preprocess_targets(targets, i):
         targets = dp(bottom, targets, hparams, target_vocab_size)[0]
       targets = common_layers.flatten4d3d(targets)
 
-      # TODO(llion): Explain! Is this even needed?
+      # GO embeddings are all zero, this is because transformer_prepare_decoder
+      # Shifts the targets along by one for the input which pads with zeros.
+      # If the modality already maps GO to the zero embeddings this is not
+      # needed.
       targets = tf.cond(
           tf.equal(i, 0), lambda: tf.zeros_like(targets), lambda: targets)
 

From e5e7d4babf9c57d943a12f10124439fc50d5e2d5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 18 Mar 2019 15:20:46 -0700
Subject: [PATCH 1798/2720] Transformer with memory in the style of
 Transformer-XL

PiperOrigin-RevId: 239072524
---
 tensor2tensor/layers/common_attention.py   |  30 ++++++
 tensor2tensor/layers/transformer_memory.py | 102 +++++++++++++++++++
 tensor2tensor/models/transformer.py        | 110 +++++++++++++++++++--
 3 files changed, 233 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index dc8344611..090e7bc39 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -3868,6 +3868,8 @@ def multihead_attention(query_antecedent,
                         dropout_broadcast_dims=None,
                         vars_3d=False,
                         layer_collection=None,
+                        recurrent_memory=None,
+                        chunk_number=None,
                         **kwargs):
   """Multihead scaled-dot-product attention with input/output transformations.
 
@@ -3921,6 +3923,10 @@ def multihead_attention(query_antecedent,
     vars_3d: use 3-dimensional variables for input/output transformations
     layer_collection: A tensorflow_kfac.LayerCollection. Only used by the
       KFAC optimizer. Default is None.
+    recurrent_memory: An optional transformer_memory.RecurrentMemory, which
+      retains state across chunks. Default is None.
+    chunk_number: an optional integer Tensor with shape [batch] used to operate
+      the recurrent_memory.
     **kwargs (dict): Parameters for the attention function
 
   Caching:
@@ -3958,9 +3964,30 @@ def multihead_attention(query_antecedent,
     if vars_3d:
       raise ValueError("KFAC implementation does not support 3d vars.")
 
+  if recurrent_memory is not None:
+    if memory_antecedent is not None:
+      raise ValueError("Recurrent memory requires memory_antecedent is None.")
+    if cache is not None:
+      raise ValueError("Cache is not supported when using recurrent memory.")
+    if vars_3d:
+      raise ValueError("3d vars are not supported when using recurrent memory.")
+    if layer_collection is not None:
+      raise ValueError("KFAC is not supported when using recurrent memory.")
+    if chunk_number is None:
+      raise ValueError("chunk_number is required when using recurrent memory.")
+
   with tf.variable_scope(name, default_name="multihead_attention",
                          values=[query_antecedent, memory_antecedent]):
 
+    if recurrent_memory is not None:
+      (
+          recurrent_memory_transaction,
+          query_antecedent, memory_antecedent, bias,
+      ) = recurrent_memory.pre_attention(
+          chunk_number,
+          query_antecedent, memory_antecedent, bias,
+      )
+
     if cache is None or memory_antecedent is None:
       q, k, v = compute_qkv(query_antecedent, memory_antecedent,
                             total_key_depth, total_value_depth, q_filter_width,
@@ -4111,6 +4138,9 @@ def multihead_attention(query_antecedent,
       x = common_layers.dense(
           x, output_depth, use_bias=False, name="output_transform",
           layer_collection=layer_collection)
+
+    if recurrent_memory is not None:
+      x = recurrent_memory.post_attention(recurrent_memory_transaction, x)
     if additional_returned_value is not None:
       return x, additional_returned_value
     return x
diff --git a/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index 383df5013..34f334dce 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -21,6 +21,108 @@
 import tensorflow as tf
 
 
+class RecurrentMemory(object):
+  """Base class for recurrent memory.
+
+  Currently implements memory in the style of Transformer-XL
+  (https://arxiv.org/abs/1901.02860)
+  """
+  # TODO(kitaev): make this a base class and then subclass for different memory
+  # types (e.g. the one defined below in this file).
+
+  def __init__(self, name, hparams):
+    hidden_size = hparams.hidden_size
+    chunk_length = hparams.split_targets_chunk_length
+    assert chunk_length > 0, "Chunking is required to use RecurrentMemory"
+
+    # TODO(kitaev): The implementation of the chunking code makes it somewhat
+    # convoluted to figure out how many actual sequences we can have per batch.
+    # The data pipeline should be revisited at some point.
+    batch_size_in_sequences = hparams.batch_size / hparams.max_length
+
+    memory_shape = [batch_size_in_sequences, chunk_length, hidden_size]
+    bias_shape = [1, 1, chunk_length, chunk_length]
+
+    with tf.variable_scope(name):
+      self.previous_segment = tf.get_variable(
+          "memsegment", (),
+          dtype=tf.int32, trainable=False,
+          initializer=tf.constant_initializer(0))
+
+      self.previous_vals = tf.get_variable(
+          "memvals", memory_shape,
+          dtype=tf.float32, trainable=False,
+          initializer=tf.constant_initializer(.0))
+
+      self.previous_bias = tf.get_variable(
+          "membias", bias_shape,
+          dtype=tf.float32, trainable=False,
+          initializer=tf.constant_initializer(.0))
+
+  def pre_attention(self, segment, query_antecedent, memory_antecedent, bias):
+    """Called prior to self-attention, to incorporate memory items.
+
+    Args:
+      segment: an integer Tensor with shape [batch]
+      query_antecedent: a Tensor with shape [batch, length_q, channels]
+      memory_antecedent: must be None. Attention normally allows this to be a
+        Tensor with shape [batch, length_m, channels], but we currently only
+        support memory for decoder-side self-attention.
+      bias: bias Tensor (see attention_bias())
+    Returns:
+      (data, new_query_antecedent, new_memory_antecedent, new_bias)
+    """
+    assert memory_antecedent is None, "We only support language modeling"
+
+    previous_vals = tf.stop_gradient(self.previous_vals)
+    # If segment id is zero, don't attend back to the memory
+    previous_bias = tf.stop_gradient(self.previous_bias) + tf.cast(
+        tf.equal(tf.reduce_sum(segment), 0), tf.float32) * -1e9
+
+    # In eval mode, batch size may be variable
+    amount_to_pad = tf.shape(previous_vals)[0] - tf.shape(query_antecedent)[0]
+    previous_vals = previous_vals[:tf.shape(query_antecedent)[0], :, :]
+    with tf.control_dependencies(
+        [tf.assert_equal(tf.shape(query_antecedent), tf.shape(previous_vals))]):
+      query_antecedent = tf.identity(query_antecedent)
+
+    new_memory_antecedent = tf.concat(
+        [tf.stop_gradient(previous_vals), query_antecedent], 1)
+    new_bias = tf.concat([previous_bias, bias], -1)
+
+    cancel_update = tf.equal(self.previous_segment, segment[0])
+    remember_segment = segment[0]
+    remember_vals = tf.cond(
+        cancel_update,
+        lambda: self.previous_vals,
+        lambda: tf.pad(query_antecedent, [[0, amount_to_pad], [0, 0], [0, 0]]))
+    remember_bias = tf.cond(
+        cancel_update,
+        lambda: self.previous_bias,
+        lambda: tf.zeros_like(bias) + tf.reduce_max(bias, -1, keep_dims=True))
+
+    token = (remember_segment, remember_vals, remember_bias)
+
+    return token, query_antecedent, new_memory_antecedent, new_bias
+
+  def post_attention(self, token, x):
+    """Called after self-attention. The memory can be updated here.
+
+    Args:
+      token: Data returned by pre_attention, which can be used to carry over
+        state related to the current memory operation.
+      x: a Tensor of data after self-attention and feed-forward
+    Returns:
+      a (possibly modified) version of the input x
+    """
+    with tf.control_dependencies([
+        self.previous_segment.assign(token[0]),
+        self.previous_vals.assign(token[1]),
+        self.previous_bias.assign(token[2]),
+        ]):
+      return tf.identity(x)
+
+
 class TransformerMemory(object):
   """Implements the Memory module.
 
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 261927e96..60b6e086f 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -34,6 +34,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import modalities
 from tensor2tensor.layers import transformer_layers
+from tensor2tensor.layers import transformer_memory
 from tensor2tensor.utils import beam_search
 from tensor2tensor.utils import expert_utils
 from tensor2tensor.utils import mlperf_log
@@ -54,7 +55,8 @@
 
 
 def transformer_encode(encoder_function, inputs, target_space, hparams,
-                       attention_weights=None, features=None, losses=None):
+                       attention_weights=None, features=None, losses=None,
+                       **kwargs):
   """Encode transformer inputs.
 
   Args:
@@ -67,6 +69,7 @@ def transformer_encode(encoder_function, inputs, target_space, hparams,
     features: optionally pass the entire features dictionary as well. This is
       needed now for "packed" datasets.
     losses: optional list onto which to append extra training losses
+    **kwargs: additional arguments to pass to encoder_function
 
   Returns:
     Tuple of:
@@ -102,7 +105,8 @@ def transformer_encode(encoder_function, inputs, target_space, hparams,
       save_weights_to=attention_weights,
       make_image_summary=not common_layers.is_xla_compiled(),
       losses=losses,
-      attn_bias_for_padding=attn_bias_for_padding)
+      attn_bias_for_padding=attn_bias_for_padding,
+      **kwargs)
 
   return encoder_output, encoder_decoder_attention_bias
 
@@ -117,7 +121,8 @@ def transformer_decode(decoder_function,
                        cache=None,
                        decode_loop_step=None,
                        nonpadding=None,
-                       losses=None):
+                       losses=None,
+                       **kwargs):
   """Decode Transformer outputs from encoder representation.
 
   Args:
@@ -138,6 +143,7 @@ def transformer_decode(decoder_function,
       for inference on TPU.
     nonpadding: optional Tensor with shape [batch_size, decoder_length]
     losses: optional list onto which to append extra training losses
+    **kwargs: additional arguments to pass to decoder_function
 
   Returns:
     Final decoder representation. [batch_size, decoder_length, hidden_dim]
@@ -159,7 +165,8 @@ def transformer_decode(decoder_function,
       decode_loop_step=decode_loop_step,
       nonpadding=nonpadding,
       save_weights_to=attention_weights,
-      losses=losses)
+      losses=losses,
+      **kwargs)
 
   if (common_layers.is_xla_compiled() and
       hparams.mode == tf.estimator.ModeKeys.TRAIN):
@@ -178,6 +185,7 @@ class Transformer(t2t_model.T2TModel):
   def __init__(self, *args, **kwargs):
     super(Transformer, self).__init__(*args, **kwargs)
     self.attention_weights = {}  # For visualizing attention heads.
+    self.recurrent_memory_by_layer = None  # Override to enable recurrent memory
     self._encoder_function = transformer_encoder
     self._decoder_function = transformer_decoder
 
@@ -197,13 +205,15 @@ def decode(self,
              cache=None,
              decode_loop_step=None,
              nonpadding=None,
-             losses=None):
+             losses=None,
+             **kwargs):
     """Decode Transformer outputs, see transformer_decode."""
     return transformer_decode(
         self._decoder_function, decoder_input, encoder_output,
         encoder_decoder_attention_bias, decoder_self_attention_bias,
         hparams, attention_weights=self.attention_weights, cache=cache,
-        decode_loop_step=decode_loop_step, nonpadding=nonpadding, losses=losses)
+        decode_loop_step=decode_loop_step, nonpadding=nonpadding, losses=losses,
+        **kwargs)
 
   def body(self, features):
     """Transformer main model_fn.
@@ -236,6 +246,28 @@ def body(self, features):
     targets = common_layers.flatten4d3d(targets)
     decoder_input, decoder_self_attention_bias = transformer_prepare_decoder(
         targets, hparams, features=features)
+
+    # Not all subclasses of Transformer support keyword arguments related to
+    # recurrent memory, so only pass these arguments if memory is enabled.
+    decode_kwargs = {}
+    if self.recurrent_memory_by_layer is not None:
+      # TODO(kitaev): The chunk_number feature currently has the same shape as
+      # "targets", but this is only for the purposes of sharing sharding code.
+      # In fact every token within the batch must have the same chunk number.
+      chunk_number_each_token = tf.squeeze(features["chunk_number"], (-1, -2))
+      chunk_number_each_batch = chunk_number_each_token[:, 0]
+      # Uncomment the code below to verify that tokens within a batch share the
+      # same chunk number:
+      # with tf.control_dependencies([
+      #     tf.assert_equal(chunk_number_each_token,
+      #                     chunk_number_each_batch[:, None])
+      # ]):
+      #   chunk_number_each_batch = tf.identity(chunk_number_each_batch)
+      decode_kwargs = dict(
+          recurrent_memory_by_layer=self.recurrent_memory_by_layer,
+          chunk_number=chunk_number_each_batch,
+          )
+
     decoder_output = self.decode(
         decoder_input,
         encoder_output,
@@ -243,7 +275,9 @@ def body(self, features):
         decoder_self_attention_bias,
         hparams,
         nonpadding=features_to_nonpadding(features, "targets"),
-        losses=losses)
+        losses=losses,
+        **decode_kwargs
+        )
 
     expected_attentions = features.get("expected_attentions")
     if expected_attentions is not None:
@@ -1330,7 +1364,10 @@ def transformer_decoder(decoder_input,
                         save_weights_to=None,
                         make_image_summary=True,
                         losses=None,
-                        layer_collection=None):
+                        layer_collection=None,
+                        recurrent_memory_by_layer=None,
+                        chunk_number=None,
+                        ):
   """A stack of transformer layers.
 
   Args:
@@ -1358,6 +1395,10 @@ def transformer_decoder(decoder_input,
     losses: optional list onto which to append extra training losses
     layer_collection: A tensorflow_kfac.LayerCollection. Only used by the
       KFAC optimizer. Default is None.
+    recurrent_memory_by_layer: Optional dict, mapping layer names to instances
+      of transformer_memory.RecurrentMemory. Default is None.
+    chunk_number: an optional integer Tensor with shape [batch] used to operate
+      the recurrent_memory.
 
   Returns:
     y: a Tensors
@@ -1388,6 +1429,10 @@ def transformer_decoder(decoder_input,
     for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
       layer_name = "layer_%d" % layer
       layer_cache = cache[layer_name] if cache is not None else None
+      if recurrent_memory_by_layer is not None:
+        recurrent_memory = recurrent_memory_by_layer[layer_name]
+      else:
+        recurrent_memory = None
       with tf.variable_scope(layer_name):
         with tf.variable_scope("self_attention"):
           y = common_attention.multihead_attention(
@@ -1414,7 +1459,10 @@ def transformer_decoder(decoder_input,
               vars_3d=hparams.get("attention_variables_3d"),
               activation_dtype=hparams.get("activation_dtype", "float32"),
               weight_dtype=hparams.get("weight_dtype", "float32"),
-              layer_collection=layer_collection)
+              layer_collection=layer_collection,
+              recurrent_memory=recurrent_memory,
+              chunk_number=chunk_number,
+              )
           x = common_layers.layer_postprocess(x, y, hparams)
         if encoder_output is not None:
           with tf.variable_scope("encdec_attention"):
@@ -1464,6 +1512,33 @@ def transformer_decoder(decoder_input,
         x, hparams, layer_collection=layer_collection)
 
 
+@registry.register_model
+class TransformerMemory(Transformer):
+  """Transformer language model with memory across chunks."""
+
+  # TODO(kitaev): consider overriding set_mode to swap out recurrent memory when
+  # switching between training and evaluation.
+
+  def __init__(self, *args, **kwargs):
+    super(TransformerMemory, self).__init__(*args, **kwargs)
+
+    hparams = self._hparams
+    self.recurrent_memory_by_layer = {}
+    for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
+      layer_name = "layer_%d" % layer
+      self.recurrent_memory_by_layer[layer_name] = transformer_memory.RecurrentMemory(
+          layer_name + "/recurrent_memory", hparams)
+
+
+  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha,
+                   use_tpu=False):
+    """Overriding beam search because for now only the slow version works with
+    memory
+    """
+    return self._beam_decode_slow(features, decode_length, beam_size,
+                                  top_beams, alpha, use_tpu)
+
+
 @registry.register_hparams
 def transformer_base_v1():
   """Set of hyperparameters."""
@@ -2541,3 +2616,20 @@ def transformer_wikitext103_l4k_v0():
   hparams.symbol_modality_num_shards = 1
 
   return hparams
+
+
+@registry.register_hparams
+def transformer_wikitext103_l4k_memory():
+  """HParams for training languagemodel_wikitext103_l4k with memory."""
+  hparams = transformer_wikitext103_l4k_v0()
+
+  hparams.split_targets_chunk_length = 8
+  hparams.split_targets_max_chunks = 512
+
+  # The hparams specify batch size *before* chunking, but we want to have a
+  # consistent 4K batch size *after* chunking to fully utilize the hardware.
+  target_tokens_per_batch = 4096
+  hparams.batch_size = target_tokens_per_batch * (
+      hparams.max_length / hparams.split_targets_chunk_length)  # 2097152
+
+  return hparams

From 886b7742ff56983866a3e30d4366fa6397f0cea8 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 19 Mar 2019 11:31:19 -0700
Subject: [PATCH 1799/2720] fix shape bugs in dilated attention

PiperOrigin-RevId: 239233677
---
 tensor2tensor/layers/common_attention.py      | 36 +++++++------
 tensor2tensor/layers/common_attention_test.py | 51 +++++++++++++++++++
 2 files changed, 68 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 090e7bc39..5dc3a58e7 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -3075,9 +3075,9 @@ def dilated_self_attention_1d(q,
   """Dilated self-attention.
 
   Args:
-    q: a Tensor with shape [batch, heads, length, depth_k]
-    k: a Tensor with shape [batch, heads, length, depth_k]
-    v: a Tensor with shape [batch, heads, length, depth_v]
+    q: a Tensor with shape [batch, heads, length, depth]
+    k: a Tensor with shape [batch, heads, length, depth]
+    v: a Tensor with shape [batch, heads, length, depth]
     query_block_size: an integer indicating size of query block
     memory_block_size: an integer indicating the size of a memory block.
     gap_size: an integer indicating the gap size
@@ -3086,11 +3086,12 @@ def dilated_self_attention_1d(q,
     name: an optional string
 
   Returns:
-    a Tensor of shape [batch, heads, length, depth_v]
+    a Tensor of shape [batch, heads, length, depth]
   """
   with tf.variable_scope(
       name, default_name="dilated_self_attention_1d", values=[q, k, v]):
     v_list_shape = v.get_shape().as_list()
+    assert v_list_shape == k.shape.as_list(), "K and V depths must be equal"
     v_shape = common_layers.shape_list(v)
     depth_v = v_shape[3]
     batch_size = v_shape[0]
@@ -3108,9 +3109,6 @@ def pad_l_and_r(x, pad_length):
     q = pad_to_multiple(q, query_block_size)
     v = pad_to_multiple(v, query_block_size)
     k = pad_to_multiple(k, query_block_size)
-    q.set_shape(v_list_shape)
-    v.set_shape(v_list_shape)
-    k.set_shape(v_list_shape)
 
     # Set up query blocks.
     new_q_shape = common_layers.shape_list(q)
@@ -3212,21 +3210,23 @@ def gather_dilated_memory_blocks(x,
   # gathering memory blocks
   for block_id in range(num_memory_blocks):
     block_end_index = -(query_block_size + gap_size *
-                        (block_id + 1) + memory_block_size * block_id) - 1
+                        (block_id + 1) + memory_block_size * block_id)
     block_start_index = (
         (memory_block_size + gap_size) * (num_memory_blocks - (block_id + 1)))
     if direction != "left":
       [block_end_index,
-       block_start_index] = [-block_start_index - 1, -block_end_index + 1]
+       block_start_index] = [-block_start_index, -block_end_index]
+    if block_end_index == 0:
+      x_block = x[block_start_index:]
+    else:
+      x_block = x[block_start_index:block_end_index]
 
     def gather_dilated_1d_blocks(x, gather_indices):
       x_new = tf.gather(x, gather_indices)
       # [batch, heads, blocks, block_length, dim]
       return tf.transpose(x_new, [2, 3, 0, 1, 4])
 
-    gathered_blocks.append(
-        gather_dilated_1d_blocks(x[block_start_index:block_end_index],
-                                 gather_indices))
+    gathered_blocks.append(gather_dilated_1d_blocks(x_block, gather_indices))
   return tf.concat(gathered_blocks, 3)
 
 
@@ -3241,9 +3241,9 @@ def masked_dilated_self_attention_1d(q,
   """Dilated self-attention. TODO(avaswani): Try it and write a paper on it.
 
   Args:
-    q: a Tensor with shape [batch, heads, length, depth_k]
-    k: a Tensor with shape [batch, heads, length, depth_k]
-    v: a Tensor with shape [batch, heads, length, depth_v]
+    q: a Tensor with shape [batch, heads, length, depth]
+    k: a Tensor with shape [batch, heads, length, depth]
+    v: a Tensor with shape [batch, heads, length, depth]
     query_block_size: an integer
     memory_block_size: an integer indicating how much to look left.
     gap_size: an integer indicating the gap size
@@ -3252,11 +3252,12 @@ def masked_dilated_self_attention_1d(q,
     name: an optional string
 
   Returns:
-    a Tensor of shape [batch, heads, length, depth_v]
+    a Tensor of shape [batch, heads, length, depth]
   """
   with tf.variable_scope(
       name, default_name="masked_dilated_self_attention_1d", values=[q, k, v]):
     v_list_shape = v.get_shape().as_list()
+    assert v_list_shape == k.shape.as_list(), "K and V depths must be equal"
     v_shape = common_layers.shape_list(v)
     depth_v = v_shape[3]
     batch_size = v_shape[0]
@@ -3274,9 +3275,6 @@ def pad_l(x, left_pad_length):
     q = pad_to_multiple(q, query_block_size)
     v = pad_to_multiple(v, query_block_size)
     k = pad_to_multiple(k, query_block_size)
-    q.set_shape(v_list_shape)
-    v.set_shape(v_list_shape)
-    k.set_shape(v_list_shape)
 
     # Set up query blocks.
     new_q_shape = common_layers.shape_list(q)
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index ae688576c..8be53be23 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -1126,6 +1126,57 @@ def testMultiheadAttentionWithLayerCollection(self):
         layer_collection=layer_collection)
     self.assertLen(layer_collection.get_blocks(), 4)
 
+  @parameterized.named_parameters(
+      ("", 1, 1, 8, 4, 3),
+      ("dynamic_batch", None, 1, 8, 4, 2),
+      ("batches", 4, 3, 8, 4, 2),
+      ("block_length", 1, 1, 8, 4, 4),
+  )
+  def testDilatedAttention(self, batch, heads, length, depth_v, block_length):
+    if batch is None:
+      batch = tf.random_uniform([], minval=0, maxval=5, dtype=tf.int32)
+    q = tf.random_normal([batch, heads, length, depth_v])
+    k = tf.random_normal([batch, heads, length, depth_v])
+    v = tf.random_normal([batch, heads, length, depth_v])
+    output = common_attention.dilated_self_attention_1d(
+        q, k, v,
+        query_block_size=block_length,
+        memory_block_size=block_length,
+        gap_size=2,
+        num_memory_blocks=2)
+    if isinstance(batch, tf.Tensor):
+      batch, res = self.evaluate([batch, output])
+    else:
+      res = self.evaluate(output)
+
+    self.assertEqual(res.shape, (batch, heads, length, depth_v))
+
+  @parameterized.named_parameters(
+      ("", 1, 1, 8, 4, 3),
+      ("dynamic_batch", None, 1, 8, 4, 2),
+      ("batches", 4, 3, 8, 4, 2),
+      ("block_length", 1, 1, 8, 4, 4),
+  )
+  def testMaskedDilatedAttention(self, batch, heads, length, depth_v,
+                                 block_length):
+    if batch is None:
+      batch = tf.random_uniform([], minval=0, maxval=5, dtype=tf.int32)
+    q = tf.random_normal([batch, heads, length, depth_v])
+    k = tf.random_normal([batch, heads, length, depth_v])
+    v = tf.random_normal([batch, heads, length, depth_v])
+    output = common_attention.masked_dilated_self_attention_1d(
+        q, k, v,
+        query_block_size=block_length,
+        memory_block_size=block_length,
+        gap_size=2,
+        num_memory_blocks=2)
+    if isinstance(batch, tf.Tensor):
+      batch, res = self.evaluate([batch, output])
+    else:
+      res = self.evaluate(output)
+
+    self.assertEqual(res.shape, (batch, heads, length, depth_v))
+
 if __name__ == "__main__":
   tf.test.main()
 

From ba35f3d0475abc2cf99441e6ee1ac35ef126c984 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 19 Mar 2019 11:54:12 -0700
Subject: [PATCH 1800/2720] Added hparam for top-k random sampling in
 Transformer.

PiperOrigin-RevId: 239238821
---
 tensor2tensor/layers/common_hparams.py |  1 +
 tensor2tensor/layers/common_layers.py  | 20 ++++++++++++++++++--
 tensor2tensor/models/transformer.py    |  9 +++++++--
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 8d3f23eaf..55b147987 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -105,6 +105,7 @@ def basic_params1():
       learning_rate=0.1,
       sampling_method="argmax",  # "argmax" or "random"
       sampling_temp=1.0,  # temperature for sampling
+      sampling_keep_top_k=-1,  # If >0, ignore all but the top k logits
       # expand the logits a piece at a time - saves memory.
       factored_logits=False,
       multiply_embedding_mode="sqrt_depth",
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 6b583549f..f30032536 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2849,13 +2849,13 @@ def list_product(els):
   return prod
 
 
-def sample_with_temperature(logits, temperature):
+def sample_with_temperature(logits, temperature, sampling_keep_top_k=-1):
   """Either argmax or random sampling.
 
   Args:
     logits: a Tensor.
     temperature: a float  0.0=argmax 1.0=random
-
+    sampling_keep_top_k: If not -1, only sample from the top k logits.
   Returns:
     a Tensor with one fewer dimension than logits.
   """
@@ -2866,6 +2866,22 @@ def sample_with_temperature(logits, temperature):
     return tf.reshape(argmax, logits_shape[:-1])
   else:
     assert temperature > 0.0
+
+    if sampling_keep_top_k != -1:
+      if sampling_keep_top_k <= 0:
+        raise ValueError("sampling_keep_top_k must either be -1 or positive.")
+
+      vocab_size = shape_list(logits)[1]
+
+      k_largest = tf.contrib.nn.nth_element(
+          logits, n=sampling_keep_top_k, reverse=True)
+      k_largest = tf.tile(tf.reshape(k_largest, [-1, 1]), [1, vocab_size])
+
+      # Force every position that is not in the top k to have probability near
+      # 0 by setting the logit to be very negative.
+      logits = tf.where(tf.less_equal(logits, k_largest),
+                        tf.ones_like(logits)*-1e6, logits)
+
     reshaped_logits = (
         tf.reshape(logits, [-1, shape_list(logits)[-1]]) / temperature)
     choices = tf.multinomial(reshaped_logits, 1)
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 60b6e086f..6c47d2d53 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -973,9 +973,12 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       logits, cache = symbols_to_logits_fn(next_id, i, cache)
       log_probs = common_layers.log_prob_from_logits(logits)
       temperature = getattr(hparams, "sampling_temp", 0.0)
+      keep_top = getattr(hparams, "sampling_keep_top_k", -1)
       if hparams.sampling_method == "argmax":
         temperature = 0.0
-      next_id = common_layers.sample_with_temperature(logits, temperature)
+      next_id = common_layers.sample_with_temperature(
+          logits, temperature, keep_top)
+
       hit_eos |= tf.equal(next_id, eos_id)
 
       log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],
@@ -1156,9 +1159,11 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       logits, cache = symbols_to_logits_fn(next_id, i, cache)
       log_probs = common_layers.log_prob_from_logits(logits)
       temperature = getattr(hparams, "sampling_temp", 0.0)
+      keep_top = getattr(hparams, "sampling_keep_top_k", -1)
       if hparams.sampling_method == "argmax":
         temperature = 0.0
-      next_id = common_layers.sample_with_temperature(logits, temperature)
+      next_id = common_layers.sample_with_temperature(
+          logits, temperature, keep_top)
       hit_eos |= tf.equal(next_id, eos_id)
 
       log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],

From 366944295a614bf9cb07bcfc0255a10046d3d331 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 21 Mar 2019 11:19:41 -0700
Subject: [PATCH 1801/2720] Allow Evolved Transformer number of decoder
 attention heads to exceed 16.

PiperOrigin-RevId: 239634336
---
 tensor2tensor/models/evolved_transformer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index a99a01e81..62098e004 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -287,7 +287,7 @@ def evolved_transformer_decoder(decoder_input,
           residual_state = hidden_state
           hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
 
-          # 16 head attention. Hard coding number of heads.
+          # Attention with at least 16 heads.
           left_state = common_attention.multihead_attention(
               hidden_state,
               None,
@@ -295,7 +295,7 @@ def evolved_transformer_decoder(decoder_input,
               hparams.attention_key_channels or hparams.hidden_size,
               hparams.attention_value_channels or hparams.hidden_size,
               hparams.hidden_size,
-              16,  # Heads are hard coded to replicate paper.
+              max(16, hparams.num_heads),
               hparams.attention_dropout,
               attention_type=hparams.self_attention_type,
               max_relative_position=hparams.max_relative_position,

From 3dce919d063fc28a8c96565719e1dfc754f49a6c Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 21 Mar 2019 19:22:50 +0100
Subject: [PATCH 1802/2720] RL fixes (#1505)

---
 tensor2tensor/rl/evaluator.py | 15 ++++++++-------
 tensor2tensor/rl/rl_utils.py  |  3 +--
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 0459da703..4b8728233 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -252,7 +252,7 @@ def make_env(env_type, real_env, sim_env_kwargs):
 
 def make_agent(
     agent_type, env, policy_hparams, policy_dir, sampling_temp,
-    sim_env_kwargs=None, frame_stack_size=None, rollout_agent_type=None,
+    sim_env_kwargs_fn=None, frame_stack_size=None, rollout_agent_type=None,
     batch_size=None, inner_batch_size=None, env_type=None, **planner_kwargs
 ):
   """Factory function for Agents."""
@@ -270,7 +270,7 @@ def make_agent(
           batch_size, make_agent(
               rollout_agent_type, env, policy_hparams, policy_dir,
               sampling_temp, batch_size=inner_batch_size
-          ), make_env(env_type, env.env, sim_env_kwargs),
+          ), make_env(env_type, env.env, sim_env_kwargs_fn()),
           lambda env: rl_utils.BatchStackWrapper(env, frame_stack_size),
           discount_factor=policy_hparams.gae_gamma, **planner_kwargs
       ),
@@ -302,17 +302,18 @@ def make_agent_from_hparams(
     planner_hparams, model_dir, policy_dir, sampling_temp, video_writers=()
 ):
   """Creates an Agent from hparams."""
-  sim_env_kwargs = rl.make_simulated_env_kwargs(
-      base_env, loop_hparams, batch_size=planner_hparams.batch_size,
-      model_dir=model_dir
-  )
+  def sim_env_kwargs_fn():
+      return rl.make_simulated_env_kwargs(
+          base_env, loop_hparams, batch_size=planner_hparams.batch_size,
+          model_dir=model_dir
+      )
   planner_kwargs = planner_hparams.values()
   planner_kwargs.pop("batch_size")
   planner_kwargs.pop("rollout_agent_type")
   planner_kwargs.pop("env_type")
   return make_agent(
       agent_type, stacked_env, policy_hparams, policy_dir, sampling_temp,
-      sim_env_kwargs, loop_hparams.frame_stack_size,
+      sim_env_kwargs_fn, loop_hparams.frame_stack_size,
       planner_hparams.rollout_agent_type,
       inner_batch_size=planner_hparams.batch_size,
       env_type=planner_hparams.env_type,
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 9eaac4f48..2e4c7660f 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -414,9 +414,8 @@ def augment_observation(
       (1, 15), "f:{:3}".format(int(frame_index)),
       fill=(255, 0, 0)
   )
-  header = np.asarray(img)
+  header = np.copy(np.asarray(img))
   del img
-  header.setflags(write=1)
   if bar_color is not None:
     header[0, :, :] = bar_color
   return np.concatenate([header, observation], axis=0)

From 13824527f68ab72722327a30fc7590934c2de2e3 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Thu, 21 Mar 2019 11:36:54 -0700
Subject: [PATCH 1803/2720] Merge of PR #1505

PiperOrigin-RevId: 239638087
---
 tensor2tensor/rl/evaluator.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 4b8728233..f86a99fc2 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -303,10 +303,10 @@ def make_agent_from_hparams(
 ):
   """Creates an Agent from hparams."""
   def sim_env_kwargs_fn():
-      return rl.make_simulated_env_kwargs(
-          base_env, loop_hparams, batch_size=planner_hparams.batch_size,
-          model_dir=model_dir
-      )
+    return rl.make_simulated_env_kwargs(
+        base_env, loop_hparams, batch_size=planner_hparams.batch_size,
+        model_dir=model_dir
+    )
   planner_kwargs = planner_hparams.values()
   planner_kwargs.pop("batch_size")
   planner_kwargs.pop("rollout_agent_type")

From 5d345a9af255812aac1726cdd691b5ec7772a03e Mon Sep 17 00:00:00 2001
From: Yongkeun Hwang <ykstyle@ykstyle.info>
Date: Fri, 22 Mar 2019 03:46:33 +0900
Subject: [PATCH 1804/2720] Allowing to use user-defined modules on t2t-eval
 (#1480)

---
 tensor2tensor/bin/t2t_eval.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/bin/t2t_eval.py b/tensor2tensor/bin/t2t_eval.py
index 2df551810..3bdb4634d 100644
--- a/tensor2tensor/bin/t2t_eval.py
+++ b/tensor2tensor/bin/t2t_eval.py
@@ -22,6 +22,7 @@
 from tensor2tensor.bin import t2t_trainer          # pylint: disable=unused-import
 from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
+from tensor2tensor.utils import usr_dir
 import tensorflow as tf
 
 flags = tf.flags
@@ -31,6 +32,7 @@
 def main(_):
   tf.logging.set_verbosity(tf.logging.INFO)
   trainer_lib.set_random_seed(FLAGS.random_seed)
+  usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
 
   hparams = trainer_lib.create_hparams(
       FLAGS.hparams_set, FLAGS.hparams, data_dir=FLAGS.data_dir,

From d5b9ba25514a6f891bb171dae0b5eeb99024a3e6 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 21 Mar 2019 11:46:03 -0700
Subject: [PATCH 1805/2720] Adding pypng package because of
 RenderedEnvProblems.

PiperOrigin-RevId: 239639943
---
 setup.py                      | 3 ++-
 tensor2tensor/bin/t2t_eval.py | 2 --
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/setup.py b/setup.py
index 3fda512dd..544b3a7e8 100644
--- a/setup.py
+++ b/setup.py
@@ -50,13 +50,14 @@
         'numpy',
         'oauth2client',
         'opencv-python',
+        'pypng',
         'requests',
         'scipy',
         'six',
         'sympy',
+        'tensorflow-datasets',
         'tensorflow-probability',
         'tf-agents',
-        'tensorflow-datasets',
         'tqdm',
     ],
     extras_require={
diff --git a/tensor2tensor/bin/t2t_eval.py b/tensor2tensor/bin/t2t_eval.py
index 3bdb4634d..2df551810 100644
--- a/tensor2tensor/bin/t2t_eval.py
+++ b/tensor2tensor/bin/t2t_eval.py
@@ -22,7 +22,6 @@
 from tensor2tensor.bin import t2t_trainer          # pylint: disable=unused-import
 from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
-from tensor2tensor.utils import usr_dir
 import tensorflow as tf
 
 flags = tf.flags
@@ -32,7 +31,6 @@
 def main(_):
   tf.logging.set_verbosity(tf.logging.INFO)
   trainer_lib.set_random_seed(FLAGS.random_seed)
-  usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
 
   hparams = trainer_lib.create_hparams(
       FLAGS.hparams_set, FLAGS.hparams, data_dir=FLAGS.data_dir,

From 642279b68bd4e72b5ce6a0750627c73a69f063e0 Mon Sep 17 00:00:00 2001
From: Yongkeun Hwang <ykstyle@ykstyle.info>
Date: Thu, 21 Mar 2019 11:46:54 -0700
Subject: [PATCH 1806/2720] Merge of PR #1480

PiperOrigin-RevId: 239640123
---
 tensor2tensor/bin/t2t_eval.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/bin/t2t_eval.py b/tensor2tensor/bin/t2t_eval.py
index 2df551810..3bdb4634d 100644
--- a/tensor2tensor/bin/t2t_eval.py
+++ b/tensor2tensor/bin/t2t_eval.py
@@ -22,6 +22,7 @@
 from tensor2tensor.bin import t2t_trainer          # pylint: disable=unused-import
 from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
+from tensor2tensor.utils import usr_dir
 import tensorflow as tf
 
 flags = tf.flags
@@ -31,6 +32,7 @@
 def main(_):
   tf.logging.set_verbosity(tf.logging.INFO)
   trainer_lib.set_random_seed(FLAGS.random_seed)
+  usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
 
   hparams = trainer_lib.create_hparams(
       FLAGS.hparams_set, FLAGS.hparams, data_dir=FLAGS.data_dir,

From a4071d62f510a3b0dace62f9fa78e2f9a60c5c40 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 21 Mar 2019 11:55:52 -0700
Subject: [PATCH 1807/2720] Remove unused tf-agents dependency for now, will
 add once we need to.

PiperOrigin-RevId: 239641868
---
 setup.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/setup.py b/setup.py
index 544b3a7e8..adb181a12 100644
--- a/setup.py
+++ b/setup.py
@@ -57,7 +57,6 @@
         'sympy',
         'tensorflow-datasets',
         'tensorflow-probability',
-        'tf-agents',
         'tqdm',
     ],
     extras_require={

From b56a5cdc2a6c0e4d66214a1f5708db50696d90b7 Mon Sep 17 00:00:00 2001
From: Akio Ohta <drunkars.p@gmail.com>
Date: Fri, 22 Mar 2019 04:29:55 +0900
Subject: [PATCH 1808/2720] Modify serving utils (#1495)

* Add decode logic for model with return_beams=True.

* Add print logic for model with return_beams=True.
---
 tensor2tensor/serving/query.py         | 18 +++++++++++++++---
 tensor2tensor/serving/serving_utils.py |  5 ++++-
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/serving/query.py b/tensor2tensor/serving/query.py
index a0ea895c7..c4978bc26 100644
--- a/tensor2tensor/serving/query.py
+++ b/tensor2tensor/serving/query.py
@@ -90,14 +90,26 @@ def main(_):
     outputs = serving_utils.predict([inputs], problem, request_fn)
     outputs, = outputs
     output, score = outputs
-    print_str = """
+    if len(score.shape) > 0:
+      print_str = """
+Input:
+{inputs}
+
+Output (Scores [{score}]):
+{output}
+        """
+      score_text = ",".join(["{:.3f}".format(s) for s in score])
+      print(print_str.format(inputs=inputs, output=output, score=score_text))
+    else:
+      print_str = """
 Input:
 {inputs}
 
 Output (Score {score:.3f}):
 {output}
-    """
-    print(print_str.format(inputs=inputs, output=output, score=score))
+        """
+      print(print_str.format(inputs=inputs, output=output, score=score))
+
     if FLAGS.inputs_once:
       break
 
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 487a312b1..263f555de 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -94,7 +94,10 @@ def _encode(inputs, encoder, add_eos=True):
 
 
 def _decode(output_ids, output_decoder):
-  return output_decoder.decode(output_ids, strip_extraneous=True)
+  if len(output_ids.shape) > 1:
+    return [output_decoder.decode(o, strip_extraneous=True) for o in output_ids]
+  else:
+    return output_decoder.decode(output_ids, strip_extraneous=True)
 
 
 

From 893f14f147035a2b4fe36922f8468f460cb60286 Mon Sep 17 00:00:00 2001
From: Akio Ohta <drunkars.p@gmail.com>
Date: Thu, 21 Mar 2019 12:30:18 -0700
Subject: [PATCH 1809/2720] Merge of PR #1495

PiperOrigin-RevId: 239648633
---
 tensor2tensor/serving/query.py         | 2 +-
 tensor2tensor/serving/serving_utils.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/serving/query.py b/tensor2tensor/serving/query.py
index c4978bc26..75fb896b7 100644
--- a/tensor2tensor/serving/query.py
+++ b/tensor2tensor/serving/query.py
@@ -90,7 +90,7 @@ def main(_):
     outputs = serving_utils.predict([inputs], problem, request_fn)
     outputs, = outputs
     output, score = outputs
-    if len(score.shape) > 0:
+    if len(score.shape) > 0:  # pylint: disable=g-explicit-length-test
       print_str = """
 Input:
 {inputs}
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 263f555de..cfc2f3b5b 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -117,7 +117,7 @@ def _make_grpc_request(examples):
     outputs = tf.make_ndarray(response.outputs["outputs"])
     scores = tf.make_ndarray(response.outputs["scores"])
     assert len(outputs) == len(scores)
-    return [{
+    return [{  # pylint: disable=g-complex-comprehension
         "outputs": output,
         "scores": score
     } for output, score in zip(outputs, scores)]
@@ -134,7 +134,7 @@ def _make_cloud_mlengine_request(examples):
     parent = "projects/%s/models/%s/versions/%s" % (cloud.default_project(),
                                                     model_name, version)
     input_data = {
-        "instances": [{
+        "instances": [{  # pylint: disable=g-complex-comprehension
             "input": {
                 "b64": base64.b64encode(ex.SerializeToString())
             }

From 45b3f4f44152184b37a4081dec2921c6e7bd15c9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 21 Mar 2019 12:55:53 -0700
Subject: [PATCH 1810/2720] Fix Travis breakage for tf != nightly.

PiperOrigin-RevId: 239653689
---
 tensor2tensor/layers/common_video.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 015bc3d97..ea94316cf 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -24,9 +24,16 @@
 from tensor2tensor.layers import common_layers
 import tensorflow as tf
 
-from tensorflow.python.distribute import summary_op_util as distribute_summary_op_util  # pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.ops import summary_op_util  # pylint: disable=g-direct-tensorflow-import
 
+# After tf-nightly 1.14.1.dev20190314 summary_op_util.skip_summary was extracted
+# out to the distribute module.
+try:
+  from tensorflow.python.distribute import summary_op_util as distribute_summary_op_util  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+except ImportError:
+  distribute_summary_op_util = summary_op_util
+
+
 tfl = tf.layers
 tfcl = tf.contrib.layers
 

From e26510c13a5d2a1ce6e6de3afd3ab4148f1283aa Mon Sep 17 00:00:00 2001
From: Ashish Vaswani <avaswani@google.com>
Date: Thu, 21 Mar 2019 13:33:04 -0700
Subject: [PATCH 1811/2720] query shape = memory shape. Allows for just one
 block extraction step for faster TPU processing. We also don't do the final
 linear transformation after attention. Includes tests

PiperOrigin-RevId: 239661165
---
 tensor2tensor/layers/common_attention.py      | 104 ++++++++++++++++--
 tensor2tensor/layers/common_attention_test.py |  68 ++++++++++++
 2 files changed, 164 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 5dc3a58e7..cef04032b 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2434,10 +2434,6 @@ def dot_product_unmasked_attention_local_2d_tpu(
 
   Returns:
     [batch, heads, height, width, depth] tensor, the output of attention.
-    height_key_relative_embeddings: a 3d or 2d tensor, depending on head sharing
-      settings, which are the relative embeddings for height.
-    width_key_relative_embeddings: a 3d or 2d tensor, depending on head sharing
-      settings, which are the relative embeddings for width.
 
   """
   if max_relative_position:
@@ -2506,6 +2502,84 @@ def dot_product_unmasked_attention_local_2d_tpu(
     return ret
 
 
+def dot_product_unmasked_attention_local_2d_tpu_simple(
+    x, bias, total_key_depth, total_value_depth, num_heads,
+    query_shape=(8, 8),
+    dropout_rate=0.0, image_shapes=None, make_image_summary=False,
+    dropout_broadcast_dims=None):
+
+  """Calculate simple unmasked dot-product local self-attention 2d on tpu.
+
+  The query, key, and value blocks are the same. We do not do a second linear
+  transformation after computing the values
+
+  Args:
+    x: a Tensor with shape [batch, height, width, depth].
+    bias: bias Tensor.
+    total_key_depth: the dimensions of the keys
+    total_value_depth: the dimensions of the values
+    num_heads: number of heads
+    query_shape: a two tuple indicating query shape
+    dropout_rate: a floating point number.
+    image_shapes: optional tuple of integer scalars.
+    make_image_summary: Whether to make an attention image summary.
+    dropout_broadcast_dims:  an optional list of integers less than 4
+      specifying in which dimensions to broadcast the dropout decisions.
+      saves memory.
+
+  Returns:
+    ret: [batch, height, width, total_value_depth] tensor,
+      the output of attention.
+    q: [batch, height, width, total_key_depth] query tensor
+    k: [batch, height, width, total_key_depth] key tensor
+    v: [batch, height, width, total_value_depth] value tensor
+
+  """
+  # This calculation only works for self attention.
+  # q, k and v must therefore have the same shape.
+  orig_x_shape = common_layers.shape_list(x)
+  # Pad query, key, value to ensure multiple of corresponding lengths if
+  # necessary
+  is_padded = False
+  if (orig_x_shape[1]%query_shape[0]) != 0 or (
+      orig_x_shape[2]%query_shape[1]) != 0:
+    x = pad_to_multiple_2d(x, query_shape)
+    is_padded = True
+  _, height, width, depth = common_layers.shape_list(x)
+  assert depth%num_heads == 0
+  num_h_blocks = height//query_shape[0]
+  num_w_blocks = width//query_shape[1]
+  # Extract center queries, keys, and values
+  x_blocks = _extract_blocks(x, query_shape[0], query_shape[1])
+  x_blocks = tf.reshape(x_blocks, [-1, query_shape[0]*query_shape[1], depth])
+  q, k, v = compute_qkv(x_blocks, None, total_key_depth, total_value_depth)
+  hsplit = lambda x: split_heads(x, num_heads)
+  q, k, v = map(hsplit, [q, k, v])
+  logits = tf.matmul(q, k, transpose_b=True)
+  if bias is not None:
+    logits += bias
+  weights = tf.nn.softmax(logits, name="attention_weights")
+  # Dropping out the attention links for each of the heads
+  weights = common_layers.dropout_with_broadcast_dims(
+      weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
+  if common_layers.should_generate_summaries() and make_image_summary:
+    attention_image_summary(weights, image_shapes)
+  output = tf.matmul(weights, v)
+  output = combine_heads(output)
+  # we need to get it back to shape [batch, height, width]
+  ret = tf.reshape(output, [-1, num_h_blocks, num_w_blocks,
+                            query_shape[0], query_shape[1], total_value_depth])
+
+  ret = tf.transpose(ret, [0, 1, 3, 2, 4, 5])
+  ret = tf.reshape(ret, [-1, num_h_blocks*query_shape[0],
+                         num_w_blocks*query_shape[1], total_value_depth])
+  # slice if padding was introduced
+  if is_padded:
+    ret = tf.slice(ret, [0, 0, 0, 0], [-1, orig_x_shape[1],
+                                       orig_x_shape[2], -1])
+  return ret, q, k, v
+
+
 def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
   """Attention to the source and a neighborhood to the left within a block.
 
@@ -3415,12 +3489,26 @@ def local_attention_2d(q,
 
 
 def pad_to_multiple_2d(x, block_shape):
-  """Making sure x is a multiple of shape. x is [batch, heads, h, w, depth]."""
+  """Making sure x is a multiple of shape.
+
+  Args:
+    x: a [batch, heads, h, w, depth] or [batch, h, w, depth] tensor
+    block_shape: a 2-d list of integer shapes
+
+  Returns:
+    padded_x: a [batch, heads, h, w, depth] or [batch, h, w, depth] tensor
+  """
   old_shape = x.get_shape().dims
   last = old_shape[-1]
-  height_padding = -common_layers.shape_list(x)[2] % block_shape[0]
-  width_padding = -common_layers.shape_list(x)[3] % block_shape[1]
-  paddings = [[0, 0], [0, 0], [0, height_padding], [0, width_padding], [0, 0]]
+  if len(old_shape) == 4:
+    height_padding = -common_layers.shape_list(x)[1] % block_shape[0]
+    width_padding = -common_layers.shape_list(x)[2] % block_shape[1]
+    paddings = [[0, 0], [0, height_padding], [0, width_padding], [0, 0]]
+  elif len(old_shape) == 5:
+    height_padding = -common_layers.shape_list(x)[2] % block_shape[0]
+    width_padding = -common_layers.shape_list(x)[3] % block_shape[1]
+    paddings = [[0, 0], [0, 0], [0, height_padding], [0, width_padding], [0, 0]]
+
   padded_x = tf.pad(x, paddings)
   padded_shape = padded_x.get_shape().as_list()
   padded_shape = padded_shape[:-1] + [last]
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 8be53be23..08cd6573d 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -815,6 +815,74 @@ def testDotProductUnmaskedAttentionLocal2dTpu(self):
     out = out[:, :, :height, :width, :]
     self.assertAllClose(res, out)
 
+  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  def testDotProductUnmaskedAttentionLocal2dTpuSimple(self):
+    batch_size = 1
+    num_heads = 3
+    height = 8
+    width = 12
+    total_depth = 15
+    num_h_blocks = 4
+    num_w_blocks = 6
+    depth = 5
+    query_shape = [2, 2]
+
+    x = np.random.rand(batch_size, height, width, total_depth)
+    a = (
+        common_attention.dot_product_unmasked_attention_local_2d_tpu_simple(
+            tf.constant(x, dtype=tf.float32),
+            None, total_depth, total_depth, num_heads,
+            query_shape=query_shape))
+    self.evaluate(tf.global_variables_initializer())
+    res, q, k, v = self.evaluate(a)
+    self.assertEqual(res.shape, (batch_size, height, width, total_depth))
+    # reshape q, k, v from batch, heads, height*width to batch, heads,
+    # num_h_blocks, num_w_blocks, query_shape[0], query_shape[1], depth
+    resh_shape = (batch_size, num_h_blocks, num_w_blocks,
+                  num_heads, query_shape[0], query_shape[1],
+                  depth)
+    resh = lambda l: np.reshape(l, resh_shape)
+    q, k, v = map(resh, [q, k, v])
+    trans = lambda l: np.transpose(l, (0, 3, 1, 2, 4, 5, 6))
+    q, k, v = map(trans, [q, k, v])
+    new_height = height + -height % query_shape[0]
+    new_width = width + -width % query_shape[1]
+    (queries, keys, values) = (q, k, v)
+    logits = np.matmul(
+        np.reshape(queries, (batch_size, num_heads,
+                             num_h_blocks, num_w_blocks,
+                             query_shape[0]*query_shape[1], depth)),
+        np.transpose(
+            np.reshape(keys, (batch_size, num_heads, num_h_blocks, num_w_blocks,
+                              query_shape[0]*query_shape[1], depth)),
+            (0, 1, 2, 3, 5, 4)))
+    # now to do a softmax across the logits
+    att = np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True)
+    att_output = np.matmul(att, np.reshape(
+        values, (batch_size, num_heads, num_h_blocks, num_w_blocks,
+                 query_shape[0]*query_shape[1], depth)))
+    att_output = np.reshape(att_output,
+                            (batch_size, num_heads, num_h_blocks, num_w_blocks,
+                             query_shape[0], query_shape[1], depth))
+    # putting the attention results back into the right place
+    out = np.zeros((batch_size, num_heads, new_height, new_width, depth))
+    for b in range(batch_size):
+      for h in range(num_heads):
+        for x in range(new_height):
+          for y in range(new_width):
+            h_block_index = x//query_shape[0]
+            w_block_index = y//query_shape[1]
+            inside_h_index = x%query_shape[0]
+            inside_w_index = y%query_shape[1]
+            out[b, h, x, y] = (
+                att_output[b, h, h_block_index, w_block_index, inside_h_index,
+                           inside_w_index])
+    out = np.transpose(out, (0, 2, 3, 1, 4))
+    out = np.reshape(out, (batch_size, new_height, new_width, total_depth))
+    out = out[:, :height, :width, :]
+
+    self.assertAllClose(res, out)
+
   def python_relative_att(self, q, k, v, batch, num_heads, height, width,
                           depth, height_key_relative_embeddings,
                           width_key_relative_embeddings,

From f41e517f813fd63ee65831de1a2522df17442efc Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Thu, 21 Mar 2019 21:16:11 +0000
Subject: [PATCH 1812/2720] Upgrade ML Engine runtime to 1.13 and remove unused
 Gcloud class (#1494)

This PR upgrades the ML Engine runtime to 1.13 in order to fix #1472. Unfortunately ML Engine doesn't support TPUs in version 1.13 yet: https://cloud.google.com/ml-engine/docs/release-notes

This also removes the unused Gcloud class to clean up the code a bit.
---
 .../wikisum/parallel_launch.py                |  2 +-
 tensor2tensor/utils/cloud_mlengine.py         | 57 ++-----------------
 2 files changed, 7 insertions(+), 52 deletions(-)

diff --git a/tensor2tensor/data_generators/wikisum/parallel_launch.py b/tensor2tensor/data_generators/wikisum/parallel_launch.py
index 194ee3dcf..08b3b172b 100644
--- a/tensor2tensor/data_generators/wikisum/parallel_launch.py
+++ b/tensor2tensor/data_generators/wikisum/parallel_launch.py
@@ -146,7 +146,7 @@ def create_instance(instance_name, cpu=1, mem=4):
 
 
 def list_vm_names_and_ips():
-  list_out = cloud.shell_output(cloud.Gcloud.LIST_VM)
+  list_out = cloud.shell_output(cloud.LIST_VM)
   lines = [l.split() for l in list_out.split("\n")[1:-1]]
   names_and_ips = [(l[0].strip(), l[-2].strip()) for l in lines]
   return names_and_ips
diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 8d2de835f..407a963cf 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -35,55 +35,10 @@
 FLAGS = tf.flags.FLAGS
 
 CONSOLE_URL = "https://console.cloud.google.com/mlengine/jobs/"
-RUNTIME_VERSION = "1.12"
-
-
-class Gcloud(object):
-  """gcloud command strings."""
-  # Note these can be modified by set_versions
-  VM_VERSION = "tf-1-12"
-  TPU_VERSION = "1.12"
-
-  @classmethod
-  def set_versions(cls, vm, tpu):
-    cls.VM_VERSION = vm
-    cls.TPU_VERSION = tpu
-
-  @classmethod
-  def create_vm(cls):
-    create_vm_str = """
-    gcloud compute instances create {name} \
-      --machine-type=n1-standard-8 \
-      --image-family=%s \
-      --image-project=ml-images \
-      --scopes=https://www.googleapis.com/auth/cloud-platform
-    """ % cls.VM_VERSION
-    return create_vm_str
-
-  DELETE_VM = "gcloud compute instances delete {name} --quiet"
-
-  @classmethod
-  def create_tpu(cls):
-    create_tpu_str = """
-    gcloud beta compute tpus create \
-      {name} \
-      --range={tpu_ip}/29 \
-      --version=%s
-    """ % cls.TPU_VERSION
-    return create_tpu_str
-
-  DELETE_TPU = "gcloud beta compute tpus delete {name} --quiet"
-
-  LIST_TPU = "gcloud beta compute tpus list"
-  LIST_VM = "gcloud compute instances list"
-
-  SSH_LOCAL_PORT_FORWARD = "-L {local_port}:{host}:{remote_port}"
-  SSH_TUNNEL = """
-  gcloud compute ssh {name} -- -N
-  """
-
-  DEFAULT_PROJECT = "gcloud config get-value project"
-  DEFAULT_REGION = "gcloud config get-value compute/region"
+RUNTIME_VERSION = "1.13"
+LIST_VM = "gcloud compute instances list"
+DEFAULT_PROJECT = "gcloud config get-value project"
+DEFAULT_REGION = "gcloud config get-value compute/region"
 
 
 def shell_output(cmd_, **kwargs):
@@ -99,11 +54,11 @@ def format_cmd(cmd_, **kwargs):
 
 
 def default_region():
-  return shell_output(Gcloud.DEFAULT_REGION).strip()
+  return shell_output(DEFAULT_REGION).strip()
 
 
 def default_project():
-  return shell_output(Gcloud.DEFAULT_PROJECT).strip()
+  return shell_output(DEFAULT_PROJECT).strip()
 
 
 def get_setup_file(name, packages=None):

From a6f8a00d96462892ac5a18260db2d425b7f52080 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 21 Mar 2019 14:29:44 -0700
Subject: [PATCH 1813/2720] Increase git depth, since there maybe a lag from
 when we trigger a build to other changes getting checked into the repo. In
 which case git checkout on a commit id (after cloning with the extra changes)
 fails with "fatal: reference is not a tree"

PiperOrigin-RevId: 239673455
---
 .travis.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.travis.yml b/.travis.yml
index 6484ecb17..dea932d82 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -2,7 +2,7 @@ sudo: required
 language: python
 cache: pip
 git:
-  depth: 1
+  depth: 3
   quiet: true
 services:
   - docker

From 0b17e188b47b50749f748c759500bd8abdd2da4c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 21 Mar 2019 15:31:46 -0700
Subject: [PATCH 1814/2720] Fix TransformerMemory model

Fixes a bug with training and enables relative attention. Absolute attention works poorly when the timing signal is reset at the start of each chunk.

PiperOrigin-RevId: 239686861
---
 tensor2tensor/layers/common_attention.py   | 26 ++++++++++++++++++++++
 tensor2tensor/layers/transformer_memory.py | 24 +++++++++++---------
 tensor2tensor/models/transformer.py        | 14 +++++++-----
 3 files changed, 49 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index cef04032b..a56962af8 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1667,6 +1667,20 @@ def dot_product_attention_relative(q,
     return _relative_attention_inner(weights, v, relations_values, False)
 
 
+def dot_product_attention_relative_memory(q, k, v, bias, *args, **kwargs):
+  """Wrapper of dot_product_attention_relative to use with recurrent memory."""
+
+  q_len = tf.shape(q)[2]
+  k_len = tf.shape(k)[2]
+  num_memory_items = k_len - q_len
+
+  q = tf.pad(q, [[0, 0], [0, 0], [num_memory_items, 0], [0, 0]])
+  bias = tf.pad(bias, [[0, 0], [0, 0], [num_memory_items, 0], [0, 0]])
+  output = dot_product_attention_relative(q, k, v, bias, *args, **kwargs)
+
+  return output[:, :, num_memory_items:, :]
+
+
 def _relative_position_to_absolute_position_masked(x):
   """Helper to dot_product_self_attention_relative_v2.
 
@@ -4152,6 +4166,18 @@ def multihead_attention(query_antecedent,
           save_weights_to=save_weights_to,
           make_image_summary=make_image_summary,
           cache=cache is not None)
+    elif attention_type == "dot_product_relative_memory":
+      x = dot_product_attention_relative_memory(
+          q,
+          k,
+          v,
+          bias,
+          max_relative_position,
+          dropout_rate,
+          image_shapes,
+          save_weights_to=save_weights_to,
+          make_image_summary=make_image_summary,
+          cache=cache is not None)
     elif attention_type == "dot_product_unmasked_relative_v2":
       x = dot_product_unmasked_self_attention_relative_v2(
           q,
diff --git a/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index 34f334dce..2d3696eac 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -90,17 +90,21 @@ def pre_attention(self, segment, query_antecedent, memory_antecedent, bias):
         [tf.stop_gradient(previous_vals), query_antecedent], 1)
     new_bias = tf.concat([previous_bias, bias], -1)
 
-    cancel_update = tf.equal(self.previous_segment, segment[0])
     remember_segment = segment[0]
-    remember_vals = tf.cond(
-        cancel_update,
-        lambda: self.previous_vals,
-        lambda: tf.pad(query_antecedent, [[0, amount_to_pad], [0, 0], [0, 0]]))
-    remember_bias = tf.cond(
-        cancel_update,
-        lambda: self.previous_bias,
-        lambda: tf.zeros_like(bias) + tf.reduce_max(bias, -1, keep_dims=True))
-
+    # TODO(kitaev): The code assumes that we always either increment the chunk
+    # number or reset it to zero, which is checked by the assertion. This
+    # assumption will not hold if we re-run the model for each token, e.g. for
+    # autoregressive greedy/beam/sampling decode.
+    with tf.control_dependencies(
+        [tf.Assert(tf.math.logical_or(
+            tf.equal(remember_segment, 0),
+            tf.equal(remember_segment, self.previous_segment + 1)),
+                   [self.previous_segment, remember_segment])]):
+      remember_segment = tf.identity(remember_segment)
+    remember_vals = tf.pad(query_antecedent,
+                           [[0, amount_to_pad], [0, 0], [0, 0]])
+    remember_bias = tf.zeros_like(bias) + tf.reduce_max(
+        bias, -1, keep_dims=True)
     token = (remember_segment, remember_vals, remember_bias)
 
     return token, query_antecedent, new_memory_antecedent, new_bias
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 6c47d2d53..e6d22b26e 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -2624,17 +2624,21 @@ def transformer_wikitext103_l4k_v0():
 
 
 @registry.register_hparams
-def transformer_wikitext103_l4k_memory():
+def transformer_wikitext103_l4k_memory_v0():
   """HParams for training languagemodel_wikitext103_l4k with memory."""
   hparams = transformer_wikitext103_l4k_v0()
 
-  hparams.split_targets_chunk_length = 8
-  hparams.split_targets_max_chunks = 512
+  hparams.split_targets_chunk_length = 64
+  hparams.split_targets_max_chunks = 64
 
   # The hparams specify batch size *before* chunking, but we want to have a
   # consistent 4K batch size *after* chunking to fully utilize the hardware.
   target_tokens_per_batch = 4096
-  hparams.batch_size = target_tokens_per_batch * (
-      hparams.max_length / hparams.split_targets_chunk_length)  # 2097152
+  hparams.batch_size = int(target_tokens_per_batch * (
+      hparams.max_length / hparams.split_targets_chunk_length))  # 262144
+
+  hparams.pos = None
+  hparams.self_attention_type = "dot_product_relative_memory"
+  hparams.max_relative_position = 2 * hparams.split_targets_chunk_length
 
   return hparams

From eedd6d76504dad6e485b2822e4189aee5c2cb939 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 21 Mar 2019 15:53:27 -0700
Subject: [PATCH 1815/2720] Bump setup.py version to 1.13.1

PiperOrigin-RevId: 239690861
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index adb181a12..0a37fcbff 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.13.0',
+    version='1.13.1',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',

From 0d840eeadcd33b31e07b650308788a121b85ddb1 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 21 Mar 2019 16:02:05 -0700
Subject: [PATCH 1816/2720] Make TransformerLM train reasonably well in trax.
 Adding loss and metric masking and dropout refactor in Transformer.

PiperOrigin-RevId: 239692595
---
 .../trax/configs/resnet50_imagenet_8gb.gin    |  2 +-
 .../trax/configs/transformer_lm1b_8gb.gin     | 18 ++++---
 tensor2tensor/trax/inputs.py                  |  6 ++-
 tensor2tensor/trax/models/transformer.py      | 53 ++++++++++---------
 tensor2tensor/trax/trax.py                    | 20 +++++--
 5 files changed, 60 insertions(+), 39 deletions(-)

diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
index 5c9bb365b..1f58f1f58 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
@@ -42,5 +42,5 @@ train.eval_steps = 20
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.Resnet50
 train.optimizer = @trax.optimizers.momentum
-train.train_steps = 500000
+train.train_steps = 1000000
 train.lr_schedule = @learning_rate.EvalAdjustingSchedule
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index 7a3f28c5a..caa237aed 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -5,28 +5,32 @@ import tensor2tensor.trax.trax
 
 # Parameters for batch_fun:
 # ==============================================================================
-batch_fun.batch_size = 32
-batch_fun.eval_batch_size = 32
+batch_fun.batch_size = 128
+batch_fun.eval_batch_size = 128
 
 # Parameters for inputs:
 # ==============================================================================
 inputs.data_dir = None
 inputs.dataset_name = 't2t_languagemodel_lm1b32k'
 
+# Parameters for mask:
+# ==============================================================================
+mask.mask_id = 0
+
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-MultifactorSchedule.constant = 0.05
+MultifactorSchedule.constant = 0.1
 MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
 MultifactorSchedule.warmup_steps = 8000
 
 # Parameters for preprocess_fun:
 # ==============================================================================
-preprocess_fun.max_target_length = 256
+preprocess_fun.max_target_length = 512
 
 # Parameters for train:
 # ==============================================================================
 train.eval_frequency = 1000
-train.eval_steps = 1
+train.eval_steps = 5
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
 train.run_debug_step = False
@@ -38,10 +42,10 @@ train_and_eval_batches.input_name = 'targets'
 
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.dropout = 0.1
+TransformerLM.dropout = 0.2
 TransformerLM.feature_depth = 512
 TransformerLM.feedforward_depth = 2048
-TransformerLM.max_len = 256
+TransformerLM.max_len = 512
 TransformerLM.mode = 'train'
 TransformerLM.num_heads = 8
 TransformerLM.num_layers = 6
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index a24cbb6ca..7a7a3c910 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -191,10 +191,12 @@ def batch_fun(dataset, training, shapes, target_names,
     if variable_target_shapes:
       bucket_boundaries = [bucket_length // 4, bucket_length // 2,
                            bucket_length, bucket_length * 2,
-                           bucket_length * 4, bucket_length * 8]
+                           bucket_length * 4, bucket_length * 8,
+                           bucket_length * 16]
       bucket_batch_sizes = [cur_batch_size * 4, cur_batch_size * 2,
                             cur_batch_size, cur_batch_size // 2,
-                            cur_batch_size // 4, cur_batch_size // 8, 1]
+                            cur_batch_size // 4, cur_batch_size // 8,
+                            max(1, cur_batch_size // 16), 1]
       buckets = (bucket_boundaries, bucket_batch_sizes)
 
   if buckets:
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 34d73e744..d2f0d40d0 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -29,7 +29,7 @@ def TransformerEncoder(mode='train',  # pylint: disable=invalid-name
                        feature_depth=512,
                        feedforward_depth=2048,
                        num_heads=8,
-                       dropout=0.9):
+                       dropout=0.1):
   """Transformer Encoder Stack.
 
   Args:
@@ -38,20 +38,22 @@ def TransformerEncoder(mode='train',  # pylint: disable=invalid-name
     feature_depth: int:  depth of embedding
     feedforward_depth: int: depth of feed-forward layer
     num_heads: int: number of attention heads
-    dropout: float: dropout rate - Stax follows TF's KEEP probability convention
+    dropout: float: dropout rate (how much to drop out; note that stax follows
+      Tensorflow's keep_rate convention, so we use 1 - dropout in calls below)
 
   Returns:
     A staxlayer for implementing a raw Transformer encoder stack.  No embedding
     or positional signals are added by this layer.
   """
+  keep_rate = 1.0 - dropout
   # Multi-headed Attention and Feed-forward layers
   multi_attention = stax.MultiHeadedAttention(
-      feature_depth, num_heads=num_heads, dropout=dropout, mode=mode)
+      feature_depth, num_heads=num_heads, dropout=keep_rate, mode=mode)
 
   feed_forward = stax.serial(
       stax.Dense(feedforward_depth, W_init=stax.xavier_uniform()),
       stax.Relu,
-      stax.Dropout(dropout, mode=mode),
+      stax.Dropout(keep_rate, mode=mode),
       stax.Dense(feature_depth, W_init=stax.xavier_uniform())
   )
 
@@ -74,11 +76,11 @@ def encoder(embedded_source, source_mask):
                                      stax.Identity,  # value
                                      source_mask),  # attention mask
                       multi_attention,
-                      stax.Dropout(dropout, mode=mode)),
+                      stax.Dropout(keep_rate, mode=mode)),
         # feed-forward
         stax.residual(stax.LayerNorm(feature_depth),
                       feed_forward,
-                      stax.Dropout(dropout, mode=mode))
+                      stax.Dropout(keep_rate, mode=mode))
     )
     return stax.serial(
         embedded_source,
@@ -95,8 +97,8 @@ def TransformerLM(vocab_size,  # pylint: disable=invalid-name
                   feature_depth=512,
                   feedforward_depth=2048,
                   num_heads=8,
-                  dropout=0.9,
-                  max_len=256):
+                  dropout=0.1,
+                  max_len=512):
   """Transformer language model (only uses the decoder part of Transformer).
 
   Args:
@@ -106,20 +108,21 @@ def TransformerLM(vocab_size,  # pylint: disable=invalid-name
     feature_depth: int:  depth of embedding
     feedforward_depth: int: depth of feed-forward layer
     num_heads: int: number of attention heads
-    dropout: float: dropout rate - Stax follows TF's KEEP probability convention
+    dropout: float: dropout rate (how much to drop out)
     max_len: int: maximum symbol length for positional encoding
 
   Returns:
     init and apply.
   """
+  keep_rate = 1.0 - dropout
   # Multi-headed Attention and Feed-forward layers
   multi_attention = stax.MultiHeadedAttention(
-      feature_depth, num_heads=num_heads, dropout=dropout, mode=mode)
+      feature_depth, num_heads=num_heads, dropout=keep_rate, mode=mode)
 
   feed_forward = stax.serial(
       stax.Dense(feedforward_depth, W_init=stax.xavier_uniform()),
       stax.Relu,
-      stax.Dropout(dropout, mode=mode),
+      stax.Dropout(keep_rate, mode=mode),
       stax.Dense(feature_depth, W_init=stax.xavier_uniform())
   )
 
@@ -132,18 +135,18 @@ def TransformerLM(vocab_size,  # pylint: disable=invalid-name
                                    stax.Identity,  # value
                                    stax.CausalMask(axis=-2)),  # attention mask
                     multi_attention,
-                    stax.Dropout(dropout, mode=mode)),
+                    stax.Dropout(keep_rate, mode=mode)),
       # feed-forward
       stax.residual(stax.LayerNorm(feature_depth),
                     feed_forward,
-                    stax.Dropout(dropout, mode=mode))
+                    stax.Dropout(keep_rate, mode=mode))
   )
 
   return stax.serial(
       stax.ShiftRight(),
       stax.Embedding(feature_depth, vocab_size),
       stax.PositionalEncoding(feature_depth, max_len=max_len),
-      stax.Dropout(dropout, mode=mode),
+      stax.Dropout(keep_rate, mode=mode),
       stax.repeat(decoder_layer, num_layers),
       stax.LayerNorm(feature_depth),
       stax.Dense(vocab_size, W_init=stax.xavier_uniform()),
@@ -158,7 +161,7 @@ def Transformer(source_vocab_size,  # pylint: disable=invalid-name
                 feature_depth=512,
                 feedforward_depth=2048,
                 num_heads=8,
-                dropout=0.9,
+                dropout=0.1,
                 shared_embedding=True,
                 max_len=200,
                 return_evals=False):
@@ -172,7 +175,7 @@ def Transformer(source_vocab_size,  # pylint: disable=invalid-name
     feature_depth: int:  depth of embedding
     feedforward_depth: int: depth of feed-forward layer
     num_heads: int: number of attention heads
-    dropout: float: dropout rate - Stax follows TF's KEEP probability convention
+    dropout: float: dropout rate (how much to drop out)
     shared_embedding: bool: specify whether source/target embeddings are tied.
     max_len: int: maximum symbol length for positional encoding
     return_evals: bool: whether to generate decode-time evaluation functions
@@ -182,11 +185,11 @@ def Transformer(source_vocab_size,  # pylint: disable=invalid-name
   the 'evals' functions that itself returns a namedtuple containing evaluation
   functions for the trained encoder, decoder, and generator substax.
   """
-
+  keep_rate = 1.0 - dropout
   # Input embedding and positional encoding
   inject_position = stax.serial(
       stax.PositionalEncoding(feature_depth, max_len=max_len),
-      stax.Dropout(dropout, mode=mode)
+      stax.Dropout(keep_rate, mode=mode)
   )
   if shared_embedding:
     assert source_vocab_size == target_vocab_size
@@ -202,12 +205,12 @@ def Transformer(source_vocab_size,  # pylint: disable=invalid-name
 
   # Multi-headed Attention and Feed-forward layers
   multi_attention = stax.MultiHeadedAttention(
-      feature_depth, num_heads=num_heads, dropout=dropout, mode=mode)
+      feature_depth, num_heads=num_heads, dropout=keep_rate, mode=mode)
 
   feed_forward = stax.serial(
       stax.Dense(feedforward_depth, W_init=stax.xavier_uniform()),
       stax.Relu,
-      stax.Dropout(dropout, mode=mode),
+      stax.Dropout(keep_rate, mode=mode),
       stax.Dense(feature_depth, W_init=stax.xavier_uniform())
   )
 
@@ -231,11 +234,11 @@ def encoder(source, source_mask):
                                      stax.Identity,  # value
                                      source_mask),  # attention mask
                       multi_attention,
-                      stax.Dropout(dropout, mode=mode)),
+                      stax.Dropout(keep_rate, mode=mode)),
         # feed-forward
         stax.residual(stax.LayerNorm(feature_depth),
                       feed_forward,
-                      stax.Dropout(dropout, mode=mode))
+                      stax.Dropout(keep_rate, mode=mode))
     )
     return stax.serial(
         source,
@@ -266,7 +269,7 @@ def decoder(memory, target, target_mask, memory_mask):
                                      stax.Identity,  # value
                                      target_mask),  # attention mask
                       multi_attention,
-                      stax.Dropout(dropout, mode=mode)),
+                      stax.Dropout(keep_rate, mode=mode)),
         # target attends to encoded source
         stax.residual(stax.LayerNorm(feature_depth),
                       stax.multiplex(stax.Identity,  # query
@@ -274,11 +277,11 @@ def decoder(memory, target, target_mask, memory_mask):
                                      memory,  # value
                                      memory_mask),  # attention mask
                       multi_attention,
-                      stax.Dropout(dropout, mode=mode)),
+                      stax.Dropout(keep_rate, mode=mode)),
         # feed-forward
         stax.residual(stax.LayerNorm(feature_depth),
                       feed_forward,
-                      stax.Dropout(dropout, mode=mode))
+                      stax.Dropout(keep_rate, mode=mode))
     )
     return stax.serial(
         target,
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 876ddc95d..d84a36275 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -46,26 +46,38 @@
 from tensorflow.io import gfile
 
 
+@gin.configurable(blacklist=["inputs", "targets"])
+def masked_mean(inputs, targets, mask_id=None):
+  """Mean of the inputs but counting only those where targets != mask_id."""
+  x = inputs.astype(np.float32)
+  if mask_id is None:
+    return np.mean(x)
+  unmask = 1.0 - np.equal(targets, mask_id).astype(np.float32)
+  return np.sum(x * unmask) / np.sum(unmask)
+
+
 def accuracy(batch, model_predictions):
   """Calculate accuracy."""
   _, targets = batch
   predicted_class = np.argmax(model_predictions, axis=-1)
-  return np.mean(predicted_class == targets)
+  correct = np.equal(predicted_class, targets)
+  return masked_mean(correct, targets)
 
 
 def neg_log_perplexity(batch, model_predictions):
   """Calculate negative log perplexity."""
   _, targets = batch
   hot_targets = stax.one_hot(targets, model_predictions.shape[-1])
-  return np.mean(np.sum(model_predictions * hot_targets, axis=-1))
+  xent = np.sum(model_predictions * hot_targets, axis=-1)
+  return masked_mean(xent, targets)
 
 
 def loss(params, batch, model_predict):
   """Calculate loss."""
   inputs, targets = batch
   preds = model_predict(params, inputs)
-  return - np.mean(np.sum(preds * stax.one_hot(targets, preds.shape[-1]),
-                          axis=-1))
+  xent = np.sum(preds * stax.one_hot(targets, preds.shape[-1]), axis=-1)
+  return - masked_mean(xent, targets)
 
 
 def log(s, stdout=True):

From 516d50eaffcff182f54fd91cf58a34ae038e0080 Mon Sep 17 00:00:00 2001
From: RJ Skerry-Ryan <rjryan@google.com>
Date: Fri, 22 Mar 2019 08:45:29 -0700
Subject: [PATCH 1817/2720] Use a TPU-compatible approach to matrix inversion,
 until tf.linalg.inv is supported.

PiperOrigin-RevId: 239801671
---
 tensor2tensor/models/research/glow_ops.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index daab40340..fdb272d2b 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -384,11 +384,12 @@ def invertible_1x1_conv(name, x, reverse=False):
       w = tf.reshape(w, [1, 1] + w_shape)
       x = tf.nn.conv2d(x, w, [1, 1, 1, 1], "SAME", data_format="NHWC")
     else:
-      u_inv = tf.matrix_inverse(u)
-      l_inv = tf.matrix_inverse(l)
-      p_inv = tf.matrix_inverse(p)
-      w_inv = tf.matmul(u_inv, tf.matmul(l_inv, p_inv))
-      w_inv = tf.reshape(w_inv, [1, 1]+w_shape)
+      # TODO(b/111271662): Remove when supported.
+      def tpu_inv(m):
+        """tf.linalg.inv workaround until it is supported on TPU."""
+        q, r = tf.linalg.qr(m)
+        return tf.linalg.triangular_solve(r, tf.transpose(q), lower=False)
+      w_inv = tf.reshape(tpu_inv(w), [1, 1]+w_shape)
       x = tf.nn.conv2d(
           x, w_inv, [1, 1, 1, 1], "SAME", data_format="NHWC")
       objective *= -1

From db2bbe48fec5a0ac71b5a38b9b35c1e648022f7f Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 22 Mar 2019 10:57:08 -0700
Subject: [PATCH 1818/2720] Update RL notebook and add a link to it in
 rl/README. Great thanks to Piotr Kozakowski for the colab!

PiperOrigin-RevId: 239825106
---
 tensor2tensor/notebooks/hello_t2t-rl.ipynb | 355 +--------------------
 tensor2tensor/rl/README.md                 |  16 +-
 2 files changed, 11 insertions(+), 360 deletions(-)

diff --git a/tensor2tensor/notebooks/hello_t2t-rl.ipynb b/tensor2tensor/notebooks/hello_t2t-rl.ipynb
index ff9fc4615..21ca2909c 100644
--- a/tensor2tensor/notebooks/hello_t2t-rl.ipynb
+++ b/tensor2tensor/notebooks/hello_t2t-rl.ipynb
@@ -1,354 +1 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 0,
-   "metadata": {
-    "colab": {
-     "autoexec": {
-      "startup": false,
-      "wait_interval": 0.0
-     }
-    },
-    "colab_type": "code",
-    "collapsed": true,
-    "id": "s19ucTii_wYb"
-   },
-   "outputs": [],
-   "source": [
-    "# Copyright 2018 Google LLC.\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",
-    "# https://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."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "collapsed": true
-   },
-   "outputs": [],
-   "source": [
-    "# Install deps\n",
-    "!pip install -q -U tensor2tensor tensorflow"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {
-    "colab": {
-     "autoexec": {
-      "startup": false,
-      "wait_interval": 0.0
-     }
-    },
-    "colab_type": "code",
-    "collapsed": true,
-    "id": "oILRLCWN_16u"
-   },
-   "outputs": [],
-   "source": [
-    "import tensorflow as tf\n",
-    "import matplotlib.pyplot as plt\n",
-    "import numpy as np\n",
-    "import os\n",
-    "import collections\n",
-    "import sys\n",
-    "import tempfile\n",
-    "\n",
-    "from tensor2tensor import models\n",
-    "from tensor2tensor import problems\n",
-    "from tensor2tensor.rl import rl_trainer_lib\n",
-    "from tensor2tensor.utils import trainer_lib\n",
-    "from tensor2tensor.utils import t2t_model\n",
-    "from tensor2tensor.utils import registry\n",
-    "\n",
-    "# Other setup\n",
-    "Modes = tf.estimator.ModeKeys\n",
-    "\n",
-    "prefix = \"~/t2t_rl_data\"\n",
-    "# Setup data directories\n",
-    "data_dir = os.path.expanduser(prefix + \"/data\")\n",
-    "tmp_dir = os.path.expanduser(prefix + \"/tmp\")\n",
-    "tf.gfile.MakeDirs(data_dir)\n",
-    "tf.gfile.MakeDirs(tmp_dir)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Train policy\n",
-    "\n",
-    "The training of the policy will take around 1h on GPU."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "INFO:tensorflow:Overriding hparams in ppo_atari_base with epochs_num=1\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "[2018-03-07 00:06:25,614] Overriding hparams in ppo_atari_base with epochs_num=1\n",
-      "[2018-03-07 00:06:25,620] Making new env: PongNoFrameskip-v4\n",
-      "[2018-03-07 00:06:25,860] Making new env: PongNoFrameskip-v4\n",
-      "[2018-03-07 00:06:25,865] Making new env: PongNoFrameskip-v4\n",
-      "[2018-03-07 00:06:25,872] Making new env: PongNoFrameskip-v4\n",
-      "[2018-03-07 00:06:25,883] Making new env: PongNoFrameskip-v4\n",
-      "[2018-03-07 00:06:25,892] Making new env: PongNoFrameskip-v4\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "WARNING:tensorflow:From /home/blazej.osinski/t2t/t2t_jupyter_kernel/local/lib/python2.7/site-packages/tensorflow/python/ops/distributions/categorical.py:310: calling argmax (from tensorflow.python.ops.math_ops) with dimension is deprecated and will be removed in a future version.\n",
-      "Instructions for updating:\n",
-      "Use the `axis` argument instead\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "[2018-03-07 00:06:26,589] From /home/blazej.osinski/t2t/t2t_jupyter_kernel/local/lib/python2.7/site-packages/tensorflow/python/ops/distributions/categorical.py:310: calling argmax (from tensorflow.python.ops.math_ops) with dimension is deprecated and will be removed in a future version.\n",
-      "Instructions for updating:\n",
-      "Use the `axis` argument instead\n",
-      "[2018-03-07 00:06:27,589] Making new env: PongNoFrameskip-v4\n",
-      "[2018-03-07 00:06:27,772] Making new env: PongNoFrameskip-v4\n",
-      "[2018-03-07 00:06:47,971] Starting new video recorder writing to /home/blazej.osinski/t2t_rl_data/data/ppo_Yr5Rjt/openaigym.video.0.144364.video000000.mp4\n",
-      "[2018-03-07 00:09:36,335] Finished writing results. You can upload them to the scoreboard via gym.upload('/home/blazej.osinski/t2t_rl_data/data/ppo_Yr5Rjt')\n"
-     ]
-    }
-   ],
-   "source": [
-    "iteration_num=300\n",
-    "hparams = trainer_lib.create_hparams(\"ppo_atari_base\", \"epochs_num={}\".format(iteration_num+1))\n",
-    "ppo_dir = tempfile.mkdtemp(dir=data_dir, prefix=\"ppo_\")\n",
-    "rl_trainer_lib.train(hparams, \"stacked_pong\", ppo_dir)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {
-    "collapsed": true
-   },
-   "outputs": [],
-   "source": [
-    "agent_policy_path = os.path.join(ppo_dir, \"model{}.ckpt.index\".format(iteration_num))[:-6]"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Generate and review frames from policy"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {
-    "collapsed": true
-   },
-   "outputs": [],
-   "source": [
-    "sys.argv = [sys.argv[0], \"--agent_policy_path\", agent_policy_path]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {
-    "collapsed": true
-   },
-   "outputs": [],
-   "source": [
-    "tf.reset_default_graph()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 14,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "[2018-03-07 00:13:20,983] Making new env: PongNoFrameskip-v4\n",
-      "[2018-03-07 00:13:21,221] Making new env: PongNoFrameskip-v4\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "INFO:tensorflow:Restoring parameters from /home/blazej.osinski/t2t_rl_data/data/ppo_Yr5Rjt/model0.ckpt\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "[2018-03-07 00:13:21,586] Restoring parameters from /home/blazej.osinski/t2t_rl_data/data/ppo_Yr5Rjt/model0.ckpt\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "INFO:tensorflow:Generated 4998 Examples\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "[2018-03-07 00:31:57,314] Generated 4998 Examples\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "INFO:tensorflow:Shuffling data...\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "[2018-03-07 00:31:57,319] Shuffling data...\n"
-     ]
-    }
-   ],
-   "source": [
-    "# This step is also time consuming - takes around 30 minutes.\n",
-    "gym_problem = problems.problem(\"gym_pong_trajectories_from_policy\")\n",
-    "gym_problem.generate_data(data_dir, tmp_dir)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "INFO:tensorflow:Reading data files from /home/blazej.osinski/t2t_rl_data/data/gym_pong_trajectories_from_policy-train*\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "[2018-03-07 00:32:00,394] Reading data files from /home/blazej.osinski/t2t_rl_data/data/gym_pong_trajectories_from_policy-train*\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "INFO:tensorflow:partition: 0 num_data_files: 10\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "[2018-03-07 00:32:00,399] partition: 0 num_data_files: 10\n"
-     ]
-    },
-    {
-     "data": {
-      "image/png": "iVBORw0KGgoAAAANSUhEUgAABG4AAARUCAYAAAAnLFWhAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMS4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvNQv5yAAAIABJREFUeJzs3U+IrNd95//P+VmJF4qQrVjIQtZEWWjj3UhNIjGzCISA7FrIKxMvxsK00CaBBGYjEkQW2nhmERgzwSCmTcsQHAwJWNAagiMCnsV18E0Y/BdbckBYRrZ0ZVAu8iIJc36L+9ykfbvrdld1VT3nnH69oOnq59btOn7ePt19v3q6qtRaAwAAAEB7/r+5FwAAAADA6QxuAAAAABplcAMAAADQKIMbAAAAgEYZ3AAAAAA0yuAGAAAAoFFbGdyUUp4opXy/lPJaKeXZbTwG26dj/zQcg47903AMOvZPwzHo2D8Nx6Dj7pRa62Y/YSnvS/KDJL+T5I0k30jyqVrrdzf6QGyVjv3TcAw69k/DMejYPw3HoGP/NByDjru1jStufiPJa7XWf6y1/nOSv0jy5BYeh+3SsX8ajkHH/mk4Bh37p+EYdOyfhmPQcYfu2MLnfCDJj459/EaS37z1TqWUZ5I8kyTvf//7H/3whz+89gPefffdJ469++67a3++i3rvvfeSJHfeeefan+Odd97J9evXy6bWtIZZOs7Z7VYDdLzUe/Fmv+PWaTnyXrzZ69ZGx/fisvvskr243O0a3krDC7vUX1OTITpe6oa+L47RMbEX19FSw8u0F7f182kLNrEPk+T111+/Vmu996z7bWNwcy611heSvJAkDz30UH3uuefW/lyLxeLEsaOjo7U/30VduXIlSfL444+v/Tmef/75TS1nqzbdcc5ut7osHUfdizf7HbdOyx4aJut1vNnr1kbH9+Ky++ySvbjc7RreSsPdGPVranJ5Oo7a0PfFMTom9uI6Wmp4mfbitn4+bcEm9mGSPP3006+f537bGNz8OMmDxz7+yHRsq279RwYXNktHNmrWvdiCm19IT/sG2RF7sX+Xfi8Ows83/dt5w+P7sIWGvi+up7WOA7j0X0/tRVa1jee4+UaSh0spv15K+eUkv5vkpS08DtulY/80HIOO/dNwDDr2T8Mx6Ng/Dceg4w5t/IqbWuu/llJ+P8lfJ3lfki/UWr+z6cdhu+bqeOsE3H8xXp+GY2ilI+trpaG9eDF+vulfK3uRi7EX+6fhGOboeJl/ttnKc9zUWl9O8vI2Pje7o2P/NByDjv3TcAw69k/DMejYPw3HoOPuzPbkxHCaln73lNX5HfAx6Ng/DaEdfrYB2IzL/PV0mMHNZYw3Ih371lq/zp/wDdbW2l5kfVr2rbV+vi9ymbW0H+3F9bTUcNe28eTEAAAAAGzAEFfctPakRBd9LffLSsf+tdRQv7Odp1cLTbVcblmfFrodp+F6dOxfSw3124wWmmq5uha63aTf7fnZ5nSl1rrTBzzN3t5evXr16tzLaMre3l6uXr1a5l7HKnQ8qbeOGp7UW8NEx9P01lHDk3prmOh4mt46anhSbw0THU/TW0cNT9JwDKWUv6+17p11P78qBQAAANAogxsAAACARhncAAAAADTK4AYAAACgUU28qtS1a9dycHAw9zKacu3atbmXsDIdT+qto4Yn9dYw0fE0vXXU8KTeGia377hYLJK09yoZ29ZbR3vxpN4aJjqepreOGp6k4eXiihsA6Mhisfi3f/RfFvv7+3MvAQBgNk1ccQMAsMyI/4Xusl1pAwCszxU3AAAAAI0yuAG4RC7Lr9n41RoAAEZhcAMAAADQKIMbAIYz4nOiAABwORncAAAAADTKq0oBQEe8GhEAwOXiihsAAACARhncAAAAADTKr0oBXCJ+zQYAAPriihsAAACARhncAAAAADTK4AYAAACgUQY3AAAAAI0yuAEAAABolMENAAAAQKMMbgAAAAAaZXADAAAA0CiDGwAAAIBGGdwAAAAAs1gsFlksFnMvo2kGNwAAAACn2N/fn3sJBjcAAAAArTK4AQAAADjFwcHB3EswuAEAAABo1ZmDm1LKF0opb5VSvn3s2D2llK+WUl6d3n9wOl5KKZ8rpbxWSvlmKeWRbS6e89OxfxqOQcf+aTgGHfun4Rh07J+GY9Cxbee54uYwyRO3HHs2ySu11oeTvDJ9nCQfS/Lw9PZMks9vZplswGF07N1hNNy4K1eu5MqVK7t8yMPo2LvDaDiCw+jYu8NouHG+L7KGw2i4cfYix505uKm1fi3Jz245/GSSF6fbLyb5xLHjX6w3fD3JB0op929qsaxPx/5pOAYd+6fhGHTsn4Zj0LF/Go5hzo5HR0c5Ojpa969fCus+x819tdY3p9s/SXLfdPuBJD86dr83pmMnlFKeKaVcLaVcvX79+prL4IJ07J+GY9CxfxqOQcf+aTgGHfun4Rgu1FHDzbnwkxPXWmuSusbfe6HWuldr3bvrrrsuugwuSMf+aTgGHfun4Rh07J+GY9CxfxqOYZ2OGm7OuoObn968FGp6/9Z0/MdJHjx2v49Mx2iTjv3TcAw69k/DMejYPw3HoGP/NByDjo1Yd3DzUpKnpttPJfnKseOfnp5l+rEk7x67tIr26Ng/DcegY/80HIOO/dNwDDr2T8Mx6NiIO866QynlS0l+K8mHSilvJPmTJJ9N8uVSyn6S15N8crr7y0k+nuS1JD9P8pktrJk16Ng/DcegY/80HIOO/dNwDDr2T8Mx6Ni2Mwc3tdZPLfmj3z7lvjXJ7110UZfBzZd2e/zxx3fyeHN3XCwWSbKzZwvf39/PwcHB1h9nlx3nbjiqXe3Bm3TcDnuxf5ft++Ko7MX++b44Bnuxf/Zi/za5Dy/85MQAAAAAbIfBDUPaxdU2AADAvBaLxb9d3b9N+/v7W38MWMbgBgAAAKBRBjcNML0FAABolyv6mZPBDQAAAECjznxVKbbj+DNLm972a9fP9g6czl7sn4Zj0BHaYC/C/Da5D5sf3Oz6ZaTZDv0AAABgdX5VCgAAAKBRzV9xAwAAAKdxZT+XgStuAAAAABplcAMAAADQKIMbAAAAgEYZ3AAAAAA0qvknJ/ZkUwAAAMBl5YobAAAAgEYZ3AAAAAA0yuAGAAAAoFEGNwAAsKLFYpHFYjH3MgC4BAxuAADgktjf3597CQCsyOAGAAAuiYODg7mXAMCKDG4AAAAAGmVwAwAAANAogxsAAACARhncAAAAADTqjrkXAAAAvTk6Opp7CQBcEq64AQAAAGiUwQ0AAABAowxuAAAAABplcAOcy2KxyGKxmHsZW7W/vz/3EgAAAH6BwQ0AAABAowxuACYHBwdzLwEAAOAXGNwAAAAANMrgBgAAAKBRBjcAAAAAraq13vYtyYNJ/jbJd5N8J8kfTMfvSfLVJK9O7z84HS9JPpfktSTfTPLIWY/x6KOPVn7RdE7O7HOet100rDqeqreOGp7UW8Oq46l666jhSb01rDqeqreOGp7UW8Oq46l666jhSRqOIcnVeo5G57ni5l+T/Nda60eTPJbk90opH03ybJJXaq0PJ3ll+jhJPpbk4entmSSfP8djsF0ajkHH/mk4Bh37p+EYdOyfhmPQsX8aNu7MwU2t9c1a6z9Mt68n+V6SB5I8meTF6W4vJvnEdPvJJF+cBkhfT/KBUsr9G18556bhGHTsn4Zj0LF/Go5Bx/5pOAYd+6dh+1Z6jptSykNJ/mOSv0tyX631zemPfpLkvun2A0l+dOyvvTEdu/VzPVNKuVpKufr222+vuGzWtcmG0+fTcQb2Yv/sxTHYi/2zF8dgL/bPXhyDvdg/Ddt07sFNKeVXkvxlkj+stf7T8T+bfjerrvLAtdYXaq17tda9e++9d5W/ypo23XD6ezrumL3YP3txDPZi/+zFMdiL/bMXx2Av9k/Ddp1rcFNK+aXcCPjntda/mg7/9OblUNP7t6bjP86NJze66SPTMWak4Rh07J+GY9CxfxqOQcf+aTgGHfunYdvOHNyUUkqSgyTfq7X+6bE/einJU9Ptp5J85djxT5cbHkvy7rHLq5iBhmPQsX8ajkHH/mk4Bh37p+EYdOyfhu0rN654us0dSvnPSf5Pkm8l+X/T4T/Kjd95+3KS/5Dk9SSfrLX+bIr+P5M8keTnST5Ta716xmNcT/L9C/zv2LUPJbm25cf4tVrrRq4n20XD6XF66riLhklnHTtrmNiLyx6np4724umP0VPDxF5c9jg9dbQXT3+Mnhom9uKyx+mpo714+mP01DCxF097jLeTvJfd/P97E5rai2cObnahlHK11ro39zrOq7f17kpP56Wnte5Sb+elt/XuSk/npae17lJv56W39e5KT+elp7XuUm/npbf17kpP56Wnte5Sb+elt/XuSk/npbW1rvSqUgAAAADsjsENAAAAQKNaGdy8MPcCVtTbenelp/PS01p3qbfz0tt6d6Wn89LTWnept/PS23p3pafz0tNad6m389Lbenelp/PS01p3qbfz0tt6d6Wn89LUWpt4jhsAAAAATmrlihsAAAAAbmFwAwAAANCo2Qc3pZQnSinfL6W8Vkp5toH1fKGU8lYp5dvHjt1TSvlqKeXV6f0Hp+OllPK5ae3fLKU8Mt/K59Naw0THdbTWUcPVtdYw0XEdrXXUcHWtNUx0XEdrHTVcXWsNEx3X0VpHDVfXWsOkv46zDm5KKe9L8mdJPpbko0k+VUr56JxrSnKY5Ilbjj2b5JVa68NJXpk+Tm6s++Hp7Zkkn9/RGpvRaMNEx5U02vEwGp5bow0THVfSaMfDaHhujTZMdFxJox0Po+G5Ndow0XEljXY8jIbn1mjDpLOOc19x8xtJXqu1/mOt9Z+T/EWSJ+dcUK31a0l+dsvhJ5O8ON1+Mcknjh3/Yr3h60k+UEq5fzcrbUZzDRMd19BcRw1X1lzDRMc1NNdRw5U11zDRcQ3NddRwZc01THRcQ3MdNVxZcw2T/jrOPbh5IMmPjn38xnSsNffVWt+cbv8kyX3T7V7Wv009nQMdl+vlHGi4XE/nQMflejkHGi7X0znQcblezoGGy/V0DnRcrpdzoOFyPZ2DZjvOPbjpTr3x+uleQ71zOvZPwzHo2D8Nx6Bj/zQcg47903AMrXWce3Dz4yQPHvv4I9Ox1vz05qVQ0/u3puO9rH+bejoHOi7XyznQcLmezoGOy/VyDjRcrqdzoONyvZwDDZfr6RzouFwv50DD5Xo6B812nHtw840kD5dSfr2U8stJfjfJSzOv6TQvJXlquv1Ukq8cO/7p6VmmH0vy7rFLqy6LXhomOt5OLx01XK6XhomOt9NLRw2X66VhouPt9NJRw+V6aZjoeDu9dNRwuV4aJi13rLVu/C03np35+0leS/LsGff9eJIfJPlhkj/exnpWXPuXkryZ5F9y43fX9pP8am48q/SrSf4myT3TfUtuPEP2D5N8K8ne3Oufo2NrDXVcvWGLHTVcvWNrDXVcvWGLHTVcvWNrDXVcvWGLHTVcvWNrDXVcvWGLHTVcvWNrDXvsWKaFbMz0cl8/SPI70wn4RpJP1Vq/u9EHYqt07J+GY9CxfxqOQcf+aTgGHfun4Rh03K1t/KpUky/3xcp07J+GY9CxfxqOQcf+aTgGHfun4Rh03KE7tvA5T3uprN+83V+466676q/+6q+e+YnvvvvuJMm777576vHjbr3Prrz33nsnjt15550rf5533nkn169fL5tY05q21nEZHTfuUu/F5N87rtPupsu6F292W9Z6V+zF9bS0FwdpmOh44liHHTW8RYcNk5l+vpnz55lbDfDzzaXei8kQDZMVO/p3xulef/31a7XWe8+63zYGN+dSSnkmyTNJcs899+S555478+8sFoskydHR0anHj7v1Prty5cqVE8cef/zxlT/P888/v4nlbN06HZfRcR6j7sXk3zuu0+6mHhomm9+LN7sta70r9uJ6WtqLl6lhouNZeuio4e310DDZ/M83c/48c6vL8vPNqHsx0fB2LtO/M5Lk6aeffv0899vG4OZcL5VVa30hyQtJ8tBDD134iXZu/UfGnG7GO+0bZEd07L/jrA3ZmFk6tsReXI+vpxunY/8dNey/YeL74ggu/V4cxJkdN9nw+L8xLmPDbTzHTU8v98VyOvZPwzHo2D8Nx6Bj/zQcg47903AMOu7Qxq+4qbX+aynl95P8dZL3JflCrfU7m34ctkvH/s3V8NYJuCtwLqaVjqzP19Mx6Ng/DcfQyvdFP9+sz14cw647XvafTbfyHDe11peTvLyNz83u6Ng/DcegY/80HIOO/dNwDDr2T8Mx6Lg7sz05MTCey/67p6PQEQB+kedHgfld5ivdhhnctPRFtPMnfJuVjn1rqR/jsBfX09J+1HB9OvZPwzG01JH1aNi/y/wri9t4cmIAAAAANqCrK26WTdRamrRd9HXcLzMd+9HDXkx03IS5m2q4nrm7Hafh+nTsn4Z98fPNuDTsh314ulLrhV5OfSP29vbq1atX515GU/b29nL16tUy9zpWoeNJvXXU8KTeGiY6nqa3jhqe1FvDRMfT9NZRw5N6a5joeJreOmp4koZjKKX8fa1176z7+VUpAAAAgEYZ3AAAAAA0yuAGAAAAoFEGNwAAAACNauJVpa5du5aDg4O5l9GUa9euzb2Elel4Um8dNTypt4aJjqfpraOGJ/XWMNHxNL111PCk3homOp6mt44anqTh5eKKGwAAAIBGGdwAAAAANMrgBgAAAKBRBjcAAAAAjTK4AQAAAGiUwQ0AAABAowxuAAAAABplcAMAAADQKIMbAAAAgEYZ3AAAAAA0yuAGAAAAoFEGNwAAAHRpsVhksVjMvQzYKoMbAAAAgEYZ3AAAAAA0yuAGAACAS2d/f3/uJcC5GNwAAAAANMrgBgAAgEvn4OBg7iXAuRjcAAAAADTqjrkXAAAAAOs4Ojqaewmwda64AYCOLBaLLBaLuZcBAMCOGNwAAAAANMrgBgAAAKBRBjcAAAAAjTK4AQC2an9/f+4lAAB0y+AGANiqg4ODuZcAANCtMwc3pZQvlFLeKqV8+9ixe0opXy2lvDq9/+B0vJRSPldKea2U8s1SyiPbXDznp2P/NByDjv3TcAw69k/DMejYPw3HoGPbznPFzWGSJ2459mySV2qtDyd5Zfo4ST6W5OHp7Zkkn9/MMsdz5cqVXLlyZZcPeRgdN27HHQ+j4QgOo+PGXaa9eHR0lKOjo4t+mub4vjiGy7QXR2UvsobDaDiCw+jYrDMHN7XWryX52S2Hn0zy4nT7xSSfOHb8i/WGryf5QCnl/k0tlvXp2D8Nx6Bj/zQcg47903AMOvZPwzHo2LZ1n+Pmvlrrm9PtnyS5b7r9QJIfHbvfG9OxE0opz5RSrpZSrl6/fn3NZXBBOvZPwzHo2D8Nx6Bj/zQcg47903AMF+qo4eZc+MmJa601SV3j771Qa92rte7dddddF10GF6Rj/zQcg47903AMOvZPwzHo2D8Nx7BORw03Z93BzU9vXgo1vX9rOv7jJA8eu99HpmO0Scf+aTgGHfun4Rh07J+GY9CxfxqOQcdGrDu4eSnJU9Ptp5J85djxT0/PMv1YknePXVpFe3Tsn4Zj0LF/Go5Bx/5pOAYd+6fhGHRsxB1n3aGU8qUkv5XkQ6WUN5L8SZLPJvlyKWU/yetJPjnd/eUkH0/yWpKfJ/nMFtbMGnTsn4Zj0LF/Go5Bx/5pOAYd+6fhGHRs25mDm1rrp5b80W+fct+a5PcuuqjL4PHHH9/p4+m4HbvsqOF23HzJ01211HE77MX++b44Bnuxf/biGHb5842G2+Fn1P5tsuGFn5wYAAAAgO0wuAEAAABoVPODm8VikcViMfcyAAAAAHau+cENAAAAwGVlcAMAAADQqDNfVQpgZLt+9QwAgG3z803/NOzfJhu64gYAAACgUQY3AAAAAI0yuAG4RFp7pb79/f25lwAAAE1r/jlujo6O5l4CAFtycHAw9xIAAKBprrhhJ1r7r/wAAADQA4MbAAAAgEYZ3AAAAAA0yuAGAAAAoFEGNwAAAACNav5VpQDYHK/UBwAAfXHFDQAAAECjXHHDTviv/AAAALA6V9wAAAAANMrgBgAAAKBRBjcAAAAAjTK4AQAAAGiUwQ0AAABAowxuAABgRYvFIovFIkmyv78/82oAGJnBDQAAAECjDG4AAOACDg4O5l4CAAMzuAEAAABolMENAAAAQKMMbgAAAAAaZXADAAAA0Kg75l4AAAD05ujoaO4lAHBJuOIGAAAAoFEGNwAAAACNMrgBAAAAaFWt9bZvSR5M8rdJvpvkO0n+YDp+T5KvJnl1ev/B6XhJ8rkkryX5ZpJHznqMRx99tPKLpnNyZp/zvO2iYdXxVL111PCk3hpWHU/VW0cNT+qtYdXxVL111PCk3hpWHU/VW0cNT9JwDEmu1nM0Os8VN/+a5L/WWj+a5LEkv1dK+WiSZ5O8Umt9OMkr08dJ8rEkD09vzyT5/Dkeg+3ScAw69k/DMejYPw3HoGP/NByDjv3TsHFnDm5qrW/WWv9hun09yfeSPJDkySQvTnd7McknpttPJvniNED6epIPlFLu3/jKOTcNx6Bj/zQcg47903AMOvZPwzHo2D8N27fSc9yUUh5K8h+T/F2S+2qtb05/9JMk9023H0jyo2N/7Y3p2K2f65lSytVSytW33357xWWzrk02nD6fjjOwF/tnL47BXuyfvTgGe7F/9uIY7MX+adimcw9uSim/kuQvk/xhrfWfjv/Z9LtZdZUHrrW+UGvdq7Xu3Xvvvav8Vda06YbT39Nxx+zF/tmLY7AX+2cvjsFe7J+9OAZ7sX8atutcg5tSyi/lRsA/r7X+1XT4pzcvh5revzUd/3FuPLnRTR+ZjjEjDcegY/80HIOO/dNwDDr2T8Mx6Ng/Ddt25uCmlFKSHCT5Xq31T4/90UtJnppuP5XkK8eOf7rc8FiSd49dXsUMNByDjv3TcAw69k/DMejYPw3HoGP/NGxfuXHF023uUMp/TvJ/knwryf+bDv9RbvzO25eT/Ickryf5ZK31Z1P0/5nkiSQ/T/KZWuvVMx7jepLvX+B/x659KMm1LT/Gr9VaN3I92S4aTo/TU8ddNEw669hZw8ReXPY4PXW0F09/jJ4aJvbissfpqaO9ePpj9NQwsReXPU5PHe3F0x+jp4aJvXjaY7yd5L3s5v/fm9DUXjxzcLMLpZSrtda9uddxXr2td1d6Oi89rXWXejsvva13V3o6Lz2tdZd6Oy+9rXdXejovPa11l3o7L72td1d6Oi89rXWXejsvva13V3o6L62tdaVXlQIAAABgdwxuAAAAABrVyuDmhbkXsKLe1rsrPZ2Xnta6S72dl97Wuys9nZee1rpLvZ2X3ta7Kz2dl57Wuku9nZfe1rsrPZ2Xnta6S72dl97Wuys9nZem1trEc9wAAAAAcFIrV9wAAAAAcAuDGwAAAIBGzT64KaU8UUr5finltVLKsw2s5wullLdKKd8+duyeUspXSymvTu8/OB0vpZTPTWv/ZinlkflWPp/WGiY6rqO1jhqurrWGiY7raK2jhqtrrWGi4zpa66jh6lprmOi4jtY6ari61hom/XWcdXBTSnlfkj9L8rEkH03yqVLKR+dcU5LDJE/ccuzZJK/UWh9O8sr0cXJj3Q9Pb88k+fyO1tiMRhsmOq6k0Y6H0fDcGm2Y6LiSRjseRsNza7RhouNKGu14GA3PrdGGiY4rabTjYTQ8t0YbJp11nPuKm99I8lqt9R9rrf+c5C+SPDnngmqtX0vys1sOP5nkxen2i0k+cez4F+sNX0/ygVLK/btZaTOaa5jouIbmOmq4suYaJjquobmOGq6suYaJjmtorqOGK2uuYaLjGprrqOHKmmuY9Ndx7sHNA0l+dOzjN6Zjrbmv1vrmdPsnSe6bbvey/m3q6RzouFwv50DD5Xo6Bzou18s50HC5ns6Bjsv1cg40XK6nc6Djcr2cAw2X6+kcNNtx7sFNd+qN10/3Guqd07F/Go5Bx/5pOAYd+6fhGHTsn4ZjaK3j3IObHyd58NjHH5mOteanNy+Fmt6/NR3vZf3b1NM50HG5Xs6Bhsv1dA50XK6Xc6Dhcj2dAx2X6+UcaLhcT+dAx+V6OQcaLtfTOWi249yDm28kebiU8uullF9O8rtJXpp5Tad5KclT0+2nknzl2PFPT88y/ViSd49dWnVZ9NIw0fF2eumo4XK9NEx0vJ1eOmq4XC8NEx1vp5eOGi7XS8NEx9vppaOGy/XSMGm5Y6111rckH0/ygyQ/TPLHDaznS0neTPIvufG7a/tJfjU3nlX61SR/k+Se6b4lN54h+4dJvpVkb+71a6jjKB017L+hjmN01LD/hjqO0VHD/hvqOEZHDftv2GPHMi1ko0opTyT5H0nel+R/1Vo/u/EHYet07J+GY9CxfxqOQcf+aTgGHfun4Rh03J2ND26m12n/QZLfyY3J1TeSfKrW+t2NPhBbpWP/NByDjv3TcAw69k/DMejYPw3HoONubeM5bpp8nXZWpmP/NByDjv3TcAw69k/DMejYPw3HoOMO3bGFz3naa5z/5q13KqU8k+SZJHn/+9//6Ic//OG1H/Duu+8+cezdd99d+/Nd1HvvvZckufPOO9f+HO+8806uX79eNrWmNejYf8dZGt5sdrPnXA1v9jtunZaXdS+20jGxF9dxvOHc7MUxOib24jr8bLMVOvbfcWsNl/3c4mebrTizo314ttdff/1arfXes+63jcHNudRaX0jyQpI89NBD9bnnnlv7cy0WixPHjo6O1v58F3XlypUkyeOPP77253j++ec3tZyt0vH2eui46YY3m93sOVfDm/2OW6dlDw2TcTsm9uI6jjecm704RsfEXlyHn23mo+Pt9dBxnYbLfm7xs8087MOzPf3006+f537bGNzM8hrnt25ELkzH/s3SsCU3v5Ce9g/Hjlz6jgO49A3tRRqx84bH/1HhZ5uN8TNq/3w9HcNOO172r6fbeI6bnl6nneV07J+GY9CxfxqOQcf+aTgGHfun4Rh03KGNX3FTa/3XUsrvJ/nr3HhZsC/UWr+z6ce51WWcum3TXB3ZHHsjl20yAAAgAElEQVRxDDr2r5WGLf3KTY907F8rDbkYHfun4Rh23fGy99vKc9zUWl9O8vI2Pje7o2P/NByDjv3TcAw69k/DMejYPw3HoOPuzPbkxJvm905hfpf9d09HoeMYfF8cg47903AMOvbNzzZjuMxXnQ4zuLEBx6Ajm9T5E6HChbT09dReXF9LHVmdfuPQEuZ3mX99eBtPTgwAAADABgxxxU1rk7aLvpb7ZaXjWObuqd9mzN0x0XIdLXS7Sb/1tdQx0XIdGo5Bx36cp1ULPTVcXQvdjtt1w1Jr3ekDnmZvb69evXp17mU0ZW9vL1evXi1zr2MVOp7UW0cNT+qtYaLjaXrrqOFJvTVMdDxNbx01PKm3homOp+mto4YnaTiGUsrf11r3zrqfX5UCAAAAaJTBDQAAAECjDG4AAAAAGmVwAwAAANCoJl5V6tq1azk4OJh7GU25du3a3EtYmY4n9dZRw5N6a5joeJreOmp4Um8Nk9t3XCwWSdp7lYxt662jvXhSbw0THU/TW0cNT9LwcnHFDQDAjh0dHV26oQ20aLFY/Nsg9bLY39+fewnAigxuAAAALglXPUB/DG4AAAAAGmVwAwAAANAogxsAAACARhncAAAAADTK4AYAAACgUXfMvQAAAIA5HB0dzb0EgDO54gbgElksFlksFnMvY6P29/fnXgIAAGyNwQ0AAABAowxuAOjawcHB3EsAAICtMbgBAAAAaJTBDQAAAECjDG4AAAAAGmVwAwAAANCoO+ZeAAC7c3R0NPcSAACAFbjiBgAAAKBRBjcAAAAAjTK4AQAAAGiUwQ0AAABAowxuAAAAABplcAMAAADQKIMbAABY0WKxyGKxmHsZp9rf3597CQBskMENAAAAQKPOHNyUUr5QSnmrlPLtY8fuKaV8tZTy6vT+g9PxUkr5XCnltVLKN0spj2xz8Zyfjv3TcAw69k/DMejYPw2XOzg4mHsJ56Zj/zQcg45tO88VN4dJnrjl2LNJXqm1PpzklenjJPlYkoent2eSfH4zy2QDDqNj7w6j4QgOo2PvDqPhCA6jY+8Oo+EIDqNj7w6j4QgOo2Ozzhzc1Fq/luRntxx+MsmL0+0Xk3zi2PEv1hu+nuQDpZT7L7LAln9/uCdzd+TiNNyOK1eu5MqVKzt7PB37p+F22IusSsMx6Ng/DcegY9vWfY6b+2qtb063f5Lkvun2A0l+dOx+b0zHTiilPFNKuVpKuXr9+vU1l8EF6dg/DcegY/80HIOO/dNwDDr2T8MxXKijhptz4ScnrrXWJHWNv/dCrXWv1rp31113XXQZXJCO/dNwDDr2T8Mx6Ng/DcegY/80HMM6HTXcnHUHNz+9eSnU9P6t6fiPkzx47H4fmY7RJh37p+EYdOyfhmPQsX87a3h0dJSjo6OLfAqWsxf7p+EYdGzEuoObl5I8Nd1+KslXjh3/9PQs048leffYpVW0R8f+aTgGHfun4Rh07J+GY9CxfxqOQcdG3HHWHUopX0ryW0k+VEp5I8mfJPlski+XUvaTvJ7kk9PdX07y8SSvJfl5ks9sYc2sYe6ON59g2n+ZWt/cDdkMHfun4Rh07J+GY9CxfxqOQce2nTm4qbV+askf/fYp961Jfu+ii2LzdOyfhmPQsX8ajkHH/mk4Bh37p+EYdGzbmYObuY16hcbNlzx9/PHHZ14JF6Fj/7Qbg73YP+3GYC/2T8Mx6Ng/Dfu3yYYXflUpAAAAALbD4AYAAACgUQY3DGV/f3/uJQAAAMDGGNwAAAAANKr5JydmDLt6kumDg4OdPA4AAADsgsHNTDw7+Bh0hDbYi9AGe7F/Go5Bx/5p2L9NNvSrUgAAAACNMrgBAAAAaJTBDQAAAECjDG4AAAAAGmVwAwAdWSwWWSwWcy/jhP39/bmXAAAwJIMbAAAAgEYZ3AAAF3ZwcDD3EgAAhmRwAwAAANAogxsAAACARhncAAAAADTK4AYAAACgUXfMvQAA4PyOjo7mXgIAADvkihvgXBaLRRaLxdzLWMv+/v7cSwAAAFiLwQ0AAABAowxugOEdHBzMvQQAAIC1GNwAAAAANMrgBgAAAKBRBjcAAAAAjTK4AQAAAGjUHXMvAOjD0dHR3EsAAAC4dFxxAwAAANAogxsAAACARhncAAAAADTK4AYAAACgUQY3AAAAAI0yuAEAAABolMENAAAAQKPOHNyUUh4spfxtKeW7pZTvlFL+YDp+Tynlq6WUV6f3H5yOl1LK50opr5VSvllKeWTb/yO4PQ3HoGP/NByDjv3TcAw69k/DMejYPw07UGu97VuS+5M8Mt2+K8kPknw0yX9P8ux0/Nkk/226/fEk/ztJSfJYkr876zEeffTRyi+azsmZfc7ztouGVcdT9dZRw5N6a1h1PFVvHTU8qbeGVcdT9dZRw5N6a1h1PFVvHTU8ScMxJLlaz9HozCtuaq1v1lr/Ybp9Pcn3kjyQ5MkkL053ezHJJ6bbTyb54rSOryf5QCnl/rMeh+3RcAw69k/DMejYPw3HoGP/NByDjv3TsH0rPcdNKeWhJP8xyd8lua/W+ub0Rz9Jct90+4EkPzr2196Yjt36uZ4ppVwtpVx9++23V1w269pkw+nz6TgDe7F/9uIY7MX+2YtjsBf7Zy+OwV7sn4ZtOvfgppTyK0n+Mskf1lr/6fifTZf41FUeuNb6Qq11r9a6d++9967yV1nTphtOf0/HHbMX+2cvjsFe7J+9OAZ7sX/24hjsxf5p2K5zDW5KKb+UGwH/vNb6V9Phn968HGp6/9Z0/MdJHjz21z8yHWNGGo5Bx/5pOAYd+6fhGHTsn4Zj0LF/GrbtPK8qVZIcJPlerfVPj/3RS0memm4/leQrx45/enqm6ceSvHvs8ipmoOEYdOyfhmPQsX8ajkHH/mk4Bh37p2H77jjHff5Tkv+S5FullP87HfujJJ9N8uVSyn6S15N8cvqzl3PjWaZfS/LzJJ/Z6IpZh4Zj0LF/Go5Bx/5pOAYd+6fhGHTsn4aNKzd+VW3mRZRyPcn3517HCj6U5NqWH+PXaq1d/SJgZx130TDprGNnDRN78VSddbQXT9FZw8RePFVnHe3FU3TWMLEXT9VZR3vxFJ01TOzFE0opbyd5L7v5//cmNLUXz3PFzS58v9a6N/cizquUcrWn9e5QNx01XKqbhomOt9FNRw2X6qZhouNtdNNRw6W6aZjoeBvddNRwqW4aJjqeptZ6b0/npbW1rvRy4AAAAADsjsENAAAAQKNaGdy8MPcCVtTbenelp/PS01p3qbfz0tt6d6Wn89LTWnept/PS23p3pafz0tNad6m389Lbenelp/PS01p3qbfz0tt6d6Wn89LUWpt4cmIAAAAATmrlihsAAAAAbjH74KaU8kQp5fullNdKKc82sJ4vlFLeKqV8+9ixe0opXy2lvDq9/+B0vJRSPjet/ZullEfmW/l8WmuY6LiO1jpquLrWGiY6rqO1jhqurrWGiY7raK2jhqtrrWGi4zpa66jh6lprmPTXcdbBTSnlfUn+LMnHknw0yadKKR+dc01JDpM8ccuxZ5O8Umt9OMkr08fJjXU/PL09k+TzO1pjMxptmOi4kkY7HkbDc2u0YaLjShrteBgNz63RhomOK2m042E0PLdGGyY6rqTRjofR8NwabZh01nHuK25+I8lrtdZ/rLX+c5K/SPLknAuqtX4tyc9uOfxkkhen2y8m+cSx41+sN3w9yQdKKffvZqXNaK5houMamuuo4cqaa5jouIbmOmq4suYaJjquobmOGq6suYaJjmtorqOGK2uuYdJfx7kHNw8k+dGxj9+YjrXmvlrrm9PtnyS5b7rdy/q3qadzoONyvZwDDZfr6RzouFwv50DD5Xo6Bzou18s50HC5ns6Bjsv1cg40XK6nc9Bsx7kHN92pN16Gy0txdU7H/mk4Bh37p+EYdOyfhmPQsX8ajqG1jnMPbn6c5MFjH39kOtaan968FGp6/9Z0vJf1b1NP50DH5Xo5Bxou19M50HG5Xs6Bhsv1dA50XK6Xc6Dhcj2dAx2X6+UcaLhcT+eg2Y5zD26+keThUsqvl1J+OcnvJnlp5jWd5qUkT023n0rylWPHPz09y/RjSd49dmnVZdFLw0TH2+mlo4bL9dIw0fF2eumo4XK9NEx0vJ1eOmq4XC8NEx1vp5eOGi7XS8Ok5Y611lnfknw8yQ+S/DDJHzewni8leTPJv+TG767tJ/nV3HhW6VeT/E2Se6b7ltx4huwfJvlWkr2516+hjqN01LD/hjqO0VHD/hvqOEZHDftvqOMYHTXsv2GPHcu0kI0qpTyR5H8keV+S/1Vr/ezGH4St07F/Go5Bx/5pOAYd+6fhGHTsn4Zj0HF3Nj64mV6n/QdJfic3JlffSPKpWut3N/pAbJWO/dNwDDr2T8Mx6Ng/DcegY/80HIOOu7WN57hp8nXaWZmO/dNwDDr2T8Mx6Ng/DcegY/80HIOOO3THFj7naa9x/pu33qmU8kySZ5Lk/e9//6Mf/vCH137Au+++O+++++6/3U7ybx/v2nvvvXfi2J133rny53nnnXdy/fr1sok1rWmWjrfS8UIudcPk3zuu0+4me/EGHS/kUjcc5Otpcsk7JvbiOjTcikvdcZCvqRreosOGyTk6rtrwdp1aaphs5utpkrz++uvXaq33nnW/bQxuzqXW+kKSF5LkoYceqs8999zan2uxWOTo6Ojfbif5t4937cqVKyeOPf744yt/nueff34Ty9m6TXe8lY7bN2rD5N87rtPuph4aJjqepYeOoza8TF9Pk3E7JvbiOjScz6gdL9PXVA1vb8SGt+vUUsNkM19Pk+Tpp59+/Tz328bgZvbXOJ/bzXinbcqOzNLxdhtz1wbouPOGx794ttBwEJd+Lw7g0jcc4Otp4mvqCC79XhzEpd+LA3xN1bD/hskWOp7V6TJ/Pd3Gc9z09DrtLKdj/zQcg47903AMOvZPwzHo2D8Nx6DjDm38ipta67+WUn4/yV/nxsuCfaHW+p1NP86tLuPUbZvm6sjmzNHQPtw8e7F/Go7B19T++Rl1DL6m9k/DMWyjo6+Xy23lOW5qrS8neXkbn5vd0bF/Go5Bx/5pOAYd+6fhGHTsn4Zj0HF3Znty4k1q7XcW4TKb80nCAEbja2rf/IwKsBrf9043xOCmNZ0/ydSsWvqhRsf13NrQF98+tbQXWU9LDX09XZ+vqcCtfE3tn4bL3e77Xks/2+zaNp6cGAAAAIANGO6Km7n/S9RFX8f9Mpu73XE6rqelhomO69Kxfy011G99LXVMtLyoFnpqeHFzd9Tw4jRs1+3azN3tVrvuWGqtO33A0+zt7dWrV6/OvYym7O3t5erVq2XudaxCx5N666jhSb01THQ8TW8dNTypt4aJjqfpraOGJ/XWMNHxNL111PAkDcdQSvn7WuveWffzq1IAAAAAjTK4AQAAAGiUwQ0AAABAowxuAAAAABrVxKtKXbt2LQcHB3MvoynXrl2bewkr0/Gk3jpqeFJvDRMdT9NbRw1P6q1houNpeuuo4Um9NUx0PE1vHTU8ScPLxRU3AAAAAI0yuAEAAABolMENAAAAQKMMbgAAAAAaZXADAAAA0CiDGwAAAIBGGdwAAAAANMrgBgAAAKBRBjcAAAAAjTK4AQAAAGiUwQ0AAABAowxuAAAAABplcAMAAADQKIMbAAAAgEYZ3AAAAAA0yuAGAAAAoFEGNwCXyGKxyGKxOPN++/v7O1gNXF7n3Yst8XUBAOZhcAMAwJkODg7mXgIAXEoGNwCc4B9oAADQBoMbAAAAgEYZ3NAtv2sPAADA6AxuAAAAABplcEO3PAcHAABA31p+pcVWfsvjjrkXAMDuHB0dzb0EIPYiwKbc/Ae/r6uM7MwrbkopXyilvFVK+faxY/eUUr5aSnl1ev/B6XgppXyulPJaKeWbpZRHtrl4zk/H/mk4Bh37p+EYdOyfhmPQsX8ajkHH07XyWx7n+VWpwyRP3HLs2SSv1FofTvLK9HGSfCzJw9PbM0k+v5llsgGHmbFjy5e/deQw9uIIDqNj7w6j4QgOo2PvDqPhCA6jY+8Oo+EIDqNjs84c3NRav5bkZ7ccfjLJi9PtF5N84tjxL9Ybvp7kA6WU+ze12JFcuXIlV65c2dnj6dg/DcegY/803A7fF1mVhmPQcTt2+TVVw+3wfZHj1n1y4vtqrW9Ot3+S5L7p9gNJfnTsfm9Mx04opTxTSrlaSrl6/fr1NZfBBenYPw3HoGP/NByDjv3TcAw69k/DMVyoo4abc+FXlaq11iR1jb/3Qq11r9a6d9ddd110GVyQjv3TcAw69k/DMejYPw3HoGP/NBzDOh013Jx1X1Xqp6WU+2utb06XRL01Hf9xkgeP3e8j0zHapGP/NByDjv3TcAw69k/DMejYv5019GpSW7WTjhqebd0rbl5K8tR0+6kkXzl2/NPTs0w/luTdY5dWrcWT2m7VzjqyNRqOQcf+aTgGHfun4Rh07J+GY9CxEWdecVNK+VKS30ryoVLKG0n+JMlnk3y5lLKf5PUkn5zu/nKSjyd5LcnPk3xmC2tmDTr2T8Mx6Ng/DcegY/80HIOO/dNwDDq27czBTa31U0v+6LdPuW9N8nsXXRSbN3dHl79d3NwN2Qwd+6fhGHTsn4Zj0LF/Go5Bx7at+xw3XNDjjz8+9xLYgJsv0adnvzQcw2XqePPXh0cbiF+GdpfBZdqLo9JwDPr1T8P+bfLr6YVfVQoAAACA7TC4AQAAAGiUwQ0AAABAo5p/jpvRfocfAAAA4LxccQMAzGp/f3/uJQAANKv5K26gZZ7tvX8ajkHHvh0cHMy9BDbEXuyfhgCbscmvpwY3ANARv0IMAHC5+FUpAAAAgEYZ3AAAAAA0yuAGAAAAoFEGNwAAAACNMrgBAAAAaJTBDQAAAECjDG4AAAAAGmVwAwAAANAogxsAAACARhncAAAAl9JischisZh7GQC3ZXADAADQif39/bmXAOyYwQ0AAEAnDg4O5l4CsGMGNwAAAACNMrgBAAAAaJTBDQAAAECjDG4AAAAAGnXH3AsAAACYw9HR0dxLADiTK24AAAAAGmVwAwAAANAogxsAAACARhncAAAAADTK4AYAAACgUQY3AAAAAI0yuAEAAABolMENAAAAQKMMbgAAAAAadebgppTyYCnlb0sp3y2lfKeU8gfT8XtKKV8tpbw6vf/gdLyUUj5XSnmtlPLNUsoj2/4fwe1pOAYd+6fhGHTsn4Zj0LF/Go5Bx/5p2IFa623fktyf5JHp9l1JfpDko0n+e5Jnp+PPJvlv0+2PJ/nfSUqSx5L83VmP8eijj1Z+0XROzuxznrddNKw6nqq3jhqe1FvDquOpeuuo4Um9Naw6nqq3jhqe1FvDquOpeuuo4UkajiHJ1XqORmdecVNrfbPW+g/T7etJvpfkgSRPJnlxutuLST4x3X4yyRendXw9yQdKKfef9Thsj4Zj0LF/Go5Bx/5pOAYd+6fhGHTsn4btW+k5bkopDyX5j0n+Lsl9tdY3pz/6SZL7ptsPJPnRsb/2xnSMBmg4Bh37p+EYdOyfhmPQsX8ajkHH/mnYpnMPbkopv5LkL5P8Ya31n47/2XSJT13lgUspz5RSrpZSrr799tur/FXWtOmG0+fUccfsxf7Zi2OwF/tnL47BXuyfvTgGe7F/GrbrXIObUsov5UbAP6+1/tV0+Kc3L4ea3r81Hf9xkgeP/fWPTMd+Qa31hVrrXq1179577113/ZzTNhomOu6avdg/e3EM9mL/7MUx2Iv9sxfHYC/2T8O2nedVpUqSgyTfq7X+6bE/einJU9Ptp5J85djxT0/PNP1YknePXV7FDDQcg47903AMOvZPwzHo2D8Nx6Bj/zRs3x3nuM9/SvJfknyrlPJ/p2N/lOSzSb5cStlP8nqST05/9nJuPMv0a0l+nuQzG10x69BwDDr2T8Mx6Ng/DcegY/80HIOO/dOwceXGr6rNvIhSrif5/tzrWMGHklzb8mP8Wq21q+vJOuu4i4ZJZx07a5jYi6fqrKO9eIrOGib24qk662gvnqKzhom9eKrOOtqLp+isYWIvnlBKeTvJe9nN/783oam9eJ4rbnbh+7XWvbkXcV6llKs9rXeHuumo4VLdNEx0vI1uOmq4VDcNEx1vo5uOGi7VTcNEx9vopqOGS3XTMNHxNLXWe3s6L62tdaWXAwcAAABgdwxuAAAAABrVyuDmhbkXsKLe1rsrPZ2Xnta6S72dl97Wuys9nZee1rpLvZ2X3ta7Kz2dl57Wuku9nZfe1rsrPZ2Xnta6S72dl97Wuys9nZem1trEkxMDAAAAcFIrV9wAAAAAcIvZBzellCdKKd8vpbxWSnm2gfV8oZTyVinl28eO3VNK+Wop5dXp/Qen46WU8rlp7d8spTwy38rn01rDRMd1tNZRw9W11jDRcR2tddRwda01THRcR2sdNVxdaw0THdfRWkcNV9daw6S/jrMObkop70vyZ0k+luSjST5VSvnonGtKcpjkiVuOPZvklVrrw0lemT5Obqz74entmSSf39Eam9Fow0THlTTa8TAanlujDRMdV9Jox8NoeG6NNkx0XEmjHQ+j4bk12jDRcSWNdjyMhufWaMOks45zX3HzG0leq7X+Y631n5P8RZIn51xQrfVrSX52y+Enk7w43X4xySeOHf9iveHrST5QSrl/NyttRnMNEx3X0FxHDVfWXMNExzU011HDlTXXMNFxDc111HBlzTVMdFxDcx01XFlzDZP+Os49uHkgyY+OffzGdKw199Va35xu/yTJfdPtXta/TT2dAx2X6+UcaLhcT+dAx+V6OQcaLtfTOdBxuV7OgYbL9XQOdFyul3Og4XI9nYNmO849uOlOvfEyXF6Kq3M69k/DMejYPw3HoGP/NByDjv3TcAytdZx7cPPjJA8e+/gj07HW/PTmpVDT+7em472sf5t6Ogc6LtfLOdBwuZ7OgY7L9XIONFyup3Og43K9nAMNl+vpHOi4XC/nQMPlejoHzXace3DzjSQPl1J+vZTyy0l+N8lLM6/pNC8leWq6/VSSrxw7/unpWaYfS/LusUurLoteGiY63k4vHTVcrpeGiY6300tHDZfrpWGi4+300lHD5XppmOh4O7101HC5XhomLXestc76luTjSX6Q5IdJ/riB9XwpyZtJ/iU3fndtP8mv5sazSr+a5G+S3DPdt+TGM2T/MMm3kuzNvX4NdRylo4b9N9RxjI4a9t9QxzE6ath/Qx3H6Khh/w177FimhWxUKeWJJP8jyfuS/K9a62c3/iBsnY7903AMOvZPwzHo2D8Nx6Bj/zQcg467s/HBzfQ67T9I8ju5Mbn6RpJP1Vq/u9EHYqt07J+GY9CxfxqOQcf+aTgGHfun4Rh03K1tPMdNk6/Tzsp07J+GY9CxfxqOQcf+aTgGHfun4Rh03KE7tvA5T3uN89+89U6llGeSPJMk73//+x/98Ic/vPYD3n333SeOvfvuu2t/vot67733kiR33nnn2p/jnXfeyfXr18um1rSGrXW82evWRjpunIb9N0wu+dfUmw2PW6fnqHtxmbvvvnvWvXcre/H2bvc1VceN8n2x/4bJJf+a6vvieuzFrTiz46gNN7UPk+T111+/Vmu996z7bWNwcy611heSvJAkDz30UH3uuefW/lyLxeLEsaOjo7U/30VduXIlSfL444+v/Tmef/75TS1nq9bpeLPXrY10nIeGt9dDw2Tcr6k3Gx63Ts8eOm664Zx771b24u3d7muqjrvn++Lt9dAwGfdrqu+L67EX5zFqw03twyR5+umnXz/P/bYxuNn5a5wfD3ZaUNaiY/80HMPOOyb/3rKFjje/EZ72TbITszRk43Tsn++LY7j0e9H3xfW09LPNIHbasbWvp7veh9t4jpueXqed5XTsn4Zj0LF/Go5Bx/5pOAYd+6fhGHTcoY1fcVNr/ddSyu8n+evceFmwL9Rav7PpxzmuhYnbaObomGi5SRqOYa6ObE4re7GVy/x7pWP/WmnIxbTS0V5cXysNuRj/7t+trTzHTa315SQvb+Nzszs69k/DMejYPw3HoGP/NByDjv3TcAw67s5sT068aabeY/C7p/3TENpgL45Bx/5pOAYd+9ba86Ownsu8D4cZ3Lh8sX+XcQOORsNxtNSy4ydfnFVLDVmfjn3TbxwttfR9kcuopT2Y7H4fbuPJiQEAAADYgCGuuGnt6pp1X8P9sljWS8d+aDi2ljpquJ6WGiY6nsXX1P5pOLaWOmp4cS301HF1LXS7aY5+pda68we91d7eXr169ercy2jK3t5erl69WuZexyp0PKm3jhqe1FvDRMfT9NZRw5N6a5joeJreOmp4Um8NEx1P01tHDU/ScAyllL+vte6ddT+/KgUAAADQKIMbAAAAgEYZ3AAAAAA0yuAGAAAAoFFNvKrUtWvXcnBwMPcymnLt2rW5l7AyHU/qraOGJ/XWMNHxNL111PCk3homOp6mt44antRbw0TH0/TWUcOTNLxcXHEDAAAA0CiDGwAAAIBGGdwAAAAANMrgBgAAAKBRBjcAAAB0abFYZLFYzL0M2CqDGwAAAGjQ/v7+3EugAQY3AAAAAI0yuAEAAIAGHRwczL0EGmBwAwAAANAogxsAAACARhncAAAAADTqjrkXAAAAAOs4Ojqaewmwda64AQAAAGiUwQ0AAABAowxuAAAAABplcAMAAADQKIMbAAAAgEYZ3AAAAAA0yuAGAAAAoFEGNwAAAACNMrgBAAAAaJTBDQAAAECjmh/cLBaLLBaLuZdxqv39/bmXAAAAAAys+cENAMBoWv4PUwBAW84c3JRSvlBKeauU8u1jx+4ppXy1lPLq9P6D0/FSSvlcKeW1Uso3SymPbHPxczs4OJh7CeemY/80HIOO/dNwDDr2T8Mx6Ng/DcegY9vOc8XNYZInbjn2bJJXaq0PJ3ll+jhJPpbk4SFGti4AACAASURBVOntmSSf38wy2YDD6Ni7w2g4gsPo2LvDaDiCw+jYu8NoOILD6Ni7w2g4gsPo2KwzBze11q8l+dkth59M8uJ0+8Uknzh2/Iv1hv+fvfsJkfQ67H/9PVhJFo6QLVsoQlLiLLTRLlaTaLhZBEJAdi3klYkX18K0mY0DCWQjEkwW3uR3F4FrEgzithkZgoMhAQnal+CIgLMYB3fCxX+xJQeEZWRLY4MyyIvE5NzFvONfe7pruqu6qt5zTj8PNF39dnW9R+9Hp6bmzNtvfSXJu0opD21qsKxPx/5pOAYd+6fhGHTsn4Zj0HE7rl+/nuvXr+9kXxqOQcfN2+Q8XPcaNw/WWl+fbv8wyYPT7YeTfP/Y/V6btp1QSrlaSjkqpRzdvHlzzWFwQTr2T8Mx6Ng/DcegY/80HIOO/dNwDBfqqOHmXPjixLXWmqSu8XPP1Vr3aq17995770WHwQXp2D8Nx6Bj/zQcg47903AMOvZPwzGs01HDzblnzZ/7USnloVrr69MpUW9M23+Q5NFj93tk2ra2w8PDi/w4d7ezjmyNhmPQsX8ajsHrm/7trOHtdwXTcis8p/ZPwzHo2Ih1z7h5Mckz0+1nkrxwbPtHp6tMP5nkrWOnVtEeHfun4Rh07J+GY9Cxf5ei4f7+/txD2LZL0XFwGo5Bx0acecZNKeXzSX4vyXtLKa8l+Yskf5nkC6WU/SSvJvnwdPcvJvlgkleS/DTJx7YwZtYwd0f/KnVxczdkM3Tsn4Zj0LF/Go5Bx/5pOAYd23bmwk2t9SNLvvX7p9y3JvnERQfF5l3Gjvv7+zk4OJh7GBtzGRuOSMf+aTgGHft3mRt6fUNLNByDjm1b9xo3XNDttwW7cuXKzCMZ1y5e1OjYPw3HoGP/NByDjv3TcAz69c9c7N8m2134XaUAAAAA2A4LNwAAAACN8qtSAACwIm+4AMCuOOMGAAAAoFHOuGEn/KsUAAAArM7CzUxcHXwMOvZPwzHo2D8Nx6Bj/zSENpiLHOdXpQAukcVikcViMfcwAACAc7JwAwAAANAoCzcAAMCl5ExUoAcWbgAAABqxv78/9xCAxli4AQAAAGiUhRsAAIBGHBwczD0EoDEWbgAAAAAadc/cAwBgdw4PD+ceAgAAsAJn3AAAAAA0yhk3AADApeRMVKAHzrgBAAAAZrFYLLJYLOYeRtMs3AAAAAA0ysINAAAAQKMs3AAAAAA0ysINAAAAQKMs3AAAAAA0ysINAAAAQKPumXsAAAAAwOV0eHg49xCa54wbAAAAgEZZuAEAAABolIUbAAAAgEZZuAEAAABolIUbAAAAgEZZuAEAAABolIUbAAAAgEZZuAEAAABolIUbAAAAgEZZuAEAAABoVa31rh9JHk3yz0m+leSbSf542n5/ki8leXn6/O5pe0ny6SSvJPlakveftY8nnnii8oumY3Jmn/N87KJh1fFUvXXU8KTeGlYdT9VbRw1P6q1h1fFUvXXU8KTeGlYdT9VbRw1P0nAMSY7qORqd54ybnyX501rr40meTPKJUsrjSZ5N8lKt9bEkL01fJ8kHkjw2fVxN8plz7IPt0nAMOvZPwzHo2D8Nx6Bj/zQcg47907BxZy7c1Fpfr7X++3T7ZpJvJ3k4ydNJnp/u9nySD023n07yuWkB6StJ3lVKeWjjI+fcNByDjv3TcAw69k/DMejYPw3HoGP/NGzfSte4KaW8L8lvJfnXJA/WWl+fvvXDJA9Otx9O8v1jP/batO3Ox7paSjkqpRy9+eabKw6bdW2y4fR4Os7AXOyfuTgGc7F/5uIYzMX+mYtjMBf7p2Gbzr1wU0r51SR/n+RPaq3/efx70+9m1VV2XGt9rta6V2vde+CBB1b5Uda06YbTz+m4Y+Zi/8zFMZiL/TMXx2Au9s9cHIO52D8N23WuhZtSyi/lVsC/rbX+w7T5R7dPh5o+vzFt/0FuXdzotkembcxIwzHo2D8Nx6Bj/zQcg47903AMOvZPw7aduXBTSilJDpJ8u9b6V8e+9WKSZ6bbzyR54dj2j5Zbnkzy1rHTq5iBhmPQsX8ajkHH/mk4Bh37p+EYdOyfhu0rt854ussdSvndJP+S5OtJ/mfa/Ge59TtvX0jy60leTfLhWutPpuh/neSpJD9N8rFa69EZ+7iZ5DsX+O/YtfcmubHlffxGrXUj55PtouG0n5467qJh0lnHzhom5uKy/fTU0Vw8fR89NUzMxWX76amjuXj6PnpqmJiLy/bTU0dz8fR99NQwMRdP28ebSd7Obv7/3oSm5uKZCze7UEo5qrXuzT2O8+ptvLvS03Hpaay71Ntx6W28u9LTcelprLvU23Hpbby70tNx6Wmsu9TbceltvLvS03Hpaay71Ntx6W28u9LTcWltrCu9qxQAAAAAu2PhBgAAAKBRrSzcPDf3AFbU23h3pafj0tNYd6m349LbeHelp+PS01h3qbfj0tt4d6Wn49LTWHept+PS23h3pafj0tNYd6m349LbeHelp+PS1FibuMYNAAAAACe1csYNAAAAAHewcAMAAADQqNkXbkopT5VSvlNKeaWU8mwD4/lsKeWNUso3jm27v5TypVLKy9Pnd0/bSynl09PYv1ZKef98I59Paw0THdfRWkcNV9daw0THdbTWUcPVtdYw0XEdrXXUcHWtNUx0XEdrHTVcXWsNk/46zrpwU0p5R5K/SfKBJI8n+Ugp5fE5x5TkWpKn7tj2bJKXaq2PJXlp+jq5Ne7Hpo+rST6zozE2o9GGiY4rabTjtWh4bo02THRcSaMdr0XDc2u0YaLjShrteC0anlujDRMdV9Jox2vR8NwabZh01nHuM25+O8krtdb/qLX+V5K/S/L0nAOqtX45yU/u2Px0kuen288n+dCx7Z+rt3wlybtKKQ/tZqTNaK5houMamuuo4cqaa5jouIbmOmq4suYaJjquobmOGq6suYaJjmtorqOGK2uuYdJfx7kXbh5O8v1jX782bWvNg7XW16fbP0zy4HS7l/FvU0/HQMflejkGGi7X0zHQcblejoGGy/V0DHRcrpdjoOFyPR0DHZfr5RhouFxPx6DZjnMv3HSn3nr/dO+h3jkd+6fhGHTsn4Zj0LF/Go5Bx/5pOIbWOs69cPODJI8e+/qRaVtrfnT7VKjp8xvT9l7Gv009HQMdl+vlGGi4XE/HQMflejkGGi7X0zHQcblejoGGy/V0DHRcrpdjoOFyPR2DZjvOvXDz1SSPlVJ+s5Tyy0n+MMmLM4/pNC8meWa6/UySF45t/+h0leknk7x17NSqy6KXhomOd9NLRw2X66VhouPd9NJRw+V6aZjoeDe9dNRwuV4aJjreTS8dNVyul4ZJyx1rrRv/yK2rM38nyStJnj3jvh9M8t0k30vy59sYz4pj/3yS15P8d2797tp+kvfk1lWlX07yT0nun+5bcusK2d9L8vUke3OPf46OrTXUcfWGLXbUcPWOrTXUcfWGLXbUcPWOrTXUcfWGLXbUcPWOrTXUcfWGLXbUcPWOrTXssWOZBrIx09t9fTfJH0wH4KtJPlJr/dZGd8RW6dg/DcegY/80HIOO/dNwDDr2T8Mx6Lhb2/hVqSbf7ouV6dg/DcegY/80HIOO/dNwDDr2T8Mx6LhD92zhMU97q6zfufNOpZSrSa4mya/8yq888Wu/9mtnPvB9992XJHnrrbdO3X7cnffZpbfffjtJ8s53vnPtx/jxj3+cmzdvlk2NaQ1b67jMfffd9/Nuy1rv0gAdL3XD2/2OW6flyHPRc+rO7HwuJvPPwdvMxYs9p97JXLwQDftvmFzyjoM8p3ptc0nmooZne/XVV2/UWh84637bWLg5l1rrc0meS5L3ve999ZOf/OSZP7NYLJIkh4eHp24/7s777NL169eTJFeuXFn7MT71qU9tajhbtU7HZRaLxc+7LWu9S5el46gNb/c7bp2WPTRMPKeepYeOm5yLyfxz8DZz8WLPqXcyF7dPw7vroWEybsfL9Jzqtc3dadh/wyT5+Mc//up57reNhZtZ3irrzr8scmGzv+UZF3bpG95+Ij3tRU5HZn1OZSPMRXNxbV7fbNTOGx5/LtVwYy59xwGeUy99w0HstONlb7iNa9z09HZfLKdj/zQcg47903AMOvZPwzHo2D8Nx6DjDm38jJta689KKX+U5B+TvCPJZ2ut39z0fu50GVfdtknH/mk4hlY6OgNnfXM1ZLNamYusb46G+m2ejv3zfDqGXXe87P22co2bWusXk3xxG4/N7ujYPw3HoGP/NByDjv3TcAw69k/DMei4O7NdnHjT/A54/y777y2OQMMx6Ajt8Pqmf85YHIOO/fN82r/LPA+HWbgxAYE7dXzRvll5PmXTzMX1mY/986unY2ipo+fU1XkuHUNL83DXtnFxYgAAAAA2oKszbpatqLW20nbR93KnjaY6XszcDfU7m+fUsbXSUb/1tdLwNi1Xp+EYWuqo4d15bdM/DU9Xaq073eFp9vb26tHR0dzDaMre3l6Ojo7K3ONYhY4n9dZRw5N6a5joeJreOmp4Um8NEx1P01tHDU/qrWGi42l666jhSRqOoZTyb7XWvbPu51elAAAAABpl4QYAAACgURZuAAAAABpl4QYAAACgUU28q9SNGzdycHAw9zCacuPGjbmHsDIdT+qto4Yn9dYwuXvHxWKRpL0r829bbx3NxZN6a5iYi6fpraO5eFJvDRMdT9NbRw1P0vByccYNAAAAQKMs3AAAAAA0ysINAAAAQKMs3AAAAAA0ysINO7FYLH5+IUYAAADgfCzcwAbs7+/PPQQAOnJ4eHjp3lEKAFhPE28HDsBu+IsiAAD0xRk3sAEHBwdzDwEAAIABWbgBAAAAaJSFGwAAAIBGWbgBAAAAaJSLE7MTLogKAAAAq3PGDQAAAECjLNwAAAAANMrCDQAAAECjLNwAAAAANMrCDQAAAECjLNwAAAAANMrCDQAAAECjLNwAAAAANMrCDQAAAECjLNwAAAAANMrCDQAAAECjLNwAAMCKFotFFovF3MM4YX9/f+4hALBhZy7clFI+W0p5o5TyjWPb7i+lfKmU8vL0+d3T9lJK+XQp5ZVSytdKKe/f5uA5Px37p+EYdOyfhmPQsX8anu7g4GDuIaxEx/7N3bDVRdTezNlRw7Od54yba0meumPbs0leqrU+luSl6esk+UCSx6aPq0k+s5lhsgHXomPvrkXDEVyLjr27Fg1HcC069u5aNBzBtejYu2vR8Oc6PuvtWnRM0mbDMxduaq1fTvKTOzY/neT56fbzST50bPvn6i1fSfKuUspDmxos69Oxfxpux/Xr13P9+vWd7U/H/mm4HeYiq9JwDDpuxy6fUzUcg45tW/caNw/WWl+fbv8wyYPT7YeTfP/Y/V6bttEmHfun4Rh07J+GY9CxfxqOQcf+XdqGvf264hkuZccWG1744sS11pqkrvpzpZSrpZSjUsrRzZs3LzoMLkjH/mk4Bh37p+EYdOyfhmPQsX8ajmGdjhpuzroLNz+6fSrU9PmNafsPkjx67H6PTNtOqLU+V2vdq7Xu3XvvvWsOgwvSsX8ajkHH/mk4Bh37t7OGh4eHOTw83MyouZO52D8Nx3ChjhpuzroLNy8meWa6/UySF45t/+h0leknk7x17NQq2qNj/zQcg47903AMOvZPwzHo2D8Nx6BjI+456w6llM8n+b0k7y2lvJbkL5L8ZZIvlFL2k7ya5MPT3b+Y5INJXkny0yQf28KYWYOO/dNwDDr2T8Mx6Ng/DcegY//mbujMt82Ys6OGZztz4abW+pEl3/r9U+5bk3ziooNi83Tsn4Zj0LF/Go5Bx/5pOAYd+6fhGHRs25kLN2zH7bfnu3Llyswj4SJ07J92YzAX+6fdGMzF/mk4Bv36Zy72b5MNL/yuUgAAAABsh4UbAAAAgEZZuAEAAABolIUbAOjIYrHIYrGYZd/7+/uz7BcA4DKzcAMAAADQKO8qNRNXBx+DjtAGc3E3Dg4O5h4CjTMX+6chtMFc7N8mGzrjBgAAAKBRFm4AAAAAGmXhBgAAAKBRFm4AAAAAGuXixADQkcPDw7mHAADADjnjBgAAAKBRFm4AAAAAGmXhBgAAAKBRFm4AAAAAGmXhBgAAAKBRFm4AAAAAGmXhBgAAAKBRFm4AAAAAGmXhBgAAAKBRFm4AAAAAGmXhBgAAAKBRFm4AAIBLabFYZLFYzD0MgLuycAOcixc2AAAAu2fhBriU9vf35x4CAADAmSzcAJfSwcHB3EMAAAA4k4UbAAAAgEZZuAEAAABolIUbAAAAgEbdM/cAAAAA5nB4eDj3EADOZOEGOBcvbAAAAHbPr0oBAAAANMrCDQAAAECjLNwAAAAAtKrWetePJI8m+eck30ryzSR/PG2/P8mXkrw8fX73tL0k+XSSV5J8Lcn7z9rHE088UflF0zE5s895PnbRsOp4qt46anhSbw2rjqfqraOGJ/XWsOp4qt46anhSbw2rjqfqraOGJ2k4hiRH9RyNznPGzc+S/Gmt9fEkTyb5RCnl8STPJnmp1vpYkpemr5PkA0kemz6uJvnMOfbBdmk4Bh37p+EYdOyfhmPQsX8ajkHH/mnYuDMXbmqtr9da/326fTPJt5M8nOTpJM9Pd3s+yYem208n+dy0gPSVJO8qpTy08ZFzbhqOQcf+aTgGHfun4Rh07J+GY9Cxfxq2b6Vr3JRS3pfkt5L8a5IHa62vT9/6YZIHp9sPJ/n+sR97bdp252NdLaUclVKO3nzzzRWHzbo22XB6PB1nYC72z1wcg7nYP3NxDOZi/8zFMZiL/dOwTedeuCml/GqSv0/yJ7XW/zz+vel3s+oqO661Pldr3au17j3wwAOr/Chr2nTD6ed03DFzsX/m4hjMxf6Zi2MwF/tnLo7BXOyfhu0618JNKeWXcivg39Za/2Ha/KPbp0NNn9+Ytv8gty5udNsj0zZmpOEYdOyfhmPQsX8ajkHH/mk4Bh37p2Hbzly4KaWUJAdJvl1r/atj33oxyTPT7WeSvHBs+0fLLU8meevY6VXMQMMx6Ng/DcegY/80HIOO/dNwDDr2T8P2lVtnPN3lDqX8bpJ/SfL1JP8zbf6z3Pqdty8k+fUkryb5cK31J1P0v07yVJKfJvlYrfXojH3cTPKdC/x37Np7k9zY8j5+o9a6kfPJdtFw2k9PHXfRMOmsY2cNE3Nx2X566mgunr6Pnhom5uKy/fTU0Vw8fR89NUzMxWX76amjuXj6PnpqmJiLp+3jzSRvZzf/f29CU3PxzIWbXSilHNVa9+Yex3n1Nt5d6em49DTWXertuPQ23l3p6bj0NNZd6u249DbeXenpuPQ01l3q7bj0Nt5d6em49DTWXertuPQ23l3p6bi0NtaV3lUKAAAAgN2xcAMAAADQqFYWbp6bewAr6m28u9LTcelprLvU23Hpbby70tNx6Wmsu9TbceltvLvS03Hpaay71Ntx6W28u9LTcelprLvU23Hpbby70tNxaWqsTVzjBgAAAICTWjnjBgAAAIA7WLgBAAAAaNTsCzellKdKKd8ppbxSSnm2gfF8tpTyRinlG8e23V9K+VIp5eXp87un7aWU8ulp7F8rpbx/vpHPp7WGiY7raK2jhqtrrWGi4zpa66jh6lprmOi4jtY6ari61homOq6jtY4arq61hkl/HWdduCmlvCPJ3yT5QJLHk3yklPL4nGNKci3JU3dsezbJS7XWx5K8NH2d3Br3Y9PH1SSf2dEYm9Fow0THlTTa8Vo0PLdGGyY6rqTRjtei4bk12jDRcSWNdrwWDc+t0YaJjitptOO1aHhujTZMOus49xk3v53klVrrf9Ra/yvJ3yV5es4B1Vq/nOQnd2x+Osnz0+3nk3zo2PbP1Vu+kuRdpZSHdjPSZjTXMNFxDc111HBlzTVMdFxDcx01XFlzDRMd19BcRw1X1lzDRMc1NNdRw5U11zDpr+PcCzcPJ/n+sa9fm7a15sFa6+vT7R8meXC63cv4t6mnY6Djcr0cAw2X6+kY6LhcL8dAw+V6OgY6LtfLMdBwuZ6OgY7L9XIMNFyup2PQbMe5F266U2+9f7r3UO+cjv3TcAw69k/DMejYPw3HoGP/NBxDax3nXrj5QZJHj339yLStNT+6fSrU9PmNaXsv49+mno6Bjsv1cgw0XK6nY6Djcr0cAw2X6+kY6LhcL8dAw+V6OgY6LtfLMdBwuZ6OQbMd5164+WqSx0opv1lK+eUkf5jkxZnHdJoXkzwz3X4myQvHtn90usr0k0neOnZq1WXRS8NEx7vppaOGy/XSMNHxbnrpqOFyvTRMdLybXjpquFwvDRMd76aXjhou10vDpOWOtdaNf+TW1Zm/k+SVJM+ecd8PJvluku8l+fNtjGfFsX8+yetJ/ju3fndtP8l7cuuq0i8n+ack90/3Lbl1hezvJfl6kr25xz9Hx9Ya6rh6wxY7arh6x9Ya6rh6wxY7arh6x9Ya6rh6wxY7arh6x9Ya6rh6wxY7arh6x9Ya9tixTAPZmOntvr6b5A+mA/DVJB+ptX5roztiq3Tsn4Zj0LF/Go5Bx/5pOAYd+6fhGHTcrW38qlSTb/fFynTsn4Zj0LF/Go5Bx/5pOAYd+6fhGHTcoXu28JinvVXW79ztB+699976nve8Z+0d3nfffSe2vfXWW2s/3kW9/fbbSZJ3vvOdaz/Gj3/849y8ebNsakxr2FrH271Oa3S37+3aAB3Nxf4bJpe84+2Gx63Tc9S5uOw5U8OtuNRzMRniOXWWhrebzf0ax1z0GvW4yzgX7+T59MJW6jhSw009nybJq6++eqPW+sBZ99vGws25lFKuJrmaJPfff38++clPrv1Yi8XixLbDw8O1H++irl+/niS5cuXK2o/xqU99alPD2ap1Ot7udVqju31v1y5LR3Px7npomIzb8XbD49bp2UPHTT6fajifUedicnmeUzfd8HazuV/jmIteox7XQ0fPp3enYf/Pp0ny8Y9//NXz3G8bvyp1rrfKqrU+V2vdq7Xu3XvvvRfa4eHh4c8/2Jidd2TjzMUxzDIXW+p45cqVC72waYCG/TdMdBzBpX9tYy7SCM+nYziz48gNd/18uo2Fm57e7ovldOyfhmPQsX8ajkHH/mk4Bh37p+EYdNyhjf+qVK31Z6WUP0ryj0nekeSztdZvbno/dzrt1CnWN1dHNsdcHIO52D8Nx6Bj//y5OAZzsX/m4hjm6HiZG27lGje11i8m+eI2Hpvd0bF/Go5Bx/5pOAYd+6fhGHTsn4Zj0HF3Zrs48abdeeE3YB7mIgDccvxaDP5chHl5jdq/y9xwiIWbyxgOWmQujqOllqdduZ+zaTiGljrSP3ORy8zzaf9aarjr59NtXJwYAAAAgA0Y4oybVt4S7LYB3mZxq+7Wq6WWOq6upX6JhutqqaOGd7eslYZjaKljouVFzd1Tv7N5jTqulvolGq6jpYZz9Cu11p3v9E57e3v16Oho7mE0ZW9vL0dHR2XucaxCx5N666jhSb01THQ8TW8dNTypt4aJjqfpraOGJ/XWMNHxNL111PAkDcdQSvm3WuveWffzq1IAAAAAjbJwAwAAANAoCzcAAAAAjbJwAwAAANCoJt5V6saNGzk4OJh7GE25cePG3ENYmY4n9dZRw5N6a5joeJreOmp4Um8NEx1P01tHDU/qrWGi42l666jhSRpeLs64AQAAAGiUhRsAAACARlm4AQAAAGiUhRsAAACARlm4AQAAAGiUhRvgXBaLRRaLxdzDAAAAuFQs3ACX3v7+/txDAAAAOJWFG+DSOzg4mHsIAAAAp7JwAwAAANAoCzcAAAAAjbJwAwAAANAoCzcAAAAAjbpn7gEAfTg8PJx7CAAAAJeOM24AAAAAGmXhBgAAAKBRFm4AAAAAGmXhBgAAAKBRFm4AAAAAGmXhBgAAAKBRFm4AAAAAGmXhBgAAAKBRFm4AAAAAGmXhBgAAAKBRFm4AAAAAGmXhBgAAAKBRZy7clFI+W0p5o5TyjWPb7i+lfKmU8vL0+d3T9lJK+XQp5ZVSytdKKe/f5uA5Px37p+EYdOyfhmPQsX8ajkHH/mk4Bh3bdp4zbq4leeqObc8meanW+liSl6avk+QDSR6bPq4m+cxmhskGXIuOvbsWDUdwLTr27lo03Ljr16/n+vXru9zltejYu2vRcOPMRdZwLRqO4Fp03KhNPp+euXBTa/1ykp/csfnpJM9Pt59P8qFj2z9Xb/lKkneVUh7ayEi5EB37p+EYdOyfhmOYu+NischisbjIQ5ywv7+/0cdr3dwN2Qwd+6fhGHRs27rXuHmw1vr6dPuHSR6cbj+c5PvH7vfatO2EUsrVUspRKeXo5s2baw6DC9KxfxqOQcf+aTgGHfun4Rh07J+GY7hQRw0358IXJ6611iR1jZ97rta6V2vdu/feey86DC5Ix/5pOAYd1zPDqf1LaTiGHjseHBzsdH+t67EhJ+nYPw3HsE5HDTdn3YWbH90+FWr6/Ma0/QdJHj12v0embbRJx/5pOAYd+6fhGHTsn4Zj0LF/Go5Bx0asu3DzYpJnptvPJHnh2PaPTleZfjLJW8dOraI9OvZPwzHo2D8Nx6Bj/zQcg47903AMOjbinrPuUEr5fJLfS/LeUsprSf4iyV8m+UIpZT/Jq0k+PN39i0k+mOSVJD9N8rEtjJk16Ng/DcegY/80HIOO/dNwDDr2T8Mx6Ni2Mxduaq0fWfKt3z/lvjXJJy46qMvg9rUYrly5spP96bgdu+yo4XaYi6xKw+3Y1Ry8Tcft8Odi/8zFMZiL/fMatX+bbHfmwg0AAJt1eHg49xAAgE5c+F2lAAAAANgOZ9wA0L1dn9oPAAC74owbAAAAgEZZuAEAAABolF+VmonT+segY/80BNgcz6nQBnOxfxpynDNuAAAAABpl4QYAAACgURZuAAAAABpl4QYAAACgURZuAAAAABpl4QYAAACgAYoyXAAAIABJREFUURZuAAAAABpl4QYAAACgURZuAAAAABpl4QYAAACgURZuAABgRYvFIovFYu5hAHAJWLgBAAAAaJSFGwAAAIBGWbgBAAAAaJSFGwAAAIBGWbgBAAAAaJSFGwAAAIBG3TP3AAAAoDeHh4dzDwGAS8IZNwAAAACNsnADAAAA0CgLNwAAAACNsnADAAAA0CgLNwDQkcVikcViMfcwAADYEQs3AAAAAI2ycAMAAADM4jKdTby/v7/Wz1m4AQAAAGiUhRsAAACALTs4OFjr5yzcAAAAALSq1nrXjySPJvnnJN9K8s0kfzxtvz/Jl5K8PH1+97S9JPl0kleSfC3J+8/axxNPPFH5RdMxObPPeT520bDqeKreOmp4Um8Nq46n6q2jhif11rDqeKreOmp4Um8Nq46n6q2jhidpOIYkR/Ucjc5zxs3PkvxprfXxJE8m+UQp5fEkzyZ5qdb6WJKXpq+T5ANJHps+rib5zDn2wXZpOAYd+6fhGHTsn4Zj0LF/Go5Bx/5p2LgzF25qra/XWv99un0zybeTPJzk6STPT3d7PsmHpttPJ/nctID0lSTvKqU8tPGRc24ajkHH/mk4Bh37p+EYdOyfhmPQsX8atm+la9yUUt6X5LeS/GuSB2utr0/f+mGSB6fbDyf5/rEfe23adudjXS2lHJVSjt58880Vh826NtlwejwdZ2Au9s9cHIO52D9zcQzmYv/MxTGYi/3TsE3nXrgppfxqkr9P8ie11v88/r3pd7PqKjuutT5Xa92rte498MADq/woa9p0w+nndNwxc7F/5uIYzMX+mYtjMBf7Zy6OwVzsn4btOtfCTSnll3Ir4N/WWv9h2vyj26dDTZ/fmLb/ILcubnTbI9M2ZqThGHTsn4Zj0LF/Go5Bx/5pOAYd+6dh285cuCmllCQHSb5da/2rY996Mckz0+1nkrxwbPtHyy1PJnnr2OlVzEDDMejYPw3HoGP/NByDjv3TcAw69k/D9pVbZzzd5Q6l/G6Sf0ny9ST/M23+s9z6nbcvJPn1JK8m+XCt9SdT9L9O8lSSnyb5WK316Ix93EzynQv8d+zae5Pc2PI+fqPWupHzyXbRcNpPTx130TDprGNnDRNzcdl+eupoLp6+j54aJubisv301NFcPH0fPTVMzMVl++mpo7l4+j56apiYi6ft480kb2c3/39vQlNz8cyFm10opRzVWvfmHsd59TbeXenpuPQ01l3q7bj0Nt5d6em49DTWXertuPQ23l3p6bj0NNZd6u249DbeXenpuPQ01l3q7bj0Nt5d6em4tDbWld5VCgAAAIDdsXADAAAA0KhWFm6em3sAK+ptvLvS03Hpaay71Ntx6W28u9LTcelprLvU23Hpbby70tNx6Wmsu9TbceltvLvS03Hpaay71Ntx6W28u9LTcWlqrE1c4wYAAACAk1o54wYAAACAO1i4AQAAAGjU7As3pZSnSinfKaW8Ukp5toHxfLaU8kYp5RvHtt1fSvlSKeXl6fO7p+2llPLpaexfK6W8f76Rz6e1homO62ito4ara61houM6Wuuo4epaa5jouI7WOmq4utYaJjquo7WOGq6utYZJfx1nXbgppbwjyd8k+UCSx5N8pJTy+JxjSnItyVN3bHs2yUu11seSvDR9ndwa92PTx9Ukn9nRGJvRaMNEx5U02vFaNDy3RhsmOq6k0Y7XouG5Ndow0XEljXa8Fg3PrdGGiY4rabTjtWh4bo02TDrrOPcZN7+d5JVa63/UWv8ryd8leXrOAdVav5zkJ3dsfjrJ89Pt55N86Nj2z9VbvpLkXaWUh3Yz0mY01zDRcQ3NddRwZc01THRcQ3MdNVxZcw0THdfQXEcNV9Zcw0THNTTXUcOVNdcw6a/j3As3Dyf5/rGvX5u2tebBWuvr0+0fJnlwut3L+Lepp2Og43K9HAMNl+vpGOi4XC/HQMPlejoGOi7XyzHQcLmejoGOy/VyDDRcrqdj0GzHuRduulNvvX+691DvnI7903AMOvZPwzHo2D8Nx6Bj/zQcQ2sd5164+UGSR499/ci0rTU/un0q1PT5jWl7L+Pfpp6OgY7L9XIMNFyup2Og43K9HAMNl+vpGOi4XC/HQMPlejoGOi7XyzHQcLmejkGzHedeuPlqksdKKb9ZSvnlJH+Y5MWZx3SaF5M8M91+JskLx7Z/dLrK9JNJ3jp2atVl0UvDRMe76aWjhsv10jDR8W566ajhcr00THS8m146arhcLw0THe+ml44aLtdLw6TljrXWjX/k1tWZv5PklSTPnnHfDyb5bpLvJfnzbYxnxbF/PsnrSf47t353bT/Je3LrqtIvJ/mnJPdP9y25dYXs7yX5epK9ucc/R8fWGuq4esMWO2q4esfWGuq4esMWO2q4esfWGuq4esMWO2q4esfWGuq4esMWO2q4esfWGvbYsUwD2Zjp7b6+m+QPpgPw1SQfqbV+a6M7Yqt07J+GY9CxfxqOQcf+aTgGHfun4Rh03K1t/KpUk2/3xcp07J+GY9CxfxqOQcf+aTgGHfun4Rh03KF7tvCYp71V1u/c7Qfuvffe+p73vGftHd53330ntr311ltrP95FvP322ye2vfOd71z5cX784x/n5s2bZRNjWtNWOt6tVUsdk//dcp1+t83c0Vy8g7n4i273urNRSx0Tc3EdLTU0F3U8btS5uOz59Kzv7doAz6eJPxdH6Oj59A4dPp8mK3b0d8XTvfrqqzdqrQ+cdb9tLNycSynlapKrSXL//ffnk5/85NqPtVgsTmw7PDxc+/Eu4vr16ye2XblyZeXH+dSnPrWJ4Wzdqh3v1qqljsn/brlOv9t66Ggu3l0PDZP1Ot7udWejljom5uI6WmpoLup4XA8dN/l8etb3du2yPJ8m/lw8Sw8dPZ/e3YgNL9vfFZPk4x//+Kvnud82Fm7O9VZZtdbnkjyXJO973/sufKGdFv4wvO12vNMmZUe20vF4p1UnJiubdS620NBcXF9Lz6kD2HnDs55rd81cXJ/n1I2apSEbd+mfUwfg+bT/59PkHB39XXFztnGNm57e7ovldOyfhmPQsX8ajkHH/mk4Bh37p+EYdNyhjZ9xU2v9WSnlj5L8Y5J3JPlsrfWbm97Pne5cdfOvxRezrY6XcXV0LnPNRTbLc2r/5mjouXbzWpmLrM+fi2MwF/un4Ri20VGj5bZyjZta6xeTfHEbj83u6Ng/DcegY/80HIOO/dNwDDr2T8Mx6Lg7s12ceJP83mlf/Ms9tM1z6hg81/bPXIR2XOZra4zA82k/vH453RALN61Nvs4vMrV1d/sVjNZasrqWGpqL62mpIetr6dfdzMUx6Mhl5c9FNs3z6XL+rni6bVycGAAAAIANGOKMm5ZOp7ro+7iP7m6tWuqYaLmOlhrqd7ZlvVrqmGi5jpYa6rcZczfV8e56eX2j4935c/FymLunfsv18lya7L5jqfXMt1Pfur29vXp0dDT3MJqyt7eXo6OjMvc4VqHjSb111PCk3homOp6mt44antRbw0TH0/TWUcOTemuY6Hia3jpqeJKGYyil/Futde+s+/lVKQAAAIBGWbgBAAAAaJSFGwAAAIBGWbgBAAAAaFQT7yp148aNHBwczD2Mpty4cWPuIaxMx5N666jhSb01THQ8TW8dNTypt4aJjqfpraOGJ/XWMNHxNL111PAkDS8XZ9wAAAAANMrCDQAAAECjLNwAAAAANMrCDQAAAECjLNwAAAAANMrCDQAArGixWGSxWMw9DAAuAQs3AAAAAI2ycAMAAADQKAs3AAAAAI2ycAMAAADQKAs3AAAAAI2ycAMAAADQqHvmHgAAAPTm8PBw7iEAcEk44wYAAACgURZuAAAAABpl4QYAAACgURZuAC6RxWKRxWIx9zBWsr+/P/cQAABgNhZuAAAAABpl4QaAph0cHMw9BAAAmI2FGwAAAIBGWbgBAAAAaJSFGwAAAIBGWbgBAAAAaNQ9cw8AgN05PDycewgAAMAKnHEDAAAA0KgzF25KKZ8tpbxRSvnGsW33l1K+VEp5efr87ml7KaV8upTySinla6WU929z8Jyfjv3TcAw69k/DMejYPw3HoGP/NByDjm07zxk315I8dce2Z5O8VGt9LMlL09dJ8oEkj00fV5N8ZjPDHM/169dz/fr1Xe7yWnTs3bVoOIJr0bF316LhxvlzcQw77ngtGo7gWnTs3bVouHH+XOS4Mxduaq1fTvKTOzY/neT56fbzST50bPvn6i1fSfKuUspDmxos69OxfxqOQcf+aTgGHfun4Rh07J+GY9Cxbete4+bBWuvr0+0fJnlwuv1wku8fu99r07YTSilXSylHpZSjmzdvrjkMLkjH/mk4Bh37p+EYdOyfhmPQsX8ajuFCHTXcnAtfnLjWWpPUNX7uuVrrXq117957773oMLigbXdcLBZZLBYXGeKF7O/vz7bvXTEXx6Bj/zQcg47903AMXqP2z1wcwzodNdycdRdufnT7VKjp8xvT9h8kefTY/R6ZttEmHfun4Rh07J+GY9CxfxqOQcf+aTgGHRux7sLNi0memW4/k+SFY9s/Ol1l+skkbx07tYr2XJqOBwcHcw9hWy5Nw8Hp2D8Nx6Bj/zQcw6Xp6DVq/w0Hp2Mj7jnrDqWUzyf5vSTvLaW8luQvkvxlki+UUvaTvJrkw9Pdv5jkg0leSfLTJB/bwphZg47903AMOvZPwzHo2D8Nx6Bj/zQcg45tO3Phptb6kSXf+v1T7luTfOKig7oMrly5stP96bgdt9+ibxc9NdyOXTZMdNwWc7F//lwcwy47argd/lwcgz8X++fPxf5tch5e+OLEAAAAAGyHhRsAAACARp35q1KwCYeHh3MPAQAAfoHXqEAPnHEDAAAA0CgLNwAAAACN8qtScAG7vto7m6fhGHQE2AzPp2PQEea3yXnojBsAAACARlm4AQAAAGhU8ws3i8Uii8Vi7mH83P7+/txDAAAAAC6J5hduWnNwcDD3EAAAAIBLwsINAAAAQKMs3AAAAAA0ysINAAAAQKMs3AAAAAA0ysINAAAAQKPumXsAZzk8PJx7CAAAAACzcMYNAAAAQKMs3AAAAAA0ysINAAAAQKMs3AAAAAA0ysINAAAAQKMs3ADnslgsslgs5h4GXHqjzMX9/f25hwAA0AULN8Da/MULAABguyzcAAA7d3BwMPcQZjXKmVMAwPZZuAHWdtn/4gUAALBtFm4AAAAAGmXhBgAAAKBRFm4AAAAAGnXP3AMA+nB4eDj3EICYiwAAl40zbgAAAAAa5YwbAIAdc+YUAHBezrgBAAAAaJSFGwAAALq0WCyyWCzmHsZO7e/vzz0EdszCDQAAAECraq13/UjyaJJ/TvKtJN9M8sfT9vuTfCnJy9Pnd0/bS5JPJ3klydeSvP+sfTzxxBOVXzQdkzP7nOdjFw2rjqfqraOGJ/XWsOp4qt46anhSbw2rjqfqraOGJ/XWsOp4qt46aniShmNIclTP0eg8Z9z8LMmf1lofT/Jkkk+UUh5P8mySl2qtjyV5afo6ST6Q5LHp42qSz5xjH2yXhmPQsX8ajkHH/mk4Bh37p+EYdOyfho07c+Gm1vp6rfXfp9s3k3w7ycNJnk7y/HS355N8aLr9dJLPTQtIX0nyrlLKQxsfOeem4Rh07J+GY9CxfxqOQcf+aTgGHfunYftWusZNKeV9SX4ryb8mebDW+vr0rR8meXC6/XCS7x/7sdembXc+1tVSylEp5ejNN99ccdisa5MNp8fTcQbmYv/MxTGYi/0zF8dgLvbPXByDudg/Ddt07oWbUsqvJvn7JH9Sa/3P49+bfjerrrLjWutztda9WuveAw88sMqPsqZNN5x+TscdMxf7Zy6OwVzsn7k4BnOxf+biGMzF/mnYrnMt3JRSfim3Av5trfUfps0/un061PT5jWn7D3Lr4ka3PTJtY0YajkHH/mk4Bh37p+EYdOyfhmPQsX8atu3MhZtSSklykOTbtda/OvatF5M8M91+JskLx7Z/tNzyZJK3jp1exQw0HIOO/dNwDDr2T8Mx6Ng/DcegY/80bF+5dcbTXe5Qyu8m+ZckX0/yP9PmP8ut33n7QpJfT/Jqkg/XWn8yRf/rJE8l+WmSj9Vaj87Yx80k37nAf8euvTfJjS3v4zdqrRs5n2wXDaf99NRxFw2Tzjp21jAxF5ftp6eO5uLp++ipYWIuLttPTx3NxdP30VPDxFxctp+eOpqLp++jp4aJuXjaPt5M8nZ28//3JjQ1F89cuNmFUspRrXVv7nGcV2/j3ZWejktPY92l3o5Lb+PdlZ6OS09j3aXejktv492Vno5LT2Pdpd6OS2/j3ZWejktPY92l3o5Lb+PdlZ6OS2tjXeldpQAAAADYHQs3AAAAAI1qZeHmubkHsKLexrsrPR2Xnsa6S70dl97Guys9HZeexrpLvR2X3sa7Kz0dl57Guku9HZfexrsrPR2Xnsa6S70dl97Guys9HZemxtrENW4AAAAAOKmVM24AAAAAuIOFGwAAAIBGzb5wU0p5qpTynVLKK6WUZxsYz2dLKW+UUr5xbNv9pZQvlVJenj6/e9peSimfnsb+tVLK++cb+Xxaa5jouI7WOmq4utYaJjquo7WOGq6utYaJjutoraOGq2utYaLjOlrrqOHqWmuY9Ndx1oWbUso7kvxNkg8keTzJR0opj885piTXkjx1x7Znk7xUa30syUvT18mtcT82fVxN8pkdjbEZjTZMdFxJox2vRcNza7RhouNKGu14LRqeW6MNEx1X0mjHa9Hw3BptmOi4kkY7XouG59Zow6SzjnOfcfPbSV6ptf5HrfW/kvxdkqfnHFCt9ctJfnLH5qeTPD/dfj7Jh45t/1y95StJ3lVKeWg3I21Gcw0THdfQXEcNV9Zcw0THNTTXUcOVNdcw0XENzXXUcGXNNUx0XENzHTVcWXMNk/46zr1w83CS7x/7+rVpW2serLW+Pt3+YZIHp9u9jH+bejoGOi7XyzHQcLmejoGOy/VyDDRcrqdjoONyvRwDDZfr6RjouFwvx0DD5Xo6Bs12nHvhpjv11vunew/1zunYPw3HoGP/NByDjv3TcAw69k/DMbTWce6Fmx8kefTY149M21rzo9unQk2f35i29zL+berpGOi4XC/HQMPlejoGOi7XyzHQcLmejoGOy/VyDDRcrqdjoONyvRwDDZfr6Rg023HuhZuvJnmslPKbpZRfTvKHSV6ceUyneTHJM9PtZ5K8cGz7R6erTD+Z5K1jp1ZdFr00THS8m146arhcLw0THe+ml44aLtdLw0THu+mlo4bL9dIw0fFueumo4XK9NExa7lhrnfUjyQeTfDfJ95L8eQPj+XyS15P8d2797tp+kvfk1lWlX07yT0nun+5bcusK2d9L8vUke3OPX0MdR+moYf8NdRyjo4b9N9RxjI4a9t9QxzE6ath/wx47lmkgG1VKeSrJ/53kHUn+n1rrX258J2ydjv3TcAw69k/DMejYPw3HoGP/NByDjruz8YWb6X3av5vkD3Jr5eqrST5Sa/3WRnfEVunYPw3HoGP/NByDjv3TcAw69k/DMei4W9u4xk2T79POynTsn4Zj0LF/Go5Bx/5pOAYd+6fhGHTcoXu28Jinvcf579x5p1LK1SRXk+RXfuVXnvi1X/u1tXd43333ndj21ltvrf14F/X2228nSd75zneu/Rg//vGPc/PmzbKpMa1Bx/47ath/w+SSd7zd8Lh1epqLGm7Ape6YDPGcqmH/DZMtdrzd685G991338+3LbvPLg3Q0Vzsv2Fyjo6jNtzUa5skefXVV2/UWh84637bWLg5l1rrc0meS5L3ve999ZOf/OTaj7VYLE5sOzw8XPvxLur69etJkitXrqz9GJ/61Kc2NZyt0vHueuio4d310DAZt+Pthset07OHjhreXQ8Nk3E7JpfnOVXDu+uhYbJex9u97my0WCx+vm3ZfXbpsnQ0F+9Ow/5f2yTJxz/+8VfPc79tLNzM8h7ndz6ZcmE69m/nDY8/eWq4MZd+Lt7+g/C0PyQ7cenn4gANk5nnIhtx6Z9PBzFLRzbKXBzDTjte9tc227jGTU/v085yOvZPwzHo2D8Nx6Bj/zQcg47903AMOu7Qxs+4qbX+rJTyR0n+MbfeFuyztdZvbno/d2ph1W0kc3VkczQcg47903AMrby+cQbO+lppyMXo2D8Nx+D1zW5t5Ro3tdYvJvniNh6b3dGxfxqOQcf+aTgGHfun4Rh07J+GY9Bxd2a7OPEmtfb7bgAAF+X1Tf80HIOO/dOQ3g2xcMM4PJFCG1qai51f0JZouK6W5iEA8L/t+rXNNi5ODAAAAMAGDHfGTQsX7Vv3PdwvuxbaHafjxbTQU8P1tNDuNg0vbu6eGq5v7nZ30vJiWuip4d2dp5GO/dOwf3M3nKNfqbXufKd32tvbq0dHR3MPoyl7e3s5Ojoqc49jFTqe1FtHDU/qrWGi42l666jhSb01THQ8TW8dNTypt4aJjqfpraOGJ2k4hlLKv9Va9866n1+VAgAAAGiUhRsAAACARlm4AQAAAGiUhRsAAACARjXxrlI3btzIwcHB3MNoyo0bN+Yewsp0PKm3jhqe1FvDRMfT9NZRw5N6a5joeJreOmp4Um8NEx1P01tHDU/S8HJxxg0AAABAoyzcAAAAADTKwg0AAABAoyzcAAAAADTKwg0AAABAoyzcAAAAADTKwg0AAHApLRaLLBaLuYcBcFcWboBzGfGFzf7+/txDAADYCq9zYBwWbgAAdsxiOLBtBwcHcw8B2BALN8Cl5QUNAADQOgs3AABcmMVwANgOCzcAAAAAjbJwAwAAANCoe+YeAAAAwBwODw/nHgLAmSzcAOfihQ0AAMDuWbgBANgxi+EAwHm5xg0AAABAoyzcAAAAADTKwg0AAABAoyzcAAAA0KXFYpHFYjH3MGCrLNwAcCH7+/tzDwEAoFleK3FRFm4AuJCDg4O5hwAAAMOycAMAACtq+dcz/Os+tMU/cnFRFm4AAAAAGnXmwk0p5bOllDdKKd84tu3+UsqXSikvT5/fPW0vpZRPl1JeKaV8rZTy/m0OnvPTsX8ajmHuji3/C3Ev5m7IZujYPw2X6+lf93Xsn4Zj0LFt5znj5lqSp+7Y9mySl2qtjyV5afo6ST6Q5LHp42qSz2xmmGzAtejYu2vRcATXomPvrkXDEVyLjr27Fg1HcC069u5aNBzBtejYrDMXbmqtX07ykzs2P53k+en280k+dGz75+otX0nyrlLKQ5saLOvTsX8ajkHH/mm4HdevX8/169d3tj8d+6fhGHTs39wNDw8Pc3h4eJGHIPN3HNEmX9use42bB2utr0+3f5jkwen2w0m+f+x+r03bTiilXC2lHJVSjm7evLnmMLggHfun4Rh07J+GY9CxfxqOQcf+aTiGC3XUcHMufHHiWmtNUtf4uedqrXu11r177733osPggnTsn4Zj0LF/Go5Bx/5pOAYd+6fhGNbpqOHmrLtw86Pbp0JNn9+Ytv8gyaPH7vfItI026dg/DcegY/80HIOO/dtZQ7+esVXmYv80HIOOjVh34ebFJM9Mt59J8sKx7R+drjL9ZJK3jp1aRXt07J+GY9CxfxqOQcf+aTgGHfun4Rh0bMQ9Z92hlPL5JL+X5L2llNeS/EWSv0zyhVLKfpJXk3x4uvsXk3wwyStJfprkY1sYM2vQsX8ajmHujv51+OLmbshm6Ng/DcegY/80HIOObTtz4abW+pEl3/r9U+5bk3ziooNi83Tsn4Zj0LF/Go5Bx/5pOAYd+6fhGHRs25kLN2zH7bcFu3Llyswj4SJ07J+GY9Cxf9qNwVzsn4Zj0LF/GvZvk+0u/K5SAAAAAGyHhRsAAACARlm4AYCOLBaLLBaLuYcBAMCOWLgBAAAAaJSFGwAAAIBGeVepmbg6+Bh07J+GY9AR2mAu9k/DMejYPw05zhk3AAAAAI2ycAMAAADMwhsvnM3CDQAAAECjXOMGADpyeHg49xAAANghZ9wAAAAANMrCDQAAAECjLNwAAAAANMrCDQAAAECjXJwYAAAAmIU3XjibM24AAAAAGmXhBgAAAKBRFm4AAAAAGmXhBgAAAKBRFm4AAAAAGmXhBgAAAKBRFm4AAAAAGmXhBgAAAKBRFm4AAAAAGmXhBgAAAKBRFm4AAAAAGmXhhp1YLBZZLBZzDwMAAAC6YuEGAAAAoFEWbgAAAAAaZeEGAAAAoFEWbgAAAAAaZeEG7rC/vz/3EAAAACCJhRsAAACAZp25cFNKebSU8s+llG+VUr5ZSvnjafv9pZQvlVJenj6/e9peSimfLqW8Ukr5Winl/dv+j+DuWmh4eHiYw8PDiz7MThwcHMw9hFO10JGL0XAMOvZPwzHo2D8Nx6Bj/zTsQK31rh9JHkry/un2vUm+m+TxJP9Xkmen7c8m+V/T7Q8m+X+TlCRPJvnXs/bxxBNPVH7RdEzO7HOej100rDqeqreOGp7UW8Oq46l666jhSb01rDqeqreOGp7UW8Oq46l666jhSRqOIclRPUejM8+4qbW+Xmv99+n2zSTfTvJwkqeTPD/d7fkkH5puP53kc9M4vpLkXaWUh87aD9uj4Rh07J+GY9CxfxqOQcf+aTgGHfunYftWusZNKeV9SX4ryb8mebDW+vr0rR8meXC6/XCS7x/7sdembTRAwzHo2D8Nx6Bj/zQcg47903AMOvZPwzade+GmlPKrSf4+yZ/UWv/z+PemU3zqKjsupVwtpRyVUo7efPPNVX6UNW264fSYOu6Yudg/c3EM5mL/zMUxmIv9MxfHYC72T8N2nWvhppTyS7kV8G9rrf8wbf7R7dOhps9vTNt/kOTRYz/+yLTtF9Ran6u17tVa9x544IF1x885baNhouOumYv9MxfHYC72z1wcg7nYP3NxDOZi/zRs23neVaokOUjy7VrrXx371otJnpluP5Oj9EtCAAAThklEQVTkhWPbPzpdafrJJG8dO72KGWg4Bh37p+EYdOyfhmPQsX8ajkHH/mnYvnvOcZ//I8n/meTrpZT/b9r2Z0n+MskXSin7SV5N8uHpe1/MratMv5Lkp/n/27uDEEnP+87jvwdr44PX2FYsZCGLOAdddFt5SCT2srAsyO6DfArxYS3MGF0S2D2KDSYHX5w9LKwhGMSOaRkWh0ACFrSXxRELvoyDZ0OwnQRZckBYRrY0NiiDckgWnhz6nWxluqu7q6ar3uf/9ucDTVe/U9316PnqqZ55eOut5POXOmK2oeEy6FifhsugY30aLoOO9Wm4DDrWp+Hg2vFL1WYeRGt3krw69zg28NEkt3f8GL/Wey91PlmxjvtomBTrWKxhYi2eqlhHa/EUxRom1uKpinW0Fk9RrGFiLZ6qWEdr8RTFGibW4gmttXeSvJf9/P99GYZaixc542YfXu29X5t7EBfVWrtVabx7VKajhmuVaZjoeIYyHTVcq0zDRMczlOmo4VplGiY6nqFMRw3XKtMw0fE0vfeHKs3LaGPd6O3AAQAAANgfGzcAAAAAgxpl4+bFuQewoWrj3ZdK81JprPtUbV6qjXdfKs1LpbHuU7V5qTbefak0L5XGuk/V5qXaePel0rxUGus+VZuXauPdl0rzMtRYh7g4MQAAAAAnjXLGDQAAAAD3mH3jprX2TGvt1dba6621FwYYz9daa2+31n64cuzB1tq3W2uvTZ8/Mh1vrbWvTGP/fmvtyflGPp/RGiY6bmO0jhpubrSGiY7bGK2jhpsbrWGi4zZG66jh5kZrmOi4jdE6ari50Rom9TrOunHTWntfkj9M8qkkTyT5bGvtiTnHlOQwyTP3HHshySu998eTvDJ9nRyP+/Hp4/kkX93TGIcxaMNEx40M2vEwGl7YoA0THTcyaMfDaHhhgzZMdNzIoB0Po+GFDdow0XEjg3Y8jIYXNmjDpFjHuc+4+Y0kr/fe/7b3/g9J/ijJs3MOqPf+nSS/vOfws0lemm6/lOQzK8e/3o99N8mHW2uP7GekwxiuYaLjFobrqOHGhmuY6LiF4TpquLHhGiY6bmG4jhpubLiGiY5bGK6jhhsbrmFSr+PcGzePJvnJytdvTsdG83Dv/a3p9s+SPDzdrjL+Xao0BzquV2UONFyv0hzouF6VOdBwvUpzoON6VeZAw/UqzYGO61WZAw3XqzQHw3ace+OmnH78Nlzeiqs4HevTcBl0rE/DZdCxPg2XQcf6NFyG0TrOvXHz0ySPrXz98enYaH5+91So6fPb0/Eq49+lSnOg43pV5kDD9SrNgY7rVZkDDderNAc6rldlDjRcr9Ic6LhelTnQcL1KczBsx7k3br6X5PHW2q+31n4lyW8neXnmMZ3m5STPTbefS/LNleOfm64y/VSSd1dOrboqqjRMdDxLlY4arlelYaLjWap01HC9Kg0THc9SpaOG61VpmOh4liodNVyvSsNk5I6991k/knw6yY+S/DjJ7w0wnm8keSvJP+b4tWvXk/xqjq8q/VqSP0vy4HTfluMrZP84yQ+SXJt7/BrquJSOGtZvqOMyOmpYv6GOy+ioYf2GOi6jo4b1G1bs2KaBXKrW2jNJ/nuS9yX5H733L1/6g7BzOtan4TLoWJ+Gy6BjfRoug471abgMOu7PpW/cTO/T/qMk/yHHO1ffS/LZ3vtfX+oDsVM61qfhMuhYn4bLoGN9Gi6DjvVpuAw67tcurnEz5Pu0szEd69NwGXSsT8Nl0LE+DZdBx/o0XAYd9+iBHfzM097j/DfvvVNr7fkkzyfJ+9///k9+7GMfO/cHf+hDH0qSvPvuuyeO33tsTu+9916S5AMf+MDWP+MXv/hF7ty50y5rTFvYWcd17vZdNVfXuw1XbdNz5o57b5isX6f7tpCGibV44ljBjrM0vNtshDXp9+J2rMVLZy1ai1sZaS0mi+g461qc20KeT5MLdLQOz/fGG2/c7r0/dN79drFxcyG99xeTvJgkn/jEJ/oXv/jFc7/n4OAgSXJ0dHTi+L3H5nTz5s0kydNPP731z/jSl750WcPZqW06rnO376q5ut5tuGqbnhU6XmbDZP063ber1DCxFs9ToeNlN7zbbIQ16ffidqzFeViLZ6vQMFnuWkyuTsddrcW5eT7dzhLXYZJ84QtfeOMi99vFxs3s73HOpZil471/uZnT3UV42pNrEVd+LS6gYTJDx9Vfgtbipbjya3Eh/F60FhnDlV+LC3Dl1+ICnk+TGf+OehXX4S6ucVPpfdpZT8f6NFwGHevTcBl0rE/DZdCxPg2XQcc9uvQzbnrv/6+19rtJ/neO3xbsa733v7rsx7nXvbtuo5wKV9UoHdneXA25XNZifRoug+fU+qzFZdCxvlEa+vfi/Zmj41Vehzu5xk3v/VtJvrWLn83+6FifhsugY30aLoOO9Wm4DDrWp+Ey6Lg/s12c+LJd5de7LcVo19WAq8xzam2eT2EM1uJy+L1Yn4a1XfXn08Vs3FzFeOxW8YuFEQ23NdrzqY5cZSOtR2uRq2yktch2Rmro+ZRN7eLixAAAAABcglJn3Ky7gNRoF5a63/dyZ/6mGm5v7nZ3abi9URomOl6GEXrquJ0R2t2l4f0boaeO2xmh3SodNzdSQ/3u3wg9992x9d73+oCnuXbtWr9169bcwxjKtWvXcuvWrTb3ODah40nVOmp4UrWGiY6nqdZRw5OqNUx0PE21jhqeVK1houNpqnXU8CQNl6G19n9779fOu5+XSgEAAAAMysYNAAAAwKBs3AAAAAAMysYNAAAAwKCGeFep27dv58aNG3MPYyi3b9+eewgb0/Gkah01PKlaw0TH01TrqOFJ1RomOp6mWkcNT6rWMNHxNNU6aniShleLM24AAAAABmXjBgAAAGBQNm4AAAAABmXjBgAAAGBQNm4AAAAABmXjBgAAAGBQNm4AAGBDBwcHOTg4mHsYAFwBNm4AAIB/dv369bmHAMAKGzcAAMA/u3HjxtxDAGCFjRsAAACAQdm4AQAAABiUjRsAAACAQdm4AQAAABjUA3MPAAAAqjk6Opp7CABcEc64AQAAABiUjRsAAACAQdm4AQAAABiUjRsAAACAQdm4AQAAABiUjRsAAACAQdm4AQAAABiUjRsAAACAQdm4AQAAABiUjRsAAACAQZ27cdNa+1pr7e3W2g9Xjj3YWvt2a+216fNHpuOttfaV1trrrbXvt9ae3OXguTgd65u74cHBQQ4ODu73x1x5c3fk/mm4DDrWp+Ey6Fifhsug49gucsbNYZJn7jn2QpJXeu+PJ3ll+jpJPpXk8enj+SRfvZxhcgkOo2N1h9FwCQ6jY3WH0XAJDqNjdYfRcAkOo2N1h9FwCQ6j47DO3bjpvX8nyS/vOfxskpem2y8l+czK8a/3Y99N8uHW2iOXNVi2p+Nu3Lx5Mzdv3tzLY2m4G/tsmOi4K9Yim9KxPg2XQcf6NNwNf0dl1bbXuHm49/7WdPtnSR6ebj+a5Ccr93tzOnZCa+351tqt1tqtO3fubDkM7pOO9Wm4DDrWp+Ey6FifhsugY30aLsN9ddTw8tz3xYl77z1J3+L7Xuy9X+u9X/vgBz94v8PgPulYn4bLoGN9Gi6DjvVpuAw61qfhMmzTUcPLs+3Gzc/vngo1fX57Ov7TJI+t3O/j07GtuSDqTu2tIzuj4TLoWJ+Gy6BjfRoug471abgMOg5i242bl5M8N91+Lsk3V45/brrK9FNJ3l05tYrx6FifhsugY30aLoOO9Wm4DDrWp+Ey6DiIB867Q2vtG0n+XZKPttbeTPL7Sb6c5I9ba9eTvJHkt6a7fyvJp5O8nuTvk3x+B2NmCzrWN3fDo6Oj+/0RZP6O3D8Nl0HH+jRcBh3r03AZdBzbuRs3vffPrvmjf3/KfXuS37nfQXH5dKxPw2XQsT4Nl0HH+jRcBh3r03AZdBzbuRs37Mbdt3Z7+umnZx4J90O/+jRcBh3r83txGXSsT8Nl0LE+7eq7zHV43+8qBQAAAMBu2LgBAAAAGNTwL5VyQVQAAADgqnLGDQAAAMCgbNwAAAAADGr4l0otlauEA8D/5/fiMuhYn4bLoCPM7zLXoTNuAAAAAAZl4wYAAABgUDZuAAAAAAZl4wYAAABgUDZuAAAAAAZl4wYAAABgUDZuAAAAAAZl4wYAAABgUDZuAAAAAAZl44a9ODg4yMHBwdzDgCuv0lq8fv363EMAAIDZ2bgBYEg3btyYewgAADA7GzcAAAAAg7JxAwAAADAoGzcAAHtW6XpTAMC8bNwAAAAADOqBuQcAwP4cHR3NPQQAAGADNm7YC/9YBAAAgM15qRQAAADAoGzcAAAAAAzKxg0AAADAoGzcAADs2dHRkeu/AVyCg4ODHBwczD2Mnbh+/frcQ2AQNm4AAK4Y/xgAgDps3AAAXDE3btyYewgAnMNzNXfZuAEAAAAYlI0bAAAAgEHZuAEAAAAYlI0bAAAAgEGdu3HTWnustfZ/Wmt/3Vr7q9baf5qOP9ha+3Zr7bXp80em46219pXW2uutte+31p7c9X8EZ9NwGXSsT8Nl0LE+DZdBx/o0XIa5Ox4dHeXo6Ogy/lOurLkbcgG99zM/kjyS5Mnp9geT/CjJE0n+a5IXpuMvJPmD6fank/yvJC3JU0n+/LzH+OQnP9n5l6Y5ObfPRT720bDreKpqHTU8qVrDruOpqnXU8KRqDbuOp6rWUcOTqjXsOp6qWkcNT9JwGZLc6hdodO4ZN733t3rvfzHdvpPkb5I8muTZJC9Nd3spyWem288m+fo0ju8m+XBr7ZHzHofd0XAZdKxPw2XQsT4Nl0HH+jRcBh3r03B8G13jprX2iST/JsmfJ3m49/7W9Ec/S/LwdPvRJD9Z+bY3p2MMQMNl0LE+DZdBx/o0XAYd69NwGXSsT8MxXXjjprX2r5P8SZL/3Hv/u9U/m07x6Zs8cGvt+dbardbarXfeeWeTb2VLl91w+pk67pm1WJ+1uAzWYn3W4jJYi/VZi8tgLdan4bgutHHTWvtXOQ74P3vvfzod/vnd06Gmz29Px3+a5LGVb//4dOxf6L2/2Hu/1nu/9tBDD207fi5oFw0THffNWqzPWlwGa7E+a3EZrMX6rMVlsBbr03BsF3lXqZbkRpK/6b3/t5U/ejnJc9Pt55J8c+X456YrTT+V5N2V06uYgYbLoGN9Gi6DjvVpuAw61qfhMuhYn4bje+AC9/m3Sf5jkh+01v5yOvZfknw5yR+31q4neSPJb01/9q0cX2X69SR/n+TzlzpitqHhMuhYn4bLoGN9Gi6DjvVpuAw61qfh4NrxS9VmHkRrd5K8Ovc4NvDRJLd3/Bi/1nsvdT5ZsY77aJgU61isYWItnqpYR2vxFMUaJtbiqYp1tBZPUaxhYi2eqlhHa/EUxRom1uIJrbV3kryX/fz/fRmGWosXOeNmH17tvV+bexAX1Vq7VWm8e1Smo4ZrlWmY6HiGMh01XKtMw0THM5TpqOFaZRomOp6hTEcN1yrTMNHxNL33hyrNy2hj3ejtwAEAAADYHxs3AAAAAIMaZePmxbkHsKFq492XSvNSaaz7VG1eqo13XyrNS6Wx7lO1eak23n2pNC+VxrpP1eal2nj3pdK8VBrrPlWbl2rj3ZdK8zLUWIe4ODEAAAAAJ41yxg0AAAAA97BxAwAAADCo2TduWmvPtNZeba293lp7YYDxfK219nZr7Ycrxx5srX27tfba9Pkj0/HWWvvKNPbvt9aenG/k8xmtYaLjNkbrqOHmRmuY6LiN0TpquLnRGiY6bmO0jhpubrSGiY7bGK2jhpsbrWFSr+OsGzettfcl+cMkn0ryRJLPttaemHNMSQ6TPHPPsReSvNJ7fzzJK9PXyfG4H58+nk/y1T2NcRiDNkx03MigHQ+j4YUN2jDRcSODdjyMhhc2aMNEx40M2vEwGl7YoA0THTcyaMfDaHhhgzZMinWc+4yb30jyeu/9b3vv/5Dkj5I8O+eAeu/fSfLLew4/m+Sl6fZLST6zcvzr/dh3k3y4tfbIfkY6jOEaJjpuYbiOGm5suIaJjlsYrqOGGxuuYaLjFobrqOHGhmuY6LiF4TpquLHhGib1Os69cfNokp+sfP3mdGw0D/fe35pu/yzJw9PtKuPfpUpzoON6VeZAw/UqzYGO61WZAw3XqzQHOq5XZQ40XK/SHOi4XpU50HC9SnMwbMe5N27K6cfvn+491IvTsT4Nl0HH+jRcBh3r03AZdKxPw2UYrePcGzc/TfLYytcfn46N5ud3T4WaPr89Ha8y/l2qNAc6rldlDjRcr9Ic6LhelTnQcL1Kc6DjelXmQMP1Ks2BjutVmQMN16s0B8N2nHvj5ntJHm+t/Xpr7VeS/HaSl2ce02leTvLcdPu5JN9cOf656SrTTyV5d+XUqquiSsNEx7NU6ajhelUaJjqepUpHDder0jDR8SxVOmq4XpWGiY5nqdJRw/WqNExG7th7n/UjyaeT/CjJj5P83gDj+UaSt5L8Y45fu3Y9ya/m+KrSryX5syQPTvdtOb5C9o+T/CDJtbnHr6GOS+moYf2GOi6jo4b1G+q4jI4a1m+o4zI6ali/YcWObRoIAAAAAIOZ+6VSAAAAAKxh4wYAAABgUDZuAAAAAAZl4wYAAABgUDZuAAAAAAZl4wYAAABgUDZuAAAAAAb1T6mwGYLR4sUnAAAAAElFTkSuQmCC\n",
-      "text/plain": [
-       "<matplotlib.figure.Figure at 0x7f52bf7af910>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "dataset = gym_problem.dataset(Modes.TRAIN, data_dir)\n",
-    "iterator = dataset.make_one_shot_iterator()\n",
-    "next_element = iterator.get_next()\n",
-    "\n",
-    "\n",
-    "fig=plt.figure(figsize=(20, 80))\n",
-    "columns = 10\n",
-    "rows = 40\n",
-    "\n",
-    "with tf.Session() as sess:\n",
-    "    for inx in range(100):\n",
-    "        value = sess.run(next_element)\n",
-    "        for i in range(10):  # skipping surplus frames.\n",
-    "            value = sess.run(next_element)\n",
-    "        fig.add_subplot(rows, columns, inx+1)        \n",
-    "        image = value[\"inputs\"].reshape([210,160,3])\n",
-    "        plt.imshow(image[:, :, 0].astype(np.float32), cmap=plt.get_cmap('gray'))\n",
-    "plt.show()"
-   ]
-  }
- ],
- "metadata": {
-  "colab": {
-   "collapsed_sections": [],
-   "default_view": {},
-   "name": "T2T with TF Eager",
-   "provenance": [
-    {
-     "file_id": "1-VScmaLkMqWiSbqgUCFWefzisSREd8l1",
-     "timestamp": 1.512175750497E12
-    }
-   ],
-   "version": "0.3.2",
-   "views": {}
-  },
-  "kernelspec": {
-   "display_name": "t2t_kernel",
-   "language": "python",
-   "name": "t2t_kernel"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 2
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython2",
-   "version": "2.7.12"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 1
-}
+{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"name":"hello_t2t-rl.ipynb","version":"0.3.2","provenance":[{"file_id":"1nQvfx1EzY3ElJUy-FVF1G16okSbkeUa2","timestamp":1553274233669}],"collapsed_sections":[]},"kernelspec":{"name":"python3","display_name":"Python 3"},"accelerator":"GPU"},"cells":[{"metadata":{"id":"xCLcAmON-m2i","colab_type":"text"},"cell_type":"markdown","source":["# Tensor2Tensor Reinforcement Learning\n","\n","The `rl` package provides the ability to run model-free and model-based reinforcement learning algorithms.\n","\n","Currently, we support the Proximal Policy Optimization ([PPO](https://arxiv.org/abs/1707.06347)) and Simulated Policy Learning ([SimPLe](https://arxiv.org/abs/1903.00374)).\n","\n","Below you will find examples of PPO training using `trainer_model_free.py` and SimPLe traning using `trainer_model_based.py`.\n"]},{"metadata":{"id":"RW7gEGp3e87G","colab_type":"code","colab":{},"cellView":"form"},"cell_type":"code","source":["#@title\n","# Copyright 2018 Google LLC.\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","# https://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."],"execution_count":0,"outputs":[]},{"metadata":{"id":"pq0BqXm4-3gJ","colab_type":"code","outputId":"6086719f-6268-4b61-8fa3-d251eda24c97","executionInfo":{"status":"ok","timestamp":1553273826475,"user_tz":-60,"elapsed":20650,"user":{"displayName":"Piotr Miłoś","photoUrl":"https://lh3.googleusercontent.com/-050ZBEGpNAA/AAAAAAAAAAI/AAAAAAAAk9g/r6cv_J6J5qA/s64/photo.jpg","userId":"12158759908531801397"}},"colab":{"base_uri":"https://localhost:8080/","height":163}},"cell_type":"code","source":["!pip install -q -U tensor2tensor==1.13.1\n","!pip install -q tensorflow==1.13.1\n","!pip install -q gym[atari]"],"execution_count":0,"outputs":[{"output_type":"stream","text":["\u001b[K    100% |████████████████████████████████| 1.3MB 9.4MB/s \n","\u001b[K    100% |████████████████████████████████| 215kB 27.3MB/s \n","\u001b[K    100% |████████████████████████████████| 143kB 29.6MB/s \n","\u001b[K    100% |████████████████████████████████| 21.1MB 1.7MB/s \n","\u001b[K    100% |████████████████████████████████| 409kB 24.7MB/s \n","\u001b[K    100% |████████████████████████████████| 296kB 25.0MB/s \n","\u001b[K    100% |████████████████████████████████| 61kB 21.5MB/s \n","\u001b[?25h  Building wheel for pypng (setup.py) ... \u001b[?25ldone\n","\u001b[?25h  Building wheel for opt-einsum (setup.py) ... \u001b[?25ldone\n","\u001b[?25h"],"name":"stdout"}]},{"metadata":{"id":"R7-Ni-39DGZW","colab_type":"code","colab":{}},"cell_type":"code","source":["# Helper function for playing videos in the colab.\n","def play_video(path):\n","  from IPython.core.magics.display import HTML\n","  display_path = \"/nbextensions/vid.mp4\"\n","  display_abs_path = \"/usr/local/share/jupyter\" + display_path\n","  !rm -f $display_abs_path\n","  !ffmpeg -loglevel error -i $path $display_abs_path\n","  return HTML(\"\"\"\n","    <video width=\"640\" height=\"480\" controls>\n","      <source src=\"{}\" type=\"video/mp4\">\n","    </video>\n","  \"\"\".format(display_path))"],"execution_count":0,"outputs":[]},{"metadata":{"id":"pueuiKUmAOUT","colab_type":"text"},"cell_type":"markdown","source":["# Play using a pre-trained policy\n","\n","We provide pretrained policies for the following games from the Atari Learning Environment ( [ALE](https://github.com/mgbellemare/Arcade-Learning-Environment)) : alien,\n","amidar,\n"," assault,\n"," asterix,\n"," asteroids,\n"," atlantis,\n"," bank_heist,\n"," battle_zone,\n"," beam_rider,\n"," bowling,\n"," boxing,\n"," breakout,\n"," chopper_command,\n"," crazy_climber,\n"," demon_attack,\n"," fishing_derby,\n"," freeway,\n"," frostbite,\n"," gopher,\n"," gravitar,\n"," hero,\n"," ice_hockey,\n"," jamesbond,\n"," kangaroo,\n"," krull,\n"," kung_fu_master,\n"," ms_pacman,\n"," name_this_game,\n"," pong,\n"," private_eye,\n"," qbert,\n"," riverraid,\n"," road_runner,\n"," seaquest,\n"," up_n_down,\n"," yars_revenge.\n"," \n"," We have 5 checkpoints for each game saved on Google Storage. To get run id to determine the storage path:"]},{"metadata":{"id":"x9pKfNbDFfVh","colab_type":"code","outputId":"97e763cc-caaa-49c8-e532-fcbde828d1a2","executionInfo":{"status":"ok","timestamp":1553274151100,"user_tz":-60,"elapsed":6162,"user":{"displayName":"Piotr Miłoś","photoUrl":"https://lh3.googleusercontent.com/-050ZBEGpNAA/AAAAAAAAAAI/AAAAAAAAk9g/r6cv_J6J5qA/s64/photo.jpg","userId":"12158759908531801397"}},"colab":{"base_uri":"https://localhost:8080/","height":147}},"cell_type":"code","source":["# experiment_id is an integer from [0, 4].\n","def get_run_dir(game, experiment_id):\n","  from tensor2tensor.data_generators.gym_env import ATARI_GAMES_WITH_HUMAN_SCORE_NICE\n","  EXPERIMENTS_PER_GAME = 5\n","  run_id = ATARI_GAMES_WITH_HUMAN_SCORE_NICE.index(game) * EXPERIMENTS_PER_GAME + experiment_id + 1\n","  return \"gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/{}\".format(run_id)\n","\n","get_run_dir('pong', 2)"],"execution_count":0,"outputs":[{"output_type":"stream","text":["\n","WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n","For more information, please see:\n","  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n","  * https://github.com/tensorflow/addons\n","If you depend on functionality not listed there, please file an issue.\n","\n"],"name":"stdout"},{"output_type":"execute_result","data":{"text/plain":["'gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/143'"]},"metadata":{"tags":[]},"execution_count":4}]},{"metadata":{"id":"77fFdm-cFEOB","colab_type":"text"},"cell_type":"markdown","source":["To evaluate and generate videos for a pretrained policy on Pong:"]},{"metadata":{"id":"X-nGlbuTAQXj","colab_type":"code","outputId":"888968f2-f551-4a0f-9fc7-074a949362d6","executionInfo":{"status":"ok","timestamp":1553271580737,"user_tz":-60,"elapsed":842128,"user":{"displayName":"Piotr Kozakowski","photoUrl":"","userId":"01014928596539690143"}},"colab":{"base_uri":"https://localhost:8080/","height":17088}},"cell_type":"code","source":["game = 'pong'\n","run_id = get_run_dir(game, 1)\n","!python -m tensor2tensor.rl.evaluator \\\n","  --loop_hparams_set=rlmb_long_stochastic_discrete \\\n","  --loop_hparams=game=$game,eval_max_num_noops=8,eval_sampling_temps=[0.5] \\\n","  --policy_dir=gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/$run_id/policy \\\n","  --eval_metrics_dir=pong_pretrained \\\n","  --debug_video_path=pong_pretrained \\\n","  --num_debug_videos=4"],"execution_count":0,"outputs":[{"output_type":"stream","text":["\n","WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n","For more information, please see:\n","  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n","  * https://github.com/tensorflow/addons\n","If you depend on functionality not listed there, please file an issue.\n","\n","INFO:tensorflow:Overriding hparams in rlmb_long_stochastic_discrete with game=pong,eval_max_num_noops=8,eval_sampling_temps=[0.5]\n","INFO:tensorflow:Evaluating metric mean_reward/eval/sampling_temp_0.5_max_noops_8_unclipped\n","2019-03-22 16:05:45.007030: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 2300000000 Hz\n","2019-03-22 16:05:45.007306: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x2697860 executing computations on platform Host. Devices:\n","2019-03-22 16:05:45.007346: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): <undefined>, <undefined>\n","2019-03-22 16:05:45.105281: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:998] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n","2019-03-22 16:05:45.105857: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x2697440 executing computations on platform CUDA. Devices:\n","2019-03-22 16:05:45.105908: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): Tesla K80, Compute Capability 3.7\n","2019-03-22 16:05:45.106380: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1433] Found device 0 with properties: \n","name: Tesla K80 major: 3 minor: 7 memoryClockRate(GHz): 0.8235\n","pciBusID: 0000:00:04.0\n","totalMemory: 11.17GiB freeMemory: 11.10GiB\n","2019-03-22 16:05:45.106420: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 16:05:45.499212: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 16:05:45.499307: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 16:05:45.499332: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 16:05:45.499671: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:42] Overriding allow_growth setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.\n","2019-03-22 16:05:45.499741: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","INFO:tensorflow:Using DummyPolicyProblem for the policy.\n","INFO:tensorflow:Setting T2TModel mode to 'train'\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Colocations handled automatically by placer.\n","INFO:tensorflow:Using variable initializer: orthogonal\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/utils/t2t_model.py:1358: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.cast instead.\n","INFO:tensorflow:Transforming feature 'input_action' with symbol_modality_6_64.bottom\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/function.py:1007: calling Graph.create_op (from tensorflow.python.framework.ops) with compute_shapes is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Shapes are always computed; don't use the compute_shapes as it has no effect.\n","INFO:tensorflow:Transforming feature 'input_reward' with symbol_modality_3_64.bottom\n","INFO:tensorflow:Transforming feature 'inputs' with video_modality.bottom\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_video.py:495: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","tf.py_func is deprecated in TF V2. Instead, use\n","    tf.py_function, which takes a python function which manipulates tf eager\n","    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to\n","    an ndarray (just call tensor.numpy()) but having access to eager tensors\n","    means `tf.py_function`s can use accelerators such as GPUs as well as\n","    being differentiable using a gradient tape.\n","    \n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:277: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.cast instead.\n","INFO:tensorflow:Transforming feature 'target_action' with symbol_modality_6_64.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_policy' with identity_modality.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_reward' with symbol_modality_3_64.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_value' with identity_modality.targets_bottom\n","INFO:tensorflow:Transforming feature 'targets' with video_modality.targets_bottom\n","INFO:tensorflow:Building model body\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:598: conv2d (from tensorflow.python.layers.convolutional) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.conv2d instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:602: flatten (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.flatten instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:603: dropout (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.dropout instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:604: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.dense instead.\n","INFO:tensorflow:Transforming body output with identity_modality.top\n","INFO:tensorflow:Transforming body output with identity_modality.top\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:2887: multinomial (from tensorflow.python.ops.random_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.random.categorical instead.\n","2019-03-22 16:06:00.352605: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 16:06:00.352688: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 16:06:00.352724: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 16:06:00.352744: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 16:06:00.353037: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","2019-03-22 16:06:00.588787: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:06:00.647797: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","INFO:tensorflow:Restoring checkpoint gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/policy/model.ckpt-171992\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/saver.py:1266: checkpoint_exists (from tensorflow.python.training.checkpoint_management) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use standard file APIs to check for files with this prefix.\n","2019-03-22 16:06:00.711910: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","INFO:tensorflow:Restoring parameters from gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/policy/model.ckpt-171992\n","2019-03-22 16:06:00.793701: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:06:00.953239: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:06:01.086594: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:06:01.259521: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:06:01.322896: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:06:03.034751: I tensorflow/stream_executor/dso_loader.cc:152] successfully opened CUDA library libcublas.so.10.0 locally\n","INFO:tensorflow:Step 5, mean_score: 0.000000\n","INFO:tensorflow:Step 10, mean_score: 0.000000\n","INFO:tensorflow:Step 15, mean_score: 0.000000\n","INFO:tensorflow:Step 20, mean_score: 0.000000\n","INFO:tensorflow:Step 25, mean_score: 0.000000\n","INFO:tensorflow:Step 30, mean_score: 0.000000\n","INFO:tensorflow:Step 35, mean_score: 0.000000\n","INFO:tensorflow:Step 40, mean_score: 0.000000\n","INFO:tensorflow:Step 45, mean_score: 0.000000\n","INFO:tensorflow:Step 50, mean_score: 0.000000\n","INFO:tensorflow:Step 55, mean_score: 0.000000\n","INFO:tensorflow:Step 60, mean_score: -0.015625\n","INFO:tensorflow:Step 65, mean_score: -0.078125\n","INFO:tensorflow:Step 70, mean_score: -0.078125\n","INFO:tensorflow:Step 75, mean_score: -0.078125\n","INFO:tensorflow:Step 80, mean_score: -0.078125\n","INFO:tensorflow:Step 85, mean_score: -0.078125\n","INFO:tensorflow:Step 90, mean_score: 0.484375\n","INFO:tensorflow:Step 95, mean_score: 0.843750\n","INFO:tensorflow:Step 100, mean_score: 0.828125\n","INFO:tensorflow:Step 105, mean_score: 0.828125\n","INFO:tensorflow:Step 110, mean_score: 0.828125\n","INFO:tensorflow:Step 115, mean_score: 0.828125\n","INFO:tensorflow:Step 120, mean_score: 0.828125\n","INFO:tensorflow:Step 125, mean_score: 0.828125\n","INFO:tensorflow:Step 130, mean_score: 0.812500\n","INFO:tensorflow:Step 135, mean_score: 0.812500\n","INFO:tensorflow:Step 140, mean_score: 0.812500\n","INFO:tensorflow:Step 145, mean_score: 0.812500\n","INFO:tensorflow:Step 150, mean_score: 0.812500\n","INFO:tensorflow:Step 155, mean_score: 0.812500\n","INFO:tensorflow:Step 160, mean_score: 0.812500\n","INFO:tensorflow:Step 165, mean_score: 0.812500\n","INFO:tensorflow:Step 170, mean_score: 0.828125\n","INFO:tensorflow:Step 175, mean_score: 0.843750\n","INFO:tensorflow:Step 180, mean_score: 0.843750\n","INFO:tensorflow:Step 185, mean_score: 0.843750\n","INFO:tensorflow:Step 190, mean_score: 1.140625\n","INFO:tensorflow:Step 195, mean_score: 1.765625\n","INFO:tensorflow:Step 200, mean_score: 1.765625\n","INFO:tensorflow:Step 205, mean_score: 1.765625\n","INFO:tensorflow:Step 210, mean_score: 1.781250\n","INFO:tensorflow:Step 215, mean_score: 1.781250\n","INFO:tensorflow:Step 220, mean_score: 1.765625\n","INFO:tensorflow:Step 225, mean_score: 1.765625\n","INFO:tensorflow:Step 230, mean_score: 1.765625\n","INFO:tensorflow:Step 235, mean_score: 1.765625\n","INFO:tensorflow:Step 240, mean_score: 1.765625\n","INFO:tensorflow:Step 245, mean_score: 1.765625\n","INFO:tensorflow:Step 250, mean_score: 1.765625\n","INFO:tensorflow:Step 255, mean_score: 1.750000\n","INFO:tensorflow:Step 260, mean_score: 1.750000\n","INFO:tensorflow:Step 265, mean_score: 1.750000\n","INFO:tensorflow:Step 270, mean_score: 2.312500\n","INFO:tensorflow:Step 275, mean_score: 2.687500\n","INFO:tensorflow:Step 280, mean_score: 2.703125\n","INFO:tensorflow:Step 285, mean_score: 2.703125\n","INFO:tensorflow:Step 290, mean_score: 2.703125\n","INFO:tensorflow:Step 295, mean_score: 2.703125\n","INFO:tensorflow:Step 300, mean_score: 2.703125\n","INFO:tensorflow:Step 305, mean_score: 2.703125\n","INFO:tensorflow:Step 310, mean_score: 2.718750\n","INFO:tensorflow:Step 315, mean_score: 2.718750\n","INFO:tensorflow:Step 320, mean_score: 2.718750\n","INFO:tensorflow:Step 325, mean_score: 2.718750\n","INFO:tensorflow:Step 330, mean_score: 2.734375\n","INFO:tensorflow:Step 335, mean_score: 2.734375\n","INFO:tensorflow:Step 340, mean_score: 2.734375\n","INFO:tensorflow:Step 345, mean_score: 2.734375\n","INFO:tensorflow:Step 350, mean_score: 2.750000\n","INFO:tensorflow:Step 355, mean_score: 2.765625\n","INFO:tensorflow:Step 360, mean_score: 2.765625\n","INFO:tensorflow:Step 365, mean_score: 2.765625\n","INFO:tensorflow:Step 370, mean_score: 3.062500\n","INFO:tensorflow:Step 375, mean_score: 3.687500\n","INFO:tensorflow:Step 380, mean_score: 3.687500\n","INFO:tensorflow:Step 385, mean_score: 3.687500\n","INFO:tensorflow:Step 390, mean_score: 3.703125\n","INFO:tensorflow:Step 395, mean_score: 3.703125\n","INFO:tensorflow:Step 400, mean_score: 3.703125\n","INFO:tensorflow:Step 405, mean_score: 3.703125\n","INFO:tensorflow:Step 410, mean_score: 3.687500\n","INFO:tensorflow:Step 415, mean_score: 3.687500\n","INFO:tensorflow:Step 420, mean_score: 3.687500\n","INFO:tensorflow:Step 425, mean_score: 3.687500\n","INFO:tensorflow:Step 430, mean_score: 3.703125\n","INFO:tensorflow:Step 435, mean_score: 3.703125\n","INFO:tensorflow:Step 440, mean_score: 3.703125\n","INFO:tensorflow:Step 445, mean_score: 3.703125\n","INFO:tensorflow:Step 450, mean_score: 4.265625\n","INFO:tensorflow:Step 455, mean_score: 4.640625\n","INFO:tensorflow:Step 460, mean_score: 4.656250\n","INFO:tensorflow:Step 465, mean_score: 4.656250\n","INFO:tensorflow:Step 470, mean_score: 4.656250\n","INFO:tensorflow:Step 475, mean_score: 4.656250\n","INFO:tensorflow:Step 480, mean_score: 4.656250\n","INFO:tensorflow:Step 485, mean_score: 4.656250\n","INFO:tensorflow:Step 490, mean_score: 4.671875\n","INFO:tensorflow:Step 495, mean_score: 4.671875\n","INFO:tensorflow:Step 500, mean_score: 4.671875\n","INFO:tensorflow:Step 505, mean_score: 4.671875\n","INFO:tensorflow:Step 510, mean_score: 4.687500\n","INFO:tensorflow:Step 515, mean_score: 4.687500\n","INFO:tensorflow:Step 520, mean_score: 4.703125\n","INFO:tensorflow:Step 525, mean_score: 4.703125\n","INFO:tensorflow:Step 530, mean_score: 4.718750\n","INFO:tensorflow:Step 535, mean_score: 4.734375\n","INFO:tensorflow:Step 540, mean_score: 4.734375\n","INFO:tensorflow:Step 545, mean_score: 4.734375\n","INFO:tensorflow:Step 550, mean_score: 5.031250\n","INFO:tensorflow:Step 555, mean_score: 5.656250\n","INFO:tensorflow:Step 560, mean_score: 5.656250\n","INFO:tensorflow:Step 565, mean_score: 5.656250\n","INFO:tensorflow:Step 570, mean_score: 5.671875\n","INFO:tensorflow:Step 575, mean_score: 5.671875\n","INFO:tensorflow:Step 580, mean_score: 5.671875\n","INFO:tensorflow:Step 585, mean_score: 5.671875\n","INFO:tensorflow:Step 590, mean_score: 5.671875\n","INFO:tensorflow:Step 595, mean_score: 5.671875\n","INFO:tensorflow:Step 600, mean_score: 5.671875\n","INFO:tensorflow:Step 605, mean_score: 5.671875\n","INFO:tensorflow:Step 610, mean_score: 5.687500\n","INFO:tensorflow:Step 615, mean_score: 5.687500\n","INFO:tensorflow:Step 620, mean_score: 5.703125\n","INFO:tensorflow:Step 625, mean_score: 5.703125\n","INFO:tensorflow:Step 630, mean_score: 6.265625\n","INFO:tensorflow:Step 635, mean_score: 6.640625\n","INFO:tensorflow:Step 640, mean_score: 6.656250\n","INFO:tensorflow:Step 645, mean_score: 6.656250\n","INFO:tensorflow:Step 650, mean_score: 6.656250\n","INFO:tensorflow:Step 655, mean_score: 6.656250\n","INFO:tensorflow:Step 660, mean_score: 6.656250\n","INFO:tensorflow:Step 665, mean_score: 6.656250\n","INFO:tensorflow:Step 670, mean_score: 6.671875\n","INFO:tensorflow:Step 675, mean_score: 6.671875\n","INFO:tensorflow:Step 680, mean_score: 6.671875\n","INFO:tensorflow:Step 685, mean_score: 6.671875\n","INFO:tensorflow:Step 690, mean_score: 6.687500\n","INFO:tensorflow:Step 695, mean_score: 6.687500\n","INFO:tensorflow:Step 700, mean_score: 6.703125\n","INFO:tensorflow:Step 705, mean_score: 6.703125\n","INFO:tensorflow:Step 710, mean_score: 6.718750\n","INFO:tensorflow:Step 715, mean_score: 6.734375\n","INFO:tensorflow:Step 720, mean_score: 6.734375\n","INFO:tensorflow:Step 725, mean_score: 6.734375\n","INFO:tensorflow:Step 730, mean_score: 7.031250\n","INFO:tensorflow:Step 735, mean_score: 7.656250\n","INFO:tensorflow:Step 740, mean_score: 7.656250\n","INFO:tensorflow:Step 745, mean_score: 7.656250\n","INFO:tensorflow:Step 750, mean_score: 7.671875\n","INFO:tensorflow:Step 755, mean_score: 7.671875\n","INFO:tensorflow:Step 760, mean_score: 7.671875\n","INFO:tensorflow:Step 765, mean_score: 7.671875\n","INFO:tensorflow:Step 770, mean_score: 7.671875\n","INFO:tensorflow:Step 775, mean_score: 7.671875\n","INFO:tensorflow:Step 780, mean_score: 7.671875\n","INFO:tensorflow:Step 785, mean_score: 7.671875\n","INFO:tensorflow:Step 790, mean_score: 7.687500\n","INFO:tensorflow:Step 795, mean_score: 7.687500\n","INFO:tensorflow:Step 800, mean_score: 7.703125\n","INFO:tensorflow:Step 805, mean_score: 7.703125\n","INFO:tensorflow:Step 810, mean_score: 8.265625\n","INFO:tensorflow:Step 815, mean_score: 8.640625\n","INFO:tensorflow:Step 820, mean_score: 8.656250\n","INFO:tensorflow:Step 825, mean_score: 8.656250\n","INFO:tensorflow:Step 830, mean_score: 8.656250\n","INFO:tensorflow:Step 835, mean_score: 8.656250\n","INFO:tensorflow:Step 840, mean_score: 8.656250\n","INFO:tensorflow:Step 845, mean_score: 8.656250\n","INFO:tensorflow:Step 850, mean_score: 8.671875\n","INFO:tensorflow:Step 855, mean_score: 8.671875\n","INFO:tensorflow:Step 860, mean_score: 8.671875\n","INFO:tensorflow:Step 865, mean_score: 8.671875\n","INFO:tensorflow:Step 870, mean_score: 8.687500\n","INFO:tensorflow:Step 875, mean_score: 8.687500\n","INFO:tensorflow:Step 880, mean_score: 8.703125\n","INFO:tensorflow:Step 885, mean_score: 8.703125\n","INFO:tensorflow:Step 890, mean_score: 8.718750\n","INFO:tensorflow:Step 895, mean_score: 8.734375\n","INFO:tensorflow:Step 900, mean_score: 8.734375\n","INFO:tensorflow:Step 905, mean_score: 8.734375\n","INFO:tensorflow:Step 910, mean_score: 9.031250\n","INFO:tensorflow:Step 915, mean_score: 9.656250\n","INFO:tensorflow:Step 920, mean_score: 9.656250\n","INFO:tensorflow:Step 925, mean_score: 9.656250\n","INFO:tensorflow:Step 930, mean_score: 9.671875\n","INFO:tensorflow:Step 935, mean_score: 9.671875\n","INFO:tensorflow:Step 940, mean_score: 9.671875\n","INFO:tensorflow:Step 945, mean_score: 9.671875\n","INFO:tensorflow:Step 950, mean_score: 9.671875\n","INFO:tensorflow:Step 955, mean_score: 9.671875\n","INFO:tensorflow:Step 960, mean_score: 9.671875\n","INFO:tensorflow:Step 965, mean_score: 9.671875\n","INFO:tensorflow:Step 970, mean_score: 9.687500\n","INFO:tensorflow:Step 975, mean_score: 9.687500\n","INFO:tensorflow:Step 980, mean_score: 9.703125\n","INFO:tensorflow:Step 985, mean_score: 9.703125\n","INFO:tensorflow:Step 990, mean_score: 10.265625\n","INFO:tensorflow:Step 995, mean_score: 10.640625\n","INFO:tensorflow:Step 1000, mean_score: 10.656250\n","INFO:tensorflow:Step 1005, mean_score: 10.656250\n","INFO:tensorflow:Step 1010, mean_score: 10.656250\n","INFO:tensorflow:Step 1015, mean_score: 10.656250\n","INFO:tensorflow:Step 1020, mean_score: 10.656250\n","INFO:tensorflow:Step 1025, mean_score: 10.656250\n","INFO:tensorflow:Step 1030, mean_score: 10.671875\n","INFO:tensorflow:Step 1035, mean_score: 10.671875\n","INFO:tensorflow:Step 1040, mean_score: 10.671875\n","INFO:tensorflow:Step 1045, mean_score: 10.671875\n","INFO:tensorflow:Step 1050, mean_score: 10.687500\n","INFO:tensorflow:Step 1055, mean_score: 10.687500\n","INFO:tensorflow:Step 1060, mean_score: 10.703125\n","INFO:tensorflow:Step 1065, mean_score: 10.703125\n","INFO:tensorflow:Step 1070, mean_score: 10.718750\n","INFO:tensorflow:Step 1075, mean_score: 10.734375\n","INFO:tensorflow:Step 1080, mean_score: 10.734375\n","INFO:tensorflow:Step 1085, mean_score: 10.734375\n","INFO:tensorflow:Step 1090, mean_score: 11.031250\n","INFO:tensorflow:Step 1095, mean_score: 11.656250\n","INFO:tensorflow:Step 1100, mean_score: 11.656250\n","INFO:tensorflow:Step 1105, mean_score: 11.656250\n","INFO:tensorflow:Step 1110, mean_score: 11.671875\n","INFO:tensorflow:Step 1115, mean_score: 11.671875\n","INFO:tensorflow:Step 1120, mean_score: 11.671875\n","INFO:tensorflow:Step 1125, mean_score: 11.671875\n","INFO:tensorflow:Step 1130, mean_score: 11.671875\n","INFO:tensorflow:Step 1135, mean_score: 11.671875\n","INFO:tensorflow:Step 1140, mean_score: 11.671875\n","INFO:tensorflow:Step 1145, mean_score: 11.671875\n","INFO:tensorflow:Step 1150, mean_score: 11.687500\n","INFO:tensorflow:Step 1155, mean_score: 11.687500\n","INFO:tensorflow:Step 1160, mean_score: 11.703125\n","INFO:tensorflow:Step 1165, mean_score: 11.703125\n","INFO:tensorflow:Step 1170, mean_score: 12.265625\n","INFO:tensorflow:Step 1175, mean_score: 12.640625\n","INFO:tensorflow:Step 1180, mean_score: 12.656250\n","INFO:tensorflow:Step 1185, mean_score: 12.656250\n","INFO:tensorflow:Step 1190, mean_score: 12.656250\n","INFO:tensorflow:Step 1195, mean_score: 12.656250\n","INFO:tensorflow:Step 1200, mean_score: 12.656250\n","INFO:tensorflow:Step 1205, mean_score: 12.656250\n","INFO:tensorflow:Step 1210, mean_score: 12.671875\n","INFO:tensorflow:Step 1215, mean_score: 12.671875\n","INFO:tensorflow:Step 1220, mean_score: 12.671875\n","INFO:tensorflow:Step 1225, mean_score: 12.671875\n","INFO:tensorflow:Step 1230, mean_score: 12.687500\n","INFO:tensorflow:Step 1235, mean_score: 12.687500\n","INFO:tensorflow:Step 1240, mean_score: 12.703125\n","INFO:tensorflow:Step 1245, mean_score: 12.703125\n","INFO:tensorflow:Step 1250, mean_score: 12.718750\n","INFO:tensorflow:Step 1255, mean_score: 12.734375\n","INFO:tensorflow:Step 1260, mean_score: 12.734375\n","INFO:tensorflow:Step 1265, mean_score: 12.734375\n","INFO:tensorflow:Step 1270, mean_score: 13.031250\n","INFO:tensorflow:Step 1275, mean_score: 13.656250\n","INFO:tensorflow:Step 1280, mean_score: 13.656250\n","INFO:tensorflow:Step 1285, mean_score: 13.656250\n","INFO:tensorflow:Step 1290, mean_score: 13.671875\n","INFO:tensorflow:Step 1295, mean_score: 13.671875\n","INFO:tensorflow:Step 1300, mean_score: 13.671875\n","INFO:tensorflow:Step 1305, mean_score: 13.671875\n","INFO:tensorflow:Step 1310, mean_score: 13.671875\n","INFO:tensorflow:Step 1315, mean_score: 13.671875\n","INFO:tensorflow:Step 1320, mean_score: 13.671875\n","INFO:tensorflow:Step 1325, mean_score: 13.671875\n","INFO:tensorflow:Step 1330, mean_score: 13.687500\n","INFO:tensorflow:Step 1335, mean_score: 13.687500\n","INFO:tensorflow:Step 1340, mean_score: 13.703125\n","INFO:tensorflow:Step 1345, mean_score: 13.703125\n","INFO:tensorflow:Step 1350, mean_score: 14.265625\n","INFO:tensorflow:Step 1355, mean_score: 14.640625\n","INFO:tensorflow:Step 1360, mean_score: 14.656250\n","INFO:tensorflow:Step 1365, mean_score: 14.656250\n","INFO:tensorflow:Step 1370, mean_score: 14.656250\n","INFO:tensorflow:Step 1375, mean_score: 14.656250\n","INFO:tensorflow:Step 1380, mean_score: 14.656250\n","INFO:tensorflow:Step 1385, mean_score: 14.656250\n","INFO:tensorflow:Step 1390, mean_score: 14.671875\n","INFO:tensorflow:Step 1395, mean_score: 14.671875\n","INFO:tensorflow:Step 1400, mean_score: 14.671875\n","INFO:tensorflow:Step 1405, mean_score: 14.671875\n","INFO:tensorflow:Step 1410, mean_score: 14.687500\n","INFO:tensorflow:Step 1415, mean_score: 14.687500\n","INFO:tensorflow:Step 1420, mean_score: 14.703125\n","INFO:tensorflow:Step 1425, mean_score: 14.703125\n","INFO:tensorflow:Step 1430, mean_score: 14.718750\n","INFO:tensorflow:Step 1435, mean_score: 14.734375\n","INFO:tensorflow:Step 1440, mean_score: 14.734375\n","INFO:tensorflow:Step 1445, mean_score: 14.734375\n","INFO:tensorflow:Step 1450, mean_score: 15.031250\n","INFO:tensorflow:Step 1455, mean_score: 15.656250\n","INFO:tensorflow:Step 1460, mean_score: 15.656250\n","INFO:tensorflow:Step 1465, mean_score: 15.656250\n","INFO:tensorflow:Step 1470, mean_score: 15.671875\n","INFO:tensorflow:Step 1475, mean_score: 15.671875\n","INFO:tensorflow:Step 1480, mean_score: 15.671875\n","INFO:tensorflow:Step 1485, mean_score: 15.671875\n","INFO:tensorflow:Step 1490, mean_score: 15.671875\n","INFO:tensorflow:Step 1495, mean_score: 15.671875\n","INFO:tensorflow:Step 1500, mean_score: 15.671875\n","INFO:tensorflow:Step 1505, mean_score: 15.671875\n","INFO:tensorflow:Step 1510, mean_score: 15.687500\n","INFO:tensorflow:Step 1515, mean_score: 15.687500\n","INFO:tensorflow:Step 1520, mean_score: 15.703125\n","INFO:tensorflow:Step 1525, mean_score: 15.703125\n","INFO:tensorflow:Step 1530, mean_score: 16.265625\n","INFO:tensorflow:Step 1535, mean_score: 16.640625\n","INFO:tensorflow:Step 1540, mean_score: 16.656250\n","INFO:tensorflow:Step 1545, mean_score: 16.656250\n","INFO:tensorflow:Step 1550, mean_score: 16.656250\n","INFO:tensorflow:Step 1555, mean_score: 16.656250\n","INFO:tensorflow:Step 1560, mean_score: 16.656250\n","INFO:tensorflow:Step 1565, mean_score: 16.656250\n","INFO:tensorflow:Step 1570, mean_score: 16.671875\n","INFO:tensorflow:Step 1575, mean_score: 16.671875\n","INFO:tensorflow:Step 1580, mean_score: 16.671875\n","INFO:tensorflow:Step 1585, mean_score: 16.671875\n","INFO:tensorflow:Step 1590, mean_score: 16.687500\n","INFO:tensorflow:Step 1595, mean_score: 16.687500\n","INFO:tensorflow:Step 1600, mean_score: 16.703125\n","INFO:tensorflow:Step 1605, mean_score: 16.703125\n","INFO:tensorflow:Step 1610, mean_score: 16.718750\n","INFO:tensorflow:Step 1615, mean_score: 16.734375\n","INFO:tensorflow:Step 1620, mean_score: 16.734375\n","INFO:tensorflow:Step 1625, mean_score: 16.734375\n","INFO:tensorflow:Step 1630, mean_score: 17.031250\n","INFO:tensorflow:Step 1635, mean_score: 17.656250\n","INFO:tensorflow:Step 1640, mean_score: 17.656250\n","INFO:tensorflow:Step 1645, mean_score: 17.656250\n","INFO:tensorflow:Step 1650, mean_score: 17.671875\n","INFO:tensorflow:Step 1655, mean_score: 17.671875\n","INFO:tensorflow:Step 1660, mean_score: 17.671875\n","INFO:tensorflow:Step 1665, mean_score: 17.671875\n","INFO:tensorflow:Step 1670, mean_score: 17.671875\n","INFO:tensorflow:Step 1675, mean_score: 17.671875\n","INFO:tensorflow:Step 1680, mean_score: 17.671875\n","INFO:tensorflow:Step 1685, mean_score: 17.671875\n","INFO:tensorflow:Step 1690, mean_score: 17.687500\n","INFO:tensorflow:Step 1695, mean_score: 17.687500\n","INFO:tensorflow:Step 1700, mean_score: 17.703125\n","INFO:tensorflow:Step 1705, mean_score: 17.703125\n","INFO:tensorflow:Step 1710, mean_score: 18.265625\n","INFO:tensorflow:Step 1715, mean_score: 18.640625\n","INFO:tensorflow:Step 1720, mean_score: 18.656250\n","INFO:tensorflow:Step 1725, mean_score: 18.656250\n","INFO:tensorflow:Step 1730, mean_score: 18.656250\n","INFO:tensorflow:Step 1735, mean_score: 18.656250\n","INFO:tensorflow:Step 1740, mean_score: 18.656250\n","INFO:tensorflow:Step 1745, mean_score: 18.656250\n","INFO:tensorflow:Step 1750, mean_score: 18.671875\n","INFO:tensorflow:Step 1755, mean_score: 18.671875\n","INFO:tensorflow:Step 1760, mean_score: 18.671875\n","INFO:tensorflow:Step 1765, mean_score: 18.671875\n","INFO:tensorflow:Step 1770, mean_score: 18.687500\n","INFO:tensorflow:Step 1775, mean_score: 18.687500\n","INFO:tensorflow:Step 1780, mean_score: 18.703125\n","INFO:tensorflow:Step 1785, mean_score: 18.703125\n","INFO:tensorflow:Step 1790, mean_score: 18.718750\n","INFO:tensorflow:Step 1795, mean_score: 18.734375\n","INFO:tensorflow:Step 1800, mean_score: 18.734375\n","INFO:tensorflow:Step 1805, mean_score: 18.734375\n","INFO:tensorflow:Step 1810, mean_score: 19.031250\n","INFO:tensorflow:Step 1815, mean_score: 19.656250\n","INFO:tensorflow:Step 1820, mean_score: 19.656250\n","INFO:tensorflow:Step 1825, mean_score: 19.656250\n","INFO:tensorflow:Step 1830, mean_score: 19.671875\n","INFO:tensorflow:Step 1835, mean_score: 19.671875\n","INFO:tensorflow:Step 1840, mean_score: 19.671875\n","INFO:tensorflow:Step 1845, mean_score: 19.671875\n","INFO:tensorflow:Step 1850, mean_score: 19.671875\n","INFO:tensorflow:Step 1855, mean_score: 19.671875\n","INFO:tensorflow:Step 1860, mean_score: 19.671875\n","INFO:tensorflow:Step 1865, mean_score: 19.671875\n","INFO:tensorflow:Step 1870, mean_score: 19.687500\n","INFO:tensorflow:Step 1875, mean_score: 19.687500\n","INFO:tensorflow:Step 1880, mean_score: 19.703125\n","INFO:tensorflow:Step 1885, mean_score: 19.703125\n","INFO:tensorflow:Step 1890, mean_score: 19.703125\n","INFO:tensorflow:Step 1895, mean_score: 19.718750\n","INFO:tensorflow:Step 1900, mean_score: 19.734375\n","INFO:tensorflow:Step 1905, mean_score: 19.734375\n","INFO:tensorflow:Step 1910, mean_score: 19.734375\n","INFO:tensorflow:Step 1915, mean_score: 19.734375\n","INFO:tensorflow:Step 1920, mean_score: 19.734375\n","INFO:tensorflow:Step 1925, mean_score: 19.734375\n","INFO:tensorflow:Step 1930, mean_score: 19.750000\n","INFO:tensorflow:Step 1935, mean_score: 19.750000\n","INFO:tensorflow:Step 1940, mean_score: 19.750000\n","INFO:tensorflow:Step 1945, mean_score: 19.750000\n","INFO:tensorflow:Step 1950, mean_score: 19.765625\n","INFO:tensorflow:Step 1955, mean_score: 19.765625\n","INFO:tensorflow:Step 1960, mean_score: 19.781250\n","INFO:tensorflow:Step 1965, mean_score: 19.781250\n","INFO:tensorflow:Step 1970, mean_score: 19.781250\n","INFO:tensorflow:Step 1975, mean_score: 19.781250\n","INFO:tensorflow:Step 1980, mean_score: 19.781250\n","INFO:tensorflow:Step 1985, mean_score: 19.781250\n","INFO:tensorflow:Step 1990, mean_score: 19.781250\n","INFO:tensorflow:Step 1995, mean_score: 19.781250\n","INFO:tensorflow:Step 2000, mean_score: 19.781250\n","INFO:tensorflow:Step 2005, mean_score: 19.781250\n","INFO:tensorflow:Step 2010, mean_score: 19.781250\n","INFO:tensorflow:Step 2015, mean_score: 19.781250\n","INFO:tensorflow:Step 2020, mean_score: 19.781250\n","INFO:tensorflow:Step 2025, mean_score: 19.781250\n","INFO:tensorflow:Step 2030, mean_score: 19.781250\n","INFO:tensorflow:Step 2035, mean_score: 19.781250\n","INFO:tensorflow:Step 2040, mean_score: 19.781250\n","INFO:tensorflow:Step 2045, mean_score: 19.781250\n","INFO:tensorflow:Step 2050, mean_score: 19.796875\n","INFO:tensorflow:Step 2055, mean_score: 19.796875\n","INFO:tensorflow:Step 2060, mean_score: 19.812500\n","INFO:tensorflow:Step 2065, mean_score: 19.812500\n","INFO:tensorflow:Step 2070, mean_score: 19.812500\n","INFO:tensorflow:Step 2075, mean_score: 19.812500\n","INFO:tensorflow:Step 2080, mean_score: 19.812500\n","INFO:tensorflow:Step 2085, mean_score: 19.812500\n","INFO:tensorflow:Step 2090, mean_score: 19.812500\n","INFO:tensorflow:Step 2095, mean_score: 19.812500\n","INFO:tensorflow:Step 2100, mean_score: 19.812500\n","INFO:tensorflow:Step 2105, mean_score: 19.812500\n","INFO:tensorflow:Step 2110, mean_score: 19.812500\n","INFO:tensorflow:Step 2115, mean_score: 19.812500\n","INFO:tensorflow:Step 2120, mean_score: 19.812500\n","INFO:tensorflow:Step 2125, mean_score: 19.812500\n","INFO:tensorflow:Step 2130, mean_score: 19.812500\n","INFO:tensorflow:Step 2135, mean_score: 19.812500\n","INFO:tensorflow:Step 2140, mean_score: 19.828125\n","INFO:tensorflow:Step 2145, mean_score: 19.828125\n","INFO:tensorflow:Step 2150, mean_score: 19.828125\n","INFO:tensorflow:Step 2155, mean_score: 19.828125\n","INFO:tensorflow:Step 2160, mean_score: 19.828125\n","INFO:tensorflow:Step 2165, mean_score: 19.828125\n","INFO:tensorflow:Step 2170, mean_score: 19.828125\n","INFO:tensorflow:Step 2175, mean_score: 19.828125\n","INFO:tensorflow:Step 2180, mean_score: 19.828125\n","INFO:tensorflow:Step 2185, mean_score: 19.828125\n","INFO:tensorflow:Step 2190, mean_score: 19.828125\n","INFO:tensorflow:Step 2195, mean_score: 19.828125\n","INFO:tensorflow:Step 2200, mean_score: 19.828125\n","INFO:tensorflow:Step 2205, mean_score: 19.828125\n","INFO:tensorflow:Step 2210, mean_score: 19.828125\n","INFO:tensorflow:Step 2215, mean_score: 19.828125\n","INFO:tensorflow:Step 2220, mean_score: 19.828125\n","INFO:tensorflow:Step 2225, mean_score: 19.828125\n","INFO:tensorflow:Step 2230, mean_score: 19.828125\n","INFO:tensorflow:Step 2235, mean_score: 19.828125\n","INFO:tensorflow:Step 2240, mean_score: 19.843750\n","INFO:tensorflow:Step 2245, mean_score: 19.843750\n","INFO:tensorflow:Step 2250, mean_score: 19.843750\n","INFO:tensorflow:Step 2255, mean_score: 19.843750\n","INFO:tensorflow:Step 2260, mean_score: 19.843750\n","INFO:tensorflow:Step 2265, mean_score: 19.843750\n","INFO:tensorflow:Step 2270, mean_score: 19.843750\n","INFO:tensorflow:Step 2275, mean_score: 19.843750\n","INFO:tensorflow:Step 2280, mean_score: 19.843750\n","INFO:tensorflow:Step 2285, mean_score: 19.843750\n","INFO:tensorflow:Step 2290, mean_score: 19.843750\n","INFO:tensorflow:Step 2295, mean_score: 19.843750\n","INFO:tensorflow:Step 2300, mean_score: 19.843750\n","INFO:tensorflow:Step 2305, mean_score: 19.843750\n","INFO:tensorflow:Step 2310, mean_score: 19.843750\n","INFO:tensorflow:Step 2315, mean_score: 19.843750\n","INFO:tensorflow:Evaluating metric mean_reward/eval/sampling_temp_0.5_max_noops_0_unclipped\n","2019-03-22 16:12:57.935045: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 16:12:57.935160: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 16:12:57.935189: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 16:12:57.935209: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 16:12:57.935553: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","INFO:tensorflow:Using DummyPolicyProblem for the policy.\n","INFO:tensorflow:Setting T2TModel mode to 'train'\n","INFO:tensorflow:Using variable initializer: orthogonal\n","INFO:tensorflow:Transforming feature 'input_action' with symbol_modality_6_64.bottom\n","INFO:tensorflow:Transforming feature 'input_reward' with symbol_modality_3_64.bottom\n","INFO:tensorflow:Transforming feature 'inputs' with video_modality.bottom\n","INFO:tensorflow:Transforming feature 'target_action' with symbol_modality_6_64.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_policy' with identity_modality.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_reward' with symbol_modality_3_64.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_value' with identity_modality.targets_bottom\n","INFO:tensorflow:Transforming feature 'targets' with video_modality.targets_bottom\n","INFO:tensorflow:Building model body\n","INFO:tensorflow:Transforming body output with identity_modality.top\n","INFO:tensorflow:Transforming body output with identity_modality.top\n","2019-03-22 16:13:12.260846: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 16:13:12.260981: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 16:13:12.261059: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 16:13:12.261099: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 16:13:12.261613: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","2019-03-22 16:13:12.493082: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","INFO:tensorflow:Restoring checkpoint gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/policy/model.ckpt-171992\n","2019-03-22 16:13:12.556955: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","INFO:tensorflow:Restoring parameters from gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/policy/model.ckpt-171992\n","2019-03-22 16:13:12.651009: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:13:12.715180: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:13:12.816774: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","INFO:tensorflow:Step 5, mean_score: 0.000000\n","INFO:tensorflow:Step 10, mean_score: 0.000000\n","INFO:tensorflow:Step 15, mean_score: 0.000000\n","INFO:tensorflow:Step 20, mean_score: 0.000000\n","INFO:tensorflow:Step 25, mean_score: 0.000000\n","INFO:tensorflow:Step 30, mean_score: 0.000000\n","INFO:tensorflow:Step 35, mean_score: 0.000000\n","INFO:tensorflow:Step 40, mean_score: 0.000000\n","INFO:tensorflow:Step 45, mean_score: 0.000000\n","INFO:tensorflow:Step 50, mean_score: 0.000000\n","INFO:tensorflow:Step 55, mean_score: 0.000000\n","INFO:tensorflow:Step 60, mean_score: 0.000000\n","INFO:tensorflow:Step 65, mean_score: -0.031250\n","INFO:tensorflow:Step 70, mean_score: -0.031250\n","INFO:tensorflow:Step 75, mean_score: -0.031250\n","INFO:tensorflow:Step 80, mean_score: -0.031250\n","INFO:tensorflow:Step 85, mean_score: -0.031250\n","INFO:tensorflow:Step 90, mean_score: -0.031250\n","INFO:tensorflow:Step 95, mean_score: 0.937500\n","INFO:tensorflow:Step 100, mean_score: 0.921875\n","INFO:tensorflow:Step 105, mean_score: 0.921875\n","INFO:tensorflow:Step 110, mean_score: 0.921875\n","INFO:tensorflow:Step 115, mean_score: 0.921875\n","INFO:tensorflow:Step 120, mean_score: 0.921875\n","INFO:tensorflow:Step 125, mean_score: 0.921875\n","INFO:tensorflow:Step 130, mean_score: 0.921875\n","INFO:tensorflow:Step 135, mean_score: 0.921875\n","INFO:tensorflow:Step 140, mean_score: 0.921875\n","INFO:tensorflow:Step 145, mean_score: 0.921875\n","INFO:tensorflow:Step 150, mean_score: 0.921875\n","INFO:tensorflow:Step 155, mean_score: 0.921875\n","INFO:tensorflow:Step 160, mean_score: 0.921875\n","INFO:tensorflow:Step 165, mean_score: 0.906250\n","INFO:tensorflow:Step 170, mean_score: 0.906250\n","INFO:tensorflow:Step 175, mean_score: 0.921875\n","INFO:tensorflow:Step 180, mean_score: 0.921875\n","INFO:tensorflow:Step 185, mean_score: 0.921875\n","INFO:tensorflow:Step 190, mean_score: 0.921875\n","INFO:tensorflow:Step 195, mean_score: 0.921875\n","INFO:tensorflow:Step 200, mean_score: 1.890625\n","INFO:tensorflow:Step 205, mean_score: 1.890625\n","INFO:tensorflow:Step 210, mean_score: 1.890625\n","INFO:tensorflow:Step 215, mean_score: 1.890625\n","INFO:tensorflow:Step 220, mean_score: 1.890625\n","INFO:tensorflow:Step 225, mean_score: 1.890625\n","INFO:tensorflow:Step 230, mean_score: 1.890625\n","INFO:tensorflow:Step 235, mean_score: 1.890625\n","INFO:tensorflow:Step 240, mean_score: 1.890625\n","INFO:tensorflow:Step 245, mean_score: 1.890625\n","INFO:tensorflow:Step 250, mean_score: 1.890625\n","INFO:tensorflow:Step 255, mean_score: 1.890625\n","INFO:tensorflow:Step 260, mean_score: 1.890625\n","INFO:tensorflow:Step 265, mean_score: 1.890625\n","INFO:tensorflow:Step 270, mean_score: 1.890625\n","INFO:tensorflow:Step 275, mean_score: 2.875000\n","INFO:tensorflow:Step 280, mean_score: 2.890625\n","INFO:tensorflow:Step 285, mean_score: 2.890625\n","INFO:tensorflow:Step 290, mean_score: 2.890625\n","INFO:tensorflow:Step 295, mean_score: 2.890625\n","INFO:tensorflow:Step 300, mean_score: 2.890625\n","INFO:tensorflow:Step 305, mean_score: 2.890625\n","INFO:tensorflow:Step 310, mean_score: 2.890625\n","INFO:tensorflow:Step 315, mean_score: 2.890625\n","INFO:tensorflow:Step 320, mean_score: 2.890625\n","INFO:tensorflow:Step 325, mean_score: 2.890625\n","INFO:tensorflow:Step 330, mean_score: 2.890625\n","INFO:tensorflow:Step 335, mean_score: 2.890625\n","INFO:tensorflow:Step 340, mean_score: 2.890625\n","INFO:tensorflow:Step 345, mean_score: 2.890625\n","INFO:tensorflow:Step 350, mean_score: 2.890625\n","INFO:tensorflow:Step 355, mean_score: 2.906250\n","INFO:tensorflow:Step 360, mean_score: 2.906250\n","INFO:tensorflow:Step 365, mean_score: 2.906250\n","INFO:tensorflow:Step 370, mean_score: 2.906250\n","INFO:tensorflow:Step 375, mean_score: 2.921875\n","INFO:tensorflow:Step 380, mean_score: 3.890625\n","INFO:tensorflow:Step 385, mean_score: 3.890625\n","INFO:tensorflow:Step 390, mean_score: 3.890625\n","INFO:tensorflow:Step 395, mean_score: 3.890625\n","INFO:tensorflow:Step 400, mean_score: 3.890625\n","INFO:tensorflow:Step 405, mean_score: 3.890625\n","INFO:tensorflow:Step 410, mean_score: 3.890625\n","INFO:tensorflow:Step 415, mean_score: 3.890625\n","INFO:tensorflow:Step 420, mean_score: 3.890625\n","INFO:tensorflow:Step 425, mean_score: 3.890625\n","INFO:tensorflow:Step 430, mean_score: 3.890625\n","INFO:tensorflow:Step 435, mean_score: 3.890625\n","INFO:tensorflow:Step 440, mean_score: 3.890625\n","INFO:tensorflow:Step 445, mean_score: 3.890625\n","INFO:tensorflow:Step 450, mean_score: 3.890625\n","INFO:tensorflow:Step 455, mean_score: 4.875000\n","INFO:tensorflow:Step 460, mean_score: 4.890625\n","INFO:tensorflow:Step 465, mean_score: 4.890625\n","INFO:tensorflow:Step 470, mean_score: 4.890625\n","INFO:tensorflow:Step 475, mean_score: 4.890625\n","INFO:tensorflow:Step 480, mean_score: 4.890625\n","INFO:tensorflow:Step 485, mean_score: 4.890625\n","INFO:tensorflow:Step 490, mean_score: 4.890625\n","INFO:tensorflow:Step 495, mean_score: 4.890625\n","INFO:tensorflow:Step 500, mean_score: 4.890625\n","INFO:tensorflow:Step 505, mean_score: 4.890625\n","INFO:tensorflow:Step 510, mean_score: 4.890625\n","INFO:tensorflow:Step 515, mean_score: 4.890625\n","INFO:tensorflow:Step 520, mean_score: 4.890625\n","INFO:tensorflow:Step 525, mean_score: 4.890625\n","INFO:tensorflow:Step 530, mean_score: 4.890625\n","INFO:tensorflow:Step 535, mean_score: 4.906250\n","INFO:tensorflow:Step 540, mean_score: 4.906250\n","INFO:tensorflow:Step 545, mean_score: 4.906250\n","INFO:tensorflow:Step 550, mean_score: 4.906250\n","INFO:tensorflow:Step 555, mean_score: 4.921875\n","INFO:tensorflow:Step 560, mean_score: 5.890625\n","INFO:tensorflow:Step 565, mean_score: 5.890625\n","INFO:tensorflow:Step 570, mean_score: 5.890625\n","INFO:tensorflow:Step 575, mean_score: 5.890625\n","INFO:tensorflow:Step 580, mean_score: 5.890625\n","INFO:tensorflow:Step 585, mean_score: 5.890625\n","INFO:tensorflow:Step 590, mean_score: 5.890625\n","INFO:tensorflow:Step 595, mean_score: 5.890625\n","INFO:tensorflow:Step 600, mean_score: 5.890625\n","INFO:tensorflow:Step 605, mean_score: 5.890625\n","INFO:tensorflow:Step 610, mean_score: 5.890625\n","INFO:tensorflow:Step 615, mean_score: 5.890625\n","INFO:tensorflow:Step 620, mean_score: 5.890625\n","INFO:tensorflow:Step 625, mean_score: 5.890625\n","INFO:tensorflow:Step 630, mean_score: 5.890625\n","INFO:tensorflow:Step 635, mean_score: 6.875000\n","INFO:tensorflow:Step 640, mean_score: 6.890625\n","INFO:tensorflow:Step 645, mean_score: 6.890625\n","INFO:tensorflow:Step 650, mean_score: 6.890625\n","INFO:tensorflow:Step 655, mean_score: 6.890625\n","INFO:tensorflow:Step 660, mean_score: 6.890625\n","INFO:tensorflow:Step 665, mean_score: 6.890625\n","INFO:tensorflow:Step 670, mean_score: 6.890625\n","INFO:tensorflow:Step 675, mean_score: 6.890625\n","INFO:tensorflow:Step 680, mean_score: 6.890625\n","INFO:tensorflow:Step 685, mean_score: 6.890625\n","INFO:tensorflow:Step 690, mean_score: 6.890625\n","INFO:tensorflow:Step 695, mean_score: 6.890625\n","INFO:tensorflow:Step 700, mean_score: 6.890625\n","INFO:tensorflow:Step 705, mean_score: 6.890625\n","INFO:tensorflow:Step 710, mean_score: 6.890625\n","INFO:tensorflow:Step 715, mean_score: 6.906250\n","INFO:tensorflow:Step 720, mean_score: 6.906250\n","INFO:tensorflow:Step 725, mean_score: 6.906250\n","INFO:tensorflow:Step 730, mean_score: 6.906250\n","INFO:tensorflow:Step 735, mean_score: 6.921875\n","INFO:tensorflow:Step 740, mean_score: 7.890625\n","INFO:tensorflow:Step 745, mean_score: 7.890625\n","INFO:tensorflow:Step 750, mean_score: 7.890625\n","INFO:tensorflow:Step 755, mean_score: 7.890625\n","INFO:tensorflow:Step 760, mean_score: 7.890625\n","INFO:tensorflow:Step 765, mean_score: 7.890625\n","INFO:tensorflow:Step 770, mean_score: 7.890625\n","INFO:tensorflow:Step 775, mean_score: 7.890625\n","INFO:tensorflow:Step 780, mean_score: 7.890625\n","INFO:tensorflow:Step 785, mean_score: 7.890625\n","INFO:tensorflow:Step 790, mean_score: 7.890625\n","INFO:tensorflow:Step 795, mean_score: 7.890625\n","INFO:tensorflow:Step 800, mean_score: 7.890625\n","INFO:tensorflow:Step 805, mean_score: 7.890625\n","INFO:tensorflow:Step 810, mean_score: 7.890625\n","INFO:tensorflow:Step 815, mean_score: 8.875000\n","INFO:tensorflow:Step 820, mean_score: 8.890625\n","INFO:tensorflow:Step 825, mean_score: 8.890625\n","INFO:tensorflow:Step 830, mean_score: 8.890625\n","INFO:tensorflow:Step 835, mean_score: 8.890625\n","INFO:tensorflow:Step 840, mean_score: 8.890625\n","INFO:tensorflow:Step 845, mean_score: 8.890625\n","INFO:tensorflow:Step 850, mean_score: 8.890625\n","INFO:tensorflow:Step 855, mean_score: 8.890625\n","INFO:tensorflow:Step 860, mean_score: 8.890625\n","INFO:tensorflow:Step 865, mean_score: 8.890625\n","INFO:tensorflow:Step 870, mean_score: 8.890625\n","INFO:tensorflow:Step 875, mean_score: 8.890625\n","INFO:tensorflow:Step 880, mean_score: 8.890625\n","INFO:tensorflow:Step 885, mean_score: 8.890625\n","INFO:tensorflow:Step 890, mean_score: 8.890625\n","INFO:tensorflow:Step 895, mean_score: 8.906250\n","INFO:tensorflow:Step 900, mean_score: 8.906250\n","INFO:tensorflow:Step 905, mean_score: 8.906250\n","INFO:tensorflow:Step 910, mean_score: 8.906250\n","INFO:tensorflow:Step 915, mean_score: 8.921875\n","INFO:tensorflow:Step 920, mean_score: 9.890625\n","INFO:tensorflow:Step 925, mean_score: 9.890625\n","INFO:tensorflow:Step 930, mean_score: 9.890625\n","INFO:tensorflow:Step 935, mean_score: 9.890625\n","INFO:tensorflow:Step 940, mean_score: 9.890625\n","INFO:tensorflow:Step 945, mean_score: 9.890625\n","INFO:tensorflow:Step 950, mean_score: 9.890625\n","INFO:tensorflow:Step 955, mean_score: 9.890625\n","INFO:tensorflow:Step 960, mean_score: 9.890625\n","INFO:tensorflow:Step 965, mean_score: 9.890625\n","INFO:tensorflow:Step 970, mean_score: 9.890625\n","INFO:tensorflow:Step 975, mean_score: 9.890625\n","INFO:tensorflow:Step 980, mean_score: 9.890625\n","INFO:tensorflow:Step 985, mean_score: 9.890625\n","INFO:tensorflow:Step 990, mean_score: 9.890625\n","INFO:tensorflow:Step 995, mean_score: 10.875000\n","INFO:tensorflow:Step 1000, mean_score: 10.890625\n","INFO:tensorflow:Step 1005, mean_score: 10.890625\n","INFO:tensorflow:Step 1010, mean_score: 10.890625\n","INFO:tensorflow:Step 1015, mean_score: 10.890625\n","INFO:tensorflow:Step 1020, mean_score: 10.890625\n","INFO:tensorflow:Step 1025, mean_score: 10.890625\n","INFO:tensorflow:Step 1030, mean_score: 10.890625\n","INFO:tensorflow:Step 1035, mean_score: 10.890625\n","INFO:tensorflow:Step 1040, mean_score: 10.890625\n","INFO:tensorflow:Step 1045, mean_score: 10.890625\n","INFO:tensorflow:Step 1050, mean_score: 10.890625\n","INFO:tensorflow:Step 1055, mean_score: 10.890625\n","INFO:tensorflow:Step 1060, mean_score: 10.890625\n","INFO:tensorflow:Step 1065, mean_score: 10.890625\n","INFO:tensorflow:Step 1070, mean_score: 10.890625\n","INFO:tensorflow:Step 1075, mean_score: 10.906250\n","INFO:tensorflow:Step 1080, mean_score: 10.906250\n","INFO:tensorflow:Step 1085, mean_score: 10.906250\n","INFO:tensorflow:Step 1090, mean_score: 10.906250\n","INFO:tensorflow:Step 1095, mean_score: 10.921875\n","INFO:tensorflow:Step 1100, mean_score: 11.890625\n","INFO:tensorflow:Step 1105, mean_score: 11.890625\n","INFO:tensorflow:Step 1110, mean_score: 11.890625\n","INFO:tensorflow:Step 1115, mean_score: 11.890625\n","INFO:tensorflow:Step 1120, mean_score: 11.890625\n","INFO:tensorflow:Step 1125, mean_score: 11.890625\n","INFO:tensorflow:Step 1130, mean_score: 11.890625\n","INFO:tensorflow:Step 1135, mean_score: 11.890625\n","INFO:tensorflow:Step 1140, mean_score: 11.890625\n","INFO:tensorflow:Step 1145, mean_score: 11.890625\n","INFO:tensorflow:Step 1150, mean_score: 11.890625\n","INFO:tensorflow:Step 1155, mean_score: 11.890625\n","INFO:tensorflow:Step 1160, mean_score: 11.890625\n","INFO:tensorflow:Step 1165, mean_score: 11.890625\n","INFO:tensorflow:Step 1170, mean_score: 11.890625\n","INFO:tensorflow:Step 1175, mean_score: 12.875000\n","INFO:tensorflow:Step 1180, mean_score: 12.890625\n","INFO:tensorflow:Step 1185, mean_score: 12.890625\n","INFO:tensorflow:Step 1190, mean_score: 12.890625\n","INFO:tensorflow:Step 1195, mean_score: 12.890625\n","INFO:tensorflow:Step 1200, mean_score: 12.890625\n","INFO:tensorflow:Step 1205, mean_score: 12.890625\n","INFO:tensorflow:Step 1210, mean_score: 12.890625\n","INFO:tensorflow:Step 1215, mean_score: 12.890625\n","INFO:tensorflow:Step 1220, mean_score: 12.890625\n","INFO:tensorflow:Step 1225, mean_score: 12.890625\n","INFO:tensorflow:Step 1230, mean_score: 12.890625\n","INFO:tensorflow:Step 1235, mean_score: 12.890625\n","INFO:tensorflow:Step 1240, mean_score: 12.890625\n","INFO:tensorflow:Step 1245, mean_score: 12.890625\n","INFO:tensorflow:Step 1250, mean_score: 12.890625\n","INFO:tensorflow:Step 1255, mean_score: 12.906250\n","INFO:tensorflow:Step 1260, mean_score: 12.906250\n","INFO:tensorflow:Step 1265, mean_score: 12.906250\n","INFO:tensorflow:Step 1270, mean_score: 12.906250\n","INFO:tensorflow:Step 1275, mean_score: 12.921875\n","INFO:tensorflow:Step 1280, mean_score: 13.890625\n","INFO:tensorflow:Step 1285, mean_score: 13.890625\n","INFO:tensorflow:Step 1290, mean_score: 13.890625\n","INFO:tensorflow:Step 1295, mean_score: 13.890625\n","INFO:tensorflow:Step 1300, mean_score: 13.890625\n","INFO:tensorflow:Step 1305, mean_score: 13.890625\n","INFO:tensorflow:Step 1310, mean_score: 13.890625\n","INFO:tensorflow:Step 1315, mean_score: 13.890625\n","INFO:tensorflow:Step 1320, mean_score: 13.890625\n","INFO:tensorflow:Step 1325, mean_score: 13.890625\n","INFO:tensorflow:Step 1330, mean_score: 13.890625\n","INFO:tensorflow:Step 1335, mean_score: 13.890625\n","INFO:tensorflow:Step 1340, mean_score: 13.890625\n","INFO:tensorflow:Step 1345, mean_score: 13.890625\n","INFO:tensorflow:Step 1350, mean_score: 13.890625\n","INFO:tensorflow:Step 1355, mean_score: 14.875000\n","INFO:tensorflow:Step 1360, mean_score: 14.890625\n","INFO:tensorflow:Step 1365, mean_score: 14.890625\n","INFO:tensorflow:Step 1370, mean_score: 14.890625\n","INFO:tensorflow:Step 1375, mean_score: 14.890625\n","INFO:tensorflow:Step 1380, mean_score: 14.890625\n","INFO:tensorflow:Step 1385, mean_score: 14.890625\n","INFO:tensorflow:Step 1390, mean_score: 14.890625\n","INFO:tensorflow:Step 1395, mean_score: 14.890625\n","INFO:tensorflow:Step 1400, mean_score: 14.890625\n","INFO:tensorflow:Step 1405, mean_score: 14.890625\n","INFO:tensorflow:Step 1410, mean_score: 14.890625\n","INFO:tensorflow:Step 1415, mean_score: 14.890625\n","INFO:tensorflow:Step 1420, mean_score: 14.890625\n","INFO:tensorflow:Step 1425, mean_score: 14.890625\n","INFO:tensorflow:Step 1430, mean_score: 14.890625\n","INFO:tensorflow:Step 1435, mean_score: 14.906250\n","INFO:tensorflow:Step 1440, mean_score: 14.906250\n","INFO:tensorflow:Step 1445, mean_score: 14.906250\n","INFO:tensorflow:Step 1450, mean_score: 14.906250\n","INFO:tensorflow:Step 1455, mean_score: 14.921875\n","INFO:tensorflow:Step 1460, mean_score: 15.890625\n","INFO:tensorflow:Step 1465, mean_score: 15.890625\n","INFO:tensorflow:Step 1470, mean_score: 15.890625\n","INFO:tensorflow:Step 1475, mean_score: 15.890625\n","INFO:tensorflow:Step 1480, mean_score: 15.890625\n","INFO:tensorflow:Step 1485, mean_score: 15.890625\n","INFO:tensorflow:Step 1490, mean_score: 15.890625\n","INFO:tensorflow:Step 1495, mean_score: 15.890625\n","INFO:tensorflow:Step 1500, mean_score: 15.890625\n","INFO:tensorflow:Step 1505, mean_score: 15.890625\n","INFO:tensorflow:Step 1510, mean_score: 15.890625\n","INFO:tensorflow:Step 1515, mean_score: 15.890625\n","INFO:tensorflow:Step 1520, mean_score: 15.890625\n","INFO:tensorflow:Step 1525, mean_score: 15.890625\n","INFO:tensorflow:Step 1530, mean_score: 15.890625\n","INFO:tensorflow:Step 1535, mean_score: 16.875000\n","INFO:tensorflow:Step 1540, mean_score: 16.890625\n","INFO:tensorflow:Step 1545, mean_score: 16.890625\n","INFO:tensorflow:Step 1550, mean_score: 16.890625\n","INFO:tensorflow:Step 1555, mean_score: 16.890625\n","INFO:tensorflow:Step 1560, mean_score: 16.890625\n","INFO:tensorflow:Step 1565, mean_score: 16.890625\n","INFO:tensorflow:Step 1570, mean_score: 16.890625\n","INFO:tensorflow:Step 1575, mean_score: 16.890625\n","INFO:tensorflow:Step 1580, mean_score: 16.890625\n","INFO:tensorflow:Step 1585, mean_score: 16.890625\n","INFO:tensorflow:Step 1590, mean_score: 16.890625\n","INFO:tensorflow:Step 1595, mean_score: 16.890625\n","INFO:tensorflow:Step 1600, mean_score: 16.890625\n","INFO:tensorflow:Step 1605, mean_score: 16.890625\n","INFO:tensorflow:Step 1610, mean_score: 16.890625\n","INFO:tensorflow:Step 1615, mean_score: 16.906250\n","INFO:tensorflow:Step 1620, mean_score: 16.906250\n","INFO:tensorflow:Step 1625, mean_score: 16.906250\n","INFO:tensorflow:Step 1630, mean_score: 16.906250\n","INFO:tensorflow:Step 1635, mean_score: 16.921875\n","INFO:tensorflow:Step 1640, mean_score: 17.890625\n","INFO:tensorflow:Step 1645, mean_score: 17.890625\n","INFO:tensorflow:Step 1650, mean_score: 17.890625\n","INFO:tensorflow:Step 1655, mean_score: 17.890625\n","INFO:tensorflow:Step 1660, mean_score: 17.890625\n","INFO:tensorflow:Step 1665, mean_score: 17.890625\n","INFO:tensorflow:Step 1670, mean_score: 17.890625\n","INFO:tensorflow:Step 1675, mean_score: 17.890625\n","INFO:tensorflow:Step 1680, mean_score: 17.890625\n","INFO:tensorflow:Step 1685, mean_score: 17.890625\n","INFO:tensorflow:Step 1690, mean_score: 17.890625\n","INFO:tensorflow:Step 1695, mean_score: 17.890625\n","INFO:tensorflow:Step 1700, mean_score: 17.890625\n","INFO:tensorflow:Step 1705, mean_score: 17.890625\n","INFO:tensorflow:Step 1710, mean_score: 17.890625\n","INFO:tensorflow:Step 1715, mean_score: 18.875000\n","INFO:tensorflow:Step 1720, mean_score: 18.890625\n","INFO:tensorflow:Step 1725, mean_score: 18.890625\n","INFO:tensorflow:Step 1730, mean_score: 18.890625\n","INFO:tensorflow:Step 1735, mean_score: 18.890625\n","INFO:tensorflow:Step 1740, mean_score: 18.890625\n","INFO:tensorflow:Step 1745, mean_score: 18.890625\n","INFO:tensorflow:Step 1750, mean_score: 18.890625\n","INFO:tensorflow:Step 1755, mean_score: 18.890625\n","INFO:tensorflow:Step 1760, mean_score: 18.890625\n","INFO:tensorflow:Step 1765, mean_score: 18.890625\n","INFO:tensorflow:Step 1770, mean_score: 18.890625\n","INFO:tensorflow:Step 1775, mean_score: 18.890625\n","INFO:tensorflow:Step 1780, mean_score: 18.890625\n","INFO:tensorflow:Step 1785, mean_score: 18.890625\n","INFO:tensorflow:Step 1790, mean_score: 18.890625\n","INFO:tensorflow:Step 1795, mean_score: 18.906250\n","INFO:tensorflow:Step 1800, mean_score: 18.906250\n","INFO:tensorflow:Step 1805, mean_score: 18.906250\n","INFO:tensorflow:Step 1810, mean_score: 18.906250\n","INFO:tensorflow:Step 1815, mean_score: 18.921875\n","INFO:tensorflow:Step 1820, mean_score: 19.890625\n","INFO:tensorflow:Step 1825, mean_score: 19.890625\n","INFO:tensorflow:Step 1830, mean_score: 19.890625\n","INFO:tensorflow:Step 1835, mean_score: 19.890625\n","INFO:tensorflow:Step 1840, mean_score: 19.890625\n","INFO:tensorflow:Step 1845, mean_score: 19.890625\n","INFO:tensorflow:Step 1850, mean_score: 19.890625\n","INFO:tensorflow:Step 1855, mean_score: 19.890625\n","INFO:tensorflow:Step 1860, mean_score: 19.890625\n","INFO:tensorflow:Step 1865, mean_score: 19.890625\n","INFO:tensorflow:Step 1870, mean_score: 19.890625\n","INFO:tensorflow:Step 1875, mean_score: 19.890625\n","INFO:tensorflow:Step 1880, mean_score: 19.890625\n","INFO:tensorflow:Step 1885, mean_score: 19.890625\n","INFO:tensorflow:Step 1890, mean_score: 19.890625\n","INFO:tensorflow:Step 1895, mean_score: 19.906250\n","INFO:tensorflow:Step 1900, mean_score: 19.921875\n","INFO:tensorflow:Step 1905, mean_score: 19.921875\n","INFO:tensorflow:Step 1910, mean_score: 19.921875\n","INFO:tensorflow:Step 1915, mean_score: 19.921875\n","INFO:tensorflow:Step 1920, mean_score: 19.921875\n","INFO:tensorflow:Step 1925, mean_score: 19.921875\n","INFO:tensorflow:Step 1930, mean_score: 19.921875\n","INFO:tensorflow:Step 1935, mean_score: 19.921875\n","INFO:tensorflow:Step 1940, mean_score: 19.921875\n","INFO:tensorflow:Step 1945, mean_score: 19.921875\n","INFO:tensorflow:Step 1950, mean_score: 19.921875\n","INFO:tensorflow:Step 1955, mean_score: 19.921875\n","INFO:tensorflow:Step 1960, mean_score: 19.921875\n","INFO:tensorflow:Step 1965, mean_score: 19.921875\n","INFO:tensorflow:Step 1970, mean_score: 19.921875\n","INFO:tensorflow:Step 1975, mean_score: 19.921875\n","INFO:tensorflow:Step 1980, mean_score: 19.921875\n","INFO:tensorflow:Step 1985, mean_score: 19.921875\n","INFO:tensorflow:Step 1990, mean_score: 19.921875\n","INFO:tensorflow:Step 1995, mean_score: 19.937500\n","INFO:tensorflow:Step 2000, mean_score: 19.937500\n","INFO:tensorflow:Step 2005, mean_score: 19.937500\n","INFO:tensorflow:Step 2010, mean_score: 19.937500\n","INFO:tensorflow:Step 2015, mean_score: 19.937500\n","INFO:tensorflow:Step 2020, mean_score: 19.937500\n","INFO:tensorflow:Step 2025, mean_score: 19.937500\n","INFO:tensorflow:Step 2030, mean_score: 19.937500\n","INFO:tensorflow:Step 2035, mean_score: 19.937500\n","INFO:tensorflow:Step 2040, mean_score: 19.937500\n","INFO:tensorflow:Step 2045, mean_score: 19.937500\n","INFO:tensorflow:Step 2050, mean_score: 19.937500\n","INFO:tensorflow:Step 2055, mean_score: 19.937500\n","INFO:tensorflow:Step 2060, mean_score: 19.937500\n","INFO:tensorflow:Step 2065, mean_score: 19.937500\n","INFO:tensorflow:Step 2070, mean_score: 19.937500\n"],"name":"stdout"}]},{"metadata":{"id":"WKWPdwP8BW_v","colab_type":"text"},"cell_type":"markdown","source":["The above command will run a single evaluation setting to get the results fast. We usually run a grid of different settings (sampling temperatures and whether to do initial no-ops). To do that, remove `eval_max_num_noops=8,eval_sampling_temps=[0.5]` from the command. You can override the evaluation settings:\n","\n","```\n","  --loop_hparams=game=pong,eval_max_num_noops=0,eval_sampling_temps=[0.0]\n"," ```\n"," \n"," The evaluator generates videos from the environment:"]},{"metadata":{"id":"At9LC5rxFyv2","colab_type":"code","outputId":"983b0e7a-2700-4e4a-d776-03c459669770","executionInfo":{"status":"ok","timestamp":1553253830168,"user_tz":-60,"elapsed":4036,"user":{"displayName":"Piotr Kozakowski","photoUrl":"","userId":"01014928596539690143"}},"colab":{"resources":{"http://localhost:8080/nbextensions/vid.mp4":{"data":"AAAAIGZ0eXBpc29tAAACAGlzb21pc28yYXZjMW1wNDEAAAAIZnJlZQACqh1tZGF0AAACrgYF//+q3EXpvebZSLeWLNgg2SPu73gyNjQgLSBjb3JlIDE1MiByMjg1NCBlOWE1OTAzIC0gSC4yNjQvTVBFRy00IEFWQyBjb2RlYyAtIENvcHlsZWZ0IDIwMDMtMjAxNyAtIGh0dHA6Ly93d3cudmlkZW9sYW4ub3JnL3gyNjQuaHRtbCAtIG9wdGlvbnM6IGNhYmFjPTEgcmVmPTMgZGVibG9jaz0xOjA6MCBhbmFseXNlPTB4MzoweDExMyBtZT1oZXggc3VibWU9NyBwc3k9MSBwc3lfcmQ9MS4wMDowLjAwIG1peGVkX3JlZj0xIG1lX3JhbmdlPTE2IGNocm9tYV9tZT0xIHRyZWxsaXM9MSA4eDhkY3Q9MSBjcW09MCBkZWFkem9uZT0yMSwxMSBmYXN0X3Bza2lwPTEgY2hyb21hX3FwX29mZnNldD0tMiB0aHJlYWRzPTMgbG9va2FoZWFkX3RocmVhZHM9MSBzbGljZWRfdGhyZWFkcz0wIG5yPTAgZGVjaW1hdGU9MSBpbnRlcmxhY2VkPTAgYmx1cmF5X2NvbXBhdD0wIGNvbnN0cmFpbmVkX2ludHJhPTAgYmZyYW1lcz0zIGJfcHlyYW1pZD0yIGJfYWRhcHQ9MSBiX2JpYXM9MCBkaXJlY3Q9MSB3ZWlnaHRiPTEgb3Blbl9nb3A9MCB3ZWlnaHRwPTIga2V5aW50PTI1MCBrZXlpbnRfbWluPTEwIHNjZW5lY3V0PTQwIGludHJhX3JlZnJlc2g9MCByY19sb29rYWhlYWQ9NDAgcmM9Y3JmIG1idHJlZT0xIGNyZj0yMy4wIHFjb21wPTAuNjAgcXBtaW49MCBxcG1heD02OSBxcHN0ZXA9NCBpcF9yYXRpbz0xLjQwIGFxPTE6MS4wMACAAAADkWWIhABvrNdXNvEPmO7lwVl73sPl0EDBzzvrz1O9Sgfa49FGnVhGNj4PrUzIEjAsiR14q5boH034au6fMfeHzW8BQIdLu5D8GWFcvhnUQvMLIDm/5fDlJWNI1pLZ0KekyKgRZvEg10IZZePvLcj64kGJzCMbJi6QZbX4WMzyM/ZwsXoWWPBmmlKBzFixHWdkptjcAYhpgDpXSILlIffpBFr5Fmv8Xdrl5eZtB/U18q6RE0tX2BrhekKyOZ5lJWnXZWIEICkLYIda8x0l/aAug9zkJAN2UJ5v8AfQgXgS7iPy41I11UQneH59QQ6r2Fy+bXVz7hKXvFUUQUW2NfwyAHSubAtKRV8FgrIBnKXwxAjc8zc/00LsdZVdehIaL1eI9qZtyap5GmVpF7ZJdkQbo7j2k9/o8Ztr6lwZrODqoujHSJK6V9bK0u9Et564zU+wWgftergJVAEl4m/D3N6/lD6Tni/a6bLzIcdcVjnfWLPAUBwAoj19NpxhAbe1VyiybbzF11k65OpExrnTpeyfXnWi2YKXmv6NMcvP6YS8WOK4pM7nWhyKetjJvO69p10oeh7Pv3PuQBq3kARIBKQ+MPYmymnbhgmxG/6w3hJ2A2Urz2k1DctVq7TiUCWnAReHSDqSpYcdQwxCm/lIpIwtl/dffgss5v+hhFs6NSNe3zqLc+wa/P6fKKBzHBPA6mZtXbiaJH0Y+5hMHtf92lFc+6I4pZ1q2XpI5Nr1V7em9lfehnp6KwZCFUTCrCle3ZgVn3/WlL0hiX3HqF/qGx1rSRBE7lqG2nGEQXx7BFJGNLF0vFVi0j2agV+lVqGOVlIxAjK3E9wGWVM0V7xAFGXQtxAYJ6qA6zMuM1AzlTqoEWcy+zkYm6Z2/Vn8RMHtpHaCW7GF05Wujcn0D05dR11MQem50GDiKxlzGighyGKWmfeex/qNBXelV1apol1nwDCUiSbC9fPUu70YI94kit+OnCdHe588u+9o5tqmvG+4ju2D1U0YtGzBJLwNtKIxTj+ycim3c7lWMz9gyNpdcRw85nQOO+UebN2j6KDuTy4XNxidtzFIcvo67EYGfl+q3WaPfQzFQLuOqvybvDRViyMxNwUidf7UCNcjzUMa0RFtd4HPTD4pR9pL0oOHG0XMOwvMlfvI/0tUl0nH8gKjxa10D+0pCAG9Sq76K3xNRb2QQ4PhDp7u3P+U7CY7JpR9qHasfUEAAAAxQZojbEb/+8M7YXuyASh7Kplen5y7UyO14JIrbok4XTbRQe3ORruR41lvOoDou8ClgAAAAB1BnkF4m/+AoGMvgzynj7iP1aLLDUqgN6Jo/suYQQAAACABnmJqRH+I4ID6PupU5REY3sf/aa003+ohQ9ie2tSAgAAAAbhBmmZJqEFomUwI3/3q7CUD4uvUkflAlpSidF3UDZUIJZsuBftR/Ot+q3GhwTn9egJDu/Q0u302gbAsB/IV/AFE9llDfK6lUW9694v3+SwMl6mllWP/WwEuWwr3bvYMZvmSan6TfqkNjfFQ2gIUyqpE9/WshNO04YsB0gSlJqYcVRMCqNuW9LOcW33er/NiW98OL0en+UGHVMigjfMRwtvQwEBx2TCEvqNHiHUZK47Ql9EBEwOC+9RX+RfTR+dz/gd9jggR1VC/W7S6VbwMD6OLZQ6pAbUGRTytAu/D3nUIHJpw2KDyU/GpGdKbj2WQLEoSWl/G1erWEOohpWvfDxHkGcHJyES1MZ1DvjVLlpUz1LfdDNETW9u5oytpQuaaon1mmF0VFggCZuK84NA/jjnYLjbiGa2mQVNuCz0xJcY1SlKU0HHZWh4xyJUcesClYPnBpMFzn1F68BZn1odfSkUsfnPO8Ozfwiiz7f3xZq0mYnuxahAz2t7bEpiZEBhl9vOvB5IwNQnQ2WMn3D1n/XqqbGvjTwCrUHeBDng8+lhZDYcGCdhZ+2y21qFiMlJaX159s9jZ9Nr1pQAAAFpBnoRFESzfiNYKja/0xDhfTAiTpHLhMdYxhLBGYN2yGAaE7v0DPQ2YrW+QLfMv2VvTlemA0fw0ITn/7H4QKf0Smi2r8j9M+wtBuP4ooX+ZG3O079gxDawGHzkAAAAeAZ6lakR/gQeVglqse6icagSNwuz60J4uGyZCePmbAAAASUGaqkmoQWyZTAjf6DCt6Yuhred9EHrqUOfXsyntdBaWSJrN7aPiRCVp6QCYQRjgAPaX306Ulc6wvo26P1UQNC4biU/od187SikAAAAxQZ7IRRUsn36T35XM0R9+0te5g6nnf972YKyQEsNo7mvZhGpibk4k6GYi5Xw1oII/gAAAACwBnud0RH+KcQxiHI/+tYvdUXZoto0LM0yuhiZoPS8TyCql7aDErRX1MckO7AAAACABnulqRH+Gdbgx/3UyLHgLBkL+GlWi7oO1l81SbcV5BQAAAC1BmutJqEFsmUwI3/QsGvLB67IEGcjaf9jQmmp2xwadCZqIc+QsmxjXnMZQAp4AAACeQZsNSeEKUmUwUVLG//1hfa7H0q8GiMhKKARrmw6Z4xdeyr7ujJme+WFMGuspP/+ISBWO3g4rVjHKMZczL9N457VGygcsdj7HL9XSQg7I9GyDcuezylwimiV/wFo/tnsP85jKo1N9u1U/c32ZRoJ6MLzOYkUW9bgpdQwc/vwurbPO6DqdSp7R4DyF3mA4Y/IOgh69Sz6dfNBC5f60MoAAAAAUAZ8sakR/iOA7CrYE++z83mxaT9kAAACCQZsxSeEOiZTAjf/7wc1PkbwRgr+fjnyNR8pfsXICQZIWSD29LO9OdoFeaCGC7kNQcI6QVUvVoNv//1awF0I8N0tfCIdJ96/qdtaHfFDnSTMhoQ9GjaS9EN6FwfTiM0n2rikHxnG2wBAjL4oKVFB1J9q4Q6oi+C8FkQ6nZ9weS8YziQAAADVBn09FFTyffW4h7bT+Vmnsg7Jwp4y5XEPjbVoz423RdhZiyhHoVMXw4WOwpIvW8VtedZMRgQAAACsBn250RH+HsPlaZzJ171H1c9c+P/yRU8UOm2VwDFLU0JopWl1F2Khtda0wAAAAHgGfcGpEf4TUK0IOqyu7XwKdBz1mDC7mcZOyimAiwAAAAGlBm3JJqEFomUwI3/wDwzHAex2BjGQbp8PmN+QZU96uDVeJsc29OEXjcldhiXFvZCjL1XZadbCRD86yiAgWQRsivW2zgAoC7DVCL9e3fSnRGKpVU/LGLClh3Y48sH47PWqbcpIEZ6IQMXUAAAByQZuUSeEKUmUwURLG//vD9/y1OQEdZZUsKUwNqfrKLfwKdmKP4MYlUfWBfMS6AB2YzPOOewKgSMfLoK01LvNA1j94yZ+2K1utCc4bPCb5fgtJtphyhALG+EJzcxQG8TmO93UCKf8P7LJwQsMlijgw4NSUAAAAMQGfs2pEf4etAFroO/Tre2AdB5O/z7sBdNC+yMviqgLALp+SqJXA7Pg/DPRDeTGlt4AAAABvQZu1SeEOiZTAjf/9dtYVwOWMQEFcGvjV4Syzg11kPr1s9125wYmt9rvuf3mx7RYJ6jF6ErfUzdz0/3v9ae/O7utr4LCBaSp+/oto7O3F1qWD/jeXpJo3T76eNJFHYJw8VuPijBUU3Vp8h5phb4yxAAAAO0Gb1knhDyZTAjf//XzJzOIH3VrgDqkUNgn86gOXkEPCMlG+9fgG6imSEzHDCVblECfv6lWOPhYa83SgAAAAbEGb90nhDyZTAjf/++FR2EKSC1to8pc6YSO3kyM61vFVTIqwBRLCgvhTW2uELbv4P1t6j3j6R5JEnH93R+6ccuNOfX+S8DB7HTWlWJcQLgqD+S2bSR1NJYKOXWfT3yjStlWid2o9fNbVb/WnQQAAAIZBmhlJ4Q8mUwURPG/74DjgdzkBXMYP2UA+265sg2RvtJWp7hbeXljZyPg5U1jpyoQAMkswdLvArqw2ReEXhiCOCghoFR9+KlVVnzlHwdrRAGNlgaJUCgaXSP04WBiefPhnJH+UNq4LFoPepHhK//kKOrVcBFaR4anoX/aA5EFLlBZC3z8HwQAAAC4BnjhqRH+HsRNxvly6e2k2NMUXC/vU7UPeiHPyDV2p/YvXZWhgk+q5Cpb6bHPQAAAASkGaOknhDyZTAjf/++A44Ilb6dcZ4EfRuTLTuDvo3J2im59tzroCzqfsO8XO0To/kC1I69bSu6OBA9I+wRfSH4s2YCwMyLJB4SXTAAAAXkGaXEnhDyZTBRE8b/wDoisgQ2L7kEBMp1YsiwFGNSG/OTV6DNluS1vSFNmiEavm3sOuILR8pgK4M6P/qHxZC276LRYUYYCavgfWkcLXP29c5JUTuaP0/h/7ytcJzS8AAAAwAZ57akR/h7Ee1u4Mf7i3wEEg6o2q6d802TgnxIfuhG+eE+Dw81v3A/7HzNqgsMZrAAAAZkGaf0nhDyZTAjf/++s+OA7zprt9FalqM30xdIjlcH2DZkkFtbEKP2KnwPRDHCudTxGL9Qj/p+X8zQUT6DsKw9cio1ehc/Va1ccw4608TJVh/tU58mts0w5trYhfnxrz7ghBbWrPgQAAAB9Bnp1FETzffU2jm3iBOi008YuPeRt0exkohagoWh6AAAAANwGevmpEf4k5No92pOx4d3tVe/jt/92QHqWDJ1gLa+IuwyWHPUwqtoP8fV9Ym+scICO5EURRBtAAAABkQZqhSahBaJlMFPG/++FR2CGqQlUqIWgPGeTAl0jkR8nWtez7chaAXKIsKU2quN90XktERKg37/m9ip4Dz5kOgzQscz6nyNY0f/sd6sj1HySlAcKWdh7zp6SDANGHziaLu7KO4QAAADABnsBqRH+JNh+VxVwUJ+Y3lV+rya5I/8TCxwR0ZLGKwtu7rocrr7JN3QZSGNke1uAAAABQQZrDSeEKUmUwUsb/++A44DvMfbctqjDdsKqNf0xOA39UYyLAzYIZyyzAIzc51s3jIS5ItuDvEtTXgkuVVUxzdDw7I8ZpdZavGnbqYY9+GhMAAAA1AZ7iakR/gQ9rhvIRB0ySHXw58101sjT3zRJJs0q1N75B5kzDs336j3kR4+OINWFbqx8fObwAAABPQZrkSeEOiZTAjf/74VHYIlYiKMdyQZmvfSW5HcZkgM1RFiY3gEH71z/3fVBzOwcraE90El70+ZDjx4HGc66IARVSqHcYq0qBrrRk30FtqQAAAFNBmwVJ4Q8mUwI3//vD7l3aeLA9vGddw8aOZ4u/aB1tWqMJJSgQk2EKksQOal+1aR5lX56hJg2jFlZ0Xy5wPHFC5sT+XuWXvNeiHLLg872OhCGzwQAAAGBBmyZJ4Q8mUwI3//vDobVTkBINCYpxqf92jigZ0Srp218/WFaj2TRlEHOqssGMM0Apc19xhhprfp/o+nyoo3nU4qIPx0toevbOZ/uMdat8h/rz3l6mAf6cx4JPPMV/58EAAABZQZtHSeEPJlMCN//7xIHZpGQEF4WlOSmv5sDvQ1HkPJ24b8mLANR7I1VevWJebie2fuXBDZeaCivuJcWRnz+sLxLjPYolWMX+OCeAR284M5BlQ+hDK45MhLkAAABfQZtpSeEPJlMFETxv++A44EJXr11WQvEfDTQVOZV8/+Xr3S8f8Sdkd5yhFoZ0Eykj3Xivg1BzZSewsPXa9nQ6gqdDgf7T2DmCjKrk7f2e2vK9ypwXpi7kgz9A0TgFjRgAAAA3AZ+IakR/h7ESVb5jShGX8VmBJn/MdoNQ3Z6fh4j/jmgTG8KfR4gLxPvHCfsFkI5Dojx4hvmkmAAAAEFBm4pJ4Q8mUwI3//vgOOAkGyUenL98r33EsuAShImg94uLBrTk6dhiDVHpVEky+DhSk5kCbedWMCnU+cpV31ziJwAAAHdBm6tJ4Q8mUwI3//vBcoj4CbBCQT2RExz0aNVXfJZgEMRjG13pvbjo3lZ6TLZP87/4osMgfYBX228XW1pfmU8k2l4pXWkF2ONrCPtLUsjWqi9NqYjz9uxWY1lh/hx5fARLNZPbPZztgBCOOKpjSsjqrRTvm2kDmwAAAGlBm8xJ4Q8mUwI3//vgPUQfdNRvnx3FP1pywn3LJO+BgkxSFcCfgS8xAc5jl7DjT9LzYEwD9isIU1yUsdzHqy/6ngxDd8Uf0UrLVqXY4UPbCNliO8n/ZegcRbOArrd0G+CXl9/2CFy7Fp0AAABdQZvtSeEPJlMCN//74D1EIHvgowka8CSMlOV1djkcwdrmN+M3w0ez2EF3Z8QgY9h2qKskrDR5Mk2+uEQHAHNJFzVoZ8OEnB6N5bVjkoQH9GCmMuS0GmqX+iknVHfRAAAAYEGaDknhDyZTAjf/++A44EPraJILnY0LuTUTVqG8Y4R+iod8qU50V48CYGzxUBWrE5/inR2Nxzjn9BK+lYDY6wizKJTKtNBOivIiofjj7puT2VHhc5kdvPo7rDnnm0xNBwAAAGpBmi9J4Q8mUwI3//vBnCJ2LA71GF25ddxrRcaDbZ6JrFKdXAbBcwenVA97XR386r26bMwGidvp2PmbPeLJeVTrLs40YO/0dkENxamY3tNtYwvFfYvB8Y8OeKM6xrpDUHsU4NOoFoGqtzCdAAAAYEGaUEnhDyZTAjf/+8N7/LIyAjCC10eRAkC4payNwDB4TaAJfH4zC+l57EAHdxnDWH/9FYP1pFGTKBZX4dIkEoguP5WochPIdwRBalFTzEmR/DYBfxJGM8ekovk7ApnUkAAAAHVBmnNJ4Q8mUwI3//wDwzHAd5j8RhmR9H0z98nDjYj/Nc7vsZGPARdvUDB2TCWaeccqi54nDBz6EstQtzy6v+wvloWKraz+uz/M18zKVlwhNHeSR4GZPnFXhk/6SFiPs3DdL9fCsVu/U6phJsgQpW9nziiIrsAAAAAlQZ6RRRE8331No5t4gTzZDbeKeBR95P1RisqLON/iq6lMAjgQ8QAAAEEBnrJqRH+HsRJVgDz377tnoNimjhTZvPY9yHf2fmv0wsOMF+wp88Zp9A47UCY9LcsoRhIaJ8uNlzZCrsZsy0xHSwAAAEdBmrRJqEFomUwI3/vgHICtV4/N+3TLBHjt+Fcn7DWjz8soUTNhnfkqkcMHOJu59OaFxhuVf+ti1/1ybHdnNcNjs9ZKiS4KwAAAAGJBmtVJ4QpSZTAjf/vB/wweEBHWjwtlVbql9dgODZ2kZPlG+D4KXaVCP4hDGqGFAR8WEBDCUpQahknTJjX5ZfkvoG4DfOZ2jNfwPVMAdpL7FdEsgo7WPCSMB38k1yCOW8g6+QAAAE1BmvZJ4Q6JlMCN//wRIX9RB98K5rsT6v40vwfGBFHd36t3Lmg6acxZXjPDoIP1nuZtPbJR67n3MUNs/jpJGr2SW9ZnVls5Jb7a9l/19AAAAFdBmxdJ4Q8mUwI3//vgR+AQ8tE14oytZoCQH1vmaiP6zviuiN8WiV39XBnBmeNGUBPj1lZw3uGVd1ASBgSvQit8rSzdGRcZWThyfaYyQbzG3uhl9FJ+f80AAABOQZs4SeEPJlMCN//74DjgOFPapEipwNQ5JLf/iFkT7vTC27zO0O1Hhr7Be9x5uGF5a9SD1TJTW/IfLAT7Nh7DJMwi4QswAkdulgB0QpQxAAAAd0GbWUnhDyZTAjf/+8Gbu08WB7ipWtNjMHVF/iPxLBqJR3dWf3VXwht118WdJ2yr8sybv5tw7Gkcen/6wXcv6hm1FN3LQa1BDmxg9YZtMWlhqjvMsaNo3uWx/3Zsn/9CB07I7ahNSdEvwIcfvP9pOIVFVaVu8QFYAAAAWkGbeknhDyZTAjf/+8Ohm1TlEJHbNU/mWRKFgYzLN2yQ8jXAB9wCzOsfbVP2lU1EY8QOgEzKqXo0R7iSjvfThP3Ejb6DhTio3axVyuxNGvXKyq/GmvvY+WeEeQAAAINBm51J4Q8mUwI3//vD7xujgdzj8Rhi1Ykq7OuqSWTi3Jyq4sFrVH+uGrfgBuxTFBDV13U4LCn/MgN+f4ZzOHXoTdRaVuXBhLEbCtA5Ms+hD9d2LkEvLXaBv0E8NnWBPx9w4ZTvBVStx4l9Na6nPeTrSSDBnaVQkpIMXYq0+x5SehxmeAAAACdBn7tFETzffVS1iQH6ubsRADaTWu5ionu4K0/0ak9zTnZggGqi+fkAAABTAZ/cakR/inEMZAKk66wt/cwHaBMWpfAc500hIeLekQWHoYVeIuwzwo3AD4zjCnVWjbJVYnBRhCSWWZm8rpySFJ3514IPVb1tZpdA21QwKGClE4UAAAB9QZvfSahBaJlMFPG/+8IGeHNgeeQ43E7kwPs1VGXGJw973890r2GIuyVyjsOqfdfCvVaSkH16R0M1/zHTQJQoPuOg2+H7r68wE56bXlYqz1zmB//ygylNGdR6NOHJuRI4gUP87HXbU2cwVemECH+lErFeByiJwCmM07PdY7YAAAArAZ/+akR/iTYflbJv66XME0GHkCF6mtJpaggCYpsevkqPWehQsYwUCZ42gAAAAHxBm+NJ4QpSZTAjf/vBm7tPFgQmOx8vn1I4MEdYOyJbxW9+T8JyO3vjlruuDeJ94kHVp6mE4InbW3mhydReOAPRHSvfloMOy1X4f1k61/T5AyQdQwSfNs/ceWh+FXJYPI5vBAAfACpz0P00Y0ZQiTyXgYR5dHA4282fpEnhAAAATEGeAUU0TJ99biHttP5WbMcugQMX/6p30Ja38txT+7rJCjWU5BS5Hu/Ped4RgDO/8NmgX1/xtKU+guus6I+GB0iyn+7zzesJ2+Dd1FQAAAA7AZ4gdER/h7D5W775k5qly41vVWL+9TyyQd/+taVv6vysDLrv4GuSfz03ASJUej48yG8+gYk5sPihZ4kAAAAvAZ4iakR/hNQrQyxiUa04Hw7yezOPlcbNtcK39Adv5vlNi03I0KSWCx97lgf8KiAAAABOQZokSahBaJlMCN/7w/fqde8Ef0LFVa6z8obTV+Z3fHanJstZCK+C9VlS5uYoIGfyFRCQyvYs8eXMoinX6wkePCQbIrdvJZtx24a1z2OBAAAASEGaRUnhClJlMCN/+8HPtQiIJq4KJFcEZqp5eZHjbfXvRNExlEoNtkHV/O1j3XsPkfmAogRZ6H3xb4fGyDO/psnXuk/zjrKbgQAAAGFBmmZJ4Q6JlMCN//vEgWa8cB7g7HnWVr5afQMBhPt/eF0XZINzgf1DTqt4LgNJ9n7jJeXS6jzqED9EpR64Rz+FTHrYHSWny0wQxBJBhnSb/iEP/JdZ86h/Wi9v9QW09VhrAAAAQEGah0nhDyZTAjf/++s+OAznGqGizRSqo3BVcpAWaqCB8ovSoNU/Q+cLNp60lGoY1XxqNFsDbFsRe3APbjIzVIEAAAA7QZqoSeEPJlMCN//74DjgJBsj6l+X0srb4l6b6JH2Ft3abZgPDnb7Plzz2pLEPDmuaBtebuDUX6hFZoAAAABcQZrJSeEPJlMCN//7x1Dlo+AmwPQ/8NXi9GKqmuY1UPDYZG3NPKYYZuf5RLSPKtiYcXru35hiQqjNXqPD4oWVOHCsaOXNK65RabsS8NeFVBf/1FIziqztdYeHaIAAAABfQZrrSeEPJlMFETxv++A44DvMfnbEhBgUx9vnA4Ho8ktSkKylpuzQvYY0CtUp4uNU/8qrDXOOYbe12K7QloChfe0qoEDZAp+Ls75A+yk9tDX6iPXAaWpTzsfeqD+dbeEAAAA4AZ8KakR/h7ETcb5ctlmh7dxoqP/yRU8UOqHP2aaGG/KPrV+Afm67tENSRXPgCr2j+g5e2szXEdIAAAA2QZsMSeEPJlMCN//74ByArVvp1VBzo6FR5mbIURGBYtSZMz2ctINnHT3rH2VjPX/zzdapOx0wAAAAW0GbLUnhDyZTAjf/+8PuXdp4sCBqPO04eytS9D46/Avaazen9hU5pU11Hmv3PBlT0PTYG549QlWOOVwWlJAXcvc1syzyvf1qFX3EhgXM1IOIcGnTwUuYkLr42jEAAABaQZtOSeEPJlMCN//7w3J3dNWQEg0JK5pzBQVfpWoOKLSI1/VmJjcsCrBle+ncG6tiC7wuDEZ5USHcZC55aKT2OIPW9YHg1aVeYKYR0wQWQDP1DG/oXsGKUwexAAAAV0Gbb0nhDyZTAjf/+8HPtP3gfFlpFrQPirkvcmukg1+owUzumItdD2e8mcIUJrYgE5Sa+gOebQM6C3UJSm7qj2KyQP6H+uiQQf6dQ3nUsv4RzQXEB9IZsQAAAHpBm5BJ4Q8mUwI3//vEgWa8cB37/p9ZF86iN25nKIigWeKhgv7YJ/UNEZWwthTtyyAJEfsOe3Z46qB8Y9u4aNguQI82wq+DpdZlQKDejU+B0T7yh/vWZYzvEjdR8mEIAxxdRZpU68SFTtrnTu7KpqA68eg72PoiJCihwAAAAEtBm7FJ4Q8mUwI3//wC0LvBNTmF1Z9IfU+yXng9Xlh5jOvxWurF0OzyKrrbAG/LhzXREODrp9/66YlkypRbzROY6MEMzZYnzbptXwMAAAAzQZvSSeEPJlMCN//74ByBNNkfODOhV2l7PuVP6jv9DFZM/xRiaZ7NjolC2zFnBs1tnrqTAAAAUUGb80nhDyZTAjf/+8dnvoPCAfKn6R60RAOrjtGLSYzJf83HxmDNtKNcZR3s7Q17gEYux4ptcFO6KrQhBRjlRY6a9irUPBzxRSHCk6HKHHiDGAAAAGlBmhRJ4Q8mUwI3//v+jz7KI4HFxHpTS0oxWxOKguK69Iqb7Pv80RE8zRIcDeLBSrbSsht1y6Ct5KnKYedMgtXkRlKks+2+B7FYsVvxecCZHOfXyHzislTlSVnRDtTX3Xm3aOyRaBXu5sYAAABAQZo1SeEPJlMCN//7wZmMlyhHAV+DQxbHCYKLdsYP8alRhrAXEGFKASYmbvEaFUC8jJbbP3+VwoeATImNqxZxmQAAADBBmlZJ4Q8mUwI3//vgHIE1raJILnXvo8sZurnfn67164Fu2/h+6+ydNTr/+OxuokgAAABKQZp3SeEPJlMCN//7w+7Y4o+oh8mWdvl4UGi7/t84ZAPcQ9+MBMDQq//WhRWSosr28nZ9S+NamaT8AaHjCriTA80H12cPmqsMBr0AAABZQZqYSeEPJlMCN//76z9RD7r4RWowH8hXR1qM2/0RPeiKZ9pQxjXugxDv5Y/OlrDSWjfi+WFQhmWEG3gYNg2K8SmfzQY7iAamgvUBbTcVy2n5t1ecNgemmcEAAABoQZq7SeEPJlMCN//77lcOB3OQuw6HD+AjOfELwUGp9SJicXV1sFnhvaKjnvCxiZwl2jex2l9TUZcaQ8RCwsTwL7Cb3QWKztEEttb0Jo9sjMWDnWrh5iANpRQyrCjmn+qlJDlANP179jgAAAAeQZ7ZRRE8331No5t4g3EwRcCB6QcQmRnRJJhqS7+1AAAAQwGe+mpEf4exElW+Y0rm0PJ/wbTD/lW+CtIwcWOFlp4l4dVOEXYZReMT7JXZAnpNCK1hAhXjScp5NYYnhBjzjGD/rSAAAABpQZr9SahBaJlMFPG/++FR2ESusZKMxurmUPsaoB3yvXwtgjhKcOw7RRICF1EOtbOnMwE+WuGfg86L92P86V4YhcXg6X7+adBswWrcPNmvGF/+VM2MJnXpcNhI8G84ECcACh3CpkqclfZfAAAAJAGfHGpEf4jgOwqrh+lSqMFWNRECeaFV48q85DBbzjXglEAaLwAAAR5BmwFJ4QpSZTAjf9P1UW2P31QTlnuH9bjbYJ7XWt/ai9e9KF19PysQ7DabFT0nEsV0yceUNAJX0BBUVv6dG9+bf4cKP6gx1WQaTSkLrl6RjZl8G8g+mu0jnqW9tdKrF6RdsgVAMdnm94bRz+PzA54Zyuy3ATHLl2cUlXEyCW9KigPQ0J5U3y/k52ZP7lyfuZ5la/1Q9Pj6hGjOXZEUFpJACy7fj69cHAm5TQWn8MCi9bYtYpHNqiSnSLn4A+l5LrMldm0cXmrlCIKI+1dB0+ItJmzaYll6630NuKKUYK9we5JspVhzVBRfoTKIYbqUfFWokuJzg/aLs0nAnaPalKUUmOI8qX3c1g70zXzdIptzRWo0lBpoXPimAkfYfsLAAAAAUkGfP0U0TJ/Tuf49tp/KzFPCIEDEvSN1rnYJbiqzzxVfJ/8OCLsf7ljUGroVfs9mei4RJMtv7Zcn6zTeStOI46ruvNTH1vZzMeIyae9K1E96hHUAAAA5AZ9edER/h7D5Wmcyc1S5cFCuXC/vU/dYydSXmT4ray0A6VWt90PLvJGDUf11AKA6FxyAGoDLyBHBAAAAMwGfQGpEf9fFlYaEHVZXMWGGCx/WUqHrS0vhwUrligXTGXJS4DE/QUk+ncgeHLqWgy+vzgAAAKVBm0NJqEFomUwU8b/oMJ46ACuOTXqhV0ljwLZShz4Tc58bkcnq9w3uogc6YHf0kLpTfoclGUCaCZ4AxrrKOqJsYXwsspdKo6DAKGat7GSKdIulJbXbhuc/IvIa1CF/DNWzV6r+cO1H8weRmlyh8KFBzm3Sb9biHpOZEdIVGO3TyFoUJ3SD0ZzyH7EROwGbMEzO1FkpRP4jj2zcTUYldf8k10XzYXEAAAA1AZ9iakR/mI64NNsipLH3uhfv8LrdLfjRclwIPczo2kTYfFSqf6TNBKKbBeuQKHqSSWNQZ4QAAAA4QZtlSeEKUmUwUsb/++A44EJXr12b4QrXF3nTB+IiC36W7K9k0PInXSBIMLwHooA0xe2aJ27xzqEAAAArAZ+EakR/h7ESVb5jShGX8VlCE/7Y/nQZhvcjseRoWw3awq8RtUipbXM1BQAAALhBm4dJ4Q6JlMFExv/9Y3sZGUqwJ/ZJQ3AZgBXt9Ufjt4ZzNxGyAfjakJd+hT55Tw9cNhg+w6ug3vz/oXT+HSXuNuSH3FG5/LjmCDKvQYF1YEcEbLkXk9fh9BECrWP0tPknccp1bt+CPAa7hxxvktQiTyB05TND7a4HjRFbvq5TqgaLwDS28Cxz/2MBBMq7yg+pQVD0bX2Sl/mRKrv1ISvuJ2v1VJxyBQ+KKcEXAsgkK1RpCQS4jATBAAAAJgGfpmpEf5BbOR6mhN42Utt0sKpEJIAR7fw4EeQJlSVwz1K/GoHlAAAAeEGbqknhDyZTAjf/++A44Hc4+26pb+I0qiyyaTY86qeqgj4kwIFKO32Zd3pJfJsEyCt3X/2gIau4QMjcRp2plA5yiGVX5wEKuxlMS7/gGuv6Qw92srH94+7Mg9EV6V3w0yihnRIaFYSL1USPSa2cDMI6VI+grv22GAAAADtBn8hFETzffU2pysalekPthsm2ZMTNyS424hytUB9PxsDoKgpOtQiLMZ0iPUVzASv533olMcsSGUtWzgAAACQBn+lqRH+BB5WG0Yk9tcQHTHiL9VohRhZ/SCKeLiSmHZaZBYEAAAC8QZvuSahBaJlMCN/HqeNQwz8YC6lddTUaq/wlNA30DyObwUgGgN9P+XYVKGiTwxMURC6nsd69EfWbXorW5UKp/iLDB0tXqGqCdJCKds+YuzihvwR1B4oMcWI3hPmbK4zhnIvnkOawDk1lzJMAFAd5Xs6In7X2yeUfKr5WJHmhlN0mNbPW0T7jDw7LEkRVdNs79QlUlWSBMCb7P0EyD2FPtivHzwZbF/QBKslG6u1N02WmHh1SzyYHC6SYhYAAAABSQZ4MRREsn812AW0R7+Pctap0S/62v/Y0HOJuqNfDpqTpeIBGPqgtfdqY1GtZkmBpYYfpCvY9UDU+DGLJFacH14P9elmz+ffaFFAfxY7tFIA+wgAAAD0Bnit0RH+ezJaXm0/z6MpMxrufrdPmfUHceEr1rbvBdYpSb3s5MV2ylnH+g7e3nIEQSIuWIYQBSs9KfJ9JAAAAKQGeLWpEf4Z1uDH/duS/8VkmZodKFFLza+EKeodh2cw31hLmuImTU+zBAAAANkGaMEmoQWyZTBRMb/QsGt4s8uUfHR6wH3Eruy7jdPNSHdGx4H2+eY45bHHjrvu+npGPfuKs4QAAACABnk9qRH+BB5WDztJXfGFrhk6Xb6COb2LMvQK/NgrJwAAAAI5BmlFJ4QpSZTAjf/0fwO8d+rKuAd3nWKoNCPGmENKEldonVkGjg98w4Me5Z48NellC0FxCEd/kdNLe1cM8GE783LvVTMHmMX1wOrayRMiu60nJZ7AhcL3MR8AOnKU7oGDehMw21vsadVsuyATMJtrcH76ypGC8FIQF3X4VlmnhpX1WTqIkRvuvbCddYWNgAAAAakGacknhDomUwI3/+/6N/UQ+1Ea1+CN+lmcvTOPgijtDvzZhN0D2G8CBJq0xATOAf5U01N9dCRuCxIu7R2vxMfg848kgrRbFR6x23LTcr2silj6PlhF7bXGnsSksL0sP7ItpzFz6DxDAfIEAAABQQZqTSeEPJlMCN//74D1EObFx0Ccz500Sa/EklQzLOUat9ZefscXQjlR3Mx55ToX4Xuw4yk5jOe10fAnmAMrV30VKTa2W5vE12YQot/5YC3AAAABHQZq0SeEPJlMCN//74DjgiVvp1xjN3F4nu5NzvpHqsfivvp/bbY0P4LsFxB1hB6kH8V7RfuzmdrHI2vT7+H8cVEjDPvhujc8AAABdQZrYSeEPJlMCN//7wcOsd4IGgTeySKCPIuwxg0YYo+96hJ2UPP5xSEcCoL3qXVqKFs/+7whBdM8r0BuQQeLuwVm53f5UyL9EOx7Lh69/Ofouy30Ld83TsPD9W9AxAAAAREGe9kURPJ+A6iW+6tW/uR9lhCZLeuQ4oZcBFeKAu5UN4mG1qmIlUSZUB+VEhmQryk+9/puz2UHFvGA+dUoS1cGBKKWIAAAAOgGfFXREf4egEy+iMAaiy7zEWbVUnq4AyqydOp/rPmZg3eGRT1NmWJ3z0oQRnaNxgE548Z1QPX7xmOUAAAA1AZ8XakR/hnW4Mf928rmxp4HnnBMUqhe0zvX0oeNpgN+7Uu+KKS3fCuekbQu5jS3lzKvhqNkAAABAQZsZSahBaJlMCN/77mtCBIDC4q6hbKW6i+2ykwSQy3Fgc2ZOrZmrvUZMVzT54+uJ2QGpl/6fQueAl+NdAyBz9AAAADRBmzpJ4QpSZTAjf/vgHIE02R84M7FU/ejsVHqM66DEeEz+EEL17DltFUqXJViZAeKAf2cfAAAAZUGbW0nhDomUwI3/++FR2EKSC1to+4xCmEjt5LncA0mqNBUks+Ym0IghLbBCqyHI+/o4VwUizGvOHaK62P8WHO1D41LvGG5uRQbBRxa2T+sO0fWvHLEeFz60KXZuVfnunSGseJ/OAAAAfEGbf0nhDyZTAjf/+8Gbu08WB7eMlZakaObAXDazwj3bsg0VKVyXEZdPl4Ml5p/16cJP/U78Iz8EJFZ2maeMrllGWqaR/jYpXDX36keDWXpidqMEWfOXkQ8fZImQl1OdRErD09YOHHaf5NGC+6km9YfOUDqGX9Pq3XQRj3EAAABXQZ+dRRE8n31uIfah+Vvv5ElY/3/K6wW7DB17bbOdxcR0rSz4be7jHQy0EgS/lLK0iaQUrBIQrvTwd5mZPmZ4ApMYBTCkxmGGY4ysfzByNNQiAJD3piF9AAAAKQGfvHREf4Bw4zB2SpmzpJNdP//8f6dskHoit65Hrw6c76E2VmXkD2UoAAAAKAGfvmpEf4TUK0MsYlS18DDnFz7Y2UwUKL7z8XXLeoLGQWrrq0RBevoAAABiQZuhSahBaJlMFPG/++5XVEH0KHETUXGSGsXNMJtH/jNeRzgaZvNQtjBuinxU/liYq+MxIz8WoClfRA/KTLqF4/Fa0SAqiv0o9OwWGtPhEpGIZo1GD+Cou/3P4e/6cRklt4EAAABAAZ/AakR/h7ESVYBQhM1ypHXwCXiUkjPjSDADDv8Q18gyx9OBwNx7koSQDcnOOQ671+kC98Jm6w05OE/cunBeYAAAAHNBm8VJ4QpSZTAjf/vhUdgh9eh1QZyogZ7+hmtlbPNslgd7qIrlLoeRIb/et4+8Vq1ePF+NNLB4Nr5y2LIdg8c54s2iT1WYUe5vxH2qfxAHd6oe3EPUU6bAiHT77lV4nOsi77nk7apWUyTxHGqurceFxTOBAAAALkGf40U0TJ9/cH7YcKBioTwOqZoa6GOmGKX7krY2+ZZJFD1Q548KXkilnIlw490AAAA1AZ4CdER/h7D4ZJjNasvMvYzkHey/9PvUWs7fgtuPBRXLAX7CfyYyOMWD0kVUTTq8sgZLsG0AAAAkAZ4EakR/hNQrQmu507TYOsLT3zhWb32toH/fgbuu1qCFxEeBAAAAVkGaBkmoQWiZTAjf/XbliF2xAJMQm7dNMvTYZ4BH/3IqFg9bdiX1lecwbkkdjr1LW6pM2UN1kfniZvPkJBXmPFV59ScvpJEqbeREB1KHz3B7jm6BYferAAAAWEGaJ0nhClJlMCN//APDMcB3mPztiP3uapuxK9w1XpJGf1ey0bVBH2OVz7uNe1T/11fYO9Iv7Fc7kkj5DXP4lwVGQc+vNvTeM4IHYwYcrs9kxs6CASSVaQ0AAABKQZpISeEOiZTAjf/74DjgJDW0SKRZBRnHc1aiJ3lEevKh7RRg1AT8zcCut9u3ECa53MXXjAVR/5TQqy9iMtme5BHk09PtAdSIfrAAAABbQZppSeEPJlMCN//7wZu7TxYEJkiWNfakae+YsHrw4VAfICk7Nb7oLk+0i+Mpuy4D8f/TgFMT1z8D+3aoAm/Am00CBytVnpVYqPkD+waLnY1NOQR2k51+vgGAegAAAExBmopJ4Q8mUwI3//vDoZtU5QkPFvbm4haJgZvT6vbhkh4T4F3xSdoAUekhgoFrQuyi5WLE6GUiQEppgDoxQGgzx7AhRZZEztPG3fTBAAAAQ0Gaq0nhDyZTAjf/++5XDgO8x+Iwxa46cT58hPz/3xXGJ1VivV3KaFJ/3WQxqtPg3LKAamhPU0NAdWGoS0oO7WTR/EAAAABfQZrNSeEPJlMFETxv++A44EJXr4MQrxHw02XeFmca+qmRNOoUcRRGPhSgHkAcuC9vwQHiT3krg6bT/BkQSItsA8CwlnnzN24kW98J7AEeP0NiNW7ZKbpdvkMpFyUz8l4AAAAyAZ7sakR/h7ESVb5ji8/rDDaLJg/7XqiPJ1B0qKo+hBZfe7U4Rdhl3xjVAHcpz+vB/FEAAAAvQZruSeEPJlMCN//74ByArMYDm/bz8wxn23X7eeY8ueYRJcmDuVxZyy1hXLUX1ucAAABgQZsPSeEPJlMCN//7x0/RSwgInmw7ggx6hfLg8Mo63KinbZ7POUyq9lLfsMaS3xKvOrwHbYeQNzz2RoRgXOv/exp6/Mde6Dums6ZRED4RzeaCX4PJ9AgJuwbmO9tXNIbBAAAAZUGbMEnhDyZTAjf/+/6N+OB6zo5Fy5eHpx0/tejGu520bLkOKYZmUYwoEKtu0gsWwF6ztlWdsn2cBXYI99l6dfY0me7vv4CJ7ndd4PxDQYUsSz0tIM6CYnR8oOrc/Gfexffa12GAAAAATkGbUUnhDyZTAjf//CnvmQIc2OrRA8MwDA9HMkp0e+MI7Pm6epVIxMw29SeVpA1ToAxIhYkHTNN5ShviDJ3BuOFRgp1C9Ah/PPMQYFk/kwAAADJBm3JJ4Q8mUwI3//vgRKhKrrGIpjJo14wE5fxn+jkJ7L0hg0Y/aAfD07llO5gOpjeyTQAAAEtBm5ZJ4Q8mUwI3//vDootTkBHWutfyL3CvbJ+vrkeH295y995Ase9ZpuUtjRd/dTZ08jGdVP/9jByukXhfEKbwAjssjjm1tYa2KIQAAABAQZ+0RRE8n4DqJb7q1BrisxfO905/2QQCdm2BOraG6khWyZ2qOCWjta0B7hKGZCu7bgEpU4a1B1yA6YtIK+nvcgAAADEBn9N0RH+HsQIbL4x/P9Mv+bVwIYmppI+PEO0zwOR0FYU9+OSfjcQXZbl+lOJ/yuuBAAAALwGf1WpEf4Z1uDH/duR8Z1p/stCcUr053SmodhihH153Yzn0yYfGEtEcJcDx5XQwAAAAR0Gb2EmoQWiZTBTxv/vuVw4DvMj6PtjWXbeCDWe8HbgEv/NcdLhFqVKB8exZhz3jCX6jioMxFUtp3r/RC8ltZJfrteOt8L2BAAAAIQGf92pEf4EHlYPO0DFLXPMyZmip/4spisSOYUIUhCPWMQAAAIxBm/lJ4QpSZTAjf/vBGT5YQIc2HKxMsdoJ483hDJQtEt1oSVY1OtNsacpcCHC1hLiyide7+igKkUYmgW63zjDEC973kPfv/kGvM9IQJRrdPvgu8NP7Gg4FonN9g1w21bXVxN2U+b+Kba5ydK3Fwptozj887wZ5U43AxRVa5ir6jhSItTpA4qmjsWn/cAAAAJVBmh1J4Q6JlMCN//vBzU+RvBF1xHc1oj+U/TZywQssct1DlyEQP3bm66D0erMGNf+QJB5DFmPKyxwSETeXY47tb9OQ+xlPePfb355uJNYeS2Og2DJDc6qxeP+J/6dc/Sv/p/LXg678yyOLMGYC22x0VzI6myko3Y2jbnwPZYHoVN3W9bTc2x7RxOC/79m7UesteT5niwAAAEpBnjtFETyffW4h7bT+QggNBgC1Tx90dvwgqWi/2bEAaoy9k8G5OWvd6LsgFQzqIuBbNMM8rzbsW0pRv4t/6wnIDQ/mCBNdzbJDDAAAAEoBnlp0RH+AcYN6mjmSXWqx+E1SBd6filHrC4NN/gK7ATVZjyRHQcmyIQAZumDsG240peDLOLW3gVbq4LmG3Njs3hJcfGYGdq1drwAAADUBnlxqRH+E1CtCDqsru19uEmwmchClU8Q9+ckZEbxoQUVCtJ39bW/wYXdNpPtmAsJQybL9wQAAAHlBml9JqEFomUwU8b/8A8MxwQo/6EiaK37P7zGFrPazCtkNNv4vIp0kq5kSnVxolGnW0VIejbpP5ZiiC+I/F0ma3+IF7CpxnIMPf4egyyXNCkjhnJHSXHusxYcAUqAOlV5uTf2eoJk7NXqf/lcR5JWs+Sf/4CuRrynQAAAASgGefmpEf4ERwYTMPCynKnF/ULnVV0i5kdCM7P0vSFSGSpVSXT+9+3FrB1//zAGZ86aGphUfLcz/y3mFRVu6iiyShLlf1IBxT4rwAAAAk0GaY0nhClJlMCN/+8HrVh2DajAEnVrgIluIofrb5gjHi7I0o+eSqEr/KXkOiwTxu2NN5bbjgK7wNu3j1ZZsWnH+OrtoEaJsgcV6iH5Pd7Uko8sgwhTol2e1ZdJ/icre4wKyTmNr5bcC5di+bdY2nLvrzZtdpa9DX7QJestuRmj2U10Cl1As2ZueLxIRqj2wGBHESQAAADpBnoFFNEyff3B+2HCdri0/UUVYHd8epT2d+YMYObZm8nW/ez46SIN8EbMlhsLlCSFaW4sFX2nF2P34AAAANAGeoHREf4ew+GS+7UvFV/w2hI9vC/vTUrzGx1QcGgUQ8rwi7DMOqBjdSkbuDSjnaADQL3kAAAA0AZ6iakR/h417ZNpNh7L591MHWZo7LvvoIcpQmohn+Xd70kfnW1Ud5bf8T7beBDmR8Xcd8AAAAGhBmqZJqEFomUwI38oIQ86YuCDOQ5tOdExAsQJusAUoBl6PXbngL0EHOm/XgoQK4cIAx16UWM0dQcvANQqKp2EBOaWAoRwWN1bqOg1bQrKKmV/ijRbo+vvcxfowunt49/emVGo9MQAXvQAAAE1BnsRFESzfg/tuGYHKjTB5ZID1kuLiL2dqNk66JpJk8w1d6Kcc2wUWTZM13JvgxIJn74EkYNkgC4PjHLlyOy0Hwh6m6YSunCOq8UCBEQAAACcBnuVqRH+BB5XFIVzyClcf37dfxipiNXufl10iycx4d2lwwnYOlmEAAACBQZrqSahBbJlMCN/7xIFnX2UJY3B5G9rwJVPFBhNUhdhx/fLZ3fI40N7ZzzVF/KfTN++yvpUue1sMAod2GLm7S5Mnr+HDpXolz9xb/2QF7HgW6iFPZ0TSu72RZcJPzQzsgubn35XTFLL3mozUEl03Plt0iKzniOced0OUW4K+kw+BAAAAXUGfCEUVLJ99PbzhStPV0DPWCp739xrVtKTUsEdCecaEM+e0OjfkkUkuNJFophvGmqzcNDPP+G6mO1enkonn7obx3Z30VXOOnNy2O0Eiwd48Dx3VZnk6OQouxGUAgAAAADUBnyd0RH+HoC9/9Y3zHp1BIl2MfTFJGRIlfr2DnVE0Sw1ikOxYlKnFwH2Ru+ncOjZ2P3qIlgAAADEBnylqRH+H5CPxj/uvJW0xJiOenPhWEISI76CVOVnP0kO3twxfh5J19iCz++FpQtdBAAAAZ0GbLEmoQWyZTBRMb/wDwzHA7nHfcLHdl3CBcWDNUIHw5qJ/uI/Rxc+UScIeicTquE85DXU/XssWAvri4jZAdY1hh7+q6h1EQeWGPd8v9kVk45hJzDq4QFLRn0armvMglMXFFZrGuXAAAABFAZ9LakR/gQeVg87ywhqQ0UZ0H9SsmVmcAjAxEPJNkzxuBCGvKgW3OfU6j6sPxhv4LrsmE68KdHBdFoxyuLi785Fg1b6+AAAAhEGbTUnhClJlMCN/+8di5MOwazCfUdt4QXA7i2vzFU3g/WQsndT77JVDSX3dyiHIJ+dcHuYgNieBhDWi40XytcT7MzqGHQ8BKd7emcKr7t9ZMy2246DgCMU6O8nfWZamMp0hGAd1hoDnPEpKtXES8lvdf8QvUpJi7H+RY6YmEvOOXCg14QAAAHRBm3FJ4Q6JlMCN//lIWPInJdToRM/Di8/wOIa7fAwqyFmQUIqzAB8UxzfJluX0+Td7Rj1n+911ubH3x5Ng8MZiV90GkBT1677/EI4xf49/XAJzyB4a0vGsKyC/rL2cdk9MBpe8qr+6ehvxBOlZU5I+9k8HgQAAAExBn49FETyffW4h7bT+VmzHLoEDF+QSPtyOibdtnOwtVf938/IoeR8/CfdhBqH8efnQNzSTcptOtFhLNhHsC9QRs6JUpZ3g6L+0o2hvAAAAPAGfrnREf4ew+VpnSNtXMyHIouF/ep2oe9EQ/UgNEJmOXn3+zo/llDpYdbcSubbJNj3j3tm/n+QkREY4VgAAADQBn7BqRH+E1CtQOK56tKoCiM8K0R3dr/Xh4xtb5DeCzfJ3ISlqMnAuma7Jff3Osd2p1kQ4AAAAeUGbs0moQWiZTBTxv/wDwzHBCj/nXShLUbtqalbXUXB4c+Jr8UsTF9vlJYy5JAsd6N3YqAwRTAVVRhIyCYy8OFra9SCQJIdvyrOlBtx6e4Nm3jz/vVxUbjh4mewKCUl3melRvDxX4XCovgeHkgRgAO0Q4EkZqLJtHfEAAABFAZ/SakR/h7DkTgDV56NbnZl/ml39qIOiEFfzBef+RWvHUibHEjD4zje6WjYo61QyA138/8GbcEZEmbUXjgjCWt9nPALYAAAApUGb10nhClJlMCN/+8HrVh2DgEzkPQSYjrWyzR/cxxfpkyJIH7uTdsyWu2DHpet18vbjP31VvB4OvU2zGwNJhH03pIxozyt5CjT0DxM8Eu6iU+hynjaNStmVahBNpSsriT7NNCrjVDXNYImUZaYR5CNKrD2nPtelij1qj2HUukGEN0p8HTXl4uhVEfUSYArW7VPcWydU5vlLZYaJeCj6iL7TK9654AAAAGFBn/VFNEyff3COo0/prpbjGuq+Yz9qQ/0UtxBbR2wSbc63al5dG7nbMBmTmMJALX+gAAHd4U3sZUwDf/j8MBK/Z+midrtvZ/dhX7BGy5kbGSKhcRoIJqsRYrNy2G/mf35bAAAARwGeFHREf4ew+GSYyDT9FKwHNGCE3zTvKrR2YTzumc5w1YU5tpz8yyBsGigtzFTb4gkcmmGk8BosAJh8sW1hPLsi/oP9XLQpAAAAKgGeFmpEf4eNe2TaahDYDlAE00B5TiGe5P3cqwgz7Jz7dAwop6UVlxL/wQAAAMZBmhtJqEFomUwI3/vBm7tPFg7wyzAD0Mnjaxg9o1oowtbJ4Iq6fMSBsDEVXAUv3Ek0V6bUA9AyXfdXplJl5OklAZrAT2DLE0CTU2Gb1ph7uMvZ6e730j1uTDGRcia83XP9lLWy0N2lgd3miSshfYKbuLgLsYsiwlykYvR7QV0d0CqGOqo55SN1Jjj3z4lpX6c0RGgzwd3ziVIoWQeFFwVeqqCbllBnRSSRaTmdGVG2XFLb5K4/MLixbmd/kI/o/AQS9ak96IEAAABVQZ45RREsn31uIe20/lZsxy6BAxfkEj7c3vTu2znWnO0o+LL2TewvhwqHeJD01QF78BcCbCuwQBKbqpl90L6uNZIPK2wJ2IYROSzDobv1Y58U1sN/QAAAAEABnlh0RH+HsPlaZzJL659u40VH/5IqeKHVDn7np90TGQtQWPbiEvzRpaGCnUR9bwO4lkoO9UNv4cqr2uoaHtjRAAAAOQGeWmpEf4TUK0IOqyRzDxju6fds63t8C0diXyWCw0nZtf4XS2WBIN2TqV3HMYRmD9ycxi04LTVA2AAAAIBBml1JqEFsmUwUTG/8A8MxwlzTJ87d34cixfMcHbv+VHu06EPCplmBHBDuntSkT2yRgALAJie3XxuB3xfOH/D7ydP8ErkN7FeqXmM3xkd7kLShI3yZeVnO9ZFHHuwsNecvxwV8jY12J7GCNBCKjcvRWa0RiLGeyRhv7w7uCpPSTQAAADwBnnxqRH+J9+qocEqY3mgUADG/qVAbloOL+4LbHt1Vueid0r6SqxaL3S0xMXyghgv+Ji1+lINFTyTiXpEAAAB/QZphSeEKUmUwI3/74VHYelCJVnhRTg1BLYiZcXT5dxyk3eX+rJNoNiJEdunbKmh+C5I7HMdbrb16UYYiG+jet6KLxaN1w2WtwbVe3TezkTPnAeJ+tInJzjZlV1FqpcNW2arifXJoFuPI/Iz2s42jmNCDbRbNxQpODftQ9RldQAAAADxBnp9FNEyfeyCnmveYuSODs+6cbniKe/vM4S56nhz+m9i8Zex7ONz5LSvfzcVpUjqdCJWt7Tx6czK0loAAAAA/AZ6+dER/h7D4ZJjHyl5mjfaYPRhVd6h/FiZQAFr6JH2QqvbB/RwwwlywmDQA/o5AxedWALtCD568UWghNBmBAAAAMQGeoGpEf4k2H5XFXAtLnyu/P7BuNQSHZxysH6gKDLB09GXBu1f/byom4T1PU9j8d7AAAABwQZqlSahBaJlMCN/7wZu7TxYPY30W6JoFSNQpwVvW1QrXI4toGmRdXr6ug5AcIzYxb3sVhpDt2R6YmUdAxjB6i/FCQrVb4p6cx70BUXqsfihn6ICLLcl0nvY1vzENxIRzj0j24J7cCRqf3p3r+wKgQQAAAE5BnsNFESyffW4h9qH5WbMXYFj/ael8uqKOT4Kdx343HScFjvKFNiDbjO2olivLVLdWGTyFi/Nn+rmW8yKVjA5v2wbZq/jG80u3A30934AAAAA4AZ7idER/h7D5WmcyS39QLC0VH/5IqeKHVEQEr8EriQ0rMLKqwWs2RvSTdcH/r+SXrvRjzQL2kGkAAAAvAZ7kakR/hNQrVcoxKlr42ol5Yisw4FbJimCD3s7LQHkiALNax3evqAJ2Z8nqxtEAAAEEQZrnSahBbJlMFExv1WRxLgVrlSU5kvRPQefQNL+hnY+5sxyM4OGP3W6gipxl+VeQcVtIs+PdwqJefHjTv2e+aZM41f823N7uJYxEIRmdHjyj4j8dVr5J5jvbLxXLbyf1aePF+es1ewhAgs3sr0LDYMZnAVaSTpMiiCjkwD85nRHwA0GmWfVDsiw1fCQ08yGxuwdcMd61SPpfRADYJpeVmWilLIkk9w7lCbACvWCJVPVxpYkJW8dIszr9QbdlVVkW+1uatXpBfwsbLty6mMFmnr5Fm2i6dwp+IR9w1Urt++4L+nMYz1LxgXCx3//Zxzgt2ANHgQF4NuaWK8l1q7TAxuxaM0EAAAA4AZ8GakR/3DTbiwrhhwSpjeaA68X7+pUBukHe3+XGyF7HhYlJIdhyOmcdef7GeSfs6o7iJKfivUUAAADJQZsJSeEKUmUwUsb/9BIkgoUyga5F+ZT3IdOhBODSowpHgJE6OOIA0hhuJSSkDCReDJDR0qEUS4FpJoQHvLAlVaAWTk//H9e/17IrpZwyJ6wCC9sEeRYyik1ullhgvUDdNUQzGaLfMevgU8An/KAX0X1/cOEMqpcK9aaJxowlmuUS21UtcIb9euUUnTFnUYwlhmrjdbVG+lNOTBID04MLbIMZIZzw6PX/9JpiKMQPPmV99BkjSWM89NuJhl5mlC5ED/6kt13QRfBgAAAAPwGfKGpEf5iZsaTq/dqTsanc/0IT3kf5Hr5m3oLGjyM1iBcL6c2qPe1QQjiri0diK6w7972kh1OttrOxLezWwAAAAJRBmytJ4Q6JlMFExv/0LBqptRAkFpHKdI1ZlaFSI2+VAhWAnIutcz+4N6enTxH30oYaYrWXeNE7ifjDCoHMMyW2gFUY+wbi/JE8adAKifO+RhxPHcvsbcbFsnmwVPaFpbwyxJYTd5QCBNnA5eacY9cU6ZwSL/KPzlL1bcR5waT+lbJg8F9k0Ygkey54qcWn8L2WvH0xAAAAIAGfSmpEf5BbOR6mhN42UttyvF0PAGvTXN8LUtVDAEJ9AAAAqkGbTknhDyZTAjf/x6oDJtkwPh8BHMKIJQeKakhBaoY1uTU1Dv+mLEdoR9c+RexgX+HTtr1kNcl3RElgObLDY/PaRJqf4jDnBLlQj1v/j+gBpH7ME2CB67cRZ0xo1T0OMscgDL2ja/1P3y/3R5yr255srgQtnSzKhY+hp0JxNKYyYJi1IFRA564tmvre91OlcSUCU/OIFZJE4oY32tHuv8tWU01eFm8j3ElGAAAAPkGfbEURPN/QgoY6wn8yMqFM8ScEy9Q41T0IvLskW7ZAQd6JL5LdsFFkc8iVrAOSANPjytFHsQEpJgFmIQehAAAAMQGfjWpEf4EHlYbRiT21xAdL3k+Vm4+8bP5+kpCxYbPyDxTyetngNkKsEo/OcSIoI20AAACGQZuSSahBaJlMCN/0LBreLNnVwoJ7HokTqcp6Q9iGVOTjfWRdX3HNLzYq0LH11Wq8Jq+uiZjebL3X4L63Z+y9aEj8POW4Tzlvo4K7tfR97PaNJboUExwnpuB7qJVSNwKDoTb48nTn0tJrvw/O3EZNUilImcvTCGE126axIMdq8DiFXn5Rl4EAAABdQZ+wRREsn309vOFLdicCm2x3Pufzb+QqT3EAn8TFTgB+03dALeaGOw1/6I6cxZ1UST7oYT8jPMi37QdASPs6q4fKsSNxFICCpLX/nIOtEt/OgZpH2dK2wf0xARgbAAAAMwGfz3REf4exAhsvjH6163w40hUbVYqr0IJX9KQl/IjfEvL74S4AfJknBl2wEwSJxGK44AAAAC4Bn9FqRH+Gdbgx/3bkh9Wq+haFwumUOl4Nq407OonaoLPLjUaPLWQnkuQZ8bOFAAAAPkGb00moQWyZTAjf9BLL/c2UM1KQX+BxGGL7bb/E+C2iJSq8ZT/fki7Xj3xHv79/3OCR4eUyro+hPAO6kwpQAAAAbkGb9UnhClJlMFFSxv/0LBqns2dXlsN5umPuNMnDIBekdIzGY3PnWipO8EQ4VO29m+HWcI89dez+iwGudwY8934b5mTOutNkupDrTYAefwpBLHDS4f6nH0oDndONGfIYoajXM6v5ZsonVMn+J7bcAAAAHgGeFGpEf4jgOwq2BPvOKJdhwabiU0ees+grZ1YMIQAAAJtBmhdJ4Q6JlMFExv/BdFX04QlYNpjTO0cfnbEhBFxg9ac2tx4t2ngcyym5CpUChqKx99i1D5EkvHHFPX9/7KMXq7jI/q3wkBGpGR/RHgt9ep8OxX3IASAi+yKKBVIw9GzZv1fxmZ6HBMZD2MbTpq0Mi3EG2QC+iK1rC7QF8ooxwG2N1JEtqo2tAd03sU98LQ9M68dpUpC/xpxxgAAAADQBnjZqRH+BD2jlniCi2imSX6//59WeyNYZ7EThdu+jrAY1bVw+gxFikWceh3l4wiMCVEWBAAAAY0GaOEnhDyZTAjf//WKRs+/QRQ3PXPg7GyNgHRyY2qAL7E/g8+2wAxx57bb9N5LU5Lwf1jU218is1FOXmWQFA+ILNlI2cuW4VuFEAUs55tZrFJbPFNXTxKQFb6zgrOGYxN6gYwAAAF5BmlpJ4Q8mUwURPG/9YpGyYu6xTlLEz1V6fELYXJYCl0J6HV2jidgU1kHU1TaIy98emaghsv30jfu//giS4u7YXeCrhV7hICMr6GVBT7WEUf+L/NNHBp/YziKg+rPOAAAAOQGeeWpEf4exHtbuDH+8TXGnxL+KjE1xi3rYQfEyAC01BI7R1nn8TQtbK9qWu/fdySehP03udb3CCQAAAFNBmn5J4Q8mUwI3//1ikbQD/zwawOGki1yA8PO6m26UL2L/9mKgReK7+ebK/al/M7IPAICEACwHoICsV2NTBXt4+cDul/OajszW51KtkLVl+aMGkAAAAERBnpxFETyffT2/FQo96LwXOn2o93ozcOD5tpYSJ7cwIJok8j0UU8aRtemPmdp7YYIUiCycIRpOWrK31Ohdggd67CoH2QAAACkBnrt0RH+HoDzW2qN3QBJ/E8gMm82gikxdFp7fAWLcEhRiEwpvi8PEkQAAACMBnr1qRH+BB5WDztAwU5oZn2uFb+hZJTgCS3z3O/PoqGYuVgAAAGpBmr9JqEFomUwI3/1ikbIV0EXMsa3YekHPSpYuprfZeuR3WOSiKkC0kFZpj7xpSGhq9ZWjmxEgbDlVL++iKztpTP/7VY798ET282ga/R1srrqOxD5u/U/+0pygtO04b7giKBp6WzCweDafAAAAekGawEnhClJlMCN//WKRs+/d5TeJbh49Ilqhs9vkFuDy5tILYlZvTBUtgGhCajr5vBfprgFsCD23kwfPBCzu+fBE9vNoGt6/fUDXQT6zOaThfR6pRtMs5SAtMZaMwqag6Y2HrhAcrJKrd7zb+2xmYXa8/28eupRQ2l2pAAAAZkGa5EnhDomUwI3//WKRsip9r0Ll4vNW49yxGm0rPx6d/W4e0Trq4nQjnHlJpUyssaFXC6OWq6h/nBwV81fzGtv5FLjjRbokIbdqcOylodWrIY6RdCrUEz4tfkGqVbBtksCTQEmc4gAAADhBnwJFETyffT2ygMiGUO/inwbX3T9EEiT+fzwLws9sywg0V8+Bo+hNRgRMgu2Mkb323JGVsThPeQAAACQBnyF0RH+AbCrDaMSfDibxDReFdC1UvC+7xMpiCX1dYFy6YMAAAABTAZ8jakR/h6Avf9Ts921mTsNzZj1+OJVQ8DVfaupr23+v6Ud85kP8cxjIPbdrvuP+lE9QsgrlQXtOD6M87SDXTNYuQnTLolxoGd6Iw9QfFrz+rYEAAABHQZsoSahBaJlMCN/KCDhlVhQaAeDWA+zilE4F4JvN0n7qpgAwW3vUlRQKCzVZ2fDY8WBCzFmBKgdcnVnxRMksKfRYMOIlGMMAAABUQZ9GRREsn309tLhK5CuvbdzT5vjjshRdA7IpVn/x5vBwKg++SIs8yQwRWo0QTWUmAKd3c5qJygjJ4vikTp3AKpczzE0vrvu10mq+f//pJzyqB/lDAAAANgGfZXREf4MOb59GU02DO+CVL8CSCBWNmo9IbuMSHjgvPw2O5k2vA2ympmv//pZzZUNT8y61kQAAACgBn2dqRH+BB5WDJUsHZv77UUMhfzqZBSad9zQgrlg0O2kNk5DK5ilQAAAAaEGbaUmoQWyZTAjfygg4ZVYMKd+rrGfciPhIiq2gCFAqMZFI7DVM2ZD3XLjJGnQmMgOIvEEsS8t+XEZfwTy76m0Qap5ds9FsZSQ6iIFZ0GU9cLjijfT2W6sRfprNkuPjiKDz/dClDTiAAAAAYkGbiknhClJlMCN//WKRs+1FSeYueQO7h1DZ7e+xqDAxtzmGG2fToXnjatvz1Tn7mSkpRsA/hIvPhdRN3kl1bvHvJTMKB9UA9ykYobd6r7+FKw/F3VVri6aWYgE36UGcq4wRAAAAZkGbq0nhDomUwI3/ygg4ZVYNdbNtBrA1dyGL1aw7/5eOTzHjgYmlvgs0LekeTrnBmIqn/lVENJBHf1RN5RqsI2gVGmV4sa2fCQ5tezOGmg2pMMQ4mlHRXTuRSMJg746BbfQilUiiLAAAAD5Bm8xJ4Q8mUwI3//1ikbPts0l1vp1ED2x1ji6bVEURYnvYIR3zetSp5m1e6jKMQTp7+VkiqKd2tavPg/v5wAAAAFRBm+1J4Q8mUwI3//1ikbHOi3JXDnst5+9/Q8YozwOIgFT629Pl7FB8I6iF0Zc7IF96heR57AZnekg18MeZ8yh2uOHSqcUTsaug9Sj4I8zr7wmNr9kAAABtQZoOSeEPJlMCN//9YpG0OvyOUdYK1lQfyPABv19dN2EuVGroHzp2zUVU0NODsPiVBimmNmF/CSLBt4F+dGOjvRT6LiulIFp5sR3+G1h1ATu9T+HuPV1x5zNOLgqLSCXQN/UuvfZ64plbVF+9zQAAAGJBmi9J4Q8mUwI3//1ikbHSvGoxINYD3SLvvyxPZxlxKc1XKrT03i5NYjyxJoTxK6gyI9UVV9jQM7CeNG8TnLBSp36B8rOa2OVfe31qsGMYmSKR4dZKBGwbKDN7WtkaEsZZZwAAAGxBmlBJ4Q8mUwI3//1ikbOuOtG+ZLt7BBIpdazWzzFa2oIXuKvsSfKEX099+kQxskSmsj0JIeAt4NsniL/81T0XG32CsxykbWR//z7P3ayX2T+O0bkNhfiOKnIVANdHE3DTcxt7oZfS/97sH3AAAABJQZpxSeEPJlMCN//9YpGyKf88uaZQCcSFhy2E1I13cQ7T4qAfZ/efBEseeX43YrqDE31Oe51xgWu3dMtxJ7PVwhztNF4CSSDjwAAAADdBmpJJ4Q8mUwI3//1ikbPts0lsarG2hQ9dcXCqWONWEzXuGsn18ByM8MdGGWrqt78s6xW7qPrhAAAAaEGas0nhDyZTAjf//WKRshXQOHBpVAT0KbRHowXnCt1oBqGLr3USERo5s9MW+PmLvdT3UmgBF+tZO4KwNstRUb0WbAoQxjGZF5h/wf/HrDjvpjl+CJ1AbSxs9y9+9mJLpBNXqaw9rkBgAAAAUUGa1EnhDyZTAjf/ygg4ZVYVaN1iGuf33zgSf6omrlwSndrkCTLPF3gaG94Se1OTYoSOIPiUGjc389Qa3QtMLmdfxJ/9EJgaN8vT7Keo1Vi14AAAAEFBmvVJ4Q8mUwI3//1ikbJi7rEg1gbbls4vToej136f50b8WMzjf3gU8vi9SoIuAXfe8+Kxcl5VytPUCNZR+0dQwQAAADxBmxZJ4Q8mUwI3//1ikbPtHlRWB+qqDdxH6C2zc77qhJ1fZ9qy55n/wDkZ1sej4ZDv6vAKtYQSdyHy4zwAAABWQZs3SeEPJlMCN//9YpGxzowPspkvEL9KenZlAqwx2krQK3l2oCa0wsWc8k2NlSHrH+b26/w5V1g8f/pnNfDFAjYmFGxsaSxgtrssgtpmXxb0wraVhN8AAABqQZtYSeEPJlMCN//9YpGzq87hxUgQYaLWrw1yOFa3/w8/9YwOlVMZthDmUgBo65bZZnRaIfs7t/B7omgkl6q/wSNF+R+xHNbXAqgPYfIrf0t/ExK3QwGLCO/7HBF9dcZigKukbsI6FQbxgQAAADFBm3lJ4Q8mUwJvFGew97V1Lcd+O+QDMDFw9wWzOYcwFZxnnmFDt6/+J6XbkhV2LCDKAAAE62WIggAJ/74Mt5iA31diPnODkHlAm45CKmrkO6tYSOdbdSitE8v4aantb94gCmriEDoY7sxccjIMexW5zP9owVi0pMMzzPblktB+I1kmpKIM33vbcRLOHPv0LC3kz7y85a64aEb5O+r1+heOtihVeJUPASxV9SaWTlNdsSrJgj4xmmjiYMl0IjoWJ9oVs4HKkTYxl/l9OJO8f5MIUmsmAuHXs9jIbJdKBA47Isns19FTdkpBDNqpHKr7q8GFdm/zMWCdfXqX41NkD8CuzgCgpicALz3yxhbW+N0bt/hFjnzm/UaPrgZPcZzqDiV5eYefi4KrgKA0ZVe3VpNSqYxcKn0235Wuxf+zZBgzqMmKLS0lLynJtcUkdDEeIyyfK4LIm48C4xhv8CoyP8GUDK8uivbXBNn7CgW7AuII2yToKdSV23wBkjD62p+pn+OELClzpVps9i3NFELoEbszOb28KZqNUr1lS8Wduuc2YwDKwxXVCCIeAzfkKMHwScs52kSwYmRCd/d8allKuQ2V4hIlH9Jl47UyFemjMfV7+rS2jrnUFGhJvwePtyEVo9P2WdscKW/cbYN2YTUkX7ujH/5UC056sGo7QCtikXJOlN9RqsmVVLXhNH901MjEGXfzKbWH5NYOmyN1mZ5lK/o17zObpovEPe3A6lXdH5FNGkd396zj8jXwQhCpcisxza7aEuADlEolAttpaS0aItC5D/12ya/vkaXhuhqDeowf/qnrMJEgsm2F4woahQkdykwDfkajV8Xoyi3jBpM/z3SsBgY6HWn3rW0ngUlpH53febs2YoTMxFBG9OvOZZ4WIWOcIb01o0qnK+4aqXMILB6bnNAae4qEqWz9+Z1pCpWendWC2cgOI5T5MSpidy6h+FiqCWdrWfJa2hloCwt6h+wVJKuhO2EUZjv5Bfuy8bQtvt4Ewrd0GT6CUsin3tuzheiKHIY7c7GARYxWOVsu7Y2YBaVScU2CERHIxqjA2+m2S99h81bJehej0iGHFz5U4SYskqcTAlUyrY13Nasq23dV54zIXVBj2HdPP/FgHKntLtIpTJhLgMLXrBf/HllTDNcthyCVB+wy+lmVyo5EGwdFVisnyRsii7TpXq3YGN3aKn3/gz7LAcKnAw5oD+vPQthR3tO2J/mgouNkoZ2GTWA2usO+gX2YiAf0fCodcAUuh6p+bh+zP/ArJYNnup5wOWspOUA0IovBT2AJTCeoNMfOAYnEEkLLP5P14x78y2wW1gWBgatjq9keNhveNxm3OId6Z/+RTUg7V/jZPZKnQJNBEy/eCv4qmUjDL8YpXPHrFThVUwNS3NoSGhVS4JZ7bpB03I1nSKSTH4DG1Lf5C7pla5OMYG3Y6n76Qbbfl9EZ86iKbiUcdGryD/jMHV12JykpDiaDb5FHPDkHH/vDAgnyicavYZBKgrx+j9OSTYucK815/uPHKsmteJ4DuUy00hL+bTACJ+49So3QPADbTTVdYUxcp68LjjJa/yyomh0a/d6qcsV/aT8s30oHmWbqIBloL+J1PNY1MzOCCU1U5ZMaoTRdCZwnZlEC/HchwPUMHg2DDEvmKsv4PqNfVSZC/GfvnrurcViZysYuuCjXGkUlYV0fNLrTBtq6ykyQkJUrZqjQdqw6kil1NFQWrOaWMeGubzMX3Bsgc3u8GFFquetJAAAAX0GaImxG//vteAAg9MZ+S0i/PC0zlHX/QQ3L0D4FiZmIOjpIFCJKKu+YprzBENnmxMhdyCs5hHE8xaJE2hp/2GNIqeNg/xvscT4V/xoCp2fXsc1CK8Yb8IojHKubQGZ8AAAAJQGeQXkR/4EHlYPO0mTyth87XGb0QiDlsCs7SOJZk1y96Bg/H+UAAACeQZpFPCGTKYRv++FR2IiwU0RA8qA1ChJxSNLn6LehKVP3kDzlM1ADrYhw7wd4Ip7AQzseLJdG5Wk6xUPebxn9TmVKfObBTfkJNs3ock1GokDdOFS6NZD0ZXPz//Ud2gir3gj4Dx1v/MhlLWLDnShmxzoy3/7AQNWKIUFk7mJHH+gLLt/06WNqUe+p7gC2E0AaYTuW0oNS7l5TDBLgt+EAAAAnQZ5jalPN/4H4Vo+TODPKsmBYbLE3RR0oNHHTburHIqYChXMkXhLgAAAAQAGehGpEf4exE2/rMWYbvveZW0hf2uqeKAaPNwi4iL5LnlJJ6paVEaSZbn+/5GomK7pT5XP9xVnXWAQw85PnL4AAAACCQZqJSahBaJlMCN/76z44S5pk3knFEhyI1Ta+1wxYdmynPZlwpDMB5imzzm+brVdgj4Ve8W9egCohBqHE8VTW+aug2DoRpkyvvsOPUbzL+H8XIBGreeSAS6btFT8tlLl2VaVtAxevDlGvVQTnr78zMiiWOQniAGGOB8j3OPC1oxhUPwAAAFdBnqdFESyffT2/dTSVQOypzfxQ2+DUil4GWJIgvIkBLArS+ydK0VasRPrGHtnNMEBi0Ye06NmDV2cc8yAn4FapQmg/tHUSVZv8nmxd+pKZ4n49OY6GrcAAAABBAZ7GdER/hRLK4pCufOnSwAxieT2ESQKanmN9zyXjppN4TTZfPjYr52MwkENEajixR/ruD0u5TZi8/TESJpQaWOcAAABSAZ7IakR/h7DkMUOCGBJh+lP5pZ7hsk9ED80Ikh1LLF5Y06UPw2TtMFuWrNVoJtNJFPcLihlZ9xJ1Td3KPSex82src3Nfa1Z7l65jvwwWZi0CQAAAAK5Bms1JqEFsmUwI3/wZX/CB9oOI26TZh6KulR6t3agG/iMaBA6B/6qMTvqGVg4k0YebbNb+LKTM12MrCqqoUI8hGiD2iXPrHlOmpgdC3Qb+yp7n1ju75DCZDB9Rh9M8MBS44veosH5GPzQfZ3g85NWC/wgf9Sj/wt2ZHm7uia3RMjx0cInK0e6gmf2UIdsg++PvXMcBn0gM8Se6JG8DKFx436aem9R9/OoLsInKDNEAAAA+QZ7rRRUsn39wfthx56FiN3xIMTuB500p7P9azlWaxHm/XjqkrOgES0XmZvka4Ctq1g1y9N9V22HlsJnZM+EAAABFAZ8KdER/h7D4ZJjIOE/G/Yh5t2iTFoHhadHs+YUVSRGmyvCLsMp2wNI81ubiOdKh6cKDbpl4cUw6CZHuVE7caJVrVL2BAAAANAGfDGpEf4k2H5XFXAtLnyqNr3TnYj9FIKhGPIf0deRjBKSCA4espm33EiqcHf4SHokHvLcAAAByQZsQSahBbJlMCN/74DjhBrqxUlYTYumv3ra8FUgL+AgXjYUzvRyU+e6HvKXgk8xL/pusoXU2ZW4/lqsHsMHeb8YKaReZyD2xnfEzDGr/au71JgRab6BdwH1e0PQ6ZYMc3yLQJm6faf8goYHDJ7uRU+gbAAAAUUGfLkUVLN+BznEOh5fTV8lzZdXNmbQtHmURs/S4z/04sX/c5yXQAjkgX0sC2qb0YByR+BiM7qb8PCLi4y3UBZ//3HtEURqKtuvuIfwOuQF+IAAAADgBn09qRH+BB5XFIVzyB9wJv/LfxiM8ksv3Sg0OqcE4Yrqo1JjCRCketlzK8503P3asAJA7OT5xlQAAAIZBm1RJqEFsmUwI3/vEgWa+ohqyK7peNimapDBP/rmBAmpQD6OSjgVVooLkr8s+Xn8exw7XhF9TA1dM1L0nWyyYDn0YYo5qCungDaPef0jmHbsa/NST2MzomuWP12pBzS/D8UW68cG/U7WtRkyLaQ2o6xwlezid/ALSuNnhZUQ8LKvoBIjMxAAAAHNBn3JFFSyffT3DpS3ZRuxv7Y2vybz+bP8qJsrnWB3uZuD6OPtDGdX3ghT2/RUlUSZnoKabDsiTKHthUxvZZN5kazH/ANj2RZ4dXRg3GoPqzpyXMF08NHSw9ukMXVXozck1ZVN27kSwitnpaQpcwBkrkivgAAAATgGfkXREf4egL3/U7PdtZa3GHl5V5zeF2LLWs0MbzXOZlb3kOt8MN0Y4o4Qo4DC8r0V9YQ7rqrIWW9ryH/Yg4OqxHG3QnithZLe466OdQQAAAEEBn5NqRH+Gdbgx/3Uy5jiTO2z6NGuYZhDosvmqOy+HJIrRZ/1j+mdlQeOSAXr9NMUAqXflcAvnLw0HhARRdO9cgQAAAGJBm5ZJqEFsmUwUTG/77lcOB3OO+4ld2XaGYPZe+y3Nf7PM2BZEj3/ak7IMgYJZ292AtkE/+fDvHwSjzWBw4ZWnozBz4y6QXsHgftapY5tl8Hx4uA9PzOWUKmjPAFe5E1HvcQAAAC4Bn7VqRH+BB5WDztAwdp20lm5s3YSJceebDQR3Or6z816nVxtIBRfgErPOT5y3AAAAe0Gbt0nhClJlMCN//XbgfUwOWc23U11yALjVVb2w/NZo+TO7IFE7jF6EPo4+YMYzRIKbgiu3F3kB85pl36VhUcvvytbvKOop5eSVkjG1CTrVgXP7WxMDh0fOf1+slEO+lIyatR6SsaZM3IzBZreQq2JodfFlKOfZwWWZLAAAAQJBm9tJ4Q6JlMCN/8FQss3+ZKWvz1Vvgwf6g9Ha3XKud4lLmF6JaAC5IPporCJUHxbKHU7vYH4HXDa3zFXkkAfldXiwDv9hY4EqQrh7laZJ6pLfk5IRDi10QhgrLa6I1cpiSQ45EGgE+Xxsrg4o4flhDnQ5Ihthw83MHZGB3B10kbD+Qt/BG0au6ZcfFXOkkOBPmx5s9/Pzof7lYV/0dpNlxpwp/tqKTFwyLX4rqRzgPi5GmFJ0mKT/M/Hz/TxmtqRDs4NYZoN35tdFAHpO0d0mGyRL2D3iMHWnkDF/LAzFtqzHaDqN77mbGjyug1vU8DKtwrcKfMEHvF/O9Vqueeym3YEAAABmQZ/5RRE8n9O9y35ttP5XA8yfAgYvyCUdxljqWjOdhas1mddNiDe/3VrNDcwN1etJ4bY18tX/hyWyP4GIwlHrtyez+ZCz5ghiUA/FV9anrOpGlw/LmUeoI015mUjorg3GPJNvCnPAAAAAPgGeGHREf4Bx3gMa3MtJzXKNV1X/TvXH0R7p0TV51BB27zemt1rU8h+OxY4zhjatnT2bUD2FU0s2nNDpkPAcAAAAOQGeGmpEf9wqXEWq5RiVMaSwAxlOO+23LmXEvXP/IGbVJ5FsJzif9EnJ3uZT0nLj1G4aGVRcuYlRjwAAALxBmh1JqEFomUwU8b9uJt+4JAB2PDGkUZKn1LVaQmhnA3PX4/yJdbhJ6+ZhTvgJKp1Mpd4XsRmIufzYB0Gjg5HjsymTtW09NUcla0xdquO7VsJlNt1J+fS07p/3CtM9jIS4PXn662Okc2oWKdm6qRq6L6au5t7Fmndyg+y3aiIWUUQe67w5+Tvad6VsoM2DsXseu9ACjX7KP9LZCac5WFt8E5lYKT+WyzrehW+bryfOIPcnHP5Q7WKqAMVWwAAAADQBnjxqRH+Yjrg02yKk3NJk1+zTlSEjhOv9WdIP223OJgbbUqmQtvMLft3pIcnpIH3fSrw9AAAAWkGaP0nhClJlMFLG//QsHcexAhSB/PMs/QCqkuzYg99CJH+9JGOv5aFabrl2r2ctLBY6M2x4Zbvd9hht9FyQ6133GQK/Ox+Ab7R/IiaxaiHq0mgOYL42zxXHdQAAAC4Bnl5qRH+HsRJVvmNK5tDyf6EJ/1EigZ555nVBvboO63s3YDIFnXPkddwMkLKQAAAAYkGaQUnhDomUwUTG//QSPy7vTRAErbpXCmQqCkFeba8BWdHebwLy/TbXEM1CAdbroo7rrsm2hvYzIsEAE/hNIfu+eyIxwlnCjcjAA2wohh/23XNHd2jo5ayVrDbcMuugD8lZAAAAHQGeYGpEf4eNe2TaTeL8zPwYEq4gvbPJGyoYjEgzAAAAeEGaZEnhDyZTAjf/++A44HeOPnbQKXKFpJcaRdAGqeg9I1987zdEm3R60S7BqCKKCkCiM2sDvOC+Lw0lYZ0ctcQkkNkRnYsQfrl2/a8fGMNco18ZFNi4twcv152XeDYkOPPOBVk/8oWKhB5mE61fKSUIalH3QzOSlQAAAEdBnoJFETzfg/vT37uAoiBAB62fY4T7UhTDnZgXNwAMuT9zHwCV7PJPIez8/S26iuBAOJekb2B0RKzV9AVe/JyD5c0nVj80UQAAACYBnqNqRH+BB5WG0Yk9tWoDpe8K6FsAkXKMufOCq5zTCtm/E5mXXQAAAJZBmqhJqEFomUwI38etDOA3Wak4jIfVLa19s6Y0Rw/msIXrJHoKvxhcLNOe8r7slL4Wa5KoaggYjByeswoCOZrcIbqHWjdx7a4x7OyGnR642XUuyo+PuqhJohJR0twDTn/jWlpUwQC7UFGO8pvAyih3nOBLZzUtnawIh01N4/zJmVIzs03ryQ/CsNNwGXnSogOB7nyq3zAAAABCQZ7GRREsn81cfmOuKtuOmYL5CdH7jVoF6ZsqPcI7do3JhGpCQH0VNgZvZikZzZ0PoPATUfW5vZTR6uB/pZC6qxM9AAAAMQGe5XREf57Mf84VVo08PAmIFtyz68/xZtlcGHAK7LcnvzsXqehJ/LT41oVmRW0o5mcAAAArAZ7nakR/hnW4Mf925RcTlmD+h0U3BboCgA52LwY4NlHSNS0Dhr7XMuQKgQAAAEpBmupJqEFsmUwUTG/KCEHK3fhHiuibLgfkSu7LpIIY3nGHB/7BUdXL8qXzP68vVCju+6mDLSp04UM3Yy5sMOop5Db5UDv6zCEmEAAAABcBnwlqRH+BB5WDztJXdtm1EV3ecxpV2QAAAGZBmwtJ4QpSZTAjf/vhUdhEluql+hMAizNOFJ5HtLg8fPT3ztEuLEr1XGD5xhtaNEj9qa8FFVnv7A21sXMqfiyoIXW19PEUwYPgxOYHv+WGMXik1z6DYg4KxCAi5t9mD/wLAUi9+oEAAABjQZssSeEOiZTAjf/74Djgegh1UZQX/WDsSXsFC7ltHlVSseL9p8ERavj1mdYkVohq6rg2aShXbm9fgcktbs7aTqfRR48AbzyAdovzp+bia6bsdbHEDLxLykgR6RTMRnNhvmvRAAAAV0GbTUnhDyZTAjf/+8PvG6OB3OPtuWzfxGhQy+SIsOLPOJtIbYUIIMS0V4qVXaqngqCi//dOlQY8F7ndPEOY7ybsB51X9jVAWXsMnnsbNepMd2SEvVMXwQAAAEtBm25J4Q8mUwI3//vEhm6OCJW+nXGeBIiklAPuU/0ZI1kWndzC0dHYpngvdDP+J144OxWFY4ccxDyLLIj6szhtQ7/kePrw8DROXeAAAABZQZuSSeEPJlMCN//7wcOsd4IFAwltfPFspkUwU+tL4SrplClPxXSDmSimJiQeXcXpRNlc4r6tw/bM9uY4UWJpP7B364iwyGkyu+lQkibJiggxD+c+9Bvsu3gAAABGQZ+wRRE8n309Vqjbpbriq8XvOdcqau8JX/6UTOUVpRY9WKhf83qV6gviCooy2oTUp7uQzw7CwPScfAIZVLOOefPzZ7OvoQAAADoBn890RH+HsQIbL4x/wcxICg+7jcceczhq0AWjcfMYx47NzTdIEGN7i0aSwUyq5LiHB0crT/GwU/SAAAAANwGf0WpEf4Z1uDH/duUXE5Zg4eut6v/g2rzmcUOWPfIVraaS0gPT3J2ns5evZ1CI1ydypxrgCnQAAAA4QZvTSahBaJlMCN/7w+8bo4Hc4/EYYvlEGmfvhZdsa87VJ6JltdXwa00H/ieFaI/Y7du5hEl+VYsAAAA8QZv0SeEKUmUwI3/74ByBNNkfUtOKnL3+oyVxxmsptyCPddx/B5KMUbKrDnuC2Vyx0n1eE4KlwYD5Vy7VAAAAZEGaFUnhDomUwI3//WN7G3P4NYxvBghU7kkcBVleyVOEARn3ED3ZepoZq5RhvWCC/yqm2nrsXEILV4frQDCdGWESAMuwcPI6wEXjGyfMZqSkgvdx7JIjxrwTiUqQdGnnY8X6eu8AAABoQZo4SeEPJlMCN//74Djgdzj87Yh9SyaadvlLvsYIaG789+bDzKp4nI7Ux9DSaBly7yJ7/BdbBlU59uboJHfQNPDqioaaRPekZqtjhMesy31GxBhA1OKS4Ps6VJGf+JLm3Aa3xC0tIAwAAAA5QZ5WRRE834P709++8xDSwYyJmXntNFlWK92Sf/fQDIC13sgr2eS1PqEUq8jnPixfgD+QQYI32zWlAAAANQGed2pEf4EHlcUhXPIH3Am/8sLtAuTjaqUxAE/k7/M900QRhPQgPEHPLe77P+AZmROrczDlAAAAbkGafEmoQWiZTAjf+/2woAB97w5AEvi75kBbZEZSk6i5LHQkcNyP+hSCGc7+2SDRFNE5x5z/5s+xx+xzk7p8O4w7+mh+W5do9ddvYMtDfUKDtyA+F85e07Sekz50aiZ+A2LJ2uLJoEQWDwl+okL4AAAAbUGemkURLJ99mioJbr7gM0pfNOj+bZFRc0eF3Hz740MUqo06KEIR1upjhqV/Gzz5amPTy3yzTR6/dNTQU+pkKnktDCZZSdmNwG8n5niOiMJYGfQ5IyXF8S4KMq7xjbHIWbNosUD2y1HHlqmypIAAAABLAZ65dER/inEMYh/V/63CvSgvZEXbfhyQx2wa0izGid6YTPbRTv0aQJw+H2w2yN2UnmHoAMN0UTFRR/aDBMgxksMmIswPXoXpwRBBAAAARwGeu2pEf4etAFVTbLs6LYf57fPpukYdfH0vO7n4NhsecsU4ANjZ8RyLB7ip5BIX0Mx0+7ZAp+laB0Jf+sCagXiPlxvuigHgAAAAnEGav0moQWyZTAjf/Xb7hr774ABr8i42lWBj0uR1Z9SPh+ggHjsNADRKfCe1SC2gCz1ZV2wfdi4gTAYitJI1nTCixMcZaP4yGm/n4CfqGc6sKSk0KSFZjxx71wspXvBNzxUtZdrOzeAHxomBMOPBhu1dzQNzjtmlggUsl7VLlU74bd3nV5k7rd2++U9+8ggj3YDa0RWl4HtEE/EBfwAAAC5Bnt1FFSzfgzrToKs1FKb08R33Z2J50qKbNlQV9+d+S4eldtRoddXTZwHI3pZ9AAAAJgGe/mpEf4jgOwq2BPvs/N5cLI/GjH7DHCev2QPFr9/a5wG4YVCQAAAAYUGa4EmoQWyZTAjfyghD0zehAPTP6iP8vI/K8ulMYZHKRgVkbItuqNJ3HEAokTCb6Pnt7yl0OLzZwaSoIONyTBCZHmwGWPv3H83nFo1lWoU+cguVFa9NsDp0gOBgSBwkMEEAAABNQZsBSeEKUmUwI3/7w+8bo4DvMfUVj6CQ3wk0ZhZRNy/Weym5JT2cFcNQkc517Tr/49Tt0A4xq3KD+uTRFAhPphx4G+TQtX/b6Q9KLoAAAABIQZsiSeEOiZTAjf/74ByBNa2ipZMyAB8rFbNzvqfypmxt6LH+u+4jWhgXuj+cO3YiKHJNJ1WQFBNgS8wGPIrNRLwTz121G6k9AAAAXkGbQ0nhDyZTAjf/+8Gbu08WB7kttaaGR98kbL2C0gBL0MGbGJD6y21jf9/9lA37OP49VLmwv/65IaqtSBP9i3lOAQ5XqdAIIYxDxwBb1v2HeMCQXhXOkOBAThSFTcEAAABkQZtkSeEPJlMCN//8AoIGyiRKThReSI30K30Lv43t8SJ+8KvIhuBxhxo/in9rIRgUOPXJvOP0aqCWU6/xuEs9JaJ8GMhw30DvHBYxT/U0Kj8r2ehB+UmTWWcR/WTYZVe6KZ5swAAAAFRBm4VJ4Q8mUwI3//wDwzHA7nHfcLHfflRrPXe74AVeC1IbYIGvxkydW+YZkCPntNjxM2qO/tU91qFpC4ImWGO5QaY+ISilW3oVGEXv9m5/sTC2qUAAAABQQZunSeEPJlMFETxv++A44IUcwuy4ITXX3nLrPA3Q//hXK5HZ7fz2fB0R5AIIhVciP1yw8k0h9Np/gyGu6xn4JoLPSopgQYiI15qUK7avBOMAAAA4AZ/GakR/gRIN2pOx7ohh37pmz7Klk5+7v+6dmyrL07Wcjao8Z7/mUQRNj6WbSSQGD97CVmEVJHQAAAAtQZvISeEPJlMCN//74ByBNNlS+GzqUsik57u8ykO74fuvsnTVRg/9Ptqbp+XBAAAAUEGb6UnhDyZTAjf/++FR2EKSC0P90owSMxGoCsS0jIjkJXAY9nBDNCe1Bk/slGqt0QC1TqM5EeT/mtDAwGnbXjDvOUDKJN+rzAy12AIq0PbWAAAAa0GaCknhDyZTAjf/+/6PPsocLdBktI9rDLAZxNlt9wqUjirtVNt+xqsPZxzK6F5vZZijVpT+u6KITqFhQSn9k+W8ILsVjI7zrB9WmSKDyuy8O409jwJKv6h/HdxfnOL2WIBLbqi82goKAz6BAAAATUGaK0nhDyZTAjf//APFuECSmz/LfleQbzi9RxtCMFA3aSBw7cfB5VdpUKlIOO9D5eaJ77NvYp1NRMnpOo8KCNP4Zz+5SwFgGPmDpyQ/AAAAOkGaTEnhDyZTAjf/++A44IldYymy5ae50jTmHGl+LJMPRkmIEvlkx6J3enH/Z1ZYnSaiWfS9Qd3bZWEAAABRQZpwSeEPJlMCN//7xIFmvqIeT19k4jZC5qgXt/92KDK0y5JmDGlXIo9RAWJ1sdV9OzVDj0ny3mjY57wjldIvC+C3BkGsD8nC+eIrXj01+wxBAAAAQUGejkURPJ+A6iW+6tW/s8hDZPRPj/ZS4R77YE3gu1EWsAgpWZFUkBfemZrEkJkVcTAOnvNlQOvZHJ6uIRbuaIJAAAAALwGerXREf4egL3/jdwY/RuCrxWt5oYNMoyXcgYWnulwzd87FVyGo2D8hOuzPX1rgAAAAMAGer2pEf4Z1uDH/deqwuVerjcUBHrhTnazCTYcm52J3Yzn0y9s0foT1etFDJ+eKgQAAAIdBmrNJqEFomUwI3/vua0IEqusHZHfP89a8nE1UwIsrlNvzfQ3IirE5WhcQKxUzthp7BUf5HkhK3zIaLvN0HYEJo/wGnymLnnzFyHPfS9TQcFpsTCptT+W+ftmIjvDK0LfLQWDU5ngMJHz/n4ZUiF0aa3S3RfkSsidDqKKs1LyafNz8fl1XPT4AAAAoQZ7RRREs34M606CdEYoROePxtpctRNT26oIcfF6VX0DqQPJvGMKc4AAAACUBnvJqRH+JNh+Vsm/raPip6RP9Og6Y/+RSUgkDHk9xc2O0jrwxAAAAnUGa90moQWyZTAjf+8Gbu2qQAjMVaMFUREfFFP+x4btt00ueekq5HGN89Wm19wbx3Kdw+o0PqxQpY0DcAXpv9K4n1LD86yV7ODCzYbxbeNispTBcvrKPCMwkXGwZIKam/tG7bdG9+lzB6J8XZvmU48gqeMwbRPYvIVNA81JXiwX5VfXEp+tOTePj9OoOWcrIF8LDW8sH5NteKrA41RcAAABMQZ8VRRUsn357kZSwzxr9U8YWP9/xl0oFI+aTRf6oWUqAXw298uTlr2goAldpwAKLQthKd2eV5xwZcFfuV6zka+DNcWzVwddT5dqn4QAAAEsBnzR0RH+HsPlbvvmSW/qDUBoMhf2usSxi8Wum4LDDW0Fu2KeEyWSegiZwWr0lVsQcu/F0Necrzgb7KR5QDSLx9yh9ScE0/rzHaJwAAAA7AZ82akR/hNfANpljEqY1fAK9wQesw2QAGiR2LAytbhpI31Aax1ACoVdfHniU9EpvumT6lsdtl9FTtjAAAACFQZs5SahBbJlMFExv/APDMcEKP+df8b2bcbvMcHhLvzdzipVTXYQpeF6xqpUhj4bNawZXtZlSHo26T+WH3zqaIllB37/NGNQ1VkY8BOUKE+zpXj8ANW28eEEEdJdDLN1PiFpb5PIMKKW9o9IlM9LDe1K216q6Y7zZmIMkCDaMmuxIIuprEQAAAEwBn1hqRH+HsRJVgFi+ZrlLDr4CkfmPNcriMt98U9nbw18dR45otYG8DOLKOeaxiW7U396pvPAr+t6Gn8NLGkSaktuhYMuySbpXKWCAAAAAp0GbXUnhClJlMCN/+8FGR0k7BwEOlhvznFDc8MKd6PlOwJwx+IkDWajbiHtKLtq12J317Zrrh1gtLttHBFn7OwW9kLbr59zfz9WSxFbcUiYfBU+4PJI46FozlfEF8jFGF4/hW1jigr3Oe+UyN6JQ7klFrYmpLuv1jvLAjYdokw6od71OZf5GD4K1jtQiVwCSoCZlTnckIlQyvMjv/ZinSucoeLKUdjlgAAAAREGfe0U0TJ97IKea95i5I4Oz7pxueJfZG4kOmA+QbzXKw7vkhvqrsiqjd/kL4XyY3lQf0TmOeb7Jxwp736HbF34x3jaBAAAANwGfmnREf4ew+GS+7XjuUTR+buTP+d/8H1HLLwuG6VVjsCWSzm1R7OCK0QXFWRrnznHKxhkYoPwAAAAuAZ+cakR/iOA7Criq74tIQStcVo0TUW87D3aTI7cUwCOn+Wxs+VAcJAD2/FWI1wAAAHBBm4BJqEFomUwI3/vgOOB3OPtuWWgGLWHjA2dK83VlKvt2+n8reBM2XzHD3zBxBIqIRM/rFRT7vg6GzXbZYVE4qWcVvmjn8+cOQkH82Z90ruTl7Fjpv4YW3KK5nKkBzEFtS3TO1V0jlDTMhHqvo4khAAAAVkGfvkURLN+D+9OfdQR6q/7IXDaBvaDpMhPEswOM1+7l4S2FsorT/pLAVOjCdAJLasClJUflp/50NXl8VQKz+m+7rpgfYkqz2G+vKYeS9jaTSkokmfCAAAAAKQGf32pEf4EHlcUhXPIuvJ32qeFdoZAknNUtwBScyRRJETefzNQclkJBAAAAcUGbxEmoQWyZTAjf+/2woAB9YANHst9xAn8Ot4dHzqUwafrp05d0vUqpjWj5to+oP2Bvpa+fUTyobSQwuvbyPiaW5y0mT1/O4V2M/c/cW/9kBex4FuohVoCvWoVQg/16KJ9s7jM9tnwTb/JeLWapEhPoAAAATkGf4kUVLJ99mioJWnVwU28WW1/7jNoqYi8ej+b1wieEOmo+V6ElGEDFkN0d6FE9/r6vZJXwBtcGjNFa7FbJe4U0C3VVH6u9ab3qqXMOQQAAADsBngF0RH+HsQIbL4x/uLfAQThciDnXlgeXNkbuTWOhTs0T8clBXVZVPYHhl/rxl3hE9I+Hrbh7KVpysQAAADEBngNqRH+Gdbgx/3Xq50pXw8mqoa5/S6osvnea87hsE9/0OGaERxf1tCTkMhfZ7ECEAAAAY0GaBkmoQWyZTBRMb8oIQcmggO5x33CtfQoNIePFJ40j0Rw5ce/FyFAT5ISrtO7pr4/ZWcfWhBygM64QbSUQzME+lubt8WOWt0M6f2A3BfB70bi4jXwkpKQm6O/cjjNNV79bYAAAAEYBniVqRH+BB5WDzvLCHMeSc6D+pWTKrSTRgYiHkmyaD8ady4xoFuAGAC4gZprGG/guu9ew6u+KEQtPsf7XkTxEL5oOibpxAAAAfUGaJ0nhClJlMCN/+8EZPlhBHWC1fwUqYD4gl8M84OnkMvjhAR1eA25Xfq3KZ7HDMrdVm+l24vSAkfG3ha5Og11+q/XKHpFWzSPOORK4/q+VjpIgK+wW3rohvnrLXSLrg00iFleonFUn2Zm8WiDb62AKQG8rPQI2jk4tFOzPAAAAd0GaS0nhDomUwI3/yghA8yNyXL4QryWtpO9TBWG+5lvHFT8rZrJwigk2mHXYghgKHt37c17joLyvASuy11M5xsDnIGswur45lcp/pg6KTMv64Bm5l7g8OAtRTsOC/qs/q8KCOxtRlop8JGxN+wVyJAecO+x+WH2JAAAATUGeaUURPJ99biHttP5WYp4RAgYvyI+p8tb+TRf6oWaAAXw4HkfPwn3YQah/Hn50Dc0k3Kbwtlf++vgNgAn442VMlPHd1szsO3j9hrRgAAAAQQGeiHREf4ew+Vu++ZJcAjosbnx/+SKnisgJBzGvhE2wYk44wff7Oj+WUOlh1tx4rNtkmx7x72zfz/ISHcplyHXRAAAAOAGeimpEf4TUK0MsYlTGr4BXuCEvYQ8yfXuACUj2t8k/d+aCycsFhAIUjm9ScUV3x2YTx3NuXIgZAAAAdkGajUmoQWiZTBTxv/vD2xXRwQo/5tqyIg0oegdJzRBsuIl+7b7ib975tSk5oO+RagBoP94bo8xTAUoC6UAV7hD2pl/WnNEYhjeekT1PkCFVNeqc3AZwsou8tNpbY/mzbsB2O5/LkQ+UJINnK+GSH8HOdUvNWTkAAABKAZ6sakR/h7ESVYBYvct3ZUH80aLSL3lILK48NjH93k+JR1JE74PB6N+HZWuHDRS04N7OCEWv96qoQw6Qs01sUJTT++6b2XHrGL0AAACqQZqxSeEKUmUwI3/7wetWHYdbgBdtDzUw0EHUbfim0S92FkIfRPj16pWKKLwRo3AU8FDsa9QId76SyIZbrKSW1k3GGqM46d3clPOrcjz5FiWx32mC+NDau6Yn5iLvV/ruKE+CYviocKuN1PAgXeEie0yP1W3y7ir7QTZDPhf3ZL2j64xZerSGI+I/vp60I4Yfhme/LvW+sG6aWg7iglDM6mb9UFBSV9g/opoAAABZQZ7PRTRMn37oQ9jXIx1lBX9mqfhAUp5p6J6k829lJWvVzf6M7ggijphbu0clU5M1vBwhyYZzt/rM9/6ylFgz0BPTGJMWpIznDxoA5gzNXt6b4LnZ35PBfoAAAABGAZ7udER/gHIYyAVJ11hb+xYFCFd58j6mfUxiy60D2axlSlNv4jbTnYFRvT61K/ZgtX21MlfVI4/OkXJuUpop7f64MnECrQAAACkBnvBqRH+I4DsKsOxDYDk/6PHIMsNIpxKsIM+ybhNlwmFnio9k9isHXAAAALZBmvVJqEFomUwI3/vBm7tPFg7we8ueSW3AsmgN1pq6VL5mONAom5p4474jhWe4CPB0/d75Ta+3AbyIfJutqpUJfL9iHeRoITsbO2cFPftnBnbG+U99/U08f+CG9I+wfsE58CY8Lf3D3w/MItP7ygZt3mLhmC8NDeSwaV013CnmXqHO8zdoI9keTJr4VJTM94XD+qoMrzWDNiUadSEvxRjUAUo8/J8dVa+AVsL8l9fEVbn49GcOQAAAAFNBnxNFESyffW4h7bT+VmzHLoEDF+QSPtyPFrba/50Rn+HndNiDb5+E/AuX0YogGD8wr8BoMCMYsRM5UoXrME0WskDjNcWx93MwmTeYntQpF1DOgQAAAEABnzJ0RH+HsPlaZzJzVMKscA/2L+9TyyQd/3O+nm90ZXl8Ksqn5Id4lYTPtFj0zEdxLJQd6obfw58aUu4rDzZRAAAANwGfNGpEf4TUKzbRXPm8OWGDdlZhb7MPjSB0cLQgO5R2NblH67E1i0ao4Dkk3g/masUQXUJUJOEAAABnQZs3SahBbJlMFExv/APDMcJc0yucaeHB6PvL/u479qPe7hcn8oUhAYtxr7RP+9d6EP7vADBZs2ne/afeqbqNMqONZ5d2u+aLATtq5FH7fBN2ApzqD2YHoBAbx7AylRx7VOtEtRa5YAAAADgBn1ZqRH+HsOQxQ4IYEJ9R380tftck9EBbFpYrpDKVIDssfLu6OZD0KIvfnYuj1dQTryD5+83wVwAAAMNBm1tJ4QpSZTAjf/vhUdhyzm1U2leLW9Qg70H6FjRAKOWGmq9x5x025+Fv3p3S/FQm5Tp0onIj4h8Zx4V1izZBK7CKWM3Sg4BQE8SIJ4khFfeHsmBs3BWM4yet7zvBdepBB46hzuc6TbCM87USRY/aAltVYBThJ2Dy1Kk+X7oELhrfv+rE7sa2+kv58rOX6d1pHhsn+G4eGGB5RP9NO1f9TZuBI2+uPWtPojfHGH2/8Aq9EEUM/xWKbPACAkCDEeGpfVcAAAA+QZ95RTRMn3sgp5r3mdEHB2fdKSdICJFhyUw9U6B6nMiDQFTH10Mf+Nz59OffzcVpUjcUpLvcP2XrpMgmGzwAAAA7AZ+YdER/h7D4ZJjHpgBMmQmFWTxuFtuKN+1nWmrvVP+cuwy6cfDaxn9EPm//fbSX/BC1BXlKAR4dXRgAAAAyAZ+aakR/iTYflcVbbqsjzADeWk/N4gygRKahtBY4GTk21u4qAH9//uy/xUzjKGN4gOEAAABkQZufSahBaJlMCN/7wZu7TxYO7yWtpO9TBWejmOuTienHcMDCHKUGeJgU8kvIN+p4Q3OK88HSnfDGE1MdkaQi+aoz5+X/+HFemTy4gYlIXyYWnSJDJo/cOHmT5qotrLEMdESHSAAAAE1Bn71FESyffW4h9qH5W+/h2dk4U8g5Pt7nJndunw3PPWIDrz9uaZyKwd46Sk4PBi2OZZ+81wbZ8oc5SG+A22lt4AZz/jYoLj5rNtjCgQAAADsBn9x0RH+AcOUFoNKCQYCGgc//6BZF/1e++Y5yM2PoTn2oKho+FzY+MYnRSCxwYlPnGESpnGT2xIEfQQAAADABn95qRH+E1Cs20Vz50jhSs7iL2EQrs9j3tB+E0OgBAtIl74/4sMyLjOTnhw5AY8AAAAEBQZvBSahBbJlMFExv08O2aXJzY9aWPq93Hlrmqv5B0QtXOYZZloiKjehd74nU1AKH52Z3+TFO92LMM98MEm40V7V+C1GenupfeqytRxcuUcO6u7TqCPTeIwqzTSEBSPq9aYeksIflaFksME00+qb4RPN3k9VCAuDXIDcWlIfhtRfV+FZNJvsVXv9aSTadgIZXTCel8E+da9dln1o4EHbsCWRxMU227KCu+aooSEKKHlx/66IpLeADW25hj/AgX0h2GHLGpU+s9OsBO2nNNGGqVU2s+i6OVvVpOmhdZkIXLFAuaV6rywH//SZMiw2+TxWyRZ7vy1v8JUsXVPcmUxHw0+EAAAA+AZ/gakR/18yldm4EIIwJMqwy/zSlRwwRYyU0uDbf23QW+RyODJVXr6oxUFynmeXdw9rHk4uHPi3FC8lMF5AAAACwQZvjSeEKUmUwUsb/9BI9qBbAgvoKrhEVvOFiUYViG/M5EJ8mdYUiEaj97nKw+nocaPSSxXVNZ1shjf+5GKqK3d7aQA54mLlWRII10yz7m8Bep30K/8/RHcJBI+3YKy6Dnp4kI0bb5XEQ/cK7bxUfpMy7cMapjRKOoRLMTOm5NDf8KQfAqQ5TeInNVpTmGv///kutmRd+MQ4u5qz5Vq8W/uVrDIc4XyeP9023KOhDiYMAAAA6AZ4CakR/mJmxpOr92pSK0PJ/oQn/bH9a0f7ds14XLC9W6fZ4EbVHvaoG3QUOZmMZfHun6TyTV7Kr/wAAAH9BmgVJ4Q6JlMFExv/74VHYcuN4xrugJWIKSCyEVY45OXNWVZpnwYhiQRi8ouWZeYuX9UrcKkoRdSfCZZoOmnAnOQp7HCaBaZsLqFrDoWpfb/ojyr+kPVjjyThmNhOz59P+kfdZc/N9o20arxf55nhcJHGZLjukhPirl/LdFsHrAAAAGQGeJGpEf4jgOwqw7EQkqKLsnOa2jmmZovYAAAChQZooSeEPJlMCN//Hr5DS9RRAa54Z7hoy0gQSdNxHb8oFhjXsrpRKZMbMTTdkBvu3TAncUA10EqZ7+4nEteFrv70v8BmFG1wgiHERfC31S38TWliBAZGrZPRYU78tMXq6ichNMahFbFt7DygIkxM0pvU/Vl+YLZqN/y1MDAS6SE19XdtA3PtOjnSNlsS0SqWCQzXThAQMqHaQZ2uGXkv+3sEAAAA7QZ5GRRE839BbCV26Byo0weWSA9YDX3+QMejMyEEDStqDfwdOOZPoewjLmsgoJTYBfGhx9VUyx/y7bNAAAAAgAZ5nakR/gQeVxSFc8gfcCb/ywu0C5M8eQRzzLW5tT1UAAACEQZpsSahBaJlMCN/7/bCgARAV0d4PVu2DIDo//3hRFBNpnloEynMkS6IgSJm98XQILX4peWuq1XgtmjZvcfov+so/e8U3LVg0FPoU/Mo4SQs54IeRwOVrQ4GMDhGOj3J+7LW8cm8BcglvzqyhewVZdBEH6P6sGuFKFltAUejrF5DmiW+AAAAAVEGeikURLJ99mioJb39gpt5ZOf/5t8yMdqhflAjF3VSF9LKRKHSP+6fBR0zNoOBXiwoFPSpm2LefrXJYQBXANQoIKzterZ//Rd7bAfQwQE2Okj3HIQAAAEIBnql0RH+HsQIbL4x/uLfAQTkb4yILaCZ3iPJCU9pVv5zI9a2kfiP3oWWAYeh919yuVlfR5mK81H/GkOepytZyXlMAAAAtAZ6rakR/hnW4Mf91Mix4JlquVA684PFShL+fqcs8dpqe6XtLIH5N1f8BA1aBAAAAOEGarUmoQWyZTAjf++5XDgdzjvuJXdlt/kTGF+RV/LdKMM3O2Izx8kplgk0r1rh9cxk0ue6JeItZAAAATEGaz0nhClJlMFFSxv/74VHYRK6xsfYetUOdE5zQVKQV9OZlZwDb2XeIsRYByM1q78Ie7kjTv+lVLWwES1iMaYWnsFBfK73puU6X9u4AAAAYAZ7uakR/iTYflbJxD5R0UYmRI1Tlb8yAAAAAWkGa8UnhDomUwUTG//vgOOB3OPztiP2UA+261TZeWO+pLr6FuxY1Gm4KJTeKb22KE1CtKd2PE/YlwvYroPfLhJg8ZrWdDVfKHDYJ8TeXliWDg+DDMHa6evgdfQAAADQBnxBqRH+BB400xy7oyzZN7DiQhYX6mHtDc+034OuRbz0duExV4R2w0AHsAUlmmJIw06XAAAAAVUGbEknhDyZTAjf/++A44Ilb6dcaU5H0bjlOnxLIo0SXrQTWgueivNvHPE/8vDiGPA/ersyEsLWjNH8iFcw0ZL+tkT9WxPFLyHI0aZt9JkWOOsnrayAAAABgQZs0SeEPJlMFETxv/AOiKyBDZVHR4Nb12r2vQbwl1ZUGpuF2dtyRYmAgNHFu9YI8DTzjaZZP2/AhZiSYSEOKyyEGzaeRtjr6LBzxI4L9G5pO+D4VbaY3srs2TWCx4mqXAAAANQGfU2pEf4pxDGIcj/67dNk+uMhrkyCin4ep171/SmuXlGTPbQXRJPBnppBZpXl2UXy4JS9RAAAAU0GbWEnhDyZTAjf//ALQu8EXYDiMMWrNiM6GWH7uYInAl2MtZ/kgekQpPORIiGhZ/nRumtGaN2e05a4CsUGjPa7Jt1HD/aPBVyqAiWxtSFal//ERAAAAOEGfdkURPJ99K3q8cj1odDChudhfcpeV0qEp95cV8xPVLYNNjT/96mXg6NNq46gQvUi2kw+4owVAAAAALwGflXREf4MOYBiW8DO/bTvzhcPmLoHmc9yW0C7QrF99GHjG17qqQEM6OaQ7qATwAAAANgGfl2pEf4EHlVcd04Is9ClW6FZCH4nYyUbEMAwSiUszmoA/R2IMPcZs3ej3LlzkBx5akMr+kQAAAJxBm5lJqEFomUwI3/1jexiPPuVF4N3+NUAUuIHEz2QO9Qjed+3c9KFSmMPrydJHG31a3UXHzy+f5ILebe40628nHuliEo6jFwOTxVv1UhoKiyPtASysBMxVAoFRxcO10B0l4c3iX8DbGNOa6io3rXjVXTO0z3QG2JpqNAWBXcHV1gHdDwEtkGhKhQAfl2adu3mdyJam0GGvKzmODyAAAABTQZu6SeEKUmUwI3/9JoP7PlJhERd+nNBdwdAGicV2sDdshxtss0WwtsFE9iV+OAFa0FSYHSTT81z38H7GxqDvlLWAdJmzfoGzta6+3HuOXIhSfDAAAACGQZveSeEOiZTAjf/9Y3sbYq0egyLMuoFUSS1zW3tAR1EhxIsXQLfJqvT5msgJUJ5lF6n+eX5oaHM3FvUq7s0ubwvX+FHL351BlWvMj+p0hEQccUBBuAKEAtp/FcKbFnScMfZEX2zhg5cgAhSWHPmnBc4H3/GBfKQU8pvzMeWNgvr9aUU5AVUAAAAzQZ/8RRE8n3wHSlhnjX6wuuQMkOxSt1zZ9hoUMwDDnzb488UvfimL2oJcGWgtG3RGIynAAAAAJAGeG3REf4BsKtkMVY91KCkU84i2Lj9E6hfd4mZLJmupc44LzQAAAEQBnh1qRH+HoC9/1Oz3bWZ3RPxWyj+fW0/FzMZb82m6SxXIuxzXGz2O9ueW8NrCaIuPQmb+HlSCvaDrTOCW1gf+t9oHuQAAAEhBmgJJqEFomUwI3/vrPjgIGEb1mbkKmaEb7B3LpV9EsRLeONLVcFxuuED+lT3/pAzfeZLEw2Kb38X2btrExLs1IQ4mp1/hywIAAAA8QZ4gRREsn309tLhK5Et0327KxGXKAtwcI4JfdoKiSSv6tpnwJ1fs0HcHDvesezzWNAerFq04zWyFWRqVAAAANgGeX3REf4MOYBiW8DPHRkut1B2MvM8gxWoMkGJG6fJNRKcDcXEbvKD2g6pI1wtMG2qZlsUlgAAAACgBnkFqRH+BB5WDztAx3ehsFB1Y2ENPbEBth3hGfB9PEXgdSLcgUejBAAAAbUGaQ0moQWyZTAjf/XbVbdFBvxTAPcrrcz6VZOCFULaB6H0iT9bpVOi048iO1wulW4BtFnQgnXqt/dwAcB0kgppQux9PGL3+oIrcHf/88oLvTn3UeHSiGxY0fUoo0QSzEULynDOOcZGpy8/0v00AAABzQZpkSeEKUmUwI3/KCEPTN6EA9ymLrly8O6B7t7VmUeTH/G0a7t2ifM3r1iMMfcBrASVNUPovKG8GYNYNCbz/ppRUblstrHKIbzRyZab8jak9fpDfmpy/m5pgm0ZFpIjmcCLOk2J7HyFBwjEKZneHITdafAAAAFpBmoVJ4Q6JlMCN/8qLkmeQgEPSPuqCmr5ArtbnLRl4fTGGkErnT9fA9ShDSntjr7GEjAb8W3TZIJnx5IZJCQVx1Awj6XTCsn7hJoDgdkQML672DCqIiq+WioAAAABHQZqmSeEPJlMCN//74ByArVvp1EDVN71GG530l/KjHXzb02zLOHiQXYfX9VEBccGpAGmsDLsvD7H889ZLz9QGv3FQvckk+bkAAABkQZrHSeEPJlMCN//7wZwidiwPbbYagjwfDcPSQg33o5EFwjljEQKjlBmdpaMBOqA7OQf5sN+AJ3xRw8XQzvvEfWu7eb+JbmKBNgl+YwIiKZxi6zn0Whj9DxvtFKFYv/L4uXFg8AAAAGNBmuhJ4Q8mUwI3//wCgK7wSDWMAKdUGMUb645fTdwYnhpkNPya+L8DWZnT0HAzto+auQYGWoyZsX3gYPKgXfkFrIfwZz5ku2kqON3Et4mdMuvJcP/wIlB3jHcb0QXsltP714EAAABiQZsJSeEPJlMCN//77lcOA7zH4jDFrjpxOzOGUd5ZawgfJTi+lCyrDkvXFWhsE8XF0nUH+68r3FUz/xnMGyUBef43CwO4Fu7mkzlFHt7U2RY0wKM3AWV7ncpWtZdHMWbpKuAAAAByQZsqSeEPJlMCN//77lcOA9yrrlSWZGlkSvhsHkUA+TYWmiPcd9erKEll/VhIYwxpcKy2emhfmAiAQ3BXo41lpahCQkJa5wBO43r2oRyWCTe4Hr4f5iHFvgO/n9/U/Wv3cJr0kF/5/7WMzkvgsE0ESqtXAAAARkGbS0nhDyZTAjf/++s+OCFI9lK+huo+FMo46iSF2IPShKnDhY+b4Nyyfn422/qBUX1fTLaA9u3h/ChZF+GVwRdZ2WRkBeEAAAA6QZtsSeEPJlMCN//74DjgQ7ZHzgzntGfTAsTgNjJN+UQhfBz3E1P/plckBND9E86yk083GNZZLnqs4QAAAFxBm41J4Q8mUwI3//vhUdhCXYakQCl1mC9vjlAE6W6K+YGIyz//JA3BetDKXvDl3VvDxmOvEBMODpgm+jCKU2Gl8xC9oBzxxspKYcHrqYH05l+P8v6Y5Qk/1pPAwQAAAE9Bm65J4Q8mUwI3//v+jfjges12kf9dZQFJRGONVgWYOoxzG5uK3Jxnwr2M+6smx4PBb//XwNxXDNlrK0yKc8dfQ3w+72OG2IyxsTTpWG8fAAAAQUGbz0nhDyZTAjf//APDMcDucfnbEfvHS6zrNEJPWVHkE15RWPV3XeHelWYNVUgciz2fa0ECdV6bGtEIrhn9uo1vAAAAP0Gb8EnhDyZTAjf/yghC9V+DzI7WDVxvgaNgcVv0j50t1Fi5NV8PQP4/2jlF1RxSfv6bfmVSxmd7SAf+mKajzQAAAFhBmhFJ4Q8mUwI3//vBm7tPFg9mlCdoqUDYZ8xiLA/58i+MrJi/Oj9iY90OsO12pTviJqIUVx7vxap1f+uydZk4zez/EtYp/4fYp4Q3XO1R4juDaEeRsAQQAAAAUkGaMknhDyZTAjf/++uE4Ai5s7YnlAoVYdUxHA6lgf+GQxBdkuL/GOcCv2WRMFlJ071NInuH+zP8ObDVzOx5D/yBFdXBjGv5oMdw9Mw36f6YbxAAAABYQZpWSeEPJlMCN//KCEPdL34Qpx+Iwxas1iPqZ7d7C263bvdkqW2ukMa2nT+82TtnjrJ/eEhfgMZHvbgwHN3G1mSeyf5ERFLmZN+rihWcq91jTRoFRIgx2QAAAEpBnnRFETyffZoqE5BKXkkXp2UOeJEo59JeY9aj9AeTC6IYWjavARsEhBIEdBw2Rve/kWbUHHscnz8S/B1NjHiIhOhIZBxTgdtyQQAAACcBnpN0RH+DDmAYl6IayBmmZNEDqZSB4/1mlPS3tDlsV/Fh1LZ9b8EAAAAoAZ6VakR/gQeVg87QeYm3bURZpmDefiWKuuiEiUDlMyxeXtBIWKCTYAAAAHRBmplJqEFomUwI3/vhUdicsgfnibHpetEFWH90rguJ3ElBMQTQH7c3ygYCf/Drrg/dtGQkavBREa3bvVMAltQMkijFYWTfqWL//8Egk401tHl0A5hfv/3NvhBuZ/FxLMknFPgLLzedIjC8QeurMrgT7aB7fgAAAC1BnrdFESzfgfhWj8hRbl2W307v17sFwmYW+jS8UpzFjs/eVgK5cyRdyT/WT7kAAABGAZ7YakR/gQ9o5bvE/zRjV0ev/5aKcO19Luh0B14WQ6ZPUsDRRYb+qb/GrGriVtRBld0p3TqxJM3rCnvICZvAAlZ5yfOX4AAAAIZBmt1JqEFsmUwI3/vrP1EjJ9ApoRLZwxrgZM9KFuXGgqknu1QvZh+YXzmCVkYA2kXu5XfTPYyy8+Cwf9E3yexwGKwVVsHQkC08DCU5fhqZfK1z8qrqNUDzTK72NXRlKxrFR1ZmiVPj2Xq4At3GtTri/A05eWs/ccw0A8+MLhzoW7U93GW4fAAAAFZBnvtFFSyffT2/dTSVQUKj9rf+0mUMU9bY6PcvN6uG7mXsKRuTE+sYe2c0wQGLRh7SjC+dfY/ykuTFxQBC9+JuGvBsR4DhRM1QspMMQ0PRXYOYT3CwOQAAAEABnxp0RH+FEsrDaMSpa8wIWWIvYRFuurkj19cl46aZDCbA0SBpIQmv8EkNyrj0XwVlw8L1phD0sXTHNQHFOHuAAAAAUgGfHGpEf4ew5DFDghAjpjpN8kQrknsoqluAN7r/wdMHwzFeUfhsnaYLctWarPhpJ1pPTcHcQt1NddNzpv6yDKOANWXO32tWe5euY78MFmYtAkEAAAC/QZsBSahBbJlMCN/8GV/wgjheSMGlboEggPQ8iUHEclFSu/AtGO3I3Vk2ER/4BHv8Tbue6BAabqGiQdgTzc4oRsQ1gh9ZCu/VbhiFCPYnJiHdwoX4YU+hRuenZpNtsObSdPTabZeCoJiif6l0xR2ZC1r2W4UwfkY/Kbu9X8RRqWwjiZ+lCEe2yO97qkKitGLRXhsDDPy6A42LbEzFS5m3Bfj9FzHAZ9IDPEUXHyZS+YwQIKAdldUuWh2n8/PuDNEAAABCQZ8/RRUsn3sgp5r3Cm/SDs+6O3nB5JhyYbLrIl3aW8UcgJDENn36FSTX+WiqBPcaJGrhj1oT+XMgOFmQ+3rs4Yq4AAAAQwGfXnREf4ew+GSYx8peZo32mD0ZxX3r90gQmxgMzNTtzm1R7ItzavQ1ubiOdKh6cKDbpl9nH+Ni2s3puThONblt6/kAAAA0AZ9AakR/iOA7CrirIwsXZRWN3XV5/opBUIx5D+UhAdCUpwQTa9Dm77iIFrwPR5D0Wos1uAAAAHZBm0RJqEFsmUwI3/vhUdhEC5b/sF1SYFY16qcgc4pX3SjdOcxNGiWTn6CYPHzCyrtX74+6hxGZS/gSlPlQe04vlBzgd9y4b9wSmNmTbl7qHsmH/8WHyl3DolsoVYqYdHMsqsE7ZL7tq75kchJTZcF0jg7A6sSBAAAATkGfYkUVLN99TaxVytK9J0LEC7e9QMMUfwZ+bYFBPy6qi3f3mYx5+pfBsels2wDlgtcp/zu436EhZrd06L3ibH1Yqz/2X7weI/SnSNFnaQAAADsBn4NqRH+BB5XFIVzyCm2kI+3PwfLld5Q/dKDRNU+AxchZjpfprwo6gFhSOxp5leYMuT7tEBqg+rCoyAAAAIlBm4hJqEFsmUwI3/vBjyyMgPRE6nSRtV0KpMIijNNSxNV6fK186yR1lZCuSqeK0zpSMgyYXrqP8vepgaucl/4jaGmvVdxyh0e9LI0kKda6Fdc80iHHn7M/V+CUeawOHDK09GYOfGXSC9g8VnDTn54KoCAuYx0MXXmNYVT7fUCj+DImYMTgYnhipAAAAGZBn6ZFFSyffZoqCW6+1T8IDFjZ/2UM4MdbOfM6cSgf0UAPojz7pI5kEFknrYw+ENpQAYUTYtr4SqEWUn7junbR8Wt56p5h7uhIcGkZL92u9QC38gUwBk1Ex5E22TDcZLcBpyUi+s0AAABQAZ/FdER/h7ECGy+MeBKdNfVxUPeZpFC63XlrzdjoJZreQchSff/+dpqAridjEZDcjRKKurU0q+oPY3sW4PFZGDT1eSQOPqlSLhDyd9rNotIAAAA+AZ/HakR/hnW4Mf925ITm3O+8SFKZ3zxUoRkrtFQyng+VH/hMQ/Sik3etRXXXNlxlZUXH5+UDoy4GCxn/BYEAAABpQZvKSahBbJlMFExv++5XDgdzjvuFjuy7jd6PFJ41DAH3l/X1jKdL2Nh/X8fAtcq/fEL8Eo81gcOGVp6Mwc+MukF7B4jwey/r9CcqRx65yLF3b31zSZv614yBGc4M9cNdQpEICcQUO1uAAAAALAGf6WpEf4EHlYPO0HmACcpM3NnBfztXc3YkJpZoDQxZYYDyvNb0Pj+EZFz7AAAAhkGb60nhClJlMCN//XbV/0CALgiH+1UnutP56iKl/F3VC1DeTp73t9y15JQ8NdWsfi5ioASn3Q9MNo/IlUuX52yG0tN0TwNd1c7MGXk7xvpMQBKyiLGoACsHChpPyimGD1eDR61bb/pFXOSc8YOX5I6sYEzmjHHEsZBjKDlilzAK/wULUzcvAAABQkGaD0nhDomUwI3/03wWPufQrGYALLUwsJoxEMXwn8/5NCn2gZUgE1FzDh4+XyZdmkDpzw9gAUC6kukl2lQx24vXEN5j5DOH1GIYPDhINtBDfyJ8WfNT/Zh9gvQkc/PgO/6SLQMtIl38tX4A4gmbreT1L7HRARz72Z1XC3C1MiMvoDJ7VVn05DVGD32R54DC+p00fXgSF34RaX9p/nhPqfqWedglnMwJtxRR8T655s+4GIfQryP4RVGWXnkJEq22REKoBSZ6ysplrhgLEE+6q2u3y8N7Ts9Gp2LifzfiCd1+mPeg3wsFSzJW1W4R7PzKVLjmOSAYJ0Z1tw5ZriRBPXA+vDYfsvkJniSnJ44XojzZkBW5ylkDkLDJmRGK798IPsVpoZydkPXtIrYMAjNlGhRi/1/7d/7lypvEWzYzVcs56JUAAABSQZ4tRRE8n9O9y35ttP5XA8tj1j/f8ilMGgI8Y0X+r/O18i9VDvuDif6vnGMNy/wHlXR9Ruw6ja3/zSdWeANSXvyVEX9mq8ty/YEARNUYIs31GQAAADkBnkx0RH+AcOZcuyaLd28wzeOfzpKa+1FVhZF6ix7MqNSnD1fychy8xNmjL0jFeLjxZFPXR4lL64AAAAA5AZ5OakR/18Wdv2q5RiSBDhQMM5SV2fbqkEFaulvczIGbVJ49MJ0WW4sHI4AZItwvvenMKbRODBhSAAAAq0GaUUmoQWiZTBTxv8IgPkZ4dbAn/C5fKq3XjZar6aAnr/IwTI6NfhYxOOdY90vHFZcpPrVgtYQEGJcIs9138OsOgln47ETCAT3wsxykodHaHt5j5zd3pHkivyreE3pXWJ39gwvIPH7BtJC2ujNyuwcpegUpqIHioXE7MREQ7wqg3InLIx6/+VYvQ5iyZJg9fQEMDyurdPwUG97r3XLTUK5xQEoR2kd86Au5sQAAAC8BnnBqRH+YjqwgmYeCDpFqcX+FVib2Lrr/3UMPLGmd1jhtC6ur496SG8bD43VmgAAAAFxBmnNJ4QpSZTBSxv/74Djga4r9eoJPe8V1sdAFa0CHxG4haVB5Ovh8u8POW4TcTMxsptzm5D7YlBYw8lG6CyBBn/nDvcnVGWnkGJ6RrJu5JKQCK1vTDPu2rBdI6QAAADABnpJqRH+HsRJVgDVevwny/m2ehVgkT1H0/+i6OSAZ+6Mbwnkk7i2l6zyDhghf5iUAAABdQZqVSeEOiZTBRMb/+8HrVh2HfgnJcCBHcdoH/EHq//8+L7kvHr9AwlItE+XW8GuuTbZfIyKey+xtzx7/z7v0mSlg6JyClysHvmOYE28a/Hm+/EFhTDgOJfpyQ5cVAAAAHQGetGpEf4jgOwqweXy2F7Du2qK5I8q22RN0g1ThAAAAbkGauEnhDyZTAjf/++A44Hc4+26pcXq1U9uXkaTwRX1v/02a9jGRPrQ2GgD4ByCHatmQBBgoB4cOdikx24qn7FSN+jELlbhvLO0Fzbb5WsWgoL++Qm4ZaBej9aGyTmk4go9Jd/7p1u9ZDw+8QmFQAAAAPUGe1kURPN99TanKxqV6ToWIFi47JibnpS3MA/q3/m8HYEriYfagkjp32QrwpuOP5xPUmZSLfxN/vcgqC/gAAAAjAZ73akR/gQeVxSFc8gpXH9+3PwF5PUSrMv/xOAUOg3reTKEAAADGQZr8SahBaJlMCN/PivaVsDLRwju4FJRtGyJld3PRfB6Xtp+0EP2ehbnkk6yo1l9S2z7XQSpm8aK8H35oc1ChfOjV/41daHgy3+x8N9GxWBhXCKBEFAFwfSFUpL6aLnnZF/WsKnCiSNGkpuCtg8KGhAifBX2jdPWMZybTgbq9FQPBnlegtHQ1gMS0JPKw8fpMr/l4VaXgPN/WiXupgFWBTWMufBFXx/pwY4hzINiNuuh8vDdV6dZqzU6E2zTAetTgZI718atmAAAAV0GfGkURLJ/NW8aI33FXDdlOIL56x/LE61rb/LEC+So3AiQ/8F2RvGBS9ER8hLuHAUZ1asXN72jsM5Y8Rsb/8m/wSfi863R2YTaypcHVPZVzRNGbKL9hgAAAADYBnzl0RH+ezH/OFUf3+Y9OhGlOXkyMreXizJcXSVs+sr5MlHfOxYn+EpqyS+pzJWKhTWIZJCkAAAArAZ87akR/hnW4Mf77wcPYM8+uXxARwgNs3mq4c8S0xyjev6Bw96W1w3E5wAAAAENBmz5JqEFsmUwUTG/8A8MxwHeY77MPL7mjRFVFR0Z/iYOC4Yp43LzKh/KoiomQ7CIxYiuEq/otbdQROB4TPZzrum1PAAAAFwGfXWpEf4EHlYPO0DB1ukwiu5JQBDnZAAAAjkGbX0nhClJlMCN/++FR2DWFm+9o4wWclxR2utVG2+azrP+x4C74zgq9slOh87PLWTihGxNCkPERnZsOwKKHGBDItM6n+UO9MhbU4gRdPCbDzqp0LYuZWoHRtYrRrXBf15GZzNVJg8jQkdGrKOzzbeXrr8TrvjpovWonk3/zv//ORDTwucm9tODFLZPZUK8AAABmQZtgSeEOiZTAjf/7/o39RD7fvlHF/LNkISXGH4z56Ai/14Fui1YLh2MuTPImlrfB6tXNm3i2ltMHwB0R0XoODBGGdW+0dQtq/9aAzDsN+WbhUEnMSSne3XEH5LffBS7TGVCwM92/AAAAXEGbgUnhDyZTAjf/yghD3S9+EKcj7DVPmcgdNOSh7FRWYDdiiT8S5ms2TgSRC1HdbXoPao/j1pyyvcmUyyOr7D8VfrorJ9Y3Ia+/Vy5I8VajatIrRzOFE3LjoTFuAAAAWkGboknhDyZTAjf/++A44Ilb6dVZAj6+5eBpPqpy0hGCly7Fd7tZOR8lBt+xP1Bp4udbg/7G0dq2F/ig4mG065Jsb5K+VasAFZ0NwTnEg3ZzHdvF7Eg7qw9SgQAAAFNBm8ZJ4Q8mUwI3//vBw6x3g997CjDRenHxJ0XtD3+ZsaWlrDYyGJRc1JoEwL+kntl1Dn61wxnH31YZTaUVFntEYkpX3nfnQ65nvul3NPCVGfp4PQAAAElBn+RFETyffV6IS33RIzDNhC6Rz4N/7Il8kZTcQoK/cQdyXIGV2ki+axJCPZb9IDd9bBgUYo7Q242WZg7cP39Ppn+R9qeii6RAAAAAMwGeA3REf4exAhsvjH6163w40hSkl0dhxqvJ28aur4o8Qw8s0CKm6YtPcSeQHz3E9C8uvAAAACgBngVqRH+HrQBVU2djmFhWg1aZbX8552eSKPQTInN9Jr83hsggwky5AAAAXkGaB0moQWiZTAjf/APDMcDucfhpJoEdgP1f82nspVGylrZAY5TqKCcynw/jcC+Jqlhay5xuF9h/wgvXJVaziw53TDP2hodN6QtZC6r5EmXQMakhypMctn19xQ16J6AAAABAQZooSeEKUmUwI3/74DjgiVfUWivtLYxGQT5W2ZTVlespsDp7BuKK9RAPQwV4UZkOEH/M1tUBScFtvFQZbdu4IQAAAGxBmklJ4Q6JlMCN//vhUdh6HHGNd0B/p9FfdgmAb2hpUl+urE83Dowm8tGYb3I7ijBirHv5xiW0qPCI8pYvAfgKKPOp/lNBaxO92h5BlDhH+jDsdWP13tgaQXw/5AYaS+Jl1e5qLi0meHE4ZMAAAABgQZpsSeEPJlMCN//74Djgdzke7Ewm/iNKmZv+eBdZSsBm3kdy+Jxax1DOGv4ctvMEmTptqD5q7dbTzcZOZfcFWec54awcbIvWYJNuhCIrSQwganEwlNR93+G0LUHO/L7HAAAALUGeikURPN99Tam3Cx80XWuScbrI8ef27x4OlvfwhN8a0R2GOpfaziVyuEYanQAAADEBnqtqRH+BB5WG0Yk+JGGXBmlhdn2n24dWagV/yd/lHGugo5EAQHiE5RPtDB43FuuRAAAAhEGasEmoQWiZTAjf+/2woACNjELAEvi75dnr9c1/IljWX13Dpl+arXwZ48mEtEWdudCKaJzjzn/zZ9jj9jnJ3T4dxh+JDYdtPLY9alEq/yvIYKIf6nHoWUKR4HUAVgyKfewwUN6TBSDcysMFrfJrLB3fH67CeBSbapkCkw0R7IsA+zrYIQAAAFhBns5FESyffZoqCW9/YL3/eJSR//l8bcyjUNJcqqPIQpL4s4lCEI7ELUU1x73ujzgFVMGgdWaoy6vKtupvoLwBbryeu7ncAuloej1ir6dhXCn/9p5M5mGAAAAAQgGe7XREf4exAhsvjH61lOxDQpSTOE5Rr6c/YMJcR40gl5ozXSNr5K0gN30y3xwpq1Doi+RBmswRSVIyyiZDT4vuRAAAAEIBnu9qRH+GdcG6fA0ab8kNOTv52rX70y5CdadgShyyECf+CNopeKFkbu/KLZSfnj1sdOK6r0ii9DLbnwNgGTiqNCEAAACbQZrzSahBbJlMCN/KCEDO8x7vwT26lNl/gUfBmtbre0IKq6E/qqdCqRddo6tWS7hBeYS8byVLG6s5bUk2/ZsQLzfaGjnk5bK2DJGat0L5NBmmSHVq8s6V5vFgGpYgYBplI9JWytw3yEcIR9mMup/7M/PMBbSjjYsrk7wqyK/gnjeQOMfm8zcWP5JRN3q8ihlin1D6q61GFw1zLmAAAAAtQZ8RRRUs34M606CqUVFCJWU4Ky0MnnSmBsKVBYKXtwQD9bO1GgwnOtqMNh0gAAAALgGfMmpEf4TUKzWWd2fqIhqL177tilPXm5J8DJFaCASw/ejbPc30eus0z8KtHbcAAABNQZs0SahBbJlMCN/8Aqdq3gjhBnFy5XA/L+1MyAlYxhiqCHpYw+5gq9qVl/jSDu/Zxz8k+rcaAMAWshx2Lt6lAsCXP1ggHOhg0h9HO+EAAABOQZtVSeEKUmUwI3/74D1EHSkI9FUIquTDM3rKaNWSFy696/Qi+BwnpHE+ukHESGyxNmYBi2gCDySgK7ykEfDx1eZ83MsyM6XidSYhkeMVAAAAR0GbdknhDomUwI3/++A44DhT2ug1cazm1y9Eq/w51RkAPdbegFjO86K9lsQdW+KZ8FITyXlyVrMcZ2dEB+mty+r1aZPOgTfAAAAAXkGbl0nhDyZTAjf/+8Gbu08WB7e+xX/RKyCUlwLsCM947CKmmivCIJDojdrhemKwj68SRMATMf/1ZjV0RX9q+NeV7iUhcKy7URTsFOrM01gL2NmAF0DuInJP62fDE0AAAABlQZu4SeEPJlMCN//8AoCu8Ic2OSncAbLPwyKakd8JvTcutKPXYQaN0OqVJNStMu+lFOFblgBZyntC3N9Z9aj1T/qbHte4IkOIPGS4G7XHAIEZZ+vYQDNER4xNgppDUTCuIneQCJ8AAABbQZvZSeEPJlMCN//KCED59yngQHs2OJROVEKCAWdcGFHRSwN1YWw4Ja+vCA1DR/wWuwzoDuSzJ8m8/s8a8vs6MregbVfZO4g5Nu/OfbvwQXQzbNm531lzo6++2AAABOdliIQAJ/++DLeYgN9XYj5zg5B5QJuOQipq5DurWEjnW3UorRPL+Gmp7W/eIApq4hA6GO7MXHIyDHsVucz/aMFYtKTDM8z25ZLQfiNZJqSiDN9723ESzhz79Cwt5M+8vOWuuGhG+Tvq9foXjrYoVXiVDwEsVfUmlk5TXbEqyYI+MZpo4mDJdCI6FifaFbOBypE2MZf5gDnzZ5v8KmDeJZU85nFu03pUJirEg0Ch4LatUJlnXu1uVKmUsxM0BfauVLDb5iKlbV8Rxe706n9ByCWTYcNmUsWgK2BOx01vNUSepY8KQudanHo3PJDxToO0OLrabPAghuBeO7puO+FrVqFKMISkHJjS3EmTdBG02wgZepLedOuw6T+jsJltu6ElokeeisTTrzRHarkaC1NCRtCBiNuKEigoGG5TW/NYmxDmIUwJrPGuuEwRd1KazxAA2rsqyTJqLM0sjQZEKi/uESROrh+l53XQlmI5Fr2TToEa1Mz9Q0myFVMzlLsQqx3/ZRQx/eaJZU1/3iRY67X7pQgN8jPG7Mu6N8yvaY+Aqj7NpykjyTLzJHyZJOS8uJC5eDrvviZZfFd8hAX6jnJ0ZJE1ceTBCU68b2mfGkwl9oA3rNyC3kk0Uu8eOhHrIGj+7a64L8d3FiwoDY47nxsn1Y9aseVrNpfuZL8I3UMDdSnrXcWoYyESX7OVr2O8rZxn9LDxM6fe5FRhF4y7x35eGeAWARki3ksb7bLxg8DWvLG/+RwVPcpcizEE8l3Xt5+avL+ZiAXLEnwFkaUZDgHR/YEmvUJL0akz3W+4HoMrkbavzzs+VqiPswLqfDbnOg90ye7LST+LLq83FG1+6MSJeC7uys0Em20StpYNB17y0kIWJY4WoqsWsZ431/5zMrMnsO/w1fUI2ocSHX8pNHwFp4puUb+kPcS/Qh09xlBE4IT8uQiu8T2UDKEJf9iCtfqONZHMdPq/oob+ISU7mDp/lbACtBQUxbSPdO3HSaLklWzLgg0sclDx8EWk7I6bx2ZO5UAMGXVQtj+CCaYjfoBtDkXhWf7nWsT0hHJZrDi2Lra0z6P95hKbBHipvDDcqX6r8JNPqXyyckj8WVa7ygUVg99xXIM3DG5RRWnB41GtWfCqWChNibXPGQ6uh1H6EF8dmXH8surDi51hLT19j9myqwkNKmT8sfWOXCSiZtHYNl9wayNXOnug23v1/WAGCRrLJeHUUtVHTIX5R9cvqnbFcs3PTXx9R74CfvppdnOxRFNeTTBjRlRfQ+cjnV6LGJ+ftcX+9UqNBGia2nKp1uKP4ccPHEr/CsMJq4q8HWHOTdHO3jV+GZWnEoLlQ6Chy7Hbj//eKkHIOL4i+yqkzkPRSbB7sfUP9vHZ0uCQVKXYKu//r0MgUKN8P5NWtC5FEPRp8HfsROGFe/2knJyGh8s53uvSdYVz93B3c9Z2acgXQNzwm2yabDfBYT4AbEZLwPZK7m0eC5LKEHAxOLjG9NU26PPjXEMfajfx/8iXqLqjmgozrPx8k41W10JV5nX0OJWGa0qBU80PviynD2x3OkSEckOt3iTHjB9RR5c5nsqTL0EgDg1YWzkmuLeJTeKxOy3pZHTLdnNqj7pt+W/2OWi1UaTdjKxFLON8eBwDadGOUrbe2ZBGiOQ8/zfFIWmtk1vwMVGZD/Iod8uAAAAATkGaIWxG//vEnVTkBINY1xBnYqV2GYveJxG6pO1ytvZdPbEXt5LeT2BV/GdImH/VH7OxTyATpFD6nkhtuovj/nla9ShYi1okzGdxi99RJwAAAC5BmkI8IZMphG/74DjgQ7aLylrdTsuPGlvjUJs/BGQzvh+6+xQ9wiS1hY7U4V8YAAAAhEGaY0nhDyZTAjf//Zwh2MSAtcqcH6DBw+kiRJyBfSLlTlMVKsKbIZvDezqc3T5Wd3NGZa+B+Pjdq0z2oeP3B1ynbNRtvd91FAPKHi9eAZQhXkzjRR1YSlxJOK3yfwo9klCCVTafgfuL5sDFzP0e62FQJQ60CKx6fANRbJhNojOKZPzzwQAAAIlBmoRJ4Q8mUwI3//1jexb3GfDeuG3QIp0iOWNrPYaRICXk20KwX4AJSK3G01g3yC/ow3Dys3hXKUnzco3wdRZFXg3LRNpOH7wjxDdO6SGSNKnYi89v4MEaiVQLPE3ixfnhiAfLznBdpEhe+byYE2PLypisngAgqtZKdabisoZIYcgmWrafiKuWvwAAAFFBmqVJ4Q8mUwI3//v9sKAA/rehi7vgA080JM8SmYpPIuCaWfrNkOSrviyYia4nOycnCZbBmW3b9POpyhuj0YMiNpTWFdH6ofSMf1BO6NsLk44AAAAvQZrGSeEPJlMCN//74ByCXW+nXGeBo12ucPu1Q3jIarCClBC9oB9Vn8fnB9XMTBkAAABQQZrqSeEPJlMCN//7w3wmMjIEMD8lMTeUbP8EWg5RZj/kvdI2pZcX8B1kNiYTAl2O0lslP+qTbr+aNjnvHbN84RS2/4Py+FdljGGrZaq0CiEAAABBQZ8IRRE8n4DqJb7q1b+zyENk9E+P9lLAQ/bAnTrrSAta8TkZGO0FbRoL9hiQMM6FUzWT+8DFzTHADnP9RRGF2gkAAAAwAZ8ndER/h6AvXSm3Bj9a9bnh63e7Lt0dwASWkqO5RTd88ehBuGdNotf7Fgu58KElAAAAMQGfKWpEf4etKGWAMuE31cSYTo8cYlrTJwO4hT1iaYKoh5yIznLWIYM+u2AdetqlZhgAAACoQZstSahBaJlMCN/77mtspIuqyAydbMmdV2hUlnnEKtJjAgm2ilJGA1nNcqvKyjyAWC+D8UjH9KvdInrfn8JalYu7So7UCAXcOewUU5OwhCi4Gn87YJgH6zDs1GEUzCJ8L4JntHmTTg/MXIc99L1NBE6e+O3vv+OYZ2RYqJ+FOkeVPph7PM7bvRsalZ9v+6j8c+a2e9FeVHNkUcRhx7xXVeu40Ckc84xcAAAAKUGfS0URLN+DOtOgq1Qm/f1v5ExIEbysTirTzQFn6S1t6UO66bJ9QaeBAAAAJgGfbGpEf4jgOwq2BOh6mMPEvE5KJWUuH45P8CXP+ZdCo9KA8njAAAAAkEGbcUmoQWyZTAjf+8dbF2niwezLbWnHkGUoGoj7QIIOkFGvROy82ieNrbjUGA+u27nmEN5OyIUtEYm1Aa29w3BQ+77exlM1D2vRT7tzMqq7x40NgyQf9UyDOeFALvKyW/UjEdd+ZTj3e7SNxPWZbzdyhtLBs/d1mrAgFnRSFTdTWE14mrJQdgV/38CNOEQBzwAAAE1Bn49FFSyffW4h7bT+QcFhQIGL8ZhHK2aOc2q4nANbZ0GFL6Hw3v96ujtmH3X1f2UGhLCU7s8t2ALYQ5jkD+tMpZBnKJdI7zOIZg3UfgAAAEYBn650RH+HsPlaZUroiY+o5658f/kip2JFnP8SLnrUdiJYATw+vVm1xpmGefahnkYvhxrqmtK90ncP5Iwqoo7sRj3ypQHBAAAANQGfsGpEf4TUKzbRXPm7/rDCb1PYRHnrxQN3OgOeuwBoQBUyoz454hkxVpE0j3mzip8afJWLAAAAgkGbs0moQWyZTBRMb/wDxbhBIn8Ka729n9RowP1GPGwmIKU7l6A++D+NCeFcwINyDokoWo5uNZep8A+tnW0VIejZtiEBw828C5obBhnpukZNMqbuoZzcKwUcdD+TUNnmG1P7VNp3N/HRWhTj24xVfemeYzegPZhwqPtf5Vu3ZOBdTWMAAABFAZ/SakR/gQeNFYTouJTIPcX9ODyG/mOY1LtDn4Z4fdloxMzM5TGxEPn1lw/NNkwpQDBeevJFEc5m+/9Go9q+nDesjfivAAAAmEGb10nhClJlMCN/++FR2Jc7HZXI1ggkQzpQbLdE3oyI2e9RpinMX6ttYauzP3zhtktbs6rBQk+2bRLN2Zii5+KnqAGy6xu+oWg28GhfC0K53s+ETlNxGnX6kbnMFmbPxkhZYhjRahFfw8hooaTDBlOVRucGi/xGwHC4t9N/ktnpgbYThbnVUl6j4xW46umuQZce+h6EgY5YAAAAP0Gf9UU0TJ97IKea95tkF1g7PulJOBBxUf3oQ8Fgj/d/EvsVRz2v92wCCfdP3jpU1XFxnM2iUhxS2IZ5eBICgAAAADUBnhR0RH+HsQBkvuNRCE5o/cA/hS/2X1Ifx8G+v32b7PS02yVwPwQRWbCtWYhN3Fq4gdnPwQAAADMBnhZqRH+I4DsKsOugF5agq8KzwVB0SOKhHPr71cVpWK3aQSfqFUhMbf8T7Lt5HmLfgm4AAACIQZoaSahBaJlMCN/74Djgdzj87Yh9SyaadyNqF883ukqwWJfvVqpKQh+Cq2+ijPdoCbCBI/sOgdzCRvS17Gg4tWxC4C6lNS9Io4e2m05/dsoHILwYVElOoJDZ/MwYJQtRJ7sMYyXzLK7leOwuqfKqCHgoXnqt0hDtESiOuqtwcgabWxgY6yvJsAAAAF5BnjhFESzfg/tuGYHKjTB5ZIQwyRxCIT4MdkzIYOd8632MtmXbBRZjmn1eHX7fw0cYLgael3Bt26A9uYW+WQfDAKavRUOn1s++6E4EarbpfnHpm4uc7noN/2Jer1MpAAAAKwGeWWpEf4EHlYbRiT21xAh/pYXaBtudfqA0qjCKgx/lAHKVbTowHw8rOoEAAAB5QZpeSahBbJlMCN/7xIFmvqIe/Z6A/pmYt10zyE29uhmz+B2Dkk2oJvashyLN7AdP2QFAd8nNOQALgCphNk8C9krr0S5+4t/7IC9jwLdRCyqMfhEx0LgdtUukX/rP3vR0dcaKekFBDBp9NJs3ELeS44zfwhtjbfJhoQAAAHRBnnxFFSyfgBLo0GKtuOp0S/62v/Y0HNZk73O1NRvXJ78CK1NMTqIBwvVi+7nrhZNFcYmSvBtS8aarNw0M8/4bqY7V6eSiefuhvHdnfNViR+1mI4DfLT0GD/NfHmcU+shsVRJ+gRZvcMuzDyIyflNTInWB4AAAAEABnpt0RH+KcQxiH9X/rcK9KC9jqj38E9Cx6kZJfJmOLXjenMkHGGnhFyfDLoqpt/GOsv4ayAmoWSNp05UVCWVgAAAALgGenWpEf4Z1uDH/dvK5xRRkl16iX0NdLBx0eflcOeGwT2wyO5F6Jk/qa44eccEAAABoQZqASahBbJlMFExv++5XDgdzj8Soxq4R8HyFDsHMPNBcfu6u0JC3yOt9Zf/oOlinfwad64b/5ZuUQET9ghJNRnLBbp2fwAk7+v9LiLGc0Wy7m/MlaKm5bEbUwKubBxg4bDV+ieC+CYAAAABHAZ6/akR/gQeVg87ypn7D6HEVU8+rDMLH00gtcABN3WFNH/SXGNGasZxuBcRvLGcib1ofmI9nJ1EIhafY/27h+f0U80Zw3TgAAABiQZqhSeEKUmUwI3/74VHYnLGdrsN8fZFN+AYUaSy1qUzsqr5angQrl4TAR4cjThoFgPzgtGOm/+rZ6RiEUZb1p/ibsShVQuItv9q99GCa7QVtMF7k5//VfdZc+IdQXtUSRRkAAAB0QZrFSeEOiZTAjf/7wZu7T4AET0n1K/wcGaW5FPPPeSEVQq1wqPmVjHHGzMMaLIFS9D9xLjG8IHT3uYZjAjm+gAJnanr133+IRxjFAFTRBkx2srdq61z07DgwIEr4PcBIbo6eVoNdPRdn3p7C6JhswPrZOxYAAABJQZ7jRRE8n31uIe20/lZinhECBiXrAbzFWKWjBooZi4YZeyeDcnLXtkoxy+PPzoG5pJuU3cZMSZ6MT5Sk+0MRyXthyQf8s2D30QAAAD8BnwJ0RH+HsPlaZzJLf1BqA0GQv7XVPFANHr6Ehnvy7et3ba61H8sodLUCczvDm+Xl28e9s38/yEiA9pPYUvcAAAA1AZ8EakR/hNQrVcoxFq3sC1exT2bccLesQAlI9raV/idtSMz02Fau0gbClXfq1HVRq3Pjj/AAAABxQZsHSahBaJlMFPG//AgTwyBIj/nZtn0aw/YuY4OjamoTtQYFPOSeyuPadadakSLJXyzkE+yUN1FhGpSW0m0mzOfSaLl6RQi3UFaI3XVG4DOFmNwv2QiaKoOLHwxouX+jWVEmhJBs6gzIfBVG2WvwKcEAAABAAZ8makR/gQeNFYTo0teWJjfvPMFcU66/zULU/AjuVpjv6JMs2v8Uk8dPu8G96ISprega9vCssVT4Py38fqdQhgAAAK5BmytJ4QpSZTAjf/vhUdg9mx2D8J1iA1xk+D/u/eWFRKPYXoR4YrnXnQ9BJv4U7EyAxknkxv1P4KGiZu1XCzyMTxrAEalziwXrOaH5CIg19HpXT8adjJuBd+g8AVU1iY+ragKeFtHO8+FXCs1iqBNVFtNR+QrIASXXiKH9ySLmtNtwkBLAWVkei24IbUCqoeu2eSmChhIvR8SYf+gjoCylIj0hQWhGkF9P9sVMkdkAAABbQZ9JRTRMn3sgp5r3m2NaOy3X8HPswTGPT1V6qqyW8TAIkb1l84u61oezFaXM31Ex6tTfZXwH8l41V+ZLVvdVcPEZOg5HCqpsYndTaGYDZU2IBrf/HbgESDFhYwAAAEYBn2h0RH+HsPhkvu1KRWh5P9CE/7YMZ8EH8FtRGt2iZZYTb+I20oUztyey5/gcWsMembQCMoi5qzr3GWB6TyRqnmRNbxfwAAAAKwGfampEf4TUK0JrudPWO9DprT1cI+LcpiA9d9Uv2SIBa8dwd7ktVb0K/sAAAAC9QZtvSahBaJlMCN/7wZwidiwO32tSx4CI9VHvf8FfC8imIOnDW7X9dLvxLE8x1BQ2+Flva/hMlBEeDnDW8Kyzs7nJ6R+vwdkQhkORE0CohKGRVTKH6TrAI0jsBYWIHEzhI/DD598H/p+WJCB1WNOv/rDnbdHgZYdaqnTDwcDWZDxml47NphnI8tuJ1LEeyLF+MD0/bGs2eXrqa0n+B7DWP332LtukSOVrs5rOy329VanveHft7k/xSHl0G6dpAAAAUUGfjUURLJ99PYWDWGUO/inwiBAxf/vE27SbnXVov+t7ciiHO3pMMbk5hOMpHCmJoAtQU+lyEMKhoGTKDStuGSi6zfSBRDqB7UFC7Ub6lRrH8AAAAEIBn6x0RH+HsPlaZzJL6599KwwEOf8I9PqjXFRQKe4Cdxrd9B4PO3J74fYU+3lo88Z9cxAIlkoQasgv4bTkDHt417sAAAA3AZ+uakR/hNQrNtFc+dI4UgJcDtsPKVP01/w7gtDdiG8weomxmMV2jRs+mYeb8YsCdKiDBK6EQAAAAG9Bm7FJqEFsmUwUTG/8A8MxwQo/6EY7tRsTvMVqrKMcemP06VoeULKvsG05ajGPmO2MoK4g9WsHj/Pw+OQXY/RfWlxRBDYU5qtmrwgOMibbsU1z+xBDOIFJfIEU7QhyGFhExFTti7BO80z66P6DnI8AAAA2AZ/QakR/h7DkMUOCIRZCQlH2r3Q6geLsF83caeCI1bZY+Xd0YL1LmfzC34lEn3p7FbZoKir3AAAAl0Gb1UnhClJlMCN/+8EZPlhAhzY4hjuP9zep6GX41gR0vpNBRWRZ0UJPip2CeOhHj9isfDJTZVLrPDCGanOAhHmGphCIAPC1fKwnrEyemLG2qRrRX0P+lii+U7+s0UA9lDzojDxhAWTZ2G0F/aKbm54hHFNSHJtKf+oJOsv2i/hXe6R9WFXNJ0qKbHaCDXE5SAk17PNkKsEAAAA3QZ/zRTRMn37oQ9jW+VIhCNZdo1DTECbzHzraqtIWwcFXlHzfzpW3DjaPNagdZrahX3k2FcjBQQAAADsBnhJ0RH+HsPhkmMfK/TRTx1yNE7nufRZv4q2ObBg0Y85tUezANeXV6ceBVnWXr0jPx3kmWjuo3pxLLwAAADABnhRqRH+I4DsKquYeXbXa80z+Rp/+iDuHWPDYZQauQjmlruP/lzQ5MNgtrxuo9+AAAABwQZoZSahBaJlMCN/7wZu7TxYO6+X+O2RPzaO3P9AZSfEckBGR/iJyC69X2mzEdsbEARVatDv39gXZOq5q1DbnQ0mecF485kMmU//5paM4iZd8h23uDrAJgM1rO1fswnjjEy9xf4szlbNzDknjKgAlhwAAAE1BnjdFESyffW4h9qH5WbMcugQMX5WbtC61gVlqffRYYu75nGloWw2+flAxCV1Cr5kLWgK0cUy5ESSi8mwPgAklHpUS2QdtwkoX7J+G7AAAADcBnlZ0RH+HsPlaZzJ17dg40xRcL+9TtQ96Ic/V7TTmS8UoawBAXq1QI1QIwW783wiVM4yfS6ugAAAAMQGeWGpEf4TUK1A4rnMTHCUTE6/ZCv+/TbCro2Jb5kNAeSIDyYf8WGZnEtBvd7y7p2EAAADtQZpbSahBbJlMFExv1TpdGQI7CrAUef/B5h6ksk8VHxtoqBOtcVtjm20cMpPx7YKCba8PtsDiRbMvl2Vzp7ViJJGJslpXSe4Da21o01Tre+PL4ouLsDE2fza/o8u0e/o193a92YPHQjxT1F+Xl837GKhxj1zCeTYoBc72ajk9wDGlXMBiyPuI/NQcbrbcGJ79AkHjk9YLxxYD/SNREJer2QOEa7kdnKAs2nctMaAcgEj7lXi2nC/Ax5gsLQKiDWNIJ6EZsLrTgM5L6GMBmxmY1G//WcuQa5/+wM7s0+v98zEhGRgkPv8Pqi9S60lbAAAANgGeempEf9wncM6KwnCyxLkxv6hw6qpeuv8JH8/pARA3aQiOdCW0jL7dvTy2HOkGJ2WI+L8RYwAAAL5Bmn1J4QpSZTBSxv/0Ej2oDXQbnmJ9xHt6Y3M7+zkeot1lPkVAjOrBk5KutJxuFikQjMMm54iVUcgq0mUJ7XmxxVFK1B6T6213KH562Y+RpgQQ2x7gidLCR4uQyyI44VQZbTdq2DELgkXsCdiNKr8x4ESrFQl2hMUKVh9TAo5qMet3aV71iz0PMWrns+RwD+p05KZ5PkbOR8UZdEHr12Xyp6MXO03cMCdxvt8XSDnM5DDAwory7MxXFStmXjinAAAAOwGenGpEf5iZsaTqsZBVHN23kGvXifJ7fiaFVng/pASKG+Iuw2tQVKHpTjjL49UFkCnlve1QSoeNybR0AAAAc0Gan0nhDomUwUTG//vB61Ydg33yGjTc8vx+XfX6HCnb3oXqwFbwLkycdi4a18wxKaCie2RwyAvsyA0uooexqJ/tYWo+LRvOoUmoxRFTOQSzQtQ02/rq9eSsBVCGtHKcyXiFeX4sl7E4a69xKNtZ4XR2mcEAAAA2AZ6+akR/h417ZNpqEQkqKLsoWa64MImq3wZOBY+C6717C9/lyBZ0g6hTD0NY1rt85FbbEW/AAAAAnUGaoknhDyZTAjf/x6o2uVDXrlL01pKNSmhQXr4P4TJQsSM1NFMlKvnjcKn1PKq2soE7idryXsHEaf0O8GkHVQjbY/noZr0ipij3QQYPsNY5BiccYwiMaOIoVLU7GYyZILtn12cl54KfPwPTwjvSw+K4bQ5yOg3G/Z0pavUu5uigl8D6ssFCj6v5AXSFHWAjabypdfzcI1hiTFK83ZYAAABNQZ7ARRE839BbCV3PROKhTO6SEYHVcap6D/bB3xdjsl1hYn21dsFTAVhbVYBtR9XbqXR7RUq+tQJh1Q7mOALEIaPB0INz3lNIRir4aIkAAAAyAZ7hakR/gQeVW1VzyCm83gTep8bM3if7zuz+nCbMLSDxTS+tvaDoEjU6XYt9NoaA70AAAACFQZrmSahBaJlMCN/Bc9IeRdMKBQUoDdmOpylyK4+kq5dpRNdskQahtc3w+l0fbnrSe3FvIzIgw10Zk3I2ld9hkChK5nRYgL3ZtOHEBaiunfTa4ieYkASLLA/NG32QGn6PIgI+rNletA0lBbBKcG5b3PN4cer0cfvtkcfPCHp+QZU01FpMCwAAAFxBnwRFESyffT284Ut2gLU6T/OF97/L42Zk7TpYejMubi0UT1LYIwyi8wna5sal5SW2vk3PndfO9Ygjv7w527MVQundwsPGKGTzDFCdJOti60VM/ivQ0/9l0RrwrQAAADMBnyN0RH+HsQIbL4x/wbzGsYYRQbbp3oKzYEUN0nPNIKpfJlJ0wD5Mk4OdIwxB+uB2eKgAAAAnAZ8lakR/h60AVVNnY0Snvntxfl4ah0IXe77u05wjItrjxSxk7twJAAAAPEGbJ0moQWyZTAjf++5XDgdzkPa1JXdlt/iHg3e20KGGbpJNaAeD/NHcMVqnpuI+rGRg85TMlkZYGQ+buAAAAGVBm0lJ4QpSZTBRUsb/++FR2DWGd4E74BcOffFFDrjSdNqIHPjqmL/A4V7zuz6iK2RJtZ67hmaIV4RT+DJzIBQaAPE+vuu5mX9yRlf+Xz3MDvkp8OYGDoPwsmZQhPHmCII0d6nKUQAAABcBn2hqRH+HjXtk2moRCSooxNiKe6PW+QAAAHlBm2tJ4Q6JlMFExv/74Djgdzj87DB+8bpj7b61MAK02m0SjItmx1gb8vC/xkdmAxbAK3mtGvA6G2vVc9Tnw2/i3sUSrGL/HbhKMa2sK/+P9jhm46fzs6O00Ue9XW7FkptH/JCazHfcLFaniM469HEARwqsesFKmIPAAAAANQGfimpEf4EHjTTHLttWmKSJqlJBcodqFPtQbfjXOzXaRUJjHdTp9FoeBf71Kn+yqxtU3pUeAAAAXEGbjEnhDyZTAjf/++FR2EAuCDGvtGowoi8nI+jfRpfapBnzNQPYo2Vr8QmdtgYjsnpaHEEUEj5fzdXZ5z8ZMEZYVB+N4bG+rt/oQ/tJOOnO8NMl+l71MxcY1QGBAAAAZ0GbrknhDyZTBRE8b/vG3I0XvCGxfcoqF4IjXpZsDkZv7oRIdkuZllH2s4RDfJ2UFZ0K+2+MrbzLGTPgjZmyPT8lOVqYMq1VU/V4XinxIGWe6yhZAC7TBvENB4AWJBowQdYdeboaXvYAAAA0AZ/NakR/h6Avf/WMAaiuyJhL+KT4JwkHNgd10OH2kTFHZuZ+4In/rkqbThwoea2xhOS5EQAAADhBm9JJ4Q8mUwI3//vrPjgO8x+Iwxa5AW96PE6JjAAg9tvdbYLxnthXvYjAICsZOG1tuY+xvSOzEAAAAEBBn/BFETyffT2/FQo8uLN+ZccYdXGUegcTqm+73SyiRqyxjBgiXqvyckNmzTA/sqcKN6IxF8rB8AB6aJm1tPGRAAAALQGeD3REf4MOYBiW8DO/c0klwuDU2bHJ6zPTXTSFXKSsxC0nU39JMiqbRsiRQQAAAB0BnhFqRH+BB5VXHdOCXHYpVjQiP78rXDhs2RTqFQAAAHNBmhNJqEFomUwI3/vhUdg1hjtBgEPsnnqReMCf9fk0pXkAepqW0GvzkT/N7d++eAjFF0LqEWTOQwynb70MqWR2LVAT3BGuJjMDsrfL3saQZmSV6KVXH2F2scQ1murg2f6hRdm+IshdKI4OTJ0vpuACyrZZAAAAWkGaNEnhClJlMCN/+/6N/USMh05oLuDoAZwruKxif239tQ1VJ4v5ZqpArC08Gz4PxiK1Zw3gJNSO9N1MaBrO4NXyS0Kk5yJ53TqLyQyUa3YaV6o9CrHZ784vgAAAAGRBmlhJ4Q6JlMCN//vDcjlriweTEJ1e1T49rsbXnEYWTbUIkOTqF9p/L4Ed+sSuPC4msb/sWIIL1Bg3A0EKKr8CzywgYo2G0BZdMZXDAvTPj7Xu2U8kKOf7J0Nmp5Nxc/b1wCJeAAAAPEGedkURPJ99PbKA1hlDv4p1vxOs/Y4zfPg6iHF0NVmxylt2/uz5ddxtZSf9odYXHgkvwvq5a8FEnAsAeQAAACgBnpV0RH+AbCrDaMSfDlBTuXnhXQUYGq59CbNAd/8CU+K1jwZOGWMmAAAARAGel2pEf4egL3/1jfMenU+Ipl63RDQu+raphKcMtMgItiKQVS5cjIUWjl5mHHdx54Jm73xhCUgr2g10y3zkK7P29MDQAAAAYUGanEmoQWiZTAjfyghBljfBFTHet0U3BKQoBgmH8wmCWbb2fhS9kf8YrXaOzQFjOPpvXVSbYEH7/91upJzbcIGzTBnw9iN8+sz32CUkThqoo1kDvR6ExRFTQPVt63gM1akAAABKQZ66RREsn309tLhK5ErqsM8w/NBncxs8+XNMfQtyU397Ah/MDPI+uVDBGq9NMLyY/eg7uJexiq401plodH0NOZVtS74swPLgbj0AAAAnAZ7ZdER/gwyVSt4Gd93RvQzgj5K3CnQvm7OjBKD3UgCK1+lg5LhBAAAAJwGe22pEf4EHlYPO0DSNXbXCGQv51Mgmje+6XpjUxwk1173/hFbKQAAAAIJBmsBJqEFsmUwI3/vhUdgSGsf7YM42S+/vjzG31Ii8r+yG9Z2rwtMeopT4f1ZXmgCKDIL/jmrbYA8wR4mFw0PT9NVLBPsNXrKLeAacgNRksycpYqdXIV7VH+6jtYtcIHq7VaKAVYULBW7vO3pgC3A4br5z1oXvv7pxKzVVrdDydj0wAAAASEGe/kUVLJ9/0qcZp1xS4Q6K/byHOdoBIV6Xps/UwNkM8yttvswZ8Jzr9I6Yffad0OJ7Pt/vu8H/euA9uoco+FMMpzX4CWnWQQAAACQBnx10RH+HvdEvcu4XZBGgCHIf5/KTqwJP2rRK48Dsj9hXDpIAAAAlAZ8fakR/gQeVxSFc8i68nfap4V2gL3VVnIbk0JUIrK52lB5NkAAAAFpBmwFJqEFsmUwI3/vBm7tPFgQM7JWZsrd+AFcbR02voqglcUnWYy1KjlCBoWGrB4/yDkcsF/yfA0JiredR3d1/zQTCNN3YSNi3+6GLuK5XeureNSQvNWSNTKEAAABLQZsiSeEKUmUwI3/76z44CE3mgfD3z/rIASyI/4ajysO1TyQ+cKUy8BCEPXwFvV7Q1kwh0NGs/guw48vBxqodJo0my+qYZKZwCVTfAAAAXkGbQ0nhDomUwI3/++5XDgIGEb1mgVp4z2caIApSFuXOZMOWUKsM6uZXxOHgwgvvNbQRkVIuaLICKhGo4osmKLu5nHbNQ2eBpkjJ0tBuuGJ+8FlhcF1OznrI2UofcQ0AAABnQZtkSeEPJlMCN//7xIFmvqIPwLXapFLrpy24N2AYpX55tMIyEhwIGCCAWtX4dRhVjqoPHEX7ZHYi/+5mDJL6QpKDldqAw59qEcj7XmhguiYHg7hxSzxJYABsNFnTFdhSjxeeeQ3HwQAAAFdBm4VJ4Q8mUwI3//vrhHYIVisLS3Cz75CdmPBNSQtJmyjSkVk8S6O1d/mV7n8Rrd+6LOzS3sEdBdAhEdUoK/3HLiqSbIKyHxx0QL+V1ET6o51gLYhrc+AAAAA3QZumSeEPJlMCN//74DjgJBsj5aW27b4bt8+9bnwCqvg57ian/4Crkwz302lGZ2Pv+rWZcpSzgQAAAHVBm8dJ4Q8mUwI3//1jexkG5DGfTzn6ACkLgiMo7PNL+8MNWKBuY2ww9P4WA81P2z1Uo6HFLULVfe0UWARzjI9fYO1TGuDiw4iAYUYp62du9UL2+zNI7rHN267ZkwLWg++4sTNfEBlf+9nqtqNbpC7ef2VGPncAAABVQZvoSeEPJlMCN//74D1EH3TVIcIrzH0/chzRECMo8UYdBgVM6BUQ/qnmrHcIvQI2T4gwNuNCHOfc8xOvd8e4RUUrEg7UZrrwxAx7PnW70QrnAkn3gQAAAEZBmglJ4Q8mUwI3//wDwzHAd5j7ZQY2ZrSRVRGS4RBORnR2QxEztLNjTbcTvgLcPBn1MFP2hKUm69thjXP4Z8ejw6yLt7eBAAAAPEGaKknhDyZTAjf/++AcgTWtokf5L9MFh6k1zvz+EhUSCxgLY7sBvcK7i6dPmedsqE34PPj3OdSB7bkjgAAAAFJBmktJ4Q8mUwI3//vgPUQhuJTLhEJqu+Apna/+hU4PzU02rtLQFGBiYXaJMHEktyfr4sDAJw5gQgiM1vOUDKdw5h9IrJLuyeXoGbf4wrVN09r3AAAAWkGabEnhDyZTAjf/+8N7/LJcoi6iLVaYzKjcM3p9XtxS6Kwh3zS27UOOJokA10Oalb5Mh9jTjqiJ2mZpX4f6frTs95x4EQuNnHB2EuOeRNjdeN33AYWtQV+GtQAAAFZBmo9J4Q8mUwI3//wj5K7wgzjvuFjuy3LAdQ+tnajIU2FYf2/S8asjPFa8Heexf9iyg+O2t1pIIQe9N50utnzQAT4GusPtoeQwLMtTfJDIuk4WzhVsVQAAACpBnq1FETzffVS1iQAjw4mKGWZ99z4TIUbS/mDFOaQ2BzRPfOEJhO2fZfkAAAA0AZ7OakR/h7ESVYA0mDs38HcgN3BlSvTaG65b/g/pfq94yFXiLsMlUCZdl+Sj6EnesriDgAAAAJxBmtFJqEFomUwU8b/9Y3sZAzW9hPGahGw7flXYE4FP6yvCNzmXPKd9sNuacBlbT71wo0upN8/BkEcSxj5QOSyFbOl4aFdbwIN/85qZNxlQT3WmJ0lyxLEqm3GKGmR/Q9dLRwbnvK1UjtfHRqg/kLxCsjH+LVN9IFW92ci8U6z0Su8pzeXhOXBhoqHLd/KR3LzdJ3Q+/pqGJRbSf5EAAAApAZ7wakR/i8dTtuWwJ99n5vKBTInol2bsI01rGqiIrOoR71X4Q081QcEAAABXQZrzSeEKUmUwUsb/++A44QawOdhgfUi1MfHqmy+I4SJvb7slhfqvjc6dx1/LRyTOBiepLebRHg/6uB0qH06DelPtYLjqPVsidGv4WOqaIDxCE334DXtxAAAARQGfEmpEf4exE2/rMauJo+rnrnx/+SKnih1QnG/vzvVzCJmpdIrzItT2aVilxCxecVUkmU6Xcdbiwv1Hi2eVfOYva6XMcQAAAHlBmxdJ4Q6JlMCN//wDwzHCXNMnztovd14vJ625TIq7dg7bjOkJiogvd97B3i0yj/I1tMjzhb+H4NpywxTMQVOk9VS90im936odpBgEAln5etaGgewn7zxeYOtk/lsAB0UYEr26ytGYTnykyW97BDgbmncu1z0/tcTAAAAAU0GfNUUVPJ99PbdJpKoJiBXMhGSlyCkS3qgHnDWoGiC2KtexepnugaGZOIX0UYhHhpVdynfptSF4E2r0Uv58jiFtJmUBvNE2oOgYwxDQ9IZDYeOAAAAAOgGfVHREf4n36qs07VZIPxf1t+QUqlPNBFREIxt7hE7ahK0iE85r2mNsOC5htsIgigROAWVjVceCw4EAAABSAZ9WakR/gQeKxIeCDSxZxf39Gix/112nHFC1PzbN92IjXL5eOaAOw4p3LVmqz/XDkVqmVp/QwtqVFTXi1vi4EzWJjc3Nfa1Z7l65jvwwWZi0CQAAAKtBm1tJqEFomUwI3/vhUdiXOx2MTspDBHmUaNxnAKtxo1ZidUG16yzOqQ6r6rwou8sSCWAoHPh0R0076A8s6id90rZISEbQSI4RHOQA08yjeqA8L6RD2GtXfiIoCm2bQHe8hkBL3gMquBV9xsgOcsX8KXQlQ4myFK/fF4oEBZAQrHDDSzotzZCZOAyfnfF5fPIXTv5TGCrdBxcbOAJoIYdUVUA1wPS5ma+yJ9AAAAA+QZ95RREsn3sgp5rjMV8XzVP+pp1YxygYLRWB2wZHB+pWeUU3ZYrSXAwe2bahlT/6TYFQTJt0HNTpEP3ubVcAAABFAZ+YdER/inEfdryR+UBQh0g+98kl5ilCOtNo1YP/LuQq8RtpQSglo/BF1WPBxIxyIr5KAxwIpISuyR2pURhW4rUvLdfxAAAAMwGfmmpEf4jgOwqw66AXnAyisbuurz/RSCoRjyFbk5dILFFLRMyoUi9I/yQ284snWqurXQAAAHlBm55JqEFsmUwI3/vhU4AmHk2jvpp39+DVM7fgz+/0mDAgC9sSz81+siWGpUWfuj4GyTOZUOB35xF0hF67DNn0ZuMEXcGIXz00pKtOL3sBXsUgc/ePA52/EzhUAvLLgjOyGZLNGqxerZSlwTR99vdiMIycknIR/ViQAAAAUkGfvEUVLN+BtiMx77QlavkuKJHZZ9Q3HsrRt1QLIHVTAm1R73+pWNsjQ0q58gETaMA5QLuokf/hrA/s22hdm3istSVmd0URqKtuvuIfwOuQF+IAAAAzAZ/dakR/gQeVxSFc8gfgZAkdinyst3l5uX7pQaHSAIaHm6zNCuVavCVTgW8m03yN9WRhAAAAnEGbwkmoQWyZTAjf+8SBZrxwNxE6nKekPQnYXJm5wBvfHrogTYU3NIBfehDfOIiS6ySK3363CBAuzRu6fA4vIPVeDAL86ZR4LtBStB/BM1+Vq+8tYxISLhvBnW8Dyw05HPMDTibFdFLhiDfH6SCO/cg2NYHDixZKTHvuNNlOBA8Eb/AVnrFkz+v5UJdJl1nEtDkUt9F4X7Ko52NhMAAAAEtBn+BFFSyffZxiPbso3KVMv0xrg/tKL0xl7/2ug2XFt9Z/8yNSEkmzVe4IR1kQbVaS6gPVxGBlsxGiPdVg0hy8V1fpzap1S8LNa54AAABQAZ4fdER/h7ECGy+Mf8HMSAoaH/GRCeAOPoTHrElxuxk6cLhDnYyUbEMAwSiUszmoA/R1gKCdAprLE4LysACZvlN3jo3Q22r11pJc7VqKCUUAAAA+AZ4BakR/h60oZYAy4R2w8B/3Wta04n3kCJAW7+Q2gBRN+jJHYMfX5uQDuUKfZkFiLp0cf+1j/fb/e24j8dQAAABxQZoESahBbJlMFExv++5XDgdzj8Rhi1ZsIZ8Sm28CkaHPAYRK2SFSACU5cWdipBQ+vrEx0RbW5ZoDsg0jJR++lEuFdaMABu0GnPbEfBPglHmsDhwytPRmDnxl0gvYO3LQZM1n54PZOnfSIOyUJsG8mYEAAAAjAZ4jakR/gQeVg87ypn6p/FGQB5EJfWf6QKy0B/hS5J4dn2EAAAB8QZooSeEKUmUwI3/74VHYRK6xsWuLkzkMFi5kUGmAexj8Ar3bWWH098HYvsoR3J76TgTz/oDQV5cpDVqE1+A6h12YunuTYfth0SsPDgk0pLNX7k4axqHCb90xZXJ6QQE0rqZvxK94nl6XIz8bJURBNnwVSR2U9urp0GW5iQAAAE5BnkZFNEyffRPrpUzJLr2/AS2nxOSJCtdNcLjn6cXpPNkqjvg+oO3azFJERI3eaUINnWDk1Poc6rrAogWMSwS1ixlRAlvzi6VgIwdfokMAAAAvAZ5ldER/h73RL3LuITCI3oDv4BPtzSRBT0ox9Zs6XKAiQZAu3sncW22aE5GNUZQAAAAzAZ5nakR/gRHB6GWMSe17aOt16nxs4vMoxLGJ6z8ZwihTVdCrHs/gv2FN45XH5F77G7EvAAABGkGabEmoQWiZTAjfzjAzJNh9+MsHWu8d7Cs0acVk7TvkhyFGrIZeP7VKZT8YSa+fqLmJANOx+HwIQKZ+hbChanD7CSavKxfEpeCd3O875McKYrY5v5SCqAR4LZnb/unYHbZG6CJQLkdmdXI3S9mC+vh3J1TAPWl+HZzfy6FGRPZ20VqNe1/AuRmfBkIEA4iz9PfeZIPIT/nIKX4skpc+PdUp5tjAsBGG8DV35+BHBUfG6nHmCnw+kVYk4zIk4Xclt1NYI/JiZT+WqrMgnGmO1KB+XctzpeUXBmT7ef+PRBXYvDjeBflA3MnR1+dcK7H8tqUmzVSB1TwBI9+qN2nnuuRbqrb0lST+s59Y8w9uzAlwZZ2ZsKrcj1TJcQAAAE9BnopFESyfzVU/m33U0DGS8eDHAR1n/6XVZZ9uylPsAlipuRHFF1836S7pQx8uES23pt0I4STTpIPdian6xejkAYoF2nWj27H7EQ1w3xZAAAAAawGeqXREf9Pb2DkjtB4in5wCGYpcxy45arMPjM909IA9gzca/O+GyZGg7wkOdd+VfzvhhMt5NvbfAVFmfDXW5Jmj20M+nvxZ2Z8uZoUmYJ/Ok7BKfgCSaQdQQUP1bKE4bTtcHvsCHyMnF6keAAAAKQGeq2pEf4Z1uDH/duUzFC4JTU6hTUz3R+lQw5ZFRkiy5ri+QAklUBnDAAAAOkGarkmoQWyZTBRMb/kp5kxTMHZmETVTIgsd+A+NbgreHIdnNX9tCFzUmdk/+DL0/JDoveQNd/vCsbAAAAAgAZ7NakR/gQeVg87ypW2l15M2hNNzDWakH6ukSDuxzYEAAACpQZrSSeEKUmUwI3/74VHYPZsdVGEb9i0qy4XOd49kR/Zd27xbxgdkhHGHJC5y1l2os18hgqRE1KaZVj5VSPxxYMpx681OhYazthcfQAXvfCWuIYtRdInpGWzy1sxMLVqgPWw2rPBeteX3bwTAuESDIworEPEl9lhf6OwWpk0mu9bIKNzmWd7gpQ9E8pnmi/jBUgxN2yeN4FcSjUT9R9wbJD4lX/zQk/2MtAAAAD1BnvBFNEyffqpK6p9FFWrePWQt+t023fow08TUcBgf6pBXtRGOvd6vNalSElK/B+orsc6pxa+DU1KIEsYhAAAAJgGfD3REf4e904PVduAufu3efP2dUs6dm0IH8JDwzh2JpxWrp7PBAAAAKAGfEWpEf4EHlYbRiSTrye8J7wrvIAVKDRGHYha2PjTkam//yFWsIsEAAACbQZsWSahBaJlMCN/Hrv6DGHttBKhEf7oitzULKk/UkvXrRGN5TSj1fcIGe9X0COkbiUIPuCUmvFu3j4oaIwdDOZcC6YE657YPavN92lNXaCH3mYlGVn/ZhGTMtcl+H3Znx2tzw6OnG8LWtRLGegqT+9PKN37VFH+g0bBpNjO5UwkVHcgMy/NYkLLWLbOHn2+rwVZaqtmV1BLeu68AAABUQZ80RREsn81bxhrs9aLlZoqp8iOSssZPF6+vuh4reMnAu5tYCzgCviZbUJoMuyNDyQ5sQrG7eozVDcPAOhyEE7YCkUrWX17txueCEEo85uzZ+398AAAAOwGfU3REf57MlpebTfh/uLfAQTkb6fSvSijofFdaiiHHloNZyWGn5J4eiSHd2ZAxcNxMqZvRdlvUiUYyAAAAKAGfVWpEf4Z1uDH/dUm3pglXxGoCPaXPuwJNhybnmpjlHL8vPo6GCD0AAABbQZtYSahBbJlMFExv++5XDgIGEb7MdX9/tJBBP77OrDPDuf6bRmi2cH4plgeyDD0btjDKcl9nj3VzhH2Dh3jpoNa7QLLWuVUen8M08hu978wsuy3OhNJ1Sm5jUgAAAB0Bn3dqRH+BB5WDztJXeH+x0rqoojqtJRdUUKnIgAAAAIpBm3lJ4QpSZTAjf/vhUdge5Ym1P1V/GNX9AIT/jkz3GQqbKNPRZXoIj0uUAbDPtH3Rh3Ek+Bu+4pMcb704ggfye904H48vWEDu5uQg/N2uSES/KbFzLPFtQ2OTa/8qYBpX1/aD+T+6hJJK1ru+MZTghCGmNpSQXWnK2Hv1rX94UdXZenTBkKmXqgUAAABnQZuaSeEOiZTAjf/7/o39RD7fprUF/NAaOCr3pu1+Sc+O4fPzCXHhcvaqDp9XfBNrwZ00TugT9n/WmQBb8BI4LwStzpt0RUudpuht1/3oLhjOlCXs3ObsVj+wmRIxGocRDVdQfy+08QAAAGNBm7tJ4Q8mUwI3//vrPjgdzj859cWvgfbdn4AHc5U1oFwreoYsjtgTR9pWZhFMOPGk65+kAs4C+VojpHiEg82clYzEV+LFhJqn3NZwdE5ihwYUpFp2IKdxJaXvho2ltDn86KEAAABMQZvcSeEPJlMCN//74DjgiV1jEgudcPZjzgKu40Dxkog8ar133QofRc/AH/ZjqW6NQl1fHvbLAdPdJrabUaT/78/X6QmNkJeXNfxPuwAAAGZBm/9J4Q8mUwI3//vgOOB3R8fPH5+zlRCyD6mVH/aMulGQkckScEuTaTId2v705sWz0b6RhD6pVrrfZlw2bn4W4jUHEnaKMk72slzxowCIHDa2FHirxY2J4Aogy/1G71Xvqzs774AAAAA5QZ4dRRE834G9C6jp7AxObQ1D5nh6xcB4XoXLYDo06V5jfpQHI0j/RblPPLPhUSYgXonQeS6joQ+hAAAAKQGePmpEf4egOXq98g4uEQsyFmb1HLaC4k+C8XMK02ZavI+WCx5//WHAAAAAQkGaIUmoQWiZTBTxv/vgOOBCV6+DEJpLiNMEuxO6P5BRQh2K+vNT+Sg5udE4JTIIQLzSP0bwLFYQSRlbfZFEgNKAxAAAADwBnkBqRH+HoDXutdsAXaYM5d5s/ldvZ1ZUZqSPwmL03HPhv2wTpeNIW0aHyufZg9R3ewLwGGEJnNYKZsEAAABJQZpCSeEKUmUwI3/74ByBNNkgWomc9ochzMRrfD4ihgnoRApgkRdI+usVmu6Y8v3e0UScczsCjQjsNfzAoG7XkUz+H1yuPskDuAAAAGdBmmNJ4Q6JlMCN//vB61Ydg8gqaxeVjbcbFf2tmleGxIAMrv7NFEHar+eha3FuY/vTZzPKjJC2Oq5sKE5oNwUknilo4J+t1ZmDCAqKelikGKtqi7gdJP35m4Kn+yOr7aYtXrd6AdZBAAAAYEGahknhDyZTAjf/++A9RDqMbrkil10LEKl6Ttr7miP6kcL0ohDF7+vXevrzG+zTLebM+8F8Mqm1WPiQigJSclUc32sO1HgWuzAPWzc3pJHtUBg400W+ishyY5Uy/xEMmQAAAEpBnqRFETzfg/tuGYHKjS/jGBEY6YCRcT42dm5owW7iDaE6eZzNyRpt3yfTYmI4OxsgAkcyVuBGb+F3A/b7UbL/gPEwyEc0Vpa2lwAAADsBnsVqRH+BB5XFIVzyB9wJwDYp8rPxDk9qpTEARkTv8z3S1UGE8PcjATyV6Zz8GAZwGPHIMrJR6gwBSwAAAINBmspJqEFomUwI3/vBjyyXKEdXz5LJiOmapC7Dh+bqzRRczmQAYnOJBgYTvI6vsixxDeHUOwDiuti47UC86r13T09kIqPRmtJYNyAq1tP4/Ft/8pL1VrpczUxkwqL5a+AJdlYjtK4aB1EcbZFat2ii0Xh5BXeTFj4VS6DQ2G/nJLgwVwAAAEpBnuhFESyffGXEt8MjyOXZ1Br/nR/99FrXZZadcoA6YG0Znt1Y61Fwz4DJCwY26xDhlnhFmPYwjeDf0M+x30DUfLS/zFrIDmZ/MQAAAD4Bnwd0RH+HoC9/9YwBqK/PwA6Zil5xHhfqFkRqWZ06PPIdbHE9bbxRGEDxKdUaS9iieGaHN/QijO9ByAGQWQAAAD8BnwlqRH+Gdbgx/3Uy5jiTLpXihCuS/L2kShyyEDBZR05rxcfD2VGCwXsTh9INhGlzQp0UwAxUx1vZR1T6xUQAAABZQZsMSahBbJlMFExv/APDMcBAwje62AS64edylmitQLrKTHGHtcvoRS05r00OvVkJCNrMh+/mVw3eQxDA3590PKiZB3nWfrN5pDxzTrKYWq9sI5F4omOgfooAAAAnAZ8rakR/gQeVg87QMgUIqzoc+IQ0ccVZvftYC2gzOa0pG5p1idXxAAAAbEGbLUnhClJlMCN/++FR2BsEzfe0cZac+UJkdgjTB26Jig5W2j5cZ26NMmefLD0RqMng2y7laLakc6lHRdZNw1oHaHFsXiqzatXDN5Pz0v6/WwXUQ/oZM3evQzDrkzK2TYV/MjXfwp98+OlJJwAAAF9Bm05J4Q6JlMCN//vgPUQfdNSB0ivMgv12kewdPs8LDHEnccStTFGw2T9LqoRXFd4GA6Lmn/C0Ae+2Ys461r7jzbWD67darQQW6fpKD70itK+HYmaanX59v4Kvjz2TrQAAAFhBm29J4Q8mUwI3//wDwzHAd5nLgo0KbNa84paAodp1tGNxaXmvV+VvBFtu1np4hpixdq01/tsJuXnBDpI6pjgYmOuYM9uC1Czm+OIeAHPNqzeKqP5ocdj4AAAASUGbkEnhDyZTAjf/++AcgTTobJBaHKECxQ6c79GRuRfeQ0KwXvcejjr76PM0aqZOg6nHibX9Iua+A36w2f//ebnY1vV75tCTbJ8AAABpQZuxSeEPJlMCN//7wZu7TxYHfZLW0nepgrO++53cE5t5rTbH/j4CuZw0tqBjVHvT/NuCxYAqZPHuW+4z/LgYIOyIRCJSyOgmlLVczhDzE5KkFP/GXuQLXa8D+P/5fNEgPj1lTUpFVO4xAAAASkGb0knhDyZTAjf/yghD0KSEAj8fRNDXd8ZwkGDcdqAIE0puIoVKKR11TsEWfh9IlRceiw1biBwH50bd5SVeqveaOBsfxwvHo+S/AAAASUGb80nhDyZTAjf/+8SB2aRkB8IFrk6Yu0T79Zx2yWX26LbBct6B2P7nMJlSNU+8im32OIA82qWr6M0WXr/wOnu9asEIe+TzxoEAAABgQZoVSeEPJlMFETxv++A44EJXr11WI/mWGmgjspyywfThqWu6vt3f6IeutLnnxyOVezpxz66ZO+XTJSNz+hSyTccsijssaFateiNeQeLGmXnXJq3lh1/FYw7UzWsF4lLwAAAALQGeNGpEf4exElW+Y0oRl/FZQhP+oVccc8fw6tVIVhtOXYZd8Y1R9+768Me1CgAAACxBmjZJ4Q8mUwI3//vgHIE02R83+/Vlx6PE7gcgcvw/dfZOmqjB/6fbU3T8uQAAAFZBmldJ4Q8mUwI3//vhUdg4BM3dCQn71waETfG5XS6Y6qV1OmJ2ZI1/p1k7hngqfLDRlYxzHulMp52MY33dNDvA4xKrjGSUX3xOzqdXe+MisKohr4vpgAAAAH5BmnhJ4Q8mUwI3//v+jf1EjId8qaxOH+NL8HxQ2fp5PizkRPrB5xmQS+tbXXrCt7darNka9SvAVPptUtLJ5w5YyvLUX1kGEExPUgNHTM6K6C2ru/0HOEs9PXbEuaYln/GuFzif4kjrgcvg+8JdtRNs5H/bTBw0pf0EAsoCf7sAAABQQZqZSeEPJlMCN//74D1EgTaKZbqNH0joJJniUx8Ruqc4abzu2rURnTM10DKgg9fFoJmD440LTpJLO9oCBujuWU7mA65jaO0wi7HremRTBMEAAAA4QZq6SeEPJlMCN//74VHYnWWrD3GQIHk5ghig/U3DKJNeEKm/nL7OOuV9SKp1OrFdiEZQH1cmsNEAAABQQZreSeEPJlMCN//7wcOsoiIc/4xpRxHtkU2H2faBzAZluyLYQW6IF6Sg2Z8+2s6OiVvP7unG/PaKjnvHx284RVK7RkYgtGcnC+eIrXj0kIEAAABCQZ78RRE8n309jmEtq7hLzM71lTx1/0xhOjUhRbpoPBZse8D9DOokAeCvfYnyUFigFVOpER0UjkrFfP8CohPIs1WAAAAAMAGfG3REf4egL3/U68Y/kJfuO5of8ZPB+vRdik27eJJcU3fOYgVo/ofNZHen2wm64AAAADIBnx1qRH+GdcG6e+81iB3HSdBA1qVClucbNobVh3bfiI53Yzn0y9s0dcFwAHbyPWxxUQAAAItBmwFJqEFomUwI3/vD26eZCCisgtY7Q4/nAy8diTjLXaDcRhbm6PzfosfVpvyYRHRlUjwp0BnQ/NWsB7lZrV0Z2KMWnG+8gj/2wN3ueG8T8c99L1NBE6r4v7E7tCr/OznxbvNYjxfXhceW8dti/xTVY8P4CvNMCqkTDHLfY5PChguysyasokCazUmXAAAAKkGfP0URLN+DOtOgnSjGpTM5Wdix7kUjmGwL/wc5dr6DxK39Y49jtKsbgAAAACgBn0BqRH+JNh+VsnD/m3MN5x6Sg5Qu8+nj/5HUWhu2O+TNUelAcVFhAAAAlkGbRUmoQWyZTAjfyghC9V+ELITH3LCIHSIvm9CG6OjHlL/31AO68jyrK9IXLK5IdMok/VYaXk303UZ7pRLjQA0f/4qS5bfqL5QiXvH0ffwKCt0Fq2vBbWwZIOms8XXulZOVX3znJGj/solkcNIBtwjIWnGGvWEb1puQ7+6i7LeJ9OTdMX37A9Yfx9MPio17AaTThn/g4AAAAEdBn2NFFSyfeyCltp/KzZHjCx/v+MwPlO+6S0X+qGwMBBl7Zd3aGLZ6jXXcpABaEUHQQndnlecT0CSIxETMb0IEO8QcA68axwAAAEUBn4J0RH+HsPlaZzJLf1BqA0GQv7XWJY6mi5DH26gpsyq8ggTdiqyPYNRnn3ZbbfzdoybbOog2JQKGI36v4G9OFjBq7XEAAAA3AZ+EakR/hNQrQyxiVMaSwCvWnQsw2ksmBu2UpY5bkEhGHFQrC+fPFAFNp4o85UVgqnxqzlGAgAAAAI1Bm4dJqEFsmUwUTG/8A8MxwlzCIUaKgXMXLU3QQJ1EpMKnk54cIKtfRAathRBVUcjjYAmeqTvTihh6Nuk08XkauvWPadL4Vijodg4AVC739ye8e1iwE/NwJHDIQv/TOUzJWUNQojoW8MDz68svtIvZyoIhTX9vvxCDlonewBE5q3SspGmeeozvc/Ou3OkAAABDAZ+makR/h7DkMU2l6PKcHO/zShtHBIsYylwPXYPe5TqpRiePe1Ts30NzxAzEqW4PxaI5KV/5p/0v4lsSu94p/NesEAAAAMlBm6tJ4QpSZTAjf/1jexkZYO6eHpSGJJB5secePzct/DA48qcDfs+3F2H5KAOIzt2Dgt7ew96sGOjWUYhosJfHZkPXPvBEDaNms1hmlF9yL7QrVVstFpOhrOq4+qlmyqh/ldeLSeBIR0MDaNW2T6lc8cDTyqgcYV+Id6ev9tWnX/Pg/+/3Bx9Euhw2lfQiFdanPgzmX/55qCVEPMNFYpIBwPcaLU3yfrivjcJks8gRYsIJlknW6guCjN+neZa8hQkWtZwBla0sYTkAAAA/QZ/JRTRMn4fNy4UC4h+oc2BifKvWVgEVteOt6GdcPcqYcktVsScBfm/exmQJXfl0JC/xI/s+EI2ZfJUQBohxAAAANQGf6HREf4ew+GS+7UpFaHk/4Emf158OcWu01ywvQyMU5k4T2GGCK0QXFWRrnznHAUmOJaj8AAAAMwGf6mpEf5Kigzx6E13On1ZVie1EO84nBtPPa0VWD3aS/Ddk0ZkbJz8T95L6crCgsixKpAAAAGVBm+5JqEFomUwI3/vgOOEGsDnbEfsgdrJ7tKwX8nHR83gyULQVhL5x+AgWYfj7sj53F1FklU/dfB0NntEcc8LAeFRHrF4HF/lXkKKa3TzSgRuUANQcRW1CFfPMX5gBxfeWFsPEkQAAAFlBngxFESzffU2ptw5lYiNEGv798jp2yP5EhQTY9tWckQ2I9YpqN5ieaWAqdGE6ASw40+H7sZMawfHoCoDply5HZaD4Q9TdMJXThHdVh2fQNrSDXLaofGcFLAAAACYBni1qRH+BB5WyGKseqwNDLE8nxiM3gCXD2Qs8x6pTnxeKoYIGqQAAAHdBmjJJqEFsmUwI38oIOGvYtN1rSBTzyT+MoTaZxH/5QP0RwoqkIo8BIYjA+yAn/0EDCt8e4sfiSZAjYvAKxVNatVPYvX26/q3BEA+lCvg82wYomT1/sm0oujidy0OvyL0i4I1stZMMRcaxZssLhgLuUf5qu0yhYAAAAGNBnlBFFSyfhIT8MY1eZNHsRSE0h+4P5gINzsuqj6GnC3epFTfnf17zrDrtInbEItFMN401WbhoZ5/w3Ux2r08lE8/dDeO7O+HJXJlcFYNAFuyz8daGCaCCNey5dBaMTuealfcAAAAxAZ5vdER/h6Avf/WN8x6dT4hVMDgeRaCIEr6D0AY2ibt4Z5BJQpnXAfZG9fF7OQKPEwAAACoBnnFqRH+Gdbgx/3bkhObG4+6VBRfnLnVvz5ibbBPbDI7nFtBHdbdYGYEAAAB+QZp0SahBbJlMFExv/R5j/x5hxQqCYDiX9FqzOLs9ly1QT91Cp7pRDMF2Jkxd2wtt+BpFU6ZBrmxwLfJiwLCcv9T3lq9NZ2wLCT+/wnYYaZFCO5IgouVXlejqeKxfunKWRud5TjL7EHetewgMnoQTi+FFP1HpNOyNxxWv3UwhAAAALAGek2pEf4EHlYPO0DFagf010Xgc4NtyHA6rTfxLOyTg0Ox5hTe92mTjN7boAAAAg0GalUnhClJlMCN//SaD/44ef0/kxstvutqsDDhlXYfdekHXbmlDX82+t1ItqO3M/ML44EVCP/8ffQY6CAnVY1XBN5ZbnYha3+46HPtD5SMAfNYPV3RND0lVvOS15Fm1CdExPBODnBhUkaWtEAbcewQEoOOq2Hen6odimOMzFQ5d9rZgAAAAg0GauUnhDomUwI3//WKRu6nVyVxRee7XpiIxlKBqJvLZu3azo2YYPGF8qjUhlKiymApjcRRZlTf7Nef6tukffI6Kd7PrbAcHwGvELaeaP0ibxo0QBjNurmKmZaKkIE6v7wE4Wpmgpar7VBmLnzu8twOoVcMvF0fNF0fSaxrmXamWCb3xAAAAXEGe10URPJ99biH2oflZsxy6BAxf/qvJQn0UmNGcVBNMrhzZROK48G5OjPWtBqH8efnPpu4XSt3r1R6HAtsxDtSSxVrh6V0/kLJR05IiNa/ZWgvQYnGCF9qmzKBvAAAARAGe9nREf4l57VUKFwSmPSmlWex1Pm065LEm1OKxbJb7BMmnONzBDc+f4xWPkczvzmTohRN3zIL6Ojmccwe1SM5TTAzFAAAAPAGe+GpEf4TUK0MsYlIRurtl/BL2J+PMXABKR7W0r/E6t29wkdoiEyOny9hdlQSa5XVyIlDO89Cpm0EBTQAAAGNBmvtJqEFomUwU8b/8A8MxwlzCHtLBjwzZdybRP3L28+w3h5H9q83W+4R7Iqot5pTdotVf866A50TMubMMKbl+05ylorAZNPpFCLdQn1pTLuSgnmK11G/Uu/D96ElwYRQfS58AAAAzAZ8aakR/h7DkMUOCGBJh+lPcfvBfkdYcocwBe9qmP8DL4C3VmEv6j6IDX8OhLFl8JqBdAAAAsEGbH0nhClJlMCN/++FR2Inm+ENt5odDIwW/GJwkPcJuckK0JIaAU3CdS/tVTXZvLR05nQOjqvn8+KStPWBFi3Gez7OgeSNwcVElqJE1rP7gdMNhbcN68HuvsKSeQvpaJEnYw8wvyoh/QW1CYyoGlOno9K1aJIr9Pk0Dj8g8+NGIQjaqBemVGjXOrl9UV8iu8KodQ7jGbm1Zq/jwidOehhp797VsqpvsPYf6ZBDtDwWEAAAAYkGfPUU0TJ97IKea4yjH9NU+Ki3CHMkcSJ9htM0rapvgZK16qh6u5dqOGGAOrakaQCu114AMCSgb/V52v4CxfpXrn5qv0WiKgUPDhAE/bLxIRa7EHZrz32McwKF2qrLBShrvAAAAVQGfXHREf4ew+GSYyDhPzPlV2bRuCDR/Fea7xly0Zv+RhCiwpzapFX2nVaD4D0ZNqSzT0q6pbaI4DoNcxC/GWSw4XfnszHz8xeBlADF65GFV4DR+M3cAAAAoAZ9eakR/h417ZNpp1xtu+fafR+HC0KcS40oDskcFqn/O0sQuxpu7lwAAAK1Bm0NJqEFomUwI3/vBm7tPFg9jGVWovZ0jmu/rAuAvcW7ubK01SqdOcWBIz02mXZObB/CLB6HlDKDI+VlNhCg/6zrUx+GcHA2xWFHBAfdbhNgkXIqJBYT8FRKoB8Q1unN40IlItyXayDFwzBeGhvGI/KN3ea1R5uXGk1x7p1Y9/VGM6iUzPd/IwqqDaaStukEQBdqZPJMfq1T+2RPjqIFRCfd5FvxLcBRYuHwYlgAAAFFBn2FFESyfeyCltp/KzZjl0CBi/IEonyzMc3bZzsIWT4CydiIbfPye6ei9io+At/mFoyePgIN2xEzlQ9KswTRayQNCfgryEsw6EmQWr1T0ADMAAABBAZ+AdER/iLlSgg3Uqa5kstRvP/YJiqA7KR90cw3LNvcyMbNM10dPyQ7u7ZQ7Pxn1vA7iWSg71Q2/hzofKx8Il0AAAAA3AZ+CakR/hNQrVcoxKlrmAEcm59sbNzD40gdHC0EziACGmA+UjRF1XzhaI48Z56F0dPEF2w0kgQAAAGRBm4VJqEFsmUwUTG/77lcOEuaZPnc+jWH7FzzFeNiG18puHuQh5C/ueG8rR5oKuQk08AWgzZtU8fG+9U3UaZPtsL8BReEOwTp6VR6Dx/WI+9cQyMaZkUwAXj5v/UdnByly7VU5AAAAOAGfpGpEf4ew5DFEIIgTtjL/DAkM6khmBhIgRPsTh1KkB2WPl3dHMh6FUA7iA9fty/cEHvdpGr2AAAAApEGbqUnhClJlMCN/++FR2JyyC0hpnDh7PIxPg6BMTLhBL1SK8o0CDfiMpJT4kRDl9tdZ62EidOxdhN0HnNX76b7eC69SCDzmuZwOgLR7EIjB4kgvtAS2pqCJfoNyhrvP0D+dIb+QRSH+77hmtbfE1K1QnwSPb3D56hVITusU1SKKSHGmPML+xyt5i7yztzLCuUxQ6GBMKC1lr+K5ma8CRUoFzeYFAAAAOUGfx0U0TJ9/cI6jT+mdVbvBbfqe0kP70WvSpDFVICoTveNkI3cjDw4DDjaPNgfOu7y52O5KIIit+AAAADwBn+Z0RH+Achjb2/3sCtNWp7oUGbK5uD2avL6oCzA40KUrmSqFJOSULMgMJ/dZcQX7xKN4kdd8aG2wieEAAAA2AZ/oakR/hNfANpNdzqjAmab/E8ax6Fvug8RPvYSSLNwtDqJx4K2Zi3etFf/xJFtiuMoa2OKpAAAAZ0Gb7UmoQWiZTAjf+8Gbu08WD2MZVh8WrSpRiUBZUWy+1Z3OOJ3h3uYEjEFhMU8BzaG88lQfmplKnzW9sLz44+YZ5qjPn5f/4kD07E+kjAoFn0kEueWpFWqaK18eV3GeAKJtOJ2gcdIAAABLQZ4LRREsn31uIfah+Vvv5ElY/3/IpzyyoeDu2znc3BzxpsJ+RQ8j5+UBxv8Bc+yc4/gYr85cjlE67KN8CH6/crTj4W8FigNzW7tQAAAANwGeKnREf4Bw4zBVesMI5kvH//3lWiD0u7E7ycYGE39qDZBe9Ez4XiaOH2ZBiU+cYRKmcZPbFu0AAAAsAZ4sakR/hNQrNtFc9WisO+KwtEafX8vb6ZXxEB5GdrbfwSZkZhFQA45ylQkAAAD2QZovSahBbJlMFExv02+yx/CHKbBQXEKSiQ8/bO2xeS1HQBGRs5Hf/V0hpgC19PEbrC94YtsEONo7PCTkPWJcLGlhvRVXDsvhZchHblbTte0GhtaJBAsYRQtNVaJ8kzaeIHkxTXGujtsMr89Zq9hCEObGEdp9ad+G+69ZLh+FJjSEVgOkHkH0PEpLJYyFAvBlFP2jIHUVm3qm3jb97P4pJCxeSwXqGAfuihTU/RGSLNPwJun8vR1Ksa+vR5Zon3K2Nw9i2s9ZxgoBryAZB90K/m6pC4pSyLW9BsQ5Qz/96LFgNPiTaKC6ucaXERHMEYPOW9aNn7YVAAAAPAGeTmpEf9wncM1iQ9Gl0i1OL+ABc9ME6axttF81n1bVRz8hEXRQGH4rWzfbm5v3nB61lX8Ctgc7vz9QgAAAAM1BmlFJ4QpSZTBSxv/CIPOnQv6FGt8l3K0kX55Ov5JRlQ03g0N8vQOD5sBbfTPpEWH0nyx5HL+fie+qID6xiOC1xEwq8zJdI2VFblTxxBMPSypkRcUoWHmHXCxUKRVZYFMjeGgiRV+6kZS1BgU0P9oD/Oidv7irP/aV0V26J1DBG2+EtBjS4lHyivwXdip0d1BQfBGsk1ugFW2Fy8/EjyelHeTPHBgZ943+z6bcQ1EJkW0tMaad2fGspFHI0IbGKScCPWX/jQ6iT/ojIy4xAAAAOQGecGpEf83u2UnV+7UnY1O5/oQn/Y6AF0MpXYTzrVvtUezCnzx+JtHpSvytH38dTVSqrrE8zYKZ6QAAAHhBmnNJ4Q6JlMFExv/0LB2GXt39+D3Cj7aj3DPZ6Pd+InbEtW1V3J2UHsXT1v31OO4EXiBsko5smAXJphYrW0fTkQJirwR8BXL92Ei4QvU8E/9ALgLDiPsylWFhiHKkipiOyh9H2qwcv9t+CUifmj23yzrIU/BAu0kAAAAkAZ6SakR/iOA7CrirItMfJpRo9cnv5zDAP6itv+Og6IlbYQGxAAAAnEGalknhDyZTAjf/yfc+3zQHuCTSl/U7GZJy/gS5f9Cy2h3VQBHdy17hrU6xaDgmNvuPdO+3nvolJXmlwIAkxED8Lx4XeaN4VE46PRHNsXwgF9trLiwzDBPhYZ8V2tPERnsdMtjrFEmrEm+drqLgNlW/gCWzTU6jf35hzyjGb6uEJOQlcGPQpDxqrtzr8Nc8vp4B1BYjdvjrlCrq/gAAAENBnrRFETzf0DZ7oRVytJSzitxXMmJRFhcbeha5ual5F1I2q2YB+zAT4PLPb+JmJ0ARNc/AIxeBO0WVPQWFlcQBs+2QAAAALgGe1WpEf4EHlcUhXPIwsFKRCeT4xUX4F3Q/slkJ9DbNgs6Z0J5ivSJcp3D4p0EAAABWQZrZSahBaJlMCI8bzhd8dK3gZ46kOJN672WFTGEpwGtbDZDet5dEQMHNg5ySdNRqv151JHg5eF5RqsxWW7PejJfmZTCQ9Q1IN/bi5e+xWsYriMFYv8AAAAA5QZ73RREsR4etQctl8Y8CYRFjDEwJCzxPJgJwzTf+tolyI5ByCNErI5NnAqLN6VSiQ3zkMiK+t/rdAAAAKwGfGGpEf4EHi7cJ0aWsBD/fwiVyP3OFrw7+rNjWzuM2AuJ/SSLVjrYN/7EAAAT1ZYiCAAn/vgy3mIDfV2I+c4OQeUCbjkIqauQ7q1hI51t1KK0Ty/hpqe1v3iAKauIQOhjuzFxyMgx7FbnM/2jBWLSkwzPM9uWS0H4jWSakogzfe9txEs4c+/QsLeTPvLzlrrhoRvk76vX6F462KFV4lQ8BLFX1JpZOU12xKsmCPjGaaOJgyXQiOhYn2hWzgcqRNjGX+X2BzGn8mEJ7aGMcvtMYUAkI8JdRks13nKu9bo7rrtiKbOLdcmLQMypUndSoHj4HIx0L94EanyxNktjsUH6UEl/L/6OZSnHAINvBlAMWXlIWyt38JXD0JUnOv+PLJ9lXJviKQHF0pLbebPoF7UwMYacH5I9nuT2nhJTj4jMvriexQBuncAXjupdykvxkUgQsRr6KKam+klVKsKAJwzYn02pUzn8ZX99Kgp0Zv6A/4AZIw+tqzSXaBRvQyPlagCtwxNENMBXBGJCaTT/M3d1NYKMBuEQdSzG/ZYjE1JEavYncvI+HARC2j+WXM1aZgz9ZL8CgLjpcv4/YrSBFGs6H2gmgzdANfJkk5Ly8dYyGKrEOA+lfembLO2OFLdnp0Kca5ZBG4NPOmAF4BBvDzXD767tWmEsGNf2xMnrIGj+8i72JWEhZfuKA2OO58X8ZDSx+O1JqcX7S0VhG6hgbqU5ktHE9a9+Ve9c4kGo1CPjPl8CTm4FxZbQjzsmwP6HqcIFYBDPVBZ9WoQ+2D7QvXIfv5HX1K2NuPVZNm9uwFtcmGsf6PuOhqK1pWSzqCgCmwApkdzlejUme6mZrR23REJ3PkQdaAyFo3gfa5LdW7PmNx7jcVqBgkY/SNldaGv+tW+xdJQrKAgqxhGSxiG5fPPeb6tkPlvE3G0CAEuGghWptSDagus6PiEnU/3yKJHGviaPLkzl10/T0zWMpIJouNnIfO3N5LeJK2C4xONQHkEUvqEC0YiEq/qK/U66LBsyfPc3Q2VPtQVRJkl/If0X6ijW2cTcVVwN2xTAqMA2dSbXoMPK3bH/VU7nJqjh5LpA1vIdkAYZ7+xljSLiYHT9oE2LHNhfDNZ2s8GSgeEYKD1VirZxu1ZbwqSl6Xuire6yZL+zwkbGlmIlsAcwVhhOfhAfiaIh4/+ogejQgeDyfWoDmEbXzSAXUx7RmQg8IZIWOZzcuX2GexReoXjhzAe2YabjMhqvs+FSJ4b2KP30cggQmhKlwsFqD1ev4SKHH0Ge/jrjEKbN1HzQ3oLtBg4GpkYoRNfImFVoXHLFg7l2LgETKlwJhboxTHUPGneYTP7AQ+K5Eb0/uCSO5Ymq3DzI1gmuGU4MVmzQSuYm/yP8e0MvRQ04RZiZtS50qyxWMrIOnk+IJ7NV3E0pPM3+ZSPpWftqVuEKawLfKtVnt/jTfRjniO7HM6pAY//32DoinfYKZ1Q8GPwjncAQdzvOS+csasvxPzi3DIy77DGfgyPztLfczbeCwixhep+9Vqr31F49wN+AMxHwVCsOJIRkR/wOscgZTuowD55INTwrue5rjZiT6CVfcKanAArmyYb3q8SvJB+n+cPqmbhydgqPgtkHij+7EGXRHJg+g2fpEgj5bDzmfXOTO2YG2fuHCeugHQkpDQ4gK1xMFUb1525vNMHHp0iO6hZxgePKmk6xBwKiD+jgXwh9euxUHrTJXmXSWmplATtcnNxFTI6sQuXnJEZ2JpBKsD+2BAAAAP0GaIWxG//vEnVTkBBzwgdJ0ZWNF1Rhjd1kKxuETzl8V2+eU76R66MAyhvqkmnkxp0vkl/+a6/BSJxSj4/j6swAAADJBmkI8IZMphG/74DjgQ7ZH1Lrefmefbtk4wS5Y2/x+DLBweoV7XRDx2U0ZTPS42sTXUAAAAF5BmmNJ4Q8mUwI3//vBGT5YQIcSq44tpvwtmFs1bxuNY3/v09IKvTPLNo54VVGqqjoQ4KZPNJThwuSr0wg0X709s5PGIgRr0KCsOg/kGruNjLFSrx/JisqEekF8neXKAAAAm0GahUnhDyZTBRE8b/vgOOB3OQFc9AXcbpj46i15PbP+Q9GQ2tJqySrnQwI7VJls29RfMrHbfsSK8eX/l7+18PjO5gb4bIXLCoP9kNhUaOMVrscc/CQyFS2DmIIbp/qLfqVQA7JSH5O/IhV43txHuIs879XU2UcXwNJG+MNUv/9SvID6qp3sUX+AyNp07MsYGwo+5nUN8BjKM9iPAAAAMgGepGpEf4exE2/rMWYsUC9Nc+P/yRU8UOqE44EY/0tBUFtXD6vVpxSXjArlMIURYAeAAAAAXEGapknhDyZTAjf/++A44Ilb6dinlSkT63y1KiWQ/IWMp8xWpc99kAKnB2h7AefDWJA658UxgGiN2hdpqZnrxVEyqj01yCAccBMH+y2STVmGjawy5zlKkJY8a3FAAAAAckGayEnhDyZTBRE8b/wDoisgQ2MspyG9MTL2twAnx4YO3N7dBgMj53wgsoD7fgMzpRUSBvxUyvU2F3yMVPcr9aQehYe7tuqCXqk4O3XxyB/oozQG6cLwffJ31HMNBW5ShOThCQwSZuY7d4iRHTfMsm/D2QAAADkBnudqRH+HoC9/9Y3zHp1PiKWpA3D3YoIQK4j0M6XxrUNNDfPHjMjM+xDhyvcRekwlGLKw4c3efGIAAAA8QZrsSeEPJlMCN//8AtC7wRdgOIxYG6HICf0tt8EYRWGccTPXL9ID3ANrHXjlUlFMeXfH4RGJz/GNKEO5AAAAOUGfCkURPJ99K3q+WhOPVadJ0Zn+q4gSRe2/zKKvWouYJMdACf73czGhSzAbCZzI2xgFyyWnZjV5cQAAACkBnyl0RH+DDJVK3gapAzXdW8ZNJVvmI9scX8y+ulkPHeyLSaA+eDUERAAAACEBnytqRH+BB5WDztAxKT4xFpAt8qSdgkigXdBpZgOPffcAAAB8QZstSahBaJlMCN/7wetWHYNorkyweRiHOPPdcXS1IITtq7AezlT6fweBHbowbC7asO2PF2/UBCbMgGRBUVHraspQzmQA6D7PUakKIGCmWZK+kVgi7Knv/moFsXlEd8QJPGaF0SHBN4ELaLTmKwUa7PPSxV2g8M1cOkFDggAAAGBBm05J4QpSZTAjf/vgPUSLvylvSJaoX67yRc1fALtzOq+SiE4b8qGjclrzocoYChnLAMbHgIHZNX73/lsLDOMLY4x13csDW904+9bFMORIwShS9SrF2hIu4YTLBzKObT8AAABsQZtySeEOiZTAjf/8AoC0REXf8qqj003CfEolq5ZHPeIt/nTlRdl6OlOcoVNgvTKHDmyYVBe5RD5gFrfZYd2HkGO874iVFT8MXHBHECxLA88R/Dor69i7i2G9sMY7GGZAq/0gNA9/XbGiCiL1AAAAL0GfkEURPJ99PbKAyIZQ79YXXIGQ9Zn3sXohriIew4GeDzxFPtkfYMRAykrz2oK+AAAAJAGfr3REf4BsKuKQrvnfCNbyOxT9/LxaMU530JspiBe0HCmuIAAAAEYBn7FqRH+HoC9/43g3bWZ3P6qT2ZqH2A9FOFkkCa9OSnFIbR6HJHG3QgSA//4qcrFGX1dvdEywSR/smSgEwGCmbp95ZNSBAAAAXEGbtkmoQWiZTAjf++s+OA7zHfcTr2WCBetLHUxD7/WNcoHXS1dnHb6Gxii8D96x0pb0d6PjC4DrB/WhVpxnxwTGDhj4JrHO2bjjfg9OQ6Gfn1d+m/OSL0v95OrEAAAARUGf1EURLJ99PbS4SuQriYYXJQH/fi36a9f+dXe/nbPtQtUH3xb5/YCrJqm5SDKlJfww5Fp9cPYaad58GXqPyd3Np0omTwAAAC4Bn/N0RH+DDmAYlvA1SBmtJcJQ6ephfzmo9IkH4wmc2d24zXHcPLKh+Lnw0TzAAAAALgGf9WpEf4EHlYMkv9E2/PMkyg5XBF3hcZ4c77PtedhhUSll/MrYkqHYxHIF000AAAB7QZv3SahBbJlMCN/9dtVqlh0uAaN0ofVjxMhZ1h5qorEdatiQ783inAG2LXHB1MiAuY0iagySzluM1H3bIodNuek635EEPQurFICs0REgQRMJVUluT+40C3yNiRn0C82e0RNNjNdsjWBfmakWZtlwAWD/Sn+XVH0p6IfBAAAAZUGaGEnhClJlMCN/+/6N/UQffCtxxis5+gpt/z2N6/IRl2ZC+fkvRLr0komX3ouJlTkB4RiYtyBbQy8957b8kr1rd495KdRJpeJJwYRiht2s6g5ssfWHOxE2mlZuF4E9aBSO+RThAAAAYkGaOUnhDomUwI3//APDMcB3mPtur2L1aw7/475bzHjZPyheciacbIxHbZ0dVmofkE1MGw1npgtJAA+Op2CJi6U6FI+Fd6w/4z91fbaUzx1UjzZgEuMaupFEOYn3mjpLK++pAAAATEGaWknhDyZTAjf/++AcgTWtokUg3fR5ZBhOWluG1KszNLcF72sEzcEsRw+v6kcUufDwbKgmv+8HgVhfNY38jxmDgIOVqtV0QZJ9wkAAAABZQZp7SeEPJlMCN//7wZu7TxYECs364K1K9zOpGtt1ZZOODFKfndyW2uhIDiGAJZLhzB4/6bupiaM+sug/+57cXUnsnlte9DZAA+U0bGdXg2VWdQHryyM7bKQAAABZQZqcSeEPJlMCN//8AoCu8Eg0IAKdPalTJwL5k3brsQ/2kMhr5e/LMbxJi9LMuGNB+zmBqO816m+p0J54fRMFrIfwddbFhUF2M8i5i2yKkZU/u9ZCP8OKKucAAABmQZq9SeEPJlMCN//7xIHZpGQEF4WuTpc9YQoczh9M1Lmum+7ob3DdhJhUVeNksMkmbU/aTy524uFw6YjqhrhnvYI9JX5S+eulWATUNNzLVA87Eb09ZoT9jApBuZ93/euYFurmkkWAAAAAhEGa3knhDyZTAjf/++A44DvTrlMIVpBxjTybeO80fsQOd462lkQixl/i8ixhGZjblZ4BIJj50LvODBaMD4t74KP0xM1A0shJj3XDsIWRvMHUMW2oF+/62kSLqLr/NQ6ysmc/R19+j6tUcDuLJBR9j9Ks4GTX/cn3YO+gyTrCwbIURfvNuAAAAEdBmv9J4Q8mUwI3//vrPjgQlevXQoH45WptzQNYqW5lvUt7k370Zb40x7Lqnj6WT8/G239QKi+rn9bt99F7eH8NCe9EpQ0PgQAAADVBmwBJ4Q8mUwI3//vgHIE02SfYN8H05IGFKI/2FJ4/9xTz/6bnIt1458gQ4K84mOGQcC5N4QAAAGVBmyFJ4Q8mUwI3//vhUdg1hSSB1aabIlvddm+MomJ80yRI4f09f1qYnZFqKwTK1GeZgBhUqgBV1BeNxJBtS3WYBxDo0ZtkTRCh095TeDwxQPJ/8LwJ+/nKDXqEa7WrGfxoouefcQAAAE1Bm0JJ4Q8mUwI3//v+jf1EPtOhzZ10VMZoqZ57G/blO7Ho+Yi3NWGkKEaOWnKTMTQBZb/3jumdLhfCXmgoAlPMUT3exw2xGWNiadKznwAAAEZBm2NJ4Q8mUwI3//wbcNERCUal7f5h0QKVTPWNlgbjmkWJ7S2NKgIHwG0iGg0786GE3lcq/pK9y34kGTeDc/yr07QOL094AAAAPEGbhEnhDyZTAjf/++A44Ilb6h2GJcOd7Cm/Xs77EAaBBuyH+UJ495yfcTaGsKqPsfkXwqaBplYMShrGWQAAAFxBm6VJ4Q8mUwI3//vBm7tPFg9mTMd/orjDHC6R/TN14aq1uH6BfRoMmsnLIs+xQljFq9ImogOyRHLL//9dg0AOfyMm3nG6BsZY4fYp4Q3XOVG/nIPZmEHHeOzYPgAAAFRBm8ZJ4Q8mUwI3//wCgK7xHVn2P8TqIU579zJBs783creb+z3a5PkBHTlb/74v4uMpPCgMQyA5H8erUDcdIt1E6LwdC+Yd5xch2XzNIEenrNxxtoAAAABbQZvqSeEPJlMCN//77mtCCdsTY5Z3Geb7Pm/QsC8ZaS5D5HX+6B/kPwTX54IPJEXG0Z05F//tLrak3nsdfXudbo0PXPSotywcbWZJ7pa7XExYa26grBmfGE1IQQAAAENBnghFETyffT3E9CkDsyrDByDhrW4cvbhQ6Som9ctl2PhFMGZxVTqRHWTPWDxP4Sfm3HbhULO9Fy6NYBKZryCailXwAAAAJAGeJ3REf4egPNbtmUKQ0b3pPCD9eBXfRv59JqxVlyBWqnu08QAAACYBnilqRH+BB5WDztJcUhExA+6Zo3n44kMJuHtPfnXh1zDcqNof5QAAAIRBmi1JqEFomUwI3/vhUdicsgtcoZFZMaO9tuG7apZo1pIvnF/gM3FczW9nxAaKJeV2vj19Cv3exUNOijnfRiR9ASN9ltcz+VQ4/0i7jHGYF+KFZP3uGFYBHkT2QRzPu+40UwUjscvCxMulW7ycJ2gzsxyoAhycqC6uUVDHWkGCd55vkPwAAAAeQZ5LRREs34H4Vo/ACpoAhQy5Ym6KOlAwH1FD/kXhAAAANwGebGpEf4exE2/rMWYbvt3Gio//JFTxWQEj0jgsGN8zykk9TVdEaSZbn/AJPKDK7pTunVhX0/AAAAB+QZpxSahBbJlMCN/76z44IUf87NrRe74hd6WSQF0BeASvxdUxL4Rm4GnonDSgvlzazo7BX+ROHA3HCwCRipy2tg6EaaSo8weDziWONE8To2CI9Nf6xMIHsssxak/PU/YdCrDJs2119wdbr4Gb4bSqtsAe2S56lY/+8eHBMDPBAAAAUUGej0UVLJ99Pb91NJVBQqP2t/+IqmE8MrKhL6lXP6xi7XAoI/PH3EdbwEwYHQ13+d+UGm6YHZjB0tU+x5EgcHwOjl3Wv0NdZKEk8tYEJAq9JwAAADoBnq50RH+FEsrDaMSpjV8Ar03X7CI6tPTzHGiBpSWUkMJVOVPbCEUeYgsYcFzDQe+eqnhHQbxBqzvTAAAAVwGesGpEf4exElWAWL5mv1zn78AGfLg+vpjLffFPe0N1Gn3GyNbeSNBNOMwp3LVlKeiRguPo5EckDmXdrHphQLAExUYkI1WZ9328MC5f+Y/xxKuGNSXp9AAAAMxBmrVJqEFsmUwI3/1jexkZYO6f0SqUBJutOylzcq7633ejQM+YcsbI01FdCSdVMPBHTcsAfhFmyKZdJ9U2koiIqT1xC6nMTtNhYSNWKTjEkyshkiwoA2kvaff/xZ7MolmxHJAhqhjc60XuqzboSWvYiUYabtLNBXbUEc94ymgQJJgQehwjl+Rj3g8ihDXo5vs8YeeuDIHNpoqyq5Z0gfHVseBv/VNMJy2baCQS8R4kw7jDJW8iWXhZUAEJXBoXny6onzxbxU7N9r71XS0AAABDQZ7TRRUsn4fNy3GU/pr7hiYIKBgdsTtsrSiKt3zM/HqMjo9nD98MgGaBM9BOcraqeIPJTzn3j7MnOe8aMxA90X+twAAAAEIBnvJ0RH+HsPhkvu1J2NTuf6EJ/2wV7RDK/mYg08fzz6vEXYZTtgVzrhrYCI12GtLgxwIpISuyR2pURhW4rUvLdf0AAAA2AZ70akR/kFs5HqaE1ZwEakhKwmD4BSiyCsJ5I7i1b/JX/yAr2QptkrDRNOf2vgGvhvSIFmCfAAAAgUGa+EmoQWyZTAjf/WJuYK1IHlx90OrSnTNXeuEGfLuro9aOPtuWzfxGlUqLMokdEgNX5hs/sSGuavwQOvVjNR8vXHeA7drDwV4vSf8z11/Kgo3tdVMNLWwgwf/pF0Ahg+eMMrSN4mNib+1ERK71wdTSsh7YgqUhZ1sRulfT/DmgQQAAAE1BnxZFFSzfiNYJXadE4qxy7VpUy9Q41T0IwKvWOe6XdY5vDoVxuSNPv5DidAVc2oMbHfP990ARqTeT6SM/eIoqojUVbdfcQ/gdcgL8QQAAADIBnzdqRH+BB5XFIVzyB9wJv/Lfxs38Dbl+6UFytIOZY0QdJSWK3tXhSOzb4vSc0LZtnQAAAIhBmzxJqEFsmUwI3/121j7DkGICHPtqvASEuNUhdr93NZ3drLM1IHCZVXsfuAYpeOu4rIQCjNeDjn/nmhlPTKoORxu4iqjlk8//dcd5KvIs5DQidbJMVc7qc3oVJpbzwdr8dbcRFN+vvSDcc7Fqa18YeCHCZHZwXF3Cz+Dy0oArWOp8iRf6FgOJAAAAR0GfWkUVLJ99mioJbr7gTovvLkf/l/YtKEL2mu5oyOikfBlzmuZ6xLlP+smIZv0LBUks+sQjAEsE5q8iWpcQhDfktWHs348gAAAATwGfeXREf4pxDGIf1f+twr0oL2RF3LEOzI8akUIhkzGFQYj1qALSEqZaaagajuk+Xra1iK8MNJ545dD2pPxL+B29q8NRoc/1g0bDynEulFgAAABCAZ97akR/hnW4Mf915MgtV8Fyl+jyhJqaffOqPIliFZXeE+XQchMQ88eJgdqLi6y+ytVR4nyOKwv8jWcpSMzkR+NpAAAAZ0GbfkmoQWyZTBRMb/wDwzHA7nHfcLHdlhDBdxj1K2iflRpBQ54lGRaeTModRU5a8siLRIR5KHbbhYv4ItU1x58f1xPglHmsDhwytPRmDnxl0gvYO3B8/ds3Yg+ydO+qQwyKDZQ62qAAAAAhAZ+dakR/gQeVg87QMGmP7dx9N0u/+sitdubSF+CATvuAAAAAf0Gbn0nhClJlMCN/+8dP0UsID4MvVXMepTuAd92+Z0g5Hm6laabBXW3H6+DO83pTGwIUl3MsSoPcrWr246Y/Qww7pWgjMh3A6DDXOd05ptlNb8IL/CBYCg/64pd6ZHCzVpOaP/3b5cwCm0Sn5E7yOZ/JXi1D/JAEeBMIy/0ioB0AAADkQZujSeEOiZTAjf/J/Y+WDSC6t+DsFl+KShs2+BRlWjHwkCbRkfk5F1YAt0JG+eWgtV3xfaTFPYBv+mxEEdHUfsAnOpwhTvDIVMNDKRQeXrvoZCgkP6BOglvBRcnVRm/QCrn8gMOw51JbsPpVerALbca6HHK4IkOueWgXj8/rw/Gq7hL81TYpKHnCgsWhhNql0edQUpWAf3cM29hrgBSKwsy3jQs6DrK17B584Fz6wiEv//JAFYwi3YgyIjyLVn3WboOj/oZw3ulrAwtsHz0duObfFaFbLBt6sED+kDCfEscSq2hhAAAASUGfwUURPJ/Tvx2ZXIhlDv4p8IgQMX5AlE+WsCstVw0AQsep3TWwxt8vGXFQoms3Q7j6KFLJHsNyi/X2PtL8B13mVJbuvM9va1sAAABDAZ/gdER/gHDjMGGbkbroMeA/+R+98sKd9OE/osjRzxqU4fU5myqG9RgJ9RKLHL1WRSCFE5PdzrAJ6RnJAR/133tzwAAAAEUBn+JqRH/XxZ2/arlGJHqHMGb4IS9k1LB2PTOkX1omkDNqjv/Hyd0h1US0g+kY2+XgTJ0qtQYnHLVn+AHiS4iyJNOTtEAAAACwQZvlSahBaJlMFPG/9BI9qA10G55ifcR7emNzO/s5HqLdZT5FQFk+OcoHWoo43SEpEIzC+YKxdwwVhbLtnL6xRqv3B4n9IDF9aNGjQC3lBkrevbFcw6jXvPXdk7lFbzKvLaUj6IFJ440Id8uOH8S2xSDBjkuw7fw7kuVtpJQ71cMUX6WsUmpYvzm182OmweSjE2W6mQahX5TDXHtuYFW+s7wqQG7mczW59/E+jHqj+GEAAAA4AZ4EakR/mJmxUG34cEMCTD9KfzSlBy5J63PV9xDXOTeu/CSqfUIWFTARN80zbyy1JouE1/9W5KAAAABpQZoHSeEKUmUwUsb/9Cwdx7ECFH/O6zLxukzGly8jJexJ12NnWZVyEH6v2GpBHI24QsdFluV6MU99hhAQ/ORGnQtP16u598A33ALGAhCiT6Fkia5Y+ko5b6BByfi+fm1pz4EM+HBx3ixAAAAAMQGeJmpEf4exElWANW9/woTjJTRwrYLgW5TCxhNGjZKsQXGThPJJ3FtL1nkHDBC/zEkAAABbQZopSeEOiZTBRMb/++FR2EKSC0QlA46ppypN08vPo2SkMTwZmYivCQot+O022lnoqxQWOoCxOpu+brG/vq1OiqImEEUqH1aSs5FB5wZJH0p8pceogOzxqbSFCAAAACEBnkhqRH+E1CtCa7nT6sPwHk1BQYUkaWxgrGWqbiIHMPMAAACMQZpNSeEPJlMCN//74Djgd5CY+5YA50g7ZWcl3e6HlM7+829ssvr2kKYX39Qy0iHGV9xtA5s2X+EIG8ksj90L//eHWxQYDvlTbZLHSXrXkPuNSEsLXk5NQmC2u4ggLAcIDZeLlKdlT77ixCPNy4BIJtROhSq40uHLaCpv/QYI7EgtcO9f8Zyycfr2+0EAAAA3QZ5rRRE8n360LEl02ntKg9DJ1qf0Sq2mRcCNvdMAMT2Deu+GU2vzNalSElMySEM6p0O8C0Zo8AAAADkBnop0RH+AbCj0Y6pj2ya9J7lxHir9K6838HLwTqvrUGcUcg/6mmd9tmwJ2olawiUUSVMN0KYF8/cAAAAhAZ6MakR/hNQrVcoxKlr4Jc/+7NWRzf8YKIY6R80MIEDgAAAAkEGaj0moQWiZTBTxv8fOfWo2KDtv87okIotgJjCiYB6nOqRdka6T2FSgrP2m37XMvHZdI8A21zD1pYWm/KGf+qjCETUhQS1RRxjZtApAg2eNLrsSASJT/M/z6WcXboIDcyflWuaJGh/CcbDrACgNgPWBxQ9Cy4GZ1z2sg/n7DG834nDz1pRrY0cd7DCoe+kfoQAAADMBnq5qRH/N1V08EzDwtoIri/39PtNtN5FyTk7/bpdx9BTbj9S5CIuX3khlff7kMGMVX4EAAACFQZqzSeEKUmUwI3/9Y3FKc0qbE9JKkUV9XeNE8I6QHJCzI1xykL6LJyJn8sMTobA325dt1+ho1fe5CgSyqyKmWLEtqPoxN/8FgOH7HVcYanqHPuM2LmQI6mevX4hqU31p6zFddFY3jjC9SJd2W8DUlE/UMqWtTbO/zdKi6N77GCUpGa8akAAAACFBntFFNEyfex1FrjBvpN37qK/wIQZCAFBrCFmyElK/aOAAAAAvAZ7wdER/h7D4ZJjIOAubxKBfeUw6T8sW9NMSWDaQ0+veb/Jd0BOAnj3BtP+dOcEAAAAaAZ7yakR/iOA7CrB7i/Mz4jo97V0uGcnVWvAAAABdQZr0SahBaJlMCN/7/o39RD7fvlTWJwGXUUqvm05+1QVmloNO9nFYywACdKqyxRSD8AZPjb9mbnY8gjIvuePK0VaUznabobdf+xj0P3e7kDpR3NBMu5vVSLNUN/ElAAAAZ0GbFUnhClJlMCN//APDMcDucj7DVPmcuWIK5vVjxlhCepo4trDOkrTmEGu/ZySN/QsIlFHz/yEpWOSVZ8qYXpP5zoMHCuoHGjtJgUJ9YY/EQHTth3rvh5r0lUIjrN6aJq1qPEo0fVYAAABHQZs2SeEOiZTAjf/74ByCXW+nW+BxzvYU369nej5XDAX92W9IvylVTvLOZPNdVa3B4/9OBz/yENcbuBsBlmdr/FOsXai19kEAAABQQZtaSeEPJlMCN//7wcOsd4IOcpHAOYB35obyD9fNfDE0mDNzUbz/oQknAvvZ4dBtywxoOtUrwhbgoCnEAdmdPKWcS/ilLOLAuTdnno4dinEAAABFQZ94RRE8n360LElvuh+CIdT9OlDp/0q7PJ+MJgJbrVboZA4kHTJ3ToffbZ9DMO5KtxdcbtUFSeygBOcNgX5zGmhe9rpJAAAANgGfl3REf4egL3/U68Y/3GAEKHH5y9IApJhkFHkWHqclHfOYIjpgTUXMpdkmgPno80wnwLqyZwAAACoBn5lqRH+Gdbgx/3baE/Aqppa+0jtFnUyKiy+Zwy1ZHyz9fyMNwfsj+4AAAAA2QZubSahBaJlMCN/9dtmbcgHeY/EYukQ7aaclKiAhW7LP0wcwt78GzVH9jEX+ul4MCqF/2B0EAAAAQ0GbvEnhClJlMCN/++A44EO2R9S/OhWXCYYnfMvGSfmF3VhHoRAookQ/B0zybEDHraNElAkFmtkn7maAS3rK5ywR/4EAAABvQZvdSeEOiZTAjf/7x0/RSwgPcTDLgPuFlZF5qsphrzQTuuhcRXT5hraAgVTlKSofuENEmB9xuPyUnh4xra8xjPnk0cC9TqCjw8oFY3hxp8p6Yv1rzk7Q/KA8AKgsEUTklNSbcJitU+GRzJIOetTuAAAAlEGb4UnhDyZTAjf/yghA8yNyXL4I6H5lYLZt7x9sYonEZViic8JINT2N7pAvBTCcrMb9wQzCFD0BtXGxWf5Rsx/8l+VOmJ7mN3Vdkcll0IUwspcRV/reK8gEFqp572XghaC8+bGg1kDUICNBmBkZxBrR0PZDAzSKR/HGRfU3wFVnoidVdywk6QM8o3fSqG0jFI6YMWAAAABdQZ4fRRE8n309hYNYZQ8K1MQRY/3//frBdkHLgzarcM4u6jySXsngsvUyy6EI7lO5ODiPJckJ9lqLmYq9/4Aow4+wuMKOMB1SJrQK2t0CbZ91mEJWOzNjbZQ4RrPhAAAALgGePnREf4ew+VpnMlKRTlxMhp2z/jXp9OUQEWoWTWTddscAANjrykkhGK8n8fEAAAAuAZ4gakR/hNQrVcoxKmNXwBAAZe2CmhnDxHLa4k62bX3bzPVMZ+4eeevEvCRnSQAAAFxBmiNJqEFomUwU8b/8A8MxwISvXq6TvqNm1NMGbNBYU3BnWrbyWWdLxOVYy1WXkGewPWIzm5q7yUvGhk/YnyoDdKJ8CrXR3ARnW/KIQtqUzAbH4jz9hM3MiI/EigAAADYBnkJqRH+HsOQxTh18+g5ZBI+1d8/B/i7BlQB+WddgPTvIA214hf/1CtuC75Mbg76dyaxml8AAAACrQZpHSeEKUmUwI3/7wRk+WEBIIWYIPFZOw56t33hsE66tFa3ZLekWImE4LvDK36TKHAgCfZYuTC9EbJYEL7N1Y6TRBVtzM2ya6R2+ErPXB2jjcjTa8RL4ox+ohOuPW7UZIiqbV/BEGgQbrsylo7OhkBXa3M4cx81PtPa5oTEEJZAzD14sh2F8HgCI7sfH9KNgsLmzFk/eBzzI7emYM5SDkSldP7xc0a0irmWBAAAASUGeZUU0TJ99EZTLsyDpEIRrKJpnoG9EXvvB3xvGEbs61IZFsRW5tRZxKSoUAv2FlyzhGJgURfX95iQnOI5xMTYX4Tm9fdpIYHAAAABCAZ6EdER/h7D4ZJjIOE7mYezedRutYU/Lr+V1IdLd155NZOE9hpCyg0nHA4eDwBg282Ojm5gxupgyr7ONyBfl+pZAAAAAKQGehmpEf4jgOwqq8yvLtAmySKIuJcaTjoN//CBYJDzeFU/Xqm2khRqBAAAAYUGaiEmoQWiZTAjf/XbliF2xAJMQm7WeNk8OMW3rdjdKBNxob1e/+NepMrYVL0nMuodAkrIMadu5mJBy5uvlxvNQp7VxX7xQevbsT781icpake6htsC9t2XleFcTwqueItgAAABQQZqpSeEKUmUwI3/74D1EHSkIO4Z9Wvyhvb+EqEIAwhAaGZld9PgYxqFm6CIw70d3crTbR84Rl1QSjXA0mnDPfQOFBboYhbHq4C0vpJXMNo0AAABQQZrKSeEOiZTAjf/74ByBB09qkSKnA0a7ky+N/3pDcHVcDQBwXvawTNvvn2svBAJZAi7D6VbLM9Eih/kRceBxSGO9APDLO5zFex2ykGUOA4EAAAByQZrrSeEPJlMCN//7wZu7TxYHfvdr3YChXTkbqSdy1lvNC4XiN9KR1PqesySw+ABEfzbkpLtzgLfVv+d1ee4PGmvseKH5sGEQg1/DG7yZPZbZVEiMJvrtqcWWCmcntEA1XNbsrhVN9bi050eM8WOF6MnOAAAAYUGbDEnhDyZTAjf/+8OhtVTlCNvyUsskUPEVWUnL9c90uNQCPoB4rg9+sYpzZa6v16m/de0qMOA/BNngL8GoEKLLIlgO5+X/dNSJs7hCfm/5KPsbwl8m3mNvXuL195ed9WEAAABaQZstSeEPJlMCN//7wc+0/eDsJNVZQtMJRUijWyOQP+twyrFR67F8Pj5Yf/Gi52beGXbajrcE2ZHcat654rgyJ1pKhO5Uy/gC8B+G9/OVL6pzknCH54OLfIXDAAAAUkGbT0nhDyZTBRE8b/vgOOBCV69dVilfxFNE78xMP/CEHWNxHiUnX6DyHq+qeebHfMqief3NvEbCPaYjoYHVRYSzyLFgSrimfr+f4tOe9AjjPHEAAAA4AZ9uakR/h7ESVYA1XP+1Kfs2fuAWfT1RiJtfaEy2T44UcRtUezFEZUPv4YZsGo4J8wK41PBEKIEAAAA0QZtwSeEPJlMCN//74DjgQ7aLyl6IAVFCZ93OMgKO7RPkPiBlzxXr1IPdsldBX9MhV6aIfAAAAFRBm5FJ4Q8mUwI3//vB61Ydg2jrrF0XGCsxft8+s5TpCXAiRHLtzCBWGDzOIi2fFiTNEfqmbdWsGe9ejcTntioJKJzrv8cb/s/5QolStYE3YQp5+/8AAAB+QZuySeEPJlMCN//7/o344HrOM5ff/3Mc7c1rlUvDf+1BYAPMg/BpElFxK2uuwTvrLu3L9d3CeuKWIiq2NiwAKzpkZXVwYoWgYQjBJkeY/r7x0HhuZX04/RYn2nHmTL/bcfYO6cjHT3FMoIt8A/Veo+J36ssJSsTbgb64suXdAAAATkGb00nhDyZTAjf/++A9RDRLfNfS154ZVlkmiUy6/g+2oGU7FxgxqSr5jI8p2G+RtB4SgTfOQhgR/tAQN0eXRO5enuY2bDLAZlbho5L7gAAAADdBm/RJ4Q8mUwI3//vgOOCJXWMSryw4h9E3CB6A0Ke4+kyCv0XYLF03+9oB9Vjy6J3L026VqcJXAAAAU0GaGEnhDyZTAjf/+/2woADtubdqT10cQKSKXeNNuaLX2/j7JafRM2QBPoIzUNui7IROPY6gQ2i2OrH/7RUn94RFukRHfq8+EKh3HJwvniK149JCAAAAR0GeNkURPJ99XohLbldl8+WVFj+Uva/8XeQ7k+3TjC2+NP1oPDpKnfim/+OozWJAwVNoYbSwY768ReLtDwbAMvTprf9JfzUVAAAALwGeVXREf4egL3/U68Y8CYofdFnLKmENVl4qCMGjiGYU3fOYIjY08srI70+2E3XBAAAAMAGeV2pEf4Z1uDH/duSE5sdhg1tyJKvPRV0pqHivmE9r1icmTrC7mlJ7RSXv2UrMMQAAAEdBmlpJqEFomUwU8b/77lcOB3OPxERFqzYjd6PE6HsucwQ5Vqvwef/OwSmT2LMOe8YS/MYKVzL0ysCkf9ELyWxaN9aVNLSQfwAAACMBnnlqRH+BB5WDztB5pfz3BcMFhRhOn/CrTtDVV5wj6Le4OwAAAJNBmntJ4QpSZTAjf/vBPT5YQOjX9DqM5HaGpFMwQ3svWbjKwnG+fOlGDVqt1ndhhMJgv4GvGjt71cvzG5oQlpabB9hxzOPik6dNzXgscBR4+cbPmz/fP5xfYfX1NBE6N+5E94WuKtbCmyYygQguamHJ0rcWbq8PZ31zvB4cN2ZTaqhzCYM9gDyDIxGEYnFFk8tLu78AAACPQZqfSeEOiZTAjf/7wZu7TxYPZR9WnPBp8pmX4bfrqeMVK4ZK+wWWuCUmWaP3lATqHHJZqYOGZWX5B1ExYwR/eLb7x7S8jkASRZ4JYh0GwZIP9CaseXY+FdZZiTo1Af9+ZTj3dfC0dKh212pB65IP6OYNrtjjWbaU5N4lw/dIVo6wfWQ1Ox2R/Ah/76cmOMEAAABSQZ69RRE8n3sgptQ/K338OzsnCnldYLd7RX/bbPagtm9AT2uZlYrwak/pO2OvTeE+MLYy+EAUJYdRjCTJ5ws02INMTcZhmyBQbhyMvzM27KZkFQAAAEkBntx0RH+AbCjmbUS83YiObdjoslRvNC+1RH/LHg47MFmS5IoLuTYHGm+zPkdYAIx1ysH28isEgGxKBbm/3QA1sfXNc14DBq7XAAAANgGe3mpEf4TUK0MsYlS18DEAHiL+PcNPmI7mQ70eRh7AhGHFQq8+OOpKoNFJ1E0xbu/KTN8mAwAAAIRBmsFJqEFomUwU8b/77lcOCFH/Oza0YIke18+jZcsifjoIbn8hLB2G1sybcgi8yUenoEyRzOvlQ9G6LQ3LjzLO/wVkdTKQGKdViQNnNx/ukcfeqyW26GuP1irc+a/5I0StQFo5b3fttkhWZLCcmvXpDHKib2Z3J8v4Mxl+oH8tYhQMmnUAAABHAZ7gakR/gQeKxIei4mGll/v6lUvt6J01kEEpB2nbmZodQK3mc4O0v1zs/0nPzjxxUZ3GOYRGobV/cAuzXdUnGujW6ZuRq5cAAACZQZrlSeEKUmUwI3/74VHYcs5s9CQn7HDcl8JfyAYtwAQoJsnV5HQm+wXQ23ptJOsGIH0+K+T4tJmGWxKq2QkvbVGAt0rBbo9MwfKNZ4T+zu+cRWcAzne7vkeYo34ehQqXPe0QPB8YoxnrjvXyG9IVsIPvTqRMCT5da0BnF7ybZ+dwHNKFrocGgyWwGl61sMCNyw7UdJNZmxywAAAAOUGfA0U0TJ97IKea4wcfTVP+rq1dhzTiOczB8P98SD4FQyfV4nTwQWnbpuSHsjXkGDVY+VFqlbqpfAAAAEABnyJ0RH+HsPhkmM1qy8bxI+teINNs77QclW9Syn7VgOysKvEXYZh1N9bkMSWtoQGKeb4wR7zVJmCMskcNH/eBAAAAMQGfJGpEf4eNe2TaTWCeCgaUGvNb9YsQ6osP4XH63aOXxG1jZCXFAnW0W+T3+SXJsEwAAABgQZsoSahBaJlMCN/74Djgeww9t1S4BqrecgJLTxVeOuPfGThsraJOF2+pS0DHf1L7/3y7WCHgywvW5UfaSeI840uFRKAdPikqUbRyEQbimIxPxFcVU7IF9oKgqilKzVfQAAAAS0GfRkURLN99UveC5swUPUef8Rxxd/DnGwoGPLEv0jI5bF/aVNigJFob+cQVyS7+9J0DBc/WkPNi4m+wyxdk5RauSIYK+bf04StobQAAACgBn2dqRH+BB5WG0Yk9tZXU7nhXQUxniPOff2if/2gJwmvSs+oXpZoQAAAAgkGbbEmoQWyZTAjf+8SBZ14QI606liNL7F0TnvmzD9NgO6qUoiqkr2dMXGBxVHEYasaZ1lgL1s3uJPXelryzxkrFiQQAJkYTJ6/ddfxquNOViqUsaM7edj3CU24qxjGC94+WxcABRaj4MjKQx/0LCKQROQaDtp8J+TfTj/pDj9bBDoEAAABiQZ+KRRUsn32aKglZlvQTovHIeNH7gDIJfCEh/db+sTywUiMMdw5UkCuutzw8i0Uw3jTVZuGhnn/DdTHavTyUTz90N47s76Kr+H9cZcMdUHeBR1waOeg55kL0Lnz5MaAKqHUAAABHAZ+pdER/h6Avf/WMAaj+02h/nyrbUOfaALfnmG2jGc9vXFwOALhj3/sE457lRtcFy2I4QbRoLfMtRwbncoZiPe5ckLTaUSAAAAAoAZ+rakR/h+Qj8Y/7qZFjwDICR967ohOCf916Jz9qNGDcGK/ZVQo0JwAAAGpBm65JqEFsmUwUTG/8CBPDIEGcd9xK7sr7PeVRUblA+A4wrZEIYDPPG4pGP8GneuG/+WblEBE/YISVAlzm8YRF1iarh9/0YM16M73ZZ8APAFET4NhFgwWEytfdD1qhPlSrycwSgeZBOZfgAAAAMAGfzWpEf4EHlYPO0DUfukzjrb6sMwf9TWr92SL5dYTUf9xex5gvecavMELDpnbvgQAAAFlBm89J4QpSZTAjf/vhUdhErrGuSF/CSLUgNGgjSozDQkuWrpzMDWZP18VLPXElhj565QB5ATyxhTtCBjXMORwH+flU+06xO0G6o+HBCuMLMHZQWc86Sgn8YQAAAHNBm/NJ4Q6JlMCN//vBm7tPFg9k/OLpHQkps0Dx8gSIgZVvNmWxh2JbcEgdj1ANWONNo0C46lGcg6+GMwWquPjmXeWwPjjFXLS5WiDMCJHZI1S8p+dhwYEHn2JemwPkCTGJT1YfRwqRnd4zFrrK7w/IrZzwAAAAUkGeEUURPJ97IKW2n8rgeWx6x/uvwym+iKn21XOuUJQvFl7Jvz/4BnZZvs/HplT9c0k3KbuOgl09GJ9Lr+6rhKi+Rn9GiuEu1IK84X22BQAPloAAAABMAZ4wdER/h7D5WmUT0hGm55/V9t/EfeHSNG4H3z8SITMcvPv9nR/LKHSw6207w5vl5dvHvbN/P8hIiRRBAbI6dHaJoqpBvvYa6Mw0GQAAADgBnjJqRH+E1CtQOK583hyuzCeT2ES3e7EAJSPa2lf4ndS5OFQQrgTaHVvvcR/0mjt9zA+CF6l0IAAAAFZBmjVJqEFomUwU8b/77lcOCFH/QjF2o2bU0D1gZNTTA5MxeI3pXbKFI+z+I3Ehsro47fLOUN7RwIz1I1j5IIQ79ZTKZXOQw7psvfz8GmJZtuqJ31wKhQAAAEYBnlRqRH+HsOQxTh16NjvxJP5pSGoyZ2MG9YV7VnkhZNujwds1WAb+DnvAXVsQBWexlXy/XhL2mMKRVlpuMTQ2rNsOIosHAAAAp0GaWUnhClJlMCN//WN7G1kj63WwyR24JHnTxkh+gYl+KzInu1p18pjfQ9WbiBY0GWq7ysDGaRRp+JW+XOLVoUSeFufVt5tWKNC3usOBK2/C6PhT+8zUu5EZNdqT60Qq4VgNXwzslbfFjnJRBDpdAh5wLLWFV14beae+G0fFOrSHDkUNpkd4sy70UWw9D1jdUuQ5rhLdlkIPJ0x98OQL5TO5gSkKid+BAAAAVkGed0U0TJ+HzcuFAuIfw9nuIuW5/0fmGPCk2buV6y+c17CysRtliEWgeuDAVuS2NEOS+YOmLSnbstF7rBcIT0xiTFqSL00ySVVFsbWwkDRuJB35PF4ZAAAARQGelnREf4pxDGQCv7OsLfYo7T8YkPot3UwsfuaooJN278sKvEbac7AqI6vugAMF0dBPcCL6ff5MKRFnV3TxohZvBW8Q2QAAACkBnphqRH+QTAw7Sa7nUoHLO0XWnIk0uc9V1EtugHZKDYDxjTdjxpe/wQAAAK1Bmp1JqEFomUwI3/vBm7tPFg7wfAZdXzxs0p4sQVbaLB0awSzslu6gp+TNyByCumJLwW5CyVRGbnfEuH5vtaS4wvsym7ljMIKBynK3Kg1MUwKdsVhXZ1B9ez8GwiegBu3bUflTwt/cPd13A1iwLV6wHvPtPLPVBufoRxLdi7fXErtqXZFkhyUPiuvfApRuTZLb6JlWmyXzVCqHWPaoL5G0ARafVrc4NRAD45EOQAAAAFFBnrtFESyffW4h7bT+VmzF2BY/3/IEonznzSaM52EK3yed02INvn4T8C5fRiiAYPyU20UWm0Zioarkwh5eCaLWUgTp5urYldHQq3kp7USwoZAAAAA/AZ7adER/h7D5Wmcyc1Gak7aI9t/EfeHMOpqmvMV7oyvL4VZVPyQ7xKvwftFj0zEdxLJQd6obfw6JTO/DrvdBAAAANgGe3GpEf4TUK1A4rnqdTB7ynQtEgcvcbzbIwk1IFg5AL8SoEuxeGvos/djDz9/TNXsT339wiAAAAGJBmt9JqEFsmUwUTG/8A8MxwQo/6EY7tRuuvUQ+IkRW5CuuQ16cI6K9UjLCHt/DFRPI0/cesLTyTUMPgo74dmiSCFWGW0NNaVwImkdPP8OS321tUl0fM2FRF/ydM1FfUQ53LAAAADsBnv5qRH+HsOQxRCCMCTKsMv80pUcMEXeFoGccbVvJvQNyqsPl3dGC9TDj68G4sLpt/4joEDUDO0HQPQAAAKZBmuNJ4QpSZTAjf/1iba5425bwbsmO+LqBQWsUydEd6pAOGPrrpQZuwRoMIuUNmIqaa/ZfsYj18F3+2O2VzASzLDk7M5bTVsxDhHKKWgIuIMjnu3J7use13jH+mTne9OXAolpGEM+YOkgGn0roRh08mvTLZT60kAdnRLZNC1MbYLQv/UtAcmYn97x3XP7Gik2RIKYbYXaybwMSmgqFRWYYz4aJwVLLAAAAPEGfAUU0TJ97IKea4yheePM3UV4BYP87QSjQ68uFMzfjZmCa2q2jTWdubpQmvIRE5Szocob6t0FwExDfgQAAAD8BnyB0RH+KcR92pSK4d0iK5pn3//DwN9n1Sf1/FycV11cqIRtpQSj1cG8kqGHiP/LJ2Q3+b3t9voEB0KrTJ0AAAAAuAZ8iakR/iOA7CrYE++z83jlf99FGL+nfLDBphBaI+WarzWjX/yqdZNNuIw/QVQAAAHhBmydJqEFomUwI3/1jd3Ei9XsYRCa6uUFLqSIqHEXps61oaZW//CrIivsMhGGZt/+wTW5oOHXAD9VovymjVgcdopGLcyH3+31aTy3MJifAOoc8WLgubo4ISngsCMnt7ldH1X98l1NIvHlX3mcEIj84dFxX5oK66OkAAABHQZ9FRREsn4SFSqCugoeFae6NWP9//36wXZWgjOWqAHC9KPG2ihYHENvn+jgjcT/OLdhLP58EwMXKROhlixI0N1Rfgfyt4FoAAAA6AZ9kdER/gGwmmmOq/Wm/D3dnucn5igsLTeqyV3HaZ/bb0TUerDFX4r8MRrARcezd2SF5uBaMc2S+4AAAACsBn2ZqRH+E1CtDLGJUtfBL4cPgesw0gfX/r7jzYBrCPcjlvYuTujdx5N+RAAABSUGbaUmoQWyZTBRMb8EdgCeRWNj1vNMstxJbRx7SX4GlzlNFlrR55BZRfwuUxlJowdEahz2muYI8I5vx8hRYpCbfVNcn/buf7/gZ9+q/Xex6CfNdIK5dm6hYw4V04WRD2xx5XfgFgFusJZ2Lu3x5KCiAJR/MKrRmuD0LpBj2HTd8vSTNdD0mqExm1jVxz11VS0JkCz1kf3+GAZ8aWDLfmsfKwTV6EgQsGOAvWjEoPcSCVw4czv5SM2yxrz2q4RWCXKP2cQ32KjlN9WUwxh1yU9H8IlwQpYkE9YLpUDTnf5VkAgRqHHtbB1bKzLqdQEBRnPDUxU3lCJ8z1ffgnZnrwQ/1jrkIVCHtxSHz+jlGJYYDk6e7WXKb6WyPh0gIrK8i81b+HquyEnpTv7rHDiPZTeoAPC5ivY3TNLqzME6cL8+xoSGHv3D8YziNAAAAOwGfiGpEf9e9wzWJD0aXSLU4v4AFyP5OmsbBEcN/WP0ljTwYR9rY1eB5C5dJcii79292VcOw01nChUeBAAAAwkGbi0nhClJlMFLG//QSJIKCLA21sITm22CQNIzL75x/coPBB2IH+CkNL45/JcPWK61azax0RKRKTm5d+BssnzMtwKy+njVCxG7G4L/0fmAAHFHQoH93jRGl/HjJy46vx0G6NyLZe8h2jDkca48Z8CHa2hCT6RIrRCpQcVFTXZInLx/gmR/ABv72Y6Rg934/cwAwDTTvn/FxI59I7ms+9ozaPN3OuXj/ZPs9jO8lHXscoqhncvzhU0B/ZiiYXSNR5B1ZAAAAQAGfqmpEf5icgpJHKscCg4dcTFi4T//E/XAXJpG32nkk08fkKUrnPGoSQKoKEiE2wpfQy29GE8eVH7/yrD8SLOoAAABWQZutSeEOiZTBRMb/++FR2EKSC4tylzphI7uNxilsvtRiMSiFLXZzkidfIAjCjC8+Yq1U+7FQos18k2wGsrHuaY/yfMMHdO+Ytp9nyOtNwN8ZSWlSsgcAAAAZAZ/MakR/iOA7CrDroBeWntKNDlGBefdNgAAAALZBm9BJ4Q8mUwI3/8b8mat5xbgYxKkPL6JuE0JrSSkGLS34qllRYb//n6qI6h+DVp72C2Bni+fBUJkeQ0aRWTwoPYVKKNVd/i/tVkLXXk1eYhf5mDfymriDx+s9bFWvnxtSq7+R30/mIRxeiq80eB3FjjMV0v9A9Ip/t0I8jdkPBxJNLY6dHX57kOj8uBOxzCEcewQyFjVwZ4U9eeiGDNW5cmamVQO+G1zTO18X6eiBVrQ1sAaKeQAAADJBn+5FETzf0FK+iADehK1fJc1pMmFCTGtLvZ658Bpk3MRU0V0znp7jGG86RAJ+kBsHtwAAACMBng9qRH+BB5WG0Yk9tcQHS94V2gbW9d0P7JZCiBx0bHdSHAAAAHRBmhRJqEFomUwI3/vEgWa8cBuzHU5S0ix1bYRHinDETekQfB+cgzhJLg7T1kDPCseQtg55Z8jXr7sf2ocT2sSw5hDEdVg0IPoU/FGMpuD08g8QhOFlplsTFkZYkjR7winNg/k/jN/HF5GJogmrWYBvwozzlAAAAFpBnjJFESyffZoqCW9/YL35uOpRg/l/iy41hTe/tE0ssOv4hnLHsHPfmKcSLTegTjkNLZVWXjO2NeDSEq6BzQ01VDjT7II0nXMOQjdC1CNOmzNCCPqLBayW02EAAAA2AZ5RdER/h6Avf9Trxjzf30y6kAOs6WLBLxUEPR8k/75zI+NLVHsG3auYqXMQfTCCvtQyg3VgAAAAKwGeU2pEf4etKGWAMuE3sYEV/anP9U6xl/ko75T+a5ob50RK54RK1XOQ4sEAAABcQZpVSahBbJlMCN/9dtmbcgHeY77hZBIuaacrM+PCy3y2hv4fgKY3CAJLaPuYPjQ9hPdbPybTfLA9ZhcESX6IM5yNtlJbVY47bfUjagQR0EA5I+ypI3tSGPRusuAAAABKQZp3SeEKUmUwUVLG//vhUdgh9eiEMPfVB8tcq9JDka35FeHzBuuUIwuu3ZTYJerdc8OfLMQc0Kb0i+lVDCTmweLq3VzwQLqGbWkAAAAeAZ6WakR/iOA7CrDrXqtfOU0OnEFoMihrTVlbhcOxAAAAckGamUnhDomUwUTG//vgOOA7zH23LaoiapuTmxATD2bJaDhWz+mrw1lKFIcbAS6u3oi38M5gOc+H7/dMAlr5FiwNQQ+MnP95HqTXfhGOyOvhpUnJ8xbBRClFaRLviLbneAYkMEKA0mVg3tgaK2CBRyPigwAAADABnrhqRH+HsRNv6zFmG77dxoqP/yRU8VkBI9JIa9XsBjVtXD6MA4j5uBIgllV5DYEAAABnQZq6SeEPJlMCN//74VHYIicdzoZjPCSvgahArEevFvdThrsLLMVaJDVNLqpZTPWNC9vVX/wO1tcwEn76IJESXUEvUSeHuARe65ZpLp38T60KexwII0tfQkOqrPOP2CqJfFhslG4y3wAAAGJBmtxJ4Q8mUwURPG/7w+3BOtghVFB1Iv2nqcWzEymAHxzkidM6m2fE6ogCs/6Xy6jtl/XzhhsHnCSDlAQ38GrqPDxXxyRtvk4SCfcKOBRFyRAOB8n6sM4vESFElpVQD/vjoAAAADABnvtqRH+HoC+APa8YzC92E2GO/nuYBco//L1SmEwKP9mqIR44ppDFXNReibXTYz0AAABWQZrgSeEPJlMCN//8AtC7wJQ477MPq5rGdC7Dcs2GbwHJuPCJyM5OIOtlfIO3cE34Ssa8cZkPDWcsRjRYjceQX64sn64LgFO/LLPWgI/X+x8DlTPlkHgAAABAQZ8eRRE8n30rer4SgxWSdC/KZ7+otOEzo+UFD9ZmniYypUuFzJJ18Je3Kzvv6f8AHcMmi6YRwakaHS3jzaoZQAAAACsBnz10RH+DDmAYlt28A+7MCxlUQzSjDlyF0aYn6fKlx9oD55kOvT73INX1AAAAIAGfP2pEf4EHlYPO0Dlcu33mg7UP+2G/OS+vSWgfJ6uzAAAAdUGbIUmoQWiZTAjf+8FGR0k7B7JF8YA+wQjft+K1l7Pl2lFva4U4iloi7jK05WAXlUvU0NMbneg0W8bipRMj+e5PR6G298PZD8Lx4/pmueJhR5QgmVsDRHocNbVfEJGisYpDRTRCNELAzBSPwHo5TkjGz1jvpQAAAFVBm0JJ4QpSZTAjf/vEhSlnIiISPCRtw/Fa1SUZHmNR9eS02TNX9cdJHHP4scrbtJbcvjnamrR2131icLwbB8mAm/JS1f8/MJ+s1WuEEA/spOfP7NBAAAAAXUGbZknhDomUwI3//AKArvBG5U5RSifFlD/Ne8Dro+FmjNGrFQ79lqWS95HaV+HspU897WfSmm3q8POHAYXGxvGACHm9MrOpAgNLObTHeON5xAh/8S0wNVUbBIo4TAAAADFBn4RFETyffT2ygu6Evr15zpphVL9ezEh177sG3lNh72dSy6+LMKsjFpOfLT5s3125AAAALAGfo3REf4BsKuKQrn4+ytpBYNz429+gOd9CbMcAZjWPrdtHWJw6lVT6DYuAAAAASwGfpWpEf4exHtbuDFFtO9aECqMpTVAgxUDNC6R6Ov/Bv45Ck/DDI++GHNS//D9vpyc2nRYpkvKI+H2S9SCYvOaW3tGooBG5951jQAAAAGBBm6lJqEFomUwI3/vrPjgITCZv0Cm/UN86vq9j2om8TSXh+Amp9UymJ2Lmx+GhPvnAhOiY8nnAk80c8/zR4tT3chvitOsi24OHlSIEosC0pK5DFjcJ9H6E07p5eo+PK4EAAABAQZ/HRREs331No5t4RtqmNt6u0B+Dhno+LpcPQ4byTbJJSghklqX8XC5+HsKV6xKkIFe7OUTxgQ0olu+pKqtaRQAAADoBn+hqRH+HsRJVgDTILm9SUFNJlrvLvY/SNT3pkk83hKOIuwzsAgzN2cA2ruLxHhJGC0nlMqRs5YvLAAAAPkGb6kmoQWyZTAjfyghC9HoFZjDr4Z3T9vVMSrvyBTETIkGEDt4rHdxFlDop46jR72snDXzSvTxci6fpWQvZAAAAbkGaC0nhClJlMCN/+8EZPlhAJgR+ivH2tTAf7wMtf7Q0lmwSp6DMeubPofusLc1pk88yiSuDMj9NPivga9JoF3g+xq0Q86goTEG6ZSyNzWHsnx25FbgqP+krw4d7Ztwy2JXZPgQfwac+bc+pWEYQAAAAa0GaLEnhDomUwI3/+/6N/UQQ5pDzZ2qOGQygvLtIfyWYFUfIE/BzO/DrGw19HGK3hwfXoRT80JapmkcAzbWSDpBJ+ucKJJwYRfQ0v3dYjGPpzjoBe3cpWbig6cBTj8apIEOyji3iKmvu0ocjAAAAaUGaTUnhDyZTAjf/yghC9YOgdCGGxMBNvyCtjziLKdC2QZ7FDESe5J9Q9AtgUrneVdY5XUnITqU7yU6iBGOU1CpzwDon0qnZ58WoQb5Yntsn3pLkDSXmlHf4GUB+Hia+8e9WqG1nEx0tyAAAADRBmm5J4Q8mUwI3//vgHICs0H6uMZza6s7teVDeMd6CYdhr9xUUK38FJYrIQPqnYUcIjdPXAAAAYkGaj0nhDyZTAjf/+8Gbu08WBCblR+AghU5I4PjruVw1nmnblE52norPEgNszfsQwVFghKI1Diu0eqa7hbfw1+Ep5QkXdJelsyIpnCozqoRLLfdal5qiCN/T5eg+PO3oDHLfAAAAaEGasEnhDyZTAjf/+8Ohm1TlBNW/tdPMTK/c3ybOs/lY+8ZeligkNk05fNM9wL12M3EuLDYgcCCXSoBpriTN/A1g6V/gLzeZ6XLH44A5B/SCX5SLpQw9q99kGKI5Kg0SBuVmiSFKDmwPAAAAXEGa0UnhDyZTAjf/++5XDgIS808950w5VkT/t3hHWXzd4j+0CV+IqCE8SuscXSYBiVDelxjqtSC3ZHtkbYDFeKciUbuZ1jlpEM/uS2KCxQtBuTVImzCSPE8oOq3AAAAAW0Ga8knhDyZTAjf/+8SBlbxwEJo+j5XM44k4C8vgBUE958Col3eoL9NbbmKsxW8m7j4xRgVojr6fB+NvuumtaBaIlIqjCRjWB+8H8xBIXkvlaBSjiwCBA42PXb8AAABJQZsTSeEPJlMCN//764R2B5oaIM9nfTCNty23j45IR1Uylna2WhBb8CUTivs1xgc8YLMXXuABpWaMZx/XUnylyMe7+O3whopv/AAAADhBmzRJ4Q8mUwI3//vgOOAkGyPnFZDh08wNcfuRYVMfPh/g57ian/0yuRdww/1n7ykQqNJMvWQ3gQAAAE1Bm1VJ4Q8mUwI3//vB61YdgbRGgzS0lo6m3KNl1T26A52Lu3lQ0zu0gITxVmY3a52vCyX+xlNajR4CQjogG549vg9JY3n9NViV9WrQgAAAAFVBm3ZJ4Q8mUwI3//v2UuUJDTVPO7uHUL9d32bbRDn5GpAyvyun5eOVCscEeVYfMq+tDarXNr6Iw0bo8A2x7veHCIxWFjrr1e/Z5ydf+rNut9UgsY3hAAAAVkGbl0nhDyZTAjf/yghBljfBFTH1dyGLzpj4+bIErKAmSdIbdehwtvdaiEU1PfP7BLv/m0QUWAKIqCCCv6cnt9U55oA+nM4cUPpNGIImms98PbEjWuRhAAAAPEGbuEnhDyZTAjf/++AcgTWtokgs1Kg+3puZmyB2zn+Ft34Z7zrHnHnpGkVel3jgu/cRJkseHDVwekGctwAAAFRBm9lJ4Q8mUwI3//vgPUQh0QPQquZqwQfnSdoeSCQCzHAsvi6Sa/8leZL4+gZ5agcDt/8ZjU6OKKewYE9oGmb98U/7Q9w6RYYgUyseVRyuZy+FRnEAAABYQZv6SeEPJlMCN//8AoC0RERAjFSuAaDWIHfxup8AQJpTK0TZjCJkqKKVH0rnE+j2G0jr/6PUOpHvIfb9HtdnRu84uQ7BUI47R1ptvzy8Ih329lBpeqV6wAAAAFNBmh1J4Q8mUwI3//vuVw4Hc4/EYZkfR9NOYQkN6IR+HLjwNnXmfJ3tFRz3h6PoPlGR9Li8C26dewveJqFV+gHMH01AqXp6nEKs3zf/A6nxbhFlYAAAACtBnjtFETzffU2jm3sH5xlgay9R1QXJfemvQPA5ZQM6hlOnPGFXo84lVDBTAAAAMQGeXGpEf4ESDdqTsacWmS87R3boOwBLmPjr9HbKZR3BTbJMNFk7UaY4D+7PnY7F7BwAAAB5QZpfSahBaJlMFPG/+8dPR8sIHRr+KlWnACcQrjiBl9PrT9yPzdGTDPem1rqoZeo+O+V63Y/a74OMloCXYtj7s5ZuipBPqAm6qHry/NiVr+IWLYfvYiQd8br9YnzFyHPfR+RTXTMKvVV/KjPD/ctM2HNXr4HAOamPMAAAACsBnn5qRH+JNiiXuXcPpCfMZDOo7tZbnRWsBJ0tsa5Lgg0pFjK/6ozWiljBAAAAY0GaYUnhClJlMFLG//vgOOEQ8B1SOgWNXwdoVV1ru5CEMBY1K6LZ98XM0dIKZIqwe5RQeu1WVBdf8q18RVrmysbabE4YEFPFSDmDTf9crfxMvgknEd2vk9sWgLG4Ej0lyc7cTQAAAEcBnoBqRH+BB400x1GiEKvi8jqlJCwv1QQndqVvnTA05w24ChMFL28XtsmkzSsUuIWLziqkkynS7jrcWF+o8Wzss8ftVv3C4QAAAG1BmoVJ4Q6JlMCN//wDwzHBCj/9cyI1iJwiC2YDq77/ovgCtBWTqAK7Qd4tMo/sMxMoLVehiq1uxY67X5FKvvfdZbk571Wc3XQkxwSdzgyBXFfYMHtvUybBlgS3PwuBgOeOBptIcLchkHRh5RWAAAAAUEGeo0UVPJ99PatKaSqB2VOb+Jj+9bQe58CI9t/mO2D7iXRZ6LVPwa0DZ6RZvvddI6Fbxs5CUyQ6tXX32iUFx46XgEEzck4WJrjn27/rI7DAAAAANwGewnREf4USysNoxKlr4GH/7yewiW8E18hvueS8dKXaY0ZRvRxC4LmG9VnnRggXFRoWNTsMGYEAAABRAZ7EakR/h7DkMUOCGBJh+lP55SdE5jTC7V7n2PYNjFca8HhrbO0wW5ZpVWgnD67LghASkduDl4tN3mv9lgWKKDVlzuDELPcvXMd+GCzMWgyAAAAAskGayUmoQWiZTAjf++FR2IObN4lvte8kZ2zczUBUIke9zcAiDpc0tIGFXyhAP1rxFfYV9NigpvIUKbs+TB+6WyVSWBBcm9MFQqwPuVSWkoeLyC9MIaVAuYZceC89lbP77DKjo077qt64djt94DKl7X3TkhZMKCZRkli9PMehVycRIJ11KVnyazgfI6QFyZBWKpZ//e1ExmOA65u9Rb7EpDZypnuNTKWfi082KgLh1cxMCfQAAABGQZ7nRREsn3sgp5ri/wKRgQ7Gu3vI9crS3yfW0cNACqpw9OIocYE6q5EFc0ZcGv9rKrcElxkw0g99wSSyz9azLUdIyByPgQAAAEQBnwZ0RH+HsPhkmMfKXjerJPWkzSzLvYro3D5CmzUrpe93QDwpOs1Z8ue1cdrhkVLaJ6fnHRj/QNrVfPJjXQ1luW3r+AAAADQBnwhqRH+E1CtCQU7NGQR/SzEljM35RLPhIf7CuGp2nH7PGqbYx4k/MVtURJNrZRoSWG+5AAAAe0GbDEmoQWyZTAjf++FR2JNg6dJRFEWpLDaHGWYdvgn2Ihhkb4tam6cKVyugm5faVmXTaDsJAu+kZAz2r1gaf/a+k5hP0Q9ZmUfdwQ9RO6sckew+0/TuDcwET5LS6b0bjlYPX3/3sW/C+vHsLkXWQdgV9S/yoRC2BzqxIQAAAFNBnypFFSzffVS161ytJSzivxBMnWEBBJbIqQ36FF83chfte2KE9wjt18/IKs6+mCtLOAckfgeDBt2owHOLwPzOn+zvI5qJcNySwzh6eafp0V1+IAAAADABn0tqRH+BB5WyGKse6lBSKcVv42bFQay/dKDQ6P/xo+hOzobC10h2D4F5zC/jvsMAAACTQZtQSahBbJlMCN/7xIFmvqIfV0ksEmI6ZqkME/+uYEFZp9EgRyfMfI9TwvOdvYfNMg96XA75e+Ug1d7MZsjl+wIS6P4ARavP29AQ7GgEU3ga/qHpTK+j4PavPhJNuvI3vAVPp+hqdWdcPhLbXeG5Bavj/tfM22MWc8DWUPgPRup92tOMWZq7zTg4fRCn0RmDMdRAAAAAZUGfbkUVLJ99nGI9vhkeUtjd0FBbV/Y0chN5PW4RPFtqyuHM6CMMdaYL9nKS2MPgWTHLmFoV57n1YIpZP6NS6yNWHqbcybKI0qEdvxQUm/iL7kE1gKC+fig+SwYYBKRaTsPgiDmBAAAARQGfjXREf4egL3/U68Y/RuCrxLwXiYQlWXYt5tY5DXkNo5EpiJppoQPvg9AOtTRsWp+ZH8D4KriCBrIbQXvLyYYUk6guCQAAADoBn49qRH+Gdbgx/3UyLHgmWsQQA67kWXbzUxYaBkUPtfk5+96VwwOUX427Zy3tdPCl2FQMk6Iaz/Z4AAAAdUGbkkmoQWyZTBRMb/vuVw4QawH5IL3wgIlYl4N9vN2/EPiB18zVQ/xYvRnR37JySPr8J67t3TEEeawOHDK09GYOfGXSC9g8VJk5uBSTrdp8SDss/hdObxpyozwaea3Hbi8nMDNDeWI3uORE/3DMAFFILF1w9wAAACMBn7FqRH+BB5WDztAwaY/r7Z4kRBTDYmo+gXVOFAHpeYiX3wAAAHRBm7ZJ4QpSZTAjf/vhUdhErlYqhggZ3H5bTa1/vToTRb7froaegfKbJwqLbBW+kSnruoRnFWo3+XuGIHxrqAtWXyVLj2RxmDnonpBATSupm/Er3ieXpcjPxjEdhAERQo1d3evkEK8xubffKr67kXpOm1rXRAAAAExBn9RFNEyff9JABS2FcQ8AcUeHDbcg6r2VMcDhBCTlaLaDqBcTdsiD8BmqZH0SKwn6VGYbJ1Q5LmOm6cs6MiJ7HuWG4vjXM+7PUM+nAAAALgGf83REf4fkGk3LuITCI39/Buk2GiLZ5Qm2lilCulAD0MHZ7IiFhpwQ4cN3dmwAAAAzAZ/1akR/gQeVhtGJPhygpFSbr+MRmIRmrZNzW2fSHOqVlo2OShwB4K5dCDy6EojzpurJAAAAhkGb+UmoQWiZTAiPjDpjcgQswobwLtrOt8Kbn50NO1b8FyU1efKfDFwXZG3fQEF2nXHsOOeAZm6hV6ynSwPWwas2HEejxuO2U7eYWeG4BCTxQsqO5lECivpfGEWizf/HBrRDxbzeeL8AMx7yZyACTy3DlgHl6NdrvffsWb+btUZdfydSDQkpAAAAekGeF0URLEfOK2yXld2veMoVsharS5jlxy1WYfGZ6w0Y2W4SyvoZW0vw2TI0HcN+K1KTShZbSy9NVAcHfLkabTxGzbYQj8WsrHVNjL1UdoGjWTy9FdZl/uN/rXdkOQFudoHdTY917al0GV88Ot6qxJTUo5a0bNZ1+ujBAAAANAGeOGpEf5iZsaTpiYK8Ckzq34eYNkr5wTjqRNgZK+2VKyFYuAaNfhszOaZxYrXPiVM5moEAAAViZYiEACf/vgy3mIDfV2I+c4OQeUCbjkIqauQ7q1hI51t1KK0Ty/hpqe1v3iAKauIQOhjuzFxyMgx7FbnM/2jBWLSkwzPM9uWS0H4jWSakogzfe9txEs4c+/QsLeTPvLzlrrhoRvk76vX6F462KFV4lQ8BLFX1JpZOU12xKsmCPjGaaOJgyXQiOhYn2hWzgcqRNjGX+V1k8oj99bE+1vc15C8Q37MxkeCbaMkDN/t8mxkCWv844zRz7jjHMSn7OeiLDTN9e7Cv32YENkK+X2fgPjTA2f9fID7IG6+92XMlKB4gG0/ewdEoHPbTbV5cbTze8c/T9wMuzhP10AuKfYQG1eXnZ1ykp0P6DRPy9Ld6m+2H6X5Ni7ssRrEN4TuoAEb5bQp2TW/e41VdMbUA37GNZ5x7DsYhq5D8keOph2B2Aoqxc+qPJPC1eN8uzF4kT9NEbttz3OMknmnGRbLqo9BolsBfy9A61ilUr1hK+lfVlTjk05qU3y3wHYEc66DTCaL4HkcZGfGwlWCtebzJGH1tSY/TGvlzwoiR6D2VHDw4Qd/QL3sXXc+s5L0t9Sw7eMZ3S47uKfNx/FSmsxh8Vu2q6ia3DDbKLB0cs1CwbSXU5GqIQSYlNKNsVMAv1OgbV1e/q0to651ASYAj8tdg/GWuSkyEBfqOW0VahPSNiWW6MzjEdYDBORzkPU+fSsEZX6EUEJWgIR6Tk5qkRkpHpcm/zO12eo3M4AhR9Cx6B/Bjk6bjXMMI1JbvT4s3OmvDMWhm4GnlKV7U+z5WfoAn+r+XEGmS5H0VzKkjL9DjNzfvJzjcKRQzSjV9OMxXeO3VBTd+WFAIv2qLLi61YRg7R2i0zYBkZiZjznSxv0Oh3XHIYCRtkOS6xY+goD6NZWh7U8+5INbjxbNW9vW0F854H74WEbjrsIE7vUh38zBz8lmukAzLqGi6lSYQVXRAstPq/sjujZpTHuizcX6q5h208Hya5GB6cWOla3tMwJs10rHPZhITBZAHx9Dfn9xOK6H7+wi//UuozeCF596Bj/46ApJUXxnQw+iPZClbVtv32bAEWwjAq5WxRo58IBlPczqDbMQpJ2v8hJZ5eL2AFeiCS8d28CkwIevJBJczJ4bl3YAfkhz6EvtCfS59R8MtyqJAHmYMxFa2BzAk4Ceko+cmThvwJsTo9TwWz6/y6nqWgaEvkLNZuyixNMc/cXcRwsDQOLI4ZyMttZLESIVUpeJZ/6cbN2LT812AnxvEtv81Oo9jf67SuSMRsG6/uLlpm7ay66WEKT2P2eyv5sPPbiv8phYtyC/hpNWM2L6InvRMw1S6W6QQw7ANE+NOljaO2+xYiaex9bOmfqoAKKcGloWbR978SVkvTCVMZ+HYiMb1wAqx/PF/B7YT3YdqVhfszlXZYvWaSPMzU5okM33PRVoLBIGHcW9rG42KMvtm5tZJTaxxpsdsz8X16wDYWJ0ZVgl2FlFrFiwpttQ+K/SBAfBbCBjc4w+9++Vy9ubBe/R1rfhdPo7EO8Upn4ojmO4i0QAQPk/+Njyc3FUAfeGBrW8UR33ylHcXQhrcCZUhYC8/6+rQTbAW1pErnFsaRgQxa84DSprqtocWc1KyDwYQLDXOQyRgEVTWwp0C4n4FCtkK+u/u1SbZZ+WlsAmGCMhbvZttIOR1dvnyda1VEmJ0q4MLNPcwxMdDoqgdP9/cDbj6E8rNAG5pJQbjF9KX6GfGfetGbQA0fy+0g/NVu53wvOeMDlX5aZKNl9hASlsYvBxoiPV+IOBqNBkpAft4IlL1zcz0cOUncvHJKuevXoD5KrbBNqehopgoUlzYTvetOhot3VJAEcPkLAAAAERBmiJsRv/7xJ1U5AfCBaOk8iMcVZeCFVxqmruwhZ1VxVh86clNufhi8lqRRs6U6DSeCULyvbJ8s2QvtuTrcAMn+2m1oAAAABwBnkF5Ef+BB5WDztAx0jJ7uwLOc/joWyLAt56hAAAA8UGaRjwhkymEb/1jexjsxsfi6hr6As1+ItIPaaxdtL/Psp4lmstRcCE43uyBptRSOlaI0NUP9+Rr2RU9EU2/XAcFdqfD0kTm2yOmA4hz5KtaLVYCIyYQrzAIpsbD5TTBRHqAEIhPQuPAFKMNcGmm6KkfE8h/rUhCTb57Tr0SCkvByHkqBQxIkeSnZE3RETT0gzu8clkUB5C2/f9gqjzOEdr0tMQy028X9g6I+jHHVIcOQqoQJqlYLTVMEPnxTm8swUhbukQQqB/V1CMcORCC/q5HxRtMLYD5qtCkgPY4iYs9BYTsiCpbpzh3Z3BHriiVt8AAAABdQZ5kalPJ/4RtMreUjvAg87ctPZkwB7tGp3pNULncn75y63gaBD9ooZ+2Z5BGJuEzAt8WFMMMtfOq7F1nffEJMK0f/ywEmY+FnZVV4/fl1h2Pxh96wGJFYB7eE78IAAAAMwGeg3REf4uRpKqqyH5pDz4TvCO+HJMaGqrdtS376viJGvqRNtOi7CrvGS9ThjJaQu0ZQQAAACABnoVqRH93OX2EwrJauMv/aYKWiM4dLgfPA6VlKtP+gQAAAOdBmopJqEFomUwI38epjsLFC/F3bWYVsVjbCSDz3/NFFz7EJcPUsjeAhu5g0HabZ/pTdPiudhVDGiYxVKcEifZtHJ0mwAP02f3KxpLwILBoE5KpKZ3ak6k96u9QX0BYIZ2zDURThz3E5Ecs7TFwmbLYgIzCpqmx3WZ2+nHSOQT2luDm0mS7tC2WUBIB1L038q70IP7mmgmcfxvHDAXuwIgQQvTDkNbDDnb7qEeNphuMepAZpH8UGL526eLVgrKBJghOlt/VEGAzajEUl3BhFpPgSsX/QRkJ/xFpACLLDGkb2Jne1e6/1z0AAABLQZ6oRREsn812DBeerx++BXNh97FhCAHBFp2YmevU1UrQGvZZfGeZt4lRrbLNgTpILEiupxGlUxfWlIrjdvxwBNEtx8ud+WR8eT/AAAAANQGex3REf9fXXAqIFgtajW1njE1bz+UortEyKI1hn9ddwwTBBszfncLM5umhE6wLGy54PXqAAAAAKQGeyWpEf4TW9Rzxc6AcnNg8AfIm2Rmp/I+vcel00x36c2AMOw8muHw9AAAAsUGazUmoQWyZTAjfbibf1bDnDmV6F6O/WnufALeelmNBoakQ+RD8No/50k5DsmJ/O8f5+6zmOB2dS7zL9WBPkesZ8MQ3/0MksLiL2syu0Os4+DX6zzVah7a75x5i3NZB9mNahlXbFzLVLgWbDFo0MPp2ztyCmOo/mDC8YFYHXaGQYAnLA1akkBw+w7m97niKcIN0hcHT0r/8VhoRhstD4zpvXRFwqgp9DQSmsJ2LaYfRqQAAACVBnutFFSzfgcsWP16iAbHfMo6XUJpab95jekM9opekt48P33/AAAAAGwGfDGpEf4RTMK9Ia6xDXExmf11mxxz6JBZgFQAAAIhBmw5JqEFsmUwI3/QSy/3Mn9LlHVQzmvoMZRAfpfzmYydz58i/v/K7UpAsceKOBkM39soM37DBOWqDpDJYUzpx6MgjQXdFpMMlD1sF3jXfAp63NUJB2HxYwKAS3+X42lams0gG7HPbT9qoElWzMJ58INMDP0k6kDCkWtzZGiljRWKipkGdSBWAAAAAb0GbMEnhClJlMFFSxv/0LB3HWhFl4F9qAxu0d+LmsKyTRhMet5l8ycg7FLJ6bMwnHDvdth6Ny3YdAsTsiP7Q356Lcu08PjhT6If5FExZc2ZvHjw+Z0tvTWfIt2djM4q2QM2UtcHgN2bP19edptftPQAAACEBn09qRH93OX3cYVktXGX/s03NlZ9nnR8Um9cfaQPg9mEAAABtQZtUSeEOiZTAjf/74D1Ed3Fo+hammI+tzPVBVB9GAMtkRSmX9W55hcG3CFWWzlEIUNNLjUOeelsGZvhADsD+Ofy3nAN5GVG/eZg36oP7ixXBkBIm0aSeFatlwVCbRiYm+S35yo4wsvnxxokQwwAAAD9Bn3JFFTyfgOmmneDD0Jh97FhDeN8cWWgHHL9ZXuuK5vJRYDuSt7EY3WZi5fyM6swhMLsBDMNgeDfcSyJYY4EAAAA7AZ+RdER/im4U4K1F63xN3y7TWMTYDhtYotw9N+3yCyCwG7hbmm/x6VomjGrPQRMq2kgD57PUgsO5+FYAAAAsAZ+TakR/hNfANQiYzYypltEIwmskWD5UNtzWBIJZzKZTZZF3JtrR7AcrpL0AAABCQZuVSahBaJlMCN/77lcOEudjsGvy5OKoLXdnzB8pMXFbfBomeiXX0yT9Bi3pzVu8acc1sJ8hB4UGgAg1X7FrVSaBAAAAPEGbtknhClJlMCN/++AchVJ/hJ7fwbdEstXIOKFMPOfow0W1J9MjZSXaU2JvcLoxTQ3MkZ+p/N6EF1h/4AAAAFlBm9dJ4Q6JlMCN//vrP1EbRInimhmXtw9PH/HLVA1X13BZGuO5tuR281DwTZtzSQ+X8/0bJQ2SirHIQ09BynMc2hGsy4/B5wCyAIX1VOyVwH2zU/wrNxXXcAAAAIxBm/pJ4Q8mUwI3//13BGajj12djqpakBtTio3DJ9DBFwgACSqtkYv5h5ytASPgOdOqkOdgnQKmgWojr1xZ/Deftcan6H9Tw3ZA+MEWlDNo1ZYBWRYty0ND9KfhocFL32qclLBL9NniRJkfFCDsnUhqrM64RBnhBgh3Fvlja3uOhq+MRbDNIEtu24O7QQAAADdBnhhFETzffU2vo9kHhP0kUckkmlcOZ1qvSaHsv7FYPzfK4q0p2kKE7Mm1B1C8Hwvzr66y7F2UAAAAMgGeOWpEf3c5fdxhWS1cZf+zUAjrwmtqgIPXmRHv2k2L/2uJj4e2MaADqH/NY1U5BRBBAAAAeEGaPkmoQWiZTAjf/WKRlJMqrSGu8YMSMo4NdVTVxbO2c0DfzHpWj3NO//rb0lcIQgLmD/JdeB2LmNOXjmIWgYrNK2njO9hbOQv+tHayjtl/SNf2g2g4kXhlY9QRMuPkMSkCrA4mXqQYXjex6M4K3kGlpFUGNYB2CAAAAFFBnlxFESyffGVVRtiU2M5XSTBlLlLpuEazvf6T0QUYwco555QlcaEH0Dj+hFfQ6SVisFfojn73EUbm8SrLs/HDJEPq2aQ8uqOTVCis92ohdCEAAABFAZ57dER/inDB4mEZFozMu4+KKCyG++hK3D037ngCnbT0Imo7MlsV5lr5kGDygb2dj6rCAP+8m8TBp/+LfwVdRRtehi0bAAAAOwGefWpEf4TUKkzU19rfRxca9Jc/qwQnAu0qRq2xIkvrtGv5GbqxB1VczyvDgln4tgvLQ64l1ZbfxKHVAAAAekGaYUmoQWyZTAjf++s+OCFIz0x01M5hRDAgQ1SNkLroyGxDyE3PDxh0h/mEHwlASDMpnpSSDpuFfd6ZN1GJRdp7/1NKZ0BSML9Ri+WpFVWYy+oWu7FUhBjVLBsl4nohvDKmz8vEbpKUMsdmHeD8TfRVnZtvTlvIPNsxAAAANUGen0UVLN+BHt1YNOwIBo8J7x78Wox8hzQych4IV8JIX4JAHVjfP6+E363LAITjUT0/2RT3AAAAJgGeoGpEf4RTMK9Ia7m5tNYSHf4/ZtvhHWvqCx8lGWv53wQ3Vc9qAAAAZEGaokmoQWyZTAjf++A44S52Owi7Ki/xHuS/BZp70AE+JQqCbzs6/5q8Bi6TLqq2chJ4O0HsA3vaJtaNjwJK08NQv3nUM1lYoOFoLcHLdvLTHCfibUiXw/+uN1EJnDUQ731KL/EAAABQQZrDSeEKUmUwI3/74ESoS0QZzPsXZCCOIIaZzh4p+rex8ji1xPYtKKYeWQssfQMp3vKo0mHw0vJQBLiKE8ApAr2l/ILTZvWEdw0SQbXmWBwAAABSQZrkSeEOiZTAjf/74ByCUGSE5NLsdewC1Yf8gDKRgA9A3XGNkTmRAIk0zuqTVUB97w21ebY1ZD4A5JuX4GJerYYG1yn3M9Fjq85VjxY3yOfXgAAAAGdBmwVJ4Q8mUwI3//vgPUREOuU+W50ENB4ZTfPgxJ4R/QAgLWfpJ4afx0kcK6OroqS7DFWyZDTIp4V4in/vBYNjSauh3e2aTeWXDJZfmmNfyGj/jeE/EeTuP4327cpUG6wUjxVSODOdAAAATEGbJknhDyZTAjf/++s+OEsdvlfnggh/lGSNWpaGXz9PDwU7DfMzDnOAUqnRhTRcsjlee5HHdMPWnOhXPk2yc35P2EwrAHuNbsWtVJsAAABJQZtHSeEPJlMCN//74ByE7aEx9iXfU0+Qi7cT6qjr7osl+ltza7QEqjKcHDYgjA8iNym4TNocc8MUDP2ywSpM0nUK6tcXIpAVoQAAAF5Bm2pJ4Q8mUwI3//vgRKic7NNky66/c817e64Ip2yRLzp3SkaICilaCnIprlv8SQe7wRAh5lXA9b8r69v0/THs/oLcawFvELgXfQ9ZMLp1MwPkpBjFsgw2mzZd4z68AAAANUGfiEURPN+EdaQm1vW/TINWOBhDWfqLGez+cStDKozqO3ZaLVmanISgcBX9Quzd3kidv0xIAAAAGgGfqWpEf3+cgskCm2JriciFWFR8KQwp4tWhAAAAXUGbq0moQWiZTAjf+8P9YCJ2LEHgSkHxWPcIhuhg99q36oJg+yoD2YMOzEXYsm24bpU7e5HDTbai/djWZ57WEvP0X74/hSDvW98TN4gQ8FYDbsgAomfxCR6enD3bgAAAAIFBm8xJ4QpSZTAjf/vgOOK5ljY+B/+L9+6Ag5nJ5NgD1rC58rPGzwaZU3TSod4Momby9YNlv2Wut5rt0gK6/Y3YwVJRjdwKaYBRIsbzg3k2OgRes4iDLyq2y1We08s8H5cePlc8GGSMvtVpkDZXgQ3ARvXgVsnai3RWNmYHRmfNrWAAAABGQZvtSeEOiZTAjf/74Dji3z+EL3mriZEn7qSTHtaSGpQSz7vu0Qa1DTfh90G8QU81BZ+J+YsG68GTxDc1tNVG2FxQvD3JmQAAADJBmg5J4Q8mUwI3//vgHI3KkJj2o4R7wAdkoxoAPl6DDEXLEUPxOrh3d5SlgY3fxMeUwAAAAHFBmjJJ4Q8mUwI3//vgOOEudjpLG75zaU3JffQF+0lnrfbIzuF4nbSX2KhtapcwdzSYt5xk/gC0dG6joJPOtvfYDyHPdosm0/eNEHkREpQU/Lx3HWTyhi/QVBnzgdkBQv8ovz6Jzex7/D3vMDIqrT0wYQAAAEBBnlBFETyfgEm4jKr15Rd2F2D4x2LPtpGUSKL3xrnxXz43NEoj/PN5GE+rLyJeghmbwz2xiT1z1CL9RYYKsz5tAAAANAGeb3REf4puFOCunliFzADLMuR8YF8uYkQJbFv5RTB0u6zvuSXxrl6rDsz+i7vbod3bBSMAAAAvAZ5xakR/im4VKd901cdgCtt0Mf0BD0pitdn5JUOOTT1aiDQ1D71X9AF0lVJsK4EAAACQQZp1SahBaJlMCN/76z9RUBUcvPliLciIvjqR73/pxCfQ8GX7jBwWgVXkAwndy884ukxbTKeablPB5O9VW5A2hak9PYq/Z45Oo6shsBUaTyA/cKmxsiL7o+G5kUuQ8MYCGvoKy2DKLSe6RNtH4L9/+rIRsEmOj8yoDy4VgxHPKb/WQm4tsTdBVn/g3+/tzcq4AAAAMkGek0URLN+BH0uHOpR/2lJ+qsy4cNxEGsg1Os89HfXB7kgKMsUuWa6wa3E9Cvz9Vu2BAAAAJAGetGpEf4RTMK9DxTPCJrwplTdrNLTARpp/H+2xqo210z7XBwAAAHRBmrhJqEFsmUwI3/vgOOK5ljY9oQx//XXAhhGlJBmz1SK5mkHqP2hvO14o18d/5fUUVISry+PxljiI49Ncla4x7BQYXDweGuz/vR/EY/E2El2sS+X64RWh5rFNorrfK7VEZzIHehE7YBLUZYyCOgseeuPE/AAAAEdBntZFFSzffU2vo9kO0Cnibp453nDAGgY00PiyDpgtljOwFy8YbTMWPWaJg1ksU2ghAt/uAQ6k2TZLipHJGD5DZkUzEzwc4AAAADUBnvdqRH+KbhXAZA8EAuKk34jx10uATWjooIQO4t7y/vwPiFVePKe96ekG16OebeqUuepNxQAAAJZBmvxJqEFsmUwI38oIQZY3yoDaqyVEGIIzrxX5up3P+yHPuwxQuIh6RmqexMGn2JW4TA6S8a45F6aXSoHDu2LRXHvNy93pz1lWEn8VRiuUaT4hawwlUHpq+vaVckWlKSp4IpFdIAZg3tMDq/eCxf/2BMIs4+tEBE2gLOsQ3IOpZN1Nbh8+Q39DhHggVbCQaqg1VnbmMLgAAABfQZ8aRRUsn4Dppp3eFORjRG7B5yFTmmIQ465a/eCYalbPok7MkL2t6szR+A0CkcE/nExx730z4edybGjkTzNz41lOg+CSJyI4X3ooGwbsjFtXwRjxwm+LP6VgAP21jRUAAABHAZ85dER/gwyVRD3jGC4pS6OxRSsTP24dYF84UyHNnutQawSpCF2+m9LZXyJCC+IxSaMtdFhc4BHjgm9gCKjgMNKJ8wtLHlYAAAA2AZ87akR/hNQrDpfZzHTty0SYpUlNuCqQ6/BRPgxApygEU5Fkje1Z2NJ1LTucnfFcqhZof1iBAAAAvUGbIEmoQWyZTAjf+8UQb2oRIDdyJBJG3p2eSRku//vGBvYLMgzmEIVmMtrZl6/1t+3GqkdEZly0y2btApk7ZYZsIsTNb8bH83WEuTPj5wetHe/hfU1579gz9eJmyWaFEN4R4kRUT5qpZ3V5aB0GqdGynWTZ4tnQlWYylh6GVVEwLVvNFqWdqS5yTosyb0fNLjd6vhS6+1ch7S8mxV5lmhCvgGRXTY5b2+5TgsBpT8dK7WAunqXgh1UZsPH2jwAAAD5Bn15FFSyffGVVSHdSb2LfjLSUp6GX0WfYywKsLovtWybvi56QWoaMzSn7o2XmUV6WjcC/zKrkS2qukaOC4QAAACoBn310RH9/iXZSQOIbeTW0680W96zcWiFcp4kIDofTn07/nwFewF7sA6UAAAA1AZ9/akR/gRCXiXa4/WbZ+NyaVYeW6Pf/aBVzqs5yIgtGflctsp+WlKrincEJHtYZiVyZ4kAAAABrQZtiSahBbJlMFExv++BEqXFF1Gh1OTCwItMBlK1vXU7Y1ChujfwgDSEjxp0EbQ9CFRRs3Ld4UN+pKee//cFrSip/MsXQzRJi6a1nf5wlbHjxkRq5Gl8WS/GeHRd97LN6POgir71VTPszaxAAAAAoAZ+BakR/dzl93GFUqBjx10gSJvD19gv7sp63PpHZI6Ie6pbZTy3WoQAAAI9Bm4ZJ4QpSZTAjf/vgPUVXdpG8vg4rlnY/ft0DGCD3fGKCJUJtOkcYpkI/rGxid5A3QTlupDhXOLGx7FAh5xXChqo0uaNM8+rYrg+39wO/wMNbPlztdpk9fui2dj0P/mWfeH1Ug09qs5t6x/FiWlyKrutmeEX6/v0IBu5vp9sVb+RmDgKXu8CyfSKKiKb4rgAAAHpBn6RFNEyfgOol5DugVzX2orzM1hG+r8MU8NyTphLAcR7riuJhLCmXsqR4r+fP38yAn52AQJ3oVcR0wEHGBUJOmK8pqozKbzUBCcuWr4/sv4PosC3bNQ7d4cLNL61+4Hajc2av2ig15vb1u1wugeDCXJaiz13UgQDQlgAAAEIBn8N0RH+KbhTd2l4QxwqahFOEir8WO4bHJp8S/K7MsINQWNL/SSuhWUuwzya5TSxVSamU7UEGgtWToCyOldc6FVEAAAAzAZ/FakR/hNb1HPFzoSznrKuNC2BC8XbeOVatSDLSeW0TBFsdg62GXXU9E7NOBwisegIBAAAAsUGbykmoQWiZTAjf/AOiKyFAmRxUiMbEE4XKUlvaqkfdnPFkuSrv/uZ3Z2K7U4O9pzkfIf7vPupN9AbVD8Dwu0MYNhQUfSc7G5a1t/833BIPTo9W4fFuf2qMVYdA1azZtGYexCZGVNJ/9/42W7lwA+c3UjwmvDt+IQ8gN/ySJxefGd+6V4IysnaZg/17TFedcsedScJPSmsaAWJV2YZr9PQrnZWmT023Hkz08fPsg3CfQQAAADpBn+hFESyffGW008aR19uICCzq5NQWZ90NfrmxIeDOiit/PCFcxaNQKPjDjQ7qd4CRyqeYilJwmJpgAAAALgGeB3REf4BsK4JrMmvsK/sAAGwlDBfyHnOGZN3YuTYe9/CpCnyrjEBzHGVjUCAAAAA/AZ4JakR/gRHB0sI2Gg5rX+0O0q2TvDPSc97UgUfQvJFcWSBSyl1AmJl5+F9Odi9A3zNqS2n8PseQabVvy6/BAAAAe0GaDkmoQWyZTAjf++A44ria4Qw4WCiSmMPHSR1OYI8sMHbiuDjzddiGMt5FFfQsG3Xg/s7xrPsrUAb7CZP61yqPvjLICQfvhRYvztPo6E8M402MdjGVc+5MoikSIp9hwj+gs/Mh6ONlkOtd6lmuN/ed6YdaokhXg6/p3AAAAD9BnixFFSyfex1K7S/556qAF2wZ6eiICEv8kLtY3rwO+XcsSPv9lHbtQilvH7t6S2w93Av2M0H+BnDw0AQGkaAAAAAyAZ5LdER/dvu+WAiZLVxl/UCxr+jB0mUBX88+mWTBAoHaD9tT3RLSGpwK3dfZUneAJAMAAABMAZ5NakR/inDB4mCq76bkuBs3tkF946bT823lmJeD00/zi7mcXVT/glUaS6jNfGEtB5R2xabRDQfa38sKu7jTT8gvmqqlDa7KdrZ/cgAAAI9BmlJJqEFsmUwI3/vrPjhLnY6p+6D40pzcI3EESZhkQEBCdNnAeRIPFFVSkNeEcVI8X5XssF+l+r0Xo+G2YobSzOdbqNWcfU6YbitpLAmr/TX0rnj2EMlad7XNyf0cLZ6ug0PPBXULiz9m54l70LAnrx45mASDWufLQkxiUHCgzbJ5tLv8P8fLTViejM1LgQAAAGZBnnBFFSyfgOolVYNR5FvP/wKasCCjnYQlxD8XXZsDY1OG4dlU0osuRUdZ9269GrNsNOazGKWnVWEXEQ356Raa/y7lNVGogDCtXx7LDbgwD5Dg7QfELHis1MwTs25skxKF8NR6CxUAAABSAZ6PdER/gwyV1/gGl4DoPPqyLS9sHc3sEMPzJO+wfSsvfs+SXVSm/M8+KeWiKtOiUBDrSQJ/dV0wHz21x3Nlknh5MfWPAzaoA5JbIyBOBuVi9AAAAC8BnpFqRH9/nILJGk5gad+eLsoOm/85p1qMApGsq5TXcltUPCfAAuJzBXY34TxlyQAAAMtBmpZJqEFsmUwI3/vBmqSMhNsPuU/Iade4QHot3LB6WrIq1REVRUWuTSSkVV07RY4SBhsJu6bBxcmmu/7p7/XlQZ81iuMoEWB0fIuqz+0F9hRriMm+Zu++c5UoVrMVXoe4liL4u9dlcciYCOKVOu20Yyvw0sSQzjqY0OXZVZDApK7OgpTk3QL5BbxEbzQlLUIc9ecQXGeD9042d6e5Frg8c5maepHMbf/XGm8v7g+lOh4PjQiOo2rPGPN5ca44oRCfG6hjuP4Xue0SoAAAAFpBnrRFFSyfeJlWxHrflwbYugGcTSy6U0LQpmxil0xIJFwT8XFx7PNwczoCH6Fhbo0X15NUcnc1rsuxZVBBQnNcbPv70rYrdQG974QEotol0nT2TDHb+Qv9E50AAAA6AZ7TdER/gw5gGJg+c863gfTJmm26qPpKYjlbgotHYmhUqPnWNt7iyfWEOKe9buyfpqlqllaZlsDSnQAAADUBntVqRH93OX2EwrJauMv/aX8LwP5QXIKtjPfqQBiLL4Nctt04kl+GBFA8tTN8WaPF6LkgQAAAAJNBmtpJqEFsmUwI3/vrP1FV2ybca4m7Qi+CcCwEN6Sm6YngKSVvXnUfM7sChF5+nlbWXM16e6m0eP3IRF/opSl4/wJ8Tmp/j912lrTy9yAirrx3u2AijNyh/O2sWcv5ZF3bAQNrmKGJ5rRtgcXJ2EyIs1FHVEKdlaUfrkIuoLml7BajyLWU7EXhV8KiyXp0EkB1IWEAAABIQZ74RRUsn4DqJeP16izBqJDziqWfWcM1ER08s1jzQXvUMa1oZMXncVBGC9ZI7nV/tGPht7tK805Y3P+T/QfODA/eDIltS+iRAAAARgGfF3REf4puFOCunliFzADLMuSGFg8oLzH6dqAnVVV0V/5aqCM05ibhk+ZkUat4TgBxzdHsk8qTta9BtubHdd0z1ZwtYcAAAAA5AZ8ZakR/hNdbIqL3l1+IwTqfo4tGdH+aq0/IiizAMo318AwUtAm1jEwsljCSWs74wPi4J4g263FBAAAAjkGbHkmoQWyZTAjf++5XDiz2R9bMf+PSkj74lbvgf+CLlKO2lnAEeCo/BovEKd8MOME92netTHHgj2Xxbi65aAN6935tsknj3pq8q8prCwziwqggVz1hNBwrQ1X6JhikyanOkZST2ZeJMNQ7Y9f7rNXp1lNxIE3MyU/nHtKei86tCmCb/jB4vfQ9xWLRyxAAAABBQZ88RRUsn3xlxPNhKxMxOguiBxRkNmyuAT/aa6QxM1PpGwiqB8JRZtraexcKZFHXdXzPREf+tWX2PRuHYEt70OEAAAAuAZ9bdER/f4l2UkDvsZGNAf/yJYoKrKSiCC4Thsosst1DE6Zt2qgCipQ/85ZSywAAAEIBn11qRH+E18A2k13OofglWp1hqldV5ObnQoJScCal5gPbTIjyW5ljrk8khorjjE0JhwAgDO8DAXOtPqLGK9BvShkAAABuQZtCSahBbJlMCN/74DjhLgyOzpkIzRKgaw5w5N1Q9CI2CzFs+eee7UKL04QlzU2OMinPtJrKhxVooB9zLXIdj1wU1RccAL/RK5Wir3p4Z5Rhz1udVe6K4shR4iHm7EJ7ua8uxY5WA3d2NucMe5EAAABDQZ9gRRUsn3sdSu0vuq3oSzNlAM8OUuryflbanNQhQdg2gcKDlBVHwMruBd8g/Ot4PFgifoE/KP6Ohum7YS9dcg57gAAAADEBn590RH92/rZMJhVKmtJ58r16K0MJGmMWg8xAuD/vkLLPx78bVItQ7HSx+t9TsrPYAAAAQQGfgWpEf4puFOCunxZgudv849O8E79Kzy8NTAnO6eR26TxQmY+gxXfNtnCzUB22XkPe8BFERNAs2Mkr21PBuWuBAAABDUGbhkmoQWyZTAjfx6pjxyne0Ht21MeCz9epidJeHVRgt2OzsUDZ9fGji/m34qapaDq3peBBbvkhLBPU43FEf8C/nr02F/kqRsg4pwf4rJ9PEH4DvxUg6fBJQ7D2PXTuvDFf4X/10q3JA3YywimHAzkd07MRAHpUkRpMLgBL8pLEMwRhbPWVpnT20qZSuLPZcdFVIiFASoAoRw2NZovY8cyAqYzbIPWrSsq65p4K76JrBR5ZOJlpYoaGRw5W0g2bASg9QEtRnRj/wfIJlboNxF8URYCsrV90PwYzTaGML8kxduJ1HEhaPv8aDti8cRlE7XHaGXPFi4RmNryyxl8CqT02GdWdUGWMt1v+HtWAAAAAUkGfpEUVLJ/NSlfleNWMCRxiLP05rGSLgkzLFEKRmQxMuJmvzGpyeNPxT3UwwrV3v/0iFzkqNLK3bHJnTrwhtjxlzfd3LB7TBGgwkOn3N18GCDgAAABaAZ/DdER/zdG63bG7KwCn5wCGa3osZLctRFh023EneSsmn3MzWC9KqCBDdw8pzx3PPwNQzEg6Gn5QdiDUHnMAGDZ3RWqH8/GR/hNAU8HnZ5GQH7ni/2AUHpBBAAAAIQGfxWpEf3+cgskDvsTAlIkk3Qv4QtRWQBJxMPzllMwDnQAAAPxBm8pJqEFsmUwI38ep3IcnxuTWqQj8g/rgSxKei33dMugniCkcN8GPABf44Tw8t1Qy3eEsRn18PRXxXi7KFfT3IERwMDz7zpA3QCQSPyLpyVwXLhM7PhboSU8kbfPO1nIY8SIm+UMfGoyE0RMWGKM0JLzYaywMKx9vosxmymYX3DSB8gROvjqj5bAU+Je92mJJE3ZzVPQnpxpwyoFzVDoLe998j0tIsq44A33luvrT6F1Iu7OaZJWo7H1e64kDqoUz6nvhWH05mQLuC2M5Y6ktRKSEcGarBYzmLNzmHh7vugEWRyOLzOV3ZQVWiOAToNJ5ySrGU9bIN9xzr1kAAAA5QZ/oRRUsn81KWB285AKVSC7RFY+doT+AHMerKL1ioOnmHy3pPwHyRwqUvphrWUAjK3ENTLIIi9eYAAAAJwGeB3REf57EiBPUsI2Gf4tmrnTNH6no8dZUfNSMntiCE3jpMKWXLwAAABwBnglqRH93OX3cYVktXmd/7NLucyKFso25IGF/AAAAl0GaDkmoQWyZTAjfyghBljfFGc2weQw5uG12FPp0eBzDiQcuXdbNVOSPW3hp8N/oJG5bAO3Zb373SP/YhHQc+p7lx0UOuTmVc4HVR7yuvu9I8eIWq8uvvC2yoCvXDnsvL6Ml7y1Qbsl4kWfJSSnIKW682NKdd78et0c3c2fAD5VjJnXMrm9zYiL5tXnYF3cZwssrIHOcmXwAAABdQZ4sRRUsn4BJuIyq9eUXdd8kUEVxBXvZDstOfWiaasOOGXz6d2BweVCpPVueVKNIzwQFXuFUU9XBKpnS68hOjBnpCQT4QYBqFRqI+fPbQVzfi6cOp2mU7sSvCMpmAAAANwGeS3REf4puFOCpcPJ2WonQTLkhtA5MCJGSAixNk7uE8XY5myUHB0SlZsop1BQ/gkY1XDSgyIEAAAApAZ5NakR/hNQq4wqUtHk0A/hBBQEMC8kIK/T0jJcco6ztGO/MojFLnPAAAAB+QZpSSahBbJlMCN/77lcOEOGGqjflHZLVcFnf06lApNkw1dyuEV9YbzJBRtMBgx2f6FglR2zHhzFYrtQrtFgGZ+DglcCrKvwM3YU7JhUBsKErOd3vij9YsCRfhH7aSRLCbVXwbVZgzfTpOsK2Ik5aDiH8Fbx0ZTn2TUf8Z05hAAAANkGecEUVLJ98ZcS8h1iXBUMZDLJz9G659Jpfnv5AxC0q3dUwbyNGgKrSkrRowhNA2+4791C3SwAAABoBno90RH9/iXZSQKtBXKXppdd1hcPZsBwcUAAAAC4BnpFqRH+BEcHSwjYaDmtfG61R47GinqanteSqsWgtBpO/fwqOA9xgYQFYm1iBAAAAbUGak0moQWyZTAjf++A44TvouN2t1o3r0TOX19VkdemfIkJQNjJtg2asYKxFD5FgdXiAetHtc+GXPVttB06G314qyFnyAAJOZ29a9/0H+3xR6ynmfY/umA2SihUcWtibuBfPpAHobrBH25owHGAAAABTQZq0SeEKUmUwI3/74DjhPbILSOAstfUwFE4bEuwQsadfXqn7plGI0oW1t1Y/wn0dDyxQqMCm+V/gNYrsp0t8yc1kYM4KZZ/Ep47KXKTj0fca5jEAAABqQZrWSeEOiZTBTRMb//vgOOEuDjiPT6VQ+iwpFIBPfM7Vc576QnIa86wIKthdGtodrF9dyzF6wRQCkAGjjRZKD9pc+qeVKQR+yM0R2LkRny6G0uUgLJfUm1EaS0RgY974VyZDjM1LWcR5wQAAADkBnvVqRH+KbhTdyySQ5ZqJ0+qpbnfk0JbeSoqS8GFf39+Y3CXtR+VXfNwhOcPWv7xsKzRkAg4LHFAAAABwQZr6SeEPJlMCN//76z44S52OVALIq71Pd2EwQJX7/N3fXhfhz3oxlU0IeuIeFSImpqXiYeUYQ8aIei7lagfq9dJeClxCf3/FRuzxdmLsPWEpBDNcG4b7eaMqmz6BvHDqlmnYyHJrEZ8IZKKB4lJ/QAAAAD5BnxhFETyfgOahWDclf403X4Twewl66rB03AAAR9ia2izX/w64YBBi7Jz+yLB2fmcHmEV1G0PRn850sHSM2QAAACoBnzd0RH+DDJXX+AYxXI2kR0rdOucxMrUy+Pe9iIFKs2gPngM17pEs73MAAAAdAZ85akR/f5zgScsRbQaQhHGzw2FWBvNSvTqWvH8AAABYQZs7SahBaJlMCN/74D1Egm5pgj7LpdrUXh+DK2shFxcZendNqyb3B1+J2NOtJQLO0BY1+ExmI8pwJ6OXctPqDeFKHLYODoMFdurTJF+a9y3NPu44TpmggAAAAJFBm1xJ4QpSZTAjf/vgOOIGX8Dywu5b6HQtuh3ftgTaLAoHUTu8zVtxvvBrwf8pq0n+w4fJtp/GtGwKM5hmO9cHCkxx0FIpHiSkAHgDSjdJ9Bu8/qv30MANfg5dNOeEtmWEqPoExrZqFj8JsKEB0TfzZtNFAYsTE01g9Z2gRAsp5MxV9DO5F2sdgezsL8ViWTohAAAAaUGbYEnhDomUwI3/++s+OEtQtJaAtNXXgd4+LqdFPoVeYoa3U/EzzhI37OlznPU7rmyUThBJLF421rLihpccajuYn/QNJFilPGf3pTu2b7OjUvIJM8X3UCvvvycXDv1blPNG3IoYzzU+0QAAAC1Bn55FETyfex1K7S+6rehPNxk1X2rNShPMmLN749hOCs32PoHwwGQnuRiu7dEAAAAiAZ+9dER/dvu+WAiZLV5nf+zUFr8AEzE8J+xc6y+E6EtaiQAAAEYBn79qRH+KbhTdyySQ5ZqKHVStxrf8zUErUDlRLkpwXJGWM5A57aENWOIeQhlA6EA+NMSo4Nf2axnu6UPFdOUMdBtyzDxoAAAAaEGbpEmoQWiZTAjf++s+OCFJBaqu5RyJqCI5ciFy+QfdViKCVpnQDNKhtjLsbBDTAxQkGtiA7qEUco1DRu3mLPro5hMkwOFSJ8jWf1cptWx5ZqYdGH7hd58lsJBKjBOOxqM2gWtMZwhsAAAAQkGfwkURLJ98BvQr/uX5kKHT5/lStdRB3EpNtMxwioGAW9Yj/IJ6BkSDVCrLUVYjVeHAMS+JXklTR7CAUKq0auNYPQAAADQBn+F0RH+DDJXX+1g7QyFsak3Ft1Oi7eT/WrVLikgNs5g3WmkArqxtoNirEsBZgIMNFCYgAAAALQGf42pEf3+cgskDvsTAnaZI7vQ1BYaJISL78CwSAW7O8MKN4TJcSzKr6QpQQAAAAI9Bm+hJqEFsmUwI3/vrPjghR6aRbs42PAWhfKG9hWLBQBYO7VhzdwkmyTb4OhpVXnuVoSIE4n58p2LyiEkgihsMcSjpZIoAbvLWcxcbbUxw9yvpC4qcgM9TIHk9bnA/m1AEfoKukgxmshP3lvq4sLFHjIb1Bb+gDAKawHICIFXSjngCSxCAgl6WAKcq3XBRUQAAAENBngZFFSyfeJlWxHr/xodW1dt25DUmu/wD5rm8TWMxQj7U8N4VJthmBHDdc8Y8kUy0uyF/QVRg6T1Nw3dOMMsOpcspAAAANgGeJXREf4MOYBqx+AWriBGIRbd89ihAU0EQj7cqlqBeUwwo+Hl7YF8BhJvIxqmsu0cqLbvOrQAAACABnidqRH93OX3cYVktXGX/s1aluoN3CrEK5rAEtDv1gAAAAG9BmilJqEFsmUwI38oIQZY3wS9iq8Ru1ZuwHge5YzmrwHbkgMduq9QhK+TbcU4LXoh7vtGatuH+orZwfaI43jayClVBfzKYUUCnemRjycGZN57cgRK9ea1OwkugTmsWQEppkBNh/JdFjjQyZ4JpTcAAAABSQZpKSeEKUmUwI3/77lcOCE/ywLc6XOt25/biPs8yLMMrhhhyESShoB+zNzfnIF0ukeMr6Ju6bgeroNShV6SlG50RR9jhmJI53KndzULEpskTgQAAAGNBmmtJ4Q6JlMCN//vgRKhGBrYv0qobhWAfXvyR/FlCNuqZwJVrpTIPYrZL+bMPL0iiMNcUQ2zWjrP0PVXRdSCVrfqzazsxwiJ/Zm9QLkBsqmq6bw0xk+Bar3U6aKO2MHQ7w54AAABjQZqMSeEPJlMCN//77lcOCEpHhvk51D3BWHT5CL5wZKLK8ZPkfOIV1GODU/wlBTykgfVfM7+jBF4bQaduzQ/+3A+AvwaK0CE5cHbjVc1pJ53H+StJxG9Ua8/DrUCHl8SGEn2wAAAASkGarUnhDyZTAjf/++tBKhInjrA/ex2nOLzwJ3o+r2vWQ3+t3tlTqjSvEVovxTIc8Jr3NY0ry7qzN4IljCb1qcgZ7YvzWMn0aYd5AAAAREGazknhDyZTAjf/++AcgmTgdQc6/hgQ6BptGAKfJDuajnh3DwMwZaWGpVqb7Mh/5bvAT6CAEuDJtJ56z+N9ZmTYlY3PAAAAZ0Ga70nhDyZTAjf/+8HNT5G8Rzz6vkY3Rq88twFJVOYwb1Xg5/k6UGi002nITY9rzS8PPaCmc5nqS0Vh2hJ5DZv/OV29OQcMFyct9AKS8kL3wc8YwFvLGykWfhhLsilALVhO0EfbQJkAAABbQZsQSeEPJlMCN//74DjhLnMdk0oGHhcQFzV4h856+SZ06+ZzX949gBdXZGd9E6QP+F7soJy+XXtQh6tUmGDVgsotechBGEMrKDbwe5/+7BdVED/4EqLPYycb4QAAAD9BmzFJ4Q8mUwI3//vgHIT56s8EtfCC7Nx7BLXzRep/pBVMb6qXu2U65llvCElOC5o5lAlBAWmxsuQ6hG/bP4AAAAA5QZtSSeEPJlMCN//74ByFT5oRDGgg8nQkCHzopVZlIYNDj50PxOrh5uMMbamN5s9P/sz2hGglU1JfAAAAWkGbc0nhDyZTAjf/++A44S5za/LdFJ3sHpSo66b88l6UQOHPKyMuIn8r+SwswHCeYG8PNQSxT+5oV0IxUUsMcQ2ooBv+/Y/FAlBB46wXopImuyCaT/Y0EtV5MAAAAGJBm5RJ4Q8mUwI3//vuVw4S4OjJsHVijyDrIVPe0J5xdGasNWKBtM9uPh9nF4cTcB0XMQNeXo2jAvyPsnd8B9C6JaPfT2A3eKrYalWHgH9tUOCy9Knavp6PwTVwZLwjY3e7jQAAAFtBm7hJ4Q8mUwI3//vuVw4S52Of/J35LxMtdpwPUYqk+ah1t5ft73SHY5tDI+UGs5XnciWB4Bckb81EGVyf92Io8RDVoOwgtO7jiYmvZj6oahKvO0TUvdWw62HlAAAAQ0Gf1kURPJ+A5qFYNR5FvuejVz6I0sUwXTmatChezYmLm4QESF4hGnzo+2N789GMiY7Hg7gSQ7H4n+uPL1I0UjWopTkAAAAiAZ/1dER/gwyV1/mv5SPC2NY+LjsomX/n03iAQv0ER9NxIAAAAB8Bn/dqRH9/nILJAn8Gzsc+imyjJea4qC7Q+WOHNZVfAAAAkkGb+0moQWiZTAjf+/2woAW0kUklgimlgFCIb///hyD//bwn+Oqnuh62tNC5+WgKA+8xn/nply3uOc/ur2f0h1fx1vLWFTcNCmIsUc8q93KPKUmAUmJhrNfrZikefSi8Pp68VDmF+//T/fZ94uc7Um4CesacTL0HWpbC4pwVw6HALwAKhczR9X+xqYErJKlS4y+AAAAAKkGeGUURLN99VLWI5lTeYG9BPyle/82GFMyCKf3wI756PFoOBiQwNLrHDQAAADMBnjpqRH9/6UAHHyCU4A1cOHRliECctKkLiOmAbCVU75F6EeKMbMG1VNe7P3F+zi6z2GAAAACKQZo/SahBbJlMCN/76z44RMcXI48OwBOWvEVAsQP9NTZluNbhD6AJVYnC7zW7VxWM/HOBQiFHH+gn6l7YBSiWoIeoGHmNKgXbtRVNg6EmvoJbA9qtyZntMZgCmaCc/vz/8jRR2HQyWOQpuYlCjrJKp0G1cCrOLzK/bZToKz16r79FSV7J1Wz3m2J9AAAAWkGeXUUVLJ+A5qFJ68onA6dsJzs09N4LvIj3cDBjRfEdWApojACyNY0OnrR86nreca/gKP/udzozcD47DfiybsBvJNDTMeXU0IZfQ5QUfUQrr8Uho5P99QOyDQAAADgBnnx0RH+Eo2cLvj4vNhzCIRfjlupd5qSQ43HazEjhH1jmlnzSvnBcw28QnLtZNA1sBrhY1L+GTQAAAFIBnn5qRH+KbhTVzL5Iwd9W8s9aJtz6PoSUC0gyKP/IlosfSF1jR7mDYgMoJzHwJmTvCykmr5d9a55uBDNKtoDcBCNVk/X7WrPcvXMd+GCzMVi/AAAAr0GaY0moQWyZTAjf+8HPtQiKzcGjVgu3kVtj50fuo8YxUJ07ui4ClGPzR4sW5Oy4COuivineLE9MqxaZ+Om25zN2GqsYwHwOjbCmp/yOZ9thUdCfxibVxPf23ClpXIFCx9pfQ3L3gMhk33abt/kuTykbkgMonPh/peH+AUaF5KyXel7QmQs8VNnRi/oq6rs/xkPZZApzDDMdJbqCvanFzHalDHtkp+lq2f+uQsCGxuQAAAA7QZ6BRRUsn3sdSq3L8Z6sV6pC5e2JBxZNJinGlzs6ATElOuuq0YT7sb+JIlwFgYRvazPMwERrZYQMplwAAABGAZ6gdER/inEMd83hLolJGCEU2A5LWgNtuQHR3ohRJdjldGiSpKiSaCj6sqpbpZtjxpR8tMmPHBZmo35RU1qsmIxgL7vhMQAAADMBnqJqRH+EaK8pIFP2GKqMXBdJoLodSSH+w/mReILD9PWSmxDPtHlmi0aTpQN6QzAxa4AAAACAQZqmSahBbJlMCN/74DjiuZY2QoaXm5FeZph47u1XtpsL5Az4WSZShugtci/xJ3x+JchRTFcSCzajdyIBFD0EmzyuRUcNgfqEDW7dKQFFCS3LIIA75c3I5bEeWMNffND2q8RNvHgFxIiQTv80jmOiJhMyx+fiURB9txK7Jr/00sAAAABRQZ7ERRUs331Nr6PZHsxt0kUckkfcgh7z3n9GmoeN10UT9FBRzEdYc5EHASutw1r3EgM7+7iDD7kqUo/s5W4KMEx5YygXQS8Tag3KI7Y8iAvxAAAALQGe5WpEf3c5fYTCslq8zvxPrc1+6JFYuIG9MO2VN6us9stILo+Cd0DlQvGUYwAAAKVBmupJqEFsmUwI3/vgPUSWjRxdIcdgBijf15+5OcRgW/d/tWcjbO3MZMV83KAk4iqqV860nMKCr1y37e9s5DWqpUgfY/zSEY7/mejd1nXV4y6ASyy7Tidqy+itoBo2ai8/KleueofaP2gYMEwzc8CsZ3c5zZFOxH/DE1t8caigPteCEIVRz7SOG5MhaZRao3+/RoS1pQJfJMw4indglMRRujCISvEAAABkQZ8IRRUsn4BNWWWPzM46t4L3m0bngUEjMFtvaLh11pCL74yxwXywy7hSWyjJhKVUryYLLI7LEmEa38rtO2jQrJ7qp+4gHE4WAIkvSUMYkTurlCHxAi5I4XvUWaZ6Iv16g8K3fAAAAE8Bnyd0RH+DDmAYYFgtadtWEF20O27XaZ+whNwYc3jC6EmY1ypgpN2m0kur0wxTdHYayalaRxghl83RSGlQP7QpaJhHZnAs8TP3vDYer5uAAAAAOwGfKWpEf4TUKkzaVNQT7JA6ND2nVupGemBJq8ChfhNtpxCX4JwNHEZ9s7UeKcUe4X4Tr5S+yFh9nZYhAAAAj0GbLEmoQWyZTBRMb/wDwzHFcyxseeHtK+jVLY3A8ftTIVqtjziN4DBisXVomoyMfd1Sdb6tJXsseTvlPzSSgUV/wZCU3d/3F5VNMl9zrx0xBHmsDhwytPRmDnxl0gvYO++K7bpjqvYn5Af5Uv9cJNeJ0/tM0VGvOcRQAHEE7GPbnUfyWt3/xg/JTgWMC3oYAAAAIwGfS2pEf3+cgskCrRYsbkBnYXyLyFu/tRS+DCNY2/rA1wIwAAAApUGbUEnhClJlMCN/++A44nVBa6M6WcdL98HeebjAAer9Qpd/rGaSXMqDHAum3dQpiQlLhO66/sRs+qjg25SFCsZ1Vo4nYX2Mb50Cn+tgIDPoTLtW+6oNkjuj9xBK8l4mOVWx/rY+B5HpiKwKANP3WwumLK5PSCAmldTN+JXvE8vS5GfjO8R4YQe0Bloy7B12ige4QGTZi7fBYPlKWdxV35hH4+gi5QAAAE1Bn25FNEyfeJmV5SQClUo6EKDR24A91N62GkhAu8Y44c0RYjbPE4+ZbvjaBpON+8Xjc6J1wglUayytVTVHCY+766iHhKM20y4SXfgwUAAAADgBn410RH+DDmAYl2uPzEbnMKed3nsUICyCyeDbyNodbOCziTs7weQQn1LnK0BhMLkLDxK2gfVQXQAAADMBn49qRH93PjJljXXeOtUIt+VecJLVYqr6QwrBQGtz901C+SNkcNsuewdjFgpR0rkVyO8AAAEIQZuUSahBaJlMCN/HqbLy8/wEVP9rejuhgZ1/hTiRpJ6ACUBjzu0rZf+rwUPqu6H7kTYGQFdQa/1tJhAwmL82YZXnHTmWT+k7HnR3DjYHPBWnvdO68MV/h6Jcec8o8RZA/ryFL10MTl1FFKOnx9c5nTtu7WVeqQ82jjFbnYiUJNgKb+Gf4/paAXxjzK3x1VOKsji7JRVvVt7IT3RxUS3bPHqqBqcQlYpSxPlruhYyaW1ZA6jay+0HRBh3+Hyob185kbI6d9pHKI+sD3Tq0UtpSvdz+/HXXIXTeB1Rm/g1ldwTI9+9Xz9EHGdGIIViutN6mdwpp2q3QburL4OaDa4rbM66+0tY5DFAAAAASUGfskURLJ/NY7+LO8GHdi60riHOu6JlBVqwNg1VZMFN4ejipDuWEsd0/ErZ0th6Pxc9Cw1Kjmuq/zyEfYtQmhQsswGsgZ1YMKkAAABrAZ/RdER/zdG63bG7KwCn5wCGa3osZLctRFh023EneSsmn3MzWC9KqCBDfOcj+nGJ9RozF7PJP3C9bOztIPLCjeboJrWZS3xCKtt7JR8Uda3/E7pTMmmqtRMFjpXMNYSOVV2dKi7GuhQUoYAAAAAjAZ/TakR/hNQq3WE1NT/3tCaEtiN3t4YxY880nJUIH9NlCBEAAABDQZvWSahBbJlMFExv++5XDhLnY7HnguBrf5Zur2mrLT+Vz/GnEwY9bAGfUJr3XeDoQFE1iAMpidl0+4X9hFsf/8NFvwAAACsBn/VqRH9/nOBJylVotztSxpovwWJ7DmTJloPNrARFySkRNw2TMBuSJKqAAAAAsEGb+knhClJlMCN/+8GajlL3k20Wujiv8rx1TQrXb7GSm0kYYGmTH9/rx3KO2kVqyvJPzuuFfEim0UAXCWVIQw97gFcZWGwEG++myxOImz7Meuw0MWDCCNZZRC2on3KRz62Ye9OQv/dYQ8UMk1DJlNJTLsnSiKauhVJx4gsD9dWa1Kjljwa7NOoMT0lw+1ORCbOeHGoMPOXwCTsZiawz1MU0E5ekYhfXVHCB3xVjdCt1AAAAQ0GeGEU0TJ94mVbEet+VNuwEtG19Vp9qEKqXEQQaYjS6ZG6EkvByqXm1cZlLTLEjoe/LYqJRdE/SF+jsKjMJt27U1bEAAAApAZ43dER/gw5gF/dlobTz7RdVKB+qPpkr6GiWwyvC7wIA+3QCH4y5k5AAAAAgAZ45akR/dzl93GFZLVxl/7NQUtFWdUUvh5Beyz+kFAkAAACkQZo+SahBaJlMCN/Hq8yjL5u9j+dOIruITmk6F0314KXaka87ZCfL7tAHNxO+iju7Wtc+zvhxob4L7mjhRqS6xh5aFdu4srY3Z9kJMkMwAoPlbslt0ED4rQc9cCKEvzZe8vP5SuEFkEJ8M+mCHUY0aVpfO6zWHHxMB8o43nhSklJXgCtKGumFXlujSGPlk4FS34jNmpqXuBhr+dGn1yKAZ3sKSYgAAABVQZ5cRREsn81VP5vFlB4brUQOGtdboyvxjRTY+YUybkkkNYTH7v9LDS9i+dPj1ZQKYWWFo2eX4uZLvCRgIBXKzUoZJdkcWfXIyanYodRMHyDg+OeXWQAAADcBnnt0RH+e0OjClgq3ZadlqC4PYwt5pyZHpbAtX2CDDw6VQUFJdMB1qD9hMsCCcnWNss5UPmCBAAAALwGefWpEf4TUKm2xG47R3Pr01Jw7Qj0y+sYt+R7i5ZfZWQaodl6QPXmStL4WVN2BAAAAckGaYEmoQWyZTBRMb/vuVw4S52Op1bIfPeM2q9fIEDP8jiZLCKuzV8ftycOzLks1iKOF1EIt3ZtTQ9F5OaVLBfL+oXlHR0aM8Z0FAD3WY96fF4f7fCu71ZyQlXJ+Xqv31x+daLq88norMrMWmEMlQ6N1gQAAABsBnp9qRH9/nILJGk9FNJxMcY1SBNTaLEKBkEAAAABtQZqBSeEKUmUwI3/74DjhOA4HWckmvMs3ndZqaLe9W6S7iB4Q7A88qrDfeboB+rjGK0bIQqpCu+nuXe0EwJabFzLCZi73OsBrQNSTgkPJi6/soSPtur9skY2PhGs7djxF0DyGMzRzfwPqmGKw3AAAAHBBmqJJ4Q6JlMCN//vgOOEuc2omt8p6qS3Ldw5989LDKzuDHw6H9tFiHKlV1hKTORwp2nJBpiWf7T+A6Gy7CdBbhR1LLJqSDyzQfYbByt6HTZQ3hiVDL/V01gSXBPgX36fCFFEJxNZwInp0CmM1ZnvBAAAAREGaxEnhDyZTBRE8b/vgHIVKDNGLdlSHwPwIRb+tFWCuSrukzMABQCTH80n679WSke43t/ig4V9E5EyVF2d/zscVwbygAAAAJwGe42pEf3c5fdxhWS1eZ3/s03ZhFuMfWe2wKdUy6MHZBxrJBq965AAAAHVBmuhJ4Q8mUwI3//vrPjghSM+5BWn2Ge5VCc0ecBc7EOmKermZ8iWDJ4BRe+G/1S2eQMXMqXpPhZIa3jeKjGvYeHBEFF/sMWQpQ0qrYWCvPbD0bfqtAROiYI0svXAdofjDVcDRuxRo9ypWwrQsM3G3a4CbkmEAAABZQZ8GRRE8n3igS8fsin25OPmCjaoPOqdQy40hqX2JsnRU3RfPgOmKeVL8IQC9iFPZaNhAD9T5boSTEP+kwtDeCH78uwP5hUVwfti/rNaWniv7PSSGTiBYuyEAAABLAZ8ldER/im4U4K6eWIXO4QmaqlkN+TQlwR1AnPBlX1lEm5Uviyk02+lXK903CJg7gqnqUC3PCmW786/bPWqAfGLvDpW5m9IU54/ZAAAAKAGfJ2pEf4TUKkzbUQAa919yFiCEgIFdd+EDagYoE1SAUggpxbI2QuQAAABSQZspSahBaJlMCN/77lcOCFIz4vqiFuZcNnUTVzCcRWfNdM98rzWozz33Q7tLgb6TakmSicu29hGCOBKuAKHMQ1NBeTV2csdtYUMxTioTCcrT4AAAADpBm0pJ4QpSZTAjf/vgHIJkCtvXBZfB+14jkra1juFx0W5C2dCXvZjLy9zVPAyjT9j980ElW/V1rmRjAAAAgEGba0nhDomUwI3/++FR2Jc7HZDkW4uJMmtI7OMrf1NHhL48zEpmHBdgRL4PFlMYTP+bOV+4pFezhcm6vv7Z7df8Mo/izx4ojiI42oNc5zuIVj9opEKt8krPnE/wfreAmiZIhYwe8cCEwpfjXXA6bJNO4q/P6CuHDPJaoZseFBbJAAAAcEGbjknhDyZTAjf/++A44S5zaiKQuEhJ4jmNCR534wzFsjpJVuFfEuiUFOAHPYvijJXldwbC09W+F5ucQhVaOaL0yaM7JqDvSBd52cPA47AOCCjU7wu4stqnihjO5ec0ETnzf0diadQ9rPIVw+r7NiQAAAA2QZ+sRRE8331Nr6PZBVkLGHLziB4VNczKxSY25csybjCivPqOIy1NJzulsNvAJVCJH0qinwjhAAAANwGfzWpEf3c5fYTCslq8zv/aYBHXhN31AQe00BnlGw3gDiP8P3ehzc2BPCbVfqZx7BaQyOGwMuAAAABwQZvSSahBaJlMCN/76z44IGSiLnArHSWpo/s1uhN/5B5ilnLO9crpBdkUCLo55nXPs/bbw0m1ni0JxFnWRtKVr2HRpNb2z2Bzw00yKTCsaNunJLxfhVSpb57Qap9q81paINA/so2j+STu6esbY+YyZwAAAFVBn/BFESyfgOol4/M1FWMKl1NWmFHIrOMqFKzWFvxvaf+3gDItP7fuN53V0wplQxqROuxwNisKkYf1A0L9vOJVudeMNQFQxSkGQDs7GIlrSq3kzQx3AAAARQGeD3REf4puFNcwyKvZaiX5qqbL35NCW5PTI+EGVfTC7EKS8nUfK0RPsMBuJRJi+OvTAxVsTOvunx8J/nO06kwqHnuQ4AAAAEIBnhFqRH+E1vUc8XsyRGeV8iG0MB/zFxYS/vkcNLBnuFwB+E7v3VvwY+kBV6T1vy286lrWGXJod8XkgjgvJuvRIJEAAACVQZoVSahBbJlMCN/7xIOhIUAB4s/hp/HFjXsfBSfVzZJCJlAqX8IqjA3hrTr26f6UjDw5+ICxhwJWzzIjlK/pZ76OT3X6i+JXOWGDxxI7H8VsRNRLv8lmJsG43V4JiuFXvawRR78NKMZC3LHXIjfHxkbHB/JcwR+f69KV2/QtRyZ2dRq7sBgyFC/D6+JgWb/L3sBmXeAAAAA0QZ4zRRUs34Ee3aUVIPxx7mDUX17S6FVLhfjTjm/q/myIrqTAit/zI5wjoSwqYdmF2r6hcQAAACQBnlRqRH+EUzCwuV86LLP/gj3H7Nt8I619QX+jP+BvQi+oayEAAABbQZpWSahBbJlMCN/74Djg6QUKvd9e947e3J2lOAkYhiVNr9cDr3ELDEevtAnqI6c+JX9/Hl6IF6osErWJfEaZznOiFMoX0003WslCugmdnvq+vcZwkx0sZcaeLAAAAE9BmndJ4QpSZTAjf8oIQvW1xSZ7x7rM2+/pGwLqLOU5uTIQepdfmTQp8FGjoLDBWrpb5S15nV/7T1MmUpm8t0TDIMYbRQkDTpqwSDDRE/k4AAAAVUGamEnhDomUwI3/++AchU+lEC2Ii5zQFVFGdU4MJVerkDyzoY2q5+KZYNo7qgNGISSgsF+cGU0UYWSZEeb+O/N9AOGI87QduAAXd1wUH2zvdsrzG8MAAAB6QZq5SeEPJlMCN//74DjhLnY6/LPx82l6Bk5pK9HvlvJS5+uXxVzkYeM+SI2nn5l8845JM2j8M31GNo9VnTQBFtuxtZgLhZRvBX9iGRYWh0U9En0SXymWgiKz8LWajDidYO5kVrzWsaI8uEAxpZex2ijK5k/IhCK3vqYAAABXQZraSeEPJlMCN//8A8MxxXEsNPU5MyGp9izFtgK5B1fgvNcD+Z43iIuSvepRhmJjGw//jbq3VYyQGndV5yIhGAhdpfOXdaDwSiD7VkSVBa0B7buOG6IhAAAAUUGa+0nhDyZTAjf/++AcjQRnx3awW7tncZ5B5Juzr9se8S1VRa6KH4UBW6LI1HAegAOeGogjZT0PeHL9xMUqvsOwVSI7eLqasA4yAdPUEaLbgAAAAGVBmx5J4Q8mUwI3//vgOOKGHX9dX0WCIwBpYFz0D6F8pUyB6eshOk+dTkDqLXSsUh+Gr9DuLFPZf17nuf/ueCIOjxOgcQScSRdP0xc1ZVu5Jqv3yslnizRuOXMfXxkx28Z7fAXjgQAAAD5BnzxFETzfhHWkJuNUHXHYaxywUjiGRpqXXHfO/nq8EmC1znAhOoVyi3LAXh5udzgkSjzv3BiqOhlYw5ff4QAAAB0Bn11qRH9/nILJAn23SE5EKnm0eUobJ6rVo9aydwAAAFpBm19JqEFomUwI3/vgOOJ1QWlW2Hq57X357pDMQnGy6t/q1J3KOnfwatDKcCKk0FUX7afaojFcA44+ouv/VS5zgOikY2Jo3gD+Kea4bPlYDa1BAowxsskYioEAAABwQZtgSeEKUmUwI3/74DjiuZY2PamuR6r++1qE53Iv1qMB7HQ4ViqPGfOl2Zby/x/4pN8jh89MtO/osbGCpJthzQcEKpTHULvYU+KEoRclqH1cLUBG25PaeWeD8zxyPBhkjHoeK7nvqLLbLzvLi4uamgAAAE5Bm4FJ4Q6JlMCN//vgHIiQ3BG1Ae3g4noIB59LbXTzyxAiwWYAA78rBEuHCDQZgSoMfigSgb1vVN+xtQOAzXzgo3pxNP2vKC9JlDb7AWAAAAAwQZuiSeEPJlMCN//74ByNynH/TQlKMDZynXLHrMF6i39xuh+J1cO7vKUsDG7+KRiZAAAAaEGbxknhDyZTAjf/++A44rmWNU+hTQglopLhhmAYLK/ZpHdfpOlqk91IgMNgFmHV5vRn8W/Xr8cuOcpnmkUCutdRPjNbPLQZXJ/3ZczxEXjRBuNotgM7N9t+vTv3LJoVu5Gxx4RvBJ+4AAAAR0Gf5EURPJ94oEvH7IgaFXdPjVnk2O6nLuXIwqxnLW/+qohFcx7b+h5swmqFMWhnXHDdgVqWqcEU8jPLgjtj7gTWitr7+s+gAAAANwGeA3REf4puFOCtReQubuf3nGJsBw2LCxyegE3x6o+mF9S3iTu6rNVSi3sIbiyLwJ2lalMxkYEAAAAwAZ4FakR/hNQq4wqUtHk0Ek9zmyjVKtbfh+Ep3L+pF8DNTunjSXvZ3i3A3VLAaZfBAAAAeUGaCUmoQWiZTAjf++s+OK5kJjsYIqxbAN9qTX7RKmcdzUOR4WmyrfmGXlSsY4VW0N0z/pBEVT3CsU34BdrTRs/Mj365dhncAfPS39tR96BzM/Bfv/33oqAmHEP2cAcTsMGtEH+NPMAl88/Okh/VCCmQF4RGqYBUBV0AAAAtQZ4nRREs34TtSIINpafvDmraZs67s4zwOz7bhVRQUEsaKdSs9EDDJCr+0tXeAAAAJAGeSGpEf4RTMK9Ia6xDVybe6bVn8VMBUydoiydjxdfTkz37RgAAAHZBmkxJqEFsmUwI3/vgOOK5lg7JrSIIIA/sNpXxrR0mCihBXocOnszuWQ+6aCY3Qf/ooqPuZPkLUBxMfTn/KI2XKfgwuHg8Ndn/ej+Ix+JsJLtXyYV1Ee5PeWlz4wtnPHfYsD2opYlEN9EtBxxEIP+F6hy+hsFBAAAAWkGeakUVLN99U+XYnHXrZhpGAxLbiGJkujNzoot//snhNxSyWew8NVYOzMd4cv+InhXj9mVM8q9zNsNntf2Lf8crPpKa3BPYJbhAof+Pw38XwU54f+mhs8MZwAAAAC8BnotqRH93PjJguRMlq8zv/ZvacNf3ZgDNS6n2w8JdDTja2l1kfVwpI5RyQrCnRgAAAJpBmpBJqEFsmUwI3/vrPjitsEkocBG1hRaRC7O1B+CbrXs8hHZmhG2DNlmpBvFXAQ8ZzaWTctxOTdptC0ofjr2GSDqDoPTD7fLpKEEg49jN3WOp2j5Odu8aGp4KdhyWrxYcLD12s3Q1/uv7HQxh4I4Z3zALxj7IUwYKgxbBRg9IlcbVNQWTb+a3gaDRS4UaCkAw/1qd0064DhTpAAAAYkGerkUVLJ+A6aad3hR4/uXSqV1B0dTjjrlZ57rkwkz5qH8A5BiiEyjQz9F4C9kFj0976Z8PPdfw/bqaUlhpDnPTBYGWwhyC6yOlloZsDr8qrxeFAV10I20MbiNT/qY9/jTAAAAARQGezXREf4puFNc5iD93O4M+je+bLdricvgJeIzf4IF7YCFGXutVSWIsMd4abDjWFaDTXZx9lcHSO7Z6bVKQQZkvgLpJIQAAADYBns9qRH+E1vUcPAJx4GZtKUaFsCF4uxw8Y6tLrD5q3QdAfoxOrS+VUyv94KTkHGklV3NC7hkAAACdQZrUSahBbJlMCN/7wZu7TxZMM/e3lqS3Hf7gR0yQfSMdS0gp+EYh3ACNJeFZWrJYWy+Tuxmmg9yENzYhMzSp0MSd2WcLpHZqtLjZjQM+Ku+MEhhVWRbHi/Vo+GTEjar+dbIOrWe8xdbaEKLpaDLCjIApOoy+j4hP/sVb741ZZK/0nw+kO0/AXBEe8Neoan+CYCnv+Ow4wt32pXRHgAAAAEFBnvJFFSyffGVVSGsrz+3b5v1Zg5L2NFJN02uV8jUcyUVVzVImpmzOJPahozMXoce68AXxo2JwibbE1C2KfdoQcQAAACMBnxF0RH9/iXZSQkB/vO6m2O+aoEGAw2+uAjDlKDyeUraXTgAAAD0BnxNqRH+BEcHSwjYaDmtf7P/3EHd73hx4e4SGlU8teL9lc635JeBPPhznR0FXQYVX63aPmrj7NT/7R+2ZAAAATUGbFkmoQWyZTBRMb/vgHIjamvtGoxs1uSLcMbcKSXfpDcfNEklZ4VeRp0H+PR6Ez0IVFGz/p3jRWk/Pf7VVg1CUVglF6oKe6mOwVL1ZAAAAJwGfNWpEf3c+MmWAiUqATXXR1gfDZs1lrr5MWJQArYUxiqhfYhdIXwAAAE5BmzlJ4QpSZTAiPxVmeOlhGwz/Fs1c6V1b8P2RQeXHGQ7kXXgq8WzG32u2Iy6GVw/lQvrM24qRUp/8JG4USCPSdqAC5av0hupmdqUQHIAAAABLQZ9XRTRMR4puFN3GPnq2c/KJzU4nbiaIdSG7X0s9TSe4iwC+u9L3527jGFEAW/pgQ1w8VKNK8KcC/2VVebquBtHGTz76FnlStsFpAAAAKgGfeGpEf4n36qqADnsqQPIJnEcRRKczAkBlIY2jZnmdmdek7gJRW9Eg0AAABZhliIIACf++DLeYgN9XYj5zg5B5QJuOQipq5DurWEjnW3UorRPL+Gmp7W/eIApq4hA6GO7MXHIyDHsVucz/aMFYtKTDM8z25ZLQfiNZJqSiDN9723ESzhz79Cwt5M+8vOWuuGhG+Tvq9foXjrYoVXiVDwEsVfUmlk5TXbEqyYI+MZpo4mDJdCI6FifaFbOBypE2MZf5XWR/BCziCfa3ua8heIb9mYyPBNtGSBm/2+TYyBLX+ccaIy4TYZdApJmzt8tHhRNNuRbzsMmREmXlyFjU2KTmRU1Z3ZVzYLYM7mBumK0HRe3QE+vr6A4gpTOdn2j9qjAKj6NNg2MmzTCieFUwql9L6ptX7F/UuuzRuFAfTt932rxOmz2BNjBs10CPMOhJ6Pu9idR/eLvVQSuID26B7YxP1Kc9QZT+tp+YkaEPC2XBY3bdZ6qbfFKSpyAAvrJhEkF1eJhOtzQzGB9M9QVORu8jkG4zlLQbGF6ZqgPB9YhIYnHkO6pyRG4ZEwQQRv3pY66HZyqmP1pcw4M/JTWg0Dc5WzUWPGIAgkAFQc/1Ta7hNEPWBuZrLM3ySme5CdFuvklwbISWRsoUpauBnsVWWN5QeWEHF6TJILZUHeFXnSgp97RhPWIEEhIlKaUd4hpJb7cGhxHmSPkyScl5dA7+PIQv2HUjD/HkIC/Ucyw3+BoGsOX7lkWmxzmjnqmrW7wAKr3betEKeTC6MuBZC2Eie0lFVnNNRkT4VOQIj4zh9jxJYkGeDLwst9Jm6zQ9hKlPa63RCfp9NnaHWjPsnJjMX9F4gwR547dWgbLvYE7IuBo2AK2pqD7u9lEBFGBQEfEy79ZlxUwHxhYbaKRUP2ScWs/5ZBg+sELMAUKpCbKjPE+vDv9EAYHgEOzeiT71raTwKS0j8v4q+LNmKEzMRQT8YYKkcXDIj+hIcNxaNKpyvuGqlycqBZerfPCw7NyeCJ1RZg+wkCwjrzypk5ge+Jt7o/7sW9DR3cOslqBHx96pj5KJC/hClCIaZB6HFYXv5qM9NlNVSZUKsifIwMWVlTFH44dQrGr49VglslLicK98bhApMsqpHLl6mvI2nl6M00yk4zvnor0ybl+78LwbD6/u2JjQ2C5StNvCBcm5T9DSRQmWz+xgw3XM2cXGR9bRyAeishmO4rIwBRYnfjW42CduoUT5uCKdwarOs6+ZReAa6Qm032qE+yuGHc+0eaRHUxlZw4pvpia7X72bi/pxC4qWFZD2iC6up9fej8EinMBLHUXfZvheqdQetWdPndbIazJcmP6J8bggLDQ4mKE7dyUF19ugTbguFq279TrvB9ynED8sJa9KoY20QSV4oLpYcqOxvEiV7vzPO9OS9m3f72BkB1HiYIFjaJXwqYGkTXiuZJWrbol4+U0Y08C/7GdlrZPe6z/onp3lz5ODKTA0KXD5NMWK1y/vxPV51UDcOCw+Hv63HNJxyUPECSyqdmyGfW9a3IZeGpmcvlwizm/fg9CiY8lZSQs8QDGJpatu/aiCPaZPOU5MgVYO8MYY3Bg83oVFT8iz4wJf9K6OnKc1nq8GDGNbbnNkX1eqV7en2c97ZImpgcP2BFIf94YIjngEbP6/uGWVEfFyMJdl81YnorD2O2896phTQJZ011W2fThMp6sugWNSD4zAbCi+DLD8su8T5rb61RQ5zpFRjgGvMFoA+Swos5EfX7Q+pWIiU3BmZVbzG6lOYqF+5VQhycWb+BBOdBy+Gpnne+KvK78cSzPaQYJ91kufnulzaeUUO2hFOH59m9d49JV2fpLZStgYBSIJTM/9Ox/Ms+HFo9hdmgYr3o87P4pp1ZyLChdnLPD4OKjtr+1825bbLYc/XTrWBAc/9du/48l2gB3e8DXTZp0m+PbqLMrA+8QljaKoWb0phbJdacof3gIbAAAAhUGaImxG//vtcWI95nYiPte4ifOiv2MWrXa6oc5GwmVX4VGEPL3MSQvvQOWz4iK1KtWLJpeEgjzPL169ZcnBgeMLsKHdUM2wjMTrsHdOygwiz5mAKMj9iGuCQFkHUsm7iXWuKx/0L39JtWxH2DB5DrjSNacwTXnDHcq4Qo51kaWnzhKTZQYAAAArAZ5BeRH/f5yCyRJTRNLqZxFl4gyBe3yn49jBqp0OksTJCkbjJxyS0eiYhQAAAJZBmkY8IZMphG/7xIFmvHEwzcpixf6g85KZsRgAFuJ83fQKmW6HunVKuoozc5sQw4UsnqcarVRces3II69oVcoWZSdF1HPDrGm+RQUode0jm5khpoTxFEJWu1alofRiI9UOP3gOWkENJTpyKfX8mUK4iEwYLYiJBNu6J7fSvAmoGVC/qd36RNT1oid+703GbBe6n57kGF8AAABZQZ5kalPJ/3iZlhm9wMbG6wiDR23j91OZf89JoXeMWa6ZGKfp7xbU0IXx54eFVJ0ZoDzq3qCKdLExoFVdKuttpjLhgTdmW67ERQDYRx2CX2eTTzYWTJl97oEAAABGAZ6DdER/gw5gGJdrj8xG5yWNjcN5RzyeYQWTwbg/YJJ06mabrNeaaAIQL3fS05zTmWWz/Ddjyn7rd/fj4m4Wz+Fu7Q7lrQAAAC8BnoVqRH93OX2EwrJauMv/aYA5jWgK/npk++GXjyk7I/nxAuL2+vMCt3X112NqjAAAAF9BmopJqEFomUwI3/vgPUR3rMO/T2Y2VhLfhP72rEslxqZ/WAOcDYyisIRnWP8zgmcr3lUb56tCV7LscA3yQHB7Y/tXggxHENASIP9bZsUZSafMEOPrKIJo6LpThykrgAAAAE9BnqhFESyfgOol5DumVJ0S0cVIM7am+NZ19h6NTyxakBR/Ofn07JLTms/iq96suWA4PuOQi/6c2rG5IgiADAoCMybzxtewQXqSlO0vzLnhAAAASAGex3REf4puFOCunxZguX9wJlyQwsHnoh0gIq3mjZDFzCFf8XeZpx63GVk7z776u7t3kmCYNdaCOmNgmx+oj+n8mM6gYJcRfwAAAE4BnslqRH+E1CsKg5vffNV6KJEaafCYgryLkW8mp8xVM0IsV7+RHJ4vavIBIrAgydB5rvlAxBRgXHTpdpwIKS0filN50G/LRYjXoqAeXMAAAACfQZrNSahBbJlMCN/7w6grIVFJtmFWBNfhha03j96D82U7Vx+2nsoSjttpRDsy/gZVh1/0E/eE4qYqKmuDqanK7zl/QXx6E6sICctx8MFYLq7fS7pEM+ZN+K8Xv+uUh4vwC1duL2ZLxK0/Cen1zyK0H9sY5VGD83b0M1fhLQs9AaMHLOe55U5g3QoShXVmVZpN6cYhIcm9HRRzqlropIiBAAAAO0Ge60UVLN+BH0uHOo1OpQs1to+yKHFwaptQ5BFtgDlNMU+icULmLGxujc2EQSoFU5ygASOOA1Zvaf/FAAAAKAGfDGpEf4RTMK9Ia6yBEGEbBgFIX6ZVi4oZ9kz+DhKMG5HG7B0dtYEAAACsQZsRSahBbJlMCN/74DjiuZYyN7bnUg1vzKqdRlP5zSFw1h8FWJhA1x/pcA6l3LhJz1U+gBgFSepwTB3UvN0Ol6E6gnujgvTnK96bFAIXqhaB1HSHC989x9BPETy+h3mNxHpOKsWF1XH3KQjb8NMCBha0VBvTpEaYlRrVz0cWvC21LWVlsZ88JWjCsQwtjBrW5+pViDUFe0ils5WWRrnU/nnwnmdT65Qet+sgwAAAAFFBny9FFSyfeKBPNcaL7+V5X1IGhoxsOwcUUvCWuXvEiV7QwC42PiHrznZ2MxQBbEYjkTmgN7viHiaaidZgm4UkJPkk5c2BG05COu3x+CtcJ/AAAAA1AZ9OdER/f6W0XexwRhdsnetPageFCn6RV/+/fWfxA49EgEMtzDU27kTwCwBlXxqw37RvOyEAAAA1AZ9QakR/hHV8g/txlNnB3G4bTYjFW6A2jIYi1+5c/6wZQtVDrFEdc65a08gfqr425KneZpgAAABuQZtVSahBbJlMCN/77lcOLBfClrEvtK1bf3ouakjSedY0d3/QQpN2eRz/kCnq0HwQhsB+x/LqObYzJyoTEkOmGoN2mLlvk6EVJl5TXjbWz/ZmHsOadqHbIr3P2nKT//Dpvb/k1/evkG6EZ6FppcAAAABDQZ9zRRUsn3xW34B8XvH3qT97ssK9VB5rRJg15lNNNG2qjubIElb3HZhXKsOnG844Dve5W64/ukmghSwP0aO7uOMrwQAAADoBn5J0RH+AbCpxyDalGnNLTih1g6/0qOdJjORclUnHAhTKEO+MwTEmokHDhXkKpprte1wLoG6IbahgAAAANwGflGpEf4n36qr5vCXRELsL+sPgx40RoVK7tlQbDF2IVvmSd+z9ZY532dnlS8XL+vQIDhQ6PYEAAACpQZuZSahBbJlMCN/KCDlDfdqihMhNsI4qRGQFbeAgSReinIWm8Q4v1leH/5Oe95uqoaeIdhN+wGZD579BhEjNcllwMDo21tloRbtSe1uLHorlQEF2gL1+3QvtZ6Ww+YbiKc0Q3snOvvYLlV6VBkpFTEa/TljT/92LFpy+EncdIYAHscQDkrRLnVMWpdwe1oCFjeZaru/rcF7oZmGko8JpJoEP9QBsOq97wQAAAE9Bn7dFFSyfeKCcVlysnd4zeF+E4wNIhpJk3CJk50y4UPNHY2/KywiCicYiTy5aF8x2ZKGFuGUKWkFqvmv0KbGkq4auUlZxFpHVEF/DCBLAAAAALwGf1nREf4USy2J4yUop700PwUbNMjzSdo+LuZe1m+BDNgF5z/+hAiVM4yhjVIaFAAAAMgGf2GpEf4As1+KMQ0SOqKb3+q2lDv62ayTIdee1EruQL1OwR7PM6bm0uHHhVSam9bCpAAABBUGb3UmoQWyZTAjfbiYOYl0PwIa4cphLVFbEikyeFvuGUtjj9rexKdFRrEbKH+W5YBn8+aNOedn0EvYpxVizQmiRCHVywZ3ukr7oPH9177x5o1M447NLUgHH2HTsuf+vyFIgHZ9FLK0jO0NIH3ib3olKIR6JsVfteJT2lw+suG01vzwcItvz385NPA29gXXf9D6eaOQ73pXz6Bx2gc4oTPpo+WdUakNFj/2HXkeAtBbWl4gLb1YcKDA4y8G1BFnlxHv88wkrXKwI6Yol6fmR3bJI1y5dVNJQ9vudagKOV49JX0x1oxZdx/+gPH+URlfuXDOCUGqDljPZF5g8ZqvpMQOAuyLrrQAAAE5Bn/tFFSyf08HdPKf1cgoaUwH1/y3y+A+ulifGqb8DB2/uEFk2vpxeiHhlya1/JJnifRcTg6+KPhDWKH2iLwmqbqc3CDxYX3kZrHMTaHkAAAAlAZ4adER/hHi2UFyJkuq62ZFmJbEncTEmcCmczdiVr+IJ1OPJgAAAADcBnhxqRH/cLDyOioAOezAE4F60h/6HDW+EY9SSjz2LoTUrQn+4MyfhBEJdSMVRLllcvr+A6MoQAAAAz0GaH0moQWyZTBRMb/QSJIKCLArV/OlT6lqtITQzgbnr8f5Eutwk9fMwp3wElU6h8UpJ4dbhFq7lbGPtQxwukMLM6TueYoDKwOsy3Kfevszyh9gQA/RGL/2Z3AABMi/YwMbw/lzJ9i4vtQYdX3S6A0UZk8DF1yH/sWuw9zQKTzfzoQmtJ+btMgMswQyVUHemhiShtbMmfnX0YTBCZIu9ZV0hmGWGErVuj/rW2+UZU9nZmMEUv1AL47Rdpv3vU+5WsKtGVwtlbafwVaVPp7mHgQAAADcBnj5qRH+Yl5YFRBc8xrboQonVRcNgLm4flPt+8unUaOv5vuV8sqY8Po1ab+Ohyjt/kuvOUGmAAAAA2kGaI0nhClJlMCN/x6qvIqNCZDKPAcy41m6RBftLAOTU/JG9a/HDlAjWHlgLUf7lw2zHlW1dN6NmvoUeaQm4Rf3oT8dK/A6+EPtUha31C9R8AGA9Pc1PHLm78eBmqhdW0/ws5Wu1XWXR1uYbEmmCJLKmCQsDRBZlUnN5H1wYtmh+CZAiqdf6mqvSNUeSmYEs+cI/4mIy1lPLWsph1V0bChaKID0e/3kgcXyOXc2Aqc80V+XWH2ucNKEcuCZxIb+YLN8Yphusto9627pXQnMmE21y2Sy3yM8koWxxAAAANEGeQUU0TJ/TtTGaqRkjlJwPbvsmSI1lkdSxXl87ZF6YN95PfBag7xnAxXc3Gmj48UCzJYAAAAAaAZ5gdER/hHi2UvUrOZ/a5c0ermQddoJtsMEAAAApAZ5iakR/zdGPifqOf3ginCBX+vXCxrjKKxMmmQAUZhX0/pxpxbOnD0EAAAB5QZpnSahBaJlMCN/9dub8vkKP19B+pw6kf9CAjfIRxnEwttHyxpVXRku3ltHlZ/lkYpN2KXe14XnfKsJvlIt7C0jFqLMM1jsU7x0dnYDgnxa3KRXsBPz9MvvWPMhxArAEWOQPLOVcPX2KE3Zj8lL7CigZAVxYZRZxRwAAAFhBnoVFESyfgOahSeu9LaQlIdYgvy8B8MWOBXFmIgdWQSr2vo4u4eMzya1/JqL8142I8jEYL77i0hisMK1dxUEyVjcD9xg19ovelUxl0c1xZjCxVSzJLLZBAAAANwGepHREf4R4tlLARMdvFFuxCgK+j1kYjkNsSNdqSwreLdD/OSM1WLdrguFezNg07LOAkL0dD/AAAAAqAZ6makR/gQeUVGdf4170agBA1IK1mKqcCUzcAxkJipPFkjr8XVNqedToAAAAb0Gaq0moQWyZTAjf++A44SZIX6myKoZpfSwx3KsEKMJToKDMqUqApD1eL5COD6ggxaIYo2oBHwnAsVolI/njGgWAuS6GnAALjGTHXTap6KqRt2wQfFm8fHkqvFwMRzV8ZoAfwTNrZTwwkwSZyDrc4QAAADFBnslFFSyfex1KrddINV5fGhaSsbM71utgNqArEne5/psmiII46ZfwBD8NIJ75jHsYAAAAMAGe6HREf4pxDHfN4S6IhdmW0Xai/zbnkB0d8faNtHQnA0/PXBa9GH5DiHe/KxjzFAAAAB0BnupqRH+EUzCwtvKZsEyQVsXmvMVAOIOnZjpXqQAAAIJBmu1JqEFsmUwUTG/KCEL1X4ozsdj2k6Ce9zYKOCXqaE6ZjAsmKXfbBn6uZEpzN2/8UmUwlCJhbW3/dILPwAkQYsQakIryqg3c7gMWCnsjcddWB3LQzuP1+p9jOVcrD3Xr4c9IecX1eIa4w8BDPtg42swGnJPlLgbQhIpfm0zFW/KBAAAALgGfDGpEf4As1+KMO/5+YX8sjGGhQz1ZkZQqK5ID4Enmz30wLs4t+OPMVJ4+oDEAAABTQZsOSeEKUmUwI3/74DjhPbILSGmfMEGXTtKl+F4XYSvJwzWWb1ZwbfPPDCkNhDyw6pPEXZYk986ThQwSCTwvFR4sqnGv09HOAFhHcX57o94v+6AAAABrQZswSeEOiZTBTRMb//vuVw4S5zapKxN+T/89Irq/c8BeyQjHqSyDZmxoDKD5I2ZER30lxBIf1WoAoNUW98QLYvYdm6zyybpZzG4HMZW8BLCyOn7D2+fhOzraEYLfhZPKgWhn6VDfppf0am8AAABBAZ9PakR/im4U4K6fFmC52/zhJFXvzNhyan1A8C0itT4SM2A+aqKJpuDqCdva6Fc3DzqWoU+0qgoASrLy+6MMhd0AAABaQZtUSeEPJlMCN//76z44IUkFqq7aTfAxQGo0WrVRPKTltKcUg/xnRhZwIe+DxhONjdOh8IzzK35Wf3PTjClIJPbgHcpwpxkfKcBfS/pqFbk+KY4ESO7kMpEgAAAAPEGfckURPJ+A5qFYNVyzptbWOl/Lio4wgDNAqOZS327gG2S1pd88/6C26jWrIjYrV1arRkukKguv1vHOvwAAAC8Bn5F0RH+DDJXX+1g1VdtIjdzQ0EXREbOPgvpkgaGuH7Q/rjrGKps7wnLqUFIdCQAAACgBn5NqRH9/nILJAn23SZGImicBQ3DUHDjZ5s7/AV1JPRKG7gT6le4kAAAAYEGblUmoQWiZTAjf++FR2JyyC1iitS+dYx0EO9MkwWF5beGuOG1ZTyHyr3lFsRHyiSoIXq4+8kJeLsXYAYR46+4k/GgvTAKNRfO38FkHkrRwtEd/5K6QMpzqI//YNmn7wQAAAH5Bm7ZJ4QpSZTAjf/vgOOEudjsWuSxRrOyCZjbojYuMdD82NYeOPzHOOkboB3OZcVutWe2p4e2RjezDJTnX7ykx++hEsasudgkCPoByJ77cStz8pSx2Mc8cb6wHv4P1NBXSxRODEmyzR6vAkTP/h4qQ+1JpzMmK543nMxeRvcEAAAB+QZvaSeEOiZTAjf/74DjhLHUdnTIr/kiBro1rbcLYxqvkewiG/h/qLASFUG7vjx8I/NQYp5JX260H1dnWv/6sKBFvjZO2Obn0YWxpuvjdc2STEdWWh1bdxgv1NBpRMD3oT8UI6ZG2HYtJKeyaW6r/H9f6pel3kaQA0Z9+aPHAAAAALEGf+EURPJ97HUrtL/mgidYi0TC+wiKzM13IdxH0hFqOnPLWN0qPOjbY0oUxAAAAHwGeF3REf3b7vlgIlKgE10AyLX4AJlNtNgJsyqIGh8wAAABHAZ4ZakR/im4U3dqjRaZYi6ufJW41oOZxfaude2YO0qLYL6p+tv2n+ZEhy2TKBugHpw+CmXFI0YFOx6UPFevBE6JTjViEV+0AAABnQZodSahBaJlMCN/9duV8/kCFIz5BMs12wfEGnFGw8X17TL2yufvVuIdnWBmq147QWoxvvU4sQmiODTxcU5IxrW48XUQ0l/UgndpZkMnXV4yvU5Wl2xRn7JOqG5tG1Wjk4im+e43neQAAADtBnjtFESzffVPfJNmH6c8uGGvN6Ua1w3dSVxyq52qBEQ5AQ+3UFiUIgQtkx2AipnIUrZN40brHaz8w/gAAADcBnlxqRH+BB5SnfnFFCFzzDarOeQ+n3x9mjYHrxdG+1G2y/RedFEpdZPCX1ehb7HH7JWFYy4bQAAAAukGaQUmoQWyZTAjf+8aEFzqPqIfb98OjTtEtGYksUtQHqef/YQxv2TlelbawxGTG5fTWV8g+MnDomIil9fbxikYH9XQmdZy8SyM4eT5UbSB3yEfjPZuf69KjR9VOuSEgKQch9dgwqBRFURNNn/ZHanJPjS32PTuCkyyVSbGosl1nzCvs5IhumOPeVjXmkR/0r0gyXtZomPanE/2OhFNpNTOQ88Qr5lABTVZ/tpyQT0FJkYsHvnwCpfD9gQAAAFJBnn9FFSyfeKBKqU691GQnRt7wT2bAwvsFlFpJYfJL3kI0Ejx2m9rfTT1n/rKrW+o3DIZl+x/z/z6dvQ5v53+xVKnddxG+xfI7Lp5qTATEEOogAAAALQGennREf4R4tlJElLvUEtQcMggGQj6hD8aXPWAOycZ1yxgmNuAfy4hKPpmOIQAAADEBnoBqRH9/t8AO1iKF8NFd5NzOsKBLGjI7tyNx83qt1R1YPIAs/dScM1VYQvLo/32AAAAAQ0GagkmoQWyZTAjf++Acgi4Oq6OXF4I6IPPdwzJFin4336rAUggpBKkw1FQCveDhbBVqTv1fDTBUjHvoqU72VIneOhEAAABfQZqjSeEKUmUwI3/76z44IUggGGDyEJHJ0zdy4sFJbfwH7K8r2+9osdhXBVubU94QlL3+0o16lzaMEUX+s5hDiLE+HU3QAsuJzTdUVHWDE2crLRMBDiXB3Nry4va26IEAAABUQZrESeEOiZTAjf/KCEHJoIEIqOD/LBqTEmdKsw0VZiIHhVnR0YHo2X8W67cycOsus55XJ0Jqd6wt0yhFaWJl16uDYWLglY5OB8KN8lyu74G+Q2atAAAAZEGa5UnhDyZTAjf/++AcgkyM+G1Ouhcn6Z2Bc0JS+SytB+cEXwbM6nE1arVgMsYki0/teQ7FCmWw4Us2D3g2l9Xae2Na6Cm9q7U3uGVww7Ls5vAowUGvf1+1PeObba2e2poP0GcAAABfQZsGSeEPJlMCN//77lcOD4iDd1SEdD3NEMCzmr6ToyL4AwseqHC/p+7NrCkSHk4I3vaFZNOZvIcubY0dF+R+0CSzGj75kO/jtMhszh/+0NqcZ6+lfDVkqMjcjmxjwAwAAABHQZsnSeEPJlMCN//8AtD5RPCcxJhhtvitPqIz74OQtZ8I0WhBpvfz/vdzPZ91H7CxyeX9rMCN0nFbY7I3EWhbNcHJ6VALu4AAAABEQZtISeEPJlMCN//74ByFVfBVBzCrec0q51pjCVag54dw81MGWphs9/vebHu4Slqm4xB0i/esr2DZbj8Jf6ynXo8mnM0AAABbQZtpSeEPJlMCN//76z9RIyID26En9fyziADdwqGitUiUfERx0K1tUeLRCixwBHGA8E5zp96UdN+aXTYSjyBXM2d52ocrCKE+B3/tBnS29ZiVSe53MesjWFUmngAAAFpBm4pJ4Q8mUwI3//vgOOEudjsIu1r9UcGS/BbdCjokhEK4p7ca9fqqtyh9/H3PZV5npSGQTV+CX0Y/0tMJZ7Ct1LC6HPc//deUSi7xe1m4BaSc4T6i42PU9PAAAAA7QZurSeEPJlMCN//74ByEPHhakORElru9VCJYbH2Yv4/MDI9PbvboGW+izcFJ6I5f62A3mPgFpsbLHqEAAAA0QZvMSeEPJlMCN//74ByFVsZ6j9D/F5iILSjxSQW3rwCXR3PeR+KppsBmj4ZYK9QEYuzckQAAAFxBm+1J4Q8mUwI3//vhUdlhbNCiiMoi7GsJj/JUSgwY7nhDk6vwDjSJ+VmDgx/g0bFuY3yAw7VdofesWO5yJ5HEM++OyZWsYg5c7kHgGfigSgYsdX+uEbqvt5ohgQAAAFBBmg5J4Q8mUwI3//wC0LvLElkFCZ+0R4VL+qEuV2to+z9d8c68NZwBRUtJQhv7rbzrKpuzMyK/+cyHm1wEmdUHgCIDOwZNtkrWOlccD0IWKAAAAF5BmjJJ4Q8mUwI3//wIE8MhYsf9rP5Kq/l8NreVCsrDsaQlU0Sws4aOq1OMlKviW1Ryzdx7P7AeQ57sIFdLR691f8uoHHExNezH1SYct2J9kEnZgyh1ie3x8zmVrr3AAAAATEGeUEURPJ+A5qFYNR5FvP+sUOFlh6OdT+5AC/ALV9jcn7A759t2aCrwOgT123/pve/kWbUHHscn0wZSsua2ijypozcsfdeqQdpwo4EAAAAmAZ5vdER/gwyV1/sg9S+lWzZCVlXJ1bZAB2pl+X6+0J1jK93r+nYAAAAnAZ5xakR/f5zgScsRbP/KCTlPetUWjSxs9oMve1ulK8unxRGQuTh3AAAAcUGadUmoQWiZTAjf++A44rmWDsmtE6da6Z88b/lOIgyX4xpZm533n9QjU4LoQqKhjYGX/tuiDl25viW0fxsXEC1UQCYY3Rxdqst2ZmTB87g9gEeOvp/+r28Wjk/lHMrim4Fl5xsal8b5UTEFgA4fhZ1xAAAALEGek0URLN99VLWJA1xrwZZgYEEgHgZAufc/vsiPVWvEmgCIWFM+Kch/8zg4AAAANgGetGpEf4As1+KMO/6oPIpubBAR5YI/mVvTOM5vFZZyOcvnNGLH+tubJVUzxsAV/kUZ8urzgQAAAIlBmrlJqEFsmUwI3/vrPjiko5FfFdReOlnYOY+aIn0vCx/RI4Xxu3GUBJIy2hCFvR4nFjOuCPWyJ9oC+4fg21JT1abT5qVCmt7KlU2DoSWZVc/093qGhFeYaBYh1VP3nhSxd3bKcoYQ++seKOZ67SE5jHbJWSiZVndbBPLb1CknRwfGQPMcrNB8gQAAAFdBntdFFSyfgOahSerkFDO4h+HT03gu9drqkmwjsbFo15C1ddns1scB1c41Z607J8TKKQBbJ7onEPnjRqY16Vmf6zqYenNHNjGFMGswbhSFbJ9SUbjgWrYAAAA1AZ72dER/hKNnC/BciTqtZU2dJJA+iXwBG47V52G2lmuPgOczguYbeITd64lNMXnwX8HviTEAAABSAZ74akR/gQeU4JrA+ZeKVXUJDRTmUyM8cptsQXz8j1l2X0x6PNzWmC3LWN9LRSLx9v3DwnzUFYMnVL55McNf8uFZblO77WrPcvXMd+GCzMVi+AAAALpBmv1JqEFsmUwI3/vgPUVnq5reM5Ssz2GdGe////9UbdFeHz6++R6IoL51iZvWWBTEo+w1/QGffKUbsN48+o5YE+HeD/8Y5yb8QquB9cRfL+hkT6gX4h7jUqOGLkv3aekdBT0u6Oz7hhVkeKCmJJIp2Op94DIntn/GHvxlz+XmeSvT4E+MDxEl2LA/adI6b5NUJrAj+zyjLOTgs/yPrd+rgMhLzwia10NmoNTcVyZRErfLqU+jEWkhGUEAAAA9QZ8bRRUsn3sdSq3XJB2xN6AjUdfvrPxLfuDa0V+3Ef2dPSVqKYe65mOrLxPdKn/0htbtlAPjz7kRvwqK+QAAAEIBnzp0RH+KbhUp35fRKSMEG8fVZtySTaTW+ISyd3w0mPrfX5tqV1xZY/iIFRY4QnnDGyfwxM3+zO/RMboY0LQswmAAAAAyAZ88akR/hFMwr0hrrENdSs5i6TQXQ6kkP9h/VZrhl6OVIn+wZmWbKSNLJEYN6QzCxG4AAAB9QZsgSahBbJlMCN/74DjiuZY2PbTAR+Dbj531bEXlChD5lrSbsGOo82bYl0bd2pDpN3WLdD1xqLg+1WuvVraz9dukzhD2dWGkCfyEavvctphVPI8P9WmurH2tugHweQ6rl55gisLfyBenMiMJ2RMJR+lzc1/3Kri59cV/bv0AAABIQZ9eRRUs331NpBIZUVg/sF8yeBOR0xrwAbP3QN4HMGL/ULx/iGepxoAV0VvJUzUyZfsAcCY/TGky6vUdHs44yyXro/pZ+OyaAAAAMAGff2pEf3c5fYTCslq8zv/aW5r90SK19ZwBG+JkrBByKG0soG9BdH2HugcqFw9T8QAAAJJBm2RJqEFsmUwI3/vrP1EmHKvUWSjtSRZlSgrv/liQAkTTh2MMOMOBFt7k9bYXOcwf59p4EuKaEKx7Y7E4B/IgnuToImCd86g1ojTwETdli09SA1Bo1HNUGbOLU83EsefzyVWJSbRzGhntADwtCJI9oXid1IkLBW7ZtDjX8pcEztbON/vJFzL5VTnKXjD+8vULEgAAAGxBn4JFFSyfeKBLyHXG9IGRaGq39J+w00CRNYhQsfHGXw+pdnBYY3Gs6lwhAdgPpg9sKS2UPZuAS9WFZHm4tDlDX2fuE3rMgJtYMXJgf90BeLeZPecojFY3k7mIB86csvefkqVua5t3fNOkykEAAABMAZ+hdER/im4U4K6eWIXO4QmBIq/I5sOTQqiQ0+cMCYfZVcNBwzmnMTca8JpuFSs22SO0wP2GDUnITSfVdndW3vFSPjn1UhcPn2zggQAAADkBn6NqRH+E1CpDUNBG/lajRglb3oTeJln85BeqE22nEJfeJpIKOlpztR4pxR7gHZelJU6RMf+Cw0EAAAB0QZumSahBbJlMFExv++5XDhLnY69ArjxbLeYiTGAFVjwNqRrhTUG41oczZ5EyBH+7dDwYd6fxBYacxJZsYzZglHmsDhw0rqZxKveLxdRtcYBSZO+liDMbXCCTHa+6LbkPiqpeWPuCFg018OjxE3T9dFSvl00AAAArAZ/FakR/f5yCyQKtFixuQGdhfIvIW7/HBKE6HEqyYkrz7gzpEkFnbRTwQAAAAJdBm8pJ4QpSZTAjf/vEgWa+oqExUYVp1olVvNbP//+8M/kBsXizdkbJTTB/BTvKR6wFbAMmXZWZPYbZpxPe4+VXweLkpoPRtXn7vfj18z8A5fxL1tFG0YfcZNJvwmwKmP55Yy9/It5uUnKSxwWtgpa0Xl7FOVHv6yBORxpBbFCVCVvuDxli2pBcO5IupgTKpGUhLJbY1ESIAAAAWkGf6EU0TJ94mVbOfACgqZwUyCA6PmiiXKLu1snrfMONhCCh+U404cQOLSIYQFmfISdXlAdVcHufy5BepvQ5vcf8R6jFC3uBu3JcyYPQL0WMxnXzTVqebdZ5fQAAADkBngd0RH+DDmAYmD5zzrhUCgMcWoHHz6RsHLXEpnx1pjFJJfd4M1l3M6R12v+1lUKb2QeenUrWvLsAAAAtAZ4JakR/dz4yZYCJktXoPfrPf/ePA889a0VEm8L/I5Oa072z9TwIkrlHV1UnAAABQUGaDkmoQWiZTAjfx6suWx+PTqHcCWZGG6ifNKCv4tUI8MKauO6cTSadburO7W9zwXrTpt9gXX3b5tdNvqHh7J5Cvv3AhRisceNgb7YMFsR1N0Rhh6/nt+TZ5PAFkbKAp6iJMn//stQ+Yz7VXAyMKB41Sngwiu8Skg+yDGuw5z0H1q2PuUD1rmITqnNCuI9uO8eg8uswHpLKHevgR52SAYvS2/Zid6FbE0GnXdftWRsiroevQ44GTyV0fPUfROdAmGxCWF/QXCmGcCNFD14g5Bu8qQ6Z+50EprnL1v3bi6SzBTqLcyvNtZPOKr8PQlWxgIVRruPOgaVr1gzPI3yA4IoHPl+VKgdU7iUksHmCQG388/DVCqccTrz+TsolCl0IqE+rbPC6A7gaixgqv0m6YoS+bdEtfIf5PoP5qHH3emI7wQAAAGJBnixFESyfzWPAWZ3eFR2Uc7F1kcb+oikzxVIfIGvLglj++hQHAONn2Im0GJpEyeStvP2EmyHct8tdtegovJLgQJsfAYqJrPchHLtBboynly/y1VWmuyJjSjqD9pkggLaG+QAAAHQBnkt0RH/N0btmFIvxlCtkLVaZxkluWs7V94CnZHEt7HGqrMRIPPdfXob/JPjyP65+5jAvNmhHfPXIgTy/h8MzRAfvoqjTEvzPxiXoq2YJ56mYweyAUgN5cy5ot0uF6KC7kwMci3mHrUqnOmZx7LD+L0r0zQAAACkBnk1qRH+E1vUcO/PuHwIOngiM/6M9xXvQGSxikpxLdsa+ZBaQwprqGgAAAHZBmlJJqEFsmUwI3/QsHiz8vl29gYzDA4k5/ybedXu7tKu7Jmy9KPWg3HFWvQIrtmp1hFrco7VDr7pRktQQWV4z7wB7QtPmuaPl7AbJCErL+KhMJjtnkudgiRpiC2pKeIZGLJlUjeD33utWpNyRcvZCNrnmrvCAAAAALEGecEUVLJ98ZVVId0czXNe9uHzxp8h3Iok3r2Td9dl1B/uClTrEreOFzvFRAAAAGAGej3REf3+JdlJGgwv+mvTSxIdWSJzF/gAAACoBnpFqRH+BEcHjkFzBJtn43EJVcXoHHh7rOnULz+Jbw70vSl2jJMWY6YAAAADcQZqWSahBbJlMCN/HqYgKv/p9ybhCWa37OXOoId/HNKBU7lW0ZvUv8UFLcBimbwNzzimRNxyV6iYYfu9yXFwXxNgztFgbdhKv7Yaw6GOdd3/mgssr3YJxxbqYri+BQ0gsKvOU/2bHepT/BrNNhFXz42GZ1RU3Z0OvhNwjsYp+atyzw6K1RQrGQzJdsjyTkZfv1NrgMxK2Kyq4o9lOjDMLVkWn0P0yn7qosC4D7s4H2EcqLGoPHxwbUKQbXH10NPUXB5VgGo7Gr2hDQFFekB222S3MELKaTdqBDIJgwwAAADFBnrRFFSyflmG9bXlI8HDsblWlpyLDNyCFSjeTVQjsKx0UBaiRpefXBhGOTqqno+v0AAAAJgGe03REf3b+tk7jCqVAx466Oq+thoM+UPfpYn/yReLuvw263iAhAAAAQgGe1WpEf9wsc7fqSm94QxwqahE+mbpNVa4sm6pOpD7UrOlV5wtZ7l4RNFEHIy7UPZEdoW6mGZQa7151EYp0yV581QAAAEVBmtpJqEFsmUwI3/vrPjhLnNsGv1BKYkfPo4cRV2egVZSmfCgBIfMKkEiLVx6WQ/+NmwmC1TmNjrOX/bDa9yyO1iT8c5kAAABEQZ74RRUsn4DmoVg1HkdZAFs1PtxQgT4vR9+7PMyPdh8OvTRv0jh6+gWn1+BLNvVSZYdeehlg2ag49VZovj1blaVazskAAAAiAZ8XdER/gwyV1/tYO0MG9HP/fzsOk/m3SBmPS7dS9yDlmAAAABkBnxlqRH9/nILJAq0W6Uz8mfnp8RC6e96BAAAAkUGbG0moQWyZTAjf/XblCEl4icEcUzKWIv5Sdz8xTnNoz+oTMnKI1DUFY37mUJgzqPkO5fZTe0TbtfncuGXCPon6ZueMlRvrNNcOK0XcUnwUA6DgCWf7RjzKphjvKjqti5kt8Itb3RMRziXZc2powyXYrDFolCRsJlmjzet99mqjeEo9j3KbahWDJHugtfI+sDEAAABqQZs8SeEKUmUwI3/74DjiuZYOyToYj+uJfx33+jD4oAI40PfC774m4iIF8KayeVa0oDX8hLv0VogcQ+NVxVP5bwWpReIJOBcCQwcgCOGRpzCDSfOHlF2gui/AgGYMXyRtA7HKV3fvBHPk2AAAAGhBm15J4Q6JlMFNExv/++A44TsoRqu9LyjXtmglYXGeW5gtQLLkwsKg8q3kkLlUIQztAr/PynJc//JwqcgO8CVKpBFGV9xrTYudyf09uTycpUDGmoGFkEx7CA8vXK/aaRGjjNBEByt6qAAAACkBn31qRH93OX3cYVktXGX/s02n7XjMhSSUiDHyWYYBLWZFRZWZz3x48QAAAFJBm2JJ4Q8mUwI3//vgPUQ2HX+ATzBMT/Pnu9JYrDfMcgGJDPQUmDYQjMezPulMGzFlZkx2U3ZoJt844xJBEP55P2VwxxHi0Kagm2yw1aPaGpOAAAAAOkGfgEURPJ+ASbiMqvfMgbI0jPmxHlOwJfTnXYgmLO8WX2zfxIhIhMrFN5/kkUGPI8Ahlw6/NHrekA0AAABDAZ+/dER/im4U4K6eWIXO4QmA1FfI5rFYZ4x7XjV9ezMO/coIVv++VYWmiu0D0vx9YATno5hyzWlfZMIOxG+n4NgyUAAAACkBn6FqRH+E1Cri6kk8ZHaVcOy1cqQKjNJ21dDGpEKstb4T2+G8eEqVGwAAAFVBm6NJqEFomUwI3/vuVw4S52M4EcFcQqXCxPpHL56kgKVqUEL4iX1OpvXCrsV1IIbzGfkz93JglG5FvhdntnM54UGgAg1vWYaNNzIhAfhLt1NhDflBAAAANUGbxEnhClJlMCN/++AcgmTgWTEr8if5xlRqXymDzn6MNQ/cQ/xqLIQJZe8l+jwFGz7ZyzAhAAAAXkGb5UnhDomUwI3/+8GQRIyCAPHVGA6jY1LoBJIpmQFHa7VcmtTnD0dyhdFuminjvN7hC4NUT3epXOwmxGLF6aq4Ii2C54UdlPXegSNE8ZyYU/M8lkhKwaBQ3c4nA8kAAAB9QZoISeEPJlMCN//74Djg46XORVi0pon+LFQxBrPNpEf7Xo0uPcjk0LUC1j144x6PKcVHoBOoaxOjtlY0/9T0DQfS7OAr7AkgvvVG0TpjF3TYvKWzgWOhJhlDHw+N//6f8iZoaO4xHGToMil4OnpnInBqExbvXmyS2Vvee0sAAAAuQZ4mRRE8331Nr6PZDrA7iBiuQ71R3LRXqWPzXk0TAwWCC9fePi3Wd1bTnw+uPAAAAC0BnkdqRH93OX2EwrJauelrxcclUb4bVAQe5fp/eXH3gDiQzy+DEM9azx07mjMAAAB7QZpMSahBaJlMCN/76z44IUkFm51/N882W7qXe+E4ycMma6jvo926yxLryVknk2weTfZ2Kk6jWgpuPKF0slNo0l2zy51T6byfECXCpBr/SZLSr2Zvw5ZFKO5Xae+BbRDytVWPc5sse+N94UDFDnnW+6PAj/4nobZzzGDAAAAAUUGeakURLJ97IKXjDBk4ZD9O3fhiXA/r/7ECLfoxJW97egIUifkFUbHvwS6RWe47OV/aXVeX/KUQSlcRcn9TJuFfK0IGrgObTvFadJzhLN17/AAAAEkBnol0RH+KbhTSvs1uWWoAkvTwy/vyakxvNOg/BBWPV+NaHNLh+5xiy5Fr5LFPB+8S8aN/EjTnuye1l0X9H1OcBtFfbZhvpHufAAAAQgGei2pEf4TUKvaaa83Ad2VtKlPt7siVdyG4UlrNxlJdXv9DrWc6/IaeuwrDtCf++mzEfYyS38DY6xAB6l6lKKK8wQAAAHdBmo9JqEFsmUwI3/vrPjggZKIvbXzqFTYjQ0VKKQs4oElaRxg0TMLtnBuolgaQrMekht8e3yiYuzVvUWS4uxk8FlkEzYMkLXoLetWyWvtKs75yWsyyQ2ot/tDY1AEK9YY5RVC4GnKBf2CZWDbyfWUm6mixHfVzgQAAAEBBnq1FFSzfhHWkMDYBGZBTsgt1r8IRHNzCaJx/FGcgoS2gz8e56BFv96s1sEBQAIJQ/tBYOx4of4RDU8GM0X5gAAAAIgGezmpEf4RTMK9Ia7m5tNYSF2IHBXygcGjvFaWjSpcJnmAAAABXQZrQSahBbJlMCN/74DjghSQPww2RGtCahPmX81LNxJQE36SVa0/HkFpDiH/yKhx8Jf2oCjWFoB3VUBDgDk4+gZhe2X6s7RXS3cJEENdXgtyyuS/MquQRAAAASUGa8UnhClJlMCN/++AcgmH6NMQeU+cJ8XbSKWyOtUbGfsUeNz5bQi7J2bKGDZFcCIS/9wonu6aEEEC/Vv0zBzYbwCfuAFMkFd4AAABTQZsSSeEOiZTAjf/74ByCYiKau+rz9PkmS/yTzhpkUb60QzRKs8rt26pMRvSkNKiKHJNJ1FXPEAY8I2SSl8M7BoAPB1/QE5fHYPJfTYyyOFiXPTAAAABoQZszSeEPJlMCN//74DjghSM+tgpIYflbxOyTaowlDPlB6SR2Rz0K+CLVW7IJQBHpZDwt/twQggloO94udHLowQ2+xoNwlqG66Pq2mXsLqfv6rBArcwHxAtHTOAqNxvb+8RDt+s4i3eEAAABSQZtUSeEPJlMCN//77lcOEuDI5sHb9JiVElFoic4mCH248Kt3rU8OlDO8sAxjxeauU3i7fx+UKBIcj59moJiQP2LUkx/p/x+YG9GV+7aTbNWtwAAAAEZBm3VJ4Q8mUwI3//vgHIUZmOEMdb1abutjrmvCM7g6DL5L3B9CMFG4FEvH5aC7uJB89RKDxxsntlPQ94dBlp2GynTalQrRAAAAaEGbmEnhDyZTAjf/++Acgjp7C+NxYI7ubVhpgmYIDTXkpEYV3zaW3e+rWLbW+n5QGSD3eCIEPH/DXrfalOt+n6Y9n9BbjQWVz2TQZgWcsZU85NeledliQG4TudhZAEF61udo0AyWq5F5AAAANkGftkURPN+EdaQm4CV0M/rYl4NgK8z84UIqrAO/czkG8l0D2j8OwsMsAK/qF8Rm4aPdZ1aSCAAAACIBn9dqRH+ALScLlpa/+GYDTEx6TsOAEaCmCVz4kvAr3WkxAAAAfEGb2UmoQWiZTAjf+8bbiClhCfpGNHNpGB3o4p+l/b+6Z4oYQfJr38IF6u7qllXxlyI04Ql52TXK1ASkOluayplXmJPCRsnnsdY6EoVDYN1QrkAehRAHXvIw8/6QTQFSPQMoerdoDhvcYrQcyeBLl6zvPm/NUvjVAocrsKwAAAB5QZv6SeEKUmUwI3/74DjiuZY2LeS976KkuiWNhYqOhk4zW/uyrwXJteQNDFt8DdlcWf/pTc81xcC8O9/HfjSyQlZjUCjsB7Da07CzCZjkJnPR1uZh2dxHaUVulf9ID7ewjMxqhMmoJ2NS8tg5XpdDvP5Tbn3fhqRHNwAAAEpBmhtJ4Q6JlMCN//vgRKlMjfz3lEv3p9Vb9LQ7v+pEQxwC/9zazAZ5tgWuhGJQVkP+XR/kRzKBKBvW9U37G1HqvI0/phIKl6ot4QAAADtBmjxJ4Q8mUwI3//vgHI3KkJVcGVSj2/3lOtKOxCchm5X08HMnVw7u8pSwMbvYNglvmYNN2GyFVjbMcAAAAG1BmkBJ4Q8mUwI3//vgOOK5ljVPnge6cjNMAzkc8L1XcMXvFe0EH49/37e6kaHHAytVBME3/FfgqqJuUggarQZYNRM/+85bFVwuRkE/Fcn/dlzM/xeteEV2bGJ+FQYk/XavueTQrg1tVZiDMduAAAAAQkGefkURPJ+A6iXj9eVzRU5pOyeyOkOK6siJQvqqq3aZI916sjslp70mJB60/u4PD1WP91IjlLAujUAFKxopKSwbqQAAADYBnp10RH+KbhTgrp5Yhc7hCV8z4dV9+Z54gcLr1jfWaZfLnK9DZ+XvNFoFhKEO39XZNyhFmPgAAAAvAZ6fakR/hNQq3WE1NaMz7BQMa3r2kDm/kxI5W0s1zZYe/BmWXhyBy4UWxPwkoakAAACjQZqDSahBaJlMCN/7w/OhfFZHbxAlIZNdK81M9DmgvjcG9HFcEgcYE3SBH+UoGuOALJjtU1dC0GdzLYklNnQHymyVpLOjxAFGP6fgPn39YtnABx6Iuq2ORMv3owIyCbWEiR2ya42PDNeEHorS3YxGwxydzC/f/trN+g10WiIaC/CCYlyr0UAogxl0oEvestbuqfBIwybgP3u9WkPg01fJr/nOwAAAAC5BnqFFESzfgcsGNr4KTYjcPlel1/Bo6Hv1V6W/hG/R9ldaJrwij6IiqixTyhYrAAAAJgGewmpEf4RTMK9Ia6xPrVbhCu9V78UOxQddtZ1H+yVxt1YlCTitAAAAeUGaxkmoQWyZTAjf++A45y/4Q3AZy8yKUZ0ax+WW8rFph6+XgDZSKK64i8BQCeecvNGQcmtcwgPyfTLBG70A7JuS0YXsmM/BbvInVHP96P4jH4mwkYmXHax7vqt2xdUMX1/uhLmr9k5kGIQN4Okuar3BDoEYDE27DDkAAABQQZ7kRRUs331Nr6W+8fh7tCKREw4UiRr64auFmklKazVIROH/WyOZMwEYrHbLo8JiX46PwpVX+28ttm84fKTzy4PcjJuhLtZLcNIyPsq30DEAAAAwAZ8FakR/hFMwssa67x1qPHT0B/ovzqciH/v7oaLqbN40xsh1mUuAuQt+HqX4eIGgAAAAmEGbCkmoQWyZTAjf++s/UeJu9eBQwVda/BhoIvFxND3yJyVrLYLCHqElIVpP4XAimELY478ot0SaWrEjaeTwo6tQAgnQW9o85VeSuV++MiiVdty8QsF1uI0JRpePl3og3F+kKHovqBhl43XtBd/3SOTYNkTJ+qnMfrxWequhajieauKGpmiI8nGWaOCva4RwfQCy5sjxyeWAAAAAZUGfKEUVLJ+A6iXj98CjRWanezX04BDJPUQppE9ePwMqecBmkgRLAxx0WEB61woNndqUocE/nExx730z4eqHj4y9zx0y2MLBK82XzU4taZ4D8bTJZKYzQismuE2WWwmBvIKWSg7hAAAAQQGfR3REf4MMlUQLMpWe9szsSyBoDD6MEYNMsQqchdl7Zl8i/O4VYyLyiKSFtAPuL9KgbCvEV7oC6OzFcVnFHup8AAAANwGfSWpEf4TUKoZwMnlUAsRqwSiK4QsOjRuBEUWYg1Fh/OCd7nuMtjeiQ5Y0p0H3UH8K47B3zCIAAAC5QZtOSahBbJlMCN/8I+Su9ssFpVth4bdS9WpvaBw5IeSHpdzxkR2ATv4jfUsYMAFeDTkERmjPr9LCIOQ65j74gzHKGy/tLT/fSolZ/CgvRfVDG/FeQDx2Sn0fsAWhLm9uQCA9XMosomEHoAokfqWTj0xGM8pbd0O/hsjd+5tP6naL8dluhRdb+uAPr/SfD6Q7T7RzN8WnGQo8foNo+PncggQZoQZ+z4TTUAfSjZP5rs1umERs2Zq5dIEAAAA9QZ9sRRUsn3xlVUfvinbaGcgB7351ebN7ARewJep+yTczsF5e4c9MzR3BPzL1klLWjrwlgKX46g/ueGxNcQAAACoBn4t0RH9/iXZSRorN2wtOvHL3rMKNXKhCQl/5dOgxZ26rVQNQnwhs5oEAAABCAZ+NakR/gRHB0lFsfre7+Bp9uH2y3Z7UlKi36/HXi/VxAc5Xd7dL3iB/+RZjRlA4K42VY9VgxSya8NR183XysSDJAAAAaUGbkEmoQWyZTBRMb/vgRKwKj6jQ6kUl2X4UvnvSwPjz38P+ukOsefshONivg0qJBk4d954f/XJfgcHAk1n/pDYI/bf9crnrP4+4WYWy892CnxvwMZGMki+3NlF1QyLUNlnRuvpjtnNZqAAAACQBn69qRH93PjJlgIlKgY8ddHWUf1/1OqaMkBqSnNFs7onwD8EAAAB2QZu0SeEKUmUwI3/74D1FV3aEE8G1Z5xdLRpbXXTFKbRJQjbqRswcA7HOC1tolGRV5dSZGwQ06qx3Sx5PBVNDREabAfzyxszynQeyjRmijN5plDb8xwdIxI2kDBBZy4ZZsO17xE3oB1eSevH0ZRSYUMODXwfK3AAAAF5Bn9JFNEyfgOol4/fAqtD33iwUqaOGcDWLURHTyerUUMGLF9/dTKp1sM6fV6ymzQ6Mw4D1HfNW8wMEkjn4rsCZqQMIPu+olAZgJJZsnp7/524VaYHVA6xNJqbPPxWhAAAAQgGf8XREf4MOYBhh7xX1+mWbR3kBqO4P31bzBBgLVw3Z4XE+EiCqAVQWrnOF2DmSkXwx2FyUnBheGRDCuJQDXml0oQAAAC4Bn/NqRH+E1CsRJUo8Q5Zy14nOo2/Zu7iAW4sc/CZgv2ou0lLXrm+dUmUgFE56AAAAfEGb+EmoQWiZTAjf++5XDnddY2O+hkpJCpGqqN+ioRd5fh6l+P/4qd3eXmdhQXWVfefzZYfFULX7dQb8xKswxzxZH07JzrnEToy3x7fjpn5UQgftlM6VOlHdwbe4xz/rGQAsG3VCvdsv9qrf3RI0iuneCBDzjBxBQQU+fHEAAAA3QZ4WRREsn3xlVUh1l0zC3bh9rNn6CiPj9vvy2i9BnSdiA+sxg5yXAeOsEXDrU7U1jqGjoUWpWQAAAC0BnjV0RH+AbCuMI/rr1UXH8O02Fm98S4fObJMRsbWn8m9/CguTUAZ/90CM84AAAAA6AZ43akR/gRHB0sI2Gg5rXgcbmg8L8jqVYtAxxuK4okCllLmEQxPSy4J9M1DJFgcgzgIfqUc+3vAZIQAAAHlBmjxJqEFsmUwI3/vgPUUr3PZU49GT22lznRKAoTW3jOcu6hD4ioK89GfSr06PHj3gvvsDv6O0stGp0b2VSw0E8X+l5NjE09e9NsWOSzj1xWS/c2SZ/lV+QRWe9Jn6Ul6stNTb4pd3Udctlx7cQDgYj8PLbS3Ij15gAAAAPEGeWkUVLJ97HStiPWYa+Sau356/Rr91SVf6JCh4jlNlHdN0d2PuQNUBv5wh3HT7qTnVQ9xZjQCXFlPLrwAAADEBnnl0RH92+75YCJkhya66QMSa6cPWkN5ju+y8d0nfVs+H50BoCjgVtakB5xqMlZAxAAAASgGee2pEf4puFN29kxLsO36+UxP+fgCnqU/gYafq9/Xg91ZfFnrFqEyyhYJCxcIhXzmEbbtohOCddejoKdPZpuX3ZH5v3KZCY6QQAAAAkkGaYEmoQWyZTAjf++s+OK5ljY27Vg48pz5a5+KYWTkP7SvlMxka68KtKhbaJTrQL/WXyvxd5weYVY1sWu6B/pmqrj/2QsNv2zSTXylPXylMXIEg2zO3tnyu/tbvMWheX6G7CPIGp+b6nvZ76NvL2x3bbS72nPa/776VwQDIVEqplwdPzD4QLQXjlie/F4jvAz+4AAAAYEGenkUVLJ+ATVljVg17/TcUb4E+TZf2O+oJvmehu9/UXCOUucMo+X7ourEbvh9gcx2HsgqBvE+aCc3LmoUYFCKGFauVj9/DXDPNwzQtla4m8OGx+lrBEXLX2daTLsmB0QAAAE8Bnr10RH+DDJXX+1g7O/ZshIy/nY4r4hQrCWJCq/V6Esd2r/5asnJk5RR4uD0AP8YUqepXEN7bk03ujsxa0LKkPI7KOCxKQmLykHBjvGJgAAAALQGev2pEf3+cgskaLZ/6LMcK/qtVrn4NRgFIs05qa7vLpR4QkwyqPFyrTb7wcQAAAMFBmqRJqEFsmUwI3/vBnCJ2LJ1QWuVMjOxOqnfxxjGCcowTZ19I1NAFB4qvy5q6Ryyp59tnEv6iWYDfMwcl+posZ/vpl/L1PLmpbuboHvedMTA2YWPl2SGbI42h5SLocWEUqLCIWgdUiQWfje697D/3h1HLpDMRw7LLhOe089llXSaMRlYUcNK3n6wBz15SFB00Q6ihaj5OREXBbwYXwrPNfyFffYIF/ieB9Rb+wytAu0366VH851MnHmh3OAwWWubcAAAAXUGewkUVLJ94mVbEetQUbYugF4NXqOohbWdmnQgq5IikjMhBE9Q8RRsBeKAYT1t6fLH/uQHXlB1P0aEmEa9WNTZMw4CcPsqSexTdEwT9nlxz5VCl/l2owvjxAPVh+wAAADsBnuF0RH+DDmAYmD8VJbh4Nh0FynC++OPAhiHknCimadJ9Fe44UHpkWLeV30b4H+zBmAv2TIXxIi3YSQAAADABnuNqRH93OX3cYVktXGXvMuRi8Jr6Agx/PtZNe1ki1o1HHbdUiBBUvLJF5NG7lEEAAAB5QZroSahBbJlMCN/KCEGWN8UZ2OksJynlUtJe/T/Ze5KZIBmbExdmqiGR5669n82a+AtLkfKGHN3Rp9oU+HakLXqfrXevKtLZuZD5XvF8lGiaSAkwjZjoDDFe5iz5YPRcz+fZww/DpgTxvDkkfcVTbleWnl2qZXgSTQAAAEVBnwZFFSyfgOmmneDDuxdaVxDngEJwiO9Wub5u3Rp6pIgiyrZeIaoVDz56LILYAvwi5cN7K66uYsQoR7hKrNuY4y7NEoAAAABIAZ8ldER/gwyVRAsMRG9wSk8dVUzVjXcbzTmj7t7tRjQ5erWnxVd83Ef81CF3ZzwpPUsI8IPWyeBOnVhve3/z/SC41nZMKfRgAAAAPQGfJ2pEf4TUKt1hNTWjMmchYe5SAUxvqnkg1Qw/OPUc3DuBY6sbZu6KUUUigXOtNJ+lIbfnS6xPA+D1gYEAAACqQZssSahBbJlMCN/7w94eHXz2QmwfLeHxPiiF3IVqCWcYhPGBobbysZlGTAT8YM63rCVG0s/TvYgSnMzoYtwQNurg7iC/yUMmiP2Uprp7HMupZaNwHNYfBppfsDg25dFnATbutk3GeDH7PkeFPJAtlp/QQQHUqqrP5NTHvXTLswfs5DENGEt86b2DoFGS/7lW7J0rnfMCbX9Sb7lVrsBZBpHc8xP6+kRPmeAAAAA8QZ9KRRUsn3xlVUh3RIT7WdkVbVtgz/uvdxDmkzwqBzSZwQ6hkamoiDiVfFO1EuO/+qDNbLAf9gNMQ6eIAAAALwGfaXREf3+JdlJHCJoZnuoVnuSxVAFzEoggTy1eBea3xrPnNe7ErbeJ/f2BanpBAAAARAGfa2pEf4ERwdJRbH63u/xqE+Tl1J8P1I+BCos4BmXSZ8BBe4HTICM52Vgu9SxEvs5LsCSMoz5HkLWd8bnJC+JCYqHBAAAAakGbcEmoQWyZTAjf++A44rmVvKoKa41M35Ch3G13lWY0bzQloWTKokZJkHadhknWNHe6v4MeLVqdkMXc1nQjTtXrx1ehZRUoti8Ga75LuX20ESKWe4YPhUknmNLVPxY4QiKLz912RQJJScEAAABBQZ+ORRUsn3sdSu0vw0w5exICsmzddDcONsXZutRG82avccQdJ8Bo2mLUJmrRj1OGq0Hhe/tpl9GLDnt8r9hgyXAAAAAuAZ+tdER/dvu+WAiZLV5nf+zTESAGqSIRhNll372WSzssDIEYXudYpMGm62vGTgAAAEEBn69qRH+KbhTdyyn2mc7fr2krca3/fg3IzUQIeXmjyL0W7kM5SHpgQjyEjdPMwM32smfci5GMHzDsab6+RotVwQAAAPdBm7RJqEFsmUwI38EzmEnMThj01oz/8HmIHDZoyWvPgj0C4DoRsDjr2VEdnSM5F+6Es/HhGkAFmMMlxRO4cBtg+v8v225tcKWm8Y+JoMo6OE9FQeGP80ZRZKcXaFITb4modkoTFB1MOoASSKLuo3eyB/nUlTi7UcHk/YyDhoKrCsj9X4UhyECyG+WaaL4nUsfz6qAEXwJcfBhGzK5jzIx4VGoXgDJKsaf2nsV0cr72ztqjUEW7uJ8MhZ4H2jvG8l5wh+DH4SjA+Y4thclgZCtEvFf6uTRjHi/V31/zDeILkYiFDRYuc/uiN22xbZOJlfR4cWsUKVWAAAAATkGf0kUVLJ/NY7p5XjclhLFPOm5V70UHjDk0y1W8WLMxUEJtyFr8LEtGYvu5dSu368Egd9yNd6Rj70InGXN93csHtMEAObKXQji9PQYG4AAAAFgBn/F0RH/N0btmHtB4in5wHCfslGyNeY6FAHgkfpAI6sJ7Jpy/UlPtc/zQjmrUm7P2PZ2IpRrB2bO7p3QSB9jQTHvmkqdIl64ZFcKRLHTYUzBmd4ons68pAAAAIgGf82pEf3+cgskaLaDSUkmmWOh29hQ0iDCVqc+7lecaq3AAAADgQZv4SahBbJlMCN/J932y9VLbElOTIQnppSdx4sQM5Zt/kb6lOpe1/qFSIxuuTjdoDeO039Y5hjcJKGrutCBLHNOGFrDitKwHUBx+wCBz7FNfIB752PfbJllVBPkMGIlohflE1kmAWoEvdKQ5RVEIJBo0NM+7Xc4QOrGwlm1kYRsr+YzMe/a8WW7js/AQvf6/YKmK3nv/WmouN6JrifggGWUpt463j+VsZPXJT1MLmlu0ZIJheCkppBM5RK9Dn0pWoHhZqpp1LS6KmH+OIyOhplBJZEO/pM5ypaWA+VX5TGEAAABDQZ4WRRUsn81KWB6R7ATs7BLkl+8C/fhmkvrMuiNDekBNsZe2yrQ9LlYwQIG0r6PzAA8CdpaaTzI+B+gavhWMMYl1MQAAACwBnjV0RH+exIgT1LCdOIZSo931hcnbNVCgfvajMBfvkpSj6y79gzRE5BB2iAAAABwBnjdqRH93OX2EwrJDk110s66jvwel9iFQh+hHAAAAjUGaPEmoQWyZTAjfyghC9V+aCM9y0tNTu1YqY7isyMrEL1ISg0j2C7zj1UUqJZwpqqKi/AnzZgAJPEx+BbF2G96bWd4yUYwWlUiGqh7yxHAdI8eIWqxQwALIgmV6FsYG0q8g7JGYi7Gagx16ZFfTcfzINdt3vHY5lmi4F/KgeU4H+REU/9BL/jlA5mhrTgAAAF1BnlpFFSyfgOol4/XldkOz+vOrEWi1Q9ySLFxK10HVuM/uqIkihgxaTiEJ55tGiib34VbPx+F1i9KTeR0sV4TCay5VXoDlJbQJ2X/dXyUkz7tlo17BgibiD+MKM8EAAAA2AZ55dER/gwyVRAsFrTC0Y1O8nU6PfpwHrj+ZB/q+LVk4xtIzN+fxT4G1VCHKhzISBKl5tZm3AAAAKQGee2pEf4TUKw5pbspc1EHDtEesCCmJf5VnNe6KBZd4GQ+7A+bHfJOAAAAAgkGaYEmoQWyZTAjf+8Pt5W8cTDD7lNWFGUFASMmfmrkidqCky5SC9+cPiAP2S/BFQOm2nlMb+dDFgmvdAP61Ih4Y1f7BfhGuZMbC8UQVSneMohC5yiCzKNIygFEz4V6khVAoCjtwHr78chXAgv8fXLr95q5eD5Lu5Nzs/iIGtMQoXUAAAAA3QZ6eRRUsn3xlxLyHWJcFQyLMUspqSzsA73febdhh0qb7NHDVptBqRvp/0KYGMiwx7j7KbYTwgQAAABkBnr10RH9/iXZS+PSr/l8r00sSHVgOJeMGAAAALAGev2pEf4ERwdJQ0Crb1298v1xh8pKagceHwIXpnt+MpKlEEEQXCRaEef9JAAAAckGaoUmoQWyZTAjf++AchU6MHD1O3v5r6pfeZSR95Rrhxom7uPZHpFqjal5BKKKMTbN6mgvg8mD64Xft0kUDE6Cue5SSiFKoyWlVEwiX5hDItaXm0RmztEkwGguvWdr4BDwmXp/9l2g6tLg2xXeYbPfQmgAAAE5BmsJJ4QpSZTAjf/wp7xIUfno1XJ+iymWGbtpLZaw7PcYW6RhtN1hZk9b9PCNe5qkJxn7A1KYQZsy24bK3hBGA5yy6riTS/6Bggyp9yOsAAABlQZrkSeEOiZTBTRMb//vuVw4S51+pMKjT792mrZCg0a9enAYDaTD0+eYSezMml8HsNrr4W4Ypl8E3Ek6kEobjk/rXZ7COVAytpQu5ZP0ZvW5Pg7D5NsaUt6+L3onc6uiUnEZYYcMAAAA5AZ8DakR/im4U3dpdyjXpuEJWyVePMxrhvDme3b6G18sYPqqk/CL+3u1bRHyeuiRn1bp/7qnfDmkRAAAAXkGbCEnhDyZTAjf//CJ/zIEiRn1jC9hpPPLN5QuFZX1Los1jhijhPS8st288jfDlg8K2nAHAViYU4SieUCj1ME+It5LvfykfCPuvpb78LHOE+CJdbH/LeiWos4RGOk0AAAA8QZ8mRRE8n4DmoVg1Hkb/e4SXAQqfAAXFApCfAhG8pf0RtikgWpTEO9+lQ2C26jfa14IxlqzJJ16I5H6AAAAAKgGfRXREf4MMldf7IOPpaY2kRu7VY9Nr6GAEHU6Xd08+AnO38OezzPbYNgAAABoBn0dqRH9/nILJAU17xAvr1/4H4dNFOjsFwQAAAGFBm0lJqEFomUwI38qLkf22+XhQO6dWZKV+0UxaeuBPWkT9AVqu95T7qYZaV8ruFgtLIWIwIFx1O4IUmJVkQy1NeM9AOftya18H1uabejHiFxrqGQv/o7wQJwgNV5M+HntAAAAAe0GbaknhClJlMCN/++A44rmWNd9junV/0ZL8GPrsR5ejMuAeKb0C2X/4TBgkmpc5cbaAmhsRGzzPOE4Fd9ectAGJfFuKvehbLb2HVVq257BpM84fIVQhLM31NAPpq1x6jVFAQ8vyPmNq8qDOjXDoMYQIaRde1lJOu2UcnQAAAG1Bm45J4Q6JlMCN//vrPjhLnY6q+y84SEeloUKkmwZvTuvQy1vBx1cbwrmnV8fPOJEA/L74o4k+P0p/PqnHXGHiFX5kRdkSTKdwmiFccfDJFoOJFk9OTzWDgiBL+BN1Qp1Rsy9KQmgSswUoMNqBAAAAKUGfrEURPJ97HStiPWYzn6tT9FPiyYP9OyFj9PCcFZvsfLF4FBcBfJ5XAAAAHgGfy3REf3b7vlgIlKgE1B7NQWvwATKbaZAMXPLJFQAAAFQBn81qRH+E1CsMICfUaEKDBpIhdx6OnKAdCl2WDjL9pRb1pMM4S55WoAWOuHs6SIZMvTEbf6TD+l89/qm6bCPBnWGYDYrn8PXt76kldgiGMd/HEyAAAABxQZvSSahBaJlMCN/76z44IUjPWzut94XBkNuv+dieP5cTiu+U4ud4qkSdtDohXcxAqxjKAO9Sq7VU/iPClB4mMSKfAXByslZgH7DJTaF0HqcQoHsXQn4FOdmjHs65BRr5LHgnCj0HknjvGWtSlUvBUBwAAABBQZ/wRREsn3wG9Ho3WAS3Ni/Hx+OpCe/8+ukUSErgBDGT0W0PMnbTfONFyGyOrLRqgi9mASDomyObYuNuslVCT28AAAA0AZ4PdER/hRLKZLTcnylaqqMzlu66hI+mhxsbMLljcUK7FzrBO7JMAaWFrqACUm2hnz5DwAAAAC4BnhFqRH+BB5cMaSRFVafg6PSCA26I6z1H3nEcDifB9PFmvQH6EfZwAD8lDAP9AAAAfEGaFkmoQWyZTAjf++s+OBCWQr1tn75c/emk9LAHs1fTdTbpvJzb3LtU233UzJwGYEL7e8Ap8RSOlZKi1Ur6TnLiFRd10qK3EQMsUKDdzkCorRRoghT23l2ZqI/4juhRRTX/HniY+Qj/rMVC30KuoLxdRYifc+0btYjQlp8AAAA6QZ40RRUsn3iZR4RsZYniIDRxKFBpXUAJNxgE5p6+BezvDoyEfR/+vl3w9ZsZ9oIbtyEjlPrxFfc9iAAAADEBnlN0RH+DDmAasfgFq4gCUD1KaJGjFNBEJVYapagXlLZLSIzJIY2+GzlTGV7riDN1AAAAJgGeVWpEf3c5fYTCslq4y/9ptS3UG7hViDeebQUTeM3L6N0COirdAAAAYEGaV0moQWyZTAjf++A44EJZY2AK6h8ndtcC/MxKotOHWqm8LCc5nxMDTrcRe8XfRAD+sULXlYsOMbyenQJfp9B0zySWMkOE1fdeLQFXHwYudsCa1bFocAlUNzU2ll5X4AAAAFVBmnhJ4QpSZTAjf/vuV1RDrkiUKryvNy50VVcIKHyfc0rVEbcQaQOHTvicnp4trDKdvq4wTZVbX2wUqr/yMRShkhjqnvVzqia/DOdlZb2Rb0wer4BhAAAAW0GamUnhDomUwI3/++AcgS6wSqZJjrGWxPv45j0u5gayBIOWxzopmY84hnH7gzzrzeXgUYnG7EGJiE34rnk/Z0xY2sxq0mDVmu0uyDqZKWyTXKp6UClgGZ0ewkwAAAVfZYiEACf/vgy3mIDfV2I+c4OQeUCbjkIqauQ7q1hI51t1KK0Ty/hpqe1v3iAKauIQOhjuzFxyMgx7FbnM/2jBWLSkwzPM9uWS0H4jWSakogzfe9txEs4c+/QsLeTPvLzlrrhoRvk76vX6F462KFV4lQ8BLFX1JpZOU12xKsmCPjGaaOJgyXQiOhYn2hWzgcqRNjGX+V1keWAHLdIU5Kt6ThBKFN9iXo020ZIGb/b5NjIEtf5xxojLXamTjllPOcFI0jqWkWbZoIKHhmS8TjKSHCuTIx9dl0FTzgEkfjAOKV+MBkOb5p4/AbH1ylTOdgDP7a7Qz9MRQalIH6M4QHnoVD53mLhP9Mm+R7+Nv1EH6Th++wBzxhKbERneXaOozvtmrCru0n5aeRSyfurbH+PPpMF9pcN0wV6BXo60O4fisdnnNQeaTOQSwEQoA9l9G4isMAZz0MMRGPJfTWi9P7ynze7cFpAtxFXNele+QfVNTyKOBwDD7HFYOx0ouGd90huFZ2/V0BetUuzhUO0amnSSpDeMcYfW1ZAlHhEdZedtfna24MPXCAzeTxN0822JnETBRQEqo5eqoRZEG9h5BpWdwsZ15MpoPzXeKTqs8tVQC1/SK0d/0/xgwE6pNhIFwxgVY7DWsiX9Xv6tLaOudQGEMJp81OTR0L3+PIQF+o5vb/GCWDAkSsOBZs62OY1NlmA//1Wvizc5pquC5tnGGabusyA/nmG7PBFvvzvk1g6c9FZTfMWn30C0nuVgD/+SL3hOUxNRRIZsd15UgRRZCQNTNuPsNF70mEMar3F5ZGy89n5mtWtCMnx6Kxmzor+SNb9cdhFvoJ0pJL6GWtPAX3YGiIajfBHoC5Yk+AsjS+erTnaEBKxjzK9GpM93N1bO7ckHjy4hYhPKAFrXZD2uS3VtXe6oTBsitQMDgoUo2V1oa/61b7CoZXWhVK4UEX6j5mQQQ37LogaqDE0eR/hnOHk0qp1/ZLliBrAn0q0l0jYC3N9Q8Ae2vkxGonmaOmrZW/AToLHGMb0emiADcOSePLr7YHGJDiVzPIBiXNYDOLMcaDwCYGfbEsBrD2uy/d1Fv59oWvDv054FejnrCwPJNhHJardC3ukYi14OF61tBTEX4on2ntlVKj6TiEl2AY60ALC2Pewdc9kxhuM53QE/DPE04ecu5tBcDncARadW/1to7/E13lltojS2w5x9whdbhPouWVm71+s5063T2CXVUoDOAGHtwV5PBLlUV5s+9rM1C1nm6Ly8oNNr5tCZZ+xgH/cb0QtGLo2/+urcmwovBsXN81PlivdxeZpmL/iI63XLhkruWXb50jY1FpX53VswqMFVEWIGAsoxcV1hcpt5SjjgtAJB9brynD19/u2V/j6WCCVV1YJc0BsMrlVTR84gzjOn9fl/YKUcx8xd5zXtGYTTjr+P/TYuSs7g+gIPIJqW3t2VM/TsVgWKT+LKDZouSIl2nMmmD4XKzO9dTJuIBhm0f3xWeCHxstiNy8dyjXq6MgdCHeN+Zcag5SxJrW/q8LCIwYDIjdfeTEv9XV8VvX1oPi87I9rfBzE/XZtG1OijeZ7AcvsTu5Q/+VzP0izbIAexsR6SN7ywadj+XkAq9hjRaN/3QPlLeJxa7wBlxl3d1FQ8pFq96dcjnDir5ouXNKkIk/zCOM+q+YORr/PMY//AhAiSwED9KgwJz8yR+k3Go6EUpMgLgQtZ73cqHm51vCETs12mqn98YrWLHREslKlJiLjFIFQpWKlBdfkllUxz2P51TqpY4sD96jVYhYQ02ZSpJu6Tqh26CeFwNmHP1DrxMk5MJAyK6KFlKGZZi96tTwAAAE5BmiFsRv/77XFgf/gTkah1R5Us0Wnk8dDFNJtGuWmlTOExBPxdbP2u422B2Lto9nlgbofDURRsHjnidXKX9OQPQ2PYSuh0BmyDHyTamf8AAAAzQZpCPCGTKYRv++AcgTjLzdvWTXYs/nqVO/kWocgnSkHh3DwMwZdqf8C8gKNCxrEXojVQAAAAkUGaY0nhDyZTAjf/++uEdiIsFrkFNTDNvjCEHv2x8z+IOmFBPu6GwLYUtDIjwQ9PmO+bcrpj4clo9KQ/KSk3neQOZHILev1Y0jHKDWKlbBJSNC7Jskuppn7Pw9tMIqVOLcYm9Im1C+17BGV5dlf9xePKMg8Jr3Ff+/cKGd6zElUtjcoq/1384Z27KW1jNv2lQ+cAAABdQZqESeEPJlMCN//74DjhLnY7IVYsSZ54teV4M9pVBpl+jrm2vgkpxJMVoCRgAjdr9o/eR/Ca4bI5DFioJpFA7jz3z3p6Z8EP6+oWPcI57n/7rtuyon5e2HiX3e/BAAAAPkGapUnhDyZTAjf/++A44TvpQ6q1oyiGyCVbHK/f15aT3JMWQoDuJk0FKSLab/1polBzKBKCBe4A/5QO3LPgAAAAM0GaxknhDyZTAjf/++AchU+Vtkc2kb2jJl/oQwyJciND8Tq4ecQix/2xNi8PSXJvoKQxIwAAAFVBmudJ4Q8mUwI3//vrPjhLnY6qAnz8Ja7nh5BGRfd02joZMmnvwGPMnY1UtK+pb8bqcPrtmFQYzPi2qj09uQT0JTvmvYDd63U5Y2jF0v2YgFFNEVesAAAAVkGbCEnhDyZTAjf/++5XDiuJrqTwQnNXojWSFTharkNbc7gdyfMlGNhg3mVGFS0alrqG5sZCfIgbqaclT6Gfe0OipRwkPfdn4m7Thq/CqsNd+JyxTOnxAAAAcEGbLEnhDyZTAjf/++5XDhLnNsgTYfDsmrw/ZxClJj+Si9cQ6o5fvdg0roRb7xCts3+qftaz09ytTf/Fu2QxxQHWzmVyf92Io8RDVoRP2XUDjiYmvZjXggNq8HwiPdQMdvmL1lhb5fPZGj6hWNxUXuEAAAA+QZ9KRRE8n4DmoVg1HkW8/6wpqyGi45GD24TJ/za6DAQ4eU9CDfguO97+RZtQcexyfOmagRCfqhL9X81Swk8AAAAkAZ9pdER/iffqq8P35xMd1IYhD9d9sBpVn7iTWzwZs4q/5ESbAAAAIQGfa2pEf3+cgskSU0kAS0c+imyjJea5l8GNLVRLfBL7wAAAAI1Bm29JqEFomUwI3/vEhS51HjidTPHSoGoQXFVXRhNlrpyj41Xz7oGDI06DVKDH46CPr5apO2EoNwcdCUxtW3hgqmJ5/H1C2ANjl0mdlryKc9SwOs08mAQdgVi9wAa3am2ilj656liiY2Pr5+C/f/vDURY0hdLqqHzbufbgWXm0W7U5ed3s5e0CGqT0N4sAAAAsQZ+NRREs331UtYjmOPqgDvwpA+dkDRfmu4CZBFY77Ij1VrxJoJU6bh56MYEAAAAzAZ+uakR/gCQt04OPppwlRhivsWAxBpyC7wzzomqfCVfnmfzfbSoDKVZgsjBm+Kt2f+3YAAAAjUGbs0moQWyZTAjfyghBljfMMuyWTw27Mmxw1DJv9ycLWX+QHLgBj6IZtikvzi8lkYN1RnpMqOZIVf/nduC4IEoAVSEkctakaQiAU2DoUeITXzdRvYi6cLHFhXwRHpVFF9boSkQf0W8GSZnPKqex9y18cqZ0/3qjnXDpyObZscNjiTUYfl8+sVcWkDfb4AAAAFlBn9FFFSyfgOahSerkFDSgMmtywjBbMGKUIhmNiwJ7hm3fI/sR/In45K/o/zgHqjSeTGmpbhIijT92cAprRPD+g2EjPnGa2bkpYhmAbzhjtjjTRAhkkbC2YQAAADgBn/B0RH+Eo2cLrfP7Es5hEIvxy3++81JIcbjtXnYbae6sJ93XFc4LmG3iE5toWPTIgZ3Rf4u8wQAAAFEBn/JqRH+J9+qqfNhEqND9NFo1QljVM1SBAde2cQSHX/EgRDnk8wp3LTUqgB6d+8T/ceO3zoPFZUv5ty3ZB/Jay3Kd32tWe5euY78MFmYrF8EAAACzQZv3SahBbJlMCN/KCEL9LNaI8OpfBjrh+prQ5Hx5F8zmd++rK6rSSpIp/Bmp0IxjryV0UhIqqY7ythB00KcAsr8fcfrBfpnJuvV+Xgc3l0ufktx7e7w3rU7jqfUydXix+/gToV3m7PqBlYaP7wIonbR+CPd7LoCHcFtKXmQsnVlI4yUAum5TYseBCN+1u3fAF9yw2dGEPlv/60B1TfKQhsyYxA36fR0Gme6stPBqWq3qMsAAAABCQZ4VRRUsn31uHUcUqyrx5k84MB/5PPciwmpzACm2Qtdwi7HXVZEde7qqTsctIRWclAzfvXhjFp+5nog7nC999IrgAAAARAGeNHREf4ByGO+by1cZ8ELMy2ttDhtzqmeTvmfBobP77YCuVEO65NLXFqkgoV2VMcV/Dme8id5gHSZgEfW0Usi+767BAAAAMgGeNmpEf4SH6CyS5Kx3XhSoUhAXrSLodSSH+wraAIhCFkqWqwzMs1mFxDnkzNVjPfe3AAAAdkGaOkmoQWyZTAjf+8SGbo4p/8Iamyi86brFs7iSG9hkRmxUtyMqDX87BjiDhTPFbYOFwbUDOtNIYN5sBPm9LHJf40gfxrHO5b+yYb+iAzQ+D0jqzD0wUE/GPk/is1vbvFMqnhMOuYRpHxFQPFZ7rMcJvsV/ceQAAABJQZ5YRRUs331NpBIZUVg/vHpzTqi1adLx96xSVQdbyZNxZ1wEI8f41eXgBXXen0pb3PYiv7Ia00pmzCGAnR7OdtwQ61Xzb3asQQAAADIBnnlqRH93OX3cYVkh8eOukCbZ+6JFa+s4Ajf1rEP6NxJzT4vwzyoFKEMzdAXFIlOjEAAAAJhBmn5JqEFsmUwI3/vrPjitsxqnuuZB/Q3LK9EB8XZHkxOb9shEL9bJjdZ+AyCA1XlQlALxJUJQknuX92nj+nF0tYg9+8IRFjYGwbn3xaqO5g3aC8XnhoUJD7JtUL/BWe7piCPNYHDhlaejMHPjLpBeweLNivYXRppq6jwaS+E6aqBmMXuESQatm6wvSHErLtCeVTRchSPIcQAAAGlBnpxFFSyfeKBLx+yIGhVJaNqiSbf3n/EhW+vNSV7ifu7le6rz/EoChwQp7PkLC2To1qCPmWejORWShLyp///Y0qa2M3jqj7c+mZFtb+jzH2ON2C5/c8hSTBB7jEfQYULfNJnrimd91mEAAABPAZ67dER/im4U4K1F5C5u5/eTKivk5Nnz0mjnt5+KqDGfDi3wS+Mc05ibi9UxpuE1o1aqd7CTGoYA5lcK6Hr2V9Y1yJwB7+JdsC2pp0re2AAAADcBnr1qRH+E1CpSs1MPhRiIa2mtfgUCd6KO7sJt69FF+hJ5oUlgGPC0EPB1ogSuuBERwapo5fuTAAAAhUGaoEmoQWyZTBRMb/vuVw4emP9PN/4ewbFLUGiXiZoXEL5lOAJuqN63bq6Ya73SLuB6bBVPnpl/y39gFl9DALFqt3PLXqzNMiQ8uGXt/+TM86YgjzWBw4ZWnozBz4y6QXsHg/Q00M24YeHSO++uLeHaT1KOziWaUuzPIKggxZQgUTyNyYAAAAA3AZ7fakR/f5yCyRJTQtCEVMtRioHnim37d2/B1xh3Ql9UxqfT9HDmdIOoZcKX+BDfTVGuDUOJ2QAAAINBmsRJ4QpSZTAjf/vrPjitsxnTnKYcmEudJHuBU34rmoZfuXG6Cx5BfKY65GVp4NsLJ81tu0fHLYXCqgIs3Opp1W8uDzLiioBJ/PdQ+EnY54EcqZA6tP3J34zBz0T0ggJpXUzfiV7xPL0uRn40C/UnuyAQslAT5p+tnCRVSaGrqcLiOQAAAE1BnuJFNEyfeJmV5SO/7e3+IXPu9DkhOBWthtp0R6laPa9QVTtdYLuaAVCQfvIfjc6J1xsnVDkt6SMF8jhI717hP9wd4YjWpEepJyiGOAAAADkBnwF0RH+DDmAYl2uPzEbmrBU7+iRoxY56o643yA9iq0lyU4HIBaatgNSaUrc7+QCg2zU0nrhV7MAAAAA3AZ8DakR/dz4yZYCJktXmd/7L+xSWqxVX0hhWCgTdFXq4mp3c/1a9hxSgDlfk8LNj5Vw2oTK18QAAAU9BmwhJqEFomUwI38erj+EOVPHy32Fbp2Yz5YGdj+kHyBt317o+i0bT5HgD8WT9xKPFw0LmVzgeNi1FS2ii0XX5T3ObkUYS35DbQo/q8fsZOm3CpwwU/Yr0/kpevNm1zY3V4E6rdm9gjhcJayG8aqoO4g4WVxhx/xMNf2JeC2NfVjzZe2vcxSFQ0+nBoYdXCz2HGhMEqn1Vj7giKBRPbFIQfXzgNot7kOQRpn9o7GlvIh2PHH8O5b9WkObUKMB0E31SJnRODa7/x1gMeVNp9x5QG1UQTgxsISHm0WiGHaHM9FndRzybD9Ti4Axz30/xzLbVD2b0F3de51FeapbwWFhm9ug8jx3oU8tVKEI+qRwwZklPZNSM5eB1TunBJCGBjo47s/DWaVQJ4etgtBC0e4W/geamug1fhr/BERy0ommj6LhH/UwJu76416mGlYt3gAAAAGRBnyZFESyfzUpvZmoxPUsjVw6GJkSBGLuJUn+Aaj9PjxAHugzMhsCOH1qsNMVFzZkiMTpbdmVc6H9d/LXbXoSjfpgIFnVOI0u4Z1YosRun2lvX2y3QNp1C04NrWmN2xpZYm7uBAAAAeAGfRXREf83Ru2Ye0HiKfnAIZilzHLjlqsw+Mz1QC7FumLzaTF+pKfa6AzSFSUgrNZ85pfDhXgDDUnT4YfcFFRuWdSKxg2T3+RWpUKLLxFjqhPOFsq9XQnstbGEYxiunguGhZ3eFb8BcZaO/8shwvJPEZqscqq/swAAAACcBn0dqRH+KbhUp33TVx2AKpAjXBFaDxlFOYkOj2jzXxGLdVqUiJCEAAAB+QZtMSahBbJlMCN/7w1h3aeLJhoTP+9zKomPaDpIjqpuU+pG82MhTY8pgjC1+/SH9odkJKpRkMPchDcbwgE0g/0Pt03UBu2cxVOCZ7XVviNWlJE6vWTCbySBPKWxVm7F5Zgym+lfh5YEmZKBo/8QqBfrF6LcDLo0R9Hd/6IWhAAAALkGfakUVLJ98ZVVIdYm9ubdI+6CZ/d+Ekcs8ct6CLr4YDU3lQFcmRYk+y35Vg6EAAAAYAZ+JdER/f4l2UkB1Zty9NLEh1ZInMX+BAAAALgGfi2pEf4ERwdI/XqLu8/7QUPdBBLv6noXG1Jbi0GH8S3h3poGiBbJjvmVilJAAAADUQZuQSahBbJlMCN/HqXjaYdAgn7QtItVYNomimA2qDY9ZCdac+T/CWtL/WPI+ea3ZhxpIBZzHmoE7ED6cmscuOpC4asmWP/ZY4D3ZofkgXRT0JIHQ5C+jN/NawCAZBlYJYeWSrnedRhv1Aq+Ezb8GwGkCKqT77Ol/D+JE8MiWDBfvfPqV5xJ6yu3+ZMsfZri2LtSEJZnT2C1JziQTBkpSnKjEt1PGYe1Ql8Dbmxq+Ajqsnbj/1Q9G0VHb1uODbJw883wCc+9y+57YbL48lccwqZcjBHAAAABhQZ+uRRUsn9O5+lrykd1W9x10yLwnbEdC7nPGYZCdulN4uKAuREGBPh4WH7u461q3qEJSnberdlJt1MwVHaJFwOjSN0QQTPgMWg4BaWK592DS+2DuO5Vpp9sLbJu4zonEwQAAAB4Bn810RH92+74LkTJauMv/KEBnzMgyQCjXgyIf5rEAAAA3AZ/PakR/19HKwpYZLtl81TUXlOY+KofzOQxBT8cp9helRzplXim5RNQlUaFCsHjzgwNANcnGkAAAAHBBm9RJqEFsmUwI3/vrPjghSM+3zdXqrzb6QyDQBh8wNH96zteG2dJP41q9qpbbpx6VfLx6pZ+hD0bxu4vInts9gfPMBlJkiYledG68AwhZomhLOtFb6Ji8an7vaJO4LtZXK+JkUFfz95QFkyWqJ2vBAAAASkGf8kUVLJ+ASbe6rPuX5kKJVKyfzTytHiClBBvNDUS9gyZ0gJuIACm6J4/MaPYlZAhSExx1G+JA3YMjZrZGz+SE8khQ7v8Q6ZsdAAAAJAGeEXREf4MMldf7IONNzwuYXTWVXiXI1jjTphyXCSJiDcl/JQAAAB4BnhNqRH9/nILJEpexGmpP0YxoFXxtFiqwczbTViAAAACUQZoVSahBbJlMCN/7xIFmvqKhxfQ7eu13hwX7Kl/0+GfUgZl6JYRtDtsk5rFajLUh7O9Y3Q/862jv9ZqodZkpyU2BPt5pwB+6u6zvaK5mzw79FkINXTazxAwa+hbdsXMsxsTU8kbKcxOzDPlpJOZmg+4BqbKzw0yHVYCbPCfiUsYQ1v/GbNQjREVqvxvErLVTetYAIAAAAHdBmjZJ4QpSZTAjf/vgOOJ6G+EORJB/i201ADDOh8gdy7TnqGzfuGC7fVt4IG74ck4fwuW5dPJSAv+ZSwOLk765EsOlxje0d6xKBDadZI86kJlJRifx+c16vpsob8hqrISDUct32W37mGuw8GN/sXV3vuo03HQdfQAAAEJBmlhJ4Q6JlMFNExv/++BEqKNR9RoJeFo2wJjCnq8GbSskuCaZk5J9oPynQ0n9SgOlT6cKPNsF1jNv7rf+cZkgbygAAAAlAZ53akR/dzl9hMKyWrjL+phdX+F0v8Q0kOKCEPls9CPUogjSggAAAF9BmnxJ4Q8mUwI3//vrPjhLnY7GuE3FgwVQBmmbcyRjHJitYGt3XJ7ttu2Rs1JgNU2hsMGCqQuf9g78vISwHq8b2HhwRAZXrups5jQ3v4TD8YIj3H3abrar6pujeVeAuQAAAFVBnppFETyffAdbc2ZCKuLu6g8BJqH4YE5YGtjEE6lGSZYtJ5hHsbl5fBMQmN9SaUIa5J0zft9yHmPwiRAJF0kKNTHMtbuNxCSscMTsy8L22gcicp5GAAAAPQGeuXREf4puFOCvsKrmMxdXS9O7u9OTUmOHpv3PBWQVCIk/VLpaFCq8y10BD+dshRrJ+UgPns8tyn3sUsEAAAArAZ67akR/inDB6JvALVyfkNmDHBAGn8T8vQbjDxek2OXAhfP9YxVeZBClgQAAAExBmr1JqEFomUwI3/vuVw4S52KAs2D87HYORWxCgFxVZ0fUYtB1vGD8zSig7GaR6PYdlvzo5m23PCr+yMcmEnzswwPvpxVY8irlJGnwAAAAQ0Ga3knhClJlMCN/++AchUq0qS3q2sujrVLxU7zB5z9GGi4Vr204wSsd+TAv/ZDLXIheOE/uGZxoJRecWc8xTdfXR0EAAABsQZr/SeEOiZTAjf/7w767baOIWn1VtSMPFg3+VI2un7hfAiMrNCPxiyraA7+7P/b4mHhjf9frw5Yf09s97+lHoLr+J40YfcfgrY6q0fAZ7vuspNsTPyaoTWu8OCda8dlFOKXttYcmZn0MYXuAAAAAdkGbAknhDyZTAjf/yghC9V+KM7GjZ3ifEeQQyURdNhHoTXnSTZGK6/FJfpI6X0y4F2LWLeuzT+DVqTb9zTTxzOILEvOnJySq5VGcY932fKbq7VpGGOVsYCTlw9TiZnaKVEnbuEjCDdLoHCPxhxVcfUPLCUueKwsAAAArQZ8gRRE8331Nr6PZBasOYw44aaDFj80ekC22ktCh57XyldfE7Ft3UzEJXwAAAC0Bn0FqRH93OX3cYVSoGPHXSBix14TW1QEOn7Ij37SbGA7W4x774czRQNnrfMQAAABuQZtGSahBaJlMCN/76z9RJhOzd1KAOnRwvW2BRZ7n6ANrAh48/h2t3g+1nVLUJRKqkD0Z9HeHn1likKz+6VR2gWMleErAnhTwmDW0YHrepMQTwNh+0/Fx5VBvqdRhjKSWumYEJreA7iMFKrvjlcAAAABbQZ9kRREsn3iZVH1hPWCBP0IcLolfj3IK062Db1vtBmUVEFxqhE2rclqJUf/kvOICwRgf+8cuGE3RT96eNk/u+JabuLcL9MSQv341jaGaaPUYvsPjkAuCYD6NMQAAAEUBn4N0RH+KbhTdyyowLli1rdXDi2OJpb2YMOsbPWx2VrALoxL353Cur+1DqAI4VUlPg/2b8RvB3K+6VuFp2qc1xr0ccmYAAABDAZ+FakR/hNfALxWamHwtC2+RNj8gRL9mIXf4wiYmsFGPVFdbv/eg0xWIw8ntIxBwu+sEvpuIcmRgof8iOS3heddgLwAAAJVBm4lJqEFsmUwI3/vuV1RItArTKe4aHtkbhVf/+Q+mhmXrHW48v4hipTzH46Ij37Wg9Db9jYlSjcF5jNQ5v1aWn65DMLo5D8XKRv2qTPmOytr/R4nJgxMjqpDiuGEXP91sGSFXXvDL5KhIcPu1Sv5rfiRKmk75SupVGhPOKnffZOhfkY7nU8WILOh3WPGj/TL7m/sJmQAAADJBn6dFFSzfgR7dVmJmGX8GHkqNZi3rEUsY+LD/L+bIktrWSlLH0rvcGVpVDuPV/0CpgQAAABoBn8hqRH+EUzCvSGusQvuCjCEcHYR+2+CswQAAAF5Bm8pJqEFsmUwI38oIQvVfijOY7Jqv/ZothAXNYBl7PM7Lu8cHbpXbUuZ1KACFF9u9XqIPlaA9xYZISQn7YtoCcDfTxu4CUZabiiBbVAJ69uAqIzwhmC7k2sQfHOcxAAAATUGb60nhClJlMCN/++AchU+aExJFTL9PPvLMNdHTufZ2ARRjtizDUCD7M3x68CS3Qr9dPeG8VhNRm6WRPsaNFzt/4aU270d35/LBN6vBAAAAU0GaDEnhDomUwI3/++AcgmDYEE7VgijEQg9aSxHaV7+OwetiVVzv9x0Bh9mAJqRJH8Dife3qBjyBpZMx0mPHKfKwSXf1DK2eIEOyj51/Ibe43gjwAAAAcEGaLUnhDyZTAjf/++s+OEudjsH36ASmbVNFHT+He3+dsKUY4kF1Fcb94L//2T/GGBOr9knUA4Fv6E+UdIJFqatHfWfL8/R/5FZzFcRTOFGht9DgyahVNF6gH94cGhO4EjBjUyOq5QvusqmUVB+YjvAAAABZQZpOSeEPJlMCN//77ldUSEVhRtuTqSuwn4ILAX/pGTDP9IL9t9LXL4gNSLMxC/6/2jK9Kp6HiBHIDwudiAEToisNqkYpE0nFGGBP2Mwq4VRrWiM0ZUYi+BcAAABRQZpvSeEPJlMCN//74ByE6AJUf3GJyJEuF8A3Qf1kQcJRR6jInzvSF1cy5rawxyX0nBOolFDiu817lzRYe7/SVVdiebEP++yydeb7tdJ3zxI4AAAAT0GakknhDyZTAjf/++BEqEiYmRvIgNxJIRznuGsny2nrjSbaGRQNjLW/f2j0oaCQe5/+54Ig6O1mLlpOJIun3/37i+STl9TMuA2Pp+8ndHAAAAA6QZ6wRRE834R1pCbTRWlD+k4+9DOjI2AHXza1raT35vBLoKirdVqNYgQnbX4DTT+ozWwyrExyRLfzhQAAABwBntFqRH9/nILJAq0W6TichzvXKAreK4xRz2yBAAAAXUGa00moQWiZTAjf/AOiKyEjuJtL+UDu/7/1ImEqkms4m590rIJHMH/lA7MFvspdRcAY8czYbgMQ6vpXSSw24CC9TXLj4CpfdwB8SVMwSgaDmTxDdOzD05PcebFnHwAAAHxBmvRJ4QpSZTAjf/vgOOK5ljYRuKx062IYgGdxv+oJsIKJXe5wzU05EYgy6CDVnrvgrKBM4IITLfdtFi9dSb58ZA2BYSlRsjpR6nFhFesKS1PpyqD9tlHbqy54P0NR0gcxVnq+CkYWBsgu1ogKQ0vpvuycY3R6AjqPjNvYAAAAR0GbFUnhDomUwI3/++BEqWWdyZKFzbFLbFx3xidN7yP1nsbxqbqZf3+mAe1LLS7fO/2lqY5lAlA3reqb9jagV1TR1CzRq2ngAAAAMEGbNknhDyZTAjf/++Aci1WcPt8e1pWqe9fATAcvlum6kTqhjA5k6uHd3lKWBjdgiQAAAGNBm1pJ4Q8mUwI3//vgOOEudjpLQF0T5/HpcMVhuo3UrgFNwvqy0NaDRsYBORhen300dcdaxgQAhyHLkqqY21QvhsGVyf926QHjO1dTQI9IklTM6dFWv11wpl628oyD4WUbTWgAAAA+QZ94RRE8n3igS8fsiwsG3BGH71JSuKL8DGQXeoxKKRzUOLM36KStM+eXk1e/WYIgappa2TFwUOXYGRkunvAAAAA0AZ+XdER/im4U4Klw8nZaidOEir35twZNCptJt12OpowRIVZ2i141MkhqsNTXZF3ZZSs7KwAAACwBn5lqRH+E1CrjX482235Jp4YGEincKNRQy+bYAzUsHzuPxfifmdaLZN0gwAAAAIFBm51JqEFomUwI3/vrPjiuZCYQyR8N4L4ZpcVyVm5GOzBlZrs3WY0M36PIs4HlCuJLEY14KxyS8bo4xIH/xQxifCHSUw0kKwQds/H/qM/Bfv/37nK8BNWssBs2UpmS7HvInr7RQkATNzNBxyIYrpPs2K5w38HW0rEzP+G0ztOo83cAAAAqQZ+7RREs34TtSIINqX8yEwwllxjEWv5dp22DsRrlE9pSWyEg90j3GhvBAAAAHgGf3GpEf4Sa2x3nqQ9dXv2XRummgfaBLvV0z7jPcAAAAHRBm8BJqEFsmUwI3/vgOOK5ljYt21R9BjygQzFkvYgmlW7E/rzUmB3aa1R/v/sRkWdXCOgIKuKUfzXiJlEyDC3eROqOf9d/4jH4mwkcdJ8kwHn1NDcqEmz68Y5bf46nu99czcvNzhvDbgh0CK9lJoeiHSXoTwAAAFJBn/5FFSzffU3evezXtC9/ZQKRQHpR1ZwveCVG8cv+8KloL2JUK/o4OcsKclrttQWDgV9mNWZ5P+28tt6XuycTD5aBLcSOUtgKm4G46iDPAKbMAAAAMAGeH2pEf3c5fdxhWS1eZ3/p8a5pR6sqQeE6lWFP5z0K+aYlfBg11Rzf+F08p1cQ5QAAAJdBmgRJqEFsmUwI3/vrPjhLnY6O54bcEV++H1QyWusE03qBnqIsvXHnWussWXieHNrrthosgcAU0Y0F9hopOWEBv0JGXS1m3o8PuW8QsFb6/rPNSkXUkO/oFNfmFTp6wgcdbxk7t64LF//YEwizj60QETaAs6xDcg6lk3U1uH0HDMl4ih5VDERKo2WeNr7DJ6PAseHZfNpBAAAAZ0GeIkUVLJ+A5qo/ZEObD72LxnQYIuEYf1ytAEWZ9RqVWhpP9byyoVJnVZyVLVwDIFybDJ9gG7viY4976Z8PP75RopC/RDB3dZ0wq5zC/a7Qp87846DKhZFGiIsYsKK9l8waIPbfwkgAAABDAZ5BdER/gwyVRD3qoTg9Jl28dVUyLW9E4j2mbSLrdlU2wMBPhJBWiaLiTejDzYOqwm38CZG2crdvRPAPjeGIWaTVVgAAADoBnkNqRH+E1CsOUDlTRB+2knMt/XIVd1w2Lu5z2obTpdRmkuyfrXpGQRJ2yHOSmMjOiyrLDlpwJ3dZAAAAvUGaSEmoQWyZTAjf+8PysOw7IqechdCyDapN16veOkycd3tl/lOY643J4iTWQW0XQRO1oHAqDTs60fM41KeAMdMg5qceLz0xHG9o7K7+QDwexR/ef2KxGk+JPpEkkQPooks4C/flF4jhBo+xQ1NvAor/CROg6sAGmqtYsKTo2RP869Xll09Z9g0/aqKgXet/XAH1/pPh9Idp9o5m+LTjIUlHJDcUkw+f1s1KFZPXPnjINmfl4g7HW8mSl7TCcAAAADpBnmZFFSyffGVVR+zevXVKzm2vgkNswDKTj5fZQrLaGvFrUNGZo7oe3bojV6WjyvmGIqww/DM1Nn65AAAAKAGehXREf3+JdlJAdWrJHHXgk96zCjVyoQkJf+XToMVT4UbfSa64B0gAAABCAZ6HakR/iLG7Hu2qMv5WSPgOGOgEf49LPHtgU6+/FKGUK0Uj2RSz8n/2jS/nuH+a5JnCsk379xSbYaLKWkzHqqhJAAAAaUGaikmoQWyZTBRMb/vgHI2b7Fd6Ws9h26Qd1WAcJFbDOuHQlDOeOVSDtthUTx2BmQWjIawleP3YJiZxUceVqUKAtVJO3MmSt4itsmeOXkS96GA/02xr8e3XWkoQ45QprFdd1gNT1R5NgQAAACEBnqlqRH93OX3cYVSoGPHXSBIm8PYCgiS/lOkz+kEEcdUAAAB/QZquSeEKUmUwI3/74VOAVE8TO2duqVlti+FmeZePWOPchj2ByBQCHzxIwA7oRji+KqKa5n4Q6OpasFlu8pZOpT7HW7iwn891D4Ifhvw+ecPpsngXgO/F02GWm9i2sspJD5g5fSDd6gF69m+tkP3mOA8J0Iqx5r9Ybt3wCsRgDQAAAFVBnsxFNEyfgOol4/dcuaO+lc5oqdzUsBCEsAWXCEd9LQnFt83M2f4mvOGuJerHkxVQwxAPUafNM6hjcUDccUEXrpfwnpDY6bBBuWi3an7M7S8Zz5MyAAAAPAGe63REf4MOYBhgWGIje2c5UIBG+OMYNcDzO/HP5o+6xdUgK/xp8ClpOu7ZqMK+GL9wf1yEuViw3Jz5MAAAAC4Bnu1qRH+E1CqJbRfRlHBZLSjlMfrSL3+KsaKQb79wkzKWLcwNd6o48CwLbJZPAAAAuUGa8kmoQWiZTAjf/AOyX9R0mJ9SwB9iksEB7ZH/nkIpajkfaK17kN8kNdFdkevjT0goBYSx1ad7E1vuUhe0fKOO67ud4jxPyVzw9UyNsylkIBxchX2bMrVXnfEsnPA9mDRN8mbHacVFgb2ToC0gaK1kxLRZDbbQ7G9HcVMynbMoma6dSOovdmZyIzwlpap73j+GNAutYydKa0z8rYc9VMPkEA+pbgbAGRXIRlXA4YjOfmY7qK4sB46wAAAAR0GfEEURLJ99nGGl1msrnfy1MSO2So7YfJXmhMb4mHw7D/nEENpn/g8dYIuHWptpGty383qQsQOmaDb1Sr7SyZujuEtrU93EAAAARwGfL3REf4BsK1ffX11Zae3WZvfEuHzoFi2UXR12VR2woUFdnWnJRyNJzEMMur5pB08F12TS5Kqh0g6hXx3s+h9IciwahuVBAAAAQwGfMWpEf4ERwdJRbH64kx+NQysevRT1D2SC34ECzZV45BjVHUr0rPEmr61jCfTNQyQZxfD2v72JxfWvi2EhWwvZnoEAAABvQZs2SahBbJlMCN/KCEL1X5n/pbtseoLLd2WSFVdinV73wKbPDDMQGOJ7HQlNf01lmQlQxpN4Xld3ENl7/XvF+385IN+ANQzbxa9MY8ghxFpqW0UYwQhQSaaY/9FtnQOc0V/gREgSgiRtrKFqyex9AAAAN0GfVEUVLJ97HUrykd1VJNIYMgfZauRWHu7JgJIJWyjt3yOsBMa11eyBZgykRdmoIPb4fmsTH6AAAAAxAZ9zdER/dvu+WAiUqATXXSBiTXTh69DK9T96yYRRTZn7NwidgRAUaKpqXh1LkSJxnAAAAEYBn3VqRH+KbhTgrp8WYLnb/OExBB1QfQOTFVBGW40ihRipF6j5TyV0K31uIo2+OzX6N9DLAIAws+5y7pu8ISS/6hJ47suBAAAAk0GbekmoQWyZTAjf++s+OK5ljXoFiiU1XPjCnQxFyu7PryXdiZtWx/wW4N5zId9wNpopmCe9CPMLQBcd5pQEY3kNzWS2+3CovTQy4TUQSeEemdF5upIcffV6vhYfM7AT24IWEWArCyFFh7pIe8+zc8VPXEM6rc+7xdmneJ3szGOTGBGxzvpFxlTTTGKv3pJIbNGDAwAAAG9Bn5hFFSyfgOahWDUeRVYeX65PXWlnysfqmajig5ypCQFzzoXCSRAp3zgjG2NG7gxdp73zP4B6ldlcHso4oxhpX24XpT/qfkQR7KKVU4NNmqjCtXvpDalQ/00eIGvukzaHkWUGHTKp/QtbKsPGwDAAAABNAZ+3dER/gwyV1/tYOwEqTd+Pf0jD+/NLCVKy6ez+vl1KUGnCmGBaIywcEmAVF6zlPqshdQhK5FM6KDE97ysl/l1xLpPEqrALZu0g6REAAAAtAZ+5akR/f5yCyQKtFiypduyg6b/z69WowFlASrlNd/l1Q8J8ACA56sgTjZmAAAAAuUGbvkmoQWyZTAjf+/4BIhQKAXMPK6lD76MX584FxB3ojNVi5Q+urU2Vg7DDnw5TYhPAcAP+JpuOWYUZvbWvqxIf4+zhikw5W4uf+nNx3QUlfGE+qQdeKFt6AVvU+gOqNlNIKHEl5X/vEOA8B/wsTOAL7//f/G5Ez96Lbtif/QxRwf+GV/m0IscN9zKokP5tFqEx9jubCAyMXQbzIGza/AdoAcgWAzpGkkEFPVZHIn262BXUMjV6hgYRAAAAZUGf3EUVLJ94mZXaX35boOk57/Jas+YqlIwzjlszmhvu1CGdNLBRaqzGIA8eJ3iguthSW3p5XKu2MOj5ceJGlYERu0+psmSaLOFXONszzF1sh6Wy11Cr5kB2FRpfgC7N0D6P4v4lAAAAOgGf+3REf4MOYBiXa4/L66IgzTpMUW3pLiTwVe2YiMlu1Od+TX4Ego0CvBosfBkHjMBfsmQviRFuwkgAAAA1AZ/9akR/dzl93GFUqBjx10gYYvCa+gIMfz7WTXtZIsyH48Q3lDcTKBK+Eqf3cTexpuqr75EAAACKQZviSahBbJlMCN/76z44rbBJLPdsY55JpbvEurvbYmuIRcWV/rJn2ygwgU9CgdRF/043si4Zr5zKodFD1SgFp7DQqYD6tTX8dwBVP474YclgRLTJThjJ7uLHj6V/CJc4In7m8OgftrGMn+Zru0NctFHKqM1bDUTtaFM/qXypCDHtWSqgegcA8LCAAAAARkGeAEUVLJ+A6iXj9eUXdhdZVK7LV1HMdFPLHxlyPlUiuab+bMNEm5whCn7SWav9ox8Nvdiza9s3scrTm/Si0xxwSXwP4ikAAABFAZ4/dER/im4U4K6eWIXMANXzPh52kzBolbrDVKhXJMRmKIXiXvz9VXDYi5ff6Qo0gHFiqU6rW9yf7KoXUhBe4Rxlo8WvAAAAOQGeIWpEf4TW9RzxeaEIBMuIawIw7r73W/EwbmR42EU8aJ8SmWMIch19eiKXWJLO+HqvXpyEDPqD4AAAAI9BmiZJqEFsmUwI3/wj5K7y3raNV+lN1e0Gmx8L5RHw0IgYfLeHCZ5xAiPYSqSZkRczptAL0OS/Eadl7UGYnx9tFSTMq/DW1+S0JjpUGH2BBi2JBcBp9gBaH1YFApjEEAfskKGqdW6ZtsgBB02//5EzIeLDw0zMvQr1qNKsH9QHaWrVtuaIvnOOjm3+a/MpYgAAAEFBnkRFFSyffGXEvIdZdMwjUHspZ5lmIjrKqSiPV9+51l3FYjTQ30SJlU6anVjo4RznACIj/57d3zRCAGeOX36D+QAAADEBnmN0RH9/iZVJcpUEaVJ+D1Qn1J1MON4UlEEFw/V4SMOqbANXYVK9+DIT1C0sj9H8AAAAQgGeZWpEf4ERwdK7XH6zbP+0KsgLE2aJGY+DvwxwgrIkeYD20yzoRSouFKQ3gK9JnSL/2w7s61b8noAFMeJsnjdrGQAAAGhBmmpJqEFsmUwI3/vgPUUrmLiGAWYzwddaBPlC1tTCKPJyTpF5x5Urt5EikjDg/Y2Ch+L/G5axssucn+L/XqrCvrKKAbuvr0oM1XhQMnNPkrUx8G7rkqNLrOkC/mXdj9DtDsSuu7cqxAAAAEJBnohFFSyfex1K7S/556pvzaNwUuIyDg22DfvrCUNWtNpRRmvgNB5PJxjDd7uUat/w3BVRmbf9xTYsRqKfuotnayEAAAAyAZ6ndER/dv62TuMKpU1pPPlXna5sbaR+8FULtAJqz/eyyWfjyJNab+zrFLDXgmE1X7sAAABCAZ6pakR/im4U1zmIP3cv/+aqmtN99AtaR/ePPWvf3Kp9eqRfWQYoVZuV8xVd50yhmTsQSs9JdwKpAboLwg3uGj6dAAABHUGarkmoQWyZTAjfzkcmRUPipKf5HPWzd5mVSXOkf2wveByGyczca4WLzC6dJy1m5B6Il/iv02CYxAuC0JMLjf4eVD09cAC4VuqW1rQdH0jYBt9fqu8NkdwlacAwc4MltNuRgUY0PZqFq9cUnMVBRelxNywmlynkI02yNJIx1oubSc9FYoxYy1GOdo4XWxNYblENWmaGnuCBIJRx19vtPp+h4vvQVBXbfsx3OZWx9A4nT7Y56PG6CxBp9NdcYCm6+YEXmiN995AfuUddw5pT8IxHOckLKvtrExmNawZ4j/qtgJQASVXghncWygED1xlvOUMsv400IbfHtGjY+xlxEauu3NbsBWE2bmc7OC5AYARFLBZS1IMRB42qDSEAMQAAAFVBnsxFFSyfzWO6eV41HkVXWetrTujpSLSxsyet/K+xuI2mJCQpfU8AiK58PkSxjDCtX3agsIW3wcJvwWe2BOb+rXnGIMtnblWi8WfxWyNGdCgOJmIQAAAAYAGe63REf83T//106i/j9g6B7W/+9pd27v4NIhMKDiG077SiuGlpjPOE7m0C1ERXhoTXOKeyMZV/Cqhwx0+nXZkTcQvFvUXkgMZhV3TQViM0RySsmTCOTLoKGrX8MDQLgAAAACMBnu1qRH9/nILJGk/YUBDweJZSU0TGSR6vSvqHjaBT1vIf4QAAAN9BmvJJqEFsmUwI38erBKuz0hHXKlCdZjiQgCATySoqsm6ltfbP7rAFGFZ822Yw4uKROyUID51JSZ2TSrmysf59dCK0wX5bIkLdjmpkPTrcM+4OA3dv70JvEV+piUKzWy4rgthcTqUOTCLotioch4OuJRHZVarI23FHHeuIXdoOx8okOazHENHgaw6X59OheNrJiEvne4/W4SXKVQV9oGUnCUDqrTQ4hxMsBqxs+I3H3zkR1F3Wqbxyocy4mEvJ67y6N0vizCEqYxGvHdXCKAGvbYr0oZ5UmemNFdqa/+jOAAAAO0GfEEUVLJ/NSlgemZ2/i+2HE/tCoU5+SvCMe9vvnjLZqsDRHgVJCIy7oHSX0WQfEACM+t3MwE6ru1+AAAAAKAGfL3REf9fXXw1RMHznnW8D6KXJTU0SNFkICnZ+oJHgBulpiD/kF6UAAAAcAZ8xakR/dzl93GFUfKEWvCRfKO/B6X2IPGTMgQAAAIpBmzZJqEFsmUwI3/vgOOCEetSB49qntGrM9qobba9AH1IQdtyfiRpDEBqUBn30cNOTlel+ROcamSYWnUPachBf2dH/7blMdcvqBzNRPhHUnxCwzkrNRlC/5GgwjKfk/a2B/o7qVamrbXYVkHxVa8bR73b7AQGyRnIHvBjdQa2Vyk2zDB2JLzShEC8AAABbQZ9URRUsn3igS8fWMzYXk9Mq+MDqIxXVuhKB5vQkgSjtJT6JSpHmpOK9dQXV8FSdFrMVLn4/C1TOlGot8odWKsZCHVCzXaMLFwpRnRtzffgvkqf4iecFvElrWAAAADYBn3N0RH+KbhTgr6xHMY/2QndZPRpMsv5JSKrqaGIUN02aOyCA/tVS67KULKPzvQ3M2ZSKp4oAAAArAZ91akR/hNQq4woHVjdAiYoVEk1dl3y4BaGjLWpTqLq69j/cZmOzoJOR/wAAAEpBm3hJqEFsmUwUTG/77lcOCFJBa/WYM+C1k6KBz1fXTj2WWIJN3lFiTnpO7wj8DQA/i7cHqWnaZRx1Bw8Vic6inCrCA4O4XOxsvgAAABsBn5dqRH9/nILJGie4fD5UaXXmEKriUh+FqfgAAACDQZubSeEKUmUwI3/7xIFmvHCDZ8t12akwl+l9pgJioglFo/qTXNJYNkUFQp4qDJhzpCYfFAIWrXQ9O8aP+CIhDK2kFEhLIikaKbSUWzNbi6H4UrqrGmAlkZo5re9RQGhCccg1Jov/2GxYYpyEq9L39zvDE2Io/i/vPk4P9W3qT6jP+oEAAAAuQZ+5RTRM331UtYkDXGvB3dApBSvf+bMEW2P5xUOFFwLXVEf/avPuP2djgg2IoAAAAC4Bn9pqRH+ALNfijDv+qEgNt2DwKYhX+aIxYXAHwEiNC9Yk6O2btseoyQdM+47xAAAAUUGb3EmoQWiZTAjf++A44T2yC1yJtTnpkDG02asj/rdoy39ht3+VB//uMshD3UPhmeAzj9IjmzSBMZH7Bn25T3xc51+ObVJBTlLroFR0mja7ywAAAG9Bm/5J4QpSZTBREsb/++5XDhI2sYlYDT94fX90cHdUWtuMJHfC8cCICr02MBGpCutBAL6UkDKLLy7zPKWJhLoLICnT12OP5gktaUg7lCh5ilnYCb2P2H4H+nX+3GiUhUVoSBX2x9WbuvhYGX30W4AAAAA3AZ4dakR/im4U3cskkOWaidJ3UzcdYIuqiRFaqXPiJJRWssAOr87hWTOg/jtyd7zuRXHH9E6yDQAAAEpBmgJJ4Q6JlMCN//vrPjghSQWWzwVzm0Fb69aOUyH2uYwkl7uzi3xrhyoi3bG6RP0a6qHuC4sp6e7D/Sqm88josE6dK/dzf/C3/gAAADVBniBFFTyfgOahWDclg+s+Jkr9gf7ZOS+stAk+S2N/q5u0Nk+SxhbPMvBKtRp9lBeEwzIS0QAAACoBnl90RH+DDJXX+ypzivEOWRG7zcM5EY2ZDjmRy2dsh3v56A+eAR5ct6EAAAAhAZ5BakR/f5yCyRJTQ0pmRtmpfeSIf6wjFYJ0wy9Ux4VAAAAAX0GaQ0moQWiZTAjf+8Pt8hO4AJKfYhcjA9rjdnAwFf7QModT7PyzpMUK+CS0HRfcsjc2DBko1hnWHk2EmhC6FeTs1eU414OhygKFL6xPmfel0xhg6rysjXFWdjRF1jxEAAAAi0GaZEnhClJlMCN/++A9RIVzu3Ghrbc0Z62jyac4b1qAl3BBQ30vUIr9JNcG3B/7thUl3gfcZJEcDA9/o5HLdtRPY8fZ+i/Rg0HkiFZFC1tsDbnRuKWHh6GA5fZne6xDlbQRKVZqTruh+xSV166g8XuXlsnFkjpjI3LHla8Lz4NACVbazHfhnbOqwIEAAABwQZqISeEOiZTAjf/76z44S52OkshAe2P0PHhV4Dei7KErq3dgxYM8IFzbxvZbXSizGnKiEKPygFN4FMdrwWYARbLCT4h4G4eB1lTgxQGMp7418j0PFqDnRSjlz6lPtttNkfg1a3moBhbWbgd7fml1UAAAADBBnqZFETyfex1K7S+Yzn6tUA1X2wuv7z+sgh4e64i0lTn5rJfpZRmkJsUMXmXEYcMAAAAdAZ7FdER/dvu+WAiUqBj0v7NQWvwATKba0oofnS8AAABJAZ7HakR/hNQqlgeCR5WmEnTtgY9GoEa2D1hVoMAPb44LseFckgc9p/mRB2QhlABhbflpFWCIhI1jRMwGxXP5gv3zZcieyZt8sQAAAJhBmsxJqEFomUwI3/vrPjghSM6GitLPhfT3QjDIXqDR2QygB3rhZ0Opwh4x7j88E+9pDmthzmQVl93MWqAlaOPoKrUNvWPRoPXTCZXTTpy36yAiQiFXV1QqO3LKvj/ZL3s4SSgWTn43YwKoppbw9S0sfviVAvMSi0iDeQvD4QNHMF6CPVgieFql05KklKluC94yEGBUepu8pQAAAEdBnupFESyfgEm3uqz7mAS3M59dNbzoYEhAUaCO6CVGZpltV3a4ulLjDLMTNTXCXVdm075xx+6WQowma6lWj8Am7X3wHh6YdQAAADUBnwl0RH+DDJXX+AaXoY5qOgBknY4s2h3Y14HLG4oExzfORWc7aR19WMqT9yojLIGmtjRfQQAAAC0BnwtqRH9/nILJA77E0uyjL+OYRlC9eMwkZdFHKSM8oLegYjG2Mku8ZesfQpAAAACeQZsQSahBbJlMCN/7xIFmvqIPoV43rFJvWX4ODfVUK4ly5WYitO5XIzfl1xd8KVyzDSy+r0BDRR//Emv+IrjHGw2ajYgXdcYjXyNQqWjfnNEm+mjBDaTGTxJzcJx7tSgOp285AZ65BM2cdcwhRN/7j+yRSrtO21pzFo5NGhAw7XafqedlhzTYP+QC5ZhN/egGmszEX1+2B4unJ0IrXIAAAABBQZ8uRRUsn3iZldpf/xqA7EEaL/5ZgaVWGAVcD315eHZ0QAYXXrP2jWx1qG9PMHNp4FYUm6nfg0+ci6cjMQFda8EAAAA0AZ9NdER/gw5gGijNiQz/5kp5JhG8oUICnRs7uFapasLfsMKPh5eAEO7oRpw9QfNO9PihmgAAACEBn09qRH93OX2EwqlQCa65akqIwjXCrER9aHr25SKikoAAAABtQZtRSahBbJlMCN/76z9REFwH6H8EqFmA4qxAVHo6OWo4ZywYj7NSM+cdJHCwf4LAdrJSFZD3glZ52eoTwrY1x19p1bxCO7ICdQGhzEiQ96Xc9Y2SETJOD2Vn5GWGsL7J+gJXntXk/X8jEs5a0QAAAGNBm3JJ4QpSZTAjf/vuVw4IT+j8dsgkolcVXPJAv/DV7tMKPOhKW0Eb6GkMP3S2+4Qotkl28worhzeCxOdJSjc6Io+xwzEknX0Ktf9XBtxbX0nSYFhFqBO/1e1O40UYGwmBSjkAAABiQZuTSeEOiZTAjf/74DjggZKH2CCRdBnEYQWZN54htqxWZjo9aIG1TMzzTiRe4WNKcKB8WgbTaefOnNjdIP2vt10+UjnumK7PI2s37KTJhbyphwMl0U5kDcZg1Yv7FQFx5mEAAAB9QZu0SeEPJlMCN//77lcOB8tM/mnq1/U/2ketb7E4WfIqxx2WSC2Pl4TSDg8o3v0A5Vjy93uLHMTZs+BBdk8ByTQwMAo1JaoqmZVq8af7jmwnevUWJQ1A4a64zgNIiFZUXmvCBDFKozh+l/Zy9LkBa4jKOr70skQOcFAQ9YAAAABGQZvVSeEPJlMCN//760EqEjYhsbGdgimO8QcnqaVCP9qbOEFZrTmsIf2JpB4TSeaxpXl3Ulw1dscvV1EvErHcSHNSQtr58AAAADhBm/ZJ4Q8mUwI3//vgOOCLONnHJ8xun5qHA3ZmgreRBzw7h4GYMtLDZh/dmKRwhM+8ymTIp4LzwQAAAF5BmhdJ4Q8mUwI3//vBkESMggDx1VvVtLMss+aF3SXvKnVrCTWWcSrFJBN0jS1tZHuusb/RraxhnYjj1+/yAsYA0rgXmW+bCEUISclYRiHSnHWCOd1ZBmdTg6KC1ZbAAAAAZkGaOEnhDyZTAjf/++A44S5zbJqFTmyU0bct5evsZEZnwCITG77hmWTV3ATPGRlTvfx5fU+QbyuvFKr5KrASAROTbHME0sqsyGlQ6eQe5/+7Q4wQpIcxz1EsaApgkTsdRh+yx8JzHAAAAD5BmllJ4Q8mUwI3//vgHIVWfpo379KAL+bC1vac0rQQmCNEFSc2ChJ69qz7GJwQcygSggLTY2XDCMS/LkX2TwAAADpBmnpJ4Q8mUwI3//vgHIVWxnyGmcsFcA9oxoBl/brkT10aH4nVw83GGNtypUTAP/suuF7gk3SK4d4YAAAAXkGam0nhDyZTAjf/++s+OEudjscNzZV5CCe/aEyxqBVLFh5VRKVEDzmEB9Dk7ScQ40u1uCiy9OuUrKIzx+gIl6P0NG1LHJzJ/qCDx1gum2N47t9UWnFNio6Aoe0ovYEAAABcQZq8SeEPJlMCN//77lcOIMTUMrHZOr0gSKOJUf3LozHu8r1h9EcXvqCO3mOg+d/VC0z5ZYdgFjL3CZniji5n4XXibsuJI5k3acNamysNf4eKfIu8h6JjblX+LUEAAABhQZrASeEPJlMCN//77lcOK5kGh6VLT1xLqj5X9bM3RBN5M3gOtGKjSsWOVbXT9w5waJDJuT3NaaL21k2UzbmVyf92Io8RDVoRP2XUDjiYmvZj67mQ91VySU1YVL5C9Sw30gAAAEJBnv5FETyfgOahWDUeRv7uaiJL9EdNC3IbVzoaH4e6aSo1n7bs0FRXQ4JWRcXvz38ocIB8WqDlMJ9G3aBO6MMl5bkAAAAmAZ8ddER/hRLKDHq0j4d47DM5gwud4/DkjlPlVFYGld+N18XNz8AAAAAlAZ8fakR/f5yCyRJHv89kYcozAl50ypnHlpKXwY51Kj7W0wQ1HwAAAHJBmwNJqEFomUwI3/vgOOIMTVxkqPJRLQpQUqKkxcew6gejhK9dEgRRiWKoB/5wIFCYAVMLtY0lhftvZnZTOruxFhOlh/7mqhOygkTIBTqYT8F+//eGeuYR96Ww48vtz7cCy86ywFxOdvLxYRAQgapIAWkAAAAwQZ8hRREs331UtYjmVOEob1vuAVgoDzC2PbDUkYLtXfYTEoiAWJS2xEv0kGt1b0TAAAAAOQGfQmpEf3/pQAeIEEpwh7Q6XdldFbTTk0g56u0PL+bxWoNXHexxjGmHvNW4Js/aa9dUwMqYLiqLmAAAAIRBm0dJqEFsmUwI3/vrQSpSocjmBEaZ+guniyl7iU1d1HS6ptlCPKacTBMZFdsmw3nQ/2SNFSQYbakk7FMBdPhxaD1VZVNg6EmmIShve7vOeawjwMsQ6qn7zs+OxVipM4voEUAHDO+9AdjYKvS5OsP1RkpHK5H8rsOaGDiUxa546nHEZkEAAABXQZ9lRRUsn4DmoUnrvS2j5/oh09N4LvXKgJJsJGF3dTgk9NBZHHzkIto5QN/1CwwGkGBV6UhxpaON36OliiklObJaIH+XV0IYQQBpF7sgk1BWy0nUlXZ4AAAANQGfhHREf4R4tlBciUsqFFuxA8UwpTFd6S8jYZCqRb3ZWGJYP1ecFzDRMlHObUR8+QdZwRJhAAAAUQGfhmpEf4n36qqIKnWUd9Njiw0QSxoVB/oX/yBhz8h1/xIEQ55PMKdy01KoAx+nvEYqAmAd9ad3OBTPQ9pOmJCNVk/X7WrPcvXMd+GCzMVi+AAAAM1Bm4tJqEFsmUwI3/vhUdkOqEdG62jueiHqqB7YeLSYauyhCAzqdiH19ZLvxGPaZ6tgKbJdVHcaPgr406CugQLvd3FPbtzm5AxeSWxsCek54+8ViKC5KSPZpGstBmWc1F5Fz8YmM3pXrKfxiE5SEFNGyFV/qkcpgvPYbfMJYpPyMe+Mok3YdDm7tJqwEmFuxpR/c7HOjEn/q78jrYgJECTdXUMmYN4j5En5F/wwQM8kA2ap5beuNPWS5eYz2P2j4BU9tNJrSFla6y+1AItxAAAAO0GfqUUVLJ97HUqt10psKM2AYDKr6M37fOlTxCJJ2VrMROQNURvJEmBHrG4u6Jh8f7q7B191ZIjfgM19AAAAQAGfyHREf4BsKlO/qXBAryvwFqs5i0G98z4g8HSYOeOjUiUe2uLVZBQrxdTZrbxwSpFGxJ6EKPGOOvOBc2c9rsEAAAA2AZ/KakR/hIfoLL7/BzcvpolKljM37OToCQ/2FbQB5A84l3GTkJR4Rl+5ZoIT+HkMEnzHYh3BAAAAg0GbzkmoQWyZTAjf++A44rmWNhGWCL72npbHi2D0rm32fIO2vw38MGXgB7/2AwMnd/lvD81EFlh6sV3nHxYoB0OZtYyq9r2rygD71P2o6DBvoGCJwkGKOl4eFiEYvK83H9emCgm3wnANdrFoijHM4zU0e2RMJh7A5zMM4OT2vz4r/00sAAAAS0Gf7EUVLN99Ta+j2QwAR62o44atyjblHXj8qcvZZupH136TPo8NbBq8vABIfCMtlgbq+SEavzCvqU9ljoozFu3hOJeYHrkwhPV+IAAAADEBng1qRH93OX2EwrJavM7/2lua/dEitfWcARv61hjXJUf8VE2rGPeKPTvdA5ULh6n5AAAAnUGaEkmoQWyZTAjf++s+OK5ljSUOA3TFai4RJE/jHgw0W4P9CLmTx1ROtoE1Ct/u5aZKpqL5UFk907WKqU+Fq87HMy5CDtuyG1GYY7+oe/qacQmKPQggXqXzfVJZxonCiuPjr8IvShkyg9jrNQn4JR5rA4cMrT0Zg58ZdIL2DxBO+Ztm7EZbuwpX6hiLDW+Fu7jZRlkzHZVGcLW+A4AAAABoQZ4wRRUsn37EZ8k4yrEJ3gXZR38ckSC4E/G9JeY6nvqw8cni+XpetBPoOPSLCtYJ2bS4fCf8qRBknduGn/zENrDlaOySIiDPyNSjlsUVOY3Be7kM82KFFJZ4p6pqGLRuc8L/WHBJbyAAAABOAZ5PdER/gw5gGGHvGKdek4JOkgzZuCPm5S4Jttsi7GrvSbpquZEjT+l8DryKb1bKgXccBTGldbxppBvtraeOctJrix9CJg+uODiF1InjAAAAPwGeUWpEf4TUKvadqH/skJdonVCLe4CSymbvWcZPcHMhqYKXxCXq7uBXK/nUvzCA7bRhXe4Org/oKbRblHoZbwAAAIxBmlRJqEFsmUwUTG/8IvJDIWLLGwa/ZHF/KLUN5/pBVYdMU+ZmVpOULfEv18kUVSeOAd3niURef0hBEmg6OBZ0i0XF+KXTRz2y9KRZZglHmsDhwytPRmDnxl0gvYO++Jtbp40k8mBlBmde7Jtob6nZkW75fdtnP5nQj7Bc2x52eeW7AOXILyoBAXy3wwAAACIBnnNqRH9/nILJCQH8el53caksCb78m2K8PektcjdV0nLgAAAAtUGaeEnhClJlMCN/++FTgKPVy25+n5usIXnTztuq86gMcZAv//lXvQMR4OFxAH4HHTYRKq4DSRdcbN2fFwL+DmtKdirEql7t068169vI3EBCCne05YqLPOllLz90IOSrfUDiyyb5NfrDwLg+wOh7qHwq2BSxmHlV8XQlZWmA7Q1K+zPZSUHgjhs8Ax1n9XnxVLQtfRLLyA/dBKYlK9a5jC+wfdcMLbf12+YuD9YhzRWeXQP6UjIAAABLQZ6WRTRMn3iZlRG9wMbG6rulz+i3xcFmicnkVfPi3vI9sl/efAzv78X6GnDlL27msiMoZuydUOS3pIwXyOEjvXuFAt/zUCd6+qXJAAAAOQGetXREf4MOYBiYRsM/4Wn9J79PPYoQFjb6GuY8jqTLyI2PeCs+yCv34RJsSxBkwViDKVa1l5B7DgAAAC8BnrdqRH93OX2EwrJavM7/2lqioT5DC3FIJv7vm//z26CikO4tVaLBSjpb/LkdYAAAAOVBmrxJqEFomUwI38n9ga29JPKIOBHSx+YsFh8iUW27g9I27uSlq/OZWnSYhIZM7ODNk5Eyw1T9OF5rr+J2dy4gl+9RDUCQuNU7dxcC1S9rJoB2wj0ZCKt9fx0BJCscbIz9DE60T8UJcZW4IXI5PQ04rPvV3fOyHU1DTpW5sLAkcKbZY9J5we4BPqx7YVed7ploglm/l2SkYkxjUmxsPzRVLOnRXmV9hlstOpUJ/aOYKohirbc2wOqCvDKXJPhu7uw6sdlwNg7ct2PF60e8zKjuPUQcgh0U4w3KxPWl+MIuQ3Ce8JUHAAAAS0Ge2kURLJ/NY7qFXvgUaFUOZCNBf1AjC4hZbIHjRp6oas+6Qqk70vXUDe5rKS9H375z+RHnXGdynCNUDFHftXvoHTlj4VbCTJvnaAAAAGoBnvl0RH/T29sxF44A/5fnAIOAgxy45Xst+FBoNPddunOVgbfDRXfDZMjQfcvkgzJf51llDHGEOPZdOqXjlDVUZbZjfZl0loJS2YOt48R/pzXYFOhFufNJzmM0UhgMuQR/zyOciCsNb0DBAAAAKAGe+2pEf4TXwDNrCamp/ew9bCI67/4OVldsPjdTTI086/CfMW5DjVUAAACXQZrgSahBbJlMCN/5KYRFpdsBNJCPfk3shMTTbZ/+m1g3BC6EDF3hvpzFAan1Nwk0GhKdlNqVrfEyKn/lndpE6h/0uKfhJLNPRapdj7uQNAbjcketIZV80xhYt/ITlWMCxK8VuHPNA3X0fatIH0MAuSlD9/rmnTCGX1rX+/0XG+Uq7kJlIC2ViD50CtLv279/vtSKlXHogAAAADNBnx5FFSyffGVVR+wTASuTnuC0n9aZFyEaIux3eFAxaKRoR6o1P4kyonm7L6VYQyuDgUEAAAAeAZ89dER/f4l2UkCrRbpOuJ2cpJvutsxXLqDlkAIuAAAALgGfP2pEf4ERwdJRbH63u/+0Ksgh/Pd611GWPgQqLOAjTiLs6CS8xHLR7EU2U8EAAADGQZskSahBbJlMCN/Hz9Z+6Xzm3/q7oYd0G2hnycvD/dneZ7oEkBQficNkv3sQWvRjB0O3J0i7f3L9TFI5xTtQclRChNrb+NvDDZvgNHtISAsf4R3ClMzNDi8Hti27fyjTiY/1PVoJovYx9gmD2NorCNGDW85eiStJX8zbjsBw/hmP8sN8jZiXmeBRjX7HLakmUV39clUJgwqNnh23yHABH+5ni2vm8+lQKHKUjJ9gK7V4lXqq7mt6OcgdnKFMyG6IrIrsSQiBAAAAOUGfQkUVLJ+WYb1tdpfdVu0F1k4kq5RhUrdGhAFDHYPp01TbmPpwo7SUbGx5vDZ6pujdLsJlfZM4VAAAAB4Bn2F0RH92+75YCJSoGPEHs1BaPe1nKVFLx3Wq4NAAAAA2AZ9jakR/19HKwpKb3hDHCpqEpgwPiuiuAowHjiceMfYzoQjP7bagp7T/MwoVg8eb8n/Krh+xAAAAREGbaEmoQWyZTAjf++s+OEuc2vQLdK/i8ClDkp0+Xtapl/Upi4Dlsp5sQ9tLIcelXy8eqVNGGFIhyPtEhgrjW/Ytb9bgAAAAMEGfhkUVLJ+A5qFYNR5HWQmzD29FNpojTcXVb5py+ZttjK0a0ojZ+6HbGv+AY1mIwQAAACUBn6V0RH+DDJXX+aBtjXyInrteZCVnwVE2lOmdE1eqH6zby3RAAAAAIQGfp2pEf3+cgskCrQfAkzm4yXvXEhlAALbstuTsEqLFPQAAAH9Bm6lJqEFsmUwI3/vgOOKjnj2kz6wE/ERkEFrp9mvc/FnvGTzuEvBxPM3+7YFlsfMrbeAOREsmq2wIOZ2exoKE6AbcyGVU2ti5k50trMZUI2gNhJRzN+j4GQsPgGo8hxfd6F/BV5QxpjTfZWDWEBGtlz/y9p06ppT6Q14s6X6BAAAAaEGbyknhClJlMCN/++A44rmWNeEapfe5ojbjO3gLbIo9IFs/sem9m8C9hrzgTXFpAoFFoez7AxD/afkFQ9GAOeRsTI8AKIit2NJwfsgAJBBdBVBQwXTcfUZuy+cNB/YDXD9y1MyNM/ElAAAAS0Gb7EnhDomUwU0TG//74ByFSpwKu/yY/TR+8rXmM/1ctJUphchJ0QahBg9/drvLVbB1uEUK1tv6c0ET7S2VOcyfJsGdD03uaOxggQAAACQBngtqRH93OX3cYVSoBNddIFIHoxaPZCAQcGLR8c5RXh9QJpAAAABcQZoQSeEPJlMCN//76z44S5zbHmusf8GcmWpiDjGPSa9PauHTPUthM4Tq2k7DqWOl1BucpTi8VQ6775shA/8Z+1n8HnIUpmW/CGfGL/3h8FtqksT/NymcW9deRJAAAABNQZ4uRRE8n3wHW3NeA+wKIZuyYfA4eD0aAfsQ5xdegmzrUzJAr6D2LS/mkx6XUnbpQmRhSBMw8zIWgJzltDGt3NdjhobjGWShq0kKsGEAAAA8AZ5NdER/im4U3cskkOWaidPqqW535NCXBHUD07pfk065jmmezMmm4OngZila6PBDQ8SA7tNzFMqQHG+AAAAAKQGeT2pEf4TUKvcv6iTIVCcF6BiozlPa/0aj/QBDSf946iRfptC1PLluAAAAXkGaUUmoQWiZTAjf++5XDhLnNr5FeZsxlZSRTEFl1hqtFA4oqhfVWwHju6NzP7Gzpz+tHjEfnDW+Oh5YrQeKavE+wrgFqrG9jsjA3/ox2THTdhWkSqy4aSCdsW/k608AAABAQZpySeEKUmUwI3/74ByE2ccN61JFkvALjEXs6KRRlHtA05Cr+i+xeGCVyVXxdoLYm1mfGZR+dej8GUfR5noE2QAAAGxBmpNJ4Q6JlMCN//vrPjhB4PuVF9LGgPQZObH+9xBHiekie60xU2mBEcs05OcSNiww67eP7NZPCqU7vyMShW9O7t9+XYJYXFlwEsy9Uhd8xicrYd7APjwDqdkyyfwfbzlstFm80HL6Q1TQ9IEAAABvQZq2SeEPJlMCN//74DjhLnNsXLgVjQBiBDF31YkR99oQqK3IUPQox0FL5Ngyqz0vr3e7zX4D7T/IzH37YlccmItA415mW5YgU9nDtlwhJalm+PnovRivR9bYdL21zzDZV93jsLfIFoLu62yIgI+HAAAAOkGe1EURPN99TaQSGVFDP8erdHfbNvCl7pmTW1uunWoHKJh+8IYvX3j4p8BO0arrZ28AEjtFV1nfwbAAAAA1AZ71akR/dzl9hMKpUAmuulsWOvCbvqAh1l1Ee3e0VIjtX+Pd6HNzYEzwaHU95p93feOAOeEAAABQQZr5SahBaJlMCI8YVKCqLqWXOXD3Rs2/i3L4bZCiXxdmPoXkYSJoYDVtTiH8gKixVFCbBh7U/5bLigBo+Lh/IlcxdgTekBHV7sq/sMm3KbwAAABJQZ8XRREsR4puFN3JW0ydhYi4T8lJqNEyweZMIfVIyGncYAQYhG3u1bYLctfKK0mJXCFln0Vl2u1VXLGsj4tNzvwBhCF9HkDygAAAADgBnzhqRH+BB5KrHlRVwgrrBwaKJDrhThiLTrF6/6zM+QOaa7ng74x86Si4GnXs5RovnBY3UKlHgQAABX9liIIACf++DLeYgN9XYj5zg5B5QJuOQipq5DurWEjnW3UorRPL+Gmp7W/eIApq4hA6GO7MXHIyDHsVucz/aMFYtKTDM8z25ZLQfiNZJqSiDN9723ESzhz79Cwt5M+8vOWuuGhG+Tvq9foXjrYoVXiVDwEsVfUmlk5TXbEqyYI+MZpo4mDJdCI6FifaFbOBypE2MZf5XWRt96wDzE+1vc15C8Q37MxkeCbaMkDN/t8mxkCWv8440Rlwk+ofcevu9fMi+Y2lGAjdq67Z7L2nX4tpcpo31hZ+kf6KZeyBIu+wE7zfBEWJkrDn+nvizONx5DVOkgfozhAUOvJbujzhGRN82gJwvYpRNAkkauwRrAkj3H3NX7/9FmB7s42z43E2AQejmKa6top87jsooGgL2yW/vkquH26p1/ZTn8L/bLpPfHN4mI2HsSxn8Xasz+QZbRKng1fvTxkt/mpfmK/kCyCd6bnIW6xeE5iW4Je98X1bSpyRG4ZEwQQRvmmA6Yna0ada8GP1pcw4M/JCmmV3O88PrasibHCahaFYs9akqjV5DdyOsA7IxWcVi1KwMrw65H0WqhITj2I/0392KBRdDeMYXdHrGwNEJhUIlQHIAJ5ZcKZWWU9TtKJlEqVRqCAyqMi4WjV04NQ1QEfwgz+T+hEehHkIC/Uc5dPLxaRa5f0rBpWg/CaNuFXLWTTpizOXU12VDPjD0PueUI5gcFxjM7UZKqEXIER8ZxBIM2uWna3B6TinIB//JF7wnKWG3KqQAoRJfs5WvYaG3M+PsmcpsEDOU5YG0Kdgf0PU4QKwIo7CjVGXnuOCeJP65DynJB1TaFN/Ec+KGgl8BEvS9O6qZ9x0NRWtKyaKyidn+wHMOsyvRqTPd0T82Zhih55cQsPNDIk78tfmfJbq26dnj6niV7W4C0RM63t1oa/61b6ejC3GPDzgPht8QbEWOmMkGWQaqDE1JlXaV0FU3nPF+k2FhLclRGSF12k3fSd0B8xrBHuCYayWiuH8ryJrtKZxJ5MQghDgBuHJO6SC1EDjEfxZGD/r6VzWAzizHGg8AmGx2xLAaw9rsvxLPUF+/7yYJueeBW5R93rr0FqleLihl5Xj2Aq5SkobABmcpWGhmxzvsHrBwKqxvkuHcaxY3VOBxUoDZhnQfEv+mZtl/1NlDKLlxwKt0ekKf+H85FMDoy087GajtU36Yx+6OhUz9cVW+TW34lW5zJzNRP+Hc8uhJDdo2lXVyYYbqCDCYtOyMVhQIpdZz2jF3WMV035t+ZxFpEnL8ijjN9Jg8JcuWyjfxs5sP97Y4nlx4+ag24HnZ+axqomYDumc0gNVyEUz6FZ9JbgAuwpXoA5ORKsBxrFSQHm0NhO5MryD+P4NGBrI8OniGQEltNCWxKO4CaPEW/ayYPnua20xl5iBviCOZOxGovBGA5TXYq/DBpZKzwKVYDEuxwEmec9Nh+C+oW4OqD94kO3pzQXD4WNt/q7X/PFLl+l9H8cdZ55u+BvzP0FXMoLDoRw7LJfWY5c+wzINL2XuOAYkfJCpNNcv1s6D4asebxSD1Oa3fRtN5j1N6T3U1HHUr8DPuALIo5ZN8ejEA5QFEH9MczGJJOxYHf2Ea6JmC49U20+9S5lgbbfIDilnnqaUN4P9cfr0GozChKSlAt+3fH1bY8x71jdkn9w2hQSZUYbU2pB9vw4KlfdleH+KnBRal9A8sdEAByt88ulnusl0zuMQqjLVJjZRTm+bVZlUB1A17g7C1j51Ik9g4dNmA5mrD6GgZNkTb5nGM3aTVSZBRdQchSdPGxKj9iMzKtky9yCgMnQWxUthLxXktBy3hagbWow9t7gvPd5DjMYYCoMNWUQ53ExMXz5i8JP+AuTFV+WVNG8AAACJQZoibEb/++1xYQAAcZSln5nDonHJ2FY8IUkCHV3Sy+GFqZcQ6LZxLbyQ6zhy/gRLW0pWGMhKeooRyMTxiJ2j2Rfe55qa3amF1vmYBco00/ApFnQ2DT4KOcgnT4ciTok2qz2OCLOIFJiYrYM4viFJA9fmIGkJAUIM20h/dudGJ62seoc2Z9AEofEAAAAnAZ5BeRH/gQeWpd9/pOade/rN+isrvDcxP1hTmlgtc64uAfKYMixRAAAAZ0GaQzwhkymEb/vEgWa+okZYoHAhrsRiqhh///5/poi1W1PBAeIgbreMG6vERAWN1PzvCg5/c1bRzTEU+q0vqY4p4w0N29AZHonm2TuYWlH/CYj3AyJ2CEW+WckK1sxFQX9Vkp3Do98AAABmQZpkSeEPJlMCN//74DjhAvudrQV/Von+LqXSU5EnN5Ici21RizZEycc6yRI7HW2mkSAr7OhvSjQAWm9VvWhylF3l+5QdacIDtyQ0c0VNiZE8gmdnua+vNHpLu+/zui78ynwbUg5fAAAARkGahUnhDyZTAjf/++AchVbGeq15oxbCanNPtandCKIk5UFi2McwnOkyw6FFIwJZCyx9BQ6ETJstOW6MqH0KmE3QQ9ScPBAAAABQQZqmSeEPJlMCN//74ByFT6UQhh76pV1tl2CE897btCzIgaESnlCrkP/rn7NokgPeW5v5Tlj6RqOhWJLwnxuEYh/nbkWklEpVsp02UeI5hcEAAABmQZrHSeEPJlMCN//76z44S52OwRnQ4IPpaB5d9PcQ3xV3g8rMCFmB5bsXTfJyvJ25sgvf46GpdxXXkPPndvnlrxgCG3+15DsD6MgtyhVrvxbNr6GQUtl2LeFZjaZQES5OWM0rLLlBAAAAVEGa6EnhDyZTAjf/++s+OK02js7X6gKTlrPgSr7UwuS/819DrDrfrMwhlkZEIEQkkAjuhvWd6ibuHYrJ9ufrJG3PFqSfmXeftKX8TepwZpbdTG4a1gAAAFxBmwlJ4Q8mUwI3//vgLpU+4B85rBOb//EpYgZFZ6eevNZXP+wPe+bIMaDyfbTJHzD4Kxd9xRs1m454YoGfxrw/BKC4vlUK6tcXIpJFyYby74nDDHCTcY26K03YRgAAAGZBmytJ4Q8mUwURPG/74ESpX+uSC7HOS2CPeetPFAoB1Gux145k4OcRwOiiCVIhzw8zcArn6HvJ2afHpwjimmjCVOlWLWAXs6/dS2mQzr5Tg+USVx+WhFO8PWOsAE5kFw2v8j/ac4EAAAA3AZ9KakR/iffqqvm8Jc9gUquRMvf9rGuOKZnd0efesGD02uCoFuLecT1P7NzuAbHBt4lonQco4AAAADNBm0xJ4Q8mUwI3//vgHITeeM3FIGYUnOaFjLcsMmQoftZNP8Bb606uLDf3eMjrSo7Sf4EAAABQQZttSeEPJlMCN//7w/CbbaisfdZegDj0M9TyOi5+HYmOurwQpxoiBko95zRoR8EUd3PPFRCKdewA0ZC3VL2iuZhfJK/KwG4Yg0/9XrmYb5AAAABtQZuOSeEPJlMCN//74DjiuZY2PsMDIvgZTyBAM69A3+mzj1NB9ReQMa5lpIjzkgpFb30Nv8USZUa4L6VDLhd3GOAcEU6/2m71HhKIVVQxQ1IYgBzydJ7TaPB/Lp71eKZ3cme7OXKbCXDl4TYBQQAAAFFBm69J4Q8mUwI3//vgRKv7oyEofCjyt3vsnfsah2kHIdELq19J2RoKoC1XL0zZTv6sAPxkSDJ4hua2mqjbC4lGiBfWgM8EmDQTaVR78G9JBOEAAAA5QZvQSeEPJlMCN//74ByjvhR0RQiScU1Fdfdc2SHU4FOuYJlXYGqlgNrIdNY2a6p25gOjejiWvqYHAAAAZEGb9EnhDyZTAjf/++A44rmUqQo2o78x9zImFz6ycgjd223QsrF0bQ9oStMycJKa04q4jD45hvqUv+ExtcHS0/sB5B/uzSniIvGiDyArKszGSLTNmQx4zT4DAS2BkYC1Ww5Py+EAAABFQZ4SRRE8n4DqJeP15Rd2F1lVcxp2jzwtUtQ9ScvCBuIZmxQraJl3DozW8juXqxWghEa2yC/V1oAhaIA2kP/2/bNYpFpmAAAANgGeMXREf4puFOCtReQubvl/VjVVRARk0FI0mb4PyIwlMKV46WFsEaeKjqMI1+R2+DtK1F/8vgAAAC8BnjNqRH+E1Cri5ko22NzWPuDSPi1danTYsunQ9icYtQPRSlq6m6AnOlFA02op4QAAAJRBmjdJqEFomUwI3/vrPjidIX8jzxcO5M8WgNUVjZvTZqttR/yezTvM/BnIL5HWxcNwKvfguadDkLdYtXjtonK3zpSm1H7UXsq/0YHnzRgNDAYnqnNKKa3EfIz15EOoRIhWAiTMZDR7m4p4cwv3/77a/ATHMyY8Lm+4XtVBteQnl4516EwPOY7Hx9ZJsncgz7c28+7ZAAAANUGeVUURLN9/4NxQllWwS6YjGgAHhfIlm6iYyMZ/pi9Hb6DSetZgTUH6s+U8uHcuB/OfkzdgAAAAHwGedmpEf4SbRPB1qbWDavfsuHn8SwfaNy3M5M/M8YEAAAB4QZp6SahBbJlMCN/74Djh7MxwQMXA4Xa1gvMmJfG5IpGYMILlDBGG9SiDcf+miA2uUriMuCZPDw2oL49uTVSBzN1sGzND8Fw8Hhrs/70fxGPxNhJdvqFB0w+DC+aj849W7UJcjUb7VkYGXQv6/CDjihFxZP7ZmzrBAAAAUUGemEUVLN99Ta+lvuryT/+7Fb855I935eQ/NKyrIpEljglVrbW4VSwnmC7RuPcxJO4HyTlFvPN0kQgJ7FjETsgCU0+IAsrdcNinbd4czNuDZwAAADABnrlqRH93OX2EwrJavM7/2lXT+mo4pwFqeXYue0BFgaMHvinTIiRN9nQe6i5heOMAAACVQZq+SahBbJlMCN/76z44rmUqw207hd7OqeG95zjobXO1BXc8GYszEgUi91ZFitoujcqtv53F3MqCQuvkILWF7cYeyLidHTm+UXkyfQp7XBSq9+bk8rzweySNLOjwjQrKhpbObcN64LF//YEwizj60QETaAs6xDcg6lk3U1uH0HDFd4ifBkn3nk+cTzIVsSPERF68h4kAAABhQZ7cRRUsn3xlVUbBHfvOCCo61ctk1vvWUTe2GOke7+ZMC+0wb8hMaFSQHz69cBeyCx6e99M+HncZK5sJxhMcQG3iUMQJQG5qp90WppVxek0grIUNEg+nMZ/OSgHVaYGxCQAAAEgBnvt0RH+KbhTdyyn2mc7b0qS5IBg49T6i1mEqW9SGK2rLIsvaSKB6hosMeUz8+LtQpC14bO5miLujSI3IhMrcLsnaQe1F9UAAAAA4AZ79akR/im4VKaciB3zfJVSkhw/Z8IUNrw+u/pOSGKeH2tujCA/SRAiTtkOg1ajX6LWUHYFkQrIAAAC0QZriSahBbJlMCN/74D1FQ4s3n7j5n5Nuu3X//9TPET7aPn8s6EoF8yfvc/MfvxyY4pGK+0qv1Wm/xldOtP5Q+S6fps0S+m/UZnQ7EwPsBycEwP1ZVg3aRa1YrfiIsrKz8S8liNH46Caw2vDUSR+GdL+B35GCtu4RO0FhgZFHJ26mYluh/0OmxLZskzny1sgm0Ygl0PCSjypQA0smNK9TeZ6vZ59Qihh8Ae3AfAr1QlZ6NWXTAAAAPEGfAEUVLJ98ZVVIdZdMwtzfMxEdTpKGRFWgnHy+yhWW0NeLWoaMzR3O/Z2SNXpaNidD21Q02s62yRuP9QAAACYBnz90RH9/iXZSP9h8cdeKVvwRfd8hZRQQgN0TPngGS7Fwab7aQQAAAEQBnyFqRH+BEcHSusAq258vWb8w5iZsvh+pHv8c6GNrxfq4gOcru6UctZ3xaies0Y6hBXGyrHqsGGLfpOkCH5M1jDKmJQAAAGdBmyRJqEFsmUwUTG/74ESpBkuxG2KSXcpc2jk8koxEI5uEi5yoS6e9xpdAt72c8MSnXbbConlHP1BzUnJ/+3z1Aoy5efyGWrhfPczqcoqfv691EvI+hHCd2LcBKTSXuxUnnL/0bEnkAAAAJAGfQ2pEf3c5fdxhVKgE11tH5m8PYCgh4Qjt0i/xdPUGMftsgQAAAINBm0hJ4QpSZTAjf/vgPUSWjRx66POWg+1VeyDdo17n1pRGCAZv0EnMYI22P5nBPtXTAZp6Xq/Jn7fkYYL6sJtWyYB/PJ+yuGPOHKNGaRXXtSKv8ZtxjnL8aUG9NP/FDr8upgFotXBQY8KM92B40pAj4Gj1imdJHCvfbU0iETKVBhp1wAAAAFpBn2ZFNEyfgOol4/XlF6ea6k7mZM3z9p4e8X2Jj46Hpmlke/qavrQOjTwS9Z47MV26uYG6NPmmXkDbOOyLBafQL9eNMNmRUBMDuSM10MkN86FV+v+G2vJvgoEAAAA6AZ+FdER/gwyVRAsFrTv5yQfi12mftmfw7u/eMLjWn2w/45OdHFn7ezFA89/eGqdczYXgEIo3louwgQAAAC0Bn4dqRH+E1CsRGdq/m1pty/GwHijW++zyvH/Gdpk3chKR+bVbMoyVnPmV0owAAABrQZuKSahBaJlMFPG/++5XDiuZY2N6PL+cwxD3ivJ/8tkqosMlZ2kuhgf0q00kNIyHwl6OQiCcnPCaq6e7v2Qvnnewqy5mzqYUh0ZzBmS09jJQtDXRECUYUTYcu7s4jMLzcEOBfNCz9U+yd4AAAAA1AZ+pakR/f5yCy9StiwNVmr5zLjqXbDKcDqtWtTICUwtkcR5hRgZuV/4Inokvtcj6Vro4MnkAAAB3QZuuSeEKUmUwI3/76z44S52Ok28EOzD5G2y58sz7giS1bNt4y1aleBVYPkGtESwSGS22z3kj7PNRXWIQBXSfnJ9k+hqAla+toGcJ/TJgevZWzmQsCjk3toczbqwsz3LPlj/Q7tpVb/p36b3WNDT2urx5wyQ2u8AAAABDQZ/MRTRMn3iZR4RsYALI0qaMPz15tf/PCJjXSCVLzyIjWEB6NIVZuFQd9IHJUmnlM+pGzTuq0jqVN7pbBnNO/83bjQAAADkBn+t0RH+J9+qrQkFHiucZngCfYVC6NrC2q1CVtl2NFIiI9/8efIy7v0ffzAkuTuoM22cIjGJ9hWAAAAAwAZ/takR/dzl9hMKyWrjL/2mAOYr5XPeIQqNEnZH84kJ2zn5AugiKzfdzkWMm0Z7BAAAAf0Gb8kmoQWiZTAjf++s+OCFJBZyyg1Y/s//daL2cmKlc5TOgqXAetAsmctqz5P823VwTaTcwPX+/sUmF5uTMjeRSbn+NbF5LxJ6JtQsih24gjtIz0ieWd7148vl0dKSJkfFHV8qbphVazHMzamX2zJ5KbX5iaHfn+RVsiaeJeH0AAABZQZ4QRREsn4DqJeP2RDmwj76BdzDR/efXogJAcIU8SeSQHCj0dqnjrDVpTgwQJI31BTaAG6INqkujzWTGU1zd1DjtXiaRFbASYt3GqCB6jgT8nxWzUXikMYEAAABKAZ4vdER/im4U1zc+vVyb5fhhMvf/N+sn8Zk5BtnQgFm5gyE+ybdRK46vySho5BSu7O5YawlwbRtBqouRndC7tT/GyEW3yt5IFoEAAABPAZ4xakR/hNQq3d0NOvTMPHejrhYTBJBnYWLUeub3UrmhxZK/nEFqJcD7/UnWAqDmNwXcWBRsIOSn8sxa0LKyMIigPgimkAATId2ZpEAtIgAAAMZBmjVJqEFsmUwI3/vrhHZDqhGUnyzSHEzswaZ2RxqKGRynYBXMBuEuEFfHmUDj9ux0sbihOOVmVeYWgA9LOoKwJ+rMjiSQ6ZUAyFQsCjjaDAUTNJtjWai9Rc+L83o1bWP1S2yuQUsJLNLVuFrlgQAymRGcNkyxEtGUkanf24vl9rYEtUPRoVcJrHcCRxamMNaxH5ea8jkqFJL3PStn/7XHDDku9uliPPRnGOvdudC3OTfeJhZuBI5KZb1w9wEivL8R2sp5CIAAAAA5QZ5TRRUs34GrPAdCUxfCqZZ2unlDcBy8LcrzCPU4ls+ywf4Ue1Eh5IMuWUEO/crgik+To25i+g7hAAAAKAGedGpEf4RorykjSfm5OcGVQs/wa5q5I3WVt0A7JQcZioy+6Cpd3cAAAADPQZp5SahBbJlMCN/74DjhLg4zcYH6PAnbkTY8iUAl5ul5AGL6HlLmdMJTDc6cLh2fLVFs26FJQNh8DpEVTfPwRt/7mERhKyfDeSIWFU5yQMCYNHv59nvZFp/a8WegCK4HUMfdTe9BBOJ+XPTTUutM2EKDdwEAeJMZxuSKyplgd9tOuySGUaYtDrIicCWJFC5wXwmU8vxDFfE04OEohATSR2A5oLuSDW3Mvf1CjoKTHqw+jE9AVMQbxRI4vzb1WUJejTxdtN9Si/S5yQWiKU8cAAAAUUGel0UVLJ94oE82Csl+UbsHxzVXEoxsMq4cxym8W9nAZxQ5JIcYizkCGXhABajDTcmqAIpfxDxLPxOswTcGYTAAIZm4AfUvHOJyB+CdHlrCgQAAAD0BnrZ0RH9/pbSDSQoUYXbJ3rSzc+kRtLUqkKeKWWa+GfiCNMxON4cEELkIjdppVpoCwBlXx1ZSOo1lcke3AAAANgGeuGpEf4R1fILWoVH4wTcI9HBEHVSvB4bqbjhbmfh2jcyxVPpzF2b1jR4X1zi6YlYACKdbQQAAAGBBmr1JqEFsmUwI3/wDwzHCe2QWuq6WGF4P0C2Zsdw46Bt6zA57Rc2TiTbjW4H7HgKfR8Gqh82t98tMHubaH8epd0G9TvymJA0ATEGErg/d+4p+TcjdLCZHjYjKV6SX5pcAAABNQZ7bRRUsn3xW3JFluUBopv+0Jh109W67tBai2wVCw1/Fli28auwIi5iiXJElON5pRsqKGOlx/dZECZCgJwmVdX4qfM5L7ofdjxymz8IAAABAAZ76dER/gGwpqdu2nr4gydT/It9EKEbgleYAPE+8ztHAa3x9WnkO6FBi5Veh1TP1eDEtNkuvOywKMxIz1SR2hwAAADsBnvxqRH+J9+qq+bxiZHCFr0gnhkOz+ad8z4gs78ck69ZSTjmsIwQTbotwcaKIZaLcxmKWlHaf9l4JxAAAAKBBmuFJqEFsmUwI3/vEhS51ZOAPrFCl1ornx2QnJRuYe9kkqMhxGOaMMLj/bNXs6VgpAhRFmLwVR09/A7837HfA1aN6O9OTeIhLPR21+RX9CY66WwmoIi5/sqTumt7qln29/+p0OIR7TaPCsms+zm3oiZ01cw+QI1KHDn7eVyJ6j4XcPg6S1r5fXKbly75nNvXLiFj5ogktvCp1SHaqEda8AAAATkGfH0UVLJ94oEqpOpYRxdsQCq36e1PDI41WDpJPcKamY+k0pWs2/KywiDKYOYI/MrTcoCQxBFm+MjLqCzLeTpaRiENFpLsVzukTIcUPCwAAACsBnz50RH+EVDZPSGusQ1m7R8O6wN1peQ/1jS1QCsK/+3Um/nowdX2k06ZAAAAAOAGfIGpEf3/ocg0UgjqvtSdJTqly3l/5RQ31A5IZqdimmnTC5MIJTLPa6Xa8Tb9vHN03oEBvLxevAAABDUGbJUmoQWyZTAjfx0V6DnpfBXTYnQk2bwKX2bwMtCO06TgUqyoBuPFXgEpXFMPAzWMC08ZZ56ZhxoN9TkSIiXVTokTYYuz3pKV7lAe/FQPey7N+dNtFiUjQWvY15ek+//wU0OU7OCIKAvTzIaRU2/1uvJ32Jwtf6MeG6udoHVMDrga2SELkoyovsSaeoA8+NLOPdIlZ9IkteJ1UZv0f5rSTvZTQ5nm5XNAnjE/38OYo1ggnhiEji+Cgk5Y2BrVp0F8ZOcAg3273NFXpfm3+x8FlaagyvTugb1BaDOo9dxr8EQG1u2ACk3eEl2m7Cw/3wccsdqmv1k5rlStFNyS7rxGwIGSghx6CdSt/U31jAAAAT0GfQ0UVLJ/Twd08p/VyChpS8T9IWptOSDvdLPszfbix+67WAXfjDNjm/bHE91nO3g6zlnZrZDZR0Gz1rsuN02ikA6mQRG4q5w3HL2ZCEhwAAAAkAZ9idER/hHi2UFyJjt4ot2IIOgG+Vq+0x5nM3Yla/iCdTgZJAAAANwGfZGpEf9fRGvKk48/Ke4CgTzeekpncWwo56eW/sFEV91i3WBPnz0ynsff3Tnr+fRiCryeLNz8AAAD4QZtnSahBbJlMFExv9Cvq/twdBueYn3Ee3pjczv7OR6i3WU+SGv/yONzlAdl0TMgBYUraQumePMz3mddrVgKe39/UH0dZMWofZ9Mqui4P/wvLxsUppGHrrazsvTR+f2rODWARjtdkCSvFbtfvMMzeLFsQ13oDAGG2vvjL5m95QTZnRKv1OnYK93SNtLC5O5AQU0iYLnxfpUikyvxOvC8fzzbFhAqGYgHL1GBJ3KgSybvxeNLFYx0n3LIy1+V+pUzTJ/EUuTURErNUiE759HI2TbRLFPHYcqtlrP1/TS0NNfBRxkAXAcx2r7lawk5xdWWia+VA11/z9cEAAAA2AZ+GakR/mJ002N+vm8JdEQuzLaLtUIOIHfz7FHdVElv3dJhMuXpT4IkFixsgNNJKwy7EeXqbAAAA4EGbi0nhClJlMCN/x0W1i/GN6GepRrlUQYVWJ5EJgT8Zlnuh2q0V4+3OOPM4/4/MCjdcBbKHis9PBt/yefxMxHgw7FScKctBHvFoR1EawxWUsaHfzfSnnalnY4m9Ak14jYSJiHQ5ToCi7Y4bqJeCGTsIs3t/Q3xTdD//Dx8soX29zf1N5xgAr/qPmK+dU6/y5sivZgEvOoh/r0LHDfok77DJX4jZRZR016pPj/OKEWHWFz2SRnFT4/7nXgVNFiCLYupDDD8vVEru70qqarHLzZrSjSEw3CIA0xIhUFDHlk5sAAAAMkGfqUU0TJ/TtTGaqT3w3TTw07mec9pdUO+qick9Inm7YO6odtgQD/RhYzoYrtniSB/IAAAAEwGfyHREf4TAtHqZwf6daDH8ihsAAAAwAZ/KakR/zdHy96S6V8otDkejDD6GAoE1dBa3yj/ms4XveSTIAKMwyhRIvm1XthnnAAAAfkGbz0moQWiZTAjf9BLVe/EEalTabFqMzxVTAPfnxJ/Zlk5dL1vhj03RqKWw3suOlU5fZYsZNbQTFyM5VPD5D34ws73lho1n4IFelgz9POTbop+0V7yvkVIdPx9z1ZXKiB6WJVdOUzkbyhc1ThMFffRcKHS2W7razGQSoY0BSQAAAFhBn+1FESyfgOahSeu9LaQmIfCcRMhY/HlRnhZaQPy6CAtyQ+LszE9EabEU3RQtsAcY3sdxq8ilGwwrV6BXotZB+DgqfBZjqOBne2bdeY+xHF3ipKcm5shAAAAAMgGeDHREf3p9slhDFwgFsyLLeY0fCtie4WUMzC6rVyKWDS2ciTpdi3cETt3ZfbuvDXChAAAAKwGeDmpEf4n36qp82ESoc+OmMW76ljVM0FgrkSvJn6xJsKLGiU64smxwxzEAAABzQZoTSahBbJlMCN/0EtUjX4qLA/VxpTkfRuOUpW531Ao8SwkqNXlLmRBGxN1IsecXQYmRrqrhzuWJINDZbuHe7a3EXT2PLlFrMUbyi7HcW331ihx5HV0pCo1UdCwPmKqK6tJGWYXzqyTyWsq03jrcYpi3wQAAADRBnjFFFSyfex1KrddINRpEPynlrUA1hgUxfwc/hqlijWK0BwD23AC6+DLASM32cUAM8GGBAAAANQGeUHREf4ByGO+by1cZ8ELLxtGybJtzqprjV2BYPrAXGxgEvx6Yxh6AyzUdErjXJBHKQFDAAAAAIwGeUmpEf4RorykjhEu7Dyr89cWrOYZIKdrrF5QUKCEwPBeAAAAAckGaVUmoQWyZTBRMb/vgOOEudjsI3Btv/VSW39BLCWXy94pOkzAzXgbOIzUOMjzvnTjsYbNLA9pAvZ9YoFxYYhx9h4eU9NZk38/IlYerTZE/6bkDdbgNWkejxbItyYDZBfXrFL08jsMcNfK3iihUtFthcQAAADUBnnRqRH9/6HINFiMdEuSVNpUqBczjZC+frF9/+VOzdYvwD3y/RUO7jtNMvopLiPWtfsfkhQAAAFZBmnZJ4QpSZTAjf/vgRKik2M+QxjuXx6xjPBvXPi+CYy/3jhDMFE5SsMyyTR0PdQ+GZ5DvFQ+4McBlrhdYQyRJoj/S1jsCJTCCYb8261AmPc/xNGhEYwAAAH1BmphJ4Q6JlMFNExv/++5XDhLmEjtL29ZAM9do3FyhDOoEv9H/h3RF9bsYANZinaVx61rvr6dcFSPOYKYVa0Ky3w6DZJzzBXVJRPeXQdWDKgQpImvJu97y8iZ+2lh5SMVTWZyUhJjoYcacZWmU6rIRgEt/EHai+6H6pE6egQAAADwBnrdqRH+KbhTgtKWIdHh3CgwGor5ORfdsM8dLegC9vy9X/NkeqmFpzE3INqVqG2xva8onEHyYu3G7LjMAAABIQZq8SeEPJlMCN//76z44IUjPWuSuXQG9n5T17mOUqH+IrXRyER2829XfIa9IKsVCuw/i9Uft+dhAIOotXMZgMVH6ntznsvIRAAAAQEGe2kURPJ+A6aXZxXJX+dIXeGpcopCzqg9ywRig3tyiB2DWkAg1Nn19JHMqupfHrhRdStTdinM2zM8a/E1bWFQAAAAvAZ75dER/hRLKKjdRBBNZ2S51o0KrXqTH62/XmFWxCqMc5WhSfqLQHzwmuHVfuIAAAAAbAZ77akR/f5yCyQJ9t0mSfiBytXO6KNT8CYfxAAAAgEGa/UmoQWiZTAjf+8PbyScRYg1RKyMpfbNX/CZRmm7Hq5VPl9Kuk8m7UOZdAbbrLswPQZQjheXgSbH4nLjepdplv51yrDISiOwhOpa5KSaf4Oh7414o8y9zXZHdmfzx+QS1lrmt24zuAg09D9q0y9iOg2/VIJnCofp+Dcv/U7ggAAAAcEGbHknhClJlMCN/++A44S52J3BJbw5nJLJfrjdwX6x1tDYZqj16z7INfCLPFEvnVSYK+EgFy5iJXZMiPitdWn9FpnmJqvFHhXzj4TTJ0U4uxYajaipJlltSfuzvOgecTBXwsPWX+UIqDZy0B/uWYEAAAABsQZsiSeEOiZTAjf/8In/MgnTsdVF4hM+h4LwTtR11DWQhRcJNRvQrtL/nWVC8Fu9eFlLAp8tLIZNEfkIXRbXsvYnQqfL8udnCweG1H8XL1oX20j7WR86CHE+9Ke8Jhtdy1bbfdOfKjqHqgB/RAAAAM0GfQEURPJ97HUrykeGrxIFqUtgUr7SYug/yEk3MgZ5o4qOnfx8yJ6ewAL36Oxl32rzUXAAAACEBn390RH92+75YCJktXmd/7NQWvwATKbauC51Qh/fxn8EAAABWAZ9hakR/im4U3dpeEMcKmoRT6qqm2/JijnAse/V8GPlbeJNnpdMFOlnqWvqie1+Z5e2AUboKtScIhZJGToOZ4DYrujj6Kul4w2TJ9Hh7ENnXsHCbwRkAAABWQZtmSahBaJlMCN/76z44EJY/u4nS+fFAtqU8yWqvWTE8GR+nbQ2pXN+qpaQkJCk9IKk89H1bE//13lfA/iJjEiHwWexUi5OP/Wk+Wl82rT8IeKMB/agAAABKQZ+ERREsn4DmoVg1HkW87stKnuJEPWC6LUKIZAqca58O77RVjao/vMVW8zN5+SBSMrUHBYqA3TCqOVCAQjR32PKFDGNqCbkULzsAAAAxAZ+jdER/gw4tmzW7rzcxfEkNyf/onzmAWVc3vPO+lyds9aaITrQ3fCOV9PNwrTxg0QAAACkBn6VqRH9/nILJAq0V22CwA26I7F68ZhwvvzudCOtW9UTJodNhQc7hgQAAAI9Bm6pJqEFsmUwI3/vEgWa+oh2wHUBNFWOi0kAxeW/oBNECdxapz6X4Ulak31gejD9icbk8DhRm0Hs2hpnHJ2o0E9aOoy4Z+ZKHeuQbkVZ3BAG80uP6pot9Q7nbRo/Sr6mvqnH8mE8zpvT0kWpYJYh5VOQG6NR3X0yMAeVL+tP9aFjzEgYaAOOYI13h/D/pQQAAAFlBn8hFFSyfeJmV5SQCgqT7iG5bhfyYEe/eRaEbB7n8Zt1Ka0FoANimziTzHCgJBfnDtij3/czGvfJsLwpDb7MAjV+SrguqOiQF3BNZwA/WhjksDFwtQJpmoAAAADcBn+d0RH+DDmAasguYEJlS+qlA/VHzkU9DE5VIEBd8g6r1fmm3k/DZypjK98baTmgzNHfZVmKFAAAAHwGf6WpEf3c5fYTCqVAJrrpdDQdDLHRuhKusF09cRKEAAABvQZvrSahBbJlMCN/74DjghLmmO2mUnbfWaW5B8Mb92g3kEwH8rIP/BLmQ9SIIzp0Z/9N2kUyz8oEJGs/Ap+sPATqqFFZLVvfg+PmnCmip/KLJjlyWyzFDoI36aRp8zoUQbAabU90Z2YMWlc5ECtZiAAAAVEGaDEnhClJlMCN/++5XVEO1lWujg8lVyQxzv7cDyj9Khzwt+GmZK7MW+odXd1sxR2cKlqoYmmu2Vz7BYmxOdJPgwJSdZxassfDNt9Cqk8Ou7kbHwQAAAGRBmi1J4Q6JlMCN//vgHII+SiBmB1jZ92ZKw687olAG63mvy4JF68C/BtdOHNqo9oOcgCDap8yVUyFaCvNPKBjx5Y12mHem5KjGuR1HE1vwiWnlm7LaNIt7dT5FKWs5/g4ZnkYzAAAAc0GaTknhDyZTAjf/++5XDghKblpByuukqCy+P01EbidExQWS9H7GxldpGPDpK7V+tBUBN0kUkqNhKVgVQcUsvuBYPIwMAl+a8Q4xkIhl/9reTev/yuiRaZox+SduCU+GYb7D5dve6JG2bETSCoHHzLVhoOEAAABGQZpvSeEPJlMCN//76zcgkiFnycdhsLEw0asP63rfGuqZvt04llNZeVygNUg8Jr3NY0ry9QaO7iXi3konyubaDZ6qMEZf4QAAADZBmpBJ4Q8mUwI3//vgHIJlMhOje7UYw2E7Ne9jyk9kHh3DzUwY2oalRKALKuqLHKNjTZEyfXEAAABdQZqxSeEPJlMCN//76z44RALE0DW553rg5eSpXu4lPlg3qaIlhzcDVr/CwZnQ+H1LWa2hx8yylsTU3OCdaOErX+9UvWtlkHhNfA7/4IxFk9dSDJPX2R0DfM6/cnV5AAAAWEGa0knhDyZTAjf/++A44S52Ox7QAqqAU+wlZbuVvsj7Y8xPJ/oIVLMj3KgFlKtIHmn1xsMKiZMfby2S2aO0xU7D3AaY+dFVaek1REHuf/uwXVRA/9Oom4AAAAA5QZrzSeEPJlMCN//74ESopM95J7cOXH+Lihd7kB1WsgMIaYuk4/P93ZFft876I1IMoEoIC02Nlj1AAAAAMkGbFEnhDyZTAi//+98ri6ylM95Xz+FrnToIvklvE5kKaL4Gyoqmk5OFcr+nIm3Aq8lFAAAAWUGbNUnhDyZTAi//++tE3K4Iv8XvdFgAw9nY/t8iJypWqoPs8SfKSv1NUm5Aud8Hq91WXCVfMYTS48S1kcJQTJb7gwxOICnLQAyqaV+gSJ/vuLhZm00gZbygAAAAV0GbVknhDyZTAi//++5Y+Bvd0edhu05xaJ4bHAypArfk42W21z0Hw9aIvP5GIJ7U1/Z42XP07srgXfJu1+InIWELVunkmOM0UPSk9U0pMQXyEMc8Ph6NIAAAAGJBm3pJ4Q8mUwIv//vuWLcrkUYnrfzKvY6YyeTO9Rk4qyfjY77IkI55iiBIrunpAkA7045O9WmOwt1TUqx7i+IPJ9eh/xnWjecTeoRQNm+CQoQo46ThJC/lHxBJ3Ma6WMJhgQAAAEVBn5hFETyfgEm1lLG500zAj+otoaEuWn4LIfRFhyaAA9CwtVZ6mp1AZb89j27+RZtQcexyfGBpeuTXTM9UC0Wa3f0yZmEAAAAlAZ+3dER/im4VKaciB4AJOzHjFym3oODn6PO6EfFl40Bc8IqeXQAAACUBn7lqRH93OX3LIRNDM9tasR3koAjrPHMOiuw7i9TDEgskYPw/AAAAiUGbvUmoQWiZTAifwAj3GUsedIbwtVPmC6UOtk0BDMV7Fztcurfnw+Ukl/88z4rxasXNQ9PFoO/5nE8qDt2fJK6LMxkaoX4RjOKmCD3wJkQpTcQNJfK33lRLTK724nIXgRJ5xYIwCOoOErOutudRLF4PZ/ViMvI+XwSlh7aUB0gxVLJooJL+QGvAAAAAL0Gf20URLN99VLWN00toRsi1JSCmlbiDfGAKxMwuqu+ejxaDbCFeBpddWwXK34eBAAAAOgGf/GpEf3+3wq4dDu7w/BsOHSau1FuQnVGdVhqVK3yiZQNkNPtat/J4UIGt5UWW8ZQo/jCP/r1Nr3AAAACHQZvhSahBbJlMCJ9uIMp0i44t0raDO9/qWe5SIVgIlo7RXxku4hPZq6uBGB4yOHnSn7I/dfcyQ7IQL8iwlPxXakCthkEjBcrQSK1+YKDZ1aeYZj1E6kswhNqi3qb+Xp8S04OXkru7TYYvfQY8IlY1dIAa1uyqVN5tJyc690K5tFol+v8sP3r8AAAAWUGeH0UVLJ+A5qFJ7VTSGZaxh09N4LvIJtQj3iBd9smy2aE0j4YJdcN+aL86nrecX6XJe9x9zozcD5zJ3+OWlY6VcS3j059iPKQYmFD66vdkRvd6PnTkL4nBAAAAMQGePnREf4R4tlBciUsnWu6mI8oSE9LvSXimcHaU1ew88XBcw2xze23nLBJw1I7DAmAAAABTAZ4gakR/gQeUyXbanD30/po8AnM4qaiU66bgB4fpv+VLJIx/KYTExAZQTmPLSKx1c76Lfl3wOXm3/JkWtAbgIRqsn6/a1Z7l65jvwwWZio2lQcEAAACtQZolSahBbJlMCf/mV66kDVGRg/gckK5KqDoGjG/Rc0jTpLhTmGpUH6p6+6Lg2rz4jmaD/nRZb7pbBze8wvCCpGUZHUhbf8B87j6IYVHLv2NWZ2HtpdsJpaK+aplUwQpQ2jyOwJc8FMrzBW66sb7D7eHQNNKX2TQDpjQOPBX2qpgf8m4FZujU9qlQGljaykfczBPPyrMOa6hPI30+YisxV90mdGge0/yWBpepB/EAAAA5QZ5DRRUsn3sdSq3KxWQxmQudZqrOmd/ZinGlzs6ATElOuuq0U0+75PEkS4CwMI3s7/bFIEfrMlMuAAAARQGeYnREf4ByGO+bxiuCCgeG77QV8cdQPoGcUaIvNKoixhEtWoRDrlt89Zpb+CvpUPEpQWLYG2sl7ltGjzmAbh7BeK7WEwAAADMBnmRqRH+EaK8pI0WvJO3oqC4LpNBdDqSQ/2H8yLxBYfp6xHMaIZZrOGpnEjIG9IZgYtcAAAB1QZpoSahBbJlMCX+Lwr9oP9EDpqimHn50oJzlNSRyBVnJc7Vs6Id8pb1K4EjPXOvWC39X4znnb5gA6pxzGbQIqrjG1MBa6dnTE5rY8vR7mZCBx+zMoAaGfd4sMr9WM2xq44dO8HqaOxPZ9945Sd/152ym2cFBAAAASkGehkUVLN99Ta+j2R91Ew92HLRzMWd5j8GehO+60tjrq8bx/jm0obpABIT+nkTNTJlaL+14ziVE2QZYV4VvHO4XQvOYtn92rvxBAAAAMAGep2pEf3c5fYTCskPjx10tNs/dEitfWcAR2pkb+X6knhPNb494o6ye6ByoXD+voAAAAIZBmqxJqEFsmUwL/waaFZ1EyQOlLmgdat7bD4DlnDKE5EZeA7wS8KuScmLjUp2EbKQvkspxKeGj6wCV4/o9afBe/9NGPieBtZO1xGmVIfxYKRK/87tsBmbY+CKUcMExGgKwo9oXid1IeYBEUO08HCareCjE9Ukys7EXjqpBhqzoxi7rDqSeDQAAAGtBnspFFSyfgOol4/fAqtD33XIsmEONOHHMdFPLHo8kgNBLKWb9Ym0a2pfYCaWEGCwtlD2bnmGnRxwnkF0/3ikz9RRWbpi5VR6qekgem9YtTpaFhaa4okfsIeRm5VyV92Bz5W5+vfvQGlzKQQAAAE4Bnul0RH+KbhTf9HyRxrjP7gTLkjfovmjjipEuxkhrSfiXFbTj+bdRKyxdbpt6VAaA+pPkrRqInGqSekN/aHLTqgaJTdQcGWwTKKS0O/AAAAA7AZ7rakR/hNQrERnav5stEZ3ol4KYJIAtAl5O8ChfhNtpxCXz+ANkDpac7UeKcUe4PCnSrmKSFibwLDUAAAB6QZruSahBbJlMFE//Ca2F/yEY9AmJs0u6nGYXXc/tlS+Ep6R9fuGO2/jOCP4GVYBDYfNYov80JOewYiS/s/vrHJa0Qz01zjkBZYXCzefVUhU29LngYko5T1hwFRoyeH8iTqTzThvYOCNdd/EJ3oSIlNm/FwdpeovOsTgAAAAwAZ8NakR/gQeWyVZmnrxk/Y85AYHkLo/YCh9kxyYfbQDb4iHTGTcQHbUUBaOiMgGBAAAAU0GbEknhClJlMCI/FWGmIo6t5jEJYH9imbohlas+DBV/nvXtxmYQEaER8fEGt1ChymPpiTYpdh51Cmh/YiLHIUzWJgypPkYhnUaomR/HCwbaiXipAAAAWUGfMEU0TJ94n01vVRnAY040HFEBoYfXPVz/vno2VX9xlArIIXFkMxlTVQQFmfCo6f94H5UswTk7N1qNbBjEf5ti7RycrMWYZI8EZmojies4I8dHu6teptD5AAAANQGfT3REf4MMlUlILcS/8oa2RUszwulzz8mTh84lHFOM1E/jfwpW/Z9ANgKIGdGUmnAXkT5BAAAALQGfUWpEf36Uk5AyZXqKVAUafPJRrljZB7JJE29//0psy6/KKcUFyP4LnFh/YAAAUAttb292AAAAbG12aGQAAAAAAAAAAAAAAAAAAAPoAALlBAABAAABAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACAABPNXRyYWsAAABcdGtoZAAAAAMAAAAAAAAAAAAAAAEAAAAAAALlBAAAAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAEAAAAAAUAAAAIQAAAAAACRlZHRzAAAAHGVsc3QAAAAAAAAAAQAC5QQAAAgAAAEAAAAATq1tZGlhAAAAIG1kaGQAAAAAAAAAAAAAAAAAACgAAB2kAFXEAAAAAAAtaGRscgAAAAAAAAAAdmlkZQAAAAAAAAAAAAAAAFZpZGVvSGFuZGxlcgAAAE5YbWluZgAAABR2bWhkAAAAAQAAAAAAAAAAAAAAJGRpbmYAAAAcZHJlZgAAAAAAAAABAAAADHVybCAAAAABAABOGHN0YmwAAACoc3RzZAAAAAAAAAABAAAAmGF2YzEAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAUACEAEgAAABIAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAY//8AAAAyYXZjQwFkAAr/4QAZZ2QACqzZRRP58BEAAAMAAQAAAwAUDxIllgEABmjr48siwAAAABBwYXNwAAAAAQAAAAEAAAAYc3R0cwAAAAAAAAABAAAHaQAABAAAAAAwc3RzcwAAAAAAAAAIAAAAAQAAAPsAAAH1AAAC7wAAA+kAAATjAAAF3QAABtcAAC84Y3R0cwAAAAAAAAXlAAAAAQAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAADAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAAEAAAIAAAAAAEAAAwAAAAAAQAABAAAAAAHAAAIAAAAAAEAABAAAAAAAgAABAAAAAAHAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAAGAAAIAAAAAAEAAAwAAAAAAQAABAAAAAANAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAABAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAAGAAAIAAAAAAEAAAwAAAAAAQAABAAAAAAFAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAASAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAQAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAADAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAAGAAAIAAAAAAEAAAwAAAAAAQAABAAAAAAFAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAEAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAABAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAA0AAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAwAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAABAAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAcAAAgAAAAAAQAADAAAAAABAAAEAAAAAAUAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAQAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAEAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAADAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAABgAACAAAAAABAAAMAAAAAAEAAAQAAAAABQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAARAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAYAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAQAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAwAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAADAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAAGAAAIAAAAAAEAABAAAAAAAgAABAAAAAAEAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAANAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAIAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAwAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAABgAACAAAAAABAAAQAAAAAAIAAAQAAAAABAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAADAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAYAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAQAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAwAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAADAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAABwAACAAAAAABAAAMAAAAAAEAAAQAAAAABQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAADAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAHHN0c2MAAAAAAAAAAQAAAAEAAAdpAAAAAQAAHbhzdHN6AAAAAAAAAAAAAAdpAAAGRwAAADUAAAAhAAAAJAAAAbwAAABeAAAAIgAAAE0AAAA1AAAAMAAAACQAAAAxAAAAogAAABgAAACGAAAAOQAAAC8AAAAiAAAAbQAAAHYAAAA1AAAAcwAAAD8AAABwAAAAigAAADIAAABOAAAAYgAAADQAAABqAAAAIwAAADsAAABoAAAANAAAAFQAAAA5AAAAUwAAAFcAAABkAAAAXQAAAGMAAAA7AAAARQAAAHsAAABtAAAAYQAAAGQAAABuAAAAZAAAAHkAAAApAAAARQAAAEsAAABmAAAAUQAAAFsAAABSAAAAewAAAF4AAACHAAAAKwAAAFcAAACBAAAALwAAAIAAAABQAAAAPwAAADMAAABSAAAATAAAAGUAAABEAAAAPwAAAGAAAABjAAAAPAAAADoAAABfAAAAXgAAAFsAAAB+AAAATwAAADcAAABVAAAAbQAAAEQAAAA0AAAATgAAAF0AAABsAAAAIgAAAEcAAABtAAAAKAAAASIAAABWAAAAPQAAADcAAACpAAAAOQAAADwAAAAvAAAAvAAAACoAAAB8AAAAPwAAACgAAADAAAAAVgAAAEEAAAAtAAAAOgAAACQAAACSAAAAbgAAAFQAAABLAAAAYQAAAEgAAAA+AAAAOQAAAEQAAAA4AAAAaQAAAIAAAABbAAAALQAAACwAAABmAAAARAAAAHcAAAAyAAAAOQAAACgAAABaAAAAXAAAAE4AAABfAAAAUAAAAEcAAABjAAAANgAAADMAAABkAAAAaQAAAFIAAAA2AAAATwAAAEQAAAA1AAAAMwAAAEsAAAAlAAAAkAAAAJkAAABOAAAATgAAADkAAAB9AAAATgAAAJcAAAA+AAAAOAAAADgAAABsAAAAUQAAACsAAACFAAAAYQAAADkAAAA1AAAAawAAAEkAAACIAAAAeAAAAFAAAABAAAAAOAAAAH0AAABJAAAAqQAAAGUAAABLAAAALgAAAMoAAABZAAAARAAAAD0AAACEAAAAQAAAAIMAAABAAAAAQwAAADUAAAB0AAAAUgAAADwAAAAzAAABCAAAADwAAADNAAAAQwAAAJgAAAAkAAAArgAAAEIAAAA1AAAAigAAAGEAAAA3AAAAMgAAAEIAAAByAAAAIgAAAJ8AAAA4AAAAZwAAAGIAAAA9AAAAVwAAAEgAAAAtAAAAJwAAAG4AAAB+AAAAagAAADwAAAAoAAAAVwAAAEsAAABYAAAAOgAAACwAAABsAAAAZgAAAGoAAABCAAAAWAAAAHEAAABmAAAAcAAAAE0AAAA7AAAAbAAAAFUAAABFAAAAQAAAAFoAAABuAAAANQAABO8AAABjAAAAKQAAAKIAAAArAAAARAAAAIYAAABbAAAARQAAAFYAAACyAAAAQgAAAEkAAAA4AAAAdgAAAFUAAAA8AAAAigAAAHcAAABSAAAARQAAAGYAAAAyAAAAfwAAAQYAAABqAAAAQgAAAD0AAADAAAAAOAAAAF4AAAAyAAAAZgAAACEAAAB8AAAASwAAACoAAACaAAAARgAAADUAAAAvAAAATgAAABsAAABqAAAAZwAAAFsAAABPAAAAXQAAAEoAAAA+AAAAOwAAADwAAABAAAAAaAAAAGwAAAA9AAAAOQAAAHIAAABxAAAATwAAAEsAAACgAAAAMgAAACoAAABlAAAAUQAAAEwAAABiAAAAaAAAAFgAAABUAAAAPAAAADEAAABUAAAAbwAAAFEAAAA+AAAAVQAAAEUAAAAzAAAANAAAAIsAAAAsAAAAKQAAAKEAAABQAAAATwAAAD8AAACJAAAAUAAAAKsAAABIAAAAOwAAADIAAAB0AAAAWgAAAC0AAAB1AAAAUgAAAD8AAAA1AAAAZwAAAEoAAACBAAAAewAAAFEAAABFAAAAPAAAAHoAAABOAAAArgAAAF0AAABKAAAALQAAALoAAABXAAAARAAAADsAAABrAAAAPAAAAMcAAABCAAAAPwAAADYAAABoAAAAUQAAAD8AAAA0AAABBQAAAEIAAAC0AAAAPgAAAIMAAAAdAAAApQAAAD8AAAAkAAAAiAAAAFgAAABGAAAAMQAAADwAAABQAAAAHAAAAF4AAAA4AAAAWQAAAGQAAAA5AAAAVwAAADwAAAAzAAAAOgAAAKAAAABXAAAAigAAADcAAAAoAAAASAAAAEwAAABAAAAAOgAAACwAAABxAAAAdwAAAF4AAABLAAAAaAAAAGcAAABmAAAAdgAAAEoAAAA+AAAAYAAAAFMAAABFAAAAQwAAAFwAAABWAAAAXAAAAE4AAAArAAAALAAAAHgAAAAxAAAASgAAAIoAAABaAAAARAAAAFYAAADDAAAARgAAAEcAAAA4AAAAegAAAFIAAAA/AAAAjQAAAGoAAABUAAAAQgAAAG0AAAAwAAAAigAAAUYAAABWAAAAPQAAAD0AAACvAAAAMwAAAGAAAAA0AAAAYQAAACEAAAByAAAAQQAAACcAAADKAAAAWwAAADoAAAAvAAAARwAAABsAAACSAAAAagAAAGAAAABeAAAAVwAAAE0AAAA3AAAALAAAAGIAAABEAAAAcAAAAGQAAAAxAAAANQAAAIgAAABcAAAARgAAAEYAAACfAAAAMQAAADIAAABRAAAAUgAAAEsAAABiAAAAaQAAAF8AAATrAAAAUgAAADIAAACIAAAAjQAAAFUAAAAzAAAAVAAAAEUAAAA0AAAANQAAAKwAAAAtAAAAKgAAAJQAAABRAAAASgAAADkAAACGAAAASQAAAJwAAABDAAAAOQAAADcAAACMAAAAYgAAAC8AAAB9AAAAeAAAAEQAAAAyAAAAbAAAAEsAAABmAAAAeAAAAE0AAABDAAAAOQAAAHUAAABEAAAAsgAAAF8AAABKAAAALwAAAMEAAABVAAAARgAAADsAAABzAAAAOgAAAJsAAAA7AAAAPwAAADQAAAB0AAAAUQAAADsAAAA1AAAA8QAAADoAAADCAAAAPwAAAHcAAAA6AAAAoQAAAFEAAAA2AAAAiQAAAGAAAAA3AAAAKwAAAEAAAABpAAAAGwAAAH0AAAA5AAAAYAAAAGsAAAA4AAAAPAAAAEQAAAAxAAAAIQAAAHcAAABeAAAAaAAAAEAAAAAsAAAASAAAAGUAAABOAAAAKwAAACsAAACGAAAATAAAACgAAAApAAAAXgAAAE8AAABiAAAAawAAAFsAAAA7AAAAeQAAAFkAAABKAAAAQAAAAFYAAABeAAAAWgAAAC4AAAA4AAAAoAAAAC0AAABbAAAASQAAAH0AAABXAAAAPgAAAFYAAACvAAAAQgAAAEkAAAA3AAAAfQAAAFYAAAA3AAAAoAAAAE8AAABUAAAAQgAAAHUAAAAnAAAAgAAAAFIAAAAzAAAANwAAAR4AAABTAAAAbwAAAC0AAAA+AAAAJAAAAK0AAABBAAAAKgAAACwAAACfAAAAWAAAAD8AAAAsAAAAXwAAACEAAACOAAAAawAAAGcAAABQAAAAagAAAD0AAAAtAAAARgAAAEAAAABNAAAAawAAAGQAAABOAAAAPwAAAIcAAABOAAAAQgAAAEMAAABdAAAAKwAAAHAAAABjAAAAXAAAAE0AAABtAAAATgAAAE0AAABkAAAAMQAAADAAAABaAAAAggAAAFQAAAA8AAAAVAAAAEYAAAA0AAAANgAAAI8AAAAuAAAALAAAAJoAAABLAAAASQAAADsAAACRAAAARwAAAM0AAABDAAAAOQAAADcAAABpAAAAXQAAACoAAAB7AAAAZwAAADUAAAAuAAAAggAAADAAAACHAAAAhwAAAGAAAABIAAAAQAAAAGcAAAA3AAAAtAAAAGYAAABZAAAALAAAALEAAABVAAAARQAAADsAAABoAAAAPAAAAKgAAAA9AAAAQAAAADoAAABrAAAATwAAADsAAAAwAAAA+gAAAEAAAADRAAAAPQAAAHwAAAAoAAAAoAAAAEcAAAAyAAAAWgAAAD0AAAAvAAAE+QAAAEMAAAA2AAAAYgAAAJ8AAAA2AAAAYAAAAHYAAAA9AAAAQAAAAD0AAAAtAAAAJQAAAIAAAABkAAAAcAAAADMAAAAoAAAASgAAAGAAAABJAAAAMgAAADIAAAB/AAAAaQAAAGYAAABQAAAAXQAAAF0AAABqAAAAiAAAAEsAAAA5AAAAaQAAAFEAAABKAAAAQAAAAGAAAABYAAAAXwAAAEcAAAAoAAAAKgAAAIgAAAAiAAAAOwAAAIIAAABVAAAAPgAAAFsAAADQAAAARwAAAEYAAAA6AAAAhQAAAFEAAAA2AAAAjAAAAEsAAABTAAAARgAAAGsAAAAlAAAAgwAAAOgAAABNAAAARwAAAEkAAAC0AAAAPAAAAG0AAAA1AAAAXwAAACUAAACQAAAAOwAAAD0AAAAlAAAAlAAAADcAAACJAAAAJQAAADMAAAAeAAAAYQAAAGsAAABLAAAAVAAAAEkAAAA6AAAALgAAADoAAABHAAAAcwAAAJgAAABhAAAAMgAAADIAAABgAAAAOgAAAK8AAABNAAAARgAAAC0AAABlAAAAVAAAAFQAAAB2AAAAZQAAAF4AAABWAAAAPAAAADgAAABYAAAAggAAAFIAAAA7AAAAVwAAAEsAAAAzAAAANAAAAEsAAAAnAAAAlwAAAJMAAABWAAAATQAAADoAAACIAAAASwAAAJ0AAAA9AAAARAAAADUAAABkAAAATwAAACwAAACGAAAAZgAAAEsAAAAsAAAAbgAAADQAAABdAAAAdwAAAFYAAABQAAAAPAAAAFoAAABKAAAAqwAAAFoAAABJAAAALQAAALEAAABVAAAAQwAAADoAAABmAAAAPwAAAKoAAABAAAAAQwAAADIAAAB8AAAASwAAAD4AAAAvAAABTQAAAD8AAADGAAAARAAAAFoAAAAdAAAAugAAADYAAAAnAAAAeAAAAF4AAAA6AAAALwAAAGAAAABOAAAAIgAAAHYAAAA0AAAAawAAAGYAAAA0AAAAWgAAAEQAAAAvAAAAJAAAAHkAAABZAAAAYQAAADUAAAAwAAAATwAAAGQAAABEAAAAPgAAAEIAAAByAAAAbwAAAG0AAAA4AAAAZgAAAGwAAABgAAAAXwAAAE0AAAA8AAAAUQAAAFkAAABaAAAAQAAAAFgAAABcAAAAVwAAAC8AAAA1AAAAfQAAAC8AAABnAAAASwAAAHEAAABUAAAAOwAAAFUAAAC2AAAASgAAAEgAAAA4AAAAfwAAAFcAAAA0AAAAlwAAAGkAAABJAAAAPgAAAHkAAAAnAAAAeAAAAFAAAAAyAAAANwAAAIoAAAB+AAAAOAAABWYAAABIAAAAIAAAAPUAAABhAAAANwAAACQAAADrAAAATwAAADkAAAAtAAAAtQAAACkAAAAfAAAAjAAAAHMAAAAlAAAAcQAAAEMAAAA/AAAAMAAAAEYAAABAAAAAXQAAAJAAAAA7AAAANgAAAHwAAABVAAAASQAAAD8AAAB+AAAAOQAAACoAAABoAAAAVAAAAFYAAABrAAAAUAAAAE0AAABiAAAAOQAAAB4AAABhAAAAhQAAAEoAAAA2AAAAdQAAAEQAAAA4AAAAMwAAAJQAAAA2AAAAKAAAAHgAAABLAAAAOQAAAJoAAABjAAAASwAAADoAAADBAAAAQgAAAC4AAAA5AAAAbwAAACwAAACTAAAAfgAAAEYAAAA3AAAAtQAAAD4AAAAyAAAAQwAAAH8AAABDAAAANgAAAFAAAACTAAAAagAAAFYAAAAzAAAAzwAAAF4AAAA+AAAAOQAAAJcAAABMAAAASgAAAD0AAACSAAAARQAAADIAAABGAAAAcgAAAEcAAAA1AAAARQAAAREAAABWAAAAXgAAACUAAAEAAAAAPQAAACsAAAAgAAAAmwAAAGEAAAA7AAAALQAAAIIAAAA6AAAAHgAAADIAAABxAAAAVwAAAG4AAAA9AAAAdAAAAEIAAAAuAAAAIQAAAFwAAACVAAAAbQAAADEAAAAmAAAASgAAAGwAAABGAAAAOAAAADEAAACTAAAARwAAADoAAAAkAAAAcwAAAFYAAABnAAAAZwAAAE4AAABIAAAAawAAAF8AAABDAAAAPQAAAF4AAABmAAAAXwAAAEcAAAAmAAAAIwAAAJYAAAAuAAAANwAAAI4AAABeAAAAPAAAAFYAAACzAAAAPwAAAEoAAAA3AAAAhAAAAFUAAAAxAAAAqQAAAGgAAABTAAAAPwAAAJMAAAAnAAAAqQAAAFEAAAA8AAAANwAAAQwAAABNAAAAbwAAACcAAABHAAAALwAAALQAAABHAAAALQAAACQAAACoAAAAWQAAADsAAAAzAAAAdgAAAB8AAABxAAAAdAAAAEgAAAArAAAAeQAAAF0AAABPAAAALAAAAFYAAAA+AAAAhAAAAHQAAAA6AAAAOwAAAHQAAABZAAAASQAAAEYAAACZAAAAOAAAACgAAABfAAAAUwAAAFkAAAB+AAAAWwAAAFUAAABpAAAAQgAAACEAAABeAAAAdAAAAFIAAAA0AAAAbAAAAEsAAAA7AAAANAAAAH0AAAAxAAAAKAAAAHoAAABeAAAAMwAAAJ4AAABmAAAASQAAADoAAAChAAAARQAAACcAAABBAAAAUQAAACsAAABSAAAATwAAAC4AAAWcAAAAiQAAAC8AAACaAAAAXQAAAEoAAAAzAAAAYwAAAFMAAABMAAAAUgAAAKMAAAA/AAAALAAAALAAAABVAAAAOQAAADkAAAByAAAARwAAAD4AAAA7AAAArQAAAFMAAAAzAAAANgAAAQkAAABSAAAAKQAAADsAAADTAAAAOwAAAN4AAAA4AAAAHgAAAC0AAAB9AAAAXAAAADsAAAAuAAAAcwAAADUAAAA0AAAAIQAAAIYAAAAyAAAAVwAAAG8AAABFAAAAXgAAAEAAAAAzAAAALAAAAGQAAACCAAAAggAAADAAAAAjAAAASwAAAGsAAAA/AAAAOwAAAL4AAABWAAAAMQAAADUAAABHAAAAYwAAAFgAAABoAAAAYwAAAEsAAABIAAAAXwAAAF4AAAA/AAAAOAAAAGAAAABUAAAAYgAAAFAAAAAqAAAAKwAAAHUAAAAwAAAAOgAAAI0AAABbAAAAOQAAAFYAAAC+AAAAQQAAAEYAAAA2AAAAgQAAAEwAAAA0AAAAlgAAAHAAAABQAAAAPQAAAHgAAAAvAAAAmwAAAF4AAAA9AAAAMQAAAUUAAABmAAAAeAAAAC0AAAB6AAAAMAAAABwAAAAuAAAA4AAAADUAAAAqAAAARgAAAEkAAABIAAAAJgAAAB0AAACVAAAAbgAAAGwAAAAtAAAAVgAAAD4AAABHAAAALQAAAFkAAAA5AAAAYgAAAIEAAAAyAAAAMQAAAH8AAABVAAAATQAAAEYAAAB7AAAARAAAACYAAABbAAAATQAAAFcAAABsAAAAVgAAAEoAAABsAAAAOgAAACYAAACAAAAAfQAAAE4AAAA/AAAAcQAAAEYAAAA6AAAAMwAAAKcAAAAyAAAAKgAAAH0AAABUAAAANAAAAJwAAABpAAAARQAAADsAAAC9AAAAQQAAAC4AAABGAAAAbQAAACgAAAB6AAAAYgAAAEYAAAAyAAAAgAAAADsAAAAxAAAAPgAAAH0AAABAAAAANQAAAE4AAACWAAAAZAAAAFMAAAAxAAAAxQAAAGEAAAA/AAAANAAAAH0AAABJAAAATAAAAEEAAACuAAAAQAAAADMAAABIAAAAbgAAAEUAAAAyAAAARQAAAPsAAABSAAAAXAAAACYAAADkAAAARwAAADAAAAAgAAAAkQAAAGEAAAA6AAAALQAAAIYAAAA7AAAAHQAAADAAAAB2AAAAUgAAAGkAAAA9AAAAYgAAAEAAAAAuAAAAHgAAAGUAAAB/AAAAcQAAAC0AAAAiAAAAWAAAAHUAAABFAAAAOAAAADIAAACAAAAAPgAAADUAAAAqAAAAZAAAAFkAAABfAAAFYwAAAFIAAAA3AAAAlQAAAGEAAABCAAAANwAAAFkAAABaAAAAdAAAAEIAAAAoAAAAJQAAAJEAAAAwAAAANwAAAJEAAABdAAAAPAAAAFUAAAC3AAAARgAAAEgAAAA2AAAAegAAAE0AAAA2AAAAnAAAAG0AAABTAAAAOwAAAIkAAAA7AAAAhwAAAFEAAAA9AAAAOwAAAVMAAABoAAAAfAAAACsAAACCAAAAMgAAABwAAAAyAAAA2AAAAGUAAAAiAAAAOwAAAHQAAABOAAAAKAAAACIAAACYAAAAewAAAEYAAAApAAAAYwAAAFkAAABBAAAALwAAAFAAAABHAAAAcAAAAHoAAAAvAAAAMQAAAHIAAABfAAAASQAAAEcAAACZAAAANgAAAB4AAABiAAAAUQAAAFcAAAB0AAAAXQAAAFUAAABTAAAAPgAAACAAAABhAAAAgAAAAEsAAAA0AAAAZwAAAEIAAAA4AAAAMAAAAIUAAAAuAAAAIgAAAHgAAABWAAAANAAAAJsAAABrAAAARwAAAD4AAADBAAAAPgAAACwAAABGAAAAbQAAACUAAACDAAAAWQAAAEAAAAAyAAAAvQAAAEsAAABLAAAARwAAAHMAAAA7AAAANQAAAEoAAACXAAAAcwAAAFEAAAAxAAAAvQAAAGkAAAA+AAAAOQAAAI4AAABKAAAASQAAAD0AAACTAAAARQAAADUAAABGAAAAbAAAAEYAAAA2AAAARgAAASEAAABZAAAAZAAAACcAAADjAAAAPwAAACwAAAAgAAAAjgAAAF8AAAA6AAAALwAAAE4AAAAfAAAAhwAAADIAAAAyAAAAVQAAAHMAAAA7AAAATgAAADkAAAAuAAAAJQAAAGMAAACPAAAAdAAAADQAAAAhAAAATQAAAJwAAABLAAAAOQAAADEAAACiAAAARQAAADgAAAAlAAAAcQAAAGcAAABmAAAAgQAAAEoAAAA8AAAAYgAAAGoAAABCAAAAPgAAAGIAAABgAAAAZQAAAEYAAAAqAAAAKQAAAHYAAAA0AAAAPQAAAIgAAABbAAAAOQAAAFUAAADRAAAAPwAAAEQAAAA6AAAAhwAAAE8AAAA1AAAAoQAAAGwAAABSAAAAQwAAAJAAAAAmAAAAuQAAAE8AAAA9AAAAMwAAAOkAAABPAAAAbgAAACwAAACbAAAANwAAACIAAAAyAAAAygAAAD0AAAAiAAAAOgAAAEgAAAA0AAAAKQAAACUAAACDAAAAbAAAAE8AAAAoAAAAYAAAAFEAAABAAAAALQAAAGIAAABEAAAAcAAAAHMAAAA+AAAAOQAAAFQAAABNAAAAPAAABYMAAACNAAAAKwAAAGsAAABqAAAASgAAAFQAAABqAAAAWAAAAGAAAABqAAAAOwAAADcAAABUAAAAcQAAAFUAAAA9AAAAaAAAAEkAAAA6AAAAMwAAAJgAAAA5AAAAIwAAAHwAAABVAAAANAAAAJkAAABlAAAATAAAADwAAAC4AAAAQAAAACoAAABIAAAAawAAACgAAACHAAAAXgAAAD4AAAAxAAAAbwAAADkAAAB7AAAARwAAAD0AAAA0AAAAgwAAAF0AAABOAAAAUwAAAMoAAAA9AAAALAAAANMAAABVAAAAQQAAADoAAABkAAAAUQAAAEQAAAA/AAAApAAAAFIAAAAvAAAAPAAAAREAAABTAAAAKAAAADsAAAD8AAAAOgAAAOQAAAA2AAAAFwAAADQAAACCAAAAXAAAADYAAAAvAAAAdwAAADgAAAA5AAAAJwAAAHYAAAA5AAAAWgAAAIEAAABAAAAATAAAAEQAAAAzAAAAHwAAAIQAAAB0AAAAcAAAADcAAAAlAAAAWgAAAFoAAABOAAAANQAAAC0AAACTAAAAXQAAADsAAAAjAAAAcwAAAFgAAABoAAAAdwAAAEoAAAA6AAAAYQAAAFwAAAA9AAAANgAAAF0AAABbAAAAZgAAAEkAAAApAAAAKQAAAI0AAAAzAAAAPgAAAIsAAABdAAAANQAAAFcAAACxAAAAPQAAAEkAAAA3AAAAeQAAAE4AAAA0AAAAigAAAG8AAABSAAAAPwAAAH4AAAA0AAAAVwAAAF0AAAA5AAAAMQAAABRzdGNvAAAAAAAAAAEAAAAwAAAAYnVkdGEAAABabWV0YQAAAAAAAAAhaGRscgAAAAAAAAAAbWRpcmFwcGwAAAAAAAAAAAAAAAAtaWxzdAAAACWpdG9vAAAAHWRhdGEAAAABAAAAAExhdmY1Ny44My4xMDA=","ok":true,"headers":[["content-type","video/mp4"]],"status":200,"status_text":""}},"base_uri":"https://localhost:8080/","height":501}},"cell_type":"code","source":["play_video('pong_pretrained/0.avi')"],"execution_count":0,"outputs":[{"output_type":"execute_result","data":{"text/html":["\n","    <video width=\"640\" height=\"480\" controls>\n","      <source src=\"/nbextensions/vid.mp4\" type=\"video/mp4\">\n","    </video>\n","  "],"text/plain":["<IPython.core.display.HTML object>"]},"metadata":{"tags":[]},"execution_count":24}]},{"metadata":{"id":"U-SyGcZBCmPn","colab_type":"text"},"cell_type":"markdown","source":["# Train your policy (model-free training)\n","Training model-free on Pong (it takes a few hours):"]},{"metadata":{"id":"WIQazd5aCocc","colab_type":"code","outputId":"0a440c18-affc-4b2a-d6e1-c3cda84465bc","executionInfo":{"status":"ok","timestamp":1553254256733,"user_tz":-60,"elapsed":19957,"user":{"displayName":"Piotr Kozakowski","photoUrl":"","userId":"01014928596539690143"}},"colab":{"base_uri":"https://localhost:8080/","height":1516}},"cell_type":"code","source":["!python -m tensor2tensor.rl.trainer_model_free \\\n","  --hparams_set=rlmf_base \\\n","  --hparams=game=pong \\\n","  --output_dir=mf_pong"],"execution_count":0,"outputs":[{"output_type":"stream","text":["\n","WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n","For more information, please see:\n","  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n","  * https://github.com/tensorflow/addons\n","If you depend on functionality not listed there, please file an issue.\n","\n","2019-03-22 11:30:42.987149: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 2300000000 Hz\n","2019-03-22 11:30:42.987392: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x30323c0 executing computations on platform Host. Devices:\n","2019-03-22 11:30:42.987491: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): <undefined>, <undefined>\n","2019-03-22 11:30:43.082876: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:998] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n","2019-03-22 11:30:43.083442: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x3032100 executing computations on platform CUDA. Devices:\n","2019-03-22 11:30:43.083493: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): Tesla K80, Compute Capability 3.7\n","2019-03-22 11:30:43.083843: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1433] Found device 0 with properties: \n","name: Tesla K80 major: 3 minor: 7 memoryClockRate(GHz): 0.8235\n","pciBusID: 0000:00:04.0\n","totalMemory: 11.17GiB freeMemory: 11.10GiB\n","2019-03-22 11:30:43.083879: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 11:30:43.475526: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 11:30:43.475601: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 11:30:43.475629: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 11:30:43.476026: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:42] Overriding allow_growth setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.\n","2019-03-22 11:30:43.476131: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Colocations handled automatically by placer.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/rl/envs/py_func_batch_env.py:122: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","tf.py_func is deprecated in TF V2. Instead, use\n","    tf.py_function, which takes a python function which manipulates tf eager\n","    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to\n","    an ndarray (just call tensor.numpy()) but having access to eager tensors\n","    means `tf.py_function`s can use accelerators such as GPUs as well as\n","    being differentiable using a gradient tape.\n","    \n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/utils/t2t_model.py:1358: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.cast instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/function.py:1007: calling Graph.create_op (from tensorflow.python.framework.ops) with compute_shapes is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Shapes are always computed; don't use the compute_shapes as it has no effect.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:277: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.cast instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:598: conv2d (from tensorflow.python.layers.convolutional) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.conv2d instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:602: flatten (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.flatten instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:603: dropout (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.dropout instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:604: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.dense instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:2887: multinomial (from tensorflow.python.ops.random_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.random.categorical instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/rl/ppo_learner.py:479: Print (from tensorflow.python.ops.logging_ops) is deprecated and will be removed after 2018-08-20.\n","Instructions for updating:\n","Use tf.print instead of tf.Print. Note that tf.print returns a no-output operator that directly prints the output. Outside of defuns or eager mode, this operator will not be executed unless it is directly specified in session.run or used as a control dependency for other operators. This is only a concern in graph mode. Below is an example of how to ensure tf.print executes in graph mode:\n","```python\n","    sess = tf.Session()\n","    with sess.as_default():\n","        tensor = tf.range(10)\n","        print_op = tf.print(tensor)\n","        with tf.control_dependencies([print_op]):\n","          out = tf.add(tensor, tensor)\n","        sess.run(out)\n","    ```\n","Additionally, to use tf.print in python 2.7, users must make sure to import\n","the following:\n","\n","  `from __future__ import print_function`\n","\n","2019-03-22 11:30:49.903512: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 11:30:49.903591: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 11:30:49.903620: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 11:30:49.903639: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 11:30:49.903898: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/saver.py:1266: checkpoint_exists (from tensorflow.python.training.checkpoint_management) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use standard file APIs to check for files with this prefix.\n","2019-03-22 11:30:51.335217: I tensorflow/stream_executor/dso_loader.cc:152] successfully opened CUDA library libcublas.so.10.0 locally\n","mean_score: [0][0][0]\n","^C\n"],"name":"stdout"}]},{"metadata":{"id":"FbSjwVAtCvLY","colab_type":"text"},"cell_type":"markdown","source":["Hyperparameter sets are defined in `tensor2tensor/models/research/rl.py`. You can override them using the hparams flag, e.g.\n","\n","```\n","--hparams=game=kung_fu_master,frame_stack_size=5\n","```\n","\n","As in model-based training, the periodic evaluation runs with timestep limit of 1000. To do full evaluation after training, run:"]},{"metadata":{"id":"jppi4FE5C2nB","colab_type":"code","outputId":"a10afb7c-edd6-4a93-eee4-e3876977e825","executionInfo":{"status":"ok","timestamp":1553254412202,"user_tz":-60,"elapsed":15104,"user":{"displayName":"Piotr Kozakowski","photoUrl":"","userId":"01014928596539690143"}},"colab":{"base_uri":"https://localhost:8080/","height":4083}},"cell_type":"code","source":["!python -m tensor2tensor.rl.evaluator \\\n","  --loop_hparams_set=rlmf_tiny \\\n","  --hparams=game=pong \\\n","  --policy_dir=mf_pong \\\n","  --debug_video_path=mf_pong \\\n","  --num_debug_videos=4 \\\n","  --eval_metrics_dir=mf_pong/full_eval_metrics"],"execution_count":0,"outputs":[{"output_type":"stream","text":["\n","WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n","For more information, please see:\n","  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n","  * https://github.com/tensorflow/addons\n","If you depend on functionality not listed there, please file an issue.\n","\n","INFO:tensorflow:Overriding hparams in rlmf_tiny with game=pong,eval_max_num_noops=0,eval_sampling_temps=[0.5]\n","INFO:tensorflow:Evaluating metric mean_reward/eval/sampling_temp_0.5_max_noops_0_unclipped\n","2019-03-22 11:33:23.214052: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 2300000000 Hz\n","2019-03-22 11:33:23.214294: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x2d07020 executing computations on platform Host. Devices:\n","2019-03-22 11:33:23.214335: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): <undefined>, <undefined>\n","2019-03-22 11:33:23.309948: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:998] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n","2019-03-22 11:33:23.310546: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x2d067e0 executing computations on platform CUDA. Devices:\n","2019-03-22 11:33:23.310585: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): Tesla K80, Compute Capability 3.7\n","2019-03-22 11:33:23.310991: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1433] Found device 0 with properties: \n","name: Tesla K80 major: 3 minor: 7 memoryClockRate(GHz): 0.8235\n","pciBusID: 0000:00:04.0\n","totalMemory: 11.17GiB freeMemory: 11.10GiB\n","2019-03-22 11:33:23.311027: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 11:33:23.707039: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 11:33:23.707114: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 11:33:23.707139: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 11:33:23.707459: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:42] Overriding allow_growth setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.\n","2019-03-22 11:33:23.707523: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","INFO:tensorflow:Using DummyPolicyProblem for the policy.\n","INFO:tensorflow:Setting T2TModel mode to 'train'\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Colocations handled automatically by placer.\n","INFO:tensorflow:Using variable initializer: orthogonal\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/utils/t2t_model.py:1358: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.cast instead.\n","INFO:tensorflow:Transforming feature 'input_action' with symbol_modality_6_64.bottom\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/function.py:1007: calling Graph.create_op (from tensorflow.python.framework.ops) with compute_shapes is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Shapes are always computed; don't use the compute_shapes as it has no effect.\n","INFO:tensorflow:Transforming feature 'input_reward' with symbol_modality_3_64.bottom\n","INFO:tensorflow:Transforming feature 'inputs' with video_modality.bottom\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_video.py:495: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","tf.py_func is deprecated in TF V2. Instead, use\n","    tf.py_function, which takes a python function which manipulates tf eager\n","    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to\n","    an ndarray (just call tensor.numpy()) but having access to eager tensors\n","    means `tf.py_function`s can use accelerators such as GPUs as well as\n","    being differentiable using a gradient tape.\n","    \n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:277: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.cast instead.\n","INFO:tensorflow:Transforming feature 'target_action' with symbol_modality_6_64.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_policy' with identity_modality.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_reward' with symbol_modality_3_64.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_value' with identity_modality.targets_bottom\n","INFO:tensorflow:Transforming feature 'targets' with video_modality.targets_bottom\n","INFO:tensorflow:Building model body\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:598: conv2d (from tensorflow.python.layers.convolutional) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.conv2d instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:602: flatten (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.flatten instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:603: dropout (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.dropout instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:604: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.dense instead.\n","INFO:tensorflow:Transforming body output with identity_modality.top\n","INFO:tensorflow:Transforming body output with identity_modality.top\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:2887: multinomial (from tensorflow.python.ops.random_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.random.categorical instead.\n","2019-03-22 11:33:24.564271: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 11:33:24.564350: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 11:33:24.564376: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 11:33:24.564410: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 11:33:24.564687: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","INFO:tensorflow:Restoring checkpoint mf_pong/model.ckpt-9\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/saver.py:1266: checkpoint_exists (from tensorflow.python.training.checkpoint_management) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use standard file APIs to check for files with this prefix.\n","INFO:tensorflow:Restoring parameters from mf_pong/model.ckpt-9\n","2019-03-22 11:33:24.985295: I tensorflow/stream_executor/dso_loader.cc:152] successfully opened CUDA library libcublas.so.10.0 locally\n","INFO:tensorflow:Step 5, mean_score: 0.000000\n","INFO:tensorflow:Step 10, mean_score: 0.000000\n","INFO:tensorflow:Step 15, mean_score: 0.000000\n","INFO:tensorflow:Step 20, mean_score: 0.000000\n","INFO:tensorflow:Step 25, mean_score: 0.000000\n","INFO:tensorflow:Step 30, mean_score: 0.000000\n","INFO:tensorflow:Step 35, mean_score: 0.000000\n","INFO:tensorflow:Step 40, mean_score: 0.000000\n","INFO:tensorflow:Step 45, mean_score: 0.000000\n","INFO:tensorflow:Step 50, mean_score: 0.000000\n","INFO:tensorflow:Step 55, mean_score: 0.000000\n","INFO:tensorflow:Step 60, mean_score: 0.000000\n","INFO:tensorflow:Step 65, mean_score: -1.000000\n","INFO:tensorflow:Step 70, mean_score: -1.000000\n","INFO:tensorflow:Step 75, mean_score: -1.000000\n","INFO:tensorflow:Step 80, mean_score: -1.000000\n","INFO:tensorflow:Step 85, mean_score: -1.000000\n","INFO:tensorflow:Step 90, mean_score: -1.000000\n","INFO:tensorflow:Step 95, mean_score: -1.000000\n","INFO:tensorflow:Step 100, mean_score: -2.000000\n","INFO:tensorflow:Step 105, mean_score: -2.000000\n","INFO:tensorflow:Step 110, mean_score: -2.000000\n","INFO:tensorflow:Step 115, mean_score: -2.000000\n","INFO:tensorflow:Step 120, mean_score: -2.000000\n","INFO:tensorflow:Step 125, mean_score: -2.000000\n","INFO:tensorflow:Step 130, mean_score: -2.000000\n","INFO:tensorflow:Step 135, mean_score: -3.000000\n","INFO:tensorflow:Step 140, mean_score: -3.000000\n","INFO:tensorflow:Step 145, mean_score: -3.000000\n","INFO:tensorflow:Step 150, mean_score: -3.000000\n","INFO:tensorflow:Step 155, mean_score: -3.000000\n","INFO:tensorflow:Step 160, mean_score: -3.000000\n","INFO:tensorflow:Step 165, mean_score: -3.000000\n","INFO:tensorflow:Step 170, mean_score: -4.000000\n","INFO:tensorflow:Step 175, mean_score: -4.000000\n","INFO:tensorflow:Step 180, mean_score: -4.000000\n","INFO:tensorflow:Step 185, mean_score: -4.000000\n","INFO:tensorflow:Step 190, mean_score: -4.000000\n","INFO:tensorflow:Step 195, mean_score: -4.000000\n","INFO:tensorflow:Step 200, mean_score: -4.000000\n","INFO:tensorflow:Step 205, mean_score: -5.000000\n","INFO:tensorflow:Step 210, mean_score: -5.000000\n","INFO:tensorflow:Step 215, mean_score: -5.000000\n","INFO:tensorflow:Step 220, mean_score: -5.000000\n","INFO:tensorflow:Step 225, mean_score: -5.000000\n","INFO:tensorflow:Step 230, mean_score: -5.000000\n","INFO:tensorflow:Step 235, mean_score: -5.000000\n","INFO:tensorflow:Step 240, mean_score: -6.000000\n","INFO:tensorflow:Step 245, mean_score: -6.000000\n","INFO:tensorflow:Step 250, mean_score: -6.000000\n","INFO:tensorflow:Step 255, mean_score: -6.000000\n","INFO:tensorflow:Step 260, mean_score: -6.000000\n","INFO:tensorflow:Step 265, mean_score: -6.000000\n","INFO:tensorflow:Step 270, mean_score: -6.000000\n","INFO:tensorflow:Step 275, mean_score: -7.000000\n","INFO:tensorflow:Step 280, mean_score: -7.000000\n","INFO:tensorflow:Step 285, mean_score: -7.000000\n","INFO:tensorflow:Step 290, mean_score: -7.000000\n","INFO:tensorflow:Step 295, mean_score: -7.000000\n","INFO:tensorflow:Step 300, mean_score: -7.000000\n","INFO:tensorflow:Step 305, mean_score: -7.000000\n","INFO:tensorflow:Step 310, mean_score: -8.000000\n","INFO:tensorflow:Step 315, mean_score: -8.000000\n","INFO:tensorflow:Step 320, mean_score: -8.000000\n","INFO:tensorflow:Step 325, mean_score: -8.000000\n","INFO:tensorflow:Step 330, mean_score: -8.000000\n","INFO:tensorflow:Step 335, mean_score: -8.000000\n","INFO:tensorflow:Step 340, mean_score: -8.000000\n","INFO:tensorflow:Step 345, mean_score: -9.000000\n","INFO:tensorflow:Step 350, mean_score: -9.000000\n","INFO:tensorflow:Step 355, mean_score: -9.000000\n","INFO:tensorflow:Step 360, mean_score: -9.000000\n","INFO:tensorflow:Step 365, mean_score: -9.000000\n","INFO:tensorflow:Step 370, mean_score: -9.000000\n","INFO:tensorflow:Step 375, mean_score: -9.000000\n","INFO:tensorflow:Step 380, mean_score: -10.000000\n","INFO:tensorflow:Step 385, mean_score: -10.000000\n","INFO:tensorflow:Step 390, mean_score: -10.000000\n","INFO:tensorflow:Step 395, mean_score: -10.000000\n","INFO:tensorflow:Step 400, mean_score: -10.000000\n","INFO:tensorflow:Step 405, mean_score: -10.000000\n","INFO:tensorflow:Step 410, mean_score: -10.000000\n","INFO:tensorflow:Step 415, mean_score: -11.000000\n","INFO:tensorflow:Step 420, mean_score: -11.000000\n","INFO:tensorflow:Step 425, mean_score: -11.000000\n","INFO:tensorflow:Step 430, mean_score: -11.000000\n","INFO:tensorflow:Step 435, mean_score: -11.000000\n","INFO:tensorflow:Step 440, mean_score: -11.000000\n","INFO:tensorflow:Step 445, mean_score: -11.000000\n","INFO:tensorflow:Step 450, mean_score: -12.000000\n","INFO:tensorflow:Step 455, mean_score: -12.000000\n","INFO:tensorflow:Step 460, mean_score: -12.000000\n","INFO:tensorflow:Step 465, mean_score: -12.000000\n","INFO:tensorflow:Step 470, mean_score: -12.000000\n","INFO:tensorflow:Step 475, mean_score: -12.000000\n","INFO:tensorflow:Step 480, mean_score: -12.000000\n","INFO:tensorflow:Step 485, mean_score: -13.000000\n","INFO:tensorflow:Step 490, mean_score: -13.000000\n","INFO:tensorflow:Step 495, mean_score: -13.000000\n","INFO:tensorflow:Step 500, mean_score: -13.000000\n","INFO:tensorflow:Step 505, mean_score: -13.000000\n","INFO:tensorflow:Step 510, mean_score: -13.000000\n","INFO:tensorflow:Step 515, mean_score: -13.000000\n","INFO:tensorflow:Step 520, mean_score: -14.000000\n","INFO:tensorflow:Step 525, mean_score: -14.000000\n","INFO:tensorflow:Step 530, mean_score: -14.000000\n","INFO:tensorflow:Step 535, mean_score: -14.000000\n","INFO:tensorflow:Step 540, mean_score: -14.000000\n","INFO:tensorflow:Step 545, mean_score: -14.000000\n","INFO:tensorflow:Step 550, mean_score: -14.000000\n","INFO:tensorflow:Step 555, mean_score: -15.000000\n","INFO:tensorflow:Step 560, mean_score: -15.000000\n","INFO:tensorflow:Step 565, mean_score: -15.000000\n","INFO:tensorflow:Step 570, mean_score: -15.000000\n","INFO:tensorflow:Step 575, mean_score: -15.000000\n","INFO:tensorflow:Step 580, mean_score: -15.000000\n","INFO:tensorflow:Step 585, mean_score: -15.000000\n","INFO:tensorflow:Step 590, mean_score: -16.000000\n","INFO:tensorflow:Step 595, mean_score: -16.000000\n","INFO:tensorflow:Step 600, mean_score: -16.000000\n","INFO:tensorflow:Step 605, mean_score: -16.000000\n","INFO:tensorflow:Step 610, mean_score: -16.000000\n","INFO:tensorflow:Step 615, mean_score: -16.000000\n","INFO:tensorflow:Step 620, mean_score: -16.000000\n","INFO:tensorflow:Step 625, mean_score: -17.000000\n","INFO:tensorflow:Step 630, mean_score: -17.000000\n","INFO:tensorflow:Step 635, mean_score: -17.000000\n","INFO:tensorflow:Step 640, mean_score: -17.000000\n","INFO:tensorflow:Step 645, mean_score: -17.000000\n","INFO:tensorflow:Step 650, mean_score: -17.000000\n","INFO:tensorflow:Step 655, mean_score: -17.000000\n","INFO:tensorflow:Step 660, mean_score: -18.000000\n","INFO:tensorflow:Step 665, mean_score: -18.000000\n","INFO:tensorflow:Step 670, mean_score: -18.000000\n","INFO:tensorflow:Step 675, mean_score: -18.000000\n","INFO:tensorflow:Step 680, mean_score: -18.000000\n","INFO:tensorflow:Step 685, mean_score: -18.000000\n","INFO:tensorflow:Step 690, mean_score: -18.000000\n","INFO:tensorflow:Step 695, mean_score: -19.000000\n","INFO:tensorflow:Step 700, mean_score: -19.000000\n","INFO:tensorflow:Step 705, mean_score: -19.000000\n","INFO:tensorflow:Step 710, mean_score: -19.000000\n","INFO:tensorflow:Step 715, mean_score: -19.000000\n","INFO:tensorflow:Step 720, mean_score: -19.000000\n","INFO:tensorflow:Step 725, mean_score: -19.000000\n","INFO:tensorflow:Step 730, mean_score: -20.000000\n","INFO:tensorflow:Step 735, mean_score: -20.000000\n","INFO:tensorflow:Step 740, mean_score: -20.000000\n","INFO:tensorflow:Step 745, mean_score: -20.000000\n","INFO:tensorflow:Step 750, mean_score: -20.000000\n","INFO:tensorflow:Step 755, mean_score: -20.000000\n","INFO:tensorflow:Step 760, mean_score: -20.000000\n"],"name":"stdout"}]},{"metadata":{"id":"mDoR0C0ZKCOn","colab_type":"code","outputId":"aba41a4d-2957-4ea0-d511-eae7ea4e238e","executionInfo":{"status":"ok","timestamp":1553254513355,"user_tz":-60,"elapsed":3908,"user":{"displayName":"Piotr Kozakowski","photoUrl":"","userId":"01014928596539690143"}},"colab":{"resources":{"http://localhost:8080/nbextensions/vid.mp4":{"data":"AAAAIGZ0eXBpc29tAAACAGlzb21pc28yYXZjMW1wNDEAAAAIZnJlZQAA6u9tZGF0AAACrgYF//+q3EXpvebZSLeWLNgg2SPu73gyNjQgLSBjb3JlIDE1MiByMjg1NCBlOWE1OTAzIC0gSC4yNjQvTVBFRy00IEFWQyBjb2RlYyAtIENvcHlsZWZ0IDIwMDMtMjAxNyAtIGh0dHA6Ly93d3cudmlkZW9sYW4ub3JnL3gyNjQuaHRtbCAtIG9wdGlvbnM6IGNhYmFjPTEgcmVmPTMgZGVibG9jaz0xOjA6MCBhbmFseXNlPTB4MzoweDExMyBtZT1oZXggc3VibWU9NyBwc3k9MSBwc3lfcmQ9MS4wMDowLjAwIG1peGVkX3JlZj0xIG1lX3JhbmdlPTE2IGNocm9tYV9tZT0xIHRyZWxsaXM9MSA4eDhkY3Q9MSBjcW09MCBkZWFkem9uZT0yMSwxMSBmYXN0X3Bza2lwPTEgY2hyb21hX3FwX29mZnNldD0tMiB0aHJlYWRzPTMgbG9va2FoZWFkX3RocmVhZHM9MSBzbGljZWRfdGhyZWFkcz0wIG5yPTAgZGVjaW1hdGU9MSBpbnRlcmxhY2VkPTAgYmx1cmF5X2NvbXBhdD0wIGNvbnN0cmFpbmVkX2ludHJhPTAgYmZyYW1lcz0zIGJfcHlyYW1pZD0yIGJfYWRhcHQ9MSBiX2JpYXM9MCBkaXJlY3Q9MSB3ZWlnaHRiPTEgb3Blbl9nb3A9MCB3ZWlnaHRwPTIga2V5aW50PTI1MCBrZXlpbnRfbWluPTEwIHNjZW5lY3V0PTQwIGludHJhX3JlZnJlc2g9MCByY19sb29rYWhlYWQ9NDAgcmM9Y3JmIG1idHJlZT0xIGNyZj0yMy4wIHFjb21wPTAuNjAgcXBtaW49MCBxcG1heD02OSBxcHN0ZXA9NCBpcF9yYXRpbz0xLjQwIGFxPTE6MS4wMACAAAADkWWIhABvrNdXNvEPmO7lwVl73sPl0EDBzzvrz1O9Sgfa49FGnVhGNj4PrUzIEjAsiR14q5boH034au6fMfeHzW8BQIdLu5D8GWFcvhnUQvMLIDm/5fDlJWNI1pLZ0KekyKgRZvEg10IZZePvLcj64kGJzCMbJi6QZbX4WMzyM/ZwsXoWWPBmmlKBzFixHWdkptjcAYhpgDpXSILlIffpBFr5Fmv8Xdrl5eZtB/U18q6RE0tX2BrhekKyOZ5lJWnXZWIEICkLYIda8x0l/aAug9zkJAN2UJ5v8AfQgXgS7iPy41I11UQneH59QQ6r2Fy+bXVz7hKXvFUUQUW2NfwyAHSubAtKRV8FgrIBnKXwxAjc8zc/00LsdZVdehIaL1eI9qZtyap5GmVpF7ZJdkQbo7j2k9/o8Ztr6lwZrODqoujHSJK6V9bK0u9Et564zU+wWgftergJVAEl4m/D3N6/lD6Tni/a6bLzIcdcVjnfWLPAUBwAoj19NpxhAbe1VyiybbzF11k65OpExrnTpeyfXnWi2YKXmv6NMcvP6YS8WOK4pM7nWhyKetjJvO69p10oeh7Pv3PuQBq3kARIBKQ+MPYmymnbhgmxG/6w3hJ2A2Urz2k1DctVq7TiUCWnAReHSDqSpYcdQwxCm/lIpIwtl/dffgss5v+hhFs6NSNe3zqLc+wa/P6fKKBzHBPA6mZtXbiaJH0Y+5hMHtf92lFc+6I4pZ1q2XpI5Nr1V7em9lfehnp6KwZCFUTCrCle3ZgVn3/WlL0hiX3HqF/qGx1rSRBE7lqG2nGEQXx7BFJGNLF0vFVi0j2agV+lVqGOVlIxAjK3E9wGWVM0V7xAFGXQtxAYJ6qA6zMuM1AzlTqoEWcy+zkYm6Z2/Vn8RMHtpHaCW7GF05Wujcn0D05dR11MQem50GDiKxlzGighyGKWmfeex/qNBXelV1apol1nwDCUiSbC9fPUu70YI94kit+OnCdHe588u+9o5tqmvG+4ju2D1U0YtGzBJLwNtKIxTj+ycim3c7lWMz9gyNpdcRw85nQOO+UebN2j6KDuTy4XNxidtzFIcvo67EYGfl+q3WaPfQzFQLuOqvybvDRViyMxNwUidf7UCNcjzUMa0RFtd4HPTD4pR9pL0oOHG0XMOwvMlfvI/0tUl0nH8gKjxa10D+0pCAG9Sq76K3xNRb2QQ4PhDp7u3P+U7CY7JpR9qHasfUEAAAAxQZojbEb/+8M7YXuyASh7Kplen5y7UyO14JIrbok4XTbRQe3ORruR41lvOoDou8ClgAAAAB1BnkF4m/+AoGMvgzynj7iP1aLLDUqgN6Jo/suYQQAAACABnmJqRH+I4ID6PupU5REY3sf/aa003+ohQ9ie2tSAgAAAAa9BmmZJqEFomUwI3/3q7CUD4uvUkflAlpSidF3UDZUIJZsuBftR/Ot+q3GhwTn9egJDu/Q0u302gbAsB/IV/AFE9llDfK6lUW9694v3+SwMl6mllWP/WwEuWwr3bvYMZvmSan6TfqkNjfFQ2gIUyqpE9/WshNO04YsB0gSlJqYcVRMCqNuW9LOcW33er/NiW98OL0en+UGHVMigjfMRwtvQwEBx2TCEvqNHiHUZK47Ql9EBEwOC+9RX+RfTR+dz/gd9jggR1VC/W7S6VbwMD6OLZQ6pAbUGRTytAu/D3nUIHJpw2KDyU/GpGdKbj2WQLEoSWl/G1erWEOohpWvfDxHkGcHJyES1MZ1DvjVLlpUz1LfdDNETW9u5oytpQuaaon1mmF0VFggCZuK84NA/jjnYLjbiGa2mQVNuCz0xJcY1SlKU0HHZWh4xyJUcesClYPnBpMFzn1F47bfO2af9hq6jk1pG03lTPGyovR1gotc0zYA9fMfxEfYVjeBaSNZ7jIelcTidTkL9rr66l4XIEEMH3QknOAml0XZzRLxPdNkS67lJssdtIGefZbu54FuzpwAAADRBnoRFESzfiNYKja/0xDhfTAiTpHLhMdYxhLBGYN2yGAaE7v0DPPtKPN4N0Bh9SMs1vQXBAAAAIgGepWpEf4EHlYJarHuonGoEjcLs+tCeMaGFsuILlYL+uWEAAABrQZqqSahBbJlMCN/oMK3pi6Gt530QeupQ59ezKe10FpZIms3to+JEJWnpAJhBGOAA9pffTpSVzrDTkYei+dF/0XoAHErlseRXMtx8drWM8lgtra4QnL5SnQQH5ZVFNtwaHUoOeGKot2pc/HkAAAA2QZ7IRRUsn36T35XM0R9+0te5g6nnf972YKyQEsNo7mvZhGpibk4k6GYi5Xw1pjTrITOOCgNgAAAALwGe53REf4pxDGIcj/61i91Rdmi2jQszTK6GJmg9LxPIKqXtoMSsvmXH+U1e8ZQgAAAAJAGe6WpEf4Z1uDH/dTIseAsGQv4aVaLug7WXzVJtvrBkNRseHwAAADNBmutJqEFsmUwI3/QsGvLB67IEGcjaf9jQmmp2xwadCZqIc+QsmxjXnMZRrdlnOWXkniAAAAC0QZsNSeEKUmUwUVLG//1hfa7H0q8GiMhKKARrmw6Z4xdeyr7ujJme+WFMGuspP/+ISBWO3g4rVjHKMZczL9N457VGygcsdj7HL9XSQg7I9GyDcuezylwimiV/wFo/tnsP85jKo1N9u1U/c32ZRoJ6MLzOYkUW9bgpdQwc/vwurbPO6DqdSp7R4DyF3mA4Y/IOgh69Sz6dcgXGl/wwDmRo72yeNgYcXZp2lUk8XqHfEbj36uBFAAAAHwGfLGpEf4jgOwq2BPvs/N5sWutoQxXrlD0LcAJHkzEAAABWQZsxSeEOiZTAjf/7wc1PkbwRgr+fjnyNR8pfsXICQZIWSD29LO9OdoFeaCGC7kNQcI6QVUvVoNv//1awF0I8OAUiwMZgFzApKdTCW820SqjcXXuKH/0AAAAtQZ9PRRU8n31uIe20/lZp7IOycKeMuVxD421aM+Nt0XYWYsoR6FTF8OFoiDbtAAAAIQGfbnREf4ew+VpnMnXvUfVz1z4//JFTxQ6bZXAMaAhfWAAAABgBn3BqRH+E1CtCDqsru18CnQc9Zgwu5C0AAABbQZtySahBaJlMCN/8A8MxwHsdgYxkG6fD5jfkGVPerg1XibHNvThF43JXYYlxb2Qoy9V2WnWwkQ/OqnQwi/YBBYBEZaAn0hrUVg4yiuTQbTao5OXYZIcHYrYRQQAAAFRBm5RJ4QpSZTBREsb/+8P3/LU5AR1llSwpTA2p+sot/Ap2Yo/gxiVR9YF8xLoAHZjM8457AqBIx7NgPEKnQG3AmhGkbhUoZZDnZk8SW2DttxUzXAwAAAAmAZ+zakR/h60AWug79Ot7YB0Hk7/PuwF00L7Iy+IMt5Tq30K+2UAAAABTQZu1SeEOiZTAjf/9dtYVwOWMQEFcGvjV4Syzg11kPr1s9125wYmt9rvuf3mx7RYJ6jF6ErfUzdz0/3v9ae/ANOnW4wjLgPCIol0QSIsPYy4da+kAAAAwQZvWSeEPJlMCN//9fMnM4gfdWuAOqRQ2CfzqA5eQQ8IrF1fJShsXtjQZPj/+EEk+AAAAWUGb90nhDyZTAjf/++FR2EKSC1to8pc6YSO3kyM61vFVTIqwBRLCgvhTW2uELbv4P1t6j3j6R5JEnH93R+6ccuNOfX9oCM+dXHFLeM1xbdq83XfyShy+a/CBAAAAXUGaGUnhDyZTBRE8b/vgOOB3OQFcxg/ZQD7brmyDZG+0lanuFt5eWNnI+DlTWOnKhAAySzB0u8CurDZF8oIWX6RlKX6AlFm3XANEhpl6OfqvPFvE4O+fARgP5ChhkQAAACgBnjhqRH+HsRNxvly6e2k2NMUXC/vU7UPeiHPyDV2qBx4cRcnapKmEAAAAOkGaOknhDyZTAjf/++A44Ilb6dcZ4EfRuTLTuDvo3J2im59tzroCzqfsO8SsZ8WggmVkE0/TN30peYEAAABDQZpcSeEPJlMFETxv/AOiKyBDYvuQQEynViyLAUY1Ib85NXoM2W5LW9IU2aIRq+bew64gtHymArgzox7IataLMEs58QAAACsBnntqRH+HsR7W7gx/uLfAQSDqjarp3zTZOCfEh+6Eb54T4PDy1sS+MHUHAAAAREGaf0nhDyZTAjf/++s+OA7zprt9FalqM30xdIjlcH2DZkkFtbEKP2KnwPRDHCudTxGL9Qj/prFtXdQZd/8kyvQgWVbBAAAAGkGenUURPN99TaObeIE6LTTxi495G3R7GPmQAAAAMgGevmpEf4k5No92pOx4d3tVe/jt/92QHqWDJ1gLa+IuwyWHPUwqtoP8fV9Ym9imBVegAAAAVEGaoUmoQWiZTBTxv/vhUdghqkJVKiFoDxnkwJdI5EfJ1rXs+3IWgFyiLClNqrjfdF5LRESoN+/5vYqeA8+ZDoM0LHM+p62tsd2sj1HySk6QgTd2RQAAAB4BnsBqRH+JNh+VxVwUJ+Y3lV+ryJZRr2sjHdcr6hMAAABLQZrDSeEKUmUwUsb/++A44DvMfbctqjDdsKqNf0xOA39UYyLAzYIZyyzAIzc51s3jIS5ItuDvEtTcwkuqq/36WLfiOMPxF1mVVotBAAAANwGe4mpEf4EPa4byEQdMkh18OfNdNbI0980SSbNKs1ZKN+RKPSfGS6ySQe8jTiCjRi2a1L2bI0AAAABTQZrkSeEOiZTAjf/74VHYIlYiKMdyQZmvfSW5HcZkgM1RFiY3gEH71z/3fVBzC8liqsmHHacePA4zlFgRGs62G32VHnZQ/DXMUNXZSe6AM+37MoEAAABCQZsFSeEPJlMCN//7w+5d2niwPbxnXcPGjmeLv2gdbVqjCSUoEJNhCpLEDmpftWkeZV+eoSYNoxZWdF8ucDxxQt1tAAAAUUGbJknhDyZTAjf/+8OhtVOQEg0JinGp/3aOKBnRKunbXz9YVqPZNGUQc6qywYwzQClzX3GGGmt+n+j6fKijedTiog/HS3groeaf1ZLhN7MSZwAAAEFBm0dJ4Q8mUwI3//vEgdmkZAQXhaU5Ka/mwO9DUeQ8nbhvyYsA1HsjVV69Yl5uJ7Z+5cENl5oKK+xgTk4zv6uteQAAAD1Bm2lJ4Q8mUwURPG/74DjgQlevXVZC8R8NNBU5lXz/5evdLx/xJ2R3nKEWhnQTJ2MhCGZTRB2wGJ2TPpnAAAAAMQGfiGpEf4exElW+Y0oRl/FZgSZ/zHaDUN2en4eI/45oExvCn0eIC8T7xwpFCCD22pAAAAA0QZuKSeEPJlMCN//74DjgJBslHpy/fK99xLLgEoSJoPeLiwa05OsKQgsTdDRtpBJScyBfZQAAAFtBm6tJ4Q8mUwI3//vBcoj4CbBCQT2RExz0aNVXfJZgEMRjG13pvbjo3lZ6TLZP87/4osMgfYBX228XW1pfmU8k2l4pXWkF2ONrCPtLhUtLA/Hya2TerrONeACzAAAAUEGbzEnhDyZTAjf/++A9RB901G+fHcU/WnLCfcsk74GCTFIVwJ+BLzEBzmOXsONP0vNgTAP2KwhTXJSx3MerL/qd/oYBMUklPzb3D+iSqTSoAAAAQEGb7UnhDyZTAjf/++A9RCB74KMJGvAkjJTldXY5HMHa5jfjN8NHs9hBd2fEIGPYdqirJKwimKLpiuNnGPY71cEAAABDQZoOSeEPJlMCN//74DjgQ+tokgudjQu5NRNWobxjhH6Kh3ypTnRXjwJga7RHFZxm4LSd8nQrtx1GfqIwsNleIlTJwQAAAF1Bmi9J4Q8mUwI3//vBnCJ2LA71GF25ddxrRcaDbZ6JrFKdXAbBcwenVA97XR386r26bMwGidvp2PmbPeLJeVTrLs40YO/0dkENxamY3tNtYwvFfXlRKOo9yyo8A4EAAABRQZpQSeEPJlMCN//7w3v8sjICMILXR5ECQLilrI3AMHhNoAl8fjML6XnsQAd3GcNYf/0Vg/WkUZMoFlfh0iQSiC4/lahyE8PO4Mqsxu0If+cpAAAASkGac0nhDyZTAjf//APDMcB3mPxGGZH0fTP3ycONiP81zu+xkY8BF29QMHZMJZp5xyqLnicMHS675xoi5qplhfLQsVW1n9dVmaldAAAAHkGekUURPN99TaObeIE82Q23ingUfeUAohKj+s7K1QAAADgBnrJqRH+HsRJVgDz377tnoNimjhTZvPY9yHf2fmv0wsOMF+wp88Zp9A47UCY9LcsoRhIHvx+VqgAAAD9BmrRJqEFomUwI3/vgHICtV4/N+3TLBHjt+Fcn7DWjz8soUTNhnfkqkcMHOJu59OaFxhuVf+ti1/1ybHdmF1AAAABfQZrVSeEKUmUwI3/7wf8MHhAR1o8LZVW6pfXYDg2dpGT5Rvg+Cl2lQj+IQxqhhQEfFhAQwlKUGoZJ0yY1+WX5L6BuA3zmdozX8D1TAHaS+xXRLIKO1jwkjAd/JNceRMEAAABNQZr2SeEOiZTAjf/8ESF/UQffCua7E+r+NL8HxgRR3d+rdy5oOmnMWV4zw6CD9Z7mbT2yUeu59zFDbP46SRq9klvWZ1ZbOSW+2vZf9fQAAABHQZsXSeEPJlMCN//74EfgEPLRNeKMrWaAkB9b5moj+s74rojfFold/VwZwZnjRlAT49c0fDe4ZV3UBIGBK9CK3ytLN0Y42oEAAABDQZs4SeEPJlMCN//74DjgOFPapEipwNQ5JLf/iFkT7vTC27zO0O1Hhr7Be9x5uGF5a9SD1TJTW/IfLAT7Nh7DJMuMeQAAAGpBm1lJ4Q8mUwI3//vBm7tPFge4qVrTYzB1Rf4j8SwaiUd3Vn91V8IbddfFnSdsq/LMm7+bcOxpHHp/+sF3L+oZtRUH+68r3BQroy9k+62TGFClrSWs36FP/9yav//StMk6n2lq+8O7DNsoAAAATEGbeknhDyZTAjf/+8Ohm1TlEJHbNU/mWRKFgYzLN2yQ8jXAB9wCzOsfbVP2lU1EY8QOgEzKqXo0R7iSjvfThP3Ejb6DhTio3axVsWsAAABXQZudSeEPJlMCN//77lcOA7zH4jDF8og00F+0xFfHR4VDZOJrgOjozJdR/bpg8saNaDwaxnfv4X/eRhAQ/CW8F20/hfAqOk7ganOxPkKbRibKABEgisd4AAAAIkGfu0URPN99VLWJAgSGMMD0PyKXcxUT3cFaf6NQH7WvfiEAAABEAZ/cakR/h7ESVb5jSubSAV8FnD8x4gGanIIpXwVIvISOGwpzao9rHxOAHmz8XBzh60OGZXgi//qWtWwRLulsxilCnuEAAABsQZvfSahBaJlMFPG//Blf8ICQaEccs93WPEeRPsjPuDf4gDC1lzaA9vG8/MDesnCCrO+qwrWg0Z041Lq6MnjtnVfjvIt0qF1hF5tlhQF/0d8e9aqpY2UcQgjUy9Hd3+CTWS0I7CWSzCsCaZnAAAAAFwGf/mpEf4jgOwqq40N+HXrqPTGH7PlEAAAAU0Gb4knhClJlMCN/++A9RBAtNsNSCFkg711//JWWvGFsLrGyK1H7zXz+lTaC4eBFVxUDg/Op1j5vlFjWLyVGxDTLWQWt7v06Slq9bNE1cmz/5OS1AAAAMEGeAEU0TN99Tam3C0rJYqxDWOyBOLJHIckz5MqYD+DjRZpqPdIIw0nQ9LZT+KG3YAAAACABniFqRH+BB5WG0Yk+JGFIqTc/MKybQhljSSoiM7L+dQAAARtBmiZJqEFomUwI39P1UQmO/CHKa5fG9B98DyBd2q//xdlw3GXkmq/tlCxKXGJhBrqImPX+Tw62HdRUoCKQH1/ixC6DIupugrK3/kza+ntSNKpOnTIb2YWwrbEfjsAUf/8Csikkf00Yy4JNpq0NT3r3G73dUSTcvN2hSQC1jpvA4PpqiSqoyt1rg4K5nEb2X3Ta2J5m01MrHsJWKap7ArXaL/fag2Yxz1DnrBGNF0N64688FzTsfy3aMFyBNTwHxDcOM5VxEYpgPbQx0UoRaoEjPSOfPfZWW0CeOiVTmPWdE40N3Fq2/t5cedvB6NP7Q1mjnmICrga1HKXTUnIgZVEERRT/lQNzJZZSENnwR7l7OYZdbPaS+0iJ0PvAAAAANkGeREURLJ/NXH5jrirbiTEhY9Y/tlMwwq1venEm0Vym9wIwyi7kBVs/cNPG6n7cE56ATPcwWQAAAHcBnmN0RH/N1YmW42fk1rr5dWrvg0hpI+7/lxaqBQ9u9qPMWMe//elYZNo+NLjg8cYBUA/d4tfmi8Qz2UXcEo8Lvkc59K0YQAe1JOntXnqOjM4McR1Gx6UrMCz14vMSat1IO5CZ6UXFyvmTyeshwVAiRewmoKI5gQAAACABnmVqRH+HrQBVU2y7OiaX26aU/4ah2bLxWm+Y90YzwQAAADxBmmdJqEFsmUwI3/wDwzHA7nIUz9EUj2ygWhqxgWMh9LPrZxgjmvIJeUBCqf7Hl5oCspr9BQrYw3VwQWEAAACvQZqJSeEKUmUwUVLG/8dHUh70tVYbmJAP5Yav75x8KgRRrVZqGBAEqyIrK6LRnrpZlSZkXNqRr0UFWwbqNCFOpYqLeS1RQ2PyEkC+dBVplJUfZiuhOQAvvreTkE7X7a6IRT6oN2yIUD3gEVfN9FOVcWA6Z8kXNa18/OFVTCJp5CwGEmxWr1WMNJva0wrVNHOS92/kbWx21BB21E9sgNNFnB189Z9RulpZSpgE8SAYOAAAACABnqhqRH/N9J+2J1XcPpDCijwctlpIFM7gAfQRhV49uAAAAHJBmq1J4Q6JlMCN//kgYHc058NIdrA4tZt28lo3D3meW1zCAqwidHRRteHMGFHB/D9M0MCX00dwKKmnXGyIjSiIo/6dF3Oz1KEMYwbbPPtUPovm6LbF9t2EyWlXZ2AhyytqaBHYUktsnfMGjBar9/t8GUEAAAA3QZ7LRRU8n36t1fJ1HlQISEIEfpaQSyOL40QRE9Vs9zvdbKrtZIXDJfPzuaeN4LIza42Lez2eIAAAACcBnup0RH+HsPlbmN8Ea52A1bNIOYqtlLjflwD1gV/uybFUQRKpUOAAAAAdAZ7sakR/hNQrQyxiUawcG/LrF0fwKrQNveY8OGEAAABTQZrvSahBaJlMFPG//APDOoh1u+/wA6/FZcjU/jvz0FUQ/BieNWgawp32ZtUYldJMEVjFO412vb48n9u19LQOw5J+xi/1yuC/hu68Dsm1661iVvEAAAApAZ8OakR/h7DkMUQbpFoFWKV9q8E5k5LkELQbh4pZCjhTQyjdPLum4gMAAABAQZsRSeEKUmUwUsb/++A44EJXr11WJCw7ECvX60Zrh07UIdiv7aDfa0prRgD6jv/0Wc4+K9FDJsKP/KqGWZJW8AAAADABnzBqRH+JfH18ZqS57WtbAe/IO8nxI2HXbfpO2Nkfx7T6UrglP5FtwO/Jq+dpJWAAAABuQZszSeEOiZTBRMb/+8dnvoPCAeCb2/rWCrRRMAZhYjQsyCGOaWbOiXjQ45Hd4KBD69SXBo/G1jQPyXicID3GsLjYV+7vn2Od6v4G+WQqMHMoRfr8vCzfPrEsm0u23BOlwjpYpIOmf//S7gManfcAAAAYAZ9SakR/iOA7CrDrXGjeuosMdCPQVofsAAAAY0GbVUnhDyZTBTxv++A44DvT323Jx0SqVNJ82IBpyyo7DXCG+RxgT1NMsMLKQax8K4yQWP/71kziHHgGscwzo5L2jLizRX09ig4TiJCfNOrU3b1tM7YaGwugP+oXNvCezwnqQAAAAC0Bn3RqRH+IvbTI3y+IN9Ohxp8NHgN3aoRoF0CyJ7Kizh4lPfrWVzPrEIihMfEAAAA1QZt2SeEPJlMCN//74DjgQ+tokgude4jclMN4O+jlwgqt2WyQlDcvOp1uGtyVo8FNYf8DxygAAABYQZuXSeEPJlMCN//7w+5d2niwPcY4fRW/HK6IFjrt0h2pzlDEeWsYZI51tgKX93CPYsiPqPq3X5B1aWcUk55oKJGxvjXmYgG2872vqJP978jdfl4d7IIZYQAAAGpBm7pJ4Q8mUwI3//vD7kjX2UR2WhZ7zwuAd6E2YZpVsaXqf/9ms8pdq2AcVeqHU+2duGYw6XN6gJJPmwjaRjOqn/963zMAo7BbLbQ0xWrehzgNm86QkFh5F4CJPtfJrXi2qfvexEjl9SQZAAAAPEGf2EURPN99VLXrSyKjUixfkxdL90ti/7QhhH5aJuIgfKR9AwA9eKagET/7LDiGl5JEDu2HE71HainiPAAAACQBn/lqRH+HrQBVU2zahqo4JDrA0Wsth2bQ6i1pu5g6vv10b/kAAABIQZv7SahBaJlMCN/7/bCFg4EB8mDknVh3Li4nF+3GJT89JkDhouv2updeuOWXeNe5p0zRxNBaBS36FPmUKGpjSDMCDJPbtDwsAAAAREGaHUnhClJlMFESxv/74VHYRK6xf9nuJ/mW8vBkEYktQx5ng60k0zpkEVATXt/+Il4+t1Y/gWE4ff2KkBugUa5RlCmVAAAAGQGePGpEf4k2H5WycQ87dp8VYGfF0/IU9Z0AAABiQZohSeEOiZTAjf/7/gOFgQm+EEcUp0ZlHiPiNyhbnVoqZ+8senN2fT/GHn0i0DJAHVHDeAAPPdb6+08ywf9ggCftEN0KTRcMP6zvyMME1GrdoLnzWjL9J+Li8Z7zUjdev4AAAABAQZ5fRRU8n3sgptQ/K33IpJWP93HYo6zTTQPcVWgvT/tt2dtXG8r9XMctnwTJYMXBZrnQX8wcobNb8TVLLKjwQAAAACsBnn50RH+HsPlaZzJzUbRAKINcL+9T91jJ1Gd5JZXKKMYQPHi5kXraXljhAAAAIgGeYGpEf4TUK0IOqy+T8lgawG4WgIsZVeplQdYOkTmpMeAAAABJQZpjSahBaJlMFPG//AgTwyAUQ/9Od+649emXNJgblrSImgV1dWIgYyUeX0WkGz1oPOdc6UKgLvfxF2udL9r+3vyOJ3zOsvfmOQAAADQBnoJqRH+HsOQxRCCMCTKsMv80pUcMEWMaBnG+bZLo/o7gAggPnWMrK71nR0MaLVpDsOS6AAAAMUGahUnhClJlMFLG//vgHIBhjoEviMguc3ZwqTYQjvB+dEBD+tTd/LD/U79/aezoU08AAAAyAZ6kakR/h7ESVYA0yC7NG/OtKYdKqWLeGmJMPcZhbIRuFc+rxG1R7C8fxWQE6y89g50AAAE9QZqoSeEOiZTAjf/V0GkzH5jjHogSzwV8Ur8f+Ec2ZqF42vRmT3+0XBKgmcsowSJtIfdoeCmqACWCkOKlAnt3SySxtVHTz5bNRvdUm7O7h2Utc/NeW1ORiPi5ahyiTp6pzvTB7gpoYxSr8OJvwgt07bwjElPFl2/DXqzJ7mlN8Ko4FuS3th/HSK3m8xjdy9Pd2HG/G/ifzmmcbp4pvEyDSaJK1W2uFie5GND/WdfCyg3jjBDikafVsZZsB9db1ELXgOz6/mtW1PO65LJeMpnL03zTJn4LERpRc9KcXLD8eihY+pbPtstJ8ymI7vFDm0fdHk9VglqCfCuOm5LOR1Xte0qY4v0ZF497FEg6J2gXBmk0KeY7+LdmgHYEcz0qRZkyA93gRlWQCcg2bdHG3R1RsZFfPnVqyGpeaMJQsD8AAABgQZ7GRRU839CfjYe8gADwWSBEZ0dUg8ycG32E01qLxeAqrHnZPUJ5Lyx+8BUnDWj0XZtSKfpUSpAhT/szG+CsJq3bkVF7YO0bZ1YvXCQjDz1cgeQteorKqspdc1Fi4XCBAAAAGQGe52pEf6jUd6g4XFW5qqvsTJAPqVNfayAAAACQQZrsSahBaJlMCN/UtXwfA78SpsbvqlpyWOlTh6QrFvdRJjjmVnCQz861RBv+whwL1CGfpCv6GfEMxGCpC9HYQQUE0dQXgCNgcSS6mD0fn/+UCPvjhBVhxi7NZ4U2QtlGo9RKKkNi+HxgY6TR164zEIP510Utju+3WmWp8LCMIzVG7g95lqRKQwTZXOyoblZEAAAAPUGfCkURLJ+megsfuYeXW7Xti51RtkOkbffzL/ASgnb6g1Hr5xJiwX7/7XlPkdirFvrO694xVrxkHWvRxaEAAAAdAZ8pdER/gGwq2QxVj1MPQwZguW2S81h08qjXqxAAAAArAZ8rakR/3Cspb5/QadRcLmNZSh6g5dIyMVdtWcxCmcKN8VSOxyyA+VX1BgAAAD1BmzBJqEFsmUwI3/1ikeaEjHOd45AIQeOuawfhOqyYDx3N0zAkiYWZShFSr60ONAwyK1zoY65cTb4j02w5AAAALkGfTkUVLJ99jGQnHgDM8ln+aBB+E2IcxAM5SUB34rzqVNrLOXd+6bM2xuY4uoEAAAAbAZ9tdER/gw4tjgY4RdE2+ViqVUsAWJkRCr/RAAAAFgGfb2pEf4EHlYPO0DBg6RcG870j464AAABUQZtxSahBbJlMCN/OlXvvG5yLTHBTaPrAMVQaCfuKj0j5S7bXX6yX7qQf5RBo8No2FXkM/qWMiQOljutkadwskYcyP7bS28Hd6oB0nybXNXLow16IAAAAR0GblUnhClJlMCN//R5j7oKUrIUIBCbU00bLa+xJg3Zc+8xHG2utLSrY8Lzn10Z4sXQk4HLrp+BfVVniEmXvgjPxoP1/+SNTAAAAM0Gfs0U0TJ9+tCxJ1Gn0jDb8mTrU/wrX4mC1oGx1U6eFX27MccNul90QPfzPIXxTBWHVoAAAACUBn9J0RH+Jee66QMi55jOocyx//JCSfKN8+wOm/wGvvJAu+yXAAAAAGwGf1GpEf4TUKzbRXPU6nraxj9iMWSV9IgbTbQAAAFBBm9ZJqEFomUwI3/wICGIsD2OziqyuzNNpqzvABec/jYnGKj/psjvJAnKx6B88nuodchjeQKU8/Sf4fxPeI3Bur+63wb/oqD2JcdFb2ce6gAAAAFFBm/hJ4QpSZTBREsb/+8PuSK8cBB148viFWw/NHqyGs+rm2ceC8lLcGprCfFWoP4uRxFnVQBO6X5Km47W1B7f7HR4MpVRMW/C/8Bu03QteuBkAAAAjAZ4XakR/h60AVVNs3VhXxKqVaJQ3kpIsIMaGLSdIpeJrRq8AAABIQZoZSeEOiZTAjf/7w6KlpQgJLKb7CVQpIKCX2GtUUOMn6NT46JbGsSJDqhqGAUlsdEP5mjnTYaodhrLt7P8/yMykEazbfVizAAAAM0GaOknhDyZTAjf/++AcgTWnNhxEl9OeuJnrXVsKR/mhF9lVb2NCffZ/IdUboGwovrVQewAAAFlBmltJ4Q8mUwI3//vhUdhAycJLcpdZgvrBHbuS0ELA2qR1GWV/sKmX9lqbykcrOwXH6Qp5yOSFSCovTOHfR4H66+O3WFhrSEGUCxI2T37abEC4N7SLdnmMQAAAAEpBmn1J4Q8mUwURPG/74Djgdzkgaa9RAdjKXG4f20PXgc0gPoX4Vj0Ucek2XSjMdrwZw3mpyUgcp/uoOq3e0VHPfh7Zubi3pdTi5QAAACoBnpxqRH+HsRNxvnokvr1IBRBrhf3qfussoHJ2FXwia8L0og/7L34YHvkAAAAxQZqeSeEPJlMCN//74ESoSq306qyBH0bjlOnxLGbrkGJ2Ws39CdVP/9jJZ6U/XdORQAAAAElBmqBJ4Q8mUwURPG/8A6IrIENlUdGDC2rGxYKvs0Zv/PoytQYJJmfkiM01RixdLevQCF9AnGptUZtShDxvMVuLKVNUIVwk3Z75AAAALQGe32pEf4egL3/1jAGor8/AQZ6N6eggOUr7LqhLAXV7FfOYIjSYAD8KRrCokQAAAD9BmsNJ4Q8mUwI3//wC0LvCDOPxGGZH0fTTrqCs2EFRlT6/L9YpDrTovL58/OT2vGOyWsqXu1aV/a7fySUgV4AAAAAeQZ7hRRE8331No5t4g23Ywzxi4929Dj8IQsijl2gJAAAALAGfAmpEf4k5No92peKr/htFkwdNuhonMp0Ro1YlnQcnCewvp2SzmqBS8yiAAAAAXUGbBUmoQWiZTBTxv/vhUdhCtFQj3CF8I0Z8KP7PO2SHsB6lgIAgPVTE9oxFKS1kdE31Ian66qP+yqf9hh734mZwin70dY7QE5chfscqW3JQE3iewy17sVw7qZM5wQAAAB0BnyRqRH+JNh+VxVvSpS/NWVqjqTRtbJ2NrINTKwAAAQVBmylJ4QpSZTAjf8EwsBpzW3+EphZT+sMjYK72pqFmMagiJ/6RxdCOrdIwWTYN0R3uDTSb/XRdzEteeehW0X8cK0bfRzlGLg54Bih5AVBjp8JIdS+frbm7i9b7XCr5CBhzbRO0uWlX13WMTPZSuvGoTUMrRVgG628K7rQSE3qBkBMgw+CcCNefjGn6Zkx+q5pfh8Mt0QqMzqj7mnSuGudADa/xLYXydiZCYQ7/ewnsEpGkOK9yVq6rK50LcAUK+I3NgBMWmbuoV9kh2ewa2IxwBTrP8/IrrKTnIq0axfEx25XaS1gBYRb6kDVOAPsog/1jZSwYZ4GatQjuLuVTjthxX68wVF8AAABLQZ9HRTRMn9O9y352oflZsxy6BAxflZu0pxOibLVcNnE+wS7cYiWlcJ/txGIxRrd+RdfvjrOjPfDoy3gMEaOA1ktKWZzpAcG2KsiDAAAAKQGfZnREf4ew+VpnMnMtuJjFWqP/yRU8Vfsm/SA+y58AHgHPfjY6Ic5gAAAAKQGfaGpEf9fFlYagcVz5uBzD/QIPWYbzU9vkHWdHRQQyeCK9nXQ4I2EsAAAAikGba0moQWiZTBTxv8gMsVZymC78zcwH990XdZZmZzJ4wtrIvyzMXIOULxsiqy3MuQ/I48r4n72WrGhLg75YH6WEsBqONLX6BjtJ8TjNYAtjYKLrQKQl6vGS9ZfD13YsDE3Jvu673GbAbsZboZ/13vUZzF8zSZOV9IwR2roprdQI0feIeu7cnUYSMQAAACgBn4pqRH+o0ZioFg2l6PKcHO/zygwvcEiwx436ZjD+vu4Dcnlr22pAAAAAQUGbjUnhClJlMFLG//vgOOCFHL75XjgbP/nLa3NdHqfl9wTp94oAO8T7lQOqg3O0Mx03n3NRR7pYlze6hVEf4hs4AAAAMgGfrGpEf4ERwY+Gt77r+yEJ4SvPAYi+RMSplok3KwSDGQpSuB95O2ZN6aV0x9y74RwxAAAAZEGbr0nhDomUwUTG//vHT9FLCBA2nxris8BwI6TDKNzQGtw992rw5prohMGR5qiMwrOY4KxeyFVh4Z5nsHZDDLBJCMqHd6qBQyfBqvHpiFjePAt/3JK/Ps50Y7xjr5QVjDmLLnEAAAAWAZ/OakR/hNQrQmu509DykZGMqHC4wQAAAJZBm9JJ4Q8mUwI3/8evureZ7lc3KXdmxnQXgMijUWF3ZgneFpdm9nM3WOKWFmTpBzrecoyXK1AK3/TzKgoznmiqaxTbi+HkZzv39k56Xd8Rwloi2NmfZ/xAvKYkrgc7Yj9489acf/vZrf+rHiwEDuf4P34wfjU87tMB0gm0+ejzliNtxzSXYbeAmmP64RveFZ816fJhWEQAAAAuQZ/wRRE839BbCo5vTqlWMLNw7JP9qz7O6uD4nu1gs/O9P9c5x1S1ytTVzTI2+AAAAB0BnhFqRH+BB5XFIVzyB9wJwDYiu0A8ZNtKI6DGeQAAAEJBmhZJqEFomUwI3/18y04SEQRzJO8XoEYWKQTqvMsTVgq6JgxNfx+qCvDs/Eb18miZtc9keAkyj4LDbnj+9YUtWYAAAAAzQZ40RREsn32aKglvf2C9/3iUkf/5fFdsKyz6B6M9ZyCIDP+BGGOtWUef8hos++ANyp2AAAAAKwGeU3REf4exAhsvjHgSnTXrK0yG7hQMOWbFoITHQQMMz45J+HIhTgvG3EEAAAAiAZ5VakR/hnW4Mf925RATlfTwJxSqw6b0//nC+VfvnPeXgAAAACpBmldJqEFsmUwI3/vuVw4CBhG+4lg3Es59Okupr9RDias/wbkJjrkoqxEAAAAuQZp4SeEKUmUwI3/74DjgJAb7CXLbzf5nq0/dretTIIHlFI1cAGN8WOxHbKwvXQAAAH5BmplJ4Q6JlMCN//vBGT5YQEgiNI4YfqwxY7pw7Y2LpcvMCmuuwZiEyCprEorTJzgFnI1fTLxSf+rSP+efSERrkuP/Gkp1S5WCPd1cY+R8nU2/Ecise1u7zn7FazbzJRLufU5Eeprb8Vk57B7YxKp6pppaYCg0Zui1dDn3N58AAABvQZq7SeEPJlMFETxv++A44DvMfbdU1hKfbdntKLt5PPLOQy9ifPEv2J3HrARgxvABi0D1MxB+8N2XXdutIkzk+EXbiNPhxpbHQRvoMb11pMEpF0dzIxjytdVY2pHHvt4vqc8m5kctbxoCGPrGzq7BAAAAOgGe2mpEf4l71gLl090r4xvvAf/rji6zKN85FiDpcbWZht1n1na/FcE7BBvTXNdoIy6yLTUk5BdVxIAAAABWQZrcSeEPJlMCN//74DjgQ+tokgubHJTmcS2bnfPkzdAbeefDjnKfV7MFzXTULGE7gZCZQWfJ6qOVkkF8AtxRAEmR69//H0c3PlHntplsbW9FUKAYTRcAAABbQZr9SeEPJlMCN//7w+5d2niwPcKLWmwiVkBeI/d3EBe4uOe41g3mCl0o4kirBlvctec2F+QjLJhaKWtN6qGqc6uLOLJ+QX+s7b6lWU0b6H0eQNGAmGh+ktnyUwAAAGJBmwBJ4Q8mUwI3//vDoZtTkBBrgFeRatJUfKwOw/xfRodt9aPbXcnC8nF6plsAvzoY3VEtvh6EDBUsUk66y8/erm9TY2A5NAgUOAFC29sW8XnubtkD7UGheuC3NYe+ACX3zAAAADlBnz5FETzfgcr6UK161fjdbXl6k8H/DlXPiuIwC6b9ntAFIpHPdd+z6GlVWgcCRI2u8oYOOGNhC4AAAAAlAZ9fakR/h60AW7YGdQtuoFCe7a0uOvlc1NbK2CvlDazwoRCj+QAAAEFBm0FJqEFomUwI3/vDoqVTkBH89pCmTZspbW1QWm3JfZZ0UQobs65Mf5LlrQRMt+qH8hg6HnjvEpO2ssyZ0Y2DKAAAAFRBm2NJ4QpSZTBREsb/+8HrVh2DcyHMmT3dUQJoZsKjKvT90DkkLfs6f0DlI9Vo6Noy8G4U70dLwbnSCcJHA+s7ZIbEbAz5i2KHdv5BjXfmYaICxMEAAAAYAZ+CakR/hNQrQmu509OcsGZ40siB/cfgAAAAT0GbhknhDomUwI3/++A44Hc4/O2JB6m19DmjiM44bKO1947whNbBWqUtCS8/5jm7dHXzg/GKHpWAQzxNsp/1CPUSWvt3lZLPoW2mRgyciF8AAAAtQZ+kRRU834HOcQ6HnuhEEjNaVKp2xQF7WlSmlAM2ebIQcM+HcEqhiiFkUxiZAAAAIAGfxWpEf4EHlcUhXPIKbaQkC9FdA677TaPsDtJppmm7AAAAa0GbykmoQWiZTAjf+8HDrHeCOtIL6oV0eH0f5bCxgOtFfLX4aVpxgeHP3XvPB/PFVi0PRBD1cFMejAR7FtsN/v1VPGtv4Q5R6r3IoLYn6BFezGJqahiRYmgCTMTFmr0X4+/IlITNyh6ZhGBjAAAAS0Gf6EURLJ99PcOlKzJo8mgZiaFQ/cRdc7LruQ9PEd+7B2bhHdjrU2WOtGZMYKSDh35QT0pidHubWQAg3xDn8inXd71eAkmkAv1mZQAAADMBngd0RH+HoC9/1Oz3bWZ3fidkWX4f6J+c3kNLc5E/BUMQnT52L5poTTPp/CeS+I7NQNgAAAAqAZ4JakR/hnW4Mf925IfVrMHCbM+zJMcVzMO8153EB9P9wj2YcBm3d5ATAAABFUGaDEmoQWyZTBRMb8Em7juH6SagJhLTIoH9PuCUpPS4Xhl1dciUnk7fRxj9of4QGmTPb0WIBMBKQd3a1O1Yc1N/4kQ5x2J4nCSJXbRNmnX5rUdd/QcMYbapL5XWBGEC8Ck5/94+EirqJFZDiSW1qDre4E/cm0Cn57dgTQx4ezV9WnDqeA8qNeZT1cxuSXLNrFWIMllZp/93Xxk4j6Boviamv91IgB12vJyIkQkUqR+e7+pn8A3J1nM0FWyTrzkO7GJqjke9Z+zeEFS+E9FCk1HBH1Z9GwPLuDpVRLSp8hvtMrExkYwEqFJ7b8rI62OQzrSoZT3+EHwb1gJEd+fqytM5WmU4uf9eEg5WPacy8CyFQCjn58sAAAAdAZ4rakR/3Cdw0sHnaByzNI/DvCcZ5/JboTOZcBgAAAC3QZotSeEKUmUwI3/IDapfcEaIOp2yhhwVS+etbAWw+lDUW+XfaSyWSVIb9I+3RUZO9EZN/RG2xvytqz5YiVBaQOmzJiMDTSFB1apCB91XRO3SZTj1OyHGJzMKmR0UueNYpWGgXCDTBd912jNrkePuUoCV/pwa+tGQblTlC/7Mntowg9/HoHjvG3CnAA6FIBd1q7bFjLrs2BRVwa6HPk3/SE5LLvUOVOCcaqjJwOK9HolGgmBx/uLZAAAAa0GaUEnhDomUwI3/+SpVI7R3snQyETVUSNJhOfrvBOYWFqawEPOShEdffbNNhj/UU5hbs9Q5x/bRP+L8yTojPljMi35wxQSoHRofUyY5upTIiGlJUlsx7SzyXtNALaPJF5KQsK7ukuM7gVr7AAAAMEGebkURPN+Wev9EAG9OqVVseJN+Srxy0ekOaNhA4RLaQQxTYVhStfIpPBuX4Jw3IQAAACwBno9qRH+BB5XFIVzyCkOiEfbn4C8Yo2klpwwoQ9Hc0Mg62lIXOy4KcuI5ZAAAAKlBmpRJqEFomUwI38epbng7UbhTYq8oWxUQirssnQMSb7P6LX7rQapmaLnTKX/YQ+mvBQNdrilXpDana8P/ipoHFw1c/59iIYR8R4DQJ3J7cU7mE0EUO/5HJiOmapC64SNTZLeKkcj/wYvW9CxQjNGOHbzdRsVXG5vsdQcdxC46I2UJ+EPeA96UgnLDLarb/ye4UGev0CmekJlKpnTVesJsW30eCABpwrUgAAAAPUGeskURLJ/TvhMx1yLmtTPxEmDZ/zbA0Szadw/YOlgT6D9Dej7fdLCUUYjrPMIemsVWLNdcv1K6VaGlVIEAAAAvAZ7RdER/h6Avf/WMAZTJgkbiUU2GBEoThRTgMuuNR0WaohOWxAZ8OhULFKorD9YAAAAqAZ7TakR/3Ck9CeH/dWEJcswcJsz6v+24f6f0uV9+xslbm5g70jXAFOIQAAAAWEGa1kmoQWyZTBRMb/wIE8MgEocd6yv5CokI+/G0GUpxgPsND4Bh5ENvi1KT8dhbDutyYfbbj0VildYU+yRxq6FTnj94kjyaTRWbO1POXfK9v3D656X4I0kAAAAtAZ71akR/gQeVg87SZ4hSEgBWkj+ZBwvgo+Jen45iep+n2pmujJBvuAKIzNpAAAAAW0Ga90nhClJlMCN/++FR2BweCG1oRIssFvf41wc63i/NokGGJ2DfzyRKo+/2yaVmu/JJW0gloXeQ5hZa5thWPWMobrf/C7m5IKYTrfpZbz3D6/1t8557AoqeOYEAAABKQZsaSeEOiZTAjf/74D1EHaAmN4EWir8LxdLPbAZPO0M2+PY2QoG/X+gqbY2OfphX+3ikmCHy9qBRinr117Pf3bBw8DsA034PWVEAAAAvQZ84RRE8331NqbcOYV8ovUPdy4IB1YpDWJQ+wSvnfZOYkxxmYn/OzBmtYmp2A4AAAAAeAZ9ZakR/gQeVxSFc8i68ngUmW/lZdC6KspEQb2WPAAAATEGbXEmoQWiZTBTxv/vD7cBveCN3NVCXGk0YJpTWI6S27wxEVPlc6/1wnjM+kCdJnnhPo14+G+H7r7GYvUMBsla4e7HrvMr8M+H/u8gAAAAvAZ97akR/h6Avf+N3Bj9G4GRPrOabGjKMwWGQh7C506fOxYn6ZCu0xmofqwUIi4EAAABTQZt+SeEKUmUwUsb/+8WeTlkZAQajP8sPFAhMqTrywyxwavokyr3SLrM+i39/I0saTn/n2QD6z2f0Tpf5o4KXxn3mPA9D17rJX36vYMwHwHjQhLEAAAAoAZ+dakR/h60oZYAy4TexgRX9qffNaY7PWMw/K4c8McCTpcC8j63SIAAAADlBm59J4Q6JlMCN/8oIQcmggO56mz4fWIMz/u1LoEQTzV5bZxKWzyYuNw0VUour5rIsDiU90umHFkAAAABPQZugSeEPJlMCN//74VHYPx3QaH7LO6y4C4GizeyIGkLKZnJtuowmRMNw+Y2sUwAe7deUwGTWQ/Q5vriZGTkIsv3ToEeu8MQHLbiiuEa7gQAAAGdBm8NJ4Q8mUwI3//wC1hCBKrZCFEPTQr0tZ5jK1ORrDiuY02WzVUIDXNEWNY+1UnXo/FFfdn9MboNzrW4gnI9f6iOOt0mZ3nsie9H+/u+vO5CY8Fj2uueEyKX8U3QhLU3/TcKrIh9gAAAAHUGf4UURPN+B+FaPyFFuXZbgYj0u945+bqtESVlZAAAAKwGeAmpEf4i9r/hMofsoVdV7ULdGYi0qwviD0eMbDVTnrmwEwqTMGF7q03gAAAAxQZoESahBaJlMCN/74DjgiVvp1xngahySXAFM70j0RNfE5vodU8e8SsZ/26Vt1OWDKQAAAEhBmiZJ4QpSZTBREsb/++5XDgeRcDFSaseJcUAQLJ6v6GMzRDVOjaHtSG54x4li8aIVzsEHeZQB/Ro7xKkMpXenDvr+1lP2ymEAAAAvAZ5FakR/inEMYhyP/rt01Hns1oE+sW1Ax0NveAZCmuIIkSUdWkHaVhqEgSKY0rkAAABIQZpJSeEOiZTAjf/KCEPOyTOHP3Ga0L1Cjd69G+Yy89VyBy1dtUAy8MuVk4n3KI+TABbfvtiT/pPhtktWSGJDgIe68yWWzEpvAAAAHkGeZ0UVPN99TaObeIE6IhnjGIo8WCgt9DhhPUr/BAAAAC4BnohqRH+HsRJVvmNKEZfxWUIT/qJm6pPNGuzqg2ax/lk4U+jw9QOaXPUFfo64AAAAj0Gai0moQWiZTBTxv/1jexkZY1WGJiT87bz5jKDKBIPIUpqiOKoFqMU9xuIH4219MWUk+QJUsSOay780Yn710gUCq/8q5pJ2Qna2z31AIaW/7r/7nw3VuvmUI84z2bsgf+2HLYAKweSx9T91NAbzXY1OYhhH/Rhoq7G9wO4NOvgpiqkSQhrt64wDHDxeNgBhAAAAGwGeqmpEf5BXF/8DaTeNmUt5+Es7Qo54X0Q20AAAAGVBmq5J4QpSZTAjf/1jexZqwShLchuDwYT5GB8R7EoFbWKry1qGHy2qItJHoYSzhYHpftREHwxHK/JzKBYIGEaS6NFr4w6720GZ4EwFtOtZcCzxJ8lwWGRhMcXbe390w6J20G6KwAAAADRBnsxFNEzfiNYJXadE4qFM8TYpcvRqNU9B/tg74vA6eRce+4rX2wUXpty0yO3D1hzy0b//AAAAJAGe7WpEf4EHlcUhXPIKbzh/PPJ8bNtdgqDOQTgASBJuMgwAoQAAAPtBmvJJqEFomUwI38HgcMuPufQ+1zSgVSxP44j8Xw0s65lEdmNlEJPiZc5zpqLY2ASG2JgZLJLCVojwLtyVBJ2Djl4lz8/HuoMrqAhvSAdjKQWYwWvku/PylYEHSBWY+TA2i9oIGaNmjGx6LG1Kj+E8emazrbQW4uXowHf6FzN0wi4FUL7pYGixlMvdCCV20iEhjFpxLFL7zMrQiJXs0FePfh470wQo3VD0x5R78tx3albjYGEvRrqiobpY8RLCYk5k9JRliWcwQKwaEOM4PxOill4WkG8tUs4AocgZR6pVN8VW+TjxJGdBFTt3zP5oF1jV/E0tbkUBVpsIQwAAADRBnxBFESyfzVx+Y64rGY1MwatWM/mzoG1t/TCEbE6JSTBCJehjrlgVbP29VWiv1os6uGraAAAAcwGfL3REf9Pc6p27yHwywlkNj5DchvZ7/p56d3iHqPJJy0TZlJpjk/4Pz69n5WmRBfZn9bTbyeeOj4mF9//0Kaye4vxJwYkk6gt36PZDZ4v1fnhhsDUliUbwZt3CzWRiFHU6yepOVv1u7QQgijOhrLbxJsAAAAAmAZ8xakR/h+Qj8Y/7ryVtMCNL7ETyKPTlzq4JQ5ZB/FO0LOIakpEAAABFQZszSahBbJlMCN/zu7OUTa9XLb3eW75wn7Kv236GCYDyaFkEi5prdIBVAhBOujVHZ91Dr3M5z7YWFsFIPaTIAQW6IvvAAAAAcEGbVUnhClJlMFFSxv/8EDLdkB8H1mLAmhRVDhMcY2mJZi7xKwfk8PJDVJGGU0diyt/1T05IeVjQbyr4y3cfABc+R55miDu8jv+ke/z5Wlghc6LtRzgp8YEXwp4pBzCx8KkBF/UR9Q5Fhyu0T3/n7sAAAAAXAZ90akR/iOA7CrYE++zmxFeP15r8ZlkAAACYQZt5SeEOiZTAjf/HqZ9anCKYumtLnkGnytBtxceH+Tab7WDld8HDYn4TZQTl+0UF3pmdah258gjA7l14kM0QSfHRYvwKPR5YYx3eazaokeo5MLOsjv1A7a/dr6fcvO8RrTfuoTK4JG2FPtdS+FD/W7GxeAJSunV30h3c7FLYa2Yf/MWZdi1pjZRZxjRsrX+SukflJqQC/GIAAAA2QZ+XRRU8n6Z5caH3IhlDPEGgwBap5BKO3zzvTuKfuLPQfl8w4HkfPygPVlVEzHop4sE0RruBAAAAJwGftnREf4Bxg3qaOZGuOib2Lt1Wma/ikk81jKfwGZSD4w4Dnqo5LQAAAB4Bn7hqRH/N5HsNDLGJUxpLAK9mT9gX6FWVFLp4DnAAAABhQZu7SahBaJlMFPG/8pGuAUUgcuCRH/OwAPFSZJ3zCRS17ZFc1C9yqHmc8/pVK7IEw/oV7yJdMcMIavNsCKHsF1NBZJBQvHr63RBRk+tra9PDpuRrNWg1DBO3hQlugGNbeQAAAC0Bn9pqRH+BEcGOCVJY/QJGN/iSBG3l/ieze7n+BOFxVwCwANjCJKx4MnZIv8AAAAA/QZvdSeEKUmUwUsb/++A44IUf87s3whHef/K9eDWAMU8ldp4l040VzlBC65n7nr79QcRWreur7WcWsMhEQu2fAAAAKgGf/GpEf4l71o92pOxqkP3CQfV7f8+wwfDv0Vwj8bTm1R7NQsii5g4zwQAAAGJBm/9J4Q6JlMFExv/74VHYQpILXs6KXOmEj6xleFO7YS9ET5DIDmczDCW5+M6UZA4DyHmCTz5FoaK0i4R5P7y599cypfH8m6UCrK8Bf/vgLazNz1wC/16tdStMHxBIjZ4P1AAAABkBnh5qRH+I4DsKsOugF12qomW/SycjPJtcAAAAZEGaAUnhDyZTBTxv++A44DvMfnbE2rCH23XL0djoBAvFeci3xno+OtABO36XM5KyvnpfzxKvle/8J3308zrG6JVwhHh4rsYQUFRysfHizQb86GKutGNCt+OY9sPNteGFKlv4dXcAAAAqAZ4gakR/h7ETcb56JL659u40VH/5IqeKyAkelvuqVBBk49Saqe7VmZuAAAAAMEGaIknhDyZTAjf/++AcgTWtokgude4jcmomud9G6Ir5iVSH+0BGgJZK3LZX0JscoQAAAFJBmkNJ4Q8mUwI3//vD7l3aeLA78Ouck9tQPwA8lswErUECEWb2BfP9eCGIuUSrHm2tubC/IOXP44SUuC0wRBPcOUZ6UwEfr45QL2uWpkMDOFPAAAAAbEGaZknhDyZTAjf/+8PuSNeECUbzFqDoHVeVghCqcIHbUioPt8t0kMbXawr+fm53In7aCpqLDhgohy3fxI7sPPRi+v4k4BfCt+ArbIn/7hr/fNgQ82kiJyjRCR88G6lkZxT3zjOqC80cK2gQDwAAADlBnoRFETzfhHJ+XgcrH6mSiINjyp//kgke8KUWO3XO8bzOwBXQN1PZPssKbb2Qy1CNfZzQq7z2F0EAAAAsAZ6lakR/h60AXDOWhD5PePsAKefz0nC6cwFfCc9G9GK1nNdoyU44ymIpUb8AAABDQZqnSahBaJlMCN/7w6grIDgahJq1KP+AGMxDG43cWoTqfOrO+8pS10UHq7xCbX4a4pJt454b+M0H1Eo62OjD2EdzxwAAAFVBmslJ4QpSZTBREsb/ygqpP0TAwFGuhMe0H+LKq01NhOFINpQwf2KjhEvSCqmG9OMZE2/Qqy0272Gt0nHLi11coL4E4EoCisDk2gIBnTbHmELqhvwwAAAAGAGe6GpEf4k2H5XFXBQn5jeU7JdWsFUrSQAAAGBBmu1J4Q6JlMCN//vBm7tPgAJMRPqV++IZbHFD8LyLSaFG1Fj3xnwOAwhoQ0X3rcSTj34Hin1WxPUwQOdGW/+/l7yZtILXAd5Z9Wgqf/aUbVivhneizf3wbIgcVChpYYEAAABCQZ8LRRU8n31uIe20/kHBd7snCnkEo5gkxddxXoKnJ5QnndNiDb5+E/ApR6QfVh6OvB748Hzqtcox/cwTiLlMZwRAAAAAKwGfKnREf4ew+Vu++ZOvbsHGmKLhf3qdqHvTmJsewon5VlSMTQ0CLsoGQswAAAAgAZ8sakR/hNQrQyxiUcxzYxkWBaHD9jMxUSoH24CdrYEAAABJQZsvSahBaJlMFPG//AgTxcoiqckUP9wRELdF44LQ+J6zoeYQmvRkGAschCcjUxhb6c6E9RwD3q5AdXgn/otCYxPZzTf5u3+JwQAAADEBn05qRH+BB40VhOFmnLMmN/Rt6KeYRcltPwhBw6Yj8gKW5Be+d4emYrBg0QKVVh+BAAAAOUGbUUnhClJlMFLG//vgOOBCV69hkAISts4hjQjQR4n/taVosJ353hMn/bQatLAc7vxjVVNjpBAL4AAAAC8Bn3BqRH+HoDXutdsAXaf9vPXmz+vq9nHMZG8uHAniuS2VQVH0wdR1pyOJp2ZkQAAAAOVBm3NJ4Q6JlMFExv/BURmgKmQ2sem8M/cS5YiFuDEKj58ESA2Ioo541n9schk3ne2Kz27RHR2Mtl7sfZ6Tk8CR5fodz+unZYWDbUKNmf3Gr1+o9Rxy9lYH9eJLD5VCN+c3svum1sTzNtsJGBK1cHItHWUqHplu24Ckc7vqxHOe0ZtkSJdUThqkuknVrcKcY1ewS9SB7XN0a/r2y0t00YwIYLkLxxLwPmaULQNo/ape3s5adm82wYXnY+bUcZi6E22itPWO92rOH7IP21AdVJhrsPeiqRAczmuYnL6vkPzP4adcC7BhAAAAWAGfkmpEf84wSggSH8mtdkNj5CRqWoGie3lqD7CWgy3dy9dUCO//W9YZNfE6Xiu5GIKYFfDcUA0T2ZBvdwLfngiXt+EDkfDWv0K7BOXReZCjX9nSExoLgSoAAABZQZuXSeEPJlMCN//74DjgOprY+5Fmn26ag8A76dO6JZsciGMnQ2pYcNEnea2hGeomPmCENu5nJeGhNjlWP/GqpX7N4n3KhFSvZLLumw3qoPc7g3O8DwTNGcAAAAA4QZ+1RRE8n309hYNYZQ7+HfhECBi/K0X4frDrq0X/KbwmZ8BThyQ2+fhu0YmOJCJCHAKCHj32xcEAAAAlAZ/UdER/h7D5Wmcz5GDrLvRDAsCvzGdqHHvdNLLc9bD46lxTMAAAAB0Bn9ZqRH+E1CtQOK5zZksA34p0LQK+zTMJUM7eYQAAADxBm9lJqEFomUwU8b/KCEPdL34SY/510oR6jdtPc2spsXUY6gpzzsPQOWCHtIKBf9RsoPGI0GagH7nZSdUAAAAlAZ/4akR/h7DkMUOCGBJh+lP5pSg5ck9EFnrQLuSi5xBprxkDQAAABQJliIIACf++DLeYgN9XYj5zg5B5QJuOQipq5DurWEjnW3UorRPL+Gmp7W/eIApq4hA6GO7MXHIyDHsVucz/aMFYtKTDM8z25ZLQfiNZJqSiDN9723ESzhz79Cwt5M+8vOWuuGhG+Tvq9foXjrYoVXiVDwEsVfUmlk5TXbEqyYI+MZpo4mDJdCI6FifaFbOBypE2MZf5ZLIBH7g8U1rXJzlzobXWQwBVEZ3H50m7+bGCm/s/TO02Npz6eSKcRnoIL6uejcV8scEPRiZKvQslaZioRDQNI3E8UN6+/c9cAekW3/GkyyVmNiyROqcKtiIhBHi7pAk2+Q/TQP/YwlHw0o4KIS19ytAyX2tYiRc1FzBcAlpBmEdwA9/TY2cHx8Zwfg5om1WDRf7F7wfrHwkDprLxXYIXoXvv2/us1OW+b+Xd99mIui98EnMS3BL3wkn/NUHomZCIfEFUCZVjLemRQTcW2M7/WlzDgz8mCOUhc7zw+tq0JUXSMkS1L/cOl+y1DmBMgMzcbhhNbzDLeKWkrUH5XaHyryAkCyr9U1Xo/mAgTLsnExuVa/P4R7zZdXvaMIk7ek35j2b1yb5zOOXYieE3ARP5O5ZjAZBq4ndpSi6jkftKCBVFGUAJHA0yJ3SOwTwLlD33y2ioFahycU0S2U1an8v7S5ljSJFgUt4xqxRzXPLMm5Q/lf3uq1K+w5f2w1oQlGAPQWlYXk6xSti7jEz02BQEwIn/1VSW1lDXiIsLQ9/68Lmyq7+KqScF8M2ttOd6ikUYm63rAEQ/1uwF2tw+vDoK58ZLDjGGY+DsvOCVmKReWQXo2mJU3Z/+dukASHHoLCtd6cudmKCocwPIOFIrZIpuWFPYP7N5w3TQNI2J7uOoaE0UlH0FyMvlyJKCSxK2KXdZbEjdWLzEYIhD/rSHAXZ83AHb4RVPmbBpaD4FR9EhjzYQ5kM67XQ4jgiM+83E3jJ1tuOTwVjEUXOTiZCx8wy+C0dX3lqduBIJOztDCxzLLUgy1pDJ11agmN4hiYEyREcCB9PLz7Gu5NIl+dxSpct15aKEYrYOog2sW/3UjxYhz6/+6vpRxNYFBe2YVFfEse/JAPrCdAPUQXwCIpcGvCAqIe9H6yayOc2/eUy00kYB9Hqrr4iHiednpjGYIa/wBQn2kawkd5eYwvCd09N//TBsRiPySS7TyptSNcNkl3kSijqYAzDU0cTjAlo7gSKsBb2500bskV+Fa6qC7WYdX+uHeAkkUfK7djG3sDH96VRFsoJGnc9eiCjPYyGt4n1nklz/pm47DKZfzdVcEQVVqXfQk182IfGwqQ6YPNZNc85BxPoa5UG0Nc7aBMBm/EEqGdEA+1bpcQALD7p+Xw3OdXvS2i01L2TZwR76H4FwPsv/kg5K4Q+gtukTkAv8aXBY908KoaNfjcQLsJI9nUsxo8CIfcaHaV3fe+aAg0mPXJT7WoUH3EqFX9KWMCffuLgk141zvZdeWb9G2OAayn61mciuEFWUB3unf//3UT9ZWH3tY8Odtz5knWw9WcA/7IUpQHs7AGXSZHFSuMbAJ7IeW5YgcJ4/2xVyVfqWwGS8VXHS84Y1Oc2iGWIYRH4DJUk2xI3a7Xnxr04yGYzMm4F14EhpEJRU19w2fQAabSuyr93+75ace0n91G5DJsmYh8MQGBLpKZioZADEdyOWOL/b3dH1hKiEdZsQsnfs/yhdAAqZAAAAPEGaIWxG//vtcWB39W5QeWKwGeLr7PMDikgXstlWyUqYOYjnV7wDPmyxfP1O0kL6PoxR7zo1Gkb+C4TG+AAAALJBmkM8IZMphG/Hqe1RFY5rt5egm2nn3JQcBwuaJRr84/ddLi9/76agb7ejonTAqd67XVYfnqN4BuUaC+fAk+1lcholPMzUb84TknMEw30GkmM7GgprsKqCgwcaNTAOWdS32AuH1x/E2QaetB+/mEk87Hq8PlJqqlYpWPshK+oBJK6KuTSGIwg/9Bv+LVgT1pALr72MIHMQXmCPnTdlDJUk4bEuuLIqk2to3LjxIlmPkYtBAAAAGAGeYmpEf830n5JzFkX3j1Qd8Q+AUGDVgQAAAEhBmmdJ4Q8mUwI3//vBm7tPFgd78v7l86Um29hR9xxADhgMD9LzlgdAOtozhLW7UEkoFYf5otHinj6odtZRtbm/yoQHLKnpnZ8AAAAzQZ6FRRE8n31uIfah+VmyyuHZOFPIEoZWWjwctTHilRP/EBaLGght8/KC14RVIcYtmzmgAAAAJwGepHREf4i5UoIN1KmuYE1Jdm1Hq0jw330ncp4Hx2M/vKMq5Z4tgQAAAB0BnqZqRH+E1CtVyjEqWYFcr+nHuyCeliO2EXDZLgAAAE9BmqhJqEFomUwI3/wDoisgI4RECc7scuWX932y6q/Gcps10izFusxOOvMniGaXl24boq3k+UZg8rPMUT3iQxdNcNfEd4A2W68ovwz4f+7zAAAAVkGayknhClJlMFESxv/7w+5IrxwHtyq44n5qzWH5tokp7F/MIurN4zLf3qsBy/LP1c33i4OFn3CMZ1TBLIwXvbEHj14Vcb6Wt30UGVr0GhXroArD8I7AAAAAJAGe6WpEf4etAFVTbNo8fLWftyKUew1Jo8zvS8XWIpeI7UlrUQAAAElBmutJ4Q6JlMCN//vE9pkb6iHbDeYYscsOWTLED8gfFsN2l+iXqUwnGAUp7SGWRWf+Tj1MgMpeHoo49mO3nGzGBc/JZpGbbwjNAAAAM0GbDEnhDyZTAjf/++A44EO2R9S7lNO7TCCSBTAoqxeSZe0BGgDXxryEEgFwFd7IWXy8YQAAAFVBmy1J4Q8mUwI3//vhUdg9fr4QWmZ38DJvZypVj9XjHyiWE0qfucSUpcKpbBOOJb2lRZHjFXqbLB/AAyfZJCzDiBXSwSe2d3nKLM/QAOEJb8mw3ifBAAAAT0GbT0nhDyZTBRE8b8oIQ86YuCDOQEpyy9UWaSGWOxDO7wiSCJDYIhWcih6YPXwwChxsQuyPkj7FvLL/pgLFAoSeOsn9+9F+L3cY6vDMfiAAAAAqAZ9uakR/gQ9o5WBUzd28w1hPr/+56V66bgLLEK74aalrWqpQrcrjTCoOAAAAM0GbcEnhDyZTAjf/++A44Ilb6pCfLUOd7CmgqVP5gaCrD6c4HElsj+qTov37vq7C0KSjKQAAAEZBm5JJ4Q8mUwURPG/8A6IrIENjLKcjAdSvTZ0AUa2HabW/T6JywFoZYZ38XBzMVTWT9Boz5CqJWPYr9zgIZJzhObYT/U75AAAALwGfsWpEf4pxDGIf1f+tYtz0i7E23tkVLGF0sTK0M/3IyamfbQgfSCSb1BPWNUB8AAAAPkGbtUnhDyZTAjf//CJ/zIEGcd9wsgkXNNOYwvx4RPo8KfLe4zM2exTFz5+cPtcVtktZUvjRO+buNFeyS0d/AAAAHUGf00URPN99TaObeEbxmT4L02e3ocfhCFkUcu0BAAAAMAGf9GpEf4exElW+Y0rm0PJ/wJM/5jtBqG7M53mIMo8kEOBG1R7FTsO3uy/8wE9rgQAAAIJBm/dJqEFomUwU8b/74VHYNYZ3gTviajBNW4911XFIgnHuPPvf6PuINmTDgKw6tl4y1gHxr/o1Eoc443t4pg43rfeXg457+frQjkKJwwuU0Dg28hgOuxXzoy3/7GNwWDiho5fxpVul4JYlW/s2wZ6+FIFz9RB3Pxla/HA9ajISXcDsAAAAGQGeFmpEf4eNe2TaTV+is9Z+r3dj/4M5IrQAAAEtQZobSeEKUmUwI3/OMDPDLD7+9LV1rvQO8dIyBgu6gBZzKIHOqwdcYW0mj4/p+KVMgoNINR6fpQkbB+e7J43ZN6Thfy8C35cIcvE6ulfkdYr0ZP5p+0EdUB5C7vhx0r6X9DkZA0F/t8XcJ1Ay1jkXzIafX5I9JMZyOPPT23eliytMDQN3Nmm15DczfsII5lelhwcySIf4ku6wzRULjnVZWRMxcaOWnCVmvPK0WKnXS6hAGrBovrD5c40CFPHBrLy9JtUVvwuELo3rnQUyLyChZip9Wl/I5F161BBhavdM3/GZPS1ybPzPYbRmbIdt1SZq/ybmP8py+DhwJs6/zeOqKGDQvWpP3NX7JQuvuIBzABOHx7UbDGk+BICall/zPVb7YjEOQFeJndkK8GmaDQAAAEdBnjlFNEyf079+wBmd8pEhgLG8Xc7r/4rqUnzQQEv0ohGUAL7kLhkvedlsNVqrHka/nxHyyL78FF+w3UmoM16jfBFXZPl4yAAAACoBnlh0RH+AbCj0Y5cPjkp7DicIuUO2PQpkoPz3B9y9RBSHaE7ieceKLYAAAAAtAZ5aakR/3CsoGURp1Gbf6xetbBvbdg7XQAomNjTNFLLYoSnN1T9WTArX1a5xAAAAf0GaXUmoQWiZTBTxv8gax37l9sNbrPn6iWwUf8RxgPTzKiAAtM7FGqChNHUkeqgcz5deTS6OtdOTIhcU10sqfJ/G7PWcIpym3Fr6VN/qN2kQwEQKlPOHuE2NC5vUo2j48TIVNP/sgwP3kyf759nBR/KFqquMTQDl4oofPV/0HkAAAAAtAZ58akR/qMvbvw4JUlkDvjD/f4ktFYfXNpM4bGfm2nM37h+CiGSGIelwUC+BAAAAQ0Gaf0nhClJlMFLG//kqHGZpISriFH/O7LMfxPF3levzQUoyR8IdbCT4Om5XQ61r7CJghz72IBV8PYi+fqlXXR+phCsAAAAuAZ6eakR/iXvWj3al4qv+G0JMwV4sy5KJ/S3ER7xINYVeI2qPYr2GLSVtouwGmAAAAGBBmoFJ4Q6JlMFExv/9fMnXUwOYhp70FTBXdRxtKOUP+0QmwlUd3WHIaIY4UFj8Yp5d9xxsyZ9jD2OI0QWf2S1Zcu0qWkg1dQDmX7aUzQyVEEdscleAHvns65gIbvQxIesAAAAWAZ6gakR/h417ZNpNX6Kz2wcSbK+6JgAAADRBmqRJ4Q8mUwI3/8oIQ86YuCDOQ6JMWiYf/iRlpY+LK2B96JBM88YxxpmUYIE8NLeGUY5dAAAAK0GewkURPN+D+9PfvvPTrHvZhI32OEWXdhkGCruwTO5SkxQve5GtKHjiB1sAAAAdAZ7jakR/gQeVxSFc8gptpCQL0Wxcfx1JTwVzldcAAACEQZroSahBaJlMCN/HrvlPtuKGl8J03skzn9tX3N+UGKgPtt5f1yW1mhtqYtpj2nU5x+dydduAMrmraSE5zwuU/zBlhXh5g8x8c38YYs2MriPKi55w4mf0Kwy5yX9U4VJwu6z5gBqEnviXVdXlHGiM76Z3+hcdzSXdSABOlyYkF+z/YctAAAAAM0GfBkURLJ/NXH5jrirbiX21Z3CtP5YcQMlKXKnVSfznQBXI1IR3YSzO0uaD7W46BIwddQAAACsBnyV0RH+vXBhX1AzTfh+ibFtqz6eGMQ/DTS0BafSOlfYuavjTxTNGlR1QAAAAIgGfJ2pEf4Z1ktqM7YEw3wBj96oC8se5g/yL18ispByD6DEAAAA3QZsqSahBbJlMFExv87zR2otPwnC4OQ4HEYtHszJf56ewlbVNYDdRLi7DaNM07knXvZOBJlAsDQAAABoBn0lqRH+BB5WDzvLCGv4OSr3ZyhW71ibSYQAAAEpBm0tJ4QpSZTAjf/vhUdhErrGwgYAPXwOxHLcvgTv8A/Zx9DjubfQaVDTVMebAKQ+RkqueYnX184COJ7NhYucgTvasYPFK4F93kQAAAFRBm21J4Q6JlMFNExv/++A44DvMfnbRI+5Ux9vnvb9fZlh0+yV0sGFmgEwKg3swNSCp7N5faAgb86AHOGr8Z8DiL5UcrnP4kp/lO962eXNLUHqjhHcAAAAmAZ+MakR/h7ETb+sxp7t4uJPRUf/kip4odeYTuzhEga5MwAislMEAAAAyQZuOSeEPJlMCN//74DjgQ+tokMdI2I2w4XB31uWN9xpSf/b7WfaAjQAslbjZLMjZGYAAAABOQZuvSeEPJlMCN//7w+5d2niwexjKrUXs6RzXh1b13XA1usrXuo4Q+o9ouvQ1D2Q2Wm0PpqZJ/HVh07u1x+848SMV9TQUPFzbtS13xBB6AAAAW0Gb0knhDyZTAjf/+8ParGvCBK2RNZj2l1x+jAR0apeuLZtFzhogEMtR9DC/p+fqK8dILLP+IbGFXeex19kT5IY4p8S3kRMYF4bXDgihVw55Ypd4XX0hDtrHn2EAAAA5QZ/wRRE834ISoy161qatNvHcx2Sov/bKprDXTckBV7PTdQhKXDndGlf3VqRlNangjqwTorzS3wjwAAAAJwGeEWpEf4etAFVTZCVL4SV8n9Kg6fE2aMda+98WXMvpKkUyXESP/AAAAElBmhNJqEFomUwI3/vDp1Hm8IcqtaXumQqIJt/IdZM/LaCqQDp5QDE0HJXXxYKzqroGplUEOXiTI56df1TOqn/+xhhrsArE1GyBAAAAX0GaNUnhClJlMFESxv/9Y3sbc/g1iqUD5e5Uke9bb1C6YGB9z3awZ4zQI0M0unHOdJnHA1rrnsh4SimxV5KpEdhAQ9gFOWZITFvHKgu1okFYOADb9R2Rozp0tWuEyuxNAAAAGwGeVGpEf5BbOR6mhp16rXzs2+e4d71TyW9KQQAAAGZBmlhJ4Q6JlMCN//vgOOB3OR7sTCb+I0045cZJEQ4U+L3XCpcnHOJwFb1e+Q0W3Q7uM0aQ6yv11+YTtxoOdsy7RQ+LYYJDbR6Rd7L5XFyxpKDB+0OWkFw+YLFioVUh95SYqxqecX8AAAA8QZ52RRU834P7bhmByo0v4xgbLCddxYxYbspCwfmTeE4h4m6gzc5yXll0xuqHmzfhwYk8ZpP+UzAb50WNAAAAIQGel2pEf4EHlcUhXPIKbaQj7c/ATAoSkel2SJNrDV4kgQAAAGxBmpxJqEFomUwI3/125QhMIg/aNp4kCmUKfLpf/QrlLIP5B/oHjyqjvBOioUOoMoL20buyufrF/UIt3sv6YVyWShlKuej/5mcBImU0X9mB2jkyIoDrWUZBjFA6haFPIB5+QhgWaDpI61usiRgAAABMQZ66RREsn32aKglvf2C9+Bdon/+X9WLOPlx7m1yaoyHmMFFCEI61yhs2ICjfrbx1xwg+UNHTSx3O+VuFL6oUO/x14wF/ouCv7NddvAAAADEBntl0RH+KcQxiH9X/rcK9KC9kRdt+HJ+G/iNZPTGZ+diB+2gxK2Dg5bXzZO90e7yZAAAAJwGe22pEf4Z1uDH/deStpiTEc9tyHnyaVWKk+d3zLNz5R2rfH3mWEAAAADpBmt1JqEFsmUwI3/wDwzHAd5jvuFkDKQEVOoZiZuzsM3wXcA8F3ucQwarCn13TkhsIYJMiT3kh+1+BAAABTEGa/0nhClJlMFFSxv/BJAf7z+EOU2CWqKManlZ0LpOsfJcBa0PMmIuvVH9dbvAcE5UzEdHVQbTrz1N1jb9WiMOxu5R2cy5A81OxB+5ZBnN5eAsC32urtdTTKK7QubeiVOo7VylTo1bcPy3DSeA7L7gshvKbqb6yZRCFyi4whKSJ7RPi0+A7uR9tDdfvNWqXsP4NZEdkTw7nJJscByribz64QmP/ICkD3ZxZqkoqOTVn6dlY+P0SHMy6N2fn26w0kDbCjyPfS1dZblvZxWjU616LYFDlf6NoIU40eHrlDbvSs4WEm5sqj6lAdmcaMt1R2DDkFv3lJfP04qVFUwKkb4diB0AgVCwqoCaqflufpSnzlUP5zpGT6XXwKzdvxhIiAGd/HKUkMg36acC0SDhcMpk2kz57Q3kYvM63o6ggTKO91qvRby/6+hO4M9+BAAAAWAGfHmpEf84wSggd11AiFSufQxM9Luw26hRYDDKsE9vSwAn7Z3kp0GuF0Eqpea5n5unP9KN0c3L3GkYNAsaXrACiTeM7/yz8z1LY9wjvbK9PNQeYtR6wMIAAAABfQZsDSeEOiZTAjf/zunZ7sMLNRD1X4hTb1ZFO384nK86xDwxI5Vpgs+jNIRsolIEfvY6lsaVpHcZvcrFP1//9wB89zGosviWcqEVLIZ7/+WmTYZWLoPL9M4Sq38lHkYEAAAAyQZ8hRRU8n3sgpbafys2Rudk4U/87NYQWWpO4rxUES/JyvoT/lfbQ9OleOWWIZeQm2IAAAAAoAZ9AdER/gHHeAx5KQ6o8Bkl+v/9P4zpRB6XdjaQVIDAbYhBhxO6UwQAAAB0Bn0JqRH+E1Cs20Vz506WBTrufbbvFjJQQkJ28wQAAAD5Bm0VJqEFomUwU8b/KCEPdL34SY/52bZ/rGfsXPf5yCTDido8G77Go9+K1BJRLL67CRLKDtmUDogd5/yp2dQAAACcBn2RqRH+BD26tG3o2sOTqJd//LOf/0exfYfKLxm7J3qgsoJYfg5gAAAAzQZtnSeEKUmUwUsb/++A44IUf87wFjEcFriJhVhBdOObtZMRA7H8L7O85Qi1Ir/23xQWBAAAALAGfhmpEf4exElWANW8vMvakG0lwp25Oh8MHyFQGujxVcqIRtUex9+Z6rsJYAAAAUUGbiUnhDomUwUTG//vhUdhCkgtD/OKNA6dG1osSc7TN7EVmJ8mHYKHJo0kuzmVJ2OcDMUMCsDRqChpG0X62F+dsmKA1zveQZSPh2EvbMNzsMQAAABUBn6hqRH+I4DsKsOugF12nJUDYjiAAAABoQZusSeEPJlMCN//07ziWU3i/kpvzHFCOn2CcItTi1yFM4xbCfhV3Hl7kKDI/9hW4d+LvO5FlZCITgCrsyrc2JS6CCKYTrtDuaRaAvGVwNn+giCSu/6I71EcdF6AAzs8YoRZvHfx+r8EAAAAuQZ/KRRE836foBN+l6dUqFIqLR7fdXGqeg/3ZJ/83TbhEn5pkA8K9kzYNkm5sAwAAAB0Bn+tqRH+BB5XFIVzyCm84fzzyfGIzavoEg2o1vQAAAFdBm+5JqEFomUwU8b/5u5inuMF7Xq8EOtQO4YzuHe2WcPkEeIzleyoHzUSRtXdZDfWefPY5FxPuebWlVj4LVrBtGv0RP79YCO7uGRp8no60SRyqygnK4c0AAAAxAZ4NakR/mJ3s34Q5H/1uFexTk3VIjFEo/fwfhqCP78QfAaOXWnodXpVWh/EbwDilHAAAAEtBmhBJ4QpSZTBSxv/8AoC0RB/Gq3N+gqfQl07//0asYn4GjHzgmabjVHeCeHimF9bt89uzqXPNoLLVfk5lK+SEWZBMH+jdIMD7NJYAAAApAZ4vakR/h+Qj8Y/7ryZBazBlCQSff9dCgZHsO8153DG/1Iu47qZ9y1EAAAA4QZoxSeEOiZTAjf/8A8O+HA7nHfcLIGMIzn/UI6nxZeLzvaZpNm4NR++zRTzqakzFTJjhdFc2sWYAAAAvQZpSSeEPJlMCN//74ByBNNlS+GOnj/iYE7uR3e+eaDjQFr415CBPX8resUSqg9AAAABEQZpzSeEPJlMCN//74VHYRK5/s8oAjSJlNV182JSgSuYDsNFFvPml8Z2TC06yUjgNGaIT+MGdO85RZn6ABwhLfk2G8T8AAABHQZqVSeEPJlMFETxv++A44Hc4/O2I/eN0x9t9al1x12ezWSFYh9FLhxu8aTceePsZOcOf/noBS072io579hr6GIKYfVDIPsEAAAArAZ60akR/gQ9vOA+UWsy1Mn3l//+f5lQBPWWuO+aJJugzYNwOMU04nd4nwQAAADtBmrZJ4Q8mUwI3//vhUdhCQoq5E2kM7tSSvPdRsTqN2Lqyx1YynzDOfYXeFZW5mbzBzov4IZCErRmMoAAAAElBmthJ4Q8mUwURPG/8A6IrIENlUdHg1vXava9lPZfzExXzXgqZrbgK0CpfLrF03F1J4CM9W6Uc/jlRwpOi+VZSpqhCt8/y775AAAAAKwGe92pEf4egL4A9rxjL4OGkdU1tP9jDyD8YL7cVkl98cgRXNXC0itKOXK8AAABMQZr8SeEPJlMCN//76z44Hc4/EYYtcgPBH9RgO0Iwi8OXMzKNpI9izDnv5+poRBie+02otvR+9Zng2UG2aLxEiEAQa2EyHRWuU4x0ZgAAADBBnxpFETyffSt6tLJT1q5nVhquTgXJSFs0fGTiWZMvJijPSmPEtnIL2xGLEmLILUAAAAAjAZ85dER/gw5gGJbwM9dZrSXxrWmk7SS7YSzCpHXtnb7SsiUAAAAXAZ87akR/gQeVg87QMSja2dC/csmeH6AAAABxQZsgSahBaJlMCN/74VHYPZq6tmtxLh1sxyPaSRfT8XGQqvC/Q4343DJxZXe1Sjk7F3JK8Ga+iJBjybN1BsuRRnEENt4FtlTjDMiyslD/QE2S+ie1BjE0ye1VekieYTcbeA5uWUCc9oAS0oHJc7DDprEAAAA9QZ9eRREsn3/Spxn+3+PADN8tr/k8dm4t2BxiP+1thpygNS5/1ulcsz7YQxXLmulOZ9bSD8leHIfywjbusQAAACQBn310RH+HvdOD1Xbhj4zYGbutZl5v6HqyOaqspaMKaK5qWeQAAAAwAZ9/akR/gQeVhtGJPhygpFSbn4CX0RYdxtJU+44KMHn0Oo8JLVH4iFlYBJyRuWktAAAAhUGbZEmoQWyZTAjfx83w+eo3udPz/wWGJzqAxttakOUQ5KIzV5rlta9UDHrdCKjDiNALKgEJA6puZDaqChxqlBxSo1lMmdEYH9x6+Iw/G7cBXYYH/vbSuijBJl7Iv1SdG2SxEz7WKOVaXQ7edDngcN7BTtsHxNL9kKKH2vdnV7h3Up4bBiUAAAA2QZ+CRRUsn81bxhrs9aLlJLIhMRyrlgrjTzbZxY1kexOB+js4nC73oqm2oS8KXxFXxDDK0y6hAAAAdQGfoXREf9PejGW7yeA3K+Ashsi9mc5rzdirkZ1i8F7plCYqBBB3h2i/Y5FflOyMyOlgUqYL/G9xIzGl8z+0IWKUH+x+Fv9B7/aREDydtwyQInEGxGayCrl3zIGiZ6TFkkOalquRg4SIEHzOpZAOwkut3sA6cQAAACYBn6NqRH+Gdbgx/3XkraYElmuoUqexzkujzvNedwvPC5NXjL0QwAAAADlBm6ZJqEFsmUwUTG/8A8MxwHeY77hY7su4Js5ePegnvwLeGgOmnRghK+XuEJ9p1ZqTpeEmvxKPv2gAAAAWAZ/FakR/gQeVg87QMT0JOabD8/BfTQAAAFtBm8dJ4QpSZTAjf/vhUdg9mx2D7TmdFz8qn2yy4xD/XK0YWq6bXtpAIDSMH89fnlAnLzoUK41nLcHaRGQ9VYKHwPVgSIejbid4zu0EWxJ7r5+L5ivAS0Wc7KiKAAAAU0Gb60nhDomUwI3/+8Gbu08WB3vy/AdqEakb767uJfNyZA4+pjSLzSTlkccggQWj3N8llt825rufA9wokyOuXzQUeglOEArfFlY1GPqmrtZTf6vhAAAAMUGeCUURPJ97IKW2n8rNmOXQIGJer6mMWBu7bBqazEhQPzV3PBt8Hffu6rEFJycnZ/AAAAAiAZ4odER/h7D5WmcyS/Um4Sed4v71PLKjq2xzxVZO257qOQAAABwBnipqRH+E1Cs20VzmQO2GD7g9FZUkcZEXqdpxAAAAR0GaLUmoQWiZTBTxv/wDwzHBCj/nXShHqN21NSADJU/Jvj4ZToB/1ZTomGKrdxyAf/iOIoo9DdUgNWCwuF0p6LBGjn9Ou8lBAAAAJgGeTGpEf4EHjJ+LiOn4U4Y39SmH/rJxyXApfMRHQp7a1U3JiFaDAAAAMkGaT0nhClJlMFLG//vgOOBCV69dVkGYj4aWcO7kQz+XVCeG9zDp5MQEvFDCbSD/y8jgAAAALQGebmpEf4n37kAo6ubDhZOY+8dpADGCpKbr5jcNPPpTe79afnKaDNgbJdD1QAAAAFpBmnFJ4Q6JlMFExv/7wetWHYG/jzvTdZyGtgH/DXdI0sh+SMrklDv0lY1O4TvP0MR8DIsygALU9Fs2WbnRqnA2PULv/h+6+vnPyCATdnYbIXw/4lIPbTxVj/MAAAAYAZ6QakR/iOA7CrB7i4FdAnsFNG/HYixgAAAAVkGak0nhDyZTBTxvyghDzpi4EXYDnbEfvG6Y+2zMQ/Hb9ovvg/0chC5MpSxOmJq9zFtkv9T+YO8D9ACgwbrIT3dU6sd/k0sYujXXO5r3U7FCF0bBVqsgAAAAJgGesmpEf4exE2/oQsw3fbuNFR/+SKnih1Q5+0sivYgWYAQIdsxBAAAAMkGatEnhDyZTAjf/++AcgTWtokgude4vE93Jud9I9whdUFDZJTfR+J6YiF8XdIzfNb/FAAAATEGa1UnhDyZTAjf/+8PuXdp4sDvx77lN0KgU/fpK8+IINR1opM0Pu4CozS8l0Ud/us9fntwO4IG0oZYvNBRI2N8a8g/0sIn+6f7QQekAAABZQZr4SeEPJlMCN//KCEPQpIQEI+hBvIFMoU+XSr3pPEb6tM+0KTmzi2hUR7j3nqTXSRaGGG+1IgCzYffvSdjj9k3eSImzEFiL89xBSAuBHuAan1qCm0SbL2AAAAA7QZ8WRRE834Ryfl4EPlupkoiD3UHAn/tyNF9cYKC+bqyZpWmU1htKRBn8v1v9le4Ox3FhOYCg5NB+ZdkAAAAsAZ83akR/hNda3bhOJUZIUUY2jVhVdrmxzmILX98ciz/mLhHdJSmfczV6U8EAAABFQZs5SahBaJlMCN/7w6grIDges6Q6N8LLKBBR/r+r01RScaGlerMw4mNRGZeM8kNv5hlT2J2qEHTq9XsVBx6y0AugLYMoAAAAZUGbW0nhClJlMFESxv/74VHYNZiODju5GRyMj9wlHqF1Lr7pV20VHiAMrHfGA6p0v/Zr0/dcP6yMHwgdbTTN5Pj6VJC/71UkDQHF2SgOkPs+SUevM06PMV4Dc/8UKpkKq5MzAfyAAAAAGgGfempEf4k2KJe5dxcZpOfMZDO0cRrAfBG/AAAAcUGbf0nhDomUwIv/+6w9u1GKETDvOv56IQEhONuu4MWW8U5koOsb6z0DMrcU8frYGb496r+0cR0j8/9YT3bwj6KnqNIzzKxPxz20387LsrGHGueynDGovxAkh6APLCz5JN0en8GfJGJgIpTpHDn4oF+AAAAAP0GfnUUVPJ99biH2ofqCeWdcOycKd3M3sDTTOXcd9oRK3Tz1jn02+G3H+x161bG+i7O4manzYReZ5zgb4/OHSQAAACoBn7x0RH+AbCaaY6JFZrmbJvY5TSCOfppisvf7jhPHXneiHxtzcLARd1EAAAAfAZ++akR/hNQrQyxiVLMCu/aBfsIZgv1fU2C7T3Ir9AAAAEZBm6FJqEFomUwU8b/8A8MxwQo/52bZ9Gsz+sY4O2RS3UuD8X0CNvCqI3uAP4dESewlMKkkQ0sYA5go9F7wlfSnLQu+tTWBAAAALQGfwGpEf4ew5DFDghAjpjv5rOLhqU4poz0tQCo2nvYk1MNk0SbxLDlMLbaRYgAAADhBm8NJ4QpSZTBSxf/730txLjWi6sL/8WkZ0SxdX+xgko4N3jmSClvhOh2sSC6GI34PiVLH5H6DQQAAAC4Bn+JqRH+HsRJVgDVvKzmQczvZIZd7ldG4fIdAbM1M+rxG1R7C8fxWQE6gkEOdAAABb0Gb5UnhDomUwUTG/2eshp/CHKaTLMVIv8yljz56U5Qg3GnKqoZs4y9mnt42jrRlYQiQvaz8jEhuvC1SoxRX3jdtVE7CW2y7mknEG5FfETEOLay9NAEvDLcXz3Kvp5o6r1Ecn5DKuRfL94ZwgC1vfOYsAR70mWBfVzPYmO5bWIXHrsivWz4TmSfJX8BTws++We/rWq+eEIZTOpQXO18Z4TBZJoJNh+NQKQ888eOePfvsvOd7ExGYc72AMj8IB5nA3qIIpZ26bWxdEigqp5WqA3CPmQ1ArglG8O23AW9a0mGskFVKIURyO7SDrMfvuSpM2Ufg2ko07dDeVY280dgvmukE3rlhobBVZ9eHnJpB/rDydTtHtcgUHxLRRN2WPaWUwfFExO2NlPxDNJNoifurDRR1TOE2xSpff831Zq9dXYJqU1zKACLUfsS7gTYxNNWcQ7fT37PoiXtJMLOR39uru4L0h58owlQKSE0dbvzeGUAAAABVAZ4EakV/6f4hhh5j4pezXGmc10booa3ca783iOnUPxnKk5Y3HfubpAwe+uDDln7x2eceSurPShjo5iSao0hBxT60iJArFLPcDKmG1k+1SrnGiBiThwAAAKNBmghJ4Q8mUwI3/2esg9XV1nn6OqvkurteiZEF/cX1g6vwMjx6DR+rv8oeaoyhuEGc4M80j7FCjt+C+gCKRU1EONaFpssJsG4r4wx6++Vl4OFVBDT83lw8NrDATf3vTsciLFGM8eGrhKs8pvOB4zbD2z/wbkYod2zVrCHakD3w/D1EFZvcK9tKs1Nr9e2ax3We6/gOaYcpR2bYNG1etHQuzTshAAAAMUGeJkURPN/smxOo9ZgcqNMHlkhDvC9wqzrKarKnUnQpTSCQZIsbYKJtUJkWtnazMvAAAAAlAZ5HakV/hwRZbYwm85/cCb2d75a9rssTV2GyXh2ERuMyRAZOQQAAAGpBmkxJqEFomUwI3/vEgWa+og3x1dnJ4PoDdBe3/3YoMTUoB3TyU4Ej/zJxdOfoW4aHRQ/I+nHKZhqGdhrBqoC886/vF7jEJplVhBx/pyxBeWOV/VwqfkCrurEj7aiFkmd0eJF8WUso2Z6AAAAALkGeakURLN+CHmGW7/MKreRIEd/8QfLllKWohrf2XGUEjDty3V/zICAlat/0cmEAAAAmAZ6JdEV/jY0N0RG/BawsJgKYUujP+rhnqWcqJ30YzX0c3+Yd3EUAAAAkAZ6LakV/j/S/wWIll3CGdQpmmnlHJJJMDXC3Dufc23w6T4+RAAAAUkGajUmoQWyZTAjf++5XDgdzj8QavWtNM5/v7SWzbYKyz+gr1lvvDdabJmin54HFA6FS5USZmz7Qe5i8GrS/OeBhEFUi7sU8j3do+iQgzZ6NN4EAAABqQZqvSeEKUmUwUVLG//vhUdhErrGwMwT2ApJmJaluur8WZXKPf3cxs7c2sP/Xbjfjlt+FJSs7mUuT13MbzM+TrSACCZSMBoDthPFifoGGtkxEnj+e8h73JprN7phM0yaCdgAO7cfoh5J18AAAABgBns5qRX+OkjwIvXb1vb5y58tHlKJ2YgMAAABUQZrTSeEOiZTAjf/7wZu7TxYHcd1Yn912Q4e6YwUip/A050OaFjhX6O2qQXp5VjUL059lIsYYgSQ2jExa1CDIwChKJheC4ssji0HjIrmDdIWl9Q3JAAAAOEGe8UUVPN9/4NuyfQi2HxhNnY/bPczyhQbbExzc0zWVB0vRCS6njJQu39UWTWCDaVRAWBLJEnaAAAAANwGfEHRFf452ZgVPqCEm6Wp074F/bJrDI8pmAHDsFuIJwlvPUbZo3jO+9e0mbvECpBl/WwaROjAAAAAkAZ8SakV/isvpAhCb6whUJeGpnwag9sUF2UycKFAAIxuGur9dAAAAVEGbFEmoQWiZTAjf/AOiKyAj+gnj0a3Gva9CNJ4js1U2x0aCm2qxyr5dY19k1pZuImuv/8peW/h4jD+21kSZLMKNT8K6Qp1NQB1Oz3dayRj8a04cwQAAAFJBmzZJ4QpSZTBREsb/+8PuSK8cB7lWURgB2UuJkbbsVfmYf6uGo0ihSHCGRIUeOkGwszI/RafeeJLVFb72xB56AeE7RHwA6KDK2FMZEiHn9cDBAAAAJAGfVWpFf42Vo2qm00Dddqve1Cwh67h8F7hLVO0VUOGCP37hFAAAAEtBm1dJ4Q6JlMCN//vDfCqSMgQ2VRj3pLESWqFJI9xvPW4RklANZ4xZHF71Z6iQq6DzlYFf0+Q/xmZ/fOpxmMDAoxgM2EezYMdxGYAAAAAyQZt4SeEPJlMCN//74DjgQ7ZH1LvBkl9O69N9KO37jJz7QEaANfGvJaL8Blb2Um8UToEAAACFQZuZSeEPJlMCN//9Y1z/dD7lReDd/jih+Ld0YtRndvMbeeMTxfw1CN28oAhAs3BFnwK+WZTwEg2f//cT8FjRJI69DXHiU0hpAJ1L4ZcWUkQNn9/F3WTm780WFIBn83SpljWZ8Vra05F6qUNyjAMjVTOd4yHsZRD4BpeH6ERAEEj9uWg3dAAAAEhBm7tJ4Q8mUwURPG/74Djgdzj87YkIIqfbdVCHiXmoEJTcLeAuy7kDg1f7oEo7JIcrQiC/U+8t/yN57HX2mDX0MQUw+qGQfYAAAAApAZ/aakV/jnhoov57JnRJzUQ582/A+KqjcuQgc99QQQxsklKC4RUIecEAAAAxQZvcSeEPJlMCN//74DjgiVvqHRYbDnew3k/WoU8O62m6cW1f4z4f2KkM/6uwtCkoygAAAGNBm/5J4Q8mUwURPG/9Y3sbYDUySihk0ZwYHPZc98GT3bkkuNK+kxG9hw6QhCPLCTOPu0cTsYadgUqI051fWxMxsY+qGfVFtfK/8hGLOVkVoaB0SjGYoDGFSY+reuTO1+te1tMAAAAvAZ4dakV/j/S/kR07/7I2ONOQxpj1EA6tGW/N30/w4dgSvirINWao4hRz+KM4mYkAAABMQZoCSeEPJlMCN//8AtC7wRdgH2cIbrCoMs/ZXvDYX8DxXmXaY94n5P78/TKQXUbvgpq5JuPge7/U+LSYd+1eMFPL7ErSRU/SQWbYYAAAADVBniBFETzfgcsI33hBxsPekc3ZDBzYdeklj7D+yHKGhCeLw2I2cmn91BAltpq2zXE35Ex0sQAAABwBnl90RX+JF8eO8NhiMY7e3XCBBxtbicm//Ab6AAAAGQGeQWpFf4cEWNCU5dbW4vKx29yJfhj97H0AAABhQZpDSahBaJlMCN/7x1VfiFwgJBqttUqkagOHWBpqzu3I3Z1ug1iDen9RqMLuWwSJns/FSjmbZC32wpaRGK22xjoTIf8DNdWNE7Xda1dD+nvhdspBUZZ+XrBk1Hqe2E0wKwAAARpBmmdJ4QpSZTAjf8FqjzHg2B+Dqry9XK4oFY4kpye5WzL+gZiytMQXb5+G1o3IxzRZhxW9aZpjLbUsacNdqbjHaYZ9CmpTlK7FK44jbJ1larK+ACOQSr42285xajESkZG6raHieCShaT5/xNHmQa+kW25//2rjWsefjIOVuKFkvGla54JhR+X5yYFqm/bC16PSKR/h99ZwXzxEyHSds+3pr2CMYM14vyO0YKx2HjFF2rVw+hoZHoyQ0JcL/ReGkwZ2fNhkLNreyaduEAjHZAAX9t8IFb0Oz30qFk8O7KgS9ZdqStAL+KUDJIhD3HAAp3xmSs+BqTgOrQ5PS+Ace+fz0i3bKIPCU8i8AJwZcWiUwBhyv7knwrO1Y4AAAABMQZ6FRTRM39foEqx/3JF1qqZd8W//5kxNE6DYERz3YbVOw/1jyVC2gk+UmHuYnrw/mUZrr6ik8rqKlUQKXKQUw5GLsaQyVoSX08YJgAAAACoBnqR0RX+Gfrvv7yJuu0tSB+P/+h4yK0TNhid5Xn9+DciKgmECSbWRRssAAAAtAZ6makV/2fIPxtyYTe2F7AUed/9kq/9NVDnFjpKyZwUbJZWypoABH7rTJteAAAAAnkGaqUmoQWiZTBTxv8gMvylDZaeIJIsMYkC0ZGm9WiRczW/gYZr7/QtVjFGHEuTg4I7VYpSYfgz+C64TnuAQXBvRwRlD7R7BrUKTMlAKgwmktl2YT4IHIivf+PBMSAKaN8pB9Tfiw85AUHSJGSwlVnsXlyHQ+LCtY7tNz+KOQQRHBQtvFemUE7/owFlAZAQX4ZhH7H8hq2bMUPNDHTHxAAAAJQGeyGpFf6zmbbzG6yAo5agB815PT4NIYmF3dRluJ6VKQoZFLbUAAACKQZrLSeEKUmUwUsb/x6mJAIU6OOaejU6jOCEaJq93ye5SDx3rwHIMSwuPbeEPkFOdU1+52yWbGid0yNY2LO00ed66IioKOPe6doVimRketNB1H1k9nKGLy24SyswHs7o8zUFs33/dimS2jZD9vA6qWcVwUdD9loJxKlVTdWWCY05iyBBnRI+HStFBAAAANwGe6mpFf9DFBS+peC/yI0i5BOXZ0Kb0Jtornv1yHk4b9Hwis05VfRcTVsB0+56ooWThOraMfVkAAABWQZrtSeEOiZTBRMb/++FR2EKSC17OilzphaOnJfb+b2v5DMpyB0LE4N5XbFR41vnnYVmX71a1TgMbvI4lkgGITPRwjYOjeSXOpCDC8xaYEVB5J45bl6EAAAAXAZ8MakV/jpI8CL128HDcV88rEWHEY1EAAABRQZsRSeEPJlMCN//7wZu36cWD2WN+3hTaxcsivWm9AgbVBHu/qp/N6ueiOGMMY/MGj/ljaJqYSQxBu7OQessdVRPOqKh231KwV4eUePFzdTULAAAANUGfL0URPN9/4LLaSW8U2BK02dj928jKfghGK1PeCa/SOciAikbQHsE/9VxvCM+EdP4P3D1QAAAAKAGfTnRFf4Z5voDwlwQuL1OkS/asB8/ACxEz3gJmWfQAj0g1qyhfBTkAAAAfAZ9QakV/isvpBhS2mK1fAdwbtVb5YgmyDpirrorcHAAAAEdBm1NJqEFomUwU8b/7w/gmRt4SJHtq/4CnNuN3iEw3xnwnXZMEbxFj861j5ny0cBWKgDehupwGrBYXAeK776eEXeWDNQa8lAAAACYBn3JqRX+NmH8JJKZfXQ2CWI5tPFL8LiOYvVF01NUI0yeIVFZaRQAAADBBm3VJ4QpSZTBSxv/74DjghR/zuyzH8Txd5XruzDOHv9RKkKOyGWsN/oiFcRlf2fEAAAAsAZ+UakV/jZiVZ6ef8q5U0fD2jgZreGAFGFOTGZSMcvYdOVWw2qIqwnqo5s0AAABPQZuXSeEOiZTBRMb/++FR2CH1+D8+1PRw2eC7cPzHLlSr4N5kdbXkTkaLMOtNZYDvN0G+BcnrfD919fNbcNCPlQajdZJesGGRw6UvTyq3mAAAABsBn7ZqRX+O5cRG3XQzJOw4cZkp7na6PmOyn1AAAABTQZu5SeEPJlMFPG/74DjgO8yArnokfcqY+3z7vgO3k81KQrKISGhM7SavcwZDj+Lf/9hb+KMuOUdKKwHApugri91Y0dc7DkP6jYD4KNCu/ZrJqskAAAAnAZ/YakV/hww5FAKjLlfq6Ljf/+1coHIsbogAeQqZ3SIn7qkD8EJiAAAAL0Gb2knhDyZTAjf/++AcgTWtokUyBqD6lhgvnfnJDd8Cx3z95yizbJFdyDlHg9jlAAAAT0Gb+0nhDyZTAjf/+8PuXcpxiwPb1g8GbiEE9IET2gffzB52f8EdNr3u9CTCJQDIJOrkhycCCgNaBG2HkcB5oPyRivqaDJ0+sBXnYINEToEAAABZQZoeSeEPJlMCN//7w+5IrxwHpn4mcnp3qC8i1PB17ucNb58sAdjpq0E8ox6I6nyB4tsTME3vaKjnv2HOUOKh1xJET0cCW0OPsSokk6kSsHBKYZjYC7L8lyAAAAA5QZ48RRE834ISoy161qatNvHcx2P0v+2schVH4rQF94qpR/q56ZD/z6y5NPUUhWN3Zgno8YGf8JZBAAAALAGeXWpFf4/0v5DikvANg1WUUB+NcTU/sSU3fSaYjuWP6ked8ltU06WZ15X9AAAAWEGaX0moQWiZTAjf/AKPrvBH9CwqRiXu1U258fBKSfqjqjBQd2g4Hb33Y/0j/xa9Lykql+WIXGWh5MAUZG4/duV0oRy89osMMauqyiZdhBQPluQQB/BxcR4AAAB9QZphSeEKUmUwURLG//vBRkdJOweuoGYmpH5xFZVLyGMWYq+kIwnpHpVXYdaSzJAWmGTPK1sFf6Znkp4MDLWKTIwxvdVJ52cC1rehXCiJGg1PAy47lCc5h0zQ7J3gBJjV4y3ubnbO8ylb6w+rjLP8+xSx/E5wYWiF5YEZn0kAAAAeAZ6AakV/isvoz+8v1Hb/P0q6S2yOKt6t2F+aO/ZDAAAAfkGahUnhDomUwI3/+8SDUbxwHf2NVhEpR83q6f+XXT0emIrLG/GjTA0ZQEsuf4ilRudQYS+kcaQP6NH/n47oeDxDgr+9I1JylC2mc1pGpx+CT9kp21tNNfIR9M3h5hKBoKhhH2mmDCHWwTlWZrU8XQRkROFjiGHTXpnI2TQPfwAAAE1BnqNFFTzff+Cy2kloXh3IfGPf7xzw0VsS6ephFl5SkLNj4sIWol/tW/UCeEA0PDXCAF0fB2wHJQCr+HvgtkKvByBuJ7pSQYYcBBA06QAAACsBnsJ0RX+GfrSKAZRt3+eUjEP/tYzGZCzSsy+iSHSWdDdvkdj7BDVckq9AAAAAHwGexGpFf4rL6TwoJvceXyel974RIYNEAeX014ZGez0AAABJQZrHSahBaJlMFPG//Xblp7QgISvXq/43tYQswnlCiPCtggpbgu8Gcd0XnodgB573BnhWBSscJCb+QLXOkHexkQ6chMI3P2elzQAAACoBnuZqRX+NmHWKNnyOWkElh6MU5PvSNQ8zujRaeaaC8dcoGksLYZPxasAAAAA1QZrpSeEKUmUwUsb/++A44EJXr12WY/1R9NcvW6Qg+sKr/ahi7wg47WJVdC+uliZS+rcEC7sAAAAwAZ8IakV/jZiVZ8+203fAhmP6SQ+/2wesW6W1rbx/GmEeycqK/2uzkQSaEM7KzV5AAAABAUGbCknhDomUwI3/wSZXK3peGISKDGuEfaX/ruGdURoMWCqwPVNEpOv+ST4slfLn0vfGnqFijtJwvKuTeo4gMwMlZ12WAx97eksVcC/xZyoYY8gEMuJoVZGXraJXsn2AWjQmljY++Yn0UGbvnrgKZ/T7K9u68fYiopDXks1XdSedLDYdCPWc87cOG2LaWaYKA+J/Rcrt93JbbJphn/RaqafF3b1J+VrUGnuLXN10mq9Ji4vhD2gcCVDgU8QP7fngj8k++0v6zvjDx4MM30Unq0Ro37cLz9Fp4dYmyVWi9EBFqKvmSgFrSgSu7qNhL23/SWg/5DKBGi7SSMBZOnX9WLFZAAAAs0GbK0nhDyZTAjf/8pF6MHN8xvhomJb9VTbd1oufHXVFvOUqiH8wtZul84FilJh+DP4Lre0QjSMquDwJB1Cd2xjlRf6f/EfPHC1oSOv7JiC13e/9H5Er1hNYw4FigKpqQ6OH9tzFrQITEYT/f3gAuicOt4g9cvzkdFsMYBTJT4tnA5QEdXJC8nEAXWHq3Dr4PL5tNfHtaPyUqNoDKx0fsKN4D/mNRoRUv/gxlV5JxljuF+HlAAAAnkGbT0nhDyZTAjf/0Y/fwfBPYIt2Ih+oVDWysEi1BfcGsNc42DZ+iX7NLlxL5cxymUJMWBhkKF68OtNiy9cj9+WIlc5Mhajqm9tEk7Q5HF5llNkQ03EIGx1nktYjyHoOPCRnfHGTijsZ7HClYxAmgkK+zsXNRz61pxBrMmmixPZw347116XCKl706m6iOHYY0uNsJzXzWYUyQe9L9FINAAAAMUGfbUURPN+sXvx5/CiDiEacNaS9U6HfVvnuP/a5UQJlGkFDYR2J/hU3JICcol6Lc4EAAAAnAZ+MdEV/jylvx+CPXBU+8gKTE/8VIElasIGY33llIOtASe54iQcmAAAAHQGfjmpFf9C2JxDQM2FrtfcernQzbgQo6cOgAiXIAAAAPEGbkUmoQWiZTBTxv/vFEDSrIEiP+dm2fRrEThE5lxbkgROmjfPtCbfhnR3G4ID1iPeJh7HY3Y16Odh6bQAAACMBn7BqRX+NmH8JJKZfXHhkLwpdR6P4XTfNRE6MatJ7PVwVpgAAADJBm7NJ4QpSZTBSxv/74DjghR/zvAWMRwWuIY0ILoghoK+2U2wGoNOA06nWHBwAcUpagAAAACsBn9JqRX+NmJVnp7PpUa8FkOcbCyJNunfz/M06nD3e8NZpyq2EIfrhbd7BAAAAT0Gb1UnhDomUwUTG//vB61Ydg3N6LfuunSF3AkvNLoQBhoTQ7aoVG6tsYzS+JiPd3mYuiPkzCyWnK401zAoTSSFfgBeqE3OpyvPlsMfUq1EAAAAVAZ/0akV/jX4K9aDUHTkIn8NGrTJxAAAAOEGb+EnhDyZTAjf/++A44DvMh5EHnmI+0qZng7g7n1G6u5/Aqyb8TemoaU8nwXBzvg6L4ceEadnwAAAALUGeFkURPN+D+24e0TinWPeueGoG4sapy+ObB3xDR00TdCay97Igcx39n3zDVgAAAB0BnjdqRX+HBFlCUE3nRNtGVVInyej3t3vz7uZ2wQAAAEtBmjpJqEFomUwU8b/8IvJDICKyxCcdzT/mboi+flPNWjFbInDCJbwHXv4CBeLXeVVCSD4rO9Def9jVQKiZ3a5b5Bie9xheg2SLH+QAAAAvAZ5ZakV/jY0N0MOVPt8h6WcoZV1S20MbM4DsMfLQUbNE8rolkKvz6B5iwJ8mvWAAAABXQZpcSeEKUmUwUsb//CJ/zICOuOx4SMR3pARtPhD3XWj3rBkO6OrvomzLoOGgKYE4i0wu3UHRLX77PtaCBTI4Tfpc3gtTPNcmsPDp9zCw//wDxv5JJCJZAAAAJwGee2pFf4x17dolDZL3wxmEFIgT8nhorh3Yq+I4oAacvjWdHciZNQAAADVBmn1J4Q6JlMCN//vuVw4Hc5yyvTItaaZzkkkNHLJkI0HEoJdIeD3dJbzkXmvtq7jUQZd45QAAADlBmp5J4Q8mUwI3//vgOOCIxgOqHGQ4dPM9u73R4ALKr9BQr/nU61eDXxry868wQq2NQQB2Cw7y8YEAAABaQZq/SeEPJlMCN//74VHYNYZ3DjtyB3KfJ/ehNVbvQ16Epn16L7ahC4CDOiOWjEm08H+S+lGobDN1/4ONf5WIIOky9RhRHLV6H2eUZ80HGf0olbdDoIPuaghAAAAATUGawUnhDyZTBRE8b/vgOOB3Ocvj1S4/IKY+GUDepFZA7JQ6D5zJgHdmU0gV4AyG3PnUUlwhcyp1sj/2sUr8E80bHPfsNfQw3EPqhkH3AAAALAGe4GpFf4cEPpZ+ayGHCrzmIFkFJJ5iOG3d7ko318enKcteEWyK6vDU0azwAAAAM0Ga4knhDyZTAjf/++A44Ilb6dcZ4GoPqU601qIovFnKC7lkIeBkBzqp//sZLPSn67pyKQAAAExBmwRJ4Q8mUwURPG/8A6IrIENlQqhLiwoaJhrnWZk18jQ7ylRVYz4q3Y43ERkjNUVHICbJU3W2fPPQUmUyKqzovlWUqaoQrc8y075BAAAAKgGfI2pFf42NDdDcJbtDdK7SWJH5HwS2s2t/lM7DAAfmUTKKRr+T4pF7/AAAAEZBmydJ4Q8mUwI3//wC0LvCDPTwaCjaXZ2btf9j1UTqkuJrlFvJUAop1+VQkGMNp2GtkRPC5V/i7vZhSjenabTSd4ZI6Z77AAAAHkGfRUURPN99TaObewep0G/CenyT8ihy+yfmI16uhQAAAC8Bn2ZqRX+NmJ1nz4vAdigsChyYtloC5NV8EcFNRJflhn+2noyyo44OPhMN8LgKTAAAAIpBm2lJqEFomUwU8b/7wetRFII6Zm6QidC0oo+c0Ns7riq1h4ZBWxZZD6JzeHG5pZuCUPFJCaRVGlMRP6GhKFXz9x3DH2KAgdWA1aMZo1pkrtviY17lc9l1KlNpbFrOjLf/sY3BYOhEQqMpFO89k6PKn8sRKcq75K2ccdT4r+p6kGpG4mZrOF/98uEAAAAZAZ+IakV/jX4K9aBpZl6QE2eOe/h3YVFcgAAAAEtBm4xJ4QpSZTAjf/vgOOB3OPztiP2ULSp7XrKrWw3l4Bc6H2W69Q6oSoF8O1hPrRXjY7we9rAJKPYETi17uxoi3v25T98McjvFdpkAAAA0QZ+qRTRM331NqbcLGtWbMNY3WScBZQDj4EVhJKd+eBPd5akdHADpGT8hVb2bsA3vvcUMnQAAADMBn8tqRX+HBFjWKW0eRGEWy1dqsM/Zh2+r0j8TMNGs/fPMii2zXF0N6KhYyAEYzOLy8UEAAADpQZvQSahBaJlMCN/HqbEHRtdjNTKn+1Ks8Fd7U1CyW72yxoAL49N0mQ/rxd6hIp3QUmPFE+CBT6W6w2fg+GvkDoA5rxlfTH5iWO7O7X1aIQrUycIpZ26bWxdE69CfvW5Ump7bO49JM3Ojp1ba+sFXZJgF2H/6c/uab0RGBZJFVPWM6/2iD4z0L1ee78rMROduc0TjmUMv4F5e8iqzbVmWZxa47HR2pX34hnX7bqBD1Y05md2DIEIMx7XWDS2zFGguo8fbZRpICoATJ0iBpYXY65NY+onyXag/KvGTn9dqOkIZN5bgNwodskEAAAAwQZ/uRREs39BT+V5BdRxLLjvRBZ/8QK0LAnU4ryUT3nxiHiAhHRZibTVIPm+0xrfhAAAAbAGeDXRFf9BX5Uw8x8UvZrjTOa6N0KNc735HzeH8dCL7LysfJAtHc6CJl64MM87h+XsJUfWPW3okcUNQiLbRiB7rS8Stpkk7rAmn2NBtaNQkEAWDZlBKfeymcqvzqnPiuE+on+yiIr3BB71tTAAAACYBng9qRX+NxKann1KePz6kv+QI+9/UjpccvAs8AqKrEhiFftqoFwAAADRBmhFJqEFsmUwI38oIRG/MCA7nH4jDF8og0z98ssIrvm0ivszdpDx+Nnn5Qb+Dz0hEkbmcAAAAj0GaM0nhClJlMFFSxv/HrNHv/cjBycDl2vih6keQw56fCiMs0YdNmpp7imPBF1y/lLKBnxMxn3M9E+3yBzgA4qoPmgxAa8dNdml0VORbm0KUx6CwEuL013ZY0JgiO+ELoSSDqSlaWOxREn4lWygn7nk/erfNPMDpvIC2jYz5OlVHNcLkNnKjrSDpnPb4eSDAAAAAGQGeUmpFf9DEtYTuNuuc8MFIfiTupN1BdFcAAABQQZpXSeEOiZTAjf/5KwLantfVy3VySM0SMHvsr0YmBrv0UL6DlEtaxQQHzUXE30YsdNChYm2O6V+hEk4F96/91dLwu/ndfGK7GATkcf47DIEAAAAxQZ51RRU833/gstpJbxTdFgmzsfvEjjhs3dxWp4ZZ4kPG8NZeKU7GYzN2OvKw1tlnqQAAACUBnpR0RX+NmIM6D8IxZFFGEALyP/EQHxDjOvu6mfZR+X/XdtvYAAAAHgGelmpFf4rL6VHGE31hHDEGQjtVce1yEtA1rtWqwAAAAE5BmplJqEFomUwU8b/77lcOCFH/OzbPo1iJwiLqg8angbnwoKqYau4/Tr0/Su9hg7+dXmJB7llgNXjWmAvP+6nn9ZyA2iVrXntcHxIEctUAAAAkAZ64akV/jZh1ijZ8jl2T9B54JFoOBhttQMJEZVln9G5lKq8BAAAFQmWIhAAn/74fcEJGW81yyeQ3yjIJuOQipq5DurWEjnW3UorRPL+Gmp7W/eH/ZijezhQ2AIPIQl4aHdQM9A6M5/MKi9xGlN69wV4R5VA4TQl0Sllg8188tv6vFr5S+P9cmRKtZwlxsQrO/PUH7ZGWYYybHVXFlEMY6Dl5808LYw/EiDr5JaoSWsHu4ebowTZo7v7qaSld+s+T31N91+CVDr6kktqJbNFOPAo3BvxQkltR4XwG+/rxvXbxkWnnZ4k7z5WpzDWYcBFPvUfTgrd7n6SLHLTp7bJf8LFMRC7GnoqmMyaPAzg75u8JidsVPIR29wsvpk/Mt9R6nnF5IJistydnX1nYc673c2swNm0ubNLgupd+k5ttRQ152rHLvBca63qjaEYz5bO6uRdZGmNbrdIR9DOFJYzEJ3T0VRVhLzRbgQhmzINEkMybfXp29t1bR5qya1mRWgLdBC5l3drBLQCcVoq19LY64FqmtYeURGFQ+dxCLISD69Irv9a0uyOxS7lVS3qeo5/gPB9YhIYgKQn6bpUQv8jq9Ep36A1DYX5VFcN9XqGkifkeY7ZbzJGH1tSe7aDoRQIiDO6HnmttmwDbytG0SzChvNEtecWmFp6+Ng7jqZmJ+LpXMqaVGcBbl1ehkdm5Vr84GAQKEq97RhEnaEuoPQsiUhi7y5CzHQA/q9/VpbR1zqAOP/78KNEfdl56fss7Y4UfTIZBj/Ef7ScMR78inLAl3j2HLP1BrnE2c1Zfp3fmi51t7IntWomhhRLlKfnfJrB056GmNbnqvtmZCvlymbrND2EqUvjJ1EJ+n02dodaM+1EsP8F2MT2Y4tjt1aBsu9gTs10CfVny7PzZrIe7ZNf33r8LTwpXSowlA4nQt9opFQ/ZJxaz8Upq6XgXe6Uqzu7x6Z4G+LEevAbdMu5y5ewVRMuV7HqXwlZX3FmRKJTwfg2u/onfLxKJGGPlejg4w6xVcIvQWpCCY94rqqRweESq52TT21mJ3ZWZPYd/bW9E4HoKBWndnm57+UuTSHIPQ35BT9W1Xnl9r9CViMvIQ8AUlOp7+ZZmJtCkSP14cMcCJhNiLGBPwta/Oyraa2eYe4xSkWML2OCXs+aPKBwLcn9R4ao5D1S55sxD7n6+Mg+z9jwufgfWX3u9ii7EiGIFM7wV1HcRH3cSIQZVgHx3bNHeN9HrMdz5R4HKc+XWQ6nfPa5qVL+jrteHlNOWPkjRWBZgYim28+DOXpSL63VQKM8wCalx/7UFzpmPLYi2k25Tn4L4RY2kTpMyq852ckgqQ/JuHNjmCM0OZxxusMdxMp9ba0q6z4Fov+fjMAzn8vKAAJKzFCfdFJg2MeZhfwlKaitFkiyEY8mnat72LBWhX8y2pFqa9wBfrk6vSpsSIPf1ItXOV5dsbkmQaOOvVHpr39w9kGJ06MPgDiO7B/qm2Bs3AlNEuL+WG8zu/vIRJdDz+RTRIk7vLQr0Rq8jEQWZ/euf8FD6BUNHSXZGM949sWmTK0eoMcmoTzlyGGCQgdVdRYqeVPDF7SHBzjIXvdTeaPeF8rdbHCdMCUN23hsQBtwTi/WxUcz0MCULeSfoCuRD9EomV9UBsyxfCUqSYYFieuAbg/IP3y2xj+T/9E9kqvzPRPQa2y1WxDCMsnMIy37FnZ/7LXXOsxfbVeew3U1YASM9XkT/y50YD3lR6WGJ8MYfYM98D8exb8s0epkXWDnGeNbeKqayRpRsYTvsgzojgHxpi5+wGkEyHCLHk/j2PjOYfUfWmKhOwj053UqWFuvky53Z+AuX9yWhjuzKAAAAOkGaIWxG//vEnVTkBINY1xBnYqV2GYveJxG6pO1ytvZdPbEXt5LdA7TGVX0iYnM5z2Z+J5AJNCh79n0AAACQQZpDPCGTKYRvygg25A/mod1FrlSV9Bg4fSRIk5ApsXKnR4Lw1hYENtNp2IdrhL36xanS2cwpV3ghaBQ0ENK/+61iJnFzrglvYgyAeUPfvh4+H91OONFJ0BQNQP7S+T1dX3IH4RbbT1dKFwTvZ3ZM7M9HZq2WxWBFVnr9ZP/zxD+foUgnXU6TFFBNN/llQnxAAAAAGQGeYmpFf5XT+htGpxm+ttzhrI2RGOltiWEAAABzQZplSeEPJlMFPG9uJuUZbUPQJIkFcwKAtahG4ohsmW9O7pJ3QXEf2moUd4qKRCCL0ZKGjcHbIVUkGfh0CDhI2ibslmbDuzOMAhXG/zv/6v/5h1pF6Fnk4OU0BdEQ7aG8XM909Bk3XHm49yHVcqbBFgXZxQAAACoBnoRqRX+RRVaIVD1YrnHAZf+m1ih7YZmVdEVTUohPyoJ0YoX7wwFZW8AAAAAvQZqGSeEPJlMCN//74ByBNa2iRTIKCh6OJLaoh8ps3rY7B+aDjQEslblmkDmrpFMAAABYQZqnSeEPJlMCN//7/0a5QmmqHJOeKUzB2eNo5H9glW5oLygWzEVYpridRvXn+2AASXaouZZC4r0hn95olSpqHwrE6msLwa+V4Ttq+yvbjDu1Y6/bkwQegAAAAGFBmspJ4Q8mUwI3/8oIQPYmOP/wRrPxM5PH+fJ0mFWAYBFh1Pe6xJT8//aEn0KlIlhE0/6putAFmsffwcLGQqaabW0OMe6iwCBaM3Kq2If4TnEQ/lSd0QI6JJrQELGDib/jAAAAO0Ge6EURPN+Ecn5eBysgKBkhDY8qsL/25Qam4wDRxrYU5Di/5y+wAXkIetqyGrdnqZuLCcwFByy7VklxAAAAKAGfCWpFf42Vo24OWiO+wC6Bf1+vC57IXLwDKAPdiUIokWez79lHX4AAAABIQZsLSahBaJlMCN/7w6dR5vCGznNtXd5bO5Zy4LxzmfFyTxf8VH6vp60ELz2qk/FRPVhaQsbXkKgrOqn/+1roIltiQuL600R9AAAAakGbLUnhClJlMFESxv/74VOAECNfo3TrZkzqu0KkrOZ9CVw2gjMRNMj1E9vxW5V1Ecvm/9qH28Rc3ffR5r+zqVMsLdHEwsswCE/8g29m7hFIyoSbEObNEnebhqoXvPeZCDCweyiSrLwsRskAAAAXAZ9MakV/jpI8CNGNvrae4VovZPXKoXMAAABnQZtRSeEOiZTAjf/KCEDu+rkuXwR0ltpBueVY4aihZ6XD3gGzDCyVqxs5vXh7SVn70nLyVRSeCD7n1+k76tXWOU8Q8cvlYby2Hnvldl8uqwEjayKEwX50N3/0maF6ny5kGq0f7Ao02QAAADxBn29FFTzff+Dbsn0Iut+RPjHv9yMXPXi4QWBtUuz84rCf2ne6MMkRmFjCTOGYGkE76dyp3SNBzeilrswAAAApAZ+OdEV/jZiDOWM1Uy9Vu0nWL0V/4HEe7fSSKWekx5hQYtiC/xgqvYEAAAAjAZ+QakV/isvo1lhN9Vf9YXwem/sU/GB4tK3zn/fFUPpEfgUAAABJQZuTSahBaJlMFPG//APDMcCEr16wAPFV1G7zGhko0fBO1LBEbSo51CNZnpYS3MHF50L4MAKNELOJvnOk8fJx9M+aEnAehC4L4QAAADcBn7JqRX+HDf42NiA1e6IoNY9dd0+pse3XDs23QkJmYymj6mByBQBlbsNBH/n+C/Zf0yypZhu9AAAANUGbtUnhClJlMFLG//vgOOBCV69dm+FVLfTWcUGQr+lSqCflwO8IOO1iVXQvrpYmUcqC/Y4LAAAALQGf1GpFf48rV2va6Z6CX7HZbnuYgcin2b1jZsuvhj8zrWfr2ZgMZiNJsrNVEAAAARJBm9dJ4Q6JlMFExv/BJu5KzDErXKYS1RifWXWjNyUkZ7LuAhIUnAIaOsVEMMQXit7waW63wzdcUIeRyXc0ns6O4xhmDwZ3ukGLdi9EU7w+8z26SXTja3sGWB+uU2PQaFr9su4rpTbgme55GdSV5Hium1b/ZY7wn0ziN7L7ptbE80Ax28UKDugysLJcSwfJeTU+FsCjXuJC35Kp+sq+nt2e0C5pikCgY4m4RGYJRGW+kGDU+yhW9YEtPxxAt90wlj4Nc2+02VtndN2Z1aSOvivKWQperyfJK+RiYH2Ua4M1Mbvt+F/gIq/S9qIVS9LvpsrQ+7+gxnXyJvs7t2yv/J7ljPiE5MYPyl8h+jNIjrv8R+MhAAAAUwGf9mpFf9Ei2Luv/RNJcaZxTnDa10UNbuQPAlqODVB2VxsIpmJnfjPkMJPHAq7sCkJAYLQ9TOcWdYY/w4T9TToTqOJl5XsxpaKH0nZtSdWfNN6EAAAAUEGb+0nhDyZTAjf/+8NYd2niwO+rMxV+zWAA+1JtHTaumj+P2U+PhhfgTzByUQbNihyr6FdaXSQST2LLyoRUr2SQPrdTZkfv31QscDSS0oRnAAAAM0GeGUURPN+E8DKIQ4/14uOVqEAMADp0zauP0LtNNslqo9KyDu+ewIdMlXYKjxCCIZ2RgQAAACYBnjh0RX+NmIM6D/Z46G0zEmw183/eHh4odUNrt8tMZebEjulWUwAAABwBnjpqRX+Ky+k78Xe04HkuIKPhhbbazcKxG4nlAAAAi0GaPUmoQWiZTBTxv8epXycoOFMLM/d1FqVzTzgCDwFATJBP1U6ge28/eu//0BsTgIFsoE4M0gMczwnHzJtkURt64dKffbzBHukviZ6wJ1DD10/gHeIjw54CCzmpIwySvy0eKrn1sxtgjxrAjk40cMg2w3WH2SwB2gwRBdRsg0MrzvNgn9GcNQO1ZTQAAAAnAZ5cakV/0J4+WxsQPcekdWueGIREStNdZlzFmXh0KEniThU6N9wlAAAASUGaQUnhClJlMCN/++FR2ESuhE+RatyVDTTzS0MEydH3+Vpq2cFMFPi5QKgi9czxuhaGemqsD7/JcIUhVdt886og+ljZbrvCIWEAAAAiQZ5/RTRM34OpqLzLMBpZQ293KhCxlhbjxhPOxHd4q2TqwAAAACsBnp50RX+NmIKsD88pu+EZBTc434H+uuuvMbHJAB0tuUvJ1rP2bn0fEH47AAAAFAGegGpFf41+CvWgaWHE5ho1aZOBAAAAOkGahEmoQWiZTAjf++A44DvMfnQGvr5aabtMPFRgZrO0PAPTbXug3OYp1ediQc1+3GfuUR2A+YYULrgAAAAtQZ6iRREs331NqcrGn20AwIJmTxeb5XHCMknUBifcKi/SGXaUJ0L8hQfqmM7vAAAAHAGew2pFf4cEWNYpbR5EYMikZF8nmrBIkdnpZdsAAABOQZrGSahBbJlMFExv/AOiKyAj+gnj0a3Gp6zI2CjNAAJIWgyCThvil0kvhxqkZ6Ht1BxxqXfrjYLInpyi9Q1SHRusH5zdhyC+Nu7rF+RwAAAALQGe5WpFf42NDdDDlJ55Cx+Tl+h4NevQbGGML2wsKXdErYWj9A80AousAs0tNwAAAFhBmuhJ4QpSZTBSxv/7xrXt8jlCNDUuxflUiLTs/g36A9cpqCK92W6XZsUT9pVvqv4iI3XXl/siKAd7P64Kbe2lHxIRZwTXchYXnFFv6PVahPdtBRwdIEdgAAAAJwGfB2pFf4x17dokp5DlQDbrpLgfa2hCiFMmKnp8wA05fGs6beVH/wAAADVBmwlJ4Q6JlMCN//wDwzHAd5j8QdbEjP/MD54/nxaJsM3etRkV7PtoitdWOIiK4oua7/OJmQAAADRBmypJ4Q8mUwI3//vgOOBDtkfUu7/INxPIL+S5L2s8+0BGgDXxrzOfaBdvcLkOu9kv7y8YAAAAhUGbTUnhDyZTAjf//CIBv4QEgt0jB6sYApnNKebF3UduVNlxp8UDV72R8VAklM98n5lOWEw+T8M4WRe//ioMQECXsBcH5k+vJghjfT5jqlmoeL/l20XeVcwcUIuafvxnJkzhBlcn/d9IuK1j3bCY8QS0oTVT1N2fu8R5Tv1RgAbc5tBoWuAAAAAhQZ9rRRE834H4ctY9gURT1dXKSAIxLkivCpURU347a3c3AAAALwGfjGpFf42YllmapQL/etyKspbtOBGV8f4EVMMX0hi4K6l6e+YACOn1K36gQd74AAAAMkGbjkmoQWiZTAjf++A44EPraJILnXuLxMWdODvz9zlGHYLFuS/WCzov37vtypupywZRAAAASkGbsEnhClJlMFESxv/77lcOA9jsDFR8Xdn5jdZFKfM583oBslJ49SZxnRGp+moI9oSjotvQRWZD8gOf558xXUMpWRXJ/cWn7ZTAAAAALwGfz2pFf42YnyFLdokqnTK0mFdGbVl5+BNzXz+jPJsmLNM19mnectkM0k+LwyFJAAAAPkGb00nhDomUwI3//ALQu8EXYDhpIvttv8TbZDOP8OLBMAO70WexTFz5+cnteMdktY7P65A4tZlIDdklK0GBAAAAH0Gf8UUVPN99TaObewetr3iFfIM29vQ4/CELIo5doCEAAAAtAZ4SakV/hw3+jEmnmFEtcwfltCMFzrX9aTK9FUCy1cKpO11Befk4O1gAKZFBAAAAe0GaFUmoQWiZTBTxv/vBGT5YQIcUNUZpARyxqc+rS5mmS+x2rMYuZzxaaNzjVUhTcMNlid1G47Ff7Gber1siDAOxLdQ+mOx+lSM//HIlVdSI36FhpMg3eSscveUe/EzN0NXsxq+B/YFxJBXKltyUDxpD4j5KmVec/I5HTgAAABoBnjRqRX+O5cRFOX/rc19dAGfoPBoZZ9JRwAAAATFBmjlJ4QpSZTAjf9RUJUwmPx6d1qkD5IiUog3fdDYG+lhD6PhB6357xt0eAhnE79OehpzS9hNg+0rld2gpMfPN78NladMHIMJnu+IgwcdY5qjgIx0rHEFP382Z+0LcY8oZUMOK+1SYe1NyN3/A/j9w4eoQ3PNSo4Ux/LOL3AlYQIiar3GTf0B95Y8t+So3rBxsftentdD21ZJ7hAbvgAAT3Nv7jjPdzJ/O1zUUPXbzz+p7nvdFPTkoO1Dqocj0bjwXKlp/xDJw9QeJ9FCKJkVoMOYhplLX41md7oQ/Y+qGl5+5p8L8BxRb1OTUS0Pldi3oqvL9d/Zrz+oryE/RRst8Em/0e7bxnWHG6Mx81D5GpO0jZK4XqbwX+2UQU4boZy2cX+Sjuvn0MsMkR236Z2K1zQAAAEVBnldFNEzf1+oL/lUB43iexxWj6oMO+zCK0ceAc+a/F/PKiFsvR2BDrZ2PrBeRSYr2NA9GNN54oqHc2LpSV/DC5yGHm8AAAAAnAZ52dEV/jZiDOV4WjRVuqfuRLr/L1bWuM4bOVk+yHw8xIsSbICVWAAAAKQGeeGpFf95p6Ti9RhN9YQqEMUYPvg2cqLZ+fsaBkkhn7KTT2jvNMkWBAAAAgEGae0moQWiZTBTxv/KSaxhdVGtNZuOiRf+JqEcNFkbAv/3toQnjkNAaWg4K8+iRR1zNbt1yIR3tJg8Hh+JLp9Fo/JoyYU3FI5nF35wWUzt3yx1MISi809oOoCzZrj//fWbdhPS2477QVN0UR6uqQ2+uEBXxub0LmMRP6aOCVtOBAAAAJwGemmpFf6zfg2l3Ygdwg0DWPIpQmdHx99XujRSHBNtWXm69NnlO/wAAAEJBmp1J4QpSZTBSxv/74DjghR/zgYSVgQLW28saHEGZfNPuk6gQX9d7Fl5UDqoN11nbQFj7zXG9wSS5vhSv3gAKUUAAAAAwAZ68akV/jZiVZ8+10ObdgxOx3183mwEhK7HkZqcvSQ1ycqtgp8jF+U3zQBWvpsMwAAAAU0Gav0nhDomUwUTG//vB61Ydg9A/wlrkMlWCPDd7QTOqj1/1NiCSa8ruC8UBK8w3Ph2mInOfYuW+cn88B5LL/toEaKTUnUAV8fIMLomlaMvI1rt7AAAAFgGe3mpFf46SPAi9fOZie0tFVOsqcMAAAACWQZrCSeEPJlMCN//HqYcZuGVy+Xdov07HU6KX5bS+1WmnLQvbAvY9UlOzUvt+HiM8nLp39Z11G9J72MQOlYTjkOf+SeXNLe5vidAHlSjT/Ore0st82H8qHv9wvch4z033oRI1qboMoq9jVzWQ3z/b49jZCWWEEv8uwxf9swldbmArV1EwuuF91pmeCP9WmkCZPwED0CvAAAAALUGe4EURPN/QWwldz0TioUzxJwTL1DjVPQjAq9Y54uZHmIdkd8NkQOEM7KiMsQAAAB4BnwFqRX+HBFltjCbzyK+FRKnf/mLOjOcjzR76r8AAAABFQZsGSahBaJlMCN/9YpHGMYUb5j/NfcgCQnnnUME/+uYEFab7R5PJRwSPzdjTXR8fgJDcETNoWNTQyIj4LAz2Alc7KxLVAAAALkGfJEURLN+CHmGW8iX1nC8et2f/B+8KsVycqqolFcdZbi0GJHLqpjMjlxYW8t0AAAAsAZ9DdEV/jZiI7Ev5h6E6AXxVvyU+Xec7Xzf+Ujz4nooqJFL612j/wKfuy7gAAAAiAZ9FakV/jZXGc7dkFSNTQS86kNC5UhVk5V114XqW+n6aNwAAAC1Bm0dJqEFsmUwI3/vuVw4DvMd9xLBIuaZ/zxsUAT+oe+xkeDynj6QmOuSirEAAAABtQZtpSeEKUmUwUVLG//vhUdgbBTvAngor8pfCekIPa/E8Y+CM1csCM1vgERqhAjVxRBDaVo8NX0IcDkRQM7Fc+VrjXHrEEHSAqDVNvfJe+IHgsFDNZNns8smESREry41It9enn5zKZn+/bjH+QQAAABgBn4hqRX+Nfgr1oNIjL185KPI+d0MY3WEAAABQQZuLSeEOiZTBRMb/++A44DvMfbdU1ETVNwuYBfSvlXeOFbA6GbqsdQXHZrecjnlZBPP40eDLe0BA3X8hhdWV4c0xWlDkOf5flpvIBUQRHYAAAAAoAZ+qakV/jnhriwfjEBC0Cyoph0KvuuH/nOiAglQQP/F8Pzo5XgfuSAAAAEBBm6xJ4Q8mUwI3//wp75kBPbfTrjPAkRSSgH3g76N+IPxHc4xKee632cn0Is7nuf/sPmuSHgjxKqRUscLPVxNhAAAAWkGbzUnhDyZTAjf/+8Gbu08WB32S1i5TV1gbt9xeW+dzaP8gfVQwdBZ2FJC6dNZtf8f0Hg+4rzvoWkTdTJuDVTqcnkDXxrzOGZQBT4SetMErEEP9qYQ3b8ROgAAAAF9Bm/BJ4Q8mUwI3//vD7kivHAemh1CS8p9HXZNoOrlrUrukxAJgTBabe3Uq+xKUtXAoE33EFB1Jjy8zf4fiuWPtUn1WH8f/CY8QlZPFmMmcQSG4fjuXN/B+qNi/1vyFEQAAADVBng5FETzfhPBPL0/8xgE055y0cUfkH7gbD0hvkO9CUgUIYF9ILO4onW0KnZIlS4Q5+PTxIAAAACgBni9qRX+NlaNuljs9U7M/SW9U7+gZciTVHqQBZINn/LlwAI3T67KhAAAAQ0GaMUmoQWiZTAjf+8GQRIyA+EC0i53Uv0awLgBrmCYaTtddMfGzDz0tG737dau/+uhw5YMVYL7inRfSxfcFiKBNgykAAABfQZpTSeEKUmUwURLG//vhUdg1hne5jlt8Hnsw/yzG9BlYA+2BCvdNQCaqq/vU4NJus/26Yiqay8cHVINjOHjrvNVNNQMyKgANVegEHkn0Atv5BheYp5J6LroUHacyJH0AAAAaAZ5yakV/isvpAo+X7EhbLzLMVxMoH0gf/oEAAABsQZp3SeEOiZTAjf/7wZu7TxYO5AYsHfwK5RdvkSVEEn21E6ZLxQXmV2KUm/dt5YseWbLYmF6cycseQ/zRHyWaDXSPfsWYc9/L5FXY8bkH+03r0tQyCwxscnMoNSAc8JSQ1LuDhPTTMWEyb+wwAAAAQEGelUUVPN+B+HGW0ktTlKD5U7sftibyzFU0itTwzBOsr8jinvAm3X+xweVleORM4VRCN4J1niSdrhdyNeGimYAAAAAoAZ60dEV/hn60g6kYGvRqVwcf/9DU0cnE/S71LZRHgHSKCWLRPYzu+wAAAB8BnrZqRX+Ky+kGFLaTu3V1h5vNMrHepiZHt++d3TYwAAAAUkGauUmoQWiZTBTxv/wIE8MgSI/5zkdt2xF39AwC2NCzH/ix26Qf9Z2is2kaMgsh4st50ymHALuUDhcWz8HLl5Mv7cFHUPup58bPZYY2bD/A0L4AAAAqAZ7YakV/jZh1ijZ8jlpBJYejDv1xMTNJL8wwhvpmAwdAZHfqRUgGccktAAAAwEGa3UnhClJlMCN/wWJmEnMThj01oz/8HmHqSy1C8Xnv2nkeHB2q5boSEX7n9tutZzT+CbMC2woNyw43AypzG01uaeEWSh0sldqT/FFT7HItLnN5/sxDby4ws3/+X5ukIG6x07NMTkKBgu9pujusd2uIIBTzr4TtqDAIoj5L/81usPjC+yPrZ6BqtdjHQj00WB66maxuzdpr2LURBzdggU7/6dRUoibhHrHu27B3o2Nx4NpnulOOXMgFyxfruOZowQAAACxBnvtFNEzf1+r+1FYhZh13lr8fS3v18bH9CE4mbV7h417LGGti+gpJ8stneQAAADEBnxp0RX+NmIKr+C9nWHyDrxwMHXm6plPWpdYz9mPCwis05VbBSSZGa7pefK68Un2SAAAAYAGfHGpFf9E0GDu+ktvFMGyuyP+/q8a08x7N3S8hze5kbadJPQ4J/VcgfzmVJ+wGbfISiHM3iOG4kNfnUv4V+1AHLRdp1XpMOeQfEcujnC8EllvBq6VYdpBNSXI+IGkPIgAAAE1BmwBJqEFomUwI3/kqHHwSjggzj87YkHr4a87jSRladhRJ0sPf1vzUl8LaQuieyIjiu6qDqBhKWXMid6Xg8mRn7LGXUL7awd9GqdtR0QAAADBBnz5FESzfg/tuGYHKjTB5ZIF5DlRZ/mG18mLNwoqUKDqUjlTWk6UxETTkhiQI1o0AAAAmAZ9fakV/hwRZbYwm8UV+fNHe+WvQ7LE1dhsl0Kkl2dgkJM2xQPQAAACVQZtESahBbJlMCN/HqjpQil7yyLFLFu10qfJCouDtF1Xmx2Qi5MZ8wiTfvWELZh8dDesXTIOc3puWHO2V2s4k4HdAKOR4ngqVZUA0LYM5DJacVR8U9Naj7I7kVfa6bltDE6u1WaZn72qoZRJsNFn5BmTYbmEO/mI8pJzVOUzeRtqmdX+yIakKz1bUKdoGm4/J27cXeasAAAAxQZ9iRRUs36xetRXUJJO0H33aYF1R/wgPQCJEkQv6Oc57u5yk1XK8xotcqZMjal7waAAAADIBn4F0RX+P9L+RFJf/ZGxnwlZxxRpnlAdZ4MTNBABowPPX1SEIZbQRP1HKfV32CEsflwAAACUBn4NqRX/Z96jvz6lPIdxMK5ZxEfdrk4JnwEbzrXk8OtLZKf2BAAAAN0GbhkmoQWyZTBRMb/vuVw4DvMhdaXt49Ohno7dC9VWUVXFssBv/tAEi1GxHEtwjRSRy8OIHuSYAAAAXAZ+lakV/hwRY0JTl1qnyRaLRk+dkPKEAAABqQZunSeEKUmUwI3/9Y3sZBuQxneh/rGW4lE9nuaCObjg+K6/Ef2UAcSMCaL0//e80Z5sMPf/FLcue0CI5qB5M2c2QWFN1gHozimkN8YvezsbDlRHLpDuEB8yJGwaUJGSNmiArYmFEwpDpPAAAAD1Bm8pJ4Q6JlMCN//1ikeX7HlErA1F/NBuaVPRthECIky6bnR/zbhvjqszzH8XsGKtqLqtEjwRPCnmhPsXXAAAAMEGf6EURPN+D+24ZgcqNMHlkgXkOV4kgZTIIkxduFFShDLukcqa0nSuPAX02qo7ZcQAAAB0BnglqRX+HBFltjCbxRX580d76sSrub4Ej3N0tuwAAAEtBmgxJqEFomUwU8b/8CAhiLA9lUJ1ZhKmJlcycBCqNW82vUrKjwfRLn3brEJAyFEiDnxP6vgG4feHLTjXsLOkn44268wtNImWMf5AAAAAtAZ4rakV/jYz8hJwlu0SinQC47vIeaSqHA97ANM9n5sjiQxvenF11Aous9wBFAAAAUkGaLknhClJlMFLG//vDoZtU5Qj04CcZokNlVJpbn+xpxn5vnomGzZJoOL36cb95LTn2GQWez/KCm44eCBSdizIJg/0bpdMf7OAdSKNj8pTKjsAAAAAoAZ5NakV/jHXt2iSnj8+pK5UPI/1KAk39osMlGVfJu+lmZ0egJ6vDwQAAADRBmk9J4Q6JlMCN//wDwzHAd5j8NJF8og0066g+rXTd5exDKM2MCnecCF4E6scY10T/iZu4AAAASEGacEnhDyZTAjf//WN7F6EssGLL9aRq9GMDYogP8Xzv0eGCplNOy5xCXRSUZ3whgiU8wFDYxIZ/nwDS8NdljQdsca56LkdmQQAAAHxBmpNJ4Q8mUwI3//vrhHYHuWOrKSbgXljf8SHjA7pt3GOVNXvZVrQnA9aKMhM1s43sy2s+zopKgZQuR3OXh19uu03zK43qCzT1h20dR/R17wlmMK6XGC4UZ+neex19jtPD/686wWjM+ofcRmWrGJuZFnTOmBpsxXPMq7oNAAAAHUGesUURPN+B+FaPuhgDkzcf9gPd1YURPpDElZWBAAAALQGe0mpFf42Yll2maqZeq3XThpLr/MMrxEP4wtMxgfA47nZIAB0Jg2xgA4T3wQAAADdBmtRJqEFomUwI3/vgOOCJW+nXGeBH0bjrWPxLGdQkrwjNTNgde2mMiCqnPmK8M/7dGPanLBlAAAAASEGa9knhClJlMFESxv/8A6IrIEOZqgng5VRuFA82zmJM8KL0rKYQRJ7Z23FpDej7QMQVdrodIfe0IE+0l4nF1TNAXWq6QYt1HQAAAC0BnxVqRX+NmJ8hS3aWVfP2c1jQ9YmStz/ocuUcfXzyBUtM57gT1Y3KU8wl/ZkAAABKQZsaSeEOiZTAjf/76z44DvMd9wsd+BXDr5TM0gGx6Y5mxrdaiR+xZhz38vjjy0wqFQLRtVm6z4tKSoPcbC90J+/nHLhkajPTm4AAAAA5QZ84RRU834HLEsIWvPinvfHxN11tVvEp3U0V5MHr+u9+xSzIWHqzyYWqp/rDE7G1eS8KR5LYfoqgAAAAJAGfV3RFf4kZCgkjcPXSmeg5eQB5cdFgpKJnWyRjJ0LC8/6RsQAAABgBn1lqRX+HBFjQlOXXoyp+tUCmnUXoUiEAAABUQZtbSahBaJlMCN/74VHYQpILVM82IUwsR3tP1BZNdPf7nkM/rWUZ7Vjytb4wtxQrFU0OUMs2AySifB/5qpsYFfAbi+7wg9W7xVeK6NYvm492GehPAAAASkGbfknhClJlMCN/++A44DvMfbctm/iNKou1bR3qaQPdbYPAwRr192Z0c7+EQI4zvYW462uiNmrbmQcLA9GkcsKw1co4BTWHqy2AAAAAMEGfnEU0TN+D+24e0TioUzxNily9Go1T0H+2EjDCXB2dGC0c6xM99Lg6Jduz8xAlJwAAACEBn71qRX+HBFltjCbzxryJ9mcJ8tdrhFmiT0jMMVk0HmEAAAD/QZuiSahBaJlMCN/BI78lscgq/Rbp6OxwvfwmtoGweCC3o0q532h4ICIwjwcptqf9EK/+79F9P8RvOlrjacSPiqIg7tu0F2k6OuQHzHZomSpzp/xnEFVc3YpD8D7mVeTAU9Xkf0URAjJ5lfHVioB1x0ui2g+17Y3BGxACMJ3vWNmNMsSB6yWO0rku1Z5c+jFLkY2xMy95ck9Xi4MpgR/rgU5kie+i3djqs+2ji7H6EQ+ogtPCUzbkhOsAieuOLwevqHshJvQds8aT+lB97TlMWAfZprX9Mrv3feABeyzfKAEPhJ3BSA1JjxJgX63hdba2srRrtpcoYk81dOBi6AEYAAAALUGfwEURLN/QU/leQXMu9VvVKf5/iBcuPswhbCCq/njMDoMNrlzDXkngFpZCXAAAAHEBn/90RX/WyxZvpLbxTBsrsj/wSTqMIPKPLDvNeh/1F4bmWrWmiWrQyZcVh27RFCrcMqxQPKxSxC/gK2atA4KwaYxGxim8cYsJdnYB2lEgpZEZ9RrWqfitTU8KyDfSoH8YyYLFXyK9azzCP++iehmJGQAAACUBn+FqRX+Ky9y9AWM9M+//Zr43fS7wb55u5yXHFaE5OBuT0DsMAAAAREGb40moQWyZTAjf87uziMWmfzdbE8J87clOeHoYJgH3Cx3Zbf5fu+tjc/rFA2s6Zalj2+4xEwd5fOTh9i8yAEFuj/ORAAAARUGaBUnhClJlMFFSxv/74VHYOCF6yTT6YYA389+T2IXKVeD0g4f5YiUfdwqkE9LXLSW0Wg9bBFbTjG9iud0FpzmEuq0RQQAAABcBniRqRX+OkjwIpFLTj3OvmKtI9JIoYAAAAJJBmilJ4Q6JlMCN/8epy8Hmt5nkvHwUaGRnirECn1F6kQRmAK+NpNIq1+wKlemlXjXEhlgMZ4GsP1pm9vvJhG7Qc9ADDBVLkvTUXiuLhFe47rYgESCog4vHJ3mr48vRHWscHlSyjaWuoADQyZ2fVR3rFtDDYeIJSS40gQKiB4RIui2G6WIFhug4v4SxqnhvpgyxgQAAADBBnkdFFTzf1+NzzhJSSLZXInxj28fVa9sZ3iOeP2ue1tve6YTVpohEK9SBnnnNjSIAAAAmAZ5mdEV/hn7CXSp3/tuQ/gVx//v4sxBQVVQo1NUSFtg3nNLgRrMAAAAgAZ5oakV/0LYnEaxS2mK1fAK4ZH+xIFCrT1VZklQbIOkAAAA2QZprSahBaJlMFPG//APDMcCEr16uk80rHrFFfOOZRse1TbxZVz1dCEySWEu+EN2piNYDrHViAAAAJwGeimpFf4cEPjJJVPrWo1K1z1x0dfojX58CSEdNp6MuCQeaYl5k4AAAACxBmo1J4QpSZTBSxv/74DjgQlevXVYpZHYgY5ILdPK/4Pa1MNf4OiS8dl1R5QAAAC4BnqxqRX+NmJVnz7XQ5t2DE0nG/9r+8NOf3WQ8rd79EUL5OVWw3o6JYHnJty2AAAAAa0Gar0nhDomUwUTG//vBQYpYQEgvj00Y4sfEOWfwyh/OqKfKANvwZZWb7+Zzvu2GXqL7tu9BEtBF1qEI4yW/avkwQegNVC2ROsxmqz3xMUT7hq6sPN5B7n/7jXSLByhN4LaYlMrIVP2Gme+pAAAAFwGezmpFf46SPAjbqplSkyFjtDCI/KdAAAAAUEGa0UnhDyZTBTxv++A44D22Hq7hOPzCa+Q/0uZ8jWnVjeMaXFL05CiAEmr94qQQKh64dWL5oOBuv5DC0B686Pn9ihyHP8aNBy14cSqPoOJZAAAAKAGe8GpFf4cEPpZCBxpCqtM230lpYDo4d21o37/cFMIRIifuqQPwQmMAAAA4QZrySeEPJlMCN//74DjgQ+vwaqsgahBQipjp6SW8umNWKvqDeH4JFmwRnU6qYAsliRyGwuOB45UAAABOQZsTSeEPJlMCN//7w+5hE7FgdvuOKx17H1EIu8P6dhdyNR3Bp44IixcplMPwMmRmTbkZKRLRSf3U7jRY95x56Yr6mgyfOofav4khhB6BAAAAV0GbNknhDyZTAjf/+8UQAxkZAR4MbGRlJS3hcrKcA7AnvEOlxAMEgQ++P4BRPl+kFk9FtV5R/aKjnv2HOToKh1xJET0cChc94nwhXYYLk8OsF9FeDw4hcgAAADhBn1RFETzfghKjLXhVTVpt3odhU1/9thHbD1IrQDd5C/7EbNj0Cg37PlRRDbIVjd2YJ6YmirkTIAAAAC0Bn3VqRX+KzvglsL6ygpfSqzjv723KBBWB5nY6qSoX/1oYogCyBnZuxHtpF3EAAABDQZt3SahBaJlMCN/KCEDzIZX/gifClN7W6rOMfzBCSLAvaxS3ajNlnvJx6yl9zKnNgrw9/FvayjPOvsK8P/j0jTsPuAAAAClBm5hJ4QpSZTAjf/vgHIEFoP7SF8/0hA40oETmSgI6eYrw0ApjXdIPoAAAAGJBm7lJ4Q6JlMCN//vBGT19kCHNXXJVpZIVW2dOsxg+4nBxXYKH8lgiSB4xOvJRCAiVFRsPHQ0NMqcJypOPU3YoZF/tZgsowrc3AiXB3rAlfqaAL7OjHoF9SqtbiXfhpU4xMQAAAGJBm91J4Q8mUwI3//121hr05LoQDvslraTvUwVhvuZbxxeVmAmt8i59vy9EIHQrpGFHmbLnzjIT69y4Yxu/aorB5LH0qX9SvpgCyK/xkpHH537XWLW78sVsfAh6XcctCpaW8QAAADxBn/tFETzfgfhy3ZPoRbD4wmzsftxxUTshAaOeG2I0yWL50v57GYufV/HImcKo1kE6CdyqZME4g3D6gpEAAAAmAZ4adEV/jZiDOg/2dUiFLHxw+J/iID4obTfYetWmDphQQWMG2L8AAAAgAZ4cakV/isvpPCgm6QyWAb6SK+HCMEJXEyPb987umxgAAABPQZofSahBaJlMFPG/yghDLjqnQ+U5Iog63cQgD0Zbv41jiJ1Bb0LhUXR3RIPLhvolGP5fCuJk+gAJDaNEK0tSdXgn/kDyA9nx+m/zdv65gQAAACgBnj5qRX+NmHWKX61xnqOPLP7m1eVao9/Jf8Igt0BHxeMZyj2uGarsAAAAQ0GaIUnhClJlMFLG//vgOOBCV69dVilkdh/3zGtCpgFkyYUdvr7GH5COYLGX7YS8knmvd86QQESRdt876rG5nJMdX+4AAAAwAZ5AakV/jY0SA6wrwXjljlqnIHhE5NeeEAmpe3+gSytfIFQvKbbmEuo179PCD/OBAAABDUGaQ0nhDomUwUTG/85HJkVD4qSGOwcA3K5VwP1c6xtTR4VzmQ0nBu1gOtLRFBdIfSP3an9P1gmsInCPe79j4zuUsL7DK//x4TUhKV8qs9aaEd6tCVykWUn5Gopj46/K++Ykslr+v5UodumlbuShmmcRvZfdNrYnmbVICpyTGYIOI3oAkq6vllklL4brVgHvyURSStEFd3HHzSm7Ib3PtR4E5I89jUSkUoRqP6CDkJUiXvc83P9RW253rviuhvzqjem5Q+hPFgGTABoRk2ZLNhrce5zzCxaYC6gToCoWrl4lzpLNEOsFnUNIz4XxjMda/67iz7TaE3NHffWa7Y/p9uDhqQF1fsC6zz4vKXmUAAAAWQGeYmpFf9Ei2Luv/RNJcaZzW23Slv1beS/YB+iADhqJnO8XGwim1UkZ/5DCTycTZV+kD7aTYfAV6mID48v5zZgCA1mNEtUUWkqJoLo+Ei7fqexybID5ptuhAAAAZUGaZknhDyZTAjf/++A44Hc4/O2JB6+GsntZ/QVS7FVbo5y6gb6CO1Sl0lysiHpH6tBqj4+RV0j1AkhNZawWUW7mw1Ydkn1vMjA/KVOpoHVQwnkQTNeO2ZHz8y7tU10XuEBdXu1VAAAALUGehEURPN99TaxHbgQG60QG5CYHG0EwwVTLS+WEIzdD/zb2JngNyA/ValMZgAAAACIBnqVqRX+HBFltjCbzxqFwcAaE+pOC4nnIXAWqYK0KnVEZAAAAWkGaqkmoQWiZTAjf+8SBZr6iDxTgWdHsgrtUrFBwt9rOo3XBJP8ewDzX8GtwlWBgrC6irWNVrTLY+WKlSir9pSwV4EguL0xV3Wcj0/zw/APuEjV4PAH1Af4Z4AAAAC5BnshFESzfgh5hlu/y+9zSPV3/9spsZbsjm+DaIPLk3e+H8QEI6F/J3CJy6zchAAAALAGe53RFf42NDc3tk/BfLpHtYrknWRGdfCfWGF2hHwQDNNRoM7TAEbExLqzBAAAAIQGe6WpFf4x17dokp4/JyHzCbIpQYkeLYHCEqpiCFaz+wAAAADJBmutJqEFsmUwI3/vD+CADeEGcfiMMXyiDTP3ydcyGMsECu6OwPlhc6nEUtvJobaLBYAAAAKdBmw1J4QpSZTBRUsb/x6w8INv7ObYOML18RRf9FGE1XcGr6KXtfZJPIpBQ3Zddf/xcPS72YLTyF1aF1OgpE73KGCe4EKRaaP6+28u2fFYqgdEy089G2c5o40yj0KCcaF2ckf3qEd7wZbN7tzArbJ3ihmBK1aiIZDAZIELsAMy+cnsNyfdeXk6vZoOzvGLtwNlOVFAWfFH7l4z8o6rut3fiMJEBlecPgQAAABgBnyxqRX/QxMwwr1BpDDFywfg+dwOr6eIAAABJQZsxSeEOiZTAjf/74DjgIHS2PuRVzH01Z0BzuHiEMxkS+es2RQpgPq4cAkl78+syHU05veIX5apNghHwp2dHWeQBuTkHc9KPKQAAADRBn09FFTzfgfhxltJIut7cibOx//YvRw99xcIjnvocl0u27U4ZQtRxhVb+2ZI+NBhFvmHQAAAAJAGfbnRFf4Z+u/AqLSLb8iSXq//onmci1775jnjU+EnU1sAD/wAAAB0Bn3BqRX+Ky+lRxhN9YRwoFckJ8K7wgeXkfQ8c+QAAAE9Bm3JJqEFomUwI3/wDoisgI4JCqEuNC9+tPU5T6akQg8IizB7OrAsSlrnkH+zL9hIrZLMG+0x3pM9DfD7xIYuRzu7XLhI5qDban3u9P/d5AAAAVUGblEnhClJlMFESxv/7w+5IrxwG7Vm0o9WyHVPRTEK7sPH3Tlc6BafENz9DO3jdvcZDYYjxJ/xezZD/wVOgNlWr1ovR60fW2n2QPDIrXRl7wbJGuBkAAAAnAZ+zakV/jZWjbgdMtmBO6kL9v/+iv11IkCbeIE0lsf2Q4YI/jwZ1AAAAREGbtUnhDomUwI3/+8OipWXVEHfdmUcs3voSdMpPrWh0B7x/buw/7bfWTU2SnRyzDOlwhY0AHOKdTjMVTnbawMzXUYswAAAALkGb1knhDyZTAjf/++AcgTTZH1Q4wq20ZHYluqF2d5yizb76mgyembcA7BYepeMAAABaQZv3SeEPJlMCN//74VHYNYZ3xEb5VWTtwzMtJdYmcfzq/85J9tfpgIKEqieYPeY/MEIfk5XhJ1koykG1DVrGfyuRE/hx6dLKPysfIvtARn9KJW3Q6CD7moIQAAAAY0GaGUnhDyZTBRE8b/vgOOB3OPztiPZIA+2600lWyYtu5etLlp5Z1h7gvaM9dZAwSh88F2MaIP+Gb/x+JzyfKCfhrWJAkMuh2vjzRb7d/N7ALJMAooNT/8nzsWXFsQPjIw4BwAAAACsBnjhqRX+OeGijP1BCTddBMbj/ecV0Or4a5xYPj/vYTAUTsVBHuE5kBqvhAAAAOkGaOknhDyZTAjf/++Acgl11jEgudcSnSrhTGl+N28Vhdn1HYhWUAuZUXRLf5VeB/cE3pIjWL90Bv0kAAABeQZpcSeEPJlMFETxv+8PaaJ1sN1uxRNIGqNtU2k3406usYCk1bEtXHnFK6NYx2++uiXE5dDpIBIRXDOXY8QkTWaxAEokog4eXK891HbBLLwiFuVsjJhcQDmLu/1BPQQAAAC8BnntqRX+NjQ3Qw5SeeQtnahQCGRD+tBtlepn/thplttEshgUAcjhCMocOSDNHoAAAAE1Bmn9J4Q8mUwI3//vDpMA3iJWBxGCoVqH005LOHvO3JnF0uPAF/1BdrAmcA6ADeTiMuCgJLUTG5+vRnVljzqS04Al1sJB4spczTDLw8AAAAB9Bnp1FETzffU2jm3iDcWTT8O+Q1NZfwE4BDhgPUoUdAAAAMAGevmpFf42YlWfPtdDm3YMTsd9f+YkkWkcv8aPIfR7QyuqOLTk5HzNqkb26HGVg0AAAAGdBmqFJqEFomUwU8b/74VHYnLGfY5FAo39Th64FzFtgZ97sVu925ubsyqHBRHBm16zii8aJWj58APV75EudDjr7xLleCLLq7/RtBuOXcv77e28R66Sg0C5+ohD0irjcvKrApu8KWKOAAAAAGAGewGpFf46SPAinGc0bJk9Hz2z1Me3AgQAAAQNBmsVJ4QpSZTAjf8fMy8FXp/MBuvzFgdo0Sga+yREhRElJ1uA9sIsY4ra6c6fc00xs5DNAd8Y8yfothafu7uajsv9X2+oe8VfvZheOdZXtqkEIPdjDT49hK59jghSaTBGRE3a2q4B8vzpyVLIVGWiucmJcTr7iO+Rs1GKykq2Hqpg/7wGE9uYxrJNr28ZDTVBHR/4CyCAPQk6jWNczPh78rCn0z6G9a1RjzePKZYANdrkvuc+8zCzahQHbPyJ3wSnp4GBi0FZoLY5fHqWlr8b307+UI2CHj0yzFgHzBe7+BMVC9y83a86LgcLsFvxDSFgPEYch+RqKUMFJwfV/KNr0okJZAAAAQUGe40U0TN/X6axmThJSSLrDjCbOx+2e33lIIgsTG2JtszOBl/PYEZ41qb1VwqSpMkCnW2Br5kIl2A3Uu0IPihiHAAAAJQGfAnRFf42Ygzli/toiEQB9cbIH+nvDxV+o1J32C6oS5Fi80REAAAAoAZ8EakV/3mnpOLtzNha7JYFMdaGbfLAP/sxzAhGKBTfBtRpxTI6dgAAAAIFBmwdJqEFomUwU8b/0EQS/wgjd+gtxcvUsjBZ/KvLmmxO81sipaPVWNzgoVqy31nJce7cJ78sV7DeSUW+b67IBOE8jk8HFII54ZV3rWpp/dqQpuWbfoc2fWtU2v+/yK+ia43/m1sEOBWkif02dhVPEQ9NMhk2dwH5CcI/zvbk/eCkAAAAoAZ8makV/rN+DaXdiJ3Sfbla5z0prYIJJdn9L5fmjpAFnOrwj7T71gAAAAEdBmylJ4QpSZTBSxv/74DjhLmEQjvo7MOz/6KwI8CxsOXVk34E9cwzHoe+ypu4pyoHVVLROmPuXefirSAYX2k1/KeHKYNsg4QAAAC0Bn0hqRX+NmJVnz7XQ5t2DE0nG/A/2tAvwsk191IaIPxh05qExET5GLoDSl9kAAACTQZtLSeEOiZTBRMb//WN7GRlg7p5tdxIOcz4ktWOw25b+IeCEiFor+loK1sW6QJ67CqjJiEE34iTVIck02sn81tu+PGvjsyHqi2o+NT4inWWVtqWkpr7f/3c+vKVzlj1mjdhK8OP4T7BUoxX0iIL7Ms5LBKBxFY04P++ZvxW8cml6mweqdCiwy+yuvk365ebqBsgsAAAAGAGfampFf5XYaTriDUHbSM4/2rI012dTgAAAADxBm25J4Q8mUwI3//vgPUQ9+HnYzR0ReHsw8528FstVpSX9I30YzTO1HQX7vIeLO4umR1DvPGSkhtbazhcAAAAqQZ+MRRE8331S+se7kbW6AxPsrNhK/dtL9AZtyTBgH5jIFFvIWlbz+Y1vAAAAHQGfrWpFf4cEWNYpbR3VxAdD0L/IP/FPtuyuyqOpAAAAgUGbskmoQWiZTAjfx9CcA6khDjKpbsCmICblPVjONIzX1mXC3smr+Dq5fXiG9Xv8kFQyUglCZwie8jt9HTy9u/pPWV6Eg7yLHBH4vZV7UNGZW/41bil6995IKcn441hEAdqxiSHI4vuUl9rENOzJKWLV2U30qhtdcyYWhWRdSq79ngAAADFBn9BFESzf0IJ5kmO08++CFl4DaPf738KLIygT84teS4UEhuc0uV5bKc3qgHgZU8MlAAAALwGf73RFf7OB2xJd3tC/Pqn8/HmIdT3BpECF6yr/lbOPB+Zy1K11SUfd4VAGF2JZAAAAIQGf8WpFf4x17dokofjXiSEcvI/1LMfeM2Env+wlhWYh/wAAAEBBm/RJqEFsmUwUTG/9YpGqI6+gv8B9xK7sh6JffgqkidB+uyLMXamYT/v0gM2Bl0cBW+zOs7+meEYWSU7mTdahAAAAFwGeE2pFf4cEWNCU6Hlhkpwoe3NC1OR+AAAAfEGaFUnhClJlMCN/x85ZyjCkCh8uzK70/iv22/5Y3C+WTfpvUCrA+KzvU4pHDm4sP40BAbPg567zdr1Wu7EEBUuR9mVRdZyCPSUp6JfGVxbeoCt4oVKZm8KmBNWkYt6ZDzaPhh8MV+sn/5rn/+5/sTVujqfTRrIRFR+q1bAAAABYQZo3SeEOiZTBTRMb//vgOOB3OPztiPYLGqbk5d2BqIQU3KblqkAke1WFGNwX7uSXHgcpobdidZRPvUHXnKBlNvG1whZJdIUyHR2g5/kj6iGYMgjOKwfAwQAAACgBnlZqRX+NmJZdpmrtcZ0NTAN84Rf1N4fqx73Qed7w+gCCXwdxAjG4AAAANEGaWEnhDyZTAjf/++A44Ilb6dVP1x6RsYOzc76QVwcX2znv5faz7QEaAFkrcI7CaujLN3AAAABVQZp5SeEPJlMCN//7w+5d2nwAIi2iZMsWHBHowYX4aPQ815YaFI7qF/qVwMK0/Txt1GjsuoBrg0JYfv/uwERjL/ROpxdPOR5x3xoTlVPE5qTT470ToQAAAFlBmpxJ4Q8mUwI3//vG3LzDXKIqZ9TYmttr5CBr315a/0BUuW50ROufAOdVNE+12YSjQ+/9Zk6dVP/79hzk6CodcSRE5KWMo3xZZ6kBcB3u5L6hIx3KSQR9gQAAADlBnrpFETzfgcr6UK161qa1NGdzHnS/SFXblAfYkkTUgGALg0oSw2I9ZdDSKhlEZD82KSrfFpA+v6EAAAAnAZ7bakV/jZWjblZD8+kx6s9vusQNFqTkiIzV2G+Y2RxdmUUlRu/AAAAAVUGa3UmoQWiZTAjf+8OnUzwggj/qGfOg+4iB+MF4Y5ta94QwdwgcgpS5V9IyYNFhVGlpN2bjMQ4DNhuvd/qEmbi/FxB9D9m8UL/JpA3iUwbQMpYlEjoAAABmQZr/SeEKUmUwURLG//vB2iYdh5ImpKHSwtHernzZCXrV5CEdMKE6oB+B0fwylF2lU+/yPss1M1i+5F7zxa3LLlnpb90pQhlGEO/tirOcGbuioWk+zu7MXqDCxg3B7BB3iJfYGJshAAAAFwGfHmpFf46SPAi9dvB8N1WaHbwo5+8wAAAAaUGbA0nhDomUwI3/+8Gbu08WIicdkBLuDqjcxTRGWFEHK1YMItB/Yxti7mXMwKkgf/cswQBjw/L2b1agYYBo9t6xn6Eetw4bxPxz38vjh5Ug36pxN1P+Jdu6y5/v+SSnY/wBOAf2anagvwAAAD1BnyFFFTzfgcrcmZwgJIut+RPjHv9yMW/NWZe0k8Bi3Zgre6X89gTtwG/OmXBXHUG6b/SiwXTP+8AECry3AAAAKAGfQHRFf42Ygzli/0aMsY2QSz83/bIH9oh1NSd2YfdN+fUQOgXqQV4AAAAgAZ9CakV/isvpBhS2mIvgYfwGKq3zW41GlvjW/evQpscAAABJQZtFSahBaJlMFPG//APDMcJc0yddeUYH6nOgY2NbwtKavnsY+nBe+fxVlILL0NiRXrOhPUiA969gXTMpBWKdBefNzN3yeWc1gQAAADABn2RqRX+NmH8N/4LEDrqsxxCUlA+/C6d/hrlaE3Eni4AjiTNfEyqi1KVhrZEtHyoAAAA8QZtnSeEKUmUwUsb/++A44S5pk6Zc8mpbajeY3FQlOl5eERU4NlZQ7KCHXnQqL1aUm+ZXyoGSi1CXUAu7AAAALQGfhmpFf48rV2va6HNuwYm2A5WJIWicynQ6BVYFp9FWnKrYE2sDGYjSbKzVRAAAASVBm4lJ4Q6JlMFExv/HSJgSC1nZxJ/p92h7iaNLjO1oWGnVb7mdLzqjQegiZZ8BFVoHleRS9opigDSoevdqf/HRbxPUT937O3gQ14DzaDIrquleqdGqhoMap2XYPTsxa/wSkDdVT0ssgNi0XqyHVXOCMqt31kydwNB3y5PU4UE0NqfePmpADjVY3oiJGyZaYPsi/X4ym21aC/XEA8g6cHPxA5l4ZzRa4mnIZk+87eHMPyVrzjVuzsYf3L5Mdth1vbIMZuHrbzmxDF4Id95FDUpJLeBKHbiVKXlKbVlORXdE3XDZHj4OQ44bzKNetjD2jssvKy2arGt8LqXFFRWGchvabmDpRlO+TJnoh6Qi34c/JAelcxWZTHtnFBN5/wh1aZMwTQYfgQAAAF4Bn6hqRX/RIti7sAopo312LMs22rFhBmyBC8ZUUUusB76jdtD7pxWchSlEaZcpo+z3LSOsqF+TjWIbKyoOjcJ0BvP5Un6zNu7CYyOjMKakeOmtBeHghnZ0Lt7ZKgmBAAAAZUGbrUnhDyZTAi//87FwbhKNgmMp3eNkrIgo+BdeIZ2RH7aVgRWSRB8VGoJF1H932TYjOcWyo9X8eYosRzKDsLYZ8OElkDGMElhddcoxE4Cuo16C400E76DzMxmRGQrfcMoXwPeAAAAANUGfy0URPN+DBf3EVFfJcIvLYYm40FJuHP41gUVpmbF9867RDcpbG+hUwmLp76j9ncDp8t5SAAAAIwGf6nRFf4Z+tIoBTm+ZzUi4P/L4NbInHV4wNqEejzvRIyMbAAAAHAGf7GpFf4rL6VHGE3tN/0VO4T4Y92cJmTGqjpgAAABFQZvvSahBaJlMFPF/+G0p3IyKrpYoACoREvabxZ2Sa1j/PX4Ig63AFh5hVcDy6lyqDJVazyv9k0A3UEbpQ4kElhDfQ/kDAAAAJQGeDmpFf42Yfwkkpl9ce7oX5djjxlYebpvmnM3if/h60uGBIJAAAAA0QZoRSeEKUmUwUsX/+99LcS41uOHpiqjPsmwGusCKQ6BF1AYtz/WtQ7b6AecL2sZs2DeFgQAAACoBnjBqRX+PK1dr2uhzbsGJtgOViSFonOVO2XvUQkyMOnNQmIoMiU2QLcEAAABnQZozSeEOiZTBRMX/++B7/JR3TzHmkys6OPVr3j60/Ducpwbn9LZ8EMRIp1AAs7+bDOf0XhisAXR+MHkgbVlZgREbbfkD3LjS6EKrXc0GPULfC2RHIyE230IZnBmWLkBNNdGdzIhkLQAAABcBnlJqRX+OkjwI26tNPqSi/HkmF/g/gQAAAHJBmlZJ4Q8mUwIn//Dz42spRFihMPfTkj2TNj0lStYXk3a+C5Wtkfypzo0uM/feYzEOqkQzy1NzoYXThBgSRvjrEcejioh4gzCm8rCbzMI2wWFHN/yLtF3p6q9/Ec94jMIwPBQYnzCY283aqNgEwQc+asIAAAAwQZ50RRE836foBRLD06goiSPhmYv9qz7O6uD4czg9Q3XDksNdfK5w0IuoG1k3GS0sAAAAHAGelWpFf4cEWNYpbQS22jK73/5IyrwB5npkP/EAAABQQZqYSahBaJlMFPE/+81+QTCAevbatPJvZfPu+5FInLiQmsx/vsDDnlk+n51NMYohIZisWJ/S3zrGQtmid1/VGv24Rk3j0Dp9U4U6eejQnKAAAAAuAZ63akV/jZifIUt2huke1h+nmbVlvSl8BsecDgLV4KKiXsuB44t0CtoWXXNbfgAAACtBmrlJ4QpSZTAivyCC9u0SU8fn1JmlqS4H0O/V2QQ8C23N8Hq2eJg+n3lBAAAEpWWIggAb/6zQYZLfQPgNc0iYMN3sPl0EDBzzvrz1O9Sgfa49FGnVhJ8qLkhVezFoQxJKbc0nUGNZR098qd7OS+y89EcOWVfDbC64tS8Y9KVBO5bX3i8BNZhBqqdfBpIeL6SUDjC2QpVPdT+4GdjwHm+HVEddJzn4mxN+ZHi1zAImnjZYMgtVcr/8KsM8y6i5AmAGJNcUB7ayE0bca3a3R6hdrr0XJr+AAAXHQ/kXcVlmP6pUrkjcKn/5MsrYsfvEq74eZ4HZz2CqbE+TyZ4bBsWkF/G8EsnWTdMeCaT0oSOggGj6F6GH2koqwzCw8EE/jxjH/fgpvYdskan80JjRYRAJC3SwAtTkEnddIDJhOwFLkmhXRf1WZtt88VL7WgZg+Do7We/RCfv4cfPN57fJ/eEcGoPIdjySktCXmVnqHKZZhx61mqXrYt2ccmalI5pFwxn7yIk4NanyDkRuoRQdzVy7V/FgNDX+/OG313omVQVM1D6TfnNy1isPsCV96//sVHChbSvRvo9Op4uXUT2Vnp2NbBUnPOlIQwYUmnls4/qAf1DSyr347ve+HuWjGSfL2uAJLxNkKGL7jaqU87H350pTIU4+n/FRrP8UOsWWuZ9oqEO+lS3qvWMNvTA00dfS9Og3CEQ1lhvXgqkfbgmWS7Q7BrqjPYxv+usZZhR4mTsySOLjIGUswJByOkD9sXOn034oNMHPSMTs90mvpwuhhZCES9QIA0MY4mVJAH4U1175JsR0JGrNT2cG+kevDwhB7L2Ps9sq35Ja1G1YIr1fCqBH1NhFZbh3r/JJ6IfmaKkwRe/x35k1btKb+ZRfDYCZwD/U+DVYV4ve4/ux5EMulf81TrZ2iS/AeNl9oIVL7FLsYt0SbM4Gdwujxy/cVa/9DfeeCmTLktfb85Jq/PPzoWfkT+3MSEL52ahvl5dyZ3P9Gjqc/3V3E3UcMDLHXaFjLji1XATHryNndFbY++diGv2rJsu8jbVZYN/lqAhbLHtMl/sxkHV0L1mnPx14QjgEVCWYb8PE3T+6bY0rrXT6WTEzJq/uaxvgC3UliUMREWCB/LHkW4/FwMNKpTt+KgL7DKMzlTFtOl01XRpLhfNYIv4yqp/8/0Ljp/cAHvBSTd0o2zGUJLgBlXpkC7hkWNvVKuNH3gZDe7aBWygB4hHI5MtXigV2BDYnOAj1gC8+xGfGlAw3JTy1PkTj3bbr+p6m1jzipdInXtuw8i53LkL/+t/cDhYKza5lgkrWZvnpedOPwE+SDsrxeadOx3w6iN4kI3PIREK2JArXFqTIikdsxGeR30xC/R4MnKZUvO9n1HHNmqCdKQY/t8N97D5MniGESaR7uA4AEebTLOvxLuJghYdjoSxajwIKjHQQZ/HnrodJI5KNoj6gM3UJuupZ/7uHrjf0Dx0Tsc7GUns7O78xmwOt1zDx4hBNRbDjREElJr9eJlRZ3FMTZgL28DJXX2Ukq+M+ZYpT8CiXX8w7XSQxcMg4cJ6c0AGluEsCz3tAkAQZLU/VxMQpKTJR8+ULaCRGp3sYhJhT5mx6NspTrLO4jl7T5GLVFULkYWq1jSnmblAgEE58LcEAAAA9QZohbE//55dAJChmGePT0n0Li66IQjlgCcvSeXFOTWduetNIn5ncKj2uvzaGAWdM4YtqdWioqkA0TyHz0QAAADBBmkI8IZMphP/nR8JCz0aek9id9vrF4lLN2RgQjyZr432P7Kp14LSDqmMhMkWW0IAAAABgQZpjSeEPJlMCf+dJU6LZgkV02QLvEfeXB5dv4nCjjykgCXjsebUglI9rrNk3mZJmTaQFGdOlZR2+jF2hfcbd4HXQBsGGhxrxbbB//345uw0C7pMaLkJqkj5rLI3WN+aIAAAAUEGahUnhDyZTBRE8/+ZYKy4kVT7wb4JNhryfYT+xBrDfh8zeOyqxLHImSnk1hEFx+vGFOK6nSbea/LiCsJT6jX4eofbWT/4UXXlaQ7ns7HmfAAAAKgGepGpFf42Yll7P81LVQddgA183/eHh4rICR6Rm/TAP+50pi9D2H4nE8AAAADBBmqZJ4Q8mUwJfi8K/Q8CzpoXb9z2cRUhHfcB/4ijSNB3DB8kDQxNhedtHQQegeFAAAABHQZrISeEPJlMFETy/jBNpyKcpBz2UWq9NhKZ8OdqT4vlDJSja2L4KsxrGpSrA2yYsC6BONsZSjEWJ7nESRHswZooLCvZWL7cAAAAtAZ7nakV/jY0N0M9lu0TcpcLEt4G2uP1I1ad/p8M0BAfmcyRE+XWn8kbnbCP/AAAAQ0Ga7EnhDyZTA/8JnI6fa/XOcPoz6Eunvm3DSqWZut6xANl//x1bLCAKYLkthnFxJePmxk5tfDOwQ/OimiJzESuj2PEAAAAyQZ8KRRE834HLEsIWvPiTtw1hFA7vKrrccI3560FD1CXoViI/ZgtgMQkFL76i8/+MS6kAAAAdAZ8pdEV/iRfHjugOkdFKyjh/gRLx8g9LelE/H6AAAAAYAZ8rakV/hwRYbMohViYe1/b50qjneEmNAAAAWEGbLUmoQWiZTAr/DOcGwr30MLcyf4RisZLdUWnlUYvXJVKV7AncbyAoLODntBs0+K4WaaoVpZfqfI2XI0wXb+fJSpC5XehdM+mTV3yL+1xrM/enVV5ivWAAAAAzQZtQSeEKUmUwIr8dln8w8D8m9uEVJ8VaXrlpwtQMUcjOxV139A/A0FOJb3veKqfxYs2BAAAALkGfbkU0TN+BznEOeLmqf8Sjo1rSbw27tYVeLjYdb2CxT8VnvpcHn1axdLPZBJ0AAAAtAZ+PakV/hwRZbYwm6h1dIRfbz8pmPJ5q1aY4QcQjnAtS8tukkKVgqkBfhK/oAAAkC21vb3YAAABsbXZoZAAAAAAAAAAAAAAAAAAAA+gAASucAAEAAAEAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAIAACM1dHJhawAAAFx0a2hkAAAAAwAAAAAAAAAAAAAAAQAAAAAAASucAAAAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAQAAAAABQAAAAhAAAAAAAJGVkdHMAAAAcZWxzdAAAAAAAAAABAAErnAAACAAAAQAAAAAirW1kaWEAAAAgbWRoZAAAAAAAAAAAAAAAAAAAKAAAC/wAVcQAAAAAAC1oZGxyAAAAAAAAAAB2aWRlAAAAAAAAAAAAAAAAVmlkZW9IYW5kbGVyAAAAIlhtaW5mAAAAFHZtaGQAAAABAAAAAAAAAAAAAAAkZGluZgAAABxkcmVmAAAAAAAAAAEAAAAMdXJsIAAAAAEAACIYc3RibAAAAKhzdHNkAAAAAAAAAAEAAACYYXZjMQAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAABQAIQASAAAAEgAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABj//wAAADJhdmNDAWQACv/hABlnZAAKrNlFE/nwEQAAAwABAAADABQPEiWWAQAGaOvjyyLAAAAAEHBhc3AAAAABAAAAAQAAABhzdHRzAAAAAAAAAAEAAAL/AAAEAAAAACBzdHNzAAAAAAAAAAQAAAABAAAA+wAAAfUAAALvAAAU8GN0dHMAAAAAAAACnAAAAAEAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAwAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAABAAACAAAAAABAAAMAAAAAAEAAAQAAAAABwAACAAAAAABAAAQAAAAAAIAAAQAAAAABwAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAADAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAwAACAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAACAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAwAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAACAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAMAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAADAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAACAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAMAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAACAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAMAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAMAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAABQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAABAAAAAAAgAABAAAAAAcc3RzYwAAAAAAAAABAAAAAQAAAv8AAAABAAAMEHN0c3oAAAAAAAAAAAAAAv8AAAZHAAAANQAAACEAAAAkAAABswAAADgAAAAmAAAAbwAAADoAAAAzAAAAKAAAADcAAAC4AAAAIwAAAFoAAAAxAAAAJQAAABwAAABfAAAAWAAAACoAAABXAAAANAAAAF0AAABhAAAALAAAAD4AAABHAAAALwAAAEgAAAAeAAAANgAAAFgAAAAiAAAATwAAADsAAABXAAAARgAAAFUAAABFAAAAQQAAADUAAAA4AAAAXwAAAFQAAABEAAAARwAAAGEAAABVAAAATgAAACIAAAA8AAAAQwAAAGMAAABRAAAASwAAAEcAAABuAAAAUAAAAFsAAAAmAAAASAAAAHAAAAAbAAAAVwAAADQAAAAkAAABHwAAADoAAAB7AAAAJAAAAEAAAACzAAAAJAAAAHYAAAA7AAAAKwAAACEAAABXAAAALQAAAEQAAAA0AAAAcgAAABwAAABnAAAAMQAAADkAAABcAAAAbgAAAEAAAAAoAAAATAAAAEgAAAAdAAAAZgAAAEQAAAAvAAAAJgAAAE0AAAA4AAAANQAAADYAAAFBAAAAZAAAAB0AAACUAAAAQQAAACEAAAAvAAAAQQAAADIAAAAfAAAAGgAAAFgAAABLAAAANwAAACkAAAAfAAAAVAAAAFUAAAAnAAAATAAAADcAAABdAAAATgAAAC4AAAA1AAAATQAAADEAAABDAAAAIgAAADAAAABhAAAAIQAAAQkAAABPAAAALQAAAC0AAACOAAAALAAAAEUAAAA2AAAAaAAAABoAAACaAAAAMgAAACEAAABGAAAANwAAAC8AAAAmAAAALgAAADIAAACCAAAAcwAAAD4AAABaAAAAXwAAAGYAAAA9AAAAKQAAAEUAAABYAAAAHAAAAFMAAAAxAAAAJAAAAG8AAABPAAAANwAAAC4AAAEZAAAAIQAAALsAAABvAAAANAAAADAAAACtAAAAQQAAADMAAAAuAAAAXAAAADEAAABfAAAATgAAADMAAAAiAAAAUAAAADMAAABXAAAALAAAAD0AAABTAAAAawAAACEAAAAvAAAANQAAAEwAAAAzAAAATAAAACIAAAAyAAAAkwAAAB8AAABpAAAAOAAAACgAAAD/AAAAOAAAAHcAAAAqAAAASQAAAHQAAAAbAAAAnAAAADoAAAArAAAAIgAAAGUAAAAxAAAAQwAAAC4AAABmAAAAHQAAAGgAAAAuAAAANAAAAFYAAABwAAAAPQAAADAAAABHAAAAWQAAABwAAABkAAAARgAAAC8AAAAkAAAATQAAADUAAAA9AAAAMwAAAOkAAABcAAAAXQAAADwAAAApAAAAIQAAAEAAAAApAAAFBgAAAEAAAAC2AAAAHAAAAEwAAAA3AAAAKwAAACEAAABTAAAAWgAAACgAAABNAAAANwAAAFkAAABTAAAALgAAADcAAABKAAAAMwAAAEIAAAAhAAAANAAAAIYAAAAdAAABMQAAAEsAAAAuAAAAMQAAAIMAAAAxAAAARwAAADIAAABkAAAAGgAAADgAAAAvAAAAIQAAAIgAAAA3AAAALwAAACYAAAA7AAAAHgAAAE4AAABYAAAAKgAAADYAAABSAAAAXwAAAD0AAAArAAAATQAAAGMAAAAfAAAAagAAAEAAAAAlAAAAcAAAAFAAAAA1AAAAKwAAAD4AAAFQAAAAXAAAAGMAAAA2AAAALAAAACEAAABCAAAAKwAAADcAAAAwAAAAVQAAABkAAABsAAAAMgAAACEAAABbAAAANQAAAE8AAAAtAAAAPAAAADMAAABIAAAASwAAAC8AAAA/AAAATQAAAC8AAABQAAAANAAAACcAAAAbAAAAdQAAAEEAAAAoAAAANAAAAIkAAAA6AAAAeQAAACoAAAA9AAAAGgAAAF8AAABXAAAANQAAACYAAAAgAAAASwAAACoAAAA2AAAAMQAAAF4AAAAcAAAAWgAAACoAAAA2AAAAUAAAAF0AAAA/AAAAMAAAAEkAAABpAAAAHgAAAHUAAABDAAAALgAAACMAAABKAAAAMQAAADwAAAAyAAABcwAAAFkAAACnAAAANQAAACkAAABuAAAAMgAAACoAAAAoAAAAVgAAAG4AAAAcAAAAWAAAADwAAAA7AAAAKAAAAFgAAABWAAAAKAAAAE8AAAA2AAAAiQAAAEwAAAAtAAAANQAAAGcAAAAzAAAAUAAAADkAAAAgAAAAHQAAAGUAAAEeAAAAUAAAAC4AAAAxAAAAogAAACkAAACOAAAAOwAAAFoAAAAbAAAAVQAAADkAAAAsAAAAIwAAAEsAAAAqAAAANAAAADAAAABTAAAAHwAAAFcAAAArAAAAMwAAAFMAAABdAAAAPQAAADAAAABcAAAAgQAAACIAAACCAAAAUQAAAC8AAAAjAAAATQAAAC4AAAA5AAAANAAAAQUAAAC3AAAAogAAADUAAAArAAAAIQAAAEAAAAAnAAAANgAAAC8AAABTAAAAGQAAADwAAAAxAAAAIQAAAE8AAAAzAAAAWwAAACsAAAA5AAAAPQAAAF4AAABRAAAAMAAAADcAAABQAAAALgAAAEoAAAAiAAAAMwAAAI4AAAAdAAAATwAAADgAAAA3AAAA7QAAADQAAABwAAAAKgAAADgAAACTAAAAHQAAAFQAAAA1AAAAKQAAACIAAABSAAAAKAAABUYAAAA+AAAAlAAAAB0AAAB3AAAALgAAADMAAABcAAAAZQAAAD8AAAAsAAAATAAAAG4AAAAbAAAAawAAAEAAAAAtAAAAJwAAAE0AAAA7AAAAOQAAADEAAAEWAAAAVwAAAFQAAAA3AAAAKgAAACAAAACPAAAAKwAAAE0AAAAmAAAALwAAABgAAAA+AAAAMQAAACAAAABSAAAAMQAAAFwAAAArAAAAOQAAADgAAACJAAAAJQAAADMAAAA2AAAATgAAADMAAABCAAAAIwAAADEAAAB/AAAAHgAAATUAAABJAAAAKwAAAC0AAACEAAAAKwAAAEYAAAA0AAAAVwAAABoAAACaAAAAMQAAACIAAABJAAAAMgAAADAAAAAmAAAAMQAAAHEAAAAcAAAAVAAAACwAAABEAAAAXgAAAGMAAAA5AAAALAAAAEcAAABjAAAAHgAAAHAAAABEAAAALAAAACMAAABWAAAALgAAAMQAAAAwAAAANQAAAGQAAABRAAAANAAAACoAAACZAAAANQAAADYAAAApAAAAOwAAABsAAABuAAAAQQAAADQAAAAhAAAATwAAADEAAABWAAAALAAAADgAAABMAAAAgAAAACEAAAAxAAAAOwAAAEwAAAAxAAAATgAAAD0AAAAoAAAAHAAAAFgAAABOAAAANAAAACUAAAEDAAAAMQAAAHUAAAApAAAASAAAAEkAAAAbAAAAlgAAADQAAAAqAAAAJAAAADoAAAArAAAAMAAAADIAAABvAAAAGwAAAFQAAAAsAAAAPAAAAFIAAABbAAAAPAAAADEAAABHAAAALQAAAGYAAABmAAAAQAAAACoAAAAkAAAAUwAAACwAAABHAAAANAAAAREAAABdAAAAaQAAADEAAAAmAAAAXgAAADIAAAAwAAAAJQAAADYAAACrAAAAHAAAAE0AAAA4AAAAKAAAACEAAABTAAAAWQAAACsAAABIAAAAMgAAAF4AAABnAAAALwAAAD4AAABiAAAAMwAAAFEAAAAjAAAANAAAAGsAAAAcAAABBwAAAEUAAAApAAAALAAAAIUAAAAsAAAASwAAADEAAACXAAAAHAAAAEAAAAAuAAAAIQAAAIUAAAA1AAAAMwAAACUAAABEAAAAGwAAAIAAAABcAAAALAAAADgAAABZAAAAXQAAAD0AAAArAAAAWQAAAGoAAAAbAAAAbQAAAEEAAAAsAAAAJAAAAE0AAAA0AAAAQAAAADEAAAEpAAAAYgAAAGkAAAA5AAAAJwAAACAAAABJAAAAKQAAADgAAAAuAAAAawAAABsAAAB2AAAANAAAACAAAABUAAAAMgAAAC8AAASpAAAAQQAAADQAAABkAAAAVAAAAC4AAAA0AAAASwAAADEAAABHAAAANgAAACEAAAAcAAAAXAAAADcAAAAyAAAAMQAAABRzdGNvAAAAAAAAAAEAAAAwAAAAYnVkdGEAAABabWV0YQAAAAAAAAAhaGRscgAAAAAAAAAAbWRpcmFwcGwAAAAAAAAAAAAAAAAtaWxzdAAAACWpdG9vAAAAHWRhdGEAAAABAAAAAExhdmY1Ny44My4xMDA=","ok":true,"headers":[["content-type","video/mp4"]],"status":200,"status_text":""}},"base_uri":"https://localhost:8080/","height":501}},"cell_type":"code","source":["play_video('mf_pong/0.avi')"],"execution_count":0,"outputs":[{"output_type":"execute_result","data":{"text/html":["\n","    <video width=\"640\" height=\"480\" controls>\n","      <source src=\"/nbextensions/vid.mp4\" type=\"video/mp4\">\n","    </video>\n","  "],"text/plain":["<IPython.core.display.HTML object>"]},"metadata":{"tags":[]},"execution_count":31}]},{"metadata":{"id":"NQmZEVKGF4Hh","colab_type":"text"},"cell_type":"markdown","source":["# Model-based training\n","\n","The `rl` package offers many more features, including model-based training. For instructions on how to use them, go to our [README](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor)."]}]}
\ No newline at end of file
diff --git a/tensor2tensor/rl/README.md b/tensor2tensor/rl/README.md
index 99c4a866f..19b318859 100644
--- a/tensor2tensor/rl/README.md
+++ b/tensor2tensor/rl/README.md
@@ -9,12 +9,16 @@ Proximal Policy Optimization (PPO). See `trainer_model_based.py`.
 As a baseline, you can also run PPO without the model using
 `trainer_model_free.py`.
 
-To use this package, you need to install the Atari dependencies for OpenAI Gym:
+To use this package, we recommend Tensorflow 1.13.1 and T2T version 1.13.1.
+You also need to install the Atari dependencies for OpenAI Gym:
 
 ```
 pip install gym[atari]
 ```
 
+[This iPython notebook](https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/hello_t2t-rl.ipynb) provides a quick start if you want to check out the videos.
+
+
 ## Play using a pre-trained policy
 
 We provide a set of pretrained policies and models you can use. To evaluate and
@@ -96,7 +100,7 @@ To train a stochastic discrete model (it will require more time and memory):
 
 ```
 python -m tensor2tensor.rl.trainer_model_based \
-  --loop_hparams_set=rlmb_long_stochastic_discrete \
+  --loop_hparams_set=rlmb_base_stochastic_discrete \
   --loop_hparams=game=pong,epochs=1,ppo_epochs_num=0 \
   --output_dir=~/t2t_train/mb_sd_pong_random
 ```
@@ -120,7 +124,7 @@ gsutil -m cp -r \
   $OUTPUT_DIR/
 python -m tensor2tensor.rl.player \
   --wm_dir=$OUTPUT_DIR/world_model \
-  --loop_hparams_set=rlmb_long_stochastic_discrete \
+  --loop_hparams_set=rlmb_base_stochastic_discrete \
   --loop_hparams=game=pong \
   --game_from_filenames=False \
   --zoom=3 \
@@ -148,7 +152,7 @@ gsutil -m cp -r \
   gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/world_model \
   $OUTPUT_DIR/
 python -m tensor2tensor.rl.trainer_model_based \
-  --loop_hparams_set=rlmb_long_stochastic_discrete \
+  --loop_hparams_set=rlmb_base_stochastic_discrete \
   --loop_hparams=game=pong,epochs=1,model_train_steps=0 \
   --eval_world_model=False \
   --output_dir=$OUTPUT_DIR
@@ -167,7 +171,7 @@ make it faster. To do full evaluation after training, run:
 
 ```
 python -m tensor2tensor.rl.evaluator \
-  --loop_hparams_set=rlmb_long_stochastic_discrete \
+  --loop_hparams_set=rlmb_base_stochastic_discrete \
   --hparams=game=pong \
   --policy_dir=$OUTPUT_DIR \
   --eval_metrics_dir=$OUTPUT_DIR/full_eval_metrics
@@ -194,7 +198,7 @@ To train a stochastic discrete model:
 
 ```
 python -m tensor2tensor.rl.trainer_model_based \
-  --loop_hparams_set=rlmb_long_stochastic_discrete \
+  --loop_hparams_set=rlmb_base_stochastic_discrete \
   --loop_hparams=game=pong \
   --output_dir ~/t2t_train/mb_sd_pong
 ```

From 012796a0658e3632acf45acebb80f62045f0397e Mon Sep 17 00:00:00 2001
From: Anudhyan Boral <anudhyan@google.com>
Date: Fri, 22 Mar 2019 11:03:05 -0700
Subject: [PATCH 1819/2720] Reshape the output explicitly in
 `sparse_message_pass`. This propagates static shape information that is lost
 during the `tf.sparse_reduce_sum` operation.

This would prevent failing downstream tests from failing after broadcasting
support is added in tf.matmul. After this addition, unknown (static) rank on either
operand of tf.matmul would result in unknown rank of the output Tensor.

PiperOrigin-RevId: 239826458
---
 tensor2tensor/layers/message_passing_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/message_passing_attention.py b/tensor2tensor/layers/message_passing_attention.py
index 949527fe2..bf00d6001 100644
--- a/tensor2tensor/layers/message_passing_attention.py
+++ b/tensor2tensor/layers/message_passing_attention.py
@@ -508,7 +508,7 @@ def sparse_message_pass(node_states,
       incoming_edges = tf.tile(incoming_edges, [1, hidden_size])
       final_node_states /= incoming_edges + 1e-7
 
-  return final_node_states
+  return tf.reshape(final_node_states, [n, hidden_size])
 
 
 def multihead_mpnn_attention(node_states,

From 765f65120ebf130030b715f5845863aa0b9b098d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 22 Mar 2019 11:52:25 -0700
Subject: [PATCH 1820/2720] Update to Paracrawl release4, and
 simplify/regularize some problem names.

PiperOrigin-RevId: 239836265
---
 .../data_generators/translate_ende.py         | 25 +++++++++++--------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index 6c62dbbce..37607d6b5 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -49,7 +49,7 @@
 ]
 _ENDE_PARACRAWL_DATASETS = [
     [
-        "https://s3.amazonaws.com/web-language-models/paracrawl/release3/en-de.bicleaner07.tmx.gz",  # pylint: disable=line-too-long
+        "https://s3.amazonaws.com/web-language-models/paracrawl/release4/en-de.bicleaner07.tmx.gz",  # pylint: disable=line-too-long
         ("tmx", "en-de.bicleaner07.tmx.gz")
     ]
 ]
@@ -76,6 +76,7 @@ def source_data_files(self, dataset_split):
 
 @registry.register_problem
 class TranslateEndeWmt32k(TranslateEndeWmt8k):
+  """En-de translation trained on WMT corpus."""
 
   @property
   def approx_vocab_size(self):
@@ -84,6 +85,7 @@ def approx_vocab_size(self):
 
 @registry.register_problem
 class TranslateEndeWmtClean32k(TranslateEndeWmt32k):
+  """En-de translation trained on WMT with further cleaning."""
 
   @property
   def use_vocab_from_other_problem(self):
@@ -95,8 +97,8 @@ def datatypes_to_clean(self):
 
 
 @registry.register_problem
-class TranslateEndeParacrawl32k(translate.TranslateProblem):
-  """Problem spec for Paracrawl en-de translation."""
+class TranslateEndePc32k(translate.TranslateProblem):
+  """En-de translation trained on Paracrawl (bicleaner corpus)."""
 
   @property
   def use_vocab_from_other_problem(self):
@@ -115,8 +117,8 @@ def source_data_files(self, dataset_split):
 
 
 @registry.register_problem
-class TranslateEndeParacrawlClean32k(TranslateEndeParacrawl32k):
-  """Paracrawl en-de Bicleaner corpus, with additional cleaning."""
+class TranslateEndePcClean32k(TranslateEndePc32k):
+  """En-de translation trained on Paracrawl with further cleaning."""
 
   @property
   def datatypes_to_clean(self):
@@ -124,8 +126,8 @@ def datatypes_to_clean(self):
 
 
 @registry.register_problem
-class TranslateEndeWmtParacrawlBicleaner32k(TranslateEndeWmt32k):
-  """WMT en-de corpus with extra data from Paracrawl, cleaned with Bicleaner."""
+class TranslateEndeWmtPc32k(TranslateEndeWmt32k):
+  """En-de translation trained on WMT plus Paracrawl."""
 
   @property
   def use_vocab_from_other_problem(self):
@@ -137,7 +139,8 @@ def additional_training_datasets(self):
 
 
 @registry.register_problem
-class TranslateEndeWmtCleanParacrawl32k(TranslateEndeWmtParacrawlBicleaner32k):
+class TranslateEndeWmtCleanPc32k(TranslateEndeWmtPc32k):
+  """En-de translation trained on cleaned WMT plus Paracrawl."""
 
   @property
   def datatypes_to_clean(self):
@@ -145,7 +148,8 @@ def datatypes_to_clean(self):
 
 
 @registry.register_problem
-class TranslateEndeWmtParacrawlClean32k(TranslateEndeWmtParacrawlBicleaner32k):
+class TranslateEndeWmtPcClean32k(TranslateEndeWmtPc32k):
+  """En-de translation trained on WMT plus cleaned Paracrawl."""
 
   @property
   def datatypes_to_clean(self):
@@ -153,7 +157,8 @@ def datatypes_to_clean(self):
 
 
 @registry.register_problem
-class TranslateEndeWmtParacrawlAllClean32k(TranslateEndeWmtParacrawlClean32k):
+class TranslateEndeWmtCleanPcClean32k(TranslateEndeWmtPcClean32k):
+  """En-de translation trained on cleaned WMT plus cleaned Paracrawl."""
 
   @property
   def datatypes_to_clean(self):

From 70b54f846868af2d9e69737b43b5e84e194e4435 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 22 Mar 2019 12:20:14 -0700
Subject: [PATCH 1821/2720] Add link to paper and website to rl/README.

PiperOrigin-RevId: 239841325
---
 README.md                  |  1 +
 docs/walkthrough.md        |  1 +
 tensor2tensor/rl/README.md | 10 +++-------
 3 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 0f2f1ff01..9e3d1491d 100644
--- a/README.md
+++ b/README.md
@@ -473,5 +473,6 @@ T2T](https://research.googleblog.com/2017/06/accelerating-deep-learning-research
 * [Universal Transformers](https://arxiv.org/abs/1807.03819)
 * [Attending to Mathematical Language with Transformers](https://arxiv.org/abs/1812.02825)
 * [The Evolved Transformer](https://arxiv.org/abs/1901.11117)
+* [Model-Based Reinforcement Learning for Atari](https://arxiv.org/abs/1903.00374)
 
 *Note: This is not an official Google product.*
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 0f2f1ff01..9e3d1491d 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -473,5 +473,6 @@ T2T](https://research.googleblog.com/2017/06/accelerating-deep-learning-research
 * [Universal Transformers](https://arxiv.org/abs/1807.03819)
 * [Attending to Mathematical Language with Transformers](https://arxiv.org/abs/1812.02825)
 * [The Evolved Transformer](https://arxiv.org/abs/1901.11117)
+* [Model-Based Reinforcement Learning for Atari](https://arxiv.org/abs/1903.00374)
 
 *Note: This is not an official Google product.*
diff --git a/tensor2tensor/rl/README.md b/tensor2tensor/rl/README.md
index 19b318859..8ca6fb3e2 100644
--- a/tensor2tensor/rl/README.md
+++ b/tensor2tensor/rl/README.md
@@ -1,13 +1,9 @@
 # Tensor2Tensor Model-Based Reinforcement Learning.
 
-The `rl` package provides the ability to run model-based reinforcement learning
-algorithms using models trained with Tensor2Tensor.
+The `rl` package allows to run reinforcement learning algorithms,
+both model-free (e.g., [Proximal Policy Optimization](https://arxiv.org/abs/1707.06347), train with `trainer_model_free.py`) and model-based ones ([SimPLe](https://arxiv.org/abs/1903.00374), train with `trainer_model_based.py`).
 
-Currently this entails alternating model training and agent training using
-Proximal Policy Optimization (PPO). See `trainer_model_based.py`.
-
-As a baseline, you can also run PPO without the model using
-`trainer_model_free.py`.
+You should be able to reproduce the [Model-Based Reinforcement Learning for Atari](https://arxiv.org/abs/1903.00374) results. [These videos](https://sites.google.com/corp/view/modelbasedrlatari/home) show what to expect from the final models.
 
 To use this package, we recommend Tensorflow 1.13.1 and T2T version 1.13.1.
 You also need to install the Atari dependencies for OpenAI Gym:

From 560644e5d23bdc9a3119826f7952c33d58350f64 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Fri, 22 Mar 2019 23:42:16 +0100
Subject: [PATCH 1822/2720] Model Based RL: Add sticky actions to model-based
 and model-free pipelines. (#1479)

* Add sticky action wrapper, symplify code for environments wrapping.

* Add sticky_actions option to model-based and model-free pipelines.
---
 tensor2tensor/data_generators/gym_env.py      |  5 +-
 tensor2tensor/envs/mujoco_problems.py         |  3 +-
 tensor2tensor/models/research/rl.py           |  1 +
 tensor2tensor/rl/gym_utils.py                 | 81 +++++++++++++------
 tensor2tensor/rl/rl_utils.py                  |  4 +-
 .../rl/trainer_model_based_params.py          |  1 +
 6 files changed, 65 insertions(+), 30 deletions(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 06989c417..46a28ce1c 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -592,7 +592,7 @@ class T2TGymEnv(T2TEnv):
   def __init__(self, base_env_name=None, batch_size=1, grayscale=False,
                resize_height_factor=2, resize_width_factor=2,
                rl_env_max_episode_steps=-1, max_num_noops=0,
-               maxskip_envs=False,
+               maxskip_envs=False, sticky_actions=False,
                should_derive_observation_space=True,
                **kwargs):
     if base_env_name is None:
@@ -606,6 +606,7 @@ def __init__(self, base_env_name=None, batch_size=1, grayscale=False,
     self.resize_width_factor = resize_width_factor
     self.rl_env_max_episode_steps = rl_env_max_episode_steps
     self.maxskip_envs = maxskip_envs
+    self.sticky_actions = sticky_actions
     self._initial_state = None
     self._initial_frames = None
     if not self.name:
@@ -615,7 +616,7 @@ def __init__(self, base_env_name=None, batch_size=1, grayscale=False,
     self._envs = [
         gym_utils.make_gym_env(
             base_env_name, rl_env_max_episode_steps=rl_env_max_episode_steps,
-            maxskip_env=maxskip_envs)
+            maxskip_env=maxskip_envs, sticky_actions=sticky_actions)
         for _ in range(self.batch_size)]
 
     # max_num_noops works only with atari envs.
diff --git a/tensor2tensor/envs/mujoco_problems.py b/tensor2tensor/envs/mujoco_problems.py
index 60f36009e..85adf7695 100644
--- a/tensor2tensor/envs/mujoco_problems.py
+++ b/tensor2tensor/envs/mujoco_problems.py
@@ -37,7 +37,8 @@ def __init__(self):
         gym_utils.gym_env_wrapper, **{
             "rl_env_max_episode_steps": -1,
             "maxskip_env": False,
-            "rendered_env": True
+            "rendered_env": True,
+            "sticky_actions": False
         })
     super(ReacherEnvProblem, self).__init__(
         base_env_name=base_env_name, env_wrapper_fn=wrapper_fn)
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 833425664..dde99d0df 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -375,6 +375,7 @@ def dqn_original_params():
 def rlmf_original():
   return HParams(
       game="pong",
+      sticky_actions=False,
       base_algo="ppo",
       base_algo_params="ppo_original_params",
       batch_size=16,
diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index 9ba522d8f..1ef101b89 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -24,6 +24,24 @@
 import tensorflow as tf
 
 
+class StickyActionEnv(gym.Wrapper):
+  """Based on openai/atari-reset implementation."""
+  def __init__(self, env, p=0.25):
+    gym.Wrapper.__init__(self, env)
+    self.p = p
+    self.last_action = 0
+
+  def step(self, action):
+    if np.random.uniform() < self.p:
+      action = self.last_action
+    self.last_action = action
+    obs, reward, done, info = self.env.step(action)
+    return obs, reward, done, info
+
+  def reset(self, **kwargs):
+    return self.env.reset(**kwargs)
+
+
 class MaxAndSkipEnv(gym.Wrapper):
   """Same wrapper as in OpenAI baselines for comparability of results."""
 
@@ -82,44 +100,55 @@ def reset(self, **kwargs):
     return self.env.render(mode=self.mode)
 
 
-def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env):
-  """Wraps a gym environment. see make_gym_environment for details."""
+def remove_time_limit_wrapper(env):
+  """Removes top level TimeLimit Wrapper.
+
+  Removes TimeLimit Wrapper from top level if exists, throws error if any other
+  TimeLimit Wrapper is present in stack.
+  """
+  if isinstance(env, gym.wrappers.TimeLimit):
+    env = env.env
+  env_ = env
+  while isinstance(env_, gym.Wrapper):
+    if isinstance(env_, gym.wrappers.TimeLimit):
+      raise ValueError("Can remove only top-level TimeLimit gym.Wrapper.")
+    env_ = env_.env
+  return env
+
+
+def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env,
+                    sticky_actions):
+  """Wraps a gym environment. see make_gym_env for details."""
   # rl_env_max_episode_steps is None or int.
   assert ((not rl_env_max_episode_steps) or
           isinstance(rl_env_max_episode_steps, int))
 
-  # If nothing to do, then return the env.
-  if rl_env_max_episode_steps and rl_env_max_episode_steps < 0:
-    if maxskip_env:
-      if isinstance(env, gym.wrappers.TimeLimit):
-        # Unwrap time limit and put it above MaxAndSkip for consistency.
-        max_episode_steps = env._max_episode_steps  # pylint: disable=protected-access
-        env = MaxAndSkipEnv(env.env)
-        env = gym.wrappers.TimeLimit(env, max_episode_steps=max_episode_steps)
-      else:
-        env = MaxAndSkipEnv(env)
-    if rendered_env:
-      env = RenderedEnv(env)
-    return env
-
-  # Sometimes (mostly?) the env is already wrapped in a TimeLimit wrapper, in
-  # which case unwrap it and wrap with the proper time limit requested.
-  if isinstance(env, gym.wrappers.TimeLimit):
-    env = env.env
+  wrap_with_time_limit = ((not rl_env_max_episode_steps) or
+                          rl_env_max_episode_steps >= 0)
+
+  if wrap_with_time_limit:
+    env = remove_time_limit_wrapper(env)
+
+  if sticky_actions:
+    env = StickyActionEnv(env)
 
   if maxskip_env:
-    env = MaxAndSkipEnv(env)
+    env = MaxAndSkipEnv(env)  # pylint: disable=redefined-variable-type
 
   if rendered_env:
     env = RenderedEnv(env)
 
-  return gym.wrappers.TimeLimit(env, max_episode_steps=rl_env_max_episode_steps)
+  if wrap_with_time_limit:
+    env = gym.wrappers.TimeLimit(env,
+                                 max_episode_steps=rl_env_max_episode_steps)
+  return env
 
 
 def make_gym_env(name,
                  rl_env_max_episode_steps=-1,
                  maxskip_env=False,
-                 rendered_env=False):
+                 rendered_env=False,
+                 sticky_actions=False):
   """Create a gym env optionally with a time limit and maxskip wrapper.
 
   NOTE: The returned env may already be wrapped with TimeLimit!
@@ -132,14 +161,14 @@ def make_gym_env(name,
     maxskip_env: whether to also use MaxAndSkip wrapper before time limit.
     rendered_env: whether to force render for observations. Use this for
       environments that are not natively rendering the scene for observations.
+    sticky_actions: whether to use sticky_actions before MaxAndSkip wrapper.
 
   Returns:
-    An instance of `gym.Env` or `gym.wrappers.TimeLimit` with the requested
-    step limit.
+    An instance of `gym.Env` or `gym.Wrapper`.
   """
   env = gym.make(name)
   return gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env,
-                         rendered_env)
+                         rendered_env, sticky_actions)
 
 
 def register_gym_env(class_entry_point, version="v0", kwargs=None):
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 2e4c7660f..ffdcb5ef0 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -307,7 +307,9 @@ def setup_env(hparams,
       resize_height_factor=hparams.resize_height_factor,
       rl_env_max_episode_steps=rl_env_max_episode_steps,
       max_num_noops=max_num_noops,
-      maxskip_envs=maxskip_envs)
+      maxskip_envs=maxskip_envs,
+      sticky_actions=hparams.sticky_actions
+  )
   return env
 
 
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 00f384507..37eb47ddc 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -89,6 +89,7 @@ def _rlmb_base():
       eval_rl_env_max_episode_steps=1000,
 
       game="pong",
+      sticky_actions=False,
       # If set, use this as the gym env name, instead of changing game mode etc.
       rl_env_name="",
       # Controls whether we should derive observation space, do some

From 6bd50f81fabefecdf70f0a79bb0a6d6340aa65f1 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Fri, 22 Mar 2019 15:53:22 -0700
Subject: [PATCH 1823/2720] Merge of PR #1479

PiperOrigin-RevId: 239880513
---
 tensor2tensor/data_generators/gym_env.py | 1 +
 tensor2tensor/models/research/rl.py      | 2 +-
 tensor2tensor/rl/gym_utils.py            | 7 +++++++
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 46a28ce1c..67aedef9b 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -44,6 +44,7 @@
 )
 
 
+# pylint: disable=g-complex-comprehension
 class Observation(object):
   """Encoded observations.
 
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index dde99d0df..90c74aeef 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -248,7 +248,7 @@ def make_simulated_env_kwargs(real_env, hparams, **extra_kwargs):
       (hparams, ["frame_stack_size", "intrinsic_reward_scale"])
   ]
   kwargs = {
-      attr: getattr(obj, attr)
+      attr: getattr(obj, attr)  # pylint: disable=g-complex-comprehension
       for (obj, attrs) in objs_and_attrs for attr in attrs
   }
   kwargs["model_name"] = hparams.generative_model
diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index 1ef101b89..e1b6032b9 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -26,6 +26,7 @@
 
 class StickyActionEnv(gym.Wrapper):
   """Based on openai/atari-reset implementation."""
+
   def __init__(self, env, p=0.25):
     gym.Wrapper.__init__(self, env)
     self.p = p
@@ -105,6 +106,12 @@ def remove_time_limit_wrapper(env):
 
   Removes TimeLimit Wrapper from top level if exists, throws error if any other
   TimeLimit Wrapper is present in stack.
+
+  Args:
+    env: environment
+
+  Returns:
+    the env with removed time limit wrapper.
   """
   if isinstance(env, gym.wrappers.TimeLimit):
     env = env.env

From ff1fd68d8ff360efedbfac6f6e81e28f4b2122ab Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 22 Mar 2019 16:19:15 -0700
Subject: [PATCH 1824/2720] Allow autoregressive output frame generation and
 relu as non-linearity in SD video models (both off by default).

PiperOrigin-RevId: 239885106
---
 tensor2tensor/layers/discretization.py        | 45 ++++++-----
 tensor2tensor/models/video/base.py            | 17 ++--
 .../models/video/basic_deterministic.py       | 78 +++++++++++++++++--
 .../video/basic_deterministic_params.py       |  6 ++
 .../models/video/basic_stochastic.py          |  5 +-
 5 files changed, 120 insertions(+), 31 deletions(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 506fd6864..970b705e4 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -18,7 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from functools import partial
+from functools import partial  # pylint: disable=g-importing-member
 
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_image_attention as cia
@@ -27,7 +27,7 @@
 import tensorflow as tf
 import tensorflow_probability as tfp
 
-from tensorflow.python.training import moving_averages
+from tensorflow.python.training import moving_averages  # pylint: disable=g-direct-tensorflow-import
 
 
 def project_hidden(x, projection_tensors, hidden_size, num_blocks):
@@ -789,8 +789,8 @@ def discrete_bottleneck(inputs,
 
 
 def predict_bits_with_lstm(prediction_source, state_size, total_num_bits,
-                           target_bits=None, bits_at_once=8, temperature=1.0,
-                           dropout=0.1):
+                           target_bits=None, extra_inputs=None,
+                           bits_at_once=8, temperature=1.0, dropout=0.1):
   """Predict a sequence of bits (a latent) with LSTM, both training and infer.
 
   Given a tensor on which the predictions are based (prediction_source), we use
@@ -807,6 +807,8 @@ def predict_bits_with_lstm(prediction_source, state_size, total_num_bits,
     total_num_bits: python integer, how many bits in total to predict.
     target_bits: a tensor of shape [batch_size, total_num_bits] used during
       training as the target to predict; each element should be -1 or 1.
+    extra_inputs: a Tensor [batch_size, total_num_bits // bits_at_once, d]
+      of additional inputs, passed as additional LSTM inputs.
     bits_at_once: pytho integer, how many bits to predict at once.
     temperature: python float, temperature used for sampling during inference.
     dropout: float, the amount of dropout to aply during training (0.1 default).
@@ -824,7 +826,7 @@ def predict_bits_with_lstm(prediction_source, state_size, total_num_bits,
     discrete_embed = tf.layers.Dense(state_size, name="discrete_embed")
     batch_size = common_layers.shape_list(prediction_source)[0]
     layer_pred = tf.layers.flatten(prediction_source)
-    prediction = tf.layers.dense(layer_pred, state_size, name="istate")
+    first_lstm_input = tf.layers.dense(layer_pred, state_size, name="istate")
     c_state = tf.layers.dense(layer_pred, state_size, name="cstate")
     m_state = tf.layers.dense(layer_pred, state_size, name="mstate")
     state = (c_state, m_state)
@@ -832,13 +834,16 @@ def predict_bits_with_lstm(prediction_source, state_size, total_num_bits,
     # Prediction mode if no targets are given.
     if target_bits is None:
       outputs = []
+      lstm_input = first_lstm_input
       for i in range(total_num_bits // bits_at_once):
-        output, state = lstm_cell(prediction, state)
+        if extra_inputs is not None:
+          lstm_input = tf.concat([lstm_input, extra_inputs[:, i, :]], axis=1)
+        output, state = lstm_cell(lstm_input, state)
         discrete_logits = discrete_predict(output)
         discrete_samples = common_layers.sample_with_temperature(
             discrete_logits, temperature)
         outputs.append(tf.expand_dims(discrete_samples, axis=1))
-        prediction = discrete_embed(tf.one_hot(discrete_samples, 256))
+        lstm_input = discrete_embed(tf.one_hot(discrete_samples, 256))
       outputs = tf.concat(outputs, axis=1)
       outputs = int_to_bit(outputs, bits_at_once)
       outputs = tf.reshape(outputs, [batch_size, total_num_bits])
@@ -846,25 +851,29 @@ def predict_bits_with_lstm(prediction_source, state_size, total_num_bits,
 
     # Training mode, calculating loss.
     assert total_num_bits % bits_at_once == 0
-    d_pred = tf.reshape(tf.maximum(tf.stop_gradient(target_bits), 0), [
+    target_bits = tf.reshape(tf.maximum(tf.stop_gradient(target_bits), 0), [
         batch_size, total_num_bits // bits_at_once, bits_at_once])
-    d_int = bit_to_int(d_pred, bits_at_once)
-    tf.summary.histogram("target_integers", tf.reshape(d_int, [-1]))
-    d_hot = tf.one_hot(d_int, 2**bits_at_once, axis=-1)
-    d_pred = discrete_embed(d_hot)
-    d_pred = tf.nn.dropout(d_pred, 1.0 - dropout)
-    pred = tf.concat([tf.expand_dims(prediction, axis=1), d_pred], axis=1)
+    target_ints = bit_to_int(target_bits, bits_at_once)
+    tf.summary.histogram("target_integers", tf.reshape(target_ints, [-1]))
+    target_hot = tf.one_hot(target_ints, 2**bits_at_once, axis=-1)
+    target_embedded = discrete_embed(target_hot)
+    target_embedded = tf.nn.dropout(target_embedded, 1.0 - dropout)
+    teacher_input = tf.concat(
+        [tf.expand_dims(first_lstm_input, axis=1), target_embedded], axis=1)
     outputs = []
     for i in range(total_num_bits // bits_at_once):
-      output, state = lstm_cell(pred[:, i, :], state)
+      lstm_input = teacher_input[:, i, :]
+      if extra_inputs is not None:
+        lstm_input = tf.concat([lstm_input, extra_inputs[:, i, :]], axis=1)
+      output, state = lstm_cell(lstm_input, state)
       outputs.append(tf.expand_dims(output, axis=1))
     outputs = tf.concat(outputs, axis=1)
     outputs = tf.nn.dropout(outputs, 1.0 - dropout)
     d_int_pred = discrete_predict(outputs)
     pred_loss = tf.losses.sparse_softmax_cross_entropy(
-        logits=d_int_pred, labels=d_int)
+        logits=d_int_pred, labels=target_ints)
     pred_loss = tf.reduce_mean(pred_loss)
-    return target_bits, pred_loss
+    return d_int_pred, pred_loss
 
 
 # New API for discretization bottlenecks:
@@ -986,7 +995,7 @@ def loss_with_update():
         return beta * e_loss
 
   # Loss, also do update if requested.
-  if do_update is True:
+  if do_update:
     loss = loss_with_update()
   else:
     loss = tf.cond(do_update, loss_with_update, lambda: beta * e_loss)
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index dbae5eb1a..bd31906c8 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -34,7 +34,7 @@
 
 
 def flat_lists(list_of_lists):
-  return [x for l in list_of_lists for x in l]
+  return [x for l in list_of_lists for x in l]  # pylint: disable=g-complex-comprehension
 
 
 def pixels_from_softmax(frame_logits, pure_sampling=False,
@@ -388,7 +388,8 @@ def get_sampled_frame(self, pred_frame):
       sampled_frame = pixels_from_softmax(
           sampled_frame, temperature=self.hparams.pixel_sampling_temperature)
       # TODO(lukaszkaiser): this should be consistent with modality.bottom()
-      sampled_frame = common_layers.standardize_images(sampled_frame)
+      # sampled_frame = common_layers.standardize_images(sampled_frame)
+      sampled_frame = tf.to_float(sampled_frame)
     else:
       x = common_layers.convert_real_to_rgb(pred_frame)
       x = x - tf.stop_gradient(x + tf.round(x))
@@ -540,9 +541,10 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
         target_frames.append(tf.identity(target_frame))
 
         with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
-          func_in = (frames, actions, rewards, target_frame,
-                     internal_states, video_features)
-          func_out = self.next_frame(*func_in)
+          float_frames = [tf.to_float(frame) for frame in frames]
+          func_out = self.next_frame(
+              float_frames, actions, rewards, tf.to_float(target_frame),
+              internal_states, video_features)
           res_frame, res_reward, res_policy, res_value, res_extra_loss, \
               internal_states = func_out
           res_frames.append(res_frame)
@@ -579,7 +581,7 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
 
         # Scheduled sampling during training.
         if self.is_training:
-          groundtruth_items = [target_frame]
+          groundtruth_items = [tf.to_float(target_frame)]
           generated_items = [sampled_frame]
           ss_frame, = self.get_scheduled_sample_inputs(
               done_warm_start, groundtruth_items, generated_items, ss_func)
@@ -608,7 +610,8 @@ def __process(self, all_frames, all_actions, all_rewards, all_raw_frames):
       sampled_frames = sampled_frames[hparams.video_num_input_frames-1:]
       target_frames = target_frames[hparams.video_num_input_frames-1:]
 
-    self.visualize_predictions(sampled_frames, target_frames)
+    self.visualize_predictions(
+        sampled_frames, [tf.to_float(f) for f in target_frames])
 
     output_frames = tf.stack(res_frames, axis=1)
     targets = output_frames
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 882fc01b8..ff712f5e8 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -22,6 +22,7 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
+from tensor2tensor.layers import discretization
 from tensor2tensor.models.video import base
 from tensor2tensor.models.video import basic_deterministic_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
@@ -43,13 +44,16 @@ def inject_latent(self, layer, inputs, target, action):
 
   def middle_network(self, layer, internal_states):
     # Run a stack of convolutions.
+    activation_fn = common_layers.belu
+    if self.hparams.activation_fn == "relu":
+      activation_fn = tf.nn.relu
     x = layer
     kernel1 = (3, 3)
     filters = common_layers.shape_list(x)[-1]
     for i in range(self.hparams.num_hidden_layers):
       with tf.variable_scope("layer%d" % i):
         y = tf.nn.dropout(x, 1.0 - self.hparams.residual_dropout)
-        y = tf.layers.conv2d(y, filters, kernel1, activation=common_layers.belu,
+        y = tf.layers.conv2d(y, filters, kernel1, activation=activation_fn,
                              strides=(1, 1), padding="SAME")
         if i == 0:
           x = y
@@ -70,6 +74,12 @@ def next_frame(self, frames, actions, rewards, target_frame,
     filters = hparams.hidden_size
     kernel2 = (4, 4)
     action = actions[-1]
+    activation_fn = common_layers.belu
+    if self.hparams.activation_fn == "relu":
+      activation_fn = tf.nn.relu
+
+    # Normalize frames.
+    frames = [common_layers.standardize_images(f) for f in frames]
 
     # Stack the inputs.
     if internal_states is not None and hparams.concat_internal_states:
@@ -102,7 +112,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
         if i < hparams.filter_double_steps:
           filters *= 2
         x = common_attention.add_timing_signal_nd(x)
-        x = tf.layers.conv2d(x, filters, kernel2, activation=common_layers.belu,
+        x = tf.layers.conv2d(x, filters, kernel2, activation=activation_fn,
                              strides=(2, 2), padding="SAME")
         x = common_layers.layer_norm(x)
 
@@ -121,7 +131,8 @@ def next_frame(self, frames, actions, rewards, target_frame,
           x, action, "action_enc", hparams.action_injection)
 
     # Inject latent if present. Only for stochastic models.
-    x, extra_loss = self.inject_latent(x, frames, target_frame, action)
+    norm_target_frame = common_layers.standardize_images(target_frame)
+    x, extra_loss = self.inject_latent(x, frames, norm_target_frame, action)
 
     x_mid = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
     x, internal_states = self.middle_network(x, internal_states)
@@ -137,7 +148,7 @@ def next_frame(self, frames, actions, rewards, target_frame,
         if i >= hparams.num_compress_steps - hparams.filter_double_steps:
           filters //= 2
         x = tf.layers.conv2d_transpose(
-            x, filters, kernel2, activation=common_layers.belu,
+            x, filters, kernel2, activation=activation_fn,
             strides=(2, 2), padding="SAME")
         y = layer_inputs[i]
         shape = common_layers.shape_list(y)
@@ -148,7 +159,64 @@ def next_frame(self, frames, actions, rewards, target_frame,
     # Cut down to original size.
     x = x[:, :inputs_shape[1], :inputs_shape[2], :]
     x_fin = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
-    if self.is_per_pixel_softmax:
+    if hparams.do_autoregressive_rnn:
+      # If enabled, we predict the target frame autoregregressively using rnns.
+      # To this end, the current prediciton is flattened into one long sequence
+      # of sub-pixels, and so is the target frame. Each sub-pixel (RGB value,
+      # from 0 to 255) is predicted with an RNN. To avoid doing as many steps
+      # as width * height * channels, we only use a number of pixels back,
+      # as many as hparams.autoregressive_rnn_lookback.
+      with tf.variable_scope("autoregressive_rnn"):
+        batch_size = common_layers.shape_list(frames[0])[0]
+        # Height, width, channels and lookback are the constants we need.
+        h, w = inputs_shape[1], inputs_shape[2]  # 105, 80 on Atari games
+        c = hparams.problem.num_channels
+        lookback = hparams.autoregressive_rnn_lookback
+        assert (h * w) % lookback == 0, "Number of pixels must divide lookback."
+        m = (h * w) // lookback  # Batch size multiplier for the RNN.
+        # These are logits that will be used as inputs to the RNN.
+        rnn_inputs = tf.layers.dense(x, c * 64, name="rnn_inputs")
+        # They are of shape [batch_size, h, w, c, 64], reshaping now.
+        rnn_inputs = tf.reshape(rnn_inputs, [batch_size * m, lookback * c, 64])
+        # Same for the target frame.
+        rnn_target = tf.reshape(target_frame, [batch_size * m, lookback * c])
+        # Construct rnn starting state: flatten rnn_inputs, apply a relu layer.
+        rnn_start_state = tf.nn.relu(tf.layers.dense(tf.nn.relu(
+            tf.layers.flatten(rnn_inputs)), 256, name="rnn_start_state"))
+        # Our RNN function API is on bits, each subpixel has 8 bits.
+        total_num_bits = lookback * c * 8
+        # We need to provide RNN targets as bits (due to the API).
+        rnn_target_bits = discretization.int_to_bit(rnn_target, 8)
+        rnn_target_bits = tf.reshape(
+            rnn_target_bits, [batch_size * m, total_num_bits])
+        if self.is_training:
+          # Run the RNN in training mode, add it's loss to the losses.
+          rnn_predict, rnn_loss = discretization.predict_bits_with_lstm(
+              rnn_start_state, 128, total_num_bits, target_bits=rnn_target_bits,
+              extra_inputs=rnn_inputs)
+          extra_loss += rnn_loss
+          # We still use non-RNN predictions too in order to guide the network.
+          x = tf.layers.dense(x, c * 256, name="logits")
+          x = tf.reshape(x, [batch_size, h, w, c, 256])
+          rnn_predict = tf.reshape(rnn_predict, [batch_size, h, w, c, 256])
+          # Mix non-RNN and RNN predictions so that after warmup the RNN is 90%.
+          x = tf.reshape(tf.nn.log_softmax(x), [batch_size, h, w, c * 256])
+          rnn_predict = tf.nn.log_softmax(rnn_predict)
+          rnn_predict = tf.reshape(rnn_predict, [batch_size, h, w, c * 256])
+          alpha = 0.9 * common_layers.inverse_lin_decay(
+              hparams.autoregressive_rnn_warmup_steps)
+          x = alpha * rnn_predict + (1.0 - alpha) * x
+        else:
+          # In prediction mode, run the RNN without any targets.
+          bits, _ = discretization.predict_bits_with_lstm(
+              rnn_start_state, 128, total_num_bits, extra_inputs=rnn_inputs,
+              temperature=0.0)  # No sampling from this RNN, just greedy.
+          # The output is in bits, get back the predicted pixels.
+          bits = tf.reshape(bits, [batch_size * m, lookback * c, 8])
+          ints = discretization.bit_to_int(tf.maximum(bits, 0), 8)
+          ints = tf.reshape(ints, [batch_size, h, w, c])
+          x = tf.reshape(tf.one_hot(ints, 256), [batch_size, h, w, c * 256])
+    elif self.is_per_pixel_softmax:
       x = tf.layers.dense(x, hparams.problem.num_channels * 256, name="logits")
     else:
       x = tf.layers.dense(x, hparams.problem.num_channels, name="logits")
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 8d8a2c6ff..b40307f7b 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -47,6 +47,12 @@ def next_frame_basic_deterministic():
   hparams.add_hparam("filter_double_steps", 2)
   hparams.add_hparam("pixel_sampling_temperature", 0.0)
   hparams.add_hparam("concat_internal_states", False)
+  hparams.add_hparam("do_autoregressive_rnn", False)
+  hparams.add_hparam("autoregressive_rnn_lookback", 8)
+  hparams.add_hparam("autoregressive_rnn_warmup_steps", 8000)
+  hparams.add_hparam("activation_fn", "belu")
+  hparams.bottom["inputs"] = modalities.video_identity_bottom
+  hparams.bottom["targets"] = modalities.video_identity_bottom
   return hparams
 
 
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index db08ae864..7fa519e1a 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -128,6 +128,9 @@ def inject_latent(self, layer, inputs, target, action):
     filters = hparams.hidden_size
     kernel = (4, 4)
     layer_shape = common_layers.shape_list(layer)
+    activation_fn = common_layers.belu
+    if hparams.activation_fn == "relu":
+      activation_fn = tf.nn.relu
 
     def add_bits(layer, bits):
       z_mul = tfl.dense(bits, final_filters, name="unbottleneck_mul")
@@ -169,7 +172,7 @@ def add_bits(layer, bits):
             filters *= 2
           x = common_attention.add_timing_signal_nd(x)
           x = tfl.conv2d(x, filters, kernel,
-                         activation=common_layers.belu,
+                         activation=activation_fn,
                          strides=(2, 2), padding="SAME")
           x = common_layers.layer_norm(x)
     else:

From d63dc7f94e77b067a10a2d3c8dd1798da7304af4 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Sun, 24 Mar 2019 19:56:45 -0700
Subject: [PATCH 1825/2720] Use pmap to make trax work in multi-device mode.

PiperOrigin-RevId: 240068295
---
 .../trax/configs/transformer_lm1b_8gb.gin     |  2 +-
 tensor2tensor/trax/trax.py                    | 41 +++++++++++++++----
 2 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index caa237aed..a4b692d65 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -15,7 +15,7 @@ inputs.dataset_name = 't2t_languagemodel_lm1b32k'
 
 # Parameters for mask:
 # ==============================================================================
-mask.mask_id = 0
+masked_mean.mask_id = 0
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index d84a36275..8ed32e64c 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -226,15 +226,36 @@ def epochs(steps=None, epoch_steps=1):
 
 def _jit_update_fun(predict_fun, loss_fun, optimizer, lr_fun):
   """Get jit-ed update function for loss, optimizer, learning rate function."""
-  @jax.jit
-  def update(i, opt_state, batch):
+  @functools.partial(jax.pmap, axis_name="batch")
+  def mapped_update(i, opt_state, batch):
     _, opt_update = optimizer(lr_fun)
     params = jax_opt.get_params(opt_state)
-    return opt_update(i, jax.grad(loss_fun)(
-        params, batch, predict_fun), opt_state)
+    grads = jax.grad(loss_fun)(params, batch, predict_fun)
+    grads = jax.tree_util.tree_map(lambda g: jax.lax.psum(g, "batch"), grads)
+    return opt_update(i, grads, opt_state)
+
+  def update(i, opt_state, batch):
+    return mapped_update(jax.replicate(i), opt_state, batch)
+
   return update
 
 
+def reshape_by_device(train_data, num_devices):
+  """Reshape the train_data into a shape [num_devices, ...]."""
+  x, y = train_data
+  x_shape, y_shape = list(x.shape), list(y.shape)
+  assert x_shape[0] == y_shape[0]  # Same batch size.
+  batch_size = x_shape[0]
+  batch_size_per_device = batch_size // num_devices
+  # We require that num_devices divides batch_size evenly.
+  assert batch_size_per_device * num_devices == batch_size
+  # New shapes.
+  new_shape_prefix = [num_devices, batch_size_per_device]
+  x = np.reshape(x, new_shape_prefix + x_shape[1:])
+  y = np.reshape(y, new_shape_prefix + y_shape[1:])
+  return x, y
+
+
 @gin.configurable(blacklist=["output_dir"])
 def train(output_dir,
           model=gin.REQUIRED,
@@ -244,6 +265,7 @@ def train(output_dir,
           train_steps=1000,
           eval_steps=10,
           eval_frequency=100,
+          num_devices=None,
           run_debug_step=False):
   """Train the model on the inputs.
 
@@ -260,12 +282,14 @@ def train(output_dir,
     eval_steps: int, num of steps per evaluation. If None or 0, eval disabled.
     eval_frequency: int, how often to run evaluation (every eval_frequency
       steps). If None or 0, eval disabled.
+    num_devices: how many devices to use (if None, default, use all available)
     run_debug_step: bool, if True, will run the model and loss without @jit for
       one step.
 
   Returns:
     trax.State
   """
+  num_devices = num_devices or jax.lib.xla_bridge.device_count()
   rng = random.PRNGKey(0)
   gfile.makedirs(output_dir)
   # Create summary writers and history.
@@ -291,7 +315,7 @@ def model_predict(x, y, **kwargs):
   step = state.step or 0
   params_initializer = lambda: model_init([-1] + list(inputs.input_shape))[1]
   params = state.params or params_initializer()
-  opt_state = opt_init(params)
+  opt_state = jax.replicate(opt_init(params))
 
   # jit model_predict and update so they're fast
   jit_model_predict = jax.jit(model_predict)  # for evaluation
@@ -302,7 +326,7 @@ def model_predict(x, y, **kwargs):
   epoch_steps = itertools.chain([1,  # first epoch only 1 step
                                  eval_frequency - 1],
                                 itertools.repeat(eval_frequency))
-  step_log(step, "Starting training")
+  step_log(step, "Starting training using %d devices" % num_devices)
 
   # Non-compiled debug step helps find problems in models easier.
   if run_debug_step:
@@ -318,7 +342,8 @@ def model_predict(x, y, **kwargs):
 
     for _ in range(epoch_steps):
       # Train
-      opt_state = jit_update_fun(step, opt_state, next(train_stream))
+      next_train_batch = reshape_by_device(next(train_stream), num_devices)
+      opt_state = jit_update_fun(step, opt_state, next_train_batch)
       step += 1
 
       # LR log
@@ -335,7 +360,7 @@ def model_predict(x, y, **kwargs):
                       epoch_steps / epoch_time, step=step)
 
     # Evaluate
-    params = jax_opt.get_params(opt_state)
+    params = jax_opt.get_params(jax.unreplicate(opt_state))
     evaluate_train_and_eval(
         step=step,
         inputs=inputs,

From 16d59abb2e80a9f06646e587a995a39aa128358a Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 25 Mar 2019 10:55:05 -0700
Subject: [PATCH 1826/2720] Queries and values don't need to have the same
 depth in dot_product_unmasked_self_attention_relative_2d.

PiperOrigin-RevId: 240174892
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index a56962af8..73735b408 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2120,7 +2120,7 @@ def dot_product_unmasked_self_attention_relative_2d(
     # This calculation only works for self attention.
     # q, k and v must therefore have the same shape.
     q.get_shape().assert_is_compatible_with(k.get_shape())
-    q.get_shape().assert_is_compatible_with(v.get_shape())
+    q.get_shape()[:-1].assert_is_compatible_with(v.get_shape()[:-1])
 
     (height, width) = (common_layers.shape_list(q)[2],
                        common_layers.shape_list(q)[3])

From 06dafa8b8a0f3b6c12473587839102b32d8b6f72 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 25 Mar 2019 12:02:06 -0700
Subject: [PATCH 1827/2720] Use GFile

PiperOrigin-RevId: 240189694
---
 tensor2tensor/data_generators/wikisum/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/wikisum/utils.py b/tensor2tensor/data_generators/wikisum/utils.py
index 09e60e7a1..396172216 100644
--- a/tensor2tensor/data_generators/wikisum/utils.py
+++ b/tensor2tensor/data_generators/wikisum/utils.py
@@ -120,7 +120,7 @@ def wet_records(wet_filepath):
   if wet_filepath.endswith('.gz'):
     fopen = gzip.open
   else:
-    fopen = tf.gfile.FastGFile
+    fopen = tf.gfile.GFile
 
   with fopen(wet_filepath) as f:
     for record in wet_records_from_file_obj(f):

From 8c9b80cf5ef9675d26010218cff9b8fed50810c6 Mon Sep 17 00:00:00 2001
From: Piotr Milos <piotr.milos@codilime.com>
Date: Mon, 25 Mar 2019 20:38:43 +0100
Subject: [PATCH 1828/2720] rl notebook fixes (#1518)

---
 tensor2tensor/notebooks/hello_t2t-rl.ipynb | 1890 +++++++++++++++++++-
 1 file changed, 1889 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/notebooks/hello_t2t-rl.ipynb b/tensor2tensor/notebooks/hello_t2t-rl.ipynb
index 21ca2909c..abd502819 100644
--- a/tensor2tensor/notebooks/hello_t2t-rl.ipynb
+++ b/tensor2tensor/notebooks/hello_t2t-rl.ipynb
@@ -1 +1,1889 @@
-{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"name":"hello_t2t-rl.ipynb","version":"0.3.2","provenance":[{"file_id":"1nQvfx1EzY3ElJUy-FVF1G16okSbkeUa2","timestamp":1553274233669}],"collapsed_sections":[]},"kernelspec":{"name":"python3","display_name":"Python 3"},"accelerator":"GPU"},"cells":[{"metadata":{"id":"xCLcAmON-m2i","colab_type":"text"},"cell_type":"markdown","source":["# Tensor2Tensor Reinforcement Learning\n","\n","The `rl` package provides the ability to run model-free and model-based reinforcement learning algorithms.\n","\n","Currently, we support the Proximal Policy Optimization ([PPO](https://arxiv.org/abs/1707.06347)) and Simulated Policy Learning ([SimPLe](https://arxiv.org/abs/1903.00374)).\n","\n","Below you will find examples of PPO training using `trainer_model_free.py` and SimPLe traning using `trainer_model_based.py`.\n"]},{"metadata":{"id":"RW7gEGp3e87G","colab_type":"code","colab":{},"cellView":"form"},"cell_type":"code","source":["#@title\n","# Copyright 2018 Google LLC.\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","# https://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."],"execution_count":0,"outputs":[]},{"metadata":{"id":"pq0BqXm4-3gJ","colab_type":"code","outputId":"6086719f-6268-4b61-8fa3-d251eda24c97","executionInfo":{"status":"ok","timestamp":1553273826475,"user_tz":-60,"elapsed":20650,"user":{"displayName":"Piotr Miłoś","photoUrl":"https://lh3.googleusercontent.com/-050ZBEGpNAA/AAAAAAAAAAI/AAAAAAAAk9g/r6cv_J6J5qA/s64/photo.jpg","userId":"12158759908531801397"}},"colab":{"base_uri":"https://localhost:8080/","height":163}},"cell_type":"code","source":["!pip install -q -U tensor2tensor==1.13.1\n","!pip install -q tensorflow==1.13.1\n","!pip install -q gym[atari]"],"execution_count":0,"outputs":[{"output_type":"stream","text":["\u001b[K    100% |████████████████████████████████| 1.3MB 9.4MB/s \n","\u001b[K    100% |████████████████████████████████| 215kB 27.3MB/s \n","\u001b[K    100% |████████████████████████████████| 143kB 29.6MB/s \n","\u001b[K    100% |████████████████████████████████| 21.1MB 1.7MB/s \n","\u001b[K    100% |████████████████████████████████| 409kB 24.7MB/s \n","\u001b[K    100% |████████████████████████████████| 296kB 25.0MB/s \n","\u001b[K    100% |████████████████████████████████| 61kB 21.5MB/s \n","\u001b[?25h  Building wheel for pypng (setup.py) ... \u001b[?25ldone\n","\u001b[?25h  Building wheel for opt-einsum (setup.py) ... \u001b[?25ldone\n","\u001b[?25h"],"name":"stdout"}]},{"metadata":{"id":"R7-Ni-39DGZW","colab_type":"code","colab":{}},"cell_type":"code","source":["# Helper function for playing videos in the colab.\n","def play_video(path):\n","  from IPython.core.magics.display import HTML\n","  display_path = \"/nbextensions/vid.mp4\"\n","  display_abs_path = \"/usr/local/share/jupyter\" + display_path\n","  !rm -f $display_abs_path\n","  !ffmpeg -loglevel error -i $path $display_abs_path\n","  return HTML(\"\"\"\n","    <video width=\"640\" height=\"480\" controls>\n","      <source src=\"{}\" type=\"video/mp4\">\n","    </video>\n","  \"\"\".format(display_path))"],"execution_count":0,"outputs":[]},{"metadata":{"id":"pueuiKUmAOUT","colab_type":"text"},"cell_type":"markdown","source":["# Play using a pre-trained policy\n","\n","We provide pretrained policies for the following games from the Atari Learning Environment ( [ALE](https://github.com/mgbellemare/Arcade-Learning-Environment)) : alien,\n","amidar,\n"," assault,\n"," asterix,\n"," asteroids,\n"," atlantis,\n"," bank_heist,\n"," battle_zone,\n"," beam_rider,\n"," bowling,\n"," boxing,\n"," breakout,\n"," chopper_command,\n"," crazy_climber,\n"," demon_attack,\n"," fishing_derby,\n"," freeway,\n"," frostbite,\n"," gopher,\n"," gravitar,\n"," hero,\n"," ice_hockey,\n"," jamesbond,\n"," kangaroo,\n"," krull,\n"," kung_fu_master,\n"," ms_pacman,\n"," name_this_game,\n"," pong,\n"," private_eye,\n"," qbert,\n"," riverraid,\n"," road_runner,\n"," seaquest,\n"," up_n_down,\n"," yars_revenge.\n"," \n"," We have 5 checkpoints for each game saved on Google Storage. To get run id to determine the storage path:"]},{"metadata":{"id":"x9pKfNbDFfVh","colab_type":"code","outputId":"97e763cc-caaa-49c8-e532-fcbde828d1a2","executionInfo":{"status":"ok","timestamp":1553274151100,"user_tz":-60,"elapsed":6162,"user":{"displayName":"Piotr Miłoś","photoUrl":"https://lh3.googleusercontent.com/-050ZBEGpNAA/AAAAAAAAAAI/AAAAAAAAk9g/r6cv_J6J5qA/s64/photo.jpg","userId":"12158759908531801397"}},"colab":{"base_uri":"https://localhost:8080/","height":147}},"cell_type":"code","source":["# experiment_id is an integer from [0, 4].\n","def get_run_dir(game, experiment_id):\n","  from tensor2tensor.data_generators.gym_env import ATARI_GAMES_WITH_HUMAN_SCORE_NICE\n","  EXPERIMENTS_PER_GAME = 5\n","  run_id = ATARI_GAMES_WITH_HUMAN_SCORE_NICE.index(game) * EXPERIMENTS_PER_GAME + experiment_id + 1\n","  return \"gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/{}\".format(run_id)\n","\n","get_run_dir('pong', 2)"],"execution_count":0,"outputs":[{"output_type":"stream","text":["\n","WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n","For more information, please see:\n","  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n","  * https://github.com/tensorflow/addons\n","If you depend on functionality not listed there, please file an issue.\n","\n"],"name":"stdout"},{"output_type":"execute_result","data":{"text/plain":["'gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/143'"]},"metadata":{"tags":[]},"execution_count":4}]},{"metadata":{"id":"77fFdm-cFEOB","colab_type":"text"},"cell_type":"markdown","source":["To evaluate and generate videos for a pretrained policy on Pong:"]},{"metadata":{"id":"X-nGlbuTAQXj","colab_type":"code","outputId":"888968f2-f551-4a0f-9fc7-074a949362d6","executionInfo":{"status":"ok","timestamp":1553271580737,"user_tz":-60,"elapsed":842128,"user":{"displayName":"Piotr Kozakowski","photoUrl":"","userId":"01014928596539690143"}},"colab":{"base_uri":"https://localhost:8080/","height":17088}},"cell_type":"code","source":["game = 'pong'\n","run_id = get_run_dir(game, 1)\n","!python -m tensor2tensor.rl.evaluator \\\n","  --loop_hparams_set=rlmb_long_stochastic_discrete \\\n","  --loop_hparams=game=$game,eval_max_num_noops=8,eval_sampling_temps=[0.5] \\\n","  --policy_dir=gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/$run_id/policy \\\n","  --eval_metrics_dir=pong_pretrained \\\n","  --debug_video_path=pong_pretrained \\\n","  --num_debug_videos=4"],"execution_count":0,"outputs":[{"output_type":"stream","text":["\n","WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n","For more information, please see:\n","  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n","  * https://github.com/tensorflow/addons\n","If you depend on functionality not listed there, please file an issue.\n","\n","INFO:tensorflow:Overriding hparams in rlmb_long_stochastic_discrete with game=pong,eval_max_num_noops=8,eval_sampling_temps=[0.5]\n","INFO:tensorflow:Evaluating metric mean_reward/eval/sampling_temp_0.5_max_noops_8_unclipped\n","2019-03-22 16:05:45.007030: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 2300000000 Hz\n","2019-03-22 16:05:45.007306: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x2697860 executing computations on platform Host. Devices:\n","2019-03-22 16:05:45.007346: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): <undefined>, <undefined>\n","2019-03-22 16:05:45.105281: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:998] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n","2019-03-22 16:05:45.105857: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x2697440 executing computations on platform CUDA. Devices:\n","2019-03-22 16:05:45.105908: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): Tesla K80, Compute Capability 3.7\n","2019-03-22 16:05:45.106380: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1433] Found device 0 with properties: \n","name: Tesla K80 major: 3 minor: 7 memoryClockRate(GHz): 0.8235\n","pciBusID: 0000:00:04.0\n","totalMemory: 11.17GiB freeMemory: 11.10GiB\n","2019-03-22 16:05:45.106420: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 16:05:45.499212: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 16:05:45.499307: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 16:05:45.499332: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 16:05:45.499671: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:42] Overriding allow_growth setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.\n","2019-03-22 16:05:45.499741: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","INFO:tensorflow:Using DummyPolicyProblem for the policy.\n","INFO:tensorflow:Setting T2TModel mode to 'train'\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Colocations handled automatically by placer.\n","INFO:tensorflow:Using variable initializer: orthogonal\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/utils/t2t_model.py:1358: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.cast instead.\n","INFO:tensorflow:Transforming feature 'input_action' with symbol_modality_6_64.bottom\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/function.py:1007: calling Graph.create_op (from tensorflow.python.framework.ops) with compute_shapes is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Shapes are always computed; don't use the compute_shapes as it has no effect.\n","INFO:tensorflow:Transforming feature 'input_reward' with symbol_modality_3_64.bottom\n","INFO:tensorflow:Transforming feature 'inputs' with video_modality.bottom\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_video.py:495: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","tf.py_func is deprecated in TF V2. Instead, use\n","    tf.py_function, which takes a python function which manipulates tf eager\n","    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to\n","    an ndarray (just call tensor.numpy()) but having access to eager tensors\n","    means `tf.py_function`s can use accelerators such as GPUs as well as\n","    being differentiable using a gradient tape.\n","    \n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:277: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.cast instead.\n","INFO:tensorflow:Transforming feature 'target_action' with symbol_modality_6_64.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_policy' with identity_modality.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_reward' with symbol_modality_3_64.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_value' with identity_modality.targets_bottom\n","INFO:tensorflow:Transforming feature 'targets' with video_modality.targets_bottom\n","INFO:tensorflow:Building model body\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:598: conv2d (from tensorflow.python.layers.convolutional) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.conv2d instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:602: flatten (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.flatten instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:603: dropout (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.dropout instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:604: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.dense instead.\n","INFO:tensorflow:Transforming body output with identity_modality.top\n","INFO:tensorflow:Transforming body output with identity_modality.top\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:2887: multinomial (from tensorflow.python.ops.random_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.random.categorical instead.\n","2019-03-22 16:06:00.352605: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 16:06:00.352688: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 16:06:00.352724: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 16:06:00.352744: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 16:06:00.353037: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","2019-03-22 16:06:00.588787: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:06:00.647797: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","INFO:tensorflow:Restoring checkpoint gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/policy/model.ckpt-171992\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/saver.py:1266: checkpoint_exists (from tensorflow.python.training.checkpoint_management) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use standard file APIs to check for files with this prefix.\n","2019-03-22 16:06:00.711910: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","INFO:tensorflow:Restoring parameters from gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/policy/model.ckpt-171992\n","2019-03-22 16:06:00.793701: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:06:00.953239: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:06:01.086594: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:06:01.259521: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:06:01.322896: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:06:03.034751: I tensorflow/stream_executor/dso_loader.cc:152] successfully opened CUDA library libcublas.so.10.0 locally\n","INFO:tensorflow:Step 5, mean_score: 0.000000\n","INFO:tensorflow:Step 10, mean_score: 0.000000\n","INFO:tensorflow:Step 15, mean_score: 0.000000\n","INFO:tensorflow:Step 20, mean_score: 0.000000\n","INFO:tensorflow:Step 25, mean_score: 0.000000\n","INFO:tensorflow:Step 30, mean_score: 0.000000\n","INFO:tensorflow:Step 35, mean_score: 0.000000\n","INFO:tensorflow:Step 40, mean_score: 0.000000\n","INFO:tensorflow:Step 45, mean_score: 0.000000\n","INFO:tensorflow:Step 50, mean_score: 0.000000\n","INFO:tensorflow:Step 55, mean_score: 0.000000\n","INFO:tensorflow:Step 60, mean_score: -0.015625\n","INFO:tensorflow:Step 65, mean_score: -0.078125\n","INFO:tensorflow:Step 70, mean_score: -0.078125\n","INFO:tensorflow:Step 75, mean_score: -0.078125\n","INFO:tensorflow:Step 80, mean_score: -0.078125\n","INFO:tensorflow:Step 85, mean_score: -0.078125\n","INFO:tensorflow:Step 90, mean_score: 0.484375\n","INFO:tensorflow:Step 95, mean_score: 0.843750\n","INFO:tensorflow:Step 100, mean_score: 0.828125\n","INFO:tensorflow:Step 105, mean_score: 0.828125\n","INFO:tensorflow:Step 110, mean_score: 0.828125\n","INFO:tensorflow:Step 115, mean_score: 0.828125\n","INFO:tensorflow:Step 120, mean_score: 0.828125\n","INFO:tensorflow:Step 125, mean_score: 0.828125\n","INFO:tensorflow:Step 130, mean_score: 0.812500\n","INFO:tensorflow:Step 135, mean_score: 0.812500\n","INFO:tensorflow:Step 140, mean_score: 0.812500\n","INFO:tensorflow:Step 145, mean_score: 0.812500\n","INFO:tensorflow:Step 150, mean_score: 0.812500\n","INFO:tensorflow:Step 155, mean_score: 0.812500\n","INFO:tensorflow:Step 160, mean_score: 0.812500\n","INFO:tensorflow:Step 165, mean_score: 0.812500\n","INFO:tensorflow:Step 170, mean_score: 0.828125\n","INFO:tensorflow:Step 175, mean_score: 0.843750\n","INFO:tensorflow:Step 180, mean_score: 0.843750\n","INFO:tensorflow:Step 185, mean_score: 0.843750\n","INFO:tensorflow:Step 190, mean_score: 1.140625\n","INFO:tensorflow:Step 195, mean_score: 1.765625\n","INFO:tensorflow:Step 200, mean_score: 1.765625\n","INFO:tensorflow:Step 205, mean_score: 1.765625\n","INFO:tensorflow:Step 210, mean_score: 1.781250\n","INFO:tensorflow:Step 215, mean_score: 1.781250\n","INFO:tensorflow:Step 220, mean_score: 1.765625\n","INFO:tensorflow:Step 225, mean_score: 1.765625\n","INFO:tensorflow:Step 230, mean_score: 1.765625\n","INFO:tensorflow:Step 235, mean_score: 1.765625\n","INFO:tensorflow:Step 240, mean_score: 1.765625\n","INFO:tensorflow:Step 245, mean_score: 1.765625\n","INFO:tensorflow:Step 250, mean_score: 1.765625\n","INFO:tensorflow:Step 255, mean_score: 1.750000\n","INFO:tensorflow:Step 260, mean_score: 1.750000\n","INFO:tensorflow:Step 265, mean_score: 1.750000\n","INFO:tensorflow:Step 270, mean_score: 2.312500\n","INFO:tensorflow:Step 275, mean_score: 2.687500\n","INFO:tensorflow:Step 280, mean_score: 2.703125\n","INFO:tensorflow:Step 285, mean_score: 2.703125\n","INFO:tensorflow:Step 290, mean_score: 2.703125\n","INFO:tensorflow:Step 295, mean_score: 2.703125\n","INFO:tensorflow:Step 300, mean_score: 2.703125\n","INFO:tensorflow:Step 305, mean_score: 2.703125\n","INFO:tensorflow:Step 310, mean_score: 2.718750\n","INFO:tensorflow:Step 315, mean_score: 2.718750\n","INFO:tensorflow:Step 320, mean_score: 2.718750\n","INFO:tensorflow:Step 325, mean_score: 2.718750\n","INFO:tensorflow:Step 330, mean_score: 2.734375\n","INFO:tensorflow:Step 335, mean_score: 2.734375\n","INFO:tensorflow:Step 340, mean_score: 2.734375\n","INFO:tensorflow:Step 345, mean_score: 2.734375\n","INFO:tensorflow:Step 350, mean_score: 2.750000\n","INFO:tensorflow:Step 355, mean_score: 2.765625\n","INFO:tensorflow:Step 360, mean_score: 2.765625\n","INFO:tensorflow:Step 365, mean_score: 2.765625\n","INFO:tensorflow:Step 370, mean_score: 3.062500\n","INFO:tensorflow:Step 375, mean_score: 3.687500\n","INFO:tensorflow:Step 380, mean_score: 3.687500\n","INFO:tensorflow:Step 385, mean_score: 3.687500\n","INFO:tensorflow:Step 390, mean_score: 3.703125\n","INFO:tensorflow:Step 395, mean_score: 3.703125\n","INFO:tensorflow:Step 400, mean_score: 3.703125\n","INFO:tensorflow:Step 405, mean_score: 3.703125\n","INFO:tensorflow:Step 410, mean_score: 3.687500\n","INFO:tensorflow:Step 415, mean_score: 3.687500\n","INFO:tensorflow:Step 420, mean_score: 3.687500\n","INFO:tensorflow:Step 425, mean_score: 3.687500\n","INFO:tensorflow:Step 430, mean_score: 3.703125\n","INFO:tensorflow:Step 435, mean_score: 3.703125\n","INFO:tensorflow:Step 440, mean_score: 3.703125\n","INFO:tensorflow:Step 445, mean_score: 3.703125\n","INFO:tensorflow:Step 450, mean_score: 4.265625\n","INFO:tensorflow:Step 455, mean_score: 4.640625\n","INFO:tensorflow:Step 460, mean_score: 4.656250\n","INFO:tensorflow:Step 465, mean_score: 4.656250\n","INFO:tensorflow:Step 470, mean_score: 4.656250\n","INFO:tensorflow:Step 475, mean_score: 4.656250\n","INFO:tensorflow:Step 480, mean_score: 4.656250\n","INFO:tensorflow:Step 485, mean_score: 4.656250\n","INFO:tensorflow:Step 490, mean_score: 4.671875\n","INFO:tensorflow:Step 495, mean_score: 4.671875\n","INFO:tensorflow:Step 500, mean_score: 4.671875\n","INFO:tensorflow:Step 505, mean_score: 4.671875\n","INFO:tensorflow:Step 510, mean_score: 4.687500\n","INFO:tensorflow:Step 515, mean_score: 4.687500\n","INFO:tensorflow:Step 520, mean_score: 4.703125\n","INFO:tensorflow:Step 525, mean_score: 4.703125\n","INFO:tensorflow:Step 530, mean_score: 4.718750\n","INFO:tensorflow:Step 535, mean_score: 4.734375\n","INFO:tensorflow:Step 540, mean_score: 4.734375\n","INFO:tensorflow:Step 545, mean_score: 4.734375\n","INFO:tensorflow:Step 550, mean_score: 5.031250\n","INFO:tensorflow:Step 555, mean_score: 5.656250\n","INFO:tensorflow:Step 560, mean_score: 5.656250\n","INFO:tensorflow:Step 565, mean_score: 5.656250\n","INFO:tensorflow:Step 570, mean_score: 5.671875\n","INFO:tensorflow:Step 575, mean_score: 5.671875\n","INFO:tensorflow:Step 580, mean_score: 5.671875\n","INFO:tensorflow:Step 585, mean_score: 5.671875\n","INFO:tensorflow:Step 590, mean_score: 5.671875\n","INFO:tensorflow:Step 595, mean_score: 5.671875\n","INFO:tensorflow:Step 600, mean_score: 5.671875\n","INFO:tensorflow:Step 605, mean_score: 5.671875\n","INFO:tensorflow:Step 610, mean_score: 5.687500\n","INFO:tensorflow:Step 615, mean_score: 5.687500\n","INFO:tensorflow:Step 620, mean_score: 5.703125\n","INFO:tensorflow:Step 625, mean_score: 5.703125\n","INFO:tensorflow:Step 630, mean_score: 6.265625\n","INFO:tensorflow:Step 635, mean_score: 6.640625\n","INFO:tensorflow:Step 640, mean_score: 6.656250\n","INFO:tensorflow:Step 645, mean_score: 6.656250\n","INFO:tensorflow:Step 650, mean_score: 6.656250\n","INFO:tensorflow:Step 655, mean_score: 6.656250\n","INFO:tensorflow:Step 660, mean_score: 6.656250\n","INFO:tensorflow:Step 665, mean_score: 6.656250\n","INFO:tensorflow:Step 670, mean_score: 6.671875\n","INFO:tensorflow:Step 675, mean_score: 6.671875\n","INFO:tensorflow:Step 680, mean_score: 6.671875\n","INFO:tensorflow:Step 685, mean_score: 6.671875\n","INFO:tensorflow:Step 690, mean_score: 6.687500\n","INFO:tensorflow:Step 695, mean_score: 6.687500\n","INFO:tensorflow:Step 700, mean_score: 6.703125\n","INFO:tensorflow:Step 705, mean_score: 6.703125\n","INFO:tensorflow:Step 710, mean_score: 6.718750\n","INFO:tensorflow:Step 715, mean_score: 6.734375\n","INFO:tensorflow:Step 720, mean_score: 6.734375\n","INFO:tensorflow:Step 725, mean_score: 6.734375\n","INFO:tensorflow:Step 730, mean_score: 7.031250\n","INFO:tensorflow:Step 735, mean_score: 7.656250\n","INFO:tensorflow:Step 740, mean_score: 7.656250\n","INFO:tensorflow:Step 745, mean_score: 7.656250\n","INFO:tensorflow:Step 750, mean_score: 7.671875\n","INFO:tensorflow:Step 755, mean_score: 7.671875\n","INFO:tensorflow:Step 760, mean_score: 7.671875\n","INFO:tensorflow:Step 765, mean_score: 7.671875\n","INFO:tensorflow:Step 770, mean_score: 7.671875\n","INFO:tensorflow:Step 775, mean_score: 7.671875\n","INFO:tensorflow:Step 780, mean_score: 7.671875\n","INFO:tensorflow:Step 785, mean_score: 7.671875\n","INFO:tensorflow:Step 790, mean_score: 7.687500\n","INFO:tensorflow:Step 795, mean_score: 7.687500\n","INFO:tensorflow:Step 800, mean_score: 7.703125\n","INFO:tensorflow:Step 805, mean_score: 7.703125\n","INFO:tensorflow:Step 810, mean_score: 8.265625\n","INFO:tensorflow:Step 815, mean_score: 8.640625\n","INFO:tensorflow:Step 820, mean_score: 8.656250\n","INFO:tensorflow:Step 825, mean_score: 8.656250\n","INFO:tensorflow:Step 830, mean_score: 8.656250\n","INFO:tensorflow:Step 835, mean_score: 8.656250\n","INFO:tensorflow:Step 840, mean_score: 8.656250\n","INFO:tensorflow:Step 845, mean_score: 8.656250\n","INFO:tensorflow:Step 850, mean_score: 8.671875\n","INFO:tensorflow:Step 855, mean_score: 8.671875\n","INFO:tensorflow:Step 860, mean_score: 8.671875\n","INFO:tensorflow:Step 865, mean_score: 8.671875\n","INFO:tensorflow:Step 870, mean_score: 8.687500\n","INFO:tensorflow:Step 875, mean_score: 8.687500\n","INFO:tensorflow:Step 880, mean_score: 8.703125\n","INFO:tensorflow:Step 885, mean_score: 8.703125\n","INFO:tensorflow:Step 890, mean_score: 8.718750\n","INFO:tensorflow:Step 895, mean_score: 8.734375\n","INFO:tensorflow:Step 900, mean_score: 8.734375\n","INFO:tensorflow:Step 905, mean_score: 8.734375\n","INFO:tensorflow:Step 910, mean_score: 9.031250\n","INFO:tensorflow:Step 915, mean_score: 9.656250\n","INFO:tensorflow:Step 920, mean_score: 9.656250\n","INFO:tensorflow:Step 925, mean_score: 9.656250\n","INFO:tensorflow:Step 930, mean_score: 9.671875\n","INFO:tensorflow:Step 935, mean_score: 9.671875\n","INFO:tensorflow:Step 940, mean_score: 9.671875\n","INFO:tensorflow:Step 945, mean_score: 9.671875\n","INFO:tensorflow:Step 950, mean_score: 9.671875\n","INFO:tensorflow:Step 955, mean_score: 9.671875\n","INFO:tensorflow:Step 960, mean_score: 9.671875\n","INFO:tensorflow:Step 965, mean_score: 9.671875\n","INFO:tensorflow:Step 970, mean_score: 9.687500\n","INFO:tensorflow:Step 975, mean_score: 9.687500\n","INFO:tensorflow:Step 980, mean_score: 9.703125\n","INFO:tensorflow:Step 985, mean_score: 9.703125\n","INFO:tensorflow:Step 990, mean_score: 10.265625\n","INFO:tensorflow:Step 995, mean_score: 10.640625\n","INFO:tensorflow:Step 1000, mean_score: 10.656250\n","INFO:tensorflow:Step 1005, mean_score: 10.656250\n","INFO:tensorflow:Step 1010, mean_score: 10.656250\n","INFO:tensorflow:Step 1015, mean_score: 10.656250\n","INFO:tensorflow:Step 1020, mean_score: 10.656250\n","INFO:tensorflow:Step 1025, mean_score: 10.656250\n","INFO:tensorflow:Step 1030, mean_score: 10.671875\n","INFO:tensorflow:Step 1035, mean_score: 10.671875\n","INFO:tensorflow:Step 1040, mean_score: 10.671875\n","INFO:tensorflow:Step 1045, mean_score: 10.671875\n","INFO:tensorflow:Step 1050, mean_score: 10.687500\n","INFO:tensorflow:Step 1055, mean_score: 10.687500\n","INFO:tensorflow:Step 1060, mean_score: 10.703125\n","INFO:tensorflow:Step 1065, mean_score: 10.703125\n","INFO:tensorflow:Step 1070, mean_score: 10.718750\n","INFO:tensorflow:Step 1075, mean_score: 10.734375\n","INFO:tensorflow:Step 1080, mean_score: 10.734375\n","INFO:tensorflow:Step 1085, mean_score: 10.734375\n","INFO:tensorflow:Step 1090, mean_score: 11.031250\n","INFO:tensorflow:Step 1095, mean_score: 11.656250\n","INFO:tensorflow:Step 1100, mean_score: 11.656250\n","INFO:tensorflow:Step 1105, mean_score: 11.656250\n","INFO:tensorflow:Step 1110, mean_score: 11.671875\n","INFO:tensorflow:Step 1115, mean_score: 11.671875\n","INFO:tensorflow:Step 1120, mean_score: 11.671875\n","INFO:tensorflow:Step 1125, mean_score: 11.671875\n","INFO:tensorflow:Step 1130, mean_score: 11.671875\n","INFO:tensorflow:Step 1135, mean_score: 11.671875\n","INFO:tensorflow:Step 1140, mean_score: 11.671875\n","INFO:tensorflow:Step 1145, mean_score: 11.671875\n","INFO:tensorflow:Step 1150, mean_score: 11.687500\n","INFO:tensorflow:Step 1155, mean_score: 11.687500\n","INFO:tensorflow:Step 1160, mean_score: 11.703125\n","INFO:tensorflow:Step 1165, mean_score: 11.703125\n","INFO:tensorflow:Step 1170, mean_score: 12.265625\n","INFO:tensorflow:Step 1175, mean_score: 12.640625\n","INFO:tensorflow:Step 1180, mean_score: 12.656250\n","INFO:tensorflow:Step 1185, mean_score: 12.656250\n","INFO:tensorflow:Step 1190, mean_score: 12.656250\n","INFO:tensorflow:Step 1195, mean_score: 12.656250\n","INFO:tensorflow:Step 1200, mean_score: 12.656250\n","INFO:tensorflow:Step 1205, mean_score: 12.656250\n","INFO:tensorflow:Step 1210, mean_score: 12.671875\n","INFO:tensorflow:Step 1215, mean_score: 12.671875\n","INFO:tensorflow:Step 1220, mean_score: 12.671875\n","INFO:tensorflow:Step 1225, mean_score: 12.671875\n","INFO:tensorflow:Step 1230, mean_score: 12.687500\n","INFO:tensorflow:Step 1235, mean_score: 12.687500\n","INFO:tensorflow:Step 1240, mean_score: 12.703125\n","INFO:tensorflow:Step 1245, mean_score: 12.703125\n","INFO:tensorflow:Step 1250, mean_score: 12.718750\n","INFO:tensorflow:Step 1255, mean_score: 12.734375\n","INFO:tensorflow:Step 1260, mean_score: 12.734375\n","INFO:tensorflow:Step 1265, mean_score: 12.734375\n","INFO:tensorflow:Step 1270, mean_score: 13.031250\n","INFO:tensorflow:Step 1275, mean_score: 13.656250\n","INFO:tensorflow:Step 1280, mean_score: 13.656250\n","INFO:tensorflow:Step 1285, mean_score: 13.656250\n","INFO:tensorflow:Step 1290, mean_score: 13.671875\n","INFO:tensorflow:Step 1295, mean_score: 13.671875\n","INFO:tensorflow:Step 1300, mean_score: 13.671875\n","INFO:tensorflow:Step 1305, mean_score: 13.671875\n","INFO:tensorflow:Step 1310, mean_score: 13.671875\n","INFO:tensorflow:Step 1315, mean_score: 13.671875\n","INFO:tensorflow:Step 1320, mean_score: 13.671875\n","INFO:tensorflow:Step 1325, mean_score: 13.671875\n","INFO:tensorflow:Step 1330, mean_score: 13.687500\n","INFO:tensorflow:Step 1335, mean_score: 13.687500\n","INFO:tensorflow:Step 1340, mean_score: 13.703125\n","INFO:tensorflow:Step 1345, mean_score: 13.703125\n","INFO:tensorflow:Step 1350, mean_score: 14.265625\n","INFO:tensorflow:Step 1355, mean_score: 14.640625\n","INFO:tensorflow:Step 1360, mean_score: 14.656250\n","INFO:tensorflow:Step 1365, mean_score: 14.656250\n","INFO:tensorflow:Step 1370, mean_score: 14.656250\n","INFO:tensorflow:Step 1375, mean_score: 14.656250\n","INFO:tensorflow:Step 1380, mean_score: 14.656250\n","INFO:tensorflow:Step 1385, mean_score: 14.656250\n","INFO:tensorflow:Step 1390, mean_score: 14.671875\n","INFO:tensorflow:Step 1395, mean_score: 14.671875\n","INFO:tensorflow:Step 1400, mean_score: 14.671875\n","INFO:tensorflow:Step 1405, mean_score: 14.671875\n","INFO:tensorflow:Step 1410, mean_score: 14.687500\n","INFO:tensorflow:Step 1415, mean_score: 14.687500\n","INFO:tensorflow:Step 1420, mean_score: 14.703125\n","INFO:tensorflow:Step 1425, mean_score: 14.703125\n","INFO:tensorflow:Step 1430, mean_score: 14.718750\n","INFO:tensorflow:Step 1435, mean_score: 14.734375\n","INFO:tensorflow:Step 1440, mean_score: 14.734375\n","INFO:tensorflow:Step 1445, mean_score: 14.734375\n","INFO:tensorflow:Step 1450, mean_score: 15.031250\n","INFO:tensorflow:Step 1455, mean_score: 15.656250\n","INFO:tensorflow:Step 1460, mean_score: 15.656250\n","INFO:tensorflow:Step 1465, mean_score: 15.656250\n","INFO:tensorflow:Step 1470, mean_score: 15.671875\n","INFO:tensorflow:Step 1475, mean_score: 15.671875\n","INFO:tensorflow:Step 1480, mean_score: 15.671875\n","INFO:tensorflow:Step 1485, mean_score: 15.671875\n","INFO:tensorflow:Step 1490, mean_score: 15.671875\n","INFO:tensorflow:Step 1495, mean_score: 15.671875\n","INFO:tensorflow:Step 1500, mean_score: 15.671875\n","INFO:tensorflow:Step 1505, mean_score: 15.671875\n","INFO:tensorflow:Step 1510, mean_score: 15.687500\n","INFO:tensorflow:Step 1515, mean_score: 15.687500\n","INFO:tensorflow:Step 1520, mean_score: 15.703125\n","INFO:tensorflow:Step 1525, mean_score: 15.703125\n","INFO:tensorflow:Step 1530, mean_score: 16.265625\n","INFO:tensorflow:Step 1535, mean_score: 16.640625\n","INFO:tensorflow:Step 1540, mean_score: 16.656250\n","INFO:tensorflow:Step 1545, mean_score: 16.656250\n","INFO:tensorflow:Step 1550, mean_score: 16.656250\n","INFO:tensorflow:Step 1555, mean_score: 16.656250\n","INFO:tensorflow:Step 1560, mean_score: 16.656250\n","INFO:tensorflow:Step 1565, mean_score: 16.656250\n","INFO:tensorflow:Step 1570, mean_score: 16.671875\n","INFO:tensorflow:Step 1575, mean_score: 16.671875\n","INFO:tensorflow:Step 1580, mean_score: 16.671875\n","INFO:tensorflow:Step 1585, mean_score: 16.671875\n","INFO:tensorflow:Step 1590, mean_score: 16.687500\n","INFO:tensorflow:Step 1595, mean_score: 16.687500\n","INFO:tensorflow:Step 1600, mean_score: 16.703125\n","INFO:tensorflow:Step 1605, mean_score: 16.703125\n","INFO:tensorflow:Step 1610, mean_score: 16.718750\n","INFO:tensorflow:Step 1615, mean_score: 16.734375\n","INFO:tensorflow:Step 1620, mean_score: 16.734375\n","INFO:tensorflow:Step 1625, mean_score: 16.734375\n","INFO:tensorflow:Step 1630, mean_score: 17.031250\n","INFO:tensorflow:Step 1635, mean_score: 17.656250\n","INFO:tensorflow:Step 1640, mean_score: 17.656250\n","INFO:tensorflow:Step 1645, mean_score: 17.656250\n","INFO:tensorflow:Step 1650, mean_score: 17.671875\n","INFO:tensorflow:Step 1655, mean_score: 17.671875\n","INFO:tensorflow:Step 1660, mean_score: 17.671875\n","INFO:tensorflow:Step 1665, mean_score: 17.671875\n","INFO:tensorflow:Step 1670, mean_score: 17.671875\n","INFO:tensorflow:Step 1675, mean_score: 17.671875\n","INFO:tensorflow:Step 1680, mean_score: 17.671875\n","INFO:tensorflow:Step 1685, mean_score: 17.671875\n","INFO:tensorflow:Step 1690, mean_score: 17.687500\n","INFO:tensorflow:Step 1695, mean_score: 17.687500\n","INFO:tensorflow:Step 1700, mean_score: 17.703125\n","INFO:tensorflow:Step 1705, mean_score: 17.703125\n","INFO:tensorflow:Step 1710, mean_score: 18.265625\n","INFO:tensorflow:Step 1715, mean_score: 18.640625\n","INFO:tensorflow:Step 1720, mean_score: 18.656250\n","INFO:tensorflow:Step 1725, mean_score: 18.656250\n","INFO:tensorflow:Step 1730, mean_score: 18.656250\n","INFO:tensorflow:Step 1735, mean_score: 18.656250\n","INFO:tensorflow:Step 1740, mean_score: 18.656250\n","INFO:tensorflow:Step 1745, mean_score: 18.656250\n","INFO:tensorflow:Step 1750, mean_score: 18.671875\n","INFO:tensorflow:Step 1755, mean_score: 18.671875\n","INFO:tensorflow:Step 1760, mean_score: 18.671875\n","INFO:tensorflow:Step 1765, mean_score: 18.671875\n","INFO:tensorflow:Step 1770, mean_score: 18.687500\n","INFO:tensorflow:Step 1775, mean_score: 18.687500\n","INFO:tensorflow:Step 1780, mean_score: 18.703125\n","INFO:tensorflow:Step 1785, mean_score: 18.703125\n","INFO:tensorflow:Step 1790, mean_score: 18.718750\n","INFO:tensorflow:Step 1795, mean_score: 18.734375\n","INFO:tensorflow:Step 1800, mean_score: 18.734375\n","INFO:tensorflow:Step 1805, mean_score: 18.734375\n","INFO:tensorflow:Step 1810, mean_score: 19.031250\n","INFO:tensorflow:Step 1815, mean_score: 19.656250\n","INFO:tensorflow:Step 1820, mean_score: 19.656250\n","INFO:tensorflow:Step 1825, mean_score: 19.656250\n","INFO:tensorflow:Step 1830, mean_score: 19.671875\n","INFO:tensorflow:Step 1835, mean_score: 19.671875\n","INFO:tensorflow:Step 1840, mean_score: 19.671875\n","INFO:tensorflow:Step 1845, mean_score: 19.671875\n","INFO:tensorflow:Step 1850, mean_score: 19.671875\n","INFO:tensorflow:Step 1855, mean_score: 19.671875\n","INFO:tensorflow:Step 1860, mean_score: 19.671875\n","INFO:tensorflow:Step 1865, mean_score: 19.671875\n","INFO:tensorflow:Step 1870, mean_score: 19.687500\n","INFO:tensorflow:Step 1875, mean_score: 19.687500\n","INFO:tensorflow:Step 1880, mean_score: 19.703125\n","INFO:tensorflow:Step 1885, mean_score: 19.703125\n","INFO:tensorflow:Step 1890, mean_score: 19.703125\n","INFO:tensorflow:Step 1895, mean_score: 19.718750\n","INFO:tensorflow:Step 1900, mean_score: 19.734375\n","INFO:tensorflow:Step 1905, mean_score: 19.734375\n","INFO:tensorflow:Step 1910, mean_score: 19.734375\n","INFO:tensorflow:Step 1915, mean_score: 19.734375\n","INFO:tensorflow:Step 1920, mean_score: 19.734375\n","INFO:tensorflow:Step 1925, mean_score: 19.734375\n","INFO:tensorflow:Step 1930, mean_score: 19.750000\n","INFO:tensorflow:Step 1935, mean_score: 19.750000\n","INFO:tensorflow:Step 1940, mean_score: 19.750000\n","INFO:tensorflow:Step 1945, mean_score: 19.750000\n","INFO:tensorflow:Step 1950, mean_score: 19.765625\n","INFO:tensorflow:Step 1955, mean_score: 19.765625\n","INFO:tensorflow:Step 1960, mean_score: 19.781250\n","INFO:tensorflow:Step 1965, mean_score: 19.781250\n","INFO:tensorflow:Step 1970, mean_score: 19.781250\n","INFO:tensorflow:Step 1975, mean_score: 19.781250\n","INFO:tensorflow:Step 1980, mean_score: 19.781250\n","INFO:tensorflow:Step 1985, mean_score: 19.781250\n","INFO:tensorflow:Step 1990, mean_score: 19.781250\n","INFO:tensorflow:Step 1995, mean_score: 19.781250\n","INFO:tensorflow:Step 2000, mean_score: 19.781250\n","INFO:tensorflow:Step 2005, mean_score: 19.781250\n","INFO:tensorflow:Step 2010, mean_score: 19.781250\n","INFO:tensorflow:Step 2015, mean_score: 19.781250\n","INFO:tensorflow:Step 2020, mean_score: 19.781250\n","INFO:tensorflow:Step 2025, mean_score: 19.781250\n","INFO:tensorflow:Step 2030, mean_score: 19.781250\n","INFO:tensorflow:Step 2035, mean_score: 19.781250\n","INFO:tensorflow:Step 2040, mean_score: 19.781250\n","INFO:tensorflow:Step 2045, mean_score: 19.781250\n","INFO:tensorflow:Step 2050, mean_score: 19.796875\n","INFO:tensorflow:Step 2055, mean_score: 19.796875\n","INFO:tensorflow:Step 2060, mean_score: 19.812500\n","INFO:tensorflow:Step 2065, mean_score: 19.812500\n","INFO:tensorflow:Step 2070, mean_score: 19.812500\n","INFO:tensorflow:Step 2075, mean_score: 19.812500\n","INFO:tensorflow:Step 2080, mean_score: 19.812500\n","INFO:tensorflow:Step 2085, mean_score: 19.812500\n","INFO:tensorflow:Step 2090, mean_score: 19.812500\n","INFO:tensorflow:Step 2095, mean_score: 19.812500\n","INFO:tensorflow:Step 2100, mean_score: 19.812500\n","INFO:tensorflow:Step 2105, mean_score: 19.812500\n","INFO:tensorflow:Step 2110, mean_score: 19.812500\n","INFO:tensorflow:Step 2115, mean_score: 19.812500\n","INFO:tensorflow:Step 2120, mean_score: 19.812500\n","INFO:tensorflow:Step 2125, mean_score: 19.812500\n","INFO:tensorflow:Step 2130, mean_score: 19.812500\n","INFO:tensorflow:Step 2135, mean_score: 19.812500\n","INFO:tensorflow:Step 2140, mean_score: 19.828125\n","INFO:tensorflow:Step 2145, mean_score: 19.828125\n","INFO:tensorflow:Step 2150, mean_score: 19.828125\n","INFO:tensorflow:Step 2155, mean_score: 19.828125\n","INFO:tensorflow:Step 2160, mean_score: 19.828125\n","INFO:tensorflow:Step 2165, mean_score: 19.828125\n","INFO:tensorflow:Step 2170, mean_score: 19.828125\n","INFO:tensorflow:Step 2175, mean_score: 19.828125\n","INFO:tensorflow:Step 2180, mean_score: 19.828125\n","INFO:tensorflow:Step 2185, mean_score: 19.828125\n","INFO:tensorflow:Step 2190, mean_score: 19.828125\n","INFO:tensorflow:Step 2195, mean_score: 19.828125\n","INFO:tensorflow:Step 2200, mean_score: 19.828125\n","INFO:tensorflow:Step 2205, mean_score: 19.828125\n","INFO:tensorflow:Step 2210, mean_score: 19.828125\n","INFO:tensorflow:Step 2215, mean_score: 19.828125\n","INFO:tensorflow:Step 2220, mean_score: 19.828125\n","INFO:tensorflow:Step 2225, mean_score: 19.828125\n","INFO:tensorflow:Step 2230, mean_score: 19.828125\n","INFO:tensorflow:Step 2235, mean_score: 19.828125\n","INFO:tensorflow:Step 2240, mean_score: 19.843750\n","INFO:tensorflow:Step 2245, mean_score: 19.843750\n","INFO:tensorflow:Step 2250, mean_score: 19.843750\n","INFO:tensorflow:Step 2255, mean_score: 19.843750\n","INFO:tensorflow:Step 2260, mean_score: 19.843750\n","INFO:tensorflow:Step 2265, mean_score: 19.843750\n","INFO:tensorflow:Step 2270, mean_score: 19.843750\n","INFO:tensorflow:Step 2275, mean_score: 19.843750\n","INFO:tensorflow:Step 2280, mean_score: 19.843750\n","INFO:tensorflow:Step 2285, mean_score: 19.843750\n","INFO:tensorflow:Step 2290, mean_score: 19.843750\n","INFO:tensorflow:Step 2295, mean_score: 19.843750\n","INFO:tensorflow:Step 2300, mean_score: 19.843750\n","INFO:tensorflow:Step 2305, mean_score: 19.843750\n","INFO:tensorflow:Step 2310, mean_score: 19.843750\n","INFO:tensorflow:Step 2315, mean_score: 19.843750\n","INFO:tensorflow:Evaluating metric mean_reward/eval/sampling_temp_0.5_max_noops_0_unclipped\n","2019-03-22 16:12:57.935045: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 16:12:57.935160: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 16:12:57.935189: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 16:12:57.935209: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 16:12:57.935553: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","INFO:tensorflow:Using DummyPolicyProblem for the policy.\n","INFO:tensorflow:Setting T2TModel mode to 'train'\n","INFO:tensorflow:Using variable initializer: orthogonal\n","INFO:tensorflow:Transforming feature 'input_action' with symbol_modality_6_64.bottom\n","INFO:tensorflow:Transforming feature 'input_reward' with symbol_modality_3_64.bottom\n","INFO:tensorflow:Transforming feature 'inputs' with video_modality.bottom\n","INFO:tensorflow:Transforming feature 'target_action' with symbol_modality_6_64.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_policy' with identity_modality.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_reward' with symbol_modality_3_64.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_value' with identity_modality.targets_bottom\n","INFO:tensorflow:Transforming feature 'targets' with video_modality.targets_bottom\n","INFO:tensorflow:Building model body\n","INFO:tensorflow:Transforming body output with identity_modality.top\n","INFO:tensorflow:Transforming body output with identity_modality.top\n","2019-03-22 16:13:12.260846: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 16:13:12.260981: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 16:13:12.261059: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 16:13:12.261099: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 16:13:12.261613: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","2019-03-22 16:13:12.493082: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","INFO:tensorflow:Restoring checkpoint gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/policy/model.ckpt-171992\n","2019-03-22 16:13:12.556955: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","INFO:tensorflow:Restoring parameters from gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/policy/model.ckpt-171992\n","2019-03-22 16:13:12.651009: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:13:12.715180: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","2019-03-22 16:13:12.816774: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n","INFO:tensorflow:Step 5, mean_score: 0.000000\n","INFO:tensorflow:Step 10, mean_score: 0.000000\n","INFO:tensorflow:Step 15, mean_score: 0.000000\n","INFO:tensorflow:Step 20, mean_score: 0.000000\n","INFO:tensorflow:Step 25, mean_score: 0.000000\n","INFO:tensorflow:Step 30, mean_score: 0.000000\n","INFO:tensorflow:Step 35, mean_score: 0.000000\n","INFO:tensorflow:Step 40, mean_score: 0.000000\n","INFO:tensorflow:Step 45, mean_score: 0.000000\n","INFO:tensorflow:Step 50, mean_score: 0.000000\n","INFO:tensorflow:Step 55, mean_score: 0.000000\n","INFO:tensorflow:Step 60, mean_score: 0.000000\n","INFO:tensorflow:Step 65, mean_score: -0.031250\n","INFO:tensorflow:Step 70, mean_score: -0.031250\n","INFO:tensorflow:Step 75, mean_score: -0.031250\n","INFO:tensorflow:Step 80, mean_score: -0.031250\n","INFO:tensorflow:Step 85, mean_score: -0.031250\n","INFO:tensorflow:Step 90, mean_score: -0.031250\n","INFO:tensorflow:Step 95, mean_score: 0.937500\n","INFO:tensorflow:Step 100, mean_score: 0.921875\n","INFO:tensorflow:Step 105, mean_score: 0.921875\n","INFO:tensorflow:Step 110, mean_score: 0.921875\n","INFO:tensorflow:Step 115, mean_score: 0.921875\n","INFO:tensorflow:Step 120, mean_score: 0.921875\n","INFO:tensorflow:Step 125, mean_score: 0.921875\n","INFO:tensorflow:Step 130, mean_score: 0.921875\n","INFO:tensorflow:Step 135, mean_score: 0.921875\n","INFO:tensorflow:Step 140, mean_score: 0.921875\n","INFO:tensorflow:Step 145, mean_score: 0.921875\n","INFO:tensorflow:Step 150, mean_score: 0.921875\n","INFO:tensorflow:Step 155, mean_score: 0.921875\n","INFO:tensorflow:Step 160, mean_score: 0.921875\n","INFO:tensorflow:Step 165, mean_score: 0.906250\n","INFO:tensorflow:Step 170, mean_score: 0.906250\n","INFO:tensorflow:Step 175, mean_score: 0.921875\n","INFO:tensorflow:Step 180, mean_score: 0.921875\n","INFO:tensorflow:Step 185, mean_score: 0.921875\n","INFO:tensorflow:Step 190, mean_score: 0.921875\n","INFO:tensorflow:Step 195, mean_score: 0.921875\n","INFO:tensorflow:Step 200, mean_score: 1.890625\n","INFO:tensorflow:Step 205, mean_score: 1.890625\n","INFO:tensorflow:Step 210, mean_score: 1.890625\n","INFO:tensorflow:Step 215, mean_score: 1.890625\n","INFO:tensorflow:Step 220, mean_score: 1.890625\n","INFO:tensorflow:Step 225, mean_score: 1.890625\n","INFO:tensorflow:Step 230, mean_score: 1.890625\n","INFO:tensorflow:Step 235, mean_score: 1.890625\n","INFO:tensorflow:Step 240, mean_score: 1.890625\n","INFO:tensorflow:Step 245, mean_score: 1.890625\n","INFO:tensorflow:Step 250, mean_score: 1.890625\n","INFO:tensorflow:Step 255, mean_score: 1.890625\n","INFO:tensorflow:Step 260, mean_score: 1.890625\n","INFO:tensorflow:Step 265, mean_score: 1.890625\n","INFO:tensorflow:Step 270, mean_score: 1.890625\n","INFO:tensorflow:Step 275, mean_score: 2.875000\n","INFO:tensorflow:Step 280, mean_score: 2.890625\n","INFO:tensorflow:Step 285, mean_score: 2.890625\n","INFO:tensorflow:Step 290, mean_score: 2.890625\n","INFO:tensorflow:Step 295, mean_score: 2.890625\n","INFO:tensorflow:Step 300, mean_score: 2.890625\n","INFO:tensorflow:Step 305, mean_score: 2.890625\n","INFO:tensorflow:Step 310, mean_score: 2.890625\n","INFO:tensorflow:Step 315, mean_score: 2.890625\n","INFO:tensorflow:Step 320, mean_score: 2.890625\n","INFO:tensorflow:Step 325, mean_score: 2.890625\n","INFO:tensorflow:Step 330, mean_score: 2.890625\n","INFO:tensorflow:Step 335, mean_score: 2.890625\n","INFO:tensorflow:Step 340, mean_score: 2.890625\n","INFO:tensorflow:Step 345, mean_score: 2.890625\n","INFO:tensorflow:Step 350, mean_score: 2.890625\n","INFO:tensorflow:Step 355, mean_score: 2.906250\n","INFO:tensorflow:Step 360, mean_score: 2.906250\n","INFO:tensorflow:Step 365, mean_score: 2.906250\n","INFO:tensorflow:Step 370, mean_score: 2.906250\n","INFO:tensorflow:Step 375, mean_score: 2.921875\n","INFO:tensorflow:Step 380, mean_score: 3.890625\n","INFO:tensorflow:Step 385, mean_score: 3.890625\n","INFO:tensorflow:Step 390, mean_score: 3.890625\n","INFO:tensorflow:Step 395, mean_score: 3.890625\n","INFO:tensorflow:Step 400, mean_score: 3.890625\n","INFO:tensorflow:Step 405, mean_score: 3.890625\n","INFO:tensorflow:Step 410, mean_score: 3.890625\n","INFO:tensorflow:Step 415, mean_score: 3.890625\n","INFO:tensorflow:Step 420, mean_score: 3.890625\n","INFO:tensorflow:Step 425, mean_score: 3.890625\n","INFO:tensorflow:Step 430, mean_score: 3.890625\n","INFO:tensorflow:Step 435, mean_score: 3.890625\n","INFO:tensorflow:Step 440, mean_score: 3.890625\n","INFO:tensorflow:Step 445, mean_score: 3.890625\n","INFO:tensorflow:Step 450, mean_score: 3.890625\n","INFO:tensorflow:Step 455, mean_score: 4.875000\n","INFO:tensorflow:Step 460, mean_score: 4.890625\n","INFO:tensorflow:Step 465, mean_score: 4.890625\n","INFO:tensorflow:Step 470, mean_score: 4.890625\n","INFO:tensorflow:Step 475, mean_score: 4.890625\n","INFO:tensorflow:Step 480, mean_score: 4.890625\n","INFO:tensorflow:Step 485, mean_score: 4.890625\n","INFO:tensorflow:Step 490, mean_score: 4.890625\n","INFO:tensorflow:Step 495, mean_score: 4.890625\n","INFO:tensorflow:Step 500, mean_score: 4.890625\n","INFO:tensorflow:Step 505, mean_score: 4.890625\n","INFO:tensorflow:Step 510, mean_score: 4.890625\n","INFO:tensorflow:Step 515, mean_score: 4.890625\n","INFO:tensorflow:Step 520, mean_score: 4.890625\n","INFO:tensorflow:Step 525, mean_score: 4.890625\n","INFO:tensorflow:Step 530, mean_score: 4.890625\n","INFO:tensorflow:Step 535, mean_score: 4.906250\n","INFO:tensorflow:Step 540, mean_score: 4.906250\n","INFO:tensorflow:Step 545, mean_score: 4.906250\n","INFO:tensorflow:Step 550, mean_score: 4.906250\n","INFO:tensorflow:Step 555, mean_score: 4.921875\n","INFO:tensorflow:Step 560, mean_score: 5.890625\n","INFO:tensorflow:Step 565, mean_score: 5.890625\n","INFO:tensorflow:Step 570, mean_score: 5.890625\n","INFO:tensorflow:Step 575, mean_score: 5.890625\n","INFO:tensorflow:Step 580, mean_score: 5.890625\n","INFO:tensorflow:Step 585, mean_score: 5.890625\n","INFO:tensorflow:Step 590, mean_score: 5.890625\n","INFO:tensorflow:Step 595, mean_score: 5.890625\n","INFO:tensorflow:Step 600, mean_score: 5.890625\n","INFO:tensorflow:Step 605, mean_score: 5.890625\n","INFO:tensorflow:Step 610, mean_score: 5.890625\n","INFO:tensorflow:Step 615, mean_score: 5.890625\n","INFO:tensorflow:Step 620, mean_score: 5.890625\n","INFO:tensorflow:Step 625, mean_score: 5.890625\n","INFO:tensorflow:Step 630, mean_score: 5.890625\n","INFO:tensorflow:Step 635, mean_score: 6.875000\n","INFO:tensorflow:Step 640, mean_score: 6.890625\n","INFO:tensorflow:Step 645, mean_score: 6.890625\n","INFO:tensorflow:Step 650, mean_score: 6.890625\n","INFO:tensorflow:Step 655, mean_score: 6.890625\n","INFO:tensorflow:Step 660, mean_score: 6.890625\n","INFO:tensorflow:Step 665, mean_score: 6.890625\n","INFO:tensorflow:Step 670, mean_score: 6.890625\n","INFO:tensorflow:Step 675, mean_score: 6.890625\n","INFO:tensorflow:Step 680, mean_score: 6.890625\n","INFO:tensorflow:Step 685, mean_score: 6.890625\n","INFO:tensorflow:Step 690, mean_score: 6.890625\n","INFO:tensorflow:Step 695, mean_score: 6.890625\n","INFO:tensorflow:Step 700, mean_score: 6.890625\n","INFO:tensorflow:Step 705, mean_score: 6.890625\n","INFO:tensorflow:Step 710, mean_score: 6.890625\n","INFO:tensorflow:Step 715, mean_score: 6.906250\n","INFO:tensorflow:Step 720, mean_score: 6.906250\n","INFO:tensorflow:Step 725, mean_score: 6.906250\n","INFO:tensorflow:Step 730, mean_score: 6.906250\n","INFO:tensorflow:Step 735, mean_score: 6.921875\n","INFO:tensorflow:Step 740, mean_score: 7.890625\n","INFO:tensorflow:Step 745, mean_score: 7.890625\n","INFO:tensorflow:Step 750, mean_score: 7.890625\n","INFO:tensorflow:Step 755, mean_score: 7.890625\n","INFO:tensorflow:Step 760, mean_score: 7.890625\n","INFO:tensorflow:Step 765, mean_score: 7.890625\n","INFO:tensorflow:Step 770, mean_score: 7.890625\n","INFO:tensorflow:Step 775, mean_score: 7.890625\n","INFO:tensorflow:Step 780, mean_score: 7.890625\n","INFO:tensorflow:Step 785, mean_score: 7.890625\n","INFO:tensorflow:Step 790, mean_score: 7.890625\n","INFO:tensorflow:Step 795, mean_score: 7.890625\n","INFO:tensorflow:Step 800, mean_score: 7.890625\n","INFO:tensorflow:Step 805, mean_score: 7.890625\n","INFO:tensorflow:Step 810, mean_score: 7.890625\n","INFO:tensorflow:Step 815, mean_score: 8.875000\n","INFO:tensorflow:Step 820, mean_score: 8.890625\n","INFO:tensorflow:Step 825, mean_score: 8.890625\n","INFO:tensorflow:Step 830, mean_score: 8.890625\n","INFO:tensorflow:Step 835, mean_score: 8.890625\n","INFO:tensorflow:Step 840, mean_score: 8.890625\n","INFO:tensorflow:Step 845, mean_score: 8.890625\n","INFO:tensorflow:Step 850, mean_score: 8.890625\n","INFO:tensorflow:Step 855, mean_score: 8.890625\n","INFO:tensorflow:Step 860, mean_score: 8.890625\n","INFO:tensorflow:Step 865, mean_score: 8.890625\n","INFO:tensorflow:Step 870, mean_score: 8.890625\n","INFO:tensorflow:Step 875, mean_score: 8.890625\n","INFO:tensorflow:Step 880, mean_score: 8.890625\n","INFO:tensorflow:Step 885, mean_score: 8.890625\n","INFO:tensorflow:Step 890, mean_score: 8.890625\n","INFO:tensorflow:Step 895, mean_score: 8.906250\n","INFO:tensorflow:Step 900, mean_score: 8.906250\n","INFO:tensorflow:Step 905, mean_score: 8.906250\n","INFO:tensorflow:Step 910, mean_score: 8.906250\n","INFO:tensorflow:Step 915, mean_score: 8.921875\n","INFO:tensorflow:Step 920, mean_score: 9.890625\n","INFO:tensorflow:Step 925, mean_score: 9.890625\n","INFO:tensorflow:Step 930, mean_score: 9.890625\n","INFO:tensorflow:Step 935, mean_score: 9.890625\n","INFO:tensorflow:Step 940, mean_score: 9.890625\n","INFO:tensorflow:Step 945, mean_score: 9.890625\n","INFO:tensorflow:Step 950, mean_score: 9.890625\n","INFO:tensorflow:Step 955, mean_score: 9.890625\n","INFO:tensorflow:Step 960, mean_score: 9.890625\n","INFO:tensorflow:Step 965, mean_score: 9.890625\n","INFO:tensorflow:Step 970, mean_score: 9.890625\n","INFO:tensorflow:Step 975, mean_score: 9.890625\n","INFO:tensorflow:Step 980, mean_score: 9.890625\n","INFO:tensorflow:Step 985, mean_score: 9.890625\n","INFO:tensorflow:Step 990, mean_score: 9.890625\n","INFO:tensorflow:Step 995, mean_score: 10.875000\n","INFO:tensorflow:Step 1000, mean_score: 10.890625\n","INFO:tensorflow:Step 1005, mean_score: 10.890625\n","INFO:tensorflow:Step 1010, mean_score: 10.890625\n","INFO:tensorflow:Step 1015, mean_score: 10.890625\n","INFO:tensorflow:Step 1020, mean_score: 10.890625\n","INFO:tensorflow:Step 1025, mean_score: 10.890625\n","INFO:tensorflow:Step 1030, mean_score: 10.890625\n","INFO:tensorflow:Step 1035, mean_score: 10.890625\n","INFO:tensorflow:Step 1040, mean_score: 10.890625\n","INFO:tensorflow:Step 1045, mean_score: 10.890625\n","INFO:tensorflow:Step 1050, mean_score: 10.890625\n","INFO:tensorflow:Step 1055, mean_score: 10.890625\n","INFO:tensorflow:Step 1060, mean_score: 10.890625\n","INFO:tensorflow:Step 1065, mean_score: 10.890625\n","INFO:tensorflow:Step 1070, mean_score: 10.890625\n","INFO:tensorflow:Step 1075, mean_score: 10.906250\n","INFO:tensorflow:Step 1080, mean_score: 10.906250\n","INFO:tensorflow:Step 1085, mean_score: 10.906250\n","INFO:tensorflow:Step 1090, mean_score: 10.906250\n","INFO:tensorflow:Step 1095, mean_score: 10.921875\n","INFO:tensorflow:Step 1100, mean_score: 11.890625\n","INFO:tensorflow:Step 1105, mean_score: 11.890625\n","INFO:tensorflow:Step 1110, mean_score: 11.890625\n","INFO:tensorflow:Step 1115, mean_score: 11.890625\n","INFO:tensorflow:Step 1120, mean_score: 11.890625\n","INFO:tensorflow:Step 1125, mean_score: 11.890625\n","INFO:tensorflow:Step 1130, mean_score: 11.890625\n","INFO:tensorflow:Step 1135, mean_score: 11.890625\n","INFO:tensorflow:Step 1140, mean_score: 11.890625\n","INFO:tensorflow:Step 1145, mean_score: 11.890625\n","INFO:tensorflow:Step 1150, mean_score: 11.890625\n","INFO:tensorflow:Step 1155, mean_score: 11.890625\n","INFO:tensorflow:Step 1160, mean_score: 11.890625\n","INFO:tensorflow:Step 1165, mean_score: 11.890625\n","INFO:tensorflow:Step 1170, mean_score: 11.890625\n","INFO:tensorflow:Step 1175, mean_score: 12.875000\n","INFO:tensorflow:Step 1180, mean_score: 12.890625\n","INFO:tensorflow:Step 1185, mean_score: 12.890625\n","INFO:tensorflow:Step 1190, mean_score: 12.890625\n","INFO:tensorflow:Step 1195, mean_score: 12.890625\n","INFO:tensorflow:Step 1200, mean_score: 12.890625\n","INFO:tensorflow:Step 1205, mean_score: 12.890625\n","INFO:tensorflow:Step 1210, mean_score: 12.890625\n","INFO:tensorflow:Step 1215, mean_score: 12.890625\n","INFO:tensorflow:Step 1220, mean_score: 12.890625\n","INFO:tensorflow:Step 1225, mean_score: 12.890625\n","INFO:tensorflow:Step 1230, mean_score: 12.890625\n","INFO:tensorflow:Step 1235, mean_score: 12.890625\n","INFO:tensorflow:Step 1240, mean_score: 12.890625\n","INFO:tensorflow:Step 1245, mean_score: 12.890625\n","INFO:tensorflow:Step 1250, mean_score: 12.890625\n","INFO:tensorflow:Step 1255, mean_score: 12.906250\n","INFO:tensorflow:Step 1260, mean_score: 12.906250\n","INFO:tensorflow:Step 1265, mean_score: 12.906250\n","INFO:tensorflow:Step 1270, mean_score: 12.906250\n","INFO:tensorflow:Step 1275, mean_score: 12.921875\n","INFO:tensorflow:Step 1280, mean_score: 13.890625\n","INFO:tensorflow:Step 1285, mean_score: 13.890625\n","INFO:tensorflow:Step 1290, mean_score: 13.890625\n","INFO:tensorflow:Step 1295, mean_score: 13.890625\n","INFO:tensorflow:Step 1300, mean_score: 13.890625\n","INFO:tensorflow:Step 1305, mean_score: 13.890625\n","INFO:tensorflow:Step 1310, mean_score: 13.890625\n","INFO:tensorflow:Step 1315, mean_score: 13.890625\n","INFO:tensorflow:Step 1320, mean_score: 13.890625\n","INFO:tensorflow:Step 1325, mean_score: 13.890625\n","INFO:tensorflow:Step 1330, mean_score: 13.890625\n","INFO:tensorflow:Step 1335, mean_score: 13.890625\n","INFO:tensorflow:Step 1340, mean_score: 13.890625\n","INFO:tensorflow:Step 1345, mean_score: 13.890625\n","INFO:tensorflow:Step 1350, mean_score: 13.890625\n","INFO:tensorflow:Step 1355, mean_score: 14.875000\n","INFO:tensorflow:Step 1360, mean_score: 14.890625\n","INFO:tensorflow:Step 1365, mean_score: 14.890625\n","INFO:tensorflow:Step 1370, mean_score: 14.890625\n","INFO:tensorflow:Step 1375, mean_score: 14.890625\n","INFO:tensorflow:Step 1380, mean_score: 14.890625\n","INFO:tensorflow:Step 1385, mean_score: 14.890625\n","INFO:tensorflow:Step 1390, mean_score: 14.890625\n","INFO:tensorflow:Step 1395, mean_score: 14.890625\n","INFO:tensorflow:Step 1400, mean_score: 14.890625\n","INFO:tensorflow:Step 1405, mean_score: 14.890625\n","INFO:tensorflow:Step 1410, mean_score: 14.890625\n","INFO:tensorflow:Step 1415, mean_score: 14.890625\n","INFO:tensorflow:Step 1420, mean_score: 14.890625\n","INFO:tensorflow:Step 1425, mean_score: 14.890625\n","INFO:tensorflow:Step 1430, mean_score: 14.890625\n","INFO:tensorflow:Step 1435, mean_score: 14.906250\n","INFO:tensorflow:Step 1440, mean_score: 14.906250\n","INFO:tensorflow:Step 1445, mean_score: 14.906250\n","INFO:tensorflow:Step 1450, mean_score: 14.906250\n","INFO:tensorflow:Step 1455, mean_score: 14.921875\n","INFO:tensorflow:Step 1460, mean_score: 15.890625\n","INFO:tensorflow:Step 1465, mean_score: 15.890625\n","INFO:tensorflow:Step 1470, mean_score: 15.890625\n","INFO:tensorflow:Step 1475, mean_score: 15.890625\n","INFO:tensorflow:Step 1480, mean_score: 15.890625\n","INFO:tensorflow:Step 1485, mean_score: 15.890625\n","INFO:tensorflow:Step 1490, mean_score: 15.890625\n","INFO:tensorflow:Step 1495, mean_score: 15.890625\n","INFO:tensorflow:Step 1500, mean_score: 15.890625\n","INFO:tensorflow:Step 1505, mean_score: 15.890625\n","INFO:tensorflow:Step 1510, mean_score: 15.890625\n","INFO:tensorflow:Step 1515, mean_score: 15.890625\n","INFO:tensorflow:Step 1520, mean_score: 15.890625\n","INFO:tensorflow:Step 1525, mean_score: 15.890625\n","INFO:tensorflow:Step 1530, mean_score: 15.890625\n","INFO:tensorflow:Step 1535, mean_score: 16.875000\n","INFO:tensorflow:Step 1540, mean_score: 16.890625\n","INFO:tensorflow:Step 1545, mean_score: 16.890625\n","INFO:tensorflow:Step 1550, mean_score: 16.890625\n","INFO:tensorflow:Step 1555, mean_score: 16.890625\n","INFO:tensorflow:Step 1560, mean_score: 16.890625\n","INFO:tensorflow:Step 1565, mean_score: 16.890625\n","INFO:tensorflow:Step 1570, mean_score: 16.890625\n","INFO:tensorflow:Step 1575, mean_score: 16.890625\n","INFO:tensorflow:Step 1580, mean_score: 16.890625\n","INFO:tensorflow:Step 1585, mean_score: 16.890625\n","INFO:tensorflow:Step 1590, mean_score: 16.890625\n","INFO:tensorflow:Step 1595, mean_score: 16.890625\n","INFO:tensorflow:Step 1600, mean_score: 16.890625\n","INFO:tensorflow:Step 1605, mean_score: 16.890625\n","INFO:tensorflow:Step 1610, mean_score: 16.890625\n","INFO:tensorflow:Step 1615, mean_score: 16.906250\n","INFO:tensorflow:Step 1620, mean_score: 16.906250\n","INFO:tensorflow:Step 1625, mean_score: 16.906250\n","INFO:tensorflow:Step 1630, mean_score: 16.906250\n","INFO:tensorflow:Step 1635, mean_score: 16.921875\n","INFO:tensorflow:Step 1640, mean_score: 17.890625\n","INFO:tensorflow:Step 1645, mean_score: 17.890625\n","INFO:tensorflow:Step 1650, mean_score: 17.890625\n","INFO:tensorflow:Step 1655, mean_score: 17.890625\n","INFO:tensorflow:Step 1660, mean_score: 17.890625\n","INFO:tensorflow:Step 1665, mean_score: 17.890625\n","INFO:tensorflow:Step 1670, mean_score: 17.890625\n","INFO:tensorflow:Step 1675, mean_score: 17.890625\n","INFO:tensorflow:Step 1680, mean_score: 17.890625\n","INFO:tensorflow:Step 1685, mean_score: 17.890625\n","INFO:tensorflow:Step 1690, mean_score: 17.890625\n","INFO:tensorflow:Step 1695, mean_score: 17.890625\n","INFO:tensorflow:Step 1700, mean_score: 17.890625\n","INFO:tensorflow:Step 1705, mean_score: 17.890625\n","INFO:tensorflow:Step 1710, mean_score: 17.890625\n","INFO:tensorflow:Step 1715, mean_score: 18.875000\n","INFO:tensorflow:Step 1720, mean_score: 18.890625\n","INFO:tensorflow:Step 1725, mean_score: 18.890625\n","INFO:tensorflow:Step 1730, mean_score: 18.890625\n","INFO:tensorflow:Step 1735, mean_score: 18.890625\n","INFO:tensorflow:Step 1740, mean_score: 18.890625\n","INFO:tensorflow:Step 1745, mean_score: 18.890625\n","INFO:tensorflow:Step 1750, mean_score: 18.890625\n","INFO:tensorflow:Step 1755, mean_score: 18.890625\n","INFO:tensorflow:Step 1760, mean_score: 18.890625\n","INFO:tensorflow:Step 1765, mean_score: 18.890625\n","INFO:tensorflow:Step 1770, mean_score: 18.890625\n","INFO:tensorflow:Step 1775, mean_score: 18.890625\n","INFO:tensorflow:Step 1780, mean_score: 18.890625\n","INFO:tensorflow:Step 1785, mean_score: 18.890625\n","INFO:tensorflow:Step 1790, mean_score: 18.890625\n","INFO:tensorflow:Step 1795, mean_score: 18.906250\n","INFO:tensorflow:Step 1800, mean_score: 18.906250\n","INFO:tensorflow:Step 1805, mean_score: 18.906250\n","INFO:tensorflow:Step 1810, mean_score: 18.906250\n","INFO:tensorflow:Step 1815, mean_score: 18.921875\n","INFO:tensorflow:Step 1820, mean_score: 19.890625\n","INFO:tensorflow:Step 1825, mean_score: 19.890625\n","INFO:tensorflow:Step 1830, mean_score: 19.890625\n","INFO:tensorflow:Step 1835, mean_score: 19.890625\n","INFO:tensorflow:Step 1840, mean_score: 19.890625\n","INFO:tensorflow:Step 1845, mean_score: 19.890625\n","INFO:tensorflow:Step 1850, mean_score: 19.890625\n","INFO:tensorflow:Step 1855, mean_score: 19.890625\n","INFO:tensorflow:Step 1860, mean_score: 19.890625\n","INFO:tensorflow:Step 1865, mean_score: 19.890625\n","INFO:tensorflow:Step 1870, mean_score: 19.890625\n","INFO:tensorflow:Step 1875, mean_score: 19.890625\n","INFO:tensorflow:Step 1880, mean_score: 19.890625\n","INFO:tensorflow:Step 1885, mean_score: 19.890625\n","INFO:tensorflow:Step 1890, mean_score: 19.890625\n","INFO:tensorflow:Step 1895, mean_score: 19.906250\n","INFO:tensorflow:Step 1900, mean_score: 19.921875\n","INFO:tensorflow:Step 1905, mean_score: 19.921875\n","INFO:tensorflow:Step 1910, mean_score: 19.921875\n","INFO:tensorflow:Step 1915, mean_score: 19.921875\n","INFO:tensorflow:Step 1920, mean_score: 19.921875\n","INFO:tensorflow:Step 1925, mean_score: 19.921875\n","INFO:tensorflow:Step 1930, mean_score: 19.921875\n","INFO:tensorflow:Step 1935, mean_score: 19.921875\n","INFO:tensorflow:Step 1940, mean_score: 19.921875\n","INFO:tensorflow:Step 1945, mean_score: 19.921875\n","INFO:tensorflow:Step 1950, mean_score: 19.921875\n","INFO:tensorflow:Step 1955, mean_score: 19.921875\n","INFO:tensorflow:Step 1960, mean_score: 19.921875\n","INFO:tensorflow:Step 1965, mean_score: 19.921875\n","INFO:tensorflow:Step 1970, mean_score: 19.921875\n","INFO:tensorflow:Step 1975, mean_score: 19.921875\n","INFO:tensorflow:Step 1980, mean_score: 19.921875\n","INFO:tensorflow:Step 1985, mean_score: 19.921875\n","INFO:tensorflow:Step 1990, mean_score: 19.921875\n","INFO:tensorflow:Step 1995, mean_score: 19.937500\n","INFO:tensorflow:Step 2000, mean_score: 19.937500\n","INFO:tensorflow:Step 2005, mean_score: 19.937500\n","INFO:tensorflow:Step 2010, mean_score: 19.937500\n","INFO:tensorflow:Step 2015, mean_score: 19.937500\n","INFO:tensorflow:Step 2020, mean_score: 19.937500\n","INFO:tensorflow:Step 2025, mean_score: 19.937500\n","INFO:tensorflow:Step 2030, mean_score: 19.937500\n","INFO:tensorflow:Step 2035, mean_score: 19.937500\n","INFO:tensorflow:Step 2040, mean_score: 19.937500\n","INFO:tensorflow:Step 2045, mean_score: 19.937500\n","INFO:tensorflow:Step 2050, mean_score: 19.937500\n","INFO:tensorflow:Step 2055, mean_score: 19.937500\n","INFO:tensorflow:Step 2060, mean_score: 19.937500\n","INFO:tensorflow:Step 2065, mean_score: 19.937500\n","INFO:tensorflow:Step 2070, mean_score: 19.937500\n"],"name":"stdout"}]},{"metadata":{"id":"WKWPdwP8BW_v","colab_type":"text"},"cell_type":"markdown","source":["The above command will run a single evaluation setting to get the results fast. We usually run a grid of different settings (sampling temperatures and whether to do initial no-ops). To do that, remove `eval_max_num_noops=8,eval_sampling_temps=[0.5]` from the command. You can override the evaluation settings:\n","\n","```\n","  --loop_hparams=game=pong,eval_max_num_noops=0,eval_sampling_temps=[0.0]\n"," ```\n"," \n"," The evaluator generates videos from the environment:"]},{"metadata":{"id":"At9LC5rxFyv2","colab_type":"code","outputId":"983b0e7a-2700-4e4a-d776-03c459669770","executionInfo":{"status":"ok","timestamp":1553253830168,"user_tz":-60,"elapsed":4036,"user":{"displayName":"Piotr Kozakowski","photoUrl":"","userId":"01014928596539690143"}},"colab":{"resources":{"http://localhost:8080/nbextensions/vid.mp4":{"data":"AAAAIGZ0eXBpc29tAAACAGlzb21pc28yYXZjMW1wNDEAAAAIZnJlZQACqh1tZGF0AAACrgYF//+q3EXpvebZSLeWLNgg2SPu73gyNjQgLSBjb3JlIDE1MiByMjg1NCBlOWE1OTAzIC0gSC4yNjQvTVBFRy00IEFWQyBjb2RlYyAtIENvcHlsZWZ0IDIwMDMtMjAxNyAtIGh0dHA6Ly93d3cudmlkZW9sYW4ub3JnL3gyNjQuaHRtbCAtIG9wdGlvbnM6IGNhYmFjPTEgcmVmPTMgZGVibG9jaz0xOjA6MCBhbmFseXNlPTB4MzoweDExMyBtZT1oZXggc3VibWU9NyBwc3k9MSBwc3lfcmQ9MS4wMDowLjAwIG1peGVkX3JlZj0xIG1lX3JhbmdlPTE2IGNocm9tYV9tZT0xIHRyZWxsaXM9MSA4eDhkY3Q9MSBjcW09MCBkZWFkem9uZT0yMSwxMSBmYXN0X3Bza2lwPTEgY2hyb21hX3FwX29mZnNldD0tMiB0aHJlYWRzPTMgbG9va2FoZWFkX3RocmVhZHM9MSBzbGljZWRfdGhyZWFkcz0wIG5yPTAgZGVjaW1hdGU9MSBpbnRlcmxhY2VkPTAgYmx1cmF5X2NvbXBhdD0wIGNvbnN0cmFpbmVkX2ludHJhPTAgYmZyYW1lcz0zIGJfcHlyYW1pZD0yIGJfYWRhcHQ9MSBiX2JpYXM9MCBkaXJlY3Q9MSB3ZWlnaHRiPTEgb3Blbl9nb3A9MCB3ZWlnaHRwPTIga2V5aW50PTI1MCBrZXlpbnRfbWluPTEwIHNjZW5lY3V0PTQwIGludHJhX3JlZnJlc2g9MCByY19sb29rYWhlYWQ9NDAgcmM9Y3JmIG1idHJlZT0xIGNyZj0yMy4wIHFjb21wPTAuNjAgcXBtaW49MCBxcG1heD02OSBxcHN0ZXA9NCBpcF9yYXRpbz0xLjQwIGFxPTE6MS4wMACAAAADkWWIhABvrNdXNvEPmO7lwVl73sPl0EDBzzvrz1O9Sgfa49FGnVhGNj4PrUzIEjAsiR14q5boH034au6fMfeHzW8BQIdLu5D8GWFcvhnUQvMLIDm/5fDlJWNI1pLZ0KekyKgRZvEg10IZZePvLcj64kGJzCMbJi6QZbX4WMzyM/ZwsXoWWPBmmlKBzFixHWdkptjcAYhpgDpXSILlIffpBFr5Fmv8Xdrl5eZtB/U18q6RE0tX2BrhekKyOZ5lJWnXZWIEICkLYIda8x0l/aAug9zkJAN2UJ5v8AfQgXgS7iPy41I11UQneH59QQ6r2Fy+bXVz7hKXvFUUQUW2NfwyAHSubAtKRV8FgrIBnKXwxAjc8zc/00LsdZVdehIaL1eI9qZtyap5GmVpF7ZJdkQbo7j2k9/o8Ztr6lwZrODqoujHSJK6V9bK0u9Et564zU+wWgftergJVAEl4m/D3N6/lD6Tni/a6bLzIcdcVjnfWLPAUBwAoj19NpxhAbe1VyiybbzF11k65OpExrnTpeyfXnWi2YKXmv6NMcvP6YS8WOK4pM7nWhyKetjJvO69p10oeh7Pv3PuQBq3kARIBKQ+MPYmymnbhgmxG/6w3hJ2A2Urz2k1DctVq7TiUCWnAReHSDqSpYcdQwxCm/lIpIwtl/dffgss5v+hhFs6NSNe3zqLc+wa/P6fKKBzHBPA6mZtXbiaJH0Y+5hMHtf92lFc+6I4pZ1q2XpI5Nr1V7em9lfehnp6KwZCFUTCrCle3ZgVn3/WlL0hiX3HqF/qGx1rSRBE7lqG2nGEQXx7BFJGNLF0vFVi0j2agV+lVqGOVlIxAjK3E9wGWVM0V7xAFGXQtxAYJ6qA6zMuM1AzlTqoEWcy+zkYm6Z2/Vn8RMHtpHaCW7GF05Wujcn0D05dR11MQem50GDiKxlzGighyGKWmfeex/qNBXelV1apol1nwDCUiSbC9fPUu70YI94kit+OnCdHe588u+9o5tqmvG+4ju2D1U0YtGzBJLwNtKIxTj+ycim3c7lWMz9gyNpdcRw85nQOO+UebN2j6KDuTy4XNxidtzFIcvo67EYGfl+q3WaPfQzFQLuOqvybvDRViyMxNwUidf7UCNcjzUMa0RFtd4HPTD4pR9pL0oOHG0XMOwvMlfvI/0tUl0nH8gKjxa10D+0pCAG9Sq76K3xNRb2QQ4PhDp7u3P+U7CY7JpR9qHasfUEAAAAxQZojbEb/+8M7YXuyASh7Kplen5y7UyO14JIrbok4XTbRQe3ORruR41lvOoDou8ClgAAAAB1BnkF4m/+AoGMvgzynj7iP1aLLDUqgN6Jo/suYQQAAACABnmJqRH+I4ID6PupU5REY3sf/aa003+ohQ9ie2tSAgAAAAbhBmmZJqEFomUwI3/3q7CUD4uvUkflAlpSidF3UDZUIJZsuBftR/Ot+q3GhwTn9egJDu/Q0u302gbAsB/IV/AFE9llDfK6lUW9694v3+SwMl6mllWP/WwEuWwr3bvYMZvmSan6TfqkNjfFQ2gIUyqpE9/WshNO04YsB0gSlJqYcVRMCqNuW9LOcW33er/NiW98OL0en+UGHVMigjfMRwtvQwEBx2TCEvqNHiHUZK47Ql9EBEwOC+9RX+RfTR+dz/gd9jggR1VC/W7S6VbwMD6OLZQ6pAbUGRTytAu/D3nUIHJpw2KDyU/GpGdKbj2WQLEoSWl/G1erWEOohpWvfDxHkGcHJyES1MZ1DvjVLlpUz1LfdDNETW9u5oytpQuaaon1mmF0VFggCZuK84NA/jjnYLjbiGa2mQVNuCz0xJcY1SlKU0HHZWh4xyJUcesClYPnBpMFzn1F68BZn1odfSkUsfnPO8Ozfwiiz7f3xZq0mYnuxahAz2t7bEpiZEBhl9vOvB5IwNQnQ2WMn3D1n/XqqbGvjTwCrUHeBDng8+lhZDYcGCdhZ+2y21qFiMlJaX159s9jZ9Nr1pQAAAFpBnoRFESzfiNYKja/0xDhfTAiTpHLhMdYxhLBGYN2yGAaE7v0DPQ2YrW+QLfMv2VvTlemA0fw0ITn/7H4QKf0Smi2r8j9M+wtBuP4ooX+ZG3O079gxDawGHzkAAAAeAZ6lakR/gQeVglqse6icagSNwuz60J4uGyZCePmbAAAASUGaqkmoQWyZTAjf6DCt6Yuhred9EHrqUOfXsyntdBaWSJrN7aPiRCVp6QCYQRjgAPaX306Ulc6wvo26P1UQNC4biU/od187SikAAAAxQZ7IRRUsn36T35XM0R9+0te5g6nnf972YKyQEsNo7mvZhGpibk4k6GYi5Xw1oII/gAAAACwBnud0RH+KcQxiHI/+tYvdUXZoto0LM0yuhiZoPS8TyCql7aDErRX1MckO7AAAACABnulqRH+Gdbgx/3UyLHgLBkL+GlWi7oO1l81SbcV5BQAAAC1BmutJqEFsmUwI3/QsGvLB67IEGcjaf9jQmmp2xwadCZqIc+QsmxjXnMZQAp4AAACeQZsNSeEKUmUwUVLG//1hfa7H0q8GiMhKKARrmw6Z4xdeyr7ujJme+WFMGuspP/+ISBWO3g4rVjHKMZczL9N457VGygcsdj7HL9XSQg7I9GyDcuezylwimiV/wFo/tnsP85jKo1N9u1U/c32ZRoJ6MLzOYkUW9bgpdQwc/vwurbPO6DqdSp7R4DyF3mA4Y/IOgh69Sz6dfNBC5f60MoAAAAAUAZ8sakR/iOA7CrYE++z83mxaT9kAAACCQZsxSeEOiZTAjf/7wc1PkbwRgr+fjnyNR8pfsXICQZIWSD29LO9OdoFeaCGC7kNQcI6QVUvVoNv//1awF0I8N0tfCIdJ96/qdtaHfFDnSTMhoQ9GjaS9EN6FwfTiM0n2rikHxnG2wBAjL4oKVFB1J9q4Q6oi+C8FkQ6nZ9weS8YziQAAADVBn09FFTyffW4h7bT+Vmnsg7Jwp4y5XEPjbVoz423RdhZiyhHoVMXw4WOwpIvW8VtedZMRgQAAACsBn250RH+HsPlaZzJ171H1c9c+P/yRU8UOm2VwDFLU0JopWl1F2Khtda0wAAAAHgGfcGpEf4TUK0IOqyu7XwKdBz1mDC7mcZOyimAiwAAAAGlBm3JJqEFomUwI3/wDwzHAex2BjGQbp8PmN+QZU96uDVeJsc29OEXjcldhiXFvZCjL1XZadbCRD86yiAgWQRsivW2zgAoC7DVCL9e3fSnRGKpVU/LGLClh3Y48sH47PWqbcpIEZ6IQMXUAAAByQZuUSeEKUmUwURLG//vD9/y1OQEdZZUsKUwNqfrKLfwKdmKP4MYlUfWBfMS6AB2YzPOOewKgSMfLoK01LvNA1j94yZ+2K1utCc4bPCb5fgtJtphyhALG+EJzcxQG8TmO93UCKf8P7LJwQsMlijgw4NSUAAAAMQGfs2pEf4etAFroO/Tre2AdB5O/z7sBdNC+yMviqgLALp+SqJXA7Pg/DPRDeTGlt4AAAABvQZu1SeEOiZTAjf/9dtYVwOWMQEFcGvjV4Syzg11kPr1s9125wYmt9rvuf3mx7RYJ6jF6ErfUzdz0/3v9ae/O7utr4LCBaSp+/oto7O3F1qWD/jeXpJo3T76eNJFHYJw8VuPijBUU3Vp8h5phb4yxAAAAO0Gb1knhDyZTAjf//XzJzOIH3VrgDqkUNgn86gOXkEPCMlG+9fgG6imSEzHDCVblECfv6lWOPhYa83SgAAAAbEGb90nhDyZTAjf/++FR2EKSC1to8pc6YSO3kyM61vFVTIqwBRLCgvhTW2uELbv4P1t6j3j6R5JEnH93R+6ccuNOfX+S8DB7HTWlWJcQLgqD+S2bSR1NJYKOXWfT3yjStlWid2o9fNbVb/WnQQAAAIZBmhlJ4Q8mUwURPG/74DjgdzkBXMYP2UA+265sg2RvtJWp7hbeXljZyPg5U1jpyoQAMkswdLvArqw2ReEXhiCOCghoFR9+KlVVnzlHwdrRAGNlgaJUCgaXSP04WBiefPhnJH+UNq4LFoPepHhK//kKOrVcBFaR4anoX/aA5EFLlBZC3z8HwQAAAC4BnjhqRH+HsRNxvly6e2k2NMUXC/vU7UPeiHPyDV2p/YvXZWhgk+q5Cpb6bHPQAAAASkGaOknhDyZTAjf/++A44Ilb6dcZ4EfRuTLTuDvo3J2im59tzroCzqfsO8XO0To/kC1I69bSu6OBA9I+wRfSH4s2YCwMyLJB4SXTAAAAXkGaXEnhDyZTBRE8b/wDoisgQ2L7kEBMp1YsiwFGNSG/OTV6DNluS1vSFNmiEavm3sOuILR8pgK4M6P/qHxZC276LRYUYYCavgfWkcLXP29c5JUTuaP0/h/7ytcJzS8AAAAwAZ57akR/h7Ee1u4Mf7i3wEEg6o2q6d802TgnxIfuhG+eE+Dw81v3A/7HzNqgsMZrAAAAZkGaf0nhDyZTAjf/++s+OA7zprt9FalqM30xdIjlcH2DZkkFtbEKP2KnwPRDHCudTxGL9Qj/p+X8zQUT6DsKw9cio1ehc/Va1ccw4608TJVh/tU58mts0w5trYhfnxrz7ghBbWrPgQAAAB9Bnp1FETzffU2jm3iBOi008YuPeRt0exkohagoWh6AAAAANwGevmpEf4k5No92pOx4d3tVe/jt/92QHqWDJ1gLa+IuwyWHPUwqtoP8fV9Ym+scICO5EURRBtAAAABkQZqhSahBaJlMFPG/++FR2CGqQlUqIWgPGeTAl0jkR8nWtez7chaAXKIsKU2quN90XktERKg37/m9ip4Dz5kOgzQscz6nyNY0f/sd6sj1HySlAcKWdh7zp6SDANGHziaLu7KO4QAAADABnsBqRH+JNh+VxVwUJ+Y3lV+rya5I/8TCxwR0ZLGKwtu7rocrr7JN3QZSGNke1uAAAABQQZrDSeEKUmUwUsb/++A44DvMfbctqjDdsKqNf0xOA39UYyLAzYIZyyzAIzc51s3jIS5ItuDvEtTXgkuVVUxzdDw7I8ZpdZavGnbqYY9+GhMAAAA1AZ7iakR/gQ9rhvIRB0ySHXw58101sjT3zRJJs0q1N75B5kzDs336j3kR4+OINWFbqx8fObwAAABPQZrkSeEOiZTAjf/74VHYIlYiKMdyQZmvfSW5HcZkgM1RFiY3gEH71z/3fVBzOwcraE90El70+ZDjx4HGc66IARVSqHcYq0qBrrRk30FtqQAAAFNBmwVJ4Q8mUwI3//vD7l3aeLA9vGddw8aOZ4u/aB1tWqMJJSgQk2EKksQOal+1aR5lX56hJg2jFlZ0Xy5wPHFC5sT+XuWXvNeiHLLg872OhCGzwQAAAGBBmyZJ4Q8mUwI3//vDobVTkBINCYpxqf92jigZ0Srp218/WFaj2TRlEHOqssGMM0Apc19xhhprfp/o+nyoo3nU4qIPx0toevbOZ/uMdat8h/rz3l6mAf6cx4JPPMV/58EAAABZQZtHSeEPJlMCN//7xIHZpGQEF4WlOSmv5sDvQ1HkPJ24b8mLANR7I1VevWJebie2fuXBDZeaCivuJcWRnz+sLxLjPYolWMX+OCeAR284M5BlQ+hDK45MhLkAAABfQZtpSeEPJlMFETxv++A44EJXr11WQvEfDTQVOZV8/+Xr3S8f8Sdkd5yhFoZ0Eykj3Xivg1BzZSewsPXa9nQ6gqdDgf7T2DmCjKrk7f2e2vK9ypwXpi7kgz9A0TgFjRgAAAA3AZ+IakR/h7ESVb5jShGX8VmBJn/MdoNQ3Z6fh4j/jmgTG8KfR4gLxPvHCfsFkI5Dojx4hvmkmAAAAEFBm4pJ4Q8mUwI3//vgOOAkGyUenL98r33EsuAShImg94uLBrTk6dhiDVHpVEky+DhSk5kCbedWMCnU+cpV31ziJwAAAHdBm6tJ4Q8mUwI3//vBcoj4CbBCQT2RExz0aNVXfJZgEMRjG13pvbjo3lZ6TLZP87/4osMgfYBX228XW1pfmU8k2l4pXWkF2ONrCPtLUsjWqi9NqYjz9uxWY1lh/hx5fARLNZPbPZztgBCOOKpjSsjqrRTvm2kDmwAAAGlBm8xJ4Q8mUwI3//vgPUQfdNRvnx3FP1pywn3LJO+BgkxSFcCfgS8xAc5jl7DjT9LzYEwD9isIU1yUsdzHqy/6ngxDd8Uf0UrLVqXY4UPbCNliO8n/ZegcRbOArrd0G+CXl9/2CFy7Fp0AAABdQZvtSeEPJlMCN//74D1EIHvgowka8CSMlOV1djkcwdrmN+M3w0ez2EF3Z8QgY9h2qKskrDR5Mk2+uEQHAHNJFzVoZ8OEnB6N5bVjkoQH9GCmMuS0GmqX+iknVHfRAAAAYEGaDknhDyZTAjf/++A44EPraJILnY0LuTUTVqG8Y4R+iod8qU50V48CYGzxUBWrE5/inR2Nxzjn9BK+lYDY6wizKJTKtNBOivIiofjj7puT2VHhc5kdvPo7rDnnm0xNBwAAAGpBmi9J4Q8mUwI3//vBnCJ2LA71GF25ddxrRcaDbZ6JrFKdXAbBcwenVA97XR386r26bMwGidvp2PmbPeLJeVTrLs40YO/0dkENxamY3tNtYwvFfYvB8Y8OeKM6xrpDUHsU4NOoFoGqtzCdAAAAYEGaUEnhDyZTAjf/+8N7/LIyAjCC10eRAkC4payNwDB4TaAJfH4zC+l57EAHdxnDWH/9FYP1pFGTKBZX4dIkEoguP5WochPIdwRBalFTzEmR/DYBfxJGM8ekovk7ApnUkAAAAHVBmnNJ4Q8mUwI3//wDwzHAd5j8RhmR9H0z98nDjYj/Nc7vsZGPARdvUDB2TCWaeccqi54nDBz6EstQtzy6v+wvloWKraz+uz/M18zKVlwhNHeSR4GZPnFXhk/6SFiPs3DdL9fCsVu/U6phJsgQpW9nziiIrsAAAAAlQZ6RRRE8331No5t4gTzZDbeKeBR95P1RisqLON/iq6lMAjgQ8QAAAEEBnrJqRH+HsRJVgDz377tnoNimjhTZvPY9yHf2fmv0wsOMF+wp88Zp9A47UCY9LcsoRhIaJ8uNlzZCrsZsy0xHSwAAAEdBmrRJqEFomUwI3/vgHICtV4/N+3TLBHjt+Fcn7DWjz8soUTNhnfkqkcMHOJu59OaFxhuVf+ti1/1ybHdnNcNjs9ZKiS4KwAAAAGJBmtVJ4QpSZTAjf/vB/wweEBHWjwtlVbql9dgODZ2kZPlG+D4KXaVCP4hDGqGFAR8WEBDCUpQahknTJjX5ZfkvoG4DfOZ2jNfwPVMAdpL7FdEsgo7WPCSMB38k1yCOW8g6+QAAAE1BmvZJ4Q6JlMCN//wRIX9RB98K5rsT6v40vwfGBFHd36t3Lmg6acxZXjPDoIP1nuZtPbJR67n3MUNs/jpJGr2SW9ZnVls5Jb7a9l/19AAAAFdBmxdJ4Q8mUwI3//vgR+AQ8tE14oytZoCQH1vmaiP6zviuiN8WiV39XBnBmeNGUBPj1lZw3uGVd1ASBgSvQit8rSzdGRcZWThyfaYyQbzG3uhl9FJ+f80AAABOQZs4SeEPJlMCN//74DjgOFPapEipwNQ5JLf/iFkT7vTC27zO0O1Hhr7Be9x5uGF5a9SD1TJTW/IfLAT7Nh7DJMwi4QswAkdulgB0QpQxAAAAd0GbWUnhDyZTAjf/+8Gbu08WB7ipWtNjMHVF/iPxLBqJR3dWf3VXwht118WdJ2yr8sybv5tw7Gkcen/6wXcv6hm1FN3LQa1BDmxg9YZtMWlhqjvMsaNo3uWx/3Zsn/9CB07I7ahNSdEvwIcfvP9pOIVFVaVu8QFYAAAAWkGbeknhDyZTAjf/+8Ohm1TlEJHbNU/mWRKFgYzLN2yQ8jXAB9wCzOsfbVP2lU1EY8QOgEzKqXo0R7iSjvfThP3Ejb6DhTio3axVyuxNGvXKyq/GmvvY+WeEeQAAAINBm51J4Q8mUwI3//vD7xujgdzj8Rhi1Ykq7OuqSWTi3Jyq4sFrVH+uGrfgBuxTFBDV13U4LCn/MgN+f4ZzOHXoTdRaVuXBhLEbCtA5Ms+hD9d2LkEvLXaBv0E8NnWBPx9w4ZTvBVStx4l9Na6nPeTrSSDBnaVQkpIMXYq0+x5SehxmeAAAACdBn7tFETzffVS1iQH6ubsRADaTWu5ionu4K0/0ak9zTnZggGqi+fkAAABTAZ/cakR/inEMZAKk66wt/cwHaBMWpfAc500hIeLekQWHoYVeIuwzwo3AD4zjCnVWjbJVYnBRhCSWWZm8rpySFJ3514IPVb1tZpdA21QwKGClE4UAAAB9QZvfSahBaJlMFPG/+8IGeHNgeeQ43E7kwPs1VGXGJw973890r2GIuyVyjsOqfdfCvVaSkH16R0M1/zHTQJQoPuOg2+H7r68wE56bXlYqz1zmB//ygylNGdR6NOHJuRI4gUP87HXbU2cwVemECH+lErFeByiJwCmM07PdY7YAAAArAZ/+akR/iTYflbJv66XME0GHkCF6mtJpaggCYpsevkqPWehQsYwUCZ42gAAAAHxBm+NJ4QpSZTAjf/vBm7tPFgQmOx8vn1I4MEdYOyJbxW9+T8JyO3vjlruuDeJ94kHVp6mE4InbW3mhydReOAPRHSvfloMOy1X4f1k61/T5AyQdQwSfNs/ceWh+FXJYPI5vBAAfACpz0P00Y0ZQiTyXgYR5dHA4282fpEnhAAAATEGeAUU0TJ99biHttP5WbMcugQMX/6p30Ja38txT+7rJCjWU5BS5Hu/Ped4RgDO/8NmgX1/xtKU+guus6I+GB0iyn+7zzesJ2+Dd1FQAAAA7AZ4gdER/h7D5W775k5qly41vVWL+9TyyQd/+taVv6vysDLrv4GuSfz03ASJUej48yG8+gYk5sPihZ4kAAAAvAZ4iakR/hNQrQyxiUa04Hw7yezOPlcbNtcK39Adv5vlNi03I0KSWCx97lgf8KiAAAABOQZokSahBaJlMCN/7w/fqde8Ef0LFVa6z8obTV+Z3fHanJstZCK+C9VlS5uYoIGfyFRCQyvYs8eXMoinX6wkePCQbIrdvJZtx24a1z2OBAAAASEGaRUnhClJlMCN/+8HPtQiIJq4KJFcEZqp5eZHjbfXvRNExlEoNtkHV/O1j3XsPkfmAogRZ6H3xb4fGyDO/psnXuk/zjrKbgQAAAGFBmmZJ4Q6JlMCN//vEgWa8cB7g7HnWVr5afQMBhPt/eF0XZINzgf1DTqt4LgNJ9n7jJeXS6jzqED9EpR64Rz+FTHrYHSWny0wQxBJBhnSb/iEP/JdZ86h/Wi9v9QW09VhrAAAAQEGah0nhDyZTAjf/++s+OAznGqGizRSqo3BVcpAWaqCB8ovSoNU/Q+cLNp60lGoY1XxqNFsDbFsRe3APbjIzVIEAAAA7QZqoSeEPJlMCN//74DjgJBsj6l+X0srb4l6b6JH2Ft3abZgPDnb7Plzz2pLEPDmuaBtebuDUX6hFZoAAAABcQZrJSeEPJlMCN//7x1Dlo+AmwPQ/8NXi9GKqmuY1UPDYZG3NPKYYZuf5RLSPKtiYcXru35hiQqjNXqPD4oWVOHCsaOXNK65RabsS8NeFVBf/1FIziqztdYeHaIAAAABfQZrrSeEPJlMFETxv++A44DvMfnbEhBgUx9vnA4Ho8ktSkKylpuzQvYY0CtUp4uNU/8qrDXOOYbe12K7QloChfe0qoEDZAp+Ls75A+yk9tDX6iPXAaWpTzsfeqD+dbeEAAAA4AZ8KakR/h7ETcb5ctlmh7dxoqP/yRU8UOqHP2aaGG/KPrV+Afm67tENSRXPgCr2j+g5e2szXEdIAAAA2QZsMSeEPJlMCN//74ByArVvp1VBzo6FR5mbIURGBYtSZMz2ctINnHT3rH2VjPX/zzdapOx0wAAAAW0GbLUnhDyZTAjf/+8PuXdp4sCBqPO04eytS9D46/Avaazen9hU5pU11Hmv3PBlT0PTYG549QlWOOVwWlJAXcvc1syzyvf1qFX3EhgXM1IOIcGnTwUuYkLr42jEAAABaQZtOSeEPJlMCN//7w3J3dNWQEg0JK5pzBQVfpWoOKLSI1/VmJjcsCrBle+ncG6tiC7wuDEZ5USHcZC55aKT2OIPW9YHg1aVeYKYR0wQWQDP1DG/oXsGKUwexAAAAV0Gbb0nhDyZTAjf/+8HPtP3gfFlpFrQPirkvcmukg1+owUzumItdD2e8mcIUJrYgE5Sa+gOebQM6C3UJSm7qj2KyQP6H+uiQQf6dQ3nUsv4RzQXEB9IZsQAAAHpBm5BJ4Q8mUwI3//vEgWa8cB37/p9ZF86iN25nKIigWeKhgv7YJ/UNEZWwthTtyyAJEfsOe3Z46qB8Y9u4aNguQI82wq+DpdZlQKDejU+B0T7yh/vWZYzvEjdR8mEIAxxdRZpU68SFTtrnTu7KpqA68eg72PoiJCihwAAAAEtBm7FJ4Q8mUwI3//wC0LvBNTmF1Z9IfU+yXng9Xlh5jOvxWurF0OzyKrrbAG/LhzXREODrp9/66YlkypRbzROY6MEMzZYnzbptXwMAAAAzQZvSSeEPJlMCN//74ByBNNkfODOhV2l7PuVP6jv9DFZM/xRiaZ7NjolC2zFnBs1tnrqTAAAAUUGb80nhDyZTAjf/+8dnvoPCAfKn6R60RAOrjtGLSYzJf83HxmDNtKNcZR3s7Q17gEYux4ptcFO6KrQhBRjlRY6a9irUPBzxRSHCk6HKHHiDGAAAAGlBmhRJ4Q8mUwI3//v+jz7KI4HFxHpTS0oxWxOKguK69Iqb7Pv80RE8zRIcDeLBSrbSsht1y6Ct5KnKYedMgtXkRlKks+2+B7FYsVvxecCZHOfXyHzislTlSVnRDtTX3Xm3aOyRaBXu5sYAAABAQZo1SeEPJlMCN//7wZmMlyhHAV+DQxbHCYKLdsYP8alRhrAXEGFKASYmbvEaFUC8jJbbP3+VwoeATImNqxZxmQAAADBBmlZJ4Q8mUwI3//vgHIE1raJILnXvo8sZurnfn67164Fu2/h+6+ydNTr/+OxuokgAAABKQZp3SeEPJlMCN//7w+7Y4o+oh8mWdvl4UGi7/t84ZAPcQ9+MBMDQq//WhRWSosr28nZ9S+NamaT8AaHjCriTA80H12cPmqsMBr0AAABZQZqYSeEPJlMCN//76z9RD7r4RWowH8hXR1qM2/0RPeiKZ9pQxjXugxDv5Y/OlrDSWjfi+WFQhmWEG3gYNg2K8SmfzQY7iAamgvUBbTcVy2n5t1ecNgemmcEAAABoQZq7SeEPJlMCN//77lcOB3OQuw6HD+AjOfELwUGp9SJicXV1sFnhvaKjnvCxiZwl2jex2l9TUZcaQ8RCwsTwL7Cb3QWKztEEttb0Jo9sjMWDnWrh5iANpRQyrCjmn+qlJDlANP179jgAAAAeQZ7ZRRE8331No5t4g3EwRcCB6QcQmRnRJJhqS7+1AAAAQwGe+mpEf4exElW+Y0rm0PJ/wbTD/lW+CtIwcWOFlp4l4dVOEXYZReMT7JXZAnpNCK1hAhXjScp5NYYnhBjzjGD/rSAAAABpQZr9SahBaJlMFPG/++FR2ESusZKMxurmUPsaoB3yvXwtgjhKcOw7RRICF1EOtbOnMwE+WuGfg86L92P86V4YhcXg6X7+adBswWrcPNmvGF/+VM2MJnXpcNhI8G84ECcACh3CpkqclfZfAAAAJAGfHGpEf4jgOwqrh+lSqMFWNRECeaFV48q85DBbzjXglEAaLwAAAR5BmwFJ4QpSZTAjf9P1UW2P31QTlnuH9bjbYJ7XWt/ai9e9KF19PysQ7DabFT0nEsV0yceUNAJX0BBUVv6dG9+bf4cKP6gx1WQaTSkLrl6RjZl8G8g+mu0jnqW9tdKrF6RdsgVAMdnm94bRz+PzA54Zyuy3ATHLl2cUlXEyCW9KigPQ0J5U3y/k52ZP7lyfuZ5la/1Q9Pj6hGjOXZEUFpJACy7fj69cHAm5TQWn8MCi9bYtYpHNqiSnSLn4A+l5LrMldm0cXmrlCIKI+1dB0+ItJmzaYll6630NuKKUYK9we5JspVhzVBRfoTKIYbqUfFWokuJzg/aLs0nAnaPalKUUmOI8qX3c1g70zXzdIptzRWo0lBpoXPimAkfYfsLAAAAAUkGfP0U0TJ/Tuf49tp/KzFPCIEDEvSN1rnYJbiqzzxVfJ/8OCLsf7ljUGroVfs9mei4RJMtv7Zcn6zTeStOI46ruvNTH1vZzMeIyae9K1E96hHUAAAA5AZ9edER/h7D5Wmcyc1S5cFCuXC/vU/dYydSXmT4ray0A6VWt90PLvJGDUf11AKA6FxyAGoDLyBHBAAAAMwGfQGpEf9fFlYaEHVZXMWGGCx/WUqHrS0vhwUrligXTGXJS4DE/QUk+ncgeHLqWgy+vzgAAAKVBm0NJqEFomUwU8b/oMJ46ACuOTXqhV0ljwLZShz4Tc58bkcnq9w3uogc6YHf0kLpTfoclGUCaCZ4AxrrKOqJsYXwsspdKo6DAKGat7GSKdIulJbXbhuc/IvIa1CF/DNWzV6r+cO1H8weRmlyh8KFBzm3Sb9biHpOZEdIVGO3TyFoUJ3SD0ZzyH7EROwGbMEzO1FkpRP4jj2zcTUYldf8k10XzYXEAAAA1AZ9iakR/mI64NNsipLH3uhfv8LrdLfjRclwIPczo2kTYfFSqf6TNBKKbBeuQKHqSSWNQZ4QAAAA4QZtlSeEKUmUwUsb/++A44EJXr12b4QrXF3nTB+IiC36W7K9k0PInXSBIMLwHooA0xe2aJ27xzqEAAAArAZ+EakR/h7ESVb5jShGX8VlCE/7Y/nQZhvcjseRoWw3awq8RtUipbXM1BQAAALhBm4dJ4Q6JlMFExv/9Y3sZGUqwJ/ZJQ3AZgBXt9Ufjt4ZzNxGyAfjakJd+hT55Tw9cNhg+w6ug3vz/oXT+HSXuNuSH3FG5/LjmCDKvQYF1YEcEbLkXk9fh9BECrWP0tPknccp1bt+CPAa7hxxvktQiTyB05TND7a4HjRFbvq5TqgaLwDS28Cxz/2MBBMq7yg+pQVD0bX2Sl/mRKrv1ISvuJ2v1VJxyBQ+KKcEXAsgkK1RpCQS4jATBAAAAJgGfpmpEf5BbOR6mhN42Utt0sKpEJIAR7fw4EeQJlSVwz1K/GoHlAAAAeEGbqknhDyZTAjf/++A44Hc4+26pb+I0qiyyaTY86qeqgj4kwIFKO32Zd3pJfJsEyCt3X/2gIau4QMjcRp2plA5yiGVX5wEKuxlMS7/gGuv6Qw92srH94+7Mg9EV6V3w0yihnRIaFYSL1USPSa2cDMI6VI+grv22GAAAADtBn8hFETzffU2pysalekPthsm2ZMTNyS424hytUB9PxsDoKgpOtQiLMZ0iPUVzASv533olMcsSGUtWzgAAACQBn+lqRH+BB5WG0Yk9tcQHTHiL9VohRhZ/SCKeLiSmHZaZBYEAAAC8QZvuSahBaJlMCN/HqeNQwz8YC6lddTUaq/wlNA30DyObwUgGgN9P+XYVKGiTwxMURC6nsd69EfWbXorW5UKp/iLDB0tXqGqCdJCKds+YuzihvwR1B4oMcWI3hPmbK4zhnIvnkOawDk1lzJMAFAd5Xs6In7X2yeUfKr5WJHmhlN0mNbPW0T7jDw7LEkRVdNs79QlUlWSBMCb7P0EyD2FPtivHzwZbF/QBKslG6u1N02WmHh1SzyYHC6SYhYAAAABSQZ4MRREsn812AW0R7+Pctap0S/62v/Y0HOJuqNfDpqTpeIBGPqgtfdqY1GtZkmBpYYfpCvY9UDU+DGLJFacH14P9elmz+ffaFFAfxY7tFIA+wgAAAD0Bnit0RH+ezJaXm0/z6MpMxrufrdPmfUHceEr1rbvBdYpSb3s5MV2ylnH+g7e3nIEQSIuWIYQBSs9KfJ9JAAAAKQGeLWpEf4Z1uDH/duS/8VkmZodKFFLza+EKeodh2cw31hLmuImTU+zBAAAANkGaMEmoQWyZTBRMb/QsGt4s8uUfHR6wH3Eruy7jdPNSHdGx4H2+eY45bHHjrvu+npGPfuKs4QAAACABnk9qRH+BB5WDztJXfGFrhk6Xb6COb2LMvQK/NgrJwAAAAI5BmlFJ4QpSZTAjf/0fwO8d+rKuAd3nWKoNCPGmENKEldonVkGjg98w4Me5Z48NellC0FxCEd/kdNLe1cM8GE783LvVTMHmMX1wOrayRMiu60nJZ7AhcL3MR8AOnKU7oGDehMw21vsadVsuyATMJtrcH76ypGC8FIQF3X4VlmnhpX1WTqIkRvuvbCddYWNgAAAAakGacknhDomUwI3/+/6N/UQ+1Ea1+CN+lmcvTOPgijtDvzZhN0D2G8CBJq0xATOAf5U01N9dCRuCxIu7R2vxMfg848kgrRbFR6x23LTcr2silj6PlhF7bXGnsSksL0sP7ItpzFz6DxDAfIEAAABQQZqTSeEPJlMCN//74D1EObFx0Ccz500Sa/EklQzLOUat9ZefscXQjlR3Mx55ToX4Xuw4yk5jOe10fAnmAMrV30VKTa2W5vE12YQot/5YC3AAAABHQZq0SeEPJlMCN//74DjgiVvp1xjN3F4nu5NzvpHqsfivvp/bbY0P4LsFxB1hB6kH8V7RfuzmdrHI2vT7+H8cVEjDPvhujc8AAABdQZrYSeEPJlMCN//7wcOsd4IGgTeySKCPIuwxg0YYo+96hJ2UPP5xSEcCoL3qXVqKFs/+7whBdM8r0BuQQeLuwVm53f5UyL9EOx7Lh69/Ofouy30Ld83TsPD9W9AxAAAAREGe9kURPJ+A6iW+6tW/uR9lhCZLeuQ4oZcBFeKAu5UN4mG1qmIlUSZUB+VEhmQryk+9/puz2UHFvGA+dUoS1cGBKKWIAAAAOgGfFXREf4egEy+iMAaiy7zEWbVUnq4AyqydOp/rPmZg3eGRT1NmWJ3z0oQRnaNxgE548Z1QPX7xmOUAAAA1AZ8XakR/hnW4Mf928rmxp4HnnBMUqhe0zvX0oeNpgN+7Uu+KKS3fCuekbQu5jS3lzKvhqNkAAABAQZsZSahBaJlMCN/77mtCBIDC4q6hbKW6i+2ykwSQy3Fgc2ZOrZmrvUZMVzT54+uJ2QGpl/6fQueAl+NdAyBz9AAAADRBmzpJ4QpSZTAjf/vgHIE02R84M7FU/ejsVHqM66DEeEz+EEL17DltFUqXJViZAeKAf2cfAAAAZUGbW0nhDomUwI3/++FR2EKSC1to+4xCmEjt5LncA0mqNBUks+Ym0IghLbBCqyHI+/o4VwUizGvOHaK62P8WHO1D41LvGG5uRQbBRxa2T+sO0fWvHLEeFz60KXZuVfnunSGseJ/OAAAAfEGbf0nhDyZTAjf/+8Gbu08WB7eMlZakaObAXDazwj3bsg0VKVyXEZdPl4Ml5p/16cJP/U78Iz8EJFZ2maeMrllGWqaR/jYpXDX36keDWXpidqMEWfOXkQ8fZImQl1OdRErD09YOHHaf5NGC+6km9YfOUDqGX9Pq3XQRj3EAAABXQZ+dRRE8n31uIfah+Vvv5ElY/3/K6wW7DB17bbOdxcR0rSz4be7jHQy0EgS/lLK0iaQUrBIQrvTwd5mZPmZ4ApMYBTCkxmGGY4ysfzByNNQiAJD3piF9AAAAKQGfvHREf4Bw4zB2SpmzpJNdP//8f6dskHoit65Hrw6c76E2VmXkD2UoAAAAKAGfvmpEf4TUK0MsYlS18DDnFz7Y2UwUKL7z8XXLeoLGQWrrq0RBevoAAABiQZuhSahBaJlMFPG/++5XVEH0KHETUXGSGsXNMJtH/jNeRzgaZvNQtjBuinxU/liYq+MxIz8WoClfRA/KTLqF4/Fa0SAqiv0o9OwWGtPhEpGIZo1GD+Cou/3P4e/6cRklt4EAAABAAZ/AakR/h7ESVYBQhM1ypHXwCXiUkjPjSDADDv8Q18gyx9OBwNx7koSQDcnOOQ671+kC98Jm6w05OE/cunBeYAAAAHNBm8VJ4QpSZTAjf/vhUdgh9eh1QZyogZ7+hmtlbPNslgd7qIrlLoeRIb/et4+8Vq1ePF+NNLB4Nr5y2LIdg8c54s2iT1WYUe5vxH2qfxAHd6oe3EPUU6bAiHT77lV4nOsi77nk7apWUyTxHGqurceFxTOBAAAALkGf40U0TJ9/cH7YcKBioTwOqZoa6GOmGKX7krY2+ZZJFD1Q548KXkilnIlw490AAAA1AZ4CdER/h7D4ZJjNasvMvYzkHey/9PvUWs7fgtuPBRXLAX7CfyYyOMWD0kVUTTq8sgZLsG0AAAAkAZ4EakR/hNQrQmu507TYOsLT3zhWb32toH/fgbuu1qCFxEeBAAAAVkGaBkmoQWiZTAjf/XbliF2xAJMQm7dNMvTYZ4BH/3IqFg9bdiX1lecwbkkdjr1LW6pM2UN1kfniZvPkJBXmPFV59ScvpJEqbeREB1KHz3B7jm6BYferAAAAWEGaJ0nhClJlMCN//APDMcB3mPztiP3uapuxK9w1XpJGf1ey0bVBH2OVz7uNe1T/11fYO9Iv7Fc7kkj5DXP4lwVGQc+vNvTeM4IHYwYcrs9kxs6CASSVaQ0AAABKQZpISeEOiZTAjf/74DjgJDW0SKRZBRnHc1aiJ3lEevKh7RRg1AT8zcCut9u3ECa53MXXjAVR/5TQqy9iMtme5BHk09PtAdSIfrAAAABbQZppSeEPJlMCN//7wZu7TxYEJkiWNfakae+YsHrw4VAfICk7Nb7oLk+0i+Mpuy4D8f/TgFMT1z8D+3aoAm/Am00CBytVnpVYqPkD+waLnY1NOQR2k51+vgGAegAAAExBmopJ4Q8mUwI3//vDoZtU5QkPFvbm4haJgZvT6vbhkh4T4F3xSdoAUekhgoFrQuyi5WLE6GUiQEppgDoxQGgzx7AhRZZEztPG3fTBAAAAQ0Gaq0nhDyZTAjf/++5XDgO8x+Iwxa46cT58hPz/3xXGJ1VivV3KaFJ/3WQxqtPg3LKAamhPU0NAdWGoS0oO7WTR/EAAAABfQZrNSeEPJlMFETxv++A44EJXr4MQrxHw02XeFmca+qmRNOoUcRRGPhSgHkAcuC9vwQHiT3krg6bT/BkQSItsA8CwlnnzN24kW98J7AEeP0NiNW7ZKbpdvkMpFyUz8l4AAAAyAZ7sakR/h7ESVb5ji8/rDDaLJg/7XqiPJ1B0qKo+hBZfe7U4Rdhl3xjVAHcpz+vB/FEAAAAvQZruSeEPJlMCN//74ByArMYDm/bz8wxn23X7eeY8ueYRJcmDuVxZyy1hXLUX1ucAAABgQZsPSeEPJlMCN//7x0/RSwgInmw7ggx6hfLg8Mo63KinbZ7POUyq9lLfsMaS3xKvOrwHbYeQNzz2RoRgXOv/exp6/Mde6Dums6ZRED4RzeaCX4PJ9AgJuwbmO9tXNIbBAAAAZUGbMEnhDyZTAjf/+/6N+OB6zo5Fy5eHpx0/tejGu520bLkOKYZmUYwoEKtu0gsWwF6ztlWdsn2cBXYI99l6dfY0me7vv4CJ7ndd4PxDQYUsSz0tIM6CYnR8oOrc/Gfexffa12GAAAAATkGbUUnhDyZTAjf//CnvmQIc2OrRA8MwDA9HMkp0e+MI7Pm6epVIxMw29SeVpA1ToAxIhYkHTNN5ShviDJ3BuOFRgp1C9Ah/PPMQYFk/kwAAADJBm3JJ4Q8mUwI3//vgRKhKrrGIpjJo14wE5fxn+jkJ7L0hg0Y/aAfD07llO5gOpjeyTQAAAEtBm5ZJ4Q8mUwI3//vDootTkBHWutfyL3CvbJ+vrkeH295y995Ase9ZpuUtjRd/dTZ08jGdVP/9jByukXhfEKbwAjssjjm1tYa2KIQAAABAQZ+0RRE8n4DqJb7q1BrisxfO905/2QQCdm2BOraG6khWyZ2qOCWjta0B7hKGZCu7bgEpU4a1B1yA6YtIK+nvcgAAADEBn9N0RH+HsQIbL4x/P9Mv+bVwIYmppI+PEO0zwOR0FYU9+OSfjcQXZbl+lOJ/yuuBAAAALwGf1WpEf4Z1uDH/duR8Z1p/stCcUr053SmodhihH153Yzn0yYfGEtEcJcDx5XQwAAAAR0Gb2EmoQWiZTBTxv/vuVw4DvMj6PtjWXbeCDWe8HbgEv/NcdLhFqVKB8exZhz3jCX6jioMxFUtp3r/RC8ltZJfrteOt8L2BAAAAIQGf92pEf4EHlYPO0DFLXPMyZmip/4spisSOYUIUhCPWMQAAAIxBm/lJ4QpSZTAjf/vBGT5YQIc2HKxMsdoJ483hDJQtEt1oSVY1OtNsacpcCHC1hLiyide7+igKkUYmgW63zjDEC973kPfv/kGvM9IQJRrdPvgu8NP7Gg4FonN9g1w21bXVxN2U+b+Kba5ydK3Fwptozj887wZ5U43AxRVa5ir6jhSItTpA4qmjsWn/cAAAAJVBmh1J4Q6JlMCN//vBzU+RvBF1xHc1oj+U/TZywQssct1DlyEQP3bm66D0erMGNf+QJB5DFmPKyxwSETeXY47tb9OQ+xlPePfb355uJNYeS2Og2DJDc6qxeP+J/6dc/Sv/p/LXg678yyOLMGYC22x0VzI6myko3Y2jbnwPZYHoVN3W9bTc2x7RxOC/79m7UesteT5niwAAAEpBnjtFETyffW4h7bT+QggNBgC1Tx90dvwgqWi/2bEAaoy9k8G5OWvd6LsgFQzqIuBbNMM8rzbsW0pRv4t/6wnIDQ/mCBNdzbJDDAAAAEoBnlp0RH+AcYN6mjmSXWqx+E1SBd6filHrC4NN/gK7ATVZjyRHQcmyIQAZumDsG240peDLOLW3gVbq4LmG3Njs3hJcfGYGdq1drwAAADUBnlxqRH+E1CtCDqsru19uEmwmchClU8Q9+ckZEbxoQUVCtJ39bW/wYXdNpPtmAsJQybL9wQAAAHlBml9JqEFomUwU8b/8A8MxwQo/6EiaK37P7zGFrPazCtkNNv4vIp0kq5kSnVxolGnW0VIejbpP5ZiiC+I/F0ma3+IF7CpxnIMPf4egyyXNCkjhnJHSXHusxYcAUqAOlV5uTf2eoJk7NXqf/lcR5JWs+Sf/4CuRrynQAAAASgGefmpEf4ERwYTMPCynKnF/ULnVV0i5kdCM7P0vSFSGSpVSXT+9+3FrB1//zAGZ86aGphUfLcz/y3mFRVu6iiyShLlf1IBxT4rwAAAAk0GaY0nhClJlMCN/+8HrVh2DajAEnVrgIluIofrb5gjHi7I0o+eSqEr/KXkOiwTxu2NN5bbjgK7wNu3j1ZZsWnH+OrtoEaJsgcV6iH5Pd7Uko8sgwhTol2e1ZdJ/icre4wKyTmNr5bcC5di+bdY2nLvrzZtdpa9DX7QJestuRmj2U10Cl1As2ZueLxIRqj2wGBHESQAAADpBnoFFNEyff3B+2HCdri0/UUVYHd8epT2d+YMYObZm8nW/ez46SIN8EbMlhsLlCSFaW4sFX2nF2P34AAAANAGeoHREf4ew+GS+7UvFV/w2hI9vC/vTUrzGx1QcGgUQ8rwi7DMOqBjdSkbuDSjnaADQL3kAAAA0AZ6iakR/h417ZNpNh7L591MHWZo7LvvoIcpQmohn+Xd70kfnW1Ud5bf8T7beBDmR8Xcd8AAAAGhBmqZJqEFomUwI38oIQ86YuCDOQ5tOdExAsQJusAUoBl6PXbngL0EHOm/XgoQK4cIAx16UWM0dQcvANQqKp2EBOaWAoRwWN1bqOg1bQrKKmV/ijRbo+vvcxfowunt49/emVGo9MQAXvQAAAE1BnsRFESzfg/tuGYHKjTB5ZID1kuLiL2dqNk66JpJk8w1d6Kcc2wUWTZM13JvgxIJn74EkYNkgC4PjHLlyOy0Hwh6m6YSunCOq8UCBEQAAACcBnuVqRH+BB5XFIVzyClcf37dfxipiNXufl10iycx4d2lwwnYOlmEAAACBQZrqSahBbJlMCN/7xIFnX2UJY3B5G9rwJVPFBhNUhdhx/fLZ3fI40N7ZzzVF/KfTN++yvpUue1sMAod2GLm7S5Mnr+HDpXolz9xb/2QF7HgW6iFPZ0TSu72RZcJPzQzsgubn35XTFLL3mozUEl03Plt0iKzniOced0OUW4K+kw+BAAAAXUGfCEUVLJ99PbzhStPV0DPWCp739xrVtKTUsEdCecaEM+e0OjfkkUkuNJFophvGmqzcNDPP+G6mO1enkonn7obx3Z30VXOOnNy2O0Eiwd48Dx3VZnk6OQouxGUAgAAAADUBnyd0RH+HoC9/9Y3zHp1BIl2MfTFJGRIlfr2DnVE0Sw1ikOxYlKnFwH2Ru+ncOjZ2P3qIlgAAADEBnylqRH+H5CPxj/uvJW0xJiOenPhWEISI76CVOVnP0kO3twxfh5J19iCz++FpQtdBAAAAZ0GbLEmoQWyZTBRMb/wDwzHA7nHfcLHdl3CBcWDNUIHw5qJ/uI/Rxc+UScIeicTquE85DXU/XssWAvri4jZAdY1hh7+q6h1EQeWGPd8v9kVk45hJzDq4QFLRn0armvMglMXFFZrGuXAAAABFAZ9LakR/gQeVg87ywhqQ0UZ0H9SsmVmcAjAxEPJNkzxuBCGvKgW3OfU6j6sPxhv4LrsmE68KdHBdFoxyuLi785Fg1b6+AAAAhEGbTUnhClJlMCN/+8di5MOwazCfUdt4QXA7i2vzFU3g/WQsndT77JVDSX3dyiHIJ+dcHuYgNieBhDWi40XytcT7MzqGHQ8BKd7emcKr7t9ZMy2246DgCMU6O8nfWZamMp0hGAd1hoDnPEpKtXES8lvdf8QvUpJi7H+RY6YmEvOOXCg14QAAAHRBm3FJ4Q6JlMCN//lIWPInJdToRM/Di8/wOIa7fAwqyFmQUIqzAB8UxzfJluX0+Td7Rj1n+911ubH3x5Ng8MZiV90GkBT1677/EI4xf49/XAJzyB4a0vGsKyC/rL2cdk9MBpe8qr+6ehvxBOlZU5I+9k8HgQAAAExBn49FETyffW4h7bT+VmzHLoEDF+QSPtyOibdtnOwtVf938/IoeR8/CfdhBqH8efnQNzSTcptOtFhLNhHsC9QRs6JUpZ3g6L+0o2hvAAAAPAGfrnREf4ew+VpnSNtXMyHIouF/ep2oe9EQ/UgNEJmOXn3+zo/llDpYdbcSubbJNj3j3tm/n+QkREY4VgAAADQBn7BqRH+E1CtQOK56tKoCiM8K0R3dr/Xh4xtb5DeCzfJ3ISlqMnAuma7Jff3Osd2p1kQ4AAAAeUGbs0moQWiZTBTxv/wDwzHBCj/nXShLUbtqalbXUXB4c+Jr8UsTF9vlJYy5JAsd6N3YqAwRTAVVRhIyCYy8OFra9SCQJIdvyrOlBtx6e4Nm3jz/vVxUbjh4mewKCUl3melRvDxX4XCovgeHkgRgAO0Q4EkZqLJtHfEAAABFAZ/SakR/h7DkTgDV56NbnZl/ml39qIOiEFfzBef+RWvHUibHEjD4zje6WjYo61QyA138/8GbcEZEmbUXjgjCWt9nPALYAAAApUGb10nhClJlMCN/+8HrVh2DgEzkPQSYjrWyzR/cxxfpkyJIH7uTdsyWu2DHpet18vbjP31VvB4OvU2zGwNJhH03pIxozyt5CjT0DxM8Eu6iU+hynjaNStmVahBNpSsriT7NNCrjVDXNYImUZaYR5CNKrD2nPtelij1qj2HUukGEN0p8HTXl4uhVEfUSYArW7VPcWydU5vlLZYaJeCj6iL7TK9654AAAAGFBn/VFNEyff3COo0/prpbjGuq+Yz9qQ/0UtxBbR2wSbc63al5dG7nbMBmTmMJALX+gAAHd4U3sZUwDf/j8MBK/Z+midrtvZ/dhX7BGy5kbGSKhcRoIJqsRYrNy2G/mf35bAAAARwGeFHREf4ew+GSYyDT9FKwHNGCE3zTvKrR2YTzumc5w1YU5tpz8yyBsGigtzFTb4gkcmmGk8BosAJh8sW1hPLsi/oP9XLQpAAAAKgGeFmpEf4eNe2TaahDYDlAE00B5TiGe5P3cqwgz7Jz7dAwop6UVlxL/wQAAAMZBmhtJqEFomUwI3/vBm7tPFg7wyzAD0Mnjaxg9o1oowtbJ4Iq6fMSBsDEVXAUv3Ek0V6bUA9AyXfdXplJl5OklAZrAT2DLE0CTU2Gb1ph7uMvZ6e730j1uTDGRcia83XP9lLWy0N2lgd3miSshfYKbuLgLsYsiwlykYvR7QV0d0CqGOqo55SN1Jjj3z4lpX6c0RGgzwd3ziVIoWQeFFwVeqqCbllBnRSSRaTmdGVG2XFLb5K4/MLixbmd/kI/o/AQS9ak96IEAAABVQZ45RREsn31uIe20/lZsxy6BAxfkEj7c3vTu2znWnO0o+LL2TewvhwqHeJD01QF78BcCbCuwQBKbqpl90L6uNZIPK2wJ2IYROSzDobv1Y58U1sN/QAAAAEABnlh0RH+HsPlaZzJL659u40VH/5IqeKHVDn7np90TGQtQWPbiEvzRpaGCnUR9bwO4lkoO9UNv4cqr2uoaHtjRAAAAOQGeWmpEf4TUK0IOqyRzDxju6fds63t8C0diXyWCw0nZtf4XS2WBIN2TqV3HMYRmD9ycxi04LTVA2AAAAIBBml1JqEFsmUwUTG/8A8MxwlzTJ87d34cixfMcHbv+VHu06EPCplmBHBDuntSkT2yRgALAJie3XxuB3xfOH/D7ydP8ErkN7FeqXmM3xkd7kLShI3yZeVnO9ZFHHuwsNecvxwV8jY12J7GCNBCKjcvRWa0RiLGeyRhv7w7uCpPSTQAAADwBnnxqRH+J9+qocEqY3mgUADG/qVAbloOL+4LbHt1Vueid0r6SqxaL3S0xMXyghgv+Ji1+lINFTyTiXpEAAAB/QZphSeEKUmUwI3/74VHYelCJVnhRTg1BLYiZcXT5dxyk3eX+rJNoNiJEdunbKmh+C5I7HMdbrb16UYYiG+jet6KLxaN1w2WtwbVe3TezkTPnAeJ+tInJzjZlV1FqpcNW2arifXJoFuPI/Iz2s42jmNCDbRbNxQpODftQ9RldQAAAADxBnp9FNEyfeyCnmveYuSODs+6cbniKe/vM4S56nhz+m9i8Zex7ONz5LSvfzcVpUjqdCJWt7Tx6czK0loAAAAA/AZ6+dER/h7D4ZJjHyl5mjfaYPRhVd6h/FiZQAFr6JH2QqvbB/RwwwlywmDQA/o5AxedWALtCD568UWghNBmBAAAAMQGeoGpEf4k2H5XFXAtLnyu/P7BuNQSHZxysH6gKDLB09GXBu1f/byom4T1PU9j8d7AAAABwQZqlSahBaJlMCN/7wZu7TxYPY30W6JoFSNQpwVvW1QrXI4toGmRdXr6ug5AcIzYxb3sVhpDt2R6YmUdAxjB6i/FCQrVb4p6cx70BUXqsfihn6ICLLcl0nvY1vzENxIRzj0j24J7cCRqf3p3r+wKgQQAAAE5BnsNFESyffW4h9qH5WbMXYFj/ael8uqKOT4Kdx343HScFjvKFNiDbjO2olivLVLdWGTyFi/Nn+rmW8yKVjA5v2wbZq/jG80u3A30934AAAAA4AZ7idER/h7D5WmcyS39QLC0VH/5IqeKHVEQEr8EriQ0rMLKqwWs2RvSTdcH/r+SXrvRjzQL2kGkAAAAvAZ7kakR/hNQrVcoxKlr42ol5Yisw4FbJimCD3s7LQHkiALNax3evqAJ2Z8nqxtEAAAEEQZrnSahBbJlMFExv1WRxLgVrlSU5kvRPQefQNL+hnY+5sxyM4OGP3W6gipxl+VeQcVtIs+PdwqJefHjTv2e+aZM41f823N7uJYxEIRmdHjyj4j8dVr5J5jvbLxXLbyf1aePF+es1ewhAgs3sr0LDYMZnAVaSTpMiiCjkwD85nRHwA0GmWfVDsiw1fCQ08yGxuwdcMd61SPpfRADYJpeVmWilLIkk9w7lCbACvWCJVPVxpYkJW8dIszr9QbdlVVkW+1uatXpBfwsbLty6mMFmnr5Fm2i6dwp+IR9w1Urt++4L+nMYz1LxgXCx3//Zxzgt2ANHgQF4NuaWK8l1q7TAxuxaM0EAAAA4AZ8GakR/3DTbiwrhhwSpjeaA68X7+pUBukHe3+XGyF7HhYlJIdhyOmcdef7GeSfs6o7iJKfivUUAAADJQZsJSeEKUmUwUsb/9BIkgoUyga5F+ZT3IdOhBODSowpHgJE6OOIA0hhuJSSkDCReDJDR0qEUS4FpJoQHvLAlVaAWTk//H9e/17IrpZwyJ6wCC9sEeRYyik1ullhgvUDdNUQzGaLfMevgU8An/KAX0X1/cOEMqpcK9aaJxowlmuUS21UtcIb9euUUnTFnUYwlhmrjdbVG+lNOTBID04MLbIMZIZzw6PX/9JpiKMQPPmV99BkjSWM89NuJhl5mlC5ED/6kt13QRfBgAAAAPwGfKGpEf5iZsaTq/dqTsanc/0IT3kf5Hr5m3oLGjyM1iBcL6c2qPe1QQjiri0diK6w7972kh1OttrOxLezWwAAAAJRBmytJ4Q6JlMFExv/0LBqptRAkFpHKdI1ZlaFSI2+VAhWAnIutcz+4N6enTxH30oYaYrWXeNE7ifjDCoHMMyW2gFUY+wbi/JE8adAKifO+RhxPHcvsbcbFsnmwVPaFpbwyxJYTd5QCBNnA5eacY9cU6ZwSL/KPzlL1bcR5waT+lbJg8F9k0Ygkey54qcWn8L2WvH0xAAAAIAGfSmpEf5BbOR6mhN42UttyvF0PAGvTXN8LUtVDAEJ9AAAAqkGbTknhDyZTAjf/x6oDJtkwPh8BHMKIJQeKakhBaoY1uTU1Dv+mLEdoR9c+RexgX+HTtr1kNcl3RElgObLDY/PaRJqf4jDnBLlQj1v/j+gBpH7ME2CB67cRZ0xo1T0OMscgDL2ja/1P3y/3R5yr255srgQtnSzKhY+hp0JxNKYyYJi1IFRA564tmvre91OlcSUCU/OIFZJE4oY32tHuv8tWU01eFm8j3ElGAAAAPkGfbEURPN/QgoY6wn8yMqFM8ScEy9Q41T0IvLskW7ZAQd6JL5LdsFFkc8iVrAOSANPjytFHsQEpJgFmIQehAAAAMQGfjWpEf4EHlYbRiT21xAdL3k+Vm4+8bP5+kpCxYbPyDxTyetngNkKsEo/OcSIoI20AAACGQZuSSahBaJlMCN/0LBreLNnVwoJ7HokTqcp6Q9iGVOTjfWRdX3HNLzYq0LH11Wq8Jq+uiZjebL3X4L63Z+y9aEj8POW4Tzlvo4K7tfR97PaNJboUExwnpuB7qJVSNwKDoTb48nTn0tJrvw/O3EZNUilImcvTCGE126axIMdq8DiFXn5Rl4EAAABdQZ+wRREsn309vOFLdicCm2x3Pufzb+QqT3EAn8TFTgB+03dALeaGOw1/6I6cxZ1UST7oYT8jPMi37QdASPs6q4fKsSNxFICCpLX/nIOtEt/OgZpH2dK2wf0xARgbAAAAMwGfz3REf4exAhsvjH6163w40hUbVYqr0IJX9KQl/IjfEvL74S4AfJknBl2wEwSJxGK44AAAAC4Bn9FqRH+Gdbgx/3bkh9Wq+haFwumUOl4Nq407OonaoLPLjUaPLWQnkuQZ8bOFAAAAPkGb00moQWyZTAjf9BLL/c2UM1KQX+BxGGL7bb/E+C2iJSq8ZT/fki7Xj3xHv79/3OCR4eUyro+hPAO6kwpQAAAAbkGb9UnhClJlMFFSxv/0LBqns2dXlsN5umPuNMnDIBekdIzGY3PnWipO8EQ4VO29m+HWcI89dez+iwGudwY8934b5mTOutNkupDrTYAefwpBLHDS4f6nH0oDndONGfIYoajXM6v5ZsonVMn+J7bcAAAAHgGeFGpEf4jgOwq2BPvOKJdhwabiU0ees+grZ1YMIQAAAJtBmhdJ4Q6JlMFExv/BdFX04QlYNpjTO0cfnbEhBFxg9ac2tx4t2ngcyym5CpUChqKx99i1D5EkvHHFPX9/7KMXq7jI/q3wkBGpGR/RHgt9ep8OxX3IASAi+yKKBVIw9GzZv1fxmZ6HBMZD2MbTpq0Mi3EG2QC+iK1rC7QF8ooxwG2N1JEtqo2tAd03sU98LQ9M68dpUpC/xpxxgAAAADQBnjZqRH+BD2jlniCi2imSX6//59WeyNYZ7EThdu+jrAY1bVw+gxFikWceh3l4wiMCVEWBAAAAY0GaOEnhDyZTAjf//WKRs+/QRQ3PXPg7GyNgHRyY2qAL7E/g8+2wAxx57bb9N5LU5Lwf1jU218is1FOXmWQFA+ILNlI2cuW4VuFEAUs55tZrFJbPFNXTxKQFb6zgrOGYxN6gYwAAAF5BmlpJ4Q8mUwURPG/9YpGyYu6xTlLEz1V6fELYXJYCl0J6HV2jidgU1kHU1TaIy98emaghsv30jfu//giS4u7YXeCrhV7hICMr6GVBT7WEUf+L/NNHBp/YziKg+rPOAAAAOQGeeWpEf4exHtbuDH+8TXGnxL+KjE1xi3rYQfEyAC01BI7R1nn8TQtbK9qWu/fdySehP03udb3CCQAAAFNBmn5J4Q8mUwI3//1ikbQD/zwawOGki1yA8PO6m26UL2L/9mKgReK7+ebK/al/M7IPAICEACwHoICsV2NTBXt4+cDul/OajszW51KtkLVl+aMGkAAAAERBnpxFETyffT2/FQo96LwXOn2o93ozcOD5tpYSJ7cwIJok8j0UU8aRtemPmdp7YYIUiCycIRpOWrK31Ohdggd67CoH2QAAACkBnrt0RH+HoDzW2qN3QBJ/E8gMm82gikxdFp7fAWLcEhRiEwpvi8PEkQAAACMBnr1qRH+BB5WDztAwU5oZn2uFb+hZJTgCS3z3O/PoqGYuVgAAAGpBmr9JqEFomUwI3/1ikbIV0EXMsa3YekHPSpYuprfZeuR3WOSiKkC0kFZpj7xpSGhq9ZWjmxEgbDlVL++iKztpTP/7VY798ET282ga/R1srrqOxD5u/U/+0pygtO04b7giKBp6WzCweDafAAAAekGawEnhClJlMCN//WKRs+/d5TeJbh49Ilqhs9vkFuDy5tILYlZvTBUtgGhCajr5vBfprgFsCD23kwfPBCzu+fBE9vNoGt6/fUDXQT6zOaThfR6pRtMs5SAtMZaMwqag6Y2HrhAcrJKrd7zb+2xmYXa8/28eupRQ2l2pAAAAZkGa5EnhDomUwI3//WKRsip9r0Ll4vNW49yxGm0rPx6d/W4e0Trq4nQjnHlJpUyssaFXC6OWq6h/nBwV81fzGtv5FLjjRbokIbdqcOylodWrIY6RdCrUEz4tfkGqVbBtksCTQEmc4gAAADhBnwJFETyffT2ygMiGUO/inwbX3T9EEiT+fzwLws9sywg0V8+Bo+hNRgRMgu2Mkb323JGVsThPeQAAACQBnyF0RH+AbCrDaMSfDibxDReFdC1UvC+7xMpiCX1dYFy6YMAAAABTAZ8jakR/h6Avf9Ts921mTsNzZj1+OJVQ8DVfaupr23+v6Ud85kP8cxjIPbdrvuP+lE9QsgrlQXtOD6M87SDXTNYuQnTLolxoGd6Iw9QfFrz+rYEAAABHQZsoSahBaJlMCN/KCDhlVhQaAeDWA+zilE4F4JvN0n7qpgAwW3vUlRQKCzVZ2fDY8WBCzFmBKgdcnVnxRMksKfRYMOIlGMMAAABUQZ9GRREsn309tLhK5CuvbdzT5vjjshRdA7IpVn/x5vBwKg++SIs8yQwRWo0QTWUmAKd3c5qJygjJ4vikTp3AKpczzE0vrvu10mq+f//pJzyqB/lDAAAANgGfZXREf4MOb59GU02DO+CVL8CSCBWNmo9IbuMSHjgvPw2O5k2vA2ympmv//pZzZUNT8y61kQAAACgBn2dqRH+BB5WDJUsHZv77UUMhfzqZBSad9zQgrlg0O2kNk5DK5ilQAAAAaEGbaUmoQWyZTAjfygg4ZVYMKd+rrGfciPhIiq2gCFAqMZFI7DVM2ZD3XLjJGnQmMgOIvEEsS8t+XEZfwTy76m0Qap5ds9FsZSQ6iIFZ0GU9cLjijfT2W6sRfprNkuPjiKDz/dClDTiAAAAAYkGbiknhClJlMCN//WKRs+1FSeYueQO7h1DZ7e+xqDAxtzmGG2fToXnjatvz1Tn7mSkpRsA/hIvPhdRN3kl1bvHvJTMKB9UA9ykYobd6r7+FKw/F3VVri6aWYgE36UGcq4wRAAAAZkGbq0nhDomUwI3/ygg4ZVYNdbNtBrA1dyGL1aw7/5eOTzHjgYmlvgs0LekeTrnBmIqn/lVENJBHf1RN5RqsI2gVGmV4sa2fCQ5tezOGmg2pMMQ4mlHRXTuRSMJg746BbfQilUiiLAAAAD5Bm8xJ4Q8mUwI3//1ikbPts0l1vp1ED2x1ji6bVEURYnvYIR3zetSp5m1e6jKMQTp7+VkiqKd2tavPg/v5wAAAAFRBm+1J4Q8mUwI3//1ikbHOi3JXDnst5+9/Q8YozwOIgFT629Pl7FB8I6iF0Zc7IF96heR57AZnekg18MeZ8yh2uOHSqcUTsaug9Sj4I8zr7wmNr9kAAABtQZoOSeEPJlMCN//9YpG0OvyOUdYK1lQfyPABv19dN2EuVGroHzp2zUVU0NODsPiVBimmNmF/CSLBt4F+dGOjvRT6LiulIFp5sR3+G1h1ATu9T+HuPV1x5zNOLgqLSCXQN/UuvfZ64plbVF+9zQAAAGJBmi9J4Q8mUwI3//1ikbHSvGoxINYD3SLvvyxPZxlxKc1XKrT03i5NYjyxJoTxK6gyI9UVV9jQM7CeNG8TnLBSp36B8rOa2OVfe31qsGMYmSKR4dZKBGwbKDN7WtkaEsZZZwAAAGxBmlBJ4Q8mUwI3//1ikbOuOtG+ZLt7BBIpdazWzzFa2oIXuKvsSfKEX099+kQxskSmsj0JIeAt4NsniL/81T0XG32CsxykbWR//z7P3ayX2T+O0bkNhfiOKnIVANdHE3DTcxt7oZfS/97sH3AAAABJQZpxSeEPJlMCN//9YpGyKf88uaZQCcSFhy2E1I13cQ7T4qAfZ/efBEseeX43YrqDE31Oe51xgWu3dMtxJ7PVwhztNF4CSSDjwAAAADdBmpJJ4Q8mUwI3//1ikbPts0lsarG2hQ9dcXCqWONWEzXuGsn18ByM8MdGGWrqt78s6xW7qPrhAAAAaEGas0nhDyZTAjf//WKRshXQOHBpVAT0KbRHowXnCt1oBqGLr3USERo5s9MW+PmLvdT3UmgBF+tZO4KwNstRUb0WbAoQxjGZF5h/wf/HrDjvpjl+CJ1AbSxs9y9+9mJLpBNXqaw9rkBgAAAAUUGa1EnhDyZTAjf/ygg4ZVYVaN1iGuf33zgSf6omrlwSndrkCTLPF3gaG94Se1OTYoSOIPiUGjc389Qa3QtMLmdfxJ/9EJgaN8vT7Keo1Vi14AAAAEFBmvVJ4Q8mUwI3//1ikbJi7rEg1gbbls4vToej136f50b8WMzjf3gU8vi9SoIuAXfe8+Kxcl5VytPUCNZR+0dQwQAAADxBmxZJ4Q8mUwI3//1ikbPtHlRWB+qqDdxH6C2zc77qhJ1fZ9qy55n/wDkZ1sej4ZDv6vAKtYQSdyHy4zwAAABWQZs3SeEPJlMCN//9YpGxzowPspkvEL9KenZlAqwx2krQK3l2oCa0wsWc8k2NlSHrH+b26/w5V1g8f/pnNfDFAjYmFGxsaSxgtrssgtpmXxb0wraVhN8AAABqQZtYSeEPJlMCN//9YpGzq87hxUgQYaLWrw1yOFa3/w8/9YwOlVMZthDmUgBo65bZZnRaIfs7t/B7omgkl6q/wSNF+R+xHNbXAqgPYfIrf0t/ExK3QwGLCO/7HBF9dcZigKukbsI6FQbxgQAAADFBm3lJ4Q8mUwJvFGew97V1Lcd+O+QDMDFw9wWzOYcwFZxnnmFDt6/+J6XbkhV2LCDKAAAE62WIggAJ/74Mt5iA31diPnODkHlAm45CKmrkO6tYSOdbdSitE8v4aantb94gCmriEDoY7sxccjIMexW5zP9owVi0pMMzzPblktB+I1kmpKIM33vbcRLOHPv0LC3kz7y85a64aEb5O+r1+heOtihVeJUPASxV9SaWTlNdsSrJgj4xmmjiYMl0IjoWJ9oVs4HKkTYxl/l9OJO8f5MIUmsmAuHXs9jIbJdKBA47Isns19FTdkpBDNqpHKr7q8GFdm/zMWCdfXqX41NkD8CuzgCgpicALz3yxhbW+N0bt/hFjnzm/UaPrgZPcZzqDiV5eYefi4KrgKA0ZVe3VpNSqYxcKn0235Wuxf+zZBgzqMmKLS0lLynJtcUkdDEeIyyfK4LIm48C4xhv8CoyP8GUDK8uivbXBNn7CgW7AuII2yToKdSV23wBkjD62p+pn+OELClzpVps9i3NFELoEbszOb28KZqNUr1lS8Wduuc2YwDKwxXVCCIeAzfkKMHwScs52kSwYmRCd/d8allKuQ2V4hIlH9Jl47UyFemjMfV7+rS2jrnUFGhJvwePtyEVo9P2WdscKW/cbYN2YTUkX7ujH/5UC056sGo7QCtikXJOlN9RqsmVVLXhNH901MjEGXfzKbWH5NYOmyN1mZ5lK/o17zObpovEPe3A6lXdH5FNGkd396zj8jXwQhCpcisxza7aEuADlEolAttpaS0aItC5D/12ya/vkaXhuhqDeowf/qnrMJEgsm2F4woahQkdykwDfkajV8Xoyi3jBpM/z3SsBgY6HWn3rW0ngUlpH53febs2YoTMxFBG9OvOZZ4WIWOcIb01o0qnK+4aqXMILB6bnNAae4qEqWz9+Z1pCpWendWC2cgOI5T5MSpidy6h+FiqCWdrWfJa2hloCwt6h+wVJKuhO2EUZjv5Bfuy8bQtvt4Ewrd0GT6CUsin3tuzheiKHIY7c7GARYxWOVsu7Y2YBaVScU2CERHIxqjA2+m2S99h81bJehej0iGHFz5U4SYskqcTAlUyrY13Nasq23dV54zIXVBj2HdPP/FgHKntLtIpTJhLgMLXrBf/HllTDNcthyCVB+wy+lmVyo5EGwdFVisnyRsii7TpXq3YGN3aKn3/gz7LAcKnAw5oD+vPQthR3tO2J/mgouNkoZ2GTWA2usO+gX2YiAf0fCodcAUuh6p+bh+zP/ArJYNnup5wOWspOUA0IovBT2AJTCeoNMfOAYnEEkLLP5P14x78y2wW1gWBgatjq9keNhveNxm3OId6Z/+RTUg7V/jZPZKnQJNBEy/eCv4qmUjDL8YpXPHrFThVUwNS3NoSGhVS4JZ7bpB03I1nSKSTH4DG1Lf5C7pla5OMYG3Y6n76Qbbfl9EZ86iKbiUcdGryD/jMHV12JykpDiaDb5FHPDkHH/vDAgnyicavYZBKgrx+j9OSTYucK815/uPHKsmteJ4DuUy00hL+bTACJ+49So3QPADbTTVdYUxcp68LjjJa/yyomh0a/d6qcsV/aT8s30oHmWbqIBloL+J1PNY1MzOCCU1U5ZMaoTRdCZwnZlEC/HchwPUMHg2DDEvmKsv4PqNfVSZC/GfvnrurcViZysYuuCjXGkUlYV0fNLrTBtq6ykyQkJUrZqjQdqw6kil1NFQWrOaWMeGubzMX3Bsgc3u8GFFquetJAAAAX0GaImxG//vteAAg9MZ+S0i/PC0zlHX/QQ3L0D4FiZmIOjpIFCJKKu+YprzBENnmxMhdyCs5hHE8xaJE2hp/2GNIqeNg/xvscT4V/xoCp2fXsc1CK8Yb8IojHKubQGZ8AAAAJQGeQXkR/4EHlYPO0mTyth87XGb0QiDlsCs7SOJZk1y96Bg/H+UAAACeQZpFPCGTKYRv++FR2IiwU0RA8qA1ChJxSNLn6LehKVP3kDzlM1ADrYhw7wd4Ip7AQzseLJdG5Wk6xUPebxn9TmVKfObBTfkJNs3ock1GokDdOFS6NZD0ZXPz//Ud2gir3gj4Dx1v/MhlLWLDnShmxzoy3/7AQNWKIUFk7mJHH+gLLt/06WNqUe+p7gC2E0AaYTuW0oNS7l5TDBLgt+EAAAAnQZ5jalPN/4H4Vo+TODPKsmBYbLE3RR0oNHHTburHIqYChXMkXhLgAAAAQAGehGpEf4exE2/rMWYbvveZW0hf2uqeKAaPNwi4iL5LnlJJ6paVEaSZbn+/5GomK7pT5XP9xVnXWAQw85PnL4AAAACCQZqJSahBaJlMCN/76z44S5pk3knFEhyI1Ta+1wxYdmynPZlwpDMB5imzzm+brVdgj4Ve8W9egCohBqHE8VTW+aug2DoRpkyvvsOPUbzL+H8XIBGreeSAS6btFT8tlLl2VaVtAxevDlGvVQTnr78zMiiWOQniAGGOB8j3OPC1oxhUPwAAAFdBnqdFESyffT2/dTSVQOypzfxQ2+DUil4GWJIgvIkBLArS+ydK0VasRPrGHtnNMEBi0Ye06NmDV2cc8yAn4FapQmg/tHUSVZv8nmxd+pKZ4n49OY6GrcAAAABBAZ7GdER/hRLK4pCufOnSwAxieT2ESQKanmN9zyXjppN4TTZfPjYr52MwkENEajixR/ruD0u5TZi8/TESJpQaWOcAAABSAZ7IakR/h7DkMUOCGBJh+lP5pZ7hsk9ED80Ikh1LLF5Y06UPw2TtMFuWrNVoJtNJFPcLihlZ9xJ1Td3KPSex82src3Nfa1Z7l65jvwwWZi0CQAAAAK5Bms1JqEFsmUwI3/wZX/CB9oOI26TZh6KulR6t3agG/iMaBA6B/6qMTvqGVg4k0YebbNb+LKTM12MrCqqoUI8hGiD2iXPrHlOmpgdC3Qb+yp7n1ju75DCZDB9Rh9M8MBS44veosH5GPzQfZ3g85NWC/wgf9Sj/wt2ZHm7uia3RMjx0cInK0e6gmf2UIdsg++PvXMcBn0gM8Se6JG8DKFx436aem9R9/OoLsInKDNEAAAA+QZ7rRRUsn39wfthx56FiN3xIMTuB500p7P9azlWaxHm/XjqkrOgES0XmZvka4Ctq1g1y9N9V22HlsJnZM+EAAABFAZ8KdER/h7D4ZJjIOE/G/Yh5t2iTFoHhadHs+YUVSRGmyvCLsMp2wNI81ubiOdKh6cKDbpl4cUw6CZHuVE7caJVrVL2BAAAANAGfDGpEf4k2H5XFXAtLnyqNr3TnYj9FIKhGPIf0deRjBKSCA4espm33EiqcHf4SHokHvLcAAAByQZsQSahBbJlMCN/74DjhBrqxUlYTYumv3ra8FUgL+AgXjYUzvRyU+e6HvKXgk8xL/pusoXU2ZW4/lqsHsMHeb8YKaReZyD2xnfEzDGr/au71JgRab6BdwH1e0PQ6ZYMc3yLQJm6faf8goYHDJ7uRU+gbAAAAUUGfLkUVLN+BznEOh5fTV8lzZdXNmbQtHmURs/S4z/04sX/c5yXQAjkgX0sC2qb0YByR+BiM7qb8PCLi4y3UBZ//3HtEURqKtuvuIfwOuQF+IAAAADgBn09qRH+BB5XFIVzyB9wJv/LfxiM8ksv3Sg0OqcE4Yrqo1JjCRCketlzK8503P3asAJA7OT5xlQAAAIZBm1RJqEFsmUwI3/vEgWa+ohqyK7peNimapDBP/rmBAmpQD6OSjgVVooLkr8s+Xn8exw7XhF9TA1dM1L0nWyyYDn0YYo5qCungDaPef0jmHbsa/NST2MzomuWP12pBzS/D8UW68cG/U7WtRkyLaQ2o6xwlezid/ALSuNnhZUQ8LKvoBIjMxAAAAHNBn3JFFSyffT3DpS3ZRuxv7Y2vybz+bP8qJsrnWB3uZuD6OPtDGdX3ghT2/RUlUSZnoKabDsiTKHthUxvZZN5kazH/ANj2RZ4dXRg3GoPqzpyXMF08NHSw9ukMXVXozck1ZVN27kSwitnpaQpcwBkrkivgAAAATgGfkXREf4egL3/U7PdtZa3GHl5V5zeF2LLWs0MbzXOZlb3kOt8MN0Y4o4Qo4DC8r0V9YQ7rqrIWW9ryH/Yg4OqxHG3QnithZLe466OdQQAAAEEBn5NqRH+Gdbgx/3Uy5jiTO2z6NGuYZhDosvmqOy+HJIrRZ/1j+mdlQeOSAXr9NMUAqXflcAvnLw0HhARRdO9cgQAAAGJBm5ZJqEFsmUwUTG/77lcOB3OO+4ld2XaGYPZe+y3Nf7PM2BZEj3/ak7IMgYJZ292AtkE/+fDvHwSjzWBw4ZWnozBz4y6QXsHgftapY5tl8Hx4uA9PzOWUKmjPAFe5E1HvcQAAAC4Bn7VqRH+BB5WDztAwdp20lm5s3YSJceebDQR3Or6z816nVxtIBRfgErPOT5y3AAAAe0Gbt0nhClJlMCN//XbgfUwOWc23U11yALjVVb2w/NZo+TO7IFE7jF6EPo4+YMYzRIKbgiu3F3kB85pl36VhUcvvytbvKOop5eSVkjG1CTrVgXP7WxMDh0fOf1+slEO+lIyatR6SsaZM3IzBZreQq2JodfFlKOfZwWWZLAAAAQJBm9tJ4Q6JlMCN/8FQss3+ZKWvz1Vvgwf6g9Ha3XKud4lLmF6JaAC5IPporCJUHxbKHU7vYH4HXDa3zFXkkAfldXiwDv9hY4EqQrh7laZJ6pLfk5IRDi10QhgrLa6I1cpiSQ45EGgE+Xxsrg4o4flhDnQ5Ihthw83MHZGB3B10kbD+Qt/BG0au6ZcfFXOkkOBPmx5s9/Pzof7lYV/0dpNlxpwp/tqKTFwyLX4rqRzgPi5GmFJ0mKT/M/Hz/TxmtqRDs4NYZoN35tdFAHpO0d0mGyRL2D3iMHWnkDF/LAzFtqzHaDqN77mbGjyug1vU8DKtwrcKfMEHvF/O9Vqueeym3YEAAABmQZ/5RRE8n9O9y35ttP5XA8yfAgYvyCUdxljqWjOdhas1mddNiDe/3VrNDcwN1etJ4bY18tX/hyWyP4GIwlHrtyez+ZCz5ghiUA/FV9anrOpGlw/LmUeoI015mUjorg3GPJNvCnPAAAAAPgGeGHREf4Bx3gMa3MtJzXKNV1X/TvXH0R7p0TV51BB27zemt1rU8h+OxY4zhjatnT2bUD2FU0s2nNDpkPAcAAAAOQGeGmpEf9wqXEWq5RiVMaSwAxlOO+23LmXEvXP/IGbVJ5FsJzif9EnJ3uZT0nLj1G4aGVRcuYlRjwAAALxBmh1JqEFomUwU8b9uJt+4JAB2PDGkUZKn1LVaQmhnA3PX4/yJdbhJ6+ZhTvgJKp1Mpd4XsRmIufzYB0Gjg5HjsymTtW09NUcla0xdquO7VsJlNt1J+fS07p/3CtM9jIS4PXn662Okc2oWKdm6qRq6L6au5t7Fmndyg+y3aiIWUUQe67w5+Tvad6VsoM2DsXseu9ACjX7KP9LZCac5WFt8E5lYKT+WyzrehW+bryfOIPcnHP5Q7WKqAMVWwAAAADQBnjxqRH+Yjrg02yKk3NJk1+zTlSEjhOv9WdIP223OJgbbUqmQtvMLft3pIcnpIH3fSrw9AAAAWkGaP0nhClJlMFLG//QsHcexAhSB/PMs/QCqkuzYg99CJH+9JGOv5aFabrl2r2ctLBY6M2x4Zbvd9hht9FyQ6133GQK/Ox+Ab7R/IiaxaiHq0mgOYL42zxXHdQAAAC4Bnl5qRH+HsRJVvmNK5tDyf6EJ/1EigZ555nVBvboO63s3YDIFnXPkddwMkLKQAAAAYkGaQUnhDomUwUTG//QSPy7vTRAErbpXCmQqCkFeba8BWdHebwLy/TbXEM1CAdbroo7rrsm2hvYzIsEAE/hNIfu+eyIxwlnCjcjAA2wohh/23XNHd2jo5ayVrDbcMuugD8lZAAAAHQGeYGpEf4eNe2TaTeL8zPwYEq4gvbPJGyoYjEgzAAAAeEGaZEnhDyZTAjf/++A44HeOPnbQKXKFpJcaRdAGqeg9I1987zdEm3R60S7BqCKKCkCiM2sDvOC+Lw0lYZ0ctcQkkNkRnYsQfrl2/a8fGMNco18ZFNi4twcv152XeDYkOPPOBVk/8oWKhB5mE61fKSUIalH3QzOSlQAAAEdBnoJFETzfg/vT37uAoiBAB62fY4T7UhTDnZgXNwAMuT9zHwCV7PJPIez8/S26iuBAOJekb2B0RKzV9AVe/JyD5c0nVj80UQAAACYBnqNqRH+BB5WG0Yk9tWoDpe8K6FsAkXKMufOCq5zTCtm/E5mXXQAAAJZBmqhJqEFomUwI38etDOA3Wak4jIfVLa19s6Y0Rw/msIXrJHoKvxhcLNOe8r7slL4Wa5KoaggYjByeswoCOZrcIbqHWjdx7a4x7OyGnR642XUuyo+PuqhJohJR0twDTn/jWlpUwQC7UFGO8pvAyih3nOBLZzUtnawIh01N4/zJmVIzs03ryQ/CsNNwGXnSogOB7nyq3zAAAABCQZ7GRREsn81cfmOuKtuOmYL5CdH7jVoF6ZsqPcI7do3JhGpCQH0VNgZvZikZzZ0PoPATUfW5vZTR6uB/pZC6qxM9AAAAMQGe5XREf57Mf84VVo08PAmIFtyz68/xZtlcGHAK7LcnvzsXqehJ/LT41oVmRW0o5mcAAAArAZ7nakR/hnW4Mf925RcTlmD+h0U3BboCgA52LwY4NlHSNS0Dhr7XMuQKgQAAAEpBmupJqEFsmUwUTG/KCEHK3fhHiuibLgfkSu7LpIIY3nGHB/7BUdXL8qXzP68vVCju+6mDLSp04UM3Yy5sMOop5Db5UDv6zCEmEAAAABcBnwlqRH+BB5WDztJXdtm1EV3ecxpV2QAAAGZBmwtJ4QpSZTAjf/vhUdhEluql+hMAizNOFJ5HtLg8fPT3ztEuLEr1XGD5xhtaNEj9qa8FFVnv7A21sXMqfiyoIXW19PEUwYPgxOYHv+WGMXik1z6DYg4KxCAi5t9mD/wLAUi9+oEAAABjQZssSeEOiZTAjf/74Djgegh1UZQX/WDsSXsFC7ltHlVSseL9p8ERavj1mdYkVohq6rg2aShXbm9fgcktbs7aTqfRR48AbzyAdovzp+bia6bsdbHEDLxLykgR6RTMRnNhvmvRAAAAV0GbTUnhDyZTAjf/+8PvG6OB3OPtuWzfxGhQy+SIsOLPOJtIbYUIIMS0V4qVXaqngqCi//dOlQY8F7ndPEOY7ybsB51X9jVAWXsMnnsbNepMd2SEvVMXwQAAAEtBm25J4Q8mUwI3//vEhm6OCJW+nXGeBIiklAPuU/0ZI1kWndzC0dHYpngvdDP+J144OxWFY4ccxDyLLIj6szhtQ7/kePrw8DROXeAAAABZQZuSSeEPJlMCN//7wcOsd4IFAwltfPFspkUwU+tL4SrplClPxXSDmSimJiQeXcXpRNlc4r6tw/bM9uY4UWJpP7B364iwyGkyu+lQkibJiggxD+c+9Bvsu3gAAABGQZ+wRRE8n309Vqjbpbriq8XvOdcqau8JX/6UTOUVpRY9WKhf83qV6gviCooy2oTUp7uQzw7CwPScfAIZVLOOefPzZ7OvoQAAADoBn890RH+HsQIbL4x/wcxICg+7jcceczhq0AWjcfMYx47NzTdIEGN7i0aSwUyq5LiHB0crT/GwU/SAAAAANwGf0WpEf4Z1uDH/duUXE5Zg4eut6v/g2rzmcUOWPfIVraaS0gPT3J2ns5evZ1CI1ydypxrgCnQAAAA4QZvTSahBaJlMCN/7w+8bo4Hc4/EYYvlEGmfvhZdsa87VJ6JltdXwa00H/ieFaI/Y7du5hEl+VYsAAAA8QZv0SeEKUmUwI3/74ByBNNkfUtOKnL3+oyVxxmsptyCPddx/B5KMUbKrDnuC2Vyx0n1eE4KlwYD5Vy7VAAAAZEGaFUnhDomUwI3//WN7G3P4NYxvBghU7kkcBVleyVOEARn3ED3ZepoZq5RhvWCC/yqm2nrsXEILV4frQDCdGWESAMuwcPI6wEXjGyfMZqSkgvdx7JIjxrwTiUqQdGnnY8X6eu8AAABoQZo4SeEPJlMCN//74Djgdzj87Yh9SyaadvlLvsYIaG789+bDzKp4nI7Ux9DSaBly7yJ7/BdbBlU59uboJHfQNPDqioaaRPekZqtjhMesy31GxBhA1OKS4Ps6VJGf+JLm3Aa3xC0tIAwAAAA5QZ5WRRE834P709++8xDSwYyJmXntNFlWK92Sf/fQDIC13sgr2eS1PqEUq8jnPixfgD+QQYI32zWlAAAANQGed2pEf4EHlcUhXPIH3Am/8sLtAuTjaqUxAE/k7/M900QRhPQgPEHPLe77P+AZmROrczDlAAAAbkGafEmoQWiZTAjf+/2woAB97w5AEvi75kBbZEZSk6i5LHQkcNyP+hSCGc7+2SDRFNE5x5z/5s+xx+xzk7p8O4w7+mh+W5do9ddvYMtDfUKDtyA+F85e07Sekz50aiZ+A2LJ2uLJoEQWDwl+okL4AAAAbUGemkURLJ99mioJbr7gM0pfNOj+bZFRc0eF3Hz740MUqo06KEIR1upjhqV/Gzz5amPTy3yzTR6/dNTQU+pkKnktDCZZSdmNwG8n5niOiMJYGfQ5IyXF8S4KMq7xjbHIWbNosUD2y1HHlqmypIAAAABLAZ65dER/inEMYh/V/63CvSgvZEXbfhyQx2wa0izGid6YTPbRTv0aQJw+H2w2yN2UnmHoAMN0UTFRR/aDBMgxksMmIswPXoXpwRBBAAAARwGeu2pEf4etAFVTbLs6LYf57fPpukYdfH0vO7n4NhsecsU4ANjZ8RyLB7ip5BIX0Mx0+7ZAp+laB0Jf+sCagXiPlxvuigHgAAAAnEGav0moQWyZTAjf/Xb7hr774ABr8i42lWBj0uR1Z9SPh+ggHjsNADRKfCe1SC2gCz1ZV2wfdi4gTAYitJI1nTCixMcZaP4yGm/n4CfqGc6sKSk0KSFZjxx71wspXvBNzxUtZdrOzeAHxomBMOPBhu1dzQNzjtmlggUsl7VLlU74bd3nV5k7rd2++U9+8ggj3YDa0RWl4HtEE/EBfwAAAC5Bnt1FFSzfgzrToKs1FKb08R33Z2J50qKbNlQV9+d+S4eldtRoddXTZwHI3pZ9AAAAJgGe/mpEf4jgOwq2BPvs/N5cLI/GjH7DHCev2QPFr9/a5wG4YVCQAAAAYUGa4EmoQWyZTAjfyghD0zehAPTP6iP8vI/K8ulMYZHKRgVkbItuqNJ3HEAokTCb6Pnt7yl0OLzZwaSoIONyTBCZHmwGWPv3H83nFo1lWoU+cguVFa9NsDp0gOBgSBwkMEEAAABNQZsBSeEKUmUwI3/7w+8bo4DvMfUVj6CQ3wk0ZhZRNy/Weym5JT2cFcNQkc517Tr/49Tt0A4xq3KD+uTRFAhPphx4G+TQtX/b6Q9KLoAAAABIQZsiSeEOiZTAjf/74ByBNa2ipZMyAB8rFbNzvqfypmxt6LH+u+4jWhgXuj+cO3YiKHJNJ1WQFBNgS8wGPIrNRLwTz121G6k9AAAAXkGbQ0nhDyZTAjf/+8Gbu08WB7kttaaGR98kbL2C0gBL0MGbGJD6y21jf9/9lA37OP49VLmwv/65IaqtSBP9i3lOAQ5XqdAIIYxDxwBb1v2HeMCQXhXOkOBAThSFTcEAAABkQZtkSeEPJlMCN//8AoIGyiRKThReSI30K30Lv43t8SJ+8KvIhuBxhxo/in9rIRgUOPXJvOP0aqCWU6/xuEs9JaJ8GMhw30DvHBYxT/U0Kj8r2ehB+UmTWWcR/WTYZVe6KZ5swAAAAFRBm4VJ4Q8mUwI3//wDwzHA7nHfcLHfflRrPXe74AVeC1IbYIGvxkydW+YZkCPntNjxM2qO/tU91qFpC4ImWGO5QaY+ISilW3oVGEXv9m5/sTC2qUAAAABQQZunSeEPJlMFETxv++A44IUcwuy4ITXX3nLrPA3Q//hXK5HZ7fz2fB0R5AIIhVciP1yw8k0h9Np/gyGu6xn4JoLPSopgQYiI15qUK7avBOMAAAA4AZ/GakR/gRIN2pOx7ohh37pmz7Klk5+7v+6dmyrL07Wcjao8Z7/mUQRNj6WbSSQGD97CVmEVJHQAAAAtQZvISeEPJlMCN//74ByBNNlS+GzqUsik57u8ykO74fuvsnTVRg/9Ptqbp+XBAAAAUEGb6UnhDyZTAjf/++FR2EKSC0P90owSMxGoCsS0jIjkJXAY9nBDNCe1Bk/slGqt0QC1TqM5EeT/mtDAwGnbXjDvOUDKJN+rzAy12AIq0PbWAAAAa0GaCknhDyZTAjf/+/6PPsocLdBktI9rDLAZxNlt9wqUjirtVNt+xqsPZxzK6F5vZZijVpT+u6KITqFhQSn9k+W8ILsVjI7zrB9WmSKDyuy8O409jwJKv6h/HdxfnOL2WIBLbqi82goKAz6BAAAATUGaK0nhDyZTAjf//APFuECSmz/LfleQbzi9RxtCMFA3aSBw7cfB5VdpUKlIOO9D5eaJ77NvYp1NRMnpOo8KCNP4Zz+5SwFgGPmDpyQ/AAAAOkGaTEnhDyZTAjf/++A44IldYymy5ae50jTmHGl+LJMPRkmIEvlkx6J3enH/Z1ZYnSaiWfS9Qd3bZWEAAABRQZpwSeEPJlMCN//7xIFmvqIeT19k4jZC5qgXt/92KDK0y5JmDGlXIo9RAWJ1sdV9OzVDj0ny3mjY57wjldIvC+C3BkGsD8nC+eIrXj01+wxBAAAAQUGejkURPJ+A6iW+6tW/s8hDZPRPj/ZS4R77YE3gu1EWsAgpWZFUkBfemZrEkJkVcTAOnvNlQOvZHJ6uIRbuaIJAAAAALwGerXREf4egL3/jdwY/RuCrxWt5oYNMoyXcgYWnulwzd87FVyGo2D8hOuzPX1rgAAAAMAGer2pEf4Z1uDH/deqwuVerjcUBHrhTnazCTYcm52J3Yzn0y9s0foT1etFDJ+eKgQAAAIdBmrNJqEFomUwI3/vua0IEqusHZHfP89a8nE1UwIsrlNvzfQ3IirE5WhcQKxUzthp7BUf5HkhK3zIaLvN0HYEJo/wGnymLnnzFyHPfS9TQcFpsTCptT+W+ftmIjvDK0LfLQWDU5ngMJHz/n4ZUiF0aa3S3RfkSsidDqKKs1LyafNz8fl1XPT4AAAAoQZ7RRREs34M606CdEYoROePxtpctRNT26oIcfF6VX0DqQPJvGMKc4AAAACUBnvJqRH+JNh+Vsm/raPip6RP9Og6Y/+RSUgkDHk9xc2O0jrwxAAAAnUGa90moQWyZTAjf+8Gbu2qQAjMVaMFUREfFFP+x4btt00ueekq5HGN89Wm19wbx3Kdw+o0PqxQpY0DcAXpv9K4n1LD86yV7ODCzYbxbeNispTBcvrKPCMwkXGwZIKam/tG7bdG9+lzB6J8XZvmU48gqeMwbRPYvIVNA81JXiwX5VfXEp+tOTePj9OoOWcrIF8LDW8sH5NteKrA41RcAAABMQZ8VRRUsn357kZSwzxr9U8YWP9/xl0oFI+aTRf6oWUqAXw298uTlr2goAldpwAKLQthKd2eV5xwZcFfuV6zka+DNcWzVwddT5dqn4QAAAEsBnzR0RH+HsPlbvvmSW/qDUBoMhf2usSxi8Wum4LDDW0Fu2KeEyWSegiZwWr0lVsQcu/F0Necrzgb7KR5QDSLx9yh9ScE0/rzHaJwAAAA7AZ82akR/hNfANpljEqY1fAK9wQesw2QAGiR2LAytbhpI31Aax1ACoVdfHniU9EpvumT6lsdtl9FTtjAAAACFQZs5SahBbJlMFExv/APDMcEKP+df8b2bcbvMcHhLvzdzipVTXYQpeF6xqpUhj4bNawZXtZlSHo26T+WH3zqaIllB37/NGNQ1VkY8BOUKE+zpXj8ANW28eEEEdJdDLN1PiFpb5PIMKKW9o9IlM9LDe1K216q6Y7zZmIMkCDaMmuxIIuprEQAAAEwBn1hqRH+HsRJVgFi+ZrlLDr4CkfmPNcriMt98U9nbw18dR45otYG8DOLKOeaxiW7U396pvPAr+t6Gn8NLGkSaktuhYMuySbpXKWCAAAAAp0GbXUnhClJlMCN/+8FGR0k7BwEOlhvznFDc8MKd6PlOwJwx+IkDWajbiHtKLtq12J317Zrrh1gtLttHBFn7OwW9kLbr59zfz9WSxFbcUiYfBU+4PJI46FozlfEF8jFGF4/hW1jigr3Oe+UyN6JQ7klFrYmpLuv1jvLAjYdokw6od71OZf5GD4K1jtQiVwCSoCZlTnckIlQyvMjv/ZinSucoeLKUdjlgAAAAREGfe0U0TJ97IKea95i5I4Oz7pxueJfZG4kOmA+QbzXKw7vkhvqrsiqjd/kL4XyY3lQf0TmOeb7Jxwp736HbF34x3jaBAAAANwGfmnREf4ew+GS+7XjuUTR+buTP+d/8H1HLLwuG6VVjsCWSzm1R7OCK0QXFWRrnznHKxhkYoPwAAAAuAZ+cakR/iOA7Criq74tIQStcVo0TUW87D3aTI7cUwCOn+Wxs+VAcJAD2/FWI1wAAAHBBm4BJqEFomUwI3/vgOOB3OPtuWWgGLWHjA2dK83VlKvt2+n8reBM2XzHD3zBxBIqIRM/rFRT7vg6GzXbZYVE4qWcVvmjn8+cOQkH82Z90ruTl7Fjpv4YW3KK5nKkBzEFtS3TO1V0jlDTMhHqvo4khAAAAVkGfvkURLN+D+9OfdQR6q/7IXDaBvaDpMhPEswOM1+7l4S2FsorT/pLAVOjCdAJLasClJUflp/50NXl8VQKz+m+7rpgfYkqz2G+vKYeS9jaTSkokmfCAAAAAKQGf32pEf4EHlcUhXPIuvJ32qeFdoZAknNUtwBScyRRJETefzNQclkJBAAAAcUGbxEmoQWyZTAjf+/2woAB9YANHst9xAn8Ot4dHzqUwafrp05d0vUqpjWj5to+oP2Bvpa+fUTyobSQwuvbyPiaW5y0mT1/O4V2M/c/cW/9kBex4FuohVoCvWoVQg/16KJ9s7jM9tnwTb/JeLWapEhPoAAAATkGf4kUVLJ99mioJWnVwU28WW1/7jNoqYi8ej+b1wieEOmo+V6ElGEDFkN0d6FE9/r6vZJXwBtcGjNFa7FbJe4U0C3VVH6u9ab3qqXMOQQAAADsBngF0RH+HsQIbL4x/uLfAQThciDnXlgeXNkbuTWOhTs0T8clBXVZVPYHhl/rxl3hE9I+Hrbh7KVpysQAAADEBngNqRH+Gdbgx/3Xq50pXw8mqoa5/S6osvnea87hsE9/0OGaERxf1tCTkMhfZ7ECEAAAAY0GaBkmoQWyZTBRMb8oIQcmggO5x33CtfQoNIePFJ40j0Rw5ce/FyFAT5ISrtO7pr4/ZWcfWhBygM64QbSUQzME+lubt8WOWt0M6f2A3BfB70bi4jXwkpKQm6O/cjjNNV79bYAAAAEYBniVqRH+BB5WDzvLCHMeSc6D+pWTKrSTRgYiHkmyaD8ady4xoFuAGAC4gZprGG/guu9ew6u+KEQtPsf7XkTxEL5oOibpxAAAAfUGaJ0nhClJlMCN/+8EZPlhBHWC1fwUqYD4gl8M84OnkMvjhAR1eA25Xfq3KZ7HDMrdVm+l24vSAkfG3ha5Og11+q/XKHpFWzSPOORK4/q+VjpIgK+wW3rohvnrLXSLrg00iFleonFUn2Zm8WiDb62AKQG8rPQI2jk4tFOzPAAAAd0GaS0nhDomUwI3/yghA8yNyXL4QryWtpO9TBWG+5lvHFT8rZrJwigk2mHXYghgKHt37c17joLyvASuy11M5xsDnIGswur45lcp/pg6KTMv64Bm5l7g8OAtRTsOC/qs/q8KCOxtRlop8JGxN+wVyJAecO+x+WH2JAAAATUGeaUURPJ99biHttP5WYp4RAgYvyI+p8tb+TRf6oWaAAXw4HkfPwn3YQah/Hn50Dc0k3Kbwtlf++vgNgAn442VMlPHd1szsO3j9hrRgAAAAQQGeiHREf4ew+Vu++ZJcAjosbnx/+SKnisgJBzGvhE2wYk44wff7Oj+WUOlh1tx4rNtkmx7x72zfz/ISHcplyHXRAAAAOAGeimpEf4TUK0MsYlTGr4BXuCEvYQ8yfXuACUj2t8k/d+aCycsFhAIUjm9ScUV3x2YTx3NuXIgZAAAAdkGajUmoQWiZTBTxv/vD2xXRwQo/5tqyIg0oegdJzRBsuIl+7b7ib975tSk5oO+RagBoP94bo8xTAUoC6UAV7hD2pl/WnNEYhjeekT1PkCFVNeqc3AZwsou8tNpbY/mzbsB2O5/LkQ+UJINnK+GSH8HOdUvNWTkAAABKAZ6sakR/h7ESVYBYvct3ZUH80aLSL3lILK48NjH93k+JR1JE74PB6N+HZWuHDRS04N7OCEWv96qoQw6Qs01sUJTT++6b2XHrGL0AAACqQZqxSeEKUmUwI3/7wetWHYdbgBdtDzUw0EHUbfim0S92FkIfRPj16pWKKLwRo3AU8FDsa9QId76SyIZbrKSW1k3GGqM46d3clPOrcjz5FiWx32mC+NDau6Yn5iLvV/ruKE+CYviocKuN1PAgXeEie0yP1W3y7ir7QTZDPhf3ZL2j64xZerSGI+I/vp60I4Yfhme/LvW+sG6aWg7iglDM6mb9UFBSV9g/opoAAABZQZ7PRTRMn37oQ9jXIx1lBX9mqfhAUp5p6J6k829lJWvVzf6M7ggijphbu0clU5M1vBwhyYZzt/rM9/6ylFgz0BPTGJMWpIznDxoA5gzNXt6b4LnZ35PBfoAAAABGAZ7udER/gHIYyAVJ11hb+xYFCFd58j6mfUxiy60D2axlSlNv4jbTnYFRvT61K/ZgtX21MlfVI4/OkXJuUpop7f64MnECrQAAACkBnvBqRH+I4DsKsOxDYDk/6PHIMsNIpxKsIM+ybhNlwmFnio9k9isHXAAAALZBmvVJqEFomUwI3/vBm7tPFg7we8ueSW3AsmgN1pq6VL5mONAom5p4474jhWe4CPB0/d75Ta+3AbyIfJutqpUJfL9iHeRoITsbO2cFPftnBnbG+U99/U08f+CG9I+wfsE58CY8Lf3D3w/MItP7ygZt3mLhmC8NDeSwaV013CnmXqHO8zdoI9keTJr4VJTM94XD+qoMrzWDNiUadSEvxRjUAUo8/J8dVa+AVsL8l9fEVbn49GcOQAAAAFNBnxNFESyffW4h7bT+VmzHLoEDF+QSPtyPFrba/50Rn+HndNiDb5+E/AuX0YogGD8wr8BoMCMYsRM5UoXrME0WskDjNcWx93MwmTeYntQpF1DOgQAAAEABnzJ0RH+HsPlaZzJzVMKscA/2L+9TyyQd/3O+nm90ZXl8Ksqn5Id4lYTPtFj0zEdxLJQd6obfw58aUu4rDzZRAAAANwGfNGpEf4TUKzbRXPm8OWGDdlZhb7MPjSB0cLQgO5R2NblH67E1i0ao4Dkk3g/masUQXUJUJOEAAABnQZs3SahBbJlMFExv/APDMcJc0yucaeHB6PvL/u479qPe7hcn8oUhAYtxr7RP+9d6EP7vADBZs2ne/afeqbqNMqONZ5d2u+aLATtq5FH7fBN2ApzqD2YHoBAbx7AylRx7VOtEtRa5YAAAADgBn1ZqRH+HsOQxQ4IYEJ9R380tftck9EBbFpYrpDKVIDssfLu6OZD0KIvfnYuj1dQTryD5+83wVwAAAMNBm1tJ4QpSZTAjf/vhUdhyzm1U2leLW9Qg70H6FjRAKOWGmq9x5x025+Fv3p3S/FQm5Tp0onIj4h8Zx4V1izZBK7CKWM3Sg4BQE8SIJ4khFfeHsmBs3BWM4yet7zvBdepBB46hzuc6TbCM87USRY/aAltVYBThJ2Dy1Kk+X7oELhrfv+rE7sa2+kv58rOX6d1pHhsn+G4eGGB5RP9NO1f9TZuBI2+uPWtPojfHGH2/8Aq9EEUM/xWKbPACAkCDEeGpfVcAAAA+QZ95RTRMn3sgp5r3mdEHB2fdKSdICJFhyUw9U6B6nMiDQFTH10Mf+Nz59OffzcVpUjcUpLvcP2XrpMgmGzwAAAA7AZ+YdER/h7D4ZJjHpgBMmQmFWTxuFtuKN+1nWmrvVP+cuwy6cfDaxn9EPm//fbSX/BC1BXlKAR4dXRgAAAAyAZ+aakR/iTYflcVbbqsjzADeWk/N4gygRKahtBY4GTk21u4qAH9//uy/xUzjKGN4gOEAAABkQZufSahBaJlMCN/7wZu7TxYO7yWtpO9TBWejmOuTienHcMDCHKUGeJgU8kvIN+p4Q3OK88HSnfDGE1MdkaQi+aoz5+X/+HFemTy4gYlIXyYWnSJDJo/cOHmT5qotrLEMdESHSAAAAE1Bn71FESyffW4h9qH5W+/h2dk4U8g5Pt7nJndunw3PPWIDrz9uaZyKwd46Sk4PBi2OZZ+81wbZ8oc5SG+A22lt4AZz/jYoLj5rNtjCgQAAADsBn9x0RH+AcOUFoNKCQYCGgc//6BZF/1e++Y5yM2PoTn2oKho+FzY+MYnRSCxwYlPnGESpnGT2xIEfQQAAADABn95qRH+E1Cs20Vz50jhSs7iL2EQrs9j3tB+E0OgBAtIl74/4sMyLjOTnhw5AY8AAAAEBQZvBSahBbJlMFExv08O2aXJzY9aWPq93Hlrmqv5B0QtXOYZZloiKjehd74nU1AKH52Z3+TFO92LMM98MEm40V7V+C1GenupfeqytRxcuUcO6u7TqCPTeIwqzTSEBSPq9aYeksIflaFksME00+qb4RPN3k9VCAuDXIDcWlIfhtRfV+FZNJvsVXv9aSTadgIZXTCel8E+da9dln1o4EHbsCWRxMU227KCu+aooSEKKHlx/66IpLeADW25hj/AgX0h2GHLGpU+s9OsBO2nNNGGqVU2s+i6OVvVpOmhdZkIXLFAuaV6rywH//SZMiw2+TxWyRZ7vy1v8JUsXVPcmUxHw0+EAAAA+AZ/gakR/18yldm4EIIwJMqwy/zSlRwwRYyU0uDbf23QW+RyODJVXr6oxUFynmeXdw9rHk4uHPi3FC8lMF5AAAACwQZvjSeEKUmUwUsb/9BI9qBbAgvoKrhEVvOFiUYViG/M5EJ8mdYUiEaj97nKw+nocaPSSxXVNZ1shjf+5GKqK3d7aQA54mLlWRII10yz7m8Bep30K/8/RHcJBI+3YKy6Dnp4kI0bb5XEQ/cK7bxUfpMy7cMapjRKOoRLMTOm5NDf8KQfAqQ5TeInNVpTmGv///kutmRd+MQ4u5qz5Vq8W/uVrDIc4XyeP9023KOhDiYMAAAA6AZ4CakR/mJmxpOr92pSK0PJ/oQn/bH9a0f7ds14XLC9W6fZ4EbVHvaoG3QUOZmMZfHun6TyTV7Kr/wAAAH9BmgVJ4Q6JlMFExv/74VHYcuN4xrugJWIKSCyEVY45OXNWVZpnwYhiQRi8ouWZeYuX9UrcKkoRdSfCZZoOmnAnOQp7HCaBaZsLqFrDoWpfb/ojyr+kPVjjyThmNhOz59P+kfdZc/N9o20arxf55nhcJHGZLjukhPirl/LdFsHrAAAAGQGeJGpEf4jgOwqw7EQkqKLsnOa2jmmZovYAAAChQZooSeEPJlMCN//Hr5DS9RRAa54Z7hoy0gQSdNxHb8oFhjXsrpRKZMbMTTdkBvu3TAncUA10EqZ7+4nEteFrv70v8BmFG1wgiHERfC31S38TWliBAZGrZPRYU78tMXq6ichNMahFbFt7DygIkxM0pvU/Vl+YLZqN/y1MDAS6SE19XdtA3PtOjnSNlsS0SqWCQzXThAQMqHaQZ2uGXkv+3sEAAAA7QZ5GRRE839BbCV26Byo0weWSA9YDX3+QMejMyEEDStqDfwdOOZPoewjLmsgoJTYBfGhx9VUyx/y7bNAAAAAgAZ5nakR/gQeVxSFc8gfcCb/ywu0C5M8eQRzzLW5tT1UAAACEQZpsSahBaJlMCN/7/bCgARAV0d4PVu2DIDo//3hRFBNpnloEynMkS6IgSJm98XQILX4peWuq1XgtmjZvcfov+so/e8U3LVg0FPoU/Mo4SQs54IeRwOVrQ4GMDhGOj3J+7LW8cm8BcglvzqyhewVZdBEH6P6sGuFKFltAUejrF5DmiW+AAAAAVEGeikURLJ99mioJb39gpt5ZOf/5t8yMdqhflAjF3VSF9LKRKHSP+6fBR0zNoOBXiwoFPSpm2LefrXJYQBXANQoIKzterZ//Rd7bAfQwQE2Okj3HIQAAAEIBnql0RH+HsQIbL4x/uLfAQTkb4yILaCZ3iPJCU9pVv5zI9a2kfiP3oWWAYeh919yuVlfR5mK81H/GkOepytZyXlMAAAAtAZ6rakR/hnW4Mf91Mix4JlquVA684PFShL+fqcs8dpqe6XtLIH5N1f8BA1aBAAAAOEGarUmoQWyZTAjf++5XDgdzjvuJXdlt/kTGF+RV/LdKMM3O2Izx8kplgk0r1rh9cxk0ue6JeItZAAAATEGaz0nhClJlMFFSxv/74VHYRK6xsfYetUOdE5zQVKQV9OZlZwDb2XeIsRYByM1q78Ie7kjTv+lVLWwES1iMaYWnsFBfK73puU6X9u4AAAAYAZ7uakR/iTYflbJxD5R0UYmRI1Tlb8yAAAAAWkGa8UnhDomUwUTG//vgOOB3OPztiP2UA+261TZeWO+pLr6FuxY1Gm4KJTeKb22KE1CtKd2PE/YlwvYroPfLhJg8ZrWdDVfKHDYJ8TeXliWDg+DDMHa6evgdfQAAADQBnxBqRH+BB400xy7oyzZN7DiQhYX6mHtDc+034OuRbz0duExV4R2w0AHsAUlmmJIw06XAAAAAVUGbEknhDyZTAjf/++A44Ilb6dcaU5H0bjlOnxLIo0SXrQTWgueivNvHPE/8vDiGPA/ersyEsLWjNH8iFcw0ZL+tkT9WxPFLyHI0aZt9JkWOOsnrayAAAABgQZs0SeEPJlMFETxv/AOiKyBDZVHR4Nb12r2vQbwl1ZUGpuF2dtyRYmAgNHFu9YI8DTzjaZZP2/AhZiSYSEOKyyEGzaeRtjr6LBzxI4L9G5pO+D4VbaY3srs2TWCx4mqXAAAANQGfU2pEf4pxDGIcj/67dNk+uMhrkyCin4ep171/SmuXlGTPbQXRJPBnppBZpXl2UXy4JS9RAAAAU0GbWEnhDyZTAjf//ALQu8EXYDiMMWrNiM6GWH7uYInAl2MtZ/kgekQpPORIiGhZ/nRumtGaN2e05a4CsUGjPa7Jt1HD/aPBVyqAiWxtSFal//ERAAAAOEGfdkURPJ99K3q8cj1odDChudhfcpeV0qEp95cV8xPVLYNNjT/96mXg6NNq46gQvUi2kw+4owVAAAAALwGflXREf4MOYBiW8DO/bTvzhcPmLoHmc9yW0C7QrF99GHjG17qqQEM6OaQ7qATwAAAANgGfl2pEf4EHlVcd04Is9ClW6FZCH4nYyUbEMAwSiUszmoA/R2IMPcZs3ej3LlzkBx5akMr+kQAAAJxBm5lJqEFomUwI3/1jexiPPuVF4N3+NUAUuIHEz2QO9Qjed+3c9KFSmMPrydJHG31a3UXHzy+f5ILebe40628nHuliEo6jFwOTxVv1UhoKiyPtASysBMxVAoFRxcO10B0l4c3iX8DbGNOa6io3rXjVXTO0z3QG2JpqNAWBXcHV1gHdDwEtkGhKhQAfl2adu3mdyJam0GGvKzmODyAAAABTQZu6SeEKUmUwI3/9JoP7PlJhERd+nNBdwdAGicV2sDdshxtss0WwtsFE9iV+OAFa0FSYHSTT81z38H7GxqDvlLWAdJmzfoGzta6+3HuOXIhSfDAAAACGQZveSeEOiZTAjf/9Y3sbYq0egyLMuoFUSS1zW3tAR1EhxIsXQLfJqvT5msgJUJ5lF6n+eX5oaHM3FvUq7s0ubwvX+FHL351BlWvMj+p0hEQccUBBuAKEAtp/FcKbFnScMfZEX2zhg5cgAhSWHPmnBc4H3/GBfKQU8pvzMeWNgvr9aUU5AVUAAAAzQZ/8RRE8n3wHSlhnjX6wuuQMkOxSt1zZ9hoUMwDDnzb488UvfimL2oJcGWgtG3RGIynAAAAAJAGeG3REf4BsKtkMVY91KCkU84i2Lj9E6hfd4mZLJmupc44LzQAAAEQBnh1qRH+HoC9/1Oz3bWZ3RPxWyj+fW0/FzMZb82m6SxXIuxzXGz2O9ueW8NrCaIuPQmb+HlSCvaDrTOCW1gf+t9oHuQAAAEhBmgJJqEFomUwI3/vrPjgIGEb1mbkKmaEb7B3LpV9EsRLeONLVcFxuuED+lT3/pAzfeZLEw2Kb38X2btrExLs1IQ4mp1/hywIAAAA8QZ4gRREsn309tLhK5Et0327KxGXKAtwcI4JfdoKiSSv6tpnwJ1fs0HcHDvesezzWNAerFq04zWyFWRqVAAAANgGeX3REf4MOYBiW8DPHRkut1B2MvM8gxWoMkGJG6fJNRKcDcXEbvKD2g6pI1wtMG2qZlsUlgAAAACgBnkFqRH+BB5WDztAx3ehsFB1Y2ENPbEBth3hGfB9PEXgdSLcgUejBAAAAbUGaQ0moQWyZTAjf/XbVbdFBvxTAPcrrcz6VZOCFULaB6H0iT9bpVOi048iO1wulW4BtFnQgnXqt/dwAcB0kgppQux9PGL3+oIrcHf/88oLvTn3UeHSiGxY0fUoo0QSzEULynDOOcZGpy8/0v00AAABzQZpkSeEKUmUwI3/KCEPTN6EA9ymLrly8O6B7t7VmUeTH/G0a7t2ifM3r1iMMfcBrASVNUPovKG8GYNYNCbz/ppRUblstrHKIbzRyZab8jak9fpDfmpy/m5pgm0ZFpIjmcCLOk2J7HyFBwjEKZneHITdafAAAAFpBmoVJ4Q6JlMCN/8qLkmeQgEPSPuqCmr5ArtbnLRl4fTGGkErnT9fA9ShDSntjr7GEjAb8W3TZIJnx5IZJCQVx1Awj6XTCsn7hJoDgdkQML672DCqIiq+WioAAAABHQZqmSeEPJlMCN//74ByArVvp1EDVN71GG530l/KjHXzb02zLOHiQXYfX9VEBccGpAGmsDLsvD7H889ZLz9QGv3FQvckk+bkAAABkQZrHSeEPJlMCN//7wZwidiwPbbYagjwfDcPSQg33o5EFwjljEQKjlBmdpaMBOqA7OQf5sN+AJ3xRw8XQzvvEfWu7eb+JbmKBNgl+YwIiKZxi6zn0Whj9DxvtFKFYv/L4uXFg8AAAAGNBmuhJ4Q8mUwI3//wCgK7wSDWMAKdUGMUb645fTdwYnhpkNPya+L8DWZnT0HAzto+auQYGWoyZsX3gYPKgXfkFrIfwZz5ku2kqON3Et4mdMuvJcP/wIlB3jHcb0QXsltP714EAAABiQZsJSeEPJlMCN//77lcOA7zH4jDFrjpxOzOGUd5ZawgfJTi+lCyrDkvXFWhsE8XF0nUH+68r3FUz/xnMGyUBef43CwO4Fu7mkzlFHt7U2RY0wKM3AWV7ncpWtZdHMWbpKuAAAAByQZsqSeEPJlMCN//77lcOA9yrrlSWZGlkSvhsHkUA+TYWmiPcd9erKEll/VhIYwxpcKy2emhfmAiAQ3BXo41lpahCQkJa5wBO43r2oRyWCTe4Hr4f5iHFvgO/n9/U/Wv3cJr0kF/5/7WMzkvgsE0ESqtXAAAARkGbS0nhDyZTAjf/++s+OCFI9lK+huo+FMo46iSF2IPShKnDhY+b4Nyyfn422/qBUX1fTLaA9u3h/ChZF+GVwRdZ2WRkBeEAAAA6QZtsSeEPJlMCN//74DjgQ7ZHzgzntGfTAsTgNjJN+UQhfBz3E1P/plckBND9E86yk083GNZZLnqs4QAAAFxBm41J4Q8mUwI3//vhUdhCXYakQCl1mC9vjlAE6W6K+YGIyz//JA3BetDKXvDl3VvDxmOvEBMODpgm+jCKU2Gl8xC9oBzxxspKYcHrqYH05l+P8v6Y5Qk/1pPAwQAAAE9Bm65J4Q8mUwI3//v+jfjges12kf9dZQFJRGONVgWYOoxzG5uK3Jxnwr2M+6smx4PBb//XwNxXDNlrK0yKc8dfQ3w+72OG2IyxsTTpWG8fAAAAQUGbz0nhDyZTAjf//APDMcDucfnbEfvHS6zrNEJPWVHkE15RWPV3XeHelWYNVUgciz2fa0ECdV6bGtEIrhn9uo1vAAAAP0Gb8EnhDyZTAjf/yghC9V+DzI7WDVxvgaNgcVv0j50t1Fi5NV8PQP4/2jlF1RxSfv6bfmVSxmd7SAf+mKajzQAAAFhBmhFJ4Q8mUwI3//vBm7tPFg9mlCdoqUDYZ8xiLA/58i+MrJi/Oj9iY90OsO12pTviJqIUVx7vxap1f+uydZk4zez/EtYp/4fYp4Q3XO1R4juDaEeRsAQQAAAAUkGaMknhDyZTAjf/++uE4Ai5s7YnlAoVYdUxHA6lgf+GQxBdkuL/GOcCv2WRMFlJ071NInuH+zP8ObDVzOx5D/yBFdXBjGv5oMdw9Mw36f6YbxAAAABYQZpWSeEPJlMCN//KCEPdL34Qpx+Iwxas1iPqZ7d7C263bvdkqW2ukMa2nT+82TtnjrJ/eEhfgMZHvbgwHN3G1mSeyf5ERFLmZN+rihWcq91jTRoFRIgx2QAAAEpBnnRFETyffZoqE5BKXkkXp2UOeJEo59JeY9aj9AeTC6IYWjavARsEhBIEdBw2Rve/kWbUHHscnz8S/B1NjHiIhOhIZBxTgdtyQQAAACcBnpN0RH+DDmAYl6IayBmmZNEDqZSB4/1mlPS3tDlsV/Fh1LZ9b8EAAAAoAZ6VakR/gQeVg87QeYm3bURZpmDefiWKuuiEiUDlMyxeXtBIWKCTYAAAAHRBmplJqEFomUwI3/vhUdicsgfnibHpetEFWH90rguJ3ElBMQTQH7c3ygYCf/Drrg/dtGQkavBREa3bvVMAltQMkijFYWTfqWL//8Egk401tHl0A5hfv/3NvhBuZ/FxLMknFPgLLzedIjC8QeurMrgT7aB7fgAAAC1BnrdFESzfgfhWj8hRbl2W307v17sFwmYW+jS8UpzFjs/eVgK5cyRdyT/WT7kAAABGAZ7YakR/gQ9o5bvE/zRjV0ev/5aKcO19Luh0B14WQ6ZPUsDRRYb+qb/GrGriVtRBld0p3TqxJM3rCnvICZvAAlZ5yfOX4AAAAIZBmt1JqEFsmUwI3/vrP1EjJ9ApoRLZwxrgZM9KFuXGgqknu1QvZh+YXzmCVkYA2kXu5XfTPYyy8+Cwf9E3yexwGKwVVsHQkC08DCU5fhqZfK1z8qrqNUDzTK72NXRlKxrFR1ZmiVPj2Xq4At3GtTri/A05eWs/ccw0A8+MLhzoW7U93GW4fAAAAFZBnvtFFSyffT2/dTSVQUKj9rf+0mUMU9bY6PcvN6uG7mXsKRuTE+sYe2c0wQGLRh7SjC+dfY/ykuTFxQBC9+JuGvBsR4DhRM1QspMMQ0PRXYOYT3CwOQAAAEABnxp0RH+FEsrDaMSpa8wIWWIvYRFuurkj19cl46aZDCbA0SBpIQmv8EkNyrj0XwVlw8L1phD0sXTHNQHFOHuAAAAAUgGfHGpEf4ew5DFDghAjpjpN8kQrknsoqluAN7r/wdMHwzFeUfhsnaYLctWarPhpJ1pPTcHcQt1NddNzpv6yDKOANWXO32tWe5euY78MFmYtAkEAAAC/QZsBSahBbJlMCN/8GV/wgjheSMGlboEggPQ8iUHEclFSu/AtGO3I3Vk2ER/4BHv8Tbue6BAabqGiQdgTzc4oRsQ1gh9ZCu/VbhiFCPYnJiHdwoX4YU+hRuenZpNtsObSdPTabZeCoJiif6l0xR2ZC1r2W4UwfkY/Kbu9X8RRqWwjiZ+lCEe2yO97qkKitGLRXhsDDPy6A42LbEzFS5m3Bfj9FzHAZ9IDPEUXHyZS+YwQIKAdldUuWh2n8/PuDNEAAABCQZ8/RRUsn3sgp5r3Cm/SDs+6O3nB5JhyYbLrIl3aW8UcgJDENn36FSTX+WiqBPcaJGrhj1oT+XMgOFmQ+3rs4Yq4AAAAQwGfXnREf4ew+GSYx8peZo32mD0ZxX3r90gQmxgMzNTtzm1R7ItzavQ1ubiOdKh6cKDbpl9nH+Ni2s3puThONblt6/kAAAA0AZ9AakR/iOA7CrirIwsXZRWN3XV5/opBUIx5D+UhAdCUpwQTa9Dm77iIFrwPR5D0Wos1uAAAAHZBm0RJqEFsmUwI3/vhUdhEC5b/sF1SYFY16qcgc4pX3SjdOcxNGiWTn6CYPHzCyrtX74+6hxGZS/gSlPlQe04vlBzgd9y4b9wSmNmTbl7qHsmH/8WHyl3DolsoVYqYdHMsqsE7ZL7tq75kchJTZcF0jg7A6sSBAAAATkGfYkUVLN99TaxVytK9J0LEC7e9QMMUfwZ+bYFBPy6qi3f3mYx5+pfBsels2wDlgtcp/zu436EhZrd06L3ibH1Yqz/2X7weI/SnSNFnaQAAADsBn4NqRH+BB5XFIVzyCm2kI+3PwfLld5Q/dKDRNU+AxchZjpfprwo6gFhSOxp5leYMuT7tEBqg+rCoyAAAAIlBm4hJqEFsmUwI3/vBjyyMgPRE6nSRtV0KpMIijNNSxNV6fK186yR1lZCuSqeK0zpSMgyYXrqP8vepgaucl/4jaGmvVdxyh0e9LI0kKda6Fdc80iHHn7M/V+CUeawOHDK09GYOfGXSC9g8VnDTn54KoCAuYx0MXXmNYVT7fUCj+DImYMTgYnhipAAAAGZBn6ZFFSyffZoqCW6+1T8IDFjZ/2UM4MdbOfM6cSgf0UAPojz7pI5kEFknrYw+ENpQAYUTYtr4SqEWUn7junbR8Wt56p5h7uhIcGkZL92u9QC38gUwBk1Ex5E22TDcZLcBpyUi+s0AAABQAZ/FdER/h7ECGy+MeBKdNfVxUPeZpFC63XlrzdjoJZreQchSff/+dpqAridjEZDcjRKKurU0q+oPY3sW4PFZGDT1eSQOPqlSLhDyd9rNotIAAAA+AZ/HakR/hnW4Mf925ITm3O+8SFKZ3zxUoRkrtFQyng+VH/hMQ/Sik3etRXXXNlxlZUXH5+UDoy4GCxn/BYEAAABpQZvKSahBbJlMFExv++5XDgdzjvuFjuy7jd6PFJ41DAH3l/X1jKdL2Nh/X8fAtcq/fEL8Eo81gcOGVp6Mwc+MukF7B4jwey/r9CcqRx65yLF3b31zSZv614yBGc4M9cNdQpEICcQUO1uAAAAALAGf6WpEf4EHlYPO0HmACcpM3NnBfztXc3YkJpZoDQxZYYDyvNb0Pj+EZFz7AAAAhkGb60nhClJlMCN//XbV/0CALgiH+1UnutP56iKl/F3VC1DeTp73t9y15JQ8NdWsfi5ioASn3Q9MNo/IlUuX52yG0tN0TwNd1c7MGXk7xvpMQBKyiLGoACsHChpPyimGD1eDR61bb/pFXOSc8YOX5I6sYEzmjHHEsZBjKDlilzAK/wULUzcvAAABQkGaD0nhDomUwI3/03wWPufQrGYALLUwsJoxEMXwn8/5NCn2gZUgE1FzDh4+XyZdmkDpzw9gAUC6kukl2lQx24vXEN5j5DOH1GIYPDhINtBDfyJ8WfNT/Zh9gvQkc/PgO/6SLQMtIl38tX4A4gmbreT1L7HRARz72Z1XC3C1MiMvoDJ7VVn05DVGD32R54DC+p00fXgSF34RaX9p/nhPqfqWedglnMwJtxRR8T655s+4GIfQryP4RVGWXnkJEq22REKoBSZ6ysplrhgLEE+6q2u3y8N7Ts9Gp2LifzfiCd1+mPeg3wsFSzJW1W4R7PzKVLjmOSAYJ0Z1tw5ZriRBPXA+vDYfsvkJniSnJ44XojzZkBW5ylkDkLDJmRGK798IPsVpoZydkPXtIrYMAjNlGhRi/1/7d/7lypvEWzYzVcs56JUAAABSQZ4tRRE8n9O9y35ttP5XA8tj1j/f8ilMGgI8Y0X+r/O18i9VDvuDif6vnGMNy/wHlXR9Ruw6ja3/zSdWeANSXvyVEX9mq8ty/YEARNUYIs31GQAAADkBnkx0RH+AcOZcuyaLd28wzeOfzpKa+1FVhZF6ix7MqNSnD1fychy8xNmjL0jFeLjxZFPXR4lL64AAAAA5AZ5OakR/18Wdv2q5RiSBDhQMM5SV2fbqkEFaulvczIGbVJ49MJ0WW4sHI4AZItwvvenMKbRODBhSAAAAq0GaUUmoQWiZTBTxv8IgPkZ4dbAn/C5fKq3XjZar6aAnr/IwTI6NfhYxOOdY90vHFZcpPrVgtYQEGJcIs9138OsOgln47ETCAT3wsxykodHaHt5j5zd3pHkivyreE3pXWJ39gwvIPH7BtJC2ujNyuwcpegUpqIHioXE7MREQ7wqg3InLIx6/+VYvQ5iyZJg9fQEMDyurdPwUG97r3XLTUK5xQEoR2kd86Au5sQAAAC8BnnBqRH+YjqwgmYeCDpFqcX+FVib2Lrr/3UMPLGmd1jhtC6ur496SG8bD43VmgAAAAFxBmnNJ4QpSZTBSxv/74Djga4r9eoJPe8V1sdAFa0CHxG4haVB5Ovh8u8POW4TcTMxsptzm5D7YlBYw8lG6CyBBn/nDvcnVGWnkGJ6RrJu5JKQCK1vTDPu2rBdI6QAAADABnpJqRH+HsRJVgDVevwny/m2ehVgkT1H0/+i6OSAZ+6Mbwnkk7i2l6zyDhghf5iUAAABdQZqVSeEOiZTBRMb/+8HrVh2HfgnJcCBHcdoH/EHq//8+L7kvHr9AwlItE+XW8GuuTbZfIyKey+xtzx7/z7v0mSlg6JyClysHvmOYE28a/Hm+/EFhTDgOJfpyQ5cVAAAAHQGetGpEf4jgOwqweXy2F7Du2qK5I8q22RN0g1ThAAAAbkGauEnhDyZTAjf/++A44Hc4+26pcXq1U9uXkaTwRX1v/02a9jGRPrQ2GgD4ByCHatmQBBgoB4cOdikx24qn7FSN+jELlbhvLO0Fzbb5WsWgoL++Qm4ZaBej9aGyTmk4go9Jd/7p1u9ZDw+8QmFQAAAAPUGe1kURPN99TanKxqV6ToWIFi47JibnpS3MA/q3/m8HYEriYfagkjp32QrwpuOP5xPUmZSLfxN/vcgqC/gAAAAjAZ73akR/gQeVxSFc8gpXH9+3PwF5PUSrMv/xOAUOg3reTKEAAADGQZr8SahBaJlMCN/PivaVsDLRwju4FJRtGyJld3PRfB6Xtp+0EP2ehbnkk6yo1l9S2z7XQSpm8aK8H35oc1ChfOjV/41daHgy3+x8N9GxWBhXCKBEFAFwfSFUpL6aLnnZF/WsKnCiSNGkpuCtg8KGhAifBX2jdPWMZybTgbq9FQPBnlegtHQ1gMS0JPKw8fpMr/l4VaXgPN/WiXupgFWBTWMufBFXx/pwY4hzINiNuuh8vDdV6dZqzU6E2zTAetTgZI718atmAAAAV0GfGkURLJ/NW8aI33FXDdlOIL56x/LE61rb/LEC+So3AiQ/8F2RvGBS9ER8hLuHAUZ1asXN72jsM5Y8Rsb/8m/wSfi863R2YTaypcHVPZVzRNGbKL9hgAAAADYBnzl0RH+ezH/OFUf3+Y9OhGlOXkyMreXizJcXSVs+sr5MlHfOxYn+EpqyS+pzJWKhTWIZJCkAAAArAZ87akR/hnW4Mf77wcPYM8+uXxARwgNs3mq4c8S0xyjev6Bw96W1w3E5wAAAAENBmz5JqEFsmUwUTG/8A8MxwHeY77MPL7mjRFVFR0Z/iYOC4Yp43LzKh/KoiomQ7CIxYiuEq/otbdQROB4TPZzrum1PAAAAFwGfXWpEf4EHlYPO0DB1ukwiu5JQBDnZAAAAjkGbX0nhClJlMCN/++FR2DWFm+9o4wWclxR2utVG2+azrP+x4C74zgq9slOh87PLWTihGxNCkPERnZsOwKKHGBDItM6n+UO9MhbU4gRdPCbDzqp0LYuZWoHRtYrRrXBf15GZzNVJg8jQkdGrKOzzbeXrr8TrvjpovWonk3/zv//ORDTwucm9tODFLZPZUK8AAABmQZtgSeEOiZTAjf/7/o39RD7fvlHF/LNkISXGH4z56Ai/14Fui1YLh2MuTPImlrfB6tXNm3i2ltMHwB0R0XoODBGGdW+0dQtq/9aAzDsN+WbhUEnMSSne3XEH5LffBS7TGVCwM92/AAAAXEGbgUnhDyZTAjf/yghD3S9+EKcj7DVPmcgdNOSh7FRWYDdiiT8S5ms2TgSRC1HdbXoPao/j1pyyvcmUyyOr7D8VfrorJ9Y3Ia+/Vy5I8VajatIrRzOFE3LjoTFuAAAAWkGboknhDyZTAjf/++A44Ilb6dVZAj6+5eBpPqpy0hGCly7Fd7tZOR8lBt+xP1Bp4udbg/7G0dq2F/ig4mG065Jsb5K+VasAFZ0NwTnEg3ZzHdvF7Eg7qw9SgQAAAFNBm8ZJ4Q8mUwI3//vBw6x3g997CjDRenHxJ0XtD3+ZsaWlrDYyGJRc1JoEwL+kntl1Dn61wxnH31YZTaUVFntEYkpX3nfnQ65nvul3NPCVGfp4PQAAAElBn+RFETyffV6IS33RIzDNhC6Rz4N/7Il8kZTcQoK/cQdyXIGV2ki+axJCPZb9IDd9bBgUYo7Q242WZg7cP39Ppn+R9qeii6RAAAAAMwGeA3REf4exAhsvjH6163w40hSkl0dhxqvJ28aur4o8Qw8s0CKm6YtPcSeQHz3E9C8uvAAAACgBngVqRH+HrQBVU2djmFhWg1aZbX8552eSKPQTInN9Jr83hsggwky5AAAAXkGaB0moQWiZTAjf/APDMcDucfhpJoEdgP1f82nspVGylrZAY5TqKCcynw/jcC+Jqlhay5xuF9h/wgvXJVaziw53TDP2hodN6QtZC6r5EmXQMakhypMctn19xQ16J6AAAABAQZooSeEKUmUwI3/74DjgiVfUWivtLYxGQT5W2ZTVlespsDp7BuKK9RAPQwV4UZkOEH/M1tUBScFtvFQZbdu4IQAAAGxBmklJ4Q6JlMCN//vhUdh6HHGNd0B/p9FfdgmAb2hpUl+urE83Dowm8tGYb3I7ijBirHv5xiW0qPCI8pYvAfgKKPOp/lNBaxO92h5BlDhH+jDsdWP13tgaQXw/5AYaS+Jl1e5qLi0meHE4ZMAAAABgQZpsSeEPJlMCN//74Djgdzke7Ewm/iNKmZv+eBdZSsBm3kdy+Jxax1DOGv4ctvMEmTptqD5q7dbTzcZOZfcFWec54awcbIvWYJNuhCIrSQwganEwlNR93+G0LUHO/L7HAAAALUGeikURPN99Tam3Cx80XWuScbrI8ef27x4OlvfwhN8a0R2GOpfaziVyuEYanQAAADEBnqtqRH+BB5WG0Yk+JGGXBmlhdn2n24dWagV/yd/lHGugo5EAQHiE5RPtDB43FuuRAAAAhEGasEmoQWiZTAjf+/2woACNjELAEvi75dnr9c1/IljWX13Dpl+arXwZ48mEtEWdudCKaJzjzn/zZ9jj9jnJ3T4dxh+JDYdtPLY9alEq/yvIYKIf6nHoWUKR4HUAVgyKfewwUN6TBSDcysMFrfJrLB3fH67CeBSbapkCkw0R7IsA+zrYIQAAAFhBns5FESyffZoqCW9/YL3/eJSR//l8bcyjUNJcqqPIQpL4s4lCEI7ELUU1x73ujzgFVMGgdWaoy6vKtupvoLwBbryeu7ncAuloej1ir6dhXCn/9p5M5mGAAAAAQgGe7XREf4exAhsvjH61lOxDQpSTOE5Rr6c/YMJcR40gl5ozXSNr5K0gN30y3xwpq1Doi+RBmswRSVIyyiZDT4vuRAAAAEIBnu9qRH+GdcG6fA0ab8kNOTv52rX70y5CdadgShyyECf+CNopeKFkbu/KLZSfnj1sdOK6r0ii9DLbnwNgGTiqNCEAAACbQZrzSahBbJlMCN/KCEDO8x7vwT26lNl/gUfBmtbre0IKq6E/qqdCqRddo6tWS7hBeYS8byVLG6s5bUk2/ZsQLzfaGjnk5bK2DJGat0L5NBmmSHVq8s6V5vFgGpYgYBplI9JWytw3yEcIR9mMup/7M/PMBbSjjYsrk7wqyK/gnjeQOMfm8zcWP5JRN3q8ihlin1D6q61GFw1zLmAAAAAtQZ8RRRUs34M606CqUVFCJWU4Ky0MnnSmBsKVBYKXtwQD9bO1GgwnOtqMNh0gAAAALgGfMmpEf4TUKzWWd2fqIhqL177tilPXm5J8DJFaCASw/ejbPc30eus0z8KtHbcAAABNQZs0SahBbJlMCN/8Aqdq3gjhBnFy5XA/L+1MyAlYxhiqCHpYw+5gq9qVl/jSDu/Zxz8k+rcaAMAWshx2Lt6lAsCXP1ggHOhg0h9HO+EAAABOQZtVSeEKUmUwI3/74D1EHSkI9FUIquTDM3rKaNWSFy696/Qi+BwnpHE+ukHESGyxNmYBi2gCDySgK7ykEfDx1eZ83MsyM6XidSYhkeMVAAAAR0GbdknhDomUwI3/++A44DhT2ug1cazm1y9Eq/w51RkAPdbegFjO86K9lsQdW+KZ8FITyXlyVrMcZ2dEB+mty+r1aZPOgTfAAAAAXkGbl0nhDyZTAjf/+8Gbu08WB7e+xX/RKyCUlwLsCM947CKmmivCIJDojdrhemKwj68SRMATMf/1ZjV0RX9q+NeV7iUhcKy7URTsFOrM01gL2NmAF0DuInJP62fDE0AAAABlQZu4SeEPJlMCN//8AoCu8Ic2OSncAbLPwyKakd8JvTcutKPXYQaN0OqVJNStMu+lFOFblgBZyntC3N9Z9aj1T/qbHte4IkOIPGS4G7XHAIEZZ+vYQDNER4xNgppDUTCuIneQCJ8AAABbQZvZSeEPJlMCN//KCED59yngQHs2OJROVEKCAWdcGFHRSwN1YWw4Ja+vCA1DR/wWuwzoDuSzJ8m8/s8a8vs6MregbVfZO4g5Nu/OfbvwQXQzbNm531lzo6++2AAABOdliIQAJ/++DLeYgN9XYj5zg5B5QJuOQipq5DurWEjnW3UorRPL+Gmp7W/eIApq4hA6GO7MXHIyDHsVucz/aMFYtKTDM8z25ZLQfiNZJqSiDN9723ESzhz79Cwt5M+8vOWuuGhG+Tvq9foXjrYoVXiVDwEsVfUmlk5TXbEqyYI+MZpo4mDJdCI6FifaFbOBypE2MZf5gDnzZ5v8KmDeJZU85nFu03pUJirEg0Ch4LatUJlnXu1uVKmUsxM0BfauVLDb5iKlbV8Rxe706n9ByCWTYcNmUsWgK2BOx01vNUSepY8KQudanHo3PJDxToO0OLrabPAghuBeO7puO+FrVqFKMISkHJjS3EmTdBG02wgZepLedOuw6T+jsJltu6ElokeeisTTrzRHarkaC1NCRtCBiNuKEigoGG5TW/NYmxDmIUwJrPGuuEwRd1KazxAA2rsqyTJqLM0sjQZEKi/uESROrh+l53XQlmI5Fr2TToEa1Mz9Q0myFVMzlLsQqx3/ZRQx/eaJZU1/3iRY67X7pQgN8jPG7Mu6N8yvaY+Aqj7NpykjyTLzJHyZJOS8uJC5eDrvviZZfFd8hAX6jnJ0ZJE1ceTBCU68b2mfGkwl9oA3rNyC3kk0Uu8eOhHrIGj+7a64L8d3FiwoDY47nxsn1Y9aseVrNpfuZL8I3UMDdSnrXcWoYyESX7OVr2O8rZxn9LDxM6fe5FRhF4y7x35eGeAWARki3ksb7bLxg8DWvLG/+RwVPcpcizEE8l3Xt5+avL+ZiAXLEnwFkaUZDgHR/YEmvUJL0akz3W+4HoMrkbavzzs+VqiPswLqfDbnOg90ye7LST+LLq83FG1+6MSJeC7uys0Em20StpYNB17y0kIWJY4WoqsWsZ431/5zMrMnsO/w1fUI2ocSHX8pNHwFp4puUb+kPcS/Qh09xlBE4IT8uQiu8T2UDKEJf9iCtfqONZHMdPq/oob+ISU7mDp/lbACtBQUxbSPdO3HSaLklWzLgg0sclDx8EWk7I6bx2ZO5UAMGXVQtj+CCaYjfoBtDkXhWf7nWsT0hHJZrDi2Lra0z6P95hKbBHipvDDcqX6r8JNPqXyyckj8WVa7ygUVg99xXIM3DG5RRWnB41GtWfCqWChNibXPGQ6uh1H6EF8dmXH8surDi51hLT19j9myqwkNKmT8sfWOXCSiZtHYNl9wayNXOnug23v1/WAGCRrLJeHUUtVHTIX5R9cvqnbFcs3PTXx9R74CfvppdnOxRFNeTTBjRlRfQ+cjnV6LGJ+ftcX+9UqNBGia2nKp1uKP4ccPHEr/CsMJq4q8HWHOTdHO3jV+GZWnEoLlQ6Chy7Hbj//eKkHIOL4i+yqkzkPRSbB7sfUP9vHZ0uCQVKXYKu//r0MgUKN8P5NWtC5FEPRp8HfsROGFe/2knJyGh8s53uvSdYVz93B3c9Z2acgXQNzwm2yabDfBYT4AbEZLwPZK7m0eC5LKEHAxOLjG9NU26PPjXEMfajfx/8iXqLqjmgozrPx8k41W10JV5nX0OJWGa0qBU80PviynD2x3OkSEckOt3iTHjB9RR5c5nsqTL0EgDg1YWzkmuLeJTeKxOy3pZHTLdnNqj7pt+W/2OWi1UaTdjKxFLON8eBwDadGOUrbe2ZBGiOQ8/zfFIWmtk1vwMVGZD/Iod8uAAAAATkGaIWxG//vEnVTkBINY1xBnYqV2GYveJxG6pO1ytvZdPbEXt5LeT2BV/GdImH/VH7OxTyATpFD6nkhtuovj/nla9ShYi1okzGdxi99RJwAAAC5BmkI8IZMphG/74DjgQ7aLylrdTsuPGlvjUJs/BGQzvh+6+xQ9wiS1hY7U4V8YAAAAhEGaY0nhDyZTAjf//Zwh2MSAtcqcH6DBw+kiRJyBfSLlTlMVKsKbIZvDezqc3T5Wd3NGZa+B+Pjdq0z2oeP3B1ynbNRtvd91FAPKHi9eAZQhXkzjRR1YSlxJOK3yfwo9klCCVTafgfuL5sDFzP0e62FQJQ60CKx6fANRbJhNojOKZPzzwQAAAIlBmoRJ4Q8mUwI3//1jexb3GfDeuG3QIp0iOWNrPYaRICXk20KwX4AJSK3G01g3yC/ow3Dys3hXKUnzco3wdRZFXg3LRNpOH7wjxDdO6SGSNKnYi89v4MEaiVQLPE3ixfnhiAfLznBdpEhe+byYE2PLypisngAgqtZKdabisoZIYcgmWrafiKuWvwAAAFFBmqVJ4Q8mUwI3//v9sKAA/rehi7vgA080JM8SmYpPIuCaWfrNkOSrviyYia4nOycnCZbBmW3b9POpyhuj0YMiNpTWFdH6ofSMf1BO6NsLk44AAAAvQZrGSeEPJlMCN//74ByCXW+nXGeBo12ucPu1Q3jIarCClBC9oB9Vn8fnB9XMTBkAAABQQZrqSeEPJlMCN//7w3wmMjIEMD8lMTeUbP8EWg5RZj/kvdI2pZcX8B1kNiYTAl2O0lslP+qTbr+aNjnvHbN84RS2/4Py+FdljGGrZaq0CiEAAABBQZ8IRRE8n4DqJb7q1b+zyENk9E+P9lLAQ/bAnTrrSAta8TkZGO0FbRoL9hiQMM6FUzWT+8DFzTHADnP9RRGF2gkAAAAwAZ8ndER/h6AvXSm3Bj9a9bnh63e7Lt0dwASWkqO5RTd88ehBuGdNotf7Fgu58KElAAAAMQGfKWpEf4etKGWAMuE31cSYTo8cYlrTJwO4hT1iaYKoh5yIznLWIYM+u2AdetqlZhgAAACoQZstSahBaJlMCN/77mtspIuqyAydbMmdV2hUlnnEKtJjAgm2ilJGA1nNcqvKyjyAWC+D8UjH9KvdInrfn8JalYu7So7UCAXcOewUU5OwhCi4Gn87YJgH6zDs1GEUzCJ8L4JntHmTTg/MXIc99L1NBE6e+O3vv+OYZ2RYqJ+FOkeVPph7PM7bvRsalZ9v+6j8c+a2e9FeVHNkUcRhx7xXVeu40Ckc84xcAAAAKUGfS0URLN+DOtOgq1Qm/f1v5ExIEbysTirTzQFn6S1t6UO66bJ9QaeBAAAAJgGfbGpEf4jgOwq2BOh6mMPEvE5KJWUuH45P8CXP+ZdCo9KA8njAAAAAkEGbcUmoQWyZTAjf+8dbF2niwezLbWnHkGUoGoj7QIIOkFGvROy82ieNrbjUGA+u27nmEN5OyIUtEYm1Aa29w3BQ+77exlM1D2vRT7tzMqq7x40NgyQf9UyDOeFALvKyW/UjEdd+ZTj3e7SNxPWZbzdyhtLBs/d1mrAgFnRSFTdTWE14mrJQdgV/38CNOEQBzwAAAE1Bn49FFSyffW4h7bT+QcFhQIGL8ZhHK2aOc2q4nANbZ0GFL6Hw3v96ujtmH3X1f2UGhLCU7s8t2ALYQ5jkD+tMpZBnKJdI7zOIZg3UfgAAAEYBn650RH+HsPlaZUroiY+o5658f/kip2JFnP8SLnrUdiJYATw+vVm1xpmGefahnkYvhxrqmtK90ncP5Iwqoo7sRj3ypQHBAAAANQGfsGpEf4TUKzbRXPm7/rDCb1PYRHnrxQN3OgOeuwBoQBUyoz454hkxVpE0j3mzip8afJWLAAAAgkGbs0moQWyZTBRMb/wDxbhBIn8Ka729n9RowP1GPGwmIKU7l6A++D+NCeFcwINyDokoWo5uNZep8A+tnW0VIejZtiEBw828C5obBhnpukZNMqbuoZzcKwUcdD+TUNnmG1P7VNp3N/HRWhTj24xVfemeYzegPZhwqPtf5Vu3ZOBdTWMAAABFAZ/SakR/gQeNFYTouJTIPcX9ODyG/mOY1LtDn4Z4fdloxMzM5TGxEPn1lw/NNkwpQDBeevJFEc5m+/9Go9q+nDesjfivAAAAmEGb10nhClJlMCN/++FR2Jc7HZXI1ggkQzpQbLdE3oyI2e9RpinMX6ttYauzP3zhtktbs6rBQk+2bRLN2Zii5+KnqAGy6xu+oWg28GhfC0K53s+ETlNxGnX6kbnMFmbPxkhZYhjRahFfw8hooaTDBlOVRucGi/xGwHC4t9N/ktnpgbYThbnVUl6j4xW46umuQZce+h6EgY5YAAAAP0Gf9UU0TJ97IKea95tkF1g7PulJOBBxUf3oQ8Fgj/d/EvsVRz2v92wCCfdP3jpU1XFxnM2iUhxS2IZ5eBICgAAAADUBnhR0RH+HsQBkvuNRCE5o/cA/hS/2X1Ifx8G+v32b7PS02yVwPwQRWbCtWYhN3Fq4gdnPwQAAADMBnhZqRH+I4DsKsOugF5agq8KzwVB0SOKhHPr71cVpWK3aQSfqFUhMbf8T7Lt5HmLfgm4AAACIQZoaSahBaJlMCN/74Djgdzj87Yh9SyaadyNqF883ukqwWJfvVqpKQh+Cq2+ijPdoCbCBI/sOgdzCRvS17Gg4tWxC4C6lNS9Io4e2m05/dsoHILwYVElOoJDZ/MwYJQtRJ7sMYyXzLK7leOwuqfKqCHgoXnqt0hDtESiOuqtwcgabWxgY6yvJsAAAAF5BnjhFESzfg/tuGYHKjTB5ZIQwyRxCIT4MdkzIYOd8632MtmXbBRZjmn1eHX7fw0cYLgael3Bt26A9uYW+WQfDAKavRUOn1s++6E4EarbpfnHpm4uc7noN/2Jer1MpAAAAKwGeWWpEf4EHlYbRiT21xAh/pYXaBtudfqA0qjCKgx/lAHKVbTowHw8rOoEAAAB5QZpeSahBbJlMCN/7xIFmvqIe/Z6A/pmYt10zyE29uhmz+B2Dkk2oJvashyLN7AdP2QFAd8nNOQALgCphNk8C9krr0S5+4t/7IC9jwLdRCyqMfhEx0LgdtUukX/rP3vR0dcaKekFBDBp9NJs3ELeS44zfwhtjbfJhoQAAAHRBnnxFFSyfgBLo0GKtuOp0S/62v/Y0HNZk73O1NRvXJ78CK1NMTqIBwvVi+7nrhZNFcYmSvBtS8aarNw0M8/4bqY7V6eSiefuhvHdnfNViR+1mI4DfLT0GD/NfHmcU+shsVRJ+gRZvcMuzDyIyflNTInWB4AAAAEABnpt0RH+KcQxiH9X/rcK9KC9jqj38E9Cx6kZJfJmOLXjenMkHGGnhFyfDLoqpt/GOsv4ayAmoWSNp05UVCWVgAAAALgGenWpEf4Z1uDH/dvK5xRRkl16iX0NdLBx0eflcOeGwT2wyO5F6Jk/qa44eccEAAABoQZqASahBbJlMFExv++5XDgdzj8Soxq4R8HyFDsHMPNBcfu6u0JC3yOt9Zf/oOlinfwad64b/5ZuUQET9ghJNRnLBbp2fwAk7+v9LiLGc0Wy7m/MlaKm5bEbUwKubBxg4bDV+ieC+CYAAAABHAZ6/akR/gQeVg87ypn7D6HEVU8+rDMLH00gtcABN3WFNH/SXGNGasZxuBcRvLGcib1ofmI9nJ1EIhafY/27h+f0U80Zw3TgAAABiQZqhSeEKUmUwI3/74VHYnLGdrsN8fZFN+AYUaSy1qUzsqr5angQrl4TAR4cjThoFgPzgtGOm/+rZ6RiEUZb1p/ibsShVQuItv9q99GCa7QVtMF7k5//VfdZc+IdQXtUSRRkAAAB0QZrFSeEOiZTAjf/7wZu7T4AET0n1K/wcGaW5FPPPeSEVQq1wqPmVjHHGzMMaLIFS9D9xLjG8IHT3uYZjAjm+gAJnanr133+IRxjFAFTRBkx2srdq61z07DgwIEr4PcBIbo6eVoNdPRdn3p7C6JhswPrZOxYAAABJQZ7jRRE8n31uIe20/lZinhECBiXrAbzFWKWjBooZi4YZeyeDcnLXtkoxy+PPzoG5pJuU3cZMSZ6MT5Sk+0MRyXthyQf8s2D30QAAAD8BnwJ0RH+HsPlaZzJLf1BqA0GQv7XVPFANHr6Ehnvy7et3ba61H8sodLUCczvDm+Xl28e9s38/yEiA9pPYUvcAAAA1AZ8EakR/hNQrVcoxFq3sC1exT2bccLesQAlI9raV/idtSMz02Fau0gbClXfq1HVRq3Pjj/AAAABxQZsHSahBaJlMFPG//AgTwyBIj/nZtn0aw/YuY4OjamoTtQYFPOSeyuPadadakSLJXyzkE+yUN1FhGpSW0m0mzOfSaLl6RQi3UFaI3XVG4DOFmNwv2QiaKoOLHwxouX+jWVEmhJBs6gzIfBVG2WvwKcEAAABAAZ8makR/gQeNFYTo0teWJjfvPMFcU66/zULU/AjuVpjv6JMs2v8Uk8dPu8G96ISprega9vCssVT4Py38fqdQhgAAAK5BmytJ4QpSZTAjf/vhUdg9mx2D8J1iA1xk+D/u/eWFRKPYXoR4YrnXnQ9BJv4U7EyAxknkxv1P4KGiZu1XCzyMTxrAEalziwXrOaH5CIg19HpXT8adjJuBd+g8AVU1iY+ragKeFtHO8+FXCs1iqBNVFtNR+QrIASXXiKH9ySLmtNtwkBLAWVkei24IbUCqoeu2eSmChhIvR8SYf+gjoCylIj0hQWhGkF9P9sVMkdkAAABbQZ9JRTRMn3sgp5r3m2NaOy3X8HPswTGPT1V6qqyW8TAIkb1l84u61oezFaXM31Ex6tTfZXwH8l41V+ZLVvdVcPEZOg5HCqpsYndTaGYDZU2IBrf/HbgESDFhYwAAAEYBn2h0RH+HsPhkvu1KRWh5P9CE/7YMZ8EH8FtRGt2iZZYTb+I20oUztyey5/gcWsMembQCMoi5qzr3GWB6TyRqnmRNbxfwAAAAKwGfampEf4TUK0JrudPWO9DprT1cI+LcpiA9d9Uv2SIBa8dwd7ktVb0K/sAAAAC9QZtvSahBaJlMCN/7wZwidiwO32tSx4CI9VHvf8FfC8imIOnDW7X9dLvxLE8x1BQ2+Flva/hMlBEeDnDW8Kyzs7nJ6R+vwdkQhkORE0CohKGRVTKH6TrAI0jsBYWIHEzhI/DD598H/p+WJCB1WNOv/rDnbdHgZYdaqnTDwcDWZDxml47NphnI8tuJ1LEeyLF+MD0/bGs2eXrqa0n+B7DWP332LtukSOVrs5rOy329VanveHft7k/xSHl0G6dpAAAAUUGfjUURLJ99PYWDWGUO/inwiBAxf/vE27SbnXVov+t7ciiHO3pMMbk5hOMpHCmJoAtQU+lyEMKhoGTKDStuGSi6zfSBRDqB7UFC7Ub6lRrH8AAAAEIBn6x0RH+HsPlaZzJL6599KwwEOf8I9PqjXFRQKe4Cdxrd9B4PO3J74fYU+3lo88Z9cxAIlkoQasgv4bTkDHt417sAAAA3AZ+uakR/hNQrNtFc+dI4UgJcDtsPKVP01/w7gtDdiG8weomxmMV2jRs+mYeb8YsCdKiDBK6EQAAAAG9Bm7FJqEFsmUwUTG/8A8MxwQo/6EY7tRsTvMVqrKMcemP06VoeULKvsG05ajGPmO2MoK4g9WsHj/Pw+OQXY/RfWlxRBDYU5qtmrwgOMibbsU1z+xBDOIFJfIEU7QhyGFhExFTti7BO80z66P6DnI8AAAA2AZ/QakR/h7DkMUOCIRZCQlH2r3Q6geLsF83caeCI1bZY+Xd0YL1LmfzC34lEn3p7FbZoKir3AAAAl0Gb1UnhClJlMCN/+8EZPlhAhzY4hjuP9zep6GX41gR0vpNBRWRZ0UJPip2CeOhHj9isfDJTZVLrPDCGanOAhHmGphCIAPC1fKwnrEyemLG2qRrRX0P+lii+U7+s0UA9lDzojDxhAWTZ2G0F/aKbm54hHFNSHJtKf+oJOsv2i/hXe6R9WFXNJ0qKbHaCDXE5SAk17PNkKsEAAAA3QZ/zRTRMn37oQ9jW+VIhCNZdo1DTECbzHzraqtIWwcFXlHzfzpW3DjaPNagdZrahX3k2FcjBQQAAADsBnhJ0RH+HsPhkmMfK/TRTx1yNE7nufRZv4q2ObBg0Y85tUezANeXV6ceBVnWXr0jPx3kmWjuo3pxLLwAAADABnhRqRH+I4DsKquYeXbXa80z+Rp/+iDuHWPDYZQauQjmlruP/lzQ5MNgtrxuo9+AAAABwQZoZSahBaJlMCN/7wZu7TxYO6+X+O2RPzaO3P9AZSfEckBGR/iJyC69X2mzEdsbEARVatDv39gXZOq5q1DbnQ0mecF485kMmU//5paM4iZd8h23uDrAJgM1rO1fswnjjEy9xf4szlbNzDknjKgAlhwAAAE1BnjdFESyffW4h9qH5WbMcugQMX5WbtC61gVlqffRYYu75nGloWw2+flAxCV1Cr5kLWgK0cUy5ESSi8mwPgAklHpUS2QdtwkoX7J+G7AAAADcBnlZ0RH+HsPlaZzJ17dg40xRcL+9TtQ96Ic/V7TTmS8UoawBAXq1QI1QIwW783wiVM4yfS6ugAAAAMQGeWGpEf4TUK1A4rnMTHCUTE6/ZCv+/TbCro2Jb5kNAeSIDyYf8WGZnEtBvd7y7p2EAAADtQZpbSahBbJlMFExv1TpdGQI7CrAUef/B5h6ksk8VHxtoqBOtcVtjm20cMpPx7YKCba8PtsDiRbMvl2Vzp7ViJJGJslpXSe4Da21o01Tre+PL4ouLsDE2fza/o8u0e/o193a92YPHQjxT1F+Xl837GKhxj1zCeTYoBc72ajk9wDGlXMBiyPuI/NQcbrbcGJ79AkHjk9YLxxYD/SNREJer2QOEa7kdnKAs2nctMaAcgEj7lXi2nC/Ax5gsLQKiDWNIJ6EZsLrTgM5L6GMBmxmY1G//WcuQa5/+wM7s0+v98zEhGRgkPv8Pqi9S60lbAAAANgGeempEf9wncM6KwnCyxLkxv6hw6qpeuv8JH8/pARA3aQiOdCW0jL7dvTy2HOkGJ2WI+L8RYwAAAL5Bmn1J4QpSZTBSxv/0Ej2oDXQbnmJ9xHt6Y3M7+zkeot1lPkVAjOrBk5KutJxuFikQjMMm54iVUcgq0mUJ7XmxxVFK1B6T6213KH562Y+RpgQQ2x7gidLCR4uQyyI44VQZbTdq2DELgkXsCdiNKr8x4ESrFQl2hMUKVh9TAo5qMet3aV71iz0PMWrns+RwD+p05KZ5PkbOR8UZdEHr12Xyp6MXO03cMCdxvt8XSDnM5DDAwory7MxXFStmXjinAAAAOwGenGpEf5iZsaTqsZBVHN23kGvXifJ7fiaFVng/pASKG+Iuw2tQVKHpTjjL49UFkCnlve1QSoeNybR0AAAAc0Gan0nhDomUwUTG//vB61Ydg33yGjTc8vx+XfX6HCnb3oXqwFbwLkycdi4a18wxKaCie2RwyAvsyA0uooexqJ/tYWo+LRvOoUmoxRFTOQSzQtQ02/rq9eSsBVCGtHKcyXiFeX4sl7E4a69xKNtZ4XR2mcEAAAA2AZ6+akR/h417ZNpqEQkqKLsoWa64MImq3wZOBY+C6717C9/lyBZ0g6hTD0NY1rt85FbbEW/AAAAAnUGaoknhDyZTAjf/x6o2uVDXrlL01pKNSmhQXr4P4TJQsSM1NFMlKvnjcKn1PKq2soE7idryXsHEaf0O8GkHVQjbY/noZr0ipij3QQYPsNY5BiccYwiMaOIoVLU7GYyZILtn12cl54KfPwPTwjvSw+K4bQ5yOg3G/Z0pavUu5uigl8D6ssFCj6v5AXSFHWAjabypdfzcI1hiTFK83ZYAAABNQZ7ARRE839BbCV3PROKhTO6SEYHVcap6D/bB3xdjsl1hYn21dsFTAVhbVYBtR9XbqXR7RUq+tQJh1Q7mOALEIaPB0INz3lNIRir4aIkAAAAyAZ7hakR/gQeVW1VzyCm83gTep8bM3if7zuz+nCbMLSDxTS+tvaDoEjU6XYt9NoaA70AAAACFQZrmSahBaJlMCN/Bc9IeRdMKBQUoDdmOpylyK4+kq5dpRNdskQahtc3w+l0fbnrSe3FvIzIgw10Zk3I2ld9hkChK5nRYgL3ZtOHEBaiunfTa4ieYkASLLA/NG32QGn6PIgI+rNletA0lBbBKcG5b3PN4cer0cfvtkcfPCHp+QZU01FpMCwAAAFxBnwRFESyffT284Ut2gLU6T/OF97/L42Zk7TpYejMubi0UT1LYIwyi8wna5sal5SW2vk3PndfO9Ygjv7w527MVQundwsPGKGTzDFCdJOti60VM/ivQ0/9l0RrwrQAAADMBnyN0RH+HsQIbL4x/wbzGsYYRQbbp3oKzYEUN0nPNIKpfJlJ0wD5Mk4OdIwxB+uB2eKgAAAAnAZ8lakR/h60AVVNnY0Snvntxfl4ah0IXe77u05wjItrjxSxk7twJAAAAPEGbJ0moQWyZTAjf++5XDgdzkPa1JXdlt/iHg3e20KGGbpJNaAeD/NHcMVqnpuI+rGRg85TMlkZYGQ+buAAAAGVBm0lJ4QpSZTBRUsb/++FR2DWGd4E74BcOffFFDrjSdNqIHPjqmL/A4V7zuz6iK2RJtZ67hmaIV4RT+DJzIBQaAPE+vuu5mX9yRlf+Xz3MDvkp8OYGDoPwsmZQhPHmCII0d6nKUQAAABcBn2hqRH+HjXtk2moRCSooxNiKe6PW+QAAAHlBm2tJ4Q6JlMFExv/74Djgdzj87DB+8bpj7b61MAK02m0SjItmx1gb8vC/xkdmAxbAK3mtGvA6G2vVc9Tnw2/i3sUSrGL/HbhKMa2sK/+P9jhm46fzs6O00Ue9XW7FkptH/JCazHfcLFaniM469HEARwqsesFKmIPAAAAANQGfimpEf4EHjTTHLttWmKSJqlJBcodqFPtQbfjXOzXaRUJjHdTp9FoeBf71Kn+yqxtU3pUeAAAAXEGbjEnhDyZTAjf/++FR2EAuCDGvtGowoi8nI+jfRpfapBnzNQPYo2Vr8QmdtgYjsnpaHEEUEj5fzdXZ5z8ZMEZYVB+N4bG+rt/oQ/tJOOnO8NMl+l71MxcY1QGBAAAAZ0GbrknhDyZTBRE8b/vG3I0XvCGxfcoqF4IjXpZsDkZv7oRIdkuZllH2s4RDfJ2UFZ0K+2+MrbzLGTPgjZmyPT8lOVqYMq1VU/V4XinxIGWe6yhZAC7TBvENB4AWJBowQdYdeboaXvYAAAA0AZ/NakR/h6Avf/WMAaiuyJhL+KT4JwkHNgd10OH2kTFHZuZ+4In/rkqbThwoea2xhOS5EQAAADhBm9JJ4Q8mUwI3//vrPjgO8x+Iwxa5AW96PE6JjAAg9tvdbYLxnthXvYjAICsZOG1tuY+xvSOzEAAAAEBBn/BFETyffT2/FQo8uLN+ZccYdXGUegcTqm+73SyiRqyxjBgiXqvyckNmzTA/sqcKN6IxF8rB8AB6aJm1tPGRAAAALQGeD3REf4MOYBiW8DO/c0klwuDU2bHJ6zPTXTSFXKSsxC0nU39JMiqbRsiRQQAAAB0BnhFqRH+BB5VXHdOCXHYpVjQiP78rXDhs2RTqFQAAAHNBmhNJqEFomUwI3/vhUdg1hjtBgEPsnnqReMCf9fk0pXkAepqW0GvzkT/N7d++eAjFF0LqEWTOQwynb70MqWR2LVAT3BGuJjMDsrfL3saQZmSV6KVXH2F2scQ1murg2f6hRdm+IshdKI4OTJ0vpuACyrZZAAAAWkGaNEnhClJlMCN/+/6N/USMh05oLuDoAZwruKxif239tQ1VJ4v5ZqpArC08Gz4PxiK1Zw3gJNSO9N1MaBrO4NXyS0Kk5yJ53TqLyQyUa3YaV6o9CrHZ784vgAAAAGRBmlhJ4Q6JlMCN//vDcjlriweTEJ1e1T49rsbXnEYWTbUIkOTqF9p/L4Ed+sSuPC4msb/sWIIL1Bg3A0EKKr8CzywgYo2G0BZdMZXDAvTPj7Xu2U8kKOf7J0Nmp5Nxc/b1wCJeAAAAPEGedkURPJ99PbKA1hlDv4p1vxOs/Y4zfPg6iHF0NVmxylt2/uz5ddxtZSf9odYXHgkvwvq5a8FEnAsAeQAAACgBnpV0RH+AbCrDaMSfDlBTuXnhXQUYGq59CbNAd/8CU+K1jwZOGWMmAAAARAGel2pEf4egL3/1jfMenU+Ipl63RDQu+raphKcMtMgItiKQVS5cjIUWjl5mHHdx54Jm73xhCUgr2g10y3zkK7P29MDQAAAAYUGanEmoQWiZTAjfyghBljfBFTHet0U3BKQoBgmH8wmCWbb2fhS9kf8YrXaOzQFjOPpvXVSbYEH7/91upJzbcIGzTBnw9iN8+sz32CUkThqoo1kDvR6ExRFTQPVt63gM1akAAABKQZ66RREsn309tLhK5ErqsM8w/NBncxs8+XNMfQtyU397Ah/MDPI+uVDBGq9NMLyY/eg7uJexiq401plodH0NOZVtS74swPLgbj0AAAAnAZ7ZdER/gwyVSt4Gd93RvQzgj5K3CnQvm7OjBKD3UgCK1+lg5LhBAAAAJwGe22pEf4EHlYPO0DSNXbXCGQv51Mgmje+6XpjUxwk1173/hFbKQAAAAIJBmsBJqEFsmUwI3/vhUdgSGsf7YM42S+/vjzG31Ii8r+yG9Z2rwtMeopT4f1ZXmgCKDIL/jmrbYA8wR4mFw0PT9NVLBPsNXrKLeAacgNRksycpYqdXIV7VH+6jtYtcIHq7VaKAVYULBW7vO3pgC3A4br5z1oXvv7pxKzVVrdDydj0wAAAASEGe/kUVLJ9/0qcZp1xS4Q6K/byHOdoBIV6Xps/UwNkM8yttvswZ8Jzr9I6Yffad0OJ7Pt/vu8H/euA9uoco+FMMpzX4CWnWQQAAACQBnx10RH+HvdEvcu4XZBGgCHIf5/KTqwJP2rRK48Dsj9hXDpIAAAAlAZ8fakR/gQeVxSFc8i68nfap4V2gL3VVnIbk0JUIrK52lB5NkAAAAFpBmwFJqEFsmUwI3/vBm7tPFgQM7JWZsrd+AFcbR02voqglcUnWYy1KjlCBoWGrB4/yDkcsF/yfA0JiredR3d1/zQTCNN3YSNi3+6GLuK5XeureNSQvNWSNTKEAAABLQZsiSeEKUmUwI3/76z44CE3mgfD3z/rIASyI/4ajysO1TyQ+cKUy8BCEPXwFvV7Q1kwh0NGs/guw48vBxqodJo0my+qYZKZwCVTfAAAAXkGbQ0nhDomUwI3/++5XDgIGEb1mgVp4z2caIApSFuXOZMOWUKsM6uZXxOHgwgvvNbQRkVIuaLICKhGo4osmKLu5nHbNQ2eBpkjJ0tBuuGJ+8FlhcF1OznrI2UofcQ0AAABnQZtkSeEPJlMCN//7xIFmvqIPwLXapFLrpy24N2AYpX55tMIyEhwIGCCAWtX4dRhVjqoPHEX7ZHYi/+5mDJL6QpKDldqAw59qEcj7XmhguiYHg7hxSzxJYABsNFnTFdhSjxeeeQ3HwQAAAFdBm4VJ4Q8mUwI3//vrhHYIVisLS3Cz75CdmPBNSQtJmyjSkVk8S6O1d/mV7n8Rrd+6LOzS3sEdBdAhEdUoK/3HLiqSbIKyHxx0QL+V1ET6o51gLYhrc+AAAAA3QZumSeEPJlMCN//74DjgJBsj5aW27b4bt8+9bnwCqvg57ian/4Crkwz302lGZ2Pv+rWZcpSzgQAAAHVBm8dJ4Q8mUwI3//1jexkG5DGfTzn6ACkLgiMo7PNL+8MNWKBuY2ww9P4WA81P2z1Uo6HFLULVfe0UWARzjI9fYO1TGuDiw4iAYUYp62du9UL2+zNI7rHN267ZkwLWg++4sTNfEBlf+9nqtqNbpC7ef2VGPncAAABVQZvoSeEPJlMCN//74D1EH3TVIcIrzH0/chzRECMo8UYdBgVM6BUQ/qnmrHcIvQI2T4gwNuNCHOfc8xOvd8e4RUUrEg7UZrrwxAx7PnW70QrnAkn3gQAAAEZBmglJ4Q8mUwI3//wDwzHAd5j7ZQY2ZrSRVRGS4RBORnR2QxEztLNjTbcTvgLcPBn1MFP2hKUm69thjXP4Z8ejw6yLt7eBAAAAPEGaKknhDyZTAjf/++AcgTWtokf5L9MFh6k1zvz+EhUSCxgLY7sBvcK7i6dPmedsqE34PPj3OdSB7bkjgAAAAFJBmktJ4Q8mUwI3//vgPUQhuJTLhEJqu+Apna/+hU4PzU02rtLQFGBiYXaJMHEktyfr4sDAJw5gQgiM1vOUDKdw5h9IrJLuyeXoGbf4wrVN09r3AAAAWkGabEnhDyZTAjf/+8N7/LJcoi6iLVaYzKjcM3p9XtxS6Kwh3zS27UOOJokA10Oalb5Mh9jTjqiJ2mZpX4f6frTs95x4EQuNnHB2EuOeRNjdeN33AYWtQV+GtQAAAFZBmo9J4Q8mUwI3//wj5K7wgzjvuFjuy3LAdQ+tnajIU2FYf2/S8asjPFa8Heexf9iyg+O2t1pIIQe9N50utnzQAT4GusPtoeQwLMtTfJDIuk4WzhVsVQAAACpBnq1FETzffVS1iQAjw4mKGWZ99z4TIUbS/mDFOaQ2BzRPfOEJhO2fZfkAAAA0AZ7OakR/h7ESVYA0mDs38HcgN3BlSvTaG65b/g/pfq94yFXiLsMlUCZdl+Sj6EnesriDgAAAAJxBmtFJqEFomUwU8b/9Y3sZAzW9hPGahGw7flXYE4FP6yvCNzmXPKd9sNuacBlbT71wo0upN8/BkEcSxj5QOSyFbOl4aFdbwIN/85qZNxlQT3WmJ0lyxLEqm3GKGmR/Q9dLRwbnvK1UjtfHRqg/kLxCsjH+LVN9IFW92ci8U6z0Su8pzeXhOXBhoqHLd/KR3LzdJ3Q+/pqGJRbSf5EAAAApAZ7wakR/i8dTtuWwJ99n5vKBTInol2bsI01rGqiIrOoR71X4Q081QcEAAABXQZrzSeEKUmUwUsb/++A44QawOdhgfUi1MfHqmy+I4SJvb7slhfqvjc6dx1/LRyTOBiepLebRHg/6uB0qH06DelPtYLjqPVsidGv4WOqaIDxCE334DXtxAAAARQGfEmpEf4exE2/rMauJo+rnrnx/+SKnih1QnG/vzvVzCJmpdIrzItT2aVilxCxecVUkmU6Xcdbiwv1Hi2eVfOYva6XMcQAAAHlBmxdJ4Q6JlMCN//wDwzHCXNMnztovd14vJ625TIq7dg7bjOkJiogvd97B3i0yj/I1tMjzhb+H4NpywxTMQVOk9VS90im936odpBgEAln5etaGgewn7zxeYOtk/lsAB0UYEr26ytGYTnykyW97BDgbmncu1z0/tcTAAAAAU0GfNUUVPJ99PbdJpKoJiBXMhGSlyCkS3qgHnDWoGiC2KtexepnugaGZOIX0UYhHhpVdynfptSF4E2r0Uv58jiFtJmUBvNE2oOgYwxDQ9IZDYeOAAAAAOgGfVHREf4n36qs07VZIPxf1t+QUqlPNBFREIxt7hE7ahK0iE85r2mNsOC5htsIgigROAWVjVceCw4EAAABSAZ9WakR/gQeKxIeCDSxZxf39Gix/112nHFC1PzbN92IjXL5eOaAOw4p3LVmqz/XDkVqmVp/QwtqVFTXi1vi4EzWJjc3Nfa1Z7l65jvwwWZi0CQAAAKtBm1tJqEFomUwI3/vhUdiXOx2MTspDBHmUaNxnAKtxo1ZidUG16yzOqQ6r6rwou8sSCWAoHPh0R0076A8s6id90rZISEbQSI4RHOQA08yjeqA8L6RD2GtXfiIoCm2bQHe8hkBL3gMquBV9xsgOcsX8KXQlQ4myFK/fF4oEBZAQrHDDSzotzZCZOAyfnfF5fPIXTv5TGCrdBxcbOAJoIYdUVUA1wPS5ma+yJ9AAAAA+QZ95RREsn3sgp5rjMV8XzVP+pp1YxygYLRWB2wZHB+pWeUU3ZYrSXAwe2bahlT/6TYFQTJt0HNTpEP3ubVcAAABFAZ+YdER/inEfdryR+UBQh0g+98kl5ilCOtNo1YP/LuQq8RtpQSglo/BF1WPBxIxyIr5KAxwIpISuyR2pURhW4rUvLdfxAAAAMwGfmmpEf4jgOwqw66AXnAyisbuurz/RSCoRjyFbk5dILFFLRMyoUi9I/yQ284snWqurXQAAAHlBm55JqEFsmUwI3/vhU4AmHk2jvpp39+DVM7fgz+/0mDAgC9sSz81+siWGpUWfuj4GyTOZUOB35xF0hF67DNn0ZuMEXcGIXz00pKtOL3sBXsUgc/ePA52/EzhUAvLLgjOyGZLNGqxerZSlwTR99vdiMIycknIR/ViQAAAAUkGfvEUVLN+BtiMx77QlavkuKJHZZ9Q3HsrRt1QLIHVTAm1R73+pWNsjQ0q58gETaMA5QLuokf/hrA/s22hdm3istSVmd0URqKtuvuIfwOuQF+IAAAAzAZ/dakR/gQeVxSFc8gfgZAkdinyst3l5uX7pQaHSAIaHm6zNCuVavCVTgW8m03yN9WRhAAAAnEGbwkmoQWyZTAjf+8SBZrxwNxE6nKekPQnYXJm5wBvfHrogTYU3NIBfehDfOIiS6ySK3363CBAuzRu6fA4vIPVeDAL86ZR4LtBStB/BM1+Vq+8tYxISLhvBnW8Dyw05HPMDTibFdFLhiDfH6SCO/cg2NYHDixZKTHvuNNlOBA8Eb/AVnrFkz+v5UJdJl1nEtDkUt9F4X7Ko52NhMAAAAEtBn+BFFSyffZxiPbso3KVMv0xrg/tKL0xl7/2ug2XFt9Z/8yNSEkmzVe4IR1kQbVaS6gPVxGBlsxGiPdVg0hy8V1fpzap1S8LNa54AAABQAZ4fdER/h7ECGy+Mf8HMSAoaH/GRCeAOPoTHrElxuxk6cLhDnYyUbEMAwSiUszmoA/R1gKCdAprLE4LysACZvlN3jo3Q22r11pJc7VqKCUUAAAA+AZ4BakR/h60oZYAy4R2w8B/3Wta04n3kCJAW7+Q2gBRN+jJHYMfX5uQDuUKfZkFiLp0cf+1j/fb/e24j8dQAAABxQZoESahBbJlMFExv++5XDgdzj8Rhi1ZsIZ8Sm28CkaHPAYRK2SFSACU5cWdipBQ+vrEx0RbW5ZoDsg0jJR++lEuFdaMABu0GnPbEfBPglHmsDhwytPRmDnxl0gvYO3LQZM1n54PZOnfSIOyUJsG8mYEAAAAjAZ4jakR/gQeVg87ypn6p/FGQB5EJfWf6QKy0B/hS5J4dn2EAAAB8QZooSeEKUmUwI3/74VHYRK6xsWuLkzkMFi5kUGmAexj8Ar3bWWH098HYvsoR3J76TgTz/oDQV5cpDVqE1+A6h12YunuTYfth0SsPDgk0pLNX7k4axqHCb90xZXJ6QQE0rqZvxK94nl6XIz8bJURBNnwVSR2U9urp0GW5iQAAAE5BnkZFNEyffRPrpUzJLr2/AS2nxOSJCtdNcLjn6cXpPNkqjvg+oO3azFJERI3eaUINnWDk1Poc6rrAogWMSwS1ixlRAlvzi6VgIwdfokMAAAAvAZ5ldER/h73RL3LuITCI3oDv4BPtzSRBT0ox9Zs6XKAiQZAu3sncW22aE5GNUZQAAAAzAZ5nakR/gRHB6GWMSe17aOt16nxs4vMoxLGJ6z8ZwihTVdCrHs/gv2FN45XH5F77G7EvAAABGkGabEmoQWiZTAjfzjAzJNh9+MsHWu8d7Cs0acVk7TvkhyFGrIZeP7VKZT8YSa+fqLmJANOx+HwIQKZ+hbChanD7CSavKxfEpeCd3O875McKYrY5v5SCqAR4LZnb/unYHbZG6CJQLkdmdXI3S9mC+vh3J1TAPWl+HZzfy6FGRPZ20VqNe1/AuRmfBkIEA4iz9PfeZIPIT/nIKX4skpc+PdUp5tjAsBGG8DV35+BHBUfG6nHmCnw+kVYk4zIk4Xclt1NYI/JiZT+WqrMgnGmO1KB+XctzpeUXBmT7ef+PRBXYvDjeBflA3MnR1+dcK7H8tqUmzVSB1TwBI9+qN2nnuuRbqrb0lST+s59Y8w9uzAlwZZ2ZsKrcj1TJcQAAAE9BnopFESyfzVU/m33U0DGS8eDHAR1n/6XVZZ9uylPsAlipuRHFF1836S7pQx8uES23pt0I4STTpIPdian6xejkAYoF2nWj27H7EQ1w3xZAAAAAawGeqXREf9Pb2DkjtB4in5wCGYpcxy45arMPjM909IA9gzca/O+GyZGg7wkOdd+VfzvhhMt5NvbfAVFmfDXW5Jmj20M+nvxZ2Z8uZoUmYJ/Ok7BKfgCSaQdQQUP1bKE4bTtcHvsCHyMnF6keAAAAKQGeq2pEf4Z1uDH/duUzFC4JTU6hTUz3R+lQw5ZFRkiy5ri+QAklUBnDAAAAOkGarkmoQWyZTBRMb/kp5kxTMHZmETVTIgsd+A+NbgreHIdnNX9tCFzUmdk/+DL0/JDoveQNd/vCsbAAAAAgAZ7NakR/gQeVg87ypW2l15M2hNNzDWakH6ukSDuxzYEAAACpQZrSSeEKUmUwI3/74VHYPZsdVGEb9i0qy4XOd49kR/Zd27xbxgdkhHGHJC5y1l2os18hgqRE1KaZVj5VSPxxYMpx681OhYazthcfQAXvfCWuIYtRdInpGWzy1sxMLVqgPWw2rPBeteX3bwTAuESDIworEPEl9lhf6OwWpk0mu9bIKNzmWd7gpQ9E8pnmi/jBUgxN2yeN4FcSjUT9R9wbJD4lX/zQk/2MtAAAAD1BnvBFNEyffqpK6p9FFWrePWQt+t023fow08TUcBgf6pBXtRGOvd6vNalSElK/B+orsc6pxa+DU1KIEsYhAAAAJgGfD3REf4e904PVduAufu3efP2dUs6dm0IH8JDwzh2JpxWrp7PBAAAAKAGfEWpEf4EHlYbRiSTrye8J7wrvIAVKDRGHYha2PjTkam//yFWsIsEAAACbQZsWSahBaJlMCN/Hrv6DGHttBKhEf7oitzULKk/UkvXrRGN5TSj1fcIGe9X0COkbiUIPuCUmvFu3j4oaIwdDOZcC6YE657YPavN92lNXaCH3mYlGVn/ZhGTMtcl+H3Znx2tzw6OnG8LWtRLGegqT+9PKN37VFH+g0bBpNjO5UwkVHcgMy/NYkLLWLbOHn2+rwVZaqtmV1BLeu68AAABUQZ80RREsn81bxhrs9aLlZoqp8iOSssZPF6+vuh4reMnAu5tYCzgCviZbUJoMuyNDyQ5sQrG7eozVDcPAOhyEE7YCkUrWX17txueCEEo85uzZ+398AAAAOwGfU3REf57MlpebTfh/uLfAQTkb6fSvSijofFdaiiHHloNZyWGn5J4eiSHd2ZAxcNxMqZvRdlvUiUYyAAAAKAGfVWpEf4Z1uDH/dUm3pglXxGoCPaXPuwJNhybnmpjlHL8vPo6GCD0AAABbQZtYSahBbJlMFExv++5XDgIGEb7MdX9/tJBBP77OrDPDuf6bRmi2cH4plgeyDD0btjDKcl9nj3VzhH2Dh3jpoNa7QLLWuVUen8M08hu978wsuy3OhNJ1Sm5jUgAAAB0Bn3dqRH+BB5WDztJXeH+x0rqoojqtJRdUUKnIgAAAAIpBm3lJ4QpSZTAjf/vhUdge5Ym1P1V/GNX9AIT/jkz3GQqbKNPRZXoIj0uUAbDPtH3Rh3Ek+Bu+4pMcb704ggfye904H48vWEDu5uQg/N2uSES/KbFzLPFtQ2OTa/8qYBpX1/aD+T+6hJJK1ru+MZTghCGmNpSQXWnK2Hv1rX94UdXZenTBkKmXqgUAAABnQZuaSeEOiZTAjf/7/o39RD7fprUF/NAaOCr3pu1+Sc+O4fPzCXHhcvaqDp9XfBNrwZ00TugT9n/WmQBb8BI4LwStzpt0RUudpuht1/3oLhjOlCXs3ObsVj+wmRIxGocRDVdQfy+08QAAAGNBm7tJ4Q8mUwI3//vrPjgdzj859cWvgfbdn4AHc5U1oFwreoYsjtgTR9pWZhFMOPGk65+kAs4C+VojpHiEg82clYzEV+LFhJqn3NZwdE5ihwYUpFp2IKdxJaXvho2ltDn86KEAAABMQZvcSeEPJlMCN//74DjgiV1jEgudcPZjzgKu40Dxkog8ar133QofRc/AH/ZjqW6NQl1fHvbLAdPdJrabUaT/78/X6QmNkJeXNfxPuwAAAGZBm/9J4Q8mUwI3//vgOOB3R8fPH5+zlRCyD6mVH/aMulGQkckScEuTaTId2v705sWz0b6RhD6pVrrfZlw2bn4W4jUHEnaKMk72slzxowCIHDa2FHirxY2J4Aogy/1G71Xvqzs774AAAAA5QZ4dRRE834G9C6jp7AxObQ1D5nh6xcB4XoXLYDo06V5jfpQHI0j/RblPPLPhUSYgXonQeS6joQ+hAAAAKQGePmpEf4egOXq98g4uEQsyFmb1HLaC4k+C8XMK02ZavI+WCx5//WHAAAAAQkGaIUmoQWiZTBTxv/vgOOBCV6+DEJpLiNMEuxO6P5BRQh2K+vNT+Sg5udE4JTIIQLzSP0bwLFYQSRlbfZFEgNKAxAAAADwBnkBqRH+HoDXutdsAXaYM5d5s/ldvZ1ZUZqSPwmL03HPhv2wTpeNIW0aHyufZg9R3ewLwGGEJnNYKZsEAAABJQZpCSeEKUmUwI3/74ByBNNkgWomc9ochzMRrfD4ihgnoRApgkRdI+usVmu6Y8v3e0UScczsCjQjsNfzAoG7XkUz+H1yuPskDuAAAAGdBmmNJ4Q6JlMCN//vB61Ydg8gqaxeVjbcbFf2tmleGxIAMrv7NFEHar+eha3FuY/vTZzPKjJC2Oq5sKE5oNwUknilo4J+t1ZmDCAqKelikGKtqi7gdJP35m4Kn+yOr7aYtXrd6AdZBAAAAYEGahknhDyZTAjf/++A9RDqMbrkil10LEKl6Ttr7miP6kcL0ohDF7+vXevrzG+zTLebM+8F8Mqm1WPiQigJSclUc32sO1HgWuzAPWzc3pJHtUBg400W+ishyY5Uy/xEMmQAAAEpBnqRFETzfg/tuGYHKjS/jGBEY6YCRcT42dm5owW7iDaE6eZzNyRpt3yfTYmI4OxsgAkcyVuBGb+F3A/b7UbL/gPEwyEc0Vpa2lwAAADsBnsVqRH+BB5XFIVzyB9wJwDYp8rPxDk9qpTEARkTv8z3S1UGE8PcjATyV6Zz8GAZwGPHIMrJR6gwBSwAAAINBmspJqEFomUwI3/vBjyyXKEdXz5LJiOmapC7Dh+bqzRRczmQAYnOJBgYTvI6vsixxDeHUOwDiuti47UC86r13T09kIqPRmtJYNyAq1tP4/Ft/8pL1VrpczUxkwqL5a+AJdlYjtK4aB1EcbZFat2ii0Xh5BXeTFj4VS6DQ2G/nJLgwVwAAAEpBnuhFESyffGXEt8MjyOXZ1Br/nR/99FrXZZadcoA6YG0Znt1Y61Fwz4DJCwY26xDhlnhFmPYwjeDf0M+x30DUfLS/zFrIDmZ/MQAAAD4Bnwd0RH+HoC9/9YwBqK/PwA6Zil5xHhfqFkRqWZ06PPIdbHE9bbxRGEDxKdUaS9iieGaHN/QijO9ByAGQWQAAAD8BnwlqRH+Gdbgx/3Uy5jiTLpXihCuS/L2kShyyEDBZR05rxcfD2VGCwXsTh9INhGlzQp0UwAxUx1vZR1T6xUQAAABZQZsMSahBbJlMFExv/APDMcBAwje62AS64edylmitQLrKTHGHtcvoRS05r00OvVkJCNrMh+/mVw3eQxDA3590PKiZB3nWfrN5pDxzTrKYWq9sI5F4omOgfooAAAAnAZ8rakR/gQeVg87QMgUIqzoc+IQ0ccVZvftYC2gzOa0pG5p1idXxAAAAbEGbLUnhClJlMCN/++FR2BsEzfe0cZac+UJkdgjTB26Jig5W2j5cZ26NMmefLD0RqMng2y7laLakc6lHRdZNw1oHaHFsXiqzatXDN5Pz0v6/WwXUQ/oZM3evQzDrkzK2TYV/MjXfwp98+OlJJwAAAF9Bm05J4Q6JlMCN//vgPUQfdNSB0ivMgv12kewdPs8LDHEnccStTFGw2T9LqoRXFd4GA6Lmn/C0Ae+2Ys461r7jzbWD67darQQW6fpKD70itK+HYmaanX59v4Kvjz2TrQAAAFhBm29J4Q8mUwI3//wDwzHAd5nLgo0KbNa84paAodp1tGNxaXmvV+VvBFtu1np4hpixdq01/tsJuXnBDpI6pjgYmOuYM9uC1Czm+OIeAHPNqzeKqP5ocdj4AAAASUGbkEnhDyZTAjf/++AcgTTobJBaHKECxQ6c79GRuRfeQ0KwXvcejjr76PM0aqZOg6nHibX9Iua+A36w2f//ebnY1vV75tCTbJ8AAABpQZuxSeEPJlMCN//7wZu7TxYHfZLW0nepgrO++53cE5t5rTbH/j4CuZw0tqBjVHvT/NuCxYAqZPHuW+4z/LgYIOyIRCJSyOgmlLVczhDzE5KkFP/GXuQLXa8D+P/5fNEgPj1lTUpFVO4xAAAASkGb0knhDyZTAjf/yghD0KSEAj8fRNDXd8ZwkGDcdqAIE0puIoVKKR11TsEWfh9IlRceiw1biBwH50bd5SVeqveaOBsfxwvHo+S/AAAASUGb80nhDyZTAjf/+8SB2aRkB8IFrk6Yu0T79Zx2yWX26LbBct6B2P7nMJlSNU+8im32OIA82qWr6M0WXr/wOnu9asEIe+TzxoEAAABgQZoVSeEPJlMFETxv++A44EJXr11WI/mWGmgjspyywfThqWu6vt3f6IeutLnnxyOVezpxz66ZO+XTJSNz+hSyTccsijssaFateiNeQeLGmXnXJq3lh1/FYw7UzWsF4lLwAAAALQGeNGpEf4exElW+Y0oRl/FZQhP+oVccc8fw6tVIVhtOXYZd8Y1R9+768Me1CgAAACxBmjZJ4Q8mUwI3//vgHIE02R83+/Vlx6PE7gcgcvw/dfZOmqjB/6fbU3T8uQAAAFZBmldJ4Q8mUwI3//vhUdg4BM3dCQn71waETfG5XS6Y6qV1OmJ2ZI1/p1k7hngqfLDRlYxzHulMp52MY33dNDvA4xKrjGSUX3xOzqdXe+MisKohr4vpgAAAAH5BmnhJ4Q8mUwI3//v+jf1EjId8qaxOH+NL8HxQ2fp5PizkRPrB5xmQS+tbXXrCt7darNka9SvAVPptUtLJ5w5YyvLUX1kGEExPUgNHTM6K6C2ru/0HOEs9PXbEuaYln/GuFzif4kjrgcvg+8JdtRNs5H/bTBw0pf0EAsoCf7sAAABQQZqZSeEPJlMCN//74D1EgTaKZbqNH0joJJniUx8Ruqc4abzu2rURnTM10DKgg9fFoJmD440LTpJLO9oCBujuWU7mA65jaO0wi7HremRTBMEAAAA4QZq6SeEPJlMCN//74VHYnWWrD3GQIHk5ghig/U3DKJNeEKm/nL7OOuV9SKp1OrFdiEZQH1cmsNEAAABQQZreSeEPJlMCN//7wcOsoiIc/4xpRxHtkU2H2faBzAZluyLYQW6IF6Sg2Z8+2s6OiVvP7unG/PaKjnvHx284RVK7RkYgtGcnC+eIrXj0kIEAAABCQZ78RRE8n309jmEtq7hLzM71lTx1/0xhOjUhRbpoPBZse8D9DOokAeCvfYnyUFigFVOpER0UjkrFfP8CohPIs1WAAAAAMAGfG3REf4egL3/U68Y/kJfuO5of8ZPB+vRdik27eJJcU3fOYgVo/ofNZHen2wm64AAAADIBnx1qRH+GdcG6e+81iB3HSdBA1qVClucbNobVh3bfiI53Yzn0y9s0dcFwAHbyPWxxUQAAAItBmwFJqEFomUwI3/vD26eZCCisgtY7Q4/nAy8diTjLXaDcRhbm6PzfosfVpvyYRHRlUjwp0BnQ/NWsB7lZrV0Z2KMWnG+8gj/2wN3ueG8T8c99L1NBE6r4v7E7tCr/OznxbvNYjxfXhceW8dti/xTVY8P4CvNMCqkTDHLfY5PChguysyasokCazUmXAAAAKkGfP0URLN+DOtOgnSjGpTM5Wdix7kUjmGwL/wc5dr6DxK39Y49jtKsbgAAAACgBn0BqRH+JNh+VsnD/m3MN5x6Sg5Qu8+nj/5HUWhu2O+TNUelAcVFhAAAAlkGbRUmoQWyZTAjfyghC9V+ELITH3LCIHSIvm9CG6OjHlL/31AO68jyrK9IXLK5IdMok/VYaXk303UZ7pRLjQA0f/4qS5bfqL5QiXvH0ffwKCt0Fq2vBbWwZIOms8XXulZOVX3znJGj/solkcNIBtwjIWnGGvWEb1puQ7+6i7LeJ9OTdMX37A9Yfx9MPio17AaTThn/g4AAAAEdBn2NFFSyfeyCltp/KzZHjCx/v+MwPlO+6S0X+qGwMBBl7Zd3aGLZ6jXXcpABaEUHQQndnlecT0CSIxETMb0IEO8QcA68axwAAAEUBn4J0RH+HsPlaZzJLf1BqA0GQv7XWJY6mi5DH26gpsyq8ggTdiqyPYNRnn3ZbbfzdoybbOog2JQKGI36v4G9OFjBq7XEAAAA3AZ+EakR/hNQrQyxiVMaSwCvWnQsw2ksmBu2UpY5bkEhGHFQrC+fPFAFNp4o85UVgqnxqzlGAgAAAAI1Bm4dJqEFsmUwUTG/8A8MxwlzCIUaKgXMXLU3QQJ1EpMKnk54cIKtfRAathRBVUcjjYAmeqTvTihh6Nuk08XkauvWPadL4Vijodg4AVC739ye8e1iwE/NwJHDIQv/TOUzJWUNQojoW8MDz68svtIvZyoIhTX9vvxCDlonewBE5q3SspGmeeozvc/Ou3OkAAABDAZ+makR/h7DkMU2l6PKcHO/zShtHBIsYylwPXYPe5TqpRiePe1Ts30NzxAzEqW4PxaI5KV/5p/0v4lsSu94p/NesEAAAAMlBm6tJ4QpSZTAjf/1jexkZYO6eHpSGJJB5secePzct/DA48qcDfs+3F2H5KAOIzt2Dgt7ew96sGOjWUYhosJfHZkPXPvBEDaNms1hmlF9yL7QrVVstFpOhrOq4+qlmyqh/ldeLSeBIR0MDaNW2T6lc8cDTyqgcYV+Id6ev9tWnX/Pg/+/3Bx9Euhw2lfQiFdanPgzmX/55qCVEPMNFYpIBwPcaLU3yfrivjcJks8gRYsIJlknW6guCjN+neZa8hQkWtZwBla0sYTkAAAA/QZ/JRTRMn4fNy4UC4h+oc2BifKvWVgEVteOt6GdcPcqYcktVsScBfm/exmQJXfl0JC/xI/s+EI2ZfJUQBohxAAAANQGf6HREf4ew+GS+7UpFaHk/4Emf158OcWu01ywvQyMU5k4T2GGCK0QXFWRrnznHAUmOJaj8AAAAMwGf6mpEf5Kigzx6E13On1ZVie1EO84nBtPPa0VWD3aS/Ddk0ZkbJz8T95L6crCgsixKpAAAAGVBm+5JqEFomUwI3/vgOOEGsDnbEfsgdrJ7tKwX8nHR83gyULQVhL5x+AgWYfj7sj53F1FklU/dfB0NntEcc8LAeFRHrF4HF/lXkKKa3TzSgRuUANQcRW1CFfPMX5gBxfeWFsPEkQAAAFlBngxFESzffU2ptw5lYiNEGv798jp2yP5EhQTY9tWckQ2I9YpqN5ieaWAqdGE6ASw40+H7sZMawfHoCoDply5HZaD4Q9TdMJXThHdVh2fQNrSDXLaofGcFLAAAACYBni1qRH+BB5WyGKseqwNDLE8nxiM3gCXD2Qs8x6pTnxeKoYIGqQAAAHdBmjJJqEFsmUwI38oIOGvYtN1rSBTzyT+MoTaZxH/5QP0RwoqkIo8BIYjA+yAn/0EDCt8e4sfiSZAjYvAKxVNatVPYvX26/q3BEA+lCvg82wYomT1/sm0oujidy0OvyL0i4I1stZMMRcaxZssLhgLuUf5qu0yhYAAAAGNBnlBFFSyfhIT8MY1eZNHsRSE0h+4P5gINzsuqj6GnC3epFTfnf17zrDrtInbEItFMN401WbhoZ5/w3Ux2r08lE8/dDeO7O+HJXJlcFYNAFuyz8daGCaCCNey5dBaMTuealfcAAAAxAZ5vdER/h6Avf/WN8x6dT4hVMDgeRaCIEr6D0AY2ibt4Z5BJQpnXAfZG9fF7OQKPEwAAACoBnnFqRH+Gdbgx/3bkhObG4+6VBRfnLnVvz5ibbBPbDI7nFtBHdbdYGYEAAAB+QZp0SahBbJlMFExv/R5j/x5hxQqCYDiX9FqzOLs9ly1QT91Cp7pRDMF2Jkxd2wtt+BpFU6ZBrmxwLfJiwLCcv9T3lq9NZ2wLCT+/wnYYaZFCO5IgouVXlejqeKxfunKWRud5TjL7EHetewgMnoQTi+FFP1HpNOyNxxWv3UwhAAAALAGek2pEf4EHlYPO0DFagf010Xgc4NtyHA6rTfxLOyTg0Ox5hTe92mTjN7boAAAAg0GalUnhClJlMCN//SaD/44ef0/kxstvutqsDDhlXYfdekHXbmlDX82+t1ItqO3M/ML44EVCP/8ffQY6CAnVY1XBN5ZbnYha3+46HPtD5SMAfNYPV3RND0lVvOS15Fm1CdExPBODnBhUkaWtEAbcewQEoOOq2Hen6odimOMzFQ5d9rZgAAAAg0GauUnhDomUwI3//WKRu6nVyVxRee7XpiIxlKBqJvLZu3azo2YYPGF8qjUhlKiymApjcRRZlTf7Nef6tukffI6Kd7PrbAcHwGvELaeaP0ibxo0QBjNurmKmZaKkIE6v7wE4Wpmgpar7VBmLnzu8twOoVcMvF0fNF0fSaxrmXamWCb3xAAAAXEGe10URPJ99biH2oflZsxy6BAxf/qvJQn0UmNGcVBNMrhzZROK48G5OjPWtBqH8efnPpu4XSt3r1R6HAtsxDtSSxVrh6V0/kLJR05IiNa/ZWgvQYnGCF9qmzKBvAAAARAGe9nREf4l57VUKFwSmPSmlWex1Pm065LEm1OKxbJb7BMmnONzBDc+f4xWPkczvzmTohRN3zIL6Ojmccwe1SM5TTAzFAAAAPAGe+GpEf4TUK0MsYlIRurtl/BL2J+PMXABKR7W0r/E6t29wkdoiEyOny9hdlQSa5XVyIlDO89Cpm0EBTQAAAGNBmvtJqEFomUwU8b/8A8MxwlzCHtLBjwzZdybRP3L28+w3h5H9q83W+4R7Iqot5pTdotVf866A50TMubMMKbl+05ylorAZNPpFCLdQn1pTLuSgnmK11G/Uu/D96ElwYRQfS58AAAAzAZ8aakR/h7DkMUOCGBJh+lPcfvBfkdYcocwBe9qmP8DL4C3VmEv6j6IDX8OhLFl8JqBdAAAAsEGbH0nhClJlMCN/++FR2Inm+ENt5odDIwW/GJwkPcJuckK0JIaAU3CdS/tVTXZvLR05nQOjqvn8+KStPWBFi3Gez7OgeSNwcVElqJE1rP7gdMNhbcN68HuvsKSeQvpaJEnYw8wvyoh/QW1CYyoGlOno9K1aJIr9Pk0Dj8g8+NGIQjaqBemVGjXOrl9UV8iu8KodQ7jGbm1Zq/jwidOehhp797VsqpvsPYf6ZBDtDwWEAAAAYkGfPUU0TJ97IKea4yjH9NU+Ki3CHMkcSJ9htM0rapvgZK16qh6u5dqOGGAOrakaQCu114AMCSgb/V52v4CxfpXrn5qv0WiKgUPDhAE/bLxIRa7EHZrz32McwKF2qrLBShrvAAAAVQGfXHREf4ew+GSYyDhPzPlV2bRuCDR/Fea7xly0Zv+RhCiwpzapFX2nVaD4D0ZNqSzT0q6pbaI4DoNcxC/GWSw4XfnszHz8xeBlADF65GFV4DR+M3cAAAAoAZ9eakR/h417ZNpp1xtu+fafR+HC0KcS40oDskcFqn/O0sQuxpu7lwAAAK1Bm0NJqEFomUwI3/vBm7tPFg9jGVWovZ0jmu/rAuAvcW7ubK01SqdOcWBIz02mXZObB/CLB6HlDKDI+VlNhCg/6zrUx+GcHA2xWFHBAfdbhNgkXIqJBYT8FRKoB8Q1unN40IlItyXayDFwzBeGhvGI/KN3ea1R5uXGk1x7p1Y9/VGM6iUzPd/IwqqDaaStukEQBdqZPJMfq1T+2RPjqIFRCfd5FvxLcBRYuHwYlgAAAFFBn2FFESyfeyCltp/KzZjl0CBi/IEonyzMc3bZzsIWT4CydiIbfPye6ei9io+At/mFoyePgIN2xEzlQ9KswTRayQNCfgryEsw6EmQWr1T0ADMAAABBAZ+AdER/iLlSgg3Uqa5kstRvP/YJiqA7KR90cw3LNvcyMbNM10dPyQ7u7ZQ7Pxn1vA7iWSg71Q2/hzofKx8Il0AAAAA3AZ+CakR/hNQrVcoxKlrmAEcm59sbNzD40gdHC0EziACGmA+UjRF1XzhaI48Z56F0dPEF2w0kgQAAAGRBm4VJqEFsmUwUTG/77lcOEuaZPnc+jWH7FzzFeNiG18puHuQh5C/ueG8rR5oKuQk08AWgzZtU8fG+9U3UaZPtsL8BReEOwTp6VR6Dx/WI+9cQyMaZkUwAXj5v/UdnByly7VU5AAAAOAGfpGpEf4ew5DFEIIgTtjL/DAkM6khmBhIgRPsTh1KkB2WPl3dHMh6FUA7iA9fty/cEHvdpGr2AAAAApEGbqUnhClJlMCN/++FR2JyyC0hpnDh7PIxPg6BMTLhBL1SK8o0CDfiMpJT4kRDl9tdZ62EidOxdhN0HnNX76b7eC69SCDzmuZwOgLR7EIjB4kgvtAS2pqCJfoNyhrvP0D+dIb+QRSH+77hmtbfE1K1QnwSPb3D56hVITusU1SKKSHGmPML+xyt5i7yztzLCuUxQ6GBMKC1lr+K5ma8CRUoFzeYFAAAAOUGfx0U0TJ9/cI6jT+mdVbvBbfqe0kP70WvSpDFVICoTveNkI3cjDw4DDjaPNgfOu7y52O5KIIit+AAAADwBn+Z0RH+Achjb2/3sCtNWp7oUGbK5uD2avL6oCzA40KUrmSqFJOSULMgMJ/dZcQX7xKN4kdd8aG2wieEAAAA2AZ/oakR/hNfANpNdzqjAmab/E8ax6Fvug8RPvYSSLNwtDqJx4K2Zi3etFf/xJFtiuMoa2OKpAAAAZ0Gb7UmoQWiZTAjf+8Gbu08WD2MZVh8WrSpRiUBZUWy+1Z3OOJ3h3uYEjEFhMU8BzaG88lQfmplKnzW9sLz44+YZ5qjPn5f/4kD07E+kjAoFn0kEueWpFWqaK18eV3GeAKJtOJ2gcdIAAABLQZ4LRREsn31uIfah+Vvv5ElY/3/IpzyyoeDu2znc3BzxpsJ+RQ8j5+UBxv8Bc+yc4/gYr85cjlE67KN8CH6/crTj4W8FigNzW7tQAAAANwGeKnREf4Bw4zBVesMI5kvH//3lWiD0u7E7ycYGE39qDZBe9Ez4XiaOH2ZBiU+cYRKmcZPbFu0AAAAsAZ4sakR/hNQrNtFc9WisO+KwtEafX8vb6ZXxEB5GdrbfwSZkZhFQA45ylQkAAAD2QZovSahBbJlMFExv02+yx/CHKbBQXEKSiQ8/bO2xeS1HQBGRs5Hf/V0hpgC19PEbrC94YtsEONo7PCTkPWJcLGlhvRVXDsvhZchHblbTte0GhtaJBAsYRQtNVaJ8kzaeIHkxTXGujtsMr89Zq9hCEObGEdp9ad+G+69ZLh+FJjSEVgOkHkH0PEpLJYyFAvBlFP2jIHUVm3qm3jb97P4pJCxeSwXqGAfuihTU/RGSLNPwJun8vR1Ksa+vR5Zon3K2Nw9i2s9ZxgoBryAZB90K/m6pC4pSyLW9BsQ5Qz/96LFgNPiTaKC6ucaXERHMEYPOW9aNn7YVAAAAPAGeTmpEf9wncM1iQ9Gl0i1OL+ABc9ME6axttF81n1bVRz8hEXRQGH4rWzfbm5v3nB61lX8Ctgc7vz9QgAAAAM1BmlFJ4QpSZTBSxv/CIPOnQv6FGt8l3K0kX55Ov5JRlQ03g0N8vQOD5sBbfTPpEWH0nyx5HL+fie+qID6xiOC1xEwq8zJdI2VFblTxxBMPSypkRcUoWHmHXCxUKRVZYFMjeGgiRV+6kZS1BgU0P9oD/Oidv7irP/aV0V26J1DBG2+EtBjS4lHyivwXdip0d1BQfBGsk1ugFW2Fy8/EjyelHeTPHBgZ943+z6bcQ1EJkW0tMaad2fGspFHI0IbGKScCPWX/jQ6iT/ojIy4xAAAAOQGecGpEf83u2UnV+7UnY1O5/oQn/Y6AF0MpXYTzrVvtUezCnzx+JtHpSvytH38dTVSqrrE8zYKZ6QAAAHhBmnNJ4Q6JlMFExv/0LB2GXt39+D3Cj7aj3DPZ6Pd+InbEtW1V3J2UHsXT1v31OO4EXiBsko5smAXJphYrW0fTkQJirwR8BXL92Ei4QvU8E/9ALgLDiPsylWFhiHKkipiOyh9H2qwcv9t+CUifmj23yzrIU/BAu0kAAAAkAZ6SakR/iOA7CrirItMfJpRo9cnv5zDAP6itv+Og6IlbYQGxAAAAnEGalknhDyZTAjf/yfc+3zQHuCTSl/U7GZJy/gS5f9Cy2h3VQBHdy17hrU6xaDgmNvuPdO+3nvolJXmlwIAkxED8Lx4XeaN4VE46PRHNsXwgF9trLiwzDBPhYZ8V2tPERnsdMtjrFEmrEm+drqLgNlW/gCWzTU6jf35hzyjGb6uEJOQlcGPQpDxqrtzr8Nc8vp4B1BYjdvjrlCrq/gAAAENBnrRFETzf0DZ7oRVytJSzitxXMmJRFhcbeha5ual5F1I2q2YB+zAT4PLPb+JmJ0ARNc/AIxeBO0WVPQWFlcQBs+2QAAAALgGe1WpEf4EHlcUhXPIwsFKRCeT4xUX4F3Q/slkJ9DbNgs6Z0J5ivSJcp3D4p0EAAABWQZrZSahBaJlMCI8bzhd8dK3gZ46kOJN672WFTGEpwGtbDZDet5dEQMHNg5ySdNRqv151JHg5eF5RqsxWW7PejJfmZTCQ9Q1IN/bi5e+xWsYriMFYv8AAAAA5QZ73RREsR4etQctl8Y8CYRFjDEwJCzxPJgJwzTf+tolyI5ByCNErI5NnAqLN6VSiQ3zkMiK+t/rdAAAAKwGfGGpEf4EHi7cJ0aWsBD/fwiVyP3OFrw7+rNjWzuM2AuJ/SSLVjrYN/7EAAAT1ZYiCAAn/vgy3mIDfV2I+c4OQeUCbjkIqauQ7q1hI51t1KK0Ty/hpqe1v3iAKauIQOhjuzFxyMgx7FbnM/2jBWLSkwzPM9uWS0H4jWSakogzfe9txEs4c+/QsLeTPvLzlrrhoRvk76vX6F462KFV4lQ8BLFX1JpZOU12xKsmCPjGaaOJgyXQiOhYn2hWzgcqRNjGX+X2BzGn8mEJ7aGMcvtMYUAkI8JdRks13nKu9bo7rrtiKbOLdcmLQMypUndSoHj4HIx0L94EanyxNktjsUH6UEl/L/6OZSnHAINvBlAMWXlIWyt38JXD0JUnOv+PLJ9lXJviKQHF0pLbebPoF7UwMYacH5I9nuT2nhJTj4jMvriexQBuncAXjupdykvxkUgQsRr6KKam+klVKsKAJwzYn02pUzn8ZX99Kgp0Zv6A/4AZIw+tqzSXaBRvQyPlagCtwxNENMBXBGJCaTT/M3d1NYKMBuEQdSzG/ZYjE1JEavYncvI+HARC2j+WXM1aZgz9ZL8CgLjpcv4/YrSBFGs6H2gmgzdANfJkk5Ly8dYyGKrEOA+lfembLO2OFLdnp0Kca5ZBG4NPOmAF4BBvDzXD767tWmEsGNf2xMnrIGj+8i72JWEhZfuKA2OO58X8ZDSx+O1JqcX7S0VhG6hgbqU5ktHE9a9+Ve9c4kGo1CPjPl8CTm4FxZbQjzsmwP6HqcIFYBDPVBZ9WoQ+2D7QvXIfv5HX1K2NuPVZNm9uwFtcmGsf6PuOhqK1pWSzqCgCmwApkdzlejUme6mZrR23REJ3PkQdaAyFo3gfa5LdW7PmNx7jcVqBgkY/SNldaGv+tW+xdJQrKAgqxhGSxiG5fPPeb6tkPlvE3G0CAEuGghWptSDagus6PiEnU/3yKJHGviaPLkzl10/T0zWMpIJouNnIfO3N5LeJK2C4xONQHkEUvqEC0YiEq/qK/U66LBsyfPc3Q2VPtQVRJkl/If0X6ijW2cTcVVwN2xTAqMA2dSbXoMPK3bH/VU7nJqjh5LpA1vIdkAYZ7+xljSLiYHT9oE2LHNhfDNZ2s8GSgeEYKD1VirZxu1ZbwqSl6Xuire6yZL+zwkbGlmIlsAcwVhhOfhAfiaIh4/+ogejQgeDyfWoDmEbXzSAXUx7RmQg8IZIWOZzcuX2GexReoXjhzAe2YabjMhqvs+FSJ4b2KP30cggQmhKlwsFqD1ev4SKHH0Ge/jrjEKbN1HzQ3oLtBg4GpkYoRNfImFVoXHLFg7l2LgETKlwJhboxTHUPGneYTP7AQ+K5Eb0/uCSO5Ymq3DzI1gmuGU4MVmzQSuYm/yP8e0MvRQ04RZiZtS50qyxWMrIOnk+IJ7NV3E0pPM3+ZSPpWftqVuEKawLfKtVnt/jTfRjniO7HM6pAY//32DoinfYKZ1Q8GPwjncAQdzvOS+csasvxPzi3DIy77DGfgyPztLfczbeCwixhep+9Vqr31F49wN+AMxHwVCsOJIRkR/wOscgZTuowD55INTwrue5rjZiT6CVfcKanAArmyYb3q8SvJB+n+cPqmbhydgqPgtkHij+7EGXRHJg+g2fpEgj5bDzmfXOTO2YG2fuHCeugHQkpDQ4gK1xMFUb1525vNMHHp0iO6hZxgePKmk6xBwKiD+jgXwh9euxUHrTJXmXSWmplATtcnNxFTI6sQuXnJEZ2JpBKsD+2BAAAAP0GaIWxG//vEnVTkBBzwgdJ0ZWNF1Rhjd1kKxuETzl8V2+eU76R66MAyhvqkmnkxp0vkl/+a6/BSJxSj4/j6swAAADJBmkI8IZMphG/74DjgQ7ZH1Lrefmefbtk4wS5Y2/x+DLBweoV7XRDx2U0ZTPS42sTXUAAAAF5BmmNJ4Q8mUwI3//vBGT5YQIcSq44tpvwtmFs1bxuNY3/v09IKvTPLNo54VVGqqjoQ4KZPNJThwuSr0wg0X709s5PGIgRr0KCsOg/kGruNjLFSrx/JisqEekF8neXKAAAAm0GahUnhDyZTBRE8b/vgOOB3OQFc9AXcbpj46i15PbP+Q9GQ2tJqySrnQwI7VJls29RfMrHbfsSK8eX/l7+18PjO5gb4bIXLCoP9kNhUaOMVrscc/CQyFS2DmIIbp/qLfqVQA7JSH5O/IhV43txHuIs879XU2UcXwNJG+MNUv/9SvID6qp3sUX+AyNp07MsYGwo+5nUN8BjKM9iPAAAAMgGepGpEf4exE2/rMWYsUC9Nc+P/yRU8UOqE44EY/0tBUFtXD6vVpxSXjArlMIURYAeAAAAAXEGapknhDyZTAjf/++A44Ilb6dinlSkT63y1KiWQ/IWMp8xWpc99kAKnB2h7AefDWJA658UxgGiN2hdpqZnrxVEyqj01yCAccBMH+y2STVmGjawy5zlKkJY8a3FAAAAAckGayEnhDyZTBRE8b/wDoisgQ2MspyG9MTL2twAnx4YO3N7dBgMj53wgsoD7fgMzpRUSBvxUyvU2F3yMVPcr9aQehYe7tuqCXqk4O3XxyB/oozQG6cLwffJ31HMNBW5ShOThCQwSZuY7d4iRHTfMsm/D2QAAADkBnudqRH+HoC9/9Y3zHp1PiKWpA3D3YoIQK4j0M6XxrUNNDfPHjMjM+xDhyvcRekwlGLKw4c3efGIAAAA8QZrsSeEPJlMCN//8AtC7wRdgOIxYG6HICf0tt8EYRWGccTPXL9ID3ANrHXjlUlFMeXfH4RGJz/GNKEO5AAAAOUGfCkURPJ99K3q+WhOPVadJ0Zn+q4gSRe2/zKKvWouYJMdACf73czGhSzAbCZzI2xgFyyWnZjV5cQAAACkBnyl0RH+DDJVK3gapAzXdW8ZNJVvmI9scX8y+ulkPHeyLSaA+eDUERAAAACEBnytqRH+BB5WDztAxKT4xFpAt8qSdgkigXdBpZgOPffcAAAB8QZstSahBaJlMCN/7wetWHYNorkyweRiHOPPdcXS1IITtq7AezlT6fweBHbowbC7asO2PF2/UBCbMgGRBUVHraspQzmQA6D7PUakKIGCmWZK+kVgi7Knv/moFsXlEd8QJPGaF0SHBN4ELaLTmKwUa7PPSxV2g8M1cOkFDggAAAGBBm05J4QpSZTAjf/vgPUSLvylvSJaoX67yRc1fALtzOq+SiE4b8qGjclrzocoYChnLAMbHgIHZNX73/lsLDOMLY4x13csDW904+9bFMORIwShS9SrF2hIu4YTLBzKObT8AAABsQZtySeEOiZTAjf/8AoC0REXf8qqj003CfEolq5ZHPeIt/nTlRdl6OlOcoVNgvTKHDmyYVBe5RD5gFrfZYd2HkGO874iVFT8MXHBHECxLA88R/Dor69i7i2G9sMY7GGZAq/0gNA9/XbGiCiL1AAAAL0GfkEURPJ99PbKAyIZQ79YXXIGQ9Zn3sXohriIew4GeDzxFPtkfYMRAykrz2oK+AAAAJAGfr3REf4BsKuKQrvnfCNbyOxT9/LxaMU530JspiBe0HCmuIAAAAEYBn7FqRH+HoC9/43g3bWZ3P6qT2ZqH2A9FOFkkCa9OSnFIbR6HJHG3QgSA//4qcrFGX1dvdEywSR/smSgEwGCmbp95ZNSBAAAAXEGbtkmoQWiZTAjf++s+OA7zHfcTr2WCBetLHUxD7/WNcoHXS1dnHb6Gxii8D96x0pb0d6PjC4DrB/WhVpxnxwTGDhj4JrHO2bjjfg9OQ6Gfn1d+m/OSL0v95OrEAAAARUGf1EURLJ99PbS4SuQriYYXJQH/fi36a9f+dXe/nbPtQtUH3xb5/YCrJqm5SDKlJfww5Fp9cPYaad58GXqPyd3Np0omTwAAAC4Bn/N0RH+DDmAYlvA1SBmtJcJQ6ephfzmo9IkH4wmc2d24zXHcPLKh+Lnw0TzAAAAALgGf9WpEf4EHlYMkv9E2/PMkyg5XBF3hcZ4c77PtedhhUSll/MrYkqHYxHIF000AAAB7QZv3SahBbJlMCN/9dtVqlh0uAaN0ofVjxMhZ1h5qorEdatiQ783inAG2LXHB1MiAuY0iagySzluM1H3bIodNuek635EEPQurFICs0REgQRMJVUluT+40C3yNiRn0C82e0RNNjNdsjWBfmakWZtlwAWD/Sn+XVH0p6IfBAAAAZUGaGEnhClJlMCN/+/6N/UQffCtxxis5+gpt/z2N6/IRl2ZC+fkvRLr0komX3ouJlTkB4RiYtyBbQy8957b8kr1rd495KdRJpeJJwYRiht2s6g5ssfWHOxE2mlZuF4E9aBSO+RThAAAAYkGaOUnhDomUwI3//APDMcB3mPtur2L1aw7/475bzHjZPyheciacbIxHbZ0dVmofkE1MGw1npgtJAA+Op2CJi6U6FI+Fd6w/4z91fbaUzx1UjzZgEuMaupFEOYn3mjpLK++pAAAATEGaWknhDyZTAjf/++AcgTWtokUg3fR5ZBhOWluG1KszNLcF72sEzcEsRw+v6kcUufDwbKgmv+8HgVhfNY38jxmDgIOVqtV0QZJ9wkAAAABZQZp7SeEPJlMCN//7wZu7TxYECs364K1K9zOpGtt1ZZOODFKfndyW2uhIDiGAJZLhzB4/6bupiaM+sug/+57cXUnsnlte9DZAA+U0bGdXg2VWdQHryyM7bKQAAABZQZqcSeEPJlMCN//8AoCu8Eg0IAKdPalTJwL5k3brsQ/2kMhr5e/LMbxJi9LMuGNB+zmBqO816m+p0J54fRMFrIfwddbFhUF2M8i5i2yKkZU/u9ZCP8OKKucAAABmQZq9SeEPJlMCN//7xIHZpGQEF4WuTpc9YQoczh9M1Lmum+7ob3DdhJhUVeNksMkmbU/aTy524uFw6YjqhrhnvYI9JX5S+eulWATUNNzLVA87Eb09ZoT9jApBuZ93/euYFurmkkWAAAAAhEGa3knhDyZTAjf/++A44DvTrlMIVpBxjTybeO80fsQOd462lkQixl/i8ixhGZjblZ4BIJj50LvODBaMD4t74KP0xM1A0shJj3XDsIWRvMHUMW2oF+/62kSLqLr/NQ6ysmc/R19+j6tUcDuLJBR9j9Ks4GTX/cn3YO+gyTrCwbIURfvNuAAAAEdBmv9J4Q8mUwI3//vrPjgQlevXQoH45WptzQNYqW5lvUt7k370Zb40x7Lqnj6WT8/G239QKi+rn9bt99F7eH8NCe9EpQ0PgQAAADVBmwBJ4Q8mUwI3//vgHIE02SfYN8H05IGFKI/2FJ4/9xTz/6bnIt1458gQ4K84mOGQcC5N4QAAAGVBmyFJ4Q8mUwI3//vhUdg1hSSB1aabIlvddm+MomJ80yRI4f09f1qYnZFqKwTK1GeZgBhUqgBV1BeNxJBtS3WYBxDo0ZtkTRCh095TeDwxQPJ/8LwJ+/nKDXqEa7WrGfxoouefcQAAAE1Bm0JJ4Q8mUwI3//v+jf1EPtOhzZ10VMZoqZ57G/blO7Ho+Yi3NWGkKEaOWnKTMTQBZb/3jumdLhfCXmgoAlPMUT3exw2xGWNiadKznwAAAEZBm2NJ4Q8mUwI3//wbcNERCUal7f5h0QKVTPWNlgbjmkWJ7S2NKgIHwG0iGg0786GE3lcq/pK9y34kGTeDc/yr07QOL094AAAAPEGbhEnhDyZTAjf/++A44Ilb6h2GJcOd7Cm/Xs77EAaBBuyH+UJ495yfcTaGsKqPsfkXwqaBplYMShrGWQAAAFxBm6VJ4Q8mUwI3//vBm7tPFg9mTMd/orjDHC6R/TN14aq1uH6BfRoMmsnLIs+xQljFq9ImogOyRHLL//9dg0AOfyMm3nG6BsZY4fYp4Q3XOVG/nIPZmEHHeOzYPgAAAFRBm8ZJ4Q8mUwI3//wCgK7xHVn2P8TqIU579zJBs783creb+z3a5PkBHTlb/74v4uMpPCgMQyA5H8erUDcdIt1E6LwdC+Yd5xch2XzNIEenrNxxtoAAAABbQZvqSeEPJlMCN//77mtCCdsTY5Z3Geb7Pm/QsC8ZaS5D5HX+6B/kPwTX54IPJEXG0Z05F//tLrak3nsdfXudbo0PXPSotywcbWZJ7pa7XExYa26grBmfGE1IQQAAAENBnghFETyffT3E9CkDsyrDByDhrW4cvbhQ6Som9ctl2PhFMGZxVTqRHWTPWDxP4Sfm3HbhULO9Fy6NYBKZryCailXwAAAAJAGeJ3REf4egPNbtmUKQ0b3pPCD9eBXfRv59JqxVlyBWqnu08QAAACYBnilqRH+BB5WDztJcUhExA+6Zo3n44kMJuHtPfnXh1zDcqNof5QAAAIRBmi1JqEFomUwI3/vhUdicsgtcoZFZMaO9tuG7apZo1pIvnF/gM3FczW9nxAaKJeV2vj19Cv3exUNOijnfRiR9ASN9ltcz+VQ4/0i7jHGYF+KFZP3uGFYBHkT2QRzPu+40UwUjscvCxMulW7ycJ2gzsxyoAhycqC6uUVDHWkGCd55vkPwAAAAeQZ5LRREs34H4Vo/ACpoAhQy5Ym6KOlAwH1FD/kXhAAAANwGebGpEf4exE2/rMWYbvt3Gio//JFTxWQEj0jgsGN8zykk9TVdEaSZbn/AJPKDK7pTunVhX0/AAAAB+QZpxSahBbJlMCN/76z44IUf87NrRe74hd6WSQF0BeASvxdUxL4Rm4GnonDSgvlzazo7BX+ROHA3HCwCRipy2tg6EaaSo8weDziWONE8To2CI9Nf6xMIHsssxak/PU/YdCrDJs2119wdbr4Gb4bSqtsAe2S56lY/+8eHBMDPBAAAAUUGej0UVLJ99Pb91NJVBQqP2t/+IqmE8MrKhL6lXP6xi7XAoI/PH3EdbwEwYHQ13+d+UGm6YHZjB0tU+x5EgcHwOjl3Wv0NdZKEk8tYEJAq9JwAAADoBnq50RH+FEsrDaMSpjV8Ar03X7CI6tPTzHGiBpSWUkMJVOVPbCEUeYgsYcFzDQe+eqnhHQbxBqzvTAAAAVwGesGpEf4exElWAWL5mv1zn78AGfLg+vpjLffFPe0N1Gn3GyNbeSNBNOMwp3LVlKeiRguPo5EckDmXdrHphQLAExUYkI1WZ9328MC5f+Y/xxKuGNSXp9AAAAMxBmrVJqEFsmUwI3/1jexkZYO6f0SqUBJutOylzcq7633ejQM+YcsbI01FdCSdVMPBHTcsAfhFmyKZdJ9U2koiIqT1xC6nMTtNhYSNWKTjEkyshkiwoA2kvaff/xZ7MolmxHJAhqhjc60XuqzboSWvYiUYabtLNBXbUEc94ymgQJJgQehwjl+Rj3g8ihDXo5vs8YeeuDIHNpoqyq5Z0gfHVseBv/VNMJy2baCQS8R4kw7jDJW8iWXhZUAEJXBoXny6onzxbxU7N9r71XS0AAABDQZ7TRRUsn4fNy3GU/pr7hiYIKBgdsTtsrSiKt3zM/HqMjo9nD98MgGaBM9BOcraqeIPJTzn3j7MnOe8aMxA90X+twAAAAEIBnvJ0RH+HsPhkvu1J2NTuf6EJ/2wV7RDK/mYg08fzz6vEXYZTtgVzrhrYCI12GtLgxwIpISuyR2pURhW4rUvLdf0AAAA2AZ70akR/kFs5HqaE1ZwEakhKwmD4BSiyCsJ5I7i1b/JX/yAr2QptkrDRNOf2vgGvhvSIFmCfAAAAgUGa+EmoQWyZTAjf/WJuYK1IHlx90OrSnTNXeuEGfLuro9aOPtuWzfxGlUqLMokdEgNX5hs/sSGuavwQOvVjNR8vXHeA7drDwV4vSf8z11/Kgo3tdVMNLWwgwf/pF0Ahg+eMMrSN4mNib+1ERK71wdTSsh7YgqUhZ1sRulfT/DmgQQAAAE1BnxZFFSzfiNYJXadE4qxy7VpUy9Q41T0IwKvWOe6XdY5vDoVxuSNPv5DidAVc2oMbHfP990ARqTeT6SM/eIoqojUVbdfcQ/gdcgL8QQAAADIBnzdqRH+BB5XFIVzyB9wJv/Lfxs38Dbl+6UFytIOZY0QdJSWK3tXhSOzb4vSc0LZtnQAAAIhBmzxJqEFsmUwI3/121j7DkGICHPtqvASEuNUhdr93NZ3drLM1IHCZVXsfuAYpeOu4rIQCjNeDjn/nmhlPTKoORxu4iqjlk8//dcd5KvIs5DQidbJMVc7qc3oVJpbzwdr8dbcRFN+vvSDcc7Fqa18YeCHCZHZwXF3Cz+Dy0oArWOp8iRf6FgOJAAAAR0GfWkUVLJ99mioJbr7gTovvLkf/l/YtKEL2mu5oyOikfBlzmuZ6xLlP+smIZv0LBUks+sQjAEsE5q8iWpcQhDfktWHs348gAAAATwGfeXREf4pxDGIf1f+twr0oL2RF3LEOzI8akUIhkzGFQYj1qALSEqZaaagajuk+Xra1iK8MNJ545dD2pPxL+B29q8NRoc/1g0bDynEulFgAAABCAZ97akR/hnW4Mf915MgtV8Fyl+jyhJqaffOqPIliFZXeE+XQchMQ88eJgdqLi6y+ytVR4nyOKwv8jWcpSMzkR+NpAAAAZ0GbfkmoQWyZTBRMb/wDwzHA7nHfcLHdlhDBdxj1K2iflRpBQ54lGRaeTModRU5a8siLRIR5KHbbhYv4ItU1x58f1xPglHmsDhwytPRmDnxl0gvYO3B8/ds3Yg+ydO+qQwyKDZQ62qAAAAAhAZ+dakR/gQeVg87QMGmP7dx9N0u/+sitdubSF+CATvuAAAAAf0Gbn0nhClJlMCN/+8dP0UsID4MvVXMepTuAd92+Z0g5Hm6laabBXW3H6+DO83pTGwIUl3MsSoPcrWr246Y/Qww7pWgjMh3A6DDXOd05ptlNb8IL/CBYCg/64pd6ZHCzVpOaP/3b5cwCm0Sn5E7yOZ/JXi1D/JAEeBMIy/0ioB0AAADkQZujSeEOiZTAjf/J/Y+WDSC6t+DsFl+KShs2+BRlWjHwkCbRkfk5F1YAt0JG+eWgtV3xfaTFPYBv+mxEEdHUfsAnOpwhTvDIVMNDKRQeXrvoZCgkP6BOglvBRcnVRm/QCrn8gMOw51JbsPpVerALbca6HHK4IkOueWgXj8/rw/Gq7hL81TYpKHnCgsWhhNql0edQUpWAf3cM29hrgBSKwsy3jQs6DrK17B584Fz6wiEv//JAFYwi3YgyIjyLVn3WboOj/oZw3ulrAwtsHz0duObfFaFbLBt6sED+kDCfEscSq2hhAAAASUGfwUURPJ/Tvx2ZXIhlDv4p8IgQMX5AlE+WsCstVw0AQsep3TWwxt8vGXFQoms3Q7j6KFLJHsNyi/X2PtL8B13mVJbuvM9va1sAAABDAZ/gdER/gHDjMGGbkbroMeA/+R+98sKd9OE/osjRzxqU4fU5myqG9RgJ9RKLHL1WRSCFE5PdzrAJ6RnJAR/133tzwAAAAEUBn+JqRH/XxZ2/arlGJHqHMGb4IS9k1LB2PTOkX1omkDNqjv/Hyd0h1US0g+kY2+XgTJ0qtQYnHLVn+AHiS4iyJNOTtEAAAACwQZvlSahBaJlMFPG/9BI9qA10G55ifcR7emNzO/s5HqLdZT5FQFk+OcoHWoo43SEpEIzC+YKxdwwVhbLtnL6xRqv3B4n9IDF9aNGjQC3lBkrevbFcw6jXvPXdk7lFbzKvLaUj6IFJ440Id8uOH8S2xSDBjkuw7fw7kuVtpJQ71cMUX6WsUmpYvzm182OmweSjE2W6mQahX5TDXHtuYFW+s7wqQG7mczW59/E+jHqj+GEAAAA4AZ4EakR/mJmxUG34cEMCTD9KfzSlBy5J63PV9xDXOTeu/CSqfUIWFTARN80zbyy1JouE1/9W5KAAAABpQZoHSeEKUmUwUsb/9Cwdx7ECFH/O6zLxukzGly8jJexJ12NnWZVyEH6v2GpBHI24QsdFluV6MU99hhAQ/ORGnQtP16u598A33ALGAhCiT6Fkia5Y+ko5b6BByfi+fm1pz4EM+HBx3ixAAAAAMQGeJmpEf4exElWANW9/woTjJTRwrYLgW5TCxhNGjZKsQXGThPJJ3FtL1nkHDBC/zEkAAABbQZopSeEOiZTBRMb/++FR2EKSC0QlA46ppypN08vPo2SkMTwZmYivCQot+O022lnoqxQWOoCxOpu+brG/vq1OiqImEEUqH1aSs5FB5wZJH0p8pceogOzxqbSFCAAAACEBnkhqRH+E1CtCa7nT6sPwHk1BQYUkaWxgrGWqbiIHMPMAAACMQZpNSeEPJlMCN//74Djgd5CY+5YA50g7ZWcl3e6HlM7+829ssvr2kKYX39Qy0iHGV9xtA5s2X+EIG8ksj90L//eHWxQYDvlTbZLHSXrXkPuNSEsLXk5NQmC2u4ggLAcIDZeLlKdlT77ixCPNy4BIJtROhSq40uHLaCpv/QYI7EgtcO9f8Zyycfr2+0EAAAA3QZ5rRRE8n360LEl02ntKg9DJ1qf0Sq2mRcCNvdMAMT2Deu+GU2vzNalSElMySEM6p0O8C0Zo8AAAADkBnop0RH+AbCj0Y6pj2ya9J7lxHir9K6838HLwTqvrUGcUcg/6mmd9tmwJ2olawiUUSVMN0KYF8/cAAAAhAZ6MakR/hNQrVcoxKlr4Jc/+7NWRzf8YKIY6R80MIEDgAAAAkEGaj0moQWiZTBTxv8fOfWo2KDtv87okIotgJjCiYB6nOqRdka6T2FSgrP2m37XMvHZdI8A21zD1pYWm/KGf+qjCETUhQS1RRxjZtApAg2eNLrsSASJT/M/z6WcXboIDcyflWuaJGh/CcbDrACgNgPWBxQ9Cy4GZ1z2sg/n7DG834nDz1pRrY0cd7DCoe+kfoQAAADMBnq5qRH/N1V08EzDwtoIri/39PtNtN5FyTk7/bpdx9BTbj9S5CIuX3khlff7kMGMVX4EAAACFQZqzSeEKUmUwI3/9Y3FKc0qbE9JKkUV9XeNE8I6QHJCzI1xykL6LJyJn8sMTobA325dt1+ho1fe5CgSyqyKmWLEtqPoxN/8FgOH7HVcYanqHPuM2LmQI6mevX4hqU31p6zFddFY3jjC9SJd2W8DUlE/UMqWtTbO/zdKi6N77GCUpGa8akAAAACFBntFFNEyfex1FrjBvpN37qK/wIQZCAFBrCFmyElK/aOAAAAAvAZ7wdER/h7D4ZJjIOAubxKBfeUw6T8sW9NMSWDaQ0+veb/Jd0BOAnj3BtP+dOcEAAAAaAZ7yakR/iOA7CrB7i/Mz4jo97V0uGcnVWvAAAABdQZr0SahBaJlMCN/7/o39RD7fvlTWJwGXUUqvm05+1QVmloNO9nFYywACdKqyxRSD8AZPjb9mbnY8gjIvuePK0VaUznabobdf+xj0P3e7kDpR3NBMu5vVSLNUN/ElAAAAZ0GbFUnhClJlMCN//APDMcDucj7DVPmcuWIK5vVjxlhCepo4trDOkrTmEGu/ZySN/QsIlFHz/yEpWOSVZ8qYXpP5zoMHCuoHGjtJgUJ9YY/EQHTth3rvh5r0lUIjrN6aJq1qPEo0fVYAAABHQZs2SeEOiZTAjf/74ByCXW+nW+BxzvYU369nej5XDAX92W9IvylVTvLOZPNdVa3B4/9OBz/yENcbuBsBlmdr/FOsXai19kEAAABQQZtaSeEPJlMCN//7wcOsd4IOcpHAOYB35obyD9fNfDE0mDNzUbz/oQknAvvZ4dBtywxoOtUrwhbgoCnEAdmdPKWcS/ilLOLAuTdnno4dinEAAABFQZ94RRE8n360LElvuh+CIdT9OlDp/0q7PJ+MJgJbrVboZA4kHTJ3ToffbZ9DMO5KtxdcbtUFSeygBOcNgX5zGmhe9rpJAAAANgGfl3REf4egL3/U68Y/3GAEKHH5y9IApJhkFHkWHqclHfOYIjpgTUXMpdkmgPno80wnwLqyZwAAACoBn5lqRH+Gdbgx/3baE/Aqppa+0jtFnUyKiy+Zwy1ZHyz9fyMNwfsj+4AAAAA2QZubSahBaJlMCN/9dtmbcgHeY/EYukQ7aaclKiAhW7LP0wcwt78GzVH9jEX+ul4MCqF/2B0EAAAAQ0GbvEnhClJlMCN/++A44EO2R9S/OhWXCYYnfMvGSfmF3VhHoRAookQ/B0zybEDHraNElAkFmtkn7maAS3rK5ywR/4EAAABvQZvdSeEOiZTAjf/7x0/RSwgPcTDLgPuFlZF5qsphrzQTuuhcRXT5hraAgVTlKSofuENEmB9xuPyUnh4xra8xjPnk0cC9TqCjw8oFY3hxp8p6Yv1rzk7Q/KA8AKgsEUTklNSbcJitU+GRzJIOetTuAAAAlEGb4UnhDyZTAjf/yghA8yNyXL4I6H5lYLZt7x9sYonEZViic8JINT2N7pAvBTCcrMb9wQzCFD0BtXGxWf5Rsx/8l+VOmJ7mN3Vdkcll0IUwspcRV/reK8gEFqp572XghaC8+bGg1kDUICNBmBkZxBrR0PZDAzSKR/HGRfU3wFVnoidVdywk6QM8o3fSqG0jFI6YMWAAAABdQZ4fRRE8n309hYNYZQ8K1MQRY/3//frBdkHLgzarcM4u6jySXsngsvUyy6EI7lO5ODiPJckJ9lqLmYq9/4Aow4+wuMKOMB1SJrQK2t0CbZ91mEJWOzNjbZQ4RrPhAAAALgGePnREf4ew+VpnMlKRTlxMhp2z/jXp9OUQEWoWTWTddscAANjrykkhGK8n8fEAAAAuAZ4gakR/hNQrVcoxKmNXwBAAZe2CmhnDxHLa4k62bX3bzPVMZ+4eeevEvCRnSQAAAFxBmiNJqEFomUwU8b/8A8MxwISvXq6TvqNm1NMGbNBYU3BnWrbyWWdLxOVYy1WXkGewPWIzm5q7yUvGhk/YnyoDdKJ8CrXR3ARnW/KIQtqUzAbH4jz9hM3MiI/EigAAADYBnkJqRH+HsOQxTh18+g5ZBI+1d8/B/i7BlQB+WddgPTvIA214hf/1CtuC75Mbg76dyaxml8AAAACrQZpHSeEKUmUwI3/7wRk+WEBIIWYIPFZOw56t33hsE66tFa3ZLekWImE4LvDK36TKHAgCfZYuTC9EbJYEL7N1Y6TRBVtzM2ya6R2+ErPXB2jjcjTa8RL4ox+ohOuPW7UZIiqbV/BEGgQbrsylo7OhkBXa3M4cx81PtPa5oTEEJZAzD14sh2F8HgCI7sfH9KNgsLmzFk/eBzzI7emYM5SDkSldP7xc0a0irmWBAAAASUGeZUU0TJ99EZTLsyDpEIRrKJpnoG9EXvvB3xvGEbs61IZFsRW5tRZxKSoUAv2FlyzhGJgURfX95iQnOI5xMTYX4Tm9fdpIYHAAAABCAZ6EdER/h7D4ZJjIOE7mYezedRutYU/Lr+V1IdLd155NZOE9hpCyg0nHA4eDwBg282Ojm5gxupgyr7ONyBfl+pZAAAAAKQGehmpEf4jgOwqq8yvLtAmySKIuJcaTjoN//CBYJDzeFU/Xqm2khRqBAAAAYUGaiEmoQWiZTAjf/XbliF2xAJMQm7WeNk8OMW3rdjdKBNxob1e/+NepMrYVL0nMuodAkrIMadu5mJBy5uvlxvNQp7VxX7xQevbsT781icpake6htsC9t2XleFcTwqueItgAAABQQZqpSeEKUmUwI3/74D1EHSkIO4Z9Wvyhvb+EqEIAwhAaGZld9PgYxqFm6CIw70d3crTbR84Rl1QSjXA0mnDPfQOFBboYhbHq4C0vpJXMNo0AAABQQZrKSeEOiZTAjf/74ByBB09qkSKnA0a7ky+N/3pDcHVcDQBwXvawTNvvn2svBAJZAi7D6VbLM9Eih/kRceBxSGO9APDLO5zFex2ykGUOA4EAAAByQZrrSeEPJlMCN//7wZu7TxYHfvdr3YChXTkbqSdy1lvNC4XiN9KR1PqesySw+ABEfzbkpLtzgLfVv+d1ee4PGmvseKH5sGEQg1/DG7yZPZbZVEiMJvrtqcWWCmcntEA1XNbsrhVN9bi050eM8WOF6MnOAAAAYUGbDEnhDyZTAjf/+8OhtVTlCNvyUsskUPEVWUnL9c90uNQCPoB4rg9+sYpzZa6v16m/de0qMOA/BNngL8GoEKLLIlgO5+X/dNSJs7hCfm/5KPsbwl8m3mNvXuL195ed9WEAAABaQZstSeEPJlMCN//7wc+0/eDsJNVZQtMJRUijWyOQP+twyrFR67F8Pj5Yf/Gi52beGXbajrcE2ZHcat654rgyJ1pKhO5Uy/gC8B+G9/OVL6pzknCH54OLfIXDAAAAUkGbT0nhDyZTBRE8b/vgOOBCV69dVilfxFNE78xMP/CEHWNxHiUnX6DyHq+qeebHfMqief3NvEbCPaYjoYHVRYSzyLFgSrimfr+f4tOe9AjjPHEAAAA4AZ9uakR/h7ESVYA1XP+1Kfs2fuAWfT1RiJtfaEy2T44UcRtUezFEZUPv4YZsGo4J8wK41PBEKIEAAAA0QZtwSeEPJlMCN//74DjgQ7aLyl6IAVFCZ93OMgKO7RPkPiBlzxXr1IPdsldBX9MhV6aIfAAAAFRBm5FJ4Q8mUwI3//vB61Ydg2jrrF0XGCsxft8+s5TpCXAiRHLtzCBWGDzOIi2fFiTNEfqmbdWsGe9ejcTntioJKJzrv8cb/s/5QolStYE3YQp5+/8AAAB+QZuySeEPJlMCN//7/o344HrOM5ff/3Mc7c1rlUvDf+1BYAPMg/BpElFxK2uuwTvrLu3L9d3CeuKWIiq2NiwAKzpkZXVwYoWgYQjBJkeY/r7x0HhuZX04/RYn2nHmTL/bcfYO6cjHT3FMoIt8A/Veo+J36ssJSsTbgb64suXdAAAATkGb00nhDyZTAjf/++A9RDRLfNfS154ZVlkmiUy6/g+2oGU7FxgxqSr5jI8p2G+RtB4SgTfOQhgR/tAQN0eXRO5enuY2bDLAZlbho5L7gAAAADdBm/RJ4Q8mUwI3//vgOOCJXWMSryw4h9E3CB6A0Ke4+kyCv0XYLF03+9oB9Vjy6J3L026VqcJXAAAAU0GaGEnhDyZTAjf/+/2woADtubdqT10cQKSKXeNNuaLX2/j7JafRM2QBPoIzUNui7IROPY6gQ2i2OrH/7RUn94RFukRHfq8+EKh3HJwvniK149JCAAAAR0GeNkURPJ99XohLbldl8+WVFj+Uva/8XeQ7k+3TjC2+NP1oPDpKnfim/+OozWJAwVNoYbSwY768ReLtDwbAMvTprf9JfzUVAAAALwGeVXREf4egL3/U68Y8CYofdFnLKmENVl4qCMGjiGYU3fOYIjY08srI70+2E3XBAAAAMAGeV2pEf4Z1uDH/duSE5sdhg1tyJKvPRV0pqHivmE9r1icmTrC7mlJ7RSXv2UrMMQAAAEdBmlpJqEFomUwU8b/77lcOB3OPxERFqzYjd6PE6HsucwQ5Vqvwef/OwSmT2LMOe8YS/MYKVzL0ysCkf9ELyWxaN9aVNLSQfwAAACMBnnlqRH+BB5WDztB5pfz3BcMFhRhOn/CrTtDVV5wj6Le4OwAAAJNBmntJ4QpSZTAjf/vBPT5YQOjX9DqM5HaGpFMwQ3svWbjKwnG+fOlGDVqt1ndhhMJgv4GvGjt71cvzG5oQlpabB9hxzOPik6dNzXgscBR4+cbPmz/fP5xfYfX1NBE6N+5E94WuKtbCmyYygQguamHJ0rcWbq8PZ31zvB4cN2ZTaqhzCYM9gDyDIxGEYnFFk8tLu78AAACPQZqfSeEOiZTAjf/7wZu7TxYPZR9WnPBp8pmX4bfrqeMVK4ZK+wWWuCUmWaP3lATqHHJZqYOGZWX5B1ExYwR/eLb7x7S8jkASRZ4JYh0GwZIP9CaseXY+FdZZiTo1Af9+ZTj3dfC0dKh212pB65IP6OYNrtjjWbaU5N4lw/dIVo6wfWQ1Ox2R/Ah/76cmOMEAAABSQZ69RRE8n3sgptQ/K338OzsnCnldYLd7RX/bbPagtm9AT2uZlYrwak/pO2OvTeE+MLYy+EAUJYdRjCTJ5ws02INMTcZhmyBQbhyMvzM27KZkFQAAAEkBntx0RH+AbCjmbUS83YiObdjoslRvNC+1RH/LHg47MFmS5IoLuTYHGm+zPkdYAIx1ysH28isEgGxKBbm/3QA1sfXNc14DBq7XAAAANgGe3mpEf4TUK0MsYlS18DEAHiL+PcNPmI7mQ70eRh7AhGHFQq8+OOpKoNFJ1E0xbu/KTN8mAwAAAIRBmsFJqEFomUwU8b/77lcOCFH/Oza0YIke18+jZcsifjoIbn8hLB2G1sybcgi8yUenoEyRzOvlQ9G6LQ3LjzLO/wVkdTKQGKdViQNnNx/ukcfeqyW26GuP1irc+a/5I0StQFo5b3fttkhWZLCcmvXpDHKib2Z3J8v4Mxl+oH8tYhQMmnUAAABHAZ7gakR/gQeKxIei4mGll/v6lUvt6J01kEEpB2nbmZodQK3mc4O0v1zs/0nPzjxxUZ3GOYRGobV/cAuzXdUnGujW6ZuRq5cAAACZQZrlSeEKUmUwI3/74VHYcs5s9CQn7HDcl8JfyAYtwAQoJsnV5HQm+wXQ23ptJOsGIH0+K+T4tJmGWxKq2QkvbVGAt0rBbo9MwfKNZ4T+zu+cRWcAzne7vkeYo34ehQqXPe0QPB8YoxnrjvXyG9IVsIPvTqRMCT5da0BnF7ybZ+dwHNKFrocGgyWwGl61sMCNyw7UdJNZmxywAAAAOUGfA0U0TJ97IKea4wcfTVP+rq1dhzTiOczB8P98SD4FQyfV4nTwQWnbpuSHsjXkGDVY+VFqlbqpfAAAAEABnyJ0RH+HsPhkmM1qy8bxI+teINNs77QclW9Syn7VgOysKvEXYZh1N9bkMSWtoQGKeb4wR7zVJmCMskcNH/eBAAAAMQGfJGpEf4eNe2TaTWCeCgaUGvNb9YsQ6osP4XH63aOXxG1jZCXFAnW0W+T3+SXJsEwAAABgQZsoSahBaJlMCN/74Djgeww9t1S4BqrecgJLTxVeOuPfGThsraJOF2+pS0DHf1L7/3y7WCHgywvW5UfaSeI840uFRKAdPikqUbRyEQbimIxPxFcVU7IF9oKgqilKzVfQAAAAS0GfRkURLN99UveC5swUPUef8Rxxd/DnGwoGPLEv0jI5bF/aVNigJFob+cQVyS7+9J0DBc/WkPNi4m+wyxdk5RauSIYK+bf04StobQAAACgBn2dqRH+BB5WG0Yk9tZXU7nhXQUxniPOff2if/2gJwmvSs+oXpZoQAAAAgkGbbEmoQWyZTAjf+8SBZ14QI606liNL7F0TnvmzD9NgO6qUoiqkr2dMXGBxVHEYasaZ1lgL1s3uJPXelryzxkrFiQQAJkYTJ6/ddfxquNOViqUsaM7edj3CU24qxjGC94+WxcABRaj4MjKQx/0LCKQROQaDtp8J+TfTj/pDj9bBDoEAAABiQZ+KRRUsn32aKglZlvQTovHIeNH7gDIJfCEh/db+sTywUiMMdw5UkCuutzw8i0Uw3jTVZuGhnn/DdTHavTyUTz90N47s76Kr+H9cZcMdUHeBR1waOeg55kL0Lnz5MaAKqHUAAABHAZ+pdER/h6Avf/WMAaj+02h/nyrbUOfaALfnmG2jGc9vXFwOALhj3/sE457lRtcFy2I4QbRoLfMtRwbncoZiPe5ckLTaUSAAAAAoAZ+rakR/h+Qj8Y/7qZFjwDICR967ohOCf916Jz9qNGDcGK/ZVQo0JwAAAGpBm65JqEFsmUwUTG/8CBPDIEGcd9xK7sr7PeVRUblA+A4wrZEIYDPPG4pGP8GneuG/+WblEBE/YISVAlzm8YRF1iarh9/0YM16M73ZZ8APAFET4NhFgwWEytfdD1qhPlSrycwSgeZBOZfgAAAAMAGfzWpEf4EHlYPO0DUfukzjrb6sMwf9TWr92SL5dYTUf9xex5gvecavMELDpnbvgQAAAFlBm89J4QpSZTAjf/vhUdhErrGuSF/CSLUgNGgjSozDQkuWrpzMDWZP18VLPXElhj565QB5ATyxhTtCBjXMORwH+flU+06xO0G6o+HBCuMLMHZQWc86Sgn8YQAAAHNBm/NJ4Q6JlMCN//vBm7tPFg9k/OLpHQkps0Dx8gSIgZVvNmWxh2JbcEgdj1ANWONNo0C46lGcg6+GMwWquPjmXeWwPjjFXLS5WiDMCJHZI1S8p+dhwYEHn2JemwPkCTGJT1YfRwqRnd4zFrrK7w/IrZzwAAAAUkGeEUURPJ97IKW2n8rgeWx6x/uvwym+iKn21XOuUJQvFl7Jvz/4BnZZvs/HplT9c0k3KbuOgl09GJ9Lr+6rhKi+Rn9GiuEu1IK84X22BQAPloAAAABMAZ4wdER/h7D5WmUT0hGm55/V9t/EfeHSNG4H3z8SITMcvPv9nR/LKHSw6207w5vl5dvHvbN/P8hIiRRBAbI6dHaJoqpBvvYa6Mw0GQAAADgBnjJqRH+E1CtQOK583hyuzCeT2ES3e7EAJSPa2lf4ndS5OFQQrgTaHVvvcR/0mjt9zA+CF6l0IAAAAFZBmjVJqEFomUwU8b/77lcOCFH/QjF2o2bU0D1gZNTTA5MxeI3pXbKFI+z+I3Ehsro47fLOUN7RwIz1I1j5IIQ79ZTKZXOQw7psvfz8GmJZtuqJ31wKhQAAAEYBnlRqRH+HsOQxTh16NjvxJP5pSGoyZ2MG9YV7VnkhZNujwds1WAb+DnvAXVsQBWexlXy/XhL2mMKRVlpuMTQ2rNsOIosHAAAAp0GaWUnhClJlMCN//WN7G1kj63WwyR24JHnTxkh+gYl+KzInu1p18pjfQ9WbiBY0GWq7ysDGaRRp+JW+XOLVoUSeFufVt5tWKNC3usOBK2/C6PhT+8zUu5EZNdqT60Qq4VgNXwzslbfFjnJRBDpdAh5wLLWFV14beae+G0fFOrSHDkUNpkd4sy70UWw9D1jdUuQ5rhLdlkIPJ0x98OQL5TO5gSkKid+BAAAAVkGed0U0TJ+HzcuFAuIfw9nuIuW5/0fmGPCk2buV6y+c17CysRtliEWgeuDAVuS2NEOS+YOmLSnbstF7rBcIT0xiTFqSL00ySVVFsbWwkDRuJB35PF4ZAAAARQGelnREf4pxDGQCv7OsLfYo7T8YkPot3UwsfuaooJN278sKvEbac7AqI6vugAMF0dBPcCL6ff5MKRFnV3TxohZvBW8Q2QAAACkBnphqRH+QTAw7Sa7nUoHLO0XWnIk0uc9V1EtugHZKDYDxjTdjxpe/wQAAAK1Bmp1JqEFomUwI3/vBm7tPFg7wfAZdXzxs0p4sQVbaLB0awSzslu6gp+TNyByCumJLwW5CyVRGbnfEuH5vtaS4wvsym7ljMIKBynK3Kg1MUwKdsVhXZ1B9ez8GwiegBu3bUflTwt/cPd13A1iwLV6wHvPtPLPVBufoRxLdi7fXErtqXZFkhyUPiuvfApRuTZLb6JlWmyXzVCqHWPaoL5G0ARafVrc4NRAD45EOQAAAAFFBnrtFESyffW4h7bT+VmzF2BY/3/IEonznzSaM52EK3yed02INvn4T8C5fRiiAYPyU20UWm0Zioarkwh5eCaLWUgTp5urYldHQq3kp7USwoZAAAAA/AZ7adER/h7D5Wmcyc1Gak7aI9t/EfeHMOpqmvMV7oyvL4VZVPyQ7xKvwftFj0zEdxLJQd6obfw6JTO/DrvdBAAAANgGe3GpEf4TUK1A4rnqdTB7ynQtEgcvcbzbIwk1IFg5AL8SoEuxeGvos/djDz9/TNXsT339wiAAAAGJBmt9JqEFsmUwUTG/8A8MxwQo/6EY7tRuuvUQ+IkRW5CuuQ16cI6K9UjLCHt/DFRPI0/cesLTyTUMPgo74dmiSCFWGW0NNaVwImkdPP8OS321tUl0fM2FRF/ydM1FfUQ53LAAAADsBnv5qRH+HsOQxRCCMCTKsMv80pUcMEXeFoGccbVvJvQNyqsPl3dGC9TDj68G4sLpt/4joEDUDO0HQPQAAAKZBmuNJ4QpSZTAjf/1iba5425bwbsmO+LqBQWsUydEd6pAOGPrrpQZuwRoMIuUNmIqaa/ZfsYj18F3+2O2VzASzLDk7M5bTVsxDhHKKWgIuIMjnu3J7use13jH+mTne9OXAolpGEM+YOkgGn0roRh08mvTLZT60kAdnRLZNC1MbYLQv/UtAcmYn97x3XP7Gik2RIKYbYXaybwMSmgqFRWYYz4aJwVLLAAAAPEGfAUU0TJ97IKea4yheePM3UV4BYP87QSjQ68uFMzfjZmCa2q2jTWdubpQmvIRE5Szocob6t0FwExDfgQAAAD8BnyB0RH+KcR92pSK4d0iK5pn3//DwN9n1Sf1/FycV11cqIRtpQSj1cG8kqGHiP/LJ2Q3+b3t9voEB0KrTJ0AAAAAuAZ8iakR/iOA7CrYE++z83jlf99FGL+nfLDBphBaI+WarzWjX/yqdZNNuIw/QVQAAAHhBmydJqEFomUwI3/1jd3Ei9XsYRCa6uUFLqSIqHEXps61oaZW//CrIivsMhGGZt/+wTW5oOHXAD9VovymjVgcdopGLcyH3+31aTy3MJifAOoc8WLgubo4ISngsCMnt7ldH1X98l1NIvHlX3mcEIj84dFxX5oK66OkAAABHQZ9FRREsn4SFSqCugoeFae6NWP9//36wXZWgjOWqAHC9KPG2ihYHENvn+jgjcT/OLdhLP58EwMXKROhlixI0N1Rfgfyt4FoAAAA6AZ9kdER/gGwmmmOq/Wm/D3dnucn5igsLTeqyV3HaZ/bb0TUerDFX4r8MRrARcezd2SF5uBaMc2S+4AAAACsBn2ZqRH+E1CtDLGJUtfBL4cPgesw0gfX/r7jzYBrCPcjlvYuTujdx5N+RAAABSUGbaUmoQWyZTBRMb8EdgCeRWNj1vNMstxJbRx7SX4GlzlNFlrR55BZRfwuUxlJowdEahz2muYI8I5vx8hRYpCbfVNcn/buf7/gZ9+q/Xex6CfNdIK5dm6hYw4V04WRD2xx5XfgFgFusJZ2Lu3x5KCiAJR/MKrRmuD0LpBj2HTd8vSTNdD0mqExm1jVxz11VS0JkCz1kf3+GAZ8aWDLfmsfKwTV6EgQsGOAvWjEoPcSCVw4czv5SM2yxrz2q4RWCXKP2cQ32KjlN9WUwxh1yU9H8IlwQpYkE9YLpUDTnf5VkAgRqHHtbB1bKzLqdQEBRnPDUxU3lCJ8z1ffgnZnrwQ/1jrkIVCHtxSHz+jlGJYYDk6e7WXKb6WyPh0gIrK8i81b+HquyEnpTv7rHDiPZTeoAPC5ivY3TNLqzME6cL8+xoSGHv3D8YziNAAAAOwGfiGpEf9e9wzWJD0aXSLU4v4AFyP5OmsbBEcN/WP0ljTwYR9rY1eB5C5dJcii79292VcOw01nChUeBAAAAwkGbi0nhClJlMFLG//QSJIKCLA21sITm22CQNIzL75x/coPBB2IH+CkNL45/JcPWK61azax0RKRKTm5d+BssnzMtwKy+njVCxG7G4L/0fmAAHFHQoH93jRGl/HjJy46vx0G6NyLZe8h2jDkca48Z8CHa2hCT6RIrRCpQcVFTXZInLx/gmR/ABv72Y6Rg934/cwAwDTTvn/FxI59I7ms+9ozaPN3OuXj/ZPs9jO8lHXscoqhncvzhU0B/ZiiYXSNR5B1ZAAAAQAGfqmpEf5icgpJHKscCg4dcTFi4T//E/XAXJpG32nkk08fkKUrnPGoSQKoKEiE2wpfQy29GE8eVH7/yrD8SLOoAAABWQZutSeEOiZTBRMb/++FR2EKSC4tylzphI7uNxilsvtRiMSiFLXZzkidfIAjCjC8+Yq1U+7FQos18k2wGsrHuaY/yfMMHdO+Ytp9nyOtNwN8ZSWlSsgcAAAAZAZ/MakR/iOA7CrDroBeWntKNDlGBefdNgAAAALZBm9BJ4Q8mUwI3/8b8mat5xbgYxKkPL6JuE0JrSSkGLS34qllRYb//n6qI6h+DVp72C2Bni+fBUJkeQ0aRWTwoPYVKKNVd/i/tVkLXXk1eYhf5mDfymriDx+s9bFWvnxtSq7+R30/mIRxeiq80eB3FjjMV0v9A9Ip/t0I8jdkPBxJNLY6dHX57kOj8uBOxzCEcewQyFjVwZ4U9eeiGDNW5cmamVQO+G1zTO18X6eiBVrQ1sAaKeQAAADJBn+5FETzf0FK+iADehK1fJc1pMmFCTGtLvZ658Bpk3MRU0V0znp7jGG86RAJ+kBsHtwAAACMBng9qRH+BB5WG0Yk9tcQHS94V2gbW9d0P7JZCiBx0bHdSHAAAAHRBmhRJqEFomUwI3/vEgWa8cBuzHU5S0ix1bYRHinDETekQfB+cgzhJLg7T1kDPCseQtg55Z8jXr7sf2ocT2sSw5hDEdVg0IPoU/FGMpuD08g8QhOFlplsTFkZYkjR7winNg/k/jN/HF5GJogmrWYBvwozzlAAAAFpBnjJFESyffZoqCW9/YL35uOpRg/l/iy41hTe/tE0ssOv4hnLHsHPfmKcSLTegTjkNLZVWXjO2NeDSEq6BzQ01VDjT7II0nXMOQjdC1CNOmzNCCPqLBayW02EAAAA2AZ5RdER/h6Avf9Trxjzf30y6kAOs6WLBLxUEPR8k/75zI+NLVHsG3auYqXMQfTCCvtQyg3VgAAAAKwGeU2pEf4etKGWAMuE3sYEV/anP9U6xl/ko75T+a5ob50RK54RK1XOQ4sEAAABcQZpVSahBbJlMCN/9dtmbcgHeY77hZBIuaacrM+PCy3y2hv4fgKY3CAJLaPuYPjQ9hPdbPybTfLA9ZhcESX6IM5yNtlJbVY47bfUjagQR0EA5I+ypI3tSGPRusuAAAABKQZp3SeEKUmUwUVLG//vhUdgh9eiEMPfVB8tcq9JDka35FeHzBuuUIwuu3ZTYJerdc8OfLMQc0Kb0i+lVDCTmweLq3VzwQLqGbWkAAAAeAZ6WakR/iOA7CrDrXqtfOU0OnEFoMihrTVlbhcOxAAAAckGamUnhDomUwUTG//vgOOA7zH23LaoiapuTmxATD2bJaDhWz+mrw1lKFIcbAS6u3oi38M5gOc+H7/dMAlr5FiwNQQ+MnP95HqTXfhGOyOvhpUnJ8xbBRClFaRLviLbneAYkMEKA0mVg3tgaK2CBRyPigwAAADABnrhqRH+HsRNv6zFmG77dxoqP/yRU8VkBI9JIa9XsBjVtXD6MA4j5uBIgllV5DYEAAABnQZq6SeEPJlMCN//74VHYIicdzoZjPCSvgahArEevFvdThrsLLMVaJDVNLqpZTPWNC9vVX/wO1tcwEn76IJESXUEvUSeHuARe65ZpLp38T60KexwII0tfQkOqrPOP2CqJfFhslG4y3wAAAGJBmtxJ4Q8mUwURPG/7w+3BOtghVFB1Iv2nqcWzEymAHxzkidM6m2fE6ogCs/6Xy6jtl/XzhhsHnCSDlAQ38GrqPDxXxyRtvk4SCfcKOBRFyRAOB8n6sM4vESFElpVQD/vjoAAAADABnvtqRH+HoC+APa8YzC92E2GO/nuYBco//L1SmEwKP9mqIR44ppDFXNReibXTYz0AAABWQZrgSeEPJlMCN//8AtC7wJQ477MPq5rGdC7Dcs2GbwHJuPCJyM5OIOtlfIO3cE34Ssa8cZkPDWcsRjRYjceQX64sn64LgFO/LLPWgI/X+x8DlTPlkHgAAABAQZ8eRRE8n30rer4SgxWSdC/KZ7+otOEzo+UFD9ZmniYypUuFzJJ18Je3Kzvv6f8AHcMmi6YRwakaHS3jzaoZQAAAACsBnz10RH+DDmAYlt28A+7MCxlUQzSjDlyF0aYn6fKlx9oD55kOvT73INX1AAAAIAGfP2pEf4EHlYPO0Dlcu33mg7UP+2G/OS+vSWgfJ6uzAAAAdUGbIUmoQWiZTAjf+8FGR0k7B7JF8YA+wQjft+K1l7Pl2lFva4U4iloi7jK05WAXlUvU0NMbneg0W8bipRMj+e5PR6G298PZD8Lx4/pmueJhR5QgmVsDRHocNbVfEJGisYpDRTRCNELAzBSPwHo5TkjGz1jvpQAAAFVBm0JJ4QpSZTAjf/vEhSlnIiISPCRtw/Fa1SUZHmNR9eS02TNX9cdJHHP4scrbtJbcvjnamrR2131icLwbB8mAm/JS1f8/MJ+s1WuEEA/spOfP7NBAAAAAXUGbZknhDomUwI3//AKArvBG5U5RSifFlD/Ne8Dro+FmjNGrFQ79lqWS95HaV+HspU897WfSmm3q8POHAYXGxvGACHm9MrOpAgNLObTHeON5xAh/8S0wNVUbBIo4TAAAADFBn4RFETyffT2ygu6Evr15zpphVL9ezEh177sG3lNh72dSy6+LMKsjFpOfLT5s3125AAAALAGfo3REf4BsKuKQrn4+ytpBYNz429+gOd9CbMcAZjWPrdtHWJw6lVT6DYuAAAAASwGfpWpEf4exHtbuDFFtO9aECqMpTVAgxUDNC6R6Ov/Bv45Ck/DDI++GHNS//D9vpyc2nRYpkvKI+H2S9SCYvOaW3tGooBG5951jQAAAAGBBm6lJqEFomUwI3/vrPjgITCZv0Cm/UN86vq9j2om8TSXh+Amp9UymJ2Lmx+GhPvnAhOiY8nnAk80c8/zR4tT3chvitOsi24OHlSIEosC0pK5DFjcJ9H6E07p5eo+PK4EAAABAQZ/HRREs331No5t4RtqmNt6u0B+Dhno+LpcPQ4byTbJJSghklqX8XC5+HsKV6xKkIFe7OUTxgQ0olu+pKqtaRQAAADoBn+hqRH+HsRJVgDTILm9SUFNJlrvLvY/SNT3pkk83hKOIuwzsAgzN2cA2ruLxHhJGC0nlMqRs5YvLAAAAPkGb6kmoQWyZTAjfyghC9HoFZjDr4Z3T9vVMSrvyBTETIkGEDt4rHdxFlDop46jR72snDXzSvTxci6fpWQvZAAAAbkGaC0nhClJlMCN/+8EZPlhAJgR+ivH2tTAf7wMtf7Q0lmwSp6DMeubPofusLc1pk88yiSuDMj9NPivga9JoF3g+xq0Q86goTEG6ZSyNzWHsnx25FbgqP+krw4d7Ztwy2JXZPgQfwac+bc+pWEYQAAAAa0GaLEnhDomUwI3/+/6N/UQQ5pDzZ2qOGQygvLtIfyWYFUfIE/BzO/DrGw19HGK3hwfXoRT80JapmkcAzbWSDpBJ+ucKJJwYRfQ0v3dYjGPpzjoBe3cpWbig6cBTj8apIEOyji3iKmvu0ocjAAAAaUGaTUnhDyZTAjf/yghC9YOgdCGGxMBNvyCtjziLKdC2QZ7FDESe5J9Q9AtgUrneVdY5XUnITqU7yU6iBGOU1CpzwDon0qnZ58WoQb5Yntsn3pLkDSXmlHf4GUB+Hia+8e9WqG1nEx0tyAAAADRBmm5J4Q8mUwI3//vgHICs0H6uMZza6s7teVDeMd6CYdhr9xUUK38FJYrIQPqnYUcIjdPXAAAAYkGaj0nhDyZTAjf/+8Gbu08WBCblR+AghU5I4PjruVw1nmnblE52norPEgNszfsQwVFghKI1Diu0eqa7hbfw1+Ep5QkXdJelsyIpnCozqoRLLfdal5qiCN/T5eg+PO3oDHLfAAAAaEGasEnhDyZTAjf/+8Ohm1TlBNW/tdPMTK/c3ybOs/lY+8ZeligkNk05fNM9wL12M3EuLDYgcCCXSoBpriTN/A1g6V/gLzeZ6XLH44A5B/SCX5SLpQw9q99kGKI5Kg0SBuVmiSFKDmwPAAAAXEGa0UnhDyZTAjf/++5XDgIS808950w5VkT/t3hHWXzd4j+0CV+IqCE8SuscXSYBiVDelxjqtSC3ZHtkbYDFeKciUbuZ1jlpEM/uS2KCxQtBuTVImzCSPE8oOq3AAAAAW0Ga8knhDyZTAjf/+8SBlbxwEJo+j5XM44k4C8vgBUE958Col3eoL9NbbmKsxW8m7j4xRgVojr6fB+NvuumtaBaIlIqjCRjWB+8H8xBIXkvlaBSjiwCBA42PXb8AAABJQZsTSeEPJlMCN//764R2B5oaIM9nfTCNty23j45IR1Uylna2WhBb8CUTivs1xgc8YLMXXuABpWaMZx/XUnylyMe7+O3whopv/AAAADhBmzRJ4Q8mUwI3//vgOOAkGyPnFZDh08wNcfuRYVMfPh/g57ian/0yuRdww/1n7ykQqNJMvWQ3gQAAAE1Bm1VJ4Q8mUwI3//vB61YdgbRGgzS0lo6m3KNl1T26A52Lu3lQ0zu0gITxVmY3a52vCyX+xlNajR4CQjogG549vg9JY3n9NViV9WrQgAAAAFVBm3ZJ4Q8mUwI3//v2UuUJDTVPO7uHUL9d32bbRDn5GpAyvyun5eOVCscEeVYfMq+tDarXNr6Iw0bo8A2x7veHCIxWFjrr1e/Z5ydf+rNut9UgsY3hAAAAVkGbl0nhDyZTAjf/yghBljfBFTH1dyGLzpj4+bIErKAmSdIbdehwtvdaiEU1PfP7BLv/m0QUWAKIqCCCv6cnt9U55oA+nM4cUPpNGIImms98PbEjWuRhAAAAPEGbuEnhDyZTAjf/++AcgTWtokgs1Kg+3puZmyB2zn+Ft34Z7zrHnHnpGkVel3jgu/cRJkseHDVwekGctwAAAFRBm9lJ4Q8mUwI3//vgPUQh0QPQquZqwQfnSdoeSCQCzHAsvi6Sa/8leZL4+gZ5agcDt/8ZjU6OKKewYE9oGmb98U/7Q9w6RYYgUyseVRyuZy+FRnEAAABYQZv6SeEPJlMCN//8AoC0RERAjFSuAaDWIHfxup8AQJpTK0TZjCJkqKKVH0rnE+j2G0jr/6PUOpHvIfb9HtdnRu84uQ7BUI47R1ptvzy8Ih329lBpeqV6wAAAAFNBmh1J4Q8mUwI3//vuVw4Hc4/EYZkfR9NOYQkN6IR+HLjwNnXmfJ3tFRz3h6PoPlGR9Li8C26dewveJqFV+gHMH01AqXp6nEKs3zf/A6nxbhFlYAAAACtBnjtFETzffU2jm3sH5xlgay9R1QXJfemvQPA5ZQM6hlOnPGFXo84lVDBTAAAAMQGeXGpEf4ESDdqTsacWmS87R3boOwBLmPjr9HbKZR3BTbJMNFk7UaY4D+7PnY7F7BwAAAB5QZpfSahBaJlMFPG/+8dPR8sIHRr+KlWnACcQrjiBl9PrT9yPzdGTDPem1rqoZeo+O+V63Y/a74OMloCXYtj7s5ZuipBPqAm6qHry/NiVr+IWLYfvYiQd8br9YnzFyHPfR+RTXTMKvVV/KjPD/ctM2HNXr4HAOamPMAAAACsBnn5qRH+JNiiXuXcPpCfMZDOo7tZbnRWsBJ0tsa5Lgg0pFjK/6ozWiljBAAAAY0GaYUnhClJlMFLG//vgOOEQ8B1SOgWNXwdoVV1ru5CEMBY1K6LZ98XM0dIKZIqwe5RQeu1WVBdf8q18RVrmysbabE4YEFPFSDmDTf9crfxMvgknEd2vk9sWgLG4Ej0lyc7cTQAAAEcBnoBqRH+BB400x1GiEKvi8jqlJCwv1QQndqVvnTA05w24ChMFL28XtsmkzSsUuIWLziqkkynS7jrcWF+o8Wzss8ftVv3C4QAAAG1BmoVJ4Q6JlMCN//wDwzHBCj/9cyI1iJwiC2YDq77/ovgCtBWTqAK7Qd4tMo/sMxMoLVehiq1uxY67X5FKvvfdZbk571Wc3XQkxwSdzgyBXFfYMHtvUybBlgS3PwuBgOeOBptIcLchkHRh5RWAAAAAUEGeo0UVPJ99PatKaSqB2VOb+Jj+9bQe58CI9t/mO2D7iXRZ6LVPwa0DZ6RZvvddI6Fbxs5CUyQ6tXX32iUFx46XgEEzck4WJrjn27/rI7DAAAAANwGewnREf4USysNoxKlr4GH/7yewiW8E18hvueS8dKXaY0ZRvRxC4LmG9VnnRggXFRoWNTsMGYEAAABRAZ7EakR/h7DkMUOCGBJh+lP55SdE5jTC7V7n2PYNjFca8HhrbO0wW5ZpVWgnD67LghASkduDl4tN3mv9lgWKKDVlzuDELPcvXMd+GCzMWgyAAAAAskGayUmoQWiZTAjf++FR2IObN4lvte8kZ2zczUBUIke9zcAiDpc0tIGFXyhAP1rxFfYV9NigpvIUKbs+TB+6WyVSWBBcm9MFQqwPuVSWkoeLyC9MIaVAuYZceC89lbP77DKjo077qt64djt94DKl7X3TkhZMKCZRkli9PMehVycRIJ11KVnyazgfI6QFyZBWKpZ//e1ExmOA65u9Rb7EpDZypnuNTKWfi082KgLh1cxMCfQAAABGQZ7nRREsn3sgp5ri/wKRgQ7Gu3vI9crS3yfW0cNACqpw9OIocYE6q5EFc0ZcGv9rKrcElxkw0g99wSSyz9azLUdIyByPgQAAAEQBnwZ0RH+HsPhkmMfKXjerJPWkzSzLvYro3D5CmzUrpe93QDwpOs1Z8ue1cdrhkVLaJ6fnHRj/QNrVfPJjXQ1luW3r+AAAADQBnwhqRH+E1CtCQU7NGQR/SzEljM35RLPhIf7CuGp2nH7PGqbYx4k/MVtURJNrZRoSWG+5AAAAe0GbDEmoQWyZTAjf++FR2JNg6dJRFEWpLDaHGWYdvgn2Ihhkb4tam6cKVyugm5faVmXTaDsJAu+kZAz2r1gaf/a+k5hP0Q9ZmUfdwQ9RO6sckew+0/TuDcwET5LS6b0bjlYPX3/3sW/C+vHsLkXWQdgV9S/yoRC2BzqxIQAAAFNBnypFFSzffVS161ytJSzivxBMnWEBBJbIqQ36FF83chfte2KE9wjt18/IKs6+mCtLOAckfgeDBt2owHOLwPzOn+zvI5qJcNySwzh6eafp0V1+IAAAADABn0tqRH+BB5WyGKse6lBSKcVv42bFQay/dKDQ6P/xo+hOzobC10h2D4F5zC/jvsMAAACTQZtQSahBbJlMCN/7xIFmvqIfV0ksEmI6ZqkME/+uYEFZp9EgRyfMfI9TwvOdvYfNMg96XA75e+Ug1d7MZsjl+wIS6P4ARavP29AQ7GgEU3ga/qHpTK+j4PavPhJNuvI3vAVPp+hqdWdcPhLbXeG5Bavj/tfM22MWc8DWUPgPRup92tOMWZq7zTg4fRCn0RmDMdRAAAAAZUGfbkUVLJ99nGI9vhkeUtjd0FBbV/Y0chN5PW4RPFtqyuHM6CMMdaYL9nKS2MPgWTHLmFoV57n1YIpZP6NS6yNWHqbcybKI0qEdvxQUm/iL7kE1gKC+fig+SwYYBKRaTsPgiDmBAAAARQGfjXREf4egL3/U68Y/RuCrxLwXiYQlWXYt5tY5DXkNo5EpiJppoQPvg9AOtTRsWp+ZH8D4KriCBrIbQXvLyYYUk6guCQAAADoBn49qRH+Gdbgx/3UyLHgmWsQQA67kWXbzUxYaBkUPtfk5+96VwwOUX427Zy3tdPCl2FQMk6Iaz/Z4AAAAdUGbkkmoQWyZTBRMb/vuVw4QawH5IL3wgIlYl4N9vN2/EPiB18zVQ/xYvRnR37JySPr8J67t3TEEeawOHDK09GYOfGXSC9g8VJk5uBSTrdp8SDss/hdObxpyozwaea3Hbi8nMDNDeWI3uORE/3DMAFFILF1w9wAAACMBn7FqRH+BB5WDztAwaY/r7Z4kRBTDYmo+gXVOFAHpeYiX3wAAAHRBm7ZJ4QpSZTAjf/vhUdhErlYqhggZ3H5bTa1/vToTRb7froaegfKbJwqLbBW+kSnruoRnFWo3+XuGIHxrqAtWXyVLj2RxmDnonpBATSupm/Er3ieXpcjPxjEdhAERQo1d3evkEK8xubffKr67kXpOm1rXRAAAAExBn9RFNEyff9JABS2FcQ8AcUeHDbcg6r2VMcDhBCTlaLaDqBcTdsiD8BmqZH0SKwn6VGYbJ1Q5LmOm6cs6MiJ7HuWG4vjXM+7PUM+nAAAALgGf83REf4fkGk3LuITCI39/Buk2GiLZ5Qm2lilCulAD0MHZ7IiFhpwQ4cN3dmwAAAAzAZ/1akR/gQeVhtGJPhygpFSbr+MRmIRmrZNzW2fSHOqVlo2OShwB4K5dCDy6EojzpurJAAAAhkGb+UmoQWiZTAiPjDpjcgQswobwLtrOt8Kbn50NO1b8FyU1efKfDFwXZG3fQEF2nXHsOOeAZm6hV6ynSwPWwas2HEejxuO2U7eYWeG4BCTxQsqO5lECivpfGEWizf/HBrRDxbzeeL8AMx7yZyACTy3DlgHl6NdrvffsWb+btUZdfydSDQkpAAAAekGeF0URLEfOK2yXld2veMoVsharS5jlxy1WYfGZ6w0Y2W4SyvoZW0vw2TI0HcN+K1KTShZbSy9NVAcHfLkabTxGzbYQj8WsrHVNjL1UdoGjWTy9FdZl/uN/rXdkOQFudoHdTY917al0GV88Ot6qxJTUo5a0bNZ1+ujBAAAANAGeOGpEf5iZsaTpiYK8Ckzq34eYNkr5wTjqRNgZK+2VKyFYuAaNfhszOaZxYrXPiVM5moEAAAViZYiEACf/vgy3mIDfV2I+c4OQeUCbjkIqauQ7q1hI51t1KK0Ty/hpqe1v3iAKauIQOhjuzFxyMgx7FbnM/2jBWLSkwzPM9uWS0H4jWSakogzfe9txEs4c+/QsLeTPvLzlrrhoRvk76vX6F462KFV4lQ8BLFX1JpZOU12xKsmCPjGaaOJgyXQiOhYn2hWzgcqRNjGX+V1k8oj99bE+1vc15C8Q37MxkeCbaMkDN/t8mxkCWv844zRz7jjHMSn7OeiLDTN9e7Cv32YENkK+X2fgPjTA2f9fID7IG6+92XMlKB4gG0/ewdEoHPbTbV5cbTze8c/T9wMuzhP10AuKfYQG1eXnZ1ykp0P6DRPy9Ld6m+2H6X5Ni7ssRrEN4TuoAEb5bQp2TW/e41VdMbUA37GNZ5x7DsYhq5D8keOph2B2Aoqxc+qPJPC1eN8uzF4kT9NEbttz3OMknmnGRbLqo9BolsBfy9A61ilUr1hK+lfVlTjk05qU3y3wHYEc66DTCaL4HkcZGfGwlWCtebzJGH1tSY/TGvlzwoiR6D2VHDw4Qd/QL3sXXc+s5L0t9Sw7eMZ3S47uKfNx/FSmsxh8Vu2q6ia3DDbKLB0cs1CwbSXU5GqIQSYlNKNsVMAv1OgbV1e/q0to651ASYAj8tdg/GWuSkyEBfqOW0VahPSNiWW6MzjEdYDBORzkPU+fSsEZX6EUEJWgIR6Tk5qkRkpHpcm/zO12eo3M4AhR9Cx6B/Bjk6bjXMMI1JbvT4s3OmvDMWhm4GnlKV7U+z5WfoAn+r+XEGmS5H0VzKkjL9DjNzfvJzjcKRQzSjV9OMxXeO3VBTd+WFAIv2qLLi61YRg7R2i0zYBkZiZjznSxv0Oh3XHIYCRtkOS6xY+goD6NZWh7U8+5INbjxbNW9vW0F854H74WEbjrsIE7vUh38zBz8lmukAzLqGi6lSYQVXRAstPq/sjujZpTHuizcX6q5h208Hya5GB6cWOla3tMwJs10rHPZhITBZAHx9Dfn9xOK6H7+wi//UuozeCF596Bj/46ApJUXxnQw+iPZClbVtv32bAEWwjAq5WxRo58IBlPczqDbMQpJ2v8hJZ5eL2AFeiCS8d28CkwIevJBJczJ4bl3YAfkhz6EvtCfS59R8MtyqJAHmYMxFa2BzAk4Ceko+cmThvwJsTo9TwWz6/y6nqWgaEvkLNZuyixNMc/cXcRwsDQOLI4ZyMttZLESIVUpeJZ/6cbN2LT812AnxvEtv81Oo9jf67SuSMRsG6/uLlpm7ay66WEKT2P2eyv5sPPbiv8phYtyC/hpNWM2L6InvRMw1S6W6QQw7ANE+NOljaO2+xYiaex9bOmfqoAKKcGloWbR978SVkvTCVMZ+HYiMb1wAqx/PF/B7YT3YdqVhfszlXZYvWaSPMzU5okM33PRVoLBIGHcW9rG42KMvtm5tZJTaxxpsdsz8X16wDYWJ0ZVgl2FlFrFiwpttQ+K/SBAfBbCBjc4w+9++Vy9ubBe/R1rfhdPo7EO8Upn4ojmO4i0QAQPk/+Njyc3FUAfeGBrW8UR33ylHcXQhrcCZUhYC8/6+rQTbAW1pErnFsaRgQxa84DSprqtocWc1KyDwYQLDXOQyRgEVTWwp0C4n4FCtkK+u/u1SbZZ+WlsAmGCMhbvZttIOR1dvnyda1VEmJ0q4MLNPcwxMdDoqgdP9/cDbj6E8rNAG5pJQbjF9KX6GfGfetGbQA0fy+0g/NVu53wvOeMDlX5aZKNl9hASlsYvBxoiPV+IOBqNBkpAft4IlL1zcz0cOUncvHJKuevXoD5KrbBNqehopgoUlzYTvetOhot3VJAEcPkLAAAAERBmiJsRv/7xJ1U5AfCBaOk8iMcVZeCFVxqmruwhZ1VxVh86clNufhi8lqRRs6U6DSeCULyvbJ8s2QvtuTrcAMn+2m1oAAAABwBnkF5Ef+BB5WDztAx0jJ7uwLOc/joWyLAt56hAAAA8UGaRjwhkymEb/1jexjsxsfi6hr6As1+ItIPaaxdtL/Psp4lmstRcCE43uyBptRSOlaI0NUP9+Rr2RU9EU2/XAcFdqfD0kTm2yOmA4hz5KtaLVYCIyYQrzAIpsbD5TTBRHqAEIhPQuPAFKMNcGmm6KkfE8h/rUhCTb57Tr0SCkvByHkqBQxIkeSnZE3RETT0gzu8clkUB5C2/f9gqjzOEdr0tMQy028X9g6I+jHHVIcOQqoQJqlYLTVMEPnxTm8swUhbukQQqB/V1CMcORCC/q5HxRtMLYD5qtCkgPY4iYs9BYTsiCpbpzh3Z3BHriiVt8AAAABdQZ5kalPJ/4RtMreUjvAg87ctPZkwB7tGp3pNULncn75y63gaBD9ooZ+2Z5BGJuEzAt8WFMMMtfOq7F1nffEJMK0f/ywEmY+FnZVV4/fl1h2Pxh96wGJFYB7eE78IAAAAMwGeg3REf4uRpKqqyH5pDz4TvCO+HJMaGqrdtS376viJGvqRNtOi7CrvGS9ThjJaQu0ZQQAAACABnoVqRH93OX2EwrJauMv/aYKWiM4dLgfPA6VlKtP+gQAAAOdBmopJqEFomUwI38epjsLFC/F3bWYVsVjbCSDz3/NFFz7EJcPUsjeAhu5g0HabZ/pTdPiudhVDGiYxVKcEifZtHJ0mwAP02f3KxpLwILBoE5KpKZ3ak6k96u9QX0BYIZ2zDURThz3E5Ecs7TFwmbLYgIzCpqmx3WZ2+nHSOQT2luDm0mS7tC2WUBIB1L038q70IP7mmgmcfxvHDAXuwIgQQvTDkNbDDnb7qEeNphuMepAZpH8UGL526eLVgrKBJghOlt/VEGAzajEUl3BhFpPgSsX/QRkJ/xFpACLLDGkb2Jne1e6/1z0AAABLQZ6oRREsn812DBeerx++BXNh97FhCAHBFp2YmevU1UrQGvZZfGeZt4lRrbLNgTpILEiupxGlUxfWlIrjdvxwBNEtx8ud+WR8eT/AAAAANQGex3REf9fXXAqIFgtajW1njE1bz+UortEyKI1hn9ddwwTBBszfncLM5umhE6wLGy54PXqAAAAAKQGeyWpEf4TW9Rzxc6AcnNg8AfIm2Rmp/I+vcel00x36c2AMOw8muHw9AAAAsUGazUmoQWyZTAjfbibf1bDnDmV6F6O/WnufALeelmNBoakQ+RD8No/50k5DsmJ/O8f5+6zmOB2dS7zL9WBPkesZ8MQ3/0MksLiL2syu0Os4+DX6zzVah7a75x5i3NZB9mNahlXbFzLVLgWbDFo0MPp2ztyCmOo/mDC8YFYHXaGQYAnLA1akkBw+w7m97niKcIN0hcHT0r/8VhoRhstD4zpvXRFwqgp9DQSmsJ2LaYfRqQAAACVBnutFFSzfgcsWP16iAbHfMo6XUJpab95jekM9opekt48P33/AAAAAGwGfDGpEf4RTMK9Ia6xDXExmf11mxxz6JBZgFQAAAIhBmw5JqEFsmUwI3/QSy/3Mn9LlHVQzmvoMZRAfpfzmYydz58i/v/K7UpAsceKOBkM39soM37DBOWqDpDJYUzpx6MgjQXdFpMMlD1sF3jXfAp63NUJB2HxYwKAS3+X42lams0gG7HPbT9qoElWzMJ58INMDP0k6kDCkWtzZGiljRWKipkGdSBWAAAAAb0GbMEnhClJlMFFSxv/0LB3HWhFl4F9qAxu0d+LmsKyTRhMet5l8ycg7FLJ6bMwnHDvdth6Ny3YdAsTsiP7Q356Lcu08PjhT6If5FExZc2ZvHjw+Z0tvTWfIt2djM4q2QM2UtcHgN2bP19edptftPQAAACEBn09qRH93OX3cYVktXGX/s03NlZ9nnR8Um9cfaQPg9mEAAABtQZtUSeEOiZTAjf/74D1Ed3Fo+hammI+tzPVBVB9GAMtkRSmX9W55hcG3CFWWzlEIUNNLjUOeelsGZvhADsD+Ofy3nAN5GVG/eZg36oP7ixXBkBIm0aSeFatlwVCbRiYm+S35yo4wsvnxxokQwwAAAD9Bn3JFFTyfgOmmneDD0Jh97FhDeN8cWWgHHL9ZXuuK5vJRYDuSt7EY3WZi5fyM6swhMLsBDMNgeDfcSyJYY4EAAAA7AZ+RdER/im4U4K1F63xN3y7TWMTYDhtYotw9N+3yCyCwG7hbmm/x6VomjGrPQRMq2kgD57PUgsO5+FYAAAAsAZ+TakR/hNfANQiYzYypltEIwmskWD5UNtzWBIJZzKZTZZF3JtrR7AcrpL0AAABCQZuVSahBaJlMCN/77lcOEudjsGvy5OKoLXdnzB8pMXFbfBomeiXX0yT9Bi3pzVu8acc1sJ8hB4UGgAg1X7FrVSaBAAAAPEGbtknhClJlMCN/++AchVJ/hJ7fwbdEstXIOKFMPOfow0W1J9MjZSXaU2JvcLoxTQ3MkZ+p/N6EF1h/4AAAAFlBm9dJ4Q6JlMCN//vrP1EbRInimhmXtw9PH/HLVA1X13BZGuO5tuR281DwTZtzSQ+X8/0bJQ2SirHIQ09BynMc2hGsy4/B5wCyAIX1VOyVwH2zU/wrNxXXcAAAAIxBm/pJ4Q8mUwI3//13BGajj12djqpakBtTio3DJ9DBFwgACSqtkYv5h5ytASPgOdOqkOdgnQKmgWojr1xZ/Deftcan6H9Tw3ZA+MEWlDNo1ZYBWRYty0ND9KfhocFL32qclLBL9NniRJkfFCDsnUhqrM64RBnhBgh3Fvlja3uOhq+MRbDNIEtu24O7QQAAADdBnhhFETzffU2vo9kHhP0kUckkmlcOZ1qvSaHsv7FYPzfK4q0p2kKE7Mm1B1C8Hwvzr66y7F2UAAAAMgGeOWpEf3c5fdxhWS1cZf+zUAjrwmtqgIPXmRHv2k2L/2uJj4e2MaADqH/NY1U5BRBBAAAAeEGaPkmoQWiZTAjf/WKRlJMqrSGu8YMSMo4NdVTVxbO2c0DfzHpWj3NO//rb0lcIQgLmD/JdeB2LmNOXjmIWgYrNK2njO9hbOQv+tHayjtl/SNf2g2g4kXhlY9QRMuPkMSkCrA4mXqQYXjex6M4K3kGlpFUGNYB2CAAAAFFBnlxFESyffGVVRtiU2M5XSTBlLlLpuEazvf6T0QUYwco555QlcaEH0Dj+hFfQ6SVisFfojn73EUbm8SrLs/HDJEPq2aQ8uqOTVCis92ohdCEAAABFAZ57dER/inDB4mEZFozMu4+KKCyG++hK3D037ngCnbT0Imo7MlsV5lr5kGDygb2dj6rCAP+8m8TBp/+LfwVdRRtehi0bAAAAOwGefWpEf4TUKkzU19rfRxca9Jc/qwQnAu0qRq2xIkvrtGv5GbqxB1VczyvDgln4tgvLQ64l1ZbfxKHVAAAAekGaYUmoQWyZTAjf++s+OCFIz0x01M5hRDAgQ1SNkLroyGxDyE3PDxh0h/mEHwlASDMpnpSSDpuFfd6ZN1GJRdp7/1NKZ0BSML9Ri+WpFVWYy+oWu7FUhBjVLBsl4nohvDKmz8vEbpKUMsdmHeD8TfRVnZtvTlvIPNsxAAAANUGen0UVLN+BHt1YNOwIBo8J7x78Wox8hzQych4IV8JIX4JAHVjfP6+E363LAITjUT0/2RT3AAAAJgGeoGpEf4RTMK9Ia7m5tNYSHf4/ZtvhHWvqCx8lGWv53wQ3Vc9qAAAAZEGaokmoQWyZTAjf++A44S52Owi7Ki/xHuS/BZp70AE+JQqCbzs6/5q8Bi6TLqq2chJ4O0HsA3vaJtaNjwJK08NQv3nUM1lYoOFoLcHLdvLTHCfibUiXw/+uN1EJnDUQ731KL/EAAABQQZrDSeEKUmUwI3/74ESoS0QZzPsXZCCOIIaZzh4p+rex8ji1xPYtKKYeWQssfQMp3vKo0mHw0vJQBLiKE8ApAr2l/ILTZvWEdw0SQbXmWBwAAABSQZrkSeEOiZTAjf/74ByCUGSE5NLsdewC1Yf8gDKRgA9A3XGNkTmRAIk0zuqTVUB97w21ebY1ZD4A5JuX4GJerYYG1yn3M9Fjq85VjxY3yOfXgAAAAGdBmwVJ4Q8mUwI3//vgPUREOuU+W50ENB4ZTfPgxJ4R/QAgLWfpJ4afx0kcK6OroqS7DFWyZDTIp4V4in/vBYNjSauh3e2aTeWXDJZfmmNfyGj/jeE/EeTuP4327cpUG6wUjxVSODOdAAAATEGbJknhDyZTAjf/++s+OEsdvlfnggh/lGSNWpaGXz9PDwU7DfMzDnOAUqnRhTRcsjlee5HHdMPWnOhXPk2yc35P2EwrAHuNbsWtVJsAAABJQZtHSeEPJlMCN//74ByE7aEx9iXfU0+Qi7cT6qjr7osl+ltza7QEqjKcHDYgjA8iNym4TNocc8MUDP2ywSpM0nUK6tcXIpAVoQAAAF5Bm2pJ4Q8mUwI3//vgRKic7NNky66/c817e64Ip2yRLzp3SkaICilaCnIprlv8SQe7wRAh5lXA9b8r69v0/THs/oLcawFvELgXfQ9ZMLp1MwPkpBjFsgw2mzZd4z68AAAANUGfiEURPN+EdaQm1vW/TINWOBhDWfqLGez+cStDKozqO3ZaLVmanISgcBX9Quzd3kidv0xIAAAAGgGfqWpEf3+cgskCm2JriciFWFR8KQwp4tWhAAAAXUGbq0moQWiZTAjf+8P9YCJ2LEHgSkHxWPcIhuhg99q36oJg+yoD2YMOzEXYsm24bpU7e5HDTbai/djWZ57WEvP0X74/hSDvW98TN4gQ8FYDbsgAomfxCR6enD3bgAAAAIFBm8xJ4QpSZTAjf/vgOOK5ljY+B/+L9+6Ag5nJ5NgD1rC58rPGzwaZU3TSod4Momby9YNlv2Wut5rt0gK6/Y3YwVJRjdwKaYBRIsbzg3k2OgRes4iDLyq2y1We08s8H5cePlc8GGSMvtVpkDZXgQ3ARvXgVsnai3RWNmYHRmfNrWAAAABGQZvtSeEOiZTAjf/74Dji3z+EL3mriZEn7qSTHtaSGpQSz7vu0Qa1DTfh90G8QU81BZ+J+YsG68GTxDc1tNVG2FxQvD3JmQAAADJBmg5J4Q8mUwI3//vgHI3KkJj2o4R7wAdkoxoAPl6DDEXLEUPxOrh3d5SlgY3fxMeUwAAAAHFBmjJJ4Q8mUwI3//vgOOEudjpLG75zaU3JffQF+0lnrfbIzuF4nbSX2KhtapcwdzSYt5xk/gC0dG6joJPOtvfYDyHPdosm0/eNEHkREpQU/Lx3HWTyhi/QVBnzgdkBQv8ovz6Jzex7/D3vMDIqrT0wYQAAAEBBnlBFETyfgEm4jKr15Rd2F2D4x2LPtpGUSKL3xrnxXz43NEoj/PN5GE+rLyJeghmbwz2xiT1z1CL9RYYKsz5tAAAANAGeb3REf4puFOCunliFzADLMuR8YF8uYkQJbFv5RTB0u6zvuSXxrl6rDsz+i7vbod3bBSMAAAAvAZ5xakR/im4VKd901cdgCtt0Mf0BD0pitdn5JUOOTT1aiDQ1D71X9AF0lVJsK4EAAACQQZp1SahBaJlMCN/76z9RUBUcvPliLciIvjqR73/pxCfQ8GX7jBwWgVXkAwndy884ukxbTKeablPB5O9VW5A2hak9PYq/Z45Oo6shsBUaTyA/cKmxsiL7o+G5kUuQ8MYCGvoKy2DKLSe6RNtH4L9/+rIRsEmOj8yoDy4VgxHPKb/WQm4tsTdBVn/g3+/tzcq4AAAAMkGek0URLN+BH0uHOpR/2lJ+qsy4cNxEGsg1Os89HfXB7kgKMsUuWa6wa3E9Cvz9Vu2BAAAAJAGetGpEf4RTMK9DxTPCJrwplTdrNLTARpp/H+2xqo210z7XBwAAAHRBmrhJqEFsmUwI3/vgOOK5ljY9oQx//XXAhhGlJBmz1SK5mkHqP2hvO14o18d/5fUUVISry+PxljiI49Ncla4x7BQYXDweGuz/vR/EY/E2El2sS+X64RWh5rFNorrfK7VEZzIHehE7YBLUZYyCOgseeuPE/AAAAEdBntZFFSzffU2vo9kO0Cnibp453nDAGgY00PiyDpgtljOwFy8YbTMWPWaJg1ksU2ghAt/uAQ6k2TZLipHJGD5DZkUzEzwc4AAAADUBnvdqRH+KbhXAZA8EAuKk34jx10uATWjooIQO4t7y/vwPiFVePKe96ekG16OebeqUuepNxQAAAJZBmvxJqEFsmUwI38oIQZY3yoDaqyVEGIIzrxX5up3P+yHPuwxQuIh6RmqexMGn2JW4TA6S8a45F6aXSoHDu2LRXHvNy93pz1lWEn8VRiuUaT4hawwlUHpq+vaVckWlKSp4IpFdIAZg3tMDq/eCxf/2BMIs4+tEBE2gLOsQ3IOpZN1Nbh8+Q39DhHggVbCQaqg1VnbmMLgAAABfQZ8aRRUsn4Dppp3eFORjRG7B5yFTmmIQ465a/eCYalbPok7MkL2t6szR+A0CkcE/nExx730z4edybGjkTzNz41lOg+CSJyI4X3ooGwbsjFtXwRjxwm+LP6VgAP21jRUAAABHAZ85dER/gwyVRD3jGC4pS6OxRSsTP24dYF84UyHNnutQawSpCF2+m9LZXyJCC+IxSaMtdFhc4BHjgm9gCKjgMNKJ8wtLHlYAAAA2AZ87akR/hNQrDpfZzHTty0SYpUlNuCqQ6/BRPgxApygEU5Fkje1Z2NJ1LTucnfFcqhZof1iBAAAAvUGbIEmoQWyZTAjf+8UQb2oRIDdyJBJG3p2eSRku//vGBvYLMgzmEIVmMtrZl6/1t+3GqkdEZly0y2btApk7ZYZsIsTNb8bH83WEuTPj5wetHe/hfU1579gz9eJmyWaFEN4R4kRUT5qpZ3V5aB0GqdGynWTZ4tnQlWYylh6GVVEwLVvNFqWdqS5yTosyb0fNLjd6vhS6+1ch7S8mxV5lmhCvgGRXTY5b2+5TgsBpT8dK7WAunqXgh1UZsPH2jwAAAD5Bn15FFSyffGVVSHdSb2LfjLSUp6GX0WfYywKsLovtWybvi56QWoaMzSn7o2XmUV6WjcC/zKrkS2qukaOC4QAAACoBn310RH9/iXZSQOIbeTW0680W96zcWiFcp4kIDofTn07/nwFewF7sA6UAAAA1AZ9/akR/gRCXiXa4/WbZ+NyaVYeW6Pf/aBVzqs5yIgtGflctsp+WlKrincEJHtYZiVyZ4kAAAABrQZtiSahBbJlMFExv++BEqXFF1Gh1OTCwItMBlK1vXU7Y1ChujfwgDSEjxp0EbQ9CFRRs3Ld4UN+pKee//cFrSip/MsXQzRJi6a1nf5wlbHjxkRq5Gl8WS/GeHRd97LN6POgir71VTPszaxAAAAAoAZ+BakR/dzl93GFUqBjx10gSJvD19gv7sp63PpHZI6Ie6pbZTy3WoQAAAI9Bm4ZJ4QpSZTAjf/vgPUVXdpG8vg4rlnY/ft0DGCD3fGKCJUJtOkcYpkI/rGxid5A3QTlupDhXOLGx7FAh5xXChqo0uaNM8+rYrg+39wO/wMNbPlztdpk9fui2dj0P/mWfeH1Ug09qs5t6x/FiWlyKrutmeEX6/v0IBu5vp9sVb+RmDgKXu8CyfSKKiKb4rgAAAHpBn6RFNEyfgOol5DugVzX2orzM1hG+r8MU8NyTphLAcR7riuJhLCmXsqR4r+fP38yAn52AQJ3oVcR0wEHGBUJOmK8pqozKbzUBCcuWr4/sv4PosC3bNQ7d4cLNL61+4Hajc2av2ig15vb1u1wugeDCXJaiz13UgQDQlgAAAEIBn8N0RH+KbhTd2l4QxwqahFOEir8WO4bHJp8S/K7MsINQWNL/SSuhWUuwzya5TSxVSamU7UEGgtWToCyOldc6FVEAAAAzAZ/FakR/hNb1HPFzoSznrKuNC2BC8XbeOVatSDLSeW0TBFsdg62GXXU9E7NOBwisegIBAAAAsUGbykmoQWiZTAjf/AOiKyFAmRxUiMbEE4XKUlvaqkfdnPFkuSrv/uZ3Z2K7U4O9pzkfIf7vPupN9AbVD8Dwu0MYNhQUfSc7G5a1t/833BIPTo9W4fFuf2qMVYdA1azZtGYexCZGVNJ/9/42W7lwA+c3UjwmvDt+IQ8gN/ySJxefGd+6V4IysnaZg/17TFedcsedScJPSmsaAWJV2YZr9PQrnZWmT023Hkz08fPsg3CfQQAAADpBn+hFESyffGW008aR19uICCzq5NQWZ90NfrmxIeDOiit/PCFcxaNQKPjDjQ7qd4CRyqeYilJwmJpgAAAALgGeB3REf4BsK4JrMmvsK/sAAGwlDBfyHnOGZN3YuTYe9/CpCnyrjEBzHGVjUCAAAAA/AZ4JakR/gRHB0sI2Gg5rX+0O0q2TvDPSc97UgUfQvJFcWSBSyl1AmJl5+F9Odi9A3zNqS2n8PseQabVvy6/BAAAAe0GaDkmoQWyZTAjf++A44ria4Qw4WCiSmMPHSR1OYI8sMHbiuDjzddiGMt5FFfQsG3Xg/s7xrPsrUAb7CZP61yqPvjLICQfvhRYvztPo6E8M402MdjGVc+5MoikSIp9hwj+gs/Mh6ONlkOtd6lmuN/ed6YdaokhXg6/p3AAAAD9BnixFFSyfex1K7S/556qAF2wZ6eiICEv8kLtY3rwO+XcsSPv9lHbtQilvH7t6S2w93Av2M0H+BnDw0AQGkaAAAAAyAZ5LdER/dvu+WAiZLVxl/UCxr+jB0mUBX88+mWTBAoHaD9tT3RLSGpwK3dfZUneAJAMAAABMAZ5NakR/inDB4mCq76bkuBs3tkF946bT823lmJeD00/zi7mcXVT/glUaS6jNfGEtB5R2xabRDQfa38sKu7jTT8gvmqqlDa7KdrZ/cgAAAI9BmlJJqEFsmUwI3/vrPjhLnY6p+6D40pzcI3EESZhkQEBCdNnAeRIPFFVSkNeEcVI8X5XssF+l+r0Xo+G2YobSzOdbqNWcfU6YbitpLAmr/TX0rnj2EMlad7XNyf0cLZ6ug0PPBXULiz9m54l70LAnrx45mASDWufLQkxiUHCgzbJ5tLv8P8fLTViejM1LgQAAAGZBnnBFFSyfgOolVYNR5FvP/wKasCCjnYQlxD8XXZsDY1OG4dlU0osuRUdZ9269GrNsNOazGKWnVWEXEQ356Raa/y7lNVGogDCtXx7LDbgwD5Dg7QfELHis1MwTs25skxKF8NR6CxUAAABSAZ6PdER/gwyV1/gGl4DoPPqyLS9sHc3sEMPzJO+wfSsvfs+SXVSm/M8+KeWiKtOiUBDrSQJ/dV0wHz21x3Nlknh5MfWPAzaoA5JbIyBOBuVi9AAAAC8BnpFqRH9/nILJGk5gad+eLsoOm/85p1qMApGsq5TXcltUPCfAAuJzBXY34TxlyQAAAMtBmpZJqEFsmUwI3/vBmqSMhNsPuU/Iade4QHot3LB6WrIq1REVRUWuTSSkVV07RY4SBhsJu6bBxcmmu/7p7/XlQZ81iuMoEWB0fIuqz+0F9hRriMm+Zu++c5UoVrMVXoe4liL4u9dlcciYCOKVOu20Yyvw0sSQzjqY0OXZVZDApK7OgpTk3QL5BbxEbzQlLUIc9ecQXGeD9042d6e5Frg8c5maepHMbf/XGm8v7g+lOh4PjQiOo2rPGPN5ca44oRCfG6hjuP4Xue0SoAAAAFpBnrRFFSyfeJlWxHrflwbYugGcTSy6U0LQpmxil0xIJFwT8XFx7PNwczoCH6Fhbo0X15NUcnc1rsuxZVBBQnNcbPv70rYrdQG974QEotol0nT2TDHb+Qv9E50AAAA6AZ7TdER/gw5gGJg+c863gfTJmm26qPpKYjlbgotHYmhUqPnWNt7iyfWEOKe9buyfpqlqllaZlsDSnQAAADUBntVqRH93OX2EwrJauMv/aX8LwP5QXIKtjPfqQBiLL4Nctt04kl+GBFA8tTN8WaPF6LkgQAAAAJNBmtpJqEFsmUwI3/vrP1FV2ybca4m7Qi+CcCwEN6Sm6YngKSVvXnUfM7sChF5+nlbWXM16e6m0eP3IRF/opSl4/wJ8Tmp/j912lrTy9yAirrx3u2AijNyh/O2sWcv5ZF3bAQNrmKGJ5rRtgcXJ2EyIs1FHVEKdlaUfrkIuoLml7BajyLWU7EXhV8KiyXp0EkB1IWEAAABIQZ74RRUsn4DqJeP16izBqJDziqWfWcM1ER08s1jzQXvUMa1oZMXncVBGC9ZI7nV/tGPht7tK805Y3P+T/QfODA/eDIltS+iRAAAARgGfF3REf4puFOCunliFzADLMuSGFg8oLzH6dqAnVVV0V/5aqCM05ibhk+ZkUat4TgBxzdHsk8qTta9BtubHdd0z1ZwtYcAAAAA5AZ8ZakR/hNdbIqL3l1+IwTqfo4tGdH+aq0/IiizAMo318AwUtAm1jEwsljCSWs74wPi4J4g263FBAAAAjkGbHkmoQWyZTAjf++5XDiz2R9bMf+PSkj74lbvgf+CLlKO2lnAEeCo/BovEKd8MOME92netTHHgj2Xxbi65aAN6935tsknj3pq8q8prCwziwqggVz1hNBwrQ1X6JhikyanOkZST2ZeJMNQ7Y9f7rNXp1lNxIE3MyU/nHtKei86tCmCb/jB4vfQ9xWLRyxAAAABBQZ88RRUsn3xlxPNhKxMxOguiBxRkNmyuAT/aa6QxM1PpGwiqB8JRZtraexcKZFHXdXzPREf+tWX2PRuHYEt70OEAAAAuAZ9bdER/f4l2UkDvsZGNAf/yJYoKrKSiCC4Thsosst1DE6Zt2qgCipQ/85ZSywAAAEIBn11qRH+E18A2k13OofglWp1hqldV5ObnQoJScCal5gPbTIjyW5ljrk8khorjjE0JhwAgDO8DAXOtPqLGK9BvShkAAABuQZtCSahBbJlMCN/74DjhLgyOzpkIzRKgaw5w5N1Q9CI2CzFs+eee7UKL04QlzU2OMinPtJrKhxVooB9zLXIdj1wU1RccAL/RK5Wir3p4Z5Rhz1udVe6K4shR4iHm7EJ7ua8uxY5WA3d2NucMe5EAAABDQZ9gRRUsn3sdSu0vuq3oSzNlAM8OUuryflbanNQhQdg2gcKDlBVHwMruBd8g/Ot4PFgifoE/KP6Ohum7YS9dcg57gAAAADEBn590RH92/rZMJhVKmtJ58r16K0MJGmMWg8xAuD/vkLLPx78bVItQ7HSx+t9TsrPYAAAAQQGfgWpEf4puFOCunxZgudv849O8E79Kzy8NTAnO6eR26TxQmY+gxXfNtnCzUB22XkPe8BFERNAs2Mkr21PBuWuBAAABDUGbhkmoQWyZTAjfx6pjxyne0Ht21MeCz9epidJeHVRgt2OzsUDZ9fGji/m34qapaDq3peBBbvkhLBPU43FEf8C/nr02F/kqRsg4pwf4rJ9PEH4DvxUg6fBJQ7D2PXTuvDFf4X/10q3JA3YywimHAzkd07MRAHpUkRpMLgBL8pLEMwRhbPWVpnT20qZSuLPZcdFVIiFASoAoRw2NZovY8cyAqYzbIPWrSsq65p4K76JrBR5ZOJlpYoaGRw5W0g2bASg9QEtRnRj/wfIJlboNxF8URYCsrV90PwYzTaGML8kxduJ1HEhaPv8aDti8cRlE7XHaGXPFi4RmNryyxl8CqT02GdWdUGWMt1v+HtWAAAAAUkGfpEUVLJ/NSlfleNWMCRxiLP05rGSLgkzLFEKRmQxMuJmvzGpyeNPxT3UwwrV3v/0iFzkqNLK3bHJnTrwhtjxlzfd3LB7TBGgwkOn3N18GCDgAAABaAZ/DdER/zdG63bG7KwCn5wCGa3osZLctRFh023EneSsmn3MzWC9KqCBDdw8pzx3PPwNQzEg6Gn5QdiDUHnMAGDZ3RWqH8/GR/hNAU8HnZ5GQH7ni/2AUHpBBAAAAIQGfxWpEf3+cgskDvsTAlIkk3Qv4QtRWQBJxMPzllMwDnQAAAPxBm8pJqEFsmUwI38ep3IcnxuTWqQj8g/rgSxKei33dMugniCkcN8GPABf44Tw8t1Qy3eEsRn18PRXxXi7KFfT3IERwMDz7zpA3QCQSPyLpyVwXLhM7PhboSU8kbfPO1nIY8SIm+UMfGoyE0RMWGKM0JLzYaywMKx9vosxmymYX3DSB8gROvjqj5bAU+Je92mJJE3ZzVPQnpxpwyoFzVDoLe998j0tIsq44A33luvrT6F1Iu7OaZJWo7H1e64kDqoUz6nvhWH05mQLuC2M5Y6ktRKSEcGarBYzmLNzmHh7vugEWRyOLzOV3ZQVWiOAToNJ5ySrGU9bIN9xzr1kAAAA5QZ/oRRUsn81KWB285AKVSC7RFY+doT+AHMerKL1ioOnmHy3pPwHyRwqUvphrWUAjK3ENTLIIi9eYAAAAJwGeB3REf57EiBPUsI2Gf4tmrnTNH6no8dZUfNSMntiCE3jpMKWXLwAAABwBnglqRH93OX3cYVktXmd/7NLucyKFso25IGF/AAAAl0GaDkmoQWyZTAjfyghBljfFGc2weQw5uG12FPp0eBzDiQcuXdbNVOSPW3hp8N/oJG5bAO3Zb373SP/YhHQc+p7lx0UOuTmVc4HVR7yuvu9I8eIWq8uvvC2yoCvXDnsvL6Ml7y1Qbsl4kWfJSSnIKW682NKdd78et0c3c2fAD5VjJnXMrm9zYiL5tXnYF3cZwssrIHOcmXwAAABdQZ4sRRUsn4BJuIyq9eUXdd8kUEVxBXvZDstOfWiaasOOGXz6d2BweVCpPVueVKNIzwQFXuFUU9XBKpnS68hOjBnpCQT4QYBqFRqI+fPbQVzfi6cOp2mU7sSvCMpmAAAANwGeS3REf4puFOCpcPJ2WonQTLkhtA5MCJGSAixNk7uE8XY5myUHB0SlZsop1BQ/gkY1XDSgyIEAAAApAZ5NakR/hNQq4wqUtHk0A/hBBQEMC8kIK/T0jJcco6ztGO/MojFLnPAAAAB+QZpSSahBbJlMCN/77lcOEOGGqjflHZLVcFnf06lApNkw1dyuEV9YbzJBRtMBgx2f6FglR2zHhzFYrtQrtFgGZ+DglcCrKvwM3YU7JhUBsKErOd3vij9YsCRfhH7aSRLCbVXwbVZgzfTpOsK2Ik5aDiH8Fbx0ZTn2TUf8Z05hAAAANkGecEUVLJ98ZcS8h1iXBUMZDLJz9G659Jpfnv5AxC0q3dUwbyNGgKrSkrRowhNA2+4791C3SwAAABoBno90RH9/iXZSQKtBXKXppdd1hcPZsBwcUAAAAC4BnpFqRH+BEcHSwjYaDmtfG61R47GinqanteSqsWgtBpO/fwqOA9xgYQFYm1iBAAAAbUGak0moQWyZTAjf++A44TvouN2t1o3r0TOX19VkdemfIkJQNjJtg2asYKxFD5FgdXiAetHtc+GXPVttB06G314qyFnyAAJOZ29a9/0H+3xR6ynmfY/umA2SihUcWtibuBfPpAHobrBH25owHGAAAABTQZq0SeEKUmUwI3/74DjhPbILSOAstfUwFE4bEuwQsadfXqn7plGI0oW1t1Y/wn0dDyxQqMCm+V/gNYrsp0t8yc1kYM4KZZ/Ep47KXKTj0fca5jEAAABqQZrWSeEOiZTBTRMb//vgOOEuDjiPT6VQ+iwpFIBPfM7Vc576QnIa86wIKthdGtodrF9dyzF6wRQCkAGjjRZKD9pc+qeVKQR+yM0R2LkRny6G0uUgLJfUm1EaS0RgY974VyZDjM1LWcR5wQAAADkBnvVqRH+KbhTdyySQ5ZqJ0+qpbnfk0JbeSoqS8GFf39+Y3CXtR+VXfNwhOcPWv7xsKzRkAg4LHFAAAABwQZr6SeEPJlMCN//76z44S52OVALIq71Pd2EwQJX7/N3fXhfhz3oxlU0IeuIeFSImpqXiYeUYQ8aIei7lagfq9dJeClxCf3/FRuzxdmLsPWEpBDNcG4b7eaMqmz6BvHDqlmnYyHJrEZ8IZKKB4lJ/QAAAAD5BnxhFETyfgOahWDclf403X4Twewl66rB03AAAR9ia2izX/w64YBBi7Jz+yLB2fmcHmEV1G0PRn850sHSM2QAAACoBnzd0RH+DDJXX+AYxXI2kR0rdOucxMrUy+Pe9iIFKs2gPngM17pEs73MAAAAdAZ85akR/f5zgScsRbQaQhHGzw2FWBvNSvTqWvH8AAABYQZs7SahBaJlMCN/74D1Egm5pgj7LpdrUXh+DK2shFxcZendNqyb3B1+J2NOtJQLO0BY1+ExmI8pwJ6OXctPqDeFKHLYODoMFdurTJF+a9y3NPu44TpmggAAAAJFBm1xJ4QpSZTAjf/vgOOIGX8Dywu5b6HQtuh3ftgTaLAoHUTu8zVtxvvBrwf8pq0n+w4fJtp/GtGwKM5hmO9cHCkxx0FIpHiSkAHgDSjdJ9Bu8/qv30MANfg5dNOeEtmWEqPoExrZqFj8JsKEB0TfzZtNFAYsTE01g9Z2gRAsp5MxV9DO5F2sdgezsL8ViWTohAAAAaUGbYEnhDomUwI3/++s+OEtQtJaAtNXXgd4+LqdFPoVeYoa3U/EzzhI37OlznPU7rmyUThBJLF421rLihpccajuYn/QNJFilPGf3pTu2b7OjUvIJM8X3UCvvvycXDv1blPNG3IoYzzU+0QAAAC1Bn55FETyfex1K7S+6rehPNxk1X2rNShPMmLN749hOCs32PoHwwGQnuRiu7dEAAAAiAZ+9dER/dvu+WAiZLV5nf+zUFr8AEzE8J+xc6y+E6EtaiQAAAEYBn79qRH+KbhTdyySQ5ZqKHVStxrf8zUErUDlRLkpwXJGWM5A57aENWOIeQhlA6EA+NMSo4Nf2axnu6UPFdOUMdBtyzDxoAAAAaEGbpEmoQWiZTAjf++s+OCFJBaqu5RyJqCI5ciFy+QfdViKCVpnQDNKhtjLsbBDTAxQkGtiA7qEUco1DRu3mLPro5hMkwOFSJ8jWf1cptWx5ZqYdGH7hd58lsJBKjBOOxqM2gWtMZwhsAAAAQkGfwkURLJ98BvQr/uX5kKHT5/lStdRB3EpNtMxwioGAW9Yj/IJ6BkSDVCrLUVYjVeHAMS+JXklTR7CAUKq0auNYPQAAADQBn+F0RH+DDJXX+1g7QyFsak3Ft1Oi7eT/WrVLikgNs5g3WmkArqxtoNirEsBZgIMNFCYgAAAALQGf42pEf3+cgskDvsTAnaZI7vQ1BYaJISL78CwSAW7O8MKN4TJcSzKr6QpQQAAAAI9Bm+hJqEFsmUwI3/vrPjghR6aRbs42PAWhfKG9hWLBQBYO7VhzdwkmyTb4OhpVXnuVoSIE4n58p2LyiEkgihsMcSjpZIoAbvLWcxcbbUxw9yvpC4qcgM9TIHk9bnA/m1AEfoKukgxmshP3lvq4sLFHjIb1Bb+gDAKawHICIFXSjngCSxCAgl6WAKcq3XBRUQAAAENBngZFFSyfeJlWxHr/xodW1dt25DUmu/wD5rm8TWMxQj7U8N4VJthmBHDdc8Y8kUy0uyF/QVRg6T1Nw3dOMMsOpcspAAAANgGeJXREf4MOYBqx+AWriBGIRbd89ihAU0EQj7cqlqBeUwwo+Hl7YF8BhJvIxqmsu0cqLbvOrQAAACABnidqRH93OX3cYVktXGX/s1aluoN3CrEK5rAEtDv1gAAAAG9BmilJqEFsmUwI38oIQZY3wS9iq8Ru1ZuwHge5YzmrwHbkgMduq9QhK+TbcU4LXoh7vtGatuH+orZwfaI43jayClVBfzKYUUCnemRjycGZN57cgRK9ea1OwkugTmsWQEppkBNh/JdFjjQyZ4JpTcAAAABSQZpKSeEKUmUwI3/77lcOCE/ywLc6XOt25/biPs8yLMMrhhhyESShoB+zNzfnIF0ukeMr6Ju6bgeroNShV6SlG50RR9jhmJI53KndzULEpskTgQAAAGNBmmtJ4Q6JlMCN//vgRKhGBrYv0qobhWAfXvyR/FlCNuqZwJVrpTIPYrZL+bMPL0iiMNcUQ2zWjrP0PVXRdSCVrfqzazsxwiJ/Zm9QLkBsqmq6bw0xk+Bar3U6aKO2MHQ7w54AAABjQZqMSeEPJlMCN//77lcOCEpHhvk51D3BWHT5CL5wZKLK8ZPkfOIV1GODU/wlBTykgfVfM7+jBF4bQaduzQ/+3A+AvwaK0CE5cHbjVc1pJ53H+StJxG9Ua8/DrUCHl8SGEn2wAAAASkGarUnhDyZTAjf/++tBKhInjrA/ex2nOLzwJ3o+r2vWQ3+t3tlTqjSvEVovxTIc8Jr3NY0ry7qzN4IljCb1qcgZ7YvzWMn0aYd5AAAAREGazknhDyZTAjf/++AcgmTgdQc6/hgQ6BptGAKfJDuajnh3DwMwZaWGpVqb7Mh/5bvAT6CAEuDJtJ56z+N9ZmTYlY3PAAAAZ0Ga70nhDyZTAjf/+8HNT5G8Rzz6vkY3Rq88twFJVOYwb1Xg5/k6UGi002nITY9rzS8PPaCmc5nqS0Vh2hJ5DZv/OV29OQcMFyct9AKS8kL3wc8YwFvLGykWfhhLsilALVhO0EfbQJkAAABbQZsQSeEPJlMCN//74DjhLnMdk0oGHhcQFzV4h856+SZ06+ZzX949gBdXZGd9E6QP+F7soJy+XXtQh6tUmGDVgsotechBGEMrKDbwe5/+7BdVED/4EqLPYycb4QAAAD9BmzFJ4Q8mUwI3//vgHIT56s8EtfCC7Nx7BLXzRep/pBVMb6qXu2U65llvCElOC5o5lAlBAWmxsuQ6hG/bP4AAAAA5QZtSSeEPJlMCN//74ByFT5oRDGgg8nQkCHzopVZlIYNDj50PxOrh5uMMbamN5s9P/sz2hGglU1JfAAAAWkGbc0nhDyZTAjf/++A44S5za/LdFJ3sHpSo66b88l6UQOHPKyMuIn8r+SwswHCeYG8PNQSxT+5oV0IxUUsMcQ2ooBv+/Y/FAlBB46wXopImuyCaT/Y0EtV5MAAAAGJBm5RJ4Q8mUwI3//vuVw4S4OjJsHVijyDrIVPe0J5xdGasNWKBtM9uPh9nF4cTcB0XMQNeXo2jAvyPsnd8B9C6JaPfT2A3eKrYalWHgH9tUOCy9Knavp6PwTVwZLwjY3e7jQAAAFtBm7hJ4Q8mUwI3//vuVw4S52Of/J35LxMtdpwPUYqk+ah1t5ft73SHY5tDI+UGs5XnciWB4Bckb81EGVyf92Io8RDVoOwgtO7jiYmvZj6oahKvO0TUvdWw62HlAAAAQ0Gf1kURPJ+A5qFYNR5FvuejVz6I0sUwXTmatChezYmLm4QESF4hGnzo+2N789GMiY7Hg7gSQ7H4n+uPL1I0UjWopTkAAAAiAZ/1dER/gwyV1/mv5SPC2NY+LjsomX/n03iAQv0ER9NxIAAAAB8Bn/dqRH9/nILJAn8Gzsc+imyjJea4qC7Q+WOHNZVfAAAAkkGb+0moQWiZTAjf+/2woAW0kUklgimlgFCIb///hyD//bwn+Oqnuh62tNC5+WgKA+8xn/nply3uOc/ur2f0h1fx1vLWFTcNCmIsUc8q93KPKUmAUmJhrNfrZikefSi8Pp68VDmF+//T/fZ94uc7Um4CesacTL0HWpbC4pwVw6HALwAKhczR9X+xqYErJKlS4y+AAAAAKkGeGUURLN99VLWI5lTeYG9BPyle/82GFMyCKf3wI756PFoOBiQwNLrHDQAAADMBnjpqRH9/6UAHHyCU4A1cOHRliECctKkLiOmAbCVU75F6EeKMbMG1VNe7P3F+zi6z2GAAAACKQZo/SahBbJlMCN/76z44RMcXI48OwBOWvEVAsQP9NTZluNbhD6AJVYnC7zW7VxWM/HOBQiFHH+gn6l7YBSiWoIeoGHmNKgXbtRVNg6EmvoJbA9qtyZntMZgCmaCc/vz/8jRR2HQyWOQpuYlCjrJKp0G1cCrOLzK/bZToKz16r79FSV7J1Wz3m2J9AAAAWkGeXUUVLJ+A5qFJ68onA6dsJzs09N4LvIj3cDBjRfEdWApojACyNY0OnrR86nreca/gKP/udzozcD47DfiybsBvJNDTMeXU0IZfQ5QUfUQrr8Uho5P99QOyDQAAADgBnnx0RH+Eo2cLvj4vNhzCIRfjlupd5qSQ43HazEjhH1jmlnzSvnBcw28QnLtZNA1sBrhY1L+GTQAAAFIBnn5qRH+KbhTVzL5Iwd9W8s9aJtz6PoSUC0gyKP/IlosfSF1jR7mDYgMoJzHwJmTvCykmr5d9a55uBDNKtoDcBCNVk/X7WrPcvXMd+GCzMVi/AAAAr0GaY0moQWyZTAjf+8HPtQiKzcGjVgu3kVtj50fuo8YxUJ07ui4ClGPzR4sW5Oy4COuivineLE9MqxaZ+Om25zN2GqsYwHwOjbCmp/yOZ9thUdCfxibVxPf23ClpXIFCx9pfQ3L3gMhk33abt/kuTykbkgMonPh/peH+AUaF5KyXel7QmQs8VNnRi/oq6rs/xkPZZApzDDMdJbqCvanFzHalDHtkp+lq2f+uQsCGxuQAAAA7QZ6BRRUsn3sdSq3L8Z6sV6pC5e2JBxZNJinGlzs6ATElOuuq0YT7sb+JIlwFgYRvazPMwERrZYQMplwAAABGAZ6gdER/inEMd83hLolJGCEU2A5LWgNtuQHR3ohRJdjldGiSpKiSaCj6sqpbpZtjxpR8tMmPHBZmo35RU1qsmIxgL7vhMQAAADMBnqJqRH+EaK8pIFP2GKqMXBdJoLodSSH+w/mReILD9PWSmxDPtHlmi0aTpQN6QzAxa4AAAACAQZqmSahBbJlMCN/74DjiuZY2QoaXm5FeZph47u1XtpsL5Az4WSZShugtci/xJ3x+JchRTFcSCzajdyIBFD0EmzyuRUcNgfqEDW7dKQFFCS3LIIA75c3I5bEeWMNffND2q8RNvHgFxIiQTv80jmOiJhMyx+fiURB9txK7Jr/00sAAAABRQZ7ERRUs331Nr6PZHsxt0kUckkfcgh7z3n9GmoeN10UT9FBRzEdYc5EHASutw1r3EgM7+7iDD7kqUo/s5W4KMEx5YygXQS8Tag3KI7Y8iAvxAAAALQGe5WpEf3c5fYTCslq8zvxPrc1+6JFYuIG9MO2VN6us9stILo+Cd0DlQvGUYwAAAKVBmupJqEFsmUwI3/vgPUSWjRxdIcdgBijf15+5OcRgW/d/tWcjbO3MZMV83KAk4iqqV860nMKCr1y37e9s5DWqpUgfY/zSEY7/mejd1nXV4y6ASyy7Tidqy+itoBo2ai8/KleueofaP2gYMEwzc8CsZ3c5zZFOxH/DE1t8caigPteCEIVRz7SOG5MhaZRao3+/RoS1pQJfJMw4indglMRRujCISvEAAABkQZ8IRRUsn4BNWWWPzM46t4L3m0bngUEjMFtvaLh11pCL74yxwXywy7hSWyjJhKVUryYLLI7LEmEa38rtO2jQrJ7qp+4gHE4WAIkvSUMYkTurlCHxAi5I4XvUWaZ6Iv16g8K3fAAAAE8Bnyd0RH+DDmAYYFgtadtWEF20O27XaZ+whNwYc3jC6EmY1ypgpN2m0kur0wxTdHYayalaRxghl83RSGlQP7QpaJhHZnAs8TP3vDYer5uAAAAAOwGfKWpEf4TUKkzaVNQT7JA6ND2nVupGemBJq8ChfhNtpxCX4JwNHEZ9s7UeKcUe4X4Tr5S+yFh9nZYhAAAAj0GbLEmoQWyZTBRMb/wDwzHFcyxseeHtK+jVLY3A8ftTIVqtjziN4DBisXVomoyMfd1Sdb6tJXsseTvlPzSSgUV/wZCU3d/3F5VNMl9zrx0xBHmsDhwytPRmDnxl0gvYO++K7bpjqvYn5Af5Uv9cJNeJ0/tM0VGvOcRQAHEE7GPbnUfyWt3/xg/JTgWMC3oYAAAAIwGfS2pEf3+cgskCrRYsbkBnYXyLyFu/tRS+DCNY2/rA1wIwAAAApUGbUEnhClJlMCN/++A44nVBa6M6WcdL98HeebjAAer9Qpd/rGaSXMqDHAum3dQpiQlLhO66/sRs+qjg25SFCsZ1Vo4nYX2Mb50Cn+tgIDPoTLtW+6oNkjuj9xBK8l4mOVWx/rY+B5HpiKwKANP3WwumLK5PSCAmldTN+JXvE8vS5GfjO8R4YQe0Bloy7B12ige4QGTZi7fBYPlKWdxV35hH4+gi5QAAAE1Bn25FNEyfeJmV5SQClUo6EKDR24A91N62GkhAu8Y44c0RYjbPE4+ZbvjaBpON+8Xjc6J1wglUayytVTVHCY+766iHhKM20y4SXfgwUAAAADgBn410RH+DDmAYl2uPzEbnMKed3nsUICyCyeDbyNodbOCziTs7weQQn1LnK0BhMLkLDxK2gfVQXQAAADMBn49qRH93PjJljXXeOtUIt+VecJLVYqr6QwrBQGtz901C+SNkcNsuewdjFgpR0rkVyO8AAAEIQZuUSahBaJlMCN/HqbLy8/wEVP9rejuhgZ1/hTiRpJ6ACUBjzu0rZf+rwUPqu6H7kTYGQFdQa/1tJhAwmL82YZXnHTmWT+k7HnR3DjYHPBWnvdO68MV/h6Jcec8o8RZA/ryFL10MTl1FFKOnx9c5nTtu7WVeqQ82jjFbnYiUJNgKb+Gf4/paAXxjzK3x1VOKsji7JRVvVt7IT3RxUS3bPHqqBqcQlYpSxPlruhYyaW1ZA6jay+0HRBh3+Hyob185kbI6d9pHKI+sD3Tq0UtpSvdz+/HXXIXTeB1Rm/g1ldwTI9+9Xz9EHGdGIIViutN6mdwpp2q3QburL4OaDa4rbM66+0tY5DFAAAAASUGfskURLJ/NY7+LO8GHdi60riHOu6JlBVqwNg1VZMFN4ejipDuWEsd0/ErZ0th6Pxc9Cw1Kjmuq/zyEfYtQmhQsswGsgZ1YMKkAAABrAZ/RdER/zdG63bG7KwCn5wCGa3osZLctRFh023EneSsmn3MzWC9KqCBDfOcj+nGJ9RozF7PJP3C9bOztIPLCjeboJrWZS3xCKtt7JR8Uda3/E7pTMmmqtRMFjpXMNYSOVV2dKi7GuhQUoYAAAAAjAZ/TakR/hNQq3WE1NT/3tCaEtiN3t4YxY880nJUIH9NlCBEAAABDQZvWSahBbJlMFExv++5XDhLnY7HnguBrf5Zur2mrLT+Vz/GnEwY9bAGfUJr3XeDoQFE1iAMpidl0+4X9hFsf/8NFvwAAACsBn/VqRH9/nOBJylVotztSxpovwWJ7DmTJloPNrARFySkRNw2TMBuSJKqAAAAAsEGb+knhClJlMCN/+8GajlL3k20Wujiv8rx1TQrXb7GSm0kYYGmTH9/rx3KO2kVqyvJPzuuFfEim0UAXCWVIQw97gFcZWGwEG++myxOImz7Meuw0MWDCCNZZRC2on3KRz62Ye9OQv/dYQ8UMk1DJlNJTLsnSiKauhVJx4gsD9dWa1Kjljwa7NOoMT0lw+1ORCbOeHGoMPOXwCTsZiawz1MU0E5ekYhfXVHCB3xVjdCt1AAAAQ0GeGEU0TJ94mVbEet+VNuwEtG19Vp9qEKqXEQQaYjS6ZG6EkvByqXm1cZlLTLEjoe/LYqJRdE/SF+jsKjMJt27U1bEAAAApAZ43dER/gw5gF/dlobTz7RdVKB+qPpkr6GiWwyvC7wIA+3QCH4y5k5AAAAAgAZ45akR/dzl93GFZLVxl/7NQUtFWdUUvh5Beyz+kFAkAAACkQZo+SahBaJlMCN/Hq8yjL5u9j+dOIruITmk6F0314KXaka87ZCfL7tAHNxO+iju7Wtc+zvhxob4L7mjhRqS6xh5aFdu4srY3Z9kJMkMwAoPlbslt0ED4rQc9cCKEvzZe8vP5SuEFkEJ8M+mCHUY0aVpfO6zWHHxMB8o43nhSklJXgCtKGumFXlujSGPlk4FS34jNmpqXuBhr+dGn1yKAZ3sKSYgAAABVQZ5cRREsn81VP5vFlB4brUQOGtdboyvxjRTY+YUybkkkNYTH7v9LDS9i+dPj1ZQKYWWFo2eX4uZLvCRgIBXKzUoZJdkcWfXIyanYodRMHyDg+OeXWQAAADcBnnt0RH+e0OjClgq3ZadlqC4PYwt5pyZHpbAtX2CDDw6VQUFJdMB1qD9hMsCCcnWNss5UPmCBAAAALwGefWpEf4TUKm2xG47R3Pr01Jw7Qj0y+sYt+R7i5ZfZWQaodl6QPXmStL4WVN2BAAAAckGaYEmoQWyZTBRMb/vuVw4S52Op1bIfPeM2q9fIEDP8jiZLCKuzV8ftycOzLks1iKOF1EIt3ZtTQ9F5OaVLBfL+oXlHR0aM8Z0FAD3WY96fF4f7fCu71ZyQlXJ+Xqv31x+daLq88norMrMWmEMlQ6N1gQAAABsBnp9qRH9/nILJGk9FNJxMcY1SBNTaLEKBkEAAAABtQZqBSeEKUmUwI3/74DjhOA4HWckmvMs3ndZqaLe9W6S7iB4Q7A88qrDfeboB+rjGK0bIQqpCu+nuXe0EwJabFzLCZi73OsBrQNSTgkPJi6/soSPtur9skY2PhGs7djxF0DyGMzRzfwPqmGKw3AAAAHBBmqJJ4Q6JlMCN//vgOOEuc2omt8p6qS3Ldw5989LDKzuDHw6H9tFiHKlV1hKTORwp2nJBpiWf7T+A6Gy7CdBbhR1LLJqSDyzQfYbByt6HTZQ3hiVDL/V01gSXBPgX36fCFFEJxNZwInp0CmM1ZnvBAAAAREGaxEnhDyZTBRE8b/vgHIVKDNGLdlSHwPwIRb+tFWCuSrukzMABQCTH80n679WSke43t/ig4V9E5EyVF2d/zscVwbygAAAAJwGe42pEf3c5fdxhWS1eZ3/s03ZhFuMfWe2wKdUy6MHZBxrJBq965AAAAHVBmuhJ4Q8mUwI3//vrPjghSM+5BWn2Ge5VCc0ecBc7EOmKermZ8iWDJ4BRe+G/1S2eQMXMqXpPhZIa3jeKjGvYeHBEFF/sMWQpQ0qrYWCvPbD0bfqtAROiYI0svXAdofjDVcDRuxRo9ypWwrQsM3G3a4CbkmEAAABZQZ8GRRE8n3igS8fsin25OPmCjaoPOqdQy40hqX2JsnRU3RfPgOmKeVL8IQC9iFPZaNhAD9T5boSTEP+kwtDeCH78uwP5hUVwfti/rNaWniv7PSSGTiBYuyEAAABLAZ8ldER/im4U4K6eWIXO4QmaqlkN+TQlwR1AnPBlX1lEm5Uviyk02+lXK903CJg7gqnqUC3PCmW786/bPWqAfGLvDpW5m9IU54/ZAAAAKAGfJ2pEf4TUKkzbUQAa919yFiCEgIFdd+EDagYoE1SAUggpxbI2QuQAAABSQZspSahBaJlMCN/77lcOCFIz4vqiFuZcNnUTVzCcRWfNdM98rzWozz33Q7tLgb6TakmSicu29hGCOBKuAKHMQ1NBeTV2csdtYUMxTioTCcrT4AAAADpBm0pJ4QpSZTAjf/vgHIJkCtvXBZfB+14jkra1juFx0W5C2dCXvZjLy9zVPAyjT9j980ElW/V1rmRjAAAAgEGba0nhDomUwI3/++FR2Jc7HZDkW4uJMmtI7OMrf1NHhL48zEpmHBdgRL4PFlMYTP+bOV+4pFezhcm6vv7Z7df8Mo/izx4ojiI42oNc5zuIVj9opEKt8krPnE/wfreAmiZIhYwe8cCEwpfjXXA6bJNO4q/P6CuHDPJaoZseFBbJAAAAcEGbjknhDyZTAjf/++A44S5zaiKQuEhJ4jmNCR534wzFsjpJVuFfEuiUFOAHPYvijJXldwbC09W+F5ucQhVaOaL0yaM7JqDvSBd52cPA47AOCCjU7wu4stqnihjO5ec0ETnzf0diadQ9rPIVw+r7NiQAAAA2QZ+sRRE8331Nr6PZBVkLGHLziB4VNczKxSY25csybjCivPqOIy1NJzulsNvAJVCJH0qinwjhAAAANwGfzWpEf3c5fYTCslq8zv/aYBHXhN31AQe00BnlGw3gDiP8P3ehzc2BPCbVfqZx7BaQyOGwMuAAAABwQZvSSahBaJlMCN/76z44IGSiLnArHSWpo/s1uhN/5B5ilnLO9crpBdkUCLo55nXPs/bbw0m1ni0JxFnWRtKVr2HRpNb2z2Bzw00yKTCsaNunJLxfhVSpb57Qap9q81paINA/so2j+STu6esbY+YyZwAAAFVBn/BFESyfgOol4/M1FWMKl1NWmFHIrOMqFKzWFvxvaf+3gDItP7fuN53V0wplQxqROuxwNisKkYf1A0L9vOJVudeMNQFQxSkGQDs7GIlrSq3kzQx3AAAARQGeD3REf4puFNcwyKvZaiX5qqbL35NCW5PTI+EGVfTC7EKS8nUfK0RPsMBuJRJi+OvTAxVsTOvunx8J/nO06kwqHnuQ4AAAAEIBnhFqRH+E1vUc8XsyRGeV8iG0MB/zFxYS/vkcNLBnuFwB+E7v3VvwY+kBV6T1vy286lrWGXJod8XkgjgvJuvRIJEAAACVQZoVSahBbJlMCN/7xIOhIUAB4s/hp/HFjXsfBSfVzZJCJlAqX8IqjA3hrTr26f6UjDw5+ICxhwJWzzIjlK/pZ76OT3X6i+JXOWGDxxI7H8VsRNRLv8lmJsG43V4JiuFXvawRR78NKMZC3LHXIjfHxkbHB/JcwR+f69KV2/QtRyZ2dRq7sBgyFC/D6+JgWb/L3sBmXeAAAAA0QZ4zRRUs34Ee3aUVIPxx7mDUX17S6FVLhfjTjm/q/myIrqTAit/zI5wjoSwqYdmF2r6hcQAAACQBnlRqRH+EUzCwuV86LLP/gj3H7Nt8I619QX+jP+BvQi+oayEAAABbQZpWSahBbJlMCN/74Djg6QUKvd9e947e3J2lOAkYhiVNr9cDr3ELDEevtAnqI6c+JX9/Hl6IF6osErWJfEaZznOiFMoX0003WslCugmdnvq+vcZwkx0sZcaeLAAAAE9BmndJ4QpSZTAjf8oIQvW1xSZ7x7rM2+/pGwLqLOU5uTIQepdfmTQp8FGjoLDBWrpb5S15nV/7T1MmUpm8t0TDIMYbRQkDTpqwSDDRE/k4AAAAVUGamEnhDomUwI3/++AchU+lEC2Ii5zQFVFGdU4MJVerkDyzoY2q5+KZYNo7qgNGISSgsF+cGU0UYWSZEeb+O/N9AOGI87QduAAXd1wUH2zvdsrzG8MAAAB6QZq5SeEPJlMCN//74DjhLnY6/LPx82l6Bk5pK9HvlvJS5+uXxVzkYeM+SI2nn5l8845JM2j8M31GNo9VnTQBFtuxtZgLhZRvBX9iGRYWh0U9En0SXymWgiKz8LWajDidYO5kVrzWsaI8uEAxpZex2ijK5k/IhCK3vqYAAABXQZraSeEPJlMCN//8A8MxxXEsNPU5MyGp9izFtgK5B1fgvNcD+Z43iIuSvepRhmJjGw//jbq3VYyQGndV5yIhGAhdpfOXdaDwSiD7VkSVBa0B7buOG6IhAAAAUUGa+0nhDyZTAjf/++AcjQRnx3awW7tncZ5B5Juzr9se8S1VRa6KH4UBW6LI1HAegAOeGogjZT0PeHL9xMUqvsOwVSI7eLqasA4yAdPUEaLbgAAAAGVBmx5J4Q8mUwI3//vgOOKGHX9dX0WCIwBpYFz0D6F8pUyB6eshOk+dTkDqLXSsUh+Gr9DuLFPZf17nuf/ueCIOjxOgcQScSRdP0xc1ZVu5Jqv3yslnizRuOXMfXxkx28Z7fAXjgQAAAD5BnzxFETzfhHWkJuNUHXHYaxywUjiGRpqXXHfO/nq8EmC1znAhOoVyi3LAXh5udzgkSjzv3BiqOhlYw5ff4QAAAB0Bn11qRH9/nILJAn23SE5EKnm0eUobJ6rVo9aydwAAAFpBm19JqEFomUwI3/vgOOJ1QWlW2Hq57X357pDMQnGy6t/q1J3KOnfwatDKcCKk0FUX7afaojFcA44+ouv/VS5zgOikY2Jo3gD+Kea4bPlYDa1BAowxsskYioEAAABwQZtgSeEKUmUwI3/74DjiuZY2PamuR6r++1qE53Iv1qMB7HQ4ViqPGfOl2Zby/x/4pN8jh89MtO/osbGCpJthzQcEKpTHULvYU+KEoRclqH1cLUBG25PaeWeD8zxyPBhkjHoeK7nvqLLbLzvLi4uamgAAAE5Bm4FJ4Q6JlMCN//vgHIiQ3BG1Ae3g4noIB59LbXTzyxAiwWYAA78rBEuHCDQZgSoMfigSgb1vVN+xtQOAzXzgo3pxNP2vKC9JlDb7AWAAAAAwQZuiSeEPJlMCN//74ByNynH/TQlKMDZynXLHrMF6i39xuh+J1cO7vKUsDG7+KRiZAAAAaEGbxknhDyZTAjf/++A44rmWNU+hTQglopLhhmAYLK/ZpHdfpOlqk91IgMNgFmHV5vRn8W/Xr8cuOcpnmkUCutdRPjNbPLQZXJ/3ZczxEXjRBuNotgM7N9t+vTv3LJoVu5Gxx4RvBJ+4AAAAR0Gf5EURPJ94oEvH7IgaFXdPjVnk2O6nLuXIwqxnLW/+qohFcx7b+h5swmqFMWhnXHDdgVqWqcEU8jPLgjtj7gTWitr7+s+gAAAANwGeA3REf4puFOCtReQubuf3nGJsBw2LCxyegE3x6o+mF9S3iTu6rNVSi3sIbiyLwJ2lalMxkYEAAAAwAZ4FakR/hNQq4wqUtHk0Ek9zmyjVKtbfh+Ep3L+pF8DNTunjSXvZ3i3A3VLAaZfBAAAAeUGaCUmoQWiZTAjf++s+OK5kJjsYIqxbAN9qTX7RKmcdzUOR4WmyrfmGXlSsY4VW0N0z/pBEVT3CsU34BdrTRs/Mj365dhncAfPS39tR96BzM/Bfv/33oqAmHEP2cAcTsMGtEH+NPMAl88/Okh/VCCmQF4RGqYBUBV0AAAAtQZ4nRREs34TtSIINpafvDmraZs67s4zwOz7bhVRQUEsaKdSs9EDDJCr+0tXeAAAAJAGeSGpEf4RTMK9Ia6xDVybe6bVn8VMBUydoiydjxdfTkz37RgAAAHZBmkxJqEFsmUwI3/vgOOK5lg7JrSIIIA/sNpXxrR0mCihBXocOnszuWQ+6aCY3Qf/ooqPuZPkLUBxMfTn/KI2XKfgwuHg8Ndn/ej+Ix+JsJLtXyYV1Ee5PeWlz4wtnPHfYsD2opYlEN9EtBxxEIP+F6hy+hsFBAAAAWkGeakUVLN99U+XYnHXrZhpGAxLbiGJkujNzoot//snhNxSyWew8NVYOzMd4cv+InhXj9mVM8q9zNsNntf2Lf8crPpKa3BPYJbhAof+Pw38XwU54f+mhs8MZwAAAAC8BnotqRH93PjJguRMlq8zv/ZvacNf3ZgDNS6n2w8JdDTja2l1kfVwpI5RyQrCnRgAAAJpBmpBJqEFsmUwI3/vrPjitsEkocBG1hRaRC7O1B+CbrXs8hHZmhG2DNlmpBvFXAQ8ZzaWTctxOTdptC0ofjr2GSDqDoPTD7fLpKEEg49jN3WOp2j5Odu8aGp4KdhyWrxYcLD12s3Q1/uv7HQxh4I4Z3zALxj7IUwYKgxbBRg9IlcbVNQWTb+a3gaDRS4UaCkAw/1qd0064DhTpAAAAYkGerkUVLJ+A6aad3hR4/uXSqV1B0dTjjrlZ57rkwkz5qH8A5BiiEyjQz9F4C9kFj0976Z8PPdfw/bqaUlhpDnPTBYGWwhyC6yOlloZsDr8qrxeFAV10I20MbiNT/qY9/jTAAAAARQGezXREf4puFNc5iD93O4M+je+bLdricvgJeIzf4IF7YCFGXutVSWIsMd4abDjWFaDTXZx9lcHSO7Z6bVKQQZkvgLpJIQAAADYBns9qRH+E1vUcPAJx4GZtKUaFsCF4uxw8Y6tLrD5q3QdAfoxOrS+VUyv94KTkHGklV3NC7hkAAACdQZrUSahBbJlMCN/7wZu7TxZMM/e3lqS3Hf7gR0yQfSMdS0gp+EYh3ACNJeFZWrJYWy+Tuxmmg9yENzYhMzSp0MSd2WcLpHZqtLjZjQM+Ku+MEhhVWRbHi/Vo+GTEjar+dbIOrWe8xdbaEKLpaDLCjIApOoy+j4hP/sVb741ZZK/0nw+kO0/AXBEe8Neoan+CYCnv+Ow4wt32pXRHgAAAAEFBnvJFFSyffGVVSGsrz+3b5v1Zg5L2NFJN02uV8jUcyUVVzVImpmzOJPahozMXoce68AXxo2JwibbE1C2KfdoQcQAAACMBnxF0RH9/iXZSQkB/vO6m2O+aoEGAw2+uAjDlKDyeUraXTgAAAD0BnxNqRH+BEcHSwjYaDmtf7P/3EHd73hx4e4SGlU8teL9lc635JeBPPhznR0FXQYVX63aPmrj7NT/7R+2ZAAAATUGbFkmoQWyZTBRMb/vgHIjamvtGoxs1uSLcMbcKSXfpDcfNEklZ4VeRp0H+PR6Ez0IVFGz/p3jRWk/Pf7VVg1CUVglF6oKe6mOwVL1ZAAAAJwGfNWpEf3c+MmWAiUqATXXR1gfDZs1lrr5MWJQArYUxiqhfYhdIXwAAAE5BmzlJ4QpSZTAiPxVmeOlhGwz/Fs1c6V1b8P2RQeXHGQ7kXXgq8WzG32u2Iy6GVw/lQvrM24qRUp/8JG4USCPSdqAC5av0hupmdqUQHIAAAABLQZ9XRTRMR4puFN3GPnq2c/KJzU4nbiaIdSG7X0s9TSe4iwC+u9L3527jGFEAW/pgQ1w8VKNK8KcC/2VVebquBtHGTz76FnlStsFpAAAAKgGfeGpEf4n36qqADnsqQPIJnEcRRKczAkBlIY2jZnmdmdek7gJRW9Eg0AAABZhliIIACf++DLeYgN9XYj5zg5B5QJuOQipq5DurWEjnW3UorRPL+Gmp7W/eIApq4hA6GO7MXHIyDHsVucz/aMFYtKTDM8z25ZLQfiNZJqSiDN9723ESzhz79Cwt5M+8vOWuuGhG+Tvq9foXjrYoVXiVDwEsVfUmlk5TXbEqyYI+MZpo4mDJdCI6FifaFbOBypE2MZf5XWR/BCziCfa3ua8heIb9mYyPBNtGSBm/2+TYyBLX+ccaIy4TYZdApJmzt8tHhRNNuRbzsMmREmXlyFjU2KTmRU1Z3ZVzYLYM7mBumK0HRe3QE+vr6A4gpTOdn2j9qjAKj6NNg2MmzTCieFUwql9L6ptX7F/UuuzRuFAfTt932rxOmz2BNjBs10CPMOhJ6Pu9idR/eLvVQSuID26B7YxP1Kc9QZT+tp+YkaEPC2XBY3bdZ6qbfFKSpyAAvrJhEkF1eJhOtzQzGB9M9QVORu8jkG4zlLQbGF6ZqgPB9YhIYnHkO6pyRG4ZEwQQRv3pY66HZyqmP1pcw4M/JTWg0Dc5WzUWPGIAgkAFQc/1Ta7hNEPWBuZrLM3ySme5CdFuvklwbISWRsoUpauBnsVWWN5QeWEHF6TJILZUHeFXnSgp97RhPWIEEhIlKaUd4hpJb7cGhxHmSPkyScl5dA7+PIQv2HUjD/HkIC/Ucyw3+BoGsOX7lkWmxzmjnqmrW7wAKr3betEKeTC6MuBZC2Eie0lFVnNNRkT4VOQIj4zh9jxJYkGeDLwst9Jm6zQ9hKlPa63RCfp9NnaHWjPsnJjMX9F4gwR547dWgbLvYE7IuBo2AK2pqD7u9lEBFGBQEfEy79ZlxUwHxhYbaKRUP2ScWs/5ZBg+sELMAUKpCbKjPE+vDv9EAYHgEOzeiT71raTwKS0j8v4q+LNmKEzMRQT8YYKkcXDIj+hIcNxaNKpyvuGqlycqBZerfPCw7NyeCJ1RZg+wkCwjrzypk5ge+Jt7o/7sW9DR3cOslqBHx96pj5KJC/hClCIaZB6HFYXv5qM9NlNVSZUKsifIwMWVlTFH44dQrGr49VglslLicK98bhApMsqpHLl6mvI2nl6M00yk4zvnor0ybl+78LwbD6/u2JjQ2C5StNvCBcm5T9DSRQmWz+xgw3XM2cXGR9bRyAeishmO4rIwBRYnfjW42CduoUT5uCKdwarOs6+ZReAa6Qm032qE+yuGHc+0eaRHUxlZw4pvpia7X72bi/pxC4qWFZD2iC6up9fej8EinMBLHUXfZvheqdQetWdPndbIazJcmP6J8bggLDQ4mKE7dyUF19ugTbguFq279TrvB9ynED8sJa9KoY20QSV4oLpYcqOxvEiV7vzPO9OS9m3f72BkB1HiYIFjaJXwqYGkTXiuZJWrbol4+U0Y08C/7GdlrZPe6z/onp3lz5ODKTA0KXD5NMWK1y/vxPV51UDcOCw+Hv63HNJxyUPECSyqdmyGfW9a3IZeGpmcvlwizm/fg9CiY8lZSQs8QDGJpatu/aiCPaZPOU5MgVYO8MYY3Bg83oVFT8iz4wJf9K6OnKc1nq8GDGNbbnNkX1eqV7en2c97ZImpgcP2BFIf94YIjngEbP6/uGWVEfFyMJdl81YnorD2O2896phTQJZ011W2fThMp6sugWNSD4zAbCi+DLD8su8T5rb61RQ5zpFRjgGvMFoA+Swos5EfX7Q+pWIiU3BmZVbzG6lOYqF+5VQhycWb+BBOdBy+Gpnne+KvK78cSzPaQYJ91kufnulzaeUUO2hFOH59m9d49JV2fpLZStgYBSIJTM/9Ox/Ms+HFo9hdmgYr3o87P4pp1ZyLChdnLPD4OKjtr+1825bbLYc/XTrWBAc/9du/48l2gB3e8DXTZp0m+PbqLMrA+8QljaKoWb0phbJdacof3gIbAAAAhUGaImxG//vtcWI95nYiPte4ifOiv2MWrXa6oc5GwmVX4VGEPL3MSQvvQOWz4iK1KtWLJpeEgjzPL169ZcnBgeMLsKHdUM2wjMTrsHdOygwiz5mAKMj9iGuCQFkHUsm7iXWuKx/0L39JtWxH2DB5DrjSNacwTXnDHcq4Qo51kaWnzhKTZQYAAAArAZ5BeRH/f5yCyRJTRNLqZxFl4gyBe3yn49jBqp0OksTJCkbjJxyS0eiYhQAAAJZBmkY8IZMphG/7xIFmvHEwzcpixf6g85KZsRgAFuJ83fQKmW6HunVKuoozc5sQw4UsnqcarVRces3II69oVcoWZSdF1HPDrGm+RQUode0jm5khpoTxFEJWu1alofRiI9UOP3gOWkENJTpyKfX8mUK4iEwYLYiJBNu6J7fSvAmoGVC/qd36RNT1oid+703GbBe6n57kGF8AAABZQZ5kalPJ/3iZlhm9wMbG6wiDR23j91OZf89JoXeMWa6ZGKfp7xbU0IXx54eFVJ0ZoDzq3qCKdLExoFVdKuttpjLhgTdmW67ERQDYRx2CX2eTTzYWTJl97oEAAABGAZ6DdER/gw5gGJdrj8xG5yWNjcN5RzyeYQWTwbg/YJJ06mabrNeaaAIQL3fS05zTmWWz/Ddjyn7rd/fj4m4Wz+Fu7Q7lrQAAAC8BnoVqRH93OX2EwrJauMv/aYA5jWgK/npk++GXjyk7I/nxAuL2+vMCt3X112NqjAAAAF9BmopJqEFomUwI3/vgPUR3rMO/T2Y2VhLfhP72rEslxqZ/WAOcDYyisIRnWP8zgmcr3lUb56tCV7LscA3yQHB7Y/tXggxHENASIP9bZsUZSafMEOPrKIJo6LpThykrgAAAAE9BnqhFESyfgOol5DumVJ0S0cVIM7am+NZ19h6NTyxakBR/Ofn07JLTms/iq96suWA4PuOQi/6c2rG5IgiADAoCMybzxtewQXqSlO0vzLnhAAAASAGex3REf4puFOCunxZguX9wJlyQwsHnoh0gIq3mjZDFzCFf8XeZpx63GVk7z776u7t3kmCYNdaCOmNgmx+oj+n8mM6gYJcRfwAAAE4BnslqRH+E1CsKg5vffNV6KJEaafCYgryLkW8mp8xVM0IsV7+RHJ4vavIBIrAgydB5rvlAxBRgXHTpdpwIKS0filN50G/LRYjXoqAeXMAAAACfQZrNSahBbJlMCN/7w6grIVFJtmFWBNfhha03j96D82U7Vx+2nsoSjttpRDsy/gZVh1/0E/eE4qYqKmuDqanK7zl/QXx6E6sICctx8MFYLq7fS7pEM+ZN+K8Xv+uUh4vwC1duL2ZLxK0/Cen1zyK0H9sY5VGD83b0M1fhLQs9AaMHLOe55U5g3QoShXVmVZpN6cYhIcm9HRRzqlropIiBAAAAO0Ge60UVLN+BH0uHOo1OpQs1to+yKHFwaptQ5BFtgDlNMU+icULmLGxujc2EQSoFU5ygASOOA1Zvaf/FAAAAKAGfDGpEf4RTMK9Ia6yBEGEbBgFIX6ZVi4oZ9kz+DhKMG5HG7B0dtYEAAACsQZsRSahBbJlMCN/74DjiuZYyN7bnUg1vzKqdRlP5zSFw1h8FWJhA1x/pcA6l3LhJz1U+gBgFSepwTB3UvN0Ol6E6gnujgvTnK96bFAIXqhaB1HSHC989x9BPETy+h3mNxHpOKsWF1XH3KQjb8NMCBha0VBvTpEaYlRrVz0cWvC21LWVlsZ88JWjCsQwtjBrW5+pViDUFe0ils5WWRrnU/nnwnmdT65Qet+sgwAAAAFFBny9FFSyfeKBPNcaL7+V5X1IGhoxsOwcUUvCWuXvEiV7QwC42PiHrznZ2MxQBbEYjkTmgN7viHiaaidZgm4UkJPkk5c2BG05COu3x+CtcJ/AAAAA1AZ9OdER/f6W0XexwRhdsnetPageFCn6RV/+/fWfxA49EgEMtzDU27kTwCwBlXxqw37RvOyEAAAA1AZ9QakR/hHV8g/txlNnB3G4bTYjFW6A2jIYi1+5c/6wZQtVDrFEdc65a08gfqr425KneZpgAAABuQZtVSahBbJlMCN/77lcOLBfClrEvtK1bf3ouakjSedY0d3/QQpN2eRz/kCnq0HwQhsB+x/LqObYzJyoTEkOmGoN2mLlvk6EVJl5TXjbWz/ZmHsOadqHbIr3P2nKT//Dpvb/k1/evkG6EZ6FppcAAAABDQZ9zRRUsn3xW34B8XvH3qT97ssK9VB5rRJg15lNNNG2qjubIElb3HZhXKsOnG844Dve5W64/ukmghSwP0aO7uOMrwQAAADoBn5J0RH+AbCpxyDalGnNLTih1g6/0qOdJjORclUnHAhTKEO+MwTEmokHDhXkKpprte1wLoG6IbahgAAAANwGflGpEf4n36qr5vCXRELsL+sPgx40RoVK7tlQbDF2IVvmSd+z9ZY532dnlS8XL+vQIDhQ6PYEAAACpQZuZSahBbJlMCN/KCDlDfdqihMhNsI4qRGQFbeAgSReinIWm8Q4v1leH/5Oe95uqoaeIdhN+wGZD579BhEjNcllwMDo21tloRbtSe1uLHorlQEF2gL1+3QvtZ6Ww+YbiKc0Q3snOvvYLlV6VBkpFTEa/TljT/92LFpy+EncdIYAHscQDkrRLnVMWpdwe1oCFjeZaru/rcF7oZmGko8JpJoEP9QBsOq97wQAAAE9Bn7dFFSyfeKCcVlysnd4zeF+E4wNIhpJk3CJk50y4UPNHY2/KywiCicYiTy5aF8x2ZKGFuGUKWkFqvmv0KbGkq4auUlZxFpHVEF/DCBLAAAAALwGf1nREf4USy2J4yUop700PwUbNMjzSdo+LuZe1m+BDNgF5z/+hAiVM4yhjVIaFAAAAMgGf2GpEf4As1+KMQ0SOqKb3+q2lDv62ayTIdee1EruQL1OwR7PM6bm0uHHhVSam9bCpAAABBUGb3UmoQWyZTAjfbiYOYl0PwIa4cphLVFbEikyeFvuGUtjj9rexKdFRrEbKH+W5YBn8+aNOedn0EvYpxVizQmiRCHVywZ3ukr7oPH9177x5o1M447NLUgHH2HTsuf+vyFIgHZ9FLK0jO0NIH3ib3olKIR6JsVfteJT2lw+suG01vzwcItvz385NPA29gXXf9D6eaOQ73pXz6Bx2gc4oTPpo+WdUakNFj/2HXkeAtBbWl4gLb1YcKDA4y8G1BFnlxHv88wkrXKwI6Yol6fmR3bJI1y5dVNJQ9vudagKOV49JX0x1oxZdx/+gPH+URlfuXDOCUGqDljPZF5g8ZqvpMQOAuyLrrQAAAE5Bn/tFFSyf08HdPKf1cgoaUwH1/y3y+A+ulifGqb8DB2/uEFk2vpxeiHhlya1/JJnifRcTg6+KPhDWKH2iLwmqbqc3CDxYX3kZrHMTaHkAAAAlAZ4adER/hHi2UFyJkuq62ZFmJbEncTEmcCmczdiVr+IJ1OPJgAAAADcBnhxqRH/cLDyOioAOezAE4F60h/6HDW+EY9SSjz2LoTUrQn+4MyfhBEJdSMVRLllcvr+A6MoQAAAAz0GaH0moQWyZTBRMb/QSJIKCLArV/OlT6lqtITQzgbnr8f5Eutwk9fMwp3wElU6h8UpJ4dbhFq7lbGPtQxwukMLM6TueYoDKwOsy3Kfevszyh9gQA/RGL/2Z3AABMi/YwMbw/lzJ9i4vtQYdX3S6A0UZk8DF1yH/sWuw9zQKTzfzoQmtJ+btMgMswQyVUHemhiShtbMmfnX0YTBCZIu9ZV0hmGWGErVuj/rW2+UZU9nZmMEUv1AL47Rdpv3vU+5WsKtGVwtlbafwVaVPp7mHgQAAADcBnj5qRH+Yl5YFRBc8xrboQonVRcNgLm4flPt+8unUaOv5vuV8sqY8Po1ab+Ohyjt/kuvOUGmAAAAA2kGaI0nhClJlMCN/x6qvIqNCZDKPAcy41m6RBftLAOTU/JG9a/HDlAjWHlgLUf7lw2zHlW1dN6NmvoUeaQm4Rf3oT8dK/A6+EPtUha31C9R8AGA9Pc1PHLm78eBmqhdW0/ws5Wu1XWXR1uYbEmmCJLKmCQsDRBZlUnN5H1wYtmh+CZAiqdf6mqvSNUeSmYEs+cI/4mIy1lPLWsph1V0bChaKID0e/3kgcXyOXc2Aqc80V+XWH2ucNKEcuCZxIb+YLN8Yphusto9627pXQnMmE21y2Sy3yM8koWxxAAAANEGeQUU0TJ/TtTGaqRkjlJwPbvsmSI1lkdSxXl87ZF6YN95PfBag7xnAxXc3Gmj48UCzJYAAAAAaAZ5gdER/hHi2UvUrOZ/a5c0ermQddoJtsMEAAAApAZ5iakR/zdGPifqOf3ginCBX+vXCxrjKKxMmmQAUZhX0/pxpxbOnD0EAAAB5QZpnSahBaJlMCN/9dub8vkKP19B+pw6kf9CAjfIRxnEwttHyxpVXRku3ltHlZ/lkYpN2KXe14XnfKsJvlIt7C0jFqLMM1jsU7x0dnYDgnxa3KRXsBPz9MvvWPMhxArAEWOQPLOVcPX2KE3Zj8lL7CigZAVxYZRZxRwAAAFhBnoVFESyfgOahSeu9LaQlIdYgvy8B8MWOBXFmIgdWQSr2vo4u4eMzya1/JqL8142I8jEYL77i0hisMK1dxUEyVjcD9xg19ovelUxl0c1xZjCxVSzJLLZBAAAANwGepHREf4R4tlLARMdvFFuxCgK+j1kYjkNsSNdqSwreLdD/OSM1WLdrguFezNg07LOAkL0dD/AAAAAqAZ6makR/gQeUVGdf4170agBA1IK1mKqcCUzcAxkJipPFkjr8XVNqedToAAAAb0Gaq0moQWyZTAjf++A44SZIX6myKoZpfSwx3KsEKMJToKDMqUqApD1eL5COD6ggxaIYo2oBHwnAsVolI/njGgWAuS6GnAALjGTHXTap6KqRt2wQfFm8fHkqvFwMRzV8ZoAfwTNrZTwwkwSZyDrc4QAAADFBnslFFSyfex1KrddINV5fGhaSsbM71utgNqArEne5/psmiII46ZfwBD8NIJ75jHsYAAAAMAGe6HREf4pxDHfN4S6IhdmW0Xai/zbnkB0d8faNtHQnA0/PXBa9GH5DiHe/KxjzFAAAAB0BnupqRH+EUzCwtvKZsEyQVsXmvMVAOIOnZjpXqQAAAIJBmu1JqEFsmUwUTG/KCEL1X4ozsdj2k6Ce9zYKOCXqaE6ZjAsmKXfbBn6uZEpzN2/8UmUwlCJhbW3/dILPwAkQYsQakIryqg3c7gMWCnsjcddWB3LQzuP1+p9jOVcrD3Xr4c9IecX1eIa4w8BDPtg42swGnJPlLgbQhIpfm0zFW/KBAAAALgGfDGpEf4As1+KMO/5+YX8sjGGhQz1ZkZQqK5ID4Enmz30wLs4t+OPMVJ4+oDEAAABTQZsOSeEKUmUwI3/74DjhPbILSGmfMEGXTtKl+F4XYSvJwzWWb1ZwbfPPDCkNhDyw6pPEXZYk986ThQwSCTwvFR4sqnGv09HOAFhHcX57o94v+6AAAABrQZswSeEOiZTBTRMb//vuVw4S5zapKxN+T/89Irq/c8BeyQjHqSyDZmxoDKD5I2ZER30lxBIf1WoAoNUW98QLYvYdm6zyybpZzG4HMZW8BLCyOn7D2+fhOzraEYLfhZPKgWhn6VDfppf0am8AAABBAZ9PakR/im4U4K6fFmC52/zhJFXvzNhyan1A8C0itT4SM2A+aqKJpuDqCdva6Fc3DzqWoU+0qgoASrLy+6MMhd0AAABaQZtUSeEPJlMCN//76z44IUkFqq7aTfAxQGo0WrVRPKTltKcUg/xnRhZwIe+DxhONjdOh8IzzK35Wf3PTjClIJPbgHcpwpxkfKcBfS/pqFbk+KY4ESO7kMpEgAAAAPEGfckURPJ+A5qFYNVyzptbWOl/Lio4wgDNAqOZS327gG2S1pd88/6C26jWrIjYrV1arRkukKguv1vHOvwAAAC8Bn5F0RH+DDJXX+1g1VdtIjdzQ0EXREbOPgvpkgaGuH7Q/rjrGKps7wnLqUFIdCQAAACgBn5NqRH9/nILJAn23SZGImicBQ3DUHDjZ5s7/AV1JPRKG7gT6le4kAAAAYEGblUmoQWiZTAjf++FR2JyyC1iitS+dYx0EO9MkwWF5beGuOG1ZTyHyr3lFsRHyiSoIXq4+8kJeLsXYAYR46+4k/GgvTAKNRfO38FkHkrRwtEd/5K6QMpzqI//YNmn7wQAAAH5Bm7ZJ4QpSZTAjf/vgOOEudjsWuSxRrOyCZjbojYuMdD82NYeOPzHOOkboB3OZcVutWe2p4e2RjezDJTnX7ykx++hEsasudgkCPoByJ77cStz8pSx2Mc8cb6wHv4P1NBXSxRODEmyzR6vAkTP/h4qQ+1JpzMmK543nMxeRvcEAAAB+QZvaSeEOiZTAjf/74DjhLHUdnTIr/kiBro1rbcLYxqvkewiG/h/qLASFUG7vjx8I/NQYp5JX260H1dnWv/6sKBFvjZO2Obn0YWxpuvjdc2STEdWWh1bdxgv1NBpRMD3oT8UI6ZG2HYtJKeyaW6r/H9f6pel3kaQA0Z9+aPHAAAAALEGf+EURPJ97HUrtL/mgidYi0TC+wiKzM13IdxH0hFqOnPLWN0qPOjbY0oUxAAAAHwGeF3REf3b7vlgIlKgE10AyLX4AJlNtNgJsyqIGh8wAAABHAZ4ZakR/im4U3dqjRaZYi6ufJW41oOZxfaude2YO0qLYL6p+tv2n+ZEhy2TKBugHpw+CmXFI0YFOx6UPFevBE6JTjViEV+0AAABnQZodSahBaJlMCN/9duV8/kCFIz5BMs12wfEGnFGw8X17TL2yufvVuIdnWBmq147QWoxvvU4sQmiODTxcU5IxrW48XUQ0l/UgndpZkMnXV4yvU5Wl2xRn7JOqG5tG1Wjk4im+e43neQAAADtBnjtFESzffVPfJNmH6c8uGGvN6Ua1w3dSVxyq52qBEQ5AQ+3UFiUIgQtkx2AipnIUrZN40brHaz8w/gAAADcBnlxqRH+BB5SnfnFFCFzzDarOeQ+n3x9mjYHrxdG+1G2y/RedFEpdZPCX1ehb7HH7JWFYy4bQAAAAukGaQUmoQWyZTAjf+8aEFzqPqIfb98OjTtEtGYksUtQHqef/YQxv2TlelbawxGTG5fTWV8g+MnDomIil9fbxikYH9XQmdZy8SyM4eT5UbSB3yEfjPZuf69KjR9VOuSEgKQch9dgwqBRFURNNn/ZHanJPjS32PTuCkyyVSbGosl1nzCvs5IhumOPeVjXmkR/0r0gyXtZomPanE/2OhFNpNTOQ88Qr5lABTVZ/tpyQT0FJkYsHvnwCpfD9gQAAAFJBnn9FFSyfeKBKqU691GQnRt7wT2bAwvsFlFpJYfJL3kI0Ejx2m9rfTT1n/rKrW+o3DIZl+x/z/z6dvQ5v53+xVKnddxG+xfI7Lp5qTATEEOogAAAALQGennREf4R4tlJElLvUEtQcMggGQj6hD8aXPWAOycZ1yxgmNuAfy4hKPpmOIQAAADEBnoBqRH9/t8AO1iKF8NFd5NzOsKBLGjI7tyNx83qt1R1YPIAs/dScM1VYQvLo/32AAAAAQ0GagkmoQWyZTAjf++Acgi4Oq6OXF4I6IPPdwzJFin4336rAUggpBKkw1FQCveDhbBVqTv1fDTBUjHvoqU72VIneOhEAAABfQZqjSeEKUmUwI3/76z44IUggGGDyEJHJ0zdy4sFJbfwH7K8r2+9osdhXBVubU94QlL3+0o16lzaMEUX+s5hDiLE+HU3QAsuJzTdUVHWDE2crLRMBDiXB3Nry4va26IEAAABUQZrESeEOiZTAjf/KCEHJoIEIqOD/LBqTEmdKsw0VZiIHhVnR0YHo2X8W67cycOsus55XJ0Jqd6wt0yhFaWJl16uDYWLglY5OB8KN8lyu74G+Q2atAAAAZEGa5UnhDyZTAjf/++AcgkyM+G1Ouhcn6Z2Bc0JS+SytB+cEXwbM6nE1arVgMsYki0/teQ7FCmWw4Us2D3g2l9Xae2Na6Cm9q7U3uGVww7Ls5vAowUGvf1+1PeObba2e2poP0GcAAABfQZsGSeEPJlMCN//77lcOD4iDd1SEdD3NEMCzmr6ToyL4AwseqHC/p+7NrCkSHk4I3vaFZNOZvIcubY0dF+R+0CSzGj75kO/jtMhszh/+0NqcZ6+lfDVkqMjcjmxjwAwAAABHQZsnSeEPJlMCN//8AtD5RPCcxJhhtvitPqIz74OQtZ8I0WhBpvfz/vdzPZ91H7CxyeX9rMCN0nFbY7I3EWhbNcHJ6VALu4AAAABEQZtISeEPJlMCN//74ByFVfBVBzCrec0q51pjCVag54dw81MGWphs9/vebHu4Slqm4xB0i/esr2DZbj8Jf6ynXo8mnM0AAABbQZtpSeEPJlMCN//76z9RIyID26En9fyziADdwqGitUiUfERx0K1tUeLRCixwBHGA8E5zp96UdN+aXTYSjyBXM2d52ocrCKE+B3/tBnS29ZiVSe53MesjWFUmngAAAFpBm4pJ4Q8mUwI3//vgOOEudjsIu1r9UcGS/BbdCjokhEK4p7ca9fqqtyh9/H3PZV5npSGQTV+CX0Y/0tMJZ7Ct1LC6HPc//deUSi7xe1m4BaSc4T6i42PU9PAAAAA7QZurSeEPJlMCN//74ByEPHhakORElru9VCJYbH2Yv4/MDI9PbvboGW+izcFJ6I5f62A3mPgFpsbLHqEAAAA0QZvMSeEPJlMCN//74ByFVsZ6j9D/F5iILSjxSQW3rwCXR3PeR+KppsBmj4ZYK9QEYuzckQAAAFxBm+1J4Q8mUwI3//vhUdlhbNCiiMoi7GsJj/JUSgwY7nhDk6vwDjSJ+VmDgx/g0bFuY3yAw7VdofesWO5yJ5HEM++OyZWsYg5c7kHgGfigSgYsdX+uEbqvt5ohgQAAAFBBmg5J4Q8mUwI3//wC0LvLElkFCZ+0R4VL+qEuV2to+z9d8c68NZwBRUtJQhv7rbzrKpuzMyK/+cyHm1wEmdUHgCIDOwZNtkrWOlccD0IWKAAAAF5BmjJJ4Q8mUwI3//wIE8MhYsf9rP5Kq/l8NreVCsrDsaQlU0Sws4aOq1OMlKviW1Ryzdx7P7AeQ57sIFdLR691f8uoHHExNezH1SYct2J9kEnZgyh1ie3x8zmVrr3AAAAATEGeUEURPJ+A5qFYNR5FvP+sUOFlh6OdT+5AC/ALV9jcn7A759t2aCrwOgT123/pve/kWbUHHscn0wZSsua2ijypozcsfdeqQdpwo4EAAAAmAZ5vdER/gwyV1/sg9S+lWzZCVlXJ1bZAB2pl+X6+0J1jK93r+nYAAAAnAZ5xakR/f5zgScsRbP/KCTlPetUWjSxs9oMve1ulK8unxRGQuTh3AAAAcUGadUmoQWiZTAjf++A44rmWDsmtE6da6Z88b/lOIgyX4xpZm533n9QjU4LoQqKhjYGX/tuiDl25viW0fxsXEC1UQCYY3Rxdqst2ZmTB87g9gEeOvp/+r28Wjk/lHMrim4Fl5xsal8b5UTEFgA4fhZ1xAAAALEGek0URLN99VLWJA1xrwZZgYEEgHgZAufc/vsiPVWvEmgCIWFM+Kch/8zg4AAAANgGetGpEf4As1+KMO/6oPIpubBAR5YI/mVvTOM5vFZZyOcvnNGLH+tubJVUzxsAV/kUZ8urzgQAAAIlBmrlJqEFsmUwI3/vrPjiko5FfFdReOlnYOY+aIn0vCx/RI4Xxu3GUBJIy2hCFvR4nFjOuCPWyJ9oC+4fg21JT1abT5qVCmt7KlU2DoSWZVc/093qGhFeYaBYh1VP3nhSxd3bKcoYQ++seKOZ67SE5jHbJWSiZVndbBPLb1CknRwfGQPMcrNB8gQAAAFdBntdFFSyfgOahSerkFDO4h+HT03gu9drqkmwjsbFo15C1ddns1scB1c41Z607J8TKKQBbJ7onEPnjRqY16Vmf6zqYenNHNjGFMGswbhSFbJ9SUbjgWrYAAAA1AZ72dER/hKNnC/BciTqtZU2dJJA+iXwBG47V52G2lmuPgOczguYbeITd64lNMXnwX8HviTEAAABSAZ74akR/gQeU4JrA+ZeKVXUJDRTmUyM8cptsQXz8j1l2X0x6PNzWmC3LWN9LRSLx9v3DwnzUFYMnVL55McNf8uFZblO77WrPcvXMd+GCzMVi+AAAALpBmv1JqEFsmUwI3/vgPUVnq5reM5Ssz2GdGe////9UbdFeHz6++R6IoL51iZvWWBTEo+w1/QGffKUbsN48+o5YE+HeD/8Y5yb8QquB9cRfL+hkT6gX4h7jUqOGLkv3aekdBT0u6Oz7hhVkeKCmJJIp2Op94DIntn/GHvxlz+XmeSvT4E+MDxEl2LA/adI6b5NUJrAj+zyjLOTgs/yPrd+rgMhLzwia10NmoNTcVyZRErfLqU+jEWkhGUEAAAA9QZ8bRRUsn3sdSq3XJB2xN6AjUdfvrPxLfuDa0V+3Ef2dPSVqKYe65mOrLxPdKn/0htbtlAPjz7kRvwqK+QAAAEIBnzp0RH+KbhUp35fRKSMEG8fVZtySTaTW+ISyd3w0mPrfX5tqV1xZY/iIFRY4QnnDGyfwxM3+zO/RMboY0LQswmAAAAAyAZ88akR/hFMwr0hrrENdSs5i6TQXQ6kkP9h/VZrhl6OVIn+wZmWbKSNLJEYN6QzCxG4AAAB9QZsgSahBbJlMCN/74DjiuZY2PbTAR+Dbj531bEXlChD5lrSbsGOo82bYl0bd2pDpN3WLdD1xqLg+1WuvVraz9dukzhD2dWGkCfyEavvctphVPI8P9WmurH2tugHweQ6rl55gisLfyBenMiMJ2RMJR+lzc1/3Kri59cV/bv0AAABIQZ9eRRUs331NpBIZUVg/sF8yeBOR0xrwAbP3QN4HMGL/ULx/iGepxoAV0VvJUzUyZfsAcCY/TGky6vUdHs44yyXro/pZ+OyaAAAAMAGff2pEf3c5fYTCslq8zv/aW5r90SK19ZwBG+JkrBByKG0soG9BdH2HugcqFw9T8QAAAJJBm2RJqEFsmUwI3/vrP1EmHKvUWSjtSRZlSgrv/liQAkTTh2MMOMOBFt7k9bYXOcwf59p4EuKaEKx7Y7E4B/IgnuToImCd86g1ojTwETdli09SA1Bo1HNUGbOLU83EsefzyVWJSbRzGhntADwtCJI9oXid1IkLBW7ZtDjX8pcEztbON/vJFzL5VTnKXjD+8vULEgAAAGxBn4JFFSyfeKBLyHXG9IGRaGq39J+w00CRNYhQsfHGXw+pdnBYY3Gs6lwhAdgPpg9sKS2UPZuAS9WFZHm4tDlDX2fuE3rMgJtYMXJgf90BeLeZPecojFY3k7mIB86csvefkqVua5t3fNOkykEAAABMAZ+hdER/im4U4K6eWIXO4QmBIq/I5sOTQqiQ0+cMCYfZVcNBwzmnMTca8JpuFSs22SO0wP2GDUnITSfVdndW3vFSPjn1UhcPn2zggQAAADkBn6NqRH+E1CpDUNBG/lajRglb3oTeJln85BeqE22nEJfeJpIKOlpztR4pxR7gHZelJU6RMf+Cw0EAAAB0QZumSahBbJlMFExv++5XDhLnY69ArjxbLeYiTGAFVjwNqRrhTUG41oczZ5EyBH+7dDwYd6fxBYacxJZsYzZglHmsDhw0rqZxKveLxdRtcYBSZO+liDMbXCCTHa+6LbkPiqpeWPuCFg018OjxE3T9dFSvl00AAAArAZ/FakR/f5yCyQKtFixuQGdhfIvIW7/HBKE6HEqyYkrz7gzpEkFnbRTwQAAAAJdBm8pJ4QpSZTAjf/vEgWa+oqExUYVp1olVvNbP//+8M/kBsXizdkbJTTB/BTvKR6wFbAMmXZWZPYbZpxPe4+VXweLkpoPRtXn7vfj18z8A5fxL1tFG0YfcZNJvwmwKmP55Yy9/It5uUnKSxwWtgpa0Xl7FOVHv6yBORxpBbFCVCVvuDxli2pBcO5IupgTKpGUhLJbY1ESIAAAAWkGf6EU0TJ94mVbOfACgqZwUyCA6PmiiXKLu1snrfMONhCCh+U404cQOLSIYQFmfISdXlAdVcHufy5BepvQ5vcf8R6jFC3uBu3JcyYPQL0WMxnXzTVqebdZ5fQAAADkBngd0RH+DDmAYmD5zzrhUCgMcWoHHz6RsHLXEpnx1pjFJJfd4M1l3M6R12v+1lUKb2QeenUrWvLsAAAAtAZ4JakR/dz4yZYCJktXoPfrPf/ePA889a0VEm8L/I5Oa072z9TwIkrlHV1UnAAABQUGaDkmoQWiZTAjfx6suWx+PTqHcCWZGG6ifNKCv4tUI8MKauO6cTSadburO7W9zwXrTpt9gXX3b5tdNvqHh7J5Cvv3AhRisceNgb7YMFsR1N0Rhh6/nt+TZ5PAFkbKAp6iJMn//stQ+Yz7VXAyMKB41Sngwiu8Skg+yDGuw5z0H1q2PuUD1rmITqnNCuI9uO8eg8uswHpLKHevgR52SAYvS2/Zid6FbE0GnXdftWRsiroevQ44GTyV0fPUfROdAmGxCWF/QXCmGcCNFD14g5Bu8qQ6Z+50EprnL1v3bi6SzBTqLcyvNtZPOKr8PQlWxgIVRruPOgaVr1gzPI3yA4IoHPl+VKgdU7iUksHmCQG388/DVCqccTrz+TsolCl0IqE+rbPC6A7gaixgqv0m6YoS+bdEtfIf5PoP5qHH3emI7wQAAAGJBnixFESyfzWPAWZ3eFR2Uc7F1kcb+oikzxVIfIGvLglj++hQHAONn2Im0GJpEyeStvP2EmyHct8tdtegovJLgQJsfAYqJrPchHLtBboynly/y1VWmuyJjSjqD9pkggLaG+QAAAHQBnkt0RH/N0btmFIvxlCtkLVaZxkluWs7V94CnZHEt7HGqrMRIPPdfXob/JPjyP65+5jAvNmhHfPXIgTy/h8MzRAfvoqjTEvzPxiXoq2YJ56mYweyAUgN5cy5ot0uF6KC7kwMci3mHrUqnOmZx7LD+L0r0zQAAACkBnk1qRH+E1vUcO/PuHwIOngiM/6M9xXvQGSxikpxLdsa+ZBaQwprqGgAAAHZBmlJJqEFsmUwI3/QsHiz8vl29gYzDA4k5/ybedXu7tKu7Jmy9KPWg3HFWvQIrtmp1hFrco7VDr7pRktQQWV4z7wB7QtPmuaPl7AbJCErL+KhMJjtnkudgiRpiC2pKeIZGLJlUjeD33utWpNyRcvZCNrnmrvCAAAAALEGecEUVLJ98ZVVId0czXNe9uHzxp8h3Iok3r2Td9dl1B/uClTrEreOFzvFRAAAAGAGej3REf3+JdlJGgwv+mvTSxIdWSJzF/gAAACoBnpFqRH+BEcHjkFzBJtn43EJVcXoHHh7rOnULz+Jbw70vSl2jJMWY6YAAAADcQZqWSahBbJlMCN/HqYgKv/p9ybhCWa37OXOoId/HNKBU7lW0ZvUv8UFLcBimbwNzzimRNxyV6iYYfu9yXFwXxNgztFgbdhKv7Yaw6GOdd3/mgssr3YJxxbqYri+BQ0gsKvOU/2bHepT/BrNNhFXz42GZ1RU3Z0OvhNwjsYp+atyzw6K1RQrGQzJdsjyTkZfv1NrgMxK2Kyq4o9lOjDMLVkWn0P0yn7qosC4D7s4H2EcqLGoPHxwbUKQbXH10NPUXB5VgGo7Gr2hDQFFekB222S3MELKaTdqBDIJgwwAAADFBnrRFFSyflmG9bXlI8HDsblWlpyLDNyCFSjeTVQjsKx0UBaiRpefXBhGOTqqno+v0AAAAJgGe03REf3b+tk7jCqVAx466Oq+thoM+UPfpYn/yReLuvw263iAhAAAAQgGe1WpEf9wsc7fqSm94QxwqahE+mbpNVa4sm6pOpD7UrOlV5wtZ7l4RNFEHIy7UPZEdoW6mGZQa7151EYp0yV581QAAAEVBmtpJqEFsmUwI3/vrPjhLnNsGv1BKYkfPo4cRV2egVZSmfCgBIfMKkEiLVx6WQ/+NmwmC1TmNjrOX/bDa9yyO1iT8c5kAAABEQZ74RRUsn4DmoVg1HkdZAFs1PtxQgT4vR9+7PMyPdh8OvTRv0jh6+gWn1+BLNvVSZYdeehlg2ag49VZovj1blaVazskAAAAiAZ8XdER/gwyV1/tYO0MG9HP/fzsOk/m3SBmPS7dS9yDlmAAAABkBnxlqRH9/nILJAq0W6Uz8mfnp8RC6e96BAAAAkUGbG0moQWyZTAjf/XblCEl4icEcUzKWIv5Sdz8xTnNoz+oTMnKI1DUFY37mUJgzqPkO5fZTe0TbtfncuGXCPon6ZueMlRvrNNcOK0XcUnwUA6DgCWf7RjzKphjvKjqti5kt8Itb3RMRziXZc2powyXYrDFolCRsJlmjzet99mqjeEo9j3KbahWDJHugtfI+sDEAAABqQZs8SeEKUmUwI3/74DjiuZYOyToYj+uJfx33+jD4oAI40PfC774m4iIF8KayeVa0oDX8hLv0VogcQ+NVxVP5bwWpReIJOBcCQwcgCOGRpzCDSfOHlF2gui/AgGYMXyRtA7HKV3fvBHPk2AAAAGhBm15J4Q6JlMFNExv/++A44TsoRqu9LyjXtmglYXGeW5gtQLLkwsKg8q3kkLlUIQztAr/PynJc//JwqcgO8CVKpBFGV9xrTYudyf09uTycpUDGmoGFkEx7CA8vXK/aaRGjjNBEByt6qAAAACkBn31qRH93OX3cYVktXGX/s02n7XjMhSSUiDHyWYYBLWZFRZWZz3x48QAAAFJBm2JJ4Q8mUwI3//vgPUQ2HX+ATzBMT/Pnu9JYrDfMcgGJDPQUmDYQjMezPulMGzFlZkx2U3ZoJt844xJBEP55P2VwxxHi0Kagm2yw1aPaGpOAAAAAOkGfgEURPJ+ASbiMqvfMgbI0jPmxHlOwJfTnXYgmLO8WX2zfxIhIhMrFN5/kkUGPI8Ahlw6/NHrekA0AAABDAZ+/dER/im4U4K6eWIXO4QmA1FfI5rFYZ4x7XjV9ezMO/coIVv++VYWmiu0D0vx9YATno5hyzWlfZMIOxG+n4NgyUAAAACkBn6FqRH+E1Cri6kk8ZHaVcOy1cqQKjNJ21dDGpEKstb4T2+G8eEqVGwAAAFVBm6NJqEFomUwI3/vuVw4S52M4EcFcQqXCxPpHL56kgKVqUEL4iX1OpvXCrsV1IIbzGfkz93JglG5FvhdntnM54UGgAg1vWYaNNzIhAfhLt1NhDflBAAAANUGbxEnhClJlMCN/++AcgmTgWTEr8if5xlRqXymDzn6MNQ/cQ/xqLIQJZe8l+jwFGz7ZyzAhAAAAXkGb5UnhDomUwI3/+8GQRIyCAPHVGA6jY1LoBJIpmQFHa7VcmtTnD0dyhdFuminjvN7hC4NUT3epXOwmxGLF6aq4Ii2C54UdlPXegSNE8ZyYU/M8lkhKwaBQ3c4nA8kAAAB9QZoISeEPJlMCN//74Djg46XORVi0pon+LFQxBrPNpEf7Xo0uPcjk0LUC1j144x6PKcVHoBOoaxOjtlY0/9T0DQfS7OAr7AkgvvVG0TpjF3TYvKWzgWOhJhlDHw+N//6f8iZoaO4xHGToMil4OnpnInBqExbvXmyS2Vvee0sAAAAuQZ4mRRE8331Nr6PZDrA7iBiuQ71R3LRXqWPzXk0TAwWCC9fePi3Wd1bTnw+uPAAAAC0BnkdqRH93OX2EwrJauelrxcclUb4bVAQe5fp/eXH3gDiQzy+DEM9azx07mjMAAAB7QZpMSahBaJlMCN/76z44IUkFm51/N882W7qXe+E4ycMma6jvo926yxLryVknk2weTfZ2Kk6jWgpuPKF0slNo0l2zy51T6byfECXCpBr/SZLSr2Zvw5ZFKO5Xae+BbRDytVWPc5sse+N94UDFDnnW+6PAj/4nobZzzGDAAAAAUUGeakURLJ97IKXjDBk4ZD9O3fhiXA/r/7ECLfoxJW97egIUifkFUbHvwS6RWe47OV/aXVeX/KUQSlcRcn9TJuFfK0IGrgObTvFadJzhLN17/AAAAEkBnol0RH+KbhTSvs1uWWoAkvTwy/vyakxvNOg/BBWPV+NaHNLh+5xiy5Fr5LFPB+8S8aN/EjTnuye1l0X9H1OcBtFfbZhvpHufAAAAQgGei2pEf4TUKvaaa83Ad2VtKlPt7siVdyG4UlrNxlJdXv9DrWc6/IaeuwrDtCf++mzEfYyS38DY6xAB6l6lKKK8wQAAAHdBmo9JqEFsmUwI3/vrPjggZKIvbXzqFTYjQ0VKKQs4oElaRxg0TMLtnBuolgaQrMekht8e3yiYuzVvUWS4uxk8FlkEzYMkLXoLetWyWvtKs75yWsyyQ2ot/tDY1AEK9YY5RVC4GnKBf2CZWDbyfWUm6mixHfVzgQAAAEBBnq1FFSzfhHWkMDYBGZBTsgt1r8IRHNzCaJx/FGcgoS2gz8e56BFv96s1sEBQAIJQ/tBYOx4of4RDU8GM0X5gAAAAIgGezmpEf4RTMK9Ia7m5tNYSF2IHBXygcGjvFaWjSpcJnmAAAABXQZrQSahBbJlMCN/74DjghSQPww2RGtCahPmX81LNxJQE36SVa0/HkFpDiH/yKhx8Jf2oCjWFoB3VUBDgDk4+gZhe2X6s7RXS3cJEENdXgtyyuS/MquQRAAAASUGa8UnhClJlMCN/++AcgmH6NMQeU+cJ8XbSKWyOtUbGfsUeNz5bQi7J2bKGDZFcCIS/9wonu6aEEEC/Vv0zBzYbwCfuAFMkFd4AAABTQZsSSeEOiZTAjf/74ByCYiKau+rz9PkmS/yTzhpkUb60QzRKs8rt26pMRvSkNKiKHJNJ1FXPEAY8I2SSl8M7BoAPB1/QE5fHYPJfTYyyOFiXPTAAAABoQZszSeEPJlMCN//74DjghSM+tgpIYflbxOyTaowlDPlB6SR2Rz0K+CLVW7IJQBHpZDwt/twQggloO94udHLowQ2+xoNwlqG66Pq2mXsLqfv6rBArcwHxAtHTOAqNxvb+8RDt+s4i3eEAAABSQZtUSeEPJlMCN//77lcOEuDI5sHb9JiVElFoic4mCH248Kt3rU8OlDO8sAxjxeauU3i7fx+UKBIcj59moJiQP2LUkx/p/x+YG9GV+7aTbNWtwAAAAEZBm3VJ4Q8mUwI3//vgHIUZmOEMdb1abutjrmvCM7g6DL5L3B9CMFG4FEvH5aC7uJB89RKDxxsntlPQ94dBlp2GynTalQrRAAAAaEGbmEnhDyZTAjf/++Acgjp7C+NxYI7ubVhpgmYIDTXkpEYV3zaW3e+rWLbW+n5QGSD3eCIEPH/DXrfalOt+n6Y9n9BbjQWVz2TQZgWcsZU85NeledliQG4TudhZAEF61udo0AyWq5F5AAAANkGftkURPN+EdaQm4CV0M/rYl4NgK8z84UIqrAO/czkG8l0D2j8OwsMsAK/qF8Rm4aPdZ1aSCAAAACIBn9dqRH+ALScLlpa/+GYDTEx6TsOAEaCmCVz4kvAr3WkxAAAAfEGb2UmoQWiZTAjf+8bbiClhCfpGNHNpGB3o4p+l/b+6Z4oYQfJr38IF6u7qllXxlyI04Ql52TXK1ASkOluayplXmJPCRsnnsdY6EoVDYN1QrkAehRAHXvIw8/6QTQFSPQMoerdoDhvcYrQcyeBLl6zvPm/NUvjVAocrsKwAAAB5QZv6SeEKUmUwI3/74DjiuZY2LeS976KkuiWNhYqOhk4zW/uyrwXJteQNDFt8DdlcWf/pTc81xcC8O9/HfjSyQlZjUCjsB7Da07CzCZjkJnPR1uZh2dxHaUVulf9ID7ewjMxqhMmoJ2NS8tg5XpdDvP5Tbn3fhqRHNwAAAEpBmhtJ4Q6JlMCN//vgRKlMjfz3lEv3p9Vb9LQ7v+pEQxwC/9zazAZ5tgWuhGJQVkP+XR/kRzKBKBvW9U37G1HqvI0/phIKl6ot4QAAADtBmjxJ4Q8mUwI3//vgHI3KkJVcGVSj2/3lOtKOxCchm5X08HMnVw7u8pSwMbvYNglvmYNN2GyFVjbMcAAAAG1BmkBJ4Q8mUwI3//vgOOK5ljVPnge6cjNMAzkc8L1XcMXvFe0EH49/37e6kaHHAytVBME3/FfgqqJuUggarQZYNRM/+85bFVwuRkE/Fcn/dlzM/xeteEV2bGJ+FQYk/XavueTQrg1tVZiDMduAAAAAQkGefkURPJ+A6iXj9eVzRU5pOyeyOkOK6siJQvqqq3aZI916sjslp70mJB60/u4PD1WP91IjlLAujUAFKxopKSwbqQAAADYBnp10RH+KbhTgrp5Yhc7hCV8z4dV9+Z54gcLr1jfWaZfLnK9DZ+XvNFoFhKEO39XZNyhFmPgAAAAvAZ6fakR/hNQq3WE1NaMz7BQMa3r2kDm/kxI5W0s1zZYe/BmWXhyBy4UWxPwkoakAAACjQZqDSahBaJlMCN/7w/OhfFZHbxAlIZNdK81M9DmgvjcG9HFcEgcYE3SBH+UoGuOALJjtU1dC0GdzLYklNnQHymyVpLOjxAFGP6fgPn39YtnABx6Iuq2ORMv3owIyCbWEiR2ya42PDNeEHorS3YxGwxydzC/f/trN+g10WiIaC/CCYlyr0UAogxl0oEvestbuqfBIwybgP3u9WkPg01fJr/nOwAAAAC5BnqFFESzfgcsGNr4KTYjcPlel1/Bo6Hv1V6W/hG/R9ldaJrwij6IiqixTyhYrAAAAJgGewmpEf4RTMK9Ia6xPrVbhCu9V78UOxQddtZ1H+yVxt1YlCTitAAAAeUGaxkmoQWyZTAjf++A45y/4Q3AZy8yKUZ0ax+WW8rFph6+XgDZSKK64i8BQCeecvNGQcmtcwgPyfTLBG70A7JuS0YXsmM/BbvInVHP96P4jH4mwkYmXHax7vqt2xdUMX1/uhLmr9k5kGIQN4Okuar3BDoEYDE27DDkAAABQQZ7kRRUs331Nr6W+8fh7tCKREw4UiRr64auFmklKazVIROH/WyOZMwEYrHbLo8JiX46PwpVX+28ttm84fKTzy4PcjJuhLtZLcNIyPsq30DEAAAAwAZ8FakR/hFMwssa67x1qPHT0B/ovzqciH/v7oaLqbN40xsh1mUuAuQt+HqX4eIGgAAAAmEGbCkmoQWyZTAjf++s/UeJu9eBQwVda/BhoIvFxND3yJyVrLYLCHqElIVpP4XAimELY478ot0SaWrEjaeTwo6tQAgnQW9o85VeSuV++MiiVdty8QsF1uI0JRpePl3og3F+kKHovqBhl43XtBd/3SOTYNkTJ+qnMfrxWequhajieauKGpmiI8nGWaOCva4RwfQCy5sjxyeWAAAAAZUGfKEUVLJ+A6iXj98CjRWanezX04BDJPUQppE9ePwMqecBmkgRLAxx0WEB61woNndqUocE/nExx730z4eqHj4y9zx0y2MLBK82XzU4taZ4D8bTJZKYzQismuE2WWwmBvIKWSg7hAAAAQQGfR3REf4MMlUQLMpWe9szsSyBoDD6MEYNMsQqchdl7Zl8i/O4VYyLyiKSFtAPuL9KgbCvEV7oC6OzFcVnFHup8AAAANwGfSWpEf4TUKoZwMnlUAsRqwSiK4QsOjRuBEUWYg1Fh/OCd7nuMtjeiQ5Y0p0H3UH8K47B3zCIAAAC5QZtOSahBbJlMCN/8I+Su9ssFpVth4bdS9WpvaBw5IeSHpdzxkR2ATv4jfUsYMAFeDTkERmjPr9LCIOQ65j74gzHKGy/tLT/fSolZ/CgvRfVDG/FeQDx2Sn0fsAWhLm9uQCA9XMosomEHoAokfqWTj0xGM8pbd0O/hsjd+5tP6naL8dluhRdb+uAPr/SfD6Q7T7RzN8WnGQo8foNo+PncggQZoQZ+z4TTUAfSjZP5rs1umERs2Zq5dIEAAAA9QZ9sRRUsn3xlVUfvinbaGcgB7351ebN7ARewJep+yTczsF5e4c9MzR3BPzL1klLWjrwlgKX46g/ueGxNcQAAACoBn4t0RH9/iXZSRorN2wtOvHL3rMKNXKhCQl/5dOgxZ26rVQNQnwhs5oEAAABCAZ+NakR/gRHB0lFsfre7+Bp9uH2y3Z7UlKi36/HXi/VxAc5Xd7dL3iB/+RZjRlA4K42VY9VgxSya8NR183XysSDJAAAAaUGbkEmoQWyZTBRMb/vgRKwKj6jQ6kUl2X4UvnvSwPjz38P+ukOsefshONivg0qJBk4d954f/XJfgcHAk1n/pDYI/bf9crnrP4+4WYWy892CnxvwMZGMki+3NlF1QyLUNlnRuvpjtnNZqAAAACQBn69qRH93PjJlgIlKgY8ddHWUf1/1OqaMkBqSnNFs7onwD8EAAAB2QZu0SeEKUmUwI3/74D1FV3aEE8G1Z5xdLRpbXXTFKbRJQjbqRswcA7HOC1tolGRV5dSZGwQ06qx3Sx5PBVNDREabAfzyxszynQeyjRmijN5plDb8xwdIxI2kDBBZy4ZZsO17xE3oB1eSevH0ZRSYUMODXwfK3AAAAF5Bn9JFNEyfgOol4/fAqtD33iwUqaOGcDWLURHTyerUUMGLF9/dTKp1sM6fV6ymzQ6Mw4D1HfNW8wMEkjn4rsCZqQMIPu+olAZgJJZsnp7/524VaYHVA6xNJqbPPxWhAAAAQgGf8XREf4MOYBhh7xX1+mWbR3kBqO4P31bzBBgLVw3Z4XE+EiCqAVQWrnOF2DmSkXwx2FyUnBheGRDCuJQDXml0oQAAAC4Bn/NqRH+E1CsRJUo8Q5Zy14nOo2/Zu7iAW4sc/CZgv2ou0lLXrm+dUmUgFE56AAAAfEGb+EmoQWiZTAjf++5XDnddY2O+hkpJCpGqqN+ioRd5fh6l+P/4qd3eXmdhQXWVfefzZYfFULX7dQb8xKswxzxZH07JzrnEToy3x7fjpn5UQgftlM6VOlHdwbe4xz/rGQAsG3VCvdsv9qrf3RI0iuneCBDzjBxBQQU+fHEAAAA3QZ4WRREsn3xlVUh1l0zC3bh9rNn6CiPj9vvy2i9BnSdiA+sxg5yXAeOsEXDrU7U1jqGjoUWpWQAAAC0BnjV0RH+AbCuMI/rr1UXH8O02Fm98S4fObJMRsbWn8m9/CguTUAZ/90CM84AAAAA6AZ43akR/gRHB0sI2Gg5rXgcbmg8L8jqVYtAxxuK4okCllLmEQxPSy4J9M1DJFgcgzgIfqUc+3vAZIQAAAHlBmjxJqEFsmUwI3/vgPUUr3PZU49GT22lznRKAoTW3jOcu6hD4ioK89GfSr06PHj3gvvsDv6O0stGp0b2VSw0E8X+l5NjE09e9NsWOSzj1xWS/c2SZ/lV+QRWe9Jn6Ul6stNTb4pd3Udctlx7cQDgYj8PLbS3Ij15gAAAAPEGeWkUVLJ97HStiPWYa+Sau356/Rr91SVf6JCh4jlNlHdN0d2PuQNUBv5wh3HT7qTnVQ9xZjQCXFlPLrwAAADEBnnl0RH92+75YCJkhya66QMSa6cPWkN5ju+y8d0nfVs+H50BoCjgVtakB5xqMlZAxAAAASgGee2pEf4puFN29kxLsO36+UxP+fgCnqU/gYafq9/Xg91ZfFnrFqEyyhYJCxcIhXzmEbbtohOCddejoKdPZpuX3ZH5v3KZCY6QQAAAAkkGaYEmoQWyZTAjf++s+OK5ljY27Vg48pz5a5+KYWTkP7SvlMxka68KtKhbaJTrQL/WXyvxd5weYVY1sWu6B/pmqrj/2QsNv2zSTXylPXylMXIEg2zO3tnyu/tbvMWheX6G7CPIGp+b6nvZ76NvL2x3bbS72nPa/776VwQDIVEqplwdPzD4QLQXjlie/F4jvAz+4AAAAYEGenkUVLJ+ATVljVg17/TcUb4E+TZf2O+oJvmehu9/UXCOUucMo+X7ourEbvh9gcx2HsgqBvE+aCc3LmoUYFCKGFauVj9/DXDPNwzQtla4m8OGx+lrBEXLX2daTLsmB0QAAAE8Bnr10RH+DDJXX+1g7O/ZshIy/nY4r4hQrCWJCq/V6Esd2r/5asnJk5RR4uD0AP8YUqepXEN7bk03ujsxa0LKkPI7KOCxKQmLykHBjvGJgAAAALQGev2pEf3+cgskaLZ/6LMcK/qtVrn4NRgFIs05qa7vLpR4QkwyqPFyrTb7wcQAAAMFBmqRJqEFsmUwI3/vBnCJ2LJ1QWuVMjOxOqnfxxjGCcowTZ19I1NAFB4qvy5q6Ryyp59tnEv6iWYDfMwcl+posZ/vpl/L1PLmpbuboHvedMTA2YWPl2SGbI42h5SLocWEUqLCIWgdUiQWfje697D/3h1HLpDMRw7LLhOe089llXSaMRlYUcNK3n6wBz15SFB00Q6ihaj5OREXBbwYXwrPNfyFffYIF/ieB9Rb+wytAu0366VH851MnHmh3OAwWWubcAAAAXUGewkUVLJ94mVbEetQUbYugF4NXqOohbWdmnQgq5IikjMhBE9Q8RRsBeKAYT1t6fLH/uQHXlB1P0aEmEa9WNTZMw4CcPsqSexTdEwT9nlxz5VCl/l2owvjxAPVh+wAAADsBnuF0RH+DDmAYmD8VJbh4Nh0FynC++OPAhiHknCimadJ9Fe44UHpkWLeV30b4H+zBmAv2TIXxIi3YSQAAADABnuNqRH93OX3cYVktXGXvMuRi8Jr6Agx/PtZNe1ki1o1HHbdUiBBUvLJF5NG7lEEAAAB5QZroSahBbJlMCN/KCEGWN8UZ2OksJynlUtJe/T/Ze5KZIBmbExdmqiGR5669n82a+AtLkfKGHN3Rp9oU+HakLXqfrXevKtLZuZD5XvF8lGiaSAkwjZjoDDFe5iz5YPRcz+fZww/DpgTxvDkkfcVTbleWnl2qZXgSTQAAAEVBnwZFFSyfgOmmneDDuxdaVxDngEJwiO9Wub5u3Rp6pIgiyrZeIaoVDz56LILYAvwi5cN7K66uYsQoR7hKrNuY4y7NEoAAAABIAZ8ldER/gwyVRAsMRG9wSk8dVUzVjXcbzTmj7t7tRjQ5erWnxVd83Ef81CF3ZzwpPUsI8IPWyeBOnVhve3/z/SC41nZMKfRgAAAAPQGfJ2pEf4TUKt1hNTWjMmchYe5SAUxvqnkg1Qw/OPUc3DuBY6sbZu6KUUUigXOtNJ+lIbfnS6xPA+D1gYEAAACqQZssSahBbJlMCN/7w94eHXz2QmwfLeHxPiiF3IVqCWcYhPGBobbysZlGTAT8YM63rCVG0s/TvYgSnMzoYtwQNurg7iC/yUMmiP2Uprp7HMupZaNwHNYfBppfsDg25dFnATbutk3GeDH7PkeFPJAtlp/QQQHUqqrP5NTHvXTLswfs5DENGEt86b2DoFGS/7lW7J0rnfMCbX9Sb7lVrsBZBpHc8xP6+kRPmeAAAAA8QZ9KRRUsn3xlVUh3RIT7WdkVbVtgz/uvdxDmkzwqBzSZwQ6hkamoiDiVfFO1EuO/+qDNbLAf9gNMQ6eIAAAALwGfaXREf3+JdlJHCJoZnuoVnuSxVAFzEoggTy1eBea3xrPnNe7ErbeJ/f2BanpBAAAARAGfa2pEf4ERwdJRbH63u/xqE+Tl1J8P1I+BCos4BmXSZ8BBe4HTICM52Vgu9SxEvs5LsCSMoz5HkLWd8bnJC+JCYqHBAAAAakGbcEmoQWyZTAjf++A44rmVvKoKa41M35Ch3G13lWY0bzQloWTKokZJkHadhknWNHe6v4MeLVqdkMXc1nQjTtXrx1ehZRUoti8Ga75LuX20ESKWe4YPhUknmNLVPxY4QiKLz912RQJJScEAAABBQZ+ORRUsn3sdSu0vw0w5exICsmzddDcONsXZutRG82avccQdJ8Bo2mLUJmrRj1OGq0Hhe/tpl9GLDnt8r9hgyXAAAAAuAZ+tdER/dvu+WAiZLV5nf+zTESAGqSIRhNll372WSzssDIEYXudYpMGm62vGTgAAAEEBn69qRH+KbhTdyyn2mc7fr2krca3/fg3IzUQIeXmjyL0W7kM5SHpgQjyEjdPMwM32smfci5GMHzDsab6+RotVwQAAAPdBm7RJqEFsmUwI38EzmEnMThj01oz/8HmIHDZoyWvPgj0C4DoRsDjr2VEdnSM5F+6Es/HhGkAFmMMlxRO4cBtg+v8v225tcKWm8Y+JoMo6OE9FQeGP80ZRZKcXaFITb4modkoTFB1MOoASSKLuo3eyB/nUlTi7UcHk/YyDhoKrCsj9X4UhyECyG+WaaL4nUsfz6qAEXwJcfBhGzK5jzIx4VGoXgDJKsaf2nsV0cr72ztqjUEW7uJ8MhZ4H2jvG8l5wh+DH4SjA+Y4thclgZCtEvFf6uTRjHi/V31/zDeILkYiFDRYuc/uiN22xbZOJlfR4cWsUKVWAAAAATkGf0kUVLJ/NY7p5XjclhLFPOm5V70UHjDk0y1W8WLMxUEJtyFr8LEtGYvu5dSu368Egd9yNd6Rj70InGXN93csHtMEAObKXQji9PQYG4AAAAFgBn/F0RH/N0btmHtB4in5wHCfslGyNeY6FAHgkfpAI6sJ7Jpy/UlPtc/zQjmrUm7P2PZ2IpRrB2bO7p3QSB9jQTHvmkqdIl64ZFcKRLHTYUzBmd4ons68pAAAAIgGf82pEf3+cgskaLaDSUkmmWOh29hQ0iDCVqc+7lecaq3AAAADgQZv4SahBbJlMCN/J932y9VLbElOTIQnppSdx4sQM5Zt/kb6lOpe1/qFSIxuuTjdoDeO039Y5hjcJKGrutCBLHNOGFrDitKwHUBx+wCBz7FNfIB752PfbJllVBPkMGIlohflE1kmAWoEvdKQ5RVEIJBo0NM+7Xc4QOrGwlm1kYRsr+YzMe/a8WW7js/AQvf6/YKmK3nv/WmouN6JrifggGWUpt463j+VsZPXJT1MLmlu0ZIJheCkppBM5RK9Dn0pWoHhZqpp1LS6KmH+OIyOhplBJZEO/pM5ypaWA+VX5TGEAAABDQZ4WRRUsn81KWB6R7ATs7BLkl+8C/fhmkvrMuiNDekBNsZe2yrQ9LlYwQIG0r6PzAA8CdpaaTzI+B+gavhWMMYl1MQAAACwBnjV0RH+exIgT1LCdOIZSo931hcnbNVCgfvajMBfvkpSj6y79gzRE5BB2iAAAABwBnjdqRH93OX2EwrJDk110s66jvwel9iFQh+hHAAAAjUGaPEmoQWyZTAjfyghC9V+aCM9y0tNTu1YqY7isyMrEL1ISg0j2C7zj1UUqJZwpqqKi/AnzZgAJPEx+BbF2G96bWd4yUYwWlUiGqh7yxHAdI8eIWqxQwALIgmV6FsYG0q8g7JGYi7Gagx16ZFfTcfzINdt3vHY5lmi4F/KgeU4H+REU/9BL/jlA5mhrTgAAAF1BnlpFFSyfgOol4/XldkOz+vOrEWi1Q9ySLFxK10HVuM/uqIkihgxaTiEJ55tGiib34VbPx+F1i9KTeR0sV4TCay5VXoDlJbQJ2X/dXyUkz7tlo17BgibiD+MKM8EAAAA2AZ55dER/gwyVRAsFrTC0Y1O8nU6PfpwHrj+ZB/q+LVk4xtIzN+fxT4G1VCHKhzISBKl5tZm3AAAAKQGee2pEf4TUKw5pbspc1EHDtEesCCmJf5VnNe6KBZd4GQ+7A+bHfJOAAAAAgkGaYEmoQWyZTAjf+8Pt5W8cTDD7lNWFGUFASMmfmrkidqCky5SC9+cPiAP2S/BFQOm2nlMb+dDFgmvdAP61Ih4Y1f7BfhGuZMbC8UQVSneMohC5yiCzKNIygFEz4V6khVAoCjtwHr78chXAgv8fXLr95q5eD5Lu5Nzs/iIGtMQoXUAAAAA3QZ6eRRUsn3xlxLyHWJcFQyLMUspqSzsA73febdhh0qb7NHDVptBqRvp/0KYGMiwx7j7KbYTwgQAAABkBnr10RH9/iXZS+PSr/l8r00sSHVgOJeMGAAAALAGev2pEf4ERwdJQ0Crb1298v1xh8pKagceHwIXpnt+MpKlEEEQXCRaEef9JAAAAckGaoUmoQWyZTAjf++AchU6MHD1O3v5r6pfeZSR95Rrhxom7uPZHpFqjal5BKKKMTbN6mgvg8mD64Xft0kUDE6Cue5SSiFKoyWlVEwiX5hDItaXm0RmztEkwGguvWdr4BDwmXp/9l2g6tLg2xXeYbPfQmgAAAE5BmsJJ4QpSZTAjf/wp7xIUfno1XJ+iymWGbtpLZaw7PcYW6RhtN1hZk9b9PCNe5qkJxn7A1KYQZsy24bK3hBGA5yy6riTS/6Bggyp9yOsAAABlQZrkSeEOiZTBTRMb//vuVw4S51+pMKjT792mrZCg0a9enAYDaTD0+eYSezMml8HsNrr4W4Ypl8E3Ek6kEobjk/rXZ7COVAytpQu5ZP0ZvW5Pg7D5NsaUt6+L3onc6uiUnEZYYcMAAAA5AZ8DakR/im4U3dpdyjXpuEJWyVePMxrhvDme3b6G18sYPqqk/CL+3u1bRHyeuiRn1bp/7qnfDmkRAAAAXkGbCEnhDyZTAjf//CJ/zIEiRn1jC9hpPPLN5QuFZX1Los1jhijhPS8st288jfDlg8K2nAHAViYU4SieUCj1ME+It5LvfykfCPuvpb78LHOE+CJdbH/LeiWos4RGOk0AAAA8QZ8mRRE8n4DmoVg1Hkb/e4SXAQqfAAXFApCfAhG8pf0RtikgWpTEO9+lQ2C26jfa14IxlqzJJ16I5H6AAAAAKgGfRXREf4MMldf7IOPpaY2kRu7VY9Nr6GAEHU6Xd08+AnO38OezzPbYNgAAABoBn0dqRH9/nILJAU17xAvr1/4H4dNFOjsFwQAAAGFBm0lJqEFomUwI38qLkf22+XhQO6dWZKV+0UxaeuBPWkT9AVqu95T7qYZaV8ruFgtLIWIwIFx1O4IUmJVkQy1NeM9AOftya18H1uabejHiFxrqGQv/o7wQJwgNV5M+HntAAAAAe0GbaknhClJlMCN/++A44rmWNd9junV/0ZL8GPrsR5ejMuAeKb0C2X/4TBgkmpc5cbaAmhsRGzzPOE4Fd9ectAGJfFuKvehbLb2HVVq257BpM84fIVQhLM31NAPpq1x6jVFAQ8vyPmNq8qDOjXDoMYQIaRde1lJOu2UcnQAAAG1Bm45J4Q6JlMCN//vrPjhLnY6q+y84SEeloUKkmwZvTuvQy1vBx1cbwrmnV8fPOJEA/L74o4k+P0p/PqnHXGHiFX5kRdkSTKdwmiFccfDJFoOJFk9OTzWDgiBL+BN1Qp1Rsy9KQmgSswUoMNqBAAAAKUGfrEURPJ97HStiPWYzn6tT9FPiyYP9OyFj9PCcFZvsfLF4FBcBfJ5XAAAAHgGfy3REf3b7vlgIlKgE1B7NQWvwATKbaZAMXPLJFQAAAFQBn81qRH+E1CsMICfUaEKDBpIhdx6OnKAdCl2WDjL9pRb1pMM4S55WoAWOuHs6SIZMvTEbf6TD+l89/qm6bCPBnWGYDYrn8PXt76kldgiGMd/HEyAAAABxQZvSSahBaJlMCN/76z44IUjPWzut94XBkNuv+dieP5cTiu+U4ud4qkSdtDohXcxAqxjKAO9Sq7VU/iPClB4mMSKfAXByslZgH7DJTaF0HqcQoHsXQn4FOdmjHs65BRr5LHgnCj0HknjvGWtSlUvBUBwAAABBQZ/wRREsn3wG9Ho3WAS3Ni/Hx+OpCe/8+ukUSErgBDGT0W0PMnbTfONFyGyOrLRqgi9mASDomyObYuNuslVCT28AAAA0AZ4PdER/hRLKZLTcnylaqqMzlu66hI+mhxsbMLljcUK7FzrBO7JMAaWFrqACUm2hnz5DwAAAAC4BnhFqRH+BB5cMaSRFVafg6PSCA26I6z1H3nEcDifB9PFmvQH6EfZwAD8lDAP9AAAAfEGaFkmoQWyZTAjf++s+OBCWQr1tn75c/emk9LAHs1fTdTbpvJzb3LtU233UzJwGYEL7e8Ap8RSOlZKi1Ur6TnLiFRd10qK3EQMsUKDdzkCorRRoghT23l2ZqI/4juhRRTX/HniY+Qj/rMVC30KuoLxdRYifc+0btYjQlp8AAAA6QZ40RRUsn3iZR4RsZYniIDRxKFBpXUAJNxgE5p6+BezvDoyEfR/+vl3w9ZsZ9oIbtyEjlPrxFfc9iAAAADEBnlN0RH+DDmAasfgFq4gCUD1KaJGjFNBEJVYapagXlLZLSIzJIY2+GzlTGV7riDN1AAAAJgGeVWpEf3c5fYTCslq4y/9ptS3UG7hViDeebQUTeM3L6N0COirdAAAAYEGaV0moQWyZTAjf++A44EJZY2AK6h8ndtcC/MxKotOHWqm8LCc5nxMDTrcRe8XfRAD+sULXlYsOMbyenQJfp9B0zySWMkOE1fdeLQFXHwYudsCa1bFocAlUNzU2ll5X4AAAAFVBmnhJ4QpSZTAjf/vuV1RDrkiUKryvNy50VVcIKHyfc0rVEbcQaQOHTvicnp4trDKdvq4wTZVbX2wUqr/yMRShkhjqnvVzqia/DOdlZb2Rb0wer4BhAAAAW0GamUnhDomUwI3/++AcgS6wSqZJjrGWxPv45j0u5gayBIOWxzopmY84hnH7gzzrzeXgUYnG7EGJiE34rnk/Z0xY2sxq0mDVmu0uyDqZKWyTXKp6UClgGZ0ewkwAAAVfZYiEACf/vgy3mIDfV2I+c4OQeUCbjkIqauQ7q1hI51t1KK0Ty/hpqe1v3iAKauIQOhjuzFxyMgx7FbnM/2jBWLSkwzPM9uWS0H4jWSakogzfe9txEs4c+/QsLeTPvLzlrrhoRvk76vX6F462KFV4lQ8BLFX1JpZOU12xKsmCPjGaaOJgyXQiOhYn2hWzgcqRNjGX+V1keWAHLdIU5Kt6ThBKFN9iXo020ZIGb/b5NjIEtf5xxojLXamTjllPOcFI0jqWkWbZoIKHhmS8TjKSHCuTIx9dl0FTzgEkfjAOKV+MBkOb5p4/AbH1ylTOdgDP7a7Qz9MRQalIH6M4QHnoVD53mLhP9Mm+R7+Nv1EH6Th++wBzxhKbERneXaOozvtmrCru0n5aeRSyfurbH+PPpMF9pcN0wV6BXo60O4fisdnnNQeaTOQSwEQoA9l9G4isMAZz0MMRGPJfTWi9P7ynze7cFpAtxFXNele+QfVNTyKOBwDD7HFYOx0ouGd90huFZ2/V0BetUuzhUO0amnSSpDeMcYfW1ZAlHhEdZedtfna24MPXCAzeTxN0822JnETBRQEqo5eqoRZEG9h5BpWdwsZ15MpoPzXeKTqs8tVQC1/SK0d/0/xgwE6pNhIFwxgVY7DWsiX9Xv6tLaOudQGEMJp81OTR0L3+PIQF+o5vb/GCWDAkSsOBZs62OY1NlmA//1Wvizc5pquC5tnGGabusyA/nmG7PBFvvzvk1g6c9FZTfMWn30C0nuVgD/+SL3hOUxNRRIZsd15UgRRZCQNTNuPsNF70mEMar3F5ZGy89n5mtWtCMnx6Kxmzor+SNb9cdhFvoJ0pJL6GWtPAX3YGiIajfBHoC5Yk+AsjS+erTnaEBKxjzK9GpM93N1bO7ckHjy4hYhPKAFrXZD2uS3VtXe6oTBsitQMDgoUo2V1oa/61b7CoZXWhVK4UEX6j5mQQQ37LogaqDE0eR/hnOHk0qp1/ZLliBrAn0q0l0jYC3N9Q8Ae2vkxGonmaOmrZW/AToLHGMb0emiADcOSePLr7YHGJDiVzPIBiXNYDOLMcaDwCYGfbEsBrD2uy/d1Fv59oWvDv054FejnrCwPJNhHJardC3ukYi14OF61tBTEX4on2ntlVKj6TiEl2AY60ALC2Pewdc9kxhuM53QE/DPE04ecu5tBcDncARadW/1to7/E13lltojS2w5x9whdbhPouWVm71+s5063T2CXVUoDOAGHtwV5PBLlUV5s+9rM1C1nm6Ly8oNNr5tCZZ+xgH/cb0QtGLo2/+urcmwovBsXN81PlivdxeZpmL/iI63XLhkruWXb50jY1FpX53VswqMFVEWIGAsoxcV1hcpt5SjjgtAJB9brynD19/u2V/j6WCCVV1YJc0BsMrlVTR84gzjOn9fl/YKUcx8xd5zXtGYTTjr+P/TYuSs7g+gIPIJqW3t2VM/TsVgWKT+LKDZouSIl2nMmmD4XKzO9dTJuIBhm0f3xWeCHxstiNy8dyjXq6MgdCHeN+Zcag5SxJrW/q8LCIwYDIjdfeTEv9XV8VvX1oPi87I9rfBzE/XZtG1OijeZ7AcvsTu5Q/+VzP0izbIAexsR6SN7ywadj+XkAq9hjRaN/3QPlLeJxa7wBlxl3d1FQ8pFq96dcjnDir5ouXNKkIk/zCOM+q+YORr/PMY//AhAiSwED9KgwJz8yR+k3Go6EUpMgLgQtZ73cqHm51vCETs12mqn98YrWLHREslKlJiLjFIFQpWKlBdfkllUxz2P51TqpY4sD96jVYhYQ02ZSpJu6Tqh26CeFwNmHP1DrxMk5MJAyK6KFlKGZZi96tTwAAAE5BmiFsRv/77XFgf/gTkah1R5Us0Wnk8dDFNJtGuWmlTOExBPxdbP2u422B2Lto9nlgbofDURRsHjnidXKX9OQPQ2PYSuh0BmyDHyTamf8AAAAzQZpCPCGTKYRv++AcgTjLzdvWTXYs/nqVO/kWocgnSkHh3DwMwZdqf8C8gKNCxrEXojVQAAAAkUGaY0nhDyZTAjf/++uEdiIsFrkFNTDNvjCEHv2x8z+IOmFBPu6GwLYUtDIjwQ9PmO+bcrpj4clo9KQ/KSk3neQOZHILev1Y0jHKDWKlbBJSNC7Jskuppn7Pw9tMIqVOLcYm9Im1C+17BGV5dlf9xePKMg8Jr3Ff+/cKGd6zElUtjcoq/1384Z27KW1jNv2lQ+cAAABdQZqESeEPJlMCN//74DjhLnY7IVYsSZ54teV4M9pVBpl+jrm2vgkpxJMVoCRgAjdr9o/eR/Ca4bI5DFioJpFA7jz3z3p6Z8EP6+oWPcI57n/7rtuyon5e2HiX3e/BAAAAPkGapUnhDyZTAjf/++A44TvpQ6q1oyiGyCVbHK/f15aT3JMWQoDuJk0FKSLab/1polBzKBKCBe4A/5QO3LPgAAAAM0GaxknhDyZTAjf/++AchU+Vtkc2kb2jJl/oQwyJciND8Tq4ecQix/2xNi8PSXJvoKQxIwAAAFVBmudJ4Q8mUwI3//vrPjhLnY6qAnz8Ja7nh5BGRfd02joZMmnvwGPMnY1UtK+pb8bqcPrtmFQYzPi2qj09uQT0JTvmvYDd63U5Y2jF0v2YgFFNEVesAAAAVkGbCEnhDyZTAjf/++5XDiuJrqTwQnNXojWSFTharkNbc7gdyfMlGNhg3mVGFS0alrqG5sZCfIgbqaclT6Gfe0OipRwkPfdn4m7Thq/CqsNd+JyxTOnxAAAAcEGbLEnhDyZTAjf/++5XDhLnNsgTYfDsmrw/ZxClJj+Si9cQ6o5fvdg0roRb7xCts3+qftaz09ytTf/Fu2QxxQHWzmVyf92Io8RDVoRP2XUDjiYmvZjXggNq8HwiPdQMdvmL1lhb5fPZGj6hWNxUXuEAAAA+QZ9KRRE8n4DmoVg1HkW8/6wpqyGi45GD24TJ/za6DAQ4eU9CDfguO97+RZtQcexyfOmagRCfqhL9X81Swk8AAAAkAZ9pdER/iffqq8P35xMd1IYhD9d9sBpVn7iTWzwZs4q/5ESbAAAAIQGfa2pEf3+cgskSU0kAS0c+imyjJea5l8GNLVRLfBL7wAAAAI1Bm29JqEFomUwI3/vEhS51HjidTPHSoGoQXFVXRhNlrpyj41Xz7oGDI06DVKDH46CPr5apO2EoNwcdCUxtW3hgqmJ5/H1C2ANjl0mdlryKc9SwOs08mAQdgVi9wAa3am2ilj656liiY2Pr5+C/f/vDURY0hdLqqHzbufbgWXm0W7U5ed3s5e0CGqT0N4sAAAAsQZ+NRREs331UtYjmOPqgDvwpA+dkDRfmu4CZBFY77Ij1VrxJoJU6bh56MYEAAAAzAZ+uakR/gCQt04OPppwlRhivsWAxBpyC7wzzomqfCVfnmfzfbSoDKVZgsjBm+Kt2f+3YAAAAjUGbs0moQWyZTAjfyghBljfMMuyWTw27Mmxw1DJv9ycLWX+QHLgBj6IZtikvzi8lkYN1RnpMqOZIVf/nduC4IEoAVSEkctakaQiAU2DoUeITXzdRvYi6cLHFhXwRHpVFF9boSkQf0W8GSZnPKqex9y18cqZ0/3qjnXDpyObZscNjiTUYfl8+sVcWkDfb4AAAAFlBn9FFFSyfgOahSerkFDSgMmtywjBbMGKUIhmNiwJ7hm3fI/sR/In45K/o/zgHqjSeTGmpbhIijT92cAprRPD+g2EjPnGa2bkpYhmAbzhjtjjTRAhkkbC2YQAAADgBn/B0RH+Eo2cLrfP7Es5hEIvxy3++81JIcbjtXnYbae6sJ93XFc4LmG3iE5toWPTIgZ3Rf4u8wQAAAFEBn/JqRH+J9+qqfNhEqND9NFo1QljVM1SBAde2cQSHX/EgRDnk8wp3LTUqgB6d+8T/ceO3zoPFZUv5ty3ZB/Jay3Kd32tWe5euY78MFmYrF8EAAACzQZv3SahBbJlMCN/KCEL9LNaI8OpfBjrh+prQ5Hx5F8zmd++rK6rSSpIp/Bmp0IxjryV0UhIqqY7ythB00KcAsr8fcfrBfpnJuvV+Xgc3l0ufktx7e7w3rU7jqfUydXix+/gToV3m7PqBlYaP7wIonbR+CPd7LoCHcFtKXmQsnVlI4yUAum5TYseBCN+1u3fAF9yw2dGEPlv/60B1TfKQhsyYxA36fR0Gme6stPBqWq3qMsAAAABCQZ4VRRUsn31uHUcUqyrx5k84MB/5PPciwmpzACm2Qtdwi7HXVZEde7qqTsctIRWclAzfvXhjFp+5nog7nC999IrgAAAARAGeNHREf4ByGO+by1cZ8ELMy2ttDhtzqmeTvmfBobP77YCuVEO65NLXFqkgoV2VMcV/Dme8id5gHSZgEfW0Usi+767BAAAAMgGeNmpEf4SH6CyS5Kx3XhSoUhAXrSLodSSH+wraAIhCFkqWqwzMs1mFxDnkzNVjPfe3AAAAdkGaOkmoQWyZTAjf+8SGbo4p/8Iamyi86brFs7iSG9hkRmxUtyMqDX87BjiDhTPFbYOFwbUDOtNIYN5sBPm9LHJf40gfxrHO5b+yYb+iAzQ+D0jqzD0wUE/GPk/is1vbvFMqnhMOuYRpHxFQPFZ7rMcJvsV/ceQAAABJQZ5YRRUs331NpBIZUVg/vHpzTqi1adLx96xSVQdbyZNxZ1wEI8f41eXgBXXen0pb3PYiv7Ia00pmzCGAnR7OdtwQ61Xzb3asQQAAADIBnnlqRH93OX3cYVkh8eOukCbZ+6JFa+s4Ajf1rEP6NxJzT4vwzyoFKEMzdAXFIlOjEAAAAJhBmn5JqEFsmUwI3/vrPjitsxqnuuZB/Q3LK9EB8XZHkxOb9shEL9bJjdZ+AyCA1XlQlALxJUJQknuX92nj+nF0tYg9+8IRFjYGwbn3xaqO5g3aC8XnhoUJD7JtUL/BWe7piCPNYHDhlaejMHPjLpBeweLNivYXRppq6jwaS+E6aqBmMXuESQatm6wvSHErLtCeVTRchSPIcQAAAGlBnpxFFSyfeKBLx+yIGhVJaNqiSbf3n/EhW+vNSV7ifu7le6rz/EoChwQp7PkLC2To1qCPmWejORWShLyp///Y0qa2M3jqj7c+mZFtb+jzH2ON2C5/c8hSTBB7jEfQYULfNJnrimd91mEAAABPAZ67dER/im4U4K1F5C5u5/eTKivk5Nnz0mjnt5+KqDGfDi3wS+Mc05ibi9UxpuE1o1aqd7CTGoYA5lcK6Hr2V9Y1yJwB7+JdsC2pp0re2AAAADcBnr1qRH+E1CpSs1MPhRiIa2mtfgUCd6KO7sJt69FF+hJ5oUlgGPC0EPB1ogSuuBERwapo5fuTAAAAhUGaoEmoQWyZTBRMb/vuVw4emP9PN/4ewbFLUGiXiZoXEL5lOAJuqN63bq6Ya73SLuB6bBVPnpl/y39gFl9DALFqt3PLXqzNMiQ8uGXt/+TM86YgjzWBw4ZWnozBz4y6QXsHg/Q00M24YeHSO++uLeHaT1KOziWaUuzPIKggxZQgUTyNyYAAAAA3AZ7fakR/f5yCyRJTQtCEVMtRioHnim37d2/B1xh3Ql9UxqfT9HDmdIOoZcKX+BDfTVGuDUOJ2QAAAINBmsRJ4QpSZTAjf/vrPjitsxnTnKYcmEudJHuBU34rmoZfuXG6Cx5BfKY65GVp4NsLJ81tu0fHLYXCqgIs3Opp1W8uDzLiioBJ/PdQ+EnY54EcqZA6tP3J34zBz0T0ggJpXUzfiV7xPL0uRn40C/UnuyAQslAT5p+tnCRVSaGrqcLiOQAAAE1BnuJFNEyfeJmV5SO/7e3+IXPu9DkhOBWthtp0R6laPa9QVTtdYLuaAVCQfvIfjc6J1xsnVDkt6SMF8jhI717hP9wd4YjWpEepJyiGOAAAADkBnwF0RH+DDmAYl2uPzEbmrBU7+iRoxY56o643yA9iq0lyU4HIBaatgNSaUrc7+QCg2zU0nrhV7MAAAAA3AZ8DakR/dz4yZYCJktXmd/7L+xSWqxVX0hhWCgTdFXq4mp3c/1a9hxSgDlfk8LNj5Vw2oTK18QAAAU9BmwhJqEFomUwI38erj+EOVPHy32Fbp2Yz5YGdj+kHyBt317o+i0bT5HgD8WT9xKPFw0LmVzgeNi1FS2ii0XX5T3ObkUYS35DbQo/q8fsZOm3CpwwU/Yr0/kpevNm1zY3V4E6rdm9gjhcJayG8aqoO4g4WVxhx/xMNf2JeC2NfVjzZe2vcxSFQ0+nBoYdXCz2HGhMEqn1Vj7giKBRPbFIQfXzgNot7kOQRpn9o7GlvIh2PHH8O5b9WkObUKMB0E31SJnRODa7/x1gMeVNp9x5QG1UQTgxsISHm0WiGHaHM9FndRzybD9Ti4Axz30/xzLbVD2b0F3de51FeapbwWFhm9ug8jx3oU8tVKEI+qRwwZklPZNSM5eB1TunBJCGBjo47s/DWaVQJ4etgtBC0e4W/geamug1fhr/BERy0ommj6LhH/UwJu76416mGlYt3gAAAAGRBnyZFESyfzUpvZmoxPUsjVw6GJkSBGLuJUn+Aaj9PjxAHugzMhsCOH1qsNMVFzZkiMTpbdmVc6H9d/LXbXoSjfpgIFnVOI0u4Z1YosRun2lvX2y3QNp1C04NrWmN2xpZYm7uBAAAAeAGfRXREf83Ru2Ye0HiKfnAIZilzHLjlqsw+Mz1QC7FumLzaTF+pKfa6AzSFSUgrNZ85pfDhXgDDUnT4YfcFFRuWdSKxg2T3+RWpUKLLxFjqhPOFsq9XQnstbGEYxiunguGhZ3eFb8BcZaO/8shwvJPEZqscqq/swAAAACcBn0dqRH+KbhUp33TVx2AKpAjXBFaDxlFOYkOj2jzXxGLdVqUiJCEAAAB+QZtMSahBbJlMCN/7w1h3aeLJhoTP+9zKomPaDpIjqpuU+pG82MhTY8pgjC1+/SH9odkJKpRkMPchDcbwgE0g/0Pt03UBu2cxVOCZ7XVviNWlJE6vWTCbySBPKWxVm7F5Zgym+lfh5YEmZKBo/8QqBfrF6LcDLo0R9Hd/6IWhAAAALkGfakUVLJ98ZVVIdYm9ubdI+6CZ/d+Ekcs8ct6CLr4YDU3lQFcmRYk+y35Vg6EAAAAYAZ+JdER/f4l2UkB1Zty9NLEh1ZInMX+BAAAALgGfi2pEf4ERwdI/XqLu8/7QUPdBBLv6noXG1Jbi0GH8S3h3poGiBbJjvmVilJAAAADUQZuQSahBbJlMCN/HqXjaYdAgn7QtItVYNomimA2qDY9ZCdac+T/CWtL/WPI+ea3ZhxpIBZzHmoE7ED6cmscuOpC4asmWP/ZY4D3ZofkgXRT0JIHQ5C+jN/NawCAZBlYJYeWSrnedRhv1Aq+Ezb8GwGkCKqT77Ol/D+JE8MiWDBfvfPqV5xJ6yu3+ZMsfZri2LtSEJZnT2C1JziQTBkpSnKjEt1PGYe1Ql8Dbmxq+Ajqsnbj/1Q9G0VHb1uODbJw883wCc+9y+57YbL48lccwqZcjBHAAAABhQZ+uRRUsn9O5+lrykd1W9x10yLwnbEdC7nPGYZCdulN4uKAuREGBPh4WH7u461q3qEJSnberdlJt1MwVHaJFwOjSN0QQTPgMWg4BaWK592DS+2DuO5Vpp9sLbJu4zonEwQAAAB4Bn810RH92+74LkTJauMv/KEBnzMgyQCjXgyIf5rEAAAA3AZ/PakR/19HKwpYZLtl81TUXlOY+KofzOQxBT8cp9helRzplXim5RNQlUaFCsHjzgwNANcnGkAAAAHBBm9RJqEFsmUwI3/vrPjghSM+3zdXqrzb6QyDQBh8wNH96zteG2dJP41q9qpbbpx6VfLx6pZ+hD0bxu4vInts9gfPMBlJkiYledG68AwhZomhLOtFb6Ji8an7vaJO4LtZXK+JkUFfz95QFkyWqJ2vBAAAASkGf8kUVLJ+ASbe6rPuX5kKJVKyfzTytHiClBBvNDUS9gyZ0gJuIACm6J4/MaPYlZAhSExx1G+JA3YMjZrZGz+SE8khQ7v8Q6ZsdAAAAJAGeEXREf4MMldf7IONNzwuYXTWVXiXI1jjTphyXCSJiDcl/JQAAAB4BnhNqRH9/nILJEpexGmpP0YxoFXxtFiqwczbTViAAAACUQZoVSahBbJlMCN/7xIFmvqKhxfQ7eu13hwX7Kl/0+GfUgZl6JYRtDtsk5rFajLUh7O9Y3Q/862jv9ZqodZkpyU2BPt5pwB+6u6zvaK5mzw79FkINXTazxAwa+hbdsXMsxsTU8kbKcxOzDPlpJOZmg+4BqbKzw0yHVYCbPCfiUsYQ1v/GbNQjREVqvxvErLVTetYAIAAAAHdBmjZJ4QpSZTAjf/vgOOJ6G+EORJB/i201ADDOh8gdy7TnqGzfuGC7fVt4IG74ck4fwuW5dPJSAv+ZSwOLk765EsOlxje0d6xKBDadZI86kJlJRifx+c16vpsob8hqrISDUct32W37mGuw8GN/sXV3vuo03HQdfQAAAEJBmlhJ4Q6JlMFNExv/++BEqKNR9RoJeFo2wJjCnq8GbSskuCaZk5J9oPynQ0n9SgOlT6cKPNsF1jNv7rf+cZkgbygAAAAlAZ53akR/dzl9hMKyWrjL+phdX+F0v8Q0kOKCEPls9CPUogjSggAAAF9BmnxJ4Q8mUwI3//vrPjhLnY7GuE3FgwVQBmmbcyRjHJitYGt3XJ7ttu2Rs1JgNU2hsMGCqQuf9g78vISwHq8b2HhwRAZXrups5jQ3v4TD8YIj3H3abrar6pujeVeAuQAAAFVBnppFETyffAdbc2ZCKuLu6g8BJqH4YE5YGtjEE6lGSZYtJ5hHsbl5fBMQmN9SaUIa5J0zft9yHmPwiRAJF0kKNTHMtbuNxCSscMTsy8L22gcicp5GAAAAPQGeuXREf4puFOCvsKrmMxdXS9O7u9OTUmOHpv3PBWQVCIk/VLpaFCq8y10BD+dshRrJ+UgPns8tyn3sUsEAAAArAZ67akR/inDB6JvALVyfkNmDHBAGn8T8vQbjDxek2OXAhfP9YxVeZBClgQAAAExBmr1JqEFomUwI3/vuVw4S52KAs2D87HYORWxCgFxVZ0fUYtB1vGD8zSig7GaR6PYdlvzo5m23PCr+yMcmEnzswwPvpxVY8irlJGnwAAAAQ0Ga3knhClJlMCN/++AchUq0qS3q2sujrVLxU7zB5z9GGi4Vr204wSsd+TAv/ZDLXIheOE/uGZxoJRecWc8xTdfXR0EAAABsQZr/SeEOiZTAjf/7w767baOIWn1VtSMPFg3+VI2un7hfAiMrNCPxiyraA7+7P/b4mHhjf9frw5Yf09s97+lHoLr+J40YfcfgrY6q0fAZ7vuspNsTPyaoTWu8OCda8dlFOKXttYcmZn0MYXuAAAAAdkGbAknhDyZTAjf/yghC9V+KM7GjZ3ifEeQQyURdNhHoTXnSTZGK6/FJfpI6X0y4F2LWLeuzT+DVqTb9zTTxzOILEvOnJySq5VGcY932fKbq7VpGGOVsYCTlw9TiZnaKVEnbuEjCDdLoHCPxhxVcfUPLCUueKwsAAAArQZ8gRRE8331Nr6PZBasOYw44aaDFj80ekC22ktCh57XyldfE7Ft3UzEJXwAAAC0Bn0FqRH93OX3cYVSoGPHXSBix14TW1QEOn7Ij37SbGA7W4x774czRQNnrfMQAAABuQZtGSahBaJlMCN/76z9RJhOzd1KAOnRwvW2BRZ7n6ANrAh48/h2t3g+1nVLUJRKqkD0Z9HeHn1likKz+6VR2gWMleErAnhTwmDW0YHrepMQTwNh+0/Fx5VBvqdRhjKSWumYEJreA7iMFKrvjlcAAAABbQZ9kRREsn3iZVH1hPWCBP0IcLolfj3IK062Db1vtBmUVEFxqhE2rclqJUf/kvOICwRgf+8cuGE3RT96eNk/u+JabuLcL9MSQv341jaGaaPUYvsPjkAuCYD6NMQAAAEUBn4N0RH+KbhTdyyowLli1rdXDi2OJpb2YMOsbPWx2VrALoxL353Cur+1DqAI4VUlPg/2b8RvB3K+6VuFp2qc1xr0ccmYAAABDAZ+FakR/hNfALxWamHwtC2+RNj8gRL9mIXf4wiYmsFGPVFdbv/eg0xWIw8ntIxBwu+sEvpuIcmRgof8iOS3heddgLwAAAJVBm4lJqEFsmUwI3/vuV1RItArTKe4aHtkbhVf/+Q+mhmXrHW48v4hipTzH46Ij37Wg9Db9jYlSjcF5jNQ5v1aWn65DMLo5D8XKRv2qTPmOytr/R4nJgxMjqpDiuGEXP91sGSFXXvDL5KhIcPu1Sv5rfiRKmk75SupVGhPOKnffZOhfkY7nU8WILOh3WPGj/TL7m/sJmQAAADJBn6dFFSzfgR7dVmJmGX8GHkqNZi3rEUsY+LD/L+bIktrWSlLH0rvcGVpVDuPV/0CpgQAAABoBn8hqRH+EUzCvSGusQvuCjCEcHYR+2+CswQAAAF5Bm8pJqEFsmUwI38oIQvVfijOY7Jqv/ZothAXNYBl7PM7Lu8cHbpXbUuZ1KACFF9u9XqIPlaA9xYZISQn7YtoCcDfTxu4CUZabiiBbVAJ69uAqIzwhmC7k2sQfHOcxAAAATUGb60nhClJlMCN/++AchU+aExJFTL9PPvLMNdHTufZ2ARRjtizDUCD7M3x68CS3Qr9dPeG8VhNRm6WRPsaNFzt/4aU270d35/LBN6vBAAAAU0GaDEnhDomUwI3/++AcgmDYEE7VgijEQg9aSxHaV7+OwetiVVzv9x0Bh9mAJqRJH8Dife3qBjyBpZMx0mPHKfKwSXf1DK2eIEOyj51/Ibe43gjwAAAAcEGaLUnhDyZTAjf/++s+OEudjsH36ASmbVNFHT+He3+dsKUY4kF1Fcb94L//2T/GGBOr9knUA4Fv6E+UdIJFqatHfWfL8/R/5FZzFcRTOFGht9DgyahVNF6gH94cGhO4EjBjUyOq5QvusqmUVB+YjvAAAABZQZpOSeEPJlMCN//77ldUSEVhRtuTqSuwn4ILAX/pGTDP9IL9t9LXL4gNSLMxC/6/2jK9Kp6HiBHIDwudiAEToisNqkYpE0nFGGBP2Mwq4VRrWiM0ZUYi+BcAAABRQZpvSeEPJlMCN//74ByE6AJUf3GJyJEuF8A3Qf1kQcJRR6jInzvSF1cy5rawxyX0nBOolFDiu817lzRYe7/SVVdiebEP++yydeb7tdJ3zxI4AAAAT0GakknhDyZTAjf/++BEqEiYmRvIgNxJIRznuGsny2nrjSbaGRQNjLW/f2j0oaCQe5/+54Ig6O1mLlpOJIun3/37i+STl9TMuA2Pp+8ndHAAAAA6QZ6wRRE834R1pCbTRWlD+k4+9DOjI2AHXza1raT35vBLoKirdVqNYgQnbX4DTT+ozWwyrExyRLfzhQAAABwBntFqRH9/nILJAq0W6TichzvXKAreK4xRz2yBAAAAXUGa00moQWiZTAjf/AOiKyEjuJtL+UDu/7/1ImEqkms4m590rIJHMH/lA7MFvspdRcAY8czYbgMQ6vpXSSw24CC9TXLj4CpfdwB8SVMwSgaDmTxDdOzD05PcebFnHwAAAHxBmvRJ4QpSZTAjf/vgOOK5ljYRuKx062IYgGdxv+oJsIKJXe5wzU05EYgy6CDVnrvgrKBM4IITLfdtFi9dSb58ZA2BYSlRsjpR6nFhFesKS1PpyqD9tlHbqy54P0NR0gcxVnq+CkYWBsgu1ogKQ0vpvuycY3R6AjqPjNvYAAAAR0GbFUnhDomUwI3/++BEqWWdyZKFzbFLbFx3xidN7yP1nsbxqbqZf3+mAe1LLS7fO/2lqY5lAlA3reqb9jagV1TR1CzRq2ngAAAAMEGbNknhDyZTAjf/++Aci1WcPt8e1pWqe9fATAcvlum6kTqhjA5k6uHd3lKWBjdgiQAAAGNBm1pJ4Q8mUwI3//vgOOEudjpLQF0T5/HpcMVhuo3UrgFNwvqy0NaDRsYBORhen300dcdaxgQAhyHLkqqY21QvhsGVyf926QHjO1dTQI9IklTM6dFWv11wpl628oyD4WUbTWgAAAA+QZ94RRE8n3igS8fsiwsG3BGH71JSuKL8DGQXeoxKKRzUOLM36KStM+eXk1e/WYIgappa2TFwUOXYGRkunvAAAAA0AZ+XdER/im4U4Klw8nZaidOEir35twZNCptJt12OpowRIVZ2i141MkhqsNTXZF3ZZSs7KwAAACwBn5lqRH+E1CrjX482235Jp4YGEincKNRQy+bYAzUsHzuPxfifmdaLZN0gwAAAAIFBm51JqEFomUwI3/vrPjiuZCYQyR8N4L4ZpcVyVm5GOzBlZrs3WY0M36PIs4HlCuJLEY14KxyS8bo4xIH/xQxifCHSUw0kKwQds/H/qM/Bfv/37nK8BNWssBs2UpmS7HvInr7RQkATNzNBxyIYrpPs2K5w38HW0rEzP+G0ztOo83cAAAAqQZ+7RREs34TtSIINqX8yEwwllxjEWv5dp22DsRrlE9pSWyEg90j3GhvBAAAAHgGf3GpEf4Sa2x3nqQ9dXv2XRummgfaBLvV0z7jPcAAAAHRBm8BJqEFsmUwI3/vgOOK5ljYt21R9BjygQzFkvYgmlW7E/rzUmB3aa1R/v/sRkWdXCOgIKuKUfzXiJlEyDC3eROqOf9d/4jH4mwkcdJ8kwHn1NDcqEmz68Y5bf46nu99czcvNzhvDbgh0CK9lJoeiHSXoTwAAAFJBn/5FFSzffU3evezXtC9/ZQKRQHpR1ZwveCVG8cv+8KloL2JUK/o4OcsKclrttQWDgV9mNWZ5P+28tt6XuycTD5aBLcSOUtgKm4G46iDPAKbMAAAAMAGeH2pEf3c5fdxhWS1eZ3/p8a5pR6sqQeE6lWFP5z0K+aYlfBg11Rzf+F08p1cQ5QAAAJdBmgRJqEFsmUwI3/vrPjhLnY6O54bcEV++H1QyWusE03qBnqIsvXHnWussWXieHNrrthosgcAU0Y0F9hopOWEBv0JGXS1m3o8PuW8QsFb6/rPNSkXUkO/oFNfmFTp6wgcdbxk7t64LF//YEwizj60QETaAs6xDcg6lk3U1uH0HDMl4ih5VDERKo2WeNr7DJ6PAseHZfNpBAAAAZ0GeIkUVLJ+A5qo/ZEObD72LxnQYIuEYf1ytAEWZ9RqVWhpP9byyoVJnVZyVLVwDIFybDJ9gG7viY4976Z8PP75RopC/RDB3dZ0wq5zC/a7Qp87846DKhZFGiIsYsKK9l8waIPbfwkgAAABDAZ5BdER/gwyVRD3qoTg9Jl28dVUyLW9E4j2mbSLrdlU2wMBPhJBWiaLiTejDzYOqwm38CZG2crdvRPAPjeGIWaTVVgAAADoBnkNqRH+E1CsOUDlTRB+2knMt/XIVd1w2Lu5z2obTpdRmkuyfrXpGQRJ2yHOSmMjOiyrLDlpwJ3dZAAAAvUGaSEmoQWyZTAjf+8PysOw7IqechdCyDapN16veOkycd3tl/lOY643J4iTWQW0XQRO1oHAqDTs60fM41KeAMdMg5qceLz0xHG9o7K7+QDwexR/ef2KxGk+JPpEkkQPooks4C/flF4jhBo+xQ1NvAor/CROg6sAGmqtYsKTo2RP869Xll09Z9g0/aqKgXet/XAH1/pPh9Idp9o5m+LTjIUlHJDcUkw+f1s1KFZPXPnjINmfl4g7HW8mSl7TCcAAAADpBnmZFFSyffGVVR+zevXVKzm2vgkNswDKTj5fZQrLaGvFrUNGZo7oe3bojV6WjyvmGIqww/DM1Nn65AAAAKAGehXREf3+JdlJAdWrJHHXgk96zCjVyoQkJf+XToMVT4UbfSa64B0gAAABCAZ6HakR/iLG7Hu2qMv5WSPgOGOgEf49LPHtgU6+/FKGUK0Uj2RSz8n/2jS/nuH+a5JnCsk379xSbYaLKWkzHqqhJAAAAaUGaikmoQWyZTBRMb/vgHI2b7Fd6Ws9h26Qd1WAcJFbDOuHQlDOeOVSDtthUTx2BmQWjIawleP3YJiZxUceVqUKAtVJO3MmSt4itsmeOXkS96GA/02xr8e3XWkoQ45QprFdd1gNT1R5NgQAAACEBnqlqRH93OX3cYVSoGPHXSBIm8PYCgiS/lOkz+kEEcdUAAAB/QZquSeEKUmUwI3/74VOAVE8TO2duqVlti+FmeZePWOPchj2ByBQCHzxIwA7oRji+KqKa5n4Q6OpasFlu8pZOpT7HW7iwn891D4Ifhvw+ecPpsngXgO/F02GWm9i2sspJD5g5fSDd6gF69m+tkP3mOA8J0Iqx5r9Ybt3wCsRgDQAAAFVBnsxFNEyfgOol4/dcuaO+lc5oqdzUsBCEsAWXCEd9LQnFt83M2f4mvOGuJerHkxVQwxAPUafNM6hjcUDccUEXrpfwnpDY6bBBuWi3an7M7S8Zz5MyAAAAPAGe63REf4MOYBhgWGIje2c5UIBG+OMYNcDzO/HP5o+6xdUgK/xp8ClpOu7ZqMK+GL9wf1yEuViw3Jz5MAAAAC4Bnu1qRH+E1CqJbRfRlHBZLSjlMfrSL3+KsaKQb79wkzKWLcwNd6o48CwLbJZPAAAAuUGa8kmoQWiZTAjf/AOyX9R0mJ9SwB9iksEB7ZH/nkIpajkfaK17kN8kNdFdkevjT0goBYSx1ad7E1vuUhe0fKOO67ud4jxPyVzw9UyNsylkIBxchX2bMrVXnfEsnPA9mDRN8mbHacVFgb2ToC0gaK1kxLRZDbbQ7G9HcVMynbMoma6dSOovdmZyIzwlpap73j+GNAutYydKa0z8rYc9VMPkEA+pbgbAGRXIRlXA4YjOfmY7qK4sB46wAAAAR0GfEEURLJ99nGGl1msrnfy1MSO2So7YfJXmhMb4mHw7D/nEENpn/g8dYIuHWptpGty383qQsQOmaDb1Sr7SyZujuEtrU93EAAAARwGfL3REf4BsK1ffX11Zae3WZvfEuHzoFi2UXR12VR2woUFdnWnJRyNJzEMMur5pB08F12TS5Kqh0g6hXx3s+h9IciwahuVBAAAAQwGfMWpEf4ERwdJRbH64kx+NQysevRT1D2SC34ECzZV45BjVHUr0rPEmr61jCfTNQyQZxfD2v72JxfWvi2EhWwvZnoEAAABvQZs2SahBbJlMCN/KCEL1X5n/pbtseoLLd2WSFVdinV73wKbPDDMQGOJ7HQlNf01lmQlQxpN4Xld3ENl7/XvF+385IN+ANQzbxa9MY8ghxFpqW0UYwQhQSaaY/9FtnQOc0V/gREgSgiRtrKFqyex9AAAAN0GfVEUVLJ97HUrykd1VJNIYMgfZauRWHu7JgJIJWyjt3yOsBMa11eyBZgykRdmoIPb4fmsTH6AAAAAxAZ9zdER/dvu+WAiUqATXXSBiTXTh69DK9T96yYRRTZn7NwidgRAUaKpqXh1LkSJxnAAAAEYBn3VqRH+KbhTgrp8WYLnb/OExBB1QfQOTFVBGW40ihRipF6j5TyV0K31uIo2+OzX6N9DLAIAws+5y7pu8ISS/6hJ47suBAAAAk0GbekmoQWyZTAjf++s+OK5ljXoFiiU1XPjCnQxFyu7PryXdiZtWx/wW4N5zId9wNpopmCe9CPMLQBcd5pQEY3kNzWS2+3CovTQy4TUQSeEemdF5upIcffV6vhYfM7AT24IWEWArCyFFh7pIe8+zc8VPXEM6rc+7xdmneJ3szGOTGBGxzvpFxlTTTGKv3pJIbNGDAwAAAG9Bn5hFFSyfgOahWDUeRVYeX65PXWlnysfqmajig5ypCQFzzoXCSRAp3zgjG2NG7gxdp73zP4B6ldlcHso4oxhpX24XpT/qfkQR7KKVU4NNmqjCtXvpDalQ/00eIGvukzaHkWUGHTKp/QtbKsPGwDAAAABNAZ+3dER/gwyV1/tYOwEqTd+Pf0jD+/NLCVKy6ez+vl1KUGnCmGBaIywcEmAVF6zlPqshdQhK5FM6KDE97ysl/l1xLpPEqrALZu0g6REAAAAtAZ+5akR/f5yCyQKtFiypduyg6b/z69WowFlASrlNd/l1Q8J8ACA56sgTjZmAAAAAuUGbvkmoQWyZTAjf+/4BIhQKAXMPK6lD76MX584FxB3ojNVi5Q+urU2Vg7DDnw5TYhPAcAP+JpuOWYUZvbWvqxIf4+zhikw5W4uf+nNx3QUlfGE+qQdeKFt6AVvU+gOqNlNIKHEl5X/vEOA8B/wsTOAL7//f/G5Ez96Lbtif/QxRwf+GV/m0IscN9zKokP5tFqEx9jubCAyMXQbzIGza/AdoAcgWAzpGkkEFPVZHIn262BXUMjV6hgYRAAAAZUGf3EUVLJ94mZXaX35boOk57/Jas+YqlIwzjlszmhvu1CGdNLBRaqzGIA8eJ3iguthSW3p5XKu2MOj5ceJGlYERu0+psmSaLOFXONszzF1sh6Wy11Cr5kB2FRpfgC7N0D6P4v4lAAAAOgGf+3REf4MOYBiXa4/L66IgzTpMUW3pLiTwVe2YiMlu1Od+TX4Ego0CvBosfBkHjMBfsmQviRFuwkgAAAA1AZ/9akR/dzl93GFUqBjx10gYYvCa+gIMfz7WTXtZIsyH48Q3lDcTKBK+Eqf3cTexpuqr75EAAACKQZviSahBbJlMCN/76z44rbBJLPdsY55JpbvEurvbYmuIRcWV/rJn2ygwgU9CgdRF/043si4Zr5zKodFD1SgFp7DQqYD6tTX8dwBVP474YclgRLTJThjJ7uLHj6V/CJc4In7m8OgftrGMn+Zru0NctFHKqM1bDUTtaFM/qXypCDHtWSqgegcA8LCAAAAARkGeAEUVLJ+A6iXj9eUXdhdZVK7LV1HMdFPLHxlyPlUiuab+bMNEm5whCn7SWav9ox8Nvdiza9s3scrTm/Si0xxwSXwP4ikAAABFAZ4/dER/im4U4K6eWIXMANXzPh52kzBolbrDVKhXJMRmKIXiXvz9VXDYi5ff6Qo0gHFiqU6rW9yf7KoXUhBe4Rxlo8WvAAAAOQGeIWpEf4TW9RzxeaEIBMuIawIw7r73W/EwbmR42EU8aJ8SmWMIch19eiKXWJLO+HqvXpyEDPqD4AAAAI9BmiZJqEFsmUwI3/wj5K7y3raNV+lN1e0Gmx8L5RHw0IgYfLeHCZ5xAiPYSqSZkRczptAL0OS/Eadl7UGYnx9tFSTMq/DW1+S0JjpUGH2BBi2JBcBp9gBaH1YFApjEEAfskKGqdW6ZtsgBB02//5EzIeLDw0zMvQr1qNKsH9QHaWrVtuaIvnOOjm3+a/MpYgAAAEFBnkRFFSyffGXEvIdZdMwjUHspZ5lmIjrKqSiPV9+51l3FYjTQ30SJlU6anVjo4RznACIj/57d3zRCAGeOX36D+QAAADEBnmN0RH9/iZVJcpUEaVJ+D1Qn1J1MON4UlEEFw/V4SMOqbANXYVK9+DIT1C0sj9H8AAAAQgGeZWpEf4ERwdK7XH6zbP+0KsgLE2aJGY+DvwxwgrIkeYD20yzoRSouFKQ3gK9JnSL/2w7s61b8noAFMeJsnjdrGQAAAGhBmmpJqEFsmUwI3/vgPUUrmLiGAWYzwddaBPlC1tTCKPJyTpF5x5Urt5EikjDg/Y2Ch+L/G5axssucn+L/XqrCvrKKAbuvr0oM1XhQMnNPkrUx8G7rkqNLrOkC/mXdj9DtDsSuu7cqxAAAAEJBnohFFSyfex1K7S/556pvzaNwUuIyDg22DfvrCUNWtNpRRmvgNB5PJxjDd7uUat/w3BVRmbf9xTYsRqKfuotnayEAAAAyAZ6ndER/dv62TuMKpU1pPPlXna5sbaR+8FULtAJqz/eyyWfjyJNab+zrFLDXgmE1X7sAAABCAZ6pakR/im4U1zmIP3cv/+aqmtN99AtaR/ePPWvf3Kp9eqRfWQYoVZuV8xVd50yhmTsQSs9JdwKpAboLwg3uGj6dAAABHUGarkmoQWyZTAjfzkcmRUPipKf5HPWzd5mVSXOkf2wveByGyczca4WLzC6dJy1m5B6Il/iv02CYxAuC0JMLjf4eVD09cAC4VuqW1rQdH0jYBt9fqu8NkdwlacAwc4MltNuRgUY0PZqFq9cUnMVBRelxNywmlynkI02yNJIx1oubSc9FYoxYy1GOdo4XWxNYblENWmaGnuCBIJRx19vtPp+h4vvQVBXbfsx3OZWx9A4nT7Y56PG6CxBp9NdcYCm6+YEXmiN995AfuUddw5pT8IxHOckLKvtrExmNawZ4j/qtgJQASVXghncWygED1xlvOUMsv400IbfHtGjY+xlxEauu3NbsBWE2bmc7OC5AYARFLBZS1IMRB42qDSEAMQAAAFVBnsxFFSyfzWO6eV41HkVXWetrTujpSLSxsyet/K+xuI2mJCQpfU8AiK58PkSxjDCtX3agsIW3wcJvwWe2BOb+rXnGIMtnblWi8WfxWyNGdCgOJmIQAAAAYAGe63REf83T//106i/j9g6B7W/+9pd27v4NIhMKDiG077SiuGlpjPOE7m0C1ERXhoTXOKeyMZV/Cqhwx0+nXZkTcQvFvUXkgMZhV3TQViM0RySsmTCOTLoKGrX8MDQLgAAAACMBnu1qRH9/nILJGk/YUBDweJZSU0TGSR6vSvqHjaBT1vIf4QAAAN9BmvJJqEFsmUwI38erBKuz0hHXKlCdZjiQgCATySoqsm6ltfbP7rAFGFZ822Yw4uKROyUID51JSZ2TSrmysf59dCK0wX5bIkLdjmpkPTrcM+4OA3dv70JvEV+piUKzWy4rgthcTqUOTCLotioch4OuJRHZVarI23FHHeuIXdoOx8okOazHENHgaw6X59OheNrJiEvne4/W4SXKVQV9oGUnCUDqrTQ4hxMsBqxs+I3H3zkR1F3Wqbxyocy4mEvJ67y6N0vizCEqYxGvHdXCKAGvbYr0oZ5UmemNFdqa/+jOAAAAO0GfEEUVLJ/NSlgemZ2/i+2HE/tCoU5+SvCMe9vvnjLZqsDRHgVJCIy7oHSX0WQfEACM+t3MwE6ru1+AAAAAKAGfL3REf9fXXw1RMHznnW8D6KXJTU0SNFkICnZ+oJHgBulpiD/kF6UAAAAcAZ8xakR/dzl93GFUfKEWvCRfKO/B6X2IPGTMgQAAAIpBmzZJqEFsmUwI3/vgOOCEetSB49qntGrM9qobba9AH1IQdtyfiRpDEBqUBn30cNOTlel+ROcamSYWnUPachBf2dH/7blMdcvqBzNRPhHUnxCwzkrNRlC/5GgwjKfk/a2B/o7qVamrbXYVkHxVa8bR73b7AQGyRnIHvBjdQa2Vyk2zDB2JLzShEC8AAABbQZ9URRUsn3igS8fWMzYXk9Mq+MDqIxXVuhKB5vQkgSjtJT6JSpHmpOK9dQXV8FSdFrMVLn4/C1TOlGot8odWKsZCHVCzXaMLFwpRnRtzffgvkqf4iecFvElrWAAAADYBn3N0RH+KbhTgr6xHMY/2QndZPRpMsv5JSKrqaGIUN02aOyCA/tVS67KULKPzvQ3M2ZSKp4oAAAArAZ91akR/hNQq4woHVjdAiYoVEk1dl3y4BaGjLWpTqLq69j/cZmOzoJOR/wAAAEpBm3hJqEFsmUwUTG/77lcOCFJBa/WYM+C1k6KBz1fXTj2WWIJN3lFiTnpO7wj8DQA/i7cHqWnaZRx1Bw8Vic6inCrCA4O4XOxsvgAAABsBn5dqRH9/nILJGie4fD5UaXXmEKriUh+FqfgAAACDQZubSeEKUmUwI3/7xIFmvHCDZ8t12akwl+l9pgJioglFo/qTXNJYNkUFQp4qDJhzpCYfFAIWrXQ9O8aP+CIhDK2kFEhLIikaKbSUWzNbi6H4UrqrGmAlkZo5re9RQGhCccg1Jov/2GxYYpyEq9L39zvDE2Io/i/vPk4P9W3qT6jP+oEAAAAuQZ+5RTRM331UtYkDXGvB3dApBSvf+bMEW2P5xUOFFwLXVEf/avPuP2djgg2IoAAAAC4Bn9pqRH+ALNfijDv+qEgNt2DwKYhX+aIxYXAHwEiNC9Yk6O2btseoyQdM+47xAAAAUUGb3EmoQWiZTAjf++A44T2yC1yJtTnpkDG02asj/rdoy39ht3+VB//uMshD3UPhmeAzj9IjmzSBMZH7Bn25T3xc51+ObVJBTlLroFR0mja7ywAAAG9Bm/5J4QpSZTBREsb/++5XDhI2sYlYDT94fX90cHdUWtuMJHfC8cCICr02MBGpCutBAL6UkDKLLy7zPKWJhLoLICnT12OP5gktaUg7lCh5ilnYCb2P2H4H+nX+3GiUhUVoSBX2x9WbuvhYGX30W4AAAAA3AZ4dakR/im4U3cskkOWaidJ3UzcdYIuqiRFaqXPiJJRWssAOr87hWTOg/jtyd7zuRXHH9E6yDQAAAEpBmgJJ4Q6JlMCN//vrPjghSQWWzwVzm0Fb69aOUyH2uYwkl7uzi3xrhyoi3bG6RP0a6qHuC4sp6e7D/Sqm88josE6dK/dzf/C3/gAAADVBniBFFTyfgOahWDclg+s+Jkr9gf7ZOS+stAk+S2N/q5u0Nk+SxhbPMvBKtRp9lBeEwzIS0QAAACoBnl90RH+DDJXX+ypzivEOWRG7zcM5EY2ZDjmRy2dsh3v56A+eAR5ct6EAAAAhAZ5BakR/f5yCyRJTQ0pmRtmpfeSIf6wjFYJ0wy9Ux4VAAAAAX0GaQ0moQWiZTAjf+8Pt8hO4AJKfYhcjA9rjdnAwFf7QModT7PyzpMUK+CS0HRfcsjc2DBko1hnWHk2EmhC6FeTs1eU414OhygKFL6xPmfel0xhg6rysjXFWdjRF1jxEAAAAi0GaZEnhClJlMCN/++A9RIVzu3Ghrbc0Z62jyac4b1qAl3BBQ30vUIr9JNcG3B/7thUl3gfcZJEcDA9/o5HLdtRPY8fZ+i/Rg0HkiFZFC1tsDbnRuKWHh6GA5fZne6xDlbQRKVZqTruh+xSV166g8XuXlsnFkjpjI3LHla8Lz4NACVbazHfhnbOqwIEAAABwQZqISeEOiZTAjf/76z44S52OkshAe2P0PHhV4Dei7KErq3dgxYM8IFzbxvZbXSizGnKiEKPygFN4FMdrwWYARbLCT4h4G4eB1lTgxQGMp7418j0PFqDnRSjlz6lPtttNkfg1a3moBhbWbgd7fml1UAAAADBBnqZFETyfex1K7S+Yzn6tUA1X2wuv7z+sgh4e64i0lTn5rJfpZRmkJsUMXmXEYcMAAAAdAZ7FdER/dvu+WAiUqBj0v7NQWvwATKba0oofnS8AAABJAZ7HakR/hNQqlgeCR5WmEnTtgY9GoEa2D1hVoMAPb44LseFckgc9p/mRB2QhlABhbflpFWCIhI1jRMwGxXP5gv3zZcieyZt8sQAAAJhBmsxJqEFomUwI3/vrPjghSM6GitLPhfT3QjDIXqDR2QygB3rhZ0Opwh4x7j88E+9pDmthzmQVl93MWqAlaOPoKrUNvWPRoPXTCZXTTpy36yAiQiFXV1QqO3LKvj/ZL3s4SSgWTn43YwKoppbw9S0sfviVAvMSi0iDeQvD4QNHMF6CPVgieFql05KklKluC94yEGBUepu8pQAAAEdBnupFESyfgEm3uqz7mAS3M59dNbzoYEhAUaCO6CVGZpltV3a4ulLjDLMTNTXCXVdm075xx+6WQowma6lWj8Am7X3wHh6YdQAAADUBnwl0RH+DDJXX+AaXoY5qOgBknY4s2h3Y14HLG4oExzfORWc7aR19WMqT9yojLIGmtjRfQQAAAC0BnwtqRH9/nILJA77E0uyjL+OYRlC9eMwkZdFHKSM8oLegYjG2Mku8ZesfQpAAAACeQZsQSahBbJlMCN/7xIFmvqIPoV43rFJvWX4ODfVUK4ly5WYitO5XIzfl1xd8KVyzDSy+r0BDRR//Emv+IrjHGw2ajYgXdcYjXyNQqWjfnNEm+mjBDaTGTxJzcJx7tSgOp285AZ65BM2cdcwhRN/7j+yRSrtO21pzFo5NGhAw7XafqedlhzTYP+QC5ZhN/egGmszEX1+2B4unJ0IrXIAAAABBQZ8uRRUsn3iZldpf/xqA7EEaL/5ZgaVWGAVcD315eHZ0QAYXXrP2jWx1qG9PMHNp4FYUm6nfg0+ci6cjMQFda8EAAAA0AZ9NdER/gw5gGijNiQz/5kp5JhG8oUICnRs7uFapasLfsMKPh5eAEO7oRpw9QfNO9PihmgAAACEBn09qRH93OX2EwqlQCa65akqIwjXCrER9aHr25SKikoAAAABtQZtRSahBbJlMCN/76z9REFwH6H8EqFmA4qxAVHo6OWo4ZywYj7NSM+cdJHCwf4LAdrJSFZD3glZ52eoTwrY1x19p1bxCO7ICdQGhzEiQ96Xc9Y2SETJOD2Vn5GWGsL7J+gJXntXk/X8jEs5a0QAAAGNBm3JJ4QpSZTAjf/vuVw4IT+j8dsgkolcVXPJAv/DV7tMKPOhKW0Eb6GkMP3S2+4Qotkl28worhzeCxOdJSjc6Io+xwzEknX0Ktf9XBtxbX0nSYFhFqBO/1e1O40UYGwmBSjkAAABiQZuTSeEOiZTAjf/74DjggZKH2CCRdBnEYQWZN54htqxWZjo9aIG1TMzzTiRe4WNKcKB8WgbTaefOnNjdIP2vt10+UjnumK7PI2s37KTJhbyphwMl0U5kDcZg1Yv7FQFx5mEAAAB9QZu0SeEPJlMCN//77lcOB8tM/mnq1/U/2ketb7E4WfIqxx2WSC2Pl4TSDg8o3v0A5Vjy93uLHMTZs+BBdk8ByTQwMAo1JaoqmZVq8af7jmwnevUWJQ1A4a64zgNIiFZUXmvCBDFKozh+l/Zy9LkBa4jKOr70skQOcFAQ9YAAAABGQZvVSeEPJlMCN//760EqEjYhsbGdgimO8QcnqaVCP9qbOEFZrTmsIf2JpB4TSeaxpXl3Ulw1dscvV1EvErHcSHNSQtr58AAAADhBm/ZJ4Q8mUwI3//vgOOCLONnHJ8xun5qHA3ZmgreRBzw7h4GYMtLDZh/dmKRwhM+8ymTIp4LzwQAAAF5BmhdJ4Q8mUwI3//vBkESMggDx1VvVtLMss+aF3SXvKnVrCTWWcSrFJBN0jS1tZHuusb/RraxhnYjj1+/yAsYA0rgXmW+bCEUISclYRiHSnHWCOd1ZBmdTg6KC1ZbAAAAAZkGaOEnhDyZTAjf/++A44S5zbJqFTmyU0bct5evsZEZnwCITG77hmWTV3ATPGRlTvfx5fU+QbyuvFKr5KrASAROTbHME0sqsyGlQ6eQe5/+7Q4wQpIcxz1EsaApgkTsdRh+yx8JzHAAAAD5BmllJ4Q8mUwI3//vgHIVWfpo379KAL+bC1vac0rQQmCNEFSc2ChJ69qz7GJwQcygSggLTY2XDCMS/LkX2TwAAADpBmnpJ4Q8mUwI3//vgHIVWxnyGmcsFcA9oxoBl/brkT10aH4nVw83GGNtypUTAP/suuF7gk3SK4d4YAAAAXkGam0nhDyZTAjf/++s+OEudjscNzZV5CCe/aEyxqBVLFh5VRKVEDzmEB9Dk7ScQ40u1uCiy9OuUrKIzx+gIl6P0NG1LHJzJ/qCDx1gum2N47t9UWnFNio6Aoe0ovYEAAABcQZq8SeEPJlMCN//77lcOIMTUMrHZOr0gSKOJUf3LozHu8r1h9EcXvqCO3mOg+d/VC0z5ZYdgFjL3CZniji5n4XXibsuJI5k3acNamysNf4eKfIu8h6JjblX+LUEAAABhQZrASeEPJlMCN//77lcOK5kGh6VLT1xLqj5X9bM3RBN5M3gOtGKjSsWOVbXT9w5waJDJuT3NaaL21k2UzbmVyf92Io8RDVoRP2XUDjiYmvZj67mQ91VySU1YVL5C9Sw30gAAAEJBnv5FETyfgOahWDUeRv7uaiJL9EdNC3IbVzoaH4e6aSo1n7bs0FRXQ4JWRcXvz38ocIB8WqDlMJ9G3aBO6MMl5bkAAAAmAZ8ddER/hRLKDHq0j4d47DM5gwud4/DkjlPlVFYGld+N18XNz8AAAAAlAZ8fakR/f5yCyRJHv89kYcozAl50ypnHlpKXwY51Kj7W0wQ1HwAAAHJBmwNJqEFomUwI3/vgOOIMTVxkqPJRLQpQUqKkxcew6gejhK9dEgRRiWKoB/5wIFCYAVMLtY0lhftvZnZTOruxFhOlh/7mqhOygkTIBTqYT8F+//eGeuYR96Ww48vtz7cCy86ywFxOdvLxYRAQgapIAWkAAAAwQZ8hRREs331UtYjmVOEob1vuAVgoDzC2PbDUkYLtXfYTEoiAWJS2xEv0kGt1b0TAAAAAOQGfQmpEf3/pQAeIEEpwh7Q6XdldFbTTk0g56u0PL+bxWoNXHexxjGmHvNW4Js/aa9dUwMqYLiqLmAAAAIRBm0dJqEFsmUwI3/vrQSpSocjmBEaZ+guniyl7iU1d1HS6ptlCPKacTBMZFdsmw3nQ/2SNFSQYbakk7FMBdPhxaD1VZVNg6EmmIShve7vOeawjwMsQ6qn7zs+OxVipM4voEUAHDO+9AdjYKvS5OsP1RkpHK5H8rsOaGDiUxa546nHEZkEAAABXQZ9lRRUsn4DmoUnrvS2j5/oh09N4LvXKgJJsJGF3dTgk9NBZHHzkIto5QN/1CwwGkGBV6UhxpaON36OliiklObJaIH+XV0IYQQBpF7sgk1BWy0nUlXZ4AAAANQGfhHREf4R4tlBciUsqFFuxA8UwpTFd6S8jYZCqRb3ZWGJYP1ecFzDRMlHObUR8+QdZwRJhAAAAUQGfhmpEf4n36qqIKnWUd9Njiw0QSxoVB/oX/yBhz8h1/xIEQ55PMKdy01KoAx+nvEYqAmAd9ad3OBTPQ9pOmJCNVk/X7WrPcvXMd+GCzMVi+AAAAM1Bm4tJqEFsmUwI3/vhUdkOqEdG62jueiHqqB7YeLSYauyhCAzqdiH19ZLvxGPaZ6tgKbJdVHcaPgr406CugQLvd3FPbtzm5AxeSWxsCek54+8ViKC5KSPZpGstBmWc1F5Fz8YmM3pXrKfxiE5SEFNGyFV/qkcpgvPYbfMJYpPyMe+Mok3YdDm7tJqwEmFuxpR/c7HOjEn/q78jrYgJECTdXUMmYN4j5En5F/wwQM8kA2ap5beuNPWS5eYz2P2j4BU9tNJrSFla6y+1AItxAAAAO0GfqUUVLJ97HUqt10psKM2AYDKr6M37fOlTxCJJ2VrMROQNURvJEmBHrG4u6Jh8f7q7B191ZIjfgM19AAAAQAGfyHREf4BsKlO/qXBAryvwFqs5i0G98z4g8HSYOeOjUiUe2uLVZBQrxdTZrbxwSpFGxJ6EKPGOOvOBc2c9rsEAAAA2AZ/KakR/hIfoLL7/BzcvpolKljM37OToCQ/2FbQB5A84l3GTkJR4Rl+5ZoIT+HkMEnzHYh3BAAAAg0GbzkmoQWyZTAjf++A44rmWNhGWCL72npbHi2D0rm32fIO2vw38MGXgB7/2AwMnd/lvD81EFlh6sV3nHxYoB0OZtYyq9r2rygD71P2o6DBvoGCJwkGKOl4eFiEYvK83H9emCgm3wnANdrFoijHM4zU0e2RMJh7A5zMM4OT2vz4r/00sAAAAS0Gf7EUVLN99Ta+j2QwAR62o44atyjblHXj8qcvZZupH136TPo8NbBq8vABIfCMtlgbq+SEavzCvqU9ljoozFu3hOJeYHrkwhPV+IAAAADEBng1qRH93OX2EwrJavM7/2lua/dEitfWcARv61hjXJUf8VE2rGPeKPTvdA5ULh6n5AAAAnUGaEkmoQWyZTAjf++s+OK5ljSUOA3TFai4RJE/jHgw0W4P9CLmTx1ROtoE1Ct/u5aZKpqL5UFk907WKqU+Fq87HMy5CDtuyG1GYY7+oe/qacQmKPQggXqXzfVJZxonCiuPjr8IvShkyg9jrNQn4JR5rA4cMrT0Zg58ZdIL2DxBO+Ztm7EZbuwpX6hiLDW+Fu7jZRlkzHZVGcLW+A4AAAABoQZ4wRRUsn37EZ8k4yrEJ3gXZR38ckSC4E/G9JeY6nvqw8cni+XpetBPoOPSLCtYJ2bS4fCf8qRBknduGn/zENrDlaOySIiDPyNSjlsUVOY3Be7kM82KFFJZ4p6pqGLRuc8L/WHBJbyAAAABOAZ5PdER/gw5gGGHvGKdek4JOkgzZuCPm5S4Jttsi7GrvSbpquZEjT+l8DryKb1bKgXccBTGldbxppBvtraeOctJrix9CJg+uODiF1InjAAAAPwGeUWpEf4TUKvadqH/skJdonVCLe4CSymbvWcZPcHMhqYKXxCXq7uBXK/nUvzCA7bRhXe4Org/oKbRblHoZbwAAAIxBmlRJqEFsmUwUTG/8IvJDIWLLGwa/ZHF/KLUN5/pBVYdMU+ZmVpOULfEv18kUVSeOAd3niURef0hBEmg6OBZ0i0XF+KXTRz2y9KRZZglHmsDhwytPRmDnxl0gvYO++Jtbp40k8mBlBmde7Jtob6nZkW75fdtnP5nQj7Bc2x52eeW7AOXILyoBAXy3wwAAACIBnnNqRH9/nILJCQH8el53caksCb78m2K8PektcjdV0nLgAAAAtUGaeEnhClJlMCN/++FTgKPVy25+n5usIXnTztuq86gMcZAv//lXvQMR4OFxAH4HHTYRKq4DSRdcbN2fFwL+DmtKdirEql7t068169vI3EBCCne05YqLPOllLz90IOSrfUDiyyb5NfrDwLg+wOh7qHwq2BSxmHlV8XQlZWmA7Q1K+zPZSUHgjhs8Ax1n9XnxVLQtfRLLyA/dBKYlK9a5jC+wfdcMLbf12+YuD9YhzRWeXQP6UjIAAABLQZ6WRTRMn3iZlRG9wMbG6rulz+i3xcFmicnkVfPi3vI9sl/efAzv78X6GnDlL27msiMoZuydUOS3pIwXyOEjvXuFAt/zUCd6+qXJAAAAOQGetXREf4MOYBiYRsM/4Wn9J79PPYoQFjb6GuY8jqTLyI2PeCs+yCv34RJsSxBkwViDKVa1l5B7DgAAAC8BnrdqRH93OX2EwrJavM7/2lqioT5DC3FIJv7vm//z26CikO4tVaLBSjpb/LkdYAAAAOVBmrxJqEFomUwI38n9ga29JPKIOBHSx+YsFh8iUW27g9I27uSlq/OZWnSYhIZM7ODNk5Eyw1T9OF5rr+J2dy4gl+9RDUCQuNU7dxcC1S9rJoB2wj0ZCKt9fx0BJCscbIz9DE60T8UJcZW4IXI5PQ04rPvV3fOyHU1DTpW5sLAkcKbZY9J5we4BPqx7YVed7ploglm/l2SkYkxjUmxsPzRVLOnRXmV9hlstOpUJ/aOYKohirbc2wOqCvDKXJPhu7uw6sdlwNg7ct2PF60e8zKjuPUQcgh0U4w3KxPWl+MIuQ3Ce8JUHAAAAS0Ge2kURLJ/NY7qFXvgUaFUOZCNBf1AjC4hZbIHjRp6oas+6Qqk70vXUDe5rKS9H375z+RHnXGdynCNUDFHftXvoHTlj4VbCTJvnaAAAAGoBnvl0RH/T29sxF44A/5fnAIOAgxy45Xst+FBoNPddunOVgbfDRXfDZMjQfcvkgzJf51llDHGEOPZdOqXjlDVUZbZjfZl0loJS2YOt48R/pzXYFOhFufNJzmM0UhgMuQR/zyOciCsNb0DBAAAAKAGe+2pEf4TXwDNrCamp/ew9bCI67/4OVldsPjdTTI086/CfMW5DjVUAAACXQZrgSahBbJlMCN/5KYRFpdsBNJCPfk3shMTTbZ/+m1g3BC6EDF3hvpzFAan1Nwk0GhKdlNqVrfEyKn/lndpE6h/0uKfhJLNPRapdj7uQNAbjcketIZV80xhYt/ITlWMCxK8VuHPNA3X0fatIH0MAuSlD9/rmnTCGX1rX+/0XG+Uq7kJlIC2ViD50CtLv279/vtSKlXHogAAAADNBnx5FFSyffGVVR+wTASuTnuC0n9aZFyEaIux3eFAxaKRoR6o1P4kyonm7L6VYQyuDgUEAAAAeAZ89dER/f4l2UkCrRbpOuJ2cpJvutsxXLqDlkAIuAAAALgGfP2pEf4ERwdJRbH63u/+0Ksgh/Pd611GWPgQqLOAjTiLs6CS8xHLR7EU2U8EAAADGQZskSahBbJlMCN/Hz9Z+6Xzm3/q7oYd0G2hnycvD/dneZ7oEkBQficNkv3sQWvRjB0O3J0i7f3L9TFI5xTtQclRChNrb+NvDDZvgNHtISAsf4R3ClMzNDi8Hti27fyjTiY/1PVoJovYx9gmD2NorCNGDW85eiStJX8zbjsBw/hmP8sN8jZiXmeBRjX7HLakmUV39clUJgwqNnh23yHABH+5ni2vm8+lQKHKUjJ9gK7V4lXqq7mt6OcgdnKFMyG6IrIrsSQiBAAAAOUGfQkUVLJ+WYb1tdpfdVu0F1k4kq5RhUrdGhAFDHYPp01TbmPpwo7SUbGx5vDZ6pujdLsJlfZM4VAAAAB4Bn2F0RH92+75YCJSoGPEHs1BaPe1nKVFLx3Wq4NAAAAA2AZ9jakR/19HKwpKb3hDHCpqEpgwPiuiuAowHjiceMfYzoQjP7bagp7T/MwoVg8eb8n/Krh+xAAAAREGbaEmoQWyZTAjf++s+OEuc2vQLdK/i8ClDkp0+Xtapl/Upi4Dlsp5sQ9tLIcelXy8eqVNGGFIhyPtEhgrjW/Ytb9bgAAAAMEGfhkUVLJ+A5qFYNR5HWQmzD29FNpojTcXVb5py+ZttjK0a0ojZ+6HbGv+AY1mIwQAAACUBn6V0RH+DDJXX+aBtjXyInrteZCVnwVE2lOmdE1eqH6zby3RAAAAAIQGfp2pEf3+cgskCrQfAkzm4yXvXEhlAALbstuTsEqLFPQAAAH9Bm6lJqEFsmUwI3/vgOOKjnj2kz6wE/ERkEFrp9mvc/FnvGTzuEvBxPM3+7YFlsfMrbeAOREsmq2wIOZ2exoKE6AbcyGVU2ti5k50trMZUI2gNhJRzN+j4GQsPgGo8hxfd6F/BV5QxpjTfZWDWEBGtlz/y9p06ppT6Q14s6X6BAAAAaEGbyknhClJlMCN/++A44rmWNeEapfe5ojbjO3gLbIo9IFs/sem9m8C9hrzgTXFpAoFFoez7AxD/afkFQ9GAOeRsTI8AKIit2NJwfsgAJBBdBVBQwXTcfUZuy+cNB/YDXD9y1MyNM/ElAAAAS0Gb7EnhDomUwU0TG//74ByFSpwKu/yY/TR+8rXmM/1ctJUphchJ0QahBg9/drvLVbB1uEUK1tv6c0ET7S2VOcyfJsGdD03uaOxggQAAACQBngtqRH93OX3cYVSoBNddIFIHoxaPZCAQcGLR8c5RXh9QJpAAAABcQZoQSeEPJlMCN//76z44S5zbHmusf8GcmWpiDjGPSa9PauHTPUthM4Tq2k7DqWOl1BucpTi8VQ6775shA/8Z+1n8HnIUpmW/CGfGL/3h8FtqksT/NymcW9deRJAAAABNQZ4uRRE8n3wHW3NeA+wKIZuyYfA4eD0aAfsQ5xdegmzrUzJAr6D2LS/mkx6XUnbpQmRhSBMw8zIWgJzltDGt3NdjhobjGWShq0kKsGEAAAA8AZ5NdER/im4U3cskkOWaidPqqW535NCXBHUD07pfk065jmmezMmm4OngZila6PBDQ8SA7tNzFMqQHG+AAAAAKQGeT2pEf4TUKvcv6iTIVCcF6BiozlPa/0aj/QBDSf946iRfptC1PLluAAAAXkGaUUmoQWiZTAjf++5XDhLnNr5FeZsxlZSRTEFl1hqtFA4oqhfVWwHju6NzP7Gzpz+tHjEfnDW+Oh5YrQeKavE+wrgFqrG9jsjA3/ox2THTdhWkSqy4aSCdsW/k608AAABAQZpySeEKUmUwI3/74ByE2ccN61JFkvALjEXs6KRRlHtA05Cr+i+xeGCVyVXxdoLYm1mfGZR+dej8GUfR5noE2QAAAGxBmpNJ4Q6JlMCN//vrPjhB4PuVF9LGgPQZObH+9xBHiekie60xU2mBEcs05OcSNiww67eP7NZPCqU7vyMShW9O7t9+XYJYXFlwEsy9Uhd8xicrYd7APjwDqdkyyfwfbzlstFm80HL6Q1TQ9IEAAABvQZq2SeEPJlMCN//74DjhLnNsXLgVjQBiBDF31YkR99oQqK3IUPQox0FL5Ngyqz0vr3e7zX4D7T/IzH37YlccmItA415mW5YgU9nDtlwhJalm+PnovRivR9bYdL21zzDZV93jsLfIFoLu62yIgI+HAAAAOkGe1EURPN99TaQSGVFDP8erdHfbNvCl7pmTW1uunWoHKJh+8IYvX3j4p8BO0arrZ28AEjtFV1nfwbAAAAA1AZ71akR/dzl9hMKpUAmuulsWOvCbvqAh1l1Ee3e0VIjtX+Pd6HNzYEzwaHU95p93feOAOeEAAABQQZr5SahBaJlMCI8YVKCqLqWXOXD3Rs2/i3L4bZCiXxdmPoXkYSJoYDVtTiH8gKixVFCbBh7U/5bLigBo+Lh/IlcxdgTekBHV7sq/sMm3KbwAAABJQZ8XRREsR4puFN3JW0ydhYi4T8lJqNEyweZMIfVIyGncYAQYhG3u1bYLctfKK0mJXCFln0Vl2u1VXLGsj4tNzvwBhCF9HkDygAAAADgBnzhqRH+BB5KrHlRVwgrrBwaKJDrhThiLTrF6/6zM+QOaa7ng74x86Si4GnXs5RovnBY3UKlHgQAABX9liIIACf++DLeYgN9XYj5zg5B5QJuOQipq5DurWEjnW3UorRPL+Gmp7W/eIApq4hA6GO7MXHIyDHsVucz/aMFYtKTDM8z25ZLQfiNZJqSiDN9723ESzhz79Cwt5M+8vOWuuGhG+Tvq9foXjrYoVXiVDwEsVfUmlk5TXbEqyYI+MZpo4mDJdCI6FifaFbOBypE2MZf5XWRt96wDzE+1vc15C8Q37MxkeCbaMkDN/t8mxkCWv8440Rlwk+ofcevu9fMi+Y2lGAjdq67Z7L2nX4tpcpo31hZ+kf6KZeyBIu+wE7zfBEWJkrDn+nvizONx5DVOkgfozhAUOvJbujzhGRN82gJwvYpRNAkkauwRrAkj3H3NX7/9FmB7s42z43E2AQejmKa6top87jsooGgL2yW/vkquH26p1/ZTn8L/bLpPfHN4mI2HsSxn8Xasz+QZbRKng1fvTxkt/mpfmK/kCyCd6bnIW6xeE5iW4Je98X1bSpyRG4ZEwQQRvmmA6Yna0ada8GP1pcw4M/JCmmV3O88PrasibHCahaFYs9akqjV5DdyOsA7IxWcVi1KwMrw65H0WqhITj2I/0392KBRdDeMYXdHrGwNEJhUIlQHIAJ5ZcKZWWU9TtKJlEqVRqCAyqMi4WjV04NQ1QEfwgz+T+hEehHkIC/Uc5dPLxaRa5f0rBpWg/CaNuFXLWTTpizOXU12VDPjD0PueUI5gcFxjM7UZKqEXIER8ZxBIM2uWna3B6TinIB//JF7wnKWG3KqQAoRJfs5WvYaG3M+PsmcpsEDOU5YG0Kdgf0PU4QKwIo7CjVGXnuOCeJP65DynJB1TaFN/Ec+KGgl8BEvS9O6qZ9x0NRWtKyaKyidn+wHMOsyvRqTPd0T82Zhih55cQsPNDIk78tfmfJbq26dnj6niV7W4C0RM63t1oa/61b6ejC3GPDzgPht8QbEWOmMkGWQaqDE1JlXaV0FU3nPF+k2FhLclRGSF12k3fSd0B8xrBHuCYayWiuH8ryJrtKZxJ5MQghDgBuHJO6SC1EDjEfxZGD/r6VzWAzizHGg8AmGx2xLAaw9rsvxLPUF+/7yYJueeBW5R93rr0FqleLihl5Xj2Aq5SkobABmcpWGhmxzvsHrBwKqxvkuHcaxY3VOBxUoDZhnQfEv+mZtl/1NlDKLlxwKt0ekKf+H85FMDoy087GajtU36Yx+6OhUz9cVW+TW34lW5zJzNRP+Hc8uhJDdo2lXVyYYbqCDCYtOyMVhQIpdZz2jF3WMV035t+ZxFpEnL8ijjN9Jg8JcuWyjfxs5sP97Y4nlx4+ag24HnZ+axqomYDumc0gNVyEUz6FZ9JbgAuwpXoA5ORKsBxrFSQHm0NhO5MryD+P4NGBrI8OniGQEltNCWxKO4CaPEW/ayYPnua20xl5iBviCOZOxGovBGA5TXYq/DBpZKzwKVYDEuxwEmec9Nh+C+oW4OqD94kO3pzQXD4WNt/q7X/PFLl+l9H8cdZ55u+BvzP0FXMoLDoRw7LJfWY5c+wzINL2XuOAYkfJCpNNcv1s6D4asebxSD1Oa3fRtN5j1N6T3U1HHUr8DPuALIo5ZN8ejEA5QFEH9MczGJJOxYHf2Ea6JmC49U20+9S5lgbbfIDilnnqaUN4P9cfr0GozChKSlAt+3fH1bY8x71jdkn9w2hQSZUYbU2pB9vw4KlfdleH+KnBRal9A8sdEAByt88ulnusl0zuMQqjLVJjZRTm+bVZlUB1A17g7C1j51Ik9g4dNmA5mrD6GgZNkTb5nGM3aTVSZBRdQchSdPGxKj9iMzKtky9yCgMnQWxUthLxXktBy3hagbWow9t7gvPd5DjMYYCoMNWUQ53ExMXz5i8JP+AuTFV+WVNG8AAACJQZoibEb/++1xYQAAcZSln5nDonHJ2FY8IUkCHV3Sy+GFqZcQ6LZxLbyQ6zhy/gRLW0pWGMhKeooRyMTxiJ2j2Rfe55qa3amF1vmYBco00/ApFnQ2DT4KOcgnT4ciTok2qz2OCLOIFJiYrYM4viFJA9fmIGkJAUIM20h/dudGJ62seoc2Z9AEofEAAAAnAZ5BeRH/gQeWpd9/pOade/rN+isrvDcxP1hTmlgtc64uAfKYMixRAAAAZ0GaQzwhkymEb/vEgWa+okZYoHAhrsRiqhh///5/poi1W1PBAeIgbreMG6vERAWN1PzvCg5/c1bRzTEU+q0vqY4p4w0N29AZHonm2TuYWlH/CYj3AyJ2CEW+WckK1sxFQX9Vkp3Do98AAABmQZpkSeEPJlMCN//74DjhAvudrQV/Von+LqXSU5EnN5Ici21RizZEycc6yRI7HW2mkSAr7OhvSjQAWm9VvWhylF3l+5QdacIDtyQ0c0VNiZE8gmdnua+vNHpLu+/zui78ynwbUg5fAAAARkGahUnhDyZTAjf/++AchVbGeq15oxbCanNPtandCKIk5UFi2McwnOkyw6FFIwJZCyx9BQ6ETJstOW6MqH0KmE3QQ9ScPBAAAABQQZqmSeEPJlMCN//74ByFT6UQhh76pV1tl2CE897btCzIgaESnlCrkP/rn7NokgPeW5v5Tlj6RqOhWJLwnxuEYh/nbkWklEpVsp02UeI5hcEAAABmQZrHSeEPJlMCN//76z44S52OwRnQ4IPpaB5d9PcQ3xV3g8rMCFmB5bsXTfJyvJ25sgvf46GpdxXXkPPndvnlrxgCG3+15DsD6MgtyhVrvxbNr6GQUtl2LeFZjaZQES5OWM0rLLlBAAAAVEGa6EnhDyZTAjf/++s+OK02js7X6gKTlrPgSr7UwuS/819DrDrfrMwhlkZEIEQkkAjuhvWd6ibuHYrJ9ufrJG3PFqSfmXeftKX8TepwZpbdTG4a1gAAAFxBmwlJ4Q8mUwI3//vgLpU+4B85rBOb//EpYgZFZ6eevNZXP+wPe+bIMaDyfbTJHzD4Kxd9xRs1m454YoGfxrw/BKC4vlUK6tcXIpJFyYby74nDDHCTcY26K03YRgAAAGZBmytJ4Q8mUwURPG/74ESpX+uSC7HOS2CPeetPFAoB1Gux145k4OcRwOiiCVIhzw8zcArn6HvJ2afHpwjimmjCVOlWLWAXs6/dS2mQzr5Tg+USVx+WhFO8PWOsAE5kFw2v8j/ac4EAAAA3AZ9KakR/iffqqvm8Jc9gUquRMvf9rGuOKZnd0efesGD02uCoFuLecT1P7NzuAbHBt4lonQco4AAAADNBm0xJ4Q8mUwI3//vgHITeeM3FIGYUnOaFjLcsMmQoftZNP8Bb606uLDf3eMjrSo7Sf4EAAABQQZttSeEPJlMCN//7w/CbbaisfdZegDj0M9TyOi5+HYmOurwQpxoiBko95zRoR8EUd3PPFRCKdewA0ZC3VL2iuZhfJK/KwG4Yg0/9XrmYb5AAAABtQZuOSeEPJlMCN//74DjiuZY2PsMDIvgZTyBAM69A3+mzj1NB9ReQMa5lpIjzkgpFb30Nv8USZUa4L6VDLhd3GOAcEU6/2m71HhKIVVQxQ1IYgBzydJ7TaPB/Lp71eKZ3cme7OXKbCXDl4TYBQQAAAFFBm69J4Q8mUwI3//vgRKv7oyEofCjyt3vsnfsah2kHIdELq19J2RoKoC1XL0zZTv6sAPxkSDJ4hua2mqjbC4lGiBfWgM8EmDQTaVR78G9JBOEAAAA5QZvQSeEPJlMCN//74ByjvhR0RQiScU1Fdfdc2SHU4FOuYJlXYGqlgNrIdNY2a6p25gOjejiWvqYHAAAAZEGb9EnhDyZTAjf/++A44rmUqQo2o78x9zImFz6ycgjd223QsrF0bQ9oStMycJKa04q4jD45hvqUv+ExtcHS0/sB5B/uzSniIvGiDyArKszGSLTNmQx4zT4DAS2BkYC1Ww5Py+EAAABFQZ4SRRE8n4DqJeP15Rd2F1lVcxp2jzwtUtQ9ScvCBuIZmxQraJl3DozW8juXqxWghEa2yC/V1oAhaIA2kP/2/bNYpFpmAAAANgGeMXREf4puFOCtReQubvl/VjVVRARk0FI0mb4PyIwlMKV46WFsEaeKjqMI1+R2+DtK1F/8vgAAAC8BnjNqRH+E1Cri5ko22NzWPuDSPi1danTYsunQ9icYtQPRSlq6m6AnOlFA02op4QAAAJRBmjdJqEFomUwI3/vrPjidIX8jzxcO5M8WgNUVjZvTZqttR/yezTvM/BnIL5HWxcNwKvfguadDkLdYtXjtonK3zpSm1H7UXsq/0YHnzRgNDAYnqnNKKa3EfIz15EOoRIhWAiTMZDR7m4p4cwv3/77a/ATHMyY8Lm+4XtVBteQnl4516EwPOY7Hx9ZJsncgz7c28+7ZAAAANUGeVUURLN9/4NxQllWwS6YjGgAHhfIlm6iYyMZ/pi9Hb6DSetZgTUH6s+U8uHcuB/OfkzdgAAAAHwGedmpEf4SbRPB1qbWDavfsuHn8SwfaNy3M5M/M8YEAAAB4QZp6SahBbJlMCN/74Djh7MxwQMXA4Xa1gvMmJfG5IpGYMILlDBGG9SiDcf+miA2uUriMuCZPDw2oL49uTVSBzN1sGzND8Fw8Hhrs/70fxGPxNhJdvqFB0w+DC+aj849W7UJcjUb7VkYGXQv6/CDjihFxZP7ZmzrBAAAAUUGemEUVLN99Ta+lvuryT/+7Fb855I935eQ/NKyrIpEljglVrbW4VSwnmC7RuPcxJO4HyTlFvPN0kQgJ7FjETsgCU0+IAsrdcNinbd4czNuDZwAAADABnrlqRH93OX2EwrJavM7/2lXT+mo4pwFqeXYue0BFgaMHvinTIiRN9nQe6i5heOMAAACVQZq+SahBbJlMCN/76z44rmUqw207hd7OqeG95zjobXO1BXc8GYszEgUi91ZFitoujcqtv53F3MqCQuvkILWF7cYeyLidHTm+UXkyfQp7XBSq9+bk8rzweySNLOjwjQrKhpbObcN64LF//YEwizj60QETaAs6xDcg6lk3U1uH0HDFd4ifBkn3nk+cTzIVsSPERF68h4kAAABhQZ7cRRUsn3xlVUbBHfvOCCo61ctk1vvWUTe2GOke7+ZMC+0wb8hMaFSQHz69cBeyCx6e99M+HncZK5sJxhMcQG3iUMQJQG5qp90WppVxek0grIUNEg+nMZ/OSgHVaYGxCQAAAEgBnvt0RH+KbhTdyyn2mc7b0qS5IBg49T6i1mEqW9SGK2rLIsvaSKB6hosMeUz8+LtQpC14bO5miLujSI3IhMrcLsnaQe1F9UAAAAA4AZ79akR/im4VKaciB3zfJVSkhw/Z8IUNrw+u/pOSGKeH2tujCA/SRAiTtkOg1ajX6LWUHYFkQrIAAAC0QZriSahBbJlMCN/74D1FQ4s3n7j5n5Nuu3X//9TPET7aPn8s6EoF8yfvc/MfvxyY4pGK+0qv1Wm/xldOtP5Q+S6fps0S+m/UZnQ7EwPsBycEwP1ZVg3aRa1YrfiIsrKz8S8liNH46Caw2vDUSR+GdL+B35GCtu4RO0FhgZFHJ26mYluh/0OmxLZskzny1sgm0Ygl0PCSjypQA0smNK9TeZ6vZ59Qihh8Ae3AfAr1QlZ6NWXTAAAAPEGfAEUVLJ98ZVVIdZdMwtzfMxEdTpKGRFWgnHy+yhWW0NeLWoaMzR3O/Z2SNXpaNidD21Q02s62yRuP9QAAACYBnz90RH9/iXZSP9h8cdeKVvwRfd8hZRQQgN0TPngGS7Fwab7aQQAAAEQBnyFqRH+BEcHSusAq258vWb8w5iZsvh+pHv8c6GNrxfq4gOcru6UctZ3xaies0Y6hBXGyrHqsGGLfpOkCH5M1jDKmJQAAAGdBmyRJqEFsmUwUTG/74ESpBkuxG2KSXcpc2jk8koxEI5uEi5yoS6e9xpdAt72c8MSnXbbConlHP1BzUnJ/+3z1Aoy5efyGWrhfPczqcoqfv691EvI+hHCd2LcBKTSXuxUnnL/0bEnkAAAAJAGfQ2pEf3c5fdxhVKgE11tH5m8PYCgh4Qjt0i/xdPUGMftsgQAAAINBm0hJ4QpSZTAjf/vgPUSWjRx66POWg+1VeyDdo17n1pRGCAZv0EnMYI22P5nBPtXTAZp6Xq/Jn7fkYYL6sJtWyYB/PJ+yuGPOHKNGaRXXtSKv8ZtxjnL8aUG9NP/FDr8upgFotXBQY8KM92B40pAj4Gj1imdJHCvfbU0iETKVBhp1wAAAAFpBn2ZFNEyfgOol4/XlF6ea6k7mZM3z9p4e8X2Jj46Hpmlke/qavrQOjTwS9Z47MV26uYG6NPmmXkDbOOyLBafQL9eNMNmRUBMDuSM10MkN86FV+v+G2vJvgoEAAAA6AZ+FdER/gwyVRAsFrTv5yQfi12mftmfw7u/eMLjWn2w/45OdHFn7ezFA89/eGqdczYXgEIo3louwgQAAAC0Bn4dqRH+E1CsRGdq/m1pty/GwHijW++zyvH/Gdpk3chKR+bVbMoyVnPmV0owAAABrQZuKSahBaJlMFPG/++5XDiuZY2N6PL+cwxD3ivJ/8tkqosMlZ2kuhgf0q00kNIyHwl6OQiCcnPCaq6e7v2Qvnnewqy5mzqYUh0ZzBmS09jJQtDXRECUYUTYcu7s4jMLzcEOBfNCz9U+yd4AAAAA1AZ+pakR/f5yCy9StiwNVmr5zLjqXbDKcDqtWtTICUwtkcR5hRgZuV/4Inokvtcj6Vro4MnkAAAB3QZuuSeEKUmUwI3/76z44S52Ok28EOzD5G2y58sz7giS1bNt4y1aleBVYPkGtESwSGS22z3kj7PNRXWIQBXSfnJ9k+hqAla+toGcJ/TJgevZWzmQsCjk3toczbqwsz3LPlj/Q7tpVb/p36b3WNDT2urx5wyQ2u8AAAABDQZ/MRTRMn3iZR4RsYALI0qaMPz15tf/PCJjXSCVLzyIjWEB6NIVZuFQd9IHJUmnlM+pGzTuq0jqVN7pbBnNO/83bjQAAADkBn+t0RH+J9+qrQkFHiucZngCfYVC6NrC2q1CVtl2NFIiI9/8efIy7v0ffzAkuTuoM22cIjGJ9hWAAAAAwAZ/takR/dzl9hMKyWrjL/2mAOYr5XPeIQqNEnZH84kJ2zn5AugiKzfdzkWMm0Z7BAAAAf0Gb8kmoQWiZTAjf++s+OCFJBZyyg1Y/s//daL2cmKlc5TOgqXAetAsmctqz5P823VwTaTcwPX+/sUmF5uTMjeRSbn+NbF5LxJ6JtQsih24gjtIz0ieWd7148vl0dKSJkfFHV8qbphVazHMzamX2zJ5KbX5iaHfn+RVsiaeJeH0AAABZQZ4QRREsn4DqJeP2RDmwj76BdzDR/efXogJAcIU8SeSQHCj0dqnjrDVpTgwQJI31BTaAG6INqkujzWTGU1zd1DjtXiaRFbASYt3GqCB6jgT8nxWzUXikMYEAAABKAZ4vdER/im4U1zc+vVyb5fhhMvf/N+sn8Zk5BtnQgFm5gyE+ybdRK46vySho5BSu7O5YawlwbRtBqouRndC7tT/GyEW3yt5IFoEAAABPAZ4xakR/hNQq3d0NOvTMPHejrhYTBJBnYWLUeub3UrmhxZK/nEFqJcD7/UnWAqDmNwXcWBRsIOSn8sxa0LKyMIigPgimkAATId2ZpEAtIgAAAMZBmjVJqEFsmUwI3/vrhHZDqhGUnyzSHEzswaZ2RxqKGRynYBXMBuEuEFfHmUDj9ux0sbihOOVmVeYWgA9LOoKwJ+rMjiSQ6ZUAyFQsCjjaDAUTNJtjWai9Rc+L83o1bWP1S2yuQUsJLNLVuFrlgQAymRGcNkyxEtGUkanf24vl9rYEtUPRoVcJrHcCRxamMNaxH5ea8jkqFJL3PStn/7XHDDku9uliPPRnGOvdudC3OTfeJhZuBI5KZb1w9wEivL8R2sp5CIAAAAA5QZ5TRRUs34GrPAdCUxfCqZZ2unlDcBy8LcrzCPU4ls+ywf4Ue1Eh5IMuWUEO/crgik+To25i+g7hAAAAKAGedGpEf4RorykjSfm5OcGVQs/wa5q5I3WVt0A7JQcZioy+6Cpd3cAAAADPQZp5SahBbJlMCN/74DjhLg4zcYH6PAnbkTY8iUAl5ul5AGL6HlLmdMJTDc6cLh2fLVFs26FJQNh8DpEVTfPwRt/7mERhKyfDeSIWFU5yQMCYNHv59nvZFp/a8WegCK4HUMfdTe9BBOJ+XPTTUutM2EKDdwEAeJMZxuSKyplgd9tOuySGUaYtDrIicCWJFC5wXwmU8vxDFfE04OEohATSR2A5oLuSDW3Mvf1CjoKTHqw+jE9AVMQbxRI4vzb1WUJejTxdtN9Si/S5yQWiKU8cAAAAUUGel0UVLJ94oE82Csl+UbsHxzVXEoxsMq4cxym8W9nAZxQ5JIcYizkCGXhABajDTcmqAIpfxDxLPxOswTcGYTAAIZm4AfUvHOJyB+CdHlrCgQAAAD0BnrZ0RH9/pbSDSQoUYXbJ3rSzc+kRtLUqkKeKWWa+GfiCNMxON4cEELkIjdppVpoCwBlXx1ZSOo1lcke3AAAANgGeuGpEf4R1fILWoVH4wTcI9HBEHVSvB4bqbjhbmfh2jcyxVPpzF2b1jR4X1zi6YlYACKdbQQAAAGBBmr1JqEFsmUwI3/wDwzHCe2QWuq6WGF4P0C2Zsdw46Bt6zA57Rc2TiTbjW4H7HgKfR8Gqh82t98tMHubaH8epd0G9TvymJA0ATEGErg/d+4p+TcjdLCZHjYjKV6SX5pcAAABNQZ7bRRUsn3xW3JFluUBopv+0Jh109W67tBai2wVCw1/Fli28auwIi5iiXJElON5pRsqKGOlx/dZECZCgJwmVdX4qfM5L7ofdjxymz8IAAABAAZ76dER/gGwpqdu2nr4gydT/It9EKEbgleYAPE+8ztHAa3x9WnkO6FBi5Veh1TP1eDEtNkuvOywKMxIz1SR2hwAAADsBnvxqRH+J9+qq+bxiZHCFr0gnhkOz+ad8z4gs78ck69ZSTjmsIwQTbotwcaKIZaLcxmKWlHaf9l4JxAAAAKBBmuFJqEFsmUwI3/vEhS51ZOAPrFCl1ornx2QnJRuYe9kkqMhxGOaMMLj/bNXs6VgpAhRFmLwVR09/A7837HfA1aN6O9OTeIhLPR21+RX9CY66WwmoIi5/sqTumt7qln29/+p0OIR7TaPCsms+zm3oiZ01cw+QI1KHDn7eVyJ6j4XcPg6S1r5fXKbly75nNvXLiFj5ogktvCp1SHaqEda8AAAATkGfH0UVLJ94oEqpOpYRxdsQCq36e1PDI41WDpJPcKamY+k0pWs2/KywiDKYOYI/MrTcoCQxBFm+MjLqCzLeTpaRiENFpLsVzukTIcUPCwAAACsBnz50RH+EVDZPSGusQ1m7R8O6wN1peQ/1jS1QCsK/+3Um/nowdX2k06ZAAAAAOAGfIGpEf3/ocg0UgjqvtSdJTqly3l/5RQ31A5IZqdimmnTC5MIJTLPa6Xa8Tb9vHN03oEBvLxevAAABDUGbJUmoQWyZTAjfx0V6DnpfBXTYnQk2bwKX2bwMtCO06TgUqyoBuPFXgEpXFMPAzWMC08ZZ56ZhxoN9TkSIiXVTokTYYuz3pKV7lAe/FQPey7N+dNtFiUjQWvY15ek+//wU0OU7OCIKAvTzIaRU2/1uvJ32Jwtf6MeG6udoHVMDrga2SELkoyovsSaeoA8+NLOPdIlZ9IkteJ1UZv0f5rSTvZTQ5nm5XNAnjE/38OYo1ggnhiEji+Cgk5Y2BrVp0F8ZOcAg3273NFXpfm3+x8FlaagyvTugb1BaDOo9dxr8EQG1u2ACk3eEl2m7Cw/3wccsdqmv1k5rlStFNyS7rxGwIGSghx6CdSt/U31jAAAAT0GfQ0UVLJ/Twd08p/VyChpS8T9IWptOSDvdLPszfbix+67WAXfjDNjm/bHE91nO3g6zlnZrZDZR0Gz1rsuN02ikA6mQRG4q5w3HL2ZCEhwAAAAkAZ9idER/hHi2UFyJjt4ot2IIOgG+Vq+0x5nM3Yla/iCdTgZJAAAANwGfZGpEf9fRGvKk48/Ke4CgTzeekpncWwo56eW/sFEV91i3WBPnz0ynsff3Tnr+fRiCryeLNz8AAAD4QZtnSahBbJlMFExv9Cvq/twdBueYn3Ee3pjczv7OR6i3WU+SGv/yONzlAdl0TMgBYUraQumePMz3mddrVgKe39/UH0dZMWofZ9Mqui4P/wvLxsUppGHrrazsvTR+f2rODWARjtdkCSvFbtfvMMzeLFsQ13oDAGG2vvjL5m95QTZnRKv1OnYK93SNtLC5O5AQU0iYLnxfpUikyvxOvC8fzzbFhAqGYgHL1GBJ3KgSybvxeNLFYx0n3LIy1+V+pUzTJ/EUuTURErNUiE759HI2TbRLFPHYcqtlrP1/TS0NNfBRxkAXAcx2r7lawk5xdWWia+VA11/z9cEAAAA2AZ+GakR/mJ002N+vm8JdEQuzLaLtUIOIHfz7FHdVElv3dJhMuXpT4IkFixsgNNJKwy7EeXqbAAAA4EGbi0nhClJlMCN/x0W1i/GN6GepRrlUQYVWJ5EJgT8Zlnuh2q0V4+3OOPM4/4/MCjdcBbKHis9PBt/yefxMxHgw7FScKctBHvFoR1EawxWUsaHfzfSnnalnY4m9Ak14jYSJiHQ5ToCi7Y4bqJeCGTsIs3t/Q3xTdD//Dx8soX29zf1N5xgAr/qPmK+dU6/y5sivZgEvOoh/r0LHDfok77DJX4jZRZR016pPj/OKEWHWFz2SRnFT4/7nXgVNFiCLYupDDD8vVEru70qqarHLzZrSjSEw3CIA0xIhUFDHlk5sAAAAMkGfqUU0TJ/TtTGaqT3w3TTw07mec9pdUO+qick9Inm7YO6odtgQD/RhYzoYrtniSB/IAAAAEwGfyHREf4TAtHqZwf6daDH8ihsAAAAwAZ/KakR/zdHy96S6V8otDkejDD6GAoE1dBa3yj/ms4XveSTIAKMwyhRIvm1XthnnAAAAfkGbz0moQWiZTAjf9BLVe/EEalTabFqMzxVTAPfnxJ/Zlk5dL1vhj03RqKWw3suOlU5fZYsZNbQTFyM5VPD5D34ws73lho1n4IFelgz9POTbop+0V7yvkVIdPx9z1ZXKiB6WJVdOUzkbyhc1ThMFffRcKHS2W7razGQSoY0BSQAAAFhBn+1FESyfgOahSeu9LaQmIfCcRMhY/HlRnhZaQPy6CAtyQ+LszE9EabEU3RQtsAcY3sdxq8ilGwwrV6BXotZB+DgqfBZjqOBne2bdeY+xHF3ipKcm5shAAAAAMgGeDHREf3p9slhDFwgFsyLLeY0fCtie4WUMzC6rVyKWDS2ciTpdi3cETt3ZfbuvDXChAAAAKwGeDmpEf4n36qp82ESoc+OmMW76ljVM0FgrkSvJn6xJsKLGiU64smxwxzEAAABzQZoTSahBbJlMCN/0EtUjX4qLA/VxpTkfRuOUpW531Ao8SwkqNXlLmRBGxN1IsecXQYmRrqrhzuWJINDZbuHe7a3EXT2PLlFrMUbyi7HcW331ihx5HV0pCo1UdCwPmKqK6tJGWYXzqyTyWsq03jrcYpi3wQAAADRBnjFFFSyfex1KrddINRpEPynlrUA1hgUxfwc/hqlijWK0BwD23AC6+DLASM32cUAM8GGBAAAANQGeUHREf4ByGO+by1cZ8ELLxtGybJtzqprjV2BYPrAXGxgEvx6Yxh6AyzUdErjXJBHKQFDAAAAAIwGeUmpEf4RorykjhEu7Dyr89cWrOYZIKdrrF5QUKCEwPBeAAAAAckGaVUmoQWyZTBRMb/vgOOEudjsI3Btv/VSW39BLCWXy94pOkzAzXgbOIzUOMjzvnTjsYbNLA9pAvZ9YoFxYYhx9h4eU9NZk38/IlYerTZE/6bkDdbgNWkejxbItyYDZBfXrFL08jsMcNfK3iihUtFthcQAAADUBnnRqRH9/6HINFiMdEuSVNpUqBczjZC+frF9/+VOzdYvwD3y/RUO7jtNMvopLiPWtfsfkhQAAAFZBmnZJ4QpSZTAjf/vgRKik2M+QxjuXx6xjPBvXPi+CYy/3jhDMFE5SsMyyTR0PdQ+GZ5DvFQ+4McBlrhdYQyRJoj/S1jsCJTCCYb8261AmPc/xNGhEYwAAAH1BmphJ4Q6JlMFNExv/++5XDhLmEjtL29ZAM9do3FyhDOoEv9H/h3RF9bsYANZinaVx61rvr6dcFSPOYKYVa0Ky3w6DZJzzBXVJRPeXQdWDKgQpImvJu97y8iZ+2lh5SMVTWZyUhJjoYcacZWmU6rIRgEt/EHai+6H6pE6egQAAADwBnrdqRH+KbhTgtKWIdHh3CgwGor5ORfdsM8dLegC9vy9X/NkeqmFpzE3INqVqG2xva8onEHyYu3G7LjMAAABIQZq8SeEPJlMCN//76z44IUjPWuSuXQG9n5T17mOUqH+IrXRyER2829XfIa9IKsVCuw/i9Uft+dhAIOotXMZgMVH6ntznsvIRAAAAQEGe2kURPJ+A6aXZxXJX+dIXeGpcopCzqg9ywRig3tyiB2DWkAg1Nn19JHMqupfHrhRdStTdinM2zM8a/E1bWFQAAAAvAZ75dER/hRLKKjdRBBNZ2S51o0KrXqTH62/XmFWxCqMc5WhSfqLQHzwmuHVfuIAAAAAbAZ77akR/f5yCyQJ9t0mSfiBytXO6KNT8CYfxAAAAgEGa/UmoQWiZTAjf+8PbyScRYg1RKyMpfbNX/CZRmm7Hq5VPl9Kuk8m7UOZdAbbrLswPQZQjheXgSbH4nLjepdplv51yrDISiOwhOpa5KSaf4Oh7414o8y9zXZHdmfzx+QS1lrmt24zuAg09D9q0y9iOg2/VIJnCofp+Dcv/U7ggAAAAcEGbHknhClJlMCN/++A44S52J3BJbw5nJLJfrjdwX6x1tDYZqj16z7INfCLPFEvnVSYK+EgFy5iJXZMiPitdWn9FpnmJqvFHhXzj4TTJ0U4uxYajaipJlltSfuzvOgecTBXwsPWX+UIqDZy0B/uWYEAAAABsQZsiSeEOiZTAjf/8In/MgnTsdVF4hM+h4LwTtR11DWQhRcJNRvQrtL/nWVC8Fu9eFlLAp8tLIZNEfkIXRbXsvYnQqfL8udnCweG1H8XL1oX20j7WR86CHE+9Ke8Jhtdy1bbfdOfKjqHqgB/RAAAAM0GfQEURPJ97HUrykeGrxIFqUtgUr7SYug/yEk3MgZ5o4qOnfx8yJ6ewAL36Oxl32rzUXAAAACEBn390RH92+75YCJktXmd/7NQWvwATKbauC51Qh/fxn8EAAABWAZ9hakR/im4U3dpeEMcKmoRT6qqm2/JijnAse/V8GPlbeJNnpdMFOlnqWvqie1+Z5e2AUboKtScIhZJGToOZ4DYrujj6Kul4w2TJ9Hh7ENnXsHCbwRkAAABWQZtmSahBaJlMCN/76z44EJY/u4nS+fFAtqU8yWqvWTE8GR+nbQ2pXN+qpaQkJCk9IKk89H1bE//13lfA/iJjEiHwWexUi5OP/Wk+Wl82rT8IeKMB/agAAABKQZ+ERREsn4DmoVg1HkW87stKnuJEPWC6LUKIZAqca58O77RVjao/vMVW8zN5+SBSMrUHBYqA3TCqOVCAQjR32PKFDGNqCbkULzsAAAAxAZ+jdER/gw4tmzW7rzcxfEkNyf/onzmAWVc3vPO+lyds9aaITrQ3fCOV9PNwrTxg0QAAACkBn6VqRH9/nILJAq0V22CwA26I7F68ZhwvvzudCOtW9UTJodNhQc7hgQAAAI9Bm6pJqEFsmUwI3/vEgWa+oh2wHUBNFWOi0kAxeW/oBNECdxapz6X4Ulak31gejD9icbk8DhRm0Hs2hpnHJ2o0E9aOoy4Z+ZKHeuQbkVZ3BAG80uP6pot9Q7nbRo/Sr6mvqnH8mE8zpvT0kWpYJYh5VOQG6NR3X0yMAeVL+tP9aFjzEgYaAOOYI13h/D/pQQAAAFlBn8hFFSyfeJmV5SQCgqT7iG5bhfyYEe/eRaEbB7n8Zt1Ka0FoANimziTzHCgJBfnDtij3/czGvfJsLwpDb7MAjV+SrguqOiQF3BNZwA/WhjksDFwtQJpmoAAAADcBn+d0RH+DDmAasguYEJlS+qlA/VHzkU9DE5VIEBd8g6r1fmm3k/DZypjK98baTmgzNHfZVmKFAAAAHwGf6WpEf3c5fYTCqVAJrrpdDQdDLHRuhKusF09cRKEAAABvQZvrSahBbJlMCN/74DjghLmmO2mUnbfWaW5B8Mb92g3kEwH8rIP/BLmQ9SIIzp0Z/9N2kUyz8oEJGs/Ap+sPATqqFFZLVvfg+PmnCmip/KLJjlyWyzFDoI36aRp8zoUQbAabU90Z2YMWlc5ECtZiAAAAVEGaDEnhClJlMCN/++5XVEO1lWujg8lVyQxzv7cDyj9Khzwt+GmZK7MW+odXd1sxR2cKlqoYmmu2Vz7BYmxOdJPgwJSdZxassfDNt9Cqk8Ou7kbHwQAAAGRBmi1J4Q6JlMCN//vgHII+SiBmB1jZ92ZKw687olAG63mvy4JF68C/BtdOHNqo9oOcgCDap8yVUyFaCvNPKBjx5Y12mHem5KjGuR1HE1vwiWnlm7LaNIt7dT5FKWs5/g4ZnkYzAAAAc0GaTknhDyZTAjf/++5XDghKblpByuukqCy+P01EbidExQWS9H7GxldpGPDpK7V+tBUBN0kUkqNhKVgVQcUsvuBYPIwMAl+a8Q4xkIhl/9reTev/yuiRaZox+SduCU+GYb7D5dve6JG2bETSCoHHzLVhoOEAAABGQZpvSeEPJlMCN//76zcgkiFnycdhsLEw0asP63rfGuqZvt04llNZeVygNUg8Jr3NY0ry9QaO7iXi3konyubaDZ6qMEZf4QAAADZBmpBJ4Q8mUwI3//vgHIJlMhOje7UYw2E7Ne9jyk9kHh3DzUwY2oalRKALKuqLHKNjTZEyfXEAAABdQZqxSeEPJlMCN//76z44RALE0DW553rg5eSpXu4lPlg3qaIlhzcDVr/CwZnQ+H1LWa2hx8yylsTU3OCdaOErX+9UvWtlkHhNfA7/4IxFk9dSDJPX2R0DfM6/cnV5AAAAWEGa0knhDyZTAjf/++A44S52Ox7QAqqAU+wlZbuVvsj7Y8xPJ/oIVLMj3KgFlKtIHmn1xsMKiZMfby2S2aO0xU7D3AaY+dFVaek1REHuf/uwXVRA/9Oom4AAAAA5QZrzSeEPJlMCN//74ESopM95J7cOXH+Lihd7kB1WsgMIaYuk4/P93ZFft876I1IMoEoIC02Nlj1AAAAAMkGbFEnhDyZTAi//+98ri6ylM95Xz+FrnToIvklvE5kKaL4Gyoqmk5OFcr+nIm3Aq8lFAAAAWUGbNUnhDyZTAi//++tE3K4Iv8XvdFgAw9nY/t8iJypWqoPs8SfKSv1NUm5Aud8Hq91WXCVfMYTS48S1kcJQTJb7gwxOICnLQAyqaV+gSJ/vuLhZm00gZbygAAAAV0GbVknhDyZTAi//++5Y+Bvd0edhu05xaJ4bHAypArfk42W21z0Hw9aIvP5GIJ7U1/Z42XP07srgXfJu1+InIWELVunkmOM0UPSk9U0pMQXyEMc8Ph6NIAAAAGJBm3pJ4Q8mUwIv//vuWLcrkUYnrfzKvY6YyeTO9Rk4qyfjY77IkI55iiBIrunpAkA7045O9WmOwt1TUqx7i+IPJ9eh/xnWjecTeoRQNm+CQoQo46ThJC/lHxBJ3Ma6WMJhgQAAAEVBn5hFETyfgEm1lLG500zAj+otoaEuWn4LIfRFhyaAA9CwtVZ6mp1AZb89j27+RZtQcexyfGBpeuTXTM9UC0Wa3f0yZmEAAAAlAZ+3dER/im4VKaciB4AJOzHjFym3oODn6PO6EfFl40Bc8IqeXQAAACUBn7lqRH93OX3LIRNDM9tasR3koAjrPHMOiuw7i9TDEgskYPw/AAAAiUGbvUmoQWiZTAifwAj3GUsedIbwtVPmC6UOtk0BDMV7Fztcurfnw+Ukl/88z4rxasXNQ9PFoO/5nE8qDt2fJK6LMxkaoX4RjOKmCD3wJkQpTcQNJfK33lRLTK724nIXgRJ5xYIwCOoOErOutudRLF4PZ/ViMvI+XwSlh7aUB0gxVLJooJL+QGvAAAAAL0Gf20URLN99VLWN00toRsi1JSCmlbiDfGAKxMwuqu+ejxaDbCFeBpddWwXK34eBAAAAOgGf/GpEf3+3wq4dDu7w/BsOHSau1FuQnVGdVhqVK3yiZQNkNPtat/J4UIGt5UWW8ZQo/jCP/r1Nr3AAAACHQZvhSahBbJlMCJ9uIMp0i44t0raDO9/qWe5SIVgIlo7RXxku4hPZq6uBGB4yOHnSn7I/dfcyQ7IQL8iwlPxXakCthkEjBcrQSK1+YKDZ1aeYZj1E6kswhNqi3qb+Xp8S04OXkru7TYYvfQY8IlY1dIAa1uyqVN5tJyc690K5tFol+v8sP3r8AAAAWUGeH0UVLJ+A5qFJ7VTSGZaxh09N4LvIJtQj3iBd9smy2aE0j4YJdcN+aL86nrecX6XJe9x9zozcD5zJ3+OWlY6VcS3j059iPKQYmFD66vdkRvd6PnTkL4nBAAAAMQGePnREf4R4tlBciUsnWu6mI8oSE9LvSXimcHaU1ew88XBcw2xze23nLBJw1I7DAmAAAABTAZ4gakR/gQeUyXbanD30/po8AnM4qaiU66bgB4fpv+VLJIx/KYTExAZQTmPLSKx1c76Lfl3wOXm3/JkWtAbgIRqsn6/a1Z7l65jvwwWZio2lQcEAAACtQZolSahBbJlMCf/mV66kDVGRg/gckK5KqDoGjG/Rc0jTpLhTmGpUH6p6+6Lg2rz4jmaD/nRZb7pbBze8wvCCpGUZHUhbf8B87j6IYVHLv2NWZ2HtpdsJpaK+aplUwQpQ2jyOwJc8FMrzBW66sb7D7eHQNNKX2TQDpjQOPBX2qpgf8m4FZujU9qlQGljaykfczBPPyrMOa6hPI30+YisxV90mdGge0/yWBpepB/EAAAA5QZ5DRRUsn3sdSq3KxWQxmQudZqrOmd/ZinGlzs6ATElOuuq0U0+75PEkS4CwMI3s7/bFIEfrMlMuAAAARQGeYnREf4ByGO+bxiuCCgeG77QV8cdQPoGcUaIvNKoixhEtWoRDrlt89Zpb+CvpUPEpQWLYG2sl7ltGjzmAbh7BeK7WEwAAADMBnmRqRH+EaK8pI0WvJO3oqC4LpNBdDqSQ/2H8yLxBYfp6xHMaIZZrOGpnEjIG9IZgYtcAAAB1QZpoSahBbJlMCX+Lwr9oP9EDpqimHn50oJzlNSRyBVnJc7Vs6Id8pb1K4EjPXOvWC39X4znnb5gA6pxzGbQIqrjG1MBa6dnTE5rY8vR7mZCBx+zMoAaGfd4sMr9WM2xq44dO8HqaOxPZ9945Sd/152ym2cFBAAAASkGehkUVLN99Ta+j2R91Ew92HLRzMWd5j8GehO+60tjrq8bx/jm0obpABIT+nkTNTJlaL+14ziVE2QZYV4VvHO4XQvOYtn92rvxBAAAAMAGep2pEf3c5fYTCskPjx10tNs/dEitfWcAR2pkb+X6knhPNb494o6ye6ByoXD+voAAAAIZBmqxJqEFsmUwL/waaFZ1EyQOlLmgdat7bD4DlnDKE5EZeA7wS8KuScmLjUp2EbKQvkspxKeGj6wCV4/o9afBe/9NGPieBtZO1xGmVIfxYKRK/87tsBmbY+CKUcMExGgKwo9oXid1IeYBEUO08HCareCjE9Ukys7EXjqpBhqzoxi7rDqSeDQAAAGtBnspFFSyfgOol4/fAqtD33XIsmEONOHHMdFPLHo8kgNBLKWb9Ym0a2pfYCaWEGCwtlD2bnmGnRxwnkF0/3ikz9RRWbpi5VR6qekgem9YtTpaFhaa4okfsIeRm5VyV92Bz5W5+vfvQGlzKQQAAAE4Bnul0RH+KbhTf9HyRxrjP7gTLkjfovmjjipEuxkhrSfiXFbTj+bdRKyxdbpt6VAaA+pPkrRqInGqSekN/aHLTqgaJTdQcGWwTKKS0O/AAAAA7AZ7rakR/hNQrERnav5stEZ3ol4KYJIAtAl5O8ChfhNtpxCXz+ANkDpac7UeKcUe4PCnSrmKSFibwLDUAAAB6QZruSahBbJlMFE//Ca2F/yEY9AmJs0u6nGYXXc/tlS+Ep6R9fuGO2/jOCP4GVYBDYfNYov80JOewYiS/s/vrHJa0Qz01zjkBZYXCzefVUhU29LngYko5T1hwFRoyeH8iTqTzThvYOCNdd/EJ3oSIlNm/FwdpeovOsTgAAAAwAZ8NakR/gQeWyVZmnrxk/Y85AYHkLo/YCh9kxyYfbQDb4iHTGTcQHbUUBaOiMgGBAAAAU0GbEknhClJlMCI/FWGmIo6t5jEJYH9imbohlas+DBV/nvXtxmYQEaER8fEGt1ChymPpiTYpdh51Cmh/YiLHIUzWJgypPkYhnUaomR/HCwbaiXipAAAAWUGfMEU0TJ94n01vVRnAY040HFEBoYfXPVz/vno2VX9xlArIIXFkMxlTVQQFmfCo6f94H5UswTk7N1qNbBjEf5ti7RycrMWYZI8EZmojies4I8dHu6teptD5AAAANQGfT3REf4MMlUlILcS/8oa2RUszwulzz8mTh84lHFOM1E/jfwpW/Z9ANgKIGdGUmnAXkT5BAAAALQGfUWpEf36Uk5AyZXqKVAUafPJRrljZB7JJE29//0psy6/KKcUFyP4LnFh/YAAAUAttb292AAAAbG12aGQAAAAAAAAAAAAAAAAAAAPoAALlBAABAAABAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACAABPNXRyYWsAAABcdGtoZAAAAAMAAAAAAAAAAAAAAAEAAAAAAALlBAAAAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAEAAAAAAUAAAAIQAAAAAACRlZHRzAAAAHGVsc3QAAAAAAAAAAQAC5QQAAAgAAAEAAAAATq1tZGlhAAAAIG1kaGQAAAAAAAAAAAAAAAAAACgAAB2kAFXEAAAAAAAtaGRscgAAAAAAAAAAdmlkZQAAAAAAAAAAAAAAAFZpZGVvSGFuZGxlcgAAAE5YbWluZgAAABR2bWhkAAAAAQAAAAAAAAAAAAAAJGRpbmYAAAAcZHJlZgAAAAAAAAABAAAADHVybCAAAAABAABOGHN0YmwAAACoc3RzZAAAAAAAAAABAAAAmGF2YzEAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAUACEAEgAAABIAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAY//8AAAAyYXZjQwFkAAr/4QAZZ2QACqzZRRP58BEAAAMAAQAAAwAUDxIllgEABmjr48siwAAAABBwYXNwAAAAAQAAAAEAAAAYc3R0cwAAAAAAAAABAAAHaQAABAAAAAAwc3RzcwAAAAAAAAAIAAAAAQAAAPsAAAH1AAAC7wAAA+kAAATjAAAF3QAABtcAAC84Y3R0cwAAAAAAAAXlAAAAAQAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAADAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAAEAAAIAAAAAAEAAAwAAAAAAQAABAAAAAAHAAAIAAAAAAEAABAAAAAAAgAABAAAAAAHAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAAGAAAIAAAAAAEAAAwAAAAAAQAABAAAAAANAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAABAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAAGAAAIAAAAAAEAAAwAAAAAAQAABAAAAAAFAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAASAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAQAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAADAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAAGAAAIAAAAAAEAAAwAAAAAAQAABAAAAAAFAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAEAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAABAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAA0AAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAwAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAABAAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAcAAAgAAAAAAQAADAAAAAABAAAEAAAAAAUAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAQAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAEAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAADAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAABgAACAAAAAABAAAMAAAAAAEAAAQAAAAABQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAARAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAYAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAQAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAwAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAADAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAAGAAAIAAAAAAEAABAAAAAAAgAABAAAAAAEAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAANAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAIAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAwAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAABgAACAAAAAABAAAQAAAAAAIAAAQAAAAABAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAADAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAYAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAQAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAwAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAADAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAABwAACAAAAAABAAAMAAAAAAEAAAQAAAAABQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAADAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAHHN0c2MAAAAAAAAAAQAAAAEAAAdpAAAAAQAAHbhzdHN6AAAAAAAAAAAAAAdpAAAGRwAAADUAAAAhAAAAJAAAAbwAAABeAAAAIgAAAE0AAAA1AAAAMAAAACQAAAAxAAAAogAAABgAAACGAAAAOQAAAC8AAAAiAAAAbQAAAHYAAAA1AAAAcwAAAD8AAABwAAAAigAAADIAAABOAAAAYgAAADQAAABqAAAAIwAAADsAAABoAAAANAAAAFQAAAA5AAAAUwAAAFcAAABkAAAAXQAAAGMAAAA7AAAARQAAAHsAAABtAAAAYQAAAGQAAABuAAAAZAAAAHkAAAApAAAARQAAAEsAAABmAAAAUQAAAFsAAABSAAAAewAAAF4AAACHAAAAKwAAAFcAAACBAAAALwAAAIAAAABQAAAAPwAAADMAAABSAAAATAAAAGUAAABEAAAAPwAAAGAAAABjAAAAPAAAADoAAABfAAAAXgAAAFsAAAB+AAAATwAAADcAAABVAAAAbQAAAEQAAAA0AAAATgAAAF0AAABsAAAAIgAAAEcAAABtAAAAKAAAASIAAABWAAAAPQAAADcAAACpAAAAOQAAADwAAAAvAAAAvAAAACoAAAB8AAAAPwAAACgAAADAAAAAVgAAAEEAAAAtAAAAOgAAACQAAACSAAAAbgAAAFQAAABLAAAAYQAAAEgAAAA+AAAAOQAAAEQAAAA4AAAAaQAAAIAAAABbAAAALQAAACwAAABmAAAARAAAAHcAAAAyAAAAOQAAACgAAABaAAAAXAAAAE4AAABfAAAAUAAAAEcAAABjAAAANgAAADMAAABkAAAAaQAAAFIAAAA2AAAATwAAAEQAAAA1AAAAMwAAAEsAAAAlAAAAkAAAAJkAAABOAAAATgAAADkAAAB9AAAATgAAAJcAAAA+AAAAOAAAADgAAABsAAAAUQAAACsAAACFAAAAYQAAADkAAAA1AAAAawAAAEkAAACIAAAAeAAAAFAAAABAAAAAOAAAAH0AAABJAAAAqQAAAGUAAABLAAAALgAAAMoAAABZAAAARAAAAD0AAACEAAAAQAAAAIMAAABAAAAAQwAAADUAAAB0AAAAUgAAADwAAAAzAAABCAAAADwAAADNAAAAQwAAAJgAAAAkAAAArgAAAEIAAAA1AAAAigAAAGEAAAA3AAAAMgAAAEIAAAByAAAAIgAAAJ8AAAA4AAAAZwAAAGIAAAA9AAAAVwAAAEgAAAAtAAAAJwAAAG4AAAB+AAAAagAAADwAAAAoAAAAVwAAAEsAAABYAAAAOgAAACwAAABsAAAAZgAAAGoAAABCAAAAWAAAAHEAAABmAAAAcAAAAE0AAAA7AAAAbAAAAFUAAABFAAAAQAAAAFoAAABuAAAANQAABO8AAABjAAAAKQAAAKIAAAArAAAARAAAAIYAAABbAAAARQAAAFYAAACyAAAAQgAAAEkAAAA4AAAAdgAAAFUAAAA8AAAAigAAAHcAAABSAAAARQAAAGYAAAAyAAAAfwAAAQYAAABqAAAAQgAAAD0AAADAAAAAOAAAAF4AAAAyAAAAZgAAACEAAAB8AAAASwAAACoAAACaAAAARgAAADUAAAAvAAAATgAAABsAAABqAAAAZwAAAFsAAABPAAAAXQAAAEoAAAA+AAAAOwAAADwAAABAAAAAaAAAAGwAAAA9AAAAOQAAAHIAAABxAAAATwAAAEsAAACgAAAAMgAAACoAAABlAAAAUQAAAEwAAABiAAAAaAAAAFgAAABUAAAAPAAAADEAAABUAAAAbwAAAFEAAAA+AAAAVQAAAEUAAAAzAAAANAAAAIsAAAAsAAAAKQAAAKEAAABQAAAATwAAAD8AAACJAAAAUAAAAKsAAABIAAAAOwAAADIAAAB0AAAAWgAAAC0AAAB1AAAAUgAAAD8AAAA1AAAAZwAAAEoAAACBAAAAewAAAFEAAABFAAAAPAAAAHoAAABOAAAArgAAAF0AAABKAAAALQAAALoAAABXAAAARAAAADsAAABrAAAAPAAAAMcAAABCAAAAPwAAADYAAABoAAAAUQAAAD8AAAA0AAABBQAAAEIAAAC0AAAAPgAAAIMAAAAdAAAApQAAAD8AAAAkAAAAiAAAAFgAAABGAAAAMQAAADwAAABQAAAAHAAAAF4AAAA4AAAAWQAAAGQAAAA5AAAAVwAAADwAAAAzAAAAOgAAAKAAAABXAAAAigAAADcAAAAoAAAASAAAAEwAAABAAAAAOgAAACwAAABxAAAAdwAAAF4AAABLAAAAaAAAAGcAAABmAAAAdgAAAEoAAAA+AAAAYAAAAFMAAABFAAAAQwAAAFwAAABWAAAAXAAAAE4AAAArAAAALAAAAHgAAAAxAAAASgAAAIoAAABaAAAARAAAAFYAAADDAAAARgAAAEcAAAA4AAAAegAAAFIAAAA/AAAAjQAAAGoAAABUAAAAQgAAAG0AAAAwAAAAigAAAUYAAABWAAAAPQAAAD0AAACvAAAAMwAAAGAAAAA0AAAAYQAAACEAAAByAAAAQQAAACcAAADKAAAAWwAAADoAAAAvAAAARwAAABsAAACSAAAAagAAAGAAAABeAAAAVwAAAE0AAAA3AAAALAAAAGIAAABEAAAAcAAAAGQAAAAxAAAANQAAAIgAAABcAAAARgAAAEYAAACfAAAAMQAAADIAAABRAAAAUgAAAEsAAABiAAAAaQAAAF8AAATrAAAAUgAAADIAAACIAAAAjQAAAFUAAAAzAAAAVAAAAEUAAAA0AAAANQAAAKwAAAAtAAAAKgAAAJQAAABRAAAASgAAADkAAACGAAAASQAAAJwAAABDAAAAOQAAADcAAACMAAAAYgAAAC8AAAB9AAAAeAAAAEQAAAAyAAAAbAAAAEsAAABmAAAAeAAAAE0AAABDAAAAOQAAAHUAAABEAAAAsgAAAF8AAABKAAAALwAAAMEAAABVAAAARgAAADsAAABzAAAAOgAAAJsAAAA7AAAAPwAAADQAAAB0AAAAUQAAADsAAAA1AAAA8QAAADoAAADCAAAAPwAAAHcAAAA6AAAAoQAAAFEAAAA2AAAAiQAAAGAAAAA3AAAAKwAAAEAAAABpAAAAGwAAAH0AAAA5AAAAYAAAAGsAAAA4AAAAPAAAAEQAAAAxAAAAIQAAAHcAAABeAAAAaAAAAEAAAAAsAAAASAAAAGUAAABOAAAAKwAAACsAAACGAAAATAAAACgAAAApAAAAXgAAAE8AAABiAAAAawAAAFsAAAA7AAAAeQAAAFkAAABKAAAAQAAAAFYAAABeAAAAWgAAAC4AAAA4AAAAoAAAAC0AAABbAAAASQAAAH0AAABXAAAAPgAAAFYAAACvAAAAQgAAAEkAAAA3AAAAfQAAAFYAAAA3AAAAoAAAAE8AAABUAAAAQgAAAHUAAAAnAAAAgAAAAFIAAAAzAAAANwAAAR4AAABTAAAAbwAAAC0AAAA+AAAAJAAAAK0AAABBAAAAKgAAACwAAACfAAAAWAAAAD8AAAAsAAAAXwAAACEAAACOAAAAawAAAGcAAABQAAAAagAAAD0AAAAtAAAARgAAAEAAAABNAAAAawAAAGQAAABOAAAAPwAAAIcAAABOAAAAQgAAAEMAAABdAAAAKwAAAHAAAABjAAAAXAAAAE0AAABtAAAATgAAAE0AAABkAAAAMQAAADAAAABaAAAAggAAAFQAAAA8AAAAVAAAAEYAAAA0AAAANgAAAI8AAAAuAAAALAAAAJoAAABLAAAASQAAADsAAACRAAAARwAAAM0AAABDAAAAOQAAADcAAABpAAAAXQAAACoAAAB7AAAAZwAAADUAAAAuAAAAggAAADAAAACHAAAAhwAAAGAAAABIAAAAQAAAAGcAAAA3AAAAtAAAAGYAAABZAAAALAAAALEAAABVAAAARQAAADsAAABoAAAAPAAAAKgAAAA9AAAAQAAAADoAAABrAAAATwAAADsAAAAwAAAA+gAAAEAAAADRAAAAPQAAAHwAAAAoAAAAoAAAAEcAAAAyAAAAWgAAAD0AAAAvAAAE+QAAAEMAAAA2AAAAYgAAAJ8AAAA2AAAAYAAAAHYAAAA9AAAAQAAAAD0AAAAtAAAAJQAAAIAAAABkAAAAcAAAADMAAAAoAAAASgAAAGAAAABJAAAAMgAAADIAAAB/AAAAaQAAAGYAAABQAAAAXQAAAF0AAABqAAAAiAAAAEsAAAA5AAAAaQAAAFEAAABKAAAAQAAAAGAAAABYAAAAXwAAAEcAAAAoAAAAKgAAAIgAAAAiAAAAOwAAAIIAAABVAAAAPgAAAFsAAADQAAAARwAAAEYAAAA6AAAAhQAAAFEAAAA2AAAAjAAAAEsAAABTAAAARgAAAGsAAAAlAAAAgwAAAOgAAABNAAAARwAAAEkAAAC0AAAAPAAAAG0AAAA1AAAAXwAAACUAAACQAAAAOwAAAD0AAAAlAAAAlAAAADcAAACJAAAAJQAAADMAAAAeAAAAYQAAAGsAAABLAAAAVAAAAEkAAAA6AAAALgAAADoAAABHAAAAcwAAAJgAAABhAAAAMgAAADIAAABgAAAAOgAAAK8AAABNAAAARgAAAC0AAABlAAAAVAAAAFQAAAB2AAAAZQAAAF4AAABWAAAAPAAAADgAAABYAAAAggAAAFIAAAA7AAAAVwAAAEsAAAAzAAAANAAAAEsAAAAnAAAAlwAAAJMAAABWAAAATQAAADoAAACIAAAASwAAAJ0AAAA9AAAARAAAADUAAABkAAAATwAAACwAAACGAAAAZgAAAEsAAAAsAAAAbgAAADQAAABdAAAAdwAAAFYAAABQAAAAPAAAAFoAAABKAAAAqwAAAFoAAABJAAAALQAAALEAAABVAAAAQwAAADoAAABmAAAAPwAAAKoAAABAAAAAQwAAADIAAAB8AAAASwAAAD4AAAAvAAABTQAAAD8AAADGAAAARAAAAFoAAAAdAAAAugAAADYAAAAnAAAAeAAAAF4AAAA6AAAALwAAAGAAAABOAAAAIgAAAHYAAAA0AAAAawAAAGYAAAA0AAAAWgAAAEQAAAAvAAAAJAAAAHkAAABZAAAAYQAAADUAAAAwAAAATwAAAGQAAABEAAAAPgAAAEIAAAByAAAAbwAAAG0AAAA4AAAAZgAAAGwAAABgAAAAXwAAAE0AAAA8AAAAUQAAAFkAAABaAAAAQAAAAFgAAABcAAAAVwAAAC8AAAA1AAAAfQAAAC8AAABnAAAASwAAAHEAAABUAAAAOwAAAFUAAAC2AAAASgAAAEgAAAA4AAAAfwAAAFcAAAA0AAAAlwAAAGkAAABJAAAAPgAAAHkAAAAnAAAAeAAAAFAAAAAyAAAANwAAAIoAAAB+AAAAOAAABWYAAABIAAAAIAAAAPUAAABhAAAANwAAACQAAADrAAAATwAAADkAAAAtAAAAtQAAACkAAAAfAAAAjAAAAHMAAAAlAAAAcQAAAEMAAAA/AAAAMAAAAEYAAABAAAAAXQAAAJAAAAA7AAAANgAAAHwAAABVAAAASQAAAD8AAAB+AAAAOQAAACoAAABoAAAAVAAAAFYAAABrAAAAUAAAAE0AAABiAAAAOQAAAB4AAABhAAAAhQAAAEoAAAA2AAAAdQAAAEQAAAA4AAAAMwAAAJQAAAA2AAAAKAAAAHgAAABLAAAAOQAAAJoAAABjAAAASwAAADoAAADBAAAAQgAAAC4AAAA5AAAAbwAAACwAAACTAAAAfgAAAEYAAAA3AAAAtQAAAD4AAAAyAAAAQwAAAH8AAABDAAAANgAAAFAAAACTAAAAagAAAFYAAAAzAAAAzwAAAF4AAAA+AAAAOQAAAJcAAABMAAAASgAAAD0AAACSAAAARQAAADIAAABGAAAAcgAAAEcAAAA1AAAARQAAAREAAABWAAAAXgAAACUAAAEAAAAAPQAAACsAAAAgAAAAmwAAAGEAAAA7AAAALQAAAIIAAAA6AAAAHgAAADIAAABxAAAAVwAAAG4AAAA9AAAAdAAAAEIAAAAuAAAAIQAAAFwAAACVAAAAbQAAADEAAAAmAAAASgAAAGwAAABGAAAAOAAAADEAAACTAAAARwAAADoAAAAkAAAAcwAAAFYAAABnAAAAZwAAAE4AAABIAAAAawAAAF8AAABDAAAAPQAAAF4AAABmAAAAXwAAAEcAAAAmAAAAIwAAAJYAAAAuAAAANwAAAI4AAABeAAAAPAAAAFYAAACzAAAAPwAAAEoAAAA3AAAAhAAAAFUAAAAxAAAAqQAAAGgAAABTAAAAPwAAAJMAAAAnAAAAqQAAAFEAAAA8AAAANwAAAQwAAABNAAAAbwAAACcAAABHAAAALwAAALQAAABHAAAALQAAACQAAACoAAAAWQAAADsAAAAzAAAAdgAAAB8AAABxAAAAdAAAAEgAAAArAAAAeQAAAF0AAABPAAAALAAAAFYAAAA+AAAAhAAAAHQAAAA6AAAAOwAAAHQAAABZAAAASQAAAEYAAACZAAAAOAAAACgAAABfAAAAUwAAAFkAAAB+AAAAWwAAAFUAAABpAAAAQgAAACEAAABeAAAAdAAAAFIAAAA0AAAAbAAAAEsAAAA7AAAANAAAAH0AAAAxAAAAKAAAAHoAAABeAAAAMwAAAJ4AAABmAAAASQAAADoAAAChAAAARQAAACcAAABBAAAAUQAAACsAAABSAAAATwAAAC4AAAWcAAAAiQAAAC8AAACaAAAAXQAAAEoAAAAzAAAAYwAAAFMAAABMAAAAUgAAAKMAAAA/AAAALAAAALAAAABVAAAAOQAAADkAAAByAAAARwAAAD4AAAA7AAAArQAAAFMAAAAzAAAANgAAAQkAAABSAAAAKQAAADsAAADTAAAAOwAAAN4AAAA4AAAAHgAAAC0AAAB9AAAAXAAAADsAAAAuAAAAcwAAADUAAAA0AAAAIQAAAIYAAAAyAAAAVwAAAG8AAABFAAAAXgAAAEAAAAAzAAAALAAAAGQAAACCAAAAggAAADAAAAAjAAAASwAAAGsAAAA/AAAAOwAAAL4AAABWAAAAMQAAADUAAABHAAAAYwAAAFgAAABoAAAAYwAAAEsAAABIAAAAXwAAAF4AAAA/AAAAOAAAAGAAAABUAAAAYgAAAFAAAAAqAAAAKwAAAHUAAAAwAAAAOgAAAI0AAABbAAAAOQAAAFYAAAC+AAAAQQAAAEYAAAA2AAAAgQAAAEwAAAA0AAAAlgAAAHAAAABQAAAAPQAAAHgAAAAvAAAAmwAAAF4AAAA9AAAAMQAAAUUAAABmAAAAeAAAAC0AAAB6AAAAMAAAABwAAAAuAAAA4AAAADUAAAAqAAAARgAAAEkAAABIAAAAJgAAAB0AAACVAAAAbgAAAGwAAAAtAAAAVgAAAD4AAABHAAAALQAAAFkAAAA5AAAAYgAAAIEAAAAyAAAAMQAAAH8AAABVAAAATQAAAEYAAAB7AAAARAAAACYAAABbAAAATQAAAFcAAABsAAAAVgAAAEoAAABsAAAAOgAAACYAAACAAAAAfQAAAE4AAAA/AAAAcQAAAEYAAAA6AAAAMwAAAKcAAAAyAAAAKgAAAH0AAABUAAAANAAAAJwAAABpAAAARQAAADsAAAC9AAAAQQAAAC4AAABGAAAAbQAAACgAAAB6AAAAYgAAAEYAAAAyAAAAgAAAADsAAAAxAAAAPgAAAH0AAABAAAAANQAAAE4AAACWAAAAZAAAAFMAAAAxAAAAxQAAAGEAAAA/AAAANAAAAH0AAABJAAAATAAAAEEAAACuAAAAQAAAADMAAABIAAAAbgAAAEUAAAAyAAAARQAAAPsAAABSAAAAXAAAACYAAADkAAAARwAAADAAAAAgAAAAkQAAAGEAAAA6AAAALQAAAIYAAAA7AAAAHQAAADAAAAB2AAAAUgAAAGkAAAA9AAAAYgAAAEAAAAAuAAAAHgAAAGUAAAB/AAAAcQAAAC0AAAAiAAAAWAAAAHUAAABFAAAAOAAAADIAAACAAAAAPgAAADUAAAAqAAAAZAAAAFkAAABfAAAFYwAAAFIAAAA3AAAAlQAAAGEAAABCAAAANwAAAFkAAABaAAAAdAAAAEIAAAAoAAAAJQAAAJEAAAAwAAAANwAAAJEAAABdAAAAPAAAAFUAAAC3AAAARgAAAEgAAAA2AAAAegAAAE0AAAA2AAAAnAAAAG0AAABTAAAAOwAAAIkAAAA7AAAAhwAAAFEAAAA9AAAAOwAAAVMAAABoAAAAfAAAACsAAACCAAAAMgAAABwAAAAyAAAA2AAAAGUAAAAiAAAAOwAAAHQAAABOAAAAKAAAACIAAACYAAAAewAAAEYAAAApAAAAYwAAAFkAAABBAAAALwAAAFAAAABHAAAAcAAAAHoAAAAvAAAAMQAAAHIAAABfAAAASQAAAEcAAACZAAAANgAAAB4AAABiAAAAUQAAAFcAAAB0AAAAXQAAAFUAAABTAAAAPgAAACAAAABhAAAAgAAAAEsAAAA0AAAAZwAAAEIAAAA4AAAAMAAAAIUAAAAuAAAAIgAAAHgAAABWAAAANAAAAJsAAABrAAAARwAAAD4AAADBAAAAPgAAACwAAABGAAAAbQAAACUAAACDAAAAWQAAAEAAAAAyAAAAvQAAAEsAAABLAAAARwAAAHMAAAA7AAAANQAAAEoAAACXAAAAcwAAAFEAAAAxAAAAvQAAAGkAAAA+AAAAOQAAAI4AAABKAAAASQAAAD0AAACTAAAARQAAADUAAABGAAAAbAAAAEYAAAA2AAAARgAAASEAAABZAAAAZAAAACcAAADjAAAAPwAAACwAAAAgAAAAjgAAAF8AAAA6AAAALwAAAE4AAAAfAAAAhwAAADIAAAAyAAAAVQAAAHMAAAA7AAAATgAAADkAAAAuAAAAJQAAAGMAAACPAAAAdAAAADQAAAAhAAAATQAAAJwAAABLAAAAOQAAADEAAACiAAAARQAAADgAAAAlAAAAcQAAAGcAAABmAAAAgQAAAEoAAAA8AAAAYgAAAGoAAABCAAAAPgAAAGIAAABgAAAAZQAAAEYAAAAqAAAAKQAAAHYAAAA0AAAAPQAAAIgAAABbAAAAOQAAAFUAAADRAAAAPwAAAEQAAAA6AAAAhwAAAE8AAAA1AAAAoQAAAGwAAABSAAAAQwAAAJAAAAAmAAAAuQAAAE8AAAA9AAAAMwAAAOkAAABPAAAAbgAAACwAAACbAAAANwAAACIAAAAyAAAAygAAAD0AAAAiAAAAOgAAAEgAAAA0AAAAKQAAACUAAACDAAAAbAAAAE8AAAAoAAAAYAAAAFEAAABAAAAALQAAAGIAAABEAAAAcAAAAHMAAAA+AAAAOQAAAFQAAABNAAAAPAAABYMAAACNAAAAKwAAAGsAAABqAAAASgAAAFQAAABqAAAAWAAAAGAAAABqAAAAOwAAADcAAABUAAAAcQAAAFUAAAA9AAAAaAAAAEkAAAA6AAAAMwAAAJgAAAA5AAAAIwAAAHwAAABVAAAANAAAAJkAAABlAAAATAAAADwAAAC4AAAAQAAAACoAAABIAAAAawAAACgAAACHAAAAXgAAAD4AAAAxAAAAbwAAADkAAAB7AAAARwAAAD0AAAA0AAAAgwAAAF0AAABOAAAAUwAAAMoAAAA9AAAALAAAANMAAABVAAAAQQAAADoAAABkAAAAUQAAAEQAAAA/AAAApAAAAFIAAAAvAAAAPAAAAREAAABTAAAAKAAAADsAAAD8AAAAOgAAAOQAAAA2AAAAFwAAADQAAACCAAAAXAAAADYAAAAvAAAAdwAAADgAAAA5AAAAJwAAAHYAAAA5AAAAWgAAAIEAAABAAAAATAAAAEQAAAAzAAAAHwAAAIQAAAB0AAAAcAAAADcAAAAlAAAAWgAAAFoAAABOAAAANQAAAC0AAACTAAAAXQAAADsAAAAjAAAAcwAAAFgAAABoAAAAdwAAAEoAAAA6AAAAYQAAAFwAAAA9AAAANgAAAF0AAABbAAAAZgAAAEkAAAApAAAAKQAAAI0AAAAzAAAAPgAAAIsAAABdAAAANQAAAFcAAACxAAAAPQAAAEkAAAA3AAAAeQAAAE4AAAA0AAAAigAAAG8AAABSAAAAPwAAAH4AAAA0AAAAVwAAAF0AAAA5AAAAMQAAABRzdGNvAAAAAAAAAAEAAAAwAAAAYnVkdGEAAABabWV0YQAAAAAAAAAhaGRscgAAAAAAAAAAbWRpcmFwcGwAAAAAAAAAAAAAAAAtaWxzdAAAACWpdG9vAAAAHWRhdGEAAAABAAAAAExhdmY1Ny44My4xMDA=","ok":true,"headers":[["content-type","video/mp4"]],"status":200,"status_text":""}},"base_uri":"https://localhost:8080/","height":501}},"cell_type":"code","source":["play_video('pong_pretrained/0.avi')"],"execution_count":0,"outputs":[{"output_type":"execute_result","data":{"text/html":["\n","    <video width=\"640\" height=\"480\" controls>\n","      <source src=\"/nbextensions/vid.mp4\" type=\"video/mp4\">\n","    </video>\n","  "],"text/plain":["<IPython.core.display.HTML object>"]},"metadata":{"tags":[]},"execution_count":24}]},{"metadata":{"id":"U-SyGcZBCmPn","colab_type":"text"},"cell_type":"markdown","source":["# Train your policy (model-free training)\n","Training model-free on Pong (it takes a few hours):"]},{"metadata":{"id":"WIQazd5aCocc","colab_type":"code","outputId":"0a440c18-affc-4b2a-d6e1-c3cda84465bc","executionInfo":{"status":"ok","timestamp":1553254256733,"user_tz":-60,"elapsed":19957,"user":{"displayName":"Piotr Kozakowski","photoUrl":"","userId":"01014928596539690143"}},"colab":{"base_uri":"https://localhost:8080/","height":1516}},"cell_type":"code","source":["!python -m tensor2tensor.rl.trainer_model_free \\\n","  --hparams_set=rlmf_base \\\n","  --hparams=game=pong \\\n","  --output_dir=mf_pong"],"execution_count":0,"outputs":[{"output_type":"stream","text":["\n","WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n","For more information, please see:\n","  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n","  * https://github.com/tensorflow/addons\n","If you depend on functionality not listed there, please file an issue.\n","\n","2019-03-22 11:30:42.987149: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 2300000000 Hz\n","2019-03-22 11:30:42.987392: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x30323c0 executing computations on platform Host. Devices:\n","2019-03-22 11:30:42.987491: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): <undefined>, <undefined>\n","2019-03-22 11:30:43.082876: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:998] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n","2019-03-22 11:30:43.083442: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x3032100 executing computations on platform CUDA. Devices:\n","2019-03-22 11:30:43.083493: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): Tesla K80, Compute Capability 3.7\n","2019-03-22 11:30:43.083843: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1433] Found device 0 with properties: \n","name: Tesla K80 major: 3 minor: 7 memoryClockRate(GHz): 0.8235\n","pciBusID: 0000:00:04.0\n","totalMemory: 11.17GiB freeMemory: 11.10GiB\n","2019-03-22 11:30:43.083879: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 11:30:43.475526: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 11:30:43.475601: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 11:30:43.475629: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 11:30:43.476026: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:42] Overriding allow_growth setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.\n","2019-03-22 11:30:43.476131: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Colocations handled automatically by placer.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/rl/envs/py_func_batch_env.py:122: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","tf.py_func is deprecated in TF V2. Instead, use\n","    tf.py_function, which takes a python function which manipulates tf eager\n","    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to\n","    an ndarray (just call tensor.numpy()) but having access to eager tensors\n","    means `tf.py_function`s can use accelerators such as GPUs as well as\n","    being differentiable using a gradient tape.\n","    \n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/utils/t2t_model.py:1358: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.cast instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/function.py:1007: calling Graph.create_op (from tensorflow.python.framework.ops) with compute_shapes is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Shapes are always computed; don't use the compute_shapes as it has no effect.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:277: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.cast instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:598: conv2d (from tensorflow.python.layers.convolutional) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.conv2d instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:602: flatten (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.flatten instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:603: dropout (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.dropout instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:604: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.dense instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:2887: multinomial (from tensorflow.python.ops.random_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.random.categorical instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/rl/ppo_learner.py:479: Print (from tensorflow.python.ops.logging_ops) is deprecated and will be removed after 2018-08-20.\n","Instructions for updating:\n","Use tf.print instead of tf.Print. Note that tf.print returns a no-output operator that directly prints the output. Outside of defuns or eager mode, this operator will not be executed unless it is directly specified in session.run or used as a control dependency for other operators. This is only a concern in graph mode. Below is an example of how to ensure tf.print executes in graph mode:\n","```python\n","    sess = tf.Session()\n","    with sess.as_default():\n","        tensor = tf.range(10)\n","        print_op = tf.print(tensor)\n","        with tf.control_dependencies([print_op]):\n","          out = tf.add(tensor, tensor)\n","        sess.run(out)\n","    ```\n","Additionally, to use tf.print in python 2.7, users must make sure to import\n","the following:\n","\n","  `from __future__ import print_function`\n","\n","2019-03-22 11:30:49.903512: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 11:30:49.903591: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 11:30:49.903620: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 11:30:49.903639: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 11:30:49.903898: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/saver.py:1266: checkpoint_exists (from tensorflow.python.training.checkpoint_management) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use standard file APIs to check for files with this prefix.\n","2019-03-22 11:30:51.335217: I tensorflow/stream_executor/dso_loader.cc:152] successfully opened CUDA library libcublas.so.10.0 locally\n","mean_score: [0][0][0]\n","^C\n"],"name":"stdout"}]},{"metadata":{"id":"FbSjwVAtCvLY","colab_type":"text"},"cell_type":"markdown","source":["Hyperparameter sets are defined in `tensor2tensor/models/research/rl.py`. You can override them using the hparams flag, e.g.\n","\n","```\n","--hparams=game=kung_fu_master,frame_stack_size=5\n","```\n","\n","As in model-based training, the periodic evaluation runs with timestep limit of 1000. To do full evaluation after training, run:"]},{"metadata":{"id":"jppi4FE5C2nB","colab_type":"code","outputId":"a10afb7c-edd6-4a93-eee4-e3876977e825","executionInfo":{"status":"ok","timestamp":1553254412202,"user_tz":-60,"elapsed":15104,"user":{"displayName":"Piotr Kozakowski","photoUrl":"","userId":"01014928596539690143"}},"colab":{"base_uri":"https://localhost:8080/","height":4083}},"cell_type":"code","source":["!python -m tensor2tensor.rl.evaluator \\\n","  --loop_hparams_set=rlmf_tiny \\\n","  --hparams=game=pong \\\n","  --policy_dir=mf_pong \\\n","  --debug_video_path=mf_pong \\\n","  --num_debug_videos=4 \\\n","  --eval_metrics_dir=mf_pong/full_eval_metrics"],"execution_count":0,"outputs":[{"output_type":"stream","text":["\n","WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n","For more information, please see:\n","  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n","  * https://github.com/tensorflow/addons\n","If you depend on functionality not listed there, please file an issue.\n","\n","INFO:tensorflow:Overriding hparams in rlmf_tiny with game=pong,eval_max_num_noops=0,eval_sampling_temps=[0.5]\n","INFO:tensorflow:Evaluating metric mean_reward/eval/sampling_temp_0.5_max_noops_0_unclipped\n","2019-03-22 11:33:23.214052: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 2300000000 Hz\n","2019-03-22 11:33:23.214294: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x2d07020 executing computations on platform Host. Devices:\n","2019-03-22 11:33:23.214335: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): <undefined>, <undefined>\n","2019-03-22 11:33:23.309948: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:998] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n","2019-03-22 11:33:23.310546: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x2d067e0 executing computations on platform CUDA. Devices:\n","2019-03-22 11:33:23.310585: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): Tesla K80, Compute Capability 3.7\n","2019-03-22 11:33:23.310991: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1433] Found device 0 with properties: \n","name: Tesla K80 major: 3 minor: 7 memoryClockRate(GHz): 0.8235\n","pciBusID: 0000:00:04.0\n","totalMemory: 11.17GiB freeMemory: 11.10GiB\n","2019-03-22 11:33:23.311027: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 11:33:23.707039: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 11:33:23.707114: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 11:33:23.707139: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 11:33:23.707459: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:42] Overriding allow_growth setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.\n","2019-03-22 11:33:23.707523: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","INFO:tensorflow:Using DummyPolicyProblem for the policy.\n","INFO:tensorflow:Setting T2TModel mode to 'train'\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Colocations handled automatically by placer.\n","INFO:tensorflow:Using variable initializer: orthogonal\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/utils/t2t_model.py:1358: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.cast instead.\n","INFO:tensorflow:Transforming feature 'input_action' with symbol_modality_6_64.bottom\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/function.py:1007: calling Graph.create_op (from tensorflow.python.framework.ops) with compute_shapes is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Shapes are always computed; don't use the compute_shapes as it has no effect.\n","INFO:tensorflow:Transforming feature 'input_reward' with symbol_modality_3_64.bottom\n","INFO:tensorflow:Transforming feature 'inputs' with video_modality.bottom\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_video.py:495: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","tf.py_func is deprecated in TF V2. Instead, use\n","    tf.py_function, which takes a python function which manipulates tf eager\n","    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to\n","    an ndarray (just call tensor.numpy()) but having access to eager tensors\n","    means `tf.py_function`s can use accelerators such as GPUs as well as\n","    being differentiable using a gradient tape.\n","    \n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:277: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.cast instead.\n","INFO:tensorflow:Transforming feature 'target_action' with symbol_modality_6_64.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_policy' with identity_modality.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_reward' with symbol_modality_3_64.targets_bottom\n","INFO:tensorflow:Transforming feature 'target_value' with identity_modality.targets_bottom\n","INFO:tensorflow:Transforming feature 'targets' with video_modality.targets_bottom\n","INFO:tensorflow:Building model body\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:598: conv2d (from tensorflow.python.layers.convolutional) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.conv2d instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:602: flatten (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.flatten instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:603: dropout (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.dropout instead.\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:604: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use keras.layers.dense instead.\n","INFO:tensorflow:Transforming body output with identity_modality.top\n","INFO:tensorflow:Transforming body output with identity_modality.top\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:2887: multinomial (from tensorflow.python.ops.random_ops) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use tf.random.categorical instead.\n","2019-03-22 11:33:24.564271: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n","2019-03-22 11:33:24.564350: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n","2019-03-22 11:33:24.564376: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n","2019-03-22 11:33:24.564410: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n","2019-03-22 11:33:24.564687: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n","INFO:tensorflow:Restoring checkpoint mf_pong/model.ckpt-9\n","WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/saver.py:1266: checkpoint_exists (from tensorflow.python.training.checkpoint_management) is deprecated and will be removed in a future version.\n","Instructions for updating:\n","Use standard file APIs to check for files with this prefix.\n","INFO:tensorflow:Restoring parameters from mf_pong/model.ckpt-9\n","2019-03-22 11:33:24.985295: I tensorflow/stream_executor/dso_loader.cc:152] successfully opened CUDA library libcublas.so.10.0 locally\n","INFO:tensorflow:Step 5, mean_score: 0.000000\n","INFO:tensorflow:Step 10, mean_score: 0.000000\n","INFO:tensorflow:Step 15, mean_score: 0.000000\n","INFO:tensorflow:Step 20, mean_score: 0.000000\n","INFO:tensorflow:Step 25, mean_score: 0.000000\n","INFO:tensorflow:Step 30, mean_score: 0.000000\n","INFO:tensorflow:Step 35, mean_score: 0.000000\n","INFO:tensorflow:Step 40, mean_score: 0.000000\n","INFO:tensorflow:Step 45, mean_score: 0.000000\n","INFO:tensorflow:Step 50, mean_score: 0.000000\n","INFO:tensorflow:Step 55, mean_score: 0.000000\n","INFO:tensorflow:Step 60, mean_score: 0.000000\n","INFO:tensorflow:Step 65, mean_score: -1.000000\n","INFO:tensorflow:Step 70, mean_score: -1.000000\n","INFO:tensorflow:Step 75, mean_score: -1.000000\n","INFO:tensorflow:Step 80, mean_score: -1.000000\n","INFO:tensorflow:Step 85, mean_score: -1.000000\n","INFO:tensorflow:Step 90, mean_score: -1.000000\n","INFO:tensorflow:Step 95, mean_score: -1.000000\n","INFO:tensorflow:Step 100, mean_score: -2.000000\n","INFO:tensorflow:Step 105, mean_score: -2.000000\n","INFO:tensorflow:Step 110, mean_score: -2.000000\n","INFO:tensorflow:Step 115, mean_score: -2.000000\n","INFO:tensorflow:Step 120, mean_score: -2.000000\n","INFO:tensorflow:Step 125, mean_score: -2.000000\n","INFO:tensorflow:Step 130, mean_score: -2.000000\n","INFO:tensorflow:Step 135, mean_score: -3.000000\n","INFO:tensorflow:Step 140, mean_score: -3.000000\n","INFO:tensorflow:Step 145, mean_score: -3.000000\n","INFO:tensorflow:Step 150, mean_score: -3.000000\n","INFO:tensorflow:Step 155, mean_score: -3.000000\n","INFO:tensorflow:Step 160, mean_score: -3.000000\n","INFO:tensorflow:Step 165, mean_score: -3.000000\n","INFO:tensorflow:Step 170, mean_score: -4.000000\n","INFO:tensorflow:Step 175, mean_score: -4.000000\n","INFO:tensorflow:Step 180, mean_score: -4.000000\n","INFO:tensorflow:Step 185, mean_score: -4.000000\n","INFO:tensorflow:Step 190, mean_score: -4.000000\n","INFO:tensorflow:Step 195, mean_score: -4.000000\n","INFO:tensorflow:Step 200, mean_score: -4.000000\n","INFO:tensorflow:Step 205, mean_score: -5.000000\n","INFO:tensorflow:Step 210, mean_score: -5.000000\n","INFO:tensorflow:Step 215, mean_score: -5.000000\n","INFO:tensorflow:Step 220, mean_score: -5.000000\n","INFO:tensorflow:Step 225, mean_score: -5.000000\n","INFO:tensorflow:Step 230, mean_score: -5.000000\n","INFO:tensorflow:Step 235, mean_score: -5.000000\n","INFO:tensorflow:Step 240, mean_score: -6.000000\n","INFO:tensorflow:Step 245, mean_score: -6.000000\n","INFO:tensorflow:Step 250, mean_score: -6.000000\n","INFO:tensorflow:Step 255, mean_score: -6.000000\n","INFO:tensorflow:Step 260, mean_score: -6.000000\n","INFO:tensorflow:Step 265, mean_score: -6.000000\n","INFO:tensorflow:Step 270, mean_score: -6.000000\n","INFO:tensorflow:Step 275, mean_score: -7.000000\n","INFO:tensorflow:Step 280, mean_score: -7.000000\n","INFO:tensorflow:Step 285, mean_score: -7.000000\n","INFO:tensorflow:Step 290, mean_score: -7.000000\n","INFO:tensorflow:Step 295, mean_score: -7.000000\n","INFO:tensorflow:Step 300, mean_score: -7.000000\n","INFO:tensorflow:Step 305, mean_score: -7.000000\n","INFO:tensorflow:Step 310, mean_score: -8.000000\n","INFO:tensorflow:Step 315, mean_score: -8.000000\n","INFO:tensorflow:Step 320, mean_score: -8.000000\n","INFO:tensorflow:Step 325, mean_score: -8.000000\n","INFO:tensorflow:Step 330, mean_score: -8.000000\n","INFO:tensorflow:Step 335, mean_score: -8.000000\n","INFO:tensorflow:Step 340, mean_score: -8.000000\n","INFO:tensorflow:Step 345, mean_score: -9.000000\n","INFO:tensorflow:Step 350, mean_score: -9.000000\n","INFO:tensorflow:Step 355, mean_score: -9.000000\n","INFO:tensorflow:Step 360, mean_score: -9.000000\n","INFO:tensorflow:Step 365, mean_score: -9.000000\n","INFO:tensorflow:Step 370, mean_score: -9.000000\n","INFO:tensorflow:Step 375, mean_score: -9.000000\n","INFO:tensorflow:Step 380, mean_score: -10.000000\n","INFO:tensorflow:Step 385, mean_score: -10.000000\n","INFO:tensorflow:Step 390, mean_score: -10.000000\n","INFO:tensorflow:Step 395, mean_score: -10.000000\n","INFO:tensorflow:Step 400, mean_score: -10.000000\n","INFO:tensorflow:Step 405, mean_score: -10.000000\n","INFO:tensorflow:Step 410, mean_score: -10.000000\n","INFO:tensorflow:Step 415, mean_score: -11.000000\n","INFO:tensorflow:Step 420, mean_score: -11.000000\n","INFO:tensorflow:Step 425, mean_score: -11.000000\n","INFO:tensorflow:Step 430, mean_score: -11.000000\n","INFO:tensorflow:Step 435, mean_score: -11.000000\n","INFO:tensorflow:Step 440, mean_score: -11.000000\n","INFO:tensorflow:Step 445, mean_score: -11.000000\n","INFO:tensorflow:Step 450, mean_score: -12.000000\n","INFO:tensorflow:Step 455, mean_score: -12.000000\n","INFO:tensorflow:Step 460, mean_score: -12.000000\n","INFO:tensorflow:Step 465, mean_score: -12.000000\n","INFO:tensorflow:Step 470, mean_score: -12.000000\n","INFO:tensorflow:Step 475, mean_score: -12.000000\n","INFO:tensorflow:Step 480, mean_score: -12.000000\n","INFO:tensorflow:Step 485, mean_score: -13.000000\n","INFO:tensorflow:Step 490, mean_score: -13.000000\n","INFO:tensorflow:Step 495, mean_score: -13.000000\n","INFO:tensorflow:Step 500, mean_score: -13.000000\n","INFO:tensorflow:Step 505, mean_score: -13.000000\n","INFO:tensorflow:Step 510, mean_score: -13.000000\n","INFO:tensorflow:Step 515, mean_score: -13.000000\n","INFO:tensorflow:Step 520, mean_score: -14.000000\n","INFO:tensorflow:Step 525, mean_score: -14.000000\n","INFO:tensorflow:Step 530, mean_score: -14.000000\n","INFO:tensorflow:Step 535, mean_score: -14.000000\n","INFO:tensorflow:Step 540, mean_score: -14.000000\n","INFO:tensorflow:Step 545, mean_score: -14.000000\n","INFO:tensorflow:Step 550, mean_score: -14.000000\n","INFO:tensorflow:Step 555, mean_score: -15.000000\n","INFO:tensorflow:Step 560, mean_score: -15.000000\n","INFO:tensorflow:Step 565, mean_score: -15.000000\n","INFO:tensorflow:Step 570, mean_score: -15.000000\n","INFO:tensorflow:Step 575, mean_score: -15.000000\n","INFO:tensorflow:Step 580, mean_score: -15.000000\n","INFO:tensorflow:Step 585, mean_score: -15.000000\n","INFO:tensorflow:Step 590, mean_score: -16.000000\n","INFO:tensorflow:Step 595, mean_score: -16.000000\n","INFO:tensorflow:Step 600, mean_score: -16.000000\n","INFO:tensorflow:Step 605, mean_score: -16.000000\n","INFO:tensorflow:Step 610, mean_score: -16.000000\n","INFO:tensorflow:Step 615, mean_score: -16.000000\n","INFO:tensorflow:Step 620, mean_score: -16.000000\n","INFO:tensorflow:Step 625, mean_score: -17.000000\n","INFO:tensorflow:Step 630, mean_score: -17.000000\n","INFO:tensorflow:Step 635, mean_score: -17.000000\n","INFO:tensorflow:Step 640, mean_score: -17.000000\n","INFO:tensorflow:Step 645, mean_score: -17.000000\n","INFO:tensorflow:Step 650, mean_score: -17.000000\n","INFO:tensorflow:Step 655, mean_score: -17.000000\n","INFO:tensorflow:Step 660, mean_score: -18.000000\n","INFO:tensorflow:Step 665, mean_score: -18.000000\n","INFO:tensorflow:Step 670, mean_score: -18.000000\n","INFO:tensorflow:Step 675, mean_score: -18.000000\n","INFO:tensorflow:Step 680, mean_score: -18.000000\n","INFO:tensorflow:Step 685, mean_score: -18.000000\n","INFO:tensorflow:Step 690, mean_score: -18.000000\n","INFO:tensorflow:Step 695, mean_score: -19.000000\n","INFO:tensorflow:Step 700, mean_score: -19.000000\n","INFO:tensorflow:Step 705, mean_score: -19.000000\n","INFO:tensorflow:Step 710, mean_score: -19.000000\n","INFO:tensorflow:Step 715, mean_score: -19.000000\n","INFO:tensorflow:Step 720, mean_score: -19.000000\n","INFO:tensorflow:Step 725, mean_score: -19.000000\n","INFO:tensorflow:Step 730, mean_score: -20.000000\n","INFO:tensorflow:Step 735, mean_score: -20.000000\n","INFO:tensorflow:Step 740, mean_score: -20.000000\n","INFO:tensorflow:Step 745, mean_score: -20.000000\n","INFO:tensorflow:Step 750, mean_score: -20.000000\n","INFO:tensorflow:Step 755, mean_score: -20.000000\n","INFO:tensorflow:Step 760, mean_score: -20.000000\n"],"name":"stdout"}]},{"metadata":{"id":"mDoR0C0ZKCOn","colab_type":"code","outputId":"aba41a4d-2957-4ea0-d511-eae7ea4e238e","executionInfo":{"status":"ok","timestamp":1553254513355,"user_tz":-60,"elapsed":3908,"user":{"displayName":"Piotr Kozakowski","photoUrl":"","userId":"01014928596539690143"}},"colab":{"resources":{"http://localhost:8080/nbextensions/vid.mp4":{"data":"AAAAIGZ0eXBpc29tAAACAGlzb21pc28yYXZjMW1wNDEAAAAIZnJlZQAA6u9tZGF0AAACrgYF//+q3EXpvebZSLeWLNgg2SPu73gyNjQgLSBjb3JlIDE1MiByMjg1NCBlOWE1OTAzIC0gSC4yNjQvTVBFRy00IEFWQyBjb2RlYyAtIENvcHlsZWZ0IDIwMDMtMjAxNyAtIGh0dHA6Ly93d3cudmlkZW9sYW4ub3JnL3gyNjQuaHRtbCAtIG9wdGlvbnM6IGNhYmFjPTEgcmVmPTMgZGVibG9jaz0xOjA6MCBhbmFseXNlPTB4MzoweDExMyBtZT1oZXggc3VibWU9NyBwc3k9MSBwc3lfcmQ9MS4wMDowLjAwIG1peGVkX3JlZj0xIG1lX3JhbmdlPTE2IGNocm9tYV9tZT0xIHRyZWxsaXM9MSA4eDhkY3Q9MSBjcW09MCBkZWFkem9uZT0yMSwxMSBmYXN0X3Bza2lwPTEgY2hyb21hX3FwX29mZnNldD0tMiB0aHJlYWRzPTMgbG9va2FoZWFkX3RocmVhZHM9MSBzbGljZWRfdGhyZWFkcz0wIG5yPTAgZGVjaW1hdGU9MSBpbnRlcmxhY2VkPTAgYmx1cmF5X2NvbXBhdD0wIGNvbnN0cmFpbmVkX2ludHJhPTAgYmZyYW1lcz0zIGJfcHlyYW1pZD0yIGJfYWRhcHQ9MSBiX2JpYXM9MCBkaXJlY3Q9MSB3ZWlnaHRiPTEgb3Blbl9nb3A9MCB3ZWlnaHRwPTIga2V5aW50PTI1MCBrZXlpbnRfbWluPTEwIHNjZW5lY3V0PTQwIGludHJhX3JlZnJlc2g9MCByY19sb29rYWhlYWQ9NDAgcmM9Y3JmIG1idHJlZT0xIGNyZj0yMy4wIHFjb21wPTAuNjAgcXBtaW49MCBxcG1heD02OSBxcHN0ZXA9NCBpcF9yYXRpbz0xLjQwIGFxPTE6MS4wMACAAAADkWWIhABvrNdXNvEPmO7lwVl73sPl0EDBzzvrz1O9Sgfa49FGnVhGNj4PrUzIEjAsiR14q5boH034au6fMfeHzW8BQIdLu5D8GWFcvhnUQvMLIDm/5fDlJWNI1pLZ0KekyKgRZvEg10IZZePvLcj64kGJzCMbJi6QZbX4WMzyM/ZwsXoWWPBmmlKBzFixHWdkptjcAYhpgDpXSILlIffpBFr5Fmv8Xdrl5eZtB/U18q6RE0tX2BrhekKyOZ5lJWnXZWIEICkLYIda8x0l/aAug9zkJAN2UJ5v8AfQgXgS7iPy41I11UQneH59QQ6r2Fy+bXVz7hKXvFUUQUW2NfwyAHSubAtKRV8FgrIBnKXwxAjc8zc/00LsdZVdehIaL1eI9qZtyap5GmVpF7ZJdkQbo7j2k9/o8Ztr6lwZrODqoujHSJK6V9bK0u9Et564zU+wWgftergJVAEl4m/D3N6/lD6Tni/a6bLzIcdcVjnfWLPAUBwAoj19NpxhAbe1VyiybbzF11k65OpExrnTpeyfXnWi2YKXmv6NMcvP6YS8WOK4pM7nWhyKetjJvO69p10oeh7Pv3PuQBq3kARIBKQ+MPYmymnbhgmxG/6w3hJ2A2Urz2k1DctVq7TiUCWnAReHSDqSpYcdQwxCm/lIpIwtl/dffgss5v+hhFs6NSNe3zqLc+wa/P6fKKBzHBPA6mZtXbiaJH0Y+5hMHtf92lFc+6I4pZ1q2XpI5Nr1V7em9lfehnp6KwZCFUTCrCle3ZgVn3/WlL0hiX3HqF/qGx1rSRBE7lqG2nGEQXx7BFJGNLF0vFVi0j2agV+lVqGOVlIxAjK3E9wGWVM0V7xAFGXQtxAYJ6qA6zMuM1AzlTqoEWcy+zkYm6Z2/Vn8RMHtpHaCW7GF05Wujcn0D05dR11MQem50GDiKxlzGighyGKWmfeex/qNBXelV1apol1nwDCUiSbC9fPUu70YI94kit+OnCdHe588u+9o5tqmvG+4ju2D1U0YtGzBJLwNtKIxTj+ycim3c7lWMz9gyNpdcRw85nQOO+UebN2j6KDuTy4XNxidtzFIcvo67EYGfl+q3WaPfQzFQLuOqvybvDRViyMxNwUidf7UCNcjzUMa0RFtd4HPTD4pR9pL0oOHG0XMOwvMlfvI/0tUl0nH8gKjxa10D+0pCAG9Sq76K3xNRb2QQ4PhDp7u3P+U7CY7JpR9qHasfUEAAAAxQZojbEb/+8M7YXuyASh7Kplen5y7UyO14JIrbok4XTbRQe3ORruR41lvOoDou8ClgAAAAB1BnkF4m/+AoGMvgzynj7iP1aLLDUqgN6Jo/suYQQAAACABnmJqRH+I4ID6PupU5REY3sf/aa003+ohQ9ie2tSAgAAAAa9BmmZJqEFomUwI3/3q7CUD4uvUkflAlpSidF3UDZUIJZsuBftR/Ot+q3GhwTn9egJDu/Q0u302gbAsB/IV/AFE9llDfK6lUW9694v3+SwMl6mllWP/WwEuWwr3bvYMZvmSan6TfqkNjfFQ2gIUyqpE9/WshNO04YsB0gSlJqYcVRMCqNuW9LOcW33er/NiW98OL0en+UGHVMigjfMRwtvQwEBx2TCEvqNHiHUZK47Ql9EBEwOC+9RX+RfTR+dz/gd9jggR1VC/W7S6VbwMD6OLZQ6pAbUGRTytAu/D3nUIHJpw2KDyU/GpGdKbj2WQLEoSWl/G1erWEOohpWvfDxHkGcHJyES1MZ1DvjVLlpUz1LfdDNETW9u5oytpQuaaon1mmF0VFggCZuK84NA/jjnYLjbiGa2mQVNuCz0xJcY1SlKU0HHZWh4xyJUcesClYPnBpMFzn1F47bfO2af9hq6jk1pG03lTPGyovR1gotc0zYA9fMfxEfYVjeBaSNZ7jIelcTidTkL9rr66l4XIEEMH3QknOAml0XZzRLxPdNkS67lJssdtIGefZbu54FuzpwAAADRBnoRFESzfiNYKja/0xDhfTAiTpHLhMdYxhLBGYN2yGAaE7v0DPPtKPN4N0Bh9SMs1vQXBAAAAIgGepWpEf4EHlYJarHuonGoEjcLs+tCeMaGFsuILlYL+uWEAAABrQZqqSahBbJlMCN/oMK3pi6Gt530QeupQ59ezKe10FpZIms3to+JEJWnpAJhBGOAA9pffTpSVzrDTkYei+dF/0XoAHErlseRXMtx8drWM8lgtra4QnL5SnQQH5ZVFNtwaHUoOeGKot2pc/HkAAAA2QZ7IRRUsn36T35XM0R9+0te5g6nnf972YKyQEsNo7mvZhGpibk4k6GYi5Xw1pjTrITOOCgNgAAAALwGe53REf4pxDGIcj/61i91Rdmi2jQszTK6GJmg9LxPIKqXtoMSsvmXH+U1e8ZQgAAAAJAGe6WpEf4Z1uDH/dTIseAsGQv4aVaLug7WXzVJtvrBkNRseHwAAADNBmutJqEFsmUwI3/QsGvLB67IEGcjaf9jQmmp2xwadCZqIc+QsmxjXnMZRrdlnOWXkniAAAAC0QZsNSeEKUmUwUVLG//1hfa7H0q8GiMhKKARrmw6Z4xdeyr7ujJme+WFMGuspP/+ISBWO3g4rVjHKMZczL9N457VGygcsdj7HL9XSQg7I9GyDcuezylwimiV/wFo/tnsP85jKo1N9u1U/c32ZRoJ6MLzOYkUW9bgpdQwc/vwurbPO6DqdSp7R4DyF3mA4Y/IOgh69Sz6dcgXGl/wwDmRo72yeNgYcXZp2lUk8XqHfEbj36uBFAAAAHwGfLGpEf4jgOwq2BPvs/N5sWutoQxXrlD0LcAJHkzEAAABWQZsxSeEOiZTAjf/7wc1PkbwRgr+fjnyNR8pfsXICQZIWSD29LO9OdoFeaCGC7kNQcI6QVUvVoNv//1awF0I8OAUiwMZgFzApKdTCW820SqjcXXuKH/0AAAAtQZ9PRRU8n31uIe20/lZp7IOycKeMuVxD421aM+Nt0XYWYsoR6FTF8OFoiDbtAAAAIQGfbnREf4ew+VpnMnXvUfVz1z4//JFTxQ6bZXAMaAhfWAAAABgBn3BqRH+E1CtCDqsru18CnQc9Zgwu5C0AAABbQZtySahBaJlMCN/8A8MxwHsdgYxkG6fD5jfkGVPerg1XibHNvThF43JXYYlxb2Qoy9V2WnWwkQ/OqnQwi/YBBYBEZaAn0hrUVg4yiuTQbTao5OXYZIcHYrYRQQAAAFRBm5RJ4QpSZTBREsb/+8P3/LU5AR1llSwpTA2p+sot/Ap2Yo/gxiVR9YF8xLoAHZjM8457AqBIx7NgPEKnQG3AmhGkbhUoZZDnZk8SW2DttxUzXAwAAAAmAZ+zakR/h60AWug79Ot7YB0Hk7/PuwF00L7Iy+IMt5Tq30K+2UAAAABTQZu1SeEOiZTAjf/9dtYVwOWMQEFcGvjV4Syzg11kPr1s9125wYmt9rvuf3mx7RYJ6jF6ErfUzdz0/3v9ae/ANOnW4wjLgPCIol0QSIsPYy4da+kAAAAwQZvWSeEPJlMCN//9fMnM4gfdWuAOqRQ2CfzqA5eQQ8IrF1fJShsXtjQZPj/+EEk+AAAAWUGb90nhDyZTAjf/++FR2EKSC1to8pc6YSO3kyM61vFVTIqwBRLCgvhTW2uELbv4P1t6j3j6R5JEnH93R+6ccuNOfX9oCM+dXHFLeM1xbdq83XfyShy+a/CBAAAAXUGaGUnhDyZTBRE8b/vgOOB3OQFcxg/ZQD7brmyDZG+0lanuFt5eWNnI+DlTWOnKhAAySzB0u8CurDZF8oIWX6RlKX6AlFm3XANEhpl6OfqvPFvE4O+fARgP5ChhkQAAACgBnjhqRH+HsRNxvly6e2k2NMUXC/vU7UPeiHPyDV2qBx4cRcnapKmEAAAAOkGaOknhDyZTAjf/++A44Ilb6dcZ4EfRuTLTuDvo3J2im59tzroCzqfsO8SsZ8WggmVkE0/TN30peYEAAABDQZpcSeEPJlMFETxv/AOiKyBDYvuQQEynViyLAUY1Ib85NXoM2W5LW9IU2aIRq+bew64gtHymArgzox7IataLMEs58QAAACsBnntqRH+HsR7W7gx/uLfAQSDqjarp3zTZOCfEh+6Eb54T4PDy1sS+MHUHAAAAREGaf0nhDyZTAjf/++s+OA7zprt9FalqM30xdIjlcH2DZkkFtbEKP2KnwPRDHCudTxGL9Qj/prFtXdQZd/8kyvQgWVbBAAAAGkGenUURPN99TaObeIE6LTTxi495G3R7GPmQAAAAMgGevmpEf4k5No92pOx4d3tVe/jt/92QHqWDJ1gLa+IuwyWHPUwqtoP8fV9Ym9imBVegAAAAVEGaoUmoQWiZTBTxv/vhUdghqkJVKiFoDxnkwJdI5EfJ1rXs+3IWgFyiLClNqrjfdF5LRESoN+/5vYqeA8+ZDoM0LHM+p62tsd2sj1HySk6QgTd2RQAAAB4BnsBqRH+JNh+VxVwUJ+Y3lV+ryJZRr2sjHdcr6hMAAABLQZrDSeEKUmUwUsb/++A44DvMfbctqjDdsKqNf0xOA39UYyLAzYIZyyzAIzc51s3jIS5ItuDvEtTcwkuqq/36WLfiOMPxF1mVVotBAAAANwGe4mpEf4EPa4byEQdMkh18OfNdNbI0980SSbNKs1ZKN+RKPSfGS6ySQe8jTiCjRi2a1L2bI0AAAABTQZrkSeEOiZTAjf/74VHYIlYiKMdyQZmvfSW5HcZkgM1RFiY3gEH71z/3fVBzC8liqsmHHacePA4zlFgRGs62G32VHnZQ/DXMUNXZSe6AM+37MoEAAABCQZsFSeEPJlMCN//7w+5d2niwPbxnXcPGjmeLv2gdbVqjCSUoEJNhCpLEDmpftWkeZV+eoSYNoxZWdF8ucDxxQt1tAAAAUUGbJknhDyZTAjf/+8OhtVOQEg0JinGp/3aOKBnRKunbXz9YVqPZNGUQc6qywYwzQClzX3GGGmt+n+j6fKijedTiog/HS3groeaf1ZLhN7MSZwAAAEFBm0dJ4Q8mUwI3//vEgdmkZAQXhaU5Ka/mwO9DUeQ8nbhvyYsA1HsjVV69Yl5uJ7Z+5cENl5oKK+xgTk4zv6uteQAAAD1Bm2lJ4Q8mUwURPG/74DjgQlevXVZC8R8NNBU5lXz/5evdLx/xJ2R3nKEWhnQTJ2MhCGZTRB2wGJ2TPpnAAAAAMQGfiGpEf4exElW+Y0oRl/FZgSZ/zHaDUN2en4eI/45oExvCn0eIC8T7xwpFCCD22pAAAAA0QZuKSeEPJlMCN//74DjgJBslHpy/fK99xLLgEoSJoPeLiwa05OsKQgsTdDRtpBJScyBfZQAAAFtBm6tJ4Q8mUwI3//vBcoj4CbBCQT2RExz0aNVXfJZgEMRjG13pvbjo3lZ6TLZP87/4osMgfYBX228XW1pfmU8k2l4pXWkF2ONrCPtLhUtLA/Hya2TerrONeACzAAAAUEGbzEnhDyZTAjf/++A9RB901G+fHcU/WnLCfcsk74GCTFIVwJ+BLzEBzmOXsONP0vNgTAP2KwhTXJSx3MerL/qd/oYBMUklPzb3D+iSqTSoAAAAQEGb7UnhDyZTAjf/++A9RCB74KMJGvAkjJTldXY5HMHa5jfjN8NHs9hBd2fEIGPYdqirJKwimKLpiuNnGPY71cEAAABDQZoOSeEPJlMCN//74DjgQ+tokgudjQu5NRNWobxjhH6Kh3ypTnRXjwJga7RHFZxm4LSd8nQrtx1GfqIwsNleIlTJwQAAAF1Bmi9J4Q8mUwI3//vBnCJ2LA71GF25ddxrRcaDbZ6JrFKdXAbBcwenVA97XR386r26bMwGidvp2PmbPeLJeVTrLs40YO/0dkENxamY3tNtYwvFfXlRKOo9yyo8A4EAAABRQZpQSeEPJlMCN//7w3v8sjICMILXR5ECQLilrI3AMHhNoAl8fjML6XnsQAd3GcNYf/0Vg/WkUZMoFlfh0iQSiC4/lahyE8PO4Mqsxu0If+cpAAAASkGac0nhDyZTAjf//APDMcB3mPxGGZH0fTP3ycONiP81zu+xkY8BF29QMHZMJZp5xyqLnicMHS675xoi5qplhfLQsVW1n9dVmaldAAAAHkGekUURPN99TaObeIE82Q23ingUfeUAohKj+s7K1QAAADgBnrJqRH+HsRJVgDz377tnoNimjhTZvPY9yHf2fmv0wsOMF+wp88Zp9A47UCY9LcsoRhIHvx+VqgAAAD9BmrRJqEFomUwI3/vgHICtV4/N+3TLBHjt+Fcn7DWjz8soUTNhnfkqkcMHOJu59OaFxhuVf+ti1/1ybHdmF1AAAABfQZrVSeEKUmUwI3/7wf8MHhAR1o8LZVW6pfXYDg2dpGT5Rvg+Cl2lQj+IQxqhhQEfFhAQwlKUGoZJ0yY1+WX5L6BuA3zmdozX8D1TAHaS+xXRLIKO1jwkjAd/JNceRMEAAABNQZr2SeEOiZTAjf/8ESF/UQffCua7E+r+NL8HxgRR3d+rdy5oOmnMWV4zw6CD9Z7mbT2yUeu59zFDbP46SRq9klvWZ1ZbOSW+2vZf9fQAAABHQZsXSeEPJlMCN//74EfgEPLRNeKMrWaAkB9b5moj+s74rojfFold/VwZwZnjRlAT49c0fDe4ZV3UBIGBK9CK3ytLN0Y42oEAAABDQZs4SeEPJlMCN//74DjgOFPapEipwNQ5JLf/iFkT7vTC27zO0O1Hhr7Be9x5uGF5a9SD1TJTW/IfLAT7Nh7DJMuMeQAAAGpBm1lJ4Q8mUwI3//vBm7tPFge4qVrTYzB1Rf4j8SwaiUd3Vn91V8IbddfFnSdsq/LMm7+bcOxpHHp/+sF3L+oZtRUH+68r3BQroy9k+62TGFClrSWs36FP/9yav//StMk6n2lq+8O7DNsoAAAATEGbeknhDyZTAjf/+8Ohm1TlEJHbNU/mWRKFgYzLN2yQ8jXAB9wCzOsfbVP2lU1EY8QOgEzKqXo0R7iSjvfThP3Ejb6DhTio3axVsWsAAABXQZudSeEPJlMCN//77lcOA7zH4jDF8og00F+0xFfHR4VDZOJrgOjozJdR/bpg8saNaDwaxnfv4X/eRhAQ/CW8F20/hfAqOk7ganOxPkKbRibKABEgisd4AAAAIkGfu0URPN99VLWJAgSGMMD0PyKXcxUT3cFaf6NQH7WvfiEAAABEAZ/cakR/h7ESVb5jSubSAV8FnD8x4gGanIIpXwVIvISOGwpzao9rHxOAHmz8XBzh60OGZXgi//qWtWwRLulsxilCnuEAAABsQZvfSahBaJlMFPG//Blf8ICQaEccs93WPEeRPsjPuDf4gDC1lzaA9vG8/MDesnCCrO+qwrWg0Z041Lq6MnjtnVfjvIt0qF1hF5tlhQF/0d8e9aqpY2UcQgjUy9Hd3+CTWS0I7CWSzCsCaZnAAAAAFwGf/mpEf4jgOwqq40N+HXrqPTGH7PlEAAAAU0Gb4knhClJlMCN/++A9RBAtNsNSCFkg711//JWWvGFsLrGyK1H7zXz+lTaC4eBFVxUDg/Op1j5vlFjWLyVGxDTLWQWt7v06Slq9bNE1cmz/5OS1AAAAMEGeAEU0TN99Tam3C0rJYqxDWOyBOLJHIckz5MqYD+DjRZpqPdIIw0nQ9LZT+KG3YAAAACABniFqRH+BB5WG0Yk+JGFIqTc/MKybQhljSSoiM7L+dQAAARtBmiZJqEFomUwI39P1UQmO/CHKa5fG9B98DyBd2q//xdlw3GXkmq/tlCxKXGJhBrqImPX+Tw62HdRUoCKQH1/ixC6DIupugrK3/kza+ntSNKpOnTIb2YWwrbEfjsAUf/8Csikkf00Yy4JNpq0NT3r3G73dUSTcvN2hSQC1jpvA4PpqiSqoyt1rg4K5nEb2X3Ta2J5m01MrHsJWKap7ArXaL/fag2Yxz1DnrBGNF0N64688FzTsfy3aMFyBNTwHxDcOM5VxEYpgPbQx0UoRaoEjPSOfPfZWW0CeOiVTmPWdE40N3Fq2/t5cedvB6NP7Q1mjnmICrga1HKXTUnIgZVEERRT/lQNzJZZSENnwR7l7OYZdbPaS+0iJ0PvAAAAANkGeREURLJ/NXH5jrirbiTEhY9Y/tlMwwq1venEm0Vym9wIwyi7kBVs/cNPG6n7cE56ATPcwWQAAAHcBnmN0RH/N1YmW42fk1rr5dWrvg0hpI+7/lxaqBQ9u9qPMWMe//elYZNo+NLjg8cYBUA/d4tfmi8Qz2UXcEo8Lvkc59K0YQAe1JOntXnqOjM4McR1Gx6UrMCz14vMSat1IO5CZ6UXFyvmTyeshwVAiRewmoKI5gQAAACABnmVqRH+HrQBVU2y7OiaX26aU/4ah2bLxWm+Y90YzwQAAADxBmmdJqEFsmUwI3/wDwzHA7nIUz9EUj2ygWhqxgWMh9LPrZxgjmvIJeUBCqf7Hl5oCspr9BQrYw3VwQWEAAACvQZqJSeEKUmUwUVLG/8dHUh70tVYbmJAP5Yav75x8KgRRrVZqGBAEqyIrK6LRnrpZlSZkXNqRr0UFWwbqNCFOpYqLeS1RQ2PyEkC+dBVplJUfZiuhOQAvvreTkE7X7a6IRT6oN2yIUD3gEVfN9FOVcWA6Z8kXNa18/OFVTCJp5CwGEmxWr1WMNJva0wrVNHOS92/kbWx21BB21E9sgNNFnB189Z9RulpZSpgE8SAYOAAAACABnqhqRH/N9J+2J1XcPpDCijwctlpIFM7gAfQRhV49uAAAAHJBmq1J4Q6JlMCN//kgYHc058NIdrA4tZt28lo3D3meW1zCAqwidHRRteHMGFHB/D9M0MCX00dwKKmnXGyIjSiIo/6dF3Oz1KEMYwbbPPtUPovm6LbF9t2EyWlXZ2AhyytqaBHYUktsnfMGjBar9/t8GUEAAAA3QZ7LRRU8n36t1fJ1HlQISEIEfpaQSyOL40QRE9Vs9zvdbKrtZIXDJfPzuaeN4LIza42Lez2eIAAAACcBnup0RH+HsPlbmN8Ea52A1bNIOYqtlLjflwD1gV/uybFUQRKpUOAAAAAdAZ7sakR/hNQrQyxiUawcG/LrF0fwKrQNveY8OGEAAABTQZrvSahBaJlMFPG//APDOoh1u+/wA6/FZcjU/jvz0FUQ/BieNWgawp32ZtUYldJMEVjFO412vb48n9u19LQOw5J+xi/1yuC/hu68Dsm1661iVvEAAAApAZ8OakR/h7DkMUQbpFoFWKV9q8E5k5LkELQbh4pZCjhTQyjdPLum4gMAAABAQZsRSeEKUmUwUsb/++A44EJXr11WJCw7ECvX60Zrh07UIdiv7aDfa0prRgD6jv/0Wc4+K9FDJsKP/KqGWZJW8AAAADABnzBqRH+JfH18ZqS57WtbAe/IO8nxI2HXbfpO2Nkfx7T6UrglP5FtwO/Jq+dpJWAAAABuQZszSeEOiZTBRMb/+8dnvoPCAeCb2/rWCrRRMAZhYjQsyCGOaWbOiXjQ45Hd4KBD69SXBo/G1jQPyXicID3GsLjYV+7vn2Od6v4G+WQqMHMoRfr8vCzfPrEsm0u23BOlwjpYpIOmf//S7gManfcAAAAYAZ9SakR/iOA7CrDrXGjeuosMdCPQVofsAAAAY0GbVUnhDyZTBTxv++A44DvT323Jx0SqVNJ82IBpyyo7DXCG+RxgT1NMsMLKQax8K4yQWP/71kziHHgGscwzo5L2jLizRX09ig4TiJCfNOrU3b1tM7YaGwugP+oXNvCezwnqQAAAAC0Bn3RqRH+IvbTI3y+IN9Ohxp8NHgN3aoRoF0CyJ7Kizh4lPfrWVzPrEIihMfEAAAA1QZt2SeEPJlMCN//74DjgQ+tokgude4jclMN4O+jlwgqt2WyQlDcvOp1uGtyVo8FNYf8DxygAAABYQZuXSeEPJlMCN//7w+5d2niwPcY4fRW/HK6IFjrt0h2pzlDEeWsYZI51tgKX93CPYsiPqPq3X5B1aWcUk55oKJGxvjXmYgG2872vqJP978jdfl4d7IIZYQAAAGpBm7pJ4Q8mUwI3//vD7kjX2UR2WhZ7zwuAd6E2YZpVsaXqf/9ms8pdq2AcVeqHU+2duGYw6XN6gJJPmwjaRjOqn/963zMAo7BbLbQ0xWrehzgNm86QkFh5F4CJPtfJrXi2qfvexEjl9SQZAAAAPEGf2EURPN99VLXrSyKjUixfkxdL90ti/7QhhH5aJuIgfKR9AwA9eKagET/7LDiGl5JEDu2HE71HainiPAAAACQBn/lqRH+HrQBVU2zahqo4JDrA0Wsth2bQ6i1pu5g6vv10b/kAAABIQZv7SahBaJlMCN/7/bCFg4EB8mDknVh3Li4nF+3GJT89JkDhouv2updeuOWXeNe5p0zRxNBaBS36FPmUKGpjSDMCDJPbtDwsAAAAREGaHUnhClJlMFESxv/74VHYRK6xf9nuJ/mW8vBkEYktQx5ng60k0zpkEVATXt/+Il4+t1Y/gWE4ff2KkBugUa5RlCmVAAAAGQGePGpEf4k2H5WycQ87dp8VYGfF0/IU9Z0AAABiQZohSeEOiZTAjf/7/gOFgQm+EEcUp0ZlHiPiNyhbnVoqZ+8senN2fT/GHn0i0DJAHVHDeAAPPdb6+08ywf9ggCftEN0KTRcMP6zvyMME1GrdoLnzWjL9J+Li8Z7zUjdev4AAAABAQZ5fRRU8n3sgptQ/K33IpJWP93HYo6zTTQPcVWgvT/tt2dtXG8r9XMctnwTJYMXBZrnQX8wcobNb8TVLLKjwQAAAACsBnn50RH+HsPlaZzJzUbRAKINcL+9T91jJ1Gd5JZXKKMYQPHi5kXraXljhAAAAIgGeYGpEf4TUK0IOqy+T8lgawG4WgIsZVeplQdYOkTmpMeAAAABJQZpjSahBaJlMFPG//AgTwyAUQ/9Od+649emXNJgblrSImgV1dWIgYyUeX0WkGz1oPOdc6UKgLvfxF2udL9r+3vyOJ3zOsvfmOQAAADQBnoJqRH+HsOQxRCCMCTKsMv80pUcMEWMaBnG+bZLo/o7gAggPnWMrK71nR0MaLVpDsOS6AAAAMUGahUnhClJlMFLG//vgHIBhjoEviMguc3ZwqTYQjvB+dEBD+tTd/LD/U79/aezoU08AAAAyAZ6kakR/h7ESVYA0yC7NG/OtKYdKqWLeGmJMPcZhbIRuFc+rxG1R7C8fxWQE6y89g50AAAE9QZqoSeEOiZTAjf/V0GkzH5jjHogSzwV8Ur8f+Ec2ZqF42vRmT3+0XBKgmcsowSJtIfdoeCmqACWCkOKlAnt3SySxtVHTz5bNRvdUm7O7h2Utc/NeW1ORiPi5ahyiTp6pzvTB7gpoYxSr8OJvwgt07bwjElPFl2/DXqzJ7mlN8Ko4FuS3th/HSK3m8xjdy9Pd2HG/G/ifzmmcbp4pvEyDSaJK1W2uFie5GND/WdfCyg3jjBDikafVsZZsB9db1ELXgOz6/mtW1PO65LJeMpnL03zTJn4LERpRc9KcXLD8eihY+pbPtstJ8ymI7vFDm0fdHk9VglqCfCuOm5LOR1Xte0qY4v0ZF497FEg6J2gXBmk0KeY7+LdmgHYEcz0qRZkyA93gRlWQCcg2bdHG3R1RsZFfPnVqyGpeaMJQsD8AAABgQZ7GRRU839CfjYe8gADwWSBEZ0dUg8ycG32E01qLxeAqrHnZPUJ5Lyx+8BUnDWj0XZtSKfpUSpAhT/szG+CsJq3bkVF7YO0bZ1YvXCQjDz1cgeQteorKqspdc1Fi4XCBAAAAGQGe52pEf6jUd6g4XFW5qqvsTJAPqVNfayAAAACQQZrsSahBaJlMCN/UtXwfA78SpsbvqlpyWOlTh6QrFvdRJjjmVnCQz861RBv+whwL1CGfpCv6GfEMxGCpC9HYQQUE0dQXgCNgcSS6mD0fn/+UCPvjhBVhxi7NZ4U2QtlGo9RKKkNi+HxgY6TR164zEIP510Utju+3WmWp8LCMIzVG7g95lqRKQwTZXOyoblZEAAAAPUGfCkURLJ+megsfuYeXW7Xti51RtkOkbffzL/ASgnb6g1Hr5xJiwX7/7XlPkdirFvrO694xVrxkHWvRxaEAAAAdAZ8pdER/gGwq2QxVj1MPQwZguW2S81h08qjXqxAAAAArAZ8rakR/3Cspb5/QadRcLmNZSh6g5dIyMVdtWcxCmcKN8VSOxyyA+VX1BgAAAD1BmzBJqEFsmUwI3/1ikeaEjHOd45AIQeOuawfhOqyYDx3N0zAkiYWZShFSr60ONAwyK1zoY65cTb4j02w5AAAALkGfTkUVLJ99jGQnHgDM8ln+aBB+E2IcxAM5SUB34rzqVNrLOXd+6bM2xuY4uoEAAAAbAZ9tdER/gw4tjgY4RdE2+ViqVUsAWJkRCr/RAAAAFgGfb2pEf4EHlYPO0DBg6RcG870j464AAABUQZtxSahBbJlMCN/OlXvvG5yLTHBTaPrAMVQaCfuKj0j5S7bXX6yX7qQf5RBo8No2FXkM/qWMiQOljutkadwskYcyP7bS28Hd6oB0nybXNXLow16IAAAAR0GblUnhClJlMCN//R5j7oKUrIUIBCbU00bLa+xJg3Zc+8xHG2utLSrY8Lzn10Z4sXQk4HLrp+BfVVniEmXvgjPxoP1/+SNTAAAAM0Gfs0U0TJ9+tCxJ1Gn0jDb8mTrU/wrX4mC1oGx1U6eFX27MccNul90QPfzPIXxTBWHVoAAAACUBn9J0RH+Jee66QMi55jOocyx//JCSfKN8+wOm/wGvvJAu+yXAAAAAGwGf1GpEf4TUKzbRXPU6nraxj9iMWSV9IgbTbQAAAFBBm9ZJqEFomUwI3/wICGIsD2OziqyuzNNpqzvABec/jYnGKj/psjvJAnKx6B88nuodchjeQKU8/Sf4fxPeI3Bur+63wb/oqD2JcdFb2ce6gAAAAFFBm/hJ4QpSZTBREsb/+8PuSK8cBB148viFWw/NHqyGs+rm2ceC8lLcGprCfFWoP4uRxFnVQBO6X5Km47W1B7f7HR4MpVRMW/C/8Bu03QteuBkAAAAjAZ4XakR/h60AVVNs3VhXxKqVaJQ3kpIsIMaGLSdIpeJrRq8AAABIQZoZSeEOiZTAjf/7w6KlpQgJLKb7CVQpIKCX2GtUUOMn6NT46JbGsSJDqhqGAUlsdEP5mjnTYaodhrLt7P8/yMykEazbfVizAAAAM0GaOknhDyZTAjf/++AcgTWnNhxEl9OeuJnrXVsKR/mhF9lVb2NCffZ/IdUboGwovrVQewAAAFlBmltJ4Q8mUwI3//vhUdhAycJLcpdZgvrBHbuS0ELA2qR1GWV/sKmX9lqbykcrOwXH6Qp5yOSFSCovTOHfR4H66+O3WFhrSEGUCxI2T37abEC4N7SLdnmMQAAAAEpBmn1J4Q8mUwURPG/74Djgdzkgaa9RAdjKXG4f20PXgc0gPoX4Vj0Ucek2XSjMdrwZw3mpyUgcp/uoOq3e0VHPfh7Zubi3pdTi5QAAACoBnpxqRH+HsRNxvnokvr1IBRBrhf3qfussoHJ2FXwia8L0og/7L34YHvkAAAAxQZqeSeEPJlMCN//74ESoSq306qyBH0bjlOnxLGbrkGJ2Ws39CdVP/9jJZ6U/XdORQAAAAElBmqBJ4Q8mUwURPG/8A6IrIENlUdGDC2rGxYKvs0Zv/PoytQYJJmfkiM01RixdLevQCF9AnGptUZtShDxvMVuLKVNUIVwk3Z75AAAALQGe32pEf4egL3/1jAGor8/AQZ6N6eggOUr7LqhLAXV7FfOYIjSYAD8KRrCokQAAAD9BmsNJ4Q8mUwI3//wC0LvCDOPxGGZH0fTTrqCs2EFRlT6/L9YpDrTovL58/OT2vGOyWsqXu1aV/a7fySUgV4AAAAAeQZ7hRRE8331No5t4g23Ywzxi4929Dj8IQsijl2gJAAAALAGfAmpEf4k5No92peKr/htFkwdNuhonMp0Ro1YlnQcnCewvp2SzmqBS8yiAAAAAXUGbBUmoQWiZTBTxv/vhUdhCtFQj3CF8I0Z8KP7PO2SHsB6lgIAgPVTE9oxFKS1kdE31Ian66qP+yqf9hh734mZwin70dY7QE5chfscqW3JQE3iewy17sVw7qZM5wQAAAB0BnyRqRH+JNh+VxVvSpS/NWVqjqTRtbJ2NrINTKwAAAQVBmylJ4QpSZTAjf8EwsBpzW3+EphZT+sMjYK72pqFmMagiJ/6RxdCOrdIwWTYN0R3uDTSb/XRdzEteeehW0X8cK0bfRzlGLg54Bih5AVBjp8JIdS+frbm7i9b7XCr5CBhzbRO0uWlX13WMTPZSuvGoTUMrRVgG628K7rQSE3qBkBMgw+CcCNefjGn6Zkx+q5pfh8Mt0QqMzqj7mnSuGudADa/xLYXydiZCYQ7/ewnsEpGkOK9yVq6rK50LcAUK+I3NgBMWmbuoV9kh2ewa2IxwBTrP8/IrrKTnIq0axfEx25XaS1gBYRb6kDVOAPsog/1jZSwYZ4GatQjuLuVTjthxX68wVF8AAABLQZ9HRTRMn9O9y352oflZsxy6BAxflZu0pxOibLVcNnE+wS7cYiWlcJ/txGIxRrd+RdfvjrOjPfDoy3gMEaOA1ktKWZzpAcG2KsiDAAAAKQGfZnREf4ew+VpnMnMtuJjFWqP/yRU8Vfsm/SA+y58AHgHPfjY6Ic5gAAAAKQGfaGpEf9fFlYagcVz5uBzD/QIPWYbzU9vkHWdHRQQyeCK9nXQ4I2EsAAAAikGba0moQWiZTBTxv8gMsVZymC78zcwH990XdZZmZzJ4wtrIvyzMXIOULxsiqy3MuQ/I48r4n72WrGhLg75YH6WEsBqONLX6BjtJ8TjNYAtjYKLrQKQl6vGS9ZfD13YsDE3Jvu673GbAbsZboZ/13vUZzF8zSZOV9IwR2roprdQI0feIeu7cnUYSMQAAACgBn4pqRH+o0ZioFg2l6PKcHO/zygwvcEiwx436ZjD+vu4Dcnlr22pAAAAAQUGbjUnhClJlMFLG//vgOOCFHL75XjgbP/nLa3NdHqfl9wTp94oAO8T7lQOqg3O0Mx03n3NRR7pYlze6hVEf4hs4AAAAMgGfrGpEf4ERwY+Gt77r+yEJ4SvPAYi+RMSplok3KwSDGQpSuB95O2ZN6aV0x9y74RwxAAAAZEGbr0nhDomUwUTG//vHT9FLCBA2nxris8BwI6TDKNzQGtw992rw5prohMGR5qiMwrOY4KxeyFVh4Z5nsHZDDLBJCMqHd6qBQyfBqvHpiFjePAt/3JK/Ps50Y7xjr5QVjDmLLnEAAAAWAZ/OakR/hNQrQmu509DykZGMqHC4wQAAAJZBm9JJ4Q8mUwI3/8evureZ7lc3KXdmxnQXgMijUWF3ZgneFpdm9nM3WOKWFmTpBzrecoyXK1AK3/TzKgoznmiqaxTbi+HkZzv39k56Xd8Rwloi2NmfZ/xAvKYkrgc7Yj9489acf/vZrf+rHiwEDuf4P34wfjU87tMB0gm0+ejzliNtxzSXYbeAmmP64RveFZ816fJhWEQAAAAuQZ/wRRE839BbCo5vTqlWMLNw7JP9qz7O6uD4nu1gs/O9P9c5x1S1ytTVzTI2+AAAAB0BnhFqRH+BB5XFIVzyB9wJwDYiu0A8ZNtKI6DGeQAAAEJBmhZJqEFomUwI3/18y04SEQRzJO8XoEYWKQTqvMsTVgq6JgxNfx+qCvDs/Eb18miZtc9keAkyj4LDbnj+9YUtWYAAAAAzQZ40RREsn32aKglvf2C9/3iUkf/5fFdsKyz6B6M9ZyCIDP+BGGOtWUef8hos++ANyp2AAAAAKwGeU3REf4exAhsvjHgSnTXrK0yG7hQMOWbFoITHQQMMz45J+HIhTgvG3EEAAAAiAZ5VakR/hnW4Mf925RATlfTwJxSqw6b0//nC+VfvnPeXgAAAACpBmldJqEFsmUwI3/vuVw4CBhG+4lg3Es59Okupr9RDias/wbkJjrkoqxEAAAAuQZp4SeEKUmUwI3/74DjgJAb7CXLbzf5nq0/dretTIIHlFI1cAGN8WOxHbKwvXQAAAH5BmplJ4Q6JlMCN//vBGT5YQEgiNI4YfqwxY7pw7Y2LpcvMCmuuwZiEyCprEorTJzgFnI1fTLxSf+rSP+efSERrkuP/Gkp1S5WCPd1cY+R8nU2/Ecise1u7zn7FazbzJRLufU5Eeprb8Vk57B7YxKp6pppaYCg0Zui1dDn3N58AAABvQZq7SeEPJlMFETxv++A44DvMfbdU1hKfbdntKLt5PPLOQy9ifPEv2J3HrARgxvABi0D1MxB+8N2XXdutIkzk+EXbiNPhxpbHQRvoMb11pMEpF0dzIxjytdVY2pHHvt4vqc8m5kctbxoCGPrGzq7BAAAAOgGe2mpEf4l71gLl090r4xvvAf/rji6zKN85FiDpcbWZht1n1na/FcE7BBvTXNdoIy6yLTUk5BdVxIAAAABWQZrcSeEPJlMCN//74DjgQ+tokgubHJTmcS2bnfPkzdAbeefDjnKfV7MFzXTULGE7gZCZQWfJ6qOVkkF8AtxRAEmR69//H0c3PlHntplsbW9FUKAYTRcAAABbQZr9SeEPJlMCN//7w+5d2niwPcKLWmwiVkBeI/d3EBe4uOe41g3mCl0o4kirBlvctec2F+QjLJhaKWtN6qGqc6uLOLJ+QX+s7b6lWU0b6H0eQNGAmGh+ktnyUwAAAGJBmwBJ4Q8mUwI3//vDoZtTkBBrgFeRatJUfKwOw/xfRodt9aPbXcnC8nF6plsAvzoY3VEtvh6EDBUsUk66y8/erm9TY2A5NAgUOAFC29sW8XnubtkD7UGheuC3NYe+ACX3zAAAADlBnz5FETzfgcr6UK161fjdbXl6k8H/DlXPiuIwC6b9ntAFIpHPdd+z6GlVWgcCRI2u8oYOOGNhC4AAAAAlAZ9fakR/h60AW7YGdQtuoFCe7a0uOvlc1NbK2CvlDazwoRCj+QAAAEFBm0FJqEFomUwI3/vDoqVTkBH89pCmTZspbW1QWm3JfZZ0UQobs65Mf5LlrQRMt+qH8hg6HnjvEpO2ssyZ0Y2DKAAAAFRBm2NJ4QpSZTBREsb/+8HrVh2DcyHMmT3dUQJoZsKjKvT90DkkLfs6f0DlI9Vo6Noy8G4U70dLwbnSCcJHA+s7ZIbEbAz5i2KHdv5BjXfmYaICxMEAAAAYAZ+CakR/hNQrQmu509OcsGZ40siB/cfgAAAAT0GbhknhDomUwI3/++A44Hc4/O2JB6m19DmjiM44bKO1947whNbBWqUtCS8/5jm7dHXzg/GKHpWAQzxNsp/1CPUSWvt3lZLPoW2mRgyciF8AAAAtQZ+kRRU834HOcQ6HnuhEEjNaVKp2xQF7WlSmlAM2ebIQcM+HcEqhiiFkUxiZAAAAIAGfxWpEf4EHlcUhXPIKbaQkC9FdA677TaPsDtJppmm7AAAAa0GbykmoQWiZTAjf+8HDrHeCOtIL6oV0eH0f5bCxgOtFfLX4aVpxgeHP3XvPB/PFVi0PRBD1cFMejAR7FtsN/v1VPGtv4Q5R6r3IoLYn6BFezGJqahiRYmgCTMTFmr0X4+/IlITNyh6ZhGBjAAAAS0Gf6EURLJ99PcOlKzJo8mgZiaFQ/cRdc7LruQ9PEd+7B2bhHdjrU2WOtGZMYKSDh35QT0pidHubWQAg3xDn8inXd71eAkmkAv1mZQAAADMBngd0RH+HoC9/1Oz3bWZ3fidkWX4f6J+c3kNLc5E/BUMQnT52L5poTTPp/CeS+I7NQNgAAAAqAZ4JakR/hnW4Mf925IfVrMHCbM+zJMcVzMO8153EB9P9wj2YcBm3d5ATAAABFUGaDEmoQWyZTBRMb8Em7juH6SagJhLTIoH9PuCUpPS4Xhl1dciUnk7fRxj9of4QGmTPb0WIBMBKQd3a1O1Yc1N/4kQ5x2J4nCSJXbRNmnX5rUdd/QcMYbapL5XWBGEC8Ck5/94+EirqJFZDiSW1qDre4E/cm0Cn57dgTQx4ezV9WnDqeA8qNeZT1cxuSXLNrFWIMllZp/93Xxk4j6Boviamv91IgB12vJyIkQkUqR+e7+pn8A3J1nM0FWyTrzkO7GJqjke9Z+zeEFS+E9FCk1HBH1Z9GwPLuDpVRLSp8hvtMrExkYwEqFJ7b8rI62OQzrSoZT3+EHwb1gJEd+fqytM5WmU4uf9eEg5WPacy8CyFQCjn58sAAAAdAZ4rakR/3Cdw0sHnaByzNI/DvCcZ5/JboTOZcBgAAAC3QZotSeEKUmUwI3/IDapfcEaIOp2yhhwVS+etbAWw+lDUW+XfaSyWSVIb9I+3RUZO9EZN/RG2xvytqz5YiVBaQOmzJiMDTSFB1apCB91XRO3SZTj1OyHGJzMKmR0UueNYpWGgXCDTBd912jNrkePuUoCV/pwa+tGQblTlC/7Mntowg9/HoHjvG3CnAA6FIBd1q7bFjLrs2BRVwa6HPk3/SE5LLvUOVOCcaqjJwOK9HolGgmBx/uLZAAAAa0GaUEnhDomUwI3/+SpVI7R3snQyETVUSNJhOfrvBOYWFqawEPOShEdffbNNhj/UU5hbs9Q5x/bRP+L8yTojPljMi35wxQSoHRofUyY5upTIiGlJUlsx7SzyXtNALaPJF5KQsK7ukuM7gVr7AAAAMEGebkURPN+Wev9EAG9OqVVseJN+Srxy0ekOaNhA4RLaQQxTYVhStfIpPBuX4Jw3IQAAACwBno9qRH+BB5XFIVzyCkOiEfbn4C8Yo2klpwwoQ9Hc0Mg62lIXOy4KcuI5ZAAAAKlBmpRJqEFomUwI38epbng7UbhTYq8oWxUQirssnQMSb7P6LX7rQapmaLnTKX/YQ+mvBQNdrilXpDana8P/ipoHFw1c/59iIYR8R4DQJ3J7cU7mE0EUO/5HJiOmapC64SNTZLeKkcj/wYvW9CxQjNGOHbzdRsVXG5vsdQcdxC46I2UJ+EPeA96UgnLDLarb/ye4UGev0CmekJlKpnTVesJsW30eCABpwrUgAAAAPUGeskURLJ/TvhMx1yLmtTPxEmDZ/zbA0Szadw/YOlgT6D9Dej7fdLCUUYjrPMIemsVWLNdcv1K6VaGlVIEAAAAvAZ7RdER/h6Avf/WMAZTJgkbiUU2GBEoThRTgMuuNR0WaohOWxAZ8OhULFKorD9YAAAAqAZ7TakR/3Ck9CeH/dWEJcswcJsz6v+24f6f0uV9+xslbm5g70jXAFOIQAAAAWEGa1kmoQWyZTBRMb/wIE8MgEocd6yv5CokI+/G0GUpxgPsND4Bh5ENvi1KT8dhbDutyYfbbj0VildYU+yRxq6FTnj94kjyaTRWbO1POXfK9v3D656X4I0kAAAAtAZ71akR/gQeVg87SZ4hSEgBWkj+ZBwvgo+Jen45iep+n2pmujJBvuAKIzNpAAAAAW0Ga90nhClJlMCN/++FR2BweCG1oRIssFvf41wc63i/NokGGJ2DfzyRKo+/2yaVmu/JJW0gloXeQ5hZa5thWPWMobrf/C7m5IKYTrfpZbz3D6/1t8557AoqeOYEAAABKQZsaSeEOiZTAjf/74D1EHaAmN4EWir8LxdLPbAZPO0M2+PY2QoG/X+gqbY2OfphX+3ikmCHy9qBRinr117Pf3bBw8DsA034PWVEAAAAvQZ84RRE8331NqbcOYV8ovUPdy4IB1YpDWJQ+wSvnfZOYkxxmYn/OzBmtYmp2A4AAAAAeAZ9ZakR/gQeVxSFc8i68ngUmW/lZdC6KspEQb2WPAAAATEGbXEmoQWiZTBTxv/vD7cBveCN3NVCXGk0YJpTWI6S27wxEVPlc6/1wnjM+kCdJnnhPo14+G+H7r7GYvUMBsla4e7HrvMr8M+H/u8gAAAAvAZ97akR/h6Avf+N3Bj9G4GRPrOabGjKMwWGQh7C506fOxYn6ZCu0xmofqwUIi4EAAABTQZt+SeEKUmUwUsb/+8WeTlkZAQajP8sPFAhMqTrywyxwavokyr3SLrM+i39/I0saTn/n2QD6z2f0Tpf5o4KXxn3mPA9D17rJX36vYMwHwHjQhLEAAAAoAZ+dakR/h60oZYAy4TexgRX9qffNaY7PWMw/K4c8McCTpcC8j63SIAAAADlBm59J4Q6JlMCN/8oIQcmggO56mz4fWIMz/u1LoEQTzV5bZxKWzyYuNw0VUour5rIsDiU90umHFkAAAABPQZugSeEPJlMCN//74VHYPx3QaH7LO6y4C4GizeyIGkLKZnJtuowmRMNw+Y2sUwAe7deUwGTWQ/Q5vriZGTkIsv3ToEeu8MQHLbiiuEa7gQAAAGdBm8NJ4Q8mUwI3//wC1hCBKrZCFEPTQr0tZ5jK1ORrDiuY02WzVUIDXNEWNY+1UnXo/FFfdn9MboNzrW4gnI9f6iOOt0mZ3nsie9H+/u+vO5CY8Fj2uueEyKX8U3QhLU3/TcKrIh9gAAAAHUGf4UURPN+B+FaPyFFuXZbgYj0u945+bqtESVlZAAAAKwGeAmpEf4i9r/hMofsoVdV7ULdGYi0qwviD0eMbDVTnrmwEwqTMGF7q03gAAAAxQZoESahBaJlMCN/74DjgiVvp1xngahySXAFM70j0RNfE5vodU8e8SsZ/26Vt1OWDKQAAAEhBmiZJ4QpSZTBREsb/++5XDgeRcDFSaseJcUAQLJ6v6GMzRDVOjaHtSG54x4li8aIVzsEHeZQB/Ro7xKkMpXenDvr+1lP2ymEAAAAvAZ5FakR/inEMYhyP/rt01Hns1oE+sW1Ax0NveAZCmuIIkSUdWkHaVhqEgSKY0rkAAABIQZpJSeEOiZTAjf/KCEPOyTOHP3Ga0L1Cjd69G+Yy89VyBy1dtUAy8MuVk4n3KI+TABbfvtiT/pPhtktWSGJDgIe68yWWzEpvAAAAHkGeZ0UVPN99TaObeIE6IhnjGIo8WCgt9DhhPUr/BAAAAC4BnohqRH+HsRJVvmNKEZfxWUIT/qJm6pPNGuzqg2ax/lk4U+jw9QOaXPUFfo64AAAAj0Gai0moQWiZTBTxv/1jexkZY1WGJiT87bz5jKDKBIPIUpqiOKoFqMU9xuIH4219MWUk+QJUsSOay780Yn710gUCq/8q5pJ2Qna2z31AIaW/7r/7nw3VuvmUI84z2bsgf+2HLYAKweSx9T91NAbzXY1OYhhH/Rhoq7G9wO4NOvgpiqkSQhrt64wDHDxeNgBhAAAAGwGeqmpEf5BXF/8DaTeNmUt5+Es7Qo54X0Q20AAAAGVBmq5J4QpSZTAjf/1jexZqwShLchuDwYT5GB8R7EoFbWKry1qGHy2qItJHoYSzhYHpftREHwxHK/JzKBYIGEaS6NFr4w6720GZ4EwFtOtZcCzxJ8lwWGRhMcXbe390w6J20G6KwAAAADRBnsxFNEzfiNYJXadE4qFM8TYpcvRqNU9B/tg74vA6eRce+4rX2wUXpty0yO3D1hzy0b//AAAAJAGe7WpEf4EHlcUhXPIKbzh/PPJ8bNtdgqDOQTgASBJuMgwAoQAAAPtBmvJJqEFomUwI38HgcMuPufQ+1zSgVSxP44j8Xw0s65lEdmNlEJPiZc5zpqLY2ASG2JgZLJLCVojwLtyVBJ2Djl4lz8/HuoMrqAhvSAdjKQWYwWvku/PylYEHSBWY+TA2i9oIGaNmjGx6LG1Kj+E8emazrbQW4uXowHf6FzN0wi4FUL7pYGixlMvdCCV20iEhjFpxLFL7zMrQiJXs0FePfh470wQo3VD0x5R78tx3albjYGEvRrqiobpY8RLCYk5k9JRliWcwQKwaEOM4PxOill4WkG8tUs4AocgZR6pVN8VW+TjxJGdBFTt3zP5oF1jV/E0tbkUBVpsIQwAAADRBnxBFESyfzVx+Y64rGY1MwatWM/mzoG1t/TCEbE6JSTBCJehjrlgVbP29VWiv1os6uGraAAAAcwGfL3REf9Pc6p27yHwywlkNj5DchvZ7/p56d3iHqPJJy0TZlJpjk/4Pz69n5WmRBfZn9bTbyeeOj4mF9//0Kaye4vxJwYkk6gt36PZDZ4v1fnhhsDUliUbwZt3CzWRiFHU6yepOVv1u7QQgijOhrLbxJsAAAAAmAZ8xakR/h+Qj8Y/7ryVtMCNL7ETyKPTlzq4JQ5ZB/FO0LOIakpEAAABFQZszSahBbJlMCN/zu7OUTa9XLb3eW75wn7Kv236GCYDyaFkEi5prdIBVAhBOujVHZ91Dr3M5z7YWFsFIPaTIAQW6IvvAAAAAcEGbVUnhClJlMFFSxv/8EDLdkB8H1mLAmhRVDhMcY2mJZi7xKwfk8PJDVJGGU0diyt/1T05IeVjQbyr4y3cfABc+R55miDu8jv+ke/z5Wlghc6LtRzgp8YEXwp4pBzCx8KkBF/UR9Q5Fhyu0T3/n7sAAAAAXAZ90akR/iOA7CrYE++zmxFeP15r8ZlkAAACYQZt5SeEOiZTAjf/HqZ9anCKYumtLnkGnytBtxceH+Tab7WDld8HDYn4TZQTl+0UF3pmdah258gjA7l14kM0QSfHRYvwKPR5YYx3eazaokeo5MLOsjv1A7a/dr6fcvO8RrTfuoTK4JG2FPtdS+FD/W7GxeAJSunV30h3c7FLYa2Yf/MWZdi1pjZRZxjRsrX+SukflJqQC/GIAAAA2QZ+XRRU8n6Z5caH3IhlDPEGgwBap5BKO3zzvTuKfuLPQfl8w4HkfPygPVlVEzHop4sE0RruBAAAAJwGftnREf4Bxg3qaOZGuOib2Lt1Wma/ikk81jKfwGZSD4w4Dnqo5LQAAAB4Bn7hqRH/N5HsNDLGJUxpLAK9mT9gX6FWVFLp4DnAAAABhQZu7SahBaJlMFPG/8pGuAUUgcuCRH/OwAPFSZJ3zCRS17ZFc1C9yqHmc8/pVK7IEw/oV7yJdMcMIavNsCKHsF1NBZJBQvHr63RBRk+tra9PDpuRrNWg1DBO3hQlugGNbeQAAAC0Bn9pqRH+BEcGOCVJY/QJGN/iSBG3l/ieze7n+BOFxVwCwANjCJKx4MnZIv8AAAAA/QZvdSeEKUmUwUsb/++A44IUf87s3whHef/K9eDWAMU8ldp4l040VzlBC65n7nr79QcRWreur7WcWsMhEQu2fAAAAKgGf/GpEf4l71o92pOxqkP3CQfV7f8+wwfDv0Vwj8bTm1R7NQsii5g4zwQAAAGJBm/9J4Q6JlMFExv/74VHYQpILXs6KXOmEj6xleFO7YS9ET5DIDmczDCW5+M6UZA4DyHmCTz5FoaK0i4R5P7y599cypfH8m6UCrK8Bf/vgLazNz1wC/16tdStMHxBIjZ4P1AAAABkBnh5qRH+I4DsKsOugF12qomW/SycjPJtcAAAAZEGaAUnhDyZTBTxv++A44DvMfnbE2rCH23XL0djoBAvFeci3xno+OtABO36XM5KyvnpfzxKvle/8J3308zrG6JVwhHh4rsYQUFRysfHizQb86GKutGNCt+OY9sPNteGFKlv4dXcAAAAqAZ4gakR/h7ETcb56JL659u40VH/5IqeKyAkelvuqVBBk49Saqe7VmZuAAAAAMEGaIknhDyZTAjf/++AcgTWtokgude4jcmomud9G6Ir5iVSH+0BGgJZK3LZX0JscoQAAAFJBmkNJ4Q8mUwI3//vD7l3aeLA78Ouck9tQPwA8lswErUECEWb2BfP9eCGIuUSrHm2tubC/IOXP44SUuC0wRBPcOUZ6UwEfr45QL2uWpkMDOFPAAAAAbEGaZknhDyZTAjf/+8PuSNeECUbzFqDoHVeVghCqcIHbUioPt8t0kMbXawr+fm53In7aCpqLDhgohy3fxI7sPPRi+v4k4BfCt+ArbIn/7hr/fNgQ82kiJyjRCR88G6lkZxT3zjOqC80cK2gQDwAAADlBnoRFETzfhHJ+XgcrH6mSiINjyp//kgke8KUWO3XO8bzOwBXQN1PZPssKbb2Qy1CNfZzQq7z2F0EAAAAsAZ6lakR/h60AXDOWhD5PePsAKefz0nC6cwFfCc9G9GK1nNdoyU44ymIpUb8AAABDQZqnSahBaJlMCN/7w6grIDgahJq1KP+AGMxDG43cWoTqfOrO+8pS10UHq7xCbX4a4pJt454b+M0H1Eo62OjD2EdzxwAAAFVBmslJ4QpSZTBREsb/ygqpP0TAwFGuhMe0H+LKq01NhOFINpQwf2KjhEvSCqmG9OMZE2/Qqy0272Gt0nHLi11coL4E4EoCisDk2gIBnTbHmELqhvwwAAAAGAGe6GpEf4k2H5XFXBQn5jeU7JdWsFUrSQAAAGBBmu1J4Q6JlMCN//vBm7tPgAJMRPqV++IZbHFD8LyLSaFG1Fj3xnwOAwhoQ0X3rcSTj34Hin1WxPUwQOdGW/+/l7yZtILXAd5Z9Wgqf/aUbVivhneizf3wbIgcVChpYYEAAABCQZ8LRRU8n31uIe20/kHBd7snCnkEo5gkxddxXoKnJ5QnndNiDb5+E/ApR6QfVh6OvB748Hzqtcox/cwTiLlMZwRAAAAAKwGfKnREf4ew+Vu++ZOvbsHGmKLhf3qdqHvTmJsewon5VlSMTQ0CLsoGQswAAAAgAZ8sakR/hNQrQyxiUcxzYxkWBaHD9jMxUSoH24CdrYEAAABJQZsvSahBaJlMFPG//AgTxcoiqckUP9wRELdF44LQ+J6zoeYQmvRkGAschCcjUxhb6c6E9RwD3q5AdXgn/otCYxPZzTf5u3+JwQAAADEBn05qRH+BB40VhOFmnLMmN/Rt6KeYRcltPwhBw6Yj8gKW5Be+d4emYrBg0QKVVh+BAAAAOUGbUUnhClJlMFLG//vgOOBCV69hkAISts4hjQjQR4n/taVosJ353hMn/bQatLAc7vxjVVNjpBAL4AAAAC8Bn3BqRH+HoDXutdsAXaf9vPXmz+vq9nHMZG8uHAniuS2VQVH0wdR1pyOJp2ZkQAAAAOVBm3NJ4Q6JlMFExv/BURmgKmQ2sem8M/cS5YiFuDEKj58ESA2Ioo541n9schk3ne2Kz27RHR2Mtl7sfZ6Tk8CR5fodz+unZYWDbUKNmf3Gr1+o9Rxy9lYH9eJLD5VCN+c3svum1sTzNtsJGBK1cHItHWUqHplu24Ckc7vqxHOe0ZtkSJdUThqkuknVrcKcY1ewS9SB7XN0a/r2y0t00YwIYLkLxxLwPmaULQNo/ape3s5adm82wYXnY+bUcZi6E22itPWO92rOH7IP21AdVJhrsPeiqRAczmuYnL6vkPzP4adcC7BhAAAAWAGfkmpEf84wSggSH8mtdkNj5CRqWoGie3lqD7CWgy3dy9dUCO//W9YZNfE6Xiu5GIKYFfDcUA0T2ZBvdwLfngiXt+EDkfDWv0K7BOXReZCjX9nSExoLgSoAAABZQZuXSeEPJlMCN//74DjgOprY+5Fmn26ag8A76dO6JZsciGMnQ2pYcNEnea2hGeomPmCENu5nJeGhNjlWP/GqpX7N4n3KhFSvZLLumw3qoPc7g3O8DwTNGcAAAAA4QZ+1RRE8n309hYNYZQ7+HfhECBi/K0X4frDrq0X/KbwmZ8BThyQ2+fhu0YmOJCJCHAKCHj32xcEAAAAlAZ/UdER/h7D5Wmcz5GDrLvRDAsCvzGdqHHvdNLLc9bD46lxTMAAAAB0Bn9ZqRH+E1CtQOK5zZksA34p0LQK+zTMJUM7eYQAAADxBm9lJqEFomUwU8b/KCEPdL34SY/510oR6jdtPc2spsXUY6gpzzsPQOWCHtIKBf9RsoPGI0GagH7nZSdUAAAAlAZ/4akR/h7DkMUOCGBJh+lP5pSg5ck9EFnrQLuSi5xBprxkDQAAABQJliIIACf++DLeYgN9XYj5zg5B5QJuOQipq5DurWEjnW3UorRPL+Gmp7W/eIApq4hA6GO7MXHIyDHsVucz/aMFYtKTDM8z25ZLQfiNZJqSiDN9723ESzhz79Cwt5M+8vOWuuGhG+Tvq9foXjrYoVXiVDwEsVfUmlk5TXbEqyYI+MZpo4mDJdCI6FifaFbOBypE2MZf5ZLIBH7g8U1rXJzlzobXWQwBVEZ3H50m7+bGCm/s/TO02Npz6eSKcRnoIL6uejcV8scEPRiZKvQslaZioRDQNI3E8UN6+/c9cAekW3/GkyyVmNiyROqcKtiIhBHi7pAk2+Q/TQP/YwlHw0o4KIS19ytAyX2tYiRc1FzBcAlpBmEdwA9/TY2cHx8Zwfg5om1WDRf7F7wfrHwkDprLxXYIXoXvv2/us1OW+b+Xd99mIui98EnMS3BL3wkn/NUHomZCIfEFUCZVjLemRQTcW2M7/WlzDgz8mCOUhc7zw+tq0JUXSMkS1L/cOl+y1DmBMgMzcbhhNbzDLeKWkrUH5XaHyryAkCyr9U1Xo/mAgTLsnExuVa/P4R7zZdXvaMIk7ek35j2b1yb5zOOXYieE3ARP5O5ZjAZBq4ndpSi6jkftKCBVFGUAJHA0yJ3SOwTwLlD33y2ioFahycU0S2U1an8v7S5ljSJFgUt4xqxRzXPLMm5Q/lf3uq1K+w5f2w1oQlGAPQWlYXk6xSti7jEz02BQEwIn/1VSW1lDXiIsLQ9/68Lmyq7+KqScF8M2ttOd6ikUYm63rAEQ/1uwF2tw+vDoK58ZLDjGGY+DsvOCVmKReWQXo2mJU3Z/+dukASHHoLCtd6cudmKCocwPIOFIrZIpuWFPYP7N5w3TQNI2J7uOoaE0UlH0FyMvlyJKCSxK2KXdZbEjdWLzEYIhD/rSHAXZ83AHb4RVPmbBpaD4FR9EhjzYQ5kM67XQ4jgiM+83E3jJ1tuOTwVjEUXOTiZCx8wy+C0dX3lqduBIJOztDCxzLLUgy1pDJ11agmN4hiYEyREcCB9PLz7Gu5NIl+dxSpct15aKEYrYOog2sW/3UjxYhz6/+6vpRxNYFBe2YVFfEse/JAPrCdAPUQXwCIpcGvCAqIe9H6yayOc2/eUy00kYB9Hqrr4iHiednpjGYIa/wBQn2kawkd5eYwvCd09N//TBsRiPySS7TyptSNcNkl3kSijqYAzDU0cTjAlo7gSKsBb2500bskV+Fa6qC7WYdX+uHeAkkUfK7djG3sDH96VRFsoJGnc9eiCjPYyGt4n1nklz/pm47DKZfzdVcEQVVqXfQk182IfGwqQ6YPNZNc85BxPoa5UG0Nc7aBMBm/EEqGdEA+1bpcQALD7p+Xw3OdXvS2i01L2TZwR76H4FwPsv/kg5K4Q+gtukTkAv8aXBY908KoaNfjcQLsJI9nUsxo8CIfcaHaV3fe+aAg0mPXJT7WoUH3EqFX9KWMCffuLgk141zvZdeWb9G2OAayn61mciuEFWUB3unf//3UT9ZWH3tY8Odtz5knWw9WcA/7IUpQHs7AGXSZHFSuMbAJ7IeW5YgcJ4/2xVyVfqWwGS8VXHS84Y1Oc2iGWIYRH4DJUk2xI3a7Xnxr04yGYzMm4F14EhpEJRU19w2fQAabSuyr93+75ace0n91G5DJsmYh8MQGBLpKZioZADEdyOWOL/b3dH1hKiEdZsQsnfs/yhdAAqZAAAAPEGaIWxG//vtcWB39W5QeWKwGeLr7PMDikgXstlWyUqYOYjnV7wDPmyxfP1O0kL6PoxR7zo1Gkb+C4TG+AAAALJBmkM8IZMphG/Hqe1RFY5rt5egm2nn3JQcBwuaJRr84/ddLi9/76agb7ejonTAqd67XVYfnqN4BuUaC+fAk+1lcholPMzUb84TknMEw30GkmM7GgprsKqCgwcaNTAOWdS32AuH1x/E2QaetB+/mEk87Hq8PlJqqlYpWPshK+oBJK6KuTSGIwg/9Bv+LVgT1pALr72MIHMQXmCPnTdlDJUk4bEuuLIqk2to3LjxIlmPkYtBAAAAGAGeYmpEf830n5JzFkX3j1Qd8Q+AUGDVgQAAAEhBmmdJ4Q8mUwI3//vBm7tPFgd78v7l86Um29hR9xxADhgMD9LzlgdAOtozhLW7UEkoFYf5otHinj6odtZRtbm/yoQHLKnpnZ8AAAAzQZ6FRRE8n31uIfah+VmyyuHZOFPIEoZWWjwctTHilRP/EBaLGght8/KC14RVIcYtmzmgAAAAJwGepHREf4i5UoIN1KmuYE1Jdm1Hq0jw330ncp4Hx2M/vKMq5Z4tgQAAAB0BnqZqRH+E1CtVyjEqWYFcr+nHuyCeliO2EXDZLgAAAE9BmqhJqEFomUwI3/wDoisgI4RECc7scuWX932y6q/Gcps10izFusxOOvMniGaXl24boq3k+UZg8rPMUT3iQxdNcNfEd4A2W68ovwz4f+7zAAAAVkGayknhClJlMFESxv/7w+5IrxwHtyq44n5qzWH5tokp7F/MIurN4zLf3qsBy/LP1c33i4OFn3CMZ1TBLIwXvbEHj14Vcb6Wt30UGVr0GhXroArD8I7AAAAAJAGe6WpEf4etAFVTbNo8fLWftyKUew1Jo8zvS8XWIpeI7UlrUQAAAElBmutJ4Q6JlMCN//vE9pkb6iHbDeYYscsOWTLED8gfFsN2l+iXqUwnGAUp7SGWRWf+Tj1MgMpeHoo49mO3nGzGBc/JZpGbbwjNAAAAM0GbDEnhDyZTAjf/++A44EO2R9S7lNO7TCCSBTAoqxeSZe0BGgDXxryEEgFwFd7IWXy8YQAAAFVBmy1J4Q8mUwI3//vhUdg9fr4QWmZ38DJvZypVj9XjHyiWE0qfucSUpcKpbBOOJb2lRZHjFXqbLB/AAyfZJCzDiBXSwSe2d3nKLM/QAOEJb8mw3ifBAAAAT0GbT0nhDyZTBRE8b8oIQ86YuCDOQEpyy9UWaSGWOxDO7wiSCJDYIhWcih6YPXwwChxsQuyPkj7FvLL/pgLFAoSeOsn9+9F+L3cY6vDMfiAAAAAqAZ9uakR/gQ9o5WBUzd28w1hPr/+56V66bgLLEK74aalrWqpQrcrjTCoOAAAAM0GbcEnhDyZTAjf/++A44Ilb6pCfLUOd7CmgqVP5gaCrD6c4HElsj+qTov37vq7C0KSjKQAAAEZBm5JJ4Q8mUwURPG/8A6IrIENjLKcjAdSvTZ0AUa2HabW/T6JywFoZYZ38XBzMVTWT9Boz5CqJWPYr9zgIZJzhObYT/U75AAAALwGfsWpEf4pxDGIf1f+tYtz0i7E23tkVLGF0sTK0M/3IyamfbQgfSCSb1BPWNUB8AAAAPkGbtUnhDyZTAjf//CJ/zIEGcd9wsgkXNNOYwvx4RPo8KfLe4zM2exTFz5+cPtcVtktZUvjRO+buNFeyS0d/AAAAHUGf00URPN99TaObeEbxmT4L02e3ocfhCFkUcu0BAAAAMAGf9GpEf4exElW+Y0rm0PJ/wJM/5jtBqG7M53mIMo8kEOBG1R7FTsO3uy/8wE9rgQAAAIJBm/dJqEFomUwU8b/74VHYNYZ3gTviajBNW4911XFIgnHuPPvf6PuINmTDgKw6tl4y1gHxr/o1Eoc443t4pg43rfeXg457+frQjkKJwwuU0Dg28hgOuxXzoy3/7GNwWDiho5fxpVul4JYlW/s2wZ6+FIFz9RB3Pxla/HA9ajISXcDsAAAAGQGeFmpEf4eNe2TaTV+is9Z+r3dj/4M5IrQAAAEtQZobSeEKUmUwI3/OMDPDLD7+9LV1rvQO8dIyBgu6gBZzKIHOqwdcYW0mj4/p+KVMgoNINR6fpQkbB+e7J43ZN6Thfy8C35cIcvE6ulfkdYr0ZP5p+0EdUB5C7vhx0r6X9DkZA0F/t8XcJ1Ay1jkXzIafX5I9JMZyOPPT23eliytMDQN3Nmm15DczfsII5lelhwcySIf4ku6wzRULjnVZWRMxcaOWnCVmvPK0WKnXS6hAGrBovrD5c40CFPHBrLy9JtUVvwuELo3rnQUyLyChZip9Wl/I5F161BBhavdM3/GZPS1ybPzPYbRmbIdt1SZq/ybmP8py+DhwJs6/zeOqKGDQvWpP3NX7JQuvuIBzABOHx7UbDGk+BICall/zPVb7YjEOQFeJndkK8GmaDQAAAEdBnjlFNEyf079+wBmd8pEhgLG8Xc7r/4rqUnzQQEv0ohGUAL7kLhkvedlsNVqrHka/nxHyyL78FF+w3UmoM16jfBFXZPl4yAAAACoBnlh0RH+AbCj0Y5cPjkp7DicIuUO2PQpkoPz3B9y9RBSHaE7ieceKLYAAAAAtAZ5aakR/3CsoGURp1Gbf6xetbBvbdg7XQAomNjTNFLLYoSnN1T9WTArX1a5xAAAAf0GaXUmoQWiZTBTxv8gax37l9sNbrPn6iWwUf8RxgPTzKiAAtM7FGqChNHUkeqgcz5deTS6OtdOTIhcU10sqfJ/G7PWcIpym3Fr6VN/qN2kQwEQKlPOHuE2NC5vUo2j48TIVNP/sgwP3kyf759nBR/KFqquMTQDl4oofPV/0HkAAAAAtAZ58akR/qMvbvw4JUlkDvjD/f4ktFYfXNpM4bGfm2nM37h+CiGSGIelwUC+BAAAAQ0Gaf0nhClJlMFLG//kqHGZpISriFH/O7LMfxPF3levzQUoyR8IdbCT4Om5XQ61r7CJghz72IBV8PYi+fqlXXR+phCsAAAAuAZ6eakR/iXvWj3al4qv+G0JMwV4sy5KJ/S3ER7xINYVeI2qPYr2GLSVtouwGmAAAAGBBmoFJ4Q6JlMFExv/9fMnXUwOYhp70FTBXdRxtKOUP+0QmwlUd3WHIaIY4UFj8Yp5d9xxsyZ9jD2OI0QWf2S1Zcu0qWkg1dQDmX7aUzQyVEEdscleAHvns65gIbvQxIesAAAAWAZ6gakR/h417ZNpNX6Kz2wcSbK+6JgAAADRBmqRJ4Q8mUwI3/8oIQ86YuCDOQ6JMWiYf/iRlpY+LK2B96JBM88YxxpmUYIE8NLeGUY5dAAAAK0GewkURPN+D+9PfvvPTrHvZhI32OEWXdhkGCruwTO5SkxQve5GtKHjiB1sAAAAdAZ7jakR/gQeVxSFc8gptpCQL0Wxcfx1JTwVzldcAAACEQZroSahBaJlMCN/HrvlPtuKGl8J03skzn9tX3N+UGKgPtt5f1yW1mhtqYtpj2nU5x+dydduAMrmraSE5zwuU/zBlhXh5g8x8c38YYs2MriPKi55w4mf0Kwy5yX9U4VJwu6z5gBqEnviXVdXlHGiM76Z3+hcdzSXdSABOlyYkF+z/YctAAAAAM0GfBkURLJ/NXH5jrirbiX21Z3CtP5YcQMlKXKnVSfznQBXI1IR3YSzO0uaD7W46BIwddQAAACsBnyV0RH+vXBhX1AzTfh+ibFtqz6eGMQ/DTS0BafSOlfYuavjTxTNGlR1QAAAAIgGfJ2pEf4Z1ktqM7YEw3wBj96oC8se5g/yL18ispByD6DEAAAA3QZsqSahBbJlMFExv87zR2otPwnC4OQ4HEYtHszJf56ewlbVNYDdRLi7DaNM07knXvZOBJlAsDQAAABoBn0lqRH+BB5WDzvLCGv4OSr3ZyhW71ibSYQAAAEpBm0tJ4QpSZTAjf/vhUdhErrGwgYAPXwOxHLcvgTv8A/Zx9DjubfQaVDTVMebAKQ+RkqueYnX184COJ7NhYucgTvasYPFK4F93kQAAAFRBm21J4Q6JlMFNExv/++A44DvMfnbRI+5Ux9vnvb9fZlh0+yV0sGFmgEwKg3swNSCp7N5faAgb86AHOGr8Z8DiL5UcrnP4kp/lO962eXNLUHqjhHcAAAAmAZ+MakR/h7ETb+sxp7t4uJPRUf/kip4odeYTuzhEga5MwAislMEAAAAyQZuOSeEPJlMCN//74DjgQ+tokMdI2I2w4XB31uWN9xpSf/b7WfaAjQAslbjZLMjZGYAAAABOQZuvSeEPJlMCN//7w+5d2niwexjKrUXs6RzXh1b13XA1usrXuo4Q+o9ouvQ1D2Q2Wm0PpqZJ/HVh07u1x+848SMV9TQUPFzbtS13xBB6AAAAW0Gb0knhDyZTAjf/+8ParGvCBK2RNZj2l1x+jAR0apeuLZtFzhogEMtR9DC/p+fqK8dILLP+IbGFXeex19kT5IY4p8S3kRMYF4bXDgihVw55Ypd4XX0hDtrHn2EAAAA5QZ/wRRE834ISoy161qatNvHcx2Sov/bKprDXTckBV7PTdQhKXDndGlf3VqRlNangjqwTorzS3wjwAAAAJwGeEWpEf4etAFVTZCVL4SV8n9Kg6fE2aMda+98WXMvpKkUyXESP/AAAAElBmhNJqEFomUwI3/vDp1Hm8IcqtaXumQqIJt/IdZM/LaCqQDp5QDE0HJXXxYKzqroGplUEOXiTI56df1TOqn/+xhhrsArE1GyBAAAAX0GaNUnhClJlMFESxv/9Y3sbc/g1iqUD5e5Uke9bb1C6YGB9z3awZ4zQI0M0unHOdJnHA1rrnsh4SimxV5KpEdhAQ9gFOWZITFvHKgu1okFYOADb9R2Rozp0tWuEyuxNAAAAGwGeVGpEf5BbOR6mhp16rXzs2+e4d71TyW9KQQAAAGZBmlhJ4Q6JlMCN//vgOOB3OR7sTCb+I0045cZJEQ4U+L3XCpcnHOJwFb1e+Q0W3Q7uM0aQ6yv11+YTtxoOdsy7RQ+LYYJDbR6Rd7L5XFyxpKDB+0OWkFw+YLFioVUh95SYqxqecX8AAAA8QZ52RRU834P7bhmByo0v4xgbLCddxYxYbspCwfmTeE4h4m6gzc5yXll0xuqHmzfhwYk8ZpP+UzAb50WNAAAAIQGel2pEf4EHlcUhXPIKbaQj7c/ATAoSkel2SJNrDV4kgQAAAGxBmpxJqEFomUwI3/125QhMIg/aNp4kCmUKfLpf/QrlLIP5B/oHjyqjvBOioUOoMoL20buyufrF/UIt3sv6YVyWShlKuej/5mcBImU0X9mB2jkyIoDrWUZBjFA6haFPIB5+QhgWaDpI61usiRgAAABMQZ66RREsn32aKglvf2C9+Bdon/+X9WLOPlx7m1yaoyHmMFFCEI61yhs2ICjfrbx1xwg+UNHTSx3O+VuFL6oUO/x14wF/ouCv7NddvAAAADEBntl0RH+KcQxiH9X/rcK9KC9kRdt+HJ+G/iNZPTGZ+diB+2gxK2Dg5bXzZO90e7yZAAAAJwGe22pEf4Z1uDH/deStpiTEc9tyHnyaVWKk+d3zLNz5R2rfH3mWEAAAADpBmt1JqEFsmUwI3/wDwzHAd5jvuFkDKQEVOoZiZuzsM3wXcA8F3ucQwarCn13TkhsIYJMiT3kh+1+BAAABTEGa/0nhClJlMFFSxv/BJAf7z+EOU2CWqKManlZ0LpOsfJcBa0PMmIuvVH9dbvAcE5UzEdHVQbTrz1N1jb9WiMOxu5R2cy5A81OxB+5ZBnN5eAsC32urtdTTKK7QubeiVOo7VylTo1bcPy3DSeA7L7gshvKbqb6yZRCFyi4whKSJ7RPi0+A7uR9tDdfvNWqXsP4NZEdkTw7nJJscByribz64QmP/ICkD3ZxZqkoqOTVn6dlY+P0SHMy6N2fn26w0kDbCjyPfS1dZblvZxWjU616LYFDlf6NoIU40eHrlDbvSs4WEm5sqj6lAdmcaMt1R2DDkFv3lJfP04qVFUwKkb4diB0AgVCwqoCaqflufpSnzlUP5zpGT6XXwKzdvxhIiAGd/HKUkMg36acC0SDhcMpk2kz57Q3kYvM63o6ggTKO91qvRby/6+hO4M9+BAAAAWAGfHmpEf84wSggd11AiFSufQxM9Luw26hRYDDKsE9vSwAn7Z3kp0GuF0Eqpea5n5unP9KN0c3L3GkYNAsaXrACiTeM7/yz8z1LY9wjvbK9PNQeYtR6wMIAAAABfQZsDSeEOiZTAjf/zunZ7sMLNRD1X4hTb1ZFO384nK86xDwxI5Vpgs+jNIRsolIEfvY6lsaVpHcZvcrFP1//9wB89zGosviWcqEVLIZ7/+WmTYZWLoPL9M4Sq38lHkYEAAAAyQZ8hRRU8n3sgpbafys2Rudk4U/87NYQWWpO4rxUES/JyvoT/lfbQ9OleOWWIZeQm2IAAAAAoAZ9AdER/gHHeAx5KQ6o8Bkl+v/9P4zpRB6XdjaQVIDAbYhBhxO6UwQAAAB0Bn0JqRH+E1Cs20Vz506WBTrufbbvFjJQQkJ28wQAAAD5Bm0VJqEFomUwU8b/KCEPdL34SY/52bZ/rGfsXPf5yCTDido8G77Go9+K1BJRLL67CRLKDtmUDogd5/yp2dQAAACcBn2RqRH+BD26tG3o2sOTqJd//LOf/0exfYfKLxm7J3qgsoJYfg5gAAAAzQZtnSeEKUmUwUsb/++A44IUf87wFjEcFriJhVhBdOObtZMRA7H8L7O85Qi1Ir/23xQWBAAAALAGfhmpEf4exElWANW8vMvakG0lwp25Oh8MHyFQGujxVcqIRtUex9+Z6rsJYAAAAUUGbiUnhDomUwUTG//vhUdhCkgtD/OKNA6dG1osSc7TN7EVmJ8mHYKHJo0kuzmVJ2OcDMUMCsDRqChpG0X62F+dsmKA1zveQZSPh2EvbMNzsMQAAABUBn6hqRH+I4DsKsOugF12nJUDYjiAAAABoQZusSeEPJlMCN//07ziWU3i/kpvzHFCOn2CcItTi1yFM4xbCfhV3Hl7kKDI/9hW4d+LvO5FlZCITgCrsyrc2JS6CCKYTrtDuaRaAvGVwNn+giCSu/6I71EcdF6AAzs8YoRZvHfx+r8EAAAAuQZ/KRRE836foBN+l6dUqFIqLR7fdXGqeg/3ZJ/83TbhEn5pkA8K9kzYNkm5sAwAAAB0Bn+tqRH+BB5XFIVzyCm84fzzyfGIzavoEg2o1vQAAAFdBm+5JqEFomUwU8b/5u5inuMF7Xq8EOtQO4YzuHe2WcPkEeIzleyoHzUSRtXdZDfWefPY5FxPuebWlVj4LVrBtGv0RP79YCO7uGRp8no60SRyqygnK4c0AAAAxAZ4NakR/mJ3s34Q5H/1uFexTk3VIjFEo/fwfhqCP78QfAaOXWnodXpVWh/EbwDilHAAAAEtBmhBJ4QpSZTBSxv/8AoC0RB/Gq3N+gqfQl07//0asYn4GjHzgmabjVHeCeHimF9bt89uzqXPNoLLVfk5lK+SEWZBMH+jdIMD7NJYAAAApAZ4vakR/h+Qj8Y/7ryZBazBlCQSff9dCgZHsO8153DG/1Iu47qZ9y1EAAAA4QZoxSeEOiZTAjf/8A8O+HA7nHfcLIGMIzn/UI6nxZeLzvaZpNm4NR++zRTzqakzFTJjhdFc2sWYAAAAvQZpSSeEPJlMCN//74ByBNNlS+GOnj/iYE7uR3e+eaDjQFr415CBPX8resUSqg9AAAABEQZpzSeEPJlMCN//74VHYRK5/s8oAjSJlNV182JSgSuYDsNFFvPml8Z2TC06yUjgNGaIT+MGdO85RZn6ABwhLfk2G8T8AAABHQZqVSeEPJlMFETxv++A44Hc4/O2I/eN0x9t9al1x12ezWSFYh9FLhxu8aTceePsZOcOf/noBS072io579hr6GIKYfVDIPsEAAAArAZ60akR/gQ9vOA+UWsy1Mn3l//+f5lQBPWWuO+aJJugzYNwOMU04nd4nwQAAADtBmrZJ4Q8mUwI3//vhUdhCQoq5E2kM7tSSvPdRsTqN2Lqyx1YynzDOfYXeFZW5mbzBzov4IZCErRmMoAAAAElBmthJ4Q8mUwURPG/8A6IrIENlUdHg1vXava9lPZfzExXzXgqZrbgK0CpfLrF03F1J4CM9W6Uc/jlRwpOi+VZSpqhCt8/y775AAAAAKwGe92pEf4egL4A9rxjL4OGkdU1tP9jDyD8YL7cVkl98cgRXNXC0itKOXK8AAABMQZr8SeEPJlMCN//76z44Hc4/EYYtcgPBH9RgO0Iwi8OXMzKNpI9izDnv5+poRBie+02otvR+9Zng2UG2aLxEiEAQa2EyHRWuU4x0ZgAAADBBnxpFETyffSt6tLJT1q5nVhquTgXJSFs0fGTiWZMvJijPSmPEtnIL2xGLEmLILUAAAAAjAZ85dER/gw5gGJbwM9dZrSXxrWmk7SS7YSzCpHXtnb7SsiUAAAAXAZ87akR/gQeVg87QMSja2dC/csmeH6AAAABxQZsgSahBaJlMCN/74VHYPZq6tmtxLh1sxyPaSRfT8XGQqvC/Q4343DJxZXe1Sjk7F3JK8Ga+iJBjybN1BsuRRnEENt4FtlTjDMiyslD/QE2S+ie1BjE0ye1VekieYTcbeA5uWUCc9oAS0oHJc7DDprEAAAA9QZ9eRREsn3/Spxn+3+PADN8tr/k8dm4t2BxiP+1thpygNS5/1ulcsz7YQxXLmulOZ9bSD8leHIfywjbusQAAACQBn310RH+HvdOD1Xbhj4zYGbutZl5v6HqyOaqspaMKaK5qWeQAAAAwAZ9/akR/gQeVhtGJPhygpFSbn4CX0RYdxtJU+44KMHn0Oo8JLVH4iFlYBJyRuWktAAAAhUGbZEmoQWyZTAjfx83w+eo3udPz/wWGJzqAxttakOUQ5KIzV5rlta9UDHrdCKjDiNALKgEJA6puZDaqChxqlBxSo1lMmdEYH9x6+Iw/G7cBXYYH/vbSuijBJl7Iv1SdG2SxEz7WKOVaXQ7edDngcN7BTtsHxNL9kKKH2vdnV7h3Up4bBiUAAAA2QZ+CRRUsn81bxhrs9aLlJLIhMRyrlgrjTzbZxY1kexOB+js4nC73oqm2oS8KXxFXxDDK0y6hAAAAdQGfoXREf9PejGW7yeA3K+Ashsi9mc5rzdirkZ1i8F7plCYqBBB3h2i/Y5FflOyMyOlgUqYL/G9xIzGl8z+0IWKUH+x+Fv9B7/aREDydtwyQInEGxGayCrl3zIGiZ6TFkkOalquRg4SIEHzOpZAOwkut3sA6cQAAACYBn6NqRH+Gdbgx/3XkraYElmuoUqexzkujzvNedwvPC5NXjL0QwAAAADlBm6ZJqEFsmUwUTG/8A8MxwHeY77hY7su4Js5ePegnvwLeGgOmnRghK+XuEJ9p1ZqTpeEmvxKPv2gAAAAWAZ/FakR/gQeVg87QMT0JOabD8/BfTQAAAFtBm8dJ4QpSZTAjf/vhUdg9mx2D7TmdFz8qn2yy4xD/XK0YWq6bXtpAIDSMH89fnlAnLzoUK41nLcHaRGQ9VYKHwPVgSIejbid4zu0EWxJ7r5+L5ivAS0Wc7KiKAAAAU0Gb60nhDomUwI3/+8Gbu08WB3vy/AdqEakb767uJfNyZA4+pjSLzSTlkccggQWj3N8llt825rufA9wokyOuXzQUeglOEArfFlY1GPqmrtZTf6vhAAAAMUGeCUURPJ97IKW2n8rNmOXQIGJer6mMWBu7bBqazEhQPzV3PBt8Hffu6rEFJycnZ/AAAAAiAZ4odER/h7D5WmcyS/Um4Sed4v71PLKjq2xzxVZO257qOQAAABwBnipqRH+E1Cs20VzmQO2GD7g9FZUkcZEXqdpxAAAAR0GaLUmoQWiZTBTxv/wDwzHBCj/nXShHqN21NSADJU/Jvj4ZToB/1ZTomGKrdxyAf/iOIoo9DdUgNWCwuF0p6LBGjn9Ou8lBAAAAJgGeTGpEf4EHjJ+LiOn4U4Y39SmH/rJxyXApfMRHQp7a1U3JiFaDAAAAMkGaT0nhClJlMFLG//vgOOBCV69dVkGYj4aWcO7kQz+XVCeG9zDp5MQEvFDCbSD/y8jgAAAALQGebmpEf4n37kAo6ubDhZOY+8dpADGCpKbr5jcNPPpTe79afnKaDNgbJdD1QAAAAFpBmnFJ4Q6JlMFExv/7wetWHYG/jzvTdZyGtgH/DXdI0sh+SMrklDv0lY1O4TvP0MR8DIsygALU9Fs2WbnRqnA2PULv/h+6+vnPyCATdnYbIXw/4lIPbTxVj/MAAAAYAZ6QakR/iOA7CrB7i4FdAnsFNG/HYixgAAAAVkGak0nhDyZTBTxvyghDzpi4EXYDnbEfvG6Y+2zMQ/Hb9ovvg/0chC5MpSxOmJq9zFtkv9T+YO8D9ACgwbrIT3dU6sd/k0sYujXXO5r3U7FCF0bBVqsgAAAAJgGesmpEf4exE2/oQsw3fbuNFR/+SKnih1Q5+0sivYgWYAQIdsxBAAAAMkGatEnhDyZTAjf/++AcgTWtokgude4vE93Jud9I9whdUFDZJTfR+J6YiF8XdIzfNb/FAAAATEGa1UnhDyZTAjf/+8PuXdp4sDvx77lN0KgU/fpK8+IINR1opM0Pu4CozS8l0Ud/us9fntwO4IG0oZYvNBRI2N8a8g/0sIn+6f7QQekAAABZQZr4SeEPJlMCN//KCEPQpIQEI+hBvIFMoU+XSr3pPEb6tM+0KTmzi2hUR7j3nqTXSRaGGG+1IgCzYffvSdjj9k3eSImzEFiL89xBSAuBHuAan1qCm0SbL2AAAAA7QZ8WRRE834Ryfl4EPlupkoiD3UHAn/tyNF9cYKC+bqyZpWmU1htKRBn8v1v9le4Ox3FhOYCg5NB+ZdkAAAAsAZ83akR/hNda3bhOJUZIUUY2jVhVdrmxzmILX98ciz/mLhHdJSmfczV6U8EAAABFQZs5SahBaJlMCN/7w6grIDges6Q6N8LLKBBR/r+r01RScaGlerMw4mNRGZeM8kNv5hlT2J2qEHTq9XsVBx6y0AugLYMoAAAAZUGbW0nhClJlMFESxv/74VHYNZiODju5GRyMj9wlHqF1Lr7pV20VHiAMrHfGA6p0v/Zr0/dcP6yMHwgdbTTN5Pj6VJC/71UkDQHF2SgOkPs+SUevM06PMV4Dc/8UKpkKq5MzAfyAAAAAGgGfempEf4k2KJe5dxcZpOfMZDO0cRrAfBG/AAAAcUGbf0nhDomUwIv/+6w9u1GKETDvOv56IQEhONuu4MWW8U5koOsb6z0DMrcU8frYGb496r+0cR0j8/9YT3bwj6KnqNIzzKxPxz20387LsrGHGueynDGovxAkh6APLCz5JN0en8GfJGJgIpTpHDn4oF+AAAAAP0GfnUUVPJ99biH2ofqCeWdcOycKd3M3sDTTOXcd9oRK3Tz1jn02+G3H+x161bG+i7O4manzYReZ5zgb4/OHSQAAACoBn7x0RH+AbCaaY6JFZrmbJvY5TSCOfppisvf7jhPHXneiHxtzcLARd1EAAAAfAZ++akR/hNQrQyxiVLMCu/aBfsIZgv1fU2C7T3Ir9AAAAEZBm6FJqEFomUwU8b/8A8MxwQo/52bZ9Gsz+sY4O2RS3UuD8X0CNvCqI3uAP4dESewlMKkkQ0sYA5go9F7wlfSnLQu+tTWBAAAALQGfwGpEf4ew5DFDghAjpjv5rOLhqU4poz0tQCo2nvYk1MNk0SbxLDlMLbaRYgAAADhBm8NJ4QpSZTBSxf/730txLjWi6sL/8WkZ0SxdX+xgko4N3jmSClvhOh2sSC6GI34PiVLH5H6DQQAAAC4Bn+JqRH+HsRJVgDVvKzmQczvZIZd7ldG4fIdAbM1M+rxG1R7C8fxWQE6gkEOdAAABb0Gb5UnhDomUwUTG/2eshp/CHKaTLMVIv8yljz56U5Qg3GnKqoZs4y9mnt42jrRlYQiQvaz8jEhuvC1SoxRX3jdtVE7CW2y7mknEG5FfETEOLay9NAEvDLcXz3Kvp5o6r1Ecn5DKuRfL94ZwgC1vfOYsAR70mWBfVzPYmO5bWIXHrsivWz4TmSfJX8BTws++We/rWq+eEIZTOpQXO18Z4TBZJoJNh+NQKQ888eOePfvsvOd7ExGYc72AMj8IB5nA3qIIpZ26bWxdEigqp5WqA3CPmQ1ArglG8O23AW9a0mGskFVKIURyO7SDrMfvuSpM2Ufg2ko07dDeVY280dgvmukE3rlhobBVZ9eHnJpB/rDydTtHtcgUHxLRRN2WPaWUwfFExO2NlPxDNJNoifurDRR1TOE2xSpff831Zq9dXYJqU1zKACLUfsS7gTYxNNWcQ7fT37PoiXtJMLOR39uru4L0h58owlQKSE0dbvzeGUAAAABVAZ4EakV/6f4hhh5j4pezXGmc10booa3ca783iOnUPxnKk5Y3HfubpAwe+uDDln7x2eceSurPShjo5iSao0hBxT60iJArFLPcDKmG1k+1SrnGiBiThwAAAKNBmghJ4Q8mUwI3/2esg9XV1nn6OqvkurteiZEF/cX1g6vwMjx6DR+rv8oeaoyhuEGc4M80j7FCjt+C+gCKRU1EONaFpssJsG4r4wx6++Vl4OFVBDT83lw8NrDATf3vTsciLFGM8eGrhKs8pvOB4zbD2z/wbkYod2zVrCHakD3w/D1EFZvcK9tKs1Nr9e2ax3We6/gOaYcpR2bYNG1etHQuzTshAAAAMUGeJkURPN/smxOo9ZgcqNMHlkhDvC9wqzrKarKnUnQpTSCQZIsbYKJtUJkWtnazMvAAAAAlAZ5HakV/hwRZbYwm85/cCb2d75a9rssTV2GyXh2ERuMyRAZOQQAAAGpBmkxJqEFomUwI3/vEgWa+og3x1dnJ4PoDdBe3/3YoMTUoB3TyU4Ej/zJxdOfoW4aHRQ/I+nHKZhqGdhrBqoC886/vF7jEJplVhBx/pyxBeWOV/VwqfkCrurEj7aiFkmd0eJF8WUso2Z6AAAAALkGeakURLN+CHmGW7/MKreRIEd/8QfLllKWohrf2XGUEjDty3V/zICAlat/0cmEAAAAmAZ6JdEV/jY0N0RG/BawsJgKYUujP+rhnqWcqJ30YzX0c3+Yd3EUAAAAkAZ6LakV/j/S/wWIll3CGdQpmmnlHJJJMDXC3Dufc23w6T4+RAAAAUkGajUmoQWyZTAjf++5XDgdzj8QavWtNM5/v7SWzbYKyz+gr1lvvDdabJmin54HFA6FS5USZmz7Qe5i8GrS/OeBhEFUi7sU8j3do+iQgzZ6NN4EAAABqQZqvSeEKUmUwUVLG//vhUdhErrGwMwT2ApJmJaluur8WZXKPf3cxs7c2sP/Xbjfjlt+FJSs7mUuT13MbzM+TrSACCZSMBoDthPFifoGGtkxEnj+e8h73JprN7phM0yaCdgAO7cfoh5J18AAAABgBns5qRX+OkjwIvXb1vb5y58tHlKJ2YgMAAABUQZrTSeEOiZTAjf/7wZu7TxYHcd1Yn912Q4e6YwUip/A050OaFjhX6O2qQXp5VjUL059lIsYYgSQ2jExa1CDIwChKJheC4ssji0HjIrmDdIWl9Q3JAAAAOEGe8UUVPN9/4NuyfQi2HxhNnY/bPczyhQbbExzc0zWVB0vRCS6njJQu39UWTWCDaVRAWBLJEnaAAAAANwGfEHRFf452ZgVPqCEm6Wp074F/bJrDI8pmAHDsFuIJwlvPUbZo3jO+9e0mbvECpBl/WwaROjAAAAAkAZ8SakV/isvpAhCb6whUJeGpnwag9sUF2UycKFAAIxuGur9dAAAAVEGbFEmoQWiZTAjf/AOiKyAj+gnj0a3Gva9CNJ4js1U2x0aCm2qxyr5dY19k1pZuImuv/8peW/h4jD+21kSZLMKNT8K6Qp1NQB1Oz3dayRj8a04cwQAAAFJBmzZJ4QpSZTBREsb/+8PuSK8cB7lWURgB2UuJkbbsVfmYf6uGo0ihSHCGRIUeOkGwszI/RafeeJLVFb72xB56AeE7RHwA6KDK2FMZEiHn9cDBAAAAJAGfVWpFf42Vo2qm00Dddqve1Cwh67h8F7hLVO0VUOGCP37hFAAAAEtBm1dJ4Q6JlMCN//vDfCqSMgQ2VRj3pLESWqFJI9xvPW4RklANZ4xZHF71Z6iQq6DzlYFf0+Q/xmZ/fOpxmMDAoxgM2EezYMdxGYAAAAAyQZt4SeEPJlMCN//74DjgQ7ZH1LvBkl9O69N9KO37jJz7QEaANfGvJaL8Blb2Um8UToEAAACFQZuZSeEPJlMCN//9Y1z/dD7lReDd/jih+Ld0YtRndvMbeeMTxfw1CN28oAhAs3BFnwK+WZTwEg2f//cT8FjRJI69DXHiU0hpAJ1L4ZcWUkQNn9/F3WTm780WFIBn83SpljWZ8Vra05F6qUNyjAMjVTOd4yHsZRD4BpeH6ERAEEj9uWg3dAAAAEhBm7tJ4Q8mUwURPG/74Djgdzj87YkIIqfbdVCHiXmoEJTcLeAuy7kDg1f7oEo7JIcrQiC/U+8t/yN57HX2mDX0MQUw+qGQfYAAAAApAZ/aakV/jnhoov57JnRJzUQ582/A+KqjcuQgc99QQQxsklKC4RUIecEAAAAxQZvcSeEPJlMCN//74DjgiVvqHRYbDnew3k/WoU8O62m6cW1f4z4f2KkM/6uwtCkoygAAAGNBm/5J4Q8mUwURPG/9Y3sbYDUySihk0ZwYHPZc98GT3bkkuNK+kxG9hw6QhCPLCTOPu0cTsYadgUqI051fWxMxsY+qGfVFtfK/8hGLOVkVoaB0SjGYoDGFSY+reuTO1+te1tMAAAAvAZ4dakV/j/S/kR07/7I2ONOQxpj1EA6tGW/N30/w4dgSvirINWao4hRz+KM4mYkAAABMQZoCSeEPJlMCN//8AtC7wRdgH2cIbrCoMs/ZXvDYX8DxXmXaY94n5P78/TKQXUbvgpq5JuPge7/U+LSYd+1eMFPL7ErSRU/SQWbYYAAAADVBniBFETzfgcsI33hBxsPekc3ZDBzYdeklj7D+yHKGhCeLw2I2cmn91BAltpq2zXE35Ex0sQAAABwBnl90RX+JF8eO8NhiMY7e3XCBBxtbicm//Ab6AAAAGQGeQWpFf4cEWNCU5dbW4vKx29yJfhj97H0AAABhQZpDSahBaJlMCN/7x1VfiFwgJBqttUqkagOHWBpqzu3I3Z1ug1iDen9RqMLuWwSJns/FSjmbZC32wpaRGK22xjoTIf8DNdWNE7Xda1dD+nvhdspBUZZ+XrBk1Hqe2E0wKwAAARpBmmdJ4QpSZTAjf8FqjzHg2B+Dqry9XK4oFY4kpye5WzL+gZiytMQXb5+G1o3IxzRZhxW9aZpjLbUsacNdqbjHaYZ9CmpTlK7FK44jbJ1larK+ACOQSr42285xajESkZG6raHieCShaT5/xNHmQa+kW25//2rjWsefjIOVuKFkvGla54JhR+X5yYFqm/bC16PSKR/h99ZwXzxEyHSds+3pr2CMYM14vyO0YKx2HjFF2rVw+hoZHoyQ0JcL/ReGkwZ2fNhkLNreyaduEAjHZAAX9t8IFb0Oz30qFk8O7KgS9ZdqStAL+KUDJIhD3HAAp3xmSs+BqTgOrQ5PS+Ace+fz0i3bKIPCU8i8AJwZcWiUwBhyv7knwrO1Y4AAAABMQZ6FRTRM39foEqx/3JF1qqZd8W//5kxNE6DYERz3YbVOw/1jyVC2gk+UmHuYnrw/mUZrr6ik8rqKlUQKXKQUw5GLsaQyVoSX08YJgAAAACoBnqR0RX+Gfrvv7yJuu0tSB+P/+h4yK0TNhid5Xn9+DciKgmECSbWRRssAAAAtAZ6makV/2fIPxtyYTe2F7AUed/9kq/9NVDnFjpKyZwUbJZWypoABH7rTJteAAAAAnkGaqUmoQWiZTBTxv8gMvylDZaeIJIsMYkC0ZGm9WiRczW/gYZr7/QtVjFGHEuTg4I7VYpSYfgz+C64TnuAQXBvRwRlD7R7BrUKTMlAKgwmktl2YT4IHIivf+PBMSAKaN8pB9Tfiw85AUHSJGSwlVnsXlyHQ+LCtY7tNz+KOQQRHBQtvFemUE7/owFlAZAQX4ZhH7H8hq2bMUPNDHTHxAAAAJQGeyGpFf6zmbbzG6yAo5agB815PT4NIYmF3dRluJ6VKQoZFLbUAAACKQZrLSeEKUmUwUsb/x6mJAIU6OOaejU6jOCEaJq93ye5SDx3rwHIMSwuPbeEPkFOdU1+52yWbGid0yNY2LO00ed66IioKOPe6doVimRketNB1H1k9nKGLy24SyswHs7o8zUFs33/dimS2jZD9vA6qWcVwUdD9loJxKlVTdWWCY05iyBBnRI+HStFBAAAANwGe6mpFf9DFBS+peC/yI0i5BOXZ0Kb0Jtornv1yHk4b9Hwis05VfRcTVsB0+56ooWThOraMfVkAAABWQZrtSeEOiZTBRMb/++FR2EKSC17OilzphaOnJfb+b2v5DMpyB0LE4N5XbFR41vnnYVmX71a1TgMbvI4lkgGITPRwjYOjeSXOpCDC8xaYEVB5J45bl6EAAAAXAZ8MakV/jpI8CL128HDcV88rEWHEY1EAAABRQZsRSeEPJlMCN//7wZu36cWD2WN+3hTaxcsivWm9AgbVBHu/qp/N6ueiOGMMY/MGj/ljaJqYSQxBu7OQessdVRPOqKh231KwV4eUePFzdTULAAAANUGfL0URPN9/4LLaSW8U2BK02dj928jKfghGK1PeCa/SOciAikbQHsE/9VxvCM+EdP4P3D1QAAAAKAGfTnRFf4Z5voDwlwQuL1OkS/asB8/ACxEz3gJmWfQAj0g1qyhfBTkAAAAfAZ9QakV/isvpBhS2mK1fAdwbtVb5YgmyDpirrorcHAAAAEdBm1NJqEFomUwU8b/7w/gmRt4SJHtq/4CnNuN3iEw3xnwnXZMEbxFj861j5ny0cBWKgDehupwGrBYXAeK776eEXeWDNQa8lAAAACYBn3JqRX+NmH8JJKZfXQ2CWI5tPFL8LiOYvVF01NUI0yeIVFZaRQAAADBBm3VJ4QpSZTBSxv/74DjghR/zuyzH8Txd5XruzDOHv9RKkKOyGWsN/oiFcRlf2fEAAAAsAZ+UakV/jZiVZ6ef8q5U0fD2jgZreGAFGFOTGZSMcvYdOVWw2qIqwnqo5s0AAABPQZuXSeEOiZTBRMb/++FR2CH1+D8+1PRw2eC7cPzHLlSr4N5kdbXkTkaLMOtNZYDvN0G+BcnrfD919fNbcNCPlQajdZJesGGRw6UvTyq3mAAAABsBn7ZqRX+O5cRG3XQzJOw4cZkp7na6PmOyn1AAAABTQZu5SeEPJlMFPG/74DjgO8yArnokfcqY+3z7vgO3k81KQrKISGhM7SavcwZDj+Lf/9hb+KMuOUdKKwHApugri91Y0dc7DkP6jYD4KNCu/ZrJqskAAAAnAZ/YakV/hww5FAKjLlfq6Ljf/+1coHIsbogAeQqZ3SIn7qkD8EJiAAAAL0Gb2knhDyZTAjf/++AcgTWtokUyBqD6lhgvnfnJDd8Cx3z95yizbJFdyDlHg9jlAAAAT0Gb+0nhDyZTAjf/+8PuXcpxiwPb1g8GbiEE9IET2gffzB52f8EdNr3u9CTCJQDIJOrkhycCCgNaBG2HkcB5oPyRivqaDJ0+sBXnYINEToEAAABZQZoeSeEPJlMCN//7w+5IrxwHpn4mcnp3qC8i1PB17ucNb58sAdjpq0E8ox6I6nyB4tsTME3vaKjnv2HOUOKh1xJET0cCW0OPsSokk6kSsHBKYZjYC7L8lyAAAAA5QZ48RRE834ISoy161qatNvHcx2P0v+2schVH4rQF94qpR/q56ZD/z6y5NPUUhWN3Zgno8YGf8JZBAAAALAGeXWpFf4/0v5DikvANg1WUUB+NcTU/sSU3fSaYjuWP6ked8ltU06WZ15X9AAAAWEGaX0moQWiZTAjf/AKPrvBH9CwqRiXu1U258fBKSfqjqjBQd2g4Hb33Y/0j/xa9Lykql+WIXGWh5MAUZG4/duV0oRy89osMMauqyiZdhBQPluQQB/BxcR4AAAB9QZphSeEKUmUwURLG//vBRkdJOweuoGYmpH5xFZVLyGMWYq+kIwnpHpVXYdaSzJAWmGTPK1sFf6Znkp4MDLWKTIwxvdVJ52cC1rehXCiJGg1PAy47lCc5h0zQ7J3gBJjV4y3ubnbO8ylb6w+rjLP8+xSx/E5wYWiF5YEZn0kAAAAeAZ6AakV/isvoz+8v1Hb/P0q6S2yOKt6t2F+aO/ZDAAAAfkGahUnhDomUwI3/+8SDUbxwHf2NVhEpR83q6f+XXT0emIrLG/GjTA0ZQEsuf4ilRudQYS+kcaQP6NH/n47oeDxDgr+9I1JylC2mc1pGpx+CT9kp21tNNfIR9M3h5hKBoKhhH2mmDCHWwTlWZrU8XQRkROFjiGHTXpnI2TQPfwAAAE1BnqNFFTzff+Cy2kloXh3IfGPf7xzw0VsS6ephFl5SkLNj4sIWol/tW/UCeEA0PDXCAF0fB2wHJQCr+HvgtkKvByBuJ7pSQYYcBBA06QAAACsBnsJ0RX+GfrSKAZRt3+eUjEP/tYzGZCzSsy+iSHSWdDdvkdj7BDVckq9AAAAAHwGexGpFf4rL6TwoJvceXyel974RIYNEAeX014ZGez0AAABJQZrHSahBaJlMFPG//Xblp7QgISvXq/43tYQswnlCiPCtggpbgu8Gcd0XnodgB573BnhWBSscJCb+QLXOkHexkQ6chMI3P2elzQAAACoBnuZqRX+NmHWKNnyOWkElh6MU5PvSNQ8zujRaeaaC8dcoGksLYZPxasAAAAA1QZrpSeEKUmUwUsb/++A44EJXr12WY/1R9NcvW6Qg+sKr/ahi7wg47WJVdC+uliZS+rcEC7sAAAAwAZ8IakV/jZiVZ8+203fAhmP6SQ+/2wesW6W1rbx/GmEeycqK/2uzkQSaEM7KzV5AAAABAUGbCknhDomUwI3/wSZXK3peGISKDGuEfaX/ruGdURoMWCqwPVNEpOv+ST4slfLn0vfGnqFijtJwvKuTeo4gMwMlZ12WAx97eksVcC/xZyoYY8gEMuJoVZGXraJXsn2AWjQmljY++Yn0UGbvnrgKZ/T7K9u68fYiopDXks1XdSedLDYdCPWc87cOG2LaWaYKA+J/Rcrt93JbbJphn/RaqafF3b1J+VrUGnuLXN10mq9Ji4vhD2gcCVDgU8QP7fngj8k++0v6zvjDx4MM30Unq0Ro37cLz9Fp4dYmyVWi9EBFqKvmSgFrSgSu7qNhL23/SWg/5DKBGi7SSMBZOnX9WLFZAAAAs0GbK0nhDyZTAjf/8pF6MHN8xvhomJb9VTbd1oufHXVFvOUqiH8wtZul84FilJh+DP4Lre0QjSMquDwJB1Cd2xjlRf6f/EfPHC1oSOv7JiC13e/9H5Er1hNYw4FigKpqQ6OH9tzFrQITEYT/f3gAuicOt4g9cvzkdFsMYBTJT4tnA5QEdXJC8nEAXWHq3Dr4PL5tNfHtaPyUqNoDKx0fsKN4D/mNRoRUv/gxlV5JxljuF+HlAAAAnkGbT0nhDyZTAjf/0Y/fwfBPYIt2Ih+oVDWysEi1BfcGsNc42DZ+iX7NLlxL5cxymUJMWBhkKF68OtNiy9cj9+WIlc5Mhajqm9tEk7Q5HF5llNkQ03EIGx1nktYjyHoOPCRnfHGTijsZ7HClYxAmgkK+zsXNRz61pxBrMmmixPZw347116XCKl706m6iOHYY0uNsJzXzWYUyQe9L9FINAAAAMUGfbUURPN+sXvx5/CiDiEacNaS9U6HfVvnuP/a5UQJlGkFDYR2J/hU3JICcol6Lc4EAAAAnAZ+MdEV/jylvx+CPXBU+8gKTE/8VIElasIGY33llIOtASe54iQcmAAAAHQGfjmpFf9C2JxDQM2FrtfcernQzbgQo6cOgAiXIAAAAPEGbkUmoQWiZTBTxv/vFEDSrIEiP+dm2fRrEThE5lxbkgROmjfPtCbfhnR3G4ID1iPeJh7HY3Y16Odh6bQAAACMBn7BqRX+NmH8JJKZfXHhkLwpdR6P4XTfNRE6MatJ7PVwVpgAAADJBm7NJ4QpSZTBSxv/74DjghR/zvAWMRwWuIY0ILoghoK+2U2wGoNOA06nWHBwAcUpagAAAACsBn9JqRX+NmJVnp7PpUa8FkOcbCyJNunfz/M06nD3e8NZpyq2EIfrhbd7BAAAAT0Gb1UnhDomUwUTG//vB61Ydg3N6LfuunSF3AkvNLoQBhoTQ7aoVG6tsYzS+JiPd3mYuiPkzCyWnK401zAoTSSFfgBeqE3OpyvPlsMfUq1EAAAAVAZ/0akV/jX4K9aDUHTkIn8NGrTJxAAAAOEGb+EnhDyZTAjf/++A44DvMh5EHnmI+0qZng7g7n1G6u5/Aqyb8TemoaU8nwXBzvg6L4ceEadnwAAAALUGeFkURPN+D+24e0TinWPeueGoG4sapy+ObB3xDR00TdCay97Igcx39n3zDVgAAAB0BnjdqRX+HBFlCUE3nRNtGVVInyej3t3vz7uZ2wQAAAEtBmjpJqEFomUwU8b/8IvJDICKyxCcdzT/mboi+flPNWjFbInDCJbwHXv4CBeLXeVVCSD4rO9Def9jVQKiZ3a5b5Bie9xheg2SLH+QAAAAvAZ5ZakV/jY0N0MOVPt8h6WcoZV1S20MbM4DsMfLQUbNE8rolkKvz6B5iwJ8mvWAAAABXQZpcSeEKUmUwUsb//CJ/zICOuOx4SMR3pARtPhD3XWj3rBkO6OrvomzLoOGgKYE4i0wu3UHRLX77PtaCBTI4Tfpc3gtTPNcmsPDp9zCw//wDxv5JJCJZAAAAJwGee2pFf4x17dolDZL3wxmEFIgT8nhorh3Yq+I4oAacvjWdHciZNQAAADVBmn1J4Q6JlMCN//vuVw4Hc5yyvTItaaZzkkkNHLJkI0HEoJdIeD3dJbzkXmvtq7jUQZd45QAAADlBmp5J4Q8mUwI3//vgOOCIxgOqHGQ4dPM9u73R4ALKr9BQr/nU61eDXxry868wQq2NQQB2Cw7y8YEAAABaQZq/SeEPJlMCN//74VHYNYZ3DjtyB3KfJ/ehNVbvQ16Epn16L7ahC4CDOiOWjEm08H+S+lGobDN1/4ONf5WIIOky9RhRHLV6H2eUZ80HGf0olbdDoIPuaghAAAAATUGawUnhDyZTBRE8b/vgOOB3Ocvj1S4/IKY+GUDepFZA7JQ6D5zJgHdmU0gV4AyG3PnUUlwhcyp1sj/2sUr8E80bHPfsNfQw3EPqhkH3AAAALAGe4GpFf4cEPpZ+ayGHCrzmIFkFJJ5iOG3d7ko318enKcteEWyK6vDU0azwAAAAM0Ga4knhDyZTAjf/++A44Ilb6dcZ4GoPqU601qIovFnKC7lkIeBkBzqp//sZLPSn67pyKQAAAExBmwRJ4Q8mUwURPG/8A6IrIENlQqhLiwoaJhrnWZk18jQ7ylRVYz4q3Y43ERkjNUVHICbJU3W2fPPQUmUyKqzovlWUqaoQrc8y075BAAAAKgGfI2pFf42NDdDcJbtDdK7SWJH5HwS2s2t/lM7DAAfmUTKKRr+T4pF7/AAAAEZBmydJ4Q8mUwI3//wC0LvCDPTwaCjaXZ2btf9j1UTqkuJrlFvJUAop1+VQkGMNp2GtkRPC5V/i7vZhSjenabTSd4ZI6Z77AAAAHkGfRUURPN99TaObewep0G/CenyT8ihy+yfmI16uhQAAAC8Bn2ZqRX+NmJ1nz4vAdigsChyYtloC5NV8EcFNRJflhn+2noyyo44OPhMN8LgKTAAAAIpBm2lJqEFomUwU8b/7wetRFII6Zm6QidC0oo+c0Ns7riq1h4ZBWxZZD6JzeHG5pZuCUPFJCaRVGlMRP6GhKFXz9x3DH2KAgdWA1aMZo1pkrtviY17lc9l1KlNpbFrOjLf/sY3BYOhEQqMpFO89k6PKn8sRKcq75K2ccdT4r+p6kGpG4mZrOF/98uEAAAAZAZ+IakV/jX4K9aBpZl6QE2eOe/h3YVFcgAAAAEtBm4xJ4QpSZTAjf/vgOOB3OPztiP2ULSp7XrKrWw3l4Bc6H2W69Q6oSoF8O1hPrRXjY7we9rAJKPYETi17uxoi3v25T98McjvFdpkAAAA0QZ+qRTRM331NqbcLGtWbMNY3WScBZQDj4EVhJKd+eBPd5akdHADpGT8hVb2bsA3vvcUMnQAAADMBn8tqRX+HBFjWKW0eRGEWy1dqsM/Zh2+r0j8TMNGs/fPMii2zXF0N6KhYyAEYzOLy8UEAAADpQZvQSahBaJlMCN/HqbEHRtdjNTKn+1Ks8Fd7U1CyW72yxoAL49N0mQ/rxd6hIp3QUmPFE+CBT6W6w2fg+GvkDoA5rxlfTH5iWO7O7X1aIQrUycIpZ26bWxdE69CfvW5Ump7bO49JM3Ojp1ba+sFXZJgF2H/6c/uab0RGBZJFVPWM6/2iD4z0L1ee78rMROduc0TjmUMv4F5e8iqzbVmWZxa47HR2pX34hnX7bqBD1Y05md2DIEIMx7XWDS2zFGguo8fbZRpICoATJ0iBpYXY65NY+onyXag/KvGTn9dqOkIZN5bgNwodskEAAAAwQZ/uRREs39BT+V5BdRxLLjvRBZ/8QK0LAnU4ryUT3nxiHiAhHRZibTVIPm+0xrfhAAAAbAGeDXRFf9BX5Uw8x8UvZrjTOa6N0KNc735HzeH8dCL7LysfJAtHc6CJl64MM87h+XsJUfWPW3okcUNQiLbRiB7rS8Stpkk7rAmn2NBtaNQkEAWDZlBKfeymcqvzqnPiuE+on+yiIr3BB71tTAAAACYBng9qRX+NxKann1KePz6kv+QI+9/UjpccvAs8AqKrEhiFftqoFwAAADRBmhFJqEFsmUwI38oIRG/MCA7nH4jDF8og0z98ssIrvm0ivszdpDx+Nnn5Qb+Dz0hEkbmcAAAAj0GaM0nhClJlMFFSxv/HrNHv/cjBycDl2vih6keQw56fCiMs0YdNmpp7imPBF1y/lLKBnxMxn3M9E+3yBzgA4qoPmgxAa8dNdml0VORbm0KUx6CwEuL013ZY0JgiO+ELoSSDqSlaWOxREn4lWygn7nk/erfNPMDpvIC2jYz5OlVHNcLkNnKjrSDpnPb4eSDAAAAAGQGeUmpFf9DEtYTuNuuc8MFIfiTupN1BdFcAAABQQZpXSeEOiZTAjf/5KwLantfVy3VySM0SMHvsr0YmBrv0UL6DlEtaxQQHzUXE30YsdNChYm2O6V+hEk4F96/91dLwu/ndfGK7GATkcf47DIEAAAAxQZ51RRU833/gstpJbxTdFgmzsfvEjjhs3dxWp4ZZ4kPG8NZeKU7GYzN2OvKw1tlnqQAAACUBnpR0RX+NmIM6D8IxZFFGEALyP/EQHxDjOvu6mfZR+X/XdtvYAAAAHgGelmpFf4rL6VHGE31hHDEGQjtVce1yEtA1rtWqwAAAAE5BmplJqEFomUwU8b/77lcOCFH/OzbPo1iJwiLqg8angbnwoKqYau4/Tr0/Su9hg7+dXmJB7llgNXjWmAvP+6nn9ZyA2iVrXntcHxIEctUAAAAkAZ64akV/jZh1ijZ8jl2T9B54JFoOBhttQMJEZVln9G5lKq8BAAAFQmWIhAAn/74fcEJGW81yyeQ3yjIJuOQipq5DurWEjnW3UorRPL+Gmp7W/eH/ZijezhQ2AIPIQl4aHdQM9A6M5/MKi9xGlN69wV4R5VA4TQl0Sllg8188tv6vFr5S+P9cmRKtZwlxsQrO/PUH7ZGWYYybHVXFlEMY6Dl5808LYw/EiDr5JaoSWsHu4ebowTZo7v7qaSld+s+T31N91+CVDr6kktqJbNFOPAo3BvxQkltR4XwG+/rxvXbxkWnnZ4k7z5WpzDWYcBFPvUfTgrd7n6SLHLTp7bJf8LFMRC7GnoqmMyaPAzg75u8JidsVPIR29wsvpk/Mt9R6nnF5IJistydnX1nYc673c2swNm0ubNLgupd+k5ttRQ152rHLvBca63qjaEYz5bO6uRdZGmNbrdIR9DOFJYzEJ3T0VRVhLzRbgQhmzINEkMybfXp29t1bR5qya1mRWgLdBC5l3drBLQCcVoq19LY64FqmtYeURGFQ+dxCLISD69Irv9a0uyOxS7lVS3qeo5/gPB9YhIYgKQn6bpUQv8jq9Ep36A1DYX5VFcN9XqGkifkeY7ZbzJGH1tSe7aDoRQIiDO6HnmttmwDbytG0SzChvNEtecWmFp6+Ng7jqZmJ+LpXMqaVGcBbl1ehkdm5Vr84GAQKEq97RhEnaEuoPQsiUhi7y5CzHQA/q9/VpbR1zqAOP/78KNEfdl56fss7Y4UfTIZBj/Ef7ScMR78inLAl3j2HLP1BrnE2c1Zfp3fmi51t7IntWomhhRLlKfnfJrB056GmNbnqvtmZCvlymbrND2EqUvjJ1EJ+n02dodaM+1EsP8F2MT2Y4tjt1aBsu9gTs10CfVny7PzZrIe7ZNf33r8LTwpXSowlA4nQt9opFQ/ZJxaz8Upq6XgXe6Uqzu7x6Z4G+LEevAbdMu5y5ewVRMuV7HqXwlZX3FmRKJTwfg2u/onfLxKJGGPlejg4w6xVcIvQWpCCY94rqqRweESq52TT21mJ3ZWZPYd/bW9E4HoKBWndnm57+UuTSHIPQ35BT9W1Xnl9r9CViMvIQ8AUlOp7+ZZmJtCkSP14cMcCJhNiLGBPwta/Oyraa2eYe4xSkWML2OCXs+aPKBwLcn9R4ao5D1S55sxD7n6+Mg+z9jwufgfWX3u9ii7EiGIFM7wV1HcRH3cSIQZVgHx3bNHeN9HrMdz5R4HKc+XWQ6nfPa5qVL+jrteHlNOWPkjRWBZgYim28+DOXpSL63VQKM8wCalx/7UFzpmPLYi2k25Tn4L4RY2kTpMyq852ckgqQ/JuHNjmCM0OZxxusMdxMp9ba0q6z4Fov+fjMAzn8vKAAJKzFCfdFJg2MeZhfwlKaitFkiyEY8mnat72LBWhX8y2pFqa9wBfrk6vSpsSIPf1ItXOV5dsbkmQaOOvVHpr39w9kGJ06MPgDiO7B/qm2Bs3AlNEuL+WG8zu/vIRJdDz+RTRIk7vLQr0Rq8jEQWZ/euf8FD6BUNHSXZGM949sWmTK0eoMcmoTzlyGGCQgdVdRYqeVPDF7SHBzjIXvdTeaPeF8rdbHCdMCUN23hsQBtwTi/WxUcz0MCULeSfoCuRD9EomV9UBsyxfCUqSYYFieuAbg/IP3y2xj+T/9E9kqvzPRPQa2y1WxDCMsnMIy37FnZ/7LXXOsxfbVeew3U1YASM9XkT/y50YD3lR6WGJ8MYfYM98D8exb8s0epkXWDnGeNbeKqayRpRsYTvsgzojgHxpi5+wGkEyHCLHk/j2PjOYfUfWmKhOwj053UqWFuvky53Z+AuX9yWhjuzKAAAAOkGaIWxG//vEnVTkBINY1xBnYqV2GYveJxG6pO1ytvZdPbEXt5LdA7TGVX0iYnM5z2Z+J5AJNCh79n0AAACQQZpDPCGTKYRvygg25A/mod1FrlSV9Bg4fSRIk5ApsXKnR4Lw1hYENtNp2IdrhL36xanS2cwpV3ghaBQ0ENK/+61iJnFzrglvYgyAeUPfvh4+H91OONFJ0BQNQP7S+T1dX3IH4RbbT1dKFwTvZ3ZM7M9HZq2WxWBFVnr9ZP/zxD+foUgnXU6TFFBNN/llQnxAAAAAGQGeYmpFf5XT+htGpxm+ttzhrI2RGOltiWEAAABzQZplSeEPJlMFPG9uJuUZbUPQJIkFcwKAtahG4ohsmW9O7pJ3QXEf2moUd4qKRCCL0ZKGjcHbIVUkGfh0CDhI2ibslmbDuzOMAhXG/zv/6v/5h1pF6Fnk4OU0BdEQ7aG8XM909Bk3XHm49yHVcqbBFgXZxQAAACoBnoRqRX+RRVaIVD1YrnHAZf+m1ih7YZmVdEVTUohPyoJ0YoX7wwFZW8AAAAAvQZqGSeEPJlMCN//74ByBNa2iRTIKCh6OJLaoh8ps3rY7B+aDjQEslblmkDmrpFMAAABYQZqnSeEPJlMCN//7/0a5QmmqHJOeKUzB2eNo5H9glW5oLygWzEVYpridRvXn+2AASXaouZZC4r0hn95olSpqHwrE6msLwa+V4Ttq+yvbjDu1Y6/bkwQegAAAAGFBmspJ4Q8mUwI3/8oIQPYmOP/wRrPxM5PH+fJ0mFWAYBFh1Pe6xJT8//aEn0KlIlhE0/6putAFmsffwcLGQqaabW0OMe6iwCBaM3Kq2If4TnEQ/lSd0QI6JJrQELGDib/jAAAAO0Ge6EURPN+Ecn5eBysgKBkhDY8qsL/25Qam4wDRxrYU5Di/5y+wAXkIetqyGrdnqZuLCcwFByy7VklxAAAAKAGfCWpFf42Vo24OWiO+wC6Bf1+vC57IXLwDKAPdiUIokWez79lHX4AAAABIQZsLSahBaJlMCN/7w6dR5vCGznNtXd5bO5Zy4LxzmfFyTxf8VH6vp60ELz2qk/FRPVhaQsbXkKgrOqn/+1roIltiQuL600R9AAAAakGbLUnhClJlMFESxv/74VOAECNfo3TrZkzqu0KkrOZ9CVw2gjMRNMj1E9vxW5V1Ecvm/9qH28Rc3ffR5r+zqVMsLdHEwsswCE/8g29m7hFIyoSbEObNEnebhqoXvPeZCDCweyiSrLwsRskAAAAXAZ9MakV/jpI8CNGNvrae4VovZPXKoXMAAABnQZtRSeEOiZTAjf/KCEDu+rkuXwR0ltpBueVY4aihZ6XD3gGzDCyVqxs5vXh7SVn70nLyVRSeCD7n1+k76tXWOU8Q8cvlYby2Hnvldl8uqwEjayKEwX50N3/0maF6ny5kGq0f7Ao02QAAADxBn29FFTzff+Dbsn0Iut+RPjHv9yMXPXi4QWBtUuz84rCf2ne6MMkRmFjCTOGYGkE76dyp3SNBzeilrswAAAApAZ+OdEV/jZiDOWM1Uy9Vu0nWL0V/4HEe7fSSKWekx5hQYtiC/xgqvYEAAAAjAZ+QakV/isvo1lhN9Vf9YXwem/sU/GB4tK3zn/fFUPpEfgUAAABJQZuTSahBaJlMFPG//APDMcCEr16wAPFV1G7zGhko0fBO1LBEbSo51CNZnpYS3MHF50L4MAKNELOJvnOk8fJx9M+aEnAehC4L4QAAADcBn7JqRX+HDf42NiA1e6IoNY9dd0+pse3XDs23QkJmYymj6mByBQBlbsNBH/n+C/Zf0yypZhu9AAAANUGbtUnhClJlMFLG//vgOOBCV69dm+FVLfTWcUGQr+lSqCflwO8IOO1iVXQvrpYmUcqC/Y4LAAAALQGf1GpFf48rV2va6Z6CX7HZbnuYgcin2b1jZsuvhj8zrWfr2ZgMZiNJsrNVEAAAARJBm9dJ4Q6JlMFExv/BJu5KzDErXKYS1RifWXWjNyUkZ7LuAhIUnAIaOsVEMMQXit7waW63wzdcUIeRyXc0ns6O4xhmDwZ3ukGLdi9EU7w+8z26SXTja3sGWB+uU2PQaFr9su4rpTbgme55GdSV5Hium1b/ZY7wn0ziN7L7ptbE80Ax28UKDugysLJcSwfJeTU+FsCjXuJC35Kp+sq+nt2e0C5pikCgY4m4RGYJRGW+kGDU+yhW9YEtPxxAt90wlj4Nc2+02VtndN2Z1aSOvivKWQperyfJK+RiYH2Ua4M1Mbvt+F/gIq/S9qIVS9LvpsrQ+7+gxnXyJvs7t2yv/J7ljPiE5MYPyl8h+jNIjrv8R+MhAAAAUwGf9mpFf9Ei2Luv/RNJcaZxTnDa10UNbuQPAlqODVB2VxsIpmJnfjPkMJPHAq7sCkJAYLQ9TOcWdYY/w4T9TToTqOJl5XsxpaKH0nZtSdWfNN6EAAAAUEGb+0nhDyZTAjf/+8NYd2niwO+rMxV+zWAA+1JtHTaumj+P2U+PhhfgTzByUQbNihyr6FdaXSQST2LLyoRUr2SQPrdTZkfv31QscDSS0oRnAAAAM0GeGUURPN+E8DKIQ4/14uOVqEAMADp0zauP0LtNNslqo9KyDu+ewIdMlXYKjxCCIZ2RgQAAACYBnjh0RX+NmIM6D/Z46G0zEmw183/eHh4odUNrt8tMZebEjulWUwAAABwBnjpqRX+Ky+k78Xe04HkuIKPhhbbazcKxG4nlAAAAi0GaPUmoQWiZTBTxv8epXycoOFMLM/d1FqVzTzgCDwFATJBP1U6ge28/eu//0BsTgIFsoE4M0gMczwnHzJtkURt64dKffbzBHukviZ6wJ1DD10/gHeIjw54CCzmpIwySvy0eKrn1sxtgjxrAjk40cMg2w3WH2SwB2gwRBdRsg0MrzvNgn9GcNQO1ZTQAAAAnAZ5cakV/0J4+WxsQPcekdWueGIREStNdZlzFmXh0KEniThU6N9wlAAAASUGaQUnhClJlMCN/++FR2ESuhE+RatyVDTTzS0MEydH3+Vpq2cFMFPi5QKgi9czxuhaGemqsD7/JcIUhVdt886og+ljZbrvCIWEAAAAiQZ5/RTRM34OpqLzLMBpZQ293KhCxlhbjxhPOxHd4q2TqwAAAACsBnp50RX+NmIKsD88pu+EZBTc434H+uuuvMbHJAB0tuUvJ1rP2bn0fEH47AAAAFAGegGpFf41+CvWgaWHE5ho1aZOBAAAAOkGahEmoQWiZTAjf++A44DvMfnQGvr5aabtMPFRgZrO0PAPTbXug3OYp1ediQc1+3GfuUR2A+YYULrgAAAAtQZ6iRREs331NqcrGn20AwIJmTxeb5XHCMknUBifcKi/SGXaUJ0L8hQfqmM7vAAAAHAGew2pFf4cEWNYpbR5EYMikZF8nmrBIkdnpZdsAAABOQZrGSahBbJlMFExv/AOiKyAj+gnj0a3Gp6zI2CjNAAJIWgyCThvil0kvhxqkZ6Ht1BxxqXfrjYLInpyi9Q1SHRusH5zdhyC+Nu7rF+RwAAAALQGe5WpFf42NDdDDlJ55Cx+Tl+h4NevQbGGML2wsKXdErYWj9A80AousAs0tNwAAAFhBmuhJ4QpSZTBSxv/7xrXt8jlCNDUuxflUiLTs/g36A9cpqCK92W6XZsUT9pVvqv4iI3XXl/siKAd7P64Kbe2lHxIRZwTXchYXnFFv6PVahPdtBRwdIEdgAAAAJwGfB2pFf4x17dokp5DlQDbrpLgfa2hCiFMmKnp8wA05fGs6beVH/wAAADVBmwlJ4Q6JlMCN//wDwzHAd5j8QdbEjP/MD54/nxaJsM3etRkV7PtoitdWOIiK4oua7/OJmQAAADRBmypJ4Q8mUwI3//vgOOBDtkfUu7/INxPIL+S5L2s8+0BGgDXxrzOfaBdvcLkOu9kv7y8YAAAAhUGbTUnhDyZTAjf//CIBv4QEgt0jB6sYApnNKebF3UduVNlxp8UDV72R8VAklM98n5lOWEw+T8M4WRe//ioMQECXsBcH5k+vJghjfT5jqlmoeL/l20XeVcwcUIuafvxnJkzhBlcn/d9IuK1j3bCY8QS0oTVT1N2fu8R5Tv1RgAbc5tBoWuAAAAAhQZ9rRRE834H4ctY9gURT1dXKSAIxLkivCpURU347a3c3AAAALwGfjGpFf42YllmapQL/etyKspbtOBGV8f4EVMMX0hi4K6l6e+YACOn1K36gQd74AAAAMkGbjkmoQWiZTAjf++A44EPraJILnXuLxMWdODvz9zlGHYLFuS/WCzov37vtypupywZRAAAASkGbsEnhClJlMFESxv/77lcOA9jsDFR8Xdn5jdZFKfM583oBslJ49SZxnRGp+moI9oSjotvQRWZD8gOf558xXUMpWRXJ/cWn7ZTAAAAALwGfz2pFf42YnyFLdokqnTK0mFdGbVl5+BNzXz+jPJsmLNM19mnectkM0k+LwyFJAAAAPkGb00nhDomUwI3//ALQu8EXYDhpIvttv8TbZDOP8OLBMAO70WexTFz5+cnteMdktY7P65A4tZlIDdklK0GBAAAAH0Gf8UUVPN99TaObewetr3iFfIM29vQ4/CELIo5doCEAAAAtAZ4SakV/hw3+jEmnmFEtcwfltCMFzrX9aTK9FUCy1cKpO11Befk4O1gAKZFBAAAAe0GaFUmoQWiZTBTxv/vBGT5YQIcUNUZpARyxqc+rS5mmS+x2rMYuZzxaaNzjVUhTcMNlid1G47Ff7Gber1siDAOxLdQ+mOx+lSM//HIlVdSI36FhpMg3eSscveUe/EzN0NXsxq+B/YFxJBXKltyUDxpD4j5KmVec/I5HTgAAABoBnjRqRX+O5cRFOX/rc19dAGfoPBoZZ9JRwAAAATFBmjlJ4QpSZTAjf9RUJUwmPx6d1qkD5IiUog3fdDYG+lhD6PhB6357xt0eAhnE79OehpzS9hNg+0rld2gpMfPN78NladMHIMJnu+IgwcdY5qjgIx0rHEFP382Z+0LcY8oZUMOK+1SYe1NyN3/A/j9w4eoQ3PNSo4Ux/LOL3AlYQIiar3GTf0B95Y8t+So3rBxsftentdD21ZJ7hAbvgAAT3Nv7jjPdzJ/O1zUUPXbzz+p7nvdFPTkoO1Dqocj0bjwXKlp/xDJw9QeJ9FCKJkVoMOYhplLX41md7oQ/Y+qGl5+5p8L8BxRb1OTUS0Pldi3oqvL9d/Zrz+oryE/RRst8Em/0e7bxnWHG6Mx81D5GpO0jZK4XqbwX+2UQU4boZy2cX+Sjuvn0MsMkR236Z2K1zQAAAEVBnldFNEzf1+oL/lUB43iexxWj6oMO+zCK0ceAc+a/F/PKiFsvR2BDrZ2PrBeRSYr2NA9GNN54oqHc2LpSV/DC5yGHm8AAAAAnAZ52dEV/jZiDOV4WjRVuqfuRLr/L1bWuM4bOVk+yHw8xIsSbICVWAAAAKQGeeGpFf95p6Ti9RhN9YQqEMUYPvg2cqLZ+fsaBkkhn7KTT2jvNMkWBAAAAgEGae0moQWiZTBTxv/KSaxhdVGtNZuOiRf+JqEcNFkbAv/3toQnjkNAaWg4K8+iRR1zNbt1yIR3tJg8Hh+JLp9Fo/JoyYU3FI5nF35wWUzt3yx1MISi809oOoCzZrj//fWbdhPS2477QVN0UR6uqQ2+uEBXxub0LmMRP6aOCVtOBAAAAJwGemmpFf6zfg2l3Ygdwg0DWPIpQmdHx99XujRSHBNtWXm69NnlO/wAAAEJBmp1J4QpSZTBSxv/74DjghR/zgYSVgQLW28saHEGZfNPuk6gQX9d7Fl5UDqoN11nbQFj7zXG9wSS5vhSv3gAKUUAAAAAwAZ68akV/jZiVZ8+10ObdgxOx3183mwEhK7HkZqcvSQ1ycqtgp8jF+U3zQBWvpsMwAAAAU0Gav0nhDomUwUTG//vB61Ydg9A/wlrkMlWCPDd7QTOqj1/1NiCSa8ruC8UBK8w3Ph2mInOfYuW+cn88B5LL/toEaKTUnUAV8fIMLomlaMvI1rt7AAAAFgGe3mpFf46SPAi9fOZie0tFVOsqcMAAAACWQZrCSeEPJlMCN//HqYcZuGVy+Xdov07HU6KX5bS+1WmnLQvbAvY9UlOzUvt+HiM8nLp39Z11G9J72MQOlYTjkOf+SeXNLe5vidAHlSjT/Ore0st82H8qHv9wvch4z033oRI1qboMoq9jVzWQ3z/b49jZCWWEEv8uwxf9swldbmArV1EwuuF91pmeCP9WmkCZPwED0CvAAAAALUGe4EURPN/QWwldz0TioUzxJwTL1DjVPQjAq9Y54uZHmIdkd8NkQOEM7KiMsQAAAB4BnwFqRX+HBFltjCbzyK+FRKnf/mLOjOcjzR76r8AAAABFQZsGSahBaJlMCN/9YpHGMYUb5j/NfcgCQnnnUME/+uYEFab7R5PJRwSPzdjTXR8fgJDcETNoWNTQyIj4LAz2Alc7KxLVAAAALkGfJEURLN+CHmGW8iX1nC8et2f/B+8KsVycqqolFcdZbi0GJHLqpjMjlxYW8t0AAAAsAZ9DdEV/jZiI7Ev5h6E6AXxVvyU+Xec7Xzf+Ujz4nooqJFL612j/wKfuy7gAAAAiAZ9FakV/jZXGc7dkFSNTQS86kNC5UhVk5V114XqW+n6aNwAAAC1Bm0dJqEFsmUwI3/vuVw4DvMd9xLBIuaZ/zxsUAT+oe+xkeDynj6QmOuSirEAAAABtQZtpSeEKUmUwUVLG//vhUdgbBTvAngor8pfCekIPa/E8Y+CM1csCM1vgERqhAjVxRBDaVo8NX0IcDkRQM7Fc+VrjXHrEEHSAqDVNvfJe+IHgsFDNZNns8smESREry41It9enn5zKZn+/bjH+QQAAABgBn4hqRX+Nfgr1oNIjL185KPI+d0MY3WEAAABQQZuLSeEOiZTBRMb/++A44DvMfbdU1ETVNwuYBfSvlXeOFbA6GbqsdQXHZrecjnlZBPP40eDLe0BA3X8hhdWV4c0xWlDkOf5flpvIBUQRHYAAAAAoAZ+qakV/jnhriwfjEBC0Cyoph0KvuuH/nOiAglQQP/F8Pzo5XgfuSAAAAEBBm6xJ4Q8mUwI3//wp75kBPbfTrjPAkRSSgH3g76N+IPxHc4xKee632cn0Is7nuf/sPmuSHgjxKqRUscLPVxNhAAAAWkGbzUnhDyZTAjf/+8Gbu08WB32S1i5TV1gbt9xeW+dzaP8gfVQwdBZ2FJC6dNZtf8f0Hg+4rzvoWkTdTJuDVTqcnkDXxrzOGZQBT4SetMErEEP9qYQ3b8ROgAAAAF9Bm/BJ4Q8mUwI3//vD7kivHAemh1CS8p9HXZNoOrlrUrukxAJgTBabe3Uq+xKUtXAoE33EFB1Jjy8zf4fiuWPtUn1WH8f/CY8QlZPFmMmcQSG4fjuXN/B+qNi/1vyFEQAAADVBng5FETzfhPBPL0/8xgE055y0cUfkH7gbD0hvkO9CUgUIYF9ILO4onW0KnZIlS4Q5+PTxIAAAACgBni9qRX+NlaNuljs9U7M/SW9U7+gZciTVHqQBZINn/LlwAI3T67KhAAAAQ0GaMUmoQWiZTAjf+8GQRIyA+EC0i53Uv0awLgBrmCYaTtddMfGzDz0tG737dau/+uhw5YMVYL7inRfSxfcFiKBNgykAAABfQZpTSeEKUmUwURLG//vhUdg1hne5jlt8Hnsw/yzG9BlYA+2BCvdNQCaqq/vU4NJus/26Yiqay8cHVINjOHjrvNVNNQMyKgANVegEHkn0Atv5BheYp5J6LroUHacyJH0AAAAaAZ5yakV/isvpAo+X7EhbLzLMVxMoH0gf/oEAAABsQZp3SeEOiZTAjf/7wZu7TxYO5AYsHfwK5RdvkSVEEn21E6ZLxQXmV2KUm/dt5YseWbLYmF6cycseQ/zRHyWaDXSPfsWYc9/L5FXY8bkH+03r0tQyCwxscnMoNSAc8JSQ1LuDhPTTMWEyb+wwAAAAQEGelUUVPN+B+HGW0ktTlKD5U7sftibyzFU0itTwzBOsr8jinvAm3X+xweVleORM4VRCN4J1niSdrhdyNeGimYAAAAAoAZ60dEV/hn60g6kYGvRqVwcf/9DU0cnE/S71LZRHgHSKCWLRPYzu+wAAAB8BnrZqRX+Ky+kGFLaTu3V1h5vNMrHepiZHt++d3TYwAAAAUkGauUmoQWiZTBTxv/wIE8MgSI/5zkdt2xF39AwC2NCzH/ix26Qf9Z2is2kaMgsh4st50ymHALuUDhcWz8HLl5Mv7cFHUPup58bPZYY2bD/A0L4AAAAqAZ7YakV/jZh1ijZ8jlpBJYejDv1xMTNJL8wwhvpmAwdAZHfqRUgGccktAAAAwEGa3UnhClJlMCN/wWJmEnMThj01oz/8HmHqSy1C8Xnv2nkeHB2q5boSEX7n9tutZzT+CbMC2woNyw43AypzG01uaeEWSh0sldqT/FFT7HItLnN5/sxDby4ws3/+X5ukIG6x07NMTkKBgu9pujusd2uIIBTzr4TtqDAIoj5L/81usPjC+yPrZ6BqtdjHQj00WB66maxuzdpr2LURBzdggU7/6dRUoibhHrHu27B3o2Nx4NpnulOOXMgFyxfruOZowQAAACxBnvtFNEzf1+r+1FYhZh13lr8fS3v18bH9CE4mbV7h417LGGti+gpJ8stneQAAADEBnxp0RX+NmIKr+C9nWHyDrxwMHXm6plPWpdYz9mPCwis05VbBSSZGa7pefK68Un2SAAAAYAGfHGpFf9E0GDu+ktvFMGyuyP+/q8a08x7N3S8hze5kbadJPQ4J/VcgfzmVJ+wGbfISiHM3iOG4kNfnUv4V+1AHLRdp1XpMOeQfEcujnC8EllvBq6VYdpBNSXI+IGkPIgAAAE1BmwBJqEFomUwI3/kqHHwSjggzj87YkHr4a87jSRladhRJ0sPf1vzUl8LaQuieyIjiu6qDqBhKWXMid6Xg8mRn7LGXUL7awd9GqdtR0QAAADBBnz5FESzfg/tuGYHKjTB5ZIF5DlRZ/mG18mLNwoqUKDqUjlTWk6UxETTkhiQI1o0AAAAmAZ9fakV/hwRZbYwm8UV+fNHe+WvQ7LE1dhsl0Kkl2dgkJM2xQPQAAACVQZtESahBbJlMCN/HqjpQil7yyLFLFu10qfJCouDtF1Xmx2Qi5MZ8wiTfvWELZh8dDesXTIOc3puWHO2V2s4k4HdAKOR4ngqVZUA0LYM5DJacVR8U9Naj7I7kVfa6bltDE6u1WaZn72qoZRJsNFn5BmTYbmEO/mI8pJzVOUzeRtqmdX+yIakKz1bUKdoGm4/J27cXeasAAAAxQZ9iRRUs36xetRXUJJO0H33aYF1R/wgPQCJEkQv6Oc57u5yk1XK8xotcqZMjal7waAAAADIBn4F0RX+P9L+RFJf/ZGxnwlZxxRpnlAdZ4MTNBABowPPX1SEIZbQRP1HKfV32CEsflwAAACUBn4NqRX/Z96jvz6lPIdxMK5ZxEfdrk4JnwEbzrXk8OtLZKf2BAAAAN0GbhkmoQWyZTBRMb/vuVw4DvMhdaXt49Ohno7dC9VWUVXFssBv/tAEi1GxHEtwjRSRy8OIHuSYAAAAXAZ+lakV/hwRY0JTl1qnyRaLRk+dkPKEAAABqQZunSeEKUmUwI3/9Y3sZBuQxneh/rGW4lE9nuaCObjg+K6/Ef2UAcSMCaL0//e80Z5sMPf/FLcue0CI5qB5M2c2QWFN1gHozimkN8YvezsbDlRHLpDuEB8yJGwaUJGSNmiArYmFEwpDpPAAAAD1Bm8pJ4Q6JlMCN//1ikeX7HlErA1F/NBuaVPRthECIky6bnR/zbhvjqszzH8XsGKtqLqtEjwRPCnmhPsXXAAAAMEGf6EURPN+D+24ZgcqNMHlkgXkOV4kgZTIIkxduFFShDLukcqa0nSuPAX02qo7ZcQAAAB0BnglqRX+HBFltjCbxRX580d76sSrub4Ej3N0tuwAAAEtBmgxJqEFomUwU8b/8CAhiLA9lUJ1ZhKmJlcycBCqNW82vUrKjwfRLn3brEJAyFEiDnxP6vgG4feHLTjXsLOkn44268wtNImWMf5AAAAAtAZ4rakV/jYz8hJwlu0SinQC47vIeaSqHA97ANM9n5sjiQxvenF11Aous9wBFAAAAUkGaLknhClJlMFLG//vDoZtU5Qj04CcZokNlVJpbn+xpxn5vnomGzZJoOL36cb95LTn2GQWez/KCm44eCBSdizIJg/0bpdMf7OAdSKNj8pTKjsAAAAAoAZ5NakV/jHXt2iSnj8+pK5UPI/1KAk39osMlGVfJu+lmZ0egJ6vDwQAAADRBmk9J4Q6JlMCN//wDwzHAd5j8NJF8og0066g+rXTd5exDKM2MCnecCF4E6scY10T/iZu4AAAASEGacEnhDyZTAjf//WN7F6EssGLL9aRq9GMDYogP8Xzv0eGCplNOy5xCXRSUZ3whgiU8wFDYxIZ/nwDS8NdljQdsca56LkdmQQAAAHxBmpNJ4Q8mUwI3//vrhHYHuWOrKSbgXljf8SHjA7pt3GOVNXvZVrQnA9aKMhM1s43sy2s+zopKgZQuR3OXh19uu03zK43qCzT1h20dR/R17wlmMK6XGC4UZ+neex19jtPD/686wWjM+ofcRmWrGJuZFnTOmBpsxXPMq7oNAAAAHUGesUURPN+B+FaPuhgDkzcf9gPd1YURPpDElZWBAAAALQGe0mpFf42Yll2maqZeq3XThpLr/MMrxEP4wtMxgfA47nZIAB0Jg2xgA4T3wQAAADdBmtRJqEFomUwI3/vgOOCJW+nXGeBH0bjrWPxLGdQkrwjNTNgde2mMiCqnPmK8M/7dGPanLBlAAAAASEGa9knhClJlMFESxv/8A6IrIEOZqgng5VRuFA82zmJM8KL0rKYQRJ7Z23FpDej7QMQVdrodIfe0IE+0l4nF1TNAXWq6QYt1HQAAAC0BnxVqRX+NmJ8hS3aWVfP2c1jQ9YmStz/ocuUcfXzyBUtM57gT1Y3KU8wl/ZkAAABKQZsaSeEOiZTAjf/76z44DvMd9wsd+BXDr5TM0gGx6Y5mxrdaiR+xZhz38vjjy0wqFQLRtVm6z4tKSoPcbC90J+/nHLhkajPTm4AAAAA5QZ84RRU834HLEsIWvPinvfHxN11tVvEp3U0V5MHr+u9+xSzIWHqzyYWqp/rDE7G1eS8KR5LYfoqgAAAAJAGfV3RFf4kZCgkjcPXSmeg5eQB5cdFgpKJnWyRjJ0LC8/6RsQAAABgBn1lqRX+HBFjQlOXXoyp+tUCmnUXoUiEAAABUQZtbSahBaJlMCN/74VHYQpILVM82IUwsR3tP1BZNdPf7nkM/rWUZ7Vjytb4wtxQrFU0OUMs2AySifB/5qpsYFfAbi+7wg9W7xVeK6NYvm492GehPAAAASkGbfknhClJlMCN/++A44DvMfbctm/iNKou1bR3qaQPdbYPAwRr192Z0c7+EQI4zvYW462uiNmrbmQcLA9GkcsKw1co4BTWHqy2AAAAAMEGfnEU0TN+D+24e0TioUzxNily9Go1T0H+2EjDCXB2dGC0c6xM99Lg6Jduz8xAlJwAAACEBn71qRX+HBFltjCbzxryJ9mcJ8tdrhFmiT0jMMVk0HmEAAAD/QZuiSahBaJlMCN/BI78lscgq/Rbp6OxwvfwmtoGweCC3o0q532h4ICIwjwcptqf9EK/+79F9P8RvOlrjacSPiqIg7tu0F2k6OuQHzHZomSpzp/xnEFVc3YpD8D7mVeTAU9Xkf0URAjJ5lfHVioB1x0ui2g+17Y3BGxACMJ3vWNmNMsSB6yWO0rku1Z5c+jFLkY2xMy95ck9Xi4MpgR/rgU5kie+i3djqs+2ji7H6EQ+ogtPCUzbkhOsAieuOLwevqHshJvQds8aT+lB97TlMWAfZprX9Mrv3feABeyzfKAEPhJ3BSA1JjxJgX63hdba2srRrtpcoYk81dOBi6AEYAAAALUGfwEURLN/QU/leQXMu9VvVKf5/iBcuPswhbCCq/njMDoMNrlzDXkngFpZCXAAAAHEBn/90RX/WyxZvpLbxTBsrsj/wSTqMIPKPLDvNeh/1F4bmWrWmiWrQyZcVh27RFCrcMqxQPKxSxC/gK2atA4KwaYxGxim8cYsJdnYB2lEgpZEZ9RrWqfitTU8KyDfSoH8YyYLFXyK9azzCP++iehmJGQAAACUBn+FqRX+Ky9y9AWM9M+//Zr43fS7wb55u5yXHFaE5OBuT0DsMAAAAREGb40moQWyZTAjf87uziMWmfzdbE8J87clOeHoYJgH3Cx3Zbf5fu+tjc/rFA2s6Zalj2+4xEwd5fOTh9i8yAEFuj/ORAAAARUGaBUnhClJlMFFSxv/74VHYOCF6yTT6YYA389+T2IXKVeD0g4f5YiUfdwqkE9LXLSW0Wg9bBFbTjG9iud0FpzmEuq0RQQAAABcBniRqRX+OkjwIpFLTj3OvmKtI9JIoYAAAAJJBmilJ4Q6JlMCN/8epy8Hmt5nkvHwUaGRnirECn1F6kQRmAK+NpNIq1+wKlemlXjXEhlgMZ4GsP1pm9vvJhG7Qc9ADDBVLkvTUXiuLhFe47rYgESCog4vHJ3mr48vRHWscHlSyjaWuoADQyZ2fVR3rFtDDYeIJSS40gQKiB4RIui2G6WIFhug4v4SxqnhvpgyxgQAAADBBnkdFFTzf1+NzzhJSSLZXInxj28fVa9sZ3iOeP2ue1tve6YTVpohEK9SBnnnNjSIAAAAmAZ5mdEV/hn7CXSp3/tuQ/gVx//v4sxBQVVQo1NUSFtg3nNLgRrMAAAAgAZ5oakV/0LYnEaxS2mK1fAK4ZH+xIFCrT1VZklQbIOkAAAA2QZprSahBaJlMFPG//APDMcCEr16uk80rHrFFfOOZRse1TbxZVz1dCEySWEu+EN2piNYDrHViAAAAJwGeimpFf4cEPjJJVPrWo1K1z1x0dfojX58CSEdNp6MuCQeaYl5k4AAAACxBmo1J4QpSZTBSxv/74DjgQlevXVYpZHYgY5ILdPK/4Pa1MNf4OiS8dl1R5QAAAC4BnqxqRX+NmJVnz7XQ5t2DE0nG/9r+8NOf3WQ8rd79EUL5OVWw3o6JYHnJty2AAAAAa0Gar0nhDomUwUTG//vBQYpYQEgvj00Y4sfEOWfwyh/OqKfKANvwZZWb7+Zzvu2GXqL7tu9BEtBF1qEI4yW/avkwQegNVC2ROsxmqz3xMUT7hq6sPN5B7n/7jXSLByhN4LaYlMrIVP2Gme+pAAAAFwGezmpFf46SPAjbqplSkyFjtDCI/KdAAAAAUEGa0UnhDyZTBTxv++A44D22Hq7hOPzCa+Q/0uZ8jWnVjeMaXFL05CiAEmr94qQQKh64dWL5oOBuv5DC0B686Pn9ihyHP8aNBy14cSqPoOJZAAAAKAGe8GpFf4cEPpZCBxpCqtM230lpYDo4d21o37/cFMIRIifuqQPwQmMAAAA4QZrySeEPJlMCN//74DjgQ+vwaqsgahBQipjp6SW8umNWKvqDeH4JFmwRnU6qYAsliRyGwuOB45UAAABOQZsTSeEPJlMCN//7w+5hE7FgdvuOKx17H1EIu8P6dhdyNR3Bp44IixcplMPwMmRmTbkZKRLRSf3U7jRY95x56Yr6mgyfOofav4khhB6BAAAAV0GbNknhDyZTAjf/+8UQAxkZAR4MbGRlJS3hcrKcA7AnvEOlxAMEgQ++P4BRPl+kFk9FtV5R/aKjnv2HOToKh1xJET0cChc94nwhXYYLk8OsF9FeDw4hcgAAADhBn1RFETzfghKjLXhVTVpt3odhU1/9thHbD1IrQDd5C/7EbNj0Cg37PlRRDbIVjd2YJ6YmirkTIAAAAC0Bn3VqRX+KzvglsL6ygpfSqzjv723KBBWB5nY6qSoX/1oYogCyBnZuxHtpF3EAAABDQZt3SahBaJlMCN/KCEDzIZX/gifClN7W6rOMfzBCSLAvaxS3ajNlnvJx6yl9zKnNgrw9/FvayjPOvsK8P/j0jTsPuAAAAClBm5hJ4QpSZTAjf/vgHIEFoP7SF8/0hA40oETmSgI6eYrw0ApjXdIPoAAAAGJBm7lJ4Q6JlMCN//vBGT19kCHNXXJVpZIVW2dOsxg+4nBxXYKH8lgiSB4xOvJRCAiVFRsPHQ0NMqcJypOPU3YoZF/tZgsowrc3AiXB3rAlfqaAL7OjHoF9SqtbiXfhpU4xMQAAAGJBm91J4Q8mUwI3//121hr05LoQDvslraTvUwVhvuZbxxeVmAmt8i59vy9EIHQrpGFHmbLnzjIT69y4Yxu/aorB5LH0qX9SvpgCyK/xkpHH537XWLW78sVsfAh6XcctCpaW8QAAADxBn/tFETzfgfhy3ZPoRbD4wmzsftxxUTshAaOeG2I0yWL50v57GYufV/HImcKo1kE6CdyqZME4g3D6gpEAAAAmAZ4adEV/jZiDOg/2dUiFLHxw+J/iID4obTfYetWmDphQQWMG2L8AAAAgAZ4cakV/isvpPCgm6QyWAb6SK+HCMEJXEyPb987umxgAAABPQZofSahBaJlMFPG/yghDLjqnQ+U5Iog63cQgD0Zbv41jiJ1Bb0LhUXR3RIPLhvolGP5fCuJk+gAJDaNEK0tSdXgn/kDyA9nx+m/zdv65gQAAACgBnj5qRX+NmHWKX61xnqOPLP7m1eVao9/Jf8Igt0BHxeMZyj2uGarsAAAAQ0GaIUnhClJlMFLG//vgOOBCV69dVilkdh/3zGtCpgFkyYUdvr7GH5COYLGX7YS8knmvd86QQESRdt876rG5nJMdX+4AAAAwAZ5AakV/jY0SA6wrwXjljlqnIHhE5NeeEAmpe3+gSytfIFQvKbbmEuo179PCD/OBAAABDUGaQ0nhDomUwUTG/85HJkVD4qSGOwcA3K5VwP1c6xtTR4VzmQ0nBu1gOtLRFBdIfSP3an9P1gmsInCPe79j4zuUsL7DK//x4TUhKV8qs9aaEd6tCVykWUn5Gopj46/K++Ykslr+v5UodumlbuShmmcRvZfdNrYnmbVICpyTGYIOI3oAkq6vllklL4brVgHvyURSStEFd3HHzSm7Ib3PtR4E5I89jUSkUoRqP6CDkJUiXvc83P9RW253rviuhvzqjem5Q+hPFgGTABoRk2ZLNhrce5zzCxaYC6gToCoWrl4lzpLNEOsFnUNIz4XxjMda/67iz7TaE3NHffWa7Y/p9uDhqQF1fsC6zz4vKXmUAAAAWQGeYmpFf9Ei2Luv/RNJcaZzW23Slv1beS/YB+iADhqJnO8XGwim1UkZ/5DCTycTZV+kD7aTYfAV6mID48v5zZgCA1mNEtUUWkqJoLo+Ei7fqexybID5ptuhAAAAZUGaZknhDyZTAjf/++A44Hc4/O2JB6+GsntZ/QVS7FVbo5y6gb6CO1Sl0lysiHpH6tBqj4+RV0j1AkhNZawWUW7mw1Ydkn1vMjA/KVOpoHVQwnkQTNeO2ZHz8y7tU10XuEBdXu1VAAAALUGehEURPN99TaxHbgQG60QG5CYHG0EwwVTLS+WEIzdD/zb2JngNyA/ValMZgAAAACIBnqVqRX+HBFltjCbzxqFwcAaE+pOC4nnIXAWqYK0KnVEZAAAAWkGaqkmoQWiZTAjf+8SBZr6iDxTgWdHsgrtUrFBwt9rOo3XBJP8ewDzX8GtwlWBgrC6irWNVrTLY+WKlSir9pSwV4EguL0xV3Wcj0/zw/APuEjV4PAH1Af4Z4AAAAC5BnshFESzfgh5hlu/y+9zSPV3/9spsZbsjm+DaIPLk3e+H8QEI6F/J3CJy6zchAAAALAGe53RFf42NDc3tk/BfLpHtYrknWRGdfCfWGF2hHwQDNNRoM7TAEbExLqzBAAAAIQGe6WpFf4x17dokp4/JyHzCbIpQYkeLYHCEqpiCFaz+wAAAADJBmutJqEFsmUwI3/vD+CADeEGcfiMMXyiDTP3ydcyGMsECu6OwPlhc6nEUtvJobaLBYAAAAKdBmw1J4QpSZTBRUsb/x6w8INv7ObYOML18RRf9FGE1XcGr6KXtfZJPIpBQ3Zddf/xcPS72YLTyF1aF1OgpE73KGCe4EKRaaP6+28u2fFYqgdEy089G2c5o40yj0KCcaF2ckf3qEd7wZbN7tzArbJ3ihmBK1aiIZDAZIELsAMy+cnsNyfdeXk6vZoOzvGLtwNlOVFAWfFH7l4z8o6rut3fiMJEBlecPgQAAABgBnyxqRX/QxMwwr1BpDDFywfg+dwOr6eIAAABJQZsxSeEOiZTAjf/74DjgIHS2PuRVzH01Z0BzuHiEMxkS+es2RQpgPq4cAkl78+syHU05veIX5apNghHwp2dHWeQBuTkHc9KPKQAAADRBn09FFTzfgfhxltJIut7cibOx//YvRw99xcIjnvocl0u27U4ZQtRxhVb+2ZI+NBhFvmHQAAAAJAGfbnRFf4Z+u/AqLSLb8iSXq//onmci1775jnjU+EnU1sAD/wAAAB0Bn3BqRX+Ky+lRxhN9YRwoFckJ8K7wgeXkfQ8c+QAAAE9Bm3JJqEFomUwI3/wDoisgI4JCqEuNC9+tPU5T6akQg8IizB7OrAsSlrnkH+zL9hIrZLMG+0x3pM9DfD7xIYuRzu7XLhI5qDban3u9P/d5AAAAVUGblEnhClJlMFESxv/7w+5IrxwG7Vm0o9WyHVPRTEK7sPH3Tlc6BafENz9DO3jdvcZDYYjxJ/xezZD/wVOgNlWr1ovR60fW2n2QPDIrXRl7wbJGuBkAAAAnAZ+zakV/jZWjbgdMtmBO6kL9v/+iv11IkCbeIE0lsf2Q4YI/jwZ1AAAAREGbtUnhDomUwI3/+8OipWXVEHfdmUcs3voSdMpPrWh0B7x/buw/7bfWTU2SnRyzDOlwhY0AHOKdTjMVTnbawMzXUYswAAAALkGb1knhDyZTAjf/++AcgTTZH1Q4wq20ZHYluqF2d5yizb76mgyembcA7BYepeMAAABaQZv3SeEPJlMCN//74VHYNYZ3xEb5VWTtwzMtJdYmcfzq/85J9tfpgIKEqieYPeY/MEIfk5XhJ1koykG1DVrGfyuRE/hx6dLKPysfIvtARn9KJW3Q6CD7moIQAAAAY0GaGUnhDyZTBRE8b/vgOOB3OPztiPZIA+2600lWyYtu5etLlp5Z1h7gvaM9dZAwSh88F2MaIP+Gb/x+JzyfKCfhrWJAkMuh2vjzRb7d/N7ALJMAooNT/8nzsWXFsQPjIw4BwAAAACsBnjhqRX+OeGijP1BCTddBMbj/ecV0Or4a5xYPj/vYTAUTsVBHuE5kBqvhAAAAOkGaOknhDyZTAjf/++Acgl11jEgudcSnSrhTGl+N28Vhdn1HYhWUAuZUXRLf5VeB/cE3pIjWL90Bv0kAAABeQZpcSeEPJlMFETxv+8PaaJ1sN1uxRNIGqNtU2k3406usYCk1bEtXHnFK6NYx2++uiXE5dDpIBIRXDOXY8QkTWaxAEokog4eXK891HbBLLwiFuVsjJhcQDmLu/1BPQQAAAC8BnntqRX+NjQ3Qw5SeeQtnahQCGRD+tBtlepn/thplttEshgUAcjhCMocOSDNHoAAAAE1Bmn9J4Q8mUwI3//vDpMA3iJWBxGCoVqH005LOHvO3JnF0uPAF/1BdrAmcA6ADeTiMuCgJLUTG5+vRnVljzqS04Al1sJB4spczTDLw8AAAAB9Bnp1FETzffU2jm3iDcWTT8O+Q1NZfwE4BDhgPUoUdAAAAMAGevmpFf42YlWfPtdDm3YMTsd9f+YkkWkcv8aPIfR7QyuqOLTk5HzNqkb26HGVg0AAAAGdBmqFJqEFomUwU8b/74VHYnLGfY5FAo39Th64FzFtgZ97sVu925ubsyqHBRHBm16zii8aJWj58APV75EudDjr7xLleCLLq7/RtBuOXcv77e28R66Sg0C5+ohD0irjcvKrApu8KWKOAAAAAGAGewGpFf46SPAinGc0bJk9Hz2z1Me3AgQAAAQNBmsVJ4QpSZTAjf8fMy8FXp/MBuvzFgdo0Sga+yREhRElJ1uA9sIsY4ra6c6fc00xs5DNAd8Y8yfothafu7uajsv9X2+oe8VfvZheOdZXtqkEIPdjDT49hK59jghSaTBGRE3a2q4B8vzpyVLIVGWiucmJcTr7iO+Rs1GKykq2Hqpg/7wGE9uYxrJNr28ZDTVBHR/4CyCAPQk6jWNczPh78rCn0z6G9a1RjzePKZYANdrkvuc+8zCzahQHbPyJ3wSnp4GBi0FZoLY5fHqWlr8b307+UI2CHj0yzFgHzBe7+BMVC9y83a86LgcLsFvxDSFgPEYch+RqKUMFJwfV/KNr0okJZAAAAQUGe40U0TN/X6axmThJSSLrDjCbOx+2e33lIIgsTG2JtszOBl/PYEZ41qb1VwqSpMkCnW2Br5kIl2A3Uu0IPihiHAAAAJQGfAnRFf42Ygzli/toiEQB9cbIH+nvDxV+o1J32C6oS5Fi80REAAAAoAZ8EakV/3mnpOLtzNha7JYFMdaGbfLAP/sxzAhGKBTfBtRpxTI6dgAAAAIFBmwdJqEFomUwU8b/0EQS/wgjd+gtxcvUsjBZ/KvLmmxO81sipaPVWNzgoVqy31nJce7cJ78sV7DeSUW+b67IBOE8jk8HFII54ZV3rWpp/dqQpuWbfoc2fWtU2v+/yK+ia43/m1sEOBWkif02dhVPEQ9NMhk2dwH5CcI/zvbk/eCkAAAAoAZ8makV/rN+DaXdiJ3Sfbla5z0prYIJJdn9L5fmjpAFnOrwj7T71gAAAAEdBmylJ4QpSZTBSxv/74DjhLmEQjvo7MOz/6KwI8CxsOXVk34E9cwzHoe+ypu4pyoHVVLROmPuXefirSAYX2k1/KeHKYNsg4QAAAC0Bn0hqRX+NmJVnz7XQ5t2DE0nG/A/2tAvwsk191IaIPxh05qExET5GLoDSl9kAAACTQZtLSeEOiZTBRMb//WN7GRlg7p5tdxIOcz4ktWOw25b+IeCEiFor+loK1sW6QJ67CqjJiEE34iTVIck02sn81tu+PGvjsyHqi2o+NT4inWWVtqWkpr7f/3c+vKVzlj1mjdhK8OP4T7BUoxX0iIL7Ms5LBKBxFY04P++ZvxW8cml6mweqdCiwy+yuvk365ebqBsgsAAAAGAGfampFf5XYaTriDUHbSM4/2rI012dTgAAAADxBm25J4Q8mUwI3//vgPUQ9+HnYzR0ReHsw8528FstVpSX9I30YzTO1HQX7vIeLO4umR1DvPGSkhtbazhcAAAAqQZ+MRRE8331S+se7kbW6AxPsrNhK/dtL9AZtyTBgH5jIFFvIWlbz+Y1vAAAAHQGfrWpFf4cEWNYpbR3VxAdD0L/IP/FPtuyuyqOpAAAAgUGbskmoQWiZTAjfx9CcA6khDjKpbsCmICblPVjONIzX1mXC3smr+Dq5fXiG9Xv8kFQyUglCZwie8jt9HTy9u/pPWV6Eg7yLHBH4vZV7UNGZW/41bil6995IKcn441hEAdqxiSHI4vuUl9rENOzJKWLV2U30qhtdcyYWhWRdSq79ngAAADFBn9BFESzf0IJ5kmO08++CFl4DaPf738KLIygT84teS4UEhuc0uV5bKc3qgHgZU8MlAAAALwGf73RFf7OB2xJd3tC/Pqn8/HmIdT3BpECF6yr/lbOPB+Zy1K11SUfd4VAGF2JZAAAAIQGf8WpFf4x17dokofjXiSEcvI/1LMfeM2Env+wlhWYh/wAAAEBBm/RJqEFsmUwUTG/9YpGqI6+gv8B9xK7sh6JffgqkidB+uyLMXamYT/v0gM2Bl0cBW+zOs7+meEYWSU7mTdahAAAAFwGeE2pFf4cEWNCU6Hlhkpwoe3NC1OR+AAAAfEGaFUnhClJlMCN/x85ZyjCkCh8uzK70/iv22/5Y3C+WTfpvUCrA+KzvU4pHDm4sP40BAbPg567zdr1Wu7EEBUuR9mVRdZyCPSUp6JfGVxbeoCt4oVKZm8KmBNWkYt6ZDzaPhh8MV+sn/5rn/+5/sTVujqfTRrIRFR+q1bAAAABYQZo3SeEOiZTBTRMb//vgOOB3OPztiPYLGqbk5d2BqIQU3KblqkAke1WFGNwX7uSXHgcpobdidZRPvUHXnKBlNvG1whZJdIUyHR2g5/kj6iGYMgjOKwfAwQAAACgBnlZqRX+NmJZdpmrtcZ0NTAN84Rf1N4fqx73Qed7w+gCCXwdxAjG4AAAANEGaWEnhDyZTAjf/++A44Ilb6dVP1x6RsYOzc76QVwcX2znv5faz7QEaAFkrcI7CaujLN3AAAABVQZp5SeEPJlMCN//7w+5d2nwAIi2iZMsWHBHowYX4aPQ815YaFI7qF/qVwMK0/Txt1GjsuoBrg0JYfv/uwERjL/ROpxdPOR5x3xoTlVPE5qTT470ToQAAAFlBmpxJ4Q8mUwI3//vG3LzDXKIqZ9TYmttr5CBr315a/0BUuW50ROufAOdVNE+12YSjQ+/9Zk6dVP/79hzk6CodcSRE5KWMo3xZZ6kBcB3u5L6hIx3KSQR9gQAAADlBnrpFETzfgcr6UK161qa1NGdzHnS/SFXblAfYkkTUgGALg0oSw2I9ZdDSKhlEZD82KSrfFpA+v6EAAAAnAZ7bakV/jZWjblZD8+kx6s9vusQNFqTkiIzV2G+Y2RxdmUUlRu/AAAAAVUGa3UmoQWiZTAjf+8OnUzwggj/qGfOg+4iB+MF4Y5ta94QwdwgcgpS5V9IyYNFhVGlpN2bjMQ4DNhuvd/qEmbi/FxB9D9m8UL/JpA3iUwbQMpYlEjoAAABmQZr/SeEKUmUwURLG//vB2iYdh5ImpKHSwtHernzZCXrV5CEdMKE6oB+B0fwylF2lU+/yPss1M1i+5F7zxa3LLlnpb90pQhlGEO/tirOcGbuioWk+zu7MXqDCxg3B7BB3iJfYGJshAAAAFwGfHmpFf46SPAi9dvB8N1WaHbwo5+8wAAAAaUGbA0nhDomUwI3/+8Gbu08WIicdkBLuDqjcxTRGWFEHK1YMItB/Yxti7mXMwKkgf/cswQBjw/L2b1agYYBo9t6xn6Eetw4bxPxz38vjh5Ug36pxN1P+Jdu6y5/v+SSnY/wBOAf2anagvwAAAD1BnyFFFTzfgcrcmZwgJIut+RPjHv9yMW/NWZe0k8Bi3Zgre6X89gTtwG/OmXBXHUG6b/SiwXTP+8AECry3AAAAKAGfQHRFf42Ygzli/0aMsY2QSz83/bIH9oh1NSd2YfdN+fUQOgXqQV4AAAAgAZ9CakV/isvpBhS2mIvgYfwGKq3zW41GlvjW/evQpscAAABJQZtFSahBaJlMFPG//APDMcJc0yddeUYH6nOgY2NbwtKavnsY+nBe+fxVlILL0NiRXrOhPUiA969gXTMpBWKdBefNzN3yeWc1gQAAADABn2RqRX+NmH8N/4LEDrqsxxCUlA+/C6d/hrlaE3Eni4AjiTNfEyqi1KVhrZEtHyoAAAA8QZtnSeEKUmUwUsb/++A44S5pk6Zc8mpbajeY3FQlOl5eERU4NlZQ7KCHXnQqL1aUm+ZXyoGSi1CXUAu7AAAALQGfhmpFf48rV2va6HNuwYm2A5WJIWicynQ6BVYFp9FWnKrYE2sDGYjSbKzVRAAAASVBm4lJ4Q6JlMFExv/HSJgSC1nZxJ/p92h7iaNLjO1oWGnVb7mdLzqjQegiZZ8BFVoHleRS9opigDSoevdqf/HRbxPUT937O3gQ14DzaDIrquleqdGqhoMap2XYPTsxa/wSkDdVT0ssgNi0XqyHVXOCMqt31kydwNB3y5PU4UE0NqfePmpADjVY3oiJGyZaYPsi/X4ym21aC/XEA8g6cHPxA5l4ZzRa4mnIZk+87eHMPyVrzjVuzsYf3L5Mdth1vbIMZuHrbzmxDF4Id95FDUpJLeBKHbiVKXlKbVlORXdE3XDZHj4OQ44bzKNetjD2jssvKy2arGt8LqXFFRWGchvabmDpRlO+TJnoh6Qi34c/JAelcxWZTHtnFBN5/wh1aZMwTQYfgQAAAF4Bn6hqRX/RIti7sAopo312LMs22rFhBmyBC8ZUUUusB76jdtD7pxWchSlEaZcpo+z3LSOsqF+TjWIbKyoOjcJ0BvP5Un6zNu7CYyOjMKakeOmtBeHghnZ0Lt7ZKgmBAAAAZUGbrUnhDyZTAi//87FwbhKNgmMp3eNkrIgo+BdeIZ2RH7aVgRWSRB8VGoJF1H932TYjOcWyo9X8eYosRzKDsLYZ8OElkDGMElhddcoxE4Cuo16C400E76DzMxmRGQrfcMoXwPeAAAAANUGfy0URPN+DBf3EVFfJcIvLYYm40FJuHP41gUVpmbF9867RDcpbG+hUwmLp76j9ncDp8t5SAAAAIwGf6nRFf4Z+tIoBTm+ZzUi4P/L4NbInHV4wNqEejzvRIyMbAAAAHAGf7GpFf4rL6VHGE3tN/0VO4T4Y92cJmTGqjpgAAABFQZvvSahBaJlMFPF/+G0p3IyKrpYoACoREvabxZ2Sa1j/PX4Ig63AFh5hVcDy6lyqDJVazyv9k0A3UEbpQ4kElhDfQ/kDAAAAJQGeDmpFf42Yfwkkpl9ce7oX5djjxlYebpvmnM3if/h60uGBIJAAAAA0QZoRSeEKUmUwUsX/+99LcS41uOHpiqjPsmwGusCKQ6BF1AYtz/WtQ7b6AecL2sZs2DeFgQAAACoBnjBqRX+PK1dr2uhzbsGJtgOViSFonOVO2XvUQkyMOnNQmIoMiU2QLcEAAABnQZozSeEOiZTBRMX/++B7/JR3TzHmkys6OPVr3j60/Ducpwbn9LZ8EMRIp1AAs7+bDOf0XhisAXR+MHkgbVlZgREbbfkD3LjS6EKrXc0GPULfC2RHIyE230IZnBmWLkBNNdGdzIhkLQAAABcBnlJqRX+OkjwI26tNPqSi/HkmF/g/gQAAAHJBmlZJ4Q8mUwIn//Dz42spRFihMPfTkj2TNj0lStYXk3a+C5Wtkfypzo0uM/feYzEOqkQzy1NzoYXThBgSRvjrEcejioh4gzCm8rCbzMI2wWFHN/yLtF3p6q9/Ec94jMIwPBQYnzCY283aqNgEwQc+asIAAAAwQZ50RRE836foBRLD06goiSPhmYv9qz7O6uD4czg9Q3XDksNdfK5w0IuoG1k3GS0sAAAAHAGelWpFf4cEWNYpbQS22jK73/5IyrwB5npkP/EAAABQQZqYSahBaJlMFPE/+81+QTCAevbatPJvZfPu+5FInLiQmsx/vsDDnlk+n51NMYohIZisWJ/S3zrGQtmid1/VGv24Rk3j0Dp9U4U6eejQnKAAAAAuAZ63akV/jZifIUt2huke1h+nmbVlvSl8BsecDgLV4KKiXsuB44t0CtoWXXNbfgAAACtBmrlJ4QpSZTAivyCC9u0SU8fn1JmlqS4H0O/V2QQ8C23N8Hq2eJg+n3lBAAAEpWWIggAb/6zQYZLfQPgNc0iYMN3sPl0EDBzzvrz1O9Sgfa49FGnVhJ8qLkhVezFoQxJKbc0nUGNZR098qd7OS+y89EcOWVfDbC64tS8Y9KVBO5bX3i8BNZhBqqdfBpIeL6SUDjC2QpVPdT+4GdjwHm+HVEddJzn4mxN+ZHi1zAImnjZYMgtVcr/8KsM8y6i5AmAGJNcUB7ayE0bca3a3R6hdrr0XJr+AAAXHQ/kXcVlmP6pUrkjcKn/5MsrYsfvEq74eZ4HZz2CqbE+TyZ4bBsWkF/G8EsnWTdMeCaT0oSOggGj6F6GH2koqwzCw8EE/jxjH/fgpvYdskan80JjRYRAJC3SwAtTkEnddIDJhOwFLkmhXRf1WZtt88VL7WgZg+Do7We/RCfv4cfPN57fJ/eEcGoPIdjySktCXmVnqHKZZhx61mqXrYt2ccmalI5pFwxn7yIk4NanyDkRuoRQdzVy7V/FgNDX+/OG313omVQVM1D6TfnNy1isPsCV96//sVHChbSvRvo9Op4uXUT2Vnp2NbBUnPOlIQwYUmnls4/qAf1DSyr347ve+HuWjGSfL2uAJLxNkKGL7jaqU87H350pTIU4+n/FRrP8UOsWWuZ9oqEO+lS3qvWMNvTA00dfS9Og3CEQ1lhvXgqkfbgmWS7Q7BrqjPYxv+usZZhR4mTsySOLjIGUswJByOkD9sXOn034oNMHPSMTs90mvpwuhhZCES9QIA0MY4mVJAH4U1175JsR0JGrNT2cG+kevDwhB7L2Ps9sq35Ja1G1YIr1fCqBH1NhFZbh3r/JJ6IfmaKkwRe/x35k1btKb+ZRfDYCZwD/U+DVYV4ve4/ux5EMulf81TrZ2iS/AeNl9oIVL7FLsYt0SbM4Gdwujxy/cVa/9DfeeCmTLktfb85Jq/PPzoWfkT+3MSEL52ahvl5dyZ3P9Gjqc/3V3E3UcMDLHXaFjLji1XATHryNndFbY++diGv2rJsu8jbVZYN/lqAhbLHtMl/sxkHV0L1mnPx14QjgEVCWYb8PE3T+6bY0rrXT6WTEzJq/uaxvgC3UliUMREWCB/LHkW4/FwMNKpTt+KgL7DKMzlTFtOl01XRpLhfNYIv4yqp/8/0Ljp/cAHvBSTd0o2zGUJLgBlXpkC7hkWNvVKuNH3gZDe7aBWygB4hHI5MtXigV2BDYnOAj1gC8+xGfGlAw3JTy1PkTj3bbr+p6m1jzipdInXtuw8i53LkL/+t/cDhYKza5lgkrWZvnpedOPwE+SDsrxeadOx3w6iN4kI3PIREK2JArXFqTIikdsxGeR30xC/R4MnKZUvO9n1HHNmqCdKQY/t8N97D5MniGESaR7uA4AEebTLOvxLuJghYdjoSxajwIKjHQQZ/HnrodJI5KNoj6gM3UJuupZ/7uHrjf0Dx0Tsc7GUns7O78xmwOt1zDx4hBNRbDjREElJr9eJlRZ3FMTZgL28DJXX2Ukq+M+ZYpT8CiXX8w7XSQxcMg4cJ6c0AGluEsCz3tAkAQZLU/VxMQpKTJR8+ULaCRGp3sYhJhT5mx6NspTrLO4jl7T5GLVFULkYWq1jSnmblAgEE58LcEAAAA9QZohbE//55dAJChmGePT0n0Li66IQjlgCcvSeXFOTWduetNIn5ncKj2uvzaGAWdM4YtqdWioqkA0TyHz0QAAADBBmkI8IZMphP/nR8JCz0aek9id9vrF4lLN2RgQjyZr432P7Kp14LSDqmMhMkWW0IAAAABgQZpjSeEPJlMCf+dJU6LZgkV02QLvEfeXB5dv4nCjjykgCXjsebUglI9rrNk3mZJmTaQFGdOlZR2+jF2hfcbd4HXQBsGGhxrxbbB//345uw0C7pMaLkJqkj5rLI3WN+aIAAAAUEGahUnhDyZTBRE8/+ZYKy4kVT7wb4JNhryfYT+xBrDfh8zeOyqxLHImSnk1hEFx+vGFOK6nSbea/LiCsJT6jX4eofbWT/4UXXlaQ7ns7HmfAAAAKgGepGpFf42Yll7P81LVQddgA183/eHh4rICR6Rm/TAP+50pi9D2H4nE8AAAADBBmqZJ4Q8mUwJfi8K/Q8CzpoXb9z2cRUhHfcB/4ijSNB3DB8kDQxNhedtHQQegeFAAAABHQZrISeEPJlMFETy/jBNpyKcpBz2UWq9NhKZ8OdqT4vlDJSja2L4KsxrGpSrA2yYsC6BONsZSjEWJ7nESRHswZooLCvZWL7cAAAAtAZ7nakV/jY0N0M9lu0TcpcLEt4G2uP1I1ad/p8M0BAfmcyRE+XWn8kbnbCP/AAAAQ0Ga7EnhDyZTA/8JnI6fa/XOcPoz6Eunvm3DSqWZut6xANl//x1bLCAKYLkthnFxJePmxk5tfDOwQ/OimiJzESuj2PEAAAAyQZ8KRRE834HLEsIWvPiTtw1hFA7vKrrccI3560FD1CXoViI/ZgtgMQkFL76i8/+MS6kAAAAdAZ8pdEV/iRfHjugOkdFKyjh/gRLx8g9LelE/H6AAAAAYAZ8rakV/hwRYbMohViYe1/b50qjneEmNAAAAWEGbLUmoQWiZTAr/DOcGwr30MLcyf4RisZLdUWnlUYvXJVKV7AncbyAoLODntBs0+K4WaaoVpZfqfI2XI0wXb+fJSpC5XehdM+mTV3yL+1xrM/enVV5ivWAAAAAzQZtQSeEKUmUwIr8dln8w8D8m9uEVJ8VaXrlpwtQMUcjOxV139A/A0FOJb3veKqfxYs2BAAAALkGfbkU0TN+BznEOeLmqf8Sjo1rSbw27tYVeLjYdb2CxT8VnvpcHn1axdLPZBJ0AAAAtAZ+PakV/hwRZbYwm6h1dIRfbz8pmPJ5q1aY4QcQjnAtS8tukkKVgqkBfhK/oAAAkC21vb3YAAABsbXZoZAAAAAAAAAAAAAAAAAAAA+gAASucAAEAAAEAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAIAACM1dHJhawAAAFx0a2hkAAAAAwAAAAAAAAAAAAAAAQAAAAAAASucAAAAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAQAAAAABQAAAAhAAAAAAAJGVkdHMAAAAcZWxzdAAAAAAAAAABAAErnAAACAAAAQAAAAAirW1kaWEAAAAgbWRoZAAAAAAAAAAAAAAAAAAAKAAAC/wAVcQAAAAAAC1oZGxyAAAAAAAAAAB2aWRlAAAAAAAAAAAAAAAAVmlkZW9IYW5kbGVyAAAAIlhtaW5mAAAAFHZtaGQAAAABAAAAAAAAAAAAAAAkZGluZgAAABxkcmVmAAAAAAAAAAEAAAAMdXJsIAAAAAEAACIYc3RibAAAAKhzdHNkAAAAAAAAAAEAAACYYXZjMQAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAABQAIQASAAAAEgAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABj//wAAADJhdmNDAWQACv/hABlnZAAKrNlFE/nwEQAAAwABAAADABQPEiWWAQAGaOvjyyLAAAAAEHBhc3AAAAABAAAAAQAAABhzdHRzAAAAAAAAAAEAAAL/AAAEAAAAACBzdHNzAAAAAAAAAAQAAAABAAAA+wAAAfUAAALvAAAU8GN0dHMAAAAAAAACnAAAAAEAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAwAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAABAAACAAAAAABAAAMAAAAAAEAAAQAAAAABwAACAAAAAABAAAQAAAAAAIAAAQAAAAABwAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAADAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAwAACAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAACAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAwAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAACAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAMAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAADAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAACAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAMAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAACAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAMAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAMAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAABQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAABAAAAAAAgAABAAAAAAcc3RzYwAAAAAAAAABAAAAAQAAAv8AAAABAAAMEHN0c3oAAAAAAAAAAAAAAv8AAAZHAAAANQAAACEAAAAkAAABswAAADgAAAAmAAAAbwAAADoAAAAzAAAAKAAAADcAAAC4AAAAIwAAAFoAAAAxAAAAJQAAABwAAABfAAAAWAAAACoAAABXAAAANAAAAF0AAABhAAAALAAAAD4AAABHAAAALwAAAEgAAAAeAAAANgAAAFgAAAAiAAAATwAAADsAAABXAAAARgAAAFUAAABFAAAAQQAAADUAAAA4AAAAXwAAAFQAAABEAAAARwAAAGEAAABVAAAATgAAACIAAAA8AAAAQwAAAGMAAABRAAAASwAAAEcAAABuAAAAUAAAAFsAAAAmAAAASAAAAHAAAAAbAAAAVwAAADQAAAAkAAABHwAAADoAAAB7AAAAJAAAAEAAAACzAAAAJAAAAHYAAAA7AAAAKwAAACEAAABXAAAALQAAAEQAAAA0AAAAcgAAABwAAABnAAAAMQAAADkAAABcAAAAbgAAAEAAAAAoAAAATAAAAEgAAAAdAAAAZgAAAEQAAAAvAAAAJgAAAE0AAAA4AAAANQAAADYAAAFBAAAAZAAAAB0AAACUAAAAQQAAACEAAAAvAAAAQQAAADIAAAAfAAAAGgAAAFgAAABLAAAANwAAACkAAAAfAAAAVAAAAFUAAAAnAAAATAAAADcAAABdAAAATgAAAC4AAAA1AAAATQAAADEAAABDAAAAIgAAADAAAABhAAAAIQAAAQkAAABPAAAALQAAAC0AAACOAAAALAAAAEUAAAA2AAAAaAAAABoAAACaAAAAMgAAACEAAABGAAAANwAAAC8AAAAmAAAALgAAADIAAACCAAAAcwAAAD4AAABaAAAAXwAAAGYAAAA9AAAAKQAAAEUAAABYAAAAHAAAAFMAAAAxAAAAJAAAAG8AAABPAAAANwAAAC4AAAEZAAAAIQAAALsAAABvAAAANAAAADAAAACtAAAAQQAAADMAAAAuAAAAXAAAADEAAABfAAAATgAAADMAAAAiAAAAUAAAADMAAABXAAAALAAAAD0AAABTAAAAawAAACEAAAAvAAAANQAAAEwAAAAzAAAATAAAACIAAAAyAAAAkwAAAB8AAABpAAAAOAAAACgAAAD/AAAAOAAAAHcAAAAqAAAASQAAAHQAAAAbAAAAnAAAADoAAAArAAAAIgAAAGUAAAAxAAAAQwAAAC4AAABmAAAAHQAAAGgAAAAuAAAANAAAAFYAAABwAAAAPQAAADAAAABHAAAAWQAAABwAAABkAAAARgAAAC8AAAAkAAAATQAAADUAAAA9AAAAMwAAAOkAAABcAAAAXQAAADwAAAApAAAAIQAAAEAAAAApAAAFBgAAAEAAAAC2AAAAHAAAAEwAAAA3AAAAKwAAACEAAABTAAAAWgAAACgAAABNAAAANwAAAFkAAABTAAAALgAAADcAAABKAAAAMwAAAEIAAAAhAAAANAAAAIYAAAAdAAABMQAAAEsAAAAuAAAAMQAAAIMAAAAxAAAARwAAADIAAABkAAAAGgAAADgAAAAvAAAAIQAAAIgAAAA3AAAALwAAACYAAAA7AAAAHgAAAE4AAABYAAAAKgAAADYAAABSAAAAXwAAAD0AAAArAAAATQAAAGMAAAAfAAAAagAAAEAAAAAlAAAAcAAAAFAAAAA1AAAAKwAAAD4AAAFQAAAAXAAAAGMAAAA2AAAALAAAACEAAABCAAAAKwAAADcAAAAwAAAAVQAAABkAAABsAAAAMgAAACEAAABbAAAANQAAAE8AAAAtAAAAPAAAADMAAABIAAAASwAAAC8AAAA/AAAATQAAAC8AAABQAAAANAAAACcAAAAbAAAAdQAAAEEAAAAoAAAANAAAAIkAAAA6AAAAeQAAACoAAAA9AAAAGgAAAF8AAABXAAAANQAAACYAAAAgAAAASwAAACoAAAA2AAAAMQAAAF4AAAAcAAAAWgAAACoAAAA2AAAAUAAAAF0AAAA/AAAAMAAAAEkAAABpAAAAHgAAAHUAAABDAAAALgAAACMAAABKAAAAMQAAADwAAAAyAAABcwAAAFkAAACnAAAANQAAACkAAABuAAAAMgAAACoAAAAoAAAAVgAAAG4AAAAcAAAAWAAAADwAAAA7AAAAKAAAAFgAAABWAAAAKAAAAE8AAAA2AAAAiQAAAEwAAAAtAAAANQAAAGcAAAAzAAAAUAAAADkAAAAgAAAAHQAAAGUAAAEeAAAAUAAAAC4AAAAxAAAAogAAACkAAACOAAAAOwAAAFoAAAAbAAAAVQAAADkAAAAsAAAAIwAAAEsAAAAqAAAANAAAADAAAABTAAAAHwAAAFcAAAArAAAAMwAAAFMAAABdAAAAPQAAADAAAABcAAAAgQAAACIAAACCAAAAUQAAAC8AAAAjAAAATQAAAC4AAAA5AAAANAAAAQUAAAC3AAAAogAAADUAAAArAAAAIQAAAEAAAAAnAAAANgAAAC8AAABTAAAAGQAAADwAAAAxAAAAIQAAAE8AAAAzAAAAWwAAACsAAAA5AAAAPQAAAF4AAABRAAAAMAAAADcAAABQAAAALgAAAEoAAAAiAAAAMwAAAI4AAAAdAAAATwAAADgAAAA3AAAA7QAAADQAAABwAAAAKgAAADgAAACTAAAAHQAAAFQAAAA1AAAAKQAAACIAAABSAAAAKAAABUYAAAA+AAAAlAAAAB0AAAB3AAAALgAAADMAAABcAAAAZQAAAD8AAAAsAAAATAAAAG4AAAAbAAAAawAAAEAAAAAtAAAAJwAAAE0AAAA7AAAAOQAAADEAAAEWAAAAVwAAAFQAAAA3AAAAKgAAACAAAACPAAAAKwAAAE0AAAAmAAAALwAAABgAAAA+AAAAMQAAACAAAABSAAAAMQAAAFwAAAArAAAAOQAAADgAAACJAAAAJQAAADMAAAA2AAAATgAAADMAAABCAAAAIwAAADEAAAB/AAAAHgAAATUAAABJAAAAKwAAAC0AAACEAAAAKwAAAEYAAAA0AAAAVwAAABoAAACaAAAAMQAAACIAAABJAAAAMgAAADAAAAAmAAAAMQAAAHEAAAAcAAAAVAAAACwAAABEAAAAXgAAAGMAAAA5AAAALAAAAEcAAABjAAAAHgAAAHAAAABEAAAALAAAACMAAABWAAAALgAAAMQAAAAwAAAANQAAAGQAAABRAAAANAAAACoAAACZAAAANQAAADYAAAApAAAAOwAAABsAAABuAAAAQQAAADQAAAAhAAAATwAAADEAAABWAAAALAAAADgAAABMAAAAgAAAACEAAAAxAAAAOwAAAEwAAAAxAAAATgAAAD0AAAAoAAAAHAAAAFgAAABOAAAANAAAACUAAAEDAAAAMQAAAHUAAAApAAAASAAAAEkAAAAbAAAAlgAAADQAAAAqAAAAJAAAADoAAAArAAAAMAAAADIAAABvAAAAGwAAAFQAAAAsAAAAPAAAAFIAAABbAAAAPAAAADEAAABHAAAALQAAAGYAAABmAAAAQAAAACoAAAAkAAAAUwAAACwAAABHAAAANAAAAREAAABdAAAAaQAAADEAAAAmAAAAXgAAADIAAAAwAAAAJQAAADYAAACrAAAAHAAAAE0AAAA4AAAAKAAAACEAAABTAAAAWQAAACsAAABIAAAAMgAAAF4AAABnAAAALwAAAD4AAABiAAAAMwAAAFEAAAAjAAAANAAAAGsAAAAcAAABBwAAAEUAAAApAAAALAAAAIUAAAAsAAAASwAAADEAAACXAAAAHAAAAEAAAAAuAAAAIQAAAIUAAAA1AAAAMwAAACUAAABEAAAAGwAAAIAAAABcAAAALAAAADgAAABZAAAAXQAAAD0AAAArAAAAWQAAAGoAAAAbAAAAbQAAAEEAAAAsAAAAJAAAAE0AAAA0AAAAQAAAADEAAAEpAAAAYgAAAGkAAAA5AAAAJwAAACAAAABJAAAAKQAAADgAAAAuAAAAawAAABsAAAB2AAAANAAAACAAAABUAAAAMgAAAC8AAASpAAAAQQAAADQAAABkAAAAVAAAAC4AAAA0AAAASwAAADEAAABHAAAANgAAACEAAAAcAAAAXAAAADcAAAAyAAAAMQAAABRzdGNvAAAAAAAAAAEAAAAwAAAAYnVkdGEAAABabWV0YQAAAAAAAAAhaGRscgAAAAAAAAAAbWRpcmFwcGwAAAAAAAAAAAAAAAAtaWxzdAAAACWpdG9vAAAAHWRhdGEAAAABAAAAAExhdmY1Ny44My4xMDA=","ok":true,"headers":[["content-type","video/mp4"]],"status":200,"status_text":""}},"base_uri":"https://localhost:8080/","height":501}},"cell_type":"code","source":["play_video('mf_pong/0.avi')"],"execution_count":0,"outputs":[{"output_type":"execute_result","data":{"text/html":["\n","    <video width=\"640\" height=\"480\" controls>\n","      <source src=\"/nbextensions/vid.mp4\" type=\"video/mp4\">\n","    </video>\n","  "],"text/plain":["<IPython.core.display.HTML object>"]},"metadata":{"tags":[]},"execution_count":31}]},{"metadata":{"id":"NQmZEVKGF4Hh","colab_type":"text"},"cell_type":"markdown","source":["# Model-based training\n","\n","The `rl` package offers many more features, including model-based training. For instructions on how to use them, go to our [README](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor)."]}]}
\ No newline at end of file
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "xCLcAmON-m2i",
+    "colab_type": "text"
+   },
+   "source": [
+    "# Tensor2Tensor Reinforcement Learning\n",
+    "\n",
+    "The `rl` package provides the ability to run model-free and model-based reinforcement learning algorithms.\n",
+    "\n",
+    "Currently, we support the Proximal Policy Optimization ([PPO](https://arxiv.org/abs/1707.06347)) and Simulated Policy Learning ([SimPLe](https://arxiv.org/abs/1903.00374)).\n",
+    "\n",
+    "Below you will find examples of PPO training using `trainer_model_free.py` and SimPLe traning using `trainer_model_based.py`.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "id": "RW7gEGp3e87G",
+    "colab_type": "code",
+    "colab": {},
+    "cellView": "form"
+   },
+   "outputs": [],
+   "source": [
+    "#@title\n",
+    "# Copyright 2018 Google LLC.\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",
+    "# https://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."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "id": "pq0BqXm4-3gJ",
+    "colab_type": "code",
+    "outputId": "6086719f-6268-4b61-8fa3-d251eda24c97",
+    "executionInfo": {
+     "status": "ok",
+     "timestamp": 1.553273826475E12,
+     "user_tz": -60.0,
+     "elapsed": 20650.0,
+     "user": {
+      "displayName": "Piotr Miłoś",
+      "photoUrl": "https://lh3.googleusercontent.com/-050ZBEGpNAA/AAAAAAAAAAI/AAAAAAAAk9g/r6cv_J6J5qA/s64/photo.jpg",
+      "userId": "12158759908531801397"
+     }
+    },
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 163.0
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[K    100% |████████████████████████████████| 1.3MB 9.4MB/s \n",
+      "\u001b[K    100% |████████████████████████████████| 215kB 27.3MB/s \n",
+      "\u001b[K    100% |████████████████████████████████| 143kB 29.6MB/s \n",
+      "\u001b[K    100% |████████████████████████████████| 21.1MB 1.7MB/s \n",
+      "\u001b[K    100% |████████████████████████████████| 409kB 24.7MB/s \n",
+      "\u001b[K    100% |████████████████████████████████| 296kB 25.0MB/s \n",
+      "\u001b[K    100% |████████████████████████████████| 61kB 21.5MB/s \n",
+      "\u001b[?25h  Building wheel for pypng (setup.py) ... \u001b[?25ldone\n",
+      "\u001b[?25h  Building wheel for opt-einsum (setup.py) ... \u001b[?25ldone\n",
+      "\u001b[?25h"
+     ]
+    }
+   ],
+   "source": [
+    "!pip install -q -U tensor2tensor==1.13.1\n",
+    "!pip install -q tensorflow==1.13.1\n",
+    "!pip install -q gym[atari]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "id": "R7-Ni-39DGZW",
+    "colab_type": "code",
+    "colab": {}
+   },
+   "outputs": [],
+   "source": [
+    "# Helper function for playing videos in the colab.\n",
+    "def play_video(path):\n",
+    "  from IPython.core.magics.display import HTML\n",
+    "  display_path = \"/nbextensions/vid.mp4\"\n",
+    "  display_abs_path = \"/usr/local/share/jupyter\" + display_path\n",
+    "  !rm -f $display_abs_path\n",
+    "  !ffmpeg -loglevel error -i $path $display_abs_path\n",
+    "  return HTML(\"\"\"\n",
+    "    <video width=\"640\" height=\"480\" controls>\n",
+    "      <source src=\"{}\" type=\"video/mp4\">\n",
+    "    </video>\n",
+    "  \"\"\".format(display_path))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "pueuiKUmAOUT",
+    "colab_type": "text"
+   },
+   "source": [
+    "# Play using a pre-trained policy\n",
+    "\n",
+    "We provide pretrained policies for the following games from the Atari Learning Environment ( [ALE](https://github.com/mgbellemare/Arcade-Learning-Environment)) : alien,\n",
+    "amidar,\n",
+    " assault,\n",
+    " asterix,\n",
+    " asteroids,\n",
+    " atlantis,\n",
+    " bank_heist,\n",
+    " battle_zone,\n",
+    " beam_rider,\n",
+    " bowling,\n",
+    " boxing,\n",
+    " breakout,\n",
+    " chopper_command,\n",
+    " crazy_climber,\n",
+    " demon_attack,\n",
+    " fishing_derby,\n",
+    " freeway,\n",
+    " frostbite,\n",
+    " gopher,\n",
+    " gravitar,\n",
+    " hero,\n",
+    " ice_hockey,\n",
+    " jamesbond,\n",
+    " kangaroo,\n",
+    " krull,\n",
+    " kung_fu_master,\n",
+    " ms_pacman,\n",
+    " name_this_game,\n",
+    " pong,\n",
+    " private_eye,\n",
+    " qbert,\n",
+    " riverraid,\n",
+    " road_runner,\n",
+    " seaquest,\n",
+    " up_n_down,\n",
+    " yars_revenge.\n",
+    " \n",
+    " We have 5 checkpoints for each game saved on Google Storage. Run the following command get the storage path:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "id": "x9pKfNbDFfVh",
+    "colab_type": "code",
+    "outputId": "97e763cc-caaa-49c8-e532-fcbde828d1a2",
+    "executionInfo": {
+     "status": "ok",
+     "timestamp": 1.5532741511E12,
+     "user_tz": -60.0,
+     "elapsed": 6162.0,
+     "user": {
+      "displayName": "Piotr Miłoś",
+      "photoUrl": "https://lh3.googleusercontent.com/-050ZBEGpNAA/AAAAAAAAAAI/AAAAAAAAk9g/r6cv_J6J5qA/s64/photo.jpg",
+      "userId": "12158759908531801397"
+     }
+    },
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 147.0
+    },
+    "collapsed": false
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
+      "For more information, please see:\n",
+      "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
+      "  * https://github.com/tensorflow/addons\n",
+      "If you depend on functionality not listed there, please file an issue.\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/143'"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {
+      "tags": []
+     },
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# experiment_id is an integer from [0, 4].\n",
+    "def get_run_dir(game, experiment_id):\n",
+    "  from tensor2tensor.data_generators.gym_env import ATARI_GAMES_WITH_HUMAN_SCORE_NICE\n",
+    "  EXPERIMENTS_PER_GAME = 5\n",
+    "  run_id = ATARI_GAMES_WITH_HUMAN_SCORE_NICE.index(game) * EXPERIMENTS_PER_GAME + experiment_id + 1\n",
+    "  return \"gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/{}\".format(run_id)\n",
+    "\n",
+    "get_run_dir('pong', 2)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "77fFdm-cFEOB",
+    "colab_type": "text"
+   },
+   "source": [
+    "To evaluate and generate videos for a pretrained policy on Pong:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "id": "X-nGlbuTAQXj",
+    "colab_type": "code",
+    "outputId": "888968f2-f551-4a0f-9fc7-074a949362d6",
+    "executionInfo": {
+     "status": "ok",
+     "timestamp": 1.553271580737E12,
+     "user_tz": -60.0,
+     "elapsed": 842128.0,
+     "user": {
+      "displayName": "Piotr Kozakowski",
+      "photoUrl": "",
+      "userId": "01014928596539690143"
+     }
+    },
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 17088.0
+    },
+    "collapsed": true
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
+      "For more information, please see:\n",
+      "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
+      "  * https://github.com/tensorflow/addons\n",
+      "If you depend on functionality not listed there, please file an issue.\n",
+      "\n",
+      "INFO:tensorflow:Overriding hparams in rlmb_long_stochastic_discrete with game=pong,eval_max_num_noops=8,eval_sampling_temps=[0.5]\n",
+      "INFO:tensorflow:Evaluating metric mean_reward/eval/sampling_temp_0.5_max_noops_8_unclipped\n",
+      "2019-03-22 16:05:45.007030: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 2300000000 Hz\n",
+      "2019-03-22 16:05:45.007306: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x2697860 executing computations on platform Host. Devices:\n",
+      "2019-03-22 16:05:45.007346: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): <undefined>, <undefined>\n",
+      "2019-03-22 16:05:45.105281: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:998] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
+      "2019-03-22 16:05:45.105857: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x2697440 executing computations on platform CUDA. Devices:\n",
+      "2019-03-22 16:05:45.105908: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): Tesla K80, Compute Capability 3.7\n",
+      "2019-03-22 16:05:45.106380: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1433] Found device 0 with properties: \n",
+      "name: Tesla K80 major: 3 minor: 7 memoryClockRate(GHz): 0.8235\n",
+      "pciBusID: 0000:00:04.0\n",
+      "totalMemory: 11.17GiB freeMemory: 11.10GiB\n",
+      "2019-03-22 16:05:45.106420: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n",
+      "2019-03-22 16:05:45.499212: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n",
+      "2019-03-22 16:05:45.499307: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n",
+      "2019-03-22 16:05:45.499332: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n",
+      "2019-03-22 16:05:45.499671: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:42] Overriding allow_growth setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.\n",
+      "2019-03-22 16:05:45.499741: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n",
+      "INFO:tensorflow:Using DummyPolicyProblem for the policy.\n",
+      "INFO:tensorflow:Setting T2TModel mode to 'train'\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Colocations handled automatically by placer.\n",
+      "INFO:tensorflow:Using variable initializer: orthogonal\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/utils/t2t_model.py:1358: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use tf.cast instead.\n",
+      "INFO:tensorflow:Transforming feature 'input_action' with symbol_modality_6_64.bottom\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/function.py:1007: calling Graph.create_op (from tensorflow.python.framework.ops) with compute_shapes is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Shapes are always computed; don't use the compute_shapes as it has no effect.\n",
+      "INFO:tensorflow:Transforming feature 'input_reward' with symbol_modality_3_64.bottom\n",
+      "INFO:tensorflow:Transforming feature 'inputs' with video_modality.bottom\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_video.py:495: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "tf.py_func is deprecated in TF V2. Instead, use\n",
+      "    tf.py_function, which takes a python function which manipulates tf eager\n",
+      "    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to\n",
+      "    an ndarray (just call tensor.numpy()) but having access to eager tensors\n",
+      "    means `tf.py_function`s can use accelerators such as GPUs as well as\n",
+      "    being differentiable using a gradient tape.\n",
+      "    \n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:277: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use tf.cast instead.\n",
+      "INFO:tensorflow:Transforming feature 'target_action' with symbol_modality_6_64.targets_bottom\n",
+      "INFO:tensorflow:Transforming feature 'target_policy' with identity_modality.targets_bottom\n",
+      "INFO:tensorflow:Transforming feature 'target_reward' with symbol_modality_3_64.targets_bottom\n",
+      "INFO:tensorflow:Transforming feature 'target_value' with identity_modality.targets_bottom\n",
+      "INFO:tensorflow:Transforming feature 'targets' with video_modality.targets_bottom\n",
+      "INFO:tensorflow:Building model body\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:598: conv2d (from tensorflow.python.layers.convolutional) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use keras.layers.conv2d instead.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:602: flatten (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use keras.layers.flatten instead.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:603: dropout (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use keras.layers.dropout instead.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:604: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use keras.layers.dense instead.\n",
+      "INFO:tensorflow:Transforming body output with identity_modality.top\n",
+      "INFO:tensorflow:Transforming body output with identity_modality.top\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:2887: multinomial (from tensorflow.python.ops.random_ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use tf.random.categorical instead.\n",
+      "2019-03-22 16:06:00.352605: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n",
+      "2019-03-22 16:06:00.352688: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n",
+      "2019-03-22 16:06:00.352724: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n",
+      "2019-03-22 16:06:00.352744: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n",
+      "2019-03-22 16:06:00.353037: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n",
+      "2019-03-22 16:06:00.588787: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n",
+      "2019-03-22 16:06:00.647797: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n",
+      "INFO:tensorflow:Restoring checkpoint gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/policy/model.ckpt-171992\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/saver.py:1266: checkpoint_exists (from tensorflow.python.training.checkpoint_management) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use standard file APIs to check for files with this prefix.\n",
+      "2019-03-22 16:06:00.711910: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n",
+      "INFO:tensorflow:Restoring parameters from gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/policy/model.ckpt-171992\n",
+      "2019-03-22 16:06:00.793701: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n",
+      "2019-03-22 16:06:00.953239: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n",
+      "2019-03-22 16:06:01.086594: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n",
+      "2019-03-22 16:06:01.259521: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n",
+      "2019-03-22 16:06:01.322896: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n",
+      "2019-03-22 16:06:03.034751: I tensorflow/stream_executor/dso_loader.cc:152] successfully opened CUDA library libcublas.so.10.0 locally\n",
+      "INFO:tensorflow:Step 5, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 10, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 15, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 20, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 25, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 30, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 35, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 40, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 45, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 50, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 55, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 60, mean_score: -0.015625\n",
+      "INFO:tensorflow:Step 65, mean_score: -0.078125\n",
+      "INFO:tensorflow:Step 70, mean_score: -0.078125\n",
+      "INFO:tensorflow:Step 75, mean_score: -0.078125\n",
+      "INFO:tensorflow:Step 80, mean_score: -0.078125\n",
+      "INFO:tensorflow:Step 85, mean_score: -0.078125\n",
+      "INFO:tensorflow:Step 90, mean_score: 0.484375\n",
+      "INFO:tensorflow:Step 95, mean_score: 0.843750\n",
+      "INFO:tensorflow:Step 100, mean_score: 0.828125\n",
+      "INFO:tensorflow:Step 105, mean_score: 0.828125\n",
+      "INFO:tensorflow:Step 110, mean_score: 0.828125\n",
+      "INFO:tensorflow:Step 115, mean_score: 0.828125\n",
+      "INFO:tensorflow:Step 120, mean_score: 0.828125\n",
+      "INFO:tensorflow:Step 125, mean_score: 0.828125\n",
+      "INFO:tensorflow:Step 130, mean_score: 0.812500\n",
+      "INFO:tensorflow:Step 135, mean_score: 0.812500\n",
+      "INFO:tensorflow:Step 140, mean_score: 0.812500\n",
+      "INFO:tensorflow:Step 145, mean_score: 0.812500\n",
+      "INFO:tensorflow:Step 150, mean_score: 0.812500\n",
+      "INFO:tensorflow:Step 155, mean_score: 0.812500\n",
+      "INFO:tensorflow:Step 160, mean_score: 0.812500\n",
+      "INFO:tensorflow:Step 165, mean_score: 0.812500\n",
+      "INFO:tensorflow:Step 170, mean_score: 0.828125\n",
+      "INFO:tensorflow:Step 175, mean_score: 0.843750\n",
+      "INFO:tensorflow:Step 180, mean_score: 0.843750\n",
+      "INFO:tensorflow:Step 185, mean_score: 0.843750\n",
+      "INFO:tensorflow:Step 190, mean_score: 1.140625\n",
+      "INFO:tensorflow:Step 195, mean_score: 1.765625\n",
+      "INFO:tensorflow:Step 200, mean_score: 1.765625\n",
+      "INFO:tensorflow:Step 205, mean_score: 1.765625\n",
+      "INFO:tensorflow:Step 210, mean_score: 1.781250\n",
+      "INFO:tensorflow:Step 215, mean_score: 1.781250\n",
+      "INFO:tensorflow:Step 220, mean_score: 1.765625\n",
+      "INFO:tensorflow:Step 225, mean_score: 1.765625\n",
+      "INFO:tensorflow:Step 230, mean_score: 1.765625\n",
+      "INFO:tensorflow:Step 235, mean_score: 1.765625\n",
+      "INFO:tensorflow:Step 240, mean_score: 1.765625\n",
+      "INFO:tensorflow:Step 245, mean_score: 1.765625\n",
+      "INFO:tensorflow:Step 250, mean_score: 1.765625\n",
+      "INFO:tensorflow:Step 255, mean_score: 1.750000\n",
+      "INFO:tensorflow:Step 260, mean_score: 1.750000\n",
+      "INFO:tensorflow:Step 265, mean_score: 1.750000\n",
+      "INFO:tensorflow:Step 270, mean_score: 2.312500\n",
+      "INFO:tensorflow:Step 275, mean_score: 2.687500\n",
+      "INFO:tensorflow:Step 280, mean_score: 2.703125\n",
+      "INFO:tensorflow:Step 285, mean_score: 2.703125\n",
+      "INFO:tensorflow:Step 290, mean_score: 2.703125\n",
+      "INFO:tensorflow:Step 295, mean_score: 2.703125\n",
+      "INFO:tensorflow:Step 300, mean_score: 2.703125\n",
+      "INFO:tensorflow:Step 305, mean_score: 2.703125\n",
+      "INFO:tensorflow:Step 310, mean_score: 2.718750\n",
+      "INFO:tensorflow:Step 315, mean_score: 2.718750\n",
+      "INFO:tensorflow:Step 320, mean_score: 2.718750\n",
+      "INFO:tensorflow:Step 325, mean_score: 2.718750\n",
+      "INFO:tensorflow:Step 330, mean_score: 2.734375\n",
+      "INFO:tensorflow:Step 335, mean_score: 2.734375\n",
+      "INFO:tensorflow:Step 340, mean_score: 2.734375\n",
+      "INFO:tensorflow:Step 345, mean_score: 2.734375\n",
+      "INFO:tensorflow:Step 350, mean_score: 2.750000\n",
+      "INFO:tensorflow:Step 355, mean_score: 2.765625\n",
+      "INFO:tensorflow:Step 360, mean_score: 2.765625\n",
+      "INFO:tensorflow:Step 365, mean_score: 2.765625\n",
+      "INFO:tensorflow:Step 370, mean_score: 3.062500\n",
+      "INFO:tensorflow:Step 375, mean_score: 3.687500\n",
+      "INFO:tensorflow:Step 380, mean_score: 3.687500\n",
+      "INFO:tensorflow:Step 385, mean_score: 3.687500\n",
+      "INFO:tensorflow:Step 390, mean_score: 3.703125\n",
+      "INFO:tensorflow:Step 395, mean_score: 3.703125\n",
+      "INFO:tensorflow:Step 400, mean_score: 3.703125\n",
+      "INFO:tensorflow:Step 405, mean_score: 3.703125\n",
+      "INFO:tensorflow:Step 410, mean_score: 3.687500\n",
+      "INFO:tensorflow:Step 415, mean_score: 3.687500\n",
+      "INFO:tensorflow:Step 420, mean_score: 3.687500\n",
+      "INFO:tensorflow:Step 425, mean_score: 3.687500\n",
+      "INFO:tensorflow:Step 430, mean_score: 3.703125\n",
+      "INFO:tensorflow:Step 435, mean_score: 3.703125\n",
+      "INFO:tensorflow:Step 440, mean_score: 3.703125\n",
+      "INFO:tensorflow:Step 445, mean_score: 3.703125\n",
+      "INFO:tensorflow:Step 450, mean_score: 4.265625\n",
+      "INFO:tensorflow:Step 455, mean_score: 4.640625\n",
+      "INFO:tensorflow:Step 460, mean_score: 4.656250\n",
+      "INFO:tensorflow:Step 465, mean_score: 4.656250\n",
+      "INFO:tensorflow:Step 470, mean_score: 4.656250\n",
+      "INFO:tensorflow:Step 475, mean_score: 4.656250\n",
+      "INFO:tensorflow:Step 480, mean_score: 4.656250\n",
+      "INFO:tensorflow:Step 485, mean_score: 4.656250\n",
+      "INFO:tensorflow:Step 490, mean_score: 4.671875\n",
+      "INFO:tensorflow:Step 495, mean_score: 4.671875\n",
+      "INFO:tensorflow:Step 500, mean_score: 4.671875\n",
+      "INFO:tensorflow:Step 505, mean_score: 4.671875\n",
+      "INFO:tensorflow:Step 510, mean_score: 4.687500\n",
+      "INFO:tensorflow:Step 515, mean_score: 4.687500\n",
+      "INFO:tensorflow:Step 520, mean_score: 4.703125\n",
+      "INFO:tensorflow:Step 525, mean_score: 4.703125\n",
+      "INFO:tensorflow:Step 530, mean_score: 4.718750\n",
+      "INFO:tensorflow:Step 535, mean_score: 4.734375\n",
+      "INFO:tensorflow:Step 540, mean_score: 4.734375\n",
+      "INFO:tensorflow:Step 545, mean_score: 4.734375\n",
+      "INFO:tensorflow:Step 550, mean_score: 5.031250\n",
+      "INFO:tensorflow:Step 555, mean_score: 5.656250\n",
+      "INFO:tensorflow:Step 560, mean_score: 5.656250\n",
+      "INFO:tensorflow:Step 565, mean_score: 5.656250\n",
+      "INFO:tensorflow:Step 570, mean_score: 5.671875\n",
+      "INFO:tensorflow:Step 575, mean_score: 5.671875\n",
+      "INFO:tensorflow:Step 580, mean_score: 5.671875\n",
+      "INFO:tensorflow:Step 585, mean_score: 5.671875\n",
+      "INFO:tensorflow:Step 590, mean_score: 5.671875\n",
+      "INFO:tensorflow:Step 595, mean_score: 5.671875\n",
+      "INFO:tensorflow:Step 600, mean_score: 5.671875\n",
+      "INFO:tensorflow:Step 605, mean_score: 5.671875\n",
+      "INFO:tensorflow:Step 610, mean_score: 5.687500\n",
+      "INFO:tensorflow:Step 615, mean_score: 5.687500\n",
+      "INFO:tensorflow:Step 620, mean_score: 5.703125\n",
+      "INFO:tensorflow:Step 625, mean_score: 5.703125\n",
+      "INFO:tensorflow:Step 630, mean_score: 6.265625\n",
+      "INFO:tensorflow:Step 635, mean_score: 6.640625\n",
+      "INFO:tensorflow:Step 640, mean_score: 6.656250\n",
+      "INFO:tensorflow:Step 645, mean_score: 6.656250\n",
+      "INFO:tensorflow:Step 650, mean_score: 6.656250\n",
+      "INFO:tensorflow:Step 655, mean_score: 6.656250\n",
+      "INFO:tensorflow:Step 660, mean_score: 6.656250\n",
+      "INFO:tensorflow:Step 665, mean_score: 6.656250\n",
+      "INFO:tensorflow:Step 670, mean_score: 6.671875\n",
+      "INFO:tensorflow:Step 675, mean_score: 6.671875\n",
+      "INFO:tensorflow:Step 680, mean_score: 6.671875\n",
+      "INFO:tensorflow:Step 685, mean_score: 6.671875\n",
+      "INFO:tensorflow:Step 690, mean_score: 6.687500\n",
+      "INFO:tensorflow:Step 695, mean_score: 6.687500\n",
+      "INFO:tensorflow:Step 700, mean_score: 6.703125\n",
+      "INFO:tensorflow:Step 705, mean_score: 6.703125\n",
+      "INFO:tensorflow:Step 710, mean_score: 6.718750\n",
+      "INFO:tensorflow:Step 715, mean_score: 6.734375\n",
+      "INFO:tensorflow:Step 720, mean_score: 6.734375\n",
+      "INFO:tensorflow:Step 725, mean_score: 6.734375\n",
+      "INFO:tensorflow:Step 730, mean_score: 7.031250\n",
+      "INFO:tensorflow:Step 735, mean_score: 7.656250\n",
+      "INFO:tensorflow:Step 740, mean_score: 7.656250\n",
+      "INFO:tensorflow:Step 745, mean_score: 7.656250\n",
+      "INFO:tensorflow:Step 750, mean_score: 7.671875\n",
+      "INFO:tensorflow:Step 755, mean_score: 7.671875\n",
+      "INFO:tensorflow:Step 760, mean_score: 7.671875\n",
+      "INFO:tensorflow:Step 765, mean_score: 7.671875\n",
+      "INFO:tensorflow:Step 770, mean_score: 7.671875\n",
+      "INFO:tensorflow:Step 775, mean_score: 7.671875\n",
+      "INFO:tensorflow:Step 780, mean_score: 7.671875\n",
+      "INFO:tensorflow:Step 785, mean_score: 7.671875\n",
+      "INFO:tensorflow:Step 790, mean_score: 7.687500\n",
+      "INFO:tensorflow:Step 795, mean_score: 7.687500\n",
+      "INFO:tensorflow:Step 800, mean_score: 7.703125\n",
+      "INFO:tensorflow:Step 805, mean_score: 7.703125\n",
+      "INFO:tensorflow:Step 810, mean_score: 8.265625\n",
+      "INFO:tensorflow:Step 815, mean_score: 8.640625\n",
+      "INFO:tensorflow:Step 820, mean_score: 8.656250\n",
+      "INFO:tensorflow:Step 825, mean_score: 8.656250\n",
+      "INFO:tensorflow:Step 830, mean_score: 8.656250\n",
+      "INFO:tensorflow:Step 835, mean_score: 8.656250\n",
+      "INFO:tensorflow:Step 840, mean_score: 8.656250\n",
+      "INFO:tensorflow:Step 845, mean_score: 8.656250\n",
+      "INFO:tensorflow:Step 850, mean_score: 8.671875\n",
+      "INFO:tensorflow:Step 855, mean_score: 8.671875\n",
+      "INFO:tensorflow:Step 860, mean_score: 8.671875\n",
+      "INFO:tensorflow:Step 865, mean_score: 8.671875\n",
+      "INFO:tensorflow:Step 870, mean_score: 8.687500\n",
+      "INFO:tensorflow:Step 875, mean_score: 8.687500\n",
+      "INFO:tensorflow:Step 880, mean_score: 8.703125\n",
+      "INFO:tensorflow:Step 885, mean_score: 8.703125\n",
+      "INFO:tensorflow:Step 890, mean_score: 8.718750\n",
+      "INFO:tensorflow:Step 895, mean_score: 8.734375\n",
+      "INFO:tensorflow:Step 900, mean_score: 8.734375\n",
+      "INFO:tensorflow:Step 905, mean_score: 8.734375\n",
+      "INFO:tensorflow:Step 910, mean_score: 9.031250\n",
+      "INFO:tensorflow:Step 915, mean_score: 9.656250\n",
+      "INFO:tensorflow:Step 920, mean_score: 9.656250\n",
+      "INFO:tensorflow:Step 925, mean_score: 9.656250\n",
+      "INFO:tensorflow:Step 930, mean_score: 9.671875\n",
+      "INFO:tensorflow:Step 935, mean_score: 9.671875\n",
+      "INFO:tensorflow:Step 940, mean_score: 9.671875\n",
+      "INFO:tensorflow:Step 945, mean_score: 9.671875\n",
+      "INFO:tensorflow:Step 950, mean_score: 9.671875\n",
+      "INFO:tensorflow:Step 955, mean_score: 9.671875\n",
+      "INFO:tensorflow:Step 960, mean_score: 9.671875\n",
+      "INFO:tensorflow:Step 965, mean_score: 9.671875\n",
+      "INFO:tensorflow:Step 970, mean_score: 9.687500\n",
+      "INFO:tensorflow:Step 975, mean_score: 9.687500\n",
+      "INFO:tensorflow:Step 980, mean_score: 9.703125\n",
+      "INFO:tensorflow:Step 985, mean_score: 9.703125\n",
+      "INFO:tensorflow:Step 990, mean_score: 10.265625\n",
+      "INFO:tensorflow:Step 995, mean_score: 10.640625\n",
+      "INFO:tensorflow:Step 1000, mean_score: 10.656250\n",
+      "INFO:tensorflow:Step 1005, mean_score: 10.656250\n",
+      "INFO:tensorflow:Step 1010, mean_score: 10.656250\n",
+      "INFO:tensorflow:Step 1015, mean_score: 10.656250\n",
+      "INFO:tensorflow:Step 1020, mean_score: 10.656250\n",
+      "INFO:tensorflow:Step 1025, mean_score: 10.656250\n",
+      "INFO:tensorflow:Step 1030, mean_score: 10.671875\n",
+      "INFO:tensorflow:Step 1035, mean_score: 10.671875\n",
+      "INFO:tensorflow:Step 1040, mean_score: 10.671875\n",
+      "INFO:tensorflow:Step 1045, mean_score: 10.671875\n",
+      "INFO:tensorflow:Step 1050, mean_score: 10.687500\n",
+      "INFO:tensorflow:Step 1055, mean_score: 10.687500\n",
+      "INFO:tensorflow:Step 1060, mean_score: 10.703125\n",
+      "INFO:tensorflow:Step 1065, mean_score: 10.703125\n",
+      "INFO:tensorflow:Step 1070, mean_score: 10.718750\n",
+      "INFO:tensorflow:Step 1075, mean_score: 10.734375\n",
+      "INFO:tensorflow:Step 1080, mean_score: 10.734375\n",
+      "INFO:tensorflow:Step 1085, mean_score: 10.734375\n",
+      "INFO:tensorflow:Step 1090, mean_score: 11.031250\n",
+      "INFO:tensorflow:Step 1095, mean_score: 11.656250\n",
+      "INFO:tensorflow:Step 1100, mean_score: 11.656250\n",
+      "INFO:tensorflow:Step 1105, mean_score: 11.656250\n",
+      "INFO:tensorflow:Step 1110, mean_score: 11.671875\n",
+      "INFO:tensorflow:Step 1115, mean_score: 11.671875\n",
+      "INFO:tensorflow:Step 1120, mean_score: 11.671875\n",
+      "INFO:tensorflow:Step 1125, mean_score: 11.671875\n",
+      "INFO:tensorflow:Step 1130, mean_score: 11.671875\n",
+      "INFO:tensorflow:Step 1135, mean_score: 11.671875\n",
+      "INFO:tensorflow:Step 1140, mean_score: 11.671875\n",
+      "INFO:tensorflow:Step 1145, mean_score: 11.671875\n",
+      "INFO:tensorflow:Step 1150, mean_score: 11.687500\n",
+      "INFO:tensorflow:Step 1155, mean_score: 11.687500\n",
+      "INFO:tensorflow:Step 1160, mean_score: 11.703125\n",
+      "INFO:tensorflow:Step 1165, mean_score: 11.703125\n",
+      "INFO:tensorflow:Step 1170, mean_score: 12.265625\n",
+      "INFO:tensorflow:Step 1175, mean_score: 12.640625\n",
+      "INFO:tensorflow:Step 1180, mean_score: 12.656250\n",
+      "INFO:tensorflow:Step 1185, mean_score: 12.656250\n",
+      "INFO:tensorflow:Step 1190, mean_score: 12.656250\n",
+      "INFO:tensorflow:Step 1195, mean_score: 12.656250\n",
+      "INFO:tensorflow:Step 1200, mean_score: 12.656250\n",
+      "INFO:tensorflow:Step 1205, mean_score: 12.656250\n",
+      "INFO:tensorflow:Step 1210, mean_score: 12.671875\n",
+      "INFO:tensorflow:Step 1215, mean_score: 12.671875\n",
+      "INFO:tensorflow:Step 1220, mean_score: 12.671875\n",
+      "INFO:tensorflow:Step 1225, mean_score: 12.671875\n",
+      "INFO:tensorflow:Step 1230, mean_score: 12.687500\n",
+      "INFO:tensorflow:Step 1235, mean_score: 12.687500\n",
+      "INFO:tensorflow:Step 1240, mean_score: 12.703125\n",
+      "INFO:tensorflow:Step 1245, mean_score: 12.703125\n",
+      "INFO:tensorflow:Step 1250, mean_score: 12.718750\n",
+      "INFO:tensorflow:Step 1255, mean_score: 12.734375\n",
+      "INFO:tensorflow:Step 1260, mean_score: 12.734375\n",
+      "INFO:tensorflow:Step 1265, mean_score: 12.734375\n",
+      "INFO:tensorflow:Step 1270, mean_score: 13.031250\n",
+      "INFO:tensorflow:Step 1275, mean_score: 13.656250\n",
+      "INFO:tensorflow:Step 1280, mean_score: 13.656250\n",
+      "INFO:tensorflow:Step 1285, mean_score: 13.656250\n",
+      "INFO:tensorflow:Step 1290, mean_score: 13.671875\n",
+      "INFO:tensorflow:Step 1295, mean_score: 13.671875\n",
+      "INFO:tensorflow:Step 1300, mean_score: 13.671875\n",
+      "INFO:tensorflow:Step 1305, mean_score: 13.671875\n",
+      "INFO:tensorflow:Step 1310, mean_score: 13.671875\n",
+      "INFO:tensorflow:Step 1315, mean_score: 13.671875\n",
+      "INFO:tensorflow:Step 1320, mean_score: 13.671875\n",
+      "INFO:tensorflow:Step 1325, mean_score: 13.671875\n",
+      "INFO:tensorflow:Step 1330, mean_score: 13.687500\n",
+      "INFO:tensorflow:Step 1335, mean_score: 13.687500\n",
+      "INFO:tensorflow:Step 1340, mean_score: 13.703125\n",
+      "INFO:tensorflow:Step 1345, mean_score: 13.703125\n",
+      "INFO:tensorflow:Step 1350, mean_score: 14.265625\n",
+      "INFO:tensorflow:Step 1355, mean_score: 14.640625\n",
+      "INFO:tensorflow:Step 1360, mean_score: 14.656250\n",
+      "INFO:tensorflow:Step 1365, mean_score: 14.656250\n",
+      "INFO:tensorflow:Step 1370, mean_score: 14.656250\n",
+      "INFO:tensorflow:Step 1375, mean_score: 14.656250\n",
+      "INFO:tensorflow:Step 1380, mean_score: 14.656250\n",
+      "INFO:tensorflow:Step 1385, mean_score: 14.656250\n",
+      "INFO:tensorflow:Step 1390, mean_score: 14.671875\n",
+      "INFO:tensorflow:Step 1395, mean_score: 14.671875\n",
+      "INFO:tensorflow:Step 1400, mean_score: 14.671875\n",
+      "INFO:tensorflow:Step 1405, mean_score: 14.671875\n",
+      "INFO:tensorflow:Step 1410, mean_score: 14.687500\n",
+      "INFO:tensorflow:Step 1415, mean_score: 14.687500\n",
+      "INFO:tensorflow:Step 1420, mean_score: 14.703125\n",
+      "INFO:tensorflow:Step 1425, mean_score: 14.703125\n",
+      "INFO:tensorflow:Step 1430, mean_score: 14.718750\n",
+      "INFO:tensorflow:Step 1435, mean_score: 14.734375\n",
+      "INFO:tensorflow:Step 1440, mean_score: 14.734375\n",
+      "INFO:tensorflow:Step 1445, mean_score: 14.734375\n",
+      "INFO:tensorflow:Step 1450, mean_score: 15.031250\n",
+      "INFO:tensorflow:Step 1455, mean_score: 15.656250\n",
+      "INFO:tensorflow:Step 1460, mean_score: 15.656250\n",
+      "INFO:tensorflow:Step 1465, mean_score: 15.656250\n",
+      "INFO:tensorflow:Step 1470, mean_score: 15.671875\n",
+      "INFO:tensorflow:Step 1475, mean_score: 15.671875\n",
+      "INFO:tensorflow:Step 1480, mean_score: 15.671875\n",
+      "INFO:tensorflow:Step 1485, mean_score: 15.671875\n",
+      "INFO:tensorflow:Step 1490, mean_score: 15.671875\n",
+      "INFO:tensorflow:Step 1495, mean_score: 15.671875\n",
+      "INFO:tensorflow:Step 1500, mean_score: 15.671875\n",
+      "INFO:tensorflow:Step 1505, mean_score: 15.671875\n",
+      "INFO:tensorflow:Step 1510, mean_score: 15.687500\n",
+      "INFO:tensorflow:Step 1515, mean_score: 15.687500\n",
+      "INFO:tensorflow:Step 1520, mean_score: 15.703125\n",
+      "INFO:tensorflow:Step 1525, mean_score: 15.703125\n",
+      "INFO:tensorflow:Step 1530, mean_score: 16.265625\n",
+      "INFO:tensorflow:Step 1535, mean_score: 16.640625\n",
+      "INFO:tensorflow:Step 1540, mean_score: 16.656250\n",
+      "INFO:tensorflow:Step 1545, mean_score: 16.656250\n",
+      "INFO:tensorflow:Step 1550, mean_score: 16.656250\n",
+      "INFO:tensorflow:Step 1555, mean_score: 16.656250\n",
+      "INFO:tensorflow:Step 1560, mean_score: 16.656250\n",
+      "INFO:tensorflow:Step 1565, mean_score: 16.656250\n",
+      "INFO:tensorflow:Step 1570, mean_score: 16.671875\n",
+      "INFO:tensorflow:Step 1575, mean_score: 16.671875\n",
+      "INFO:tensorflow:Step 1580, mean_score: 16.671875\n",
+      "INFO:tensorflow:Step 1585, mean_score: 16.671875\n",
+      "INFO:tensorflow:Step 1590, mean_score: 16.687500\n",
+      "INFO:tensorflow:Step 1595, mean_score: 16.687500\n",
+      "INFO:tensorflow:Step 1600, mean_score: 16.703125\n",
+      "INFO:tensorflow:Step 1605, mean_score: 16.703125\n",
+      "INFO:tensorflow:Step 1610, mean_score: 16.718750\n",
+      "INFO:tensorflow:Step 1615, mean_score: 16.734375\n",
+      "INFO:tensorflow:Step 1620, mean_score: 16.734375\n",
+      "INFO:tensorflow:Step 1625, mean_score: 16.734375\n",
+      "INFO:tensorflow:Step 1630, mean_score: 17.031250\n",
+      "INFO:tensorflow:Step 1635, mean_score: 17.656250\n",
+      "INFO:tensorflow:Step 1640, mean_score: 17.656250\n",
+      "INFO:tensorflow:Step 1645, mean_score: 17.656250\n",
+      "INFO:tensorflow:Step 1650, mean_score: 17.671875\n",
+      "INFO:tensorflow:Step 1655, mean_score: 17.671875\n",
+      "INFO:tensorflow:Step 1660, mean_score: 17.671875\n",
+      "INFO:tensorflow:Step 1665, mean_score: 17.671875\n",
+      "INFO:tensorflow:Step 1670, mean_score: 17.671875\n",
+      "INFO:tensorflow:Step 1675, mean_score: 17.671875\n",
+      "INFO:tensorflow:Step 1680, mean_score: 17.671875\n",
+      "INFO:tensorflow:Step 1685, mean_score: 17.671875\n",
+      "INFO:tensorflow:Step 1690, mean_score: 17.687500\n",
+      "INFO:tensorflow:Step 1695, mean_score: 17.687500\n",
+      "INFO:tensorflow:Step 1700, mean_score: 17.703125\n",
+      "INFO:tensorflow:Step 1705, mean_score: 17.703125\n",
+      "INFO:tensorflow:Step 1710, mean_score: 18.265625\n",
+      "INFO:tensorflow:Step 1715, mean_score: 18.640625\n",
+      "INFO:tensorflow:Step 1720, mean_score: 18.656250\n",
+      "INFO:tensorflow:Step 1725, mean_score: 18.656250\n",
+      "INFO:tensorflow:Step 1730, mean_score: 18.656250\n",
+      "INFO:tensorflow:Step 1735, mean_score: 18.656250\n",
+      "INFO:tensorflow:Step 1740, mean_score: 18.656250\n",
+      "INFO:tensorflow:Step 1745, mean_score: 18.656250\n",
+      "INFO:tensorflow:Step 1750, mean_score: 18.671875\n",
+      "INFO:tensorflow:Step 1755, mean_score: 18.671875\n",
+      "INFO:tensorflow:Step 1760, mean_score: 18.671875\n",
+      "INFO:tensorflow:Step 1765, mean_score: 18.671875\n",
+      "INFO:tensorflow:Step 1770, mean_score: 18.687500\n",
+      "INFO:tensorflow:Step 1775, mean_score: 18.687500\n",
+      "INFO:tensorflow:Step 1780, mean_score: 18.703125\n",
+      "INFO:tensorflow:Step 1785, mean_score: 18.703125\n",
+      "INFO:tensorflow:Step 1790, mean_score: 18.718750\n",
+      "INFO:tensorflow:Step 1795, mean_score: 18.734375\n",
+      "INFO:tensorflow:Step 1800, mean_score: 18.734375\n",
+      "INFO:tensorflow:Step 1805, mean_score: 18.734375\n",
+      "INFO:tensorflow:Step 1810, mean_score: 19.031250\n",
+      "INFO:tensorflow:Step 1815, mean_score: 19.656250\n",
+      "INFO:tensorflow:Step 1820, mean_score: 19.656250\n",
+      "INFO:tensorflow:Step 1825, mean_score: 19.656250\n",
+      "INFO:tensorflow:Step 1830, mean_score: 19.671875\n",
+      "INFO:tensorflow:Step 1835, mean_score: 19.671875\n",
+      "INFO:tensorflow:Step 1840, mean_score: 19.671875\n",
+      "INFO:tensorflow:Step 1845, mean_score: 19.671875\n",
+      "INFO:tensorflow:Step 1850, mean_score: 19.671875\n",
+      "INFO:tensorflow:Step 1855, mean_score: 19.671875\n",
+      "INFO:tensorflow:Step 1860, mean_score: 19.671875\n",
+      "INFO:tensorflow:Step 1865, mean_score: 19.671875\n",
+      "INFO:tensorflow:Step 1870, mean_score: 19.687500\n",
+      "INFO:tensorflow:Step 1875, mean_score: 19.687500\n",
+      "INFO:tensorflow:Step 1880, mean_score: 19.703125\n",
+      "INFO:tensorflow:Step 1885, mean_score: 19.703125\n",
+      "INFO:tensorflow:Step 1890, mean_score: 19.703125\n",
+      "INFO:tensorflow:Step 1895, mean_score: 19.718750\n",
+      "INFO:tensorflow:Step 1900, mean_score: 19.734375\n",
+      "INFO:tensorflow:Step 1905, mean_score: 19.734375\n",
+      "INFO:tensorflow:Step 1910, mean_score: 19.734375\n",
+      "INFO:tensorflow:Step 1915, mean_score: 19.734375\n",
+      "INFO:tensorflow:Step 1920, mean_score: 19.734375\n",
+      "INFO:tensorflow:Step 1925, mean_score: 19.734375\n",
+      "INFO:tensorflow:Step 1930, mean_score: 19.750000\n",
+      "INFO:tensorflow:Step 1935, mean_score: 19.750000\n",
+      "INFO:tensorflow:Step 1940, mean_score: 19.750000\n",
+      "INFO:tensorflow:Step 1945, mean_score: 19.750000\n",
+      "INFO:tensorflow:Step 1950, mean_score: 19.765625\n",
+      "INFO:tensorflow:Step 1955, mean_score: 19.765625\n",
+      "INFO:tensorflow:Step 1960, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 1965, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 1970, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 1975, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 1980, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 1985, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 1990, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 1995, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 2000, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 2005, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 2010, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 2015, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 2020, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 2025, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 2030, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 2035, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 2040, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 2045, mean_score: 19.781250\n",
+      "INFO:tensorflow:Step 2050, mean_score: 19.796875\n",
+      "INFO:tensorflow:Step 2055, mean_score: 19.796875\n",
+      "INFO:tensorflow:Step 2060, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2065, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2070, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2075, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2080, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2085, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2090, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2095, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2100, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2105, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2110, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2115, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2120, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2125, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2130, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2135, mean_score: 19.812500\n",
+      "INFO:tensorflow:Step 2140, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2145, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2150, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2155, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2160, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2165, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2170, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2175, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2180, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2185, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2190, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2195, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2200, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2205, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2210, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2215, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2220, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2225, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2230, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2235, mean_score: 19.828125\n",
+      "INFO:tensorflow:Step 2240, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2245, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2250, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2255, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2260, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2265, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2270, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2275, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2280, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2285, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2290, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2295, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2300, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2305, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2310, mean_score: 19.843750\n",
+      "INFO:tensorflow:Step 2315, mean_score: 19.843750\n",
+      "INFO:tensorflow:Evaluating metric mean_reward/eval/sampling_temp_0.5_max_noops_0_unclipped\n",
+      "2019-03-22 16:12:57.935045: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n",
+      "2019-03-22 16:12:57.935160: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n",
+      "2019-03-22 16:12:57.935189: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n",
+      "2019-03-22 16:12:57.935209: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n",
+      "2019-03-22 16:12:57.935553: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n",
+      "INFO:tensorflow:Using DummyPolicyProblem for the policy.\n",
+      "INFO:tensorflow:Setting T2TModel mode to 'train'\n",
+      "INFO:tensorflow:Using variable initializer: orthogonal\n",
+      "INFO:tensorflow:Transforming feature 'input_action' with symbol_modality_6_64.bottom\n",
+      "INFO:tensorflow:Transforming feature 'input_reward' with symbol_modality_3_64.bottom\n",
+      "INFO:tensorflow:Transforming feature 'inputs' with video_modality.bottom\n",
+      "INFO:tensorflow:Transforming feature 'target_action' with symbol_modality_6_64.targets_bottom\n",
+      "INFO:tensorflow:Transforming feature 'target_policy' with identity_modality.targets_bottom\n",
+      "INFO:tensorflow:Transforming feature 'target_reward' with symbol_modality_3_64.targets_bottom\n",
+      "INFO:tensorflow:Transforming feature 'target_value' with identity_modality.targets_bottom\n",
+      "INFO:tensorflow:Transforming feature 'targets' with video_modality.targets_bottom\n",
+      "INFO:tensorflow:Building model body\n",
+      "INFO:tensorflow:Transforming body output with identity_modality.top\n",
+      "INFO:tensorflow:Transforming body output with identity_modality.top\n",
+      "2019-03-22 16:13:12.260846: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n",
+      "2019-03-22 16:13:12.260981: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n",
+      "2019-03-22 16:13:12.261059: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n",
+      "2019-03-22 16:13:12.261099: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n",
+      "2019-03-22 16:13:12.261613: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n",
+      "2019-03-22 16:13:12.493082: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n",
+      "INFO:tensorflow:Restoring checkpoint gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/policy/model.ckpt-171992\n",
+      "2019-03-22 16:13:12.556955: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n",
+      "INFO:tensorflow:Restoring parameters from gs://tensor2tensor-checkpoints/modelrl_experiments/train_sd/142/policy/model.ckpt-171992\n",
+      "2019-03-22 16:13:12.651009: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n",
+      "2019-03-22 16:13:12.715180: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n",
+      "2019-03-22 16:13:12.816774: W tensorflow/core/platform/cloud/google_auth_provider.cc:178] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Cancelled: GCE check skipped due to presence of $NO_GCE_CHECK environment variable.\".\n",
+      "INFO:tensorflow:Step 5, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 10, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 15, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 20, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 25, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 30, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 35, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 40, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 45, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 50, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 55, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 60, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 65, mean_score: -0.031250\n",
+      "INFO:tensorflow:Step 70, mean_score: -0.031250\n",
+      "INFO:tensorflow:Step 75, mean_score: -0.031250\n",
+      "INFO:tensorflow:Step 80, mean_score: -0.031250\n",
+      "INFO:tensorflow:Step 85, mean_score: -0.031250\n",
+      "INFO:tensorflow:Step 90, mean_score: -0.031250\n",
+      "INFO:tensorflow:Step 95, mean_score: 0.937500\n",
+      "INFO:tensorflow:Step 100, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 105, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 110, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 115, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 120, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 125, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 130, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 135, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 140, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 145, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 150, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 155, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 160, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 165, mean_score: 0.906250\n",
+      "INFO:tensorflow:Step 170, mean_score: 0.906250\n",
+      "INFO:tensorflow:Step 175, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 180, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 185, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 190, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 195, mean_score: 0.921875\n",
+      "INFO:tensorflow:Step 200, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 205, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 210, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 215, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 220, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 225, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 230, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 235, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 240, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 245, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 250, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 255, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 260, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 265, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 270, mean_score: 1.890625\n",
+      "INFO:tensorflow:Step 275, mean_score: 2.875000\n",
+      "INFO:tensorflow:Step 280, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 285, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 290, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 295, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 300, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 305, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 310, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 315, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 320, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 325, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 330, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 335, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 340, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 345, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 350, mean_score: 2.890625\n",
+      "INFO:tensorflow:Step 355, mean_score: 2.906250\n",
+      "INFO:tensorflow:Step 360, mean_score: 2.906250\n",
+      "INFO:tensorflow:Step 365, mean_score: 2.906250\n",
+      "INFO:tensorflow:Step 370, mean_score: 2.906250\n",
+      "INFO:tensorflow:Step 375, mean_score: 2.921875\n",
+      "INFO:tensorflow:Step 380, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 385, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 390, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 395, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 400, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 405, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 410, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 415, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 420, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 425, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 430, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 435, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 440, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 445, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 450, mean_score: 3.890625\n",
+      "INFO:tensorflow:Step 455, mean_score: 4.875000\n",
+      "INFO:tensorflow:Step 460, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 465, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 470, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 475, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 480, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 485, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 490, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 495, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 500, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 505, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 510, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 515, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 520, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 525, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 530, mean_score: 4.890625\n",
+      "INFO:tensorflow:Step 535, mean_score: 4.906250\n",
+      "INFO:tensorflow:Step 540, mean_score: 4.906250\n",
+      "INFO:tensorflow:Step 545, mean_score: 4.906250\n",
+      "INFO:tensorflow:Step 550, mean_score: 4.906250\n",
+      "INFO:tensorflow:Step 555, mean_score: 4.921875\n",
+      "INFO:tensorflow:Step 560, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 565, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 570, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 575, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 580, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 585, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 590, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 595, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 600, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 605, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 610, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 615, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 620, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 625, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 630, mean_score: 5.890625\n",
+      "INFO:tensorflow:Step 635, mean_score: 6.875000\n",
+      "INFO:tensorflow:Step 640, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 645, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 650, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 655, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 660, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 665, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 670, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 675, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 680, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 685, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 690, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 695, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 700, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 705, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 710, mean_score: 6.890625\n",
+      "INFO:tensorflow:Step 715, mean_score: 6.906250\n",
+      "INFO:tensorflow:Step 720, mean_score: 6.906250\n",
+      "INFO:tensorflow:Step 725, mean_score: 6.906250\n",
+      "INFO:tensorflow:Step 730, mean_score: 6.906250\n",
+      "INFO:tensorflow:Step 735, mean_score: 6.921875\n",
+      "INFO:tensorflow:Step 740, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 745, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 750, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 755, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 760, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 765, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 770, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 775, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 780, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 785, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 790, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 795, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 800, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 805, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 810, mean_score: 7.890625\n",
+      "INFO:tensorflow:Step 815, mean_score: 8.875000\n",
+      "INFO:tensorflow:Step 820, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 825, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 830, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 835, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 840, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 845, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 850, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 855, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 860, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 865, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 870, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 875, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 880, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 885, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 890, mean_score: 8.890625\n",
+      "INFO:tensorflow:Step 895, mean_score: 8.906250\n",
+      "INFO:tensorflow:Step 900, mean_score: 8.906250\n",
+      "INFO:tensorflow:Step 905, mean_score: 8.906250\n",
+      "INFO:tensorflow:Step 910, mean_score: 8.906250\n",
+      "INFO:tensorflow:Step 915, mean_score: 8.921875\n",
+      "INFO:tensorflow:Step 920, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 925, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 930, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 935, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 940, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 945, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 950, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 955, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 960, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 965, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 970, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 975, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 980, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 985, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 990, mean_score: 9.890625\n",
+      "INFO:tensorflow:Step 995, mean_score: 10.875000\n",
+      "INFO:tensorflow:Step 1000, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1005, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1010, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1015, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1020, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1025, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1030, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1035, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1040, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1045, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1050, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1055, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1060, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1065, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1070, mean_score: 10.890625\n",
+      "INFO:tensorflow:Step 1075, mean_score: 10.906250\n",
+      "INFO:tensorflow:Step 1080, mean_score: 10.906250\n",
+      "INFO:tensorflow:Step 1085, mean_score: 10.906250\n",
+      "INFO:tensorflow:Step 1090, mean_score: 10.906250\n",
+      "INFO:tensorflow:Step 1095, mean_score: 10.921875\n",
+      "INFO:tensorflow:Step 1100, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1105, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1110, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1115, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1120, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1125, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1130, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1135, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1140, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1145, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1150, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1155, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1160, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1165, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1170, mean_score: 11.890625\n",
+      "INFO:tensorflow:Step 1175, mean_score: 12.875000\n",
+      "INFO:tensorflow:Step 1180, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1185, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1190, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1195, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1200, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1205, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1210, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1215, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1220, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1225, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1230, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1235, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1240, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1245, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1250, mean_score: 12.890625\n",
+      "INFO:tensorflow:Step 1255, mean_score: 12.906250\n",
+      "INFO:tensorflow:Step 1260, mean_score: 12.906250\n",
+      "INFO:tensorflow:Step 1265, mean_score: 12.906250\n",
+      "INFO:tensorflow:Step 1270, mean_score: 12.906250\n",
+      "INFO:tensorflow:Step 1275, mean_score: 12.921875\n",
+      "INFO:tensorflow:Step 1280, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1285, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1290, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1295, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1300, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1305, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1310, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1315, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1320, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1325, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1330, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1335, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1340, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1345, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1350, mean_score: 13.890625\n",
+      "INFO:tensorflow:Step 1355, mean_score: 14.875000\n",
+      "INFO:tensorflow:Step 1360, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1365, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1370, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1375, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1380, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1385, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1390, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1395, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1400, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1405, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1410, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1415, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1420, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1425, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1430, mean_score: 14.890625\n",
+      "INFO:tensorflow:Step 1435, mean_score: 14.906250\n",
+      "INFO:tensorflow:Step 1440, mean_score: 14.906250\n",
+      "INFO:tensorflow:Step 1445, mean_score: 14.906250\n",
+      "INFO:tensorflow:Step 1450, mean_score: 14.906250\n",
+      "INFO:tensorflow:Step 1455, mean_score: 14.921875\n",
+      "INFO:tensorflow:Step 1460, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1465, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1470, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1475, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1480, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1485, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1490, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1495, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1500, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1505, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1510, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1515, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1520, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1525, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1530, mean_score: 15.890625\n",
+      "INFO:tensorflow:Step 1535, mean_score: 16.875000\n",
+      "INFO:tensorflow:Step 1540, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1545, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1550, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1555, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1560, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1565, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1570, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1575, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1580, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1585, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1590, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1595, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1600, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1605, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1610, mean_score: 16.890625\n",
+      "INFO:tensorflow:Step 1615, mean_score: 16.906250\n",
+      "INFO:tensorflow:Step 1620, mean_score: 16.906250\n",
+      "INFO:tensorflow:Step 1625, mean_score: 16.906250\n",
+      "INFO:tensorflow:Step 1630, mean_score: 16.906250\n",
+      "INFO:tensorflow:Step 1635, mean_score: 16.921875\n",
+      "INFO:tensorflow:Step 1640, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1645, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1650, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1655, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1660, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1665, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1670, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1675, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1680, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1685, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1690, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1695, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1700, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1705, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1710, mean_score: 17.890625\n",
+      "INFO:tensorflow:Step 1715, mean_score: 18.875000\n",
+      "INFO:tensorflow:Step 1720, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1725, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1730, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1735, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1740, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1745, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1750, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1755, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1760, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1765, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1770, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1775, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1780, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1785, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1790, mean_score: 18.890625\n",
+      "INFO:tensorflow:Step 1795, mean_score: 18.906250\n",
+      "INFO:tensorflow:Step 1800, mean_score: 18.906250\n",
+      "INFO:tensorflow:Step 1805, mean_score: 18.906250\n",
+      "INFO:tensorflow:Step 1810, mean_score: 18.906250\n",
+      "INFO:tensorflow:Step 1815, mean_score: 18.921875\n",
+      "INFO:tensorflow:Step 1820, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1825, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1830, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1835, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1840, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1845, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1850, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1855, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1860, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1865, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1870, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1875, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1880, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1885, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1890, mean_score: 19.890625\n",
+      "INFO:tensorflow:Step 1895, mean_score: 19.906250\n",
+      "INFO:tensorflow:Step 1900, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1905, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1910, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1915, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1920, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1925, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1930, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1935, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1940, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1945, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1950, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1955, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1960, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1965, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1970, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1975, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1980, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1985, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1990, mean_score: 19.921875\n",
+      "INFO:tensorflow:Step 1995, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2000, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2005, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2010, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2015, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2020, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2025, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2030, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2035, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2040, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2045, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2050, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2055, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2060, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2065, mean_score: 19.937500\n",
+      "INFO:tensorflow:Step 2070, mean_score: 19.937500\n"
+     ]
+    }
+   ],
+   "source": [
+    "game = 'pong'\n",
+    "run_dir = get_run_dir(game, 1)\n",
+    "!python -m tensor2tensor.rl.evaluator \\\n",
+    "  --loop_hparams_set=rlmb_long_stochastic_discrete \\\n",
+    "  --loop_hparams=game=$game,eval_max_num_noops=8,eval_sampling_temps=[0.5] \\\n",
+    "  --policy_dir=$run_dir/policy \\\n",
+    "  --eval_metrics_dir=pong_pretrained \\\n",
+    "  --debug_video_path=pong_pretrained \\\n",
+    "  --num_debug_videos=4"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "WKWPdwP8BW_v",
+    "colab_type": "text"
+   },
+   "source": [
+    "The above command will run a single evaluation setting to get the results fast. We usually run a grid of different settings (sampling temperatures and whether to do initial no-ops). To do that, remove `eval_max_num_noops=8,eval_sampling_temps=[0.5]` from the command. You can override the evaluation settings:\n",
+    "\n",
+    "```\n",
+    "  --loop_hparams=game=pong,eval_max_num_noops=0,eval_sampling_temps=[0.0]\n",
+    " ```\n",
+    " \n",
+    " The evaluator generates videos from the environment:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "id": "At9LC5rxFyv2",
+    "colab_type": "code",
+    "outputId": "983b0e7a-2700-4e4a-d776-03c459669770",
+    "executionInfo": {
+     "status": "ok",
+     "timestamp": 1.553253830168E12,
+     "user_tz": -60.0,
+     "elapsed": 4036.0,
+     "user": {
+      "displayName": "Piotr Kozakowski",
+      "photoUrl": "",
+      "userId": "01014928596539690143"
+     }
+    },
+    "colab": {
+     "resources": {
+      "http://localhost:8080/nbextensions/vid.mp4": {
+       "data": "AAAAIGZ0eXBpc29tAAACAGlzb21pc28yYXZjMW1wNDEAAAAIZnJlZQACqh1tZGF0AAACrgYF//+q3EXpvebZSLeWLNgg2SPu73gyNjQgLSBjb3JlIDE1MiByMjg1NCBlOWE1OTAzIC0gSC4yNjQvTVBFRy00IEFWQyBjb2RlYyAtIENvcHlsZWZ0IDIwMDMtMjAxNyAtIGh0dHA6Ly93d3cudmlkZW9sYW4ub3JnL3gyNjQuaHRtbCAtIG9wdGlvbnM6IGNhYmFjPTEgcmVmPTMgZGVibG9jaz0xOjA6MCBhbmFseXNlPTB4MzoweDExMyBtZT1oZXggc3VibWU9NyBwc3k9MSBwc3lfcmQ9MS4wMDowLjAwIG1peGVkX3JlZj0xIG1lX3JhbmdlPTE2IGNocm9tYV9tZT0xIHRyZWxsaXM9MSA4eDhkY3Q9MSBjcW09MCBkZWFkem9uZT0yMSwxMSBmYXN0X3Bza2lwPTEgY2hyb21hX3FwX29mZnNldD0tMiB0aHJlYWRzPTMgbG9va2FoZWFkX3RocmVhZHM9MSBzbGljZWRfdGhyZWFkcz0wIG5yPTAgZGVjaW1hdGU9MSBpbnRlcmxhY2VkPTAgYmx1cmF5X2NvbXBhdD0wIGNvbnN0cmFpbmVkX2ludHJhPTAgYmZyYW1lcz0zIGJfcHlyYW1pZD0yIGJfYWRhcHQ9MSBiX2JpYXM9MCBkaXJlY3Q9MSB3ZWlnaHRiPTEgb3Blbl9nb3A9MCB3ZWlnaHRwPTIga2V5aW50PTI1MCBrZXlpbnRfbWluPTEwIHNjZW5lY3V0PTQwIGludHJhX3JlZnJlc2g9MCByY19sb29rYWhlYWQ9NDAgcmM9Y3JmIG1idHJlZT0xIGNyZj0yMy4wIHFjb21wPTAuNjAgcXBtaW49MCBxcG1heD02OSBxcHN0ZXA9NCBpcF9yYXRpbz0xLjQwIGFxPTE6MS4wMACAAAADkWWIhABvrNdXNvEPmO7lwVl73sPl0EDBzzvrz1O9Sgfa49FGnVhGNj4PrUzIEjAsiR14q5boH034au6fMfeHzW8BQIdLu5D8GWFcvhnUQvMLIDm/5fDlJWNI1pLZ0KekyKgRZvEg10IZZePvLcj64kGJzCMbJi6QZbX4WMzyM/ZwsXoWWPBmmlKBzFixHWdkptjcAYhpgDpXSILlIffpBFr5Fmv8Xdrl5eZtB/U18q6RE0tX2BrhekKyOZ5lJWnXZWIEICkLYIda8x0l/aAug9zkJAN2UJ5v8AfQgXgS7iPy41I11UQneH59QQ6r2Fy+bXVz7hKXvFUUQUW2NfwyAHSubAtKRV8FgrIBnKXwxAjc8zc/00LsdZVdehIaL1eI9qZtyap5GmVpF7ZJdkQbo7j2k9/o8Ztr6lwZrODqoujHSJK6V9bK0u9Et564zU+wWgftergJVAEl4m/D3N6/lD6Tni/a6bLzIcdcVjnfWLPAUBwAoj19NpxhAbe1VyiybbzF11k65OpExrnTpeyfXnWi2YKXmv6NMcvP6YS8WOK4pM7nWhyKetjJvO69p10oeh7Pv3PuQBq3kARIBKQ+MPYmymnbhgmxG/6w3hJ2A2Urz2k1DctVq7TiUCWnAReHSDqSpYcdQwxCm/lIpIwtl/dffgss5v+hhFs6NSNe3zqLc+wa/P6fKKBzHBPA6mZtXbiaJH0Y+5hMHtf92lFc+6I4pZ1q2XpI5Nr1V7em9lfehnp6KwZCFUTCrCle3ZgVn3/WlL0hiX3HqF/qGx1rSRBE7lqG2nGEQXx7BFJGNLF0vFVi0j2agV+lVqGOVlIxAjK3E9wGWVM0V7xAFGXQtxAYJ6qA6zMuM1AzlTqoEWcy+zkYm6Z2/Vn8RMHtpHaCW7GF05Wujcn0D05dR11MQem50GDiKxlzGighyGKWmfeex/qNBXelV1apol1nwDCUiSbC9fPUu70YI94kit+OnCdHe588u+9o5tqmvG+4ju2D1U0YtGzBJLwNtKIxTj+ycim3c7lWMz9gyNpdcRw85nQOO+UebN2j6KDuTy4XNxidtzFIcvo67EYGfl+q3WaPfQzFQLuOqvybvDRViyMxNwUidf7UCNcjzUMa0RFtd4HPTD4pR9pL0oOHG0XMOwvMlfvI/0tUl0nH8gKjxa10D+0pCAG9Sq76K3xNRb2QQ4PhDp7u3P+U7CY7JpR9qHasfUEAAAAxQZojbEb/+8M7YXuyASh7Kplen5y7UyO14JIrbok4XTbRQe3ORruR41lvOoDou8ClgAAAAB1BnkF4m/+AoGMvgzynj7iP1aLLDUqgN6Jo/suYQQAAACABnmJqRH+I4ID6PupU5REY3sf/aa003+ohQ9ie2tSAgAAAAbhBmmZJqEFomUwI3/3q7CUD4uvUkflAlpSidF3UDZUIJZsuBftR/Ot+q3GhwTn9egJDu/Q0u302gbAsB/IV/AFE9llDfK6lUW9694v3+SwMl6mllWP/WwEuWwr3bvYMZvmSan6TfqkNjfFQ2gIUyqpE9/WshNO04YsB0gSlJqYcVRMCqNuW9LOcW33er/NiW98OL0en+UGHVMigjfMRwtvQwEBx2TCEvqNHiHUZK47Ql9EBEwOC+9RX+RfTR+dz/gd9jggR1VC/W7S6VbwMD6OLZQ6pAbUGRTytAu/D3nUIHJpw2KDyU/GpGdKbj2WQLEoSWl/G1erWEOohpWvfDxHkGcHJyES1MZ1DvjVLlpUz1LfdDNETW9u5oytpQuaaon1mmF0VFggCZuK84NA/jjnYLjbiGa2mQVNuCz0xJcY1SlKU0HHZWh4xyJUcesClYPnBpMFzn1F68BZn1odfSkUsfnPO8Ozfwiiz7f3xZq0mYnuxahAz2t7bEpiZEBhl9vOvB5IwNQnQ2WMn3D1n/XqqbGvjTwCrUHeBDng8+lhZDYcGCdhZ+2y21qFiMlJaX159s9jZ9Nr1pQAAAFpBnoRFESzfiNYKja/0xDhfTAiTpHLhMdYxhLBGYN2yGAaE7v0DPQ2YrW+QLfMv2VvTlemA0fw0ITn/7H4QKf0Smi2r8j9M+wtBuP4ooX+ZG3O079gxDawGHzkAAAAeAZ6lakR/gQeVglqse6icagSNwuz60J4uGyZCePmbAAAASUGaqkmoQWyZTAjf6DCt6Yuhred9EHrqUOfXsyntdBaWSJrN7aPiRCVp6QCYQRjgAPaX306Ulc6wvo26P1UQNC4biU/od187SikAAAAxQZ7IRRUsn36T35XM0R9+0te5g6nnf972YKyQEsNo7mvZhGpibk4k6GYi5Xw1oII/gAAAACwBnud0RH+KcQxiHI/+tYvdUXZoto0LM0yuhiZoPS8TyCql7aDErRX1MckO7AAAACABnulqRH+Gdbgx/3UyLHgLBkL+GlWi7oO1l81SbcV5BQAAAC1BmutJqEFsmUwI3/QsGvLB67IEGcjaf9jQmmp2xwadCZqIc+QsmxjXnMZQAp4AAACeQZsNSeEKUmUwUVLG//1hfa7H0q8GiMhKKARrmw6Z4xdeyr7ujJme+WFMGuspP/+ISBWO3g4rVjHKMZczL9N457VGygcsdj7HL9XSQg7I9GyDcuezylwimiV/wFo/tnsP85jKo1N9u1U/c32ZRoJ6MLzOYkUW9bgpdQwc/vwurbPO6DqdSp7R4DyF3mA4Y/IOgh69Sz6dfNBC5f60MoAAAAAUAZ8sakR/iOA7CrYE++z83mxaT9kAAACCQZsxSeEOiZTAjf/7wc1PkbwRgr+fjnyNR8pfsXICQZIWSD29LO9OdoFeaCGC7kNQcI6QVUvVoNv//1awF0I8N0tfCIdJ96/qdtaHfFDnSTMhoQ9GjaS9EN6FwfTiM0n2rikHxnG2wBAjL4oKVFB1J9q4Q6oi+C8FkQ6nZ9weS8YziQAAADVBn09FFTyffW4h7bT+Vmnsg7Jwp4y5XEPjbVoz423RdhZiyhHoVMXw4WOwpIvW8VtedZMRgQAAACsBn250RH+HsPlaZzJ171H1c9c+P/yRU8UOm2VwDFLU0JopWl1F2Khtda0wAAAAHgGfcGpEf4TUK0IOqyu7XwKdBz1mDC7mcZOyimAiwAAAAGlBm3JJqEFomUwI3/wDwzHAex2BjGQbp8PmN+QZU96uDVeJsc29OEXjcldhiXFvZCjL1XZadbCRD86yiAgWQRsivW2zgAoC7DVCL9e3fSnRGKpVU/LGLClh3Y48sH47PWqbcpIEZ6IQMXUAAAByQZuUSeEKUmUwURLG//vD9/y1OQEdZZUsKUwNqfrKLfwKdmKP4MYlUfWBfMS6AB2YzPOOewKgSMfLoK01LvNA1j94yZ+2K1utCc4bPCb5fgtJtphyhALG+EJzcxQG8TmO93UCKf8P7LJwQsMlijgw4NSUAAAAMQGfs2pEf4etAFroO/Tre2AdB5O/z7sBdNC+yMviqgLALp+SqJXA7Pg/DPRDeTGlt4AAAABvQZu1SeEOiZTAjf/9dtYVwOWMQEFcGvjV4Syzg11kPr1s9125wYmt9rvuf3mx7RYJ6jF6ErfUzdz0/3v9ae/O7utr4LCBaSp+/oto7O3F1qWD/jeXpJo3T76eNJFHYJw8VuPijBUU3Vp8h5phb4yxAAAAO0Gb1knhDyZTAjf//XzJzOIH3VrgDqkUNgn86gOXkEPCMlG+9fgG6imSEzHDCVblECfv6lWOPhYa83SgAAAAbEGb90nhDyZTAjf/++FR2EKSC1to8pc6YSO3kyM61vFVTIqwBRLCgvhTW2uELbv4P1t6j3j6R5JEnH93R+6ccuNOfX+S8DB7HTWlWJcQLgqD+S2bSR1NJYKOXWfT3yjStlWid2o9fNbVb/WnQQAAAIZBmhlJ4Q8mUwURPG/74DjgdzkBXMYP2UA+265sg2RvtJWp7hbeXljZyPg5U1jpyoQAMkswdLvArqw2ReEXhiCOCghoFR9+KlVVnzlHwdrRAGNlgaJUCgaXSP04WBiefPhnJH+UNq4LFoPepHhK//kKOrVcBFaR4anoX/aA5EFLlBZC3z8HwQAAAC4BnjhqRH+HsRNxvly6e2k2NMUXC/vU7UPeiHPyDV2p/YvXZWhgk+q5Cpb6bHPQAAAASkGaOknhDyZTAjf/++A44Ilb6dcZ4EfRuTLTuDvo3J2im59tzroCzqfsO8XO0To/kC1I69bSu6OBA9I+wRfSH4s2YCwMyLJB4SXTAAAAXkGaXEnhDyZTBRE8b/wDoisgQ2L7kEBMp1YsiwFGNSG/OTV6DNluS1vSFNmiEavm3sOuILR8pgK4M6P/qHxZC276LRYUYYCavgfWkcLXP29c5JUTuaP0/h/7ytcJzS8AAAAwAZ57akR/h7Ee1u4Mf7i3wEEg6o2q6d802TgnxIfuhG+eE+Dw81v3A/7HzNqgsMZrAAAAZkGaf0nhDyZTAjf/++s+OA7zprt9FalqM30xdIjlcH2DZkkFtbEKP2KnwPRDHCudTxGL9Qj/p+X8zQUT6DsKw9cio1ehc/Va1ccw4608TJVh/tU58mts0w5trYhfnxrz7ghBbWrPgQAAAB9Bnp1FETzffU2jm3iBOi008YuPeRt0exkohagoWh6AAAAANwGevmpEf4k5No92pOx4d3tVe/jt/92QHqWDJ1gLa+IuwyWHPUwqtoP8fV9Ym+scICO5EURRBtAAAABkQZqhSahBaJlMFPG/++FR2CGqQlUqIWgPGeTAl0jkR8nWtez7chaAXKIsKU2quN90XktERKg37/m9ip4Dz5kOgzQscz6nyNY0f/sd6sj1HySlAcKWdh7zp6SDANGHziaLu7KO4QAAADABnsBqRH+JNh+VxVwUJ+Y3lV+rya5I/8TCxwR0ZLGKwtu7rocrr7JN3QZSGNke1uAAAABQQZrDSeEKUmUwUsb/++A44DvMfbctqjDdsKqNf0xOA39UYyLAzYIZyyzAIzc51s3jIS5ItuDvEtTXgkuVVUxzdDw7I8ZpdZavGnbqYY9+GhMAAAA1AZ7iakR/gQ9rhvIRB0ySHXw58101sjT3zRJJs0q1N75B5kzDs336j3kR4+OINWFbqx8fObwAAABPQZrkSeEOiZTAjf/74VHYIlYiKMdyQZmvfSW5HcZkgM1RFiY3gEH71z/3fVBzOwcraE90El70+ZDjx4HGc66IARVSqHcYq0qBrrRk30FtqQAAAFNBmwVJ4Q8mUwI3//vD7l3aeLA9vGddw8aOZ4u/aB1tWqMJJSgQk2EKksQOal+1aR5lX56hJg2jFlZ0Xy5wPHFC5sT+XuWXvNeiHLLg872OhCGzwQAAAGBBmyZJ4Q8mUwI3//vDobVTkBINCYpxqf92jigZ0Srp218/WFaj2TRlEHOqssGMM0Apc19xhhprfp/o+nyoo3nU4qIPx0toevbOZ/uMdat8h/rz3l6mAf6cx4JPPMV/58EAAABZQZtHSeEPJlMCN//7xIHZpGQEF4WlOSmv5sDvQ1HkPJ24b8mLANR7I1VevWJebie2fuXBDZeaCivuJcWRnz+sLxLjPYolWMX+OCeAR284M5BlQ+hDK45MhLkAAABfQZtpSeEPJlMFETxv++A44EJXr11WQvEfDTQVOZV8/+Xr3S8f8Sdkd5yhFoZ0Eykj3Xivg1BzZSewsPXa9nQ6gqdDgf7T2DmCjKrk7f2e2vK9ypwXpi7kgz9A0TgFjRgAAAA3AZ+IakR/h7ESVb5jShGX8VmBJn/MdoNQ3Z6fh4j/jmgTG8KfR4gLxPvHCfsFkI5Dojx4hvmkmAAAAEFBm4pJ4Q8mUwI3//vgOOAkGyUenL98r33EsuAShImg94uLBrTk6dhiDVHpVEky+DhSk5kCbedWMCnU+cpV31ziJwAAAHdBm6tJ4Q8mUwI3//vBcoj4CbBCQT2RExz0aNVXfJZgEMRjG13pvbjo3lZ6TLZP87/4osMgfYBX228XW1pfmU8k2l4pXWkF2ONrCPtLUsjWqi9NqYjz9uxWY1lh/hx5fARLNZPbPZztgBCOOKpjSsjqrRTvm2kDmwAAAGlBm8xJ4Q8mUwI3//vgPUQfdNRvnx3FP1pywn3LJO+BgkxSFcCfgS8xAc5jl7DjT9LzYEwD9isIU1yUsdzHqy/6ngxDd8Uf0UrLVqXY4UPbCNliO8n/ZegcRbOArrd0G+CXl9/2CFy7Fp0AAABdQZvtSeEPJlMCN//74D1EIHvgowka8CSMlOV1djkcwdrmN+M3w0ez2EF3Z8QgY9h2qKskrDR5Mk2+uEQHAHNJFzVoZ8OEnB6N5bVjkoQH9GCmMuS0GmqX+iknVHfRAAAAYEGaDknhDyZTAjf/++A44EPraJILnY0LuTUTVqG8Y4R+iod8qU50V48CYGzxUBWrE5/inR2Nxzjn9BK+lYDY6wizKJTKtNBOivIiofjj7puT2VHhc5kdvPo7rDnnm0xNBwAAAGpBmi9J4Q8mUwI3//vBnCJ2LA71GF25ddxrRcaDbZ6JrFKdXAbBcwenVA97XR386r26bMwGidvp2PmbPeLJeVTrLs40YO/0dkENxamY3tNtYwvFfYvB8Y8OeKM6xrpDUHsU4NOoFoGqtzCdAAAAYEGaUEnhDyZTAjf/+8N7/LIyAjCC10eRAkC4payNwDB4TaAJfH4zC+l57EAHdxnDWH/9FYP1pFGTKBZX4dIkEoguP5WochPIdwRBalFTzEmR/DYBfxJGM8ekovk7ApnUkAAAAHVBmnNJ4Q8mUwI3//wDwzHAd5j8RhmR9H0z98nDjYj/Nc7vsZGPARdvUDB2TCWaeccqi54nDBz6EstQtzy6v+wvloWKraz+uz/M18zKVlwhNHeSR4GZPnFXhk/6SFiPs3DdL9fCsVu/U6phJsgQpW9nziiIrsAAAAAlQZ6RRRE8331No5t4gTzZDbeKeBR95P1RisqLON/iq6lMAjgQ8QAAAEEBnrJqRH+HsRJVgDz377tnoNimjhTZvPY9yHf2fmv0wsOMF+wp88Zp9A47UCY9LcsoRhIaJ8uNlzZCrsZsy0xHSwAAAEdBmrRJqEFomUwI3/vgHICtV4/N+3TLBHjt+Fcn7DWjz8soUTNhnfkqkcMHOJu59OaFxhuVf+ti1/1ybHdnNcNjs9ZKiS4KwAAAAGJBmtVJ4QpSZTAjf/vB/wweEBHWjwtlVbql9dgODZ2kZPlG+D4KXaVCP4hDGqGFAR8WEBDCUpQahknTJjX5ZfkvoG4DfOZ2jNfwPVMAdpL7FdEsgo7WPCSMB38k1yCOW8g6+QAAAE1BmvZJ4Q6JlMCN//wRIX9RB98K5rsT6v40vwfGBFHd36t3Lmg6acxZXjPDoIP1nuZtPbJR67n3MUNs/jpJGr2SW9ZnVls5Jb7a9l/19AAAAFdBmxdJ4Q8mUwI3//vgR+AQ8tE14oytZoCQH1vmaiP6zviuiN8WiV39XBnBmeNGUBPj1lZw3uGVd1ASBgSvQit8rSzdGRcZWThyfaYyQbzG3uhl9FJ+f80AAABOQZs4SeEPJlMCN//74DjgOFPapEipwNQ5JLf/iFkT7vTC27zO0O1Hhr7Be9x5uGF5a9SD1TJTW/IfLAT7Nh7DJMwi4QswAkdulgB0QpQxAAAAd0GbWUnhDyZTAjf/+8Gbu08WB7ipWtNjMHVF/iPxLBqJR3dWf3VXwht118WdJ2yr8sybv5tw7Gkcen/6wXcv6hm1FN3LQa1BDmxg9YZtMWlhqjvMsaNo3uWx/3Zsn/9CB07I7ahNSdEvwIcfvP9pOIVFVaVu8QFYAAAAWkGbeknhDyZTAjf/+8Ohm1TlEJHbNU/mWRKFgYzLN2yQ8jXAB9wCzOsfbVP2lU1EY8QOgEzKqXo0R7iSjvfThP3Ejb6DhTio3axVyuxNGvXKyq/GmvvY+WeEeQAAAINBm51J4Q8mUwI3//vD7xujgdzj8Rhi1Ykq7OuqSWTi3Jyq4sFrVH+uGrfgBuxTFBDV13U4LCn/MgN+f4ZzOHXoTdRaVuXBhLEbCtA5Ms+hD9d2LkEvLXaBv0E8NnWBPx9w4ZTvBVStx4l9Na6nPeTrSSDBnaVQkpIMXYq0+x5SehxmeAAAACdBn7tFETzffVS1iQH6ubsRADaTWu5ionu4K0/0ak9zTnZggGqi+fkAAABTAZ/cakR/inEMZAKk66wt/cwHaBMWpfAc500hIeLekQWHoYVeIuwzwo3AD4zjCnVWjbJVYnBRhCSWWZm8rpySFJ3514IPVb1tZpdA21QwKGClE4UAAAB9QZvfSahBaJlMFPG/+8IGeHNgeeQ43E7kwPs1VGXGJw973890r2GIuyVyjsOqfdfCvVaSkH16R0M1/zHTQJQoPuOg2+H7r68wE56bXlYqz1zmB//ygylNGdR6NOHJuRI4gUP87HXbU2cwVemECH+lErFeByiJwCmM07PdY7YAAAArAZ/+akR/iTYflbJv66XME0GHkCF6mtJpaggCYpsevkqPWehQsYwUCZ42gAAAAHxBm+NJ4QpSZTAjf/vBm7tPFgQmOx8vn1I4MEdYOyJbxW9+T8JyO3vjlruuDeJ94kHVp6mE4InbW3mhydReOAPRHSvfloMOy1X4f1k61/T5AyQdQwSfNs/ceWh+FXJYPI5vBAAfACpz0P00Y0ZQiTyXgYR5dHA4282fpEnhAAAATEGeAUU0TJ99biHttP5WbMcugQMX/6p30Ja38txT+7rJCjWU5BS5Hu/Ped4RgDO/8NmgX1/xtKU+guus6I+GB0iyn+7zzesJ2+Dd1FQAAAA7AZ4gdER/h7D5W775k5qly41vVWL+9TyyQd/+taVv6vysDLrv4GuSfz03ASJUej48yG8+gYk5sPihZ4kAAAAvAZ4iakR/hNQrQyxiUa04Hw7yezOPlcbNtcK39Adv5vlNi03I0KSWCx97lgf8KiAAAABOQZokSahBaJlMCN/7w/fqde8Ef0LFVa6z8obTV+Z3fHanJstZCK+C9VlS5uYoIGfyFRCQyvYs8eXMoinX6wkePCQbIrdvJZtx24a1z2OBAAAASEGaRUnhClJlMCN/+8HPtQiIJq4KJFcEZqp5eZHjbfXvRNExlEoNtkHV/O1j3XsPkfmAogRZ6H3xb4fGyDO/psnXuk/zjrKbgQAAAGFBmmZJ4Q6JlMCN//vEgWa8cB7g7HnWVr5afQMBhPt/eF0XZINzgf1DTqt4LgNJ9n7jJeXS6jzqED9EpR64Rz+FTHrYHSWny0wQxBJBhnSb/iEP/JdZ86h/Wi9v9QW09VhrAAAAQEGah0nhDyZTAjf/++s+OAznGqGizRSqo3BVcpAWaqCB8ovSoNU/Q+cLNp60lGoY1XxqNFsDbFsRe3APbjIzVIEAAAA7QZqoSeEPJlMCN//74DjgJBsj6l+X0srb4l6b6JH2Ft3abZgPDnb7Plzz2pLEPDmuaBtebuDUX6hFZoAAAABcQZrJSeEPJlMCN//7x1Dlo+AmwPQ/8NXi9GKqmuY1UPDYZG3NPKYYZuf5RLSPKtiYcXru35hiQqjNXqPD4oWVOHCsaOXNK65RabsS8NeFVBf/1FIziqztdYeHaIAAAABfQZrrSeEPJlMFETxv++A44DvMfnbEhBgUx9vnA4Ho8ktSkKylpuzQvYY0CtUp4uNU/8qrDXOOYbe12K7QloChfe0qoEDZAp+Ls75A+yk9tDX6iPXAaWpTzsfeqD+dbeEAAAA4AZ8KakR/h7ETcb5ctlmh7dxoqP/yRU8UOqHP2aaGG/KPrV+Afm67tENSRXPgCr2j+g5e2szXEdIAAAA2QZsMSeEPJlMCN//74ByArVvp1VBzo6FR5mbIURGBYtSZMz2ctINnHT3rH2VjPX/zzdapOx0wAAAAW0GbLUnhDyZTAjf/+8PuXdp4sCBqPO04eytS9D46/Avaazen9hU5pU11Hmv3PBlT0PTYG549QlWOOVwWlJAXcvc1syzyvf1qFX3EhgXM1IOIcGnTwUuYkLr42jEAAABaQZtOSeEPJlMCN//7w3J3dNWQEg0JK5pzBQVfpWoOKLSI1/VmJjcsCrBle+ncG6tiC7wuDEZ5USHcZC55aKT2OIPW9YHg1aVeYKYR0wQWQDP1DG/oXsGKUwexAAAAV0Gbb0nhDyZTAjf/+8HPtP3gfFlpFrQPirkvcmukg1+owUzumItdD2e8mcIUJrYgE5Sa+gOebQM6C3UJSm7qj2KyQP6H+uiQQf6dQ3nUsv4RzQXEB9IZsQAAAHpBm5BJ4Q8mUwI3//vEgWa8cB37/p9ZF86iN25nKIigWeKhgv7YJ/UNEZWwthTtyyAJEfsOe3Z46qB8Y9u4aNguQI82wq+DpdZlQKDejU+B0T7yh/vWZYzvEjdR8mEIAxxdRZpU68SFTtrnTu7KpqA68eg72PoiJCihwAAAAEtBm7FJ4Q8mUwI3//wC0LvBNTmF1Z9IfU+yXng9Xlh5jOvxWurF0OzyKrrbAG/LhzXREODrp9/66YlkypRbzROY6MEMzZYnzbptXwMAAAAzQZvSSeEPJlMCN//74ByBNNkfODOhV2l7PuVP6jv9DFZM/xRiaZ7NjolC2zFnBs1tnrqTAAAAUUGb80nhDyZTAjf/+8dnvoPCAfKn6R60RAOrjtGLSYzJf83HxmDNtKNcZR3s7Q17gEYux4ptcFO6KrQhBRjlRY6a9irUPBzxRSHCk6HKHHiDGAAAAGlBmhRJ4Q8mUwI3//v+jz7KI4HFxHpTS0oxWxOKguK69Iqb7Pv80RE8zRIcDeLBSrbSsht1y6Ct5KnKYedMgtXkRlKks+2+B7FYsVvxecCZHOfXyHzislTlSVnRDtTX3Xm3aOyRaBXu5sYAAABAQZo1SeEPJlMCN//7wZmMlyhHAV+DQxbHCYKLdsYP8alRhrAXEGFKASYmbvEaFUC8jJbbP3+VwoeATImNqxZxmQAAADBBmlZJ4Q8mUwI3//vgHIE1raJILnXvo8sZurnfn67164Fu2/h+6+ydNTr/+OxuokgAAABKQZp3SeEPJlMCN//7w+7Y4o+oh8mWdvl4UGi7/t84ZAPcQ9+MBMDQq//WhRWSosr28nZ9S+NamaT8AaHjCriTA80H12cPmqsMBr0AAABZQZqYSeEPJlMCN//76z9RD7r4RWowH8hXR1qM2/0RPeiKZ9pQxjXugxDv5Y/OlrDSWjfi+WFQhmWEG3gYNg2K8SmfzQY7iAamgvUBbTcVy2n5t1ecNgemmcEAAABoQZq7SeEPJlMCN//77lcOB3OQuw6HD+AjOfELwUGp9SJicXV1sFnhvaKjnvCxiZwl2jex2l9TUZcaQ8RCwsTwL7Cb3QWKztEEttb0Jo9sjMWDnWrh5iANpRQyrCjmn+qlJDlANP179jgAAAAeQZ7ZRRE8331No5t4g3EwRcCB6QcQmRnRJJhqS7+1AAAAQwGe+mpEf4exElW+Y0rm0PJ/wbTD/lW+CtIwcWOFlp4l4dVOEXYZReMT7JXZAnpNCK1hAhXjScp5NYYnhBjzjGD/rSAAAABpQZr9SahBaJlMFPG/++FR2ESusZKMxurmUPsaoB3yvXwtgjhKcOw7RRICF1EOtbOnMwE+WuGfg86L92P86V4YhcXg6X7+adBswWrcPNmvGF/+VM2MJnXpcNhI8G84ECcACh3CpkqclfZfAAAAJAGfHGpEf4jgOwqrh+lSqMFWNRECeaFV48q85DBbzjXglEAaLwAAAR5BmwFJ4QpSZTAjf9P1UW2P31QTlnuH9bjbYJ7XWt/ai9e9KF19PysQ7DabFT0nEsV0yceUNAJX0BBUVv6dG9+bf4cKP6gx1WQaTSkLrl6RjZl8G8g+mu0jnqW9tdKrF6RdsgVAMdnm94bRz+PzA54Zyuy3ATHLl2cUlXEyCW9KigPQ0J5U3y/k52ZP7lyfuZ5la/1Q9Pj6hGjOXZEUFpJACy7fj69cHAm5TQWn8MCi9bYtYpHNqiSnSLn4A+l5LrMldm0cXmrlCIKI+1dB0+ItJmzaYll6630NuKKUYK9we5JspVhzVBRfoTKIYbqUfFWokuJzg/aLs0nAnaPalKUUmOI8qX3c1g70zXzdIptzRWo0lBpoXPimAkfYfsLAAAAAUkGfP0U0TJ/Tuf49tp/KzFPCIEDEvSN1rnYJbiqzzxVfJ/8OCLsf7ljUGroVfs9mei4RJMtv7Zcn6zTeStOI46ruvNTH1vZzMeIyae9K1E96hHUAAAA5AZ9edER/h7D5Wmcyc1S5cFCuXC/vU/dYydSXmT4ray0A6VWt90PLvJGDUf11AKA6FxyAGoDLyBHBAAAAMwGfQGpEf9fFlYaEHVZXMWGGCx/WUqHrS0vhwUrligXTGXJS4DE/QUk+ncgeHLqWgy+vzgAAAKVBm0NJqEFomUwU8b/oMJ46ACuOTXqhV0ljwLZShz4Tc58bkcnq9w3uogc6YHf0kLpTfoclGUCaCZ4AxrrKOqJsYXwsspdKo6DAKGat7GSKdIulJbXbhuc/IvIa1CF/DNWzV6r+cO1H8weRmlyh8KFBzm3Sb9biHpOZEdIVGO3TyFoUJ3SD0ZzyH7EROwGbMEzO1FkpRP4jj2zcTUYldf8k10XzYXEAAAA1AZ9iakR/mI64NNsipLH3uhfv8LrdLfjRclwIPczo2kTYfFSqf6TNBKKbBeuQKHqSSWNQZ4QAAAA4QZtlSeEKUmUwUsb/++A44EJXr12b4QrXF3nTB+IiC36W7K9k0PInXSBIMLwHooA0xe2aJ27xzqEAAAArAZ+EakR/h7ESVb5jShGX8VlCE/7Y/nQZhvcjseRoWw3awq8RtUipbXM1BQAAALhBm4dJ4Q6JlMFExv/9Y3sZGUqwJ/ZJQ3AZgBXt9Ufjt4ZzNxGyAfjakJd+hT55Tw9cNhg+w6ug3vz/oXT+HSXuNuSH3FG5/LjmCDKvQYF1YEcEbLkXk9fh9BECrWP0tPknccp1bt+CPAa7hxxvktQiTyB05TND7a4HjRFbvq5TqgaLwDS28Cxz/2MBBMq7yg+pQVD0bX2Sl/mRKrv1ISvuJ2v1VJxyBQ+KKcEXAsgkK1RpCQS4jATBAAAAJgGfpmpEf5BbOR6mhN42Utt0sKpEJIAR7fw4EeQJlSVwz1K/GoHlAAAAeEGbqknhDyZTAjf/++A44Hc4+26pb+I0qiyyaTY86qeqgj4kwIFKO32Zd3pJfJsEyCt3X/2gIau4QMjcRp2plA5yiGVX5wEKuxlMS7/gGuv6Qw92srH94+7Mg9EV6V3w0yihnRIaFYSL1USPSa2cDMI6VI+grv22GAAAADtBn8hFETzffU2pysalekPthsm2ZMTNyS424hytUB9PxsDoKgpOtQiLMZ0iPUVzASv533olMcsSGUtWzgAAACQBn+lqRH+BB5WG0Yk9tcQHTHiL9VohRhZ/SCKeLiSmHZaZBYEAAAC8QZvuSahBaJlMCN/HqeNQwz8YC6lddTUaq/wlNA30DyObwUgGgN9P+XYVKGiTwxMURC6nsd69EfWbXorW5UKp/iLDB0tXqGqCdJCKds+YuzihvwR1B4oMcWI3hPmbK4zhnIvnkOawDk1lzJMAFAd5Xs6In7X2yeUfKr5WJHmhlN0mNbPW0T7jDw7LEkRVdNs79QlUlWSBMCb7P0EyD2FPtivHzwZbF/QBKslG6u1N02WmHh1SzyYHC6SYhYAAAABSQZ4MRREsn812AW0R7+Pctap0S/62v/Y0HOJuqNfDpqTpeIBGPqgtfdqY1GtZkmBpYYfpCvY9UDU+DGLJFacH14P9elmz+ffaFFAfxY7tFIA+wgAAAD0Bnit0RH+ezJaXm0/z6MpMxrufrdPmfUHceEr1rbvBdYpSb3s5MV2ylnH+g7e3nIEQSIuWIYQBSs9KfJ9JAAAAKQGeLWpEf4Z1uDH/duS/8VkmZodKFFLza+EKeodh2cw31hLmuImTU+zBAAAANkGaMEmoQWyZTBRMb/QsGt4s8uUfHR6wH3Eruy7jdPNSHdGx4H2+eY45bHHjrvu+npGPfuKs4QAAACABnk9qRH+BB5WDztJXfGFrhk6Xb6COb2LMvQK/NgrJwAAAAI5BmlFJ4QpSZTAjf/0fwO8d+rKuAd3nWKoNCPGmENKEldonVkGjg98w4Me5Z48NellC0FxCEd/kdNLe1cM8GE783LvVTMHmMX1wOrayRMiu60nJZ7AhcL3MR8AOnKU7oGDehMw21vsadVsuyATMJtrcH76ypGC8FIQF3X4VlmnhpX1WTqIkRvuvbCddYWNgAAAAakGacknhDomUwI3/+/6N/UQ+1Ea1+CN+lmcvTOPgijtDvzZhN0D2G8CBJq0xATOAf5U01N9dCRuCxIu7R2vxMfg848kgrRbFR6x23LTcr2silj6PlhF7bXGnsSksL0sP7ItpzFz6DxDAfIEAAABQQZqTSeEPJlMCN//74D1EObFx0Ccz500Sa/EklQzLOUat9ZefscXQjlR3Mx55ToX4Xuw4yk5jOe10fAnmAMrV30VKTa2W5vE12YQot/5YC3AAAABHQZq0SeEPJlMCN//74DjgiVvp1xjN3F4nu5NzvpHqsfivvp/bbY0P4LsFxB1hB6kH8V7RfuzmdrHI2vT7+H8cVEjDPvhujc8AAABdQZrYSeEPJlMCN//7wcOsd4IGgTeySKCPIuwxg0YYo+96hJ2UPP5xSEcCoL3qXVqKFs/+7whBdM8r0BuQQeLuwVm53f5UyL9EOx7Lh69/Ofouy30Ld83TsPD9W9AxAAAAREGe9kURPJ+A6iW+6tW/uR9lhCZLeuQ4oZcBFeKAu5UN4mG1qmIlUSZUB+VEhmQryk+9/puz2UHFvGA+dUoS1cGBKKWIAAAAOgGfFXREf4egEy+iMAaiy7zEWbVUnq4AyqydOp/rPmZg3eGRT1NmWJ3z0oQRnaNxgE548Z1QPX7xmOUAAAA1AZ8XakR/hnW4Mf928rmxp4HnnBMUqhe0zvX0oeNpgN+7Uu+KKS3fCuekbQu5jS3lzKvhqNkAAABAQZsZSahBaJlMCN/77mtCBIDC4q6hbKW6i+2ykwSQy3Fgc2ZOrZmrvUZMVzT54+uJ2QGpl/6fQueAl+NdAyBz9AAAADRBmzpJ4QpSZTAjf/vgHIE02R84M7FU/ejsVHqM66DEeEz+EEL17DltFUqXJViZAeKAf2cfAAAAZUGbW0nhDomUwI3/++FR2EKSC1to+4xCmEjt5LncA0mqNBUks+Ym0IghLbBCqyHI+/o4VwUizGvOHaK62P8WHO1D41LvGG5uRQbBRxa2T+sO0fWvHLEeFz60KXZuVfnunSGseJ/OAAAAfEGbf0nhDyZTAjf/+8Gbu08WB7eMlZakaObAXDazwj3bsg0VKVyXEZdPl4Ml5p/16cJP/U78Iz8EJFZ2maeMrllGWqaR/jYpXDX36keDWXpidqMEWfOXkQ8fZImQl1OdRErD09YOHHaf5NGC+6km9YfOUDqGX9Pq3XQRj3EAAABXQZ+dRRE8n31uIfah+Vvv5ElY/3/K6wW7DB17bbOdxcR0rSz4be7jHQy0EgS/lLK0iaQUrBIQrvTwd5mZPmZ4ApMYBTCkxmGGY4ysfzByNNQiAJD3piF9AAAAKQGfvHREf4Bw4zB2SpmzpJNdP//8f6dskHoit65Hrw6c76E2VmXkD2UoAAAAKAGfvmpEf4TUK0MsYlS18DDnFz7Y2UwUKL7z8XXLeoLGQWrrq0RBevoAAABiQZuhSahBaJlMFPG/++5XVEH0KHETUXGSGsXNMJtH/jNeRzgaZvNQtjBuinxU/liYq+MxIz8WoClfRA/KTLqF4/Fa0SAqiv0o9OwWGtPhEpGIZo1GD+Cou/3P4e/6cRklt4EAAABAAZ/AakR/h7ESVYBQhM1ypHXwCXiUkjPjSDADDv8Q18gyx9OBwNx7koSQDcnOOQ671+kC98Jm6w05OE/cunBeYAAAAHNBm8VJ4QpSZTAjf/vhUdgh9eh1QZyogZ7+hmtlbPNslgd7qIrlLoeRIb/et4+8Vq1ePF+NNLB4Nr5y2LIdg8c54s2iT1WYUe5vxH2qfxAHd6oe3EPUU6bAiHT77lV4nOsi77nk7apWUyTxHGqurceFxTOBAAAALkGf40U0TJ9/cH7YcKBioTwOqZoa6GOmGKX7krY2+ZZJFD1Q548KXkilnIlw490AAAA1AZ4CdER/h7D4ZJjNasvMvYzkHey/9PvUWs7fgtuPBRXLAX7CfyYyOMWD0kVUTTq8sgZLsG0AAAAkAZ4EakR/hNQrQmu507TYOsLT3zhWb32toH/fgbuu1qCFxEeBAAAAVkGaBkmoQWiZTAjf/XbliF2xAJMQm7dNMvTYZ4BH/3IqFg9bdiX1lecwbkkdjr1LW6pM2UN1kfniZvPkJBXmPFV59ScvpJEqbeREB1KHz3B7jm6BYferAAAAWEGaJ0nhClJlMCN//APDMcB3mPztiP3uapuxK9w1XpJGf1ey0bVBH2OVz7uNe1T/11fYO9Iv7Fc7kkj5DXP4lwVGQc+vNvTeM4IHYwYcrs9kxs6CASSVaQ0AAABKQZpISeEOiZTAjf/74DjgJDW0SKRZBRnHc1aiJ3lEevKh7RRg1AT8zcCut9u3ECa53MXXjAVR/5TQqy9iMtme5BHk09PtAdSIfrAAAABbQZppSeEPJlMCN//7wZu7TxYEJkiWNfakae+YsHrw4VAfICk7Nb7oLk+0i+Mpuy4D8f/TgFMT1z8D+3aoAm/Am00CBytVnpVYqPkD+waLnY1NOQR2k51+vgGAegAAAExBmopJ4Q8mUwI3//vDoZtU5QkPFvbm4haJgZvT6vbhkh4T4F3xSdoAUekhgoFrQuyi5WLE6GUiQEppgDoxQGgzx7AhRZZEztPG3fTBAAAAQ0Gaq0nhDyZTAjf/++5XDgO8x+Iwxa46cT58hPz/3xXGJ1VivV3KaFJ/3WQxqtPg3LKAamhPU0NAdWGoS0oO7WTR/EAAAABfQZrNSeEPJlMFETxv++A44EJXr4MQrxHw02XeFmca+qmRNOoUcRRGPhSgHkAcuC9vwQHiT3krg6bT/BkQSItsA8CwlnnzN24kW98J7AEeP0NiNW7ZKbpdvkMpFyUz8l4AAAAyAZ7sakR/h7ESVb5ji8/rDDaLJg/7XqiPJ1B0qKo+hBZfe7U4Rdhl3xjVAHcpz+vB/FEAAAAvQZruSeEPJlMCN//74ByArMYDm/bz8wxn23X7eeY8ueYRJcmDuVxZyy1hXLUX1ucAAABgQZsPSeEPJlMCN//7x0/RSwgInmw7ggx6hfLg8Mo63KinbZ7POUyq9lLfsMaS3xKvOrwHbYeQNzz2RoRgXOv/exp6/Mde6Dums6ZRED4RzeaCX4PJ9AgJuwbmO9tXNIbBAAAAZUGbMEnhDyZTAjf/+/6N+OB6zo5Fy5eHpx0/tejGu520bLkOKYZmUYwoEKtu0gsWwF6ztlWdsn2cBXYI99l6dfY0me7vv4CJ7ndd4PxDQYUsSz0tIM6CYnR8oOrc/Gfexffa12GAAAAATkGbUUnhDyZTAjf//CnvmQIc2OrRA8MwDA9HMkp0e+MI7Pm6epVIxMw29SeVpA1ToAxIhYkHTNN5ShviDJ3BuOFRgp1C9Ah/PPMQYFk/kwAAADJBm3JJ4Q8mUwI3//vgRKhKrrGIpjJo14wE5fxn+jkJ7L0hg0Y/aAfD07llO5gOpjeyTQAAAEtBm5ZJ4Q8mUwI3//vDootTkBHWutfyL3CvbJ+vrkeH295y995Ase9ZpuUtjRd/dTZ08jGdVP/9jByukXhfEKbwAjssjjm1tYa2KIQAAABAQZ+0RRE8n4DqJb7q1BrisxfO905/2QQCdm2BOraG6khWyZ2qOCWjta0B7hKGZCu7bgEpU4a1B1yA6YtIK+nvcgAAADEBn9N0RH+HsQIbL4x/P9Mv+bVwIYmppI+PEO0zwOR0FYU9+OSfjcQXZbl+lOJ/yuuBAAAALwGf1WpEf4Z1uDH/duR8Z1p/stCcUr053SmodhihH153Yzn0yYfGEtEcJcDx5XQwAAAAR0Gb2EmoQWiZTBTxv/vuVw4DvMj6PtjWXbeCDWe8HbgEv/NcdLhFqVKB8exZhz3jCX6jioMxFUtp3r/RC8ltZJfrteOt8L2BAAAAIQGf92pEf4EHlYPO0DFLXPMyZmip/4spisSOYUIUhCPWMQAAAIxBm/lJ4QpSZTAjf/vBGT5YQIc2HKxMsdoJ483hDJQtEt1oSVY1OtNsacpcCHC1hLiyide7+igKkUYmgW63zjDEC973kPfv/kGvM9IQJRrdPvgu8NP7Gg4FonN9g1w21bXVxN2U+b+Kba5ydK3Fwptozj887wZ5U43AxRVa5ir6jhSItTpA4qmjsWn/cAAAAJVBmh1J4Q6JlMCN//vBzU+RvBF1xHc1oj+U/TZywQssct1DlyEQP3bm66D0erMGNf+QJB5DFmPKyxwSETeXY47tb9OQ+xlPePfb355uJNYeS2Og2DJDc6qxeP+J/6dc/Sv/p/LXg678yyOLMGYC22x0VzI6myko3Y2jbnwPZYHoVN3W9bTc2x7RxOC/79m7UesteT5niwAAAEpBnjtFETyffW4h7bT+QggNBgC1Tx90dvwgqWi/2bEAaoy9k8G5OWvd6LsgFQzqIuBbNMM8rzbsW0pRv4t/6wnIDQ/mCBNdzbJDDAAAAEoBnlp0RH+AcYN6mjmSXWqx+E1SBd6filHrC4NN/gK7ATVZjyRHQcmyIQAZumDsG240peDLOLW3gVbq4LmG3Njs3hJcfGYGdq1drwAAADUBnlxqRH+E1CtCDqsru19uEmwmchClU8Q9+ckZEbxoQUVCtJ39bW/wYXdNpPtmAsJQybL9wQAAAHlBml9JqEFomUwU8b/8A8MxwQo/6EiaK37P7zGFrPazCtkNNv4vIp0kq5kSnVxolGnW0VIejbpP5ZiiC+I/F0ma3+IF7CpxnIMPf4egyyXNCkjhnJHSXHusxYcAUqAOlV5uTf2eoJk7NXqf/lcR5JWs+Sf/4CuRrynQAAAASgGefmpEf4ERwYTMPCynKnF/ULnVV0i5kdCM7P0vSFSGSpVSXT+9+3FrB1//zAGZ86aGphUfLcz/y3mFRVu6iiyShLlf1IBxT4rwAAAAk0GaY0nhClJlMCN/+8HrVh2DajAEnVrgIluIofrb5gjHi7I0o+eSqEr/KXkOiwTxu2NN5bbjgK7wNu3j1ZZsWnH+OrtoEaJsgcV6iH5Pd7Uko8sgwhTol2e1ZdJ/icre4wKyTmNr5bcC5di+bdY2nLvrzZtdpa9DX7QJestuRmj2U10Cl1As2ZueLxIRqj2wGBHESQAAADpBnoFFNEyff3B+2HCdri0/UUVYHd8epT2d+YMYObZm8nW/ez46SIN8EbMlhsLlCSFaW4sFX2nF2P34AAAANAGeoHREf4ew+GS+7UvFV/w2hI9vC/vTUrzGx1QcGgUQ8rwi7DMOqBjdSkbuDSjnaADQL3kAAAA0AZ6iakR/h417ZNpNh7L591MHWZo7LvvoIcpQmohn+Xd70kfnW1Ud5bf8T7beBDmR8Xcd8AAAAGhBmqZJqEFomUwI38oIQ86YuCDOQ5tOdExAsQJusAUoBl6PXbngL0EHOm/XgoQK4cIAx16UWM0dQcvANQqKp2EBOaWAoRwWN1bqOg1bQrKKmV/ijRbo+vvcxfowunt49/emVGo9MQAXvQAAAE1BnsRFESzfg/tuGYHKjTB5ZID1kuLiL2dqNk66JpJk8w1d6Kcc2wUWTZM13JvgxIJn74EkYNkgC4PjHLlyOy0Hwh6m6YSunCOq8UCBEQAAACcBnuVqRH+BB5XFIVzyClcf37dfxipiNXufl10iycx4d2lwwnYOlmEAAACBQZrqSahBbJlMCN/7xIFnX2UJY3B5G9rwJVPFBhNUhdhx/fLZ3fI40N7ZzzVF/KfTN++yvpUue1sMAod2GLm7S5Mnr+HDpXolz9xb/2QF7HgW6iFPZ0TSu72RZcJPzQzsgubn35XTFLL3mozUEl03Plt0iKzniOced0OUW4K+kw+BAAAAXUGfCEUVLJ99PbzhStPV0DPWCp739xrVtKTUsEdCecaEM+e0OjfkkUkuNJFophvGmqzcNDPP+G6mO1enkonn7obx3Z30VXOOnNy2O0Eiwd48Dx3VZnk6OQouxGUAgAAAADUBnyd0RH+HoC9/9Y3zHp1BIl2MfTFJGRIlfr2DnVE0Sw1ikOxYlKnFwH2Ru+ncOjZ2P3qIlgAAADEBnylqRH+H5CPxj/uvJW0xJiOenPhWEISI76CVOVnP0kO3twxfh5J19iCz++FpQtdBAAAAZ0GbLEmoQWyZTBRMb/wDwzHA7nHfcLHdl3CBcWDNUIHw5qJ/uI/Rxc+UScIeicTquE85DXU/XssWAvri4jZAdY1hh7+q6h1EQeWGPd8v9kVk45hJzDq4QFLRn0armvMglMXFFZrGuXAAAABFAZ9LakR/gQeVg87ywhqQ0UZ0H9SsmVmcAjAxEPJNkzxuBCGvKgW3OfU6j6sPxhv4LrsmE68KdHBdFoxyuLi785Fg1b6+AAAAhEGbTUnhClJlMCN/+8di5MOwazCfUdt4QXA7i2vzFU3g/WQsndT77JVDSX3dyiHIJ+dcHuYgNieBhDWi40XytcT7MzqGHQ8BKd7emcKr7t9ZMy2246DgCMU6O8nfWZamMp0hGAd1hoDnPEpKtXES8lvdf8QvUpJi7H+RY6YmEvOOXCg14QAAAHRBm3FJ4Q6JlMCN//lIWPInJdToRM/Di8/wOIa7fAwqyFmQUIqzAB8UxzfJluX0+Td7Rj1n+911ubH3x5Ng8MZiV90GkBT1677/EI4xf49/XAJzyB4a0vGsKyC/rL2cdk9MBpe8qr+6ehvxBOlZU5I+9k8HgQAAAExBn49FETyffW4h7bT+VmzHLoEDF+QSPtyOibdtnOwtVf938/IoeR8/CfdhBqH8efnQNzSTcptOtFhLNhHsC9QRs6JUpZ3g6L+0o2hvAAAAPAGfrnREf4ew+VpnSNtXMyHIouF/ep2oe9EQ/UgNEJmOXn3+zo/llDpYdbcSubbJNj3j3tm/n+QkREY4VgAAADQBn7BqRH+E1CtQOK56tKoCiM8K0R3dr/Xh4xtb5DeCzfJ3ISlqMnAuma7Jff3Osd2p1kQ4AAAAeUGbs0moQWiZTBTxv/wDwzHBCj/nXShLUbtqalbXUXB4c+Jr8UsTF9vlJYy5JAsd6N3YqAwRTAVVRhIyCYy8OFra9SCQJIdvyrOlBtx6e4Nm3jz/vVxUbjh4mewKCUl3melRvDxX4XCovgeHkgRgAO0Q4EkZqLJtHfEAAABFAZ/SakR/h7DkTgDV56NbnZl/ml39qIOiEFfzBef+RWvHUibHEjD4zje6WjYo61QyA138/8GbcEZEmbUXjgjCWt9nPALYAAAApUGb10nhClJlMCN/+8HrVh2DgEzkPQSYjrWyzR/cxxfpkyJIH7uTdsyWu2DHpet18vbjP31VvB4OvU2zGwNJhH03pIxozyt5CjT0DxM8Eu6iU+hynjaNStmVahBNpSsriT7NNCrjVDXNYImUZaYR5CNKrD2nPtelij1qj2HUukGEN0p8HTXl4uhVEfUSYArW7VPcWydU5vlLZYaJeCj6iL7TK9654AAAAGFBn/VFNEyff3COo0/prpbjGuq+Yz9qQ/0UtxBbR2wSbc63al5dG7nbMBmTmMJALX+gAAHd4U3sZUwDf/j8MBK/Z+midrtvZ/dhX7BGy5kbGSKhcRoIJqsRYrNy2G/mf35bAAAARwGeFHREf4ew+GSYyDT9FKwHNGCE3zTvKrR2YTzumc5w1YU5tpz8yyBsGigtzFTb4gkcmmGk8BosAJh8sW1hPLsi/oP9XLQpAAAAKgGeFmpEf4eNe2TaahDYDlAE00B5TiGe5P3cqwgz7Jz7dAwop6UVlxL/wQAAAMZBmhtJqEFomUwI3/vBm7tPFg7wyzAD0Mnjaxg9o1oowtbJ4Iq6fMSBsDEVXAUv3Ek0V6bUA9AyXfdXplJl5OklAZrAT2DLE0CTU2Gb1ph7uMvZ6e730j1uTDGRcia83XP9lLWy0N2lgd3miSshfYKbuLgLsYsiwlykYvR7QV0d0CqGOqo55SN1Jjj3z4lpX6c0RGgzwd3ziVIoWQeFFwVeqqCbllBnRSSRaTmdGVG2XFLb5K4/MLixbmd/kI/o/AQS9ak96IEAAABVQZ45RREsn31uIe20/lZsxy6BAxfkEj7c3vTu2znWnO0o+LL2TewvhwqHeJD01QF78BcCbCuwQBKbqpl90L6uNZIPK2wJ2IYROSzDobv1Y58U1sN/QAAAAEABnlh0RH+HsPlaZzJL659u40VH/5IqeKHVDn7np90TGQtQWPbiEvzRpaGCnUR9bwO4lkoO9UNv4cqr2uoaHtjRAAAAOQGeWmpEf4TUK0IOqyRzDxju6fds63t8C0diXyWCw0nZtf4XS2WBIN2TqV3HMYRmD9ycxi04LTVA2AAAAIBBml1JqEFsmUwUTG/8A8MxwlzTJ87d34cixfMcHbv+VHu06EPCplmBHBDuntSkT2yRgALAJie3XxuB3xfOH/D7ydP8ErkN7FeqXmM3xkd7kLShI3yZeVnO9ZFHHuwsNecvxwV8jY12J7GCNBCKjcvRWa0RiLGeyRhv7w7uCpPSTQAAADwBnnxqRH+J9+qocEqY3mgUADG/qVAbloOL+4LbHt1Vueid0r6SqxaL3S0xMXyghgv+Ji1+lINFTyTiXpEAAAB/QZphSeEKUmUwI3/74VHYelCJVnhRTg1BLYiZcXT5dxyk3eX+rJNoNiJEdunbKmh+C5I7HMdbrb16UYYiG+jet6KLxaN1w2WtwbVe3TezkTPnAeJ+tInJzjZlV1FqpcNW2arifXJoFuPI/Iz2s42jmNCDbRbNxQpODftQ9RldQAAAADxBnp9FNEyfeyCnmveYuSODs+6cbniKe/vM4S56nhz+m9i8Zex7ONz5LSvfzcVpUjqdCJWt7Tx6czK0loAAAAA/AZ6+dER/h7D4ZJjHyl5mjfaYPRhVd6h/FiZQAFr6JH2QqvbB/RwwwlywmDQA/o5AxedWALtCD568UWghNBmBAAAAMQGeoGpEf4k2H5XFXAtLnyu/P7BuNQSHZxysH6gKDLB09GXBu1f/byom4T1PU9j8d7AAAABwQZqlSahBaJlMCN/7wZu7TxYPY30W6JoFSNQpwVvW1QrXI4toGmRdXr6ug5AcIzYxb3sVhpDt2R6YmUdAxjB6i/FCQrVb4p6cx70BUXqsfihn6ICLLcl0nvY1vzENxIRzj0j24J7cCRqf3p3r+wKgQQAAAE5BnsNFESyffW4h9qH5WbMXYFj/ael8uqKOT4Kdx343HScFjvKFNiDbjO2olivLVLdWGTyFi/Nn+rmW8yKVjA5v2wbZq/jG80u3A30934AAAAA4AZ7idER/h7D5WmcyS39QLC0VH/5IqeKHVEQEr8EriQ0rMLKqwWs2RvSTdcH/r+SXrvRjzQL2kGkAAAAvAZ7kakR/hNQrVcoxKlr42ol5Yisw4FbJimCD3s7LQHkiALNax3evqAJ2Z8nqxtEAAAEEQZrnSahBbJlMFExv1WRxLgVrlSU5kvRPQefQNL+hnY+5sxyM4OGP3W6gipxl+VeQcVtIs+PdwqJefHjTv2e+aZM41f823N7uJYxEIRmdHjyj4j8dVr5J5jvbLxXLbyf1aePF+es1ewhAgs3sr0LDYMZnAVaSTpMiiCjkwD85nRHwA0GmWfVDsiw1fCQ08yGxuwdcMd61SPpfRADYJpeVmWilLIkk9w7lCbACvWCJVPVxpYkJW8dIszr9QbdlVVkW+1uatXpBfwsbLty6mMFmnr5Fm2i6dwp+IR9w1Urt++4L+nMYz1LxgXCx3//Zxzgt2ANHgQF4NuaWK8l1q7TAxuxaM0EAAAA4AZ8GakR/3DTbiwrhhwSpjeaA68X7+pUBukHe3+XGyF7HhYlJIdhyOmcdef7GeSfs6o7iJKfivUUAAADJQZsJSeEKUmUwUsb/9BIkgoUyga5F+ZT3IdOhBODSowpHgJE6OOIA0hhuJSSkDCReDJDR0qEUS4FpJoQHvLAlVaAWTk//H9e/17IrpZwyJ6wCC9sEeRYyik1ullhgvUDdNUQzGaLfMevgU8An/KAX0X1/cOEMqpcK9aaJxowlmuUS21UtcIb9euUUnTFnUYwlhmrjdbVG+lNOTBID04MLbIMZIZzw6PX/9JpiKMQPPmV99BkjSWM89NuJhl5mlC5ED/6kt13QRfBgAAAAPwGfKGpEf5iZsaTq/dqTsanc/0IT3kf5Hr5m3oLGjyM1iBcL6c2qPe1QQjiri0diK6w7972kh1OttrOxLezWwAAAAJRBmytJ4Q6JlMFExv/0LBqptRAkFpHKdI1ZlaFSI2+VAhWAnIutcz+4N6enTxH30oYaYrWXeNE7ifjDCoHMMyW2gFUY+wbi/JE8adAKifO+RhxPHcvsbcbFsnmwVPaFpbwyxJYTd5QCBNnA5eacY9cU6ZwSL/KPzlL1bcR5waT+lbJg8F9k0Ygkey54qcWn8L2WvH0xAAAAIAGfSmpEf5BbOR6mhN42UttyvF0PAGvTXN8LUtVDAEJ9AAAAqkGbTknhDyZTAjf/x6oDJtkwPh8BHMKIJQeKakhBaoY1uTU1Dv+mLEdoR9c+RexgX+HTtr1kNcl3RElgObLDY/PaRJqf4jDnBLlQj1v/j+gBpH7ME2CB67cRZ0xo1T0OMscgDL2ja/1P3y/3R5yr255srgQtnSzKhY+hp0JxNKYyYJi1IFRA564tmvre91OlcSUCU/OIFZJE4oY32tHuv8tWU01eFm8j3ElGAAAAPkGfbEURPN/QgoY6wn8yMqFM8ScEy9Q41T0IvLskW7ZAQd6JL5LdsFFkc8iVrAOSANPjytFHsQEpJgFmIQehAAAAMQGfjWpEf4EHlYbRiT21xAdL3k+Vm4+8bP5+kpCxYbPyDxTyetngNkKsEo/OcSIoI20AAACGQZuSSahBaJlMCN/0LBreLNnVwoJ7HokTqcp6Q9iGVOTjfWRdX3HNLzYq0LH11Wq8Jq+uiZjebL3X4L63Z+y9aEj8POW4Tzlvo4K7tfR97PaNJboUExwnpuB7qJVSNwKDoTb48nTn0tJrvw/O3EZNUilImcvTCGE126axIMdq8DiFXn5Rl4EAAABdQZ+wRREsn309vOFLdicCm2x3Pufzb+QqT3EAn8TFTgB+03dALeaGOw1/6I6cxZ1UST7oYT8jPMi37QdASPs6q4fKsSNxFICCpLX/nIOtEt/OgZpH2dK2wf0xARgbAAAAMwGfz3REf4exAhsvjH6163w40hUbVYqr0IJX9KQl/IjfEvL74S4AfJknBl2wEwSJxGK44AAAAC4Bn9FqRH+Gdbgx/3bkh9Wq+haFwumUOl4Nq407OonaoLPLjUaPLWQnkuQZ8bOFAAAAPkGb00moQWyZTAjf9BLL/c2UM1KQX+BxGGL7bb/E+C2iJSq8ZT/fki7Xj3xHv79/3OCR4eUyro+hPAO6kwpQAAAAbkGb9UnhClJlMFFSxv/0LBqns2dXlsN5umPuNMnDIBekdIzGY3PnWipO8EQ4VO29m+HWcI89dez+iwGudwY8934b5mTOutNkupDrTYAefwpBLHDS4f6nH0oDndONGfIYoajXM6v5ZsonVMn+J7bcAAAAHgGeFGpEf4jgOwq2BPvOKJdhwabiU0ees+grZ1YMIQAAAJtBmhdJ4Q6JlMFExv/BdFX04QlYNpjTO0cfnbEhBFxg9ac2tx4t2ngcyym5CpUChqKx99i1D5EkvHHFPX9/7KMXq7jI/q3wkBGpGR/RHgt9ep8OxX3IASAi+yKKBVIw9GzZv1fxmZ6HBMZD2MbTpq0Mi3EG2QC+iK1rC7QF8ooxwG2N1JEtqo2tAd03sU98LQ9M68dpUpC/xpxxgAAAADQBnjZqRH+BD2jlniCi2imSX6//59WeyNYZ7EThdu+jrAY1bVw+gxFikWceh3l4wiMCVEWBAAAAY0GaOEnhDyZTAjf//WKRs+/QRQ3PXPg7GyNgHRyY2qAL7E/g8+2wAxx57bb9N5LU5Lwf1jU218is1FOXmWQFA+ILNlI2cuW4VuFEAUs55tZrFJbPFNXTxKQFb6zgrOGYxN6gYwAAAF5BmlpJ4Q8mUwURPG/9YpGyYu6xTlLEz1V6fELYXJYCl0J6HV2jidgU1kHU1TaIy98emaghsv30jfu//giS4u7YXeCrhV7hICMr6GVBT7WEUf+L/NNHBp/YziKg+rPOAAAAOQGeeWpEf4exHtbuDH+8TXGnxL+KjE1xi3rYQfEyAC01BI7R1nn8TQtbK9qWu/fdySehP03udb3CCQAAAFNBmn5J4Q8mUwI3//1ikbQD/zwawOGki1yA8PO6m26UL2L/9mKgReK7+ebK/al/M7IPAICEACwHoICsV2NTBXt4+cDul/OajszW51KtkLVl+aMGkAAAAERBnpxFETyffT2/FQo96LwXOn2o93ozcOD5tpYSJ7cwIJok8j0UU8aRtemPmdp7YYIUiCycIRpOWrK31Ohdggd67CoH2QAAACkBnrt0RH+HoDzW2qN3QBJ/E8gMm82gikxdFp7fAWLcEhRiEwpvi8PEkQAAACMBnr1qRH+BB5WDztAwU5oZn2uFb+hZJTgCS3z3O/PoqGYuVgAAAGpBmr9JqEFomUwI3/1ikbIV0EXMsa3YekHPSpYuprfZeuR3WOSiKkC0kFZpj7xpSGhq9ZWjmxEgbDlVL++iKztpTP/7VY798ET282ga/R1srrqOxD5u/U/+0pygtO04b7giKBp6WzCweDafAAAAekGawEnhClJlMCN//WKRs+/d5TeJbh49Ilqhs9vkFuDy5tILYlZvTBUtgGhCajr5vBfprgFsCD23kwfPBCzu+fBE9vNoGt6/fUDXQT6zOaThfR6pRtMs5SAtMZaMwqag6Y2HrhAcrJKrd7zb+2xmYXa8/28eupRQ2l2pAAAAZkGa5EnhDomUwI3//WKRsip9r0Ll4vNW49yxGm0rPx6d/W4e0Trq4nQjnHlJpUyssaFXC6OWq6h/nBwV81fzGtv5FLjjRbokIbdqcOylodWrIY6RdCrUEz4tfkGqVbBtksCTQEmc4gAAADhBnwJFETyffT2ygMiGUO/inwbX3T9EEiT+fzwLws9sywg0V8+Bo+hNRgRMgu2Mkb323JGVsThPeQAAACQBnyF0RH+AbCrDaMSfDibxDReFdC1UvC+7xMpiCX1dYFy6YMAAAABTAZ8jakR/h6Avf9Ts921mTsNzZj1+OJVQ8DVfaupr23+v6Ud85kP8cxjIPbdrvuP+lE9QsgrlQXtOD6M87SDXTNYuQnTLolxoGd6Iw9QfFrz+rYEAAABHQZsoSahBaJlMCN/KCDhlVhQaAeDWA+zilE4F4JvN0n7qpgAwW3vUlRQKCzVZ2fDY8WBCzFmBKgdcnVnxRMksKfRYMOIlGMMAAABUQZ9GRREsn309tLhK5CuvbdzT5vjjshRdA7IpVn/x5vBwKg++SIs8yQwRWo0QTWUmAKd3c5qJygjJ4vikTp3AKpczzE0vrvu10mq+f//pJzyqB/lDAAAANgGfZXREf4MOb59GU02DO+CVL8CSCBWNmo9IbuMSHjgvPw2O5k2vA2ympmv//pZzZUNT8y61kQAAACgBn2dqRH+BB5WDJUsHZv77UUMhfzqZBSad9zQgrlg0O2kNk5DK5ilQAAAAaEGbaUmoQWyZTAjfygg4ZVYMKd+rrGfciPhIiq2gCFAqMZFI7DVM2ZD3XLjJGnQmMgOIvEEsS8t+XEZfwTy76m0Qap5ds9FsZSQ6iIFZ0GU9cLjijfT2W6sRfprNkuPjiKDz/dClDTiAAAAAYkGbiknhClJlMCN//WKRs+1FSeYueQO7h1DZ7e+xqDAxtzmGG2fToXnjatvz1Tn7mSkpRsA/hIvPhdRN3kl1bvHvJTMKB9UA9ykYobd6r7+FKw/F3VVri6aWYgE36UGcq4wRAAAAZkGbq0nhDomUwI3/ygg4ZVYNdbNtBrA1dyGL1aw7/5eOTzHjgYmlvgs0LekeTrnBmIqn/lVENJBHf1RN5RqsI2gVGmV4sa2fCQ5tezOGmg2pMMQ4mlHRXTuRSMJg746BbfQilUiiLAAAAD5Bm8xJ4Q8mUwI3//1ikbPts0l1vp1ED2x1ji6bVEURYnvYIR3zetSp5m1e6jKMQTp7+VkiqKd2tavPg/v5wAAAAFRBm+1J4Q8mUwI3//1ikbHOi3JXDnst5+9/Q8YozwOIgFT629Pl7FB8I6iF0Zc7IF96heR57AZnekg18MeZ8yh2uOHSqcUTsaug9Sj4I8zr7wmNr9kAAABtQZoOSeEPJlMCN//9YpG0OvyOUdYK1lQfyPABv19dN2EuVGroHzp2zUVU0NODsPiVBimmNmF/CSLBt4F+dGOjvRT6LiulIFp5sR3+G1h1ATu9T+HuPV1x5zNOLgqLSCXQN/UuvfZ64plbVF+9zQAAAGJBmi9J4Q8mUwI3//1ikbHSvGoxINYD3SLvvyxPZxlxKc1XKrT03i5NYjyxJoTxK6gyI9UVV9jQM7CeNG8TnLBSp36B8rOa2OVfe31qsGMYmSKR4dZKBGwbKDN7WtkaEsZZZwAAAGxBmlBJ4Q8mUwI3//1ikbOuOtG+ZLt7BBIpdazWzzFa2oIXuKvsSfKEX099+kQxskSmsj0JIeAt4NsniL/81T0XG32CsxykbWR//z7P3ayX2T+O0bkNhfiOKnIVANdHE3DTcxt7oZfS/97sH3AAAABJQZpxSeEPJlMCN//9YpGyKf88uaZQCcSFhy2E1I13cQ7T4qAfZ/efBEseeX43YrqDE31Oe51xgWu3dMtxJ7PVwhztNF4CSSDjwAAAADdBmpJJ4Q8mUwI3//1ikbPts0lsarG2hQ9dcXCqWONWEzXuGsn18ByM8MdGGWrqt78s6xW7qPrhAAAAaEGas0nhDyZTAjf//WKRshXQOHBpVAT0KbRHowXnCt1oBqGLr3USERo5s9MW+PmLvdT3UmgBF+tZO4KwNstRUb0WbAoQxjGZF5h/wf/HrDjvpjl+CJ1AbSxs9y9+9mJLpBNXqaw9rkBgAAAAUUGa1EnhDyZTAjf/ygg4ZVYVaN1iGuf33zgSf6omrlwSndrkCTLPF3gaG94Se1OTYoSOIPiUGjc389Qa3QtMLmdfxJ/9EJgaN8vT7Keo1Vi14AAAAEFBmvVJ4Q8mUwI3//1ikbJi7rEg1gbbls4vToej136f50b8WMzjf3gU8vi9SoIuAXfe8+Kxcl5VytPUCNZR+0dQwQAAADxBmxZJ4Q8mUwI3//1ikbPtHlRWB+qqDdxH6C2zc77qhJ1fZ9qy55n/wDkZ1sej4ZDv6vAKtYQSdyHy4zwAAABWQZs3SeEPJlMCN//9YpGxzowPspkvEL9KenZlAqwx2krQK3l2oCa0wsWc8k2NlSHrH+b26/w5V1g8f/pnNfDFAjYmFGxsaSxgtrssgtpmXxb0wraVhN8AAABqQZtYSeEPJlMCN//9YpGzq87hxUgQYaLWrw1yOFa3/w8/9YwOlVMZthDmUgBo65bZZnRaIfs7t/B7omgkl6q/wSNF+R+xHNbXAqgPYfIrf0t/ExK3QwGLCO/7HBF9dcZigKukbsI6FQbxgQAAADFBm3lJ4Q8mUwJvFGew97V1Lcd+O+QDMDFw9wWzOYcwFZxnnmFDt6/+J6XbkhV2LCDKAAAE62WIggAJ/74Mt5iA31diPnODkHlAm45CKmrkO6tYSOdbdSitE8v4aantb94gCmriEDoY7sxccjIMexW5zP9owVi0pMMzzPblktB+I1kmpKIM33vbcRLOHPv0LC3kz7y85a64aEb5O+r1+heOtihVeJUPASxV9SaWTlNdsSrJgj4xmmjiYMl0IjoWJ9oVs4HKkTYxl/l9OJO8f5MIUmsmAuHXs9jIbJdKBA47Isns19FTdkpBDNqpHKr7q8GFdm/zMWCdfXqX41NkD8CuzgCgpicALz3yxhbW+N0bt/hFjnzm/UaPrgZPcZzqDiV5eYefi4KrgKA0ZVe3VpNSqYxcKn0235Wuxf+zZBgzqMmKLS0lLynJtcUkdDEeIyyfK4LIm48C4xhv8CoyP8GUDK8uivbXBNn7CgW7AuII2yToKdSV23wBkjD62p+pn+OELClzpVps9i3NFELoEbszOb28KZqNUr1lS8Wduuc2YwDKwxXVCCIeAzfkKMHwScs52kSwYmRCd/d8allKuQ2V4hIlH9Jl47UyFemjMfV7+rS2jrnUFGhJvwePtyEVo9P2WdscKW/cbYN2YTUkX7ujH/5UC056sGo7QCtikXJOlN9RqsmVVLXhNH901MjEGXfzKbWH5NYOmyN1mZ5lK/o17zObpovEPe3A6lXdH5FNGkd396zj8jXwQhCpcisxza7aEuADlEolAttpaS0aItC5D/12ya/vkaXhuhqDeowf/qnrMJEgsm2F4woahQkdykwDfkajV8Xoyi3jBpM/z3SsBgY6HWn3rW0ngUlpH53febs2YoTMxFBG9OvOZZ4WIWOcIb01o0qnK+4aqXMILB6bnNAae4qEqWz9+Z1pCpWendWC2cgOI5T5MSpidy6h+FiqCWdrWfJa2hloCwt6h+wVJKuhO2EUZjv5Bfuy8bQtvt4Ewrd0GT6CUsin3tuzheiKHIY7c7GARYxWOVsu7Y2YBaVScU2CERHIxqjA2+m2S99h81bJehej0iGHFz5U4SYskqcTAlUyrY13Nasq23dV54zIXVBj2HdPP/FgHKntLtIpTJhLgMLXrBf/HllTDNcthyCVB+wy+lmVyo5EGwdFVisnyRsii7TpXq3YGN3aKn3/gz7LAcKnAw5oD+vPQthR3tO2J/mgouNkoZ2GTWA2usO+gX2YiAf0fCodcAUuh6p+bh+zP/ArJYNnup5wOWspOUA0IovBT2AJTCeoNMfOAYnEEkLLP5P14x78y2wW1gWBgatjq9keNhveNxm3OId6Z/+RTUg7V/jZPZKnQJNBEy/eCv4qmUjDL8YpXPHrFThVUwNS3NoSGhVS4JZ7bpB03I1nSKSTH4DG1Lf5C7pla5OMYG3Y6n76Qbbfl9EZ86iKbiUcdGryD/jMHV12JykpDiaDb5FHPDkHH/vDAgnyicavYZBKgrx+j9OSTYucK815/uPHKsmteJ4DuUy00hL+bTACJ+49So3QPADbTTVdYUxcp68LjjJa/yyomh0a/d6qcsV/aT8s30oHmWbqIBloL+J1PNY1MzOCCU1U5ZMaoTRdCZwnZlEC/HchwPUMHg2DDEvmKsv4PqNfVSZC/GfvnrurcViZysYuuCjXGkUlYV0fNLrTBtq6ykyQkJUrZqjQdqw6kil1NFQWrOaWMeGubzMX3Bsgc3u8GFFquetJAAAAX0GaImxG//vteAAg9MZ+S0i/PC0zlHX/QQ3L0D4FiZmIOjpIFCJKKu+YprzBENnmxMhdyCs5hHE8xaJE2hp/2GNIqeNg/xvscT4V/xoCp2fXsc1CK8Yb8IojHKubQGZ8AAAAJQGeQXkR/4EHlYPO0mTyth87XGb0QiDlsCs7SOJZk1y96Bg/H+UAAACeQZpFPCGTKYRv++FR2IiwU0RA8qA1ChJxSNLn6LehKVP3kDzlM1ADrYhw7wd4Ip7AQzseLJdG5Wk6xUPebxn9TmVKfObBTfkJNs3ock1GokDdOFS6NZD0ZXPz//Ud2gir3gj4Dx1v/MhlLWLDnShmxzoy3/7AQNWKIUFk7mJHH+gLLt/06WNqUe+p7gC2E0AaYTuW0oNS7l5TDBLgt+EAAAAnQZ5jalPN/4H4Vo+TODPKsmBYbLE3RR0oNHHTburHIqYChXMkXhLgAAAAQAGehGpEf4exE2/rMWYbvveZW0hf2uqeKAaPNwi4iL5LnlJJ6paVEaSZbn+/5GomK7pT5XP9xVnXWAQw85PnL4AAAACCQZqJSahBaJlMCN/76z44S5pk3knFEhyI1Ta+1wxYdmynPZlwpDMB5imzzm+brVdgj4Ve8W9egCohBqHE8VTW+aug2DoRpkyvvsOPUbzL+H8XIBGreeSAS6btFT8tlLl2VaVtAxevDlGvVQTnr78zMiiWOQniAGGOB8j3OPC1oxhUPwAAAFdBnqdFESyffT2/dTSVQOypzfxQ2+DUil4GWJIgvIkBLArS+ydK0VasRPrGHtnNMEBi0Ye06NmDV2cc8yAn4FapQmg/tHUSVZv8nmxd+pKZ4n49OY6GrcAAAABBAZ7GdER/hRLK4pCufOnSwAxieT2ESQKanmN9zyXjppN4TTZfPjYr52MwkENEajixR/ruD0u5TZi8/TESJpQaWOcAAABSAZ7IakR/h7DkMUOCGBJh+lP5pZ7hsk9ED80Ikh1LLF5Y06UPw2TtMFuWrNVoJtNJFPcLihlZ9xJ1Td3KPSex82src3Nfa1Z7l65jvwwWZi0CQAAAAK5Bms1JqEFsmUwI3/wZX/CB9oOI26TZh6KulR6t3agG/iMaBA6B/6qMTvqGVg4k0YebbNb+LKTM12MrCqqoUI8hGiD2iXPrHlOmpgdC3Qb+yp7n1ju75DCZDB9Rh9M8MBS44veosH5GPzQfZ3g85NWC/wgf9Sj/wt2ZHm7uia3RMjx0cInK0e6gmf2UIdsg++PvXMcBn0gM8Se6JG8DKFx436aem9R9/OoLsInKDNEAAAA+QZ7rRRUsn39wfthx56FiN3xIMTuB500p7P9azlWaxHm/XjqkrOgES0XmZvka4Ctq1g1y9N9V22HlsJnZM+EAAABFAZ8KdER/h7D4ZJjIOE/G/Yh5t2iTFoHhadHs+YUVSRGmyvCLsMp2wNI81ubiOdKh6cKDbpl4cUw6CZHuVE7caJVrVL2BAAAANAGfDGpEf4k2H5XFXAtLnyqNr3TnYj9FIKhGPIf0deRjBKSCA4espm33EiqcHf4SHokHvLcAAAByQZsQSahBbJlMCN/74DjhBrqxUlYTYumv3ra8FUgL+AgXjYUzvRyU+e6HvKXgk8xL/pusoXU2ZW4/lqsHsMHeb8YKaReZyD2xnfEzDGr/au71JgRab6BdwH1e0PQ6ZYMc3yLQJm6faf8goYHDJ7uRU+gbAAAAUUGfLkUVLN+BznEOh5fTV8lzZdXNmbQtHmURs/S4z/04sX/c5yXQAjkgX0sC2qb0YByR+BiM7qb8PCLi4y3UBZ//3HtEURqKtuvuIfwOuQF+IAAAADgBn09qRH+BB5XFIVzyB9wJv/LfxiM8ksv3Sg0OqcE4Yrqo1JjCRCketlzK8503P3asAJA7OT5xlQAAAIZBm1RJqEFsmUwI3/vEgWa+ohqyK7peNimapDBP/rmBAmpQD6OSjgVVooLkr8s+Xn8exw7XhF9TA1dM1L0nWyyYDn0YYo5qCungDaPef0jmHbsa/NST2MzomuWP12pBzS/D8UW68cG/U7WtRkyLaQ2o6xwlezid/ALSuNnhZUQ8LKvoBIjMxAAAAHNBn3JFFSyffT3DpS3ZRuxv7Y2vybz+bP8qJsrnWB3uZuD6OPtDGdX3ghT2/RUlUSZnoKabDsiTKHthUxvZZN5kazH/ANj2RZ4dXRg3GoPqzpyXMF08NHSw9ukMXVXozck1ZVN27kSwitnpaQpcwBkrkivgAAAATgGfkXREf4egL3/U7PdtZa3GHl5V5zeF2LLWs0MbzXOZlb3kOt8MN0Y4o4Qo4DC8r0V9YQ7rqrIWW9ryH/Yg4OqxHG3QnithZLe466OdQQAAAEEBn5NqRH+Gdbgx/3Uy5jiTO2z6NGuYZhDosvmqOy+HJIrRZ/1j+mdlQeOSAXr9NMUAqXflcAvnLw0HhARRdO9cgQAAAGJBm5ZJqEFsmUwUTG/77lcOB3OO+4ld2XaGYPZe+y3Nf7PM2BZEj3/ak7IMgYJZ292AtkE/+fDvHwSjzWBw4ZWnozBz4y6QXsHgftapY5tl8Hx4uA9PzOWUKmjPAFe5E1HvcQAAAC4Bn7VqRH+BB5WDztAwdp20lm5s3YSJceebDQR3Or6z816nVxtIBRfgErPOT5y3AAAAe0Gbt0nhClJlMCN//XbgfUwOWc23U11yALjVVb2w/NZo+TO7IFE7jF6EPo4+YMYzRIKbgiu3F3kB85pl36VhUcvvytbvKOop5eSVkjG1CTrVgXP7WxMDh0fOf1+slEO+lIyatR6SsaZM3IzBZreQq2JodfFlKOfZwWWZLAAAAQJBm9tJ4Q6JlMCN/8FQss3+ZKWvz1Vvgwf6g9Ha3XKud4lLmF6JaAC5IPporCJUHxbKHU7vYH4HXDa3zFXkkAfldXiwDv9hY4EqQrh7laZJ6pLfk5IRDi10QhgrLa6I1cpiSQ45EGgE+Xxsrg4o4flhDnQ5Ihthw83MHZGB3B10kbD+Qt/BG0au6ZcfFXOkkOBPmx5s9/Pzof7lYV/0dpNlxpwp/tqKTFwyLX4rqRzgPi5GmFJ0mKT/M/Hz/TxmtqRDs4NYZoN35tdFAHpO0d0mGyRL2D3iMHWnkDF/LAzFtqzHaDqN77mbGjyug1vU8DKtwrcKfMEHvF/O9Vqueeym3YEAAABmQZ/5RRE8n9O9y35ttP5XA8yfAgYvyCUdxljqWjOdhas1mddNiDe/3VrNDcwN1etJ4bY18tX/hyWyP4GIwlHrtyez+ZCz5ghiUA/FV9anrOpGlw/LmUeoI015mUjorg3GPJNvCnPAAAAAPgGeGHREf4Bx3gMa3MtJzXKNV1X/TvXH0R7p0TV51BB27zemt1rU8h+OxY4zhjatnT2bUD2FU0s2nNDpkPAcAAAAOQGeGmpEf9wqXEWq5RiVMaSwAxlOO+23LmXEvXP/IGbVJ5FsJzif9EnJ3uZT0nLj1G4aGVRcuYlRjwAAALxBmh1JqEFomUwU8b9uJt+4JAB2PDGkUZKn1LVaQmhnA3PX4/yJdbhJ6+ZhTvgJKp1Mpd4XsRmIufzYB0Gjg5HjsymTtW09NUcla0xdquO7VsJlNt1J+fS07p/3CtM9jIS4PXn662Okc2oWKdm6qRq6L6au5t7Fmndyg+y3aiIWUUQe67w5+Tvad6VsoM2DsXseu9ACjX7KP9LZCac5WFt8E5lYKT+WyzrehW+bryfOIPcnHP5Q7WKqAMVWwAAAADQBnjxqRH+Yjrg02yKk3NJk1+zTlSEjhOv9WdIP223OJgbbUqmQtvMLft3pIcnpIH3fSrw9AAAAWkGaP0nhClJlMFLG//QsHcexAhSB/PMs/QCqkuzYg99CJH+9JGOv5aFabrl2r2ctLBY6M2x4Zbvd9hht9FyQ6133GQK/Ox+Ab7R/IiaxaiHq0mgOYL42zxXHdQAAAC4Bnl5qRH+HsRJVvmNK5tDyf6EJ/1EigZ555nVBvboO63s3YDIFnXPkddwMkLKQAAAAYkGaQUnhDomUwUTG//QSPy7vTRAErbpXCmQqCkFeba8BWdHebwLy/TbXEM1CAdbroo7rrsm2hvYzIsEAE/hNIfu+eyIxwlnCjcjAA2wohh/23XNHd2jo5ayVrDbcMuugD8lZAAAAHQGeYGpEf4eNe2TaTeL8zPwYEq4gvbPJGyoYjEgzAAAAeEGaZEnhDyZTAjf/++A44HeOPnbQKXKFpJcaRdAGqeg9I1987zdEm3R60S7BqCKKCkCiM2sDvOC+Lw0lYZ0ctcQkkNkRnYsQfrl2/a8fGMNco18ZFNi4twcv152XeDYkOPPOBVk/8oWKhB5mE61fKSUIalH3QzOSlQAAAEdBnoJFETzfg/vT37uAoiBAB62fY4T7UhTDnZgXNwAMuT9zHwCV7PJPIez8/S26iuBAOJekb2B0RKzV9AVe/JyD5c0nVj80UQAAACYBnqNqRH+BB5WG0Yk9tWoDpe8K6FsAkXKMufOCq5zTCtm/E5mXXQAAAJZBmqhJqEFomUwI38etDOA3Wak4jIfVLa19s6Y0Rw/msIXrJHoKvxhcLNOe8r7slL4Wa5KoaggYjByeswoCOZrcIbqHWjdx7a4x7OyGnR642XUuyo+PuqhJohJR0twDTn/jWlpUwQC7UFGO8pvAyih3nOBLZzUtnawIh01N4/zJmVIzs03ryQ/CsNNwGXnSogOB7nyq3zAAAABCQZ7GRREsn81cfmOuKtuOmYL5CdH7jVoF6ZsqPcI7do3JhGpCQH0VNgZvZikZzZ0PoPATUfW5vZTR6uB/pZC6qxM9AAAAMQGe5XREf57Mf84VVo08PAmIFtyz68/xZtlcGHAK7LcnvzsXqehJ/LT41oVmRW0o5mcAAAArAZ7nakR/hnW4Mf925RcTlmD+h0U3BboCgA52LwY4NlHSNS0Dhr7XMuQKgQAAAEpBmupJqEFsmUwUTG/KCEHK3fhHiuibLgfkSu7LpIIY3nGHB/7BUdXL8qXzP68vVCju+6mDLSp04UM3Yy5sMOop5Db5UDv6zCEmEAAAABcBnwlqRH+BB5WDztJXdtm1EV3ecxpV2QAAAGZBmwtJ4QpSZTAjf/vhUdhEluql+hMAizNOFJ5HtLg8fPT3ztEuLEr1XGD5xhtaNEj9qa8FFVnv7A21sXMqfiyoIXW19PEUwYPgxOYHv+WGMXik1z6DYg4KxCAi5t9mD/wLAUi9+oEAAABjQZssSeEOiZTAjf/74Djgegh1UZQX/WDsSXsFC7ltHlVSseL9p8ERavj1mdYkVohq6rg2aShXbm9fgcktbs7aTqfRR48AbzyAdovzp+bia6bsdbHEDLxLykgR6RTMRnNhvmvRAAAAV0GbTUnhDyZTAjf/+8PvG6OB3OPtuWzfxGhQy+SIsOLPOJtIbYUIIMS0V4qVXaqngqCi//dOlQY8F7ndPEOY7ybsB51X9jVAWXsMnnsbNepMd2SEvVMXwQAAAEtBm25J4Q8mUwI3//vEhm6OCJW+nXGeBIiklAPuU/0ZI1kWndzC0dHYpngvdDP+J144OxWFY4ccxDyLLIj6szhtQ7/kePrw8DROXeAAAABZQZuSSeEPJlMCN//7wcOsd4IFAwltfPFspkUwU+tL4SrplClPxXSDmSimJiQeXcXpRNlc4r6tw/bM9uY4UWJpP7B364iwyGkyu+lQkibJiggxD+c+9Bvsu3gAAABGQZ+wRRE8n309Vqjbpbriq8XvOdcqau8JX/6UTOUVpRY9WKhf83qV6gviCooy2oTUp7uQzw7CwPScfAIZVLOOefPzZ7OvoQAAADoBn890RH+HsQIbL4x/wcxICg+7jcceczhq0AWjcfMYx47NzTdIEGN7i0aSwUyq5LiHB0crT/GwU/SAAAAANwGf0WpEf4Z1uDH/duUXE5Zg4eut6v/g2rzmcUOWPfIVraaS0gPT3J2ns5evZ1CI1ydypxrgCnQAAAA4QZvTSahBaJlMCN/7w+8bo4Hc4/EYYvlEGmfvhZdsa87VJ6JltdXwa00H/ieFaI/Y7du5hEl+VYsAAAA8QZv0SeEKUmUwI3/74ByBNNkfUtOKnL3+oyVxxmsptyCPddx/B5KMUbKrDnuC2Vyx0n1eE4KlwYD5Vy7VAAAAZEGaFUnhDomUwI3//WN7G3P4NYxvBghU7kkcBVleyVOEARn3ED3ZepoZq5RhvWCC/yqm2nrsXEILV4frQDCdGWESAMuwcPI6wEXjGyfMZqSkgvdx7JIjxrwTiUqQdGnnY8X6eu8AAABoQZo4SeEPJlMCN//74Djgdzj87Yh9SyaadvlLvsYIaG789+bDzKp4nI7Ux9DSaBly7yJ7/BdbBlU59uboJHfQNPDqioaaRPekZqtjhMesy31GxBhA1OKS4Ps6VJGf+JLm3Aa3xC0tIAwAAAA5QZ5WRRE834P709++8xDSwYyJmXntNFlWK92Sf/fQDIC13sgr2eS1PqEUq8jnPixfgD+QQYI32zWlAAAANQGed2pEf4EHlcUhXPIH3Am/8sLtAuTjaqUxAE/k7/M900QRhPQgPEHPLe77P+AZmROrczDlAAAAbkGafEmoQWiZTAjf+/2woAB97w5AEvi75kBbZEZSk6i5LHQkcNyP+hSCGc7+2SDRFNE5x5z/5s+xx+xzk7p8O4w7+mh+W5do9ddvYMtDfUKDtyA+F85e07Sekz50aiZ+A2LJ2uLJoEQWDwl+okL4AAAAbUGemkURLJ99mioJbr7gM0pfNOj+bZFRc0eF3Hz740MUqo06KEIR1upjhqV/Gzz5amPTy3yzTR6/dNTQU+pkKnktDCZZSdmNwG8n5niOiMJYGfQ5IyXF8S4KMq7xjbHIWbNosUD2y1HHlqmypIAAAABLAZ65dER/inEMYh/V/63CvSgvZEXbfhyQx2wa0izGid6YTPbRTv0aQJw+H2w2yN2UnmHoAMN0UTFRR/aDBMgxksMmIswPXoXpwRBBAAAARwGeu2pEf4etAFVTbLs6LYf57fPpukYdfH0vO7n4NhsecsU4ANjZ8RyLB7ip5BIX0Mx0+7ZAp+laB0Jf+sCagXiPlxvuigHgAAAAnEGav0moQWyZTAjf/Xb7hr774ABr8i42lWBj0uR1Z9SPh+ggHjsNADRKfCe1SC2gCz1ZV2wfdi4gTAYitJI1nTCixMcZaP4yGm/n4CfqGc6sKSk0KSFZjxx71wspXvBNzxUtZdrOzeAHxomBMOPBhu1dzQNzjtmlggUsl7VLlU74bd3nV5k7rd2++U9+8ggj3YDa0RWl4HtEE/EBfwAAAC5Bnt1FFSzfgzrToKs1FKb08R33Z2J50qKbNlQV9+d+S4eldtRoddXTZwHI3pZ9AAAAJgGe/mpEf4jgOwq2BPvs/N5cLI/GjH7DHCev2QPFr9/a5wG4YVCQAAAAYUGa4EmoQWyZTAjfyghD0zehAPTP6iP8vI/K8ulMYZHKRgVkbItuqNJ3HEAokTCb6Pnt7yl0OLzZwaSoIONyTBCZHmwGWPv3H83nFo1lWoU+cguVFa9NsDp0gOBgSBwkMEEAAABNQZsBSeEKUmUwI3/7w+8bo4DvMfUVj6CQ3wk0ZhZRNy/Weym5JT2cFcNQkc517Tr/49Tt0A4xq3KD+uTRFAhPphx4G+TQtX/b6Q9KLoAAAABIQZsiSeEOiZTAjf/74ByBNa2ipZMyAB8rFbNzvqfypmxt6LH+u+4jWhgXuj+cO3YiKHJNJ1WQFBNgS8wGPIrNRLwTz121G6k9AAAAXkGbQ0nhDyZTAjf/+8Gbu08WB7kttaaGR98kbL2C0gBL0MGbGJD6y21jf9/9lA37OP49VLmwv/65IaqtSBP9i3lOAQ5XqdAIIYxDxwBb1v2HeMCQXhXOkOBAThSFTcEAAABkQZtkSeEPJlMCN//8AoIGyiRKThReSI30K30Lv43t8SJ+8KvIhuBxhxo/in9rIRgUOPXJvOP0aqCWU6/xuEs9JaJ8GMhw30DvHBYxT/U0Kj8r2ehB+UmTWWcR/WTYZVe6KZ5swAAAAFRBm4VJ4Q8mUwI3//wDwzHA7nHfcLHfflRrPXe74AVeC1IbYIGvxkydW+YZkCPntNjxM2qO/tU91qFpC4ImWGO5QaY+ISilW3oVGEXv9m5/sTC2qUAAAABQQZunSeEPJlMFETxv++A44IUcwuy4ITXX3nLrPA3Q//hXK5HZ7fz2fB0R5AIIhVciP1yw8k0h9Np/gyGu6xn4JoLPSopgQYiI15qUK7avBOMAAAA4AZ/GakR/gRIN2pOx7ohh37pmz7Klk5+7v+6dmyrL07Wcjao8Z7/mUQRNj6WbSSQGD97CVmEVJHQAAAAtQZvISeEPJlMCN//74ByBNNlS+GzqUsik57u8ykO74fuvsnTVRg/9Ptqbp+XBAAAAUEGb6UnhDyZTAjf/++FR2EKSC0P90owSMxGoCsS0jIjkJXAY9nBDNCe1Bk/slGqt0QC1TqM5EeT/mtDAwGnbXjDvOUDKJN+rzAy12AIq0PbWAAAAa0GaCknhDyZTAjf/+/6PPsocLdBktI9rDLAZxNlt9wqUjirtVNt+xqsPZxzK6F5vZZijVpT+u6KITqFhQSn9k+W8ILsVjI7zrB9WmSKDyuy8O409jwJKv6h/HdxfnOL2WIBLbqi82goKAz6BAAAATUGaK0nhDyZTAjf//APFuECSmz/LfleQbzi9RxtCMFA3aSBw7cfB5VdpUKlIOO9D5eaJ77NvYp1NRMnpOo8KCNP4Zz+5SwFgGPmDpyQ/AAAAOkGaTEnhDyZTAjf/++A44IldYymy5ae50jTmHGl+LJMPRkmIEvlkx6J3enH/Z1ZYnSaiWfS9Qd3bZWEAAABRQZpwSeEPJlMCN//7xIFmvqIeT19k4jZC5qgXt/92KDK0y5JmDGlXIo9RAWJ1sdV9OzVDj0ny3mjY57wjldIvC+C3BkGsD8nC+eIrXj01+wxBAAAAQUGejkURPJ+A6iW+6tW/s8hDZPRPj/ZS4R77YE3gu1EWsAgpWZFUkBfemZrEkJkVcTAOnvNlQOvZHJ6uIRbuaIJAAAAALwGerXREf4egL3/jdwY/RuCrxWt5oYNMoyXcgYWnulwzd87FVyGo2D8hOuzPX1rgAAAAMAGer2pEf4Z1uDH/deqwuVerjcUBHrhTnazCTYcm52J3Yzn0y9s0foT1etFDJ+eKgQAAAIdBmrNJqEFomUwI3/vua0IEqusHZHfP89a8nE1UwIsrlNvzfQ3IirE5WhcQKxUzthp7BUf5HkhK3zIaLvN0HYEJo/wGnymLnnzFyHPfS9TQcFpsTCptT+W+ftmIjvDK0LfLQWDU5ngMJHz/n4ZUiF0aa3S3RfkSsidDqKKs1LyafNz8fl1XPT4AAAAoQZ7RRREs34M606CdEYoROePxtpctRNT26oIcfF6VX0DqQPJvGMKc4AAAACUBnvJqRH+JNh+Vsm/raPip6RP9Og6Y/+RSUgkDHk9xc2O0jrwxAAAAnUGa90moQWyZTAjf+8Gbu2qQAjMVaMFUREfFFP+x4btt00ueekq5HGN89Wm19wbx3Kdw+o0PqxQpY0DcAXpv9K4n1LD86yV7ODCzYbxbeNispTBcvrKPCMwkXGwZIKam/tG7bdG9+lzB6J8XZvmU48gqeMwbRPYvIVNA81JXiwX5VfXEp+tOTePj9OoOWcrIF8LDW8sH5NteKrA41RcAAABMQZ8VRRUsn357kZSwzxr9U8YWP9/xl0oFI+aTRf6oWUqAXw298uTlr2goAldpwAKLQthKd2eV5xwZcFfuV6zka+DNcWzVwddT5dqn4QAAAEsBnzR0RH+HsPlbvvmSW/qDUBoMhf2usSxi8Wum4LDDW0Fu2KeEyWSegiZwWr0lVsQcu/F0Necrzgb7KR5QDSLx9yh9ScE0/rzHaJwAAAA7AZ82akR/hNfANpljEqY1fAK9wQesw2QAGiR2LAytbhpI31Aax1ACoVdfHniU9EpvumT6lsdtl9FTtjAAAACFQZs5SahBbJlMFExv/APDMcEKP+df8b2bcbvMcHhLvzdzipVTXYQpeF6xqpUhj4bNawZXtZlSHo26T+WH3zqaIllB37/NGNQ1VkY8BOUKE+zpXj8ANW28eEEEdJdDLN1PiFpb5PIMKKW9o9IlM9LDe1K216q6Y7zZmIMkCDaMmuxIIuprEQAAAEwBn1hqRH+HsRJVgFi+ZrlLDr4CkfmPNcriMt98U9nbw18dR45otYG8DOLKOeaxiW7U396pvPAr+t6Gn8NLGkSaktuhYMuySbpXKWCAAAAAp0GbXUnhClJlMCN/+8FGR0k7BwEOlhvznFDc8MKd6PlOwJwx+IkDWajbiHtKLtq12J317Zrrh1gtLttHBFn7OwW9kLbr59zfz9WSxFbcUiYfBU+4PJI46FozlfEF8jFGF4/hW1jigr3Oe+UyN6JQ7klFrYmpLuv1jvLAjYdokw6od71OZf5GD4K1jtQiVwCSoCZlTnckIlQyvMjv/ZinSucoeLKUdjlgAAAAREGfe0U0TJ97IKea95i5I4Oz7pxueJfZG4kOmA+QbzXKw7vkhvqrsiqjd/kL4XyY3lQf0TmOeb7Jxwp736HbF34x3jaBAAAANwGfmnREf4ew+GS+7XjuUTR+buTP+d/8H1HLLwuG6VVjsCWSzm1R7OCK0QXFWRrnznHKxhkYoPwAAAAuAZ+cakR/iOA7Criq74tIQStcVo0TUW87D3aTI7cUwCOn+Wxs+VAcJAD2/FWI1wAAAHBBm4BJqEFomUwI3/vgOOB3OPtuWWgGLWHjA2dK83VlKvt2+n8reBM2XzHD3zBxBIqIRM/rFRT7vg6GzXbZYVE4qWcVvmjn8+cOQkH82Z90ruTl7Fjpv4YW3KK5nKkBzEFtS3TO1V0jlDTMhHqvo4khAAAAVkGfvkURLN+D+9OfdQR6q/7IXDaBvaDpMhPEswOM1+7l4S2FsorT/pLAVOjCdAJLasClJUflp/50NXl8VQKz+m+7rpgfYkqz2G+vKYeS9jaTSkokmfCAAAAAKQGf32pEf4EHlcUhXPIuvJ32qeFdoZAknNUtwBScyRRJETefzNQclkJBAAAAcUGbxEmoQWyZTAjf+/2woAB9YANHst9xAn8Ot4dHzqUwafrp05d0vUqpjWj5to+oP2Bvpa+fUTyobSQwuvbyPiaW5y0mT1/O4V2M/c/cW/9kBex4FuohVoCvWoVQg/16KJ9s7jM9tnwTb/JeLWapEhPoAAAATkGf4kUVLJ99mioJWnVwU28WW1/7jNoqYi8ej+b1wieEOmo+V6ElGEDFkN0d6FE9/r6vZJXwBtcGjNFa7FbJe4U0C3VVH6u9ab3qqXMOQQAAADsBngF0RH+HsQIbL4x/uLfAQThciDnXlgeXNkbuTWOhTs0T8clBXVZVPYHhl/rxl3hE9I+Hrbh7KVpysQAAADEBngNqRH+Gdbgx/3Xq50pXw8mqoa5/S6osvnea87hsE9/0OGaERxf1tCTkMhfZ7ECEAAAAY0GaBkmoQWyZTBRMb8oIQcmggO5x33CtfQoNIePFJ40j0Rw5ce/FyFAT5ISrtO7pr4/ZWcfWhBygM64QbSUQzME+lubt8WOWt0M6f2A3BfB70bi4jXwkpKQm6O/cjjNNV79bYAAAAEYBniVqRH+BB5WDzvLCHMeSc6D+pWTKrSTRgYiHkmyaD8ady4xoFuAGAC4gZprGG/guu9ew6u+KEQtPsf7XkTxEL5oOibpxAAAAfUGaJ0nhClJlMCN/+8EZPlhBHWC1fwUqYD4gl8M84OnkMvjhAR1eA25Xfq3KZ7HDMrdVm+l24vSAkfG3ha5Og11+q/XKHpFWzSPOORK4/q+VjpIgK+wW3rohvnrLXSLrg00iFleonFUn2Zm8WiDb62AKQG8rPQI2jk4tFOzPAAAAd0GaS0nhDomUwI3/yghA8yNyXL4QryWtpO9TBWG+5lvHFT8rZrJwigk2mHXYghgKHt37c17joLyvASuy11M5xsDnIGswur45lcp/pg6KTMv64Bm5l7g8OAtRTsOC/qs/q8KCOxtRlop8JGxN+wVyJAecO+x+WH2JAAAATUGeaUURPJ99biHttP5WYp4RAgYvyI+p8tb+TRf6oWaAAXw4HkfPwn3YQah/Hn50Dc0k3Kbwtlf++vgNgAn442VMlPHd1szsO3j9hrRgAAAAQQGeiHREf4ew+Vu++ZJcAjosbnx/+SKnisgJBzGvhE2wYk44wff7Oj+WUOlh1tx4rNtkmx7x72zfz/ISHcplyHXRAAAAOAGeimpEf4TUK0MsYlTGr4BXuCEvYQ8yfXuACUj2t8k/d+aCycsFhAIUjm9ScUV3x2YTx3NuXIgZAAAAdkGajUmoQWiZTBTxv/vD2xXRwQo/5tqyIg0oegdJzRBsuIl+7b7ib975tSk5oO+RagBoP94bo8xTAUoC6UAV7hD2pl/WnNEYhjeekT1PkCFVNeqc3AZwsou8tNpbY/mzbsB2O5/LkQ+UJINnK+GSH8HOdUvNWTkAAABKAZ6sakR/h7ESVYBYvct3ZUH80aLSL3lILK48NjH93k+JR1JE74PB6N+HZWuHDRS04N7OCEWv96qoQw6Qs01sUJTT++6b2XHrGL0AAACqQZqxSeEKUmUwI3/7wetWHYdbgBdtDzUw0EHUbfim0S92FkIfRPj16pWKKLwRo3AU8FDsa9QId76SyIZbrKSW1k3GGqM46d3clPOrcjz5FiWx32mC+NDau6Yn5iLvV/ruKE+CYviocKuN1PAgXeEie0yP1W3y7ir7QTZDPhf3ZL2j64xZerSGI+I/vp60I4Yfhme/LvW+sG6aWg7iglDM6mb9UFBSV9g/opoAAABZQZ7PRTRMn37oQ9jXIx1lBX9mqfhAUp5p6J6k829lJWvVzf6M7ggijphbu0clU5M1vBwhyYZzt/rM9/6ylFgz0BPTGJMWpIznDxoA5gzNXt6b4LnZ35PBfoAAAABGAZ7udER/gHIYyAVJ11hb+xYFCFd58j6mfUxiy60D2axlSlNv4jbTnYFRvT61K/ZgtX21MlfVI4/OkXJuUpop7f64MnECrQAAACkBnvBqRH+I4DsKsOxDYDk/6PHIMsNIpxKsIM+ybhNlwmFnio9k9isHXAAAALZBmvVJqEFomUwI3/vBm7tPFg7we8ueSW3AsmgN1pq6VL5mONAom5p4474jhWe4CPB0/d75Ta+3AbyIfJutqpUJfL9iHeRoITsbO2cFPftnBnbG+U99/U08f+CG9I+wfsE58CY8Lf3D3w/MItP7ygZt3mLhmC8NDeSwaV013CnmXqHO8zdoI9keTJr4VJTM94XD+qoMrzWDNiUadSEvxRjUAUo8/J8dVa+AVsL8l9fEVbn49GcOQAAAAFNBnxNFESyffW4h7bT+VmzHLoEDF+QSPtyPFrba/50Rn+HndNiDb5+E/AuX0YogGD8wr8BoMCMYsRM5UoXrME0WskDjNcWx93MwmTeYntQpF1DOgQAAAEABnzJ0RH+HsPlaZzJzVMKscA/2L+9TyyQd/3O+nm90ZXl8Ksqn5Id4lYTPtFj0zEdxLJQd6obfw58aUu4rDzZRAAAANwGfNGpEf4TUKzbRXPm8OWGDdlZhb7MPjSB0cLQgO5R2NblH67E1i0ao4Dkk3g/masUQXUJUJOEAAABnQZs3SahBbJlMFExv/APDMcJc0yucaeHB6PvL/u479qPe7hcn8oUhAYtxr7RP+9d6EP7vADBZs2ne/afeqbqNMqONZ5d2u+aLATtq5FH7fBN2ApzqD2YHoBAbx7AylRx7VOtEtRa5YAAAADgBn1ZqRH+HsOQxQ4IYEJ9R380tftck9EBbFpYrpDKVIDssfLu6OZD0KIvfnYuj1dQTryD5+83wVwAAAMNBm1tJ4QpSZTAjf/vhUdhyzm1U2leLW9Qg70H6FjRAKOWGmq9x5x025+Fv3p3S/FQm5Tp0onIj4h8Zx4V1izZBK7CKWM3Sg4BQE8SIJ4khFfeHsmBs3BWM4yet7zvBdepBB46hzuc6TbCM87USRY/aAltVYBThJ2Dy1Kk+X7oELhrfv+rE7sa2+kv58rOX6d1pHhsn+G4eGGB5RP9NO1f9TZuBI2+uPWtPojfHGH2/8Aq9EEUM/xWKbPACAkCDEeGpfVcAAAA+QZ95RTRMn3sgp5r3mdEHB2fdKSdICJFhyUw9U6B6nMiDQFTH10Mf+Nz59OffzcVpUjcUpLvcP2XrpMgmGzwAAAA7AZ+YdER/h7D4ZJjHpgBMmQmFWTxuFtuKN+1nWmrvVP+cuwy6cfDaxn9EPm//fbSX/BC1BXlKAR4dXRgAAAAyAZ+aakR/iTYflcVbbqsjzADeWk/N4gygRKahtBY4GTk21u4qAH9//uy/xUzjKGN4gOEAAABkQZufSahBaJlMCN/7wZu7TxYO7yWtpO9TBWejmOuTienHcMDCHKUGeJgU8kvIN+p4Q3OK88HSnfDGE1MdkaQi+aoz5+X/+HFemTy4gYlIXyYWnSJDJo/cOHmT5qotrLEMdESHSAAAAE1Bn71FESyffW4h9qH5W+/h2dk4U8g5Pt7nJndunw3PPWIDrz9uaZyKwd46Sk4PBi2OZZ+81wbZ8oc5SG+A22lt4AZz/jYoLj5rNtjCgQAAADsBn9x0RH+AcOUFoNKCQYCGgc//6BZF/1e++Y5yM2PoTn2oKho+FzY+MYnRSCxwYlPnGESpnGT2xIEfQQAAADABn95qRH+E1Cs20Vz50jhSs7iL2EQrs9j3tB+E0OgBAtIl74/4sMyLjOTnhw5AY8AAAAEBQZvBSahBbJlMFExv08O2aXJzY9aWPq93Hlrmqv5B0QtXOYZZloiKjehd74nU1AKH52Z3+TFO92LMM98MEm40V7V+C1GenupfeqytRxcuUcO6u7TqCPTeIwqzTSEBSPq9aYeksIflaFksME00+qb4RPN3k9VCAuDXIDcWlIfhtRfV+FZNJvsVXv9aSTadgIZXTCel8E+da9dln1o4EHbsCWRxMU227KCu+aooSEKKHlx/66IpLeADW25hj/AgX0h2GHLGpU+s9OsBO2nNNGGqVU2s+i6OVvVpOmhdZkIXLFAuaV6rywH//SZMiw2+TxWyRZ7vy1v8JUsXVPcmUxHw0+EAAAA+AZ/gakR/18yldm4EIIwJMqwy/zSlRwwRYyU0uDbf23QW+RyODJVXr6oxUFynmeXdw9rHk4uHPi3FC8lMF5AAAACwQZvjSeEKUmUwUsb/9BI9qBbAgvoKrhEVvOFiUYViG/M5EJ8mdYUiEaj97nKw+nocaPSSxXVNZ1shjf+5GKqK3d7aQA54mLlWRII10yz7m8Bep30K/8/RHcJBI+3YKy6Dnp4kI0bb5XEQ/cK7bxUfpMy7cMapjRKOoRLMTOm5NDf8KQfAqQ5TeInNVpTmGv///kutmRd+MQ4u5qz5Vq8W/uVrDIc4XyeP9023KOhDiYMAAAA6AZ4CakR/mJmxpOr92pSK0PJ/oQn/bH9a0f7ds14XLC9W6fZ4EbVHvaoG3QUOZmMZfHun6TyTV7Kr/wAAAH9BmgVJ4Q6JlMFExv/74VHYcuN4xrugJWIKSCyEVY45OXNWVZpnwYhiQRi8ouWZeYuX9UrcKkoRdSfCZZoOmnAnOQp7HCaBaZsLqFrDoWpfb/ojyr+kPVjjyThmNhOz59P+kfdZc/N9o20arxf55nhcJHGZLjukhPirl/LdFsHrAAAAGQGeJGpEf4jgOwqw7EQkqKLsnOa2jmmZovYAAAChQZooSeEPJlMCN//Hr5DS9RRAa54Z7hoy0gQSdNxHb8oFhjXsrpRKZMbMTTdkBvu3TAncUA10EqZ7+4nEteFrv70v8BmFG1wgiHERfC31S38TWliBAZGrZPRYU78tMXq6ichNMahFbFt7DygIkxM0pvU/Vl+YLZqN/y1MDAS6SE19XdtA3PtOjnSNlsS0SqWCQzXThAQMqHaQZ2uGXkv+3sEAAAA7QZ5GRRE839BbCV26Byo0weWSA9YDX3+QMejMyEEDStqDfwdOOZPoewjLmsgoJTYBfGhx9VUyx/y7bNAAAAAgAZ5nakR/gQeVxSFc8gfcCb/ywu0C5M8eQRzzLW5tT1UAAACEQZpsSahBaJlMCN/7/bCgARAV0d4PVu2DIDo//3hRFBNpnloEynMkS6IgSJm98XQILX4peWuq1XgtmjZvcfov+so/e8U3LVg0FPoU/Mo4SQs54IeRwOVrQ4GMDhGOj3J+7LW8cm8BcglvzqyhewVZdBEH6P6sGuFKFltAUejrF5DmiW+AAAAAVEGeikURLJ99mioJb39gpt5ZOf/5t8yMdqhflAjF3VSF9LKRKHSP+6fBR0zNoOBXiwoFPSpm2LefrXJYQBXANQoIKzterZ//Rd7bAfQwQE2Okj3HIQAAAEIBnql0RH+HsQIbL4x/uLfAQTkb4yILaCZ3iPJCU9pVv5zI9a2kfiP3oWWAYeh919yuVlfR5mK81H/GkOepytZyXlMAAAAtAZ6rakR/hnW4Mf91Mix4JlquVA684PFShL+fqcs8dpqe6XtLIH5N1f8BA1aBAAAAOEGarUmoQWyZTAjf++5XDgdzjvuJXdlt/kTGF+RV/LdKMM3O2Izx8kplgk0r1rh9cxk0ue6JeItZAAAATEGaz0nhClJlMFFSxv/74VHYRK6xsfYetUOdE5zQVKQV9OZlZwDb2XeIsRYByM1q78Ie7kjTv+lVLWwES1iMaYWnsFBfK73puU6X9u4AAAAYAZ7uakR/iTYflbJxD5R0UYmRI1Tlb8yAAAAAWkGa8UnhDomUwUTG//vgOOB3OPztiP2UA+261TZeWO+pLr6FuxY1Gm4KJTeKb22KE1CtKd2PE/YlwvYroPfLhJg8ZrWdDVfKHDYJ8TeXliWDg+DDMHa6evgdfQAAADQBnxBqRH+BB400xy7oyzZN7DiQhYX6mHtDc+034OuRbz0duExV4R2w0AHsAUlmmJIw06XAAAAAVUGbEknhDyZTAjf/++A44Ilb6dcaU5H0bjlOnxLIo0SXrQTWgueivNvHPE/8vDiGPA/ersyEsLWjNH8iFcw0ZL+tkT9WxPFLyHI0aZt9JkWOOsnrayAAAABgQZs0SeEPJlMFETxv/AOiKyBDZVHR4Nb12r2vQbwl1ZUGpuF2dtyRYmAgNHFu9YI8DTzjaZZP2/AhZiSYSEOKyyEGzaeRtjr6LBzxI4L9G5pO+D4VbaY3srs2TWCx4mqXAAAANQGfU2pEf4pxDGIcj/67dNk+uMhrkyCin4ep171/SmuXlGTPbQXRJPBnppBZpXl2UXy4JS9RAAAAU0GbWEnhDyZTAjf//ALQu8EXYDiMMWrNiM6GWH7uYInAl2MtZ/kgekQpPORIiGhZ/nRumtGaN2e05a4CsUGjPa7Jt1HD/aPBVyqAiWxtSFal//ERAAAAOEGfdkURPJ99K3q8cj1odDChudhfcpeV0qEp95cV8xPVLYNNjT/96mXg6NNq46gQvUi2kw+4owVAAAAALwGflXREf4MOYBiW8DO/bTvzhcPmLoHmc9yW0C7QrF99GHjG17qqQEM6OaQ7qATwAAAANgGfl2pEf4EHlVcd04Is9ClW6FZCH4nYyUbEMAwSiUszmoA/R2IMPcZs3ej3LlzkBx5akMr+kQAAAJxBm5lJqEFomUwI3/1jexiPPuVF4N3+NUAUuIHEz2QO9Qjed+3c9KFSmMPrydJHG31a3UXHzy+f5ILebe40628nHuliEo6jFwOTxVv1UhoKiyPtASysBMxVAoFRxcO10B0l4c3iX8DbGNOa6io3rXjVXTO0z3QG2JpqNAWBXcHV1gHdDwEtkGhKhQAfl2adu3mdyJam0GGvKzmODyAAAABTQZu6SeEKUmUwI3/9JoP7PlJhERd+nNBdwdAGicV2sDdshxtss0WwtsFE9iV+OAFa0FSYHSTT81z38H7GxqDvlLWAdJmzfoGzta6+3HuOXIhSfDAAAACGQZveSeEOiZTAjf/9Y3sbYq0egyLMuoFUSS1zW3tAR1EhxIsXQLfJqvT5msgJUJ5lF6n+eX5oaHM3FvUq7s0ubwvX+FHL351BlWvMj+p0hEQccUBBuAKEAtp/FcKbFnScMfZEX2zhg5cgAhSWHPmnBc4H3/GBfKQU8pvzMeWNgvr9aUU5AVUAAAAzQZ/8RRE8n3wHSlhnjX6wuuQMkOxSt1zZ9hoUMwDDnzb488UvfimL2oJcGWgtG3RGIynAAAAAJAGeG3REf4BsKtkMVY91KCkU84i2Lj9E6hfd4mZLJmupc44LzQAAAEQBnh1qRH+HoC9/1Oz3bWZ3RPxWyj+fW0/FzMZb82m6SxXIuxzXGz2O9ueW8NrCaIuPQmb+HlSCvaDrTOCW1gf+t9oHuQAAAEhBmgJJqEFomUwI3/vrPjgIGEb1mbkKmaEb7B3LpV9EsRLeONLVcFxuuED+lT3/pAzfeZLEw2Kb38X2btrExLs1IQ4mp1/hywIAAAA8QZ4gRREsn309tLhK5Et0327KxGXKAtwcI4JfdoKiSSv6tpnwJ1fs0HcHDvesezzWNAerFq04zWyFWRqVAAAANgGeX3REf4MOYBiW8DPHRkut1B2MvM8gxWoMkGJG6fJNRKcDcXEbvKD2g6pI1wtMG2qZlsUlgAAAACgBnkFqRH+BB5WDztAx3ehsFB1Y2ENPbEBth3hGfB9PEXgdSLcgUejBAAAAbUGaQ0moQWyZTAjf/XbVbdFBvxTAPcrrcz6VZOCFULaB6H0iT9bpVOi048iO1wulW4BtFnQgnXqt/dwAcB0kgppQux9PGL3+oIrcHf/88oLvTn3UeHSiGxY0fUoo0QSzEULynDOOcZGpy8/0v00AAABzQZpkSeEKUmUwI3/KCEPTN6EA9ymLrly8O6B7t7VmUeTH/G0a7t2ifM3r1iMMfcBrASVNUPovKG8GYNYNCbz/ppRUblstrHKIbzRyZab8jak9fpDfmpy/m5pgm0ZFpIjmcCLOk2J7HyFBwjEKZneHITdafAAAAFpBmoVJ4Q6JlMCN/8qLkmeQgEPSPuqCmr5ArtbnLRl4fTGGkErnT9fA9ShDSntjr7GEjAb8W3TZIJnx5IZJCQVx1Awj6XTCsn7hJoDgdkQML672DCqIiq+WioAAAABHQZqmSeEPJlMCN//74ByArVvp1EDVN71GG530l/KjHXzb02zLOHiQXYfX9VEBccGpAGmsDLsvD7H889ZLz9QGv3FQvckk+bkAAABkQZrHSeEPJlMCN//7wZwidiwPbbYagjwfDcPSQg33o5EFwjljEQKjlBmdpaMBOqA7OQf5sN+AJ3xRw8XQzvvEfWu7eb+JbmKBNgl+YwIiKZxi6zn0Whj9DxvtFKFYv/L4uXFg8AAAAGNBmuhJ4Q8mUwI3//wCgK7wSDWMAKdUGMUb645fTdwYnhpkNPya+L8DWZnT0HAzto+auQYGWoyZsX3gYPKgXfkFrIfwZz5ku2kqON3Et4mdMuvJcP/wIlB3jHcb0QXsltP714EAAABiQZsJSeEPJlMCN//77lcOA7zH4jDFrjpxOzOGUd5ZawgfJTi+lCyrDkvXFWhsE8XF0nUH+68r3FUz/xnMGyUBef43CwO4Fu7mkzlFHt7U2RY0wKM3AWV7ncpWtZdHMWbpKuAAAAByQZsqSeEPJlMCN//77lcOA9yrrlSWZGlkSvhsHkUA+TYWmiPcd9erKEll/VhIYwxpcKy2emhfmAiAQ3BXo41lpahCQkJa5wBO43r2oRyWCTe4Hr4f5iHFvgO/n9/U/Wv3cJr0kF/5/7WMzkvgsE0ESqtXAAAARkGbS0nhDyZTAjf/++s+OCFI9lK+huo+FMo46iSF2IPShKnDhY+b4Nyyfn422/qBUX1fTLaA9u3h/ChZF+GVwRdZ2WRkBeEAAAA6QZtsSeEPJlMCN//74DjgQ7ZHzgzntGfTAsTgNjJN+UQhfBz3E1P/plckBND9E86yk083GNZZLnqs4QAAAFxBm41J4Q8mUwI3//vhUdhCXYakQCl1mC9vjlAE6W6K+YGIyz//JA3BetDKXvDl3VvDxmOvEBMODpgm+jCKU2Gl8xC9oBzxxspKYcHrqYH05l+P8v6Y5Qk/1pPAwQAAAE9Bm65J4Q8mUwI3//v+jfjges12kf9dZQFJRGONVgWYOoxzG5uK3Jxnwr2M+6smx4PBb//XwNxXDNlrK0yKc8dfQ3w+72OG2IyxsTTpWG8fAAAAQUGbz0nhDyZTAjf//APDMcDucfnbEfvHS6zrNEJPWVHkE15RWPV3XeHelWYNVUgciz2fa0ECdV6bGtEIrhn9uo1vAAAAP0Gb8EnhDyZTAjf/yghC9V+DzI7WDVxvgaNgcVv0j50t1Fi5NV8PQP4/2jlF1RxSfv6bfmVSxmd7SAf+mKajzQAAAFhBmhFJ4Q8mUwI3//vBm7tPFg9mlCdoqUDYZ8xiLA/58i+MrJi/Oj9iY90OsO12pTviJqIUVx7vxap1f+uydZk4zez/EtYp/4fYp4Q3XO1R4juDaEeRsAQQAAAAUkGaMknhDyZTAjf/++uE4Ai5s7YnlAoVYdUxHA6lgf+GQxBdkuL/GOcCv2WRMFlJ071NInuH+zP8ObDVzOx5D/yBFdXBjGv5oMdw9Mw36f6YbxAAAABYQZpWSeEPJlMCN//KCEPdL34Qpx+Iwxas1iPqZ7d7C263bvdkqW2ukMa2nT+82TtnjrJ/eEhfgMZHvbgwHN3G1mSeyf5ERFLmZN+rihWcq91jTRoFRIgx2QAAAEpBnnRFETyffZoqE5BKXkkXp2UOeJEo59JeY9aj9AeTC6IYWjavARsEhBIEdBw2Rve/kWbUHHscnz8S/B1NjHiIhOhIZBxTgdtyQQAAACcBnpN0RH+DDmAYl6IayBmmZNEDqZSB4/1mlPS3tDlsV/Fh1LZ9b8EAAAAoAZ6VakR/gQeVg87QeYm3bURZpmDefiWKuuiEiUDlMyxeXtBIWKCTYAAAAHRBmplJqEFomUwI3/vhUdicsgfnibHpetEFWH90rguJ3ElBMQTQH7c3ygYCf/Drrg/dtGQkavBREa3bvVMAltQMkijFYWTfqWL//8Egk401tHl0A5hfv/3NvhBuZ/FxLMknFPgLLzedIjC8QeurMrgT7aB7fgAAAC1BnrdFESzfgfhWj8hRbl2W307v17sFwmYW+jS8UpzFjs/eVgK5cyRdyT/WT7kAAABGAZ7YakR/gQ9o5bvE/zRjV0ev/5aKcO19Luh0B14WQ6ZPUsDRRYb+qb/GrGriVtRBld0p3TqxJM3rCnvICZvAAlZ5yfOX4AAAAIZBmt1JqEFsmUwI3/vrP1EjJ9ApoRLZwxrgZM9KFuXGgqknu1QvZh+YXzmCVkYA2kXu5XfTPYyy8+Cwf9E3yexwGKwVVsHQkC08DCU5fhqZfK1z8qrqNUDzTK72NXRlKxrFR1ZmiVPj2Xq4At3GtTri/A05eWs/ccw0A8+MLhzoW7U93GW4fAAAAFZBnvtFFSyffT2/dTSVQUKj9rf+0mUMU9bY6PcvN6uG7mXsKRuTE+sYe2c0wQGLRh7SjC+dfY/ykuTFxQBC9+JuGvBsR4DhRM1QspMMQ0PRXYOYT3CwOQAAAEABnxp0RH+FEsrDaMSpa8wIWWIvYRFuurkj19cl46aZDCbA0SBpIQmv8EkNyrj0XwVlw8L1phD0sXTHNQHFOHuAAAAAUgGfHGpEf4ew5DFDghAjpjpN8kQrknsoqluAN7r/wdMHwzFeUfhsnaYLctWarPhpJ1pPTcHcQt1NddNzpv6yDKOANWXO32tWe5euY78MFmYtAkEAAAC/QZsBSahBbJlMCN/8GV/wgjheSMGlboEggPQ8iUHEclFSu/AtGO3I3Vk2ER/4BHv8Tbue6BAabqGiQdgTzc4oRsQ1gh9ZCu/VbhiFCPYnJiHdwoX4YU+hRuenZpNtsObSdPTabZeCoJiif6l0xR2ZC1r2W4UwfkY/Kbu9X8RRqWwjiZ+lCEe2yO97qkKitGLRXhsDDPy6A42LbEzFS5m3Bfj9FzHAZ9IDPEUXHyZS+YwQIKAdldUuWh2n8/PuDNEAAABCQZ8/RRUsn3sgp5r3Cm/SDs+6O3nB5JhyYbLrIl3aW8UcgJDENn36FSTX+WiqBPcaJGrhj1oT+XMgOFmQ+3rs4Yq4AAAAQwGfXnREf4ew+GSYx8peZo32mD0ZxX3r90gQmxgMzNTtzm1R7ItzavQ1ubiOdKh6cKDbpl9nH+Ni2s3puThONblt6/kAAAA0AZ9AakR/iOA7CrirIwsXZRWN3XV5/opBUIx5D+UhAdCUpwQTa9Dm77iIFrwPR5D0Wos1uAAAAHZBm0RJqEFsmUwI3/vhUdhEC5b/sF1SYFY16qcgc4pX3SjdOcxNGiWTn6CYPHzCyrtX74+6hxGZS/gSlPlQe04vlBzgd9y4b9wSmNmTbl7qHsmH/8WHyl3DolsoVYqYdHMsqsE7ZL7tq75kchJTZcF0jg7A6sSBAAAATkGfYkUVLN99TaxVytK9J0LEC7e9QMMUfwZ+bYFBPy6qi3f3mYx5+pfBsels2wDlgtcp/zu436EhZrd06L3ibH1Yqz/2X7weI/SnSNFnaQAAADsBn4NqRH+BB5XFIVzyCm2kI+3PwfLld5Q/dKDRNU+AxchZjpfprwo6gFhSOxp5leYMuT7tEBqg+rCoyAAAAIlBm4hJqEFsmUwI3/vBjyyMgPRE6nSRtV0KpMIijNNSxNV6fK186yR1lZCuSqeK0zpSMgyYXrqP8vepgaucl/4jaGmvVdxyh0e9LI0kKda6Fdc80iHHn7M/V+CUeawOHDK09GYOfGXSC9g8VnDTn54KoCAuYx0MXXmNYVT7fUCj+DImYMTgYnhipAAAAGZBn6ZFFSyffZoqCW6+1T8IDFjZ/2UM4MdbOfM6cSgf0UAPojz7pI5kEFknrYw+ENpQAYUTYtr4SqEWUn7junbR8Wt56p5h7uhIcGkZL92u9QC38gUwBk1Ex5E22TDcZLcBpyUi+s0AAABQAZ/FdER/h7ECGy+MeBKdNfVxUPeZpFC63XlrzdjoJZreQchSff/+dpqAridjEZDcjRKKurU0q+oPY3sW4PFZGDT1eSQOPqlSLhDyd9rNotIAAAA+AZ/HakR/hnW4Mf925ITm3O+8SFKZ3zxUoRkrtFQyng+VH/hMQ/Sik3etRXXXNlxlZUXH5+UDoy4GCxn/BYEAAABpQZvKSahBbJlMFExv++5XDgdzjvuFjuy7jd6PFJ41DAH3l/X1jKdL2Nh/X8fAtcq/fEL8Eo81gcOGVp6Mwc+MukF7B4jwey/r9CcqRx65yLF3b31zSZv614yBGc4M9cNdQpEICcQUO1uAAAAALAGf6WpEf4EHlYPO0HmACcpM3NnBfztXc3YkJpZoDQxZYYDyvNb0Pj+EZFz7AAAAhkGb60nhClJlMCN//XbV/0CALgiH+1UnutP56iKl/F3VC1DeTp73t9y15JQ8NdWsfi5ioASn3Q9MNo/IlUuX52yG0tN0TwNd1c7MGXk7xvpMQBKyiLGoACsHChpPyimGD1eDR61bb/pFXOSc8YOX5I6sYEzmjHHEsZBjKDlilzAK/wULUzcvAAABQkGaD0nhDomUwI3/03wWPufQrGYALLUwsJoxEMXwn8/5NCn2gZUgE1FzDh4+XyZdmkDpzw9gAUC6kukl2lQx24vXEN5j5DOH1GIYPDhINtBDfyJ8WfNT/Zh9gvQkc/PgO/6SLQMtIl38tX4A4gmbreT1L7HRARz72Z1XC3C1MiMvoDJ7VVn05DVGD32R54DC+p00fXgSF34RaX9p/nhPqfqWedglnMwJtxRR8T655s+4GIfQryP4RVGWXnkJEq22REKoBSZ6ysplrhgLEE+6q2u3y8N7Ts9Gp2LifzfiCd1+mPeg3wsFSzJW1W4R7PzKVLjmOSAYJ0Z1tw5ZriRBPXA+vDYfsvkJniSnJ44XojzZkBW5ylkDkLDJmRGK798IPsVpoZydkPXtIrYMAjNlGhRi/1/7d/7lypvEWzYzVcs56JUAAABSQZ4tRRE8n9O9y35ttP5XA8tj1j/f8ilMGgI8Y0X+r/O18i9VDvuDif6vnGMNy/wHlXR9Ruw6ja3/zSdWeANSXvyVEX9mq8ty/YEARNUYIs31GQAAADkBnkx0RH+AcOZcuyaLd28wzeOfzpKa+1FVhZF6ix7MqNSnD1fychy8xNmjL0jFeLjxZFPXR4lL64AAAAA5AZ5OakR/18Wdv2q5RiSBDhQMM5SV2fbqkEFaulvczIGbVJ49MJ0WW4sHI4AZItwvvenMKbRODBhSAAAAq0GaUUmoQWiZTBTxv8IgPkZ4dbAn/C5fKq3XjZar6aAnr/IwTI6NfhYxOOdY90vHFZcpPrVgtYQEGJcIs9138OsOgln47ETCAT3wsxykodHaHt5j5zd3pHkivyreE3pXWJ39gwvIPH7BtJC2ujNyuwcpegUpqIHioXE7MREQ7wqg3InLIx6/+VYvQ5iyZJg9fQEMDyurdPwUG97r3XLTUK5xQEoR2kd86Au5sQAAAC8BnnBqRH+YjqwgmYeCDpFqcX+FVib2Lrr/3UMPLGmd1jhtC6ur496SG8bD43VmgAAAAFxBmnNJ4QpSZTBSxv/74Djga4r9eoJPe8V1sdAFa0CHxG4haVB5Ovh8u8POW4TcTMxsptzm5D7YlBYw8lG6CyBBn/nDvcnVGWnkGJ6RrJu5JKQCK1vTDPu2rBdI6QAAADABnpJqRH+HsRJVgDVevwny/m2ehVgkT1H0/+i6OSAZ+6Mbwnkk7i2l6zyDhghf5iUAAABdQZqVSeEOiZTBRMb/+8HrVh2HfgnJcCBHcdoH/EHq//8+L7kvHr9AwlItE+XW8GuuTbZfIyKey+xtzx7/z7v0mSlg6JyClysHvmOYE28a/Hm+/EFhTDgOJfpyQ5cVAAAAHQGetGpEf4jgOwqweXy2F7Du2qK5I8q22RN0g1ThAAAAbkGauEnhDyZTAjf/++A44Hc4+26pcXq1U9uXkaTwRX1v/02a9jGRPrQ2GgD4ByCHatmQBBgoB4cOdikx24qn7FSN+jELlbhvLO0Fzbb5WsWgoL++Qm4ZaBej9aGyTmk4go9Jd/7p1u9ZDw+8QmFQAAAAPUGe1kURPN99TanKxqV6ToWIFi47JibnpS3MA/q3/m8HYEriYfagkjp32QrwpuOP5xPUmZSLfxN/vcgqC/gAAAAjAZ73akR/gQeVxSFc8gpXH9+3PwF5PUSrMv/xOAUOg3reTKEAAADGQZr8SahBaJlMCN/PivaVsDLRwju4FJRtGyJld3PRfB6Xtp+0EP2ehbnkk6yo1l9S2z7XQSpm8aK8H35oc1ChfOjV/41daHgy3+x8N9GxWBhXCKBEFAFwfSFUpL6aLnnZF/WsKnCiSNGkpuCtg8KGhAifBX2jdPWMZybTgbq9FQPBnlegtHQ1gMS0JPKw8fpMr/l4VaXgPN/WiXupgFWBTWMufBFXx/pwY4hzINiNuuh8vDdV6dZqzU6E2zTAetTgZI718atmAAAAV0GfGkURLJ/NW8aI33FXDdlOIL56x/LE61rb/LEC+So3AiQ/8F2RvGBS9ER8hLuHAUZ1asXN72jsM5Y8Rsb/8m/wSfi863R2YTaypcHVPZVzRNGbKL9hgAAAADYBnzl0RH+ezH/OFUf3+Y9OhGlOXkyMreXizJcXSVs+sr5MlHfOxYn+EpqyS+pzJWKhTWIZJCkAAAArAZ87akR/hnW4Mf77wcPYM8+uXxARwgNs3mq4c8S0xyjev6Bw96W1w3E5wAAAAENBmz5JqEFsmUwUTG/8A8MxwHeY77MPL7mjRFVFR0Z/iYOC4Yp43LzKh/KoiomQ7CIxYiuEq/otbdQROB4TPZzrum1PAAAAFwGfXWpEf4EHlYPO0DB1ukwiu5JQBDnZAAAAjkGbX0nhClJlMCN/++FR2DWFm+9o4wWclxR2utVG2+azrP+x4C74zgq9slOh87PLWTihGxNCkPERnZsOwKKHGBDItM6n+UO9MhbU4gRdPCbDzqp0LYuZWoHRtYrRrXBf15GZzNVJg8jQkdGrKOzzbeXrr8TrvjpovWonk3/zv//ORDTwucm9tODFLZPZUK8AAABmQZtgSeEOiZTAjf/7/o39RD7fvlHF/LNkISXGH4z56Ai/14Fui1YLh2MuTPImlrfB6tXNm3i2ltMHwB0R0XoODBGGdW+0dQtq/9aAzDsN+WbhUEnMSSne3XEH5LffBS7TGVCwM92/AAAAXEGbgUnhDyZTAjf/yghD3S9+EKcj7DVPmcgdNOSh7FRWYDdiiT8S5ms2TgSRC1HdbXoPao/j1pyyvcmUyyOr7D8VfrorJ9Y3Ia+/Vy5I8VajatIrRzOFE3LjoTFuAAAAWkGboknhDyZTAjf/++A44Ilb6dVZAj6+5eBpPqpy0hGCly7Fd7tZOR8lBt+xP1Bp4udbg/7G0dq2F/ig4mG065Jsb5K+VasAFZ0NwTnEg3ZzHdvF7Eg7qw9SgQAAAFNBm8ZJ4Q8mUwI3//vBw6x3g997CjDRenHxJ0XtD3+ZsaWlrDYyGJRc1JoEwL+kntl1Dn61wxnH31YZTaUVFntEYkpX3nfnQ65nvul3NPCVGfp4PQAAAElBn+RFETyffV6IS33RIzDNhC6Rz4N/7Il8kZTcQoK/cQdyXIGV2ki+axJCPZb9IDd9bBgUYo7Q242WZg7cP39Ppn+R9qeii6RAAAAAMwGeA3REf4exAhsvjH6163w40hSkl0dhxqvJ28aur4o8Qw8s0CKm6YtPcSeQHz3E9C8uvAAAACgBngVqRH+HrQBVU2djmFhWg1aZbX8552eSKPQTInN9Jr83hsggwky5AAAAXkGaB0moQWiZTAjf/APDMcDucfhpJoEdgP1f82nspVGylrZAY5TqKCcynw/jcC+Jqlhay5xuF9h/wgvXJVaziw53TDP2hodN6QtZC6r5EmXQMakhypMctn19xQ16J6AAAABAQZooSeEKUmUwI3/74DjgiVfUWivtLYxGQT5W2ZTVlespsDp7BuKK9RAPQwV4UZkOEH/M1tUBScFtvFQZbdu4IQAAAGxBmklJ4Q6JlMCN//vhUdh6HHGNd0B/p9FfdgmAb2hpUl+urE83Dowm8tGYb3I7ijBirHv5xiW0qPCI8pYvAfgKKPOp/lNBaxO92h5BlDhH+jDsdWP13tgaQXw/5AYaS+Jl1e5qLi0meHE4ZMAAAABgQZpsSeEPJlMCN//74Djgdzke7Ewm/iNKmZv+eBdZSsBm3kdy+Jxax1DOGv4ctvMEmTptqD5q7dbTzcZOZfcFWec54awcbIvWYJNuhCIrSQwganEwlNR93+G0LUHO/L7HAAAALUGeikURPN99Tam3Cx80XWuScbrI8ef27x4OlvfwhN8a0R2GOpfaziVyuEYanQAAADEBnqtqRH+BB5WG0Yk+JGGXBmlhdn2n24dWagV/yd/lHGugo5EAQHiE5RPtDB43FuuRAAAAhEGasEmoQWiZTAjf+/2woACNjELAEvi75dnr9c1/IljWX13Dpl+arXwZ48mEtEWdudCKaJzjzn/zZ9jj9jnJ3T4dxh+JDYdtPLY9alEq/yvIYKIf6nHoWUKR4HUAVgyKfewwUN6TBSDcysMFrfJrLB3fH67CeBSbapkCkw0R7IsA+zrYIQAAAFhBns5FESyffZoqCW9/YL3/eJSR//l8bcyjUNJcqqPIQpL4s4lCEI7ELUU1x73ujzgFVMGgdWaoy6vKtupvoLwBbryeu7ncAuloej1ir6dhXCn/9p5M5mGAAAAAQgGe7XREf4exAhsvjH61lOxDQpSTOE5Rr6c/YMJcR40gl5ozXSNr5K0gN30y3xwpq1Doi+RBmswRSVIyyiZDT4vuRAAAAEIBnu9qRH+GdcG6fA0ab8kNOTv52rX70y5CdadgShyyECf+CNopeKFkbu/KLZSfnj1sdOK6r0ii9DLbnwNgGTiqNCEAAACbQZrzSahBbJlMCN/KCEDO8x7vwT26lNl/gUfBmtbre0IKq6E/qqdCqRddo6tWS7hBeYS8byVLG6s5bUk2/ZsQLzfaGjnk5bK2DJGat0L5NBmmSHVq8s6V5vFgGpYgYBplI9JWytw3yEcIR9mMup/7M/PMBbSjjYsrk7wqyK/gnjeQOMfm8zcWP5JRN3q8ihlin1D6q61GFw1zLmAAAAAtQZ8RRRUs34M606CqUVFCJWU4Ky0MnnSmBsKVBYKXtwQD9bO1GgwnOtqMNh0gAAAALgGfMmpEf4TUKzWWd2fqIhqL177tilPXm5J8DJFaCASw/ejbPc30eus0z8KtHbcAAABNQZs0SahBbJlMCN/8Aqdq3gjhBnFy5XA/L+1MyAlYxhiqCHpYw+5gq9qVl/jSDu/Zxz8k+rcaAMAWshx2Lt6lAsCXP1ggHOhg0h9HO+EAAABOQZtVSeEKUmUwI3/74D1EHSkI9FUIquTDM3rKaNWSFy696/Qi+BwnpHE+ukHESGyxNmYBi2gCDySgK7ykEfDx1eZ83MsyM6XidSYhkeMVAAAAR0GbdknhDomUwI3/++A44DhT2ug1cazm1y9Eq/w51RkAPdbegFjO86K9lsQdW+KZ8FITyXlyVrMcZ2dEB+mty+r1aZPOgTfAAAAAXkGbl0nhDyZTAjf/+8Gbu08WB7e+xX/RKyCUlwLsCM947CKmmivCIJDojdrhemKwj68SRMATMf/1ZjV0RX9q+NeV7iUhcKy7URTsFOrM01gL2NmAF0DuInJP62fDE0AAAABlQZu4SeEPJlMCN//8AoCu8Ic2OSncAbLPwyKakd8JvTcutKPXYQaN0OqVJNStMu+lFOFblgBZyntC3N9Z9aj1T/qbHte4IkOIPGS4G7XHAIEZZ+vYQDNER4xNgppDUTCuIneQCJ8AAABbQZvZSeEPJlMCN//KCED59yngQHs2OJROVEKCAWdcGFHRSwN1YWw4Ja+vCA1DR/wWuwzoDuSzJ8m8/s8a8vs6MregbVfZO4g5Nu/OfbvwQXQzbNm531lzo6++2AAABOdliIQAJ/++DLeYgN9XYj5zg5B5QJuOQipq5DurWEjnW3UorRPL+Gmp7W/eIApq4hA6GO7MXHIyDHsVucz/aMFYtKTDM8z25ZLQfiNZJqSiDN9723ESzhz79Cwt5M+8vOWuuGhG+Tvq9foXjrYoVXiVDwEsVfUmlk5TXbEqyYI+MZpo4mDJdCI6FifaFbOBypE2MZf5gDnzZ5v8KmDeJZU85nFu03pUJirEg0Ch4LatUJlnXu1uVKmUsxM0BfauVLDb5iKlbV8Rxe706n9ByCWTYcNmUsWgK2BOx01vNUSepY8KQudanHo3PJDxToO0OLrabPAghuBeO7puO+FrVqFKMISkHJjS3EmTdBG02wgZepLedOuw6T+jsJltu6ElokeeisTTrzRHarkaC1NCRtCBiNuKEigoGG5TW/NYmxDmIUwJrPGuuEwRd1KazxAA2rsqyTJqLM0sjQZEKi/uESROrh+l53XQlmI5Fr2TToEa1Mz9Q0myFVMzlLsQqx3/ZRQx/eaJZU1/3iRY67X7pQgN8jPG7Mu6N8yvaY+Aqj7NpykjyTLzJHyZJOS8uJC5eDrvviZZfFd8hAX6jnJ0ZJE1ceTBCU68b2mfGkwl9oA3rNyC3kk0Uu8eOhHrIGj+7a64L8d3FiwoDY47nxsn1Y9aseVrNpfuZL8I3UMDdSnrXcWoYyESX7OVr2O8rZxn9LDxM6fe5FRhF4y7x35eGeAWARki3ksb7bLxg8DWvLG/+RwVPcpcizEE8l3Xt5+avL+ZiAXLEnwFkaUZDgHR/YEmvUJL0akz3W+4HoMrkbavzzs+VqiPswLqfDbnOg90ye7LST+LLq83FG1+6MSJeC7uys0Em20StpYNB17y0kIWJY4WoqsWsZ431/5zMrMnsO/w1fUI2ocSHX8pNHwFp4puUb+kPcS/Qh09xlBE4IT8uQiu8T2UDKEJf9iCtfqONZHMdPq/oob+ISU7mDp/lbACtBQUxbSPdO3HSaLklWzLgg0sclDx8EWk7I6bx2ZO5UAMGXVQtj+CCaYjfoBtDkXhWf7nWsT0hHJZrDi2Lra0z6P95hKbBHipvDDcqX6r8JNPqXyyckj8WVa7ygUVg99xXIM3DG5RRWnB41GtWfCqWChNibXPGQ6uh1H6EF8dmXH8surDi51hLT19j9myqwkNKmT8sfWOXCSiZtHYNl9wayNXOnug23v1/WAGCRrLJeHUUtVHTIX5R9cvqnbFcs3PTXx9R74CfvppdnOxRFNeTTBjRlRfQ+cjnV6LGJ+ftcX+9UqNBGia2nKp1uKP4ccPHEr/CsMJq4q8HWHOTdHO3jV+GZWnEoLlQ6Chy7Hbj//eKkHIOL4i+yqkzkPRSbB7sfUP9vHZ0uCQVKXYKu//r0MgUKN8P5NWtC5FEPRp8HfsROGFe/2knJyGh8s53uvSdYVz93B3c9Z2acgXQNzwm2yabDfBYT4AbEZLwPZK7m0eC5LKEHAxOLjG9NU26PPjXEMfajfx/8iXqLqjmgozrPx8k41W10JV5nX0OJWGa0qBU80PviynD2x3OkSEckOt3iTHjB9RR5c5nsqTL0EgDg1YWzkmuLeJTeKxOy3pZHTLdnNqj7pt+W/2OWi1UaTdjKxFLON8eBwDadGOUrbe2ZBGiOQ8/zfFIWmtk1vwMVGZD/Iod8uAAAAATkGaIWxG//vEnVTkBINY1xBnYqV2GYveJxG6pO1ytvZdPbEXt5LeT2BV/GdImH/VH7OxTyATpFD6nkhtuovj/nla9ShYi1okzGdxi99RJwAAAC5BmkI8IZMphG/74DjgQ7aLylrdTsuPGlvjUJs/BGQzvh+6+xQ9wiS1hY7U4V8YAAAAhEGaY0nhDyZTAjf//Zwh2MSAtcqcH6DBw+kiRJyBfSLlTlMVKsKbIZvDezqc3T5Wd3NGZa+B+Pjdq0z2oeP3B1ynbNRtvd91FAPKHi9eAZQhXkzjRR1YSlxJOK3yfwo9klCCVTafgfuL5sDFzP0e62FQJQ60CKx6fANRbJhNojOKZPzzwQAAAIlBmoRJ4Q8mUwI3//1jexb3GfDeuG3QIp0iOWNrPYaRICXk20KwX4AJSK3G01g3yC/ow3Dys3hXKUnzco3wdRZFXg3LRNpOH7wjxDdO6SGSNKnYi89v4MEaiVQLPE3ixfnhiAfLznBdpEhe+byYE2PLypisngAgqtZKdabisoZIYcgmWrafiKuWvwAAAFFBmqVJ4Q8mUwI3//v9sKAA/rehi7vgA080JM8SmYpPIuCaWfrNkOSrviyYia4nOycnCZbBmW3b9POpyhuj0YMiNpTWFdH6ofSMf1BO6NsLk44AAAAvQZrGSeEPJlMCN//74ByCXW+nXGeBo12ucPu1Q3jIarCClBC9oB9Vn8fnB9XMTBkAAABQQZrqSeEPJlMCN//7w3wmMjIEMD8lMTeUbP8EWg5RZj/kvdI2pZcX8B1kNiYTAl2O0lslP+qTbr+aNjnvHbN84RS2/4Py+FdljGGrZaq0CiEAAABBQZ8IRRE8n4DqJb7q1b+zyENk9E+P9lLAQ/bAnTrrSAta8TkZGO0FbRoL9hiQMM6FUzWT+8DFzTHADnP9RRGF2gkAAAAwAZ8ndER/h6AvXSm3Bj9a9bnh63e7Lt0dwASWkqO5RTd88ehBuGdNotf7Fgu58KElAAAAMQGfKWpEf4etKGWAMuE31cSYTo8cYlrTJwO4hT1iaYKoh5yIznLWIYM+u2AdetqlZhgAAACoQZstSahBaJlMCN/77mtspIuqyAydbMmdV2hUlnnEKtJjAgm2ilJGA1nNcqvKyjyAWC+D8UjH9KvdInrfn8JalYu7So7UCAXcOewUU5OwhCi4Gn87YJgH6zDs1GEUzCJ8L4JntHmTTg/MXIc99L1NBE6e+O3vv+OYZ2RYqJ+FOkeVPph7PM7bvRsalZ9v+6j8c+a2e9FeVHNkUcRhx7xXVeu40Ckc84xcAAAAKUGfS0URLN+DOtOgq1Qm/f1v5ExIEbysTirTzQFn6S1t6UO66bJ9QaeBAAAAJgGfbGpEf4jgOwq2BOh6mMPEvE5KJWUuH45P8CXP+ZdCo9KA8njAAAAAkEGbcUmoQWyZTAjf+8dbF2niwezLbWnHkGUoGoj7QIIOkFGvROy82ieNrbjUGA+u27nmEN5OyIUtEYm1Aa29w3BQ+77exlM1D2vRT7tzMqq7x40NgyQf9UyDOeFALvKyW/UjEdd+ZTj3e7SNxPWZbzdyhtLBs/d1mrAgFnRSFTdTWE14mrJQdgV/38CNOEQBzwAAAE1Bn49FFSyffW4h7bT+QcFhQIGL8ZhHK2aOc2q4nANbZ0GFL6Hw3v96ujtmH3X1f2UGhLCU7s8t2ALYQ5jkD+tMpZBnKJdI7zOIZg3UfgAAAEYBn650RH+HsPlaZUroiY+o5658f/kip2JFnP8SLnrUdiJYATw+vVm1xpmGefahnkYvhxrqmtK90ncP5Iwqoo7sRj3ypQHBAAAANQGfsGpEf4TUKzbRXPm7/rDCb1PYRHnrxQN3OgOeuwBoQBUyoz454hkxVpE0j3mzip8afJWLAAAAgkGbs0moQWyZTBRMb/wDxbhBIn8Ka729n9RowP1GPGwmIKU7l6A++D+NCeFcwINyDokoWo5uNZep8A+tnW0VIejZtiEBw828C5obBhnpukZNMqbuoZzcKwUcdD+TUNnmG1P7VNp3N/HRWhTj24xVfemeYzegPZhwqPtf5Vu3ZOBdTWMAAABFAZ/SakR/gQeNFYTouJTIPcX9ODyG/mOY1LtDn4Z4fdloxMzM5TGxEPn1lw/NNkwpQDBeevJFEc5m+/9Go9q+nDesjfivAAAAmEGb10nhClJlMCN/++FR2Jc7HZXI1ggkQzpQbLdE3oyI2e9RpinMX6ttYauzP3zhtktbs6rBQk+2bRLN2Zii5+KnqAGy6xu+oWg28GhfC0K53s+ETlNxGnX6kbnMFmbPxkhZYhjRahFfw8hooaTDBlOVRucGi/xGwHC4t9N/ktnpgbYThbnVUl6j4xW46umuQZce+h6EgY5YAAAAP0Gf9UU0TJ97IKea95tkF1g7PulJOBBxUf3oQ8Fgj/d/EvsVRz2v92wCCfdP3jpU1XFxnM2iUhxS2IZ5eBICgAAAADUBnhR0RH+HsQBkvuNRCE5o/cA/hS/2X1Ifx8G+v32b7PS02yVwPwQRWbCtWYhN3Fq4gdnPwQAAADMBnhZqRH+I4DsKsOugF5agq8KzwVB0SOKhHPr71cVpWK3aQSfqFUhMbf8T7Lt5HmLfgm4AAACIQZoaSahBaJlMCN/74Djgdzj87Yh9SyaadyNqF883ukqwWJfvVqpKQh+Cq2+ijPdoCbCBI/sOgdzCRvS17Gg4tWxC4C6lNS9Io4e2m05/dsoHILwYVElOoJDZ/MwYJQtRJ7sMYyXzLK7leOwuqfKqCHgoXnqt0hDtESiOuqtwcgabWxgY6yvJsAAAAF5BnjhFESzfg/tuGYHKjTB5ZIQwyRxCIT4MdkzIYOd8632MtmXbBRZjmn1eHX7fw0cYLgael3Bt26A9uYW+WQfDAKavRUOn1s++6E4EarbpfnHpm4uc7noN/2Jer1MpAAAAKwGeWWpEf4EHlYbRiT21xAh/pYXaBtudfqA0qjCKgx/lAHKVbTowHw8rOoEAAAB5QZpeSahBbJlMCN/7xIFmvqIe/Z6A/pmYt10zyE29uhmz+B2Dkk2oJvashyLN7AdP2QFAd8nNOQALgCphNk8C9krr0S5+4t/7IC9jwLdRCyqMfhEx0LgdtUukX/rP3vR0dcaKekFBDBp9NJs3ELeS44zfwhtjbfJhoQAAAHRBnnxFFSyfgBLo0GKtuOp0S/62v/Y0HNZk73O1NRvXJ78CK1NMTqIBwvVi+7nrhZNFcYmSvBtS8aarNw0M8/4bqY7V6eSiefuhvHdnfNViR+1mI4DfLT0GD/NfHmcU+shsVRJ+gRZvcMuzDyIyflNTInWB4AAAAEABnpt0RH+KcQxiH9X/rcK9KC9jqj38E9Cx6kZJfJmOLXjenMkHGGnhFyfDLoqpt/GOsv4ayAmoWSNp05UVCWVgAAAALgGenWpEf4Z1uDH/dvK5xRRkl16iX0NdLBx0eflcOeGwT2wyO5F6Jk/qa44eccEAAABoQZqASahBbJlMFExv++5XDgdzj8Soxq4R8HyFDsHMPNBcfu6u0JC3yOt9Zf/oOlinfwad64b/5ZuUQET9ghJNRnLBbp2fwAk7+v9LiLGc0Wy7m/MlaKm5bEbUwKubBxg4bDV+ieC+CYAAAABHAZ6/akR/gQeVg87ypn7D6HEVU8+rDMLH00gtcABN3WFNH/SXGNGasZxuBcRvLGcib1ofmI9nJ1EIhafY/27h+f0U80Zw3TgAAABiQZqhSeEKUmUwI3/74VHYnLGdrsN8fZFN+AYUaSy1qUzsqr5angQrl4TAR4cjThoFgPzgtGOm/+rZ6RiEUZb1p/ibsShVQuItv9q99GCa7QVtMF7k5//VfdZc+IdQXtUSRRkAAAB0QZrFSeEOiZTAjf/7wZu7T4AET0n1K/wcGaW5FPPPeSEVQq1wqPmVjHHGzMMaLIFS9D9xLjG8IHT3uYZjAjm+gAJnanr133+IRxjFAFTRBkx2srdq61z07DgwIEr4PcBIbo6eVoNdPRdn3p7C6JhswPrZOxYAAABJQZ7jRRE8n31uIe20/lZinhECBiXrAbzFWKWjBooZi4YZeyeDcnLXtkoxy+PPzoG5pJuU3cZMSZ6MT5Sk+0MRyXthyQf8s2D30QAAAD8BnwJ0RH+HsPlaZzJLf1BqA0GQv7XVPFANHr6Ehnvy7et3ba61H8sodLUCczvDm+Xl28e9s38/yEiA9pPYUvcAAAA1AZ8EakR/hNQrVcoxFq3sC1exT2bccLesQAlI9raV/idtSMz02Fau0gbClXfq1HVRq3Pjj/AAAABxQZsHSahBaJlMFPG//AgTwyBIj/nZtn0aw/YuY4OjamoTtQYFPOSeyuPadadakSLJXyzkE+yUN1FhGpSW0m0mzOfSaLl6RQi3UFaI3XVG4DOFmNwv2QiaKoOLHwxouX+jWVEmhJBs6gzIfBVG2WvwKcEAAABAAZ8makR/gQeNFYTo0teWJjfvPMFcU66/zULU/AjuVpjv6JMs2v8Uk8dPu8G96ISprega9vCssVT4Py38fqdQhgAAAK5BmytJ4QpSZTAjf/vhUdg9mx2D8J1iA1xk+D/u/eWFRKPYXoR4YrnXnQ9BJv4U7EyAxknkxv1P4KGiZu1XCzyMTxrAEalziwXrOaH5CIg19HpXT8adjJuBd+g8AVU1iY+ragKeFtHO8+FXCs1iqBNVFtNR+QrIASXXiKH9ySLmtNtwkBLAWVkei24IbUCqoeu2eSmChhIvR8SYf+gjoCylIj0hQWhGkF9P9sVMkdkAAABbQZ9JRTRMn3sgp5r3m2NaOy3X8HPswTGPT1V6qqyW8TAIkb1l84u61oezFaXM31Ex6tTfZXwH8l41V+ZLVvdVcPEZOg5HCqpsYndTaGYDZU2IBrf/HbgESDFhYwAAAEYBn2h0RH+HsPhkvu1KRWh5P9CE/7YMZ8EH8FtRGt2iZZYTb+I20oUztyey5/gcWsMembQCMoi5qzr3GWB6TyRqnmRNbxfwAAAAKwGfampEf4TUK0JrudPWO9DprT1cI+LcpiA9d9Uv2SIBa8dwd7ktVb0K/sAAAAC9QZtvSahBaJlMCN/7wZwidiwO32tSx4CI9VHvf8FfC8imIOnDW7X9dLvxLE8x1BQ2+Flva/hMlBEeDnDW8Kyzs7nJ6R+vwdkQhkORE0CohKGRVTKH6TrAI0jsBYWIHEzhI/DD598H/p+WJCB1WNOv/rDnbdHgZYdaqnTDwcDWZDxml47NphnI8tuJ1LEeyLF+MD0/bGs2eXrqa0n+B7DWP332LtukSOVrs5rOy329VanveHft7k/xSHl0G6dpAAAAUUGfjUURLJ99PYWDWGUO/inwiBAxf/vE27SbnXVov+t7ciiHO3pMMbk5hOMpHCmJoAtQU+lyEMKhoGTKDStuGSi6zfSBRDqB7UFC7Ub6lRrH8AAAAEIBn6x0RH+HsPlaZzJL6599KwwEOf8I9PqjXFRQKe4Cdxrd9B4PO3J74fYU+3lo88Z9cxAIlkoQasgv4bTkDHt417sAAAA3AZ+uakR/hNQrNtFc+dI4UgJcDtsPKVP01/w7gtDdiG8weomxmMV2jRs+mYeb8YsCdKiDBK6EQAAAAG9Bm7FJqEFsmUwUTG/8A8MxwQo/6EY7tRsTvMVqrKMcemP06VoeULKvsG05ajGPmO2MoK4g9WsHj/Pw+OQXY/RfWlxRBDYU5qtmrwgOMibbsU1z+xBDOIFJfIEU7QhyGFhExFTti7BO80z66P6DnI8AAAA2AZ/QakR/h7DkMUOCIRZCQlH2r3Q6geLsF83caeCI1bZY+Xd0YL1LmfzC34lEn3p7FbZoKir3AAAAl0Gb1UnhClJlMCN/+8EZPlhAhzY4hjuP9zep6GX41gR0vpNBRWRZ0UJPip2CeOhHj9isfDJTZVLrPDCGanOAhHmGphCIAPC1fKwnrEyemLG2qRrRX0P+lii+U7+s0UA9lDzojDxhAWTZ2G0F/aKbm54hHFNSHJtKf+oJOsv2i/hXe6R9WFXNJ0qKbHaCDXE5SAk17PNkKsEAAAA3QZ/zRTRMn37oQ9jW+VIhCNZdo1DTECbzHzraqtIWwcFXlHzfzpW3DjaPNagdZrahX3k2FcjBQQAAADsBnhJ0RH+HsPhkmMfK/TRTx1yNE7nufRZv4q2ObBg0Y85tUezANeXV6ceBVnWXr0jPx3kmWjuo3pxLLwAAADABnhRqRH+I4DsKquYeXbXa80z+Rp/+iDuHWPDYZQauQjmlruP/lzQ5MNgtrxuo9+AAAABwQZoZSahBaJlMCN/7wZu7TxYO6+X+O2RPzaO3P9AZSfEckBGR/iJyC69X2mzEdsbEARVatDv39gXZOq5q1DbnQ0mecF485kMmU//5paM4iZd8h23uDrAJgM1rO1fswnjjEy9xf4szlbNzDknjKgAlhwAAAE1BnjdFESyffW4h9qH5WbMcugQMX5WbtC61gVlqffRYYu75nGloWw2+flAxCV1Cr5kLWgK0cUy5ESSi8mwPgAklHpUS2QdtwkoX7J+G7AAAADcBnlZ0RH+HsPlaZzJ17dg40xRcL+9TtQ96Ic/V7TTmS8UoawBAXq1QI1QIwW783wiVM4yfS6ugAAAAMQGeWGpEf4TUK1A4rnMTHCUTE6/ZCv+/TbCro2Jb5kNAeSIDyYf8WGZnEtBvd7y7p2EAAADtQZpbSahBbJlMFExv1TpdGQI7CrAUef/B5h6ksk8VHxtoqBOtcVtjm20cMpPx7YKCba8PtsDiRbMvl2Vzp7ViJJGJslpXSe4Da21o01Tre+PL4ouLsDE2fza/o8u0e/o193a92YPHQjxT1F+Xl837GKhxj1zCeTYoBc72ajk9wDGlXMBiyPuI/NQcbrbcGJ79AkHjk9YLxxYD/SNREJer2QOEa7kdnKAs2nctMaAcgEj7lXi2nC/Ax5gsLQKiDWNIJ6EZsLrTgM5L6GMBmxmY1G//WcuQa5/+wM7s0+v98zEhGRgkPv8Pqi9S60lbAAAANgGeempEf9wncM6KwnCyxLkxv6hw6qpeuv8JH8/pARA3aQiOdCW0jL7dvTy2HOkGJ2WI+L8RYwAAAL5Bmn1J4QpSZTBSxv/0Ej2oDXQbnmJ9xHt6Y3M7+zkeot1lPkVAjOrBk5KutJxuFikQjMMm54iVUcgq0mUJ7XmxxVFK1B6T6213KH562Y+RpgQQ2x7gidLCR4uQyyI44VQZbTdq2DELgkXsCdiNKr8x4ESrFQl2hMUKVh9TAo5qMet3aV71iz0PMWrns+RwD+p05KZ5PkbOR8UZdEHr12Xyp6MXO03cMCdxvt8XSDnM5DDAwory7MxXFStmXjinAAAAOwGenGpEf5iZsaTqsZBVHN23kGvXifJ7fiaFVng/pASKG+Iuw2tQVKHpTjjL49UFkCnlve1QSoeNybR0AAAAc0Gan0nhDomUwUTG//vB61Ydg33yGjTc8vx+XfX6HCnb3oXqwFbwLkycdi4a18wxKaCie2RwyAvsyA0uooexqJ/tYWo+LRvOoUmoxRFTOQSzQtQ02/rq9eSsBVCGtHKcyXiFeX4sl7E4a69xKNtZ4XR2mcEAAAA2AZ6+akR/h417ZNpqEQkqKLsoWa64MImq3wZOBY+C6717C9/lyBZ0g6hTD0NY1rt85FbbEW/AAAAAnUGaoknhDyZTAjf/x6o2uVDXrlL01pKNSmhQXr4P4TJQsSM1NFMlKvnjcKn1PKq2soE7idryXsHEaf0O8GkHVQjbY/noZr0ipij3QQYPsNY5BiccYwiMaOIoVLU7GYyZILtn12cl54KfPwPTwjvSw+K4bQ5yOg3G/Z0pavUu5uigl8D6ssFCj6v5AXSFHWAjabypdfzcI1hiTFK83ZYAAABNQZ7ARRE839BbCV3PROKhTO6SEYHVcap6D/bB3xdjsl1hYn21dsFTAVhbVYBtR9XbqXR7RUq+tQJh1Q7mOALEIaPB0INz3lNIRir4aIkAAAAyAZ7hakR/gQeVW1VzyCm83gTep8bM3if7zuz+nCbMLSDxTS+tvaDoEjU6XYt9NoaA70AAAACFQZrmSahBaJlMCN/Bc9IeRdMKBQUoDdmOpylyK4+kq5dpRNdskQahtc3w+l0fbnrSe3FvIzIgw10Zk3I2ld9hkChK5nRYgL3ZtOHEBaiunfTa4ieYkASLLA/NG32QGn6PIgI+rNletA0lBbBKcG5b3PN4cer0cfvtkcfPCHp+QZU01FpMCwAAAFxBnwRFESyffT284Ut2gLU6T/OF97/L42Zk7TpYejMubi0UT1LYIwyi8wna5sal5SW2vk3PndfO9Ygjv7w527MVQundwsPGKGTzDFCdJOti60VM/ivQ0/9l0RrwrQAAADMBnyN0RH+HsQIbL4x/wbzGsYYRQbbp3oKzYEUN0nPNIKpfJlJ0wD5Mk4OdIwxB+uB2eKgAAAAnAZ8lakR/h60AVVNnY0Snvntxfl4ah0IXe77u05wjItrjxSxk7twJAAAAPEGbJ0moQWyZTAjf++5XDgdzkPa1JXdlt/iHg3e20KGGbpJNaAeD/NHcMVqnpuI+rGRg85TMlkZYGQ+buAAAAGVBm0lJ4QpSZTBRUsb/++FR2DWGd4E74BcOffFFDrjSdNqIHPjqmL/A4V7zuz6iK2RJtZ67hmaIV4RT+DJzIBQaAPE+vuu5mX9yRlf+Xz3MDvkp8OYGDoPwsmZQhPHmCII0d6nKUQAAABcBn2hqRH+HjXtk2moRCSooxNiKe6PW+QAAAHlBm2tJ4Q6JlMFExv/74Djgdzj87DB+8bpj7b61MAK02m0SjItmx1gb8vC/xkdmAxbAK3mtGvA6G2vVc9Tnw2/i3sUSrGL/HbhKMa2sK/+P9jhm46fzs6O00Ue9XW7FkptH/JCazHfcLFaniM469HEARwqsesFKmIPAAAAANQGfimpEf4EHjTTHLttWmKSJqlJBcodqFPtQbfjXOzXaRUJjHdTp9FoeBf71Kn+yqxtU3pUeAAAAXEGbjEnhDyZTAjf/++FR2EAuCDGvtGowoi8nI+jfRpfapBnzNQPYo2Vr8QmdtgYjsnpaHEEUEj5fzdXZ5z8ZMEZYVB+N4bG+rt/oQ/tJOOnO8NMl+l71MxcY1QGBAAAAZ0GbrknhDyZTBRE8b/vG3I0XvCGxfcoqF4IjXpZsDkZv7oRIdkuZllH2s4RDfJ2UFZ0K+2+MrbzLGTPgjZmyPT8lOVqYMq1VU/V4XinxIGWe6yhZAC7TBvENB4AWJBowQdYdeboaXvYAAAA0AZ/NakR/h6Avf/WMAaiuyJhL+KT4JwkHNgd10OH2kTFHZuZ+4In/rkqbThwoea2xhOS5EQAAADhBm9JJ4Q8mUwI3//vrPjgO8x+Iwxa5AW96PE6JjAAg9tvdbYLxnthXvYjAICsZOG1tuY+xvSOzEAAAAEBBn/BFETyffT2/FQo8uLN+ZccYdXGUegcTqm+73SyiRqyxjBgiXqvyckNmzTA/sqcKN6IxF8rB8AB6aJm1tPGRAAAALQGeD3REf4MOYBiW8DO/c0klwuDU2bHJ6zPTXTSFXKSsxC0nU39JMiqbRsiRQQAAAB0BnhFqRH+BB5VXHdOCXHYpVjQiP78rXDhs2RTqFQAAAHNBmhNJqEFomUwI3/vhUdg1hjtBgEPsnnqReMCf9fk0pXkAepqW0GvzkT/N7d++eAjFF0LqEWTOQwynb70MqWR2LVAT3BGuJjMDsrfL3saQZmSV6KVXH2F2scQ1murg2f6hRdm+IshdKI4OTJ0vpuACyrZZAAAAWkGaNEnhClJlMCN/+/6N/USMh05oLuDoAZwruKxif239tQ1VJ4v5ZqpArC08Gz4PxiK1Zw3gJNSO9N1MaBrO4NXyS0Kk5yJ53TqLyQyUa3YaV6o9CrHZ784vgAAAAGRBmlhJ4Q6JlMCN//vDcjlriweTEJ1e1T49rsbXnEYWTbUIkOTqF9p/L4Ed+sSuPC4msb/sWIIL1Bg3A0EKKr8CzywgYo2G0BZdMZXDAvTPj7Xu2U8kKOf7J0Nmp5Nxc/b1wCJeAAAAPEGedkURPJ99PbKA1hlDv4p1vxOs/Y4zfPg6iHF0NVmxylt2/uz5ddxtZSf9odYXHgkvwvq5a8FEnAsAeQAAACgBnpV0RH+AbCrDaMSfDlBTuXnhXQUYGq59CbNAd/8CU+K1jwZOGWMmAAAARAGel2pEf4egL3/1jfMenU+Ipl63RDQu+raphKcMtMgItiKQVS5cjIUWjl5mHHdx54Jm73xhCUgr2g10y3zkK7P29MDQAAAAYUGanEmoQWiZTAjfyghBljfBFTHet0U3BKQoBgmH8wmCWbb2fhS9kf8YrXaOzQFjOPpvXVSbYEH7/91upJzbcIGzTBnw9iN8+sz32CUkThqoo1kDvR6ExRFTQPVt63gM1akAAABKQZ66RREsn309tLhK5ErqsM8w/NBncxs8+XNMfQtyU397Ah/MDPI+uVDBGq9NMLyY/eg7uJexiq401plodH0NOZVtS74swPLgbj0AAAAnAZ7ZdER/gwyVSt4Gd93RvQzgj5K3CnQvm7OjBKD3UgCK1+lg5LhBAAAAJwGe22pEf4EHlYPO0DSNXbXCGQv51Mgmje+6XpjUxwk1173/hFbKQAAAAIJBmsBJqEFsmUwI3/vhUdgSGsf7YM42S+/vjzG31Ii8r+yG9Z2rwtMeopT4f1ZXmgCKDIL/jmrbYA8wR4mFw0PT9NVLBPsNXrKLeAacgNRksycpYqdXIV7VH+6jtYtcIHq7VaKAVYULBW7vO3pgC3A4br5z1oXvv7pxKzVVrdDydj0wAAAASEGe/kUVLJ9/0qcZp1xS4Q6K/byHOdoBIV6Xps/UwNkM8yttvswZ8Jzr9I6Yffad0OJ7Pt/vu8H/euA9uoco+FMMpzX4CWnWQQAAACQBnx10RH+HvdEvcu4XZBGgCHIf5/KTqwJP2rRK48Dsj9hXDpIAAAAlAZ8fakR/gQeVxSFc8i68nfap4V2gL3VVnIbk0JUIrK52lB5NkAAAAFpBmwFJqEFsmUwI3/vBm7tPFgQM7JWZsrd+AFcbR02voqglcUnWYy1KjlCBoWGrB4/yDkcsF/yfA0JiredR3d1/zQTCNN3YSNi3+6GLuK5XeureNSQvNWSNTKEAAABLQZsiSeEKUmUwI3/76z44CE3mgfD3z/rIASyI/4ajysO1TyQ+cKUy8BCEPXwFvV7Q1kwh0NGs/guw48vBxqodJo0my+qYZKZwCVTfAAAAXkGbQ0nhDomUwI3/++5XDgIGEb1mgVp4z2caIApSFuXOZMOWUKsM6uZXxOHgwgvvNbQRkVIuaLICKhGo4osmKLu5nHbNQ2eBpkjJ0tBuuGJ+8FlhcF1OznrI2UofcQ0AAABnQZtkSeEPJlMCN//7xIFmvqIPwLXapFLrpy24N2AYpX55tMIyEhwIGCCAWtX4dRhVjqoPHEX7ZHYi/+5mDJL6QpKDldqAw59qEcj7XmhguiYHg7hxSzxJYABsNFnTFdhSjxeeeQ3HwQAAAFdBm4VJ4Q8mUwI3//vrhHYIVisLS3Cz75CdmPBNSQtJmyjSkVk8S6O1d/mV7n8Rrd+6LOzS3sEdBdAhEdUoK/3HLiqSbIKyHxx0QL+V1ET6o51gLYhrc+AAAAA3QZumSeEPJlMCN//74DjgJBsj5aW27b4bt8+9bnwCqvg57ian/4Crkwz302lGZ2Pv+rWZcpSzgQAAAHVBm8dJ4Q8mUwI3//1jexkG5DGfTzn6ACkLgiMo7PNL+8MNWKBuY2ww9P4WA81P2z1Uo6HFLULVfe0UWARzjI9fYO1TGuDiw4iAYUYp62du9UL2+zNI7rHN267ZkwLWg++4sTNfEBlf+9nqtqNbpC7ef2VGPncAAABVQZvoSeEPJlMCN//74D1EH3TVIcIrzH0/chzRECMo8UYdBgVM6BUQ/qnmrHcIvQI2T4gwNuNCHOfc8xOvd8e4RUUrEg7UZrrwxAx7PnW70QrnAkn3gQAAAEZBmglJ4Q8mUwI3//wDwzHAd5j7ZQY2ZrSRVRGS4RBORnR2QxEztLNjTbcTvgLcPBn1MFP2hKUm69thjXP4Z8ejw6yLt7eBAAAAPEGaKknhDyZTAjf/++AcgTWtokf5L9MFh6k1zvz+EhUSCxgLY7sBvcK7i6dPmedsqE34PPj3OdSB7bkjgAAAAFJBmktJ4Q8mUwI3//vgPUQhuJTLhEJqu+Apna/+hU4PzU02rtLQFGBiYXaJMHEktyfr4sDAJw5gQgiM1vOUDKdw5h9IrJLuyeXoGbf4wrVN09r3AAAAWkGabEnhDyZTAjf/+8N7/LJcoi6iLVaYzKjcM3p9XtxS6Kwh3zS27UOOJokA10Oalb5Mh9jTjqiJ2mZpX4f6frTs95x4EQuNnHB2EuOeRNjdeN33AYWtQV+GtQAAAFZBmo9J4Q8mUwI3//wj5K7wgzjvuFjuy3LAdQ+tnajIU2FYf2/S8asjPFa8Heexf9iyg+O2t1pIIQe9N50utnzQAT4GusPtoeQwLMtTfJDIuk4WzhVsVQAAACpBnq1FETzffVS1iQAjw4mKGWZ99z4TIUbS/mDFOaQ2BzRPfOEJhO2fZfkAAAA0AZ7OakR/h7ESVYA0mDs38HcgN3BlSvTaG65b/g/pfq94yFXiLsMlUCZdl+Sj6EnesriDgAAAAJxBmtFJqEFomUwU8b/9Y3sZAzW9hPGahGw7flXYE4FP6yvCNzmXPKd9sNuacBlbT71wo0upN8/BkEcSxj5QOSyFbOl4aFdbwIN/85qZNxlQT3WmJ0lyxLEqm3GKGmR/Q9dLRwbnvK1UjtfHRqg/kLxCsjH+LVN9IFW92ci8U6z0Su8pzeXhOXBhoqHLd/KR3LzdJ3Q+/pqGJRbSf5EAAAApAZ7wakR/i8dTtuWwJ99n5vKBTInol2bsI01rGqiIrOoR71X4Q081QcEAAABXQZrzSeEKUmUwUsb/++A44QawOdhgfUi1MfHqmy+I4SJvb7slhfqvjc6dx1/LRyTOBiepLebRHg/6uB0qH06DelPtYLjqPVsidGv4WOqaIDxCE334DXtxAAAARQGfEmpEf4exE2/rMauJo+rnrnx/+SKnih1QnG/vzvVzCJmpdIrzItT2aVilxCxecVUkmU6Xcdbiwv1Hi2eVfOYva6XMcQAAAHlBmxdJ4Q6JlMCN//wDwzHCXNMnztovd14vJ625TIq7dg7bjOkJiogvd97B3i0yj/I1tMjzhb+H4NpywxTMQVOk9VS90im936odpBgEAln5etaGgewn7zxeYOtk/lsAB0UYEr26ytGYTnykyW97BDgbmncu1z0/tcTAAAAAU0GfNUUVPJ99PbdJpKoJiBXMhGSlyCkS3qgHnDWoGiC2KtexepnugaGZOIX0UYhHhpVdynfptSF4E2r0Uv58jiFtJmUBvNE2oOgYwxDQ9IZDYeOAAAAAOgGfVHREf4n36qs07VZIPxf1t+QUqlPNBFREIxt7hE7ahK0iE85r2mNsOC5htsIgigROAWVjVceCw4EAAABSAZ9WakR/gQeKxIeCDSxZxf39Gix/112nHFC1PzbN92IjXL5eOaAOw4p3LVmqz/XDkVqmVp/QwtqVFTXi1vi4EzWJjc3Nfa1Z7l65jvwwWZi0CQAAAKtBm1tJqEFomUwI3/vhUdiXOx2MTspDBHmUaNxnAKtxo1ZidUG16yzOqQ6r6rwou8sSCWAoHPh0R0076A8s6id90rZISEbQSI4RHOQA08yjeqA8L6RD2GtXfiIoCm2bQHe8hkBL3gMquBV9xsgOcsX8KXQlQ4myFK/fF4oEBZAQrHDDSzotzZCZOAyfnfF5fPIXTv5TGCrdBxcbOAJoIYdUVUA1wPS5ma+yJ9AAAAA+QZ95RREsn3sgp5rjMV8XzVP+pp1YxygYLRWB2wZHB+pWeUU3ZYrSXAwe2bahlT/6TYFQTJt0HNTpEP3ubVcAAABFAZ+YdER/inEfdryR+UBQh0g+98kl5ilCOtNo1YP/LuQq8RtpQSglo/BF1WPBxIxyIr5KAxwIpISuyR2pURhW4rUvLdfxAAAAMwGfmmpEf4jgOwqw66AXnAyisbuurz/RSCoRjyFbk5dILFFLRMyoUi9I/yQ284snWqurXQAAAHlBm55JqEFsmUwI3/vhU4AmHk2jvpp39+DVM7fgz+/0mDAgC9sSz81+siWGpUWfuj4GyTOZUOB35xF0hF67DNn0ZuMEXcGIXz00pKtOL3sBXsUgc/ePA52/EzhUAvLLgjOyGZLNGqxerZSlwTR99vdiMIycknIR/ViQAAAAUkGfvEUVLN+BtiMx77QlavkuKJHZZ9Q3HsrRt1QLIHVTAm1R73+pWNsjQ0q58gETaMA5QLuokf/hrA/s22hdm3istSVmd0URqKtuvuIfwOuQF+IAAAAzAZ/dakR/gQeVxSFc8gfgZAkdinyst3l5uX7pQaHSAIaHm6zNCuVavCVTgW8m03yN9WRhAAAAnEGbwkmoQWyZTAjf+8SBZrxwNxE6nKekPQnYXJm5wBvfHrogTYU3NIBfehDfOIiS6ySK3363CBAuzRu6fA4vIPVeDAL86ZR4LtBStB/BM1+Vq+8tYxISLhvBnW8Dyw05HPMDTibFdFLhiDfH6SCO/cg2NYHDixZKTHvuNNlOBA8Eb/AVnrFkz+v5UJdJl1nEtDkUt9F4X7Ko52NhMAAAAEtBn+BFFSyffZxiPbso3KVMv0xrg/tKL0xl7/2ug2XFt9Z/8yNSEkmzVe4IR1kQbVaS6gPVxGBlsxGiPdVg0hy8V1fpzap1S8LNa54AAABQAZ4fdER/h7ECGy+Mf8HMSAoaH/GRCeAOPoTHrElxuxk6cLhDnYyUbEMAwSiUszmoA/R1gKCdAprLE4LysACZvlN3jo3Q22r11pJc7VqKCUUAAAA+AZ4BakR/h60oZYAy4R2w8B/3Wta04n3kCJAW7+Q2gBRN+jJHYMfX5uQDuUKfZkFiLp0cf+1j/fb/e24j8dQAAABxQZoESahBbJlMFExv++5XDgdzj8Rhi1ZsIZ8Sm28CkaHPAYRK2SFSACU5cWdipBQ+vrEx0RbW5ZoDsg0jJR++lEuFdaMABu0GnPbEfBPglHmsDhwytPRmDnxl0gvYO3LQZM1n54PZOnfSIOyUJsG8mYEAAAAjAZ4jakR/gQeVg87ypn6p/FGQB5EJfWf6QKy0B/hS5J4dn2EAAAB8QZooSeEKUmUwI3/74VHYRK6xsWuLkzkMFi5kUGmAexj8Ar3bWWH098HYvsoR3J76TgTz/oDQV5cpDVqE1+A6h12YunuTYfth0SsPDgk0pLNX7k4axqHCb90xZXJ6QQE0rqZvxK94nl6XIz8bJURBNnwVSR2U9urp0GW5iQAAAE5BnkZFNEyffRPrpUzJLr2/AS2nxOSJCtdNcLjn6cXpPNkqjvg+oO3azFJERI3eaUINnWDk1Poc6rrAogWMSwS1ixlRAlvzi6VgIwdfokMAAAAvAZ5ldER/h73RL3LuITCI3oDv4BPtzSRBT0ox9Zs6XKAiQZAu3sncW22aE5GNUZQAAAAzAZ5nakR/gRHB6GWMSe17aOt16nxs4vMoxLGJ6z8ZwihTVdCrHs/gv2FN45XH5F77G7EvAAABGkGabEmoQWiZTAjfzjAzJNh9+MsHWu8d7Cs0acVk7TvkhyFGrIZeP7VKZT8YSa+fqLmJANOx+HwIQKZ+hbChanD7CSavKxfEpeCd3O875McKYrY5v5SCqAR4LZnb/unYHbZG6CJQLkdmdXI3S9mC+vh3J1TAPWl+HZzfy6FGRPZ20VqNe1/AuRmfBkIEA4iz9PfeZIPIT/nIKX4skpc+PdUp5tjAsBGG8DV35+BHBUfG6nHmCnw+kVYk4zIk4Xclt1NYI/JiZT+WqrMgnGmO1KB+XctzpeUXBmT7ef+PRBXYvDjeBflA3MnR1+dcK7H8tqUmzVSB1TwBI9+qN2nnuuRbqrb0lST+s59Y8w9uzAlwZZ2ZsKrcj1TJcQAAAE9BnopFESyfzVU/m33U0DGS8eDHAR1n/6XVZZ9uylPsAlipuRHFF1836S7pQx8uES23pt0I4STTpIPdian6xejkAYoF2nWj27H7EQ1w3xZAAAAAawGeqXREf9Pb2DkjtB4in5wCGYpcxy45arMPjM909IA9gzca/O+GyZGg7wkOdd+VfzvhhMt5NvbfAVFmfDXW5Jmj20M+nvxZ2Z8uZoUmYJ/Ok7BKfgCSaQdQQUP1bKE4bTtcHvsCHyMnF6keAAAAKQGeq2pEf4Z1uDH/duUzFC4JTU6hTUz3R+lQw5ZFRkiy5ri+QAklUBnDAAAAOkGarkmoQWyZTBRMb/kp5kxTMHZmETVTIgsd+A+NbgreHIdnNX9tCFzUmdk/+DL0/JDoveQNd/vCsbAAAAAgAZ7NakR/gQeVg87ypW2l15M2hNNzDWakH6ukSDuxzYEAAACpQZrSSeEKUmUwI3/74VHYPZsdVGEb9i0qy4XOd49kR/Zd27xbxgdkhHGHJC5y1l2os18hgqRE1KaZVj5VSPxxYMpx681OhYazthcfQAXvfCWuIYtRdInpGWzy1sxMLVqgPWw2rPBeteX3bwTAuESDIworEPEl9lhf6OwWpk0mu9bIKNzmWd7gpQ9E8pnmi/jBUgxN2yeN4FcSjUT9R9wbJD4lX/zQk/2MtAAAAD1BnvBFNEyffqpK6p9FFWrePWQt+t023fow08TUcBgf6pBXtRGOvd6vNalSElK/B+orsc6pxa+DU1KIEsYhAAAAJgGfD3REf4e904PVduAufu3efP2dUs6dm0IH8JDwzh2JpxWrp7PBAAAAKAGfEWpEf4EHlYbRiSTrye8J7wrvIAVKDRGHYha2PjTkam//yFWsIsEAAACbQZsWSahBaJlMCN/Hrv6DGHttBKhEf7oitzULKk/UkvXrRGN5TSj1fcIGe9X0COkbiUIPuCUmvFu3j4oaIwdDOZcC6YE657YPavN92lNXaCH3mYlGVn/ZhGTMtcl+H3Znx2tzw6OnG8LWtRLGegqT+9PKN37VFH+g0bBpNjO5UwkVHcgMy/NYkLLWLbOHn2+rwVZaqtmV1BLeu68AAABUQZ80RREsn81bxhrs9aLlZoqp8iOSssZPF6+vuh4reMnAu5tYCzgCviZbUJoMuyNDyQ5sQrG7eozVDcPAOhyEE7YCkUrWX17txueCEEo85uzZ+398AAAAOwGfU3REf57MlpebTfh/uLfAQTkb6fSvSijofFdaiiHHloNZyWGn5J4eiSHd2ZAxcNxMqZvRdlvUiUYyAAAAKAGfVWpEf4Z1uDH/dUm3pglXxGoCPaXPuwJNhybnmpjlHL8vPo6GCD0AAABbQZtYSahBbJlMFExv++5XDgIGEb7MdX9/tJBBP77OrDPDuf6bRmi2cH4plgeyDD0btjDKcl9nj3VzhH2Dh3jpoNa7QLLWuVUen8M08hu978wsuy3OhNJ1Sm5jUgAAAB0Bn3dqRH+BB5WDztJXeH+x0rqoojqtJRdUUKnIgAAAAIpBm3lJ4QpSZTAjf/vhUdge5Ym1P1V/GNX9AIT/jkz3GQqbKNPRZXoIj0uUAbDPtH3Rh3Ek+Bu+4pMcb704ggfye904H48vWEDu5uQg/N2uSES/KbFzLPFtQ2OTa/8qYBpX1/aD+T+6hJJK1ru+MZTghCGmNpSQXWnK2Hv1rX94UdXZenTBkKmXqgUAAABnQZuaSeEOiZTAjf/7/o39RD7fprUF/NAaOCr3pu1+Sc+O4fPzCXHhcvaqDp9XfBNrwZ00TugT9n/WmQBb8BI4LwStzpt0RUudpuht1/3oLhjOlCXs3ObsVj+wmRIxGocRDVdQfy+08QAAAGNBm7tJ4Q8mUwI3//vrPjgdzj859cWvgfbdn4AHc5U1oFwreoYsjtgTR9pWZhFMOPGk65+kAs4C+VojpHiEg82clYzEV+LFhJqn3NZwdE5ihwYUpFp2IKdxJaXvho2ltDn86KEAAABMQZvcSeEPJlMCN//74DjgiV1jEgudcPZjzgKu40Dxkog8ar133QofRc/AH/ZjqW6NQl1fHvbLAdPdJrabUaT/78/X6QmNkJeXNfxPuwAAAGZBm/9J4Q8mUwI3//vgOOB3R8fPH5+zlRCyD6mVH/aMulGQkckScEuTaTId2v705sWz0b6RhD6pVrrfZlw2bn4W4jUHEnaKMk72slzxowCIHDa2FHirxY2J4Aogy/1G71Xvqzs774AAAAA5QZ4dRRE834G9C6jp7AxObQ1D5nh6xcB4XoXLYDo06V5jfpQHI0j/RblPPLPhUSYgXonQeS6joQ+hAAAAKQGePmpEf4egOXq98g4uEQsyFmb1HLaC4k+C8XMK02ZavI+WCx5//WHAAAAAQkGaIUmoQWiZTBTxv/vgOOBCV6+DEJpLiNMEuxO6P5BRQh2K+vNT+Sg5udE4JTIIQLzSP0bwLFYQSRlbfZFEgNKAxAAAADwBnkBqRH+HoDXutdsAXaYM5d5s/ldvZ1ZUZqSPwmL03HPhv2wTpeNIW0aHyufZg9R3ewLwGGEJnNYKZsEAAABJQZpCSeEKUmUwI3/74ByBNNkgWomc9ochzMRrfD4ihgnoRApgkRdI+usVmu6Y8v3e0UScczsCjQjsNfzAoG7XkUz+H1yuPskDuAAAAGdBmmNJ4Q6JlMCN//vB61Ydg8gqaxeVjbcbFf2tmleGxIAMrv7NFEHar+eha3FuY/vTZzPKjJC2Oq5sKE5oNwUknilo4J+t1ZmDCAqKelikGKtqi7gdJP35m4Kn+yOr7aYtXrd6AdZBAAAAYEGahknhDyZTAjf/++A9RDqMbrkil10LEKl6Ttr7miP6kcL0ohDF7+vXevrzG+zTLebM+8F8Mqm1WPiQigJSclUc32sO1HgWuzAPWzc3pJHtUBg400W+ishyY5Uy/xEMmQAAAEpBnqRFETzfg/tuGYHKjS/jGBEY6YCRcT42dm5owW7iDaE6eZzNyRpt3yfTYmI4OxsgAkcyVuBGb+F3A/b7UbL/gPEwyEc0Vpa2lwAAADsBnsVqRH+BB5XFIVzyB9wJwDYp8rPxDk9qpTEARkTv8z3S1UGE8PcjATyV6Zz8GAZwGPHIMrJR6gwBSwAAAINBmspJqEFomUwI3/vBjyyXKEdXz5LJiOmapC7Dh+bqzRRczmQAYnOJBgYTvI6vsixxDeHUOwDiuti47UC86r13T09kIqPRmtJYNyAq1tP4/Ft/8pL1VrpczUxkwqL5a+AJdlYjtK4aB1EcbZFat2ii0Xh5BXeTFj4VS6DQ2G/nJLgwVwAAAEpBnuhFESyffGXEt8MjyOXZ1Br/nR/99FrXZZadcoA6YG0Znt1Y61Fwz4DJCwY26xDhlnhFmPYwjeDf0M+x30DUfLS/zFrIDmZ/MQAAAD4Bnwd0RH+HoC9/9YwBqK/PwA6Zil5xHhfqFkRqWZ06PPIdbHE9bbxRGEDxKdUaS9iieGaHN/QijO9ByAGQWQAAAD8BnwlqRH+Gdbgx/3Uy5jiTLpXihCuS/L2kShyyEDBZR05rxcfD2VGCwXsTh9INhGlzQp0UwAxUx1vZR1T6xUQAAABZQZsMSahBbJlMFExv/APDMcBAwje62AS64edylmitQLrKTHGHtcvoRS05r00OvVkJCNrMh+/mVw3eQxDA3590PKiZB3nWfrN5pDxzTrKYWq9sI5F4omOgfooAAAAnAZ8rakR/gQeVg87QMgUIqzoc+IQ0ccVZvftYC2gzOa0pG5p1idXxAAAAbEGbLUnhClJlMCN/++FR2BsEzfe0cZac+UJkdgjTB26Jig5W2j5cZ26NMmefLD0RqMng2y7laLakc6lHRdZNw1oHaHFsXiqzatXDN5Pz0v6/WwXUQ/oZM3evQzDrkzK2TYV/MjXfwp98+OlJJwAAAF9Bm05J4Q6JlMCN//vgPUQfdNSB0ivMgv12kewdPs8LDHEnccStTFGw2T9LqoRXFd4GA6Lmn/C0Ae+2Ys461r7jzbWD67darQQW6fpKD70itK+HYmaanX59v4Kvjz2TrQAAAFhBm29J4Q8mUwI3//wDwzHAd5nLgo0KbNa84paAodp1tGNxaXmvV+VvBFtu1np4hpixdq01/tsJuXnBDpI6pjgYmOuYM9uC1Czm+OIeAHPNqzeKqP5ocdj4AAAASUGbkEnhDyZTAjf/++AcgTTobJBaHKECxQ6c79GRuRfeQ0KwXvcejjr76PM0aqZOg6nHibX9Iua+A36w2f//ebnY1vV75tCTbJ8AAABpQZuxSeEPJlMCN//7wZu7TxYHfZLW0nepgrO++53cE5t5rTbH/j4CuZw0tqBjVHvT/NuCxYAqZPHuW+4z/LgYIOyIRCJSyOgmlLVczhDzE5KkFP/GXuQLXa8D+P/5fNEgPj1lTUpFVO4xAAAASkGb0knhDyZTAjf/yghD0KSEAj8fRNDXd8ZwkGDcdqAIE0puIoVKKR11TsEWfh9IlRceiw1biBwH50bd5SVeqveaOBsfxwvHo+S/AAAASUGb80nhDyZTAjf/+8SB2aRkB8IFrk6Yu0T79Zx2yWX26LbBct6B2P7nMJlSNU+8im32OIA82qWr6M0WXr/wOnu9asEIe+TzxoEAAABgQZoVSeEPJlMFETxv++A44EJXr11WI/mWGmgjspyywfThqWu6vt3f6IeutLnnxyOVezpxz66ZO+XTJSNz+hSyTccsijssaFateiNeQeLGmXnXJq3lh1/FYw7UzWsF4lLwAAAALQGeNGpEf4exElW+Y0oRl/FZQhP+oVccc8fw6tVIVhtOXYZd8Y1R9+768Me1CgAAACxBmjZJ4Q8mUwI3//vgHIE02R83+/Vlx6PE7gcgcvw/dfZOmqjB/6fbU3T8uQAAAFZBmldJ4Q8mUwI3//vhUdg4BM3dCQn71waETfG5XS6Y6qV1OmJ2ZI1/p1k7hngqfLDRlYxzHulMp52MY33dNDvA4xKrjGSUX3xOzqdXe+MisKohr4vpgAAAAH5BmnhJ4Q8mUwI3//v+jf1EjId8qaxOH+NL8HxQ2fp5PizkRPrB5xmQS+tbXXrCt7darNka9SvAVPptUtLJ5w5YyvLUX1kGEExPUgNHTM6K6C2ru/0HOEs9PXbEuaYln/GuFzif4kjrgcvg+8JdtRNs5H/bTBw0pf0EAsoCf7sAAABQQZqZSeEPJlMCN//74D1EgTaKZbqNH0joJJniUx8Ruqc4abzu2rURnTM10DKgg9fFoJmD440LTpJLO9oCBujuWU7mA65jaO0wi7HremRTBMEAAAA4QZq6SeEPJlMCN//74VHYnWWrD3GQIHk5ghig/U3DKJNeEKm/nL7OOuV9SKp1OrFdiEZQH1cmsNEAAABQQZreSeEPJlMCN//7wcOsoiIc/4xpRxHtkU2H2faBzAZluyLYQW6IF6Sg2Z8+2s6OiVvP7unG/PaKjnvHx284RVK7RkYgtGcnC+eIrXj0kIEAAABCQZ78RRE8n309jmEtq7hLzM71lTx1/0xhOjUhRbpoPBZse8D9DOokAeCvfYnyUFigFVOpER0UjkrFfP8CohPIs1WAAAAAMAGfG3REf4egL3/U68Y/kJfuO5of8ZPB+vRdik27eJJcU3fOYgVo/ofNZHen2wm64AAAADIBnx1qRH+GdcG6e+81iB3HSdBA1qVClucbNobVh3bfiI53Yzn0y9s0dcFwAHbyPWxxUQAAAItBmwFJqEFomUwI3/vD26eZCCisgtY7Q4/nAy8diTjLXaDcRhbm6PzfosfVpvyYRHRlUjwp0BnQ/NWsB7lZrV0Z2KMWnG+8gj/2wN3ueG8T8c99L1NBE6r4v7E7tCr/OznxbvNYjxfXhceW8dti/xTVY8P4CvNMCqkTDHLfY5PChguysyasokCazUmXAAAAKkGfP0URLN+DOtOgnSjGpTM5Wdix7kUjmGwL/wc5dr6DxK39Y49jtKsbgAAAACgBn0BqRH+JNh+VsnD/m3MN5x6Sg5Qu8+nj/5HUWhu2O+TNUelAcVFhAAAAlkGbRUmoQWyZTAjfyghC9V+ELITH3LCIHSIvm9CG6OjHlL/31AO68jyrK9IXLK5IdMok/VYaXk303UZ7pRLjQA0f/4qS5bfqL5QiXvH0ffwKCt0Fq2vBbWwZIOms8XXulZOVX3znJGj/solkcNIBtwjIWnGGvWEb1puQ7+6i7LeJ9OTdMX37A9Yfx9MPio17AaTThn/g4AAAAEdBn2NFFSyfeyCltp/KzZHjCx/v+MwPlO+6S0X+qGwMBBl7Zd3aGLZ6jXXcpABaEUHQQndnlecT0CSIxETMb0IEO8QcA68axwAAAEUBn4J0RH+HsPlaZzJLf1BqA0GQv7XWJY6mi5DH26gpsyq8ggTdiqyPYNRnn3ZbbfzdoybbOog2JQKGI36v4G9OFjBq7XEAAAA3AZ+EakR/hNQrQyxiVMaSwCvWnQsw2ksmBu2UpY5bkEhGHFQrC+fPFAFNp4o85UVgqnxqzlGAgAAAAI1Bm4dJqEFsmUwUTG/8A8MxwlzCIUaKgXMXLU3QQJ1EpMKnk54cIKtfRAathRBVUcjjYAmeqTvTihh6Nuk08XkauvWPadL4Vijodg4AVC739ye8e1iwE/NwJHDIQv/TOUzJWUNQojoW8MDz68svtIvZyoIhTX9vvxCDlonewBE5q3SspGmeeozvc/Ou3OkAAABDAZ+makR/h7DkMU2l6PKcHO/zShtHBIsYylwPXYPe5TqpRiePe1Ts30NzxAzEqW4PxaI5KV/5p/0v4lsSu94p/NesEAAAAMlBm6tJ4QpSZTAjf/1jexkZYO6eHpSGJJB5secePzct/DA48qcDfs+3F2H5KAOIzt2Dgt7ew96sGOjWUYhosJfHZkPXPvBEDaNms1hmlF9yL7QrVVstFpOhrOq4+qlmyqh/ldeLSeBIR0MDaNW2T6lc8cDTyqgcYV+Id6ev9tWnX/Pg/+/3Bx9Euhw2lfQiFdanPgzmX/55qCVEPMNFYpIBwPcaLU3yfrivjcJks8gRYsIJlknW6guCjN+neZa8hQkWtZwBla0sYTkAAAA/QZ/JRTRMn4fNy4UC4h+oc2BifKvWVgEVteOt6GdcPcqYcktVsScBfm/exmQJXfl0JC/xI/s+EI2ZfJUQBohxAAAANQGf6HREf4ew+GS+7UpFaHk/4Emf158OcWu01ywvQyMU5k4T2GGCK0QXFWRrnznHAUmOJaj8AAAAMwGf6mpEf5Kigzx6E13On1ZVie1EO84nBtPPa0VWD3aS/Ddk0ZkbJz8T95L6crCgsixKpAAAAGVBm+5JqEFomUwI3/vgOOEGsDnbEfsgdrJ7tKwX8nHR83gyULQVhL5x+AgWYfj7sj53F1FklU/dfB0NntEcc8LAeFRHrF4HF/lXkKKa3TzSgRuUANQcRW1CFfPMX5gBxfeWFsPEkQAAAFlBngxFESzffU2ptw5lYiNEGv798jp2yP5EhQTY9tWckQ2I9YpqN5ieaWAqdGE6ASw40+H7sZMawfHoCoDply5HZaD4Q9TdMJXThHdVh2fQNrSDXLaofGcFLAAAACYBni1qRH+BB5WyGKseqwNDLE8nxiM3gCXD2Qs8x6pTnxeKoYIGqQAAAHdBmjJJqEFsmUwI38oIOGvYtN1rSBTzyT+MoTaZxH/5QP0RwoqkIo8BIYjA+yAn/0EDCt8e4sfiSZAjYvAKxVNatVPYvX26/q3BEA+lCvg82wYomT1/sm0oujidy0OvyL0i4I1stZMMRcaxZssLhgLuUf5qu0yhYAAAAGNBnlBFFSyfhIT8MY1eZNHsRSE0h+4P5gINzsuqj6GnC3epFTfnf17zrDrtInbEItFMN401WbhoZ5/w3Ux2r08lE8/dDeO7O+HJXJlcFYNAFuyz8daGCaCCNey5dBaMTuealfcAAAAxAZ5vdER/h6Avf/WN8x6dT4hVMDgeRaCIEr6D0AY2ibt4Z5BJQpnXAfZG9fF7OQKPEwAAACoBnnFqRH+Gdbgx/3bkhObG4+6VBRfnLnVvz5ibbBPbDI7nFtBHdbdYGYEAAAB+QZp0SahBbJlMFExv/R5j/x5hxQqCYDiX9FqzOLs9ly1QT91Cp7pRDMF2Jkxd2wtt+BpFU6ZBrmxwLfJiwLCcv9T3lq9NZ2wLCT+/wnYYaZFCO5IgouVXlejqeKxfunKWRud5TjL7EHetewgMnoQTi+FFP1HpNOyNxxWv3UwhAAAALAGek2pEf4EHlYPO0DFagf010Xgc4NtyHA6rTfxLOyTg0Ox5hTe92mTjN7boAAAAg0GalUnhClJlMCN//SaD/44ef0/kxstvutqsDDhlXYfdekHXbmlDX82+t1ItqO3M/ML44EVCP/8ffQY6CAnVY1XBN5ZbnYha3+46HPtD5SMAfNYPV3RND0lVvOS15Fm1CdExPBODnBhUkaWtEAbcewQEoOOq2Hen6odimOMzFQ5d9rZgAAAAg0GauUnhDomUwI3//WKRu6nVyVxRee7XpiIxlKBqJvLZu3azo2YYPGF8qjUhlKiymApjcRRZlTf7Nef6tukffI6Kd7PrbAcHwGvELaeaP0ibxo0QBjNurmKmZaKkIE6v7wE4Wpmgpar7VBmLnzu8twOoVcMvF0fNF0fSaxrmXamWCb3xAAAAXEGe10URPJ99biH2oflZsxy6BAxf/qvJQn0UmNGcVBNMrhzZROK48G5OjPWtBqH8efnPpu4XSt3r1R6HAtsxDtSSxVrh6V0/kLJR05IiNa/ZWgvQYnGCF9qmzKBvAAAARAGe9nREf4l57VUKFwSmPSmlWex1Pm065LEm1OKxbJb7BMmnONzBDc+f4xWPkczvzmTohRN3zIL6Ojmccwe1SM5TTAzFAAAAPAGe+GpEf4TUK0MsYlIRurtl/BL2J+PMXABKR7W0r/E6t29wkdoiEyOny9hdlQSa5XVyIlDO89Cpm0EBTQAAAGNBmvtJqEFomUwU8b/8A8MxwlzCHtLBjwzZdybRP3L28+w3h5H9q83W+4R7Iqot5pTdotVf866A50TMubMMKbl+05ylorAZNPpFCLdQn1pTLuSgnmK11G/Uu/D96ElwYRQfS58AAAAzAZ8aakR/h7DkMUOCGBJh+lPcfvBfkdYcocwBe9qmP8DL4C3VmEv6j6IDX8OhLFl8JqBdAAAAsEGbH0nhClJlMCN/++FR2Inm+ENt5odDIwW/GJwkPcJuckK0JIaAU3CdS/tVTXZvLR05nQOjqvn8+KStPWBFi3Gez7OgeSNwcVElqJE1rP7gdMNhbcN68HuvsKSeQvpaJEnYw8wvyoh/QW1CYyoGlOno9K1aJIr9Pk0Dj8g8+NGIQjaqBemVGjXOrl9UV8iu8KodQ7jGbm1Zq/jwidOehhp797VsqpvsPYf6ZBDtDwWEAAAAYkGfPUU0TJ97IKea4yjH9NU+Ki3CHMkcSJ9htM0rapvgZK16qh6u5dqOGGAOrakaQCu114AMCSgb/V52v4CxfpXrn5qv0WiKgUPDhAE/bLxIRa7EHZrz32McwKF2qrLBShrvAAAAVQGfXHREf4ew+GSYyDhPzPlV2bRuCDR/Fea7xly0Zv+RhCiwpzapFX2nVaD4D0ZNqSzT0q6pbaI4DoNcxC/GWSw4XfnszHz8xeBlADF65GFV4DR+M3cAAAAoAZ9eakR/h417ZNpp1xtu+fafR+HC0KcS40oDskcFqn/O0sQuxpu7lwAAAK1Bm0NJqEFomUwI3/vBm7tPFg9jGVWovZ0jmu/rAuAvcW7ubK01SqdOcWBIz02mXZObB/CLB6HlDKDI+VlNhCg/6zrUx+GcHA2xWFHBAfdbhNgkXIqJBYT8FRKoB8Q1unN40IlItyXayDFwzBeGhvGI/KN3ea1R5uXGk1x7p1Y9/VGM6iUzPd/IwqqDaaStukEQBdqZPJMfq1T+2RPjqIFRCfd5FvxLcBRYuHwYlgAAAFFBn2FFESyfeyCltp/KzZjl0CBi/IEonyzMc3bZzsIWT4CydiIbfPye6ei9io+At/mFoyePgIN2xEzlQ9KswTRayQNCfgryEsw6EmQWr1T0ADMAAABBAZ+AdER/iLlSgg3Uqa5kstRvP/YJiqA7KR90cw3LNvcyMbNM10dPyQ7u7ZQ7Pxn1vA7iWSg71Q2/hzofKx8Il0AAAAA3AZ+CakR/hNQrVcoxKlrmAEcm59sbNzD40gdHC0EziACGmA+UjRF1XzhaI48Z56F0dPEF2w0kgQAAAGRBm4VJqEFsmUwUTG/77lcOEuaZPnc+jWH7FzzFeNiG18puHuQh5C/ueG8rR5oKuQk08AWgzZtU8fG+9U3UaZPtsL8BReEOwTp6VR6Dx/WI+9cQyMaZkUwAXj5v/UdnByly7VU5AAAAOAGfpGpEf4ew5DFEIIgTtjL/DAkM6khmBhIgRPsTh1KkB2WPl3dHMh6FUA7iA9fty/cEHvdpGr2AAAAApEGbqUnhClJlMCN/++FR2JyyC0hpnDh7PIxPg6BMTLhBL1SK8o0CDfiMpJT4kRDl9tdZ62EidOxdhN0HnNX76b7eC69SCDzmuZwOgLR7EIjB4kgvtAS2pqCJfoNyhrvP0D+dIb+QRSH+77hmtbfE1K1QnwSPb3D56hVITusU1SKKSHGmPML+xyt5i7yztzLCuUxQ6GBMKC1lr+K5ma8CRUoFzeYFAAAAOUGfx0U0TJ9/cI6jT+mdVbvBbfqe0kP70WvSpDFVICoTveNkI3cjDw4DDjaPNgfOu7y52O5KIIit+AAAADwBn+Z0RH+Achjb2/3sCtNWp7oUGbK5uD2avL6oCzA40KUrmSqFJOSULMgMJ/dZcQX7xKN4kdd8aG2wieEAAAA2AZ/oakR/hNfANpNdzqjAmab/E8ax6Fvug8RPvYSSLNwtDqJx4K2Zi3etFf/xJFtiuMoa2OKpAAAAZ0Gb7UmoQWiZTAjf+8Gbu08WD2MZVh8WrSpRiUBZUWy+1Z3OOJ3h3uYEjEFhMU8BzaG88lQfmplKnzW9sLz44+YZ5qjPn5f/4kD07E+kjAoFn0kEueWpFWqaK18eV3GeAKJtOJ2gcdIAAABLQZ4LRREsn31uIfah+Vvv5ElY/3/IpzyyoeDu2znc3BzxpsJ+RQ8j5+UBxv8Bc+yc4/gYr85cjlE67KN8CH6/crTj4W8FigNzW7tQAAAANwGeKnREf4Bw4zBVesMI5kvH//3lWiD0u7E7ycYGE39qDZBe9Ez4XiaOH2ZBiU+cYRKmcZPbFu0AAAAsAZ4sakR/hNQrNtFc9WisO+KwtEafX8vb6ZXxEB5GdrbfwSZkZhFQA45ylQkAAAD2QZovSahBbJlMFExv02+yx/CHKbBQXEKSiQ8/bO2xeS1HQBGRs5Hf/V0hpgC19PEbrC94YtsEONo7PCTkPWJcLGlhvRVXDsvhZchHblbTte0GhtaJBAsYRQtNVaJ8kzaeIHkxTXGujtsMr89Zq9hCEObGEdp9ad+G+69ZLh+FJjSEVgOkHkH0PEpLJYyFAvBlFP2jIHUVm3qm3jb97P4pJCxeSwXqGAfuihTU/RGSLNPwJun8vR1Ksa+vR5Zon3K2Nw9i2s9ZxgoBryAZB90K/m6pC4pSyLW9BsQ5Qz/96LFgNPiTaKC6ucaXERHMEYPOW9aNn7YVAAAAPAGeTmpEf9wncM1iQ9Gl0i1OL+ABc9ME6axttF81n1bVRz8hEXRQGH4rWzfbm5v3nB61lX8Ctgc7vz9QgAAAAM1BmlFJ4QpSZTBSxv/CIPOnQv6FGt8l3K0kX55Ov5JRlQ03g0N8vQOD5sBbfTPpEWH0nyx5HL+fie+qID6xiOC1xEwq8zJdI2VFblTxxBMPSypkRcUoWHmHXCxUKRVZYFMjeGgiRV+6kZS1BgU0P9oD/Oidv7irP/aV0V26J1DBG2+EtBjS4lHyivwXdip0d1BQfBGsk1ugFW2Fy8/EjyelHeTPHBgZ943+z6bcQ1EJkW0tMaad2fGspFHI0IbGKScCPWX/jQ6iT/ojIy4xAAAAOQGecGpEf83u2UnV+7UnY1O5/oQn/Y6AF0MpXYTzrVvtUezCnzx+JtHpSvytH38dTVSqrrE8zYKZ6QAAAHhBmnNJ4Q6JlMFExv/0LB2GXt39+D3Cj7aj3DPZ6Pd+InbEtW1V3J2UHsXT1v31OO4EXiBsko5smAXJphYrW0fTkQJirwR8BXL92Ei4QvU8E/9ALgLDiPsylWFhiHKkipiOyh9H2qwcv9t+CUifmj23yzrIU/BAu0kAAAAkAZ6SakR/iOA7CrirItMfJpRo9cnv5zDAP6itv+Og6IlbYQGxAAAAnEGalknhDyZTAjf/yfc+3zQHuCTSl/U7GZJy/gS5f9Cy2h3VQBHdy17hrU6xaDgmNvuPdO+3nvolJXmlwIAkxED8Lx4XeaN4VE46PRHNsXwgF9trLiwzDBPhYZ8V2tPERnsdMtjrFEmrEm+drqLgNlW/gCWzTU6jf35hzyjGb6uEJOQlcGPQpDxqrtzr8Nc8vp4B1BYjdvjrlCrq/gAAAENBnrRFETzf0DZ7oRVytJSzitxXMmJRFhcbeha5ual5F1I2q2YB+zAT4PLPb+JmJ0ARNc/AIxeBO0WVPQWFlcQBs+2QAAAALgGe1WpEf4EHlcUhXPIwsFKRCeT4xUX4F3Q/slkJ9DbNgs6Z0J5ivSJcp3D4p0EAAABWQZrZSahBaJlMCI8bzhd8dK3gZ46kOJN672WFTGEpwGtbDZDet5dEQMHNg5ySdNRqv151JHg5eF5RqsxWW7PejJfmZTCQ9Q1IN/bi5e+xWsYriMFYv8AAAAA5QZ73RREsR4etQctl8Y8CYRFjDEwJCzxPJgJwzTf+tolyI5ByCNErI5NnAqLN6VSiQ3zkMiK+t/rdAAAAKwGfGGpEf4EHi7cJ0aWsBD/fwiVyP3OFrw7+rNjWzuM2AuJ/SSLVjrYN/7EAAAT1ZYiCAAn/vgy3mIDfV2I+c4OQeUCbjkIqauQ7q1hI51t1KK0Ty/hpqe1v3iAKauIQOhjuzFxyMgx7FbnM/2jBWLSkwzPM9uWS0H4jWSakogzfe9txEs4c+/QsLeTPvLzlrrhoRvk76vX6F462KFV4lQ8BLFX1JpZOU12xKsmCPjGaaOJgyXQiOhYn2hWzgcqRNjGX+X2BzGn8mEJ7aGMcvtMYUAkI8JdRks13nKu9bo7rrtiKbOLdcmLQMypUndSoHj4HIx0L94EanyxNktjsUH6UEl/L/6OZSnHAINvBlAMWXlIWyt38JXD0JUnOv+PLJ9lXJviKQHF0pLbebPoF7UwMYacH5I9nuT2nhJTj4jMvriexQBuncAXjupdykvxkUgQsRr6KKam+klVKsKAJwzYn02pUzn8ZX99Kgp0Zv6A/4AZIw+tqzSXaBRvQyPlagCtwxNENMBXBGJCaTT/M3d1NYKMBuEQdSzG/ZYjE1JEavYncvI+HARC2j+WXM1aZgz9ZL8CgLjpcv4/YrSBFGs6H2gmgzdANfJkk5Ly8dYyGKrEOA+lfembLO2OFLdnp0Kca5ZBG4NPOmAF4BBvDzXD767tWmEsGNf2xMnrIGj+8i72JWEhZfuKA2OO58X8ZDSx+O1JqcX7S0VhG6hgbqU5ktHE9a9+Ve9c4kGo1CPjPl8CTm4FxZbQjzsmwP6HqcIFYBDPVBZ9WoQ+2D7QvXIfv5HX1K2NuPVZNm9uwFtcmGsf6PuOhqK1pWSzqCgCmwApkdzlejUme6mZrR23REJ3PkQdaAyFo3gfa5LdW7PmNx7jcVqBgkY/SNldaGv+tW+xdJQrKAgqxhGSxiG5fPPeb6tkPlvE3G0CAEuGghWptSDagus6PiEnU/3yKJHGviaPLkzl10/T0zWMpIJouNnIfO3N5LeJK2C4xONQHkEUvqEC0YiEq/qK/U66LBsyfPc3Q2VPtQVRJkl/If0X6ijW2cTcVVwN2xTAqMA2dSbXoMPK3bH/VU7nJqjh5LpA1vIdkAYZ7+xljSLiYHT9oE2LHNhfDNZ2s8GSgeEYKD1VirZxu1ZbwqSl6Xuire6yZL+zwkbGlmIlsAcwVhhOfhAfiaIh4/+ogejQgeDyfWoDmEbXzSAXUx7RmQg8IZIWOZzcuX2GexReoXjhzAe2YabjMhqvs+FSJ4b2KP30cggQmhKlwsFqD1ev4SKHH0Ge/jrjEKbN1HzQ3oLtBg4GpkYoRNfImFVoXHLFg7l2LgETKlwJhboxTHUPGneYTP7AQ+K5Eb0/uCSO5Ymq3DzI1gmuGU4MVmzQSuYm/yP8e0MvRQ04RZiZtS50qyxWMrIOnk+IJ7NV3E0pPM3+ZSPpWftqVuEKawLfKtVnt/jTfRjniO7HM6pAY//32DoinfYKZ1Q8GPwjncAQdzvOS+csasvxPzi3DIy77DGfgyPztLfczbeCwixhep+9Vqr31F49wN+AMxHwVCsOJIRkR/wOscgZTuowD55INTwrue5rjZiT6CVfcKanAArmyYb3q8SvJB+n+cPqmbhydgqPgtkHij+7EGXRHJg+g2fpEgj5bDzmfXOTO2YG2fuHCeugHQkpDQ4gK1xMFUb1525vNMHHp0iO6hZxgePKmk6xBwKiD+jgXwh9euxUHrTJXmXSWmplATtcnNxFTI6sQuXnJEZ2JpBKsD+2BAAAAP0GaIWxG//vEnVTkBBzwgdJ0ZWNF1Rhjd1kKxuETzl8V2+eU76R66MAyhvqkmnkxp0vkl/+a6/BSJxSj4/j6swAAADJBmkI8IZMphG/74DjgQ7ZH1Lrefmefbtk4wS5Y2/x+DLBweoV7XRDx2U0ZTPS42sTXUAAAAF5BmmNJ4Q8mUwI3//vBGT5YQIcSq44tpvwtmFs1bxuNY3/v09IKvTPLNo54VVGqqjoQ4KZPNJThwuSr0wg0X709s5PGIgRr0KCsOg/kGruNjLFSrx/JisqEekF8neXKAAAAm0GahUnhDyZTBRE8b/vgOOB3OQFc9AXcbpj46i15PbP+Q9GQ2tJqySrnQwI7VJls29RfMrHbfsSK8eX/l7+18PjO5gb4bIXLCoP9kNhUaOMVrscc/CQyFS2DmIIbp/qLfqVQA7JSH5O/IhV43txHuIs879XU2UcXwNJG+MNUv/9SvID6qp3sUX+AyNp07MsYGwo+5nUN8BjKM9iPAAAAMgGepGpEf4exE2/rMWYsUC9Nc+P/yRU8UOqE44EY/0tBUFtXD6vVpxSXjArlMIURYAeAAAAAXEGapknhDyZTAjf/++A44Ilb6dinlSkT63y1KiWQ/IWMp8xWpc99kAKnB2h7AefDWJA658UxgGiN2hdpqZnrxVEyqj01yCAccBMH+y2STVmGjawy5zlKkJY8a3FAAAAAckGayEnhDyZTBRE8b/wDoisgQ2MspyG9MTL2twAnx4YO3N7dBgMj53wgsoD7fgMzpRUSBvxUyvU2F3yMVPcr9aQehYe7tuqCXqk4O3XxyB/oozQG6cLwffJ31HMNBW5ShOThCQwSZuY7d4iRHTfMsm/D2QAAADkBnudqRH+HoC9/9Y3zHp1PiKWpA3D3YoIQK4j0M6XxrUNNDfPHjMjM+xDhyvcRekwlGLKw4c3efGIAAAA8QZrsSeEPJlMCN//8AtC7wRdgOIxYG6HICf0tt8EYRWGccTPXL9ID3ANrHXjlUlFMeXfH4RGJz/GNKEO5AAAAOUGfCkURPJ99K3q+WhOPVadJ0Zn+q4gSRe2/zKKvWouYJMdACf73czGhSzAbCZzI2xgFyyWnZjV5cQAAACkBnyl0RH+DDJVK3gapAzXdW8ZNJVvmI9scX8y+ulkPHeyLSaA+eDUERAAAACEBnytqRH+BB5WDztAxKT4xFpAt8qSdgkigXdBpZgOPffcAAAB8QZstSahBaJlMCN/7wetWHYNorkyweRiHOPPdcXS1IITtq7AezlT6fweBHbowbC7asO2PF2/UBCbMgGRBUVHraspQzmQA6D7PUakKIGCmWZK+kVgi7Knv/moFsXlEd8QJPGaF0SHBN4ELaLTmKwUa7PPSxV2g8M1cOkFDggAAAGBBm05J4QpSZTAjf/vgPUSLvylvSJaoX67yRc1fALtzOq+SiE4b8qGjclrzocoYChnLAMbHgIHZNX73/lsLDOMLY4x13csDW904+9bFMORIwShS9SrF2hIu4YTLBzKObT8AAABsQZtySeEOiZTAjf/8AoC0REXf8qqj003CfEolq5ZHPeIt/nTlRdl6OlOcoVNgvTKHDmyYVBe5RD5gFrfZYd2HkGO874iVFT8MXHBHECxLA88R/Dor69i7i2G9sMY7GGZAq/0gNA9/XbGiCiL1AAAAL0GfkEURPJ99PbKAyIZQ79YXXIGQ9Zn3sXohriIew4GeDzxFPtkfYMRAykrz2oK+AAAAJAGfr3REf4BsKuKQrvnfCNbyOxT9/LxaMU530JspiBe0HCmuIAAAAEYBn7FqRH+HoC9/43g3bWZ3P6qT2ZqH2A9FOFkkCa9OSnFIbR6HJHG3QgSA//4qcrFGX1dvdEywSR/smSgEwGCmbp95ZNSBAAAAXEGbtkmoQWiZTAjf++s+OA7zHfcTr2WCBetLHUxD7/WNcoHXS1dnHb6Gxii8D96x0pb0d6PjC4DrB/WhVpxnxwTGDhj4JrHO2bjjfg9OQ6Gfn1d+m/OSL0v95OrEAAAARUGf1EURLJ99PbS4SuQriYYXJQH/fi36a9f+dXe/nbPtQtUH3xb5/YCrJqm5SDKlJfww5Fp9cPYaad58GXqPyd3Np0omTwAAAC4Bn/N0RH+DDmAYlvA1SBmtJcJQ6ephfzmo9IkH4wmc2d24zXHcPLKh+Lnw0TzAAAAALgGf9WpEf4EHlYMkv9E2/PMkyg5XBF3hcZ4c77PtedhhUSll/MrYkqHYxHIF000AAAB7QZv3SahBbJlMCN/9dtVqlh0uAaN0ofVjxMhZ1h5qorEdatiQ783inAG2LXHB1MiAuY0iagySzluM1H3bIodNuek635EEPQurFICs0REgQRMJVUluT+40C3yNiRn0C82e0RNNjNdsjWBfmakWZtlwAWD/Sn+XVH0p6IfBAAAAZUGaGEnhClJlMCN/+/6N/UQffCtxxis5+gpt/z2N6/IRl2ZC+fkvRLr0komX3ouJlTkB4RiYtyBbQy8957b8kr1rd495KdRJpeJJwYRiht2s6g5ssfWHOxE2mlZuF4E9aBSO+RThAAAAYkGaOUnhDomUwI3//APDMcB3mPtur2L1aw7/475bzHjZPyheciacbIxHbZ0dVmofkE1MGw1npgtJAA+Op2CJi6U6FI+Fd6w/4z91fbaUzx1UjzZgEuMaupFEOYn3mjpLK++pAAAATEGaWknhDyZTAjf/++AcgTWtokUg3fR5ZBhOWluG1KszNLcF72sEzcEsRw+v6kcUufDwbKgmv+8HgVhfNY38jxmDgIOVqtV0QZJ9wkAAAABZQZp7SeEPJlMCN//7wZu7TxYECs364K1K9zOpGtt1ZZOODFKfndyW2uhIDiGAJZLhzB4/6bupiaM+sug/+57cXUnsnlte9DZAA+U0bGdXg2VWdQHryyM7bKQAAABZQZqcSeEPJlMCN//8AoCu8Eg0IAKdPalTJwL5k3brsQ/2kMhr5e/LMbxJi9LMuGNB+zmBqO816m+p0J54fRMFrIfwddbFhUF2M8i5i2yKkZU/u9ZCP8OKKucAAABmQZq9SeEPJlMCN//7xIHZpGQEF4WuTpc9YQoczh9M1Lmum+7ob3DdhJhUVeNksMkmbU/aTy524uFw6YjqhrhnvYI9JX5S+eulWATUNNzLVA87Eb09ZoT9jApBuZ93/euYFurmkkWAAAAAhEGa3knhDyZTAjf/++A44DvTrlMIVpBxjTybeO80fsQOd462lkQixl/i8ixhGZjblZ4BIJj50LvODBaMD4t74KP0xM1A0shJj3XDsIWRvMHUMW2oF+/62kSLqLr/NQ6ysmc/R19+j6tUcDuLJBR9j9Ks4GTX/cn3YO+gyTrCwbIURfvNuAAAAEdBmv9J4Q8mUwI3//vrPjgQlevXQoH45WptzQNYqW5lvUt7k370Zb40x7Lqnj6WT8/G239QKi+rn9bt99F7eH8NCe9EpQ0PgQAAADVBmwBJ4Q8mUwI3//vgHIE02SfYN8H05IGFKI/2FJ4/9xTz/6bnIt1458gQ4K84mOGQcC5N4QAAAGVBmyFJ4Q8mUwI3//vhUdg1hSSB1aabIlvddm+MomJ80yRI4f09f1qYnZFqKwTK1GeZgBhUqgBV1BeNxJBtS3WYBxDo0ZtkTRCh095TeDwxQPJ/8LwJ+/nKDXqEa7WrGfxoouefcQAAAE1Bm0JJ4Q8mUwI3//v+jf1EPtOhzZ10VMZoqZ57G/blO7Ho+Yi3NWGkKEaOWnKTMTQBZb/3jumdLhfCXmgoAlPMUT3exw2xGWNiadKznwAAAEZBm2NJ4Q8mUwI3//wbcNERCUal7f5h0QKVTPWNlgbjmkWJ7S2NKgIHwG0iGg0786GE3lcq/pK9y34kGTeDc/yr07QOL094AAAAPEGbhEnhDyZTAjf/++A44Ilb6h2GJcOd7Cm/Xs77EAaBBuyH+UJ495yfcTaGsKqPsfkXwqaBplYMShrGWQAAAFxBm6VJ4Q8mUwI3//vBm7tPFg9mTMd/orjDHC6R/TN14aq1uH6BfRoMmsnLIs+xQljFq9ImogOyRHLL//9dg0AOfyMm3nG6BsZY4fYp4Q3XOVG/nIPZmEHHeOzYPgAAAFRBm8ZJ4Q8mUwI3//wCgK7xHVn2P8TqIU579zJBs783creb+z3a5PkBHTlb/74v4uMpPCgMQyA5H8erUDcdIt1E6LwdC+Yd5xch2XzNIEenrNxxtoAAAABbQZvqSeEPJlMCN//77mtCCdsTY5Z3Geb7Pm/QsC8ZaS5D5HX+6B/kPwTX54IPJEXG0Z05F//tLrak3nsdfXudbo0PXPSotywcbWZJ7pa7XExYa26grBmfGE1IQQAAAENBnghFETyffT3E9CkDsyrDByDhrW4cvbhQ6Som9ctl2PhFMGZxVTqRHWTPWDxP4Sfm3HbhULO9Fy6NYBKZryCailXwAAAAJAGeJ3REf4egPNbtmUKQ0b3pPCD9eBXfRv59JqxVlyBWqnu08QAAACYBnilqRH+BB5WDztJcUhExA+6Zo3n44kMJuHtPfnXh1zDcqNof5QAAAIRBmi1JqEFomUwI3/vhUdicsgtcoZFZMaO9tuG7apZo1pIvnF/gM3FczW9nxAaKJeV2vj19Cv3exUNOijnfRiR9ASN9ltcz+VQ4/0i7jHGYF+KFZP3uGFYBHkT2QRzPu+40UwUjscvCxMulW7ycJ2gzsxyoAhycqC6uUVDHWkGCd55vkPwAAAAeQZ5LRREs34H4Vo/ACpoAhQy5Ym6KOlAwH1FD/kXhAAAANwGebGpEf4exE2/rMWYbvt3Gio//JFTxWQEj0jgsGN8zykk9TVdEaSZbn/AJPKDK7pTunVhX0/AAAAB+QZpxSahBbJlMCN/76z44IUf87NrRe74hd6WSQF0BeASvxdUxL4Rm4GnonDSgvlzazo7BX+ROHA3HCwCRipy2tg6EaaSo8weDziWONE8To2CI9Nf6xMIHsssxak/PU/YdCrDJs2119wdbr4Gb4bSqtsAe2S56lY/+8eHBMDPBAAAAUUGej0UVLJ99Pb91NJVBQqP2t/+IqmE8MrKhL6lXP6xi7XAoI/PH3EdbwEwYHQ13+d+UGm6YHZjB0tU+x5EgcHwOjl3Wv0NdZKEk8tYEJAq9JwAAADoBnq50RH+FEsrDaMSpjV8Ar03X7CI6tPTzHGiBpSWUkMJVOVPbCEUeYgsYcFzDQe+eqnhHQbxBqzvTAAAAVwGesGpEf4exElWAWL5mv1zn78AGfLg+vpjLffFPe0N1Gn3GyNbeSNBNOMwp3LVlKeiRguPo5EckDmXdrHphQLAExUYkI1WZ9328MC5f+Y/xxKuGNSXp9AAAAMxBmrVJqEFsmUwI3/1jexkZYO6f0SqUBJutOylzcq7633ejQM+YcsbI01FdCSdVMPBHTcsAfhFmyKZdJ9U2koiIqT1xC6nMTtNhYSNWKTjEkyshkiwoA2kvaff/xZ7MolmxHJAhqhjc60XuqzboSWvYiUYabtLNBXbUEc94ymgQJJgQehwjl+Rj3g8ihDXo5vs8YeeuDIHNpoqyq5Z0gfHVseBv/VNMJy2baCQS8R4kw7jDJW8iWXhZUAEJXBoXny6onzxbxU7N9r71XS0AAABDQZ7TRRUsn4fNy3GU/pr7hiYIKBgdsTtsrSiKt3zM/HqMjo9nD98MgGaBM9BOcraqeIPJTzn3j7MnOe8aMxA90X+twAAAAEIBnvJ0RH+HsPhkvu1J2NTuf6EJ/2wV7RDK/mYg08fzz6vEXYZTtgVzrhrYCI12GtLgxwIpISuyR2pURhW4rUvLdf0AAAA2AZ70akR/kFs5HqaE1ZwEakhKwmD4BSiyCsJ5I7i1b/JX/yAr2QptkrDRNOf2vgGvhvSIFmCfAAAAgUGa+EmoQWyZTAjf/WJuYK1IHlx90OrSnTNXeuEGfLuro9aOPtuWzfxGlUqLMokdEgNX5hs/sSGuavwQOvVjNR8vXHeA7drDwV4vSf8z11/Kgo3tdVMNLWwgwf/pF0Ahg+eMMrSN4mNib+1ERK71wdTSsh7YgqUhZ1sRulfT/DmgQQAAAE1BnxZFFSzfiNYJXadE4qxy7VpUy9Q41T0IwKvWOe6XdY5vDoVxuSNPv5DidAVc2oMbHfP990ARqTeT6SM/eIoqojUVbdfcQ/gdcgL8QQAAADIBnzdqRH+BB5XFIVzyB9wJv/Lfxs38Dbl+6UFytIOZY0QdJSWK3tXhSOzb4vSc0LZtnQAAAIhBmzxJqEFsmUwI3/121j7DkGICHPtqvASEuNUhdr93NZ3drLM1IHCZVXsfuAYpeOu4rIQCjNeDjn/nmhlPTKoORxu4iqjlk8//dcd5KvIs5DQidbJMVc7qc3oVJpbzwdr8dbcRFN+vvSDcc7Fqa18YeCHCZHZwXF3Cz+Dy0oArWOp8iRf6FgOJAAAAR0GfWkUVLJ99mioJbr7gTovvLkf/l/YtKEL2mu5oyOikfBlzmuZ6xLlP+smIZv0LBUks+sQjAEsE5q8iWpcQhDfktWHs348gAAAATwGfeXREf4pxDGIf1f+twr0oL2RF3LEOzI8akUIhkzGFQYj1qALSEqZaaagajuk+Xra1iK8MNJ545dD2pPxL+B29q8NRoc/1g0bDynEulFgAAABCAZ97akR/hnW4Mf915MgtV8Fyl+jyhJqaffOqPIliFZXeE+XQchMQ88eJgdqLi6y+ytVR4nyOKwv8jWcpSMzkR+NpAAAAZ0GbfkmoQWyZTBRMb/wDwzHA7nHfcLHdlhDBdxj1K2iflRpBQ54lGRaeTModRU5a8siLRIR5KHbbhYv4ItU1x58f1xPglHmsDhwytPRmDnxl0gvYO3B8/ds3Yg+ydO+qQwyKDZQ62qAAAAAhAZ+dakR/gQeVg87QMGmP7dx9N0u/+sitdubSF+CATvuAAAAAf0Gbn0nhClJlMCN/+8dP0UsID4MvVXMepTuAd92+Z0g5Hm6laabBXW3H6+DO83pTGwIUl3MsSoPcrWr246Y/Qww7pWgjMh3A6DDXOd05ptlNb8IL/CBYCg/64pd6ZHCzVpOaP/3b5cwCm0Sn5E7yOZ/JXi1D/JAEeBMIy/0ioB0AAADkQZujSeEOiZTAjf/J/Y+WDSC6t+DsFl+KShs2+BRlWjHwkCbRkfk5F1YAt0JG+eWgtV3xfaTFPYBv+mxEEdHUfsAnOpwhTvDIVMNDKRQeXrvoZCgkP6BOglvBRcnVRm/QCrn8gMOw51JbsPpVerALbca6HHK4IkOueWgXj8/rw/Gq7hL81TYpKHnCgsWhhNql0edQUpWAf3cM29hrgBSKwsy3jQs6DrK17B584Fz6wiEv//JAFYwi3YgyIjyLVn3WboOj/oZw3ulrAwtsHz0duObfFaFbLBt6sED+kDCfEscSq2hhAAAASUGfwUURPJ/Tvx2ZXIhlDv4p8IgQMX5AlE+WsCstVw0AQsep3TWwxt8vGXFQoms3Q7j6KFLJHsNyi/X2PtL8B13mVJbuvM9va1sAAABDAZ/gdER/gHDjMGGbkbroMeA/+R+98sKd9OE/osjRzxqU4fU5myqG9RgJ9RKLHL1WRSCFE5PdzrAJ6RnJAR/133tzwAAAAEUBn+JqRH/XxZ2/arlGJHqHMGb4IS9k1LB2PTOkX1omkDNqjv/Hyd0h1US0g+kY2+XgTJ0qtQYnHLVn+AHiS4iyJNOTtEAAAACwQZvlSahBaJlMFPG/9BI9qA10G55ifcR7emNzO/s5HqLdZT5FQFk+OcoHWoo43SEpEIzC+YKxdwwVhbLtnL6xRqv3B4n9IDF9aNGjQC3lBkrevbFcw6jXvPXdk7lFbzKvLaUj6IFJ440Id8uOH8S2xSDBjkuw7fw7kuVtpJQ71cMUX6WsUmpYvzm182OmweSjE2W6mQahX5TDXHtuYFW+s7wqQG7mczW59/E+jHqj+GEAAAA4AZ4EakR/mJmxUG34cEMCTD9KfzSlBy5J63PV9xDXOTeu/CSqfUIWFTARN80zbyy1JouE1/9W5KAAAABpQZoHSeEKUmUwUsb/9Cwdx7ECFH/O6zLxukzGly8jJexJ12NnWZVyEH6v2GpBHI24QsdFluV6MU99hhAQ/ORGnQtP16u598A33ALGAhCiT6Fkia5Y+ko5b6BByfi+fm1pz4EM+HBx3ixAAAAAMQGeJmpEf4exElWANW9/woTjJTRwrYLgW5TCxhNGjZKsQXGThPJJ3FtL1nkHDBC/zEkAAABbQZopSeEOiZTBRMb/++FR2EKSC0QlA46ppypN08vPo2SkMTwZmYivCQot+O022lnoqxQWOoCxOpu+brG/vq1OiqImEEUqH1aSs5FB5wZJH0p8pceogOzxqbSFCAAAACEBnkhqRH+E1CtCa7nT6sPwHk1BQYUkaWxgrGWqbiIHMPMAAACMQZpNSeEPJlMCN//74Djgd5CY+5YA50g7ZWcl3e6HlM7+829ssvr2kKYX39Qy0iHGV9xtA5s2X+EIG8ksj90L//eHWxQYDvlTbZLHSXrXkPuNSEsLXk5NQmC2u4ggLAcIDZeLlKdlT77ixCPNy4BIJtROhSq40uHLaCpv/QYI7EgtcO9f8Zyycfr2+0EAAAA3QZ5rRRE8n360LEl02ntKg9DJ1qf0Sq2mRcCNvdMAMT2Deu+GU2vzNalSElMySEM6p0O8C0Zo8AAAADkBnop0RH+AbCj0Y6pj2ya9J7lxHir9K6838HLwTqvrUGcUcg/6mmd9tmwJ2olawiUUSVMN0KYF8/cAAAAhAZ6MakR/hNQrVcoxKlr4Jc/+7NWRzf8YKIY6R80MIEDgAAAAkEGaj0moQWiZTBTxv8fOfWo2KDtv87okIotgJjCiYB6nOqRdka6T2FSgrP2m37XMvHZdI8A21zD1pYWm/KGf+qjCETUhQS1RRxjZtApAg2eNLrsSASJT/M/z6WcXboIDcyflWuaJGh/CcbDrACgNgPWBxQ9Cy4GZ1z2sg/n7DG834nDz1pRrY0cd7DCoe+kfoQAAADMBnq5qRH/N1V08EzDwtoIri/39PtNtN5FyTk7/bpdx9BTbj9S5CIuX3khlff7kMGMVX4EAAACFQZqzSeEKUmUwI3/9Y3FKc0qbE9JKkUV9XeNE8I6QHJCzI1xykL6LJyJn8sMTobA325dt1+ho1fe5CgSyqyKmWLEtqPoxN/8FgOH7HVcYanqHPuM2LmQI6mevX4hqU31p6zFddFY3jjC9SJd2W8DUlE/UMqWtTbO/zdKi6N77GCUpGa8akAAAACFBntFFNEyfex1FrjBvpN37qK/wIQZCAFBrCFmyElK/aOAAAAAvAZ7wdER/h7D4ZJjIOAubxKBfeUw6T8sW9NMSWDaQ0+veb/Jd0BOAnj3BtP+dOcEAAAAaAZ7yakR/iOA7CrB7i/Mz4jo97V0uGcnVWvAAAABdQZr0SahBaJlMCN/7/o39RD7fvlTWJwGXUUqvm05+1QVmloNO9nFYywACdKqyxRSD8AZPjb9mbnY8gjIvuePK0VaUznabobdf+xj0P3e7kDpR3NBMu5vVSLNUN/ElAAAAZ0GbFUnhClJlMCN//APDMcDucj7DVPmcuWIK5vVjxlhCepo4trDOkrTmEGu/ZySN/QsIlFHz/yEpWOSVZ8qYXpP5zoMHCuoHGjtJgUJ9YY/EQHTth3rvh5r0lUIjrN6aJq1qPEo0fVYAAABHQZs2SeEOiZTAjf/74ByCXW+nW+BxzvYU369nej5XDAX92W9IvylVTvLOZPNdVa3B4/9OBz/yENcbuBsBlmdr/FOsXai19kEAAABQQZtaSeEPJlMCN//7wcOsd4IOcpHAOYB35obyD9fNfDE0mDNzUbz/oQknAvvZ4dBtywxoOtUrwhbgoCnEAdmdPKWcS/ilLOLAuTdnno4dinEAAABFQZ94RRE8n360LElvuh+CIdT9OlDp/0q7PJ+MJgJbrVboZA4kHTJ3ToffbZ9DMO5KtxdcbtUFSeygBOcNgX5zGmhe9rpJAAAANgGfl3REf4egL3/U68Y/3GAEKHH5y9IApJhkFHkWHqclHfOYIjpgTUXMpdkmgPno80wnwLqyZwAAACoBn5lqRH+Gdbgx/3baE/Aqppa+0jtFnUyKiy+Zwy1ZHyz9fyMNwfsj+4AAAAA2QZubSahBaJlMCN/9dtmbcgHeY/EYukQ7aaclKiAhW7LP0wcwt78GzVH9jEX+ul4MCqF/2B0EAAAAQ0GbvEnhClJlMCN/++A44EO2R9S/OhWXCYYnfMvGSfmF3VhHoRAookQ/B0zybEDHraNElAkFmtkn7maAS3rK5ywR/4EAAABvQZvdSeEOiZTAjf/7x0/RSwgPcTDLgPuFlZF5qsphrzQTuuhcRXT5hraAgVTlKSofuENEmB9xuPyUnh4xra8xjPnk0cC9TqCjw8oFY3hxp8p6Yv1rzk7Q/KA8AKgsEUTklNSbcJitU+GRzJIOetTuAAAAlEGb4UnhDyZTAjf/yghA8yNyXL4I6H5lYLZt7x9sYonEZViic8JINT2N7pAvBTCcrMb9wQzCFD0BtXGxWf5Rsx/8l+VOmJ7mN3Vdkcll0IUwspcRV/reK8gEFqp572XghaC8+bGg1kDUICNBmBkZxBrR0PZDAzSKR/HGRfU3wFVnoidVdywk6QM8o3fSqG0jFI6YMWAAAABdQZ4fRRE8n309hYNYZQ8K1MQRY/3//frBdkHLgzarcM4u6jySXsngsvUyy6EI7lO5ODiPJckJ9lqLmYq9/4Aow4+wuMKOMB1SJrQK2t0CbZ91mEJWOzNjbZQ4RrPhAAAALgGePnREf4ew+VpnMlKRTlxMhp2z/jXp9OUQEWoWTWTddscAANjrykkhGK8n8fEAAAAuAZ4gakR/hNQrVcoxKmNXwBAAZe2CmhnDxHLa4k62bX3bzPVMZ+4eeevEvCRnSQAAAFxBmiNJqEFomUwU8b/8A8MxwISvXq6TvqNm1NMGbNBYU3BnWrbyWWdLxOVYy1WXkGewPWIzm5q7yUvGhk/YnyoDdKJ8CrXR3ARnW/KIQtqUzAbH4jz9hM3MiI/EigAAADYBnkJqRH+HsOQxTh18+g5ZBI+1d8/B/i7BlQB+WddgPTvIA214hf/1CtuC75Mbg76dyaxml8AAAACrQZpHSeEKUmUwI3/7wRk+WEBIIWYIPFZOw56t33hsE66tFa3ZLekWImE4LvDK36TKHAgCfZYuTC9EbJYEL7N1Y6TRBVtzM2ya6R2+ErPXB2jjcjTa8RL4ox+ohOuPW7UZIiqbV/BEGgQbrsylo7OhkBXa3M4cx81PtPa5oTEEJZAzD14sh2F8HgCI7sfH9KNgsLmzFk/eBzzI7emYM5SDkSldP7xc0a0irmWBAAAASUGeZUU0TJ99EZTLsyDpEIRrKJpnoG9EXvvB3xvGEbs61IZFsRW5tRZxKSoUAv2FlyzhGJgURfX95iQnOI5xMTYX4Tm9fdpIYHAAAABCAZ6EdER/h7D4ZJjIOE7mYezedRutYU/Lr+V1IdLd155NZOE9hpCyg0nHA4eDwBg282Ojm5gxupgyr7ONyBfl+pZAAAAAKQGehmpEf4jgOwqq8yvLtAmySKIuJcaTjoN//CBYJDzeFU/Xqm2khRqBAAAAYUGaiEmoQWiZTAjf/XbliF2xAJMQm7WeNk8OMW3rdjdKBNxob1e/+NepMrYVL0nMuodAkrIMadu5mJBy5uvlxvNQp7VxX7xQevbsT781icpake6htsC9t2XleFcTwqueItgAAABQQZqpSeEKUmUwI3/74D1EHSkIO4Z9Wvyhvb+EqEIAwhAaGZld9PgYxqFm6CIw70d3crTbR84Rl1QSjXA0mnDPfQOFBboYhbHq4C0vpJXMNo0AAABQQZrKSeEOiZTAjf/74ByBB09qkSKnA0a7ky+N/3pDcHVcDQBwXvawTNvvn2svBAJZAi7D6VbLM9Eih/kRceBxSGO9APDLO5zFex2ykGUOA4EAAAByQZrrSeEPJlMCN//7wZu7TxYHfvdr3YChXTkbqSdy1lvNC4XiN9KR1PqesySw+ABEfzbkpLtzgLfVv+d1ee4PGmvseKH5sGEQg1/DG7yZPZbZVEiMJvrtqcWWCmcntEA1XNbsrhVN9bi050eM8WOF6MnOAAAAYUGbDEnhDyZTAjf/+8OhtVTlCNvyUsskUPEVWUnL9c90uNQCPoB4rg9+sYpzZa6v16m/de0qMOA/BNngL8GoEKLLIlgO5+X/dNSJs7hCfm/5KPsbwl8m3mNvXuL195ed9WEAAABaQZstSeEPJlMCN//7wc+0/eDsJNVZQtMJRUijWyOQP+twyrFR67F8Pj5Yf/Gi52beGXbajrcE2ZHcat654rgyJ1pKhO5Uy/gC8B+G9/OVL6pzknCH54OLfIXDAAAAUkGbT0nhDyZTBRE8b/vgOOBCV69dVilfxFNE78xMP/CEHWNxHiUnX6DyHq+qeebHfMqief3NvEbCPaYjoYHVRYSzyLFgSrimfr+f4tOe9AjjPHEAAAA4AZ9uakR/h7ESVYA1XP+1Kfs2fuAWfT1RiJtfaEy2T44UcRtUezFEZUPv4YZsGo4J8wK41PBEKIEAAAA0QZtwSeEPJlMCN//74DjgQ7aLyl6IAVFCZ93OMgKO7RPkPiBlzxXr1IPdsldBX9MhV6aIfAAAAFRBm5FJ4Q8mUwI3//vB61Ydg2jrrF0XGCsxft8+s5TpCXAiRHLtzCBWGDzOIi2fFiTNEfqmbdWsGe9ejcTntioJKJzrv8cb/s/5QolStYE3YQp5+/8AAAB+QZuySeEPJlMCN//7/o344HrOM5ff/3Mc7c1rlUvDf+1BYAPMg/BpElFxK2uuwTvrLu3L9d3CeuKWIiq2NiwAKzpkZXVwYoWgYQjBJkeY/r7x0HhuZX04/RYn2nHmTL/bcfYO6cjHT3FMoIt8A/Veo+J36ssJSsTbgb64suXdAAAATkGb00nhDyZTAjf/++A9RDRLfNfS154ZVlkmiUy6/g+2oGU7FxgxqSr5jI8p2G+RtB4SgTfOQhgR/tAQN0eXRO5enuY2bDLAZlbho5L7gAAAADdBm/RJ4Q8mUwI3//vgOOCJXWMSryw4h9E3CB6A0Ke4+kyCv0XYLF03+9oB9Vjy6J3L026VqcJXAAAAU0GaGEnhDyZTAjf/+/2woADtubdqT10cQKSKXeNNuaLX2/j7JafRM2QBPoIzUNui7IROPY6gQ2i2OrH/7RUn94RFukRHfq8+EKh3HJwvniK149JCAAAAR0GeNkURPJ99XohLbldl8+WVFj+Uva/8XeQ7k+3TjC2+NP1oPDpKnfim/+OozWJAwVNoYbSwY768ReLtDwbAMvTprf9JfzUVAAAALwGeVXREf4egL3/U68Y8CYofdFnLKmENVl4qCMGjiGYU3fOYIjY08srI70+2E3XBAAAAMAGeV2pEf4Z1uDH/duSE5sdhg1tyJKvPRV0pqHivmE9r1icmTrC7mlJ7RSXv2UrMMQAAAEdBmlpJqEFomUwU8b/77lcOB3OPxERFqzYjd6PE6HsucwQ5Vqvwef/OwSmT2LMOe8YS/MYKVzL0ysCkf9ELyWxaN9aVNLSQfwAAACMBnnlqRH+BB5WDztB5pfz3BcMFhRhOn/CrTtDVV5wj6Le4OwAAAJNBmntJ4QpSZTAjf/vBPT5YQOjX9DqM5HaGpFMwQ3svWbjKwnG+fOlGDVqt1ndhhMJgv4GvGjt71cvzG5oQlpabB9hxzOPik6dNzXgscBR4+cbPmz/fP5xfYfX1NBE6N+5E94WuKtbCmyYygQguamHJ0rcWbq8PZ31zvB4cN2ZTaqhzCYM9gDyDIxGEYnFFk8tLu78AAACPQZqfSeEOiZTAjf/7wZu7TxYPZR9WnPBp8pmX4bfrqeMVK4ZK+wWWuCUmWaP3lATqHHJZqYOGZWX5B1ExYwR/eLb7x7S8jkASRZ4JYh0GwZIP9CaseXY+FdZZiTo1Af9+ZTj3dfC0dKh212pB65IP6OYNrtjjWbaU5N4lw/dIVo6wfWQ1Ox2R/Ah/76cmOMEAAABSQZ69RRE8n3sgptQ/K338OzsnCnldYLd7RX/bbPagtm9AT2uZlYrwak/pO2OvTeE+MLYy+EAUJYdRjCTJ5ws02INMTcZhmyBQbhyMvzM27KZkFQAAAEkBntx0RH+AbCjmbUS83YiObdjoslRvNC+1RH/LHg47MFmS5IoLuTYHGm+zPkdYAIx1ysH28isEgGxKBbm/3QA1sfXNc14DBq7XAAAANgGe3mpEf4TUK0MsYlS18DEAHiL+PcNPmI7mQ70eRh7AhGHFQq8+OOpKoNFJ1E0xbu/KTN8mAwAAAIRBmsFJqEFomUwU8b/77lcOCFH/Oza0YIke18+jZcsifjoIbn8hLB2G1sybcgi8yUenoEyRzOvlQ9G6LQ3LjzLO/wVkdTKQGKdViQNnNx/ukcfeqyW26GuP1irc+a/5I0StQFo5b3fttkhWZLCcmvXpDHKib2Z3J8v4Mxl+oH8tYhQMmnUAAABHAZ7gakR/gQeKxIei4mGll/v6lUvt6J01kEEpB2nbmZodQK3mc4O0v1zs/0nPzjxxUZ3GOYRGobV/cAuzXdUnGujW6ZuRq5cAAACZQZrlSeEKUmUwI3/74VHYcs5s9CQn7HDcl8JfyAYtwAQoJsnV5HQm+wXQ23ptJOsGIH0+K+T4tJmGWxKq2QkvbVGAt0rBbo9MwfKNZ4T+zu+cRWcAzne7vkeYo34ehQqXPe0QPB8YoxnrjvXyG9IVsIPvTqRMCT5da0BnF7ybZ+dwHNKFrocGgyWwGl61sMCNyw7UdJNZmxywAAAAOUGfA0U0TJ97IKea4wcfTVP+rq1dhzTiOczB8P98SD4FQyfV4nTwQWnbpuSHsjXkGDVY+VFqlbqpfAAAAEABnyJ0RH+HsPhkmM1qy8bxI+teINNs77QclW9Syn7VgOysKvEXYZh1N9bkMSWtoQGKeb4wR7zVJmCMskcNH/eBAAAAMQGfJGpEf4eNe2TaTWCeCgaUGvNb9YsQ6osP4XH63aOXxG1jZCXFAnW0W+T3+SXJsEwAAABgQZsoSahBaJlMCN/74Djgeww9t1S4BqrecgJLTxVeOuPfGThsraJOF2+pS0DHf1L7/3y7WCHgywvW5UfaSeI840uFRKAdPikqUbRyEQbimIxPxFcVU7IF9oKgqilKzVfQAAAAS0GfRkURLN99UveC5swUPUef8Rxxd/DnGwoGPLEv0jI5bF/aVNigJFob+cQVyS7+9J0DBc/WkPNi4m+wyxdk5RauSIYK+bf04StobQAAACgBn2dqRH+BB5WG0Yk9tZXU7nhXQUxniPOff2if/2gJwmvSs+oXpZoQAAAAgkGbbEmoQWyZTAjf+8SBZ14QI606liNL7F0TnvmzD9NgO6qUoiqkr2dMXGBxVHEYasaZ1lgL1s3uJPXelryzxkrFiQQAJkYTJ6/ddfxquNOViqUsaM7edj3CU24qxjGC94+WxcABRaj4MjKQx/0LCKQROQaDtp8J+TfTj/pDj9bBDoEAAABiQZ+KRRUsn32aKglZlvQTovHIeNH7gDIJfCEh/db+sTywUiMMdw5UkCuutzw8i0Uw3jTVZuGhnn/DdTHavTyUTz90N47s76Kr+H9cZcMdUHeBR1waOeg55kL0Lnz5MaAKqHUAAABHAZ+pdER/h6Avf/WMAaj+02h/nyrbUOfaALfnmG2jGc9vXFwOALhj3/sE457lRtcFy2I4QbRoLfMtRwbncoZiPe5ckLTaUSAAAAAoAZ+rakR/h+Qj8Y/7qZFjwDICR967ohOCf916Jz9qNGDcGK/ZVQo0JwAAAGpBm65JqEFsmUwUTG/8CBPDIEGcd9xK7sr7PeVRUblA+A4wrZEIYDPPG4pGP8GneuG/+WblEBE/YISVAlzm8YRF1iarh9/0YM16M73ZZ8APAFET4NhFgwWEytfdD1qhPlSrycwSgeZBOZfgAAAAMAGfzWpEf4EHlYPO0DUfukzjrb6sMwf9TWr92SL5dYTUf9xex5gvecavMELDpnbvgQAAAFlBm89J4QpSZTAjf/vhUdhErrGuSF/CSLUgNGgjSozDQkuWrpzMDWZP18VLPXElhj565QB5ATyxhTtCBjXMORwH+flU+06xO0G6o+HBCuMLMHZQWc86Sgn8YQAAAHNBm/NJ4Q6JlMCN//vBm7tPFg9k/OLpHQkps0Dx8gSIgZVvNmWxh2JbcEgdj1ANWONNo0C46lGcg6+GMwWquPjmXeWwPjjFXLS5WiDMCJHZI1S8p+dhwYEHn2JemwPkCTGJT1YfRwqRnd4zFrrK7w/IrZzwAAAAUkGeEUURPJ97IKW2n8rgeWx6x/uvwym+iKn21XOuUJQvFl7Jvz/4BnZZvs/HplT9c0k3KbuOgl09GJ9Lr+6rhKi+Rn9GiuEu1IK84X22BQAPloAAAABMAZ4wdER/h7D5WmUT0hGm55/V9t/EfeHSNG4H3z8SITMcvPv9nR/LKHSw6207w5vl5dvHvbN/P8hIiRRBAbI6dHaJoqpBvvYa6Mw0GQAAADgBnjJqRH+E1CtQOK583hyuzCeT2ES3e7EAJSPa2lf4ndS5OFQQrgTaHVvvcR/0mjt9zA+CF6l0IAAAAFZBmjVJqEFomUwU8b/77lcOCFH/QjF2o2bU0D1gZNTTA5MxeI3pXbKFI+z+I3Ehsro47fLOUN7RwIz1I1j5IIQ79ZTKZXOQw7psvfz8GmJZtuqJ31wKhQAAAEYBnlRqRH+HsOQxTh16NjvxJP5pSGoyZ2MG9YV7VnkhZNujwds1WAb+DnvAXVsQBWexlXy/XhL2mMKRVlpuMTQ2rNsOIosHAAAAp0GaWUnhClJlMCN//WN7G1kj63WwyR24JHnTxkh+gYl+KzInu1p18pjfQ9WbiBY0GWq7ysDGaRRp+JW+XOLVoUSeFufVt5tWKNC3usOBK2/C6PhT+8zUu5EZNdqT60Qq4VgNXwzslbfFjnJRBDpdAh5wLLWFV14beae+G0fFOrSHDkUNpkd4sy70UWw9D1jdUuQ5rhLdlkIPJ0x98OQL5TO5gSkKid+BAAAAVkGed0U0TJ+HzcuFAuIfw9nuIuW5/0fmGPCk2buV6y+c17CysRtliEWgeuDAVuS2NEOS+YOmLSnbstF7rBcIT0xiTFqSL00ySVVFsbWwkDRuJB35PF4ZAAAARQGelnREf4pxDGQCv7OsLfYo7T8YkPot3UwsfuaooJN278sKvEbac7AqI6vugAMF0dBPcCL6ff5MKRFnV3TxohZvBW8Q2QAAACkBnphqRH+QTAw7Sa7nUoHLO0XWnIk0uc9V1EtugHZKDYDxjTdjxpe/wQAAAK1Bmp1JqEFomUwI3/vBm7tPFg7wfAZdXzxs0p4sQVbaLB0awSzslu6gp+TNyByCumJLwW5CyVRGbnfEuH5vtaS4wvsym7ljMIKBynK3Kg1MUwKdsVhXZ1B9ez8GwiegBu3bUflTwt/cPd13A1iwLV6wHvPtPLPVBufoRxLdi7fXErtqXZFkhyUPiuvfApRuTZLb6JlWmyXzVCqHWPaoL5G0ARafVrc4NRAD45EOQAAAAFFBnrtFESyffW4h7bT+VmzF2BY/3/IEonznzSaM52EK3yed02INvn4T8C5fRiiAYPyU20UWm0Zioarkwh5eCaLWUgTp5urYldHQq3kp7USwoZAAAAA/AZ7adER/h7D5Wmcyc1Gak7aI9t/EfeHMOpqmvMV7oyvL4VZVPyQ7xKvwftFj0zEdxLJQd6obfw6JTO/DrvdBAAAANgGe3GpEf4TUK1A4rnqdTB7ynQtEgcvcbzbIwk1IFg5AL8SoEuxeGvos/djDz9/TNXsT339wiAAAAGJBmt9JqEFsmUwUTG/8A8MxwQo/6EY7tRuuvUQ+IkRW5CuuQ16cI6K9UjLCHt/DFRPI0/cesLTyTUMPgo74dmiSCFWGW0NNaVwImkdPP8OS321tUl0fM2FRF/ydM1FfUQ53LAAAADsBnv5qRH+HsOQxRCCMCTKsMv80pUcMEXeFoGccbVvJvQNyqsPl3dGC9TDj68G4sLpt/4joEDUDO0HQPQAAAKZBmuNJ4QpSZTAjf/1iba5425bwbsmO+LqBQWsUydEd6pAOGPrrpQZuwRoMIuUNmIqaa/ZfsYj18F3+2O2VzASzLDk7M5bTVsxDhHKKWgIuIMjnu3J7use13jH+mTne9OXAolpGEM+YOkgGn0roRh08mvTLZT60kAdnRLZNC1MbYLQv/UtAcmYn97x3XP7Gik2RIKYbYXaybwMSmgqFRWYYz4aJwVLLAAAAPEGfAUU0TJ97IKea4yheePM3UV4BYP87QSjQ68uFMzfjZmCa2q2jTWdubpQmvIRE5Szocob6t0FwExDfgQAAAD8BnyB0RH+KcR92pSK4d0iK5pn3//DwN9n1Sf1/FycV11cqIRtpQSj1cG8kqGHiP/LJ2Q3+b3t9voEB0KrTJ0AAAAAuAZ8iakR/iOA7CrYE++z83jlf99FGL+nfLDBphBaI+WarzWjX/yqdZNNuIw/QVQAAAHhBmydJqEFomUwI3/1jd3Ei9XsYRCa6uUFLqSIqHEXps61oaZW//CrIivsMhGGZt/+wTW5oOHXAD9VovymjVgcdopGLcyH3+31aTy3MJifAOoc8WLgubo4ISngsCMnt7ldH1X98l1NIvHlX3mcEIj84dFxX5oK66OkAAABHQZ9FRREsn4SFSqCugoeFae6NWP9//36wXZWgjOWqAHC9KPG2ihYHENvn+jgjcT/OLdhLP58EwMXKROhlixI0N1Rfgfyt4FoAAAA6AZ9kdER/gGwmmmOq/Wm/D3dnucn5igsLTeqyV3HaZ/bb0TUerDFX4r8MRrARcezd2SF5uBaMc2S+4AAAACsBn2ZqRH+E1CtDLGJUtfBL4cPgesw0gfX/r7jzYBrCPcjlvYuTujdx5N+RAAABSUGbaUmoQWyZTBRMb8EdgCeRWNj1vNMstxJbRx7SX4GlzlNFlrR55BZRfwuUxlJowdEahz2muYI8I5vx8hRYpCbfVNcn/buf7/gZ9+q/Xex6CfNdIK5dm6hYw4V04WRD2xx5XfgFgFusJZ2Lu3x5KCiAJR/MKrRmuD0LpBj2HTd8vSTNdD0mqExm1jVxz11VS0JkCz1kf3+GAZ8aWDLfmsfKwTV6EgQsGOAvWjEoPcSCVw4czv5SM2yxrz2q4RWCXKP2cQ32KjlN9WUwxh1yU9H8IlwQpYkE9YLpUDTnf5VkAgRqHHtbB1bKzLqdQEBRnPDUxU3lCJ8z1ffgnZnrwQ/1jrkIVCHtxSHz+jlGJYYDk6e7WXKb6WyPh0gIrK8i81b+HquyEnpTv7rHDiPZTeoAPC5ivY3TNLqzME6cL8+xoSGHv3D8YziNAAAAOwGfiGpEf9e9wzWJD0aXSLU4v4AFyP5OmsbBEcN/WP0ljTwYR9rY1eB5C5dJcii79292VcOw01nChUeBAAAAwkGbi0nhClJlMFLG//QSJIKCLA21sITm22CQNIzL75x/coPBB2IH+CkNL45/JcPWK61azax0RKRKTm5d+BssnzMtwKy+njVCxG7G4L/0fmAAHFHQoH93jRGl/HjJy46vx0G6NyLZe8h2jDkca48Z8CHa2hCT6RIrRCpQcVFTXZInLx/gmR/ABv72Y6Rg934/cwAwDTTvn/FxI59I7ms+9ozaPN3OuXj/ZPs9jO8lHXscoqhncvzhU0B/ZiiYXSNR5B1ZAAAAQAGfqmpEf5icgpJHKscCg4dcTFi4T//E/XAXJpG32nkk08fkKUrnPGoSQKoKEiE2wpfQy29GE8eVH7/yrD8SLOoAAABWQZutSeEOiZTBRMb/++FR2EKSC4tylzphI7uNxilsvtRiMSiFLXZzkidfIAjCjC8+Yq1U+7FQos18k2wGsrHuaY/yfMMHdO+Ytp9nyOtNwN8ZSWlSsgcAAAAZAZ/MakR/iOA7CrDroBeWntKNDlGBefdNgAAAALZBm9BJ4Q8mUwI3/8b8mat5xbgYxKkPL6JuE0JrSSkGLS34qllRYb//n6qI6h+DVp72C2Bni+fBUJkeQ0aRWTwoPYVKKNVd/i/tVkLXXk1eYhf5mDfymriDx+s9bFWvnxtSq7+R30/mIRxeiq80eB3FjjMV0v9A9Ip/t0I8jdkPBxJNLY6dHX57kOj8uBOxzCEcewQyFjVwZ4U9eeiGDNW5cmamVQO+G1zTO18X6eiBVrQ1sAaKeQAAADJBn+5FETzf0FK+iADehK1fJc1pMmFCTGtLvZ658Bpk3MRU0V0znp7jGG86RAJ+kBsHtwAAACMBng9qRH+BB5WG0Yk9tcQHS94V2gbW9d0P7JZCiBx0bHdSHAAAAHRBmhRJqEFomUwI3/vEgWa8cBuzHU5S0ix1bYRHinDETekQfB+cgzhJLg7T1kDPCseQtg55Z8jXr7sf2ocT2sSw5hDEdVg0IPoU/FGMpuD08g8QhOFlplsTFkZYkjR7winNg/k/jN/HF5GJogmrWYBvwozzlAAAAFpBnjJFESyffZoqCW9/YL35uOpRg/l/iy41hTe/tE0ssOv4hnLHsHPfmKcSLTegTjkNLZVWXjO2NeDSEq6BzQ01VDjT7II0nXMOQjdC1CNOmzNCCPqLBayW02EAAAA2AZ5RdER/h6Avf9Trxjzf30y6kAOs6WLBLxUEPR8k/75zI+NLVHsG3auYqXMQfTCCvtQyg3VgAAAAKwGeU2pEf4etKGWAMuE3sYEV/anP9U6xl/ko75T+a5ob50RK54RK1XOQ4sEAAABcQZpVSahBbJlMCN/9dtmbcgHeY77hZBIuaacrM+PCy3y2hv4fgKY3CAJLaPuYPjQ9hPdbPybTfLA9ZhcESX6IM5yNtlJbVY47bfUjagQR0EA5I+ypI3tSGPRusuAAAABKQZp3SeEKUmUwUVLG//vhUdgh9eiEMPfVB8tcq9JDka35FeHzBuuUIwuu3ZTYJerdc8OfLMQc0Kb0i+lVDCTmweLq3VzwQLqGbWkAAAAeAZ6WakR/iOA7CrDrXqtfOU0OnEFoMihrTVlbhcOxAAAAckGamUnhDomUwUTG//vgOOA7zH23LaoiapuTmxATD2bJaDhWz+mrw1lKFIcbAS6u3oi38M5gOc+H7/dMAlr5FiwNQQ+MnP95HqTXfhGOyOvhpUnJ8xbBRClFaRLviLbneAYkMEKA0mVg3tgaK2CBRyPigwAAADABnrhqRH+HsRNv6zFmG77dxoqP/yRU8VkBI9JIa9XsBjVtXD6MA4j5uBIgllV5DYEAAABnQZq6SeEPJlMCN//74VHYIicdzoZjPCSvgahArEevFvdThrsLLMVaJDVNLqpZTPWNC9vVX/wO1tcwEn76IJESXUEvUSeHuARe65ZpLp38T60KexwII0tfQkOqrPOP2CqJfFhslG4y3wAAAGJBmtxJ4Q8mUwURPG/7w+3BOtghVFB1Iv2nqcWzEymAHxzkidM6m2fE6ogCs/6Xy6jtl/XzhhsHnCSDlAQ38GrqPDxXxyRtvk4SCfcKOBRFyRAOB8n6sM4vESFElpVQD/vjoAAAADABnvtqRH+HoC+APa8YzC92E2GO/nuYBco//L1SmEwKP9mqIR44ppDFXNReibXTYz0AAABWQZrgSeEPJlMCN//8AtC7wJQ477MPq5rGdC7Dcs2GbwHJuPCJyM5OIOtlfIO3cE34Ssa8cZkPDWcsRjRYjceQX64sn64LgFO/LLPWgI/X+x8DlTPlkHgAAABAQZ8eRRE8n30rer4SgxWSdC/KZ7+otOEzo+UFD9ZmniYypUuFzJJ18Je3Kzvv6f8AHcMmi6YRwakaHS3jzaoZQAAAACsBnz10RH+DDmAYlt28A+7MCxlUQzSjDlyF0aYn6fKlx9oD55kOvT73INX1AAAAIAGfP2pEf4EHlYPO0Dlcu33mg7UP+2G/OS+vSWgfJ6uzAAAAdUGbIUmoQWiZTAjf+8FGR0k7B7JF8YA+wQjft+K1l7Pl2lFva4U4iloi7jK05WAXlUvU0NMbneg0W8bipRMj+e5PR6G298PZD8Lx4/pmueJhR5QgmVsDRHocNbVfEJGisYpDRTRCNELAzBSPwHo5TkjGz1jvpQAAAFVBm0JJ4QpSZTAjf/vEhSlnIiISPCRtw/Fa1SUZHmNR9eS02TNX9cdJHHP4scrbtJbcvjnamrR2131icLwbB8mAm/JS1f8/MJ+s1WuEEA/spOfP7NBAAAAAXUGbZknhDomUwI3//AKArvBG5U5RSifFlD/Ne8Dro+FmjNGrFQ79lqWS95HaV+HspU897WfSmm3q8POHAYXGxvGACHm9MrOpAgNLObTHeON5xAh/8S0wNVUbBIo4TAAAADFBn4RFETyffT2ygu6Evr15zpphVL9ezEh177sG3lNh72dSy6+LMKsjFpOfLT5s3125AAAALAGfo3REf4BsKuKQrn4+ytpBYNz429+gOd9CbMcAZjWPrdtHWJw6lVT6DYuAAAAASwGfpWpEf4exHtbuDFFtO9aECqMpTVAgxUDNC6R6Ov/Bv45Ck/DDI++GHNS//D9vpyc2nRYpkvKI+H2S9SCYvOaW3tGooBG5951jQAAAAGBBm6lJqEFomUwI3/vrPjgITCZv0Cm/UN86vq9j2om8TSXh+Amp9UymJ2Lmx+GhPvnAhOiY8nnAk80c8/zR4tT3chvitOsi24OHlSIEosC0pK5DFjcJ9H6E07p5eo+PK4EAAABAQZ/HRREs331No5t4RtqmNt6u0B+Dhno+LpcPQ4byTbJJSghklqX8XC5+HsKV6xKkIFe7OUTxgQ0olu+pKqtaRQAAADoBn+hqRH+HsRJVgDTILm9SUFNJlrvLvY/SNT3pkk83hKOIuwzsAgzN2cA2ruLxHhJGC0nlMqRs5YvLAAAAPkGb6kmoQWyZTAjfyghC9HoFZjDr4Z3T9vVMSrvyBTETIkGEDt4rHdxFlDop46jR72snDXzSvTxci6fpWQvZAAAAbkGaC0nhClJlMCN/+8EZPlhAJgR+ivH2tTAf7wMtf7Q0lmwSp6DMeubPofusLc1pk88yiSuDMj9NPivga9JoF3g+xq0Q86goTEG6ZSyNzWHsnx25FbgqP+krw4d7Ztwy2JXZPgQfwac+bc+pWEYQAAAAa0GaLEnhDomUwI3/+/6N/UQQ5pDzZ2qOGQygvLtIfyWYFUfIE/BzO/DrGw19HGK3hwfXoRT80JapmkcAzbWSDpBJ+ucKJJwYRfQ0v3dYjGPpzjoBe3cpWbig6cBTj8apIEOyji3iKmvu0ocjAAAAaUGaTUnhDyZTAjf/yghC9YOgdCGGxMBNvyCtjziLKdC2QZ7FDESe5J9Q9AtgUrneVdY5XUnITqU7yU6iBGOU1CpzwDon0qnZ58WoQb5Yntsn3pLkDSXmlHf4GUB+Hia+8e9WqG1nEx0tyAAAADRBmm5J4Q8mUwI3//vgHICs0H6uMZza6s7teVDeMd6CYdhr9xUUK38FJYrIQPqnYUcIjdPXAAAAYkGaj0nhDyZTAjf/+8Gbu08WBCblR+AghU5I4PjruVw1nmnblE52norPEgNszfsQwVFghKI1Diu0eqa7hbfw1+Ep5QkXdJelsyIpnCozqoRLLfdal5qiCN/T5eg+PO3oDHLfAAAAaEGasEnhDyZTAjf/+8Ohm1TlBNW/tdPMTK/c3ybOs/lY+8ZeligkNk05fNM9wL12M3EuLDYgcCCXSoBpriTN/A1g6V/gLzeZ6XLH44A5B/SCX5SLpQw9q99kGKI5Kg0SBuVmiSFKDmwPAAAAXEGa0UnhDyZTAjf/++5XDgIS808950w5VkT/t3hHWXzd4j+0CV+IqCE8SuscXSYBiVDelxjqtSC3ZHtkbYDFeKciUbuZ1jlpEM/uS2KCxQtBuTVImzCSPE8oOq3AAAAAW0Ga8knhDyZTAjf/+8SBlbxwEJo+j5XM44k4C8vgBUE958Col3eoL9NbbmKsxW8m7j4xRgVojr6fB+NvuumtaBaIlIqjCRjWB+8H8xBIXkvlaBSjiwCBA42PXb8AAABJQZsTSeEPJlMCN//764R2B5oaIM9nfTCNty23j45IR1Uylna2WhBb8CUTivs1xgc8YLMXXuABpWaMZx/XUnylyMe7+O3whopv/AAAADhBmzRJ4Q8mUwI3//vgOOAkGyPnFZDh08wNcfuRYVMfPh/g57ian/0yuRdww/1n7ykQqNJMvWQ3gQAAAE1Bm1VJ4Q8mUwI3//vB61YdgbRGgzS0lo6m3KNl1T26A52Lu3lQ0zu0gITxVmY3a52vCyX+xlNajR4CQjogG549vg9JY3n9NViV9WrQgAAAAFVBm3ZJ4Q8mUwI3//v2UuUJDTVPO7uHUL9d32bbRDn5GpAyvyun5eOVCscEeVYfMq+tDarXNr6Iw0bo8A2x7veHCIxWFjrr1e/Z5ydf+rNut9UgsY3hAAAAVkGbl0nhDyZTAjf/yghBljfBFTH1dyGLzpj4+bIErKAmSdIbdehwtvdaiEU1PfP7BLv/m0QUWAKIqCCCv6cnt9U55oA+nM4cUPpNGIImms98PbEjWuRhAAAAPEGbuEnhDyZTAjf/++AcgTWtokgs1Kg+3puZmyB2zn+Ft34Z7zrHnHnpGkVel3jgu/cRJkseHDVwekGctwAAAFRBm9lJ4Q8mUwI3//vgPUQh0QPQquZqwQfnSdoeSCQCzHAsvi6Sa/8leZL4+gZ5agcDt/8ZjU6OKKewYE9oGmb98U/7Q9w6RYYgUyseVRyuZy+FRnEAAABYQZv6SeEPJlMCN//8AoC0RERAjFSuAaDWIHfxup8AQJpTK0TZjCJkqKKVH0rnE+j2G0jr/6PUOpHvIfb9HtdnRu84uQ7BUI47R1ptvzy8Ih329lBpeqV6wAAAAFNBmh1J4Q8mUwI3//vuVw4Hc4/EYZkfR9NOYQkN6IR+HLjwNnXmfJ3tFRz3h6PoPlGR9Li8C26dewveJqFV+gHMH01AqXp6nEKs3zf/A6nxbhFlYAAAACtBnjtFETzffU2jm3sH5xlgay9R1QXJfemvQPA5ZQM6hlOnPGFXo84lVDBTAAAAMQGeXGpEf4ESDdqTsacWmS87R3boOwBLmPjr9HbKZR3BTbJMNFk7UaY4D+7PnY7F7BwAAAB5QZpfSahBaJlMFPG/+8dPR8sIHRr+KlWnACcQrjiBl9PrT9yPzdGTDPem1rqoZeo+O+V63Y/a74OMloCXYtj7s5ZuipBPqAm6qHry/NiVr+IWLYfvYiQd8br9YnzFyHPfR+RTXTMKvVV/KjPD/ctM2HNXr4HAOamPMAAAACsBnn5qRH+JNiiXuXcPpCfMZDOo7tZbnRWsBJ0tsa5Lgg0pFjK/6ozWiljBAAAAY0GaYUnhClJlMFLG//vgOOEQ8B1SOgWNXwdoVV1ru5CEMBY1K6LZ98XM0dIKZIqwe5RQeu1WVBdf8q18RVrmysbabE4YEFPFSDmDTf9crfxMvgknEd2vk9sWgLG4Ej0lyc7cTQAAAEcBnoBqRH+BB400x1GiEKvi8jqlJCwv1QQndqVvnTA05w24ChMFL28XtsmkzSsUuIWLziqkkynS7jrcWF+o8Wzss8ftVv3C4QAAAG1BmoVJ4Q6JlMCN//wDwzHBCj/9cyI1iJwiC2YDq77/ovgCtBWTqAK7Qd4tMo/sMxMoLVehiq1uxY67X5FKvvfdZbk571Wc3XQkxwSdzgyBXFfYMHtvUybBlgS3PwuBgOeOBptIcLchkHRh5RWAAAAAUEGeo0UVPJ99PatKaSqB2VOb+Jj+9bQe58CI9t/mO2D7iXRZ6LVPwa0DZ6RZvvddI6Fbxs5CUyQ6tXX32iUFx46XgEEzck4WJrjn27/rI7DAAAAANwGewnREf4USysNoxKlr4GH/7yewiW8E18hvueS8dKXaY0ZRvRxC4LmG9VnnRggXFRoWNTsMGYEAAABRAZ7EakR/h7DkMUOCGBJh+lP55SdE5jTC7V7n2PYNjFca8HhrbO0wW5ZpVWgnD67LghASkduDl4tN3mv9lgWKKDVlzuDELPcvXMd+GCzMWgyAAAAAskGayUmoQWiZTAjf++FR2IObN4lvte8kZ2zczUBUIke9zcAiDpc0tIGFXyhAP1rxFfYV9NigpvIUKbs+TB+6WyVSWBBcm9MFQqwPuVSWkoeLyC9MIaVAuYZceC89lbP77DKjo077qt64djt94DKl7X3TkhZMKCZRkli9PMehVycRIJ11KVnyazgfI6QFyZBWKpZ//e1ExmOA65u9Rb7EpDZypnuNTKWfi082KgLh1cxMCfQAAABGQZ7nRREsn3sgp5ri/wKRgQ7Gu3vI9crS3yfW0cNACqpw9OIocYE6q5EFc0ZcGv9rKrcElxkw0g99wSSyz9azLUdIyByPgQAAAEQBnwZ0RH+HsPhkmMfKXjerJPWkzSzLvYro3D5CmzUrpe93QDwpOs1Z8ue1cdrhkVLaJ6fnHRj/QNrVfPJjXQ1luW3r+AAAADQBnwhqRH+E1CtCQU7NGQR/SzEljM35RLPhIf7CuGp2nH7PGqbYx4k/MVtURJNrZRoSWG+5AAAAe0GbDEmoQWyZTAjf++FR2JNg6dJRFEWpLDaHGWYdvgn2Ihhkb4tam6cKVyugm5faVmXTaDsJAu+kZAz2r1gaf/a+k5hP0Q9ZmUfdwQ9RO6sckew+0/TuDcwET5LS6b0bjlYPX3/3sW/C+vHsLkXWQdgV9S/yoRC2BzqxIQAAAFNBnypFFSzffVS161ytJSzivxBMnWEBBJbIqQ36FF83chfte2KE9wjt18/IKs6+mCtLOAckfgeDBt2owHOLwPzOn+zvI5qJcNySwzh6eafp0V1+IAAAADABn0tqRH+BB5WyGKse6lBSKcVv42bFQay/dKDQ6P/xo+hOzobC10h2D4F5zC/jvsMAAACTQZtQSahBbJlMCN/7xIFmvqIfV0ksEmI6ZqkME/+uYEFZp9EgRyfMfI9TwvOdvYfNMg96XA75e+Ug1d7MZsjl+wIS6P4ARavP29AQ7GgEU3ga/qHpTK+j4PavPhJNuvI3vAVPp+hqdWdcPhLbXeG5Bavj/tfM22MWc8DWUPgPRup92tOMWZq7zTg4fRCn0RmDMdRAAAAAZUGfbkUVLJ99nGI9vhkeUtjd0FBbV/Y0chN5PW4RPFtqyuHM6CMMdaYL9nKS2MPgWTHLmFoV57n1YIpZP6NS6yNWHqbcybKI0qEdvxQUm/iL7kE1gKC+fig+SwYYBKRaTsPgiDmBAAAARQGfjXREf4egL3/U68Y/RuCrxLwXiYQlWXYt5tY5DXkNo5EpiJppoQPvg9AOtTRsWp+ZH8D4KriCBrIbQXvLyYYUk6guCQAAADoBn49qRH+Gdbgx/3UyLHgmWsQQA67kWXbzUxYaBkUPtfk5+96VwwOUX427Zy3tdPCl2FQMk6Iaz/Z4AAAAdUGbkkmoQWyZTBRMb/vuVw4QawH5IL3wgIlYl4N9vN2/EPiB18zVQ/xYvRnR37JySPr8J67t3TEEeawOHDK09GYOfGXSC9g8VJk5uBSTrdp8SDss/hdObxpyozwaea3Hbi8nMDNDeWI3uORE/3DMAFFILF1w9wAAACMBn7FqRH+BB5WDztAwaY/r7Z4kRBTDYmo+gXVOFAHpeYiX3wAAAHRBm7ZJ4QpSZTAjf/vhUdhErlYqhggZ3H5bTa1/vToTRb7froaegfKbJwqLbBW+kSnruoRnFWo3+XuGIHxrqAtWXyVLj2RxmDnonpBATSupm/Er3ieXpcjPxjEdhAERQo1d3evkEK8xubffKr67kXpOm1rXRAAAAExBn9RFNEyff9JABS2FcQ8AcUeHDbcg6r2VMcDhBCTlaLaDqBcTdsiD8BmqZH0SKwn6VGYbJ1Q5LmOm6cs6MiJ7HuWG4vjXM+7PUM+nAAAALgGf83REf4fkGk3LuITCI39/Buk2GiLZ5Qm2lilCulAD0MHZ7IiFhpwQ4cN3dmwAAAAzAZ/1akR/gQeVhtGJPhygpFSbr+MRmIRmrZNzW2fSHOqVlo2OShwB4K5dCDy6EojzpurJAAAAhkGb+UmoQWiZTAiPjDpjcgQswobwLtrOt8Kbn50NO1b8FyU1efKfDFwXZG3fQEF2nXHsOOeAZm6hV6ynSwPWwas2HEejxuO2U7eYWeG4BCTxQsqO5lECivpfGEWizf/HBrRDxbzeeL8AMx7yZyACTy3DlgHl6NdrvffsWb+btUZdfydSDQkpAAAAekGeF0URLEfOK2yXld2veMoVsharS5jlxy1WYfGZ6w0Y2W4SyvoZW0vw2TI0HcN+K1KTShZbSy9NVAcHfLkabTxGzbYQj8WsrHVNjL1UdoGjWTy9FdZl/uN/rXdkOQFudoHdTY917al0GV88Ot6qxJTUo5a0bNZ1+ujBAAAANAGeOGpEf5iZsaTpiYK8Ckzq34eYNkr5wTjqRNgZK+2VKyFYuAaNfhszOaZxYrXPiVM5moEAAAViZYiEACf/vgy3mIDfV2I+c4OQeUCbjkIqauQ7q1hI51t1KK0Ty/hpqe1v3iAKauIQOhjuzFxyMgx7FbnM/2jBWLSkwzPM9uWS0H4jWSakogzfe9txEs4c+/QsLeTPvLzlrrhoRvk76vX6F462KFV4lQ8BLFX1JpZOU12xKsmCPjGaaOJgyXQiOhYn2hWzgcqRNjGX+V1k8oj99bE+1vc15C8Q37MxkeCbaMkDN/t8mxkCWv844zRz7jjHMSn7OeiLDTN9e7Cv32YENkK+X2fgPjTA2f9fID7IG6+92XMlKB4gG0/ewdEoHPbTbV5cbTze8c/T9wMuzhP10AuKfYQG1eXnZ1ykp0P6DRPy9Ld6m+2H6X5Ni7ssRrEN4TuoAEb5bQp2TW/e41VdMbUA37GNZ5x7DsYhq5D8keOph2B2Aoqxc+qPJPC1eN8uzF4kT9NEbttz3OMknmnGRbLqo9BolsBfy9A61ilUr1hK+lfVlTjk05qU3y3wHYEc66DTCaL4HkcZGfGwlWCtebzJGH1tSY/TGvlzwoiR6D2VHDw4Qd/QL3sXXc+s5L0t9Sw7eMZ3S47uKfNx/FSmsxh8Vu2q6ia3DDbKLB0cs1CwbSXU5GqIQSYlNKNsVMAv1OgbV1e/q0to651ASYAj8tdg/GWuSkyEBfqOW0VahPSNiWW6MzjEdYDBORzkPU+fSsEZX6EUEJWgIR6Tk5qkRkpHpcm/zO12eo3M4AhR9Cx6B/Bjk6bjXMMI1JbvT4s3OmvDMWhm4GnlKV7U+z5WfoAn+r+XEGmS5H0VzKkjL9DjNzfvJzjcKRQzSjV9OMxXeO3VBTd+WFAIv2qLLi61YRg7R2i0zYBkZiZjznSxv0Oh3XHIYCRtkOS6xY+goD6NZWh7U8+5INbjxbNW9vW0F854H74WEbjrsIE7vUh38zBz8lmukAzLqGi6lSYQVXRAstPq/sjujZpTHuizcX6q5h208Hya5GB6cWOla3tMwJs10rHPZhITBZAHx9Dfn9xOK6H7+wi//UuozeCF596Bj/46ApJUXxnQw+iPZClbVtv32bAEWwjAq5WxRo58IBlPczqDbMQpJ2v8hJZ5eL2AFeiCS8d28CkwIevJBJczJ4bl3YAfkhz6EvtCfS59R8MtyqJAHmYMxFa2BzAk4Ceko+cmThvwJsTo9TwWz6/y6nqWgaEvkLNZuyixNMc/cXcRwsDQOLI4ZyMttZLESIVUpeJZ/6cbN2LT812AnxvEtv81Oo9jf67SuSMRsG6/uLlpm7ay66WEKT2P2eyv5sPPbiv8phYtyC/hpNWM2L6InvRMw1S6W6QQw7ANE+NOljaO2+xYiaex9bOmfqoAKKcGloWbR978SVkvTCVMZ+HYiMb1wAqx/PF/B7YT3YdqVhfszlXZYvWaSPMzU5okM33PRVoLBIGHcW9rG42KMvtm5tZJTaxxpsdsz8X16wDYWJ0ZVgl2FlFrFiwpttQ+K/SBAfBbCBjc4w+9++Vy9ubBe/R1rfhdPo7EO8Upn4ojmO4i0QAQPk/+Njyc3FUAfeGBrW8UR33ylHcXQhrcCZUhYC8/6+rQTbAW1pErnFsaRgQxa84DSprqtocWc1KyDwYQLDXOQyRgEVTWwp0C4n4FCtkK+u/u1SbZZ+WlsAmGCMhbvZttIOR1dvnyda1VEmJ0q4MLNPcwxMdDoqgdP9/cDbj6E8rNAG5pJQbjF9KX6GfGfetGbQA0fy+0g/NVu53wvOeMDlX5aZKNl9hASlsYvBxoiPV+IOBqNBkpAft4IlL1zcz0cOUncvHJKuevXoD5KrbBNqehopgoUlzYTvetOhot3VJAEcPkLAAAAERBmiJsRv/7xJ1U5AfCBaOk8iMcVZeCFVxqmruwhZ1VxVh86clNufhi8lqRRs6U6DSeCULyvbJ8s2QvtuTrcAMn+2m1oAAAABwBnkF5Ef+BB5WDztAx0jJ7uwLOc/joWyLAt56hAAAA8UGaRjwhkymEb/1jexjsxsfi6hr6As1+ItIPaaxdtL/Psp4lmstRcCE43uyBptRSOlaI0NUP9+Rr2RU9EU2/XAcFdqfD0kTm2yOmA4hz5KtaLVYCIyYQrzAIpsbD5TTBRHqAEIhPQuPAFKMNcGmm6KkfE8h/rUhCTb57Tr0SCkvByHkqBQxIkeSnZE3RETT0gzu8clkUB5C2/f9gqjzOEdr0tMQy028X9g6I+jHHVIcOQqoQJqlYLTVMEPnxTm8swUhbukQQqB/V1CMcORCC/q5HxRtMLYD5qtCkgPY4iYs9BYTsiCpbpzh3Z3BHriiVt8AAAABdQZ5kalPJ/4RtMreUjvAg87ctPZkwB7tGp3pNULncn75y63gaBD9ooZ+2Z5BGJuEzAt8WFMMMtfOq7F1nffEJMK0f/ywEmY+FnZVV4/fl1h2Pxh96wGJFYB7eE78IAAAAMwGeg3REf4uRpKqqyH5pDz4TvCO+HJMaGqrdtS376viJGvqRNtOi7CrvGS9ThjJaQu0ZQQAAACABnoVqRH93OX2EwrJauMv/aYKWiM4dLgfPA6VlKtP+gQAAAOdBmopJqEFomUwI38epjsLFC/F3bWYVsVjbCSDz3/NFFz7EJcPUsjeAhu5g0HabZ/pTdPiudhVDGiYxVKcEifZtHJ0mwAP02f3KxpLwILBoE5KpKZ3ak6k96u9QX0BYIZ2zDURThz3E5Ecs7TFwmbLYgIzCpqmx3WZ2+nHSOQT2luDm0mS7tC2WUBIB1L038q70IP7mmgmcfxvHDAXuwIgQQvTDkNbDDnb7qEeNphuMepAZpH8UGL526eLVgrKBJghOlt/VEGAzajEUl3BhFpPgSsX/QRkJ/xFpACLLDGkb2Jne1e6/1z0AAABLQZ6oRREsn812DBeerx++BXNh97FhCAHBFp2YmevU1UrQGvZZfGeZt4lRrbLNgTpILEiupxGlUxfWlIrjdvxwBNEtx8ud+WR8eT/AAAAANQGex3REf9fXXAqIFgtajW1njE1bz+UortEyKI1hn9ddwwTBBszfncLM5umhE6wLGy54PXqAAAAAKQGeyWpEf4TW9Rzxc6AcnNg8AfIm2Rmp/I+vcel00x36c2AMOw8muHw9AAAAsUGazUmoQWyZTAjfbibf1bDnDmV6F6O/WnufALeelmNBoakQ+RD8No/50k5DsmJ/O8f5+6zmOB2dS7zL9WBPkesZ8MQ3/0MksLiL2syu0Os4+DX6zzVah7a75x5i3NZB9mNahlXbFzLVLgWbDFo0MPp2ztyCmOo/mDC8YFYHXaGQYAnLA1akkBw+w7m97niKcIN0hcHT0r/8VhoRhstD4zpvXRFwqgp9DQSmsJ2LaYfRqQAAACVBnutFFSzfgcsWP16iAbHfMo6XUJpab95jekM9opekt48P33/AAAAAGwGfDGpEf4RTMK9Ia6xDXExmf11mxxz6JBZgFQAAAIhBmw5JqEFsmUwI3/QSy/3Mn9LlHVQzmvoMZRAfpfzmYydz58i/v/K7UpAsceKOBkM39soM37DBOWqDpDJYUzpx6MgjQXdFpMMlD1sF3jXfAp63NUJB2HxYwKAS3+X42lams0gG7HPbT9qoElWzMJ58INMDP0k6kDCkWtzZGiljRWKipkGdSBWAAAAAb0GbMEnhClJlMFFSxv/0LB3HWhFl4F9qAxu0d+LmsKyTRhMet5l8ycg7FLJ6bMwnHDvdth6Ny3YdAsTsiP7Q356Lcu08PjhT6If5FExZc2ZvHjw+Z0tvTWfIt2djM4q2QM2UtcHgN2bP19edptftPQAAACEBn09qRH93OX3cYVktXGX/s03NlZ9nnR8Um9cfaQPg9mEAAABtQZtUSeEOiZTAjf/74D1Ed3Fo+hammI+tzPVBVB9GAMtkRSmX9W55hcG3CFWWzlEIUNNLjUOeelsGZvhADsD+Ofy3nAN5GVG/eZg36oP7ixXBkBIm0aSeFatlwVCbRiYm+S35yo4wsvnxxokQwwAAAD9Bn3JFFTyfgOmmneDD0Jh97FhDeN8cWWgHHL9ZXuuK5vJRYDuSt7EY3WZi5fyM6swhMLsBDMNgeDfcSyJYY4EAAAA7AZ+RdER/im4U4K1F63xN3y7TWMTYDhtYotw9N+3yCyCwG7hbmm/x6VomjGrPQRMq2kgD57PUgsO5+FYAAAAsAZ+TakR/hNfANQiYzYypltEIwmskWD5UNtzWBIJZzKZTZZF3JtrR7AcrpL0AAABCQZuVSahBaJlMCN/77lcOEudjsGvy5OKoLXdnzB8pMXFbfBomeiXX0yT9Bi3pzVu8acc1sJ8hB4UGgAg1X7FrVSaBAAAAPEGbtknhClJlMCN/++AchVJ/hJ7fwbdEstXIOKFMPOfow0W1J9MjZSXaU2JvcLoxTQ3MkZ+p/N6EF1h/4AAAAFlBm9dJ4Q6JlMCN//vrP1EbRInimhmXtw9PH/HLVA1X13BZGuO5tuR281DwTZtzSQ+X8/0bJQ2SirHIQ09BynMc2hGsy4/B5wCyAIX1VOyVwH2zU/wrNxXXcAAAAIxBm/pJ4Q8mUwI3//13BGajj12djqpakBtTio3DJ9DBFwgACSqtkYv5h5ytASPgOdOqkOdgnQKmgWojr1xZ/Deftcan6H9Tw3ZA+MEWlDNo1ZYBWRYty0ND9KfhocFL32qclLBL9NniRJkfFCDsnUhqrM64RBnhBgh3Fvlja3uOhq+MRbDNIEtu24O7QQAAADdBnhhFETzffU2vo9kHhP0kUckkmlcOZ1qvSaHsv7FYPzfK4q0p2kKE7Mm1B1C8Hwvzr66y7F2UAAAAMgGeOWpEf3c5fdxhWS1cZf+zUAjrwmtqgIPXmRHv2k2L/2uJj4e2MaADqH/NY1U5BRBBAAAAeEGaPkmoQWiZTAjf/WKRlJMqrSGu8YMSMo4NdVTVxbO2c0DfzHpWj3NO//rb0lcIQgLmD/JdeB2LmNOXjmIWgYrNK2njO9hbOQv+tHayjtl/SNf2g2g4kXhlY9QRMuPkMSkCrA4mXqQYXjex6M4K3kGlpFUGNYB2CAAAAFFBnlxFESyffGVVRtiU2M5XSTBlLlLpuEazvf6T0QUYwco555QlcaEH0Dj+hFfQ6SVisFfojn73EUbm8SrLs/HDJEPq2aQ8uqOTVCis92ohdCEAAABFAZ57dER/inDB4mEZFozMu4+KKCyG++hK3D037ngCnbT0Imo7MlsV5lr5kGDygb2dj6rCAP+8m8TBp/+LfwVdRRtehi0bAAAAOwGefWpEf4TUKkzU19rfRxca9Jc/qwQnAu0qRq2xIkvrtGv5GbqxB1VczyvDgln4tgvLQ64l1ZbfxKHVAAAAekGaYUmoQWyZTAjf++s+OCFIz0x01M5hRDAgQ1SNkLroyGxDyE3PDxh0h/mEHwlASDMpnpSSDpuFfd6ZN1GJRdp7/1NKZ0BSML9Ri+WpFVWYy+oWu7FUhBjVLBsl4nohvDKmz8vEbpKUMsdmHeD8TfRVnZtvTlvIPNsxAAAANUGen0UVLN+BHt1YNOwIBo8J7x78Wox8hzQych4IV8JIX4JAHVjfP6+E363LAITjUT0/2RT3AAAAJgGeoGpEf4RTMK9Ia7m5tNYSHf4/ZtvhHWvqCx8lGWv53wQ3Vc9qAAAAZEGaokmoQWyZTAjf++A44S52Owi7Ki/xHuS/BZp70AE+JQqCbzs6/5q8Bi6TLqq2chJ4O0HsA3vaJtaNjwJK08NQv3nUM1lYoOFoLcHLdvLTHCfibUiXw/+uN1EJnDUQ731KL/EAAABQQZrDSeEKUmUwI3/74ESoS0QZzPsXZCCOIIaZzh4p+rex8ji1xPYtKKYeWQssfQMp3vKo0mHw0vJQBLiKE8ApAr2l/ILTZvWEdw0SQbXmWBwAAABSQZrkSeEOiZTAjf/74ByCUGSE5NLsdewC1Yf8gDKRgA9A3XGNkTmRAIk0zuqTVUB97w21ebY1ZD4A5JuX4GJerYYG1yn3M9Fjq85VjxY3yOfXgAAAAGdBmwVJ4Q8mUwI3//vgPUREOuU+W50ENB4ZTfPgxJ4R/QAgLWfpJ4afx0kcK6OroqS7DFWyZDTIp4V4in/vBYNjSauh3e2aTeWXDJZfmmNfyGj/jeE/EeTuP4327cpUG6wUjxVSODOdAAAATEGbJknhDyZTAjf/++s+OEsdvlfnggh/lGSNWpaGXz9PDwU7DfMzDnOAUqnRhTRcsjlee5HHdMPWnOhXPk2yc35P2EwrAHuNbsWtVJsAAABJQZtHSeEPJlMCN//74ByE7aEx9iXfU0+Qi7cT6qjr7osl+ltza7QEqjKcHDYgjA8iNym4TNocc8MUDP2ywSpM0nUK6tcXIpAVoQAAAF5Bm2pJ4Q8mUwI3//vgRKic7NNky66/c817e64Ip2yRLzp3SkaICilaCnIprlv8SQe7wRAh5lXA9b8r69v0/THs/oLcawFvELgXfQ9ZMLp1MwPkpBjFsgw2mzZd4z68AAAANUGfiEURPN+EdaQm1vW/TINWOBhDWfqLGez+cStDKozqO3ZaLVmanISgcBX9Quzd3kidv0xIAAAAGgGfqWpEf3+cgskCm2JriciFWFR8KQwp4tWhAAAAXUGbq0moQWiZTAjf+8P9YCJ2LEHgSkHxWPcIhuhg99q36oJg+yoD2YMOzEXYsm24bpU7e5HDTbai/djWZ57WEvP0X74/hSDvW98TN4gQ8FYDbsgAomfxCR6enD3bgAAAAIFBm8xJ4QpSZTAjf/vgOOK5ljY+B/+L9+6Ag5nJ5NgD1rC58rPGzwaZU3TSod4Momby9YNlv2Wut5rt0gK6/Y3YwVJRjdwKaYBRIsbzg3k2OgRes4iDLyq2y1We08s8H5cePlc8GGSMvtVpkDZXgQ3ARvXgVsnai3RWNmYHRmfNrWAAAABGQZvtSeEOiZTAjf/74Dji3z+EL3mriZEn7qSTHtaSGpQSz7vu0Qa1DTfh90G8QU81BZ+J+YsG68GTxDc1tNVG2FxQvD3JmQAAADJBmg5J4Q8mUwI3//vgHI3KkJj2o4R7wAdkoxoAPl6DDEXLEUPxOrh3d5SlgY3fxMeUwAAAAHFBmjJJ4Q8mUwI3//vgOOEudjpLG75zaU3JffQF+0lnrfbIzuF4nbSX2KhtapcwdzSYt5xk/gC0dG6joJPOtvfYDyHPdosm0/eNEHkREpQU/Lx3HWTyhi/QVBnzgdkBQv8ovz6Jzex7/D3vMDIqrT0wYQAAAEBBnlBFETyfgEm4jKr15Rd2F2D4x2LPtpGUSKL3xrnxXz43NEoj/PN5GE+rLyJeghmbwz2xiT1z1CL9RYYKsz5tAAAANAGeb3REf4puFOCunliFzADLMuR8YF8uYkQJbFv5RTB0u6zvuSXxrl6rDsz+i7vbod3bBSMAAAAvAZ5xakR/im4VKd901cdgCtt0Mf0BD0pitdn5JUOOTT1aiDQ1D71X9AF0lVJsK4EAAACQQZp1SahBaJlMCN/76z9RUBUcvPliLciIvjqR73/pxCfQ8GX7jBwWgVXkAwndy884ukxbTKeablPB5O9VW5A2hak9PYq/Z45Oo6shsBUaTyA/cKmxsiL7o+G5kUuQ8MYCGvoKy2DKLSe6RNtH4L9/+rIRsEmOj8yoDy4VgxHPKb/WQm4tsTdBVn/g3+/tzcq4AAAAMkGek0URLN+BH0uHOpR/2lJ+qsy4cNxEGsg1Os89HfXB7kgKMsUuWa6wa3E9Cvz9Vu2BAAAAJAGetGpEf4RTMK9DxTPCJrwplTdrNLTARpp/H+2xqo210z7XBwAAAHRBmrhJqEFsmUwI3/vgOOK5ljY9oQx//XXAhhGlJBmz1SK5mkHqP2hvO14o18d/5fUUVISry+PxljiI49Ncla4x7BQYXDweGuz/vR/EY/E2El2sS+X64RWh5rFNorrfK7VEZzIHehE7YBLUZYyCOgseeuPE/AAAAEdBntZFFSzffU2vo9kO0Cnibp453nDAGgY00PiyDpgtljOwFy8YbTMWPWaJg1ksU2ghAt/uAQ6k2TZLipHJGD5DZkUzEzwc4AAAADUBnvdqRH+KbhXAZA8EAuKk34jx10uATWjooIQO4t7y/vwPiFVePKe96ekG16OebeqUuepNxQAAAJZBmvxJqEFsmUwI38oIQZY3yoDaqyVEGIIzrxX5up3P+yHPuwxQuIh6RmqexMGn2JW4TA6S8a45F6aXSoHDu2LRXHvNy93pz1lWEn8VRiuUaT4hawwlUHpq+vaVckWlKSp4IpFdIAZg3tMDq/eCxf/2BMIs4+tEBE2gLOsQ3IOpZN1Nbh8+Q39DhHggVbCQaqg1VnbmMLgAAABfQZ8aRRUsn4Dppp3eFORjRG7B5yFTmmIQ465a/eCYalbPok7MkL2t6szR+A0CkcE/nExx730z4edybGjkTzNz41lOg+CSJyI4X3ooGwbsjFtXwRjxwm+LP6VgAP21jRUAAABHAZ85dER/gwyVRD3jGC4pS6OxRSsTP24dYF84UyHNnutQawSpCF2+m9LZXyJCC+IxSaMtdFhc4BHjgm9gCKjgMNKJ8wtLHlYAAAA2AZ87akR/hNQrDpfZzHTty0SYpUlNuCqQ6/BRPgxApygEU5Fkje1Z2NJ1LTucnfFcqhZof1iBAAAAvUGbIEmoQWyZTAjf+8UQb2oRIDdyJBJG3p2eSRku//vGBvYLMgzmEIVmMtrZl6/1t+3GqkdEZly0y2btApk7ZYZsIsTNb8bH83WEuTPj5wetHe/hfU1579gz9eJmyWaFEN4R4kRUT5qpZ3V5aB0GqdGynWTZ4tnQlWYylh6GVVEwLVvNFqWdqS5yTosyb0fNLjd6vhS6+1ch7S8mxV5lmhCvgGRXTY5b2+5TgsBpT8dK7WAunqXgh1UZsPH2jwAAAD5Bn15FFSyffGVVSHdSb2LfjLSUp6GX0WfYywKsLovtWybvi56QWoaMzSn7o2XmUV6WjcC/zKrkS2qukaOC4QAAACoBn310RH9/iXZSQOIbeTW0680W96zcWiFcp4kIDofTn07/nwFewF7sA6UAAAA1AZ9/akR/gRCXiXa4/WbZ+NyaVYeW6Pf/aBVzqs5yIgtGflctsp+WlKrincEJHtYZiVyZ4kAAAABrQZtiSahBbJlMFExv++BEqXFF1Gh1OTCwItMBlK1vXU7Y1ChujfwgDSEjxp0EbQ9CFRRs3Ld4UN+pKee//cFrSip/MsXQzRJi6a1nf5wlbHjxkRq5Gl8WS/GeHRd97LN6POgir71VTPszaxAAAAAoAZ+BakR/dzl93GFUqBjx10gSJvD19gv7sp63PpHZI6Ie6pbZTy3WoQAAAI9Bm4ZJ4QpSZTAjf/vgPUVXdpG8vg4rlnY/ft0DGCD3fGKCJUJtOkcYpkI/rGxid5A3QTlupDhXOLGx7FAh5xXChqo0uaNM8+rYrg+39wO/wMNbPlztdpk9fui2dj0P/mWfeH1Ug09qs5t6x/FiWlyKrutmeEX6/v0IBu5vp9sVb+RmDgKXu8CyfSKKiKb4rgAAAHpBn6RFNEyfgOol5DugVzX2orzM1hG+r8MU8NyTphLAcR7riuJhLCmXsqR4r+fP38yAn52AQJ3oVcR0wEHGBUJOmK8pqozKbzUBCcuWr4/sv4PosC3bNQ7d4cLNL61+4Hajc2av2ig15vb1u1wugeDCXJaiz13UgQDQlgAAAEIBn8N0RH+KbhTd2l4QxwqahFOEir8WO4bHJp8S/K7MsINQWNL/SSuhWUuwzya5TSxVSamU7UEGgtWToCyOldc6FVEAAAAzAZ/FakR/hNb1HPFzoSznrKuNC2BC8XbeOVatSDLSeW0TBFsdg62GXXU9E7NOBwisegIBAAAAsUGbykmoQWiZTAjf/AOiKyFAmRxUiMbEE4XKUlvaqkfdnPFkuSrv/uZ3Z2K7U4O9pzkfIf7vPupN9AbVD8Dwu0MYNhQUfSc7G5a1t/833BIPTo9W4fFuf2qMVYdA1azZtGYexCZGVNJ/9/42W7lwA+c3UjwmvDt+IQ8gN/ySJxefGd+6V4IysnaZg/17TFedcsedScJPSmsaAWJV2YZr9PQrnZWmT023Hkz08fPsg3CfQQAAADpBn+hFESyffGW008aR19uICCzq5NQWZ90NfrmxIeDOiit/PCFcxaNQKPjDjQ7qd4CRyqeYilJwmJpgAAAALgGeB3REf4BsK4JrMmvsK/sAAGwlDBfyHnOGZN3YuTYe9/CpCnyrjEBzHGVjUCAAAAA/AZ4JakR/gRHB0sI2Gg5rX+0O0q2TvDPSc97UgUfQvJFcWSBSyl1AmJl5+F9Odi9A3zNqS2n8PseQabVvy6/BAAAAe0GaDkmoQWyZTAjf++A44ria4Qw4WCiSmMPHSR1OYI8sMHbiuDjzddiGMt5FFfQsG3Xg/s7xrPsrUAb7CZP61yqPvjLICQfvhRYvztPo6E8M402MdjGVc+5MoikSIp9hwj+gs/Mh6ONlkOtd6lmuN/ed6YdaokhXg6/p3AAAAD9BnixFFSyfex1K7S/556qAF2wZ6eiICEv8kLtY3rwO+XcsSPv9lHbtQilvH7t6S2w93Av2M0H+BnDw0AQGkaAAAAAyAZ5LdER/dvu+WAiZLVxl/UCxr+jB0mUBX88+mWTBAoHaD9tT3RLSGpwK3dfZUneAJAMAAABMAZ5NakR/inDB4mCq76bkuBs3tkF946bT823lmJeD00/zi7mcXVT/glUaS6jNfGEtB5R2xabRDQfa38sKu7jTT8gvmqqlDa7KdrZ/cgAAAI9BmlJJqEFsmUwI3/vrPjhLnY6p+6D40pzcI3EESZhkQEBCdNnAeRIPFFVSkNeEcVI8X5XssF+l+r0Xo+G2YobSzOdbqNWcfU6YbitpLAmr/TX0rnj2EMlad7XNyf0cLZ6ug0PPBXULiz9m54l70LAnrx45mASDWufLQkxiUHCgzbJ5tLv8P8fLTViejM1LgQAAAGZBnnBFFSyfgOolVYNR5FvP/wKasCCjnYQlxD8XXZsDY1OG4dlU0osuRUdZ9269GrNsNOazGKWnVWEXEQ356Raa/y7lNVGogDCtXx7LDbgwD5Dg7QfELHis1MwTs25skxKF8NR6CxUAAABSAZ6PdER/gwyV1/gGl4DoPPqyLS9sHc3sEMPzJO+wfSsvfs+SXVSm/M8+KeWiKtOiUBDrSQJ/dV0wHz21x3Nlknh5MfWPAzaoA5JbIyBOBuVi9AAAAC8BnpFqRH9/nILJGk5gad+eLsoOm/85p1qMApGsq5TXcltUPCfAAuJzBXY34TxlyQAAAMtBmpZJqEFsmUwI3/vBmqSMhNsPuU/Iade4QHot3LB6WrIq1REVRUWuTSSkVV07RY4SBhsJu6bBxcmmu/7p7/XlQZ81iuMoEWB0fIuqz+0F9hRriMm+Zu++c5UoVrMVXoe4liL4u9dlcciYCOKVOu20Yyvw0sSQzjqY0OXZVZDApK7OgpTk3QL5BbxEbzQlLUIc9ecQXGeD9042d6e5Frg8c5maepHMbf/XGm8v7g+lOh4PjQiOo2rPGPN5ca44oRCfG6hjuP4Xue0SoAAAAFpBnrRFFSyfeJlWxHrflwbYugGcTSy6U0LQpmxil0xIJFwT8XFx7PNwczoCH6Fhbo0X15NUcnc1rsuxZVBBQnNcbPv70rYrdQG974QEotol0nT2TDHb+Qv9E50AAAA6AZ7TdER/gw5gGJg+c863gfTJmm26qPpKYjlbgotHYmhUqPnWNt7iyfWEOKe9buyfpqlqllaZlsDSnQAAADUBntVqRH93OX2EwrJauMv/aX8LwP5QXIKtjPfqQBiLL4Nctt04kl+GBFA8tTN8WaPF6LkgQAAAAJNBmtpJqEFsmUwI3/vrP1FV2ybca4m7Qi+CcCwEN6Sm6YngKSVvXnUfM7sChF5+nlbWXM16e6m0eP3IRF/opSl4/wJ8Tmp/j912lrTy9yAirrx3u2AijNyh/O2sWcv5ZF3bAQNrmKGJ5rRtgcXJ2EyIs1FHVEKdlaUfrkIuoLml7BajyLWU7EXhV8KiyXp0EkB1IWEAAABIQZ74RRUsn4DqJeP16izBqJDziqWfWcM1ER08s1jzQXvUMa1oZMXncVBGC9ZI7nV/tGPht7tK805Y3P+T/QfODA/eDIltS+iRAAAARgGfF3REf4puFOCunliFzADLMuSGFg8oLzH6dqAnVVV0V/5aqCM05ibhk+ZkUat4TgBxzdHsk8qTta9BtubHdd0z1ZwtYcAAAAA5AZ8ZakR/hNdbIqL3l1+IwTqfo4tGdH+aq0/IiizAMo318AwUtAm1jEwsljCSWs74wPi4J4g263FBAAAAjkGbHkmoQWyZTAjf++5XDiz2R9bMf+PSkj74lbvgf+CLlKO2lnAEeCo/BovEKd8MOME92netTHHgj2Xxbi65aAN6935tsknj3pq8q8prCwziwqggVz1hNBwrQ1X6JhikyanOkZST2ZeJMNQ7Y9f7rNXp1lNxIE3MyU/nHtKei86tCmCb/jB4vfQ9xWLRyxAAAABBQZ88RRUsn3xlxPNhKxMxOguiBxRkNmyuAT/aa6QxM1PpGwiqB8JRZtraexcKZFHXdXzPREf+tWX2PRuHYEt70OEAAAAuAZ9bdER/f4l2UkDvsZGNAf/yJYoKrKSiCC4Thsosst1DE6Zt2qgCipQ/85ZSywAAAEIBn11qRH+E18A2k13OofglWp1hqldV5ObnQoJScCal5gPbTIjyW5ljrk8khorjjE0JhwAgDO8DAXOtPqLGK9BvShkAAABuQZtCSahBbJlMCN/74DjhLgyOzpkIzRKgaw5w5N1Q9CI2CzFs+eee7UKL04QlzU2OMinPtJrKhxVooB9zLXIdj1wU1RccAL/RK5Wir3p4Z5Rhz1udVe6K4shR4iHm7EJ7ua8uxY5WA3d2NucMe5EAAABDQZ9gRRUsn3sdSu0vuq3oSzNlAM8OUuryflbanNQhQdg2gcKDlBVHwMruBd8g/Ot4PFgifoE/KP6Ohum7YS9dcg57gAAAADEBn590RH92/rZMJhVKmtJ58r16K0MJGmMWg8xAuD/vkLLPx78bVItQ7HSx+t9TsrPYAAAAQQGfgWpEf4puFOCunxZgudv849O8E79Kzy8NTAnO6eR26TxQmY+gxXfNtnCzUB22XkPe8BFERNAs2Mkr21PBuWuBAAABDUGbhkmoQWyZTAjfx6pjxyne0Ht21MeCz9epidJeHVRgt2OzsUDZ9fGji/m34qapaDq3peBBbvkhLBPU43FEf8C/nr02F/kqRsg4pwf4rJ9PEH4DvxUg6fBJQ7D2PXTuvDFf4X/10q3JA3YywimHAzkd07MRAHpUkRpMLgBL8pLEMwRhbPWVpnT20qZSuLPZcdFVIiFASoAoRw2NZovY8cyAqYzbIPWrSsq65p4K76JrBR5ZOJlpYoaGRw5W0g2bASg9QEtRnRj/wfIJlboNxF8URYCsrV90PwYzTaGML8kxduJ1HEhaPv8aDti8cRlE7XHaGXPFi4RmNryyxl8CqT02GdWdUGWMt1v+HtWAAAAAUkGfpEUVLJ/NSlfleNWMCRxiLP05rGSLgkzLFEKRmQxMuJmvzGpyeNPxT3UwwrV3v/0iFzkqNLK3bHJnTrwhtjxlzfd3LB7TBGgwkOn3N18GCDgAAABaAZ/DdER/zdG63bG7KwCn5wCGa3osZLctRFh023EneSsmn3MzWC9KqCBDdw8pzx3PPwNQzEg6Gn5QdiDUHnMAGDZ3RWqH8/GR/hNAU8HnZ5GQH7ni/2AUHpBBAAAAIQGfxWpEf3+cgskDvsTAlIkk3Qv4QtRWQBJxMPzllMwDnQAAAPxBm8pJqEFsmUwI38ep3IcnxuTWqQj8g/rgSxKei33dMugniCkcN8GPABf44Tw8t1Qy3eEsRn18PRXxXi7KFfT3IERwMDz7zpA3QCQSPyLpyVwXLhM7PhboSU8kbfPO1nIY8SIm+UMfGoyE0RMWGKM0JLzYaywMKx9vosxmymYX3DSB8gROvjqj5bAU+Je92mJJE3ZzVPQnpxpwyoFzVDoLe998j0tIsq44A33luvrT6F1Iu7OaZJWo7H1e64kDqoUz6nvhWH05mQLuC2M5Y6ktRKSEcGarBYzmLNzmHh7vugEWRyOLzOV3ZQVWiOAToNJ5ySrGU9bIN9xzr1kAAAA5QZ/oRRUsn81KWB285AKVSC7RFY+doT+AHMerKL1ioOnmHy3pPwHyRwqUvphrWUAjK3ENTLIIi9eYAAAAJwGeB3REf57EiBPUsI2Gf4tmrnTNH6no8dZUfNSMntiCE3jpMKWXLwAAABwBnglqRH93OX3cYVktXmd/7NLucyKFso25IGF/AAAAl0GaDkmoQWyZTAjfyghBljfFGc2weQw5uG12FPp0eBzDiQcuXdbNVOSPW3hp8N/oJG5bAO3Zb373SP/YhHQc+p7lx0UOuTmVc4HVR7yuvu9I8eIWq8uvvC2yoCvXDnsvL6Ml7y1Qbsl4kWfJSSnIKW682NKdd78et0c3c2fAD5VjJnXMrm9zYiL5tXnYF3cZwssrIHOcmXwAAABdQZ4sRRUsn4BJuIyq9eUXdd8kUEVxBXvZDstOfWiaasOOGXz6d2BweVCpPVueVKNIzwQFXuFUU9XBKpnS68hOjBnpCQT4QYBqFRqI+fPbQVzfi6cOp2mU7sSvCMpmAAAANwGeS3REf4puFOCpcPJ2WonQTLkhtA5MCJGSAixNk7uE8XY5myUHB0SlZsop1BQ/gkY1XDSgyIEAAAApAZ5NakR/hNQq4wqUtHk0A/hBBQEMC8kIK/T0jJcco6ztGO/MojFLnPAAAAB+QZpSSahBbJlMCN/77lcOEOGGqjflHZLVcFnf06lApNkw1dyuEV9YbzJBRtMBgx2f6FglR2zHhzFYrtQrtFgGZ+DglcCrKvwM3YU7JhUBsKErOd3vij9YsCRfhH7aSRLCbVXwbVZgzfTpOsK2Ik5aDiH8Fbx0ZTn2TUf8Z05hAAAANkGecEUVLJ98ZcS8h1iXBUMZDLJz9G659Jpfnv5AxC0q3dUwbyNGgKrSkrRowhNA2+4791C3SwAAABoBno90RH9/iXZSQKtBXKXppdd1hcPZsBwcUAAAAC4BnpFqRH+BEcHSwjYaDmtfG61R47GinqanteSqsWgtBpO/fwqOA9xgYQFYm1iBAAAAbUGak0moQWyZTAjf++A44TvouN2t1o3r0TOX19VkdemfIkJQNjJtg2asYKxFD5FgdXiAetHtc+GXPVttB06G314qyFnyAAJOZ29a9/0H+3xR6ynmfY/umA2SihUcWtibuBfPpAHobrBH25owHGAAAABTQZq0SeEKUmUwI3/74DjhPbILSOAstfUwFE4bEuwQsadfXqn7plGI0oW1t1Y/wn0dDyxQqMCm+V/gNYrsp0t8yc1kYM4KZZ/Ep47KXKTj0fca5jEAAABqQZrWSeEOiZTBTRMb//vgOOEuDjiPT6VQ+iwpFIBPfM7Vc576QnIa86wIKthdGtodrF9dyzF6wRQCkAGjjRZKD9pc+qeVKQR+yM0R2LkRny6G0uUgLJfUm1EaS0RgY974VyZDjM1LWcR5wQAAADkBnvVqRH+KbhTdyySQ5ZqJ0+qpbnfk0JbeSoqS8GFf39+Y3CXtR+VXfNwhOcPWv7xsKzRkAg4LHFAAAABwQZr6SeEPJlMCN//76z44S52OVALIq71Pd2EwQJX7/N3fXhfhz3oxlU0IeuIeFSImpqXiYeUYQ8aIei7lagfq9dJeClxCf3/FRuzxdmLsPWEpBDNcG4b7eaMqmz6BvHDqlmnYyHJrEZ8IZKKB4lJ/QAAAAD5BnxhFETyfgOahWDclf403X4Twewl66rB03AAAR9ia2izX/w64YBBi7Jz+yLB2fmcHmEV1G0PRn850sHSM2QAAACoBnzd0RH+DDJXX+AYxXI2kR0rdOucxMrUy+Pe9iIFKs2gPngM17pEs73MAAAAdAZ85akR/f5zgScsRbQaQhHGzw2FWBvNSvTqWvH8AAABYQZs7SahBaJlMCN/74D1Egm5pgj7LpdrUXh+DK2shFxcZendNqyb3B1+J2NOtJQLO0BY1+ExmI8pwJ6OXctPqDeFKHLYODoMFdurTJF+a9y3NPu44TpmggAAAAJFBm1xJ4QpSZTAjf/vgOOIGX8Dywu5b6HQtuh3ftgTaLAoHUTu8zVtxvvBrwf8pq0n+w4fJtp/GtGwKM5hmO9cHCkxx0FIpHiSkAHgDSjdJ9Bu8/qv30MANfg5dNOeEtmWEqPoExrZqFj8JsKEB0TfzZtNFAYsTE01g9Z2gRAsp5MxV9DO5F2sdgezsL8ViWTohAAAAaUGbYEnhDomUwI3/++s+OEtQtJaAtNXXgd4+LqdFPoVeYoa3U/EzzhI37OlznPU7rmyUThBJLF421rLihpccajuYn/QNJFilPGf3pTu2b7OjUvIJM8X3UCvvvycXDv1blPNG3IoYzzU+0QAAAC1Bn55FETyfex1K7S+6rehPNxk1X2rNShPMmLN749hOCs32PoHwwGQnuRiu7dEAAAAiAZ+9dER/dvu+WAiZLV5nf+zUFr8AEzE8J+xc6y+E6EtaiQAAAEYBn79qRH+KbhTdyySQ5ZqKHVStxrf8zUErUDlRLkpwXJGWM5A57aENWOIeQhlA6EA+NMSo4Nf2axnu6UPFdOUMdBtyzDxoAAAAaEGbpEmoQWiZTAjf++s+OCFJBaqu5RyJqCI5ciFy+QfdViKCVpnQDNKhtjLsbBDTAxQkGtiA7qEUco1DRu3mLPro5hMkwOFSJ8jWf1cptWx5ZqYdGH7hd58lsJBKjBOOxqM2gWtMZwhsAAAAQkGfwkURLJ98BvQr/uX5kKHT5/lStdRB3EpNtMxwioGAW9Yj/IJ6BkSDVCrLUVYjVeHAMS+JXklTR7CAUKq0auNYPQAAADQBn+F0RH+DDJXX+1g7QyFsak3Ft1Oi7eT/WrVLikgNs5g3WmkArqxtoNirEsBZgIMNFCYgAAAALQGf42pEf3+cgskDvsTAnaZI7vQ1BYaJISL78CwSAW7O8MKN4TJcSzKr6QpQQAAAAI9Bm+hJqEFsmUwI3/vrPjghR6aRbs42PAWhfKG9hWLBQBYO7VhzdwkmyTb4OhpVXnuVoSIE4n58p2LyiEkgihsMcSjpZIoAbvLWcxcbbUxw9yvpC4qcgM9TIHk9bnA/m1AEfoKukgxmshP3lvq4sLFHjIb1Bb+gDAKawHICIFXSjngCSxCAgl6WAKcq3XBRUQAAAENBngZFFSyfeJlWxHr/xodW1dt25DUmu/wD5rm8TWMxQj7U8N4VJthmBHDdc8Y8kUy0uyF/QVRg6T1Nw3dOMMsOpcspAAAANgGeJXREf4MOYBqx+AWriBGIRbd89ihAU0EQj7cqlqBeUwwo+Hl7YF8BhJvIxqmsu0cqLbvOrQAAACABnidqRH93OX3cYVktXGX/s1aluoN3CrEK5rAEtDv1gAAAAG9BmilJqEFsmUwI38oIQZY3wS9iq8Ru1ZuwHge5YzmrwHbkgMduq9QhK+TbcU4LXoh7vtGatuH+orZwfaI43jayClVBfzKYUUCnemRjycGZN57cgRK9ea1OwkugTmsWQEppkBNh/JdFjjQyZ4JpTcAAAABSQZpKSeEKUmUwI3/77lcOCE/ywLc6XOt25/biPs8yLMMrhhhyESShoB+zNzfnIF0ukeMr6Ju6bgeroNShV6SlG50RR9jhmJI53KndzULEpskTgQAAAGNBmmtJ4Q6JlMCN//vgRKhGBrYv0qobhWAfXvyR/FlCNuqZwJVrpTIPYrZL+bMPL0iiMNcUQ2zWjrP0PVXRdSCVrfqzazsxwiJ/Zm9QLkBsqmq6bw0xk+Bar3U6aKO2MHQ7w54AAABjQZqMSeEPJlMCN//77lcOCEpHhvk51D3BWHT5CL5wZKLK8ZPkfOIV1GODU/wlBTykgfVfM7+jBF4bQaduzQ/+3A+AvwaK0CE5cHbjVc1pJ53H+StJxG9Ua8/DrUCHl8SGEn2wAAAASkGarUnhDyZTAjf/++tBKhInjrA/ex2nOLzwJ3o+r2vWQ3+t3tlTqjSvEVovxTIc8Jr3NY0ry7qzN4IljCb1qcgZ7YvzWMn0aYd5AAAAREGazknhDyZTAjf/++AcgmTgdQc6/hgQ6BptGAKfJDuajnh3DwMwZaWGpVqb7Mh/5bvAT6CAEuDJtJ56z+N9ZmTYlY3PAAAAZ0Ga70nhDyZTAjf/+8HNT5G8Rzz6vkY3Rq88twFJVOYwb1Xg5/k6UGi002nITY9rzS8PPaCmc5nqS0Vh2hJ5DZv/OV29OQcMFyct9AKS8kL3wc8YwFvLGykWfhhLsilALVhO0EfbQJkAAABbQZsQSeEPJlMCN//74DjhLnMdk0oGHhcQFzV4h856+SZ06+ZzX949gBdXZGd9E6QP+F7soJy+XXtQh6tUmGDVgsotechBGEMrKDbwe5/+7BdVED/4EqLPYycb4QAAAD9BmzFJ4Q8mUwI3//vgHIT56s8EtfCC7Nx7BLXzRep/pBVMb6qXu2U65llvCElOC5o5lAlBAWmxsuQ6hG/bP4AAAAA5QZtSSeEPJlMCN//74ByFT5oRDGgg8nQkCHzopVZlIYNDj50PxOrh5uMMbamN5s9P/sz2hGglU1JfAAAAWkGbc0nhDyZTAjf/++A44S5za/LdFJ3sHpSo66b88l6UQOHPKyMuIn8r+SwswHCeYG8PNQSxT+5oV0IxUUsMcQ2ooBv+/Y/FAlBB46wXopImuyCaT/Y0EtV5MAAAAGJBm5RJ4Q8mUwI3//vuVw4S4OjJsHVijyDrIVPe0J5xdGasNWKBtM9uPh9nF4cTcB0XMQNeXo2jAvyPsnd8B9C6JaPfT2A3eKrYalWHgH9tUOCy9Knavp6PwTVwZLwjY3e7jQAAAFtBm7hJ4Q8mUwI3//vuVw4S52Of/J35LxMtdpwPUYqk+ah1t5ft73SHY5tDI+UGs5XnciWB4Bckb81EGVyf92Io8RDVoOwgtO7jiYmvZj6oahKvO0TUvdWw62HlAAAAQ0Gf1kURPJ+A5qFYNR5FvuejVz6I0sUwXTmatChezYmLm4QESF4hGnzo+2N789GMiY7Hg7gSQ7H4n+uPL1I0UjWopTkAAAAiAZ/1dER/gwyV1/mv5SPC2NY+LjsomX/n03iAQv0ER9NxIAAAAB8Bn/dqRH9/nILJAn8Gzsc+imyjJea4qC7Q+WOHNZVfAAAAkkGb+0moQWiZTAjf+/2woAW0kUklgimlgFCIb///hyD//bwn+Oqnuh62tNC5+WgKA+8xn/nply3uOc/ur2f0h1fx1vLWFTcNCmIsUc8q93KPKUmAUmJhrNfrZikefSi8Pp68VDmF+//T/fZ94uc7Um4CesacTL0HWpbC4pwVw6HALwAKhczR9X+xqYErJKlS4y+AAAAAKkGeGUURLN99VLWI5lTeYG9BPyle/82GFMyCKf3wI756PFoOBiQwNLrHDQAAADMBnjpqRH9/6UAHHyCU4A1cOHRliECctKkLiOmAbCVU75F6EeKMbMG1VNe7P3F+zi6z2GAAAACKQZo/SahBbJlMCN/76z44RMcXI48OwBOWvEVAsQP9NTZluNbhD6AJVYnC7zW7VxWM/HOBQiFHH+gn6l7YBSiWoIeoGHmNKgXbtRVNg6EmvoJbA9qtyZntMZgCmaCc/vz/8jRR2HQyWOQpuYlCjrJKp0G1cCrOLzK/bZToKz16r79FSV7J1Wz3m2J9AAAAWkGeXUUVLJ+A5qFJ68onA6dsJzs09N4LvIj3cDBjRfEdWApojACyNY0OnrR86nreca/gKP/udzozcD47DfiybsBvJNDTMeXU0IZfQ5QUfUQrr8Uho5P99QOyDQAAADgBnnx0RH+Eo2cLvj4vNhzCIRfjlupd5qSQ43HazEjhH1jmlnzSvnBcw28QnLtZNA1sBrhY1L+GTQAAAFIBnn5qRH+KbhTVzL5Iwd9W8s9aJtz6PoSUC0gyKP/IlosfSF1jR7mDYgMoJzHwJmTvCykmr5d9a55uBDNKtoDcBCNVk/X7WrPcvXMd+GCzMVi/AAAAr0GaY0moQWyZTAjf+8HPtQiKzcGjVgu3kVtj50fuo8YxUJ07ui4ClGPzR4sW5Oy4COuivineLE9MqxaZ+Om25zN2GqsYwHwOjbCmp/yOZ9thUdCfxibVxPf23ClpXIFCx9pfQ3L3gMhk33abt/kuTykbkgMonPh/peH+AUaF5KyXel7QmQs8VNnRi/oq6rs/xkPZZApzDDMdJbqCvanFzHalDHtkp+lq2f+uQsCGxuQAAAA7QZ6BRRUsn3sdSq3L8Z6sV6pC5e2JBxZNJinGlzs6ATElOuuq0YT7sb+JIlwFgYRvazPMwERrZYQMplwAAABGAZ6gdER/inEMd83hLolJGCEU2A5LWgNtuQHR3ohRJdjldGiSpKiSaCj6sqpbpZtjxpR8tMmPHBZmo35RU1qsmIxgL7vhMQAAADMBnqJqRH+EaK8pIFP2GKqMXBdJoLodSSH+w/mReILD9PWSmxDPtHlmi0aTpQN6QzAxa4AAAACAQZqmSahBbJlMCN/74DjiuZY2QoaXm5FeZph47u1XtpsL5Az4WSZShugtci/xJ3x+JchRTFcSCzajdyIBFD0EmzyuRUcNgfqEDW7dKQFFCS3LIIA75c3I5bEeWMNffND2q8RNvHgFxIiQTv80jmOiJhMyx+fiURB9txK7Jr/00sAAAABRQZ7ERRUs331Nr6PZHsxt0kUckkfcgh7z3n9GmoeN10UT9FBRzEdYc5EHASutw1r3EgM7+7iDD7kqUo/s5W4KMEx5YygXQS8Tag3KI7Y8iAvxAAAALQGe5WpEf3c5fYTCslq8zvxPrc1+6JFYuIG9MO2VN6us9stILo+Cd0DlQvGUYwAAAKVBmupJqEFsmUwI3/vgPUSWjRxdIcdgBijf15+5OcRgW/d/tWcjbO3MZMV83KAk4iqqV860nMKCr1y37e9s5DWqpUgfY/zSEY7/mejd1nXV4y6ASyy7Tidqy+itoBo2ai8/KleueofaP2gYMEwzc8CsZ3c5zZFOxH/DE1t8caigPteCEIVRz7SOG5MhaZRao3+/RoS1pQJfJMw4indglMRRujCISvEAAABkQZ8IRRUsn4BNWWWPzM46t4L3m0bngUEjMFtvaLh11pCL74yxwXywy7hSWyjJhKVUryYLLI7LEmEa38rtO2jQrJ7qp+4gHE4WAIkvSUMYkTurlCHxAi5I4XvUWaZ6Iv16g8K3fAAAAE8Bnyd0RH+DDmAYYFgtadtWEF20O27XaZ+whNwYc3jC6EmY1ypgpN2m0kur0wxTdHYayalaRxghl83RSGlQP7QpaJhHZnAs8TP3vDYer5uAAAAAOwGfKWpEf4TUKkzaVNQT7JA6ND2nVupGemBJq8ChfhNtpxCX4JwNHEZ9s7UeKcUe4X4Tr5S+yFh9nZYhAAAAj0GbLEmoQWyZTBRMb/wDwzHFcyxseeHtK+jVLY3A8ftTIVqtjziN4DBisXVomoyMfd1Sdb6tJXsseTvlPzSSgUV/wZCU3d/3F5VNMl9zrx0xBHmsDhwytPRmDnxl0gvYO++K7bpjqvYn5Af5Uv9cJNeJ0/tM0VGvOcRQAHEE7GPbnUfyWt3/xg/JTgWMC3oYAAAAIwGfS2pEf3+cgskCrRYsbkBnYXyLyFu/tRS+DCNY2/rA1wIwAAAApUGbUEnhClJlMCN/++A44nVBa6M6WcdL98HeebjAAer9Qpd/rGaSXMqDHAum3dQpiQlLhO66/sRs+qjg25SFCsZ1Vo4nYX2Mb50Cn+tgIDPoTLtW+6oNkjuj9xBK8l4mOVWx/rY+B5HpiKwKANP3WwumLK5PSCAmldTN+JXvE8vS5GfjO8R4YQe0Bloy7B12ige4QGTZi7fBYPlKWdxV35hH4+gi5QAAAE1Bn25FNEyfeJmV5SQClUo6EKDR24A91N62GkhAu8Y44c0RYjbPE4+ZbvjaBpON+8Xjc6J1wglUayytVTVHCY+766iHhKM20y4SXfgwUAAAADgBn410RH+DDmAYl2uPzEbnMKed3nsUICyCyeDbyNodbOCziTs7weQQn1LnK0BhMLkLDxK2gfVQXQAAADMBn49qRH93PjJljXXeOtUIt+VecJLVYqr6QwrBQGtz901C+SNkcNsuewdjFgpR0rkVyO8AAAEIQZuUSahBaJlMCN/HqbLy8/wEVP9rejuhgZ1/hTiRpJ6ACUBjzu0rZf+rwUPqu6H7kTYGQFdQa/1tJhAwmL82YZXnHTmWT+k7HnR3DjYHPBWnvdO68MV/h6Jcec8o8RZA/ryFL10MTl1FFKOnx9c5nTtu7WVeqQ82jjFbnYiUJNgKb+Gf4/paAXxjzK3x1VOKsji7JRVvVt7IT3RxUS3bPHqqBqcQlYpSxPlruhYyaW1ZA6jay+0HRBh3+Hyob185kbI6d9pHKI+sD3Tq0UtpSvdz+/HXXIXTeB1Rm/g1ldwTI9+9Xz9EHGdGIIViutN6mdwpp2q3QburL4OaDa4rbM66+0tY5DFAAAAASUGfskURLJ/NY7+LO8GHdi60riHOu6JlBVqwNg1VZMFN4ejipDuWEsd0/ErZ0th6Pxc9Cw1Kjmuq/zyEfYtQmhQsswGsgZ1YMKkAAABrAZ/RdER/zdG63bG7KwCn5wCGa3osZLctRFh023EneSsmn3MzWC9KqCBDfOcj+nGJ9RozF7PJP3C9bOztIPLCjeboJrWZS3xCKtt7JR8Uda3/E7pTMmmqtRMFjpXMNYSOVV2dKi7GuhQUoYAAAAAjAZ/TakR/hNQq3WE1NT/3tCaEtiN3t4YxY880nJUIH9NlCBEAAABDQZvWSahBbJlMFExv++5XDhLnY7HnguBrf5Zur2mrLT+Vz/GnEwY9bAGfUJr3XeDoQFE1iAMpidl0+4X9hFsf/8NFvwAAACsBn/VqRH9/nOBJylVotztSxpovwWJ7DmTJloPNrARFySkRNw2TMBuSJKqAAAAAsEGb+knhClJlMCN/+8GajlL3k20Wujiv8rx1TQrXb7GSm0kYYGmTH9/rx3KO2kVqyvJPzuuFfEim0UAXCWVIQw97gFcZWGwEG++myxOImz7Meuw0MWDCCNZZRC2on3KRz62Ye9OQv/dYQ8UMk1DJlNJTLsnSiKauhVJx4gsD9dWa1Kjljwa7NOoMT0lw+1ORCbOeHGoMPOXwCTsZiawz1MU0E5ekYhfXVHCB3xVjdCt1AAAAQ0GeGEU0TJ94mVbEet+VNuwEtG19Vp9qEKqXEQQaYjS6ZG6EkvByqXm1cZlLTLEjoe/LYqJRdE/SF+jsKjMJt27U1bEAAAApAZ43dER/gw5gF/dlobTz7RdVKB+qPpkr6GiWwyvC7wIA+3QCH4y5k5AAAAAgAZ45akR/dzl93GFZLVxl/7NQUtFWdUUvh5Beyz+kFAkAAACkQZo+SahBaJlMCN/Hq8yjL5u9j+dOIruITmk6F0314KXaka87ZCfL7tAHNxO+iju7Wtc+zvhxob4L7mjhRqS6xh5aFdu4srY3Z9kJMkMwAoPlbslt0ED4rQc9cCKEvzZe8vP5SuEFkEJ8M+mCHUY0aVpfO6zWHHxMB8o43nhSklJXgCtKGumFXlujSGPlk4FS34jNmpqXuBhr+dGn1yKAZ3sKSYgAAABVQZ5cRREsn81VP5vFlB4brUQOGtdboyvxjRTY+YUybkkkNYTH7v9LDS9i+dPj1ZQKYWWFo2eX4uZLvCRgIBXKzUoZJdkcWfXIyanYodRMHyDg+OeXWQAAADcBnnt0RH+e0OjClgq3ZadlqC4PYwt5pyZHpbAtX2CDDw6VQUFJdMB1qD9hMsCCcnWNss5UPmCBAAAALwGefWpEf4TUKm2xG47R3Pr01Jw7Qj0y+sYt+R7i5ZfZWQaodl6QPXmStL4WVN2BAAAAckGaYEmoQWyZTBRMb/vuVw4S52Op1bIfPeM2q9fIEDP8jiZLCKuzV8ftycOzLks1iKOF1EIt3ZtTQ9F5OaVLBfL+oXlHR0aM8Z0FAD3WY96fF4f7fCu71ZyQlXJ+Xqv31x+daLq88norMrMWmEMlQ6N1gQAAABsBnp9qRH9/nILJGk9FNJxMcY1SBNTaLEKBkEAAAABtQZqBSeEKUmUwI3/74DjhOA4HWckmvMs3ndZqaLe9W6S7iB4Q7A88qrDfeboB+rjGK0bIQqpCu+nuXe0EwJabFzLCZi73OsBrQNSTgkPJi6/soSPtur9skY2PhGs7djxF0DyGMzRzfwPqmGKw3AAAAHBBmqJJ4Q6JlMCN//vgOOEuc2omt8p6qS3Ldw5989LDKzuDHw6H9tFiHKlV1hKTORwp2nJBpiWf7T+A6Gy7CdBbhR1LLJqSDyzQfYbByt6HTZQ3hiVDL/V01gSXBPgX36fCFFEJxNZwInp0CmM1ZnvBAAAAREGaxEnhDyZTBRE8b/vgHIVKDNGLdlSHwPwIRb+tFWCuSrukzMABQCTH80n679WSke43t/ig4V9E5EyVF2d/zscVwbygAAAAJwGe42pEf3c5fdxhWS1eZ3/s03ZhFuMfWe2wKdUy6MHZBxrJBq965AAAAHVBmuhJ4Q8mUwI3//vrPjghSM+5BWn2Ge5VCc0ecBc7EOmKermZ8iWDJ4BRe+G/1S2eQMXMqXpPhZIa3jeKjGvYeHBEFF/sMWQpQ0qrYWCvPbD0bfqtAROiYI0svXAdofjDVcDRuxRo9ypWwrQsM3G3a4CbkmEAAABZQZ8GRRE8n3igS8fsin25OPmCjaoPOqdQy40hqX2JsnRU3RfPgOmKeVL8IQC9iFPZaNhAD9T5boSTEP+kwtDeCH78uwP5hUVwfti/rNaWniv7PSSGTiBYuyEAAABLAZ8ldER/im4U4K6eWIXO4QmaqlkN+TQlwR1AnPBlX1lEm5Uviyk02+lXK903CJg7gqnqUC3PCmW786/bPWqAfGLvDpW5m9IU54/ZAAAAKAGfJ2pEf4TUKkzbUQAa919yFiCEgIFdd+EDagYoE1SAUggpxbI2QuQAAABSQZspSahBaJlMCN/77lcOCFIz4vqiFuZcNnUTVzCcRWfNdM98rzWozz33Q7tLgb6TakmSicu29hGCOBKuAKHMQ1NBeTV2csdtYUMxTioTCcrT4AAAADpBm0pJ4QpSZTAjf/vgHIJkCtvXBZfB+14jkra1juFx0W5C2dCXvZjLy9zVPAyjT9j980ElW/V1rmRjAAAAgEGba0nhDomUwI3/++FR2Jc7HZDkW4uJMmtI7OMrf1NHhL48zEpmHBdgRL4PFlMYTP+bOV+4pFezhcm6vv7Z7df8Mo/izx4ojiI42oNc5zuIVj9opEKt8krPnE/wfreAmiZIhYwe8cCEwpfjXXA6bJNO4q/P6CuHDPJaoZseFBbJAAAAcEGbjknhDyZTAjf/++A44S5zaiKQuEhJ4jmNCR534wzFsjpJVuFfEuiUFOAHPYvijJXldwbC09W+F5ucQhVaOaL0yaM7JqDvSBd52cPA47AOCCjU7wu4stqnihjO5ec0ETnzf0diadQ9rPIVw+r7NiQAAAA2QZ+sRRE8331Nr6PZBVkLGHLziB4VNczKxSY25csybjCivPqOIy1NJzulsNvAJVCJH0qinwjhAAAANwGfzWpEf3c5fYTCslq8zv/aYBHXhN31AQe00BnlGw3gDiP8P3ehzc2BPCbVfqZx7BaQyOGwMuAAAABwQZvSSahBaJlMCN/76z44IGSiLnArHSWpo/s1uhN/5B5ilnLO9crpBdkUCLo55nXPs/bbw0m1ni0JxFnWRtKVr2HRpNb2z2Bzw00yKTCsaNunJLxfhVSpb57Qap9q81paINA/so2j+STu6esbY+YyZwAAAFVBn/BFESyfgOol4/M1FWMKl1NWmFHIrOMqFKzWFvxvaf+3gDItP7fuN53V0wplQxqROuxwNisKkYf1A0L9vOJVudeMNQFQxSkGQDs7GIlrSq3kzQx3AAAARQGeD3REf4puFNcwyKvZaiX5qqbL35NCW5PTI+EGVfTC7EKS8nUfK0RPsMBuJRJi+OvTAxVsTOvunx8J/nO06kwqHnuQ4AAAAEIBnhFqRH+E1vUc8XsyRGeV8iG0MB/zFxYS/vkcNLBnuFwB+E7v3VvwY+kBV6T1vy286lrWGXJod8XkgjgvJuvRIJEAAACVQZoVSahBbJlMCN/7xIOhIUAB4s/hp/HFjXsfBSfVzZJCJlAqX8IqjA3hrTr26f6UjDw5+ICxhwJWzzIjlK/pZ76OT3X6i+JXOWGDxxI7H8VsRNRLv8lmJsG43V4JiuFXvawRR78NKMZC3LHXIjfHxkbHB/JcwR+f69KV2/QtRyZ2dRq7sBgyFC/D6+JgWb/L3sBmXeAAAAA0QZ4zRRUs34Ee3aUVIPxx7mDUX17S6FVLhfjTjm/q/myIrqTAit/zI5wjoSwqYdmF2r6hcQAAACQBnlRqRH+EUzCwuV86LLP/gj3H7Nt8I619QX+jP+BvQi+oayEAAABbQZpWSahBbJlMCN/74Djg6QUKvd9e947e3J2lOAkYhiVNr9cDr3ELDEevtAnqI6c+JX9/Hl6IF6osErWJfEaZznOiFMoX0003WslCugmdnvq+vcZwkx0sZcaeLAAAAE9BmndJ4QpSZTAjf8oIQvW1xSZ7x7rM2+/pGwLqLOU5uTIQepdfmTQp8FGjoLDBWrpb5S15nV/7T1MmUpm8t0TDIMYbRQkDTpqwSDDRE/k4AAAAVUGamEnhDomUwI3/++AchU+lEC2Ii5zQFVFGdU4MJVerkDyzoY2q5+KZYNo7qgNGISSgsF+cGU0UYWSZEeb+O/N9AOGI87QduAAXd1wUH2zvdsrzG8MAAAB6QZq5SeEPJlMCN//74DjhLnY6/LPx82l6Bk5pK9HvlvJS5+uXxVzkYeM+SI2nn5l8845JM2j8M31GNo9VnTQBFtuxtZgLhZRvBX9iGRYWh0U9En0SXymWgiKz8LWajDidYO5kVrzWsaI8uEAxpZex2ijK5k/IhCK3vqYAAABXQZraSeEPJlMCN//8A8MxxXEsNPU5MyGp9izFtgK5B1fgvNcD+Z43iIuSvepRhmJjGw//jbq3VYyQGndV5yIhGAhdpfOXdaDwSiD7VkSVBa0B7buOG6IhAAAAUUGa+0nhDyZTAjf/++AcjQRnx3awW7tncZ5B5Juzr9se8S1VRa6KH4UBW6LI1HAegAOeGogjZT0PeHL9xMUqvsOwVSI7eLqasA4yAdPUEaLbgAAAAGVBmx5J4Q8mUwI3//vgOOKGHX9dX0WCIwBpYFz0D6F8pUyB6eshOk+dTkDqLXSsUh+Gr9DuLFPZf17nuf/ueCIOjxOgcQScSRdP0xc1ZVu5Jqv3yslnizRuOXMfXxkx28Z7fAXjgQAAAD5BnzxFETzfhHWkJuNUHXHYaxywUjiGRpqXXHfO/nq8EmC1znAhOoVyi3LAXh5udzgkSjzv3BiqOhlYw5ff4QAAAB0Bn11qRH9/nILJAn23SE5EKnm0eUobJ6rVo9aydwAAAFpBm19JqEFomUwI3/vgOOJ1QWlW2Hq57X357pDMQnGy6t/q1J3KOnfwatDKcCKk0FUX7afaojFcA44+ouv/VS5zgOikY2Jo3gD+Kea4bPlYDa1BAowxsskYioEAAABwQZtgSeEKUmUwI3/74DjiuZY2PamuR6r++1qE53Iv1qMB7HQ4ViqPGfOl2Zby/x/4pN8jh89MtO/osbGCpJthzQcEKpTHULvYU+KEoRclqH1cLUBG25PaeWeD8zxyPBhkjHoeK7nvqLLbLzvLi4uamgAAAE5Bm4FJ4Q6JlMCN//vgHIiQ3BG1Ae3g4noIB59LbXTzyxAiwWYAA78rBEuHCDQZgSoMfigSgb1vVN+xtQOAzXzgo3pxNP2vKC9JlDb7AWAAAAAwQZuiSeEPJlMCN//74ByNynH/TQlKMDZynXLHrMF6i39xuh+J1cO7vKUsDG7+KRiZAAAAaEGbxknhDyZTAjf/++A44rmWNU+hTQglopLhhmAYLK/ZpHdfpOlqk91IgMNgFmHV5vRn8W/Xr8cuOcpnmkUCutdRPjNbPLQZXJ/3ZczxEXjRBuNotgM7N9t+vTv3LJoVu5Gxx4RvBJ+4AAAAR0Gf5EURPJ94oEvH7IgaFXdPjVnk2O6nLuXIwqxnLW/+qohFcx7b+h5swmqFMWhnXHDdgVqWqcEU8jPLgjtj7gTWitr7+s+gAAAANwGeA3REf4puFOCtReQubuf3nGJsBw2LCxyegE3x6o+mF9S3iTu6rNVSi3sIbiyLwJ2lalMxkYEAAAAwAZ4FakR/hNQq4wqUtHk0Ek9zmyjVKtbfh+Ep3L+pF8DNTunjSXvZ3i3A3VLAaZfBAAAAeUGaCUmoQWiZTAjf++s+OK5kJjsYIqxbAN9qTX7RKmcdzUOR4WmyrfmGXlSsY4VW0N0z/pBEVT3CsU34BdrTRs/Mj365dhncAfPS39tR96BzM/Bfv/33oqAmHEP2cAcTsMGtEH+NPMAl88/Okh/VCCmQF4RGqYBUBV0AAAAtQZ4nRREs34TtSIINpafvDmraZs67s4zwOz7bhVRQUEsaKdSs9EDDJCr+0tXeAAAAJAGeSGpEf4RTMK9Ia6xDVybe6bVn8VMBUydoiydjxdfTkz37RgAAAHZBmkxJqEFsmUwI3/vgOOK5lg7JrSIIIA/sNpXxrR0mCihBXocOnszuWQ+6aCY3Qf/ooqPuZPkLUBxMfTn/KI2XKfgwuHg8Ndn/ej+Ix+JsJLtXyYV1Ee5PeWlz4wtnPHfYsD2opYlEN9EtBxxEIP+F6hy+hsFBAAAAWkGeakUVLN99U+XYnHXrZhpGAxLbiGJkujNzoot//snhNxSyWew8NVYOzMd4cv+InhXj9mVM8q9zNsNntf2Lf8crPpKa3BPYJbhAof+Pw38XwU54f+mhs8MZwAAAAC8BnotqRH93PjJguRMlq8zv/ZvacNf3ZgDNS6n2w8JdDTja2l1kfVwpI5RyQrCnRgAAAJpBmpBJqEFsmUwI3/vrPjitsEkocBG1hRaRC7O1B+CbrXs8hHZmhG2DNlmpBvFXAQ8ZzaWTctxOTdptC0ofjr2GSDqDoPTD7fLpKEEg49jN3WOp2j5Odu8aGp4KdhyWrxYcLD12s3Q1/uv7HQxh4I4Z3zALxj7IUwYKgxbBRg9IlcbVNQWTb+a3gaDRS4UaCkAw/1qd0064DhTpAAAAYkGerkUVLJ+A6aad3hR4/uXSqV1B0dTjjrlZ57rkwkz5qH8A5BiiEyjQz9F4C9kFj0976Z8PPdfw/bqaUlhpDnPTBYGWwhyC6yOlloZsDr8qrxeFAV10I20MbiNT/qY9/jTAAAAARQGezXREf4puFNc5iD93O4M+je+bLdricvgJeIzf4IF7YCFGXutVSWIsMd4abDjWFaDTXZx9lcHSO7Z6bVKQQZkvgLpJIQAAADYBns9qRH+E1vUcPAJx4GZtKUaFsCF4uxw8Y6tLrD5q3QdAfoxOrS+VUyv94KTkHGklV3NC7hkAAACdQZrUSahBbJlMCN/7wZu7TxZMM/e3lqS3Hf7gR0yQfSMdS0gp+EYh3ACNJeFZWrJYWy+Tuxmmg9yENzYhMzSp0MSd2WcLpHZqtLjZjQM+Ku+MEhhVWRbHi/Vo+GTEjar+dbIOrWe8xdbaEKLpaDLCjIApOoy+j4hP/sVb741ZZK/0nw+kO0/AXBEe8Neoan+CYCnv+Ow4wt32pXRHgAAAAEFBnvJFFSyffGVVSGsrz+3b5v1Zg5L2NFJN02uV8jUcyUVVzVImpmzOJPahozMXoce68AXxo2JwibbE1C2KfdoQcQAAACMBnxF0RH9/iXZSQkB/vO6m2O+aoEGAw2+uAjDlKDyeUraXTgAAAD0BnxNqRH+BEcHSwjYaDmtf7P/3EHd73hx4e4SGlU8teL9lc635JeBPPhznR0FXQYVX63aPmrj7NT/7R+2ZAAAATUGbFkmoQWyZTBRMb/vgHIjamvtGoxs1uSLcMbcKSXfpDcfNEklZ4VeRp0H+PR6Ez0IVFGz/p3jRWk/Pf7VVg1CUVglF6oKe6mOwVL1ZAAAAJwGfNWpEf3c+MmWAiUqATXXR1gfDZs1lrr5MWJQArYUxiqhfYhdIXwAAAE5BmzlJ4QpSZTAiPxVmeOlhGwz/Fs1c6V1b8P2RQeXHGQ7kXXgq8WzG32u2Iy6GVw/lQvrM24qRUp/8JG4USCPSdqAC5av0hupmdqUQHIAAAABLQZ9XRTRMR4puFN3GPnq2c/KJzU4nbiaIdSG7X0s9TSe4iwC+u9L3527jGFEAW/pgQ1w8VKNK8KcC/2VVebquBtHGTz76FnlStsFpAAAAKgGfeGpEf4n36qqADnsqQPIJnEcRRKczAkBlIY2jZnmdmdek7gJRW9Eg0AAABZhliIIACf++DLeYgN9XYj5zg5B5QJuOQipq5DurWEjnW3UorRPL+Gmp7W/eIApq4hA6GO7MXHIyDHsVucz/aMFYtKTDM8z25ZLQfiNZJqSiDN9723ESzhz79Cwt5M+8vOWuuGhG+Tvq9foXjrYoVXiVDwEsVfUmlk5TXbEqyYI+MZpo4mDJdCI6FifaFbOBypE2MZf5XWR/BCziCfa3ua8heIb9mYyPBNtGSBm/2+TYyBLX+ccaIy4TYZdApJmzt8tHhRNNuRbzsMmREmXlyFjU2KTmRU1Z3ZVzYLYM7mBumK0HRe3QE+vr6A4gpTOdn2j9qjAKj6NNg2MmzTCieFUwql9L6ptX7F/UuuzRuFAfTt932rxOmz2BNjBs10CPMOhJ6Pu9idR/eLvVQSuID26B7YxP1Kc9QZT+tp+YkaEPC2XBY3bdZ6qbfFKSpyAAvrJhEkF1eJhOtzQzGB9M9QVORu8jkG4zlLQbGF6ZqgPB9YhIYnHkO6pyRG4ZEwQQRv3pY66HZyqmP1pcw4M/JTWg0Dc5WzUWPGIAgkAFQc/1Ta7hNEPWBuZrLM3ySme5CdFuvklwbISWRsoUpauBnsVWWN5QeWEHF6TJILZUHeFXnSgp97RhPWIEEhIlKaUd4hpJb7cGhxHmSPkyScl5dA7+PIQv2HUjD/HkIC/Ucyw3+BoGsOX7lkWmxzmjnqmrW7wAKr3betEKeTC6MuBZC2Eie0lFVnNNRkT4VOQIj4zh9jxJYkGeDLwst9Jm6zQ9hKlPa63RCfp9NnaHWjPsnJjMX9F4gwR547dWgbLvYE7IuBo2AK2pqD7u9lEBFGBQEfEy79ZlxUwHxhYbaKRUP2ScWs/5ZBg+sELMAUKpCbKjPE+vDv9EAYHgEOzeiT71raTwKS0j8v4q+LNmKEzMRQT8YYKkcXDIj+hIcNxaNKpyvuGqlycqBZerfPCw7NyeCJ1RZg+wkCwjrzypk5ge+Jt7o/7sW9DR3cOslqBHx96pj5KJC/hClCIaZB6HFYXv5qM9NlNVSZUKsifIwMWVlTFH44dQrGr49VglslLicK98bhApMsqpHLl6mvI2nl6M00yk4zvnor0ybl+78LwbD6/u2JjQ2C5StNvCBcm5T9DSRQmWz+xgw3XM2cXGR9bRyAeishmO4rIwBRYnfjW42CduoUT5uCKdwarOs6+ZReAa6Qm032qE+yuGHc+0eaRHUxlZw4pvpia7X72bi/pxC4qWFZD2iC6up9fej8EinMBLHUXfZvheqdQetWdPndbIazJcmP6J8bggLDQ4mKE7dyUF19ugTbguFq279TrvB9ynED8sJa9KoY20QSV4oLpYcqOxvEiV7vzPO9OS9m3f72BkB1HiYIFjaJXwqYGkTXiuZJWrbol4+U0Y08C/7GdlrZPe6z/onp3lz5ODKTA0KXD5NMWK1y/vxPV51UDcOCw+Hv63HNJxyUPECSyqdmyGfW9a3IZeGpmcvlwizm/fg9CiY8lZSQs8QDGJpatu/aiCPaZPOU5MgVYO8MYY3Bg83oVFT8iz4wJf9K6OnKc1nq8GDGNbbnNkX1eqV7en2c97ZImpgcP2BFIf94YIjngEbP6/uGWVEfFyMJdl81YnorD2O2896phTQJZ011W2fThMp6sugWNSD4zAbCi+DLD8su8T5rb61RQ5zpFRjgGvMFoA+Swos5EfX7Q+pWIiU3BmZVbzG6lOYqF+5VQhycWb+BBOdBy+Gpnne+KvK78cSzPaQYJ91kufnulzaeUUO2hFOH59m9d49JV2fpLZStgYBSIJTM/9Ox/Ms+HFo9hdmgYr3o87P4pp1ZyLChdnLPD4OKjtr+1825bbLYc/XTrWBAc/9du/48l2gB3e8DXTZp0m+PbqLMrA+8QljaKoWb0phbJdacof3gIbAAAAhUGaImxG//vtcWI95nYiPte4ifOiv2MWrXa6oc5GwmVX4VGEPL3MSQvvQOWz4iK1KtWLJpeEgjzPL169ZcnBgeMLsKHdUM2wjMTrsHdOygwiz5mAKMj9iGuCQFkHUsm7iXWuKx/0L39JtWxH2DB5DrjSNacwTXnDHcq4Qo51kaWnzhKTZQYAAAArAZ5BeRH/f5yCyRJTRNLqZxFl4gyBe3yn49jBqp0OksTJCkbjJxyS0eiYhQAAAJZBmkY8IZMphG/7xIFmvHEwzcpixf6g85KZsRgAFuJ83fQKmW6HunVKuoozc5sQw4UsnqcarVRces3II69oVcoWZSdF1HPDrGm+RQUode0jm5khpoTxFEJWu1alofRiI9UOP3gOWkENJTpyKfX8mUK4iEwYLYiJBNu6J7fSvAmoGVC/qd36RNT1oid+703GbBe6n57kGF8AAABZQZ5kalPJ/3iZlhm9wMbG6wiDR23j91OZf89JoXeMWa6ZGKfp7xbU0IXx54eFVJ0ZoDzq3qCKdLExoFVdKuttpjLhgTdmW67ERQDYRx2CX2eTTzYWTJl97oEAAABGAZ6DdER/gw5gGJdrj8xG5yWNjcN5RzyeYQWTwbg/YJJ06mabrNeaaAIQL3fS05zTmWWz/Ddjyn7rd/fj4m4Wz+Fu7Q7lrQAAAC8BnoVqRH93OX2EwrJauMv/aYA5jWgK/npk++GXjyk7I/nxAuL2+vMCt3X112NqjAAAAF9BmopJqEFomUwI3/vgPUR3rMO/T2Y2VhLfhP72rEslxqZ/WAOcDYyisIRnWP8zgmcr3lUb56tCV7LscA3yQHB7Y/tXggxHENASIP9bZsUZSafMEOPrKIJo6LpThykrgAAAAE9BnqhFESyfgOol5DumVJ0S0cVIM7am+NZ19h6NTyxakBR/Ofn07JLTms/iq96suWA4PuOQi/6c2rG5IgiADAoCMybzxtewQXqSlO0vzLnhAAAASAGex3REf4puFOCunxZguX9wJlyQwsHnoh0gIq3mjZDFzCFf8XeZpx63GVk7z776u7t3kmCYNdaCOmNgmx+oj+n8mM6gYJcRfwAAAE4BnslqRH+E1CsKg5vffNV6KJEaafCYgryLkW8mp8xVM0IsV7+RHJ4vavIBIrAgydB5rvlAxBRgXHTpdpwIKS0filN50G/LRYjXoqAeXMAAAACfQZrNSahBbJlMCN/7w6grIVFJtmFWBNfhha03j96D82U7Vx+2nsoSjttpRDsy/gZVh1/0E/eE4qYqKmuDqanK7zl/QXx6E6sICctx8MFYLq7fS7pEM+ZN+K8Xv+uUh4vwC1duL2ZLxK0/Cen1zyK0H9sY5VGD83b0M1fhLQs9AaMHLOe55U5g3QoShXVmVZpN6cYhIcm9HRRzqlropIiBAAAAO0Ge60UVLN+BH0uHOo1OpQs1to+yKHFwaptQ5BFtgDlNMU+icULmLGxujc2EQSoFU5ygASOOA1Zvaf/FAAAAKAGfDGpEf4RTMK9Ia6yBEGEbBgFIX6ZVi4oZ9kz+DhKMG5HG7B0dtYEAAACsQZsRSahBbJlMCN/74DjiuZYyN7bnUg1vzKqdRlP5zSFw1h8FWJhA1x/pcA6l3LhJz1U+gBgFSepwTB3UvN0Ol6E6gnujgvTnK96bFAIXqhaB1HSHC989x9BPETy+h3mNxHpOKsWF1XH3KQjb8NMCBha0VBvTpEaYlRrVz0cWvC21LWVlsZ88JWjCsQwtjBrW5+pViDUFe0ils5WWRrnU/nnwnmdT65Qet+sgwAAAAFFBny9FFSyfeKBPNcaL7+V5X1IGhoxsOwcUUvCWuXvEiV7QwC42PiHrznZ2MxQBbEYjkTmgN7viHiaaidZgm4UkJPkk5c2BG05COu3x+CtcJ/AAAAA1AZ9OdER/f6W0XexwRhdsnetPageFCn6RV/+/fWfxA49EgEMtzDU27kTwCwBlXxqw37RvOyEAAAA1AZ9QakR/hHV8g/txlNnB3G4bTYjFW6A2jIYi1+5c/6wZQtVDrFEdc65a08gfqr425KneZpgAAABuQZtVSahBbJlMCN/77lcOLBfClrEvtK1bf3ouakjSedY0d3/QQpN2eRz/kCnq0HwQhsB+x/LqObYzJyoTEkOmGoN2mLlvk6EVJl5TXjbWz/ZmHsOadqHbIr3P2nKT//Dpvb/k1/evkG6EZ6FppcAAAABDQZ9zRRUsn3xW34B8XvH3qT97ssK9VB5rRJg15lNNNG2qjubIElb3HZhXKsOnG844Dve5W64/ukmghSwP0aO7uOMrwQAAADoBn5J0RH+AbCpxyDalGnNLTih1g6/0qOdJjORclUnHAhTKEO+MwTEmokHDhXkKpprte1wLoG6IbahgAAAANwGflGpEf4n36qr5vCXRELsL+sPgx40RoVK7tlQbDF2IVvmSd+z9ZY532dnlS8XL+vQIDhQ6PYEAAACpQZuZSahBbJlMCN/KCDlDfdqihMhNsI4qRGQFbeAgSReinIWm8Q4v1leH/5Oe95uqoaeIdhN+wGZD579BhEjNcllwMDo21tloRbtSe1uLHorlQEF2gL1+3QvtZ6Ww+YbiKc0Q3snOvvYLlV6VBkpFTEa/TljT/92LFpy+EncdIYAHscQDkrRLnVMWpdwe1oCFjeZaru/rcF7oZmGko8JpJoEP9QBsOq97wQAAAE9Bn7dFFSyfeKCcVlysnd4zeF+E4wNIhpJk3CJk50y4UPNHY2/KywiCicYiTy5aF8x2ZKGFuGUKWkFqvmv0KbGkq4auUlZxFpHVEF/DCBLAAAAALwGf1nREf4USy2J4yUop700PwUbNMjzSdo+LuZe1m+BDNgF5z/+hAiVM4yhjVIaFAAAAMgGf2GpEf4As1+KMQ0SOqKb3+q2lDv62ayTIdee1EruQL1OwR7PM6bm0uHHhVSam9bCpAAABBUGb3UmoQWyZTAjfbiYOYl0PwIa4cphLVFbEikyeFvuGUtjj9rexKdFRrEbKH+W5YBn8+aNOedn0EvYpxVizQmiRCHVywZ3ukr7oPH9177x5o1M447NLUgHH2HTsuf+vyFIgHZ9FLK0jO0NIH3ib3olKIR6JsVfteJT2lw+suG01vzwcItvz385NPA29gXXf9D6eaOQ73pXz6Bx2gc4oTPpo+WdUakNFj/2HXkeAtBbWl4gLb1YcKDA4y8G1BFnlxHv88wkrXKwI6Yol6fmR3bJI1y5dVNJQ9vudagKOV49JX0x1oxZdx/+gPH+URlfuXDOCUGqDljPZF5g8ZqvpMQOAuyLrrQAAAE5Bn/tFFSyf08HdPKf1cgoaUwH1/y3y+A+ulifGqb8DB2/uEFk2vpxeiHhlya1/JJnifRcTg6+KPhDWKH2iLwmqbqc3CDxYX3kZrHMTaHkAAAAlAZ4adER/hHi2UFyJkuq62ZFmJbEncTEmcCmczdiVr+IJ1OPJgAAAADcBnhxqRH/cLDyOioAOezAE4F60h/6HDW+EY9SSjz2LoTUrQn+4MyfhBEJdSMVRLllcvr+A6MoQAAAAz0GaH0moQWyZTBRMb/QSJIKCLArV/OlT6lqtITQzgbnr8f5Eutwk9fMwp3wElU6h8UpJ4dbhFq7lbGPtQxwukMLM6TueYoDKwOsy3Kfevszyh9gQA/RGL/2Z3AABMi/YwMbw/lzJ9i4vtQYdX3S6A0UZk8DF1yH/sWuw9zQKTzfzoQmtJ+btMgMswQyVUHemhiShtbMmfnX0YTBCZIu9ZV0hmGWGErVuj/rW2+UZU9nZmMEUv1AL47Rdpv3vU+5WsKtGVwtlbafwVaVPp7mHgQAAADcBnj5qRH+Yl5YFRBc8xrboQonVRcNgLm4flPt+8unUaOv5vuV8sqY8Po1ab+Ohyjt/kuvOUGmAAAAA2kGaI0nhClJlMCN/x6qvIqNCZDKPAcy41m6RBftLAOTU/JG9a/HDlAjWHlgLUf7lw2zHlW1dN6NmvoUeaQm4Rf3oT8dK/A6+EPtUha31C9R8AGA9Pc1PHLm78eBmqhdW0/ws5Wu1XWXR1uYbEmmCJLKmCQsDRBZlUnN5H1wYtmh+CZAiqdf6mqvSNUeSmYEs+cI/4mIy1lPLWsph1V0bChaKID0e/3kgcXyOXc2Aqc80V+XWH2ucNKEcuCZxIb+YLN8Yphusto9627pXQnMmE21y2Sy3yM8koWxxAAAANEGeQUU0TJ/TtTGaqRkjlJwPbvsmSI1lkdSxXl87ZF6YN95PfBag7xnAxXc3Gmj48UCzJYAAAAAaAZ5gdER/hHi2UvUrOZ/a5c0ermQddoJtsMEAAAApAZ5iakR/zdGPifqOf3ginCBX+vXCxrjKKxMmmQAUZhX0/pxpxbOnD0EAAAB5QZpnSahBaJlMCN/9dub8vkKP19B+pw6kf9CAjfIRxnEwttHyxpVXRku3ltHlZ/lkYpN2KXe14XnfKsJvlIt7C0jFqLMM1jsU7x0dnYDgnxa3KRXsBPz9MvvWPMhxArAEWOQPLOVcPX2KE3Zj8lL7CigZAVxYZRZxRwAAAFhBnoVFESyfgOahSeu9LaQlIdYgvy8B8MWOBXFmIgdWQSr2vo4u4eMzya1/JqL8142I8jEYL77i0hisMK1dxUEyVjcD9xg19ovelUxl0c1xZjCxVSzJLLZBAAAANwGepHREf4R4tlLARMdvFFuxCgK+j1kYjkNsSNdqSwreLdD/OSM1WLdrguFezNg07LOAkL0dD/AAAAAqAZ6makR/gQeUVGdf4170agBA1IK1mKqcCUzcAxkJipPFkjr8XVNqedToAAAAb0Gaq0moQWyZTAjf++A44SZIX6myKoZpfSwx3KsEKMJToKDMqUqApD1eL5COD6ggxaIYo2oBHwnAsVolI/njGgWAuS6GnAALjGTHXTap6KqRt2wQfFm8fHkqvFwMRzV8ZoAfwTNrZTwwkwSZyDrc4QAAADFBnslFFSyfex1KrddINV5fGhaSsbM71utgNqArEne5/psmiII46ZfwBD8NIJ75jHsYAAAAMAGe6HREf4pxDHfN4S6IhdmW0Xai/zbnkB0d8faNtHQnA0/PXBa9GH5DiHe/KxjzFAAAAB0BnupqRH+EUzCwtvKZsEyQVsXmvMVAOIOnZjpXqQAAAIJBmu1JqEFsmUwUTG/KCEL1X4ozsdj2k6Ce9zYKOCXqaE6ZjAsmKXfbBn6uZEpzN2/8UmUwlCJhbW3/dILPwAkQYsQakIryqg3c7gMWCnsjcddWB3LQzuP1+p9jOVcrD3Xr4c9IecX1eIa4w8BDPtg42swGnJPlLgbQhIpfm0zFW/KBAAAALgGfDGpEf4As1+KMO/5+YX8sjGGhQz1ZkZQqK5ID4Enmz30wLs4t+OPMVJ4+oDEAAABTQZsOSeEKUmUwI3/74DjhPbILSGmfMEGXTtKl+F4XYSvJwzWWb1ZwbfPPDCkNhDyw6pPEXZYk986ThQwSCTwvFR4sqnGv09HOAFhHcX57o94v+6AAAABrQZswSeEOiZTBTRMb//vuVw4S5zapKxN+T/89Irq/c8BeyQjHqSyDZmxoDKD5I2ZER30lxBIf1WoAoNUW98QLYvYdm6zyybpZzG4HMZW8BLCyOn7D2+fhOzraEYLfhZPKgWhn6VDfppf0am8AAABBAZ9PakR/im4U4K6fFmC52/zhJFXvzNhyan1A8C0itT4SM2A+aqKJpuDqCdva6Fc3DzqWoU+0qgoASrLy+6MMhd0AAABaQZtUSeEPJlMCN//76z44IUkFqq7aTfAxQGo0WrVRPKTltKcUg/xnRhZwIe+DxhONjdOh8IzzK35Wf3PTjClIJPbgHcpwpxkfKcBfS/pqFbk+KY4ESO7kMpEgAAAAPEGfckURPJ+A5qFYNVyzptbWOl/Lio4wgDNAqOZS327gG2S1pd88/6C26jWrIjYrV1arRkukKguv1vHOvwAAAC8Bn5F0RH+DDJXX+1g1VdtIjdzQ0EXREbOPgvpkgaGuH7Q/rjrGKps7wnLqUFIdCQAAACgBn5NqRH9/nILJAn23SZGImicBQ3DUHDjZ5s7/AV1JPRKG7gT6le4kAAAAYEGblUmoQWiZTAjf++FR2JyyC1iitS+dYx0EO9MkwWF5beGuOG1ZTyHyr3lFsRHyiSoIXq4+8kJeLsXYAYR46+4k/GgvTAKNRfO38FkHkrRwtEd/5K6QMpzqI//YNmn7wQAAAH5Bm7ZJ4QpSZTAjf/vgOOEudjsWuSxRrOyCZjbojYuMdD82NYeOPzHOOkboB3OZcVutWe2p4e2RjezDJTnX7ykx++hEsasudgkCPoByJ77cStz8pSx2Mc8cb6wHv4P1NBXSxRODEmyzR6vAkTP/h4qQ+1JpzMmK543nMxeRvcEAAAB+QZvaSeEOiZTAjf/74DjhLHUdnTIr/kiBro1rbcLYxqvkewiG/h/qLASFUG7vjx8I/NQYp5JX260H1dnWv/6sKBFvjZO2Obn0YWxpuvjdc2STEdWWh1bdxgv1NBpRMD3oT8UI6ZG2HYtJKeyaW6r/H9f6pel3kaQA0Z9+aPHAAAAALEGf+EURPJ97HUrtL/mgidYi0TC+wiKzM13IdxH0hFqOnPLWN0qPOjbY0oUxAAAAHwGeF3REf3b7vlgIlKgE10AyLX4AJlNtNgJsyqIGh8wAAABHAZ4ZakR/im4U3dqjRaZYi6ufJW41oOZxfaude2YO0qLYL6p+tv2n+ZEhy2TKBugHpw+CmXFI0YFOx6UPFevBE6JTjViEV+0AAABnQZodSahBaJlMCN/9duV8/kCFIz5BMs12wfEGnFGw8X17TL2yufvVuIdnWBmq147QWoxvvU4sQmiODTxcU5IxrW48XUQ0l/UgndpZkMnXV4yvU5Wl2xRn7JOqG5tG1Wjk4im+e43neQAAADtBnjtFESzffVPfJNmH6c8uGGvN6Ua1w3dSVxyq52qBEQ5AQ+3UFiUIgQtkx2AipnIUrZN40brHaz8w/gAAADcBnlxqRH+BB5SnfnFFCFzzDarOeQ+n3x9mjYHrxdG+1G2y/RedFEpdZPCX1ehb7HH7JWFYy4bQAAAAukGaQUmoQWyZTAjf+8aEFzqPqIfb98OjTtEtGYksUtQHqef/YQxv2TlelbawxGTG5fTWV8g+MnDomIil9fbxikYH9XQmdZy8SyM4eT5UbSB3yEfjPZuf69KjR9VOuSEgKQch9dgwqBRFURNNn/ZHanJPjS32PTuCkyyVSbGosl1nzCvs5IhumOPeVjXmkR/0r0gyXtZomPanE/2OhFNpNTOQ88Qr5lABTVZ/tpyQT0FJkYsHvnwCpfD9gQAAAFJBnn9FFSyfeKBKqU691GQnRt7wT2bAwvsFlFpJYfJL3kI0Ejx2m9rfTT1n/rKrW+o3DIZl+x/z/z6dvQ5v53+xVKnddxG+xfI7Lp5qTATEEOogAAAALQGennREf4R4tlJElLvUEtQcMggGQj6hD8aXPWAOycZ1yxgmNuAfy4hKPpmOIQAAADEBnoBqRH9/t8AO1iKF8NFd5NzOsKBLGjI7tyNx83qt1R1YPIAs/dScM1VYQvLo/32AAAAAQ0GagkmoQWyZTAjf++Acgi4Oq6OXF4I6IPPdwzJFin4336rAUggpBKkw1FQCveDhbBVqTv1fDTBUjHvoqU72VIneOhEAAABfQZqjSeEKUmUwI3/76z44IUggGGDyEJHJ0zdy4sFJbfwH7K8r2+9osdhXBVubU94QlL3+0o16lzaMEUX+s5hDiLE+HU3QAsuJzTdUVHWDE2crLRMBDiXB3Nry4va26IEAAABUQZrESeEOiZTAjf/KCEHJoIEIqOD/LBqTEmdKsw0VZiIHhVnR0YHo2X8W67cycOsus55XJ0Jqd6wt0yhFaWJl16uDYWLglY5OB8KN8lyu74G+Q2atAAAAZEGa5UnhDyZTAjf/++AcgkyM+G1Ouhcn6Z2Bc0JS+SytB+cEXwbM6nE1arVgMsYki0/teQ7FCmWw4Us2D3g2l9Xae2Na6Cm9q7U3uGVww7Ls5vAowUGvf1+1PeObba2e2poP0GcAAABfQZsGSeEPJlMCN//77lcOD4iDd1SEdD3NEMCzmr6ToyL4AwseqHC/p+7NrCkSHk4I3vaFZNOZvIcubY0dF+R+0CSzGj75kO/jtMhszh/+0NqcZ6+lfDVkqMjcjmxjwAwAAABHQZsnSeEPJlMCN//8AtD5RPCcxJhhtvitPqIz74OQtZ8I0WhBpvfz/vdzPZ91H7CxyeX9rMCN0nFbY7I3EWhbNcHJ6VALu4AAAABEQZtISeEPJlMCN//74ByFVfBVBzCrec0q51pjCVag54dw81MGWphs9/vebHu4Slqm4xB0i/esr2DZbj8Jf6ynXo8mnM0AAABbQZtpSeEPJlMCN//76z9RIyID26En9fyziADdwqGitUiUfERx0K1tUeLRCixwBHGA8E5zp96UdN+aXTYSjyBXM2d52ocrCKE+B3/tBnS29ZiVSe53MesjWFUmngAAAFpBm4pJ4Q8mUwI3//vgOOEudjsIu1r9UcGS/BbdCjokhEK4p7ca9fqqtyh9/H3PZV5npSGQTV+CX0Y/0tMJZ7Ct1LC6HPc//deUSi7xe1m4BaSc4T6i42PU9PAAAAA7QZurSeEPJlMCN//74ByEPHhakORElru9VCJYbH2Yv4/MDI9PbvboGW+izcFJ6I5f62A3mPgFpsbLHqEAAAA0QZvMSeEPJlMCN//74ByFVsZ6j9D/F5iILSjxSQW3rwCXR3PeR+KppsBmj4ZYK9QEYuzckQAAAFxBm+1J4Q8mUwI3//vhUdlhbNCiiMoi7GsJj/JUSgwY7nhDk6vwDjSJ+VmDgx/g0bFuY3yAw7VdofesWO5yJ5HEM++OyZWsYg5c7kHgGfigSgYsdX+uEbqvt5ohgQAAAFBBmg5J4Q8mUwI3//wC0LvLElkFCZ+0R4VL+qEuV2to+z9d8c68NZwBRUtJQhv7rbzrKpuzMyK/+cyHm1wEmdUHgCIDOwZNtkrWOlccD0IWKAAAAF5BmjJJ4Q8mUwI3//wIE8MhYsf9rP5Kq/l8NreVCsrDsaQlU0Sws4aOq1OMlKviW1Ryzdx7P7AeQ57sIFdLR691f8uoHHExNezH1SYct2J9kEnZgyh1ie3x8zmVrr3AAAAATEGeUEURPJ+A5qFYNR5FvP+sUOFlh6OdT+5AC/ALV9jcn7A759t2aCrwOgT123/pve/kWbUHHscn0wZSsua2ijypozcsfdeqQdpwo4EAAAAmAZ5vdER/gwyV1/sg9S+lWzZCVlXJ1bZAB2pl+X6+0J1jK93r+nYAAAAnAZ5xakR/f5zgScsRbP/KCTlPetUWjSxs9oMve1ulK8unxRGQuTh3AAAAcUGadUmoQWiZTAjf++A44rmWDsmtE6da6Z88b/lOIgyX4xpZm533n9QjU4LoQqKhjYGX/tuiDl25viW0fxsXEC1UQCYY3Rxdqst2ZmTB87g9gEeOvp/+r28Wjk/lHMrim4Fl5xsal8b5UTEFgA4fhZ1xAAAALEGek0URLN99VLWJA1xrwZZgYEEgHgZAufc/vsiPVWvEmgCIWFM+Kch/8zg4AAAANgGetGpEf4As1+KMO/6oPIpubBAR5YI/mVvTOM5vFZZyOcvnNGLH+tubJVUzxsAV/kUZ8urzgQAAAIlBmrlJqEFsmUwI3/vrPjiko5FfFdReOlnYOY+aIn0vCx/RI4Xxu3GUBJIy2hCFvR4nFjOuCPWyJ9oC+4fg21JT1abT5qVCmt7KlU2DoSWZVc/093qGhFeYaBYh1VP3nhSxd3bKcoYQ++seKOZ67SE5jHbJWSiZVndbBPLb1CknRwfGQPMcrNB8gQAAAFdBntdFFSyfgOahSerkFDO4h+HT03gu9drqkmwjsbFo15C1ddns1scB1c41Z607J8TKKQBbJ7onEPnjRqY16Vmf6zqYenNHNjGFMGswbhSFbJ9SUbjgWrYAAAA1AZ72dER/hKNnC/BciTqtZU2dJJA+iXwBG47V52G2lmuPgOczguYbeITd64lNMXnwX8HviTEAAABSAZ74akR/gQeU4JrA+ZeKVXUJDRTmUyM8cptsQXz8j1l2X0x6PNzWmC3LWN9LRSLx9v3DwnzUFYMnVL55McNf8uFZblO77WrPcvXMd+GCzMVi+AAAALpBmv1JqEFsmUwI3/vgPUVnq5reM5Ssz2GdGe////9UbdFeHz6++R6IoL51iZvWWBTEo+w1/QGffKUbsN48+o5YE+HeD/8Y5yb8QquB9cRfL+hkT6gX4h7jUqOGLkv3aekdBT0u6Oz7hhVkeKCmJJIp2Op94DIntn/GHvxlz+XmeSvT4E+MDxEl2LA/adI6b5NUJrAj+zyjLOTgs/yPrd+rgMhLzwia10NmoNTcVyZRErfLqU+jEWkhGUEAAAA9QZ8bRRUsn3sdSq3XJB2xN6AjUdfvrPxLfuDa0V+3Ef2dPSVqKYe65mOrLxPdKn/0htbtlAPjz7kRvwqK+QAAAEIBnzp0RH+KbhUp35fRKSMEG8fVZtySTaTW+ISyd3w0mPrfX5tqV1xZY/iIFRY4QnnDGyfwxM3+zO/RMboY0LQswmAAAAAyAZ88akR/hFMwr0hrrENdSs5i6TQXQ6kkP9h/VZrhl6OVIn+wZmWbKSNLJEYN6QzCxG4AAAB9QZsgSahBbJlMCN/74DjiuZY2PbTAR+Dbj531bEXlChD5lrSbsGOo82bYl0bd2pDpN3WLdD1xqLg+1WuvVraz9dukzhD2dWGkCfyEavvctphVPI8P9WmurH2tugHweQ6rl55gisLfyBenMiMJ2RMJR+lzc1/3Kri59cV/bv0AAABIQZ9eRRUs331NpBIZUVg/sF8yeBOR0xrwAbP3QN4HMGL/ULx/iGepxoAV0VvJUzUyZfsAcCY/TGky6vUdHs44yyXro/pZ+OyaAAAAMAGff2pEf3c5fYTCslq8zv/aW5r90SK19ZwBG+JkrBByKG0soG9BdH2HugcqFw9T8QAAAJJBm2RJqEFsmUwI3/vrP1EmHKvUWSjtSRZlSgrv/liQAkTTh2MMOMOBFt7k9bYXOcwf59p4EuKaEKx7Y7E4B/IgnuToImCd86g1ojTwETdli09SA1Bo1HNUGbOLU83EsefzyVWJSbRzGhntADwtCJI9oXid1IkLBW7ZtDjX8pcEztbON/vJFzL5VTnKXjD+8vULEgAAAGxBn4JFFSyfeKBLyHXG9IGRaGq39J+w00CRNYhQsfHGXw+pdnBYY3Gs6lwhAdgPpg9sKS2UPZuAS9WFZHm4tDlDX2fuE3rMgJtYMXJgf90BeLeZPecojFY3k7mIB86csvefkqVua5t3fNOkykEAAABMAZ+hdER/im4U4K6eWIXO4QmBIq/I5sOTQqiQ0+cMCYfZVcNBwzmnMTca8JpuFSs22SO0wP2GDUnITSfVdndW3vFSPjn1UhcPn2zggQAAADkBn6NqRH+E1CpDUNBG/lajRglb3oTeJln85BeqE22nEJfeJpIKOlpztR4pxR7gHZelJU6RMf+Cw0EAAAB0QZumSahBbJlMFExv++5XDhLnY69ArjxbLeYiTGAFVjwNqRrhTUG41oczZ5EyBH+7dDwYd6fxBYacxJZsYzZglHmsDhw0rqZxKveLxdRtcYBSZO+liDMbXCCTHa+6LbkPiqpeWPuCFg018OjxE3T9dFSvl00AAAArAZ/FakR/f5yCyQKtFixuQGdhfIvIW7/HBKE6HEqyYkrz7gzpEkFnbRTwQAAAAJdBm8pJ4QpSZTAjf/vEgWa+oqExUYVp1olVvNbP//+8M/kBsXizdkbJTTB/BTvKR6wFbAMmXZWZPYbZpxPe4+VXweLkpoPRtXn7vfj18z8A5fxL1tFG0YfcZNJvwmwKmP55Yy9/It5uUnKSxwWtgpa0Xl7FOVHv6yBORxpBbFCVCVvuDxli2pBcO5IupgTKpGUhLJbY1ESIAAAAWkGf6EU0TJ94mVbOfACgqZwUyCA6PmiiXKLu1snrfMONhCCh+U404cQOLSIYQFmfISdXlAdVcHufy5BepvQ5vcf8R6jFC3uBu3JcyYPQL0WMxnXzTVqebdZ5fQAAADkBngd0RH+DDmAYmD5zzrhUCgMcWoHHz6RsHLXEpnx1pjFJJfd4M1l3M6R12v+1lUKb2QeenUrWvLsAAAAtAZ4JakR/dz4yZYCJktXoPfrPf/ePA889a0VEm8L/I5Oa072z9TwIkrlHV1UnAAABQUGaDkmoQWiZTAjfx6suWx+PTqHcCWZGG6ifNKCv4tUI8MKauO6cTSadburO7W9zwXrTpt9gXX3b5tdNvqHh7J5Cvv3AhRisceNgb7YMFsR1N0Rhh6/nt+TZ5PAFkbKAp6iJMn//stQ+Yz7VXAyMKB41Sngwiu8Skg+yDGuw5z0H1q2PuUD1rmITqnNCuI9uO8eg8uswHpLKHevgR52SAYvS2/Zid6FbE0GnXdftWRsiroevQ44GTyV0fPUfROdAmGxCWF/QXCmGcCNFD14g5Bu8qQ6Z+50EprnL1v3bi6SzBTqLcyvNtZPOKr8PQlWxgIVRruPOgaVr1gzPI3yA4IoHPl+VKgdU7iUksHmCQG388/DVCqccTrz+TsolCl0IqE+rbPC6A7gaixgqv0m6YoS+bdEtfIf5PoP5qHH3emI7wQAAAGJBnixFESyfzWPAWZ3eFR2Uc7F1kcb+oikzxVIfIGvLglj++hQHAONn2Im0GJpEyeStvP2EmyHct8tdtegovJLgQJsfAYqJrPchHLtBboynly/y1VWmuyJjSjqD9pkggLaG+QAAAHQBnkt0RH/N0btmFIvxlCtkLVaZxkluWs7V94CnZHEt7HGqrMRIPPdfXob/JPjyP65+5jAvNmhHfPXIgTy/h8MzRAfvoqjTEvzPxiXoq2YJ56mYweyAUgN5cy5ot0uF6KC7kwMci3mHrUqnOmZx7LD+L0r0zQAAACkBnk1qRH+E1vUcO/PuHwIOngiM/6M9xXvQGSxikpxLdsa+ZBaQwprqGgAAAHZBmlJJqEFsmUwI3/QsHiz8vl29gYzDA4k5/ybedXu7tKu7Jmy9KPWg3HFWvQIrtmp1hFrco7VDr7pRktQQWV4z7wB7QtPmuaPl7AbJCErL+KhMJjtnkudgiRpiC2pKeIZGLJlUjeD33utWpNyRcvZCNrnmrvCAAAAALEGecEUVLJ98ZVVId0czXNe9uHzxp8h3Iok3r2Td9dl1B/uClTrEreOFzvFRAAAAGAGej3REf3+JdlJGgwv+mvTSxIdWSJzF/gAAACoBnpFqRH+BEcHjkFzBJtn43EJVcXoHHh7rOnULz+Jbw70vSl2jJMWY6YAAAADcQZqWSahBbJlMCN/HqYgKv/p9ybhCWa37OXOoId/HNKBU7lW0ZvUv8UFLcBimbwNzzimRNxyV6iYYfu9yXFwXxNgztFgbdhKv7Yaw6GOdd3/mgssr3YJxxbqYri+BQ0gsKvOU/2bHepT/BrNNhFXz42GZ1RU3Z0OvhNwjsYp+atyzw6K1RQrGQzJdsjyTkZfv1NrgMxK2Kyq4o9lOjDMLVkWn0P0yn7qosC4D7s4H2EcqLGoPHxwbUKQbXH10NPUXB5VgGo7Gr2hDQFFekB222S3MELKaTdqBDIJgwwAAADFBnrRFFSyflmG9bXlI8HDsblWlpyLDNyCFSjeTVQjsKx0UBaiRpefXBhGOTqqno+v0AAAAJgGe03REf3b+tk7jCqVAx466Oq+thoM+UPfpYn/yReLuvw263iAhAAAAQgGe1WpEf9wsc7fqSm94QxwqahE+mbpNVa4sm6pOpD7UrOlV5wtZ7l4RNFEHIy7UPZEdoW6mGZQa7151EYp0yV581QAAAEVBmtpJqEFsmUwI3/vrPjhLnNsGv1BKYkfPo4cRV2egVZSmfCgBIfMKkEiLVx6WQ/+NmwmC1TmNjrOX/bDa9yyO1iT8c5kAAABEQZ74RRUsn4DmoVg1HkdZAFs1PtxQgT4vR9+7PMyPdh8OvTRv0jh6+gWn1+BLNvVSZYdeehlg2ag49VZovj1blaVazskAAAAiAZ8XdER/gwyV1/tYO0MG9HP/fzsOk/m3SBmPS7dS9yDlmAAAABkBnxlqRH9/nILJAq0W6Uz8mfnp8RC6e96BAAAAkUGbG0moQWyZTAjf/XblCEl4icEcUzKWIv5Sdz8xTnNoz+oTMnKI1DUFY37mUJgzqPkO5fZTe0TbtfncuGXCPon6ZueMlRvrNNcOK0XcUnwUA6DgCWf7RjzKphjvKjqti5kt8Itb3RMRziXZc2powyXYrDFolCRsJlmjzet99mqjeEo9j3KbahWDJHugtfI+sDEAAABqQZs8SeEKUmUwI3/74DjiuZYOyToYj+uJfx33+jD4oAI40PfC774m4iIF8KayeVa0oDX8hLv0VogcQ+NVxVP5bwWpReIJOBcCQwcgCOGRpzCDSfOHlF2gui/AgGYMXyRtA7HKV3fvBHPk2AAAAGhBm15J4Q6JlMFNExv/++A44TsoRqu9LyjXtmglYXGeW5gtQLLkwsKg8q3kkLlUIQztAr/PynJc//JwqcgO8CVKpBFGV9xrTYudyf09uTycpUDGmoGFkEx7CA8vXK/aaRGjjNBEByt6qAAAACkBn31qRH93OX3cYVktXGX/s02n7XjMhSSUiDHyWYYBLWZFRZWZz3x48QAAAFJBm2JJ4Q8mUwI3//vgPUQ2HX+ATzBMT/Pnu9JYrDfMcgGJDPQUmDYQjMezPulMGzFlZkx2U3ZoJt844xJBEP55P2VwxxHi0Kagm2yw1aPaGpOAAAAAOkGfgEURPJ+ASbiMqvfMgbI0jPmxHlOwJfTnXYgmLO8WX2zfxIhIhMrFN5/kkUGPI8Ahlw6/NHrekA0AAABDAZ+/dER/im4U4K6eWIXO4QmA1FfI5rFYZ4x7XjV9ezMO/coIVv++VYWmiu0D0vx9YATno5hyzWlfZMIOxG+n4NgyUAAAACkBn6FqRH+E1Cri6kk8ZHaVcOy1cqQKjNJ21dDGpEKstb4T2+G8eEqVGwAAAFVBm6NJqEFomUwI3/vuVw4S52M4EcFcQqXCxPpHL56kgKVqUEL4iX1OpvXCrsV1IIbzGfkz93JglG5FvhdntnM54UGgAg1vWYaNNzIhAfhLt1NhDflBAAAANUGbxEnhClJlMCN/++AcgmTgWTEr8if5xlRqXymDzn6MNQ/cQ/xqLIQJZe8l+jwFGz7ZyzAhAAAAXkGb5UnhDomUwI3/+8GQRIyCAPHVGA6jY1LoBJIpmQFHa7VcmtTnD0dyhdFuminjvN7hC4NUT3epXOwmxGLF6aq4Ii2C54UdlPXegSNE8ZyYU/M8lkhKwaBQ3c4nA8kAAAB9QZoISeEPJlMCN//74Djg46XORVi0pon+LFQxBrPNpEf7Xo0uPcjk0LUC1j144x6PKcVHoBOoaxOjtlY0/9T0DQfS7OAr7AkgvvVG0TpjF3TYvKWzgWOhJhlDHw+N//6f8iZoaO4xHGToMil4OnpnInBqExbvXmyS2Vvee0sAAAAuQZ4mRRE8331Nr6PZDrA7iBiuQ71R3LRXqWPzXk0TAwWCC9fePi3Wd1bTnw+uPAAAAC0BnkdqRH93OX2EwrJauelrxcclUb4bVAQe5fp/eXH3gDiQzy+DEM9azx07mjMAAAB7QZpMSahBaJlMCN/76z44IUkFm51/N882W7qXe+E4ycMma6jvo926yxLryVknk2weTfZ2Kk6jWgpuPKF0slNo0l2zy51T6byfECXCpBr/SZLSr2Zvw5ZFKO5Xae+BbRDytVWPc5sse+N94UDFDnnW+6PAj/4nobZzzGDAAAAAUUGeakURLJ97IKXjDBk4ZD9O3fhiXA/r/7ECLfoxJW97egIUifkFUbHvwS6RWe47OV/aXVeX/KUQSlcRcn9TJuFfK0IGrgObTvFadJzhLN17/AAAAEkBnol0RH+KbhTSvs1uWWoAkvTwy/vyakxvNOg/BBWPV+NaHNLh+5xiy5Fr5LFPB+8S8aN/EjTnuye1l0X9H1OcBtFfbZhvpHufAAAAQgGei2pEf4TUKvaaa83Ad2VtKlPt7siVdyG4UlrNxlJdXv9DrWc6/IaeuwrDtCf++mzEfYyS38DY6xAB6l6lKKK8wQAAAHdBmo9JqEFsmUwI3/vrPjggZKIvbXzqFTYjQ0VKKQs4oElaRxg0TMLtnBuolgaQrMekht8e3yiYuzVvUWS4uxk8FlkEzYMkLXoLetWyWvtKs75yWsyyQ2ot/tDY1AEK9YY5RVC4GnKBf2CZWDbyfWUm6mixHfVzgQAAAEBBnq1FFSzfhHWkMDYBGZBTsgt1r8IRHNzCaJx/FGcgoS2gz8e56BFv96s1sEBQAIJQ/tBYOx4of4RDU8GM0X5gAAAAIgGezmpEf4RTMK9Ia7m5tNYSF2IHBXygcGjvFaWjSpcJnmAAAABXQZrQSahBbJlMCN/74DjghSQPww2RGtCahPmX81LNxJQE36SVa0/HkFpDiH/yKhx8Jf2oCjWFoB3VUBDgDk4+gZhe2X6s7RXS3cJEENdXgtyyuS/MquQRAAAASUGa8UnhClJlMCN/++AcgmH6NMQeU+cJ8XbSKWyOtUbGfsUeNz5bQi7J2bKGDZFcCIS/9wonu6aEEEC/Vv0zBzYbwCfuAFMkFd4AAABTQZsSSeEOiZTAjf/74ByCYiKau+rz9PkmS/yTzhpkUb60QzRKs8rt26pMRvSkNKiKHJNJ1FXPEAY8I2SSl8M7BoAPB1/QE5fHYPJfTYyyOFiXPTAAAABoQZszSeEPJlMCN//74DjghSM+tgpIYflbxOyTaowlDPlB6SR2Rz0K+CLVW7IJQBHpZDwt/twQggloO94udHLowQ2+xoNwlqG66Pq2mXsLqfv6rBArcwHxAtHTOAqNxvb+8RDt+s4i3eEAAABSQZtUSeEPJlMCN//77lcOEuDI5sHb9JiVElFoic4mCH248Kt3rU8OlDO8sAxjxeauU3i7fx+UKBIcj59moJiQP2LUkx/p/x+YG9GV+7aTbNWtwAAAAEZBm3VJ4Q8mUwI3//vgHIUZmOEMdb1abutjrmvCM7g6DL5L3B9CMFG4FEvH5aC7uJB89RKDxxsntlPQ94dBlp2GynTalQrRAAAAaEGbmEnhDyZTAjf/++Acgjp7C+NxYI7ubVhpgmYIDTXkpEYV3zaW3e+rWLbW+n5QGSD3eCIEPH/DXrfalOt+n6Y9n9BbjQWVz2TQZgWcsZU85NeledliQG4TudhZAEF61udo0AyWq5F5AAAANkGftkURPN+EdaQm4CV0M/rYl4NgK8z84UIqrAO/czkG8l0D2j8OwsMsAK/qF8Rm4aPdZ1aSCAAAACIBn9dqRH+ALScLlpa/+GYDTEx6TsOAEaCmCVz4kvAr3WkxAAAAfEGb2UmoQWiZTAjf+8bbiClhCfpGNHNpGB3o4p+l/b+6Z4oYQfJr38IF6u7qllXxlyI04Ql52TXK1ASkOluayplXmJPCRsnnsdY6EoVDYN1QrkAehRAHXvIw8/6QTQFSPQMoerdoDhvcYrQcyeBLl6zvPm/NUvjVAocrsKwAAAB5QZv6SeEKUmUwI3/74DjiuZY2LeS976KkuiWNhYqOhk4zW/uyrwXJteQNDFt8DdlcWf/pTc81xcC8O9/HfjSyQlZjUCjsB7Da07CzCZjkJnPR1uZh2dxHaUVulf9ID7ewjMxqhMmoJ2NS8tg5XpdDvP5Tbn3fhqRHNwAAAEpBmhtJ4Q6JlMCN//vgRKlMjfz3lEv3p9Vb9LQ7v+pEQxwC/9zazAZ5tgWuhGJQVkP+XR/kRzKBKBvW9U37G1HqvI0/phIKl6ot4QAAADtBmjxJ4Q8mUwI3//vgHI3KkJVcGVSj2/3lOtKOxCchm5X08HMnVw7u8pSwMbvYNglvmYNN2GyFVjbMcAAAAG1BmkBJ4Q8mUwI3//vgOOK5ljVPnge6cjNMAzkc8L1XcMXvFe0EH49/37e6kaHHAytVBME3/FfgqqJuUggarQZYNRM/+85bFVwuRkE/Fcn/dlzM/xeteEV2bGJ+FQYk/XavueTQrg1tVZiDMduAAAAAQkGefkURPJ+A6iXj9eVzRU5pOyeyOkOK6siJQvqqq3aZI916sjslp70mJB60/u4PD1WP91IjlLAujUAFKxopKSwbqQAAADYBnp10RH+KbhTgrp5Yhc7hCV8z4dV9+Z54gcLr1jfWaZfLnK9DZ+XvNFoFhKEO39XZNyhFmPgAAAAvAZ6fakR/hNQq3WE1NaMz7BQMa3r2kDm/kxI5W0s1zZYe/BmWXhyBy4UWxPwkoakAAACjQZqDSahBaJlMCN/7w/OhfFZHbxAlIZNdK81M9DmgvjcG9HFcEgcYE3SBH+UoGuOALJjtU1dC0GdzLYklNnQHymyVpLOjxAFGP6fgPn39YtnABx6Iuq2ORMv3owIyCbWEiR2ya42PDNeEHorS3YxGwxydzC/f/trN+g10WiIaC/CCYlyr0UAogxl0oEvestbuqfBIwybgP3u9WkPg01fJr/nOwAAAAC5BnqFFESzfgcsGNr4KTYjcPlel1/Bo6Hv1V6W/hG/R9ldaJrwij6IiqixTyhYrAAAAJgGewmpEf4RTMK9Ia6xPrVbhCu9V78UOxQddtZ1H+yVxt1YlCTitAAAAeUGaxkmoQWyZTAjf++A45y/4Q3AZy8yKUZ0ax+WW8rFph6+XgDZSKK64i8BQCeecvNGQcmtcwgPyfTLBG70A7JuS0YXsmM/BbvInVHP96P4jH4mwkYmXHax7vqt2xdUMX1/uhLmr9k5kGIQN4Okuar3BDoEYDE27DDkAAABQQZ7kRRUs331Nr6W+8fh7tCKREw4UiRr64auFmklKazVIROH/WyOZMwEYrHbLo8JiX46PwpVX+28ttm84fKTzy4PcjJuhLtZLcNIyPsq30DEAAAAwAZ8FakR/hFMwssa67x1qPHT0B/ovzqciH/v7oaLqbN40xsh1mUuAuQt+HqX4eIGgAAAAmEGbCkmoQWyZTAjf++s/UeJu9eBQwVda/BhoIvFxND3yJyVrLYLCHqElIVpP4XAimELY478ot0SaWrEjaeTwo6tQAgnQW9o85VeSuV++MiiVdty8QsF1uI0JRpePl3og3F+kKHovqBhl43XtBd/3SOTYNkTJ+qnMfrxWequhajieauKGpmiI8nGWaOCva4RwfQCy5sjxyeWAAAAAZUGfKEUVLJ+A6iXj98CjRWanezX04BDJPUQppE9ePwMqecBmkgRLAxx0WEB61woNndqUocE/nExx730z4eqHj4y9zx0y2MLBK82XzU4taZ4D8bTJZKYzQismuE2WWwmBvIKWSg7hAAAAQQGfR3REf4MMlUQLMpWe9szsSyBoDD6MEYNMsQqchdl7Zl8i/O4VYyLyiKSFtAPuL9KgbCvEV7oC6OzFcVnFHup8AAAANwGfSWpEf4TUKoZwMnlUAsRqwSiK4QsOjRuBEUWYg1Fh/OCd7nuMtjeiQ5Y0p0H3UH8K47B3zCIAAAC5QZtOSahBbJlMCN/8I+Su9ssFpVth4bdS9WpvaBw5IeSHpdzxkR2ATv4jfUsYMAFeDTkERmjPr9LCIOQ65j74gzHKGy/tLT/fSolZ/CgvRfVDG/FeQDx2Sn0fsAWhLm9uQCA9XMosomEHoAokfqWTj0xGM8pbd0O/hsjd+5tP6naL8dluhRdb+uAPr/SfD6Q7T7RzN8WnGQo8foNo+PncggQZoQZ+z4TTUAfSjZP5rs1umERs2Zq5dIEAAAA9QZ9sRRUsn3xlVUfvinbaGcgB7351ebN7ARewJep+yTczsF5e4c9MzR3BPzL1klLWjrwlgKX46g/ueGxNcQAAACoBn4t0RH9/iXZSRorN2wtOvHL3rMKNXKhCQl/5dOgxZ26rVQNQnwhs5oEAAABCAZ+NakR/gRHB0lFsfre7+Bp9uH2y3Z7UlKi36/HXi/VxAc5Xd7dL3iB/+RZjRlA4K42VY9VgxSya8NR183XysSDJAAAAaUGbkEmoQWyZTBRMb/vgRKwKj6jQ6kUl2X4UvnvSwPjz38P+ukOsefshONivg0qJBk4d954f/XJfgcHAk1n/pDYI/bf9crnrP4+4WYWy892CnxvwMZGMki+3NlF1QyLUNlnRuvpjtnNZqAAAACQBn69qRH93PjJlgIlKgY8ddHWUf1/1OqaMkBqSnNFs7onwD8EAAAB2QZu0SeEKUmUwI3/74D1FV3aEE8G1Z5xdLRpbXXTFKbRJQjbqRswcA7HOC1tolGRV5dSZGwQ06qx3Sx5PBVNDREabAfzyxszynQeyjRmijN5plDb8xwdIxI2kDBBZy4ZZsO17xE3oB1eSevH0ZRSYUMODXwfK3AAAAF5Bn9JFNEyfgOol4/fAqtD33iwUqaOGcDWLURHTyerUUMGLF9/dTKp1sM6fV6ymzQ6Mw4D1HfNW8wMEkjn4rsCZqQMIPu+olAZgJJZsnp7/524VaYHVA6xNJqbPPxWhAAAAQgGf8XREf4MOYBhh7xX1+mWbR3kBqO4P31bzBBgLVw3Z4XE+EiCqAVQWrnOF2DmSkXwx2FyUnBheGRDCuJQDXml0oQAAAC4Bn/NqRH+E1CsRJUo8Q5Zy14nOo2/Zu7iAW4sc/CZgv2ou0lLXrm+dUmUgFE56AAAAfEGb+EmoQWiZTAjf++5XDnddY2O+hkpJCpGqqN+ioRd5fh6l+P/4qd3eXmdhQXWVfefzZYfFULX7dQb8xKswxzxZH07JzrnEToy3x7fjpn5UQgftlM6VOlHdwbe4xz/rGQAsG3VCvdsv9qrf3RI0iuneCBDzjBxBQQU+fHEAAAA3QZ4WRREsn3xlVUh1l0zC3bh9rNn6CiPj9vvy2i9BnSdiA+sxg5yXAeOsEXDrU7U1jqGjoUWpWQAAAC0BnjV0RH+AbCuMI/rr1UXH8O02Fm98S4fObJMRsbWn8m9/CguTUAZ/90CM84AAAAA6AZ43akR/gRHB0sI2Gg5rXgcbmg8L8jqVYtAxxuK4okCllLmEQxPSy4J9M1DJFgcgzgIfqUc+3vAZIQAAAHlBmjxJqEFsmUwI3/vgPUUr3PZU49GT22lznRKAoTW3jOcu6hD4ioK89GfSr06PHj3gvvsDv6O0stGp0b2VSw0E8X+l5NjE09e9NsWOSzj1xWS/c2SZ/lV+QRWe9Jn6Ul6stNTb4pd3Udctlx7cQDgYj8PLbS3Ij15gAAAAPEGeWkUVLJ97HStiPWYa+Sau356/Rr91SVf6JCh4jlNlHdN0d2PuQNUBv5wh3HT7qTnVQ9xZjQCXFlPLrwAAADEBnnl0RH92+75YCJkhya66QMSa6cPWkN5ju+y8d0nfVs+H50BoCjgVtakB5xqMlZAxAAAASgGee2pEf4puFN29kxLsO36+UxP+fgCnqU/gYafq9/Xg91ZfFnrFqEyyhYJCxcIhXzmEbbtohOCddejoKdPZpuX3ZH5v3KZCY6QQAAAAkkGaYEmoQWyZTAjf++s+OK5ljY27Vg48pz5a5+KYWTkP7SvlMxka68KtKhbaJTrQL/WXyvxd5weYVY1sWu6B/pmqrj/2QsNv2zSTXylPXylMXIEg2zO3tnyu/tbvMWheX6G7CPIGp+b6nvZ76NvL2x3bbS72nPa/776VwQDIVEqplwdPzD4QLQXjlie/F4jvAz+4AAAAYEGenkUVLJ+ATVljVg17/TcUb4E+TZf2O+oJvmehu9/UXCOUucMo+X7ourEbvh9gcx2HsgqBvE+aCc3LmoUYFCKGFauVj9/DXDPNwzQtla4m8OGx+lrBEXLX2daTLsmB0QAAAE8Bnr10RH+DDJXX+1g7O/ZshIy/nY4r4hQrCWJCq/V6Esd2r/5asnJk5RR4uD0AP8YUqepXEN7bk03ujsxa0LKkPI7KOCxKQmLykHBjvGJgAAAALQGev2pEf3+cgskaLZ/6LMcK/qtVrn4NRgFIs05qa7vLpR4QkwyqPFyrTb7wcQAAAMFBmqRJqEFsmUwI3/vBnCJ2LJ1QWuVMjOxOqnfxxjGCcowTZ19I1NAFB4qvy5q6Ryyp59tnEv6iWYDfMwcl+posZ/vpl/L1PLmpbuboHvedMTA2YWPl2SGbI42h5SLocWEUqLCIWgdUiQWfje697D/3h1HLpDMRw7LLhOe089llXSaMRlYUcNK3n6wBz15SFB00Q6ihaj5OREXBbwYXwrPNfyFffYIF/ieB9Rb+wytAu0366VH851MnHmh3OAwWWubcAAAAXUGewkUVLJ94mVbEetQUbYugF4NXqOohbWdmnQgq5IikjMhBE9Q8RRsBeKAYT1t6fLH/uQHXlB1P0aEmEa9WNTZMw4CcPsqSexTdEwT9nlxz5VCl/l2owvjxAPVh+wAAADsBnuF0RH+DDmAYmD8VJbh4Nh0FynC++OPAhiHknCimadJ9Fe44UHpkWLeV30b4H+zBmAv2TIXxIi3YSQAAADABnuNqRH93OX3cYVktXGXvMuRi8Jr6Agx/PtZNe1ki1o1HHbdUiBBUvLJF5NG7lEEAAAB5QZroSahBbJlMCN/KCEGWN8UZ2OksJynlUtJe/T/Ze5KZIBmbExdmqiGR5669n82a+AtLkfKGHN3Rp9oU+HakLXqfrXevKtLZuZD5XvF8lGiaSAkwjZjoDDFe5iz5YPRcz+fZww/DpgTxvDkkfcVTbleWnl2qZXgSTQAAAEVBnwZFFSyfgOmmneDDuxdaVxDngEJwiO9Wub5u3Rp6pIgiyrZeIaoVDz56LILYAvwi5cN7K66uYsQoR7hKrNuY4y7NEoAAAABIAZ8ldER/gwyVRAsMRG9wSk8dVUzVjXcbzTmj7t7tRjQ5erWnxVd83Ef81CF3ZzwpPUsI8IPWyeBOnVhve3/z/SC41nZMKfRgAAAAPQGfJ2pEf4TUKt1hNTWjMmchYe5SAUxvqnkg1Qw/OPUc3DuBY6sbZu6KUUUigXOtNJ+lIbfnS6xPA+D1gYEAAACqQZssSahBbJlMCN/7w94eHXz2QmwfLeHxPiiF3IVqCWcYhPGBobbysZlGTAT8YM63rCVG0s/TvYgSnMzoYtwQNurg7iC/yUMmiP2Uprp7HMupZaNwHNYfBppfsDg25dFnATbutk3GeDH7PkeFPJAtlp/QQQHUqqrP5NTHvXTLswfs5DENGEt86b2DoFGS/7lW7J0rnfMCbX9Sb7lVrsBZBpHc8xP6+kRPmeAAAAA8QZ9KRRUsn3xlVUh3RIT7WdkVbVtgz/uvdxDmkzwqBzSZwQ6hkamoiDiVfFO1EuO/+qDNbLAf9gNMQ6eIAAAALwGfaXREf3+JdlJHCJoZnuoVnuSxVAFzEoggTy1eBea3xrPnNe7ErbeJ/f2BanpBAAAARAGfa2pEf4ERwdJRbH63u/xqE+Tl1J8P1I+BCos4BmXSZ8BBe4HTICM52Vgu9SxEvs5LsCSMoz5HkLWd8bnJC+JCYqHBAAAAakGbcEmoQWyZTAjf++A44rmVvKoKa41M35Ch3G13lWY0bzQloWTKokZJkHadhknWNHe6v4MeLVqdkMXc1nQjTtXrx1ehZRUoti8Ga75LuX20ESKWe4YPhUknmNLVPxY4QiKLz912RQJJScEAAABBQZ+ORRUsn3sdSu0vw0w5exICsmzddDcONsXZutRG82avccQdJ8Bo2mLUJmrRj1OGq0Hhe/tpl9GLDnt8r9hgyXAAAAAuAZ+tdER/dvu+WAiZLV5nf+zTESAGqSIRhNll372WSzssDIEYXudYpMGm62vGTgAAAEEBn69qRH+KbhTdyyn2mc7fr2krca3/fg3IzUQIeXmjyL0W7kM5SHpgQjyEjdPMwM32smfci5GMHzDsab6+RotVwQAAAPdBm7RJqEFsmUwI38EzmEnMThj01oz/8HmIHDZoyWvPgj0C4DoRsDjr2VEdnSM5F+6Es/HhGkAFmMMlxRO4cBtg+v8v225tcKWm8Y+JoMo6OE9FQeGP80ZRZKcXaFITb4modkoTFB1MOoASSKLuo3eyB/nUlTi7UcHk/YyDhoKrCsj9X4UhyECyG+WaaL4nUsfz6qAEXwJcfBhGzK5jzIx4VGoXgDJKsaf2nsV0cr72ztqjUEW7uJ8MhZ4H2jvG8l5wh+DH4SjA+Y4thclgZCtEvFf6uTRjHi/V31/zDeILkYiFDRYuc/uiN22xbZOJlfR4cWsUKVWAAAAATkGf0kUVLJ/NY7p5XjclhLFPOm5V70UHjDk0y1W8WLMxUEJtyFr8LEtGYvu5dSu368Egd9yNd6Rj70InGXN93csHtMEAObKXQji9PQYG4AAAAFgBn/F0RH/N0btmHtB4in5wHCfslGyNeY6FAHgkfpAI6sJ7Jpy/UlPtc/zQjmrUm7P2PZ2IpRrB2bO7p3QSB9jQTHvmkqdIl64ZFcKRLHTYUzBmd4ons68pAAAAIgGf82pEf3+cgskaLaDSUkmmWOh29hQ0iDCVqc+7lecaq3AAAADgQZv4SahBbJlMCN/J932y9VLbElOTIQnppSdx4sQM5Zt/kb6lOpe1/qFSIxuuTjdoDeO039Y5hjcJKGrutCBLHNOGFrDitKwHUBx+wCBz7FNfIB752PfbJllVBPkMGIlohflE1kmAWoEvdKQ5RVEIJBo0NM+7Xc4QOrGwlm1kYRsr+YzMe/a8WW7js/AQvf6/YKmK3nv/WmouN6JrifggGWUpt463j+VsZPXJT1MLmlu0ZIJheCkppBM5RK9Dn0pWoHhZqpp1LS6KmH+OIyOhplBJZEO/pM5ypaWA+VX5TGEAAABDQZ4WRRUsn81KWB6R7ATs7BLkl+8C/fhmkvrMuiNDekBNsZe2yrQ9LlYwQIG0r6PzAA8CdpaaTzI+B+gavhWMMYl1MQAAACwBnjV0RH+exIgT1LCdOIZSo931hcnbNVCgfvajMBfvkpSj6y79gzRE5BB2iAAAABwBnjdqRH93OX2EwrJDk110s66jvwel9iFQh+hHAAAAjUGaPEmoQWyZTAjfyghC9V+aCM9y0tNTu1YqY7isyMrEL1ISg0j2C7zj1UUqJZwpqqKi/AnzZgAJPEx+BbF2G96bWd4yUYwWlUiGqh7yxHAdI8eIWqxQwALIgmV6FsYG0q8g7JGYi7Gagx16ZFfTcfzINdt3vHY5lmi4F/KgeU4H+REU/9BL/jlA5mhrTgAAAF1BnlpFFSyfgOol4/XldkOz+vOrEWi1Q9ySLFxK10HVuM/uqIkihgxaTiEJ55tGiib34VbPx+F1i9KTeR0sV4TCay5VXoDlJbQJ2X/dXyUkz7tlo17BgibiD+MKM8EAAAA2AZ55dER/gwyVRAsFrTC0Y1O8nU6PfpwHrj+ZB/q+LVk4xtIzN+fxT4G1VCHKhzISBKl5tZm3AAAAKQGee2pEf4TUKw5pbspc1EHDtEesCCmJf5VnNe6KBZd4GQ+7A+bHfJOAAAAAgkGaYEmoQWyZTAjf+8Pt5W8cTDD7lNWFGUFASMmfmrkidqCky5SC9+cPiAP2S/BFQOm2nlMb+dDFgmvdAP61Ih4Y1f7BfhGuZMbC8UQVSneMohC5yiCzKNIygFEz4V6khVAoCjtwHr78chXAgv8fXLr95q5eD5Lu5Nzs/iIGtMQoXUAAAAA3QZ6eRRUsn3xlxLyHWJcFQyLMUspqSzsA73febdhh0qb7NHDVptBqRvp/0KYGMiwx7j7KbYTwgQAAABkBnr10RH9/iXZS+PSr/l8r00sSHVgOJeMGAAAALAGev2pEf4ERwdJQ0Crb1298v1xh8pKagceHwIXpnt+MpKlEEEQXCRaEef9JAAAAckGaoUmoQWyZTAjf++AchU6MHD1O3v5r6pfeZSR95Rrhxom7uPZHpFqjal5BKKKMTbN6mgvg8mD64Xft0kUDE6Cue5SSiFKoyWlVEwiX5hDItaXm0RmztEkwGguvWdr4BDwmXp/9l2g6tLg2xXeYbPfQmgAAAE5BmsJJ4QpSZTAjf/wp7xIUfno1XJ+iymWGbtpLZaw7PcYW6RhtN1hZk9b9PCNe5qkJxn7A1KYQZsy24bK3hBGA5yy6riTS/6Bggyp9yOsAAABlQZrkSeEOiZTBTRMb//vuVw4S51+pMKjT792mrZCg0a9enAYDaTD0+eYSezMml8HsNrr4W4Ypl8E3Ek6kEobjk/rXZ7COVAytpQu5ZP0ZvW5Pg7D5NsaUt6+L3onc6uiUnEZYYcMAAAA5AZ8DakR/im4U3dpdyjXpuEJWyVePMxrhvDme3b6G18sYPqqk/CL+3u1bRHyeuiRn1bp/7qnfDmkRAAAAXkGbCEnhDyZTAjf//CJ/zIEiRn1jC9hpPPLN5QuFZX1Los1jhijhPS8st288jfDlg8K2nAHAViYU4SieUCj1ME+It5LvfykfCPuvpb78LHOE+CJdbH/LeiWos4RGOk0AAAA8QZ8mRRE8n4DmoVg1Hkb/e4SXAQqfAAXFApCfAhG8pf0RtikgWpTEO9+lQ2C26jfa14IxlqzJJ16I5H6AAAAAKgGfRXREf4MMldf7IOPpaY2kRu7VY9Nr6GAEHU6Xd08+AnO38OezzPbYNgAAABoBn0dqRH9/nILJAU17xAvr1/4H4dNFOjsFwQAAAGFBm0lJqEFomUwI38qLkf22+XhQO6dWZKV+0UxaeuBPWkT9AVqu95T7qYZaV8ruFgtLIWIwIFx1O4IUmJVkQy1NeM9AOftya18H1uabejHiFxrqGQv/o7wQJwgNV5M+HntAAAAAe0GbaknhClJlMCN/++A44rmWNd9junV/0ZL8GPrsR5ejMuAeKb0C2X/4TBgkmpc5cbaAmhsRGzzPOE4Fd9ectAGJfFuKvehbLb2HVVq257BpM84fIVQhLM31NAPpq1x6jVFAQ8vyPmNq8qDOjXDoMYQIaRde1lJOu2UcnQAAAG1Bm45J4Q6JlMCN//vrPjhLnY6q+y84SEeloUKkmwZvTuvQy1vBx1cbwrmnV8fPOJEA/L74o4k+P0p/PqnHXGHiFX5kRdkSTKdwmiFccfDJFoOJFk9OTzWDgiBL+BN1Qp1Rsy9KQmgSswUoMNqBAAAAKUGfrEURPJ97HStiPWYzn6tT9FPiyYP9OyFj9PCcFZvsfLF4FBcBfJ5XAAAAHgGfy3REf3b7vlgIlKgE1B7NQWvwATKbaZAMXPLJFQAAAFQBn81qRH+E1CsMICfUaEKDBpIhdx6OnKAdCl2WDjL9pRb1pMM4S55WoAWOuHs6SIZMvTEbf6TD+l89/qm6bCPBnWGYDYrn8PXt76kldgiGMd/HEyAAAABxQZvSSahBaJlMCN/76z44IUjPWzut94XBkNuv+dieP5cTiu+U4ud4qkSdtDohXcxAqxjKAO9Sq7VU/iPClB4mMSKfAXByslZgH7DJTaF0HqcQoHsXQn4FOdmjHs65BRr5LHgnCj0HknjvGWtSlUvBUBwAAABBQZ/wRREsn3wG9Ho3WAS3Ni/Hx+OpCe/8+ukUSErgBDGT0W0PMnbTfONFyGyOrLRqgi9mASDomyObYuNuslVCT28AAAA0AZ4PdER/hRLKZLTcnylaqqMzlu66hI+mhxsbMLljcUK7FzrBO7JMAaWFrqACUm2hnz5DwAAAAC4BnhFqRH+BB5cMaSRFVafg6PSCA26I6z1H3nEcDifB9PFmvQH6EfZwAD8lDAP9AAAAfEGaFkmoQWyZTAjf++s+OBCWQr1tn75c/emk9LAHs1fTdTbpvJzb3LtU233UzJwGYEL7e8Ap8RSOlZKi1Ur6TnLiFRd10qK3EQMsUKDdzkCorRRoghT23l2ZqI/4juhRRTX/HniY+Qj/rMVC30KuoLxdRYifc+0btYjQlp8AAAA6QZ40RRUsn3iZR4RsZYniIDRxKFBpXUAJNxgE5p6+BezvDoyEfR/+vl3w9ZsZ9oIbtyEjlPrxFfc9iAAAADEBnlN0RH+DDmAasfgFq4gCUD1KaJGjFNBEJVYapagXlLZLSIzJIY2+GzlTGV7riDN1AAAAJgGeVWpEf3c5fYTCslq4y/9ptS3UG7hViDeebQUTeM3L6N0COirdAAAAYEGaV0moQWyZTAjf++A44EJZY2AK6h8ndtcC/MxKotOHWqm8LCc5nxMDTrcRe8XfRAD+sULXlYsOMbyenQJfp9B0zySWMkOE1fdeLQFXHwYudsCa1bFocAlUNzU2ll5X4AAAAFVBmnhJ4QpSZTAjf/vuV1RDrkiUKryvNy50VVcIKHyfc0rVEbcQaQOHTvicnp4trDKdvq4wTZVbX2wUqr/yMRShkhjqnvVzqia/DOdlZb2Rb0wer4BhAAAAW0GamUnhDomUwI3/++AcgS6wSqZJjrGWxPv45j0u5gayBIOWxzopmY84hnH7gzzrzeXgUYnG7EGJiE34rnk/Z0xY2sxq0mDVmu0uyDqZKWyTXKp6UClgGZ0ewkwAAAVfZYiEACf/vgy3mIDfV2I+c4OQeUCbjkIqauQ7q1hI51t1KK0Ty/hpqe1v3iAKauIQOhjuzFxyMgx7FbnM/2jBWLSkwzPM9uWS0H4jWSakogzfe9txEs4c+/QsLeTPvLzlrrhoRvk76vX6F462KFV4lQ8BLFX1JpZOU12xKsmCPjGaaOJgyXQiOhYn2hWzgcqRNjGX+V1keWAHLdIU5Kt6ThBKFN9iXo020ZIGb/b5NjIEtf5xxojLXamTjllPOcFI0jqWkWbZoIKHhmS8TjKSHCuTIx9dl0FTzgEkfjAOKV+MBkOb5p4/AbH1ylTOdgDP7a7Qz9MRQalIH6M4QHnoVD53mLhP9Mm+R7+Nv1EH6Th++wBzxhKbERneXaOozvtmrCru0n5aeRSyfurbH+PPpMF9pcN0wV6BXo60O4fisdnnNQeaTOQSwEQoA9l9G4isMAZz0MMRGPJfTWi9P7ynze7cFpAtxFXNele+QfVNTyKOBwDD7HFYOx0ouGd90huFZ2/V0BetUuzhUO0amnSSpDeMcYfW1ZAlHhEdZedtfna24MPXCAzeTxN0822JnETBRQEqo5eqoRZEG9h5BpWdwsZ15MpoPzXeKTqs8tVQC1/SK0d/0/xgwE6pNhIFwxgVY7DWsiX9Xv6tLaOudQGEMJp81OTR0L3+PIQF+o5vb/GCWDAkSsOBZs62OY1NlmA//1Wvizc5pquC5tnGGabusyA/nmG7PBFvvzvk1g6c9FZTfMWn30C0nuVgD/+SL3hOUxNRRIZsd15UgRRZCQNTNuPsNF70mEMar3F5ZGy89n5mtWtCMnx6Kxmzor+SNb9cdhFvoJ0pJL6GWtPAX3YGiIajfBHoC5Yk+AsjS+erTnaEBKxjzK9GpM93N1bO7ckHjy4hYhPKAFrXZD2uS3VtXe6oTBsitQMDgoUo2V1oa/61b7CoZXWhVK4UEX6j5mQQQ37LogaqDE0eR/hnOHk0qp1/ZLliBrAn0q0l0jYC3N9Q8Ae2vkxGonmaOmrZW/AToLHGMb0emiADcOSePLr7YHGJDiVzPIBiXNYDOLMcaDwCYGfbEsBrD2uy/d1Fv59oWvDv054FejnrCwPJNhHJardC3ukYi14OF61tBTEX4on2ntlVKj6TiEl2AY60ALC2Pewdc9kxhuM53QE/DPE04ecu5tBcDncARadW/1to7/E13lltojS2w5x9whdbhPouWVm71+s5063T2CXVUoDOAGHtwV5PBLlUV5s+9rM1C1nm6Ly8oNNr5tCZZ+xgH/cb0QtGLo2/+urcmwovBsXN81PlivdxeZpmL/iI63XLhkruWXb50jY1FpX53VswqMFVEWIGAsoxcV1hcpt5SjjgtAJB9brynD19/u2V/j6WCCVV1YJc0BsMrlVTR84gzjOn9fl/YKUcx8xd5zXtGYTTjr+P/TYuSs7g+gIPIJqW3t2VM/TsVgWKT+LKDZouSIl2nMmmD4XKzO9dTJuIBhm0f3xWeCHxstiNy8dyjXq6MgdCHeN+Zcag5SxJrW/q8LCIwYDIjdfeTEv9XV8VvX1oPi87I9rfBzE/XZtG1OijeZ7AcvsTu5Q/+VzP0izbIAexsR6SN7ywadj+XkAq9hjRaN/3QPlLeJxa7wBlxl3d1FQ8pFq96dcjnDir5ouXNKkIk/zCOM+q+YORr/PMY//AhAiSwED9KgwJz8yR+k3Go6EUpMgLgQtZ73cqHm51vCETs12mqn98YrWLHREslKlJiLjFIFQpWKlBdfkllUxz2P51TqpY4sD96jVYhYQ02ZSpJu6Tqh26CeFwNmHP1DrxMk5MJAyK6KFlKGZZi96tTwAAAE5BmiFsRv/77XFgf/gTkah1R5Us0Wnk8dDFNJtGuWmlTOExBPxdbP2u422B2Lto9nlgbofDURRsHjnidXKX9OQPQ2PYSuh0BmyDHyTamf8AAAAzQZpCPCGTKYRv++AcgTjLzdvWTXYs/nqVO/kWocgnSkHh3DwMwZdqf8C8gKNCxrEXojVQAAAAkUGaY0nhDyZTAjf/++uEdiIsFrkFNTDNvjCEHv2x8z+IOmFBPu6GwLYUtDIjwQ9PmO+bcrpj4clo9KQ/KSk3neQOZHILev1Y0jHKDWKlbBJSNC7Jskuppn7Pw9tMIqVOLcYm9Im1C+17BGV5dlf9xePKMg8Jr3Ff+/cKGd6zElUtjcoq/1384Z27KW1jNv2lQ+cAAABdQZqESeEPJlMCN//74DjhLnY7IVYsSZ54teV4M9pVBpl+jrm2vgkpxJMVoCRgAjdr9o/eR/Ca4bI5DFioJpFA7jz3z3p6Z8EP6+oWPcI57n/7rtuyon5e2HiX3e/BAAAAPkGapUnhDyZTAjf/++A44TvpQ6q1oyiGyCVbHK/f15aT3JMWQoDuJk0FKSLab/1polBzKBKCBe4A/5QO3LPgAAAAM0GaxknhDyZTAjf/++AchU+Vtkc2kb2jJl/oQwyJciND8Tq4ecQix/2xNi8PSXJvoKQxIwAAAFVBmudJ4Q8mUwI3//vrPjhLnY6qAnz8Ja7nh5BGRfd02joZMmnvwGPMnY1UtK+pb8bqcPrtmFQYzPi2qj09uQT0JTvmvYDd63U5Y2jF0v2YgFFNEVesAAAAVkGbCEnhDyZTAjf/++5XDiuJrqTwQnNXojWSFTharkNbc7gdyfMlGNhg3mVGFS0alrqG5sZCfIgbqaclT6Gfe0OipRwkPfdn4m7Thq/CqsNd+JyxTOnxAAAAcEGbLEnhDyZTAjf/++5XDhLnNsgTYfDsmrw/ZxClJj+Si9cQ6o5fvdg0roRb7xCts3+qftaz09ytTf/Fu2QxxQHWzmVyf92Io8RDVoRP2XUDjiYmvZjXggNq8HwiPdQMdvmL1lhb5fPZGj6hWNxUXuEAAAA+QZ9KRRE8n4DmoVg1HkW8/6wpqyGi45GD24TJ/za6DAQ4eU9CDfguO97+RZtQcexyfOmagRCfqhL9X81Swk8AAAAkAZ9pdER/iffqq8P35xMd1IYhD9d9sBpVn7iTWzwZs4q/5ESbAAAAIQGfa2pEf3+cgskSU0kAS0c+imyjJea5l8GNLVRLfBL7wAAAAI1Bm29JqEFomUwI3/vEhS51HjidTPHSoGoQXFVXRhNlrpyj41Xz7oGDI06DVKDH46CPr5apO2EoNwcdCUxtW3hgqmJ5/H1C2ANjl0mdlryKc9SwOs08mAQdgVi9wAa3am2ilj656liiY2Pr5+C/f/vDURY0hdLqqHzbufbgWXm0W7U5ed3s5e0CGqT0N4sAAAAsQZ+NRREs331UtYjmOPqgDvwpA+dkDRfmu4CZBFY77Ij1VrxJoJU6bh56MYEAAAAzAZ+uakR/gCQt04OPppwlRhivsWAxBpyC7wzzomqfCVfnmfzfbSoDKVZgsjBm+Kt2f+3YAAAAjUGbs0moQWyZTAjfyghBljfMMuyWTw27Mmxw1DJv9ycLWX+QHLgBj6IZtikvzi8lkYN1RnpMqOZIVf/nduC4IEoAVSEkctakaQiAU2DoUeITXzdRvYi6cLHFhXwRHpVFF9boSkQf0W8GSZnPKqex9y18cqZ0/3qjnXDpyObZscNjiTUYfl8+sVcWkDfb4AAAAFlBn9FFFSyfgOahSerkFDSgMmtywjBbMGKUIhmNiwJ7hm3fI/sR/In45K/o/zgHqjSeTGmpbhIijT92cAprRPD+g2EjPnGa2bkpYhmAbzhjtjjTRAhkkbC2YQAAADgBn/B0RH+Eo2cLrfP7Es5hEIvxy3++81JIcbjtXnYbae6sJ93XFc4LmG3iE5toWPTIgZ3Rf4u8wQAAAFEBn/JqRH+J9+qqfNhEqND9NFo1QljVM1SBAde2cQSHX/EgRDnk8wp3LTUqgB6d+8T/ceO3zoPFZUv5ty3ZB/Jay3Kd32tWe5euY78MFmYrF8EAAACzQZv3SahBbJlMCN/KCEL9LNaI8OpfBjrh+prQ5Hx5F8zmd++rK6rSSpIp/Bmp0IxjryV0UhIqqY7ythB00KcAsr8fcfrBfpnJuvV+Xgc3l0ufktx7e7w3rU7jqfUydXix+/gToV3m7PqBlYaP7wIonbR+CPd7LoCHcFtKXmQsnVlI4yUAum5TYseBCN+1u3fAF9yw2dGEPlv/60B1TfKQhsyYxA36fR0Gme6stPBqWq3qMsAAAABCQZ4VRRUsn31uHUcUqyrx5k84MB/5PPciwmpzACm2Qtdwi7HXVZEde7qqTsctIRWclAzfvXhjFp+5nog7nC999IrgAAAARAGeNHREf4ByGO+by1cZ8ELMy2ttDhtzqmeTvmfBobP77YCuVEO65NLXFqkgoV2VMcV/Dme8id5gHSZgEfW0Usi+767BAAAAMgGeNmpEf4SH6CyS5Kx3XhSoUhAXrSLodSSH+wraAIhCFkqWqwzMs1mFxDnkzNVjPfe3AAAAdkGaOkmoQWyZTAjf+8SGbo4p/8Iamyi86brFs7iSG9hkRmxUtyMqDX87BjiDhTPFbYOFwbUDOtNIYN5sBPm9LHJf40gfxrHO5b+yYb+iAzQ+D0jqzD0wUE/GPk/is1vbvFMqnhMOuYRpHxFQPFZ7rMcJvsV/ceQAAABJQZ5YRRUs331NpBIZUVg/vHpzTqi1adLx96xSVQdbyZNxZ1wEI8f41eXgBXXen0pb3PYiv7Ia00pmzCGAnR7OdtwQ61Xzb3asQQAAADIBnnlqRH93OX3cYVkh8eOukCbZ+6JFa+s4Ajf1rEP6NxJzT4vwzyoFKEMzdAXFIlOjEAAAAJhBmn5JqEFsmUwI3/vrPjitsxqnuuZB/Q3LK9EB8XZHkxOb9shEL9bJjdZ+AyCA1XlQlALxJUJQknuX92nj+nF0tYg9+8IRFjYGwbn3xaqO5g3aC8XnhoUJD7JtUL/BWe7piCPNYHDhlaejMHPjLpBeweLNivYXRppq6jwaS+E6aqBmMXuESQatm6wvSHErLtCeVTRchSPIcQAAAGlBnpxFFSyfeKBLx+yIGhVJaNqiSbf3n/EhW+vNSV7ifu7le6rz/EoChwQp7PkLC2To1qCPmWejORWShLyp///Y0qa2M3jqj7c+mZFtb+jzH2ON2C5/c8hSTBB7jEfQYULfNJnrimd91mEAAABPAZ67dER/im4U4K1F5C5u5/eTKivk5Nnz0mjnt5+KqDGfDi3wS+Mc05ibi9UxpuE1o1aqd7CTGoYA5lcK6Hr2V9Y1yJwB7+JdsC2pp0re2AAAADcBnr1qRH+E1CpSs1MPhRiIa2mtfgUCd6KO7sJt69FF+hJ5oUlgGPC0EPB1ogSuuBERwapo5fuTAAAAhUGaoEmoQWyZTBRMb/vuVw4emP9PN/4ewbFLUGiXiZoXEL5lOAJuqN63bq6Ya73SLuB6bBVPnpl/y39gFl9DALFqt3PLXqzNMiQ8uGXt/+TM86YgjzWBw4ZWnozBz4y6QXsHg/Q00M24YeHSO++uLeHaT1KOziWaUuzPIKggxZQgUTyNyYAAAAA3AZ7fakR/f5yCyRJTQtCEVMtRioHnim37d2/B1xh3Ql9UxqfT9HDmdIOoZcKX+BDfTVGuDUOJ2QAAAINBmsRJ4QpSZTAjf/vrPjitsxnTnKYcmEudJHuBU34rmoZfuXG6Cx5BfKY65GVp4NsLJ81tu0fHLYXCqgIs3Opp1W8uDzLiioBJ/PdQ+EnY54EcqZA6tP3J34zBz0T0ggJpXUzfiV7xPL0uRn40C/UnuyAQslAT5p+tnCRVSaGrqcLiOQAAAE1BnuJFNEyfeJmV5SO/7e3+IXPu9DkhOBWthtp0R6laPa9QVTtdYLuaAVCQfvIfjc6J1xsnVDkt6SMF8jhI717hP9wd4YjWpEepJyiGOAAAADkBnwF0RH+DDmAYl2uPzEbmrBU7+iRoxY56o643yA9iq0lyU4HIBaatgNSaUrc7+QCg2zU0nrhV7MAAAAA3AZ8DakR/dz4yZYCJktXmd/7L+xSWqxVX0hhWCgTdFXq4mp3c/1a9hxSgDlfk8LNj5Vw2oTK18QAAAU9BmwhJqEFomUwI38erj+EOVPHy32Fbp2Yz5YGdj+kHyBt317o+i0bT5HgD8WT9xKPFw0LmVzgeNi1FS2ii0XX5T3ObkUYS35DbQo/q8fsZOm3CpwwU/Yr0/kpevNm1zY3V4E6rdm9gjhcJayG8aqoO4g4WVxhx/xMNf2JeC2NfVjzZe2vcxSFQ0+nBoYdXCz2HGhMEqn1Vj7giKBRPbFIQfXzgNot7kOQRpn9o7GlvIh2PHH8O5b9WkObUKMB0E31SJnRODa7/x1gMeVNp9x5QG1UQTgxsISHm0WiGHaHM9FndRzybD9Ti4Axz30/xzLbVD2b0F3de51FeapbwWFhm9ug8jx3oU8tVKEI+qRwwZklPZNSM5eB1TunBJCGBjo47s/DWaVQJ4etgtBC0e4W/geamug1fhr/BERy0ommj6LhH/UwJu76416mGlYt3gAAAAGRBnyZFESyfzUpvZmoxPUsjVw6GJkSBGLuJUn+Aaj9PjxAHugzMhsCOH1qsNMVFzZkiMTpbdmVc6H9d/LXbXoSjfpgIFnVOI0u4Z1YosRun2lvX2y3QNp1C04NrWmN2xpZYm7uBAAAAeAGfRXREf83Ru2Ye0HiKfnAIZilzHLjlqsw+Mz1QC7FumLzaTF+pKfa6AzSFSUgrNZ85pfDhXgDDUnT4YfcFFRuWdSKxg2T3+RWpUKLLxFjqhPOFsq9XQnstbGEYxiunguGhZ3eFb8BcZaO/8shwvJPEZqscqq/swAAAACcBn0dqRH+KbhUp33TVx2AKpAjXBFaDxlFOYkOj2jzXxGLdVqUiJCEAAAB+QZtMSahBbJlMCN/7w1h3aeLJhoTP+9zKomPaDpIjqpuU+pG82MhTY8pgjC1+/SH9odkJKpRkMPchDcbwgE0g/0Pt03UBu2cxVOCZ7XVviNWlJE6vWTCbySBPKWxVm7F5Zgym+lfh5YEmZKBo/8QqBfrF6LcDLo0R9Hd/6IWhAAAALkGfakUVLJ98ZVVIdYm9ubdI+6CZ/d+Ekcs8ct6CLr4YDU3lQFcmRYk+y35Vg6EAAAAYAZ+JdER/f4l2UkB1Zty9NLEh1ZInMX+BAAAALgGfi2pEf4ERwdI/XqLu8/7QUPdBBLv6noXG1Jbi0GH8S3h3poGiBbJjvmVilJAAAADUQZuQSahBbJlMCN/HqXjaYdAgn7QtItVYNomimA2qDY9ZCdac+T/CWtL/WPI+ea3ZhxpIBZzHmoE7ED6cmscuOpC4asmWP/ZY4D3ZofkgXRT0JIHQ5C+jN/NawCAZBlYJYeWSrnedRhv1Aq+Ezb8GwGkCKqT77Ol/D+JE8MiWDBfvfPqV5xJ6yu3+ZMsfZri2LtSEJZnT2C1JziQTBkpSnKjEt1PGYe1Ql8Dbmxq+Ajqsnbj/1Q9G0VHb1uODbJw883wCc+9y+57YbL48lccwqZcjBHAAAABhQZ+uRRUsn9O5+lrykd1W9x10yLwnbEdC7nPGYZCdulN4uKAuREGBPh4WH7u461q3qEJSnberdlJt1MwVHaJFwOjSN0QQTPgMWg4BaWK592DS+2DuO5Vpp9sLbJu4zonEwQAAAB4Bn810RH92+74LkTJauMv/KEBnzMgyQCjXgyIf5rEAAAA3AZ/PakR/19HKwpYZLtl81TUXlOY+KofzOQxBT8cp9helRzplXim5RNQlUaFCsHjzgwNANcnGkAAAAHBBm9RJqEFsmUwI3/vrPjghSM+3zdXqrzb6QyDQBh8wNH96zteG2dJP41q9qpbbpx6VfLx6pZ+hD0bxu4vInts9gfPMBlJkiYledG68AwhZomhLOtFb6Ji8an7vaJO4LtZXK+JkUFfz95QFkyWqJ2vBAAAASkGf8kUVLJ+ASbe6rPuX5kKJVKyfzTytHiClBBvNDUS9gyZ0gJuIACm6J4/MaPYlZAhSExx1G+JA3YMjZrZGz+SE8khQ7v8Q6ZsdAAAAJAGeEXREf4MMldf7IONNzwuYXTWVXiXI1jjTphyXCSJiDcl/JQAAAB4BnhNqRH9/nILJEpexGmpP0YxoFXxtFiqwczbTViAAAACUQZoVSahBbJlMCN/7xIFmvqKhxfQ7eu13hwX7Kl/0+GfUgZl6JYRtDtsk5rFajLUh7O9Y3Q/862jv9ZqodZkpyU2BPt5pwB+6u6zvaK5mzw79FkINXTazxAwa+hbdsXMsxsTU8kbKcxOzDPlpJOZmg+4BqbKzw0yHVYCbPCfiUsYQ1v/GbNQjREVqvxvErLVTetYAIAAAAHdBmjZJ4QpSZTAjf/vgOOJ6G+EORJB/i201ADDOh8gdy7TnqGzfuGC7fVt4IG74ck4fwuW5dPJSAv+ZSwOLk765EsOlxje0d6xKBDadZI86kJlJRifx+c16vpsob8hqrISDUct32W37mGuw8GN/sXV3vuo03HQdfQAAAEJBmlhJ4Q6JlMFNExv/++BEqKNR9RoJeFo2wJjCnq8GbSskuCaZk5J9oPynQ0n9SgOlT6cKPNsF1jNv7rf+cZkgbygAAAAlAZ53akR/dzl9hMKyWrjL+phdX+F0v8Q0kOKCEPls9CPUogjSggAAAF9BmnxJ4Q8mUwI3//vrPjhLnY7GuE3FgwVQBmmbcyRjHJitYGt3XJ7ttu2Rs1JgNU2hsMGCqQuf9g78vISwHq8b2HhwRAZXrups5jQ3v4TD8YIj3H3abrar6pujeVeAuQAAAFVBnppFETyffAdbc2ZCKuLu6g8BJqH4YE5YGtjEE6lGSZYtJ5hHsbl5fBMQmN9SaUIa5J0zft9yHmPwiRAJF0kKNTHMtbuNxCSscMTsy8L22gcicp5GAAAAPQGeuXREf4puFOCvsKrmMxdXS9O7u9OTUmOHpv3PBWQVCIk/VLpaFCq8y10BD+dshRrJ+UgPns8tyn3sUsEAAAArAZ67akR/inDB6JvALVyfkNmDHBAGn8T8vQbjDxek2OXAhfP9YxVeZBClgQAAAExBmr1JqEFomUwI3/vuVw4S52KAs2D87HYORWxCgFxVZ0fUYtB1vGD8zSig7GaR6PYdlvzo5m23PCr+yMcmEnzswwPvpxVY8irlJGnwAAAAQ0Ga3knhClJlMCN/++AchUq0qS3q2sujrVLxU7zB5z9GGi4Vr204wSsd+TAv/ZDLXIheOE/uGZxoJRecWc8xTdfXR0EAAABsQZr/SeEOiZTAjf/7w767baOIWn1VtSMPFg3+VI2un7hfAiMrNCPxiyraA7+7P/b4mHhjf9frw5Yf09s97+lHoLr+J40YfcfgrY6q0fAZ7vuspNsTPyaoTWu8OCda8dlFOKXttYcmZn0MYXuAAAAAdkGbAknhDyZTAjf/yghC9V+KM7GjZ3ifEeQQyURdNhHoTXnSTZGK6/FJfpI6X0y4F2LWLeuzT+DVqTb9zTTxzOILEvOnJySq5VGcY932fKbq7VpGGOVsYCTlw9TiZnaKVEnbuEjCDdLoHCPxhxVcfUPLCUueKwsAAAArQZ8gRRE8331Nr6PZBasOYw44aaDFj80ekC22ktCh57XyldfE7Ft3UzEJXwAAAC0Bn0FqRH93OX3cYVSoGPHXSBix14TW1QEOn7Ij37SbGA7W4x774czRQNnrfMQAAABuQZtGSahBaJlMCN/76z9RJhOzd1KAOnRwvW2BRZ7n6ANrAh48/h2t3g+1nVLUJRKqkD0Z9HeHn1likKz+6VR2gWMleErAnhTwmDW0YHrepMQTwNh+0/Fx5VBvqdRhjKSWumYEJreA7iMFKrvjlcAAAABbQZ9kRREsn3iZVH1hPWCBP0IcLolfj3IK062Db1vtBmUVEFxqhE2rclqJUf/kvOICwRgf+8cuGE3RT96eNk/u+JabuLcL9MSQv341jaGaaPUYvsPjkAuCYD6NMQAAAEUBn4N0RH+KbhTdyyowLli1rdXDi2OJpb2YMOsbPWx2VrALoxL353Cur+1DqAI4VUlPg/2b8RvB3K+6VuFp2qc1xr0ccmYAAABDAZ+FakR/hNfALxWamHwtC2+RNj8gRL9mIXf4wiYmsFGPVFdbv/eg0xWIw8ntIxBwu+sEvpuIcmRgof8iOS3heddgLwAAAJVBm4lJqEFsmUwI3/vuV1RItArTKe4aHtkbhVf/+Q+mhmXrHW48v4hipTzH46Ij37Wg9Db9jYlSjcF5jNQ5v1aWn65DMLo5D8XKRv2qTPmOytr/R4nJgxMjqpDiuGEXP91sGSFXXvDL5KhIcPu1Sv5rfiRKmk75SupVGhPOKnffZOhfkY7nU8WILOh3WPGj/TL7m/sJmQAAADJBn6dFFSzfgR7dVmJmGX8GHkqNZi3rEUsY+LD/L+bIktrWSlLH0rvcGVpVDuPV/0CpgQAAABoBn8hqRH+EUzCvSGusQvuCjCEcHYR+2+CswQAAAF5Bm8pJqEFsmUwI38oIQvVfijOY7Jqv/ZothAXNYBl7PM7Lu8cHbpXbUuZ1KACFF9u9XqIPlaA9xYZISQn7YtoCcDfTxu4CUZabiiBbVAJ69uAqIzwhmC7k2sQfHOcxAAAATUGb60nhClJlMCN/++AchU+aExJFTL9PPvLMNdHTufZ2ARRjtizDUCD7M3x68CS3Qr9dPeG8VhNRm6WRPsaNFzt/4aU270d35/LBN6vBAAAAU0GaDEnhDomUwI3/++AcgmDYEE7VgijEQg9aSxHaV7+OwetiVVzv9x0Bh9mAJqRJH8Dife3qBjyBpZMx0mPHKfKwSXf1DK2eIEOyj51/Ibe43gjwAAAAcEGaLUnhDyZTAjf/++s+OEudjsH36ASmbVNFHT+He3+dsKUY4kF1Fcb94L//2T/GGBOr9knUA4Fv6E+UdIJFqatHfWfL8/R/5FZzFcRTOFGht9DgyahVNF6gH94cGhO4EjBjUyOq5QvusqmUVB+YjvAAAABZQZpOSeEPJlMCN//77ldUSEVhRtuTqSuwn4ILAX/pGTDP9IL9t9LXL4gNSLMxC/6/2jK9Kp6HiBHIDwudiAEToisNqkYpE0nFGGBP2Mwq4VRrWiM0ZUYi+BcAAABRQZpvSeEPJlMCN//74ByE6AJUf3GJyJEuF8A3Qf1kQcJRR6jInzvSF1cy5rawxyX0nBOolFDiu817lzRYe7/SVVdiebEP++yydeb7tdJ3zxI4AAAAT0GakknhDyZTAjf/++BEqEiYmRvIgNxJIRznuGsny2nrjSbaGRQNjLW/f2j0oaCQe5/+54Ig6O1mLlpOJIun3/37i+STl9TMuA2Pp+8ndHAAAAA6QZ6wRRE834R1pCbTRWlD+k4+9DOjI2AHXza1raT35vBLoKirdVqNYgQnbX4DTT+ozWwyrExyRLfzhQAAABwBntFqRH9/nILJAq0W6TichzvXKAreK4xRz2yBAAAAXUGa00moQWiZTAjf/AOiKyEjuJtL+UDu/7/1ImEqkms4m590rIJHMH/lA7MFvspdRcAY8czYbgMQ6vpXSSw24CC9TXLj4CpfdwB8SVMwSgaDmTxDdOzD05PcebFnHwAAAHxBmvRJ4QpSZTAjf/vgOOK5ljYRuKx062IYgGdxv+oJsIKJXe5wzU05EYgy6CDVnrvgrKBM4IITLfdtFi9dSb58ZA2BYSlRsjpR6nFhFesKS1PpyqD9tlHbqy54P0NR0gcxVnq+CkYWBsgu1ogKQ0vpvuycY3R6AjqPjNvYAAAAR0GbFUnhDomUwI3/++BEqWWdyZKFzbFLbFx3xidN7yP1nsbxqbqZf3+mAe1LLS7fO/2lqY5lAlA3reqb9jagV1TR1CzRq2ngAAAAMEGbNknhDyZTAjf/++Aci1WcPt8e1pWqe9fATAcvlum6kTqhjA5k6uHd3lKWBjdgiQAAAGNBm1pJ4Q8mUwI3//vgOOEudjpLQF0T5/HpcMVhuo3UrgFNwvqy0NaDRsYBORhen300dcdaxgQAhyHLkqqY21QvhsGVyf926QHjO1dTQI9IklTM6dFWv11wpl628oyD4WUbTWgAAAA+QZ94RRE8n3igS8fsiwsG3BGH71JSuKL8DGQXeoxKKRzUOLM36KStM+eXk1e/WYIgappa2TFwUOXYGRkunvAAAAA0AZ+XdER/im4U4Klw8nZaidOEir35twZNCptJt12OpowRIVZ2i141MkhqsNTXZF3ZZSs7KwAAACwBn5lqRH+E1CrjX482235Jp4YGEincKNRQy+bYAzUsHzuPxfifmdaLZN0gwAAAAIFBm51JqEFomUwI3/vrPjiuZCYQyR8N4L4ZpcVyVm5GOzBlZrs3WY0M36PIs4HlCuJLEY14KxyS8bo4xIH/xQxifCHSUw0kKwQds/H/qM/Bfv/37nK8BNWssBs2UpmS7HvInr7RQkATNzNBxyIYrpPs2K5w38HW0rEzP+G0ztOo83cAAAAqQZ+7RREs34TtSIINqX8yEwwllxjEWv5dp22DsRrlE9pSWyEg90j3GhvBAAAAHgGf3GpEf4Sa2x3nqQ9dXv2XRummgfaBLvV0z7jPcAAAAHRBm8BJqEFsmUwI3/vgOOK5ljYt21R9BjygQzFkvYgmlW7E/rzUmB3aa1R/v/sRkWdXCOgIKuKUfzXiJlEyDC3eROqOf9d/4jH4mwkcdJ8kwHn1NDcqEmz68Y5bf46nu99czcvNzhvDbgh0CK9lJoeiHSXoTwAAAFJBn/5FFSzffU3evezXtC9/ZQKRQHpR1ZwveCVG8cv+8KloL2JUK/o4OcsKclrttQWDgV9mNWZ5P+28tt6XuycTD5aBLcSOUtgKm4G46iDPAKbMAAAAMAGeH2pEf3c5fdxhWS1eZ3/p8a5pR6sqQeE6lWFP5z0K+aYlfBg11Rzf+F08p1cQ5QAAAJdBmgRJqEFsmUwI3/vrPjhLnY6O54bcEV++H1QyWusE03qBnqIsvXHnWussWXieHNrrthosgcAU0Y0F9hopOWEBv0JGXS1m3o8PuW8QsFb6/rPNSkXUkO/oFNfmFTp6wgcdbxk7t64LF//YEwizj60QETaAs6xDcg6lk3U1uH0HDMl4ih5VDERKo2WeNr7DJ6PAseHZfNpBAAAAZ0GeIkUVLJ+A5qo/ZEObD72LxnQYIuEYf1ytAEWZ9RqVWhpP9byyoVJnVZyVLVwDIFybDJ9gG7viY4976Z8PP75RopC/RDB3dZ0wq5zC/a7Qp87846DKhZFGiIsYsKK9l8waIPbfwkgAAABDAZ5BdER/gwyVRD3qoTg9Jl28dVUyLW9E4j2mbSLrdlU2wMBPhJBWiaLiTejDzYOqwm38CZG2crdvRPAPjeGIWaTVVgAAADoBnkNqRH+E1CsOUDlTRB+2knMt/XIVd1w2Lu5z2obTpdRmkuyfrXpGQRJ2yHOSmMjOiyrLDlpwJ3dZAAAAvUGaSEmoQWyZTAjf+8PysOw7IqechdCyDapN16veOkycd3tl/lOY643J4iTWQW0XQRO1oHAqDTs60fM41KeAMdMg5qceLz0xHG9o7K7+QDwexR/ef2KxGk+JPpEkkQPooks4C/flF4jhBo+xQ1NvAor/CROg6sAGmqtYsKTo2RP869Xll09Z9g0/aqKgXet/XAH1/pPh9Idp9o5m+LTjIUlHJDcUkw+f1s1KFZPXPnjINmfl4g7HW8mSl7TCcAAAADpBnmZFFSyffGVVR+zevXVKzm2vgkNswDKTj5fZQrLaGvFrUNGZo7oe3bojV6WjyvmGIqww/DM1Nn65AAAAKAGehXREf3+JdlJAdWrJHHXgk96zCjVyoQkJf+XToMVT4UbfSa64B0gAAABCAZ6HakR/iLG7Hu2qMv5WSPgOGOgEf49LPHtgU6+/FKGUK0Uj2RSz8n/2jS/nuH+a5JnCsk379xSbYaLKWkzHqqhJAAAAaUGaikmoQWyZTBRMb/vgHI2b7Fd6Ws9h26Qd1WAcJFbDOuHQlDOeOVSDtthUTx2BmQWjIawleP3YJiZxUceVqUKAtVJO3MmSt4itsmeOXkS96GA/02xr8e3XWkoQ45QprFdd1gNT1R5NgQAAACEBnqlqRH93OX3cYVSoGPHXSBIm8PYCgiS/lOkz+kEEcdUAAAB/QZquSeEKUmUwI3/74VOAVE8TO2duqVlti+FmeZePWOPchj2ByBQCHzxIwA7oRji+KqKa5n4Q6OpasFlu8pZOpT7HW7iwn891D4Ifhvw+ecPpsngXgO/F02GWm9i2sspJD5g5fSDd6gF69m+tkP3mOA8J0Iqx5r9Ybt3wCsRgDQAAAFVBnsxFNEyfgOol4/dcuaO+lc5oqdzUsBCEsAWXCEd9LQnFt83M2f4mvOGuJerHkxVQwxAPUafNM6hjcUDccUEXrpfwnpDY6bBBuWi3an7M7S8Zz5MyAAAAPAGe63REf4MOYBhgWGIje2c5UIBG+OMYNcDzO/HP5o+6xdUgK/xp8ClpOu7ZqMK+GL9wf1yEuViw3Jz5MAAAAC4Bnu1qRH+E1CqJbRfRlHBZLSjlMfrSL3+KsaKQb79wkzKWLcwNd6o48CwLbJZPAAAAuUGa8kmoQWiZTAjf/AOyX9R0mJ9SwB9iksEB7ZH/nkIpajkfaK17kN8kNdFdkevjT0goBYSx1ad7E1vuUhe0fKOO67ud4jxPyVzw9UyNsylkIBxchX2bMrVXnfEsnPA9mDRN8mbHacVFgb2ToC0gaK1kxLRZDbbQ7G9HcVMynbMoma6dSOovdmZyIzwlpap73j+GNAutYydKa0z8rYc9VMPkEA+pbgbAGRXIRlXA4YjOfmY7qK4sB46wAAAAR0GfEEURLJ99nGGl1msrnfy1MSO2So7YfJXmhMb4mHw7D/nEENpn/g8dYIuHWptpGty383qQsQOmaDb1Sr7SyZujuEtrU93EAAAARwGfL3REf4BsK1ffX11Zae3WZvfEuHzoFi2UXR12VR2woUFdnWnJRyNJzEMMur5pB08F12TS5Kqh0g6hXx3s+h9IciwahuVBAAAAQwGfMWpEf4ERwdJRbH64kx+NQysevRT1D2SC34ECzZV45BjVHUr0rPEmr61jCfTNQyQZxfD2v72JxfWvi2EhWwvZnoEAAABvQZs2SahBbJlMCN/KCEL1X5n/pbtseoLLd2WSFVdinV73wKbPDDMQGOJ7HQlNf01lmQlQxpN4Xld3ENl7/XvF+385IN+ANQzbxa9MY8ghxFpqW0UYwQhQSaaY/9FtnQOc0V/gREgSgiRtrKFqyex9AAAAN0GfVEUVLJ97HUrykd1VJNIYMgfZauRWHu7JgJIJWyjt3yOsBMa11eyBZgykRdmoIPb4fmsTH6AAAAAxAZ9zdER/dvu+WAiUqATXXSBiTXTh69DK9T96yYRRTZn7NwidgRAUaKpqXh1LkSJxnAAAAEYBn3VqRH+KbhTgrp8WYLnb/OExBB1QfQOTFVBGW40ihRipF6j5TyV0K31uIo2+OzX6N9DLAIAws+5y7pu8ISS/6hJ47suBAAAAk0GbekmoQWyZTAjf++s+OK5ljXoFiiU1XPjCnQxFyu7PryXdiZtWx/wW4N5zId9wNpopmCe9CPMLQBcd5pQEY3kNzWS2+3CovTQy4TUQSeEemdF5upIcffV6vhYfM7AT24IWEWArCyFFh7pIe8+zc8VPXEM6rc+7xdmneJ3szGOTGBGxzvpFxlTTTGKv3pJIbNGDAwAAAG9Bn5hFFSyfgOahWDUeRVYeX65PXWlnysfqmajig5ypCQFzzoXCSRAp3zgjG2NG7gxdp73zP4B6ldlcHso4oxhpX24XpT/qfkQR7KKVU4NNmqjCtXvpDalQ/00eIGvukzaHkWUGHTKp/QtbKsPGwDAAAABNAZ+3dER/gwyV1/tYOwEqTd+Pf0jD+/NLCVKy6ez+vl1KUGnCmGBaIywcEmAVF6zlPqshdQhK5FM6KDE97ysl/l1xLpPEqrALZu0g6REAAAAtAZ+5akR/f5yCyQKtFiypduyg6b/z69WowFlASrlNd/l1Q8J8ACA56sgTjZmAAAAAuUGbvkmoQWyZTAjf+/4BIhQKAXMPK6lD76MX584FxB3ojNVi5Q+urU2Vg7DDnw5TYhPAcAP+JpuOWYUZvbWvqxIf4+zhikw5W4uf+nNx3QUlfGE+qQdeKFt6AVvU+gOqNlNIKHEl5X/vEOA8B/wsTOAL7//f/G5Ez96Lbtif/QxRwf+GV/m0IscN9zKokP5tFqEx9jubCAyMXQbzIGza/AdoAcgWAzpGkkEFPVZHIn262BXUMjV6hgYRAAAAZUGf3EUVLJ94mZXaX35boOk57/Jas+YqlIwzjlszmhvu1CGdNLBRaqzGIA8eJ3iguthSW3p5XKu2MOj5ceJGlYERu0+psmSaLOFXONszzF1sh6Wy11Cr5kB2FRpfgC7N0D6P4v4lAAAAOgGf+3REf4MOYBiXa4/L66IgzTpMUW3pLiTwVe2YiMlu1Od+TX4Ego0CvBosfBkHjMBfsmQviRFuwkgAAAA1AZ/9akR/dzl93GFUqBjx10gYYvCa+gIMfz7WTXtZIsyH48Q3lDcTKBK+Eqf3cTexpuqr75EAAACKQZviSahBbJlMCN/76z44rbBJLPdsY55JpbvEurvbYmuIRcWV/rJn2ygwgU9CgdRF/043si4Zr5zKodFD1SgFp7DQqYD6tTX8dwBVP474YclgRLTJThjJ7uLHj6V/CJc4In7m8OgftrGMn+Zru0NctFHKqM1bDUTtaFM/qXypCDHtWSqgegcA8LCAAAAARkGeAEUVLJ+A6iXj9eUXdhdZVK7LV1HMdFPLHxlyPlUiuab+bMNEm5whCn7SWav9ox8Nvdiza9s3scrTm/Si0xxwSXwP4ikAAABFAZ4/dER/im4U4K6eWIXMANXzPh52kzBolbrDVKhXJMRmKIXiXvz9VXDYi5ff6Qo0gHFiqU6rW9yf7KoXUhBe4Rxlo8WvAAAAOQGeIWpEf4TW9RzxeaEIBMuIawIw7r73W/EwbmR42EU8aJ8SmWMIch19eiKXWJLO+HqvXpyEDPqD4AAAAI9BmiZJqEFsmUwI3/wj5K7y3raNV+lN1e0Gmx8L5RHw0IgYfLeHCZ5xAiPYSqSZkRczptAL0OS/Eadl7UGYnx9tFSTMq/DW1+S0JjpUGH2BBi2JBcBp9gBaH1YFApjEEAfskKGqdW6ZtsgBB02//5EzIeLDw0zMvQr1qNKsH9QHaWrVtuaIvnOOjm3+a/MpYgAAAEFBnkRFFSyffGXEvIdZdMwjUHspZ5lmIjrKqSiPV9+51l3FYjTQ30SJlU6anVjo4RznACIj/57d3zRCAGeOX36D+QAAADEBnmN0RH9/iZVJcpUEaVJ+D1Qn1J1MON4UlEEFw/V4SMOqbANXYVK9+DIT1C0sj9H8AAAAQgGeZWpEf4ERwdK7XH6zbP+0KsgLE2aJGY+DvwxwgrIkeYD20yzoRSouFKQ3gK9JnSL/2w7s61b8noAFMeJsnjdrGQAAAGhBmmpJqEFsmUwI3/vgPUUrmLiGAWYzwddaBPlC1tTCKPJyTpF5x5Urt5EikjDg/Y2Ch+L/G5axssucn+L/XqrCvrKKAbuvr0oM1XhQMnNPkrUx8G7rkqNLrOkC/mXdj9DtDsSuu7cqxAAAAEJBnohFFSyfex1K7S/556pvzaNwUuIyDg22DfvrCUNWtNpRRmvgNB5PJxjDd7uUat/w3BVRmbf9xTYsRqKfuotnayEAAAAyAZ6ndER/dv62TuMKpU1pPPlXna5sbaR+8FULtAJqz/eyyWfjyJNab+zrFLDXgmE1X7sAAABCAZ6pakR/im4U1zmIP3cv/+aqmtN99AtaR/ePPWvf3Kp9eqRfWQYoVZuV8xVd50yhmTsQSs9JdwKpAboLwg3uGj6dAAABHUGarkmoQWyZTAjfzkcmRUPipKf5HPWzd5mVSXOkf2wveByGyczca4WLzC6dJy1m5B6Il/iv02CYxAuC0JMLjf4eVD09cAC4VuqW1rQdH0jYBt9fqu8NkdwlacAwc4MltNuRgUY0PZqFq9cUnMVBRelxNywmlynkI02yNJIx1oubSc9FYoxYy1GOdo4XWxNYblENWmaGnuCBIJRx19vtPp+h4vvQVBXbfsx3OZWx9A4nT7Y56PG6CxBp9NdcYCm6+YEXmiN995AfuUddw5pT8IxHOckLKvtrExmNawZ4j/qtgJQASVXghncWygED1xlvOUMsv400IbfHtGjY+xlxEauu3NbsBWE2bmc7OC5AYARFLBZS1IMRB42qDSEAMQAAAFVBnsxFFSyfzWO6eV41HkVXWetrTujpSLSxsyet/K+xuI2mJCQpfU8AiK58PkSxjDCtX3agsIW3wcJvwWe2BOb+rXnGIMtnblWi8WfxWyNGdCgOJmIQAAAAYAGe63REf83T//106i/j9g6B7W/+9pd27v4NIhMKDiG077SiuGlpjPOE7m0C1ERXhoTXOKeyMZV/Cqhwx0+nXZkTcQvFvUXkgMZhV3TQViM0RySsmTCOTLoKGrX8MDQLgAAAACMBnu1qRH9/nILJGk/YUBDweJZSU0TGSR6vSvqHjaBT1vIf4QAAAN9BmvJJqEFsmUwI38erBKuz0hHXKlCdZjiQgCATySoqsm6ltfbP7rAFGFZ822Yw4uKROyUID51JSZ2TSrmysf59dCK0wX5bIkLdjmpkPTrcM+4OA3dv70JvEV+piUKzWy4rgthcTqUOTCLotioch4OuJRHZVarI23FHHeuIXdoOx8okOazHENHgaw6X59OheNrJiEvne4/W4SXKVQV9oGUnCUDqrTQ4hxMsBqxs+I3H3zkR1F3Wqbxyocy4mEvJ67y6N0vizCEqYxGvHdXCKAGvbYr0oZ5UmemNFdqa/+jOAAAAO0GfEEUVLJ/NSlgemZ2/i+2HE/tCoU5+SvCMe9vvnjLZqsDRHgVJCIy7oHSX0WQfEACM+t3MwE6ru1+AAAAAKAGfL3REf9fXXw1RMHznnW8D6KXJTU0SNFkICnZ+oJHgBulpiD/kF6UAAAAcAZ8xakR/dzl93GFUfKEWvCRfKO/B6X2IPGTMgQAAAIpBmzZJqEFsmUwI3/vgOOCEetSB49qntGrM9qobba9AH1IQdtyfiRpDEBqUBn30cNOTlel+ROcamSYWnUPachBf2dH/7blMdcvqBzNRPhHUnxCwzkrNRlC/5GgwjKfk/a2B/o7qVamrbXYVkHxVa8bR73b7AQGyRnIHvBjdQa2Vyk2zDB2JLzShEC8AAABbQZ9URRUsn3igS8fWMzYXk9Mq+MDqIxXVuhKB5vQkgSjtJT6JSpHmpOK9dQXV8FSdFrMVLn4/C1TOlGot8odWKsZCHVCzXaMLFwpRnRtzffgvkqf4iecFvElrWAAAADYBn3N0RH+KbhTgr6xHMY/2QndZPRpMsv5JSKrqaGIUN02aOyCA/tVS67KULKPzvQ3M2ZSKp4oAAAArAZ91akR/hNQq4woHVjdAiYoVEk1dl3y4BaGjLWpTqLq69j/cZmOzoJOR/wAAAEpBm3hJqEFsmUwUTG/77lcOCFJBa/WYM+C1k6KBz1fXTj2WWIJN3lFiTnpO7wj8DQA/i7cHqWnaZRx1Bw8Vic6inCrCA4O4XOxsvgAAABsBn5dqRH9/nILJGie4fD5UaXXmEKriUh+FqfgAAACDQZubSeEKUmUwI3/7xIFmvHCDZ8t12akwl+l9pgJioglFo/qTXNJYNkUFQp4qDJhzpCYfFAIWrXQ9O8aP+CIhDK2kFEhLIikaKbSUWzNbi6H4UrqrGmAlkZo5re9RQGhCccg1Jov/2GxYYpyEq9L39zvDE2Io/i/vPk4P9W3qT6jP+oEAAAAuQZ+5RTRM331UtYkDXGvB3dApBSvf+bMEW2P5xUOFFwLXVEf/avPuP2djgg2IoAAAAC4Bn9pqRH+ALNfijDv+qEgNt2DwKYhX+aIxYXAHwEiNC9Yk6O2btseoyQdM+47xAAAAUUGb3EmoQWiZTAjf++A44T2yC1yJtTnpkDG02asj/rdoy39ht3+VB//uMshD3UPhmeAzj9IjmzSBMZH7Bn25T3xc51+ObVJBTlLroFR0mja7ywAAAG9Bm/5J4QpSZTBREsb/++5XDhI2sYlYDT94fX90cHdUWtuMJHfC8cCICr02MBGpCutBAL6UkDKLLy7zPKWJhLoLICnT12OP5gktaUg7lCh5ilnYCb2P2H4H+nX+3GiUhUVoSBX2x9WbuvhYGX30W4AAAAA3AZ4dakR/im4U3cskkOWaidJ3UzcdYIuqiRFaqXPiJJRWssAOr87hWTOg/jtyd7zuRXHH9E6yDQAAAEpBmgJJ4Q6JlMCN//vrPjghSQWWzwVzm0Fb69aOUyH2uYwkl7uzi3xrhyoi3bG6RP0a6qHuC4sp6e7D/Sqm88josE6dK/dzf/C3/gAAADVBniBFFTyfgOahWDclg+s+Jkr9gf7ZOS+stAk+S2N/q5u0Nk+SxhbPMvBKtRp9lBeEwzIS0QAAACoBnl90RH+DDJXX+ypzivEOWRG7zcM5EY2ZDjmRy2dsh3v56A+eAR5ct6EAAAAhAZ5BakR/f5yCyRJTQ0pmRtmpfeSIf6wjFYJ0wy9Ux4VAAAAAX0GaQ0moQWiZTAjf+8Pt8hO4AJKfYhcjA9rjdnAwFf7QModT7PyzpMUK+CS0HRfcsjc2DBko1hnWHk2EmhC6FeTs1eU414OhygKFL6xPmfel0xhg6rysjXFWdjRF1jxEAAAAi0GaZEnhClJlMCN/++A9RIVzu3Ghrbc0Z62jyac4b1qAl3BBQ30vUIr9JNcG3B/7thUl3gfcZJEcDA9/o5HLdtRPY8fZ+i/Rg0HkiFZFC1tsDbnRuKWHh6GA5fZne6xDlbQRKVZqTruh+xSV166g8XuXlsnFkjpjI3LHla8Lz4NACVbazHfhnbOqwIEAAABwQZqISeEOiZTAjf/76z44S52OkshAe2P0PHhV4Dei7KErq3dgxYM8IFzbxvZbXSizGnKiEKPygFN4FMdrwWYARbLCT4h4G4eB1lTgxQGMp7418j0PFqDnRSjlz6lPtttNkfg1a3moBhbWbgd7fml1UAAAADBBnqZFETyfex1K7S+Yzn6tUA1X2wuv7z+sgh4e64i0lTn5rJfpZRmkJsUMXmXEYcMAAAAdAZ7FdER/dvu+WAiUqBj0v7NQWvwATKba0oofnS8AAABJAZ7HakR/hNQqlgeCR5WmEnTtgY9GoEa2D1hVoMAPb44LseFckgc9p/mRB2QhlABhbflpFWCIhI1jRMwGxXP5gv3zZcieyZt8sQAAAJhBmsxJqEFomUwI3/vrPjghSM6GitLPhfT3QjDIXqDR2QygB3rhZ0Opwh4x7j88E+9pDmthzmQVl93MWqAlaOPoKrUNvWPRoPXTCZXTTpy36yAiQiFXV1QqO3LKvj/ZL3s4SSgWTn43YwKoppbw9S0sfviVAvMSi0iDeQvD4QNHMF6CPVgieFql05KklKluC94yEGBUepu8pQAAAEdBnupFESyfgEm3uqz7mAS3M59dNbzoYEhAUaCO6CVGZpltV3a4ulLjDLMTNTXCXVdm075xx+6WQowma6lWj8Am7X3wHh6YdQAAADUBnwl0RH+DDJXX+AaXoY5qOgBknY4s2h3Y14HLG4oExzfORWc7aR19WMqT9yojLIGmtjRfQQAAAC0BnwtqRH9/nILJA77E0uyjL+OYRlC9eMwkZdFHKSM8oLegYjG2Mku8ZesfQpAAAACeQZsQSahBbJlMCN/7xIFmvqIPoV43rFJvWX4ODfVUK4ly5WYitO5XIzfl1xd8KVyzDSy+r0BDRR//Emv+IrjHGw2ajYgXdcYjXyNQqWjfnNEm+mjBDaTGTxJzcJx7tSgOp285AZ65BM2cdcwhRN/7j+yRSrtO21pzFo5NGhAw7XafqedlhzTYP+QC5ZhN/egGmszEX1+2B4unJ0IrXIAAAABBQZ8uRRUsn3iZldpf/xqA7EEaL/5ZgaVWGAVcD315eHZ0QAYXXrP2jWx1qG9PMHNp4FYUm6nfg0+ci6cjMQFda8EAAAA0AZ9NdER/gw5gGijNiQz/5kp5JhG8oUICnRs7uFapasLfsMKPh5eAEO7oRpw9QfNO9PihmgAAACEBn09qRH93OX2EwqlQCa65akqIwjXCrER9aHr25SKikoAAAABtQZtRSahBbJlMCN/76z9REFwH6H8EqFmA4qxAVHo6OWo4ZywYj7NSM+cdJHCwf4LAdrJSFZD3glZ52eoTwrY1x19p1bxCO7ICdQGhzEiQ96Xc9Y2SETJOD2Vn5GWGsL7J+gJXntXk/X8jEs5a0QAAAGNBm3JJ4QpSZTAjf/vuVw4IT+j8dsgkolcVXPJAv/DV7tMKPOhKW0Eb6GkMP3S2+4Qotkl28worhzeCxOdJSjc6Io+xwzEknX0Ktf9XBtxbX0nSYFhFqBO/1e1O40UYGwmBSjkAAABiQZuTSeEOiZTAjf/74DjggZKH2CCRdBnEYQWZN54htqxWZjo9aIG1TMzzTiRe4WNKcKB8WgbTaefOnNjdIP2vt10+UjnumK7PI2s37KTJhbyphwMl0U5kDcZg1Yv7FQFx5mEAAAB9QZu0SeEPJlMCN//77lcOB8tM/mnq1/U/2ketb7E4WfIqxx2WSC2Pl4TSDg8o3v0A5Vjy93uLHMTZs+BBdk8ByTQwMAo1JaoqmZVq8af7jmwnevUWJQ1A4a64zgNIiFZUXmvCBDFKozh+l/Zy9LkBa4jKOr70skQOcFAQ9YAAAABGQZvVSeEPJlMCN//760EqEjYhsbGdgimO8QcnqaVCP9qbOEFZrTmsIf2JpB4TSeaxpXl3Ulw1dscvV1EvErHcSHNSQtr58AAAADhBm/ZJ4Q8mUwI3//vgOOCLONnHJ8xun5qHA3ZmgreRBzw7h4GYMtLDZh/dmKRwhM+8ymTIp4LzwQAAAF5BmhdJ4Q8mUwI3//vBkESMggDx1VvVtLMss+aF3SXvKnVrCTWWcSrFJBN0jS1tZHuusb/RraxhnYjj1+/yAsYA0rgXmW+bCEUISclYRiHSnHWCOd1ZBmdTg6KC1ZbAAAAAZkGaOEnhDyZTAjf/++A44S5zbJqFTmyU0bct5evsZEZnwCITG77hmWTV3ATPGRlTvfx5fU+QbyuvFKr5KrASAROTbHME0sqsyGlQ6eQe5/+7Q4wQpIcxz1EsaApgkTsdRh+yx8JzHAAAAD5BmllJ4Q8mUwI3//vgHIVWfpo379KAL+bC1vac0rQQmCNEFSc2ChJ69qz7GJwQcygSggLTY2XDCMS/LkX2TwAAADpBmnpJ4Q8mUwI3//vgHIVWxnyGmcsFcA9oxoBl/brkT10aH4nVw83GGNtypUTAP/suuF7gk3SK4d4YAAAAXkGam0nhDyZTAjf/++s+OEudjscNzZV5CCe/aEyxqBVLFh5VRKVEDzmEB9Dk7ScQ40u1uCiy9OuUrKIzx+gIl6P0NG1LHJzJ/qCDx1gum2N47t9UWnFNio6Aoe0ovYEAAABcQZq8SeEPJlMCN//77lcOIMTUMrHZOr0gSKOJUf3LozHu8r1h9EcXvqCO3mOg+d/VC0z5ZYdgFjL3CZniji5n4XXibsuJI5k3acNamysNf4eKfIu8h6JjblX+LUEAAABhQZrASeEPJlMCN//77lcOK5kGh6VLT1xLqj5X9bM3RBN5M3gOtGKjSsWOVbXT9w5waJDJuT3NaaL21k2UzbmVyf92Io8RDVoRP2XUDjiYmvZj67mQ91VySU1YVL5C9Sw30gAAAEJBnv5FETyfgOahWDUeRv7uaiJL9EdNC3IbVzoaH4e6aSo1n7bs0FRXQ4JWRcXvz38ocIB8WqDlMJ9G3aBO6MMl5bkAAAAmAZ8ddER/hRLKDHq0j4d47DM5gwud4/DkjlPlVFYGld+N18XNz8AAAAAlAZ8fakR/f5yCyRJHv89kYcozAl50ypnHlpKXwY51Kj7W0wQ1HwAAAHJBmwNJqEFomUwI3/vgOOIMTVxkqPJRLQpQUqKkxcew6gejhK9dEgRRiWKoB/5wIFCYAVMLtY0lhftvZnZTOruxFhOlh/7mqhOygkTIBTqYT8F+//eGeuYR96Ww48vtz7cCy86ywFxOdvLxYRAQgapIAWkAAAAwQZ8hRREs331UtYjmVOEob1vuAVgoDzC2PbDUkYLtXfYTEoiAWJS2xEv0kGt1b0TAAAAAOQGfQmpEf3/pQAeIEEpwh7Q6XdldFbTTk0g56u0PL+bxWoNXHexxjGmHvNW4Js/aa9dUwMqYLiqLmAAAAIRBm0dJqEFsmUwI3/vrQSpSocjmBEaZ+guniyl7iU1d1HS6ptlCPKacTBMZFdsmw3nQ/2SNFSQYbakk7FMBdPhxaD1VZVNg6EmmIShve7vOeawjwMsQ6qn7zs+OxVipM4voEUAHDO+9AdjYKvS5OsP1RkpHK5H8rsOaGDiUxa546nHEZkEAAABXQZ9lRRUsn4DmoUnrvS2j5/oh09N4LvXKgJJsJGF3dTgk9NBZHHzkIto5QN/1CwwGkGBV6UhxpaON36OliiklObJaIH+XV0IYQQBpF7sgk1BWy0nUlXZ4AAAANQGfhHREf4R4tlBciUsqFFuxA8UwpTFd6S8jYZCqRb3ZWGJYP1ecFzDRMlHObUR8+QdZwRJhAAAAUQGfhmpEf4n36qqIKnWUd9Njiw0QSxoVB/oX/yBhz8h1/xIEQ55PMKdy01KoAx+nvEYqAmAd9ad3OBTPQ9pOmJCNVk/X7WrPcvXMd+GCzMVi+AAAAM1Bm4tJqEFsmUwI3/vhUdkOqEdG62jueiHqqB7YeLSYauyhCAzqdiH19ZLvxGPaZ6tgKbJdVHcaPgr406CugQLvd3FPbtzm5AxeSWxsCek54+8ViKC5KSPZpGstBmWc1F5Fz8YmM3pXrKfxiE5SEFNGyFV/qkcpgvPYbfMJYpPyMe+Mok3YdDm7tJqwEmFuxpR/c7HOjEn/q78jrYgJECTdXUMmYN4j5En5F/wwQM8kA2ap5beuNPWS5eYz2P2j4BU9tNJrSFla6y+1AItxAAAAO0GfqUUVLJ97HUqt10psKM2AYDKr6M37fOlTxCJJ2VrMROQNURvJEmBHrG4u6Jh8f7q7B191ZIjfgM19AAAAQAGfyHREf4BsKlO/qXBAryvwFqs5i0G98z4g8HSYOeOjUiUe2uLVZBQrxdTZrbxwSpFGxJ6EKPGOOvOBc2c9rsEAAAA2AZ/KakR/hIfoLL7/BzcvpolKljM37OToCQ/2FbQB5A84l3GTkJR4Rl+5ZoIT+HkMEnzHYh3BAAAAg0GbzkmoQWyZTAjf++A44rmWNhGWCL72npbHi2D0rm32fIO2vw38MGXgB7/2AwMnd/lvD81EFlh6sV3nHxYoB0OZtYyq9r2rygD71P2o6DBvoGCJwkGKOl4eFiEYvK83H9emCgm3wnANdrFoijHM4zU0e2RMJh7A5zMM4OT2vz4r/00sAAAAS0Gf7EUVLN99Ta+j2QwAR62o44atyjblHXj8qcvZZupH136TPo8NbBq8vABIfCMtlgbq+SEavzCvqU9ljoozFu3hOJeYHrkwhPV+IAAAADEBng1qRH93OX2EwrJavM7/2lua/dEitfWcARv61hjXJUf8VE2rGPeKPTvdA5ULh6n5AAAAnUGaEkmoQWyZTAjf++s+OK5ljSUOA3TFai4RJE/jHgw0W4P9CLmTx1ROtoE1Ct/u5aZKpqL5UFk907WKqU+Fq87HMy5CDtuyG1GYY7+oe/qacQmKPQggXqXzfVJZxonCiuPjr8IvShkyg9jrNQn4JR5rA4cMrT0Zg58ZdIL2DxBO+Ztm7EZbuwpX6hiLDW+Fu7jZRlkzHZVGcLW+A4AAAABoQZ4wRRUsn37EZ8k4yrEJ3gXZR38ckSC4E/G9JeY6nvqw8cni+XpetBPoOPSLCtYJ2bS4fCf8qRBknduGn/zENrDlaOySIiDPyNSjlsUVOY3Be7kM82KFFJZ4p6pqGLRuc8L/WHBJbyAAAABOAZ5PdER/gw5gGGHvGKdek4JOkgzZuCPm5S4Jttsi7GrvSbpquZEjT+l8DryKb1bKgXccBTGldbxppBvtraeOctJrix9CJg+uODiF1InjAAAAPwGeUWpEf4TUKvadqH/skJdonVCLe4CSymbvWcZPcHMhqYKXxCXq7uBXK/nUvzCA7bRhXe4Org/oKbRblHoZbwAAAIxBmlRJqEFsmUwUTG/8IvJDIWLLGwa/ZHF/KLUN5/pBVYdMU+ZmVpOULfEv18kUVSeOAd3niURef0hBEmg6OBZ0i0XF+KXTRz2y9KRZZglHmsDhwytPRmDnxl0gvYO++Jtbp40k8mBlBmde7Jtob6nZkW75fdtnP5nQj7Bc2x52eeW7AOXILyoBAXy3wwAAACIBnnNqRH9/nILJCQH8el53caksCb78m2K8PektcjdV0nLgAAAAtUGaeEnhClJlMCN/++FTgKPVy25+n5usIXnTztuq86gMcZAv//lXvQMR4OFxAH4HHTYRKq4DSRdcbN2fFwL+DmtKdirEql7t068169vI3EBCCne05YqLPOllLz90IOSrfUDiyyb5NfrDwLg+wOh7qHwq2BSxmHlV8XQlZWmA7Q1K+zPZSUHgjhs8Ax1n9XnxVLQtfRLLyA/dBKYlK9a5jC+wfdcMLbf12+YuD9YhzRWeXQP6UjIAAABLQZ6WRTRMn3iZlRG9wMbG6rulz+i3xcFmicnkVfPi3vI9sl/efAzv78X6GnDlL27msiMoZuydUOS3pIwXyOEjvXuFAt/zUCd6+qXJAAAAOQGetXREf4MOYBiYRsM/4Wn9J79PPYoQFjb6GuY8jqTLyI2PeCs+yCv34RJsSxBkwViDKVa1l5B7DgAAAC8BnrdqRH93OX2EwrJavM7/2lqioT5DC3FIJv7vm//z26CikO4tVaLBSjpb/LkdYAAAAOVBmrxJqEFomUwI38n9ga29JPKIOBHSx+YsFh8iUW27g9I27uSlq/OZWnSYhIZM7ODNk5Eyw1T9OF5rr+J2dy4gl+9RDUCQuNU7dxcC1S9rJoB2wj0ZCKt9fx0BJCscbIz9DE60T8UJcZW4IXI5PQ04rPvV3fOyHU1DTpW5sLAkcKbZY9J5we4BPqx7YVed7ploglm/l2SkYkxjUmxsPzRVLOnRXmV9hlstOpUJ/aOYKohirbc2wOqCvDKXJPhu7uw6sdlwNg7ct2PF60e8zKjuPUQcgh0U4w3KxPWl+MIuQ3Ce8JUHAAAAS0Ge2kURLJ/NY7qFXvgUaFUOZCNBf1AjC4hZbIHjRp6oas+6Qqk70vXUDe5rKS9H375z+RHnXGdynCNUDFHftXvoHTlj4VbCTJvnaAAAAGoBnvl0RH/T29sxF44A/5fnAIOAgxy45Xst+FBoNPddunOVgbfDRXfDZMjQfcvkgzJf51llDHGEOPZdOqXjlDVUZbZjfZl0loJS2YOt48R/pzXYFOhFufNJzmM0UhgMuQR/zyOciCsNb0DBAAAAKAGe+2pEf4TXwDNrCamp/ew9bCI67/4OVldsPjdTTI086/CfMW5DjVUAAACXQZrgSahBbJlMCN/5KYRFpdsBNJCPfk3shMTTbZ/+m1g3BC6EDF3hvpzFAan1Nwk0GhKdlNqVrfEyKn/lndpE6h/0uKfhJLNPRapdj7uQNAbjcketIZV80xhYt/ITlWMCxK8VuHPNA3X0fatIH0MAuSlD9/rmnTCGX1rX+/0XG+Uq7kJlIC2ViD50CtLv279/vtSKlXHogAAAADNBnx5FFSyffGVVR+wTASuTnuC0n9aZFyEaIux3eFAxaKRoR6o1P4kyonm7L6VYQyuDgUEAAAAeAZ89dER/f4l2UkCrRbpOuJ2cpJvutsxXLqDlkAIuAAAALgGfP2pEf4ERwdJRbH63u/+0Ksgh/Pd611GWPgQqLOAjTiLs6CS8xHLR7EU2U8EAAADGQZskSahBbJlMCN/Hz9Z+6Xzm3/q7oYd0G2hnycvD/dneZ7oEkBQficNkv3sQWvRjB0O3J0i7f3L9TFI5xTtQclRChNrb+NvDDZvgNHtISAsf4R3ClMzNDi8Hti27fyjTiY/1PVoJovYx9gmD2NorCNGDW85eiStJX8zbjsBw/hmP8sN8jZiXmeBRjX7HLakmUV39clUJgwqNnh23yHABH+5ni2vm8+lQKHKUjJ9gK7V4lXqq7mt6OcgdnKFMyG6IrIrsSQiBAAAAOUGfQkUVLJ+WYb1tdpfdVu0F1k4kq5RhUrdGhAFDHYPp01TbmPpwo7SUbGx5vDZ6pujdLsJlfZM4VAAAAB4Bn2F0RH92+75YCJSoGPEHs1BaPe1nKVFLx3Wq4NAAAAA2AZ9jakR/19HKwpKb3hDHCpqEpgwPiuiuAowHjiceMfYzoQjP7bagp7T/MwoVg8eb8n/Krh+xAAAAREGbaEmoQWyZTAjf++s+OEuc2vQLdK/i8ClDkp0+Xtapl/Upi4Dlsp5sQ9tLIcelXy8eqVNGGFIhyPtEhgrjW/Ytb9bgAAAAMEGfhkUVLJ+A5qFYNR5HWQmzD29FNpojTcXVb5py+ZttjK0a0ojZ+6HbGv+AY1mIwQAAACUBn6V0RH+DDJXX+aBtjXyInrteZCVnwVE2lOmdE1eqH6zby3RAAAAAIQGfp2pEf3+cgskCrQfAkzm4yXvXEhlAALbstuTsEqLFPQAAAH9Bm6lJqEFsmUwI3/vgOOKjnj2kz6wE/ERkEFrp9mvc/FnvGTzuEvBxPM3+7YFlsfMrbeAOREsmq2wIOZ2exoKE6AbcyGVU2ti5k50trMZUI2gNhJRzN+j4GQsPgGo8hxfd6F/BV5QxpjTfZWDWEBGtlz/y9p06ppT6Q14s6X6BAAAAaEGbyknhClJlMCN/++A44rmWNeEapfe5ojbjO3gLbIo9IFs/sem9m8C9hrzgTXFpAoFFoez7AxD/afkFQ9GAOeRsTI8AKIit2NJwfsgAJBBdBVBQwXTcfUZuy+cNB/YDXD9y1MyNM/ElAAAAS0Gb7EnhDomUwU0TG//74ByFSpwKu/yY/TR+8rXmM/1ctJUphchJ0QahBg9/drvLVbB1uEUK1tv6c0ET7S2VOcyfJsGdD03uaOxggQAAACQBngtqRH93OX3cYVSoBNddIFIHoxaPZCAQcGLR8c5RXh9QJpAAAABcQZoQSeEPJlMCN//76z44S5zbHmusf8GcmWpiDjGPSa9PauHTPUthM4Tq2k7DqWOl1BucpTi8VQ6775shA/8Z+1n8HnIUpmW/CGfGL/3h8FtqksT/NymcW9deRJAAAABNQZ4uRRE8n3wHW3NeA+wKIZuyYfA4eD0aAfsQ5xdegmzrUzJAr6D2LS/mkx6XUnbpQmRhSBMw8zIWgJzltDGt3NdjhobjGWShq0kKsGEAAAA8AZ5NdER/im4U3cskkOWaidPqqW535NCXBHUD07pfk065jmmezMmm4OngZila6PBDQ8SA7tNzFMqQHG+AAAAAKQGeT2pEf4TUKvcv6iTIVCcF6BiozlPa/0aj/QBDSf946iRfptC1PLluAAAAXkGaUUmoQWiZTAjf++5XDhLnNr5FeZsxlZSRTEFl1hqtFA4oqhfVWwHju6NzP7Gzpz+tHjEfnDW+Oh5YrQeKavE+wrgFqrG9jsjA3/ox2THTdhWkSqy4aSCdsW/k608AAABAQZpySeEKUmUwI3/74ByE2ccN61JFkvALjEXs6KRRlHtA05Cr+i+xeGCVyVXxdoLYm1mfGZR+dej8GUfR5noE2QAAAGxBmpNJ4Q6JlMCN//vrPjhB4PuVF9LGgPQZObH+9xBHiekie60xU2mBEcs05OcSNiww67eP7NZPCqU7vyMShW9O7t9+XYJYXFlwEsy9Uhd8xicrYd7APjwDqdkyyfwfbzlstFm80HL6Q1TQ9IEAAABvQZq2SeEPJlMCN//74DjhLnNsXLgVjQBiBDF31YkR99oQqK3IUPQox0FL5Ngyqz0vr3e7zX4D7T/IzH37YlccmItA415mW5YgU9nDtlwhJalm+PnovRivR9bYdL21zzDZV93jsLfIFoLu62yIgI+HAAAAOkGe1EURPN99TaQSGVFDP8erdHfbNvCl7pmTW1uunWoHKJh+8IYvX3j4p8BO0arrZ28AEjtFV1nfwbAAAAA1AZ71akR/dzl9hMKpUAmuulsWOvCbvqAh1l1Ee3e0VIjtX+Pd6HNzYEzwaHU95p93feOAOeEAAABQQZr5SahBaJlMCI8YVKCqLqWXOXD3Rs2/i3L4bZCiXxdmPoXkYSJoYDVtTiH8gKixVFCbBh7U/5bLigBo+Lh/IlcxdgTekBHV7sq/sMm3KbwAAABJQZ8XRREsR4puFN3JW0ydhYi4T8lJqNEyweZMIfVIyGncYAQYhG3u1bYLctfKK0mJXCFln0Vl2u1VXLGsj4tNzvwBhCF9HkDygAAAADgBnzhqRH+BB5KrHlRVwgrrBwaKJDrhThiLTrF6/6zM+QOaa7ng74x86Si4GnXs5RovnBY3UKlHgQAABX9liIIACf++DLeYgN9XYj5zg5B5QJuOQipq5DurWEjnW3UorRPL+Gmp7W/eIApq4hA6GO7MXHIyDHsVucz/aMFYtKTDM8z25ZLQfiNZJqSiDN9723ESzhz79Cwt5M+8vOWuuGhG+Tvq9foXjrYoVXiVDwEsVfUmlk5TXbEqyYI+MZpo4mDJdCI6FifaFbOBypE2MZf5XWRt96wDzE+1vc15C8Q37MxkeCbaMkDN/t8mxkCWv8440Rlwk+ofcevu9fMi+Y2lGAjdq67Z7L2nX4tpcpo31hZ+kf6KZeyBIu+wE7zfBEWJkrDn+nvizONx5DVOkgfozhAUOvJbujzhGRN82gJwvYpRNAkkauwRrAkj3H3NX7/9FmB7s42z43E2AQejmKa6top87jsooGgL2yW/vkquH26p1/ZTn8L/bLpPfHN4mI2HsSxn8Xasz+QZbRKng1fvTxkt/mpfmK/kCyCd6bnIW6xeE5iW4Je98X1bSpyRG4ZEwQQRvmmA6Yna0ada8GP1pcw4M/JCmmV3O88PrasibHCahaFYs9akqjV5DdyOsA7IxWcVi1KwMrw65H0WqhITj2I/0392KBRdDeMYXdHrGwNEJhUIlQHIAJ5ZcKZWWU9TtKJlEqVRqCAyqMi4WjV04NQ1QEfwgz+T+hEehHkIC/Uc5dPLxaRa5f0rBpWg/CaNuFXLWTTpizOXU12VDPjD0PueUI5gcFxjM7UZKqEXIER8ZxBIM2uWna3B6TinIB//JF7wnKWG3KqQAoRJfs5WvYaG3M+PsmcpsEDOU5YG0Kdgf0PU4QKwIo7CjVGXnuOCeJP65DynJB1TaFN/Ec+KGgl8BEvS9O6qZ9x0NRWtKyaKyidn+wHMOsyvRqTPd0T82Zhih55cQsPNDIk78tfmfJbq26dnj6niV7W4C0RM63t1oa/61b6ejC3GPDzgPht8QbEWOmMkGWQaqDE1JlXaV0FU3nPF+k2FhLclRGSF12k3fSd0B8xrBHuCYayWiuH8ryJrtKZxJ5MQghDgBuHJO6SC1EDjEfxZGD/r6VzWAzizHGg8AmGx2xLAaw9rsvxLPUF+/7yYJueeBW5R93rr0FqleLihl5Xj2Aq5SkobABmcpWGhmxzvsHrBwKqxvkuHcaxY3VOBxUoDZhnQfEv+mZtl/1NlDKLlxwKt0ekKf+H85FMDoy087GajtU36Yx+6OhUz9cVW+TW34lW5zJzNRP+Hc8uhJDdo2lXVyYYbqCDCYtOyMVhQIpdZz2jF3WMV035t+ZxFpEnL8ijjN9Jg8JcuWyjfxs5sP97Y4nlx4+ag24HnZ+axqomYDumc0gNVyEUz6FZ9JbgAuwpXoA5ORKsBxrFSQHm0NhO5MryD+P4NGBrI8OniGQEltNCWxKO4CaPEW/ayYPnua20xl5iBviCOZOxGovBGA5TXYq/DBpZKzwKVYDEuxwEmec9Nh+C+oW4OqD94kO3pzQXD4WNt/q7X/PFLl+l9H8cdZ55u+BvzP0FXMoLDoRw7LJfWY5c+wzINL2XuOAYkfJCpNNcv1s6D4asebxSD1Oa3fRtN5j1N6T3U1HHUr8DPuALIo5ZN8ejEA5QFEH9MczGJJOxYHf2Ea6JmC49U20+9S5lgbbfIDilnnqaUN4P9cfr0GozChKSlAt+3fH1bY8x71jdkn9w2hQSZUYbU2pB9vw4KlfdleH+KnBRal9A8sdEAByt88ulnusl0zuMQqjLVJjZRTm+bVZlUB1A17g7C1j51Ik9g4dNmA5mrD6GgZNkTb5nGM3aTVSZBRdQchSdPGxKj9iMzKtky9yCgMnQWxUthLxXktBy3hagbWow9t7gvPd5DjMYYCoMNWUQ53ExMXz5i8JP+AuTFV+WVNG8AAACJQZoibEb/++1xYQAAcZSln5nDonHJ2FY8IUkCHV3Sy+GFqZcQ6LZxLbyQ6zhy/gRLW0pWGMhKeooRyMTxiJ2j2Rfe55qa3amF1vmYBco00/ApFnQ2DT4KOcgnT4ciTok2qz2OCLOIFJiYrYM4viFJA9fmIGkJAUIM20h/dudGJ62seoc2Z9AEofEAAAAnAZ5BeRH/gQeWpd9/pOade/rN+isrvDcxP1hTmlgtc64uAfKYMixRAAAAZ0GaQzwhkymEb/vEgWa+okZYoHAhrsRiqhh///5/poi1W1PBAeIgbreMG6vERAWN1PzvCg5/c1bRzTEU+q0vqY4p4w0N29AZHonm2TuYWlH/CYj3AyJ2CEW+WckK1sxFQX9Vkp3Do98AAABmQZpkSeEPJlMCN//74DjhAvudrQV/Von+LqXSU5EnN5Ici21RizZEycc6yRI7HW2mkSAr7OhvSjQAWm9VvWhylF3l+5QdacIDtyQ0c0VNiZE8gmdnua+vNHpLu+/zui78ynwbUg5fAAAARkGahUnhDyZTAjf/++AchVbGeq15oxbCanNPtandCKIk5UFi2McwnOkyw6FFIwJZCyx9BQ6ETJstOW6MqH0KmE3QQ9ScPBAAAABQQZqmSeEPJlMCN//74ByFT6UQhh76pV1tl2CE897btCzIgaESnlCrkP/rn7NokgPeW5v5Tlj6RqOhWJLwnxuEYh/nbkWklEpVsp02UeI5hcEAAABmQZrHSeEPJlMCN//76z44S52OwRnQ4IPpaB5d9PcQ3xV3g8rMCFmB5bsXTfJyvJ25sgvf46GpdxXXkPPndvnlrxgCG3+15DsD6MgtyhVrvxbNr6GQUtl2LeFZjaZQES5OWM0rLLlBAAAAVEGa6EnhDyZTAjf/++s+OK02js7X6gKTlrPgSr7UwuS/819DrDrfrMwhlkZEIEQkkAjuhvWd6ibuHYrJ9ufrJG3PFqSfmXeftKX8TepwZpbdTG4a1gAAAFxBmwlJ4Q8mUwI3//vgLpU+4B85rBOb//EpYgZFZ6eevNZXP+wPe+bIMaDyfbTJHzD4Kxd9xRs1m454YoGfxrw/BKC4vlUK6tcXIpJFyYby74nDDHCTcY26K03YRgAAAGZBmytJ4Q8mUwURPG/74ESpX+uSC7HOS2CPeetPFAoB1Gux145k4OcRwOiiCVIhzw8zcArn6HvJ2afHpwjimmjCVOlWLWAXs6/dS2mQzr5Tg+USVx+WhFO8PWOsAE5kFw2v8j/ac4EAAAA3AZ9KakR/iffqqvm8Jc9gUquRMvf9rGuOKZnd0efesGD02uCoFuLecT1P7NzuAbHBt4lonQco4AAAADNBm0xJ4Q8mUwI3//vgHITeeM3FIGYUnOaFjLcsMmQoftZNP8Bb606uLDf3eMjrSo7Sf4EAAABQQZttSeEPJlMCN//7w/CbbaisfdZegDj0M9TyOi5+HYmOurwQpxoiBko95zRoR8EUd3PPFRCKdewA0ZC3VL2iuZhfJK/KwG4Yg0/9XrmYb5AAAABtQZuOSeEPJlMCN//74DjiuZY2PsMDIvgZTyBAM69A3+mzj1NB9ReQMa5lpIjzkgpFb30Nv8USZUa4L6VDLhd3GOAcEU6/2m71HhKIVVQxQ1IYgBzydJ7TaPB/Lp71eKZ3cme7OXKbCXDl4TYBQQAAAFFBm69J4Q8mUwI3//vgRKv7oyEofCjyt3vsnfsah2kHIdELq19J2RoKoC1XL0zZTv6sAPxkSDJ4hua2mqjbC4lGiBfWgM8EmDQTaVR78G9JBOEAAAA5QZvQSeEPJlMCN//74ByjvhR0RQiScU1Fdfdc2SHU4FOuYJlXYGqlgNrIdNY2a6p25gOjejiWvqYHAAAAZEGb9EnhDyZTAjf/++A44rmUqQo2o78x9zImFz6ycgjd223QsrF0bQ9oStMycJKa04q4jD45hvqUv+ExtcHS0/sB5B/uzSniIvGiDyArKszGSLTNmQx4zT4DAS2BkYC1Ww5Py+EAAABFQZ4SRRE8n4DqJeP15Rd2F1lVcxp2jzwtUtQ9ScvCBuIZmxQraJl3DozW8juXqxWghEa2yC/V1oAhaIA2kP/2/bNYpFpmAAAANgGeMXREf4puFOCtReQubvl/VjVVRARk0FI0mb4PyIwlMKV46WFsEaeKjqMI1+R2+DtK1F/8vgAAAC8BnjNqRH+E1Cri5ko22NzWPuDSPi1danTYsunQ9icYtQPRSlq6m6AnOlFA02op4QAAAJRBmjdJqEFomUwI3/vrPjidIX8jzxcO5M8WgNUVjZvTZqttR/yezTvM/BnIL5HWxcNwKvfguadDkLdYtXjtonK3zpSm1H7UXsq/0YHnzRgNDAYnqnNKKa3EfIz15EOoRIhWAiTMZDR7m4p4cwv3/77a/ATHMyY8Lm+4XtVBteQnl4516EwPOY7Hx9ZJsncgz7c28+7ZAAAANUGeVUURLN9/4NxQllWwS6YjGgAHhfIlm6iYyMZ/pi9Hb6DSetZgTUH6s+U8uHcuB/OfkzdgAAAAHwGedmpEf4SbRPB1qbWDavfsuHn8SwfaNy3M5M/M8YEAAAB4QZp6SahBbJlMCN/74Djh7MxwQMXA4Xa1gvMmJfG5IpGYMILlDBGG9SiDcf+miA2uUriMuCZPDw2oL49uTVSBzN1sGzND8Fw8Hhrs/70fxGPxNhJdvqFB0w+DC+aj849W7UJcjUb7VkYGXQv6/CDjihFxZP7ZmzrBAAAAUUGemEUVLN99Ta+lvuryT/+7Fb855I935eQ/NKyrIpEljglVrbW4VSwnmC7RuPcxJO4HyTlFvPN0kQgJ7FjETsgCU0+IAsrdcNinbd4czNuDZwAAADABnrlqRH93OX2EwrJavM7/2lXT+mo4pwFqeXYue0BFgaMHvinTIiRN9nQe6i5heOMAAACVQZq+SahBbJlMCN/76z44rmUqw207hd7OqeG95zjobXO1BXc8GYszEgUi91ZFitoujcqtv53F3MqCQuvkILWF7cYeyLidHTm+UXkyfQp7XBSq9+bk8rzweySNLOjwjQrKhpbObcN64LF//YEwizj60QETaAs6xDcg6lk3U1uH0HDFd4ifBkn3nk+cTzIVsSPERF68h4kAAABhQZ7cRRUsn3xlVUbBHfvOCCo61ctk1vvWUTe2GOke7+ZMC+0wb8hMaFSQHz69cBeyCx6e99M+HncZK5sJxhMcQG3iUMQJQG5qp90WppVxek0grIUNEg+nMZ/OSgHVaYGxCQAAAEgBnvt0RH+KbhTdyyn2mc7b0qS5IBg49T6i1mEqW9SGK2rLIsvaSKB6hosMeUz8+LtQpC14bO5miLujSI3IhMrcLsnaQe1F9UAAAAA4AZ79akR/im4VKaciB3zfJVSkhw/Z8IUNrw+u/pOSGKeH2tujCA/SRAiTtkOg1ajX6LWUHYFkQrIAAAC0QZriSahBbJlMCN/74D1FQ4s3n7j5n5Nuu3X//9TPET7aPn8s6EoF8yfvc/MfvxyY4pGK+0qv1Wm/xldOtP5Q+S6fps0S+m/UZnQ7EwPsBycEwP1ZVg3aRa1YrfiIsrKz8S8liNH46Caw2vDUSR+GdL+B35GCtu4RO0FhgZFHJ26mYluh/0OmxLZskzny1sgm0Ygl0PCSjypQA0smNK9TeZ6vZ59Qihh8Ae3AfAr1QlZ6NWXTAAAAPEGfAEUVLJ98ZVVIdZdMwtzfMxEdTpKGRFWgnHy+yhWW0NeLWoaMzR3O/Z2SNXpaNidD21Q02s62yRuP9QAAACYBnz90RH9/iXZSP9h8cdeKVvwRfd8hZRQQgN0TPngGS7Fwab7aQQAAAEQBnyFqRH+BEcHSusAq258vWb8w5iZsvh+pHv8c6GNrxfq4gOcru6UctZ3xaies0Y6hBXGyrHqsGGLfpOkCH5M1jDKmJQAAAGdBmyRJqEFsmUwUTG/74ESpBkuxG2KSXcpc2jk8koxEI5uEi5yoS6e9xpdAt72c8MSnXbbConlHP1BzUnJ/+3z1Aoy5efyGWrhfPczqcoqfv691EvI+hHCd2LcBKTSXuxUnnL/0bEnkAAAAJAGfQ2pEf3c5fdxhVKgE11tH5m8PYCgh4Qjt0i/xdPUGMftsgQAAAINBm0hJ4QpSZTAjf/vgPUSWjRx66POWg+1VeyDdo17n1pRGCAZv0EnMYI22P5nBPtXTAZp6Xq/Jn7fkYYL6sJtWyYB/PJ+yuGPOHKNGaRXXtSKv8ZtxjnL8aUG9NP/FDr8upgFotXBQY8KM92B40pAj4Gj1imdJHCvfbU0iETKVBhp1wAAAAFpBn2ZFNEyfgOol4/XlF6ea6k7mZM3z9p4e8X2Jj46Hpmlke/qavrQOjTwS9Z47MV26uYG6NPmmXkDbOOyLBafQL9eNMNmRUBMDuSM10MkN86FV+v+G2vJvgoEAAAA6AZ+FdER/gwyVRAsFrTv5yQfi12mftmfw7u/eMLjWn2w/45OdHFn7ezFA89/eGqdczYXgEIo3louwgQAAAC0Bn4dqRH+E1CsRGdq/m1pty/GwHijW++zyvH/Gdpk3chKR+bVbMoyVnPmV0owAAABrQZuKSahBaJlMFPG/++5XDiuZY2N6PL+cwxD3ivJ/8tkqosMlZ2kuhgf0q00kNIyHwl6OQiCcnPCaq6e7v2Qvnnewqy5mzqYUh0ZzBmS09jJQtDXRECUYUTYcu7s4jMLzcEOBfNCz9U+yd4AAAAA1AZ+pakR/f5yCy9StiwNVmr5zLjqXbDKcDqtWtTICUwtkcR5hRgZuV/4Inokvtcj6Vro4MnkAAAB3QZuuSeEKUmUwI3/76z44S52Ok28EOzD5G2y58sz7giS1bNt4y1aleBVYPkGtESwSGS22z3kj7PNRXWIQBXSfnJ9k+hqAla+toGcJ/TJgevZWzmQsCjk3toczbqwsz3LPlj/Q7tpVb/p36b3WNDT2urx5wyQ2u8AAAABDQZ/MRTRMn3iZR4RsYALI0qaMPz15tf/PCJjXSCVLzyIjWEB6NIVZuFQd9IHJUmnlM+pGzTuq0jqVN7pbBnNO/83bjQAAADkBn+t0RH+J9+qrQkFHiucZngCfYVC6NrC2q1CVtl2NFIiI9/8efIy7v0ffzAkuTuoM22cIjGJ9hWAAAAAwAZ/takR/dzl9hMKyWrjL/2mAOYr5XPeIQqNEnZH84kJ2zn5AugiKzfdzkWMm0Z7BAAAAf0Gb8kmoQWiZTAjf++s+OCFJBZyyg1Y/s//daL2cmKlc5TOgqXAetAsmctqz5P823VwTaTcwPX+/sUmF5uTMjeRSbn+NbF5LxJ6JtQsih24gjtIz0ieWd7148vl0dKSJkfFHV8qbphVazHMzamX2zJ5KbX5iaHfn+RVsiaeJeH0AAABZQZ4QRREsn4DqJeP2RDmwj76BdzDR/efXogJAcIU8SeSQHCj0dqnjrDVpTgwQJI31BTaAG6INqkujzWTGU1zd1DjtXiaRFbASYt3GqCB6jgT8nxWzUXikMYEAAABKAZ4vdER/im4U1zc+vVyb5fhhMvf/N+sn8Zk5BtnQgFm5gyE+ybdRK46vySho5BSu7O5YawlwbRtBqouRndC7tT/GyEW3yt5IFoEAAABPAZ4xakR/hNQq3d0NOvTMPHejrhYTBJBnYWLUeub3UrmhxZK/nEFqJcD7/UnWAqDmNwXcWBRsIOSn8sxa0LKyMIigPgimkAATId2ZpEAtIgAAAMZBmjVJqEFsmUwI3/vrhHZDqhGUnyzSHEzswaZ2RxqKGRynYBXMBuEuEFfHmUDj9ux0sbihOOVmVeYWgA9LOoKwJ+rMjiSQ6ZUAyFQsCjjaDAUTNJtjWai9Rc+L83o1bWP1S2yuQUsJLNLVuFrlgQAymRGcNkyxEtGUkanf24vl9rYEtUPRoVcJrHcCRxamMNaxH5ea8jkqFJL3PStn/7XHDDku9uliPPRnGOvdudC3OTfeJhZuBI5KZb1w9wEivL8R2sp5CIAAAAA5QZ5TRRUs34GrPAdCUxfCqZZ2unlDcBy8LcrzCPU4ls+ywf4Ue1Eh5IMuWUEO/crgik+To25i+g7hAAAAKAGedGpEf4RorykjSfm5OcGVQs/wa5q5I3WVt0A7JQcZioy+6Cpd3cAAAADPQZp5SahBbJlMCN/74DjhLg4zcYH6PAnbkTY8iUAl5ul5AGL6HlLmdMJTDc6cLh2fLVFs26FJQNh8DpEVTfPwRt/7mERhKyfDeSIWFU5yQMCYNHv59nvZFp/a8WegCK4HUMfdTe9BBOJ+XPTTUutM2EKDdwEAeJMZxuSKyplgd9tOuySGUaYtDrIicCWJFC5wXwmU8vxDFfE04OEohATSR2A5oLuSDW3Mvf1CjoKTHqw+jE9AVMQbxRI4vzb1WUJejTxdtN9Si/S5yQWiKU8cAAAAUUGel0UVLJ94oE82Csl+UbsHxzVXEoxsMq4cxym8W9nAZxQ5JIcYizkCGXhABajDTcmqAIpfxDxLPxOswTcGYTAAIZm4AfUvHOJyB+CdHlrCgQAAAD0BnrZ0RH9/pbSDSQoUYXbJ3rSzc+kRtLUqkKeKWWa+GfiCNMxON4cEELkIjdppVpoCwBlXx1ZSOo1lcke3AAAANgGeuGpEf4R1fILWoVH4wTcI9HBEHVSvB4bqbjhbmfh2jcyxVPpzF2b1jR4X1zi6YlYACKdbQQAAAGBBmr1JqEFsmUwI3/wDwzHCe2QWuq6WGF4P0C2Zsdw46Bt6zA57Rc2TiTbjW4H7HgKfR8Gqh82t98tMHubaH8epd0G9TvymJA0ATEGErg/d+4p+TcjdLCZHjYjKV6SX5pcAAABNQZ7bRRUsn3xW3JFluUBopv+0Jh109W67tBai2wVCw1/Fli28auwIi5iiXJElON5pRsqKGOlx/dZECZCgJwmVdX4qfM5L7ofdjxymz8IAAABAAZ76dER/gGwpqdu2nr4gydT/It9EKEbgleYAPE+8ztHAa3x9WnkO6FBi5Veh1TP1eDEtNkuvOywKMxIz1SR2hwAAADsBnvxqRH+J9+qq+bxiZHCFr0gnhkOz+ad8z4gs78ck69ZSTjmsIwQTbotwcaKIZaLcxmKWlHaf9l4JxAAAAKBBmuFJqEFsmUwI3/vEhS51ZOAPrFCl1ornx2QnJRuYe9kkqMhxGOaMMLj/bNXs6VgpAhRFmLwVR09/A7837HfA1aN6O9OTeIhLPR21+RX9CY66WwmoIi5/sqTumt7qln29/+p0OIR7TaPCsms+zm3oiZ01cw+QI1KHDn7eVyJ6j4XcPg6S1r5fXKbly75nNvXLiFj5ogktvCp1SHaqEda8AAAATkGfH0UVLJ94oEqpOpYRxdsQCq36e1PDI41WDpJPcKamY+k0pWs2/KywiDKYOYI/MrTcoCQxBFm+MjLqCzLeTpaRiENFpLsVzukTIcUPCwAAACsBnz50RH+EVDZPSGusQ1m7R8O6wN1peQ/1jS1QCsK/+3Um/nowdX2k06ZAAAAAOAGfIGpEf3/ocg0UgjqvtSdJTqly3l/5RQ31A5IZqdimmnTC5MIJTLPa6Xa8Tb9vHN03oEBvLxevAAABDUGbJUmoQWyZTAjfx0V6DnpfBXTYnQk2bwKX2bwMtCO06TgUqyoBuPFXgEpXFMPAzWMC08ZZ56ZhxoN9TkSIiXVTokTYYuz3pKV7lAe/FQPey7N+dNtFiUjQWvY15ek+//wU0OU7OCIKAvTzIaRU2/1uvJ32Jwtf6MeG6udoHVMDrga2SELkoyovsSaeoA8+NLOPdIlZ9IkteJ1UZv0f5rSTvZTQ5nm5XNAnjE/38OYo1ggnhiEji+Cgk5Y2BrVp0F8ZOcAg3273NFXpfm3+x8FlaagyvTugb1BaDOo9dxr8EQG1u2ACk3eEl2m7Cw/3wccsdqmv1k5rlStFNyS7rxGwIGSghx6CdSt/U31jAAAAT0GfQ0UVLJ/Twd08p/VyChpS8T9IWptOSDvdLPszfbix+67WAXfjDNjm/bHE91nO3g6zlnZrZDZR0Gz1rsuN02ikA6mQRG4q5w3HL2ZCEhwAAAAkAZ9idER/hHi2UFyJjt4ot2IIOgG+Vq+0x5nM3Yla/iCdTgZJAAAANwGfZGpEf9fRGvKk48/Ke4CgTzeekpncWwo56eW/sFEV91i3WBPnz0ynsff3Tnr+fRiCryeLNz8AAAD4QZtnSahBbJlMFExv9Cvq/twdBueYn3Ee3pjczv7OR6i3WU+SGv/yONzlAdl0TMgBYUraQumePMz3mddrVgKe39/UH0dZMWofZ9Mqui4P/wvLxsUppGHrrazsvTR+f2rODWARjtdkCSvFbtfvMMzeLFsQ13oDAGG2vvjL5m95QTZnRKv1OnYK93SNtLC5O5AQU0iYLnxfpUikyvxOvC8fzzbFhAqGYgHL1GBJ3KgSybvxeNLFYx0n3LIy1+V+pUzTJ/EUuTURErNUiE759HI2TbRLFPHYcqtlrP1/TS0NNfBRxkAXAcx2r7lawk5xdWWia+VA11/z9cEAAAA2AZ+GakR/mJ002N+vm8JdEQuzLaLtUIOIHfz7FHdVElv3dJhMuXpT4IkFixsgNNJKwy7EeXqbAAAA4EGbi0nhClJlMCN/x0W1i/GN6GepRrlUQYVWJ5EJgT8Zlnuh2q0V4+3OOPM4/4/MCjdcBbKHis9PBt/yefxMxHgw7FScKctBHvFoR1EawxWUsaHfzfSnnalnY4m9Ak14jYSJiHQ5ToCi7Y4bqJeCGTsIs3t/Q3xTdD//Dx8soX29zf1N5xgAr/qPmK+dU6/y5sivZgEvOoh/r0LHDfok77DJX4jZRZR016pPj/OKEWHWFz2SRnFT4/7nXgVNFiCLYupDDD8vVEru70qqarHLzZrSjSEw3CIA0xIhUFDHlk5sAAAAMkGfqUU0TJ/TtTGaqT3w3TTw07mec9pdUO+qick9Inm7YO6odtgQD/RhYzoYrtniSB/IAAAAEwGfyHREf4TAtHqZwf6daDH8ihsAAAAwAZ/KakR/zdHy96S6V8otDkejDD6GAoE1dBa3yj/ms4XveSTIAKMwyhRIvm1XthnnAAAAfkGbz0moQWiZTAjf9BLVe/EEalTabFqMzxVTAPfnxJ/Zlk5dL1vhj03RqKWw3suOlU5fZYsZNbQTFyM5VPD5D34ws73lho1n4IFelgz9POTbop+0V7yvkVIdPx9z1ZXKiB6WJVdOUzkbyhc1ThMFffRcKHS2W7razGQSoY0BSQAAAFhBn+1FESyfgOahSeu9LaQmIfCcRMhY/HlRnhZaQPy6CAtyQ+LszE9EabEU3RQtsAcY3sdxq8ilGwwrV6BXotZB+DgqfBZjqOBne2bdeY+xHF3ipKcm5shAAAAAMgGeDHREf3p9slhDFwgFsyLLeY0fCtie4WUMzC6rVyKWDS2ciTpdi3cETt3ZfbuvDXChAAAAKwGeDmpEf4n36qp82ESoc+OmMW76ljVM0FgrkSvJn6xJsKLGiU64smxwxzEAAABzQZoTSahBbJlMCN/0EtUjX4qLA/VxpTkfRuOUpW531Ao8SwkqNXlLmRBGxN1IsecXQYmRrqrhzuWJINDZbuHe7a3EXT2PLlFrMUbyi7HcW331ihx5HV0pCo1UdCwPmKqK6tJGWYXzqyTyWsq03jrcYpi3wQAAADRBnjFFFSyfex1KrddINRpEPynlrUA1hgUxfwc/hqlijWK0BwD23AC6+DLASM32cUAM8GGBAAAANQGeUHREf4ByGO+by1cZ8ELLxtGybJtzqprjV2BYPrAXGxgEvx6Yxh6AyzUdErjXJBHKQFDAAAAAIwGeUmpEf4RorykjhEu7Dyr89cWrOYZIKdrrF5QUKCEwPBeAAAAAckGaVUmoQWyZTBRMb/vgOOEudjsI3Btv/VSW39BLCWXy94pOkzAzXgbOIzUOMjzvnTjsYbNLA9pAvZ9YoFxYYhx9h4eU9NZk38/IlYerTZE/6bkDdbgNWkejxbItyYDZBfXrFL08jsMcNfK3iihUtFthcQAAADUBnnRqRH9/6HINFiMdEuSVNpUqBczjZC+frF9/+VOzdYvwD3y/RUO7jtNMvopLiPWtfsfkhQAAAFZBmnZJ4QpSZTAjf/vgRKik2M+QxjuXx6xjPBvXPi+CYy/3jhDMFE5SsMyyTR0PdQ+GZ5DvFQ+4McBlrhdYQyRJoj/S1jsCJTCCYb8261AmPc/xNGhEYwAAAH1BmphJ4Q6JlMFNExv/++5XDhLmEjtL29ZAM9do3FyhDOoEv9H/h3RF9bsYANZinaVx61rvr6dcFSPOYKYVa0Ky3w6DZJzzBXVJRPeXQdWDKgQpImvJu97y8iZ+2lh5SMVTWZyUhJjoYcacZWmU6rIRgEt/EHai+6H6pE6egQAAADwBnrdqRH+KbhTgtKWIdHh3CgwGor5ORfdsM8dLegC9vy9X/NkeqmFpzE3INqVqG2xva8onEHyYu3G7LjMAAABIQZq8SeEPJlMCN//76z44IUjPWuSuXQG9n5T17mOUqH+IrXRyER2829XfIa9IKsVCuw/i9Uft+dhAIOotXMZgMVH6ntznsvIRAAAAQEGe2kURPJ+A6aXZxXJX+dIXeGpcopCzqg9ywRig3tyiB2DWkAg1Nn19JHMqupfHrhRdStTdinM2zM8a/E1bWFQAAAAvAZ75dER/hRLKKjdRBBNZ2S51o0KrXqTH62/XmFWxCqMc5WhSfqLQHzwmuHVfuIAAAAAbAZ77akR/f5yCyQJ9t0mSfiBytXO6KNT8CYfxAAAAgEGa/UmoQWiZTAjf+8PbyScRYg1RKyMpfbNX/CZRmm7Hq5VPl9Kuk8m7UOZdAbbrLswPQZQjheXgSbH4nLjepdplv51yrDISiOwhOpa5KSaf4Oh7414o8y9zXZHdmfzx+QS1lrmt24zuAg09D9q0y9iOg2/VIJnCofp+Dcv/U7ggAAAAcEGbHknhClJlMCN/++A44S52J3BJbw5nJLJfrjdwX6x1tDYZqj16z7INfCLPFEvnVSYK+EgFy5iJXZMiPitdWn9FpnmJqvFHhXzj4TTJ0U4uxYajaipJlltSfuzvOgecTBXwsPWX+UIqDZy0B/uWYEAAAABsQZsiSeEOiZTAjf/8In/MgnTsdVF4hM+h4LwTtR11DWQhRcJNRvQrtL/nWVC8Fu9eFlLAp8tLIZNEfkIXRbXsvYnQqfL8udnCweG1H8XL1oX20j7WR86CHE+9Ke8Jhtdy1bbfdOfKjqHqgB/RAAAAM0GfQEURPJ97HUrykeGrxIFqUtgUr7SYug/yEk3MgZ5o4qOnfx8yJ6ewAL36Oxl32rzUXAAAACEBn390RH92+75YCJktXmd/7NQWvwATKbauC51Qh/fxn8EAAABWAZ9hakR/im4U3dpeEMcKmoRT6qqm2/JijnAse/V8GPlbeJNnpdMFOlnqWvqie1+Z5e2AUboKtScIhZJGToOZ4DYrujj6Kul4w2TJ9Hh7ENnXsHCbwRkAAABWQZtmSahBaJlMCN/76z44EJY/u4nS+fFAtqU8yWqvWTE8GR+nbQ2pXN+qpaQkJCk9IKk89H1bE//13lfA/iJjEiHwWexUi5OP/Wk+Wl82rT8IeKMB/agAAABKQZ+ERREsn4DmoVg1HkW87stKnuJEPWC6LUKIZAqca58O77RVjao/vMVW8zN5+SBSMrUHBYqA3TCqOVCAQjR32PKFDGNqCbkULzsAAAAxAZ+jdER/gw4tmzW7rzcxfEkNyf/onzmAWVc3vPO+lyds9aaITrQ3fCOV9PNwrTxg0QAAACkBn6VqRH9/nILJAq0V22CwA26I7F68ZhwvvzudCOtW9UTJodNhQc7hgQAAAI9Bm6pJqEFsmUwI3/vEgWa+oh2wHUBNFWOi0kAxeW/oBNECdxapz6X4Ulak31gejD9icbk8DhRm0Hs2hpnHJ2o0E9aOoy4Z+ZKHeuQbkVZ3BAG80uP6pot9Q7nbRo/Sr6mvqnH8mE8zpvT0kWpYJYh5VOQG6NR3X0yMAeVL+tP9aFjzEgYaAOOYI13h/D/pQQAAAFlBn8hFFSyfeJmV5SQCgqT7iG5bhfyYEe/eRaEbB7n8Zt1Ka0FoANimziTzHCgJBfnDtij3/czGvfJsLwpDb7MAjV+SrguqOiQF3BNZwA/WhjksDFwtQJpmoAAAADcBn+d0RH+DDmAasguYEJlS+qlA/VHzkU9DE5VIEBd8g6r1fmm3k/DZypjK98baTmgzNHfZVmKFAAAAHwGf6WpEf3c5fYTCqVAJrrpdDQdDLHRuhKusF09cRKEAAABvQZvrSahBbJlMCN/74DjghLmmO2mUnbfWaW5B8Mb92g3kEwH8rIP/BLmQ9SIIzp0Z/9N2kUyz8oEJGs/Ap+sPATqqFFZLVvfg+PmnCmip/KLJjlyWyzFDoI36aRp8zoUQbAabU90Z2YMWlc5ECtZiAAAAVEGaDEnhClJlMCN/++5XVEO1lWujg8lVyQxzv7cDyj9Khzwt+GmZK7MW+odXd1sxR2cKlqoYmmu2Vz7BYmxOdJPgwJSdZxassfDNt9Cqk8Ou7kbHwQAAAGRBmi1J4Q6JlMCN//vgHII+SiBmB1jZ92ZKw687olAG63mvy4JF68C/BtdOHNqo9oOcgCDap8yVUyFaCvNPKBjx5Y12mHem5KjGuR1HE1vwiWnlm7LaNIt7dT5FKWs5/g4ZnkYzAAAAc0GaTknhDyZTAjf/++5XDghKblpByuukqCy+P01EbidExQWS9H7GxldpGPDpK7V+tBUBN0kUkqNhKVgVQcUsvuBYPIwMAl+a8Q4xkIhl/9reTev/yuiRaZox+SduCU+GYb7D5dve6JG2bETSCoHHzLVhoOEAAABGQZpvSeEPJlMCN//76zcgkiFnycdhsLEw0asP63rfGuqZvt04llNZeVygNUg8Jr3NY0ry9QaO7iXi3konyubaDZ6qMEZf4QAAADZBmpBJ4Q8mUwI3//vgHIJlMhOje7UYw2E7Ne9jyk9kHh3DzUwY2oalRKALKuqLHKNjTZEyfXEAAABdQZqxSeEPJlMCN//76z44RALE0DW553rg5eSpXu4lPlg3qaIlhzcDVr/CwZnQ+H1LWa2hx8yylsTU3OCdaOErX+9UvWtlkHhNfA7/4IxFk9dSDJPX2R0DfM6/cnV5AAAAWEGa0knhDyZTAjf/++A44S52Ox7QAqqAU+wlZbuVvsj7Y8xPJ/oIVLMj3KgFlKtIHmn1xsMKiZMfby2S2aO0xU7D3AaY+dFVaek1REHuf/uwXVRA/9Oom4AAAAA5QZrzSeEPJlMCN//74ESopM95J7cOXH+Lihd7kB1WsgMIaYuk4/P93ZFft876I1IMoEoIC02Nlj1AAAAAMkGbFEnhDyZTAi//+98ri6ylM95Xz+FrnToIvklvE5kKaL4Gyoqmk5OFcr+nIm3Aq8lFAAAAWUGbNUnhDyZTAi//++tE3K4Iv8XvdFgAw9nY/t8iJypWqoPs8SfKSv1NUm5Aud8Hq91WXCVfMYTS48S1kcJQTJb7gwxOICnLQAyqaV+gSJ/vuLhZm00gZbygAAAAV0GbVknhDyZTAi//++5Y+Bvd0edhu05xaJ4bHAypArfk42W21z0Hw9aIvP5GIJ7U1/Z42XP07srgXfJu1+InIWELVunkmOM0UPSk9U0pMQXyEMc8Ph6NIAAAAGJBm3pJ4Q8mUwIv//vuWLcrkUYnrfzKvY6YyeTO9Rk4qyfjY77IkI55iiBIrunpAkA7045O9WmOwt1TUqx7i+IPJ9eh/xnWjecTeoRQNm+CQoQo46ThJC/lHxBJ3Ma6WMJhgQAAAEVBn5hFETyfgEm1lLG500zAj+otoaEuWn4LIfRFhyaAA9CwtVZ6mp1AZb89j27+RZtQcexyfGBpeuTXTM9UC0Wa3f0yZmEAAAAlAZ+3dER/im4VKaciB4AJOzHjFym3oODn6PO6EfFl40Bc8IqeXQAAACUBn7lqRH93OX3LIRNDM9tasR3koAjrPHMOiuw7i9TDEgskYPw/AAAAiUGbvUmoQWiZTAifwAj3GUsedIbwtVPmC6UOtk0BDMV7Fztcurfnw+Ukl/88z4rxasXNQ9PFoO/5nE8qDt2fJK6LMxkaoX4RjOKmCD3wJkQpTcQNJfK33lRLTK724nIXgRJ5xYIwCOoOErOutudRLF4PZ/ViMvI+XwSlh7aUB0gxVLJooJL+QGvAAAAAL0Gf20URLN99VLWN00toRsi1JSCmlbiDfGAKxMwuqu+ejxaDbCFeBpddWwXK34eBAAAAOgGf/GpEf3+3wq4dDu7w/BsOHSau1FuQnVGdVhqVK3yiZQNkNPtat/J4UIGt5UWW8ZQo/jCP/r1Nr3AAAACHQZvhSahBbJlMCJ9uIMp0i44t0raDO9/qWe5SIVgIlo7RXxku4hPZq6uBGB4yOHnSn7I/dfcyQ7IQL8iwlPxXakCthkEjBcrQSK1+YKDZ1aeYZj1E6kswhNqi3qb+Xp8S04OXkru7TYYvfQY8IlY1dIAa1uyqVN5tJyc690K5tFol+v8sP3r8AAAAWUGeH0UVLJ+A5qFJ7VTSGZaxh09N4LvIJtQj3iBd9smy2aE0j4YJdcN+aL86nrecX6XJe9x9zozcD5zJ3+OWlY6VcS3j059iPKQYmFD66vdkRvd6PnTkL4nBAAAAMQGePnREf4R4tlBciUsnWu6mI8oSE9LvSXimcHaU1ew88XBcw2xze23nLBJw1I7DAmAAAABTAZ4gakR/gQeUyXbanD30/po8AnM4qaiU66bgB4fpv+VLJIx/KYTExAZQTmPLSKx1c76Lfl3wOXm3/JkWtAbgIRqsn6/a1Z7l65jvwwWZio2lQcEAAACtQZolSahBbJlMCf/mV66kDVGRg/gckK5KqDoGjG/Rc0jTpLhTmGpUH6p6+6Lg2rz4jmaD/nRZb7pbBze8wvCCpGUZHUhbf8B87j6IYVHLv2NWZ2HtpdsJpaK+aplUwQpQ2jyOwJc8FMrzBW66sb7D7eHQNNKX2TQDpjQOPBX2qpgf8m4FZujU9qlQGljaykfczBPPyrMOa6hPI30+YisxV90mdGge0/yWBpepB/EAAAA5QZ5DRRUsn3sdSq3KxWQxmQudZqrOmd/ZinGlzs6ATElOuuq0U0+75PEkS4CwMI3s7/bFIEfrMlMuAAAARQGeYnREf4ByGO+bxiuCCgeG77QV8cdQPoGcUaIvNKoixhEtWoRDrlt89Zpb+CvpUPEpQWLYG2sl7ltGjzmAbh7BeK7WEwAAADMBnmRqRH+EaK8pI0WvJO3oqC4LpNBdDqSQ/2H8yLxBYfp6xHMaIZZrOGpnEjIG9IZgYtcAAAB1QZpoSahBbJlMCX+Lwr9oP9EDpqimHn50oJzlNSRyBVnJc7Vs6Id8pb1K4EjPXOvWC39X4znnb5gA6pxzGbQIqrjG1MBa6dnTE5rY8vR7mZCBx+zMoAaGfd4sMr9WM2xq44dO8HqaOxPZ9945Sd/152ym2cFBAAAASkGehkUVLN99Ta+j2R91Ew92HLRzMWd5j8GehO+60tjrq8bx/jm0obpABIT+nkTNTJlaL+14ziVE2QZYV4VvHO4XQvOYtn92rvxBAAAAMAGep2pEf3c5fYTCskPjx10tNs/dEitfWcAR2pkb+X6knhPNb494o6ye6ByoXD+voAAAAIZBmqxJqEFsmUwL/waaFZ1EyQOlLmgdat7bD4DlnDKE5EZeA7wS8KuScmLjUp2EbKQvkspxKeGj6wCV4/o9afBe/9NGPieBtZO1xGmVIfxYKRK/87tsBmbY+CKUcMExGgKwo9oXid1IeYBEUO08HCareCjE9Ukys7EXjqpBhqzoxi7rDqSeDQAAAGtBnspFFSyfgOol4/fAqtD33XIsmEONOHHMdFPLHo8kgNBLKWb9Ym0a2pfYCaWEGCwtlD2bnmGnRxwnkF0/3ikz9RRWbpi5VR6qekgem9YtTpaFhaa4okfsIeRm5VyV92Bz5W5+vfvQGlzKQQAAAE4Bnul0RH+KbhTf9HyRxrjP7gTLkjfovmjjipEuxkhrSfiXFbTj+bdRKyxdbpt6VAaA+pPkrRqInGqSekN/aHLTqgaJTdQcGWwTKKS0O/AAAAA7AZ7rakR/hNQrERnav5stEZ3ol4KYJIAtAl5O8ChfhNtpxCXz+ANkDpac7UeKcUe4PCnSrmKSFibwLDUAAAB6QZruSahBbJlMFE//Ca2F/yEY9AmJs0u6nGYXXc/tlS+Ep6R9fuGO2/jOCP4GVYBDYfNYov80JOewYiS/s/vrHJa0Qz01zjkBZYXCzefVUhU29LngYko5T1hwFRoyeH8iTqTzThvYOCNdd/EJ3oSIlNm/FwdpeovOsTgAAAAwAZ8NakR/gQeWyVZmnrxk/Y85AYHkLo/YCh9kxyYfbQDb4iHTGTcQHbUUBaOiMgGBAAAAU0GbEknhClJlMCI/FWGmIo6t5jEJYH9imbohlas+DBV/nvXtxmYQEaER8fEGt1ChymPpiTYpdh51Cmh/YiLHIUzWJgypPkYhnUaomR/HCwbaiXipAAAAWUGfMEU0TJ94n01vVRnAY040HFEBoYfXPVz/vno2VX9xlArIIXFkMxlTVQQFmfCo6f94H5UswTk7N1qNbBjEf5ti7RycrMWYZI8EZmojies4I8dHu6teptD5AAAANQGfT3REf4MMlUlILcS/8oa2RUszwulzz8mTh84lHFOM1E/jfwpW/Z9ANgKIGdGUmnAXkT5BAAAALQGfUWpEf36Uk5AyZXqKVAUafPJRrljZB7JJE29//0psy6/KKcUFyP4LnFh/YAAAUAttb292AAAAbG12aGQAAAAAAAAAAAAAAAAAAAPoAALlBAABAAABAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACAABPNXRyYWsAAABcdGtoZAAAAAMAAAAAAAAAAAAAAAEAAAAAAALlBAAAAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAEAAAAAAUAAAAIQAAAAAACRlZHRzAAAAHGVsc3QAAAAAAAAAAQAC5QQAAAgAAAEAAAAATq1tZGlhAAAAIG1kaGQAAAAAAAAAAAAAAAAAACgAAB2kAFXEAAAAAAAtaGRscgAAAAAAAAAAdmlkZQAAAAAAAAAAAAAAAFZpZGVvSGFuZGxlcgAAAE5YbWluZgAAABR2bWhkAAAAAQAAAAAAAAAAAAAAJGRpbmYAAAAcZHJlZgAAAAAAAAABAAAADHVybCAAAAABAABOGHN0YmwAAACoc3RzZAAAAAAAAAABAAAAmGF2YzEAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAUACEAEgAAABIAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAY//8AAAAyYXZjQwFkAAr/4QAZZ2QACqzZRRP58BEAAAMAAQAAAwAUDxIllgEABmjr48siwAAAABBwYXNwAAAAAQAAAAEAAAAYc3R0cwAAAAAAAAABAAAHaQAABAAAAAAwc3RzcwAAAAAAAAAIAAAAAQAAAPsAAAH1AAAC7wAAA+kAAATjAAAF3QAABtcAAC84Y3R0cwAAAAAAAAXlAAAAAQAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAADAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAAEAAAIAAAAAAEAAAwAAAAAAQAABAAAAAAHAAAIAAAAAAEAABAAAAAAAgAABAAAAAAHAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAAGAAAIAAAAAAEAAAwAAAAAAQAABAAAAAANAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAABAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAAGAAAIAAAAAAEAAAwAAAAAAQAABAAAAAAFAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAASAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAQAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAADAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAAGAAAIAAAAAAEAAAwAAAAAAQAABAAAAAAFAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAEAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAABAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAA0AAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAwAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAABAAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAcAAAgAAAAAAQAADAAAAAABAAAEAAAAAAUAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAQAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAEAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAADAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAABgAACAAAAAABAAAMAAAAAAEAAAQAAAAABQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAARAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAYAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAQAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAwAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAADAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAAGAAAIAAAAAAEAABAAAAAAAgAABAAAAAAEAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAANAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAIAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAwAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAABgAACAAAAAABAAAQAAAAAAIAAAQAAAAABAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAADAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAMAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAYAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAQAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAACAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAwAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAgAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAADAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAABwAACAAAAAABAAAMAAAAAAEAAAQAAAAABQAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAADAAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAHHN0c2MAAAAAAAAAAQAAAAEAAAdpAAAAAQAAHbhzdHN6AAAAAAAAAAAAAAdpAAAGRwAAADUAAAAhAAAAJAAAAbwAAABeAAAAIgAAAE0AAAA1AAAAMAAAACQAAAAxAAAAogAAABgAAACGAAAAOQAAAC8AAAAiAAAAbQAAAHYAAAA1AAAAcwAAAD8AAABwAAAAigAAADIAAABOAAAAYgAAADQAAABqAAAAIwAAADsAAABoAAAANAAAAFQAAAA5AAAAUwAAAFcAAABkAAAAXQAAAGMAAAA7AAAARQAAAHsAAABtAAAAYQAAAGQAAABuAAAAZAAAAHkAAAApAAAARQAAAEsAAABmAAAAUQAAAFsAAABSAAAAewAAAF4AAACHAAAAKwAAAFcAAACBAAAALwAAAIAAAABQAAAAPwAAADMAAABSAAAATAAAAGUAAABEAAAAPwAAAGAAAABjAAAAPAAAADoAAABfAAAAXgAAAFsAAAB+AAAATwAAADcAAABVAAAAbQAAAEQAAAA0AAAATgAAAF0AAABsAAAAIgAAAEcAAABtAAAAKAAAASIAAABWAAAAPQAAADcAAACpAAAAOQAAADwAAAAvAAAAvAAAACoAAAB8AAAAPwAAACgAAADAAAAAVgAAAEEAAAAtAAAAOgAAACQAAACSAAAAbgAAAFQAAABLAAAAYQAAAEgAAAA+AAAAOQAAAEQAAAA4AAAAaQAAAIAAAABbAAAALQAAACwAAABmAAAARAAAAHcAAAAyAAAAOQAAACgAAABaAAAAXAAAAE4AAABfAAAAUAAAAEcAAABjAAAANgAAADMAAABkAAAAaQAAAFIAAAA2AAAATwAAAEQAAAA1AAAAMwAAAEsAAAAlAAAAkAAAAJkAAABOAAAATgAAADkAAAB9AAAATgAAAJcAAAA+AAAAOAAAADgAAABsAAAAUQAAACsAAACFAAAAYQAAADkAAAA1AAAAawAAAEkAAACIAAAAeAAAAFAAAABAAAAAOAAAAH0AAABJAAAAqQAAAGUAAABLAAAALgAAAMoAAABZAAAARAAAAD0AAACEAAAAQAAAAIMAAABAAAAAQwAAADUAAAB0AAAAUgAAADwAAAAzAAABCAAAADwAAADNAAAAQwAAAJgAAAAkAAAArgAAAEIAAAA1AAAAigAAAGEAAAA3AAAAMgAAAEIAAAByAAAAIgAAAJ8AAAA4AAAAZwAAAGIAAAA9AAAAVwAAAEgAAAAtAAAAJwAAAG4AAAB+AAAAagAAADwAAAAoAAAAVwAAAEsAAABYAAAAOgAAACwAAABsAAAAZgAAAGoAAABCAAAAWAAAAHEAAABmAAAAcAAAAE0AAAA7AAAAbAAAAFUAAABFAAAAQAAAAFoAAABuAAAANQAABO8AAABjAAAAKQAAAKIAAAArAAAARAAAAIYAAABbAAAARQAAAFYAAACyAAAAQgAAAEkAAAA4AAAAdgAAAFUAAAA8AAAAigAAAHcAAABSAAAARQAAAGYAAAAyAAAAfwAAAQYAAABqAAAAQgAAAD0AAADAAAAAOAAAAF4AAAAyAAAAZgAAACEAAAB8AAAASwAAACoAAACaAAAARgAAADUAAAAvAAAATgAAABsAAABqAAAAZwAAAFsAAABPAAAAXQAAAEoAAAA+AAAAOwAAADwAAABAAAAAaAAAAGwAAAA9AAAAOQAAAHIAAABxAAAATwAAAEsAAACgAAAAMgAAACoAAABlAAAAUQAAAEwAAABiAAAAaAAAAFgAAABUAAAAPAAAADEAAABUAAAAbwAAAFEAAAA+AAAAVQAAAEUAAAAzAAAANAAAAIsAAAAsAAAAKQAAAKEAAABQAAAATwAAAD8AAACJAAAAUAAAAKsAAABIAAAAOwAAADIAAAB0AAAAWgAAAC0AAAB1AAAAUgAAAD8AAAA1AAAAZwAAAEoAAACBAAAAewAAAFEAAABFAAAAPAAAAHoAAABOAAAArgAAAF0AAABKAAAALQAAALoAAABXAAAARAAAADsAAABrAAAAPAAAAMcAAABCAAAAPwAAADYAAABoAAAAUQAAAD8AAAA0AAABBQAAAEIAAAC0AAAAPgAAAIMAAAAdAAAApQAAAD8AAAAkAAAAiAAAAFgAAABGAAAAMQAAADwAAABQAAAAHAAAAF4AAAA4AAAAWQAAAGQAAAA5AAAAVwAAADwAAAAzAAAAOgAAAKAAAABXAAAAigAAADcAAAAoAAAASAAAAEwAAABAAAAAOgAAACwAAABxAAAAdwAAAF4AAABLAAAAaAAAAGcAAABmAAAAdgAAAEoAAAA+AAAAYAAAAFMAAABFAAAAQwAAAFwAAABWAAAAXAAAAE4AAAArAAAALAAAAHgAAAAxAAAASgAAAIoAAABaAAAARAAAAFYAAADDAAAARgAAAEcAAAA4AAAAegAAAFIAAAA/AAAAjQAAAGoAAABUAAAAQgAAAG0AAAAwAAAAigAAAUYAAABWAAAAPQAAAD0AAACvAAAAMwAAAGAAAAA0AAAAYQAAACEAAAByAAAAQQAAACcAAADKAAAAWwAAADoAAAAvAAAARwAAABsAAACSAAAAagAAAGAAAABeAAAAVwAAAE0AAAA3AAAALAAAAGIAAABEAAAAcAAAAGQAAAAxAAAANQAAAIgAAABcAAAARgAAAEYAAACfAAAAMQAAADIAAABRAAAAUgAAAEsAAABiAAAAaQAAAF8AAATrAAAAUgAAADIAAACIAAAAjQAAAFUAAAAzAAAAVAAAAEUAAAA0AAAANQAAAKwAAAAtAAAAKgAAAJQAAABRAAAASgAAADkAAACGAAAASQAAAJwAAABDAAAAOQAAADcAAACMAAAAYgAAAC8AAAB9AAAAeAAAAEQAAAAyAAAAbAAAAEsAAABmAAAAeAAAAE0AAABDAAAAOQAAAHUAAABEAAAAsgAAAF8AAABKAAAALwAAAMEAAABVAAAARgAAADsAAABzAAAAOgAAAJsAAAA7AAAAPwAAADQAAAB0AAAAUQAAADsAAAA1AAAA8QAAADoAAADCAAAAPwAAAHcAAAA6AAAAoQAAAFEAAAA2AAAAiQAAAGAAAAA3AAAAKwAAAEAAAABpAAAAGwAAAH0AAAA5AAAAYAAAAGsAAAA4AAAAPAAAAEQAAAAxAAAAIQAAAHcAAABeAAAAaAAAAEAAAAAsAAAASAAAAGUAAABOAAAAKwAAACsAAACGAAAATAAAACgAAAApAAAAXgAAAE8AAABiAAAAawAAAFsAAAA7AAAAeQAAAFkAAABKAAAAQAAAAFYAAABeAAAAWgAAAC4AAAA4AAAAoAAAAC0AAABbAAAASQAAAH0AAABXAAAAPgAAAFYAAACvAAAAQgAAAEkAAAA3AAAAfQAAAFYAAAA3AAAAoAAAAE8AAABUAAAAQgAAAHUAAAAnAAAAgAAAAFIAAAAzAAAANwAAAR4AAABTAAAAbwAAAC0AAAA+AAAAJAAAAK0AAABBAAAAKgAAACwAAACfAAAAWAAAAD8AAAAsAAAAXwAAACEAAACOAAAAawAAAGcAAABQAAAAagAAAD0AAAAtAAAARgAAAEAAAABNAAAAawAAAGQAAABOAAAAPwAAAIcAAABOAAAAQgAAAEMAAABdAAAAKwAAAHAAAABjAAAAXAAAAE0AAABtAAAATgAAAE0AAABkAAAAMQAAADAAAABaAAAAggAAAFQAAAA8AAAAVAAAAEYAAAA0AAAANgAAAI8AAAAuAAAALAAAAJoAAABLAAAASQAAADsAAACRAAAARwAAAM0AAABDAAAAOQAAADcAAABpAAAAXQAAACoAAAB7AAAAZwAAADUAAAAuAAAAggAAADAAAACHAAAAhwAAAGAAAABIAAAAQAAAAGcAAAA3AAAAtAAAAGYAAABZAAAALAAAALEAAABVAAAARQAAADsAAABoAAAAPAAAAKgAAAA9AAAAQAAAADoAAABrAAAATwAAADsAAAAwAAAA+gAAAEAAAADRAAAAPQAAAHwAAAAoAAAAoAAAAEcAAAAyAAAAWgAAAD0AAAAvAAAE+QAAAEMAAAA2AAAAYgAAAJ8AAAA2AAAAYAAAAHYAAAA9AAAAQAAAAD0AAAAtAAAAJQAAAIAAAABkAAAAcAAAADMAAAAoAAAASgAAAGAAAABJAAAAMgAAADIAAAB/AAAAaQAAAGYAAABQAAAAXQAAAF0AAABqAAAAiAAAAEsAAAA5AAAAaQAAAFEAAABKAAAAQAAAAGAAAABYAAAAXwAAAEcAAAAoAAAAKgAAAIgAAAAiAAAAOwAAAIIAAABVAAAAPgAAAFsAAADQAAAARwAAAEYAAAA6AAAAhQAAAFEAAAA2AAAAjAAAAEsAAABTAAAARgAAAGsAAAAlAAAAgwAAAOgAAABNAAAARwAAAEkAAAC0AAAAPAAAAG0AAAA1AAAAXwAAACUAAACQAAAAOwAAAD0AAAAlAAAAlAAAADcAAACJAAAAJQAAADMAAAAeAAAAYQAAAGsAAABLAAAAVAAAAEkAAAA6AAAALgAAADoAAABHAAAAcwAAAJgAAABhAAAAMgAAADIAAABgAAAAOgAAAK8AAABNAAAARgAAAC0AAABlAAAAVAAAAFQAAAB2AAAAZQAAAF4AAABWAAAAPAAAADgAAABYAAAAggAAAFIAAAA7AAAAVwAAAEsAAAAzAAAANAAAAEsAAAAnAAAAlwAAAJMAAABWAAAATQAAADoAAACIAAAASwAAAJ0AAAA9AAAARAAAADUAAABkAAAATwAAACwAAACGAAAAZgAAAEsAAAAsAAAAbgAAADQAAABdAAAAdwAAAFYAAABQAAAAPAAAAFoAAABKAAAAqwAAAFoAAABJAAAALQAAALEAAABVAAAAQwAAADoAAABmAAAAPwAAAKoAAABAAAAAQwAAADIAAAB8AAAASwAAAD4AAAAvAAABTQAAAD8AAADGAAAARAAAAFoAAAAdAAAAugAAADYAAAAnAAAAeAAAAF4AAAA6AAAALwAAAGAAAABOAAAAIgAAAHYAAAA0AAAAawAAAGYAAAA0AAAAWgAAAEQAAAAvAAAAJAAAAHkAAABZAAAAYQAAADUAAAAwAAAATwAAAGQAAABEAAAAPgAAAEIAAAByAAAAbwAAAG0AAAA4AAAAZgAAAGwAAABgAAAAXwAAAE0AAAA8AAAAUQAAAFkAAABaAAAAQAAAAFgAAABcAAAAVwAAAC8AAAA1AAAAfQAAAC8AAABnAAAASwAAAHEAAABUAAAAOwAAAFUAAAC2AAAASgAAAEgAAAA4AAAAfwAAAFcAAAA0AAAAlwAAAGkAAABJAAAAPgAAAHkAAAAnAAAAeAAAAFAAAAAyAAAANwAAAIoAAAB+AAAAOAAABWYAAABIAAAAIAAAAPUAAABhAAAANwAAACQAAADrAAAATwAAADkAAAAtAAAAtQAAACkAAAAfAAAAjAAAAHMAAAAlAAAAcQAAAEMAAAA/AAAAMAAAAEYAAABAAAAAXQAAAJAAAAA7AAAANgAAAHwAAABVAAAASQAAAD8AAAB+AAAAOQAAACoAAABoAAAAVAAAAFYAAABrAAAAUAAAAE0AAABiAAAAOQAAAB4AAABhAAAAhQAAAEoAAAA2AAAAdQAAAEQAAAA4AAAAMwAAAJQAAAA2AAAAKAAAAHgAAABLAAAAOQAAAJoAAABjAAAASwAAADoAAADBAAAAQgAAAC4AAAA5AAAAbwAAACwAAACTAAAAfgAAAEYAAAA3AAAAtQAAAD4AAAAyAAAAQwAAAH8AAABDAAAANgAAAFAAAACTAAAAagAAAFYAAAAzAAAAzwAAAF4AAAA+AAAAOQAAAJcAAABMAAAASgAAAD0AAACSAAAARQAAADIAAABGAAAAcgAAAEcAAAA1AAAARQAAAREAAABWAAAAXgAAACUAAAEAAAAAPQAAACsAAAAgAAAAmwAAAGEAAAA7AAAALQAAAIIAAAA6AAAAHgAAADIAAABxAAAAVwAAAG4AAAA9AAAAdAAAAEIAAAAuAAAAIQAAAFwAAACVAAAAbQAAADEAAAAmAAAASgAAAGwAAABGAAAAOAAAADEAAACTAAAARwAAADoAAAAkAAAAcwAAAFYAAABnAAAAZwAAAE4AAABIAAAAawAAAF8AAABDAAAAPQAAAF4AAABmAAAAXwAAAEcAAAAmAAAAIwAAAJYAAAAuAAAANwAAAI4AAABeAAAAPAAAAFYAAACzAAAAPwAAAEoAAAA3AAAAhAAAAFUAAAAxAAAAqQAAAGgAAABTAAAAPwAAAJMAAAAnAAAAqQAAAFEAAAA8AAAANwAAAQwAAABNAAAAbwAAACcAAABHAAAALwAAALQAAABHAAAALQAAACQAAACoAAAAWQAAADsAAAAzAAAAdgAAAB8AAABxAAAAdAAAAEgAAAArAAAAeQAAAF0AAABPAAAALAAAAFYAAAA+AAAAhAAAAHQAAAA6AAAAOwAAAHQAAABZAAAASQAAAEYAAACZAAAAOAAAACgAAABfAAAAUwAAAFkAAAB+AAAAWwAAAFUAAABpAAAAQgAAACEAAABeAAAAdAAAAFIAAAA0AAAAbAAAAEsAAAA7AAAANAAAAH0AAAAxAAAAKAAAAHoAAABeAAAAMwAAAJ4AAABmAAAASQAAADoAAAChAAAARQAAACcAAABBAAAAUQAAACsAAABSAAAATwAAAC4AAAWcAAAAiQAAAC8AAACaAAAAXQAAAEoAAAAzAAAAYwAAAFMAAABMAAAAUgAAAKMAAAA/AAAALAAAALAAAABVAAAAOQAAADkAAAByAAAARwAAAD4AAAA7AAAArQAAAFMAAAAzAAAANgAAAQkAAABSAAAAKQAAADsAAADTAAAAOwAAAN4AAAA4AAAAHgAAAC0AAAB9AAAAXAAAADsAAAAuAAAAcwAAADUAAAA0AAAAIQAAAIYAAAAyAAAAVwAAAG8AAABFAAAAXgAAAEAAAAAzAAAALAAAAGQAAACCAAAAggAAADAAAAAjAAAASwAAAGsAAAA/AAAAOwAAAL4AAABWAAAAMQAAADUAAABHAAAAYwAAAFgAAABoAAAAYwAAAEsAAABIAAAAXwAAAF4AAAA/AAAAOAAAAGAAAABUAAAAYgAAAFAAAAAqAAAAKwAAAHUAAAAwAAAAOgAAAI0AAABbAAAAOQAAAFYAAAC+AAAAQQAAAEYAAAA2AAAAgQAAAEwAAAA0AAAAlgAAAHAAAABQAAAAPQAAAHgAAAAvAAAAmwAAAF4AAAA9AAAAMQAAAUUAAABmAAAAeAAAAC0AAAB6AAAAMAAAABwAAAAuAAAA4AAAADUAAAAqAAAARgAAAEkAAABIAAAAJgAAAB0AAACVAAAAbgAAAGwAAAAtAAAAVgAAAD4AAABHAAAALQAAAFkAAAA5AAAAYgAAAIEAAAAyAAAAMQAAAH8AAABVAAAATQAAAEYAAAB7AAAARAAAACYAAABbAAAATQAAAFcAAABsAAAAVgAAAEoAAABsAAAAOgAAACYAAACAAAAAfQAAAE4AAAA/AAAAcQAAAEYAAAA6AAAAMwAAAKcAAAAyAAAAKgAAAH0AAABUAAAANAAAAJwAAABpAAAARQAAADsAAAC9AAAAQQAAAC4AAABGAAAAbQAAACgAAAB6AAAAYgAAAEYAAAAyAAAAgAAAADsAAAAxAAAAPgAAAH0AAABAAAAANQAAAE4AAACWAAAAZAAAAFMAAAAxAAAAxQAAAGEAAAA/AAAANAAAAH0AAABJAAAATAAAAEEAAACuAAAAQAAAADMAAABIAAAAbgAAAEUAAAAyAAAARQAAAPsAAABSAAAAXAAAACYAAADkAAAARwAAADAAAAAgAAAAkQAAAGEAAAA6AAAALQAAAIYAAAA7AAAAHQAAADAAAAB2AAAAUgAAAGkAAAA9AAAAYgAAAEAAAAAuAAAAHgAAAGUAAAB/AAAAcQAAAC0AAAAiAAAAWAAAAHUAAABFAAAAOAAAADIAAACAAAAAPgAAADUAAAAqAAAAZAAAAFkAAABfAAAFYwAAAFIAAAA3AAAAlQAAAGEAAABCAAAANwAAAFkAAABaAAAAdAAAAEIAAAAoAAAAJQAAAJEAAAAwAAAANwAAAJEAAABdAAAAPAAAAFUAAAC3AAAARgAAAEgAAAA2AAAAegAAAE0AAAA2AAAAnAAAAG0AAABTAAAAOwAAAIkAAAA7AAAAhwAAAFEAAAA9AAAAOwAAAVMAAABoAAAAfAAAACsAAACCAAAAMgAAABwAAAAyAAAA2AAAAGUAAAAiAAAAOwAAAHQAAABOAAAAKAAAACIAAACYAAAAewAAAEYAAAApAAAAYwAAAFkAAABBAAAALwAAAFAAAABHAAAAcAAAAHoAAAAvAAAAMQAAAHIAAABfAAAASQAAAEcAAACZAAAANgAAAB4AAABiAAAAUQAAAFcAAAB0AAAAXQAAAFUAAABTAAAAPgAAACAAAABhAAAAgAAAAEsAAAA0AAAAZwAAAEIAAAA4AAAAMAAAAIUAAAAuAAAAIgAAAHgAAABWAAAANAAAAJsAAABrAAAARwAAAD4AAADBAAAAPgAAACwAAABGAAAAbQAAACUAAACDAAAAWQAAAEAAAAAyAAAAvQAAAEsAAABLAAAARwAAAHMAAAA7AAAANQAAAEoAAACXAAAAcwAAAFEAAAAxAAAAvQAAAGkAAAA+AAAAOQAAAI4AAABKAAAASQAAAD0AAACTAAAARQAAADUAAABGAAAAbAAAAEYAAAA2AAAARgAAASEAAABZAAAAZAAAACcAAADjAAAAPwAAACwAAAAgAAAAjgAAAF8AAAA6AAAALwAAAE4AAAAfAAAAhwAAADIAAAAyAAAAVQAAAHMAAAA7AAAATgAAADkAAAAuAAAAJQAAAGMAAACPAAAAdAAAADQAAAAhAAAATQAAAJwAAABLAAAAOQAAADEAAACiAAAARQAAADgAAAAlAAAAcQAAAGcAAABmAAAAgQAAAEoAAAA8AAAAYgAAAGoAAABCAAAAPgAAAGIAAABgAAAAZQAAAEYAAAAqAAAAKQAAAHYAAAA0AAAAPQAAAIgAAABbAAAAOQAAAFUAAADRAAAAPwAAAEQAAAA6AAAAhwAAAE8AAAA1AAAAoQAAAGwAAABSAAAAQwAAAJAAAAAmAAAAuQAAAE8AAAA9AAAAMwAAAOkAAABPAAAAbgAAACwAAACbAAAANwAAACIAAAAyAAAAygAAAD0AAAAiAAAAOgAAAEgAAAA0AAAAKQAAACUAAACDAAAAbAAAAE8AAAAoAAAAYAAAAFEAAABAAAAALQAAAGIAAABEAAAAcAAAAHMAAAA+AAAAOQAAAFQAAABNAAAAPAAABYMAAACNAAAAKwAAAGsAAABqAAAASgAAAFQAAABqAAAAWAAAAGAAAABqAAAAOwAAADcAAABUAAAAcQAAAFUAAAA9AAAAaAAAAEkAAAA6AAAAMwAAAJgAAAA5AAAAIwAAAHwAAABVAAAANAAAAJkAAABlAAAATAAAADwAAAC4AAAAQAAAACoAAABIAAAAawAAACgAAACHAAAAXgAAAD4AAAAxAAAAbwAAADkAAAB7AAAARwAAAD0AAAA0AAAAgwAAAF0AAABOAAAAUwAAAMoAAAA9AAAALAAAANMAAABVAAAAQQAAADoAAABkAAAAUQAAAEQAAAA/AAAApAAAAFIAAAAvAAAAPAAAAREAAABTAAAAKAAAADsAAAD8AAAAOgAAAOQAAAA2AAAAFwAAADQAAACCAAAAXAAAADYAAAAvAAAAdwAAADgAAAA5AAAAJwAAAHYAAAA5AAAAWgAAAIEAAABAAAAATAAAAEQAAAAzAAAAHwAAAIQAAAB0AAAAcAAAADcAAAAlAAAAWgAAAFoAAABOAAAANQAAAC0AAACTAAAAXQAAADsAAAAjAAAAcwAAAFgAAABoAAAAdwAAAEoAAAA6AAAAYQAAAFwAAAA9AAAANgAAAF0AAABbAAAAZgAAAEkAAAApAAAAKQAAAI0AAAAzAAAAPgAAAIsAAABdAAAANQAAAFcAAACxAAAAPQAAAEkAAAA3AAAAeQAAAE4AAAA0AAAAigAAAG8AAABSAAAAPwAAAH4AAAA0AAAAVwAAAF0AAAA5AAAAMQAAABRzdGNvAAAAAAAAAAEAAAAwAAAAYnVkdGEAAABabWV0YQAAAAAAAAAhaGRscgAAAAAAAAAAbWRpcmFwcGwAAAAAAAAAAAAAAAAtaWxzdAAAACWpdG9vAAAAHWRhdGEAAAABAAAAAExhdmY1Ny44My4xMDA=",
+       "ok": true,
+       "headers": [
+        [
+         "content-type",
+         "video/mp4"
+        ]
+       ],
+       "status": 200.0,
+       "status_text": ""
+      }
+     },
+     "base_uri": "https://localhost:8080/",
+     "height": 501.0
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "\n",
+       "    <video width=\"640\" height=\"480\" controls>\n",
+       "      <source src=\"/nbextensions/vid.mp4\" type=\"video/mp4\">\n",
+       "    </video>\n",
+       "  "
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "execution_count": 24,
+     "metadata": {
+      "tags": []
+     },
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "play_video('pong_pretrained/0.avi')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "U-SyGcZBCmPn",
+    "colab_type": "text"
+   },
+   "source": [
+    "# Train your policy (model-free training)\n",
+    "Training model-free on Pong (it takes a few hours):"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "id": "WIQazd5aCocc",
+    "colab_type": "code",
+    "outputId": "0a440c18-affc-4b2a-d6e1-c3cda84465bc",
+    "executionInfo": {
+     "status": "ok",
+     "timestamp": 1.553254256733E12,
+     "user_tz": -60.0,
+     "elapsed": 19957.0,
+     "user": {
+      "displayName": "Piotr Kozakowski",
+      "photoUrl": "",
+      "userId": "01014928596539690143"
+     }
+    },
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 1516.0
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
+      "For more information, please see:\n",
+      "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
+      "  * https://github.com/tensorflow/addons\n",
+      "If you depend on functionality not listed there, please file an issue.\n",
+      "\n",
+      "2019-03-22 11:30:42.987149: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 2300000000 Hz\n",
+      "2019-03-22 11:30:42.987392: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x30323c0 executing computations on platform Host. Devices:\n",
+      "2019-03-22 11:30:42.987491: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): <undefined>, <undefined>\n",
+      "2019-03-22 11:30:43.082876: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:998] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
+      "2019-03-22 11:30:43.083442: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x3032100 executing computations on platform CUDA. Devices:\n",
+      "2019-03-22 11:30:43.083493: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): Tesla K80, Compute Capability 3.7\n",
+      "2019-03-22 11:30:43.083843: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1433] Found device 0 with properties: \n",
+      "name: Tesla K80 major: 3 minor: 7 memoryClockRate(GHz): 0.8235\n",
+      "pciBusID: 0000:00:04.0\n",
+      "totalMemory: 11.17GiB freeMemory: 11.10GiB\n",
+      "2019-03-22 11:30:43.083879: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n",
+      "2019-03-22 11:30:43.475526: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n",
+      "2019-03-22 11:30:43.475601: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n",
+      "2019-03-22 11:30:43.475629: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n",
+      "2019-03-22 11:30:43.476026: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:42] Overriding allow_growth setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.\n",
+      "2019-03-22 11:30:43.476131: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Colocations handled automatically by placer.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/rl/envs/py_func_batch_env.py:122: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "tf.py_func is deprecated in TF V2. Instead, use\n",
+      "    tf.py_function, which takes a python function which manipulates tf eager\n",
+      "    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to\n",
+      "    an ndarray (just call tensor.numpy()) but having access to eager tensors\n",
+      "    means `tf.py_function`s can use accelerators such as GPUs as well as\n",
+      "    being differentiable using a gradient tape.\n",
+      "    \n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/utils/t2t_model.py:1358: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use tf.cast instead.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/function.py:1007: calling Graph.create_op (from tensorflow.python.framework.ops) with compute_shapes is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Shapes are always computed; don't use the compute_shapes as it has no effect.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:277: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use tf.cast instead.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:598: conv2d (from tensorflow.python.layers.convolutional) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use keras.layers.conv2d instead.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:602: flatten (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use keras.layers.flatten instead.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:603: dropout (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use keras.layers.dropout instead.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:604: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use keras.layers.dense instead.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:2887: multinomial (from tensorflow.python.ops.random_ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use tf.random.categorical instead.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/rl/ppo_learner.py:479: Print (from tensorflow.python.ops.logging_ops) is deprecated and will be removed after 2018-08-20.\n",
+      "Instructions for updating:\n",
+      "Use tf.print instead of tf.Print. Note that tf.print returns a no-output operator that directly prints the output. Outside of defuns or eager mode, this operator will not be executed unless it is directly specified in session.run or used as a control dependency for other operators. This is only a concern in graph mode. Below is an example of how to ensure tf.print executes in graph mode:\n",
+      "```python\n",
+      "    sess = tf.Session()\n",
+      "    with sess.as_default():\n",
+      "        tensor = tf.range(10)\n",
+      "        print_op = tf.print(tensor)\n",
+      "        with tf.control_dependencies([print_op]):\n",
+      "          out = tf.add(tensor, tensor)\n",
+      "        sess.run(out)\n",
+      "    ```\n",
+      "Additionally, to use tf.print in python 2.7, users must make sure to import\n",
+      "the following:\n",
+      "\n",
+      "  `from __future__ import print_function`\n",
+      "\n",
+      "2019-03-22 11:30:49.903512: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n",
+      "2019-03-22 11:30:49.903591: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n",
+      "2019-03-22 11:30:49.903620: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n",
+      "2019-03-22 11:30:49.903639: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n",
+      "2019-03-22 11:30:49.903898: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/saver.py:1266: checkpoint_exists (from tensorflow.python.training.checkpoint_management) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use standard file APIs to check for files with this prefix.\n",
+      "2019-03-22 11:30:51.335217: I tensorflow/stream_executor/dso_loader.cc:152] successfully opened CUDA library libcublas.so.10.0 locally\n",
+      "mean_score: [0][0][0]\n",
+      "^C\n"
+     ]
+    }
+   ],
+   "source": [
+    "!python -m tensor2tensor.rl.trainer_model_free \\\n",
+    "  --hparams_set=rlmf_base \\\n",
+    "  --hparams=game=pong \\\n",
+    "  --output_dir=mf_pong"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "FbSjwVAtCvLY",
+    "colab_type": "text"
+   },
+   "source": [
+    "Hyperparameter sets are defined in `tensor2tensor/models/research/rl.py`. You can override them using the hparams flag, e.g.\n",
+    "\n",
+    "```\n",
+    "--hparams=game=kung_fu_master,frame_stack_size=5\n",
+    "```\n",
+    "\n",
+    "As in model-based training, the periodic evaluation runs with timestep limit of 1000. To do full evaluation after training, run:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "id": "jppi4FE5C2nB",
+    "colab_type": "code",
+    "outputId": "a10afb7c-edd6-4a93-eee4-e3876977e825",
+    "executionInfo": {
+     "status": "ok",
+     "timestamp": 1.553254412202E12,
+     "user_tz": -60.0,
+     "elapsed": 15104.0,
+     "user": {
+      "displayName": "Piotr Kozakowski",
+      "photoUrl": "",
+      "userId": "01014928596539690143"
+     }
+    },
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 4083.0
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
+      "For more information, please see:\n",
+      "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
+      "  * https://github.com/tensorflow/addons\n",
+      "If you depend on functionality not listed there, please file an issue.\n",
+      "\n",
+      "INFO:tensorflow:Overriding hparams in rlmf_tiny with game=pong,eval_max_num_noops=0,eval_sampling_temps=[0.5]\n",
+      "INFO:tensorflow:Evaluating metric mean_reward/eval/sampling_temp_0.5_max_noops_0_unclipped\n",
+      "2019-03-22 11:33:23.214052: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 2300000000 Hz\n",
+      "2019-03-22 11:33:23.214294: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x2d07020 executing computations on platform Host. Devices:\n",
+      "2019-03-22 11:33:23.214335: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): <undefined>, <undefined>\n",
+      "2019-03-22 11:33:23.309948: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:998] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
+      "2019-03-22 11:33:23.310546: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x2d067e0 executing computations on platform CUDA. Devices:\n",
+      "2019-03-22 11:33:23.310585: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): Tesla K80, Compute Capability 3.7\n",
+      "2019-03-22 11:33:23.310991: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1433] Found device 0 with properties: \n",
+      "name: Tesla K80 major: 3 minor: 7 memoryClockRate(GHz): 0.8235\n",
+      "pciBusID: 0000:00:04.0\n",
+      "totalMemory: 11.17GiB freeMemory: 11.10GiB\n",
+      "2019-03-22 11:33:23.311027: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n",
+      "2019-03-22 11:33:23.707039: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n",
+      "2019-03-22 11:33:23.707114: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n",
+      "2019-03-22 11:33:23.707139: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n",
+      "2019-03-22 11:33:23.707459: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:42] Overriding allow_growth setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.\n",
+      "2019-03-22 11:33:23.707523: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n",
+      "INFO:tensorflow:Using DummyPolicyProblem for the policy.\n",
+      "INFO:tensorflow:Setting T2TModel mode to 'train'\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Colocations handled automatically by placer.\n",
+      "INFO:tensorflow:Using variable initializer: orthogonal\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/utils/t2t_model.py:1358: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use tf.cast instead.\n",
+      "INFO:tensorflow:Transforming feature 'input_action' with symbol_modality_6_64.bottom\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/function.py:1007: calling Graph.create_op (from tensorflow.python.framework.ops) with compute_shapes is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Shapes are always computed; don't use the compute_shapes as it has no effect.\n",
+      "INFO:tensorflow:Transforming feature 'input_reward' with symbol_modality_3_64.bottom\n",
+      "INFO:tensorflow:Transforming feature 'inputs' with video_modality.bottom\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_video.py:495: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "tf.py_func is deprecated in TF V2. Instead, use\n",
+      "    tf.py_function, which takes a python function which manipulates tf eager\n",
+      "    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to\n",
+      "    an ndarray (just call tensor.numpy()) but having access to eager tensors\n",
+      "    means `tf.py_function`s can use accelerators such as GPUs as well as\n",
+      "    being differentiable using a gradient tape.\n",
+      "    \n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:277: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use tf.cast instead.\n",
+      "INFO:tensorflow:Transforming feature 'target_action' with symbol_modality_6_64.targets_bottom\n",
+      "INFO:tensorflow:Transforming feature 'target_policy' with identity_modality.targets_bottom\n",
+      "INFO:tensorflow:Transforming feature 'target_reward' with symbol_modality_3_64.targets_bottom\n",
+      "INFO:tensorflow:Transforming feature 'target_value' with identity_modality.targets_bottom\n",
+      "INFO:tensorflow:Transforming feature 'targets' with video_modality.targets_bottom\n",
+      "INFO:tensorflow:Building model body\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:598: conv2d (from tensorflow.python.layers.convolutional) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use keras.layers.conv2d instead.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:602: flatten (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use keras.layers.flatten instead.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:603: dropout (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use keras.layers.dropout instead.\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/models/research/rl.py:604: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use keras.layers.dense instead.\n",
+      "INFO:tensorflow:Transforming body output with identity_modality.top\n",
+      "INFO:tensorflow:Transforming body output with identity_modality.top\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensor2tensor/layers/common_layers.py:2887: multinomial (from tensorflow.python.ops.random_ops) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use tf.random.categorical instead.\n",
+      "2019-03-22 11:33:24.564271: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0\n",
+      "2019-03-22 11:33:24.564350: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:\n",
+      "2019-03-22 11:33:24.564376: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 \n",
+      "2019-03-22 11:33:24.564410: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N \n",
+      "2019-03-22 11:33:24.564687: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10754 MB memory) -> physical GPU (device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7)\n",
+      "INFO:tensorflow:Restoring checkpoint mf_pong/model.ckpt-9\n",
+      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/saver.py:1266: checkpoint_exists (from tensorflow.python.training.checkpoint_management) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use standard file APIs to check for files with this prefix.\n",
+      "INFO:tensorflow:Restoring parameters from mf_pong/model.ckpt-9\n",
+      "2019-03-22 11:33:24.985295: I tensorflow/stream_executor/dso_loader.cc:152] successfully opened CUDA library libcublas.so.10.0 locally\n",
+      "INFO:tensorflow:Step 5, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 10, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 15, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 20, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 25, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 30, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 35, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 40, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 45, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 50, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 55, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 60, mean_score: 0.000000\n",
+      "INFO:tensorflow:Step 65, mean_score: -1.000000\n",
+      "INFO:tensorflow:Step 70, mean_score: -1.000000\n",
+      "INFO:tensorflow:Step 75, mean_score: -1.000000\n",
+      "INFO:tensorflow:Step 80, mean_score: -1.000000\n",
+      "INFO:tensorflow:Step 85, mean_score: -1.000000\n",
+      "INFO:tensorflow:Step 90, mean_score: -1.000000\n",
+      "INFO:tensorflow:Step 95, mean_score: -1.000000\n",
+      "INFO:tensorflow:Step 100, mean_score: -2.000000\n",
+      "INFO:tensorflow:Step 105, mean_score: -2.000000\n",
+      "INFO:tensorflow:Step 110, mean_score: -2.000000\n",
+      "INFO:tensorflow:Step 115, mean_score: -2.000000\n",
+      "INFO:tensorflow:Step 120, mean_score: -2.000000\n",
+      "INFO:tensorflow:Step 125, mean_score: -2.000000\n",
+      "INFO:tensorflow:Step 130, mean_score: -2.000000\n",
+      "INFO:tensorflow:Step 135, mean_score: -3.000000\n",
+      "INFO:tensorflow:Step 140, mean_score: -3.000000\n",
+      "INFO:tensorflow:Step 145, mean_score: -3.000000\n",
+      "INFO:tensorflow:Step 150, mean_score: -3.000000\n",
+      "INFO:tensorflow:Step 155, mean_score: -3.000000\n",
+      "INFO:tensorflow:Step 160, mean_score: -3.000000\n",
+      "INFO:tensorflow:Step 165, mean_score: -3.000000\n",
+      "INFO:tensorflow:Step 170, mean_score: -4.000000\n",
+      "INFO:tensorflow:Step 175, mean_score: -4.000000\n",
+      "INFO:tensorflow:Step 180, mean_score: -4.000000\n",
+      "INFO:tensorflow:Step 185, mean_score: -4.000000\n",
+      "INFO:tensorflow:Step 190, mean_score: -4.000000\n",
+      "INFO:tensorflow:Step 195, mean_score: -4.000000\n",
+      "INFO:tensorflow:Step 200, mean_score: -4.000000\n",
+      "INFO:tensorflow:Step 205, mean_score: -5.000000\n",
+      "INFO:tensorflow:Step 210, mean_score: -5.000000\n",
+      "INFO:tensorflow:Step 215, mean_score: -5.000000\n",
+      "INFO:tensorflow:Step 220, mean_score: -5.000000\n",
+      "INFO:tensorflow:Step 225, mean_score: -5.000000\n",
+      "INFO:tensorflow:Step 230, mean_score: -5.000000\n",
+      "INFO:tensorflow:Step 235, mean_score: -5.000000\n",
+      "INFO:tensorflow:Step 240, mean_score: -6.000000\n",
+      "INFO:tensorflow:Step 245, mean_score: -6.000000\n",
+      "INFO:tensorflow:Step 250, mean_score: -6.000000\n",
+      "INFO:tensorflow:Step 255, mean_score: -6.000000\n",
+      "INFO:tensorflow:Step 260, mean_score: -6.000000\n",
+      "INFO:tensorflow:Step 265, mean_score: -6.000000\n",
+      "INFO:tensorflow:Step 270, mean_score: -6.000000\n",
+      "INFO:tensorflow:Step 275, mean_score: -7.000000\n",
+      "INFO:tensorflow:Step 280, mean_score: -7.000000\n",
+      "INFO:tensorflow:Step 285, mean_score: -7.000000\n",
+      "INFO:tensorflow:Step 290, mean_score: -7.000000\n",
+      "INFO:tensorflow:Step 295, mean_score: -7.000000\n",
+      "INFO:tensorflow:Step 300, mean_score: -7.000000\n",
+      "INFO:tensorflow:Step 305, mean_score: -7.000000\n",
+      "INFO:tensorflow:Step 310, mean_score: -8.000000\n",
+      "INFO:tensorflow:Step 315, mean_score: -8.000000\n",
+      "INFO:tensorflow:Step 320, mean_score: -8.000000\n",
+      "INFO:tensorflow:Step 325, mean_score: -8.000000\n",
+      "INFO:tensorflow:Step 330, mean_score: -8.000000\n",
+      "INFO:tensorflow:Step 335, mean_score: -8.000000\n",
+      "INFO:tensorflow:Step 340, mean_score: -8.000000\n",
+      "INFO:tensorflow:Step 345, mean_score: -9.000000\n",
+      "INFO:tensorflow:Step 350, mean_score: -9.000000\n",
+      "INFO:tensorflow:Step 355, mean_score: -9.000000\n",
+      "INFO:tensorflow:Step 360, mean_score: -9.000000\n",
+      "INFO:tensorflow:Step 365, mean_score: -9.000000\n",
+      "INFO:tensorflow:Step 370, mean_score: -9.000000\n",
+      "INFO:tensorflow:Step 375, mean_score: -9.000000\n",
+      "INFO:tensorflow:Step 380, mean_score: -10.000000\n",
+      "INFO:tensorflow:Step 385, mean_score: -10.000000\n",
+      "INFO:tensorflow:Step 390, mean_score: -10.000000\n",
+      "INFO:tensorflow:Step 395, mean_score: -10.000000\n",
+      "INFO:tensorflow:Step 400, mean_score: -10.000000\n",
+      "INFO:tensorflow:Step 405, mean_score: -10.000000\n",
+      "INFO:tensorflow:Step 410, mean_score: -10.000000\n",
+      "INFO:tensorflow:Step 415, mean_score: -11.000000\n",
+      "INFO:tensorflow:Step 420, mean_score: -11.000000\n",
+      "INFO:tensorflow:Step 425, mean_score: -11.000000\n",
+      "INFO:tensorflow:Step 430, mean_score: -11.000000\n",
+      "INFO:tensorflow:Step 435, mean_score: -11.000000\n",
+      "INFO:tensorflow:Step 440, mean_score: -11.000000\n",
+      "INFO:tensorflow:Step 445, mean_score: -11.000000\n",
+      "INFO:tensorflow:Step 450, mean_score: -12.000000\n",
+      "INFO:tensorflow:Step 455, mean_score: -12.000000\n",
+      "INFO:tensorflow:Step 460, mean_score: -12.000000\n",
+      "INFO:tensorflow:Step 465, mean_score: -12.000000\n",
+      "INFO:tensorflow:Step 470, mean_score: -12.000000\n",
+      "INFO:tensorflow:Step 475, mean_score: -12.000000\n",
+      "INFO:tensorflow:Step 480, mean_score: -12.000000\n",
+      "INFO:tensorflow:Step 485, mean_score: -13.000000\n",
+      "INFO:tensorflow:Step 490, mean_score: -13.000000\n",
+      "INFO:tensorflow:Step 495, mean_score: -13.000000\n",
+      "INFO:tensorflow:Step 500, mean_score: -13.000000\n",
+      "INFO:tensorflow:Step 505, mean_score: -13.000000\n",
+      "INFO:tensorflow:Step 510, mean_score: -13.000000\n",
+      "INFO:tensorflow:Step 515, mean_score: -13.000000\n",
+      "INFO:tensorflow:Step 520, mean_score: -14.000000\n",
+      "INFO:tensorflow:Step 525, mean_score: -14.000000\n",
+      "INFO:tensorflow:Step 530, mean_score: -14.000000\n",
+      "INFO:tensorflow:Step 535, mean_score: -14.000000\n",
+      "INFO:tensorflow:Step 540, mean_score: -14.000000\n",
+      "INFO:tensorflow:Step 545, mean_score: -14.000000\n",
+      "INFO:tensorflow:Step 550, mean_score: -14.000000\n",
+      "INFO:tensorflow:Step 555, mean_score: -15.000000\n",
+      "INFO:tensorflow:Step 560, mean_score: -15.000000\n",
+      "INFO:tensorflow:Step 565, mean_score: -15.000000\n",
+      "INFO:tensorflow:Step 570, mean_score: -15.000000\n",
+      "INFO:tensorflow:Step 575, mean_score: -15.000000\n",
+      "INFO:tensorflow:Step 580, mean_score: -15.000000\n",
+      "INFO:tensorflow:Step 585, mean_score: -15.000000\n",
+      "INFO:tensorflow:Step 590, mean_score: -16.000000\n",
+      "INFO:tensorflow:Step 595, mean_score: -16.000000\n",
+      "INFO:tensorflow:Step 600, mean_score: -16.000000\n",
+      "INFO:tensorflow:Step 605, mean_score: -16.000000\n",
+      "INFO:tensorflow:Step 610, mean_score: -16.000000\n",
+      "INFO:tensorflow:Step 615, mean_score: -16.000000\n",
+      "INFO:tensorflow:Step 620, mean_score: -16.000000\n",
+      "INFO:tensorflow:Step 625, mean_score: -17.000000\n",
+      "INFO:tensorflow:Step 630, mean_score: -17.000000\n",
+      "INFO:tensorflow:Step 635, mean_score: -17.000000\n",
+      "INFO:tensorflow:Step 640, mean_score: -17.000000\n",
+      "INFO:tensorflow:Step 645, mean_score: -17.000000\n",
+      "INFO:tensorflow:Step 650, mean_score: -17.000000\n",
+      "INFO:tensorflow:Step 655, mean_score: -17.000000\n",
+      "INFO:tensorflow:Step 660, mean_score: -18.000000\n",
+      "INFO:tensorflow:Step 665, mean_score: -18.000000\n",
+      "INFO:tensorflow:Step 670, mean_score: -18.000000\n",
+      "INFO:tensorflow:Step 675, mean_score: -18.000000\n",
+      "INFO:tensorflow:Step 680, mean_score: -18.000000\n",
+      "INFO:tensorflow:Step 685, mean_score: -18.000000\n",
+      "INFO:tensorflow:Step 690, mean_score: -18.000000\n",
+      "INFO:tensorflow:Step 695, mean_score: -19.000000\n",
+      "INFO:tensorflow:Step 700, mean_score: -19.000000\n",
+      "INFO:tensorflow:Step 705, mean_score: -19.000000\n",
+      "INFO:tensorflow:Step 710, mean_score: -19.000000\n",
+      "INFO:tensorflow:Step 715, mean_score: -19.000000\n",
+      "INFO:tensorflow:Step 720, mean_score: -19.000000\n",
+      "INFO:tensorflow:Step 725, mean_score: -19.000000\n",
+      "INFO:tensorflow:Step 730, mean_score: -20.000000\n",
+      "INFO:tensorflow:Step 735, mean_score: -20.000000\n",
+      "INFO:tensorflow:Step 740, mean_score: -20.000000\n",
+      "INFO:tensorflow:Step 745, mean_score: -20.000000\n",
+      "INFO:tensorflow:Step 750, mean_score: -20.000000\n",
+      "INFO:tensorflow:Step 755, mean_score: -20.000000\n",
+      "INFO:tensorflow:Step 760, mean_score: -20.000000\n"
+     ]
+    }
+   ],
+   "source": [
+    "!python -m tensor2tensor.rl.evaluator \\\n",
+    "  --loop_hparams_set=rlmf_tiny \\\n",
+    "  --hparams=game=pong \\\n",
+    "  --policy_dir=mf_pong \\\n",
+    "  --debug_video_path=mf_pong \\\n",
+    "  --num_debug_videos=4 \\\n",
+    "  --eval_metrics_dir=mf_pong/full_eval_metrics"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "id": "mDoR0C0ZKCOn",
+    "colab_type": "code",
+    "outputId": "aba41a4d-2957-4ea0-d511-eae7ea4e238e",
+    "executionInfo": {
+     "status": "ok",
+     "timestamp": 1.553254513355E12,
+     "user_tz": -60.0,
+     "elapsed": 3908.0,
+     "user": {
+      "displayName": "Piotr Kozakowski",
+      "photoUrl": "",
+      "userId": "01014928596539690143"
+     }
+    },
+    "colab": {
+     "resources": {
+      "http://localhost:8080/nbextensions/vid.mp4": {
+       "data": "AAAAIGZ0eXBpc29tAAACAGlzb21pc28yYXZjMW1wNDEAAAAIZnJlZQAA6u9tZGF0AAACrgYF//+q3EXpvebZSLeWLNgg2SPu73gyNjQgLSBjb3JlIDE1MiByMjg1NCBlOWE1OTAzIC0gSC4yNjQvTVBFRy00IEFWQyBjb2RlYyAtIENvcHlsZWZ0IDIwMDMtMjAxNyAtIGh0dHA6Ly93d3cudmlkZW9sYW4ub3JnL3gyNjQuaHRtbCAtIG9wdGlvbnM6IGNhYmFjPTEgcmVmPTMgZGVibG9jaz0xOjA6MCBhbmFseXNlPTB4MzoweDExMyBtZT1oZXggc3VibWU9NyBwc3k9MSBwc3lfcmQ9MS4wMDowLjAwIG1peGVkX3JlZj0xIG1lX3JhbmdlPTE2IGNocm9tYV9tZT0xIHRyZWxsaXM9MSA4eDhkY3Q9MSBjcW09MCBkZWFkem9uZT0yMSwxMSBmYXN0X3Bza2lwPTEgY2hyb21hX3FwX29mZnNldD0tMiB0aHJlYWRzPTMgbG9va2FoZWFkX3RocmVhZHM9MSBzbGljZWRfdGhyZWFkcz0wIG5yPTAgZGVjaW1hdGU9MSBpbnRlcmxhY2VkPTAgYmx1cmF5X2NvbXBhdD0wIGNvbnN0cmFpbmVkX2ludHJhPTAgYmZyYW1lcz0zIGJfcHlyYW1pZD0yIGJfYWRhcHQ9MSBiX2JpYXM9MCBkaXJlY3Q9MSB3ZWlnaHRiPTEgb3Blbl9nb3A9MCB3ZWlnaHRwPTIga2V5aW50PTI1MCBrZXlpbnRfbWluPTEwIHNjZW5lY3V0PTQwIGludHJhX3JlZnJlc2g9MCByY19sb29rYWhlYWQ9NDAgcmM9Y3JmIG1idHJlZT0xIGNyZj0yMy4wIHFjb21wPTAuNjAgcXBtaW49MCBxcG1heD02OSBxcHN0ZXA9NCBpcF9yYXRpbz0xLjQwIGFxPTE6MS4wMACAAAADkWWIhABvrNdXNvEPmO7lwVl73sPl0EDBzzvrz1O9Sgfa49FGnVhGNj4PrUzIEjAsiR14q5boH034au6fMfeHzW8BQIdLu5D8GWFcvhnUQvMLIDm/5fDlJWNI1pLZ0KekyKgRZvEg10IZZePvLcj64kGJzCMbJi6QZbX4WMzyM/ZwsXoWWPBmmlKBzFixHWdkptjcAYhpgDpXSILlIffpBFr5Fmv8Xdrl5eZtB/U18q6RE0tX2BrhekKyOZ5lJWnXZWIEICkLYIda8x0l/aAug9zkJAN2UJ5v8AfQgXgS7iPy41I11UQneH59QQ6r2Fy+bXVz7hKXvFUUQUW2NfwyAHSubAtKRV8FgrIBnKXwxAjc8zc/00LsdZVdehIaL1eI9qZtyap5GmVpF7ZJdkQbo7j2k9/o8Ztr6lwZrODqoujHSJK6V9bK0u9Et564zU+wWgftergJVAEl4m/D3N6/lD6Tni/a6bLzIcdcVjnfWLPAUBwAoj19NpxhAbe1VyiybbzF11k65OpExrnTpeyfXnWi2YKXmv6NMcvP6YS8WOK4pM7nWhyKetjJvO69p10oeh7Pv3PuQBq3kARIBKQ+MPYmymnbhgmxG/6w3hJ2A2Urz2k1DctVq7TiUCWnAReHSDqSpYcdQwxCm/lIpIwtl/dffgss5v+hhFs6NSNe3zqLc+wa/P6fKKBzHBPA6mZtXbiaJH0Y+5hMHtf92lFc+6I4pZ1q2XpI5Nr1V7em9lfehnp6KwZCFUTCrCle3ZgVn3/WlL0hiX3HqF/qGx1rSRBE7lqG2nGEQXx7BFJGNLF0vFVi0j2agV+lVqGOVlIxAjK3E9wGWVM0V7xAFGXQtxAYJ6qA6zMuM1AzlTqoEWcy+zkYm6Z2/Vn8RMHtpHaCW7GF05Wujcn0D05dR11MQem50GDiKxlzGighyGKWmfeex/qNBXelV1apol1nwDCUiSbC9fPUu70YI94kit+OnCdHe588u+9o5tqmvG+4ju2D1U0YtGzBJLwNtKIxTj+ycim3c7lWMz9gyNpdcRw85nQOO+UebN2j6KDuTy4XNxidtzFIcvo67EYGfl+q3WaPfQzFQLuOqvybvDRViyMxNwUidf7UCNcjzUMa0RFtd4HPTD4pR9pL0oOHG0XMOwvMlfvI/0tUl0nH8gKjxa10D+0pCAG9Sq76K3xNRb2QQ4PhDp7u3P+U7CY7JpR9qHasfUEAAAAxQZojbEb/+8M7YXuyASh7Kplen5y7UyO14JIrbok4XTbRQe3ORruR41lvOoDou8ClgAAAAB1BnkF4m/+AoGMvgzynj7iP1aLLDUqgN6Jo/suYQQAAACABnmJqRH+I4ID6PupU5REY3sf/aa003+ohQ9ie2tSAgAAAAa9BmmZJqEFomUwI3/3q7CUD4uvUkflAlpSidF3UDZUIJZsuBftR/Ot+q3GhwTn9egJDu/Q0u302gbAsB/IV/AFE9llDfK6lUW9694v3+SwMl6mllWP/WwEuWwr3bvYMZvmSan6TfqkNjfFQ2gIUyqpE9/WshNO04YsB0gSlJqYcVRMCqNuW9LOcW33er/NiW98OL0en+UGHVMigjfMRwtvQwEBx2TCEvqNHiHUZK47Ql9EBEwOC+9RX+RfTR+dz/gd9jggR1VC/W7S6VbwMD6OLZQ6pAbUGRTytAu/D3nUIHJpw2KDyU/GpGdKbj2WQLEoSWl/G1erWEOohpWvfDxHkGcHJyES1MZ1DvjVLlpUz1LfdDNETW9u5oytpQuaaon1mmF0VFggCZuK84NA/jjnYLjbiGa2mQVNuCz0xJcY1SlKU0HHZWh4xyJUcesClYPnBpMFzn1F47bfO2af9hq6jk1pG03lTPGyovR1gotc0zYA9fMfxEfYVjeBaSNZ7jIelcTidTkL9rr66l4XIEEMH3QknOAml0XZzRLxPdNkS67lJssdtIGefZbu54FuzpwAAADRBnoRFESzfiNYKja/0xDhfTAiTpHLhMdYxhLBGYN2yGAaE7v0DPPtKPN4N0Bh9SMs1vQXBAAAAIgGepWpEf4EHlYJarHuonGoEjcLs+tCeMaGFsuILlYL+uWEAAABrQZqqSahBbJlMCN/oMK3pi6Gt530QeupQ59ezKe10FpZIms3to+JEJWnpAJhBGOAA9pffTpSVzrDTkYei+dF/0XoAHErlseRXMtx8drWM8lgtra4QnL5SnQQH5ZVFNtwaHUoOeGKot2pc/HkAAAA2QZ7IRRUsn36T35XM0R9+0te5g6nnf972YKyQEsNo7mvZhGpibk4k6GYi5Xw1pjTrITOOCgNgAAAALwGe53REf4pxDGIcj/61i91Rdmi2jQszTK6GJmg9LxPIKqXtoMSsvmXH+U1e8ZQgAAAAJAGe6WpEf4Z1uDH/dTIseAsGQv4aVaLug7WXzVJtvrBkNRseHwAAADNBmutJqEFsmUwI3/QsGvLB67IEGcjaf9jQmmp2xwadCZqIc+QsmxjXnMZRrdlnOWXkniAAAAC0QZsNSeEKUmUwUVLG//1hfa7H0q8GiMhKKARrmw6Z4xdeyr7ujJme+WFMGuspP/+ISBWO3g4rVjHKMZczL9N457VGygcsdj7HL9XSQg7I9GyDcuezylwimiV/wFo/tnsP85jKo1N9u1U/c32ZRoJ6MLzOYkUW9bgpdQwc/vwurbPO6DqdSp7R4DyF3mA4Y/IOgh69Sz6dcgXGl/wwDmRo72yeNgYcXZp2lUk8XqHfEbj36uBFAAAAHwGfLGpEf4jgOwq2BPvs/N5sWutoQxXrlD0LcAJHkzEAAABWQZsxSeEOiZTAjf/7wc1PkbwRgr+fjnyNR8pfsXICQZIWSD29LO9OdoFeaCGC7kNQcI6QVUvVoNv//1awF0I8OAUiwMZgFzApKdTCW820SqjcXXuKH/0AAAAtQZ9PRRU8n31uIe20/lZp7IOycKeMuVxD421aM+Nt0XYWYsoR6FTF8OFoiDbtAAAAIQGfbnREf4ew+VpnMnXvUfVz1z4//JFTxQ6bZXAMaAhfWAAAABgBn3BqRH+E1CtCDqsru18CnQc9Zgwu5C0AAABbQZtySahBaJlMCN/8A8MxwHsdgYxkG6fD5jfkGVPerg1XibHNvThF43JXYYlxb2Qoy9V2WnWwkQ/OqnQwi/YBBYBEZaAn0hrUVg4yiuTQbTao5OXYZIcHYrYRQQAAAFRBm5RJ4QpSZTBREsb/+8P3/LU5AR1llSwpTA2p+sot/Ap2Yo/gxiVR9YF8xLoAHZjM8457AqBIx7NgPEKnQG3AmhGkbhUoZZDnZk8SW2DttxUzXAwAAAAmAZ+zakR/h60AWug79Ot7YB0Hk7/PuwF00L7Iy+IMt5Tq30K+2UAAAABTQZu1SeEOiZTAjf/9dtYVwOWMQEFcGvjV4Syzg11kPr1s9125wYmt9rvuf3mx7RYJ6jF6ErfUzdz0/3v9ae/ANOnW4wjLgPCIol0QSIsPYy4da+kAAAAwQZvWSeEPJlMCN//9fMnM4gfdWuAOqRQ2CfzqA5eQQ8IrF1fJShsXtjQZPj/+EEk+AAAAWUGb90nhDyZTAjf/++FR2EKSC1to8pc6YSO3kyM61vFVTIqwBRLCgvhTW2uELbv4P1t6j3j6R5JEnH93R+6ccuNOfX9oCM+dXHFLeM1xbdq83XfyShy+a/CBAAAAXUGaGUnhDyZTBRE8b/vgOOB3OQFcxg/ZQD7brmyDZG+0lanuFt5eWNnI+DlTWOnKhAAySzB0u8CurDZF8oIWX6RlKX6AlFm3XANEhpl6OfqvPFvE4O+fARgP5ChhkQAAACgBnjhqRH+HsRNxvly6e2k2NMUXC/vU7UPeiHPyDV2qBx4cRcnapKmEAAAAOkGaOknhDyZTAjf/++A44Ilb6dcZ4EfRuTLTuDvo3J2im59tzroCzqfsO8SsZ8WggmVkE0/TN30peYEAAABDQZpcSeEPJlMFETxv/AOiKyBDYvuQQEynViyLAUY1Ib85NXoM2W5LW9IU2aIRq+bew64gtHymArgzox7IataLMEs58QAAACsBnntqRH+HsR7W7gx/uLfAQSDqjarp3zTZOCfEh+6Eb54T4PDy1sS+MHUHAAAAREGaf0nhDyZTAjf/++s+OA7zprt9FalqM30xdIjlcH2DZkkFtbEKP2KnwPRDHCudTxGL9Qj/prFtXdQZd/8kyvQgWVbBAAAAGkGenUURPN99TaObeIE6LTTxi495G3R7GPmQAAAAMgGevmpEf4k5No92pOx4d3tVe/jt/92QHqWDJ1gLa+IuwyWHPUwqtoP8fV9Ym9imBVegAAAAVEGaoUmoQWiZTBTxv/vhUdghqkJVKiFoDxnkwJdI5EfJ1rXs+3IWgFyiLClNqrjfdF5LRESoN+/5vYqeA8+ZDoM0LHM+p62tsd2sj1HySk6QgTd2RQAAAB4BnsBqRH+JNh+VxVwUJ+Y3lV+ryJZRr2sjHdcr6hMAAABLQZrDSeEKUmUwUsb/++A44DvMfbctqjDdsKqNf0xOA39UYyLAzYIZyyzAIzc51s3jIS5ItuDvEtTcwkuqq/36WLfiOMPxF1mVVotBAAAANwGe4mpEf4EPa4byEQdMkh18OfNdNbI0980SSbNKs1ZKN+RKPSfGS6ySQe8jTiCjRi2a1L2bI0AAAABTQZrkSeEOiZTAjf/74VHYIlYiKMdyQZmvfSW5HcZkgM1RFiY3gEH71z/3fVBzC8liqsmHHacePA4zlFgRGs62G32VHnZQ/DXMUNXZSe6AM+37MoEAAABCQZsFSeEPJlMCN//7w+5d2niwPbxnXcPGjmeLv2gdbVqjCSUoEJNhCpLEDmpftWkeZV+eoSYNoxZWdF8ucDxxQt1tAAAAUUGbJknhDyZTAjf/+8OhtVOQEg0JinGp/3aOKBnRKunbXz9YVqPZNGUQc6qywYwzQClzX3GGGmt+n+j6fKijedTiog/HS3groeaf1ZLhN7MSZwAAAEFBm0dJ4Q8mUwI3//vEgdmkZAQXhaU5Ka/mwO9DUeQ8nbhvyYsA1HsjVV69Yl5uJ7Z+5cENl5oKK+xgTk4zv6uteQAAAD1Bm2lJ4Q8mUwURPG/74DjgQlevXVZC8R8NNBU5lXz/5evdLx/xJ2R3nKEWhnQTJ2MhCGZTRB2wGJ2TPpnAAAAAMQGfiGpEf4exElW+Y0oRl/FZgSZ/zHaDUN2en4eI/45oExvCn0eIC8T7xwpFCCD22pAAAAA0QZuKSeEPJlMCN//74DjgJBslHpy/fK99xLLgEoSJoPeLiwa05OsKQgsTdDRtpBJScyBfZQAAAFtBm6tJ4Q8mUwI3//vBcoj4CbBCQT2RExz0aNVXfJZgEMRjG13pvbjo3lZ6TLZP87/4osMgfYBX228XW1pfmU8k2l4pXWkF2ONrCPtLhUtLA/Hya2TerrONeACzAAAAUEGbzEnhDyZTAjf/++A9RB901G+fHcU/WnLCfcsk74GCTFIVwJ+BLzEBzmOXsONP0vNgTAP2KwhTXJSx3MerL/qd/oYBMUklPzb3D+iSqTSoAAAAQEGb7UnhDyZTAjf/++A9RCB74KMJGvAkjJTldXY5HMHa5jfjN8NHs9hBd2fEIGPYdqirJKwimKLpiuNnGPY71cEAAABDQZoOSeEPJlMCN//74DjgQ+tokgudjQu5NRNWobxjhH6Kh3ypTnRXjwJga7RHFZxm4LSd8nQrtx1GfqIwsNleIlTJwQAAAF1Bmi9J4Q8mUwI3//vBnCJ2LA71GF25ddxrRcaDbZ6JrFKdXAbBcwenVA97XR386r26bMwGidvp2PmbPeLJeVTrLs40YO/0dkENxamY3tNtYwvFfXlRKOo9yyo8A4EAAABRQZpQSeEPJlMCN//7w3v8sjICMILXR5ECQLilrI3AMHhNoAl8fjML6XnsQAd3GcNYf/0Vg/WkUZMoFlfh0iQSiC4/lahyE8PO4Mqsxu0If+cpAAAASkGac0nhDyZTAjf//APDMcB3mPxGGZH0fTP3ycONiP81zu+xkY8BF29QMHZMJZp5xyqLnicMHS675xoi5qplhfLQsVW1n9dVmaldAAAAHkGekUURPN99TaObeIE82Q23ingUfeUAohKj+s7K1QAAADgBnrJqRH+HsRJVgDz377tnoNimjhTZvPY9yHf2fmv0wsOMF+wp88Zp9A47UCY9LcsoRhIHvx+VqgAAAD9BmrRJqEFomUwI3/vgHICtV4/N+3TLBHjt+Fcn7DWjz8soUTNhnfkqkcMHOJu59OaFxhuVf+ti1/1ybHdmF1AAAABfQZrVSeEKUmUwI3/7wf8MHhAR1o8LZVW6pfXYDg2dpGT5Rvg+Cl2lQj+IQxqhhQEfFhAQwlKUGoZJ0yY1+WX5L6BuA3zmdozX8D1TAHaS+xXRLIKO1jwkjAd/JNceRMEAAABNQZr2SeEOiZTAjf/8ESF/UQffCua7E+r+NL8HxgRR3d+rdy5oOmnMWV4zw6CD9Z7mbT2yUeu59zFDbP46SRq9klvWZ1ZbOSW+2vZf9fQAAABHQZsXSeEPJlMCN//74EfgEPLRNeKMrWaAkB9b5moj+s74rojfFold/VwZwZnjRlAT49c0fDe4ZV3UBIGBK9CK3ytLN0Y42oEAAABDQZs4SeEPJlMCN//74DjgOFPapEipwNQ5JLf/iFkT7vTC27zO0O1Hhr7Be9x5uGF5a9SD1TJTW/IfLAT7Nh7DJMuMeQAAAGpBm1lJ4Q8mUwI3//vBm7tPFge4qVrTYzB1Rf4j8SwaiUd3Vn91V8IbddfFnSdsq/LMm7+bcOxpHHp/+sF3L+oZtRUH+68r3BQroy9k+62TGFClrSWs36FP/9yav//StMk6n2lq+8O7DNsoAAAATEGbeknhDyZTAjf/+8Ohm1TlEJHbNU/mWRKFgYzLN2yQ8jXAB9wCzOsfbVP2lU1EY8QOgEzKqXo0R7iSjvfThP3Ejb6DhTio3axVsWsAAABXQZudSeEPJlMCN//77lcOA7zH4jDF8og00F+0xFfHR4VDZOJrgOjozJdR/bpg8saNaDwaxnfv4X/eRhAQ/CW8F20/hfAqOk7ganOxPkKbRibKABEgisd4AAAAIkGfu0URPN99VLWJAgSGMMD0PyKXcxUT3cFaf6NQH7WvfiEAAABEAZ/cakR/h7ESVb5jSubSAV8FnD8x4gGanIIpXwVIvISOGwpzao9rHxOAHmz8XBzh60OGZXgi//qWtWwRLulsxilCnuEAAABsQZvfSahBaJlMFPG//Blf8ICQaEccs93WPEeRPsjPuDf4gDC1lzaA9vG8/MDesnCCrO+qwrWg0Z041Lq6MnjtnVfjvIt0qF1hF5tlhQF/0d8e9aqpY2UcQgjUy9Hd3+CTWS0I7CWSzCsCaZnAAAAAFwGf/mpEf4jgOwqq40N+HXrqPTGH7PlEAAAAU0Gb4knhClJlMCN/++A9RBAtNsNSCFkg711//JWWvGFsLrGyK1H7zXz+lTaC4eBFVxUDg/Op1j5vlFjWLyVGxDTLWQWt7v06Slq9bNE1cmz/5OS1AAAAMEGeAEU0TN99Tam3C0rJYqxDWOyBOLJHIckz5MqYD+DjRZpqPdIIw0nQ9LZT+KG3YAAAACABniFqRH+BB5WG0Yk+JGFIqTc/MKybQhljSSoiM7L+dQAAARtBmiZJqEFomUwI39P1UQmO/CHKa5fG9B98DyBd2q//xdlw3GXkmq/tlCxKXGJhBrqImPX+Tw62HdRUoCKQH1/ixC6DIupugrK3/kza+ntSNKpOnTIb2YWwrbEfjsAUf/8Csikkf00Yy4JNpq0NT3r3G73dUSTcvN2hSQC1jpvA4PpqiSqoyt1rg4K5nEb2X3Ta2J5m01MrHsJWKap7ArXaL/fag2Yxz1DnrBGNF0N64688FzTsfy3aMFyBNTwHxDcOM5VxEYpgPbQx0UoRaoEjPSOfPfZWW0CeOiVTmPWdE40N3Fq2/t5cedvB6NP7Q1mjnmICrga1HKXTUnIgZVEERRT/lQNzJZZSENnwR7l7OYZdbPaS+0iJ0PvAAAAANkGeREURLJ/NXH5jrirbiTEhY9Y/tlMwwq1venEm0Vym9wIwyi7kBVs/cNPG6n7cE56ATPcwWQAAAHcBnmN0RH/N1YmW42fk1rr5dWrvg0hpI+7/lxaqBQ9u9qPMWMe//elYZNo+NLjg8cYBUA/d4tfmi8Qz2UXcEo8Lvkc59K0YQAe1JOntXnqOjM4McR1Gx6UrMCz14vMSat1IO5CZ6UXFyvmTyeshwVAiRewmoKI5gQAAACABnmVqRH+HrQBVU2y7OiaX26aU/4ah2bLxWm+Y90YzwQAAADxBmmdJqEFsmUwI3/wDwzHA7nIUz9EUj2ygWhqxgWMh9LPrZxgjmvIJeUBCqf7Hl5oCspr9BQrYw3VwQWEAAACvQZqJSeEKUmUwUVLG/8dHUh70tVYbmJAP5Yav75x8KgRRrVZqGBAEqyIrK6LRnrpZlSZkXNqRr0UFWwbqNCFOpYqLeS1RQ2PyEkC+dBVplJUfZiuhOQAvvreTkE7X7a6IRT6oN2yIUD3gEVfN9FOVcWA6Z8kXNa18/OFVTCJp5CwGEmxWr1WMNJva0wrVNHOS92/kbWx21BB21E9sgNNFnB189Z9RulpZSpgE8SAYOAAAACABnqhqRH/N9J+2J1XcPpDCijwctlpIFM7gAfQRhV49uAAAAHJBmq1J4Q6JlMCN//kgYHc058NIdrA4tZt28lo3D3meW1zCAqwidHRRteHMGFHB/D9M0MCX00dwKKmnXGyIjSiIo/6dF3Oz1KEMYwbbPPtUPovm6LbF9t2EyWlXZ2AhyytqaBHYUktsnfMGjBar9/t8GUEAAAA3QZ7LRRU8n36t1fJ1HlQISEIEfpaQSyOL40QRE9Vs9zvdbKrtZIXDJfPzuaeN4LIza42Lez2eIAAAACcBnup0RH+HsPlbmN8Ea52A1bNIOYqtlLjflwD1gV/uybFUQRKpUOAAAAAdAZ7sakR/hNQrQyxiUawcG/LrF0fwKrQNveY8OGEAAABTQZrvSahBaJlMFPG//APDOoh1u+/wA6/FZcjU/jvz0FUQ/BieNWgawp32ZtUYldJMEVjFO412vb48n9u19LQOw5J+xi/1yuC/hu68Dsm1661iVvEAAAApAZ8OakR/h7DkMUQbpFoFWKV9q8E5k5LkELQbh4pZCjhTQyjdPLum4gMAAABAQZsRSeEKUmUwUsb/++A44EJXr11WJCw7ECvX60Zrh07UIdiv7aDfa0prRgD6jv/0Wc4+K9FDJsKP/KqGWZJW8AAAADABnzBqRH+JfH18ZqS57WtbAe/IO8nxI2HXbfpO2Nkfx7T6UrglP5FtwO/Jq+dpJWAAAABuQZszSeEOiZTBRMb/+8dnvoPCAeCb2/rWCrRRMAZhYjQsyCGOaWbOiXjQ45Hd4KBD69SXBo/G1jQPyXicID3GsLjYV+7vn2Od6v4G+WQqMHMoRfr8vCzfPrEsm0u23BOlwjpYpIOmf//S7gManfcAAAAYAZ9SakR/iOA7CrDrXGjeuosMdCPQVofsAAAAY0GbVUnhDyZTBTxv++A44DvT323Jx0SqVNJ82IBpyyo7DXCG+RxgT1NMsMLKQax8K4yQWP/71kziHHgGscwzo5L2jLizRX09ig4TiJCfNOrU3b1tM7YaGwugP+oXNvCezwnqQAAAAC0Bn3RqRH+IvbTI3y+IN9Ohxp8NHgN3aoRoF0CyJ7Kizh4lPfrWVzPrEIihMfEAAAA1QZt2SeEPJlMCN//74DjgQ+tokgude4jclMN4O+jlwgqt2WyQlDcvOp1uGtyVo8FNYf8DxygAAABYQZuXSeEPJlMCN//7w+5d2niwPcY4fRW/HK6IFjrt0h2pzlDEeWsYZI51tgKX93CPYsiPqPq3X5B1aWcUk55oKJGxvjXmYgG2872vqJP978jdfl4d7IIZYQAAAGpBm7pJ4Q8mUwI3//vD7kjX2UR2WhZ7zwuAd6E2YZpVsaXqf/9ms8pdq2AcVeqHU+2duGYw6XN6gJJPmwjaRjOqn/963zMAo7BbLbQ0xWrehzgNm86QkFh5F4CJPtfJrXi2qfvexEjl9SQZAAAAPEGf2EURPN99VLXrSyKjUixfkxdL90ti/7QhhH5aJuIgfKR9AwA9eKagET/7LDiGl5JEDu2HE71HainiPAAAACQBn/lqRH+HrQBVU2zahqo4JDrA0Wsth2bQ6i1pu5g6vv10b/kAAABIQZv7SahBaJlMCN/7/bCFg4EB8mDknVh3Li4nF+3GJT89JkDhouv2updeuOWXeNe5p0zRxNBaBS36FPmUKGpjSDMCDJPbtDwsAAAAREGaHUnhClJlMFESxv/74VHYRK6xf9nuJ/mW8vBkEYktQx5ng60k0zpkEVATXt/+Il4+t1Y/gWE4ff2KkBugUa5RlCmVAAAAGQGePGpEf4k2H5WycQ87dp8VYGfF0/IU9Z0AAABiQZohSeEOiZTAjf/7/gOFgQm+EEcUp0ZlHiPiNyhbnVoqZ+8senN2fT/GHn0i0DJAHVHDeAAPPdb6+08ywf9ggCftEN0KTRcMP6zvyMME1GrdoLnzWjL9J+Li8Z7zUjdev4AAAABAQZ5fRRU8n3sgptQ/K33IpJWP93HYo6zTTQPcVWgvT/tt2dtXG8r9XMctnwTJYMXBZrnQX8wcobNb8TVLLKjwQAAAACsBnn50RH+HsPlaZzJzUbRAKINcL+9T91jJ1Gd5JZXKKMYQPHi5kXraXljhAAAAIgGeYGpEf4TUK0IOqy+T8lgawG4WgIsZVeplQdYOkTmpMeAAAABJQZpjSahBaJlMFPG//AgTwyAUQ/9Od+649emXNJgblrSImgV1dWIgYyUeX0WkGz1oPOdc6UKgLvfxF2udL9r+3vyOJ3zOsvfmOQAAADQBnoJqRH+HsOQxRCCMCTKsMv80pUcMEWMaBnG+bZLo/o7gAggPnWMrK71nR0MaLVpDsOS6AAAAMUGahUnhClJlMFLG//vgHIBhjoEviMguc3ZwqTYQjvB+dEBD+tTd/LD/U79/aezoU08AAAAyAZ6kakR/h7ESVYA0yC7NG/OtKYdKqWLeGmJMPcZhbIRuFc+rxG1R7C8fxWQE6y89g50AAAE9QZqoSeEOiZTAjf/V0GkzH5jjHogSzwV8Ur8f+Ec2ZqF42vRmT3+0XBKgmcsowSJtIfdoeCmqACWCkOKlAnt3SySxtVHTz5bNRvdUm7O7h2Utc/NeW1ORiPi5ahyiTp6pzvTB7gpoYxSr8OJvwgt07bwjElPFl2/DXqzJ7mlN8Ko4FuS3th/HSK3m8xjdy9Pd2HG/G/ifzmmcbp4pvEyDSaJK1W2uFie5GND/WdfCyg3jjBDikafVsZZsB9db1ELXgOz6/mtW1PO65LJeMpnL03zTJn4LERpRc9KcXLD8eihY+pbPtstJ8ymI7vFDm0fdHk9VglqCfCuOm5LOR1Xte0qY4v0ZF497FEg6J2gXBmk0KeY7+LdmgHYEcz0qRZkyA93gRlWQCcg2bdHG3R1RsZFfPnVqyGpeaMJQsD8AAABgQZ7GRRU839CfjYe8gADwWSBEZ0dUg8ycG32E01qLxeAqrHnZPUJ5Lyx+8BUnDWj0XZtSKfpUSpAhT/szG+CsJq3bkVF7YO0bZ1YvXCQjDz1cgeQteorKqspdc1Fi4XCBAAAAGQGe52pEf6jUd6g4XFW5qqvsTJAPqVNfayAAAACQQZrsSahBaJlMCN/UtXwfA78SpsbvqlpyWOlTh6QrFvdRJjjmVnCQz861RBv+whwL1CGfpCv6GfEMxGCpC9HYQQUE0dQXgCNgcSS6mD0fn/+UCPvjhBVhxi7NZ4U2QtlGo9RKKkNi+HxgY6TR164zEIP510Utju+3WmWp8LCMIzVG7g95lqRKQwTZXOyoblZEAAAAPUGfCkURLJ+megsfuYeXW7Xti51RtkOkbffzL/ASgnb6g1Hr5xJiwX7/7XlPkdirFvrO694xVrxkHWvRxaEAAAAdAZ8pdER/gGwq2QxVj1MPQwZguW2S81h08qjXqxAAAAArAZ8rakR/3Cspb5/QadRcLmNZSh6g5dIyMVdtWcxCmcKN8VSOxyyA+VX1BgAAAD1BmzBJqEFsmUwI3/1ikeaEjHOd45AIQeOuawfhOqyYDx3N0zAkiYWZShFSr60ONAwyK1zoY65cTb4j02w5AAAALkGfTkUVLJ99jGQnHgDM8ln+aBB+E2IcxAM5SUB34rzqVNrLOXd+6bM2xuY4uoEAAAAbAZ9tdER/gw4tjgY4RdE2+ViqVUsAWJkRCr/RAAAAFgGfb2pEf4EHlYPO0DBg6RcG870j464AAABUQZtxSahBbJlMCN/OlXvvG5yLTHBTaPrAMVQaCfuKj0j5S7bXX6yX7qQf5RBo8No2FXkM/qWMiQOljutkadwskYcyP7bS28Hd6oB0nybXNXLow16IAAAAR0GblUnhClJlMCN//R5j7oKUrIUIBCbU00bLa+xJg3Zc+8xHG2utLSrY8Lzn10Z4sXQk4HLrp+BfVVniEmXvgjPxoP1/+SNTAAAAM0Gfs0U0TJ9+tCxJ1Gn0jDb8mTrU/wrX4mC1oGx1U6eFX27MccNul90QPfzPIXxTBWHVoAAAACUBn9J0RH+Jee66QMi55jOocyx//JCSfKN8+wOm/wGvvJAu+yXAAAAAGwGf1GpEf4TUKzbRXPU6nraxj9iMWSV9IgbTbQAAAFBBm9ZJqEFomUwI3/wICGIsD2OziqyuzNNpqzvABec/jYnGKj/psjvJAnKx6B88nuodchjeQKU8/Sf4fxPeI3Bur+63wb/oqD2JcdFb2ce6gAAAAFFBm/hJ4QpSZTBREsb/+8PuSK8cBB148viFWw/NHqyGs+rm2ceC8lLcGprCfFWoP4uRxFnVQBO6X5Km47W1B7f7HR4MpVRMW/C/8Bu03QteuBkAAAAjAZ4XakR/h60AVVNs3VhXxKqVaJQ3kpIsIMaGLSdIpeJrRq8AAABIQZoZSeEOiZTAjf/7w6KlpQgJLKb7CVQpIKCX2GtUUOMn6NT46JbGsSJDqhqGAUlsdEP5mjnTYaodhrLt7P8/yMykEazbfVizAAAAM0GaOknhDyZTAjf/++AcgTWnNhxEl9OeuJnrXVsKR/mhF9lVb2NCffZ/IdUboGwovrVQewAAAFlBmltJ4Q8mUwI3//vhUdhAycJLcpdZgvrBHbuS0ELA2qR1GWV/sKmX9lqbykcrOwXH6Qp5yOSFSCovTOHfR4H66+O3WFhrSEGUCxI2T37abEC4N7SLdnmMQAAAAEpBmn1J4Q8mUwURPG/74Djgdzkgaa9RAdjKXG4f20PXgc0gPoX4Vj0Ucek2XSjMdrwZw3mpyUgcp/uoOq3e0VHPfh7Zubi3pdTi5QAAACoBnpxqRH+HsRNxvnokvr1IBRBrhf3qfussoHJ2FXwia8L0og/7L34YHvkAAAAxQZqeSeEPJlMCN//74ESoSq306qyBH0bjlOnxLGbrkGJ2Ws39CdVP/9jJZ6U/XdORQAAAAElBmqBJ4Q8mUwURPG/8A6IrIENlUdGDC2rGxYKvs0Zv/PoytQYJJmfkiM01RixdLevQCF9AnGptUZtShDxvMVuLKVNUIVwk3Z75AAAALQGe32pEf4egL3/1jAGor8/AQZ6N6eggOUr7LqhLAXV7FfOYIjSYAD8KRrCokQAAAD9BmsNJ4Q8mUwI3//wC0LvCDOPxGGZH0fTTrqCs2EFRlT6/L9YpDrTovL58/OT2vGOyWsqXu1aV/a7fySUgV4AAAAAeQZ7hRRE8331No5t4g23Ywzxi4929Dj8IQsijl2gJAAAALAGfAmpEf4k5No92peKr/htFkwdNuhonMp0Ro1YlnQcnCewvp2SzmqBS8yiAAAAAXUGbBUmoQWiZTBTxv/vhUdhCtFQj3CF8I0Z8KP7PO2SHsB6lgIAgPVTE9oxFKS1kdE31Ian66qP+yqf9hh734mZwin70dY7QE5chfscqW3JQE3iewy17sVw7qZM5wQAAAB0BnyRqRH+JNh+VxVvSpS/NWVqjqTRtbJ2NrINTKwAAAQVBmylJ4QpSZTAjf8EwsBpzW3+EphZT+sMjYK72pqFmMagiJ/6RxdCOrdIwWTYN0R3uDTSb/XRdzEteeehW0X8cK0bfRzlGLg54Bih5AVBjp8JIdS+frbm7i9b7XCr5CBhzbRO0uWlX13WMTPZSuvGoTUMrRVgG628K7rQSE3qBkBMgw+CcCNefjGn6Zkx+q5pfh8Mt0QqMzqj7mnSuGudADa/xLYXydiZCYQ7/ewnsEpGkOK9yVq6rK50LcAUK+I3NgBMWmbuoV9kh2ewa2IxwBTrP8/IrrKTnIq0axfEx25XaS1gBYRb6kDVOAPsog/1jZSwYZ4GatQjuLuVTjthxX68wVF8AAABLQZ9HRTRMn9O9y352oflZsxy6BAxflZu0pxOibLVcNnE+wS7cYiWlcJ/txGIxRrd+RdfvjrOjPfDoy3gMEaOA1ktKWZzpAcG2KsiDAAAAKQGfZnREf4ew+VpnMnMtuJjFWqP/yRU8Vfsm/SA+y58AHgHPfjY6Ic5gAAAAKQGfaGpEf9fFlYagcVz5uBzD/QIPWYbzU9vkHWdHRQQyeCK9nXQ4I2EsAAAAikGba0moQWiZTBTxv8gMsVZymC78zcwH990XdZZmZzJ4wtrIvyzMXIOULxsiqy3MuQ/I48r4n72WrGhLg75YH6WEsBqONLX6BjtJ8TjNYAtjYKLrQKQl6vGS9ZfD13YsDE3Jvu673GbAbsZboZ/13vUZzF8zSZOV9IwR2roprdQI0feIeu7cnUYSMQAAACgBn4pqRH+o0ZioFg2l6PKcHO/zygwvcEiwx436ZjD+vu4Dcnlr22pAAAAAQUGbjUnhClJlMFLG//vgOOCFHL75XjgbP/nLa3NdHqfl9wTp94oAO8T7lQOqg3O0Mx03n3NRR7pYlze6hVEf4hs4AAAAMgGfrGpEf4ERwY+Gt77r+yEJ4SvPAYi+RMSplok3KwSDGQpSuB95O2ZN6aV0x9y74RwxAAAAZEGbr0nhDomUwUTG//vHT9FLCBA2nxris8BwI6TDKNzQGtw992rw5prohMGR5qiMwrOY4KxeyFVh4Z5nsHZDDLBJCMqHd6qBQyfBqvHpiFjePAt/3JK/Ps50Y7xjr5QVjDmLLnEAAAAWAZ/OakR/hNQrQmu509DykZGMqHC4wQAAAJZBm9JJ4Q8mUwI3/8evureZ7lc3KXdmxnQXgMijUWF3ZgneFpdm9nM3WOKWFmTpBzrecoyXK1AK3/TzKgoznmiqaxTbi+HkZzv39k56Xd8Rwloi2NmfZ/xAvKYkrgc7Yj9489acf/vZrf+rHiwEDuf4P34wfjU87tMB0gm0+ejzliNtxzSXYbeAmmP64RveFZ816fJhWEQAAAAuQZ/wRRE839BbCo5vTqlWMLNw7JP9qz7O6uD4nu1gs/O9P9c5x1S1ytTVzTI2+AAAAB0BnhFqRH+BB5XFIVzyB9wJwDYiu0A8ZNtKI6DGeQAAAEJBmhZJqEFomUwI3/18y04SEQRzJO8XoEYWKQTqvMsTVgq6JgxNfx+qCvDs/Eb18miZtc9keAkyj4LDbnj+9YUtWYAAAAAzQZ40RREsn32aKglvf2C9/3iUkf/5fFdsKyz6B6M9ZyCIDP+BGGOtWUef8hos++ANyp2AAAAAKwGeU3REf4exAhsvjHgSnTXrK0yG7hQMOWbFoITHQQMMz45J+HIhTgvG3EEAAAAiAZ5VakR/hnW4Mf925RATlfTwJxSqw6b0//nC+VfvnPeXgAAAACpBmldJqEFsmUwI3/vuVw4CBhG+4lg3Es59Okupr9RDias/wbkJjrkoqxEAAAAuQZp4SeEKUmUwI3/74DjgJAb7CXLbzf5nq0/dretTIIHlFI1cAGN8WOxHbKwvXQAAAH5BmplJ4Q6JlMCN//vBGT5YQEgiNI4YfqwxY7pw7Y2LpcvMCmuuwZiEyCprEorTJzgFnI1fTLxSf+rSP+efSERrkuP/Gkp1S5WCPd1cY+R8nU2/Ecise1u7zn7FazbzJRLufU5Eeprb8Vk57B7YxKp6pppaYCg0Zui1dDn3N58AAABvQZq7SeEPJlMFETxv++A44DvMfbdU1hKfbdntKLt5PPLOQy9ifPEv2J3HrARgxvABi0D1MxB+8N2XXdutIkzk+EXbiNPhxpbHQRvoMb11pMEpF0dzIxjytdVY2pHHvt4vqc8m5kctbxoCGPrGzq7BAAAAOgGe2mpEf4l71gLl090r4xvvAf/rji6zKN85FiDpcbWZht1n1na/FcE7BBvTXNdoIy6yLTUk5BdVxIAAAABWQZrcSeEPJlMCN//74DjgQ+tokgubHJTmcS2bnfPkzdAbeefDjnKfV7MFzXTULGE7gZCZQWfJ6qOVkkF8AtxRAEmR69//H0c3PlHntplsbW9FUKAYTRcAAABbQZr9SeEPJlMCN//7w+5d2niwPcKLWmwiVkBeI/d3EBe4uOe41g3mCl0o4kirBlvctec2F+QjLJhaKWtN6qGqc6uLOLJ+QX+s7b6lWU0b6H0eQNGAmGh+ktnyUwAAAGJBmwBJ4Q8mUwI3//vDoZtTkBBrgFeRatJUfKwOw/xfRodt9aPbXcnC8nF6plsAvzoY3VEtvh6EDBUsUk66y8/erm9TY2A5NAgUOAFC29sW8XnubtkD7UGheuC3NYe+ACX3zAAAADlBnz5FETzfgcr6UK161fjdbXl6k8H/DlXPiuIwC6b9ntAFIpHPdd+z6GlVWgcCRI2u8oYOOGNhC4AAAAAlAZ9fakR/h60AW7YGdQtuoFCe7a0uOvlc1NbK2CvlDazwoRCj+QAAAEFBm0FJqEFomUwI3/vDoqVTkBH89pCmTZspbW1QWm3JfZZ0UQobs65Mf5LlrQRMt+qH8hg6HnjvEpO2ssyZ0Y2DKAAAAFRBm2NJ4QpSZTBREsb/+8HrVh2DcyHMmT3dUQJoZsKjKvT90DkkLfs6f0DlI9Vo6Noy8G4U70dLwbnSCcJHA+s7ZIbEbAz5i2KHdv5BjXfmYaICxMEAAAAYAZ+CakR/hNQrQmu509OcsGZ40siB/cfgAAAAT0GbhknhDomUwI3/++A44Hc4/O2JB6m19DmjiM44bKO1947whNbBWqUtCS8/5jm7dHXzg/GKHpWAQzxNsp/1CPUSWvt3lZLPoW2mRgyciF8AAAAtQZ+kRRU834HOcQ6HnuhEEjNaVKp2xQF7WlSmlAM2ebIQcM+HcEqhiiFkUxiZAAAAIAGfxWpEf4EHlcUhXPIKbaQkC9FdA677TaPsDtJppmm7AAAAa0GbykmoQWiZTAjf+8HDrHeCOtIL6oV0eH0f5bCxgOtFfLX4aVpxgeHP3XvPB/PFVi0PRBD1cFMejAR7FtsN/v1VPGtv4Q5R6r3IoLYn6BFezGJqahiRYmgCTMTFmr0X4+/IlITNyh6ZhGBjAAAAS0Gf6EURLJ99PcOlKzJo8mgZiaFQ/cRdc7LruQ9PEd+7B2bhHdjrU2WOtGZMYKSDh35QT0pidHubWQAg3xDn8inXd71eAkmkAv1mZQAAADMBngd0RH+HoC9/1Oz3bWZ3fidkWX4f6J+c3kNLc5E/BUMQnT52L5poTTPp/CeS+I7NQNgAAAAqAZ4JakR/hnW4Mf925IfVrMHCbM+zJMcVzMO8153EB9P9wj2YcBm3d5ATAAABFUGaDEmoQWyZTBRMb8Em7juH6SagJhLTIoH9PuCUpPS4Xhl1dciUnk7fRxj9of4QGmTPb0WIBMBKQd3a1O1Yc1N/4kQ5x2J4nCSJXbRNmnX5rUdd/QcMYbapL5XWBGEC8Ck5/94+EirqJFZDiSW1qDre4E/cm0Cn57dgTQx4ezV9WnDqeA8qNeZT1cxuSXLNrFWIMllZp/93Xxk4j6Boviamv91IgB12vJyIkQkUqR+e7+pn8A3J1nM0FWyTrzkO7GJqjke9Z+zeEFS+E9FCk1HBH1Z9GwPLuDpVRLSp8hvtMrExkYwEqFJ7b8rI62OQzrSoZT3+EHwb1gJEd+fqytM5WmU4uf9eEg5WPacy8CyFQCjn58sAAAAdAZ4rakR/3Cdw0sHnaByzNI/DvCcZ5/JboTOZcBgAAAC3QZotSeEKUmUwI3/IDapfcEaIOp2yhhwVS+etbAWw+lDUW+XfaSyWSVIb9I+3RUZO9EZN/RG2xvytqz5YiVBaQOmzJiMDTSFB1apCB91XRO3SZTj1OyHGJzMKmR0UueNYpWGgXCDTBd912jNrkePuUoCV/pwa+tGQblTlC/7Mntowg9/HoHjvG3CnAA6FIBd1q7bFjLrs2BRVwa6HPk3/SE5LLvUOVOCcaqjJwOK9HolGgmBx/uLZAAAAa0GaUEnhDomUwI3/+SpVI7R3snQyETVUSNJhOfrvBOYWFqawEPOShEdffbNNhj/UU5hbs9Q5x/bRP+L8yTojPljMi35wxQSoHRofUyY5upTIiGlJUlsx7SzyXtNALaPJF5KQsK7ukuM7gVr7AAAAMEGebkURPN+Wev9EAG9OqVVseJN+Srxy0ekOaNhA4RLaQQxTYVhStfIpPBuX4Jw3IQAAACwBno9qRH+BB5XFIVzyCkOiEfbn4C8Yo2klpwwoQ9Hc0Mg62lIXOy4KcuI5ZAAAAKlBmpRJqEFomUwI38epbng7UbhTYq8oWxUQirssnQMSb7P6LX7rQapmaLnTKX/YQ+mvBQNdrilXpDana8P/ipoHFw1c/59iIYR8R4DQJ3J7cU7mE0EUO/5HJiOmapC64SNTZLeKkcj/wYvW9CxQjNGOHbzdRsVXG5vsdQcdxC46I2UJ+EPeA96UgnLDLarb/ye4UGev0CmekJlKpnTVesJsW30eCABpwrUgAAAAPUGeskURLJ/TvhMx1yLmtTPxEmDZ/zbA0Szadw/YOlgT6D9Dej7fdLCUUYjrPMIemsVWLNdcv1K6VaGlVIEAAAAvAZ7RdER/h6Avf/WMAZTJgkbiUU2GBEoThRTgMuuNR0WaohOWxAZ8OhULFKorD9YAAAAqAZ7TakR/3Ck9CeH/dWEJcswcJsz6v+24f6f0uV9+xslbm5g70jXAFOIQAAAAWEGa1kmoQWyZTBRMb/wIE8MgEocd6yv5CokI+/G0GUpxgPsND4Bh5ENvi1KT8dhbDutyYfbbj0VildYU+yRxq6FTnj94kjyaTRWbO1POXfK9v3D656X4I0kAAAAtAZ71akR/gQeVg87SZ4hSEgBWkj+ZBwvgo+Jen45iep+n2pmujJBvuAKIzNpAAAAAW0Ga90nhClJlMCN/++FR2BweCG1oRIssFvf41wc63i/NokGGJ2DfzyRKo+/2yaVmu/JJW0gloXeQ5hZa5thWPWMobrf/C7m5IKYTrfpZbz3D6/1t8557AoqeOYEAAABKQZsaSeEOiZTAjf/74D1EHaAmN4EWir8LxdLPbAZPO0M2+PY2QoG/X+gqbY2OfphX+3ikmCHy9qBRinr117Pf3bBw8DsA034PWVEAAAAvQZ84RRE8331NqbcOYV8ovUPdy4IB1YpDWJQ+wSvnfZOYkxxmYn/OzBmtYmp2A4AAAAAeAZ9ZakR/gQeVxSFc8i68ngUmW/lZdC6KspEQb2WPAAAATEGbXEmoQWiZTBTxv/vD7cBveCN3NVCXGk0YJpTWI6S27wxEVPlc6/1wnjM+kCdJnnhPo14+G+H7r7GYvUMBsla4e7HrvMr8M+H/u8gAAAAvAZ97akR/h6Avf+N3Bj9G4GRPrOabGjKMwWGQh7C506fOxYn6ZCu0xmofqwUIi4EAAABTQZt+SeEKUmUwUsb/+8WeTlkZAQajP8sPFAhMqTrywyxwavokyr3SLrM+i39/I0saTn/n2QD6z2f0Tpf5o4KXxn3mPA9D17rJX36vYMwHwHjQhLEAAAAoAZ+dakR/h60oZYAy4TexgRX9qffNaY7PWMw/K4c8McCTpcC8j63SIAAAADlBm59J4Q6JlMCN/8oIQcmggO56mz4fWIMz/u1LoEQTzV5bZxKWzyYuNw0VUour5rIsDiU90umHFkAAAABPQZugSeEPJlMCN//74VHYPx3QaH7LO6y4C4GizeyIGkLKZnJtuowmRMNw+Y2sUwAe7deUwGTWQ/Q5vriZGTkIsv3ToEeu8MQHLbiiuEa7gQAAAGdBm8NJ4Q8mUwI3//wC1hCBKrZCFEPTQr0tZ5jK1ORrDiuY02WzVUIDXNEWNY+1UnXo/FFfdn9MboNzrW4gnI9f6iOOt0mZ3nsie9H+/u+vO5CY8Fj2uueEyKX8U3QhLU3/TcKrIh9gAAAAHUGf4UURPN+B+FaPyFFuXZbgYj0u945+bqtESVlZAAAAKwGeAmpEf4i9r/hMofsoVdV7ULdGYi0qwviD0eMbDVTnrmwEwqTMGF7q03gAAAAxQZoESahBaJlMCN/74DjgiVvp1xngahySXAFM70j0RNfE5vodU8e8SsZ/26Vt1OWDKQAAAEhBmiZJ4QpSZTBREsb/++5XDgeRcDFSaseJcUAQLJ6v6GMzRDVOjaHtSG54x4li8aIVzsEHeZQB/Ro7xKkMpXenDvr+1lP2ymEAAAAvAZ5FakR/inEMYhyP/rt01Hns1oE+sW1Ax0NveAZCmuIIkSUdWkHaVhqEgSKY0rkAAABIQZpJSeEOiZTAjf/KCEPOyTOHP3Ga0L1Cjd69G+Yy89VyBy1dtUAy8MuVk4n3KI+TABbfvtiT/pPhtktWSGJDgIe68yWWzEpvAAAAHkGeZ0UVPN99TaObeIE6IhnjGIo8WCgt9DhhPUr/BAAAAC4BnohqRH+HsRJVvmNKEZfxWUIT/qJm6pPNGuzqg2ax/lk4U+jw9QOaXPUFfo64AAAAj0Gai0moQWiZTBTxv/1jexkZY1WGJiT87bz5jKDKBIPIUpqiOKoFqMU9xuIH4219MWUk+QJUsSOay780Yn710gUCq/8q5pJ2Qna2z31AIaW/7r/7nw3VuvmUI84z2bsgf+2HLYAKweSx9T91NAbzXY1OYhhH/Rhoq7G9wO4NOvgpiqkSQhrt64wDHDxeNgBhAAAAGwGeqmpEf5BXF/8DaTeNmUt5+Es7Qo54X0Q20AAAAGVBmq5J4QpSZTAjf/1jexZqwShLchuDwYT5GB8R7EoFbWKry1qGHy2qItJHoYSzhYHpftREHwxHK/JzKBYIGEaS6NFr4w6720GZ4EwFtOtZcCzxJ8lwWGRhMcXbe390w6J20G6KwAAAADRBnsxFNEzfiNYJXadE4qFM8TYpcvRqNU9B/tg74vA6eRce+4rX2wUXpty0yO3D1hzy0b//AAAAJAGe7WpEf4EHlcUhXPIKbzh/PPJ8bNtdgqDOQTgASBJuMgwAoQAAAPtBmvJJqEFomUwI38HgcMuPufQ+1zSgVSxP44j8Xw0s65lEdmNlEJPiZc5zpqLY2ASG2JgZLJLCVojwLtyVBJ2Djl4lz8/HuoMrqAhvSAdjKQWYwWvku/PylYEHSBWY+TA2i9oIGaNmjGx6LG1Kj+E8emazrbQW4uXowHf6FzN0wi4FUL7pYGixlMvdCCV20iEhjFpxLFL7zMrQiJXs0FePfh470wQo3VD0x5R78tx3albjYGEvRrqiobpY8RLCYk5k9JRliWcwQKwaEOM4PxOill4WkG8tUs4AocgZR6pVN8VW+TjxJGdBFTt3zP5oF1jV/E0tbkUBVpsIQwAAADRBnxBFESyfzVx+Y64rGY1MwatWM/mzoG1t/TCEbE6JSTBCJehjrlgVbP29VWiv1os6uGraAAAAcwGfL3REf9Pc6p27yHwywlkNj5DchvZ7/p56d3iHqPJJy0TZlJpjk/4Pz69n5WmRBfZn9bTbyeeOj4mF9//0Kaye4vxJwYkk6gt36PZDZ4v1fnhhsDUliUbwZt3CzWRiFHU6yepOVv1u7QQgijOhrLbxJsAAAAAmAZ8xakR/h+Qj8Y/7ryVtMCNL7ETyKPTlzq4JQ5ZB/FO0LOIakpEAAABFQZszSahBbJlMCN/zu7OUTa9XLb3eW75wn7Kv236GCYDyaFkEi5prdIBVAhBOujVHZ91Dr3M5z7YWFsFIPaTIAQW6IvvAAAAAcEGbVUnhClJlMFFSxv/8EDLdkB8H1mLAmhRVDhMcY2mJZi7xKwfk8PJDVJGGU0diyt/1T05IeVjQbyr4y3cfABc+R55miDu8jv+ke/z5Wlghc6LtRzgp8YEXwp4pBzCx8KkBF/UR9Q5Fhyu0T3/n7sAAAAAXAZ90akR/iOA7CrYE++zmxFeP15r8ZlkAAACYQZt5SeEOiZTAjf/HqZ9anCKYumtLnkGnytBtxceH+Tab7WDld8HDYn4TZQTl+0UF3pmdah258gjA7l14kM0QSfHRYvwKPR5YYx3eazaokeo5MLOsjv1A7a/dr6fcvO8RrTfuoTK4JG2FPtdS+FD/W7GxeAJSunV30h3c7FLYa2Yf/MWZdi1pjZRZxjRsrX+SukflJqQC/GIAAAA2QZ+XRRU8n6Z5caH3IhlDPEGgwBap5BKO3zzvTuKfuLPQfl8w4HkfPygPVlVEzHop4sE0RruBAAAAJwGftnREf4Bxg3qaOZGuOib2Lt1Wma/ikk81jKfwGZSD4w4Dnqo5LQAAAB4Bn7hqRH/N5HsNDLGJUxpLAK9mT9gX6FWVFLp4DnAAAABhQZu7SahBaJlMFPG/8pGuAUUgcuCRH/OwAPFSZJ3zCRS17ZFc1C9yqHmc8/pVK7IEw/oV7yJdMcMIavNsCKHsF1NBZJBQvHr63RBRk+tra9PDpuRrNWg1DBO3hQlugGNbeQAAAC0Bn9pqRH+BEcGOCVJY/QJGN/iSBG3l/ieze7n+BOFxVwCwANjCJKx4MnZIv8AAAAA/QZvdSeEKUmUwUsb/++A44IUf87s3whHef/K9eDWAMU8ldp4l040VzlBC65n7nr79QcRWreur7WcWsMhEQu2fAAAAKgGf/GpEf4l71o92pOxqkP3CQfV7f8+wwfDv0Vwj8bTm1R7NQsii5g4zwQAAAGJBm/9J4Q6JlMFExv/74VHYQpILXs6KXOmEj6xleFO7YS9ET5DIDmczDCW5+M6UZA4DyHmCTz5FoaK0i4R5P7y599cypfH8m6UCrK8Bf/vgLazNz1wC/16tdStMHxBIjZ4P1AAAABkBnh5qRH+I4DsKsOugF12qomW/SycjPJtcAAAAZEGaAUnhDyZTBTxv++A44DvMfnbE2rCH23XL0djoBAvFeci3xno+OtABO36XM5KyvnpfzxKvle/8J3308zrG6JVwhHh4rsYQUFRysfHizQb86GKutGNCt+OY9sPNteGFKlv4dXcAAAAqAZ4gakR/h7ETcb56JL659u40VH/5IqeKyAkelvuqVBBk49Saqe7VmZuAAAAAMEGaIknhDyZTAjf/++AcgTWtokgude4jcmomud9G6Ir5iVSH+0BGgJZK3LZX0JscoQAAAFJBmkNJ4Q8mUwI3//vD7l3aeLA78Ouck9tQPwA8lswErUECEWb2BfP9eCGIuUSrHm2tubC/IOXP44SUuC0wRBPcOUZ6UwEfr45QL2uWpkMDOFPAAAAAbEGaZknhDyZTAjf/+8PuSNeECUbzFqDoHVeVghCqcIHbUioPt8t0kMbXawr+fm53In7aCpqLDhgohy3fxI7sPPRi+v4k4BfCt+ArbIn/7hr/fNgQ82kiJyjRCR88G6lkZxT3zjOqC80cK2gQDwAAADlBnoRFETzfhHJ+XgcrH6mSiINjyp//kgke8KUWO3XO8bzOwBXQN1PZPssKbb2Qy1CNfZzQq7z2F0EAAAAsAZ6lakR/h60AXDOWhD5PePsAKefz0nC6cwFfCc9G9GK1nNdoyU44ymIpUb8AAABDQZqnSahBaJlMCN/7w6grIDgahJq1KP+AGMxDG43cWoTqfOrO+8pS10UHq7xCbX4a4pJt454b+M0H1Eo62OjD2EdzxwAAAFVBmslJ4QpSZTBREsb/ygqpP0TAwFGuhMe0H+LKq01NhOFINpQwf2KjhEvSCqmG9OMZE2/Qqy0272Gt0nHLi11coL4E4EoCisDk2gIBnTbHmELqhvwwAAAAGAGe6GpEf4k2H5XFXBQn5jeU7JdWsFUrSQAAAGBBmu1J4Q6JlMCN//vBm7tPgAJMRPqV++IZbHFD8LyLSaFG1Fj3xnwOAwhoQ0X3rcSTj34Hin1WxPUwQOdGW/+/l7yZtILXAd5Z9Wgqf/aUbVivhneizf3wbIgcVChpYYEAAABCQZ8LRRU8n31uIe20/kHBd7snCnkEo5gkxddxXoKnJ5QnndNiDb5+E/ApR6QfVh6OvB748Hzqtcox/cwTiLlMZwRAAAAAKwGfKnREf4ew+Vu++ZOvbsHGmKLhf3qdqHvTmJsewon5VlSMTQ0CLsoGQswAAAAgAZ8sakR/hNQrQyxiUcxzYxkWBaHD9jMxUSoH24CdrYEAAABJQZsvSahBaJlMFPG//AgTxcoiqckUP9wRELdF44LQ+J6zoeYQmvRkGAschCcjUxhb6c6E9RwD3q5AdXgn/otCYxPZzTf5u3+JwQAAADEBn05qRH+BB40VhOFmnLMmN/Rt6KeYRcltPwhBw6Yj8gKW5Be+d4emYrBg0QKVVh+BAAAAOUGbUUnhClJlMFLG//vgOOBCV69hkAISts4hjQjQR4n/taVosJ353hMn/bQatLAc7vxjVVNjpBAL4AAAAC8Bn3BqRH+HoDXutdsAXaf9vPXmz+vq9nHMZG8uHAniuS2VQVH0wdR1pyOJp2ZkQAAAAOVBm3NJ4Q6JlMFExv/BURmgKmQ2sem8M/cS5YiFuDEKj58ESA2Ioo541n9schk3ne2Kz27RHR2Mtl7sfZ6Tk8CR5fodz+unZYWDbUKNmf3Gr1+o9Rxy9lYH9eJLD5VCN+c3svum1sTzNtsJGBK1cHItHWUqHplu24Ckc7vqxHOe0ZtkSJdUThqkuknVrcKcY1ewS9SB7XN0a/r2y0t00YwIYLkLxxLwPmaULQNo/ape3s5adm82wYXnY+bUcZi6E22itPWO92rOH7IP21AdVJhrsPeiqRAczmuYnL6vkPzP4adcC7BhAAAAWAGfkmpEf84wSggSH8mtdkNj5CRqWoGie3lqD7CWgy3dy9dUCO//W9YZNfE6Xiu5GIKYFfDcUA0T2ZBvdwLfngiXt+EDkfDWv0K7BOXReZCjX9nSExoLgSoAAABZQZuXSeEPJlMCN//74DjgOprY+5Fmn26ag8A76dO6JZsciGMnQ2pYcNEnea2hGeomPmCENu5nJeGhNjlWP/GqpX7N4n3KhFSvZLLumw3qoPc7g3O8DwTNGcAAAAA4QZ+1RRE8n309hYNYZQ7+HfhECBi/K0X4frDrq0X/KbwmZ8BThyQ2+fhu0YmOJCJCHAKCHj32xcEAAAAlAZ/UdER/h7D5Wmcz5GDrLvRDAsCvzGdqHHvdNLLc9bD46lxTMAAAAB0Bn9ZqRH+E1CtQOK5zZksA34p0LQK+zTMJUM7eYQAAADxBm9lJqEFomUwU8b/KCEPdL34SY/510oR6jdtPc2spsXUY6gpzzsPQOWCHtIKBf9RsoPGI0GagH7nZSdUAAAAlAZ/4akR/h7DkMUOCGBJh+lP5pSg5ck9EFnrQLuSi5xBprxkDQAAABQJliIIACf++DLeYgN9XYj5zg5B5QJuOQipq5DurWEjnW3UorRPL+Gmp7W/eIApq4hA6GO7MXHIyDHsVucz/aMFYtKTDM8z25ZLQfiNZJqSiDN9723ESzhz79Cwt5M+8vOWuuGhG+Tvq9foXjrYoVXiVDwEsVfUmlk5TXbEqyYI+MZpo4mDJdCI6FifaFbOBypE2MZf5ZLIBH7g8U1rXJzlzobXWQwBVEZ3H50m7+bGCm/s/TO02Npz6eSKcRnoIL6uejcV8scEPRiZKvQslaZioRDQNI3E8UN6+/c9cAekW3/GkyyVmNiyROqcKtiIhBHi7pAk2+Q/TQP/YwlHw0o4KIS19ytAyX2tYiRc1FzBcAlpBmEdwA9/TY2cHx8Zwfg5om1WDRf7F7wfrHwkDprLxXYIXoXvv2/us1OW+b+Xd99mIui98EnMS3BL3wkn/NUHomZCIfEFUCZVjLemRQTcW2M7/WlzDgz8mCOUhc7zw+tq0JUXSMkS1L/cOl+y1DmBMgMzcbhhNbzDLeKWkrUH5XaHyryAkCyr9U1Xo/mAgTLsnExuVa/P4R7zZdXvaMIk7ek35j2b1yb5zOOXYieE3ARP5O5ZjAZBq4ndpSi6jkftKCBVFGUAJHA0yJ3SOwTwLlD33y2ioFahycU0S2U1an8v7S5ljSJFgUt4xqxRzXPLMm5Q/lf3uq1K+w5f2w1oQlGAPQWlYXk6xSti7jEz02BQEwIn/1VSW1lDXiIsLQ9/68Lmyq7+KqScF8M2ttOd6ikUYm63rAEQ/1uwF2tw+vDoK58ZLDjGGY+DsvOCVmKReWQXo2mJU3Z/+dukASHHoLCtd6cudmKCocwPIOFIrZIpuWFPYP7N5w3TQNI2J7uOoaE0UlH0FyMvlyJKCSxK2KXdZbEjdWLzEYIhD/rSHAXZ83AHb4RVPmbBpaD4FR9EhjzYQ5kM67XQ4jgiM+83E3jJ1tuOTwVjEUXOTiZCx8wy+C0dX3lqduBIJOztDCxzLLUgy1pDJ11agmN4hiYEyREcCB9PLz7Gu5NIl+dxSpct15aKEYrYOog2sW/3UjxYhz6/+6vpRxNYFBe2YVFfEse/JAPrCdAPUQXwCIpcGvCAqIe9H6yayOc2/eUy00kYB9Hqrr4iHiednpjGYIa/wBQn2kawkd5eYwvCd09N//TBsRiPySS7TyptSNcNkl3kSijqYAzDU0cTjAlo7gSKsBb2500bskV+Fa6qC7WYdX+uHeAkkUfK7djG3sDH96VRFsoJGnc9eiCjPYyGt4n1nklz/pm47DKZfzdVcEQVVqXfQk182IfGwqQ6YPNZNc85BxPoa5UG0Nc7aBMBm/EEqGdEA+1bpcQALD7p+Xw3OdXvS2i01L2TZwR76H4FwPsv/kg5K4Q+gtukTkAv8aXBY908KoaNfjcQLsJI9nUsxo8CIfcaHaV3fe+aAg0mPXJT7WoUH3EqFX9KWMCffuLgk141zvZdeWb9G2OAayn61mciuEFWUB3unf//3UT9ZWH3tY8Odtz5knWw9WcA/7IUpQHs7AGXSZHFSuMbAJ7IeW5YgcJ4/2xVyVfqWwGS8VXHS84Y1Oc2iGWIYRH4DJUk2xI3a7Xnxr04yGYzMm4F14EhpEJRU19w2fQAabSuyr93+75ace0n91G5DJsmYh8MQGBLpKZioZADEdyOWOL/b3dH1hKiEdZsQsnfs/yhdAAqZAAAAPEGaIWxG//vtcWB39W5QeWKwGeLr7PMDikgXstlWyUqYOYjnV7wDPmyxfP1O0kL6PoxR7zo1Gkb+C4TG+AAAALJBmkM8IZMphG/Hqe1RFY5rt5egm2nn3JQcBwuaJRr84/ddLi9/76agb7ejonTAqd67XVYfnqN4BuUaC+fAk+1lcholPMzUb84TknMEw30GkmM7GgprsKqCgwcaNTAOWdS32AuH1x/E2QaetB+/mEk87Hq8PlJqqlYpWPshK+oBJK6KuTSGIwg/9Bv+LVgT1pALr72MIHMQXmCPnTdlDJUk4bEuuLIqk2to3LjxIlmPkYtBAAAAGAGeYmpEf830n5JzFkX3j1Qd8Q+AUGDVgQAAAEhBmmdJ4Q8mUwI3//vBm7tPFgd78v7l86Um29hR9xxADhgMD9LzlgdAOtozhLW7UEkoFYf5otHinj6odtZRtbm/yoQHLKnpnZ8AAAAzQZ6FRRE8n31uIfah+VmyyuHZOFPIEoZWWjwctTHilRP/EBaLGght8/KC14RVIcYtmzmgAAAAJwGepHREf4i5UoIN1KmuYE1Jdm1Hq0jw330ncp4Hx2M/vKMq5Z4tgQAAAB0BnqZqRH+E1CtVyjEqWYFcr+nHuyCeliO2EXDZLgAAAE9BmqhJqEFomUwI3/wDoisgI4RECc7scuWX932y6q/Gcps10izFusxOOvMniGaXl24boq3k+UZg8rPMUT3iQxdNcNfEd4A2W68ovwz4f+7zAAAAVkGayknhClJlMFESxv/7w+5IrxwHtyq44n5qzWH5tokp7F/MIurN4zLf3qsBy/LP1c33i4OFn3CMZ1TBLIwXvbEHj14Vcb6Wt30UGVr0GhXroArD8I7AAAAAJAGe6WpEf4etAFVTbNo8fLWftyKUew1Jo8zvS8XWIpeI7UlrUQAAAElBmutJ4Q6JlMCN//vE9pkb6iHbDeYYscsOWTLED8gfFsN2l+iXqUwnGAUp7SGWRWf+Tj1MgMpeHoo49mO3nGzGBc/JZpGbbwjNAAAAM0GbDEnhDyZTAjf/++A44EO2R9S7lNO7TCCSBTAoqxeSZe0BGgDXxryEEgFwFd7IWXy8YQAAAFVBmy1J4Q8mUwI3//vhUdg9fr4QWmZ38DJvZypVj9XjHyiWE0qfucSUpcKpbBOOJb2lRZHjFXqbLB/AAyfZJCzDiBXSwSe2d3nKLM/QAOEJb8mw3ifBAAAAT0GbT0nhDyZTBRE8b8oIQ86YuCDOQEpyy9UWaSGWOxDO7wiSCJDYIhWcih6YPXwwChxsQuyPkj7FvLL/pgLFAoSeOsn9+9F+L3cY6vDMfiAAAAAqAZ9uakR/gQ9o5WBUzd28w1hPr/+56V66bgLLEK74aalrWqpQrcrjTCoOAAAAM0GbcEnhDyZTAjf/++A44Ilb6pCfLUOd7CmgqVP5gaCrD6c4HElsj+qTov37vq7C0KSjKQAAAEZBm5JJ4Q8mUwURPG/8A6IrIENjLKcjAdSvTZ0AUa2HabW/T6JywFoZYZ38XBzMVTWT9Boz5CqJWPYr9zgIZJzhObYT/U75AAAALwGfsWpEf4pxDGIf1f+tYtz0i7E23tkVLGF0sTK0M/3IyamfbQgfSCSb1BPWNUB8AAAAPkGbtUnhDyZTAjf//CJ/zIEGcd9wsgkXNNOYwvx4RPo8KfLe4zM2exTFz5+cPtcVtktZUvjRO+buNFeyS0d/AAAAHUGf00URPN99TaObeEbxmT4L02e3ocfhCFkUcu0BAAAAMAGf9GpEf4exElW+Y0rm0PJ/wJM/5jtBqG7M53mIMo8kEOBG1R7FTsO3uy/8wE9rgQAAAIJBm/dJqEFomUwU8b/74VHYNYZ3gTviajBNW4911XFIgnHuPPvf6PuINmTDgKw6tl4y1gHxr/o1Eoc443t4pg43rfeXg457+frQjkKJwwuU0Dg28hgOuxXzoy3/7GNwWDiho5fxpVul4JYlW/s2wZ6+FIFz9RB3Pxla/HA9ajISXcDsAAAAGQGeFmpEf4eNe2TaTV+is9Z+r3dj/4M5IrQAAAEtQZobSeEKUmUwI3/OMDPDLD7+9LV1rvQO8dIyBgu6gBZzKIHOqwdcYW0mj4/p+KVMgoNINR6fpQkbB+e7J43ZN6Thfy8C35cIcvE6ulfkdYr0ZP5p+0EdUB5C7vhx0r6X9DkZA0F/t8XcJ1Ay1jkXzIafX5I9JMZyOPPT23eliytMDQN3Nmm15DczfsII5lelhwcySIf4ku6wzRULjnVZWRMxcaOWnCVmvPK0WKnXS6hAGrBovrD5c40CFPHBrLy9JtUVvwuELo3rnQUyLyChZip9Wl/I5F161BBhavdM3/GZPS1ybPzPYbRmbIdt1SZq/ybmP8py+DhwJs6/zeOqKGDQvWpP3NX7JQuvuIBzABOHx7UbDGk+BICall/zPVb7YjEOQFeJndkK8GmaDQAAAEdBnjlFNEyf079+wBmd8pEhgLG8Xc7r/4rqUnzQQEv0ohGUAL7kLhkvedlsNVqrHka/nxHyyL78FF+w3UmoM16jfBFXZPl4yAAAACoBnlh0RH+AbCj0Y5cPjkp7DicIuUO2PQpkoPz3B9y9RBSHaE7ieceKLYAAAAAtAZ5aakR/3CsoGURp1Gbf6xetbBvbdg7XQAomNjTNFLLYoSnN1T9WTArX1a5xAAAAf0GaXUmoQWiZTBTxv8gax37l9sNbrPn6iWwUf8RxgPTzKiAAtM7FGqChNHUkeqgcz5deTS6OtdOTIhcU10sqfJ/G7PWcIpym3Fr6VN/qN2kQwEQKlPOHuE2NC5vUo2j48TIVNP/sgwP3kyf759nBR/KFqquMTQDl4oofPV/0HkAAAAAtAZ58akR/qMvbvw4JUlkDvjD/f4ktFYfXNpM4bGfm2nM37h+CiGSGIelwUC+BAAAAQ0Gaf0nhClJlMFLG//kqHGZpISriFH/O7LMfxPF3levzQUoyR8IdbCT4Om5XQ61r7CJghz72IBV8PYi+fqlXXR+phCsAAAAuAZ6eakR/iXvWj3al4qv+G0JMwV4sy5KJ/S3ER7xINYVeI2qPYr2GLSVtouwGmAAAAGBBmoFJ4Q6JlMFExv/9fMnXUwOYhp70FTBXdRxtKOUP+0QmwlUd3WHIaIY4UFj8Yp5d9xxsyZ9jD2OI0QWf2S1Zcu0qWkg1dQDmX7aUzQyVEEdscleAHvns65gIbvQxIesAAAAWAZ6gakR/h417ZNpNX6Kz2wcSbK+6JgAAADRBmqRJ4Q8mUwI3/8oIQ86YuCDOQ6JMWiYf/iRlpY+LK2B96JBM88YxxpmUYIE8NLeGUY5dAAAAK0GewkURPN+D+9PfvvPTrHvZhI32OEWXdhkGCruwTO5SkxQve5GtKHjiB1sAAAAdAZ7jakR/gQeVxSFc8gptpCQL0Wxcfx1JTwVzldcAAACEQZroSahBaJlMCN/HrvlPtuKGl8J03skzn9tX3N+UGKgPtt5f1yW1mhtqYtpj2nU5x+dydduAMrmraSE5zwuU/zBlhXh5g8x8c38YYs2MriPKi55w4mf0Kwy5yX9U4VJwu6z5gBqEnviXVdXlHGiM76Z3+hcdzSXdSABOlyYkF+z/YctAAAAAM0GfBkURLJ/NXH5jrirbiX21Z3CtP5YcQMlKXKnVSfznQBXI1IR3YSzO0uaD7W46BIwddQAAACsBnyV0RH+vXBhX1AzTfh+ibFtqz6eGMQ/DTS0BafSOlfYuavjTxTNGlR1QAAAAIgGfJ2pEf4Z1ktqM7YEw3wBj96oC8se5g/yL18ispByD6DEAAAA3QZsqSahBbJlMFExv87zR2otPwnC4OQ4HEYtHszJf56ewlbVNYDdRLi7DaNM07knXvZOBJlAsDQAAABoBn0lqRH+BB5WDzvLCGv4OSr3ZyhW71ibSYQAAAEpBm0tJ4QpSZTAjf/vhUdhErrGwgYAPXwOxHLcvgTv8A/Zx9DjubfQaVDTVMebAKQ+RkqueYnX184COJ7NhYucgTvasYPFK4F93kQAAAFRBm21J4Q6JlMFNExv/++A44DvMfnbRI+5Ux9vnvb9fZlh0+yV0sGFmgEwKg3swNSCp7N5faAgb86AHOGr8Z8DiL5UcrnP4kp/lO962eXNLUHqjhHcAAAAmAZ+MakR/h7ETb+sxp7t4uJPRUf/kip4odeYTuzhEga5MwAislMEAAAAyQZuOSeEPJlMCN//74DjgQ+tokMdI2I2w4XB31uWN9xpSf/b7WfaAjQAslbjZLMjZGYAAAABOQZuvSeEPJlMCN//7w+5d2niwexjKrUXs6RzXh1b13XA1usrXuo4Q+o9ouvQ1D2Q2Wm0PpqZJ/HVh07u1x+848SMV9TQUPFzbtS13xBB6AAAAW0Gb0knhDyZTAjf/+8ParGvCBK2RNZj2l1x+jAR0apeuLZtFzhogEMtR9DC/p+fqK8dILLP+IbGFXeex19kT5IY4p8S3kRMYF4bXDgihVw55Ypd4XX0hDtrHn2EAAAA5QZ/wRRE834ISoy161qatNvHcx2Sov/bKprDXTckBV7PTdQhKXDndGlf3VqRlNangjqwTorzS3wjwAAAAJwGeEWpEf4etAFVTZCVL4SV8n9Kg6fE2aMda+98WXMvpKkUyXESP/AAAAElBmhNJqEFomUwI3/vDp1Hm8IcqtaXumQqIJt/IdZM/LaCqQDp5QDE0HJXXxYKzqroGplUEOXiTI56df1TOqn/+xhhrsArE1GyBAAAAX0GaNUnhClJlMFESxv/9Y3sbc/g1iqUD5e5Uke9bb1C6YGB9z3awZ4zQI0M0unHOdJnHA1rrnsh4SimxV5KpEdhAQ9gFOWZITFvHKgu1okFYOADb9R2Rozp0tWuEyuxNAAAAGwGeVGpEf5BbOR6mhp16rXzs2+e4d71TyW9KQQAAAGZBmlhJ4Q6JlMCN//vgOOB3OR7sTCb+I0045cZJEQ4U+L3XCpcnHOJwFb1e+Q0W3Q7uM0aQ6yv11+YTtxoOdsy7RQ+LYYJDbR6Rd7L5XFyxpKDB+0OWkFw+YLFioVUh95SYqxqecX8AAAA8QZ52RRU834P7bhmByo0v4xgbLCddxYxYbspCwfmTeE4h4m6gzc5yXll0xuqHmzfhwYk8ZpP+UzAb50WNAAAAIQGel2pEf4EHlcUhXPIKbaQj7c/ATAoSkel2SJNrDV4kgQAAAGxBmpxJqEFomUwI3/125QhMIg/aNp4kCmUKfLpf/QrlLIP5B/oHjyqjvBOioUOoMoL20buyufrF/UIt3sv6YVyWShlKuej/5mcBImU0X9mB2jkyIoDrWUZBjFA6haFPIB5+QhgWaDpI61usiRgAAABMQZ66RREsn32aKglvf2C9+Bdon/+X9WLOPlx7m1yaoyHmMFFCEI61yhs2ICjfrbx1xwg+UNHTSx3O+VuFL6oUO/x14wF/ouCv7NddvAAAADEBntl0RH+KcQxiH9X/rcK9KC9kRdt+HJ+G/iNZPTGZ+diB+2gxK2Dg5bXzZO90e7yZAAAAJwGe22pEf4Z1uDH/deStpiTEc9tyHnyaVWKk+d3zLNz5R2rfH3mWEAAAADpBmt1JqEFsmUwI3/wDwzHAd5jvuFkDKQEVOoZiZuzsM3wXcA8F3ucQwarCn13TkhsIYJMiT3kh+1+BAAABTEGa/0nhClJlMFFSxv/BJAf7z+EOU2CWqKManlZ0LpOsfJcBa0PMmIuvVH9dbvAcE5UzEdHVQbTrz1N1jb9WiMOxu5R2cy5A81OxB+5ZBnN5eAsC32urtdTTKK7QubeiVOo7VylTo1bcPy3DSeA7L7gshvKbqb6yZRCFyi4whKSJ7RPi0+A7uR9tDdfvNWqXsP4NZEdkTw7nJJscByribz64QmP/ICkD3ZxZqkoqOTVn6dlY+P0SHMy6N2fn26w0kDbCjyPfS1dZblvZxWjU616LYFDlf6NoIU40eHrlDbvSs4WEm5sqj6lAdmcaMt1R2DDkFv3lJfP04qVFUwKkb4diB0AgVCwqoCaqflufpSnzlUP5zpGT6XXwKzdvxhIiAGd/HKUkMg36acC0SDhcMpk2kz57Q3kYvM63o6ggTKO91qvRby/6+hO4M9+BAAAAWAGfHmpEf84wSggd11AiFSufQxM9Luw26hRYDDKsE9vSwAn7Z3kp0GuF0Eqpea5n5unP9KN0c3L3GkYNAsaXrACiTeM7/yz8z1LY9wjvbK9PNQeYtR6wMIAAAABfQZsDSeEOiZTAjf/zunZ7sMLNRD1X4hTb1ZFO384nK86xDwxI5Vpgs+jNIRsolIEfvY6lsaVpHcZvcrFP1//9wB89zGosviWcqEVLIZ7/+WmTYZWLoPL9M4Sq38lHkYEAAAAyQZ8hRRU8n3sgpbafys2Rudk4U/87NYQWWpO4rxUES/JyvoT/lfbQ9OleOWWIZeQm2IAAAAAoAZ9AdER/gHHeAx5KQ6o8Bkl+v/9P4zpRB6XdjaQVIDAbYhBhxO6UwQAAAB0Bn0JqRH+E1Cs20Vz506WBTrufbbvFjJQQkJ28wQAAAD5Bm0VJqEFomUwU8b/KCEPdL34SY/52bZ/rGfsXPf5yCTDido8G77Go9+K1BJRLL67CRLKDtmUDogd5/yp2dQAAACcBn2RqRH+BD26tG3o2sOTqJd//LOf/0exfYfKLxm7J3qgsoJYfg5gAAAAzQZtnSeEKUmUwUsb/++A44IUf87wFjEcFriJhVhBdOObtZMRA7H8L7O85Qi1Ir/23xQWBAAAALAGfhmpEf4exElWANW8vMvakG0lwp25Oh8MHyFQGujxVcqIRtUex9+Z6rsJYAAAAUUGbiUnhDomUwUTG//vhUdhCkgtD/OKNA6dG1osSc7TN7EVmJ8mHYKHJo0kuzmVJ2OcDMUMCsDRqChpG0X62F+dsmKA1zveQZSPh2EvbMNzsMQAAABUBn6hqRH+I4DsKsOugF12nJUDYjiAAAABoQZusSeEPJlMCN//07ziWU3i/kpvzHFCOn2CcItTi1yFM4xbCfhV3Hl7kKDI/9hW4d+LvO5FlZCITgCrsyrc2JS6CCKYTrtDuaRaAvGVwNn+giCSu/6I71EcdF6AAzs8YoRZvHfx+r8EAAAAuQZ/KRRE836foBN+l6dUqFIqLR7fdXGqeg/3ZJ/83TbhEn5pkA8K9kzYNkm5sAwAAAB0Bn+tqRH+BB5XFIVzyCm84fzzyfGIzavoEg2o1vQAAAFdBm+5JqEFomUwU8b/5u5inuMF7Xq8EOtQO4YzuHe2WcPkEeIzleyoHzUSRtXdZDfWefPY5FxPuebWlVj4LVrBtGv0RP79YCO7uGRp8no60SRyqygnK4c0AAAAxAZ4NakR/mJ3s34Q5H/1uFexTk3VIjFEo/fwfhqCP78QfAaOXWnodXpVWh/EbwDilHAAAAEtBmhBJ4QpSZTBSxv/8AoC0RB/Gq3N+gqfQl07//0asYn4GjHzgmabjVHeCeHimF9bt89uzqXPNoLLVfk5lK+SEWZBMH+jdIMD7NJYAAAApAZ4vakR/h+Qj8Y/7ryZBazBlCQSff9dCgZHsO8153DG/1Iu47qZ9y1EAAAA4QZoxSeEOiZTAjf/8A8O+HA7nHfcLIGMIzn/UI6nxZeLzvaZpNm4NR++zRTzqakzFTJjhdFc2sWYAAAAvQZpSSeEPJlMCN//74ByBNNlS+GOnj/iYE7uR3e+eaDjQFr415CBPX8resUSqg9AAAABEQZpzSeEPJlMCN//74VHYRK5/s8oAjSJlNV182JSgSuYDsNFFvPml8Z2TC06yUjgNGaIT+MGdO85RZn6ABwhLfk2G8T8AAABHQZqVSeEPJlMFETxv++A44Hc4/O2I/eN0x9t9al1x12ezWSFYh9FLhxu8aTceePsZOcOf/noBS072io579hr6GIKYfVDIPsEAAAArAZ60akR/gQ9vOA+UWsy1Mn3l//+f5lQBPWWuO+aJJugzYNwOMU04nd4nwQAAADtBmrZJ4Q8mUwI3//vhUdhCQoq5E2kM7tSSvPdRsTqN2Lqyx1YynzDOfYXeFZW5mbzBzov4IZCErRmMoAAAAElBmthJ4Q8mUwURPG/8A6IrIENlUdHg1vXava9lPZfzExXzXgqZrbgK0CpfLrF03F1J4CM9W6Uc/jlRwpOi+VZSpqhCt8/y775AAAAAKwGe92pEf4egL4A9rxjL4OGkdU1tP9jDyD8YL7cVkl98cgRXNXC0itKOXK8AAABMQZr8SeEPJlMCN//76z44Hc4/EYYtcgPBH9RgO0Iwi8OXMzKNpI9izDnv5+poRBie+02otvR+9Zng2UG2aLxEiEAQa2EyHRWuU4x0ZgAAADBBnxpFETyffSt6tLJT1q5nVhquTgXJSFs0fGTiWZMvJijPSmPEtnIL2xGLEmLILUAAAAAjAZ85dER/gw5gGJbwM9dZrSXxrWmk7SS7YSzCpHXtnb7SsiUAAAAXAZ87akR/gQeVg87QMSja2dC/csmeH6AAAABxQZsgSahBaJlMCN/74VHYPZq6tmtxLh1sxyPaSRfT8XGQqvC/Q4343DJxZXe1Sjk7F3JK8Ga+iJBjybN1BsuRRnEENt4FtlTjDMiyslD/QE2S+ie1BjE0ye1VekieYTcbeA5uWUCc9oAS0oHJc7DDprEAAAA9QZ9eRREsn3/Spxn+3+PADN8tr/k8dm4t2BxiP+1thpygNS5/1ulcsz7YQxXLmulOZ9bSD8leHIfywjbusQAAACQBn310RH+HvdOD1Xbhj4zYGbutZl5v6HqyOaqspaMKaK5qWeQAAAAwAZ9/akR/gQeVhtGJPhygpFSbn4CX0RYdxtJU+44KMHn0Oo8JLVH4iFlYBJyRuWktAAAAhUGbZEmoQWyZTAjfx83w+eo3udPz/wWGJzqAxttakOUQ5KIzV5rlta9UDHrdCKjDiNALKgEJA6puZDaqChxqlBxSo1lMmdEYH9x6+Iw/G7cBXYYH/vbSuijBJl7Iv1SdG2SxEz7WKOVaXQ7edDngcN7BTtsHxNL9kKKH2vdnV7h3Up4bBiUAAAA2QZ+CRRUsn81bxhrs9aLlJLIhMRyrlgrjTzbZxY1kexOB+js4nC73oqm2oS8KXxFXxDDK0y6hAAAAdQGfoXREf9PejGW7yeA3K+Ashsi9mc5rzdirkZ1i8F7plCYqBBB3h2i/Y5FflOyMyOlgUqYL/G9xIzGl8z+0IWKUH+x+Fv9B7/aREDydtwyQInEGxGayCrl3zIGiZ6TFkkOalquRg4SIEHzOpZAOwkut3sA6cQAAACYBn6NqRH+Gdbgx/3XkraYElmuoUqexzkujzvNedwvPC5NXjL0QwAAAADlBm6ZJqEFsmUwUTG/8A8MxwHeY77hY7su4Js5ePegnvwLeGgOmnRghK+XuEJ9p1ZqTpeEmvxKPv2gAAAAWAZ/FakR/gQeVg87QMT0JOabD8/BfTQAAAFtBm8dJ4QpSZTAjf/vhUdg9mx2D7TmdFz8qn2yy4xD/XK0YWq6bXtpAIDSMH89fnlAnLzoUK41nLcHaRGQ9VYKHwPVgSIejbid4zu0EWxJ7r5+L5ivAS0Wc7KiKAAAAU0Gb60nhDomUwI3/+8Gbu08WB3vy/AdqEakb767uJfNyZA4+pjSLzSTlkccggQWj3N8llt825rufA9wokyOuXzQUeglOEArfFlY1GPqmrtZTf6vhAAAAMUGeCUURPJ97IKW2n8rNmOXQIGJer6mMWBu7bBqazEhQPzV3PBt8Hffu6rEFJycnZ/AAAAAiAZ4odER/h7D5WmcyS/Um4Sed4v71PLKjq2xzxVZO257qOQAAABwBnipqRH+E1Cs20VzmQO2GD7g9FZUkcZEXqdpxAAAAR0GaLUmoQWiZTBTxv/wDwzHBCj/nXShHqN21NSADJU/Jvj4ZToB/1ZTomGKrdxyAf/iOIoo9DdUgNWCwuF0p6LBGjn9Ou8lBAAAAJgGeTGpEf4EHjJ+LiOn4U4Y39SmH/rJxyXApfMRHQp7a1U3JiFaDAAAAMkGaT0nhClJlMFLG//vgOOBCV69dVkGYj4aWcO7kQz+XVCeG9zDp5MQEvFDCbSD/y8jgAAAALQGebmpEf4n37kAo6ubDhZOY+8dpADGCpKbr5jcNPPpTe79afnKaDNgbJdD1QAAAAFpBmnFJ4Q6JlMFExv/7wetWHYG/jzvTdZyGtgH/DXdI0sh+SMrklDv0lY1O4TvP0MR8DIsygALU9Fs2WbnRqnA2PULv/h+6+vnPyCATdnYbIXw/4lIPbTxVj/MAAAAYAZ6QakR/iOA7CrB7i4FdAnsFNG/HYixgAAAAVkGak0nhDyZTBTxvyghDzpi4EXYDnbEfvG6Y+2zMQ/Hb9ovvg/0chC5MpSxOmJq9zFtkv9T+YO8D9ACgwbrIT3dU6sd/k0sYujXXO5r3U7FCF0bBVqsgAAAAJgGesmpEf4exE2/oQsw3fbuNFR/+SKnih1Q5+0sivYgWYAQIdsxBAAAAMkGatEnhDyZTAjf/++AcgTWtokgude4vE93Jud9I9whdUFDZJTfR+J6YiF8XdIzfNb/FAAAATEGa1UnhDyZTAjf/+8PuXdp4sDvx77lN0KgU/fpK8+IINR1opM0Pu4CozS8l0Ud/us9fntwO4IG0oZYvNBRI2N8a8g/0sIn+6f7QQekAAABZQZr4SeEPJlMCN//KCEPQpIQEI+hBvIFMoU+XSr3pPEb6tM+0KTmzi2hUR7j3nqTXSRaGGG+1IgCzYffvSdjj9k3eSImzEFiL89xBSAuBHuAan1qCm0SbL2AAAAA7QZ8WRRE834Ryfl4EPlupkoiD3UHAn/tyNF9cYKC+bqyZpWmU1htKRBn8v1v9le4Ox3FhOYCg5NB+ZdkAAAAsAZ83akR/hNda3bhOJUZIUUY2jVhVdrmxzmILX98ciz/mLhHdJSmfczV6U8EAAABFQZs5SahBaJlMCN/7w6grIDges6Q6N8LLKBBR/r+r01RScaGlerMw4mNRGZeM8kNv5hlT2J2qEHTq9XsVBx6y0AugLYMoAAAAZUGbW0nhClJlMFESxv/74VHYNZiODju5GRyMj9wlHqF1Lr7pV20VHiAMrHfGA6p0v/Zr0/dcP6yMHwgdbTTN5Pj6VJC/71UkDQHF2SgOkPs+SUevM06PMV4Dc/8UKpkKq5MzAfyAAAAAGgGfempEf4k2KJe5dxcZpOfMZDO0cRrAfBG/AAAAcUGbf0nhDomUwIv/+6w9u1GKETDvOv56IQEhONuu4MWW8U5koOsb6z0DMrcU8frYGb496r+0cR0j8/9YT3bwj6KnqNIzzKxPxz20387LsrGHGueynDGovxAkh6APLCz5JN0en8GfJGJgIpTpHDn4oF+AAAAAP0GfnUUVPJ99biH2ofqCeWdcOycKd3M3sDTTOXcd9oRK3Tz1jn02+G3H+x161bG+i7O4manzYReZ5zgb4/OHSQAAACoBn7x0RH+AbCaaY6JFZrmbJvY5TSCOfppisvf7jhPHXneiHxtzcLARd1EAAAAfAZ++akR/hNQrQyxiVLMCu/aBfsIZgv1fU2C7T3Ir9AAAAEZBm6FJqEFomUwU8b/8A8MxwQo/52bZ9Gsz+sY4O2RS3UuD8X0CNvCqI3uAP4dESewlMKkkQ0sYA5go9F7wlfSnLQu+tTWBAAAALQGfwGpEf4ew5DFDghAjpjv5rOLhqU4poz0tQCo2nvYk1MNk0SbxLDlMLbaRYgAAADhBm8NJ4QpSZTBSxf/730txLjWi6sL/8WkZ0SxdX+xgko4N3jmSClvhOh2sSC6GI34PiVLH5H6DQQAAAC4Bn+JqRH+HsRJVgDVvKzmQczvZIZd7ldG4fIdAbM1M+rxG1R7C8fxWQE6gkEOdAAABb0Gb5UnhDomUwUTG/2eshp/CHKaTLMVIv8yljz56U5Qg3GnKqoZs4y9mnt42jrRlYQiQvaz8jEhuvC1SoxRX3jdtVE7CW2y7mknEG5FfETEOLay9NAEvDLcXz3Kvp5o6r1Ecn5DKuRfL94ZwgC1vfOYsAR70mWBfVzPYmO5bWIXHrsivWz4TmSfJX8BTws++We/rWq+eEIZTOpQXO18Z4TBZJoJNh+NQKQ888eOePfvsvOd7ExGYc72AMj8IB5nA3qIIpZ26bWxdEigqp5WqA3CPmQ1ArglG8O23AW9a0mGskFVKIURyO7SDrMfvuSpM2Ufg2ko07dDeVY280dgvmukE3rlhobBVZ9eHnJpB/rDydTtHtcgUHxLRRN2WPaWUwfFExO2NlPxDNJNoifurDRR1TOE2xSpff831Zq9dXYJqU1zKACLUfsS7gTYxNNWcQ7fT37PoiXtJMLOR39uru4L0h58owlQKSE0dbvzeGUAAAABVAZ4EakV/6f4hhh5j4pezXGmc10booa3ca783iOnUPxnKk5Y3HfubpAwe+uDDln7x2eceSurPShjo5iSao0hBxT60iJArFLPcDKmG1k+1SrnGiBiThwAAAKNBmghJ4Q8mUwI3/2esg9XV1nn6OqvkurteiZEF/cX1g6vwMjx6DR+rv8oeaoyhuEGc4M80j7FCjt+C+gCKRU1EONaFpssJsG4r4wx6++Vl4OFVBDT83lw8NrDATf3vTsciLFGM8eGrhKs8pvOB4zbD2z/wbkYod2zVrCHakD3w/D1EFZvcK9tKs1Nr9e2ax3We6/gOaYcpR2bYNG1etHQuzTshAAAAMUGeJkURPN/smxOo9ZgcqNMHlkhDvC9wqzrKarKnUnQpTSCQZIsbYKJtUJkWtnazMvAAAAAlAZ5HakV/hwRZbYwm85/cCb2d75a9rssTV2GyXh2ERuMyRAZOQQAAAGpBmkxJqEFomUwI3/vEgWa+og3x1dnJ4PoDdBe3/3YoMTUoB3TyU4Ej/zJxdOfoW4aHRQ/I+nHKZhqGdhrBqoC886/vF7jEJplVhBx/pyxBeWOV/VwqfkCrurEj7aiFkmd0eJF8WUso2Z6AAAAALkGeakURLN+CHmGW7/MKreRIEd/8QfLllKWohrf2XGUEjDty3V/zICAlat/0cmEAAAAmAZ6JdEV/jY0N0RG/BawsJgKYUujP+rhnqWcqJ30YzX0c3+Yd3EUAAAAkAZ6LakV/j/S/wWIll3CGdQpmmnlHJJJMDXC3Dufc23w6T4+RAAAAUkGajUmoQWyZTAjf++5XDgdzj8QavWtNM5/v7SWzbYKyz+gr1lvvDdabJmin54HFA6FS5USZmz7Qe5i8GrS/OeBhEFUi7sU8j3do+iQgzZ6NN4EAAABqQZqvSeEKUmUwUVLG//vhUdhErrGwMwT2ApJmJaluur8WZXKPf3cxs7c2sP/Xbjfjlt+FJSs7mUuT13MbzM+TrSACCZSMBoDthPFifoGGtkxEnj+e8h73JprN7phM0yaCdgAO7cfoh5J18AAAABgBns5qRX+OkjwIvXb1vb5y58tHlKJ2YgMAAABUQZrTSeEOiZTAjf/7wZu7TxYHcd1Yn912Q4e6YwUip/A050OaFjhX6O2qQXp5VjUL059lIsYYgSQ2jExa1CDIwChKJheC4ssji0HjIrmDdIWl9Q3JAAAAOEGe8UUVPN9/4NuyfQi2HxhNnY/bPczyhQbbExzc0zWVB0vRCS6njJQu39UWTWCDaVRAWBLJEnaAAAAANwGfEHRFf452ZgVPqCEm6Wp074F/bJrDI8pmAHDsFuIJwlvPUbZo3jO+9e0mbvECpBl/WwaROjAAAAAkAZ8SakV/isvpAhCb6whUJeGpnwag9sUF2UycKFAAIxuGur9dAAAAVEGbFEmoQWiZTAjf/AOiKyAj+gnj0a3Gva9CNJ4js1U2x0aCm2qxyr5dY19k1pZuImuv/8peW/h4jD+21kSZLMKNT8K6Qp1NQB1Oz3dayRj8a04cwQAAAFJBmzZJ4QpSZTBREsb/+8PuSK8cB7lWURgB2UuJkbbsVfmYf6uGo0ihSHCGRIUeOkGwszI/RafeeJLVFb72xB56AeE7RHwA6KDK2FMZEiHn9cDBAAAAJAGfVWpFf42Vo2qm00Dddqve1Cwh67h8F7hLVO0VUOGCP37hFAAAAEtBm1dJ4Q6JlMCN//vDfCqSMgQ2VRj3pLESWqFJI9xvPW4RklANZ4xZHF71Z6iQq6DzlYFf0+Q/xmZ/fOpxmMDAoxgM2EezYMdxGYAAAAAyQZt4SeEPJlMCN//74DjgQ7ZH1LvBkl9O69N9KO37jJz7QEaANfGvJaL8Blb2Um8UToEAAACFQZuZSeEPJlMCN//9Y1z/dD7lReDd/jih+Ld0YtRndvMbeeMTxfw1CN28oAhAs3BFnwK+WZTwEg2f//cT8FjRJI69DXHiU0hpAJ1L4ZcWUkQNn9/F3WTm780WFIBn83SpljWZ8Vra05F6qUNyjAMjVTOd4yHsZRD4BpeH6ERAEEj9uWg3dAAAAEhBm7tJ4Q8mUwURPG/74Djgdzj87YkIIqfbdVCHiXmoEJTcLeAuy7kDg1f7oEo7JIcrQiC/U+8t/yN57HX2mDX0MQUw+qGQfYAAAAApAZ/aakV/jnhoov57JnRJzUQ582/A+KqjcuQgc99QQQxsklKC4RUIecEAAAAxQZvcSeEPJlMCN//74DjgiVvqHRYbDnew3k/WoU8O62m6cW1f4z4f2KkM/6uwtCkoygAAAGNBm/5J4Q8mUwURPG/9Y3sbYDUySihk0ZwYHPZc98GT3bkkuNK+kxG9hw6QhCPLCTOPu0cTsYadgUqI051fWxMxsY+qGfVFtfK/8hGLOVkVoaB0SjGYoDGFSY+reuTO1+te1tMAAAAvAZ4dakV/j/S/kR07/7I2ONOQxpj1EA6tGW/N30/w4dgSvirINWao4hRz+KM4mYkAAABMQZoCSeEPJlMCN//8AtC7wRdgH2cIbrCoMs/ZXvDYX8DxXmXaY94n5P78/TKQXUbvgpq5JuPge7/U+LSYd+1eMFPL7ErSRU/SQWbYYAAAADVBniBFETzfgcsI33hBxsPekc3ZDBzYdeklj7D+yHKGhCeLw2I2cmn91BAltpq2zXE35Ex0sQAAABwBnl90RX+JF8eO8NhiMY7e3XCBBxtbicm//Ab6AAAAGQGeQWpFf4cEWNCU5dbW4vKx29yJfhj97H0AAABhQZpDSahBaJlMCN/7x1VfiFwgJBqttUqkagOHWBpqzu3I3Z1ug1iDen9RqMLuWwSJns/FSjmbZC32wpaRGK22xjoTIf8DNdWNE7Xda1dD+nvhdspBUZZ+XrBk1Hqe2E0wKwAAARpBmmdJ4QpSZTAjf8FqjzHg2B+Dqry9XK4oFY4kpye5WzL+gZiytMQXb5+G1o3IxzRZhxW9aZpjLbUsacNdqbjHaYZ9CmpTlK7FK44jbJ1larK+ACOQSr42285xajESkZG6raHieCShaT5/xNHmQa+kW25//2rjWsefjIOVuKFkvGla54JhR+X5yYFqm/bC16PSKR/h99ZwXzxEyHSds+3pr2CMYM14vyO0YKx2HjFF2rVw+hoZHoyQ0JcL/ReGkwZ2fNhkLNreyaduEAjHZAAX9t8IFb0Oz30qFk8O7KgS9ZdqStAL+KUDJIhD3HAAp3xmSs+BqTgOrQ5PS+Ace+fz0i3bKIPCU8i8AJwZcWiUwBhyv7knwrO1Y4AAAABMQZ6FRTRM39foEqx/3JF1qqZd8W//5kxNE6DYERz3YbVOw/1jyVC2gk+UmHuYnrw/mUZrr6ik8rqKlUQKXKQUw5GLsaQyVoSX08YJgAAAACoBnqR0RX+Gfrvv7yJuu0tSB+P/+h4yK0TNhid5Xn9+DciKgmECSbWRRssAAAAtAZ6makV/2fIPxtyYTe2F7AUed/9kq/9NVDnFjpKyZwUbJZWypoABH7rTJteAAAAAnkGaqUmoQWiZTBTxv8gMvylDZaeIJIsMYkC0ZGm9WiRczW/gYZr7/QtVjFGHEuTg4I7VYpSYfgz+C64TnuAQXBvRwRlD7R7BrUKTMlAKgwmktl2YT4IHIivf+PBMSAKaN8pB9Tfiw85AUHSJGSwlVnsXlyHQ+LCtY7tNz+KOQQRHBQtvFemUE7/owFlAZAQX4ZhH7H8hq2bMUPNDHTHxAAAAJQGeyGpFf6zmbbzG6yAo5agB815PT4NIYmF3dRluJ6VKQoZFLbUAAACKQZrLSeEKUmUwUsb/x6mJAIU6OOaejU6jOCEaJq93ye5SDx3rwHIMSwuPbeEPkFOdU1+52yWbGid0yNY2LO00ed66IioKOPe6doVimRketNB1H1k9nKGLy24SyswHs7o8zUFs33/dimS2jZD9vA6qWcVwUdD9loJxKlVTdWWCY05iyBBnRI+HStFBAAAANwGe6mpFf9DFBS+peC/yI0i5BOXZ0Kb0Jtornv1yHk4b9Hwis05VfRcTVsB0+56ooWThOraMfVkAAABWQZrtSeEOiZTBRMb/++FR2EKSC17OilzphaOnJfb+b2v5DMpyB0LE4N5XbFR41vnnYVmX71a1TgMbvI4lkgGITPRwjYOjeSXOpCDC8xaYEVB5J45bl6EAAAAXAZ8MakV/jpI8CL128HDcV88rEWHEY1EAAABRQZsRSeEPJlMCN//7wZu36cWD2WN+3hTaxcsivWm9AgbVBHu/qp/N6ueiOGMMY/MGj/ljaJqYSQxBu7OQessdVRPOqKh231KwV4eUePFzdTULAAAANUGfL0URPN9/4LLaSW8U2BK02dj928jKfghGK1PeCa/SOciAikbQHsE/9VxvCM+EdP4P3D1QAAAAKAGfTnRFf4Z5voDwlwQuL1OkS/asB8/ACxEz3gJmWfQAj0g1qyhfBTkAAAAfAZ9QakV/isvpBhS2mK1fAdwbtVb5YgmyDpirrorcHAAAAEdBm1NJqEFomUwU8b/7w/gmRt4SJHtq/4CnNuN3iEw3xnwnXZMEbxFj861j5ny0cBWKgDehupwGrBYXAeK776eEXeWDNQa8lAAAACYBn3JqRX+NmH8JJKZfXQ2CWI5tPFL8LiOYvVF01NUI0yeIVFZaRQAAADBBm3VJ4QpSZTBSxv/74DjghR/zuyzH8Txd5XruzDOHv9RKkKOyGWsN/oiFcRlf2fEAAAAsAZ+UakV/jZiVZ6ef8q5U0fD2jgZreGAFGFOTGZSMcvYdOVWw2qIqwnqo5s0AAABPQZuXSeEOiZTBRMb/++FR2CH1+D8+1PRw2eC7cPzHLlSr4N5kdbXkTkaLMOtNZYDvN0G+BcnrfD919fNbcNCPlQajdZJesGGRw6UvTyq3mAAAABsBn7ZqRX+O5cRG3XQzJOw4cZkp7na6PmOyn1AAAABTQZu5SeEPJlMFPG/74DjgO8yArnokfcqY+3z7vgO3k81KQrKISGhM7SavcwZDj+Lf/9hb+KMuOUdKKwHApugri91Y0dc7DkP6jYD4KNCu/ZrJqskAAAAnAZ/YakV/hww5FAKjLlfq6Ljf/+1coHIsbogAeQqZ3SIn7qkD8EJiAAAAL0Gb2knhDyZTAjf/++AcgTWtokUyBqD6lhgvnfnJDd8Cx3z95yizbJFdyDlHg9jlAAAAT0Gb+0nhDyZTAjf/+8PuXcpxiwPb1g8GbiEE9IET2gffzB52f8EdNr3u9CTCJQDIJOrkhycCCgNaBG2HkcB5oPyRivqaDJ0+sBXnYINEToEAAABZQZoeSeEPJlMCN//7w+5IrxwHpn4mcnp3qC8i1PB17ucNb58sAdjpq0E8ox6I6nyB4tsTME3vaKjnv2HOUOKh1xJET0cCW0OPsSokk6kSsHBKYZjYC7L8lyAAAAA5QZ48RRE834ISoy161qatNvHcx2P0v+2schVH4rQF94qpR/q56ZD/z6y5NPUUhWN3Zgno8YGf8JZBAAAALAGeXWpFf4/0v5DikvANg1WUUB+NcTU/sSU3fSaYjuWP6ked8ltU06WZ15X9AAAAWEGaX0moQWiZTAjf/AKPrvBH9CwqRiXu1U258fBKSfqjqjBQd2g4Hb33Y/0j/xa9Lykql+WIXGWh5MAUZG4/duV0oRy89osMMauqyiZdhBQPluQQB/BxcR4AAAB9QZphSeEKUmUwURLG//vBRkdJOweuoGYmpH5xFZVLyGMWYq+kIwnpHpVXYdaSzJAWmGTPK1sFf6Znkp4MDLWKTIwxvdVJ52cC1rehXCiJGg1PAy47lCc5h0zQ7J3gBJjV4y3ubnbO8ylb6w+rjLP8+xSx/E5wYWiF5YEZn0kAAAAeAZ6AakV/isvoz+8v1Hb/P0q6S2yOKt6t2F+aO/ZDAAAAfkGahUnhDomUwI3/+8SDUbxwHf2NVhEpR83q6f+XXT0emIrLG/GjTA0ZQEsuf4ilRudQYS+kcaQP6NH/n47oeDxDgr+9I1JylC2mc1pGpx+CT9kp21tNNfIR9M3h5hKBoKhhH2mmDCHWwTlWZrU8XQRkROFjiGHTXpnI2TQPfwAAAE1BnqNFFTzff+Cy2kloXh3IfGPf7xzw0VsS6ephFl5SkLNj4sIWol/tW/UCeEA0PDXCAF0fB2wHJQCr+HvgtkKvByBuJ7pSQYYcBBA06QAAACsBnsJ0RX+GfrSKAZRt3+eUjEP/tYzGZCzSsy+iSHSWdDdvkdj7BDVckq9AAAAAHwGexGpFf4rL6TwoJvceXyel974RIYNEAeX014ZGez0AAABJQZrHSahBaJlMFPG//Xblp7QgISvXq/43tYQswnlCiPCtggpbgu8Gcd0XnodgB573BnhWBSscJCb+QLXOkHexkQ6chMI3P2elzQAAACoBnuZqRX+NmHWKNnyOWkElh6MU5PvSNQ8zujRaeaaC8dcoGksLYZPxasAAAAA1QZrpSeEKUmUwUsb/++A44EJXr12WY/1R9NcvW6Qg+sKr/ahi7wg47WJVdC+uliZS+rcEC7sAAAAwAZ8IakV/jZiVZ8+203fAhmP6SQ+/2wesW6W1rbx/GmEeycqK/2uzkQSaEM7KzV5AAAABAUGbCknhDomUwI3/wSZXK3peGISKDGuEfaX/ruGdURoMWCqwPVNEpOv+ST4slfLn0vfGnqFijtJwvKuTeo4gMwMlZ12WAx97eksVcC/xZyoYY8gEMuJoVZGXraJXsn2AWjQmljY++Yn0UGbvnrgKZ/T7K9u68fYiopDXks1XdSedLDYdCPWc87cOG2LaWaYKA+J/Rcrt93JbbJphn/RaqafF3b1J+VrUGnuLXN10mq9Ji4vhD2gcCVDgU8QP7fngj8k++0v6zvjDx4MM30Unq0Ro37cLz9Fp4dYmyVWi9EBFqKvmSgFrSgSu7qNhL23/SWg/5DKBGi7SSMBZOnX9WLFZAAAAs0GbK0nhDyZTAjf/8pF6MHN8xvhomJb9VTbd1oufHXVFvOUqiH8wtZul84FilJh+DP4Lre0QjSMquDwJB1Cd2xjlRf6f/EfPHC1oSOv7JiC13e/9H5Er1hNYw4FigKpqQ6OH9tzFrQITEYT/f3gAuicOt4g9cvzkdFsMYBTJT4tnA5QEdXJC8nEAXWHq3Dr4PL5tNfHtaPyUqNoDKx0fsKN4D/mNRoRUv/gxlV5JxljuF+HlAAAAnkGbT0nhDyZTAjf/0Y/fwfBPYIt2Ih+oVDWysEi1BfcGsNc42DZ+iX7NLlxL5cxymUJMWBhkKF68OtNiy9cj9+WIlc5Mhajqm9tEk7Q5HF5llNkQ03EIGx1nktYjyHoOPCRnfHGTijsZ7HClYxAmgkK+zsXNRz61pxBrMmmixPZw347116XCKl706m6iOHYY0uNsJzXzWYUyQe9L9FINAAAAMUGfbUURPN+sXvx5/CiDiEacNaS9U6HfVvnuP/a5UQJlGkFDYR2J/hU3JICcol6Lc4EAAAAnAZ+MdEV/jylvx+CPXBU+8gKTE/8VIElasIGY33llIOtASe54iQcmAAAAHQGfjmpFf9C2JxDQM2FrtfcernQzbgQo6cOgAiXIAAAAPEGbkUmoQWiZTBTxv/vFEDSrIEiP+dm2fRrEThE5lxbkgROmjfPtCbfhnR3G4ID1iPeJh7HY3Y16Odh6bQAAACMBn7BqRX+NmH8JJKZfXHhkLwpdR6P4XTfNRE6MatJ7PVwVpgAAADJBm7NJ4QpSZTBSxv/74DjghR/zvAWMRwWuIY0ILoghoK+2U2wGoNOA06nWHBwAcUpagAAAACsBn9JqRX+NmJVnp7PpUa8FkOcbCyJNunfz/M06nD3e8NZpyq2EIfrhbd7BAAAAT0Gb1UnhDomUwUTG//vB61Ydg3N6LfuunSF3AkvNLoQBhoTQ7aoVG6tsYzS+JiPd3mYuiPkzCyWnK401zAoTSSFfgBeqE3OpyvPlsMfUq1EAAAAVAZ/0akV/jX4K9aDUHTkIn8NGrTJxAAAAOEGb+EnhDyZTAjf/++A44DvMh5EHnmI+0qZng7g7n1G6u5/Aqyb8TemoaU8nwXBzvg6L4ceEadnwAAAALUGeFkURPN+D+24e0TinWPeueGoG4sapy+ObB3xDR00TdCay97Igcx39n3zDVgAAAB0BnjdqRX+HBFlCUE3nRNtGVVInyej3t3vz7uZ2wQAAAEtBmjpJqEFomUwU8b/8IvJDICKyxCcdzT/mboi+flPNWjFbInDCJbwHXv4CBeLXeVVCSD4rO9Def9jVQKiZ3a5b5Bie9xheg2SLH+QAAAAvAZ5ZakV/jY0N0MOVPt8h6WcoZV1S20MbM4DsMfLQUbNE8rolkKvz6B5iwJ8mvWAAAABXQZpcSeEKUmUwUsb//CJ/zICOuOx4SMR3pARtPhD3XWj3rBkO6OrvomzLoOGgKYE4i0wu3UHRLX77PtaCBTI4Tfpc3gtTPNcmsPDp9zCw//wDxv5JJCJZAAAAJwGee2pFf4x17dolDZL3wxmEFIgT8nhorh3Yq+I4oAacvjWdHciZNQAAADVBmn1J4Q6JlMCN//vuVw4Hc5yyvTItaaZzkkkNHLJkI0HEoJdIeD3dJbzkXmvtq7jUQZd45QAAADlBmp5J4Q8mUwI3//vgOOCIxgOqHGQ4dPM9u73R4ALKr9BQr/nU61eDXxry868wQq2NQQB2Cw7y8YEAAABaQZq/SeEPJlMCN//74VHYNYZ3DjtyB3KfJ/ehNVbvQ16Epn16L7ahC4CDOiOWjEm08H+S+lGobDN1/4ONf5WIIOky9RhRHLV6H2eUZ80HGf0olbdDoIPuaghAAAAATUGawUnhDyZTBRE8b/vgOOB3Ocvj1S4/IKY+GUDepFZA7JQ6D5zJgHdmU0gV4AyG3PnUUlwhcyp1sj/2sUr8E80bHPfsNfQw3EPqhkH3AAAALAGe4GpFf4cEPpZ+ayGHCrzmIFkFJJ5iOG3d7ko318enKcteEWyK6vDU0azwAAAAM0Ga4knhDyZTAjf/++A44Ilb6dcZ4GoPqU601qIovFnKC7lkIeBkBzqp//sZLPSn67pyKQAAAExBmwRJ4Q8mUwURPG/8A6IrIENlQqhLiwoaJhrnWZk18jQ7ylRVYz4q3Y43ERkjNUVHICbJU3W2fPPQUmUyKqzovlWUqaoQrc8y075BAAAAKgGfI2pFf42NDdDcJbtDdK7SWJH5HwS2s2t/lM7DAAfmUTKKRr+T4pF7/AAAAEZBmydJ4Q8mUwI3//wC0LvCDPTwaCjaXZ2btf9j1UTqkuJrlFvJUAop1+VQkGMNp2GtkRPC5V/i7vZhSjenabTSd4ZI6Z77AAAAHkGfRUURPN99TaObewep0G/CenyT8ihy+yfmI16uhQAAAC8Bn2ZqRX+NmJ1nz4vAdigsChyYtloC5NV8EcFNRJflhn+2noyyo44OPhMN8LgKTAAAAIpBm2lJqEFomUwU8b/7wetRFII6Zm6QidC0oo+c0Ns7riq1h4ZBWxZZD6JzeHG5pZuCUPFJCaRVGlMRP6GhKFXz9x3DH2KAgdWA1aMZo1pkrtviY17lc9l1KlNpbFrOjLf/sY3BYOhEQqMpFO89k6PKn8sRKcq75K2ccdT4r+p6kGpG4mZrOF/98uEAAAAZAZ+IakV/jX4K9aBpZl6QE2eOe/h3YVFcgAAAAEtBm4xJ4QpSZTAjf/vgOOB3OPztiP2ULSp7XrKrWw3l4Bc6H2W69Q6oSoF8O1hPrRXjY7we9rAJKPYETi17uxoi3v25T98McjvFdpkAAAA0QZ+qRTRM331NqbcLGtWbMNY3WScBZQDj4EVhJKd+eBPd5akdHADpGT8hVb2bsA3vvcUMnQAAADMBn8tqRX+HBFjWKW0eRGEWy1dqsM/Zh2+r0j8TMNGs/fPMii2zXF0N6KhYyAEYzOLy8UEAAADpQZvQSahBaJlMCN/HqbEHRtdjNTKn+1Ks8Fd7U1CyW72yxoAL49N0mQ/rxd6hIp3QUmPFE+CBT6W6w2fg+GvkDoA5rxlfTH5iWO7O7X1aIQrUycIpZ26bWxdE69CfvW5Ump7bO49JM3Ojp1ba+sFXZJgF2H/6c/uab0RGBZJFVPWM6/2iD4z0L1ee78rMROduc0TjmUMv4F5e8iqzbVmWZxa47HR2pX34hnX7bqBD1Y05md2DIEIMx7XWDS2zFGguo8fbZRpICoATJ0iBpYXY65NY+onyXag/KvGTn9dqOkIZN5bgNwodskEAAAAwQZ/uRREs39BT+V5BdRxLLjvRBZ/8QK0LAnU4ryUT3nxiHiAhHRZibTVIPm+0xrfhAAAAbAGeDXRFf9BX5Uw8x8UvZrjTOa6N0KNc735HzeH8dCL7LysfJAtHc6CJl64MM87h+XsJUfWPW3okcUNQiLbRiB7rS8Stpkk7rAmn2NBtaNQkEAWDZlBKfeymcqvzqnPiuE+on+yiIr3BB71tTAAAACYBng9qRX+NxKann1KePz6kv+QI+9/UjpccvAs8AqKrEhiFftqoFwAAADRBmhFJqEFsmUwI38oIRG/MCA7nH4jDF8og0z98ssIrvm0ivszdpDx+Nnn5Qb+Dz0hEkbmcAAAAj0GaM0nhClJlMFFSxv/HrNHv/cjBycDl2vih6keQw56fCiMs0YdNmpp7imPBF1y/lLKBnxMxn3M9E+3yBzgA4qoPmgxAa8dNdml0VORbm0KUx6CwEuL013ZY0JgiO+ELoSSDqSlaWOxREn4lWygn7nk/erfNPMDpvIC2jYz5OlVHNcLkNnKjrSDpnPb4eSDAAAAAGQGeUmpFf9DEtYTuNuuc8MFIfiTupN1BdFcAAABQQZpXSeEOiZTAjf/5KwLantfVy3VySM0SMHvsr0YmBrv0UL6DlEtaxQQHzUXE30YsdNChYm2O6V+hEk4F96/91dLwu/ndfGK7GATkcf47DIEAAAAxQZ51RRU833/gstpJbxTdFgmzsfvEjjhs3dxWp4ZZ4kPG8NZeKU7GYzN2OvKw1tlnqQAAACUBnpR0RX+NmIM6D8IxZFFGEALyP/EQHxDjOvu6mfZR+X/XdtvYAAAAHgGelmpFf4rL6VHGE31hHDEGQjtVce1yEtA1rtWqwAAAAE5BmplJqEFomUwU8b/77lcOCFH/OzbPo1iJwiLqg8angbnwoKqYau4/Tr0/Su9hg7+dXmJB7llgNXjWmAvP+6nn9ZyA2iVrXntcHxIEctUAAAAkAZ64akV/jZh1ijZ8jl2T9B54JFoOBhttQMJEZVln9G5lKq8BAAAFQmWIhAAn/74fcEJGW81yyeQ3yjIJuOQipq5DurWEjnW3UorRPL+Gmp7W/eH/ZijezhQ2AIPIQl4aHdQM9A6M5/MKi9xGlN69wV4R5VA4TQl0Sllg8188tv6vFr5S+P9cmRKtZwlxsQrO/PUH7ZGWYYybHVXFlEMY6Dl5808LYw/EiDr5JaoSWsHu4ebowTZo7v7qaSld+s+T31N91+CVDr6kktqJbNFOPAo3BvxQkltR4XwG+/rxvXbxkWnnZ4k7z5WpzDWYcBFPvUfTgrd7n6SLHLTp7bJf8LFMRC7GnoqmMyaPAzg75u8JidsVPIR29wsvpk/Mt9R6nnF5IJistydnX1nYc673c2swNm0ubNLgupd+k5ttRQ152rHLvBca63qjaEYz5bO6uRdZGmNbrdIR9DOFJYzEJ3T0VRVhLzRbgQhmzINEkMybfXp29t1bR5qya1mRWgLdBC5l3drBLQCcVoq19LY64FqmtYeURGFQ+dxCLISD69Irv9a0uyOxS7lVS3qeo5/gPB9YhIYgKQn6bpUQv8jq9Ep36A1DYX5VFcN9XqGkifkeY7ZbzJGH1tSe7aDoRQIiDO6HnmttmwDbytG0SzChvNEtecWmFp6+Ng7jqZmJ+LpXMqaVGcBbl1ehkdm5Vr84GAQKEq97RhEnaEuoPQsiUhi7y5CzHQA/q9/VpbR1zqAOP/78KNEfdl56fss7Y4UfTIZBj/Ef7ScMR78inLAl3j2HLP1BrnE2c1Zfp3fmi51t7IntWomhhRLlKfnfJrB056GmNbnqvtmZCvlymbrND2EqUvjJ1EJ+n02dodaM+1EsP8F2MT2Y4tjt1aBsu9gTs10CfVny7PzZrIe7ZNf33r8LTwpXSowlA4nQt9opFQ/ZJxaz8Upq6XgXe6Uqzu7x6Z4G+LEevAbdMu5y5ewVRMuV7HqXwlZX3FmRKJTwfg2u/onfLxKJGGPlejg4w6xVcIvQWpCCY94rqqRweESq52TT21mJ3ZWZPYd/bW9E4HoKBWndnm57+UuTSHIPQ35BT9W1Xnl9r9CViMvIQ8AUlOp7+ZZmJtCkSP14cMcCJhNiLGBPwta/Oyraa2eYe4xSkWML2OCXs+aPKBwLcn9R4ao5D1S55sxD7n6+Mg+z9jwufgfWX3u9ii7EiGIFM7wV1HcRH3cSIQZVgHx3bNHeN9HrMdz5R4HKc+XWQ6nfPa5qVL+jrteHlNOWPkjRWBZgYim28+DOXpSL63VQKM8wCalx/7UFzpmPLYi2k25Tn4L4RY2kTpMyq852ckgqQ/JuHNjmCM0OZxxusMdxMp9ba0q6z4Fov+fjMAzn8vKAAJKzFCfdFJg2MeZhfwlKaitFkiyEY8mnat72LBWhX8y2pFqa9wBfrk6vSpsSIPf1ItXOV5dsbkmQaOOvVHpr39w9kGJ06MPgDiO7B/qm2Bs3AlNEuL+WG8zu/vIRJdDz+RTRIk7vLQr0Rq8jEQWZ/euf8FD6BUNHSXZGM949sWmTK0eoMcmoTzlyGGCQgdVdRYqeVPDF7SHBzjIXvdTeaPeF8rdbHCdMCUN23hsQBtwTi/WxUcz0MCULeSfoCuRD9EomV9UBsyxfCUqSYYFieuAbg/IP3y2xj+T/9E9kqvzPRPQa2y1WxDCMsnMIy37FnZ/7LXXOsxfbVeew3U1YASM9XkT/y50YD3lR6WGJ8MYfYM98D8exb8s0epkXWDnGeNbeKqayRpRsYTvsgzojgHxpi5+wGkEyHCLHk/j2PjOYfUfWmKhOwj053UqWFuvky53Z+AuX9yWhjuzKAAAAOkGaIWxG//vEnVTkBINY1xBnYqV2GYveJxG6pO1ytvZdPbEXt5LdA7TGVX0iYnM5z2Z+J5AJNCh79n0AAACQQZpDPCGTKYRvygg25A/mod1FrlSV9Bg4fSRIk5ApsXKnR4Lw1hYENtNp2IdrhL36xanS2cwpV3ghaBQ0ENK/+61iJnFzrglvYgyAeUPfvh4+H91OONFJ0BQNQP7S+T1dX3IH4RbbT1dKFwTvZ3ZM7M9HZq2WxWBFVnr9ZP/zxD+foUgnXU6TFFBNN/llQnxAAAAAGQGeYmpFf5XT+htGpxm+ttzhrI2RGOltiWEAAABzQZplSeEPJlMFPG9uJuUZbUPQJIkFcwKAtahG4ohsmW9O7pJ3QXEf2moUd4qKRCCL0ZKGjcHbIVUkGfh0CDhI2ibslmbDuzOMAhXG/zv/6v/5h1pF6Fnk4OU0BdEQ7aG8XM909Bk3XHm49yHVcqbBFgXZxQAAACoBnoRqRX+RRVaIVD1YrnHAZf+m1ih7YZmVdEVTUohPyoJ0YoX7wwFZW8AAAAAvQZqGSeEPJlMCN//74ByBNa2iRTIKCh6OJLaoh8ps3rY7B+aDjQEslblmkDmrpFMAAABYQZqnSeEPJlMCN//7/0a5QmmqHJOeKUzB2eNo5H9glW5oLygWzEVYpridRvXn+2AASXaouZZC4r0hn95olSpqHwrE6msLwa+V4Ttq+yvbjDu1Y6/bkwQegAAAAGFBmspJ4Q8mUwI3/8oIQPYmOP/wRrPxM5PH+fJ0mFWAYBFh1Pe6xJT8//aEn0KlIlhE0/6putAFmsffwcLGQqaabW0OMe6iwCBaM3Kq2If4TnEQ/lSd0QI6JJrQELGDib/jAAAAO0Ge6EURPN+Ecn5eBysgKBkhDY8qsL/25Qam4wDRxrYU5Di/5y+wAXkIetqyGrdnqZuLCcwFByy7VklxAAAAKAGfCWpFf42Vo24OWiO+wC6Bf1+vC57IXLwDKAPdiUIokWez79lHX4AAAABIQZsLSahBaJlMCN/7w6dR5vCGznNtXd5bO5Zy4LxzmfFyTxf8VH6vp60ELz2qk/FRPVhaQsbXkKgrOqn/+1roIltiQuL600R9AAAAakGbLUnhClJlMFESxv/74VOAECNfo3TrZkzqu0KkrOZ9CVw2gjMRNMj1E9vxW5V1Ecvm/9qH28Rc3ffR5r+zqVMsLdHEwsswCE/8g29m7hFIyoSbEObNEnebhqoXvPeZCDCweyiSrLwsRskAAAAXAZ9MakV/jpI8CNGNvrae4VovZPXKoXMAAABnQZtRSeEOiZTAjf/KCEDu+rkuXwR0ltpBueVY4aihZ6XD3gGzDCyVqxs5vXh7SVn70nLyVRSeCD7n1+k76tXWOU8Q8cvlYby2Hnvldl8uqwEjayKEwX50N3/0maF6ny5kGq0f7Ao02QAAADxBn29FFTzff+Dbsn0Iut+RPjHv9yMXPXi4QWBtUuz84rCf2ne6MMkRmFjCTOGYGkE76dyp3SNBzeilrswAAAApAZ+OdEV/jZiDOWM1Uy9Vu0nWL0V/4HEe7fSSKWekx5hQYtiC/xgqvYEAAAAjAZ+QakV/isvo1lhN9Vf9YXwem/sU/GB4tK3zn/fFUPpEfgUAAABJQZuTSahBaJlMFPG//APDMcCEr16wAPFV1G7zGhko0fBO1LBEbSo51CNZnpYS3MHF50L4MAKNELOJvnOk8fJx9M+aEnAehC4L4QAAADcBn7JqRX+HDf42NiA1e6IoNY9dd0+pse3XDs23QkJmYymj6mByBQBlbsNBH/n+C/Zf0yypZhu9AAAANUGbtUnhClJlMFLG//vgOOBCV69dm+FVLfTWcUGQr+lSqCflwO8IOO1iVXQvrpYmUcqC/Y4LAAAALQGf1GpFf48rV2va6Z6CX7HZbnuYgcin2b1jZsuvhj8zrWfr2ZgMZiNJsrNVEAAAARJBm9dJ4Q6JlMFExv/BJu5KzDErXKYS1RifWXWjNyUkZ7LuAhIUnAIaOsVEMMQXit7waW63wzdcUIeRyXc0ns6O4xhmDwZ3ukGLdi9EU7w+8z26SXTja3sGWB+uU2PQaFr9su4rpTbgme55GdSV5Hium1b/ZY7wn0ziN7L7ptbE80Ax28UKDugysLJcSwfJeTU+FsCjXuJC35Kp+sq+nt2e0C5pikCgY4m4RGYJRGW+kGDU+yhW9YEtPxxAt90wlj4Nc2+02VtndN2Z1aSOvivKWQperyfJK+RiYH2Ua4M1Mbvt+F/gIq/S9qIVS9LvpsrQ+7+gxnXyJvs7t2yv/J7ljPiE5MYPyl8h+jNIjrv8R+MhAAAAUwGf9mpFf9Ei2Luv/RNJcaZxTnDa10UNbuQPAlqODVB2VxsIpmJnfjPkMJPHAq7sCkJAYLQ9TOcWdYY/w4T9TToTqOJl5XsxpaKH0nZtSdWfNN6EAAAAUEGb+0nhDyZTAjf/+8NYd2niwO+rMxV+zWAA+1JtHTaumj+P2U+PhhfgTzByUQbNihyr6FdaXSQST2LLyoRUr2SQPrdTZkfv31QscDSS0oRnAAAAM0GeGUURPN+E8DKIQ4/14uOVqEAMADp0zauP0LtNNslqo9KyDu+ewIdMlXYKjxCCIZ2RgQAAACYBnjh0RX+NmIM6D/Z46G0zEmw183/eHh4odUNrt8tMZebEjulWUwAAABwBnjpqRX+Ky+k78Xe04HkuIKPhhbbazcKxG4nlAAAAi0GaPUmoQWiZTBTxv8epXycoOFMLM/d1FqVzTzgCDwFATJBP1U6ge28/eu//0BsTgIFsoE4M0gMczwnHzJtkURt64dKffbzBHukviZ6wJ1DD10/gHeIjw54CCzmpIwySvy0eKrn1sxtgjxrAjk40cMg2w3WH2SwB2gwRBdRsg0MrzvNgn9GcNQO1ZTQAAAAnAZ5cakV/0J4+WxsQPcekdWueGIREStNdZlzFmXh0KEniThU6N9wlAAAASUGaQUnhClJlMCN/++FR2ESuhE+RatyVDTTzS0MEydH3+Vpq2cFMFPi5QKgi9czxuhaGemqsD7/JcIUhVdt886og+ljZbrvCIWEAAAAiQZ5/RTRM34OpqLzLMBpZQ293KhCxlhbjxhPOxHd4q2TqwAAAACsBnp50RX+NmIKsD88pu+EZBTc434H+uuuvMbHJAB0tuUvJ1rP2bn0fEH47AAAAFAGegGpFf41+CvWgaWHE5ho1aZOBAAAAOkGahEmoQWiZTAjf++A44DvMfnQGvr5aabtMPFRgZrO0PAPTbXug3OYp1ediQc1+3GfuUR2A+YYULrgAAAAtQZ6iRREs331NqcrGn20AwIJmTxeb5XHCMknUBifcKi/SGXaUJ0L8hQfqmM7vAAAAHAGew2pFf4cEWNYpbR5EYMikZF8nmrBIkdnpZdsAAABOQZrGSahBbJlMFExv/AOiKyAj+gnj0a3Gp6zI2CjNAAJIWgyCThvil0kvhxqkZ6Ht1BxxqXfrjYLInpyi9Q1SHRusH5zdhyC+Nu7rF+RwAAAALQGe5WpFf42NDdDDlJ55Cx+Tl+h4NevQbGGML2wsKXdErYWj9A80AousAs0tNwAAAFhBmuhJ4QpSZTBSxv/7xrXt8jlCNDUuxflUiLTs/g36A9cpqCK92W6XZsUT9pVvqv4iI3XXl/siKAd7P64Kbe2lHxIRZwTXchYXnFFv6PVahPdtBRwdIEdgAAAAJwGfB2pFf4x17dokp5DlQDbrpLgfa2hCiFMmKnp8wA05fGs6beVH/wAAADVBmwlJ4Q6JlMCN//wDwzHAd5j8QdbEjP/MD54/nxaJsM3etRkV7PtoitdWOIiK4oua7/OJmQAAADRBmypJ4Q8mUwI3//vgOOBDtkfUu7/INxPIL+S5L2s8+0BGgDXxrzOfaBdvcLkOu9kv7y8YAAAAhUGbTUnhDyZTAjf//CIBv4QEgt0jB6sYApnNKebF3UduVNlxp8UDV72R8VAklM98n5lOWEw+T8M4WRe//ioMQECXsBcH5k+vJghjfT5jqlmoeL/l20XeVcwcUIuafvxnJkzhBlcn/d9IuK1j3bCY8QS0oTVT1N2fu8R5Tv1RgAbc5tBoWuAAAAAhQZ9rRRE834H4ctY9gURT1dXKSAIxLkivCpURU347a3c3AAAALwGfjGpFf42YllmapQL/etyKspbtOBGV8f4EVMMX0hi4K6l6e+YACOn1K36gQd74AAAAMkGbjkmoQWiZTAjf++A44EPraJILnXuLxMWdODvz9zlGHYLFuS/WCzov37vtypupywZRAAAASkGbsEnhClJlMFESxv/77lcOA9jsDFR8Xdn5jdZFKfM583oBslJ49SZxnRGp+moI9oSjotvQRWZD8gOf558xXUMpWRXJ/cWn7ZTAAAAALwGfz2pFf42YnyFLdokqnTK0mFdGbVl5+BNzXz+jPJsmLNM19mnectkM0k+LwyFJAAAAPkGb00nhDomUwI3//ALQu8EXYDhpIvttv8TbZDOP8OLBMAO70WexTFz5+cnteMdktY7P65A4tZlIDdklK0GBAAAAH0Gf8UUVPN99TaObewetr3iFfIM29vQ4/CELIo5doCEAAAAtAZ4SakV/hw3+jEmnmFEtcwfltCMFzrX9aTK9FUCy1cKpO11Befk4O1gAKZFBAAAAe0GaFUmoQWiZTBTxv/vBGT5YQIcUNUZpARyxqc+rS5mmS+x2rMYuZzxaaNzjVUhTcMNlid1G47Ff7Gber1siDAOxLdQ+mOx+lSM//HIlVdSI36FhpMg3eSscveUe/EzN0NXsxq+B/YFxJBXKltyUDxpD4j5KmVec/I5HTgAAABoBnjRqRX+O5cRFOX/rc19dAGfoPBoZZ9JRwAAAATFBmjlJ4QpSZTAjf9RUJUwmPx6d1qkD5IiUog3fdDYG+lhD6PhB6357xt0eAhnE79OehpzS9hNg+0rld2gpMfPN78NladMHIMJnu+IgwcdY5qjgIx0rHEFP382Z+0LcY8oZUMOK+1SYe1NyN3/A/j9w4eoQ3PNSo4Ux/LOL3AlYQIiar3GTf0B95Y8t+So3rBxsftentdD21ZJ7hAbvgAAT3Nv7jjPdzJ/O1zUUPXbzz+p7nvdFPTkoO1Dqocj0bjwXKlp/xDJw9QeJ9FCKJkVoMOYhplLX41md7oQ/Y+qGl5+5p8L8BxRb1OTUS0Pldi3oqvL9d/Zrz+oryE/RRst8Em/0e7bxnWHG6Mx81D5GpO0jZK4XqbwX+2UQU4boZy2cX+Sjuvn0MsMkR236Z2K1zQAAAEVBnldFNEzf1+oL/lUB43iexxWj6oMO+zCK0ceAc+a/F/PKiFsvR2BDrZ2PrBeRSYr2NA9GNN54oqHc2LpSV/DC5yGHm8AAAAAnAZ52dEV/jZiDOV4WjRVuqfuRLr/L1bWuM4bOVk+yHw8xIsSbICVWAAAAKQGeeGpFf95p6Ti9RhN9YQqEMUYPvg2cqLZ+fsaBkkhn7KTT2jvNMkWBAAAAgEGae0moQWiZTBTxv/KSaxhdVGtNZuOiRf+JqEcNFkbAv/3toQnjkNAaWg4K8+iRR1zNbt1yIR3tJg8Hh+JLp9Fo/JoyYU3FI5nF35wWUzt3yx1MISi809oOoCzZrj//fWbdhPS2477QVN0UR6uqQ2+uEBXxub0LmMRP6aOCVtOBAAAAJwGemmpFf6zfg2l3Ygdwg0DWPIpQmdHx99XujRSHBNtWXm69NnlO/wAAAEJBmp1J4QpSZTBSxv/74DjghR/zgYSVgQLW28saHEGZfNPuk6gQX9d7Fl5UDqoN11nbQFj7zXG9wSS5vhSv3gAKUUAAAAAwAZ68akV/jZiVZ8+10ObdgxOx3183mwEhK7HkZqcvSQ1ycqtgp8jF+U3zQBWvpsMwAAAAU0Gav0nhDomUwUTG//vB61Ydg9A/wlrkMlWCPDd7QTOqj1/1NiCSa8ruC8UBK8w3Ph2mInOfYuW+cn88B5LL/toEaKTUnUAV8fIMLomlaMvI1rt7AAAAFgGe3mpFf46SPAi9fOZie0tFVOsqcMAAAACWQZrCSeEPJlMCN//HqYcZuGVy+Xdov07HU6KX5bS+1WmnLQvbAvY9UlOzUvt+HiM8nLp39Z11G9J72MQOlYTjkOf+SeXNLe5vidAHlSjT/Ore0st82H8qHv9wvch4z033oRI1qboMoq9jVzWQ3z/b49jZCWWEEv8uwxf9swldbmArV1EwuuF91pmeCP9WmkCZPwED0CvAAAAALUGe4EURPN/QWwldz0TioUzxJwTL1DjVPQjAq9Y54uZHmIdkd8NkQOEM7KiMsQAAAB4BnwFqRX+HBFltjCbzyK+FRKnf/mLOjOcjzR76r8AAAABFQZsGSahBaJlMCN/9YpHGMYUb5j/NfcgCQnnnUME/+uYEFab7R5PJRwSPzdjTXR8fgJDcETNoWNTQyIj4LAz2Alc7KxLVAAAALkGfJEURLN+CHmGW8iX1nC8et2f/B+8KsVycqqolFcdZbi0GJHLqpjMjlxYW8t0AAAAsAZ9DdEV/jZiI7Ev5h6E6AXxVvyU+Xec7Xzf+Ujz4nooqJFL612j/wKfuy7gAAAAiAZ9FakV/jZXGc7dkFSNTQS86kNC5UhVk5V114XqW+n6aNwAAAC1Bm0dJqEFsmUwI3/vuVw4DvMd9xLBIuaZ/zxsUAT+oe+xkeDynj6QmOuSirEAAAABtQZtpSeEKUmUwUVLG//vhUdgbBTvAngor8pfCekIPa/E8Y+CM1csCM1vgERqhAjVxRBDaVo8NX0IcDkRQM7Fc+VrjXHrEEHSAqDVNvfJe+IHgsFDNZNns8smESREry41It9enn5zKZn+/bjH+QQAAABgBn4hqRX+Nfgr1oNIjL185KPI+d0MY3WEAAABQQZuLSeEOiZTBRMb/++A44DvMfbdU1ETVNwuYBfSvlXeOFbA6GbqsdQXHZrecjnlZBPP40eDLe0BA3X8hhdWV4c0xWlDkOf5flpvIBUQRHYAAAAAoAZ+qakV/jnhriwfjEBC0Cyoph0KvuuH/nOiAglQQP/F8Pzo5XgfuSAAAAEBBm6xJ4Q8mUwI3//wp75kBPbfTrjPAkRSSgH3g76N+IPxHc4xKee632cn0Is7nuf/sPmuSHgjxKqRUscLPVxNhAAAAWkGbzUnhDyZTAjf/+8Gbu08WB32S1i5TV1gbt9xeW+dzaP8gfVQwdBZ2FJC6dNZtf8f0Hg+4rzvoWkTdTJuDVTqcnkDXxrzOGZQBT4SetMErEEP9qYQ3b8ROgAAAAF9Bm/BJ4Q8mUwI3//vD7kivHAemh1CS8p9HXZNoOrlrUrukxAJgTBabe3Uq+xKUtXAoE33EFB1Jjy8zf4fiuWPtUn1WH8f/CY8QlZPFmMmcQSG4fjuXN/B+qNi/1vyFEQAAADVBng5FETzfhPBPL0/8xgE055y0cUfkH7gbD0hvkO9CUgUIYF9ILO4onW0KnZIlS4Q5+PTxIAAAACgBni9qRX+NlaNuljs9U7M/SW9U7+gZciTVHqQBZINn/LlwAI3T67KhAAAAQ0GaMUmoQWiZTAjf+8GQRIyA+EC0i53Uv0awLgBrmCYaTtddMfGzDz0tG737dau/+uhw5YMVYL7inRfSxfcFiKBNgykAAABfQZpTSeEKUmUwURLG//vhUdg1hne5jlt8Hnsw/yzG9BlYA+2BCvdNQCaqq/vU4NJus/26Yiqay8cHVINjOHjrvNVNNQMyKgANVegEHkn0Atv5BheYp5J6LroUHacyJH0AAAAaAZ5yakV/isvpAo+X7EhbLzLMVxMoH0gf/oEAAABsQZp3SeEOiZTAjf/7wZu7TxYO5AYsHfwK5RdvkSVEEn21E6ZLxQXmV2KUm/dt5YseWbLYmF6cycseQ/zRHyWaDXSPfsWYc9/L5FXY8bkH+03r0tQyCwxscnMoNSAc8JSQ1LuDhPTTMWEyb+wwAAAAQEGelUUVPN+B+HGW0ktTlKD5U7sftibyzFU0itTwzBOsr8jinvAm3X+xweVleORM4VRCN4J1niSdrhdyNeGimYAAAAAoAZ60dEV/hn60g6kYGvRqVwcf/9DU0cnE/S71LZRHgHSKCWLRPYzu+wAAAB8BnrZqRX+Ky+kGFLaTu3V1h5vNMrHepiZHt++d3TYwAAAAUkGauUmoQWiZTBTxv/wIE8MgSI/5zkdt2xF39AwC2NCzH/ix26Qf9Z2is2kaMgsh4st50ymHALuUDhcWz8HLl5Mv7cFHUPup58bPZYY2bD/A0L4AAAAqAZ7YakV/jZh1ijZ8jlpBJYejDv1xMTNJL8wwhvpmAwdAZHfqRUgGccktAAAAwEGa3UnhClJlMCN/wWJmEnMThj01oz/8HmHqSy1C8Xnv2nkeHB2q5boSEX7n9tutZzT+CbMC2woNyw43AypzG01uaeEWSh0sldqT/FFT7HItLnN5/sxDby4ws3/+X5ukIG6x07NMTkKBgu9pujusd2uIIBTzr4TtqDAIoj5L/81usPjC+yPrZ6BqtdjHQj00WB66maxuzdpr2LURBzdggU7/6dRUoibhHrHu27B3o2Nx4NpnulOOXMgFyxfruOZowQAAACxBnvtFNEzf1+r+1FYhZh13lr8fS3v18bH9CE4mbV7h417LGGti+gpJ8stneQAAADEBnxp0RX+NmIKr+C9nWHyDrxwMHXm6plPWpdYz9mPCwis05VbBSSZGa7pefK68Un2SAAAAYAGfHGpFf9E0GDu+ktvFMGyuyP+/q8a08x7N3S8hze5kbadJPQ4J/VcgfzmVJ+wGbfISiHM3iOG4kNfnUv4V+1AHLRdp1XpMOeQfEcujnC8EllvBq6VYdpBNSXI+IGkPIgAAAE1BmwBJqEFomUwI3/kqHHwSjggzj87YkHr4a87jSRladhRJ0sPf1vzUl8LaQuieyIjiu6qDqBhKWXMid6Xg8mRn7LGXUL7awd9GqdtR0QAAADBBnz5FESzfg/tuGYHKjTB5ZIF5DlRZ/mG18mLNwoqUKDqUjlTWk6UxETTkhiQI1o0AAAAmAZ9fakV/hwRZbYwm8UV+fNHe+WvQ7LE1dhsl0Kkl2dgkJM2xQPQAAACVQZtESahBbJlMCN/HqjpQil7yyLFLFu10qfJCouDtF1Xmx2Qi5MZ8wiTfvWELZh8dDesXTIOc3puWHO2V2s4k4HdAKOR4ngqVZUA0LYM5DJacVR8U9Naj7I7kVfa6bltDE6u1WaZn72qoZRJsNFn5BmTYbmEO/mI8pJzVOUzeRtqmdX+yIakKz1bUKdoGm4/J27cXeasAAAAxQZ9iRRUs36xetRXUJJO0H33aYF1R/wgPQCJEkQv6Oc57u5yk1XK8xotcqZMjal7waAAAADIBn4F0RX+P9L+RFJf/ZGxnwlZxxRpnlAdZ4MTNBABowPPX1SEIZbQRP1HKfV32CEsflwAAACUBn4NqRX/Z96jvz6lPIdxMK5ZxEfdrk4JnwEbzrXk8OtLZKf2BAAAAN0GbhkmoQWyZTBRMb/vuVw4DvMhdaXt49Ohno7dC9VWUVXFssBv/tAEi1GxHEtwjRSRy8OIHuSYAAAAXAZ+lakV/hwRY0JTl1qnyRaLRk+dkPKEAAABqQZunSeEKUmUwI3/9Y3sZBuQxneh/rGW4lE9nuaCObjg+K6/Ef2UAcSMCaL0//e80Z5sMPf/FLcue0CI5qB5M2c2QWFN1gHozimkN8YvezsbDlRHLpDuEB8yJGwaUJGSNmiArYmFEwpDpPAAAAD1Bm8pJ4Q6JlMCN//1ikeX7HlErA1F/NBuaVPRthECIky6bnR/zbhvjqszzH8XsGKtqLqtEjwRPCnmhPsXXAAAAMEGf6EURPN+D+24ZgcqNMHlkgXkOV4kgZTIIkxduFFShDLukcqa0nSuPAX02qo7ZcQAAAB0BnglqRX+HBFltjCbxRX580d76sSrub4Ej3N0tuwAAAEtBmgxJqEFomUwU8b/8CAhiLA9lUJ1ZhKmJlcycBCqNW82vUrKjwfRLn3brEJAyFEiDnxP6vgG4feHLTjXsLOkn44268wtNImWMf5AAAAAtAZ4rakV/jYz8hJwlu0SinQC47vIeaSqHA97ANM9n5sjiQxvenF11Aous9wBFAAAAUkGaLknhClJlMFLG//vDoZtU5Qj04CcZokNlVJpbn+xpxn5vnomGzZJoOL36cb95LTn2GQWez/KCm44eCBSdizIJg/0bpdMf7OAdSKNj8pTKjsAAAAAoAZ5NakV/jHXt2iSnj8+pK5UPI/1KAk39osMlGVfJu+lmZ0egJ6vDwQAAADRBmk9J4Q6JlMCN//wDwzHAd5j8NJF8og0066g+rXTd5exDKM2MCnecCF4E6scY10T/iZu4AAAASEGacEnhDyZTAjf//WN7F6EssGLL9aRq9GMDYogP8Xzv0eGCplNOy5xCXRSUZ3whgiU8wFDYxIZ/nwDS8NdljQdsca56LkdmQQAAAHxBmpNJ4Q8mUwI3//vrhHYHuWOrKSbgXljf8SHjA7pt3GOVNXvZVrQnA9aKMhM1s43sy2s+zopKgZQuR3OXh19uu03zK43qCzT1h20dR/R17wlmMK6XGC4UZ+neex19jtPD/686wWjM+ofcRmWrGJuZFnTOmBpsxXPMq7oNAAAAHUGesUURPN+B+FaPuhgDkzcf9gPd1YURPpDElZWBAAAALQGe0mpFf42Yll2maqZeq3XThpLr/MMrxEP4wtMxgfA47nZIAB0Jg2xgA4T3wQAAADdBmtRJqEFomUwI3/vgOOCJW+nXGeBH0bjrWPxLGdQkrwjNTNgde2mMiCqnPmK8M/7dGPanLBlAAAAASEGa9knhClJlMFESxv/8A6IrIEOZqgng5VRuFA82zmJM8KL0rKYQRJ7Z23FpDej7QMQVdrodIfe0IE+0l4nF1TNAXWq6QYt1HQAAAC0BnxVqRX+NmJ8hS3aWVfP2c1jQ9YmStz/ocuUcfXzyBUtM57gT1Y3KU8wl/ZkAAABKQZsaSeEOiZTAjf/76z44DvMd9wsd+BXDr5TM0gGx6Y5mxrdaiR+xZhz38vjjy0wqFQLRtVm6z4tKSoPcbC90J+/nHLhkajPTm4AAAAA5QZ84RRU834HLEsIWvPinvfHxN11tVvEp3U0V5MHr+u9+xSzIWHqzyYWqp/rDE7G1eS8KR5LYfoqgAAAAJAGfV3RFf4kZCgkjcPXSmeg5eQB5cdFgpKJnWyRjJ0LC8/6RsQAAABgBn1lqRX+HBFjQlOXXoyp+tUCmnUXoUiEAAABUQZtbSahBaJlMCN/74VHYQpILVM82IUwsR3tP1BZNdPf7nkM/rWUZ7Vjytb4wtxQrFU0OUMs2AySifB/5qpsYFfAbi+7wg9W7xVeK6NYvm492GehPAAAASkGbfknhClJlMCN/++A44DvMfbctm/iNKou1bR3qaQPdbYPAwRr192Z0c7+EQI4zvYW462uiNmrbmQcLA9GkcsKw1co4BTWHqy2AAAAAMEGfnEU0TN+D+24e0TioUzxNily9Go1T0H+2EjDCXB2dGC0c6xM99Lg6Jduz8xAlJwAAACEBn71qRX+HBFltjCbzxryJ9mcJ8tdrhFmiT0jMMVk0HmEAAAD/QZuiSahBaJlMCN/BI78lscgq/Rbp6OxwvfwmtoGweCC3o0q532h4ICIwjwcptqf9EK/+79F9P8RvOlrjacSPiqIg7tu0F2k6OuQHzHZomSpzp/xnEFVc3YpD8D7mVeTAU9Xkf0URAjJ5lfHVioB1x0ui2g+17Y3BGxACMJ3vWNmNMsSB6yWO0rku1Z5c+jFLkY2xMy95ck9Xi4MpgR/rgU5kie+i3djqs+2ji7H6EQ+ogtPCUzbkhOsAieuOLwevqHshJvQds8aT+lB97TlMWAfZprX9Mrv3feABeyzfKAEPhJ3BSA1JjxJgX63hdba2srRrtpcoYk81dOBi6AEYAAAALUGfwEURLN/QU/leQXMu9VvVKf5/iBcuPswhbCCq/njMDoMNrlzDXkngFpZCXAAAAHEBn/90RX/WyxZvpLbxTBsrsj/wSTqMIPKPLDvNeh/1F4bmWrWmiWrQyZcVh27RFCrcMqxQPKxSxC/gK2atA4KwaYxGxim8cYsJdnYB2lEgpZEZ9RrWqfitTU8KyDfSoH8YyYLFXyK9azzCP++iehmJGQAAACUBn+FqRX+Ky9y9AWM9M+//Zr43fS7wb55u5yXHFaE5OBuT0DsMAAAAREGb40moQWyZTAjf87uziMWmfzdbE8J87clOeHoYJgH3Cx3Zbf5fu+tjc/rFA2s6Zalj2+4xEwd5fOTh9i8yAEFuj/ORAAAARUGaBUnhClJlMFFSxv/74VHYOCF6yTT6YYA389+T2IXKVeD0g4f5YiUfdwqkE9LXLSW0Wg9bBFbTjG9iud0FpzmEuq0RQQAAABcBniRqRX+OkjwIpFLTj3OvmKtI9JIoYAAAAJJBmilJ4Q6JlMCN/8epy8Hmt5nkvHwUaGRnirECn1F6kQRmAK+NpNIq1+wKlemlXjXEhlgMZ4GsP1pm9vvJhG7Qc9ADDBVLkvTUXiuLhFe47rYgESCog4vHJ3mr48vRHWscHlSyjaWuoADQyZ2fVR3rFtDDYeIJSS40gQKiB4RIui2G6WIFhug4v4SxqnhvpgyxgQAAADBBnkdFFTzf1+NzzhJSSLZXInxj28fVa9sZ3iOeP2ue1tve6YTVpohEK9SBnnnNjSIAAAAmAZ5mdEV/hn7CXSp3/tuQ/gVx//v4sxBQVVQo1NUSFtg3nNLgRrMAAAAgAZ5oakV/0LYnEaxS2mK1fAK4ZH+xIFCrT1VZklQbIOkAAAA2QZprSahBaJlMFPG//APDMcCEr16uk80rHrFFfOOZRse1TbxZVz1dCEySWEu+EN2piNYDrHViAAAAJwGeimpFf4cEPjJJVPrWo1K1z1x0dfojX58CSEdNp6MuCQeaYl5k4AAAACxBmo1J4QpSZTBSxv/74DjgQlevXVYpZHYgY5ILdPK/4Pa1MNf4OiS8dl1R5QAAAC4BnqxqRX+NmJVnz7XQ5t2DE0nG/9r+8NOf3WQ8rd79EUL5OVWw3o6JYHnJty2AAAAAa0Gar0nhDomUwUTG//vBQYpYQEgvj00Y4sfEOWfwyh/OqKfKANvwZZWb7+Zzvu2GXqL7tu9BEtBF1qEI4yW/avkwQegNVC2ROsxmqz3xMUT7hq6sPN5B7n/7jXSLByhN4LaYlMrIVP2Gme+pAAAAFwGezmpFf46SPAjbqplSkyFjtDCI/KdAAAAAUEGa0UnhDyZTBTxv++A44D22Hq7hOPzCa+Q/0uZ8jWnVjeMaXFL05CiAEmr94qQQKh64dWL5oOBuv5DC0B686Pn9ihyHP8aNBy14cSqPoOJZAAAAKAGe8GpFf4cEPpZCBxpCqtM230lpYDo4d21o37/cFMIRIifuqQPwQmMAAAA4QZrySeEPJlMCN//74DjgQ+vwaqsgahBQipjp6SW8umNWKvqDeH4JFmwRnU6qYAsliRyGwuOB45UAAABOQZsTSeEPJlMCN//7w+5hE7FgdvuOKx17H1EIu8P6dhdyNR3Bp44IixcplMPwMmRmTbkZKRLRSf3U7jRY95x56Yr6mgyfOofav4khhB6BAAAAV0GbNknhDyZTAjf/+8UQAxkZAR4MbGRlJS3hcrKcA7AnvEOlxAMEgQ++P4BRPl+kFk9FtV5R/aKjnv2HOToKh1xJET0cChc94nwhXYYLk8OsF9FeDw4hcgAAADhBn1RFETzfghKjLXhVTVpt3odhU1/9thHbD1IrQDd5C/7EbNj0Cg37PlRRDbIVjd2YJ6YmirkTIAAAAC0Bn3VqRX+KzvglsL6ygpfSqzjv723KBBWB5nY6qSoX/1oYogCyBnZuxHtpF3EAAABDQZt3SahBaJlMCN/KCEDzIZX/gifClN7W6rOMfzBCSLAvaxS3ajNlnvJx6yl9zKnNgrw9/FvayjPOvsK8P/j0jTsPuAAAAClBm5hJ4QpSZTAjf/vgHIEFoP7SF8/0hA40oETmSgI6eYrw0ApjXdIPoAAAAGJBm7lJ4Q6JlMCN//vBGT19kCHNXXJVpZIVW2dOsxg+4nBxXYKH8lgiSB4xOvJRCAiVFRsPHQ0NMqcJypOPU3YoZF/tZgsowrc3AiXB3rAlfqaAL7OjHoF9SqtbiXfhpU4xMQAAAGJBm91J4Q8mUwI3//121hr05LoQDvslraTvUwVhvuZbxxeVmAmt8i59vy9EIHQrpGFHmbLnzjIT69y4Yxu/aorB5LH0qX9SvpgCyK/xkpHH537XWLW78sVsfAh6XcctCpaW8QAAADxBn/tFETzfgfhy3ZPoRbD4wmzsftxxUTshAaOeG2I0yWL50v57GYufV/HImcKo1kE6CdyqZME4g3D6gpEAAAAmAZ4adEV/jZiDOg/2dUiFLHxw+J/iID4obTfYetWmDphQQWMG2L8AAAAgAZ4cakV/isvpPCgm6QyWAb6SK+HCMEJXEyPb987umxgAAABPQZofSahBaJlMFPG/yghDLjqnQ+U5Iog63cQgD0Zbv41jiJ1Bb0LhUXR3RIPLhvolGP5fCuJk+gAJDaNEK0tSdXgn/kDyA9nx+m/zdv65gQAAACgBnj5qRX+NmHWKX61xnqOPLP7m1eVao9/Jf8Igt0BHxeMZyj2uGarsAAAAQ0GaIUnhClJlMFLG//vgOOBCV69dVilkdh/3zGtCpgFkyYUdvr7GH5COYLGX7YS8knmvd86QQESRdt876rG5nJMdX+4AAAAwAZ5AakV/jY0SA6wrwXjljlqnIHhE5NeeEAmpe3+gSytfIFQvKbbmEuo179PCD/OBAAABDUGaQ0nhDomUwUTG/85HJkVD4qSGOwcA3K5VwP1c6xtTR4VzmQ0nBu1gOtLRFBdIfSP3an9P1gmsInCPe79j4zuUsL7DK//x4TUhKV8qs9aaEd6tCVykWUn5Gopj46/K++Ykslr+v5UodumlbuShmmcRvZfdNrYnmbVICpyTGYIOI3oAkq6vllklL4brVgHvyURSStEFd3HHzSm7Ib3PtR4E5I89jUSkUoRqP6CDkJUiXvc83P9RW253rviuhvzqjem5Q+hPFgGTABoRk2ZLNhrce5zzCxaYC6gToCoWrl4lzpLNEOsFnUNIz4XxjMda/67iz7TaE3NHffWa7Y/p9uDhqQF1fsC6zz4vKXmUAAAAWQGeYmpFf9Ei2Luv/RNJcaZzW23Slv1beS/YB+iADhqJnO8XGwim1UkZ/5DCTycTZV+kD7aTYfAV6mID48v5zZgCA1mNEtUUWkqJoLo+Ei7fqexybID5ptuhAAAAZUGaZknhDyZTAjf/++A44Hc4/O2JB6+GsntZ/QVS7FVbo5y6gb6CO1Sl0lysiHpH6tBqj4+RV0j1AkhNZawWUW7mw1Ydkn1vMjA/KVOpoHVQwnkQTNeO2ZHz8y7tU10XuEBdXu1VAAAALUGehEURPN99TaxHbgQG60QG5CYHG0EwwVTLS+WEIzdD/zb2JngNyA/ValMZgAAAACIBnqVqRX+HBFltjCbzxqFwcAaE+pOC4nnIXAWqYK0KnVEZAAAAWkGaqkmoQWiZTAjf+8SBZr6iDxTgWdHsgrtUrFBwt9rOo3XBJP8ewDzX8GtwlWBgrC6irWNVrTLY+WKlSir9pSwV4EguL0xV3Wcj0/zw/APuEjV4PAH1Af4Z4AAAAC5BnshFESzfgh5hlu/y+9zSPV3/9spsZbsjm+DaIPLk3e+H8QEI6F/J3CJy6zchAAAALAGe53RFf42NDc3tk/BfLpHtYrknWRGdfCfWGF2hHwQDNNRoM7TAEbExLqzBAAAAIQGe6WpFf4x17dokp4/JyHzCbIpQYkeLYHCEqpiCFaz+wAAAADJBmutJqEFsmUwI3/vD+CADeEGcfiMMXyiDTP3ydcyGMsECu6OwPlhc6nEUtvJobaLBYAAAAKdBmw1J4QpSZTBRUsb/x6w8INv7ObYOML18RRf9FGE1XcGr6KXtfZJPIpBQ3Zddf/xcPS72YLTyF1aF1OgpE73KGCe4EKRaaP6+28u2fFYqgdEy089G2c5o40yj0KCcaF2ckf3qEd7wZbN7tzArbJ3ihmBK1aiIZDAZIELsAMy+cnsNyfdeXk6vZoOzvGLtwNlOVFAWfFH7l4z8o6rut3fiMJEBlecPgQAAABgBnyxqRX/QxMwwr1BpDDFywfg+dwOr6eIAAABJQZsxSeEOiZTAjf/74DjgIHS2PuRVzH01Z0BzuHiEMxkS+es2RQpgPq4cAkl78+syHU05veIX5apNghHwp2dHWeQBuTkHc9KPKQAAADRBn09FFTzfgfhxltJIut7cibOx//YvRw99xcIjnvocl0u27U4ZQtRxhVb+2ZI+NBhFvmHQAAAAJAGfbnRFf4Z+u/AqLSLb8iSXq//onmci1775jnjU+EnU1sAD/wAAAB0Bn3BqRX+Ky+lRxhN9YRwoFckJ8K7wgeXkfQ8c+QAAAE9Bm3JJqEFomUwI3/wDoisgI4JCqEuNC9+tPU5T6akQg8IizB7OrAsSlrnkH+zL9hIrZLMG+0x3pM9DfD7xIYuRzu7XLhI5qDban3u9P/d5AAAAVUGblEnhClJlMFESxv/7w+5IrxwG7Vm0o9WyHVPRTEK7sPH3Tlc6BafENz9DO3jdvcZDYYjxJ/xezZD/wVOgNlWr1ovR60fW2n2QPDIrXRl7wbJGuBkAAAAnAZ+zakV/jZWjbgdMtmBO6kL9v/+iv11IkCbeIE0lsf2Q4YI/jwZ1AAAAREGbtUnhDomUwI3/+8OipWXVEHfdmUcs3voSdMpPrWh0B7x/buw/7bfWTU2SnRyzDOlwhY0AHOKdTjMVTnbawMzXUYswAAAALkGb1knhDyZTAjf/++AcgTTZH1Q4wq20ZHYluqF2d5yizb76mgyembcA7BYepeMAAABaQZv3SeEPJlMCN//74VHYNYZ3xEb5VWTtwzMtJdYmcfzq/85J9tfpgIKEqieYPeY/MEIfk5XhJ1koykG1DVrGfyuRE/hx6dLKPysfIvtARn9KJW3Q6CD7moIQAAAAY0GaGUnhDyZTBRE8b/vgOOB3OPztiPZIA+2600lWyYtu5etLlp5Z1h7gvaM9dZAwSh88F2MaIP+Gb/x+JzyfKCfhrWJAkMuh2vjzRb7d/N7ALJMAooNT/8nzsWXFsQPjIw4BwAAAACsBnjhqRX+OeGijP1BCTddBMbj/ecV0Or4a5xYPj/vYTAUTsVBHuE5kBqvhAAAAOkGaOknhDyZTAjf/++Acgl11jEgudcSnSrhTGl+N28Vhdn1HYhWUAuZUXRLf5VeB/cE3pIjWL90Bv0kAAABeQZpcSeEPJlMFETxv+8PaaJ1sN1uxRNIGqNtU2k3406usYCk1bEtXHnFK6NYx2++uiXE5dDpIBIRXDOXY8QkTWaxAEokog4eXK891HbBLLwiFuVsjJhcQDmLu/1BPQQAAAC8BnntqRX+NjQ3Qw5SeeQtnahQCGRD+tBtlepn/thplttEshgUAcjhCMocOSDNHoAAAAE1Bmn9J4Q8mUwI3//vDpMA3iJWBxGCoVqH005LOHvO3JnF0uPAF/1BdrAmcA6ADeTiMuCgJLUTG5+vRnVljzqS04Al1sJB4spczTDLw8AAAAB9Bnp1FETzffU2jm3iDcWTT8O+Q1NZfwE4BDhgPUoUdAAAAMAGevmpFf42YlWfPtdDm3YMTsd9f+YkkWkcv8aPIfR7QyuqOLTk5HzNqkb26HGVg0AAAAGdBmqFJqEFomUwU8b/74VHYnLGfY5FAo39Th64FzFtgZ97sVu925ubsyqHBRHBm16zii8aJWj58APV75EudDjr7xLleCLLq7/RtBuOXcv77e28R66Sg0C5+ohD0irjcvKrApu8KWKOAAAAAGAGewGpFf46SPAinGc0bJk9Hz2z1Me3AgQAAAQNBmsVJ4QpSZTAjf8fMy8FXp/MBuvzFgdo0Sga+yREhRElJ1uA9sIsY4ra6c6fc00xs5DNAd8Y8yfothafu7uajsv9X2+oe8VfvZheOdZXtqkEIPdjDT49hK59jghSaTBGRE3a2q4B8vzpyVLIVGWiucmJcTr7iO+Rs1GKykq2Hqpg/7wGE9uYxrJNr28ZDTVBHR/4CyCAPQk6jWNczPh78rCn0z6G9a1RjzePKZYANdrkvuc+8zCzahQHbPyJ3wSnp4GBi0FZoLY5fHqWlr8b307+UI2CHj0yzFgHzBe7+BMVC9y83a86LgcLsFvxDSFgPEYch+RqKUMFJwfV/KNr0okJZAAAAQUGe40U0TN/X6axmThJSSLrDjCbOx+2e33lIIgsTG2JtszOBl/PYEZ41qb1VwqSpMkCnW2Br5kIl2A3Uu0IPihiHAAAAJQGfAnRFf42Ygzli/toiEQB9cbIH+nvDxV+o1J32C6oS5Fi80REAAAAoAZ8EakV/3mnpOLtzNha7JYFMdaGbfLAP/sxzAhGKBTfBtRpxTI6dgAAAAIFBmwdJqEFomUwU8b/0EQS/wgjd+gtxcvUsjBZ/KvLmmxO81sipaPVWNzgoVqy31nJce7cJ78sV7DeSUW+b67IBOE8jk8HFII54ZV3rWpp/dqQpuWbfoc2fWtU2v+/yK+ia43/m1sEOBWkif02dhVPEQ9NMhk2dwH5CcI/zvbk/eCkAAAAoAZ8makV/rN+DaXdiJ3Sfbla5z0prYIJJdn9L5fmjpAFnOrwj7T71gAAAAEdBmylJ4QpSZTBSxv/74DjhLmEQjvo7MOz/6KwI8CxsOXVk34E9cwzHoe+ypu4pyoHVVLROmPuXefirSAYX2k1/KeHKYNsg4QAAAC0Bn0hqRX+NmJVnz7XQ5t2DE0nG/A/2tAvwsk191IaIPxh05qExET5GLoDSl9kAAACTQZtLSeEOiZTBRMb//WN7GRlg7p5tdxIOcz4ktWOw25b+IeCEiFor+loK1sW6QJ67CqjJiEE34iTVIck02sn81tu+PGvjsyHqi2o+NT4inWWVtqWkpr7f/3c+vKVzlj1mjdhK8OP4T7BUoxX0iIL7Ms5LBKBxFY04P++ZvxW8cml6mweqdCiwy+yuvk365ebqBsgsAAAAGAGfampFf5XYaTriDUHbSM4/2rI012dTgAAAADxBm25J4Q8mUwI3//vgPUQ9+HnYzR0ReHsw8528FstVpSX9I30YzTO1HQX7vIeLO4umR1DvPGSkhtbazhcAAAAqQZ+MRRE8331S+se7kbW6AxPsrNhK/dtL9AZtyTBgH5jIFFvIWlbz+Y1vAAAAHQGfrWpFf4cEWNYpbR3VxAdD0L/IP/FPtuyuyqOpAAAAgUGbskmoQWiZTAjfx9CcA6khDjKpbsCmICblPVjONIzX1mXC3smr+Dq5fXiG9Xv8kFQyUglCZwie8jt9HTy9u/pPWV6Eg7yLHBH4vZV7UNGZW/41bil6995IKcn441hEAdqxiSHI4vuUl9rENOzJKWLV2U30qhtdcyYWhWRdSq79ngAAADFBn9BFESzf0IJ5kmO08++CFl4DaPf738KLIygT84teS4UEhuc0uV5bKc3qgHgZU8MlAAAALwGf73RFf7OB2xJd3tC/Pqn8/HmIdT3BpECF6yr/lbOPB+Zy1K11SUfd4VAGF2JZAAAAIQGf8WpFf4x17dokofjXiSEcvI/1LMfeM2Env+wlhWYh/wAAAEBBm/RJqEFsmUwUTG/9YpGqI6+gv8B9xK7sh6JffgqkidB+uyLMXamYT/v0gM2Bl0cBW+zOs7+meEYWSU7mTdahAAAAFwGeE2pFf4cEWNCU6Hlhkpwoe3NC1OR+AAAAfEGaFUnhClJlMCN/x85ZyjCkCh8uzK70/iv22/5Y3C+WTfpvUCrA+KzvU4pHDm4sP40BAbPg567zdr1Wu7EEBUuR9mVRdZyCPSUp6JfGVxbeoCt4oVKZm8KmBNWkYt6ZDzaPhh8MV+sn/5rn/+5/sTVujqfTRrIRFR+q1bAAAABYQZo3SeEOiZTBTRMb//vgOOB3OPztiPYLGqbk5d2BqIQU3KblqkAke1WFGNwX7uSXHgcpobdidZRPvUHXnKBlNvG1whZJdIUyHR2g5/kj6iGYMgjOKwfAwQAAACgBnlZqRX+NmJZdpmrtcZ0NTAN84Rf1N4fqx73Qed7w+gCCXwdxAjG4AAAANEGaWEnhDyZTAjf/++A44Ilb6dVP1x6RsYOzc76QVwcX2znv5faz7QEaAFkrcI7CaujLN3AAAABVQZp5SeEPJlMCN//7w+5d2nwAIi2iZMsWHBHowYX4aPQ815YaFI7qF/qVwMK0/Txt1GjsuoBrg0JYfv/uwERjL/ROpxdPOR5x3xoTlVPE5qTT470ToQAAAFlBmpxJ4Q8mUwI3//vG3LzDXKIqZ9TYmttr5CBr315a/0BUuW50ROufAOdVNE+12YSjQ+/9Zk6dVP/79hzk6CodcSRE5KWMo3xZZ6kBcB3u5L6hIx3KSQR9gQAAADlBnrpFETzfgcr6UK161qa1NGdzHnS/SFXblAfYkkTUgGALg0oSw2I9ZdDSKhlEZD82KSrfFpA+v6EAAAAnAZ7bakV/jZWjblZD8+kx6s9vusQNFqTkiIzV2G+Y2RxdmUUlRu/AAAAAVUGa3UmoQWiZTAjf+8OnUzwggj/qGfOg+4iB+MF4Y5ta94QwdwgcgpS5V9IyYNFhVGlpN2bjMQ4DNhuvd/qEmbi/FxB9D9m8UL/JpA3iUwbQMpYlEjoAAABmQZr/SeEKUmUwURLG//vB2iYdh5ImpKHSwtHernzZCXrV5CEdMKE6oB+B0fwylF2lU+/yPss1M1i+5F7zxa3LLlnpb90pQhlGEO/tirOcGbuioWk+zu7MXqDCxg3B7BB3iJfYGJshAAAAFwGfHmpFf46SPAi9dvB8N1WaHbwo5+8wAAAAaUGbA0nhDomUwI3/+8Gbu08WIicdkBLuDqjcxTRGWFEHK1YMItB/Yxti7mXMwKkgf/cswQBjw/L2b1agYYBo9t6xn6Eetw4bxPxz38vjh5Ug36pxN1P+Jdu6y5/v+SSnY/wBOAf2anagvwAAAD1BnyFFFTzfgcrcmZwgJIut+RPjHv9yMW/NWZe0k8Bi3Zgre6X89gTtwG/OmXBXHUG6b/SiwXTP+8AECry3AAAAKAGfQHRFf42Ygzli/0aMsY2QSz83/bIH9oh1NSd2YfdN+fUQOgXqQV4AAAAgAZ9CakV/isvpBhS2mIvgYfwGKq3zW41GlvjW/evQpscAAABJQZtFSahBaJlMFPG//APDMcJc0yddeUYH6nOgY2NbwtKavnsY+nBe+fxVlILL0NiRXrOhPUiA969gXTMpBWKdBefNzN3yeWc1gQAAADABn2RqRX+NmH8N/4LEDrqsxxCUlA+/C6d/hrlaE3Eni4AjiTNfEyqi1KVhrZEtHyoAAAA8QZtnSeEKUmUwUsb/++A44S5pk6Zc8mpbajeY3FQlOl5eERU4NlZQ7KCHXnQqL1aUm+ZXyoGSi1CXUAu7AAAALQGfhmpFf48rV2va6HNuwYm2A5WJIWicynQ6BVYFp9FWnKrYE2sDGYjSbKzVRAAAASVBm4lJ4Q6JlMFExv/HSJgSC1nZxJ/p92h7iaNLjO1oWGnVb7mdLzqjQegiZZ8BFVoHleRS9opigDSoevdqf/HRbxPUT937O3gQ14DzaDIrquleqdGqhoMap2XYPTsxa/wSkDdVT0ssgNi0XqyHVXOCMqt31kydwNB3y5PU4UE0NqfePmpADjVY3oiJGyZaYPsi/X4ym21aC/XEA8g6cHPxA5l4ZzRa4mnIZk+87eHMPyVrzjVuzsYf3L5Mdth1vbIMZuHrbzmxDF4Id95FDUpJLeBKHbiVKXlKbVlORXdE3XDZHj4OQ44bzKNetjD2jssvKy2arGt8LqXFFRWGchvabmDpRlO+TJnoh6Qi34c/JAelcxWZTHtnFBN5/wh1aZMwTQYfgQAAAF4Bn6hqRX/RIti7sAopo312LMs22rFhBmyBC8ZUUUusB76jdtD7pxWchSlEaZcpo+z3LSOsqF+TjWIbKyoOjcJ0BvP5Un6zNu7CYyOjMKakeOmtBeHghnZ0Lt7ZKgmBAAAAZUGbrUnhDyZTAi//87FwbhKNgmMp3eNkrIgo+BdeIZ2RH7aVgRWSRB8VGoJF1H932TYjOcWyo9X8eYosRzKDsLYZ8OElkDGMElhddcoxE4Cuo16C400E76DzMxmRGQrfcMoXwPeAAAAANUGfy0URPN+DBf3EVFfJcIvLYYm40FJuHP41gUVpmbF9867RDcpbG+hUwmLp76j9ncDp8t5SAAAAIwGf6nRFf4Z+tIoBTm+ZzUi4P/L4NbInHV4wNqEejzvRIyMbAAAAHAGf7GpFf4rL6VHGE3tN/0VO4T4Y92cJmTGqjpgAAABFQZvvSahBaJlMFPF/+G0p3IyKrpYoACoREvabxZ2Sa1j/PX4Ig63AFh5hVcDy6lyqDJVazyv9k0A3UEbpQ4kElhDfQ/kDAAAAJQGeDmpFf42Yfwkkpl9ce7oX5djjxlYebpvmnM3if/h60uGBIJAAAAA0QZoRSeEKUmUwUsX/+99LcS41uOHpiqjPsmwGusCKQ6BF1AYtz/WtQ7b6AecL2sZs2DeFgQAAACoBnjBqRX+PK1dr2uhzbsGJtgOViSFonOVO2XvUQkyMOnNQmIoMiU2QLcEAAABnQZozSeEOiZTBRMX/++B7/JR3TzHmkys6OPVr3j60/Ducpwbn9LZ8EMRIp1AAs7+bDOf0XhisAXR+MHkgbVlZgREbbfkD3LjS6EKrXc0GPULfC2RHIyE230IZnBmWLkBNNdGdzIhkLQAAABcBnlJqRX+OkjwI26tNPqSi/HkmF/g/gQAAAHJBmlZJ4Q8mUwIn//Dz42spRFihMPfTkj2TNj0lStYXk3a+C5Wtkfypzo0uM/feYzEOqkQzy1NzoYXThBgSRvjrEcejioh4gzCm8rCbzMI2wWFHN/yLtF3p6q9/Ec94jMIwPBQYnzCY283aqNgEwQc+asIAAAAwQZ50RRE836foBRLD06goiSPhmYv9qz7O6uD4czg9Q3XDksNdfK5w0IuoG1k3GS0sAAAAHAGelWpFf4cEWNYpbQS22jK73/5IyrwB5npkP/EAAABQQZqYSahBaJlMFPE/+81+QTCAevbatPJvZfPu+5FInLiQmsx/vsDDnlk+n51NMYohIZisWJ/S3zrGQtmid1/VGv24Rk3j0Dp9U4U6eejQnKAAAAAuAZ63akV/jZifIUt2huke1h+nmbVlvSl8BsecDgLV4KKiXsuB44t0CtoWXXNbfgAAACtBmrlJ4QpSZTAivyCC9u0SU8fn1JmlqS4H0O/V2QQ8C23N8Hq2eJg+n3lBAAAEpWWIggAb/6zQYZLfQPgNc0iYMN3sPl0EDBzzvrz1O9Sgfa49FGnVhJ8qLkhVezFoQxJKbc0nUGNZR098qd7OS+y89EcOWVfDbC64tS8Y9KVBO5bX3i8BNZhBqqdfBpIeL6SUDjC2QpVPdT+4GdjwHm+HVEddJzn4mxN+ZHi1zAImnjZYMgtVcr/8KsM8y6i5AmAGJNcUB7ayE0bca3a3R6hdrr0XJr+AAAXHQ/kXcVlmP6pUrkjcKn/5MsrYsfvEq74eZ4HZz2CqbE+TyZ4bBsWkF/G8EsnWTdMeCaT0oSOggGj6F6GH2koqwzCw8EE/jxjH/fgpvYdskan80JjRYRAJC3SwAtTkEnddIDJhOwFLkmhXRf1WZtt88VL7WgZg+Do7We/RCfv4cfPN57fJ/eEcGoPIdjySktCXmVnqHKZZhx61mqXrYt2ccmalI5pFwxn7yIk4NanyDkRuoRQdzVy7V/FgNDX+/OG313omVQVM1D6TfnNy1isPsCV96//sVHChbSvRvo9Op4uXUT2Vnp2NbBUnPOlIQwYUmnls4/qAf1DSyr347ve+HuWjGSfL2uAJLxNkKGL7jaqU87H350pTIU4+n/FRrP8UOsWWuZ9oqEO+lS3qvWMNvTA00dfS9Og3CEQ1lhvXgqkfbgmWS7Q7BrqjPYxv+usZZhR4mTsySOLjIGUswJByOkD9sXOn034oNMHPSMTs90mvpwuhhZCES9QIA0MY4mVJAH4U1175JsR0JGrNT2cG+kevDwhB7L2Ps9sq35Ja1G1YIr1fCqBH1NhFZbh3r/JJ6IfmaKkwRe/x35k1btKb+ZRfDYCZwD/U+DVYV4ve4/ux5EMulf81TrZ2iS/AeNl9oIVL7FLsYt0SbM4Gdwujxy/cVa/9DfeeCmTLktfb85Jq/PPzoWfkT+3MSEL52ahvl5dyZ3P9Gjqc/3V3E3UcMDLHXaFjLji1XATHryNndFbY++diGv2rJsu8jbVZYN/lqAhbLHtMl/sxkHV0L1mnPx14QjgEVCWYb8PE3T+6bY0rrXT6WTEzJq/uaxvgC3UliUMREWCB/LHkW4/FwMNKpTt+KgL7DKMzlTFtOl01XRpLhfNYIv4yqp/8/0Ljp/cAHvBSTd0o2zGUJLgBlXpkC7hkWNvVKuNH3gZDe7aBWygB4hHI5MtXigV2BDYnOAj1gC8+xGfGlAw3JTy1PkTj3bbr+p6m1jzipdInXtuw8i53LkL/+t/cDhYKza5lgkrWZvnpedOPwE+SDsrxeadOx3w6iN4kI3PIREK2JArXFqTIikdsxGeR30xC/R4MnKZUvO9n1HHNmqCdKQY/t8N97D5MniGESaR7uA4AEebTLOvxLuJghYdjoSxajwIKjHQQZ/HnrodJI5KNoj6gM3UJuupZ/7uHrjf0Dx0Tsc7GUns7O78xmwOt1zDx4hBNRbDjREElJr9eJlRZ3FMTZgL28DJXX2Ukq+M+ZYpT8CiXX8w7XSQxcMg4cJ6c0AGluEsCz3tAkAQZLU/VxMQpKTJR8+ULaCRGp3sYhJhT5mx6NspTrLO4jl7T5GLVFULkYWq1jSnmblAgEE58LcEAAAA9QZohbE//55dAJChmGePT0n0Li66IQjlgCcvSeXFOTWduetNIn5ncKj2uvzaGAWdM4YtqdWioqkA0TyHz0QAAADBBmkI8IZMphP/nR8JCz0aek9id9vrF4lLN2RgQjyZr432P7Kp14LSDqmMhMkWW0IAAAABgQZpjSeEPJlMCf+dJU6LZgkV02QLvEfeXB5dv4nCjjykgCXjsebUglI9rrNk3mZJmTaQFGdOlZR2+jF2hfcbd4HXQBsGGhxrxbbB//345uw0C7pMaLkJqkj5rLI3WN+aIAAAAUEGahUnhDyZTBRE8/+ZYKy4kVT7wb4JNhryfYT+xBrDfh8zeOyqxLHImSnk1hEFx+vGFOK6nSbea/LiCsJT6jX4eofbWT/4UXXlaQ7ns7HmfAAAAKgGepGpFf42Yll7P81LVQddgA183/eHh4rICR6Rm/TAP+50pi9D2H4nE8AAAADBBmqZJ4Q8mUwJfi8K/Q8CzpoXb9z2cRUhHfcB/4ijSNB3DB8kDQxNhedtHQQegeFAAAABHQZrISeEPJlMFETy/jBNpyKcpBz2UWq9NhKZ8OdqT4vlDJSja2L4KsxrGpSrA2yYsC6BONsZSjEWJ7nESRHswZooLCvZWL7cAAAAtAZ7nakV/jY0N0M9lu0TcpcLEt4G2uP1I1ad/p8M0BAfmcyRE+XWn8kbnbCP/AAAAQ0Ga7EnhDyZTA/8JnI6fa/XOcPoz6Eunvm3DSqWZut6xANl//x1bLCAKYLkthnFxJePmxk5tfDOwQ/OimiJzESuj2PEAAAAyQZ8KRRE834HLEsIWvPiTtw1hFA7vKrrccI3560FD1CXoViI/ZgtgMQkFL76i8/+MS6kAAAAdAZ8pdEV/iRfHjugOkdFKyjh/gRLx8g9LelE/H6AAAAAYAZ8rakV/hwRYbMohViYe1/b50qjneEmNAAAAWEGbLUmoQWiZTAr/DOcGwr30MLcyf4RisZLdUWnlUYvXJVKV7AncbyAoLODntBs0+K4WaaoVpZfqfI2XI0wXb+fJSpC5XehdM+mTV3yL+1xrM/enVV5ivWAAAAAzQZtQSeEKUmUwIr8dln8w8D8m9uEVJ8VaXrlpwtQMUcjOxV139A/A0FOJb3veKqfxYs2BAAAALkGfbkU0TN+BznEOeLmqf8Sjo1rSbw27tYVeLjYdb2CxT8VnvpcHn1axdLPZBJ0AAAAtAZ+PakV/hwRZbYwm6h1dIRfbz8pmPJ5q1aY4QcQjnAtS8tukkKVgqkBfhK/oAAAkC21vb3YAAABsbXZoZAAAAAAAAAAAAAAAAAAAA+gAASucAAEAAAEAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAIAACM1dHJhawAAAFx0a2hkAAAAAwAAAAAAAAAAAAAAAQAAAAAAASucAAAAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAQAAAAABQAAAAhAAAAAAAJGVkdHMAAAAcZWxzdAAAAAAAAAABAAErnAAACAAAAQAAAAAirW1kaWEAAAAgbWRoZAAAAAAAAAAAAAAAAAAAKAAAC/wAVcQAAAAAAC1oZGxyAAAAAAAAAAB2aWRlAAAAAAAAAAAAAAAAVmlkZW9IYW5kbGVyAAAAIlhtaW5mAAAAFHZtaGQAAAABAAAAAAAAAAAAAAAkZGluZgAAABxkcmVmAAAAAAAAAAEAAAAMdXJsIAAAAAEAACIYc3RibAAAAKhzdHNkAAAAAAAAAAEAAACYYXZjMQAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAABQAIQASAAAAEgAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABj//wAAADJhdmNDAWQACv/hABlnZAAKrNlFE/nwEQAAAwABAAADABQPEiWWAQAGaOvjyyLAAAAAEHBhc3AAAAABAAAAAQAAABhzdHRzAAAAAAAAAAEAAAL/AAAEAAAAACBzdHNzAAAAAAAAAAQAAAABAAAA+wAAAfUAAALvAAAU8GN0dHMAAAAAAAACnAAAAAEAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAwAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAABAAACAAAAAABAAAMAAAAAAEAAAQAAAAABwAACAAAAAABAAAQAAAAAAIAAAQAAAAABwAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAADAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAwAACAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAACAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAwAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAACAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAMAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAADAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAACAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAMAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAACAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAAEAAAAAACAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAABAAAAAAAgAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABQAAAAAAQAACAAAAAABAAAAAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAIAAAgAAAAAAQAAEAAAAAACAAAEAAAAAAMAAAgAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAMAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAAAgAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAQAAAAAAIAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAgAACAAAAAABAAAQAAAAAAIAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAUAAAAAAEAAAgAAAAAAQAAAAAAAAABAAAEAAAAAAEAAAwAAAAAAQAABAAAAAABAAAMAAAAAAEAAAQAAAAAAQAADAAAAAABAAAEAAAAAAEAABAAAAAAAgAABAAAAAABAAAMAAAAAAEAAAQAAAAABQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAACAAAAAABAAAMAAAAAAEAAAQAAAAAAQAAFAAAAAABAAAIAAAAAAEAAAAAAAAAAQAABAAAAAABAAAIAAAAAAEAABAAAAAAAgAABAAAAAAcc3RzYwAAAAAAAAABAAAAAQAAAv8AAAABAAAMEHN0c3oAAAAAAAAAAAAAAv8AAAZHAAAANQAAACEAAAAkAAABswAAADgAAAAmAAAAbwAAADoAAAAzAAAAKAAAADcAAAC4AAAAIwAAAFoAAAAxAAAAJQAAABwAAABfAAAAWAAAACoAAABXAAAANAAAAF0AAABhAAAALAAAAD4AAABHAAAALwAAAEgAAAAeAAAANgAAAFgAAAAiAAAATwAAADsAAABXAAAARgAAAFUAAABFAAAAQQAAADUAAAA4AAAAXwAAAFQAAABEAAAARwAAAGEAAABVAAAATgAAACIAAAA8AAAAQwAAAGMAAABRAAAASwAAAEcAAABuAAAAUAAAAFsAAAAmAAAASAAAAHAAAAAbAAAAVwAAADQAAAAkAAABHwAAADoAAAB7AAAAJAAAAEAAAACzAAAAJAAAAHYAAAA7AAAAKwAAACEAAABXAAAALQAAAEQAAAA0AAAAcgAAABwAAABnAAAAMQAAADkAAABcAAAAbgAAAEAAAAAoAAAATAAAAEgAAAAdAAAAZgAAAEQAAAAvAAAAJgAAAE0AAAA4AAAANQAAADYAAAFBAAAAZAAAAB0AAACUAAAAQQAAACEAAAAvAAAAQQAAADIAAAAfAAAAGgAAAFgAAABLAAAANwAAACkAAAAfAAAAVAAAAFUAAAAnAAAATAAAADcAAABdAAAATgAAAC4AAAA1AAAATQAAADEAAABDAAAAIgAAADAAAABhAAAAIQAAAQkAAABPAAAALQAAAC0AAACOAAAALAAAAEUAAAA2AAAAaAAAABoAAACaAAAAMgAAACEAAABGAAAANwAAAC8AAAAmAAAALgAAADIAAACCAAAAcwAAAD4AAABaAAAAXwAAAGYAAAA9AAAAKQAAAEUAAABYAAAAHAAAAFMAAAAxAAAAJAAAAG8AAABPAAAANwAAAC4AAAEZAAAAIQAAALsAAABvAAAANAAAADAAAACtAAAAQQAAADMAAAAuAAAAXAAAADEAAABfAAAATgAAADMAAAAiAAAAUAAAADMAAABXAAAALAAAAD0AAABTAAAAawAAACEAAAAvAAAANQAAAEwAAAAzAAAATAAAACIAAAAyAAAAkwAAAB8AAABpAAAAOAAAACgAAAD/AAAAOAAAAHcAAAAqAAAASQAAAHQAAAAbAAAAnAAAADoAAAArAAAAIgAAAGUAAAAxAAAAQwAAAC4AAABmAAAAHQAAAGgAAAAuAAAANAAAAFYAAABwAAAAPQAAADAAAABHAAAAWQAAABwAAABkAAAARgAAAC8AAAAkAAAATQAAADUAAAA9AAAAMwAAAOkAAABcAAAAXQAAADwAAAApAAAAIQAAAEAAAAApAAAFBgAAAEAAAAC2AAAAHAAAAEwAAAA3AAAAKwAAACEAAABTAAAAWgAAACgAAABNAAAANwAAAFkAAABTAAAALgAAADcAAABKAAAAMwAAAEIAAAAhAAAANAAAAIYAAAAdAAABMQAAAEsAAAAuAAAAMQAAAIMAAAAxAAAARwAAADIAAABkAAAAGgAAADgAAAAvAAAAIQAAAIgAAAA3AAAALwAAACYAAAA7AAAAHgAAAE4AAABYAAAAKgAAADYAAABSAAAAXwAAAD0AAAArAAAATQAAAGMAAAAfAAAAagAAAEAAAAAlAAAAcAAAAFAAAAA1AAAAKwAAAD4AAAFQAAAAXAAAAGMAAAA2AAAALAAAACEAAABCAAAAKwAAADcAAAAwAAAAVQAAABkAAABsAAAAMgAAACEAAABbAAAANQAAAE8AAAAtAAAAPAAAADMAAABIAAAASwAAAC8AAAA/AAAATQAAAC8AAABQAAAANAAAACcAAAAbAAAAdQAAAEEAAAAoAAAANAAAAIkAAAA6AAAAeQAAACoAAAA9AAAAGgAAAF8AAABXAAAANQAAACYAAAAgAAAASwAAACoAAAA2AAAAMQAAAF4AAAAcAAAAWgAAACoAAAA2AAAAUAAAAF0AAAA/AAAAMAAAAEkAAABpAAAAHgAAAHUAAABDAAAALgAAACMAAABKAAAAMQAAADwAAAAyAAABcwAAAFkAAACnAAAANQAAACkAAABuAAAAMgAAACoAAAAoAAAAVgAAAG4AAAAcAAAAWAAAADwAAAA7AAAAKAAAAFgAAABWAAAAKAAAAE8AAAA2AAAAiQAAAEwAAAAtAAAANQAAAGcAAAAzAAAAUAAAADkAAAAgAAAAHQAAAGUAAAEeAAAAUAAAAC4AAAAxAAAAogAAACkAAACOAAAAOwAAAFoAAAAbAAAAVQAAADkAAAAsAAAAIwAAAEsAAAAqAAAANAAAADAAAABTAAAAHwAAAFcAAAArAAAAMwAAAFMAAABdAAAAPQAAADAAAABcAAAAgQAAACIAAACCAAAAUQAAAC8AAAAjAAAATQAAAC4AAAA5AAAANAAAAQUAAAC3AAAAogAAADUAAAArAAAAIQAAAEAAAAAnAAAANgAAAC8AAABTAAAAGQAAADwAAAAxAAAAIQAAAE8AAAAzAAAAWwAAACsAAAA5AAAAPQAAAF4AAABRAAAAMAAAADcAAABQAAAALgAAAEoAAAAiAAAAMwAAAI4AAAAdAAAATwAAADgAAAA3AAAA7QAAADQAAABwAAAAKgAAADgAAACTAAAAHQAAAFQAAAA1AAAAKQAAACIAAABSAAAAKAAABUYAAAA+AAAAlAAAAB0AAAB3AAAALgAAADMAAABcAAAAZQAAAD8AAAAsAAAATAAAAG4AAAAbAAAAawAAAEAAAAAtAAAAJwAAAE0AAAA7AAAAOQAAADEAAAEWAAAAVwAAAFQAAAA3AAAAKgAAACAAAACPAAAAKwAAAE0AAAAmAAAALwAAABgAAAA+AAAAMQAAACAAAABSAAAAMQAAAFwAAAArAAAAOQAAADgAAACJAAAAJQAAADMAAAA2AAAATgAAADMAAABCAAAAIwAAADEAAAB/AAAAHgAAATUAAABJAAAAKwAAAC0AAACEAAAAKwAAAEYAAAA0AAAAVwAAABoAAACaAAAAMQAAACIAAABJAAAAMgAAADAAAAAmAAAAMQAAAHEAAAAcAAAAVAAAACwAAABEAAAAXgAAAGMAAAA5AAAALAAAAEcAAABjAAAAHgAAAHAAAABEAAAALAAAACMAAABWAAAALgAAAMQAAAAwAAAANQAAAGQAAABRAAAANAAAACoAAACZAAAANQAAADYAAAApAAAAOwAAABsAAABuAAAAQQAAADQAAAAhAAAATwAAADEAAABWAAAALAAAADgAAABMAAAAgAAAACEAAAAxAAAAOwAAAEwAAAAxAAAATgAAAD0AAAAoAAAAHAAAAFgAAABOAAAANAAAACUAAAEDAAAAMQAAAHUAAAApAAAASAAAAEkAAAAbAAAAlgAAADQAAAAqAAAAJAAAADoAAAArAAAAMAAAADIAAABvAAAAGwAAAFQAAAAsAAAAPAAAAFIAAABbAAAAPAAAADEAAABHAAAALQAAAGYAAABmAAAAQAAAACoAAAAkAAAAUwAAACwAAABHAAAANAAAAREAAABdAAAAaQAAADEAAAAmAAAAXgAAADIAAAAwAAAAJQAAADYAAACrAAAAHAAAAE0AAAA4AAAAKAAAACEAAABTAAAAWQAAACsAAABIAAAAMgAAAF4AAABnAAAALwAAAD4AAABiAAAAMwAAAFEAAAAjAAAANAAAAGsAAAAcAAABBwAAAEUAAAApAAAALAAAAIUAAAAsAAAASwAAADEAAACXAAAAHAAAAEAAAAAuAAAAIQAAAIUAAAA1AAAAMwAAACUAAABEAAAAGwAAAIAAAABcAAAALAAAADgAAABZAAAAXQAAAD0AAAArAAAAWQAAAGoAAAAbAAAAbQAAAEEAAAAsAAAAJAAAAE0AAAA0AAAAQAAAADEAAAEpAAAAYgAAAGkAAAA5AAAAJwAAACAAAABJAAAAKQAAADgAAAAuAAAAawAAABsAAAB2AAAANAAAACAAAABUAAAAMgAAAC8AAASpAAAAQQAAADQAAABkAAAAVAAAAC4AAAA0AAAASwAAADEAAABHAAAANgAAACEAAAAcAAAAXAAAADcAAAAyAAAAMQAAABRzdGNvAAAAAAAAAAEAAAAwAAAAYnVkdGEAAABabWV0YQAAAAAAAAAhaGRscgAAAAAAAAAAbWRpcmFwcGwAAAAAAAAAAAAAAAAtaWxzdAAAACWpdG9vAAAAHWRhdGEAAAABAAAAAExhdmY1Ny44My4xMDA=",
+       "ok": true,
+       "headers": [
+        [
+         "content-type",
+         "video/mp4"
+        ]
+       ],
+       "status": 200.0,
+       "status_text": ""
+      }
+     },
+     "base_uri": "https://localhost:8080/",
+     "height": 501.0
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "\n",
+       "    <video width=\"640\" height=\"480\" controls>\n",
+       "      <source src=\"/nbextensions/vid.mp4\" type=\"video/mp4\">\n",
+       "    </video>\n",
+       "  "
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "execution_count": 31,
+     "metadata": {
+      "tags": []
+     },
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "play_video('mf_pong/0.avi')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "NQmZEVKGF4Hh",
+    "colab_type": "text"
+   },
+   "source": [
+    "# Model-based training\n",
+    "\n",
+    "The `rl` package offers many more features, including model-based training. For instructions on how to use them, go to our [README](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/rl)."
+   ]
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "name": "hello_t2t-rl.ipynb",
+   "version": "0.3.2",
+   "provenance": [
+    {
+     "file_id": "1nQvfx1EzY3ElJUy-FVF1G16okSbkeUa2",
+     "timestamp": 1.553274233669E12
+    }
+   ],
+   "collapsed_sections": []
+  },
+  "kernelspec": {
+   "name": "python3",
+   "display_name": "Python 3"
+  },
+  "accelerator": "GPU"
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

From 623f5cbbdb66bebe1c49659faabf6bbe1588acae Mon Sep 17 00:00:00 2001
From: Piotr Milos <piotr.milos@codilime.com>
Date: Mon, 25 Mar 2019 12:38:59 -0700
Subject: [PATCH 1829/2720] Merge of PR #1518

PiperOrigin-RevId: 240196506
---
 tensor2tensor/notebooks/hello_t2t-rl.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/notebooks/hello_t2t-rl.ipynb b/tensor2tensor/notebooks/hello_t2t-rl.ipynb
index abd502819..ed5e71331 100644
--- a/tensor2tensor/notebooks/hello_t2t-rl.ipynb
+++ b/tensor2tensor/notebooks/hello_t2t-rl.ipynb
@@ -1862,7 +1862,7 @@
    "source": [
     "# Model-based training\n",
     "\n",
-    "The `rl` package offers many more features, including model-based training. For instructions on how to use them, go to our [README](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/rl)."
+    "The `rl` package offers many more features, including model-based training. For instructions on how to use them, go to our [README](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/rl/README.md)."
    ]
   }
  ],

From 3744017ba8ce609186d2a91dbe3d6b8ddea9203b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 25 Mar 2019 13:56:11 -0700
Subject: [PATCH 1830/2720] Add LanguagemodelWikitext103L16k problem

PiperOrigin-RevId: 240211764
---
 tensor2tensor/data_generators/wikitext103.py | 10 ++++++++++
 tensor2tensor/models/transformer.py          | 21 ++++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/tensor2tensor/data_generators/wikitext103.py b/tensor2tensor/data_generators/wikitext103.py
index 6c6f9eb1f..1e53eefbc 100644
--- a/tensor2tensor/data_generators/wikitext103.py
+++ b/tensor2tensor/data_generators/wikitext103.py
@@ -191,3 +191,13 @@ def max_length(self, model_hparams):
   def sequence_length(self):
     """Length of each example (in tokens)."""
     return 4096
+
+
+@registry.register_problem
+class LanguagemodelWikitext103L16k(LanguagemodelWikitext103L4k):
+  """Wikitext-103, token-level, with examples up to 16,384 tokens long."""
+
+  @property
+  def sequence_length(self):
+    """Length of each example (in tokens)."""
+    return 16384
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index e6d22b26e..0dfdca3ad 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -2642,3 +2642,24 @@ def transformer_wikitext103_l4k_memory_v0():
   hparams.max_relative_position = 2 * hparams.split_targets_chunk_length
 
   return hparams
+
+
+@registry.register_hparams
+def transformer_wikitext103_l16k_memory_v0():
+  """HParams for training languagemodel_wikitext103_l16k with memory."""
+  hparams = transformer_wikitext103_l4k_memory_v0()
+
+  hparams.max_length = 16384
+  hparams.split_targets_chunk_length = 64
+  hparams.split_targets_max_chunks = int(
+      hparams.max_length / hparams.split_targets_chunk_length)
+
+  # The hparams specify batch size *before* chunking, but we want to have a
+  # consistent 4K batch size *after* chunking to fully utilize the hardware.
+  target_tokens_per_batch = 4096
+  hparams.batch_size = int(target_tokens_per_batch * (
+      hparams.max_length / hparams.split_targets_chunk_length))
+
+  hparams.max_relative_position = 2 * hparams.split_targets_chunk_length
+
+  return hparams

From 150aad3be93c68f9424c0769e9e710c754ebaaea Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Tue, 26 Mar 2019 00:28:50 +0100
Subject: [PATCH 1831/2720] Model-Based RL: batched environments for DQN
 (#1500)

* MBRL: batched dopamine, runner and agent

* MBRL: Fix _observation usage in BatchDQNAgent; clean up tests; some minor changes.

* MBRL: Perform multiple _train_steps per env_step in batched dqn, to keep the same _train_steps:env_steps ratio as in non-batched version.

* MBRL: Use batched dopamine in MBRL pipeline.

* Minor fixes, including dopamine 1.0.4 compatibility.

* Batched Dopamine: Fix current_rollouts reset in BatchedAgent

* Padded BatchEnv, prints.

* Assert batch_size=1 for dopamine evaluation.

* Enable model-free with dqn.

* Move to dopamine 2.0.1.

* Remove unused functions, add documentation.

* Remove deprecated TODOs.

* Fix SimulatedBatchEnv closing.

* Fix closing environment in dopamine.

* Improve batch size inference.

* Add test for model-free and model-based dqn, reduce model-free ppo test time.

* Parameter for model-based dqn number of evaluation episodes.

* Unify batch_env attribute name for dopamine environment and wrappers.

* Remove PaddedBatchEnv from default model-based dqn pipeline.

* Linting.

* Update tests for batch dqn runner and agent.
---
 tensor2tensor/layers/common_video.py          |   2 +
 tensor2tensor/models/research/rl.py           |  52 ++-
 tensor2tensor/rl/batch_dqn_agent_test.py      | 157 +++++++
 tensor2tensor/rl/batch_runner_test.py         | 297 +++++++++++++
 tensor2tensor/rl/dopamine_connector.py        | 399 ++++++++++++++----
 tensor2tensor/rl/envs/simulated_batch_env.py  |   3 +-
 .../rl/envs/simulated_batch_gym_env.py        |   4 +-
 tensor2tensor/rl/policy_learner.py            |   2 -
 tensor2tensor/rl/trainer_model_based.py       |   2 +
 .../rl/trainer_model_based_dqn_test.py        |  41 --
 .../rl/trainer_model_based_params.py          |   8 +-
 tensor2tensor/rl/trainer_model_based_test.py  |  10 +-
 tensor2tensor/rl/trainer_model_free.py        |  78 ++--
 tensor2tensor/rl/trainer_model_free_test.py   |  14 +-
 14 files changed, 894 insertions(+), 175 deletions(-)
 create mode 100644 tensor2tensor/rl/batch_dqn_agent_test.py
 create mode 100644 tensor2tensor/rl/batch_runner_test.py
 delete mode 100644 tensor2tensor/rl/trainer_model_based_dqn_test.py

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index ea94316cf..10bc5d7ab 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -790,6 +790,8 @@ def finish(self):
     (out, err) = [
         b"".join(chunks) for chunks in (self._out_chunks, self._err_chunks)
     ]
+    self.proc.stdout.close()
+    self.proc.stderr.close()
     if self.proc.returncode:
       err = "\n".join([" ".join(self.cmd), err.decode("utf8")])
       raise IOError(err)
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 90c74aeef..249e87a2c 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -354,12 +354,18 @@ def dqn_atari_base():
       optimizer_epsilon=0.00001,
       optimizer_centered=True,
 
+      # TODO: change names maybe replay_buffer -> agent? Also batch_size is now
+      # buffer_batch_size in _DQNAgent.
       replay_buffer_replay_capacity=1000000,
-      replay_buffer_batch_size=32,
+      replay_buffer_buffer_batch_size=32,
 
       time_limit=27000,
       save_every_steps=50000,
       num_frames=int(20 * 1e6),
+
+      # TODO(konradczechowski) this is not used in trainer_model_free, clean
+      # this up after evaluation refactor
+      eval_episodes_num=3,
   )
 
 
@@ -370,6 +376,16 @@ def dqn_original_params():
   hparams.set_hparam("num_frames", int(1e6))
   return hparams
 
+def rlmf_tiny_overrides():
+  """Parameters to override for tiny setting excluding agent-related hparams."""
+  return dict(
+      max_num_noops=1,
+      eval_max_num_noops=1,
+      rl_env_max_episode_steps=7,
+      eval_rl_env_max_episode_steps=7,
+      eval_sampling_temps=[0.0, 1.0],
+  )
+
 
 @registry.register_hparams
 def rlmf_original():
@@ -382,6 +398,7 @@ def rlmf_original():
       eval_batch_size=2,
       frame_stack_size=4,
       eval_sampling_temps=[0.0, 0.2, 0.5, 0.8, 1.0, 2.0],
+      max_num_noops=8,
       eval_max_num_noops=8,
       eval_rl_env_max_episode_steps=1000,
       resize_height_factor=2,
@@ -426,6 +443,31 @@ def rlmf_base():
   return hparams
 
 
+@registry.register_hparams
+def rlmf_tiny():
+  """Tiny set of hparams for model-free PPO."""
+  hparams = rlmf_original()
+  hparams = hparams.override_from_dict(rlmf_tiny_overrides())
+  hparams.batch_size = 2
+  hparams.add_hparam("ppo_epochs_num", 3)
+  hparams.add_hparam("ppo_epoch_length", 2)
+  return hparams
+
+
+@registry.register_hparams
+def rlmf_dqn_tiny():
+  hparams = rlmf_original()
+  hparams = hparams.override_from_dict(rlmf_tiny_overrides())
+  hparams.batch_size = 1
+  hparams.base_algo = "dqn"
+  hparams.base_algo_params = "dqn_original_params"
+  hparams.add_hparam("dqn_num_frames", 128)
+  hparams.add_hparam("dqn_save_every_steps", 128)
+  hparams.add_hparam("dqn_replay_buffer_replay_capacity", 100)
+  hparams.add_hparam("dqn_agent_min_replay_history", 10)
+  return hparams
+
+
 @registry.register_hparams
 def rlmf_eval():
   """Eval set of hparams for model-free PPO."""
@@ -442,14 +484,6 @@ def rlmf_eval():
   return hparams
 
 
-@registry.register_hparams
-def rlmf_tiny():
-  hparams = rlmf_base()
-  hparams.ppo_epochs_num = 100
-  hparams.ppo_eval_every_epochs = 10
-  return hparams
-
-
 class PolicyBase(t2t_model.T2TModel):
 
   def loss(self, *args, **kwargs):
diff --git a/tensor2tensor/rl/batch_dqn_agent_test.py b/tensor2tensor/rl/batch_dqn_agent_test.py
new file mode 100644
index 000000000..6e70c8e54
--- /dev/null
+++ b/tensor2tensor/rl/batch_dqn_agent_test.py
@@ -0,0 +1,157 @@
+# Copyright 2018 The Dopamine Authors.
+#
+# 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 BatchDQNAgent."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import shutil
+
+from absl import flags
+from dopamine.agents.dqn import dqn_agent
+import numpy as np
+
+from tensor2tensor.rl import dopamine_connector
+
+import tensorflow as tf
+
+
+slim = tf.contrib.slim
+
+FLAGS = flags.FLAGS
+
+
+class BatchDQNAgentTest(tf.test.TestCase):
+  # TODO: maybe add testStepTrain (and possibly some other tests) from dopamine
+  # dqn_agent_test.py
+
+  def setUp(self):
+    self._test_subdir = os.path.join('/tmp/dopamine_tests', 'ckpts')
+    shutil.rmtree(self._test_subdir, ignore_errors=True)
+    os.makedirs(self._test_subdir)
+    self.num_actions = 4
+    self.min_replay_history = 6
+    self.update_period = 2
+    self.target_update_period = 4
+    self.epsilon_decay_period = 90
+    self.epsilon_train = 0.05
+    self.observation_shape = dqn_agent.NATURE_DQN_OBSERVATION_SHAPE
+    self.stack_size = dqn_agent.NATURE_DQN_STACK_SIZE
+    self.env_batch_size = 4
+
+    self.zero_state = np.zeros(
+        [self.env_batch_size, self.observation_shape[0],
+         self.observation_shape[1], self.stack_size])
+
+
+  def _create_test_agent(self, sess):
+    stack_size = self.stack_size
+
+    class MockDQNAgent(dopamine_connector.BatchDQNAgent):
+
+      def _network_template(self, state):
+        # This dummy network allows us to deterministically anticipate that
+        # action 0 will be selected by an argmax.
+        inputs = tf.constant(
+            np.zeros((state.shape[0], stack_size)), dtype=tf.float32)
+        # This weights_initializer gives action 0 a higher weight, ensuring
+        # that it gets picked by the argmax.
+        weights_initializer = np.tile(
+            np.arange(self.num_actions, 0, -1), (stack_size, 1))
+        q = slim.fully_connected(
+            inputs,
+            self.num_actions,
+            weights_initializer=tf.constant_initializer(weights_initializer),
+            biases_initializer=tf.ones_initializer(),
+            activation_fn=None)
+        return self._get_network_type()(q)
+
+    agent = MockDQNAgent(
+        replay_capacity=100,
+        buffer_batch_size=8,
+        generates_trainable_dones=True,
+        sess=sess,
+        env_batch_size=self.env_batch_size,
+        num_actions=self.num_actions,
+        min_replay_history=self.min_replay_history,
+        epsilon_fn=lambda w, x, y, z: 0.0,  # No exploration.
+        update_period=self.update_period,
+        target_update_period=self.target_update_period,
+        epsilon_eval=0.0)  # No exploration during evaluation.
+    # This ensures non-random action choices (since epsilon_eval = 0.0) and
+    # skips the train_step.
+    agent.eval_mode = True
+    sess.run(tf.global_variables_initializer())
+    return agent
+
+  def testCreateAgentWithDefaults(self):
+    # Verifies that we can create and train an agent with the default values.
+    with tf.Session() as sess:
+      agent = self._create_test_agent(sess)
+      sess.run(tf.global_variables_initializer())
+      observation = np.ones([84, 84, 1])
+      agent.begin_episode([observation])
+      agent.step(reward=[1], observation=[observation])
+      agent.end_episode(reward=[1])
+
+  def testBeginEpisode(self):
+    """Test the functionality of agent.begin_episode.
+
+    Specifically, the action returned and its effect on state.
+    """
+    with tf.Session() as sess:
+      agent = self._create_test_agent(sess)
+      # We fill up the state with 9s. On calling agent.begin_episode the state
+      # should be reset to all 0s.
+      agent.state_batch.fill(9)
+      first_observation = np.ones(
+          [self.env_batch_size, self.observation_shape[0],
+           self.observation_shape[1], 1])
+      self.assertTrue((agent.begin_episode(first_observation) == 0).all())
+      # When the all-1s observation is received, it will be placed at the end of
+      # the state.
+      expected_state = self.zero_state
+      expected_state[:, :, :, -1] = np.ones(
+          [self.env_batch_size, self.observation_shape[0],
+           self.observation_shape[1]])
+      self.assertAllEqual(agent.state_batch, expected_state)
+      self.assertAllEqual(agent._observation_batch, first_observation[..., 0])
+      # No training happens in eval mode.
+      self.assertEqual(agent.training_steps, 0)
+
+      # This will now cause training to happen.
+      agent.eval_mode = False
+      # Having a low replay memory add_count will prevent any of the
+      # train/prefetch/sync ops from being called.
+      agent._replay.memory.add_count = 0
+      second_observation = np.ones(
+          [self.env_batch_size, self.observation_shape[0],
+           self.observation_shape[1], 1]) * 2
+      agent.begin_episode(second_observation)
+      # The agent's state will be reset, so we will only be left with the all-2s
+      # observation.
+      expected_state[:, :, :, -1] = np.full(
+          (self.env_batch_size, self.observation_shape[0],
+           self.observation_shape[1]), 2
+      )
+      self.assertAllEqual(agent.state_batch, expected_state)
+      self.assertAllEqual(agent._observation_batch,
+                          second_observation[:, :, :, 0])
+      # training_steps is incremented since we set eval_mode to False.
+      self.assertEqual(agent.training_steps, 1)
+
+
+if __name__ == '__main__':
+  tf.test.main()
diff --git a/tensor2tensor/rl/batch_runner_test.py b/tensor2tensor/rl/batch_runner_test.py
new file mode 100644
index 000000000..86d1d6645
--- /dev/null
+++ b/tensor2tensor/rl/batch_runner_test.py
@@ -0,0 +1,297 @@
+# Copyright 2018 The Dopamine Authors.
+#
+# 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 BatchRunner."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import shutil
+
+from absl import flags
+from dopamine.discrete_domains import run_experiment
+from dopamine.discrete_domains import logger
+import gin.tf
+import numpy as np
+import mock
+
+from tensor2tensor.rl import dopamine_connector
+
+import tensorflow as tf
+
+
+FLAGS = flags.FLAGS
+
+
+def _create_mock_checkpointer():
+  mock_checkpointer = mock.Mock()
+  test_dictionary = {'current_iteration': 1729,
+                     'logs': 'logs'}
+  mock_checkpointer.load_checkpoint.return_value = test_dictionary
+  return mock_checkpointer
+
+
+class MockEnvironment(object):
+  """Mock environment for testing."""
+
+  def __init__(self, max_steps=10, reward_multiplier=1):
+    self._observation = 0
+    self.max_steps = max_steps
+    self.reward_multiplier = reward_multiplier
+    self.game_over = False
+
+  def reset(self):
+    self._observation = 0
+    return self._observation
+
+  def step(self, action):
+    self._observation += 1
+    action_reward_multiplier = -1 if action > 0 else 1
+    reward_multiplier = self.reward_multiplier * action_reward_multiplier
+    reward = self._observation * reward_multiplier
+    is_terminal = self._observation >= self.max_steps
+    self.game_over = is_terminal
+
+    unused = 0
+    return (self._observation, reward, is_terminal, unused)
+
+  def render(self, mode):
+    pass
+
+
+class BatchEnv(object):
+  """
+
+  Batch of environments. Assumes that all throws 'done' on the same step.
+
+  Observations and rewards are returned as arrays, done as single value.
+  """
+  # TODO: this can be used for mbrl pipeline (for both simulated and real env),
+  #  move it to dopamine_connector.py (rename it?)
+  def __init__(self, envs):
+    self.env_batch = envs
+    self.batch_size = len(self.env_batch)
+    self.max_steps = self.env_batch[0].max_steps
+    assert np.all(self.max_steps == env.max_steps for env in self.env_batch)
+
+  def step(self, actions):
+    ret = [env.step(action) for env, action in zip(self.env_batch, actions)]
+    obs, rewards, dones, infos = [np.array(r) for r in zip(*ret)]
+    done = dones[0]
+    assert np.all(done == dones)
+    self.game_over = done
+    return obs, rewards, done, infos
+
+  def reset(self):
+    return np.array([env.reset() for env in self.env_batch])
+
+  def render(self, mode):
+    pass
+
+
+class MockLogger(object):
+  """Class to mock the experiment logger."""
+
+  def __init__(self, test_cls=None, run_asserts=True, data=None):
+    self._test_cls = test_cls
+    self._run_asserts = run_asserts
+    self._iter = 0
+    self._calls_to_set = 0
+    self._calls_to_log = 0
+    self.data = data
+
+  def __setitem__(self, key, val):
+    if self._run_asserts:
+      self._test_cls.assertEqual('iteration_{:d}'.format(self._iter), key)
+      self._test_cls.assertEqual('statistics', val)
+      self._iter += 1
+    self._calls_to_set += 1
+
+  def log_to_file(self, filename_prefix, iteration_number):
+    if self._run_asserts:
+      self._test_cls.assertEqual(
+          'prefix_{}'.format(self._iter - 1),
+          '{}_{}'.format(filename_prefix, iteration_number))
+    self._calls_to_log += 1
+
+
+class RunExperimentTest(tf.test.TestCase):
+
+  @mock.patch.object(gin, 'parse_config_files_and_bindings')
+  def testLoadGinConfigs(self, mock_parse_config_files_and_bindings):
+    gin_files = ['file1', 'file2', 'file3']
+    gin_bindings = ['binding1', 'binding2']
+    run_experiment.load_gin_configs(gin_files, gin_bindings)
+    self.assertEqual(1, mock_parse_config_files_and_bindings.call_count)
+    mock_args, mock_kwargs = mock_parse_config_files_and_bindings.call_args
+    self.assertEqual(gin_files, mock_args[0])
+    self.assertEqual(gin_bindings, mock_kwargs['bindings'])
+    self.assertFalse(mock_kwargs['skip_unknown'])
+
+
+class BatchedRunnerTest(tf.test.TestCase):
+  """Modified tests from dopamine run_experiment_test.py."""
+
+  # TODO: decide if we want to use and modify more tests from
+  # dopamine/tests/atari/run_experiment_test.py (e.g.  testRunExperiment.py)
+
+  def _agent_step(self, rewards, observations):
+    # We verify that rewards are clipped (and set by MockEnvironment as a
+    # function of observation)
+    # observation = observations[0]
+    # expected_rewards = [1 if observation % 2 else -1]
+    # self.assertEqual(expected_reward, reward)
+    actions = [ob % 2 for ob in observations]
+    return actions
+
+  def prepare_mock_agent(self, batch_size):
+    assert batch_size % 2 == 0, "Some of tests assume that batch_size % 2 == 0"
+    self.batch_size = batch_size
+    self._agent = mock.Mock()
+    self._agent.begin_episode.side_effect = \
+      lambda x: np.repeat(0, self.batch_size)
+    self._agent.step.side_effect = self._agent_step
+    self._create_agent_fn = lambda x, y, summary_writer: self._agent
+
+  def setUp(self):
+    super(BatchedRunnerTest, self).setUp()
+    self._test_subdir = "/tmp/dopamine_tests"
+    shutil.rmtree(self._test_subdir, ignore_errors=True)
+    os.makedirs(self._test_subdir)
+    self.prepare_mock_agent(batch_size=4)
+
+  def testRunEpisodeBatch(self):
+    max_steps_per_episode = 11
+    batch_size = self.batch_size
+    reward_multipliers = [-1, 1] * int(batch_size / 2)
+    envs = [MockEnvironment(reward_multiplier=rm) for rm in reward_multipliers]
+    environment = BatchEnv(envs)
+    runner = dopamine_connector.BatchRunner(
+        self._test_subdir, self._create_agent_fn,
+        create_environment_fn=lambda: environment,
+        max_steps_per_episode=max_steps_per_episode)
+    step_number, total_rewards = runner._run_one_episode()
+
+    self.assertEqual(self._agent.step.call_count, environment.max_steps - 1)
+    self.assertEqual(self._agent.end_episode.call_count, 1)
+    self.assertEqual(environment.max_steps, step_number / batch_size)
+    # Expected reward will be \sum_{i=0}^{9} (-1)**i * i = -5 when reward
+    # multiplier=1
+    self.assertAllEqual(np.array(reward_multipliers) * -5, total_rewards)
+
+  def testRunOneEpisodeWithLowMaxSteps(self):
+    max_steps_per_episode = 2
+    batch_size = self.batch_size
+    reward_multipliers = [-1, 1] * int(batch_size / 2)
+    envs = [MockEnvironment(reward_multiplier=rm) for rm in reward_multipliers]
+    environment = BatchEnv(envs)
+    runner = dopamine_connector.BatchRunner(
+        self._test_subdir, self._create_agent_fn,
+        create_environment_fn=lambda: environment,
+        max_steps_per_episode=max_steps_per_episode)
+    step_number, total_rewards = runner._run_one_episode()
+
+    self.assertEqual(self._agent.step.call_count, max_steps_per_episode - 1)
+    self.assertEqual(self._agent.end_episode.call_count, 1)
+    self.assertEqual(max_steps_per_episode, step_number / batch_size)
+    self.assertAllEqual(np.array(reward_multipliers) * -1, total_rewards)
+
+  def testRunOnePhase(self):
+    batch_size = self.batch_size
+    environment_steps = 2
+    max_steps = environment_steps * batch_size * 10
+
+    envs = [MockEnvironment(max_steps=environment_steps)
+            for _ in range(batch_size)]
+
+    environment = BatchEnv(envs)
+    runner = dopamine_connector.BatchRunner(
+        self._test_subdir, self._create_agent_fn,
+        create_environment_fn=lambda: environment)
+
+    statistics = []
+
+    step_number, sum_returns, num_episodes = runner._run_one_phase(
+        max_steps, statistics, "test")
+    calls_to_run_episode = int(max_steps / (environment_steps * batch_size))
+    self.assertEqual(self._agent.step.call_count, calls_to_run_episode)
+    self.assertEqual(self._agent.end_episode.call_count, calls_to_run_episode)
+    self.assertEqual(max_steps, step_number)
+    self.assertEqual(-1 * calls_to_run_episode * batch_size, sum_returns)
+    self.assertEqual(calls_to_run_episode, num_episodes / batch_size)
+    expected_statistics = []
+    for _ in range(calls_to_run_episode * batch_size):
+      expected_statistics.append({
+          "test_episode_lengths": 2,
+          "test_episode_returns": -1
+      })
+    self.assertEqual(len(expected_statistics), len(statistics))
+    for expected_stats, stats in zip(expected_statistics, statistics):
+      self.assertDictEqual(expected_stats, stats)
+
+  def testRunOneIteration(self):
+    environment_steps = 2
+    batch_size = self.batch_size
+    envs = [MockEnvironment(max_steps=environment_steps)
+            for _ in range(batch_size)]
+
+    environment = BatchEnv(envs)
+
+    training_steps = 20 * batch_size
+    evaluation_steps = 10 * batch_size
+
+    runner = dopamine_connector.BatchRunner(
+        self._test_subdir, self._create_agent_fn,
+        create_environment_fn=lambda: environment,
+        training_steps=training_steps, evaluation_steps=evaluation_steps
+    )
+
+    dictionary = runner._run_one_iteration(1)
+    train_rollouts = int(training_steps / environment_steps)
+    eval_rollouts = int(evaluation_steps / environment_steps)
+    expected_dictionary = {
+        "train_episode_lengths": [2 for _ in range(train_rollouts)],
+        "train_episode_returns": [-1 for _ in range(train_rollouts)],
+        "train_average_return": [-1],
+        "eval_episode_lengths": [2 for _ in range(eval_rollouts)],
+        "eval_episode_returns": [-1 for _ in range(eval_rollouts)],
+        "eval_average_return": [-1]
+    }
+    self.assertDictEqual(expected_dictionary, dictionary)
+
+  @mock.patch.object(logger, "Logger")
+  def testLogExperiment(self, mock_logger_constructor):
+    #TODO: We probably do not need this test, dopamine test for Runner is enugh
+    # here. Remove this?
+    log_every_n = 2
+    logging_file_prefix = "prefix"
+    statistics = "statistics"
+    experiment_logger = MockLogger(test_cls=self)
+    mock_logger_constructor.return_value = experiment_logger
+    runner = dopamine_connector.BatchRunner(
+        self._test_subdir, self._create_agent_fn,
+        create_environment_fn=mock.Mock,
+        logging_file_prefix=logging_file_prefix,
+        log_every_n=log_every_n)
+    num_iterations = 10
+    for i in range(num_iterations):
+      runner._log_experiment(i, statistics)
+    self.assertEqual(num_iterations, experiment_logger._calls_to_set)
+    self.assertEqual((num_iterations / log_every_n),
+                     experiment_logger._calls_to_log)
+
+
+if __name__ == '__main__':
+  tf.test.main()
diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 9e1ad3765..c9dab6bbb 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -20,6 +20,8 @@
 from __future__ import print_function
 
 import copy
+import random
+import sys
 
 from dopamine.agents.dqn import dqn_agent
 from dopamine.replay_memory import circular_replay_buffer
@@ -46,52 +48,6 @@
 # pylint: enable=g-import-not-at-top
 
 
-class ResizeObservation(gym.ObservationWrapper):
-  """TODO(konradczechowski): Add doc-string."""
-
-  def __init__(self, env, size=84):
-    """Based on WarpFrame from openai baselines atari_wrappers.py.
-
-    Dopamine also uses cv2.resize(..., interpolation=cv2.INTER_AREA).
-
-    Args:
-      env: TODO(konradczechowski): Add doc-string.
-      size: TODO(konradczechowski): Add doc-string.
-    """
-    gym.ObservationWrapper.__init__(self, env)
-    self.width = size
-    self.height = size
-    assert env.observation_space.dtype == np.uint8
-    self.observation_space = spaces.Box(
-        low=0,
-        high=255,
-        shape=(self.height, self.width, env.observation_space.shape[2]),
-        dtype=np.uint8)
-
-  def observation(self, frame):
-    if not cv2:
-      return frame
-    return cv2.resize(
-        frame, (self.width, self.height), interpolation=cv2.INTER_AREA)
-
-
-class GameOverOnDone(Wrapper):
-  """TODO(konradczechowski): Add doc-string."""
-
-  def __init__(self, env):
-    Wrapper.__init__(self, env)
-    self.game_over = False
-
-  def reset(self, **kwargs):
-    self.game_over = False
-    return self.env.reset(**kwargs)
-
-  def step(self, action):
-    ob, reward, done, info = self.env.step(action)
-    self.game_over = done
-    return ob, reward, done, info
-
-
 class _DQNAgent(dqn_agent.DQNAgent):
   """Modify dopamine DQNAgent to match our needs.
 
@@ -99,10 +55,10 @@ class _DQNAgent(dqn_agent.DQNAgent):
   (some of) terminal episode transitions in training.
   """
 
-  def __init__(self, replay_capacity, batch_size, generates_trainable_dones,
-               **kwargs):
+  def __init__(self, replay_capacity, buffer_batch_size,
+               generates_trainable_dones, **kwargs):
     self._replay_capacity = replay_capacity
-    self._batch_size = batch_size
+    self._buffer_batch_size = buffer_batch_size
     self._generates_trainable_dones = generates_trainable_dones
     super(_DQNAgent, self).__init__(**kwargs)
 
@@ -112,7 +68,7 @@ def _build_replay_buffer(self, use_staging):
         observation_shape=dqn_agent.NATURE_DQN_OBSERVATION_SHAPE,
         stack_size=dqn_agent.NATURE_DQN_STACK_SIZE,
         replay_capacity=self._replay_capacity,
-        batch_size=self._batch_size,
+        batch_size=self._buffer_batch_size,
         update_horizon=self.update_horizon,
         gamma=self.gamma,
         extra_storage_types=None,
@@ -128,6 +84,145 @@ def _build_replay_buffer(self, use_staging):
         **replay_buffer_kwargs)
 
 
+class BatchDQNAgent(_DQNAgent):
+  """
+  Episodes are stored on done.
+
+  Assumes that all rollouts in batch would end at the same moment.
+  """
+
+  def __init__(self, env_batch_size, *args, **kwargs):
+    super(BatchDQNAgent, self).__init__(*args, **kwargs)
+    self.env_batch_size = env_batch_size
+    obs_size = dqn_agent.NATURE_DQN_OBSERVATION_SHAPE
+    state_shape = [self.env_batch_size, obs_size[0], obs_size[1],
+                   dqn_agent.NATURE_DQN_STACK_SIZE]
+    self.state_batch = np.zeros(state_shape)
+    self.state = None  # assure it will be not used
+    self._observation = None  # assure it will be not used
+    self.reset_current_rollouts()
+
+  def reset_current_rollouts(self):
+    self._current_rollouts = [[] for _ in range(self.env_batch_size)]
+
+  def _record_observation(self, observation_batch):
+    # Set current observation. Represents an (batch_size x 84 x 84 x 1) image
+    # frame.
+    observation_batch = np.array(observation_batch)
+    self._observation_batch = observation_batch[:, :, :, 0]
+    # Swap out the oldest frames with the current frames.
+    self.state_batch = np.roll(self.state_batch, -1, axis=3)
+    self.state_batch[:, :, :, -1] = self._observation_batch
+
+  def _reset_state(self):
+    self.state_batch.fill(0)
+
+  def begin_episode(self, observation):
+    self._reset_state()
+    self._record_observation(observation)
+
+    if not self.eval_mode:
+      self._train_step()
+
+    self.action = self._select_action()
+    return self.action
+
+  def _update_current_rollouts(self, last_observation, action, reward,
+                              are_terminal):
+    transitions = zip(last_observation, action, reward, are_terminal)
+    for transition, rollout in zip(transitions, self._current_rollouts):
+      rollout.append(transition)
+
+  def _store_current_rollouts(self):
+    for rollout in self._current_rollouts:
+      for transition in rollout:
+        self._store_transition(*transition)
+    self.reset_current_rollouts()
+
+  def step(self, reward, observation):
+    self._last_observation = self._observation_batch
+    self._record_observation(observation)
+
+    if not self.eval_mode:
+      self._update_current_rollouts(self._last_observation, self.action, reward,
+                                    [False] * self.env_batch_size)
+      # We want to have the same train_step:env_step ratio not depending on
+      # batch size.
+      for _ in range(self.env_batch_size):
+        self._train_step()
+
+    self.action = self._select_action()
+    return self.action
+
+  def end_episode(self, reward):
+    if not self.eval_mode:
+      self._update_current_rollouts(self._observation_batch, self.action, reward,
+                                    [True] * self.env_batch_size)
+      self._store_current_rollouts()
+
+  def _select_action(self):
+    epsilon = self.epsilon_eval if self.eval_mode else self.epsilon_fn(
+        self.epsilon_decay_period,
+        self.training_steps,
+        self.min_replay_history,
+        self.epsilon_train)
+
+    def choose_action(ix):
+      if random.random() <= epsilon:
+        # Choose a random action with probability epsilon.
+        return random.randint(0, self.num_actions - 1)
+      else:
+        # Choose the action with highest Q-value at the current state.
+        return self._sess.run(self._q_argmax,
+                              {self.state_ph: self.state_batch[ix:ix+1]})
+
+    return np.array([choose_action(ix) for ix in range(self.env_batch_size)])
+
+
+class BatchRunner(run_experiment.Runner):
+  """
+
+  Assumes that all environments would end at the same moment.
+  """
+  def __init__(self, base_dir, create_agent_fn, **kwargs):
+    super(BatchRunner, self).__init__(base_dir, create_agent_fn, **kwargs)
+    self.batch_size = self._environment.batch_size
+
+  def _run_one_episode(self):
+    # This assumes that everything inside _run_one_episode works on batches,
+    # which is risky for future.
+    steps_number, total_rewards = super(BatchRunner, self)._run_one_episode()
+    return steps_number * self.batch_size, total_rewards
+
+  def _run_one_phase(self, min_steps, statistics, run_mode_str):
+    # Mostly copy of parent method.
+    step_count = 0
+    num_episodes = 0
+    sum_returns = 0.
+
+    while step_count < min_steps:
+      num_steps, episode_returns = self._run_one_episode()
+      for episode_return in episode_returns:
+        statistics.append({
+            '{}_episode_lengths'.format(run_mode_str):
+                num_steps / self.batch_size,
+            '{}_episode_returns'.format(run_mode_str): episode_return
+        })
+      step_count += num_steps
+      sum_returns += sum(episode_returns)
+      num_episodes += self.batch_size
+      # We use sys.stdout.write instead of tf.logging so as to flush frequently
+      # without generating a line break.
+      sys.stdout.write('Steps executed: {} '.format(step_count) +
+                       'Batch episodes steps: {} '.format(num_steps) +
+                       'Returns: {}\r'.format(episode_returns))
+      sys.stdout.flush()
+    return step_count, sum_returns, num_episodes
+
+  def close(self):
+    self._environment.close()
+
+
 class _OutOfGraphReplayBuffer(OutOfGraphReplayBuffer):
   """Replay not sampling artificial_terminal transition.
 
@@ -177,7 +272,14 @@ def load(self, *args, **kwargs):
 
 
 def get_create_agent(agent_kwargs):
-  """TODO(): Document."""
+  """Factory for dopamine agent initialization.
+
+  Args:
+    agent_kwargs: dict of BatchDQNAgent parameters
+
+  Returns:
+    Function(sess, environment, summary_writer) -> BatchDQNAgent instance.
+  """
 
   def create_agent(sess, environment, summary_writer=None):
     """Creates a DQN agent.
@@ -192,7 +294,8 @@ def create_agent(sess, environment, summary_writer=None):
     Returns:
       a DQN agent.
     """
-    return _DQNAgent(
+    return BatchDQNAgent(
+        env_batch_size=environment.batch_size,
         sess=sess,
         num_actions=environment.action_space.n,
         summary_writer=summary_writer,
@@ -202,23 +305,171 @@ def create_agent(sess, environment, summary_writer=None):
   return create_agent
 
 
-def get_create_env_fun(batch_env_fn, time_limit):
-  """TODO(konradczechowski): Add doc-string."""
+class ResizeBatchObservation(object):
+  """Wrapper resizing observations for batched environment.
+
+  Dopamine also uses cv2.resize(..., interpolation=cv2.INTER_AREA).
+
+  Attributes:
+    batch_env: batched environment
+    size: size of width and height for returned observations
+  """
+
+  def __init__(self, batch_env, size=84):
+    self.size = size
+    self.batch_env = batch_env
+
+  def observation(self, frames):
+    if not cv2:
+      return frames
+    return np.array([cv2.resize(
+        frame, (self.size, self.size), interpolation=cv2.INTER_AREA)
+        for frame in frames])
+
+  def step(self, actions):
+    obs, rewards, dones = self.batch_env.step(actions)
+    obs = self.observation(obs)
+    return obs, rewards, dones
+
+  def reset(self, *args, **kwargs):
+    return self.observation(self.batch_env.reset(*args, **kwargs))
 
-  def create_env_fun(game_name, sticky_actions=True):
+  @property
+  def action_space(self):
+    return self.batch_env.action_space
+
+  @property
+  def batch_size(self):
+    return self.batch_env.batch_size
+
+  def close(self):
+    self.batch_env.close()
+
+
+class DopamineBatchEnv(object):
+  """Batch of environments.
+
+  Assumes that all given environments finishes at the same time.
+
+  Observations and rewards are returned as batches (arrays). Done is returned
+  as single boolean.
+  """
+  def __init__(self, batch_env, max_episode_steps):
+    self.batch_env = batch_env
+    self._max_episode_steps = max_episode_steps
+    self.game_over = None
+    self._elapsed_steps = 0
+
+  def reset(self):
+    self.game_over = False
+    self._elapsed_steps = 0
+    return np.array(self.batch_env.reset())
+
+  def step(self, actions):
+    self._elapsed_steps += 1
+    obs, rewards, dones = \
+        [np.array(r) for r in self.batch_env.step(actions)]
+    if self._elapsed_steps > self._max_episode_steps:
+      done = True
+      if self._elapsed_steps > self._max_episode_steps + 1:
+        rewards.fill(0)
+    else:
+      done = dones[0]
+      assert np.all(done == dones), "Current modifications of Dopamine " \
+                                    "require same number of steps for each " \
+                                    "environment in batch"
+      del dones
+
+    self.game_over = done
+    return obs, rewards, done, {}
+
+  def render(self, mode):
+    pass
+
+  def close(self):
+    self.batch_env.close()
+
+  @property
+  def action_space(self):
+    return self.batch_env.action_space
+
+  @property
+  def batch_size(self):
+    return self.batch_env.batch_size
+
+
+class PaddedTrajectoriesEnv(DopamineBatchEnv):
+  """ Padd finished episodes with zeros.
+
+  Allow episodes in batch to end on different timesteps, return zero
+  observations and rewards for finished ones. Return done=True when all
+  episodes are finished.
+
+  Note that output of this class might be misleading - the agent/evaluator
+  which uses this environment gets false information about when episodes have
+  ended. This class is used for informal check of Batched dopamine
+  implementation in model-free pipeline.
+  """
+
+  def reset(self):
+    self.done_envs = [False] * self.batch_size
+    self.game_over = False
+    self._elapsed_steps = 0
+    return np.array(self.batch_env.reset())
+
+  def step(self, actions):
+    if any(self.done_envs):
+      print("Warning, some environments already ended, using mocked data.")
+
+    self._elapsed_steps += 1
+    obs, rewards, dones = \
+        [np.array(r) for r in self.batch_env.step(actions)]
+    for i, ignore in enumerate(self.done_envs):
+      if ignore:
+        obs[i] = np.zeros(obs[i].shape, dtype=obs.dtype)
+        rewards[i] = 0
+      if dones[i]:
+        self.batch_env.reset([i])
+        self.done_envs[i] = True
+
+    all_done = all(self.done_envs)
+
+    if self._elapsed_steps > self._max_episode_steps:
+      all_done = True
+      if self._elapsed_steps > self._max_episode_steps + 1:
+        rewards.fill(0)
+
+    self.game_over = all_done
+    return obs, rewards, all_done, {}
+
+
+def get_create_batch_env_fun(batch_env_fn, time_limit):
+  """Factory for dopamine environment initialization function.
+
+  Args:
+    batch_env_fn: function(in_graph: bool) -> batch environment.
+    time_limit: time steps limit for environment.
+
+  Returns:
+    function (with optional, unused parameters) initializing environment.
+  """
+
+  def create_env_fun(game_name=None, sticky_actions=None):
     del game_name, sticky_actions
     batch_env = batch_env_fn(in_graph=False)
-    env = FlatBatchEnv(batch_env)
-    env = TimeLimit(env, max_episode_steps=time_limit)
-    env = ResizeObservation(env)  # pylint: disable=redefined-variable-type
-    env = GameOverOnDone(env)
-    return env
+    batch_env = ResizeBatchObservation(batch_env)  # pylint: disable=redefined-variable-type
+    batch_env = DopamineBatchEnv(batch_env, max_episode_steps=time_limit)
+    return batch_env
 
   return create_env_fun
 
 
 def _parse_hparams(hparams):
-  """TODO(konradczechowski): Add doc-string."""
+  """Split hparams, based on key prefixes.
+
+  Returns:
+    Tuple of hparams for respectably: agent, optimizer, runner, replay_buffer.
+  """
   prefixes = ["agent_", "optimizer_", "runner_", "replay_buffer_"]
   ret = []
 
@@ -242,9 +493,8 @@ def _get_optimizer(params):
 class DQNLearner(PolicyLearner):
   """Interface for learning dqn implemented in dopamine."""
 
-  def __init__(self, frame_stack_size, base_event_dir, agent_model_dir):
-    super(DQNLearner, self).__init__(frame_stack_size, base_event_dir,
-                                     agent_model_dir)
+  def __init__(self, *args, **kwargs):
+    super(DQNLearner, self).__init__(*args, **kwargs)
     self.completed_iterations = 0
 
   def _target_iteractions_and_steps(self, num_env_steps, save_continuously,
@@ -269,10 +519,10 @@ def create_runner(self, env_fn, hparams, target_iterations,
     agent_params["optimizer"] = optimizer
     agent_params.update(replay_buffer_params)
     create_agent_fn = get_create_agent(agent_params)
-    runner = run_experiment.Runner(
+    runner = BatchRunner(
         base_dir=self.agent_model_dir,
         create_agent_fn=create_agent_fn,
-        create_environment_fn=get_create_env_fun(
+        create_environment_fn=get_create_batch_env_fun(
             env_fn, time_limit=hparams.time_limit),
         evaluation_steps=0,
         num_iterations=target_iterations,
@@ -290,9 +540,10 @@ def train(self,
             num_env_steps=None,
             env_step_multiplier=1,
             eval_env_fn=None,
-            report_fn=None):
+            report_fn=None,
+            model_save_fn=None):
     # TODO(konradczechowski): evaluation during training (with eval_env_fun)
-    del epoch, eval_env_fn, simulated, report_fn
+    del epoch, eval_env_fn, simulated, report_fn, model_save_fn
     if num_env_steps is None:
       num_env_steps = hparams.num_frames
 
@@ -305,13 +556,13 @@ def train(self,
       self._target_iteractions_and_steps(
           num_env_steps=num_env_steps * env_step_multiplier,
           save_continuously=save_continuously,
-          save_every_steps=hparams.save_every_steps,)
+          save_every_steps=hparams.save_every_steps)
 
     with tf.Graph().as_default():
       runner = self.create_runner(env_fn, hparams, target_iterations,
                                   training_steps_per_iteration)
       runner.run_experiment()
-
+      runner.close()
     self.completed_iterations = target_iterations
 
   def evaluate(self, env_fn, hparams, sampling_temp):
@@ -323,7 +574,7 @@ def evaluate(self, env_fn, hparams, sampling_temp):
         "agent_epsilon_eval", min(hparams.agent_epsilon_eval * sampling_temp, 1)
     )
 
-    create_environment_fn = get_create_env_fun(
+    create_environment_fn = get_create_batch_env_fun(
         env_fn, time_limit=hparams.time_limit)
     env = create_environment_fn(
         game_name="unused_arg", sticky_actions="unused_arg")
@@ -331,13 +582,13 @@ def evaluate(self, env_fn, hparams, sampling_temp):
     with tf.Graph().as_default():
       runner = self.create_runner(env_fn, hparams, target_iterations,
                                   training_steps_per_iteration)
+      assert runner.batch_size == 1
       agent = runner._agent  # pylint: disable=protected-access
+      runner.close()
       del runner
       agent.eval = True
 
-      # TODO(konradczechowski): correct number of episodes, when this will
-      # be hparam
-      for _ in range(30):
+      for _ in range(hparams.eval_episodes_num):
         # Run single episode
         ob = env.reset()
         action = agent.begin_episode(ob)
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index eac13e012..bfee48de0 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -294,6 +294,5 @@ def _video_reset_writer(self):
       self._video_writer.finish_to_disk()
     self._video_writer = None
 
-  def __del__(self):
+  def close(self):
     self._video_reset_writer()
-    super(SimulatedBatchEnv, self).__del__()
diff --git a/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index e0ffd9805..fd220837f 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -86,9 +86,6 @@ def reset(self, indices=None):
     if indices is None:
       indices = np.array(range(self.batch_size))
     obs = self._sess.run(self._reset_op, feed_dict={self._indices_t: indices})
-    # TODO(pmilos): remove if possible
-    # obs[:, 0, 0, 0] = 0
-    # obs[:, 0, 0, 1] = 255
     return obs
 
   def step(self, actions):
@@ -99,3 +96,4 @@ def step(self, actions):
 
   def close(self):
     self._sess.close()
+    self._batch_env.close()
diff --git a/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
index 82051f159..ec59e4c8a 100644
--- a/tensor2tensor/rl/policy_learner.py
+++ b/tensor2tensor/rl/policy_learner.py
@@ -45,8 +45,6 @@ def train(
       report_fn=None
   ):
     """Train."""
-    # TODO(konradczechowski): pass name_scope instead of epoch?
-    # TODO(konradczechowski): move 'simulated' to  batch_env
     raise NotImplementedError()
 
   def evaluate(self, env_fn, hparams, sampling_temp):
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index f34303d92..625c42916 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -39,6 +39,7 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl import trainer_model_based_params
+from tensor2tensor.rl.dopamine_connector import DQNLearner
 from tensor2tensor.rl.restarter import Restarter
 from tensor2tensor.utils import trainer_lib
 
@@ -287,6 +288,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   metrics = {}
 
   # Collect data from the real environment.
+  policy_model_dir = directories["policy"]
   tf.logging.info("Initial training of the policy in real environment.")
   train_agent_real_env(env, learner, hparams, epoch)
   metrics["mean_reward/train/clipped"] = rl_utils.compute_mean_reward(
diff --git a/tensor2tensor/rl/trainer_model_based_dqn_test.py b/tensor2tensor/rl/trainer_model_based_dqn_test.py
deleted file mode 100644
index 50c056649..000000000
--- a/tensor2tensor/rl/trainer_model_based_dqn_test.py
+++ /dev/null
@@ -1,41 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Tiny run of trainer_model_based. Smoke test."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-# from tensor2tensor.rl import trainer_model_based
-
-import tensorflow as tf
-
-FLAGS = tf.flags.FLAGS
-
-
-class ModelRLExperimentTest(tf.test.TestCase):
-
-  def test_dqn_basic(self):
-    # TODO(afrozm): The latest changes in Dopamine break this test, so
-    # temporarily disabling this test.
-    pass
-    # FLAGS.output_dir = tf.test.get_temp_dir()
-    # FLAGS.loop_hparams_set = "rlmb_dqn_tiny"
-    # FLAGS.schedule = "train"  # skip evaluation for world model training
-    # trainer_model_based.main(None)
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 37eb47ddc..a60efef1e 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -111,6 +111,7 @@ def _rlmb_base():
       # This is only used for world-model evaluation currently, PolicyLearner
       # uses algorithm specific hparams to set this during training.
       simulated_rollout_length=50,
+      wm_policy_param_sharing=False,
 
       # To be overridden.
       base_algo="",
@@ -145,7 +146,6 @@ def rlmb_ppo_base():
       # Number of simulated environments to train on simultaneously.
       simulated_batch_size=16,
       eval_batch_size=32,
-      wm_policy_param_sharing=False,
 
       # Unused; number of PPO epochs is calculated from the real frame limit.
       real_ppo_epochs_num=0,
@@ -194,13 +194,17 @@ def rlmb_dqn_base():
       base_algo="dqn",
       base_algo_params="dqn_original_params",
       real_batch_size=1,
-      simulated_batch_size=1,
+      simulated_batch_size=16,
       dqn_agent_generates_trainable_dones=False,
       eval_batch_size=1,
       # Must be equal to dqn_time_limit for now
       simulated_rollout_length=simulated_rollout_length,
       dqn_time_limit=simulated_rollout_length,
       simulation_flip_first_random_for_beginning=False,
+      dqn_eval_episodes_num=3,
+
+      # TODO(kc): only for model-free compatibility, remove this
+      epochs_num=-1,
   )
   update_hparams(hparams, dqn_params)
   return hparams
diff --git a/tensor2tensor/rl/trainer_model_based_test.py b/tensor2tensor/rl/trainer_model_based_test.py
index 0c836162c..87f5e61e3 100644
--- a/tensor2tensor/rl/trainer_model_based_test.py
+++ b/tensor2tensor/rl/trainer_model_based_test.py
@@ -27,12 +27,18 @@
 
 class ModelRLExperimentTest(tf.test.TestCase):
 
-  def test_basic(self):
+  def _test_hparams_skip_evaluation(self, hparams_set):
     FLAGS.output_dir = tf.test.get_temp_dir()
-    FLAGS.loop_hparams_set = "rlmb_tiny"
+    FLAGS.loop_hparams_set = hparams_set
     FLAGS.schedule = "train"  # skip evaluation for world model training
     trainer_model_based.main(None)
 
+  def test_basic(self):
+    self._test_hparams_skip_evaluation("rlmb_tiny")
+
+  def test_dqn_basic(self):
+    self._test_hparams_skip_evaluation("rlmb_dqn_tiny")
+
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 5d7a0974a..0d48a372d 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -103,41 +103,51 @@ def train(hparams, output_dir, env_problem_name, report_fn=None):
   tf.logging.vlog(1, "Policy HParams : %s",
                   misc_utils.pprint_hparams(policy_hparams))
 
-  total_steps = policy_hparams.epochs_num
-  tf.logging.vlog(2, "total_steps: %d", total_steps)
-
-  eval_every_epochs = policy_hparams.eval_every_epochs
-  tf.logging.vlog(2, "eval_every_epochs: %d", eval_every_epochs)
-
-  if eval_every_epochs == 0:
-    eval_every_epochs = total_steps
-  policy_hparams.eval_every_epochs = 0
-
-  metric_name = rl_utils.get_metric_name(
-      sampling_temp=hparams.eval_sampling_temps[0],
-      max_num_noops=hparams.eval_max_num_noops,
-      clipped=False
-  )
+  # TODO(konradczechowski): remove base_algo dependance, when evaluation method
+  # will be decided
+  if hparams.base_algo == "ppo":
+    print("\n\n\npolicy_hparams {}\n\n\n".format(policy_hparams))
+    total_steps = policy_hparams.epochs_num
+    tf.logging.vlog(2, "total_steps: %d", total_steps)
+
+    eval_every_epochs = policy_hparams.eval_every_epochs
+    tf.logging.vlog(2, "eval_every_epochs: %d", eval_every_epochs)
+
+    if eval_every_epochs == 0:
+      eval_every_epochs = total_steps
+    policy_hparams.eval_every_epochs = 0
+
+    metric_name = rl_utils.get_metric_name(
+        sampling_temp=hparams.eval_sampling_temps[0],
+        max_num_noops=hparams.eval_max_num_noops,
+        clipped=False
+    )
+
+    tf.logging.vlog(1, "metric_name: %s", metric_name)
+
+    eval_metrics_dir = os.path.join(output_dir, "eval_metrics")
+    eval_metrics_dir = os.path.expanduser(eval_metrics_dir)
+    tf.gfile.MakeDirs(eval_metrics_dir)
+    eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
+
+    def evaluate_on_new_model(model_dir_path):
+      global step
+      eval_metrics = rl_utils.evaluate_all_configs(hparams, model_dir_path)
+      tf.logging.info(
+          "Agent eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
+      rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, step)
+      if report_fn:
+        report_fn(eval_metrics[metric_name], step)
+      step += 1
+
+    policy_hparams.epochs_num = total_steps
+    policy_hparams.save_models_every_epochs = eval_every_epochs
+  else:
+    def evaluate_on_new_model(model_dir_path):
+      del model_dir_path
+      raise NotImplementedError(
+          "This function is currently implemented only for ppo")
 
-  tf.logging.vlog(1, "metric_name: %s", metric_name)
-
-  eval_metrics_dir = os.path.join(output_dir, "eval_metrics")
-  eval_metrics_dir = os.path.expanduser(eval_metrics_dir)
-  tf.gfile.MakeDirs(eval_metrics_dir)
-  eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
-
-  def evaluate_on_new_model(model_dir_path):
-    global step
-    eval_metrics = rl_utils.evaluate_all_configs(hparams, model_dir_path)
-    tf.logging.info(
-        "Agent eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
-    rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, step)
-    if report_fn:
-      report_fn(eval_metrics[metric_name], step)
-    step += 1
-
-  policy_hparams.epochs_num = total_steps
-  policy_hparams.save_models_every_epochs = eval_every_epochs
   learner.train(env_fn,
                 policy_hparams,
                 simulated=False,
diff --git a/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index 98c28feca..3e21a1387 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -28,16 +28,18 @@
 
 class TrainTest(tf.test.TestCase):
 
-  def test_train_pong(self):
-    hparams = registry.hparams("rlmf_original")
-    hparams.batch_size = 2
-    hparams.eval_sampling_temps = [0.0, 1.0]
-    hparams.add_hparam("ppo_epochs_num", 2)
-    hparams.add_hparam("ppo_epoch_length", 3)
+  def _test_hparams_set(self, hparams_set):
+    hparams = registry.hparams(hparams_set)
     FLAGS.output_dir = tf.test.get_temp_dir()
     trainer_model_free.train(hparams, FLAGS.output_dir,
                              env_problem_name=None)
 
+  def test_train_pong(self):
+    self._test_hparams_set("rlmf_tiny")
+
+  def test_train_pong_dqn(self):
+    self._test_hparams_set("rlmf_dqn_tiny")
+
 
 if __name__ == "__main__":
   tf.test.main()

From 151dc27eb1b9f169c7e08e9e1b660f011ea99796 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 25 Mar 2019 15:21:18 -0700
Subject: [PATCH 1832/2720] "Adding mixture transformer"

PiperOrigin-RevId: 240229309
---
 tensor2tensor/layers/common_layers.py         |  92 ++++
 tensor2tensor/layers/common_video.py          |   2 -
 tensor2tensor/models/research/rl.py           |  52 +--
 tensor2tensor/rl/batch_dqn_agent_test.py      | 157 -------
 tensor2tensor/rl/batch_runner_test.py         | 297 -------------
 tensor2tensor/rl/dopamine_connector.py        | 399 ++++--------------
 tensor2tensor/rl/envs/simulated_batch_env.py  |   3 +-
 .../rl/envs/simulated_batch_gym_env.py        |   4 +-
 tensor2tensor/rl/policy_learner.py            |   2 +
 tensor2tensor/rl/trainer_model_based.py       |   2 -
 .../rl/trainer_model_based_dqn_test.py        |  41 ++
 .../rl/trainer_model_based_params.py          |   8 +-
 tensor2tensor/rl/trainer_model_based_test.py  |  10 +-
 tensor2tensor/rl/trainer_model_free.py        |  78 ++--
 tensor2tensor/rl/trainer_model_free_test.py   |  14 +-
 15 files changed, 267 insertions(+), 894 deletions(-)
 delete mode 100644 tensor2tensor/rl/batch_dqn_agent_test.py
 delete mode 100644 tensor2tensor/rl/batch_runner_test.py
 create mode 100644 tensor2tensor/rl/trainer_model_based_dqn_test.py

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index f30032536..97a1ab3a1 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1790,6 +1790,98 @@ def padded_cross_entropy(logits,
     return tf.reduce_sum(xent * weights), tf.reduce_sum(weights)
 
 
+def padded_cross_entropy_mixture(logits,
+                                 labels,
+                                 label_smoothing,
+                                 num_mixtures,
+                                 weights_fn=weights_nonzero,
+                                 reduce_sum=False,
+                                 cutoff=0.0,
+                                 gaussian=False,
+                                 return_best_logits=False):
+  """Compute cross-entropy assuming 0s are padding.
+
+  Computes a loss numerator (the sum of losses), and loss denominator
+  (the number of non-padding tokens).
+
+  Computes cross-entropy for each mixture, and returns the corresponding values
+  for the mixture with the highest probability
+
+  Args:
+    logits: `Tensor` with shape `[batch * num_mixtures, timesteps, vocab_size]`.
+      optionally a FactoredTensor.
+    labels: an integer `Tensor` with shape `[batch, timesteps]`.
+    label_smoothing: a floating point `Scalar`.
+    num_mixtures: an integer.
+    weights_fn: A function from labels to weights.
+    reduce_sum: a Boolean, whether to sum at the end or not.
+    cutoff: a float, at which point to have no loss.
+    gaussian: If true, use a Gaussian distribution for label smoothing
+    return_best_logits: If true, return the logits of the mixture with highest
+    probabilities for an example
+
+  Returns:
+    loss_numerator: a `Scalar`.  Sum of losses.
+    loss_denominator: a `Scalar.  The number of non-padding target tokens.
+
+  Raises:
+    ValueError: in case of unsupported argument types.
+  """
+  logit_shapes = shape_list(
+      logits)  # batch_size * num_mixtures, timesteps, 1, 1, vocab_size
+  batch_size = tf.cast(logit_shapes[0] / num_mixtures, dtype=tf.int32)
+  timesteps = logit_shapes[1]
+  vocab_size = logit_shapes[4]
+
+  new_shape_for_xent = [num_mixtures] + shape_list(labels)
+  labels = tf.tile(labels, [num_mixtures, 1, 1, 1])
+
+  xent, weights = padded_cross_entropy(
+      logits, labels, label_smoothing, weights_fn, reduce_sum, cutoff, gaussian)
+
+  # reshape xent and weights to have the num_mixtures as first dimension
+  xent = tf.reshape(xent, new_shape_for_xent)
+  weights = tf.reshape(weights, new_shape_for_xent[:-1])
+
+  # sum up sentence neg log probs
+  xent = tf.reduce_sum(xent, axis=2)
+
+  # if we need to compute the best logits
+  if return_best_logits:
+    best_mixture_indices = tf.cast(tf.argmin(xent, 0), dtype=tf.int32)
+    individual_element_indices = tf.range(batch_size)
+    stacked_mixture_element_indices = tf.stack(
+        (tf.squeeze(best_mixture_indices), individual_element_indices), -1)
+    best_logits = tf.reshape(logits,
+                             [num_mixtures, -1, timesteps, 1, 1, vocab_size])
+    best_logits = tf.gather_nd(best_logits, stacked_mixture_element_indices)
+    best_logits = tf.reshape(best_logits,
+                             [batch_size, timesteps, 1, 1, vocab_size])
+
+  with tf.control_dependencies([
+      tf.assert_equal(
+          tf.shape(xent)[:3], [num_mixtures, batch_size, 1],
+          message="Each batch element should have a probability value for each mixture element"
+      )
+  ]):
+    xent = tf.reduce_min(xent, axis=0)
+    weights = tf.reduce_mean(weights, axis=0)
+
+  with tf.control_dependencies([
+      tf.assert_equal(
+          tf.shape(xent)[0], [batch_size],
+          message="There should be batch_size elements after selecting best mixture probabilities"
+      )
+  ]):
+    summed_xent = tf.reduce_sum(xent)
+    summed_weights = tf.reduce_sum(weights)
+
+  if return_best_logits:
+    return summed_xent, summed_weights, best_logits
+  else:
+    return summed_xent, summed_weights
+
+
 def _weights_one_third(labels):
   """Returns Tensor of shape [batch, height, width]. Each element is 1/3."""
   return tf.ones(tf.shape(labels)[:-1]) / 3.
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 10bc5d7ab..ea94316cf 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -790,8 +790,6 @@ def finish(self):
     (out, err) = [
         b"".join(chunks) for chunks in (self._out_chunks, self._err_chunks)
     ]
-    self.proc.stdout.close()
-    self.proc.stderr.close()
     if self.proc.returncode:
       err = "\n".join([" ".join(self.cmd), err.decode("utf8")])
       raise IOError(err)
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 249e87a2c..90c74aeef 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -354,18 +354,12 @@ def dqn_atari_base():
       optimizer_epsilon=0.00001,
       optimizer_centered=True,
 
-      # TODO: change names maybe replay_buffer -> agent? Also batch_size is now
-      # buffer_batch_size in _DQNAgent.
       replay_buffer_replay_capacity=1000000,
-      replay_buffer_buffer_batch_size=32,
+      replay_buffer_batch_size=32,
 
       time_limit=27000,
       save_every_steps=50000,
       num_frames=int(20 * 1e6),
-
-      # TODO(konradczechowski) this is not used in trainer_model_free, clean
-      # this up after evaluation refactor
-      eval_episodes_num=3,
   )
 
 
@@ -376,16 +370,6 @@ def dqn_original_params():
   hparams.set_hparam("num_frames", int(1e6))
   return hparams
 
-def rlmf_tiny_overrides():
-  """Parameters to override for tiny setting excluding agent-related hparams."""
-  return dict(
-      max_num_noops=1,
-      eval_max_num_noops=1,
-      rl_env_max_episode_steps=7,
-      eval_rl_env_max_episode_steps=7,
-      eval_sampling_temps=[0.0, 1.0],
-  )
-
 
 @registry.register_hparams
 def rlmf_original():
@@ -398,7 +382,6 @@ def rlmf_original():
       eval_batch_size=2,
       frame_stack_size=4,
       eval_sampling_temps=[0.0, 0.2, 0.5, 0.8, 1.0, 2.0],
-      max_num_noops=8,
       eval_max_num_noops=8,
       eval_rl_env_max_episode_steps=1000,
       resize_height_factor=2,
@@ -443,31 +426,6 @@ def rlmf_base():
   return hparams
 
 
-@registry.register_hparams
-def rlmf_tiny():
-  """Tiny set of hparams for model-free PPO."""
-  hparams = rlmf_original()
-  hparams = hparams.override_from_dict(rlmf_tiny_overrides())
-  hparams.batch_size = 2
-  hparams.add_hparam("ppo_epochs_num", 3)
-  hparams.add_hparam("ppo_epoch_length", 2)
-  return hparams
-
-
-@registry.register_hparams
-def rlmf_dqn_tiny():
-  hparams = rlmf_original()
-  hparams = hparams.override_from_dict(rlmf_tiny_overrides())
-  hparams.batch_size = 1
-  hparams.base_algo = "dqn"
-  hparams.base_algo_params = "dqn_original_params"
-  hparams.add_hparam("dqn_num_frames", 128)
-  hparams.add_hparam("dqn_save_every_steps", 128)
-  hparams.add_hparam("dqn_replay_buffer_replay_capacity", 100)
-  hparams.add_hparam("dqn_agent_min_replay_history", 10)
-  return hparams
-
-
 @registry.register_hparams
 def rlmf_eval():
   """Eval set of hparams for model-free PPO."""
@@ -484,6 +442,14 @@ def rlmf_eval():
   return hparams
 
 
+@registry.register_hparams
+def rlmf_tiny():
+  hparams = rlmf_base()
+  hparams.ppo_epochs_num = 100
+  hparams.ppo_eval_every_epochs = 10
+  return hparams
+
+
 class PolicyBase(t2t_model.T2TModel):
 
   def loss(self, *args, **kwargs):
diff --git a/tensor2tensor/rl/batch_dqn_agent_test.py b/tensor2tensor/rl/batch_dqn_agent_test.py
deleted file mode 100644
index 6e70c8e54..000000000
--- a/tensor2tensor/rl/batch_dqn_agent_test.py
+++ /dev/null
@@ -1,157 +0,0 @@
-# Copyright 2018 The Dopamine Authors.
-#
-# 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 BatchDQNAgent."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-import shutil
-
-from absl import flags
-from dopamine.agents.dqn import dqn_agent
-import numpy as np
-
-from tensor2tensor.rl import dopamine_connector
-
-import tensorflow as tf
-
-
-slim = tf.contrib.slim
-
-FLAGS = flags.FLAGS
-
-
-class BatchDQNAgentTest(tf.test.TestCase):
-  # TODO: maybe add testStepTrain (and possibly some other tests) from dopamine
-  # dqn_agent_test.py
-
-  def setUp(self):
-    self._test_subdir = os.path.join('/tmp/dopamine_tests', 'ckpts')
-    shutil.rmtree(self._test_subdir, ignore_errors=True)
-    os.makedirs(self._test_subdir)
-    self.num_actions = 4
-    self.min_replay_history = 6
-    self.update_period = 2
-    self.target_update_period = 4
-    self.epsilon_decay_period = 90
-    self.epsilon_train = 0.05
-    self.observation_shape = dqn_agent.NATURE_DQN_OBSERVATION_SHAPE
-    self.stack_size = dqn_agent.NATURE_DQN_STACK_SIZE
-    self.env_batch_size = 4
-
-    self.zero_state = np.zeros(
-        [self.env_batch_size, self.observation_shape[0],
-         self.observation_shape[1], self.stack_size])
-
-
-  def _create_test_agent(self, sess):
-    stack_size = self.stack_size
-
-    class MockDQNAgent(dopamine_connector.BatchDQNAgent):
-
-      def _network_template(self, state):
-        # This dummy network allows us to deterministically anticipate that
-        # action 0 will be selected by an argmax.
-        inputs = tf.constant(
-            np.zeros((state.shape[0], stack_size)), dtype=tf.float32)
-        # This weights_initializer gives action 0 a higher weight, ensuring
-        # that it gets picked by the argmax.
-        weights_initializer = np.tile(
-            np.arange(self.num_actions, 0, -1), (stack_size, 1))
-        q = slim.fully_connected(
-            inputs,
-            self.num_actions,
-            weights_initializer=tf.constant_initializer(weights_initializer),
-            biases_initializer=tf.ones_initializer(),
-            activation_fn=None)
-        return self._get_network_type()(q)
-
-    agent = MockDQNAgent(
-        replay_capacity=100,
-        buffer_batch_size=8,
-        generates_trainable_dones=True,
-        sess=sess,
-        env_batch_size=self.env_batch_size,
-        num_actions=self.num_actions,
-        min_replay_history=self.min_replay_history,
-        epsilon_fn=lambda w, x, y, z: 0.0,  # No exploration.
-        update_period=self.update_period,
-        target_update_period=self.target_update_period,
-        epsilon_eval=0.0)  # No exploration during evaluation.
-    # This ensures non-random action choices (since epsilon_eval = 0.0) and
-    # skips the train_step.
-    agent.eval_mode = True
-    sess.run(tf.global_variables_initializer())
-    return agent
-
-  def testCreateAgentWithDefaults(self):
-    # Verifies that we can create and train an agent with the default values.
-    with tf.Session() as sess:
-      agent = self._create_test_agent(sess)
-      sess.run(tf.global_variables_initializer())
-      observation = np.ones([84, 84, 1])
-      agent.begin_episode([observation])
-      agent.step(reward=[1], observation=[observation])
-      agent.end_episode(reward=[1])
-
-  def testBeginEpisode(self):
-    """Test the functionality of agent.begin_episode.
-
-    Specifically, the action returned and its effect on state.
-    """
-    with tf.Session() as sess:
-      agent = self._create_test_agent(sess)
-      # We fill up the state with 9s. On calling agent.begin_episode the state
-      # should be reset to all 0s.
-      agent.state_batch.fill(9)
-      first_observation = np.ones(
-          [self.env_batch_size, self.observation_shape[0],
-           self.observation_shape[1], 1])
-      self.assertTrue((agent.begin_episode(first_observation) == 0).all())
-      # When the all-1s observation is received, it will be placed at the end of
-      # the state.
-      expected_state = self.zero_state
-      expected_state[:, :, :, -1] = np.ones(
-          [self.env_batch_size, self.observation_shape[0],
-           self.observation_shape[1]])
-      self.assertAllEqual(agent.state_batch, expected_state)
-      self.assertAllEqual(agent._observation_batch, first_observation[..., 0])
-      # No training happens in eval mode.
-      self.assertEqual(agent.training_steps, 0)
-
-      # This will now cause training to happen.
-      agent.eval_mode = False
-      # Having a low replay memory add_count will prevent any of the
-      # train/prefetch/sync ops from being called.
-      agent._replay.memory.add_count = 0
-      second_observation = np.ones(
-          [self.env_batch_size, self.observation_shape[0],
-           self.observation_shape[1], 1]) * 2
-      agent.begin_episode(second_observation)
-      # The agent's state will be reset, so we will only be left with the all-2s
-      # observation.
-      expected_state[:, :, :, -1] = np.full(
-          (self.env_batch_size, self.observation_shape[0],
-           self.observation_shape[1]), 2
-      )
-      self.assertAllEqual(agent.state_batch, expected_state)
-      self.assertAllEqual(agent._observation_batch,
-                          second_observation[:, :, :, 0])
-      # training_steps is incremented since we set eval_mode to False.
-      self.assertEqual(agent.training_steps, 1)
-
-
-if __name__ == '__main__':
-  tf.test.main()
diff --git a/tensor2tensor/rl/batch_runner_test.py b/tensor2tensor/rl/batch_runner_test.py
deleted file mode 100644
index 86d1d6645..000000000
--- a/tensor2tensor/rl/batch_runner_test.py
+++ /dev/null
@@ -1,297 +0,0 @@
-# Copyright 2018 The Dopamine Authors.
-#
-# 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 BatchRunner."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-import shutil
-
-from absl import flags
-from dopamine.discrete_domains import run_experiment
-from dopamine.discrete_domains import logger
-import gin.tf
-import numpy as np
-import mock
-
-from tensor2tensor.rl import dopamine_connector
-
-import tensorflow as tf
-
-
-FLAGS = flags.FLAGS
-
-
-def _create_mock_checkpointer():
-  mock_checkpointer = mock.Mock()
-  test_dictionary = {'current_iteration': 1729,
-                     'logs': 'logs'}
-  mock_checkpointer.load_checkpoint.return_value = test_dictionary
-  return mock_checkpointer
-
-
-class MockEnvironment(object):
-  """Mock environment for testing."""
-
-  def __init__(self, max_steps=10, reward_multiplier=1):
-    self._observation = 0
-    self.max_steps = max_steps
-    self.reward_multiplier = reward_multiplier
-    self.game_over = False
-
-  def reset(self):
-    self._observation = 0
-    return self._observation
-
-  def step(self, action):
-    self._observation += 1
-    action_reward_multiplier = -1 if action > 0 else 1
-    reward_multiplier = self.reward_multiplier * action_reward_multiplier
-    reward = self._observation * reward_multiplier
-    is_terminal = self._observation >= self.max_steps
-    self.game_over = is_terminal
-
-    unused = 0
-    return (self._observation, reward, is_terminal, unused)
-
-  def render(self, mode):
-    pass
-
-
-class BatchEnv(object):
-  """
-
-  Batch of environments. Assumes that all throws 'done' on the same step.
-
-  Observations and rewards are returned as arrays, done as single value.
-  """
-  # TODO: this can be used for mbrl pipeline (for both simulated and real env),
-  #  move it to dopamine_connector.py (rename it?)
-  def __init__(self, envs):
-    self.env_batch = envs
-    self.batch_size = len(self.env_batch)
-    self.max_steps = self.env_batch[0].max_steps
-    assert np.all(self.max_steps == env.max_steps for env in self.env_batch)
-
-  def step(self, actions):
-    ret = [env.step(action) for env, action in zip(self.env_batch, actions)]
-    obs, rewards, dones, infos = [np.array(r) for r in zip(*ret)]
-    done = dones[0]
-    assert np.all(done == dones)
-    self.game_over = done
-    return obs, rewards, done, infos
-
-  def reset(self):
-    return np.array([env.reset() for env in self.env_batch])
-
-  def render(self, mode):
-    pass
-
-
-class MockLogger(object):
-  """Class to mock the experiment logger."""
-
-  def __init__(self, test_cls=None, run_asserts=True, data=None):
-    self._test_cls = test_cls
-    self._run_asserts = run_asserts
-    self._iter = 0
-    self._calls_to_set = 0
-    self._calls_to_log = 0
-    self.data = data
-
-  def __setitem__(self, key, val):
-    if self._run_asserts:
-      self._test_cls.assertEqual('iteration_{:d}'.format(self._iter), key)
-      self._test_cls.assertEqual('statistics', val)
-      self._iter += 1
-    self._calls_to_set += 1
-
-  def log_to_file(self, filename_prefix, iteration_number):
-    if self._run_asserts:
-      self._test_cls.assertEqual(
-          'prefix_{}'.format(self._iter - 1),
-          '{}_{}'.format(filename_prefix, iteration_number))
-    self._calls_to_log += 1
-
-
-class RunExperimentTest(tf.test.TestCase):
-
-  @mock.patch.object(gin, 'parse_config_files_and_bindings')
-  def testLoadGinConfigs(self, mock_parse_config_files_and_bindings):
-    gin_files = ['file1', 'file2', 'file3']
-    gin_bindings = ['binding1', 'binding2']
-    run_experiment.load_gin_configs(gin_files, gin_bindings)
-    self.assertEqual(1, mock_parse_config_files_and_bindings.call_count)
-    mock_args, mock_kwargs = mock_parse_config_files_and_bindings.call_args
-    self.assertEqual(gin_files, mock_args[0])
-    self.assertEqual(gin_bindings, mock_kwargs['bindings'])
-    self.assertFalse(mock_kwargs['skip_unknown'])
-
-
-class BatchedRunnerTest(tf.test.TestCase):
-  """Modified tests from dopamine run_experiment_test.py."""
-
-  # TODO: decide if we want to use and modify more tests from
-  # dopamine/tests/atari/run_experiment_test.py (e.g.  testRunExperiment.py)
-
-  def _agent_step(self, rewards, observations):
-    # We verify that rewards are clipped (and set by MockEnvironment as a
-    # function of observation)
-    # observation = observations[0]
-    # expected_rewards = [1 if observation % 2 else -1]
-    # self.assertEqual(expected_reward, reward)
-    actions = [ob % 2 for ob in observations]
-    return actions
-
-  def prepare_mock_agent(self, batch_size):
-    assert batch_size % 2 == 0, "Some of tests assume that batch_size % 2 == 0"
-    self.batch_size = batch_size
-    self._agent = mock.Mock()
-    self._agent.begin_episode.side_effect = \
-      lambda x: np.repeat(0, self.batch_size)
-    self._agent.step.side_effect = self._agent_step
-    self._create_agent_fn = lambda x, y, summary_writer: self._agent
-
-  def setUp(self):
-    super(BatchedRunnerTest, self).setUp()
-    self._test_subdir = "/tmp/dopamine_tests"
-    shutil.rmtree(self._test_subdir, ignore_errors=True)
-    os.makedirs(self._test_subdir)
-    self.prepare_mock_agent(batch_size=4)
-
-  def testRunEpisodeBatch(self):
-    max_steps_per_episode = 11
-    batch_size = self.batch_size
-    reward_multipliers = [-1, 1] * int(batch_size / 2)
-    envs = [MockEnvironment(reward_multiplier=rm) for rm in reward_multipliers]
-    environment = BatchEnv(envs)
-    runner = dopamine_connector.BatchRunner(
-        self._test_subdir, self._create_agent_fn,
-        create_environment_fn=lambda: environment,
-        max_steps_per_episode=max_steps_per_episode)
-    step_number, total_rewards = runner._run_one_episode()
-
-    self.assertEqual(self._agent.step.call_count, environment.max_steps - 1)
-    self.assertEqual(self._agent.end_episode.call_count, 1)
-    self.assertEqual(environment.max_steps, step_number / batch_size)
-    # Expected reward will be \sum_{i=0}^{9} (-1)**i * i = -5 when reward
-    # multiplier=1
-    self.assertAllEqual(np.array(reward_multipliers) * -5, total_rewards)
-
-  def testRunOneEpisodeWithLowMaxSteps(self):
-    max_steps_per_episode = 2
-    batch_size = self.batch_size
-    reward_multipliers = [-1, 1] * int(batch_size / 2)
-    envs = [MockEnvironment(reward_multiplier=rm) for rm in reward_multipliers]
-    environment = BatchEnv(envs)
-    runner = dopamine_connector.BatchRunner(
-        self._test_subdir, self._create_agent_fn,
-        create_environment_fn=lambda: environment,
-        max_steps_per_episode=max_steps_per_episode)
-    step_number, total_rewards = runner._run_one_episode()
-
-    self.assertEqual(self._agent.step.call_count, max_steps_per_episode - 1)
-    self.assertEqual(self._agent.end_episode.call_count, 1)
-    self.assertEqual(max_steps_per_episode, step_number / batch_size)
-    self.assertAllEqual(np.array(reward_multipliers) * -1, total_rewards)
-
-  def testRunOnePhase(self):
-    batch_size = self.batch_size
-    environment_steps = 2
-    max_steps = environment_steps * batch_size * 10
-
-    envs = [MockEnvironment(max_steps=environment_steps)
-            for _ in range(batch_size)]
-
-    environment = BatchEnv(envs)
-    runner = dopamine_connector.BatchRunner(
-        self._test_subdir, self._create_agent_fn,
-        create_environment_fn=lambda: environment)
-
-    statistics = []
-
-    step_number, sum_returns, num_episodes = runner._run_one_phase(
-        max_steps, statistics, "test")
-    calls_to_run_episode = int(max_steps / (environment_steps * batch_size))
-    self.assertEqual(self._agent.step.call_count, calls_to_run_episode)
-    self.assertEqual(self._agent.end_episode.call_count, calls_to_run_episode)
-    self.assertEqual(max_steps, step_number)
-    self.assertEqual(-1 * calls_to_run_episode * batch_size, sum_returns)
-    self.assertEqual(calls_to_run_episode, num_episodes / batch_size)
-    expected_statistics = []
-    for _ in range(calls_to_run_episode * batch_size):
-      expected_statistics.append({
-          "test_episode_lengths": 2,
-          "test_episode_returns": -1
-      })
-    self.assertEqual(len(expected_statistics), len(statistics))
-    for expected_stats, stats in zip(expected_statistics, statistics):
-      self.assertDictEqual(expected_stats, stats)
-
-  def testRunOneIteration(self):
-    environment_steps = 2
-    batch_size = self.batch_size
-    envs = [MockEnvironment(max_steps=environment_steps)
-            for _ in range(batch_size)]
-
-    environment = BatchEnv(envs)
-
-    training_steps = 20 * batch_size
-    evaluation_steps = 10 * batch_size
-
-    runner = dopamine_connector.BatchRunner(
-        self._test_subdir, self._create_agent_fn,
-        create_environment_fn=lambda: environment,
-        training_steps=training_steps, evaluation_steps=evaluation_steps
-    )
-
-    dictionary = runner._run_one_iteration(1)
-    train_rollouts = int(training_steps / environment_steps)
-    eval_rollouts = int(evaluation_steps / environment_steps)
-    expected_dictionary = {
-        "train_episode_lengths": [2 for _ in range(train_rollouts)],
-        "train_episode_returns": [-1 for _ in range(train_rollouts)],
-        "train_average_return": [-1],
-        "eval_episode_lengths": [2 for _ in range(eval_rollouts)],
-        "eval_episode_returns": [-1 for _ in range(eval_rollouts)],
-        "eval_average_return": [-1]
-    }
-    self.assertDictEqual(expected_dictionary, dictionary)
-
-  @mock.patch.object(logger, "Logger")
-  def testLogExperiment(self, mock_logger_constructor):
-    #TODO: We probably do not need this test, dopamine test for Runner is enugh
-    # here. Remove this?
-    log_every_n = 2
-    logging_file_prefix = "prefix"
-    statistics = "statistics"
-    experiment_logger = MockLogger(test_cls=self)
-    mock_logger_constructor.return_value = experiment_logger
-    runner = dopamine_connector.BatchRunner(
-        self._test_subdir, self._create_agent_fn,
-        create_environment_fn=mock.Mock,
-        logging_file_prefix=logging_file_prefix,
-        log_every_n=log_every_n)
-    num_iterations = 10
-    for i in range(num_iterations):
-      runner._log_experiment(i, statistics)
-    self.assertEqual(num_iterations, experiment_logger._calls_to_set)
-    self.assertEqual((num_iterations / log_every_n),
-                     experiment_logger._calls_to_log)
-
-
-if __name__ == '__main__':
-  tf.test.main()
diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index c9dab6bbb..9e1ad3765 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -20,8 +20,6 @@
 from __future__ import print_function
 
 import copy
-import random
-import sys
 
 from dopamine.agents.dqn import dqn_agent
 from dopamine.replay_memory import circular_replay_buffer
@@ -48,6 +46,52 @@
 # pylint: enable=g-import-not-at-top
 
 
+class ResizeObservation(gym.ObservationWrapper):
+  """TODO(konradczechowski): Add doc-string."""
+
+  def __init__(self, env, size=84):
+    """Based on WarpFrame from openai baselines atari_wrappers.py.
+
+    Dopamine also uses cv2.resize(..., interpolation=cv2.INTER_AREA).
+
+    Args:
+      env: TODO(konradczechowski): Add doc-string.
+      size: TODO(konradczechowski): Add doc-string.
+    """
+    gym.ObservationWrapper.__init__(self, env)
+    self.width = size
+    self.height = size
+    assert env.observation_space.dtype == np.uint8
+    self.observation_space = spaces.Box(
+        low=0,
+        high=255,
+        shape=(self.height, self.width, env.observation_space.shape[2]),
+        dtype=np.uint8)
+
+  def observation(self, frame):
+    if not cv2:
+      return frame
+    return cv2.resize(
+        frame, (self.width, self.height), interpolation=cv2.INTER_AREA)
+
+
+class GameOverOnDone(Wrapper):
+  """TODO(konradczechowski): Add doc-string."""
+
+  def __init__(self, env):
+    Wrapper.__init__(self, env)
+    self.game_over = False
+
+  def reset(self, **kwargs):
+    self.game_over = False
+    return self.env.reset(**kwargs)
+
+  def step(self, action):
+    ob, reward, done, info = self.env.step(action)
+    self.game_over = done
+    return ob, reward, done, info
+
+
 class _DQNAgent(dqn_agent.DQNAgent):
   """Modify dopamine DQNAgent to match our needs.
 
@@ -55,10 +99,10 @@ class _DQNAgent(dqn_agent.DQNAgent):
   (some of) terminal episode transitions in training.
   """
 
-  def __init__(self, replay_capacity, buffer_batch_size,
-               generates_trainable_dones, **kwargs):
+  def __init__(self, replay_capacity, batch_size, generates_trainable_dones,
+               **kwargs):
     self._replay_capacity = replay_capacity
-    self._buffer_batch_size = buffer_batch_size
+    self._batch_size = batch_size
     self._generates_trainable_dones = generates_trainable_dones
     super(_DQNAgent, self).__init__(**kwargs)
 
@@ -68,7 +112,7 @@ def _build_replay_buffer(self, use_staging):
         observation_shape=dqn_agent.NATURE_DQN_OBSERVATION_SHAPE,
         stack_size=dqn_agent.NATURE_DQN_STACK_SIZE,
         replay_capacity=self._replay_capacity,
-        batch_size=self._buffer_batch_size,
+        batch_size=self._batch_size,
         update_horizon=self.update_horizon,
         gamma=self.gamma,
         extra_storage_types=None,
@@ -84,145 +128,6 @@ def _build_replay_buffer(self, use_staging):
         **replay_buffer_kwargs)
 
 
-class BatchDQNAgent(_DQNAgent):
-  """
-  Episodes are stored on done.
-
-  Assumes that all rollouts in batch would end at the same moment.
-  """
-
-  def __init__(self, env_batch_size, *args, **kwargs):
-    super(BatchDQNAgent, self).__init__(*args, **kwargs)
-    self.env_batch_size = env_batch_size
-    obs_size = dqn_agent.NATURE_DQN_OBSERVATION_SHAPE
-    state_shape = [self.env_batch_size, obs_size[0], obs_size[1],
-                   dqn_agent.NATURE_DQN_STACK_SIZE]
-    self.state_batch = np.zeros(state_shape)
-    self.state = None  # assure it will be not used
-    self._observation = None  # assure it will be not used
-    self.reset_current_rollouts()
-
-  def reset_current_rollouts(self):
-    self._current_rollouts = [[] for _ in range(self.env_batch_size)]
-
-  def _record_observation(self, observation_batch):
-    # Set current observation. Represents an (batch_size x 84 x 84 x 1) image
-    # frame.
-    observation_batch = np.array(observation_batch)
-    self._observation_batch = observation_batch[:, :, :, 0]
-    # Swap out the oldest frames with the current frames.
-    self.state_batch = np.roll(self.state_batch, -1, axis=3)
-    self.state_batch[:, :, :, -1] = self._observation_batch
-
-  def _reset_state(self):
-    self.state_batch.fill(0)
-
-  def begin_episode(self, observation):
-    self._reset_state()
-    self._record_observation(observation)
-
-    if not self.eval_mode:
-      self._train_step()
-
-    self.action = self._select_action()
-    return self.action
-
-  def _update_current_rollouts(self, last_observation, action, reward,
-                              are_terminal):
-    transitions = zip(last_observation, action, reward, are_terminal)
-    for transition, rollout in zip(transitions, self._current_rollouts):
-      rollout.append(transition)
-
-  def _store_current_rollouts(self):
-    for rollout in self._current_rollouts:
-      for transition in rollout:
-        self._store_transition(*transition)
-    self.reset_current_rollouts()
-
-  def step(self, reward, observation):
-    self._last_observation = self._observation_batch
-    self._record_observation(observation)
-
-    if not self.eval_mode:
-      self._update_current_rollouts(self._last_observation, self.action, reward,
-                                    [False] * self.env_batch_size)
-      # We want to have the same train_step:env_step ratio not depending on
-      # batch size.
-      for _ in range(self.env_batch_size):
-        self._train_step()
-
-    self.action = self._select_action()
-    return self.action
-
-  def end_episode(self, reward):
-    if not self.eval_mode:
-      self._update_current_rollouts(self._observation_batch, self.action, reward,
-                                    [True] * self.env_batch_size)
-      self._store_current_rollouts()
-
-  def _select_action(self):
-    epsilon = self.epsilon_eval if self.eval_mode else self.epsilon_fn(
-        self.epsilon_decay_period,
-        self.training_steps,
-        self.min_replay_history,
-        self.epsilon_train)
-
-    def choose_action(ix):
-      if random.random() <= epsilon:
-        # Choose a random action with probability epsilon.
-        return random.randint(0, self.num_actions - 1)
-      else:
-        # Choose the action with highest Q-value at the current state.
-        return self._sess.run(self._q_argmax,
-                              {self.state_ph: self.state_batch[ix:ix+1]})
-
-    return np.array([choose_action(ix) for ix in range(self.env_batch_size)])
-
-
-class BatchRunner(run_experiment.Runner):
-  """
-
-  Assumes that all environments would end at the same moment.
-  """
-  def __init__(self, base_dir, create_agent_fn, **kwargs):
-    super(BatchRunner, self).__init__(base_dir, create_agent_fn, **kwargs)
-    self.batch_size = self._environment.batch_size
-
-  def _run_one_episode(self):
-    # This assumes that everything inside _run_one_episode works on batches,
-    # which is risky for future.
-    steps_number, total_rewards = super(BatchRunner, self)._run_one_episode()
-    return steps_number * self.batch_size, total_rewards
-
-  def _run_one_phase(self, min_steps, statistics, run_mode_str):
-    # Mostly copy of parent method.
-    step_count = 0
-    num_episodes = 0
-    sum_returns = 0.
-
-    while step_count < min_steps:
-      num_steps, episode_returns = self._run_one_episode()
-      for episode_return in episode_returns:
-        statistics.append({
-            '{}_episode_lengths'.format(run_mode_str):
-                num_steps / self.batch_size,
-            '{}_episode_returns'.format(run_mode_str): episode_return
-        })
-      step_count += num_steps
-      sum_returns += sum(episode_returns)
-      num_episodes += self.batch_size
-      # We use sys.stdout.write instead of tf.logging so as to flush frequently
-      # without generating a line break.
-      sys.stdout.write('Steps executed: {} '.format(step_count) +
-                       'Batch episodes steps: {} '.format(num_steps) +
-                       'Returns: {}\r'.format(episode_returns))
-      sys.stdout.flush()
-    return step_count, sum_returns, num_episodes
-
-  def close(self):
-    self._environment.close()
-
-
 class _OutOfGraphReplayBuffer(OutOfGraphReplayBuffer):
   """Replay not sampling artificial_terminal transition.
 
@@ -272,14 +177,7 @@ def load(self, *args, **kwargs):
 
 
 def get_create_agent(agent_kwargs):
-  """Factory for dopamine agent initialization.
-
-  Args:
-    agent_kwargs: dict of BatchDQNAgent parameters
-
-  Returns:
-    Function(sess, environment, summary_writer) -> BatchDQNAgent instance.
-  """
+  """TODO(): Document."""
 
   def create_agent(sess, environment, summary_writer=None):
     """Creates a DQN agent.
@@ -294,8 +192,7 @@ def create_agent(sess, environment, summary_writer=None):
     Returns:
       a DQN agent.
     """
-    return BatchDQNAgent(
-        env_batch_size=environment.batch_size,
+    return _DQNAgent(
         sess=sess,
         num_actions=environment.action_space.n,
         summary_writer=summary_writer,
@@ -305,171 +202,23 @@ def create_agent(sess, environment, summary_writer=None):
   return create_agent
 
 
-class ResizeBatchObservation(object):
-  """Wrapper resizing observations for batched environment.
-
-  Dopamine also uses cv2.resize(..., interpolation=cv2.INTER_AREA).
-
-  Attributes:
-    batch_env: batched environment
-    size: size of width and height for returned observations
-  """
-
-  def __init__(self, batch_env, size=84):
-    self.size = size
-    self.batch_env = batch_env
-
-  def observation(self, frames):
-    if not cv2:
-      return frames
-    return np.array([cv2.resize(
-        frame, (self.size, self.size), interpolation=cv2.INTER_AREA)
-        for frame in frames])
-
-  def step(self, actions):
-    obs, rewards, dones = self.batch_env.step(actions)
-    obs = self.observation(obs)
-    return obs, rewards, dones
-
-  def reset(self, *args, **kwargs):
-    return self.observation(self.batch_env.reset(*args, **kwargs))
+def get_create_env_fun(batch_env_fn, time_limit):
+  """TODO(konradczechowski): Add doc-string."""
 
-  @property
-  def action_space(self):
-    return self.batch_env.action_space
-
-  @property
-  def batch_size(self):
-    return self.batch_env.batch_size
-
-  def close(self):
-    self.batch_env.close()
-
-
-class DopamineBatchEnv(object):
-  """Batch of environments.
-
-  Assumes that all given environments finishes at the same time.
-
-  Observations and rewards are returned as batches (arrays). Done is returned
-  as single boolean.
-  """
-  def __init__(self, batch_env, max_episode_steps):
-    self.batch_env = batch_env
-    self._max_episode_steps = max_episode_steps
-    self.game_over = None
-    self._elapsed_steps = 0
-
-  def reset(self):
-    self.game_over = False
-    self._elapsed_steps = 0
-    return np.array(self.batch_env.reset())
-
-  def step(self, actions):
-    self._elapsed_steps += 1
-    obs, rewards, dones = \
-        [np.array(r) for r in self.batch_env.step(actions)]
-    if self._elapsed_steps > self._max_episode_steps:
-      done = True
-      if self._elapsed_steps > self._max_episode_steps + 1:
-        rewards.fill(0)
-    else:
-      done = dones[0]
-      assert np.all(done == dones), "Current modifications of Dopamine " \
-                                    "require same number of steps for each " \
-                                    "environment in batch"
-      del dones
-
-    self.game_over = done
-    return obs, rewards, done, {}
-
-  def render(self, mode):
-    pass
-
-  def close(self):
-    self.batch_env.close()
-
-  @property
-  def action_space(self):
-    return self.batch_env.action_space
-
-  @property
-  def batch_size(self):
-    return self.batch_env.batch_size
-
-
-class PaddedTrajectoriesEnv(DopamineBatchEnv):
-  """ Padd finished episodes with zeros.
-
-  Allow episodes in batch to end on different timesteps, return zero
-  observations and rewards for finished ones. Return done=True when all
-  episodes are finished.
-
-  Note that output of this class might be misleading - the agent/evaluator
-  which uses this environment gets false information about when episodes have
-  ended. This class is used for informal check of Batched dopamine
-  implementation in model-free pipeline.
-  """
-
-  def reset(self):
-    self.done_envs = [False] * self.batch_size
-    self.game_over = False
-    self._elapsed_steps = 0
-    return np.array(self.batch_env.reset())
-
-  def step(self, actions):
-    if any(self.done_envs):
-      print("Warning, some environments already ended, using mocked data.")
-
-    self._elapsed_steps += 1
-    obs, rewards, dones = \
-        [np.array(r) for r in self.batch_env.step(actions)]
-    for i, ignore in enumerate(self.done_envs):
-      if ignore:
-        obs[i] = np.zeros(obs[i].shape, dtype=obs.dtype)
-        rewards[i] = 0
-      if dones[i]:
-        self.batch_env.reset([i])
-        self.done_envs[i] = True
-
-    all_done = all(self.done_envs)
-
-    if self._elapsed_steps > self._max_episode_steps:
-      all_done = True
-      if self._elapsed_steps > self._max_episode_steps + 1:
-        rewards.fill(0)
-
-    self.game_over = all_done
-    return obs, rewards, all_done, {}
-
-
-def get_create_batch_env_fun(batch_env_fn, time_limit):
-  """Factory for dopamine environment initialization function.
-
-  Args:
-    batch_env_fn: function(in_graph: bool) -> batch environment.
-    time_limit: time steps limit for environment.
-
-  Returns:
-    function (with optional, unused parameters) initializing environment.
-  """
-
-  def create_env_fun(game_name=None, sticky_actions=None):
+  def create_env_fun(game_name, sticky_actions=True):
     del game_name, sticky_actions
     batch_env = batch_env_fn(in_graph=False)
-    batch_env = ResizeBatchObservation(batch_env)  # pylint: disable=redefined-variable-type
-    batch_env = DopamineBatchEnv(batch_env, max_episode_steps=time_limit)
-    return batch_env
+    env = FlatBatchEnv(batch_env)
+    env = TimeLimit(env, max_episode_steps=time_limit)
+    env = ResizeObservation(env)  # pylint: disable=redefined-variable-type
+    env = GameOverOnDone(env)
+    return env
 
   return create_env_fun
 
 
 def _parse_hparams(hparams):
-  """Split hparams, based on key prefixes.
-
-  Returns:
-    Tuple of hparams for respectably: agent, optimizer, runner, replay_buffer.
-  """
+  """TODO(konradczechowski): Add doc-string."""
   prefixes = ["agent_", "optimizer_", "runner_", "replay_buffer_"]
   ret = []
 
@@ -493,8 +242,9 @@ def _get_optimizer(params):
 class DQNLearner(PolicyLearner):
   """Interface for learning dqn implemented in dopamine."""
 
-  def __init__(self, *args, **kwargs):
-    super(DQNLearner, self).__init__(*args, **kwargs)
+  def __init__(self, frame_stack_size, base_event_dir, agent_model_dir):
+    super(DQNLearner, self).__init__(frame_stack_size, base_event_dir,
+                                     agent_model_dir)
     self.completed_iterations = 0
 
   def _target_iteractions_and_steps(self, num_env_steps, save_continuously,
@@ -519,10 +269,10 @@ def create_runner(self, env_fn, hparams, target_iterations,
     agent_params["optimizer"] = optimizer
     agent_params.update(replay_buffer_params)
     create_agent_fn = get_create_agent(agent_params)
-    runner = BatchRunner(
+    runner = run_experiment.Runner(
         base_dir=self.agent_model_dir,
         create_agent_fn=create_agent_fn,
-        create_environment_fn=get_create_batch_env_fun(
+        create_environment_fn=get_create_env_fun(
             env_fn, time_limit=hparams.time_limit),
         evaluation_steps=0,
         num_iterations=target_iterations,
@@ -540,10 +290,9 @@ def train(self,
             num_env_steps=None,
             env_step_multiplier=1,
             eval_env_fn=None,
-            report_fn=None,
-            model_save_fn=None):
+            report_fn=None):
     # TODO(konradczechowski): evaluation during training (with eval_env_fun)
-    del epoch, eval_env_fn, simulated, report_fn, model_save_fn
+    del epoch, eval_env_fn, simulated, report_fn
     if num_env_steps is None:
       num_env_steps = hparams.num_frames
 
@@ -556,13 +305,13 @@ def train(self,
       self._target_iteractions_and_steps(
           num_env_steps=num_env_steps * env_step_multiplier,
           save_continuously=save_continuously,
-          save_every_steps=hparams.save_every_steps)
+          save_every_steps=hparams.save_every_steps,)
 
     with tf.Graph().as_default():
       runner = self.create_runner(env_fn, hparams, target_iterations,
                                   training_steps_per_iteration)
       runner.run_experiment()
-      runner.close()
+
     self.completed_iterations = target_iterations
 
   def evaluate(self, env_fn, hparams, sampling_temp):
@@ -574,7 +323,7 @@ def evaluate(self, env_fn, hparams, sampling_temp):
         "agent_epsilon_eval", min(hparams.agent_epsilon_eval * sampling_temp, 1)
     )
 
-    create_environment_fn = get_create_batch_env_fun(
+    create_environment_fn = get_create_env_fun(
         env_fn, time_limit=hparams.time_limit)
     env = create_environment_fn(
         game_name="unused_arg", sticky_actions="unused_arg")
@@ -582,13 +331,13 @@ def evaluate(self, env_fn, hparams, sampling_temp):
     with tf.Graph().as_default():
       runner = self.create_runner(env_fn, hparams, target_iterations,
                                   training_steps_per_iteration)
-      assert runner.batch_size == 1
       agent = runner._agent  # pylint: disable=protected-access
-      runner.close()
       del runner
       agent.eval = True
 
-      for _ in range(hparams.eval_episodes_num):
+      # TODO(konradczechowski): correct number of episodes, when this will
+      # be hparam
+      for _ in range(30):
         # Run single episode
         ob = env.reset()
         action = agent.begin_episode(ob)
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index bfee48de0..eac13e012 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -294,5 +294,6 @@ def _video_reset_writer(self):
       self._video_writer.finish_to_disk()
     self._video_writer = None
 
-  def close(self):
+  def __del__(self):
     self._video_reset_writer()
+    super(SimulatedBatchEnv, self).__del__()
diff --git a/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index fd220837f..e0ffd9805 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -86,6 +86,9 @@ def reset(self, indices=None):
     if indices is None:
       indices = np.array(range(self.batch_size))
     obs = self._sess.run(self._reset_op, feed_dict={self._indices_t: indices})
+    # TODO(pmilos): remove if possible
+    # obs[:, 0, 0, 0] = 0
+    # obs[:, 0, 0, 1] = 255
     return obs
 
   def step(self, actions):
@@ -96,4 +99,3 @@ def step(self, actions):
 
   def close(self):
     self._sess.close()
-    self._batch_env.close()
diff --git a/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
index ec59e4c8a..82051f159 100644
--- a/tensor2tensor/rl/policy_learner.py
+++ b/tensor2tensor/rl/policy_learner.py
@@ -45,6 +45,8 @@ def train(
       report_fn=None
   ):
     """Train."""
+    # TODO(konradczechowski): pass name_scope instead of epoch?
+    # TODO(konradczechowski): move 'simulated' to  batch_env
     raise NotImplementedError()
 
   def evaluate(self, env_fn, hparams, sampling_temp):
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 625c42916..f34303d92 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -39,7 +39,6 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl import trainer_model_based_params
-from tensor2tensor.rl.dopamine_connector import DQNLearner
 from tensor2tensor.rl.restarter import Restarter
 from tensor2tensor.utils import trainer_lib
 
@@ -288,7 +287,6 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   metrics = {}
 
   # Collect data from the real environment.
-  policy_model_dir = directories["policy"]
   tf.logging.info("Initial training of the policy in real environment.")
   train_agent_real_env(env, learner, hparams, epoch)
   metrics["mean_reward/train/clipped"] = rl_utils.compute_mean_reward(
diff --git a/tensor2tensor/rl/trainer_model_based_dqn_test.py b/tensor2tensor/rl/trainer_model_based_dqn_test.py
new file mode 100644
index 000000000..50c056649
--- /dev/null
+++ b/tensor2tensor/rl/trainer_model_based_dqn_test.py
@@ -0,0 +1,41 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Tiny run of trainer_model_based. Smoke test."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# from tensor2tensor.rl import trainer_model_based
+
+import tensorflow as tf
+
+FLAGS = tf.flags.FLAGS
+
+
+class ModelRLExperimentTest(tf.test.TestCase):
+
+  def test_dqn_basic(self):
+    # TODO(afrozm): The latest changes in Dopamine break this test, so
+    # temporarily disabling this test.
+    pass
+    # FLAGS.output_dir = tf.test.get_temp_dir()
+    # FLAGS.loop_hparams_set = "rlmb_dqn_tiny"
+    # FLAGS.schedule = "train"  # skip evaluation for world model training
+    # trainer_model_based.main(None)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index a60efef1e..37eb47ddc 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -111,7 +111,6 @@ def _rlmb_base():
       # This is only used for world-model evaluation currently, PolicyLearner
       # uses algorithm specific hparams to set this during training.
       simulated_rollout_length=50,
-      wm_policy_param_sharing=False,
 
       # To be overridden.
       base_algo="",
@@ -146,6 +145,7 @@ def rlmb_ppo_base():
       # Number of simulated environments to train on simultaneously.
       simulated_batch_size=16,
       eval_batch_size=32,
+      wm_policy_param_sharing=False,
 
       # Unused; number of PPO epochs is calculated from the real frame limit.
       real_ppo_epochs_num=0,
@@ -194,17 +194,13 @@ def rlmb_dqn_base():
       base_algo="dqn",
       base_algo_params="dqn_original_params",
       real_batch_size=1,
-      simulated_batch_size=16,
+      simulated_batch_size=1,
       dqn_agent_generates_trainable_dones=False,
       eval_batch_size=1,
       # Must be equal to dqn_time_limit for now
       simulated_rollout_length=simulated_rollout_length,
       dqn_time_limit=simulated_rollout_length,
       simulation_flip_first_random_for_beginning=False,
-      dqn_eval_episodes_num=3,
-
-      # TODO(kc): only for model-free compatibility, remove this
-      epochs_num=-1,
   )
   update_hparams(hparams, dqn_params)
   return hparams
diff --git a/tensor2tensor/rl/trainer_model_based_test.py b/tensor2tensor/rl/trainer_model_based_test.py
index 87f5e61e3..0c836162c 100644
--- a/tensor2tensor/rl/trainer_model_based_test.py
+++ b/tensor2tensor/rl/trainer_model_based_test.py
@@ -27,18 +27,12 @@
 
 class ModelRLExperimentTest(tf.test.TestCase):
 
-  def _test_hparams_skip_evaluation(self, hparams_set):
+  def test_basic(self):
     FLAGS.output_dir = tf.test.get_temp_dir()
-    FLAGS.loop_hparams_set = hparams_set
+    FLAGS.loop_hparams_set = "rlmb_tiny"
     FLAGS.schedule = "train"  # skip evaluation for world model training
     trainer_model_based.main(None)
 
-  def test_basic(self):
-    self._test_hparams_skip_evaluation("rlmb_tiny")
-
-  def test_dqn_basic(self):
-    self._test_hparams_skip_evaluation("rlmb_dqn_tiny")
-
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 0d48a372d..5d7a0974a 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -103,51 +103,41 @@ def train(hparams, output_dir, env_problem_name, report_fn=None):
   tf.logging.vlog(1, "Policy HParams : %s",
                   misc_utils.pprint_hparams(policy_hparams))
 
-  # TODO(konradczechowski): remove base_algo dependance, when evaluation method
-  # will be decided
-  if hparams.base_algo == "ppo":
-    print("\n\n\npolicy_hparams {}\n\n\n".format(policy_hparams))
-    total_steps = policy_hparams.epochs_num
-    tf.logging.vlog(2, "total_steps: %d", total_steps)
-
-    eval_every_epochs = policy_hparams.eval_every_epochs
-    tf.logging.vlog(2, "eval_every_epochs: %d", eval_every_epochs)
-
-    if eval_every_epochs == 0:
-      eval_every_epochs = total_steps
-    policy_hparams.eval_every_epochs = 0
-
-    metric_name = rl_utils.get_metric_name(
-        sampling_temp=hparams.eval_sampling_temps[0],
-        max_num_noops=hparams.eval_max_num_noops,
-        clipped=False
-    )
-
-    tf.logging.vlog(1, "metric_name: %s", metric_name)
-
-    eval_metrics_dir = os.path.join(output_dir, "eval_metrics")
-    eval_metrics_dir = os.path.expanduser(eval_metrics_dir)
-    tf.gfile.MakeDirs(eval_metrics_dir)
-    eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
-
-    def evaluate_on_new_model(model_dir_path):
-      global step
-      eval_metrics = rl_utils.evaluate_all_configs(hparams, model_dir_path)
-      tf.logging.info(
-          "Agent eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
-      rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, step)
-      if report_fn:
-        report_fn(eval_metrics[metric_name], step)
-      step += 1
-
-    policy_hparams.epochs_num = total_steps
-    policy_hparams.save_models_every_epochs = eval_every_epochs
-  else:
-    def evaluate_on_new_model(model_dir_path):
-      del model_dir_path
-      raise NotImplementedError(
-          "This function is currently implemented only for ppo")
+  total_steps = policy_hparams.epochs_num
+  tf.logging.vlog(2, "total_steps: %d", total_steps)
+
+  eval_every_epochs = policy_hparams.eval_every_epochs
+  tf.logging.vlog(2, "eval_every_epochs: %d", eval_every_epochs)
+
+  if eval_every_epochs == 0:
+    eval_every_epochs = total_steps
+  policy_hparams.eval_every_epochs = 0
+
+  metric_name = rl_utils.get_metric_name(
+      sampling_temp=hparams.eval_sampling_temps[0],
+      max_num_noops=hparams.eval_max_num_noops,
+      clipped=False
+  )
 
+  tf.logging.vlog(1, "metric_name: %s", metric_name)
+
+  eval_metrics_dir = os.path.join(output_dir, "eval_metrics")
+  eval_metrics_dir = os.path.expanduser(eval_metrics_dir)
+  tf.gfile.MakeDirs(eval_metrics_dir)
+  eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
+
+  def evaluate_on_new_model(model_dir_path):
+    global step
+    eval_metrics = rl_utils.evaluate_all_configs(hparams, model_dir_path)
+    tf.logging.info(
+        "Agent eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
+    rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, step)
+    if report_fn:
+      report_fn(eval_metrics[metric_name], step)
+    step += 1
+
+  policy_hparams.epochs_num = total_steps
+  policy_hparams.save_models_every_epochs = eval_every_epochs
   learner.train(env_fn,
                 policy_hparams,
                 simulated=False,
diff --git a/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index 3e21a1387..98c28feca 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -28,18 +28,16 @@
 
 class TrainTest(tf.test.TestCase):
 
-  def _test_hparams_set(self, hparams_set):
-    hparams = registry.hparams(hparams_set)
+  def test_train_pong(self):
+    hparams = registry.hparams("rlmf_original")
+    hparams.batch_size = 2
+    hparams.eval_sampling_temps = [0.0, 1.0]
+    hparams.add_hparam("ppo_epochs_num", 2)
+    hparams.add_hparam("ppo_epoch_length", 3)
     FLAGS.output_dir = tf.test.get_temp_dir()
     trainer_model_free.train(hparams, FLAGS.output_dir,
                              env_problem_name=None)
 
-  def test_train_pong(self):
-    self._test_hparams_set("rlmf_tiny")
-
-  def test_train_pong_dqn(self):
-    self._test_hparams_set("rlmf_dqn_tiny")
-
 
 if __name__ == "__main__":
   tf.test.main()

From 3c7f7ca01d81d3dc1bf61b0ac1cde68ce05e229f Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 25 Mar 2019 16:23:45 -0700
Subject: [PATCH 1833/2720] Correct LayerNorm implementation and remove
 slax.multiplex to be more basic stax in Transformer (which now trains
 better).

PiperOrigin-RevId: 240241257
---
 .../trax/configs/transformer_lm1b_8gb.gin     |  2 +-
 tensor2tensor/trax/models/transformer.py      | 77 ++++++++++---------
 tensor2tensor/trax/stax/attention.py          | 16 ++--
 tensor2tensor/trax/stax/slax.py               | 16 ----
 tensor2tensor/trax/stax/slax_test.py          |  9 ++-
 5 files changed, 57 insertions(+), 63 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index a4b692d65..cded73df9 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -25,7 +25,7 @@ MultifactorSchedule.warmup_steps = 8000
 
 # Parameters for preprocess_fun:
 # ==============================================================================
-preprocess_fun.max_target_length = 512
+preprocess_fun.max_target_length = 511
 
 # Parameters for train:
 # ==============================================================================
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index d2f0d40d0..afe00d609 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -70,22 +70,23 @@ def encoder(embedded_source, source_mask):
     """
     encoder_layer = stax.serial(
         # input attends to self
-        stax.residual(stax.LayerNorm(feature_depth),
-                      stax.multiplex(stax.Identity,  # query
-                                     stax.Identity,  # key
-                                     stax.Identity,  # value
-                                     source_mask),  # attention mask
+        stax.residual(stax.LayerNorm(),
+                      stax.FanOut(4),
+                      stax.parallel(stax.Identity,  # query
+                                    stax.Identity,  # key
+                                    stax.Identity,  # value
+                                    source_mask),  # attention mask
                       multi_attention,
                       stax.Dropout(keep_rate, mode=mode)),
         # feed-forward
-        stax.residual(stax.LayerNorm(feature_depth),
+        stax.residual(stax.LayerNorm(),
                       feed_forward,
                       stax.Dropout(keep_rate, mode=mode))
     )
     return stax.serial(
         embedded_source,
         stax.repeat(encoder_layer, num_layers),
-        stax.LayerNorm(feature_depth),
+        stax.LayerNorm(),
     )
 
   return encoder
@@ -129,15 +130,16 @@ def TransformerLM(vocab_size,  # pylint: disable=invalid-name
   # Single decoder layer
   decoder_layer = stax.serial(
       # target attends to self
-      stax.residual(stax.LayerNorm(feature_depth),
-                    stax.multiplex(stax.Identity,  # query
-                                   stax.Identity,  # key
-                                   stax.Identity,  # value
-                                   stax.CausalMask(axis=-2)),  # attention mask
+      stax.residual(stax.LayerNorm(),
+                    stax.FanOut(4),
+                    stax.parallel(stax.Identity,  # query
+                                  stax.Identity,  # key
+                                  stax.Identity,  # value
+                                  stax.CausalMask(axis=-2)),  # attention mask
                     multi_attention,
                     stax.Dropout(keep_rate, mode=mode)),
       # feed-forward
-      stax.residual(stax.LayerNorm(feature_depth),
+      stax.residual(stax.LayerNorm(),
                     feed_forward,
                     stax.Dropout(keep_rate, mode=mode))
   )
@@ -145,10 +147,10 @@ def TransformerLM(vocab_size,  # pylint: disable=invalid-name
   return stax.serial(
       stax.ShiftRight(),
       stax.Embedding(feature_depth, vocab_size),
-      stax.PositionalEncoding(feature_depth, max_len=max_len),
       stax.Dropout(keep_rate, mode=mode),
+      stax.PositionalEncoding(feature_depth, max_len=max_len),
       stax.repeat(decoder_layer, num_layers),
-      stax.LayerNorm(feature_depth),
+      stax.LayerNorm(),
       stax.Dense(vocab_size, W_init=stax.xavier_uniform()),
       stax.LogSoftmax
   )
@@ -188,8 +190,8 @@ def Transformer(source_vocab_size,  # pylint: disable=invalid-name
   keep_rate = 1.0 - dropout
   # Input embedding and positional encoding
   inject_position = stax.serial(
-      stax.PositionalEncoding(feature_depth, max_len=max_len),
-      stax.Dropout(keep_rate, mode=mode)
+      stax.Dropout(keep_rate, mode=mode),
+      stax.PositionalEncoding(feature_depth, max_len=max_len)
   )
   if shared_embedding:
     assert source_vocab_size == target_vocab_size
@@ -228,15 +230,16 @@ def encoder(source, source_mask):
     """
     encoder_layer = stax.serial(
         # input attends to self
-        stax.residual(stax.LayerNorm(feature_depth),
-                      stax.multiplex(stax.Identity,  # query
-                                     stax.Identity,  # key
-                                     stax.Identity,  # value
-                                     source_mask),  # attention mask
+        stax.residual(stax.LayerNorm(),
+                      stax.FanOut(4),
+                      stax.parallel(stax.Identity,  # query
+                                    stax.Identity,  # key
+                                    stax.Identity,  # value
+                                    source_mask),  # attention mask
                       multi_attention,
                       stax.Dropout(keep_rate, mode=mode)),
         # feed-forward
-        stax.residual(stax.LayerNorm(feature_depth),
+        stax.residual(stax.LayerNorm(),
                       feed_forward,
                       stax.Dropout(keep_rate, mode=mode))
     )
@@ -244,7 +247,7 @@ def encoder(source, source_mask):
         source,
         source_embedding_layer,
         stax.repeat(encoder_layer, num_layers),
-        stax.LayerNorm(feature_depth),
+        stax.LayerNorm(),
     )
 
   # Decoder
@@ -263,23 +266,25 @@ def decoder(memory, target, target_mask, memory_mask):
     """
     decoder_layer = stax.serial(
         # target attends to self
-        stax.residual(stax.LayerNorm(feature_depth),
-                      stax.multiplex(stax.Identity,  # query
-                                     stax.Identity,  # key
-                                     stax.Identity,  # value
-                                     target_mask),  # attention mask
+        stax.residual(stax.LayerNorm(),
+                      stax.FanOut(4),
+                      stax.parallel(stax.Identity,  # query
+                                    stax.Identity,  # key
+                                    stax.Identity,  # value
+                                    target_mask),  # attention mask
                       multi_attention,
                       stax.Dropout(keep_rate, mode=mode)),
         # target attends to encoded source
-        stax.residual(stax.LayerNorm(feature_depth),
-                      stax.multiplex(stax.Identity,  # query
-                                     memory,  # key
-                                     memory,  # value
-                                     memory_mask),  # attention mask
+        stax.residual(stax.LayerNorm(),
+                      stax.FanOut(4),
+                      stax.parallel(stax.Identity,  # query
+                                    memory,  # key
+                                    memory,  # value
+                                    memory_mask),  # attention mask
                       multi_attention,
                       stax.Dropout(keep_rate, mode=mode)),
         # feed-forward
-        stax.residual(stax.LayerNorm(feature_depth),
+        stax.residual(stax.LayerNorm(),
                       feed_forward,
                       stax.Dropout(keep_rate, mode=mode))
     )
@@ -287,7 +292,7 @@ def decoder(memory, target, target_mask, memory_mask):
         target,
         target_embedding_layer,
         stax.repeat(decoder_layer, num_layers),
-        stax.LayerNorm(feature_depth),
+        stax.LayerNorm(),
     )
 
   # The Transformer
diff --git a/tensor2tensor/trax/stax/attention.py b/tensor2tensor/trax/stax/attention.py
index d9f2fd52c..afa444338 100644
--- a/tensor2tensor/trax/stax/attention.py
+++ b/tensor2tensor/trax/stax/attention.py
@@ -82,19 +82,21 @@ def init(shape):
   return init
 
 
-def LayerNorm(features, epsilon=1e-5):  # pylint: disable=invalid-name
+def LayerNorm(epsilon=1e-6):  # pylint: disable=invalid-name
   """Layer construction function for Layer Normalization layer.."""
   def init_fun(input_shape):
-    a_2 = np.ones(features)
-    b_2 = np.zeros(features)
-    return input_shape, (a_2, b_2)
+    features = input_shape[-1]
+    scale = np.ones(features)
+    bias = np.zeros(features)
+    return input_shape, (scale, bias)
 
   def apply_fun(params, inputs, **kwargs):
     del kwargs
-    (a_2, b_2) = params
+    (scale, bias) = params
     mean = np.mean(inputs, axis=-1, keepdims=True)
-    std = np.std(inputs, axis=-1, keepdims=True)
-    return a_2 * (inputs - mean) / (std + epsilon) + b_2
+    variance = np.mean((inputs - mean)**2, axis=-1, keepdims=True)
+    norm_inputs = (inputs - mean) / np.sqrt(variance + epsilon)
+    return norm_inputs * scale + bias
 
   return init_fun, apply_fun
 
diff --git a/tensor2tensor/trax/stax/slax.py b/tensor2tensor/trax/stax/slax.py
index d092f4387..a99631241 100644
--- a/tensor2tensor/trax/stax/slax.py
+++ b/tensor2tensor/trax/stax/slax.py
@@ -73,22 +73,6 @@ def residual(*layers, **kwargs):
     raise ValueError('Empty residual combinator.')
 
 
-def multiplex(*args):
-  """Helper to form input argument lists of bound variables.
-
-  Args:
-    *args: list of bound layers or raw stax Identity layers.
-
-  Returns:
-    A layer returning in parallel the bound variables as well as
-  (multiple) copies of this layer's input wherever Identity has been specified.
-  """
-  return stax.serial(
-      stax.FanOut(len(args)),
-      stax.parallel(*args)
-  )
-
-
 # Utility Layers
 # ------------------------------------------------------------------------------
 def Take(*args):  # pylint: disable=invalid-name
diff --git a/tensor2tensor/trax/stax/slax_test.py b/tensor2tensor/trax/stax/slax_test.py
index 49cca2890..bdb7783d0 100644
--- a/tensor2tensor/trax/stax/slax_test.py
+++ b/tensor2tensor/trax/stax/slax_test.py
@@ -184,7 +184,8 @@ def testLambda_5_args_2_post_input_tree(self):
       def lambda_fun1(x, y, z, w, v):
         input_tree = _build_combinator_tree(tree_spec, (x, y, z))
         return stax.serial(input_tree,
-                           stax.multiplex(stax.Identity, w, v),
+                           stax.FanOut(3),
+                           stax.parallel(stax.Identity, w, v),
                            stax.FanInSum)
       check_staxlayer(self, lambda_fun1, [(1, 5, 7, 11),]*5)
 
@@ -192,7 +193,8 @@ def lambda_fun1(x, y, z, w, v):
       def lambda_fun2(x, y, z, w, v):
         input_tree = _build_combinator_tree(tree_spec, (x, y, z))
         return stax.serial(input_tree,
-                           stax.multiplex(w, stax.Identity, v),
+                           stax.FanOut(3),
+                           stax.parallel(w, stax.Identity, v),
                            stax.FanInSum)
       check_staxlayer(self, lambda_fun2, [(1, 5, 7, 11),]*5)
 
@@ -200,7 +202,8 @@ def lambda_fun2(x, y, z, w, v):
       def lambda_fun3(x, y, z, w, v):
         input_tree = _build_combinator_tree(tree_spec, (x, y, z))
         return stax.serial(input_tree,
-                           stax.multiplex(w, v, stax.Identity),
+                           stax.FanOut(3),
+                           stax.parallel(w, v, stax.Identity),
                            stax.FanInSum)
       check_staxlayer(self, lambda_fun3, [(1, 5, 7, 11),]*5)
   # pylint: enable=cell-var-from-loop

From f28a5e9cec63fb2e0af575fdca754318b0cbdbca Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Mon, 25 Mar 2019 16:29:11 -0700
Subject: [PATCH 1834/2720] Merge of PR #1500

PiperOrigin-RevId: 240242179
---
 tensor2tensor/layers/common_video.py          |   4 +-
 tensor2tensor/models/research/rl.py           |  56 ++-
 tensor2tensor/rl/batch_dqn_agent_test.py      | 158 +++++++
 tensor2tensor/rl/batch_runner_test.py         | 284 ++++++++++++
 tensor2tensor/rl/dopamine_connector.py        | 418 ++++++++++++++----
 tensor2tensor/rl/envs/simulated_batch_env.py  |   3 +-
 .../rl/envs/simulated_batch_gym_env.py        |   4 +-
 tensor2tensor/rl/policy_learner.py            |   2 -
 tensor2tensor/rl/trainer_model_based.py       |   2 +
 .../rl/trainer_model_based_dqn_test.py        |  41 --
 .../rl/trainer_model_based_params.py          |   8 +-
 tensor2tensor/rl/trainer_model_based_test.py  |  11 +-
 tensor2tensor/rl/trainer_model_free.py        |  78 ++--
 tensor2tensor/rl/trainer_model_free_test.py   |  14 +-
 14 files changed, 900 insertions(+), 183 deletions(-)
 create mode 100644 tensor2tensor/rl/batch_dqn_agent_test.py
 create mode 100644 tensor2tensor/rl/batch_runner_test.py
 delete mode 100644 tensor2tensor/rl/trainer_model_based_dqn_test.py

diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index ea94316cf..0d2a1af4c 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -790,6 +790,8 @@ def finish(self):
     (out, err) = [
         b"".join(chunks) for chunks in (self._out_chunks, self._err_chunks)
     ]
+    self.proc.stdout.close()
+    self.proc.stderr.close()
     if self.proc.returncode:
       err = "\n".join([" ".join(self.cmd), err.decode("utf8")])
       raise IOError(err)
@@ -820,7 +822,7 @@ def write(self, batch_frame, batch_encoded_frame=None):
     del batch_encoded_frame
     if self.writers is None:
       self.writers = [
-          WholeVideoWriter(
+          WholeVideoWriter(  # pylint: disable=g-complex-comprehension
               self.fps, self.path_template.format(i), self.file_format
           )
           for i in range(len(batch_frame))
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 90c74aeef..6025a132c 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -354,12 +354,18 @@ def dqn_atari_base():
       optimizer_epsilon=0.00001,
       optimizer_centered=True,
 
+      # TODO(kozak): change names maybe replay_buffer -> agent?
+      # Also batch_size is now buffer_batch_size in _DQNAgent.
       replay_buffer_replay_capacity=1000000,
-      replay_buffer_batch_size=32,
+      replay_buffer_buffer_batch_size=32,
 
       time_limit=27000,
       save_every_steps=50000,
       num_frames=int(20 * 1e6),
+
+      # TODO(konradczechowski) this is not used in trainer_model_free, clean
+      # this up after evaluation refactor
+      eval_episodes_num=3,
   )
 
 
@@ -371,6 +377,17 @@ def dqn_original_params():
   return hparams
 
 
+def rlmf_tiny_overrides():
+  """Parameters to override for tiny setting excluding agent-related hparams."""
+  return dict(
+      max_num_noops=1,
+      eval_max_num_noops=1,
+      rl_env_max_episode_steps=7,
+      eval_rl_env_max_episode_steps=7,
+      eval_sampling_temps=[0.0, 1.0],
+  )
+
+
 @registry.register_hparams
 def rlmf_original():
   return HParams(
@@ -382,6 +399,7 @@ def rlmf_original():
       eval_batch_size=2,
       frame_stack_size=4,
       eval_sampling_temps=[0.0, 0.2, 0.5, 0.8, 1.0, 2.0],
+      max_num_noops=8,
       eval_max_num_noops=8,
       eval_rl_env_max_episode_steps=1000,
       resize_height_factor=2,
@@ -405,7 +423,7 @@ def rlmf_tictactoe():
   # Since we don't have any no-op actions, otherwise we have to have an
   # attribute called `get_action_meanings`.
   hparams.eval_max_num_noops = 0
-  hparams.add_hparam("max_num_noops", 0)
+  hparams.max_num_noops = 0
   hparams.rl_should_derive_observation_space = False
 
   hparams.policy_network = "feed_forward_categorical_policy"
@@ -426,6 +444,32 @@ def rlmf_base():
   return hparams
 
 
+@registry.register_hparams
+def rlmf_tiny():
+  """Tiny set of hparams for model-free PPO."""
+  hparams = rlmf_original()
+  hparams = hparams.override_from_dict(rlmf_tiny_overrides())
+  hparams.batch_size = 2
+  hparams.add_hparam("ppo_epochs_num", 3)
+  hparams.add_hparam("ppo_epoch_length", 2)
+  return hparams
+
+
+@registry.register_hparams
+def rlmf_dqn_tiny():
+  """Tiny DQN params."""
+  hparams = rlmf_original()
+  hparams = hparams.override_from_dict(rlmf_tiny_overrides())
+  hparams.batch_size = 1
+  hparams.base_algo = "dqn"
+  hparams.base_algo_params = "dqn_original_params"
+  hparams.add_hparam("dqn_num_frames", 128)
+  hparams.add_hparam("dqn_save_every_steps", 128)
+  hparams.add_hparam("dqn_replay_buffer_replay_capacity", 100)
+  hparams.add_hparam("dqn_agent_min_replay_history", 10)
+  return hparams
+
+
 @registry.register_hparams
 def rlmf_eval():
   """Eval set of hparams for model-free PPO."""
@@ -442,14 +486,6 @@ def rlmf_eval():
   return hparams
 
 
-@registry.register_hparams
-def rlmf_tiny():
-  hparams = rlmf_base()
-  hparams.ppo_epochs_num = 100
-  hparams.ppo_eval_every_epochs = 10
-  return hparams
-
-
 class PolicyBase(t2t_model.T2TModel):
 
   def loss(self, *args, **kwargs):
diff --git a/tensor2tensor/rl/batch_dqn_agent_test.py b/tensor2tensor/rl/batch_dqn_agent_test.py
new file mode 100644
index 000000000..365d7b054
--- /dev/null
+++ b/tensor2tensor/rl/batch_dqn_agent_test.py
@@ -0,0 +1,158 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 BatchDQNAgent."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import shutil
+
+from absl import flags
+from dopamine.agents.dqn import dqn_agent
+import numpy as np
+
+from tensor2tensor.rl import dopamine_connector
+
+import tensorflow as tf
+
+
+slim = tf.contrib.slim
+
+FLAGS = flags.FLAGS
+
+
+class BatchDQNAgentTest(tf.test.TestCase):
+  # TODO(kozak): add testStepTrain (and possibly other tests) from dopamine
+  # dqn_agent_test.py
+
+  def setUp(self):
+    self._test_subdir = os.path.join('/tmp/dopamine_tests', 'ckpts')
+    shutil.rmtree(self._test_subdir, ignore_errors=True)
+    os.makedirs(self._test_subdir)
+    self.num_actions = 4
+    self.min_replay_history = 6
+    self.update_period = 2
+    self.target_update_period = 4
+    self.epsilon_decay_period = 90
+    self.epsilon_train = 0.05
+    self.observation_shape = dqn_agent.NATURE_DQN_OBSERVATION_SHAPE
+    self.stack_size = dqn_agent.NATURE_DQN_STACK_SIZE
+    self.env_batch_size = 4
+
+    self.zero_state = np.zeros(
+        [self.env_batch_size, self.observation_shape[0],
+         self.observation_shape[1], self.stack_size])
+
+  def _create_test_agent(self, sess):
+    stack_size = self.stack_size
+
+    class MockDQNAgent(dopamine_connector.BatchDQNAgent):
+
+      def _network_template(self, state):
+        # This dummy network allows us to deterministically anticipate that
+        # action 0 will be selected by an argmax.
+        inputs = tf.constant(
+            np.zeros((state.shape[0], stack_size)), dtype=tf.float32)
+        # This weights_initializer gives action 0 a higher weight, ensuring
+        # that it gets picked by the argmax.
+        weights_initializer = np.tile(
+            np.arange(self.num_actions, 0, -1), (stack_size, 1))
+        q = slim.fully_connected(
+            inputs,
+            self.num_actions,
+            weights_initializer=tf.constant_initializer(weights_initializer),
+            biases_initializer=tf.ones_initializer(),
+            activation_fn=None)
+        return self._get_network_type()(q)
+
+    agent = MockDQNAgent(
+        replay_capacity=100,
+        buffer_batch_size=8,
+        generates_trainable_dones=True,
+        sess=sess,
+        env_batch_size=self.env_batch_size,
+        num_actions=self.num_actions,
+        min_replay_history=self.min_replay_history,
+        epsilon_fn=lambda w, x, y, z: 0.0,  # No exploration.
+        update_period=self.update_period,
+        target_update_period=self.target_update_period,
+        epsilon_eval=0.0)  # No exploration during evaluation.
+    # This ensures non-random action choices (since epsilon_eval = 0.0) and
+    # skips the train_step.
+    agent.eval_mode = True
+    sess.run(tf.global_variables_initializer())
+    return agent
+
+  def testCreateAgentWithDefaults(self):
+    # Verifies that we can create and train an agent with the default values.
+    with tf.Session() as sess:
+      agent = self._create_test_agent(sess)
+      sess.run(tf.global_variables_initializer())
+      observation = np.ones([84, 84, 1])
+      agent.begin_episode([observation])
+      agent.step(reward=[1], observation=[observation])
+      agent.end_episode(reward=[1])
+
+  def testBeginEpisode(self):
+    """Test the functionality of agent.begin_episode.
+
+    Specifically, the action returned and its effect on state.
+    """
+    with tf.Session() as sess:
+      agent = self._create_test_agent(sess)
+      # We fill up the state with 9s. On calling agent.begin_episode the state
+      # should be reset to all 0s.
+      agent.state_batch.fill(9)
+      first_observation = np.ones(
+          [self.env_batch_size, self.observation_shape[0],
+           self.observation_shape[1], 1])
+      self.assertTrue((agent.begin_episode(first_observation) == 0).all())
+      # When the all-1s observation is received, it will be placed at the end of
+      # the state.
+      expected_state = self.zero_state
+      expected_state[:, :, :, -1] = np.ones(
+          [self.env_batch_size, self.observation_shape[0],
+           self.observation_shape[1]])
+      self.assertAllEqual(agent.state_batch, expected_state)
+      self.assertAllEqual(agent._observation_batch, first_observation[..., 0])
+      # No training happens in eval mode.
+      self.assertEqual(agent.training_steps, 0)
+
+      # This will now cause training to happen.
+      agent.eval_mode = False
+      # Having a low replay memory add_count will prevent any of the
+      # train/prefetch/sync ops from being called.
+      agent._replay.memory.add_count = 0
+      second_observation = np.ones(
+          [self.env_batch_size, self.observation_shape[0],
+           self.observation_shape[1], 1]) * 2
+      agent.begin_episode(second_observation)
+      # The agent's state will be reset, so we will only be left with the all-2s
+      # observation.
+      expected_state[:, :, :, -1] = np.full(
+          (self.env_batch_size, self.observation_shape[0],
+           self.observation_shape[1]), 2
+      )
+      self.assertAllEqual(agent.state_batch, expected_state)
+      self.assertAllEqual(agent._observation_batch,
+                          second_observation[:, :, :, 0])
+      # training_steps is incremented since we set eval_mode to False.
+      self.assertEqual(agent.training_steps, 1)
+
+
+if __name__ == '__main__':
+  tf.test.main()
diff --git a/tensor2tensor/rl/batch_runner_test.py b/tensor2tensor/rl/batch_runner_test.py
new file mode 100644
index 000000000..3b3ee21d4
--- /dev/null
+++ b/tensor2tensor/rl/batch_runner_test.py
@@ -0,0 +1,284 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 BatchRunner."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import shutil
+
+from absl import flags
+from dopamine.discrete_domains import logger
+import mock
+import numpy as np
+
+from tensor2tensor.rl import dopamine_connector
+
+import tensorflow as tf
+
+
+FLAGS = flags.FLAGS
+
+
+def _create_mock_checkpointer():
+  mock_checkpointer = mock.Mock()
+  test_dictionary = {"current_iteration": 1729,
+                     "logs": "logs"}
+  mock_checkpointer.load_checkpoint.return_value = test_dictionary
+  return mock_checkpointer
+
+
+class MockEnvironment(object):
+  """Mock environment for testing."""
+
+  def __init__(self, max_steps=10, reward_multiplier=1):
+    self._observation = 0
+    self.max_steps = max_steps
+    self.reward_multiplier = reward_multiplier
+    self.game_over = False
+
+  def reset(self):
+    self._observation = 0
+    return self._observation
+
+  def step(self, action):
+    self._observation += 1
+    action_reward_multiplier = -1 if action > 0 else 1
+    reward_multiplier = self.reward_multiplier * action_reward_multiplier
+    reward = self._observation * reward_multiplier
+    is_terminal = self._observation >= self.max_steps
+    self.game_over = is_terminal
+
+    unused = 0
+    return (self._observation, reward, is_terminal, unused)
+
+  def render(self, mode):
+    pass
+
+
+class BatchEnv(object):
+  """Batch env.
+
+  Batch of environments. Assumes that all throws "done" on the same step.
+
+  Observations and rewards are returned as arrays, done as single value.
+  """
+
+  # TODO(kozak): this can be used for mbrl pipeline (for both simulated and
+  # real env), move it to dopamine_connector.py (rename it?)
+  def __init__(self, envs):
+    self.env_batch = envs
+    self.batch_size = len(self.env_batch)
+    self.max_steps = self.env_batch[0].max_steps
+    assert np.all(self.max_steps == env.max_steps for env in self.env_batch)
+
+  def step(self, actions):
+    ret = [env.step(action) for env, action in zip(self.env_batch, actions)]
+    obs, rewards, dones, infos = [np.array(r) for r in zip(*ret)]
+    done = dones[0]
+    assert np.all(done == dones)
+    self.game_over = done
+    return obs, rewards, done, infos
+
+  def reset(self):
+    return np.array([env.reset() for env in self.env_batch])
+
+  def render(self, mode):
+    pass
+
+
+class MockLogger(object):
+  """Class to mock the experiment logger."""
+
+  def __init__(self, test_cls=None, run_asserts=True, data=None):
+    self._test_cls = test_cls
+    self._run_asserts = run_asserts
+    self._iter = 0
+    self._calls_to_set = 0
+    self._calls_to_log = 0
+    self.data = data
+
+  def __setitem__(self, key, val):
+    if self._run_asserts:
+      self._test_cls.assertEqual("iteration_{:d}".format(self._iter), key)
+      self._test_cls.assertEqual("statistics", val)
+      self._iter += 1
+    self._calls_to_set += 1
+
+  def log_to_file(self, filename_prefix, iteration_number):
+    if self._run_asserts:
+      self._test_cls.assertEqual(
+          "prefix_{}".format(self._iter - 1),
+          "{}_{}".format(filename_prefix, iteration_number))
+    self._calls_to_log += 1
+
+
+class BatchedRunnerTest(tf.test.TestCase):
+  """Modified tests from dopamine run_experiment_test.py."""
+
+  # TODO(kozak): decide if we want to use and modify more tests from
+  # dopamine/tests/atari/run_experiment_test.py (e.g.  testRunExperiment.py)
+
+  def _agent_step(self, rewards, observations):
+    # We verify that rewards are clipped (and set by MockEnvironment as a
+    # function of observation)
+    # observation = observations[0]
+    # expected_rewards = [1 if observation % 2 else -1]
+    # self.assertEqual(expected_reward, reward)
+    actions = [ob % 2 for ob in observations]
+    return actions
+
+  def prepare_mock_agent(self, batch_size):
+    assert batch_size % 2 == 0, "Some of tests assume that batch_size % 2 == 0"
+    self.batch_size = batch_size
+    self._agent = mock.Mock()
+    self._agent.begin_episode.side_effect = \
+      lambda x: np.repeat(0, self.batch_size)
+    self._agent.step.side_effect = self._agent_step
+    self._create_agent_fn = lambda x, y, summary_writer: self._agent
+
+  def setUp(self):
+    super(BatchedRunnerTest, self).setUp()
+    self._test_subdir = "/tmp/dopamine_tests"
+    shutil.rmtree(self._test_subdir, ignore_errors=True)
+    os.makedirs(self._test_subdir)
+    self.prepare_mock_agent(batch_size=4)
+
+  def testRunEpisodeBatch(self):
+    max_steps_per_episode = 11
+    batch_size = self.batch_size
+    reward_multipliers = [-1, 1] * int(batch_size / 2)
+    envs = [MockEnvironment(reward_multiplier=rm) for rm in reward_multipliers]
+    environment = BatchEnv(envs)
+    runner = dopamine_connector.BatchRunner(
+        self._test_subdir, self._create_agent_fn,
+        create_environment_fn=lambda: environment,
+        max_steps_per_episode=max_steps_per_episode)
+    step_number, total_rewards = runner._run_one_episode()
+
+    self.assertEqual(self._agent.step.call_count, environment.max_steps - 1)
+    self.assertEqual(self._agent.end_episode.call_count, 1)
+    self.assertEqual(environment.max_steps, step_number / batch_size)
+    # Expected reward will be \sum_{i=0}^{9} (-1)**i * i = -5 when reward
+    # multiplier=1
+    self.assertAllEqual(np.array(reward_multipliers) * -5, total_rewards)
+
+  def testRunOneEpisodeWithLowMaxSteps(self):
+    max_steps_per_episode = 2
+    batch_size = self.batch_size
+    reward_multipliers = [-1, 1] * int(batch_size / 2)
+    envs = [MockEnvironment(reward_multiplier=rm) for rm in reward_multipliers]
+    environment = BatchEnv(envs)
+    runner = dopamine_connector.BatchRunner(
+        self._test_subdir, self._create_agent_fn,
+        create_environment_fn=lambda: environment,
+        max_steps_per_episode=max_steps_per_episode)
+    step_number, total_rewards = runner._run_one_episode()
+
+    self.assertEqual(self._agent.step.call_count, max_steps_per_episode - 1)
+    self.assertEqual(self._agent.end_episode.call_count, 1)
+    self.assertEqual(max_steps_per_episode, step_number / batch_size)
+    self.assertAllEqual(np.array(reward_multipliers) * -1, total_rewards)
+
+  def testRunOnePhase(self):
+    batch_size = self.batch_size
+    environment_steps = 2
+    max_steps = environment_steps * batch_size * 10
+
+    envs = [MockEnvironment(max_steps=environment_steps)
+            for _ in range(batch_size)]
+
+    environment = BatchEnv(envs)
+    runner = dopamine_connector.BatchRunner(
+        self._test_subdir, self._create_agent_fn,
+        create_environment_fn=lambda: environment)
+
+    statistics = []
+
+    step_number, sum_returns, num_episodes = runner._run_one_phase(
+        max_steps, statistics, "test")
+    calls_to_run_episode = int(max_steps / (environment_steps * batch_size))
+    self.assertEqual(self._agent.step.call_count, calls_to_run_episode)
+    self.assertEqual(self._agent.end_episode.call_count, calls_to_run_episode)
+    self.assertEqual(max_steps, step_number)
+    self.assertEqual(-1 * calls_to_run_episode * batch_size, sum_returns)
+    self.assertEqual(calls_to_run_episode, num_episodes / batch_size)
+    expected_statistics = []
+    for _ in range(calls_to_run_episode * batch_size):
+      expected_statistics.append({
+          "test_episode_lengths": 2,
+          "test_episode_returns": -1
+      })
+    self.assertEqual(len(expected_statistics), len(statistics))
+    for expected_stats, stats in zip(expected_statistics, statistics):
+      self.assertDictEqual(expected_stats, stats)
+
+  def testRunOneIteration(self):
+    environment_steps = 2
+    batch_size = self.batch_size
+    envs = [MockEnvironment(max_steps=environment_steps)
+            for _ in range(batch_size)]
+
+    environment = BatchEnv(envs)
+
+    training_steps = 20 * batch_size
+    evaluation_steps = 10 * batch_size
+
+    runner = dopamine_connector.BatchRunner(
+        self._test_subdir, self._create_agent_fn,
+        create_environment_fn=lambda: environment,
+        training_steps=training_steps, evaluation_steps=evaluation_steps
+    )
+
+    dictionary = runner._run_one_iteration(1)
+    train_rollouts = int(training_steps / environment_steps)
+    eval_rollouts = int(evaluation_steps / environment_steps)
+    expected_dictionary = {
+        "train_episode_lengths": [2 for _ in range(train_rollouts)],
+        "train_episode_returns": [-1 for _ in range(train_rollouts)],
+        "train_average_return": [-1],
+        "eval_episode_lengths": [2 for _ in range(eval_rollouts)],
+        "eval_episode_returns": [-1 for _ in range(eval_rollouts)],
+        "eval_average_return": [-1]
+    }
+    self.assertDictEqual(expected_dictionary, dictionary)
+
+  @mock.patch.object(logger, "Logger")
+  def testLogExperiment(self, mock_logger_constructor):
+    # TODO(kozak): We probably do not need this test, dopamine test
+    # for Runner is enough here. Remove this?
+    log_every_n = 2
+    logging_file_prefix = "prefix"
+    statistics = "statistics"
+    experiment_logger = MockLogger(test_cls=self)
+    mock_logger_constructor.return_value = experiment_logger
+    runner = dopamine_connector.BatchRunner(
+        self._test_subdir, self._create_agent_fn,
+        create_environment_fn=mock.Mock,
+        logging_file_prefix=logging_file_prefix,
+        log_every_n=log_every_n)
+    num_iterations = 10
+    for i in range(num_iterations):
+      runner._log_experiment(i, statistics)
+    self.assertEqual(num_iterations, experiment_logger._calls_to_set)
+    self.assertEqual((num_iterations / log_every_n),
+                     experiment_logger._calls_to_log)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 9e1ad3765..1686c86f4 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -20,17 +20,14 @@
 from __future__ import print_function
 
 import copy
+import random
+import sys
 
 from dopamine.agents.dqn import dqn_agent
 from dopamine.replay_memory import circular_replay_buffer
 from dopamine.replay_memory.circular_replay_buffer import OutOfGraphReplayBuffer
 from dopamine.replay_memory.circular_replay_buffer import ReplayElement
-import gym
-from gym import spaces
-from gym import Wrapper
-from gym.wrappers import TimeLimit
 import numpy as np
-from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
 from tensor2tensor.rl.policy_learner import PolicyLearner
 import tensorflow as tf
 
@@ -46,52 +43,6 @@
 # pylint: enable=g-import-not-at-top
 
 
-class ResizeObservation(gym.ObservationWrapper):
-  """TODO(konradczechowski): Add doc-string."""
-
-  def __init__(self, env, size=84):
-    """Based on WarpFrame from openai baselines atari_wrappers.py.
-
-    Dopamine also uses cv2.resize(..., interpolation=cv2.INTER_AREA).
-
-    Args:
-      env: TODO(konradczechowski): Add doc-string.
-      size: TODO(konradczechowski): Add doc-string.
-    """
-    gym.ObservationWrapper.__init__(self, env)
-    self.width = size
-    self.height = size
-    assert env.observation_space.dtype == np.uint8
-    self.observation_space = spaces.Box(
-        low=0,
-        high=255,
-        shape=(self.height, self.width, env.observation_space.shape[2]),
-        dtype=np.uint8)
-
-  def observation(self, frame):
-    if not cv2:
-      return frame
-    return cv2.resize(
-        frame, (self.width, self.height), interpolation=cv2.INTER_AREA)
-
-
-class GameOverOnDone(Wrapper):
-  """TODO(konradczechowski): Add doc-string."""
-
-  def __init__(self, env):
-    Wrapper.__init__(self, env)
-    self.game_over = False
-
-  def reset(self, **kwargs):
-    self.game_over = False
-    return self.env.reset(**kwargs)
-
-  def step(self, action):
-    ob, reward, done, info = self.env.step(action)
-    self.game_over = done
-    return ob, reward, done, info
-
-
 class _DQNAgent(dqn_agent.DQNAgent):
   """Modify dopamine DQNAgent to match our needs.
 
@@ -99,10 +50,10 @@ class _DQNAgent(dqn_agent.DQNAgent):
   (some of) terminal episode transitions in training.
   """
 
-  def __init__(self, replay_capacity, batch_size, generates_trainable_dones,
-               **kwargs):
+  def __init__(self, replay_capacity, buffer_batch_size,
+               generates_trainable_dones, **kwargs):
     self._replay_capacity = replay_capacity
-    self._batch_size = batch_size
+    self._buffer_batch_size = buffer_batch_size
     self._generates_trainable_dones = generates_trainable_dones
     super(_DQNAgent, self).__init__(**kwargs)
 
@@ -112,7 +63,7 @@ def _build_replay_buffer(self, use_staging):
         observation_shape=dqn_agent.NATURE_DQN_OBSERVATION_SHAPE,
         stack_size=dqn_agent.NATURE_DQN_STACK_SIZE,
         replay_capacity=self._replay_capacity,
-        batch_size=self._batch_size,
+        batch_size=self._buffer_batch_size,
         update_horizon=self.update_horizon,
         gamma=self.gamma,
         extra_storage_types=None,
@@ -128,10 +79,154 @@ def _build_replay_buffer(self, use_staging):
         **replay_buffer_kwargs)
 
 
+class BatchDQNAgent(_DQNAgent):
+  """Batch agent for DQN.
+
+  Episodes are stored on done.
+
+  Assumes that all rollouts in batch would end at the same moment.
+  """
+
+  def __init__(self, env_batch_size, *args, **kwargs):
+    super(BatchDQNAgent, self).__init__(*args, **kwargs)
+    self.env_batch_size = env_batch_size
+    obs_size = dqn_agent.NATURE_DQN_OBSERVATION_SHAPE
+    state_shape = [self.env_batch_size, obs_size[0], obs_size[1],
+                   dqn_agent.NATURE_DQN_STACK_SIZE]
+    self.state_batch = np.zeros(state_shape)
+    self.state = None  # assure it will be not used
+    self._observation = None  # assure it will be not used
+    self.reset_current_rollouts()
+
+  def reset_current_rollouts(self):
+    self._current_rollouts = [[] for _ in range(self.env_batch_size)]
+
+  def _record_observation(self, observation_batch):
+    # Set current observation. Represents an (batch_size x 84 x 84 x 1) image
+    # frame.
+    observation_batch = np.array(observation_batch)
+    self._observation_batch = observation_batch[:, :, :, 0]
+    # Swap out the oldest frames with the current frames.
+    self.state_batch = np.roll(self.state_batch, -1, axis=3)
+    self.state_batch[:, :, :, -1] = self._observation_batch
+
+  def _reset_state(self):
+    self.state_batch.fill(0)
+
+  def begin_episode(self, observation):
+    self._reset_state()
+    self._record_observation(observation)
+
+    if not self.eval_mode:
+      self._train_step()
+
+    self.action = self._select_action()
+    return self.action
+
+  def _update_current_rollouts(self, last_observation, action, reward,
+                               are_terminal):
+    transitions = zip(last_observation, action, reward, are_terminal)
+    for transition, rollout in zip(transitions, self._current_rollouts):
+      rollout.append(transition)
+
+  def _store_current_rollouts(self):
+    for rollout in self._current_rollouts:
+      for transition in rollout:
+        self._store_transition(*transition)
+    self.reset_current_rollouts()
+
+  def step(self, reward, observation):
+    self._last_observation = self._observation_batch
+    self._record_observation(observation)
+
+    if not self.eval_mode:
+      self._update_current_rollouts(self._last_observation, self.action, reward,
+                                    [False] * self.env_batch_size)
+      # We want to have the same train_step:env_step ratio not depending on
+      # batch size.
+      for _ in range(self.env_batch_size):
+        self._train_step()
+
+    self.action = self._select_action()
+    return self.action
+
+  def end_episode(self, reward):
+    if not self.eval_mode:
+      self._update_current_rollouts(
+          self._observation_batch, self.action, reward,
+          [True] * self.env_batch_size)
+      self._store_current_rollouts()
+
+  def _select_action(self):
+    epsilon = self.epsilon_eval
+    if not self.eval_mode:
+      epsilon = self.epsilon_fn(
+          self.epsilon_decay_period,
+          self.training_steps,
+          self.min_replay_history,
+          self.epsilon_train)
+
+    def choose_action(ix):
+      if random.random() <= epsilon:
+        # Choose a random action with probability epsilon.
+        return random.randint(0, self.num_actions - 1)
+      else:
+        # Choose the action with highest Q-value at the current state.
+        return self._sess.run(self._q_argmax,
+                              {self.state_ph: self.state_batch[ix:ix+1]})
+
+    return np.array([choose_action(ix) for ix in range(self.env_batch_size)])
+
+
+class BatchRunner(run_experiment.Runner):
+  """Run a batch of environments.
+
+  Assumes that all environments would end at the same moment.
+  """
+
+  def __init__(self, base_dir, create_agent_fn, **kwargs):
+    super(BatchRunner, self).__init__(base_dir, create_agent_fn, **kwargs)
+    self.batch_size = self._environment.batch_size
+
+  def _run_one_episode(self):
+    # This assumes that everything inside _run_one_episode works on batches,
+    # which is risky for future.
+    steps_number, total_rewards = super(BatchRunner, self)._run_one_episode()
+    return steps_number * self.batch_size, total_rewards
+
+  def _run_one_phase(self, min_steps, statistics, run_mode_str):
+    # Mostly copy of parent method.
+    step_count = 0
+    num_episodes = 0
+    sum_returns = 0.
+
+    while step_count < min_steps:
+      num_steps, episode_returns = self._run_one_episode()
+      for episode_return in episode_returns:
+        statistics.append({
+            "{}_episode_lengths".format(run_mode_str):
+                num_steps / self.batch_size,
+            "{}_episode_returns".format(run_mode_str): episode_return
+        })
+      step_count += num_steps
+      sum_returns += sum(episode_returns)
+      num_episodes += self.batch_size
+      # We use sys.stdout.write instead of tf.logging so as to flush frequently
+      # without generating a line break.
+      sys.stdout.write("Steps executed: {} ".format(step_count) +
+                       "Batch episodes steps: {} ".format(num_steps) +
+                       "Returns: {}\r".format(episode_returns))
+      sys.stdout.flush()
+    return step_count, sum_returns, num_episodes
+
+  def close(self):
+    self._environment.close()
+
+
 class _OutOfGraphReplayBuffer(OutOfGraphReplayBuffer):
   """Replay not sampling artificial_terminal transition.
 
-  Adds to stored tuples 'artificial_done' field (as last ReplayElement).
+  Adds to stored tuples "artificial_done" field (as last ReplayElement).
   When sampling, ignores tuples for which artificial_done is True.
 
   When adding new attributes check if there are loaded from disk, when using
@@ -177,7 +272,14 @@ def load(self, *args, **kwargs):
 
 
 def get_create_agent(agent_kwargs):
-  """TODO(): Document."""
+  """Factory for dopamine agent initialization.
+
+  Args:
+    agent_kwargs: dict of BatchDQNAgent parameters
+
+  Returns:
+    Function(sess, environment, summary_writer) -> BatchDQNAgent instance.
+  """
 
   def create_agent(sess, environment, summary_writer=None):
     """Creates a DQN agent.
@@ -192,7 +294,8 @@ def create_agent(sess, environment, summary_writer=None):
     Returns:
       a DQN agent.
     """
-    return _DQNAgent(
+    return BatchDQNAgent(
+        env_batch_size=environment.batch_size,
         sess=sess,
         num_actions=environment.action_space.n,
         summary_writer=summary_writer,
@@ -202,23 +305,178 @@ def create_agent(sess, environment, summary_writer=None):
   return create_agent
 
 
-def get_create_env_fun(batch_env_fn, time_limit):
-  """TODO(konradczechowski): Add doc-string."""
+class ResizeBatchObservation(object):
+  """Wrapper resizing observations for batched environment.
+
+  Dopamine also uses cv2.resize(..., interpolation=cv2.INTER_AREA).
+
+  Attributes:
+    batch_env: batched environment
+    batch_size: batch size
+    action_space: the action space
+    size: size of width and height for returned observations
+  """
+
+  def __init__(self, batch_env, size=84):
+    self.size = size
+    self.batch_env = batch_env
+
+  def observation(self, frames):
+    if not cv2:
+      return frames
+    return np.array([cv2.resize(
+        frame, (self.size, self.size), interpolation=cv2.INTER_AREA)
+                     for frame in frames])
+
+  def step(self, actions):
+    obs, rewards, dones = self.batch_env.step(actions)
+    obs = self.observation(obs)
+    return obs, rewards, dones
+
+  def reset(self, *args, **kwargs):
+    return self.observation(self.batch_env.reset(*args, **kwargs))
+
+  @property
+  def action_space(self):
+    return self.batch_env.action_space
+
+  @property
+  def batch_size(self):
+    return self.batch_env.batch_size
 
-  def create_env_fun(game_name, sticky_actions=True):
+  def close(self):
+    self.batch_env.close()
+
+
+class DopamineBatchEnv(object):
+  """Batch of environments.
+
+  Assumes that all given environments finishes at the same time.
+
+  Observations and rewards are returned as batches (arrays). Done is returned
+  as single boolean.
+  """
+
+  def __init__(self, batch_env, max_episode_steps):
+    self.batch_env = batch_env
+    self._max_episode_steps = max_episode_steps
+    self.game_over = None
+    self._elapsed_steps = 0
+
+  def reset(self):
+    self.game_over = False
+    self._elapsed_steps = 0
+    return np.array(self.batch_env.reset())
+
+  def step(self, actions):
+    """Step."""
+    self._elapsed_steps += 1
+    obs, rewards, dones = \
+        [np.array(r) for r in self.batch_env.step(actions)]
+    if self._elapsed_steps > self._max_episode_steps:
+      done = True
+      if self._elapsed_steps > self._max_episode_steps + 1:
+        rewards.fill(0)
+    else:
+      done = dones[0]
+      assert np.all(done == dones), ("Current modifications of Dopamine "
+                                     "require same number of steps for each "
+                                     "environment in batch")
+      del dones
+
+    self.game_over = done
+    return obs, rewards, done, {}
+
+  def render(self, mode):
+    pass
+
+  def close(self):
+    self.batch_env.close()
+
+  @property
+  def action_space(self):
+    return self.batch_env.action_space
+
+  @property
+  def batch_size(self):
+    return self.batch_env.batch_size
+
+
+class PaddedTrajectoriesEnv(DopamineBatchEnv):
+  """Pad finished episodes with zeros.
+
+  Allow episodes in batch to end on different timesteps, return zero
+  observations and rewards for finished ones. Return done=True when all
+  episodes are finished.
+
+  Note that output of this class might be misleading - the agent/evaluator
+  which uses this environment gets false information about when episodes have
+  ended. This class is used for informal check of Batched dopamine
+  implementation in model-free pipeline.
+  """
+
+  def reset(self):
+    self.done_envs = [False] * self.batch_size
+    self.game_over = False
+    self._elapsed_steps = 0
+    return np.array(self.batch_env.reset())
+
+  def step(self, actions):
+    if any(self.done_envs):
+      print("Warning, some environments already ended, using mocked data.")
+
+    self._elapsed_steps += 1
+    obs, rewards, dones = \
+        [np.array(r) for r in self.batch_env.step(actions)]
+    for i, ignore in enumerate(self.done_envs):
+      if ignore:
+        obs[i] = np.zeros(obs[i].shape, dtype=obs.dtype)
+        rewards[i] = 0
+      if dones[i]:
+        self.batch_env.reset([i])
+        self.done_envs[i] = True
+
+    all_done = all(self.done_envs)
+
+    if self._elapsed_steps > self._max_episode_steps:
+      all_done = True
+      if self._elapsed_steps > self._max_episode_steps + 1:
+        rewards.fill(0)
+
+    self.game_over = all_done
+    return obs, rewards, all_done, {}
+
+
+def get_create_batch_env_fun(batch_env_fn, time_limit):
+  """Factory for dopamine environment initialization function.
+
+  Args:
+    batch_env_fn: function(in_graph: bool) -> batch environment.
+    time_limit: time steps limit for environment.
+
+  Returns:
+    function (with optional, unused parameters) initializing environment.
+  """
+
+  def create_env_fun(game_name=None, sticky_actions=None):
     del game_name, sticky_actions
     batch_env = batch_env_fn(in_graph=False)
-    env = FlatBatchEnv(batch_env)
-    env = TimeLimit(env, max_episode_steps=time_limit)
-    env = ResizeObservation(env)  # pylint: disable=redefined-variable-type
-    env = GameOverOnDone(env)
-    return env
+    batch_env = ResizeBatchObservation(batch_env)  # pylint: disable=redefined-variable-type
+    batch_env = DopamineBatchEnv(batch_env, max_episode_steps=time_limit)
+    return batch_env
 
   return create_env_fun
 
 
 def _parse_hparams(hparams):
-  """TODO(konradczechowski): Add doc-string."""
+  """Split hparams, based on key prefixes.
+
+  Args:
+    hparams: hyperparameters
+
+  Returns:
+    Tuple of hparams for respectably: agent, optimizer, runner, replay_buffer.
+  """
   prefixes = ["agent_", "optimizer_", "runner_", "replay_buffer_"]
   ret = []
 
@@ -242,9 +500,8 @@ def _get_optimizer(params):
 class DQNLearner(PolicyLearner):
   """Interface for learning dqn implemented in dopamine."""
 
-  def __init__(self, frame_stack_size, base_event_dir, agent_model_dir):
-    super(DQNLearner, self).__init__(frame_stack_size, base_event_dir,
-                                     agent_model_dir)
+  def __init__(self, *args, **kwargs):
+    super(DQNLearner, self).__init__(*args, **kwargs)
     self.completed_iterations = 0
 
   def _target_iteractions_and_steps(self, num_env_steps, save_continuously,
@@ -269,10 +526,10 @@ def create_runner(self, env_fn, hparams, target_iterations,
     agent_params["optimizer"] = optimizer
     agent_params.update(replay_buffer_params)
     create_agent_fn = get_create_agent(agent_params)
-    runner = run_experiment.Runner(
+    runner = BatchRunner(
         base_dir=self.agent_model_dir,
         create_agent_fn=create_agent_fn,
-        create_environment_fn=get_create_env_fun(
+        create_environment_fn=get_create_batch_env_fun(
             env_fn, time_limit=hparams.time_limit),
         evaluation_steps=0,
         num_iterations=target_iterations,
@@ -290,9 +547,10 @@ def train(self,
             num_env_steps=None,
             env_step_multiplier=1,
             eval_env_fn=None,
-            report_fn=None):
+            report_fn=None,
+            model_save_fn=None):
     # TODO(konradczechowski): evaluation during training (with eval_env_fun)
-    del epoch, eval_env_fn, simulated, report_fn
+    del epoch, eval_env_fn, simulated, report_fn, model_save_fn
     if num_env_steps is None:
       num_env_steps = hparams.num_frames
 
@@ -305,13 +563,13 @@ def train(self,
       self._target_iteractions_and_steps(
           num_env_steps=num_env_steps * env_step_multiplier,
           save_continuously=save_continuously,
-          save_every_steps=hparams.save_every_steps,)
+          save_every_steps=hparams.save_every_steps)
 
     with tf.Graph().as_default():
       runner = self.create_runner(env_fn, hparams, target_iterations,
                                   training_steps_per_iteration)
       runner.run_experiment()
-
+      runner.close()
     self.completed_iterations = target_iterations
 
   def evaluate(self, env_fn, hparams, sampling_temp):
@@ -323,7 +581,7 @@ def evaluate(self, env_fn, hparams, sampling_temp):
         "agent_epsilon_eval", min(hparams.agent_epsilon_eval * sampling_temp, 1)
     )
 
-    create_environment_fn = get_create_env_fun(
+    create_environment_fn = get_create_batch_env_fun(
         env_fn, time_limit=hparams.time_limit)
     env = create_environment_fn(
         game_name="unused_arg", sticky_actions="unused_arg")
@@ -331,13 +589,13 @@ def evaluate(self, env_fn, hparams, sampling_temp):
     with tf.Graph().as_default():
       runner = self.create_runner(env_fn, hparams, target_iterations,
                                   training_steps_per_iteration)
+      assert runner.batch_size == 1
       agent = runner._agent  # pylint: disable=protected-access
+      runner.close()
       del runner
       agent.eval = True
 
-      # TODO(konradczechowski): correct number of episodes, when this will
-      # be hparam
-      for _ in range(30):
+      for _ in range(hparams.eval_episodes_num):
         # Run single episode
         ob = env.reset()
         action = agent.begin_episode(ob)
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index eac13e012..bfee48de0 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -294,6 +294,5 @@ def _video_reset_writer(self):
       self._video_writer.finish_to_disk()
     self._video_writer = None
 
-  def __del__(self):
+  def close(self):
     self._video_reset_writer()
-    super(SimulatedBatchEnv, self).__del__()
diff --git a/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index e0ffd9805..fd220837f 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -86,9 +86,6 @@ def reset(self, indices=None):
     if indices is None:
       indices = np.array(range(self.batch_size))
     obs = self._sess.run(self._reset_op, feed_dict={self._indices_t: indices})
-    # TODO(pmilos): remove if possible
-    # obs[:, 0, 0, 0] = 0
-    # obs[:, 0, 0, 1] = 255
     return obs
 
   def step(self, actions):
@@ -99,3 +96,4 @@ def step(self, actions):
 
   def close(self):
     self._sess.close()
+    self._batch_env.close()
diff --git a/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
index 82051f159..ec59e4c8a 100644
--- a/tensor2tensor/rl/policy_learner.py
+++ b/tensor2tensor/rl/policy_learner.py
@@ -45,8 +45,6 @@ def train(
       report_fn=None
   ):
     """Train."""
-    # TODO(konradczechowski): pass name_scope instead of epoch?
-    # TODO(konradczechowski): move 'simulated' to  batch_env
     raise NotImplementedError()
 
   def evaluate(self, env_fn, hparams, sampling_temp):
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index f34303d92..37b4504dc 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -39,6 +39,7 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl import trainer_model_based_params
+from tensor2tensor.rl.dopamine_connector import DQNLearner  # pylint: disable=unused-import
 from tensor2tensor.rl.restarter import Restarter
 from tensor2tensor.utils import trainer_lib
 
@@ -287,6 +288,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   metrics = {}
 
   # Collect data from the real environment.
+  policy_model_dir = directories["policy"]
   tf.logging.info("Initial training of the policy in real environment.")
   train_agent_real_env(env, learner, hparams, epoch)
   metrics["mean_reward/train/clipped"] = rl_utils.compute_mean_reward(
diff --git a/tensor2tensor/rl/trainer_model_based_dqn_test.py b/tensor2tensor/rl/trainer_model_based_dqn_test.py
deleted file mode 100644
index 50c056649..000000000
--- a/tensor2tensor/rl/trainer_model_based_dqn_test.py
+++ /dev/null
@@ -1,41 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Tiny run of trainer_model_based. Smoke test."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-# from tensor2tensor.rl import trainer_model_based
-
-import tensorflow as tf
-
-FLAGS = tf.flags.FLAGS
-
-
-class ModelRLExperimentTest(tf.test.TestCase):
-
-  def test_dqn_basic(self):
-    # TODO(afrozm): The latest changes in Dopamine break this test, so
-    # temporarily disabling this test.
-    pass
-    # FLAGS.output_dir = tf.test.get_temp_dir()
-    # FLAGS.loop_hparams_set = "rlmb_dqn_tiny"
-    # FLAGS.schedule = "train"  # skip evaluation for world model training
-    # trainer_model_based.main(None)
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 37eb47ddc..a60efef1e 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -111,6 +111,7 @@ def _rlmb_base():
       # This is only used for world-model evaluation currently, PolicyLearner
       # uses algorithm specific hparams to set this during training.
       simulated_rollout_length=50,
+      wm_policy_param_sharing=False,
 
       # To be overridden.
       base_algo="",
@@ -145,7 +146,6 @@ def rlmb_ppo_base():
       # Number of simulated environments to train on simultaneously.
       simulated_batch_size=16,
       eval_batch_size=32,
-      wm_policy_param_sharing=False,
 
       # Unused; number of PPO epochs is calculated from the real frame limit.
       real_ppo_epochs_num=0,
@@ -194,13 +194,17 @@ def rlmb_dqn_base():
       base_algo="dqn",
       base_algo_params="dqn_original_params",
       real_batch_size=1,
-      simulated_batch_size=1,
+      simulated_batch_size=16,
       dqn_agent_generates_trainable_dones=False,
       eval_batch_size=1,
       # Must be equal to dqn_time_limit for now
       simulated_rollout_length=simulated_rollout_length,
       dqn_time_limit=simulated_rollout_length,
       simulation_flip_first_random_for_beginning=False,
+      dqn_eval_episodes_num=3,
+
+      # TODO(kc): only for model-free compatibility, remove this
+      epochs_num=-1,
   )
   update_hparams(hparams, dqn_params)
   return hparams
diff --git a/tensor2tensor/rl/trainer_model_based_test.py b/tensor2tensor/rl/trainer_model_based_test.py
index 0c836162c..90e6ad260 100644
--- a/tensor2tensor/rl/trainer_model_based_test.py
+++ b/tensor2tensor/rl/trainer_model_based_test.py
@@ -27,12 +27,19 @@
 
 class ModelRLExperimentTest(tf.test.TestCase):
 
-  def test_basic(self):
+  def _test_hparams_skip_evaluation(self, hparams_set):
     FLAGS.output_dir = tf.test.get_temp_dir()
-    FLAGS.loop_hparams_set = "rlmb_tiny"
+    FLAGS.loop_hparams_set = hparams_set
     FLAGS.schedule = "train"  # skip evaluation for world model training
     trainer_model_based.main(None)
 
+  def test_basic(self):
+    self._test_hparams_skip_evaluation("rlmb_tiny")
+
+  # TODO(kozak): enable when it works.
+  # def test_dqn_basic(self):
+  #   self._test_hparams_skip_evaluation("rlmb_dqn_tiny")
+
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 5d7a0974a..0d48a372d 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -103,41 +103,51 @@ def train(hparams, output_dir, env_problem_name, report_fn=None):
   tf.logging.vlog(1, "Policy HParams : %s",
                   misc_utils.pprint_hparams(policy_hparams))
 
-  total_steps = policy_hparams.epochs_num
-  tf.logging.vlog(2, "total_steps: %d", total_steps)
-
-  eval_every_epochs = policy_hparams.eval_every_epochs
-  tf.logging.vlog(2, "eval_every_epochs: %d", eval_every_epochs)
-
-  if eval_every_epochs == 0:
-    eval_every_epochs = total_steps
-  policy_hparams.eval_every_epochs = 0
-
-  metric_name = rl_utils.get_metric_name(
-      sampling_temp=hparams.eval_sampling_temps[0],
-      max_num_noops=hparams.eval_max_num_noops,
-      clipped=False
-  )
+  # TODO(konradczechowski): remove base_algo dependance, when evaluation method
+  # will be decided
+  if hparams.base_algo == "ppo":
+    print("\n\n\npolicy_hparams {}\n\n\n".format(policy_hparams))
+    total_steps = policy_hparams.epochs_num
+    tf.logging.vlog(2, "total_steps: %d", total_steps)
+
+    eval_every_epochs = policy_hparams.eval_every_epochs
+    tf.logging.vlog(2, "eval_every_epochs: %d", eval_every_epochs)
+
+    if eval_every_epochs == 0:
+      eval_every_epochs = total_steps
+    policy_hparams.eval_every_epochs = 0
+
+    metric_name = rl_utils.get_metric_name(
+        sampling_temp=hparams.eval_sampling_temps[0],
+        max_num_noops=hparams.eval_max_num_noops,
+        clipped=False
+    )
+
+    tf.logging.vlog(1, "metric_name: %s", metric_name)
+
+    eval_metrics_dir = os.path.join(output_dir, "eval_metrics")
+    eval_metrics_dir = os.path.expanduser(eval_metrics_dir)
+    tf.gfile.MakeDirs(eval_metrics_dir)
+    eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
+
+    def evaluate_on_new_model(model_dir_path):
+      global step
+      eval_metrics = rl_utils.evaluate_all_configs(hparams, model_dir_path)
+      tf.logging.info(
+          "Agent eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
+      rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, step)
+      if report_fn:
+        report_fn(eval_metrics[metric_name], step)
+      step += 1
+
+    policy_hparams.epochs_num = total_steps
+    policy_hparams.save_models_every_epochs = eval_every_epochs
+  else:
+    def evaluate_on_new_model(model_dir_path):
+      del model_dir_path
+      raise NotImplementedError(
+          "This function is currently implemented only for ppo")
 
-  tf.logging.vlog(1, "metric_name: %s", metric_name)
-
-  eval_metrics_dir = os.path.join(output_dir, "eval_metrics")
-  eval_metrics_dir = os.path.expanduser(eval_metrics_dir)
-  tf.gfile.MakeDirs(eval_metrics_dir)
-  eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
-
-  def evaluate_on_new_model(model_dir_path):
-    global step
-    eval_metrics = rl_utils.evaluate_all_configs(hparams, model_dir_path)
-    tf.logging.info(
-        "Agent eval metrics:\n{}".format(pprint.pformat(eval_metrics)))
-    rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, step)
-    if report_fn:
-      report_fn(eval_metrics[metric_name], step)
-    step += 1
-
-  policy_hparams.epochs_num = total_steps
-  policy_hparams.save_models_every_epochs = eval_every_epochs
   learner.train(env_fn,
                 policy_hparams,
                 simulated=False,
diff --git a/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index 98c28feca..3e21a1387 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -28,16 +28,18 @@
 
 class TrainTest(tf.test.TestCase):
 
-  def test_train_pong(self):
-    hparams = registry.hparams("rlmf_original")
-    hparams.batch_size = 2
-    hparams.eval_sampling_temps = [0.0, 1.0]
-    hparams.add_hparam("ppo_epochs_num", 2)
-    hparams.add_hparam("ppo_epoch_length", 3)
+  def _test_hparams_set(self, hparams_set):
+    hparams = registry.hparams(hparams_set)
     FLAGS.output_dir = tf.test.get_temp_dir()
     trainer_model_free.train(hparams, FLAGS.output_dir,
                              env_problem_name=None)
 
+  def test_train_pong(self):
+    self._test_hparams_set("rlmf_tiny")
+
+  def test_train_pong_dqn(self):
+    self._test_hparams_set("rlmf_dqn_tiny")
+
 
 if __name__ == "__main__":
   tf.test.main()

From 8521bf71559f47ac7e4667e797266802b389745a Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 25 Mar 2019 16:58:33 -0700
Subject: [PATCH 1835/2720] Minor change to README.md

PiperOrigin-RevId: 240247338
---
 README.md           | 2 +-
 docs/walkthrough.md | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 9e3d1491d..a104f1bdc 100644
--- a/README.md
+++ b/README.md
@@ -475,4 +475,4 @@ T2T](https://research.googleblog.com/2017/06/accelerating-deep-learning-research
 * [The Evolved Transformer](https://arxiv.org/abs/1901.11117)
 * [Model-Based Reinforcement Learning for Atari](https://arxiv.org/abs/1903.00374)
 
-*Note: This is not an official Google product.*
+*NOTE: This is not an official Google product.*
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 9e3d1491d..a104f1bdc 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -475,4 +475,4 @@ T2T](https://research.googleblog.com/2017/06/accelerating-deep-learning-research
 * [The Evolved Transformer](https://arxiv.org/abs/1901.11117)
 * [Model-Based Reinforcement Learning for Atari](https://arxiv.org/abs/1903.00374)
 
-*Note: This is not an official Google product.*
+*NOTE: This is not an official Google product.*

From 36894c6ff0553ec0e182af485cedc96a4c6196df Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 26 Mar 2019 09:51:55 -0700
Subject: [PATCH 1836/2720] pass tpu_job_name to TPUConfig.

PiperOrigin-RevId: 240367576
---
 tensor2tensor/bin/t2t_trainer.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index c82043b3a..586f56331 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -56,6 +56,9 @@
                     "available to the t2t-trainer.")
 flags.DEFINE_integer("random_seed", None, "Random seed.")
 flags.DEFINE_integer("tpu_num_shards", 8, "Number of tpu shards.")
+flags.DEFINE_string("tpu_job_name", None,
+                    "TPU job name. TPUEstimator can auto-infer this but if the "
+                    "configuration is esoteric it should be provided here.")
 flags.DEFINE_integer("iterations_per_loop", 100,
                      "Number of iterations in a TPU training loop.")
 flags.DEFINE_bool("use_tpu", False, "Whether to use TPU.")
@@ -224,6 +227,8 @@ def create_run_config(hp, output_dir=None):
     save_ckpt_steps = None
   assert FLAGS.output_dir or FLAGS.checkpoint_path
   tpu_config_extra_kwargs = {}
+  if FLAGS.tpu_job_name is not None:
+    tpu_config_extra_kwargs["tpu_job_name"] = FLAGS.tpu_job_name
 
   if getattr(hp, "mtf_mode", False):
     save_ckpt_steps = None  # Disable the default saver

From 9fa015b586853d0142519d78980057eb353b6e78 Mon Sep 17 00:00:00 2001
From: Piotr Milos <piotr.milos@codilime.com>
Date: Tue, 26 Mar 2019 23:03:25 +0100
Subject: [PATCH 1837/2720] removing datasets for serving data (#1511)

---
 tensor2tensor/rl/ppo.py                       | 44 ++++++++++---------
 .../rl/trainer_model_based_params.py          |  2 +-
 2 files changed, 25 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index a40a46f59..b393aa4c5 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -95,32 +95,36 @@ def define_ppo_epoch(memory, hparams, action_space, batch_size):
   add_lists_elementwise = lambda l1, l2: [x + y for x, y in zip(l1, l2)]
 
   number_of_batches = ((hparams.epoch_length-1) * hparams.optimization_epochs
-                       / hparams.optimization_batch_size)
+                       // hparams.optimization_batch_size)
 
+  epoch_length = hparams.epoch_length
   if hparams.effective_num_agents is not None:
     number_of_batches *= batch_size
-    number_of_batches /= hparams.effective_num_agents
-
-  dataset = tf.data.Dataset.from_tensor_slices(
-      (observation[:-1], action[:-1], discounted_reward, advantage_normalized,
-       old_pdf[:-1]))
-  dataset = dataset.shuffle(buffer_size=hparams.epoch_length-1,
-                            reshuffle_each_iteration=True)
-  dataset = dataset.repeat(-1)
-  dataset = dataset.batch(hparams.optimization_batch_size, drop_remainder=True)
-  iterator = dataset.make_initializable_iterator()
+    number_of_batches //= hparams.effective_num_agents
+    epoch_length //= hparams.effective_num_agents
 
+  assert number_of_batches > 0, "Set the paremeters so that number_of_batches>0"
   lr = learning_rate.learning_rate_schedule(hparams)
 
-  with tf.control_dependencies([iterator.initializer]):
-    ppo_step_rets = tf.scan(
-        lambda a, i: add_lists_elementwise(  # pylint: disable=g-long-lambda
-            a, define_ppo_step(
-                iterator.get_next(), hparams, action_space, lr
-            )),
-        tf.range(number_of_batches),
-        [0., 0., 0.],
-        parallel_iterations=1)
+  shuffled_indices = [tf.random.shuffle(tf.range(epoch_length - 1))
+                      for _ in range(hparams.optimization_epochs)]
+  shuffled_indices = tf.concat(shuffled_indices, axis=0)
+  shuffled_indices = shuffled_indices[:number_of_batches *
+                                      hparams.optimization_batch_size]
+  indices_of_batches = tf.reshape(shuffled_indices,
+                                  shape=(-1, hparams.optimization_batch_size))
+  input_tensors = [observation, action, discounted_reward,
+                   advantage_normalized, old_pdf]
+
+  ppo_step_rets = tf.scan(
+      lambda a, i: add_lists_elementwise(  # pylint: disable=g-long-lambda
+          a, define_ppo_step([tf.gather(t, indices_of_batches[i, :])
+                              for t in input_tensors],
+                             hparams, action_space, lr
+                            )),
+      tf.range(number_of_batches),
+      [0., 0., 0.],
+      parallel_iterations=1)
 
   ppo_summaries = [tf.reduce_mean(ret) / number_of_batches
                    for ret in ppo_step_rets]
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index a60efef1e..515282678 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -521,7 +521,7 @@ def rlmb_ppo_tiny():
   hparams = hparams.override_from_dict(_rlmb_tiny_overrides())
   update_hparams(hparams, dict(
       ppo_epochs_num=2,
-      ppo_epoch_length=hparams.simulated_rollout_length,
+      ppo_epoch_length=10,
       real_ppo_epoch_length=36,
       real_ppo_effective_num_agents=2,
       real_batch_size=1,

From b2cc9f219269aec8dab8c54122610506f16670cc Mon Sep 17 00:00:00 2001
From: Piotr Milos <piotr.milos@codilime.com>
Date: Tue, 26 Mar 2019 15:03:44 -0700
Subject: [PATCH 1838/2720] Merge of PR #1511

PiperOrigin-RevId: 240435415
---
 tensor2tensor/models/research/rl.py                 | 13 +++++++++++--
 .../models/video/basic_deterministic_params.py      |  2 +-
 tensor2tensor/rl/ppo.py                             |  1 -
 tensor2tensor/rl/trainer_model_free.py              |  3 ---
 4 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 6025a132c..7e0ac4455 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -134,10 +134,19 @@ def ppo_original_params():
   return hparams
 
 
+@registry.register_hparams
+def ppo_original_tiny():
+  """Parameters based on the original PPO paper, tiny version."""
+  hparams = ppo_original_params()
+  hparams.epoch_length = 5
+  hparams.optimization_batch_size = 1
+  return hparams
+
+
 @registry.register_hparams
 def ppo_ttt_params():
   """Parameters based on the original PPO paper."""
-  hparams = ppo_original_params()
+  hparams = ppo_original_tiny()
   hparams.policy_network = "feed_forward_categorical_policy"
   hparams.policy_problem_name = "dummy_policy_problem_ttt"
   return hparams
@@ -431,7 +440,6 @@ def rlmf_tictactoe():
 
   # Number of last observations to feed to the agent
   hparams.frame_stack_size = 1
-
   return hparams
 
 
@@ -450,6 +458,7 @@ def rlmf_tiny():
   hparams = rlmf_original()
   hparams = hparams.override_from_dict(rlmf_tiny_overrides())
   hparams.batch_size = 2
+  hparams.base_algo_params = "ppo_original_tiny"
   hparams.add_hparam("ppo_epochs_num", 3)
   hparams.add_hparam("ppo_epoch_length", 2)
   return hparams
diff --git a/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index b40307f7b..cebe148c9 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -50,7 +50,7 @@ def next_frame_basic_deterministic():
   hparams.add_hparam("do_autoregressive_rnn", False)
   hparams.add_hparam("autoregressive_rnn_lookback", 8)
   hparams.add_hparam("autoregressive_rnn_warmup_steps", 8000)
-  hparams.add_hparam("activation_fn", "belu")
+  hparams.add_hparam("activation_fn", "relu")
   hparams.bottom["inputs"] = modalities.video_identity_bottom
   hparams.bottom["targets"] = modalities.video_identity_bottom
   return hparams
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index b393aa4c5..249be9e0f 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -96,7 +96,6 @@ def define_ppo_epoch(memory, hparams, action_space, batch_size):
 
   number_of_batches = ((hparams.epoch_length-1) * hparams.optimization_epochs
                        // hparams.optimization_batch_size)
-
   epoch_length = hparams.epoch_length
   if hparams.effective_num_agents is not None:
     number_of_batches *= batch_size
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 0d48a372d..fdcc6ec3b 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -88,14 +88,12 @@ def train(hparams, output_dir, env_problem_name, report_fn=None):
 
   tf.logging.vlog(1, "HParams in trainer_model_free.train : %s",
                   misc_utils.pprint_hparams(hparams))
-
   tf.logging.vlog(1, "Using hparams.base_algo: %s", hparams.base_algo)
   learner = rl_utils.LEARNERS[hparams.base_algo](
       hparams.frame_stack_size, output_dir, output_dir, total_num_epochs=1
   )
 
   policy_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
-
   rl_utils.update_hparams_from_hparams(
       policy_hparams, hparams, hparams.base_algo + "_"
   )
@@ -106,7 +104,6 @@ def train(hparams, output_dir, env_problem_name, report_fn=None):
   # TODO(konradczechowski): remove base_algo dependance, when evaluation method
   # will be decided
   if hparams.base_algo == "ppo":
-    print("\n\n\npolicy_hparams {}\n\n\n".format(policy_hparams))
     total_steps = policy_hparams.epochs_num
     tf.logging.vlog(2, "total_steps: %d", total_steps)
 

From 7561ead411fcd4328de0e65d1e24bacf101747ed Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 26 Mar 2019 16:20:16 -0700
Subject: [PATCH 1839/2720] Decouple recurrent memory size from chunk size

PiperOrigin-RevId: 240450826
---
 tensor2tensor/layers/transformer_memory.py | 75 +++++++++++++++++-----
 tensor2tensor/models/transformer.py        |  9 ++-
 2 files changed, 66 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index 2d3696eac..1f78213e8 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -24,24 +24,64 @@
 class RecurrentMemory(object):
   """Base class for recurrent memory.
 
-  Currently implements memory in the style of Transformer-XL
+  This class defines the memory interface, but behaves like a no-op.
+  """
+
+  def pre_attention(self, segment, query_antecedent, memory_antecedent, bias):
+    """Called prior to self-attention, to incorporate memory items.
+
+    Args:
+      segment: an integer Tensor with shape [batch]
+      query_antecedent: a Tensor with shape [batch, length_q, channels]
+      memory_antecedent: must be None. Attention normally allows this to be a
+        Tensor with shape [batch, length_m, channels], but we currently only
+        support memory for decoder-side self-attention.
+      bias: bias Tensor (see attention_bias())
+    Returns:
+      (data, new_query_antecedent, new_memory_antecedent, new_bias)
+    """
+    del segment
+    return None, query_antecedent, memory_antecedent, bias
+
+  def post_attention(self, token, x):
+    """Called after self-attention. The memory can be updated here.
+
+    Args:
+      token: Data returned by pre_attention, which can be used to carry over
+        state related to the current memory operation.
+      x: a Tensor of data after self-attention and feed-forward
+    Returns:
+      a (possibly modified) version of the input x
+    """
+    assert token is None
+    return x
+
+
+class RecentTokensMemory(RecurrentMemory):
+  """A memory module that caches features for recent tokens.
+
+  When the number of tokens cached is equal to the chunk size, this is
+  equivalent to the memory used by Transformer-XL
   (https://arxiv.org/abs/1901.02860)
   """
-  # TODO(kitaev): make this a base class and then subclass for different memory
-  # types (e.g. the one defined below in this file).
 
   def __init__(self, name, hparams):
     hidden_size = hparams.hidden_size
-    chunk_length = hparams.split_targets_chunk_length
-    assert chunk_length > 0, "Chunking is required to use RecurrentMemory"
+    self.chunk_length = hparams.split_targets_chunk_length
+    assert self.chunk_length > 0, "Chunking is required to use recurrent memory"
+
+    if hasattr(hparams, "num_memory_items") and hparams.num_memory_items > 0:
+      self.tokens_to_cache = hparams.num_memory_items
+    else:
+      self.tokens_to_cache = self.chunk_length
 
     # TODO(kitaev): The implementation of the chunking code makes it somewhat
     # convoluted to figure out how many actual sequences we can have per batch.
     # The data pipeline should be revisited at some point.
     batch_size_in_sequences = hparams.batch_size / hparams.max_length
 
-    memory_shape = [batch_size_in_sequences, chunk_length, hidden_size]
-    bias_shape = [1, 1, chunk_length, chunk_length]
+    memory_shape = [batch_size_in_sequences, self.tokens_to_cache, hidden_size]
+    bias_shape = [1, 1, self.chunk_length, self.tokens_to_cache]
 
     with tf.variable_scope(name):
       self.previous_segment = tf.get_variable(
@@ -74,21 +114,18 @@ def pre_attention(self, segment, query_antecedent, memory_antecedent, bias):
     """
     assert memory_antecedent is None, "We only support language modeling"
 
-    previous_vals = tf.stop_gradient(self.previous_vals)
+    previous_vals = self.previous_vals
     # If segment id is zero, don't attend back to the memory
-    previous_bias = tf.stop_gradient(self.previous_bias) + tf.cast(
+    previous_bias = self.previous_bias + tf.cast(
         tf.equal(tf.reduce_sum(segment), 0), tf.float32) * -1e9
 
     # In eval mode, batch size may be variable
     amount_to_pad = tf.shape(previous_vals)[0] - tf.shape(query_antecedent)[0]
-    previous_vals = previous_vals[:tf.shape(query_antecedent)[0], :, :]
-    with tf.control_dependencies(
-        [tf.assert_equal(tf.shape(query_antecedent), tf.shape(previous_vals))]):
-      query_antecedent = tf.identity(query_antecedent)
+    sliced_previous_vals = previous_vals[:tf.shape(query_antecedent)[0], :, :]
 
     new_memory_antecedent = tf.concat(
-        [tf.stop_gradient(previous_vals), query_antecedent], 1)
-    new_bias = tf.concat([previous_bias, bias], -1)
+        [tf.stop_gradient(sliced_previous_vals), query_antecedent], 1)
+    new_bias = tf.concat([tf.stop_gradient(previous_bias), bias], -1)
 
     remember_segment = segment[0]
     # TODO(kitaev): The code assumes that we always either increment the chunk
@@ -101,10 +138,18 @@ def pre_attention(self, segment, query_antecedent, memory_antecedent, bias):
             tf.equal(remember_segment, self.previous_segment + 1)),
                    [self.previous_segment, remember_segment])]):
       remember_segment = tf.identity(remember_segment)
+
     remember_vals = tf.pad(query_antecedent,
                            [[0, amount_to_pad], [0, 0], [0, 0]])
     remember_bias = tf.zeros_like(bias) + tf.reduce_max(
         bias, -1, keep_dims=True)
+    # Assume that query_antecedent is always a full chunk (i.e. not truncated)
+    if self.chunk_length < self.tokens_to_cache:
+      remember_vals = tf.concat([previous_vals, remember_vals], 1)
+      remember_bias = tf.concat([previous_bias, remember_bias], -1)
+    if self.chunk_length != self.tokens_to_cache:
+      remember_vals = remember_vals[:, -self.tokens_to_cache:, :]
+      remember_bias = remember_bias[:, :, :, -self.tokens_to_cache:]
     token = (remember_segment, remember_vals, remember_bias)
 
     return token, query_antecedent, new_memory_antecedent, new_bias
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 0dfdca3ad..d63ed85c1 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1531,9 +1531,9 @@ def __init__(self, *args, **kwargs):
     self.recurrent_memory_by_layer = {}
     for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
       layer_name = "layer_%d" % layer
-      self.recurrent_memory_by_layer[layer_name] = transformer_memory.RecurrentMemory(
-          layer_name + "/recurrent_memory", hparams)
-
+      self.recurrent_memory_by_layer[
+          layer_name] = transformer_memory.RecentTokensMemory(
+              layer_name + "/recurrent_memory", hparams)
 
   def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha,
                    use_tpu=False):
@@ -2641,6 +2641,9 @@ def transformer_wikitext103_l4k_memory_v0():
   hparams.self_attention_type = "dot_product_relative_memory"
   hparams.max_relative_position = 2 * hparams.split_targets_chunk_length
 
+  # By default, cache one chunk only (like Transformer-XL)
+  hparams.add_hparam("num_memory_items", hparams.split_targets_chunk_length)
+
   return hparams
 
 

From 0a251efcf8c95c0c6d7a457ee7d0809ddb254997 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 27 Mar 2019 07:40:41 -0700
Subject: [PATCH 1840/2720] CIFAR-10 flat subpixel generation

PiperOrigin-RevId: 240555453
---
 tensor2tensor/data_generators/cifar.py     | 22 ++++++++++++++
 tensor2tensor/layers/transformer_memory.py |  5 +++-
 tensor2tensor/models/transformer.py        | 35 ++++++++++++++++++++++
 3 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index 92bcb97da..7f895663e 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -29,6 +29,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import mnist
+from tensor2tensor.data_generators import problem
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
@@ -179,6 +180,27 @@ def preprocess_example(self, example, mode, unused_hparams):
     return example
 
 
+@registry.register_problem
+class ImageCifar10PlainGenFlat(ImageCifar10PlainGen):
+  """CIFAR-10 for image generation as a flat array of 64*64*3=12228 elements."""
+
+  def preprocess_example(self, example, mode, unused_hparams):
+    example["inputs"].set_shape([_CIFAR10_IMAGE_SIZE, _CIFAR10_IMAGE_SIZE, 3])
+    example["inputs"] = tf.to_int64(example["inputs"])
+    example["inputs"] = tf.reshape(example["inputs"], (-1,))
+
+    del example["targets"]  # Ensure unconditional generation
+
+    return example
+
+  def hparams(self, defaults, model_hparams):
+    super(ImageCifar10PlainGenFlat, self).hparams(defaults, model_hparams)
+    # Switch to symbol modality
+    p = defaults
+    p.modality["inputs"] = modalities.ModalityType.SYMBOL
+    p.input_space_id = problem.SpaceID.GENERIC
+
+
 @registry.register_problem
 class ImageCifar10PlainRandomShift(ImageCifar10Plain):
   """CIFAR-10 32x32 for image generation with random shift data-augmentation."""
diff --git a/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index 1f78213e8..2e815ddec 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -78,7 +78,10 @@ def __init__(self, name, hparams):
     # TODO(kitaev): The implementation of the chunking code makes it somewhat
     # convoluted to figure out how many actual sequences we can have per batch.
     # The data pipeline should be revisited at some point.
-    batch_size_in_sequences = hparams.batch_size / hparams.max_length
+    if hasattr(hparams, "recurrent_memory_batch_size"):
+      batch_size_in_sequences = hparams.recurrent_memory_batch_size
+    else:
+      batch_size_in_sequences = hparams.batch_size / hparams.max_length
 
     memory_shape = [batch_size_in_sequences, self.tokens_to_cache, hidden_size]
     bias_shape = [1, 1, self.chunk_length, self.tokens_to_cache]
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index d63ed85c1..0cba09701 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1535,6 +1535,12 @@ def __init__(self, *args, **kwargs):
           layer_name] = transformer_memory.RecentTokensMemory(
               layer_name + "/recurrent_memory", hparams)
 
+  @property
+  def has_input(self):
+    if hasattr(self._hparams, "unconditional") and self._hparams.unconditional:
+      return False
+    return super(TransformerMemory, self).has_input
+
   def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha,
                    use_tpu=False):
     """Overriding beam search because for now only the slow version works with
@@ -2641,6 +2647,8 @@ def transformer_wikitext103_l4k_memory_v0():
   hparams.self_attention_type = "dot_product_relative_memory"
   hparams.max_relative_position = 2 * hparams.split_targets_chunk_length
 
+  hparams.add_hparam("unconditional", True)
+  hparams.add_hparam("recurrent_memory_batch_size", 0)  # 0 = try to guess
   # By default, cache one chunk only (like Transformer-XL)
   hparams.add_hparam("num_memory_items", hparams.split_targets_chunk_length)
 
@@ -2666,3 +2674,30 @@ def transformer_wikitext103_l16k_memory_v0():
   hparams.max_relative_position = 2 * hparams.split_targets_chunk_length
 
   return hparams
+
+
+@registry.register_hparams
+def transformer_cifar10_memory_v0():
+  """HParams for training image_cifar10_plain_gen_flat_rev with memory."""
+  hparams = transformer_wikitext103_l4k_memory_v0()
+
+  hparams.num_hidden_layers = 6
+
+  hparams.max_length = 32 * 32 * 3
+  hparams.split_targets_chunk_length = 64 * 3
+  hparams.split_targets_max_chunks = int(
+      hparams.max_length / hparams.split_targets_chunk_length)
+  hparams.num_memory_items = 128 * 3
+
+  # Since this is an image problem, batch size refers to examples (not tokens)
+  target_images_per_batch = 4
+  hparams.batch_size = int(target_images_per_batch * (
+      hparams.max_length / hparams.split_targets_chunk_length))
+
+  # The recurrent memory needs to know the actual batch size (in sequences)
+  hparams.recurrent_memory_batch_size = hparams.batch_size
+
+  hparams.max_relative_position = (
+      hparams.num_memory_items + hparams.split_targets_chunk_length)
+
+  return hparams

From 211c8245bb4303a6a2519fa570ae170b26c99801 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 27 Mar 2019 14:40:16 -0700
Subject: [PATCH 1841/2720] enable concrete models to override the default tpu
 host call.

PiperOrigin-RevId: 240643874
---
 tensor2tensor/utils/t2t_model.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 010f619d1..8a1a1c7b0 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1532,6 +1532,9 @@ def estimator_model_fn(cls,
   def initialize_from_ckpt(self, ckpt_dir):
     return initialize_from_ckpt(ckpt_dir=ckpt_dir, hparams=self._hparams)
 
+  def create_host_call(self):
+    return create_host_call(self.hparams.model_dir)
+
   def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
     """Constructs `tf.estimator.EstimatorSpec` for TRAIN (training) mode."""
     train_op = self.optimize(loss, num_async_replicas=num_async_replicas,
@@ -1547,7 +1550,7 @@ def scaffold_fn():
 
       # Note: important to call this before remove_summaries()
       if self.hparams.tpu_enable_host_call:
-        host_call = create_host_call(self.hparams.model_dir)
+        host_call = self.create_host_call()
       else:
         host_call = None
 
@@ -1582,7 +1585,14 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
     problem = hparams.problem
 
     if common_layers.is_xla_compiled():
+      # Note: important to call this before remove_summaries()
+      if self.hparams.tpu_enable_host_call:
+        host_call = self.create_host_call()
+      else:
+        host_call = None
+
       remove_summaries()
+
       eval_metrics_fn = create_tpu_eval_metrics_fn(problem, hparams)
 
       batch_size = [feature.shape.as_list()[0] for _, feature
@@ -1606,6 +1616,7 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
       return tf.contrib.tpu.TPUEstimatorSpec(
           tf.estimator.ModeKeys.EVAL,
           eval_metrics=(eval_metrics_fn, eval_metrics_fn_flat_args),
+          host_call=host_call,
           loss=loss)
     else:
       task_list = [problem]

From 2d2d160c4773e38ecdac03d9862b2a90e0170ef6 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 27 Mar 2019 17:47:59 -0700
Subject: [PATCH 1842/2720] Correct flat CIFAR modality to not consider 0 as
 padding.

PiperOrigin-RevId: 240682373
---
 tensor2tensor/data_generators/cifar.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index 7f895663e..22bdfc831 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -197,7 +197,7 @@ def hparams(self, defaults, model_hparams):
     super(ImageCifar10PlainGenFlat, self).hparams(defaults, model_hparams)
     # Switch to symbol modality
     p = defaults
-    p.modality["inputs"] = modalities.ModalityType.SYMBOL
+    p.modality["inputs"] = modalities.ModalityType.SYMBOL_WEIGHTS_ALL
     p.input_space_id = problem.SpaceID.GENERIC
 
 

From 6be7d5c83fb6c34f2ec73c02fb1d689899cfd1d0 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 28 Mar 2019 10:17:16 -0700
Subject: [PATCH 1843/2720] Add a basic type of hard attention to Transformer;
 set hparams="hard_attention_k=16" to try.

PiperOrigin-RevId: 240798075
---
 tensor2tensor/layers/common_attention.py      |  50 +++++++--
 tensor2tensor/layers/common_attention_test.py |  17 +++
 tensor2tensor/layers/common_layers.py         | 106 ++++++++++++------
 tensor2tensor/layers/common_layers_test.py    |  14 +++
 tensor2tensor/layers/transformer_layers.py    |   3 +-
 tensor2tensor/models/transformer.py           |   6 +-
 6 files changed, 150 insertions(+), 46 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 73735b408..e1b926c9e 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -34,8 +34,10 @@
 import tensorflow as tf
 import tensorflow_probability as tfp
 
+# pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.framework import function
 from tensorflow.python.ops import inplace_ops
+# pylint: enable=g-direct-tensorflow-import
 
 
 # TODO(lukaszkaiser): remove this function when not needed any more.
@@ -1468,6 +1470,20 @@ def grouped_attention_multihead(query_antecedent,
     return o, extra_loss
 
 
+def harden_attention_weights(weights, hard_attention_k):
+  """Make attention weights non-0 only on the top-hard_attention_k ones."""
+  # Subtract the top-kth weight and zero-out all lower ones.
+  # Note that currently in case of numerical ties it will retain more
+  # than k elements. In the future, we may want to avoid this.
+  weights -= common_layers.top_kth_iterative(weights, hard_attention_k)
+  weights = tf.nn.relu(weights)
+  # Re-normalize the weights.
+  weights_sum = tf.reduce_sum(weights, axis=-1, keep_dims=True)
+  weights_sum = tf.maximum(weights_sum, 1e-6)  # Avoid division by 0.
+  weights /= weights_sum
+  return weights
+
+
 def dot_product_attention(q,
                           k,
                           v,
@@ -1479,7 +1495,8 @@ def dot_product_attention(q,
                           save_weights_to=None,
                           dropout_broadcast_dims=None,
                           activation_dtype=None,
-                          weight_dtype=None):
+                          weight_dtype=None,
+                          hard_attention_k=0):
   """Dot-product attention.
 
   Args:
@@ -1502,6 +1519,7 @@ def dot_product_attention(q,
     activation_dtype: Used to define function activation dtype when using
       mixed precision.
     weight_dtype: The dtype weights are stored in when using mixed precision
+    hard_attention_k: integer, if > 0 triggers hard attention (picking top-k)
 
   Returns:
     Tensor with shape [..., length_q, depth_v].
@@ -1515,6 +1533,8 @@ def dot_product_attention(q,
     # If logits are fp16, upcast before softmax
     logits = maybe_upcast(logits, activation_dtype, weight_dtype)
     weights = tf.nn.softmax(logits, name="attention_weights")
+    if hard_attention_k > 0:
+      weights = harden_attention_weights(weights, hard_attention_k)
     weights = common_layers.cast_like(weights, q)
     if save_weights_to is not None:
       save_weights_to[scope.name] = weights
@@ -1602,7 +1622,8 @@ def dot_product_attention_relative(q,
                                    save_weights_to=None,
                                    name=None,
                                    make_image_summary=True,
-                                   cache=False):
+                                   cache=False,
+                                   hard_attention_k=0):
   """Calculate relative position-aware dot-product self-attention.
 
   The attention calculation is augmented with learned representations for the
@@ -1623,6 +1644,7 @@ def dot_product_attention_relative(q,
     name: an optional string.
     make_image_summary: Whether to make an attention image summary.
     cache: whether use cache mode
+    hard_attention_k: integer, if > 0 triggers hard attention (picking top-k)
 
   Returns:
     A Tensor.
@@ -1658,6 +1680,8 @@ def dot_product_attention_relative(q,
     if bias is not None:
       logits += bias
     weights = tf.nn.softmax(logits, name="attention_weights")
+    if hard_attention_k > 0:
+      weights = harden_attention_weights(weights, hard_attention_k)
     if save_weights_to is not None:
       save_weights_to[scope.name] = weights
       save_weights_to[scope.name + "/logits"] = logits
@@ -3970,6 +3994,7 @@ def multihead_attention(query_antecedent,
                         layer_collection=None,
                         recurrent_memory=None,
                         chunk_number=None,
+                        hard_attention_k=0,
                         **kwargs):
   """Multihead scaled-dot-product attention with input/output transformations.
 
@@ -4027,7 +4052,8 @@ def multihead_attention(query_antecedent,
       retains state across chunks. Default is None.
     chunk_number: an optional integer Tensor with shape [batch] used to operate
       the recurrent_memory.
-    **kwargs (dict): Parameters for the attention function
+    hard_attention_k: integer, if > 0 triggers hard attention (picking top-k).
+    **kwargs (dict): Parameters for the attention function.
 
   Caching:
     WARNING: For decoder self-attention, i.e. when memory_antecedent == None,
@@ -4149,11 +4175,13 @@ def multihead_attention(query_antecedent,
       if isinstance(x, tuple):
         x, additional_returned_value = x  # Unpack
     elif attention_type == "dot_product":
-      x = dot_product_attention(q, k, v, bias, dropout_rate, image_shapes,
-                                save_weights_to=save_weights_to,
-                                make_image_summary=make_image_summary,
-                                dropout_broadcast_dims=dropout_broadcast_dims,
-                                activation_dtype=kwargs.get("activation_dtype"))
+      x = dot_product_attention(
+          q, k, v, bias, dropout_rate, image_shapes,
+          save_weights_to=save_weights_to,
+          make_image_summary=make_image_summary,
+          dropout_broadcast_dims=dropout_broadcast_dims,
+          activation_dtype=kwargs.get("activation_dtype"),
+          hard_attention_k=hard_attention_k)
     elif attention_type == "dot_product_relative":
       x = dot_product_attention_relative(
           q,
@@ -4165,7 +4193,8 @@ def multihead_attention(query_antecedent,
           image_shapes,
           save_weights_to=save_weights_to,
           make_image_summary=make_image_summary,
-          cache=cache is not None)
+          cache=cache is not None,
+          hard_attention_k=hard_attention_k)
     elif attention_type == "dot_product_relative_memory":
       x = dot_product_attention_relative_memory(
           q,
@@ -4177,7 +4206,8 @@ def multihead_attention(query_antecedent,
           image_shapes,
           save_weights_to=save_weights_to,
           make_image_summary=make_image_summary,
-          cache=cache is not None)
+          cache=cache is not None,
+          hard_attention_k=hard_attention_k)
     elif attention_type == "dot_product_unmasked_relative_v2":
       x = dot_product_unmasked_self_attention_relative_v2(
           q,
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 08cd6573d..174dceb20 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -44,6 +44,23 @@ def testAddPositionalEmbedding(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 3, 12))
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testHardenAttentionWeights(self):
+    x = np.random.rand(5, 3, 12)
+    y = common_attention.harden_attention_weights(
+        tf.nn.softmax(tf.constant(x, dtype=tf.float32)), 3)
+    res = self.evaluate(y)
+    self.assertEqual(res.shape, (5, 3, 12))
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testHardenAttentionAllZeros(self):
+    """Check if the hardening code does not divide by zero for all zeros."""
+    x = np.zeros((5, 3, 12), dtype=np.float32)
+    y = common_attention.harden_attention_weights(
+        tf.constant(x, dtype=tf.float32), 3)
+    res = self.evaluate(y)
+    self.assertAllClose(res, x)
+
   @parameterized.parameters(
       {"input_shape": (5, 3, 12)},
       {"input_shape": (5, 5, 5, 12)},
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 97a1ab3a1..cce124b3c 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -95,6 +95,11 @@ def is_xla_compiled():
   return control_flow_util.GetContainingXLAContext(ctxt) is not None
 
 
+def to_float(x):
+  """Cast x to float; created because tf.to_float is deprecated."""
+  return tf.cast(x, tf.float32)
+
+
 def dropout_with_broadcast_dims(x, keep_prob, broadcast_dims=None, **kwargs):
   """Like tf.nn.dropout but takes broadcast_dims instead of noise_shape.
 
@@ -154,7 +159,7 @@ def inverse_exp_decay(max_step, min_value=0.01, step=None):
     step = tf.train.get_global_step()
   if step is None:
     return 1.0
-  step = tf.to_float(step)
+  step = to_float(step)
   return inv_base**tf.maximum(float(max_step) - step, 0.0)
 
 
@@ -164,7 +169,7 @@ def inverse_lin_decay(max_step, min_value=0.01, step=None):
     step = tf.train.get_global_step()
   if step is None:
     return 1.0
-  step = tf.to_float(step)
+  step = to_float(step)
   progress = tf.minimum(step / float(max_step), 1.0)
   return progress * (1.0 - min_value) + min_value
 
@@ -237,7 +242,7 @@ def shakeshake(xs, equal_grad=False):
 def convert_rgb_to_real(x):
   """Conversion of pixel values to real numbers."""
   with tf.name_scope("rgb_to_real", values=[x]):
-    x = tf.to_float(x)
+    x = to_float(x)
     x /= 255.0
     return x
 
@@ -245,7 +250,7 @@ def convert_rgb_to_real(x):
 def convert_rgb_to_symmetric_real(x):
   """Conversion of pixel values to real numbers."""
   with tf.name_scope("rgb_to_real", values=[x]):
-    x = tf.to_float(x)
+    x = to_float(x)
     # Convert each pixel intensity in [0, 1, 2, ..., 255] into a real number in
     # the range [-1, 1].
     x = (x / 127.5) - 1
@@ -274,11 +279,11 @@ def standardize_images(x):
   """Image standardization on batches and videos."""
   with tf.name_scope("standardize_images", values=[x]):
     x_shape = shape_list(x)
-    x = tf.to_float(tf.reshape(x, [-1] + x_shape[-3:]))
+    x = to_float(tf.reshape(x, [-1] + x_shape[-3:]))
     x_mean = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
     x_variance = tf.reduce_mean(
         tf.squared_difference(x, x_mean), axis=[1, 2], keepdims=True)
-    num_pixels = tf.to_float(x_shape[-2] * x_shape[-3])
+    num_pixels = to_float(x_shape[-2] * x_shape[-3])
     x = (x - x_mean) / tf.maximum(tf.sqrt(x_variance), tf.rsqrt(num_pixels))
     return tf.reshape(x, x_shape)
 
@@ -708,7 +713,7 @@ def noam_norm(x, epsilon=1.0, name=None):
     shape = x.get_shape()
     ndims = len(shape)
     return (tf.nn.l2_normalize(x, ndims - 1, epsilon=epsilon) * tf.sqrt(
-        tf.to_float(shape[-1])))
+        to_float(shape[-1])))
 
 
 def l2_norm(x, filters=None, epsilon=1e-6, name=None, reuse=None):
@@ -1135,11 +1140,11 @@ def get_timing_signal(length,
   Returns:
     Tensor of shape (length, 2*num_timescales)
   """
-  positions = tf.to_float(tf.range(length))
+  positions = to_float(tf.range(length))
   log_timescale_increment = (
       math.log(max_timescale / min_timescale) / (num_timescales - 1))
   inv_timescales = min_timescale * tf.exp(
-      tf.to_float(tf.range(num_timescales)) * -log_timescale_increment)
+      to_float(tf.range(num_timescales)) * -log_timescale_increment)
   scaled_time = tf.expand_dims(positions, 1) * tf.expand_dims(inv_timescales, 0)
   return tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
 
@@ -1232,7 +1237,7 @@ def relu_density_logit(x, reduce_dims):
   Returns:
     a Tensor
   """
-  frac = tf.reduce_mean(tf.to_float(x > 0.0), reduce_dims)
+  frac = tf.reduce_mean(to_float(x > 0.0), reduce_dims)
   scaled = tf.log(frac + math.exp(-10)) - tf.log((1.0 - frac) + math.exp(-10))
   return scaled
 
@@ -1614,7 +1619,7 @@ def pad_with_zeros(logits, labels):
 
 def weights_nonzero(labels):
   """Assign weight 1.0 to all labels except for padding (id=0)."""
-  return tf.to_float(tf.not_equal(labels, 0))
+  return to_float(tf.not_equal(labels, 0))
 
 
 def weights_prepend_inputs_to_targets(labels):
@@ -1629,9 +1634,9 @@ def weights_prepend_inputs_to_targets(labels):
   Returns:
     A Tensor of floats.
   """
-  past_first_zero = tf.cumsum(tf.to_float(tf.equal(labels, 0)), axis=1)
-  nonzero = tf.to_float(labels)
-  return tf.to_float(tf.not_equal(past_first_zero * nonzero, 0))
+  past_first_zero = tf.cumsum(to_float(tf.equal(labels, 0)), axis=1)
+  nonzero = to_float(labels)
+  return to_float(tf.not_equal(past_first_zero * nonzero, 0))
 
 
 def check_nonnegative(value):
@@ -1660,24 +1665,24 @@ def weights_multi_problem(labels, taskid=-1):
     ValueError: The Task ID must be valid.
   """
   taskid = check_nonnegative(taskid)
-  past_taskid = tf.cumsum(tf.to_float(tf.equal(labels, taskid)), axis=1)
+  past_taskid = tf.cumsum(to_float(tf.equal(labels, taskid)), axis=1)
   # Additionally zero out the task id location
-  past_taskid *= tf.to_float(tf.not_equal(labels, taskid))
-  non_taskid = tf.to_float(labels)
-  return tf.to_float(tf.not_equal(past_taskid * non_taskid, 0))
+  past_taskid *= to_float(tf.not_equal(labels, taskid))
+  non_taskid = to_float(labels)
+  return to_float(tf.not_equal(past_taskid * non_taskid, 0))
 
 
 def weights_multi_problem_all(labels, taskid=-1):
   """Assign weight 1.0 to only examples from the given task."""
   taskid = check_nonnegative(taskid)
-  weights = tf.to_float(tf.not_equal(labels, 0))
-  past_taskid = tf.cumsum(tf.to_float(tf.equal(labels, taskid)), axis=1)
+  weights = to_float(tf.not_equal(labels, 0))
+  past_taskid = tf.cumsum(to_float(tf.equal(labels, taskid)), axis=1)
   # Additionally zero out the task id location
-  past_taskid *= tf.to_float(tf.not_equal(labels, taskid))
-  non_taskid = tf.to_float(labels)
-  example_mask = tf.to_float(tf.not_equal(past_taskid * non_taskid, 0))
+  past_taskid *= to_float(tf.not_equal(labels, taskid))
+  non_taskid = to_float(labels)
+  example_mask = to_float(tf.not_equal(past_taskid * non_taskid, 0))
   example_mask = tf.reduce_sum(example_mask, axis=1)
-  example_mask = tf.to_float(
+  example_mask = to_float(
       tf.greater(example_mask, tf.zeros_like(example_mask)))
 
   return weights * tf.expand_dims(example_mask, axis=-1)
@@ -1721,7 +1726,7 @@ def weights_concatenated(labels):
   shifted = tf.pad(sentence_num_plus_one,
                    [[0, 0], [2, 0], [0, 0], [0, 0]])[:, :-2, :, :]
   nonboilerplate = tf.equal(sentence_num_plus_one, shifted)
-  ret = tf.to_float(tf.logical_and(nonboilerplate, in_target))
+  ret = to_float(tf.logical_and(nonboilerplate, in_target))
   return ret
 
 
@@ -2106,11 +2111,11 @@ def smoothing_cross_entropy(logits,
   """
   with tf.name_scope("smoothing_cross_entropy", values=[logits, labels]):
     # Low confidence is given to all non-true labels, uniformly.
-    low_confidence = (1.0 - confidence) / tf.to_float(vocab_size - 1)
+    low_confidence = (1.0 - confidence) / to_float(vocab_size - 1)
     # Normalizing constant is the best cross-entropy value with soft targets.
     # We subtract it just for readability, makes no difference on learning.
     normalizing = -(
-        confidence * tf.log(confidence) + tf.to_float(vocab_size - 1) *
+        confidence * tf.log(confidence) + to_float(vocab_size - 1) *
         low_confidence * tf.log(low_confidence + 1e-20))
 
     if gaussian and confidence > 0.0:
@@ -3198,7 +3203,7 @@ def mix(x1,
       if broadcast_last:
         alpha_shape = alpha_shape[:-1] + [1]
       alpha = tf.random_uniform(alpha_shape)
-      alpha = tf.to_float(tf.less(alpha, max_prob))
+      alpha = to_float(tf.less(alpha, max_prob))
       return alpha * x1 + (1.0 - alpha) * x2
 
     def get_res():
@@ -3220,7 +3225,7 @@ def get_res():
       if broadcast_last:
         alpha_shape = alpha_shape[:-1] + [1]
       alpha = tf.random_uniform(alpha_shape)
-      alpha = tf.to_float(tf.less(alpha, alpha_p))
+      alpha = to_float(tf.less(alpha, alpha_p))
       return alpha * x1 + (1.0 - alpha) * x2
 
     if max_prob < 1.0:
@@ -3327,6 +3332,39 @@ def log_prob_from_logits(logits, reduce_axis=-1):
   return logits - tf.reduce_logsumexp(logits, axis=reduce_axis, keepdims=True)
 
 
+def top_kth_iterative(x, k):
+  """Compute the k-th top element of x on the last axis iteratively.
+
+  This assumes values in x are non-negative, rescale if needed.
+  It is often faster than tf.nn.top_k for small k, especially if k < 30.
+  Note: this does not support back-propagation, it stops gradients!
+
+  Args:
+    x: a Tensor of non-negative numbers of type float.
+    k: a python integer.
+
+  Returns:
+    a float tensor of the same shape as x but with 1 on the last axis
+    that contains the k-th largest number in x.
+  """
+  # The iterative computation is as follows:
+  #
+  # cur_x = x
+  # for _ in range(k):
+  #   top_x = maximum of elements of cur_x on the last axis
+  #   cur_x = cur_x where cur_x < top_x and 0 everywhere else (top elements)
+  #
+  # We encode this computation in a TF graph using tf.foldl, so the inner
+  # part of the above loop is called "next_x" and tf.foldl does the loop.
+  def next_x(cur_x, _):
+    top_x = tf.reduce_max(cur_x, axis=-1, keep_dims=True)
+    return cur_x * to_float(cur_x < top_x)
+  # We only do k-1 steps of the loop and compute the final max separately.
+  fin_x = tf.foldl(next_x, tf.range(k - 1), initializer=tf.stop_gradient(x),
+                   parallel_iterations=2, back_prop=False)
+  return tf.stop_gradient(tf.reduce_max(fin_x, axis=-1, keep_dims=True))
+
+
 def top_1_tpu(inputs):
   """find max and argmax over the last dimension.
 
@@ -3723,7 +3761,7 @@ def tpu_safe_image_summary(image):
   if is_xla_compiled():
     # We only support float32 images at the moment due to casting complications.
     if image.dtype != tf.float32:
-      image = tf.to_float(image)
+      image = to_float(image)
   else:
     image = tf.cast(image, tf.uint8)
   return image
@@ -3802,7 +3840,7 @@ def weight_targeting(w, k):
 
   transpose_w = tf.transpose(w)
   thres = tf.contrib.framework.sort(tf.abs(transpose_w), axis=1)[:, k]
-  mask = tf.to_float(thres[None, :] >= tf.abs(w))
+  mask = to_float(thres[None, :] >= tf.abs(w))
 
   return tf.reshape(mask, w_shape)
 
@@ -3816,7 +3854,7 @@ def unit_targeting(w, k):
 
   norm = tf.norm(w, axis=0)
   thres = tf.contrib.framework.sort(norm, axis=0)[k]
-  mask = tf.to_float(thres >= norm)[None, :]
+  mask = to_float(thres >= norm)[None, :]
   mask = tf.tile(mask, [size, 1])
 
   return tf.reshape(mask, w_shape)
@@ -3921,7 +3959,7 @@ def targeted_dropout(inputs,
     Tensor, same shape and dtype as `inputs`.
   """
   if not is_training and do_prune:
-    k = tf.round(tf.to_float(k) * tf.to_float(1. - keep_prob))
+    k = tf.round(to_float(k) * to_float(1. - keep_prob))
 
   mask = targeting_fn(inputs, k)
   mask = tf.cast(mask, inputs.dtype)
@@ -3953,7 +3991,7 @@ def kl_divergence(mu, log_var, mu_p=0.0, log_var_p=0.0):
       mu, tf.exp(tf.multiply(0.5, log_var)))
   kld = tfp.distributions.kl_divergence(posterior_distribution,
                                         prior_distribution)
-  return tf.reduce_sum(kld) / tf.to_float(batch_size)
+  return tf.reduce_sum(kld) / to_float(batch_size)
 
 
 def sparse_equals_constant(constant, tensor):
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index a11aff86e..cc3b8bb73 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -389,6 +389,20 @@ def testRavanbakhshSetLayer(self):
     actual = self.evaluate(layer)
     self.assertEqual(actual.shape, (5, 4, 32))
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testTopKthIterativeShape(self):
+    x = np.random.rand(5, 2, 1, 12)
+    y = common_layers.top_kth_iterative(tf.constant(x, dtype=tf.float32), 3)
+    actual = self.evaluate(y)
+    self.assertEqual(actual.shape, (5, 2, 1, 1))
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testTopKthIterativeValue(self):
+    x = [1.0, 2.0, 3.0, 4.0]
+    y = common_layers.top_kth_iterative(tf.constant(x, dtype=tf.float32), 3)
+    actual = self.evaluate(y)
+    self.assertEqual(int(actual[0]), 2.0)
+
   @test_utils.run_in_graph_and_eager_modes()
   def testBReLU(self):
     x = np.random.rand(5, 2, 1, 12)
diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 2202fdefb..3a62021f6 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -203,7 +203,8 @@ def transformer_encoder(encoder_input,
               max_length=hparams.get("max_length"),
               vars_3d=hparams.get("attention_variables_3d"),
               activation_dtype=hparams.get("activation_dtype", "float32"),
-              weight_dtype=hparams.get("weight_dtype", "float32"))
+              weight_dtype=hparams.get("weight_dtype", "float32"),
+              hard_attention_k=hparams.get("hard_attention_k", 0))
           x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 0cba09701..252be5385 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1467,6 +1467,7 @@ def transformer_decoder(decoder_input,
               layer_collection=layer_collection,
               recurrent_memory=recurrent_memory,
               chunk_number=chunk_number,
+              hard_attention_k=hparams.get("hard_attention_k", 0)
               )
           x = common_layers.layer_postprocess(x, y, hparams)
         if encoder_output is not None:
@@ -1493,7 +1494,8 @@ def transformer_decoder(decoder_input,
                 vars_3d=hparams.get("attention_variables_3d"),
                 activation_dtype=hparams.get("activation_dtype", "float32"),
                 weight_dtype=hparams.get("weight_dtype", "float32"),
-                layer_collection=layer_collection)
+                layer_collection=layer_collection,
+                hard_attention_k=hparams.get("hard_attention_k", 0))
             x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(
@@ -1614,6 +1616,8 @@ def transformer_base_v1():
   # For making a transformer encoder unidirectional by using masked
   # attention.
   hparams.add_hparam("unidirectional_encoder", False)
+  # For hard attention.
+  hparams.add_hparam("hard_attention_k", 0)
   return hparams
 
 

From dd58574f6218333dd62ff1a74dad9970a34cc7c7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 28 Mar 2019 14:57:15 -0700
Subject: [PATCH 1844/2720] Fix recurrent memory batch size inference

PiperOrigin-RevId: 240859652
---
 tensor2tensor/layers/transformer_memory.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index 2e815ddec..67b4b5078 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -78,7 +78,8 @@ def __init__(self, name, hparams):
     # TODO(kitaev): The implementation of the chunking code makes it somewhat
     # convoluted to figure out how many actual sequences we can have per batch.
     # The data pipeline should be revisited at some point.
-    if hasattr(hparams, "recurrent_memory_batch_size"):
+    if (hasattr(hparams, "recurrent_memory_batch_size")
+        and hparams.recurrent_memory_batch_size > 0):
       batch_size_in_sequences = hparams.recurrent_memory_batch_size
     else:
       batch_size_in_sequences = hparams.batch_size / hparams.max_length

From 75611fcd5066c277da15ddace1e3a1d0692618d9 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 28 Mar 2019 20:48:12 -0700
Subject: [PATCH 1845/2720] Do not save recurrent memory state in checkpoints

This ensures that checkpoint files are compatible across different chunk and memory sizes.

PiperOrigin-RevId: 240912219
---
 tensor2tensor/layers/transformer_memory.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index 67b4b5078..59132d0e2 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -91,16 +91,19 @@ def __init__(self, name, hparams):
       self.previous_segment = tf.get_variable(
           "memsegment", (),
           dtype=tf.int32, trainable=False,
+          collections=[tf.GraphKeys.LOCAL_VARIABLES],
           initializer=tf.constant_initializer(0))
 
       self.previous_vals = tf.get_variable(
           "memvals", memory_shape,
           dtype=tf.float32, trainable=False,
+          collections=[tf.GraphKeys.LOCAL_VARIABLES],
           initializer=tf.constant_initializer(.0))
 
       self.previous_bias = tf.get_variable(
           "membias", bias_shape,
           dtype=tf.float32, trainable=False,
+          collections=[tf.GraphKeys.LOCAL_VARIABLES],
           initializer=tf.constant_initializer(.0))
 
   def pre_attention(self, segment, query_antecedent, memory_antecedent, bias):

From de821fcce3f3965b80843b8779dd80cbd0c2fd18 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 29 Mar 2019 13:43:41 -0700
Subject: [PATCH 1846/2720] Correct random seed handling in TRAX, make input
 pipeline more aligned with T2T and set defaults better for colab ease of use.

PiperOrigin-RevId: 241042919
---
 tensor2tensor/trax/inputs.py  | 19 ++++++++++++++++---
 tensor2tensor/trax/trainer.py |  4 ----
 tensor2tensor/trax/trax.py    | 28 ++++++++++++++++++++++------
 3 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 7a7a3c910..cf7e42e9a 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -21,6 +21,7 @@
 
 import collections
 import os
+import random
 
 import gin
 
@@ -33,9 +34,14 @@
 Inputs = collections.namedtuple(
     "_Inputs", ["train_stream", "eval_stream", "input_shape"])
 
+# How many examples from the stream to skip at random during training.
+# For now, we skip at most 1M examples.
+# TODO(lukaszkaiser): does it matter for efficiency, should that be changed?
+_MAX_SKIP_EXAMPLES = 1e6
+
 
 @gin.configurable()
-def inputs(dataset_name, data_dir):
+def inputs(dataset_name, data_dir=None):
   """Make Inputs for built-in datasets.
 
   Args:
@@ -175,7 +181,8 @@ def target_right_length(_, target):
 @gin.configurable(blacklist=["dataset", "training", "shapes", "target_names"])
 def batch_fun(dataset, training, shapes, target_names,
               batch_size=32, eval_batch_size=32,
-              bucket_length=32, buckets=None):
+              bucket_length=32, buckets=None,
+              batch_shuffle_size=512):
   """Batching function."""
   del target_names
   # If bucketing is not specified, check if target shapes are variable.
@@ -209,6 +216,8 @@ def example_length(_, target):
         pad_to_bucket_boundary=training))
   else:
     dataset = dataset.padded_batch(cur_batch_size, shapes)
+  if training:
+    return dataset.shuffle(batch_shuffle_size)
   return dataset
 
 
@@ -225,12 +234,16 @@ def append_targets(example):
   dataset = dataset.map(append_targets)
   if training:
     dataset = dataset.repeat()
+    # Skip a random fraction at the beginning of the stream.  The skip is
+    # essential for synchronous highly-parallel training to avoid multiple
+    # replicas reading the same data in lock-step.
+    dataset = dataset.skip(random.randint(0, _MAX_SKIP_EXAMPLES))
   shapes = {k: features_info[k].shape for k in features_info}
   shapes = (shapes, shapes[target_names[0]])
   dataset = dataset.shuffle(1024)
   dataset = preprocess_fun(dataset, training)
   dataset = batch_fun(dataset, training, shapes, target_names)
-  return dataset.prefetch(32)
+  return dataset.prefetch(2)
 
 
 @gin.configurable(whitelist=["input_name"])
diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index 9f4cfb6db..221d01d8f 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -61,7 +61,6 @@ def _setup_gin():
   """Setup gin configuration."""
   # Imports for configurables
   # pylint: disable=g-import-not-at-top,unused-import,g-bad-import-order,reimported,unused-variable
-  from tensor2tensor.trax import inputs as _trax_inputs
   from tensor2tensor.trax import models as _trax_models
   from tensor2tensor.trax import optimizers as _trax_opt
   # pylint: disable=g-import-not-at-top,unused-import,g-bad-import-order,reimported,unused-variable
@@ -72,9 +71,6 @@ def _setup_gin():
     configs.append("inputs.dataset_name='%s'" % FLAGS.dataset)
     if FLAGS.data_dir:
       configs.append("inputs.data_dir='%s'" % FLAGS.data_dir)
-    else:
-      configs.append("inputs.data_dir=None")
-    configs.append("train.inputs=@trax.inputs.inputs")
   if FLAGS.model:
     configs.append("train.model=@trax.models.%s" % FLAGS.model)
   gin.parse_config_files_and_bindings(FLAGS.config_file, configs)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 8ed32e64c..4b829eced 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -24,6 +24,7 @@
 import itertools
 import os
 import pickle
+import random
 import time
 
 from absl import logging
@@ -33,16 +34,17 @@
 import jax
 from jax.experimental import optimizers as jax_opt
 import jax.numpy as np
-import jax.random as random
-
+import numpy
 import six
 
 from tensor2tensor.trax import history as trax_history
+from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import jaxboard
 from tensor2tensor.trax import learning_rate as lr
-from tensor2tensor.trax import optimizers as trax_opt
+from tensor2tensor.trax import optimizers as trax_optimizers
 import tensor2tensor.trax.stax as stax
 
+import tensorflow as tf
 from tensorflow.io import gfile
 
 
@@ -187,6 +189,18 @@ def log_metrics(metrics, summ_writer, log_prefix, step, history=None):
       summ_writer.scalar(full_name, value, step)
 
 
+def get_random_number_generator_and_set_seed(seed=None):
+  """Get a JAX random number generator and set random seed everywhere."""
+  random.seed(seed)
+  # While python random accepts None as seed and uses time/os seed then,
+  # some other functions expect integers so we create one here.
+  if seed is None:
+    seed = random.randint(0, 2**31 - 1)
+  tf.set_random_seed(seed)
+  numpy.random.seed(seed)
+  return jax.random.PRNGKey(seed)
+
+
 # TODO(trax):
 # * Make configurable:
 #   * loss
@@ -259,13 +273,14 @@ def reshape_by_device(train_data, num_devices):
 @gin.configurable(blacklist=["output_dir"])
 def train(output_dir,
           model=gin.REQUIRED,
-          inputs=gin.REQUIRED,
-          optimizer=trax_opt.adam,
+          inputs=trax_inputs.inputs,
+          optimizer=trax_optimizers.adam,
           lr_schedule=lr.MultifactorSchedule,
           train_steps=1000,
           eval_steps=10,
           eval_frequency=100,
           num_devices=None,
+          random_seed=None,
           run_debug_step=False):
   """Train the model on the inputs.
 
@@ -283,6 +298,7 @@ def train(output_dir,
     eval_frequency: int, how often to run evaluation (every eval_frequency
       steps). If None or 0, eval disabled.
     num_devices: how many devices to use (if None, default, use all available)
+    random_seed: the random seed to use; time/os dependent if None (default).
     run_debug_step: bool, if True, will run the model and loss without @jit for
       one step.
 
@@ -290,7 +306,7 @@ def train(output_dir,
     trax.State
   """
   num_devices = num_devices or jax.lib.xla_bridge.device_count()
-  rng = random.PRNGKey(0)
+  rng = get_random_number_generator_and_set_seed(random_seed)
   gfile.makedirs(output_dir)
   # Create summary writers and history.
   train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))

From d69128328ab9942cd4009c1ff1fe2492743afa23 Mon Sep 17 00:00:00 2001
From: Zi Yang <ziy@google.com>
Date: Fri, 29 Mar 2019 14:44:23 -0700
Subject: [PATCH 1847/2720] Added key checks for "inputs" and "targets".

PiperOrigin-RevId: 241054946
---
 tensor2tensor/data_generators/problem.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index d83ef5bbe..f9998c850 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -141,7 +141,7 @@ def default_model_hparams():
 
 def preprocess_example_common(example, mode, hparams):
   """Preprocessing steps common to all models."""
-  if hparams.max_input_seq_length > 0:
+  if "inputs" in example and hparams.max_input_seq_length > 0:
     example["inputs"] = example["inputs"][:hparams.max_input_seq_length]
   if hparams.prepend_mode != "none":
     if mode == tf.estimator.ModeKeys.PREDICT:
@@ -149,7 +149,7 @@ def preprocess_example_common(example, mode, hparams):
     else:
       example["targets"] = tf.concat(
           [example["inputs"], [0], example["targets"]], 0)
-  if hparams.max_target_seq_length > 0:
+  if "targets" in example and hparams.max_target_seq_length > 0:
     example["targets"] = example["targets"][:hparams.max_target_seq_length]
   if hparams.split_to_length:
     new_example = {}

From 164b3422c929f3598f200a7022f68521feb5a0e1 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 29 Mar 2019 18:31:49 -0700
Subject: [PATCH 1848/2720] internal

PiperOrigin-RevId: 241086998
---
 tensor2tensor/trax/trainer.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index 221d01d8f..f02f275fc 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -27,7 +27,7 @@
 from absl import logging
 
 import gin
-
+import jax
 from tensor2tensor.trax import trax
 
 FLAGS = flags.FLAGS
@@ -42,6 +42,8 @@
 flags.DEFINE_multi_string("config", None,
                           "Configuration parameters (gin string).")
 flags.DEFINE_integer("log_level", logging.INFO, "Log level.")
+flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
+
 
 
 def _default_output_dir():
@@ -86,6 +88,10 @@ def main(_):
   trax.log("Using --output_dir %s" % output_dir)
   output_dir = os.path.expanduser(output_dir)
 
+  # If on TPU, let JAX know.
+  if FLAGS.use_tpu:
+    jax.config.update("jax_platform_name", "tpu")
+
   trax.train(output_dir=output_dir)
 
 

From bb6440d8dcb394153c45ab44886bf71b3d6f5799 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sun, 31 Mar 2019 08:31:14 -0700
Subject: [PATCH 1849/2720] Integrated the neural memory with transformer.

PiperOrigin-RevId: 241213370
---
 tensor2tensor/layers/transformer_memory.py    | 117 ++++++++++++++----
 .../layers/transformer_memory_test.py         |  23 ++++
 tensor2tensor/models/transformer.py           |  19 ++-
 3 files changed, 135 insertions(+), 24 deletions(-)

diff --git a/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index 59132d0e2..92eaeadd5 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.layers import common_layers
 import tensorflow as tf
 
 
@@ -186,7 +187,7 @@ class TransformerMemory(object):
   """
 
   def __init__(self, batch_size, key_depth, val_depth, memory_size,
-               sharpen_factor=1.):
+               sharpen_factor=1., name="neural_memory"):
     """Initialize the memory object.
 
     Args:
@@ -195,20 +196,31 @@ def __init__(self, batch_size, key_depth, val_depth, memory_size,
       val_depth: the depth of the memory values.
       memory_size: the number of items in the memory.
       sharpen_factor: the sharpen_factor for addressing the memory.
+      name: the optional variable scope.
     """
+    self.name = name
     self.batch_size = batch_size
     self.key_depth = key_depth
     self.val_depth = val_depth
     self.memory_size = memory_size
     self.sharpen_factor = sharpen_factor
-    self.mem_vals = tf.get_variable(
-        "memvals", [self.batch_size, self.memory_size, self.val_depth],
-        dtype=tf.float32, trainable=False,
-        initializer=tf.constant_initializer(.0))
-    self.mean_logits = tf.get_variable(
-        "meanlogits", [self.batch_size, self.memory_size],
-        dtype=tf.float32, trainable=False,
-        initializer=tf.constant_initializer(.0))
+    with tf.variable_scope(name):
+      self.segment_number = tf.get_variable(
+          "segment_number", [self.batch_size],
+          dtype=tf.int32, trainable=False,
+          initializer=tf.constant_initializer(100000))
+      self.mem_vals = tf.get_variable(
+          "memvals", [self.batch_size, self.memory_size, self.val_depth],
+          dtype=tf.float32, trainable=False,
+          initializer=tf.constant_initializer(.0))
+      self.mean_logits = tf.get_variable(
+          "meanlogits", [self.batch_size, self.memory_size],
+          dtype=tf.float32, trainable=False,
+          initializer=tf.constant_initializer(.0))
+
+  def _norm(self, x):
+    """Compute the safe norm."""
+    return tf.sqrt(tf.reduce_sum(tf.square(x), keepdims=True, axis=-1) + 1e-7)
 
   def _address_content(self, x):
     """Address the memory based on content similarity.
@@ -218,14 +230,16 @@ def _address_content(self, x):
     Returns:
       the logits for each memory entry [batch_size, length, memory_size].
     """
-    mem_keys = tf.layers.dense(self.mem_vals, self.key_depth, name="mem_key")
-    mem_query = tf.layers.dense(x, self.key_depth, name="mem_query")
-    norm = tf.matmul(
-        tf.norm(mem_query, axis=-1, keepdims=True),
-        tf.norm(mem_keys, axis=-1, keepdims=True), transpose_b=True)
-    cos_dist = tf.div(
-        tf.matmul(mem_query, mem_keys, transpose_b=True), norm,
-        name="cos_dist")
+    mem_keys = tf.layers.dense(self.mem_vals, self.key_depth,
+                               bias_initializer=tf.constant_initializer(1.0),
+                               name="mem_key")
+    mem_query = tf.layers.dense(x, self.key_depth,
+                                bias_initializer=tf.constant_initializer(1.0),
+                                name="mem_query")
+    norm = tf.matmul(self._norm(mem_query), self._norm(mem_keys),
+                     transpose_b=True)
+    dot_product = tf.matmul(mem_query, mem_keys, transpose_b=True)
+    cos_dist = tf.div(dot_product, norm + 1e-7, name="cos_dist")
     access_logits = self.sharpen_factor * cos_dist
     return access_logits
 
@@ -270,15 +284,17 @@ def write(self, x, access_logits):
                                   activation=tf.nn.sigmoid,
                                   name="erase")
     write_weights = tf.nn.softmax(write_logits)
-    erase = tf.multiply(tf.expand_dims(1 - erase_gates * write_weights, 3),
+    erase_weights = tf.expand_dims(1 - erase_gates * write_weights, 3)
+    erase = tf.multiply(erase_weights,
                         tf.expand_dims(self.mem_vals, 1))
     addition = tf.multiply(
-        tf.expand_dims(write_weights, 3), tf.expand_dims(candidate_value, 2))
+        tf.expand_dims(write_weights, 3),
+        tf.expand_dims(candidate_value, 2))
     update_value_op = self.mem_vals.assign(
-        tf.reduce_sum(erase + addition, axis=1))
+        tf.reduce_mean(erase + addition, axis=1))
     with tf.control_dependencies([update_value_op]):
       write_op = self.mean_logits.assign(
-          self.mean_logits * 0.1 + tf.reduce_sum(write_logits * 0.9, axis=1))
+          self.mean_logits * 0.1 + tf.reduce_mean(write_logits * 0.9, axis=1))
       return write_op
 
   def set(self, mem_vals, mean_logits):
@@ -290,6 +306,9 @@ def set(self, mem_vals, mean_logits):
   def get(self):
     return self.mem_vals, self.mean_logits
 
+  def update_segment_number(self, segment_number):
+    return self.segment_number.assign(segment_number)
+
   def reset(self, entries_to_reset):
     """Reset the entries in the memory.
 
@@ -311,3 +330,59 @@ def reset(self, entries_to_reset):
                 [num_updates, 1]))
     reset_op = tf.group([update_vals, update_logits])
     return reset_op
+
+  def pre_attention(self, segment_number, query_antecedent,
+                    memory_antecedent, bias):
+    """Called prior to self-attention, to incorporate memory items.
+
+    Args:
+      segment_number: an integer Tensor with shape [batch]
+      query_antecedent: a Tensor with shape [batch, length_q, channels]
+      memory_antecedent: must be None. Attention normally allows this to be a
+        Tensor with shape [batch, length_m, channels], but we currently only
+        support memory for decoder-side self-attention.
+      bias: bias Tensor (see attention_bias())
+    Returns:
+      (data, new_query_antecedent, new_memory_antecedent, new_bias)
+    """
+    with tf.variable_scope(self.name + "/pre_attention", reuse=tf.AUTO_REUSE):
+      assert memory_antecedent is None, "We only support language modeling"
+      with tf.control_dependencies([
+          tf.assert_greater_equal(self.batch_size, tf.size(segment_number))]):
+        difference = self.batch_size - tf.size(segment_number)
+        segment_number = tf.pad(segment_number, [[0, difference]])
+        reset_op = self.reset(tf.reshape(tf.where(
+            tf.less(segment_number, self.segment_number)), [-1]))
+      memory_results = {}
+      with tf.control_dependencies([reset_op]):
+        with tf.control_dependencies([
+            self.update_segment_number(segment_number)]):
+          x = tf.pad(query_antecedent, [
+              [0, difference], [0, 0], [0, 0]])
+          access_logits, retrieved_mem = self.read(x)
+      memory_results["x"] = x
+      memory_results["access_logits"] = access_logits
+      memory_results["retrieved_mem"] = retrieved_mem
+      return memory_results, query_antecedent, memory_antecedent, bias
+
+  def post_attention(self, token, x):
+    """Called after self-attention. The memory can be updated here.
+
+    Args:
+      token: Data returned by pre_attention, which can be used to carry over
+        state related to the current memory operation.
+      x: a Tensor of data after self-attention and feed-forward
+    Returns:
+      a (possibly modified) version of the input x
+    """
+    with tf.variable_scope(self.name + "/post_attention", reuse=tf.AUTO_REUSE):
+      depth = common_layers.shape_list(x)[-1]
+      actual_batch_size = common_layers.shape_list(x)[0]
+      memory_output = tf.gather(token["retrieved_mem"],
+                                tf.range(actual_batch_size))
+      output = tf.add(tf.layers.dense(x, depth, use_bias=False),
+                      tf.layers.dense(memory_output, depth))
+      with tf.control_dependencies([output]):
+        with tf.control_dependencies([
+            self.write(token["x"], token["access_logits"])]):
+          return tf.identity(output)
diff --git a/tensor2tensor/layers/transformer_memory_test.py b/tensor2tensor/layers/transformer_memory_test.py
index b5c365c41..cd86dae95 100644
--- a/tensor2tensor/layers/transformer_memory_test.py
+++ b/tensor2tensor/layers/transformer_memory_test.py
@@ -98,5 +98,28 @@ def testReset(self):
     self.assertAllEqual(0, zero1)
     self.assertAllEqual(0, zero2)
 
+  def testLoss(self):
+    batch_size = 2
+    key_depth = 5
+    val_depth = 5
+    memory_size = 4
+    window_size = 3
+    x_depth = 5
+    memory = transformer_memory.TransformerMemory(
+        batch_size, key_depth, val_depth, memory_size)
+    x = tf.random_uniform([batch_size, window_size, x_depth], minval=.0)
+    memory_results, _, _, _ = (
+        memory.pre_attention(
+            tf.random_uniform([batch_size], minval=0, maxval=1, dtype=tf.int32),
+            x, None, None))
+    x = memory.post_attention(memory_results, x)
+    with tf.control_dependencies([tf.print("x", x)]):
+      is_nan = tf.reduce_any(tf.math.is_nan(x))
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      for _ in range(100):
+        is_nan_value, _ = session.run([is_nan, x])
+    self.assertEqual(is_nan_value, False)
+
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 252be5385..536352ff4 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1533,9 +1533,20 @@ def __init__(self, *args, **kwargs):
     self.recurrent_memory_by_layer = {}
     for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
       layer_name = "layer_%d" % layer
-      self.recurrent_memory_by_layer[
-          layer_name] = transformer_memory.RecentTokensMemory(
-              layer_name + "/recurrent_memory", hparams)
+      if hparams.memory_type == "neural_memory":
+        memory = transformer_memory.TransformerMemory(
+            batch_size=int(hparams.batch_size / hparams.max_length),
+            key_depth=hparams.hidden_size,
+            val_depth=hparams.hidden_size,
+            memory_size=hparams.split_targets_chunk_length,
+            sharpen_factor=1.,
+            name=layer_name + "/recurrent_memory")
+      elif hparams.memory_type == "transformer_xl":
+        memory = transformer_memory.RecentTokensMemory(
+            layer_name + "/recurrent_memory", hparams)
+      else:
+        raise ValueError("Unsupported memory type: %s" % hparams.memory_type)
+      self.recurrent_memory_by_layer[layer_name] = memory
 
   @property
   def has_input(self):
@@ -2640,6 +2651,7 @@ def transformer_wikitext103_l4k_memory_v0():
 
   hparams.split_targets_chunk_length = 64
   hparams.split_targets_max_chunks = 64
+  hparams.add_hparam("memory_type", "transformer_xl")
 
   # The hparams specify batch size *before* chunking, but we want to have a
   # consistent 4K batch size *after* chunking to fully utilize the hardware.
@@ -2705,3 +2717,4 @@ def transformer_cifar10_memory_v0():
       hparams.num_memory_items + hparams.split_targets_chunk_length)
 
   return hparams
+

From 2f1380d40751757e0d933fd6e4734faf525eccf2 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 1 Apr 2019 09:42:27 -0700
Subject: [PATCH 1850/2720] Enable customizing host_call for eval/train modes.

PiperOrigin-RevId: 241341729
---
 tensor2tensor/utils/t2t_model.py | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 8a1a1c7b0..fb336c52a 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1532,9 +1532,12 @@ def estimator_model_fn(cls,
   def initialize_from_ckpt(self, ckpt_dir):
     return initialize_from_ckpt(ckpt_dir=ckpt_dir, hparams=self._hparams)
 
-  def create_host_call(self):
+  def create_train_host_call(self):
     return create_host_call(self.hparams.model_dir)
 
+  def create_eval_host_call(self):
+    return self.create_train_host_call()
+
   def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
     """Constructs `tf.estimator.EstimatorSpec` for TRAIN (training) mode."""
     train_op = self.optimize(loss, num_async_replicas=num_async_replicas,
@@ -1550,7 +1553,7 @@ def scaffold_fn():
 
       # Note: important to call this before remove_summaries()
       if self.hparams.tpu_enable_host_call:
-        host_call = self.create_host_call()
+        host_call = self.create_train_host_call()
       else:
         host_call = None
 
@@ -1587,7 +1590,7 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
     if common_layers.is_xla_compiled():
       # Note: important to call this before remove_summaries()
       if self.hparams.tpu_enable_host_call:
-        host_call = self.create_host_call()
+        host_call = self.create_eval_host_call()
       else:
         host_call = None
 
@@ -1718,16 +1721,23 @@ def estimator_spec_predict(self, features, use_tpu=False):
     if "batch_prediction_key" in predictions:
       export_out["batch_prediction_key"] = predictions["batch_prediction_key"]
 
-    remove_summaries()
-
     export_outputs = {
         tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
             tf.estimator.export.PredictOutput(export_out)
     }
     if use_tpu:
+      # Note: important to call this before remove_summaries()
+      if self.hparams.tpu_enable_host_call:
+        host_call = self.create_eval_host_call()
+      else:
+        host_call = None
+
+      remove_summaries()
+
       return tf.contrib.tpu.TPUEstimatorSpec(
           tf.estimator.ModeKeys.PREDICT,
           predictions=predictions,
+          host_call=host_call,
           export_outputs=export_outputs)
     else:
       return tf.estimator.EstimatorSpec(

From 241a315ec906fa50526e4d3b458576f4c5d608b5 Mon Sep 17 00:00:00 2001
From: James Martens <jamesmartens@google.com>
Date: Mon, 1 Apr 2019 13:01:42 -0700
Subject: [PATCH 1851/2720] Modifying the Layer Norm parameters to use the new
 features of K-FAC

PiperOrigin-RevId: 241382675
---
 tensor2tensor/layers/common_layers.py | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index cce124b3c..68f2b534d 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -657,14 +657,22 @@ def layer_norm_vars(filters):
   return scale, bias
 
 
-def layer_norm_compute(x, epsilon, scale, bias):
+def layer_norm_compute(x, epsilon, scale, bias, layer_collection=None):
   """Layer norm raw computation."""
+
+  # Save these before they get converted to tensors by the casting below
+  params = (scale, bias)
+
   epsilon, scale, bias = [cast_like(t, x) for t in [epsilon, scale, bias]]
   mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
   variance = tf.reduce_mean(
       tf.squared_difference(x, mean), axis=[-1], keepdims=True)
   norm_x = (x - mean) * tf.rsqrt(variance + epsilon)
-  return norm_x * scale + bias
+
+  output = norm_x * scale + bias
+
+
+  return output
 
 
 def layer_norm(x,
@@ -679,11 +687,8 @@ def layer_norm(x,
   with tf.variable_scope(
       name, default_name="layer_norm", values=[x], reuse=reuse):
     scale, bias = layer_norm_vars(filters)
-    if layer_collection:
-      tf.logging.info("Registering layer norm to collection with (scale, bias):"
-                      " ({}, {})".format(scale, bias))
-      layer_collection.register_generic((scale, bias), shape_list(x)[0])
-    return layer_norm_compute(x, epsilon, scale, bias)
+    return layer_norm_compute(x, epsilon, scale, bias,
+                              layer_collection=layer_collection)
 
 
 def group_norm(x, filters=None, num_groups=8, epsilon=1e-5):

From de7f0b34087e9173d52c09b97f18ec21565b89f6 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 1 Apr 2019 13:55:18 -0700
Subject: [PATCH 1852/2720] Demonstrate how to use posterior mean (or any other
 value) on forward pass.

PiperOrigin-RevId: 241392931
---
 tensor2tensor/layers/bayes.py      | 17 +++++++++--------
 tensor2tensor/layers/bayes_test.py | 28 +++++++++++++++++++++++++---
 2 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index d6e2162f1..c1a99ccc9 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -452,16 +452,15 @@ def __init__(self,
         activity_regularizer=get(activity_regularizer),
         **kwargs)
 
-  # TODO(trandustin): This name is not accurate. Rename or move functionality
-  # into random variables to resample/recreate their init ops.
-  def sample_weights(self):
+  def call_weights(self):
+    """Calls any weights if the initializer is itself a layer."""
     if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
       self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
     if isinstance(self.bias_initializer, tf.keras.layers.Layer):
       self.bias = self.bias_initializer(self.bias.shape, self.dtype)
 
   def call(self, *args, **kwargs):
-    self.sample_weights()
+    self.call_weights()
     return super(DenseReparameterization, self).call(*args, **kwargs)
 
 
@@ -520,14 +519,15 @@ def __init__(self,
         bias_constraint=get(bias_constraint),
         **kwargs)
 
-  def sample_weights(self):
+  def call_weights(self):
+    """Calls any weights if the initializer is itself a layer."""
     if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
       self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
     if isinstance(self.bias_initializer, tf.keras.layers.Layer):
       self.bias = self.bias_initializer(self.bias.shape, self.dtype)
 
   def call(self, *args, **kwargs):
-    self.sample_weights()
+    self.call_weights()
     return super(Conv2DReparameterization, self).call(*args, **kwargs)
 
 
@@ -796,7 +796,8 @@ def bias_initializer(_, *args, **kwargs):
       self.bias = None
     self.built = True
 
-  def sample_weights(self):
+  def call_weights(self):
+    """Calls any weights if the initializer is itself a layer."""
     if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
       self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
     if isinstance(self.recurrent_initializer, tf.keras.layers.Layer):
@@ -809,7 +810,7 @@ def sample_weights(self):
   def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
     """Get the initial state and side-effect sampling of stochastic weights."""
     if self.built:
-      self.sample_weights()
+      self.call_weights()
     return super(LSTMCellReparameterization, self).get_initial_state(
         inputs=inputs, batch_size=batch_size, dtype=dtype)
 
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index bcde4ecfa..85a092528 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow as tf
+from tensorflow_probability import edward2 as ed
 tf.compat.v1.enable_eager_execution()
 
 
@@ -75,6 +76,27 @@ def testDenseReparameterizationKernel(
       self.assertNotAllClose(res1, res2)
     layer.get_config()
 
+  @test_utils.run_in_graph_and_eager_modes
+  def testDenseReparameterizationMean(self):
+    """Tests that forward pass can use other values, e.g., posterior mean."""
+    def take_mean(f, *args, **kwargs):
+      """Sets random variable value to its mean."""
+      rv = f(*args, **kwargs)
+      rv._value = rv.distribution.mean()
+      return rv
+    inputs = tf.to_float(np.random.rand(5, 3, 7))
+    layer = bayes.DenseReparameterization(4,
+                                          activation=tf.nn.relu,
+                                          use_bias=False)
+    outputs1 = layer(inputs)
+    with ed.interception(take_mean):
+      outputs2 = layer(inputs)
+    self.evaluate(tf.global_variables_initializer())
+    res1, res2 = self.evaluate([outputs1, outputs2])
+    self.assertEqual(res1.shape, (5, 3, 4))
+    self.assertNotAllClose(res1, res2)
+    self.assertAllClose(res2, np.zeros((5, 3, 4)), atol=1e-4)
+
   @test_utils.run_in_graph_and_eager_modes()
   def testDenseReparameterizationLoss(self):
     features = tf.to_float(np.random.rand(5, 12))
@@ -289,7 +311,7 @@ def testLSTMCellReparameterization(
     state = (h0 + noise, c0)
     outputs1, _ = cell(inputs[:, 0, :], state)
     outputs2, _ = cell(inputs[:, 0, :], state)
-    cell.sample_weights()
+    cell.call_weights()
     outputs3, _ = cell(inputs[:, 0, :], state)
     self.evaluate(tf.global_variables_initializer())
     res1, res2, res3 = self.evaluate([outputs1, outputs2, outputs3])
@@ -378,10 +400,10 @@ def testLSTMCellReparameterizationModel(self):
     res1, res2, res3 = self.evaluate([outputs1, outputs2, outputs3])
     self.assertEqual(res1.shape, (batch_size, timesteps, hidden_size))
     self.assertEqual(res3.shape, (batch_size, timesteps, hidden_size))
-    # NOTE: `cell.sample_weights` should have been called at the beginning of
+    # NOTE: `cell.call_weights` should have been called at the beginning of
     # each call, so these should be different.
     self.assertNotAllClose(res1, res2)
-    # NOTE: We didn't call `cell.sample_weights` again before computing
+    # NOTE: We didn't call `cell.call_weights` again before computing
     # `outputs3`, so the cell should have had the same weights as it did during
     # computation of `outputs2`, and thus yielded the same output tensor.
     self.assertAllClose(res2, res3)

From 861ead837c077493446f34d3d3c1553f778afae3 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 1 Apr 2019 14:38:52 -0700
Subject: [PATCH 1853/2720] Add T2T constraints, initializers, & regularizers
 following Keras.

PiperOrigin-RevId: 241402234
---
 tensor2tensor/keras/constraints.py  |  78 +++
 tensor2tensor/keras/initializers.py | 298 +++++++++
 tensor2tensor/keras/regularizers.py |  94 +++
 tensor2tensor/layers/bayes.py       | 918 +++++++++-------------------
 tensor2tensor/layers/bayes_test.py  |   2 +-
 5 files changed, 776 insertions(+), 614 deletions(-)
 create mode 100644 tensor2tensor/keras/constraints.py
 create mode 100644 tensor2tensor/keras/initializers.py
 create mode 100644 tensor2tensor/keras/regularizers.py

diff --git a/tensor2tensor/keras/constraints.py b/tensor2tensor/keras/constraints.py
new file mode 100644
index 000000000..95283c694
--- /dev/null
+++ b/tensor2tensor/keras/constraints.py
@@ -0,0 +1,78 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Constraints."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import six
+import tensorflow as tf
+
+
+class Positive(tf.keras.constraints.Constraint):
+  """Positive constraint."""
+
+  def __init__(self, epsilon=tf.keras.backend.epsilon()):
+    self.epsilon = epsilon
+
+  def __call__(self, w):
+    return tf.maximum(w, self.epsilon)
+
+  def get_config(self):
+    return {'epsilon': self.epsilon}
+
+
+# Compatibility aliases, following tf.keras
+
+
+positive = Positive  # pylint: disable=invalid-name
+
+# Utility functions, following tf.keras
+
+
+def serialize(initializer):
+  return tf.keras.utils.serialize_keras_object(initializer)
+
+
+def deserialize(config, custom_objects=None):
+  return tf.keras.utils.deserialize_keras_object(
+      config,
+      module_objects=globals(),
+      custom_objects=custom_objects,
+      printable_module_name='constraints')
+
+
+def get(identifier, value=None):
+  """Getter for loading from strings; returns value if can't load."""
+  if value is None:
+    value = identifier
+  if identifier is None:
+    return None
+  elif isinstance(identifier, dict):
+    try:
+      return deserialize(identifier)
+    except ValueError:
+      return value
+  elif isinstance(identifier, six.string_types):
+    config = {'class_name': str(identifier), 'config': {}}
+    try:
+      return deserialize(config)
+    except ValueError:
+      return value
+  elif callable(identifier):
+    return identifier
+  return value
diff --git a/tensor2tensor/keras/initializers.py b/tensor2tensor/keras/initializers.py
new file mode 100644
index 000000000..92c1bf450
--- /dev/null
+++ b/tensor2tensor/keras/initializers.py
@@ -0,0 +1,298 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Initializers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import six
+
+from tensor2tensor.keras import constraints
+from tensor2tensor.keras import regularizers
+import tensorflow as tf
+from tensorflow_probability import edward2 as ed
+
+
+# From `tensorflow/python/ops/init_ops.py`
+def _compute_fans(shape):
+  """Computes the number of input and output units for a weight shape.
+
+  Args:
+    shape: Integer shape tuple or TF tensor shape.
+
+  Returns:
+    A tuple of scalars (fan_in, fan_out).
+  """
+  if len(shape) < 1:  # Just to avoid errors for constants.
+    fan_in = fan_out = 1
+  elif len(shape) == 1:
+    fan_in = fan_out = shape[0]
+  elif len(shape) == 2:
+    fan_in = shape[0]
+    fan_out = shape[1]
+  else:
+    # Assuming convolution kernels (2D, 3D, or more).
+    # kernel shape: (..., input_depth, depth)
+    receptive_field_size = 1.
+    for dim in shape[:-2]:
+      receptive_field_size *= dim
+    fan_in = shape[-2] * receptive_field_size
+    fan_out = shape[-1] * receptive_field_size
+  if isinstance(fan_in, tf.Dimension):
+    fan_in = fan_in.value
+  if isinstance(fan_out, tf.Dimension):
+    fan_out = fan_out.value
+  return fan_in, fan_out
+
+
+class ScaledNormalStdDev(tf.keras.initializers.VarianceScaling):
+  """Initializer capable of adapting its scale to the shape of weights tensors.
+
+  This initializes the standard deviation parameter of a Trainable Normal
+  distribution with a scale based on the shape of the weights tensor.
+  Additionally, A small amount of noise will be added to break weigh symmetry.
+
+  With `distribution="truncated_normal" or "untruncated_normal"`, the standard
+  deviation (after truncation, if used) is `stddev = sqrt(scale / n)`, where n
+  is:
+    - number of input units in the weight tensor, if mode = "fan_in"
+    - number of output units, if mode = "fan_out"
+    - average of the numbers of input and output units, if mode = "fan_avg"
+
+  Args:
+    scale: Scaling factor (positive float).
+    mode: One of "fan_in", "fan_out", "fan_avg".
+    distribution: Random distribution to use. One of "truncated_normal", or
+      "untruncated_normal".
+    seed: A Python integer. Used to create random seeds. See
+      `tf.set_random_seed`
+      for behavior.
+    dtype: The data type. Only floating point types are supported.
+
+  Raises:
+    ValueError: In case of an invalid value for the "scale", mode" or
+      "distribution" arguments.
+  """
+
+  def __init__(self,
+               scale=1.0,
+               mode='fan_in',
+               distribution='untruncated_normal',
+               seed=None,
+               dtype=tf.float32):
+    distribution = distribution.lower()
+    if distribution not in {'truncated_normal', 'untruncated_normal'}:
+      raise ValueError('Invalid `distribution` argument:', distribution)
+    super(ScaledNormalStdDev, self).__init__(scale=scale, mode=mode,
+                                             distribution=distribution,
+                                             seed=seed, dtype=dtype)
+
+  def __call__(self, shape, dtype=None, partition_info=None):
+    if dtype is None:
+      dtype = self.dtype
+    scale = self.scale
+    scale_shape = shape
+    if partition_info is not None:
+      scale_shape = partition_info.full_shape
+    fan_in, fan_out = _compute_fans(scale_shape)
+    if self.mode == 'fan_in':
+      scale /= max(1., fan_in)
+    elif self.mode == 'fan_out':
+      scale /= max(1., fan_out)
+    else:
+      scale /= max(1., (fan_in + fan_out) / 2.)
+    if self.distribution == 'truncated_normal':
+      # constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.)
+      stddev = math.sqrt(scale) / .87962566103423978
+    else:  # self.distribution == 'untruncated_normal':
+      stddev = math.sqrt(scale)
+    return tf.random.truncated_normal(shape, mean=stddev, stddev=stddev*0.1,
+                                      dtype=dtype)
+
+
+class TrainableNormal(tf.keras.layers.Layer):
+  """Random normal op as an initializer with trainable mean and stddev."""
+
+  def __init__(self,
+               mean_initializer=tf.keras.initializers.truncated_normal(
+                   stddev=1e-5),
+               stddev_initializer='scaled_normal_std_dev',
+               mean_regularizer=None,
+               stddev_regularizer=None,
+               mean_constraint=None,
+               stddev_constraint='positive',
+               seed=None,
+               dtype=tf.float32,
+               **kwargs):
+    """Constructs the initializer."""
+    super(TrainableNormal, self).__init__(dtype=dtype, **kwargs)
+    self.mean_initializer = get(mean_initializer)
+    self.stddev_initializer = get(stddev_initializer)
+    self.mean_regularizer = regularizers.get(mean_regularizer)
+    self.stddev_regularizer = regularizers.get(stddev_regularizer)
+    self.mean_constraint = constraints.get(mean_constraint)
+    self.stddev_constraint = constraints.get(stddev_constraint)
+    self.seed = seed
+
+  def build(self, shape, dtype=None):
+    if dtype is None:
+      dtype = self.dtype
+
+    self.mean = self.add_weight(
+        'mean',
+        shape=shape,
+        initializer=self.mean_initializer,
+        regularizer=self.mean_regularizer,
+        constraint=self.mean_constraint,
+        dtype=dtype,
+        trainable=True)
+    self.stddev = self.add_weight(
+        'stddev',
+        shape=shape,
+        initializer=self.stddev_initializer,
+        regularizer=self.stddev_regularizer,
+        constraint=self.stddev_constraint,
+        dtype=dtype,
+        trainable=True)
+    self.built = True
+
+  def __call__(self, shape, dtype=None, partition_info=None):
+    del partition_info  # unused arg
+    if not self.built:
+      self.build(shape, dtype)
+    return ed.Independent(
+        ed.Normal(loc=self.mean, scale=self.stddev).distribution,
+        reinterpreted_batch_ndims=len(shape))
+
+  def get_config(self):
+    return {
+        'mean_initializer':
+            tf.keras.initializers.serialize(self.mean_initializer),
+        'stddev_initializer':
+            tf.keras.initializers.serialize(self.stddev_initializer),
+        'mean_regularizer':
+            tf.keras.regularizers.serialize(self.mean_regularizer),
+        'stddev_regularizer':
+            tf.keras.regularizers.serialize(self.stddev_regularizer),
+        'mean_constraint':
+            tf.keras.constraints.serialize(self.mean_constraint),
+        'stddev_constraint':
+            tf.keras.constraints.serialize(self.stddev_constraint),
+        'seed': self.seed,
+        'dtype': self.dtype,
+    }
+
+
+class TrainableHeNormal(TrainableNormal):
+  """Trainable normal initialized per He et al. 2015, given a ReLU nonlinearity.
+
+  The distribution is initialized to a Normal scaled by `sqrt(2 / fan_in)`,
+  where `fan_in` is the number of input units. A ReLU nonlinearity is assumed
+  for this initialization scheme.
+
+  References:
+    He K, Zhang X, Ren S, Sun J. Delving deep into rectifiers: Surpassing
+    human-level performance on imagenet classification. In Proceedings of the
+    IEEE international conference on computer vision 2015 (pp. 1026-1034).
+    https://arxiv.org/abs/1502.01852
+  """
+
+  def __init__(self, seed=None, dtype=tf.float32):
+    super(TrainableHeNormal, self).__init__(
+        stddev_initializer=ScaledNormalStdDev(scale=2.0, seed=seed,
+                                              dtype=dtype),
+        seed=seed, dtype=dtype)
+
+  def get_config(self):
+    return {
+        'seed': self.seed,
+        'dtype': self.dtype,
+    }
+
+
+class TrainableGlorotNormal(TrainableNormal):
+  """Trainable normal initialized per Glorot and Bengio, 2010.
+
+  The distribution is initialized to a Normal scaled by `sqrt(2 / fan_in +
+  fan_out)`, where `fan_in` is the number of input units and `fan_out` is the
+  number of output units.
+
+  References:
+    Glorot X, Bengio Y. Understanding the difficulty of training deep
+    feedforward neural networks. In Proceedings of the thirteenth international
+    conference on artificial intelligence and statistics 2010 Mar 31 (pp.
+    249-256). http://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf
+  """
+
+  def __init__(self, seed=None, dtype=tf.float32):
+    super(TrainableGlorotNormal, self).__init__(
+        stddev_initializer=ScaledNormalStdDev(mode='fan_avg', seed=seed,
+                                              dtype=dtype),
+        seed=seed, dtype=dtype)
+
+  def get_config(self):
+    return {
+        'seed': self.seed,
+        'dtype': self.dtype
+    }
+
+
+# Compatibility aliases, following tf.keras
+
+# pylint: disable=invalid-name
+scaled_normal_std_dev = ScaledNormalStdDev
+trainable_normal = TrainableNormal
+trainable_he_normal = TrainableHeNormal
+trainable_glorot_normal = TrainableGlorotNormal
+# pylint: enable=invalid-name
+
+# Utility functions, following tf.keras
+
+
+def serialize(initializer):
+  return tf.keras.utils.serialize_keras_object(initializer)
+
+
+def deserialize(config, custom_objects=None):
+  return tf.keras.utils.deserialize_keras_object(
+      config,
+      module_objects=globals(),
+      custom_objects=custom_objects,
+      printable_module_name='initializers')
+
+
+def get(identifier, value=None):
+  """Getter for loading from strings; returns value if can't load."""
+  if value is None:
+    value = identifier
+  if identifier is None:
+    return None
+  elif isinstance(identifier, dict):
+    try:
+      return deserialize(identifier)
+    except ValueError:
+      return value
+  elif isinstance(identifier, six.string_types):
+    config = {'class_name': str(identifier), 'config': {}}
+    try:
+      return deserialize(config)
+    except ValueError:
+      return value
+  elif callable(identifier):
+    return identifier
+  return value
diff --git a/tensor2tensor/keras/regularizers.py b/tensor2tensor/keras/regularizers.py
new file mode 100644
index 000000000..635d8a794
--- /dev/null
+++ b/tensor2tensor/keras/regularizers.py
@@ -0,0 +1,94 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Regularizers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import six
+import tensorflow as tf
+
+from tensorflow_probability import edward2 as ed
+
+
+class NormalKLDivergence(tf.keras.regularizers.Regularizer):
+  """KL divergence regularizer from one normal distribution to another."""
+
+  def __init__(self, mean=0., stddev=1.):
+    """Construct regularizer where default is a KL towards the std normal."""
+    self.mean = mean
+    self.stddev = stddev
+
+  def __call__(self, x):
+    """Computes regularization given an ed.Normal random variable as input."""
+    if not isinstance(x, ed.RandomVariable):
+      raise ValueError('Input must be an ed.RandomVariable.')
+    random_variable = ed.Independent(
+        ed.Normal(
+            loc=tf.broadcast_to(self.mean, x.distribution.event_shape),
+            scale=tf.broadcast_to(self.stddev, x.distribution.event_shape)
+        ).distribution,
+        reinterpreted_batch_ndims=len(x.distribution.event_shape))
+    return random_variable.distribution.kl_divergence(x.distribution)
+
+  def get_config(self):
+    return {
+        'mean': self.mean,
+        'stddev': self.stddev,
+    }
+
+
+# Compatibility aliases, following tf.keras
+
+
+normal_kl_divergence = NormalKLDivergence  # pylint: disable=invalid-name
+
+# Utility functions, following tf.keras
+
+
+def serialize(initializer):
+  return tf.keras.utils.serialize_keras_object(initializer)
+
+
+def deserialize(config, custom_objects=None):
+  return tf.keras.utils.deserialize_keras_object(
+      config,
+      module_objects=globals(),
+      custom_objects=custom_objects,
+      printable_module_name='regularizers')
+
+
+def get(identifier, value=None):
+  """Getter for loading from strings; returns value if can't load."""
+  if value is None:
+    value = identifier
+  if identifier is None:
+    return None
+  elif isinstance(identifier, dict):
+    try:
+      return deserialize(identifier)
+    except ValueError:
+      return value
+  elif isinstance(identifier, six.string_types):
+    config = {'class_name': str(identifier), 'config': {}}
+    try:
+      return deserialize(config)
+    except ValueError:
+      return value
+  elif callable(identifier):
+    return identifier
+  return value
diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index c1a99ccc9..69ae77430 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -20,362 +20,15 @@
 from __future__ import print_function
 
 import functools
-import math
+from tensor2tensor.keras import constraints
+from tensor2tensor.keras import initializers
+from tensor2tensor.keras import regularizers
 
-import six
 import tensorflow as tf
 import tensorflow_probability as tfp
-
 from tensorflow_probability import edward2 as ed
 
 
-class Positive(tf.keras.constraints.Constraint):
-  """Positive constraint."""
-
-  def __init__(self, epsilon=tf.keras.backend.epsilon()):
-    self.epsilon = epsilon
-
-  def __call__(self, w):
-    return tf.maximum(w, self.epsilon)
-
-  def get_config(self):
-    return {'epsilon': self.epsilon}
-
-
-class Zeros(object):
-  """Function returning zeros tensor of same shape excluding the last dim."""
-
-  def __call__(self, inputs):
-    return tf.zeros(tf.shape(inputs)[:-1], inputs.dtype)
-
-  def get_config(self):
-    return {}
-
-
-class ExponentiatedQuadratic(object):
-  """Exponentiated quadratic kernel."""
-
-  def __init__(self, variance, lengthscale):
-    self.variance = variance
-    self.lengthscale = lengthscale
-
-  def __call__(self, x1, x2):
-    """Computes exponentiated quadratic over all pairs of inputs.
-
-    Args:
-      x1: Tensor of shape [batch_x1, ...]. Slices along the batch axis denote an
-        individual input to be passed to the kernel. It is computed pairwise
-        with each input sliced from x2.
-      x2: Tensor of shape [batch_x2, ...]. Slices along the batch axis denote an
-        individual input passed to the kernel function. It is computed pairwise
-        with each input sliced from x1.
-
-    Returns:
-      Tensor of shape [batch_x1, batch_x2].
-    """
-    size = tf.convert_to_tensor(x1).shape.ndims
-    if size > 2:
-      raise NotImplementedError('Multiple feature dimensions is not yet '
-                                'supported.')
-    x1 = x1 / self.lengthscale
-    x2 = x2 / self.lengthscale
-    x1_squared = tf.reduce_sum(tf.square(x1), list(range(1, len(x1.shape))))
-    x2_squared = tf.reduce_sum(tf.square(x2), list(range(1, len(x2.shape))))
-    square = (x1_squared[:, tf.newaxis] +
-              x2_squared[tf.newaxis, :] -
-              2 * tf.matmul(x1, x2, transpose_b=True))
-    return self.variance * tf.exp(-square / 2)
-
-  def get_config(self):
-    return {'variance': self.variance, 'lengthscale': self.lengthscale}
-
-
-class LinearKernel(object):
-  """Linear kernel, optionally on top of a feature extractor (e.g., encoder)."""
-
-  def __init__(self, variance, bias, encoder=tf.identity):
-    self.variance = variance
-    self.bias = bias
-    self.encoder = encoder
-
-  def __call__(self, x1, x2):
-    """Computes scaled dot product of over all pairs of encoded inputs.
-
-    Args:
-      x1: Tensor of shape [batch_x1] + encoder domain. Slices along the batch
-        axis denote an individual input to be passed to the kernel. It is
-        computed pairwise with each input sliced from x2.
-      x2: Tensor of shape [batch_x2] + encoder domain. Slices along the batch
-        axis denote an individual input to be passed to the kernel. It is
-        computed pairwise with each input sliced from x1.
-
-    Returns:
-      Tensor of shape [batch_x1, batch_x2].
-    """
-    encoded_x1 = self.encoder(x1)
-    encoded_x2 = self.encoder(x2)
-    dot_product = tf.matmul(encoded_x1, encoded_x2, transpose_b=True)
-    return self.variance * dot_product + self.bias
-
-  def get_config(self):
-    return {
-        'variance': self.variance,
-        'bias': self.bias,
-        'encoder': tf.keras.utils.serialize_keras_object(self.encoder),
-    }
-
-
-# From `tensorflow/python/ops/init_ops.py`
-def _compute_fans(shape):
-  """Computes the number of input and output units for a weight shape.
-
-  Args:
-    shape: Integer shape tuple or TF tensor shape.
-
-  Returns:
-    A tuple of scalars (fan_in, fan_out).
-  """
-  if len(shape) < 1:  # Just to avoid errors for constants.
-    fan_in = fan_out = 1
-  elif len(shape) == 1:
-    fan_in = fan_out = shape[0]
-  elif len(shape) == 2:
-    fan_in = shape[0]
-    fan_out = shape[1]
-  else:
-    # Assuming convolution kernels (2D, 3D, or more).
-    # kernel shape: (..., input_depth, depth)
-    receptive_field_size = 1.
-    for dim in shape[:-2]:
-      receptive_field_size *= dim
-    fan_in = shape[-2] * receptive_field_size
-    fan_out = shape[-1] * receptive_field_size
-  if isinstance(fan_in, tf.Dimension):
-    fan_in = fan_in.value
-  if isinstance(fan_out, tf.Dimension):
-    fan_out = fan_out.value
-  return fan_in, fan_out
-
-
-class ScaledNormalStdDev(tf.keras.initializers.VarianceScaling):
-  """Initializer capable of adapting its scale to the shape of weights tensors.
-
-  This initializes the standard deviation parameter of a Trainable Normal
-  distribution with a scale based on the shape of the weights tensor.
-  Additionally, A small amount of noise will be added to break weigh symmetry.
-
-  With `distribution="truncated_normal" or "untruncated_normal"`, the standard
-  deviation (after truncation, if used) is `stddev = sqrt(scale / n)`, where n
-  is:
-    - number of input units in the weight tensor, if mode = "fan_in"
-    - number of output units, if mode = "fan_out"
-    - average of the numbers of input and output units, if mode = "fan_avg"
-
-  Args:
-    scale: Scaling factor (positive float).
-    mode: One of "fan_in", "fan_out", "fan_avg".
-    distribution: Random distribution to use. One of "truncated_normal", or
-      "untruncated_normal".
-    seed: A Python integer. Used to create random seeds. See
-      `tf.set_random_seed`
-      for behavior.
-    dtype: The data type. Only floating point types are supported.
-
-  Raises:
-    ValueError: In case of an invalid value for the "scale", mode" or
-      "distribution" arguments.
-  """
-
-  def __init__(self,
-               scale=1.0,
-               mode='fan_in',
-               distribution='untruncated_normal',
-               seed=None,
-               dtype=tf.float32):
-    distribution = distribution.lower()
-    if distribution not in {'truncated_normal', 'untruncated_normal'}:
-      raise ValueError('Invalid `distribution` argument:', distribution)
-    super(ScaledNormalStdDev, self).__init__(scale=scale, mode=mode,
-                                             distribution=distribution,
-                                             seed=seed, dtype=dtype)
-
-  def __call__(self, shape, dtype=None, partition_info=None):
-    if dtype is None:
-      dtype = self.dtype
-    scale = self.scale
-    scale_shape = shape
-    if partition_info is not None:
-      scale_shape = partition_info.full_shape
-    fan_in, fan_out = _compute_fans(scale_shape)
-    if self.mode == 'fan_in':
-      scale /= max(1., fan_in)
-    elif self.mode == 'fan_out':
-      scale /= max(1., fan_out)
-    else:
-      scale /= max(1., (fan_in + fan_out) / 2.)
-    if self.distribution == 'truncated_normal':
-      # constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.)
-      stddev = math.sqrt(scale) / .87962566103423978
-    else:  # self.distribution == 'untruncated_normal':
-      stddev = math.sqrt(scale)
-    return tf.random.truncated_normal(shape, mean=stddev, stddev=stddev*0.1,
-                                      dtype=dtype)
-
-
-class TrainableNormal(tf.keras.layers.Layer):
-  """Random normal op as an initializer with trainable mean and stddev."""
-
-  def __init__(self,
-               mean_initializer=tf.keras.initializers.truncated_normal(
-                   stddev=1e-5),
-               stddev_initializer='scaled_normal_std_dev',
-               mean_regularizer=None,
-               stddev_regularizer=None,
-               mean_constraint=None,
-               stddev_constraint='positive',
-               seed=None,
-               dtype=tf.float32,
-               **kwargs):
-    """Constructs the initializer."""
-    super(TrainableNormal, self).__init__(dtype=dtype, **kwargs)
-    self.mean_initializer = get(mean_initializer)
-    self.stddev_initializer = get(stddev_initializer)
-    self.mean_regularizer = get(mean_regularizer)
-    self.stddev_regularizer = get(stddev_regularizer)
-    self.mean_constraint = get(mean_constraint)
-    self.stddev_constraint = get(stddev_constraint)
-    self.seed = seed
-
-  def build(self, shape, dtype=None):
-    if dtype is None:
-      dtype = self.dtype
-
-    self.mean = self.add_weight(
-        'mean',
-        shape=shape,
-        initializer=self.mean_initializer,
-        regularizer=self.mean_regularizer,
-        constraint=self.mean_constraint,
-        dtype=dtype,
-        trainable=True)
-    self.stddev = self.add_weight(
-        'stddev',
-        shape=shape,
-        initializer=self.stddev_initializer,
-        regularizer=self.stddev_regularizer,
-        constraint=self.stddev_constraint,
-        dtype=dtype,
-        trainable=True)
-    self.built = True
-
-  def __call__(self, shape, dtype=None, partition_info=None):
-    del partition_info  # unused arg
-    if not self.built:
-      self.build(shape, dtype)
-    return ed.Independent(
-        ed.Normal(loc=self.mean, scale=self.stddev).distribution,
-        reinterpreted_batch_ndims=len(shape))
-
-  def get_config(self):
-    return {
-        'mean_initializer':
-            tf.keras.initializers.serialize(self.mean_initializer),
-        'stddev_initializer':
-            tf.keras.initializers.serialize(self.stddev_initializer),
-        'mean_regularizer':
-            tf.keras.regularizers.serialize(self.mean_regularizer),
-        'stddev_regularizer':
-            tf.keras.regularizers.serialize(self.stddev_regularizer),
-        'mean_constraint':
-            tf.keras.constraints.serialize(self.mean_constraint),
-        'stddev_constraint':
-            tf.keras.constraints.serialize(self.stddev_constraint),
-        'seed': self.seed,
-        'dtype': self.dtype,
-    }
-
-
-class TrainableHeNormal(TrainableNormal):
-  """Trainable normal initialized per He et al. 2015, given a ReLU nonlinearity.
-
-  The distribution is initialized to a Normal scaled by `sqrt(2 / fan_in)`,
-  where `fan_in` is the number of input units. A ReLU nonlinearity is assumed
-  for this initialization scheme.
-
-  References:
-    He K, Zhang X, Ren S, Sun J. Delving deep into rectifiers: Surpassing
-    human-level performance on imagenet classification. In Proceedings of the
-    IEEE international conference on computer vision 2015 (pp. 1026-1034).
-    https://arxiv.org/abs/1502.01852
-  """
-
-  def __init__(self, seed=None, dtype=tf.float32):
-    super(TrainableHeNormal, self).__init__(
-        stddev_initializer=ScaledNormalStdDev(scale=2.0, seed=seed,
-                                              dtype=dtype),
-        seed=seed, dtype=dtype)
-
-  def get_config(self):
-    return {
-        'seed': self.seed,
-        'dtype': self.dtype,
-    }
-
-
-class TrainableGlorotNormal(TrainableNormal):
-  """Trainable normal initialized per Glorot and Bengio, 2010.
-
-  The distribution is initialized to a Normal scaled by `sqrt(2 / fan_in +
-  fan_out)`, where `fan_in` is the number of input units and `fan_out` is the
-  number of output units.
-
-  References:
-    Glorot X, Bengio Y. Understanding the difficulty of training deep
-    feedforward neural networks. In Proceedings of the thirteenth international
-    conference on artificial intelligence and statistics 2010 Mar 31 (pp.
-    249-256). http://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf
-  """
-
-  def __init__(self, seed=None, dtype=tf.float32):
-    super(TrainableGlorotNormal, self).__init__(
-        stddev_initializer=ScaledNormalStdDev(mode='fan_avg', seed=seed,
-                                              dtype=dtype),
-        seed=seed, dtype=dtype)
-
-  def get_config(self):
-    return {
-        'seed': self.seed,
-        'dtype': self.dtype
-    }
-
-
-class NormalKLDivergence(tf.keras.regularizers.Regularizer):
-  """KL divergence regularizer from one normal distribution to another."""
-
-  def __init__(self, mean=0., stddev=1.):
-    """Construct regularizer where default is a KL towards the std normal."""
-    self.mean = mean
-    self.stddev = stddev
-
-  def __call__(self, x):
-    """Computes regularization given an ed.Normal random variable as input."""
-    if not isinstance(x, ed.RandomVariable):
-      raise ValueError('Input must be an ed.RandomVariable.')
-    random_variable = ed.Independent(
-        ed.Normal(
-            loc=tf.broadcast_to(self.mean, x.distribution.event_shape),
-            scale=tf.broadcast_to(self.stddev, x.distribution.event_shape)
-        ).distribution,
-        reinterpreted_batch_ndims=len(x.distribution.event_shape))
-    return random_variable.distribution.kl_divergence(x.distribution)
-
-  def get_config(self):
-    return {
-        'mean': self.mean,
-        'stddev': self.stddev,
-    }
-
-
 def add_weight(cls):
   """Decorator for Layers, overriding add_weight for trainable initializers."""
   @functools.wraps(cls.add_weight)
@@ -411,6 +64,73 @@ def loss_fn():
   return cls
 
 
+@add_weight
+class Conv2DReparameterization(tf.keras.layers.Conv2D):
+  """2D convolution layer (e.g. spatial convolution over images).
+
+  The layer computes a variational Bayesian approximation to the distribution
+  over convolutional layers,
+
+  ```
+  p(outputs | inputs) = int conv2d(inputs; weights, bias) p(weights, bias)
+    dweights dbias.
+  ```
+
+  It does this with a stochastic forward pass, sampling from learnable
+  distributions on the kernel and bias. Gradients with respect to the
+  distributions' learnable parameters backpropagate via reparameterization.
+  Minimizing cross-entropy plus the layer's losses performs variational
+  minimum description length, i.e., it minimizes an upper bound to the negative
+  marginal likelihood.
+  """
+
+  def __init__(self,
+               filters,
+               kernel_size,
+               strides=(1, 1),
+               padding='valid',
+               data_format=None,
+               dilation_rate=(1, 1),
+               activation=None,
+               use_bias=True,
+               kernel_initializer='trainable_normal',
+               bias_initializer='zeros',
+               kernel_regularizer='normal_kl_divergence',
+               bias_regularizer=None,
+               activity_regularizer=None,
+               kernel_constraint=None,
+               bias_constraint=None,
+               **kwargs):
+    super(Conv2DReparameterization, self).__init__(
+        filters=filters,
+        kernel_size=kernel_size,
+        strides=strides,
+        padding=padding,
+        data_format=data_format,
+        dilation_rate=dilation_rate,
+        activation=activation,
+        use_bias=use_bias,
+        kernel_initializer=initializers.get(kernel_initializer),
+        bias_initializer=initializers.get(bias_initializer),
+        kernel_regularizer=regularizers.get(kernel_regularizer),
+        bias_regularizer=regularizers.get(bias_regularizer),
+        activity_regularizer=regularizers.get(activity_regularizer),
+        kernel_constraint=constraints.get(kernel_constraint),
+        bias_constraint=constraints.get(bias_constraint),
+        **kwargs)
+
+  def call_weights(self):
+    """Calls any weights if the initializer is itself a layer."""
+    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
+      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
+    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
+      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
+
+  def call(self, *args, **kwargs):
+    self.call_weights()
+    return super(Conv2DReparameterization, self).call(*args, **kwargs)
+
+
 @add_weight
 class DenseReparameterization(tf.keras.layers.Dense):
   """Bayesian densely-connected layer estimated via reparameterization.
@@ -431,104 +151,251 @@ class DenseReparameterization(tf.keras.layers.Dense):
   marginal likelihood.
   """
 
-  def __init__(self,
-               units,
-               activation=None,
-               use_bias=True,
-               kernel_initializer='trainable_normal',
-               bias_initializer='zero',
-               kernel_regularizer='normal_kl_divergence',
-               bias_regularizer=None,
-               activity_regularizer=None,
-               **kwargs):
-    super(DenseReparameterization, self).__init__(
-        units=units,
-        activation=get(activation),
-        use_bias=use_bias,
-        kernel_initializer=get(kernel_initializer),
-        bias_initializer=get(bias_initializer),
-        kernel_regularizer=get(kernel_regularizer),
-        bias_regularizer=get(bias_regularizer),
-        activity_regularizer=get(activity_regularizer),
-        **kwargs)
+  def __init__(self,
+               units,
+               activation=None,
+               use_bias=True,
+               kernel_initializer='trainable_normal',
+               bias_initializer='zero',
+               kernel_regularizer='normal_kl_divergence',
+               bias_regularizer=None,
+               activity_regularizer=None,
+               **kwargs):
+    super(DenseReparameterization, self).__init__(
+        units=units,
+        activation=activation,
+        use_bias=use_bias,
+        kernel_initializer=initializers.get(kernel_initializer),
+        bias_initializer=initializers.get(bias_initializer),
+        kernel_regularizer=regularizers.get(kernel_regularizer),
+        bias_regularizer=regularizers.get(bias_regularizer),
+        activity_regularizer=regularizers.get(activity_regularizer),
+        **kwargs)
+
+  def call_weights(self):
+    """Calls any weights if the initializer is itself a layer."""
+    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
+      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
+    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
+      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
+
+  def call(self, *args, **kwargs):
+    self.call_weights()
+    return super(DenseReparameterization, self).call(*args, **kwargs)
+
+
+@add_weight
+class LSTMCellReparameterization(tf.keras.layers.LSTMCell):
+  """Bayesian LSTM cell class estimated via reparameterization.
+
+  The layer computes a variational Bayesian approximation to the distribution
+  over LSTM cell functions,
+
+  ```
+  p(outputs | inputs) = int lstm_cell(inputs; weights, bias) p(weights, bias)
+    dweights dbias,
+  ```
+
+  where the weights consist of both input and recurrent weights.
+
+  It does this with a stochastic forward pass, sampling from learnable
+  distributions on the kernel, recurrent kernel, and bias. Gradients with
+  respect to the distributions' learnable parameters backpropagate via
+  reparameterization.  Minimizing cross-entropy plus the layer's losses performs
+  variational minimum description length, i.e., it minimizes an upper bound to
+  the negative marginal likelihood.
+  """
+
+  def __init__(self,
+               units,
+               activation='tanh',
+               recurrent_activation='hard_sigmoid',
+               use_bias=True,
+               kernel_initializer='trainable_normal',
+               recurrent_initializer='trainable_normal',
+               bias_initializer='zeros',
+               unit_forget_bias=True,
+               kernel_regularizer='normal_kl_divergence',
+               recurrent_regularizer='normal_kl_divergence',
+               bias_regularizer=None,
+               kernel_constraint=None,
+               recurrent_constraint=None,
+               bias_constraint=None,
+               dropout=0.,
+               recurrent_dropout=0.,
+               implementation=1,
+               **kwargs):
+    super(LSTMCellReparameterization, self).__init__(
+        units=units,
+        activation=activation,
+        recurrent_activation=recurrent_activation,
+        use_bias=use_bias,
+        kernel_initializer=initializers.get(kernel_initializer),
+        recurrent_initializer=initializers.get(recurrent_initializer),
+        bias_initializer=initializers.get(bias_initializer),
+        unit_forget_bias=unit_forget_bias,
+        kernel_regularizer=regularizers.get(kernel_regularizer),
+        recurrent_regularizer=regularizers.get(recurrent_regularizer),
+        bias_regularizer=regularizers.get(bias_regularizer),
+        kernel_constraint=constraints.get(kernel_constraint),
+        recurrent_constraint=constraints.get(recurrent_constraint),
+        bias_constraint=constraints.get(bias_constraint),
+        dropout=dropout,
+        recurrent_dropout=recurrent_dropout,
+        implementation=implementation,
+        **kwargs)
+
+  def build(self, input_shape):
+    input_shape = tf.TensorShape(input_shape)
+    input_dim = input_shape[-1]
+    if isinstance(input_dim, tf.Dimension):
+      input_dim = input_dim.value
+    self.kernel = self.add_weight(
+        shape=(input_dim, self.units * 4),
+        name='kernel',
+        initializer=self.kernel_initializer,
+        regularizer=self.kernel_regularizer,
+        constraint=self.kernel_constraint)
+    self.recurrent_kernel = self.add_weight(
+        shape=(self.units, self.units * 4),
+        name='recurrent_kernel',
+        initializer=self.recurrent_initializer,
+        regularizer=self.recurrent_regularizer,
+        constraint=self.recurrent_constraint)
+
+    if self.use_bias:
+      if self.unit_forget_bias:
+        if isinstance(self.bias_initializer, tf.keras.layers.Layer):
+          def bias_mean_initializer(_, *args, **kwargs):
+            return tf.concat([
+                tf.keras.initializers.truncated_normal(
+                    stddev=1e-5)((self.units,), *args, **kwargs),
+                tf.keras.initializers.truncated_normal(
+                    mean=1., stddev=1e-5)((self.units,), *args, **kwargs),
+                tf.keras.initializers.truncated_normal(
+                    stddev=1e-5)((self.units * 2,), *args, **kwargs),
+            ], axis=0)
+          bias_initializer = initializers.TrainableNormal(
+              mean_initializer=bias_mean_initializer)
+        else:
+          def bias_initializer(_, *args, **kwargs):
+            return tf.keras.backend.concatenate([
+                self.bias_initializer((self.units,), *args, **kwargs),
+                tf.keras.initializers.Ones()((self.units,), *args, **kwargs),
+                self.bias_initializer((self.units * 2,), *args, **kwargs),
+            ])
+      else:
+        bias_initializer = self.bias_initializer
+      self.bias = self.add_weight(
+          shape=(self.units * 4,),
+          name='bias',
+          initializer=bias_initializer,
+          regularizer=self.bias_regularizer,
+          constraint=self.bias_constraint)
+    else:
+      self.bias = None
+    self.built = True
+
+  def call_weights(self):
+    """Calls any weights if the initializer is itself a layer."""
+    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
+      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
+    if isinstance(self.recurrent_initializer, tf.keras.layers.Layer):
+      self.recurrent_kernel = self.recurrent_initializer(
+          self.recurrent_kernel.shape, self.dtype)
+    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
+      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
+
+  # NOTE: This will not be called in TF < 1.11.
+  def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
+    """Get the initial state and side-effect sampling of stochastic weights."""
+    if self.built:
+      self.call_weights()
+    return super(LSTMCellReparameterization, self).get_initial_state(
+        inputs=inputs, batch_size=batch_size, dtype=dtype)
+
+
+class Zeros(object):
+  """Function returning zeros tensor of same shape excluding the last dim."""
+
+  def __call__(self, inputs):
+    return tf.zeros(tf.shape(inputs)[:-1], inputs.dtype)
+
+  def get_config(self):
+    return {}
+
+
+class ExponentiatedQuadratic(object):
+  """Exponentiated quadratic kernel."""
 
-  def call_weights(self):
-    """Calls any weights if the initializer is itself a layer."""
-    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
-      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
-    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
+  def __init__(self, variance, lengthscale):
+    self.variance = variance
+    self.lengthscale = lengthscale
 
-  def call(self, *args, **kwargs):
-    self.call_weights()
-    return super(DenseReparameterization, self).call(*args, **kwargs)
+  def __call__(self, x1, x2):
+    """Computes exponentiated quadratic over all pairs of inputs.
 
+    Args:
+      x1: Tensor of shape [batch_x1, ...]. Slices along the batch axis denote an
+        individual input to be passed to the kernel. It is computed pairwise
+        with each input sliced from x2.
+      x2: Tensor of shape [batch_x2, ...]. Slices along the batch axis denote an
+        individual input passed to the kernel function. It is computed pairwise
+        with each input sliced from x1.
 
-@add_weight
-class Conv2DReparameterization(tf.keras.layers.Conv2D):
-  """2D convolution layer (e.g. spatial convolution over images).
+    Returns:
+      Tensor of shape [batch_x1, batch_x2].
+    """
+    size = tf.convert_to_tensor(x1).shape.ndims
+    if size > 2:
+      raise NotImplementedError('Multiple feature dimensions is not yet '
+                                'supported.')
+    x1 = x1 / self.lengthscale
+    x2 = x2 / self.lengthscale
+    x1_squared = tf.reduce_sum(tf.square(x1), list(range(1, len(x1.shape))))
+    x2_squared = tf.reduce_sum(tf.square(x2), list(range(1, len(x2.shape))))
+    square = (x1_squared[:, tf.newaxis] +
+              x2_squared[tf.newaxis, :] -
+              2 * tf.matmul(x1, x2, transpose_b=True))
+    return self.variance * tf.exp(-square / 2)
 
-  The layer computes a variational Bayesian approximation to the distribution
-  over convolutional layers,
+  def get_config(self):
+    return {'variance': self.variance, 'lengthscale': self.lengthscale}
 
-  ```
-  p(outputs | inputs) = int conv2d(inputs; weights, bias) p(weights, bias)
-    dweights dbias.
-  ```
 
-  It does this with a stochastic forward pass, sampling from learnable
-  distributions on the kernel and bias. Gradients with respect to the
-  distributions' learnable parameters backpropagate via reparameterization.
-  Minimizing cross-entropy plus the layer's losses performs variational
-  minimum description length, i.e., it minimizes an upper bound to the negative
-  marginal likelihood.
-  """
+class LinearKernel(object):
+  """Linear kernel, optionally on top of a feature extractor (e.g., encoder)."""
 
-  def __init__(self,
-               filters,
-               kernel_size,
-               strides=(1, 1),
-               padding='valid',
-               data_format=None,
-               dilation_rate=(1, 1),
-               activation=None,
-               use_bias=True,
-               kernel_initializer='trainable_normal',
-               bias_initializer='zeros',
-               kernel_regularizer='normal_kl_divergence',
-               bias_regularizer=None,
-               activity_regularizer=None,
-               kernel_constraint=None,
-               bias_constraint=None,
-               **kwargs):
-    super(Conv2DReparameterization, self).__init__(
-        filters=filters,
-        kernel_size=kernel_size,
-        strides=strides,
-        padding=padding,
-        data_format=data_format,
-        dilation_rate=dilation_rate,
-        activation=get(activation),
-        use_bias=use_bias,
-        kernel_initializer=get(kernel_initializer),
-        bias_initializer=get(bias_initializer),
-        kernel_regularizer=get(kernel_regularizer),
-        bias_regularizer=get(bias_regularizer),
-        activity_regularizer=get(activity_regularizer),
-        kernel_constraint=get(kernel_constraint),
-        bias_constraint=get(bias_constraint),
-        **kwargs)
+  def __init__(self, variance, bias, encoder=tf.identity):
+    self.variance = variance
+    self.bias = bias
+    self.encoder = encoder
 
-  def call_weights(self):
-    """Calls any weights if the initializer is itself a layer."""
-    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
-      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
-    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
+  def __call__(self, x1, x2):
+    """Computes scaled dot product of over all pairs of encoded inputs.
 
-  def call(self, *args, **kwargs):
-    self.call_weights()
-    return super(Conv2DReparameterization, self).call(*args, **kwargs)
+    Args:
+      x1: Tensor of shape [batch_x1] + encoder domain. Slices along the batch
+        axis denote an individual input to be passed to the kernel. It is
+        computed pairwise with each input sliced from x2.
+      x2: Tensor of shape [batch_x2] + encoder domain. Slices along the batch
+        axis denote an individual input to be passed to the kernel. It is
+        computed pairwise with each input sliced from x1.
+
+    Returns:
+      Tensor of shape [batch_x1, batch_x2].
+    """
+    encoded_x1 = self.encoder(x1)
+    encoded_x2 = self.encoder(x2)
+    dot_product = tf.matmul(encoded_x1, encoded_x2, transpose_b=True)
+    return self.variance * dot_product + self.bias
+
+  def get_config(self):
+    return {
+        'variance': self.variance,
+        'bias': self.bias,
+        'encoder': tf.keras.utils.serialize_keras_object(self.encoder),
+    }
 
 
 class GaussianProcess(tf.keras.layers.Layer):
@@ -684,137 +551,6 @@ def get_config(self):
     return dict(list(base_config.items()) + list(config.items()))
 
 
-@add_weight
-class LSTMCellReparameterization(tf.keras.layers.LSTMCell):
-  """Bayesian LSTM cell class estimated via reparameterization.
-
-  The layer computes a variational Bayesian approximation to the distribution
-  over LSTM cell functions,
-
-  ```
-  p(outputs | inputs) = int lstm_cell(inputs; weights, bias) p(weights, bias)
-    dweights dbias,
-  ```
-
-  where the weights consist of both input and recurrent weights.
-
-  It does this with a stochastic forward pass, sampling from learnable
-  distributions on the kernel, recurrent kernel, and bias. Gradients with
-  respect to the distributions' learnable parameters backpropagate via
-  reparameterization.  Minimizing cross-entropy plus the layer's losses performs
-  variational minimum description length, i.e., it minimizes an upper bound to
-  the negative marginal likelihood.
-  """
-
-  def __init__(self,
-               units,
-               activation='tanh',
-               recurrent_activation='hard_sigmoid',
-               use_bias=True,
-               kernel_initializer='trainable_normal',
-               recurrent_initializer='trainable_normal',
-               bias_initializer='zeros',
-               unit_forget_bias=True,
-               kernel_regularizer='normal_kl_divergence',
-               recurrent_regularizer='normal_kl_divergence',
-               bias_regularizer=None,
-               kernel_constraint=None,
-               recurrent_constraint=None,
-               bias_constraint=None,
-               dropout=0.,
-               recurrent_dropout=0.,
-               implementation=1,
-               **kwargs):
-    super(LSTMCellReparameterization, self).__init__(
-        units=units,
-        activation=get(activation),
-        recurrent_activation=get(recurrent_activation),
-        use_bias=use_bias,
-        kernel_initializer=get(kernel_initializer),
-        recurrent_initializer=get(recurrent_initializer),
-        bias_initializer=get(bias_initializer),
-        unit_forget_bias=unit_forget_bias,
-        kernel_regularizer=get(kernel_regularizer),
-        recurrent_regularizer=get(recurrent_regularizer),
-        bias_regularizer=get(bias_regularizer),
-        kernel_constraint=get(kernel_constraint),
-        recurrent_constraint=get(recurrent_constraint),
-        bias_constraint=get(bias_constraint),
-        dropout=dropout,
-        recurrent_dropout=recurrent_dropout,
-        implementation=implementation,
-        **kwargs)
-
-  def build(self, input_shape):
-    input_shape = tf.TensorShape(input_shape)
-    input_dim = input_shape[-1]
-    if isinstance(input_dim, tf.Dimension):
-      input_dim = input_dim.value
-    self.kernel = self.add_weight(
-        shape=(input_dim, self.units * 4),
-        name='kernel',
-        initializer=self.kernel_initializer,
-        regularizer=self.kernel_regularizer,
-        constraint=self.kernel_constraint)
-    self.recurrent_kernel = self.add_weight(
-        shape=(self.units, self.units * 4),
-        name='recurrent_kernel',
-        initializer=self.recurrent_initializer,
-        regularizer=self.recurrent_regularizer,
-        constraint=self.recurrent_constraint)
-
-    if self.use_bias:
-      if self.unit_forget_bias:
-        if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-          def bias_mean_initializer(_, *args, **kwargs):
-            return tf.concat([
-                tf.keras.initializers.truncated_normal(
-                    stddev=1e-5)((self.units,), *args, **kwargs),
-                tf.keras.initializers.truncated_normal(
-                    mean=1., stddev=1e-5)((self.units,), *args, **kwargs),
-                tf.keras.initializers.truncated_normal(
-                    stddev=1e-5)((self.units * 2,), *args, **kwargs),
-            ], axis=0)
-          bias_initializer = TrainableNormal(
-              mean_initializer=bias_mean_initializer)
-        else:
-          def bias_initializer(_, *args, **kwargs):
-            return tf.keras.backend.concatenate([
-                self.bias_initializer((self.units,), *args, **kwargs),
-                tf.keras.initializers.Ones()((self.units,), *args, **kwargs),
-                self.bias_initializer((self.units * 2,), *args, **kwargs),
-            ])
-      else:
-        bias_initializer = self.bias_initializer
-      self.bias = self.add_weight(
-          shape=(self.units * 4,),
-          name='bias',
-          initializer=bias_initializer,
-          regularizer=self.bias_regularizer,
-          constraint=self.bias_constraint)
-    else:
-      self.bias = None
-    self.built = True
-
-  def call_weights(self):
-    """Calls any weights if the initializer is itself a layer."""
-    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
-      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
-    if isinstance(self.recurrent_initializer, tf.keras.layers.Layer):
-      self.recurrent_kernel = self.recurrent_initializer(
-          self.recurrent_kernel.shape, self.dtype)
-    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
-
-  # NOTE: This will not be called in TF < 1.11.
-  def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
-    """Get the initial state and side-effect sampling of stochastic weights."""
-    if self.built:
-      self.call_weights()
-    return super(LSTMCellReparameterization, self).get_initial_state(
-        inputs=inputs, batch_size=batch_size, dtype=dtype)
-
-
 class BayesianLinearModel(tf.keras.Model):
   r"""Bayesian linear model with standard normal prior over its coefficients.
 
@@ -901,47 +637,3 @@ def get_config(self):
     config = {'num_components': self.num_components}
     base_config = super(MixtureLogistic, self).get_config()
     return dict(list(base_config.items()) + list(config.items()))
-
-
-# Compatibility aliases, following tf.keras
-
-# pylint: disable=invalid-name
-positive = Positive
-scaled_normal_std_dev = ScaledNormalStdDev
-trainable_normal = TrainableNormal
-trainable_he_normal = TrainableHeNormal
-trainable_glorot_normal = TrainableGlorotNormal
-normal_kl_divergence = NormalKLDivergence
-# pylint: enable=invalid-name
-
-# Utility functions, following tf.keras
-
-
-def deserialize(config, custom_objects=None):
-  return tf.keras.utils.deserialize_keras_object(
-      config,
-      module_objects=globals(),
-      custom_objects=custom_objects,
-      printable_module_name='bayes')
-
-
-def get(identifier, value=None):
-  """Getter for loading from strings; returns value if can't load."""
-  if value is None:
-    value = identifier
-  if identifier is None:
-    return None
-  elif isinstance(identifier, dict):
-    try:
-      return deserialize(identifier)
-    except ValueError:
-      return value
-  elif isinstance(identifier, six.string_types):
-    config = {'class_name': str(identifier), 'config': {}}
-    try:
-      return deserialize(config)
-    except ValueError:
-      return value
-  elif callable(identifier):
-    return identifier
-  return value
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 85a092528..bb8090bb0 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -35,7 +35,7 @@ class BayesTest(parameterized.TestCase, tf.test.TestCase):
   @test_utils.run_in_graph_and_eager_modes
   def testTrainableNormalStddevConstraint(self):
     layer = bayes.DenseReparameterization(
-        100, kernel_initializer=bayes.TrainableNormal())
+        100, kernel_initializer="trainable_normal")
     inputs = tf.random_normal([1, 1])
     out = layer(inputs)
     stddev = layer.kernel.distribution.stddev()

From c84eafe3722035dc632b5b912039427c03d90c20 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 1 Apr 2019 16:47:24 -0700
Subject: [PATCH 1854/2720] Make 1-device mode not call pmap for now, correct
 rng handling in TRAX.

PiperOrigin-RevId: 241427191
---
 tensor2tensor/trax/inputs.py |  6 ++--
 tensor2tensor/trax/trax.py   | 59 +++++++++++++++++++++++-------------
 2 files changed, 41 insertions(+), 24 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index cf7e42e9a..99e7aa71d 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -35,9 +35,9 @@
     "_Inputs", ["train_stream", "eval_stream", "input_shape"])
 
 # How many examples from the stream to skip at random during training.
-# For now, we skip at most 1M examples.
-# TODO(lukaszkaiser): does it matter for efficiency, should that be changed?
-_MAX_SKIP_EXAMPLES = 1e6
+# For now, we skip at most 100K examples for efficiency.
+# TODO(lukaszkaiser): can we improve efficiency, should that be changed?
+_MAX_SKIP_EXAMPLES = 1e5
 
 
 @gin.configurable()
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 4b829eced..75e3cbab5 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -74,10 +74,10 @@ def neg_log_perplexity(batch, model_predictions):
   return masked_mean(xent, targets)
 
 
-def loss(params, batch, model_predict):
+def loss(params, batch, model_predict, rng):
   """Calculate loss."""
   inputs, targets = batch
-  preds = model_predict(params, inputs)
+  preds = model_predict(params, inputs, rng=rng)
   xent = np.sum(preds * stax.one_hot(targets, preds.shape[-1]), axis=-1)
   return - masked_mean(xent, targets)
 
@@ -134,7 +134,7 @@ def save_state(state, output_dir):
 }
 
 
-def evaluate_train_and_eval(step, inputs, predict_fun, eval_steps,
+def evaluate_train_and_eval(step, inputs, predict_fun, eval_steps, rng,
                             train_sw=None, eval_sw=None, history=None):
   """Evalaute on train and eval data, and log metrics."""
   step_log(step, "Evaluation")
@@ -142,7 +142,8 @@ def evaluate_train_and_eval(step, inputs, predict_fun, eval_steps,
       evaluate(  # pylint: disable=g-complex-comprehension
           itertools.islice(input_stream(), eval_steps),
           predict_fun,
-          _METRICS)
+          _METRICS,
+          rng)
       for input_stream in
       [inputs.train_stream, inputs.eval_stream]]
   if train_sw:
@@ -152,7 +153,7 @@ def evaluate_train_and_eval(step, inputs, predict_fun, eval_steps,
   return train_metrics, eval_metrics
 
 
-def evaluate(inputs_stream, predict_fun, metric_funs):
+def evaluate(inputs_stream, predict_fun, metric_funs, rng):
   """Evaluate.
 
   Args:
@@ -161,6 +162,7 @@ def evaluate(inputs_stream, predict_fun, metric_funs):
       partially applied.
     metric_funs: dict from metric name to metric function, which takes inputs
       and predictions and returns a scalar metric value.
+    rng: random number generator.
 
   Returns:
     metrics: dict from metric name to metric value averaged over the number of
@@ -170,7 +172,8 @@ def evaluate(inputs_stream, predict_fun, metric_funs):
   count = 0
   for inp in inputs_stream:
     count += 1
-    preds = predict_fun(inp[0])
+    rng, subrng = jax.random.split(rng)
+    preds = predict_fun(inp[0], rng=subrng)
     for m, f in six.iteritems(metric_funs):
       metrics[m] += f(inp, preds)
   return {m: v / count for (m, v) in six.iteritems(metrics)}
@@ -238,8 +241,17 @@ def epochs(steps=None, epoch_steps=1):
       break
 
 
-def _jit_update_fun(predict_fun, loss_fun, optimizer, lr_fun):
+def _jit_update_fun(predict_fun, loss_fun, optimizer, lr_fun, num_devices):
   """Get jit-ed update function for loss, optimizer, learning rate function."""
+  if num_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
+    @jax.jit
+    def single_update(i, opt_state, batch, rng):
+      _, opt_update = optimizer(lr_fun)
+      params = jax_opt.get_params(opt_state)
+      return opt_update(i, jax.grad(loss_fun)(
+          params, batch, predict_fun, rng), opt_state)
+    return single_update
+
   @functools.partial(jax.pmap, axis_name="batch")
   def mapped_update(i, opt_state, batch):
     _, opt_update = optimizer(lr_fun)
@@ -319,23 +331,20 @@ def train(output_dir,
   history = state.history
   lr_fun = lr_schedule(history)
   opt_init, _ = optimizer(lr_fun)
-  model_init, model_predict_original = model()
-  # We need a model_predict that fills in the random generator if needed.
-  def model_predict(x, y, **kwargs):
-    """Same as model_predict_original but fill in rng if it isn't passed."""
-    if "rng" in kwargs:
-      return model_predict_original(x, y, **kwargs)
-    return model_predict_original(x, y, rng=rng, **kwargs)
+  model_init, model_predict = model()
 
   # Setup state
   step = state.step or 0
   params_initializer = lambda: model_init([-1] + list(inputs.input_shape))[1]
   params = state.params or params_initializer()
-  opt_state = jax.replicate(opt_init(params))
+  opt_state = opt_init(params)
+  if num_devices > 1:  # TODO(lukaszkaiser): use everywhere when pmap is stable.
+    opt_state = jax.replicate(opt_state)
 
   # jit model_predict and update so they're fast
   jit_model_predict = jax.jit(model_predict)  # for evaluation
-  jit_update_fun = _jit_update_fun(model_predict, loss, optimizer, lr_fun)
+  jit_update_fun = _jit_update_fun(model_predict, loss, optimizer, lr_fun,
+                                   num_devices)
 
   print()
   train_stream = inputs.train_stream()
@@ -346,7 +355,7 @@ def model_predict(x, y, **kwargs):
 
   # Non-compiled debug step helps find problems in models easier.
   if run_debug_step:
-    debug_loss = loss(params, next(train_stream), model_predict)
+    debug_loss = loss(params, next(train_stream), model_predict, rng)
     step_log(step, "Debug step loss %.8f" % debug_loss)
 
   for epoch, epoch_steps in epochs(train_steps, epoch_steps):
@@ -358,8 +367,11 @@ def model_predict(x, y, **kwargs):
 
     for _ in range(epoch_steps):
       # Train
-      next_train_batch = reshape_by_device(next(train_stream), num_devices)
-      opt_state = jit_update_fun(step, opt_state, next_train_batch)
+      next_train_batch = next(train_stream)
+      if num_devices > 1:  # TODO(lukaszkaiser): use everywhere when possible.
+        next_train_batch = reshape_by_device(next_train_batch, num_devices)
+      rng, subrng = jax.random.split(rng)
+      opt_state = jit_update_fun(step, opt_state, next_train_batch, subrng)
       step += 1
 
       # LR log
@@ -376,12 +388,16 @@ def model_predict(x, y, **kwargs):
                       epoch_steps / epoch_time, step=step)
 
     # Evaluate
-    params = jax_opt.get_params(jax.unreplicate(opt_state))
+    if num_devices > 1:   # TODO(lukaszkaiser): remove branch when possible.
+      params = jax_opt.get_params(jax.unreplicate(opt_state))
+    else:
+      params = jax_opt.get_params(opt_state)
     evaluate_train_and_eval(
         step=step,
         inputs=inputs,
         predict_fun=functools.partial(jit_model_predict, params),
         eval_steps=eval_steps,
+        rng=rng,
         train_sw=train_sw,
         eval_sw=eval_sw,
         history=history)
@@ -398,7 +414,8 @@ def model_predict(x, y, **kwargs):
     old_lr_fun = lr_fun
     lr_fun = lr_schedule(history)
     if lr_fun != old_lr_fun:  # For performance, only jit if there is a change.
-      jit_update_fun = _jit_update_fun(model_predict, loss, optimizer, lr_fun)
+      jit_update_fun = _jit_update_fun(model_predict, loss, optimizer, lr_fun,
+                                       num_devices)
 
     # Flush summary writers
     train_sw.writer.flush()

From 87a8784a22e954b159d170a9969e1c21273dea77 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 1 Apr 2019 21:49:01 -0700
Subject: [PATCH 1855/2720] Decouple TRAX from t2t_trainer (use trax/trainer
 instead).

PiperOrigin-RevId: 241461284
---
 tensor2tensor/bin/t2t_trainer.py | 44 --------------------------------
 tensor2tensor/trax/trainer.py    |  1 -
 2 files changed, 45 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 586f56331..e317f8078 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -21,7 +21,6 @@
 import contextlib
 import os
 import sys
-import gin
 from tensor2tensor import models  # pylint: disable=unused-import
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
 from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
@@ -38,11 +37,6 @@
 
 from tensorflow.contrib.tpu.python.tpu import tpu_config
 
-try:
-  from tensor2tensor.trax import trax  # pylint: disable=g-import-not-at-top
-except (TypeError, ImportError):
-  pass
-
 
 flags = tf.flags
 FLAGS = flags.FLAGS
@@ -80,7 +74,6 @@
 flags.DEFINE_integer("intra_op_parallelism_threads", 0,
                      "Number of intra_op_parallelism_threads to use for CPU. "
                      "See TensorFlow config.proto for details.")
-flags.DEFINE_bool("jax", False, "Whether to use trax.")
 # TODO(lukaszkaiser): resolve memory and variable assign issues and set to True.
 flags.DEFINE_bool(
     "optionally_use_dist_strat", False,
@@ -371,43 +364,6 @@ def run_std_server():
 def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
 
-  if FLAGS.jax:
-    # Setup trax FLAGS
-    dataset = FLAGS.problem
-    model = FLAGS.model
-    data_dir = FLAGS.data_dir
-    output_dir = FLAGS.output_dir
-    config_file = [FLAGS.hparams_set]
-    config = [
-        "train.train_steps=%d" % FLAGS.train_steps,
-        "train.eval_steps=%d" % FLAGS.eval_steps,
-        "train.eval_frequency=%d" % FLAGS.local_eval_frequency,
-    ] + str(FLAGS.hparams).split(",")
-
-    # Copied _setup_gin exactly from trax/trainer.py and removed "FLAGS."
-
-    def _setup_gin():
-      """Setup gin configuration."""
-      # Imports for configurables
-      # pylint: disable=g-import-not-at-top,unused-import,g-bad-import-order,reimported,unused-variable
-      from tensor2tensor.trax import inputs as _trax_inputs
-      from tensor2tensor.trax import models as _trax_models
-      from tensor2tensor.trax import optimizers as _trax_opt
-      # pylint: enable=g-import-not-at-top,unused-import,g-bad-import-order,reimported,unused-variable
-
-      configs = config or []
-      # Override with --dataset and --model
-      if dataset:
-        configs.append("inputs.dataset_name='%s'" % dataset)
-        configs.append("inputs.data_dir='%s'" % data_dir)
-        configs.append("train.inputs=@trax.inputs.inputs")
-      if model:
-        configs.append("train.model=@trax.models.%s" % model)
-      gin.parse_config_files_and_bindings(config_file, configs)
-
-    _setup_gin()
-    trax.train(output_dir=output_dir)
-
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
 
   # If we just have to print the registry, do that and exit early.
diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index f02f275fc..de6895423 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -45,7 +45,6 @@
 flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
 
 
-
 def _default_output_dir():
   """Default output directory."""
   dir_name = "{model_name}_{dataset_name}_{timestamp}".format(

From 8b41452ce8f8c7aab3b1bc9b5c6a591ad577f17a Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 2 Apr 2019 11:26:25 -0700
Subject: [PATCH 1856/2720] Fix wasteful relative attention when using memory

PiperOrigin-RevId: 241567826
---
 tensor2tensor/layers/common_attention.py | 67 +++++++++---------------
 tensor2tensor/models/transformer.py      |  2 +-
 2 files changed, 26 insertions(+), 43 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index e1b926c9e..982cb0a72 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1547,15 +1547,19 @@ def dot_product_attention(q,
     return tf.matmul(weights, v)
 
 
-def _generate_relative_positions_matrix(length, max_relative_position,
+def _generate_relative_positions_matrix(length_q, length_k,
+                                        max_relative_position,
                                         cache=False):
   """Generates matrix of relative positions between inputs."""
   if not cache:
-    range_vec = tf.range(length)
-    range_mat = tf.reshape(tf.tile(range_vec, [length]), [length, length])
-    distance_mat = range_mat - tf.transpose(range_mat)
+    if length_q == length_k:
+      range_vec_q = range_vec_k = tf.range(length_q)
+    else:
+      range_vec_k = tf.range(length_k)
+      range_vec_q = range_vec_k[-length_q:]
+    distance_mat = range_vec_k[None, :] - range_vec_q[:, None]
   else:
-    distance_mat = tf.expand_dims(tf.range(-length+1, 1, 1), 0)
+    distance_mat = tf.expand_dims(tf.range(-length_k+1, 1, 1), 0)
   distance_mat_clipped = tf.clip_by_value(distance_mat, -max_relative_position,
                                           max_relative_position)
   # Shift values to be >= 0. Each integer still uniquely identifies a relative
@@ -1564,13 +1568,13 @@ def _generate_relative_positions_matrix(length, max_relative_position,
   return final_mat
 
 
-def _generate_relative_positions_embeddings(length, depth,
+def _generate_relative_positions_embeddings(length_q, length_k, depth,
                                             max_relative_position, name,
                                             cache=False):
-  """Generates tensor of size [1 if cache else length, length, depth]."""
+  """Generates tensor of size [1 if cache else length_q, length_k, depth]."""
   with tf.variable_scope(name):
     relative_positions_matrix = _generate_relative_positions_matrix(
-        length, max_relative_position, cache=cache)
+        length_q, length_k, max_relative_position, cache=cache)
     vocab_size = max_relative_position * 2 + 1
     # Generates embedding for each relative position of dimension depth.
     embeddings_table = tf.get_variable("embeddings", [vocab_size, depth])
@@ -1623,6 +1627,7 @@ def dot_product_attention_relative(q,
                                    name=None,
                                    make_image_summary=True,
                                    cache=False,
+                                   allow_memory=False,
                                    hard_attention_k=0):
   """Calculate relative position-aware dot-product self-attention.
 
@@ -1644,6 +1649,9 @@ def dot_product_attention_relative(q,
     name: an optional string.
     make_image_summary: Whether to make an attention image summary.
     cache: whether use cache mode
+    allow_memory: whether to assume that recurrent memory is in use. If True,
+      the length dimension of k/v/bias may be longer than the queries, and it is
+      assumed that the extra memory entries precede the non-memory entries.
     hard_attention_k: integer, if > 0 triggers hard attention (picking top-k)
 
   Returns:
@@ -1660,20 +1668,21 @@ def dot_product_attention_relative(q,
       values=[q, k, v]) as scope:
 
     # This calculation only works for self attention.
-    # q, k and v must therefore have the same shape.
-    if not cache:
+    # q, k and v must therefore have the same shape, unless memory is enabled.
+    if not cache and not allow_memory:
       q.get_shape().assert_is_compatible_with(k.get_shape())
       q.get_shape().assert_is_compatible_with(v.get_shape())
 
     # Use separate embeddings suitable for keys and values.
     depth = k.get_shape().as_list()[3]
-    length = common_layers.shape_list(k)[2]
+    length_k = common_layers.shape_list(k)[2]
+    length_q = common_layers.shape_list(q)[2] if allow_memory else length_k
     relations_keys = _generate_relative_positions_embeddings(
-        length, depth, max_relative_position, "relative_positions_keys",
-        cache=cache)
+        length_q, length_k, depth, max_relative_position,
+        "relative_positions_keys", cache=cache)
     relations_values = _generate_relative_positions_embeddings(
-        length, depth, max_relative_position, "relative_positions_values",
-        cache=cache)
+        length_q, length_k, depth, max_relative_position,
+        "relative_positions_values", cache=cache)
 
     # Compute self attention considering the relative position embeddings.
     logits = _relative_attention_inner(q, k, relations_keys, True)
@@ -1691,20 +1700,6 @@ def dot_product_attention_relative(q,
     return _relative_attention_inner(weights, v, relations_values, False)
 
 
-def dot_product_attention_relative_memory(q, k, v, bias, *args, **kwargs):
-  """Wrapper of dot_product_attention_relative to use with recurrent memory."""
-
-  q_len = tf.shape(q)[2]
-  k_len = tf.shape(k)[2]
-  num_memory_items = k_len - q_len
-
-  q = tf.pad(q, [[0, 0], [0, 0], [num_memory_items, 0], [0, 0]])
-  bias = tf.pad(bias, [[0, 0], [0, 0], [num_memory_items, 0], [0, 0]])
-  output = dot_product_attention_relative(q, k, v, bias, *args, **kwargs)
-
-  return output[:, :, num_memory_items:, :]
-
-
 def _relative_position_to_absolute_position_masked(x):
   """Helper to dot_product_self_attention_relative_v2.
 
@@ -4194,19 +4189,7 @@ def multihead_attention(query_antecedent,
           save_weights_to=save_weights_to,
           make_image_summary=make_image_summary,
           cache=cache is not None,
-          hard_attention_k=hard_attention_k)
-    elif attention_type == "dot_product_relative_memory":
-      x = dot_product_attention_relative_memory(
-          q,
-          k,
-          v,
-          bias,
-          max_relative_position,
-          dropout_rate,
-          image_shapes,
-          save_weights_to=save_weights_to,
-          make_image_summary=make_image_summary,
-          cache=cache is not None,
+          allow_memory=recurrent_memory is not None,
           hard_attention_k=hard_attention_k)
     elif attention_type == "dot_product_unmasked_relative_v2":
       x = dot_product_unmasked_self_attention_relative_v2(
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 536352ff4..e3d5ed624 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -2660,7 +2660,7 @@ def transformer_wikitext103_l4k_memory_v0():
       hparams.max_length / hparams.split_targets_chunk_length))  # 262144
 
   hparams.pos = None
-  hparams.self_attention_type = "dot_product_relative_memory"
+  hparams.self_attention_type = "dot_product_relative"
   hparams.max_relative_position = 2 * hparams.split_targets_chunk_length
 
   hparams.add_hparam("unconditional", True)

From 0fc87aae0d9a8ac89d8e051d5d8963f660100213 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 2 Apr 2019 15:31:34 -0700
Subject: [PATCH 1857/2720] Minor documentation.

PiperOrigin-RevId: 241617014
---
 tensor2tensor/models/research/rl.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 7e0ac4455..da9fcedef 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -634,6 +634,10 @@ class FeedForwardCnnSmallCategoricalPolicy(PolicyBase):
 
   def body(self, features):
     observations = features["inputs_raw"]
+    # Axis 0    - Batch.
+    # Axis 1    - Input Frames, 4 frames.
+    # Axis 2, 3 - Height & Width.
+    # Axis 4    - Channels RGB, 3 colours.
     x = tf.transpose(observations, [0, 2, 3, 1, 4])
     x_shape = common_layers.shape_list(x)
     x = tf.reshape(x, x_shape[:-2] + [-1])

From b6a9bbbd7c04e69ccfbf8f8d9c4b5b8947729bea Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Tue, 2 Apr 2019 17:26:54 -0700
Subject: [PATCH 1858/2720] Add VideoFlow paper to T2T Readme

PiperOrigin-RevId: 241637251
---
 README.md           | 1 +
 docs/walkthrough.md | 1 +
 2 files changed, 2 insertions(+)

diff --git a/README.md b/README.md
index a104f1bdc..27c992491 100644
--- a/README.md
+++ b/README.md
@@ -474,5 +474,6 @@ T2T](https://research.googleblog.com/2017/06/accelerating-deep-learning-research
 * [Attending to Mathematical Language with Transformers](https://arxiv.org/abs/1812.02825)
 * [The Evolved Transformer](https://arxiv.org/abs/1901.11117)
 * [Model-Based Reinforcement Learning for Atari](https://arxiv.org/abs/1903.00374)
+* [VideoFlow: A Flow-Based Generative Model for Video](https://arxiv.org/abs/1903.01434)
 
 *NOTE: This is not an official Google product.*
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index a104f1bdc..27c992491 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -474,5 +474,6 @@ T2T](https://research.googleblog.com/2017/06/accelerating-deep-learning-research
 * [Attending to Mathematical Language with Transformers](https://arxiv.org/abs/1812.02825)
 * [The Evolved Transformer](https://arxiv.org/abs/1901.11117)
 * [Model-Based Reinforcement Learning for Atari](https://arxiv.org/abs/1903.00374)
+* [VideoFlow: A Flow-Based Generative Model for Video](https://arxiv.org/abs/1903.01434)
 
 *NOTE: This is not an official Google product.*

From e79589f197ec9de2de829c9d4926fec6a4b34f3c Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 2 Apr 2019 19:14:52 -0700
Subject: [PATCH 1859/2720] Correct rng passing in multi-device mode. ResNet
 trains on a TPU donut now.

PiperOrigin-RevId: 241649966
---
 tensor2tensor/trax/configs/transformer_lm1b_8gb.gin |  2 +-
 tensor2tensor/trax/trax.py                          | 11 +++++++----
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index cded73df9..c0e9839b6 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -42,7 +42,7 @@ train_and_eval_batches.input_name = 'targets'
 
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.dropout = 0.2
+TransformerLM.dropout = 0.1
 TransformerLM.feature_depth = 512
 TransformerLM.feedforward_depth = 2048
 TransformerLM.max_len = 512
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 75e3cbab5..afe3b0478 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -253,15 +253,18 @@ def single_update(i, opt_state, batch, rng):
     return single_update
 
   @functools.partial(jax.pmap, axis_name="batch")
-  def mapped_update(i, opt_state, batch):
+  def mapped_update(i, opt_state, batch, rng):
+    """This is a multi-device version of the update function above."""
+    # We assume all tensors have the first dimension = num_devices.
     _, opt_update = optimizer(lr_fun)
     params = jax_opt.get_params(opt_state)
-    grads = jax.grad(loss_fun)(params, batch, predict_fun)
+    grads = jax.grad(loss_fun)(params, batch, predict_fun, rng)
     grads = jax.tree_util.tree_map(lambda g: jax.lax.psum(g, "batch"), grads)
     return opt_update(i, grads, opt_state)
 
-  def update(i, opt_state, batch):
-    return mapped_update(jax.replicate(i), opt_state, batch)
+  def update(i, opt_state, batch, rng):
+    # TODO(lukaszkaiser): investigate how to replicate rng and correct.
+    return mapped_update(jax.replicate(i), opt_state, batch, jax.replicate(rng))
 
   return update
 

From c2a89e8fef64d49710c1c0083da4a8e7ac0ab9ab Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 3 Apr 2019 13:19:25 -0700
Subject: [PATCH 1860/2720] Allow skip_eos_postprocess to happen when decoding
 as well.

PiperOrigin-RevId: 241792692
---
 tensor2tensor/utils/decoding.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index b54390cf6..511787942 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -110,7 +110,8 @@ def log_decode_results(inputs,
                        save_images=False,
                        output_dir=None,
                        identity_output=False,
-                       log_results=True):
+                       log_results=True,
+                       skip_eos_postprocess=False):
   """Log inference results."""
 
   # TODO(lukaszkaiser) refactor this into feature_encoder
@@ -132,7 +133,7 @@ def fix_and_save_video(vid, prefix):
   is_image = "image" in problem_name
   is_text2class = isinstance(registry.problem(problem_name),
                              text_problems.Text2ClassProblem)
-  skip_eos_postprocess = is_image or is_text2class
+  skip_eos_postprocess = is_image or is_text2class or skip_eos_postprocess
 
   decoded_inputs = None
   if is_image and save_images:
@@ -356,7 +357,8 @@ def decode_once(estimator,
           output_dir=output_dir,
           identity_output=decode_hp.identity_output,
           targets=targets,
-          log_results=log_results)
+          log_results=log_results,
+          skip_eos_postprocess=decode_hp.skip_eos_postprocess)
       decoded_outputs.append(decoded)
 
     # Write out predictions if decode_to_file passed
@@ -488,7 +490,8 @@ def timer(gen):
             None,
             inputs_vocab,
             targets_vocab,
-            log_results=decode_hp.log_results)
+            log_results=decode_hp.log_results,
+            skip_eos_postprocess=decode_hp.skip_eos_postprocess)
         beam_decodes.append(decoded_outputs)
         if decode_hp.write_beam_scores:
           beam_scores.append(score)
@@ -507,7 +510,8 @@ def timer(gen):
           None,
           inputs_vocab,
           targets_vocab,
-          log_results=decode_hp.log_results)
+          log_results=decode_hp.log_results,
+          skip_eos_postprocess=decode_hp.skip_eos_postprocess)
       decodes.append(decoded_outputs)
     total_time_per_step += elapsed_time
     total_cnt += result["outputs"].shape[-1]

From 4de825477a3708d80842f649bc07cb8c62eb8ee5 Mon Sep 17 00:00:00 2001
From: cbockman <c.bockman@gmail.com>
Date: Wed, 3 Apr 2019 14:18:17 -0700
Subject: [PATCH 1861/2720] fix get_standardized_layers spelling (#1529)

---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 982cb0a72..3d3322635 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -299,7 +299,7 @@ def memeff_attention_fn(*args, **kwargs):
 
 
 def add_standard_attention_hparams(hparams):
-  """Adds the hparams used by get_standadized_layers."""
+  """Adds the hparams used by get_standardized_layers."""
   # All hyperparameters ending in "dropout" are automatically set to 0.0
   # when not in training mode.
 

From 64c4e2ce5369c67ccefd6bf54233f0ebe2747c77 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 3 Apr 2019 17:03:57 -0700
Subject: [PATCH 1862/2720] Remove internal implementation's reliance on
 Bijectors.

This cleans up dependencies, showing how TFP Bijectors aren't necessary in the world of Keras layers. Future CLs may:

+ Move TransformedRandomVariable upstream to Edward2.
+ Refactor Edward2 to follow Edward1's mix-in approach to wrap Distributions. This lets us implement TransformedRandomVariable without having to implement a new TransformedDistribution.

PiperOrigin-RevId: 241835590
---
 tensor2tensor/layers/reversible_layers.py     | 108 ++++++++++++++++--
 .../layers/reversible_layers_test.py          |  26 +++++
 2 files changed, 127 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/layers/reversible_layers.py b/tensor2tensor/layers/reversible_layers.py
index 4c7003c15..4082d49e6 100644
--- a/tensor2tensor/layers/reversible_layers.py
+++ b/tensor2tensor/layers/reversible_layers.py
@@ -83,13 +83,7 @@ def __call__(self, inputs, *args, **kwargs):
 
     if not isinstance(inputs, ed.RandomVariable):
       return super(ActNorm, self).__call__(inputs, *args, **kwargs)
-
-    bijector = tfp.bijectors.Inline(
-        forward_fn=self.__call__,
-        inverse_fn=self.reverse,
-        inverse_log_det_jacobian_fn=lambda y: -self.log_det_jacobian(y),
-        forward_min_event_ndims=0)
-    return ed.TransformedDistribution(inputs.distribution, bijector=bijector)
+    return TransformedRandomVariable(inputs, self)
 
   def call(self, inputs):
     return (inputs + self.bias) * tf.exp(self.log_scale)
@@ -362,3 +356,103 @@ def sinkhorn(inputs, n_iters=20):
                             [-1, 1, vocab_size])
   outputs = tf.exp(log_alpha)
   return outputs
+
+
+class TransformedDistribution(tfp.distributions.Distribution):
+  """Distribution of f(x), where x ~ p(x) and f is reversible."""
+
+  def __init__(self, base, reversible_layer, name=None):
+    """Constructs a transformed distribution.
+
+    Args:
+      base: Base distribution.
+      reversible_layer: Callable with methods `reverse` and `log_det_jacobian`.
+      name: Name for scoping operations in the class.
+    """
+    self.base = base
+    self.reversible_layer = reversible_layer
+    if name is None:
+      name = reversible_layer.name + base.name
+    super(TransformedDistribution, self).__init__(
+        base.dtype,
+        base.reparameterization_type,
+        base.validate_args,
+        base.allow_nan_stats,
+        parameters=dict(locals()),
+        name=name)
+
+  def _event_shape_tensor(self):
+    return self.base.event_shape_tensor()
+
+  def _event_shape(self):
+    return self.base.event_shape
+
+  def _batch_shape_tensor(self):
+    return self.base.batch_shape_tensor()
+
+  def _batch_shape(self):
+    return self.base.batch_shape
+
+  def __getitem__(self, slices):
+    overrides = {'base': self.base[slices]}
+    return self.copy(**overrides)
+
+  def _call_sample_n(self, sample_shape, seed, name, **kwargs):
+    x = self.base.sample(sample_shape, seed, **kwargs)
+    y = self.reversible_layer(x)
+    return y
+
+  def _log_prob(self, value):
+    x = self.reversible_layer.reverse(value)
+    log_det_jacobian = self.reversible_layer.log_det_jacobian(value)
+    return self.base.log_prob(x) + log_det_jacobian
+
+  def _prob(self, value):
+    if not hasattr(self.base, '_prob'):
+      return tf.exp(self.log_prob(value))
+    x = self.reversible_layer.reverse(value)
+    log_det_jacobian = self.reversible_layer.log_det_jacobian(value)
+    return self.base.prob(x) * tf.exp(log_det_jacobian)
+
+  def _log_cdf(self, value):
+    x = self.reversible_layer.reverse(value)
+    return self.base.log_cdf(x)
+
+  def _cdf(self, value):
+    x = self.reversible_layer.reverse(value)
+    return self.base.cdf(x)
+
+  def _log_survival_function(self, value):
+    x = self.reversible_layer.reverse(value)
+    return self.base.log_survival_function(x)
+
+  def _survival_function(self, value):
+    x = self.reversible_layer.reverse(value)
+    return self.base.survival_function(x)
+
+  def _quantile(self, value):
+    inverse_cdf = self.base.quantile(value)
+    return self.reversible_layer(inverse_cdf)
+
+  def _entropy(self):
+    dummy = tf.zeros(
+        tf.concat([self.batch_shape_tensor(), self.event_shape_tensor()], 0),
+        dtype=self.dtype)
+    log_det_jacobian = self.reversible_layer.log_det_jacobian(dummy)
+    entropy = self.base.entropy() - log_det_jacobian
+    return entropy
+
+
+@ed.interceptable
+def TransformedRandomVariable(random_variable,  # pylint: disable=invalid-name
+                              reversible_layer,
+                              name=None,
+                              sample_shape=(),
+                              value=None):
+  """Random variable for f(x), where x ~ p(x) and f is reversible."""
+  return ed.RandomVariable(
+      distribution=TransformedDistribution(random_variable.distribution,
+                                           reversible_layer,
+                                           name=name),
+      sample_shape=sample_shape,
+      value=value)
diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
index bc0e73870..aaf7c6970 100644
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -127,6 +127,32 @@ def testMADENoHidden(self):
     self.assertAllEqual(outputs_val[:, 0, :], np.zeros((batch_size, units)))
     self.assertEqual(outputs_val.shape, (batch_size, length, units))
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testTransformedRandomVariable(self):
+    class Exp(tf.keras.layers.Layer):
+      """Exponential activation function for reversible networks."""
+
+      def __call__(self, inputs, *args, **kwargs):
+        if not isinstance(inputs, ed.RandomVariable):
+          return super(Exp, self).__call__(inputs, *args, **kwargs)
+        return reversible.TransformedRandomVariable(inputs, self)
+
+      def call(self, inputs):
+        return tf.exp(inputs)
+
+      def reverse(self, inputs):
+        return tf.log(inputs)
+
+      def log_det_jacobian(self, inputs):
+        return -tf.log(inputs)
+
+    x = ed.Normal(0., 1.)
+    y = Exp()(x)
+    y_sample = self.evaluate(y.distribution.sample())
+    y_log_prob = self.evaluate(y.distribution.log_prob(y_sample))
+    self.assertGreater(y_sample, 0.)
+    self.assertTrue(np.isfinite(y_log_prob))
+
 
 if __name__ == '__main__':
   tf.test.main()

From 614794aeae591d3eaf9215bc722ca7f470b6dfaa Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 4 Apr 2019 08:30:35 -0700
Subject: [PATCH 1863/2720] internal change

PiperOrigin-RevId: 241934922
---
 tensor2tensor/trax/trax_test.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index 8eadee4fe..a73b90169 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -23,6 +23,8 @@
 import functools
 import tempfile
 
+from jax import test_util  # pylint: disable=unused-import
+from jax.config import config
 import numpy as np
 
 from tensor2tensor.trax import inputs as inputs_lib
@@ -91,4 +93,5 @@ def test_train_eval_predict(self):
 
 
 if __name__ == "__main__":
+  config.config_with_absl()
   test.main()

From 2182ee5cb1657a746d8a160e2ddeb1f29792d85d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 4 Apr 2019 14:06:56 -0700
Subject: [PATCH 1864/2720] Add __init__.py to t2t/keras to make Travis happy.

PiperOrigin-RevId: 242004246
---
 tensor2tensor/keras/__init__.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)
 create mode 100644 tensor2tensor/keras/__init__.py

diff --git a/tensor2tensor/keras/__init__.py b/tensor2tensor/keras/__init__.py
new file mode 100644
index 000000000..b775a72bd
--- /dev/null
+++ b/tensor2tensor/keras/__init__.py
@@ -0,0 +1,16 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+

From 9dc3d1274ce8cb25513adb071262cadb4ba7e5d3 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 4 Apr 2019 14:27:28 -0700
Subject: [PATCH 1865/2720] Add recent DeepMind math dataset to T2T.

PiperOrigin-RevId: 242008722
---
 .../algorithmic_math_deepmind.py              | 104 ++++++++++++++++++
 tensor2tensor/data_generators/all_problems.py |   1 +
 2 files changed, 105 insertions(+)
 create mode 100644 tensor2tensor/data_generators/algorithmic_math_deepmind.py

diff --git a/tensor2tensor/data_generators/algorithmic_math_deepmind.py b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
new file mode 100644
index 000000000..18841b2f0
--- /dev/null
+++ b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
@@ -0,0 +1,104 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+r"""Data generators for the DeepMind Mathematics Dataset.
+
+See https://github.com/deepmind/mathematics_dataset for the original repository.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import tarfile
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+_URL = "https://storage.googleapis.com/mathematics-dataset/v1.0.tar.gz"
+
+
+@registry.register_problem
+class AlgorithmicMathDeepmindAll(text_problems.Text2TextProblem):
+  """DeepMind Mathematics Problem, v1.0, all data."""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 128,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    """Downloads and extracts the dataset and generates examples.
+
+    Args:
+      data_dir: The base directory where data and vocab files are stored.
+      tmp_dir: temp directory to download and extract the dataset.
+      dataset_split: split of the data-set.
+
+    Yields:
+      The data examples.
+    """
+    # Create directories if needed.
+    if not tf.gfile.Exists(tmp_dir):
+      tf.gfile.MakeDirs(tmp_dir)
+    if not tf.gfile.Exists(data_dir):
+      tf.gfile.MakeDirs(data_dir)
+
+    # Download and extract the data.
+    filename = os.path.basename(_URL)
+    path = generator_utils.maybe_download(tmp_dir, filename, _URL)
+    tarfile.open(path, "r:gz").extractall(tmp_dir)
+
+    # Create the list of directories with data files.
+    train_dirs = ["v1.0/train-easy", "v1.0/train-medium", "v1.0/train-hard"]
+    eval_dirs = ["v1.0/interpolate", "v1.0/extrapolate"]
+    dirs = eval_dirs
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      dirs = train_dirs
+    dirs = [os.path.join(tmp_dir, d) for d in dirs]
+
+    # Iterate over directories and files generating examples.
+    for d in dirs:
+      files = tf.gfile.Glob(d + "/*.txt")
+      for fname in files:
+        # In each text file, the first line is the input, the next the answer,
+        # and so on until the end of the file.
+        cur_input = None
+        with tf.gfile.Open(fname, "rb") as f:
+          for line in f:
+            if cur_input is None:
+              cur_input = line.strip()
+            else:
+              yield {"inputs": cur_input, "targets": line.strip()}
+              cur_input = None
diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 006a881be..f636867c9 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -24,6 +24,7 @@
 MODULES = [
     "tensor2tensor.data_generators.algorithmic",
     "tensor2tensor.data_generators.algorithmic_math",
+    "tensor2tensor.data_generators.algorithmic_math_deepmind",
     "tensor2tensor.data_generators.algorithmic_math_two_variables",
     "tensor2tensor.data_generators.allen_brain",
     "tensor2tensor.data_generators.audio",

From f6c024be1b1129d9d83fcd96828d0aea3b3d67ac Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 4 Apr 2019 15:48:35 -0700
Subject: [PATCH 1866/2720] Enable frame resizing for rendered gym
 environments.

PiperOrigin-RevId: 242024877
---
 tensor2tensor/envs/mujoco_problems.py |  1 +
 tensor2tensor/rl/gym_utils.py         | 38 +++++++++++++++++++++------
 tensor2tensor/rl/gym_utils_test.py    | 19 +++++++++-----
 3 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/envs/mujoco_problems.py b/tensor2tensor/envs/mujoco_problems.py
index 85adf7695..927321d1d 100644
--- a/tensor2tensor/envs/mujoco_problems.py
+++ b/tensor2tensor/envs/mujoco_problems.py
@@ -38,6 +38,7 @@ def __init__(self):
             "rl_env_max_episode_steps": -1,
             "maxskip_env": False,
             "rendered_env": True,
+            "rendered_env_resize_to": None,  # Do not resize frames
             "sticky_actions": False
         })
     super(ReacherEnvProblem, self).__init__(
diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index e1b6032b9..a7c339b12 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -21,6 +21,7 @@
 
 import gym
 import numpy as np
+from PIL import Image
 import tensorflow as tf
 
 
@@ -82,18 +83,36 @@ def reset(self, **kwargs):
 class RenderedEnv(gym.Wrapper):
   """Simple Env wrapper to override observations with rendered rgb values."""
 
-  def __init__(self, env, mode="rgb_array", low=0, high=255):
+  def __init__(self, env, mode="rgb_array", low=0, high=255, resize_to=None):
     gym.Wrapper.__init__(self, env)
     # Get a sample frame to correctly set observation space
     self.mode = mode
     sample_frame = self.render(mode=self.mode)
     assert sample_frame is not None
-    self.observation_space = gym.spaces.Box(
-        low=low, high=high, shape=sample_frame.shape, dtype=sample_frame.dtype)
+    self.should_resize = False
+    if resize_to is None:
+      self.observation_space = gym.spaces.Box(
+          low=low,
+          high=high,
+          shape=sample_frame.shape,
+          dtype=sample_frame.dtype)
+    else:
+      assert len(resize_to) == 2
+      self.should_resize = True
+      self.observation_space = gym.spaces.Box(
+          low=low,
+          high=high,
+          shape=list(resize_to) + list(sample_frame.shape[-1:]),
+          dtype=sample_frame.dtype)
 
   def step(self, action):
     _, reward, done, info = self.env.step(action)
     obs = self.env.render(mode=self.mode)
+    if self.should_resize:
+      img = Image.fromarray(obs)
+      img = img.resize(
+          self.observation_space.shape[:-1], resample=Image.ANTIALIAS)
+      obs = np.array(img)
     return obs, reward, done, info
 
   def reset(self, **kwargs):
@@ -124,7 +143,7 @@ def remove_time_limit_wrapper(env):
 
 
 def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env,
-                    sticky_actions):
+                    rendered_env_resize_to, sticky_actions):
   """Wraps a gym environment. see make_gym_env for details."""
   # rl_env_max_episode_steps is None or int.
   assert ((not rl_env_max_episode_steps) or
@@ -143,11 +162,11 @@ def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env,
     env = MaxAndSkipEnv(env)  # pylint: disable=redefined-variable-type
 
   if rendered_env:
-    env = RenderedEnv(env)
+    env = RenderedEnv(env, resize_to=rendered_env_resize_to)
 
   if wrap_with_time_limit:
-    env = gym.wrappers.TimeLimit(env,
-                                 max_episode_steps=rl_env_max_episode_steps)
+    env = gym.wrappers.TimeLimit(
+        env, max_episode_steps=rl_env_max_episode_steps)
   return env
 
 
@@ -155,6 +174,7 @@ def make_gym_env(name,
                  rl_env_max_episode_steps=-1,
                  maxskip_env=False,
                  rendered_env=False,
+                 rendered_env_resize_to=None,
                  sticky_actions=False):
   """Create a gym env optionally with a time limit and maxskip wrapper.
 
@@ -168,6 +188,8 @@ def make_gym_env(name,
     maxskip_env: whether to also use MaxAndSkip wrapper before time limit.
     rendered_env: whether to force render for observations. Use this for
       environments that are not natively rendering the scene for observations.
+    rendered_env_resize_to: a list of [height, width] to change the original
+      resolution of the native environment render.
     sticky_actions: whether to use sticky_actions before MaxAndSkip wrapper.
 
   Returns:
@@ -175,7 +197,7 @@ def make_gym_env(name,
   """
   env = gym.make(name)
   return gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env,
-                         rendered_env, sticky_actions)
+                         rendered_env, rendered_env_resize_to, sticky_actions)
 
 
 def register_gym_env(class_entry_point, version="v0", kwargs=None):
diff --git a/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index 95d92faf7..267f846b0 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -43,6 +43,10 @@ def step(self, action):
     else:
       return self.observation_space.high, +1.0, True, {}
 
+  def render(self, mode="human"):
+    del mode  # Unused
+    return np.zeros([640, 480, 3], np.uint8)
+
 
 class EnvWithOptions(SimpleEnv):
   """A simple env that takes arguments on init."""
@@ -79,6 +83,11 @@ def test_unlimited_env(self):
     self.assertTrue(isinstance(env, gym.wrappers.TimeLimit))
     self.assertTrue(env._max_episode_steps is None)
 
+  def test_rendered_env(self):
+    env = gym_utils.RenderedEnv(SimpleEnv(), resize_to=(64, 64))
+    obs, _, _, _ = env.step(1)
+    self.assertTrue(np.allclose(np.zeros([64, 64, 3], np.uint8), obs))
+
   def test_gym_registration(self):
     reg_id, env = gym_utils.register_gym_env(
         "tensor2tensor.rl.gym_utils_test:SimpleEnv")
@@ -89,8 +98,8 @@ def test_gym_registration(self):
     self.assertTrue(isinstance(env, gym.Env))
 
     # Just make sure we got the same environment.
-    self.assertTrue(np.allclose(env.reset(),
-                                np.zeros(shape=(3, 3), dtype=np.uint8)))
+    self.assertTrue(
+        np.allclose(env.reset(), np.zeros(shape=(3, 3), dtype=np.uint8)))
 
     _, _, done, _ = env.step(1)
     self.assertTrue(done)
@@ -98,8 +107,7 @@ def test_gym_registration(self):
   def test_gym_registration_with_kwargs(self):
     reg_id, env = gym_utils.register_gym_env(
         "tensor2tensor.rl.gym_utils_test:EnvWithOptions",
-        kwargs={"done_action": 2}
-    )
+        kwargs={"done_action": 2})
 
     self.assertEqual("T2TEnv-EnvWithOptions-v0", reg_id)
 
@@ -120,8 +128,7 @@ def test_gym_registration_with_kwargs(self):
     reg_id, env = gym_utils.register_gym_env(
         "tensor2tensor.rl.gym_utils_test:EnvWithOptions",
         version="v1",
-        kwargs={"done_action": 1}
-    )
+        kwargs={"done_action": 1})
 
     self.assertEqual("T2TEnv-EnvWithOptions-v1", reg_id)
 

From 7e06bd6e05311815fe94583c890fab4067a7f341 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 4 Apr 2019 17:45:22 -0700
Subject: [PATCH 1867/2720] Allow hard attention in Universal Transformer.

PiperOrigin-RevId: 242044506
---
 .../models/research/universal_transformer.py        | 13 +++++++++++++
 .../models/research/universal_transformer_util.py   |  9 ++++++---
 tensor2tensor/trax/configs/transformer_lm1b_8gb.gin |  2 +-
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 9f13e4125..ab93d54b6 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -555,6 +555,19 @@ def adaptive_universal_transformer_multilayer_tpu():
   return hparams
 
 
+@registry.register_hparams
+def adaptive_universal_transformer_multilayer_hard():
+  """Multi-layer config for adaptive Transformer with hard attention."""
+  hparams = adaptive_universal_transformer_multilayer_tpu()
+  hparams.batch_size = 256
+  hparams.hard_attention_k = 8
+  hparams.add_step_timing_signal = True
+  # hparams.add_sru = True  # This is very slow on GPUs, does it help?
+  hparams.self_attention_type = "dot_product_relative_v2"
+  hparams.max_relative_position = 256
+  return hparams
+
+
 @registry.register_hparams
 def adaptive_universal_transformer_small():
   hparams = universal_transformer_small()
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 2937e7dea..062e397a4 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -440,7 +440,8 @@ def transformer_encoder_attention_unit(x,
         save_weights_to=save_weights_to,
         max_relative_position=hparams.max_relative_position,
         make_image_summary=make_image_summary,
-        dropout_broadcast_dims=attention_dropout_broadcast_dims)
+        dropout_broadcast_dims=attention_dropout_broadcast_dims,
+        hard_attention_k=hparams.hard_attention_k)
     x = common_layers.layer_postprocess(x, y, hparams)
   return x
 
@@ -533,7 +534,8 @@ def transformer_decoder_attention_unit(x,
         max_relative_position=hparams.max_relative_position,
         cache=None,
         make_image_summary=make_image_summary,
-        dropout_broadcast_dims=attention_dropout_broadcast_dims)
+        dropout_broadcast_dims=attention_dropout_broadcast_dims,
+        hard_attention_k=hparams.hard_attention_k)
     x = common_layers.layer_postprocess(x, y, hparams)
   if encoder_output is not None:
     with tf.variable_scope("encdec_attention"):
@@ -548,7 +550,8 @@ def transformer_decoder_attention_unit(x,
           hparams.attention_dropout,
           save_weights_to=save_weights_to,
           make_image_summary=make_image_summary,
-          dropout_broadcast_dims=attention_dropout_broadcast_dims)
+          dropout_broadcast_dims=attention_dropout_broadcast_dims,
+          hard_attention_k=hparams.hard_attention_k)
       x = common_layers.layer_postprocess(x, y, hparams)
   return x
 
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index c0e9839b6..0d07d1083 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -34,7 +34,7 @@ train.eval_steps = 5
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
 train.run_debug_step = False
-train.train_steps = 50000
+train.train_steps = 500000
 
 # Parameters for train_and_eval_batches:
 # ==============================================================================

From 9c557c2eaa10d454f11abb2ab3c8d2dd2a6582af Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 4 Apr 2019 19:18:14 -0700
Subject: [PATCH 1868/2720] Fork stax to allow more experimentation.

PiperOrigin-RevId: 242055396
---
 tensor2tensor/trax/stax/__init__.py  |   2 +-
 tensor2tensor/trax/stax/attention.py |   3 +-
 tensor2tensor/trax/stax/slax.py      |   3 +-
 tensor2tensor/trax/stax/stax_base.py | 331 +++++++++++++++++++++++++++
 4 files changed, 336 insertions(+), 3 deletions(-)
 create mode 100644 tensor2tensor/trax/stax/stax_base.py

diff --git a/tensor2tensor/trax/stax/__init__.py b/tensor2tensor/trax/stax/__init__.py
index aca079496..eb0c51790 100644
--- a/tensor2tensor/trax/stax/__init__.py
+++ b/tensor2tensor/trax/stax/__init__.py
@@ -21,7 +21,7 @@
 # We create a flat stax.* namespace for uniform calling conventions as we
 # upstream changes.
 # pylint: disable=wildcard-import
-from jax.experimental.stax import *
 from tensor2tensor.trax.stax.attention import *
 from tensor2tensor.trax.stax.losses import *
 from tensor2tensor.trax.stax.slax import *
+from tensor2tensor.trax.stax.stax_base import *
diff --git a/tensor2tensor/trax/stax/attention.py b/tensor2tensor/trax/stax/attention.py
index afa444338..b771c1076 100644
--- a/tensor2tensor/trax/stax/attention.py
+++ b/tensor2tensor/trax/stax/attention.py
@@ -19,11 +19,12 @@
 from __future__ import print_function
 
 from jax import random
-import jax.experimental.stax as stax
 import jax.numpy as np
 import numpy as onp
 import numpy.random as npr
 
+from tensor2tensor.trax.stax import stax_base as stax
+
 
 def causal_mask(size, dtype=np.uint8):
   """Causal attention mask."""
diff --git a/tensor2tensor/trax/stax/slax.py b/tensor2tensor/trax/stax/slax.py
index a99631241..da4ed7559 100644
--- a/tensor2tensor/trax/stax/slax.py
+++ b/tensor2tensor/trax/stax/slax.py
@@ -20,10 +20,11 @@
 
 import inspect
 from absl import logging
-import jax.experimental.stax as stax
 import jax.numpy as np
 from jax.tree_util import register_pytree_node as _register_pytree_node
 
+from tensor2tensor.trax.stax import stax_base as stax
+
 
 # Utility functions
 # ------------------------------------------------------------------------------
diff --git a/tensor2tensor/trax/stax/stax_base.py b/tensor2tensor/trax/stax/stax_base.py
new file mode 100644
index 000000000..f67a5bfaf
--- /dev/null
+++ b/tensor2tensor/trax/stax/stax_base.py
@@ -0,0 +1,331 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Stax is a small flexible neural net specification library from scratch."""
+
+# Forked from JAX for more experimentation on syntax and back-ends.
+# See JAX version at https://github.com/google/jax/tree/master/jax/experimental
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+import itertools
+import operator as op
+
+import numpy as onp
+import numpy.random as npr
+from six.moves import reduce
+
+from jax import lax
+from jax import random
+from jax.scipy.special import logsumexp
+import jax.numpy as np
+
+
+# Following the convention used in Keras and tf.layers, we use CamelCase for the
+# names of layer constructors, like Conv and Relu, while using snake_case for
+# other functions, like lax.conv and relu.
+
+
+def relu(x): return np.maximum(x, 0.)
+def softplus(x): return np.logaddexp(x, 0.)
+
+def logsoftmax(x, axis=-1):
+  """Apply log softmax to an array of logits, log-normalizing along an axis."""
+  return x - logsumexp(x, axis, keepdims=True)
+
+def softmax(x, axis=-1):
+  """Apply softmax to an array of logits, exponentiating and normalizing along an axis."""
+  unnormalized = np.exp(x - x.max(axis, keepdims=True))
+  return unnormalized / unnormalized.sum(axis, keepdims=True)
+
+def fastvar(x, axis, keepdims):
+  """A fast but less numerically-stable variance calculation than np.var."""
+  return np.mean(x**2, axis, keepdims=keepdims) - np.mean(x, axis, keepdims=keepdims)**2
+
+
+# Initializers
+
+def randn(stddev=1e-2, rng=npr):
+  """An initializer function for random normal coefficients."""
+  def init(shape):
+    return rng.normal(size=shape, scale=stddev).astype('float32')
+  return init
+
+def glorot(out_dim=0, in_dim=1, scale=onp.sqrt(2), rng=npr):
+  """An initializer function for random Glorot-scaled coefficients."""
+  def init(shape):
+    fan_in, fan_out = shape[in_dim], shape[out_dim]
+    size = onp.prod(onp.delete(shape, [in_dim, out_dim]))
+    std = scale / np.sqrt((fan_in + fan_out) / 2. * size)
+    return rng.normal(size=shape, scale=std).astype('float32')
+  return init
+
+zeros = functools.partial(np.zeros, dtype='float32')
+ones = functools.partial(np.ones, dtype='float32')
+
+
+# Layers
+
+# Each layer constructor function returns an (init_fun, apply_fun) pair, where
+#   init_fun: takes an input shape and returns an (output_shape, params) pair,
+#   apply_fun: takes params, inputs, and an rng key and applies the layer.
+
+
+def Dense(out_dim, W_init=glorot(), b_init=randn()):
+  """Layer constructor function for a dense (fully-connected) layer."""
+  def init_fun(input_shape):
+    output_shape = input_shape[:-1] + (out_dim,)
+    W, b = W_init((input_shape[-1], out_dim)), b_init((out_dim,))
+    return output_shape, (W, b)
+  def apply_fun(params, inputs, **kwargs):
+    W, b = params
+    return np.dot(inputs, W) + b
+  return init_fun, apply_fun
+
+
+def GeneralConv(dimension_numbers, out_chan, filter_shape,
+                strides=None, padding='VALID', W_init=None, b_init=randn(1e-6)):
+  """Layer construction function for a general convolution layer."""
+  lhs_spec, rhs_spec, out_spec = dimension_numbers
+  one = (1,) * len(filter_shape)
+  strides = strides or one
+  W_init = W_init or glorot(rhs_spec.index('O'), rhs_spec.index('I'))
+  def init_fun(input_shape):
+    filter_shape_iter = iter(filter_shape)
+    kernel_shape = [out_chan if c == 'O' else
+                    input_shape[lhs_spec.index('C')] if c == 'I' else
+                    next(filter_shape_iter) for c in rhs_spec]
+    output_shape = lax.conv_general_shape_tuple(
+        input_shape, kernel_shape, strides, padding, dimension_numbers)
+    bias_shape = [out_chan if c == 'C' else 1 for c in out_spec]
+    bias_shape = tuple(itertools.dropwhile(lambda x: x == 1, bias_shape))
+    W, b = W_init(kernel_shape), b_init(bias_shape)
+    return output_shape, (W, b)
+  def apply_fun(params, inputs, **kwargs):
+    W, b = params
+    return lax.conv_general_dilated(inputs, W, strides, padding, one, one,
+                                    dimension_numbers) + b
+  return init_fun, apply_fun
+Conv = functools.partial(GeneralConv, ('NHWC', 'HWIO', 'NHWC'))
+
+
+def BatchNorm(axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True,
+              beta_init=zeros, gamma_init=ones):
+  """Layer construction function for a batch normalization layer."""
+  _beta_init = lambda shape: beta_init(shape) if center else ()
+  _gamma_init = lambda shape: gamma_init(shape) if scale else ()
+  axis = (axis,) if np.isscalar(axis) else axis
+  def init_fun(input_shape):
+    shape = tuple(d for i, d in enumerate(input_shape) if i not in axis)
+    beta, gamma = _beta_init(shape), _gamma_init(shape)
+    return input_shape, (beta, gamma)
+  def apply_fun(params, x, **kwargs):
+    beta, gamma = params
+    # TODO(phawkins): np.expand_dims should accept an axis tuple.
+    # (https://github.com/numpy/numpy/issues/12290)
+    ed = tuple(None if i in axis else slice(None) for i in range(np.ndim(x)))
+    beta = beta[ed]
+    gamma = gamma[ed]
+    mean, var = np.mean(x, axis, keepdims=True), fastvar(x, axis, keepdims=True)
+    z = (x - mean) / np.sqrt(var + epsilon)
+    if center and scale: return gamma * z + beta
+    if center: return z + beta
+    if scale: return gamma * z
+    return z
+  return init_fun, apply_fun
+
+
+def _elemwise_no_params(fun, **fun_kwargs):
+  init_fun = lambda input_shape: (input_shape, ())
+  apply_fun = lambda params, inputs, **kwargs: fun(inputs, **fun_kwargs)
+  return init_fun, apply_fun
+Tanh = _elemwise_no_params(np.tanh)
+Relu = _elemwise_no_params(relu)
+Exp = _elemwise_no_params(np.exp)
+LogSoftmax = _elemwise_no_params(logsoftmax, axis=-1)
+Softmax = _elemwise_no_params(softmax, axis=-1)
+Softplus = _elemwise_no_params(softplus)
+
+
+def _pooling_layer(reducer, init_val, rescaler=None):
+  def PoolingLayer(window_shape, strides=None, padding='VALID'):
+    """Layer construction function for a pooling layer."""
+    strides = strides or (1,) * len(window_shape)
+    rescale = rescaler(window_shape, strides, padding) if rescaler else None
+    dims = (1,) + window_shape + (1,)  # NHWC
+    strides = (1,) + strides + (1,)
+    def init_fun(input_shape):
+      out_shape = lax.reduce_window_shape_tuple(input_shape, dims, strides, padding)
+      return out_shape, ()
+    def apply_fun(params, inputs, **kwargs):
+      out = lax.reduce_window(inputs, init_val, reducer, dims, strides, padding)
+      return rescale(out, inputs) if rescale else out
+    return init_fun, apply_fun
+  return PoolingLayer
+MaxPool = _pooling_layer(lax.max, -np.inf)
+SumPool = _pooling_layer(lax.add, 0.)
+
+
+def _normalize_by_window_size(dims, strides, padding):
+  def rescale(outputs, inputs):
+    one = np.ones(inputs.shape[1:-1], dtype=inputs.dtype)
+    window_sizes = lax.reduce_window(one, 0., lax.add, dims, strides, padding)
+    return outputs / window_sizes[..., np.newaxis]
+  return rescale
+AvgPool = _pooling_layer(lax.add, 0., _normalize_by_window_size)
+
+
+def Flatten():
+  """Layer construction function for flattening all but the leading dim."""
+  def init_fun(input_shape):
+    output_shape = input_shape[0], reduce(op.mul, input_shape[1:], 1)
+    return output_shape, ()
+  def apply_fun(params, inputs, **kwargs):
+    return np.reshape(inputs, (inputs.shape[0], -1))
+  return init_fun, apply_fun
+Flatten = Flatten()
+
+
+def Identity():
+  """Layer construction function for an identity layer."""
+  init_fun = lambda input_shape: (input_shape, ())
+  apply_fun = lambda params, inputs, **kwargs: inputs
+  return init_fun, apply_fun
+Identity = Identity()
+
+
+def FanOut(num):
+  """Layer construction function for a fan-out layer."""
+  init_fun = lambda input_shape: ([input_shape] * num, ())
+  apply_fun = lambda params, inputs, **kwargs: [inputs] * num
+  return init_fun, apply_fun
+
+
+def FanInSum():
+  """Layer construction function for a fan-in sum layer."""
+  init_fun = lambda input_shape: (input_shape[0], ())
+  apply_fun = lambda params, inputs, **kwargs: sum(inputs)
+  return init_fun, apply_fun
+FanInSum = FanInSum()
+
+
+def FanInConcat(axis=-1):
+  """Layer construction function for a fan-in concatenation layer."""
+  def init_fun(input_shape):
+    ax = axis % len(input_shape[0])
+    concat_size = sum(shape[ax] for shape in input_shape)
+    out_shape = input_shape[0][:ax] + (concat_size,) + input_shape[0][ax+1:]
+    return out_shape, ()
+  def apply_fun(params, inputs, **kwargs):
+    return np.concatenate(inputs, axis)
+  return init_fun, apply_fun
+
+
+def Dropout(rate, mode='train'):
+  """Layer construction function for a dropout layer with given rate."""
+  def init_fun(input_shape):
+    return input_shape, ()
+  def apply_fun(params, inputs, **kwargs):
+    rng = kwargs.get('rng', None)
+    if rng is None:
+      msg = ("Dropout layer requires apply_fun to be called with a PRNG key "
+             "argument. That is, instead of `apply_fun(params, inputs)`, call "
+             "it like `apply_fun(params, inputs, key)` where `key` is a "
+             "jax.random.PRNGKey value.")
+      raise ValueError(msg)
+    if mode == 'train':
+      keep = random.bernoulli(rng, rate, inputs.shape)
+      return np.where(keep, inputs / rate, 0)
+    else:
+      return inputs
+  return init_fun, apply_fun
+
+
+# Composing layers via combinators
+
+
+def serial(*layers):
+  """Combinator for composing layers in serial.
+
+  Args:
+    *layers: a sequence of layers, each an (init_fun, apply_fun) pair.
+
+  Returns:
+    A new layer, meaning an (init_fun, apply_fun) pair, representing the serial
+    composition of the given sequence of layers.
+  """
+  nlayers = len(layers)
+  init_funs, apply_funs = zip(*layers)
+  def init_fun(input_shape):
+    params = []
+    for init_fun in init_funs:
+      input_shape, param = init_fun(input_shape)
+      params.append(param)
+    return input_shape, params
+  def apply_fun(params, inputs, **kwargs):
+    rng = kwargs.pop('rng', None)
+    rngs = random.split(rng, nlayers) if rng is not None else (None,) * nlayers
+    for fun, param, rng in zip(apply_funs, params, rngs):
+      inputs = fun(param, inputs, rng=rng, **kwargs)
+    return inputs
+  return init_fun, apply_fun
+
+
+def parallel(*layers):
+  """Combinator for composing layers in parallel.
+
+  The layer resulting from this combinator is often used with the FanOut and
+  FanInSum layers.
+
+  Args:
+    *layers: a sequence of layers, each an (init_fun, apply_fun) pair.
+
+  Returns:
+    A new layer, meaning an (init_fun, apply_fun) pair, representing the
+    parallel composition of the given sequence of layers. In particular, the
+    returned layer takes a sequence of inputs and returns a sequence of outputs
+    with the same length as the argument `layers`.
+  """
+  nlayers = len(layers)
+  init_funs, apply_funs = zip(*layers)
+  def init_fun(input_shape):
+    return zip(*[init(shape) for init, shape in zip(init_funs, input_shape)])
+  def apply_fun(params, inputs, **kwargs):
+    rng = kwargs.pop('rng', None)
+    rngs = random.split(rng, nlayers) if rng is not None else (None,) * nlayers
+    return [f(p, x, rng=r, **kwargs) for f, p, x, r in zip(apply_funs, params, inputs, rngs)]
+  return init_fun, apply_fun
+
+
+def shape_dependent(make_layer):
+  """Combinator to delay layer constructor pair until input shapes are known.
+
+  Args:
+    make_layer: a one-argument function that takes an input shape as an argument
+      (a tuple of positive integers) and returns an (init_fun, apply_fun) pair.
+
+  Returns:
+    A new layer, meaning an (init_fun, apply_fun) pair, representing the same
+    layer as returned by `make_layer` but with its construction delayed until
+    input shapes are known.
+  """
+  def init_fun(input_shape):
+    return make_layer(input_shape)[0](input_shape)
+  def apply_fun(params, inputs, **kwargs):
+    return make_layer(inputs.shape)[1](params, inputs, **kwargs)
+  return init_fun, apply_fun

From d7c9c6078fe4a56ba8df0ff71f99cb7cea94083c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 5 Apr 2019 09:48:02 -0700
Subject: [PATCH 1869/2720] Adding Pillow package dependency.

PiperOrigin-RevId: 242140706
---
 setup.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/setup.py b/setup.py
index 0a37fcbff..3a3996921 100644
--- a/setup.py
+++ b/setup.py
@@ -50,6 +50,7 @@
         'numpy',
         'oauth2client',
         'opencv-python',
+        'Pillow',
         'pypng',
         'requests',
         'scipy',

From 3562111a30d3ede91345e277a3a8566ca0300fb6 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 5 Apr 2019 11:11:08 -0700
Subject: [PATCH 1870/2720] Fix OSS test for reversible layers.

PiperOrigin-RevId: 242157332
---
 tensor2tensor/layers/reversible_layers_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
index aaf7c6970..335623fb5 100644
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -23,10 +23,10 @@
 import numpy as np
 
 from tensor2tensor.layers import reversible_layers as reversible
-
 from tensor2tensor.utils import test_utils
 
 import tensorflow as tf
+from tensorflow_probability import edward2 as ed
 tf.compat.v1.enable_eager_execution()
 
 

From 42b2d348e7601dd5bbe5590521464083b19b2b44 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 5 Apr 2019 12:11:02 -0700
Subject: [PATCH 1871/2720] Internal change

PiperOrigin-RevId: 242169130
---
 tensor2tensor/insights/server.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/insights/server.py b/tensor2tensor/insights/server.py
index 942e87504..122064f97 100644
--- a/tensor2tensor/insights/server.py
+++ b/tensor2tensor/insights/server.py
@@ -184,7 +184,6 @@ def root(path):  # pylint: disable=unused-variable
       The landing page html text.
     """
     if (path == "index.js" or
-        path == "webcomponentsjs/custom-elements-es5-adapter.js" or
         path == "webcomponentsjs/webcomponents-lite.js"):
       # Some vulcanizing methods bundle the javascript into a index.js file
       # paired with index.html but leave two important webcomponents js files

From 5ddd46da958806453872429f1f565c0061710608 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Fri, 5 Apr 2019 17:18:05 -0700
Subject: [PATCH 1872/2720] Move the unfortunately-named
 examples/transformer_standalone.py to transformer/main.py make model and
 layers hyperparameters gin-configurable so as to cut out a lot of plumbing.
 Rename "encoder" to "vocabulary" to avoid confusion with the encoder in
 encoder/decoder architectures. Add functionality to run transformer on t2t
 datasets. Modify some of the default model hyperparameters to better match
 previous transformer experiments. Enable layer-postprocess-dropout, which had
 been mistakenly omitted. Add options for sharing embedding and softmax layer
 weights. Several changes to speed up input pipeline:  - custom op for
 sequence-packing (borrowed from t2t)    (requires custom-built tf binary, so
 won't work yet on cloud-tpu)  - enable parallelism in dataset.map calls  -
 add prefetching  - input pipeline seems to run fine now at >1M tokens/sec

PiperOrigin-RevId: 242222578
---
 tensor2tensor/models/mtf_transformer2.py | 29 +++++++++++++++++-------
 1 file changed, 21 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index c8cb5354d..f4b37752b 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -218,21 +218,34 @@ def model(self):
     hparams = self._hparams
     encoder_layer_stack = layer_stack_from_hparams(hparams, "encoder_")
     decoder_layer_stack = layer_stack_from_hparams(hparams, "decoder_")
-    return transformer.Bitransformer(
-        encoder_layer_stack=encoder_layer_stack,
-        decoder_layer_stack=decoder_layer_stack,
-        encoder_d_model=hparams.d_model,
-        decoder_d_model=hparams.d_model,
+    encoder = transformer.Unitransformer(
+        layer_stack=encoder_layer_stack,
+        d_model=hparams.d_model,
         input_vocab_size=self._inputs_vocab_size,
+        output_vocab_size=None,
+        autoregressive=False,
+        max_length=hparams.max_length,
+        name="encoder",
+        layout=hparams.layout,
+        mesh_shape=hparams.mesh_shape,
+    )
+    decoder = transformer.Unitransformer(
+        layer_stack=decoder_layer_stack,
+        d_model=hparams.d_model,
+        input_vocab_size=self._targets_vocab_size,
         output_vocab_size=self._targets_vocab_size,
+        autoregressive=True,
         max_length=hparams.max_length,
-        shared_embedding=hparams.shared_embedding,
+        label_smoothing=hparams.label_smoothing,
         shared_embedding_and_softmax_weights=(
             hparams.shared_embedding_and_softmax_weights),
-        label_smoothing=hparams.label_smoothing,
         z_loss=hparams.z_loss,
+        name="decoder",
         layout=hparams.layout,
-        mesh_shape=hparams.mesh_shape)
+        mesh_shape=hparams.mesh_shape,
+    )
+    return transformer.Bitransformer(
+        encoder, decoder, shared_embedding=hparams.shared_embedding)
 
   def _mtf_model_fn(self, features, mesh):
     self._original_features = features

From 22b964cfce40f6b1ab2cd215a583d17c6d1e40d7 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 5 Apr 2019 21:00:24 -0700
Subject: [PATCH 1873/2720] A PPO implementation in JAX.

Much work remains to be done.

PiperOrigin-RevId: 242240039
---
 tensor2tensor/trax/rlax/ppo.py | 676 +++++++++++++++++++++++++++++++++
 1 file changed, 676 insertions(+)
 create mode 100644 tensor2tensor/trax/rlax/ppo.py

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
new file mode 100644
index 000000000..607c141fc
--- /dev/null
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -0,0 +1,676 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""PPO in JAX.
+
+Automatically generated by Colaboratory.
+
+Original file is located at
+    https://colab.research.google.com/drive/10TTS11vHoZYD4n1h7EXNX5jhLgfqZK6r
+
+# [go/ppo-jax](http://go/ppo-jax)
+
+ - Full batched implementation
+ - Policy implementations
+  - Greedy
+  - Epsilon Greedy
+  - Categorical
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+import time
+from absl import app
+import gym
+
+from jax import grad
+from jax import jit
+from jax import lax
+from jax import numpy as np
+from jax import vmap
+from jax.experimental import optimizers
+from jax.experimental import stax
+from jax.experimental.stax import Dense
+from jax.experimental.stax import Relu
+from jax.experimental.stax import Softmax
+
+import numpy as onp
+
+DEBUG_LOGGING = False
+GAMMA = 0.99
+LAMBDA = 0.95
+EPOCHS = 50  # 100
+NUM_OPTIMIZER_STEPS = 100
+PRINT_EVERY_OPTIMIZER_STEP = 20
+BATCH_TRAJECTORIES = 32
+POLICY = "categorical-sampling"
+
+
+# TODO(afrozm): Have a single net for both policy and value.
+def initialize_policy_and_value_nets(num_actions, batch_observations_shape):
+  """Setup and initialize the policy and value networks."""
+  policy_net_init, policy_net_apply = stax.serial(
+      Dense(16),
+      Relu,
+      Dense(4),
+      Relu,
+      Dense(num_actions),
+      Softmax,
+  )
+
+  _, policy_net_params = policy_net_init(
+      batch_observations_shape)
+
+  value_net_init, value_net_apply = stax.serial(
+      Dense(16),
+      Relu,
+      Dense(4),
+      Relu,
+      Dense(1),  # 1 since we want to predict reward using value network.
+  )
+
+  _, value_net_params = value_net_init(
+      batch_observations_shape)
+
+  return ((policy_net_params, policy_net_apply), (value_net_params,
+                                                  value_net_apply))
+
+
+def initialize_optimizers(policy_net_params, value_net_params):
+  """Initialize optimizers for the policy and value params."""
+  # ppo_opt_init, ppo_opt_update = optimizers.sgd(step_size=1e-3)
+  # val_opt_init, val_opt_update = optimizers.sgd(step_size=1e-3)
+  ppo_opt_init, ppo_opt_update = optimizers.adam(
+      step_size=1e-3, b1=0.9, b2=0.999, eps=1e-08)
+  value_opt_init, value_opt_update = optimizers.adam(
+      step_size=1e-3, b1=0.9, b2=0.999, eps=1e-08)
+
+  ppo_opt_state = ppo_opt_init(policy_net_params)
+  value_opt_state = value_opt_init(value_net_params)
+
+  return (ppo_opt_state, ppo_opt_update), (value_opt_state, value_opt_update)
+
+
+# Should this be collect 'n' trajectories, or
+# Run the env for 'n' steps and take completed trajectories, or
+# Any other option?
+def collect_trajectories(env,
+                         policy_net_apply,
+                         policy_net_params,
+                         num_trajectories=1,
+                         policy="greedy",
+                         epsilon=0.1):
+  """Collect trajectories with the given policy net and behaviour."""
+  trajectories = []
+
+  for _ in range(num_trajectories):
+    observations = []
+    rewards = []
+    actions = []
+    done = False
+
+    observation = env.reset()
+    observations.append(observation)
+    while not done:
+      # Run the policy, pick an action.
+      predictions = policy_net_apply(policy_net_params, observation)
+
+      # Greedy policy.
+      action = np.argmax(predictions)
+      if policy == "epsilon-greedy":
+        # A schedule for epsilon is 1/k where k is the episode number sampled.
+        if onp.random.random() < epsilon:
+          # Choose an action at random.
+          action = onp.random.randint(0, high=len(predictions))
+        else:
+          # Return the best action.
+          action = np.argmax(predictions)
+      elif policy == "categorical-sampling":
+        action = int(onp.argwhere(onp.random.multinomial(1, predictions) == 1))
+
+      if DEBUG_LOGGING:
+        print("With predictions: ", predictions, " chose action: ", action)
+
+      # NOTE: Assumption, single batch.
+      action = int(action)
+
+      observation, reward, done, _ = env.step(action)
+
+      observations.append(observation)
+      rewards.append(reward)
+      actions.append(action)
+
+    # This means we are done
+    assert done
+    trajectories.append((np.stack(observations), np.stack(actions),
+                         np.stack(rewards)))
+
+  return trajectories
+
+
+# This function can probably be simplified, ask how?
+# Can we do something much simpler than lax.pad, maybe np.pad?
+# Others?
+def pad_trajectories(trajectories, boundary=10):
+  """Pad trajectories to a bucket length that is a multiple of boundary."""
+
+  # trajectories is a list of tuples of (observations, actions, rewards)
+  # observations's length is one more than actions and rewards
+  #
+  # i.e. observations = (o_0, o_1, ... o_{T-1}, o_T)
+  #           actions = (a_0, a_1, ... a_{T-1})
+  #           rewards = (r_0, r_1, ... r_{T-1})
+
+  # Given the above, let's compute max(T) over all trajectories.
+  t_max = max(o.shape[0] for (o, a, r) in trajectories)
+  if DEBUG_LOGGING:
+    print("t_max: %s" % t_max)
+
+  # t_max - 1 is rounded to the next multiple of `boundary`
+  boundary = int(boundary)
+  bucket_length = boundary * int(np.ceil(float(t_max - 1) / boundary))
+  if DEBUG_LOGGING:
+    print("bucket_length: %s" % bucket_length)
+
+  # So all obs will be padded to t_max and actions and rewards to t_max - 1.
+  padded_observations = []
+  padded_actions = []
+  padded_rewards = []
+  padded_lengths = []
+  reward_masks = []
+  for (o, a, r) in trajectories:
+    # Determine the amount to pad, this holds true for obs, actions and rewards.
+    num_to_pad = bucket_length + 1 - o.shape[0]
+    padded_lengths.append(num_to_pad)
+    if num_to_pad == 0:
+      padded_observations.append(o)
+      padded_actions.append(a)
+      padded_rewards.append(r)
+      reward_masks.append(onp.ones_like(r, dtype=np.int32))
+      continue
+
+    # First pad observations.
+    if DEBUG_LOGGING:
+      print("num_to_pad: %s" % num_to_pad)
+    padding_config = [(0, num_to_pad, 0)]
+    for _ in range(o.ndim - 1):
+      padding_config.append((0, 0, 0))
+    padding_config = tuple(padding_config)
+    if DEBUG_LOGGING:
+      print("padding_config: %s" % str(padding_config))
+    padding_value = 0.0 if o.dtype == np.float32 else 0
+    if DEBUG_LOGGING:
+      print("padding_value: %s" % padding_value)
+    padded_obs = lax.pad(o, padding_value, padding_config)
+    padded_observations.append(padded_obs)
+
+    # Now pad actions and rewards.
+    assert a.ndim == 1 and r.ndim == 1
+    padding_config = ((0, num_to_pad, 0),)
+    if DEBUG_LOGGING:
+      print("action/reward padding_config: %s" % str(padding_config))
+    action_padding_value = 0.0 if a.dtype == np.float32 else 0
+    reward_padding_value = 0.0 if r.dtype == np.float32 else 0
+    if DEBUG_LOGGING:
+      print("action_padding_value: %s" % action_padding_value)
+    padded_action = lax.pad(a, action_padding_value, padding_config)
+    padded_actions.append(padded_action)
+    if DEBUG_LOGGING:
+      print("reward_padding_value: %s" % reward_padding_value)
+    padded_reward = lax.pad(r, reward_padding_value, padding_config)
+    padded_rewards.append(padded_reward)
+
+    # Also create the mask to use later.
+    reward_mask = onp.ones_like(r, dtype=np.int32)
+    reward_masks.append(lax.pad(reward_mask, 0, padding_config))
+
+  return padded_lengths, np.stack(reward_masks), np.stack(
+      padded_observations), np.stack(padded_actions), np.stack(padded_rewards)
+
+
+def rewards_to_go_discounted(rewards, reward_mask=1.0, gamma=0.99):
+  r"""r2g[t] = \sum_{l=0}^{\infty}(\gamma^l * r_{t+l})."""
+  time_steps = len(rewards)
+  # r2g[t] = r[t] + (gamma * r2g[t+1])
+
+  # First initialize like:
+  # r2g[t] = r[t], for t = 0 to T-1
+  rewards_to_go = list(rewards)
+
+  # Then add the discounted version of the next time-step.
+  # i = [T-2 .. 0]
+  for i in range(time_steps - 2, -1, -1):
+    rewards_to_go[i] += gamma * rewards_to_go[i + 1]
+
+  # Makes this back into JAX's DeviceArray
+  rewards_to_go = np.stack(list(rewards_to_go))
+
+  return rewards_to_go * reward_mask
+
+
+def batched_avg_value_function_loss(value_net_apply,
+                                    value_net_params,
+                                    observations,
+                                    rewards,
+                                    reward_mask=1.0,
+                                    gamma=0.99):
+  """L2 loss on the value function's outputs."""
+  # Capturing the value_net_apply from the parent function's scope.
+  # See: https://github.com/google/jax/issues/183
+  def _value_function_loss_trajectory(value_net_params,
+                                      observations,
+                                      rewards,
+                                      reward_mask=1.0,
+                                      gamma=0.99):
+    """Compute the actual loss for a trajectory."""
+    r2g = rewards_to_go_discounted(
+        rewards, reward_mask=reward_mask, gamma=gamma)
+    v = value_net_apply(value_net_params, observations[:-1])
+    v = np.squeeze(v) * reward_mask
+    loss = v - r2g
+    return np.sum(loss**2)
+
+  batched_value_function_loss_trajectory = vmap(
+      _value_function_loss_trajectory, in_axes=(None, 0, 0, 0), out_axes=0)
+
+  return np.mean(
+      batched_value_function_loss_trajectory(
+          value_net_params, observations, rewards, reward_mask, gamma=gamma))
+
+
+def batched_deltas(predicted_values, rewards, reward_mask, gamma=0.99):
+  r"""\delta_t = \sum_{l = 0}^{\infty}(r_t + \gamma * V(s_{t+1}) - V(s_t))."""
+  # predicted_values are application of value net only the observations.
+  # B x T+1
+
+  deltas = []
+  _, T = rewards.shape  # pylint: disable=invalid-name
+  for t in range(T):
+    deltas.append(rewards[:, t] + (gamma * predicted_values[:, t + 1]) -
+                  predicted_values[:, t])
+
+  return np.array(deltas).T * reward_mask
+
+
+def batched_gae_advantages(deltas, reward_mask, lamda=0.95,  # NOTYPO
+                           gamma=0.99):
+  r"""A_t = \sum_{l=0}^{\infty}(\gamma * \lambda)^{l}(\delta_{t+l})."""
+  _, T = deltas.shape  # pylint: disable=invalid-name
+  gl = lamda * gamma  # NOTYPO
+
+  # [[1, gl, gl**2, ... gl**T-1]]
+  # Not jittable, T should be a compile time constant.
+  # gl_gp = np.geomspace(1, gl**T, T, endpoint=False).reshape(1, T)
+  gl_geometric_progression = [1]
+  for _ in range(1, T):
+    gl_geometric_progression.append(gl_geometric_progression[-1] * gl)
+  gl_gp = np.array(gl_geometric_progression)
+  gl_gp = gl_gp.reshape((1, T))
+
+  # deltas * gl_gp
+  deltas_gl_gp = deltas * gl_gp
+
+  # A0 - advantage for 0th time-step, across all batches.
+  As = []  # pylint: disable=invalid-name
+  A0 = np.sum(deltas_gl_gp, axis=1)  # (B,)  # pylint: disable=invalid-name
+  As.append(A0)
+
+  # Now compute the other advantages.
+  for t in range(1, T):
+    As.append((As[-1] - deltas[:, t - 1]) / gl)
+
+  return np.stack(As).T * reward_mask
+
+
+def batched_probabs(probab_observations, actions):
+  b, t = actions.shape
+  return probab_observations[np.arange(b)[:, None], np.arange(t), actions]
+
+
+def batched_probab_ratios(policy_net_apply, old_policy_params,
+                          new_policy_params, observations, actions,
+                          reward_mask):
+  """Calculates the probaility ratios for each time-step in a trajectory."""
+  p_old = policy_net_apply(old_policy_params, observations)
+  p_new = policy_net_apply(new_policy_params, observations)
+
+  bp_old = batched_probabs(p_old, actions)
+  bp_new = batched_probabs(p_new, actions)
+
+  if DEBUG_LOGGING:
+    print("bp_old: ", bp_old)
+    print("bp_new: ", bp_new)
+
+  # Add a small number to bp_old, where reward_mask is 0, this is just to help
+  # never to divide by 0.
+  bp_old = bp_old + (0.1 * np.abs(reward_mask - 1))
+
+  if DEBUG_LOGGING:
+    print("masked bp_old: ", bp_old)
+
+  ret_val = (bp_new * reward_mask) / bp_old
+
+  if DEBUG_LOGGING:
+    print("ret_val: ", ret_val)
+
+  return ret_val
+
+
+def batched_clipped_probab_ratios(bpr, reward_mask, epsilon=0.2):
+  return reward_mask * np.clip(bpr, 1 - epsilon, 1 + epsilon)
+
+
+def batched_clipped_objective(bpr, adv, reward_mask, epsilon=0.2):
+  c1 = bpr * adv
+  c2 = batched_clipped_probab_ratios(bpr, reward_mask, epsilon=epsilon) * adv
+  return np.minimum(c1, c2)
+
+
+def batched_ppo_loss(policy_net_apply,
+                     new_policy_params,
+                     old_policy_params,
+                     value_net_apply,
+                     value_net_params,
+                     padded_observations,
+                     padded_actions,
+                     padded_rewards,
+                     reward_mask,
+                     gamma=0.99,
+                     lamda=0.95,  # NOTYPO
+                     epsilon=0.2):
+  """PPO objective, with an eventual minus sign."""
+  # V(s_t) forall s & t
+  value_function = np.squeeze(
+      value_net_apply(value_net_params, padded_observations))
+  deltas = batched_deltas(
+      value_function, padded_rewards, reward_mask, gamma=gamma)
+  advantages = batched_gae_advantages(
+      deltas, reward_mask, lamda=lamda, gamma=gamma)  # NOTYPO
+  ratios = batched_probab_ratios(policy_net_apply, old_policy_params,
+                                 new_policy_params, padded_observations,
+                                 padded_actions, reward_mask)
+  clipped_loss = batched_clipped_objective(ratios, advantages, reward_mask,
+                                           epsilon=epsilon)
+  return -np.sum(clipped_loss)
+
+
+@functools.partial(jit, static_argnums=(2, 3, 5))
+def ppo_opt_step(i,
+                 opt_state,
+                 ppo_opt_update,
+                 policy_net_apply,
+                 old_policy_params,
+                 value_net_apply,
+                 value_net_params,
+                 padded_observations,
+                 padded_actions,
+                 padded_rewards,
+                 reward_mask,
+                 gamma=0.99,
+                 lamda=0.95,  # NOTYPO
+                 epsilon=0.1):
+  """PPO optimizer step."""
+  new_policy_params = optimizers.get_params(opt_state)
+  g = grad(
+      batched_ppo_loss, argnums=1)(
+          policy_net_apply,
+          new_policy_params,
+          old_policy_params,
+          value_net_apply,
+          value_net_params,
+          padded_observations,
+          padded_actions,
+          padded_rewards,
+          reward_mask,
+          gamma=gamma,
+          lamda=lamda,  # NOTYPO
+          epsilon=epsilon)
+  return ppo_opt_update(i, g, opt_state)
+
+
+@functools.partial(jit, static_argnums=(2, 3))
+def value_opt_step(i,
+                   opt_state,
+                   opt_update,
+                   value_net_apply,
+                   padded_observations,
+                   padded_rewards,
+                   reward_mask,
+                   gamma=0.99):
+  """Value optimizer step."""
+  value_params = optimizers.get_params(opt_state)
+  # Note this partial application here and argnums above in ppo_opt_step.
+  g = grad(functools.partial(batched_avg_value_function_loss, value_net_apply))(
+      value_params,
+      padded_observations,
+      padded_rewards,
+      reward_mask,
+      gamma=gamma)
+  return opt_update(i, g, opt_state)
+
+
+def main(unused_argv):
+  onp.random.seed(0)
+
+  value_losses = []
+  ppo_objective = []
+  average_rewards = []
+
+  env = gym.make("CartPole-v0")
+
+  print("Initial observation: ", env.reset())
+
+  for i in range(100):
+    random_action = env.action_space.sample()
+    obs, rew, done, _ = env.step(random_action)
+    print("[%s] reward [%s], done [%s] and obs [%s]" % (i, rew, done, obs))
+    if done:
+      print("Done, so exiting, step: ", i)
+      env.close()
+      break
+
+  print("action_space.shape", env.action_space.shape)
+  print("observation_space.shape", env.observation_space.shape)
+
+  batch_observations_shape = (-1,) + env.observation_space.shape
+
+  assert isinstance(env.action_space, gym.spaces.Discrete)
+  num_actions = env.action_space.n
+
+  print("batch_observations_shape: ", batch_observations_shape)
+  print("num_actions: ", num_actions)
+
+  ((policy_net_params, policy_net_apply),
+   (value_net_params, value_net_apply)) = initialize_policy_and_value_nets(
+       num_actions, batch_observations_shape)
+
+  (ppo_opt_state, ppo_opt_update), (value_opt_state,
+                                    value_opt_update) = initialize_optimizers(
+                                        policy_net_params, value_net_params)
+
+  for i in range(EPOCHS):
+    t = time.time()
+    t0 = t
+    trajs = collect_trajectories(
+        env,
+        policy_net_apply,
+        policy_net_params,
+        num_trajectories=BATCH_TRAJECTORIES,
+        policy=POLICY,
+        epsilon=(10.0 / (i + 10.0)))  # this is a different epsilon.
+    print("Took ", round((time.time() - t) * 1000, 2),
+          "msec to collect trajectories.")
+
+    print("Average Trajectory size: ",
+          float(sum(len(traj[0]) for traj in trajs)) / len(trajs))
+    avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
+    average_rewards.append(avg_reward)
+    print("Average sum rewards: ", avg_reward)
+
+    if (avg_reward > 190.0) and (i % 5 == 0):
+      print("policy_net_params:\n", policy_net_params)
+      print("value_net_params:\n", value_net_params)
+
+    t = time.time()
+    (_, reward_mask, padded_observations, padded_actions,
+     padded_rewards) = pad_trajectories(trajs, boundary=20)
+    print("Took ", round((time.time() - t) * 1000, 2),
+          "msec to pad trajectories.")
+
+    print("Padded Observations' shape: ", padded_observations.shape)
+    print("Padded Actions' shape:      ", padded_actions.shape)
+    print("Padded Rewards' shape:      ", padded_rewards.shape)
+
+    # Linear annealing from 0.1 to 0.0
+    epsilon = 0.1 if EPOCHS == 1 else 0.1 * (1.0 - (i / (EPOCHS - 1)))
+
+    t = time.time()
+    val_loss = jit(
+        batched_avg_value_function_loss, static_argnums=(0,))(
+            value_net_apply,
+            value_net_params,
+            padded_observations,
+            padded_rewards,
+            reward_mask,
+            gamma=GAMMA)
+
+    print("Took ", round((time.time() - t) * 1000, 2),
+          "msec to calculate value loss = ", val_loss)
+    value_losses.append(val_loss)
+
+    t = time.time()
+    ppo_loss = jit(
+        batched_ppo_loss, static_argnums=(0,
+                                          3))(policy_net_apply,
+                                              policy_net_params,
+                                              policy_net_params,
+                                              value_net_apply,
+                                              value_net_params,
+                                              padded_observations,
+                                              padded_actions,
+                                              padded_rewards,
+                                              reward_mask,
+                                              gamma=GAMMA,
+                                              lamda=LAMBDA,  # NOTYPO
+                                              epsilon=epsilon)
+    # ppo_loss = 11.00110011
+    print("Took ", round((time.time() - t) * 1000, 2),
+          "msec to calculate ppo loss = ", ppo_loss)
+    ppo_objective.append(-ppo_loss)
+
+    # Run optimizers.
+    t1 = time.time()
+
+    print("PPO objective optimization.")
+
+    for j in range(NUM_OPTIMIZER_STEPS):
+      t = time.time()
+      # Update the optimizer state.
+      ppo_opt_state = ppo_opt_step(
+          j,
+          ppo_opt_state,
+          ppo_opt_update,
+          policy_net_apply,
+          policy_net_params,
+          value_net_apply,
+          value_net_params,
+          padded_observations,
+          padded_actions,
+          padded_rewards,
+          reward_mask,
+          gamma=GAMMA,
+          lamda=LAMBDA,  # NOTYPO
+          epsilon=epsilon)
+      t2 = time.time()
+      # Get the new params.
+      new_policy_net_params = optimizers.get_params(ppo_opt_state)
+      if ((j + 1) %
+          PRINT_EVERY_OPTIMIZER_STEP == 0) or (j == NUM_OPTIMIZER_STEPS - 1):
+        new_ppo_loss = jit(
+            batched_ppo_loss, static_argnums=(0,
+                                              3))(policy_net_apply,
+                                                  new_policy_net_params,
+                                                  policy_net_params,
+                                                  value_net_apply,
+                                                  value_net_params,
+                                                  padded_observations,
+                                                  padded_actions,
+                                                  padded_rewards,
+                                                  reward_mask,
+                                                  gamma=GAMMA,
+                                                  lamda=LAMBDA,  # NOTYPO
+                                                  epsilon=epsilon)
+        print("Took ", round((t2 - t) * 1000, 2),
+              "msec to do one step ppo grad desc")
+        print("New ppo loss[", j, "]: ", new_ppo_loss, " vs old ppo loss: ",
+              ppo_loss)
+      # Update the params.
+      policy_net_params = new_policy_net_params
+
+    print("Total ppo loss reduction: ",
+          100 * (ppo_loss - new_ppo_loss) / np.abs(ppo_loss), "%")
+
+    print("Value optimization.")
+
+    for j in range(NUM_OPTIMIZER_STEPS):
+      t = time.time()
+      value_opt_state = value_opt_step(
+          j,
+          value_opt_state,
+          value_opt_update,
+          value_net_apply,
+          padded_observations,
+          padded_rewards,
+          reward_mask,
+          gamma=GAMMA)
+      t2 = time.time()
+      value_net_params = optimizers.get_params(value_opt_state)
+      if ((j + 1) %
+          PRINT_EVERY_OPTIMIZER_STEP == 0) or (j == NUM_OPTIMIZER_STEPS - 1):
+        new_val_loss = jit(
+            batched_avg_value_function_loss, static_argnums=(0,))(
+                value_net_apply,
+                value_net_params,
+                padded_observations,
+                padded_rewards,
+                reward_mask,
+                gamma=GAMMA)
+        print("Took ", round((t2 - t) * 1000, 2),
+              "msec to do one step value grad desc")
+        print("New value loss[", j, "]: ", new_val_loss, " vs old value loss: ",
+              val_loss)
+    print("Total value loss reduction: ",
+          100 * (val_loss - new_val_loss) / val_loss, "%")
+
+    print("Took ", round((time.time() - t1) * 1000, 2), "msec to do grad desc")
+
+    # Set the optimized params to new params.
+    policy_net_params = optimizers.get_params(ppo_opt_state)
+    value_net_params = optimizers.get_params(value_opt_state)
+
+    print("Epoch [%s] took [%s]msec." % (i, round(
+        (time.time() - t0) * 1000, 2)))
+    print()
+
+  print("value_losses: ", np.stack(value_losses))
+  print("ppo_objective: ", np.stack(ppo_objective))
+  print("average_rewards: ", average_rewards)
+
+
+if __name__ == "__main__":
+  app.run(main)

From ef283ac1b1f0e66b3abfab175cf90ee1fbbf84ed Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 6 Apr 2019 19:22:20 -0700
Subject: [PATCH 1874/2720] Change comments.

PiperOrigin-RevId: 242312909
---
 tensor2tensor/trax/rlax/ppo.py | 16 +---------------
 1 file changed, 1 insertion(+), 15 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 607c141fc..87cb1935c 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -13,21 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""PPO in JAX.
-
-Automatically generated by Colaboratory.
-
-Original file is located at
-    https://colab.research.google.com/drive/10TTS11vHoZYD4n1h7EXNX5jhLgfqZK6r
-
-# [go/ppo-jax](http://go/ppo-jax)
-
- - Full batched implementation
- - Policy implementations
-  - Greedy
-  - Epsilon Greedy
-  - Categorical
-"""
+"""PPO in JAX."""
 
 from __future__ import absolute_import
 from __future__ import division

From 3fb0eee9887374c66a1d570adacbd48d65f4010d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 8 Apr 2019 08:34:56 -0700
Subject: [PATCH 1875/2720] Internal change.

PiperOrigin-RevId: 242467951
---
 tensor2tensor/layers/common_layers_test.py | 11 -----------
 1 file changed, 11 deletions(-)

diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index cc3b8bb73..fe20c2956 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -214,10 +214,6 @@ def testLayerNorm(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 7, 11))
 
-    # Testing layer collection.
-    layer_collection = kfac.LayerCollection()
-    common_layers.layer_norm(x, layer_collection=layer_collection)
-    self.assertLen(layer_collection.get_blocks(), 1)
 
   @test_utils.run_in_graph_and_eager_modes()
   def testGroupNorm(self):
@@ -317,13 +313,6 @@ def testApplyNormNone(self):
     self.assertEqual(actual.shape, (5, 2, 1, 11))
     self.assertAllClose(actual, x1, atol=1e-03)
 
-  @test_utils.run_in_graph_mode_only()
-  def testApplyNormWithLayerCollection(self):
-    x = np.random.rand(5, 2, 1, 11)
-    layer_collection = kfac.LayerCollection()
-    common_layers.apply_norm(x, "layer", depth=11, epsilon=1e-6,
-                             layer_collection=layer_collection)
-    self.assertLen(layer_collection.get_blocks(), 1)
 
   @test_utils.run_in_graph_mode_only()
   def testDenseWithLayerCollection(self):

From 79b05828f714f7eea5ad4b4a7714e54ce35a329b Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 8 Apr 2019 10:01:21 -0700
Subject: [PATCH 1876/2720] Don't require installing trax dependencies (jax,
 jaxlib) since they aren't available for Windows.

PiperOrigin-RevId: 242481869
---
 oss_scripts/oss_pip_install.sh | 2 +-
 setup.py                       | 6 ++++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/oss_scripts/oss_pip_install.sh b/oss_scripts/oss_pip_install.sh
index 333f49b41..9db9854a4 100755
--- a/oss_scripts/oss_pip_install.sh
+++ b/oss_scripts/oss_pip_install.sh
@@ -22,6 +22,6 @@ t2t-trainer --registry_help 2>&1 >/dev/null
 t2t-datagen 2>&1 | grep translate_ende 2>&1 >/dev/null && echo passed
 
 # Then install the test dependencies
-pip install -q -e .[tests,allen]
+pip install -q -e .[tests,allen,trax]
 # Make sure to install the atari extras for gym
 pip install "gym[atari]"
diff --git a/setup.py b/setup.py
index 3a3996921..e6ac6bf1b 100644
--- a/setup.py
+++ b/setup.py
@@ -43,8 +43,6 @@
         'gunicorn',
         'gym',
         'h5py',
-        'jax',
-        'jaxlib',
         'kfac',
         'mesh-tensorflow',
         'numpy',
@@ -77,6 +75,10 @@
             # explicit pip install gym[atari] for the tests.
             # 'gym[atari]',
         ],
+        'trax': [
+            'jax',
+            'jaxlib',
+        ],
         'allen': ['Pillow==5.1.0', 'pandas==0.23.0'],
     },
     classifiers=[

From 012d521144e45b8c8d19615266f503cc6c21312e Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 8 Apr 2019 11:07:46 -0700
Subject: [PATCH 1877/2720] Bump setup.py to 1.13.2 -- Travis is green.

PiperOrigin-RevId: 242497797
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index e6ac6bf1b..5b895b689 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.13.1',
+    version='1.13.2',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',

From cfaf657210749979fb31797451056f51bfde6831 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 8 Apr 2019 12:22:33 -0700
Subject: [PATCH 1878/2720] Enable multiple chunk numbers per batch with
 recurrent memory

This also fixes a bug where after a new sequence started, the memory bias term would allow attending to memory items from the previous sequence.

PiperOrigin-RevId: 242513103
---
 tensor2tensor/layers/common_hparams.py     |  1 +
 tensor2tensor/layers/transformer_memory.py | 52 ++++++++++++----------
 tensor2tensor/models/transformer.py        | 11 ++---
 tensor2tensor/utils/data_reader.py         | 32 +++++++++++++
 4 files changed, 68 insertions(+), 28 deletions(-)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 55b147987..9988283ff 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -158,6 +158,7 @@ def basic_params1():
       # Split targets on the first axis into chunks of this length.
       split_targets_chunk_length=0,
       split_targets_max_chunks=100,
+      split_targets_strided_training=False,
       # Maximum length in the smallest length bucket.  Setting this
       # flag too high will result in wasteful padding of short
       # sequences.  Due to some (hopefully) temporary hacks in the
diff --git a/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index 92eaeadd5..002e4004c 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -86,11 +86,11 @@ def __init__(self, name, hparams):
       batch_size_in_sequences = hparams.batch_size / hparams.max_length
 
     memory_shape = [batch_size_in_sequences, self.tokens_to_cache, hidden_size]
-    bias_shape = [1, 1, self.chunk_length, self.tokens_to_cache]
+    bias_shape = [batch_size_in_sequences, 1, 1, self.tokens_to_cache]
 
     with tf.variable_scope(name):
       self.previous_segment = tf.get_variable(
-          "memsegment", (),
+          "memsegment", (batch_size_in_sequences,),
           dtype=tf.int32, trainable=False,
           collections=[tf.GraphKeys.LOCAL_VARIABLES],
           initializer=tf.constant_initializer(0))
@@ -105,7 +105,7 @@ def __init__(self, name, hparams):
           "membias", bias_shape,
           dtype=tf.float32, trainable=False,
           collections=[tf.GraphKeys.LOCAL_VARIABLES],
-          initializer=tf.constant_initializer(.0))
+          initializer=tf.constant_initializer(-1e9))
 
   def pre_attention(self, segment, query_antecedent, memory_antecedent, bias):
     """Called prior to self-attention, to incorporate memory items.
@@ -122,39 +122,45 @@ def pre_attention(self, segment, query_antecedent, memory_antecedent, bias):
     """
     assert memory_antecedent is None, "We only support language modeling"
 
+    # In eval mode, batch size may be variable
+    memory_batch_size = tf.shape(self.previous_vals)[0]
+    current_batch_size = tf.shape(query_antecedent)[0]
+    amount_to_pad = memory_batch_size - current_batch_size
+
     previous_vals = self.previous_vals
     # If segment id is zero, don't attend back to the memory
-    previous_bias = self.previous_bias + tf.cast(
-        tf.equal(tf.reduce_sum(segment), 0), tf.float32) * -1e9
+    previous_bias = self.previous_bias[:current_batch_size, :, :, :] + tf.cast(
+        tf.equal(segment[:, None, None, None], 0), tf.float32) * -1e9
 
-    # In eval mode, batch size may be variable
-    amount_to_pad = tf.shape(previous_vals)[0] - tf.shape(query_antecedent)[0]
-    sliced_previous_vals = previous_vals[:tf.shape(query_antecedent)[0], :, :]
+    sliced_previous_vals = previous_vals[:current_batch_size, :, :]
 
     new_memory_antecedent = tf.concat(
         [tf.stop_gradient(sliced_previous_vals), query_antecedent], 1)
-    new_bias = tf.concat([tf.stop_gradient(previous_bias), bias], -1)
+    new_bias = tf.concat([
+        tf.tile(tf.stop_gradient(previous_bias), [1, 1, self.chunk_length, 1]),
+        tf.tile(bias, [current_batch_size, 1, 1, 1]),
+    ], -1)
 
-    remember_segment = segment[0]
+    remember_segment = tf.pad(segment, [[0, amount_to_pad]])
     # TODO(kitaev): The code assumes that we always either increment the chunk
-    # number or reset it to zero, which is checked by the assertion. This
-    # assumption will not hold if we re-run the model for each token, e.g. for
-    # autoregressive greedy/beam/sampling decode.
-    with tf.control_dependencies(
-        [tf.Assert(tf.math.logical_or(
-            tf.equal(remember_segment, 0),
-            tf.equal(remember_segment, self.previous_segment + 1)),
-                   [self.previous_segment, remember_segment])]):
-      remember_segment = tf.identity(remember_segment)
-
+    # number or reset it to zero. This assumption will not hold if we re-run the
+    # model for each token, e.g. for autoregressive greedy/beam/sampling decode.
     remember_vals = tf.pad(query_antecedent,
                            [[0, amount_to_pad], [0, 0], [0, 0]])
-    remember_bias = tf.zeros_like(bias) + tf.reduce_max(
-        bias, -1, keep_dims=True)
+    # Query position is on axis -2 for bias: as long as a token can be attended
+    # to from at least one query position (i.e. it's not padding), memorize it.
+    remember_bias = tf.tile(
+        tf.reduce_max(bias, -2, keepdims=True), [memory_batch_size, 1, 1, 1])
     # Assume that query_antecedent is always a full chunk (i.e. not truncated)
     if self.chunk_length < self.tokens_to_cache:
       remember_vals = tf.concat([previous_vals, remember_vals], 1)
-      remember_bias = tf.concat([previous_bias, remember_bias], -1)
+      remember_bias = tf.concat([
+          previous_bias - 1e9 * tf.cast(
+              tf.equal(
+                  tf.pad(segment, [[0, amount_to_pad]])[:, None, None, None],
+                  0), tf.float32),
+          remember_bias
+      ], -1)
     if self.chunk_length != self.tokens_to_cache:
       remember_vals = remember_vals[:, -self.tokens_to_cache:, :]
       remember_bias = remember_bias[:, :, :, -self.tokens_to_cache:]
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index e3d5ed624..dc1cb1b40 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -253,19 +253,19 @@ def body(self, features):
     if self.recurrent_memory_by_layer is not None:
       # TODO(kitaev): The chunk_number feature currently has the same shape as
       # "targets", but this is only for the purposes of sharing sharding code.
-      # In fact every token within the batch must have the same chunk number.
+      # In fact every token within an example must have the same chunk number.
       chunk_number_each_token = tf.squeeze(features["chunk_number"], (-1, -2))
-      chunk_number_each_batch = chunk_number_each_token[:, 0]
+      chunk_number_each_example = chunk_number_each_token[:, 0]
       # Uncomment the code below to verify that tokens within a batch share the
       # same chunk number:
       # with tf.control_dependencies([
       #     tf.assert_equal(chunk_number_each_token,
-      #                     chunk_number_each_batch[:, None])
+      #                     chunk_number_each_example[:, None])
       # ]):
-      #   chunk_number_each_batch = tf.identity(chunk_number_each_batch)
+      #   chunk_number_each_example = tf.identity(chunk_number_each_example)
       decode_kwargs = dict(
           recurrent_memory_by_layer=self.recurrent_memory_by_layer,
-          chunk_number=chunk_number_each_batch,
+          chunk_number=chunk_number_each_example,
           )
 
     decoder_output = self.decode(
@@ -2651,6 +2651,7 @@ def transformer_wikitext103_l4k_memory_v0():
 
   hparams.split_targets_chunk_length = 64
   hparams.split_targets_max_chunks = 64
+  hparams.split_targets_strided_training = True
   hparams.add_hparam("memory_type", "transformer_xl")
 
   # The hparams specify batch size *before* chunking, but we want to have a
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index c623ea47b..7f29944e6 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -519,6 +519,38 @@ def split_on_length(example):
     dataset = dataset.flat_map(split_on_length)
     dataset = dataset.filter(is_nonzero_chunk)
 
+    # The chunking data pipeline thus far creates batches of examples where all
+    # of the examples have the same chunk number. This can lead to periodic
+    # fluctuations in the loss; for example, when all examples in the batch have
+    # chunk number 0 the loss may be higher than midway through a sequence.
+    # Enabling split_targets_strided_training adjusts the data so that each
+    # batch includes examples at various points within a sequence.
+    if is_training and hparams.split_targets_strided_training:
+      # TODO(kitaev): make sure that shape inference works on GPU, not just TPU.
+      inferred_batch_size = dataset.output_shapes["targets"].as_list()[0]
+      if inferred_batch_size is None:
+        raise ValueError(
+            "Strided training is only implemented when the batch size can be "
+            "inferred statically, for example when training on TPU."
+        )
+      chunk_stride = inferred_batch_size * max(
+          1, max_chunks // inferred_batch_size) + 1
+
+      def collapse_nested_datasets(example):
+        """Converts a dataset of datasets to a dataset of tensor features."""
+        new_example = {}
+        for k, v in example.items():
+          v = tf.data.experimental.get_single_element(
+              v.batch(inferred_batch_size, drop_remainder=True))
+          new_example[k] = v
+        return tf.data.Dataset.from_tensor_slices(new_example)
+
+      dataset = dataset.apply(tf.data.experimental.unbatch())
+      dataset = dataset.window(inferred_batch_size, inferred_batch_size,
+                               chunk_stride)
+      dataset = dataset.flat_map(collapse_nested_datasets)
+      dataset = dataset.batch(inferred_batch_size, drop_remainder=True)
+
   def prepare_for_output(example):
     if not config or not config.use_tpu:
       _summarize_features(example, num_shards)

From ccc4d0781e774ceb105cfc8496a61a543bf5ca44 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 8 Apr 2019 13:33:08 -0700
Subject: [PATCH 1879/2720] Allow random inputs in TRAX (mostly for debugging)
 and no-eval training.

PiperOrigin-RevId: 242527716
---
 tensor2tensor/trax/inputs.py  | 43 +++++++++++++++++++++++++++++++++++
 tensor2tensor/trax/trainer.py |  6 ++++-
 tensor2tensor/trax/trax.py    |  8 ++++---
 3 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 99e7aa71d..6355d7b0e 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -26,6 +26,7 @@
 import gin
 
 import jax.numpy as np
+import numpy as onp
 
 import tensorflow as tf
 import tensorflow_datasets as tfds
@@ -70,6 +71,48 @@ def eval_input_fun():
                 input_shape=input_shape)
 
 
+@gin.configurable()
+def random_inputs(
+    input_shape=gin.REQUIRED, input_dtype=onp.int32, input_range=(0, 255),
+    output_shape=gin.REQUIRED, output_dtype=onp.int32, output_range=(0, 9)):
+  """Make random Inputs for debugging.
+
+  Args:
+    input_shape: the shape of inputs (including batch dimension).
+    input_dtype: the type of the inputs (int32 by default).
+    input_range: the range of inputs (defaults to (0, 255)).
+    output_shape: the shape of outputs (including batch dimension).
+    output_dtype: the type of the outputs (int32 by default).
+    output_range: the range of outputs (defaults to (0, 9)).
+
+  Returns:
+    trax.inputs.Inputs
+  """
+  def random_minibatches():
+    """Generate a stream of random mini-batches."""
+    if input_dtype in [onp.float16, onp.float32, onp.float64]:
+      rand = onp.random.uniform
+    else:
+      rand = onp.random.random_integers
+    while True:
+      inp = rand(input_range[0], input_range[1], input_shape)
+      inp = inp.astype(input_dtype)
+      out = rand(output_range[0], output_range[1], output_shape)
+      out = out.astype(output_dtype)
+      yield inp, out
+
+  def train_input_fun():
+    return random_minibatches()
+
+  def eval_input_fun():
+    return random_minibatches()
+
+  input_shape_without_batch = list(input_shape)[1:]
+  return Inputs(train_stream=train_input_fun,
+                eval_stream=eval_input_fun,
+                input_shape=input_shape_without_batch)
+
+
 def dataset_to_stream(dataset, input_name):
   """Takes a tf.Dataset and creates a numpy stream of ready batches."""
   for example in tfds.as_numpy(dataset):
diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index de6895423..b5f9f5ffe 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -47,9 +47,13 @@
 
 def _default_output_dir():
   """Default output directory."""
+  try:
+    dataset_name = gin.query_parameter("inputs.dataset_name")
+  except ValueError:
+    dataset_name = "random"
   dir_name = "{model_name}_{dataset_name}_{timestamp}".format(
       model_name=gin.query_parameter("train.model").configurable.name,
-      dataset_name=gin.query_parameter("inputs.dataset_name"),
+      dataset_name=dataset_name,
       timestamp=datetime.datetime.now().strftime("%Y%m%d_%H%M"),
   )
   dir_path = os.path.join("~", "trax", dir_name)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index afe3b0478..082d2183a 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -351,9 +351,11 @@ def train(output_dir,
 
   print()
   train_stream = inputs.train_stream()
-  epoch_steps = itertools.chain([1,  # first epoch only 1 step
-                                 eval_frequency - 1],
-                                itertools.repeat(eval_frequency))
+  epoch_steps = [train_steps]  # Only training if eval_frequency is 0 or None.
+  if eval_frequency:
+    epoch_steps = itertools.chain([1,  # first epoch only 1 step
+                                   eval_frequency - 1],
+                                  itertools.repeat(eval_frequency))
   step_log(step, "Starting training using %d devices" % num_devices)
 
   # Non-compiled debug step helps find problems in models easier.

From 271a2077279a3e7ea79140b46bdc659fe5dcf439 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 8 Apr 2019 14:58:42 -0700
Subject: [PATCH 1880/2720] Abstract backend for trax so we can swap functions
 (part 1, more to come).

PiperOrigin-RevId: 242545287
---
 tensor2tensor/trax/backend.py        | 69 ++++++++++++++++++++++++++++
 tensor2tensor/trax/inputs.py         | 14 +++---
 tensor2tensor/trax/learning_rate.py  |  3 +-
 tensor2tensor/trax/optimizers.py     |  3 ++
 tensor2tensor/trax/stax/attention.py |  2 +-
 tensor2tensor/trax/stax/losses.py    |  2 +-
 tensor2tensor/trax/stax/slax.py      |  2 +-
 tensor2tensor/trax/stax/stax_base.py | 30 +++++++-----
 tensor2tensor/trax/trax.py           | 37 +++++++--------
 9 files changed, 120 insertions(+), 42 deletions(-)
 create mode 100644 tensor2tensor/trax/backend.py

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
new file mode 100644
index 000000000..c8634cf2c
--- /dev/null
+++ b/tensor2tensor/trax/backend.py
@@ -0,0 +1,69 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Trax backend: all the primitive functions needed."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gin
+
+import jax
+import jax.numpy as jnp
+import jax.scipy.special as jax_special
+import numpy as onp
+
+
+@gin.configurable()
+def backend(name="jax"):
+  if name == "numpy":
+    return _NUMPY_BACKEND
+  return _JAX_BACKEND
+
+
+_JAX_BACKEND = {
+    "np": jnp,
+    "logsumexp": jax_special.logsumexp,
+    "jit": jax.jit,
+    "grad": jax.grad,
+    "pmap": jax.pmap,
+}
+
+
+_NUMPY_BACKEND = {
+    "np": onp,
+    "jit": (lambda f: f),
+}
+
+
+# TODO(lukaszkaiser): make this lazy so we can switch backends on the fly.
+numpy = backend()["np"]
+
+
+def logsumexp(*args, **kwargs):
+  return backend()["logsumexp"](*args, **kwargs)
+
+
+def jit(*args, **kwargs):
+  return backend()["jit"](*args, **kwargs)
+
+
+def grad(*args, **kwargs):
+  return backend()["grad"](*args, **kwargs)
+
+
+def pmap(*args, **kwargs):
+  return backend()["pmap"](*args, **kwargs)
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 6355d7b0e..0890f7e90 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -25,13 +25,11 @@
 
 import gin
 
-import jax.numpy as np
-import numpy as onp
+import numpy as np
 
 import tensorflow as tf
 import tensorflow_datasets as tfds
 
-
 Inputs = collections.namedtuple(
     "_Inputs", ["train_stream", "eval_stream", "input_shape"])
 
@@ -73,8 +71,8 @@ def eval_input_fun():
 
 @gin.configurable()
 def random_inputs(
-    input_shape=gin.REQUIRED, input_dtype=onp.int32, input_range=(0, 255),
-    output_shape=gin.REQUIRED, output_dtype=onp.int32, output_range=(0, 9)):
+    input_shape=gin.REQUIRED, input_dtype=np.int32, input_range=(0, 255),
+    output_shape=gin.REQUIRED, output_dtype=np.int32, output_range=(0, 9)):
   """Make random Inputs for debugging.
 
   Args:
@@ -90,10 +88,10 @@ def random_inputs(
   """
   def random_minibatches():
     """Generate a stream of random mini-batches."""
-    if input_dtype in [onp.float16, onp.float32, onp.float64]:
-      rand = onp.random.uniform
+    if input_dtype in [np.float16, np.float32, np.float64]:
+      rand = np.random.uniform
     else:
-      rand = onp.random.random_integers
+      rand = np.random.random_integers
     while True:
       inp = rand(input_range[0], input_range[1], input_shape)
       inp = inp.astype(input_dtype)
diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index e14f43efb..eb7017f12 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -27,8 +27,7 @@
 from __future__ import print_function
 
 import gin
-
-import jax.numpy as np
+from tensor2tensor.trax.backend import numpy as np
 
 
 # A dictionary to memoize results of the MultifactorSchedule below.
diff --git a/tensor2tensor/trax/optimizers.py b/tensor2tensor/trax/optimizers.py
index f835d1d49..e0c5c9bc9 100644
--- a/tensor2tensor/trax/optimizers.py
+++ b/tensor2tensor/trax/optimizers.py
@@ -39,3 +39,6 @@ def opt_configure(*args, **kwargs):
 exponential_decay = opt_configure(opt.exponential_decay)
 inverse_time_decay = opt_configure(opt.inverse_time_decay)
 piecewise_constant = opt_configure(opt.piecewise_constant)
+
+# Get params
+get_params = opt.get_params
diff --git a/tensor2tensor/trax/stax/attention.py b/tensor2tensor/trax/stax/attention.py
index b771c1076..d00de0ffd 100644
--- a/tensor2tensor/trax/stax/attention.py
+++ b/tensor2tensor/trax/stax/attention.py
@@ -19,10 +19,10 @@
 from __future__ import print_function
 
 from jax import random
-import jax.numpy as np
 import numpy as onp
 import numpy.random as npr
 
+from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.stax import stax_base as stax
 
 
diff --git a/tensor2tensor/trax/stax/losses.py b/tensor2tensor/trax/stax/losses.py
index 68c3112c3..56551266f 100644
--- a/tensor2tensor/trax/stax/losses.py
+++ b/tensor2tensor/trax/stax/losses.py
@@ -19,7 +19,7 @@
 from __future__ import print_function
 
 import gin
-import jax.numpy as np
+from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.stax import slax
 
 
diff --git a/tensor2tensor/trax/stax/slax.py b/tensor2tensor/trax/stax/slax.py
index da4ed7559..c4fc096a8 100644
--- a/tensor2tensor/trax/stax/slax.py
+++ b/tensor2tensor/trax/stax/slax.py
@@ -20,9 +20,9 @@
 
 import inspect
 from absl import logging
-import jax.numpy as np
 from jax.tree_util import register_pytree_node as _register_pytree_node
 
+from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.stax import stax_base as stax
 
 
diff --git a/tensor2tensor/trax/stax/stax_base.py b/tensor2tensor/trax/stax/stax_base.py
index f67a5bfaf..3a13911e6 100644
--- a/tensor2tensor/trax/stax/stax_base.py
+++ b/tensor2tensor/trax/stax/stax_base.py
@@ -26,46 +26,54 @@
 import itertools
 import operator as op
 
-import numpy as onp
-import numpy.random as npr
-from six.moves import reduce
-
 from jax import lax
 from jax import random
-from jax.scipy.special import logsumexp
-import jax.numpy as np
 
+import numpy as onp
+import numpy.random as npr
+from six.moves import reduce
+from tensor2tensor.trax import backend
+from tensor2tensor.trax.backend import numpy as np
 
 # Following the convention used in Keras and tf.layers, we use CamelCase for the
 # names of layer constructors, like Conv and Relu, while using snake_case for
 # other functions, like lax.conv and relu.
 
 
-def relu(x): return np.maximum(x, 0.)
-def softplus(x): return np.logaddexp(x, 0.)
+def relu(x):
+  return np.maximum(x, 0.)
+
+
+def softplus(x):
+  return np.logaddexp(x, 0.)
+
 
 def logsoftmax(x, axis=-1):
   """Apply log softmax to an array of logits, log-normalizing along an axis."""
-  return x - logsumexp(x, axis, keepdims=True)
+  return x - backend.logsumexp(x, axis, keepdims=True)
+
 
 def softmax(x, axis=-1):
   """Apply softmax to an array of logits, exponentiating and normalizing along an axis."""
   unnormalized = np.exp(x - x.max(axis, keepdims=True))
   return unnormalized / unnormalized.sum(axis, keepdims=True)
 
+
 def fastvar(x, axis, keepdims):
   """A fast but less numerically-stable variance calculation than np.var."""
-  return np.mean(x**2, axis, keepdims=keepdims) - np.mean(x, axis, keepdims=keepdims)**2
+  m1 = np.mean(x**2, axis, keepdims=keepdims)
+  m2 = np.mean(x, axis, keepdims=keepdims)**2
+  return m1 - m2
 
 
 # Initializers
-
 def randn(stddev=1e-2, rng=npr):
   """An initializer function for random normal coefficients."""
   def init(shape):
     return rng.normal(size=shape, scale=stddev).astype('float32')
   return init
 
+
 def glorot(out_dim=0, in_dim=1, scale=onp.sqrt(2), rng=npr):
   """An initializer function for random Glorot-scaled coefficients."""
   def init(shape):
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 082d2183a..443936a86 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -32,16 +32,17 @@
 import gin
 
 import jax
-from jax.experimental import optimizers as jax_opt
-import jax.numpy as np
+from jax import random as jax_random
 import numpy
 import six
 
+from tensor2tensor.trax import backend
 from tensor2tensor.trax import history as trax_history
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import jaxboard
 from tensor2tensor.trax import learning_rate as lr
-from tensor2tensor.trax import optimizers as trax_optimizers
+from tensor2tensor.trax import optimizers as trax_opt
+from tensor2tensor.trax.backend import numpy as np
 import tensor2tensor.trax.stax as stax
 
 import tensorflow as tf
@@ -172,7 +173,7 @@ def evaluate(inputs_stream, predict_fun, metric_funs, rng):
   count = 0
   for inp in inputs_stream:
     count += 1
-    rng, subrng = jax.random.split(rng)
+    rng, subrng = jax_random.split(rng)
     preds = predict_fun(inp[0], rng=subrng)
     for m, f in six.iteritems(metric_funs):
       metrics[m] += f(inp, preds)
@@ -201,7 +202,7 @@ def get_random_number_generator_and_set_seed(seed=None):
     seed = random.randint(0, 2**31 - 1)
   tf.set_random_seed(seed)
   numpy.random.seed(seed)
-  return jax.random.PRNGKey(seed)
+  return jax_random.PRNGKey(seed)
 
 
 # TODO(trax):
@@ -244,27 +245,27 @@ def epochs(steps=None, epoch_steps=1):
 def _jit_update_fun(predict_fun, loss_fun, optimizer, lr_fun, num_devices):
   """Get jit-ed update function for loss, optimizer, learning rate function."""
   if num_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
-    @jax.jit
     def single_update(i, opt_state, batch, rng):
       _, opt_update = optimizer(lr_fun)
-      params = jax_opt.get_params(opt_state)
-      return opt_update(i, jax.grad(loss_fun)(
+      params = trax_opt.get_params(opt_state)
+      return opt_update(i, backend.grad(loss_fun)(
           params, batch, predict_fun, rng), opt_state)
-    return single_update
+    return backend.jit(single_update)
 
-  @functools.partial(jax.pmap, axis_name="batch")
   def mapped_update(i, opt_state, batch, rng):
     """This is a multi-device version of the update function above."""
     # We assume all tensors have the first dimension = num_devices.
     _, opt_update = optimizer(lr_fun)
-    params = jax_opt.get_params(opt_state)
-    grads = jax.grad(loss_fun)(params, batch, predict_fun, rng)
+    params = trax_opt.get_params(opt_state)
+    grads = backend.grad(loss_fun)(params, batch, predict_fun, rng)
     grads = jax.tree_util.tree_map(lambda g: jax.lax.psum(g, "batch"), grads)
     return opt_update(i, grads, opt_state)
 
   def update(i, opt_state, batch, rng):
     # TODO(lukaszkaiser): investigate how to replicate rng and correct.
-    return mapped_update(jax.replicate(i), opt_state, batch, jax.replicate(rng))
+    return backend.pmap(mapped_update(
+        jax.replicate(i), opt_state, batch, jax.replicate(rng)),
+                        axis_name="batch")
 
   return update
 
@@ -289,7 +290,7 @@ def reshape_by_device(train_data, num_devices):
 def train(output_dir,
           model=gin.REQUIRED,
           inputs=trax_inputs.inputs,
-          optimizer=trax_optimizers.adam,
+          optimizer=trax_opt.adam,
           lr_schedule=lr.MultifactorSchedule,
           train_steps=1000,
           eval_steps=10,
@@ -345,7 +346,7 @@ def train(output_dir,
     opt_state = jax.replicate(opt_state)
 
   # jit model_predict and update so they're fast
-  jit_model_predict = jax.jit(model_predict)  # for evaluation
+  jit_model_predict = backend.jit(model_predict)  # for evaluation
   jit_update_fun = _jit_update_fun(model_predict, loss, optimizer, lr_fun,
                                    num_devices)
 
@@ -375,7 +376,7 @@ def train(output_dir,
       next_train_batch = next(train_stream)
       if num_devices > 1:  # TODO(lukaszkaiser): use everywhere when possible.
         next_train_batch = reshape_by_device(next_train_batch, num_devices)
-      rng, subrng = jax.random.split(rng)
+      rng, subrng = jax_random.split(rng)
       opt_state = jit_update_fun(step, opt_state, next_train_batch, subrng)
       step += 1
 
@@ -394,9 +395,9 @@ def train(output_dir,
 
     # Evaluate
     if num_devices > 1:   # TODO(lukaszkaiser): remove branch when possible.
-      params = jax_opt.get_params(jax.unreplicate(opt_state))
+      params = trax_opt.get_params(jax.unreplicate(opt_state))
     else:
-      params = jax_opt.get_params(opt_state)
+      params = trax_opt.get_params(opt_state)
     evaluate_train_and_eval(
         step=step,
         inputs=inputs,

From a89ce94e3463198dc705a6618d7b84cd9b4b57d8 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 8 Apr 2019 16:39:44 -0700
Subject: [PATCH 1881/2720] Re-enable scheduled sampling.

PiperOrigin-RevId: 242563990
---
 tensor2tensor/utils/t2t_model.py | 217 +++++++++++++++++++------------
 1 file changed, 133 insertions(+), 84 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index fb336c52a..de79f2b75 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -20,6 +20,7 @@
 
 import collections
 import contextlib
+import copy
 import functools
 import math
 import os
@@ -345,9 +346,25 @@ def body_sharded(self, sharded_features):
                               "and set use_body_sharded to True.")
 
   def model_fn_sharded(self, sharded_features):
+    """Estimator model_fn sharded along batch dimension.
+
+    Args:
+      sharded_features: {str: [Tensor]}. Features sharded along batch dimension.
+        Each list is the same length (== number of shards).
+
+    Returns:
+      sharded_logits: [Tensor]. Logits for each shard of examples.
+      losses: {str: 0-D Tensor}. Loss averaged across shards.
+    """
     dp = self._data_parallelism
+
+    # [{str: Tensor}]. Transpose of 'sharded_features'.
     datashard_to_features = self._to_features_per_datashard(sharded_features)
     if self.use_body_sharded():
+      if  self.hparams.scheduled_sampling_prob > 0.0:
+        raise NotImplementedError(
+            "Scheduled sampling for non-sharded body only.")
+
       # MoE models override body_sharded
       transformed_features = dp(self.bottom, datashard_to_features)
       body_out = self.body_sharded(
@@ -381,6 +398,9 @@ def model_fn_sharded(self, sharded_features):
           losses.update(training_loss_dict)
     else:
       sharded_logits, sharded_losses = dp(self.model_fn, datashard_to_features)
+      sharded_logits, sharded_losses = dp(
+          self.maybe_scheduled_sampling,
+          datashard_to_features, sharded_logits, sharded_losses)
       if isinstance(sharded_logits[0], dict):
         temp_dict = {k: [] for k, _ in six.iteritems(sharded_logits[0])}
         for k, _ in six.iteritems(sharded_logits[0]):
@@ -389,18 +409,6 @@ def model_fn_sharded(self, sharded_features):
         sharded_logits = temp_dict
       losses = average_sharded_losses(sharded_losses)
 
-    # TODO(rsepassi): Reenable scheduled sampling
-    # Disabled because of model_fn_sharded refactor
-    #
-    # do_scheduled_sampling = (  # Only do it if training and set for it.
-    #     self.hparams.scheduled_sampling_prob > 0.0 and
-    #     self.hparams.mode == tf.estimator.ModeKeys.TRAIN)
-    # if do_scheduled_sampling:
-    #   sharded_logits, losses = scheduled_sampling(
-    #       self.hparams, self._problem_hparams, dp,
-    #       sharded_logits, losses, sharded_features,
-    #       transformed_features, self)
-
     return sharded_logits, losses
 
   def model_fn(self, features):
@@ -1767,6 +1775,119 @@ def _summarize_losses(self, losses_dict):
         for loss_name, loss_val in sorted(losses_dict.items()):
           tf.summary.scalar(loss_name, loss_val)
 
+  def maybe_scheduled_sampling(self, features, logits, losses):
+    """Scheduled sampling.
+
+    Performs forward inference again with "targets" feature replaced with values
+    sampled from the model.
+
+    This is the identity unless self.hparams.scheduled_sampling_prob > 0
+    (default).
+
+    **WARNING**: This is not a faithful implementation of scheduled sampling.
+    This implementation samples tokens for timestep t condtioned on gold tokens
+    1...t-1. A proper implementation must condition on a mix of gold and
+    sampled tokens. Doing so is not efficient for models such like Transformer.
+
+    Args:
+      features: {str: Tensor}. Features sharded along batch dimension.
+      logits: Tensor. Logits for each shard of data.
+      losses: 0-D Tensor or (num: 0-D Tensor, denom: 0-D Tensor). Loss Tensor
+
+    Returns:
+      new_logits: Tensor.
+      new_losses: {str: loss} where loss is one of (i) a 0-D Tensor or
+        (ii) a (num: 0-D Tensor, denom: 0-D Tensor) pair to be used in a
+        weighted average.
+    """
+    hparams = self.hparams
+    problem_hparams = self._problem_hparams
+
+    # Only do scheduled sampling if requested.
+    if hparams.scheduled_sampling_prob == 0.0:
+      return (logits, losses)
+
+    # Only do scheduled sampling on language tasks.
+    modality = problem_hparams.modality["targets"]
+    if modality != modalities.ModalityType.SYMBOL:
+      assert hparams.scheduled_sampling_prob == 0, (
+          "Scheduled sampling only applies to ModalityType.SYMBOL. Set "
+          "hparams.scheduled_sampling_prob == 0.0.")
+      return (logits, losses)
+
+    # Only do scheduled sampling when training.
+    is_training = (hparams.mode == tf.estimator.ModeKeys.TRAIN)
+    if not is_training:
+      tf.logging.info("Running in %s mode. Not using scheduled sampling.",
+                      hparams.mode)
+      return (logits, losses)
+
+    # Pad vocabulary if vocab size must be evenly divisible by vocab_divisor.
+    vocab_size = problem_hparams.vocab_size["targets"]
+    assert vocab_size is not None
+    assert hparams.vocab_divisor == 1
+
+    def sample(x):
+      """Multinomial sampling from a n-dimensional tensor."""
+      samples = tf.multinomial(tf.reshape(x, [-1, vocab_size]), 1)
+      reshaped_samples = tf.reshape(samples, common_layers.shape_list(x)[:-1])
+      return tf.to_int32(reshaped_samples)
+
+    def mix_gold_sampled(gold_targets, sampled_targets, mixin_prob):
+      """Interleave sampled and gold tokens randomly."""
+      return tf.where(
+          tf.less(
+              tf.random_uniform(common_layers.shape_list(sampled_targets)),
+              mixin_prob),
+          sampled_targets,
+          gold_targets)
+
+    def sampled_results(mixin_prob):
+      """Generate scheduled sampling results."""
+      sampled_targets = sample(logits)
+      new_targets = mix_gold_sampled(features["targets"],
+                                     sampled_targets,
+                                     mixin_prob)
+      new_targets = tf.stop_gradient(new_targets)  # Treat new_targets as given.
+      new_features = copy.copy(features)
+      new_features["targets"] = new_targets
+      with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+        # Compute bottom() for new_targets.
+        #
+        # TODO(duckworthd): Only apply bottom to 'new_targets'.
+        new_transformed_features = self.bottom(new_features)
+
+        # Compute body.
+        with tf.variable_scope("body"):
+          new_body_outputs, new_losses = self._normalize_body_output(
+              self.body(new_transformed_features))
+        assert "training" not in new_losses
+
+        # Compute top.
+        new_logits = self.top(new_body_outputs, new_features)
+
+        # Compute loss. Use original features (== labels).
+        if (hparams.mode != tf.estimator.ModeKeys.PREDICT and
+            hparams.mode != "attack"):
+          new_losses["training"] = self.loss(new_logits, features)
+        else:
+          new_losses["training"] = 0.0
+
+      return new_logits, new_losses
+
+    tf.logging.info("Using scheduled sampling.")
+    assert hparams.scheduled_sampling_prob == 1.0, (
+        "hparams.scheduled_sampling_prob must be 0 or 1.")
+    # Gradually increase over a warmup period. Lower numbers mean more gold
+    # tokens.
+    mixin_prob = (
+        hparams.scheduled_sampling_gold_mixin_prob *
+        common_layers.inverse_exp_decay(
+            hparams.scheduled_sampling_warmup_steps,
+            min_value=0.001)
+    )
+    return sampled_results(mixin_prob)
+
 
 def _with_timing(fn, msg, silent=False):
 
@@ -1985,78 +2106,6 @@ def create_eager_var_store():
     return DummyVariableStore()
 
 
-def scheduled_sampling(hparams, problem_hparams, dp, sharded_logits, losses,
-                       sharded_features, transformed_features, model):
-  """Scheduled sampling."""
-  modality = problem_hparams.modality["targets"]
-  vocab_size = problem_hparams.vocab_size["targets"]
-  if vocab_size is not None and hasattr(hparams, "vocab_divisor"):
-    vocab_size += (-vocab_size) % hparams.vocab_divisor
-
-  def sample(x):
-    """Multinomial sampling from a n-dimensional tensor."""
-    samples = tf.multinomial(tf.reshape(x, [-1, vocab_size]), 1)
-    reshaped_samples = tf.reshape(samples, common_layers.shape_list(x)[:-1])
-    return tf.to_int32(reshaped_samples)
-
-  def mix_gold_sampled(gold_targets, sampled_targets):
-    return tf.where(
-        tf.less(
-            tf.random_uniform(common_layers.shape_list(sampled_targets)),
-            hparams.scheduled_sampling_gold_mixin_prob), gold_targets,
-        sampled_targets)
-
-  def sampled_results():
-    """Generate scheduled sampling results."""
-    sampled_targets = dp(sample, sharded_logits)
-    new_targets = dp(mix_gold_sampled, sharded_features["targets"],
-                     sampled_targets)
-    new_features = transformed_features
-    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
-      modality_name = hparams.name.get(
-          "targets",
-          modalities.get_name(modality))(hparams, vocab_size)
-      with tf.variable_scope(modality_name):
-        bottom = hparams.bottom.get(
-            "targets", modalities.get_targets_bottom(modality))
-        new_features["targets"] = dp(bottom, new_targets, hparams, vocab_size)
-      with tf.variable_scope("body"):
-        body_outputs, losses = model.model_fn_sharded(new_features)
-        if not isinstance(losses, dict):  # If it's a single extra loss.
-          losses = {"extra": losses}
-      with tf.variable_scope(modality_name):
-        top = hparams.top.get("targets", modalities.get_top(modality))
-        new_sharded_logits = dp(top,
-                                body_outputs,
-                                sharded_features["targets"],
-                                hparams,
-                                vocab_size)
-        if "training" not in losses:
-          loss = hparams.loss.get("targets", modalities.get_loss(modality))
-          weights_fn = hparams.weights_fn.get(
-              "targets", modalities.get_weights_fn(modality))
-          sharded_loss_num, sharded_loss_den = dp(loss,
-                                                  sharded_logits,
-                                                  sharded_features["targets"],
-                                                  hparams,
-                                                  vocab_size,
-                                                  weights_fn=weights_fn)
-          training_loss = (tf.add_n(sharded_loss_num) /
-                           tf.maximum(1.0, tf.add_n(sharded_loss_den)))
-          training_loss *= problem_hparams.loss_multiplier
-          losses["training"] = training_loss
-    return new_sharded_logits, losses
-
-  # Run the above conditionally.
-  prob = hparams.scheduled_sampling_prob
-  prob *= common_layers.inverse_exp_decay(
-      hparams.scheduled_sampling_warmup_steps, min_value=0.001)
-  sharded_logits, losses = tf.cond(
-      tf.less(tf.random_uniform([]), prob), sampled_results,
-      lambda: (sharded_logits, losses))
-  return sharded_logits, losses
-
-
 def average_sharded_losses(sharded_losses):
   """Average losses across datashards.
 

From 22ecfab75cfe306978bbc279da71afc576287f93 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 8 Apr 2019 16:59:40 -0700
Subject: [PATCH 1882/2720] Add CIFAR-10 and WideResnet. Changed how
 preprocessing is configured, I believe it would be cleaner moving forward to
 allow the user to write an entire preprocess function for their task.

PiperOrigin-RevId: 242567423
---
 .../trax/configs/resnet50_imagenet_8gb.gin    |  4 +-
 .../trax/configs/transformer_lm1b_8gb.gin     |  3 +-
 .../trax/configs/wide_resnet_cifar10_8gb.gin  | 47 +++++++++++++++++++
 tensor2tensor/trax/inputs.py                  | 46 +++++++++++++-----
 tensor2tensor/trax/models/__init__.py         |  1 +
 tensor2tensor/trax/models/resnet.py           | 40 ++++++++++++++++
 6 files changed, 126 insertions(+), 15 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin

diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
index 1f58f1f58..01cac65ac 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
@@ -26,9 +26,6 @@ MultifactorSchedule.warmup_steps = 400
 # ==============================================================================
 momentum.mass = 0.9
 
-# Parameters for preprocess_fun:
-# ==============================================================================
-preprocess_fun.max_target_length = -1
 
 # Parameters for Resnet50:
 # ==============================================================================
@@ -44,3 +41,4 @@ train.model = @trax.models.Resnet50
 train.optimizer = @trax.optimizers.momentum
 train.train_steps = 1000000
 train.lr_schedule = @learning_rate.EvalAdjustingSchedule
+
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index 0d07d1083..3aeeb77ca 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -25,7 +25,8 @@ MultifactorSchedule.warmup_steps = 8000
 
 # Parameters for preprocess_fun:
 # ==============================================================================
-preprocess_fun.max_target_length = 511
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
+lm1b.max_target_length = 511
 
 # Parameters for train:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
new file mode 100644
index 000000000..c2e8a605f
--- /dev/null
+++ b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
@@ -0,0 +1,47 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.learning_rate
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size = 32
+batch_fun.bucket_length = 32
+batch_fun.buckets = None
+batch_fun.eval_batch_size = 32
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 'cifar10'
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+EvalAdjustingSchedule.constant = 1.0
+MultifactorSchedule.factors = 'constant * linear_warmup'
+MultifactorSchedule.warmup_steps = 400
+
+# Parameters for momentum:
+# ==============================================================================
+momentum.mass = 0.9
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.cifar10_no_augmentation_preprocess
+
+# Parameters for WideResnet:
+# ==============================================================================
+WideResnet.num_blocks = 3
+WideResnet.hidden_size = 64
+WideResnet.num_output_classes = 10
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 2000
+train.eval_steps = 20
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.WideResnet
+train.optimizer = @trax.optimizers.momentum
+train.train_steps = 1000000
+train.lr_schedule = @learning_rate.EvalAdjustingSchedule
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 0890f7e90..bd2a7c144 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -210,15 +210,6 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
   return train_dataset, eval_dataset, info, supervised_keys
 
 
-@gin.configurable(blacklist=["dataset", "training"])
-def preprocess_fun(dataset, training, max_target_length=-1):
-  def target_right_length(_, target):
-    return tf.less(tf.shape(target)[0], max_target_length + 1)
-  if max_target_length > 0 and training:
-    dataset = dataset.filter(target_right_length)
-  return dataset
-
-
 @gin.configurable(blacklist=["dataset", "training", "shapes", "target_names"])
 def batch_fun(dataset, training, shapes, target_names,
               batch_size=32, eval_batch_size=32,
@@ -262,7 +253,40 @@ def example_length(_, target):
   return dataset
 
 
-def shuffle_and_batch_data(dataset, target_names, features_info, training):
+# pylint: disable=unused-argument
+@gin.configurable(blacklist=["dataset", "training"])
+def cifar10_no_augmentation_preprocess(dataset, training):
+
+  def cast_image(features, targets):
+    features["image"] = tf.cast(features["image"], tf.float32) / 255.0
+    return features, targets
+
+  dataset = dataset.map(cast_image)
+  return dataset
+
+
+# pylint: disable=unused-argument
+def no_preprocess(dataset, training):
+  return dataset
+
+
+@gin.configurable(blacklist=["dataset", "training"])
+def lm1b_preprocess(dataset, training, max_target_length=-1):
+
+  def target_right_length(_, target):
+    return tf.less(tf.shape(target)[0], max_target_length + 1)
+
+  if max_target_length > 0 and training:
+    dataset = dataset.filter(target_right_length)
+  return dataset
+
+
+@gin.configurable(whitelist=["preprocess_fun"])
+def shuffle_and_batch_data(dataset,
+                           target_names,
+                           features_info,
+                           training,
+                           preprocess_fun=no_preprocess):
   """Shuffle and batch the given dataset."""
   def append_targets(example):
     """Append targets to the example dictionary. Needed for Keras."""
@@ -279,10 +303,10 @@ def append_targets(example):
     # essential for synchronous highly-parallel training to avoid multiple
     # replicas reading the same data in lock-step.
     dataset = dataset.skip(random.randint(0, _MAX_SKIP_EXAMPLES))
+  dataset = preprocess_fun(dataset, training)
   shapes = {k: features_info[k].shape for k in features_info}
   shapes = (shapes, shapes[target_names[0]])
   dataset = dataset.shuffle(1024)
-  dataset = preprocess_fun(dataset, training)
   dataset = batch_fun(dataset, training, shapes, target_names)
   return dataset.prefetch(2)
 
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index ceaf732bf..bde69f33d 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -34,4 +34,5 @@ def model_configure(*args, **kwargs):
 # pylint: disable=invalid-name
 MLP = model_configure(mlp.MLP)
 Resnet50 = model_configure(resnet.Resnet50)
+WideResnet = model_configure(resnet.WideResnet)
 TransformerLM = model_configure(transformer.TransformerLM)
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index 9de5c4232..b731861b7 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -92,3 +92,43 @@ def Resnet50(hidden_size=64, num_output_classes=1001):
       IdentityBlock(3, [8 * hidden_size, 8 * hidden_size]),
       stax.AvgPool((7, 7)), stax.Flatten,
       stax.Dense(num_output_classes), stax.LogSoftmax)
+
+
+def WideResnetBlock(channels, strides=(1, 1), channel_mismatch=False):
+  """WideResnet convolutational block."""
+  main = stax.serial(stax.BatchNorm(), stax.Relu,
+                     stax.Conv(channels, (3, 3), strides, padding='SAME'),
+                     stax.BatchNorm(), stax.Relu,
+                     stax.Conv(channels, (3, 3), padding='SAME'))
+  shortcut = stax.Identity if not channel_mismatch else stax.Conv(
+      channels, (3, 3), strides, padding='SAME')
+  return stax.serial(
+      stax.FanOut(2), stax.parallel(main, shortcut), stax.FanInSum)
+
+
+def WideResnetGroup(n, channels, strides=(1, 1)):
+  blocks = []
+  blocks += [WideResnetBlock(channels, strides, channel_mismatch=True)]
+  for _ in range(n - 1):
+    blocks += [WideResnetBlock(channels, (1, 1))]
+  return stax.serial(*blocks)
+
+
+def WideResnet(num_blocks=3, hidden_size=64, num_output_classes=10):
+  """WideResnet from https://arxiv.org/pdf/1605.07146.pdf.
+
+  Args:
+    num_blocks: int, number of blocks in a group.
+    hidden_size: the size of the first hidden layer (multiplied later).
+    num_output_classes: int, number of classes to distinguish.
+
+  Returns:
+    The WideResnet model with given layer and output sizes.
+  """
+  return stax.serial(
+      stax.Conv(hidden_size, (3, 3), padding='SAME'),
+      WideResnetGroup(num_blocks, hidden_size),
+      WideResnetGroup(num_blocks, hidden_size * 2, (2, 2)),
+      WideResnetGroup(num_blocks, hidden_size * 4, (2, 2)), stax.BatchNorm(),
+      stax.Relu, stax.AvgPool((8, 8)), stax.Flatten,
+      stax.Dense(num_output_classes), stax.LogSoftmax)

From 13c82c63f2d192ffd7b3a24c3c856afe1b61d63d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 9 Apr 2019 07:25:02 -0700
Subject: [PATCH 1883/2720] * Factored out main from ppo.py into ppo_main.py *
 Testing the training_loop in ppo_training_loop_test.py * Testing the other
 functions in ppo_test.py   * tests will slowly be added for the rest of the
 functions.

Many TODO(s):
* Change the training loops to be trax.train invocations.
* Make things gin configurable:
  * Default params for gamma, epsilon etc.
  * But mostly make policy and value functions configurable too.
* Unify value and policy function.
* Batchify value loss computation (currently uses vmap).

PiperOrigin-RevId: 242662317
---
 tensor2tensor/trax/rlax/ppo.py                | 361 ++++++++----------
 tensor2tensor/trax/rlax/ppo_main.py           |  40 ++
 tensor2tensor/trax/rlax/ppo_test.py           |  42 ++
 .../trax/rlax/ppo_training_loop_test.py       |  47 +++
 4 files changed, 285 insertions(+), 205 deletions(-)
 create mode 100644 tensor2tensor/trax/rlax/ppo_main.py
 create mode 100644 tensor2tensor/trax/rlax/ppo_test.py
 create mode 100644 tensor2tensor/trax/rlax/ppo_training_loop_test.py

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 87cb1935c..aa73dd3ad 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -21,21 +21,16 @@
 
 import functools
 import time
-from absl import app
+from absl import logging
 import gym
-
 from jax import grad
 from jax import jit
 from jax import lax
 from jax import numpy as np
 from jax import vmap
 from jax.experimental import optimizers
-from jax.experimental import stax
-from jax.experimental.stax import Dense
-from jax.experimental.stax import Relu
-from jax.experimental.stax import Softmax
-
 import numpy as onp
+from tensor2tensor.trax.stax import stax_base as stax
 
 DEBUG_LOGGING = False
 GAMMA = 0.99
@@ -51,27 +46,25 @@
 def initialize_policy_and_value_nets(num_actions, batch_observations_shape):
   """Setup and initialize the policy and value networks."""
   policy_net_init, policy_net_apply = stax.serial(
-      Dense(16),
-      Relu,
-      Dense(4),
-      Relu,
-      Dense(num_actions),
-      Softmax,
+      stax.Dense(16),
+      stax.Relu,
+      stax.Dense(4),
+      stax.Relu,
+      stax.Dense(num_actions),
+      stax.Softmax,
   )
 
-  _, policy_net_params = policy_net_init(
-      batch_observations_shape)
+  _, policy_net_params = policy_net_init(batch_observations_shape)
 
   value_net_init, value_net_apply = stax.serial(
-      Dense(16),
-      Relu,
-      Dense(4),
-      Relu,
-      Dense(1),  # 1 since we want to predict reward using value network.
+      stax.Dense(16),
+      stax.Relu,
+      stax.Dense(4),
+      stax.Relu,
+      stax.Dense(1),  # 1 since we want to predict reward using value network.
   )
 
-  _, value_net_params = value_net_init(
-      batch_observations_shape)
+  _, value_net_params = value_net_init(batch_observations_shape)
 
   return ((policy_net_params, policy_net_apply), (value_net_params,
                                                   value_net_apply))
@@ -95,6 +88,7 @@ def initialize_optimizers(policy_net_params, value_net_params):
 # Should this be collect 'n' trajectories, or
 # Run the env for 'n' steps and take completed trajectories, or
 # Any other option?
+# TODO(afrozm): Replace this with EnvProblem?
 def collect_trajectories(env,
                          policy_net_apply,
                          policy_net_params,
@@ -129,9 +123,6 @@ def collect_trajectories(env,
       elif policy == "categorical-sampling":
         action = int(onp.argwhere(onp.random.multinomial(1, predictions) == 1))
 
-      if DEBUG_LOGGING:
-        print("With predictions: ", predictions, " chose action: ", action)
-
       # NOTE: Assumption, single batch.
       action = int(action)
 
@@ -143,8 +134,8 @@ def collect_trajectories(env,
 
     # This means we are done
     assert done
-    trajectories.append((np.stack(observations), np.stack(actions),
-                         np.stack(rewards)))
+    trajectories.append(
+        (np.stack(observations), np.stack(actions), np.stack(rewards)))
 
   return trajectories
 
@@ -164,14 +155,10 @@ def pad_trajectories(trajectories, boundary=10):
 
   # Given the above, let's compute max(T) over all trajectories.
   t_max = max(o.shape[0] for (o, a, r) in trajectories)
-  if DEBUG_LOGGING:
-    print("t_max: %s" % t_max)
 
   # t_max - 1 is rounded to the next multiple of `boundary`
   boundary = int(boundary)
   bucket_length = boundary * int(np.ceil(float(t_max - 1) / boundary))
-  if DEBUG_LOGGING:
-    print("bucket_length: %s" % bucket_length)
 
   # So all obs will be padded to t_max and actions and rewards to t_max - 1.
   padded_observations = []
@@ -191,33 +178,21 @@ def pad_trajectories(trajectories, boundary=10):
       continue
 
     # First pad observations.
-    if DEBUG_LOGGING:
-      print("num_to_pad: %s" % num_to_pad)
     padding_config = [(0, num_to_pad, 0)]
     for _ in range(o.ndim - 1):
       padding_config.append((0, 0, 0))
     padding_config = tuple(padding_config)
-    if DEBUG_LOGGING:
-      print("padding_config: %s" % str(padding_config))
     padding_value = 0.0 if o.dtype == np.float32 else 0
-    if DEBUG_LOGGING:
-      print("padding_value: %s" % padding_value)
     padded_obs = lax.pad(o, padding_value, padding_config)
     padded_observations.append(padded_obs)
 
     # Now pad actions and rewards.
     assert a.ndim == 1 and r.ndim == 1
     padding_config = ((0, num_to_pad, 0),)
-    if DEBUG_LOGGING:
-      print("action/reward padding_config: %s" % str(padding_config))
     action_padding_value = 0.0 if a.dtype == np.float32 else 0
     reward_padding_value = 0.0 if r.dtype == np.float32 else 0
-    if DEBUG_LOGGING:
-      print("action_padding_value: %s" % action_padding_value)
     padded_action = lax.pad(a, action_padding_value, padding_config)
     padded_actions.append(padded_action)
-    if DEBUG_LOGGING:
-      print("reward_padding_value: %s" % reward_padding_value)
     padded_reward = lax.pad(r, reward_padding_value, padding_config)
     padded_rewards.append(padded_reward)
 
@@ -229,26 +204,27 @@ def pad_trajectories(trajectories, boundary=10):
       padded_observations), np.stack(padded_actions), np.stack(padded_rewards)
 
 
-def rewards_to_go_discounted(rewards, reward_mask=1.0, gamma=0.99):
+def rewards_to_go(rewards, reward_mask=1.0, gamma=0.99):
   r"""r2g[t] = \sum_{l=0}^{\infty}(\gamma^l * r_{t+l})."""
   time_steps = len(rewards)
   # r2g[t] = r[t] + (gamma * r2g[t+1])
 
   # First initialize like:
   # r2g[t] = r[t], for t = 0 to T-1
-  rewards_to_go = list(rewards)
+  r2g = list(rewards)
 
   # Then add the discounted version of the next time-step.
   # i = [T-2 .. 0]
   for i in range(time_steps - 2, -1, -1):
-    rewards_to_go[i] += gamma * rewards_to_go[i + 1]
+    r2g[i] += gamma * r2g[i + 1]
 
   # Makes this back into JAX's DeviceArray
-  rewards_to_go = np.stack(list(rewards_to_go))
+  r2g = np.stack(list(r2g))
 
-  return rewards_to_go * reward_mask
+  return r2g * reward_mask
 
 
+@functools.partial(jit, static_argnums=(0,))
 def batched_avg_value_function_loss(value_net_apply,
                                     value_net_params,
                                     observations,
@@ -256,6 +232,7 @@ def batched_avg_value_function_loss(value_net_apply,
                                     reward_mask=1.0,
                                     gamma=0.99):
   """L2 loss on the value function's outputs."""
+
   # Capturing the value_net_apply from the parent function's scope.
   # See: https://github.com/google/jax/issues/183
   def _value_function_loss_trajectory(value_net_params,
@@ -264,7 +241,7 @@ def _value_function_loss_trajectory(value_net_params,
                                       reward_mask=1.0,
                                       gamma=0.99):
     """Compute the actual loss for a trajectory."""
-    r2g = rewards_to_go_discounted(
+    r2g = rewards_to_go(
         rewards, reward_mask=reward_mask, gamma=gamma)
     v = value_net_apply(value_net_params, observations[:-1])
     v = np.squeeze(v) * reward_mask
@@ -293,8 +270,11 @@ def batched_deltas(predicted_values, rewards, reward_mask, gamma=0.99):
   return np.array(deltas).T * reward_mask
 
 
-def batched_gae_advantages(deltas, reward_mask, lamda=0.95,  # NOTYPO
-                           gamma=0.99):
+def batched_gae_advantages(
+    deltas,
+    reward_mask,
+    lamda=0.95,  # NOTYPO
+    gamma=0.99):
   r"""A_t = \sum_{l=0}^{\infty}(\gamma * \lambda)^{l}(\delta_{t+l})."""
   _, T = deltas.shape  # pylint: disable=invalid-name
   gl = lamda * gamma  # NOTYPO
@@ -338,22 +318,12 @@ def batched_probab_ratios(policy_net_apply, old_policy_params,
   bp_old = batched_probabs(p_old, actions)
   bp_new = batched_probabs(p_new, actions)
 
-  if DEBUG_LOGGING:
-    print("bp_old: ", bp_old)
-    print("bp_new: ", bp_new)
-
   # Add a small number to bp_old, where reward_mask is 0, this is just to help
   # never to divide by 0.
   bp_old = bp_old + (0.1 * np.abs(reward_mask - 1))
 
-  if DEBUG_LOGGING:
-    print("masked bp_old: ", bp_old)
-
   ret_val = (bp_new * reward_mask) / bp_old
 
-  if DEBUG_LOGGING:
-    print("ret_val: ", ret_val)
-
   return ret_val
 
 
@@ -367,18 +337,20 @@ def batched_clipped_objective(bpr, adv, reward_mask, epsilon=0.2):
   return np.minimum(c1, c2)
 
 
-def batched_ppo_loss(policy_net_apply,
-                     new_policy_params,
-                     old_policy_params,
-                     value_net_apply,
-                     value_net_params,
-                     padded_observations,
-                     padded_actions,
-                     padded_rewards,
-                     reward_mask,
-                     gamma=0.99,
-                     lamda=0.95,  # NOTYPO
-                     epsilon=0.2):
+@functools.partial(jit, static_argnums=(0, 3))
+def batched_ppo_loss(
+    policy_net_apply,
+    new_policy_params,
+    old_policy_params,
+    value_net_apply,
+    value_net_params,
+    padded_observations,
+    padded_actions,
+    padded_rewards,
+    reward_mask,
+    gamma=0.99,
+    lamda=0.95,  # NOTYPO
+    epsilon=0.2):
   """PPO objective, with an eventual minus sign."""
   # V(s_t) forall s & t
   value_function = np.squeeze(
@@ -390,26 +362,27 @@ def batched_ppo_loss(policy_net_apply,
   ratios = batched_probab_ratios(policy_net_apply, old_policy_params,
                                  new_policy_params, padded_observations,
                                  padded_actions, reward_mask)
-  clipped_loss = batched_clipped_objective(ratios, advantages, reward_mask,
-                                           epsilon=epsilon)
+  clipped_loss = batched_clipped_objective(
+      ratios, advantages, reward_mask, epsilon=epsilon)
   return -np.sum(clipped_loss)
 
 
 @functools.partial(jit, static_argnums=(2, 3, 5))
-def ppo_opt_step(i,
-                 opt_state,
-                 ppo_opt_update,
-                 policy_net_apply,
-                 old_policy_params,
-                 value_net_apply,
-                 value_net_params,
-                 padded_observations,
-                 padded_actions,
-                 padded_rewards,
-                 reward_mask,
-                 gamma=0.99,
-                 lamda=0.95,  # NOTYPO
-                 epsilon=0.1):
+def ppo_opt_step(
+    i,
+    opt_state,
+    ppo_opt_update,
+    policy_net_apply,
+    old_policy_params,
+    value_net_apply,
+    value_net_params,
+    padded_observations,
+    padded_actions,
+    padded_rewards,
+    reward_mask,
+    gamma=0.99,
+    lamda=0.95,  # NOTYPO
+    epsilon=0.1):
   """PPO optimizer step."""
   new_policy_params = optimizers.get_params(opt_state)
   g = grad(
@@ -450,37 +423,33 @@ def value_opt_step(i,
   return opt_update(i, g, opt_state)
 
 
-def main(unused_argv):
-  onp.random.seed(0)
+def get_time(t1, t2=None):
+  if t2 is None:
+    t2 = time.time()
+  return round((t2 - t1) * 1000, 2)
+
+
+def training_loop(env=None,
+                  env_name="CartPole-v0",
+                  epochs=EPOCHS,
+                  batch_size=BATCH_TRAJECTORIES,
+                  num_optimizer_steps=NUM_OPTIMIZER_STEPS,
+                  print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
+                  random_seed=None):
+  """Runs the training loop for PPO, with fixed policy and value nets."""
+  onp.random.seed(random_seed)
 
   value_losses = []
   ppo_objective = []
   average_rewards = []
 
-  env = gym.make("CartPole-v0")
-
-  print("Initial observation: ", env.reset())
-
-  for i in range(100):
-    random_action = env.action_space.sample()
-    obs, rew, done, _ = env.step(random_action)
-    print("[%s] reward [%s], done [%s] and obs [%s]" % (i, rew, done, obs))
-    if done:
-      print("Done, so exiting, step: ", i)
-      env.close()
-      break
-
-  print("action_space.shape", env.action_space.shape)
-  print("observation_space.shape", env.observation_space.shape)
+  env = env if env is not None else gym.make(env_name)
 
   batch_observations_shape = (-1,) + env.observation_space.shape
 
   assert isinstance(env.action_space, gym.spaces.Discrete)
   num_actions = env.action_space.n
 
-  print("batch_observations_shape: ", batch_observations_shape)
-  print("num_actions: ", num_actions)
-
   ((policy_net_params, policy_net_apply),
    (value_net_params, value_net_apply)) = initialize_policy_and_value_nets(
        num_actions, batch_observations_shape)
@@ -489,82 +458,71 @@ def main(unused_argv):
                                     value_opt_update) = initialize_optimizers(
                                         policy_net_params, value_net_params)
 
-  for i in range(EPOCHS):
+  for i in range(epochs):
     t = time.time()
     t0 = t
     trajs = collect_trajectories(
         env,
         policy_net_apply,
         policy_net_params,
-        num_trajectories=BATCH_TRAJECTORIES,
+        num_trajectories=batch_size,
         policy=POLICY,
         epsilon=(10.0 / (i + 10.0)))  # this is a different epsilon.
-    print("Took ", round((time.time() - t) * 1000, 2),
-          "msec to collect trajectories.")
 
-    print("Average Trajectory size: ",
-          float(sum(len(traj[0]) for traj in trajs)) / len(trajs))
     avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
     average_rewards.append(avg_reward)
-    print("Average sum rewards: ", avg_reward)
 
-    if (avg_reward > 190.0) and (i % 5 == 0):
-      print("policy_net_params:\n", policy_net_params)
-      print("value_net_params:\n", value_net_params)
+    logging.debug("Average sum rewards [%0.2f]", avg_reward)
+    logging.debug("Collecting trajectories took %0.2f msec.", get_time(t))
+    logging.debug("Average Trajectory size [%0.2f]",
+                  float(sum(len(traj[0]) for traj in trajs)) / len(trajs))
 
     t = time.time()
     (_, reward_mask, padded_observations, padded_actions,
-     padded_rewards) = pad_trajectories(trajs, boundary=20)
-    print("Took ", round((time.time() - t) * 1000, 2),
-          "msec to pad trajectories.")
+     padded_rewards) = pad_trajectories(
+         trajs, boundary=20)
 
-    print("Padded Observations' shape: ", padded_observations.shape)
-    print("Padded Actions' shape:      ", padded_actions.shape)
-    print("Padded Rewards' shape:      ", padded_rewards.shape)
+    logging.debug("Padding trajectories took %0.2f msec.", get_time(t))
+    logging.debug("Padded Actions' shape [%s]", str(padded_actions.shape))
 
     # Linear annealing from 0.1 to 0.0
-    epsilon = 0.1 if EPOCHS == 1 else 0.1 * (1.0 - (i / (EPOCHS - 1)))
+    epsilon = 0.1 if epochs == 1 else 0.1 * (1.0 - (i / (epochs - 1)))
 
     t = time.time()
-    val_loss = jit(
-        batched_avg_value_function_loss, static_argnums=(0,))(
-            value_net_apply,
-            value_net_params,
-            padded_observations,
-            padded_rewards,
-            reward_mask,
-            gamma=GAMMA)
-
-    print("Took ", round((time.time() - t) * 1000, 2),
-          "msec to calculate value loss = ", val_loss)
+    val_loss = batched_avg_value_function_loss(
+        value_net_apply,
+        value_net_params,
+        padded_observations,
+        padded_rewards,
+        reward_mask,
+        gamma=GAMMA)
+
+    logging.debug("Calculating value loss took %0.2f msec.", get_time(t))
     value_losses.append(val_loss)
 
     t = time.time()
-    ppo_loss = jit(
-        batched_ppo_loss, static_argnums=(0,
-                                          3))(policy_net_apply,
-                                              policy_net_params,
-                                              policy_net_params,
-                                              value_net_apply,
-                                              value_net_params,
-                                              padded_observations,
-                                              padded_actions,
-                                              padded_rewards,
-                                              reward_mask,
-                                              gamma=GAMMA,
-                                              lamda=LAMBDA,  # NOTYPO
-                                              epsilon=epsilon)
+    ppo_loss = batched_ppo_loss(
+        policy_net_apply,
+        policy_net_params,
+        policy_net_params,
+        value_net_apply,
+        value_net_params,
+        padded_observations,
+        padded_actions,
+        padded_rewards,
+        reward_mask,
+        gamma=GAMMA,
+        lamda=LAMBDA,  # NOTYPO
+        epsilon=epsilon)
     # ppo_loss = 11.00110011
-    print("Took ", round((time.time() - t) * 1000, 2),
-          "msec to calculate ppo loss = ", ppo_loss)
+    logging.debug("Calculating PPO loss took %0.2f msec.", get_time(t))
     ppo_objective.append(-ppo_loss)
 
     # Run optimizers.
+    logging.debug("PPO Optimization")
     t1 = time.time()
 
-    print("PPO objective optimization.")
-
-    for j in range(NUM_OPTIMIZER_STEPS):
+    for j in range(num_optimizer_steps):
       t = time.time()
       # Update the optimizer state.
       ppo_opt_state = ppo_opt_step(
@@ -586,34 +544,31 @@ def main(unused_argv):
       # Get the new params.
       new_policy_net_params = optimizers.get_params(ppo_opt_state)
       if ((j + 1) %
-          PRINT_EVERY_OPTIMIZER_STEP == 0) or (j == NUM_OPTIMIZER_STEPS - 1):
-        new_ppo_loss = jit(
-            batched_ppo_loss, static_argnums=(0,
-                                              3))(policy_net_apply,
-                                                  new_policy_net_params,
-                                                  policy_net_params,
-                                                  value_net_apply,
-                                                  value_net_params,
-                                                  padded_observations,
-                                                  padded_actions,
-                                                  padded_rewards,
-                                                  reward_mask,
-                                                  gamma=GAMMA,
-                                                  lamda=LAMBDA,  # NOTYPO
-                                                  epsilon=epsilon)
-        print("Took ", round((t2 - t) * 1000, 2),
-              "msec to do one step ppo grad desc")
-        print("New ppo loss[", j, "]: ", new_ppo_loss, " vs old ppo loss: ",
-              ppo_loss)
+          print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
+        new_ppo_loss = batched_ppo_loss(
+            policy_net_apply,
+            new_policy_net_params,
+            policy_net_params,
+            value_net_apply,
+            value_net_params,
+            padded_observations,
+            padded_actions,
+            padded_rewards,
+            reward_mask,
+            gamma=GAMMA,
+            lamda=LAMBDA,  # NOTYPO
+            epsilon=epsilon)
+        logging.debug("One PPO grad desc took: %0.2f msec", get_time(t, t2))
+        logging.debug("PPO loss [%10.2f] -> [%10.2f]", ppo_loss, new_ppo_loss)
       # Update the params.
       policy_net_params = new_policy_net_params
 
-    print("Total ppo loss reduction: ",
-          100 * (ppo_loss - new_ppo_loss) / np.abs(ppo_loss), "%")
+    logging.debug("Total PPO loss reduction [%0.2f]%%",
+                  (100 * (ppo_loss - new_ppo_loss) / np.abs(ppo_loss)))
 
-    print("Value optimization.")
+    logging.debug("Value Optimization")
 
-    for j in range(NUM_OPTIMIZER_STEPS):
+    for j in range(num_optimizer_steps):
       t = time.time()
       value_opt_state = value_opt_step(
           j,
@@ -627,36 +582,32 @@ def main(unused_argv):
       t2 = time.time()
       value_net_params = optimizers.get_params(value_opt_state)
       if ((j + 1) %
-          PRINT_EVERY_OPTIMIZER_STEP == 0) or (j == NUM_OPTIMIZER_STEPS - 1):
-        new_val_loss = jit(
-            batched_avg_value_function_loss, static_argnums=(0,))(
-                value_net_apply,
-                value_net_params,
-                padded_observations,
-                padded_rewards,
-                reward_mask,
-                gamma=GAMMA)
-        print("Took ", round((t2 - t) * 1000, 2),
-              "msec to do one step value grad desc")
-        print("New value loss[", j, "]: ", new_val_loss, " vs old value loss: ",
-              val_loss)
-    print("Total value loss reduction: ",
-          100 * (val_loss - new_val_loss) / val_loss, "%")
-
-    print("Took ", round((time.time() - t1) * 1000, 2), "msec to do grad desc")
+          print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
+        new_val_loss = batched_avg_value_function_loss(
+            value_net_apply,
+            value_net_params,
+            padded_observations,
+            padded_rewards,
+            reward_mask,
+            gamma=GAMMA)
+        logging.debug("One value grad desc took: %0.2f msec", get_time(t, t2))
+        logging.debug("Value loss [%10.2f] -> [%10.2f]", val_loss, new_val_loss)
+    logging.debug("Total value loss reduction [%0.2f]%%",
+                  (100 * (val_loss - new_val_loss) / np.abs(val_loss)))
+
+    logging.debug("Grad desc took %0.2f msec", get_time(t1))
 
     # Set the optimized params to new params.
     policy_net_params = optimizers.get_params(ppo_opt_state)
     value_net_params = optimizers.get_params(value_opt_state)
 
-    print("Epoch [%s] took [%s]msec." % (i, round(
-        (time.time() - t0) * 1000, 2)))
-    print()
-
-  print("value_losses: ", np.stack(value_losses))
-  print("ppo_objective: ", np.stack(ppo_objective))
-  print("average_rewards: ", average_rewards)
+    logging.info("Epoch [% 6d], average reward [%10.2f], ppo loss [%10.2f], "
+                 "value loss [%10.2f], took [%10.2f msec]",
+                 i, avg_reward, new_ppo_loss, new_val_loss, get_time(t0))
 
+  logging.debug("value_losses: %s", np.stack(value_losses))
+  logging.debug("ppo_objective: %s", np.stack(ppo_objective))
+  logging.debug("average_rewards: %s", average_rewards)
 
-if __name__ == "__main__":
-  app.run(main)
+  return ((policy_net_params, value_net_params), average_rewards,
+          np.stack(value_losses), np.stack(ppo_objective))
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
new file mode 100644
index 000000000..be2edac14
--- /dev/null
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -0,0 +1,40 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""PPO binary over a gym env."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl import app
+from absl import flags
+from absl import logging
+from tensor2tensor.trax.rlax import ppo
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("env", "CartPole-v0", "Name of the environment to make.")
+flags.DEFINE_integer("epochs", 100, "Number of epochs to run for.")
+flags.DEFINE_integer("log_level", logging.INFO, "Log level.")
+
+
+def main(unused_argv):
+  logging.set_verbosity(FLAGS.log_level)
+  ppo.training_loop(env_name=FLAGS.env, epochs=FLAGS.epochs)
+
+
+if __name__ == "__main__":
+  app.run(main)
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
new file mode 100644
index 000000000..1d5154688
--- /dev/null
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -0,0 +1,42 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.rlax.ppo."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.trax.rlax import ppo
+from tensorflow import test
+
+
+class PpoTest(test.TestCase):
+
+  def test_rewards_to_go(self):
+    time_steps = 4
+    # [1., 1., 1., 1.]
+    rewards = np.ones((time_steps,))
+    # No discounting.
+    self.assertAllEqual(ppo.rewards_to_go(rewards, gamma=1.0),
+                        np.array([4., 3., 2., 1.]))
+    # Discounting.
+    self.assertAllEqual(ppo.rewards_to_go(rewards, gamma=0.5),
+                        np.array([1.875, 1.75, 1.5, 1.]))
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
new file mode 100644
index 000000000..42fcd7876
--- /dev/null
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -0,0 +1,47 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.rlax.ppo's training_loop."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gym
+from tensor2tensor.rl import gym_utils
+from tensor2tensor.trax.rlax import ppo
+from tensorflow import test
+
+
+class PpoTrainingLoopTest(test.TestCase):
+
+  def test_training_loop(self):
+    env = gym.make("CartPole-v0")
+    # Usually gym envs are wrapped in TimeLimit wrapper.
+    env = gym_utils.remove_time_limit_wrapper(env)
+    # Limit this to a small number for tests.
+    env = gym.wrappers.TimeLimit(env, max_episode_steps=2)
+    num_epochs = 2
+    batch_size = 2
+    _, rewards, val_losses, ppo_objectives = ppo.training_loop(
+        env=env, epochs=num_epochs, batch_size=batch_size,
+        num_optimizer_steps=1)
+    self.assertLen(rewards, num_epochs)
+    self.assertLen(val_losses, num_epochs)
+    self.assertLen(ppo_objectives, num_epochs)
+
+
+if __name__ == "__main__":
+  test.main()

From efea4ab1ca262092dde6e357ade29c353c37ed9f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 9 Apr 2019 11:48:05 -0700
Subject: [PATCH 1884/2720] [trax] Update trax with new stax init_fun signature
 c.f. https://github.com/google/jax/pull/566

PiperOrigin-RevId: 242710075
---
 tensor2tensor/trax/rlax/ppo.py       |  13 ++--
 tensor2tensor/trax/stax/attention.py |  12 ++--
 tensor2tensor/trax/stax/slax.py      |  18 ++---
 tensor2tensor/trax/stax/slax_test.py |   9 +--
 tensor2tensor/trax/stax/stax_base.py | 102 +++++++++++++++------------
 tensor2tensor/trax/trax.py           |   9 +--
 6 files changed, 91 insertions(+), 72 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index aa73dd3ad..9dd7825b0 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -27,6 +27,7 @@
 from jax import jit
 from jax import lax
 from jax import numpy as np
+from jax import random as jax_random
 from jax import vmap
 from jax.experimental import optimizers
 import numpy as onp
@@ -43,8 +44,11 @@
 
 
 # TODO(afrozm): Have a single net for both policy and value.
-def initialize_policy_and_value_nets(num_actions, batch_observations_shape):
+def initialize_policy_and_value_nets(rng_key, num_actions,
+                                     batch_observations_shape):
   """Setup and initialize the policy and value networks."""
+  key1, key2 = jax_random.split(rng_key)
+
   policy_net_init, policy_net_apply = stax.serial(
       stax.Dense(16),
       stax.Relu,
@@ -54,7 +58,7 @@ def initialize_policy_and_value_nets(num_actions, batch_observations_shape):
       stax.Softmax,
   )
 
-  _, policy_net_params = policy_net_init(batch_observations_shape)
+  _, policy_net_params = policy_net_init(key1, batch_observations_shape)
 
   value_net_init, value_net_apply = stax.serial(
       stax.Dense(16),
@@ -64,7 +68,7 @@ def initialize_policy_and_value_nets(num_actions, batch_observations_shape):
       stax.Dense(1),  # 1 since we want to predict reward using value network.
   )
 
-  _, value_net_params = value_net_init(batch_observations_shape)
+  _, value_net_params = value_net_init(key2, batch_observations_shape)
 
   return ((policy_net_params, policy_net_apply), (value_net_params,
                                                   value_net_apply))
@@ -450,9 +454,10 @@ def training_loop(env=None,
   assert isinstance(env.action_space, gym.spaces.Discrete)
   num_actions = env.action_space.n
 
+  rng_key = jax_random.PRNGKey(0)
   ((policy_net_params, policy_net_apply),
    (value_net_params, value_net_apply)) = initialize_policy_and_value_nets(
-       num_actions, batch_observations_shape)
+       rng_key, num_actions, batch_observations_shape)
 
   (ppo_opt_state, ppo_opt_update), (value_opt_state,
                                     value_opt_update) = initialize_optimizers(
diff --git a/tensor2tensor/trax/stax/attention.py b/tensor2tensor/trax/stax/attention.py
index d00de0ffd..309f30140 100644
--- a/tensor2tensor/trax/stax/attention.py
+++ b/tensor2tensor/trax/stax/attention.py
@@ -33,7 +33,7 @@ def causal_mask(size, dtype=np.uint8):
 
 def CausalMask(axis=-1):  # pylint: disable=invalid-name
   """Layer to create a causal mask for its inputs."""
-  init_fun = lambda input_shape: (input_shape, ())
+  init_fun = lambda _, input_shape: (input_shape, ())
   def apply_fun(params, inputs, **kwargs):
     del params, kwargs
     return causal_mask(inputs.shape[axis], dtype=inputs.dtype)
@@ -85,7 +85,7 @@ def init(shape):
 
 def LayerNorm(epsilon=1e-6):  # pylint: disable=invalid-name
   """Layer construction function for Layer Normalization layer.."""
-  def init_fun(input_shape):
+  def init_fun(_, input_shape):
     features = input_shape[-1]
     scale = np.ones(features)
     bias = np.zeros(features)
@@ -104,7 +104,7 @@ def apply_fun(params, inputs, **kwargs):
 
 def Embedding(feature_depth, vocab_size):  # pylint: disable=invalid-name
   """Layer constructor function for a dense embedding layer."""
-  def init_fun(input_shape):
+  def init_fun(_, input_shape):
     output_shape = tuple(input_shape) + (feature_depth,)
     dense_embedding = xavier_uniform()((vocab_size, feature_depth))
     return output_shape, dense_embedding
@@ -117,7 +117,7 @@ def apply_fun(params, inputs, **kwargs):
 
 def PositionalEncoding(feature_depth, max_len):  # pylint: disable=invalid-name
   """Implements bare positional encoding."""
-  def init_fun(input_shape):
+  def init_fun(_, input_shape):
     # Compute the positional encodings once in log space.
     pe = onp.zeros((max_len, feature_depth), dtype=onp.float32)
     position = onp.arange(0, max_len)[:, onp.newaxis]
@@ -174,7 +174,7 @@ def PureDotProductAttention(dropout=1.0, mode='train'):  # pylint: disable=inval
   Returns:
     Pure single-headed attention layer. (No Dense transforms on input.)
   """
-  def init_fun(input_shapes):
+  def init_fun(_, input_shapes):
     q_shape, _, v_shape, _ = input_shapes
     output_shape = q_shape[:-1] + (v_shape[-1],)
     return output_shape, ()
@@ -200,7 +200,7 @@ def PureMultiHeadedAttention(  # pylint: disable=invalid-name
   Returns:
     Pure Multi-headed attention layer. (No Dense transforms on input.)
   """
-  def init_fun(input_shapes):
+  def init_fun(_, input_shapes):
     input_shape = input_shapes[0]
     output_shape = input_shape[:-1] + (feature_depth,)
     return output_shape, ()
diff --git a/tensor2tensor/trax/stax/slax.py b/tensor2tensor/trax/stax/slax.py
index c4fc096a8..37142fee0 100644
--- a/tensor2tensor/trax/stax/slax.py
+++ b/tensor2tensor/trax/stax/slax.py
@@ -35,7 +35,7 @@ def one_hot(x, size, dtype=np.float32):
 
 def ShiftRight():  # pylint: disable=invalid-name
   """Layer to shift the tensor to the right by padding on axis 1."""
-  init_fun = lambda input_shape: (input_shape, ())
+  init_fun = lambda _, input_shape: (input_shape, ())
   def apply_fun(params, inputs, **kwargs):
     del params, kwargs
     pad_widths = [(0, 0), (1, 0)]
@@ -88,7 +88,7 @@ def Take(*args):  # pylint: disable=invalid-name
     If the resulting output list has only one member, it is automatically
     unwrapped and the contents are passed on directly.
   """
-  def init_fun(input_shape):
+  def init_fun(_, input_shape):
     output_shape = []
     for arg in args:
       output_shape.append(input_shape[arg])
@@ -128,7 +128,7 @@ def return_shapes(inputs):
       return [inputs[k].shape for k in inputs.keys()]
     else:
       return inputs.shape
-  def init_fun(input_shape):
+  def init_fun(_, input_shape):
     if debug:
       logging.info('%s [init]: %s', prefix, input_shape)
     return input_shape, ()
@@ -170,16 +170,16 @@ def __init__(self, staxlayer):  # pylint: disable=super-init-not-called
     self._first_init = True
     self.params = None  # cached staxlayer params
 
-  def _init_fun(self, input_shape):  # pylint: disable=missing-docstring
+  def _init_fun(self, rng_key, input_shape):  # pylint: disable=missing-docstring
     if self._first_init:
       # point of first subgraph initialization call: sets params, output_shape
       self._first_init = False
-      out_shape, self.params = self._orig_init_fun(input_shape)
+      out_shape, self.params = self._orig_init_fun(rng_key, input_shape)
       return out_shape, self.params
     else:
       # point of subgraph reuse:
       # params are just a marker to apply_funs signalling subgraph params reuse
-      out_shape, _ = self._orig_init_fun(input_shape)
+      out_shape, _ = self._orig_init_fun(rng_key, input_shape)
       return out_shape, _TreeMarker()
 
   def _apply_fun(self, params, inputs, **kwargs):
@@ -215,11 +215,11 @@ def __init__(self, staxlayer):  # pylint: disable=super-init-not-called
     self.params = None  # cached staxlayer params
     self.value = None  # cached staxlayer output value
 
-  def _init_fun(self, input_shape):
+  def _init_fun(self, rng_key, input_shape):
     if self._first_init:
       # point of first subgraph initialization call: sets params, output_shape
       self._first_init = False
-      self._out_shape, self.params = self._orig_init_fun(input_shape)
+      self._out_shape, self.params = self._orig_init_fun(rng_key, input_shape)
       return self._out_shape, self.params
     else:
       # point of subgraph reuse:
@@ -344,7 +344,7 @@ def next(self):  # PY2
 
 def _PlaceholderInputs():  # pylint: disable=invalid-name
   """Feeds placeholders into input combinators of a Lambda-bound staxlayer."""
-  init_fun = lambda input_shape: iter((_PlaceholderTree(), _PlaceholderTree()))
+  init_fun = lambda _, shape: iter((_PlaceholderTree(), _PlaceholderTree()))
   apply_fun = lambda params, inputs, **kwargs: _PlaceholderTree()
   return init_fun, apply_fun
 _PlaceholderInputs = _PlaceholderInputs()  # pylint: disable=invalid-name
diff --git a/tensor2tensor/trax/stax/slax_test.py b/tensor2tensor/trax/stax/slax_test.py
index bdb7783d0..dd15a378b 100644
--- a/tensor2tensor/trax/stax/slax_test.py
+++ b/tensor2tensor/trax/stax/slax_test.py
@@ -35,10 +35,10 @@ def random_inputs(rng, input_shape):
 
 
 def check_shape_agreement(test_case, init_fun, apply_fun, input_shape):
-  result_shape, params = init_fun(input_shape)
+  rng_key1, rng_key2 = random.split(random.PRNGKey(0))
+  result_shape, params = init_fun(rng_key1, input_shape)
   inputs = random_inputs(onp.random.RandomState(0), input_shape)
-  rng_key = random.PRNGKey(0)
-  result = apply_fun(params, inputs, rng=rng_key)
+  result = apply_fun(params, inputs, rng=rng_key2)
   test_case.assertEqual(result.shape, result_shape)
 
 
@@ -149,7 +149,8 @@ def lambda_fun(x, y, z, w, v):
         return _build_combinator_tree(tree_spec, (x, y, z, w, v))
       check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*5)
 
-  def testLambda_6_args(self):
+  # TODO(mattjj,levskaya): timing out, re-enable with longer timeout?
+  def DISABLED_testLambda_6_args(self):  # pylint: disable=invalid-name
     for tree_spec in _enumerate_trees_w_leaves(6):
       @stax.Lambda
       def lambda_fun(x, y, z, w, v, u):
diff --git a/tensor2tensor/trax/stax/stax_base.py b/tensor2tensor/trax/stax/stax_base.py
index 3a13911e6..b368f230f 100644
--- a/tensor2tensor/trax/stax/stax_base.py
+++ b/tensor2tensor/trax/stax/stax_base.py
@@ -30,7 +30,6 @@
 from jax import random
 
 import numpy as onp
-import numpy.random as npr
 from six.moves import reduce
 from tensor2tensor.trax import backend
 from tensor2tensor.trax.backend import numpy as np
@@ -39,6 +38,15 @@
 # names of layer constructors, like Conv and Relu, while using snake_case for
 # other functions, like lax.conv and relu.
 
+# use CamelCase for layer constructors
+# pylint: disable=invalid-name
+
+# don't have docstrings for init_fun / apply_fun pairs
+# pylint: disable=missing-docstring
+
+# ignore unused arguments
+# pylint: disable=unused-argument
+
 
 def relu(x):
   return np.maximum(x, 0.)
@@ -67,24 +75,25 @@ def fastvar(x, axis, keepdims):
 
 
 # Initializers
-def randn(stddev=1e-2, rng=npr):
+def randn(stddev=1e-2):
   """An initializer function for random normal coefficients."""
-  def init(shape):
-    return rng.normal(size=shape, scale=stddev).astype('float32')
+  def init(rng, shape):
+    return (stddev * random.normal(rng, shape)).astype('float32')
   return init
 
 
-def glorot(out_dim=0, in_dim=1, scale=onp.sqrt(2), rng=npr):
+def glorot(out_dim=0, in_dim=1, scale=onp.sqrt(2)):
   """An initializer function for random Glorot-scaled coefficients."""
-  def init(shape):
+  def init(rng, shape):
     fan_in, fan_out = shape[in_dim], shape[out_dim]
     size = onp.prod(onp.delete(shape, [in_dim, out_dim]))
     std = scale / np.sqrt((fan_in + fan_out) / 2. * size)
-    return rng.normal(size=shape, scale=std).astype('float32')
+    return (std * random.normal(rng, shape)).astype('float32')
   return init
 
-zeros = functools.partial(np.zeros, dtype='float32')
-ones = functools.partial(np.ones, dtype='float32')
+
+zeros = lambda rng, shape: np.zeros(shape, dtype='float32')
+ones = lambda rng, shape: np.ones(shape, dtype='float32')
 
 
 # Layers
@@ -96,13 +105,14 @@ def init(shape):
 
 def Dense(out_dim, W_init=glorot(), b_init=randn()):
   """Layer constructor function for a dense (fully-connected) layer."""
-  def init_fun(input_shape):
+  def init_fun(rng, input_shape):
     output_shape = input_shape[:-1] + (out_dim,)
-    W, b = W_init((input_shape[-1], out_dim)), b_init((out_dim,))
-    return output_shape, (W, b)
+    w, b = W_init(rng, (input_shape[-1], out_dim)), b_init(rng, (out_dim,))
+    return output_shape, (w, b)
   def apply_fun(params, inputs, **kwargs):
-    W, b = params
-    return np.dot(inputs, W) + b
+    del kwargs  # unused
+    w, b = params
+    return np.dot(inputs, w) + b
   return init_fun, apply_fun
 
 
@@ -113,7 +123,7 @@ def GeneralConv(dimension_numbers, out_chan, filter_shape,
   one = (1,) * len(filter_shape)
   strides = strides or one
   W_init = W_init or glorot(rhs_spec.index('O'), rhs_spec.index('I'))
-  def init_fun(input_shape):
+  def init_fun(rng, input_shape):
     filter_shape_iter = iter(filter_shape)
     kernel_shape = [out_chan if c == 'O' else
                     input_shape[lhs_spec.index('C')] if c == 'I' else
@@ -122,7 +132,7 @@ def init_fun(input_shape):
         input_shape, kernel_shape, strides, padding, dimension_numbers)
     bias_shape = [out_chan if c == 'C' else 1 for c in out_spec]
     bias_shape = tuple(itertools.dropwhile(lambda x: x == 1, bias_shape))
-    W, b = W_init(kernel_shape), b_init(bias_shape)
+    W, b = W_init(rng, kernel_shape), b_init(rng, bias_shape)
     return output_shape, (W, b)
   def apply_fun(params, inputs, **kwargs):
     W, b = params
@@ -135,12 +145,12 @@ def apply_fun(params, inputs, **kwargs):
 def BatchNorm(axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True,
               beta_init=zeros, gamma_init=ones):
   """Layer construction function for a batch normalization layer."""
-  _beta_init = lambda shape: beta_init(shape) if center else ()
-  _gamma_init = lambda shape: gamma_init(shape) if scale else ()
+  _beta_init = lambda rng, shape: beta_init(rng, shape) if center else ()
+  _gamma_init = lambda rng, shape: gamma_init(rng, shape) if scale else ()
   axis = (axis,) if np.isscalar(axis) else axis
-  def init_fun(input_shape):
+  def init_fun(rng, input_shape):
     shape = tuple(d for i, d in enumerate(input_shape) if i not in axis)
-    beta, gamma = _beta_init(shape), _gamma_init(shape)
+    beta, gamma = _beta_init(rng, shape), _gamma_init(rng, shape)
     return input_shape, (beta, gamma)
   def apply_fun(params, x, **kwargs):
     beta, gamma = params
@@ -159,7 +169,7 @@ def apply_fun(params, x, **kwargs):
 
 
 def _elemwise_no_params(fun, **fun_kwargs):
-  init_fun = lambda input_shape: (input_shape, ())
+  init_fun = lambda rng, input_shape: (input_shape, ())
   apply_fun = lambda params, inputs, **kwargs: fun(inputs, **fun_kwargs)
   return init_fun, apply_fun
 Tanh = _elemwise_no_params(np.tanh)
@@ -177,8 +187,9 @@ def PoolingLayer(window_shape, strides=None, padding='VALID'):
     rescale = rescaler(window_shape, strides, padding) if rescaler else None
     dims = (1,) + window_shape + (1,)  # NHWC
     strides = (1,) + strides + (1,)
-    def init_fun(input_shape):
-      out_shape = lax.reduce_window_shape_tuple(input_shape, dims, strides, padding)
+    def init_fun(rng, input_shape):
+      out_shape = lax.reduce_window_shape_tuple(input_shape, dims, strides,
+                                                padding)
       return out_shape, ()
     def apply_fun(params, inputs, **kwargs):
       out = lax.reduce_window(inputs, init_val, reducer, dims, strides, padding)
@@ -200,7 +211,7 @@ def rescale(outputs, inputs):
 
 def Flatten():
   """Layer construction function for flattening all but the leading dim."""
-  def init_fun(input_shape):
+  def init_fun(rng, input_shape):
     output_shape = input_shape[0], reduce(op.mul, input_shape[1:], 1)
     return output_shape, ()
   def apply_fun(params, inputs, **kwargs):
@@ -211,7 +222,7 @@ def apply_fun(params, inputs, **kwargs):
 
 def Identity():
   """Layer construction function for an identity layer."""
-  init_fun = lambda input_shape: (input_shape, ())
+  init_fun = lambda rng, input_shape: (input_shape, ())
   apply_fun = lambda params, inputs, **kwargs: inputs
   return init_fun, apply_fun
 Identity = Identity()
@@ -219,14 +230,14 @@ def Identity():
 
 def FanOut(num):
   """Layer construction function for a fan-out layer."""
-  init_fun = lambda input_shape: ([input_shape] * num, ())
+  init_fun = lambda rng, input_shape: ([input_shape] * num, ())
   apply_fun = lambda params, inputs, **kwargs: [inputs] * num
   return init_fun, apply_fun
 
 
 def FanInSum():
   """Layer construction function for a fan-in sum layer."""
-  init_fun = lambda input_shape: (input_shape[0], ())
+  init_fun = lambda rng, input_shape: (input_shape[0], ())
   apply_fun = lambda params, inputs, **kwargs: sum(inputs)
   return init_fun, apply_fun
 FanInSum = FanInSum()
@@ -234,7 +245,7 @@ def FanInSum():
 
 def FanInConcat(axis=-1):
   """Layer construction function for a fan-in concatenation layer."""
-  def init_fun(input_shape):
+  def init_fun(rng, input_shape):
     ax = axis % len(input_shape[0])
     concat_size = sum(shape[ax] for shape in input_shape)
     out_shape = input_shape[0][:ax] + (concat_size,) + input_shape[0][ax+1:]
@@ -246,15 +257,16 @@ def apply_fun(params, inputs, **kwargs):
 
 def Dropout(rate, mode='train'):
   """Layer construction function for a dropout layer with given rate."""
-  def init_fun(input_shape):
+  def init_fun(_, input_shape):
     return input_shape, ()
-  def apply_fun(params, inputs, **kwargs):
+  def apply_fun(params, inputs, **kwargs):  # pylint: disable=missing-docstring
+    del params  # Unused.
     rng = kwargs.get('rng', None)
     if rng is None:
-      msg = ("Dropout layer requires apply_fun to be called with a PRNG key "
-             "argument. That is, instead of `apply_fun(params, inputs)`, call "
-             "it like `apply_fun(params, inputs, key)` where `key` is a "
-             "jax.random.PRNGKey value.")
+      msg = ('Dropout layer requires apply_fun to be called with a PRNG key '
+             'argument. That is, instead of `apply_fun(params, inputs)`, call '
+             'it like `apply_fun(params, inputs, key)` where `key` is a '
+             'jax.random.PRNGKey value.')
       raise ValueError(msg)
     if mode == 'train':
       keep = random.bernoulli(rng, rate, inputs.shape)
@@ -272,17 +284,17 @@ def serial(*layers):
 
   Args:
     *layers: a sequence of layers, each an (init_fun, apply_fun) pair.
-
   Returns:
     A new layer, meaning an (init_fun, apply_fun) pair, representing the serial
     composition of the given sequence of layers.
   """
   nlayers = len(layers)
   init_funs, apply_funs = zip(*layers)
-  def init_fun(input_shape):
+  def init_fun(rng, input_shape):
     params = []
     for init_fun in init_funs:
-      input_shape, param = init_fun(input_shape)
+      rng, layer_rng = random.split(rng)
+      input_shape, param = init_fun(layer_rng, input_shape)
       params.append(param)
     return input_shape, params
   def apply_fun(params, inputs, **kwargs):
@@ -299,10 +311,8 @@ def parallel(*layers):
 
   The layer resulting from this combinator is often used with the FanOut and
   FanInSum layers.
-
   Args:
     *layers: a sequence of layers, each an (init_fun, apply_fun) pair.
-
   Returns:
     A new layer, meaning an (init_fun, apply_fun) pair, representing the
     parallel composition of the given sequence of layers. In particular, the
@@ -311,12 +321,15 @@ def parallel(*layers):
   """
   nlayers = len(layers)
   init_funs, apply_funs = zip(*layers)
-  def init_fun(input_shape):
-    return zip(*[init(shape) for init, shape in zip(init_funs, input_shape)])
+  def init_fun(rng, input_shape):
+    rngs = random.split(rng, nlayers)
+    return zip(*[init(rng, shape) for init, rng, shape
+                 in zip(init_funs, rngs, input_shape)])
   def apply_fun(params, inputs, **kwargs):
     rng = kwargs.pop('rng', None)
     rngs = random.split(rng, nlayers) if rng is not None else (None,) * nlayers
-    return [f(p, x, rng=r, **kwargs) for f, p, x, r in zip(apply_funs, params, inputs, rngs)]
+    return [f(p, x, rng=r, **kwargs)
+            for f, p, x, r in zip(apply_funs, params, inputs, rngs)]
   return init_fun, apply_fun
 
 
@@ -326,14 +339,13 @@ def shape_dependent(make_layer):
   Args:
     make_layer: a one-argument function that takes an input shape as an argument
       (a tuple of positive integers) and returns an (init_fun, apply_fun) pair.
-
   Returns:
     A new layer, meaning an (init_fun, apply_fun) pair, representing the same
     layer as returned by `make_layer` but with its construction delayed until
     input shapes are known.
   """
-  def init_fun(input_shape):
-    return make_layer(input_shape)[0](input_shape)
+  def init_fun(rng, input_shape):
+    return make_layer(input_shape)[0](rng, input_shape)
   def apply_fun(params, inputs, **kwargs):
     return make_layer(inputs.shape)[1](params, inputs, **kwargs)
   return init_fun, apply_fun
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 443936a86..e8da33b7d 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -252,6 +252,7 @@ def single_update(i, opt_state, batch, rng):
           params, batch, predict_fun, rng), opt_state)
     return backend.jit(single_update)
 
+  @functools.partial(backend.pmap, axis_name="batch")
   def mapped_update(i, opt_state, batch, rng):
     """This is a multi-device version of the update function above."""
     # We assume all tensors have the first dimension = num_devices.
@@ -263,9 +264,7 @@ def mapped_update(i, opt_state, batch, rng):
 
   def update(i, opt_state, batch, rng):
     # TODO(lukaszkaiser): investigate how to replicate rng and correct.
-    return backend.pmap(mapped_update(
-        jax.replicate(i), opt_state, batch, jax.replicate(rng)),
-                        axis_name="batch")
+    return mapped_update(jax.replicate(i), opt_state, batch, jax.replicate(rng))
 
   return update
 
@@ -339,7 +338,9 @@ def train(output_dir,
 
   # Setup state
   step = state.step or 0
-  params_initializer = lambda: model_init([-1] + list(inputs.input_shape))[1]
+  rng, init_key = jax_random.split(rng)
+  params_initializer = \
+      lambda: model_init(init_key, [-1] + list(inputs.input_shape))[1]
   params = state.params or params_initializer()
   opt_state = opt_init(params)
   if num_devices > 1:  # TODO(lukaszkaiser): use everywhere when pmap is stable.

From a20133d79039a2daa1c03623498b6ece8e628fd0 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 9 Apr 2019 12:22:47 -0700
Subject: [PATCH 1885/2720] Bunch of renamings and adding some docstrings.

PiperOrigin-RevId: 242716550
---
 tensor2tensor/trax/rlax/ppo.py | 183 +++++++++++++++++----------------
 1 file changed, 96 insertions(+), 87 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 9dd7825b0..bc27195a0 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -208,6 +208,7 @@ def pad_trajectories(trajectories, boundary=10):
       padded_observations), np.stack(padded_actions), np.stack(padded_rewards)
 
 
+# TODO(afrozm): Make this batched by default.
 def rewards_to_go(rewards, reward_mask=1.0, gamma=0.99):
   r"""r2g[t] = \sum_{l=0}^{\infty}(\gamma^l * r_{t+l})."""
   time_steps = len(rewards)
@@ -228,13 +229,14 @@ def rewards_to_go(rewards, reward_mask=1.0, gamma=0.99):
   return r2g * reward_mask
 
 
+# TODO(afrozm): Make this batched by default.
 @functools.partial(jit, static_argnums=(0,))
-def batched_avg_value_function_loss(value_net_apply,
-                                    value_net_params,
-                                    observations,
-                                    rewards,
-                                    reward_mask=1.0,
-                                    gamma=0.99):
+def value_loss(value_net_apply,
+               value_net_params,
+               observations,
+               rewards,
+               reward_mask=1.0,
+               gamma=0.99):
   """L2 loss on the value function's outputs."""
 
   # Capturing the value_net_apply from the parent function's scope.
@@ -245,8 +247,7 @@ def _value_function_loss_trajectory(value_net_params,
                                       reward_mask=1.0,
                                       gamma=0.99):
     """Compute the actual loss for a trajectory."""
-    r2g = rewards_to_go(
-        rewards, reward_mask=reward_mask, gamma=gamma)
+    r2g = rewards_to_go(rewards, reward_mask=reward_mask, gamma=gamma)
     v = value_net_apply(value_net_params, observations[:-1])
     v = np.squeeze(v) * reward_mask
     loss = v - r2g
@@ -260,28 +261,25 @@ def _value_function_loss_trajectory(value_net_params,
           value_net_params, observations, rewards, reward_mask, gamma=gamma))
 
 
-def batched_deltas(predicted_values, rewards, reward_mask, gamma=0.99):
+def deltas(predicted_values, rewards, reward_mask, gamma=0.99):
   r"""\delta_t = \sum_{l = 0}^{\infty}(r_t + \gamma * V(s_{t+1}) - V(s_t))."""
   # predicted_values are application of value net only the observations.
   # B x T+1
 
-  deltas = []
+  # `d`s are basically one-step TD residuals.
+  d = []
   _, T = rewards.shape  # pylint: disable=invalid-name
   for t in range(T):
-    deltas.append(rewards[:, t] + (gamma * predicted_values[:, t + 1]) -
-                  predicted_values[:, t])
+    d.append(rewards[:, t] + (gamma * predicted_values[:, t + 1]) -
+             predicted_values[:, t])
 
-  return np.array(deltas).T * reward_mask
+  return np.array(d).T * reward_mask
 
 
-def batched_gae_advantages(
-    deltas,
-    reward_mask,
-    lamda=0.95,  # NOTYPO
-    gamma=0.99):
+def gae_advantages(td_deltas, reward_mask, lambda_=0.95, gamma=0.99):
   r"""A_t = \sum_{l=0}^{\infty}(\gamma * \lambda)^{l}(\delta_{t+l})."""
-  _, T = deltas.shape  # pylint: disable=invalid-name
-  gl = lamda * gamma  # NOTYPO
+  _, T = td_deltas.shape  # pylint: disable=invalid-name
+  gl = lambda_ * gamma
 
   # [[1, gl, gl**2, ... gl**T-1]]
   # Not jittable, T should be a compile time constant.
@@ -292,8 +290,8 @@ def batched_gae_advantages(
   gl_gp = np.array(gl_geometric_progression)
   gl_gp = gl_gp.reshape((1, T))
 
-  # deltas * gl_gp
-  deltas_gl_gp = deltas * gl_gp
+  # td_deltas * gl_gp
+  deltas_gl_gp = td_deltas * gl_gp
 
   # A0 - advantage for 0th time-step, across all batches.
   As = []  # pylint: disable=invalid-name
@@ -302,25 +300,36 @@ def batched_gae_advantages(
 
   # Now compute the other advantages.
   for t in range(1, T):
-    As.append((As[-1] - deltas[:, t - 1]) / gl)
+    As.append((As[-1] - td_deltas[:, t - 1]) / gl)
 
   return np.stack(As).T * reward_mask
 
 
-def batched_probabs(probab_observations, actions):
+def chosen_probabs(probab_observations, actions):
+  """Picks out the probabilities of the actions along batch and time-steps.
+
+  Args:
+    probab_observations: `[B, T, #actions]` ndarray, where
+      probab_observations[b, t, i] contains the probability of action = i at the
+      t^th time-step in the b^th trajectory.
+    actions: `[B, T]` ndarray, with each entry in [0, #actions) denoting which
+      action was chosen in the b^th trajectory's t^th time-step.
+
+  Returns:
+    `[B, T]` ndarray with the probabilities of the chosen actions.
+  """
   b, t = actions.shape
   return probab_observations[np.arange(b)[:, None], np.arange(t), actions]
 
 
-def batched_probab_ratios(policy_net_apply, old_policy_params,
-                          new_policy_params, observations, actions,
-                          reward_mask):
+def probab_ratios(policy_net_apply, old_policy_params, new_policy_params,
+                  observations, actions, reward_mask):
   """Calculates the probaility ratios for each time-step in a trajectory."""
   p_old = policy_net_apply(old_policy_params, observations)
   p_new = policy_net_apply(new_policy_params, observations)
 
-  bp_old = batched_probabs(p_old, actions)
-  bp_new = batched_probabs(p_new, actions)
+  bp_old = chosen_probabs(p_old, actions)
+  bp_new = chosen_probabs(p_new, actions)
 
   # Add a small number to bp_old, where reward_mask is 0, this is just to help
   # never to divide by 0.
@@ -331,66 +340,62 @@ def batched_probab_ratios(policy_net_apply, old_policy_params,
   return ret_val
 
 
-def batched_clipped_probab_ratios(bpr, reward_mask, epsilon=0.2):
+def clipped_probab_ratios(bpr, reward_mask, epsilon=0.2):
   return reward_mask * np.clip(bpr, 1 - epsilon, 1 + epsilon)
 
 
-def batched_clipped_objective(bpr, adv, reward_mask, epsilon=0.2):
+def clipped_objective(bpr, adv, reward_mask, epsilon=0.2):
   c1 = bpr * adv
-  c2 = batched_clipped_probab_ratios(bpr, reward_mask, epsilon=epsilon) * adv
+  c2 = clipped_probab_ratios(bpr, reward_mask, epsilon=epsilon) * adv
   return np.minimum(c1, c2)
 
 
 @functools.partial(jit, static_argnums=(0, 3))
-def batched_ppo_loss(
-    policy_net_apply,
-    new_policy_params,
-    old_policy_params,
-    value_net_apply,
-    value_net_params,
-    padded_observations,
-    padded_actions,
-    padded_rewards,
-    reward_mask,
-    gamma=0.99,
-    lamda=0.95,  # NOTYPO
-    epsilon=0.2):
+def ppo_loss(policy_net_apply,
+             new_policy_params,
+             old_policy_params,
+             value_net_apply,
+             value_net_params,
+             padded_observations,
+             padded_actions,
+             padded_rewards,
+             reward_mask,
+             gamma=0.99,
+             lambda_=0.95,
+             epsilon=0.2):
   """PPO objective, with an eventual minus sign."""
   # V(s_t) forall s & t
   value_function = np.squeeze(
       value_net_apply(value_net_params, padded_observations))
-  deltas = batched_deltas(
-      value_function, padded_rewards, reward_mask, gamma=gamma)
-  advantages = batched_gae_advantages(
-      deltas, reward_mask, lamda=lamda, gamma=gamma)  # NOTYPO
-  ratios = batched_probab_ratios(policy_net_apply, old_policy_params,
-                                 new_policy_params, padded_observations,
-                                 padded_actions, reward_mask)
-  clipped_loss = batched_clipped_objective(
+  td_deltas = deltas(value_function, padded_rewards, reward_mask, gamma=gamma)
+  advantages = gae_advantages(
+      td_deltas, reward_mask, lambda_=lambda_, gamma=gamma)
+  ratios = probab_ratios(policy_net_apply, old_policy_params, new_policy_params,
+                         padded_observations, padded_actions, reward_mask)
+  clipped_loss = clipped_objective(
       ratios, advantages, reward_mask, epsilon=epsilon)
   return -np.sum(clipped_loss)
 
 
 @functools.partial(jit, static_argnums=(2, 3, 5))
-def ppo_opt_step(
-    i,
-    opt_state,
-    ppo_opt_update,
-    policy_net_apply,
-    old_policy_params,
-    value_net_apply,
-    value_net_params,
-    padded_observations,
-    padded_actions,
-    padded_rewards,
-    reward_mask,
-    gamma=0.99,
-    lamda=0.95,  # NOTYPO
-    epsilon=0.1):
+def ppo_opt_step(i,
+                 opt_state,
+                 ppo_opt_update,
+                 policy_net_apply,
+                 old_policy_params,
+                 value_net_apply,
+                 value_net_params,
+                 padded_observations,
+                 padded_actions,
+                 padded_rewards,
+                 reward_mask,
+                 gamma=0.99,
+                 lambda_=0.95,
+                 epsilon=0.1):
   """PPO optimizer step."""
   new_policy_params = optimizers.get_params(opt_state)
   g = grad(
-      batched_ppo_loss, argnums=1)(
+      ppo_loss, argnums=1)(
           policy_net_apply,
           new_policy_params,
           old_policy_params,
@@ -401,7 +406,7 @@ def ppo_opt_step(
           padded_rewards,
           reward_mask,
           gamma=gamma,
-          lamda=lamda,  # NOTYPO
+          lambda_=lambda_,
           epsilon=epsilon)
   return ppo_opt_update(i, g, opt_state)
 
@@ -418,7 +423,7 @@ def value_opt_step(i,
   """Value optimizer step."""
   value_params = optimizers.get_params(opt_state)
   # Note this partial application here and argnums above in ppo_opt_step.
-  g = grad(functools.partial(batched_avg_value_function_loss, value_net_apply))(
+  g = grad(functools.partial(value_loss, value_net_apply))(
       value_params,
       padded_observations,
       padded_rewards,
@@ -494,7 +499,7 @@ def training_loop(env=None,
     epsilon = 0.1 if epochs == 1 else 0.1 * (1.0 - (i / (epochs - 1)))
 
     t = time.time()
-    val_loss = batched_avg_value_function_loss(
+    cur_value_loss = value_loss(
         value_net_apply,
         value_net_params,
         padded_observations,
@@ -503,10 +508,10 @@ def training_loop(env=None,
         gamma=GAMMA)
 
     logging.debug("Calculating value loss took %0.2f msec.", get_time(t))
-    value_losses.append(val_loss)
+    value_losses.append(cur_value_loss)
 
     t = time.time()
-    ppo_loss = batched_ppo_loss(
+    cur_ppo_loss = ppo_loss(
         policy_net_apply,
         policy_net_params,
         policy_net_params,
@@ -517,11 +522,11 @@ def training_loop(env=None,
         padded_rewards,
         reward_mask,
         gamma=GAMMA,
-        lamda=LAMBDA,  # NOTYPO
+        lambda_=LAMBDA,
         epsilon=epsilon)
     # ppo_loss = 11.00110011
     logging.debug("Calculating PPO loss took %0.2f msec.", get_time(t))
-    ppo_objective.append(-ppo_loss)
+    ppo_objective.append(-cur_ppo_loss)
 
     # Run optimizers.
     logging.debug("PPO Optimization")
@@ -543,14 +548,14 @@ def training_loop(env=None,
           padded_rewards,
           reward_mask,
           gamma=GAMMA,
-          lamda=LAMBDA,  # NOTYPO
+          lambda_=LAMBDA,
           epsilon=epsilon)
       t2 = time.time()
       # Get the new params.
       new_policy_net_params = optimizers.get_params(ppo_opt_state)
       if ((j + 1) %
           print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
-        new_ppo_loss = batched_ppo_loss(
+        new_ppo_loss = ppo_loss(
             policy_net_apply,
             new_policy_net_params,
             policy_net_params,
@@ -561,15 +566,16 @@ def training_loop(env=None,
             padded_rewards,
             reward_mask,
             gamma=GAMMA,
-            lamda=LAMBDA,  # NOTYPO
+            lambda_=LAMBDA,
             epsilon=epsilon)
         logging.debug("One PPO grad desc took: %0.2f msec", get_time(t, t2))
-        logging.debug("PPO loss [%10.2f] -> [%10.2f]", ppo_loss, new_ppo_loss)
+        logging.debug("PPO loss [%10.2f] -> [%10.2f]", cur_ppo_loss,
+                      new_ppo_loss)
       # Update the params.
       policy_net_params = new_policy_net_params
 
     logging.debug("Total PPO loss reduction [%0.2f]%%",
-                  (100 * (ppo_loss - new_ppo_loss) / np.abs(ppo_loss)))
+                  (100 * (cur_ppo_loss - new_ppo_loss) / np.abs(cur_ppo_loss)))
 
     logging.debug("Value Optimization")
 
@@ -588,7 +594,7 @@ def training_loop(env=None,
       value_net_params = optimizers.get_params(value_opt_state)
       if ((j + 1) %
           print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
-        new_val_loss = batched_avg_value_function_loss(
+        new_value_loss = value_loss(
             value_net_apply,
             value_net_params,
             padded_observations,
@@ -596,9 +602,11 @@ def training_loop(env=None,
             reward_mask,
             gamma=GAMMA)
         logging.debug("One value grad desc took: %0.2f msec", get_time(t, t2))
-        logging.debug("Value loss [%10.2f] -> [%10.2f]", val_loss, new_val_loss)
+        logging.debug("Value loss [%10.2f] -> [%10.2f]", cur_value_loss,
+                      new_value_loss)
     logging.debug("Total value loss reduction [%0.2f]%%",
-                  (100 * (val_loss - new_val_loss) / np.abs(val_loss)))
+                  (100 *
+                   (cur_value_loss - new_value_loss) / np.abs(cur_value_loss)))
 
     logging.debug("Grad desc took %0.2f msec", get_time(t1))
 
@@ -606,9 +614,10 @@ def training_loop(env=None,
     policy_net_params = optimizers.get_params(ppo_opt_state)
     value_net_params = optimizers.get_params(value_opt_state)
 
-    logging.info("Epoch [% 6d], average reward [%10.2f], ppo loss [%10.2f], "
-                 "value loss [%10.2f], took [%10.2f msec]",
-                 i, avg_reward, new_ppo_loss, new_val_loss, get_time(t0))
+    logging.info(
+        "Epoch [% 6d], average reward [%10.2f], ppo loss [%10.2f], "
+        "value loss [%10.2f], took [%10.2f msec]", i, avg_reward, new_ppo_loss,
+        new_value_loss, get_time(t0))
 
   logging.debug("value_losses: %s", np.stack(value_losses))
   logging.debug("ppo_objective: %s", np.stack(ppo_objective))

From 4a751a165aabdef112eae408d5cd23cfd3de32ec Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 9 Apr 2019 12:28:03 -0700
Subject: [PATCH 1886/2720] Make loss function configurable in trax.

Was thinking of using trax in PPO so did it, but will probably integrate it
later on.

PiperOrigin-RevId: 242717394
---
 tensor2tensor/trax/trax.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index e8da33b7d..b12bec351 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -288,6 +288,7 @@ def reshape_by_device(train_data, num_devices):
 @gin.configurable(blacklist=["output_dir"])
 def train(output_dir,
           model=gin.REQUIRED,
+          loss_fun=loss,
           inputs=trax_inputs.inputs,
           optimizer=trax_opt.adam,
           lr_schedule=lr.MultifactorSchedule,
@@ -303,6 +304,8 @@ def train(output_dir,
     output_dir: Directory where to put the logs and checkpoints.
     model: The model to train as a callable returning 2 callables, an init_fun
       and apply_fun.
+    loss_fun: callable with signature: params, trax.inputs.Inputs, model, rng
+      -> loss.
     inputs: callable returning trax.inputs.Inputs.
     optimizer: The optimizer as a callable taking a learning_rate callable and
       returning 2 callables, opt_init and opt_update.
@@ -348,7 +351,7 @@ def train(output_dir,
 
   # jit model_predict and update so they're fast
   jit_model_predict = backend.jit(model_predict)  # for evaluation
-  jit_update_fun = _jit_update_fun(model_predict, loss, optimizer, lr_fun,
+  jit_update_fun = _jit_update_fun(model_predict, loss_fun, optimizer, lr_fun,
                                    num_devices)
 
   print()
@@ -362,7 +365,7 @@ def train(output_dir,
 
   # Non-compiled debug step helps find problems in models easier.
   if run_debug_step:
-    debug_loss = loss(params, next(train_stream), model_predict, rng)
+    debug_loss = loss_fun(params, next(train_stream), model_predict, rng)
     step_log(step, "Debug step loss %.8f" % debug_loss)
 
   for epoch, epoch_steps in epochs(train_steps, epoch_steps):
@@ -421,8 +424,8 @@ def train(output_dir,
     old_lr_fun = lr_fun
     lr_fun = lr_schedule(history)
     if lr_fun != old_lr_fun:  # For performance, only jit if there is a change.
-      jit_update_fun = _jit_update_fun(model_predict, loss, optimizer, lr_fun,
-                                       num_devices)
+      jit_update_fun = _jit_update_fun(model_predict, loss_fun, optimizer,
+                                       lr_fun, num_devices)
 
     # Flush summary writers
     train_sw.writer.flush()

From c0b2b6dedab82fc38cd2ffe739509381ce9baff5 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 9 Apr 2019 16:48:26 -0700
Subject: [PATCH 1887/2720]  * Make policy and value functions configurable in
 ppo training loop.  * Using trax optimizers, instead of jax.  * Use a smaller
 network in the test, this is still slow to run.

TODO
 * Use Gin, since its still a lot more verbose than I'd like.

PiperOrigin-RevId: 242766615
---
 tensor2tensor/trax/rlax/ppo.py                | 131 ++++++++++--------
 tensor2tensor/trax/rlax/ppo_main.py           |  21 ++-
 .../trax/rlax/ppo_training_loop_test.py       |  16 ++-
 3 files changed, 107 insertions(+), 61 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index bc27195a0..e532b03e1 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -21,6 +21,7 @@
 
 import functools
 import time
+
 from absl import logging
 import gym
 from jax import grad
@@ -29,8 +30,9 @@
 from jax import numpy as np
 from jax import random as jax_random
 from jax import vmap
-from jax.experimental import optimizers
 import numpy as onp
+from tensor2tensor.trax import optimizers as trax_opt
+from tensor2tensor.trax import trax
 from tensor2tensor.trax.stax import stax_base as stax
 
 DEBUG_LOGGING = False
@@ -43,50 +45,50 @@
 POLICY = "categorical-sampling"
 
 
-# TODO(afrozm): Have a single net for both policy and value.
-def initialize_policy_and_value_nets(rng_key, num_actions,
-                                     batch_observations_shape):
-  """Setup and initialize the policy and value networks."""
-  key1, key2 = jax_random.split(rng_key)
+def policy_net(jax_rng_key,
+               batch_observations_shape,
+               num_actions,
+               bottom_layers=None):
+  """A policy net function."""
+  key1, _ = jax_random.split(jax_rng_key)
 
-  policy_net_init, policy_net_apply = stax.serial(
-      stax.Dense(16),
-      stax.Relu,
-      stax.Dense(4),
-      stax.Relu,
-      stax.Dense(num_actions),
-      stax.Softmax,
-  )
+  # Use the bottom_layers as the bottom part of the network and just add the
+  # required layers on top of it.
+  if bottom_layers is None:
+    bottom_layers = []
+  bottom_layers.extend([stax.Dense(num_actions), stax.Softmax])
 
-  _, policy_net_params = policy_net_init(key1, batch_observations_shape)
+  net_init, net_apply = stax.serial(*bottom_layers)
 
-  value_net_init, value_net_apply = stax.serial(
-      stax.Dense(16),
-      stax.Relu,
-      stax.Dense(4),
-      stax.Relu,
-      stax.Dense(1),  # 1 since we want to predict reward using value network.
-  )
+  _, net_params = net_init(key1, batch_observations_shape)
+  return net_params, net_apply
 
-  _, value_net_params = value_net_init(key2, batch_observations_shape)
 
-  return ((policy_net_params, policy_net_apply), (value_net_params,
-                                                  value_net_apply))
+def value_net(jax_rng_key,
+              batch_observations_shape,
+              num_actions,
+              bottom_layers=None):
+  """A value net function."""
+  del num_actions
+  key1, _ = jax_random.split(jax_rng_key)
 
+  if bottom_layers is None:
+    bottom_layers = []
+  bottom_layers.extend([
+      stax.Dense(1),
+  ])
 
-def initialize_optimizers(policy_net_params, value_net_params):
-  """Initialize optimizers for the policy and value params."""
-  # ppo_opt_init, ppo_opt_update = optimizers.sgd(step_size=1e-3)
-  # val_opt_init, val_opt_update = optimizers.sgd(step_size=1e-3)
-  ppo_opt_init, ppo_opt_update = optimizers.adam(
-      step_size=1e-3, b1=0.9, b2=0.999, eps=1e-08)
-  value_opt_init, value_opt_update = optimizers.adam(
-      step_size=1e-3, b1=0.9, b2=0.999, eps=1e-08)
+  net_init, net_apply = stax.serial(*bottom_layers)
+
+  _, net_params = net_init(key1, batch_observations_shape)
+  return net_params, net_apply
 
-  ppo_opt_state = ppo_opt_init(policy_net_params)
-  value_opt_state = value_opt_init(value_net_params)
 
-  return (ppo_opt_state, ppo_opt_update), (value_opt_state, value_opt_update)
+def optimizer_fun(net_params):
+  opt_init, opt_update = trax_opt.adam(
+      step_size=1e-3, b1=0.9, b2=0.999, eps=1e-08)
+  opt_state = opt_init(net_params)
+  return opt_state, opt_update
 
 
 # Should this be collect 'n' trajectories, or
@@ -393,7 +395,7 @@ def ppo_opt_step(i,
                  lambda_=0.95,
                  epsilon=0.1):
   """PPO optimizer step."""
-  new_policy_params = optimizers.get_params(opt_state)
+  new_policy_params = trax_opt.get_params(opt_state)
   g = grad(
       ppo_loss, argnums=1)(
           policy_net_apply,
@@ -421,7 +423,7 @@ def value_opt_step(i,
                    reward_mask,
                    gamma=0.99):
   """Value optimizer step."""
-  value_params = optimizers.get_params(opt_state)
+  value_params = trax_opt.get_params(opt_state)
   # Note this partial application here and argnums above in ppo_opt_step.
   g = grad(functools.partial(value_loss, value_net_apply))(
       value_params,
@@ -438,15 +440,21 @@ def get_time(t1, t2=None):
   return round((t2 - t1) * 1000, 2)
 
 
-def training_loop(env=None,
-                  env_name="CartPole-v0",
-                  epochs=EPOCHS,
-                  batch_size=BATCH_TRAJECTORIES,
-                  num_optimizer_steps=NUM_OPTIMIZER_STEPS,
-                  print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
-                  random_seed=None):
+def training_loop(
+    env=None,
+    env_name="CartPole-v0",
+    epochs=EPOCHS,
+    policy_net_fun=None,
+    value_net_fun=None,
+    policy_and_value_net_fun=None,  # TODO(afrozm): Implement.
+    policy_optimizer_fun=optimizer_fun,
+    value_optimizer_fun=optimizer_fun,
+    batch_size=BATCH_TRAJECTORIES,
+    num_optimizer_steps=NUM_OPTIMIZER_STEPS,
+    print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
+    random_seed=None):
   """Runs the training loop for PPO, with fixed policy and value nets."""
-  onp.random.seed(random_seed)
+  jax_rng_key = trax.get_random_number_generator_and_set_seed(random_seed)
 
   value_losses = []
   ppo_objective = []
@@ -459,14 +467,23 @@ def training_loop(env=None,
   assert isinstance(env.action_space, gym.spaces.Discrete)
   num_actions = env.action_space.n
 
-  rng_key = jax_random.PRNGKey(0)
-  ((policy_net_params, policy_net_apply),
-   (value_net_params, value_net_apply)) = initialize_policy_and_value_nets(
-       rng_key, num_actions, batch_observations_shape)
+  # TODO(afrozm): Have a single net for both policy and action.
+  assert policy_and_value_net_fun is None
+
+  # Initialize the policy and value functions.
+  assert policy_net_fun and value_net_fun
+  jax_rng_key, key1, key2 = jax_random.split(jax_rng_key, num=3)
+
+  policy_net_params, policy_net_apply = policy_net_fun(
+      key1, batch_observations_shape, num_actions)
+  value_net_params, value_net_apply = value_net_fun(
+      key2, batch_observations_shape, num_actions)
+
+  # Initialize the optimizers.
+  assert policy_optimizer_fun and value_optimizer_fun
 
-  (ppo_opt_state, ppo_opt_update), (value_opt_state,
-                                    value_opt_update) = initialize_optimizers(
-                                        policy_net_params, value_net_params)
+  ppo_opt_state, ppo_opt_update = policy_optimizer_fun(policy_net_params)
+  value_opt_state, value_opt_update = value_optimizer_fun(value_net_params)
 
   for i in range(epochs):
     t = time.time()
@@ -552,7 +569,7 @@ def training_loop(env=None,
           epsilon=epsilon)
       t2 = time.time()
       # Get the new params.
-      new_policy_net_params = optimizers.get_params(ppo_opt_state)
+      new_policy_net_params = trax_opt.get_params(ppo_opt_state)
       if ((j + 1) %
           print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
         new_ppo_loss = ppo_loss(
@@ -591,7 +608,7 @@ def training_loop(env=None,
           reward_mask,
           gamma=GAMMA)
       t2 = time.time()
-      value_net_params = optimizers.get_params(value_opt_state)
+      value_net_params = trax_opt.get_params(value_opt_state)
       if ((j + 1) %
           print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
         new_value_loss = value_loss(
@@ -611,8 +628,8 @@ def training_loop(env=None,
     logging.debug("Grad desc took %0.2f msec", get_time(t1))
 
     # Set the optimized params to new params.
-    policy_net_params = optimizers.get_params(ppo_opt_state)
-    value_net_params = optimizers.get_params(value_opt_state)
+    policy_net_params = trax_opt.get_params(ppo_opt_state)
+    value_net_params = trax_opt.get_params(value_opt_state)
 
     logging.info(
         "Epoch [% 6d], average reward [%10.2f], ppo loss [%10.2f], "
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index be2edac14..11708bcc6 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -19,21 +19,38 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
+
 from absl import app
 from absl import flags
 from absl import logging
 from tensor2tensor.trax.rlax import ppo
+from tensor2tensor.trax.stax import stax_base as stax
 
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string("env", "CartPole-v0", "Name of the environment to make.")
 flags.DEFINE_integer("epochs", 100, "Number of epochs to run for.")
+flags.DEFINE_integer("random_seed", 0, "Random seed.")
 flags.DEFINE_integer("log_level", logging.INFO, "Log level.")
 
 
-def main(unused_argv):
+def common_stax_layers():
+  return [stax.Dense(16), stax.Relu, stax.Dense(4), stax.Relu]
+
+
+def main(argv):
+  del argv
   logging.set_verbosity(FLAGS.log_level)
-  ppo.training_loop(env_name=FLAGS.env, epochs=FLAGS.epochs)
+  bottom_layers = common_stax_layers()
+  ppo.training_loop(
+      env_name=FLAGS.env,
+      epochs=FLAGS.epochs,
+      policy_net_fun=functools.partial(
+          ppo.policy_net, bottom_layers=bottom_layers),
+      value_net_fun=functools.partial(
+          ppo.value_net, bottom_layers=bottom_layers),
+      random_seed=FLAGS.random_seed)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index 42fcd7876..7731f600a 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -19,9 +19,11 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
 import gym
 from tensor2tensor.rl import gym_utils
 from tensor2tensor.trax.rlax import ppo
+from tensor2tensor.trax.stax import stax_base as stax
 from tensorflow import test
 
 
@@ -35,9 +37,19 @@ def test_training_loop(self):
     env = gym.wrappers.TimeLimit(env, max_episode_steps=2)
     num_epochs = 2
     batch_size = 2
+    # Common bottom layer(s).
+    bottom_layers = [stax.Dense(1)]
+    # Run the training loop.
     _, rewards, val_losses, ppo_objectives = ppo.training_loop(
-        env=env, epochs=num_epochs, batch_size=batch_size,
-        num_optimizer_steps=1)
+        env=env,
+        epochs=num_epochs,
+        policy_net_fun=functools.partial(
+            ppo.policy_net, bottom_layers=bottom_layers),
+        value_net_fun=functools.partial(
+            ppo.value_net, bottom_layers=bottom_layers),
+        batch_size=batch_size,
+        num_optimizer_steps=1,
+        random_seed=0)
     self.assertLen(rewards, num_epochs)
     self.assertLen(val_losses, num_epochs)
     self.assertLen(ppo_objectives, num_epochs)

From 242e089ce1b5e7187f7f01a93b3fcd4cd0dd5565 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 10 Apr 2019 13:21:30 -0700
Subject: [PATCH 1888/2720] Make jaxboard scalar summaries work with tf eager,
 remove stale demo.

PiperOrigin-RevId: 242928265
---
 tensor2tensor/trax/jaxboard.py      |  73 ++++++++++-------
 tensor2tensor/trax/jaxboard_demo.py | 122 ----------------------------
 tensor2tensor/trax/trax.py          |   4 +-
 3 files changed, 45 insertions(+), 154 deletions(-)
 delete mode 100644 tensor2tensor/trax/jaxboard_demo.py

diff --git a/tensor2tensor/trax/jaxboard.py b/tensor2tensor/trax/jaxboard.py
index 69ac5afce..5ea33cee6 100644
--- a/tensor2tensor/trax/jaxboard.py
+++ b/tensor2tensor/trax/jaxboard.py
@@ -23,6 +23,7 @@
 
 import io
 import struct
+import time
 import warnings
 import wave
 import matplotlib as mpl
@@ -39,6 +40,11 @@
 from tensorflow import SummaryMetadata
 from tensorflow.io import gfile
 
+# pylint: disable=g-direct-tensorflow-import
+from tensorflow.core.util import event_pb2
+from tensorflow.python.summary.writer.event_file_writer import EventFileWriter
+# pylint: enable=g-direct-tensorflow-import
+
 
 def _pack_images(images, rows, cols):
   """Helper utility to make a tiled field of images from numpy arrays.
@@ -78,23 +84,30 @@ def __init__(self, log_dir):
     if not gfile.isdir(log_dir):
       gfile.makedirs(log_dir)
 
-    self.writer = tf.summary.FileWriter(log_dir, graph=None)
-    self.end_summaries = []
-    self.step = 0
-    self.closed = False
+    self._event_writer = EventFileWriter(log_dir, 10, 120, None)
+    self._step = 0
+    self._closed = False
+
+  def add_summary(self, summary, step):
+    event = event_pb2.Event(summary=summary)
+    event.wall_time = time.time()
+    if step is not None:
+      event.step = int(step)
+    self._event_writer.add_event(event)
 
   def close(self):
     """Close SummaryWriter. Final!"""
-    if not self.closed:
-      for summary in self.end_summaries:
-        self.writer.add_summary(summary, self.step)
-      self.writer.close()
-      self.closed = True
-      del self.writer
+    if not self._closed:
+      self._event_writer.close()
+      self._closed = True
+      del self._event_writer
 
   def __del__(self):  # safe?
     self.close()
 
+  def flush(self):
+    self._event_writer.flush()
+
   def scalar(self, tag, value, step=None):
     """Saves scalar value.
 
@@ -105,11 +118,11 @@ def scalar(self, tag, value, step=None):
     """
     value = float(onp.array(value))
     if step is None:
-      step = self.step
+      step = self._step
     else:
-      self.step = step
+      self._step = step
     summary = Summary(value=[Summary.Value(tag=tag, simple_value=value)])
-    self.writer.add_summary(summary, step)
+    self.add_summary(summary, step)
 
   def image(self, tag, image, step=None):
     """Saves RGB image summary from onp.ndarray [H,W], [H,W,1], or [H,W,3].
@@ -121,9 +134,9 @@ def image(self, tag, image, step=None):
     """
     image = onp.array(image)
     if step is None:
-      step = self.step
+      step = self._step
     else:
-      self.step = step
+      self._step = step
     if len(onp.shape(image)) == 2:
       image = image[:, :, onp.newaxis]
     if onp.shape(image)[-1] == 1:
@@ -136,7 +149,7 @@ def image(self, tag, image, step=None):
         height=image.shape[0],
         width=image.shape[1])
     summary = Summary(value=[Summary.Value(tag=tag, image=image_summary)])
-    self.writer.add_summary(summary, step)
+    self.add_summary(summary, step)
 
   def images(self, tag, images, step=None, rows=None, cols=None):
     """Saves (rows, cols) tiled images from onp.ndarray.
@@ -155,9 +168,9 @@ def images(self, tag, images, step=None, rows=None, cols=None):
     """
     images = onp.array(images)
     if step is None:
-      step = self.step
+      step = self._step
     else:
-      self.step = step
+      self._step = step
     n_images = onp.shape(images)[0]
     if rows is None and cols is None:
       rows = 1
@@ -179,9 +192,9 @@ def plot(self, tag, mpl_plt, step=None, close_plot=True):
       close_plot: bool: automatically closes plot
     """
     if step is None:
-      step = self.step
+      step = self._step
     else:
-      self.step = step
+      self._step = step
     fig = mpl_plt.get_current_fig_manager()
     img_w, img_h = fig.canvas.get_width_height()
     image_buf = io.BytesIO()
@@ -192,7 +205,7 @@ def plot(self, tag, mpl_plt, step=None, close_plot=True):
         height=img_h,
         width=img_w)
     summary = Summary(value=[Summary.Value(tag=tag, image=image_summary)])
-    self.writer.add_summary(summary, step)
+    self.add_summary(summary, step)
     if close_plot:
       mpl_plt.close()
 
@@ -209,9 +222,9 @@ def audio(self, tag, audiodata, step=None, sample_rate=44100):
     """
     audiodata = onp.array(audiodata)
     if step is None:
-      step = self.step
+      step = self._step
     else:
-      self.step = step
+      self._step = step
     audiodata = onp.clip(onp.squeeze(audiodata), -1, 1)
     if audiodata.ndim != 1:
       raise ValueError('Audio data must be 1D.')
@@ -233,7 +246,7 @@ def audio(self, tag, audiodata, step=None, sample_rate=44100):
         encoded_audio_string=encoded_audio_bytes,
         content_type='audio/wav')
     summary = Summary(value=[Summary.Value(tag=tag, audio=audio)])
-    self.writer.add_summary(summary, step)
+    self.add_summary(summary, step)
 
   def histogram(self, tag, values, bins, step=None):
     """Saves histogram of values.
@@ -245,9 +258,9 @@ def histogram(self, tag, values, bins, step=None):
       step: int: training step
     """
     if step is None:
-      step = self.step
+      step = self._step
     else:
-      self.step = step
+      self._step = step
     values = onp.array(values)
     bins = onp.array(bins)
     values = onp.reshape(values, -1)
@@ -271,7 +284,7 @@ def histogram(self, tag, values, bins, step=None):
         bucket_limit=limits.tolist(),
         bucket=counts.tolist())
     summary = Summary(value=[Summary.Value(tag=tag, histo=histo)])
-    self.writer.add_summary(summary, step)
+    self.add_summary(summary, step)
 
   def text(self, tag, textdata, step=None):
     """Saves a text summary.
@@ -283,9 +296,9 @@ def text(self, tag, textdata, step=None):
     Note: markdown formatting is rendered by tensorboard.
     """
     if step is None:
-      step = self.step
+      step = self._step
     else:
-      self.step = step
+      self._step = step
     smd = SummaryMetadata(
         plugin_data=SummaryMetadata.PluginData(plugin_name='text'))
     if isinstance(textdata, (str, bytes)):
@@ -306,7 +319,7 @@ def text(self, tag, textdata, step=None):
             shape=(datashape[0], datashape[1]))
     summary = Summary(
         value=[Summary.Value(tag=tag, metadata=smd, tensor=tensor)])
-    self.writer.add_summary(summary, step)
+    self.add_summary(summary, step)
 
 
 # Copied from gin/tf/utils.py:GinConfigSaverHook
diff --git a/tensor2tensor/trax/jaxboard_demo.py b/tensor2tensor/trax/jaxboard_demo.py
deleted file mode 100644
index 53b740415..000000000
--- a/tensor2tensor/trax/jaxboard_demo.py
+++ /dev/null
@@ -1,122 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-r"""Jaxboard Summary Types Demo."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-from absl import app
-from absl import flags
-
-import warnings  # pylint: disable=g-bad-import-order
-import matplotlib as mpl
-with warnings.catch_warnings():
-  warnings.simplefilter('ignore')
-  mpl.use('Agg')
-# pylint: disable=g-import-not-at-top
-from matplotlib import pyplot as plt
-import numpy as onp
-
-from tensor2tensor.jax import jaxboard
-
-flags.DEFINE_string('tb_log_dir', '/tmp/tb_logs',
-                    'Path where we store summaries.')
-FLAGS = flags.FLAGS
-
-
-def demo():
-  """Run Summary Types Demo."""
-  sw = jaxboard.SummaryWriter(
-      os.path.join(FLAGS.tb_log_dir, 'demo', 'summarydemo'))
-
-  # Scalars.  We pass in step explicitly.
-  for i, v in enumerate(onp.sin(onp.linspace(0.0, 1.0, 100))):
-    sw.scalar('summarydemo_loss', v + 0.1 * onp.random.random(), step=i)
-
-  # SummaryWriter stores last step variable passed in, we can also set it
-  # explicitly for a set of exports to avoid providing the kwarg.
-  sw.step = 2
-
-  # Images. [H,W] or [H,W,C] with C = 1 or 3
-  sw.image('pic_c0', onp.random.random((100, 100)))
-  sw.image('pic_c1', onp.random.random((100, 100, 1)))
-  sw.image('pic_c3', onp.random.random((100, 100, 3)))
-
-  # Tiled sets of images. Must be [N,H,W,C] with C = 1 or 3
-  bw_tiles = onp.stack([
-      0.1 * onp.random.random((100, 100, 1)), 0.2 * onp.random.random(
-          (100, 100, 1)), 0.4 * onp.random.random((100, 100, 1)),
-      0.8 * onp.random.random((100, 100, 1))
-  ])
-  sw.images('pics_tiled_c1', bw_tiles, rows=2, cols=2)
-  clr_tiles = onp.stack([
-      0.1 * onp.random.random((100, 100, 3)), 0.2 * onp.random.random(
-          (100, 100, 3)), 0.4 * onp.random.random((100, 100, 3)),
-      0.8 * onp.random.random((100, 100, 3))
-  ])
-  sw.images('pics_tiled_c3', clr_tiles, rows=2, cols=2)
-
-  # Matplotlib plots. Just pass in prepared stateful pyplot object.
-  # -- scatter
-  plt.figure(figsize=(4, 4))
-  plt.scatter(
-      onp.random.randint(size=(10,), low=0, high=10),
-      onp.random.randint(size=(10,), low=0, high=10))
-  sw.plot('plot1', plt)
-
-  # -- imshow
-  plt.figure(figsize=(4, 4))
-  plt.imshow(
-      onp.random.randint(size=(50, 50, 3), low=0, high=255),
-      cmap='viridis',
-      interpolation='nearest')
-  sw.plot('plot2', plt)
-
-  # Audio.
-  t = onp.linspace(0, 1.0, 44100)
-  sinwave = (
-      0.1 * onp.sin(440. * onp.pi * t) *
-      # slow ramp-up to prevent 'pop'
-      onp.where(t > 0.2, 1.0, t / 0.2))
-  sw.audio('audio', sinwave)
-
-  # Text.
-  # -- tensorboard text plugin supports some markdown formatting!
-  sw.text('text', 'Colorless _green_ __ideas__ sleep furiously.')
-
-  # -- 1d/2d arrays of strings rendered as tables by plugin:
-  sw.text('text1d', ['Colorless', 'green', 'ideas', 'sleep', 'furiously.'])
-  sw.text('text2d', onp.array([['foo', 'bar'], ['baz', 'qup']]))
-
-  # Histograms / Distributions.
-  # (bins can be int or array - passed into onp.histogram bins arg)
-  sw.histogram('histo', onp.random.normal(size=(1000,)), 25, step=3)
-  sw.histogram('histo', onp.random.normal(size=(1000,)), 25, step=4)
-  sw.histogram('histo', onp.random.normal(size=(1000,)), 25, step=5)
-
-  # Fin.
-  sw.close()
-
-
-def main(argv):
-  del argv
-  demo()
-
-
-if __name__ == '__main__':
-  app.run(main)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index b12bec351..9f0ae0334 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -428,8 +428,8 @@ def train(output_dir,
                                        lr_fun, num_devices)
 
     # Flush summary writers
-    train_sw.writer.flush()
-    eval_sw.writer.flush()
+    train_sw.flush()
+    eval_sw.flush()
 
   step_log(step, "Training done")
   return State(params=params, step=step, history=history)

From 2f32f612dc77922027b0ce460fcfcc65d81683c4 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 10 Apr 2019 17:31:11 -0700
Subject: [PATCH 1889/2720] Make transformer run again after API changes in
 jax/stax.

PiperOrigin-RevId: 242974692
---
 tensor2tensor/trax/configs/transformer_lm1b_8gb.gin |  4 ++--
 tensor2tensor/trax/models/mlp.py                    |  2 +-
 tensor2tensor/trax/models/resnet.py                 |  2 +-
 tensor2tensor/trax/stax/attention.py                | 13 ++++++-------
 tensor2tensor/trax/trax.py                          |  1 +
 5 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index 3aeeb77ca..9f00817d9 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -26,12 +26,12 @@ MultifactorSchedule.warmup_steps = 8000
 # Parameters for preprocess_fun:
 # ==============================================================================
 shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
-lm1b.max_target_length = 511
+lm1b_preprocess.max_target_length = 511
 
 # Parameters for train:
 # ==============================================================================
 train.eval_frequency = 1000
-train.eval_steps = 5
+train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
 train.run_debug_step = False
diff --git a/tensor2tensor/trax/models/mlp.py b/tensor2tensor/trax/models/mlp.py
index 62b679199..a2f4ea351 100644
--- a/tensor2tensor/trax/models/mlp.py
+++ b/tensor2tensor/trax/models/mlp.py
@@ -19,7 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from jax.experimental import stax
+import tensor2tensor.trax.stax as stax
 
 
 def MLP(num_hidden_layers=2,
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index b731861b7..0b52145eb 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -19,7 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from jax.experimental import stax
+import tensor2tensor.trax.stax as stax
 
 
 def ConvBlock(kernel_size, filters, strides):
diff --git a/tensor2tensor/trax/stax/attention.py b/tensor2tensor/trax/stax/attention.py
index 309f30140..d3a24e754 100644
--- a/tensor2tensor/trax/stax/attention.py
+++ b/tensor2tensor/trax/stax/attention.py
@@ -20,7 +20,6 @@
 
 from jax import random
 import numpy as onp
-import numpy.random as npr
 
 from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.stax import stax_base as stax
@@ -73,13 +72,13 @@ def prepare_paired_sequence_batch(source, target_in, pad=0):
           source_mask, target_mask, memory_mask, ntokens)
 
 
-def xavier_uniform(out_dim=0, in_dim=1, rng=npr):
+def xavier_uniform(out_dim=0, in_dim=1):
   """An initializer function for random uniform xavier-scaled coefficients."""
-  def init(shape):
+  def init(rng, shape):
     fan_in, fan_out = shape[in_dim], shape[out_dim]
     std = np.sqrt(2.0 / (fan_in + fan_out))
-    a = onp.sqrt(3.0) * std
-    return rng.uniform(low=-a, high=a, size=shape).astype('float32')
+    a = np.sqrt(3.0) * std
+    return random.uniform(rng, shape, minval=-a, maxval=a)
   return init
 
 
@@ -104,9 +103,9 @@ def apply_fun(params, inputs, **kwargs):
 
 def Embedding(feature_depth, vocab_size):  # pylint: disable=invalid-name
   """Layer constructor function for a dense embedding layer."""
-  def init_fun(_, input_shape):
+  def init_fun(rng, input_shape):
     output_shape = tuple(input_shape) + (feature_depth,)
-    dense_embedding = xavier_uniform()((vocab_size, feature_depth))
+    dense_embedding = xavier_uniform()(rng, (vocab_size, feature_depth))
     return output_shape, dense_embedding
   def apply_fun(params, inputs, **kwargs):
     del kwargs
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 9f0ae0334..e44a4a6a4 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -151,6 +151,7 @@ def evaluate_train_and_eval(step, inputs, predict_fun, eval_steps, rng,
     log_metrics(train_metrics, train_sw, "train", step, history=history)
   if eval_sw:
     log_metrics(eval_metrics, eval_sw, "eval", step, history=history)
+  step_log(step, "Finished evaluation")
   return train_metrics, eval_metrics
 
 

From cccafb00bd2c0b7358057da071161df671a1037b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 11 Apr 2019 09:53:28 -0700
Subject: [PATCH 1890/2720] [trax] update import of lax.psum

PiperOrigin-RevId: 243088280
---
 tensor2tensor/trax/trax.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index e44a4a6a4..278ee85cc 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -32,6 +32,7 @@
 import gin
 
 import jax
+from jax import lax_parallel as lax_para
 from jax import random as jax_random
 import numpy
 import six
@@ -260,7 +261,7 @@ def mapped_update(i, opt_state, batch, rng):
     _, opt_update = optimizer(lr_fun)
     params = trax_opt.get_params(opt_state)
     grads = backend.grad(loss_fun)(params, batch, predict_fun, rng)
-    grads = jax.tree_util.tree_map(lambda g: jax.lax.psum(g, "batch"), grads)
+    grads = jax.tree_util.tree_map(lambda g: lax_para.psum(g, "batch"), grads)
     return opt_update(i, grads, opt_state)
 
   def update(i, opt_state, batch, rng):

From 0e5160e58c2fe3ae658ed368ba892b8f001ac903 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 11 Apr 2019 14:24:38 -0700
Subject: [PATCH 1891/2720] First trax Transformer config for TPU donuts,
 tuning to follow.

PiperOrigin-RevId: 243141831
---
 .../trax/configs/transformer_lm1b_tpu.gin     | 54 +++++++++++++++++++
 tensor2tensor/trax/trax.py                    |  6 ++-
 2 files changed, 59 insertions(+), 1 deletion(-)
 create mode 100644 tensor2tensor/trax/configs/transformer_lm1b_tpu.gin

diff --git a/tensor2tensor/trax/configs/transformer_lm1b_tpu.gin b/tensor2tensor/trax/configs/transformer_lm1b_tpu.gin
new file mode 100644
index 000000000..2681e23cc
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_lm1b_tpu.gin
@@ -0,0 +1,54 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size = 256
+batch_fun.eval_batch_size = 128
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_languagemodel_lm1b32k'
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 0.1
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
+lm1b_preprocess.max_target_length = 511
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 1000
+train.eval_steps = 10
+train.jit_eval = False  # Eval lengths vary a lot, compiling each time is slow.
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerLM
+train.run_debug_step = False
+train.train_steps = 500000
+
+# Parameters for train_and_eval_batches:
+# ==============================================================================
+train_and_eval_batches.input_name = 'targets'
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.dropout = 0.1
+TransformerLM.feature_depth = 512
+TransformerLM.feedforward_depth = 2048
+TransformerLM.max_len = 512
+TransformerLM.mode = 'train'
+TransformerLM.num_heads = 8
+TransformerLM.num_layers = 6
+TransformerLM.vocab_size = 32000
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 278ee85cc..8f896f119 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -299,6 +299,7 @@ def train(output_dir,
           eval_frequency=100,
           num_devices=None,
           random_seed=None,
+          jit_eval=True,
           run_debug_step=False):
   """Train the model on the inputs.
 
@@ -319,6 +320,7 @@ def train(output_dir,
       steps). If None or 0, eval disabled.
     num_devices: how many devices to use (if None, default, use all available)
     random_seed: the random seed to use; time/os dependent if None (default).
+    jit_eval: whether to compile the evaulation function (true by default).
     run_debug_step: bool, if True, will run the model and loss without @jit for
       one step.
 
@@ -352,7 +354,9 @@ def train(output_dir,
     opt_state = jax.replicate(opt_state)
 
   # jit model_predict and update so they're fast
-  jit_model_predict = backend.jit(model_predict)  # for evaluation
+  jit_model_predict = model_predict
+  if jit_eval:
+    jit_model_predict = backend.jit(model_predict)  # for evaluation
   jit_update_fun = _jit_update_fun(model_predict, loss_fun, optimizer, lr_fun,
                                    num_devices)
 

From bb074176e504d895bcb764e8782d760841684f7d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 11 Apr 2019 14:36:36 -0700
Subject: [PATCH 1892/2720] Modify Flatten to take in an optional argument
 `num_axis_to_keep` of the leading axis to keep and flattens the rest.

PiperOrigin-RevId: 243144265
---
 tensor2tensor/trax/models/mlp.py     |  2 +-
 tensor2tensor/trax/models/resnet.py  |  4 ++--
 tensor2tensor/trax/stax/slax_test.py | 25 ++++++++++++++++++++++++-
 tensor2tensor/trax/stax/stax_base.py | 16 +++++++++++-----
 4 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/trax/models/mlp.py b/tensor2tensor/trax/models/mlp.py
index a2f4ea351..c1296120a 100644
--- a/tensor2tensor/trax/models/mlp.py
+++ b/tensor2tensor/trax/models/mlp.py
@@ -26,7 +26,7 @@ def MLP(num_hidden_layers=2,
         hidden_size=512,
         activation_fn=stax.Relu,
         num_output_classes=10):
-  layers = [stax.Flatten]
+  layers = [stax.Flatten()]
   layers += [stax.Dense(hidden_size), activation_fn] * num_hidden_layers
   layers += [stax.Dense(num_output_classes), stax.LogSoftmax]
   return stax.serial(*layers)
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index 0b52145eb..dc0db56d9 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -90,7 +90,7 @@ def Resnet50(hidden_size=64, num_output_classes=1001):
       ConvBlock(3, [8 * hidden_size, 8 * hidden_size, 32*hidden_size], (2, 2)),
       IdentityBlock(3, [8 * hidden_size, 8 * hidden_size]),
       IdentityBlock(3, [8 * hidden_size, 8 * hidden_size]),
-      stax.AvgPool((7, 7)), stax.Flatten,
+      stax.AvgPool((7, 7)), stax.Flatten(),
       stax.Dense(num_output_classes), stax.LogSoftmax)
 
 
@@ -130,5 +130,5 @@ def WideResnet(num_blocks=3, hidden_size=64, num_output_classes=10):
       WideResnetGroup(num_blocks, hidden_size),
       WideResnetGroup(num_blocks, hidden_size * 2, (2, 2)),
       WideResnetGroup(num_blocks, hidden_size * 4, (2, 2)), stax.BatchNorm(),
-      stax.Relu, stax.AvgPool((8, 8)), stax.Flatten,
+      stax.Relu, stax.AvgPool((8, 8)), stax.Flatten(),
       stax.Dense(num_output_classes), stax.LogSoftmax)
diff --git a/tensor2tensor/trax/stax/slax_test.py b/tensor2tensor/trax/stax/slax_test.py
index dd15a378b..25ef39bf7 100644
--- a/tensor2tensor/trax/stax/slax_test.py
+++ b/tensor2tensor/trax/stax/slax_test.py
@@ -40,11 +40,12 @@ def check_shape_agreement(test_case, init_fun, apply_fun, input_shape):
   inputs = random_inputs(onp.random.RandomState(0), input_shape)
   result = apply_fun(params, inputs, rng=rng_key2)
   test_case.assertEqual(result.shape, result_shape)
+  return result_shape
 
 
 def check_staxlayer(test_case, staxlayer, input_shape):
   init_fun, apply_fun = staxlayer
-  check_shape_agreement(test_case, init_fun, apply_fun, input_shape)
+  return check_shape_agreement(test_case, init_fun, apply_fun, input_shape)
 
 
 # Helper functions for testing Lambda wrapper against functions involving
@@ -94,6 +95,28 @@ def _build_combinator_tree(input_treespec, in_vars):
 
 class SlaxTest(absltest.TestCase):
 
+  def test_flatten_n(self):
+    input_shape = (29, 87, 10, 20, 30)
+
+    actual_shape = check_staxlayer(self, stax.Flatten(1), input_shape)
+    self.assertEqual(actual_shape, (29, 87 * 10 * 20 * 30))
+
+    actual_shape = check_staxlayer(self, stax.Flatten(2), input_shape)
+    self.assertEqual(actual_shape, (29, 87, 10 * 20 * 30))
+
+    actual_shape = check_staxlayer(self, stax.Flatten(3), input_shape)
+    self.assertEqual(actual_shape, (29, 87, 10, 20 * 30))
+
+    actual_shape = check_staxlayer(self, stax.Flatten(4), input_shape)
+    self.assertEqual(actual_shape, (29, 87, 10, 20, 30))
+
+    # Not enough dimensions.
+    with self.assertRaises(ValueError):
+      check_staxlayer(self, stax.Flatten(5), input_shape)
+
+    with self.assertRaises(ValueError):
+      check_staxlayer(self, stax.Flatten(6), input_shape)
+
   # Lambdas replace the staxlayer input stream with a placeholder that
   # _should_ break any use of unbound variables in the input stream.
   def testLambda_forbidden_access(self):
diff --git a/tensor2tensor/trax/stax/stax_base.py b/tensor2tensor/trax/stax/stax_base.py
index b368f230f..35f8465e9 100644
--- a/tensor2tensor/trax/stax/stax_base.py
+++ b/tensor2tensor/trax/stax/stax_base.py
@@ -209,15 +209,21 @@ def rescale(outputs, inputs):
 AvgPool = _pooling_layer(lax.add, 0., _normalize_by_window_size)
 
 
-def Flatten():
-  """Layer construction function for flattening all but the leading dim."""
+def Flatten(num_axis_to_keep=1):
+  """Layer construction function for flattening all but the leading dims."""
   def init_fun(rng, input_shape):
-    output_shape = input_shape[0], reduce(op.mul, input_shape[1:], 1)
+    del rng
+    if num_axis_to_keep >= len(input_shape):
+      raise ValueError(
+          "num_axis_to_keep[%d] should be less than input's rank[%d]" %
+          (num_axis_to_keep, len(input_shape)))
+    output_shape = tuple(input_shape[:num_axis_to_keep]) + (
+        reduce(op.mul, input_shape[num_axis_to_keep:], 1),)
     return output_shape, ()
   def apply_fun(params, inputs, **kwargs):
-    return np.reshape(inputs, (inputs.shape[0], -1))
+    del params, kwargs
+    return np.reshape(inputs, (inputs.shape[:num_axis_to_keep] + (-1,)))
   return init_fun, apply_fun
-Flatten = Flatten()
 
 
 def Identity():

From ada626e6c879fbece1358c1d29e3f453af42f94d Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 11 Apr 2019 16:08:25 -0700
Subject: [PATCH 1893/2720] Abstract random functions in trax backend.

PiperOrigin-RevId: 243162092
---
 tensor2tensor/trax/backend.py        | 46 +++++++++++++++++++++++++---
 tensor2tensor/trax/stax/stax_base.py | 19 +++++++-----
 tensor2tensor/trax/trax.py           |  9 +++---
 3 files changed, 58 insertions(+), 16 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index c8634cf2c..ffd078501 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -22,10 +22,14 @@
 import gin
 
 import jax
+from jax import random as jax_random
 import jax.numpy as jnp
 import jax.scipy.special as jax_special
 import numpy as onp
 
+import tensorflow as tf
+
+
 
 @gin.configurable()
 def backend(name="jax"):
@@ -40,6 +44,11 @@ def backend(name="jax"):
     "jit": jax.jit,
     "grad": jax.grad,
     "pmap": jax.pmap,
+    "random_uniform": jax_random.uniform,
+    "random_normal": jax_random.normal,
+    "random_bernoulli": jax_random.bernoulli,
+    "random_get_prng": jax_random.PRNGKey,
+    "random_split": jax_random.split,
 }
 
 
@@ -49,10 +58,6 @@ def backend(name="jax"):
 }
 
 
-# TODO(lukaszkaiser): make this lazy so we can switch backends on the fly.
-numpy = backend()["np"]
-
-
 def logsumexp(*args, **kwargs):
   return backend()["logsumexp"](*args, **kwargs)
 
@@ -67,3 +72,36 @@ def grad(*args, **kwargs):
 
 def pmap(*args, **kwargs):
   return backend()["pmap"](*args, **kwargs)
+
+
+# For numpy and random modules, we need to call "backend()" lazily, only when
+# the function is called -- so that it can be set by gin configs.
+# (Otherwise, backend() is called on import before gin-config is parsed.)
+# To do that, we make objects to encapsulated these modules.
+
+
+class RandomBackend(object):
+  """Backend providing random functions."""
+
+  def get_prng(self, seed):
+    return backend()["random_get_prng"](seed)
+
+  def split(self, prng, num=2):
+    return backend()["random_split"](prng, num)
+
+  def uniform(self, *args, **kwargs):
+    return backend()["random_uniform"](*args, **kwargs)
+
+  def normal(self, *args, **kwargs):
+    return backend()["random_normal"](*args, **kwargs)
+
+  def bernoulli(self, *args, **kwargs):
+    return backend()["random_bernoulli"](*args, **kwargs)
+
+
+random = RandomBackend()
+
+# TODO(lukaszkaiser): make this lazy as random above so gin-config works.
+numpy = backend()["np"]
+
+
diff --git a/tensor2tensor/trax/stax/stax_base.py b/tensor2tensor/trax/stax/stax_base.py
index 35f8465e9..aa1ceb2e4 100644
--- a/tensor2tensor/trax/stax/stax_base.py
+++ b/tensor2tensor/trax/stax/stax_base.py
@@ -27,7 +27,6 @@
 import operator as op
 
 from jax import lax
-from jax import random
 
 import numpy as onp
 from six.moves import reduce
@@ -78,7 +77,7 @@ def fastvar(x, axis, keepdims):
 def randn(stddev=1e-2):
   """An initializer function for random normal coefficients."""
   def init(rng, shape):
-    return (stddev * random.normal(rng, shape)).astype('float32')
+    return (stddev * backend.random.normal(rng, shape)).astype('float32')
   return init
 
 
@@ -88,7 +87,7 @@ def init(rng, shape):
     fan_in, fan_out = shape[in_dim], shape[out_dim]
     size = onp.prod(onp.delete(shape, [in_dim, out_dim]))
     std = scale / np.sqrt((fan_in + fan_out) / 2. * size)
-    return (std * random.normal(rng, shape)).astype('float32')
+    return (std * backend.random.normal(rng, shape)).astype('float32')
   return init
 
 
@@ -275,7 +274,7 @@ def apply_fun(params, inputs, **kwargs):  # pylint: disable=missing-docstring
              'jax.random.PRNGKey value.')
       raise ValueError(msg)
     if mode == 'train':
-      keep = random.bernoulli(rng, rate, inputs.shape)
+      keep = backend.random.bernoulli(rng, rate, inputs.shape)
       return np.where(keep, inputs / rate, 0)
     else:
       return inputs
@@ -299,13 +298,15 @@ def serial(*layers):
   def init_fun(rng, input_shape):
     params = []
     for init_fun in init_funs:
-      rng, layer_rng = random.split(rng)
+      rng, layer_rng = backend.random.split(rng)
       input_shape, param = init_fun(layer_rng, input_shape)
       params.append(param)
     return input_shape, params
   def apply_fun(params, inputs, **kwargs):
     rng = kwargs.pop('rng', None)
-    rngs = random.split(rng, nlayers) if rng is not None else (None,) * nlayers
+    rngs = (None,) * nlayers
+    if rng is not None:
+      rngs = backend.random.split(rng, nlayers)
     for fun, param, rng in zip(apply_funs, params, rngs):
       inputs = fun(param, inputs, rng=rng, **kwargs)
     return inputs
@@ -328,12 +329,14 @@ def parallel(*layers):
   nlayers = len(layers)
   init_funs, apply_funs = zip(*layers)
   def init_fun(rng, input_shape):
-    rngs = random.split(rng, nlayers)
+    rngs = backend.random.split(rng, nlayers)
     return zip(*[init(rng, shape) for init, rng, shape
                  in zip(init_funs, rngs, input_shape)])
   def apply_fun(params, inputs, **kwargs):
     rng = kwargs.pop('rng', None)
-    rngs = random.split(rng, nlayers) if rng is not None else (None,) * nlayers
+    rngs = (None,) * nlayers
+    if rng is not None:
+      rngs = backend.random.split(rng, nlayers)
     return [f(p, x, rng=r, **kwargs)
             for f, p, x, r in zip(apply_funs, params, inputs, rngs)]
   return init_fun, apply_fun
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 8f896f119..e4977b82c 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -32,8 +32,7 @@
 import gin
 
 import jax
-from jax import lax_parallel as lax_para
-from jax import random as jax_random
+from jax import lax_parallel as lax_parallel
 import numpy
 import six
 
@@ -44,6 +43,7 @@
 from tensor2tensor.trax import learning_rate as lr
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax.backend import numpy as np
+from tensor2tensor.trax.backend import random as jax_random
 import tensor2tensor.trax.stax as stax
 
 import tensorflow as tf
@@ -204,7 +204,7 @@ def get_random_number_generator_and_set_seed(seed=None):
     seed = random.randint(0, 2**31 - 1)
   tf.set_random_seed(seed)
   numpy.random.seed(seed)
-  return jax_random.PRNGKey(seed)
+  return jax_random.get_prng(seed)
 
 
 # TODO(trax):
@@ -261,7 +261,8 @@ def mapped_update(i, opt_state, batch, rng):
     _, opt_update = optimizer(lr_fun)
     params = trax_opt.get_params(opt_state)
     grads = backend.grad(loss_fun)(params, batch, predict_fun, rng)
-    grads = jax.tree_util.tree_map(lambda g: lax_para.psum(g, "batch"), grads)
+    grads = jax.tree_util.tree_map(
+        lambda g: lax_parallel.psum(g, "batch"), grads)
     return opt_update(i, grads, opt_state)
 
   def update(i, opt_state, batch, rng):

From 4432a90f202f861a016c43be05d22e0dbbd5c400 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 12 Apr 2019 10:12:05 -0700
Subject: [PATCH 1894/2720] Implement fast decoding for the Evolved
 Transformer.

PiperOrigin-RevId: 243285212
---
 tensor2tensor/models/evolved_transformer.py   | 218 ++++++++++++++----
 .../models/evolved_transformer_test.py        | 211 ++++++++++++++++-
 tensor2tensor/models/transformer.py           | 183 +++++++--------
 3 files changed, 465 insertions(+), 147 deletions(-)

diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 62098e004..4d27d2416 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -29,6 +29,16 @@
 
 import tensorflow as tf
 
+# pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.ops import inplace_ops
+# pylint: enable=g-direct-tensorflow-import
+
+_CONV_BRANCHES_NAME = "conv_branches"
+_FIRST_ATTEND_TO_ENCODER_NAME = "first_attend_to_encoder"
+_SECOND_ATTEND_TO_ENCODER_NAME = "second_attend_to_encoder"
+_SIXTEEN_HEAD_ATTENTION_NAME = "16_head_self_attention"
+_VANILLA_ATTENTION_NAME = "self_attention"
+
 
 @registry.register_model
 class EvolvedTransformer(transformer.Transformer):
@@ -38,31 +48,7 @@ def __init__(self, *args, **kwargs):
     super(EvolvedTransformer, self).__init__(*args, **kwargs)
     self._encoder_function = evolved_transformer_encoder
     self._decoder_function = evolved_transformer_decoder
-
-  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha,
-                   use_tpu):
-    """Forced slow beam decode because cache is not supported.
-
-    Args:
-      features: an map of string to `Tensor`.
-      decode_length: an integer.  How many additional timesteps to decode.
-      beam_size: number of beams.
-      top_beams: an integer. How many of the beams to return.
-      alpha: Float that controls the length penalty. larger the alpha, stronger
-        the preference for longer translations.
-      use_tpu: Whether or not TPU is being used.
-
-    Returns:
-      A dict of decoding results {
-          "outputs": integer `Tensor` of decoded ids of shape
-              [batch_size, <= decode_length] if beam_size == 1 or
-              [batch_size, top_beams, <= decode_length].
-          "scores": decoding log probs from the beam search,
-              None if using greedy decoding (beam_size=1).
-      }
-    """
-    return self._beam_decode_slow(features, decode_length, beam_size, top_beams,
-                                  alpha, use_tpu)
+    self._init_cache_fn = _init_evolved_transformer_cache
 
 
 def evolved_transformer_encoder(encoder_input,
@@ -252,7 +238,8 @@ def evolved_transformer_decoder(decoder_input,
     encoder_decoder_attention_bias: bias Tensor for encoder-decoder attention
       (see common_attention.attention_bias()).
     hparams: hyperparameters for model.
-    cache: Not supported.
+    cache: dict, containing tensors which are the results of previous
+      layers, used for fast decoding.
     decode_loop_step: An integer, step number of the decoding loop. Only used
       for inference on TPU.
     name: a string.
@@ -270,7 +257,7 @@ def evolved_transformer_decoder(decoder_input,
   Returns:
     Decoder output tensor.
   """
-  del cache, losses
+  del losses
 
   attention_dropout_broadcast_dims = (
       common_layers.comma_separated_string_to_integer_list(
@@ -278,16 +265,19 @@ def evolved_transformer_decoder(decoder_input,
 
   with tf.variable_scope(name):
     hidden_state = decoder_input
-    layer_cache = None
 
     for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
-      with tf.variable_scope("layer_%d" % layer):
+      layer_name = "layer_%d" % layer
+      layer_cache = cache[layer_name] if cache is not None else None
+      with tf.variable_scope(layer_name):
 
-        with tf.variable_scope("16_head_self_attention"):
+        with tf.variable_scope(_SIXTEEN_HEAD_ATTENTION_NAME):
           residual_state = hidden_state
           hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
 
           # Attention with at least 16 heads.
+          attention_cache = layer_cache[
+              _SIXTEEN_HEAD_ATTENTION_NAME] if layer_cache is not None else None
           left_state = common_attention.multihead_attention(
               hidden_state,
               None,
@@ -303,7 +293,7 @@ def evolved_transformer_decoder(decoder_input,
                   hparams.heads_share_relative_embedding),
               add_relative_to_values=hparams.add_relative_to_values,
               save_weights_to=save_weights_to,
-              cache=layer_cache,
+              cache=attention_cache,
               make_image_summary=make_image_summary,
               dropout_broadcast_dims=attention_dropout_broadcast_dims,
               max_length=hparams.get("max_length"),
@@ -313,7 +303,10 @@ def evolved_transformer_decoder(decoder_input,
               weight_dtype=hparams.get("weight_dtype", "float32"))
 
         if encoder_output is not None:
-          with tf.variable_scope("first_attend_to_encoder"):
+          with tf.variable_scope(_FIRST_ATTEND_TO_ENCODER_NAME):
+            attention_cache = (
+                layer_cache[_FIRST_ATTEND_TO_ENCODER_NAME]
+                if layer_cache is not None else None)
             right_state = common_attention.multihead_attention(
                 hidden_state,
                 encoder_output,
@@ -328,7 +321,7 @@ def evolved_transformer_decoder(decoder_input,
                     hparams.heads_share_relative_embedding),
                 add_relative_to_values=hparams.add_relative_to_values,
                 save_weights_to=save_weights_to,
-                cache=layer_cache,
+                cache=attention_cache,
                 make_image_summary=make_image_summary,
                 dropout_broadcast_dims=attention_dropout_broadcast_dims,
                 max_length=hparams.get("max_length"),
@@ -347,7 +340,7 @@ def evolved_transformer_decoder(decoder_input,
           hidden_state = common_layers.layer_postprocess(
               residual_state, left_state, hparams)
 
-        with tf.variable_scope("conv_branches"):
+        with tf.variable_scope(_CONV_BRANCHES_NAME):
           residual_state = hidden_state
           hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
 
@@ -357,8 +350,38 @@ def evolved_transformer_decoder(decoder_input,
                 tf.expand_dims(nonpadding, 2), [1, 1, hparams.hidden_size])
             hidden_state *= mask
 
-          # Shift inputs so that future tokens cannot be seen.
-          left_state = tf.pad(hidden_state, paddings=[[0, 0], [10, 0], [0, 0]])
+          # TODO(davidso): This needlessly recomputes past positions. Limit
+          # the module inputs to only include positions that are convered by
+          # the new output position's receptive field.
+          if layer_cache:
+            if decode_loop_step is None:
+              hidden_state = layer_cache[_CONV_BRANCHES_NAME] = tf.concat(
+                  [layer_cache[_CONV_BRANCHES_NAME], hidden_state], axis=1)
+              left_state = hidden_state
+              right_state = hidden_state
+
+            else:
+              # Inplace update is required for inference on TPU.
+              # Inplace_ops only supports inplace_update on the first dimension.
+              tmp = tf.transpose(
+                  layer_cache[_CONV_BRANCHES_NAME], perm=[1, 0, 2])
+              tmp = tf.expand_dims(tmp, axis=1)
+              tmp = inplace_ops.alias_inplace_update(
+                  tmp,
+                  decode_loop_step * tf.shape(hidden_state)[1],
+                  tf.transpose(hidden_state, perm=[1, 0, 2]))
+              tmp = tf.squeeze(tmp, axis=1)
+              hidden_state = layer_cache[_CONV_BRANCHES_NAME] = tf.transpose(
+                  tmp, perm=[1, 0, 2])
+
+              read_to_index = decode_loop_step + 1
+              left_state = hidden_state[:, :read_to_index, :]
+              right_state = hidden_state[:, :read_to_index, :]
+          else:
+            left_state = hidden_state
+            right_state = hidden_state
+
+          left_state = tf.pad(left_state, paddings=[[0, 0], [10, 0], [0, 0]])
           left_output_dim = int(hparams.hidden_size * 2)
           separable_conv_11x1 = tf.layers.SeparableConv1D(
               left_output_dim,
@@ -370,7 +393,7 @@ def evolved_transformer_decoder(decoder_input,
           left_state = tf.nn.dropout(left_state,
                                      1 - hparams.layer_prepostprocess_dropout)
 
-          right_state = tf.pad(hidden_state, paddings=[[0, 0], [6, 0], [0, 0]])
+          right_state = tf.pad(right_state, paddings=[[0, 0], [6, 0], [0, 0]])
           right_output_dim = int(hparams.hidden_size / 2)
           separable_conv_7x1_1 = tf.layers.SeparableConv1D(
               right_output_dim, 7, padding="VALID", name="separable_conv_7x1_1")
@@ -399,13 +422,18 @@ def evolved_transformer_decoder(decoder_input,
               name="separable_conv_7x1_2")
           hidden_state = separable_conv_7x1_2.apply(hidden_state)
 
+          if layer_cache:
+            hidden_state = hidden_state[:, -1:, :]
+
           hidden_state = common_layers.layer_postprocess(
               residual_state, hidden_state, hparams)
 
-        with tf.variable_scope("self_attention"):
+        with tf.variable_scope(_VANILLA_ATTENTION_NAME):
           residual_state = hidden_state
           hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
 
+          attention_cache = layer_cache[
+              _VANILLA_ATTENTION_NAME] if layer_cache is not None else None
           hidden_state = common_attention.multihead_attention(
               hidden_state,
               None,
@@ -421,7 +449,7 @@ def evolved_transformer_decoder(decoder_input,
                   hparams.heads_share_relative_embedding),
               add_relative_to_values=hparams.add_relative_to_values,
               save_weights_to=save_weights_to,
-              cache=layer_cache,
+              cache=attention_cache,
               make_image_summary=make_image_summary,
               dropout_broadcast_dims=attention_dropout_broadcast_dims,
               max_length=hparams.get("max_length"),
@@ -433,10 +461,13 @@ def evolved_transformer_decoder(decoder_input,
               residual_state, hidden_state, hparams)
 
         if encoder_output is not None:
-          with tf.variable_scope("second_attend_to_encoder"):
+          with tf.variable_scope(_SECOND_ATTEND_TO_ENCODER_NAME):
             residual_state = hidden_state
             hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
 
+            attention_cache = (
+                layer_cache[_SECOND_ATTEND_TO_ENCODER_NAME]
+                if layer_cache is not None else None)
             hidden_state = common_attention.multihead_attention(
                 hidden_state,
                 encoder_output,
@@ -451,7 +482,7 @@ def evolved_transformer_decoder(decoder_input,
                     hparams.heads_share_relative_embedding),
                 add_relative_to_values=hparams.add_relative_to_values,
                 save_weights_to=save_weights_to,
-                cache=layer_cache,
+                cache=attention_cache,
                 make_image_summary=make_image_summary,
                 dropout_broadcast_dims=attention_dropout_broadcast_dims,
                 max_length=hparams.get("max_length"),
@@ -481,6 +512,111 @@ def evolved_transformer_decoder(decoder_input,
     return common_layers.layer_preprocess(hidden_state, hparams)
 
 
+def _add_attend_to_encoder_cache(cache, attention_name, hparams, num_layers,
+                                 key_channels, value_channels,
+                                 vars_3d_num_heads, scope_prefix,
+                                 encoder_output):
+  """Add attend-to-encoder layers to cache."""
+  for layer in range(num_layers):
+    layer_name = "layer_%d" % layer
+    with tf.variable_scope("%sdecoder/%s/%s/multihead_attention" %
+                           (scope_prefix, layer_name, attention_name)):
+      k_encdec = common_attention.compute_attention_component(
+          encoder_output,
+          key_channels,
+          name="k",
+          vars_3d_num_heads=vars_3d_num_heads)
+      k_encdec = common_attention.split_heads(k_encdec, hparams.num_heads)
+      v_encdec = common_attention.compute_attention_component(
+          encoder_output,
+          value_channels,
+          name="v",
+          vars_3d_num_heads=vars_3d_num_heads)
+      v_encdec = common_attention.split_heads(v_encdec, hparams.num_heads)
+    cache[layer_name][attention_name] = {
+        "k_encdec": k_encdec,
+        "v_encdec": v_encdec
+    }
+  return cache
+
+
+def _init_evolved_transformer_cache(cache, hparams, batch_size,
+                                    attention_init_length, encoder_output,
+                                    encoder_decoder_attention_bias,
+                                    scope_prefix):
+  """Create the initial cache for Evolved Transformer fast decoding."""
+  key_channels = hparams.attention_key_channels or hparams.hidden_size
+  value_channels = hparams.attention_value_channels or hparams.hidden_size
+  num_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
+  vars_3d_num_heads = (
+      hparams.num_heads if hparams.get("attention_variables_3d") else 0)
+
+  # Add self-attentions.
+  if cache is None:
+    cache = {}
+  cache.update({
+      "layer_%d" % layer: {  # pylint: disable=g-complex-comprehension
+          _SIXTEEN_HEAD_ATTENTION_NAME: {
+              "k":
+                  common_attention.split_heads(
+                      tf.zeros(
+                          [batch_size, attention_init_length, key_channels]),
+                      max(16, hparams.num_heads)),
+              "v":
+                  common_attention.split_heads(
+                      tf.zeros(
+                          [batch_size, attention_init_length, value_channels]),
+                      max(16, hparams.num_heads)),
+          },
+          _VANILLA_ATTENTION_NAME: {
+              "k":
+                  common_attention.split_heads(
+                      tf.zeros(
+                          [batch_size, attention_init_length, key_channels]),
+                      hparams.num_heads),
+              "v":
+                  common_attention.split_heads(
+                      tf.zeros(
+                          [batch_size, attention_init_length, value_channels]),
+                      hparams.num_heads),
+          }
+      } for layer in range(num_layers)
+  })
+
+  # Add branched layers.
+  for layer in range(num_layers):
+    cache["layer_%d" % layer][_CONV_BRANCHES_NAME] = tf.zeros(
+        [batch_size, attention_init_length, hparams.hidden_size])
+
+  # Add encoder embedding attentions.
+  if encoder_output is not None:
+    cache = _add_attend_to_encoder_cache(
+        cache=cache,
+        attention_name=_FIRST_ATTEND_TO_ENCODER_NAME,
+        hparams=hparams,
+        num_layers=num_layers,
+        key_channels=key_channels,
+        value_channels=value_channels,
+        vars_3d_num_heads=vars_3d_num_heads,
+        scope_prefix=scope_prefix,
+        encoder_output=encoder_output)
+    cache = _add_attend_to_encoder_cache(
+        cache=cache,
+        attention_name=_SECOND_ATTEND_TO_ENCODER_NAME,
+        hparams=hparams,
+        num_layers=num_layers,
+        key_channels=key_channels,
+        value_channels=value_channels,
+        vars_3d_num_heads=vars_3d_num_heads,
+        scope_prefix=scope_prefix,
+        encoder_output=encoder_output)
+
+    cache["encoder_output"] = encoder_output
+    cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias
+
+  return cache
+
+
 # TODO(davidso): Update optimizer, learning rate, and decay to match paper.
 def add_evolved_transformer_hparams(hparams):
   """Add Evolved Transformer hparams.
diff --git a/tensor2tensor/models/evolved_transformer_test.py b/tensor2tensor/models/evolved_transformer_test.py
index ae29f1c83..8691d162c 100644
--- a/tensor2tensor/models/evolved_transformer_test.py
+++ b/tensor2tensor/models/evolved_transformer_test.py
@@ -32,12 +32,15 @@
 VOCAB_SIZE = 10
 
 
-def get_model():
-  hparams = transformer.transformer_tiny()
+def get_model(hparams, has_input=True):
   hparams.layer_prepostprocess_dropout = 0.0
+  hparams.hidden_size = 16
+  hparams.num_heads = 1
 
   p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE, VOCAB_SIZE,
                                                    hparams)
+  if not has_input:
+    del p_hparams.modality["inputs"]
   hparams.problem_hparams = p_hparams
 
   inputs = np.random.randint(VOCAB_SIZE, size=(BATCH_SIZE, INPUT_LENGTH, 1, 1))
@@ -46,8 +49,9 @@ def get_model():
   features = {
       "targets": tf.constant(targets, dtype=tf.int32, name="targets"),
       "target_space_id": tf.constant(1, dtype=tf.int32),
-      "inputs": tf.constant(inputs, dtype=tf.int32, name="inputs"),
   }
+  if has_input:
+    features["inputs"] = tf.constant(inputs, dtype=tf.int32, name="inputs")
 
   return (evolved_transformer.EvolvedTransformer(
       hparams, tf.estimator.ModeKeys.TRAIN, p_hparams), features)
@@ -56,13 +60,212 @@ def get_model():
 class EvolvedTransformerTest(tf.test.TestCase):
 
   def testEvolvedTransformer(self):
-    model, features = get_model()
+    model, features = get_model(hparams=transformer.transformer_tiny())
     logits, _ = model(features)
     with self.test_session() as session:
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
     self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, 1, 1, VOCAB_SIZE))
 
+  def testSlowVsFast(self):
+    model, features = get_model(transformer.transformer_tiny())
+
+    decode_length = 30
+
+    out_logits, _ = model(features)
+    out_logits = tf.squeeze(out_logits, axis=[2, 3])
+    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
+        logits=tf.reshape(out_logits, [-1, VOCAB_SIZE]),
+        labels=tf.reshape(features["targets"], [-1]))
+    loss = tf.reduce_mean(loss)
+    apply_grad = tf.train.AdamOptimizer(0.001).minimize(loss)
+
+    with self.test_session():
+      tf.global_variables_initializer().run()
+      for _ in range(10):
+        apply_grad.run()
+
+    model.set_mode(tf.estimator.ModeKeys.PREDICT)
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+      greedy_result = model._slow_greedy_infer(
+          features, decode_length)["outputs"]
+      greedy_result = tf.squeeze(greedy_result, axis=[2, 3])
+
+      fast_result = model._greedy_infer(features, decode_length)["outputs"]
+
+    with self.test_session():
+      greedy_res = greedy_result.eval()
+      fast_res = fast_result.eval()
+
+    self.assertEqual(fast_res.shape, (BATCH_SIZE, INPUT_LENGTH + decode_length))
+    self.assertAllClose(greedy_res, fast_res)
+
+  def testSlowVsFastNoInput(self):
+    model, features = get_model(
+        transformer.transformer_tiny(), has_input=False)
+
+    decode_length = 30
+
+    out_logits, _ = model(features)
+    out_logits = tf.squeeze(out_logits, axis=[2, 3])
+    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
+        logits=tf.reshape(out_logits, [-1, VOCAB_SIZE]),
+        labels=tf.reshape(features["targets"], [-1]))
+    loss = tf.reduce_mean(loss)
+    apply_grad = tf.train.AdamOptimizer(0.001).minimize(loss)
+
+    with self.test_session():
+      tf.global_variables_initializer().run()
+      for _ in range(10):
+        apply_grad.run()
+
+    model.set_mode(tf.estimator.ModeKeys.PREDICT)
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+      slow_result = model._slow_greedy_infer(
+          features, decode_length)["outputs"]
+      slow_result = tf.squeeze(slow_result, axis=[2, 3])
+
+      fast_result = model._greedy_infer(features, decode_length)["outputs"]
+
+    with self.test_session():
+      slow_res = slow_result.eval()
+      fast_res = fast_result.eval()
+
+    self.assertEqual(slow_res.shape, (BATCH_SIZE, decode_length))
+    self.assertAllClose(slow_res, fast_res)
+
+  def testBeamVsFast(self):
+    model, features = get_model(transformer.transformer_tiny())
+
+    decode_length = 30
+
+    out_logits, _ = model(features)
+    out_logits = tf.squeeze(out_logits, axis=[2, 3])
+    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
+        logits=tf.reshape(out_logits, [-1, VOCAB_SIZE]),
+        labels=tf.reshape(features["targets"], [-1]))
+    loss = tf.reduce_mean(loss)
+    apply_grad = tf.train.AdamOptimizer(0.001).minimize(loss)
+
+    with self.test_session():
+      tf.global_variables_initializer().run()
+      for _ in range(10):
+        apply_grad.run()
+
+    model.set_mode(tf.estimator.ModeKeys.PREDICT)
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+      beam_result = model._beam_decode_slow(
+          features,
+          decode_length,
+          beam_size=4,
+          top_beams=1,
+          alpha=1.0)["outputs"]
+
+      fast_result = model._beam_decode(
+          features,
+          decode_length,
+          beam_size=4,
+          top_beams=1,
+          alpha=1.0)["outputs"]
+
+    with self.test_session():
+      beam_res = beam_result.eval()
+      fast_res = fast_result.eval()
+
+    self.assertAllClose(beam_res, fast_res)
+
+  def _create_greedy_infer_model(self):
+    """Creates model for greedy inference testing.
+
+    Returns:
+      model: A t2t model.
+      features: An map of string to tensor.
+    """
+    model, features = get_model(transformer.transformer_tiny())
+
+    out_logits, _ = model(features)
+    out_logits = tf.squeeze(out_logits, axis=[2, 3])
+    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
+        logits=tf.reshape(out_logits, [-1, VOCAB_SIZE]),
+        labels=tf.reshape(features["targets"], [-1]))
+    loss = tf.reduce_mean(loss)
+    apply_grad = tf.train.AdamOptimizer(0.001).minimize(loss)
+
+    with self.test_session():
+      tf.global_variables_initializer().run()
+      for _ in range(10):
+        apply_grad.run()
+
+    model.set_mode(tf.estimator.ModeKeys.PREDICT)
+
+    return model, features
+
+  def testGreedySlowTPUVsNonTPU(self):
+    decode_length = 30
+
+    model, features = self._create_greedy_infer_model()
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+      slow_result_non_tpu = model._slow_greedy_infer(
+          features, decode_length)["outputs"]
+      slow_result_non_tpu = tf.squeeze(slow_result_non_tpu, axis=[2, 3])
+
+      slow_result_tpu = model._slow_greedy_infer_tpu(
+          features, decode_length)["outputs"]
+      slow_result_tpu = tf.squeeze(slow_result_tpu, axis=[2, 3])
+
+    with self.test_session():
+      slow_non_tpu_res = slow_result_non_tpu.eval()
+      slow_tpu_res = slow_result_tpu.eval()
+
+    self.assertEqual(slow_tpu_res.shape,
+                     (BATCH_SIZE, INPUT_LENGTH + decode_length))
+    self.assertAllClose(slow_tpu_res, slow_non_tpu_res)
+
+  def testGreedyFastTPUVsNonTPU(self):
+    decode_length = 30
+
+    model, features = self._create_greedy_infer_model()
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+      fast_result_non_tpu = model._greedy_infer(
+          features, decode_length, use_tpu=False)["outputs"]
+
+      fast_result_tpu = model._greedy_infer(
+          features, decode_length, use_tpu=True)["outputs"]
+
+    with self.test_session():
+      fast_non_tpu_res = fast_result_non_tpu.eval()
+      fast_tpu_res = fast_result_tpu.eval()
+
+    self.assertEqual(fast_tpu_res.shape,
+                     (BATCH_SIZE, INPUT_LENGTH + decode_length))
+    self.assertAllClose(fast_tpu_res, fast_non_tpu_res)
+
+  def testGreedyTPUSlowVsFast(self):
+    decode_length = 30
+
+    model, features = self._create_greedy_infer_model()
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+      slow_result = model._slow_greedy_infer_tpu(
+          features, decode_length)["outputs"]
+      slow_result = tf.squeeze(slow_result, axis=[2, 3])
+
+      fast_result = model._greedy_infer(
+          features, decode_length, use_tpu=True)["outputs"]
+
+    with self.test_session():
+      slow_res = slow_result.eval()
+      fast_res = fast_result.eval()
+
+    self.assertEqual(fast_res.shape,
+                     (BATCH_SIZE, INPUT_LENGTH + decode_length))
+    self.assertAllClose(fast_res, slow_res)
+
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index dc1cb1b40..40ecc977f 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -188,6 +188,7 @@ def __init__(self, *args, **kwargs):
     self.recurrent_memory_by_layer = None  # Override to enable recurrent memory
     self._encoder_function = transformer_encoder
     self._decoder_function = transformer_decoder
+    self._init_cache_fn = _init_transformer_cache
 
   def encode(self, inputs, target_space, hparams, features=None, losses=None):
     """Encode transformer inputs, see transformer_encode."""
@@ -594,6 +595,7 @@ def forced_logits():
         hparams=hparams,
         decode_length=decode_length,
         vocab_size=target_vocab_size,
+        init_cache_fn=self._init_cache_fn,
         beam_size=beam_size,
         top_beams=top_beams,
         alpha=alpha,
@@ -813,6 +815,7 @@ def forced_logits():
         hparams=hparams,
         decode_length=decode_length,
         vocab_size=target_vocab_size,
+        init_cache_fn=self._init_cache_fn,
         beam_size=beam_size,
         top_beams=top_beams,
         alpha=alpha,
@@ -826,12 +829,76 @@ def forced_logits():
     return ret
 
 
+def _init_transformer_cache(cache, hparams, batch_size, attention_init_length,
+                            encoder_output, encoder_decoder_attention_bias,
+                            scope_prefix):
+  """Create the initial cache for Transformer fast decoding."""
+  key_channels = hparams.attention_key_channels or hparams.hidden_size
+  value_channels = hparams.attention_value_channels or hparams.hidden_size
+  num_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
+  vars_3d_num_heads = (
+      hparams.num_heads if hparams.get("attention_variables_3d") else 0)
+
+  if cache is None:
+    cache = {}
+  cache.update({
+      "layer_%d" % layer: {  # pylint: disable=g-complex-comprehension
+          "k":
+              common_attention.split_heads(
+                  tf.zeros([batch_size,
+                            attention_init_length,
+                            key_channels]), hparams.num_heads),
+          "v":
+              common_attention.split_heads(
+                  tf.zeros([batch_size,
+                            attention_init_length,
+                            value_channels]), hparams.num_heads),
+      } for layer in range(num_layers)
+  })
+
+  # If `ffn_layer` is in `["dense_relu_dense" or "conv_hidden_relu"]`, then the
+  # cache key "f" won't be used, which means that the` shape of cache["f"]`
+  # won't be changed to
+  # `[beamsize*batch_size, decode_length, hparams.hidden_size]` and may cause
+  # error when applying `nest.map reshape function` on it.
+  if hparams.ffn_layer not in ["dense_relu_dense", "conv_hidden_relu"]:
+    for layer in range(num_layers):
+      cache["layer_%d" % layer]["f"] = tf.zeros(
+          [batch_size, 0, hparams.hidden_size])
+
+  if encoder_output is not None:
+    for layer in range(num_layers):
+      layer_name = "layer_%d" % layer
+      with tf.variable_scope(
+          "%sdecoder/%s/encdec_attention/multihead_attention" %
+          (scope_prefix, layer_name)):
+        k_encdec = common_attention.compute_attention_component(
+            encoder_output,
+            key_channels,
+            name="k",
+            vars_3d_num_heads=vars_3d_num_heads)
+        k_encdec = common_attention.split_heads(k_encdec, hparams.num_heads)
+        v_encdec = common_attention.compute_attention_component(
+            encoder_output,
+            value_channels,
+            name="v",
+            vars_3d_num_heads=vars_3d_num_heads)
+        v_encdec = common_attention.split_heads(v_encdec, hparams.num_heads)
+      cache[layer_name]["k_encdec"] = k_encdec
+      cache[layer_name]["v_encdec"] = v_encdec
+
+    cache["encoder_output"] = encoder_output
+    cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias
+  return cache
+
+
 def fast_decode_tpu(encoder_output,
                     encoder_decoder_attention_bias,
                     symbols_to_logits_fn,
                     hparams,
                     decode_length,
                     vocab_size,
+                    init_cache_fn=_init_transformer_cache,
                     beam_size=1,
                     top_beams=1,
                     alpha=1.0,
@@ -855,6 +922,7 @@ def fast_decode_tpu(encoder_output,
     hparams: Run hyperparameters.
     decode_length: An integer, how many additional timesteps to decode.
     vocab_size: Output vocabulary size.
+    init_cache_fn: Function that returns the initial cache dict.
     beam_size: An integer, number of beams.
     top_beams: An integer, how many of the beams to return.
     alpha: A float that controls the length penalty. Larger the alpha, stronger
@@ -883,57 +951,9 @@ def fast_decode_tpu(encoder_output,
   if encoder_output is not None:
     batch_size = common_layers.shape_list(encoder_output)[0]
 
-  key_channels = hparams.attention_key_channels or hparams.hidden_size
-  value_channels = hparams.attention_value_channels or hparams.hidden_size
-  num_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
-  vars_3d_num_heads = (
-      hparams.num_heads if hparams.get("attention_variables_3d") else 0)
-
-  cache = {
-      "layer_%d" % layer: {  # pylint: disable=g-complex-comprehension
-          "k":
-          common_attention.split_heads(
-              tf.zeros([batch_size, decode_length, key_channels]),
-              hparams.num_heads),
-          "v":
-          common_attention.split_heads(
-              tf.zeros([batch_size, decode_length, value_channels]),
-              hparams.num_heads),
-      } for layer in range(num_layers)
-  }
-
-  # If `ffn_layer` is in `["dense_relu_dense" or "conv_hidden_relu"]`, then the
-  # cache key "f" won't be used, which means that the` shape of cache["f"]`
-  # won't be changed to
-  # `[beamsize*batch_size, decode_length, hparams.hidden_size]` and may cause
-  # error when applying `nest.map reshape function` on it.
-  if hparams.ffn_layer not in ["dense_relu_dense", "conv_hidden_relu"]:
-    for layer in range(num_layers):
-      cache["layer_%d" % layer]["f"] = tf.zeros(
-          [batch_size, 0, hparams.hidden_size])
-
-  if encoder_output is not None:
-    for layer in range(num_layers):
-      layer_name = "layer_%d" % layer
-      with tf.variable_scope("%sdecoder/%s/encdec_attention/multihead_attention"
-                             % (scope_prefix, layer_name)):
-        k_encdec = common_attention.compute_attention_component(
-            encoder_output,
-            key_channels,
-            name="k",
-            vars_3d_num_heads=vars_3d_num_heads)
-        k_encdec = common_attention.split_heads(k_encdec, hparams.num_heads)
-        v_encdec = common_attention.compute_attention_component(
-            encoder_output,
-            value_channels,
-            name="v",
-            vars_3d_num_heads=vars_3d_num_heads)
-        v_encdec = common_attention.split_heads(v_encdec, hparams.num_heads)
-      cache[layer_name]["k_encdec"] = k_encdec
-      cache[layer_name]["v_encdec"] = v_encdec
-
-    cache["encoder_output"] = encoder_output
-    cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias
+  cache = init_cache_fn(None, hparams, batch_size, decode_length,
+                        encoder_output, encoder_decoder_attention_bias,
+                        scope_prefix)
 
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_SEQ_BEAM_SEARCH,
@@ -1031,6 +1051,7 @@ def fast_decode(encoder_output,
                 hparams,
                 decode_length,
                 vocab_size,
+                init_cache_fn=_init_transformer_cache,
                 beam_size=1,
                 top_beams=1,
                 alpha=1.0,
@@ -1054,6 +1075,7 @@ def fast_decode(encoder_output,
     hparams: run hyperparameters
     decode_length: an integer.  How many additional timesteps to decode.
     vocab_size: Output vocabulary size.
+    init_cache_fn: Function that returns the initial cache dict.
     beam_size: number of beams.
     top_beams: an integer. How many of the beams to return.
     alpha: Float that controls the length penalty. larger the alpha, stronger
@@ -1081,57 +1103,14 @@ def fast_decode(encoder_output,
   if encoder_output is not None:
     batch_size = common_layers.shape_list(encoder_output)[0]
 
-  key_channels = hparams.attention_key_channels or hparams.hidden_size
-  value_channels = hparams.attention_value_channels or hparams.hidden_size
-  num_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
-  vars_3d_num_heads = (
-      hparams.num_heads if hparams.get("attention_variables_3d") else 0)
-
-  if cache is None:
-    cache = {}
-  cache.update({
-      "layer_%d" % layer: {  # pylint: disable=g-complex-comprehension
-          "k":
-              common_attention.split_heads(
-                  tf.zeros([batch_size, 0, key_channels]), hparams.num_heads),
-          "v":
-              common_attention.split_heads(
-                  tf.zeros([batch_size, 0, value_channels]), hparams.num_heads),
-      } for layer in range(num_layers)
-  })
-
-  # If `ffn_layer` is in `["dense_relu_dense" or "conv_hidden_relu"]`, then the
-  # cache key "f" won't be used, which means that the` shape of cache["f"]`
-  # won't be changed to
-  # `[beamsize*batch_size, decode_length, hparams.hidden_size]` and may cause
-  # error when applying `nest.map reshape function` on it.
-  if hparams.ffn_layer not in ["dense_relu_dense", "conv_hidden_relu"]:
-    for layer in range(num_layers):
-      cache["layer_%d" % layer]["f"] = tf.zeros(
-          [batch_size, 0, hparams.hidden_size])
-
-  if encoder_output is not None:
-    for layer in range(num_layers):
-      layer_name = "layer_%d" % layer
-      with tf.variable_scope("%sdecoder/%s/encdec_attention/multihead_attention"
-                             % (scope_prefix, layer_name)):
-        k_encdec = common_attention.compute_attention_component(
-            encoder_output,
-            key_channels,
-            name="k",
-            vars_3d_num_heads=vars_3d_num_heads)
-        k_encdec = common_attention.split_heads(k_encdec, hparams.num_heads)
-        v_encdec = common_attention.compute_attention_component(
-            encoder_output,
-            value_channels,
-            name="v",
-            vars_3d_num_heads=vars_3d_num_heads)
-        v_encdec = common_attention.split_heads(v_encdec, hparams.num_heads)
-      cache[layer_name]["k_encdec"] = k_encdec
-      cache[layer_name]["v_encdec"] = v_encdec
-
-    cache["encoder_output"] = encoder_output
-    cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias
+  cache = init_cache_fn(
+      cache=cache,
+      hparams=hparams,
+      batch_size=batch_size,
+      attention_init_length=0,
+      encoder_output=encoder_output,
+      encoder_decoder_attention_bias=encoder_decoder_attention_bias,
+      scope_prefix=scope_prefix)
 
   if beam_size > 1:  # Beam Search
     initial_ids = sos_id * tf.ones([batch_size], dtype=tf.int32)

From 86283ea534497b2f97d775abd69d8c57cf59a754 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 12 Apr 2019 15:24:46 -0700
Subject: [PATCH 1895/2720] Internal

PiperOrigin-RevId: 243343517
---
 tensor2tensor/data_generators/all_problems.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index f636867c9..5ac52724c 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -19,6 +19,7 @@
 from __future__ import print_function
 
 import importlib
+import six
 from six.moves import range  # pylint: disable=redefined-builtin
 
 MODULES = [

From f4ad26d03dc9361ee2db6f5059824fe22baa939d Mon Sep 17 00:00:00 2001
From: Peng Wang <wangpeng@google.com>
Date: Fri, 12 Apr 2019 15:51:55 -0700
Subject: [PATCH 1896/2720] Internal changes.

PiperOrigin-RevId: 243348205
---
 tensor2tensor/trax/backend.py | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index ffd078501..776ea44db 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -31,13 +31,6 @@
 
 
 
-@gin.configurable()
-def backend(name="jax"):
-  if name == "numpy":
-    return _NUMPY_BACKEND
-  return _JAX_BACKEND
-
-
 _JAX_BACKEND = {
     "np": jnp,
     "logsumexp": jax_special.logsumexp,
@@ -101,7 +94,15 @@ def bernoulli(self, *args, **kwargs):
 
 random = RandomBackend()
 
-# TODO(lukaszkaiser): make this lazy as random above so gin-config works.
-numpy = backend()["np"]
 
 
+
+@gin.configurable()
+def backend(name="jax"):
+  if name == "numpy":
+    return _NUMPY_BACKEND
+  return _JAX_BACKEND
+
+
+# TODO(lukaszkaiser): make this lazy as random above so gin-config works.
+numpy = backend()["np"]

From f155f5a5deca1e8c0380eabca8223cd6a13d4681 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 12 Apr 2019 16:30:04 -0700
Subject: [PATCH 1897/2720] Fix bug with certain memory configurations

PiperOrigin-RevId: 243354693
---
 tensor2tensor/layers/transformer_memory.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index 002e4004c..be2410fbb 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -127,12 +127,11 @@ def pre_attention(self, segment, query_antecedent, memory_antecedent, bias):
     current_batch_size = tf.shape(query_antecedent)[0]
     amount_to_pad = memory_batch_size - current_batch_size
 
-    previous_vals = self.previous_vals
     # If segment id is zero, don't attend back to the memory
     previous_bias = self.previous_bias[:current_batch_size, :, :, :] + tf.cast(
         tf.equal(segment[:, None, None, None], 0), tf.float32) * -1e9
 
-    sliced_previous_vals = previous_vals[:current_batch_size, :, :]
+    sliced_previous_vals = self.previous_vals[:current_batch_size, :, :]
 
     new_memory_antecedent = tf.concat(
         [tf.stop_gradient(sliced_previous_vals), query_antecedent], 1)
@@ -153,9 +152,9 @@ def pre_attention(self, segment, query_antecedent, memory_antecedent, bias):
         tf.reduce_max(bias, -2, keepdims=True), [memory_batch_size, 1, 1, 1])
     # Assume that query_antecedent is always a full chunk (i.e. not truncated)
     if self.chunk_length < self.tokens_to_cache:
-      remember_vals = tf.concat([previous_vals, remember_vals], 1)
+      remember_vals = tf.concat([self.previous_vals, remember_vals], 1)
       remember_bias = tf.concat([
-          previous_bias - 1e9 * tf.cast(
+          self.previous_bias - 1e9 * tf.cast(
               tf.equal(
                   tf.pad(segment, [[0, amount_to_pad]])[:, None, None, None],
                   0), tf.float32),

From d41a4b8ed70b44dfd9cd42410f079829d8a16f98 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 12 Apr 2019 17:42:24 -0700
Subject: [PATCH 1898/2720] Parallel evaluation and eval_batch_size correction
 in trax; refactoring inputs.

PiperOrigin-RevId: 243364919
---
 .../trax/configs/transformer_lm1b_8gb.gin     |  5 +-
 .../trax/configs/transformer_lm1b_tpu.gin     |  7 +-
 tensor2tensor/trax/inputs.py                  | 53 +++++++++------
 tensor2tensor/trax/trax.py                    | 64 ++++++++++++++-----
 tensor2tensor/trax/trax_test.py               |  1 +
 5 files changed, 88 insertions(+), 42 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index 9f00817d9..5b99bd2f7 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -12,6 +12,7 @@ batch_fun.eval_batch_size = 128
 # ==============================================================================
 inputs.data_dir = None
 inputs.dataset_name = 't2t_languagemodel_lm1b32k'
+inputs.input_name = 'targets'
 
 # Parameters for mask:
 # ==============================================================================
@@ -37,10 +38,6 @@ train.model = @trax.models.TransformerLM
 train.run_debug_step = False
 train.train_steps = 500000
 
-# Parameters for train_and_eval_batches:
-# ==============================================================================
-train_and_eval_batches.input_name = 'targets'
-
 # Parameters for TransformerLM:
 # ==============================================================================
 TransformerLM.dropout = 0.1
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_tpu.gin b/tensor2tensor/trax/configs/transformer_lm1b_tpu.gin
index 2681e23cc..7ad8ddfa6 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_tpu.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_tpu.gin
@@ -5,13 +5,14 @@ import tensor2tensor.trax.trax
 
 # Parameters for batch_fun:
 # ==============================================================================
-batch_fun.batch_size = 256
+batch_fun.batch_size = 1024
 batch_fun.eval_batch_size = 128
 
 # Parameters for inputs:
 # ==============================================================================
 inputs.data_dir = None
 inputs.dataset_name = 't2t_languagemodel_lm1b32k'
+inputs.input_name = 'targets'
 
 # Parameters for mask:
 # ==============================================================================
@@ -38,10 +39,6 @@ train.model = @trax.models.TransformerLM
 train.run_debug_step = False
 train.train_steps = 500000
 
-# Parameters for train_and_eval_batches:
-# ==============================================================================
-train_and_eval_batches.input_name = 'targets'
-
 # Parameters for TransformerLM:
 # ==============================================================================
 TransformerLM.dropout = 0.1
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index bd2a7c144..afa528797 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -30,8 +30,22 @@
 import tensorflow as tf
 import tensorflow_datasets as tfds
 
+# Inputs is the trax tuple defining the input streams and shapes.
+# * train_stream: training data that will be used for training
+#     may include all the augmentation or selection the training wants
+#     the shape of examples is [batch_fun.batch_size, ...]
+# * train_eval_stream: training data used for evaluation
+#     examples from training data but usually without augmentation
+#     the shape of examples is [batch_fun.eval_batch_size, ...]
+# * eval_stream: evaluation data stream
+#     examples from evaluation data, usually without augmentation
+#     the shape of examples is [batch_fun.eval_batch_size, ...]
+# * input_shape: the shape of inputs
+#     the [...] above, without batch size
+
 Inputs = collections.namedtuple(
-    "_Inputs", ["train_stream", "eval_stream", "input_shape"])
+    "_Inputs",
+    ["train_stream", "train_eval_stream", "eval_stream", "input_shape"])
 
 # How many examples from the stream to skip at random during training.
 # For now, we skip at most 100K examples for efficiency.
@@ -40,13 +54,14 @@
 
 
 @gin.configurable()
-def inputs(dataset_name, data_dir=None):
+def inputs(dataset_name, data_dir=None, input_name=None):
   """Make Inputs for built-in datasets.
 
   Args:
     dataset_name: a TFDS or T2T dataset name. If it's a T2T dataset name, prefix
       with "t2t_".
     data_dir: data directory.
+    input_name: optional, name of the inputs from the dictionary.
 
   Returns:
     trax.inputs.Inputs
@@ -54,17 +69,21 @@ def inputs(dataset_name, data_dir=None):
   assert data_dir, "Must provide a data directory"
   data_dir = os.path.expanduser(data_dir)
 
-  (train_batches, eval_batches,
-   input_name, input_shape) = train_and_eval_batches(
-       dataset_name, data_dir)
+  (train_batches, train_eval_batches, eval_batches,
+   input_name, input_shape) = _train_and_eval_batches(
+       dataset_name, data_dir, input_name)
 
   def train_input_fun():
     return dataset_to_stream(train_batches, input_name)
 
+  def train_eval_input_fun():
+    return dataset_to_stream(train_eval_batches, input_name)
+
   def eval_input_fun():
     return dataset_to_stream(eval_batches, input_name)
 
   return Inputs(train_stream=train_input_fun,
+                train_eval_stream=train_eval_input_fun,
                 eval_stream=eval_input_fun,
                 input_shape=input_shape)
 
@@ -99,15 +118,10 @@ def random_minibatches():
       out = out.astype(output_dtype)
       yield inp, out
 
-  def train_input_fun():
-    return random_minibatches()
-
-  def eval_input_fun():
-    return random_minibatches()
-
   input_shape_without_batch = list(input_shape)[1:]
-  return Inputs(train_stream=train_input_fun,
-                eval_stream=eval_input_fun,
+  return Inputs(train_stream=random_minibatches,
+                train_eval_stream=random_minibatches,
+                eval_stream=random_minibatches,
                 input_shape=input_shape_without_batch)
 
 
@@ -281,11 +295,12 @@ def target_right_length(_, target):
   return dataset
 
 
-@gin.configurable(whitelist=["preprocess_fun"])
+@gin.configurable(whitelist=["preprocess_fun", "shuffle_buffer_size"])
 def shuffle_and_batch_data(dataset,
                            target_names,
                            features_info,
                            training,
+                           shuffle_buffer_size=1024,
                            preprocess_fun=no_preprocess):
   """Shuffle and batch the given dataset."""
   def append_targets(example):
@@ -306,21 +321,23 @@ def append_targets(example):
   dataset = preprocess_fun(dataset, training)
   shapes = {k: features_info[k].shape for k in features_info}
   shapes = (shapes, shapes[target_names[0]])
-  dataset = dataset.shuffle(1024)
+  dataset = dataset.shuffle(shuffle_buffer_size)
   dataset = batch_fun(dataset, training, shapes, target_names)
   return dataset.prefetch(2)
 
 
-@gin.configurable(whitelist=["input_name"])
-def train_and_eval_batches(dataset, data_dir, input_name=None):
+def _train_and_eval_batches(dataset, data_dir, input_name):
   """Return train and eval batches with input name and shape."""
   (train_data, eval_data, features_info, keys) = train_and_eval_dataset(
       dataset, data_dir)
   input_names, target_names = keys[0], keys[1]
   train_batches = shuffle_and_batch_data(
       train_data, target_names, features_info, training=True)
+  train_eval_batches = shuffle_and_batch_data(  # Data for eval-on-train.
+      train_data, target_names, features_info, training=False)
   eval_batches = shuffle_and_batch_data(
       eval_data, target_names, features_info, training=False)
   input_name = input_name or input_names[0]
   input_shape = features_info[input_name].shape
-  return train_batches, eval_batches, input_name, list(input_shape)
+  return (train_batches, train_eval_batches, eval_batches,
+          input_name, list(input_shape))
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index e4977b82c..48a100ce4 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -147,7 +147,7 @@ def evaluate_train_and_eval(step, inputs, predict_fun, eval_steps, rng,
           _METRICS,
           rng)
       for input_stream in
-      [inputs.train_stream, inputs.eval_stream]]
+      [inputs.train_eval_stream, inputs.eval_stream]]
   if train_sw:
     log_metrics(train_metrics, train_sw, "train", step, history=history)
   if eval_sw:
@@ -244,6 +244,32 @@ def epochs(steps=None, epoch_steps=1):
       break
 
 
+def _jit_predict_fun(model_predict, num_devices, jit_eval):
+  """Use jit on model_predict if required."""
+  def predict(params, batch, rng=None):
+    """Predict function jited and parallelized as requested."""
+    # If not jit'ing, just run the function.
+    if not jit_eval:
+      return model_predict(params, batch, rng=rng)
+
+    # On one device, jit and run.
+    if num_devices == 1:
+      return backend.jit(model_predict)(params, batch, rng=rng)
+
+    # Multi-devices, pmap and run.
+    @functools.partial(backend.pmap, axis_name="batch")
+    def mapped_predict(params, batch, rng):
+      return model_predict(params, batch, rng=rng)
+    pred = mapped_predict(
+        jax.replicate(params),
+        reshape_by_device(batch, num_devices),
+        jax.replicate(rng))
+    batch_size = batch.shape[0]
+    return np.reshape(pred, [batch_size] + list(pred.shape[2:]))
+
+  return predict
+
+
 def _jit_update_fun(predict_fun, loss_fun, optimizer, lr_fun, num_devices):
   """Get jit-ed update function for loss, optimizer, learning rate function."""
   if num_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
@@ -272,20 +298,30 @@ def update(i, opt_state, batch, rng):
   return update
 
 
-def reshape_by_device(train_data, num_devices):
-  """Reshape the train_data into a shape [num_devices, ...]."""
-  x, y = train_data
-  x_shape, y_shape = list(x.shape), list(y.shape)
-  assert x_shape[0] == y_shape[0]  # Same batch size.
+def reshape_by_device(x, num_devices):
+  """Reshape x into a shape [num_devices, ...]."""
+  x_shape = list(x.shape)
   batch_size = x_shape[0]
   batch_size_per_device = batch_size // num_devices
   # We require that num_devices divides batch_size evenly.
-  assert batch_size_per_device * num_devices == batch_size
-  # New shapes.
+  if batch_size_per_device * num_devices != batch_size:
+    logging.fatal(
+        "We require that num_devices[%d] divides batch_size[%d] evenly.",
+        num_devices, batch_size)
+  # New shape.
   new_shape_prefix = [num_devices, batch_size_per_device]
-  x = np.reshape(x, new_shape_prefix + x_shape[1:])
-  y = np.reshape(y, new_shape_prefix + y_shape[1:])
-  return x, y
+  return np.reshape(x, new_shape_prefix + x_shape[1:])
+
+
+def reshape_by_device_pair(train_data, num_devices):
+  """Reshape by device for a pair."""
+  x, y = train_data
+  x_shape, y_shape = list(x.shape), list(y.shape)
+  if x_shape[0] != y_shape[0]:  # Same batch size.
+    logging.fatal(
+        "Batch size is not the same for train_data pair: [%d] vs [%d]",
+        x_shape[0], y_shape[0])
+  return reshape_by_device(x, num_devices), reshape_by_device(y, num_devices)
 
 
 @gin.configurable(blacklist=["output_dir"])
@@ -355,9 +391,7 @@ def train(output_dir,
     opt_state = jax.replicate(opt_state)
 
   # jit model_predict and update so they're fast
-  jit_model_predict = model_predict
-  if jit_eval:
-    jit_model_predict = backend.jit(model_predict)  # for evaluation
+  jit_model_predict = _jit_predict_fun(model_predict, num_devices, jit_eval)
   jit_update_fun = _jit_update_fun(model_predict, loss_fun, optimizer, lr_fun,
                                    num_devices)
 
@@ -386,7 +420,7 @@ def train(output_dir,
       # Train
       next_train_batch = next(train_stream)
       if num_devices > 1:  # TODO(lukaszkaiser): use everywhere when possible.
-        next_train_batch = reshape_by_device(next_train_batch, num_devices)
+        next_train_batch = reshape_by_device_pair(next_train_batch, num_devices)
       rng, subrng = jax_random.split(rng)
       opt_state = jit_update_fun(step, opt_state, next_train_batch, subrng)
       step += 1
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index a73b90169..cb375d55a 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -47,6 +47,7 @@ def input_stream():
 
   return inputs_lib.Inputs(
       train_stream=input_stream,
+      train_eval_stream=input_stream,
       eval_stream=input_stream,
       input_shape=input_shape)
 

From 2d43d8c812a72d77205af3bf4f5dc7388a6d5747 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 15 Apr 2019 10:07:30 -0700
Subject: [PATCH 1899/2720] [trax] update for jax.lax.psum module structure
 change

PiperOrigin-RevId: 243637442
---
 tensor2tensor/trax/trax.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 48a100ce4..595bac5bc 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -32,7 +32,7 @@
 import gin
 
 import jax
-from jax import lax_parallel as lax_parallel
+from jax import lax
 import numpy
 import six
 
@@ -288,7 +288,7 @@ def mapped_update(i, opt_state, batch, rng):
     params = trax_opt.get_params(opt_state)
     grads = backend.grad(loss_fun)(params, batch, predict_fun, rng)
     grads = jax.tree_util.tree_map(
-        lambda g: lax_parallel.psum(g, "batch"), grads)
+        lambda g: lax.psum(g, "batch"), grads)
     return opt_update(i, grads, opt_state)
 
   def update(i, opt_state, batch, rng):

From abf268a1d353d75d257e14e1a73dcea112337559 Mon Sep 17 00:00:00 2001
From: Aidan Gomez <aidangomez@google.com>
Date: Mon, 15 Apr 2019 10:14:05 -0700
Subject: [PATCH 1900/2720] Fix distillation code to use temperature on student
 and scale objective.

PiperOrigin-RevId: 243638717
---
 tensor2tensor/models/distillation.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/distillation.py b/tensor2tensor/models/distillation.py
index 1dbe46bc7..6600f48d6 100644
--- a/tensor2tensor/models/distillation.py
+++ b/tensor2tensor/models/distillation.py
@@ -105,7 +105,10 @@ def body(self, features):
             labels=one_hot_targets, logits=student_logits)
         teacher_targets = tf.nn.softmax(teacher_logits / hp.distill_temperature)
         student_distill_xent = tf.nn.softmax_cross_entropy_with_logits_v2(
-            labels=tf.stop_gradient(teacher_targets), logits=student_logits)
+            labels=tf.stop_gradient(teacher_targets),
+            logits=student_logits / hp.distill_temperature)
+        # scale soft target obj. to match hard target obj. scale
+        student_distill_xent *= hp.distill_temperature**2
 
         outputs = student_logits
 

From 7f5ae3e76249b24f3ea7448899cb3e375ab9d433 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 15 Apr 2019 12:01:01 -0700
Subject: [PATCH 1901/2720] Make an object in backend.py for numpy.

PiperOrigin-RevId: 243661125
---
 tensor2tensor/trax/backend.py      | 13 ++++--
 tensor2tensor/trax/backend_test.py | 75 ++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 3 deletions(-)
 create mode 100644 tensor2tensor/trax/backend_test.py

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index 776ea44db..872352c0f 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -95,6 +95,16 @@ def bernoulli(self, *args, **kwargs):
 random = RandomBackend()
 
 
+# A class that just forwards attribute accesses to backend's numpy object.
+class NumpyBackend(object):
+
+  def __getattr__(self, attr):
+    return getattr(backend()["np"], attr)
+
+
+numpy = NumpyBackend()
+
+
 
 
 @gin.configurable()
@@ -103,6 +113,3 @@ def backend(name="jax"):
     return _NUMPY_BACKEND
   return _JAX_BACKEND
 
-
-# TODO(lukaszkaiser): make this lazy as random above so gin-config works.
-numpy = backend()["np"]
diff --git a/tensor2tensor/trax/backend_test.py b/tensor2tensor/trax/backend_test.py
new file mode 100644
index 000000000..a0939e540
--- /dev/null
+++ b/tensor2tensor/trax/backend_test.py
@@ -0,0 +1,75 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.backend."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gin
+import jax.numpy as jnp
+import numpy as onp
+from tensor2tensor.trax import backend as backend_lib
+from tensorflow import test
+
+
+class BackendTest(test.TestCase):
+
+  def setUp(self):
+    gin.clear_config()
+
+  def override_gin(self, bindings):
+    gin.parse_config_files_and_bindings(None, bindings)
+
+  def test_backend_imports_correctly(self):
+    backend = backend_lib.backend()
+    self.assertEqual(jnp, backend["np"])
+    self.assertNotEqual(onp, backend["np"])
+
+    self.override_gin("backend.name = 'numpy'")
+
+    backend = backend_lib.backend()
+    self.assertNotEqual(jnp, backend["np"])
+    self.assertEqual(onp, backend["np"])
+
+  def test_numpy_backend_delegation(self):
+    # Assert that we are getting JAX's numpy backend.
+    backend = backend_lib.backend()
+    numpy = backend_lib.numpy
+    self.assertEqual(jnp, backend["np"])
+
+    # Assert that `numpy` calls the appropriate gin configured functions and
+    # properties.
+    self.assertTrue(numpy.isinf(numpy.inf))
+    self.assertEqual(jnp.isinf, numpy.isinf)
+    self.assertEqual(jnp.inf, numpy.inf)
+
+    # Assert that we will now get the pure numpy backend.
+
+    self.override_gin("backend.name = 'numpy'")
+
+    backend = backend_lib.backend()
+    numpy = backend_lib.numpy
+    self.assertEqual(onp, backend["np"])
+
+    # Assert that `numpy` calls the appropriate gin configured functions and
+    # properties.
+    self.assertTrue(numpy.isinf(numpy.inf))
+    self.assertEqual(onp.isinf, numpy.isinf)
+    self.assertEqual(onp.inf, numpy.inf)
+
+if __name__ == "__main__":
+  test.main()

From 917bc0e6f5fff29deedfa0f5ab09e6036364eea9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 15 Apr 2019 12:03:32 -0700
Subject: [PATCH 1902/2720] A fake env (and its test) for testing. Will be used
 in PPO tests.

PiperOrigin-RevId: 243661833
---
 tensor2tensor/trax/rlax/fake_env.py      | 63 +++++++++++++++++++++++
 tensor2tensor/trax/rlax/fake_env_test.py | 65 ++++++++++++++++++++++++
 2 files changed, 128 insertions(+)
 create mode 100644 tensor2tensor/trax/rlax/fake_env.py
 create mode 100644 tensor2tensor/trax/rlax/fake_env_test.py

diff --git a/tensor2tensor/trax/rlax/fake_env.py b/tensor2tensor/trax/rlax/fake_env.py
new file mode 100644
index 000000000..cc4b6fd98
--- /dev/null
+++ b/tensor2tensor/trax/rlax/fake_env.py
@@ -0,0 +1,63 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""A fake gym environment.
+
+Can specify either:
+1. A done action, i.e. the action on which the environment returns done.
+2. A done time-step, i.e. the time step at which the environment returns done.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gym
+import numpy as np
+
+
+class FakeEnv(object):
+  """A fake env which is either done with a specific action or a time-step."""
+
+  def __init__(self,
+               input_shape=(4,),
+               num_actions=2,
+               done_time_step=None,
+               done_action=None):
+    self._input_shape = input_shape
+    self._done_time_step = done_time_step
+    self._done_action = done_action
+    self._t = 0
+    self.action_space = gym.spaces.Discrete(num_actions)
+    self.observation_space = gym.spaces.Box(
+        low=-1.0, high=1.0, shape=input_shape)
+
+  def _get_random_observation(self):
+    return np.random.random(self._input_shape)
+
+  def reset(self):
+    self._t = 0
+    return self._get_random_observation()
+
+  def step(self, action):
+    done = False
+    if self._done_action is not None:
+      done = action == self._done_action
+    elif self._done_time_step is not None:
+      done = self._t == self._done_time_step
+
+    reward = -1.0 if not done else 1.0
+    self._t += 1
+    return self._get_random_observation(), reward, done, {}
diff --git a/tensor2tensor/trax/rlax/fake_env_test.py b/tensor2tensor/trax/rlax/fake_env_test.py
new file mode 100644
index 000000000..34aba74d7
--- /dev/null
+++ b/tensor2tensor/trax/rlax/fake_env_test.py
@@ -0,0 +1,65 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.rlax.fake_env."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.trax.rlax import fake_env
+from tensorflow import test
+
+
+class FakeEnvTest(test.TestCase):
+
+  def test_done_action(self):
+    env = fake_env.FakeEnv(input_shape=(2, 3),
+                           num_actions=10,
+                           done_time_step=None,
+                           done_action=9)
+    env.reset()
+
+    # Actions 0 to 8
+    for action in range(9):
+      _, reward, done, _ = env.step(action)
+      self.assertFalse(done)
+      self.assertEqual(-1.0, reward)
+
+    _, reward, done, _ = env.step(9)
+    self.assertTrue(done)
+    self.assertEqual(1.0, reward)
+
+  def test_done_time_step(self):
+    env = fake_env.FakeEnv(input_shape=(2, 3),
+                           num_actions=10,
+                           done_time_step=10,
+                           done_action=None)
+    env.reset()
+
+    # Take 10 steps.
+    for _ in range(10):
+      _, reward, done, _ = env.step(0)
+      self.assertFalse(done)
+      self.assertEqual(-1.0, reward)
+
+    # Take final time-step, this is the time-step numbered 10 since time-steps
+    # are 0 indexed.
+    _, reward, done, _ = env.step(0)
+    self.assertTrue(done)
+    self.assertEqual(1.0, reward)
+
+if __name__ == '__main__':
+  test.main()

From a0f387d32eb810104c0d11a7f020a88565ca72d0 Mon Sep 17 00:00:00 2001
From: Zi Yang <ziy@google.com>
Date: Mon, 15 Apr 2019 13:37:06 -0700
Subject: [PATCH 1903/2720] Default experimental_export_device_assignment to
 True.

This is a workaround for fixing the default graph mismatch (ValueError: prediction values with "inputs: DatasetToSingleElement:0" must be from the default graph.) raised during exporting TPU models.

PiperOrigin-RevId: 243679157
---
 tensor2tensor/utils/trainer_lib.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index ec8606c82..d890d02e7 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -314,7 +314,8 @@ def create_estimator(model_name,
         use_tpu=use_tpu,
         train_batch_size=batch_size,
         eval_batch_size=batch_size if "eval" in schedule else None,
-        predict_batch_size=predict_batch_size)
+        predict_batch_size=predict_batch_size,
+        experimental_export_device_assignment=True)
   else:
     estimator = tf.estimator.Estimator(
         model_fn=model_fn,

From f4e330ea0587a0eae3032e0486f6430a246c669f Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 16 Apr 2019 08:52:24 -0700
Subject: [PATCH 1904/2720] Make decode_x arguments strings by default to avoid
 hparams parse errors.

PiperOrigin-RevId: 243815594
---
 tensor2tensor/data_generators/translate.py |  3 +--
 tensor2tensor/utils/decoding.py            |  3 ++-
 tensor2tensor/utils/flags.py               |  6 +++---
 tensor2tensor/utils/trainer_lib.py         | 12 +++++++-----
 4 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 448d970f1..1128ba63d 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -91,8 +91,7 @@ def compute_bleu_summaries(hook_args):
   """
   decode_hparams = hook_args.decode_hparams
 
-  if (decode_hparams.decode_reference is None or
-      decode_hparams.decode_to_file is None):
+  if not (decode_hparams.decode_reference and decode_hparams.decode_to_file):
     return None
 
   values = []
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 511787942..78fec4b79 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -65,7 +65,8 @@ def decode_hparams(overrides=""):
       identity_output=False,
       num_samples=-1,  # Number of examples to decode.
       delimiter="\n",
-      decode_to_file=None,  # str. Prefix for filename to write decodings to.
+      decode_to_file="",  # str. Prefix for filename to write decodings to.
+      decode_reference="",  # str. Filename to read references from.
       decode_in_memory=False,
       # How much decode should wait for the next checkpoint
       decode_timeout_mins=240,
diff --git a/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index 5525a1796..1359fd140 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -113,12 +113,12 @@
                     "Comma-separated list of name=value pairs to control "
                     "decode behavior. See decoding.decode_hparams for "
                     "defaults.")
-flags.DEFINE_string("decode_from_file", None,
+flags.DEFINE_string("decode_from_file", "",
                     "Path to the source file for decoding, used by "
                     "continuous_decode_from_file.")
-flags.DEFINE_string("decode_to_file", None,
+flags.DEFINE_string("decode_to_file", "",
                     "Path to the decoded file generated by decoding, used by "
                     "continuous_decode_from_file.")
-flags.DEFINE_string("decode_reference", None,
+flags.DEFINE_string("decode_reference", "",
                     "Path to the reference file for decoding, used by "
                     "continuous_decode_from_file to compute BLEU score.")
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index d890d02e7..2788bbd09 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -637,9 +637,9 @@ def create_experiment(
     additional_train_hooks=None,
     additional_eval_hooks=None,
     warm_start_from=None,
-    decode_from_file=None,
-    decode_to_file=None,
-    decode_reference=None,
+    decode_from_file="",
+    decode_to_file="",
+    decode_reference="",
     std_server_protocol=None):
   """Create Experiment."""
   # HParams
@@ -654,8 +654,10 @@ def create_experiment(
   hparams.add_hparam("eval_timeout_mins", eval_timeout_mins)
   if decode_hparams is not None:
     decode_hparams.add_hparam("decode_from_file", decode_from_file)
-    decode_hparams.add_hparam("decode_to_file", decode_to_file)
-    decode_hparams.add_hparam("decode_reference", decode_reference)
+    if decode_to_file and not decode_hparams.decode_to_file:
+      decode_hparams.decode_to_file = decode_to_file
+    if decode_reference and not decode_hparams.decode_reference:
+      decode_hparams.decode_reference = decode_reference
   add_problem_hparams(hparams, problem_name)
 
   # Estimator

From 63d677cd2801a58967540646596ff2296de0329a Mon Sep 17 00:00:00 2001
From: Noah Fiedel <nfiedel@google.com>
Date: Wed, 17 Apr 2019 09:40:23 -0700
Subject: [PATCH 1905/2720] Fix typo.

PiperOrigin-RevId: 244012639
---
 tensor2tensor/data_generators/ops/pack_sequences_ops.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
index 4c07c7bdd..13b5906a1 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
@@ -15,7 +15,7 @@ using ::tensorflow::Tensor;
 using ::tensorflow::TensorShape;
 using ::tensorflow::shape_inference::InferenceContext;
 
-// TODO(noam): this op packs a dataset of pairs of sequaneces (inputs, targets)
+// TODO(noam): this op packs a dataset of pairs of sequences (inputs, targets)
 // Generalize later to an arbitrary number of sequences.
 REGISTER_OP("PackSequences2")
     .Input("inputs: int64")

From 02b4dcf5b032fc5245b7f80728866efbb2ea9f62 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 17 Apr 2019 11:49:43 -0700
Subject: [PATCH 1906/2720] Add output type option to
 attention_bias_to_padding() so that Evolved Transformer can support bfloat16.

PiperOrigin-RevId: 244038573
---
 tensor2tensor/layers/common_attention.py    |  7 ++++---
 tensor2tensor/models/evolved_transformer.py | 10 +++++++++-
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 3d3322635..2c4067056 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -947,20 +947,21 @@ def attention_bias_ignore_padding(memory_padding):
 
 
 @expert_utils.add_name_scope()
-def attention_bias_to_padding(attention_bias):
+def attention_bias_to_padding(attention_bias, cast_fn=tf.to_float):
   """Inverse of attention_bias_ignore_padding().
 
   Args:
     attention_bias: a `Tensor` with shape [batch, 1, 1, memory_length], as
       returned by attention_bias_ignore_padding().
+    cast_fn: function used to cast to output type.
 
   Returns:
     a Tensor with shape [batch, memory_length] with 1.0 in padding positions
-    and 0.0 in non-padding positions.
+    and 0.0 in non-padding positions. Type is determined by cast_fn.
   """
   # `attention_bias` is a large negative number in padding positions and 0.0
   # elsewhere.
-  return tf.squeeze(tf.to_float(tf.less(attention_bias, -1)), axis=[1, 2])
+  return tf.squeeze(cast_fn(tf.less(attention_bias, -1)), axis=[1, 2])
 
 
 @expert_utils.add_name_scope()
diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 4d27d2416..e39f12583 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -100,7 +100,15 @@ def evolved_transformer_encoder(encoder_input,
       attention_bias = encoder_self_attention_bias
       if attn_bias_for_padding is not None:
         attention_bias = attn_bias_for_padding
-      padding = common_attention.attention_bias_to_padding(attention_bias)
+      # Only bfloat16 and float32 supported.
+      float_type = hparams.get("activation_dtype", "float32")
+      if float_type == "bfloat16":
+        cast_fn = tf.to_bfloat16
+      else:
+        assert float_type == "float32"
+        cast_fn = tf.to_float
+      padding = common_attention.attention_bias_to_padding(
+          attention_bias, cast_fn)
       nonpadding = 1.0 - padding
 
     for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):

From 56f8d063b1928f1328d776efea3a7b72769d2f18 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 17 Apr 2019 12:04:14 -0700
Subject: [PATCH 1907/2720] Improve efficiency of Evolved Transformer fast
 decoding by only caching and computing convolutions over the receptive field.

PiperOrigin-RevId: 244041280
---
 tensor2tensor/models/evolved_transformer.py | 109 +++++++++++++++-----
 1 file changed, 84 insertions(+), 25 deletions(-)

diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index e39f12583..cf6468277 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -34,11 +34,17 @@
 # pylint: enable=g-direct-tensorflow-import
 
 _CONV_BRANCHES_NAME = "conv_branches"
+_CONV_BRANCHES_FIRST_LAYER_NAME = _CONV_BRANCHES_NAME + "_first"
+_CONV_BRANCHES_SECOND_LAYER_NAME = _CONV_BRANCHES_NAME + "_second"
 _FIRST_ATTEND_TO_ENCODER_NAME = "first_attend_to_encoder"
 _SECOND_ATTEND_TO_ENCODER_NAME = "second_attend_to_encoder"
 _SIXTEEN_HEAD_ATTENTION_NAME = "16_head_self_attention"
 _VANILLA_ATTENTION_NAME = "self_attention"
 
+_DECODER_LEFT_CONV_PADDING = 10
+_DECODER_RIGHT_CONV_PADDING = 6
+_DECODER_FINAL_CONV_PADDING = 6
+
 
 @registry.register_model
 class EvolvedTransformer(transformer.Transformer):
@@ -358,38 +364,54 @@ def evolved_transformer_decoder(decoder_input,
                 tf.expand_dims(nonpadding, 2), [1, 1, hparams.hidden_size])
             hidden_state *= mask
 
-          # TODO(davidso): This needlessly recomputes past positions. Limit
-          # the module inputs to only include positions that are convered by
-          # the new output position's receptive field.
           if layer_cache:
             if decode_loop_step is None:
-              hidden_state = layer_cache[_CONV_BRANCHES_NAME] = tf.concat(
-                  [layer_cache[_CONV_BRANCHES_NAME], hidden_state], axis=1)
+              hidden_state = layer_cache[
+                  _CONV_BRANCHES_FIRST_LAYER_NAME] = tf.concat(
+                      [
+                          layer_cache[_CONV_BRANCHES_FIRST_LAYER_NAME],
+                          hidden_state
+                      ],
+                      axis=1)[:, -1 * _DECODER_LEFT_CONV_PADDING - 1:, :]
               left_state = hidden_state
-              right_state = hidden_state
+              right_state = hidden_state[:, _DECODER_LEFT_CONV_PADDING -
+                                         _DECODER_RIGHT_CONV_PADDING:, :]
 
             else:
               # Inplace update is required for inference on TPU.
               # Inplace_ops only supports inplace_update on the first dimension.
               tmp = tf.transpose(
-                  layer_cache[_CONV_BRANCHES_NAME], perm=[1, 0, 2])
+                  layer_cache[_CONV_BRANCHES_FIRST_LAYER_NAME], perm=[1, 0, 2])
               tmp = tf.expand_dims(tmp, axis=1)
               tmp = inplace_ops.alias_inplace_update(
                   tmp,
                   decode_loop_step * tf.shape(hidden_state)[1],
                   tf.transpose(hidden_state, perm=[1, 0, 2]))
               tmp = tf.squeeze(tmp, axis=1)
-              hidden_state = layer_cache[_CONV_BRANCHES_NAME] = tf.transpose(
-                  tmp, perm=[1, 0, 2])
-
-              read_to_index = decode_loop_step + 1
-              left_state = hidden_state[:, :read_to_index, :]
-              right_state = hidden_state[:, :read_to_index, :]
-          else:
-            left_state = hidden_state
-            right_state = hidden_state
+              hidden_state = layer_cache[
+                  _CONV_BRANCHES_FIRST_LAYER_NAME] = tf.transpose(
+                      tmp, perm=[1, 0, 2])
+
+              left_state_indexes = [
+                  decode_loop_step + i
+                  for i in range(_DECODER_LEFT_CONV_PADDING + 1)
+              ]
+              left_state = tf.gather(hidden_state, left_state_indexes, axis=1)
+              right_state_indexes = [
+                  decode_loop_step + i +
+                  (_DECODER_LEFT_CONV_PADDING - _DECODER_RIGHT_CONV_PADDING)
+                  for i in range(_DECODER_RIGHT_CONV_PADDING + 1)
+              ]
+              right_state = tf.gather(hidden_state, right_state_indexes, axis=1)
+
+          else:  # No caching.
+            left_state = tf.pad(
+                hidden_state,
+                paddings=[[0, 0], [_DECODER_LEFT_CONV_PADDING, 0], [0, 0]])
+            right_state = tf.pad(
+                hidden_state,
+                paddings=[[0, 0], [_DECODER_RIGHT_CONV_PADDING, 0], [0, 0]])
 
-          left_state = tf.pad(left_state, paddings=[[0, 0], [10, 0], [0, 0]])
           left_output_dim = int(hparams.hidden_size * 2)
           separable_conv_11x1 = tf.layers.SeparableConv1D(
               left_output_dim,
@@ -401,7 +423,6 @@ def evolved_transformer_decoder(decoder_input,
           left_state = tf.nn.dropout(left_state,
                                      1 - hparams.layer_prepostprocess_dropout)
 
-          right_state = tf.pad(right_state, paddings=[[0, 0], [6, 0], [0, 0]])
           right_output_dim = int(hparams.hidden_size / 2)
           separable_conv_7x1_1 = tf.layers.SeparableConv1D(
               right_output_dim, 7, padding="VALID", name="separable_conv_7x1_1")
@@ -422,7 +443,42 @@ def evolved_transformer_decoder(decoder_input,
                 tf.expand_dims(nonpadding, 2), [1, 1, hparams.hidden_size * 2])
             hidden_state *= mask
 
-          hidden_state = tf.pad(hidden_state, paddings=[[0, 0], [6, 0], [0, 0]])
+          if layer_cache:
+            if decode_loop_step is None:
+              hidden_state = layer_cache[
+                  _CONV_BRANCHES_SECOND_LAYER_NAME] = tf.concat(
+                      [
+                          layer_cache[_CONV_BRANCHES_SECOND_LAYER_NAME],
+                          hidden_state
+                      ],
+                      axis=1)[:, -1 * _DECODER_FINAL_CONV_PADDING - 1:, :]
+
+            else:
+              # Inplace update is required for inference on TPU.
+              # Inplace_ops only supports inplace_update on the first dimension.
+              tmp = tf.transpose(
+                  layer_cache[_CONV_BRANCHES_SECOND_LAYER_NAME], perm=[1, 0, 2])
+              tmp = tf.expand_dims(tmp, axis=1)
+              tmp = inplace_ops.alias_inplace_update(
+                  tmp, (decode_loop_step + _DECODER_FINAL_CONV_PADDING) *
+                  tf.shape(hidden_state)[1],
+                  tf.transpose(hidden_state, perm=[1, 0, 2]))
+              tmp = tf.squeeze(tmp, axis=1)
+              hidden_state = layer_cache[
+                  _CONV_BRANCHES_SECOND_LAYER_NAME] = tf.transpose(
+                      tmp, perm=[1, 0, 2])
+
+              hidden_state_indexes = [
+                  decode_loop_step + i
+                  for i in range(_DECODER_FINAL_CONV_PADDING + 1)
+              ]
+              hidden_state = tf.gather(
+                  hidden_state, hidden_state_indexes, axis=1)
+          else:
+            hidden_state = tf.pad(
+                hidden_state,
+                paddings=[[0, 0], [_DECODER_FINAL_CONV_PADDING, 0], [0, 0]])
+
           separable_conv_7x1_2 = tf.layers.SeparableConv1D(
               hparams.hidden_size,
               7,
@@ -430,9 +486,6 @@ def evolved_transformer_decoder(decoder_input,
               name="separable_conv_7x1_2")
           hidden_state = separable_conv_7x1_2.apply(hidden_state)
 
-          if layer_cache:
-            hidden_state = hidden_state[:, -1:, :]
-
           hidden_state = common_layers.layer_postprocess(
               residual_state, hidden_state, hparams)
 
@@ -591,10 +644,16 @@ def _init_evolved_transformer_cache(cache, hparams, batch_size,
       } for layer in range(num_layers)
   })
 
-  # Add branched layers.
+  # Add branched layers. Pad with additional zeros for causal convolution.
   for layer in range(num_layers):
-    cache["layer_%d" % layer][_CONV_BRANCHES_NAME] = tf.zeros(
-        [batch_size, attention_init_length, hparams.hidden_size])
+    cache["layer_%d" % layer][_CONV_BRANCHES_FIRST_LAYER_NAME] = tf.zeros([
+        batch_size, attention_init_length + _DECODER_LEFT_CONV_PADDING,
+        hparams.hidden_size
+    ])
+    cache["layer_%d" % layer][_CONV_BRANCHES_SECOND_LAYER_NAME] = tf.zeros([
+        batch_size, attention_init_length + _DECODER_FINAL_CONV_PADDING,
+        hparams.hidden_size * 2
+    ])
 
   # Add encoder embedding attentions.
   if encoder_output is not None:

From de95c49826308fc01f1b2702271da26defcf2277 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 17 Apr 2019 12:29:46 -0700
Subject: [PATCH 1908/2720] Add DenseVariationalDropout and
 Conv2DVariationalDropout.

Implementation is based on https://github.com/google-research/google-research/tree/master/state_of_sparsity and  https://github.com/senya-ashukha/variational-dropout-sparsifies-dnn.

Note it does not implement thresholding units for sparsity at eval time.

PiperOrigin-RevId: 244045603
---
 tensor2tensor/keras/regularizers.py |  40 ++++-
 tensor2tensor/layers/bayes.py       | 180 ++++++++++++++++++++++
 tensor2tensor/layers/bayes_test.py  | 227 +++++++++++++++++-----------
 3 files changed, 354 insertions(+), 93 deletions(-)

diff --git a/tensor2tensor/keras/regularizers.py b/tensor2tensor/keras/regularizers.py
index 635d8a794..4220ababb 100644
--- a/tensor2tensor/keras/regularizers.py
+++ b/tensor2tensor/keras/regularizers.py
@@ -25,11 +25,41 @@
 from tensorflow_probability import edward2 as ed
 
 
+class LogUniformKLDivergence(tf.keras.regularizers.Regularizer):
+  """KL divergence regularizer from an input to the log-uniform distribution."""
+
+  def __call__(self, x):
+    """Computes regularization given an ed.Normal random variable as input."""
+    if not isinstance(x, ed.RandomVariable):
+      raise ValueError('Input must be an ed.RandomVariable (for correct math, '
+                       'an ed.Normal random variable).')
+    # Clip magnitude of dropout rate, where we get the dropout rate alpha from
+    # the additive parameterization (Molchanov et al., 2017): for weight ~
+    # Normal(mu, sigma**2), the variance `sigma**2 = alpha * mu**2`.
+    mean = x.distribution.mean()
+    log_variance = tf.log(x.distribution.variance())
+    log_alpha = log_variance - tf.log(tf.square(mean) +
+                                      tf.keras.backend.epsilon())
+    log_alpha = tf.clip_by_value(log_alpha, -8., 8.)
+
+    # Set magic numbers for cubic polynomial approx. (Molchanov et al., 2017).
+    k1 = 0.63576
+    k2 = 1.8732
+    k3 = 1.48695
+    c = -k1
+    output = tf.reduce_sum(k1 * tf.nn.sigmoid(k2 + k3 * log_alpha) +
+                           -0.5 * tf.log1p(tf.exp(-log_alpha)) + c)
+    return output
+
+  def get_config(self):
+    return {}
+
+
 class NormalKLDivergence(tf.keras.regularizers.Regularizer):
-  """KL divergence regularizer from one normal distribution to another."""
+  """KL divergence regularizer from an input to the normal distribution."""
 
   def __init__(self, mean=0., stddev=1.):
-    """Construct regularizer where default is a KL towards the std normal."""
+    """Constructs regularizer where default is a KL towards the std normal."""
     self.mean = mean
     self.stddev = stddev
 
@@ -54,8 +84,10 @@ def get_config(self):
 
 # Compatibility aliases, following tf.keras
 
-
-normal_kl_divergence = NormalKLDivergence  # pylint: disable=invalid-name
+# pylint: disable=invalid-name
+log_uniform_kl_divergence = LogUniformKLDivergence
+normal_kl_divergence = NormalKLDivergence
+# pylint: enable=invalid-name
 
 # Utility functions, following tf.keras
 
diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 69ae77430..2752baa00 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -131,6 +131,100 @@ def call(self, *args, **kwargs):
     return super(Conv2DReparameterization, self).call(*args, **kwargs)
 
 
+@add_weight
+class Conv2DVariationalDropout(tf.keras.layers.Conv2D):
+  """2D convolution layer with variational dropout (Kingma et al., 2015).
+
+  Implementation follows the additive parameterization of
+  Molchanov et al. (2017).
+  """
+
+  def __init__(self,
+               filters,
+               kernel_size,
+               strides=(1, 1),
+               padding='valid',
+               data_format=None,
+               dilation_rate=(1, 1),
+               activation=None,
+               use_bias=True,
+               kernel_initializer='trainable_normal',
+               bias_initializer='zeros',
+               kernel_regularizer='log_uniform_kl_divergence',
+               bias_regularizer=None,
+               activity_regularizer=None,
+               kernel_constraint=None,
+               bias_constraint=None,
+               **kwargs):
+    super(Conv2DVariationalDropout, self).__init__(
+        filters=filters,
+        kernel_size=kernel_size,
+        strides=strides,
+        padding=padding,
+        data_format=data_format,
+        dilation_rate=dilation_rate,
+        activation=activation,
+        use_bias=use_bias,
+        kernel_initializer=initializers.get(kernel_initializer),
+        bias_initializer=initializers.get(bias_initializer),
+        kernel_regularizer=regularizers.get(kernel_regularizer),
+        bias_regularizer=regularizers.get(bias_regularizer),
+        activity_regularizer=regularizers.get(activity_regularizer),
+        kernel_constraint=constraints.get(kernel_constraint),
+        bias_constraint=constraints.get(bias_constraint),
+        **kwargs)
+
+  def call_weights(self):
+    """Calls any weights if the initializer is itself a layer."""
+    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
+      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
+    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
+      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
+
+  def call(self, inputs, training=None):
+    self.call_weights()
+    if training is None:
+      training = tf.keras.backend.learning_phase()
+
+    def dropped_inputs():
+      """Forward pass with dropout."""
+      # Clip magnitude of dropout rate, where we get the dropout rate alpha from
+      # the additive parameterization (Molchanov et al., 2017): for weight ~
+      # Normal(mu, sigma**2), the variance `sigma**2 = alpha * mu**2`.
+      mean = self.kernel.distribution.mean()
+      log_variance = tf.log(self.kernel.distribution.variance())
+      log_alpha = log_variance - tf.log(tf.square(mean) +
+                                        tf.keras.backend.epsilon())
+      log_alpha = tf.clip_by_value(log_alpha, -8., 8.)
+      log_variance = log_alpha + tf.log(tf.square(mean) +
+                                        tf.keras.backend.epsilon())
+
+      means = self._convolution_op(inputs, mean)
+      stddevs = tf.sqrt(
+          self._convolution_op(tf.square(inputs), tf.exp(log_variance)) +
+          tf.keras.backend.epsilon())
+      outputs = means + stddevs * tf.random_normal(tf.shape(stddevs))
+      if self.use_bias:
+        outputs = tf.nn.bias_add(outputs, self.bias)
+      if self.activation is not None:
+        outputs = self.activation(outputs)
+      return outputs
+
+    # Following tf.keras.Dropout, only apply variational dropout if training
+    # flag is True. The kernel must also be a random variable.
+    training_value = tf.contrib.util.constant_value(training)
+    if training_value is not None:
+      if training_value and isinstance(self.kernel, ed.RandomVariable):
+        return dropped_inputs()
+      else:
+        return super(Conv2DVariationalDropout, self).call(inputs)
+    else:
+      return tf.cond(tf.logical_and(training,
+                                    isinstance(self.kernel, ed.RandomVariable)),
+                     dropped_inputs,
+                     lambda: super(Conv2DVariationalDropout, self).call(inputs))
+
+
 @add_weight
 class DenseReparameterization(tf.keras.layers.Dense):
   """Bayesian densely-connected layer estimated via reparameterization.
@@ -184,6 +278,92 @@ def call(self, *args, **kwargs):
     return super(DenseReparameterization, self).call(*args, **kwargs)
 
 
+@add_weight
+class DenseVariationalDropout(tf.keras.layers.Dense):
+  """Densely-connected layer with variational dropout (Kingma et al., 2015).
+
+  Implementation follows the additive parameterization of
+  Molchanov et al. (2017).
+  """
+
+  def __init__(self,
+               units,
+               activation=None,
+               use_bias=True,
+               kernel_initializer='trainable_normal',
+               bias_initializer='zero',
+               kernel_regularizer='log_uniform_kl_divergence',
+               bias_regularizer=None,
+               activity_regularizer=None,
+               **kwargs):
+    super(DenseVariationalDropout, self).__init__(
+        units=units,
+        activation=activation,
+        use_bias=use_bias,
+        kernel_initializer=initializers.get(kernel_initializer),
+        bias_initializer=initializers.get(bias_initializer),
+        kernel_regularizer=regularizers.get(kernel_regularizer),
+        bias_regularizer=regularizers.get(bias_regularizer),
+        activity_regularizer=regularizers.get(activity_regularizer),
+        **kwargs)
+
+  def call_weights(self):
+    """Calls any weights if the initializer is itself a layer."""
+    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
+      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
+    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
+      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
+
+  def call(self, inputs, training=None):
+    self.call_weights()
+    if training is None:
+      training = tf.keras.backend.learning_phase()
+
+    def dropped_inputs():
+      """Forward pass with dropout."""
+      # Clip magnitude of dropout rate, where we get the dropout rate alpha from
+      # the additive parameterization (Molchanov et al., 2017): for weight ~
+      # Normal(mu, sigma**2), the variance `sigma**2 = alpha * mu**2`.
+      mean = self.kernel.distribution.mean()
+      log_variance = tf.log(self.kernel.distribution.variance())
+      log_alpha = log_variance - tf.log(tf.square(mean) +
+                                        tf.keras.backend.epsilon())
+      log_alpha = tf.clip_by_value(log_alpha, -8., 8.)
+      log_variance = log_alpha + tf.log(tf.square(mean) +
+                                        tf.keras.backend.epsilon())
+
+      if inputs.shape.ndims <= 2:
+        means = tf.matmul(inputs, mean)
+        stddevs = tf.sqrt(
+            tf.matmul(tf.square(inputs), tf.exp(log_variance)) +
+            tf.keras.backend.epsilon())
+      else:
+        means = tf.tensordot(inputs, mean, [[-1], [0]])
+        stddevs = tf.sqrt(
+            tf.tensordot(tf.square(inputs), tf.exp(log_variance), [[-1], [0]]) +
+            tf.keras.backend.epsilon())
+      outputs = means + stddevs * tf.random_normal(tf.shape(stddevs))
+      if self.use_bias:
+        outputs = tf.nn.bias_add(outputs, self.bias)
+      if self.activation is not None:
+        outputs = self.activation(outputs)
+      return outputs
+
+    # Following tf.keras.Dropout, only apply variational dropout if training
+    # flag is True. The kernel must also be a random variable.
+    training_value = tf.contrib.util.constant_value(training)
+    if training_value is not None:
+      if training_value and isinstance(self.kernel, ed.RandomVariable):
+        return dropped_inputs()
+      else:
+        return super(DenseVariationalDropout, self).call(inputs)
+    else:
+      return tf.cond(tf.logical_and(training,
+                                    isinstance(self.kernel, ed.RandomVariable)),
+                     dropped_inputs,
+                     lambda: super(DenseVariationalDropout, self).call(inputs))
+
+
 @add_weight
 class LSTMCellReparameterization(tf.keras.layers.LSTMCell):
   """Bayesian LSTM cell class estimated via reparameterization.
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index bb8090bb0..a53b0ebb4 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -32,6 +32,75 @@
 
 class BayesTest(parameterized.TestCase, tf.test.TestCase):
 
+  @parameterized.parameters(
+      {"layer": bayes.Conv2DReparameterization,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "zeros",
+       "all_close": True},
+      {"layer": bayes.Conv2DReparameterization,
+       "kernel_initializer": "trainable_normal",
+       "bias_initializer": "zeros",
+       "all_close": False},
+      {"layer": bayes.Conv2DReparameterization,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "trainable_normal",
+       "all_close": False},
+      {"layer": bayes.Conv2DVariationalDropout,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "zeros",
+       "all_close": True},
+      {"layer": bayes.Conv2DVariationalDropout,
+       "kernel_initializer": "trainable_normal",
+       "bias_initializer": "zeros",
+       "all_close": False},
+      {"layer": bayes.Conv2DVariationalDropout,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "trainable_normal",
+       "all_close": False},
+  )
+  @test_utils.run_in_graph_and_eager_modes
+  def testConv2DKernel(self,
+                       layer,
+                       kernel_initializer,
+                       bias_initializer,
+                       all_close):
+    tf.keras.backend.set_learning_phase(1)  # training time
+    inputs = tf.to_float(np.random.rand(5, 4, 4, 12))
+    model = layer(4,
+                  kernel_size=2,
+                  kernel_initializer=kernel_initializer,
+                  bias_initializer=bias_initializer,
+                  activation=tf.nn.relu)
+    outputs1 = model(inputs)
+    outputs2 = model(inputs)
+    self.evaluate(tf.global_variables_initializer())
+    res1, res2 = self.evaluate([outputs1, outputs2])
+    self.assertEqual(res1.shape, (5, 3, 3, 4))
+    self.assertAllGreaterEqual(res1, 0.)
+    if all_close:
+      self.assertAllClose(res1, res2)
+    else:
+      self.assertNotAllClose(res1, res2)
+    model.get_config()
+
+  @parameterized.parameters(
+      {"layer": bayes.Conv2DReparameterization},
+      {"layer": bayes.Conv2DVariationalDropout},
+  )
+  @test_utils.run_in_graph_and_eager_modes()
+  def testConv2DModel(self, layer):
+    inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
+    model = tf.keras.Sequential([
+        layer(3, kernel_size=2, padding="SAME", activation=tf.nn.relu),
+        tf.keras.layers.Flatten(),
+        tf.keras.layers.Dense(2, activation=None),
+    ])
+    outputs = model(inputs, training=True)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(outputs)
+    self.assertEqual(res.shape, (3, 2))
+    self.assertLen(model.losses, 1)
+
   @test_utils.run_in_graph_and_eager_modes
   def testTrainableNormalStddevConstraint(self):
     layer = bayes.DenseReparameterization(
@@ -43,29 +112,46 @@ def testTrainableNormalStddevConstraint(self):
     res, _ = self.evaluate([stddev, out])
     self.assertAllGreater(res, 0.)
 
-  @parameterized.named_parameters(
-      {"testcase_name": "_no_uncertainty",
+  @parameterized.parameters(
+      {"layer": bayes.DenseReparameterization,
        "kernel_initializer": "zeros",
        "bias_initializer": "zeros",
        "all_close": True},
-      {"testcase_name": "_kernel_uncertainty",
+      {"layer": bayes.DenseReparameterization,
        "kernel_initializer": "trainable_normal",
        "bias_initializer": "zeros",
        "all_close": False},
-      {"testcase_name": "_bias_uncertainty",
+      {"layer": bayes.DenseReparameterization,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "trainable_normal",
+       "all_close": False},
+      {"layer": bayes.DenseVariationalDropout,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "zeros",
+       "all_close": True},
+      {"layer": bayes.DenseVariationalDropout,
+       "kernel_initializer": "trainable_normal",
+       "bias_initializer": "zeros",
+       "all_close": False},
+      {"layer": bayes.DenseVariationalDropout,
        "kernel_initializer": "zeros",
        "bias_initializer": "trainable_normal",
        "all_close": False},
   )
   @test_utils.run_in_graph_and_eager_modes
-  def testDenseReparameterizationKernel(
-      self, kernel_initializer, bias_initializer, all_close):
+  def testDenseKernel(self,
+                      layer,
+                      kernel_initializer,
+                      bias_initializer,
+                      all_close):
+    tf.keras.backend.set_learning_phase(1)  # training time
     inputs = tf.to_float(np.random.rand(5, 3, 12))
-    layer = bayes.DenseReparameterization(
-        4, kernel_initializer=kernel_initializer,
-        bias_initializer=bias_initializer, activation=tf.nn.relu)
-    outputs1 = layer(inputs)
-    outputs2 = layer(inputs)
+    model = layer(4,
+                  kernel_initializer=kernel_initializer,
+                  bias_initializer=bias_initializer,
+                  activation=tf.nn.relu)
+    outputs1 = model(inputs)
+    outputs2 = model(inputs)
     self.evaluate(tf.global_variables_initializer())
     res1, res2 = self.evaluate([outputs1, outputs2])
     self.assertEqual(res1.shape, (5, 3, 4))
@@ -74,45 +160,53 @@ def testDenseReparameterizationKernel(
       self.assertAllClose(res1, res2)
     else:
       self.assertNotAllClose(res1, res2)
-    layer.get_config()
+    model.get_config()
 
+  @parameterized.parameters(
+      {"layer": bayes.DenseReparameterization},
+      {"layer": bayes.DenseVariationalDropout},
+  )
   @test_utils.run_in_graph_and_eager_modes
-  def testDenseReparameterizationMean(self):
+  def testDenseMean(self, layer):
     """Tests that forward pass can use other values, e.g., posterior mean."""
+    tf.keras.backend.set_learning_phase(0)  # test time
     def take_mean(f, *args, **kwargs):
       """Sets random variable value to its mean."""
       rv = f(*args, **kwargs)
       rv._value = rv.distribution.mean()
       return rv
     inputs = tf.to_float(np.random.rand(5, 3, 7))
-    layer = bayes.DenseReparameterization(4,
-                                          activation=tf.nn.relu,
-                                          use_bias=False)
-    outputs1 = layer(inputs)
+    model = layer(4, activation=tf.nn.relu, use_bias=False)
+    outputs1 = model(inputs)
     with ed.interception(take_mean):
-      outputs2 = layer(inputs)
+      outputs2 = model(inputs)
     self.evaluate(tf.global_variables_initializer())
     res1, res2 = self.evaluate([outputs1, outputs2])
     self.assertEqual(res1.shape, (5, 3, 4))
     self.assertNotAllClose(res1, res2)
     self.assertAllClose(res2, np.zeros((5, 3, 4)), atol=1e-4)
 
+  @parameterized.parameters(
+      {"layer": bayes.DenseReparameterization},
+      {"layer": bayes.DenseVariationalDropout},
+  )
   @test_utils.run_in_graph_and_eager_modes()
-  def testDenseReparameterizationLoss(self):
+  def testDenseLoss(self, layer):
+    tf.keras.backend.set_learning_phase(1)  # training time
     features = tf.to_float(np.random.rand(5, 12))
     labels = tf.to_float(np.random.rand(5, 10))
-    layer = bayes.DenseReparameterization(10)
+    model = layer(10)
 
     # Imagine this is the 1st epoch.
     with tf.GradientTape(persistent=True) as tape:
-      predictions = layer(features)  # first call forces build
-      layer(features)  # ensure robustness after multiple calls
+      predictions = model(features)  # first call forces build
+      model(features)  # ensure robustness after multiple calls
       nll = tf.losses.mean_squared_error(labels, predictions)
-      kl = sum(layer.losses)
+      kl = sum(model.losses)
 
-    variables = [layer.kernel_initializer.mean, layer.kernel_initializer.stddev]
+    variables = [model.kernel_initializer.mean, model.kernel_initializer.stddev]
     for v in variables:
-      self.assertIn(v, layer.variables)
+      self.assertIn(v, model.variables)
 
     # This will be fine, since the layer was built inside this tape, and thus
     # the distribution init ops were inside this tape.
@@ -125,13 +219,13 @@ def testDenseReparameterizationLoss(self):
 
     # Imagine this is the 2nd epoch.
     with tf.GradientTape(persistent=True) as tape:
-      predictions = layer(features)  # build is not called
+      predictions = model(features)  # build is not called
       nll = tf.losses.mean_squared_error(labels, predictions)
-      kl = sum(layer.losses)
+      kl = sum(model.losses)
 
-    variables = [layer.kernel_initializer.mean, layer.kernel_initializer.stddev]
+    variables = [model.kernel_initializer.mean, model.kernel_initializer.stddev]
     for v in variables:
-      self.assertIn(v, layer.variables)
+      self.assertIn(v, model.variables)
 
     # This would fail, since the layer was built inside the tape from the 1st
     # epoch, and thus the distribution init ops were inside that tape instead of
@@ -143,8 +237,12 @@ def testDenseReparameterizationLoss(self):
     for grad in grads:
       self.assertIsNotNone(grad)
 
+  @parameterized.parameters(
+      {"layer": bayes.DenseReparameterization},
+      {"layer": bayes.DenseVariationalDropout},
+  )
   @test_utils.run_in_graph_and_eager_modes()
-  def testDenseReparameterizationModel(self):
+  def testDenseModel(self, layer):
     inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
     model = tf.keras.Sequential([
         tf.keras.layers.Conv2D(3,
@@ -152,17 +250,21 @@ def testDenseReparameterizationModel(self):
                                padding="SAME",
                                activation=tf.nn.relu),
         tf.keras.layers.Flatten(),
-        bayes.DenseReparameterization(2, activation=None),
+        layer(2, activation=None),
     ])
-    outputs = model(inputs)
+    outputs = model(inputs, training=True)
     self.evaluate(tf.global_variables_initializer())
     res = self.evaluate(outputs)
     self.assertEqual(res.shape, (3, 2))
     self.assertLen(model.losses, 1)
 
+  @parameterized.parameters(
+      {"layer": bayes.DenseReparameterization},
+      {"layer": bayes.DenseVariationalDropout},
+  )
   @test_utils.run_in_graph_and_eager_modes()
-  def testDenseReparameterizationSubclass(self):
-    class DenseReparameterizationSubclass(bayes.DenseReparameterization):
+  def testDenseSubclass(self, layer):
+    class DenseSubclass(layer):
       pass
 
     inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
@@ -172,62 +274,9 @@ class DenseReparameterizationSubclass(bayes.DenseReparameterization):
                                padding="SAME",
                                activation=tf.nn.relu),
         tf.keras.layers.Flatten(),
-        DenseReparameterizationSubclass(2, activation=None),
-    ])
-    outputs = model(inputs)
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(outputs)
-    self.assertEqual(res.shape, (3, 2))
-    self.assertLen(model.losses, 1)
-
-  @parameterized.named_parameters(
-      {"testcase_name": "_no_uncertainty",
-       "kernel_initializer": "zeros",
-       "bias_initializer": "zeros",
-       "all_close": True},
-      {"testcase_name": "_kernel_uncertainty",
-       "kernel_initializer": "trainable_normal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"testcase_name": "_bias_uncertainty",
-       "kernel_initializer": "zeros",
-       "bias_initializer": "trainable_normal",
-       "all_close": False},
-  )
-  @test_utils.run_in_graph_and_eager_modes
-  def testConv2DReparameterizationKernel(
-      self, kernel_initializer, bias_initializer, all_close):
-    inputs = tf.to_float(np.random.rand(5, 4, 4, 12))
-    layer = bayes.Conv2DReparameterization(
-        4,
-        kernel_size=2,
-        kernel_initializer=kernel_initializer,
-        bias_initializer=bias_initializer,
-        activation=tf.nn.relu)
-    outputs1 = layer(inputs)
-    outputs2 = layer(inputs)
-    self.evaluate(tf.global_variables_initializer())
-    res1, res2 = self.evaluate([outputs1, outputs2])
-    self.assertEqual(res1.shape, (5, 3, 3, 4))
-    self.assertAllGreaterEqual(res1, 0.)
-    if all_close:
-      self.assertAllClose(res1, res2)
-    else:
-      self.assertNotAllClose(res1, res2)
-    layer.get_config()
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testConv2DReparameterizationModel(self):
-    inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
-    model = tf.keras.Sequential([
-        bayes.Conv2DReparameterization(3,
-                                       kernel_size=2,
-                                       padding="SAME",
-                                       activation=tf.nn.relu),
-        tf.keras.layers.Flatten(),
-        tf.keras.layers.Dense(2, activation=None),
+        DenseSubclass(2, activation=None),
     ])
-    outputs = model(inputs)
+    outputs = model(inputs, training=True)
     self.evaluate(tf.global_variables_initializer())
     res = self.evaluate(outputs)
     self.assertEqual(res.shape, (3, 2))

From 72d3d3a49e29690adcd8a06c5b7b14642dacd955 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 17 Apr 2019 14:46:46 -0700
Subject: [PATCH 1909/2720]  - Added documentation to PPO implementation.  -
 PPO tests for non-optimizer functions.  - Added loads of vlogs, these are
 indispensable for debugging.  - NOTE: Pong still doesn't run with these
 changes, but goes further than before.

PiperOrigin-RevId: 244071832
---
 tensor2tensor/trax/rlax/ppo.py      | 448 ++++++++++++++++++----------
 tensor2tensor/trax/rlax/ppo_main.py |  11 +-
 tensor2tensor/trax/rlax/ppo_test.py | 405 ++++++++++++++++++++++++-
 3 files changed, 702 insertions(+), 162 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index e532b03e1..6a4cafc56 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -13,7 +13,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""PPO in JAX."""
+"""PPO in JAX.
+
+Notation:
+
+B, scalar  - batch size
+T, scalar  - number of time-steps in a trajectory, or the value of the padded
+             time-step dimension.
+OBS, tuple - shape of a singular observation from the environment.
+             Ex: For CartPole-v0 this is (4,) and Pong-v0 it's (210, 160, 3)
+A, scalar  - Number of actions, assuming a discrete space.
+
+Policy and Value function signatures:
+
+Policy Function :: [B, T] + OBS -> [B, T, A]
+Value  Function :: [B, T] + OBS -> [B, T, 1]
+
+i.e. the policy net should take a batch of *trajectories* and at each time-step
+in each batch deliver a probability distribution over actions.
+
+NOTE: It doesn't return logits, rather the expectation is that it return a
+normalized distribution instead.
+
+NOTE: The policy and value functions need to take care to not take into account
+future time-steps while deciding the actions (or value) for the current
+time-step.
+
+"""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -29,7 +55,6 @@
 from jax import lax
 from jax import numpy as np
 from jax import random as jax_random
-from jax import vmap
 import numpy as onp
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
@@ -45,13 +70,11 @@
 POLICY = "categorical-sampling"
 
 
-def policy_net(jax_rng_key,
+def policy_net(rng_key,
                batch_observations_shape,
                num_actions,
                bottom_layers=None):
   """A policy net function."""
-  key1, _ = jax_random.split(jax_rng_key)
-
   # Use the bottom_layers as the bottom part of the network and just add the
   # required layers on top of it.
   if bottom_layers is None:
@@ -60,17 +83,16 @@ def policy_net(jax_rng_key,
 
   net_init, net_apply = stax.serial(*bottom_layers)
 
-  _, net_params = net_init(key1, batch_observations_shape)
+  _, net_params = net_init(rng_key, batch_observations_shape)
   return net_params, net_apply
 
 
-def value_net(jax_rng_key,
+def value_net(rng_key,
               batch_observations_shape,
               num_actions,
               bottom_layers=None):
   """A value net function."""
   del num_actions
-  key1, _ = jax_random.split(jax_rng_key)
 
   if bottom_layers is None:
     bottom_layers = []
@@ -80,7 +102,7 @@ def value_net(jax_rng_key,
 
   net_init, net_apply = stax.serial(*bottom_layers)
 
-  _, net_params = net_init(key1, batch_observations_shape)
+  _, net_params = net_init(rng_key, batch_observations_shape)
   return net_params, net_apply
 
 
@@ -113,12 +135,23 @@ def collect_trajectories(env,
     observation = env.reset()
     observations.append(observation)
     while not done:
+      # Add a batch dimension and time dimension.
+      observation = observation[np.newaxis, np.newaxis, :]
+
       # Run the policy, pick an action.
       predictions = policy_net_apply(policy_net_params, observation)
 
-      # Greedy policy.
-      action = np.argmax(predictions)
-      if policy == "epsilon-greedy":
+      # Squeeze the added dimension.
+      predictions = np.squeeze(predictions)
+
+      # Policy can be run in one of the following ways:
+      #  - Greedy
+      #  - Epsilon-Greedy
+      #  - Categorical-Sampling
+      action = None
+      if policy == "greedy":
+        action = np.argmax(predictions)
+      elif policy == "epsilon-greedy":
         # A schedule for epsilon is 1/k where k is the episode number sampled.
         if onp.random.random() < epsilon:
           # Choose an action at random.
@@ -127,7 +160,10 @@ def collect_trajectories(env,
           # Return the best action.
           action = np.argmax(predictions)
       elif policy == "categorical-sampling":
+        # import pdb; pdb.set_trace()
         action = int(onp.argwhere(onp.random.multinomial(1, predictions) == 1))
+      else:
+        raise ValueError("Unknown policy: %s" % policy)
 
       # NOTE: Assumption, single batch.
       action = int(action)
@@ -149,24 +185,51 @@ def collect_trajectories(env,
 # This function can probably be simplified, ask how?
 # Can we do something much simpler than lax.pad, maybe np.pad?
 # Others?
+
+
+def get_padding_value(dtype):
+  """Returns the padding value given a dtype."""
+  padding_value = None
+  if dtype == np.uint8:
+    padding_value = np.uint8(0)
+  elif dtype == np.uint16:
+    padding_value = np.uint16(0)
+  elif dtype == np.float32:
+    padding_value = 0.0
+  else:
+    padding_value = 0
+  assert padding_value is not None
+  return padding_value
+
+
+# TODO(afrozm): Use np.pad instead and make jittable?
 def pad_trajectories(trajectories, boundary=10):
-  """Pad trajectories to a bucket length that is a multiple of boundary."""
+  """Pad trajectories to a bucket length that is a multiple of boundary.
 
-  # trajectories is a list of tuples of (observations, actions, rewards)
-  # observations's length is one more than actions and rewards
-  #
-  # i.e. observations = (o_0, o_1, ... o_{T-1}, o_T)
-  #           actions = (a_0, a_1, ... a_{T-1})
-  #           rewards = (r_0, r_1, ... r_{T-1})
+  Args:
+    trajectories: list[(observation, actions, rewards)], where each observation
+      is shaped (t+1,) + OBS and actions & rewards are shaped (t,), with the
+      length of the list being B (batch size).
+    boundary: int, bucket length, the actions and rewards are padded to integer
+      multiples of boundary.
 
-  # Given the above, let's compute max(T) over all trajectories.
-  t_max = max(o.shape[0] for (o, a, r) in trajectories)
+  Returns:
+    tuple: (padding lengths, reward_mask, padded_observations, padded_actions,
+        padded_rewards) where padded_observations is shaped (B, T+1) + OBS and
+        padded_actions, padded_rewards & reward_mask are shaped (B, T).
+        Where T is max(t) rounded up to an integer multiple of boundary.
+        padded_length is how much padding we've added and
+        reward_mask is 1s for actual rewards and 0s for the padding.
+  """
 
-  # t_max - 1 is rounded to the next multiple of `boundary`
+  # Let's compute max(t) over all trajectories.
+  t_max = max(r.shape[0] for (_, _, r) in trajectories)
+
+  # t_max is rounded to the next multiple of `boundary`
   boundary = int(boundary)
-  bucket_length = boundary * int(np.ceil(float(t_max - 1) / boundary))
+  bucket_length = boundary * int(np.ceil(float(t_max) / boundary))
 
-  # So all obs will be padded to t_max and actions and rewards to t_max - 1.
+  # So all obs will be padded to t_max + 1 and actions and rewards to t_max.
   padded_observations = []
   padded_actions = []
   padded_rewards = []
@@ -188,15 +251,18 @@ def pad_trajectories(trajectories, boundary=10):
     for _ in range(o.ndim - 1):
       padding_config.append((0, 0, 0))
     padding_config = tuple(padding_config)
-    padding_value = 0.0 if o.dtype == np.float32 else 0
+
+    padding_value = get_padding_value(o.dtype)
+    action_padding_value = get_padding_value(a.dtype)
+    reward_padding_value = get_padding_value(r.dtype)
+
     padded_obs = lax.pad(o, padding_value, padding_config)
     padded_observations.append(padded_obs)
 
     # Now pad actions and rewards.
     assert a.ndim == 1 and r.ndim == 1
     padding_config = ((0, num_to_pad, 0),)
-    action_padding_value = 0.0 if a.dtype == np.float32 else 0
-    reward_padding_value = 0.0 if r.dtype == np.float32 else 0
+
     padded_action = lax.pad(a, action_padding_value, padding_config)
     padded_actions.append(padded_action)
     padded_reward = lax.pad(r, reward_padding_value, padding_config)
@@ -210,63 +276,104 @@ def pad_trajectories(trajectories, boundary=10):
       padded_observations), np.stack(padded_actions), np.stack(padded_rewards)
 
 
-# TODO(afrozm): Make this batched by default.
-def rewards_to_go(rewards, reward_mask=1.0, gamma=0.99):
-  r"""r2g[t] = \sum_{l=0}^{\infty}(\gamma^l * r_{t+l})."""
-  time_steps = len(rewards)
-  # r2g[t] = r[t] + (gamma * r2g[t+1])
+def rewards_to_go(rewards, mask, gamma=0.99):
+  r"""Computes rewards to go.
+
+  Reward to go is defined as follows, the discounted reward that we have to
+  yet collect, going forward from this point, i.e.:
+
+  r2g_t = \sum_{l=0}^{\infty} (\gamma^{l} * reward_{t+l})
+
+  Args:
+    rewards: np.ndarray of shape (B, T) of rewards.
+    mask: np.ndarray of shape (B, T) of mask for the rewards.
+    gamma: float, discount factor.
 
-  # First initialize like:
-  # r2g[t] = r[t], for t = 0 to T-1
-  r2g = list(rewards)
+  Returns:
+    rewards to go, np.ndarray of shape (B, T).
+  """
+  B, T = rewards.shape  # pylint: disable=invalid-name
+
+  # [[1, g, g**2, ... g**T-1]]
+  # Not jittable, T should be a compile time constant.
+  # gammas = np.geomspace(1, g**T, T, endpoint=False).reshape(1, T)
 
-  # Then add the discounted version of the next time-step.
-  # i = [T-2 .. 0]
-  for i in range(time_steps - 2, -1, -1):
-    r2g[i] += gamma * r2g[i + 1]
+  # Get a geometric progression of gamma, of length T.
+  gammas = [gamma**t for t in range(T)]
+  gammas = np.array(gammas).reshape((1, T))
 
-  # Makes this back into JAX's DeviceArray
-  r2g = np.stack(list(r2g))
+  # Discounted rewards.
+  undiscounted_rewards = rewards * mask  # (B, T)
+  discounted_rewards = undiscounted_rewards * gammas  # (B, T)
 
-  return r2g * reward_mask
+  # Get rewards to go at first time-step.
+  r0 = np.sum(discounted_rewards, axis=1)  # (B,)
+  assert r0.shape == (B,)
+
+  rs = [r0]
+
+  # Now compute the other advantages wrt the first one.
+  for t in range(1, T):
+    rs.append((rs[-1] - undiscounted_rewards[:, t - 1]) / gamma)
+
+  # len(rs) is T and each element is (B,), this makes it (B, T)
+  return np.stack(rs, axis=1)
 
 
-# TODO(afrozm): Make this batched by default.
 @functools.partial(jit, static_argnums=(0,))
 def value_loss(value_net_apply,
                value_net_params,
                observations,
                rewards,
-               reward_mask=1.0,
+               reward_mask,
                gamma=0.99):
-  """L2 loss on the value function's outputs."""
-
-  # Capturing the value_net_apply from the parent function's scope.
-  # See: https://github.com/google/jax/issues/183
-  def _value_function_loss_trajectory(value_net_params,
-                                      observations,
-                                      rewards,
-                                      reward_mask=1.0,
-                                      gamma=0.99):
-    """Compute the actual loss for a trajectory."""
-    r2g = rewards_to_go(rewards, reward_mask=reward_mask, gamma=gamma)
-    v = value_net_apply(value_net_params, observations[:-1])
-    v = np.squeeze(v) * reward_mask
-    loss = v - r2g
-    return np.sum(loss**2)
-
-  batched_value_function_loss_trajectory = vmap(
-      _value_function_loss_trajectory, in_axes=(None, 0, 0, 0), out_axes=0)
-
-  return np.mean(
-      batched_value_function_loss_trajectory(
-          value_net_params, observations, rewards, reward_mask, gamma=gamma))
-
-
-def deltas(predicted_values, rewards, reward_mask, gamma=0.99):
-  r"""\delta_t = \sum_{l = 0}^{\infty}(r_t + \gamma * V(s_{t+1}) - V(s_t))."""
-  # predicted_values are application of value net only the observations.
-  # B x T+1
+  """Computes the value loss.
+
+  Args:
+    value_net_apply: value net apply function with signature (params, ndarray of
+      shape (B, T+1) + OBS) -> ndarray(B, T+1, 1)
+    value_net_params: params of value_net_apply.
+    observations: np.ndarray of shape (B, T+1) + OBS
+    rewards: np.ndarray of shape (B, T) of rewards.
+    reward_mask: np.ndarray of shape (B, T), the mask over rewards.
+    gamma: float, discount factor.
+
+  Returns:
+    The average L2 value loss, averaged over instances where reward_mask is 1.
+  """
+
+  B, T = rewards.shape  # pylint: disable=invalid-name
+  assert (B, T) == reward_mask.shape
+  assert (B, T + 1) == observations.shape[:2]
+
+  r2g = rewards_to_go(rewards, reward_mask, gamma=gamma)  # (B, T)
+  # NOTE: observations is (B, T+1) + OBS, value_prediction is (B, T+1, 1)
+  value_prediction = value_net_apply(value_net_params, observations)
+  assert (B, T + 1, 1) == value_prediction.shape
+  value_prediction = np.squeeze(value_prediction, axis=2)  # (B, T+1)
+  value_prediction = value_prediction[:, :-1] * reward_mask  # (B, T)
+  loss = (value_prediction - r2g)**2
+  # Take an average on only the points where mask != 0.
+  return np.sum(loss) / np.sum(reward_mask)
+
+
+def deltas(predicted_values, rewards, mask, gamma=0.99):
+  r"""Computes TD-residuals from V(s) and rewards.
+
+  Where a `delta`, i.e. a td-residual is defined as:
+
+  delta_{b,t} = r_{b,t} + \gamma * v_{b,t+1} - v_{b,t}.
+
+  Args:
+    predicted_values: ndarray of shape (B, T+1). NOTE: Expects axis 2 was
+      squeezed. These represent V(s_bt) for b < B and t < T+1
+    rewards: ndarray of shape (B, T) of rewards.
+    mask: ndarray of shape (B, T) of mask for rewards.
+    gamma: float, discount factor.
+
+  Returns:
+    ndarray of shape (B, T) of one-step TD-residuals.
+  """
 
   # `d`s are basically one-step TD residuals.
   d = []
@@ -275,46 +382,41 @@ def deltas(predicted_values, rewards, reward_mask, gamma=0.99):
     d.append(rewards[:, t] + (gamma * predicted_values[:, t + 1]) -
              predicted_values[:, t])
 
-  return np.array(d).T * reward_mask
+  return np.array(d).T * mask
 
 
-def gae_advantages(td_deltas, reward_mask, lambda_=0.95, gamma=0.99):
-  r"""A_t = \sum_{l=0}^{\infty}(\gamma * \lambda)^{l}(\delta_{t+l})."""
-  _, T = td_deltas.shape  # pylint: disable=invalid-name
-  gl = lambda_ * gamma
+def gae_advantages(td_deltas, mask, lambda_=0.95, gamma=0.99):
+  r"""Computes the GAE advantages given the one step TD-residuals.
 
-  # [[1, gl, gl**2, ... gl**T-1]]
-  # Not jittable, T should be a compile time constant.
-  # gl_gp = np.geomspace(1, gl**T, T, endpoint=False).reshape(1, T)
-  gl_geometric_progression = [1]
-  for _ in range(1, T):
-    gl_geometric_progression.append(gl_geometric_progression[-1] * gl)
-  gl_gp = np.array(gl_geometric_progression)
-  gl_gp = gl_gp.reshape((1, T))
-
-  # td_deltas * gl_gp
-  deltas_gl_gp = td_deltas * gl_gp
-
-  # A0 - advantage for 0th time-step, across all batches.
-  As = []  # pylint: disable=invalid-name
-  A0 = np.sum(deltas_gl_gp, axis=1)  # (B,)  # pylint: disable=invalid-name
-  As.append(A0)
-
-  # Now compute the other advantages.
-  for t in range(1, T):
-    As.append((As[-1] - td_deltas[:, t - 1]) / gl)
+  The formula for a GAE advantage estimator is as follows:
+
+  A_{bt} = \sum_{l=0}^{\infty}(\gamma * \lambda)^{l}(\delta_{b,t+l}).
+
+  Internally we just call rewards_to_go, since it is the same computation.
 
-  return np.stack(As).T * reward_mask
+  Args:
+    td_deltas: np.ndarray of shape (B, T) of one step TD-residuals.
+    mask: np.ndarray of shape (B, T) of mask for the residuals. It maybe the
+      case that the `td_deltas` are already masked correctly since they are
+      produced by `deltas(...)`
+    lambda_: float, lambda parameter for GAE estimators.
+    gamma: float, lambda parameter for GAE estimators.
+
+  Returns:
+    GAE advantage estimates.
+  """
+
+  return rewards_to_go(td_deltas, mask, lambda_ * gamma)
 
 
 def chosen_probabs(probab_observations, actions):
   """Picks out the probabilities of the actions along batch and time-steps.
 
   Args:
-    probab_observations: `[B, T, #actions]` ndarray, where
+    probab_observations: ndarray of shape `[B, T, A]`, where
       probab_observations[b, t, i] contains the probability of action = i at the
       t^th time-step in the b^th trajectory.
-    actions: `[B, T]` ndarray, with each entry in [0, #actions) denoting which
+    actions: ndarray of shape `[B, T]`, with each entry in [0, A) denoting which
       action was chosen in the b^th trajectory's t^th time-step.
 
   Returns:
@@ -324,32 +426,40 @@ def chosen_probabs(probab_observations, actions):
   return probab_observations[np.arange(b)[:, None], np.arange(t), actions]
 
 
-def probab_ratios(policy_net_apply, old_policy_params, new_policy_params,
-                  observations, actions, reward_mask):
-  """Calculates the probaility ratios for each time-step in a trajectory."""
-  p_old = policy_net_apply(old_policy_params, observations)
-  p_new = policy_net_apply(new_policy_params, observations)
+def compute_probab_ratios(p_old, p_new, actions, reward_mask):
+  """Computes the probability ratios for each time-step in a trajectory.
+
+  Args:
+    p_old: ndarray of shape [B, T, A] of the probabilities that the policy
+      network assigns to all the actions at each time-step in each batch using
+      the old parameters.
+    p_new: ndarray of shape [B, T, A], same as above, but using new policy
+      network parameters.
+    actions: ndarray of shape [B, T] where each element is from [0, A).
+    reward_mask: ndarray of shape [B, T] masking over probabilities.
 
+  Returns:
+    probab_ratios: ndarray of shape [B, T], where
+    probab_ratios_{b,t} = p_new_{b,t,action_{b,t}} / p_old_{b,t,action_{b,t}}
+  """
   bp_old = chosen_probabs(p_old, actions)
   bp_new = chosen_probabs(p_new, actions)
 
   # Add a small number to bp_old, where reward_mask is 0, this is just to help
   # never to divide by 0.
   bp_old = bp_old + (0.1 * np.abs(reward_mask - 1))
+  probab_ratios = (bp_new * reward_mask) / bp_old
+  return probab_ratios
 
-  ret_val = (bp_new * reward_mask) / bp_old
 
-  return ret_val
+def clipped_probab_ratios(probab_ratios, reward_mask, epsilon=0.2):
+  return reward_mask * np.clip(probab_ratios, 1 - epsilon, 1 + epsilon)
 
 
-def clipped_probab_ratios(bpr, reward_mask, epsilon=0.2):
-  return reward_mask * np.clip(bpr, 1 - epsilon, 1 + epsilon)
-
-
-def clipped_objective(bpr, adv, reward_mask, epsilon=0.2):
-  c1 = bpr * adv
-  c2 = clipped_probab_ratios(bpr, reward_mask, epsilon=epsilon) * adv
-  return np.minimum(c1, c2)
+def clipped_objective(probab_ratios, advantages, reward_mask, epsilon=0.2):
+  c1 = probab_ratios * reward_mask
+  c2 = clipped_probab_ratios(probab_ratios, reward_mask, epsilon=epsilon)
+  return np.minimum(c1, c2) * advantages
 
 
 @functools.partial(jit, static_argnums=(0, 3))
@@ -366,17 +476,37 @@ def ppo_loss(policy_net_apply,
              lambda_=0.95,
              epsilon=0.2):
   """PPO objective, with an eventual minus sign."""
-  # V(s_t) forall s & t
-  value_function = np.squeeze(
-      value_net_apply(value_net_params, padded_observations))
-  td_deltas = deltas(value_function, padded_rewards, reward_mask, gamma=gamma)
+  # (B, T+1, 1)
+  predicted_values = value_net_apply(value_net_params, padded_observations)
+
+  # (B, T)
+  td_deltas = deltas(
+      np.squeeze(predicted_values, axis=2),
+      padded_rewards,
+      reward_mask,
+      gamma=gamma)
+
+  # (B, T)
   advantages = gae_advantages(
       td_deltas, reward_mask, lambda_=lambda_, gamma=gamma)
-  ratios = probab_ratios(policy_net_apply, old_policy_params, new_policy_params,
-                         padded_observations, padded_actions, reward_mask)
-  clipped_loss = clipped_objective(
+
+  # probab_actions_{old,new} are both (B, T, A)
+  probab_actions_old = policy_net_apply(old_policy_params, padded_observations)
+  probab_actions_new = policy_net_apply(new_policy_params, padded_observations)
+
+  # (B, T)
+  ratios = compute_probab_ratios(probab_actions_old, probab_actions_new,
+                                 padded_actions, reward_mask)
+
+  # (B, T)
+  objective = clipped_objective(
       ratios, advantages, reward_mask, epsilon=epsilon)
-  return -np.sum(clipped_loss)
+
+  # ()
+  average_objective = np.sum(objective) / np.sum(reward_mask)
+
+  # Loss is negative objective.
+  return -average_objective
 
 
 @functools.partial(jit, static_argnums=(2, 3, 5))
@@ -462,7 +592,9 @@ def training_loop(
 
   env = env if env is not None else gym.make(env_name)
 
-  batch_observations_shape = (-1,) + env.observation_space.shape
+  # Batch Observations Shape = [-1, -1] + OBS, because we will eventually call
+  # policy and value networks on shape [B, T] +_OBS
+  batch_observations_shape = (-1, -1) + env.observation_space.shape
 
   assert isinstance(env.action_space, gym.spaces.Discrete)
   num_actions = env.action_space.n
@@ -476,8 +608,9 @@ def training_loop(
 
   policy_net_params, policy_net_apply = policy_net_fun(
       key1, batch_observations_shape, num_actions)
-  value_net_params, value_net_apply = value_net_fun(
-      key2, batch_observations_shape, num_actions)
+  value_net_params, value_net_apply = value_net_fun(key2,
+                                                    batch_observations_shape,
+                                                    num_actions)
 
   # Initialize the optimizers.
   assert policy_optimizer_fun and value_optimizer_fun
@@ -497,20 +630,33 @@ def training_loop(
         epsilon=(10.0 / (i + 10.0)))  # this is a different epsilon.
 
     avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
+    max_reward = max(np.sum(traj[2]) for traj in trajs)
     average_rewards.append(avg_reward)
 
-    logging.debug("Average sum rewards [%0.2f]", avg_reward)
-    logging.debug("Collecting trajectories took %0.2f msec.", get_time(t))
-    logging.debug("Average Trajectory size [%0.2f]",
-                  float(sum(len(traj[0]) for traj in trajs)) / len(trajs))
+    logging.vlog(1, "Rewards average=[%0.2f], max=[%0.2f]", avg_reward,
+                 max_reward)
+    logging.vlog(1, "Collecting trajectories took %0.2f msec.", get_time(t))
+    logging.vlog(1, "Trajectory Length average=[%0.2f], max=[%0.2f]",
+                 float(sum(len(traj[0]) for traj in trajs)) / len(trajs),
+                 max(len(traj[0]) for traj in trajs))
 
     t = time.time()
     (_, reward_mask, padded_observations, padded_actions,
      padded_rewards) = pad_trajectories(
          trajs, boundary=20)
 
-    logging.debug("Padding trajectories took %0.2f msec.", get_time(t))
-    logging.debug("Padded Actions' shape [%s]", str(padded_actions.shape))
+    logging.vlog(1, "Padding trajectories took %0.2f msec.", get_time(t))
+    logging.vlog(1, "Padded Observations' shape [%s]",
+                 str(padded_observations.shape))
+    logging.vlog(1, "Padded Actions' shape [%s]", str(padded_actions.shape))
+    logging.vlog(1, "Padded Rewards' shape [%s]", str(padded_rewards.shape))
+
+    # Some assertions.
+    B, T = padded_actions.shape  # pylint: disable=invalid-name
+    assert (B, T) == padded_rewards.shape
+    assert (B, T) == reward_mask.shape
+    assert (B, T + 1) == padded_observations.shape[:2]
+    assert (B, T + 1) + env.observation_space.shape == padded_observations.shape
 
     # Linear annealing from 0.1 to 0.0
     epsilon = 0.1 if epochs == 1 else 0.1 * (1.0 - (i / (epochs - 1)))
@@ -524,7 +670,7 @@ def training_loop(
         reward_mask,
         gamma=GAMMA)
 
-    logging.debug("Calculating value loss took %0.2f msec.", get_time(t))
+    logging.vlog(1, "Calculating value loss took %0.2f msec.", get_time(t))
     value_losses.append(cur_value_loss)
 
     t = time.time()
@@ -542,11 +688,11 @@ def training_loop(
         lambda_=LAMBDA,
         epsilon=epsilon)
     # ppo_loss = 11.00110011
-    logging.debug("Calculating PPO loss took %0.2f msec.", get_time(t))
+    logging.vlog(1, "Calculating PPO loss took %0.2f msec.", get_time(t))
     ppo_objective.append(-cur_ppo_loss)
 
     # Run optimizers.
-    logging.debug("PPO Optimization")
+    logging.vlog(1, "PPO Optimization")
     t1 = time.time()
 
     for j in range(num_optimizer_steps):
@@ -585,16 +731,16 @@ def training_loop(
             gamma=GAMMA,
             lambda_=LAMBDA,
             epsilon=epsilon)
-        logging.debug("One PPO grad desc took: %0.2f msec", get_time(t, t2))
-        logging.debug("PPO loss [%10.2f] -> [%10.2f]", cur_ppo_loss,
-                      new_ppo_loss)
+        logging.vlog(1, "One PPO grad desc took: %0.2f msec", get_time(t, t2))
+        logging.vlog(1, "PPO loss [%10.2f] -> [%10.2f]", cur_ppo_loss,
+                     new_ppo_loss)
       # Update the params.
       policy_net_params = new_policy_net_params
 
-    logging.debug("Total PPO loss reduction [%0.2f]%%",
-                  (100 * (cur_ppo_loss - new_ppo_loss) / np.abs(cur_ppo_loss)))
+    logging.vlog(1, "Total PPO loss reduction [%0.2f]%%",
+                 (100 * (cur_ppo_loss - new_ppo_loss) / np.abs(cur_ppo_loss)))
 
-    logging.debug("Value Optimization")
+    logging.vlog(1, "Value Optimization")
 
     for j in range(num_optimizer_steps):
       t = time.time()
@@ -618,14 +764,14 @@ def training_loop(
             padded_rewards,
             reward_mask,
             gamma=GAMMA)
-        logging.debug("One value grad desc took: %0.2f msec", get_time(t, t2))
-        logging.debug("Value loss [%10.2f] -> [%10.2f]", cur_value_loss,
-                      new_value_loss)
-    logging.debug("Total value loss reduction [%0.2f]%%",
-                  (100 *
-                   (cur_value_loss - new_value_loss) / np.abs(cur_value_loss)))
+        logging.vlog(1, "One value grad desc took: %0.2f msec", get_time(t, t2))
+        logging.vlog(1, "Value loss [%10.2f] -> [%10.2f]", cur_value_loss,
+                     new_value_loss)
+    logging.vlog(1, "Total value loss reduction [%0.2f]%%",
+                 (100 *
+                  (cur_value_loss - new_value_loss) / np.abs(cur_value_loss)))
 
-    logging.debug("Grad desc took %0.2f msec", get_time(t1))
+    logging.vlog(1, "Grad desc took %0.2f msec", get_time(t1))
 
     # Set the optimized params to new params.
     policy_net_params = trax_opt.get_params(ppo_opt_state)
@@ -636,9 +782,9 @@ def training_loop(
         "value loss [%10.2f], took [%10.2f msec]", i, avg_reward, new_ppo_loss,
         new_value_loss, get_time(t0))
 
-  logging.debug("value_losses: %s", np.stack(value_losses))
-  logging.debug("ppo_objective: %s", np.stack(ppo_objective))
-  logging.debug("average_rewards: %s", average_rewards)
+  logging.vlog(1, "value_losses: %s", np.stack(value_losses))
+  logging.vlog(1, "ppo_objective: %s", np.stack(ppo_objective))
+  logging.vlog(1, "average_rewards: %s", average_rewards)
 
   return ((policy_net_params, value_net_params), average_rewards,
           np.stack(value_losses), np.stack(ppo_objective))
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 11708bcc6..e4d58cf89 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -29,10 +29,12 @@
 
 FLAGS = flags.FLAGS
 
-flags.DEFINE_string("env", "CartPole-v0", "Name of the environment to make.")
+flags.DEFINE_string("env_name", None, "Name of the environment to make.")
+flags.DEFINE_string("t2t_gym_env", None, "Name of the T2TGymEnv to make.")
 flags.DEFINE_integer("epochs", 100, "Number of epochs to run for.")
 flags.DEFINE_integer("random_seed", 0, "Random seed.")
 flags.DEFINE_integer("log_level", logging.INFO, "Log level.")
+flags.DEFINE_integer("batch_size", 32, "Batch of trajectories needed.")
 
 
 def common_stax_layers():
@@ -43,13 +45,18 @@ def main(argv):
   del argv
   logging.set_verbosity(FLAGS.log_level)
   bottom_layers = common_stax_layers()
+
+  if FLAGS.env_name == "Pong-v0":
+    bottom_layers = [stax.Flatten(2)] + bottom_layers
+
   ppo.training_loop(
-      env_name=FLAGS.env,
+      env_name=FLAGS.env_name,
       epochs=FLAGS.epochs,
       policy_net_fun=functools.partial(
           ppo.policy_net, bottom_layers=bottom_layers),
       value_net_fun=functools.partial(
           ppo.value_net, bottom_layers=bottom_layers),
+      batch_size=FLAGS.batch_size,
       random_seed=FLAGS.random_seed)
 
 
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 1d5154688..758555e7f 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -19,23 +19,410 @@
 from __future__ import division
 from __future__ import print_function
 
+import jax
+from jax import random as jax_random
 import numpy as np
+from tensor2tensor.trax import trax
+from tensor2tensor.trax.rlax import fake_env
 from tensor2tensor.trax.rlax import ppo
+from tensor2tensor.trax.stax import stax_base as stax
 from tensorflow import test
 
 
 class PpoTest(test.TestCase):
 
+  def setUp(self):
+    self.rng_key = trax.get_random_number_generator_and_set_seed(0)
+
+  def test_policy_net(self):
+    observation_shape = (3, 4)
+    num_actions = 2
+    policy_params, policy_apply = ppo.policy_net(
+        self.rng_key,
+        (-1, -1) + observation_shape,
+        num_actions,
+        # flatten except batch and time
+        # step dimensions.
+        [stax.Flatten(2)])
+
+    # Generate a batch of observations.
+    batch = 2
+    time_steps = 10
+    batch_of_observations = np.random.uniform(
+        size=(batch, time_steps) + observation_shape)
+
+    # Apply the policy net on observations
+    policy_output = policy_apply(policy_params, batch_of_observations)
+
+    # Verify certain expectations on the output.
+    self.assertEqual((batch, time_steps, num_actions), policy_output.shape)
+
+    # Also last axis normalizes to 1, since these are probabilities.
+    sum_actions = np.sum(policy_output, axis=-1)
+    self.assertAllClose(np.ones_like(sum_actions), sum_actions)
+
+  def test_value_net(self):
+    observation_shape = (3, 4, 5)
+    num_actions = 2
+    value_params, value_apply = ppo.value_net(self.rng_key,
+                                              (-1, -1) + observation_shape,
+                                              num_actions, [stax.Flatten(2)])
+    batch = 2
+    time_steps = 10
+    batch_of_observations = np.random.uniform(
+        size=(batch, time_steps) + observation_shape)
+    value_output = value_apply(value_params, batch_of_observations)
+
+    # NOTE: The extra dimension at the end because of Dense(1).
+    self.assertEqual((batch, time_steps, 1), value_output.shape)
+
+  def test_collect_trajectories(self):
+    observation_shape = (2, 3, 4)
+    num_actions = 2
+    policy_params, policy_apply = ppo.policy_net(
+        self.rng_key,
+        (-1, -1) + observation_shape,
+        num_actions,
+        # flatten except batch and time
+        # step dimensions.
+        [stax.Flatten(2)])
+
+    # We'll get done at time-step #10, starting from 0, therefore in 11 steps.
+    done_time_step = 5
+    env = fake_env.FakeEnv(
+        observation_shape, num_actions, done_time_step=done_time_step)
+
+    num_trajectories = 5
+    trajectories = ppo.collect_trajectories(
+        env,
+        policy_apply,
+        policy_params,
+        num_trajectories,
+        policy="categorical-sampling")
+
+    # Number of trajectories is as expected.
+    self.assertEqual(num_trajectories, len(trajectories))
+
+    # Shapes of observations, actions and rewards are as expected.
+    for observations, actions, rewards in trajectories:
+      # observations are one more in number than rewards or actions.
+      self.assertEqual((done_time_step + 2,) + observation_shape,
+                       observations.shape)
+      self.assertEqual((done_time_step + 1,), actions.shape)
+      self.assertEqual((done_time_step + 1,), rewards.shape)
+
+  def test_pad_trajectories(self):
+    observation_shape = (2, 3, 4)
+    trajectories = []
+    num_trajectories = 7
+    num_actions = 10
+
+    # Time-steps are between [min_allowable_time_step, max_allowable_time_step]
+    max_allowable_time_step = 19
+    min_allowable_time_step = 5
+
+    # The actual max we see in the data.
+    max_time_step = -1
+
+    # Bucket length.
+    bucket_length = 15
+
+    # Make `num_trajectories` random trajectories.
+    for i in range(num_trajectories):
+      time_steps = np.random.randint(min_allowable_time_step,
+                                     max_allowable_time_step + 1)
+      if time_steps > max_time_step:
+        max_time_step = time_steps
+      observations = np.random.randint(
+          0, 255, size=(time_steps + 1,) + observation_shape).astype(np.uint8)
+      rewards = np.random.uniform(size=(time_steps,)).astype(np.float32)
+      actions = np.random.randint(
+          0, num_actions, size=(time_steps,)).astype(np.int32)
+      trajectories.append((observations, rewards, actions))
+
+    # Now pad these trajectories.
+    padded_trajectories = ppo.pad_trajectories(
+        trajectories, boundary=bucket_length)
+
+    # Expected padding.
+    i = 1
+    while i * bucket_length < max_time_step:
+      i += 1
+    expected_padding = i * bucket_length
+
+    # Get the padded objects.
+    (pad_lengths, reward_mask, padded_observations, padded_actions,
+     padded_rewards) = padded_trajectories
+
+    # Expectations on the padded shapes.
+    self.assertEqual(padded_observations.shape, (
+        num_trajectories,
+        expected_padding + 1,
+    ) + observation_shape)
+    self.assertEqual(padded_actions.shape, (num_trajectories, expected_padding))
+    self.assertEqual(padded_rewards.shape, (num_trajectories, expected_padding))
+    self.assertEqual(reward_mask.shape, (num_trajectories, expected_padding))
+
+    # Assert that the padding lengths and reward mask are consistent.
+    self.assertAllEqual(
+        np.full((num_trajectories,), expected_padding),
+        np.array(np.sum(reward_mask, axis=1)) + pad_lengths)
+
   def test_rewards_to_go(self):
-    time_steps = 4
-    # [1., 1., 1., 1.]
-    rewards = np.ones((time_steps,))
-    # No discounting.
-    self.assertAllEqual(ppo.rewards_to_go(rewards, gamma=1.0),
-                        np.array([4., 3., 2., 1.]))
-    # Discounting.
-    self.assertAllEqual(ppo.rewards_to_go(rewards, gamma=0.5),
-                        np.array([1.875, 1.75, 1.5, 1.]))
+    rewards = np.array([
+        [1, 2, 4, 8, 16, 32, 64, 128],
+        [1, 1, 1, 1, 1, 1, 1, 1],
+    ])
+
+    rewards_mask = np.array([
+        [1, 1, 1, 1, 1, 0, 0, 0],
+        [1, 1, 1, 1, 1, 1, 1, 0],
+    ])
+
+    gamma = 0.5
+
+    rewards_to_go = ppo.rewards_to_go(rewards, rewards_mask, gamma)
+
+    self.assertAllEqual(
+        np.array([
+            [5, 8, 12, 16, 16, 0, 0, 0],
+            [1.984375, 1.96875, 1.9375, 1.875, 1.75, 1.5, 1.0, 0],
+        ]), rewards_to_go)
+
+  def test_value_loss(self):
+    rewards = np.array([
+        [1, 2, 4, 8, 16, 32, 64, 128],
+        [1, 1, 1, 1, 1, 1, 1, 1],
+    ])
+
+    rewards_mask = np.array([
+        [1, 1, 1, 1, 1, 0, 0, 0],
+        [1, 1, 1, 1, 1, 1, 1, 0],
+    ])
+
+    gamma = 0.5
+
+    # Random observations and a value function that returns a constant value.
+    # NOTE: Observations have an extra time-step.
+    B, T = rewards.shape  # pylint: disable=invalid-name
+    observation_shape = (210, 160, 3)  # atari pong
+    random_observations = np.random.uniform(size=(B, T + 1) + observation_shape)
+
+    def value_net_apply(params, observations):
+      del params
+      # pylint: disable=invalid-name
+      B, T_p_1, OBS = (observations.shape[0], observations.shape[1],
+                       observations.shape[2:])
+      del OBS
+      return np.ones((B, T_p_1, 1))
+      # pylint: enable=invalid-name
+
+    with jax.disable_jit():
+      value_loss = ppo.value_loss(
+          value_net_apply, [],
+          random_observations,
+          rewards,
+          rewards_mask,
+          gamma=gamma)
+
+    self.assertNear(53.3637084961, value_loss, 1e-6)
+
+  def test_deltas(self):
+    rewards = np.array([
+        [1, 2, 4, 8, 16, 32, 64, 128],
+        [1, 1, 1, 1, 1, 1, 1, 1],
+    ])
+
+    rewards_mask = np.array([
+        [1, 1, 1, 1, 1, 0, 0, 0],
+        [1, 1, 1, 1, 1, 1, 1, 0],
+    ])
+
+    B, T = rewards.shape  # pylint: disable=invalid-name
+
+    # Say, all predicted values are 1.
+    predicted_values = np.ones((B, T + 1))
+
+    gamma = 1.0
+
+    td_residuals = ppo.deltas(predicted_values, rewards, rewards_mask, gamma)
+
+    # With V(s) being the same for all s, td_residuals should be
+    # equal to the rewards + (\gamma - 1)*v(s), masked in the right places.
+    truncated_pv = predicted_values[:, :-1]
+    masked_rewards = rewards * rewards_mask
+    expected_residuals = (masked_rewards +
+                          (gamma - 1) * truncated_pv) * rewards_mask
+    self.assertAllEqual(expected_residuals, td_residuals)
+
+    gamma = 0.5
+    td_residuals = ppo.deltas(predicted_values, rewards, rewards_mask, gamma)
+    expected_residuals = (masked_rewards +
+                          (gamma - 1) * truncated_pv) * rewards_mask
+    self.assertAllEqual(expected_residuals, td_residuals)
+
+  def test_gae_advantages(self):
+    td_deltas = np.array([
+        [1, 2, 4, 8, 16, 32, 64, 128],
+        [1, 1, 1, 1, 1, 1, 1, 1],
+    ])
+
+    rewards_mask = np.array([
+        [1, 1, 1, 1, 1, 0, 0, 0],
+        [1, 1, 1, 1, 1, 1, 1, 0],
+    ])
+
+    gamma = 0.5
+    lambda_ = 1.0
+
+    expected_gae_advantages = np.array([
+        [5, 8, 12, 16, 16, 0, 0, 0],
+        [1.984375, 1.96875, 1.9375, 1.875, 1.75, 1.5, 1.0, 0],
+    ])
+
+    gae_advantages = ppo.gae_advantages(td_deltas * rewards_mask, rewards_mask,
+                                        lambda_, gamma)
+    self.assertAllEqual(expected_gae_advantages, gae_advantages)
+
+    gamma = 1.0
+    lambda_ = 0.5
+
+    gae_advantages = ppo.gae_advantages(td_deltas * rewards_mask, rewards_mask,
+                                        lambda_, gamma)
+    self.assertAllEqual(expected_gae_advantages, gae_advantages)
+
+  def test_chosen_probabs(self):
+    # Shape (2, 2, 3)
+    probab_observations = np.array([[[0.1, 0.2, 0.7], [0.4, 0.1, 0.5]],
+                                    [[0.3, 0.1, 0.6], [0.1, 0.1, 0.8]]])
+
+    # Shape (2, 2)
+    actions = np.array([[1, 2], [0, 1]])
+
+    chosen_probabs = ppo.chosen_probabs(probab_observations, actions)
+
+    self.assertAllEqual(np.array([[0.2, 0.5], [0.3, 0.1]]), chosen_probabs)
+
+  def test_compute_probab_ratios(self):
+    p_old = np.array([[
+        [0.1, 0.2, 0.6, 0.1],
+        [0.4, 0.1, 0.4, 0.1],
+        [0.3, 0.1, 0.5, 0.1],
+        [0.1, 0.2, 0.6, 0.1],
+    ],
+                      [
+                          [0.3, 0.1, 0.5, 0.1],
+                          [0.1, 0.1, 0.4, 0.4],
+                          [0.3, 0.1, 0.5, 0.1],
+                          [0.1, 0.2, 0.6, 0.1],
+                      ]])
+
+    p_new = np.array([[
+        [0.3, 0.1, 0.5, 0.1],
+        [0.4, 0.1, 0.1, 0.3],
+        [0.1, 0.2, 0.1, 0.6],
+        [0.3, 0.1, 0.5, 0.1],
+    ],
+                      [
+                          [0.1, 0.2, 0.1, 0.6],
+                          [0.1, 0.1, 0.2, 0.6],
+                          [0.3, 0.1, 0.3, 0.3],
+                          [0.1, 0.2, 0.1, 0.6],
+                      ]])
+
+    actions = np.array([[1, 2, 0, 1], [0, 3, 3, 0]])
+
+    mask = np.array([[1, 1, 0, 0], [1, 1, 1, 0]])
+
+    probab_ratios = ppo.compute_probab_ratios(p_old, p_new, actions, mask)
+
+    self.assertAllClose(
+        np.array([
+            [0.1 / 0.2, 0.1 / 0.4, 0.0, 0.0],
+            [0.1 / 0.3, 0.6 / 0.4, 0.3 / 0.1, 0.0],
+        ]), probab_ratios)
+
+  def test_clipped_probab_ratios(self):
+    probab_ratios = np.array([
+        [1.5, 1.0, 0.5, 0.7],
+        [2.5, 2.0, 0.1, 1.0],
+    ])
+
+    mask = np.array([[1, 1, 0, 0], [1, 1, 1, 0]])
+
+    clipped_probab_ratios = ppo.clipped_probab_ratios(probab_ratios, mask, 0.1)
+
+    self.assertAllClose(
+        np.array([
+            [1.1, 1.0, 0, 0],
+            [1.1, 1.1, 0.9, 0],
+        ]), clipped_probab_ratios)
+
+  def test_clipped_objective(self):
+    probab_ratios = np.array([
+        [1.5, 2.0, 0.5, 0.7],
+        [2.5, 2.0, 0.1, 1.0],
+    ])
+
+    advantages = np.array([
+        [0.1, 0.1, 0.5, 0.7],
+        [2.0, 2.0, 2.0, 2.0],
+    ])
+
+    mask = np.array([[1, 1, 0, 0], [1, 1, 1, 0]])
+
+    epsilon = 0.1
+
+    unused_clipped_probab_ratios = np.array([
+        [1.1, 1.1, 0.9, 0.9],
+        [1.1, 1.1, 0.9, 1.0],
+    ])
+
+    minimums = np.array([
+        [1.1, 1.1, 0.5, 0.7],
+        [1.1, 1.1, 0.1, 1.0],
+    ])
+
+    # advantages * minimums * mask
+    objective = np.array([
+        [0.11, 0.11, 0.0, 0.0],
+        [2.2, 2.2, 0.2, 0.0],
+    ])
+
+    # Assert that we computed things correctly in this test.
+    self.assertAllClose(advantages * mask * minimums, objective)
+
+    self.assertAllClose(
+        objective,
+        ppo.clipped_objective(probab_ratios, advantages, mask, epsilon))
+
+  def test_ppo_loss(self):
+    self.rng_key, key1, key2, key3 = jax_random.split(self.rng_key, num=4)
+
+    B, T, A, OBS = 2, 10, 2, (28, 28, 3)  # pylint: disable=invalid-name
+    batch_observation_shape = (-1, -1) + OBS
+
+    old_policy_params, _ = ppo.policy_net(key1, batch_observation_shape, A,
+                                          [stax.Flatten(2)])
+
+    new_policy_params, policy_apply = ppo.policy_net(key2,
+                                                     batch_observation_shape, A,
+                                                     [stax.Flatten(2)])
+
+    value_params, value_apply = ppo.value_net(key3, batch_observation_shape, A,
+                                              [stax.Flatten(2)])
+
+    # Generate a batch of observations.
+
+    observations = np.random.uniform(size=(B, T + 1) + OBS)
+    actions = np.random.randint(0, A, size=(B, T))
+    rewards = np.random.uniform(0, 1, size=(B, T))
+    mask = np.ones_like(rewards)
+
+    # Just test that this computes at all.
+    _ = ppo.ppo_loss(policy_apply, new_policy_params, old_policy_params,
+                     value_apply, value_params, observations, actions, rewards,
+                     mask)
 
 
 if __name__ == "__main__":

From 3c4b36fe234a392ad702de43d1c5b311770adbeb Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Wed, 17 Apr 2019 16:31:10 -0700
Subject: [PATCH 1910/2720] Add missing documentation to functions in
 glow_ops.py

PiperOrigin-RevId: 244090444
---
 tensor2tensor/models/research/glow_ops.py     | 73 ++++++++++++++++---
 tensor2tensor/models/video/next_frame_glow.py |  1 +
 2 files changed, 63 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index fdb272d2b..318fb05a2 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -35,11 +35,13 @@ def linear_interpolate(tensor1, tensor2, coeffs):
   """Linearly interpolate between two tensors at coeff.
 
   Args:
-    tensor1: 3-D Tensor, NHWC
-    tensor2: 3-D Tensor, NHWC
+    tensor1: 4-D Tensor, shape=(NHWC)
+    tensor2: 4-D Tensor, shape=(NHWC)
     coeffs: list of floats.
   Returns:
-    interp_latents: list of interpolated 4-D Tensors, shape=(1HWC)
+    interp_latents: 5-D Tensor, with interp_latents[i] representing
+                    interpolations at coeffs[i].
+                    shape=(len(coeffs), NHWC)
   """
   interp_tensors = []
   for coeff in coeffs:
@@ -167,7 +169,8 @@ def check_cond_latents(cond_latents, hparams):
 def get_variable_ddi(name, shape, initial_value, dtype=tf.float32, init=False,
                      trainable=True):
   """Wrapper for data-dependent initialization."""
-  # If init is a tensor bool, w is returned dynamically.
+  # If init is a tf bool: w is assigned dynamically at runtime.
+  # If init is a python bool: then w is determined during graph construction.
   w = tf.get_variable(name, shape, dtype, None, trainable=trainable)
   if isinstance(init, bool):
     if init:
@@ -179,7 +182,9 @@ def get_variable_ddi(name, shape, initial_value, dtype=tf.float32, init=False,
 
 @add_arg_scope
 def get_dropout(x, rate=0.0, init=True):
-  """Zero dropout during init or prediction time.
+  """Dropout x with dropout_rate = rate.
+
+  Apply zero dropout during init or prediction time.
 
   Args:
     x: 4-D Tensor, shape=(NHWC).
@@ -667,13 +672,13 @@ def additive_coupling(name, x, mid_channels=512, reverse=False,
 
   Args:
     name: variable scope.
-    x: 4-D Tensor.
+    x: 4-D Tensor, shape=(NHWC).
     mid_channels: number of channels in the coupling layer.
     reverse: Forward or reverse operation.
     activation: "relu" or "gatu"
     dropout: default, 0.0
   Returns:
-    output:
+    output: 4-D Tensor, shape=(NHWC)
     objective: 0.0
   """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
@@ -704,7 +709,7 @@ def affine_coupling(name, x, mid_channels=512, activation="relu",
     reverse: Forward or reverse operation.
     dropout: default, 0.0
   Returns:
-    output: input s
+    output: x shifted and scaled by an affine transformation.
     objective: log-determinant of the jacobian
   """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
@@ -1098,7 +1103,7 @@ def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
     name: variable scope.
     x: 4-D Tensor, shape (NHWC).
     reverse: Forward or reverse pass.
-    eps: If eps is provided, x2 is set to be
+    eps: If eps is provided, x2 is set to be mu(x1) + eps * sigma(x1).
     eps_std: Sample x2 with the provided eps_std.
     cond_latents: optionally condition x2 on cond_latents.
     hparams: next_frame_glow hparams.
@@ -1109,6 +1114,16 @@ def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
     temperature: Temperature with which to sample from the gaussian.
 
   Returns:
+    If reverse:
+      x: 4-D Tensor, concats input and x2 across channels.
+      x2: 4-D Tensor, a sample from N(mu(x1), sigma(x1))
+    Else:
+      x1: 4-D Tensor, Output of the split operation.
+      logpb: log-probability of x2 belonging to mu(x1), sigma(x1)
+      eps: 4-D Tensor, (x2 - mu(x1)) / sigma(x1)
+      x2: 4-D Tensor, Latent representation at the current level.
+    state: Current LSTM state.
+           4-D Tensor, only if hparams.latent_dist_encoder is set to conv_lstm.
   Raises:
     ValueError: If latent is provided and shape is not equal to NHW(C/2)
                 where (NHWC) is the size of x.
@@ -1179,7 +1194,17 @@ def revnet_step(name, x, hparams, reverse=True):
 
 
 def revnet(name, x, hparams, reverse=True):
-  """'hparams.depth' steps of generative flow."""
+  """'hparams.depth' steps of generative flow.
+
+  Args:
+    name: variable scope for the revnet block.
+    x: 4-D Tensor, shape=(NHWC).
+    hparams: tf.contrib.training.HParams.
+    reverse: bool, forward or backward pass.
+  Returns:
+    x: 4-D Tensor, shape=(NHWC).
+    objective: float.
+  """
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
     steps = np.arange(hparams.depth)
     if reverse:
@@ -1276,7 +1301,33 @@ def uniform_binning_correction(x, n_bits=8):
 def encoder_decoder(name, x, hparams, eps=None, reverse=False,
                     cond_latents=None, condition=False, states=None,
                     temperature=1.0):
-  """Glow encoder-decoder. n_levels of (Squeeze + Flow + Split.) operations."""
+  """Glow encoder-decoder. n_levels of (Squeeze + Flow + Split.) operations.
+
+  Args:
+    name: variable scope.
+    x: 4-D Tensor, shape=(NHWC).
+    hparams: tf.contrib.training.HParams.
+    eps: Stores (glow(x) - mu) / sigma during the forward pass.
+         Used only to test if the network is reversible.
+    reverse: Forward or reverse pass.
+    cond_latents: list of lists of tensors.
+                  outer length equals hparams.num_cond_latents
+                  innter length equals hparams.num_levels - 1.
+    condition: If set to True, condition the encoder/decoder on cond_latents.
+    states: LSTM states, used only if hparams.latent_dist_encoder is set
+            to "conv_lstm.
+    temperature: Temperature set during sampling.
+  Returns:
+    x: If reverse, decoded image, else the encoded glow latent representation.
+    objective: log-likelihood.
+    eps: list of tensors, shape=(num_levels-1).
+         Stores (glow(x) - mu_level(x)) / sigma_level(x)) for each level.
+    all_latents: list of tensors, shape=(num_levels-1).
+                 Latent representatios for each level.
+    new_states: list of tensors, shape=(num_levels-1).
+                useful only if hparams.latent_dist_encoder="conv_lstm", returns
+                the current state of each level.
+  """
   # TODO(mechcoder) Change return_type to a dict to be backward compatible.
   with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
 
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index cc29c6b02..e82ed2876 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -50,6 +50,7 @@ def next_frame_glow_hparams():
   # This function is used to model the prior over z_{t}. Can be,
   # Pointwise -> point-wise multiplication of z_{t-1}.
   # conv_net -> one-layer convolution over z_{t-1} .. z_{t - num_cond_latents}
+  # conv3d_net or conv_lstm
   hparams.add_hparam("latent_dist_encoder", "conv_net")
   # Number of latents used in the encoder above.
   hparams.add_hparam("num_cond_latents", 1)

From 157b4d392b3762bb49dfc5b3fedc61a08381a9e6 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 17 Apr 2019 17:24:14 -0700
Subject: [PATCH 1911/2720] - Add a `Div` layer to stax_base. - Use `Div` as
 one of the layers in the Pong model. - PPO on Pong atleast makes progress
 with this.   - Although quite slow: Since the trajectory size changes a lot
 between     epochs, so might need to adjust bucket size for lesser
 re-jitting.

TODOs:
- Use a single network for policy and value.
- In collect_trajectories, give the whole history of time-steps.
PiperOrigin-RevId: 244098994
---
 tensor2tensor/trax/rlax/ppo.py       | 26 ++++++++++++++------------
 tensor2tensor/trax/rlax/ppo_main.py  |  7 +++++--
 tensor2tensor/trax/stax/slax_test.py | 13 +++++++++++++
 tensor2tensor/trax/stax/stax_base.py |  8 ++++++++
 4 files changed, 40 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 6a4cafc56..572c69bb0 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -135,13 +135,13 @@ def collect_trajectories(env,
     observation = env.reset()
     observations.append(observation)
     while not done:
-      # Add a batch dimension and time dimension.
+      # Add a batch dimension and time dimension, so shape is (1, 1) + OBS
       observation = observation[np.newaxis, np.newaxis, :]
 
-      # Run the policy, pick an action.
+      # Run the policy, to pick an action, shape is (1, 1, A)
       predictions = policy_net_apply(policy_net_params, observation)
 
-      # Squeeze the added dimension.
+      # Squeeze the added dimension, shape is (A,)
       predictions = np.squeeze(predictions)
 
       # Policy can be run in one of the following ways:
@@ -160,8 +160,7 @@ def collect_trajectories(env,
           # Return the best action.
           action = np.argmax(predictions)
       elif policy == "categorical-sampling":
-        # import pdb; pdb.set_trace()
-        action = int(onp.argwhere(onp.random.multinomial(1, predictions) == 1))
+        action = onp.argwhere(onp.random.multinomial(1, predictions) == 1)
       else:
         raise ValueError("Unknown policy: %s" % policy)
 
@@ -203,7 +202,7 @@ def get_padding_value(dtype):
 
 
 # TODO(afrozm): Use np.pad instead and make jittable?
-def pad_trajectories(trajectories, boundary=10):
+def pad_trajectories(trajectories, boundary=20):
   """Pad trajectories to a bucket length that is a multiple of boundary.
 
   Args:
@@ -582,6 +581,7 @@ def training_loop(
     batch_size=BATCH_TRAJECTORIES,
     num_optimizer_steps=NUM_OPTIMIZER_STEPS,
     print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
+    boundary=20,
     random_seed=None):
   """Runs the training loop for PPO, with fixed policy and value nets."""
   jax_rng_key = trax.get_random_number_generator_and_set_seed(random_seed)
@@ -631,19 +631,21 @@ def training_loop(
 
     avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
     max_reward = max(np.sum(traj[2]) for traj in trajs)
+    min_reward = min(np.sum(traj[2]) for traj in trajs)
     average_rewards.append(avg_reward)
 
-    logging.vlog(1, "Rewards average=[%0.2f], max=[%0.2f]", avg_reward,
-                 max_reward)
+    logging.vlog(1, "Rewards average=[%0.2f], max=[%0.2f], min=[%0.2f]",
+                 avg_reward, max_reward, min_reward)
     logging.vlog(1, "Collecting trajectories took %0.2f msec.", get_time(t))
-    logging.vlog(1, "Trajectory Length average=[%0.2f], max=[%0.2f]",
+    logging.vlog(1,
+                 "Trajectory Length average=[%0.2f], max=[%0.2f], min=[%0.2f]",
                  float(sum(len(traj[0]) for traj in trajs)) / len(trajs),
-                 max(len(traj[0]) for traj in trajs))
+                 max(len(traj[0]) for traj in trajs),
+                 min(len(traj[0]) for traj in trajs))
 
     t = time.time()
     (_, reward_mask, padded_observations, padded_actions,
-     padded_rewards) = pad_trajectories(
-         trajs, boundary=20)
+     padded_rewards) = pad_trajectories(trajs, boundary=boundary)
 
     logging.vlog(1, "Padding trajectories took %0.2f msec.", get_time(t))
     logging.vlog(1, "Padded Observations' shape [%s]",
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index e4d58cf89..aa343f42f 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -24,8 +24,8 @@
 from absl import app
 from absl import flags
 from absl import logging
+from tensor2tensor.trax import stax
 from tensor2tensor.trax.rlax import ppo
-from tensor2tensor.trax.stax import stax_base as stax
 
 FLAGS = flags.FLAGS
 
@@ -35,6 +35,8 @@
 flags.DEFINE_integer("random_seed", 0, "Random seed.")
 flags.DEFINE_integer("log_level", logging.INFO, "Log level.")
 flags.DEFINE_integer("batch_size", 32, "Batch of trajectories needed.")
+flags.DEFINE_integer("boundary", 20,
+                     "We pad trajectories at integer multiples of this number.")
 
 
 def common_stax_layers():
@@ -47,7 +49,7 @@ def main(argv):
   bottom_layers = common_stax_layers()
 
   if FLAGS.env_name == "Pong-v0":
-    bottom_layers = [stax.Flatten(2)] + bottom_layers
+    bottom_layers = [stax.Div(255.0), stax.Flatten(2)] + bottom_layers
 
   ppo.training_loop(
       env_name=FLAGS.env_name,
@@ -57,6 +59,7 @@ def main(argv):
       value_net_fun=functools.partial(
           ppo.value_net, bottom_layers=bottom_layers),
       batch_size=FLAGS.batch_size,
+      boundary=FLAGS.boundary,
       random_seed=FLAGS.random_seed)
 
 
diff --git a/tensor2tensor/trax/stax/slax_test.py b/tensor2tensor/trax/stax/slax_test.py
index 25ef39bf7..fa258e3ba 100644
--- a/tensor2tensor/trax/stax/slax_test.py
+++ b/tensor2tensor/trax/stax/slax_test.py
@@ -117,6 +117,19 @@ def test_flatten_n(self):
     with self.assertRaises(ValueError):
       check_staxlayer(self, stax.Flatten(6), input_shape)
 
+  def test_div(self):
+    init_fun, apply_fun = stax.Div(2)
+    input_np = onp.array([[1, 2, 3], [4, 5, 6]], dtype=onp.float32)
+    input_shape = input_np.shape
+    _, _ = init_fun(None, input_shape)
+    output_np = apply_fun(None, input_np)
+    # absltest doesn't have ndarray equalities.
+    expected_output_np = input_np / 2.0
+    self.assertAlmostEqual(
+        0.0,
+        onp.sum((output_np - expected_output_np) ** 2),
+        delta=1e-6)
+
   # Lambdas replace the staxlayer input stream with a placeholder that
   # _should_ break any use of unbound variables in the input stream.
   def testLambda_forbidden_access(self):
diff --git a/tensor2tensor/trax/stax/stax_base.py b/tensor2tensor/trax/stax/stax_base.py
index aa1ceb2e4..f690c982a 100644
--- a/tensor2tensor/trax/stax/stax_base.py
+++ b/tensor2tensor/trax/stax/stax_base.py
@@ -281,6 +281,14 @@ def apply_fun(params, inputs, **kwargs):  # pylint: disable=missing-docstring
   return init_fun, apply_fun
 
 
+def Div(divisor):
+  def init_fun(_, input_shape):
+    return input_shape, ()
+  def apply_fun(params, inputs, **kwargs):
+    return inputs / divisor
+  return init_fun, apply_fun
+
+
 # Composing layers via combinators
 
 

From f87c54291013a824bebefdd46ef4233e568496cd Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 18 Apr 2019 09:55:13 -0700
Subject: [PATCH 1912/2720] While collecting trajectories have access to all
 history of observations.

This does slow down collect, because now we are calling policy over different shapes each time, i.e. t = 1 to T  batch of observations is shaped (b, t) + OBS and t increases every iteration.

But this sems to be unavoidable, unless we know we just want the last time-step's predictions.

PiperOrigin-RevId: 244204123
---
 tensor2tensor/trax/rlax/ppo.py | 32 ++++++++++++++++++++------------
 1 file changed, 20 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 572c69bb0..6fc2cf924 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -127,22 +127,24 @@ def collect_trajectories(env,
   trajectories = []
 
   for _ in range(num_trajectories):
-    observations = []
     rewards = []
     actions = []
     done = False
 
     observation = env.reset()
-    observations.append(observation)
-    while not done:
-      # Add a batch dimension and time dimension, so shape is (1, 1) + OBS
-      observation = observation[np.newaxis, np.newaxis, :]
 
-      # Run the policy, to pick an action, shape is (1, 1, A)
-      predictions = policy_net_apply(policy_net_params, observation)
+    # This is currently shaped (1, 1) + OBS, but new observations will keep
+    # getting added to it, making it eventually (1, T+1) + OBS
+    observation_history = observation[np.newaxis, np.newaxis, :]
+
+    while not done:
+      # Run the policy, to pick an action, shape is (1, t, A) because
+      # observation_history is shaped (1, t) + OBS
+      predictions = policy_net_apply(policy_net_params, observation_history)
 
-      # Squeeze the added dimension, shape is (A,)
-      predictions = np.squeeze(predictions)
+      # We need the predictions for the last time-step, so squeeze the batch
+      # dimension and take the last time-step.
+      predictions = np.squeeze(predictions, axis=0)[-1]
 
       # Policy can be run in one of the following ways:
       #  - Greedy
@@ -169,14 +171,20 @@ def collect_trajectories(env,
 
       observation, reward, done, _ = env.step(action)
 
-      observations.append(observation)
+      # observation is of shape OBS, so add extra dims and concatenate on the
+      # time dimension.
+      observation_history = np.concatenate(
+          [observation_history, observation[np.newaxis, np.newaxis, :]], axis=1)
+
       rewards.append(reward)
       actions.append(action)
 
     # This means we are done
     assert done
+    # observation_history is (1, T+1) + OBS, lets squeeze out the batch dim.
+    observation_history = np.squeeze(observation_history, axis=0)
     trajectories.append(
-        (np.stack(observations), np.stack(actions), np.stack(rewards)))
+        (observation_history, np.stack(actions), np.stack(rewards)))
 
   return trajectories
 
@@ -480,7 +488,7 @@ def ppo_loss(policy_net_apply,
 
   # (B, T)
   td_deltas = deltas(
-      np.squeeze(predicted_values, axis=2),
+      np.squeeze(predicted_values, axis=2),  # (B, T)
       padded_rewards,
       reward_mask,
       gamma=gamma)

From 13b48761a180d2b040096fe09e361054fe473907 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 18 Apr 2019 11:22:18 -0700
Subject: [PATCH 1913/2720] Trax debugging: pass mode to models, always jit
 eval and do multi-device evals; correct input pipeline for that.

PiperOrigin-RevId: 244224385
---
 .../trax/configs/transformer_lm1b_8gb.gin     |  6 +-
 .../trax/configs/transformer_lm1b_tpu.gin     |  7 +-
 tensor2tensor/trax/inputs.py                  | 69 ++++++++++++++----
 tensor2tensor/trax/inputs_test.py             | 72 +++++++++++++++++++
 tensor2tensor/trax/models/mlp.py              |  4 +-
 tensor2tensor/trax/models/resnet.py           |  4 +-
 tensor2tensor/trax/models/transformer.py      |  2 +-
 tensor2tensor/trax/trax.py                    | 36 ++++------
 tensor2tensor/trax/trax_test.py               |  4 +-
 9 files changed, 157 insertions(+), 47 deletions(-)
 create mode 100644 tensor2tensor/trax/inputs_test.py

diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index 5b99bd2f7..cf47eaee0 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -7,6 +7,7 @@ import tensor2tensor.trax.trax
 # ==============================================================================
 batch_fun.batch_size = 128
 batch_fun.eval_batch_size = 128
+batch_fun.max_eval_length = 2048
 
 # Parameters for inputs:
 # ==============================================================================
@@ -27,7 +28,8 @@ MultifactorSchedule.warmup_steps = 8000
 # Parameters for preprocess_fun:
 # ==============================================================================
 shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
-lm1b_preprocess.max_target_length = 511
+lm1b_preprocess.max_target_length = 512
+lm1b_preprocess.max_eval_target_length = 2048
 
 # Parameters for train:
 # ==============================================================================
@@ -43,7 +45,7 @@ train.train_steps = 500000
 TransformerLM.dropout = 0.1
 TransformerLM.feature_depth = 512
 TransformerLM.feedforward_depth = 2048
-TransformerLM.max_len = 512
+TransformerLM.max_len = 2048
 TransformerLM.mode = 'train'
 TransformerLM.num_heads = 8
 TransformerLM.num_layers = 6
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_tpu.gin b/tensor2tensor/trax/configs/transformer_lm1b_tpu.gin
index 7ad8ddfa6..fd264af49 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_tpu.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_tpu.gin
@@ -7,6 +7,7 @@ import tensor2tensor.trax.trax
 # ==============================================================================
 batch_fun.batch_size = 1024
 batch_fun.eval_batch_size = 128
+batch_fun.max_eval_length = 2048
 
 # Parameters for inputs:
 # ==============================================================================
@@ -27,13 +28,13 @@ MultifactorSchedule.warmup_steps = 8000
 # Parameters for preprocess_fun:
 # ==============================================================================
 shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
-lm1b_preprocess.max_target_length = 511
+lm1b_preprocess.max_target_length = 512
+lm1b_preprocess.max_eval_target_length = 2048
 
 # Parameters for train:
 # ==============================================================================
 train.eval_frequency = 1000
 train.eval_steps = 10
-train.jit_eval = False  # Eval lengths vary a lot, compiling each time is slow.
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
 train.run_debug_step = False
@@ -44,7 +45,7 @@ train.train_steps = 500000
 TransformerLM.dropout = 0.1
 TransformerLM.feature_depth = 512
 TransformerLM.feedforward_depth = 2048
-TransformerLM.max_len = 512
+TransformerLM.max_len = 2048
 TransformerLM.mode = 'train'
 TransformerLM.num_heads = 8
 TransformerLM.num_layers = 6
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index afa528797..8b33738a1 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -53,11 +53,12 @@
 _MAX_SKIP_EXAMPLES = 1e5
 
 
-@gin.configurable()
-def inputs(dataset_name, data_dir=None, input_name=None):
+@gin.configurable(blacklist=["num_devices"])
+def inputs(num_devices, dataset_name, data_dir=None, input_name=None):
   """Make Inputs for built-in datasets.
 
   Args:
+    num_devices: how many devices to build the inputs for.
     dataset_name: a TFDS or T2T dataset name. If it's a T2T dataset name, prefix
       with "t2t_".
     data_dir: data directory.
@@ -71,7 +72,7 @@ def inputs(dataset_name, data_dir=None, input_name=None):
 
   (train_batches, train_eval_batches, eval_batches,
    input_name, input_shape) = _train_and_eval_batches(
-       dataset_name, data_dir, input_name)
+       dataset_name, data_dir, input_name, num_devices)
 
   def train_input_fun():
     return dataset_to_stream(train_batches, input_name)
@@ -88,13 +89,15 @@ def eval_input_fun():
                 input_shape=input_shape)
 
 
-@gin.configurable()
+@gin.configurable(blacklist=["num_devices"])
 def random_inputs(
+    num_devices,
     input_shape=gin.REQUIRED, input_dtype=np.int32, input_range=(0, 255),
     output_shape=gin.REQUIRED, output_dtype=np.int32, output_range=(0, 9)):
   """Make random Inputs for debugging.
 
   Args:
+    num_devices: how many devices to build the inputs for.
     input_shape: the shape of inputs (including batch dimension).
     input_dtype: the type of the inputs (int32 by default).
     input_range: the range of inputs (defaults to (0, 255)).
@@ -105,6 +108,15 @@ def random_inputs(
   Returns:
     trax.inputs.Inputs
   """
+  if input_shape[0] % num_devices != 0:
+    tf.logging.fatal(
+        "num_devices[%d] should divide the first dimension of input_shape[%s]",
+        num_devices, input_shape)
+  if output_shape[0] % num_devices != 0:
+    tf.logging.fatal(
+        "num_devices[%d] should divide the first dimension of output_shape[%s]",
+        num_devices, output_shape)
+
   def random_minibatches():
     """Generate a stream of random mini-batches."""
     if input_dtype in [np.float16, np.float32, np.float64]:
@@ -224,15 +236,19 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
   return train_dataset, eval_dataset, info, supervised_keys
 
 
-@gin.configurable(blacklist=["dataset", "training", "shapes", "target_names"])
-def batch_fun(dataset, training, shapes, target_names,
+@gin.configurable(blacklist=["dataset", "training", "shapes",
+                             "target_names", "num_devices"])
+def batch_fun(dataset, training, shapes, target_names, num_devices,
               batch_size=32, eval_batch_size=32,
               bucket_length=32, buckets=None,
-              batch_shuffle_size=512):
+              batch_shuffle_size=128, max_eval_length=None):
   """Batching function."""
   del target_names
   # If bucketing is not specified, check if target shapes are variable.
   cur_batch_size = batch_size if training else eval_batch_size
+  # Make cur_batch_size divisible by num_devices.
+  cur_batch_size = max(cur_batch_size // num_devices, 1) * num_devices
+  # Create heuristic buckets is none are specified.
   if buckets is None:
     variable_target_shapes = False
     target_shape = shapes[1]
@@ -246,10 +262,20 @@ def batch_fun(dataset, training, shapes, target_names,
                            bucket_length, bucket_length * 2,
                            bucket_length * 4, bucket_length * 8,
                            bucket_length * 16]
+      # We will pad to boundaries which pads to bucket_boundary - 1: add 1 here.
+      bucket_boundaries = [b + 1 for b in bucket_boundaries]
+      if not training:
+        max_eval_length = max_eval_length or bucket_length * 32
+        bucket_boundaries[-1] = max_eval_length
       bucket_batch_sizes = [cur_batch_size * 4, cur_batch_size * 2,
                             cur_batch_size, cur_batch_size // 2,
                             cur_batch_size // 4, cur_batch_size // 8,
-                            max(1, cur_batch_size // 16), 1]
+                            cur_batch_size // 16, 1]
+      if not training:
+        bucket_batch_sizes[-2] = cur_batch_size // max_eval_length
+      # Make batch sizes divisible by num_devices.
+      bucket_batch_sizes = [max(b // num_devices, 1) * num_devices
+                            for b in bucket_batch_sizes]
       buckets = (bucket_boundaries, bucket_batch_sizes)
 
   if buckets:
@@ -259,7 +285,7 @@ def example_length(_, target):
     boundaries, batch_sizes = buckets
     dataset = dataset.apply(tf.data.experimental.bucket_by_sequence_length(
         example_length, boundaries, batch_sizes,
-        pad_to_bucket_boundary=training))
+        pad_to_bucket_boundary=True))
   else:
     dataset = dataset.padded_batch(cur_batch_size, shapes)
   if training:
@@ -285,13 +311,22 @@ def no_preprocess(dataset, training):
 
 
 @gin.configurable(blacklist=["dataset", "training"])
-def lm1b_preprocess(dataset, training, max_target_length=-1):
+def lm1b_preprocess(dataset, training,
+                    max_target_length=-1, max_eval_target_length=-1):
+  """Preprocessing for LM1B: filter out targets exceeding maximum length."""
 
   def target_right_length(_, target):
     return tf.less(tf.shape(target)[0], max_target_length + 1)
 
+  def eval_target_right_length(_, target):
+    return tf.less(tf.shape(target)[0], max_eval_target_length + 1)
+
   if max_target_length > 0 and training:
     dataset = dataset.filter(target_right_length)
+
+  if max_eval_target_length > 0 and not training:
+    dataset = dataset.filter(eval_target_right_length)
+
   return dataset
 
 
@@ -300,6 +335,7 @@ def shuffle_and_batch_data(dataset,
                            target_names,
                            features_info,
                            training,
+                           num_devices,
                            shuffle_buffer_size=1024,
                            preprocess_fun=no_preprocess):
   """Shuffle and batch the given dataset."""
@@ -322,21 +358,24 @@ def append_targets(example):
   shapes = {k: features_info[k].shape for k in features_info}
   shapes = (shapes, shapes[target_names[0]])
   dataset = dataset.shuffle(shuffle_buffer_size)
-  dataset = batch_fun(dataset, training, shapes, target_names)
+  dataset = batch_fun(dataset, training, shapes, target_names, num_devices)
   return dataset.prefetch(2)
 
 
-def _train_and_eval_batches(dataset, data_dir, input_name):
+def _train_and_eval_batches(dataset, data_dir, input_name, num_devices):
   """Return train and eval batches with input name and shape."""
   (train_data, eval_data, features_info, keys) = train_and_eval_dataset(
       dataset, data_dir)
   input_names, target_names = keys[0], keys[1]
   train_batches = shuffle_and_batch_data(
-      train_data, target_names, features_info, training=True)
+      train_data, target_names, features_info, training=True,
+      num_devices=num_devices)
   train_eval_batches = shuffle_and_batch_data(  # Data for eval-on-train.
-      train_data, target_names, features_info, training=False)
+      train_data, target_names, features_info, training=False,
+      num_devices=num_devices)
   eval_batches = shuffle_and_batch_data(
-      eval_data, target_names, features_info, training=False)
+      eval_data, target_names, features_info, training=False,
+      num_devices=num_devices)
   input_name = input_name or input_names[0]
   input_shape = features_info[input_name].shape
   return (train_batches, train_eval_batches, eval_batches,
diff --git a/tensor2tensor/trax/inputs_test.py b/tensor2tensor/trax/inputs_test.py
new file mode 100644
index 000000000..1d9f353d2
--- /dev/null
+++ b/tensor2tensor/trax/inputs_test.py
@@ -0,0 +1,72 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.inputs."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gin
+import numpy as np
+from tensor2tensor.trax import inputs
+import tensorflow as tf
+import tensorflow_datasets as tfds
+
+
+def test_dataset_ints(lengths):
+  """Create a test dataset of int64 tensors of shape [length]."""
+  def generator():
+    """Sample generator of sequences of shape [length] of type int64."""
+    for length in lengths:
+      x = np.zeros([length], dtype=np.int64)
+      yield (x, x)  # Inputs and targets are the same here.
+  types = (tf.int64, tf.int64)
+  shapes = (tf.TensorShape([None]), tf.TensorShape([None]))
+  return tf.data.Dataset.from_generator(
+      generator, output_types=types, output_shapes=shapes)
+
+
+class InputsTest(tf.test.TestCase):
+
+  def setUp(self):
+    gin.clear_config()
+
+  def test_batch_fun(self):
+    dataset = test_dataset_ints([32])
+    dataset = dataset.repeat(10)
+    batches = inputs.batch_fun(
+        dataset, True, ([None], [None]), [], 1, batch_size=10)
+    count = 0
+    for example in tfds.as_numpy(batches):
+      count += 1
+      self.assertEqual(example[0].shape[0], 10)  # Batch size = 10.
+    self.assertEqual(count, 1)  # Just one batch here.
+
+  def test_batch_fun_num_devices(self):
+    dataset = test_dataset_ints([32])
+    dataset = dataset.repeat(9)
+    batches = inputs.batch_fun(
+        dataset, True, ([None], [None]), [], 9, batch_size=10)
+    count = 0
+    for example in tfds.as_numpy(batches):
+      count += 1
+      # Batch size adjusted to be divisible by num_devices.
+      self.assertEqual(example[0].shape[0], 9)
+    self.assertEqual(count, 1)  # Just one batch here.
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/trax/models/mlp.py b/tensor2tensor/trax/models/mlp.py
index c1296120a..7618c6612 100644
--- a/tensor2tensor/trax/models/mlp.py
+++ b/tensor2tensor/trax/models/mlp.py
@@ -25,7 +25,9 @@
 def MLP(num_hidden_layers=2,
         hidden_size=512,
         activation_fn=stax.Relu,
-        num_output_classes=10):
+        num_output_classes=10,
+        mode="train"):
+  del mode
   layers = [stax.Flatten()]
   layers += [stax.Dense(hidden_size), activation_fn] * num_hidden_layers
   layers += [stax.Dense(num_output_classes), stax.LogSoftmax]
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index dc0db56d9..45a9e4d98 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -60,16 +60,18 @@ def MakeMain(input_shape):
       stax.FanInSum, stax.Relu)
 
 
-def Resnet50(hidden_size=64, num_output_classes=1001):
+def Resnet50(hidden_size=64, num_output_classes=1001, mode='train'):
   """ResNet.
 
   Args:
     hidden_size: the size of the first hidden layer (multiplied later).
     num_output_classes: how many classes to distinguish.
+    mode: whether we are training or evaluating or doing inference.
 
   Returns:
     The ResNet model with the given layer and output sizes.
   """
+  del mode
   return stax.serial(
       stax.Conv(hidden_size, (7, 7), (2, 2), 'SAME'),
       stax.BatchNorm(), stax.Relu,
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index afe00d609..4eb4b45e9 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -99,7 +99,7 @@ def TransformerLM(vocab_size,  # pylint: disable=invalid-name
                   feedforward_depth=2048,
                   num_heads=8,
                   dropout=0.1,
-                  max_len=512):
+                  max_len=2048):
   """Transformer language model (only uses the decoder part of Transformer).
 
   Args:
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 595bac5bc..b59fb07af 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -244,14 +244,10 @@ def epochs(steps=None, epoch_steps=1):
       break
 
 
-def _jit_predict_fun(model_predict, num_devices, jit_eval):
+def _jit_predict_fun(model_predict, num_devices):
   """Use jit on model_predict if required."""
   def predict(params, batch, rng=None):
     """Predict function jited and parallelized as requested."""
-    # If not jit'ing, just run the function.
-    if not jit_eval:
-      return model_predict(params, batch, rng=rng)
-
     # On one device, jit and run.
     if num_devices == 1:
       return backend.jit(model_predict)(params, batch, rng=rng)
@@ -261,7 +257,7 @@ def predict(params, batch, rng=None):
     def mapped_predict(params, batch, rng):
       return model_predict(params, batch, rng=rng)
     pred = mapped_predict(
-        jax.replicate(params),
+        params,
         reshape_by_device(batch, num_devices),
         jax.replicate(rng))
     batch_size = batch.shape[0]
@@ -336,7 +332,6 @@ def train(output_dir,
           eval_frequency=100,
           num_devices=None,
           random_seed=None,
-          jit_eval=True,
           run_debug_step=False):
   """Train the model on the inputs.
 
@@ -357,7 +352,6 @@ def train(output_dir,
       steps). If None or 0, eval disabled.
     num_devices: how many devices to use (if None, default, use all available)
     random_seed: the random seed to use; time/os dependent if None (default).
-    jit_eval: whether to compile the evaulation function (true by default).
     run_debug_step: bool, if True, will run the model and loss without @jit for
       one step.
 
@@ -371,14 +365,15 @@ def train(output_dir,
   train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
   eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
 
-  inputs = inputs()
+  inputs = inputs(num_devices)
 
   # Setup optimizer and model
   state = restore_state(output_dir)
   history = state.history
   lr_fun = lr_schedule(history)
   opt_init, _ = optimizer(lr_fun)
-  model_init, model_predict = model()
+  model_init, model_predict_train = model(mode="train")
+  _, model_predict_eval = model(mode="eval")
 
   # Setup state
   step = state.step or 0
@@ -391,14 +386,14 @@ def train(output_dir,
     opt_state = jax.replicate(opt_state)
 
   # jit model_predict and update so they're fast
-  jit_model_predict = _jit_predict_fun(model_predict, num_devices, jit_eval)
-  jit_update_fun = _jit_update_fun(model_predict, loss_fun, optimizer, lr_fun,
-                                   num_devices)
+  jit_model_predict_eval = _jit_predict_fun(model_predict_eval, num_devices)
+  jit_update_fun = _jit_update_fun(
+      model_predict_train, loss_fun, optimizer, lr_fun, num_devices)
 
   print()
   train_stream = inputs.train_stream()
   epoch_steps = [train_steps]  # Only training if eval_frequency is 0 or None.
-  if eval_frequency:
+  if eval_frequency and eval_steps > 0:
     epoch_steps = itertools.chain([1,  # first epoch only 1 step
                                    eval_frequency - 1],
                                   itertools.repeat(eval_frequency))
@@ -406,7 +401,7 @@ def train(output_dir,
 
   # Non-compiled debug step helps find problems in models easier.
   if run_debug_step:
-    debug_loss = loss_fun(params, next(train_stream), model_predict, rng)
+    debug_loss = loss_fun(params, next(train_stream), model_predict_train, rng)
     step_log(step, "Debug step loss %.8f" % debug_loss)
 
   for epoch, epoch_steps in epochs(train_steps, epoch_steps):
@@ -439,14 +434,11 @@ def train(output_dir,
                       epoch_steps / epoch_time, step=step)
 
     # Evaluate
-    if num_devices > 1:   # TODO(lukaszkaiser): remove branch when possible.
-      params = trax_opt.get_params(jax.unreplicate(opt_state))
-    else:
-      params = trax_opt.get_params(opt_state)
+    params = trax_opt.get_params(opt_state)
     evaluate_train_and_eval(
         step=step,
         inputs=inputs,
-        predict_fun=functools.partial(jit_model_predict, params),
+        predict_fun=functools.partial(jit_model_predict_eval, params),
         eval_steps=eval_steps,
         rng=rng,
         train_sw=train_sw,
@@ -465,8 +457,8 @@ def train(output_dir,
     old_lr_fun = lr_fun
     lr_fun = lr_schedule(history)
     if lr_fun != old_lr_fun:  # For performance, only jit if there is a change.
-      jit_update_fun = _jit_update_fun(model_predict, loss_fun, optimizer,
-                                       lr_fun, num_devices)
+      jit_update_fun = _jit_update_fun(
+          model_predict_train, loss_fun, optimizer, lr_fun, num_devices)
 
     # Flush summary writers
     train_sw.flush()
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index cb375d55a..6f0032059 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -69,7 +69,7 @@ def test_train_eval_predict(self):
       model = functools.partial(models.MLP,
                                 hidden_size=16,
                                 num_output_classes=num_classes)
-      inputs = lambda: test_inputs(num_classes)
+      inputs = lambda _: test_inputs(num_classes)
 
       # Train and evaluate
       state = trax.train(output_dir,
@@ -89,7 +89,7 @@ def test_train_eval_predict(self):
 
       # Predict with final params
       _, predict_fun = model()
-      inputs = inputs().train_stream()
+      inputs = inputs(1).train_stream()
       predict_fun(state.params, next(inputs)[0])
 
 

From 693c41b649db9658ddedea4dcbfc3d705e5c287e Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 18 Apr 2019 14:32:41 -0700
Subject: [PATCH 1914/2720] Bigger config for Transformer in trax.

PiperOrigin-RevId: 244259605
---
 .../trax/configs/transformer_lm1b_tpu_big.gin | 52 +++++++++++++++++++
 1 file changed, 52 insertions(+)
 create mode 100644 tensor2tensor/trax/configs/transformer_lm1b_tpu_big.gin

diff --git a/tensor2tensor/trax/configs/transformer_lm1b_tpu_big.gin b/tensor2tensor/trax/configs/transformer_lm1b_tpu_big.gin
new file mode 100644
index 000000000..64cc2a2ed
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_lm1b_tpu_big.gin
@@ -0,0 +1,52 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size = 256
+batch_fun.eval_batch_size = 128
+batch_fun.max_eval_length = 2048
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_languagemodel_lm1b32k'
+inputs.input_name = 'targets'
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 0.1
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
+lm1b_preprocess.max_target_length = 512
+lm1b_preprocess.max_eval_target_length = 2048
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 1000
+train.eval_steps = 10
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerLM
+train.run_debug_step = False
+train.train_steps = 500000
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.dropout = 0.1
+TransformerLM.feature_depth = 1024
+TransformerLM.feedforward_depth = 8192
+TransformerLM.max_len = 2048
+TransformerLM.mode = 'train'
+TransformerLM.num_heads = 8
+TransformerLM.num_layers = 8
+TransformerLM.vocab_size = 32000

From ed760b119bb2ed794e51dd13fb4adecbf7590c13 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 18 Apr 2019 17:11:58 -0700
Subject: [PATCH 1915/2720] Fix a numerical overflow issue in PPO.

This amounts to running a recurrence relation to compute rewards_to_go and
gae_advantages in reverse.

Added a test which fails on the current code.

PiperOrigin-RevId: 244288574
---
 tensor2tensor/trax/rlax/ppo.py      | 50 +++++++++++++++++------------
 tensor2tensor/trax/rlax/ppo_test.py | 24 ++++++++++++++
 2 files changed, 54 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 6fc2cf924..ec52f98c4 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -283,6 +283,7 @@ def pad_trajectories(trajectories, boundary=20):
       padded_observations), np.stack(padded_actions), np.stack(padded_rewards)
 
 
+# TODO(afrozm): JAX-ify this, this is too slow for pong.
 def rewards_to_go(rewards, mask, gamma=0.99):
   r"""Computes rewards to go.
 
@@ -299,32 +300,39 @@ def rewards_to_go(rewards, mask, gamma=0.99):
   Returns:
     rewards to go, np.ndarray of shape (B, T).
   """
-  B, T = rewards.shape  # pylint: disable=invalid-name
-
-  # [[1, g, g**2, ... g**T-1]]
-  # Not jittable, T should be a compile time constant.
-  # gammas = np.geomspace(1, g**T, T, endpoint=False).reshape(1, T)
+  B, T = rewards.shape  # pylint: disable=invalid-name,unused-variable
 
-  # Get a geometric progression of gamma, of length T.
-  gammas = [gamma**t for t in range(T)]
-  gammas = np.array(gammas).reshape((1, T))
+  masked_rewards = rewards * mask  # (B, T)
 
-  # Discounted rewards.
-  undiscounted_rewards = rewards * mask  # (B, T)
-  discounted_rewards = undiscounted_rewards * gammas  # (B, T)
+  # We use the following recurrence relation, derived from the equation above:
+  #
+  # r2g[t+1] = (r2g[t] - r[t]) / gamma
+  #
+  # This means we'll need to calculate r2g[0] first and then r2g[1] and so on ..
+  #
+  # **However** this leads to overflows for long sequences: r2g[t] - r[t] > 0
+  # and gamma < 1.0, so the division keeps increasing.
+  #
+  # So we just run the recurrence in reverse, i.e.
+  #
+  # r2g[t] = r[t] + (gamma*r2g[t+1])
+  #
+  # This is much better, but might have lost updates since the (small) rewards
+  # at earlier time-steps may get added to a (very?) large sum.
 
-  # Get rewards to go at first time-step.
-  r0 = np.sum(discounted_rewards, axis=1)  # (B,)
-  assert r0.shape == (B,)
+  # Compute r2g_{T-1} at the start and then compute backwards in time.
+  r2gs = [masked_rewards[:, -1]]
 
-  rs = [r0]
+  # Go from T-2 down to 0.
+  for t in reversed(range(T - 1)):
+    r2gs.append(masked_rewards[:, t] + (gamma * r2gs[-1]))
 
-  # Now compute the other advantages wrt the first one.
-  for t in range(1, T):
-    rs.append((rs[-1] - undiscounted_rewards[:, t - 1]) / gamma)
+  # The list should have length T.
+  assert T == len(r2gs)
 
-  # len(rs) is T and each element is (B,), this makes it (B, T)
-  return np.stack(rs, axis=1)
+  # First we stack them in the correct way to make it (B, T), but these are
+  # still from newest (T-1) to oldest (0), so then we flip it on time axis.
+  return np.flip(np.stack(r2gs, axis=1), axis=1)
 
 
 @functools.partial(jit, static_argnums=(0,))
@@ -364,6 +372,7 @@ def value_loss(value_net_apply,
   return np.sum(loss) / np.sum(reward_mask)
 
 
+# TODO(afrozm): JAX-ify this, this is too slow for pong.
 def deltas(predicted_values, rewards, mask, gamma=0.99):
   r"""Computes TD-residuals from V(s) and rewards.
 
@@ -629,6 +638,7 @@ def training_loop(
   for i in range(epochs):
     t = time.time()
     t0 = t
+    logging.vlog(1, "Epoch [% 6d] collecting trajectories.", i)
     trajs = collect_trajectories(
         env,
         policy_net_apply,
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 758555e7f..f51d64c5a 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -189,6 +189,30 @@ def test_rewards_to_go(self):
             [1.984375, 1.96875, 1.9375, 1.875, 1.75, 1.5, 1.0, 0],
         ]), rewards_to_go)
 
+  def test_rewards_to_go_really_long_sequences(self):
+    T = 1200  # pylint: disable=invalid-name
+
+    rewards = np.random.uniform(1e-3, 1e-2, (1, T))
+
+    # Make a mask, clear out a fixed number `L` of 1s from the end.
+    L = 36  # pylint: disable=invalid-name
+    assert L < T
+    rewards_mask = np.ones_like(rewards)
+    rewards_mask[0, L:] = 0
+
+    gamma = 0.94
+
+    actual_r2g = ppo.rewards_to_go(rewards, rewards_mask, gamma).reshape(-1)
+
+    # Let's compute r2g the slow way.
+    masked_rewards = (rewards_mask * rewards).reshape(-1)
+    expected_r2g = np.zeros_like(masked_rewards)
+    for t in range(T):
+      for j in range(t, T):
+        expected_r2g[t] += (gamma**(j-t)) * masked_rewards[j]
+
+    self.assertAllClose(expected_r2g, actual_r2g)
+
   def test_value_loss(self):
     rewards = np.array([
         [1, 2, 4, 8, 16, 32, 64, 128],

From 33104de3e928255ccb348643f6d1c5e3994d0adf Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 19 Apr 2019 05:09:22 -0700
Subject: [PATCH 1916/2720] In multi-core TRAX, split random generator properly
 across cores.

PiperOrigin-RevId: 244346333
---
 tensor2tensor/trax/trax.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index b59fb07af..d72d65da5 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -259,7 +259,7 @@ def mapped_predict(params, batch, rng):
     pred = mapped_predict(
         params,
         reshape_by_device(batch, num_devices),
-        jax.replicate(rng))
+        jax_random.split(rng, num_devices))
     batch_size = batch.shape[0]
     return np.reshape(pred, [batch_size] + list(pred.shape[2:]))
 
@@ -289,7 +289,8 @@ def mapped_update(i, opt_state, batch, rng):
 
   def update(i, opt_state, batch, rng):
     # TODO(lukaszkaiser): investigate how to replicate rng and correct.
-    return mapped_update(jax.replicate(i), opt_state, batch, jax.replicate(rng))
+    rngs = jax_random.split(rng, num_devices)
+    return mapped_update(jax.replicate(i), opt_state, batch, rngs)
 
   return update
 

From e48cf23c505565fd63378286d9722a1632f4bef7 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 19 Apr 2019 05:18:13 -0700
Subject: [PATCH 1917/2720] Cleaning TRAX Transformer: use dropout everywhere,
 pull out residual feed forward layer; code readability improvements.

PiperOrigin-RevId: 244346833
---
 tensor2tensor/trax/models/transformer.py | 195 ++++++++++-------------
 tensor2tensor/trax/stax/attention.py     |  22 +--
 tensor2tensor/trax/stax/stax_base.py     |   8 +-
 3 files changed, 104 insertions(+), 121 deletions(-)

diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 4eb4b45e9..0a4677d26 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -18,13 +18,25 @@
 from __future__ import division
 from __future__ import print_function
 
-import collections
-from jax import random
-import jax.numpy as np
 import tensor2tensor.trax.stax as stax
 
 
-def TransformerEncoder(mode='train',  # pylint: disable=invalid-name
+def ResidualFeedForward(feature_depth,
+                        feedforward_depth,
+                        dropout,
+                        mode):
+  """Residual feed-forward layer with normalization at start."""
+  return stax.residual(
+      stax.LayerNorm(),
+      stax.Dense(feedforward_depth, W_init=stax.xavier_uniform()),
+      stax.Relu,
+      stax.Dropout(dropout, mode=mode),
+      stax.Dense(feature_depth, W_init=stax.xavier_uniform()),
+      stax.Dropout(dropout, mode=mode)
+  )
+
+
+def TransformerEncoder(mode='train',
                        num_layers=6,
                        feature_depth=512,
                        feedforward_depth=2048,
@@ -45,20 +57,12 @@ def TransformerEncoder(mode='train',  # pylint: disable=invalid-name
     A staxlayer for implementing a raw Transformer encoder stack.  No embedding
     or positional signals are added by this layer.
   """
-  keep_rate = 1.0 - dropout
   # Multi-headed Attention and Feed-forward layers
   multi_attention = stax.MultiHeadedAttention(
-      feature_depth, num_heads=num_heads, dropout=keep_rate, mode=mode)
-
-  feed_forward = stax.serial(
-      stax.Dense(feedforward_depth, W_init=stax.xavier_uniform()),
-      stax.Relu,
-      stax.Dropout(keep_rate, mode=mode),
-      stax.Dense(feature_depth, W_init=stax.xavier_uniform())
-  )
+      feature_depth, num_heads=num_heads, dropout=dropout, mode=mode)
 
   @stax.Lambda
-  def encoder(embedded_source, source_mask):
+  def Encoder(embedded_source, source_mask):
     """Transformer encoder stack.
 
     Args:
@@ -77,11 +81,10 @@ def encoder(embedded_source, source_mask):
                                     stax.Identity,  # value
                                     source_mask),  # attention mask
                       multi_attention,
-                      stax.Dropout(keep_rate, mode=mode)),
+                      stax.Dropout(dropout, mode=mode)),
         # feed-forward
-        stax.residual(stax.LayerNorm(),
-                      feed_forward,
-                      stax.Dropout(keep_rate, mode=mode))
+        ResidualFeedForward(
+            feature_depth, feedforward_depth, dropout, mode=mode)
     )
     return stax.serial(
         embedded_source,
@@ -89,74 +92,84 @@ def encoder(embedded_source, source_mask):
         stax.LayerNorm(),
     )
 
-  return encoder
+  return Encoder
 
 
-def TransformerLM(vocab_size,  # pylint: disable=invalid-name
-                  mode='train',
-                  num_layers=6,
+def DecoderLayer(feature_depth,
+                 feedforward_depth,
+                 num_heads,
+                 dropout,
+                 mode):
+  """Transformer decoder layer.
+
+  Args:
+    feature_depth: int:  depth of embedding
+    feedforward_depth: int: depth of feed-forward layer
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate (how much to drop out)
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    init and apply.
+  """
+  return stax.serial(
+      stax.residual(  # Self-attention block.
+          stax.LayerNorm(),
+          stax.FanOut(4),
+          stax.parallel(stax.Identity,  # query
+                        stax.Identity,  # key
+                        stax.Identity,  # value
+                        stax.CausalMask(axis=-2)),  # attention mask
+          stax.MultiHeadedAttention(feature_depth, num_heads=num_heads,
+                                    dropout=dropout, mode=mode),
+          stax.Dropout(dropout, mode=mode)
+      ),
+      ResidualFeedForward(feature_depth, feedforward_depth, dropout, mode=mode)
+  )
+
+
+def TransformerLM(vocab_size,
                   feature_depth=512,
                   feedforward_depth=2048,
+                  num_layers=6,
                   num_heads=8,
                   dropout=0.1,
-                  max_len=2048):
+                  max_len=2048,
+                  mode='train'):
   """Transformer language model (only uses the decoder part of Transformer).
 
   Args:
     vocab_size: int: vocab size
-    mode: str: 'train' or 'eval'
-    num_layers: int: number of encoder/decoder layers
     feature_depth: int:  depth of embedding
     feedforward_depth: int: depth of feed-forward layer
+    num_layers: int: number of encoder/decoder layers
     num_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
     max_len: int: maximum symbol length for positional encoding
+    mode: str: 'train' or 'eval'
 
   Returns:
     init and apply.
   """
-  keep_rate = 1.0 - dropout
-  # Multi-headed Attention and Feed-forward layers
-  multi_attention = stax.MultiHeadedAttention(
-      feature_depth, num_heads=num_heads, dropout=keep_rate, mode=mode)
-
-  feed_forward = stax.serial(
-      stax.Dense(feedforward_depth, W_init=stax.xavier_uniform()),
-      stax.Relu,
-      stax.Dropout(keep_rate, mode=mode),
-      stax.Dense(feature_depth, W_init=stax.xavier_uniform())
-  )
-
-  # Single decoder layer
-  decoder_layer = stax.serial(
-      # target attends to self
-      stax.residual(stax.LayerNorm(),
-                    stax.FanOut(4),
-                    stax.parallel(stax.Identity,  # query
-                                  stax.Identity,  # key
-                                  stax.Identity,  # value
-                                  stax.CausalMask(axis=-2)),  # attention mask
-                    multi_attention,
-                    stax.Dropout(keep_rate, mode=mode)),
-      # feed-forward
-      stax.residual(stax.LayerNorm(),
-                    feed_forward,
-                    stax.Dropout(keep_rate, mode=mode))
-  )
-
   return stax.serial(
       stax.ShiftRight(),
       stax.Embedding(feature_depth, vocab_size),
-      stax.Dropout(keep_rate, mode=mode),
+      stax.Dropout(dropout, mode=mode),
       stax.PositionalEncoding(feature_depth, max_len=max_len),
-      stax.repeat(decoder_layer, num_layers),
+      stax.repeat(
+          DecoderLayer(
+              feature_depth, feedforward_depth, num_heads, dropout, mode),
+          num_layers),
       stax.LayerNorm(),
       stax.Dense(vocab_size, W_init=stax.xavier_uniform()),
       stax.LogSoftmax
   )
 
 
-def Transformer(source_vocab_size,  # pylint: disable=invalid-name
+# TODO(lukaszkaiser): rewrite the model below.
+
+
+def Transformer(source_vocab_size,
                 target_vocab_size,
                 mode='train',
                 num_layers=6,
@@ -187,10 +200,9 @@ def Transformer(source_vocab_size,  # pylint: disable=invalid-name
   the 'evals' functions that itself returns a namedtuple containing evaluation
   functions for the trained encoder, decoder, and generator substax.
   """
-  keep_rate = 1.0 - dropout
   # Input embedding and positional encoding
   inject_position = stax.serial(
-      stax.Dropout(keep_rate, mode=mode),
+      stax.Dropout(dropout, mode=mode),
       stax.PositionalEncoding(feature_depth, max_len=max_len)
   )
   if shared_embedding:
@@ -207,18 +219,11 @@ def Transformer(source_vocab_size,  # pylint: disable=invalid-name
 
   # Multi-headed Attention and Feed-forward layers
   multi_attention = stax.MultiHeadedAttention(
-      feature_depth, num_heads=num_heads, dropout=keep_rate, mode=mode)
-
-  feed_forward = stax.serial(
-      stax.Dense(feedforward_depth, W_init=stax.xavier_uniform()),
-      stax.Relu,
-      stax.Dropout(keep_rate, mode=mode),
-      stax.Dense(feature_depth, W_init=stax.xavier_uniform())
-  )
+      feature_depth, num_heads=num_heads, dropout=dropout, mode=mode)
 
   # Encoder
   @stax.Lambda
-  def encoder(source, source_mask):
+  def Encoder(source, source_mask):
     """Transformer encoder stack.
 
     Args:
@@ -237,11 +242,10 @@ def encoder(source, source_mask):
                                     stax.Identity,  # value
                                     source_mask),  # attention mask
                       multi_attention,
-                      stax.Dropout(keep_rate, mode=mode)),
+                      stax.Dropout(dropout, mode=mode)),
         # feed-forward
-        stax.residual(stax.LayerNorm(),
-                      feed_forward,
-                      stax.Dropout(keep_rate, mode=mode))
+        ResidualFeedForward(
+            feature_depth, feedforward_depth, dropout, mode=mode),
     )
     return stax.serial(
         source,
@@ -252,7 +256,7 @@ def encoder(source, source_mask):
 
   # Decoder
   @stax.Lambda
-  def decoder(memory, target, target_mask, memory_mask):
+  def Decoder(memory, target, target_mask, memory_mask):
     """Transformer decoder stack.
 
     Args:
@@ -273,7 +277,7 @@ def decoder(memory, target, target_mask, memory_mask):
                                     stax.Identity,  # value
                                     target_mask),  # attention mask
                       multi_attention,
-                      stax.Dropout(keep_rate, mode=mode)),
+                      stax.Dropout(dropout, mode=mode)),
         # target attends to encoded source
         stax.residual(stax.LayerNorm(),
                       stax.FanOut(4),
@@ -282,11 +286,10 @@ def decoder(memory, target, target_mask, memory_mask):
                                     memory,  # value
                                     memory_mask),  # attention mask
                       multi_attention,
-                      stax.Dropout(keep_rate, mode=mode)),
+                      stax.Dropout(dropout, mode=mode)),
         # feed-forward
-        stax.residual(stax.LayerNorm(),
-                      feed_forward,
-                      stax.Dropout(keep_rate, mode=mode))
+        ResidualFeedForward(
+            feature_depth, feedforward_depth, dropout, mode=mode)
     )
     return stax.serial(
         target,
@@ -297,13 +300,13 @@ def decoder(memory, target, target_mask, memory_mask):
 
   # The Transformer
   @stax.Lambda
-  def transformer(source, target, source_mask, target_mask, memory_mask):
-    encoded_source = encoder(source, source_mask)
-    return decoder(encoded_source, target, target_mask, memory_mask)
+  def transformer(source, target, source_mask, target_mask, memory_mask):  # pylint: disable=invalid-name
+    encoded_source = Encoder(source, source_mask)
+    return Decoder(encoded_source, target, target_mask, memory_mask)
 
   # Finally, bind the generator transform to use later for inference.
   @stax.Lambda
-  def generator(encoded_target):
+  def Generator(encoded_target):
     return stax.serial(
         encoded_target,
         stax.Dense(target_vocab_size, W_init=stax.xavier_uniform()),
@@ -312,34 +315,10 @@ def generator(encoded_target):
 
   # Model-Building and Evaluation Functions
   # Get entire model's init and apply pair
-  top_init, top_apply = generator(transformer)
+  top_init, top_apply = Generator(transformer)
 
   # By default act as a normal Stax constructor and emit an (init, apply) pair.
   if not return_evals:
     return (top_init, top_apply)
   else:
-    # Inference-time function for binding trained params to model and returning
-    # the python-bound sub-expressions for evaluation and sequence generation.
-    def make_namedtuple(**kwargs):
-      return collections.namedtuple('Model', kwargs.keys())(**kwargs)
-
-    def get_evals(params):
-      # We need to feed _concrete_ trained parameters through the network once.
-      # Otherwise the bound parameters point to abstract tracer values.
-      # The inputs don't matter.
-      fake_inputs = 5 * (np.ones((1), dtype=np.int32),)
-      fake_key = random.PRNGKey(1)
-      top_apply(params, fake_inputs, rng=fake_key)
-      # We can now return eval functions from the bound pieces of the model.
-      return make_namedtuple(
-          encoder=stax.make_apply_fun(encoder),
-          generator=stax.make_apply_fun(generator),
-          decoder=stax.make_apply_fun(decoder),
-      )
-
-    # We return the functions needed to train and evaluate the Transformer.
-    return make_namedtuple(
-        init=top_init,
-        apply=top_apply,
-        evals=get_evals,
-    )
+    raise ValueError('inference in this model is still a work in progress')
diff --git a/tensor2tensor/trax/stax/attention.py b/tensor2tensor/trax/stax/attention.py
index d3a24e754..cdb52eb9c 100644
--- a/tensor2tensor/trax/stax/attention.py
+++ b/tensor2tensor/trax/stax/attention.py
@@ -144,7 +144,7 @@ def dot_product_attention(query, key, value, mask, dropout, mode, rng):
     key: array of representations
     value: array of representations
     mask: attention-mask, gates attention
-    dropout: float: dropout rate - keep probability
+    dropout: float: dropout rate
     mode: 'eval' or 'train': whether to use dropout
     rng: JAX PRNGKey: subkey for disposable use
 
@@ -156,18 +156,20 @@ def dot_product_attention(query, key, value, mask, dropout, mode, rng):
   if mask is not None:
     dots = np.where(mask, dots, -1e9)
   dots = stax.softmax(dots, axis=-1)
-  if dropout is not None and mode == 'train':
-    keep = random.bernoulli(rng, dropout, dots.shape)
-    dots = np.where(keep, dots / dropout, 0)
+  if dropout >= 1.0:
+    raise ValueError('Dropout rates must be lower than 1.')
+  if dropout is not None and dropout > 0.0 and mode == 'train':
+    keep = random.bernoulli(rng, 1.0 - dropout, dots.shape)
+    dots = np.where(keep, dots / (1.0 - dropout), 0)
   out = np.matmul(dots, value)
   return out
 
 
-def PureDotProductAttention(dropout=1.0, mode='train'):  # pylint: disable=invalid-name
+def PureDotProductAttention(dropout=0.0, mode='train'):  # pylint: disable=invalid-name
   """Pure single-headed self-attention.
 
   Args:
-    dropout: float: dropout rate - keep probability
+    dropout: float: dropout rate
     mode: str: 'train' or 'eval'
 
   Returns:
@@ -187,13 +189,13 @@ def apply_fun(params, inputs, **kwargs):
 
 
 def PureMultiHeadedAttention(  # pylint: disable=invalid-name
-    feature_depth, num_heads=8, dropout=1.0, mode='train'):
+    feature_depth, num_heads=8, dropout=0.0, mode='train'):
   """Pure transformer-style multi-headed attention.
 
   Args:
     feature_depth: int:  depth of embedding
     num_heads: int: number of attention heads
-    dropout: float: dropout rate - keep probability
+    dropout: float: dropout rate
     mode: str: 'train' or 'eval'
 
   Returns:
@@ -227,13 +229,13 @@ def join_heads(x):
 
 
 def MultiHeadedAttention(  # pylint: disable=invalid-name
-    feature_depth, num_heads=8, dropout=1.0, mode='train'):
+    feature_depth, num_heads=8, dropout=0.0, mode='train'):
   """Transformer-style multi-headed attention.
 
   Args:
     feature_depth: int:  depth of embedding
     num_heads: int: number of attention heads
-    dropout: float: dropout rate - keep probability
+    dropout: float: dropout rate
     mode: str: 'train' or 'eval'
 
   Returns:
diff --git a/tensor2tensor/trax/stax/stax_base.py b/tensor2tensor/trax/stax/stax_base.py
index f690c982a..8a1ece9f8 100644
--- a/tensor2tensor/trax/stax/stax_base.py
+++ b/tensor2tensor/trax/stax/stax_base.py
@@ -273,9 +273,11 @@ def apply_fun(params, inputs, **kwargs):  # pylint: disable=missing-docstring
              'it like `apply_fun(params, inputs, key)` where `key` is a '
              'jax.random.PRNGKey value.')
       raise ValueError(msg)
-    if mode == 'train':
-      keep = backend.random.bernoulli(rng, rate, inputs.shape)
-      return np.where(keep, inputs / rate, 0)
+    if rate >= 1.0:
+      raise ValueError('Dropout rates must be lower than 1.')
+    if mode == 'train' and rate > 0.0:
+      keep = backend.random.bernoulli(rng, 1.0 - rate, inputs.shape)
+      return np.where(keep, inputs / (1.0 - rate), 0)
     else:
       return inputs
   return init_fun, apply_fun

From 9624e410a6f5a051a19dee93001d550a2ad2c263 Mon Sep 17 00:00:00 2001
From: Dumitru Erhan <dumitru@google.com>
Date: Fri, 19 Apr 2019 11:33:12 -0700
Subject: [PATCH 1918/2720] Adding a few new RLMB configs

PiperOrigin-RevId: 244393153
---
 .../rl/trainer_model_based_params.py          | 32 ++++++++++++++++---
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 515282678..7975ca25a 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -324,6 +324,30 @@ def rlmb_base_stochastic_discrete():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_base_stochastic_discrete_200k():
+  """Base setting with stochastic discrete model with 200k steps."""
+  hparams = rlmb_base_stochastic_discrete()
+  hparams.num_real_env_frames = 96000 * 2
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic_discrete_500k():
+  """Base setting with stochastic discrete model with 500k steps."""
+  hparams = rlmb_base_stochastic_discrete()
+  hparams.num_real_env_frames = 96000 * 5
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic_discrete_1m():
+  """Base setting with stochastic discrete model with 1M steps."""
+  hparams = rlmb_base_stochastic_discrete()
+  hparams.num_real_env_frames = 96000 * 10
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_base_stochastic_discrete_param_sharing():
   """Base setting with stochastic discrete model with parameter sharing."""
@@ -406,18 +430,18 @@ def rlmb_long_stochastic_discrete_gamma90():
 
 
 @registry.register_hparams
-def rlmb_long_stochastic_discrete_3epochs():
+def rlmb_base_stochastic_discrete_3epochs():
   """Long setting with stochastic discrete model, changed epochs."""
-  hparams = rlmb_long_stochastic_discrete()
+  hparams = rlmb_base_stochastic_discrete()
   hparams.epochs = 3
   hparams.ppo_epochs_num = 2000
   return hparams
 
 
 @registry.register_hparams
-def rlmb_long_stochastic_discrete_1epoch():
+def rlmb_base_stochastic_discrete_1epoch():
   """Long setting with stochastic discrete model, changed epochs."""
-  hparams = rlmb_long_stochastic_discrete()
+  hparams = rlmb_base_stochastic_discrete()
   hparams.epochs = 1
   hparams.ppo_epochs_num = 3000
   return hparams

From ac01285fff09655b277c9e370e266542d8576934 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 19 Apr 2019 11:33:40 -0700
Subject: [PATCH 1919/2720] Allow to specify batch size divided by number of
 devices in trax, remove now-redundant tpu config.

PiperOrigin-RevId: 244393232
---
 .../trax/configs/resnet50_imagenet_8gb.gin    |  2 +-
 ...u_big.gin => transformer_big_lm1b_8gb.gin} |  2 +-
 .../trax/configs/transformer_lm1b_8gb.gin     |  2 +-
 .../trax/configs/transformer_lm1b_tpu.gin     | 52 -------------------
 tensor2tensor/trax/inputs.py                  |  4 +-
 5 files changed, 6 insertions(+), 56 deletions(-)
 rename tensor2tensor/trax/configs/{transformer_lm1b_tpu_big.gin => transformer_big_lm1b_8gb.gin} (98%)
 delete mode 100644 tensor2tensor/trax/configs/transformer_lm1b_tpu.gin

diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
index 01cac65ac..243ea5a39 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
@@ -6,7 +6,7 @@ import tensor2tensor.trax.trax
 
 # Parameters for batch_fun:
 # ==============================================================================
-batch_fun.batch_size = 32
+batch_fun.batch_size_per_device = 32
 batch_fun.bucket_length = 32
 batch_fun.buckets = None
 batch_fun.eval_batch_size = 32
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_tpu_big.gin b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
similarity index 98%
rename from tensor2tensor/trax/configs/transformer_lm1b_tpu_big.gin
rename to tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
index 64cc2a2ed..99ee6ea5d 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_tpu_big.gin
+++ b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
@@ -5,7 +5,7 @@ import tensor2tensor.trax.trax
 
 # Parameters for batch_fun:
 # ==============================================================================
-batch_fun.batch_size = 256
+batch_fun.batch_size_per_device = 16
 batch_fun.eval_batch_size = 128
 batch_fun.max_eval_length = 2048
 
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index cf47eaee0..89aa4bd48 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -5,7 +5,7 @@ import tensor2tensor.trax.trax
 
 # Parameters for batch_fun:
 # ==============================================================================
-batch_fun.batch_size = 128
+batch_fun.batch_size_per_device = 128
 batch_fun.eval_batch_size = 128
 batch_fun.max_eval_length = 2048
 
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_tpu.gin b/tensor2tensor/trax/configs/transformer_lm1b_tpu.gin
deleted file mode 100644
index fd264af49..000000000
--- a/tensor2tensor/trax/configs/transformer_lm1b_tpu.gin
+++ /dev/null
@@ -1,52 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fun:
-# ==============================================================================
-batch_fun.batch_size = 1024
-batch_fun.eval_batch_size = 128
-batch_fun.max_eval_length = 2048
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_languagemodel_lm1b32k'
-inputs.input_name = 'targets'
-
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 0.1
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
-lm1b_preprocess.max_target_length = 512
-lm1b_preprocess.max_eval_target_length = 2048
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 1000
-train.eval_steps = 10
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerLM
-train.run_debug_step = False
-train.train_steps = 500000
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.dropout = 0.1
-TransformerLM.feature_depth = 512
-TransformerLM.feedforward_depth = 2048
-TransformerLM.max_len = 2048
-TransformerLM.mode = 'train'
-TransformerLM.num_heads = 8
-TransformerLM.num_layers = 6
-TransformerLM.vocab_size = 32000
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 8b33738a1..ffeb4d639 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -239,11 +239,13 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
 @gin.configurable(blacklist=["dataset", "training", "shapes",
                              "target_names", "num_devices"])
 def batch_fun(dataset, training, shapes, target_names, num_devices,
-              batch_size=32, eval_batch_size=32,
+              batch_size_per_device=32, batch_size=None, eval_batch_size=32,
               bucket_length=32, buckets=None,
               batch_shuffle_size=128, max_eval_length=None):
   """Batching function."""
   del target_names
+  # Batch size is batch_size_per_device * num_devices unless given directly.
+  batch_size = batch_size or batch_size_per_device * num_devices
   # If bucketing is not specified, check if target shapes are variable.
   cur_batch_size = batch_size if training else eval_batch_size
   # Make cur_batch_size divisible by num_devices.

From c5faadef0406718311c4ab1bae02f0d5cc251b09 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 19 Apr 2019 11:33:43 -0700
Subject: [PATCH 1920/2720] Add some debugging logs to collect that dump out a
 whole bunch of stuff before crashing.

PiperOrigin-RevId: 244393243
---
 tensor2tensor/trax/rlax/ppo.py | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index ec52f98c4..0cb0f1ff9 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -167,7 +167,29 @@ def collect_trajectories(env,
         raise ValueError("Unknown policy: %s" % policy)
 
       # NOTE: Assumption, single batch.
-      action = int(action)
+      try:
+        action = int(action)
+      except TypeError as err:
+        # Let's dump some information before we die off.
+        logging.error("Cannot convert action into an integer: [%s]", err)
+        logging.error("action.shape: [%s]", action.shape)
+        logging.error("action: [%s]", action)
+        logging.error("predictions.shape: [%s]", predictions.shape)
+        logging.error("predictions: [%s]", predictions)
+        logging.error("observation_history: [%s]", observation_history)
+        logging.error("policy_net_params: [%s]", policy_net_params)
+        for i, param in enumerate(policy_net_params):
+          if not param:
+            # Empty tuple.
+            continue
+          if not isinstance(param, tuple):
+            logging.error(
+                "Param[%d] : (%s) = [%s]", i, param.shape, onp.array(param))
+          else:
+            for j, p in enumerate(param):
+              logging.error(
+                  "\tParam[%d, %d] : (%s) = [%s]", i, j, p.shape, onp.array(p))
+        raise err
 
       observation, reward, done, _ = env.step(action)
 
@@ -808,3 +830,4 @@ def training_loop(
 
   return ((policy_net_params, value_net_params), average_rewards,
           np.stack(value_losses), np.stack(ppo_objective))
+

From a0158006b65e1d6a6c88378de6efee8e24e32fea Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 19 Apr 2019 11:50:11 -0700
Subject: [PATCH 1921/2720] Split rngs per device in the main loop, not the
 update function.

PiperOrigin-RevId: 244396062
---
 tensor2tensor/trax/trax.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index d72d65da5..1545d1cf6 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -270,27 +270,27 @@ def _jit_update_fun(predict_fun, loss_fun, optimizer, lr_fun, num_devices):
   """Get jit-ed update function for loss, optimizer, learning rate function."""
   if num_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
     def single_update(i, opt_state, batch, rng):
+      rng, subrng = jax_random.split(rng[0])
       _, opt_update = optimizer(lr_fun)
       params = trax_opt.get_params(opt_state)
       return opt_update(i, backend.grad(loss_fun)(
-          params, batch, predict_fun, rng), opt_state)
+          params, batch, predict_fun, rng), opt_state), [subrng]
     return backend.jit(single_update)
 
   @functools.partial(backend.pmap, axis_name="batch")
   def mapped_update(i, opt_state, batch, rng):
     """This is a multi-device version of the update function above."""
     # We assume all tensors have the first dimension = num_devices.
+    rng, subrng = jax_random.split(rng)
     _, opt_update = optimizer(lr_fun)
     params = trax_opt.get_params(opt_state)
     grads = backend.grad(loss_fun)(params, batch, predict_fun, rng)
     grads = jax.tree_util.tree_map(
         lambda g: lax.psum(g, "batch"), grads)
-    return opt_update(i, grads, opt_state)
+    return opt_update(i, grads, opt_state), subrng
 
   def update(i, opt_state, batch, rng):
-    # TODO(lukaszkaiser): investigate how to replicate rng and correct.
-    rngs = jax_random.split(rng, num_devices)
-    return mapped_update(jax.replicate(i), opt_state, batch, rngs)
+    return mapped_update(jax.replicate(i), opt_state, batch, rng)
 
   return update
 
@@ -379,6 +379,7 @@ def train(output_dir,
   # Setup state
   step = state.step or 0
   rng, init_key = jax_random.split(rng)
+  rngs = jax_random.split(rng, num_devices)
   params_initializer = \
       lambda: model_init(init_key, [-1] + list(inputs.input_shape))[1]
   params = state.params or params_initializer()
@@ -417,8 +418,7 @@ def train(output_dir,
       next_train_batch = next(train_stream)
       if num_devices > 1:  # TODO(lukaszkaiser): use everywhere when possible.
         next_train_batch = reshape_by_device_pair(next_train_batch, num_devices)
-      rng, subrng = jax_random.split(rng)
-      opt_state = jit_update_fun(step, opt_state, next_train_batch, subrng)
+      opt_state, rngs = jit_update_fun(step, opt_state, next_train_batch, rngs)
       step += 1
 
       # LR log

From 70c73803eba9b14095b0f065db84df68fdf69297 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 19 Apr 2019 12:02:34 -0700
Subject: [PATCH 1922/2720] Allow to configure learning rate and number of
 optimizer steps from flags.

TODO(afrozm): Use gin :)

PiperOrigin-RevId: 244398248
---
 tensor2tensor/trax/rlax/ppo.py      | 4 ++--
 tensor2tensor/trax/rlax/ppo_main.py | 8 ++++++++
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 0cb0f1ff9..6c9dbc326 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -106,9 +106,9 @@ def value_net(rng_key,
   return net_params, net_apply
 
 
-def optimizer_fun(net_params):
+def optimizer_fun(net_params, step_size=1e-3):
   opt_init, opt_update = trax_opt.adam(
-      step_size=1e-3, b1=0.9, b2=0.999, eps=1e-08)
+      step_size=step_size, b1=0.9, b2=0.999, eps=1e-08)
   opt_state = opt_init(net_params)
   return opt_state, opt_update
 
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index aa343f42f..a5b9d31af 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -35,8 +35,10 @@
 flags.DEFINE_integer("random_seed", 0, "Random seed.")
 flags.DEFINE_integer("log_level", logging.INFO, "Log level.")
 flags.DEFINE_integer("batch_size", 32, "Batch of trajectories needed.")
+flags.DEFINE_integer("num_optimizer_steps", 100, "Number of optimizer steps.")
 flags.DEFINE_integer("boundary", 20,
                      "We pad trajectories at integer multiples of this number.")
+flags.DEFINE_float("learning_rate", 1e-3, "Learning rate.")
 
 
 def common_stax_layers():
@@ -51,6 +53,9 @@ def main(argv):
   if FLAGS.env_name == "Pong-v0":
     bottom_layers = [stax.Div(255.0), stax.Flatten(2)] + bottom_layers
 
+  optimizer_fun = functools.partial(ppo.optimizer_fun,
+                                    step_size=FLAGS.learning_rate)
+
   ppo.training_loop(
       env_name=FLAGS.env_name,
       epochs=FLAGS.epochs,
@@ -58,7 +63,10 @@ def main(argv):
           ppo.policy_net, bottom_layers=bottom_layers),
       value_net_fun=functools.partial(
           ppo.value_net, bottom_layers=bottom_layers),
+      policy_optimizer_fun=optimizer_fun,
+      value_optimizer_fun=optimizer_fun,
       batch_size=FLAGS.batch_size,
+      num_optimizer_steps=FLAGS.num_optimizer_steps,
       boundary=FLAGS.boundary,
       random_seed=FLAGS.random_seed)
 

From f22ab2a43a294f48aa4c71c1b4bce6a989575a99 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 19 Apr 2019 12:31:18 -0700
Subject: [PATCH 1923/2720] Adding a policy and value function that produces
 both policy and value.

NOTE: This isn't used anywhere right now though.
PiperOrigin-RevId: 244402893
---
 tensor2tensor/trax/rlax/ppo.py      | 24 ++++++++++++++++++++++++
 tensor2tensor/trax/rlax/ppo_test.py | 17 +++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 6c9dbc326..fdecb4f65 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -106,6 +106,30 @@ def value_net(rng_key,
   return net_params, net_apply
 
 
+def policy_and_value_net(rng_key,
+                         batch_observations_shape,
+                         num_actions,
+                         bottom_layers=None):
+  """A policy and value net function."""
+
+  # Layers.
+  layers = []
+  if bottom_layers is not None:
+    layers.extend(bottom_layers)
+
+  # Now, with the current logits, one head computes action probabilities and the
+  # other computes the value function.
+  layers.extend([stax.FanOut(2), stax.parallel(
+      stax.serial(stax.Dense(num_actions), stax.Softmax),
+      stax.Dense(1)
+  )])
+
+  net_init, net_apply = stax.serial(*layers)
+
+  _, net_params = net_init(rng_key, batch_observations_shape)
+  return net_params, net_apply
+
+
 def optimizer_fun(net_params, step_size=1e-3):
   opt_init, opt_update = trax_opt.adam(
       step_size=step_size, b1=0.9, b2=0.999, eps=1e-08)
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index f51d64c5a..6a0599a33 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -76,6 +76,23 @@ def test_value_net(self):
     # NOTE: The extra dimension at the end because of Dense(1).
     self.assertEqual((batch, time_steps, 1), value_output.shape)
 
+  def test_policy_and_value_net(self):
+    observation_shape = (3, 4, 5)
+    batch_observation_shape = (-1, -1) + observation_shape
+    num_actions = 2
+    pnv_params, pnv_apply = ppo.policy_and_value_net(
+        self.rng_key, batch_observation_shape, num_actions, [stax.Flatten(2)])
+    batch = 2
+    time_steps = 10
+    batch_of_observations = np.random.uniform(
+        size=(batch, time_steps) + observation_shape)
+    pnv_output = pnv_apply(pnv_params, batch_of_observations)
+
+    # Output is a list, first is probab of actions and the next is value output.
+    self.assertEqual(2, len(pnv_output))
+    self.assertEqual((batch, time_steps, num_actions), pnv_output[0].shape)
+    self.assertEqual((batch, time_steps, 1), pnv_output[1].shape)
+
   def test_collect_trajectories(self):
     observation_shape = (2, 3, 4)
     num_actions = 2

From 07a599b8db68ae8ab9a90f136a1853e4c0bce90a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 19 Apr 2019 14:51:53 -0700
Subject: [PATCH 1924/2720] Cleaning TRAX layers, step 1: remove unused parts.

We may re-add some of the parts we are removing now, but it will be easier to refactor with less parts.

All models (MLP, ResNet, TransformerLM) work fine, other Transformers need refactoring anyway.

Keeping the slax Shared layer so as to use it as the base layer later
(we will want to enable sharing weights by object as that's very natural).

PiperOrigin-RevId: 244427175
---
 tensor2tensor/trax/models/resnet.py  |   5 +-
 tensor2tensor/trax/stax/__init__.py  |   3 +-
 tensor2tensor/trax/stax/attention.py |   2 +-
 tensor2tensor/trax/stax/base.py      |  78 ++++++
 tensor2tensor/trax/stax/base_test.py | 135 ++++++++++
 tensor2tensor/trax/stax/losses.py    |  49 ----
 tensor2tensor/trax/stax/slax.py      | 383 ---------------------------
 tensor2tensor/trax/stax/slax_test.py | 250 -----------------
 tensor2tensor/trax/stax/stax_base.py |  48 ++++
 9 files changed, 267 insertions(+), 686 deletions(-)
 create mode 100644 tensor2tensor/trax/stax/base.py
 create mode 100644 tensor2tensor/trax/stax/base_test.py
 delete mode 100644 tensor2tensor/trax/stax/losses.py
 delete mode 100644 tensor2tensor/trax/stax/slax.py
 delete mode 100644 tensor2tensor/trax/stax/slax_test.py

diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index 45a9e4d98..0c81dee0e 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -116,17 +116,20 @@ def WideResnetGroup(n, channels, strides=(1, 1)):
   return stax.serial(*blocks)
 
 
-def WideResnet(num_blocks=3, hidden_size=64, num_output_classes=10):
+def WideResnet(num_blocks=3, hidden_size=64, num_output_classes=10,
+               mode='train'):
   """WideResnet from https://arxiv.org/pdf/1605.07146.pdf.
 
   Args:
     num_blocks: int, number of blocks in a group.
     hidden_size: the size of the first hidden layer (multiplied later).
     num_output_classes: int, number of classes to distinguish.
+    mode: is it training or eval.
 
   Returns:
     The WideResnet model with given layer and output sizes.
   """
+  del mode
   return stax.serial(
       stax.Conv(hidden_size, (3, 3), padding='SAME'),
       WideResnetGroup(num_blocks, hidden_size),
diff --git a/tensor2tensor/trax/stax/__init__.py b/tensor2tensor/trax/stax/__init__.py
index eb0c51790..88fd1cd96 100644
--- a/tensor2tensor/trax/stax/__init__.py
+++ b/tensor2tensor/trax/stax/__init__.py
@@ -22,6 +22,5 @@
 # upstream changes.
 # pylint: disable=wildcard-import
 from tensor2tensor.trax.stax.attention import *
-from tensor2tensor.trax.stax.losses import *
-from tensor2tensor.trax.stax.slax import *
+from tensor2tensor.trax.stax.base import *
 from tensor2tensor.trax.stax.stax_base import *
diff --git a/tensor2tensor/trax/stax/attention.py b/tensor2tensor/trax/stax/attention.py
index cdb52eb9c..0879295b8 100644
--- a/tensor2tensor/trax/stax/attention.py
+++ b/tensor2tensor/trax/stax/attention.py
@@ -18,10 +18,10 @@
 from __future__ import division
 from __future__ import print_function
 
-from jax import random
 import numpy as onp
 
 from tensor2tensor.trax.backend import numpy as np
+from tensor2tensor.trax.backend import random
 from tensor2tensor.trax.stax import stax_base as stax
 
 
diff --git a/tensor2tensor/trax/stax/base.py b/tensor2tensor/trax/stax/base.py
new file mode 100644
index 000000000..1b1ab2300
--- /dev/null
+++ b/tensor2tensor/trax/stax/base.py
@@ -0,0 +1,78 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 layer class."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from jax.tree_util import register_pytree_node as _register_pytree_node
+
+
+# Staxlayer binding to python variables
+# ------------------------------------------------------------------------------
+# Stax params-tree leaf type to mark bound subtrees references.
+class _TreeMarker(dict):
+  pass
+# Add this leaf-type to JAX's tree-walker.
+_register_pytree_node(_TreeMarker,
+                      lambda xs: (tuple(), None),
+                      lambda _, xs: _TreeMarker())
+
+
+# TODO(lukaszkaiser): make this the base layer class (share by object).
+class Share(tuple):
+  """Layer parameter caching function to allow weight sharing.
+
+  Args:
+    A staxlayer: an (init_fun, apply_fun) pair.
+
+  Returns:
+    A 'parameter-bound' staxlayer that can be assigned to a python variable.
+  Wherever this value is needed elsewhere in the stax tree, call this bound
+  variable and all occurrences will share parameters that will automatically
+  be updated by Stax optimizers.
+  """
+
+  def __init__(self, staxlayer):  # pylint: disable=super-init-not-called
+    self._orig_init_fun, self._orig_apply_fun = staxlayer
+    self._first_init = True
+    self.params = None  # cached staxlayer params
+
+  def _init_fun(self, rng_key, input_shape):  # pylint: disable=missing-docstring
+    if self._first_init:
+      # point of first subgraph initialization call: sets params, output_shape
+      self._first_init = False
+      out_shape, self.params = self._orig_init_fun(rng_key, input_shape)
+      return out_shape, self.params
+    else:
+      # point of subgraph reuse:
+      # params are just a marker to apply_funs signalling subgraph params reuse
+      out_shape, _ = self._orig_init_fun(rng_key, input_shape)
+      return out_shape, _TreeMarker()
+
+  def _apply_fun(self, params, inputs, **kwargs):
+    if isinstance(params, _TreeMarker):
+      # point of subgraph reuse: calculate new value with cached params
+      return self._orig_apply_fun(self.params, inputs, **kwargs)
+    else:
+      # point of first subgraph application to params: cache params
+      self.params = params
+      return self._orig_apply_fun(params, inputs, **kwargs)
+
+  # when unpacking this (init, apply) pair we return the wrapped funs
+  def __iter__(self):
+    return iter((self._init_fun, self._apply_fun))
diff --git a/tensor2tensor/trax/stax/base_test.py b/tensor2tensor/trax/stax/base_test.py
new file mode 100644
index 000000000..5902ec7ef
--- /dev/null
+++ b/tensor2tensor/trax/stax/base_test.py
@@ -0,0 +1,135 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 Stax base layer."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import itertools
+from absl.testing import absltest
+from jax import random
+import numpy as onp
+import tensor2tensor.trax.stax as stax
+
+
+def random_inputs(rng, input_shape):
+  if isinstance(input_shape, tuple):
+    return rng.randn(*input_shape).astype(onp.float32)
+  elif isinstance(input_shape, list):
+    return [random_inputs(rng, shape) for shape in input_shape]
+  else:
+    raise TypeError(type(input_shape))
+
+
+def check_shape_agreement(test_case, init_fun, apply_fun, input_shape):
+  rng_key1, rng_key2 = random.split(random.PRNGKey(0))
+  result_shape, params = init_fun(rng_key1, input_shape)
+  inputs = random_inputs(onp.random.RandomState(0), input_shape)
+  result = apply_fun(params, inputs, rng=rng_key2)
+  test_case.assertEqual(result.shape, result_shape)
+  return result_shape
+
+
+def check_staxlayer(test_case, staxlayer, input_shape):
+  init_fun, apply_fun = staxlayer
+  return check_shape_agreement(test_case, init_fun, apply_fun, input_shape)
+
+
+# Helper functions for testing Lambda wrapper against functions involving
+# complicated input trees:
+def _enumerate_trees_w_leaves(n_leaves):
+  """Construct all rooted trees with n leaves."""
+  def enumtree(*args):
+    n_args = len(args)
+    # trivial cases:
+    if n_args == 0:
+      return []
+    if n_args == 1:
+      return args
+    # general case of 2 or more args:
+    # build index array
+    idxs = range(0, n_args)
+    trees = []
+    # we consider all possible subsets of size n_set to gather
+    for n_set in range(2, n_args+1):
+      idxsets = list(itertools.combinations(idxs, n_set))
+      for idxset in idxsets:
+        # recurse by joining all subtrees with
+        # n_set leaves and (n_args - n_set) leaves
+        arg_set = tuple(args[i] for i in idxs if i in idxset)
+        arg_coset = tuple(args[i] for i in idxs if i not in idxset)
+        if arg_coset:
+          trees.extend(tuple(itertools.product(enumtree(*arg_set),
+                                               enumtree(*arg_coset))))
+        else:
+          # trivial case where arg_set is entire set
+          trees.append(arg_set)
+    return trees
+  # return enumerated trees with integers as leaves
+  return enumtree(*range(n_leaves))
+
+
+def _build_combinator_tree(input_treespec, in_vars):
+  """Build a trivial Staxlayer that takes a complicated tree of inputs."""
+  parallel_args = []
+  for e in input_treespec:
+    if isinstance(e, int):
+      parallel_args.append(in_vars[e])
+    elif isinstance(e, tuple):
+      parallel_args.append(_build_combinator_tree(e, in_vars))
+  return stax.serial(stax.parallel(*parallel_args), stax.FanInSum)
+
+
+class SlaxTest(absltest.TestCase):
+
+  def test_flatten_n(self):
+    input_shape = (29, 87, 10, 20, 30)
+
+    actual_shape = check_staxlayer(self, stax.Flatten(1), input_shape)
+    self.assertEqual(actual_shape, (29, 87 * 10 * 20 * 30))
+
+    actual_shape = check_staxlayer(self, stax.Flatten(2), input_shape)
+    self.assertEqual(actual_shape, (29, 87, 10 * 20 * 30))
+
+    actual_shape = check_staxlayer(self, stax.Flatten(3), input_shape)
+    self.assertEqual(actual_shape, (29, 87, 10, 20 * 30))
+
+    actual_shape = check_staxlayer(self, stax.Flatten(4), input_shape)
+    self.assertEqual(actual_shape, (29, 87, 10, 20, 30))
+
+    # Not enough dimensions.
+    with self.assertRaises(ValueError):
+      check_staxlayer(self, stax.Flatten(5), input_shape)
+
+    with self.assertRaises(ValueError):
+      check_staxlayer(self, stax.Flatten(6), input_shape)
+
+  def test_div(self):
+    init_fun, apply_fun = stax.Div(2)
+    input_np = onp.array([[1, 2, 3], [4, 5, 6]], dtype=onp.float32)
+    input_shape = input_np.shape
+    _, _ = init_fun(None, input_shape)
+    output_np = apply_fun(None, input_np)
+    # absltest doesn't have ndarray equalities.
+    expected_output_np = input_np / 2.0
+    self.assertAlmostEqual(
+        0.0,
+        onp.sum((output_np - expected_output_np) ** 2),
+        delta=1e-6)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tensor2tensor/trax/stax/losses.py b/tensor2tensor/trax/stax/losses.py
deleted file mode 100644
index 56551266f..000000000
--- a/tensor2tensor/trax/stax/losses.py
+++ /dev/null
@@ -1,49 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Loss functions and layers."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gin
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.stax import slax
-
-
-@gin.configurable(blacklist=['logpred', 'target'])
-def kl_div(logpred, target, eps=np.finfo(np.float32).eps):
-  """Calculate KL-divergence."""
-  return np.sum(target * (np.log(target + eps) - logpred))
-
-
-def crossentropy_loss(logpred, target):
-  """Calculate crossentropy loss."""
-  return - np.mean(
-      np.sum(logpred * slax.one_hot(target, logpred.shape[-1]), axis=-1))
-
-
-@gin.configurable(blacklist=['logpred', 'target', 'size'])
-def label_smoothed_loss(logpred, target, size, padding_idx=0, smoothing=0.0):
-  """Returns a label-smoothing loss-criterion function."""
-  confidence = 1.0 - smoothing
-  zerosmoothed = smoothing / (size - 2)
-  delta = confidence - zerosmoothed
-  assert logpred.shape[1] == size
-  truedist = (np.full_like(logpred, zerosmoothed) +
-              delta * slax.one_hot(target, size))
-  truedist *= (1 - (np.arange(size) == padding_idx))
-  truedist *= (1 - (target == padding_idx))[:, np.newaxis]
-  return kl_div(logpred, truedist, eps=1e-6)
diff --git a/tensor2tensor/trax/stax/slax.py b/tensor2tensor/trax/stax/slax.py
deleted file mode 100644
index 37142fee0..000000000
--- a/tensor2tensor/trax/stax/slax.py
+++ /dev/null
@@ -1,383 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""SLAX - Layer eXtensions to Stax."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import inspect
-from absl import logging
-from jax.tree_util import register_pytree_node as _register_pytree_node
-
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.stax import stax_base as stax
-
-
-# Utility functions
-# ------------------------------------------------------------------------------
-def one_hot(x, size, dtype=np.float32):
-  """Make a n+1 dim one-hot array from n dim int-categorical array."""
-  return np.array(x[..., np.newaxis] == np.arange(size), dtype)
-
-
-def ShiftRight():  # pylint: disable=invalid-name
-  """Layer to shift the tensor to the right by padding on axis 1."""
-  init_fun = lambda _, input_shape: (input_shape, ())
-  def apply_fun(params, inputs, **kwargs):
-    del params, kwargs
-    pad_widths = [(0, 0), (1, 0)]
-    pad_widths += [(0, 0) for _ in range(len(inputs.shape) - 2)]
-    padded = np.pad(inputs, pad_widths, mode='constant')
-    return padded[:, :-1, ...]
-  return init_fun, apply_fun
-
-
-# Utility Combinators
-# ------------------------------------------------------------------------------
-def repeat(layer, num_repeats):
-  """Repeats layers serially num_repeats times."""
-  if num_repeats < 1:
-    raise ValueError('Repeat combinator num_repeats must be >= 1.')
-  layers = num_repeats * (layer,)
-  return stax.serial(*layers)
-
-
-def residual(*layers, **kwargs):
-  """Constructs a residual version of layers, summing input to layers output."""
-  res = kwargs.get('res', stax.Identity)
-  if len(layers) > 1:
-    return stax.serial(
-        stax.FanOut(2),
-        stax.parallel(stax.serial(*layers), res),
-        stax.FanInSum
-    )
-  elif len(layers) == 1:
-    return stax.serial(
-        stax.FanOut(2),
-        stax.parallel(layers[0], res),
-        stax.FanInSum
-    )
-  else:
-    raise ValueError('Empty residual combinator.')
-
-
-# Utility Layers
-# ------------------------------------------------------------------------------
-def Take(*args):  # pylint: disable=invalid-name
-  """Layer to pick subset of inputs from parallel input stream.
-
-  Args:
-    *args: a sequence of ints
-
-  Returns:
-    A new layer that selects inputs from an incoming parallel stream.
-    In numpy notation: outputs = parallel_inputs[args]
-    If the resulting output list has only one member, it is automatically
-    unwrapped and the contents are passed on directly.
-  """
-  def init_fun(_, input_shape):
-    output_shape = []
-    for arg in args:
-      output_shape.append(input_shape[arg])
-    if len(output_shape) == 1:
-      output_shape = output_shape[0]
-    return (output_shape, ())
-  def apply_fun(params, inputs, **kwargs):
-    del params, kwargs
-    outputs = []
-    for arg in args:
-      outputs.append(inputs[arg])
-    if len(outputs) == 1:
-      outputs = outputs[0]
-    return outputs
-  return init_fun, apply_fun
-
-
-def LogInputs(prefix='', debug=True):  # pylint: disable=invalid-name
-  """Logging side-effects layer, equivalent to Identity.
-
-  Args:
-    prefix: string: logging prefix
-    debug: bool: if True this will print logs, otherwise not.
-
-  Returns:
-    An Identity layer with log-printing side-effects. This
-  prints the types and shapes of the inputs.  NB: at the moment
-  this doesn't handle printing nested tuple/list shapes!
-  """
-  def return_shapes(inputs):
-    """Return shape information of inputs."""
-    if isinstance(inputs, _PlaceholderTree):
-      return []
-    if isinstance(inputs, (list, tuple)):
-      return [x.shape for x in inputs]
-    elif isinstance(inputs, dict):
-      return [inputs[k].shape for k in inputs.keys()]
-    else:
-      return inputs.shape
-  def init_fun(_, input_shape):
-    if debug:
-      logging.info('%s [init]: %s', prefix, input_shape)
-    return input_shape, ()
-  def apply_fun(params, inputs, **kwargs):
-    del params, kwargs
-    if debug:
-      logging.info('%s: %s %s', prefix, type(inputs), return_shapes(inputs))
-    return inputs
-  return init_fun, apply_fun
-
-
-# Staxlayer binding to python variables
-# ------------------------------------------------------------------------------
-# Stax params-tree leaf type to mark bound subtrees references.
-class _TreeMarker(dict):
-  pass
-# Add this leaf-type to JAX's tree-walker.
-_register_pytree_node(_TreeMarker,
-                      lambda xs: (tuple(), None),
-                      lambda _, xs: _TreeMarker())
-
-
-# TODO(levskaya, rsepassi): abstract away tuple-subclassing to StaxLayer?
-class Share(tuple):
-  """Layer parameter caching function to allow weight sharing.
-
-  Args:
-    A staxlayer: an (init_fun, apply_fun) pair.
-
-  Returns:
-    A 'parameter-bound' staxlayer that can be assigned to a python variable.
-  Wherever this value is needed elsewhere in the stax tree, call this bound
-  variable and all occurrences will share parameters that will automatically
-  be updated by Stax optimizers.
-  """
-
-  def __init__(self, staxlayer):  # pylint: disable=super-init-not-called
-    self._orig_init_fun, self._orig_apply_fun = staxlayer
-    self._first_init = True
-    self.params = None  # cached staxlayer params
-
-  def _init_fun(self, rng_key, input_shape):  # pylint: disable=missing-docstring
-    if self._first_init:
-      # point of first subgraph initialization call: sets params, output_shape
-      self._first_init = False
-      out_shape, self.params = self._orig_init_fun(rng_key, input_shape)
-      return out_shape, self.params
-    else:
-      # point of subgraph reuse:
-      # params are just a marker to apply_funs signalling subgraph params reuse
-      out_shape, _ = self._orig_init_fun(rng_key, input_shape)
-      return out_shape, _TreeMarker()
-
-  def _apply_fun(self, params, inputs, **kwargs):
-    if isinstance(params, _TreeMarker):
-      # point of subgraph reuse: calculate new value with cached params
-      return self._orig_apply_fun(self.params, inputs, **kwargs)
-    else:
-      # point of first subgraph application to params: cache params
-      self.params = params
-      return self._orig_apply_fun(params, inputs, **kwargs)
-
-  # when unpacking this (init, apply) pair we return the wrapped funs
-  def __iter__(self):
-    return iter((self._init_fun, self._apply_fun))
-
-
-class Bind(tuple):
-  """Layer/variable caching function to allow name binding.
-
-  Args:
-    A staxlayer: an (init_fun, apply_fun) pair.
-
-  Returns:
-    A 'bound' staxlayer that can be assigned to a python variable.
-  Wherever this value is needed elsewhere in the stax tree, call this bound
-  variable and all occurrences will share output values.
-  """
-
-  def __init__(self, staxlayer):  # pylint: disable=super-init-not-called
-    self._orig_init_fun, self._orig_apply_fun = staxlayer
-    self._first_init = True
-    self._out_shape = None  # cached staxlayer output shape
-    self.params = None  # cached staxlayer params
-    self.value = None  # cached staxlayer output value
-
-  def _init_fun(self, rng_key, input_shape):
-    if self._first_init:
-      # point of first subgraph initialization call: sets params, output_shape
-      self._first_init = False
-      self._out_shape, self.params = self._orig_init_fun(rng_key, input_shape)
-      return self._out_shape, self.params
-    else:
-      # point of subgraph reuse:
-      # params are just a marker to apply_funs signalling subgraph value reuse
-      return self._out_shape, _TreeMarker()
-
-  def _apply_fun(self, params, inputs, **kwargs):
-    if isinstance(params, _TreeMarker):
-      # point of subgraph reuse: return cached value
-      return self.value
-    else:
-      # point of first subgraph application to params: cache value
-      self.params = params
-      self.value = self._orig_apply_fun(params, inputs, **kwargs)
-      return self.value
-
-  # when unpacking this (init, apply) pair we return the wrapped funs
-  def __iter__(self):
-    return iter((self._init_fun, self._apply_fun))
-
-
-# Convenience methods for common use-case of input variable capture and reuse.
-Var = lambda: Bind(stax.Identity)  # pylint: disable=invalid-name,
-Vars = lambda num_vars: tuple(Bind(stax.Identity) for _ in range(num_vars))  # pylint: disable=invalid-name,
-
-
-def make_apply_fun(bound_layer):
-  """Returns an apply function partially applied to bound params.
-
-  Requires that the top-level model apply_fun be fed params with
-  concrete values for these bound params to be numerically meaningful!
-  (e.g. not JaxprTrace arrays from a JAX JIT pass!)
-
-  Args:
-    bound_layer: Share/Bind/Lambda-bound staxlayer
-
-  Returns:
-    An apply function for this subgraph.
-  """
-  if not isinstance(bound_layer, (Share, Bind)):
-    raise ValueError('Can only create apply function from bound layer.')
-  def partial_apply_fun(inputs, **kwargs):
-    return bound_layer._orig_apply_fun(  # pylint: disable=protected-access
-        bound_layer.params, inputs, **kwargs)
-  return partial_apply_fun
-
-
-# Lambda
-# ------------------------------------------------------------------------------
-# The below provide a nicer syntax for 'pointy' function definition than using
-# raw bound variables.
-class LambdaBind(Bind):
-  """Layer/variable caching function to allow name binding for Lambda layers.
-
-  Args:
-    A staxlayer: an (init_fun, apply_fun) pair.
-
-  Returns:
-    A 'bound' staxlayer that can be assigned to a python variable.
-  Wherever this value is needed elsewhere in the stax tree, call this bound
-  variable and all occurrences will share output values.  Overloads __call__
-  to provide syntactic sugar for Lambda-like invocation.
-  """
-
-  # Syntactic sugar for applying this Lambda to other staxlayers
-  # NB: we do not bind the result by default here!
-  def __call__(self, *args):
-    if len(args) > 1:
-      return stax.serial(stax.parallel(*args), self)
-    elif len(args) == 1:
-      return stax.serial(args[0], self)
-    else:
-      return self
-
-
-class _PlaceholderTree(tuple):
-  """Placeholder tree object for 'initializing' combinators inside Lambdas.
-
-  When we create a Lambda, we're cutting off normal Stax data flow into
-  the subgraph that Lambda wraps with its bound inputs.  This is a
-  problem for any (potentially nested) parallel/serial combinators that
-  are input-facing, as they'll try to unpack the input_shape, inputs, and
-  params trees to feed their sub-layers.  We can't easily know what series
-  of nested access patterns are in a function, so we instead provide
-  recursive placeholder trees to placate the combinators. These placeholders
-  should feed into Lambda input nodes that completely ignore their inputs
-  anyway, but they'll break immediately if the user tries to use unbound
-  inputs from the Stax chain, which is a useful way to force the semantics
-  of Lambda.  This is aggressively tested for correctness in our unit tests.
-  """
-
-  def __init__(self):  # pylint: disable=super-init-not-called
-    self.shape = 0
-    # set generous safety limits for placeholder tree recursion and traversal
-    self.iterator_limit = 1000
-    self.recursion_limit = 30
-
-  def __getitem__(self, _):
-    if self.recursion_limit > 0:
-      self.recursion_limit -= 1
-      return self
-    else:
-      raise IndexError('_PlaceholderTree reached maximum depth')
-
-  def __iter__(self):
-    return self
-
-  def __next__(self):  # PY3
-    return self.next()
-
-  def next(self):  # PY2
-    if self.iterator_limit > 0:
-      self.iterator_limit -= 1
-      return self
-    else:
-      raise StopIteration
-# Register this class with tree-walker to be ignored by optimizers' init fns.
-_register_pytree_node(_PlaceholderTree,
-                      lambda xs: (tuple(), None),
-                      lambda _, xs: _PlaceholderTree())
-
-
-def _PlaceholderInputs():  # pylint: disable=invalid-name
-  """Feeds placeholders into input combinators of a Lambda-bound staxlayer."""
-  init_fun = lambda _, shape: iter((_PlaceholderTree(), _PlaceholderTree()))
-  apply_fun = lambda params, inputs, **kwargs: _PlaceholderTree()
-  return init_fun, apply_fun
-_PlaceholderInputs = _PlaceholderInputs()  # pylint: disable=invalid-name
-
-
-def Lambda(fn):  # pylint: disable=invalid-name
-  """Turn a normal function into a bound, callable Stax layer.
-
-  Args:
-    fn: a python function with _named_ args (i.e. no *args) and no kwargs.
-
-  Returns:
-    A callable, 'bound' staxlayer that can be assigned to a python variable and
-    called like a function with other staxlayers as arguments.  Like Bind,
-    wherever this value is placed in the stax tree, it will always output the
-    same cached value.
-  """
-  # fn's args are just symbolic names that we fill with Vars.
-  num_args = len(inspect.getargspec(fn).args)
-  if num_args > 1:
-    bound_args = Vars(num_args)
-    return LambdaBind(stax.serial(
-        stax.parallel(*bound_args),  # capture inputs
-        _PlaceholderInputs,  # placeholders for input combinators inside fn
-        fn(*bound_args)  # feed captured inputs into fn's args
-    ))
-  elif num_args == 1:
-    bound_arg = Var()
-    return LambdaBind(stax.serial(
-        bound_arg,  # capture input
-        _PlaceholderInputs,  # placeholders for input combinators inside fn
-        fn(bound_arg)  # feed captured inputs into fn's args
-    ))
-  # LambdaBind when no args are given:
-  else:
-    return LambdaBind(fn())
diff --git a/tensor2tensor/trax/stax/slax_test.py b/tensor2tensor/trax/stax/slax_test.py
deleted file mode 100644
index fa258e3ba..000000000
--- a/tensor2tensor/trax/stax/slax_test.py
+++ /dev/null
@@ -1,250 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 Stax Extensions."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import itertools
-from absl.testing import absltest
-from jax import random
-import numpy as onp
-import tensor2tensor.trax.stax as stax
-
-
-def random_inputs(rng, input_shape):
-  if isinstance(input_shape, tuple):
-    return rng.randn(*input_shape).astype(onp.float32)
-  elif isinstance(input_shape, list):
-    return [random_inputs(rng, shape) for shape in input_shape]
-  else:
-    raise TypeError(type(input_shape))
-
-
-def check_shape_agreement(test_case, init_fun, apply_fun, input_shape):
-  rng_key1, rng_key2 = random.split(random.PRNGKey(0))
-  result_shape, params = init_fun(rng_key1, input_shape)
-  inputs = random_inputs(onp.random.RandomState(0), input_shape)
-  result = apply_fun(params, inputs, rng=rng_key2)
-  test_case.assertEqual(result.shape, result_shape)
-  return result_shape
-
-
-def check_staxlayer(test_case, staxlayer, input_shape):
-  init_fun, apply_fun = staxlayer
-  return check_shape_agreement(test_case, init_fun, apply_fun, input_shape)
-
-
-# Helper functions for testing Lambda wrapper against functions involving
-# complicated input trees:
-def _enumerate_trees_w_leaves(n_leaves):
-  """Construct all rooted trees with n leaves."""
-  def enumtree(*args):
-    n_args = len(args)
-    # trivial cases:
-    if n_args == 0:
-      return []
-    if n_args == 1:
-      return args
-    # general case of 2 or more args:
-    # build index array
-    idxs = range(0, n_args)
-    trees = []
-    # we consider all possible subsets of size n_set to gather
-    for n_set in range(2, n_args+1):
-      idxsets = list(itertools.combinations(idxs, n_set))
-      for idxset in idxsets:
-        # recurse by joining all subtrees with
-        # n_set leaves and (n_args - n_set) leaves
-        arg_set = tuple(args[i] for i in idxs if i in idxset)
-        arg_coset = tuple(args[i] for i in idxs if i not in idxset)
-        if arg_coset:
-          trees.extend(tuple(itertools.product(enumtree(*arg_set),
-                                               enumtree(*arg_coset))))
-        else:
-          # trivial case where arg_set is entire set
-          trees.append(arg_set)
-    return trees
-  # return enumerated trees with integers as leaves
-  return enumtree(*range(n_leaves))
-
-
-def _build_combinator_tree(input_treespec, in_vars):
-  """Build a trivial Staxlayer that takes a complicated tree of inputs."""
-  parallel_args = []
-  for e in input_treespec:
-    if isinstance(e, int):
-      parallel_args.append(in_vars[e])
-    elif isinstance(e, tuple):
-      parallel_args.append(_build_combinator_tree(e, in_vars))
-  return stax.serial(stax.parallel(*parallel_args), stax.FanInSum)
-
-
-class SlaxTest(absltest.TestCase):
-
-  def test_flatten_n(self):
-    input_shape = (29, 87, 10, 20, 30)
-
-    actual_shape = check_staxlayer(self, stax.Flatten(1), input_shape)
-    self.assertEqual(actual_shape, (29, 87 * 10 * 20 * 30))
-
-    actual_shape = check_staxlayer(self, stax.Flatten(2), input_shape)
-    self.assertEqual(actual_shape, (29, 87, 10 * 20 * 30))
-
-    actual_shape = check_staxlayer(self, stax.Flatten(3), input_shape)
-    self.assertEqual(actual_shape, (29, 87, 10, 20 * 30))
-
-    actual_shape = check_staxlayer(self, stax.Flatten(4), input_shape)
-    self.assertEqual(actual_shape, (29, 87, 10, 20, 30))
-
-    # Not enough dimensions.
-    with self.assertRaises(ValueError):
-      check_staxlayer(self, stax.Flatten(5), input_shape)
-
-    with self.assertRaises(ValueError):
-      check_staxlayer(self, stax.Flatten(6), input_shape)
-
-  def test_div(self):
-    init_fun, apply_fun = stax.Div(2)
-    input_np = onp.array([[1, 2, 3], [4, 5, 6]], dtype=onp.float32)
-    input_shape = input_np.shape
-    _, _ = init_fun(None, input_shape)
-    output_np = apply_fun(None, input_np)
-    # absltest doesn't have ndarray equalities.
-    expected_output_np = input_np / 2.0
-    self.assertAlmostEqual(
-        0.0,
-        onp.sum((output_np - expected_output_np) ** 2),
-        delta=1e-6)
-
-  # Lambdas replace the staxlayer input stream with a placeholder that
-  # _should_ break any use of unbound variables in the input stream.
-  def testLambda_forbidden_access(self):
-    with self.assertRaises(ValueError):
-      for tree_spec in _enumerate_trees_w_leaves(2):
-        @stax.Lambda
-        def lambda_fun(x, y):  # pylint: disable=unused-argument
-          return _build_combinator_tree(tree_spec,  # pylint: disable=cell-var-from-loop
-                                        # try to read from input stream
-                                        # rather than bound vars
-                                        (x, stax.Identity))
-        check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*2)
-
-  # Exhaustively test the tricky part of Lambda - input combinator
-  # "initialization" for all 2412 trees of stax serial and parallel
-  # combinators of up to six variables.  This probably covers most
-  # practical use patterns!
-
-  # The variables in for loops below are used immediately, disable lint warning
-  # for this section:
-  # pylint: disable=cell-var-from-loop
-  def testLambda_1_arg(self):
-    @stax.Lambda
-    def lambda_fun(x):
-      return _build_combinator_tree((0,), (x,))
-    check_staxlayer(self, lambda_fun, (1, 5, 7, 11))
-
-  def testLambda_2_args(self):
-    for tree_spec in _enumerate_trees_w_leaves(2):
-      @stax.Lambda
-      def lambda_fun(x, y):
-        return _build_combinator_tree(tree_spec, (x, y))
-      check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*2)
-
-  def testLambda_3_args(self):
-    for tree_spec in _enumerate_trees_w_leaves(3):
-      @stax.Lambda
-      def lambda_fun(x, y, z):
-        return _build_combinator_tree(tree_spec, (x, y, z))
-      check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*3)
-
-  def testLambda_4_args(self):
-    for tree_spec in _enumerate_trees_w_leaves(4):
-      @stax.Lambda
-      def lambda_fun(x, y, z, w):
-        return _build_combinator_tree(tree_spec, (x, y, z, w))
-      check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*4)
-
-  def testLambda_5_args(self):
-    for tree_spec in _enumerate_trees_w_leaves(5):
-      @stax.Lambda
-      def lambda_fun(x, y, z, w, v):
-        return _build_combinator_tree(tree_spec, (x, y, z, w, v))
-      check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*5)
-
-  # TODO(mattjj,levskaya): timing out, re-enable with longer timeout?
-  def DISABLED_testLambda_6_args(self):  # pylint: disable=invalid-name
-    for tree_spec in _enumerate_trees_w_leaves(6):
-      @stax.Lambda
-      def lambda_fun(x, y, z, w, v, u):
-        return _build_combinator_tree(tree_spec, (x, y, z, w, v, u))
-      check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*6)
-
-  # Test a few other cases, unused variables, non-input-tree use of
-  # bound Lambda input variables.
-  def testLambda_4_args_only_3_used(self):
-    for tree_spec in _enumerate_trees_w_leaves(3):
-      @stax.Lambda
-      def lambda_fun(x, y, z, w):  # pylint: disable=unused-argument
-        return _build_combinator_tree(tree_spec, (x, y, z))
-      check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*4)
-
-  def testLambda_4_args_only_2_used(self):
-    for tree_spec in _enumerate_trees_w_leaves(2):
-      @stax.Lambda
-      def lambda_fun(x, y, z, w):  # pylint: disable=unused-argument
-        return _build_combinator_tree(tree_spec, (x, y))
-      check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*4)
-
-  def testLambda_4_args_only_1_used(self):
-    @stax.Lambda
-    def lambda_fun(x, y, z, w):  # pylint: disable=unused-argument
-      return _build_combinator_tree((0,), (x,))
-    check_staxlayer(self, lambda_fun, [(1, 5, 7, 11),]*4)
-
-  def testLambda_5_args_2_post_input_tree(self):
-    for tree_spec in _enumerate_trees_w_leaves(3):
-      @stax.Lambda
-      def lambda_fun1(x, y, z, w, v):
-        input_tree = _build_combinator_tree(tree_spec, (x, y, z))
-        return stax.serial(input_tree,
-                           stax.FanOut(3),
-                           stax.parallel(stax.Identity, w, v),
-                           stax.FanInSum)
-      check_staxlayer(self, lambda_fun1, [(1, 5, 7, 11),]*5)
-
-      @stax.Lambda
-      def lambda_fun2(x, y, z, w, v):
-        input_tree = _build_combinator_tree(tree_spec, (x, y, z))
-        return stax.serial(input_tree,
-                           stax.FanOut(3),
-                           stax.parallel(w, stax.Identity, v),
-                           stax.FanInSum)
-      check_staxlayer(self, lambda_fun2, [(1, 5, 7, 11),]*5)
-
-      @stax.Lambda
-      def lambda_fun3(x, y, z, w, v):
-        input_tree = _build_combinator_tree(tree_spec, (x, y, z))
-        return stax.serial(input_tree,
-                           stax.FanOut(3),
-                           stax.parallel(w, v, stax.Identity),
-                           stax.FanInSum)
-      check_staxlayer(self, lambda_fun3, [(1, 5, 7, 11),]*5)
-  # pylint: enable=cell-var-from-loop
-
-
-if __name__ == "__main__":
-  absltest.main()
diff --git a/tensor2tensor/trax/stax/stax_base.py b/tensor2tensor/trax/stax/stax_base.py
index 8a1ece9f8..d9e9e4c6e 100644
--- a/tensor2tensor/trax/stax/stax_base.py
+++ b/tensor2tensor/trax/stax/stax_base.py
@@ -368,3 +368,51 @@ def init_fun(rng, input_shape):
   def apply_fun(params, inputs, **kwargs):
     return make_layer(inputs.shape)[1](params, inputs, **kwargs)
   return init_fun, apply_fun
+
+
+# Utility functions
+# ------------------------------------------------------------------------------
+def one_hot(x, size, dtype=np.float32):
+  """Make a n+1 dim one-hot array from n dim int-categorical array."""
+  return np.array(x[..., np.newaxis] == np.arange(size), dtype)
+
+
+def ShiftRight():  # pylint: disable=invalid-name
+  """Layer to shift the tensor to the right by padding on axis 1."""
+  init_fun = lambda _, input_shape: (input_shape, ())
+  def apply_fun(params, inputs, **kwargs):
+    del params, kwargs
+    pad_widths = [(0, 0), (1, 0)]
+    pad_widths += [(0, 0) for _ in range(len(inputs.shape) - 2)]
+    padded = np.pad(inputs, pad_widths, mode='constant')
+    return padded[:, :-1, ...]
+  return init_fun, apply_fun
+
+
+# Utility Combinators
+# ------------------------------------------------------------------------------
+def repeat(layer, num_repeats):
+  """Repeats layers serially num_repeats times."""
+  if num_repeats < 1:
+    raise ValueError('Repeat combinator num_repeats must be >= 1.')
+  layers = num_repeats * (layer,)
+  return serial(*layers)
+
+
+def residual(*layers, **kwargs):
+  """Constructs a residual version of layers, summing input to layers output."""
+  res = kwargs.get('res', Identity)
+  if len(layers) > 1:
+    return serial(
+        FanOut(2),
+        parallel(serial(*layers), res),
+        FanInSum
+    )
+  elif len(layers) == 1:
+    return serial(
+        FanOut(2),
+        parallel(layers[0], res),
+        FanInSum
+    )
+  else:
+    raise ValueError('Empty residual combinator.')

From 40ef465feb10d55198c1555f75722c54bf66a57d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 19 Apr 2019 16:02:31 -0700
Subject: [PATCH 1925/2720] Extract out dumping params into a function. Then it
 can be used from pdb interactively and in other places if need be.

PiperOrigin-RevId: 244438577
---
 tensor2tensor/trax/rlax/ppo.py | 27 ++++++++++++++++-----------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index fdecb4f65..45f3c0cfe 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -137,6 +137,21 @@ def optimizer_fun(net_params, step_size=1e-3):
   return opt_state, opt_update
 
 
+def log_params(params, name="params"):
+  """Dumps the params with `logging.error`."""
+  for i, param in enumerate(params):
+    if not param:
+      # Empty tuple.
+      continue
+    if not isinstance(param, tuple):
+      logging.error(
+          "%s[%d] : (%s) = [%s]", name, i, param.shape, onp.array(param))
+    else:
+      for j, p in enumerate(param):
+        logging.error(
+            "\t%s[%d, %d] : (%s) = [%s]", name, i, j, p.shape, onp.array(p))
+
+
 # Should this be collect 'n' trajectories, or
 # Run the env for 'n' steps and take completed trajectories, or
 # Any other option?
@@ -202,17 +217,7 @@ def collect_trajectories(env,
         logging.error("predictions: [%s]", predictions)
         logging.error("observation_history: [%s]", observation_history)
         logging.error("policy_net_params: [%s]", policy_net_params)
-        for i, param in enumerate(policy_net_params):
-          if not param:
-            # Empty tuple.
-            continue
-          if not isinstance(param, tuple):
-            logging.error(
-                "Param[%d] : (%s) = [%s]", i, param.shape, onp.array(param))
-          else:
-            for j, p in enumerate(param):
-              logging.error(
-                  "\tParam[%d, %d] : (%s) = [%s]", i, j, p.shape, onp.array(p))
+        log_params(policy_net_params, "policy_net_params")
         raise err
 
       observation, reward, done, _ = env.step(action)

From e76f6fafef73d90862da082cac10bcbf1a71ca21 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 19 Apr 2019 16:08:14 -0700
Subject: [PATCH 1926/2720] Add flag to debug nans.

PiperOrigin-RevId: 244439471
---
 tensor2tensor/trax/rlax/ppo_main.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index a5b9d31af..32a2e68b1 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -23,7 +23,7 @@
 
 from absl import app
 from absl import flags
-from absl import logging
+from jax.config import config
 from tensor2tensor.trax import stax
 from tensor2tensor.trax.rlax import ppo
 
@@ -33,12 +33,13 @@
 flags.DEFINE_string("t2t_gym_env", None, "Name of the T2TGymEnv to make.")
 flags.DEFINE_integer("epochs", 100, "Number of epochs to run for.")
 flags.DEFINE_integer("random_seed", 0, "Random seed.")
-flags.DEFINE_integer("log_level", logging.INFO, "Log level.")
 flags.DEFINE_integer("batch_size", 32, "Batch of trajectories needed.")
 flags.DEFINE_integer("num_optimizer_steps", 100, "Number of optimizer steps.")
 flags.DEFINE_integer("boundary", 20,
                      "We pad trajectories at integer multiples of this number.")
 flags.DEFINE_float("learning_rate", 1e-3, "Learning rate.")
+flags.DEFINE_boolean("jax_debug_nans", False,
+                     "Setting to true will help to debug nans.")
 
 
 def common_stax_layers():
@@ -47,7 +48,10 @@ def common_stax_layers():
 
 def main(argv):
   del argv
-  logging.set_verbosity(FLAGS.log_level)
+
+  if FLAGS.jax_debug_nans:
+    config.update("jax_debug_nans", True)
+
   bottom_layers = common_stax_layers()
 
   if FLAGS.env_name == "Pong-v0":

From e56072d9cf2f2fc97399208c5116372cb4a425af Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 19 Apr 2019 16:27:13 -0700
Subject: [PATCH 1927/2720] Add imagenet64 sequence generation task

PiperOrigin-RevId: 244441999
---
 tensor2tensor/data_generators/imagenet.py | 26 +++++++++++++++++++++++
 tensor2tensor/models/transformer.py       | 23 ++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 9d724fe6e..5e1b3d358 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -23,6 +23,7 @@
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import image_utils
+from tensor2tensor.data_generators import problem
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
@@ -330,6 +331,31 @@ def preprocess_example(self, example, mode, hparams):
     return example
 
 
+@registry.register_problem
+class ImageImagenet64GenFlat(ImageImagenet64Gen):
+  """Imagenet 64 from the pixen cnn paper, as a flat array."""
+
+  def dataset_filename(self):
+    return "image_imagenet64_gen"  # Reuse data.
+
+  def preprocess_example(self, example, mode, unused_hparams):
+    example["inputs"].set_shape(
+        [_IMAGENET_MEDIUM_IMAGE_SIZE, _IMAGENET_MEDIUM_IMAGE_SIZE, 3])
+    example["inputs"] = tf.to_int64(example["inputs"])
+    example["inputs"] = tf.reshape(example["inputs"], (-1,))
+
+    del example["targets"]  # Ensure unconditional generation
+
+    return example
+
+  def hparams(self, defaults, model_hparams):
+    super(ImageImagenet64GenFlat, self).hparams(defaults, model_hparams)
+    # Switch to symbol modality
+    p = defaults
+    p.modality["inputs"] = modalities.ModalityType.SYMBOL_WEIGHTS_ALL
+    p.input_space_id = problem.SpaceID.GENERIC
+
+
 @registry.register_problem
 class ImageImagenet32Small(ImageImagenet):
   """Imagenet small from the pixel cnn paper."""
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 40ecc977f..e3e6317a0 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -2698,3 +2698,26 @@ def transformer_cifar10_memory_v0():
 
   return hparams
 
+
+@registry.register_hparams
+def transformer_imagenet64_memory_v0():
+  """HParams for training image_imagenet64_gen_flat_rev with memory."""
+  hparams = transformer_cifar10_memory_v0()
+
+  hparams.max_length = 64 * 64 * 3
+  hparams.split_targets_chunk_length = 64 * 3
+  hparams.split_targets_max_chunks = int(
+      hparams.max_length / hparams.split_targets_chunk_length)
+  hparams.num_memory_items = 128 * 3
+
+  # Since this is an image problem, batch size refers to examples (not tokens)
+  target_images_per_batch = 2
+  hparams.batch_size = int(target_images_per_batch * (
+      hparams.max_length / hparams.split_targets_chunk_length))
+
+  # The recurrent memory needs to know the actual batch size (in sequences)
+  hparams.recurrent_memory_batch_size = hparams.batch_size
+
+  hparams.max_relative_position = 3072
+
+  return hparams

From e6d2a3e4f2cd6e51dffa229d9c3fdcabf31bdfea Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 22 Apr 2019 01:43:08 -0700
Subject: [PATCH 1928/2720] TRAX layer refactor, make our layers classes.

PiperOrigin-RevId: 244631294
---
 tensor2tensor/trax/models/mlp.py              |   8 +-
 tensor2tensor/trax/models/resnet.py           |  93 ++-
 tensor2tensor/trax/models/transformer.py      |  97 ++-
 tensor2tensor/trax/rlax/ppo.py                |  15 +-
 tensor2tensor/trax/rlax/ppo_main.py           |  14 +-
 tensor2tensor/trax/rlax/ppo_test.py           |  28 +-
 .../trax/rlax/ppo_training_loop_test.py       |   8 +-
 tensor2tensor/trax/stax/README.md             | 134 +---
 tensor2tensor/trax/stax/__init__.py           |   1 +
 tensor2tensor/trax/stax/attention.py          | 228 +++----
 tensor2tensor/trax/stax/base.py               | 248 +++++--
 tensor2tensor/trax/stax/base_test.py          |  86 +--
 tensor2tensor/trax/stax/combinators.py        | 154 +++++
 tensor2tensor/trax/stax/stax_base.py          | 609 ++++++++----------
 14 files changed, 916 insertions(+), 807 deletions(-)
 create mode 100644 tensor2tensor/trax/stax/combinators.py

diff --git a/tensor2tensor/trax/models/mlp.py b/tensor2tensor/trax/models/mlp.py
index 7618c6612..ab39d6a11 100644
--- a/tensor2tensor/trax/models/mlp.py
+++ b/tensor2tensor/trax/models/mlp.py
@@ -27,8 +27,10 @@ def MLP(num_hidden_layers=2,
         activation_fn=stax.Relu,
         num_output_classes=10,
         mode="train"):
+  """Multi-layer feed-forward neural network with non-linear activations."""
   del mode
   layers = [stax.Flatten()]
-  layers += [stax.Dense(hidden_size), activation_fn] * num_hidden_layers
-  layers += [stax.Dense(num_output_classes), stax.LogSoftmax]
-  return stax.serial(*layers)
+  for _ in range(num_hidden_layers):
+    layers += [stax.Dense(hidden_size), activation_fn()]
+  layers += [stax.Dense(num_output_classes), stax.LogSoftmax()]
+  return stax.Serial(*layers)
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index 0c81dee0e..1c88108d3 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -26,38 +26,35 @@ def ConvBlock(kernel_size, filters, strides):
   """ResNet convolutional striding block."""
   ks = kernel_size
   filters1, filters2, filters3 = filters
-  main = stax.serial(
+  main = stax.Serial(
       stax.Conv(filters1, (1, 1), strides),
-      stax.BatchNorm(), stax.Relu,
+      stax.BatchNorm(), stax.Relu(),
       stax.Conv(filters2, (ks, ks), padding='SAME'),
-      stax.BatchNorm(), stax.Relu,
+      stax.BatchNorm(), stax.Relu(),
       stax.Conv(filters3, (1, 1)), stax.BatchNorm())
-  shortcut = stax.serial(
+  shortcut = stax.Serial(
       stax.Conv(filters3, (1, 1), strides),
       stax.BatchNorm())
-  return stax.serial(
-      stax.FanOut(2),
-      stax.parallel(main, shortcut),
-      stax.FanInSum, stax.Relu)
+  return stax.Serial(
+      stax.FanOut(),
+      stax.Parallel(main, shortcut),
+      stax.FanInSum(), stax.Relu())
 
 
 def IdentityBlock(kernel_size, filters):
   """ResNet identical size block."""
   ks = kernel_size
-  filters1, filters2 = filters
-  def MakeMain(input_shape):
-    # the number of output channels depends on the number of input channels
-    return stax.serial(
-        stax.Conv(filters1, (1, 1)),
-        stax.BatchNorm(), stax.Relu,
-        stax.Conv(filters2, (ks, ks), padding='SAME'),
-        stax.BatchNorm(), stax.Relu,
-        stax.Conv(input_shape[3], (1, 1)), stax.BatchNorm())
-  main = stax.shape_dependent(MakeMain)
-  return stax.serial(
-      stax.FanOut(2),
-      stax.parallel(main, stax.Identity),
-      stax.FanInSum, stax.Relu)
+  filters1, filters2, filters3 = filters
+  main = stax.Serial(
+      stax.Conv(filters1, (1, 1)),
+      stax.BatchNorm(), stax.Relu(),
+      stax.Conv(filters2, (ks, ks), padding='SAME'),
+      stax.BatchNorm(), stax.Relu(),
+      stax.Conv(filters3, (1, 1)), stax.BatchNorm())
+  return stax.Serial(
+      stax.FanOut(),
+      stax.Parallel(main, stax.Identity()),
+      stax.FanInSum(), stax.Relu())
 
 
 def Resnet50(hidden_size=64, num_output_classes=1001, mode='train'):
@@ -72,40 +69,40 @@ def Resnet50(hidden_size=64, num_output_classes=1001, mode='train'):
     The ResNet model with the given layer and output sizes.
   """
   del mode
-  return stax.serial(
+  return stax.Serial(
       stax.Conv(hidden_size, (7, 7), (2, 2), 'SAME'),
-      stax.BatchNorm(), stax.Relu,
-      stax.MaxPool((3, 3), strides=(2, 2)),
+      stax.BatchNorm(), stax.Relu(),
+      stax.MaxPool(pool_size=(3, 3), strides=(2, 2)),
       ConvBlock(3, [hidden_size, hidden_size, 4 * hidden_size], (1, 1)),
-      IdentityBlock(3, [hidden_size, hidden_size]),
-      IdentityBlock(3, [hidden_size, hidden_size]),
+      IdentityBlock(3, [hidden_size, hidden_size, 4 * hidden_size]),
+      IdentityBlock(3, [hidden_size, hidden_size, 4 * hidden_size]),
       ConvBlock(3, [2 * hidden_size, 2 * hidden_size, 8 * hidden_size], (2, 2)),
-      IdentityBlock(3, [2 * hidden_size, 2 * hidden_size]),
-      IdentityBlock(3, [2 * hidden_size, 2 * hidden_size]),
-      IdentityBlock(3, [2 * hidden_size, 2 * hidden_size]),
+      IdentityBlock(3, [2 * hidden_size, 2 * hidden_size, 8 * hidden_size]),
+      IdentityBlock(3, [2 * hidden_size, 2 * hidden_size, 8 * hidden_size]),
+      IdentityBlock(3, [2 * hidden_size, 2 * hidden_size, 8 * hidden_size]),
       ConvBlock(3, [4 * hidden_size, 4 * hidden_size, 16*hidden_size], (2, 2)),
-      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size]),
-      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size]),
-      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size]),
-      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size]),
-      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size]),
+      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size, 16 * hidden_size]),
+      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size, 16 * hidden_size]),
+      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size, 16 * hidden_size]),
+      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size, 16 * hidden_size]),
+      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size, 16 * hidden_size]),
       ConvBlock(3, [8 * hidden_size, 8 * hidden_size, 32*hidden_size], (2, 2)),
-      IdentityBlock(3, [8 * hidden_size, 8 * hidden_size]),
-      IdentityBlock(3, [8 * hidden_size, 8 * hidden_size]),
-      stax.AvgPool((7, 7)), stax.Flatten(),
-      stax.Dense(num_output_classes), stax.LogSoftmax)
+      IdentityBlock(3, [8 * hidden_size, 8 * hidden_size, 32 * hidden_size]),
+      IdentityBlock(3, [8 * hidden_size, 8 * hidden_size, 32 * hidden_size]),
+      stax.AvgPool(pool_size=(7, 7)), stax.Flatten(),
+      stax.Dense(num_output_classes), stax.LogSoftmax())
 
 
 def WideResnetBlock(channels, strides=(1, 1), channel_mismatch=False):
   """WideResnet convolutational block."""
-  main = stax.serial(stax.BatchNorm(), stax.Relu,
+  main = stax.Serial(stax.BatchNorm(), stax.Relu(),
                      stax.Conv(channels, (3, 3), strides, padding='SAME'),
-                     stax.BatchNorm(), stax.Relu,
+                     stax.BatchNorm(), stax.Relu(),
                      stax.Conv(channels, (3, 3), padding='SAME'))
-  shortcut = stax.Identity if not channel_mismatch else stax.Conv(
+  shortcut = stax.Identity() if not channel_mismatch else stax.Conv(
       channels, (3, 3), strides, padding='SAME')
-  return stax.serial(
-      stax.FanOut(2), stax.parallel(main, shortcut), stax.FanInSum)
+  return stax.Serial(
+      stax.FanOut(), stax.Parallel(main, shortcut), stax.FanInSum())
 
 
 def WideResnetGroup(n, channels, strides=(1, 1)):
@@ -113,7 +110,7 @@ def WideResnetGroup(n, channels, strides=(1, 1)):
   blocks += [WideResnetBlock(channels, strides, channel_mismatch=True)]
   for _ in range(n - 1):
     blocks += [WideResnetBlock(channels, (1, 1))]
-  return stax.serial(*blocks)
+  return stax.Serial(*blocks)
 
 
 def WideResnet(num_blocks=3, hidden_size=64, num_output_classes=10,
@@ -130,10 +127,10 @@ def WideResnet(num_blocks=3, hidden_size=64, num_output_classes=10,
     The WideResnet model with given layer and output sizes.
   """
   del mode
-  return stax.serial(
+  return stax.Serial(
       stax.Conv(hidden_size, (3, 3), padding='SAME'),
       WideResnetGroup(num_blocks, hidden_size),
       WideResnetGroup(num_blocks, hidden_size * 2, (2, 2)),
       WideResnetGroup(num_blocks, hidden_size * 4, (2, 2)), stax.BatchNorm(),
-      stax.Relu, stax.AvgPool((8, 8)), stax.Flatten(),
-      stax.Dense(num_output_classes), stax.LogSoftmax)
+      stax.Relu(), stax.AvgPool(pool_size=(8, 8)), stax.Flatten(),
+      stax.Dense(num_output_classes), stax.LogSoftmax())
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 0a4677d26..eee2ac4dd 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -26,13 +26,13 @@ def ResidualFeedForward(feature_depth,
                         dropout,
                         mode):
   """Residual feed-forward layer with normalization at start."""
-  return stax.residual(
+  return stax.Residual(
       stax.LayerNorm(),
       stax.Dense(feedforward_depth, W_init=stax.xavier_uniform()),
-      stax.Relu,
-      stax.Dropout(dropout, mode=mode),
+      stax.Relu(),
+      stax.Dropout(rate=dropout, mode=mode),
       stax.Dense(feature_depth, W_init=stax.xavier_uniform()),
-      stax.Dropout(dropout, mode=mode)
+      stax.Dropout(rate=dropout, mode=mode)
   )
 
 
@@ -72,21 +72,21 @@ def Encoder(embedded_source, source_mask):
     Returns:
       Staxlayer variable that outputs encoded source.
     """
-    encoder_layer = stax.serial(
+    encoder_layer = stax.Serial(
         # input attends to self
-        stax.residual(stax.LayerNorm(),
-                      stax.FanOut(4),
-                      stax.parallel(stax.Identity,  # query
-                                    stax.Identity,  # key
-                                    stax.Identity,  # value
+        stax.Residual(stax.LayerNorm(),
+                      stax.FanOut(size=4),
+                      stax.Parallel(stax.Identity(),  # query
+                                    stax.Identity(),  # key
+                                    stax.Identity(),  # value
                                     source_mask),  # attention mask
                       multi_attention,
-                      stax.Dropout(dropout, mode=mode)),
+                      stax.Dropout(rate=dropout, mode=mode)),
         # feed-forward
         ResidualFeedForward(
             feature_depth, feedforward_depth, dropout, mode=mode)
     )
-    return stax.serial(
+    return stax.Serial(
         embedded_source,
         stax.repeat(encoder_layer, num_layers),
         stax.LayerNorm(),
@@ -112,17 +112,17 @@ def DecoderLayer(feature_depth,
   Returns:
     init and apply.
   """
-  return stax.serial(
-      stax.residual(  # Self-attention block.
+  return stax.Serial(
+      stax.Residual(  # Self-attention block.
           stax.LayerNorm(),
-          stax.FanOut(4),
-          stax.parallel(stax.Identity,  # query
-                        stax.Identity,  # key
-                        stax.Identity,  # value
+          stax.FanOut(size=4),
+          stax.Parallel(stax.Identity(),  # query
+                        stax.Identity(),  # key
+                        stax.Identity(),  # value
                         stax.CausalMask(axis=-2)),  # attention mask
           stax.MultiHeadedAttention(feature_depth, num_heads=num_heads,
                                     dropout=dropout, mode=mode),
-          stax.Dropout(dropout, mode=mode)
+          stax.Dropout(rate=dropout, mode=mode)
       ),
       ResidualFeedForward(feature_depth, feedforward_depth, dropout, mode=mode)
   )
@@ -151,18 +151,17 @@ def TransformerLM(vocab_size,
   Returns:
     init and apply.
   """
-  return stax.serial(
+  return stax.Serial(
       stax.ShiftRight(),
       stax.Embedding(feature_depth, vocab_size),
-      stax.Dropout(dropout, mode=mode),
-      stax.PositionalEncoding(feature_depth, max_len=max_len),
-      stax.repeat(
-          DecoderLayer(
-              feature_depth, feedforward_depth, num_heads, dropout, mode),
-          num_layers),
+      stax.Dropout(rate=dropout, mode=mode),
+      stax.PositionalEncoding(max_len=max_len),
+      stax.Serial([DecoderLayer(feature_depth, feedforward_depth, num_heads,
+                                dropout, mode)
+                   for _ in range(num_layers)]),
       stax.LayerNorm(),
       stax.Dense(vocab_size, W_init=stax.xavier_uniform()),
-      stax.LogSoftmax
+      stax.LogSoftmax()
   )
 
 
@@ -201,7 +200,7 @@ def Transformer(source_vocab_size,
   functions for the trained encoder, decoder, and generator substax.
   """
   # Input embedding and positional encoding
-  inject_position = stax.serial(
+  inject_position = stax.Serial(
       stax.Dropout(dropout, mode=mode),
       stax.PositionalEncoding(feature_depth, max_len=max_len)
   )
@@ -209,13 +208,13 @@ def Transformer(source_vocab_size,
     assert source_vocab_size == target_vocab_size
     # Weight-shared Embedding
     embedding = stax.Share(stax.Embedding(feature_depth, source_vocab_size))
-    source_embedding_layer = stax.serial(embedding, inject_position)
+    source_embedding_layer = stax.Serial(embedding, inject_position)
     target_embedding_layer = source_embedding_layer
   else:
     source_embedding = stax.Embedding(feature_depth, source_vocab_size)
     target_embedding = stax.Embedding(feature_depth, target_vocab_size)
-    source_embedding_layer = stax.serial(source_embedding, inject_position)
-    target_embedding_layer = stax.serial(target_embedding, inject_position)
+    source_embedding_layer = stax.Serial(source_embedding, inject_position)
+    target_embedding_layer = stax.Serial(target_embedding, inject_position)
 
   # Multi-headed Attention and Feed-forward layers
   multi_attention = stax.MultiHeadedAttention(
@@ -233,13 +232,13 @@ def Encoder(source, source_mask):
     Returns:
       Staxlayer variable that outputs encoded source.
     """
-    encoder_layer = stax.serial(
+    encoder_layer = stax.Serial(
         # input attends to self
-        stax.residual(stax.LayerNorm(),
-                      stax.FanOut(4),
-                      stax.parallel(stax.Identity,  # query
-                                    stax.Identity,  # key
-                                    stax.Identity,  # value
+        stax.Residual(stax.LayerNorm(),
+                      stax.FanOut(size=4),
+                      stax.Parallel(stax.Identity(),  # query
+                                    stax.Identity(),  # key
+                                    stax.Identity(),  # value
                                     source_mask),  # attention mask
                       multi_attention,
                       stax.Dropout(dropout, mode=mode)),
@@ -247,7 +246,7 @@ def Encoder(source, source_mask):
         ResidualFeedForward(
             feature_depth, feedforward_depth, dropout, mode=mode),
     )
-    return stax.serial(
+    return stax.Serial(
         source,
         source_embedding_layer,
         stax.repeat(encoder_layer, num_layers),
@@ -268,20 +267,20 @@ def Decoder(memory, target, target_mask, memory_mask):
     Returns:
       Staxlayer variable that outputs encoded source.
     """
-    decoder_layer = stax.serial(
+    decoder_layer = stax.Serial(
         # target attends to self
-        stax.residual(stax.LayerNorm(),
-                      stax.FanOut(4),
-                      stax.parallel(stax.Identity,  # query
-                                    stax.Identity,  # key
-                                    stax.Identity,  # value
+        stax.Residual(stax.LayerNorm(),
+                      stax.FanOut(size=4),
+                      stax.Parallel(stax.Identity(),  # query
+                                    stax.Identity(),  # key
+                                    stax.Identity(),  # value
                                     target_mask),  # attention mask
                       multi_attention,
                       stax.Dropout(dropout, mode=mode)),
         # target attends to encoded source
-        stax.residual(stax.LayerNorm(),
-                      stax.FanOut(4),
-                      stax.parallel(stax.Identity,  # query
+        stax.Residual(stax.LayerNorm(),
+                      stax.FanOut(size=4),
+                      stax.Parallel(stax.Identity(),  # query
                                     memory,  # key
                                     memory,  # value
                                     memory_mask),  # attention mask
@@ -291,7 +290,7 @@ def Decoder(memory, target, target_mask, memory_mask):
         ResidualFeedForward(
             feature_depth, feedforward_depth, dropout, mode=mode)
     )
-    return stax.serial(
+    return stax.Serial(
         target,
         target_embedding_layer,
         stax.repeat(decoder_layer, num_layers),
@@ -307,7 +306,7 @@ def transformer(source, target, source_mask, target_mask, memory_mask):  # pylin
   # Finally, bind the generator transform to use later for inference.
   @stax.Lambda
   def Generator(encoded_target):
-    return stax.serial(
+    return stax.Serial(
         encoded_target,
         stax.Dense(target_vocab_size, W_init=stax.xavier_uniform()),
         stax.LogSoftmax
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 45f3c0cfe..2de65b959 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -57,8 +57,8 @@
 from jax import random as jax_random
 import numpy as onp
 from tensor2tensor.trax import optimizers as trax_opt
+from tensor2tensor.trax import stax
 from tensor2tensor.trax import trax
-from tensor2tensor.trax.stax import stax_base as stax
 
 DEBUG_LOGGING = False
 GAMMA = 0.99
@@ -79,9 +79,9 @@ def policy_net(rng_key,
   # required layers on top of it.
   if bottom_layers is None:
     bottom_layers = []
-  bottom_layers.extend([stax.Dense(num_actions), stax.Softmax])
+  bottom_layers.extend([stax.Dense(num_actions), stax.Softmax()])
 
-  net_init, net_apply = stax.serial(*bottom_layers)
+  net_init, net_apply = stax.Serial(bottom_layers)
 
   _, net_params = net_init(rng_key, batch_observations_shape)
   return net_params, net_apply
@@ -100,7 +100,7 @@ def value_net(rng_key,
       stax.Dense(1),
   ])
 
-  net_init, net_apply = stax.serial(*bottom_layers)
+  net_init, net_apply = stax.Serial(bottom_layers)
 
   _, net_params = net_init(rng_key, batch_observations_shape)
   return net_params, net_apply
@@ -119,12 +119,12 @@ def policy_and_value_net(rng_key,
 
   # Now, with the current logits, one head computes action probabilities and the
   # other computes the value function.
-  layers.extend([stax.FanOut(2), stax.parallel(
-      stax.serial(stax.Dense(num_actions), stax.Softmax),
+  layers.extend([stax.FanOut(), stax.Parallel(
+      stax.Serial(stax.Dense(num_actions), stax.Softmax()),
       stax.Dense(1)
   )])
 
-  net_init, net_apply = stax.serial(*layers)
+  net_init, net_apply = stax.Serial(layers)
 
   _, net_params = net_init(rng_key, batch_observations_shape)
   return net_params, net_apply
@@ -859,4 +859,3 @@ def training_loop(
 
   return ((policy_net_params, value_net_params), average_rewards,
           np.stack(value_losses), np.stack(ppo_objective))
-
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 32a2e68b1..0738e758a 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -43,7 +43,10 @@
 
 
 def common_stax_layers():
-  return [stax.Dense(16), stax.Relu, stax.Dense(4), stax.Relu]
+  layers = []
+  if FLAGS.env_name == "Pong-v0":
+    layers = [stax.Div(divisor=255.0), stax.Flatten(num_axis_to_keep=2)]
+  return layers + [stax.Dense(16), stax.Relu(), stax.Dense(4), stax.Relu()]
 
 
 def main(argv):
@@ -52,11 +55,6 @@ def main(argv):
   if FLAGS.jax_debug_nans:
     config.update("jax_debug_nans", True)
 
-  bottom_layers = common_stax_layers()
-
-  if FLAGS.env_name == "Pong-v0":
-    bottom_layers = [stax.Div(255.0), stax.Flatten(2)] + bottom_layers
-
   optimizer_fun = functools.partial(ppo.optimizer_fun,
                                     step_size=FLAGS.learning_rate)
 
@@ -64,9 +62,9 @@ def main(argv):
       env_name=FLAGS.env_name,
       epochs=FLAGS.epochs,
       policy_net_fun=functools.partial(
-          ppo.policy_net, bottom_layers=bottom_layers),
+          ppo.policy_net, bottom_layers=common_stax_layers()),
       value_net_fun=functools.partial(
-          ppo.value_net, bottom_layers=bottom_layers),
+          ppo.value_net, bottom_layers=common_stax_layers()),
       policy_optimizer_fun=optimizer_fun,
       value_optimizer_fun=optimizer_fun,
       batch_size=FLAGS.batch_size,
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 6a0599a33..8dbff285c 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -43,7 +43,7 @@ def test_policy_net(self):
         num_actions,
         # flatten except batch and time
         # step dimensions.
-        [stax.Flatten(2)])
+        [stax.Flatten(num_axis_to_keep=2)])
 
     # Generate a batch of observations.
     batch = 2
@@ -64,9 +64,10 @@ def test_policy_net(self):
   def test_value_net(self):
     observation_shape = (3, 4, 5)
     num_actions = 2
-    value_params, value_apply = ppo.value_net(self.rng_key,
-                                              (-1, -1) + observation_shape,
-                                              num_actions, [stax.Flatten(2)])
+    value_params, value_apply = ppo.value_net(
+        self.rng_key,
+        (-1, -1) + observation_shape,
+        num_actions, [stax.Flatten(num_axis_to_keep=2)])
     batch = 2
     time_steps = 10
     batch_of_observations = np.random.uniform(
@@ -81,7 +82,8 @@ def test_policy_and_value_net(self):
     batch_observation_shape = (-1, -1) + observation_shape
     num_actions = 2
     pnv_params, pnv_apply = ppo.policy_and_value_net(
-        self.rng_key, batch_observation_shape, num_actions, [stax.Flatten(2)])
+        self.rng_key, batch_observation_shape, num_actions,
+        [stax.Flatten(num_axis_to_keep=2)])
     batch = 2
     time_steps = 10
     batch_of_observations = np.random.uniform(
@@ -102,7 +104,7 @@ def test_collect_trajectories(self):
         num_actions,
         # flatten except batch and time
         # step dimensions.
-        [stax.Flatten(2)])
+        [stax.Flatten(num_axis_to_keep=2)])
 
     # We'll get done at time-step #10, starting from 0, therefore in 11 steps.
     done_time_step = 5
@@ -444,14 +446,16 @@ def test_ppo_loss(self):
     batch_observation_shape = (-1, -1) + OBS
 
     old_policy_params, _ = ppo.policy_net(key1, batch_observation_shape, A,
-                                          [stax.Flatten(2)])
+                                          [stax.Flatten(num_axis_to_keep=2)])
 
-    new_policy_params, policy_apply = ppo.policy_net(key2,
-                                                     batch_observation_shape, A,
-                                                     [stax.Flatten(2)])
+    new_policy_params, policy_apply = ppo.policy_net(
+        key2,
+        batch_observation_shape, A,
+        [stax.Flatten(num_axis_to_keep=2)])
 
-    value_params, value_apply = ppo.value_net(key3, batch_observation_shape, A,
-                                              [stax.Flatten(2)])
+    value_params, value_apply = ppo.value_net(
+        key3, batch_observation_shape, A,
+        [stax.Flatten(num_axis_to_keep=2)])
 
     # Generate a batch of observations.
 
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index 7731f600a..2ac2b5912 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -22,8 +22,8 @@
 import functools
 import gym
 from tensor2tensor.rl import gym_utils
+from tensor2tensor.trax import stax
 from tensor2tensor.trax.rlax import ppo
-from tensor2tensor.trax.stax import stax_base as stax
 from tensorflow import test
 
 
@@ -37,16 +37,14 @@ def test_training_loop(self):
     env = gym.wrappers.TimeLimit(env, max_episode_steps=2)
     num_epochs = 2
     batch_size = 2
-    # Common bottom layer(s).
-    bottom_layers = [stax.Dense(1)]
     # Run the training loop.
     _, rewards, val_losses, ppo_objectives = ppo.training_loop(
         env=env,
         epochs=num_epochs,
         policy_net_fun=functools.partial(
-            ppo.policy_net, bottom_layers=bottom_layers),
+            ppo.policy_net, bottom_layers=[stax.Dense(1)]),
         value_net_fun=functools.partial(
-            ppo.value_net, bottom_layers=bottom_layers),
+            ppo.value_net, bottom_layers=[stax.Dense(1)]),
         batch_size=batch_size,
         num_optimizer_steps=1,
         random_seed=0)
diff --git a/tensor2tensor/trax/stax/README.md b/tensor2tensor/trax/stax/README.md
index d92f8520a..7c96d6e2b 100644
--- a/tensor2tensor/trax/stax/README.md
+++ b/tensor2tensor/trax/stax/README.md
@@ -1,124 +1,48 @@
-# Stax - Layer Extensions
+# Trax Layers
 
-# Convenience layers and combinators
 
-SLAX implements repeat, residual, and multiplex combinators, parallel
-input-tuple sub-selection with Take layer, and graph inputs shape-logging with
-LogInputs layer for debugging.
 
-# Name Binding
+## Base layer structure
 
-SLAX implements Share, Bind, Var, Vars, Lambda, and make_apply_fun.
-These operators augment the point-free Stax API with name-binding operations.
-This provides a concise way of pointifying Stax notation when needed for
-complicated neural net models while retaining a very functional style overall.
-
-### Bind
-
-Layer name-binding. Caches the results of the layer on its first application
-in the computation DAG so that it can be referred to elsewhere in a model
-definition and used as though it were a pointer to the cached variable.
-
-We use the name-bound layer inside the main model just like a normal stax layer:
+All layers inherit form the Layer class and need to implement 3 functions:
 
 ```python
-# bind a layer with Bind:
-encoder = Bind(serial(Dense(10), Relu))
-
-# elsewhere in stax definition:
-model = serial(
-    # ...
-    encoder, # evaluated and cached here
-    # ...
-    encoder, # this always returns the same value
-    #...
-)
-
-# after training, we can access its params:
-encoder.params
-
-# or its last activations:
-encoder.value
-
-# or we can re-evaluate it with its trained set of params:
-eval_time_result = make_apply_fun(encoder)(inputs, **kwargs)
-```
+def call(self, params, inputs, **kwargs):
+"""Call this layer using the given parameters on the given inputs."""
 
-Also note the convenience functions __Var__ and __Vars__(_N_), which are just
-bound Identity layers. This is convenient for capturing input values to be used
-elsewhere in the model. These can be used with __parallel__ and the helper
-__multiplex__ combinators to easily route data around inside a stax model.
+def output_shape(self, input_shape):
+"""The shape of the output given the shape of the input."""
 
-### Share
+def new_parameters(self, input_shape, rng):
+"""Create new parameters given the shape of the input."""
+```
 
-Parameter name-binding, for shared parameters. Just like __Bind__, but __Share__
-doesn't bind the cached _results_ of a layer, but only it's _parameters_. This
-allows us to create a weight-sharing layer by name. This works transparently
-with jax.grad and optimizers as they only ever see one set of real parameters
-from the state tree in the traced computations, so there's no inefficiency
-introduced.
+The base layer class wraps these functions and provides initialization
+and call functions to be used as follows.
 
 ```python
-# bind a layer with Share:
-shared_layer = Share(serial(Dense(10), Relu))
+input = np.zeros(10)
+layer = MyLayer()
+params = layer.initialize()
+output = layer(params, input)
+```
 
-# elsewhere in stax definition:
-tower_A = serial(..., shared_layer, ...)
-tower_B = serial(..., shared_layer, ...)
+## Parameter sharing
 
-# after training, we can access its params:
-shared_layer.params
+Parameters are shared when the same layer object is used.
 
-# or we can re-evaluate it with its trained set of params:
-eval_time_result = make_apply_fun(shared_layer)(inputs, **kwargs)
+```python
+standard_mlp = layers.Serial(layers.Dense(10), layers.Dense(10))
+layer = Dense(10)
+shared_parameters_mlp = layers.Serial(layer, layer)
 ```
 
-### Lambda
+## Core layers
 
-A function wrapper to allow concise function definitions of model layers.
-This uses __Bind__ behind the scenes to fill in the values of the named
-arguments with an input layer that captures the tuple of inputs, finally it
-wraps the output of the function with a special form of __Bind__ that overloads
-the `__call__` operator to make it easy to couple this subgraph to inputs as if
-it were a normal function call.
+* Dense
+* Conv
 
-```python
-# we wrap a normal python function (*args only no **kwargs supported, but they
-# always can be fed in from an outer scope.) e.g.:
-
-some_layer_outside = serial(...)
-@Lambda
-def fun(x, y):
-    tmp = serial(x, serial(Dense(10), Relu)))
-    return serial(parallel(tmp, y), FanInSum, some_layer_outside)
-
-# Later we can simply call the function with staxlayer arguments, even within
-# another Lambda wrapped function:
-result = fun(input1, input2)
-
-# or chain them:
-result = fun(input3, fun(input1, input2))
-
-# Lambda is doing the "spiritual equivalent" to following:
-x, y = Var(), Var()
-Bind(
-  serial(
-    parallel(x, y),
-    fun(x, y)
-  )
-)
-# But Lambda also takes care of some annoying technical issues with
-# combinators behind the scenes to make this work as well using a
-# special Bind that overloads __call__ to make the result act like a
-# function.
-
-# after training, we can access its params:
-fun.params
-
-# or its last activations:
-fun.value
-
-# or we can re-evaluate it with its trained set of params:
-eval_time_result = make_apply_fun(fun)(inputs, **kwargs)
+## Layer composition
 
-```
+* Serial
+* Parallel
diff --git a/tensor2tensor/trax/stax/__init__.py b/tensor2tensor/trax/stax/__init__.py
index 88fd1cd96..263fab640 100644
--- a/tensor2tensor/trax/stax/__init__.py
+++ b/tensor2tensor/trax/stax/__init__.py
@@ -23,4 +23,5 @@
 # pylint: disable=wildcard-import
 from tensor2tensor.trax.stax.attention import *
 from tensor2tensor.trax.stax.base import *
+from tensor2tensor.trax.stax.combinators import *
 from tensor2tensor.trax.stax.stax_base import *
diff --git a/tensor2tensor/trax/stax/attention.py b/tensor2tensor/trax/stax/attention.py
index 0879295b8..4b631a597 100644
--- a/tensor2tensor/trax/stax/attention.py
+++ b/tensor2tensor/trax/stax/attention.py
@@ -20,26 +20,21 @@
 
 import numpy as onp
 
+from tensor2tensor.trax import backend
 from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.backend import random
+from tensor2tensor.trax.stax import base
+from tensor2tensor.trax.stax import combinators
 from tensor2tensor.trax.stax import stax_base as stax
 
 
-def causal_mask(size, dtype=np.uint8):
-  """Causal attention mask."""
-  return onp.tril(onp.ones((1, size, size), dtype=dtype), k=0)
+@base.layer(output_shape=lambda shape, axis=-1: (1, shape[axis], shape[axis]))
+def CausalMask(params, x, axis=-1, **kwargs):
+  del params, kwargs
+  size = x.shape[axis]
+  return onp.tril(onp.ones((1, size, size), dtype=x.dtype), k=0)
 
 
-def CausalMask(axis=-1):  # pylint: disable=invalid-name
-  """Layer to create a causal mask for its inputs."""
-  init_fun = lambda _, input_shape: (input_shape, ())
-  def apply_fun(params, inputs, **kwargs):
-    del params, kwargs
-    return causal_mask(inputs.shape[axis], dtype=inputs.dtype)
-  return init_fun, apply_fun
-
-
-def make_target_mask(target, pad=0):
+def MakeTargetMask(target, pad=0):
   """Create an attention mask to hide padding and future words."""
   target_mask = (target != pad)[ :, np.newaxis, :]
   target_dtype = target_mask.dtype
@@ -48,7 +43,7 @@ def make_target_mask(target, pad=0):
   return np.expand_dims(target_mask, axis=1)
 
 
-def prepare_paired_sequence_batch(source, target_in, pad=0):
+def PreparePairedSequenceBatch(source, target_in, pad=0):
   """Build masks for this batch.
 
   Args:
@@ -64,7 +59,7 @@ def prepare_paired_sequence_batch(source, target_in, pad=0):
   target_y = target_in[:, 1:]
   source_mask = np.reshape(source != pad,
                            (source.shape[0], 1, 1, source.shape[-1]))
-  target_mask = make_target_mask(target, pad)
+  target_mask = MakeTargetMask(target, pad)
   memory_mask = (
       np.reshape(np.arange(target.shape[-1]) < source.shape[-1], [-1, 1]))
   ntokens = np.sum(target_y != pad)
@@ -72,71 +67,47 @@ def prepare_paired_sequence_batch(source, target_in, pad=0):
           source_mask, target_mask, memory_mask, ntokens)
 
 
-def xavier_uniform(out_dim=0, in_dim=1):
-  """An initializer function for random uniform xavier-scaled coefficients."""
-  def init(rng, shape):
-    fan_in, fan_out = shape[in_dim], shape[out_dim]
-    std = np.sqrt(2.0 / (fan_in + fan_out))
-    a = np.sqrt(3.0) * std
-    return random.uniform(rng, shape, minval=-a, maxval=a)
-  return init
-
-
-def LayerNorm(epsilon=1e-6):  # pylint: disable=invalid-name
-  """Layer construction function for Layer Normalization layer.."""
-  def init_fun(_, input_shape):
-    features = input_shape[-1]
-    scale = np.ones(features)
-    bias = np.zeros(features)
-    return input_shape, (scale, bias)
-
-  def apply_fun(params, inputs, **kwargs):
-    del kwargs
-    (scale, bias) = params
-    mean = np.mean(inputs, axis=-1, keepdims=True)
-    variance = np.mean((inputs - mean)**2, axis=-1, keepdims=True)
-    norm_inputs = (inputs - mean) / np.sqrt(variance + epsilon)
-    return norm_inputs * scale + bias
-
-  return init_fun, apply_fun
-
-
-def Embedding(feature_depth, vocab_size):  # pylint: disable=invalid-name
-  """Layer constructor function for a dense embedding layer."""
-  def init_fun(rng, input_shape):
-    output_shape = tuple(input_shape) + (feature_depth,)
-    dense_embedding = xavier_uniform()(rng, (vocab_size, feature_depth))
-    return output_shape, dense_embedding
-  def apply_fun(params, inputs, **kwargs):
-    del kwargs
-    dense_embedding = params
-    return np.take(dense_embedding, inputs, axis=0)
-  return init_fun, apply_fun
-
-
-def PositionalEncoding(feature_depth, max_len):  # pylint: disable=invalid-name
+# Layer normalization.
+def _layer_norm_new_params(input_shape, rng, epsilon=1e-6):  # pylint: disable=invalid-name
+  """Helper: create layer norm parameters."""
+  del rng, epsilon
+  features = input_shape[-1]
+  scale = np.ones(features)
+  bias = np.zeros(features)
+  return (scale, bias)
+
+
+@base.layer(new_parameters=_layer_norm_new_params)
+def LayerNorm(params, x, epsilon=1e-6, **unused_kwargs):
+  (scale, bias) = params
+  mean = np.mean(x, axis=-1, keepdims=True)
+  variance = np.mean((x - mean)**2, axis=-1, keepdims=True)
+  norm_inputs = (x - mean) / np.sqrt(variance + epsilon)
+  return norm_inputs * scale + bias
+
+
+# Positional encoding.
+def _positional_encoding_new_params(input_shape, rng, max_len=2048):  # pylint: disable=invalid-name
+  """Helper: create positional encoding parameters."""
+  del rng
+  feature_depth = input_shape[-1]
+  pe = onp.zeros((max_len, feature_depth), dtype=onp.float32)
+  position = onp.arange(0, max_len)[:, onp.newaxis]
+  div_term = onp.exp(
+      onp.arange(0, feature_depth, 2) * -(onp.log(10000.0) / feature_depth))
+  pe[:, 0::2] = onp.sin(position * div_term)
+  pe[:, 1::2] = onp.cos(position * div_term)
+  return np.array(pe[onp.newaxis, :])  # send to device
+
+
+@base.layer(new_parameters=_positional_encoding_new_params)
+def PositionalEncoding(params, x, **unused_kwargs):
   """Implements bare positional encoding."""
-  def init_fun(_, input_shape):
-    # Compute the positional encodings once in log space.
-    pe = onp.zeros((max_len, feature_depth), dtype=onp.float32)
-    position = onp.arange(0, max_len)[:, onp.newaxis]
-    div_term = onp.exp(
-        onp.arange(0, feature_depth, 2) * -(onp.log(10000.0) / feature_depth))
-    pe[:, 0::2] = onp.sin(position * div_term)
-    pe[:, 1::2] = onp.cos(position * div_term)
-    pe = np.array(pe[onp.newaxis, :])  # send to device
-    return input_shape, pe
-
-  def apply_fun(params, inputs, **kwargs):
-    del kwargs
-    pe = params
-    symbol_size = np.shape(inputs)[1]
-    return inputs + pe[:, :symbol_size]
-
-  return init_fun, apply_fun
+  symbol_size = np.shape(x)[1]
+  return x + params[:, :symbol_size]
 
 
-def dot_product_attention(query, key, value, mask, dropout, mode, rng):
+def DotProductAttention(query, key, value, mask, dropout, mode, rng):
   """Core dot product self-attention.
 
   Args:
@@ -155,17 +126,19 @@ def dot_product_attention(query, key, value, mask, dropout, mode, rng):
   dots = np.matmul(query, np.swapaxes(key, -1, -2)) / np.sqrt(depth)
   if mask is not None:
     dots = np.where(mask, dots, -1e9)
-  dots = stax.softmax(dots, axis=-1)
+  # Softmax.
+  dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))
   if dropout >= 1.0:
     raise ValueError('Dropout rates must be lower than 1.')
   if dropout is not None and dropout > 0.0 and mode == 'train':
-    keep = random.bernoulli(rng, 1.0 - dropout, dots.shape)
+    keep = backend.random.bernoulli(rng, 1.0 - dropout, dots.shape)
     dots = np.where(keep, dots / (1.0 - dropout), 0)
   out = np.matmul(dots, value)
   return out
 
 
-def PureDotProductAttention(dropout=0.0, mode='train'):  # pylint: disable=invalid-name
+# TODO(lukaszkaiser): make this a layer.
+def PureDotProductAttention(dropout=0.0, mode='train'):
   """Pure single-headed self-attention.
 
   Args:
@@ -175,60 +148,66 @@ def PureDotProductAttention(dropout=0.0, mode='train'):  # pylint: disable=inval
   Returns:
     Pure single-headed attention layer. (No Dense transforms on input.)
   """
-  def init_fun(_, input_shapes):
+  def init_fun(_, input_shapes):  # pylint: disable=invalid-name
     q_shape, _, v_shape, _ = input_shapes
     output_shape = q_shape[:-1] + (v_shape[-1],)
     return output_shape, ()
-  def apply_fun(params, inputs, **kwargs):
+  def apply_fun(params, inputs, **kwargs):  # pylint: disable=invalid-name
     del params
     q, k, v, mask = inputs
     rng = kwargs.get('rng', None)
-    return dot_product_attention(q, k, v, mask,
-                                 dropout=dropout, mode=mode, rng=rng)
+    return DotProductAttention(q, k, v, mask,
+                               dropout=dropout, mode=mode, rng=rng)
   return init_fun, apply_fun
 
 
-def PureMultiHeadedAttention(  # pylint: disable=invalid-name
-    feature_depth, num_heads=8, dropout=0.0, mode='train'):
+def _multihead_attention_output_shape(  # pylint: disable=invalid-name
+    input_shapes, feature_depth=None, **unused_kwargs):
+  """Helper: calculate multihead attention output shape."""
+  input_shape = input_shapes[0]  # Inputs are (q, k, v, mask).
+  return input_shape[:-1] + (feature_depth,)
+
+
+@base.layer(output_shape=_multihead_attention_output_shape)
+def PureMultiHeadedAttention(
+    params, x, feature_depth=None, num_heads=8, dropout=0.0, mode='train',
+    **kwargs):
   """Pure transformer-style multi-headed attention.
 
   Args:
+    params: parameters (none)
+    x: inputs (q, k, v, mask)
     feature_depth: int:  depth of embedding
     num_heads: int: number of attention heads
     dropout: float: dropout rate
     mode: str: 'train' or 'eval'
+    **kwargs: other arguments including the rng
 
   Returns:
     Pure Multi-headed attention layer. (No Dense transforms on input.)
   """
-  def init_fun(_, input_shapes):
-    input_shape = input_shapes[0]
-    output_shape = input_shape[:-1] + (feature_depth,)
-    return output_shape, ()
-  def apply_fun(params, inputs, **kwargs):  # pylint: disable=missing-docstring
-    del params
-    rng = kwargs.get('rng', None)
-    q, k, v, mask = inputs
-    assert feature_depth % num_heads == 0
-    head_depth = feature_depth // num_heads
-    nbatch = np.shape(q)[0]
-    # nbatch, seqlen, feature_depth --> nbatch, num_heads, seqlen, head_depth
-    def split_heads(x):
-      return np.transpose(
-          np.reshape(x, (nbatch, -1, num_heads, head_depth)), (0, 2, 1, 3))
-    # nbatch, num_heads, seqlen, head_depth --> nbatch, seqlen, feature_depth
-    def join_heads(x):
-      return np.reshape(
-          np.transpose(x, (0, 2, 1, 3)), (nbatch, -1, num_heads*head_depth))
-    # Split heads, dot-product attention, rejoin heads.
-    return join_heads(
-        dot_product_attention(
-            split_heads(q), split_heads(k), split_heads(v), mask,
-            dropout=dropout, mode=mode, rng=rng))
-  return init_fun, apply_fun
-
-
-def MultiHeadedAttention(  # pylint: disable=invalid-name
+  del params
+  rng = kwargs.get('rng', None)
+  q, k, v, mask = x
+  assert feature_depth % num_heads == 0
+  head_depth = feature_depth // num_heads
+  nbatch = np.shape(q)[0]
+  # nbatch, seqlen, feature_depth --> nbatch, num_heads, seqlen, head_depth
+  def SplitHeads(x):
+    return np.transpose(
+        np.reshape(x, (nbatch, -1, num_heads, head_depth)), (0, 2, 1, 3))
+  # nbatch, num_heads, seqlen, head_depth --> nbatch, seqlen, feature_depth
+  def JoinHeads(x):  # pylint: disable=invalid-name
+    return np.reshape(
+        np.transpose(x, (0, 2, 1, 3)), (nbatch, -1, num_heads*head_depth))
+  # Split heads, dot-product attention, rejoin heads.
+  return JoinHeads(
+      DotProductAttention(
+          SplitHeads(q), SplitHeads(k), SplitHeads(v), mask,
+          dropout=dropout, mode=mode, rng=rng))
+
+
+def MultiHeadedAttention(
     feature_depth, num_heads=8, dropout=0.0, mode='train'):
   """Transformer-style multi-headed attention.
 
@@ -241,14 +220,15 @@ def MultiHeadedAttention(  # pylint: disable=invalid-name
   Returns:
     Multi-headed self-attention layer.
   """
-  return stax.serial(
-      stax.parallel(
-          stax.Dense(feature_depth, W_init=xavier_uniform()),
-          stax.Dense(feature_depth, W_init=xavier_uniform()),
-          stax.Dense(feature_depth, W_init=xavier_uniform()),
-          stax.Identity
+  return combinators.Serial(
+      combinators.Parallel(
+          stax.Dense(feature_depth, W_init=stax.xavier_uniform()),
+          stax.Dense(feature_depth, W_init=stax.xavier_uniform()),
+          stax.Dense(feature_depth, W_init=stax.xavier_uniform()),
+          combinators.Identity()
       ),
-      PureMultiHeadedAttention(
-          feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),
-      stax.Dense(feature_depth, W_init=xavier_uniform()),
+      PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
+          feature_depth=feature_depth, num_heads=num_heads,
+          dropout=dropout, mode=mode),
+      stax.Dense(feature_depth, W_init=stax.xavier_uniform()),
   )
diff --git a/tensor2tensor/trax/stax/base.py b/tensor2tensor/trax/stax/base.py
index 1b1ab2300..71228036a 100644
--- a/tensor2tensor/trax/stax/base.py
+++ b/tensor2tensor/trax/stax/base.py
@@ -19,60 +19,214 @@
 from __future__ import division
 from __future__ import print_function
 
-from jax.tree_util import register_pytree_node as _register_pytree_node
+import inspect
+import traceback
 
 
-# Staxlayer binding to python variables
-# ------------------------------------------------------------------------------
-# Stax params-tree leaf type to mark bound subtrees references.
-class _TreeMarker(dict):
-  pass
-# Add this leaf-type to JAX's tree-walker.
-_register_pytree_node(_TreeMarker,
-                      lambda xs: (tuple(), None),
-                      lambda _, xs: _TreeMarker())
+class Layer(object):
+  """Layer object, base class. Handles parameter sharing."""
 
+  def __init__(self, **kwargs):
+    # We store kwargs by default, used below in creating a generic decorator.
+    self._init_kwargs = kwargs
+    # This field says if this layer's init has already been called or not.
+    self._first_init = True
+    # Cache parameters here, defaults empty params (we use () for that).
+    self._params = ()  # cached parameters
+    # Caller field storing file name and line number where the caller class
+    # was created. Since most layers inherit directly from this class, they
+    # call this init (stack 0) in their init (stack 1) and we want the caller
+    # that created them, so we take stack position 2 here.
+    # TODO(lukaszkaiser): this will break with more inheritance, correct.
+    self._caller = inspect.getframeinfo(inspect.stack()[2][0])
 
-# TODO(lukaszkaiser): make this the base layer class (share by object).
-class Share(tuple):
-  """Layer parameter caching function to allow weight sharing.
+  def call(self, params, inputs, **kwargs):
+    """Call this layer using the given parameters on the given inputs."""
+    raise NotImplementedError
 
-  Args:
-    A staxlayer: an (init_fun, apply_fun) pair.
+  def output_shape(self, input_shape):
+    """The shape of the output of this layer given the shape of the input.
 
-  Returns:
-    A 'parameter-bound' staxlayer that can be assigned to a python variable.
-  Wherever this value is needed elsewhere in the stax tree, call this bound
-  variable and all occurrences will share parameters that will automatically
-  be updated by Stax optimizers.
-  """
+    Note that all arguments and return values can be tuples or dictionaries
+    or arbitraty nested structures composed of tuples and dictionaries.
 
-  def __init__(self, staxlayer):  # pylint: disable=super-init-not-called
-    self._orig_init_fun, self._orig_apply_fun = staxlayer
-    self._first_init = True
-    self.params = None  # cached staxlayer params
-
-  def _init_fun(self, rng_key, input_shape):  # pylint: disable=missing-docstring
-    if self._first_init:
-      # point of first subgraph initialization call: sets params, output_shape
-      self._first_init = False
-      out_shape, self.params = self._orig_init_fun(rng_key, input_shape)
-      return out_shape, self.params
-    else:
-      # point of subgraph reuse:
-      # params are just a marker to apply_funs signalling subgraph params reuse
-      out_shape, _ = self._orig_init_fun(rng_key, input_shape)
-      return out_shape, _TreeMarker()
-
-  def _apply_fun(self, params, inputs, **kwargs):
-    if isinstance(params, _TreeMarker):
-      # point of subgraph reuse: calculate new value with cached params
-      return self._orig_apply_fun(self.params, inputs, **kwargs)
-    else:
-      # point of first subgraph application to params: cache params
-      self.params = params
-      return self._orig_apply_fun(params, inputs, **kwargs)
+    Args:
+      input_shape: a tuple representing the shape of the input.
+
+    Returns:
+      The shape of the output.
+    """
+    raise NotImplementedError
+
+  def new_parameters(self, input_shape, rng):
+    """Initialize parameters given input shape and return with output shape.
+
+    Note that all arguments and return values can be tuples or dictionaries
+    or arbitraty nested structures composed of tuples and dictionaries.
+
+    Args:
+      input_shape: a tuple representing the shape of the input.
+      rng: random number generator.
+
+    Returns:
+      The newly created parameters for this layer.
+    """
+    raise NotImplementedError
+
+  # End of subclassing interface, all functions below are internal.
+
+  def _init_fun(self, rng, input_shape):
+    """Internal modification of init_fun, saves variables."""
+    out_shape = self.output_shape(input_shape)
+
+    # Re-using this layer, no new parameters.
+    if not self._first_init:
+      return out_shape, ()
+
+    # First call of this layer, create parameters.
+    self._first_init = False
+    self._params = self.new_parameters(input_shape, rng)
+    return out_shape, self._params
+
+  def __call__(self, params, inputs, **kwargs):
+    try:
+      # If params are nothing, we may be reusing this layer.
+      # Use the cached parameters to calculate the value.
+      # Note: to make sure jit tracers can decide this branch in python we
+      #   use "params is ()" instead of, e.g., "not params" or "params == ()".
+      if params is ():  # pylint: disable=literal-comparison
+        return self.call(self._params, inputs, **kwargs)
+      # In this case, we're called for the first time: cache parameters.
+      self._params = params
+      return self.call(params, inputs, **kwargs)
+    except Exception:
+      name, trace = self.__class__.__name__, _short_traceback()
+      raise LayerError(name, self._caller, shapes(inputs), trace)
 
   # when unpacking this (init, apply) pair we return the wrapped funs
   def __iter__(self):
-    return iter((self._init_fun, self._apply_fun))
+    return iter((self._init_fun, self.__call__))
+
+
+class LayerError(Exception):
+  """Exception raised in the layer stack.
+
+  Attributes:
+    message: the message corresponding to this exception.
+  """
+
+  def __init__(self, layer_name, caller, input_shapes, traceback_string):
+    self._layer_name = layer_name
+    self._caller = caller  # Python inspect object with init caller info.
+    self._traceback = traceback_string
+    self._input_shapes = input_shapes
+    super(LayerError, self).__init__(self.message)
+
+  @property
+  def message(self):
+    prefix = 'Exception passing through layer %s:\n' % self._layer_name
+    short_path = '[...]/' + '/'.join(self._caller.filename.split('/')[-3:])
+    caller = '  layer created in file %s, line %d\n' % (short_path,
+                                                        self._caller.lineno)
+    shapes_str = '  layer input shapes: %s\n\n' % str(self._input_shapes)
+    return prefix + caller + shapes_str + self._traceback
+
+
+def nested_map(x, f):
+  """Map the function f to the nested structure x (dicts, tuples, lists)."""
+  if isinstance(x, list):
+    return [nested_map(y, f) for y in x]
+  if isinstance(x, tuple):
+    return tuple([nested_map(y, f) for y in x])
+  if isinstance(x, dict):
+    return {k: nested_map(x[k], f) for k in x}
+  return f(x)
+
+
+def shapes(x):
+  """Get a structure of shapes for a structure of nested arrays."""
+  def shape(x):
+    try:
+      return x.shape
+    except Exception:  # pylint: disable=broad-except
+      return []
+  return nested_map(x, shape)
+
+
+def _shorten_file_path(line):
+  """Shorten file path in error lines for more readable tracebacks."""
+  start = line.lower().find('file')
+  if start < 0:
+    return line
+  first_quote = line.find('"', start)
+  if first_quote < 0:
+    return line
+  second_quote = line.find('"', first_quote + 1)
+  if second_quote < 0:
+    return line
+  path = line[first_quote + 1:second_quote]
+  new_path = '/'.join(path.split('/')[-3:])
+  return line[:first_quote] + '[...]/' + new_path + line[second_quote + 1:]
+
+
+def _short_traceback(skip=3):
+  """Cleaned-up form of traceback."""
+  counter, res = 0, []
+  # Skipping 3 lines by default: the top (useless) and self-call.
+  lines = traceback.format_exc().splitlines()[skip:]
+  for l in lines:
+    res.append(_shorten_file_path(l))
+    if counter % 2 == 1:
+      res.append('')
+    counter += 1
+    # If we see a LayerError, the traceback has already been processed.
+    if l.startswith('LayerError'):
+      # Skip 4 back except last as these are internal base-layer calls.
+      res = res[:-4] + [res[-1]]
+      res += lines[counter:]
+      break
+  return '\n'.join(res)
+
+
+# Decorator for making layers from functions.
+
+
+def layer(output_shape=None, new_parameters=None):
+  """Create a layer class from a function."""
+  def layer_decorator(call):
+    """Decorating the call function."""
+    def output_shape_fun(self, input_shape):
+      if output_shape is None:
+        return input_shape
+      kwargs = self._init_kwargs  # pylint: disable=protected-access
+      return output_shape(input_shape, **kwargs)
+
+    def new_parameters_fun(self, input_shape, rng):
+      if new_parameters is None:
+        return ()
+      kwargs = self._init_kwargs  # pylint: disable=protected-access
+      return new_parameters(input_shape, rng, **kwargs)
+
+    def call_fun(self, params, inputs, **kwargs):
+      """The call function of the created class, derived from call."""
+      # Merge on-call kwargs with class-kwargs.
+      call_kwargs = kwargs.copy()
+      call_kwargs.update(self._init_kwargs)  # pylint: disable=protected-access
+      # Call with the merged kwargs.
+      return call(params, inputs, **call_kwargs)
+
+    # Set doc for python help.
+    call_fun.__doc__ = call.__doc__
+    if output_shape is None:
+      output_shape_fun.__doc__ = output_shape.__doc__
+    if new_parameters is None:
+      new_parameters_fun.__doc__ = new_parameters.__doc__
+
+    # Create the class.
+    cls = type(call.__name__, (Layer,),
+               {'call': call_fun,
+                'output_shape': output_shape_fun,
+                'new_parameters': new_parameters_fun})
+
+    return cls
+  return layer_decorator
diff --git a/tensor2tensor/trax/stax/base_test.py b/tensor2tensor/trax/stax/base_test.py
index 5902ec7ef..5d886635f 100644
--- a/tensor2tensor/trax/stax/base_test.py
+++ b/tensor2tensor/trax/stax/base_test.py
@@ -18,10 +18,9 @@
 from __future__ import division
 from __future__ import print_function
 
-import itertools
 from absl.testing import absltest
-from jax import random
 import numpy as onp
+from tensor2tensor.trax.backend import random
 import tensor2tensor.trax.stax as stax
 
 
@@ -35,7 +34,7 @@ def random_inputs(rng, input_shape):
 
 
 def check_shape_agreement(test_case, init_fun, apply_fun, input_shape):
-  rng_key1, rng_key2 = random.split(random.PRNGKey(0))
+  rng_key1, rng_key2 = random.split(random.get_prng(0))
   result_shape, params = init_fun(rng_key1, input_shape)
   inputs = random_inputs(onp.random.RandomState(0), input_shape)
   result = apply_fun(params, inputs, rng=rng_key2)
@@ -48,80 +47,39 @@ def check_staxlayer(test_case, staxlayer, input_shape):
   return check_shape_agreement(test_case, init_fun, apply_fun, input_shape)
 
 
-# Helper functions for testing Lambda wrapper against functions involving
-# complicated input trees:
-def _enumerate_trees_w_leaves(n_leaves):
-  """Construct all rooted trees with n leaves."""
-  def enumtree(*args):
-    n_args = len(args)
-    # trivial cases:
-    if n_args == 0:
-      return []
-    if n_args == 1:
-      return args
-    # general case of 2 or more args:
-    # build index array
-    idxs = range(0, n_args)
-    trees = []
-    # we consider all possible subsets of size n_set to gather
-    for n_set in range(2, n_args+1):
-      idxsets = list(itertools.combinations(idxs, n_set))
-      for idxset in idxsets:
-        # recurse by joining all subtrees with
-        # n_set leaves and (n_args - n_set) leaves
-        arg_set = tuple(args[i] for i in idxs if i in idxset)
-        arg_coset = tuple(args[i] for i in idxs if i not in idxset)
-        if arg_coset:
-          trees.extend(tuple(itertools.product(enumtree(*arg_set),
-                                               enumtree(*arg_coset))))
-        else:
-          # trivial case where arg_set is entire set
-          trees.append(arg_set)
-    return trees
-  # return enumerated trees with integers as leaves
-  return enumtree(*range(n_leaves))
-
-
-def _build_combinator_tree(input_treespec, in_vars):
-  """Build a trivial Staxlayer that takes a complicated tree of inputs."""
-  parallel_args = []
-  for e in input_treespec:
-    if isinstance(e, int):
-      parallel_args.append(in_vars[e])
-    elif isinstance(e, tuple):
-      parallel_args.append(_build_combinator_tree(e, in_vars))
-  return stax.serial(stax.parallel(*parallel_args), stax.FanInSum)
-
-
 class SlaxTest(absltest.TestCase):
 
   def test_flatten_n(self):
     input_shape = (29, 87, 10, 20, 30)
 
-    actual_shape = check_staxlayer(self, stax.Flatten(1), input_shape)
+    actual_shape = check_staxlayer(self, stax.Flatten(), input_shape)
     self.assertEqual(actual_shape, (29, 87 * 10 * 20 * 30))
 
-    actual_shape = check_staxlayer(self, stax.Flatten(2), input_shape)
+    actual_shape = check_staxlayer(self, stax.Flatten(num_axis_to_keep=2),
+                                   input_shape)
     self.assertEqual(actual_shape, (29, 87, 10 * 20 * 30))
 
-    actual_shape = check_staxlayer(self, stax.Flatten(3), input_shape)
+    actual_shape = check_staxlayer(self, stax.Flatten(num_axis_to_keep=3),
+                                   input_shape)
     self.assertEqual(actual_shape, (29, 87, 10, 20 * 30))
 
-    actual_shape = check_staxlayer(self, stax.Flatten(4), input_shape)
+    actual_shape = check_staxlayer(self, stax.Flatten(num_axis_to_keep=4),
+                                   input_shape)
     self.assertEqual(actual_shape, (29, 87, 10, 20, 30))
 
     # Not enough dimensions.
     with self.assertRaises(ValueError):
-      check_staxlayer(self, stax.Flatten(5), input_shape)
+      check_staxlayer(self, stax.Flatten(num_axis_to_keep=5), input_shape)
 
     with self.assertRaises(ValueError):
-      check_staxlayer(self, stax.Flatten(6), input_shape)
+      check_staxlayer(self, stax.Flatten(num_axis_to_keep=6), input_shape)
 
   def test_div(self):
-    init_fun, apply_fun = stax.Div(2)
+    init_fun, apply_fun = stax.Div(divisor=2.0)
     input_np = onp.array([[1, 2, 3], [4, 5, 6]], dtype=onp.float32)
     input_shape = input_np.shape
-    _, _ = init_fun(None, input_shape)
+    rng = random.get_prng(0)
+    _, _ = init_fun(rng, input_shape)
     output_np = apply_fun(None, input_np)
     # absltest doesn't have ndarray equalities.
     expected_output_np = input_np / 2.0
@@ -130,6 +88,22 @@ def test_div(self):
         onp.sum((output_np - expected_output_np) ** 2),
         delta=1e-6)
 
+  def test_dense_param_sharing(self):
+    model1 = stax.Serial(stax.Dense(32), stax.Dense(32))
+    layer = stax.Dense(32)
+    model2 = stax.Serial(layer, layer)
+    init_fun1, _ = model1
+    init_fun2, _ = model2
+    rng = random.get_prng(0)
+    _, params1 = init_fun1(rng, [-1, 32])
+    _, params2 = init_fun2(rng, [-1, 32])
+    # The first parameters have 2 kernels of size (32, 32).
+    self.assertEqual((32, 32), params1[0][0].shape)
+    self.assertEqual((32, 32), params1[1][0].shape)
+    # The second parameters have 1 kernel of size (32, 32) and an empty dict.
+    self.assertEqual((32, 32), params2[0][0].shape)
+    self.assertEqual((), params2[1])
+
 
 if __name__ == "__main__":
   absltest.main()
diff --git a/tensor2tensor/trax/stax/combinators.py b/tensor2tensor/trax/stax/combinators.py
new file mode 100644
index 000000000..044cbca95
--- /dev/null
+++ b/tensor2tensor/trax/stax/combinators.py
@@ -0,0 +1,154 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Combinators for composing layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.trax import backend
+from tensor2tensor.trax.stax import base
+
+
+class Serial(base.Layer):
+  """Layer composing a number of sub-layers in a serial way.."""
+
+  def __init__(self, *layers):
+    super(Serial, self).__init__()
+    # If called with one list argument, treat it as layers.
+    if len(layers) == 1 and isinstance(layers[0], list):
+      layers = layers[0]
+    self._nlayers = len(layers)
+    self._layers = layers
+    self._init_funs, self._apply_funs = zip(*layers)
+
+  def call(self, params, inputs, **kwargs):
+    rng = kwargs.pop('rng', None)
+    rngs = (None,) * self._nlayers
+    if rng is not None:
+      rngs = backend.random.split(rng, self._nlayers)
+    for fun, param, rng in zip(self._apply_funs, params, rngs):
+      inputs = fun(param, inputs, rng=rng, **kwargs)
+    return inputs
+
+  def output_shape(self, input_shape):
+    cur_shape = input_shape
+    for layer in self._layers:
+      cur_shape = layer.output_shape(cur_shape)
+    return cur_shape
+
+  def new_parameters(self, input_shape, rng):
+    params = []
+    for init_fun in self._init_funs:
+      rng, layer_rng = backend.random.split(rng)
+      input_shape, param = init_fun(layer_rng, input_shape)
+      params.append(param)
+    return params
+
+
+@base.layer()
+def Identity(params, x, **kwargs):
+  del params, kwargs
+  return x
+
+
+@base.layer(output_shape=lambda input_shape, size=2: [input_shape] * size)
+def FanOut(params, x, size=2, **kwargs):
+  del params, kwargs
+  return [x] * size
+
+
+@base.layer(output_shape=lambda input_shape_list: input_shape_list[0])
+def FanInSum(params, x, **kwargs):
+  del params, kwargs
+  return sum(x)  # Here x is a list of tensors of the same shape, we add them.
+
+
+def _fan_in_concat_shape(input_shape, axis=-1):  # pylint: disable=invalid-name
+  """Helper to determine the shape of FanInConcat output."""
+  ax = axis % len(input_shape[0])
+  concat_size = sum(shape[ax] for shape in input_shape)
+  out_shape = input_shape[0][:ax] + (concat_size,) + input_shape[0][ax+1:]
+  return out_shape
+
+
+@base.layer(output_shape=_fan_in_concat_shape)
+def FanInConcat(params, x, axis=-1, **kwargs):
+  del params, kwargs
+  return backend.numpy.concatenate(x, axis)
+
+
+class Parallel(base.Layer):
+  """Combinator for composing layers in parallel.
+
+  The layer resulting from this combinator is often used with the FanOut and
+  FanInSum layers.
+
+  Args:
+    *layers: a sequence of layers, each an (init_fun, apply_fun) pair.
+
+  Returns:
+    A new layer, meaning an (init_fun, apply_fun) pair, representing the
+    parallel composition of the given sequence of layers. In particular, the
+    returned layer takes a sequence of inputs and returns a sequence of outputs
+    with the same length as the argument `layers`.
+  """
+
+  def __init__(self, *layers):
+    super(Parallel, self).__init__()
+    # If called with one list argument, treat it as layers.
+    if len(layers) == 1 and isinstance(layers[0], list):
+      layers = layers[0]
+    self._nlayers = len(layers)
+    self._layers = layers
+    self._init_funs, self._apply_funs = zip(*layers)
+
+  def call(self, params, inputs, **kwargs):
+    rng = kwargs.pop('rng', None)
+    rngs = (None,) * self._nlayers
+    if rng is not None:
+      rngs = backend.random.split(rng, self._nlayers)
+    return [f(p, x, rng=r, **kwargs)
+            for f, p, x, r in zip(self._apply_funs, params, inputs, rngs)]
+
+  def output_shape(self, input_shapes):
+    return tuple([layer.output_shape(shape)
+                  for layer, shape in zip(self._layers, input_shapes)])
+
+  def new_parameters(self, input_shape, rng):
+    rngs = backend.random.split(rng, self._nlayers)
+    _, p = zip(*[init(rng, shape) for init, rng, shape
+                 in zip(self._init_funs, rngs, input_shape)])
+    return p
+
+
+def Residual(*layers, **kwargs):
+  """Constructs a residual version of layers, summing input to layers output."""
+  res = kwargs.get('res', Identity())  # pylint: disable=no-value-for-parameter
+  if len(layers) > 1:
+    return Serial(
+        FanOut(),  # pylint: disable=no-value-for-parameter
+        Parallel(Serial(*layers), res),
+        FanInSum()  # pylint: disable=no-value-for-parameter
+    )
+  elif len(layers) == 1:
+    return Serial(
+        FanOut(),  # pylint: disable=no-value-for-parameter
+        Parallel(layers[0], res),
+        FanInSum()  # pylint: disable=no-value-for-parameter
+    )
+  else:
+    raise ValueError('Empty residual combinator.')
diff --git a/tensor2tensor/trax/stax/stax_base.py b/tensor2tensor/trax/stax/stax_base.py
index d9e9e4c6e..1d9115f0b 100644
--- a/tensor2tensor/trax/stax/stax_base.py
+++ b/tensor2tensor/trax/stax/stax_base.py
@@ -13,16 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Stax is a small flexible neural net specification library from scratch."""
-
-# Forked from JAX for more experimentation on syntax and back-ends.
-# See JAX version at https://github.com/google/jax/tree/master/jax/experimental
+"""Trax layers library."""
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import functools
 import itertools
 import operator as op
 
@@ -32,48 +28,17 @@
 from six.moves import reduce
 from tensor2tensor.trax import backend
 from tensor2tensor.trax.backend import numpy as np
+from tensor2tensor.trax.stax import base
 
 # Following the convention used in Keras and tf.layers, we use CamelCase for the
 # names of layer constructors, like Conv and Relu, while using snake_case for
-# other functions, like lax.conv and relu.
-
-# use CamelCase for layer constructors
+# other functions, like lax.conv and relu. To allow this, we disable below.
 # pylint: disable=invalid-name
 
-# don't have docstrings for init_fun / apply_fun pairs
-# pylint: disable=missing-docstring
-
-# ignore unused arguments
-# pylint: disable=unused-argument
-
-
-def relu(x):
-  return np.maximum(x, 0.)
-
 
-def softplus(x):
-  return np.logaddexp(x, 0.)
+# Initializers.
 
 
-def logsoftmax(x, axis=-1):
-  """Apply log softmax to an array of logits, log-normalizing along an axis."""
-  return x - backend.logsumexp(x, axis, keepdims=True)
-
-
-def softmax(x, axis=-1):
-  """Apply softmax to an array of logits, exponentiating and normalizing along an axis."""
-  unnormalized = np.exp(x - x.max(axis, keepdims=True))
-  return unnormalized / unnormalized.sum(axis, keepdims=True)
-
-
-def fastvar(x, axis, keepdims):
-  """A fast but less numerically-stable variance calculation than np.var."""
-  m1 = np.mean(x**2, axis, keepdims=keepdims)
-  m2 = np.mean(x, axis, keepdims=keepdims)**2
-  return m1 - m2
-
-
-# Initializers
 def randn(stddev=1e-2):
   """An initializer function for random normal coefficients."""
   def init(rng, shape):
@@ -91,112 +56,242 @@ def init(rng, shape):
   return init
 
 
-zeros = lambda rng, shape: np.zeros(shape, dtype='float32')
-ones = lambda rng, shape: np.ones(shape, dtype='float32')
+def xavier_uniform(out_dim=0, in_dim=1):
+  """An initializer function for random uniform xavier-scaled coefficients."""
+  def init(rng, shape):
+    fan_in, fan_out = shape[in_dim], shape[out_dim]
+    std = np.sqrt(2.0 / (fan_in + fan_out))
+    a = np.sqrt(3.0) * std
+    return backend.random.uniform(rng, shape, minval=-a, maxval=a)
+  return init
 
 
-# Layers
+def one_hot(x, size, dtype=np.float32):
+  """Make a n+1 dim one-hot array from n dim int-categorical array."""
+  return np.array(x[..., np.newaxis] == np.arange(size), dtype)
+
+
+# Layers.
+
+
+@base.layer()
+def Relu(params, x, **kwargs):
+  del params, kwargs
+  return np.maximum(x, 0.)
+
+
+@base.layer()
+def Tanh(params, x, **kwargs):
+  del params, kwargs
+  return np.tanh(x)
 
-# Each layer constructor function returns an (init_fun, apply_fun) pair, where
-#   init_fun: takes an input shape and returns an (output_shape, params) pair,
-#   apply_fun: takes params, inputs, and an rng key and applies the layer.
+
+@base.layer()
+def Exp(params, x, **kwargs):
+  del params, kwargs
+  return np.exp(x)
+
+
+@base.layer()
+def LogSoftmax(params, x, axis=-1, **kwargs):
+  """Apply log softmax to x: log-normalize along the given axis."""
+  del params, kwargs
+  return x - backend.logsumexp(x, axis, keepdims=True)
+
+
+@base.layer()
+def Softmax(params, x, axis=-1, **kwargs):
+  """Apply softmax to x: exponentiate and normalize along the given axis."""
+  del params, kwargs
+  return np.exp(x - backend.logsumexp(x, axis, keepdims=True))
+
+
+@base.layer()
+def Softplus(params, x, **kwargs):
+  del params, kwargs
+  return np.logaddexp(x, 0.)
 
 
-def Dense(out_dim, W_init=glorot(), b_init=randn()):
+class Dense(base.Layer):
   """Layer constructor function for a dense (fully-connected) layer."""
-  def init_fun(rng, input_shape):
-    output_shape = input_shape[:-1] + (out_dim,)
-    w, b = W_init(rng, (input_shape[-1], out_dim)), b_init(rng, (out_dim,))
-    return output_shape, (w, b)
-  def apply_fun(params, inputs, **kwargs):
-    del kwargs  # unused
+
+  def __init__(self, out_dim, W_init=glorot(), b_init=randn()):
+    super(Dense, self).__init__()
+    self._out_dim = out_dim
+    self._W_init = W_init
+    self._b_init = b_init
+
+  def call(self, params, inputs, **kwargs):
+    del kwargs
     w, b = params
     return np.dot(inputs, w) + b
-  return init_fun, apply_fun
-
-
-def GeneralConv(dimension_numbers, out_chan, filter_shape,
-                strides=None, padding='VALID', W_init=None, b_init=randn(1e-6)):
-  """Layer construction function for a general convolution layer."""
-  lhs_spec, rhs_spec, out_spec = dimension_numbers
-  one = (1,) * len(filter_shape)
-  strides = strides or one
-  W_init = W_init or glorot(rhs_spec.index('O'), rhs_spec.index('I'))
-  def init_fun(rng, input_shape):
-    filter_shape_iter = iter(filter_shape)
-    kernel_shape = [out_chan if c == 'O' else
-                    input_shape[lhs_spec.index('C')] if c == 'I' else
-                    next(filter_shape_iter) for c in rhs_spec]
-    output_shape = lax.conv_general_shape_tuple(
-        input_shape, kernel_shape, strides, padding, dimension_numbers)
-    bias_shape = [out_chan if c == 'C' else 1 for c in out_spec]
+
+  def output_shape(self, input_shape):
+    return tuple(input_shape[:-1]) + (self._out_dim,)
+
+  def new_parameters(self, input_shape, rng):
+    w = self._W_init(rng, (input_shape[-1], self._out_dim))
+    b = self._b_init(rng, (self._out_dim,))
+    return (w, b)
+
+
+class Embedding(base.Layer):
+  """Layer constructor function for an embedding layer."""
+
+  def __init__(self, feature_depth, vocab_size, W_init=xavier_uniform()):
+    super(Embedding, self).__init__()
+    self._feature_depth = feature_depth
+    self._vocab_size = vocab_size
+    self._W_init = W_init
+
+  def call(self, params, inputs, **kwargs):
+    del kwargs
+    return np.take(params, inputs, axis=0)
+
+  def output_shape(self, input_shape):
+    return tuple(input_shape) + (self._feature_depth,)
+
+  def new_parameters(self, input_shape, rng):
+    return self._W_init(rng, (self._vocab_size, self._feature_depth))
+
+
+class Conv(base.Layer):
+  """Layer constructor function for a general convolution layer."""
+
+  def __init__(self, out_chan, filter_shape, strides=None, padding='VALID',
+               dimension_numbers=('NHWC', 'HWIO', 'NHWC'),
+               W_init=None, b_init=randn(1e-6)):
+    super(Conv, self).__init__()
+    self._out_chan = out_chan
+    self._filter_shape = filter_shape
+    self._padding = padding
+    self._dimension_numbers = dimension_numbers
+    self._lhs_spec, self._rhs_spec, self._out_spec = dimension_numbers
+    self._one = (1,) * len(filter_shape)
+    self._strides = strides or self._one
+    self._b_init = b_init
+    rhs_spec = self._rhs_spec
+    self._W_init = W_init or glorot(rhs_spec.index('O'), rhs_spec.index('I'))
+
+  def call(self, params, inputs, **kwargs):
+    del kwargs
+    w, b = params
+    return lax.conv_general_dilated(
+        inputs, w, self._strides, self._padding, self._one, self._one,
+        self._dimension_numbers) + b
+
+  def _kernel_shape(self, input_shape):
+    """Helper to calculate the kernel shape."""
+    filter_shape_iter = iter(self._filter_shape)
+    return [self._out_chan if c == 'O' else
+            input_shape[self._lhs_spec.index('C')] if c == 'I' else
+            next(filter_shape_iter) for c in self._rhs_spec]
+
+  def output_shape(self, input_shape):
+    kernel_shape = self._kernel_shape(input_shape)
+    return lax.conv_general_shape_tuple(
+        input_shape, kernel_shape,
+        self._strides, self._padding, self._dimension_numbers)
+
+  def new_parameters(self, input_shape, rng):
+    kernel_shape = self._kernel_shape(input_shape)
+    bias_shape = [self._out_chan if c == 'C' else 1 for c in self._out_spec]
     bias_shape = tuple(itertools.dropwhile(lambda x: x == 1, bias_shape))
-    W, b = W_init(rng, kernel_shape), b_init(rng, bias_shape)
-    return output_shape, (W, b)
-  def apply_fun(params, inputs, **kwargs):
-    W, b = params
-    return lax.conv_general_dilated(inputs, W, strides, padding, one, one,
-                                    dimension_numbers) + b
-  return init_fun, apply_fun
-Conv = functools.partial(GeneralConv, ('NHWC', 'HWIO', 'NHWC'))
-
-
-def BatchNorm(axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True,
-              beta_init=zeros, gamma_init=ones):
-  """Layer construction function for a batch normalization layer."""
-  _beta_init = lambda rng, shape: beta_init(rng, shape) if center else ()
-  _gamma_init = lambda rng, shape: gamma_init(rng, shape) if scale else ()
+    w = self._W_init(rng, kernel_shape)
+    b = self._b_init(rng, bias_shape)
+    return (w, b)
+
+
+# Flatten.
+def _flatten_output_shape(input_shape, num_axis_to_keep=1):
+  """Output shape of a flatten layer."""
+  if num_axis_to_keep >= len(input_shape):
+    raise ValueError(
+        "num_axis_to_keep[%d] should be less than input's rank[%d]" %
+        (num_axis_to_keep, len(input_shape)))
+  return tuple(input_shape[:num_axis_to_keep]) + (
+      reduce(op.mul, input_shape[num_axis_to_keep:], 1),)
+
+
+@base.layer(output_shape=_flatten_output_shape)
+def Flatten(params, inputs, num_axis_to_keep=1, **kwargs):
+  del params, kwargs
+  return np.reshape(inputs, (inputs.shape[:num_axis_to_keep] + (-1,)))
+
+
+# Batch normalization.
+def _batch_norm_new_params(input_shape, rng, axis=(0, 1, 2),
+                           center=True, scale=True, **kwargs):
+  """Helper to initialize batch norm params."""
+  del rng, kwargs
   axis = (axis,) if np.isscalar(axis) else axis
-  def init_fun(rng, input_shape):
-    shape = tuple(d for i, d in enumerate(input_shape) if i not in axis)
-    beta, gamma = _beta_init(rng, shape), _gamma_init(rng, shape)
-    return input_shape, (beta, gamma)
-  def apply_fun(params, x, **kwargs):
-    beta, gamma = params
-    # TODO(phawkins): np.expand_dims should accept an axis tuple.
-    # (https://github.com/numpy/numpy/issues/12290)
-    ed = tuple(None if i in axis else slice(None) for i in range(np.ndim(x)))
-    beta = beta[ed]
-    gamma = gamma[ed]
-    mean, var = np.mean(x, axis, keepdims=True), fastvar(x, axis, keepdims=True)
-    z = (x - mean) / np.sqrt(var + epsilon)
-    if center and scale: return gamma * z + beta
-    if center: return z + beta
-    if scale: return gamma * z
-    return z
-  return init_fun, apply_fun
-
-
-def _elemwise_no_params(fun, **fun_kwargs):
-  init_fun = lambda rng, input_shape: (input_shape, ())
-  apply_fun = lambda params, inputs, **kwargs: fun(inputs, **fun_kwargs)
-  return init_fun, apply_fun
-Tanh = _elemwise_no_params(np.tanh)
-Relu = _elemwise_no_params(relu)
-Exp = _elemwise_no_params(np.exp)
-LogSoftmax = _elemwise_no_params(logsoftmax, axis=-1)
-Softmax = _elemwise_no_params(softmax, axis=-1)
-Softplus = _elemwise_no_params(softplus)
-
-
-def _pooling_layer(reducer, init_val, rescaler=None):
-  def PoolingLayer(window_shape, strides=None, padding='VALID'):
-    """Layer construction function for a pooling layer."""
-    strides = strides or (1,) * len(window_shape)
-    rescale = rescaler(window_shape, strides, padding) if rescaler else None
-    dims = (1,) + window_shape + (1,)  # NHWC
-    strides = (1,) + strides + (1,)
-    def init_fun(rng, input_shape):
-      out_shape = lax.reduce_window_shape_tuple(input_shape, dims, strides,
-                                                padding)
-      return out_shape, ()
-    def apply_fun(params, inputs, **kwargs):
-      out = lax.reduce_window(inputs, init_val, reducer, dims, strides, padding)
-      return rescale(out, inputs) if rescale else out
-    return init_fun, apply_fun
-  return PoolingLayer
-MaxPool = _pooling_layer(lax.max, -np.inf)
-SumPool = _pooling_layer(lax.add, 0.)
+  shape = tuple(d for i, d in enumerate(input_shape) if i not in axis)
+  beta = np.zeros(shape, dtype='float32') if center else ()
+  gamma = np.ones(shape, dtype='float32') if scale else ()
+  return (beta, gamma)
+
+
+@base.layer(new_parameters=_batch_norm_new_params)
+def BatchNorm(params, x, axis=(0, 1, 2), epsilon=1e-5,
+              center=True, scale=True, **unused_kwargs):
+  """Layer construction function for a batch normalization layer."""
+  mean = np.mean(x, axis, keepdims=True)
+  # Fast but less numerically-stable variance calculation than np.var.
+  m1 = np.mean(x**2, axis, keepdims=True)
+  var = m1 - mean**2
+  z = (x - mean) / np.sqrt(var + epsilon)
+
+  # Expand the parameters to have the right axes.
+  beta, gamma = params
+  # TODO(phawkins): np.expand_dims should accept an axis tuple.
+  # (https://github.com/numpy/numpy/issues/12290)
+  ed = tuple(None if i in axis else slice(None) for i in range(np.ndim(x)))
+  beta = beta[ed]
+  gamma = gamma[ed]
+
+  # Return the z rescaled by the parameters if requested.
+  if center and scale:
+    return gamma * z + beta
+  if center:
+    return z + beta
+  if scale:
+    return gamma * z
+  return z
+
+
+# Pooling.
+def _pooling_output_shape(input_shape, pool_size=(2, 2),
+                          strides=None, padding='VALID'):
+  """Helper: compute the output shape for the pooling layer."""
+  dims = (1,) + pool_size + (1,)  # NHWC
+  strides = strides or (1,) * len(pool_size)
+  strides = (1,) + strides + (1,)
+  return lax.reduce_window_shape_tuple(input_shape, dims, strides, padding)
+
+
+def _pooling_general(inputs, reducer, init_val, rescaler=None,
+                     pool_size=(2, 2), strides=None, padding='VALID'):
+  """Helper: general pooling computation used in pooling layers later."""
+  strides = strides or (1,) * len(pool_size)
+  rescale = rescaler(pool_size, strides, padding) if rescaler else None
+  dims = (1,) + pool_size + (1,)  # NHWC
+  strides = (1,) + strides + (1,)
+  out = lax.reduce_window(inputs, init_val, reducer, dims, strides, padding)
+  return rescale(out, inputs) if rescale else out
+
+
+@base.layer(output_shape=_pooling_output_shape)
+def MaxPool(params, x, pool_size=(2, 2), strides=None, padding='VALID', **kw):
+  del params, kw
+  return _pooling_general(x, lax.max, -np.inf, pool_size=pool_size,
+                          strides=strides, padding=padding)
+
+
+@base.layer(output_shape=_pooling_output_shape)
+def SumPool(params, x, pool_size=(2, 2), strides=None, padding='VALID', **kw):
+  del params, kw
+  return _pooling_general(x, lax.add, 0., pool_size=pool_size,
+                          strides=strides, padding=padding)
 
 
 def _normalize_by_window_size(dims, strides, padding):
@@ -205,214 +300,44 @@ def rescale(outputs, inputs):
     window_sizes = lax.reduce_window(one, 0., lax.add, dims, strides, padding)
     return outputs / window_sizes[..., np.newaxis]
   return rescale
-AvgPool = _pooling_layer(lax.add, 0., _normalize_by_window_size)
-
-
-def Flatten(num_axis_to_keep=1):
-  """Layer construction function for flattening all but the leading dims."""
-  def init_fun(rng, input_shape):
-    del rng
-    if num_axis_to_keep >= len(input_shape):
-      raise ValueError(
-          "num_axis_to_keep[%d] should be less than input's rank[%d]" %
-          (num_axis_to_keep, len(input_shape)))
-    output_shape = tuple(input_shape[:num_axis_to_keep]) + (
-        reduce(op.mul, input_shape[num_axis_to_keep:], 1),)
-    return output_shape, ()
-  def apply_fun(params, inputs, **kwargs):
-    del params, kwargs
-    return np.reshape(inputs, (inputs.shape[:num_axis_to_keep] + (-1,)))
-  return init_fun, apply_fun
-
-
-def Identity():
-  """Layer construction function for an identity layer."""
-  init_fun = lambda rng, input_shape: (input_shape, ())
-  apply_fun = lambda params, inputs, **kwargs: inputs
-  return init_fun, apply_fun
-Identity = Identity()
-
-
-def FanOut(num):
-  """Layer construction function for a fan-out layer."""
-  init_fun = lambda rng, input_shape: ([input_shape] * num, ())
-  apply_fun = lambda params, inputs, **kwargs: [inputs] * num
-  return init_fun, apply_fun
-
-
-def FanInSum():
-  """Layer construction function for a fan-in sum layer."""
-  init_fun = lambda rng, input_shape: (input_shape[0], ())
-  apply_fun = lambda params, inputs, **kwargs: sum(inputs)
-  return init_fun, apply_fun
-FanInSum = FanInSum()
-
-
-def FanInConcat(axis=-1):
-  """Layer construction function for a fan-in concatenation layer."""
-  def init_fun(rng, input_shape):
-    ax = axis % len(input_shape[0])
-    concat_size = sum(shape[ax] for shape in input_shape)
-    out_shape = input_shape[0][:ax] + (concat_size,) + input_shape[0][ax+1:]
-    return out_shape, ()
-  def apply_fun(params, inputs, **kwargs):
-    return np.concatenate(inputs, axis)
-  return init_fun, apply_fun
-
-
-def Dropout(rate, mode='train'):
+
+
+@base.layer(output_shape=_pooling_output_shape)
+def AvgPool(params, x, pool_size=(2, 2), strides=None, padding='VALID', **kw):
+  del params, kw
+  return _pooling_general(x, lax.add, 0., _normalize_by_window_size,
+                          pool_size, strides=strides, padding=padding)
+
+
+@base.layer()
+def Dropout(params, x, rate=0.0, mode='train', rng=None, **kwargs):
   """Layer construction function for a dropout layer with given rate."""
-  def init_fun(_, input_shape):
-    return input_shape, ()
-  def apply_fun(params, inputs, **kwargs):  # pylint: disable=missing-docstring
-    del params  # Unused.
-    rng = kwargs.get('rng', None)
-    if rng is None:
-      msg = ('Dropout layer requires apply_fun to be called with a PRNG key '
-             'argument. That is, instead of `apply_fun(params, inputs)`, call '
-             'it like `apply_fun(params, inputs, key)` where `key` is a '
-             'jax.random.PRNGKey value.')
-      raise ValueError(msg)
-    if rate >= 1.0:
-      raise ValueError('Dropout rates must be lower than 1.')
-    if mode == 'train' and rate > 0.0:
-      keep = backend.random.bernoulli(rng, 1.0 - rate, inputs.shape)
-      return np.where(keep, inputs / (1.0 - rate), 0)
-    else:
-      return inputs
-  return init_fun, apply_fun
-
-
-def Div(divisor):
-  def init_fun(_, input_shape):
-    return input_shape, ()
-  def apply_fun(params, inputs, **kwargs):
-    return inputs / divisor
-  return init_fun, apply_fun
-
-
-# Composing layers via combinators
-
-
-def serial(*layers):
-  """Combinator for composing layers in serial.
-
-  Args:
-    *layers: a sequence of layers, each an (init_fun, apply_fun) pair.
-  Returns:
-    A new layer, meaning an (init_fun, apply_fun) pair, representing the serial
-    composition of the given sequence of layers.
-  """
-  nlayers = len(layers)
-  init_funs, apply_funs = zip(*layers)
-  def init_fun(rng, input_shape):
-    params = []
-    for init_fun in init_funs:
-      rng, layer_rng = backend.random.split(rng)
-      input_shape, param = init_fun(layer_rng, input_shape)
-      params.append(param)
-    return input_shape, params
-  def apply_fun(params, inputs, **kwargs):
-    rng = kwargs.pop('rng', None)
-    rngs = (None,) * nlayers
-    if rng is not None:
-      rngs = backend.random.split(rng, nlayers)
-    for fun, param, rng in zip(apply_funs, params, rngs):
-      inputs = fun(param, inputs, rng=rng, **kwargs)
-    return inputs
-  return init_fun, apply_fun
-
-
-def parallel(*layers):
-  """Combinator for composing layers in parallel.
-
-  The layer resulting from this combinator is often used with the FanOut and
-  FanInSum layers.
-  Args:
-    *layers: a sequence of layers, each an (init_fun, apply_fun) pair.
-  Returns:
-    A new layer, meaning an (init_fun, apply_fun) pair, representing the
-    parallel composition of the given sequence of layers. In particular, the
-    returned layer takes a sequence of inputs and returns a sequence of outputs
-    with the same length as the argument `layers`.
-  """
-  nlayers = len(layers)
-  init_funs, apply_funs = zip(*layers)
-  def init_fun(rng, input_shape):
-    rngs = backend.random.split(rng, nlayers)
-    return zip(*[init(rng, shape) for init, rng, shape
-                 in zip(init_funs, rngs, input_shape)])
-  def apply_fun(params, inputs, **kwargs):
-    rng = kwargs.pop('rng', None)
-    rngs = (None,) * nlayers
-    if rng is not None:
-      rngs = backend.random.split(rng, nlayers)
-    return [f(p, x, rng=r, **kwargs)
-            for f, p, x, r in zip(apply_funs, params, inputs, rngs)]
-  return init_fun, apply_fun
-
-
-def shape_dependent(make_layer):
-  """Combinator to delay layer constructor pair until input shapes are known.
-
-  Args:
-    make_layer: a one-argument function that takes an input shape as an argument
-      (a tuple of positive integers) and returns an (init_fun, apply_fun) pair.
-  Returns:
-    A new layer, meaning an (init_fun, apply_fun) pair, representing the same
-    layer as returned by `make_layer` but with its construction delayed until
-    input shapes are known.
-  """
-  def init_fun(rng, input_shape):
-    return make_layer(input_shape)[0](rng, input_shape)
-  def apply_fun(params, inputs, **kwargs):
-    return make_layer(inputs.shape)[1](params, inputs, **kwargs)
-  return init_fun, apply_fun
-
-
-# Utility functions
-# ------------------------------------------------------------------------------
-def one_hot(x, size, dtype=np.float32):
-  """Make a n+1 dim one-hot array from n dim int-categorical array."""
-  return np.array(x[..., np.newaxis] == np.arange(size), dtype)
+  del params, kwargs
+  if rng is None:
+    msg = ('Dropout layer requires apply_fun to be called with a rng keyword '
+           'argument. That is, instead of `Dropout(params, inputs)`, call '
+           'it like `Dropout(params, inputs, rng=key)`.')
+    raise ValueError(msg)
+  if rate >= 1.0:
+    raise ValueError('Dropout rate (%f) must be lower than 1.' % rate)
+  if mode == 'train' and rate > 0.0:
+    keep = backend.random.bernoulli(rng, 1.0 - rate, x.shape)
+    return np.where(keep, x / (1.0 - rate), 0)
+  else:
+    return x
+
 
+@base.layer()
+def Div(params, x, divisor=1.0, **kwargs):
+  del params, kwargs
+  return x / divisor
 
-def ShiftRight():  # pylint: disable=invalid-name
+
+@base.layer()
+def ShiftRight(params, inputs, **kwargs):
   """Layer to shift the tensor to the right by padding on axis 1."""
-  init_fun = lambda _, input_shape: (input_shape, ())
-  def apply_fun(params, inputs, **kwargs):
-    del params, kwargs
-    pad_widths = [(0, 0), (1, 0)]
-    pad_widths += [(0, 0) for _ in range(len(inputs.shape) - 2)]
-    padded = np.pad(inputs, pad_widths, mode='constant')
-    return padded[:, :-1, ...]
-  return init_fun, apply_fun
-
-
-# Utility Combinators
-# ------------------------------------------------------------------------------
-def repeat(layer, num_repeats):
-  """Repeats layers serially num_repeats times."""
-  if num_repeats < 1:
-    raise ValueError('Repeat combinator num_repeats must be >= 1.')
-  layers = num_repeats * (layer,)
-  return serial(*layers)
-
-
-def residual(*layers, **kwargs):
-  """Constructs a residual version of layers, summing input to layers output."""
-  res = kwargs.get('res', Identity)
-  if len(layers) > 1:
-    return serial(
-        FanOut(2),
-        parallel(serial(*layers), res),
-        FanInSum
-    )
-  elif len(layers) == 1:
-    return serial(
-        FanOut(2),
-        parallel(layers[0], res),
-        FanInSum
-    )
-  else:
-    raise ValueError('Empty residual combinator.')
+  del params, kwargs
+  pad_widths = [(0, 0), (1, 0)]
+  pad_widths += [(0, 0) for _ in range(len(inputs.shape) - 2)]
+  padded = np.pad(inputs, pad_widths, mode='constant')
+  return padded[:, :-1, ...]

From f1e6f37a8d8932cc11df3af5a06e8ddb077febbb Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 22 Apr 2019 10:46:14 -0700
Subject: [PATCH 1929/2720] Fix imports

PiperOrigin-RevId: 244689120
---
 tensor2tensor/trax/rlax/ppo_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 8dbff285c..3f6d09999 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -22,10 +22,10 @@
 import jax
 from jax import random as jax_random
 import numpy as np
+from tensor2tensor.trax import stax
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.rlax import fake_env
 from tensor2tensor.trax.rlax import ppo
-from tensor2tensor.trax.stax import stax_base as stax
 from tensorflow import test
 
 

From 8372e5df3ae0ceb65c575daa06034e396151d46d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 22 Apr 2019 11:20:13 -0700
Subject: [PATCH 1930/2720] Fix bug where tfds was not deterministic even when
 random seed was set.

PiperOrigin-RevId: 244696681
---
 tensor2tensor/trax/inputs.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index ffeb4d639..0ddcba04e 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -146,7 +146,9 @@ def dataset_to_stream(dataset, input_name):
     yield inp, out
 
 
-def train_and_eval_dataset(dataset_name, data_dir):
+@gin.configurable(whitelist=["train_shuffle_files", "test_shuffle_files"])
+def train_and_eval_dataset(dataset_name, data_dir, train_shuffle_files=True,
+                           test_shuffle_files=False):
   """Return train and evaluation datasets, feature info and supervised keys.
 
   Args:
@@ -154,6 +156,10 @@ def train_and_eval_dataset(dataset_name, data_dir):
       then we'll search T2T Problem registry for it, otherwise we assume it
       is a dataset from TFDS and load it from there.
     data_dir: directory where the data is located.
+    train_shuffle_files: Boolean determining whether or not to shuffle the train
+      files at startup. Set to False if you want data determinism.
+    test_shuffle_files: Boolean determining whether or not to shuffle the test
+      files at startup. Set to False if you want data determinism.
 
   Returns:
     a 4-tuple consisting of:
@@ -176,8 +182,12 @@ def train_and_eval_dataset(dataset_name, data_dir):
   eval_split = tfds.Split.VALIDATION
   if tfds.Split.VALIDATION not in splits:
     eval_split = tfds.Split.TEST
-  train, valid = tfds.load(
-      name=dataset_name, split=[tfds.Split.TRAIN, eval_split])
+  train = tfds.load(
+      name=dataset_name, split=tfds.Split.TRAIN,
+      as_dataset_kwargs={"shuffle_files": train_shuffle_files})
+  valid = tfds.load(
+      name=dataset_name, split=eval_split,
+      as_dataset_kwargs={"shuffle_files": test_shuffle_files})
   keys = None
   if info.supervised_keys:
     keys = ([info.supervised_keys[0]], [info.supervised_keys[1]])
@@ -208,6 +218,7 @@ def _select_features(example, feature_list=None):
 def _train_and_eval_dataset_v1(problem_name, data_dir):
   """Return train and evaluation datasets, feature info and supervised keys."""
   from tensor2tensor import problems  # pylint: disable=g-import-not-at-top
+  assert not tf.executing_eagerly(), "tf.eager mode must be turned off."
   problem = problems.problem(problem_name)
   train_dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, data_dir)
   train_dataset = train_dataset.map(_select_features)

From f20c6a22b7cd17853a8e1c948aee7c1eda9787f4 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 22 Apr 2019 11:22:49 -0700
Subject: [PATCH 1931/2720]  * Do jax nan crashing without jitting so as to get
 a better stack trace.  * Errors out while computing gradients even with a
 constant small learning rate (1e-4) and cutting trajectories at 100 steps
 only (this early stopping is not in this change).  * This happens with both
 SGD and Adam.    * Any ideas why?    * It goes slightly further with SGD than
 Adam, but this is possibly random.

PiperOrigin-RevId: 244697209
---
 tensor2tensor/trax/rlax/ppo_main.py | 37 +++++++++++++++++------------
 1 file changed, 22 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 0738e758a..4f1054d72 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -23,6 +23,7 @@
 
 from absl import app
 from absl import flags
+import jax
 from jax.config import config
 from tensor2tensor.trax import stax
 from tensor2tensor.trax.rlax import ppo
@@ -55,23 +56,29 @@ def main(argv):
   if FLAGS.jax_debug_nans:
     config.update("jax_debug_nans", True)
 
-  optimizer_fun = functools.partial(ppo.optimizer_fun,
-                                    step_size=FLAGS.learning_rate)
+  def run_training_loop():
+    optimizer_fun = functools.partial(
+        ppo.optimizer_fun, step_size=FLAGS.learning_rate)
 
-  ppo.training_loop(
-      env_name=FLAGS.env_name,
-      epochs=FLAGS.epochs,
-      policy_net_fun=functools.partial(
-          ppo.policy_net, bottom_layers=common_stax_layers()),
-      value_net_fun=functools.partial(
-          ppo.value_net, bottom_layers=common_stax_layers()),
-      policy_optimizer_fun=optimizer_fun,
-      value_optimizer_fun=optimizer_fun,
-      batch_size=FLAGS.batch_size,
-      num_optimizer_steps=FLAGS.num_optimizer_steps,
-      boundary=FLAGS.boundary,
-      random_seed=FLAGS.random_seed)
+    ppo.training_loop(
+        env_name=FLAGS.env_name,
+        epochs=FLAGS.epochs,
+        policy_net_fun=functools.partial(
+            ppo.policy_net, bottom_layers=common_stax_layers()),
+        value_net_fun=functools.partial(
+            ppo.value_net, bottom_layers=common_stax_layers()),
+        policy_optimizer_fun=optimizer_fun,
+        value_optimizer_fun=optimizer_fun,
+        batch_size=FLAGS.batch_size,
+        num_optimizer_steps=FLAGS.num_optimizer_steps,
+        boundary=FLAGS.boundary,
+        random_seed=FLAGS.random_seed)
 
+  if FLAGS.jax_debug_nans:
+    with jax.disable_jit():
+      run_training_loop()
+  else:
+    run_training_loop()
 
 if __name__ == "__main__":
   app.run(main)

From aecd5d6fed6ca6de4e4988cb60fa66e4a9213a8c Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 22 Apr 2019 11:33:12 -0700
Subject: [PATCH 1932/2720] Rename stax_base to core and fork lax shape
 computations as they are pure python.

PiperOrigin-RevId: 244699352
---
 tensor2tensor/trax/stax/__init__.py           |  2 +-
 tensor2tensor/trax/stax/attention.py          | 15 +--
 .../trax/stax/{stax_base.py => core.py}       | 95 +++++++++++++++++--
 3 files changed, 95 insertions(+), 17 deletions(-)
 rename tensor2tensor/trax/stax/{stax_base.py => core.py} (72%)

diff --git a/tensor2tensor/trax/stax/__init__.py b/tensor2tensor/trax/stax/__init__.py
index 263fab640..ac85bf1da 100644
--- a/tensor2tensor/trax/stax/__init__.py
+++ b/tensor2tensor/trax/stax/__init__.py
@@ -24,4 +24,4 @@
 from tensor2tensor.trax.stax.attention import *
 from tensor2tensor.trax.stax.base import *
 from tensor2tensor.trax.stax.combinators import *
-from tensor2tensor.trax.stax.stax_base import *
+from tensor2tensor.trax.stax.core import *
diff --git a/tensor2tensor/trax/stax/attention.py b/tensor2tensor/trax/stax/attention.py
index 4b631a597..cbf42d701 100644
--- a/tensor2tensor/trax/stax/attention.py
+++ b/tensor2tensor/trax/stax/attention.py
@@ -24,7 +24,7 @@
 from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.stax import base
 from tensor2tensor.trax.stax import combinators
-from tensor2tensor.trax.stax import stax_base as stax
+from tensor2tensor.trax.stax import core
 
 
 @base.layer(output_shape=lambda shape, axis=-1: (1, shape[axis], shape[axis]))
@@ -38,8 +38,9 @@ def MakeTargetMask(target, pad=0):
   """Create an attention mask to hide padding and future words."""
   target_mask = (target != pad)[ :, np.newaxis, :]
   target_dtype = target_mask.dtype
-  target_mask = (
-      (target_mask & stax.causal_mask(target.shape[-1])).astype(target_dtype))
+  causal_mask = onp.tril(onp.ones((1, target.shape[-1], target.shape[-1]),
+                                  dtype=target_dtype), k=0)
+  target_mask = target_mask & causal_mask
   return np.expand_dims(target_mask, axis=1)
 
 
@@ -222,13 +223,13 @@ def MultiHeadedAttention(
   """
   return combinators.Serial(
       combinators.Parallel(
-          stax.Dense(feature_depth, W_init=stax.xavier_uniform()),
-          stax.Dense(feature_depth, W_init=stax.xavier_uniform()),
-          stax.Dense(feature_depth, W_init=stax.xavier_uniform()),
+          core.Dense(feature_depth, W_init=core.xavier_uniform()),
+          core.Dense(feature_depth, W_init=core.xavier_uniform()),
+          core.Dense(feature_depth, W_init=core.xavier_uniform()),
           combinators.Identity()
       ),
       PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
           feature_depth=feature_depth, num_heads=num_heads,
           dropout=dropout, mode=mode),
-      stax.Dense(feature_depth, W_init=stax.xavier_uniform()),
+      core.Dense(feature_depth, W_init=core.xavier_uniform()),
   )
diff --git a/tensor2tensor/trax/stax/stax_base.py b/tensor2tensor/trax/stax/core.py
similarity index 72%
rename from tensor2tensor/trax/stax/stax_base.py
rename to tensor2tensor/trax/stax/core.py
index 1d9115f0b..810eab332 100644
--- a/tensor2tensor/trax/stax/stax_base.py
+++ b/tensor2tensor/trax/stax/core.py
@@ -155,6 +155,24 @@ def new_parameters(self, input_shape, rng):
     return self._W_init(rng, (self._vocab_size, self._feature_depth))
 
 
+def padtype_to_pads(in_shape, window_shape, window_strides, padding):
+  """Convert padding string to list of pairs of pad values."""
+  padding = padding.upper()
+  if padding == 'SAME':
+    out_shape = onp.ceil(
+        onp.true_divide(in_shape, window_strides)).astype(int)
+    pad_sizes = [max((out_size - 1) * stride + window_shape - in_size, 0)
+                 for out_size, stride, window_shape, in_size
+                 in zip(out_shape, window_strides, window_shape, in_shape)]
+    return [(pad_size // 2, pad_size - pad_size // 2)
+            for pad_size in pad_sizes]
+  elif padding == 'VALID':
+    return [(0, 0)] * len(in_shape)
+  else:
+    msg = 'Unknown padding type: {}.'
+    raise TypeError(msg.format(padding))
+
+
 class Conv(base.Layer):
   """Layer constructor function for a general convolution layer."""
 
@@ -187,9 +205,64 @@ def _kernel_shape(self, input_shape):
             input_shape[self._lhs_spec.index('C')] if c == 'I' else
             next(filter_shape_iter) for c in self._rhs_spec]
 
+  def _conv_shape_tuple(self, lhs_shape, rhs_shape, strides, pads):
+    """Compute the shape of a conv given input shapes in canonical order."""
+    if isinstance(pads, str):
+      pads = padtype_to_pads(lhs_shape[2:], rhs_shape[2:], strides, pads)
+    if len(pads) != len(lhs_shape) - 2:
+      msg = 'Wrong number of explicit pads for conv: expected {}, got {}.'
+      raise TypeError(msg.format(len(lhs_shape) - 2, len(pads)))
+    lhs_padded = onp.add(lhs_shape[2:], onp.add(*zip(*pads)))
+    out_space = onp.floor_divide(
+        onp.subtract(lhs_padded, rhs_shape[2:]), strides) + 1
+    out_space = onp.maximum(0, out_space)
+    out_shape = (lhs_shape[0], rhs_shape[0]) + tuple(out_space)
+    return tuple(out_shape)
+
+  def _conv_general_permutations(self, dimension_numbers):
+    """Utility for convolution dimension permutations relative to Conv HLO."""
+    lhs_spec, rhs_spec, out_spec = dimension_numbers
+    lhs_char, rhs_char, out_char = ('N', 'C'), ('O', 'I'), ('N', 'C')
+    charpairs = (lhs_char, rhs_char, out_char)
+    for i, (a, b) in enumerate(charpairs):
+      if not (dimension_numbers[i].count(a) == 1 and
+              dimension_numbers[i].count(b) == 1):
+        msg = ('convolution dimension_numbers[{}] must contain the characters '
+               '"{}" and "{}" exatly once, got {}.')
+        raise TypeError(msg.format(i, a, b, dimension_numbers[i]))
+      if len(dimension_numbers[i]) != len(set(dimension_numbers[i])):
+        msg = ('convolution dimension_numbers[{}] cannot have duplicate '
+               'characters, got {}.')
+        raise TypeError(msg.format(i, dimension_numbers[i]))
+    if not (set(lhs_spec) - set(lhs_char) == set(rhs_spec) - set(rhs_char) ==
+            set(out_spec) - set(out_char)):
+      msg = ('convolution dimension_numbers elements must each have the same '
+             'set of spatial characters, got {}.')
+      raise TypeError(msg.format(dimension_numbers))
+
+    def getperm(spec, charpair):
+      spatial = (i for i, c in enumerate(spec) if c not in charpair)
+      if spec is not rhs_spec:
+        spatial = sorted(spatial, key=lambda i: rhs_spec.index(spec[i]))
+      return (spec.index(charpair[0]), spec.index(charpair[1])) + tuple(spatial)
+
+    lhs_perm, rhs_perm, out_perm = map(getperm, dimension_numbers, charpairs)
+    return lhs_perm, rhs_perm, out_perm
+
+  def _conv_general_shape_tuple(self, lhs_shape, rhs_shape, window_strides,
+                                padding, dimension_numbers):
+    """Generalized computation of conv shape."""
+    lhs_perm, rhs_perm, out_perm = self._conv_general_permutations(
+        dimension_numbers)
+    lhs_trans = onp.take(lhs_shape, lhs_perm)
+    rhs_trans = onp.take(rhs_shape, rhs_perm)
+    out_trans = self._conv_shape_tuple(
+        lhs_trans, rhs_trans, window_strides, padding)
+    return tuple(onp.take(out_trans, onp.argsort(out_perm)))
+
   def output_shape(self, input_shape):
     kernel_shape = self._kernel_shape(input_shape)
-    return lax.conv_general_shape_tuple(
+    return self._conv_general_shape_tuple(
         input_shape, kernel_shape,
         self._strides, self._padding, self._dimension_numbers)
 
@@ -264,18 +337,21 @@ def _pooling_output_shape(input_shape, pool_size=(2, 2),
                           strides=None, padding='VALID'):
   """Helper: compute the output shape for the pooling layer."""
   dims = (1,) + pool_size + (1,)  # NHWC
-  strides = strides or (1,) * len(pool_size)
-  strides = (1,) + strides + (1,)
-  return lax.reduce_window_shape_tuple(input_shape, dims, strides, padding)
+  spatial_strides = strides or (1,) * len(pool_size)
+  strides = (1,) + spatial_strides + (1,)
+  pads = padtype_to_pads(input_shape, dims, strides, padding)
+  operand_padded = onp.add(input_shape, onp.add(*zip(*pads)))
+  t = onp.floor_divide(onp.subtract(operand_padded, dims), strides) + 1
+  return tuple(t)
 
 
 def _pooling_general(inputs, reducer, init_val, rescaler=None,
                      pool_size=(2, 2), strides=None, padding='VALID'):
   """Helper: general pooling computation used in pooling layers later."""
-  strides = strides or (1,) * len(pool_size)
-  rescale = rescaler(pool_size, strides, padding) if rescaler else None
+  spatial_strides = strides or (1,) * len(pool_size)
+  rescale = rescaler(pool_size, spatial_strides, padding) if rescaler else None
   dims = (1,) + pool_size + (1,)  # NHWC
-  strides = (1,) + strides + (1,)
+  strides = (1,) + spatial_strides + (1,)
   out = lax.reduce_window(inputs, init_val, reducer, dims, strides, padding)
   return rescale(out, inputs) if rescale else out
 
@@ -294,10 +370,11 @@ def SumPool(params, x, pool_size=(2, 2), strides=None, padding='VALID', **kw):
                           strides=strides, padding=padding)
 
 
-def _normalize_by_window_size(dims, strides, padding):
+def _normalize_by_window_size(dims, spatial_strides, padding):
   def rescale(outputs, inputs):
     one = np.ones(inputs.shape[1:-1], dtype=inputs.dtype)
-    window_sizes = lax.reduce_window(one, 0., lax.add, dims, strides, padding)
+    window_sizes = lax.reduce_window(
+        one, 0., lax.add, dims, spatial_strides, padding)
     return outputs / window_sizes[..., np.newaxis]
   return rescale
 

From 809bc09c3e2d9369080c1ec4e5c8de6b49b5d6f9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 22 Apr 2019 16:17:19 -0700
Subject: [PATCH 1933/2720]  - Add the option to terminate collect at a
 specified maximum time-step.  - Additional logging in ppo.py for min/max/avg
 rewards.

PiperOrigin-RevId: 244756017
---
 tensor2tensor/trax/rlax/ppo.py      | 20 +++++++++----
 tensor2tensor/trax/rlax/ppo_main.py |  4 +++
 tensor2tensor/trax/rlax/ppo_test.py | 45 ++++++++++++++++++++++++++++-
 3 files changed, 62 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 2de65b959..365b6a1b0 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -161,6 +161,7 @@ def collect_trajectories(env,
                          policy_net_params,
                          num_trajectories=1,
                          policy="greedy",
+                         max_timestep=None,
                          epsilon=0.1):
   """Collect trajectories with the given policy net and behaviour."""
   trajectories = []
@@ -176,7 +177,10 @@ def collect_trajectories(env,
     # getting added to it, making it eventually (1, T+1) + OBS
     observation_history = observation[np.newaxis, np.newaxis, :]
 
-    while not done:
+    # Run either till we're done OR if max_timestep is defined only till that
+    # timestep.
+    while ((not done) and
+           (not max_timestep or observation_history.shape[1] < max_timestep)):
       # Run the policy, to pick an action, shape is (1, t, A) because
       # observation_history is shaped (1, t) + OBS
       predictions = policy_net_apply(policy_net_params, observation_history)
@@ -230,8 +234,9 @@ def collect_trajectories(env,
       rewards.append(reward)
       actions.append(action)
 
-    # This means we are done
-    assert done
+    # This means we are done we're been terminated early.
+    assert done or (
+        max_timestep and max_timestep >= observation_history.shape[1])
     # observation_history is (1, T+1) + OBS, lets squeeze out the batch dim.
     observation_history = np.squeeze(observation_history, axis=0)
     trajectories.append(
@@ -650,6 +655,7 @@ def training_loop(
     num_optimizer_steps=NUM_OPTIMIZER_STEPS,
     print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
     boundary=20,
+    max_timestep=None,
     random_seed=None):
   """Runs the training loop for PPO, with fixed policy and value nets."""
   jax_rng_key = trax.get_random_number_generator_and_set_seed(random_seed)
@@ -696,6 +702,7 @@ def training_loop(
         policy_net_params,
         num_trajectories=batch_size,
         policy=POLICY,
+        max_timestep=max_timestep,
         epsilon=(10.0 / (i + 10.0)))  # this is a different epsilon.
 
     avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
@@ -849,9 +856,10 @@ def training_loop(
     value_net_params = trax_opt.get_params(value_opt_state)
 
     logging.info(
-        "Epoch [% 6d], average reward [%10.2f], ppo loss [%10.2f], "
-        "value loss [%10.2f], took [%10.2f msec]", i, avg_reward, new_ppo_loss,
-        new_value_loss, get_time(t0))
+        "Epoch [% 6d], Reward[min, max, avg] [%10.2f,%10.2f,%10.2f], "
+        "ppo loss [%10.2f], value loss [%10.2f], took [%10.2f msec]",
+        i, min_reward, max_reward, avg_reward, new_ppo_loss, new_value_loss,
+        get_time(t0))
 
   logging.vlog(1, "value_losses: %s", np.stack(value_losses))
   logging.vlog(1, "ppo_objective: %s", np.stack(ppo_objective))
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 4f1054d72..fa0028e7d 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -38,6 +38,9 @@
 flags.DEFINE_integer("num_optimizer_steps", 100, "Number of optimizer steps.")
 flags.DEFINE_integer("boundary", 20,
                      "We pad trajectories at integer multiples of this number.")
+flags.DEFINE_integer("max_timestep", None,
+                     "If set to an integer, maximum number of time-steps in a "
+                     "trajectory.")
 flags.DEFINE_float("learning_rate", 1e-3, "Learning rate.")
 flags.DEFINE_boolean("jax_debug_nans", False,
                      "Setting to true will help to debug nans.")
@@ -72,6 +75,7 @@ def run_training_loop():
         batch_size=FLAGS.batch_size,
         num_optimizer_steps=FLAGS.num_optimizer_steps,
         boundary=FLAGS.boundary,
+        max_timestep=FLAGS.max_timestep,
         random_seed=FLAGS.random_seed)
 
   if FLAGS.jax_debug_nans:
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 3f6d09999..6f7f0f132 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -106,7 +106,7 @@ def test_collect_trajectories(self):
         # step dimensions.
         [stax.Flatten(num_axis_to_keep=2)])
 
-    # We'll get done at time-step #10, starting from 0, therefore in 11 steps.
+    # We'll get done at time-step #5, starting from 0, therefore in 6 steps.
     done_time_step = 5
     env = fake_env.FakeEnv(
         observation_shape, num_actions, done_time_step=done_time_step)
@@ -130,6 +130,49 @@ def test_collect_trajectories(self):
       self.assertEqual((done_time_step + 1,), actions.shape)
       self.assertEqual((done_time_step + 1,), rewards.shape)
 
+  def test_collect_trajectories_max_timestep(self):
+    observation_shape = (2, 3, 4)
+    num_actions = 2
+    policy_params, policy_apply = ppo.policy_net(
+        self.rng_key,
+        (-1, -1) + observation_shape,
+        num_actions,
+        # flatten except batch and time
+        # step dimensions.
+        [stax.Flatten(num_axis_to_keep=2)])
+
+    # We'll get done at time-step #5, starting from 0, therefore in 6 steps.
+    done_time_step = 5
+    env = fake_env.FakeEnv(
+        observation_shape, num_actions, done_time_step=done_time_step)
+
+    num_trajectories = 5
+
+    # Let's collect trajectories only till `max_timestep`.
+    max_timestep = 3
+
+    # we're testing when we early stop the trajectory.
+    assert max_timestep < done_time_step
+
+    trajectories = ppo.collect_trajectories(
+        env,
+        policy_apply,
+        policy_params,
+        num_trajectories,
+        policy="categorical-sampling",
+        max_timestep=max_timestep)
+
+    # Number of trajectories is as expected.
+    self.assertEqual(num_trajectories, len(trajectories))
+
+    # Shapes of observations, actions and rewards are as expected.
+    for observations, actions, rewards in trajectories:
+      # observations are one more in number than rewards or actions.
+      self.assertEqual((max_timestep,) + observation_shape,
+                       observations.shape)
+      self.assertEqual((max_timestep - 1,), actions.shape)
+      self.assertEqual((max_timestep - 1,), rewards.shape)
+
   def test_pad_trajectories(self):
     observation_shape = (2, 3, 4)
     trajectories = []

From 698088b09a963c807ac27abb79c26e3edfd21727 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 22 Apr 2019 17:38:30 -0700
Subject: [PATCH 1934/2720] TRAX: Change stax API to net.params(shape, rng) and
 net(x, params) instead of (apply, init) = net.

PiperOrigin-RevId: 244769076
---
 tensor2tensor/trax/models/transformer.py |  18 ++--
 tensor2tensor/trax/rlax/ppo.py           |  30 +++---
 tensor2tensor/trax/rlax/ppo_test.py      |   8 +-
 tensor2tensor/trax/stax/README.md        |  19 +++-
 tensor2tensor/trax/stax/attention.py     |  25 ++---
 tensor2tensor/trax/stax/base.py          |  66 ++++++++-----
 tensor2tensor/trax/stax/base_test.py     |  23 ++---
 tensor2tensor/trax/stax/combinators.py   |  43 ++++-----
 tensor2tensor/trax/stax/core.py          | 115 ++++++++++++-----------
 tensor2tensor/trax/trax.py               |  33 ++++---
 tensor2tensor/trax/trax_test.py          |   3 +-
 11 files changed, 198 insertions(+), 185 deletions(-)

diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index eee2ac4dd..026d54a4e 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -28,10 +28,12 @@ def ResidualFeedForward(feature_depth,
   """Residual feed-forward layer with normalization at start."""
   return stax.Residual(
       stax.LayerNorm(),
-      stax.Dense(feedforward_depth, W_init=stax.xavier_uniform()),
+      stax.Dense(feedforward_depth,
+                 kernel_initializer=stax.XavierUniformInitializer()),
       stax.Relu(),
       stax.Dropout(rate=dropout, mode=mode),
-      stax.Dense(feature_depth, W_init=stax.xavier_uniform()),
+      stax.Dense(feature_depth,
+                 kernel_initializer=stax.XavierUniformInitializer()),
       stax.Dropout(rate=dropout, mode=mode)
   )
 
@@ -156,11 +158,12 @@ def TransformerLM(vocab_size,
       stax.Embedding(feature_depth, vocab_size),
       stax.Dropout(rate=dropout, mode=mode),
       stax.PositionalEncoding(max_len=max_len),
-      stax.Serial([DecoderLayer(feature_depth, feedforward_depth, num_heads,
-                                dropout, mode)
-                   for _ in range(num_layers)]),
+      stax.Serial(*[DecoderLayer(feature_depth, feedforward_depth, num_heads,
+                                 dropout, mode)
+                    for _ in range(num_layers)]),
       stax.LayerNorm(),
-      stax.Dense(vocab_size, W_init=stax.xavier_uniform()),
+      stax.Dense(vocab_size,
+                 kernel_initializer=stax.XavierUniformInitializer()),
       stax.LogSoftmax()
   )
 
@@ -308,7 +311,8 @@ def transformer(source, target, source_mask, target_mask, memory_mask):  # pylin
   def Generator(encoded_target):
     return stax.Serial(
         encoded_target,
-        stax.Dense(target_vocab_size, W_init=stax.xavier_uniform()),
+        stax.Dense(target_vocab_size,
+                   kernel_initializer=stax.XavierUniformInitializer()),
         stax.LogSoftmax
     )
 
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 365b6a1b0..5c724d4cf 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -80,11 +80,8 @@ def policy_net(rng_key,
   if bottom_layers is None:
     bottom_layers = []
   bottom_layers.extend([stax.Dense(num_actions), stax.Softmax()])
-
-  net_init, net_apply = stax.Serial(bottom_layers)
-
-  _, net_params = net_init(rng_key, batch_observations_shape)
-  return net_params, net_apply
+  net = stax.Serial(*bottom_layers)
+  return net.initialize(batch_observations_shape, rng_key), net
 
 
 def value_net(rng_key,
@@ -99,11 +96,8 @@ def value_net(rng_key,
   bottom_layers.extend([
       stax.Dense(1),
   ])
-
-  net_init, net_apply = stax.Serial(bottom_layers)
-
-  _, net_params = net_init(rng_key, batch_observations_shape)
-  return net_params, net_apply
+  net = stax.Serial(*bottom_layers)
+  return net.initialize(batch_observations_shape, rng_key), net
 
 
 def policy_and_value_net(rng_key,
@@ -124,10 +118,8 @@ def policy_and_value_net(rng_key,
       stax.Dense(1)
   )])
 
-  net_init, net_apply = stax.Serial(layers)
-
-  _, net_params = net_init(rng_key, batch_observations_shape)
-  return net_params, net_apply
+  net = stax.Serial(*layers)
+  return net.initialize(batch_observations_shape, rng_key), net
 
 
 def optimizer_fun(net_params, step_size=1e-3):
@@ -183,7 +175,7 @@ def collect_trajectories(env,
            (not max_timestep or observation_history.shape[1] < max_timestep)):
       # Run the policy, to pick an action, shape is (1, t, A) because
       # observation_history is shaped (1, t) + OBS
-      predictions = policy_net_apply(policy_net_params, observation_history)
+      predictions = policy_net_apply(observation_history, policy_net_params)
 
       # We need the predictions for the last time-step, so squeeze the batch
       # dimension and take the last time-step.
@@ -419,7 +411,7 @@ def value_loss(value_net_apply,
 
   r2g = rewards_to_go(rewards, reward_mask, gamma=gamma)  # (B, T)
   # NOTE: observations is (B, T+1) + OBS, value_prediction is (B, T+1, 1)
-  value_prediction = value_net_apply(value_net_params, observations)
+  value_prediction = value_net_apply(observations, value_net_params)
   assert (B, T + 1, 1) == value_prediction.shape
   value_prediction = np.squeeze(value_prediction, axis=2)  # (B, T+1)
   value_prediction = value_prediction[:, :-1] * reward_mask  # (B, T)
@@ -549,7 +541,7 @@ def ppo_loss(policy_net_apply,
              epsilon=0.2):
   """PPO objective, with an eventual minus sign."""
   # (B, T+1, 1)
-  predicted_values = value_net_apply(value_net_params, padded_observations)
+  predicted_values = value_net_apply(padded_observations, value_net_params)
 
   # (B, T)
   td_deltas = deltas(
@@ -563,8 +555,8 @@ def ppo_loss(policy_net_apply,
       td_deltas, reward_mask, lambda_=lambda_, gamma=gamma)
 
   # probab_actions_{old,new} are both (B, T, A)
-  probab_actions_old = policy_net_apply(old_policy_params, padded_observations)
-  probab_actions_new = policy_net_apply(new_policy_params, padded_observations)
+  probab_actions_old = policy_net_apply(padded_observations, old_policy_params)
+  probab_actions_new = policy_net_apply(padded_observations, new_policy_params)
 
   # (B, T)
   ratios = compute_probab_ratios(probab_actions_old, probab_actions_new,
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 6f7f0f132..197297ef6 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -52,7 +52,7 @@ def test_policy_net(self):
         size=(batch, time_steps) + observation_shape)
 
     # Apply the policy net on observations
-    policy_output = policy_apply(policy_params, batch_of_observations)
+    policy_output = policy_apply(batch_of_observations, policy_params)
 
     # Verify certain expectations on the output.
     self.assertEqual((batch, time_steps, num_actions), policy_output.shape)
@@ -72,7 +72,7 @@ def test_value_net(self):
     time_steps = 10
     batch_of_observations = np.random.uniform(
         size=(batch, time_steps) + observation_shape)
-    value_output = value_apply(value_params, batch_of_observations)
+    value_output = value_apply(batch_of_observations, value_params)
 
     # NOTE: The extra dimension at the end because of Dense(1).
     self.assertEqual((batch, time_steps, 1), value_output.shape)
@@ -88,7 +88,7 @@ def test_policy_and_value_net(self):
     time_steps = 10
     batch_of_observations = np.random.uniform(
         size=(batch, time_steps) + observation_shape)
-    pnv_output = pnv_apply(pnv_params, batch_of_observations)
+    pnv_output = pnv_apply(batch_of_observations, pnv_params)
 
     # Output is a list, first is probab of actions and the next is value output.
     self.assertEqual(2, len(pnv_output))
@@ -294,7 +294,7 @@ def test_value_loss(self):
     observation_shape = (210, 160, 3)  # atari pong
     random_observations = np.random.uniform(size=(B, T + 1) + observation_shape)
 
-    def value_net_apply(params, observations):
+    def value_net_apply(observations, params):
       del params
       # pylint: disable=invalid-name
       B, T_p_1, OBS = (observations.shape[0], observations.shape[1],
diff --git a/tensor2tensor/trax/stax/README.md b/tensor2tensor/trax/stax/README.md
index 7c96d6e2b..679552bbc 100644
--- a/tensor2tensor/trax/stax/README.md
+++ b/tensor2tensor/trax/stax/README.md
@@ -21,10 +21,21 @@ The base layer class wraps these functions and provides initialization
 and call functions to be used as follows.
 
 ```python
-input = np.zeros(10)
 layer = MyLayer()
-params = layer.initialize()
-output = layer(params, input)
+x = np.zeros(10)
+params = layer.initialize(x.shape)
+output = layer(x, params)
+```
+
+## Decorator
+
+To create simple layers, especially ones without parameters and where
+the output shape is the same as the input shape, use the layer decorator.
+
+```python
+@base.layer()
+def Relu(x, **unused_kwargs):
+  return np.maximum(x, 0.)
 ```
 
 ## Parameter sharing
@@ -36,6 +47,8 @@ standard_mlp = layers.Serial(layers.Dense(10), layers.Dense(10))
 layer = Dense(10)
 shared_parameters_mlp = layers.Serial(layer, layer)
 ```
+For this reason, if you call `layer.initialize(...)` for the second time
+on an already initialized layer, it will return `()`.
 
 ## Core layers
 
diff --git a/tensor2tensor/trax/stax/attention.py b/tensor2tensor/trax/stax/attention.py
index cbf42d701..ae3188238 100644
--- a/tensor2tensor/trax/stax/attention.py
+++ b/tensor2tensor/trax/stax/attention.py
@@ -28,7 +28,7 @@
 
 
 @base.layer(output_shape=lambda shape, axis=-1: (1, shape[axis], shape[axis]))
-def CausalMask(params, x, axis=-1, **kwargs):
+def CausalMask(x, params, axis=-1, **kwargs):
   del params, kwargs
   size = x.shape[axis]
   return onp.tril(onp.ones((1, size, size), dtype=x.dtype), k=0)
@@ -79,7 +79,7 @@ def _layer_norm_new_params(input_shape, rng, epsilon=1e-6):  # pylint: disable=i
 
 
 @base.layer(new_parameters=_layer_norm_new_params)
-def LayerNorm(params, x, epsilon=1e-6, **unused_kwargs):
+def LayerNorm(x, params, epsilon=1e-6, **unused_kwargs):
   (scale, bias) = params
   mean = np.mean(x, axis=-1, keepdims=True)
   variance = np.mean((x - mean)**2, axis=-1, keepdims=True)
@@ -102,7 +102,7 @@ def _positional_encoding_new_params(input_shape, rng, max_len=2048):  # pylint:
 
 
 @base.layer(new_parameters=_positional_encoding_new_params)
-def PositionalEncoding(params, x, **unused_kwargs):
+def PositionalEncoding(x, params, **unused_kwargs):
   """Implements bare positional encoding."""
   symbol_size = np.shape(x)[1]
   return x + params[:, :symbol_size]
@@ -170,14 +170,13 @@ def _multihead_attention_output_shape(  # pylint: disable=invalid-name
 
 
 @base.layer(output_shape=_multihead_attention_output_shape)
-def PureMultiHeadedAttention(
-    params, x, feature_depth=None, num_heads=8, dropout=0.0, mode='train',
-    **kwargs):
+def PureMultiHeadedAttention(x, params, feature_depth=None,
+                             num_heads=8, dropout=0.0, mode='train', **kwargs):
   """Pure transformer-style multi-headed attention.
 
   Args:
-    params: parameters (none)
     x: inputs (q, k, v, mask)
+    params: parameters (none)
     feature_depth: int:  depth of embedding
     num_heads: int: number of attention heads
     dropout: float: dropout rate
@@ -223,13 +222,17 @@ def MultiHeadedAttention(
   """
   return combinators.Serial(
       combinators.Parallel(
-          core.Dense(feature_depth, W_init=core.xavier_uniform()),
-          core.Dense(feature_depth, W_init=core.xavier_uniform()),
-          core.Dense(feature_depth, W_init=core.xavier_uniform()),
+          core.Dense(feature_depth,
+                     kernel_initializer=core.XavierUniformInitializer()),
+          core.Dense(feature_depth,
+                     kernel_initializer=core.XavierUniformInitializer()),
+          core.Dense(feature_depth,
+                     kernel_initializer=core.XavierUniformInitializer()),
           combinators.Identity()
       ),
       PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
           feature_depth=feature_depth, num_heads=num_heads,
           dropout=dropout, mode=mode),
-      core.Dense(feature_depth, W_init=core.xavier_uniform()),
+      core.Dense(feature_depth,
+                 kernel_initializer=core.XavierUniformInitializer()),
   )
diff --git a/tensor2tensor/trax/stax/base.py b/tensor2tensor/trax/stax/base.py
index 71228036a..f2155f808 100644
--- a/tensor2tensor/trax/stax/base.py
+++ b/tensor2tensor/trax/stax/base.py
@@ -33,15 +33,11 @@ def __init__(self, **kwargs):
     self._first_init = True
     # Cache parameters here, defaults empty params (we use () for that).
     self._params = ()  # cached parameters
-    # Caller field storing file name and line number where the caller class
-    # was created. Since most layers inherit directly from this class, they
-    # call this init (stack 0) in their init (stack 1) and we want the caller
-    # that created them, so we take stack position 2 here.
-    # TODO(lukaszkaiser): this will break with more inheritance, correct.
-    self._caller = inspect.getframeinfo(inspect.stack()[2][0])
-
-  def call(self, params, inputs, **kwargs):
-    """Call this layer using the given parameters on the given inputs."""
+    # Caller field storing info on where the caller class was created.
+    self._caller = _find_frame(inspect.stack())
+
+  def call(self, x, params=(), **kwargs):
+    """Call this layer in input x using the given parameters."""
     raise NotImplementedError
 
   def output_shape(self, input_shape):
@@ -59,7 +55,7 @@ def output_shape(self, input_shape):
     raise NotImplementedError
 
   def new_parameters(self, input_shape, rng):
-    """Initialize parameters given input shape and return with output shape.
+    """Create new parameters for the layer given an input shape and rng.
 
     Note that all arguments and return values can be tuples or dictionaries
     or arbitraty nested structures composed of tuples and dictionaries.
@@ -75,37 +71,46 @@ def new_parameters(self, input_shape, rng):
 
   # End of subclassing interface, all functions below are internal.
 
-  def _init_fun(self, rng, input_shape):
-    """Internal modification of init_fun, saves variables."""
-    out_shape = self.output_shape(input_shape)
+  def initialize(self, input_shape, rng):
+    """Initialize the layer given an input shape and rng.
+
+    Returns new_parameters(input_shape, rng) on the first call and () on any
+    subsequent call, as the layer is already initialized. This is used for
+    networks that share parameters, so the layer only produces them once.
+
+    Note that all arguments and return values can be tuples or dictionaries
+    or arbitraty nested structures composed of tuples and dictionaries.
+
+    Args:
+      input_shape: a tuple representing the shape of the input.
+      rng: random number generator.
 
+    Returns:
+      Newly created parameters on the first call and () on all subsequent calls.
+    """
     # Re-using this layer, no new parameters.
     if not self._first_init:
-      return out_shape, ()
+      return ()
 
     # First call of this layer, create parameters.
     self._first_init = False
     self._params = self.new_parameters(input_shape, rng)
-    return out_shape, self._params
+    return self._params
 
-  def __call__(self, params, inputs, **kwargs):
+  def __call__(self, x, params=(), **kwargs):
     try:
       # If params are nothing, we may be reusing this layer.
       # Use the cached parameters to calculate the value.
       # Note: to make sure jit tracers can decide this branch in python we
       #   use "params is ()" instead of, e.g., "not params" or "params == ()".
       if params is ():  # pylint: disable=literal-comparison
-        return self.call(self._params, inputs, **kwargs)
+        return self.call(x, params=self._params, **kwargs)
       # In this case, we're called for the first time: cache parameters.
       self._params = params
-      return self.call(params, inputs, **kwargs)
+      return self.call(x, params=params, **kwargs)
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback()
-      raise LayerError(name, self._caller, shapes(inputs), trace)
-
-  # when unpacking this (init, apply) pair we return the wrapped funs
-  def __iter__(self):
-    return iter((self._init_fun, self.__call__))
+      raise LayerError(name, self._caller, shapes(x), trace)
 
 
 class LayerError(Exception):
@@ -153,6 +158,17 @@ def shape(x):
   return nested_map(x, shape)
 
 
+def _find_frame(stack, start=0):
+  """Find the frame with the caller on the stack."""
+  # We want to find the first place where the layer was called
+  # that is *not* an __init__ function of an inheriting layer.
+  frame = inspect.getframeinfo(stack[start][0])
+  # If we are in an init, move on.
+  if frame.function == '__init__':
+    return _find_frame(stack, start + 1)
+  return frame
+
+
 def _shorten_file_path(line):
   """Shorten file path in error lines for more readable tracebacks."""
   start = line.lower().find('file')
@@ -207,13 +223,13 @@ def new_parameters_fun(self, input_shape, rng):
       kwargs = self._init_kwargs  # pylint: disable=protected-access
       return new_parameters(input_shape, rng, **kwargs)
 
-    def call_fun(self, params, inputs, **kwargs):
+    def call_fun(self, x, params=(), **kwargs):
       """The call function of the created class, derived from call."""
       # Merge on-call kwargs with class-kwargs.
       call_kwargs = kwargs.copy()
       call_kwargs.update(self._init_kwargs)  # pylint: disable=protected-access
       # Call with the merged kwargs.
-      return call(params, inputs, **call_kwargs)
+      return call(x, params=params, **call_kwargs)
 
     # Set doc for python help.
     call_fun.__doc__ = call.__doc__
diff --git a/tensor2tensor/trax/stax/base_test.py b/tensor2tensor/trax/stax/base_test.py
index 5d886635f..07e02e575 100644
--- a/tensor2tensor/trax/stax/base_test.py
+++ b/tensor2tensor/trax/stax/base_test.py
@@ -33,18 +33,18 @@ def random_inputs(rng, input_shape):
     raise TypeError(type(input_shape))
 
 
-def check_shape_agreement(test_case, init_fun, apply_fun, input_shape):
+def check_shape_agreement(test_case, layer, input_shape):
   rng_key1, rng_key2 = random.split(random.get_prng(0))
-  result_shape, params = init_fun(rng_key1, input_shape)
+  result_shape = layer.output_shape(input_shape)
+  params = layer.initialize(input_shape, rng_key1)
   inputs = random_inputs(onp.random.RandomState(0), input_shape)
-  result = apply_fun(params, inputs, rng=rng_key2)
+  result = layer(inputs, params, rng=rng_key2)
   test_case.assertEqual(result.shape, result_shape)
   return result_shape
 
 
 def check_staxlayer(test_case, staxlayer, input_shape):
-  init_fun, apply_fun = staxlayer
-  return check_shape_agreement(test_case, init_fun, apply_fun, input_shape)
+  return check_shape_agreement(test_case, staxlayer, input_shape)
 
 
 class SlaxTest(absltest.TestCase):
@@ -75,12 +75,9 @@ def test_flatten_n(self):
       check_staxlayer(self, stax.Flatten(num_axis_to_keep=6), input_shape)
 
   def test_div(self):
-    init_fun, apply_fun = stax.Div(divisor=2.0)
+    layer = stax.Div(divisor=2.0)
     input_np = onp.array([[1, 2, 3], [4, 5, 6]], dtype=onp.float32)
-    input_shape = input_np.shape
-    rng = random.get_prng(0)
-    _, _ = init_fun(rng, input_shape)
-    output_np = apply_fun(None, input_np)
+    output_np = layer(input_np)
     # absltest doesn't have ndarray equalities.
     expected_output_np = input_np / 2.0
     self.assertAlmostEqual(
@@ -92,11 +89,9 @@ def test_dense_param_sharing(self):
     model1 = stax.Serial(stax.Dense(32), stax.Dense(32))
     layer = stax.Dense(32)
     model2 = stax.Serial(layer, layer)
-    init_fun1, _ = model1
-    init_fun2, _ = model2
     rng = random.get_prng(0)
-    _, params1 = init_fun1(rng, [-1, 32])
-    _, params2 = init_fun2(rng, [-1, 32])
+    params1 = model1.initialize((-1, 32), rng)
+    params2 = model2.initialize((-1, 32), rng)
     # The first parameters have 2 kernels of size (32, 32).
     self.assertEqual((32, 32), params1[0][0].shape)
     self.assertEqual((32, 32), params1[1][0].shape)
diff --git a/tensor2tensor/trax/stax/combinators.py b/tensor2tensor/trax/stax/combinators.py
index 044cbca95..47885ca3c 100644
--- a/tensor2tensor/trax/stax/combinators.py
+++ b/tensor2tensor/trax/stax/combinators.py
@@ -28,21 +28,17 @@ class Serial(base.Layer):
 
   def __init__(self, *layers):
     super(Serial, self).__init__()
-    # If called with one list argument, treat it as layers.
-    if len(layers) == 1 and isinstance(layers[0], list):
-      layers = layers[0]
     self._nlayers = len(layers)
     self._layers = layers
-    self._init_funs, self._apply_funs = zip(*layers)
 
-  def call(self, params, inputs, **kwargs):
+  def call(self, x, params=(), **kwargs):
     rng = kwargs.pop('rng', None)
     rngs = (None,) * self._nlayers
     if rng is not None:
       rngs = backend.random.split(rng, self._nlayers)
-    for fun, param, rng in zip(self._apply_funs, params, rngs):
-      inputs = fun(param, inputs, rng=rng, **kwargs)
-    return inputs
+    for layer, p, rng in zip(self._layers, params, rngs):
+      x = layer(x, p, rng=rng, **kwargs)
+    return x
 
   def output_shape(self, input_shape):
     cur_shape = input_shape
@@ -52,28 +48,28 @@ def output_shape(self, input_shape):
 
   def new_parameters(self, input_shape, rng):
     params = []
-    for init_fun in self._init_funs:
+    cur_shape = input_shape
+    for layer in self._layers:
       rng, layer_rng = backend.random.split(rng)
-      input_shape, param = init_fun(layer_rng, input_shape)
+      param = layer.initialize(cur_shape, layer_rng)
+      cur_shape = layer.output_shape(cur_shape)
       params.append(param)
     return params
 
 
 @base.layer()
-def Identity(params, x, **kwargs):
-  del params, kwargs
+def Identity(x, **unused_kwargs):
   return x
 
 
 @base.layer(output_shape=lambda input_shape, size=2: [input_shape] * size)
-def FanOut(params, x, size=2, **kwargs):
+def FanOut(x, params, size=2, **kwargs):
   del params, kwargs
   return [x] * size
 
 
 @base.layer(output_shape=lambda input_shape_list: input_shape_list[0])
-def FanInSum(params, x, **kwargs):
-  del params, kwargs
+def FanInSum(x, **unused_kwargs):
   return sum(x)  # Here x is a list of tensors of the same shape, we add them.
 
 
@@ -86,7 +82,7 @@ def _fan_in_concat_shape(input_shape, axis=-1):  # pylint: disable=invalid-name
 
 
 @base.layer(output_shape=_fan_in_concat_shape)
-def FanInConcat(params, x, axis=-1, **kwargs):
+def FanInConcat(x, params, axis=-1, **kwargs):
   del params, kwargs
   return backend.numpy.concatenate(x, axis)
 
@@ -109,20 +105,16 @@ class Parallel(base.Layer):
 
   def __init__(self, *layers):
     super(Parallel, self).__init__()
-    # If called with one list argument, treat it as layers.
-    if len(layers) == 1 and isinstance(layers[0], list):
-      layers = layers[0]
     self._nlayers = len(layers)
     self._layers = layers
-    self._init_funs, self._apply_funs = zip(*layers)
 
-  def call(self, params, inputs, **kwargs):
+  def call(self, inputs, params=(), **kwargs):
     rng = kwargs.pop('rng', None)
     rngs = (None,) * self._nlayers
     if rng is not None:
       rngs = backend.random.split(rng, self._nlayers)
-    return [f(p, x, rng=r, **kwargs)
-            for f, p, x, r in zip(self._apply_funs, params, inputs, rngs)]
+    return [layer(x, params=p, rng=r, **kwargs)
+            for layer, x, p, r in zip(self._layers, inputs, params, rngs)]
 
   def output_shape(self, input_shapes):
     return tuple([layer.output_shape(shape)
@@ -130,9 +122,8 @@ def output_shape(self, input_shapes):
 
   def new_parameters(self, input_shape, rng):
     rngs = backend.random.split(rng, self._nlayers)
-    _, p = zip(*[init(rng, shape) for init, rng, shape
-                 in zip(self._init_funs, rngs, input_shape)])
-    return p
+    return [layer.initialize(shape, rng) for layer, shape, rng
+            in zip(self._layers, input_shape, rngs)]
 
 
 def Residual(*layers, **kwargs):
diff --git a/tensor2tensor/trax/stax/core.py b/tensor2tensor/trax/stax/core.py
index 810eab332..e3924fa63 100644
--- a/tensor2tensor/trax/stax/core.py
+++ b/tensor2tensor/trax/stax/core.py
@@ -39,16 +39,16 @@
 # Initializers.
 
 
-def randn(stddev=1e-2):
+def RandomNormalInitializer(stddev=1e-2):
   """An initializer function for random normal coefficients."""
-  def init(rng, shape):
+  def init(shape, rng):
     return (stddev * backend.random.normal(rng, shape)).astype('float32')
   return init
 
 
-def glorot(out_dim=0, in_dim=1, scale=onp.sqrt(2)):
+def GlorotNormalInitializer(out_dim=0, in_dim=1, scale=onp.sqrt(2)):
   """An initializer function for random Glorot-scaled coefficients."""
-  def init(rng, shape):
+  def init(shape, rng):
     fan_in, fan_out = shape[in_dim], shape[out_dim]
     size = onp.prod(onp.delete(shape, [in_dim, out_dim]))
     std = scale / np.sqrt((fan_in + fan_out) / 2. * size)
@@ -56,9 +56,9 @@ def init(rng, shape):
   return init
 
 
-def xavier_uniform(out_dim=0, in_dim=1):
+def XavierUniformInitializer(out_dim=0, in_dim=1):
   """An initializer function for random uniform xavier-scaled coefficients."""
-  def init(rng, shape):
+  def init(shape, rng):
     fan_in, fan_out = shape[in_dim], shape[out_dim]
     std = np.sqrt(2.0 / (fan_in + fan_out))
     a = np.sqrt(3.0) * std
@@ -75,84 +75,84 @@ def one_hot(x, size, dtype=np.float32):
 
 
 @base.layer()
-def Relu(params, x, **kwargs):
-  del params, kwargs
+def Relu(x, **unused_kwargs):
   return np.maximum(x, 0.)
 
 
 @base.layer()
-def Tanh(params, x, **kwargs):
-  del params, kwargs
+def Tanh(x, **unused_kwargs):
   return np.tanh(x)
 
 
 @base.layer()
-def Exp(params, x, **kwargs):
-  del params, kwargs
+def Exp(x, **unused_kwargs):
   return np.exp(x)
 
 
 @base.layer()
-def LogSoftmax(params, x, axis=-1, **kwargs):
+def LogSoftmax(x, params, axis=-1, **kwargs):
   """Apply log softmax to x: log-normalize along the given axis."""
   del params, kwargs
   return x - backend.logsumexp(x, axis, keepdims=True)
 
 
 @base.layer()
-def Softmax(params, x, axis=-1, **kwargs):
+def Softmax(x, params, axis=-1, **kwargs):
   """Apply softmax to x: exponentiate and normalize along the given axis."""
   del params, kwargs
   return np.exp(x - backend.logsumexp(x, axis, keepdims=True))
 
 
 @base.layer()
-def Softplus(params, x, **kwargs):
-  del params, kwargs
+def Softplus(x, **unused_kwargs):
   return np.logaddexp(x, 0.)
 
 
 class Dense(base.Layer):
   """Layer constructor function for a dense (fully-connected) layer."""
 
-  def __init__(self, out_dim, W_init=glorot(), b_init=randn()):
+  def __init__(self, units,
+               kernel_initializer=GlorotNormalInitializer(),
+               bias_initializer=RandomNormalInitializer(1e-6)):
     super(Dense, self).__init__()
-    self._out_dim = out_dim
-    self._W_init = W_init
-    self._b_init = b_init
+    self._units = units
+    self._kernel_initializer = kernel_initializer
+    self._bias_initializer = bias_initializer
 
-  def call(self, params, inputs, **kwargs):
+  def call(self, x, params, **kwargs):
     del kwargs
     w, b = params
-    return np.dot(inputs, w) + b
+    return np.dot(x, w) + b
 
   def output_shape(self, input_shape):
-    return tuple(input_shape[:-1]) + (self._out_dim,)
+    return tuple(input_shape[:-1]) + (self._units,)
 
   def new_parameters(self, input_shape, rng):
-    w = self._W_init(rng, (input_shape[-1], self._out_dim))
-    b = self._b_init(rng, (self._out_dim,))
+    w = self._kernel_initializer((input_shape[-1], self._units), rng)
+    b = self._bias_initializer((self._units,), rng)
     return (w, b)
 
 
 class Embedding(base.Layer):
   """Layer constructor function for an embedding layer."""
 
-  def __init__(self, feature_depth, vocab_size, W_init=xavier_uniform()):
+  def __init__(self, feature_depth, vocab_size,
+               kernel_initializer=XavierUniformInitializer()):
     super(Embedding, self).__init__()
     self._feature_depth = feature_depth
     self._vocab_size = vocab_size
-    self._W_init = W_init
+    self._kernel_initializer = kernel_initializer
 
-  def call(self, params, inputs, **kwargs):
+  def call(self, x, params, **kwargs):
     del kwargs
-    return np.take(params, inputs, axis=0)
+    return np.take(params, x, axis=0)
 
   def output_shape(self, input_shape):
     return tuple(input_shape) + (self._feature_depth,)
 
   def new_parameters(self, input_shape, rng):
-    return self._W_init(rng, (self._vocab_size, self._feature_depth))
+    return self._kernel_initializer(
+        (self._vocab_size, self._feature_depth), rng)
 
 
 def padtype_to_pads(in_shape, window_shape, window_strides, padding):
@@ -176,34 +176,36 @@ def padtype_to_pads(in_shape, window_shape, window_strides, padding):
 class Conv(base.Layer):
   """Layer constructor function for a general convolution layer."""
 
-  def __init__(self, out_chan, filter_shape, strides=None, padding='VALID',
+  def __init__(self, filters, kernel_size, strides=None, padding='VALID',
                dimension_numbers=('NHWC', 'HWIO', 'NHWC'),
-               W_init=None, b_init=randn(1e-6)):
+               kernel_initializer=None,
+               bias_initializer=RandomNormalInitializer(1e-6)):
     super(Conv, self).__init__()
-    self._out_chan = out_chan
-    self._filter_shape = filter_shape
+    self._filters = filters
+    self._kernel_size = kernel_size
     self._padding = padding
     self._dimension_numbers = dimension_numbers
     self._lhs_spec, self._rhs_spec, self._out_spec = dimension_numbers
-    self._one = (1,) * len(filter_shape)
+    self._one = (1,) * len(kernel_size)
     self._strides = strides or self._one
-    self._b_init = b_init
+    self._bias_initializer = bias_initializer
     rhs_spec = self._rhs_spec
-    self._W_init = W_init or glorot(rhs_spec.index('O'), rhs_spec.index('I'))
+    self._kernel_initializer = kernel_initializer or GlorotNormalInitializer(
+        rhs_spec.index('O'), rhs_spec.index('I'))
 
-  def call(self, params, inputs, **kwargs):
+  def call(self, x, params=(), **kwargs):
     del kwargs
     w, b = params
     return lax.conv_general_dilated(
-        inputs, w, self._strides, self._padding, self._one, self._one,
+        x, w, self._strides, self._padding, self._one, self._one,
         self._dimension_numbers) + b
 
   def _kernel_shape(self, input_shape):
     """Helper to calculate the kernel shape."""
-    filter_shape_iter = iter(self._filter_shape)
-    return [self._out_chan if c == 'O' else
+    kernel_size_iter = iter(self._kernel_size)
+    return [self._filters if c == 'O' else
             input_shape[self._lhs_spec.index('C')] if c == 'I' else
-            next(filter_shape_iter) for c in self._rhs_spec]
+            next(kernel_size_iter) for c in self._rhs_spec]
 
   def _conv_shape_tuple(self, lhs_shape, rhs_shape, strides, pads):
     """Compute the shape of a conv given input shapes in canonical order."""
@@ -268,10 +270,10 @@ def output_shape(self, input_shape):
 
   def new_parameters(self, input_shape, rng):
     kernel_shape = self._kernel_shape(input_shape)
-    bias_shape = [self._out_chan if c == 'C' else 1 for c in self._out_spec]
+    bias_shape = [self._filters if c == 'C' else 1 for c in self._out_spec]
     bias_shape = tuple(itertools.dropwhile(lambda x: x == 1, bias_shape))
-    w = self._W_init(rng, kernel_shape)
-    b = self._b_init(rng, bias_shape)
+    w = self._kernel_initializer(kernel_shape, rng)
+    b = self._bias_initializer(bias_shape, rng)
     return (w, b)
 
 
@@ -287,9 +289,9 @@ def _flatten_output_shape(input_shape, num_axis_to_keep=1):
 
 
 @base.layer(output_shape=_flatten_output_shape)
-def Flatten(params, inputs, num_axis_to_keep=1, **kwargs):
+def Flatten(x, params, num_axis_to_keep=1, **kwargs):
   del params, kwargs
-  return np.reshape(inputs, (inputs.shape[:num_axis_to_keep] + (-1,)))
+  return np.reshape(x, (x.shape[:num_axis_to_keep] + (-1,)))
 
 
 # Batch normalization.
@@ -305,7 +307,7 @@ def _batch_norm_new_params(input_shape, rng, axis=(0, 1, 2),
 
 
 @base.layer(new_parameters=_batch_norm_new_params)
-def BatchNorm(params, x, axis=(0, 1, 2), epsilon=1e-5,
+def BatchNorm(x, params, axis=(0, 1, 2), epsilon=1e-5,
               center=True, scale=True, **unused_kwargs):
   """Layer construction function for a batch normalization layer."""
   mean = np.mean(x, axis, keepdims=True)
@@ -357,14 +359,14 @@ def _pooling_general(inputs, reducer, init_val, rescaler=None,
 
 
 @base.layer(output_shape=_pooling_output_shape)
-def MaxPool(params, x, pool_size=(2, 2), strides=None, padding='VALID', **kw):
+def MaxPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
   del params, kw
   return _pooling_general(x, lax.max, -np.inf, pool_size=pool_size,
                           strides=strides, padding=padding)
 
 
 @base.layer(output_shape=_pooling_output_shape)
-def SumPool(params, x, pool_size=(2, 2), strides=None, padding='VALID', **kw):
+def SumPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
   del params, kw
   return _pooling_general(x, lax.add, 0., pool_size=pool_size,
                           strides=strides, padding=padding)
@@ -380,14 +382,14 @@ def rescale(outputs, inputs):
 
 
 @base.layer(output_shape=_pooling_output_shape)
-def AvgPool(params, x, pool_size=(2, 2), strides=None, padding='VALID', **kw):
+def AvgPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
   del params, kw
   return _pooling_general(x, lax.add, 0., _normalize_by_window_size,
                           pool_size, strides=strides, padding=padding)
 
 
 @base.layer()
-def Dropout(params, x, rate=0.0, mode='train', rng=None, **kwargs):
+def Dropout(x, params, rate=0.0, mode='train', rng=None, **kwargs):
   """Layer construction function for a dropout layer with given rate."""
   del params, kwargs
   if rng is None:
@@ -405,16 +407,15 @@ def Dropout(params, x, rate=0.0, mode='train', rng=None, **kwargs):
 
 
 @base.layer()
-def Div(params, x, divisor=1.0, **kwargs):
+def Div(x, params, divisor=1.0, **kwargs):
   del params, kwargs
   return x / divisor
 
 
 @base.layer()
-def ShiftRight(params, inputs, **kwargs):
+def ShiftRight(x, **unused_kwargs):
   """Layer to shift the tensor to the right by padding on axis 1."""
-  del params, kwargs
   pad_widths = [(0, 0), (1, 0)]
-  pad_widths += [(0, 0) for _ in range(len(inputs.shape) - 2)]
-  padded = np.pad(inputs, pad_widths, mode='constant')
+  pad_widths += [(0, 0) for _ in range(len(x.shape) - 2)]
+  padded = np.pad(x, pad_widths, mode='constant')
   return padded[:, :-1, ...]
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 1545d1cf6..d15526e2c 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -79,7 +79,7 @@ def neg_log_perplexity(batch, model_predictions):
 def loss(params, batch, model_predict, rng):
   """Calculate loss."""
   inputs, targets = batch
-  preds = model_predict(params, inputs, rng=rng)
+  preds = model_predict(inputs, params, rng=rng)
   xent = np.sum(preds * stax.one_hot(targets, preds.shape[-1]), axis=-1)
   return - masked_mean(xent, targets)
 
@@ -246,21 +246,21 @@ def epochs(steps=None, epoch_steps=1):
 
 def _jit_predict_fun(model_predict, num_devices):
   """Use jit on model_predict if required."""
-  def predict(params, batch, rng=None):
+  def predict(x, params=(), rng=None):
     """Predict function jited and parallelized as requested."""
     # On one device, jit and run.
     if num_devices == 1:
-      return backend.jit(model_predict)(params, batch, rng=rng)
+      return backend.jit(model_predict)(x, params, rng=rng)
 
     # Multi-devices, pmap and run.
     @functools.partial(backend.pmap, axis_name="batch")
-    def mapped_predict(params, batch, rng):
-      return model_predict(params, batch, rng=rng)
+    def mapped_predict(x, params, rng):
+      return model_predict(x, params, rng=rng)
     pred = mapped_predict(
+        reshape_by_device(x, num_devices),
         params,
-        reshape_by_device(batch, num_devices),
         jax_random.split(rng, num_devices))
-    batch_size = batch.shape[0]
+    batch_size = x.shape[0]
     return np.reshape(pred, [batch_size] + list(pred.shape[2:]))
 
   return predict
@@ -373,16 +373,15 @@ def train(output_dir,
   history = state.history
   lr_fun = lr_schedule(history)
   opt_init, _ = optimizer(lr_fun)
-  model_init, model_predict_train = model(mode="train")
-  _, model_predict_eval = model(mode="eval")
+  model_train = model(mode="train")
+  model_predict_eval = model(mode="eval")
 
   # Setup state
   step = state.step or 0
-  rng, init_key = jax_random.split(rng)
+  rng, init_rng = jax_random.split(rng)
   rngs = jax_random.split(rng, num_devices)
-  params_initializer = \
-      lambda: model_init(init_key, [-1] + list(inputs.input_shape))[1]
-  params = state.params or params_initializer()
+  model_input_shape = tuple([-1] + list(inputs.input_shape))
+  params = state.params or model_train.initialize(model_input_shape, init_rng)
   opt_state = opt_init(params)
   if num_devices > 1:  # TODO(lukaszkaiser): use everywhere when pmap is stable.
     opt_state = jax.replicate(opt_state)
@@ -390,7 +389,7 @@ def train(output_dir,
   # jit model_predict and update so they're fast
   jit_model_predict_eval = _jit_predict_fun(model_predict_eval, num_devices)
   jit_update_fun = _jit_update_fun(
-      model_predict_train, loss_fun, optimizer, lr_fun, num_devices)
+      model_train, loss_fun, optimizer, lr_fun, num_devices)
 
   print()
   train_stream = inputs.train_stream()
@@ -403,7 +402,7 @@ def train(output_dir,
 
   # Non-compiled debug step helps find problems in models easier.
   if run_debug_step:
-    debug_loss = loss_fun(params, next(train_stream), model_predict_train, rng)
+    debug_loss = loss_fun(params, next(train_stream), model_train, rng)
     step_log(step, "Debug step loss %.8f" % debug_loss)
 
   for epoch, epoch_steps in epochs(train_steps, epoch_steps):
@@ -439,7 +438,7 @@ def train(output_dir,
     evaluate_train_and_eval(
         step=step,
         inputs=inputs,
-        predict_fun=functools.partial(jit_model_predict_eval, params),
+        predict_fun=functools.partial(jit_model_predict_eval, params=params),
         eval_steps=eval_steps,
         rng=rng,
         train_sw=train_sw,
@@ -459,7 +458,7 @@ def train(output_dir,
     lr_fun = lr_schedule(history)
     if lr_fun != old_lr_fun:  # For performance, only jit if there is a change.
       jit_update_fun = _jit_update_fun(
-          model_predict_train, loss_fun, optimizer, lr_fun, num_devices)
+          model_train, loss_fun, optimizer, lr_fun, num_devices)
 
     # Flush summary writers
     train_sw.flush()
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index 6f0032059..682635404 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -88,9 +88,8 @@ def test_train_eval_predict(self):
       self.assertEqual(2, len(eval_acc))
 
       # Predict with final params
-      _, predict_fun = model()
       inputs = inputs(1).train_stream()
-      predict_fun(state.params, next(inputs)[0])
+      model()(next(inputs)[0], state.params)
 
 
 if __name__ == "__main__":

From 169b3b470719895f60b53870d954ad8422ad8e4a Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 22 Apr 2019 19:27:28 -0700
Subject: [PATCH 1935/2720] Use LogSoftmax instead of Softmax in the policy
 network.

PiperOrigin-RevId: 244780716
---
 tensor2tensor/trax/rlax/ppo.py      | 37 +++++++++++++----------
 tensor2tensor/trax/rlax/ppo_test.py | 46 ++++++++++++++---------------
 2 files changed, 44 insertions(+), 39 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 5c724d4cf..28f7c949d 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -79,7 +79,8 @@ def policy_net(rng_key,
   # required layers on top of it.
   if bottom_layers is None:
     bottom_layers = []
-  bottom_layers.extend([stax.Dense(num_actions), stax.Softmax()])
+  # NOTE: The LogSoftmax instead of the Softmax.
+  bottom_layers.extend([stax.Dense(num_actions), stax.LogSoftmax()])
   net = stax.Serial(*bottom_layers)
   return net.initialize(batch_observations_shape, rng_key), net
 
@@ -113,8 +114,9 @@ def policy_and_value_net(rng_key,
 
   # Now, with the current logits, one head computes action probabilities and the
   # other computes the value function.
+  # NOTE: The LogSoftmax instead of the Softmax.
   layers.extend([stax.FanOut(), stax.Parallel(
-      stax.Serial(stax.Dense(num_actions), stax.Softmax()),
+      stax.Serial(stax.Dense(num_actions), stax.LogSoftmax()),
       stax.Dense(1)
   )])
 
@@ -197,6 +199,10 @@ def collect_trajectories(env,
           # Return the best action.
           action = np.argmax(predictions)
       elif policy == "categorical-sampling":
+        # NOTE: The predictions aren't probabilities but log-probabilities
+        # instead, since they were computed with LogSoftmax.
+        # So just np.exp them to make them probabilities.
+        predictions = np.exp(predictions)
         action = onp.argwhere(onp.random.multinomial(1, predictions) == 1)
       else:
         raise ValueError("Unknown policy: %s" % policy)
@@ -478,13 +484,13 @@ def chosen_probabs(probab_observations, actions):
 
   Args:
     probab_observations: ndarray of shape `[B, T, A]`, where
-      probab_observations[b, t, i] contains the probability of action = i at the
-      t^th time-step in the b^th trajectory.
+      probab_observations[b, t, i] contains the log-probability of action = i at
+      the t^th time-step in the b^th trajectory.
     actions: ndarray of shape `[B, T]`, with each entry in [0, A) denoting which
       action was chosen in the b^th trajectory's t^th time-step.
 
   Returns:
-    `[B, T]` ndarray with the probabilities of the chosen actions.
+    `[B, T]` ndarray with the log-probabilities of the chosen actions.
   """
   b, t = actions.shape
   return probab_observations[np.arange(b)[:, None], np.arange(t), actions]
@@ -494,7 +500,7 @@ def compute_probab_ratios(p_old, p_new, actions, reward_mask):
   """Computes the probability ratios for each time-step in a trajectory.
 
   Args:
-    p_old: ndarray of shape [B, T, A] of the probabilities that the policy
+    p_old: ndarray of shape [B, T, A] of the log-probabilities that the policy
       network assigns to all the actions at each time-step in each batch using
       the old parameters.
     p_new: ndarray of shape [B, T, A], same as above, but using new policy
@@ -506,13 +512,11 @@ def compute_probab_ratios(p_old, p_new, actions, reward_mask):
     probab_ratios: ndarray of shape [B, T], where
     probab_ratios_{b,t} = p_new_{b,t,action_{b,t}} / p_old_{b,t,action_{b,t}}
   """
-  bp_old = chosen_probabs(p_old, actions)
-  bp_new = chosen_probabs(p_new, actions)
+  logp_old = chosen_probabs(p_old, actions)
+  logp_new = chosen_probabs(p_new, actions)
 
-  # Add a small number to bp_old, where reward_mask is 0, this is just to help
-  # never to divide by 0.
-  bp_old = bp_old + (0.1 * np.abs(reward_mask - 1))
-  probab_ratios = (bp_new * reward_mask) / bp_old
+  # Since these are log-probabilities, we just subtract them.
+  probab_ratios = np.exp(logp_new - logp_old) * reward_mask
   return probab_ratios
 
 
@@ -555,11 +559,14 @@ def ppo_loss(policy_net_apply,
       td_deltas, reward_mask, lambda_=lambda_, gamma=gamma)
 
   # probab_actions_{old,new} are both (B, T, A)
-  probab_actions_old = policy_net_apply(padded_observations, old_policy_params)
-  probab_actions_new = policy_net_apply(padded_observations, new_policy_params)
+  log_probab_actions_old = policy_net_apply(padded_observations,
+                                            old_policy_params)
+  log_probab_actions_new = policy_net_apply(padded_observations,
+                                            new_policy_params)
 
   # (B, T)
-  ratios = compute_probab_ratios(probab_actions_old, probab_actions_new,
+  ratios = compute_probab_ratios(log_probab_actions_old,
+                                 log_probab_actions_new,
                                  padded_actions, reward_mask)
 
   # (B, T)
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 197297ef6..c025c5253 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -57,8 +57,8 @@ def test_policy_net(self):
     # Verify certain expectations on the output.
     self.assertEqual((batch, time_steps, num_actions), policy_output.shape)
 
-    # Also last axis normalizes to 1, since these are probabilities.
-    sum_actions = np.sum(policy_output, axis=-1)
+    # Also exp of last axis normalizes to 1, since these are log-probabilities.
+    sum_actions = np.sum(np.exp(policy_output), axis=-1)
     self.assertAllClose(np.ones_like(sum_actions), sum_actions)
 
   def test_value_net(self):
@@ -391,30 +391,28 @@ def test_chosen_probabs(self):
 
   def test_compute_probab_ratios(self):
     p_old = np.array([[
-        [0.1, 0.2, 0.6, 0.1],
-        [0.4, 0.1, 0.4, 0.1],
-        [0.3, 0.1, 0.5, 0.1],
-        [0.1, 0.2, 0.6, 0.1],
-    ],
-                      [
-                          [0.3, 0.1, 0.5, 0.1],
-                          [0.1, 0.1, 0.4, 0.4],
-                          [0.3, 0.1, 0.5, 0.1],
-                          [0.1, 0.2, 0.6, 0.1],
-                      ]])
+        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
+        [np.log(0.4), np.log(0.1), np.log(0.4), np.log(0.1)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
+    ], [
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [np.log(0.1), np.log(0.1), np.log(0.4), np.log(0.4)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
+    ]])
 
     p_new = np.array([[
-        [0.3, 0.1, 0.5, 0.1],
-        [0.4, 0.1, 0.1, 0.3],
-        [0.1, 0.2, 0.1, 0.6],
-        [0.3, 0.1, 0.5, 0.1],
-    ],
-                      [
-                          [0.1, 0.2, 0.1, 0.6],
-                          [0.1, 0.1, 0.2, 0.6],
-                          [0.3, 0.1, 0.3, 0.3],
-                          [0.1, 0.2, 0.1, 0.6],
-                      ]])
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [np.log(0.4), np.log(0.1), np.log(0.1), np.log(0.3)],
+        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+    ], [
+        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
+        [np.log(0.1), np.log(0.1), np.log(0.2), np.log(0.6)],
+        [np.log(0.3), np.log(0.1), np.log(0.3), np.log(0.3)],
+        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
+    ]])
 
     actions = np.array([[1, 2, 0, 1], [0, 3, 3, 0]])
 

From e6f32c7563b3c2a1442bf7bca49b86eb3e916a00 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 22 Apr 2019 22:51:58 -0700
Subject: [PATCH 1936/2720] Rename trax/stax --> trax/layers since trax layers
 are not compatible with stax any more.

PiperOrigin-RevId: 244798179
---
 tensor2tensor/trax/{stax => layers}/README.md |   0
 .../trax/{stax => layers}/__init__.py         |  10 +-
 .../trax/{stax => layers}/attention.py        |   6 +-
 tensor2tensor/trax/{stax => layers}/base.py   |   0
 .../trax/{stax => layers}/base_test.py        |  34 +--
 .../trax/{stax => layers}/combinators.py      |   2 +-
 tensor2tensor/trax/{stax => layers}/core.py   |   2 +-
 tensor2tensor/trax/models/mlp.py              |  12 +-
 tensor2tensor/trax/models/resnet.py           |  86 +++----
 tensor2tensor/trax/models/transformer.py      | 223 +++++++++---------
 tensor2tensor/trax/rlax/ppo.py                |  25 +-
 tensor2tensor/trax/rlax/ppo_main.py           |  15 +-
 tensor2tensor/trax/rlax/ppo_test.py           |  18 +-
 .../trax/rlax/ppo_training_loop_test.py       |   6 +-
 tensor2tensor/trax/trax.py                    |   6 +-
 15 files changed, 223 insertions(+), 222 deletions(-)
 rename tensor2tensor/trax/{stax => layers}/README.md (100%)
 rename tensor2tensor/trax/{stax => layers}/__init__.py (74%)
 rename tensor2tensor/trax/{stax => layers}/attention.py (98%)
 rename tensor2tensor/trax/{stax => layers}/base.py (100%)
 rename tensor2tensor/trax/{stax => layers}/base_test.py (75%)
 rename tensor2tensor/trax/{stax => layers}/combinators.py (99%)
 rename tensor2tensor/trax/{stax => layers}/core.py (99%)

diff --git a/tensor2tensor/trax/stax/README.md b/tensor2tensor/trax/layers/README.md
similarity index 100%
rename from tensor2tensor/trax/stax/README.md
rename to tensor2tensor/trax/layers/README.md
diff --git a/tensor2tensor/trax/stax/__init__.py b/tensor2tensor/trax/layers/__init__.py
similarity index 74%
rename from tensor2tensor/trax/stax/__init__.py
rename to tensor2tensor/trax/layers/__init__.py
index ac85bf1da..b359e2a1a 100644
--- a/tensor2tensor/trax/stax/__init__.py
+++ b/tensor2tensor/trax/layers/__init__.py
@@ -18,10 +18,10 @@
 from __future__ import division
 from __future__ import print_function
 
-# We create a flat stax.* namespace for uniform calling conventions as we
+# We create a flat layers.* namespace for uniform calling conventions as we
 # upstream changes.
 # pylint: disable=wildcard-import
-from tensor2tensor.trax.stax.attention import *
-from tensor2tensor.trax.stax.base import *
-from tensor2tensor.trax.stax.combinators import *
-from tensor2tensor.trax.stax.core import *
+from tensor2tensor.trax.layers.attention import *
+from tensor2tensor.trax.layers.base import *
+from tensor2tensor.trax.layers.combinators import *
+from tensor2tensor.trax.layers.core import *
diff --git a/tensor2tensor/trax/stax/attention.py b/tensor2tensor/trax/layers/attention.py
similarity index 98%
rename from tensor2tensor/trax/stax/attention.py
rename to tensor2tensor/trax/layers/attention.py
index ae3188238..b6d7d37db 100644
--- a/tensor2tensor/trax/stax/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -22,9 +22,9 @@
 
 from tensor2tensor.trax import backend
 from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.stax import base
-from tensor2tensor.trax.stax import combinators
-from tensor2tensor.trax.stax import core
+from tensor2tensor.trax.layers import base
+from tensor2tensor.trax.layers import combinators
+from tensor2tensor.trax.layers import core
 
 
 @base.layer(output_shape=lambda shape, axis=-1: (1, shape[axis], shape[axis]))
diff --git a/tensor2tensor/trax/stax/base.py b/tensor2tensor/trax/layers/base.py
similarity index 100%
rename from tensor2tensor/trax/stax/base.py
rename to tensor2tensor/trax/layers/base.py
diff --git a/tensor2tensor/trax/stax/base_test.py b/tensor2tensor/trax/layers/base_test.py
similarity index 75%
rename from tensor2tensor/trax/stax/base_test.py
rename to tensor2tensor/trax/layers/base_test.py
index 07e02e575..ed0baf65d 100644
--- a/tensor2tensor/trax/stax/base_test.py
+++ b/tensor2tensor/trax/layers/base_test.py
@@ -13,15 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for Stax base layer."""
+"""Tests for base layers."""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
 from absl.testing import absltest
 import numpy as onp
+from tensor2tensor.trax import layers
 from tensor2tensor.trax.backend import random
-import tensor2tensor.trax.stax as stax
 
 
 def random_inputs(rng, input_shape):
@@ -43,8 +43,8 @@ def check_shape_agreement(test_case, layer, input_shape):
   return result_shape
 
 
-def check_staxlayer(test_case, staxlayer, input_shape):
-  return check_shape_agreement(test_case, staxlayer, input_shape)
+def check_layer(test_case, layer, input_shape):
+  return check_shape_agreement(test_case, layer, input_shape)
 
 
 class SlaxTest(absltest.TestCase):
@@ -52,30 +52,30 @@ class SlaxTest(absltest.TestCase):
   def test_flatten_n(self):
     input_shape = (29, 87, 10, 20, 30)
 
-    actual_shape = check_staxlayer(self, stax.Flatten(), input_shape)
+    actual_shape = check_layer(self, layers.Flatten(), input_shape)
     self.assertEqual(actual_shape, (29, 87 * 10 * 20 * 30))
 
-    actual_shape = check_staxlayer(self, stax.Flatten(num_axis_to_keep=2),
-                                   input_shape)
+    actual_shape = check_layer(self, layers.Flatten(num_axis_to_keep=2),
+                               input_shape)
     self.assertEqual(actual_shape, (29, 87, 10 * 20 * 30))
 
-    actual_shape = check_staxlayer(self, stax.Flatten(num_axis_to_keep=3),
-                                   input_shape)
+    actual_shape = check_layer(self, layers.Flatten(num_axis_to_keep=3),
+                               input_shape)
     self.assertEqual(actual_shape, (29, 87, 10, 20 * 30))
 
-    actual_shape = check_staxlayer(self, stax.Flatten(num_axis_to_keep=4),
-                                   input_shape)
+    actual_shape = check_layer(self, layers.Flatten(num_axis_to_keep=4),
+                               input_shape)
     self.assertEqual(actual_shape, (29, 87, 10, 20, 30))
 
     # Not enough dimensions.
     with self.assertRaises(ValueError):
-      check_staxlayer(self, stax.Flatten(num_axis_to_keep=5), input_shape)
+      check_layer(self, layers.Flatten(num_axis_to_keep=5), input_shape)
 
     with self.assertRaises(ValueError):
-      check_staxlayer(self, stax.Flatten(num_axis_to_keep=6), input_shape)
+      check_layer(self, layers.Flatten(num_axis_to_keep=6), input_shape)
 
   def test_div(self):
-    layer = stax.Div(divisor=2.0)
+    layer = layers.Div(divisor=2.0)
     input_np = onp.array([[1, 2, 3], [4, 5, 6]], dtype=onp.float32)
     output_np = layer(input_np)
     # absltest doesn't have ndarray equalities.
@@ -86,9 +86,9 @@ def test_div(self):
         delta=1e-6)
 
   def test_dense_param_sharing(self):
-    model1 = stax.Serial(stax.Dense(32), stax.Dense(32))
-    layer = stax.Dense(32)
-    model2 = stax.Serial(layer, layer)
+    model1 = layers.Serial(layers.Dense(32), layers.Dense(32))
+    layer = layers.Dense(32)
+    model2 = layers.Serial(layer, layer)
     rng = random.get_prng(0)
     params1 = model1.initialize((-1, 32), rng)
     params2 = model2.initialize((-1, 32), rng)
diff --git a/tensor2tensor/trax/stax/combinators.py b/tensor2tensor/trax/layers/combinators.py
similarity index 99%
rename from tensor2tensor/trax/stax/combinators.py
rename to tensor2tensor/trax/layers/combinators.py
index 47885ca3c..a4f58e793 100644
--- a/tensor2tensor/trax/stax/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -20,7 +20,7 @@
 from __future__ import print_function
 
 from tensor2tensor.trax import backend
-from tensor2tensor.trax.stax import base
+from tensor2tensor.trax.layers import base
 
 
 class Serial(base.Layer):
diff --git a/tensor2tensor/trax/stax/core.py b/tensor2tensor/trax/layers/core.py
similarity index 99%
rename from tensor2tensor/trax/stax/core.py
rename to tensor2tensor/trax/layers/core.py
index e3924fa63..57e3c9b1c 100644
--- a/tensor2tensor/trax/stax/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -28,7 +28,7 @@
 from six.moves import reduce
 from tensor2tensor.trax import backend
 from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.stax import base
+from tensor2tensor.trax.layers import base
 
 # Following the convention used in Keras and tf.layers, we use CamelCase for the
 # names of layer constructors, like Conv and Relu, while using snake_case for
diff --git a/tensor2tensor/trax/models/mlp.py b/tensor2tensor/trax/models/mlp.py
index ab39d6a11..f93536821 100644
--- a/tensor2tensor/trax/models/mlp.py
+++ b/tensor2tensor/trax/models/mlp.py
@@ -19,18 +19,18 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensor2tensor.trax.stax as stax
+from tensor2tensor.trax import layers
 
 
 def MLP(num_hidden_layers=2,
         hidden_size=512,
-        activation_fn=stax.Relu,
+        activation_fn=layers.Relu,
         num_output_classes=10,
         mode="train"):
   """Multi-layer feed-forward neural network with non-linear activations."""
   del mode
-  layers = [stax.Flatten()]
+  cur_layers = [layers.Flatten()]
   for _ in range(num_hidden_layers):
-    layers += [stax.Dense(hidden_size), activation_fn()]
-  layers += [stax.Dense(num_output_classes), stax.LogSoftmax()]
-  return stax.Serial(*layers)
+    cur_layers += [layers.Dense(hidden_size), activation_fn()]
+  cur_layers += [layers.Dense(num_output_classes), layers.LogSoftmax()]
+  return layers.Serial(*cur_layers)
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index 1c88108d3..aa3caae5c 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -19,42 +19,42 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensor2tensor.trax.stax as stax
+from tensor2tensor.trax import layers
 
 
 def ConvBlock(kernel_size, filters, strides):
   """ResNet convolutional striding block."""
   ks = kernel_size
   filters1, filters2, filters3 = filters
-  main = stax.Serial(
-      stax.Conv(filters1, (1, 1), strides),
-      stax.BatchNorm(), stax.Relu(),
-      stax.Conv(filters2, (ks, ks), padding='SAME'),
-      stax.BatchNorm(), stax.Relu(),
-      stax.Conv(filters3, (1, 1)), stax.BatchNorm())
-  shortcut = stax.Serial(
-      stax.Conv(filters3, (1, 1), strides),
-      stax.BatchNorm())
-  return stax.Serial(
-      stax.FanOut(),
-      stax.Parallel(main, shortcut),
-      stax.FanInSum(), stax.Relu())
+  main = layers.Serial(
+      layers.Conv(filters1, (1, 1), strides),
+      layers.BatchNorm(), layers.Relu(),
+      layers.Conv(filters2, (ks, ks), padding='SAME'),
+      layers.BatchNorm(), layers.Relu(),
+      layers.Conv(filters3, (1, 1)), layers.BatchNorm())
+  shortcut = layers.Serial(
+      layers.Conv(filters3, (1, 1), strides),
+      layers.BatchNorm())
+  return layers.Serial(
+      layers.FanOut(),
+      layers.Parallel(main, shortcut),
+      layers.FanInSum(), layers.Relu())
 
 
 def IdentityBlock(kernel_size, filters):
   """ResNet identical size block."""
   ks = kernel_size
   filters1, filters2, filters3 = filters
-  main = stax.Serial(
-      stax.Conv(filters1, (1, 1)),
-      stax.BatchNorm(), stax.Relu(),
-      stax.Conv(filters2, (ks, ks), padding='SAME'),
-      stax.BatchNorm(), stax.Relu(),
-      stax.Conv(filters3, (1, 1)), stax.BatchNorm())
-  return stax.Serial(
-      stax.FanOut(),
-      stax.Parallel(main, stax.Identity()),
-      stax.FanInSum(), stax.Relu())
+  main = layers.Serial(
+      layers.Conv(filters1, (1, 1)),
+      layers.BatchNorm(), layers.Relu(),
+      layers.Conv(filters2, (ks, ks), padding='SAME'),
+      layers.BatchNorm(), layers.Relu(),
+      layers.Conv(filters3, (1, 1)), layers.BatchNorm())
+  return layers.Serial(
+      layers.FanOut(),
+      layers.Parallel(main, layers.Identity()),
+      layers.FanInSum(), layers.Relu())
 
 
 def Resnet50(hidden_size=64, num_output_classes=1001, mode='train'):
@@ -69,10 +69,10 @@ def Resnet50(hidden_size=64, num_output_classes=1001, mode='train'):
     The ResNet model with the given layer and output sizes.
   """
   del mode
-  return stax.Serial(
-      stax.Conv(hidden_size, (7, 7), (2, 2), 'SAME'),
-      stax.BatchNorm(), stax.Relu(),
-      stax.MaxPool(pool_size=(3, 3), strides=(2, 2)),
+  return layers.Serial(
+      layers.Conv(hidden_size, (7, 7), (2, 2), 'SAME'),
+      layers.BatchNorm(), layers.Relu(),
+      layers.MaxPool(pool_size=(3, 3), strides=(2, 2)),
       ConvBlock(3, [hidden_size, hidden_size, 4 * hidden_size], (1, 1)),
       IdentityBlock(3, [hidden_size, hidden_size, 4 * hidden_size]),
       IdentityBlock(3, [hidden_size, hidden_size, 4 * hidden_size]),
@@ -89,20 +89,20 @@ def Resnet50(hidden_size=64, num_output_classes=1001, mode='train'):
       ConvBlock(3, [8 * hidden_size, 8 * hidden_size, 32*hidden_size], (2, 2)),
       IdentityBlock(3, [8 * hidden_size, 8 * hidden_size, 32 * hidden_size]),
       IdentityBlock(3, [8 * hidden_size, 8 * hidden_size, 32 * hidden_size]),
-      stax.AvgPool(pool_size=(7, 7)), stax.Flatten(),
-      stax.Dense(num_output_classes), stax.LogSoftmax())
+      layers.AvgPool(pool_size=(7, 7)), layers.Flatten(),
+      layers.Dense(num_output_classes), layers.LogSoftmax())
 
 
 def WideResnetBlock(channels, strides=(1, 1), channel_mismatch=False):
   """WideResnet convolutational block."""
-  main = stax.Serial(stax.BatchNorm(), stax.Relu(),
-                     stax.Conv(channels, (3, 3), strides, padding='SAME'),
-                     stax.BatchNorm(), stax.Relu(),
-                     stax.Conv(channels, (3, 3), padding='SAME'))
-  shortcut = stax.Identity() if not channel_mismatch else stax.Conv(
+  main = layers.Serial(layers.BatchNorm(), layers.Relu(),
+                       layers.Conv(channels, (3, 3), strides, padding='SAME'),
+                       layers.BatchNorm(), layers.Relu(),
+                       layers.Conv(channels, (3, 3), padding='SAME'))
+  shortcut = layers.Identity() if not channel_mismatch else layers.Conv(
       channels, (3, 3), strides, padding='SAME')
-  return stax.Serial(
-      stax.FanOut(), stax.Parallel(main, shortcut), stax.FanInSum())
+  return layers.Serial(
+      layers.FanOut(), layers.Parallel(main, shortcut), layers.FanInSum())
 
 
 def WideResnetGroup(n, channels, strides=(1, 1)):
@@ -110,7 +110,7 @@ def WideResnetGroup(n, channels, strides=(1, 1)):
   blocks += [WideResnetBlock(channels, strides, channel_mismatch=True)]
   for _ in range(n - 1):
     blocks += [WideResnetBlock(channels, (1, 1))]
-  return stax.Serial(*blocks)
+  return layers.Serial(*blocks)
 
 
 def WideResnet(num_blocks=3, hidden_size=64, num_output_classes=10,
@@ -127,10 +127,10 @@ def WideResnet(num_blocks=3, hidden_size=64, num_output_classes=10,
     The WideResnet model with given layer and output sizes.
   """
   del mode
-  return stax.Serial(
-      stax.Conv(hidden_size, (3, 3), padding='SAME'),
+  return layers.Serial(
+      layers.Conv(hidden_size, (3, 3), padding='SAME'),
       WideResnetGroup(num_blocks, hidden_size),
       WideResnetGroup(num_blocks, hidden_size * 2, (2, 2)),
-      WideResnetGroup(num_blocks, hidden_size * 4, (2, 2)), stax.BatchNorm(),
-      stax.Relu(), stax.AvgPool(pool_size=(8, 8)), stax.Flatten(),
-      stax.Dense(num_output_classes), stax.LogSoftmax())
+      WideResnetGroup(num_blocks, hidden_size * 4, (2, 2)), layers.BatchNorm(),
+      layers.Relu(), layers.AvgPool(pool_size=(8, 8)), layers.Flatten(),
+      layers.Dense(num_output_classes), layers.LogSoftmax())
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 026d54a4e..e2026a4ed 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -18,7 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensor2tensor.trax.stax as stax
+from tensor2tensor.trax import layers
 
 
 def ResidualFeedForward(feature_depth,
@@ -26,15 +26,15 @@ def ResidualFeedForward(feature_depth,
                         dropout,
                         mode):
   """Residual feed-forward layer with normalization at start."""
-  return stax.Residual(
-      stax.LayerNorm(),
-      stax.Dense(feedforward_depth,
-                 kernel_initializer=stax.XavierUniformInitializer()),
-      stax.Relu(),
-      stax.Dropout(rate=dropout, mode=mode),
-      stax.Dense(feature_depth,
-                 kernel_initializer=stax.XavierUniformInitializer()),
-      stax.Dropout(rate=dropout, mode=mode)
+  return layers.Residual(
+      layers.LayerNorm(),
+      layers.Dense(feedforward_depth,
+                   kernel_initializer=layers.XavierUniformInitializer()),
+      layers.Relu(),
+      layers.Dropout(rate=dropout, mode=mode),
+      layers.Dense(feature_depth,
+                   kernel_initializer=layers.XavierUniformInitializer()),
+      layers.Dropout(rate=dropout, mode=mode)
   )
 
 
@@ -52,46 +52,45 @@ def TransformerEncoder(mode='train',
     feature_depth: int:  depth of embedding
     feedforward_depth: int: depth of feed-forward layer
     num_heads: int: number of attention heads
-    dropout: float: dropout rate (how much to drop out; note that stax follows
-      Tensorflow's keep_rate convention, so we use 1 - dropout in calls below)
+    dropout: float: dropout rate
 
   Returns:
-    A staxlayer for implementing a raw Transformer encoder stack.  No embedding
+    A layer for implementing a raw Transformer encoder stack.  No embedding
     or positional signals are added by this layer.
   """
   # Multi-headed Attention and Feed-forward layers
-  multi_attention = stax.MultiHeadedAttention(
+  multi_attention = layers.MultiHeadedAttention(
       feature_depth, num_heads=num_heads, dropout=dropout, mode=mode)
 
-  @stax.Lambda
+  @layers.Lambda
   def Encoder(embedded_source, source_mask):
     """Transformer encoder stack.
 
     Args:
-      embedded_source: staxlayer variable: embedded source sequences
-      source_mask: staxlayer variable: self-attention mask
+      embedded_source: layer variable: embedded source sequences
+      source_mask: layer variable: self-attention mask
 
     Returns:
-      Staxlayer variable that outputs encoded source.
+      Layer variable that outputs encoded source.
     """
-    encoder_layer = stax.Serial(
+    encoder_layer = layers.Serial(
         # input attends to self
-        stax.Residual(stax.LayerNorm(),
-                      stax.FanOut(size=4),
-                      stax.Parallel(stax.Identity(),  # query
-                                    stax.Identity(),  # key
-                                    stax.Identity(),  # value
-                                    source_mask),  # attention mask
-                      multi_attention,
-                      stax.Dropout(rate=dropout, mode=mode)),
+        layers.Residual(layers.LayerNorm(),
+                        layers.FanOut(size=4),
+                        layers.Parallel(layers.Identity(),  # query
+                                        layers.Identity(),  # key
+                                        layers.Identity(),  # value
+                                        source_mask),  # attention mask
+                        multi_attention,
+                        layers.Dropout(rate=dropout, mode=mode)),
         # feed-forward
         ResidualFeedForward(
             feature_depth, feedforward_depth, dropout, mode=mode)
     )
-    return stax.Serial(
+    return layers.Serial(
         embedded_source,
-        stax.repeat(encoder_layer, num_layers),
-        stax.LayerNorm(),
+        layers.repeat(encoder_layer, num_layers),
+        layers.LayerNorm(),
     )
 
   return Encoder
@@ -114,17 +113,17 @@ def DecoderLayer(feature_depth,
   Returns:
     init and apply.
   """
-  return stax.Serial(
-      stax.Residual(  # Self-attention block.
-          stax.LayerNorm(),
-          stax.FanOut(size=4),
-          stax.Parallel(stax.Identity(),  # query
-                        stax.Identity(),  # key
-                        stax.Identity(),  # value
-                        stax.CausalMask(axis=-2)),  # attention mask
-          stax.MultiHeadedAttention(feature_depth, num_heads=num_heads,
-                                    dropout=dropout, mode=mode),
-          stax.Dropout(rate=dropout, mode=mode)
+  return layers.Serial(
+      layers.Residual(  # Self-attention block.
+          layers.LayerNorm(),
+          layers.FanOut(size=4),
+          layers.Parallel(layers.Identity(),  # query
+                          layers.Identity(),  # key
+                          layers.Identity(),  # value
+                          layers.CausalMask(axis=-2)),  # attention mask
+          layers.MultiHeadedAttention(feature_depth, num_heads=num_heads,
+                                      dropout=dropout, mode=mode),
+          layers.Dropout(rate=dropout, mode=mode)
       ),
       ResidualFeedForward(feature_depth, feedforward_depth, dropout, mode=mode)
   )
@@ -153,18 +152,18 @@ def TransformerLM(vocab_size,
   Returns:
     init and apply.
   """
-  return stax.Serial(
-      stax.ShiftRight(),
-      stax.Embedding(feature_depth, vocab_size),
-      stax.Dropout(rate=dropout, mode=mode),
-      stax.PositionalEncoding(max_len=max_len),
-      stax.Serial(*[DecoderLayer(feature_depth, feedforward_depth, num_heads,
-                                 dropout, mode)
-                    for _ in range(num_layers)]),
-      stax.LayerNorm(),
-      stax.Dense(vocab_size,
-                 kernel_initializer=stax.XavierUniformInitializer()),
-      stax.LogSoftmax()
+  return layers.Serial(
+      layers.ShiftRight(),
+      layers.Embedding(feature_depth, vocab_size),
+      layers.Dropout(rate=dropout, mode=mode),
+      layers.PositionalEncoding(max_len=max_len),
+      layers.Serial(*[DecoderLayer(feature_depth, feedforward_depth, num_heads,
+                                   dropout, mode)
+                      for _ in range(num_layers)]),
+      layers.LayerNorm(),
+      layers.Dense(vocab_size,
+                   kernel_initializer=layers.XavierUniformInitializer()),
+      layers.LogSoftmax()
   )
 
 
@@ -203,124 +202,124 @@ def Transformer(source_vocab_size,
   functions for the trained encoder, decoder, and generator substax.
   """
   # Input embedding and positional encoding
-  inject_position = stax.Serial(
-      stax.Dropout(dropout, mode=mode),
-      stax.PositionalEncoding(feature_depth, max_len=max_len)
+  inject_position = layers.Serial(
+      layers.Dropout(dropout, mode=mode),
+      layers.PositionalEncoding(feature_depth, max_len=max_len)
   )
   if shared_embedding:
     assert source_vocab_size == target_vocab_size
     # Weight-shared Embedding
-    embedding = stax.Share(stax.Embedding(feature_depth, source_vocab_size))
-    source_embedding_layer = stax.Serial(embedding, inject_position)
+    embedding = layers.Share(layers.Embedding(feature_depth, source_vocab_size))
+    source_embedding_layer = layers.Serial(embedding, inject_position)
     target_embedding_layer = source_embedding_layer
   else:
-    source_embedding = stax.Embedding(feature_depth, source_vocab_size)
-    target_embedding = stax.Embedding(feature_depth, target_vocab_size)
-    source_embedding_layer = stax.Serial(source_embedding, inject_position)
-    target_embedding_layer = stax.Serial(target_embedding, inject_position)
+    source_embedding = layers.Embedding(feature_depth, source_vocab_size)
+    target_embedding = layers.Embedding(feature_depth, target_vocab_size)
+    source_embedding_layer = layers.Serial(source_embedding, inject_position)
+    target_embedding_layer = layers.Serial(target_embedding, inject_position)
 
   # Multi-headed Attention and Feed-forward layers
-  multi_attention = stax.MultiHeadedAttention(
+  multi_attention = layers.MultiHeadedAttention(
       feature_depth, num_heads=num_heads, dropout=dropout, mode=mode)
 
   # Encoder
-  @stax.Lambda
+  @layers.Lambda
   def Encoder(source, source_mask):
     """Transformer encoder stack.
 
     Args:
-      source: staxlayer variable: raw source sequences
-      source_mask: staxlayer variable: self-attention mask
+      source: layer variable: raw source sequences
+      source_mask: layer variable: self-attention mask
 
     Returns:
-      Staxlayer variable that outputs encoded source.
+      Layer variable that outputs encoded source.
     """
-    encoder_layer = stax.Serial(
+    encoder_layer = layers.Serial(
         # input attends to self
-        stax.Residual(stax.LayerNorm(),
-                      stax.FanOut(size=4),
-                      stax.Parallel(stax.Identity(),  # query
-                                    stax.Identity(),  # key
-                                    stax.Identity(),  # value
-                                    source_mask),  # attention mask
-                      multi_attention,
-                      stax.Dropout(dropout, mode=mode)),
+        layers.Residual(layers.LayerNorm(),
+                        layers.FanOut(size=4),
+                        layers.Parallel(layers.Identity(),  # query
+                                        layers.Identity(),  # key
+                                        layers.Identity(),  # value
+                                        source_mask),  # attention mask
+                        multi_attention,
+                        layers.Dropout(dropout, mode=mode)),
         # feed-forward
         ResidualFeedForward(
             feature_depth, feedforward_depth, dropout, mode=mode),
     )
-    return stax.Serial(
+    return layers.Serial(
         source,
         source_embedding_layer,
-        stax.repeat(encoder_layer, num_layers),
-        stax.LayerNorm(),
+        layers.repeat(encoder_layer, num_layers),
+        layers.LayerNorm(),
     )
 
   # Decoder
-  @stax.Lambda
+  @layers.Lambda
   def Decoder(memory, target, target_mask, memory_mask):
     """Transformer decoder stack.
 
     Args:
-      memory: staxlayer variable: encoded source sequences
-      target: staxlayer variable: raw target sequences
-      target_mask: staxlayer variable: self-attention mask
-      memory_mask: staxlayer variable: memory attention mask
+      memory: layer variable: encoded source sequences
+      target: layer variable: raw target sequences
+      target_mask: layer variable: self-attention mask
+      memory_mask: layer variable: memory attention mask
 
     Returns:
-      Staxlayer variable that outputs encoded source.
+      Layer variable that outputs encoded source.
     """
-    decoder_layer = stax.Serial(
+    decoder_layer = layers.Serial(
         # target attends to self
-        stax.Residual(stax.LayerNorm(),
-                      stax.FanOut(size=4),
-                      stax.Parallel(stax.Identity(),  # query
-                                    stax.Identity(),  # key
-                                    stax.Identity(),  # value
-                                    target_mask),  # attention mask
-                      multi_attention,
-                      stax.Dropout(dropout, mode=mode)),
+        layers.Residual(layers.LayerNorm(),
+                        layers.FanOut(size=4),
+                        layers.Parallel(layers.Identity(),  # query
+                                        layers.Identity(),  # key
+                                        layers.Identity(),  # value
+                                        target_mask),  # attention mask
+                        multi_attention,
+                        layers.Dropout(dropout, mode=mode)),
         # target attends to encoded source
-        stax.Residual(stax.LayerNorm(),
-                      stax.FanOut(size=4),
-                      stax.Parallel(stax.Identity(),  # query
-                                    memory,  # key
-                                    memory,  # value
-                                    memory_mask),  # attention mask
-                      multi_attention,
-                      stax.Dropout(dropout, mode=mode)),
+        layers.Residual(layers.LayerNorm(),
+                        layers.FanOut(size=4),
+                        layers.Parallel(layers.Identity(),  # query
+                                        memory,  # key
+                                        memory,  # value
+                                        memory_mask),  # attention mask
+                        multi_attention,
+                        layers.Dropout(dropout, mode=mode)),
         # feed-forward
         ResidualFeedForward(
             feature_depth, feedforward_depth, dropout, mode=mode)
     )
-    return stax.Serial(
+    return layers.Serial(
         target,
         target_embedding_layer,
-        stax.repeat(decoder_layer, num_layers),
-        stax.LayerNorm(),
+        layers.repeat(decoder_layer, num_layers),
+        layers.LayerNorm(),
     )
 
   # The Transformer
-  @stax.Lambda
+  @layers.Lambda
   def transformer(source, target, source_mask, target_mask, memory_mask):  # pylint: disable=invalid-name
     encoded_source = Encoder(source, source_mask)
     return Decoder(encoded_source, target, target_mask, memory_mask)
 
   # Finally, bind the generator transform to use later for inference.
-  @stax.Lambda
+  @layers.Lambda
   def Generator(encoded_target):
-    return stax.Serial(
+    return layers.Serial(
         encoded_target,
-        stax.Dense(target_vocab_size,
-                   kernel_initializer=stax.XavierUniformInitializer()),
-        stax.LogSoftmax
+        layers.Dense(target_vocab_size,
+                     kernel_initializer=layers.XavierUniformInitializer()),
+        layers.LogSoftmax
     )
 
   # Model-Building and Evaluation Functions
   # Get entire model's init and apply pair
   top_init, top_apply = Generator(transformer)
 
-  # By default act as a normal Stax constructor and emit an (init, apply) pair.
+  # By default act as a normal constructor and emit an (init, apply) pair.
   if not return_evals:
     return (top_init, top_apply)
   else:
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 28f7c949d..65e2c4947 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -56,8 +56,8 @@
 from jax import numpy as np
 from jax import random as jax_random
 import numpy as onp
+from tensor2tensor.trax import layers
 from tensor2tensor.trax import optimizers as trax_opt
-from tensor2tensor.trax import stax
 from tensor2tensor.trax import trax
 
 DEBUG_LOGGING = False
@@ -79,9 +79,11 @@ def policy_net(rng_key,
   # required layers on top of it.
   if bottom_layers is None:
     bottom_layers = []
+
   # NOTE: The LogSoftmax instead of the Softmax.
-  bottom_layers.extend([stax.Dense(num_actions), stax.LogSoftmax()])
-  net = stax.Serial(*bottom_layers)
+  bottom_layers.extend([layers.Dense(num_actions), layers.LogSoftmax()])
+  net = layers.Serial(*bottom_layers)
+
   return net.initialize(batch_observations_shape, rng_key), net
 
 
@@ -95,9 +97,9 @@ def value_net(rng_key,
   if bottom_layers is None:
     bottom_layers = []
   bottom_layers.extend([
-      stax.Dense(1),
+      layers.Dense(1),
   ])
-  net = stax.Serial(*bottom_layers)
+  net = layers.Serial(*bottom_layers)
   return net.initialize(batch_observations_shape, rng_key), net
 
 
@@ -108,19 +110,18 @@ def policy_and_value_net(rng_key,
   """A policy and value net function."""
 
   # Layers.
-  layers = []
+  cur_layers = []
   if bottom_layers is not None:
-    layers.extend(bottom_layers)
+    cur_layers.extend(bottom_layers)
 
   # Now, with the current logits, one head computes action probabilities and the
   # other computes the value function.
   # NOTE: The LogSoftmax instead of the Softmax.
-  layers.extend([stax.FanOut(), stax.Parallel(
-      stax.Serial(stax.Dense(num_actions), stax.LogSoftmax()),
-      stax.Dense(1)
+  cur_layers.extend([layers.FanOut(), layers.Parallel(
+      layers.Serial(layers.Dense(num_actions), layers.LogSoftmax()),
+      layers.Dense(1)
   )])
-
-  net = stax.Serial(*layers)
+  net = layers.Serial(*cur_layers)
   return net.initialize(batch_observations_shape, rng_key), net
 
 
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index fa0028e7d..1e8968c39 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -25,7 +25,7 @@
 from absl import flags
 import jax
 from jax.config import config
-from tensor2tensor.trax import stax
+from tensor2tensor.trax import layers
 from tensor2tensor.trax.rlax import ppo
 
 FLAGS = flags.FLAGS
@@ -46,11 +46,12 @@
                      "Setting to true will help to debug nans.")
 
 
-def common_stax_layers():
-  layers = []
+def common_layers():
+  cur_layers = []
   if FLAGS.env_name == "Pong-v0":
-    layers = [stax.Div(divisor=255.0), stax.Flatten(num_axis_to_keep=2)]
-  return layers + [stax.Dense(16), stax.Relu(), stax.Dense(4), stax.Relu()]
+    cur_layers = [layers.Div(divisor=255.0), layers.Flatten(num_axis_to_keep=2)]
+  return cur_layers + [layers.Dense(16), layers.Relu(),
+                       layers.Dense(4), layers.Relu()]
 
 
 def main(argv):
@@ -67,9 +68,9 @@ def run_training_loop():
         env_name=FLAGS.env_name,
         epochs=FLAGS.epochs,
         policy_net_fun=functools.partial(
-            ppo.policy_net, bottom_layers=common_stax_layers()),
+            ppo.policy_net, bottom_layers=common_layers()),
         value_net_fun=functools.partial(
-            ppo.value_net, bottom_layers=common_stax_layers()),
+            ppo.value_net, bottom_layers=common_layers()),
         policy_optimizer_fun=optimizer_fun,
         value_optimizer_fun=optimizer_fun,
         batch_size=FLAGS.batch_size,
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index c025c5253..c388ee736 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -22,7 +22,7 @@
 import jax
 from jax import random as jax_random
 import numpy as np
-from tensor2tensor.trax import stax
+from tensor2tensor.trax import layers
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.rlax import fake_env
 from tensor2tensor.trax.rlax import ppo
@@ -43,7 +43,7 @@ def test_policy_net(self):
         num_actions,
         # flatten except batch and time
         # step dimensions.
-        [stax.Flatten(num_axis_to_keep=2)])
+        [layers.Flatten(num_axis_to_keep=2)])
 
     # Generate a batch of observations.
     batch = 2
@@ -67,7 +67,7 @@ def test_value_net(self):
     value_params, value_apply = ppo.value_net(
         self.rng_key,
         (-1, -1) + observation_shape,
-        num_actions, [stax.Flatten(num_axis_to_keep=2)])
+        num_actions, [layers.Flatten(num_axis_to_keep=2)])
     batch = 2
     time_steps = 10
     batch_of_observations = np.random.uniform(
@@ -83,7 +83,7 @@ def test_policy_and_value_net(self):
     num_actions = 2
     pnv_params, pnv_apply = ppo.policy_and_value_net(
         self.rng_key, batch_observation_shape, num_actions,
-        [stax.Flatten(num_axis_to_keep=2)])
+        [layers.Flatten(num_axis_to_keep=2)])
     batch = 2
     time_steps = 10
     batch_of_observations = np.random.uniform(
@@ -104,7 +104,7 @@ def test_collect_trajectories(self):
         num_actions,
         # flatten except batch and time
         # step dimensions.
-        [stax.Flatten(num_axis_to_keep=2)])
+        [layers.Flatten(num_axis_to_keep=2)])
 
     # We'll get done at time-step #5, starting from 0, therefore in 6 steps.
     done_time_step = 5
@@ -139,7 +139,7 @@ def test_collect_trajectories_max_timestep(self):
         num_actions,
         # flatten except batch and time
         # step dimensions.
-        [stax.Flatten(num_axis_to_keep=2)])
+        [layers.Flatten(num_axis_to_keep=2)])
 
     # We'll get done at time-step #5, starting from 0, therefore in 6 steps.
     done_time_step = 5
@@ -487,16 +487,16 @@ def test_ppo_loss(self):
     batch_observation_shape = (-1, -1) + OBS
 
     old_policy_params, _ = ppo.policy_net(key1, batch_observation_shape, A,
-                                          [stax.Flatten(num_axis_to_keep=2)])
+                                          [layers.Flatten(num_axis_to_keep=2)])
 
     new_policy_params, policy_apply = ppo.policy_net(
         key2,
         batch_observation_shape, A,
-        [stax.Flatten(num_axis_to_keep=2)])
+        [layers.Flatten(num_axis_to_keep=2)])
 
     value_params, value_apply = ppo.value_net(
         key3, batch_observation_shape, A,
-        [stax.Flatten(num_axis_to_keep=2)])
+        [layers.Flatten(num_axis_to_keep=2)])
 
     # Generate a batch of observations.
 
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index 2ac2b5912..a0aabc0da 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -22,7 +22,7 @@
 import functools
 import gym
 from tensor2tensor.rl import gym_utils
-from tensor2tensor.trax import stax
+from tensor2tensor.trax import layers
 from tensor2tensor.trax.rlax import ppo
 from tensorflow import test
 
@@ -42,9 +42,9 @@ def test_training_loop(self):
         env=env,
         epochs=num_epochs,
         policy_net_fun=functools.partial(
-            ppo.policy_net, bottom_layers=[stax.Dense(1)]),
+            ppo.policy_net, bottom_layers=[layers.Dense(1)]),
         value_net_fun=functools.partial(
-            ppo.value_net, bottom_layers=[stax.Dense(1)]),
+            ppo.value_net, bottom_layers=[layers.Dense(1)]),
         batch_size=batch_size,
         num_optimizer_steps=1,
         random_seed=0)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index d15526e2c..3b83c7fee 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -40,11 +40,11 @@
 from tensor2tensor.trax import history as trax_history
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import jaxboard
+from tensor2tensor.trax import layers
 from tensor2tensor.trax import learning_rate as lr
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.backend import random as jax_random
-import tensor2tensor.trax.stax as stax
 
 import tensorflow as tf
 from tensorflow.io import gfile
@@ -71,7 +71,7 @@ def accuracy(batch, model_predictions):
 def neg_log_perplexity(batch, model_predictions):
   """Calculate negative log perplexity."""
   _, targets = batch
-  hot_targets = stax.one_hot(targets, model_predictions.shape[-1])
+  hot_targets = layers.one_hot(targets, model_predictions.shape[-1])
   xent = np.sum(model_predictions * hot_targets, axis=-1)
   return masked_mean(xent, targets)
 
@@ -80,7 +80,7 @@ def loss(params, batch, model_predict, rng):
   """Calculate loss."""
   inputs, targets = batch
   preds = model_predict(inputs, params, rng=rng)
-  xent = np.sum(preds * stax.one_hot(targets, preds.shape[-1]), axis=-1)
+  xent = np.sum(preds * layers.one_hot(targets, preds.shape[-1]), axis=-1)
   return - masked_mean(xent, targets)
 
 

From 9abac7c61501b11215e31f006920bb5e78710185 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 23 Apr 2019 11:06:29 -0700
Subject: [PATCH 1937/2720] t2t_decoding with option score_file now takes
 FLAGS.checkpoint_path into account.

PiperOrigin-RevId: 244887513
---
 tensor2tensor/bin/t2t_decoder.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 768fdac7b..72bda7b2b 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -129,8 +129,11 @@ def score_file(filename):
 
   with tf.Session() as sess:
     # Load weights from checkpoint.
-    ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
-    ckpt = ckpts.model_checkpoint_path
+    if FLAGS.checkpoint_path is None:
+      ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
+      ckpt = ckpts.model_checkpoint_path
+    else:
+      ckpt = FLAGS.checkpoint_path
     saver.restore(sess, ckpt)
     # Run on each line.
     with tf.gfile.Open(filename) as f:

From 1fb0dd3e008c4a2a8dc159a8153bcb2d0b86a1d8 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 23 Apr 2019 13:38:45 -0700
Subject: [PATCH 1938/2720] Adding summaries for min and max xent Rename
 hparam_sets for mixture transformer Bugfix: Fixing tf.squeeze call to handle
 case where only 1 element is present in batch Change implementation to add
 mixture embeddings to original vocab embedding matrix, and use bottom method
 to retrieve them Change implementation to add mixture embeddings directly to
 decoder_input Added new problem for multi-spelling dataset

PiperOrigin-RevId: 244915958
---
 .../data_generators/translate_enro.py         |  2 +-
 tensor2tensor/layers/common_layers.py         | 23 +++++++++++--------
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/data_generators/translate_enro.py b/tensor2tensor/data_generators/translate_enro.py
index f022f638d..35115570f 100644
--- a/tensor2tensor/data_generators/translate_enro.py
+++ b/tensor2tensor/data_generators/translate_enro.py
@@ -32,7 +32,7 @@
     [
         "http://www.statmt.org/europarl/v7/ro-en.tgz",
         ("europarl-v7.ro-en.en", "europarl-v7.ro-en.ro")
-    ],
+    ]
 ]
 _ENRO_TEST_DATASETS = [
     [
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 68f2b534d..4a6b16865 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1846,8 +1846,8 @@ def padded_cross_entropy_mixture(logits,
   new_shape_for_xent = [num_mixtures] + shape_list(labels)
   labels = tf.tile(labels, [num_mixtures, 1, 1, 1])
 
-  xent, weights = padded_cross_entropy(
-      logits, labels, label_smoothing, weights_fn, reduce_sum, cutoff, gaussian)
+  xent, weights = padded_cross_entropy(logits, labels, label_smoothing,
+                                       weights_fn, reduce_sum, cutoff, gaussian)
 
   # reshape xent and weights to have the num_mixtures as first dimension
   xent = tf.reshape(xent, new_shape_for_xent)
@@ -1860,8 +1860,8 @@ def padded_cross_entropy_mixture(logits,
   if return_best_logits:
     best_mixture_indices = tf.cast(tf.argmin(xent, 0), dtype=tf.int32)
     individual_element_indices = tf.range(batch_size)
-    stacked_mixture_element_indices = tf.stack(
-        (tf.squeeze(best_mixture_indices), individual_element_indices), -1)
+    stacked_mixture_element_indices = tf.stack((tf.squeeze(
+        best_mixture_indices, axis=[1, 2]), individual_element_indices), -1)
     best_logits = tf.reshape(logits,
                              [num_mixtures, -1, timesteps, 1, 1, vocab_size])
     best_logits = tf.gather_nd(best_logits, stacked_mixture_element_indices)
@@ -1874,22 +1874,27 @@ def padded_cross_entropy_mixture(logits,
           message="Each batch element should have a probability value for each mixture element"
       )
   ]):
-    xent = tf.reduce_min(xent, axis=0)
+    xent_min = tf.reduce_min(xent, axis=0)
+    xent_max = tf.reduce_max(xent, axis=0)
     weights = tf.reduce_mean(weights, axis=0)
 
   with tf.control_dependencies([
       tf.assert_equal(
-          tf.shape(xent)[0], [batch_size],
+          tf.shape(xent_min)[0], [batch_size],
           message="There should be batch_size elements after selecting best mixture probabilities"
       )
   ]):
-    summed_xent = tf.reduce_sum(xent)
+    summed_xent_min = tf.reduce_sum(xent_min)
+    summed_xent_max = tf.reduce_sum(xent_max)
     summed_weights = tf.reduce_sum(weights)
 
+    tf.summary.scalar("mixture_xents_min", summed_xent_min / summed_weights)
+    tf.summary.scalar("mixture_xents_max", summed_xent_max / summed_weights)
+
   if return_best_logits:
-    return summed_xent, summed_weights, best_logits
+    return summed_xent_min, summed_weights, best_logits
   else:
-    return summed_xent, summed_weights
+    return summed_xent_min, summed_weights
 
 
 def _weights_one_third(labels):

From a68e020ae9cf419fc42fc6e0f4161cfbc80c2c55 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 23 Apr 2019 14:31:26 -0700
Subject: [PATCH 1939/2720] Fixed typo referencing tf.Dataset

PiperOrigin-RevId: 244926626
---
 tensor2tensor/trax/inputs.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 0ddcba04e..7e9b88177 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -163,8 +163,8 @@ def train_and_eval_dataset(dataset_name, data_dir, train_shuffle_files=True,
 
   Returns:
     a 4-tuple consisting of:
-     * the train tf.Daataset
-     * the eval tf.Daataset
+     * the train tf.Dataset
+     * the eval tf.Dataset
      * information about features: a python dictionary with feature names
          as keys and an object as value that provides .shape and .num_classes.
      * supervised_keys: information what's the input and what's the target,

From e9bf9a592bca5446fc996d057a7d20416e2b9967 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 23 Apr 2019 16:24:50 -0700
Subject: [PATCH 1940/2720]  - Correct some shapes in the test.     - Previous
 tests pass observations as (B, T) + OBS.     - But actually should be (B,
 T+1) + OBS.  - More shape checking in the code.     - We now assert on those
 shapes in the code.  - More logging in collect.     - Mainly for timing, this
 is on vlog 2, rest of the logging is vlog 1.

PiperOrigin-RevId: 244947311
---
 tensor2tensor/trax/rlax/ppo.py      | 62 +++++++++++++++++++++++------
 tensor2tensor/trax/rlax/ppo_test.py | 12 ++++--
 2 files changed, 59 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 65e2c4947..f2e26441c 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -26,19 +26,23 @@
 
 Policy and Value function signatures:
 
-Policy Function :: [B, T] + OBS -> [B, T, A]
-Value  Function :: [B, T] + OBS -> [B, T, 1]
+Policy            Function :: [B, T] + OBS ->  [B, T, A]
+Value             Function :: [B, T] + OBS ->  [B, T, 1]
+Policy and Value  Function :: [B, T] + OBS -> ([B, T, A], [B, T, 1])
 
 i.e. the policy net should take a batch of *trajectories* and at each time-step
 in each batch deliver a probability distribution over actions.
 
-NOTE: It doesn't return logits, rather the expectation is that it return a
-normalized distribution instead.
+NOTE: It doesn't return logits, rather the expectation is that it returns
+log-probabilities instead.
 
 NOTE: The policy and value functions need to take care to not take into account
 future time-steps while deciding the actions (or value) for the current
 time-step.
 
+Policy and Value Function produces a tuple of the expected output of a policy
+function and a value function.
+
 """
 
 from __future__ import absolute_import
@@ -161,7 +165,8 @@ def collect_trajectories(env,
   """Collect trajectories with the given policy net and behaviour."""
   trajectories = []
 
-  for _ in range(num_trajectories):
+  for t in range(num_trajectories):
+    t_start = time.time()
     rewards = []
     actions = []
     done = False
@@ -174,8 +179,10 @@ def collect_trajectories(env,
 
     # Run either till we're done OR if max_timestep is defined only till that
     # timestep.
+    ts = 0
     while ((not done) and
            (not max_timestep or observation_history.shape[1] < max_timestep)):
+      ts_start = time.time()
       # Run the policy, to pick an action, shape is (1, t, A) because
       # observation_history is shaped (1, t) + OBS
       predictions = policy_net_apply(observation_history, policy_net_params)
@@ -233,6 +240,13 @@ def collect_trajectories(env,
       rewards.append(reward)
       actions.append(action)
 
+      ts += 1
+      logging.vlog(
+          2, "  Collected time-step[ %5d] of trajectory[ %5d] in [%0.2f] msec.",
+          ts, t, get_time(ts_start))
+    logging.vlog(
+        2, " Collected trajectory[ %5d] in [%0.2f] msec.", t, get_time(t_start))
+
     # This means we are done we're been terminated early.
     assert done or (
         max_timestep and max_timestep >= observation_history.shape[1])
@@ -484,7 +498,7 @@ def chosen_probabs(probab_observations, actions):
   """Picks out the probabilities of the actions along batch and time-steps.
 
   Args:
-    probab_observations: ndarray of shape `[B, T, A]`, where
+    probab_observations: ndarray of shape `[B, T+1, A]`, where
       probab_observations[b, t, i] contains the log-probability of action = i at
       the t^th time-step in the b^th trajectory.
     actions: ndarray of shape `[B, T]`, with each entry in [0, A) denoting which
@@ -493,18 +507,19 @@ def chosen_probabs(probab_observations, actions):
   Returns:
     `[B, T]` ndarray with the log-probabilities of the chosen actions.
   """
-  b, t = actions.shape
-  return probab_observations[np.arange(b)[:, None], np.arange(t), actions]
+  B, T = actions.shape  # pylint: disable=invalid-name
+  assert (B, T+1) == probab_observations.shape[:2]
+  return probab_observations[np.arange(B)[:, None], np.arange(T), actions]
 
 
 def compute_probab_ratios(p_old, p_new, actions, reward_mask):
   """Computes the probability ratios for each time-step in a trajectory.
 
   Args:
-    p_old: ndarray of shape [B, T, A] of the log-probabilities that the policy
+    p_old: ndarray of shape [B, T+1, A] of the log-probabilities that the policy
       network assigns to all the actions at each time-step in each batch using
       the old parameters.
-    p_new: ndarray of shape [B, T, A], same as above, but using new policy
+    p_new: ndarray of shape [B, T+1, A], same as above, but using new policy
       network parameters.
     actions: ndarray of shape [B, T] where each element is from [0, A).
     reward_mask: ndarray of shape [B, T] masking over probabilities.
@@ -513,11 +528,20 @@ def compute_probab_ratios(p_old, p_new, actions, reward_mask):
     probab_ratios: ndarray of shape [B, T], where
     probab_ratios_{b,t} = p_new_{b,t,action_{b,t}} / p_old_{b,t,action_{b,t}}
   """
+
+  B, T = actions.shape  # pylint: disable=invalid-name
+  assert (B, T+1) == p_old.shape[:2]
+  assert (B, T+1) == p_new.shape[:2]
+
   logp_old = chosen_probabs(p_old, actions)
   logp_new = chosen_probabs(p_new, actions)
 
+  assert (B, T) == logp_old.shape
+  assert (B, T) == logp_new.shape
+
   # Since these are log-probabilities, we just subtract them.
   probab_ratios = np.exp(logp_new - logp_old) * reward_mask
+  assert (B, T) == probab_ratios.shape
   return probab_ratios
 
 
@@ -545,34 +569,48 @@ def ppo_loss(policy_net_apply,
              lambda_=0.95,
              epsilon=0.2):
   """PPO objective, with an eventual minus sign."""
+  B, T = padded_rewards.shape  # pylint: disable=invalid-name
+  assert (B, T+1) == padded_observations.shape[:2]
+  assert (B, T) == padded_actions.shape
+  assert (B, T) == padded_rewards.shape
+  assert (B, T) == reward_mask.shape
+
   # (B, T+1, 1)
   predicted_values = value_net_apply(padded_observations, value_net_params)
+  assert (B, T+1, 1) == predicted_values.shape
 
   # (B, T)
   td_deltas = deltas(
-      np.squeeze(predicted_values, axis=2),  # (B, T)
+      np.squeeze(predicted_values, axis=2),  # (B, T+1)
       padded_rewards,
       reward_mask,
       gamma=gamma)
+  assert (B, T) == td_deltas.shape
 
   # (B, T)
   advantages = gae_advantages(
       td_deltas, reward_mask, lambda_=lambda_, gamma=gamma)
+  assert (B, T) == advantages.shape
 
-  # probab_actions_{old,new} are both (B, T, A)
+  # probab_actions_{old,new} are both (B, T+1, A)
   log_probab_actions_old = policy_net_apply(padded_observations,
                                             old_policy_params)
   log_probab_actions_new = policy_net_apply(padded_observations,
                                             new_policy_params)
+  assert (B, T+1) == log_probab_actions_old.shape[:2]
+  assert (B, T+1) == log_probab_actions_new.shape[:2]
+  assert log_probab_actions_old.shape[-1] == log_probab_actions_new.shape[-1]
 
   # (B, T)
   ratios = compute_probab_ratios(log_probab_actions_old,
                                  log_probab_actions_new,
                                  padded_actions, reward_mask)
+  assert (B, T) == ratios.shape
 
   # (B, T)
   objective = clipped_objective(
       ratios, advantages, reward_mask, epsilon=epsilon)
+  assert (B, T) == objective.shape
 
   # ()
   average_objective = np.sum(objective) / np.sum(reward_mask)
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index c388ee736..cbd9b1d6d 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -378,9 +378,11 @@ def test_gae_advantages(self):
     self.assertAllEqual(expected_gae_advantages, gae_advantages)
 
   def test_chosen_probabs(self):
-    # Shape (2, 2, 3)
-    probab_observations = np.array([[[0.1, 0.2, 0.7], [0.4, 0.1, 0.5]],
-                                    [[0.3, 0.1, 0.6], [0.1, 0.1, 0.8]]])
+    # Shape (2, 2+1, 3)
+    probab_observations = np.array(
+        [[[0.1, 0.2, 0.7], [0.4, 0.1, 0.5], [0.2, 0.4, 0.4]],
+         [[0.3, 0.1, 0.6], [0.1, 0.1, 0.8], [0.2, 0.4, 0.4]]]
+    )
 
     # Shape (2, 2)
     actions = np.array([[1, 2], [0, 1]])
@@ -395,11 +397,13 @@ def test_compute_probab_ratios(self):
         [np.log(0.4), np.log(0.1), np.log(0.4), np.log(0.1)],
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
         [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
     ], [
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
         [np.log(0.1), np.log(0.1), np.log(0.4), np.log(0.4)],
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
         [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
     ]])
 
     p_new = np.array([[
@@ -407,11 +411,13 @@ def test_compute_probab_ratios(self):
         [np.log(0.4), np.log(0.1), np.log(0.1), np.log(0.3)],
         [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
     ], [
         [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
         [np.log(0.1), np.log(0.1), np.log(0.2), np.log(0.6)],
         [np.log(0.3), np.log(0.1), np.log(0.3), np.log(0.3)],
         [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
+        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
     ]])
 
     actions = np.array([[1, 2, 0, 1], [0, 3, 3, 0]])

From e8c8192dcd985a412a9a33dc6aac06c5e80c4d02 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 23 Apr 2019 16:59:27 -0700
Subject: [PATCH 1941/2720] Add support for multiple passes of scheduled
 sampling.

PiperOrigin-RevId: 244952793
---
 tensor2tensor/layers/common_hparams.py |  8 +++++---
 tensor2tensor/utils/t2t_model.py       | 14 ++++++++++++--
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 9988283ff..502bbbc61 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -242,12 +242,14 @@ def basic_params1():
       # will such additional step be run. It's turned off (0.0) by default.
       # This probability will exponentially warm up for the number of
       # steps determined by scheduled_sampling_warmup_steps.
-      # The tensor used for the second step will consist of outputs from
-      # the first step mixed with gold truth, with the proportion of gold
-      # determined by scheduled_sampling_gold_mixin_prob.
+      # The tensor used for the n-th pass will consist of outputs from
+      # the (n-1)-th pass mixed with gold truth, with the proportion of gold
+      # determined by scheduled_sampling_gold_mixin_prob. Control the number
+      # of passes with scheduled_sampling_num_passes.
       scheduled_sampling_prob=0.0,
       scheduled_sampling_warmup_steps=50000,
       scheduled_sampling_gold_mixin_prob=0.5,
+      scheduled_sampling_num_passes=1,
       # This setting controls whether to copy variables around in a daisy chain
       # (if true) or leave their placement to TensorFlow. It only affects multi
       # device training and mostly should be turned on for performance. One
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index de79f2b75..88988574b 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1842,7 +1842,7 @@ def mix_gold_sampled(gold_targets, sampled_targets, mixin_prob):
           sampled_targets,
           gold_targets)
 
-    def sampled_results(mixin_prob):
+    def sampled_results(features, logits, mixin_prob):
       """Generate scheduled sampling results."""
       sampled_targets = sample(logits)
       new_targets = mix_gold_sampled(features["targets"],
@@ -1886,7 +1886,17 @@ def sampled_results(mixin_prob):
             hparams.scheduled_sampling_warmup_steps,
             min_value=0.001)
     )
-    return sampled_results(mixin_prob)
+
+    # Apply scheduled sampling over N passes. The logits from the (n-1)-th pass
+    # will be mixed with gold tokens for conditioning in the n-th pass.
+    assert hparams.scheduled_sampling_num_passes > 0, (
+        "hparams.scheduled_sampling_num_passes must be > 0 if "
+        "hparams.scheduled_sampling_prob > 0.0")
+    new_logits = logits
+    new_losses = losses
+    for _ in range(hparams.scheduled_sampling_num_passes):
+      new_logits, new_losses = sampled_results(features, new_logits, mixin_prob)
+    return new_logits, new_losses
 
 
 def _with_timing(fn, msg, silent=False):

From 95aa6092cbc569508e5f51e45c2e29e04cbc1fab Mon Sep 17 00:00:00 2001
From: Harini Kannan <hkannan@google.com>
Date: Tue, 23 Apr 2019 17:06:31 -0700
Subject: [PATCH 1942/2720] Fix bug causing t2t_datagen to crash when
 generating resized images for more than one episode.

PiperOrigin-RevId: 244954085
---
 tensor2tensor/rl/gym_utils.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index a7c339b12..c4cd1da1d 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -117,7 +117,13 @@ def step(self, action):
 
   def reset(self, **kwargs):
     self.env.reset(**kwargs)
-    return self.env.render(mode=self.mode)
+    obs = self.env.render(mode=self.mode)
+    if self.should_resize:
+      img = Image.fromarray(obs)
+      img = img.resize(self.observation_space.shape[:-1],
+                       resample=Image.ANTIALIAS)
+      obs = np.asarray(img)
+    return obs
 
 
 def remove_time_limit_wrapper(env):

From b93fc036fdbacfddcadad8fb781f5b670533384e Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 23 Apr 2019 22:51:42 -0700
Subject: [PATCH 1943/2720] Rename FanIn --> Branch and make TransformerEncoder
 work.

PiperOrigin-RevId: 244986961
---
 .../trax/configs/transformer_imdb_8gb.gin     |  50 +++++++
 tensor2tensor/trax/layers/attention.py        |  51 +++++--
 tensor2tensor/trax/layers/combinators.py      |  93 +++++++++---
 tensor2tensor/trax/layers/core.py             |  14 ++
 tensor2tensor/trax/models/__init__.py         |   1 +
 tensor2tensor/trax/models/resnet.py           |  37 +++--
 tensor2tensor/trax/models/transformer.py      | 132 +++++++++++-------
 tensor2tensor/trax/rlax/ppo.py                |   2 +-
 8 files changed, 284 insertions(+), 96 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/transformer_imdb_8gb.gin

diff --git a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
new file mode 100644
index 000000000..3fdac10aa
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
@@ -0,0 +1,50 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 128
+batch_fun.eval_batch_size = 128
+batch_fun.max_eval_length = 2048
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_sentiment_imdb'
+inputs.input_name = 'targets'
+
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 0.1
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
+lm1b_preprocess.max_target_length = 512
+lm1b_preprocess.max_eval_target_length = 2048
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 100
+train.eval_steps = 10
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerEncoder
+train.run_debug_step = False
+train.train_steps = 1000
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerEncoder.dropout = 0.1
+TransformerEncoder.feature_depth = 512
+TransformerEncoder.feedforward_depth = 2048
+TransformerEncoder.max_len = 2048
+TransformerEncoder.mode = 'train'
+TransformerEncoder.num_classes = 10
+TransformerEncoder.num_heads = 8
+TransformerEncoder.num_layers = 6
+TransformerEncoder.vocab_size = 32000
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index b6d7d37db..23f487b6a 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -34,6 +34,12 @@ def CausalMask(x, params, axis=-1, **kwargs):
   return onp.tril(onp.ones((1, size, size), dtype=x.dtype), k=0)
 
 
+@base.layer(output_shape=lambda shape, pad=0: (shape[0], 1, 1, shape[-1]))
+def PaddingMask(x, params, pad=0, **kwargs):
+  del params, kwargs
+  return np.reshape(x != pad, (x.shape[0], 1, 1, x.shape[-1]))
+
+
 def MakeTargetMask(target, pad=0):
   """Create an attention mask to hide padding and future words."""
   target_mask = (target != pad)[ :, np.newaxis, :]
@@ -188,7 +194,7 @@ def PureMultiHeadedAttention(x, params, feature_depth=None,
   """
   del params
   rng = kwargs.get('rng', None)
-  q, k, v, mask = x
+  (q, k, v), mask = x
   assert feature_depth % num_heads == 0
   head_depth = feature_depth // num_heads
   nbatch = np.shape(q)[0]
@@ -207,10 +213,12 @@ def JoinHeads(x):  # pylint: disable=invalid-name
           dropout=dropout, mode=mode, rng=rng))
 
 
-def MultiHeadedAttention(
+def MultiHeadedAttentionQKV(
     feature_depth, num_heads=8, dropout=0.0, mode='train'):
   """Transformer-style multi-headed attention.
 
+  Accepts inputs of the form (q, k, v), mask.
+
   Args:
     feature_depth: int:  depth of embedding
     num_heads: int: number of attention heads
@@ -222,12 +230,14 @@ def MultiHeadedAttention(
   """
   return combinators.Serial(
       combinators.Parallel(
-          core.Dense(feature_depth,
-                     kernel_initializer=core.XavierUniformInitializer()),
-          core.Dense(feature_depth,
-                     kernel_initializer=core.XavierUniformInitializer()),
-          core.Dense(feature_depth,
-                     kernel_initializer=core.XavierUniformInitializer()),
+          combinators.Parallel(
+              core.Dense(feature_depth,
+                         kernel_initializer=core.XavierUniformInitializer()),
+              core.Dense(feature_depth,
+                         kernel_initializer=core.XavierUniformInitializer()),
+              core.Dense(feature_depth,
+                         kernel_initializer=core.XavierUniformInitializer()),
+          ),
           combinators.Identity()
       ),
       PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
@@ -236,3 +246,28 @@ def MultiHeadedAttention(
       core.Dense(feature_depth,
                  kernel_initializer=core.XavierUniformInitializer()),
   )
+
+
+def MultiHeadedAttention(
+    feature_depth, num_heads=8, dropout=0.0, mode='train'):
+  """Transformer-style multi-headed attention.
+
+  Accepts inputs of the form (x, mask) and constructs (q, k, v) from x.
+
+  Args:
+    feature_depth: int:  depth of embedding
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    Multi-headed self-attention layer.
+  """
+  return combinators.Serial(
+      combinators.Parallel(
+          combinators.Branch(num_branches=3),  # q = k = v = first input
+          combinators.Identity()  # pass the mask
+      ),
+      MultiHeadedAttentionQKV(  # pylint: disable=no-value-for-parameter
+          feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),
+  )
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index a4f58e793..e12f7ccc4 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -59,30 +59,83 @@ def new_parameters(self, input_shape, rng):
 
 @base.layer()
 def Identity(x, **unused_kwargs):
+  """Identity layer, return the inputs."""
   return x
 
 
-@base.layer(output_shape=lambda input_shape, size=2: [input_shape] * size)
-def FanOut(x, params, size=2, **kwargs):
+# Re-ordering layer.
+def _reorder_shape(input_shape, output=None):  # pylint: disable=invalid-name
+  """Helper to determine the shape of reorder output."""
+  if output is None:
+    return input_shape
+  return base.nested_map(output, lambda i: input_shape[i])
+
+
+@base.layer(output_shape=_reorder_shape)
+def Reorder(x, params, output=None, **kwargs):
+  """Reorder a tuple into another tuple.
+
+  For example, we can re-order (x, y) into (y, x) or even (y, (x, y), y).
+  The output argument specifies how to re-order, using integers that refer
+  to indices in the input tuple. For example, if
+
+    input = (x, y, z)
+
+  then
+
+    Reorder(input, output=(1, 0, 2))   = (y, x, z)
+    Reorder(input, output=(0, 0))      = (x, x)
+    Reorder(input, output=(0, (1, 1))) = (x, (y, y))
+    Reorder(input, output=((2, 0), (1, 1))) = ((z, x), (y, y))
+
+  By default (if no output is given) Reorder does nothing (Identity).
+
+  Args:
+    x: the input tuple to re-order.
+    params: layer parameters (unused).
+    output: the specification of the output tuple: a nested tuple of ints.
+    **kwargs: other arguments (unused).
+
+  Returns:
+    The re-ordered tuple with the same shape as output.
+  """
+  del params, kwargs
+  if output is None:
+    return x
+  return base.nested_map(output, lambda i: x[i])
+
+
+@base.layer(output_shape=lambda shape, num_branches=2: [shape] * num_branches)
+def Branch(x, params, num_branches=2, **kwargs):
   del params, kwargs
-  return [x] * size
+  return [x] * num_branches
+
+
+@base.layer(output_shape=lambda input_shape_list: input_shape_list[0])
+def FirstBranch(x, **unused_kwargs):
+  return x[0]  # Here x is a list of tensors, we select the first.
+
+
+@base.layer(output_shape=lambda input_shape_list: input_shape_list[1])
+def SecondBranch(x, **unused_kwargs):
+  return x[1]  # Here x is a list of tensors, we select the second.
 
 
 @base.layer(output_shape=lambda input_shape_list: input_shape_list[0])
-def FanInSum(x, **unused_kwargs):
+def SumBranches(x, **unused_kwargs):
   return sum(x)  # Here x is a list of tensors of the same shape, we add them.
 
 
-def _fan_in_concat_shape(input_shape, axis=-1):  # pylint: disable=invalid-name
-  """Helper to determine the shape of FanInConcat output."""
+def _concatenate_shape(input_shape, axis=-1):  # pylint: disable=invalid-name
+  """Helper to determine the shape of Concatenate output."""
   ax = axis % len(input_shape[0])
   concat_size = sum(shape[ax] for shape in input_shape)
   out_shape = input_shape[0][:ax] + (concat_size,) + input_shape[0][ax+1:]
   return out_shape
 
 
-@base.layer(output_shape=_fan_in_concat_shape)
-def FanInConcat(x, params, axis=-1, **kwargs):
+@base.layer(output_shape=_concatenate_shape)
+def Concatenate(x, params, axis=-1, **kwargs):
   del params, kwargs
   return backend.numpy.concatenate(x, axis)
 
@@ -90,16 +143,14 @@ def FanInConcat(x, params, axis=-1, **kwargs):
 class Parallel(base.Layer):
   """Combinator for composing layers in parallel.
 
-  The layer resulting from this combinator is often used with the FanOut and
-  FanInSum layers.
+  This layer is often used with the Branch and SumBranches layers.
 
   Args:
-    *layers: a sequence of layers, each an (init_fun, apply_fun) pair.
+    *layers: a sequence of layers.
 
   Returns:
-    A new layer, meaning an (init_fun, apply_fun) pair, representing the
-    parallel composition of the given sequence of layers. In particular, the
-    returned layer takes a sequence of inputs and returns a sequence of outputs
+    A new layer representing parallel composition of the given layers.
+    The new layer takes a sequence of inputs and returns a sequence of outputs
     with the same length as the argument `layers`.
   """
 
@@ -128,18 +179,18 @@ def new_parameters(self, input_shape, rng):
 
 def Residual(*layers, **kwargs):
   """Constructs a residual version of layers, summing input to layers output."""
-  res = kwargs.get('res', Identity())  # pylint: disable=no-value-for-parameter
+  shortcut = kwargs.get('shortcut', Identity())  # pylint: disable=no-value-for-parameter
   if len(layers) > 1:
     return Serial(
-        FanOut(),  # pylint: disable=no-value-for-parameter
-        Parallel(Serial(*layers), res),
-        FanInSum()  # pylint: disable=no-value-for-parameter
+        Branch(),  # pylint: disable=no-value-for-parameter
+        Parallel(Serial(*layers), shortcut),
+        SumBranches()  # pylint: disable=no-value-for-parameter
     )
   elif len(layers) == 1:
     return Serial(
-        FanOut(),  # pylint: disable=no-value-for-parameter
-        Parallel(layers[0], res),
-        FanInSum()  # pylint: disable=no-value-for-parameter
+        Branch(),  # pylint: disable=no-value-for-parameter
+        Parallel(layers[0], shortcut),
+        SumBranches()  # pylint: disable=no-value-for-parameter
     )
   else:
     raise ValueError('Empty residual combinator.')
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 57e3c9b1c..f75c127c9 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -412,6 +412,20 @@ def Div(x, params, divisor=1.0, **kwargs):
   return x / divisor
 
 
+# Mean.
+def _mean_output_shape(input_shape, axis=-1, keepdims=False):
+  shape1 = list(input_shape)[:axis]  # Shape before axis.
+  shape2 = list(input_shape)[axis:][1:]  # Shape after axis.
+  mid_shape = [1] if keepdims else []
+  return tuple(shape1 + mid_shape + shape2)
+
+
+@base.layer(output_shape=_mean_output_shape)
+def Mean(x, params, axis=-1, keepdims=False, **kwargs):
+  del params, kwargs
+  return np.mean(x, axis=axis, keepdims=keepdims)
+
+
 @base.layer()
 def ShiftRight(x, **unused_kwargs):
   """Layer to shift the tensor to the right by padding on axis 1."""
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index bde69f33d..ff07113fc 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -35,4 +35,5 @@ def model_configure(*args, **kwargs):
 MLP = model_configure(mlp.MLP)
 Resnet50 = model_configure(resnet.Resnet50)
 WideResnet = model_configure(resnet.WideResnet)
+TransformerEncoder = model_configure(transformer.TransformerEncoder)
 TransformerLM = model_configure(transformer.TransformerLM)
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index aa3caae5c..aa6316792 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -28,17 +28,24 @@ def ConvBlock(kernel_size, filters, strides):
   filters1, filters2, filters3 = filters
   main = layers.Serial(
       layers.Conv(filters1, (1, 1), strides),
-      layers.BatchNorm(), layers.Relu(),
+      layers.BatchNorm(),
+      layers.Relu(),
       layers.Conv(filters2, (ks, ks), padding='SAME'),
-      layers.BatchNorm(), layers.Relu(),
-      layers.Conv(filters3, (1, 1)), layers.BatchNorm())
+      layers.BatchNorm(),
+      layers.Relu(),
+      layers.Conv(filters3, (1, 1)),
+      layers.BatchNorm()
+  )
   shortcut = layers.Serial(
       layers.Conv(filters3, (1, 1), strides),
-      layers.BatchNorm())
+      layers.BatchNorm()
+  )
   return layers.Serial(
-      layers.FanOut(),
+      layers.Branch(),
       layers.Parallel(main, shortcut),
-      layers.FanInSum(), layers.Relu())
+      layers.SumBranches(),
+      layers.Relu()
+  )
 
 
 def IdentityBlock(kernel_size, filters):
@@ -47,14 +54,20 @@ def IdentityBlock(kernel_size, filters):
   filters1, filters2, filters3 = filters
   main = layers.Serial(
       layers.Conv(filters1, (1, 1)),
-      layers.BatchNorm(), layers.Relu(),
+      layers.BatchNorm(),
+      layers.Relu(),
       layers.Conv(filters2, (ks, ks), padding='SAME'),
-      layers.BatchNorm(), layers.Relu(),
-      layers.Conv(filters3, (1, 1)), layers.BatchNorm())
+      layers.BatchNorm(),
+      layers.Relu(),
+      layers.Conv(filters3, (1, 1)),
+      layers.BatchNorm()
+  )
   return layers.Serial(
-      layers.FanOut(),
+      layers.Branch(),
       layers.Parallel(main, layers.Identity()),
-      layers.FanInSum(), layers.Relu())
+      layers.SumBranches(),
+      layers.Relu()
+  )
 
 
 def Resnet50(hidden_size=64, num_output_classes=1001, mode='train'):
@@ -102,7 +115,7 @@ def WideResnetBlock(channels, strides=(1, 1), channel_mismatch=False):
   shortcut = layers.Identity() if not channel_mismatch else layers.Conv(
       channels, (3, 3), strides, padding='SAME')
   return layers.Serial(
-      layers.FanOut(), layers.Parallel(main, shortcut), layers.FanInSum())
+      layers.Branch(), layers.Parallel(main, shortcut), layers.SumBranches())
 
 
 def WideResnetGroup(n, channels, strides=(1, 1)):
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index e2026a4ed..3e06f3ec6 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -38,62 +38,88 @@ def ResidualFeedForward(feature_depth,
   )
 
 
-def TransformerEncoder(mode='train',
-                       num_layers=6,
-                       feature_depth=512,
-                       feedforward_depth=2048,
-                       num_heads=8,
-                       dropout=0.1):
-  """Transformer Encoder Stack.
+def EncoderLayer(feature_depth,
+                 feedforward_depth,
+                 num_heads,
+                 dropout,
+                 mode):
+  """Transformer encoder layer.
+
+  The input to the encoder is a pair (embedded source, mask) where
+  the mask is created from the original source to prevent attending
+  to the padding part of the input.
 
   Args:
-    mode: str: 'train' or 'eval'
-    num_layers: int: number of encoder/decoder layers
     feature_depth: int:  depth of embedding
     feedforward_depth: int: depth of feed-forward layer
     num_heads: int: number of attention heads
-    dropout: float: dropout rate
+    dropout: float: dropout rate (how much to drop out)
+    mode: str: 'train' or 'eval'
 
   Returns:
-    A layer for implementing a raw Transformer encoder stack.  No embedding
-    or positional signals are added by this layer.
+    the layer, returning a pair (actiavtions, mask).
   """
-  # Multi-headed Attention and Feed-forward layers
-  multi_attention = layers.MultiHeadedAttention(
-      feature_depth, num_heads=num_heads, dropout=dropout, mode=mode)
+  # The encoder block expects (activation, mask) as input and returns
+  # the new activations only, we add the mask back to output next.
+  encoder_block = layers.Serial(
+      layers.Residual(  # Attention block here.
+          layers.Parallel(layers.LayerNorm(), layers.Identity()),
+          layers.MultiHeadedAttention(feature_depth, num_heads=num_heads,
+                                      dropout=dropout, mode=mode),
+          layers.Dropout(rate=dropout, mode=mode),
+          shortcut=layers.FirstBranch()
+      ),
+      ResidualFeedForward(feature_depth, feedforward_depth, dropout, mode=mode)
+  )
+  # Now we add the mask back.
+  return layers.Serial(
+      layers.Reorder(output=((0, 1), 1)),  # (x, mask) --> ((x, mask), mask)
+      layers.Parallel(encoder_block, layers.Identity())
+  )
 
-  @layers.Lambda
-  def Encoder(embedded_source, source_mask):
-    """Transformer encoder stack.
 
-    Args:
-      embedded_source: layer variable: embedded source sequences
-      source_mask: layer variable: self-attention mask
+def TransformerEncoder(vocab_size,
+                       num_classes=10,
+                       feature_depth=512,
+                       feedforward_depth=2048,
+                       num_layers=6,
+                       num_heads=8,
+                       dropout=0.1,
+                       max_len=2048,
+                       mode='train'):
+  """Transformer encoder.
 
-    Returns:
-      Layer variable that outputs encoded source.
-    """
-    encoder_layer = layers.Serial(
-        # input attends to self
-        layers.Residual(layers.LayerNorm(),
-                        layers.FanOut(size=4),
-                        layers.Parallel(layers.Identity(),  # query
-                                        layers.Identity(),  # key
-                                        layers.Identity(),  # value
-                                        source_mask),  # attention mask
-                        multi_attention,
-                        layers.Dropout(rate=dropout, mode=mode)),
-        # feed-forward
-        ResidualFeedForward(
-            feature_depth, feedforward_depth, dropout, mode=mode)
-    )
-    return layers.Serial(
-        embedded_source,
-        layers.repeat(encoder_layer, num_layers),
-        layers.LayerNorm(),
-    )
+  Args:
+    vocab_size: int: vocab size
+    num_classes: how many classes on output
+    feature_depth: int:  depth of embedding
+    feedforward_depth: int: depth of feed-forward layer
+    num_layers: int: number of encoder/decoder layers
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate (how much to drop out)
+    max_len: int: maximum symbol length for positional encoding
+    mode: str: 'train' or 'eval'
 
-  return Encoder
+  Returns:
+    the Transformer encoder layer.
+  """
+  input_embedding = layers.Serial(
+      layers.Embedding(feature_depth, vocab_size),
+      layers.Dropout(rate=dropout, mode=mode),
+      layers.PositionalEncoding(max_len=max_len)
+  )
+  return layers.Serial(
+      layers.Branch(),  # Branch input to create embedding and mask.
+      layers.Parallel(input_embedding, layers.PaddingMask()),
+      layers.Serial(*[EncoderLayer(feature_depth, feedforward_depth, num_heads,
+                                   dropout, mode)
+                      for _ in range(num_layers)]),
+      layers.FirstBranch(),  # Drop the mask.
+      layers.LayerNorm(),
+      layers.Mean(axis=1),  # Average on length.
+      layers.Dense(num_classes),
+      layers.LogSoftmax()
+  )
 
 
 def DecoderLayer(feature_depth,
@@ -111,15 +137,13 @@ def DecoderLayer(feature_depth,
     mode: str: 'train' or 'eval'
 
   Returns:
-    init and apply.
+    the layer.
   """
   return layers.Serial(
       layers.Residual(  # Self-attention block.
           layers.LayerNorm(),
-          layers.FanOut(size=4),
-          layers.Parallel(layers.Identity(),  # query
-                          layers.Identity(),  # key
-                          layers.Identity(),  # value
+          layers.Branch(),
+          layers.Parallel(layers.Identity(),  # activation for (q, k, v)
                           layers.CausalMask(axis=-2)),  # attention mask
           layers.MultiHeadedAttention(feature_depth, num_heads=num_heads,
                                       dropout=dropout, mode=mode),
@@ -150,7 +174,7 @@ def TransformerLM(vocab_size,
     mode: str: 'train' or 'eval'
 
   Returns:
-    init and apply.
+    the layer.
   """
   return layers.Serial(
       layers.ShiftRight(),
@@ -237,7 +261,7 @@ def Encoder(source, source_mask):
     encoder_layer = layers.Serial(
         # input attends to self
         layers.Residual(layers.LayerNorm(),
-                        layers.FanOut(size=4),
+                        layers.Branch(size=4),
                         layers.Parallel(layers.Identity(),  # query
                                         layers.Identity(),  # key
                                         layers.Identity(),  # value
@@ -272,7 +296,7 @@ def Decoder(memory, target, target_mask, memory_mask):
     decoder_layer = layers.Serial(
         # target attends to self
         layers.Residual(layers.LayerNorm(),
-                        layers.FanOut(size=4),
+                        layers.Branch(size=4),
                         layers.Parallel(layers.Identity(),  # query
                                         layers.Identity(),  # key
                                         layers.Identity(),  # value
@@ -281,7 +305,7 @@ def Decoder(memory, target, target_mask, memory_mask):
                         layers.Dropout(dropout, mode=mode)),
         # target attends to encoded source
         layers.Residual(layers.LayerNorm(),
-                        layers.FanOut(size=4),
+                        layers.Branch(size=4),
                         layers.Parallel(layers.Identity(),  # query
                                         memory,  # key
                                         memory,  # value
@@ -316,7 +340,7 @@ def Generator(encoded_target):
     )
 
   # Model-Building and Evaluation Functions
-  # Get entire model's init and apply pair
+  # Get entire model's the layer pair
   top_init, top_apply = Generator(transformer)
 
   # By default act as a normal constructor and emit an (init, apply) pair.
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index f2e26441c..db21ad7af 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -121,7 +121,7 @@ def policy_and_value_net(rng_key,
   # Now, with the current logits, one head computes action probabilities and the
   # other computes the value function.
   # NOTE: The LogSoftmax instead of the Softmax.
-  cur_layers.extend([layers.FanOut(), layers.Parallel(
+  cur_layers.extend([layers.Branch(), layers.Parallel(
       layers.Serial(layers.Dense(num_actions), layers.LogSoftmax()),
       layers.Dense(1)
   )])

From f84945adc3d91d83df9943a6eada034bb46eec43 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 24 Apr 2019 15:50:18 -0700
Subject: [PATCH 1944/2720] Import T2T problems at top to ease use in colab.

PiperOrigin-RevId: 245133138
---
 tensor2tensor/trax/inputs.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 7e9b88177..ee242507d 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -27,6 +27,7 @@
 
 import numpy as np
 
+from tensor2tensor import problems_colab as t2t_problems
 import tensorflow as tf
 import tensorflow_datasets as tfds
 
@@ -217,9 +218,8 @@ def _select_features(example, feature_list=None):
 
 def _train_and_eval_dataset_v1(problem_name, data_dir):
   """Return train and evaluation datasets, feature info and supervised keys."""
-  from tensor2tensor import problems  # pylint: disable=g-import-not-at-top
   assert not tf.executing_eagerly(), "tf.eager mode must be turned off."
-  problem = problems.problem(problem_name)
+  problem = t2t_problems.problem(problem_name)
   train_dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, data_dir)
   train_dataset = train_dataset.map(_select_features)
   eval_dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, data_dir)

From 6d3444284105c80631f5c4511df8247cfe8e6fce Mon Sep 17 00:00:00 2001
From: Bairen Yi <byi@connect.ust.hk>
Date: Fri, 26 Apr 2019 02:35:05 +0800
Subject: [PATCH 1945/2720] Update t2t_trainer.py (#1557)

Fixes #1528.

Signed-off-by: Bairen Yi <byi@connect.ust.hk>
---
 tensor2tensor/bin/t2t_trainer.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index e317f8078..0c76819c6 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -372,7 +372,8 @@ def main(argv):
   # Create HParams.
   if argv:
     set_hparams_from_args(argv[1:])
-  hparams = create_hparams()
+  if FLAGS.schedule != "run_std_server":
+    hparams = create_hparams()
 
   if FLAGS.schedule == "train" or FLAGS.schedule == "train_eval_and_decode":
     mlperf_log.transformer_print(key=mlperf_log.RUN_START, hparams=hparams)

From e640eb09ff961f3becf7234b871a655afa332edc Mon Sep 17 00:00:00 2001
From: Bairen Yi <byi@connect.ust.hk>
Date: Fri, 26 Apr 2019 02:35:19 +0800
Subject: [PATCH 1946/2720] Update trainer_lib.py (#1556)

Fixes #1299.

Signed-off-by: Bairen Yi <byi@connect.ust.hk>
---
 tensor2tensor/utils/trainer_lib.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 2788bbd09..f6016cfbb 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -532,7 +532,7 @@ def run_std_server(self):
         config.cluster_spec,
         job_name=config.task_type,
         task_index=config.task_id,
-        protocol=self._hparams.std_server_protocol)
+        protocol=config.protocol)
     server.join()
 
   def decode(self,

From 99c24d581f718b210999baeca573c797087a4ff0 Mon Sep 17 00:00:00 2001
From: Randall Lin <randall@fathomhealth.co>
Date: Thu, 25 Apr 2019 11:35:36 -0700
Subject: [PATCH 1947/2720] Update problem.py (#1555)

---
 tensor2tensor/data_generators/problem.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index f9998c850..f1a4da9da 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -356,18 +356,22 @@ def eval_metrics(self):
         metrics.Metrics.ACC, metrics.Metrics.ACC_TOP5,
         metrics.Metrics.ACC_PER_SEQ, metrics.Metrics.NEG_LOG_PERPLEXITY
     ]
+  
+  @property
+  def all_metrics_fns(self):
+    return metrics.METRICS_FNS
 
   def eval_metric_fns(self, model_hparams):
     del model_hparams
     metric_names = self.eval_metrics()
-    if not all([m in metrics.METRICS_FNS for m in metric_names]):
+    if not all([m in self.all_metrics_fns for m in metric_names]):
       error_str = ("Unrecognized metric. Problem %s specified metrics "
                    "%s. Recognized metrics are %s.")
       raise ValueError(error_str % (self.name,
                                     metric_names,
-                                    list(metrics.METRICS_FNS.keys())))
+                                    list(self.all_metrics_fns.keys())))
     return {
-        metric_name: metrics.METRICS_FNS[metric_name]
+        metric_name: self.all_metrics_fns[metric_name]
         for metric_name in metric_names
     }
 

From 6654134b9f72e262738cdc8eff6b0c410dc0ee85 Mon Sep 17 00:00:00 2001
From: ksboy <weihouks@gmail.com>
Date: Fri, 26 Apr 2019 02:36:20 +0800
Subject: [PATCH 1948/2720] Add Yelp_Review Dataset (#1551)

* add yelp_review_polarity

yelp_review_polarity is a classical dataset for sentiment analysis, which is proposed by Xiang Zhang in Paper Character-level Convolutional Networks for Text Classification.
Here is the url: https://s3.amazonaws.com/fast-ai-nlp/yelp_review_polarity_csv.tgz

* add  yelp_review_full

yelp_review_full is a classical dataset for sentiment analysis, which is proposed by Xiang Zhang in Paper Character-level Convolutional Networks for Text Classification.
Here is the url: https://s3.amazonaws.com/fast-ai-nlp/yelp_review_full_csv.tgz

* Update all_problems.py
---
 tensor2tensor/data_generators/all_problems.py |   2 +
 tensor2tensor/data_generators/yelp_full.py    | 107 ++++++++++++++++++
 .../data_generators/yelp_polarity.py          | 107 ++++++++++++++++++
 3 files changed, 216 insertions(+)
 create mode 100644 tensor2tensor/data_generators/yelp_full.py
 create mode 100644 tensor2tensor/data_generators/yelp_polarity.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 5ac52724c..0315d918f 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -93,6 +93,8 @@
     "tensor2tensor.data_generators.wikitext103",
     "tensor2tensor.data_generators.wsj_parsing",
     "tensor2tensor.data_generators.wnli",
+    "tensor2tensor.data_generators.yelp_polarity",
+    "tensor2tensor.data_generators.yelp_full",
     "tensor2tensor.envs.mujoco_problems",
     "tensor2tensor.envs.tic_tac_toe_env_problem",
 ]
diff --git a/tensor2tensor/data_generators/yelp_full.py b/tensor2tensor/data_generators/yelp_full.py
new file mode 100644
index 000000000..f32c7a3f9
--- /dev/null
+++ b/tensor2tensor/data_generators/yelp_full.py
@@ -0,0 +1,107 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""IMDB Sentiment Classification Problem."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import tarfile
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+@registry.register_problem
+class SentimentYelpFull(text_problems.Text2ClassProblem):
+  """IMDB sentiment classification."""
+  URL = "https://s3.amazonaws.com/fast-ai-nlp/yelp_review_full_csv.tgz"
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+      "split": problem.DatasetSplit.TRAIN,
+      "shards": 10,
+    }, {
+      "split": problem.DatasetSplit.EVAL,
+      "shards": 1,
+    }]
+
+  @property
+  def approx_vocab_size(self):
+    return 2 ** 13  # 8k vocab suffices for this small dataset.
+
+  @property
+  def num_classes(self):
+    return 5
+
+  def class_labels(self, data_dir):
+    del data_dir
+    return ["1", "2", "3", "4", "5"]
+
+  def doc_generator(self, yelp_dir, dataset, include_label=False):
+
+    file_path = os.path.join(yelp_dir, dataset + ".csv")
+    with tf.gfile.Open(file_path) as yelp_f:
+      lines = yelp_f.readlines()  # 接收数据
+      for line in lines:  # 遍历数据
+        label = line[1]
+        doc = line[5:-2].strip()
+        # print(line[1], line[5:-2].strip())
+        if include_label:
+          yield doc, label
+        else:
+          yield doc
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    """Generate examples."""
+    # Download and extract
+    compressed_filename = os.path.basename(self.URL)
+    download_path = generator_utils.maybe_download(tmp_dir, compressed_filename,
+                                                   self.URL)
+    yelp_dir = os.path.join(tmp_dir, "yelp_review_full_csv")
+    if not tf.gfile.Exists(yelp_dir):
+      with tarfile.open(download_path, "r:gz") as tar:
+        tar.extractall(tmp_dir)
+
+    # Generate examples
+    train = dataset_split == problem.DatasetSplit.TRAIN
+    dataset = "train" if train else "test"
+    for doc, label in self.doc_generator(yelp_dir, dataset, include_label=True):
+      yield {
+        "inputs": doc,
+        "label": int(label),
+      }
+
+
+@registry.register_problem
+class SentimentYelpFullCharacters(SentimentYelpFull):
+  """IMDB sentiment classification, character level."""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  def global_task_id(self):
+    return problem.TaskID.EN_CHR_SENT
diff --git a/tensor2tensor/data_generators/yelp_polarity.py b/tensor2tensor/data_generators/yelp_polarity.py
new file mode 100644
index 000000000..aba5ae117
--- /dev/null
+++ b/tensor2tensor/data_generators/yelp_polarity.py
@@ -0,0 +1,107 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""IMDB Sentiment Classification Problem."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import tarfile
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+@registry.register_problem
+class SentimentYelpPolarity(text_problems.Text2ClassProblem):
+  """IMDB sentiment classification."""
+  URL = "https://s3.amazonaws.com/fast-ai-nlp/yelp_review_polarity_csv.tgz"
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+      "split": problem.DatasetSplit.TRAIN,
+      "shards": 10,
+    }, {
+      "split": problem.DatasetSplit.EVAL,
+      "shards": 1,
+    }]
+
+  @property
+  def approx_vocab_size(self):
+    return 2 ** 13  # 8k vocab suffices for this small dataset.
+
+  @property
+  def num_classes(self):
+    return 2
+
+  def class_labels(self, data_dir):
+    del data_dir
+    return ["1", "2"]
+
+  def doc_generator(self, yelp_dir, dataset, include_label=False):
+
+    file_path = os.path.join(yelp_dir, dataset + ".csv")
+    with tf.gfile.Open(file_path) as yelp_f:
+      lines = yelp_f.readlines()  # 接收数据
+      for line in lines:  # 遍历数据
+        label = line[1]
+        doc = line[5:-2].strip()
+        # print(line[1], line[5:-2].strip())
+        if include_label:
+          yield doc, label
+        else:
+          yield doc
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    """Generate examples."""
+    # Download and extract
+    compressed_filename = os.path.basename(self.URL)
+    download_path = generator_utils.maybe_download(tmp_dir, compressed_filename,
+                                                   self.URL)
+    yelp_dir = os.path.join(tmp_dir, "yelp_review_polarity_csv")
+    if not tf.gfile.Exists(yelp_dir):
+      with tarfile.open(download_path, "r:gz") as tar:
+        tar.extractall(tmp_dir)
+
+    # Generate examples
+    train = dataset_split == problem.DatasetSplit.TRAIN
+    dataset = "train" if train else "test"
+    for doc, label in self.doc_generator(yelp_dir, dataset, include_label=True):
+      yield {
+        "inputs": doc,
+        "label": int(label),
+      }
+
+
+@registry.register_problem
+class SentimentYelpPolarityCharacters(SentimentYelpPolarity):
+  """IMDB sentiment classification, character level."""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  def global_task_id(self):
+    return problem.TaskID.EN_CHR_SENT

From f58e378bbf0835f6fa67451cae1b7e852859fbde Mon Sep 17 00:00:00 2001
From: Souradip Mookerjee <souramoo@gmx.com>
Date: Thu, 25 Apr 2019 19:36:27 +0100
Subject: [PATCH 1949/2720] Fixed typo in _target_modality_is_real - was
 checking for Real_ when should have checked for real_ (#1550)

---
 tensor2tensor/utils/t2t_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 88988574b..7fd4a14d4 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -308,7 +308,7 @@ def _target_modality_is_real(self):
     modality_name = self._hparams.name.get(
         "targets",
         modalities.get_name(modality))(self._hparams, vocab_size)
-    return modality_name.startswith("Real")
+    return modality_name.startswith("real")
 
   def call(self, inputs, **kwargs):
     del kwargs

From 7b2947d40912637d5f52b805414e0e3a8d764f97 Mon Sep 17 00:00:00 2001
From: Joost van Amersfoort <joost.van.amersfoort@gmail.com>
Date: Thu, 25 Apr 2019 19:37:34 +0100
Subject: [PATCH 1950/2720] Fix glow (#1545)

Currently running `glow` fails with `assert data_dir`

This fixes the code to pass data_dir and make it run again.
---
 tensor2tensor/models/research/glow.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 9fa39d742..ec32525d9 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -128,7 +128,8 @@ def create_init_batch(self, features):
       init_features: initialization features.
     """
     train_dataset = self.hparams.problem.dataset(
-        tf.estimator.ModeKeys.TRAIN, hparams=self.hparams)
+        tf.estimator.ModeKeys.TRAIN, data_dir=self.hparams.data_dir,
+        hparams=self.hparams)
     train_dataset = train_dataset.batch(self.hparams.init_batch_size)
     train_dataset = self.init_preprocess(train_dataset)
     return train_dataset.make_one_shot_iterator().get_next()

From 5ee6f134e330c1a674e2833ec1ab7e3812cd7b23 Mon Sep 17 00:00:00 2001
From: Bairen Yi <byi@connect.ust.hk>
Date: Thu, 25 Apr 2019 11:38:12 -0700
Subject: [PATCH 1951/2720] Merge of PR #1557

PiperOrigin-RevId: 245276372
---
 tensor2tensor/data_generators/all_problems.py |   2 -
 tensor2tensor/data_generators/problem.py      |  10 +-
 tensor2tensor/data_generators/yelp_full.py    | 107 ------------------
 .../data_generators/yelp_polarity.py          | 107 ------------------
 tensor2tensor/models/research/glow.py         |   3 +-
 tensor2tensor/utils/t2t_model.py              |   2 +-
 tensor2tensor/utils/trainer_lib.py            |   2 +-
 7 files changed, 6 insertions(+), 227 deletions(-)
 delete mode 100644 tensor2tensor/data_generators/yelp_full.py
 delete mode 100644 tensor2tensor/data_generators/yelp_polarity.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 0315d918f..5ac52724c 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -93,8 +93,6 @@
     "tensor2tensor.data_generators.wikitext103",
     "tensor2tensor.data_generators.wsj_parsing",
     "tensor2tensor.data_generators.wnli",
-    "tensor2tensor.data_generators.yelp_polarity",
-    "tensor2tensor.data_generators.yelp_full",
     "tensor2tensor.envs.mujoco_problems",
     "tensor2tensor.envs.tic_tac_toe_env_problem",
 ]
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index f1a4da9da..f9998c850 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -356,22 +356,18 @@ def eval_metrics(self):
         metrics.Metrics.ACC, metrics.Metrics.ACC_TOP5,
         metrics.Metrics.ACC_PER_SEQ, metrics.Metrics.NEG_LOG_PERPLEXITY
     ]
-  
-  @property
-  def all_metrics_fns(self):
-    return metrics.METRICS_FNS
 
   def eval_metric_fns(self, model_hparams):
     del model_hparams
     metric_names = self.eval_metrics()
-    if not all([m in self.all_metrics_fns for m in metric_names]):
+    if not all([m in metrics.METRICS_FNS for m in metric_names]):
       error_str = ("Unrecognized metric. Problem %s specified metrics "
                    "%s. Recognized metrics are %s.")
       raise ValueError(error_str % (self.name,
                                     metric_names,
-                                    list(self.all_metrics_fns.keys())))
+                                    list(metrics.METRICS_FNS.keys())))
     return {
-        metric_name: self.all_metrics_fns[metric_name]
+        metric_name: metrics.METRICS_FNS[metric_name]
         for metric_name in metric_names
     }
 
diff --git a/tensor2tensor/data_generators/yelp_full.py b/tensor2tensor/data_generators/yelp_full.py
deleted file mode 100644
index f32c7a3f9..000000000
--- a/tensor2tensor/data_generators/yelp_full.py
+++ /dev/null
@@ -1,107 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""IMDB Sentiment Classification Problem."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-import tarfile
-from tensor2tensor.data_generators import generator_utils
-from tensor2tensor.data_generators import problem
-from tensor2tensor.data_generators import text_problems
-from tensor2tensor.utils import registry
-
-import tensorflow as tf
-
-
-@registry.register_problem
-class SentimentYelpFull(text_problems.Text2ClassProblem):
-  """IMDB sentiment classification."""
-  URL = "https://s3.amazonaws.com/fast-ai-nlp/yelp_review_full_csv.tgz"
-
-  @property
-  def is_generate_per_split(self):
-    return True
-
-  @property
-  def dataset_splits(self):
-    return [{
-      "split": problem.DatasetSplit.TRAIN,
-      "shards": 10,
-    }, {
-      "split": problem.DatasetSplit.EVAL,
-      "shards": 1,
-    }]
-
-  @property
-  def approx_vocab_size(self):
-    return 2 ** 13  # 8k vocab suffices for this small dataset.
-
-  @property
-  def num_classes(self):
-    return 5
-
-  def class_labels(self, data_dir):
-    del data_dir
-    return ["1", "2", "3", "4", "5"]
-
-  def doc_generator(self, yelp_dir, dataset, include_label=False):
-
-    file_path = os.path.join(yelp_dir, dataset + ".csv")
-    with tf.gfile.Open(file_path) as yelp_f:
-      lines = yelp_f.readlines()  # 接收数据
-      for line in lines:  # 遍历数据
-        label = line[1]
-        doc = line[5:-2].strip()
-        # print(line[1], line[5:-2].strip())
-        if include_label:
-          yield doc, label
-        else:
-          yield doc
-
-  def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    """Generate examples."""
-    # Download and extract
-    compressed_filename = os.path.basename(self.URL)
-    download_path = generator_utils.maybe_download(tmp_dir, compressed_filename,
-                                                   self.URL)
-    yelp_dir = os.path.join(tmp_dir, "yelp_review_full_csv")
-    if not tf.gfile.Exists(yelp_dir):
-      with tarfile.open(download_path, "r:gz") as tar:
-        tar.extractall(tmp_dir)
-
-    # Generate examples
-    train = dataset_split == problem.DatasetSplit.TRAIN
-    dataset = "train" if train else "test"
-    for doc, label in self.doc_generator(yelp_dir, dataset, include_label=True):
-      yield {
-        "inputs": doc,
-        "label": int(label),
-      }
-
-
-@registry.register_problem
-class SentimentYelpFullCharacters(SentimentYelpFull):
-  """IMDB sentiment classification, character level."""
-
-  @property
-  def vocab_type(self):
-    return text_problems.VocabType.CHARACTER
-
-  def global_task_id(self):
-    return problem.TaskID.EN_CHR_SENT
diff --git a/tensor2tensor/data_generators/yelp_polarity.py b/tensor2tensor/data_generators/yelp_polarity.py
deleted file mode 100644
index aba5ae117..000000000
--- a/tensor2tensor/data_generators/yelp_polarity.py
+++ /dev/null
@@ -1,107 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""IMDB Sentiment Classification Problem."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-import tarfile
-from tensor2tensor.data_generators import generator_utils
-from tensor2tensor.data_generators import problem
-from tensor2tensor.data_generators import text_problems
-from tensor2tensor.utils import registry
-
-import tensorflow as tf
-
-
-@registry.register_problem
-class SentimentYelpPolarity(text_problems.Text2ClassProblem):
-  """IMDB sentiment classification."""
-  URL = "https://s3.amazonaws.com/fast-ai-nlp/yelp_review_polarity_csv.tgz"
-
-  @property
-  def is_generate_per_split(self):
-    return True
-
-  @property
-  def dataset_splits(self):
-    return [{
-      "split": problem.DatasetSplit.TRAIN,
-      "shards": 10,
-    }, {
-      "split": problem.DatasetSplit.EVAL,
-      "shards": 1,
-    }]
-
-  @property
-  def approx_vocab_size(self):
-    return 2 ** 13  # 8k vocab suffices for this small dataset.
-
-  @property
-  def num_classes(self):
-    return 2
-
-  def class_labels(self, data_dir):
-    del data_dir
-    return ["1", "2"]
-
-  def doc_generator(self, yelp_dir, dataset, include_label=False):
-
-    file_path = os.path.join(yelp_dir, dataset + ".csv")
-    with tf.gfile.Open(file_path) as yelp_f:
-      lines = yelp_f.readlines()  # 接收数据
-      for line in lines:  # 遍历数据
-        label = line[1]
-        doc = line[5:-2].strip()
-        # print(line[1], line[5:-2].strip())
-        if include_label:
-          yield doc, label
-        else:
-          yield doc
-
-  def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    """Generate examples."""
-    # Download and extract
-    compressed_filename = os.path.basename(self.URL)
-    download_path = generator_utils.maybe_download(tmp_dir, compressed_filename,
-                                                   self.URL)
-    yelp_dir = os.path.join(tmp_dir, "yelp_review_polarity_csv")
-    if not tf.gfile.Exists(yelp_dir):
-      with tarfile.open(download_path, "r:gz") as tar:
-        tar.extractall(tmp_dir)
-
-    # Generate examples
-    train = dataset_split == problem.DatasetSplit.TRAIN
-    dataset = "train" if train else "test"
-    for doc, label in self.doc_generator(yelp_dir, dataset, include_label=True):
-      yield {
-        "inputs": doc,
-        "label": int(label),
-      }
-
-
-@registry.register_problem
-class SentimentYelpPolarityCharacters(SentimentYelpPolarity):
-  """IMDB sentiment classification, character level."""
-
-  @property
-  def vocab_type(self):
-    return text_problems.VocabType.CHARACTER
-
-  def global_task_id(self):
-    return problem.TaskID.EN_CHR_SENT
diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index ec32525d9..9fa39d742 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -128,8 +128,7 @@ def create_init_batch(self, features):
       init_features: initialization features.
     """
     train_dataset = self.hparams.problem.dataset(
-        tf.estimator.ModeKeys.TRAIN, data_dir=self.hparams.data_dir,
-        hparams=self.hparams)
+        tf.estimator.ModeKeys.TRAIN, hparams=self.hparams)
     train_dataset = train_dataset.batch(self.hparams.init_batch_size)
     train_dataset = self.init_preprocess(train_dataset)
     return train_dataset.make_one_shot_iterator().get_next()
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 7fd4a14d4..88988574b 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -308,7 +308,7 @@ def _target_modality_is_real(self):
     modality_name = self._hparams.name.get(
         "targets",
         modalities.get_name(modality))(self._hparams, vocab_size)
-    return modality_name.startswith("real")
+    return modality_name.startswith("Real")
 
   def call(self, inputs, **kwargs):
     del kwargs
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index f6016cfbb..2788bbd09 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -532,7 +532,7 @@ def run_std_server(self):
         config.cluster_spec,
         job_name=config.task_type,
         task_index=config.task_id,
-        protocol=config.protocol)
+        protocol=self._hparams.std_server_protocol)
     server.join()
 
   def decode(self,

From 46bd771ea08ba5ae8dc90a58967bc58d024a41d4 Mon Sep 17 00:00:00 2001
From: Bairen Yi <byi@connect.ust.hk>
Date: Thu, 25 Apr 2019 11:38:29 -0700
Subject: [PATCH 1952/2720] Merge of PR #1556

PiperOrigin-RevId: 245276428
---
 tensor2tensor/utils/trainer_lib.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 2788bbd09..f6016cfbb 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -532,7 +532,7 @@ def run_std_server(self):
         config.cluster_spec,
         job_name=config.task_type,
         task_index=config.task_id,
-        protocol=self._hparams.std_server_protocol)
+        protocol=config.protocol)
     server.join()
 
   def decode(self,

From 93ba3a76244f04864070326d9d5c1f590674d27c Mon Sep 17 00:00:00 2001
From: Randall Lin <randall@fathomhealth.co>
Date: Thu, 25 Apr 2019 11:52:16 -0700
Subject: [PATCH 1953/2720] Merge of PR #1555

PiperOrigin-RevId: 245278920
---
 tensor2tensor/data_generators/problem.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index f9998c850..eb92b94ed 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -357,17 +357,21 @@ def eval_metrics(self):
         metrics.Metrics.ACC_PER_SEQ, metrics.Metrics.NEG_LOG_PERPLEXITY
     ]
 
+  @property
+  def all_metrics_fns(self):
+    return metrics.METRICS_FNS
+
   def eval_metric_fns(self, model_hparams):
     del model_hparams
     metric_names = self.eval_metrics()
-    if not all([m in metrics.METRICS_FNS for m in metric_names]):
+    if not all([m in self.all_metrics_fns for m in metric_names]):
       error_str = ("Unrecognized metric. Problem %s specified metrics "
                    "%s. Recognized metrics are %s.")
       raise ValueError(error_str % (self.name,
                                     metric_names,
-                                    list(metrics.METRICS_FNS.keys())))
+                                    list(self.all_metrics_fns.keys())))
     return {
-        metric_name: metrics.METRICS_FNS[metric_name]
+        metric_name: self.all_metrics_fns[metric_name]
         for metric_name in metric_names
     }
 

From 5c6ec1041fa7a0753f57e68455ac6b54a63cb8db Mon Sep 17 00:00:00 2001
From: ksboy <weihouks@gmail.com>
Date: Thu, 25 Apr 2019 11:54:29 -0700
Subject: [PATCH 1954/2720] Merge of PR #1551

PiperOrigin-RevId: 245279344
---
 tensor2tensor/data_generators/all_problems.py |   2 +
 tensor2tensor/data_generators/yelp_full.py    | 106 ++++++++++++++++++
 .../data_generators/yelp_polarity.py          | 106 ++++++++++++++++++
 3 files changed, 214 insertions(+)
 create mode 100644 tensor2tensor/data_generators/yelp_full.py
 create mode 100644 tensor2tensor/data_generators/yelp_polarity.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 5ac52724c..0315d918f 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -93,6 +93,8 @@
     "tensor2tensor.data_generators.wikitext103",
     "tensor2tensor.data_generators.wsj_parsing",
     "tensor2tensor.data_generators.wnli",
+    "tensor2tensor.data_generators.yelp_polarity",
+    "tensor2tensor.data_generators.yelp_full",
     "tensor2tensor.envs.mujoco_problems",
     "tensor2tensor.envs.tic_tac_toe_env_problem",
 ]
diff --git a/tensor2tensor/data_generators/yelp_full.py b/tensor2tensor/data_generators/yelp_full.py
new file mode 100644
index 000000000..1ce222482
--- /dev/null
+++ b/tensor2tensor/data_generators/yelp_full.py
@@ -0,0 +1,106 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Yelp dataset."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import tarfile
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+@registry.register_problem
+class SentimentYelpFull(text_problems.Text2ClassProblem):
+  """Yelp dataset."""
+  URL = "https://s3.amazonaws.com/fast-ai-nlp/yelp_review_full_csv.tgz"
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 10,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  @property
+  def approx_vocab_size(self):
+    return 2**13  # 8k vocab suffices for this small dataset.
+
+  @property
+  def num_classes(self):
+    return 5
+
+  def class_labels(self, data_dir):
+    del data_dir
+    return ["1", "2", "3", "4", "5"]
+
+  def doc_generator(self, yelp_dir, dataset, include_label=False):
+
+    file_path = os.path.join(yelp_dir, dataset + ".csv")
+    with tf.gfile.Open(file_path) as yelp_f:
+      lines = yelp_f.readlines()
+      for line in lines:
+        label = line[1]
+        doc = line[5:-2].strip()
+        if include_label:
+          yield doc, label
+        else:
+          yield doc
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    """Generate examples."""
+    # Download and extract
+    compressed_filename = os.path.basename(self.URL)
+    download_path = generator_utils.maybe_download(tmp_dir, compressed_filename,
+                                                   self.URL)
+    yelp_dir = os.path.join(tmp_dir, "yelp_review_full_csv")
+    if not tf.gfile.Exists(yelp_dir):
+      with tarfile.open(download_path, "r:gz") as tar:
+        tar.extractall(tmp_dir)
+
+    # Generate examples
+    train = dataset_split == problem.DatasetSplit.TRAIN
+    dataset = "train" if train else "test"
+    for doc, label in self.doc_generator(yelp_dir, dataset, include_label=True):
+      yield {
+          "inputs": doc,
+          "label": int(label),
+      }
+
+
+@registry.register_problem
+class SentimentYelpFullCharacters(SentimentYelpFull):
+  """Yelp dataset, character level."""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  def global_task_id(self):
+    return problem.TaskID.EN_CHR_SENT
diff --git a/tensor2tensor/data_generators/yelp_polarity.py b/tensor2tensor/data_generators/yelp_polarity.py
new file mode 100644
index 000000000..ae1e9dce3
--- /dev/null
+++ b/tensor2tensor/data_generators/yelp_polarity.py
@@ -0,0 +1,106 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Yelp dataset."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import tarfile
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+@registry.register_problem
+class SentimentYelpPolarity(text_problems.Text2ClassProblem):
+  """Yelp dataset."""
+  URL = "https://s3.amazonaws.com/fast-ai-nlp/yelp_review_polarity_csv.tgz"
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def dataset_splits(self):
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 10,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  @property
+  def approx_vocab_size(self):
+    return 2**13  # 8k vocab suffices for this small dataset.
+
+  @property
+  def num_classes(self):
+    return 2
+
+  def class_labels(self, data_dir):
+    del data_dir
+    return ["1", "2"]
+
+  def doc_generator(self, yelp_dir, dataset, include_label=False):
+
+    file_path = os.path.join(yelp_dir, dataset + ".csv")
+    with tf.gfile.Open(file_path) as yelp_f:
+      lines = yelp_f.readlines()
+      for line in lines:
+        label = line[1]
+        doc = line[5:-2].strip()
+        if include_label:
+          yield doc, label
+        else:
+          yield doc
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    """Generate examples."""
+    # Download and extract
+    compressed_filename = os.path.basename(self.URL)
+    download_path = generator_utils.maybe_download(tmp_dir, compressed_filename,
+                                                   self.URL)
+    yelp_dir = os.path.join(tmp_dir, "yelp_review_polarity_csv")
+    if not tf.gfile.Exists(yelp_dir):
+      with tarfile.open(download_path, "r:gz") as tar:
+        tar.extractall(tmp_dir)
+
+    # Generate examples
+    train = dataset_split == problem.DatasetSplit.TRAIN
+    dataset = "train" if train else "test"
+    for doc, label in self.doc_generator(yelp_dir, dataset, include_label=True):
+      yield {
+          "inputs": doc,
+          "label": int(label),
+      }
+
+
+@registry.register_problem
+class SentimentYelpPolarityCharacters(SentimentYelpPolarity):
+  """Yelp dataset, character level."""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  def global_task_id(self):
+    return problem.TaskID.EN_CHR_SENT

From 1dbd48439247ebf317c28ae62506db2935f5340e Mon Sep 17 00:00:00 2001
From: Souradip Mookerjee <souramoo@gmx.com>
Date: Thu, 25 Apr 2019 11:55:23 -0700
Subject: [PATCH 1955/2720] Merge of PR #1550

PiperOrigin-RevId: 245279496
---
 tensor2tensor/utils/t2t_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 88988574b..7fd4a14d4 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -308,7 +308,7 @@ def _target_modality_is_real(self):
     modality_name = self._hparams.name.get(
         "targets",
         modalities.get_name(modality))(self._hparams, vocab_size)
-    return modality_name.startswith("Real")
+    return modality_name.startswith("real")
 
   def call(self, inputs, **kwargs):
     del kwargs

From d529be36662d7687483d3f5d803090c377137fbc Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 25 Apr 2019 14:14:32 -0700
Subject: [PATCH 1956/2720] Add fixed schedule lr decay. Also allow for train
 to specify a set of steps to save the model.

PiperOrigin-RevId: 245305162
---
 tensor2tensor/trax/learning_rate.py |  9 ++++++++-
 tensor2tensor/trax/trax.py          | 17 +++++++++++++++--
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index eb7017f12..9146ec176 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -42,19 +42,24 @@
 def MultifactorSchedule(history=None,
                         factors="constant * linear_warmup * rsqrt_decay",
                         constant=0.1,
-                        warmup_steps=100):
+                        warmup_steps=100,
+                        decay_factor=0.5,
+                        steps_per_decay=20000):
   """Factor-based learning rate schedule.
 
   Interprets factors in the factors string which can consist of:
   * constant: interpreted as the constant value,
   * linear_warmup: interpreted as linear warmup until warmup_steps,
   * rsqrt_decay: divide by square root of max(step, warmup_steps)
+  * decay_every: Every k steps decay the learning rate by decay_factor.
 
   Args:
     history: the history of training and evaluation (History object).
     factors: a string with factors separated by "*" that defines the schedule.
     constant: float, the starting constant for the learning rate schedule.
     warmup_steps: how many steps to warm up for in the warmup schedule.
+    decay_factor: The amount to decay the learning rate by.
+    steps_per_decay: How often to decay the learning rate.
 
   Returns:
     a function learning_rate(step): float -> float, the step-dependent lr.
@@ -77,6 +82,8 @@ def learning_rate(step):  # pylint: disable=invalid-name
         ret *= np.minimum(1.0, step / warmup_steps)
       elif name == "rsqrt_decay":
         ret /= np.sqrt(np.maximum(step, warmup_steps))
+      elif name == "decay_every":
+        ret *= (decay_factor ** (step//steps_per_decay))
       else:
         raise ValueError("Unknown factor %s." % name)
     return ret
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 3b83c7fee..64772d6c3 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -120,11 +120,15 @@ def save_gin(output_dir, sw=None):
             jaxboard.markdownify_operative_config_str(config_str))
 
 
-def save_state(state, output_dir):
+def save_state(state, output_dir, keep=False):
   """Save State and optionally gin config."""
   params_file = os.path.join(output_dir, "model.pkl")
   with gfile.GFile(params_file, "wb") as f:
     pickle.dump((state.params, state.step, state.history), f)
+  if keep:
+    params_file = os.path.join(output_dir, "model_{}.pkl".format(state.step))
+    with gfile.GFile(params_file, "wb") as f:
+      pickle.dump((state.params, state.step, state.history), f)
   log("Model saved to %s" % params_file, stdout=False)
 
 
@@ -329,6 +333,7 @@ def train(output_dir,
           optimizer=trax_opt.adam,
           lr_schedule=lr.MultifactorSchedule,
           train_steps=1000,
+          save_steps=None,
           eval_steps=10,
           eval_frequency=100,
           num_devices=None,
@@ -348,6 +353,8 @@ def train(output_dir,
     lr_schedule: A learning rate schedule as a function that takes history and
       returns a function from step to learning rate (a float).
     train_steps: int, total number of training steps.
+    save_steps: list of integers. Keep a model file at each of the supplied save
+      steps.
     eval_steps: int, num of steps per evaluation. If None or 0, eval disabled.
     eval_frequency: int, how often to run evaluation (every eval_frequency
       steps). If None or 0, eval disabled.
@@ -359,6 +366,8 @@ def train(output_dir,
   Returns:
     trax.State
   """
+  if save_steps is None:
+    save_steps = []
   num_devices = num_devices or jax.lib.xla_bridge.device_count()
   rng = get_random_number_generator_and_set_seed(random_seed)
   gfile.makedirs(output_dir)
@@ -391,7 +400,6 @@ def train(output_dir,
   jit_update_fun = _jit_update_fun(
       model_train, loss_fun, optimizer, lr_fun, num_devices)
 
-  print()
   train_stream = inputs.train_stream()
   epoch_steps = [train_steps]  # Only training if eval_frequency is 0 or None.
   if eval_frequency and eval_steps > 0:
@@ -420,6 +428,11 @@ def train(output_dir,
       opt_state, rngs = jit_update_fun(step, opt_state, next_train_batch, rngs)
       step += 1
 
+      if step in save_steps:
+        save_state(State(params=params, step=step, history=history),
+                   output_dir,
+                   keep=True)
+
       # LR log
       if step == 1 or step % 10 == 0:
         train_sw.scalar("training/learning rate",

From e9ec51fb36ea082ee6af2085fe94118d52331709 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 25 Apr 2019 15:47:07 -0700
Subject: [PATCH 1957/2720] Disable 'scheduled_sampling_num_passes' pending
 debugging

PiperOrigin-RevId: 245321647
---
 tensor2tensor/layers/common_hparams.py |  5 ++++-
 tensor2tensor/utils/t2t_model.py       | 11 +++++++++--
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 502bbbc61..9b6a63d73 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -249,7 +249,10 @@ def basic_params1():
       scheduled_sampling_prob=0.0,
       scheduled_sampling_warmup_steps=50000,
       scheduled_sampling_gold_mixin_prob=0.5,
-      scheduled_sampling_num_passes=1,
+      # TODO(duckworthd): Uncomment when we can ascertain why adding an
+      # extra field to HParam causes test failures.
+      # scheduled_sampling_num_passes=1,
+
       # This setting controls whether to copy variables around in a daisy chain
       # (if true) or leave their placement to TensorFlow. It only affects multi
       # device training and mostly should be turned on for performance. One
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 7fd4a14d4..4bcd9e8eb 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1889,12 +1889,19 @@ def sampled_results(features, logits, mixin_prob):
 
     # Apply scheduled sampling over N passes. The logits from the (n-1)-th pass
     # will be mixed with gold tokens for conditioning in the n-th pass.
-    assert hparams.scheduled_sampling_num_passes > 0, (
+    if hasattr(hparams, "scheduled_sampling_num_passes"):
+      scheduled_sampling_num_passes = hparams.scheduled_sampling_num_passes
+    else:
+      # TODO(duckworthd): Delete once scheduled_sampling_num_passes is added to
+      # common_hparams.py.
+      scheduled_sampling_num_passes = 1
+
+    assert scheduled_sampling_num_passes > 0, (
         "hparams.scheduled_sampling_num_passes must be > 0 if "
         "hparams.scheduled_sampling_prob > 0.0")
     new_logits = logits
     new_losses = losses
-    for _ in range(hparams.scheduled_sampling_num_passes):
+    for _ in range(scheduled_sampling_num_passes):
       new_logits, new_losses = sampled_results(features, new_logits, mixin_prob)
     return new_logits, new_losses
 

From d8e8fb2be11f214406a2c952cc83e386c58192ae Mon Sep 17 00:00:00 2001
From: Dumitru Erhan <dumitru@google.com>
Date: Fri, 26 Apr 2019 13:34:16 -0700
Subject: [PATCH 1958/2720] Added longer model training option for RLMB

PiperOrigin-RevId: 245478060
---
 tensor2tensor/rl/trainer_model_based_params.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 7975ca25a..f7c56911e 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -324,6 +324,14 @@ def rlmb_base_stochastic_discrete():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_base_stochastic_discrete_75k_model_steps():
+  """Base setting with stochastic discrete model with 75k WM steps."""
+  hparams = rlmb_base_stochastic_discrete()
+  hparams.model_train_steps = 15000 * 5
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_base_stochastic_discrete_200k():
   """Base setting with stochastic discrete model with 200k steps."""

From 9c215db7aabdcece6c7dcdb2a7aa279a7731dddb Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 26 Apr 2019 13:40:28 -0700
Subject: [PATCH 1959/2720] Chunked Transformer in TRAX together with the
 needed layers and error reporting improvements.

PiperOrigin-RevId: 245479113
---
 .../chunked_transformer_imagenet64_8gb.gin    |  43 +++++
 tensor2tensor/trax/inputs.py                  |  34 ++--
 tensor2tensor/trax/layers/attention.py        | 168 ++++++++++++++++--
 tensor2tensor/trax/layers/base.py             |  27 +--
 tensor2tensor/trax/layers/combinators.py      |  72 +++++++-
 tensor2tensor/trax/layers/core.py             |   9 -
 tensor2tensor/trax/models/__init__.py         |   1 +
 tensor2tensor/trax/models/transformer.py      |  86 +++++++++
 tensor2tensor/trax/trax.py                    |  72 +++++---
 9 files changed, 445 insertions(+), 67 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin

diff --git a/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin b/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin
new file mode 100644
index 000000000..2c7fd1fb9
--- /dev/null
+++ b/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin
@@ -0,0 +1,43 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 1
+batch_fun.eval_batch_size = 8
+batch_fun.max_eval_length = 12288  # 64 * 64 * 3
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.num_chunks = 64
+inputs.data_dir = None
+inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
+inputs.input_name = 'targets'
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 0.1
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 1000
+train.eval_steps = 10
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.ChunkedTransformerLM
+train.run_debug_step = False
+train.train_steps = 500000
+
+# Parameters for ChunkedTransformerLM:
+# ==============================================================================
+ChunkedTransformerLM.dropout = 0.1
+ChunkedTransformerLM.feature_depth = 1024
+ChunkedTransformerLM.feedforward_depth = 4096
+ChunkedTransformerLM.max_len = 12288  # 64 * 64 * 3
+ChunkedTransformerLM.mode = 'train'
+ChunkedTransformerLM.num_heads = 4
+ChunkedTransformerLM.num_layers = 3
+ChunkedTransformerLM.vocab_size = 256
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index ee242507d..2142f88a0 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -55,7 +55,8 @@
 
 
 @gin.configurable(blacklist=["num_devices"])
-def inputs(num_devices, dataset_name, data_dir=None, input_name=None):
+def inputs(num_devices, dataset_name, data_dir=None, input_name=None,
+           num_chunks=0, append_targets=False):
   """Make Inputs for built-in datasets.
 
   Args:
@@ -64,6 +65,9 @@ def inputs(num_devices, dataset_name, data_dir=None, input_name=None):
       with "t2t_".
     data_dir: data directory.
     input_name: optional, name of the inputs from the dictionary.
+    num_chunks: optional, into how many pieces should we chunk (large inputs).
+    append_targets: optional, instead of inputs return a pair (inputs, targets)
+      which is useful for autoregressive models.
 
   Returns:
     trax.inputs.Inputs
@@ -75,18 +79,19 @@ def inputs(num_devices, dataset_name, data_dir=None, input_name=None):
    input_name, input_shape) = _train_and_eval_batches(
        dataset_name, data_dir, input_name, num_devices)
 
-  def train_input_fun():
-    return dataset_to_stream(train_batches, input_name)
+  def numpy_stream(dataset):
+    return dataset_to_stream(
+        dataset, input_name,
+        num_chunks=num_chunks, append_targets=append_targets)
 
-  def train_eval_input_fun():
-    return dataset_to_stream(train_eval_batches, input_name)
+  if num_chunks > 0:
+    length = input_shape[0]
+    input_shape = tuple(
+        [tuple([length // num_chunks] + list(input_shape)[1:])] * num_chunks)
 
-  def eval_input_fun():
-    return dataset_to_stream(eval_batches, input_name)
-
-  return Inputs(train_stream=train_input_fun,
-                train_eval_stream=train_eval_input_fun,
-                eval_stream=eval_input_fun,
+  return Inputs(train_stream=lambda: numpy_stream(train_batches),
+                train_eval_stream=lambda: numpy_stream(train_eval_batches),
+                eval_stream=lambda: numpy_stream(eval_batches),
                 input_shape=input_shape)
 
 
@@ -138,12 +143,17 @@ def random_minibatches():
                 input_shape=input_shape_without_batch)
 
 
-def dataset_to_stream(dataset, input_name):
+def dataset_to_stream(dataset, input_name, num_chunks=0, append_targets=False):
   """Takes a tf.Dataset and creates a numpy stream of ready batches."""
   for example in tfds.as_numpy(dataset):
     inp, out = example[0][input_name], example[1]
     if len(out.shape) > 1 and out.shape[-1] == 1:
       out = np.squeeze(out, axis=-1)
+    if num_chunks > 0:
+      inp = np.split(inp, num_chunks, axis=1)
+      out = np.split(out, num_chunks, axis=1)
+    if append_targets:
+      inp = (inp, out)
     yield inp, out
 
 
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 23f487b6a..f5440fb64 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -97,21 +97,34 @@ def LayerNorm(x, params, epsilon=1e-6, **unused_kwargs):
 def _positional_encoding_new_params(input_shape, rng, max_len=2048):  # pylint: disable=invalid-name
   """Helper: create positional encoding parameters."""
   del rng
-  feature_depth = input_shape[-1]
+  # Check if we are operating on chunked inputs by checking if the first
+  # shape is a list/tuple of shapes (otherwise it's an int or numpy array).
+  is_chunked = isinstance(input_shape[0], (list, tuple))
+  feature_depth = input_shape[0][-1] if is_chunked else input_shape[-1]
   pe = onp.zeros((max_len, feature_depth), dtype=onp.float32)
   position = onp.arange(0, max_len)[:, onp.newaxis]
   div_term = onp.exp(
       onp.arange(0, feature_depth, 2) * -(onp.log(10000.0) / feature_depth))
   pe[:, 0::2] = onp.sin(position * div_term)
   pe[:, 1::2] = onp.cos(position * div_term)
-  return np.array(pe[onp.newaxis, :])  # send to device
+  pe = pe[onp.newaxis, :, :]  # [1, max_len, feature_depth]
+  return np.array(pe)  # These are trainable parameters, initialized as above.
 
 
 @base.layer(new_parameters=_positional_encoding_new_params)
 def PositionalEncoding(x, params, **unused_kwargs):
   """Implements bare positional encoding."""
-  symbol_size = np.shape(x)[1]
-  return x + params[:, :symbol_size]
+  if not isinstance(x, (list, tuple)):  # non-chunked inputs
+    symbol_size = np.shape(x)[1]
+    return x + params[:, :symbol_size, :]
+  # Chunked case: apply to all chunks selecting as much as needed.
+  offset = 0
+  results = []
+  for chunk in x:
+    symbol_size = np.shape(chunk)[1]
+    results.append(chunk + params[:, offset:offset + symbol_size, :])
+    offset += symbol_size
+  return results
 
 
 def DotProductAttention(query, key, value, mask, dropout, mode, rng):
@@ -169,32 +182,32 @@ def apply_fun(params, inputs, **kwargs):  # pylint: disable=invalid-name
 
 
 def _multihead_attention_output_shape(  # pylint: disable=invalid-name
-    input_shapes, feature_depth=None, **unused_kwargs):
+    input_shapes, **unused_kwargs):
   """Helper: calculate multihead attention output shape."""
-  input_shape = input_shapes[0]  # Inputs are (q, k, v, mask).
-  return input_shape[:-1] + (feature_depth,)
+  q_shape = input_shapes[0][0]  # Inputs are ((q, k, v), mask).
+  return q_shape
 
 
 @base.layer(output_shape=_multihead_attention_output_shape)
-def PureMultiHeadedAttention(x, params, feature_depth=None,
-                             num_heads=8, dropout=0.0, mode='train', **kwargs):
+def PureMultiHeadedAttention(x, params, num_heads=8, dropout=0.0,
+                             mode='train', **kwargs):
   """Pure transformer-style multi-headed attention.
 
   Args:
-    x: inputs (q, k, v, mask)
+    x: inputs ((q, k, v), mask)
     params: parameters (none)
-    feature_depth: int:  depth of embedding
     num_heads: int: number of attention heads
     dropout: float: dropout rate
     mode: str: 'train' or 'eval'
     **kwargs: other arguments including the rng
 
   Returns:
-    Pure Multi-headed attention layer. (No Dense transforms on input.)
+    Pure Multi-headed attention layer (no Dense transforms on input).
   """
   del params
   rng = kwargs.get('rng', None)
   (q, k, v), mask = x
+  feature_depth = q.shape[-1]
   assert feature_depth % num_heads == 0
   head_depth = feature_depth // num_heads
   nbatch = np.shape(q)[0]
@@ -271,3 +284,134 @@ def MultiHeadedAttention(
       MultiHeadedAttentionQKV(  # pylint: disable=no-value-for-parameter
           feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),
   )
+
+
+# Chunked attention.
+def _chunked_selector_output_shape(  # pylint: disable=invalid-name
+    input_shapes, selector=None, **unused_kwargs):
+  """Helper: calculate output shape for chunked key selector (see below)."""
+  # Read the main function below first, the shape logic just follows the ops.
+  selector = selector or (lambda x: [] if x < 1 else [x-1])
+  triples, _ = zip(*input_shapes)
+  (query_shapes, key_shapes, value_shapes) = zip(*triples)
+  result = []
+  for i in range(len(input_shapes)):
+    selected = selector(i)
+    cur_key_shape, cur_value_shape = key_shapes[i], value_shapes[i]
+    # Since keys and values are [batch, length, depth] we concatenate on axis=1.
+    new_key_len = sum([key_shapes[j][1] for j in selected]) + cur_key_shape[1]
+    new_key_shape = (cur_key_shape[0], new_key_len, cur_key_shape[2])
+    new_value_len = sum(
+        [value_shapes[j][1] for j in selected]) + cur_value_shape[1]
+    new_value_shape = (cur_value_shape[0], new_value_len, cur_value_shape[2])
+    # Masks are (1, query-len, key-len).
+    new_mask_shape = (1, query_shapes[i][1], new_key_len)
+    new_shape = ((query_shapes[i], new_key_shape, new_value_shape),
+                 new_mask_shape)
+    result.append(new_shape)
+  return tuple(result)
+
+
+@base.layer(output_shape=_chunked_selector_output_shape)
+def ChunkedAttentionSelector(x, params, selector=None, **kwargs):
+  """Select which chunks to attend to in chunked attention.
+
+  Args:
+    x: inputs, a list of elements of the form (q, k, v), mask for each chunk.
+    params: parameters (unused).
+    selector: a function from chunk_number -> list of chunk numbers that says
+      which other chunks should be appended to the given one (previous if None).
+    **kwargs: unused other arguments.
+
+  Returns:
+    a list of elements of the form (q, k', v'), mask' where k', v' and mask' are
+    concatenations of k, v and identity-extended masks from selected chunks.
+  """
+  del params, kwargs
+  selector = selector or (lambda x: [] if x < 1 else [x-1])
+  triples, masks = zip(*x)
+  (queries, keys, values) = zip(*triples)
+  result = []
+  for i in range(len(x)):
+    selected = selector(i)
+    # Since keys and values are [batch, length, depth] we concatenate on axis=1.
+    # We also always include the current key or value at the end.
+    new_key_list = [keys[j] for j in selected]
+    new_key = np.concatenate(new_key_list + [keys[i]], axis=1)
+    new_value = np.concatenate(
+        [values[j] for j in selected] + [values[i]], axis=1)
+    # Masks are (1, query-len, key-len) so we concatenate on axis=2.
+    new_mask_shapes = [(1, queries[i].shape[1], key.shape[1])
+                       for key in new_key_list]
+    cur_mask = masks[i]
+    # Masks are all-1 for the added chunks (no masking).
+    new_mask_list = [np.ones(s, dtype=cur_mask.dtype) for s in new_mask_shapes]
+    # We still use the current (often causal) mask for the final chunk.
+    new_mask = np.concatenate(new_mask_list + [cur_mask], axis=2)
+    result.append(((queries[i], new_key, new_value), new_mask))
+  return tuple(result)
+
+
+def ChunkedCausalMultiHeadedAttention(
+    feature_depth, num_heads=8, dropout=0.0, chunk_selector=None, mode='train'):
+  """Transformer-style causal multi-headed attention operating on chunks.
+
+  Accepts inputs that are a list of chunks and applies causal attention.
+
+  Args:
+    feature_depth: int:  depth of embedding
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate
+    chunk_selector: a function from chunk number to list of chunks to attend.
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    Multi-headed self-attention layer.
+  """
+  prepare_attention_input = combinators.Serial(
+      combinators.Branch(),
+      combinators.Parallel(
+          combinators.Branch(num_branches=3),  # q = k = v = first input
+          CausalMask(axis=-2),  # pylint: disable=no-value-for-parameter
+      ),
+      combinators.Parallel(
+          combinators.Parallel(
+              core.Dense(feature_depth,
+                         kernel_initializer=core.XavierUniformInitializer()),
+              core.Dense(feature_depth,
+                         kernel_initializer=core.XavierUniformInitializer()),
+              core.Dense(feature_depth,
+                         kernel_initializer=core.XavierUniformInitializer()),
+          ),
+          combinators.Identity()
+      )
+  )
+  return combinators.Serial(
+      combinators.Map(prepare_attention_input),
+      ChunkedAttentionSelector(selector=chunk_selector),  # pylint: disable=no-value-for-parameter
+      combinators.Map(PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
+          feature_depth=feature_depth, num_heads=num_heads,
+          dropout=dropout, mode=mode), check_shapes=False),
+      combinators.Map(core.Dense(
+          feature_depth, kernel_initializer=core.XavierUniformInitializer())),
+  )
+
+
+@base.layer()
+def ShiftRight(x, **unused_kwargs):
+  """Layer to shift the tensor to the right by padding on axis 1."""
+  if not isinstance(x, (list, tuple)):  # non-chunked inputs
+    pad_widths = [(0, 0), (1, 0)]
+    padded = np.pad(x, pad_widths, mode='constant')
+    return padded[:, :-1]
+  # Handling chunked inputs. Recall that the list of chunks represents a big
+  # sequence (the concatenation of the chunks). We want to shift that sequence,
+  # so we put a 0 in the beginning of the first chunk and the last element of
+  # that chunk is used as the new first element of the next chunk, and so on.
+  padded = []
+  last_value = np.zeros_like(x[0][:, -1])
+  for chunk in x:
+    padded_chunk = np.concatenate([last_value[:, np.newaxis], chunk], axis=1)
+    last_value = chunk[:, -1]
+    padded.append(padded_chunk[:, :-1])
+  return padded
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index f2155f808..b7e45a641 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -88,14 +88,18 @@ def initialize(self, input_shape, rng):
     Returns:
       Newly created parameters on the first call and () on all subsequent calls.
     """
-    # Re-using this layer, no new parameters.
-    if not self._first_init:
-      return ()
+    try:
+      # Re-using this layer, no new parameters.
+      if not self._first_init:
+        return ()
 
-    # First call of this layer, create parameters.
-    self._first_init = False
-    self._params = self.new_parameters(input_shape, rng)
-    return self._params
+      # First call of this layer, create parameters.
+      self._first_init = False
+      self._params = self.new_parameters(input_shape, rng)
+      return self._params
+    except Exception:
+      name, trace = self.__class__.__name__, _short_traceback()
+      raise LayerError(name, 'initialize', self._caller, input_shape, trace)
 
   def __call__(self, x, params=(), **kwargs):
     try:
@@ -110,7 +114,7 @@ def __call__(self, x, params=(), **kwargs):
       return self.call(x, params=params, **kwargs)
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback()
-      raise LayerError(name, self._caller, shapes(x), trace)
+      raise LayerError(name, 'call', self._caller, shapes(x), trace)
 
 
 class LayerError(Exception):
@@ -120,8 +124,10 @@ class LayerError(Exception):
     message: the message corresponding to this exception.
   """
 
-  def __init__(self, layer_name, caller, input_shapes, traceback_string):
+  def __init__(self, layer_name, function_name, caller,
+               input_shapes, traceback_string):
     self._layer_name = layer_name
+    self._function_name = function_name  # Is it call or initialize?
     self._caller = caller  # Python inspect object with init caller info.
     self._traceback = traceback_string
     self._input_shapes = input_shapes
@@ -129,7 +135,8 @@ def __init__(self, layer_name, caller, input_shapes, traceback_string):
 
   @property
   def message(self):
-    prefix = 'Exception passing through layer %s:\n' % self._layer_name
+    prefix = 'Exception passing through layer '
+    prefix += '%s (in %s):\n' % (self._layer_name, self._function_name)
     short_path = '[...]/' + '/'.join(self._caller.filename.split('/')[-3:])
     caller = '  layer created in file %s, line %d\n' % (short_path,
                                                         self._caller.lineno)
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index e12f7ccc4..67d9458f2 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -43,7 +43,16 @@ def call(self, x, params=(), **kwargs):
   def output_shape(self, input_shape):
     cur_shape = input_shape
     for layer in self._layers:
-      cur_shape = layer.output_shape(cur_shape)
+      try:
+        cur_shape = layer.output_shape(cur_shape)
+      except Exception:
+        # Since this is a widely used combinator, we improve errors here.
+        # Private methods are accessed as an exception for that reason.
+        name, trace = layer.__class__.__name__, base._short_traceback()  # pylint: disable=protected-access
+        raise base.LayerError(
+            name, 'output_shape',
+            layer._caller, input_shape, trace)  # pylint: disable=protected-access
+
     return cur_shape
 
   def new_parameters(self, input_shape, rng):
@@ -121,9 +130,24 @@ def SecondBranch(x, **unused_kwargs):
   return x[1]  # Here x is a list of tensors, we select the second.
 
 
+def _nested_sum(inputs):  # pylint: disable=invalid-name
+  """Helper: sum a list of arrays or nested arrays."""
+  # First the simple non-nested case.
+  if not isinstance(inputs[0], (list, tuple)):
+    return sum(inputs)
+  # In the nested case, sum on each axis separately.
+  result_list = []
+  for i in range(len(inputs[0])):
+    result_list.append(_nested_sum([x[i] for x in inputs]))
+  if isinstance(inputs[0], list):
+    return result_list
+  return tuple(result_list)
+
+
 @base.layer(output_shape=lambda input_shape_list: input_shape_list[0])
 def SumBranches(x, **unused_kwargs):
-  return sum(x)  # Here x is a list of tensors of the same shape, we add them.
+  # Here x is a list of tensors of the same shape, or nested structures.
+  return _nested_sum(x)
 
 
 def _concatenate_shape(input_shape, axis=-1):  # pylint: disable=invalid-name
@@ -194,3 +218,47 @@ def Residual(*layers, **kwargs):
     )
   else:
     raise ValueError('Empty residual combinator.')
+
+
+class Map(base.Layer):
+  """Combinator for applying a layer to a list or tuple.
+
+  Args:
+    layer: a layer to apply to each element.
+
+  Returns:
+    A new layer representing mapping layer to all elements of the input.
+  """
+
+  def __init__(self, layer, check_shapes=True):
+    super(Map, self).__init__()
+    self._layer = layer
+    # Generally a Map should be applied to lists where all elements have
+    # the same shape -- because self._layer will only be initialized once
+    # and it could have different parameters for different shapes. But there
+    # are valid cases -- e.g., when self._layer has no parameters -- where we
+    # can apply Map to different shapes -- set check_shapes=False in such cases.
+    self._check_shapes = check_shapes
+
+  def call(self, inputs, params=(), **kwargs):
+    rng = kwargs.pop('rng', None)
+    rngs = (None,) * len(inputs)
+    if rng is not None:
+      rngs = backend.random.split(rng, len(inputs))
+    result = [self._layer(x, params=params, rng=r, **kwargs)
+              for x, r in zip(inputs, rngs)]
+    if isinstance(inputs, list):
+      return result
+    return tuple(result)
+
+  def output_shape(self, input_shapes):
+    return tuple([self._layer.output_shape(shape) for shape in input_shapes])
+
+  def new_parameters(self, input_shape, rng):
+    first_shape = input_shape[0]
+    if self._check_shapes:
+      for shape in input_shape:
+        if shape != first_shape:
+          raise ValueError('Map layer can only be applied to list of elements '
+                           'with the same shapes. Shapes: %s' % str(shape))
+    return self._layer.initialize(first_shape, rng)
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index f75c127c9..6cc541188 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -424,12 +424,3 @@ def _mean_output_shape(input_shape, axis=-1, keepdims=False):
 def Mean(x, params, axis=-1, keepdims=False, **kwargs):
   del params, kwargs
   return np.mean(x, axis=axis, keepdims=keepdims)
-
-
-@base.layer()
-def ShiftRight(x, **unused_kwargs):
-  """Layer to shift the tensor to the right by padding on axis 1."""
-  pad_widths = [(0, 0), (1, 0)]
-  pad_widths += [(0, 0) for _ in range(len(x.shape) - 2)]
-  padded = np.pad(x, pad_widths, mode='constant')
-  return padded[:, :-1, ...]
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index ff07113fc..8dd46c0de 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -37,3 +37,4 @@ def model_configure(*args, **kwargs):
 WideResnet = model_configure(resnet.WideResnet)
 TransformerEncoder = model_configure(transformer.TransformerEncoder)
 TransformerLM = model_configure(transformer.TransformerLM)
+ChunkedTransformerLM = model_configure(transformer.ChunkedTransformerLM)
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 3e06f3ec6..2ada531d2 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -191,6 +191,92 @@ def TransformerLM(vocab_size,
   )
 
 
+def ChunkedDecoderLayer(feature_depth,
+                        feedforward_depth,
+                        num_heads,
+                        dropout,
+                        chunk_selector,
+                        mode):
+  """Transformer decoder layer operating on chunks.
+
+  Args:
+    feature_depth: int:  depth of embedding
+    feedforward_depth: int: depth of feed-forward layer
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate (how much to drop out)
+    chunk_selector: a function from chunk number to list of chunks to attend.
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    the layer.
+  """
+  return layers.Serial(
+      layers.Residual(  # Self-attention block.
+          layers.Map(layers.LayerNorm()),
+          layers.ChunkedCausalMultiHeadedAttention(
+              feature_depth, num_heads=num_heads, dropout=dropout,
+              chunk_selector=chunk_selector, mode=mode),
+          layers.Map(layers.Dropout(rate=dropout, mode=mode)),
+      ),
+      layers.Map(ResidualFeedForward(
+          feature_depth, feedforward_depth, dropout, mode=mode))
+  )
+
+
+def ChunkedTransformerLM(vocab_size,
+                         feature_depth=512,
+                         feedforward_depth=2048,
+                         num_layers=6,
+                         num_heads=8,
+                         dropout=0.1,
+                         chunk_selector=None,
+                         max_len=2048,
+                         mode='train'):
+  """Transformer language model operating on chunks.
+
+  The input to this  model is a sequence presented as a list or tuple of chunks:
+    (chunk1, chunk2, chunks3, ..., chunkN).
+  Each chunk should have the same shape (batch, chunk-length) and together they
+  represent a long sequence that's a concatenation chunk1,chunk2,...,chunkN.
+
+  Chunked Transformer emulates the operation of a Transformer on this long
+  sequence except for the chunked attention layer, which may attend to only
+  a subset of the chunks to reduce memory use.
+
+  Args:
+    vocab_size: int: vocab size
+    feature_depth: int:  depth of embedding
+    feedforward_depth: int: depth of feed-forward layer
+    num_layers: int: number of encoder/decoder layers
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate (how much to drop out)
+    chunk_selector: a function from chunk number to list of chunks to attend
+      (if None, attends to the previous chunks which is equivalent to setting
+       chunk_selector(x) = [] if x < 1 else [x-1] (TransformerXL); we attend
+       to the current chunk with a causal mask too, selected chunks unmasked).
+    max_len: int: maximum symbol length for positional encoding
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    the layer.
+  """
+  stack = [ChunkedDecoderLayer(feature_depth, feedforward_depth, num_heads,
+                               dropout, chunk_selector, mode)
+           for _ in range(num_layers)]
+  # Below each Map(L) applies the layer L to each chunk independently.
+  return layers.Serial(
+      layers.ShiftRight(),
+      layers.Map(layers.Embedding(feature_depth, vocab_size)),
+      layers.Map(layers.Dropout(rate=dropout, mode=mode)),
+      layers.PositionalEncoding(max_len=max_len),
+      layers.Serial(*stack),
+      layers.Map(layers.LayerNorm()),
+      layers.Map(layers.Dense(
+          vocab_size, kernel_initializer=layers.XavierUniformInitializer())),
+      layers.Map(layers.LogSoftmax()),
+  )
+
+
 # TODO(lukaszkaiser): rewrite the model below.
 
 
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 64772d6c3..0c7295981 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -50,37 +50,64 @@
 from tensorflow.io import gfile
 
 
+def _make_list(predictions, targets):
+  """Helper: make predictions and targets lists, check they match on length."""
+  #  Our models sometimes return predictions in lists, make it a list always.
+  # TODO(lukaszkaiser): make abstractions for nested structures and refactor.
+  if not isinstance(predictions, (list, tuple)):
+    if isinstance(targets, (list, tuple)):
+      raise ValueError("Targets are a list or tuple but predictions are not.")
+    predictions, targets = [predictions], [targets]
+  if len(predictions) != len(targets):
+    raise ValueError("Predictions and targets have different lengths.")
+  return list(predictions), list(targets)
+
+
 @gin.configurable(blacklist=["inputs", "targets"])
 def masked_mean(inputs, targets, mask_id=None):
   """Mean of the inputs but counting only those where targets != mask_id."""
-  x = inputs.astype(np.float32)
+  inputs = [x.astype(np.float32) for x in inputs]
+  # We assume all elements in the list contribute equally.
+  # TODO(lukaszkaiser): remove this assumption (e.g., when masks differ).
+  length = len(inputs)
   if mask_id is None:
-    return np.mean(x)
-  unmask = 1.0 - np.equal(targets, mask_id).astype(np.float32)
-  return np.sum(x * unmask) / np.sum(unmask)
+    # TODO(lukaszkaiser): can we just divide the sum by length? XLA optimizes?
+    return sum([np.mean(x) / length for x in inputs])
+  unmask = [1.0 - np.equal(t, mask_id).astype(np.float32) for t in targets]
+  return sum([np.sum(x * m) / (length * np.sum(m))
+              for x, m in zip(inputs, unmask)])
 
 
 def accuracy(batch, model_predictions):
   """Calculate accuracy."""
   _, targets = batch
-  predicted_class = np.argmax(model_predictions, axis=-1)
-  correct = np.equal(predicted_class, targets)
+  model_predictions, targets = _make_list(model_predictions, targets)
+  correct = []
+  for (prediction, target) in zip(model_predictions, targets):
+    predicted_class = np.argmax(prediction, axis=-1)
+    correct.append(np.equal(predicted_class, target))
   return masked_mean(correct, targets)
 
 
 def neg_log_perplexity(batch, model_predictions):
   """Calculate negative log perplexity."""
   _, targets = batch
-  hot_targets = layers.one_hot(targets, model_predictions.shape[-1])
-  xent = np.sum(model_predictions * hot_targets, axis=-1)
+  model_predictions, targets = _make_list(model_predictions, targets)
+  xent = []
+  for (prediction, target) in zip(model_predictions, targets):
+    hot_target = layers.one_hot(target, prediction.shape[-1])
+    xent.append(np.sum(prediction * hot_target, axis=-1))
   return masked_mean(xent, targets)
 
 
 def loss(params, batch, model_predict, rng):
   """Calculate loss."""
   inputs, targets = batch
-  preds = model_predict(inputs, params, rng=rng)
-  xent = np.sum(preds * layers.one_hot(targets, preds.shape[-1]), axis=-1)
+  predictions = model_predict(inputs, params, rng=rng)
+  predictions, targets = _make_list(predictions, targets)
+  xent = []
+  for (pred, target) in zip(predictions, targets):
+    xent.append(np.sum(pred * layers.one_hot(target, pred.shape[-1]), axis=-1))
   return - masked_mean(xent, targets)
 
 
@@ -299,7 +326,7 @@ def update(i, opt_state, batch, rng):
   return update
 
 
-def reshape_by_device(x, num_devices):
+def _reshape_by_device_single(x, num_devices):
   """Reshape x into a shape [num_devices, ...]."""
   x_shape = list(x.shape)
   batch_size = x_shape[0]
@@ -314,15 +341,10 @@ def reshape_by_device(x, num_devices):
   return np.reshape(x, new_shape_prefix + x_shape[1:])
 
 
-def reshape_by_device_pair(train_data, num_devices):
-  """Reshape by device for a pair."""
-  x, y = train_data
-  x_shape, y_shape = list(x.shape), list(y.shape)
-  if x_shape[0] != y_shape[0]:  # Same batch size.
-    logging.fatal(
-        "Batch size is not the same for train_data pair: [%d] vs [%d]",
-        x_shape[0], y_shape[0])
-  return reshape_by_device(x, num_devices), reshape_by_device(y, num_devices)
+def reshape_by_device(x, num_devices):
+  """Reshape possibly nested x into a shape [num_devices, ...]."""
+  return layers.nested_map(
+      x, lambda x: _reshape_by_device_single(x, num_devices))
 
 
 @gin.configurable(blacklist=["output_dir"])
@@ -389,7 +411,13 @@ def train(output_dir,
   step = state.step or 0
   rng, init_rng = jax_random.split(rng)
   rngs = jax_random.split(rng, num_devices)
-  model_input_shape = tuple([-1] + list(inputs.input_shape))
+  first_shape = inputs.input_shape[0]
+  # If the inputs are a tuple/list, add [-1] (batch) to each element.
+  if isinstance(first_shape, (list, tuple)):
+    model_input_shape = tuple(
+        [tuple([-1] + list(shape)) for shape in inputs.input_shape])
+  else:  # Otherwise just add [-1] to the input shape.
+    model_input_shape = tuple([-1] + list(inputs.input_shape))
   params = state.params or model_train.initialize(model_input_shape, init_rng)
   opt_state = opt_init(params)
   if num_devices > 1:  # TODO(lukaszkaiser): use everywhere when pmap is stable.
@@ -424,7 +452,7 @@ def train(output_dir,
       # Train
       next_train_batch = next(train_stream)
       if num_devices > 1:  # TODO(lukaszkaiser): use everywhere when possible.
-        next_train_batch = reshape_by_device_pair(next_train_batch, num_devices)
+        next_train_batch = reshape_by_device(next_train_batch, num_devices)
       opt_state, rngs = jit_update_fun(step, opt_state, next_train_batch, rngs)
       step += 1
 

From d84e1baa4c2ac6b6056ee9478957018825f7c315 Mon Sep 17 00:00:00 2001
From: Niki Parmar <nikip@google.com>
Date: Fri, 26 Apr 2019 14:01:27 -0700
Subject: [PATCH 1960/2720] Dilated attention 1d

PiperOrigin-RevId: 245482925
---
 .../layers/common_image_attention.py          | 28 +++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index ddc0ecd9d..c4af3bc1a 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import numpy as np
+
 from six.moves import range  # pylint: disable=redefined-builtin
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
@@ -159,6 +161,32 @@ def local_attention_1d(x,
     return y
 
 
+def get_dilated_1d_attention_mask(
+    num_heads, block_size,
+    num_blocks, memory_size, gap_size,
+    name="dilated_mask"):
+  """Dilated attention with a masking strategy."""
+  mask = np.ones((num_heads, block_size, 2*block_size), np.bool)
+
+  # now going over every row to do the right assignment of
+  # memory blocks
+  for i in range(block_size):
+    visible = 2*block_size  - (block_size-i)
+    # You always attend to yourself, set the mask for that
+    mask[:, i, -(block_size - i)] = 0
+    # Maybe num_blocks can be automatically calculated?
+    for j in range(num_blocks):
+      for k in range(memory_size):
+        index = ((gap_size + memory_size)*j) + k
+        if index >= visible:
+          break
+        mask[:, i, -(index + block_size - i + 1)] = 0  # Verify
+
+  # adding a num blocks dimension
+  mask = np.expand_dims(mask, axis=1)
+  return tf.constant(mask, dtype=tf.int32, name=name)
+
+
 def dilated_attention_1d(x,
                          hparams,
                          attention_type="masked_dilated_1d",

From a0bf3b90b13f75e77fdacf5da025d09309165b92 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 26 Apr 2019 14:15:30 -0700
Subject: [PATCH 1961/2720] * s/from tensor2tensor.utils.hparam import
 HParams/from tensor2tensor.utils import hparams   * We should only import
 packages and not specific classes, this has led to   * hard to track bugs in
 the past. * Minor change in the scheduled_sampling code, access hparam with a
 default value.

PiperOrigin-RevId: 245485611
---
 tensor2tensor/data_generators/celeba_test.py  |  4 ++--
 .../data_generators/imagenet_test.py          |  4 ++--
 tensor2tensor/data_generators/mscoco_test.py  |  4 ++--
 tensor2tensor/data_generators/problem.py      |  6 +++---
 tensor2tensor/layers/common_hparams.py        |  6 +++---
 .../layers/common_image_attention_test.py     |  8 ++++----
 .../layers/message_passing_attention.py       |  2 +-
 tensor2tensor/layers/modalities.py            | 20 +++++++++----------
 tensor2tensor/models/research/glow_ops.py     |  4 ++--
 .../models/research/glow_ops_test.py          |  4 ++--
 tensor2tensor/models/research/rl.py           |  8 ++++----
 tensor2tensor/models/resnet.py                |  6 +++---
 tensor2tensor/models/shake_shake.py           |  4 ++--
 tensor2tensor/models/video/nfg_interpolate.py |  8 ++++----
 tensor2tensor/rl/evaluator.py                 |  8 ++++----
 tensor2tensor/rl/player_utils.py              |  4 ++--
 .../rl/trainer_model_based_params.py          | 10 +++++-----
 tensor2tensor/serving/query.py                |  4 ++--
 tensor2tensor/utils/decoding.py               |  8 ++++----
 tensor2tensor/utils/diet.py                   |  4 ++--
 tensor2tensor/utils/hparams_lib.py            |  6 +++---
 tensor2tensor/utils/misc_utils_test.py        |  4 ++--
 tensor2tensor/utils/t2t_model.py              |  9 ++-------
 tensor2tensor/utils/t2t_model_test.py         |  6 +++---
 24 files changed, 73 insertions(+), 78 deletions(-)

diff --git a/tensor2tensor/data_generators/celeba_test.py b/tensor2tensor/data_generators/celeba_test.py
index 598a081a6..a782bac16 100644
--- a/tensor2tensor/data_generators/celeba_test.py
+++ b/tensor2tensor/data_generators/celeba_test.py
@@ -21,7 +21,7 @@
 
 from absl.testing import parameterized
 from tensor2tensor.data_generators import celeba
-from tensor2tensor.utils.hparam import HParams
+from tensor2tensor.utils import hparam
 
 import tensorflow as tf
 
@@ -35,7 +35,7 @@ class CelebaTest(parameterized.TestCase, tf.test.TestCase):
   def testCelebaMultiResolutionPreprocessExample(self, resize_method):
     example = {"inputs": tf.random_uniform([218, 178, 3], minval=-1.)}
     mode = tf.estimator.ModeKeys.TRAIN
-    hparams = HParams(resolutions=[8, 16, 32])
+    hparams = hparam.HParams(resolutions=[8, 16, 32])
     if resize_method is not None:
       hparams.resize_method = resize_method
 
diff --git a/tensor2tensor/data_generators/imagenet_test.py b/tensor2tensor/data_generators/imagenet_test.py
index 66569cc0b..9ea07a43c 100644
--- a/tensor2tensor/data_generators/imagenet_test.py
+++ b/tensor2tensor/data_generators/imagenet_test.py
@@ -21,7 +21,7 @@
 
 from absl.testing import parameterized
 from tensor2tensor.data_generators import imagenet
-from tensor2tensor.utils.hparam import HParams
+from tensor2tensor.utils import hparam
 
 import tensorflow as tf
 
@@ -35,7 +35,7 @@ class ImagenetTest(parameterized.TestCase, tf.test.TestCase):
   def testImagenetMultiResolutionPreprocessExample(self, resize_method):
     example = {"inputs": tf.random_uniform([64, 64, 3], minval=-1.)}
     mode = tf.estimator.ModeKeys.TRAIN
-    hparams = HParams(resolutions=[8, 16, 32])
+    hparams = hparam.HParams(resolutions=[8, 16, 32])
     if resize_method is not None:
       hparams.resize_method = resize_method
 
diff --git a/tensor2tensor/data_generators/mscoco_test.py b/tensor2tensor/data_generators/mscoco_test.py
index ce87332a2..e2ffadd75 100644
--- a/tensor2tensor/data_generators/mscoco_test.py
+++ b/tensor2tensor/data_generators/mscoco_test.py
@@ -21,7 +21,7 @@
 
 from absl.testing import parameterized
 from tensor2tensor.data_generators import mscoco
-from tensor2tensor.utils.hparam import HParams
+from tensor2tensor.utils import hparam
 
 import tensorflow as tf
 
@@ -35,7 +35,7 @@ class MscocoTest(parameterized.TestCase, tf.test.TestCase):
   def testMsCocoMultiResolutionPreprocessExample(self, resize_method):
     example = {"inputs": tf.random_uniform([400, 400, 3], minval=-1.)}
     mode = tf.estimator.ModeKeys.TRAIN
-    hparams = HParams(resolutions=[8, 16, 32])
+    hparams = hparam.HParams(resolutions=[8, 16, 32])
     if resize_method is not None:
       hparams.resize_method = resize_method
 
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index eb92b94ed..a8f12ee77 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -27,9 +27,9 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.utils import data_reader
+from tensor2tensor.utils import hparam
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import mlperf_log
-from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 from tensorflow.contrib.tpu.python.tpu import tpu_config
@@ -131,7 +131,7 @@ class TaskID(object):
 
 
 def default_model_hparams():
-  return HParams(
+  return hparam.HParams(
       max_input_seq_length=0,
       max_target_seq_length=0,
       prepend_mode="none",
@@ -1016,7 +1016,7 @@ def _reverse_problem_hparams(p_hparams):
 
 def _default_hparams():
   """A set of basic model hyperparameters."""
-  return HParams(
+  return hparam.HParams(
       # Use this parameter to get comparable perplexity numbers with different
       # tokenizations.  This value should be set to the ratio of the number of
       # tokens in the test set according to the tokenization used to the number
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 9b6a63d73..c25381d4e 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -19,8 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 from six.moves import zip  # pylint: disable=redefined-builtin
+from tensor2tensor.utils import hparam
 from tensor2tensor.utils import registry
-from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -28,7 +28,7 @@
 @registry.register_hparams("basic_1")
 def basic_params1():
   """A set of basic hyperparameters."""
-  return HParams(
+  return hparam.HParams(
       # If the problem consists of variable-length sequences
       # (see problem.batch_size_means_tokens()), then this is the number
       # of tokens per batch per GPU or per TPU core.  Otherwise, this is
@@ -339,7 +339,7 @@ def basic_params1():
       # Load weights from a second model. For instance, when using
       # pre-trained weights, you might want to initialize the encoder
       # and decoder by loading different models.
-      warm_start_from_second=""
+      warm_start_from_second="",
   )
 
 
diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index fae6806fa..26f60c68f 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -22,7 +22,7 @@
 from absl.testing import parameterized
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_image_attention
-from tensor2tensor.utils.hparam import HParams
+from tensor2tensor.utils import hparam
 
 import tensorflow as tf
 
@@ -37,7 +37,7 @@ def testPostProcessImageTrainMode(self, likelihood, num_mixtures, depth):
     batch = 1
     rows = 8
     cols = 24
-    hparams = HParams(
+    hparams = hparam.HParams(
         hidden_size=2,
         likelihood=likelihood,
         mode=tf.estimator.ModeKeys.TRAIN,
@@ -59,7 +59,7 @@ def testPostProcessImageInferMode(self, likelihood, num_mixtures, depth):
     cols = 24
     block_length = 4
     block_width = 2
-    hparams = HParams(
+    hparams = hparam.HParams(
         block_raster_scan=True,
         hidden_size=2,
         likelihood=likelihood,
@@ -91,7 +91,7 @@ def testCreateOutputTrainMode(self, likelihood, num_mixtures, depth):
       cols = channels * width
     else:
       cols = width
-    hparams = HParams(
+    hparams = hparam.HParams(
         hidden_size=2,
         likelihood=likelihood,
         num_channels=channels,
diff --git a/tensor2tensor/layers/message_passing_attention.py b/tensor2tensor/layers/message_passing_attention.py
index bf00d6001..a87304f73 100644
--- a/tensor2tensor/layers/message_passing_attention.py
+++ b/tensor2tensor/layers/message_passing_attention.py
@@ -866,7 +866,7 @@ def precompute_edge_matrices(adjacency, hparams):
   (we don't want to add to the graph everytime _fprop is called)
   Args:
     adjacency: placeholder of real valued vectors of shape [B, L, L, E]
-    hparams: tf.HParams object
+    hparams: HParams object
   Returns:
     edge_matrices: [batch, L * D, L * D] the dense matrix for message passing
     viewed as a block matrix (L,L) blocks of size (D,D). Each plot is a function
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 21409ad73..eaf56b120 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -130,7 +130,7 @@ def audio_bottom(x, model_hparams, vocab_size):
 
   Args:
     x: A Tensor with shape [batch, ...]
-    model_hparams: tf.HParams, model hyperparmeters.
+    model_hparams: HParams, model hyperparmeters.
     vocab_size: int, vocabulary size.
 
   Returns:
@@ -178,7 +178,7 @@ def audio_spectral_bottom(x, model_hparams, vocab_size):
 
   Args:
     x: A Tensor with shape [batch, ...]
-    model_hparams: tf.HParams, model hyperparmeters.
+    model_hparams: HParams, model hyperparmeters.
     vocab_size: int, vocabulary size.
 
   Returns:
@@ -298,7 +298,7 @@ def _image_channel_compress_bottom(inputs, model_hparams, name="bottom"):
   Args:
     inputs: Tensor representing RGB pixel intensities as integers, of shape
       [batch, img_len, img_len, channels].
-    model_hparams: tf.HParams, model hyperparmeters.
+    model_hparams: HParams, model hyperparmeters.
     name: string, scope.
 
   Returns:
@@ -375,7 +375,7 @@ def speech_recognition_bottom(x, model_hparams, vocab_size):
 
   Args:
     x: float32 tensor with shape [batch_size, len, 1, freqs * channels]
-    model_hparams: tf.HParams, model hyperparmeters.
+    model_hparams: HParams, model hyperparmeters.
     vocab_size: int, vocabulary size.
 
   Returns:
@@ -456,7 +456,7 @@ def get_weights(model_hparams, vocab_size, hidden_dim=None):
   """Create or get concatenated embedding or softmax variable.
 
   Args:
-    model_hparams: tf.HParams, model hyperparmeters.
+    model_hparams: HParams, model hyperparmeters.
     vocab_size: int, vocabulary size.
     hidden_dim: dim of the variable. Defaults to _model_hparams' hidden_size
 
@@ -721,7 +721,7 @@ def one_hot_class_label_loss(top_out,
   Args:
     top_out: logits Tensor with shape [batch, ?, ?, num_classes]
     targets: one-hot encoding Tensor with shape [batch, ?, ?, num_classes]
-    model_hparams: tf.HParams, model hyperparmeters.
+    model_hparams: HParams, model hyperparmeters.
     vocab_size: int, vocabulary size.
     weights_fn:
 
@@ -932,7 +932,7 @@ def class_label_top(body_output, targets, model_hparams, vocab_size):
   Args:
     body_output: A Tensor with shape [batch, ?, ?, body_output_size].
     targets:
-    model_hparams: tf.HParams, model hyperparmeters.
+    model_hparams: HParams, model hyperparmeters.
     vocab_size: int, vocabulary size.
 
   Returns:
@@ -978,7 +978,7 @@ def image_channel_compress_top(body_output, targets, model_hparams, vocab_size):
   Args:
     body_output: Tensor of shape [batch, img_len, img_len, depth].
     targets:
-    model_hparams: tf.HParams, model hyperparmeters.
+    model_hparams: HParams, model hyperparmeters.
     vocab_size: int, vocabulary size.
 
   Returns:
@@ -1044,7 +1044,7 @@ def sigmoid_max_pooling_class_label_top(body_output,
   Args:
     body_output: A Tensor with shape [batch, timesteps, 1, body_output_size].
     targets:
-    model_hparams: tf.HParams, model hyperparmeters.
+    model_hparams: HParams, model hyperparmeters.
     vocab_size: int, vocabulary size.
 
   Returns:
@@ -1109,7 +1109,7 @@ def symbol_top(body_output, targets, model_hparams, vocab_size):
     body_output: A Tensor with shape
       [batch, p0, p1, model_hparams.hidden_size].
     targets: Unused.
-    model_hparams: tf.HParams, model hyperparmeters.
+    model_hparams: HParams, model hyperparmeters.
     vocab_size: int, vocabulary size.
 
   Returns:
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 318fb05a2..888055040 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -1199,7 +1199,7 @@ def revnet(name, x, hparams, reverse=True):
   Args:
     name: variable scope for the revnet block.
     x: 4-D Tensor, shape=(NHWC).
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     reverse: bool, forward or backward pass.
   Returns:
     x: 4-D Tensor, shape=(NHWC).
@@ -1306,7 +1306,7 @@ def encoder_decoder(name, x, hparams, eps=None, reverse=False,
   Args:
     name: variable scope.
     x: 4-D Tensor, shape=(NHWC).
-    hparams: tf.contrib.training.HParams.
+    hparams: HParams.
     eps: Stores (glow(x) - mu) / sigma during the forward pass.
          Used only to test if the network is reversible.
     reverse: Forward or reverse pass.
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index ef01ff9b7..46914ec56 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -25,7 +25,7 @@
 import numpy as np
 from tensor2tensor.models.research import glow
 from tensor2tensor.models.research import glow_ops
-from tensor2tensor.utils.hparam import HParams
+from tensor2tensor.utils import hparam
 import tensorflow as tf
 
 arg_scope = tf.contrib.framework.arg_scope
@@ -157,7 +157,7 @@ def test_conv_stack(self, activation="relu"):
   def check_latent_to_dist(self, architecture):
     with tf.Graph().as_default():
       x = tf.random_uniform(shape=(16, 5, 5, 32))
-      hparams = HParams(architecture=architecture)
+      hparams = hparam.HParams(architecture=architecture)
       x_prior = glow_ops.latent_to_dist("split_prior", x, hparams=hparams,
                                         output_channels=64)
       mean_t, scale_t = x_prior.loc, x_prior.scale
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index da9fcedef..350410a54 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -33,10 +33,10 @@
 from tensor2tensor.rl.envs.py_func_batch_env import PyFuncBatchEnv
 from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
 from tensor2tensor.rl.envs.simulated_batch_gym_env import SimulatedBatchGymEnv
+from tensor2tensor.utils import hparam
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 from tensor2tensor.utils import trainer_lib
-from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 import tensorflow_probability as tfp
@@ -79,7 +79,7 @@ def ppo_base_v1():
 @registry.register_hparams
 def basic_policy_parameters():
   wrappers = None
-  return HParams(wrappers=wrappers)
+  return hparam.HParams(wrappers=wrappers)
 
 
 @registry.register_hparams
@@ -345,7 +345,7 @@ def ppo_pong_ae_base():
 def dqn_atari_base():
   # These params are based on agents/dqn/configs/dqn.gin
   # with some modifications taking into account our code
-  return HParams(
+  return hparam.HParams(
       agent_gamma=0.99,
       agent_update_horizon=1,
       agent_min_replay_history=20000,  # agent steps
@@ -399,7 +399,7 @@ def rlmf_tiny_overrides():
 
 @registry.register_hparams
 def rlmf_original():
-  return HParams(
+  return hparam.HParams(
       game="pong",
       sticky_actions=False,
       base_algo="ppo",
diff --git a/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index 19ace57af..ee0c3bc72 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -22,9 +22,9 @@
 
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import hparam
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
-from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -792,7 +792,7 @@ def resnet_200():
 # Pruning parameters
 @registry.register_pruning_params
 def resnet_weight():
-  hp = HParams()
+  hp = hparam.HParams()
   hp.add_hparam("strategy", "weight")
   hp.add_hparam("black_list", ["logits", "bias"])
   hp.add_hparam("white_list", ["td_conv"])
@@ -810,7 +810,7 @@ def resnet_unit():
 # Adversarial attack parameters
 @registry.register_attack_params
 def resnet_fgsm():
-  aparams = HParams()
+  aparams = hparam.HParams()
   aparams.attack = "fgsm"
   aparams.epsilon_name = "eps"
   aparams.attack_epsilons = [i * 0.8 for i in range(20)]
diff --git a/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index 51b52af2e..3594a9f56 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -20,9 +20,9 @@
 from __future__ import print_function
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import hparam
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
-from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -215,7 +215,7 @@ def shakeshake_tpu():
 
 @registry.register_attack_params
 def shake_shake_fgsm():
-  aparams = HParams()
+  aparams = hparam.HParams()
   aparams.attack = "fgsm"
   aparams.attack_epsilons = [(i+1) * 0.1 for i in range(12)]
   aparams.add_hparam("clip_min", 0.0)
diff --git a/tensor2tensor/models/video/nfg_interpolate.py b/tensor2tensor/models/video/nfg_interpolate.py
index e7284dde7..db665853c 100644
--- a/tensor2tensor/models/video/nfg_interpolate.py
+++ b/tensor2tensor/models/video/nfg_interpolate.py
@@ -111,8 +111,8 @@ def interpolate(features, hparams, decode_hp):
 
   Args:
     features: dict of tensors
-    hparams: tf.contrib.training.HParams, training hparams.
-    decode_hp: tf.contrib.training.HParams, decode hparams.
+    hparams: HParams, training hparams.
+    decode_hp: HParams, decode hparams.
   Returns:
     images: interpolated images, 4-D Tensor, shape=(num_interp, H, W, C)
     first_frame: image, 3-D Tensor, shape=(1, H, W, C)
@@ -178,8 +178,8 @@ def interpolations_to_summary(sample_ind, interpolations, first_frame,
     interpolations: Numpy array, shape=(num_interp, H, W, 3)
     first_frame: Numpy array, shape=(HWC)
     last_frame: Numpy array, shape=(HWC)
-    hparams: tf.contrib.training.HParams, train hparams
-    decode_hp: tf.contrib.training.HParams, decode hparams
+    hparams: HParams, train hparams
+    decode_hp: HParams, decode hparams
   Returns:
     summaries: list of tf Summary Values.
   """
diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index f86a99fc2..2928665f3 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -37,9 +37,9 @@
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl import trainer_model_based_params  # pylint: disable=unused-import
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
+from tensor2tensor.utils import hparam
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
-from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -106,7 +106,7 @@
 
 @registry.register_hparams
 def planner_tiny():
-  return HParams(
+  return hparam.HParams(
       num_rollouts=1,
       planning_horizon=2,
       rollout_agent_type="random",
@@ -119,7 +119,7 @@ def planner_tiny():
 
 @registry.register_hparams
 def planner_small():
-  return HParams(
+  return hparam.HParams(
       num_rollouts=64,
       planning_horizon=16,
       rollout_agent_type="policy",
@@ -132,7 +132,7 @@ def planner_small():
 
 @registry.register_hparams
 def planner_base():
-  return HParams(
+  return hparam.HParams(
       num_rollouts=96,
       batch_size=96,
       planning_horizon=8,
diff --git a/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index 59317f80f..d53066091 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -31,8 +31,8 @@
 from tensor2tensor.models.research.rl import make_simulated_env_fn_from_hparams
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
+from tensor2tensor.utils import hparam
 from tensor2tensor.utils import trainer_lib
-from tensor2tensor.utils.hparam import HParams
 from tensor2tensor.utils.misc_utils import camelcase_to_snakecase
 
 import tensorflow as tf
@@ -283,7 +283,7 @@ def create_simulated_env(
     if key not in other_hparams:
       other_hparams[key] = a_bit_risky_defaults[key]
 
-  hparams = HParams(
+  hparams = hparam.HParams(
       grayscale=grayscale,
       resize_width_factor=resize_width_factor,
       resize_height_factor=resize_height_factor,
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index f7c56911e..2776084f3 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -22,8 +22,8 @@
 
 
 from tensor2tensor.data_generators import gym_env
+from tensor2tensor.utils import hparam
 from tensor2tensor.utils import registry
-from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
@@ -45,7 +45,7 @@
 
 
 def _rlmb_base():
-  return HParams(
+  return hparam.HParams(
       epochs=15,
       # Total frames used for training. This will be distributed evenly across
       # hparams.epochs.
@@ -852,7 +852,7 @@ def merge_unscoped_hparams(scopes_and_hparams):
       scoped_key = "%s.%s" % (scope, key)
       merged_values[scoped_key] = value
 
-  return HParams(**merged_values)
+  return hparam.HParams(**merged_values)
 
 
 def split_scoped_hparams(scopes, merged_hparams):
@@ -865,7 +865,7 @@ def split_scoped_hparams(scopes, merged_hparams):
     split_values[scope][key] = value
 
   return [
-      HParams(**split_values[scope]) for scope in scopes
+      hparam.HParams(**split_values[scope]) for scope in scopes
   ]
 
 
@@ -919,7 +919,7 @@ def dynamic_register_hparams(name, hparams):
 
   @registry.register_hparams(name)
   def new_hparams_set():
-    return HParams(**hparams.values())
+    return hparam.HParams(**hparams.values())
 
   return new_hparams_set
 
diff --git a/tensor2tensor/serving/query.py b/tensor2tensor/serving/query.py
index 75fb896b7..424cd5581 100644
--- a/tensor2tensor/serving/query.py
+++ b/tensor2tensor/serving/query.py
@@ -25,9 +25,9 @@
 
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
 from tensor2tensor.serving import serving_utils
+from tensor2tensor.utils import hparam
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import usr_dir
-from tensor2tensor.utils.hparam import HParams
 import tensorflow as tf
 
 flags = tf.flags
@@ -81,7 +81,7 @@ def main(_):
   validate_flags()
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
   problem = registry.problem(FLAGS.problem)
-  hparams = HParams(
+  hparams = hparam.HParams(
       data_dir=os.path.expanduser(FLAGS.data_dir))
   problem.get_hparams(hparams)
   request_fn = make_request_fn()
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 78fec4b79..656edfdfd 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -33,9 +33,9 @@
 from tensor2tensor.data_generators import problem as problem_lib
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import hparam
 from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import registry
-from tensor2tensor.utils.hparam import HParams
 import tensorflow as tf
 
 FLAGS = tf.flags.FLAGS
@@ -46,7 +46,7 @@
 
 def decode_hparams(overrides=""):
   """Hyperparameters for decoding."""
-  hp = HParams(
+  hp = hparam.HParams(
       save_images=False,
       log_results=True,
       extra_length=100,
@@ -257,9 +257,9 @@ def decode_once(estimator,
     estimator: tf.estimator.Estimator instance. Used to generate encoded
       predictions.
     problem_name: str. Name of problem.
-    hparams: tf.HParams instance. HParams for model training.
+    hparams: HParams instance. HParams for model training.
     infer_input_fn: zero-arg function. Input function for estimator.
-    decode_hp: tf.HParams instance. See decode_hparams() above.
+    decode_hp: HParams instance. See decode_hparams() above.
     decode_to_file: str. Prefix for filenames. Used to generated filenames to
       which decoded predictions are written.
     output_dir: str. Output directory. Only used for writing images.
diff --git a/tensor2tensor/utils/diet.py b/tensor2tensor/utils/diet.py
index 73d4ffeda..fa06b9902 100644
--- a/tensor2tensor/utils/diet.py
+++ b/tensor2tensor/utils/diet.py
@@ -27,7 +27,7 @@
 import math
 
 from tensor2tensor.layers import common_layers
-from tensor2tensor.utils.hparam import HParams
+from tensor2tensor.utils import hparam
 import tensorflow as tf
 
 
@@ -37,7 +37,7 @@ def diet_adam_optimizer_params():
   Returns:
     a hyperparameters object.
   """
-  return HParams(
+  return hparam.HParams(
       quantize=True,  # use 16-bit fixed-point
       quantization_scale=10.0 / tf.int16.max,
       optimizer="DietAdam",
diff --git a/tensor2tensor/utils/hparams_lib.py b/tensor2tensor/utils/hparams_lib.py
index 3db112e96..ff7f8756e 100644
--- a/tensor2tensor/utils/hparams_lib.py
+++ b/tensor2tensor/utils/hparams_lib.py
@@ -22,15 +22,15 @@
 import json
 
 from tensor2tensor.data_generators import problem as problem_lib
+from tensor2tensor.utils import hparam
 from tensor2tensor.utils import registry
-from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 
 
 def copy_hparams(hparams):
   hp_vals = hparams.values()
-  new_hparams = HParams(**hp_vals)
+  new_hparams = hparam.HParams(**hp_vals)
   other_attrs = ["problem", "problem_hparams"]
   for attr in other_attrs:
     attr_val = getattr(hparams, attr, None)
@@ -72,7 +72,7 @@ def create_hparams_from_json(json_path, hparams=None):
     hparams_values.pop("name", None)
     hparams_values.pop("top", None)
     hparams_values.pop("weights_fn", None)
-    new_hparams = HParams(**hparams_values)
+    new_hparams = hparam.HParams(**hparams_values)
     # Some keys are in new_hparams but not hparams, so we need to be more
     #   careful than simply using parse_json() from HParams
     if hparams:  # hparams specified, so update values from json
diff --git a/tensor2tensor/utils/misc_utils_test.py b/tensor2tensor/utils/misc_utils_test.py
index 5d5ae805d..6beedb098 100644
--- a/tensor2tensor/utils/misc_utils_test.py
+++ b/tensor2tensor/utils/misc_utils_test.py
@@ -19,8 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensor2tensor.utils import hparam
 from tensor2tensor.utils import misc_utils
-from tensor2tensor.utils.hparam import HParams
 import tensorflow as tf
 
 
@@ -57,7 +57,7 @@ def test_snakecase_to_camelcase(self):
                      misc_utils.snakecase_to_camelcase("lstm_seq2_seq"))
 
   def test_pprint_hparams(self):
-    hparams = HParams(
+    hparams = hparam.HParams(
         int_=1, str_="str", bool_=True, float_=1.1, list_int=[1, 2], none=None)
 
     # pylint: disable=g-inconsistent-quotes
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 4bcd9e8eb..41c2252d9 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1889,13 +1889,8 @@ def sampled_results(features, logits, mixin_prob):
 
     # Apply scheduled sampling over N passes. The logits from the (n-1)-th pass
     # will be mixed with gold tokens for conditioning in the n-th pass.
-    if hasattr(hparams, "scheduled_sampling_num_passes"):
-      scheduled_sampling_num_passes = hparams.scheduled_sampling_num_passes
-    else:
-      # TODO(duckworthd): Delete once scheduled_sampling_num_passes is added to
-      # common_hparams.py.
-      scheduled_sampling_num_passes = 1
-
+    scheduled_sampling_num_passes = getattr(
+        hparams, "scheduled_sampling_num_passes", 1)
     assert scheduled_sampling_num_passes > 0, (
         "hparams.scheduled_sampling_num_passes must be > 0 if "
         "hparams.scheduled_sampling_prob > 0.0")
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index 229846832..b914d6607 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -20,9 +20,9 @@
 from __future__ import print_function
 
 from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.utils import hparam
 from tensor2tensor.utils import t2t_model
 from tensor2tensor.utils import test_utils
-from tensor2tensor.utils.hparam import HParams
 
 import tensorflow as tf
 tf.compat.v1.enable_eager_execution()
@@ -33,7 +33,7 @@ class T2TModelTest(tf.test.TestCase):
   @test_utils.run_in_graph_and_eager_modes()
   def testSummarizeLosses(self):
     with tf.Graph().as_default():
-      model = t2t_model.T2TModel(HParams())
+      model = t2t_model.T2TModel(hparam.HParams())
       losses = {"training": tf.random_normal([]),
                 "extra": tf.random_normal([])}
       outputs = model._summarize_losses(losses)
@@ -50,7 +50,7 @@ def testLossSingleWeights(self):
         sequence_size = 16
         vocab_size = 3
 
-        model_hparams = HParams(
+        model_hparams = hparam.HParams(
             prepend_mode="none",
             loss={},
             weights_fn={},

From 7259104114ebe4b493e9c082942bfffd5a8db55f Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 26 Apr 2019 16:35:30 -0700
Subject: [PATCH 1962/2720] Make default Dense initializer be glorot-uniform
 instead of glorot-normal (as in Keras); rename Xavier to Glorot for
 consistency (it was "Xavier Glorot", let's use the surname).

PiperOrigin-RevId: 245508572
---
 tensor2tensor/trax/layers/attention.py   | 24 ++++++++----------------
 tensor2tensor/trax/layers/core.py        |  8 ++++----
 tensor2tensor/trax/models/transformer.py | 15 +++++----------
 tensor2tensor/trax/trax.py               |  2 ++
 4 files changed, 19 insertions(+), 30 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index f5440fb64..b2fe70814 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -244,20 +244,16 @@ def MultiHeadedAttentionQKV(
   return combinators.Serial(
       combinators.Parallel(
           combinators.Parallel(
-              core.Dense(feature_depth,
-                         kernel_initializer=core.XavierUniformInitializer()),
-              core.Dense(feature_depth,
-                         kernel_initializer=core.XavierUniformInitializer()),
-              core.Dense(feature_depth,
-                         kernel_initializer=core.XavierUniformInitializer()),
+              core.Dense(feature_depth),
+              core.Dense(feature_depth),
+              core.Dense(feature_depth),
           ),
           combinators.Identity()
       ),
       PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
           feature_depth=feature_depth, num_heads=num_heads,
           dropout=dropout, mode=mode),
-      core.Dense(feature_depth,
-                 kernel_initializer=core.XavierUniformInitializer()),
+      core.Dense(feature_depth),
   )
 
 
@@ -376,12 +372,9 @@ def ChunkedCausalMultiHeadedAttention(
       ),
       combinators.Parallel(
           combinators.Parallel(
-              core.Dense(feature_depth,
-                         kernel_initializer=core.XavierUniformInitializer()),
-              core.Dense(feature_depth,
-                         kernel_initializer=core.XavierUniformInitializer()),
-              core.Dense(feature_depth,
-                         kernel_initializer=core.XavierUniformInitializer()),
+              core.Dense(feature_depth),
+              core.Dense(feature_depth),
+              core.Dense(feature_depth),
           ),
           combinators.Identity()
       )
@@ -392,8 +385,7 @@ def ChunkedCausalMultiHeadedAttention(
       combinators.Map(PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
           feature_depth=feature_depth, num_heads=num_heads,
           dropout=dropout, mode=mode), check_shapes=False),
-      combinators.Map(core.Dense(
-          feature_depth, kernel_initializer=core.XavierUniformInitializer())),
+      combinators.Map(core.Dense(feature_depth))
   )
 
 
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 6cc541188..31999fdce 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -56,8 +56,8 @@ def init(shape, rng):
   return init
 
 
-def XavierUniformInitializer(out_dim=0, in_dim=1):
-  """An initializer function for random uniform xavier-scaled coefficients."""
+def GlorotUniformInitializer(out_dim=0, in_dim=1):
+  """An initializer function for random uniform Glorot-scaled coefficients."""
   def init(shape, rng):
     fan_in, fan_out = shape[in_dim], shape[out_dim]
     std = np.sqrt(2.0 / (fan_in + fan_out))
@@ -112,7 +112,7 @@ class Dense(base.Layer):
   """Layer constructor function for a dense (fully-connected) layer."""
 
   def __init__(self, units,
-               kernel_initializer=GlorotNormalInitializer(),
+               kernel_initializer=GlorotUniformInitializer(),
                bias_initializer=RandomNormalInitializer(1e-6)):
     super(Dense, self).__init__()
     self._units = units
@@ -137,7 +137,7 @@ class Embedding(base.Layer):
   """Layer constructor function for an embedding layer."""
 
   def __init__(self, feature_depth, vocab_size,
-               kernel_initializer=XavierUniformInitializer()):
+               kernel_initializer=GlorotUniformInitializer()):
     super(Embedding, self).__init__()
     self._feature_depth = feature_depth
     self._vocab_size = vocab_size
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 2ada531d2..a2b082a1c 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -28,12 +28,10 @@ def ResidualFeedForward(feature_depth,
   """Residual feed-forward layer with normalization at start."""
   return layers.Residual(
       layers.LayerNorm(),
-      layers.Dense(feedforward_depth,
-                   kernel_initializer=layers.XavierUniformInitializer()),
+      layers.Dense(feedforward_depth),
       layers.Relu(),
       layers.Dropout(rate=dropout, mode=mode),
-      layers.Dense(feature_depth,
-                   kernel_initializer=layers.XavierUniformInitializer()),
+      layers.Dense(feature_depth),
       layers.Dropout(rate=dropout, mode=mode)
   )
 
@@ -185,8 +183,7 @@ def TransformerLM(vocab_size,
                                    dropout, mode)
                       for _ in range(num_layers)]),
       layers.LayerNorm(),
-      layers.Dense(vocab_size,
-                   kernel_initializer=layers.XavierUniformInitializer()),
+      layers.Dense(vocab_size),
       layers.LogSoftmax()
   )
 
@@ -271,8 +268,7 @@ def ChunkedTransformerLM(vocab_size,
       layers.PositionalEncoding(max_len=max_len),
       layers.Serial(*stack),
       layers.Map(layers.LayerNorm()),
-      layers.Map(layers.Dense(
-          vocab_size, kernel_initializer=layers.XavierUniformInitializer())),
+      layers.Map(layers.Dense(vocab_size)),
       layers.Map(layers.LogSoftmax()),
   )
 
@@ -420,8 +416,7 @@ def transformer(source, target, source_mask, target_mask, memory_mask):  # pylin
   def Generator(encoded_target):
     return layers.Serial(
         encoded_target,
-        layers.Dense(target_vocab_size,
-                     kernel_initializer=layers.XavierUniformInitializer()),
+        layers.Dense(target_vocab_size),
         layers.LogSoftmax
     )
 
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 0c7295981..9f1a3eb37 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -25,6 +25,7 @@
 import os
 import pickle
 import random
+import sys
 import time
 
 from absl import logging
@@ -115,6 +116,7 @@ def log(s, stdout=True):
   logging.info(s)
   if stdout:
     print(s)
+    sys.stdout.flush()
 
 
 def step_log(step, s):

From 5ab32cd61430aef2afca666bf41fb5fedeec402f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 26 Apr 2019 20:39:03 -0700
Subject: [PATCH 1963/2720] Added area attention to Transformer and LSTM.

PiperOrigin-RevId: 245530442
---
 tensor2tensor/layers/area_attention.py     | 73 +++++++++---------
 tensor2tensor/layers/common_attention.py   | 43 +++++++++--
 tensor2tensor/layers/common_hparams.py     |  8 ++
 tensor2tensor/layers/transformer_layers.py | 17 ++++-
 tensor2tensor/models/lstm.py               | 86 +++++++++++++++++++++-
 tensor2tensor/models/transformer.py        | 28 ++++++-
 6 files changed, 206 insertions(+), 49 deletions(-)

diff --git a/tensor2tensor/layers/area_attention.py b/tensor2tensor/layers/area_attention.py
index 8372718f2..2635608a4 100644
--- a/tensor2tensor/layers/area_attention.py
+++ b/tensor2tensor/layers/area_attention.py
@@ -44,23 +44,24 @@ def lengths_to_area_mask(feature_length, length, max_area_size):
   return mask
 
 
-def _max_pool_one_shape(features_2d, area_width, area_height, batch_size,
-                        width, height, depth, name=None):
-  """Computes area max for features_2d.
+def _pool_one_shape(features_2d, area_width, area_height, batch_size,
+                    width, height, depth, fn=tf.reduce_max, name=None):
+  """Pools for an area in features_2d.
 
   Args:
-    features_2d: a Tensor in a shape of [batch_size, height * width, depth].
+    features_2d: a Tensor in a shape of [batch_size, height, width, depth].
     area_width: the max width allowed for an area.
     area_height: the max height allowed for an area.
     batch_size: the batch size.
     width: the width of the memory.
     height: the height of the memory.
     depth: the depth of the features.
+    fn: the TF function for the pooling.
     name: the op name.
   Returns:
-    max_tensor: A Tensor of shape [batch_size, num_areas, depth]
+    pool_tensor: A Tensor of shape [batch_size, num_areas, depth]
   """
-  with tf.name_scope(name, default_name="max_pool_one_shape"):
+  with tf.name_scope(name, default_name="pool_one_shape"):
     images = []
     for y_shift in range(area_height):
       image_height = tf.maximum(height - area_height + 1 + y_shift, 0)
@@ -70,26 +71,27 @@ def _max_pool_one_shape(features_2d, area_width, area_height, batch_size,
         flatten_area = tf.reshape(area, [batch_size, -1, depth, 1])
         images.append(flatten_area)
     image_tensor = tf.concat(images, axis=3)
-    max_tensor = tf.reduce_max(image_tensor, axis=3)
+    max_tensor = fn(image_tensor, axis=3)
   return max_tensor
 
 
-def max_pool(features, max_area_width, max_area_height=1, height=1,
-             name=None):
-  """Computes area max for features.
+def basic_pool(features, max_area_width, max_area_height=1, height=1,
+               fn=tf.reduce_max, name=None):
+  """Pools for each area based on a given pooling function (fn).
 
   Args:
     features: a Tensor in a shape of [batch_size, height * width, depth].
     max_area_width: the max width allowed for an area.
     max_area_height: the max height allowed for an area.
     height: the height of the image.
+    fn: the TF function for the pooling.
     name: the namescope.
   Returns:
-    max_results: A Tensor of shape [batch_size, num_areas, depth]
+    pool_results: A Tensor of shape [batch_size, num_areas, depth]
     area_heights: A Tensor of shape [batch_size, num_areas, 1]
     area_widths: A Tensor of shape [batch_size, num_areas, 1]
   """
-  with tf.name_scope(name, default_name="max_pool"):
+  with tf.name_scope(name, default_name="basic_pool"):
     feature_shape = common_layers.shape_list(features)
     batch_size = feature_shape[0]
     length = feature_shape[-2]
@@ -98,19 +100,20 @@ def max_pool(features, max_area_width, max_area_height=1, height=1,
     features_2d = tf.reshape(features, [batch_size, height, width, depth])
     height_list = []
     width_list = []
-    max_list = []
+    pool_list = []
     size_tensor = tf.ones_like(features_2d[:, :, :, 0], dtype=tf.int32)
     for area_height in range(max_area_height):
       for area_width in range(max_area_width):
-        max_tensor = _max_pool_one_shape(features_2d,
-                                         area_width=area_width + 1,
-                                         area_height=area_height + 1,
-                                         batch_size=batch_size,
-                                         width=width,
-                                         height=height,
-                                         depth=depth)
-        max_list.append(
-            tf.reshape(max_tensor, [batch_size, -1, depth]))
+        pool_tensor = _pool_one_shape(features_2d,
+                                      area_width=area_width + 1,
+                                      area_height=area_height + 1,
+                                      batch_size=batch_size,
+                                      width=width,
+                                      height=height,
+                                      depth=depth,
+                                      fn=fn)
+        pool_list.append(
+            tf.reshape(pool_tensor, [batch_size, -1, depth]))
         height_list.append(
             tf.reshape(
                 size_tensor[:, area_height:, area_width:] *\
@@ -119,10 +122,10 @@ def max_pool(features, max_area_width, max_area_height=1, height=1,
             tf.reshape(
                 size_tensor[:, area_height:, area_width:] *\
                 (area_width + 1), [batch_size, -1]))
-    max_results = tf.concat(max_list, axis=1)
+    pool_results = tf.concat(pool_list, axis=1)
     area_heights = tf.expand_dims(tf.concat(height_list, axis=1), 2)
     area_widths = tf.expand_dims(tf.concat(width_list, axis=1), 2)
-  return max_results, area_heights, area_widths
+  return pool_results, area_heights, area_widths
 
 
 def _compute_sum_image(features, max_area_width, max_area_height=1, height=1,
@@ -253,8 +256,8 @@ def compute_area_key(features, max_area_width, max_area_height=1, height=1,
   if mode == "mean":
     return area_mean
   elif mode == "max":
-    area_max, _, _ = max_pool(features, max_area_width=max_area_width,
-                              max_area_height=max_area_height, height=height)
+    area_max, _, _ = basic_pool(features, max_area_width=max_area_width,
+                                max_area_height=max_area_height, height=height)
     return area_max
   elif mode == "sample":
     if training:
@@ -276,8 +279,9 @@ def compute_area_key(features, max_area_width, max_area_height=1, height=1,
     if mode == "concat":
       feature_concat = tf.concat([area_mean, area_std, size_embed], -1)
     elif mode == "max_concat":
-      area_max, _, _ = max_pool(features, max_area_width=max_area_width,
-                                max_area_height=max_area_height, height=height)
+      area_max, _, _ = basic_pool(features, max_area_width=max_area_width,
+                                  max_area_height=max_area_height,
+                                  height=height)
       feature_concat = tf.concat([area_max, size_embed], -1)
     elif mode == "sum":
       feature_concat = size_embed + area_mean + area_std
@@ -352,10 +356,10 @@ def dot_product_area_attention(q,
   tf.logging.info("dot_product_area_attention: "
                   "area_h=%d, area_w=%d, mem_h=%d, "
                   "area_key_mode=%s, area_value_mode=%s, "
-                  "area_temperature=%f, top_k_areas=%d",
+                  "area_temperature=%f",
                   max_area_height, max_area_width, memory_height,
                   area_key_mode, area_value_mode,
-                  area_temperature, top_k_areas)
+                  area_temperature)
   with tf.variable_scope(
       name, default_name="dot_product_area_attention",
       values=[q, k, v]) as scope:
@@ -376,10 +380,11 @@ def dot_product_area_attention(q,
           tf.reshape(v, [-1, length, depth]), max_area_width=max_area_width,
           max_area_height=max_area_height, height=memory_height)
     elif area_value_mode == "max":
-      v_area, _, _ = max_pool(tf.reshape(v, [-1, length, depth]),
-                              max_area_width=max_area_width,
-                              max_area_height=max_area_height,
-                              height=memory_height)
+      v_area, _, _ = basic_pool(tf.reshape(v, [-1, length, depth]),
+                                max_area_width=max_area_width,
+                                max_area_height=max_area_height,
+                                height=memory_height,
+                                fn=tf.reduce_max)
     elif area_value_mode == "sum":
       _, _, v_area, _, _ = compute_area_features(
           tf.reshape(v, [-1, length, depth]), max_area_width=max_area_width,
diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 2c4067056..d49d8a6d0 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -28,6 +28,7 @@
 from six.moves import range  # pylint: disable=redefined-builtin
 from six.moves import zip  # pylint: disable=redefined-builtin
 
+from tensor2tensor.layers import area_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import expert_utils
 
@@ -3991,6 +3992,12 @@ def multihead_attention(query_antecedent,
                         recurrent_memory=None,
                         chunk_number=None,
                         hard_attention_k=0,
+                        max_area_width=1,
+                        max_area_height=1,
+                        memory_height=1,
+                        area_key_mode="mean",
+                        area_value_mode="sum",
+                        training=True,
                         **kwargs):
   """Multihead scaled-dot-product attention with input/output transformations.
 
@@ -4049,6 +4056,14 @@ def multihead_attention(query_antecedent,
     chunk_number: an optional integer Tensor with shape [batch] used to operate
       the recurrent_memory.
     hard_attention_k: integer, if > 0 triggers hard attention (picking top-k).
+    max_area_width: the max width allowed for an area.
+    max_area_height: the max height allowed for an area.
+    memory_height: the height of the memory.
+    area_key_mode: the mode for computing area keys, which can be "mean",
+      "concat", "sum", "sample_concat", and "sample_sum".
+    area_value_mode: the mode for computing area values, which can be either
+      "mean", or "sum".
+    training: indicating if it is in the training mode.
     **kwargs (dict): Parameters for the attention function.
 
   Caching:
@@ -4171,13 +4186,27 @@ def multihead_attention(query_antecedent,
       if isinstance(x, tuple):
         x, additional_returned_value = x  # Unpack
     elif attention_type == "dot_product":
-      x = dot_product_attention(
-          q, k, v, bias, dropout_rate, image_shapes,
-          save_weights_to=save_weights_to,
-          make_image_summary=make_image_summary,
-          dropout_broadcast_dims=dropout_broadcast_dims,
-          activation_dtype=kwargs.get("activation_dtype"),
-          hard_attention_k=hard_attention_k)
+      tf.logging.info("max_area_width=%d, max_area_height=%d",
+                      max_area_width, max_area_height)
+      if max_area_width > 1 or max_area_height > 1:
+        x = area_attention.dot_product_area_attention(
+            q, k, v, bias, dropout_rate, image_shapes,
+            save_weights_to=save_weights_to,
+            dropout_broadcast_dims=dropout_broadcast_dims,
+            max_area_width=max_area_width,
+            max_area_height=max_area_height,
+            memory_height=memory_height,
+            area_key_mode=area_key_mode,
+            area_value_mode=area_value_mode,
+            training=training)
+      else:
+        x = dot_product_attention(q, k, v, bias, dropout_rate, image_shapes,
+                                  save_weights_to=save_weights_to,
+                                  make_image_summary=make_image_summary,
+                                  dropout_broadcast_dims=dropout_broadcast_dims,
+                                  activation_dtype=kwargs.get(
+                                      "activation_dtype"),
+                                  hard_attention_k=hard_attention_k)
     elif attention_type == "dot_product_relative":
       x = dot_product_attention_relative(
           q,
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index c25381d4e..5aba1af63 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -340,6 +340,14 @@ def basic_params1():
       # pre-trained weights, you might want to initialize the encoder
       # and decoder by loading different models.
       warm_start_from_second="",
+      # Area attention hyper parameters
+      area_value_mode="none",
+      area_key_mode="none",
+      # Using area attention for the number of layers from the bottom
+      num_area_layers=0,
+      max_area_width=1,
+      max_area_height=1,
+      memory_height=1
   )
 
 
diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 3a62021f6..9b23ecad5 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -183,6 +183,14 @@ def transformer_encoder(encoder_input,
     for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):
       with tf.variable_scope("layer_%d" % layer):
         with tf.variable_scope("self_attention"):
+          if layer < hparams.get("num_area_layers", 0):
+            max_area_width = hparams.get("max_area_width", 1)
+            max_area_height = hparams.get("max_area_height", 1)
+            memory_height = hparams.get("memory_height", 1)
+          else:
+            max_area_width = 1
+            max_area_height = 1
+            memory_height = 1
           y = common_attention.multihead_attention(
               common_layers.layer_preprocess(x, hparams),
               None,
@@ -204,7 +212,14 @@ def transformer_encoder(encoder_input,
               vars_3d=hparams.get("attention_variables_3d"),
               activation_dtype=hparams.get("activation_dtype", "float32"),
               weight_dtype=hparams.get("weight_dtype", "float32"),
-              hard_attention_k=hparams.get("hard_attention_k", 0))
+              hard_attention_k=hparams.get("hard_attention_k", 0),
+              max_area_width=max_area_width,
+              max_area_height=max_area_height,
+              memory_height=memory_height,
+              area_key_mode=hparams.get("area_key_mode", "none"),
+              area_value_mode=hparams.get("area_value_mode", "none"),
+              training=(hparams.get("mode", tf.estimator.ModeKeys.TRAIN)
+                        == tf.estimator.ModeKeys.TRAIN))
           x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(
diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 9d343cdb9..f0ef2aae6 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import copy
+from tensor2tensor.layers import area_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import registry
@@ -102,10 +103,46 @@ def lstm_attention_decoder(inputs, hparams, train, name, initial_state,
   else:
     raise ValueError("Unknown hparams.attention_mechanism = %s, must be "
                      "luong or bahdanau." % hparams.attention_mechanism)
-  attention_mechanism = attention_mechanism_class(
-      hparams.hidden_size, encoder_outputs,
-      memory_sequence_length=encoder_output_length)
-
+  if hparams.get("max_area_width", 1) > 1:
+    def _area_key_value_fn(keys, values):
+      """Custom fn for computing area keys and values."""
+      tf.logging.info("max_area_width=%d, area_key_mode=%s, area_value_mode=%s",
+                      hparams.get("max_area_width", 1),
+                      hparams.get("area_key_mode", "none"),
+                      hparams.get("area_value_mode", "none"))
+      keys = area_attention.compute_area_key(
+          keys, max_area_width=hparams.get("max_area_width", 1),
+          mode=hparams.get("area_key_mode", "none"), name="decoder_encoder",
+          training=(hparams.mode == tf.estimator.ModeKeys.TRAIN))
+      if hparams.get("area_value_mode", "none") == "sum":
+        _, _, values, _, _ = area_attention.compute_area_features(
+            values, max_area_width=hparams.get("max_area_width", 1))
+      elif hparams.get("area_value_mode", "none") == "mean":
+        values, _, _, _, _ = area_attention.compute_area_features(
+            values, max_area_width=hparams.get("max_area_width", 1))
+      else:
+        raise ValueError(
+            "Unsupported area_value_mode: %s" % hparams.get(
+                "area_value_mode", "none"))
+      return keys, values
+    area_mask = area_attention.lengths_to_area_mask(
+        feature_length=encoder_output_length,
+        length=common_layers.shape_list(encoder_outputs)[1],
+        max_area_size=hparams.get("max_area_width", "1"))
+    def _area_prob_fn(score):
+      alignments = tf.nn.softmax(score)
+      alignments = tf.where(area_mask, alignments, tf.zeros_like(alignments))
+      alignments = tf.div(alignments, tf.reduce_sum(
+          alignments, axis=-1, keepdims=True))
+      return alignments
+    attention_mechanism = attention_mechanism_class(
+        hparams.hidden_size, encoder_outputs,
+        memory_sequence_length=None,
+        probability_fn=_area_prob_fn,
+        custom_key_value_fn=_area_key_value_fn)
+  else:
+    attention_mechanism = attention_mechanism_class(hparams.hidden_size,
+                                                    encoder_outputs)
   cell = tf.contrib.seq2seq.AttentionWrapper(
       tf.nn.rnn_cell.MultiRNNCell(layers),
       [attention_mechanism]*hparams.num_heads,
@@ -339,6 +376,7 @@ def body(self, features):
     flat_target = tf.reshape(features["targets_raw"],
                              [target_shape[0], target_shape[1]])
     targets_length = tf.reduce_sum(tf.minimum(flat_target, 1), -1)
+    tf.logging.info(self._hparams)
     return lstm_seq2seq_internal_attention(
         features["inputs"], features["targets"], self._hparams, train,
         inputs_length, targets_length)
@@ -442,3 +480,43 @@ def lstm_asr_v1():
   hparams.min_length_bucket = hparams.max_input_seq_length // 2
   hparams.learning_rate = 0.05
   return hparams
+
+
+@registry.register_hparams
+def lstm_area_attention_base():
+  """Hparams for LSTM with area attention."""
+  hparams = lstm_luong_attention()
+  hparams.batch_size = 16384
+  hparams.num_hidden_layers = 2
+  hparams.hidden_size = 1024
+  hparams.num_heads = 4
+  hparams.dropout = 0.2
+  hparams.learning_rate = 0.1
+  hparams.max_area_width = 2
+  hparams.area_key_mode = "mean"
+  hparams.area_value_mode = "sum"
+  return hparams
+
+
+@registry.register_hparams
+def lstm_area_attention_enfr():
+  """Hparams for LSTM with area attention."""
+  hparams = lstm_area_attention_base()
+  hparams.dropout = 0.1
+  return hparams
+
+
+@registry.register_hparams
+def lstm_area_attention_char():
+  """Hparams for LSTM with area attention."""
+  hparams = lstm_area_attention_base()
+  hparams.batch_size = 20480
+  return hparams
+
+
+@registry.register_hparams
+def lstm_area_attention_char_enfr():
+  """Hparams for LSTM with area attention."""
+  hparams = lstm_area_attention_char()
+  hparams.dropout = 0.1
+  return hparams
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index e3e6317a0..e87ba9c0d 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1417,6 +1417,15 @@ def transformer_decoder(decoder_input,
         recurrent_memory = recurrent_memory_by_layer[layer_name]
       else:
         recurrent_memory = None
+
+      if layer < hparams.get("num_area_layers", 0):
+        max_area_width = hparams.get("max_area_width", 1)
+        max_area_height = hparams.get("max_area_height", 1)
+        memory_height = hparams.get("max_area_height", 1)
+      else:
+        max_area_width = 1
+        max_area_height = 1
+        memory_height = 1
       with tf.variable_scope(layer_name):
         with tf.variable_scope("self_attention"):
           y = common_attention.multihead_attention(
@@ -1446,8 +1455,14 @@ def transformer_decoder(decoder_input,
               layer_collection=layer_collection,
               recurrent_memory=recurrent_memory,
               chunk_number=chunk_number,
-              hard_attention_k=hparams.get("hard_attention_k", 0)
-              )
+              hard_attention_k=hparams.get("hard_attention_k", 0),
+              max_area_width=max_area_width,
+              max_area_height=max_area_height,
+              memory_height=memory_height,
+              area_key_mode=hparams.get("area_key_mode", "none"),
+              area_value_mode=hparams.get("area_value_mode", "none"),
+              training=(hparams.get("mode", tf.estimator.ModeKeys.TRAIN)
+                        == tf.estimator.ModeKeys.TRAIN))
           x = common_layers.layer_postprocess(x, y, hparams)
         if encoder_output is not None:
           with tf.variable_scope("encdec_attention"):
@@ -1474,7 +1489,14 @@ def transformer_decoder(decoder_input,
                 activation_dtype=hparams.get("activation_dtype", "float32"),
                 weight_dtype=hparams.get("weight_dtype", "float32"),
                 layer_collection=layer_collection,
-                hard_attention_k=hparams.get("hard_attention_k", 0))
+                hard_attention_k=hparams.get("hard_attention_k", 0),
+                max_area_width=max_area_width,
+                max_area_height=max_area_height,
+                memory_height=memory_height,
+                area_key_mode=hparams.get("area_key_mode", "none"),
+                area_value_mode=hparams.get("area_value_mode", "none"),
+                training=(hparams.get("mode", tf.estimator.ModeKeys.TRAIN)
+                          == tf.estimator.ModeKeys.TRAIN))
             x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(

From a0cc5541460414e38e5be06a019738208372be84 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 29 Apr 2019 09:04:27 -0700
Subject: [PATCH 1964/2720] Make multi-step Adam work again.

PiperOrigin-RevId: 245756861
---
 tensor2tensor/layers/common_attention.py | 2 --
 tensor2tensor/layers/common_hparams.py   | 2 +-
 tensor2tensor/models/transformer.py      | 2 +-
 3 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index d49d8a6d0..07c66c73c 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -4186,8 +4186,6 @@ def multihead_attention(query_antecedent,
       if isinstance(x, tuple):
         x, additional_returned_value = x  # Unpack
     elif attention_type == "dot_product":
-      tf.logging.info("max_area_width=%d, max_area_height=%d",
-                      max_area_width, max_area_height)
       if max_area_width > 1 or max_area_height > 1:
         x = area_attention.dot_product_area_attention(
             q, k, v, bias, dropout_rate, image_shapes,
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 5aba1af63..26703e78e 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -70,7 +70,7 @@ def basic_params1():
       optimizer_adafactor_clipping_threshold=1.0,
       optimizer_adafactor_multiply_by_parameter_scale=True,
       # Number of accumulating steps for multi step optimizers.
-      optimizer_multistep_accumulate_steps=None,
+      optimizer_multistep_accumulate_steps=0,
       # Loss scaling used.
       # Generally only necessary with mixed precision training.
       # Mixed precision training only supports exponential scaling currently
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index e87ba9c0d..5b5a1b1a4 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1974,7 +1974,7 @@ def transformer_base_single_gpu():
 def transformer_base_multistep8():
   """HParams for simulating 8 GPUs with MultistepAdam optimizer."""
   hparams = transformer_base()
-  hparams.optimizer = "MultistepAdam"
+  hparams.optimizer = "multistep_adam"
   hparams.optimizer_multistep_accumulate_steps = 8
   return hparams
 

From 4e805cc13b9d36b2778e97ad2315532887061201 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 29 Apr 2019 12:42:10 -0700
Subject: [PATCH 1965/2720] Print the number of parameters in trax and add an
 option to save XLA graph.

PiperOrigin-RevId: 245803347
---
 tensor2tensor/trax/layers/base.py | 21 +++++++++++++++++++++
 tensor2tensor/trax/trax.py        | 29 +++++++++++++++++++++++++----
 2 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index b7e45a641..5f9285c9a 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -155,6 +155,17 @@ def nested_map(x, f):
   return f(x)
 
 
+def nested_reduce(x, f):
+  """Fold the function f to the nested structure x (dicts, tuples, lists)."""
+  if isinstance(x, list):
+    return f([nested_reduce(y, f) for y in x])
+  if isinstance(x, tuple):
+    return f(tuple([nested_reduce(y, f) for y in x]))
+  if isinstance(x, dict):
+    return f({k: nested_reduce(x[k], f) for k in x})
+  return x
+
+
 def shapes(x):
   """Get a structure of shapes for a structure of nested arrays."""
   def shape(x):
@@ -165,6 +176,16 @@ def shape(x):
   return nested_map(x, shape)
 
 
+def sizes(x):
+  """Get a structure of sizes for a structure of nested arrays."""
+  def size(x):
+    try:
+      return x.size
+    except Exception:  # pylint: disable=broad-except
+      return 0
+  return nested_map(x, size)
+
+
 def _find_frame(stack, start=0):
   """Find the frame with the caller on the stack."""
   # We want to find the first place where the layer was called
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 9f1a3eb37..7908aabba 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -293,8 +293,13 @@ def mapped_predict(x, params, rng):
         reshape_by_device(x, num_devices),
         params,
         jax_random.split(rng, num_devices))
-    batch_size = x.shape[0]
-    return np.reshape(pred, [batch_size] + list(pred.shape[2:]))
+    # Need to reduce the [device, per-device-batch, ...] tensors back to
+    # a [batch, ...] tensor. The tensors may be nested.
+    if not isinstance(x, (list, tuple)):  # Not nested.
+      batch_size = x.shape[0]
+      return np.reshape(pred, [batch_size] + list(pred.shape[2:]))
+    batch_size = x[0].shape[0]
+    return [np.reshape(p, [batch_size] + list(p.shape[2:])) for p in pred]
 
   return predict
 
@@ -362,7 +367,8 @@ def train(output_dir,
           eval_frequency=100,
           num_devices=None,
           random_seed=None,
-          run_debug_step=False):
+          run_debug_step=False,
+          save_forward_graph=False):
   """Train the model on the inputs.
 
   Args:
@@ -386,6 +392,7 @@ def train(output_dir,
     random_seed: the random seed to use; time/os dependent if None (default).
     run_debug_step: bool, if True, will run the model and loss without @jit for
       one step.
+    save_forward_graph: bool, if True, save forward computation graph to file.
 
   Returns:
     trax.State
@@ -476,8 +483,14 @@ def train(output_dir,
       train_sw.scalar("training/steps per second",
                       epoch_steps / epoch_time, step=step)
 
-    # Evaluate
+    # Print number of parameters
     params = trax_opt.get_params(opt_state)
+    if step == 1:
+      sizes = layers.sizes(params)
+      total_size = layers.nested_reduce(sizes, sum)
+      step_log(step, "Total trainable parameters size: %d" % total_size)
+
+    # Evaluate
     evaluate_train_and_eval(
         step=step,
         inputs=inputs,
@@ -488,6 +501,14 @@ def train(output_dir,
         eval_sw=eval_sw,
         history=history)
 
+    # Save computation graph
+    if save_forward_graph and step == 1:
+      # Dump forward computation graph to file.
+      computation = jax.xla_computation(model_predict_eval)(
+          next_train_batch[0], params=params, rng=rng)
+      with gfile.GFile(os.path.join(output_dir, "forward_graph.dot"), "w") as f:
+        f.write(computation.GetHloDotGraph())
+
     # Save state
     save_state(State(params=params, step=step, history=history), output_dir)
 

From 50c883f12edeb907193422843998c299dbe0cb53 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Mon, 29 Apr 2019 23:07:06 +0200
Subject: [PATCH 1966/2720] Parameter set for sticky actions. (#1561)

---
 tensor2tensor/rl/trainer_model_based_params.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 2776084f3..a583febd5 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -323,6 +323,14 @@ def rlmb_base_stochastic_discrete():
   hparams.simulated_batch_size = 16
   return hparams
 
+@registry.register_hparams
+def rlmb_base_stochastic_discrete_sticky_actions():
+  """Base setting with stochastic discrete model with sticky action
+  environment.
+  """
+  hparams = rlmb_base_stochastic_discrete()
+  hparams.sticky_actions = True
+  return hparams
 
 @registry.register_hparams
 def rlmb_base_stochastic_discrete_75k_model_steps():

From a0e8429ec06d14ae0624e3c51b04f015c13985a1 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Mon, 29 Apr 2019 14:07:22 -0700
Subject: [PATCH 1967/2720] Merge of PR #1561

PiperOrigin-RevId: 245819502
---
 tensor2tensor/rl/trainer_model_based_params.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index a583febd5..65d8e6046 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -323,15 +323,15 @@ def rlmb_base_stochastic_discrete():
   hparams.simulated_batch_size = 16
   return hparams
 
+
 @registry.register_hparams
 def rlmb_base_stochastic_discrete_sticky_actions():
-  """Base setting with stochastic discrete model with sticky action
-  environment.
-  """
+  """Base setting, stochastic discrete model with sticky action environment."""
   hparams = rlmb_base_stochastic_discrete()
   hparams.sticky_actions = True
   return hparams
 
+
 @registry.register_hparams
 def rlmb_base_stochastic_discrete_75k_model_steps():
   """Base setting with stochastic discrete model with 75k WM steps."""

From 04020ef659840890af71fd33cc7f18e4c5fe26ea Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 29 Apr 2019 14:24:15 -0700
Subject: [PATCH 1968/2720] Policy and Value Net changes.  - Extract functions
 out of value_loss and ppo_loss that just take predictions, i.e. don't call a
 net.  - Combine these into a "combined_loss" function, which becomes the loss
 for the combined net.  - Optimization of the combined net against the
 combined loss.  - Standardize to "new" then "old" in function arguments.  -
 collect_trajectories now just takes a function to generate policy outputs.

PiperOrigin-RevId: 245823014
---
 tensor2tensor/trax/rlax/ppo.py                | 648 +++++++++++++-----
 tensor2tensor/trax/rlax/ppo_main.py           |  14 +-
 tensor2tensor/trax/rlax/ppo_test.py           | 110 ++-
 .../trax/rlax/ppo_training_loop_test.py       |  29 +-
 4 files changed, 604 insertions(+), 197 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index db21ad7af..5968d6da6 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -67,6 +67,7 @@
 DEBUG_LOGGING = False
 GAMMA = 0.99
 LAMBDA = 0.95
+EPSILON = 0.1
 EPOCHS = 50  # 100
 NUM_OPTIMIZER_STEPS = 100
 PRINT_EVERY_OPTIMIZER_STEP = 20
@@ -120,7 +121,7 @@ def policy_and_value_net(rng_key,
 
   # Now, with the current logits, one head computes action probabilities and the
   # other computes the value function.
-  # NOTE: The LogSoftmax instead of the Softmax.
+  # NOTE: The LogSoftmax instead of the Softmax because of numerical stability.
   cur_layers.extend([layers.Branch(), layers.Parallel(
       layers.Serial(layers.Dense(num_actions), layers.LogSoftmax()),
       layers.Dense(1)
@@ -142,13 +143,13 @@ def log_params(params, name="params"):
     if not param:
       # Empty tuple.
       continue
-    if not isinstance(param, tuple):
+    if not isinstance(param, (list, tuple)):
       logging.error(
           "%s[%d] : (%s) = [%s]", name, i, param.shape, onp.array(param))
     else:
       for j, p in enumerate(param):
         logging.error(
-            "\t%s[%d, %d] : (%s) = [%s]", name, i, j, p.shape, onp.array(p))
+            "\t%s[%d, %d] = [%s]", name, i, j, onp.array(p))
 
 
 # Should this be collect 'n' trajectories, or
@@ -156,13 +157,31 @@ def log_params(params, name="params"):
 # Any other option?
 # TODO(afrozm): Replace this with EnvProblem?
 def collect_trajectories(env,
-                         policy_net_apply,
-                         policy_net_params,
+                         policy_fun,
                          num_trajectories=1,
                          policy="greedy",
                          max_timestep=None,
                          epsilon=0.1):
-  """Collect trajectories with the given policy net and behaviour."""
+  """Collect trajectories with the given policy net and behaviour.
+
+  Args:
+    env: A gym env interface, for now this is not-batched.
+    policy_fun: observations(B,T+1) -> log-probabs(B,T+1, A) callable.
+    num_trajectories: int, number of trajectories.
+    policy: string, "greedy", "epsilon-greedy", or "categorical-sampling" i.e.
+        how to use the policy_fun to return an action.
+    max_timestep: int or None, the index of the maximum time-step at which we
+        return the trajectory, None for ending a trajectory only when env
+        returns done.
+    epsilon: float, the epsilon for `epsilon-greedy` policy.
+
+  Returns:
+    trajectory: list of (observation, action, reward) tuples, where each element
+    `i` is a tuple of numpy arrays with shapes as follows:
+    observation[i] = (B, T_i + 1)
+    action[i] = (B, T_i)
+    reward[i] = (B, T_i)
+  """
   trajectories = []
 
   for t in range(num_trajectories):
@@ -185,7 +204,7 @@ def collect_trajectories(env,
       ts_start = time.time()
       # Run the policy, to pick an action, shape is (1, t, A) because
       # observation_history is shaped (1, t) + OBS
-      predictions = policy_net_apply(observation_history, policy_net_params)
+      predictions = policy_fun(observation_history)
 
       # We need the predictions for the last time-step, so squeeze the batch
       # dimension and take the last time-step.
@@ -226,8 +245,6 @@ def collect_trajectories(env,
         logging.error("predictions.shape: [%s]", predictions.shape)
         logging.error("predictions: [%s]", predictions)
         logging.error("observation_history: [%s]", observation_history)
-        logging.error("policy_net_params: [%s]", policy_net_params)
-        log_params(policy_net_params, "policy_net_params")
         raise err
 
       observation, reward, done, _ = env.step(action)
@@ -427,16 +444,42 @@ def value_loss(value_net_apply,
   """
 
   B, T = rewards.shape  # pylint: disable=invalid-name
-  assert (B, T) == reward_mask.shape
   assert (B, T + 1) == observations.shape[:2]
 
-  r2g = rewards_to_go(rewards, reward_mask, gamma=gamma)  # (B, T)
   # NOTE: observations is (B, T+1) + OBS, value_prediction is (B, T+1, 1)
   value_prediction = value_net_apply(observations, value_net_params)
   assert (B, T + 1, 1) == value_prediction.shape
+
+  return value_loss_given_predictions(value_prediction, rewards, reward_mask,
+                                      gamma)
+
+
+@jit
+def value_loss_given_predictions(value_prediction,
+                                 rewards,
+                                 reward_mask,
+                                 gamma=0.99):
+  """Computes the value loss given the prediction of the value function.
+
+  Args:
+    value_prediction: np.ndarray of shape (B, T+1, 1)
+    rewards: np.ndarray of shape (B, T) of rewards.
+    reward_mask: np.ndarray of shape (B, T), the mask over rewards.
+    gamma: float, discount factor.
+
+  Returns:
+    The average L2 value loss, averaged over instances where reward_mask is 1.
+  """
+
+  B, T = rewards.shape  # pylint: disable=invalid-name
+  assert (B, T) == reward_mask.shape
+  assert (B, T + 1, 1) == value_prediction.shape
+
   value_prediction = np.squeeze(value_prediction, axis=2)  # (B, T+1)
   value_prediction = value_prediction[:, :-1] * reward_mask  # (B, T)
+  r2g = rewards_to_go(rewards, reward_mask, gamma=gamma)  # (B, T)
   loss = (value_prediction - r2g)**2
+
   # Take an average on only the points where mask != 0.
   return np.sum(loss) / np.sum(reward_mask)
 
@@ -508,18 +551,18 @@ def chosen_probabs(probab_observations, actions):
     `[B, T]` ndarray with the log-probabilities of the chosen actions.
   """
   B, T = actions.shape  # pylint: disable=invalid-name
-  assert (B, T+1) == probab_observations.shape[:2]
+  assert (B, T + 1) == probab_observations.shape[:2]
   return probab_observations[np.arange(B)[:, None], np.arange(T), actions]
 
 
-def compute_probab_ratios(p_old, p_new, actions, reward_mask):
+def compute_probab_ratios(p_new, p_old, actions, reward_mask):
   """Computes the probability ratios for each time-step in a trajectory.
 
   Args:
-    p_old: ndarray of shape [B, T+1, A] of the log-probabilities that the policy
+    p_new: ndarray of shape [B, T+1, A] of the log-probabilities that the policy
       network assigns to all the actions at each time-step in each batch using
       the old parameters.
-    p_new: ndarray of shape [B, T+1, A], same as above, but using new policy
+    p_old: ndarray of shape [B, T+1, A], same as above, but using old policy
       network parameters.
     actions: ndarray of shape [B, T] where each element is from [0, A).
     reward_mask: ndarray of shape [B, T] masking over probabilities.
@@ -530,8 +573,8 @@ def compute_probab_ratios(p_old, p_new, actions, reward_mask):
   """
 
   B, T = actions.shape  # pylint: disable=invalid-name
-  assert (B, T+1) == p_old.shape[:2]
-  assert (B, T+1) == p_new.shape[:2]
+  assert (B, T + 1) == p_old.shape[:2]
+  assert (B, T + 1) == p_new.shape[:2]
 
   logp_old = chosen_probabs(p_old, actions)
   logp_new = chosen_probabs(p_new, actions)
@@ -568,16 +611,59 @@ def ppo_loss(policy_net_apply,
              gamma=0.99,
              lambda_=0.95,
              epsilon=0.2):
-  """PPO objective, with an eventual minus sign."""
+  """PPO objective, with an eventual minus sign, given observations."""
   B, T = padded_rewards.shape  # pylint: disable=invalid-name
-  assert (B, T+1) == padded_observations.shape[:2]
+  assert (B, T + 1) == padded_observations.shape[:2]
   assert (B, T) == padded_actions.shape
   assert (B, T) == padded_rewards.shape
   assert (B, T) == reward_mask.shape
 
+  # Compute predicted values and predicted log-probs and hand it over to
+  # `ppo_loss_given_predictions`.
+
   # (B, T+1, 1)
   predicted_values = value_net_apply(padded_observations, value_net_params)
-  assert (B, T+1, 1) == predicted_values.shape
+  assert (B, T + 1, 1) == predicted_values.shape
+
+  # log_probab_actions_{old,new} are both (B, T+1, A)
+  log_probab_actions_old = policy_net_apply(padded_observations,
+                                            old_policy_params)
+  log_probab_actions_new = policy_net_apply(padded_observations,
+                                            new_policy_params)
+  assert (B, T + 1) == log_probab_actions_old.shape[:2]
+  assert (B, T + 1) == log_probab_actions_new.shape[:2]
+  assert log_probab_actions_old.shape[-1] == log_probab_actions_new.shape[-1]
+
+  return ppo_loss_given_predictions(log_probab_actions_new,
+                                    log_probab_actions_old,
+                                    predicted_values,
+                                    padded_actions,
+                                    padded_rewards,
+                                    reward_mask,
+                                    gamma=gamma,
+                                    lambda_=lambda_,
+                                    epsilon=epsilon)
+
+
+@jit
+def ppo_loss_given_predictions(log_probab_actions_new,
+                               log_probab_actions_old,
+                               predicted_values,
+                               padded_actions,
+                               padded_rewards,
+                               reward_mask,
+                               gamma=0.99,
+                               lambda_=0.95,
+                               epsilon=0.2):
+  """PPO objective, with an eventual minus sign, given predictions."""
+  B, T = padded_rewards.shape  # pylint: disable=invalid-name
+  assert (B, T) == padded_actions.shape
+  assert (B, T) == reward_mask.shape
+
+  _, _, A = log_probab_actions_old.shape  # pylint: disable=invalid-name
+  assert (B, T + 1, 1) == predicted_values.shape
+  assert (B, T + 1, A) == log_probab_actions_old.shape
+  assert (B, T + 1, A) == log_probab_actions_new.shape
 
   # (B, T)
   td_deltas = deltas(
@@ -585,26 +671,16 @@ def ppo_loss(policy_net_apply,
       padded_rewards,
       reward_mask,
       gamma=gamma)
-  assert (B, T) == td_deltas.shape
 
   # (B, T)
   advantages = gae_advantages(
       td_deltas, reward_mask, lambda_=lambda_, gamma=gamma)
-  assert (B, T) == advantages.shape
-
-  # probab_actions_{old,new} are both (B, T+1, A)
-  log_probab_actions_old = policy_net_apply(padded_observations,
-                                            old_policy_params)
-  log_probab_actions_new = policy_net_apply(padded_observations,
-                                            new_policy_params)
-  assert (B, T+1) == log_probab_actions_old.shape[:2]
-  assert (B, T+1) == log_probab_actions_new.shape[:2]
-  assert log_probab_actions_old.shape[-1] == log_probab_actions_new.shape[-1]
 
   # (B, T)
-  ratios = compute_probab_ratios(log_probab_actions_old,
-                                 log_probab_actions_new,
-                                 padded_actions, reward_mask)
+  ratios = compute_probab_ratios(log_probab_actions_new,
+                                 log_probab_actions_old,
+                                 padded_actions,
+                                 reward_mask)
   assert (B, T) == ratios.shape
 
   # (B, T)
@@ -619,6 +695,72 @@ def ppo_loss(policy_net_apply,
   return -average_objective
 
 
+@jit
+def combined_loss_given_predictions(log_probab_actions_new,
+                                    log_probab_actions_old,
+                                    value_prediction,
+                                    padded_actions,
+                                    padded_rewards,
+                                    reward_mask,
+                                    gamma=0.99,
+                                    lambda_=0.95,
+                                    epsilon=0.2,
+                                    c1=1.0,
+                                    c2=0.01):
+  """Computes the combined (clipped loss + value loss) given predictions."""
+  loss_value = value_loss_given_predictions(
+      value_prediction, padded_rewards, reward_mask, gamma=gamma)
+  loss_ppo = ppo_loss_given_predictions(log_probab_actions_new,
+                                        log_probab_actions_old,
+                                        value_prediction,
+                                        padded_actions,
+                                        padded_rewards,
+                                        reward_mask,
+                                        gamma=gamma,
+                                        lambda_=lambda_,
+                                        epsilon=epsilon)
+  # TODO(afrozm): Add the entropy bonus, but since we don't do that in T2T
+  # we'll skip if for now.
+  entropy_bonus = 0.0
+  return (loss_ppo + (c1 * loss_value) - (c2 * entropy_bonus), loss_ppo,
+          loss_value, entropy_bonus)
+
+
+# TODO(afrozm): Pass in `log_probab_actions_old` instead of re=computing it.
+@functools.partial(jit, static_argnums=(2,))
+def combined_loss(new_params,
+                  old_params,
+                  policy_and_value_net_apply,
+                  padded_observations,
+                  padded_actions,
+                  padded_rewards,
+                  reward_mask,
+                  gamma=0.99,
+                  lambda_=0.95,
+                  epsilon=0.2,
+                  c1=1.0,
+                  c2=0.01):
+  """Computes the combined (clipped loss + value loss) given observations."""
+  log_probab_actions_new, value_predictions = policy_and_value_net_apply(
+      padded_observations, new_params)
+
+  log_probab_actions_old, _ = policy_and_value_net_apply(
+      padded_observations, old_params)
+
+  # (combined_loss, ppo_loss, value_loss, entropy_bonus)
+  return combined_loss_given_predictions(log_probab_actions_new,
+                                         log_probab_actions_old,
+                                         value_predictions,
+                                         padded_actions,
+                                         padded_rewards,
+                                         reward_mask,
+                                         c1=c1,
+                                         c2=c2,
+                                         gamma=gamma,
+                                         lambda_=lambda_,
+                                         epsilon=epsilon)
+
+
 @functools.partial(jit, static_argnums=(2, 3, 5))
 def ppo_opt_step(i,
                  opt_state,
@@ -674,32 +816,77 @@ def value_opt_step(i,
   return opt_update(i, g, opt_state)
 
 
+@functools.partial(jit, static_argnums=(2, 3))
+def policy_and_value_opt_step(i,
+                              opt_state,
+                              opt_update,
+                              policy_and_value_net_apply,
+                              old_params,
+                              padded_observations,
+                              padded_actions,
+                              padded_rewards,
+                              reward_mask,
+                              c1=1.0,
+                              c2=0.01,
+                              gamma=0.99,
+                              lambda_=0.95,
+                              epsilon=0.1):
+  """Policy and Value optimizer step."""
+  # Combined loss function given the new params.
+  def policy_and_value_loss(params):
+    """Returns the combined loss given just parameters."""
+    (loss, _, _, _) = combined_loss(
+        params,
+        old_params,
+        policy_and_value_net_apply,
+        padded_observations,
+        padded_actions,
+        padded_rewards,
+        reward_mask,
+        c1=c1,
+        c2=c2,
+        gamma=gamma,
+        lambda_=lambda_,
+        epsilon=epsilon)
+    return loss
+
+  new_params = trax_opt.get_params(opt_state)
+  g = grad(policy_and_value_loss)(new_params)
+  return opt_update(i, g, opt_state)
+
+
 def get_time(t1, t2=None):
   if t2 is None:
     t2 = time.time()
   return round((t2 - t1) * 1000, 2)
 
 
-def training_loop(
-    env=None,
-    env_name="CartPole-v0",
-    epochs=EPOCHS,
-    policy_net_fun=None,
-    value_net_fun=None,
-    policy_and_value_net_fun=None,  # TODO(afrozm): Implement.
-    policy_optimizer_fun=optimizer_fun,
-    value_optimizer_fun=optimizer_fun,
-    batch_size=BATCH_TRAJECTORIES,
-    num_optimizer_steps=NUM_OPTIMIZER_STEPS,
-    print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
-    boundary=20,
-    max_timestep=None,
-    random_seed=None):
+def training_loop(env=None,
+                  env_name="CartPole-v0",
+                  epochs=EPOCHS,
+                  policy_net_fun=None,
+                  value_net_fun=None,
+                  policy_and_value_net_fun=None,
+                  policy_optimizer_fun=None,
+                  value_optimizer_fun=None,
+                  policy_and_value_optimizer_fun=None,
+                  batch_size=BATCH_TRAJECTORIES,
+                  num_optimizer_steps=NUM_OPTIMIZER_STEPS,
+                  print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
+                  boundary=20,
+                  max_timestep=None,
+                  random_seed=None,
+                  gamma=GAMMA,
+                  lambda_=LAMBDA,
+                  epsilon=EPSILON,
+                  c1=1.0,
+                  c2=0.01):
   """Runs the training loop for PPO, with fixed policy and value nets."""
   jax_rng_key = trax.get_random_number_generator_and_set_seed(random_seed)
 
   value_losses = []
   ppo_objective = []
+  combined_losses = []
   average_rewards = []
 
   env = env if env is not None else gym.make(env_name)
@@ -711,24 +898,45 @@ def training_loop(
   assert isinstance(env.action_space, gym.spaces.Discrete)
   num_actions = env.action_space.n
 
-  # TODO(afrozm): Have a single net for both policy and action.
-  assert policy_and_value_net_fun is None
-
-  # Initialize the policy and value functions.
-  assert policy_net_fun and value_net_fun
-  jax_rng_key, key1, key2 = jax_random.split(jax_rng_key, num=3)
+  policy_and_value_net_params, policy_and_value_net_apply = None, None
+  policy_and_value_opt_state, policy_and_value_opt_update = None, None
+  policy_net_params, policy_net_apply = None, None
+  value_net_params, value_net_apply = None, None
+  if policy_and_value_net_fun is not None:
+    jax_rng_key, subkey = jax_random.split(jax_rng_key)
 
-  policy_net_params, policy_net_apply = policy_net_fun(
-      key1, batch_observations_shape, num_actions)
-  value_net_params, value_net_apply = value_net_fun(key2,
-                                                    batch_observations_shape,
-                                                    num_actions)
+    # Initialize the policy and value network.
+    policy_and_value_net_params, policy_and_value_net_apply = (
+        policy_and_value_net_fun(subkey, batch_observations_shape, num_actions))
 
-  # Initialize the optimizers.
-  assert policy_optimizer_fun and value_optimizer_fun
-
-  ppo_opt_state, ppo_opt_update = policy_optimizer_fun(policy_net_params)
-  value_opt_state, value_opt_update = value_optimizer_fun(value_net_params)
+    # Initialize the optimizers.
+    policy_and_value_opt_state, policy_and_value_opt_update = (
+        policy_and_value_optimizer_fun(policy_and_value_net_params))
+  else:
+    # Initialize the policy and value functions.
+    assert policy_net_fun and value_net_fun
+    jax_rng_key, key1, key2 = jax_random.split(jax_rng_key, num=3)
+
+    policy_net_params, policy_net_apply = policy_net_fun(
+        key1, batch_observations_shape, num_actions)
+    value_net_params, value_net_apply = value_net_fun(key2,
+                                                      batch_observations_shape,
+                                                      num_actions)
+
+    # Initialize the optimizers.
+    ppo_opt_state, ppo_opt_update = policy_optimizer_fun(policy_net_params)
+    value_opt_state, value_opt_update = value_optimizer_fun(value_net_params)
+
+  # A function that will call the appropriate policy function with parameters.
+  def get_policy_output(observations):
+    if policy_net_apply is not None:
+      assert policy_net_params
+      return policy_net_apply(observations, policy_net_params)
+
+    assert policy_and_value_net_apply and policy_and_value_net_params
+    policy_predictions, unused_value_predictions = policy_and_value_net_apply(
+        observations, policy_and_value_net_params)
+    return policy_predictions
 
   for i in range(epochs):
     t = time.time()
@@ -736,8 +944,7 @@ def training_loop(
     logging.vlog(1, "Epoch [% 6d] collecting trajectories.", i)
     trajs = collect_trajectories(
         env,
-        policy_net_apply,
-        policy_net_params,
+        policy_fun=get_policy_output,
         num_trajectories=batch_size,
         policy=POLICY,
         max_timestep=max_timestep,
@@ -775,68 +982,137 @@ def training_loop(
     assert (B, T + 1) + env.observation_space.shape == padded_observations.shape
 
     # Linear annealing from 0.1 to 0.0
-    epsilon = 0.1 if epochs == 1 else 0.1 * (1.0 - (i / (epochs - 1)))
-
-    t = time.time()
-    cur_value_loss = value_loss(
-        value_net_apply,
-        value_net_params,
-        padded_observations,
-        padded_rewards,
-        reward_mask,
-        gamma=GAMMA)
+    epsilon_schedule = epsilon if epochs == 1 else epsilon * (1.0 -
+                                                              (i /
+                                                               (epochs - 1)))
 
-    logging.vlog(1, "Calculating value loss took %0.2f msec.", get_time(t))
-    value_losses.append(cur_value_loss)
-
-    t = time.time()
-    cur_ppo_loss = ppo_loss(
-        policy_net_apply,
-        policy_net_params,
-        policy_net_params,
-        value_net_apply,
-        value_net_params,
-        padded_observations,
-        padded_actions,
-        padded_rewards,
-        reward_mask,
-        gamma=GAMMA,
-        lambda_=LAMBDA,
-        epsilon=epsilon)
-    # ppo_loss = 11.00110011
-    logging.vlog(1, "Calculating PPO loss took %0.2f msec.", get_time(t))
-    ppo_objective.append(-cur_ppo_loss)
+    # Compute value and ppo losses.
+    cur_value_loss, cur_ppo_loss, cur_combined_loss = None, None, None
+    if policy_and_value_net_apply is not None:
+      t = time.time()
+      cur_combined_loss, cur_ppo_loss, cur_value_loss, _ = (
+          combined_loss(
+              policy_and_value_net_params,
+              policy_and_value_net_params,
+              policy_and_value_net_apply,
+              padded_observations,
+              padded_actions,
+              padded_rewards,
+              reward_mask,
+              gamma=gamma,
+              lambda_=lambda_,
+              epsilon=epsilon_schedule,
+              c1=c1,
+              c2=c2))
+      logging.vlog(
+          1, "Calculating P&V loss [%10.2f(%10.2f, %10.2f)] took %0.2f msec.",
+          cur_combined_loss, cur_value_loss, cur_ppo_loss, get_time(t))
+    else:
+      t = time.time()
+      cur_value_loss = value_loss(
+          value_net_apply,
+          value_net_params,
+          padded_observations,
+          padded_rewards,
+          reward_mask,
+          gamma=gamma)
 
-    # Run optimizers.
-    logging.vlog(1, "PPO Optimization")
-    t1 = time.time()
+      logging.vlog(1, "Calculating value loss took %0.2f msec.", get_time(t))
 
-    for j in range(num_optimizer_steps):
       t = time.time()
-      # Update the optimizer state.
-      ppo_opt_state = ppo_opt_step(
-          j,
-          ppo_opt_state,
-          ppo_opt_update,
+      cur_ppo_loss = ppo_loss(
           policy_net_apply,
           policy_net_params,
+          policy_net_params,
           value_net_apply,
           value_net_params,
           padded_observations,
           padded_actions,
           padded_rewards,
           reward_mask,
-          gamma=GAMMA,
-          lambda_=LAMBDA,
-          epsilon=epsilon)
-      t2 = time.time()
-      # Get the new params.
-      new_policy_net_params = trax_opt.get_params(ppo_opt_state)
-      if ((j + 1) %
-          print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
-        new_ppo_loss = ppo_loss(
+          gamma=gamma,
+          lambda_=lambda_,
+          epsilon=epsilon_schedule)
+      logging.vlog(1, "Calculating PPO loss took %0.2f msec.", get_time(t))
+
+    value_losses.append(cur_value_loss)
+    ppo_objective.append(-1.0 * cur_ppo_loss)
+    combined_losses.append(cur_combined_loss)
+
+    if policy_and_value_net_apply:
+      logging.vlog(1, "Policy and Value Optimization")
+      t1 = time.time()
+      for j in range(num_optimizer_steps):
+        t = time.time()
+        # Update the optimizer state.
+        policy_and_value_opt_state = policy_and_value_opt_step(
+            j,
+            policy_and_value_opt_state,
+            policy_and_value_opt_update,
+            policy_and_value_net_apply,
+            policy_and_value_net_params,
+            padded_observations,
+            padded_actions,
+            padded_rewards,
+            reward_mask,
+            c1=c1,
+            c2=c2,
+            gamma=gamma,
+            lambda_=lambda_,
+            epsilon=epsilon_schedule)
+        t2 = time.time()
+        # Get the new params.
+        new_policy_and_value_net_params = trax_opt.get_params(
+            policy_and_value_opt_state)
+        if ((j + 1) %
+            print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
+          # Compute and log the loss.
+          (loss_combined, loss_ppo, loss_value, unused_entropy_bonus) = (
+              combined_loss(
+                  new_policy_and_value_net_params,
+                  policy_and_value_net_params,  # old params
+                  policy_and_value_net_apply,
+                  padded_observations,
+                  padded_actions,
+                  padded_rewards,
+                  reward_mask,
+                  gamma=gamma,
+                  lambda_=lambda_,
+                  epsilon=epsilon_schedule,
+                  c1=c1,
+                  c2=c2))
+          logging.vlog(1, "One Policy and Value grad desc took: %0.2f msec",
+                       get_time(t, t2))
+          logging.vlog(
+              1,
+              "Combined Loss(value, ppo) [%10.2f] -> [%10.2f(%10.2f,%10.2f)]",
+              cur_combined_loss, loss_combined, loss_value, loss_ppo)
+        # Update the params.
+        policy_and_value_net_params = new_policy_and_value_net_params
+
+      logging.vlog(
+          1, "Total PPO loss reduction [%0.2f]%%",
+          (100 *
+           (cur_combined_loss - loss_combined) / np.abs(cur_combined_loss)))
+
+      logging.info(
+          "Epoch [% 6d], Reward[min, max, avg] [%10.2f,%10.2f,%10.2f], Combined"
+          " Loss(value, ppo) [%10.2f(%10.2f,%10.2f)], took [%10.2f msec]",
+          i, min_reward, max_reward, avg_reward, loss_combined, loss_value,
+          loss_ppo, get_time(t1))
+    else:
+      # Run optimizers.
+      logging.vlog(1, "PPO Optimization")
+      t1 = time.time()
+
+      for j in range(num_optimizer_steps):
+        t = time.time()
+        # Update the optimizer state.
+        ppo_opt_state = ppo_opt_step(
+            j,
+            ppo_opt_state,
+            ppo_opt_update,
             policy_net_apply,
-            new_policy_net_params,
             policy_net_params,
             value_net_apply,
             value_net_params,
@@ -844,64 +1120,96 @@ def training_loop(
             padded_actions,
             padded_rewards,
             reward_mask,
-            gamma=GAMMA,
-            lambda_=LAMBDA,
-            epsilon=epsilon)
-        logging.vlog(1, "One PPO grad desc took: %0.2f msec", get_time(t, t2))
-        logging.vlog(1, "PPO loss [%10.2f] -> [%10.2f]", cur_ppo_loss,
-                     new_ppo_loss)
-      # Update the params.
-      policy_net_params = new_policy_net_params
-
-    logging.vlog(1, "Total PPO loss reduction [%0.2f]%%",
-                 (100 * (cur_ppo_loss - new_ppo_loss) / np.abs(cur_ppo_loss)))
-
-    logging.vlog(1, "Value Optimization")
-
-    for j in range(num_optimizer_steps):
-      t = time.time()
-      value_opt_state = value_opt_step(
-          j,
-          value_opt_state,
-          value_opt_update,
-          value_net_apply,
-          padded_observations,
-          padded_rewards,
-          reward_mask,
-          gamma=GAMMA)
-      t2 = time.time()
-      value_net_params = trax_opt.get_params(value_opt_state)
-      if ((j + 1) %
-          print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
-        new_value_loss = value_loss(
+            gamma=gamma,
+            lambda_=lambda_,
+            epsilon=epsilon_schedule,
+        )
+        t2 = time.time()
+        # Get the new params.
+        new_policy_net_params = trax_opt.get_params(ppo_opt_state)
+        if ((j + 1) %
+            print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
+          new_ppo_loss = ppo_loss(
+              policy_net_apply,
+              new_policy_net_params,
+              policy_net_params,
+              value_net_apply,
+              value_net_params,
+              padded_observations,
+              padded_actions,
+              padded_rewards,
+              reward_mask,
+              gamma=gamma,
+              lambda_=lambda_,
+              epsilon=epsilon_schedule,
+          )
+          logging.vlog(1, "One PPO grad desc took: %0.2f msec", get_time(t, t2))
+          logging.vlog(1, "PPO loss [%10.2f] -> [%10.2f]", cur_ppo_loss,
+                       new_ppo_loss)
+        # Update the params.
+        policy_net_params = new_policy_net_params
+
+      logging.vlog(1, "Total PPO loss reduction [%0.2f]%%",
+                   (100 * (cur_ppo_loss - new_ppo_loss) / np.abs(cur_ppo_loss)))
+
+      logging.vlog(1, "Value Optimization")
+
+      for j in range(num_optimizer_steps):
+        t = time.time()
+        value_opt_state = value_opt_step(
+            j,
+            value_opt_state,
+            value_opt_update,
             value_net_apply,
-            value_net_params,
             padded_observations,
             padded_rewards,
             reward_mask,
-            gamma=GAMMA)
-        logging.vlog(1, "One value grad desc took: %0.2f msec", get_time(t, t2))
-        logging.vlog(1, "Value loss [%10.2f] -> [%10.2f]", cur_value_loss,
-                     new_value_loss)
-    logging.vlog(1, "Total value loss reduction [%0.2f]%%",
-                 (100 *
-                  (cur_value_loss - new_value_loss) / np.abs(cur_value_loss)))
-
-    logging.vlog(1, "Grad desc took %0.2f msec", get_time(t1))
-
-    # Set the optimized params to new params.
-    policy_net_params = trax_opt.get_params(ppo_opt_state)
-    value_net_params = trax_opt.get_params(value_opt_state)
-
-    logging.info(
-        "Epoch [% 6d], Reward[min, max, avg] [%10.2f,%10.2f,%10.2f], "
-        "ppo loss [%10.2f], value loss [%10.2f], took [%10.2f msec]",
-        i, min_reward, max_reward, avg_reward, new_ppo_loss, new_value_loss,
-        get_time(t0))
-
-  logging.vlog(1, "value_losses: %s", np.stack(value_losses))
-  logging.vlog(1, "ppo_objective: %s", np.stack(ppo_objective))
-  logging.vlog(1, "average_rewards: %s", average_rewards)
+            gamma=gamma)
+        t2 = time.time()
+        value_net_params = trax_opt.get_params(value_opt_state)
+        if ((j + 1) %
+            print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
+          new_value_loss = value_loss(
+              value_net_apply,
+              value_net_params,
+              padded_observations,
+              padded_rewards,
+              reward_mask,
+              gamma=gamma)
+          logging.vlog(1, "One value grad desc took: %0.2f msec",
+                       get_time(t, t2))
+          logging.vlog(1, "Value loss [%10.2f] -> [%10.2f]", cur_value_loss,
+                       new_value_loss)
+      logging.vlog(1, "Total value loss reduction [%0.2f]%%",
+                   (100 *
+                    (cur_value_loss - new_value_loss) / np.abs(cur_value_loss)))
+
+      logging.vlog(1, "Grad desc took %0.2f msec", get_time(t1))
+
+      # Set the optimized params to new params.
+      policy_net_params = trax_opt.get_params(ppo_opt_state)
+      value_net_params = trax_opt.get_params(value_opt_state)
+
+      logging.info(
+          "Epoch [% 6d], Reward[min, max, avg] [%10.2f,%10.2f,%10.2f], "
+          "ppo loss [%10.2f], value loss [%10.2f], took [%10.2f msec]", i,
+          min_reward, max_reward, avg_reward, new_ppo_loss, new_value_loss,
+          get_time(t0))
+
+  # Log the parameters, just for the sake of it.
+  if policy_net_params:
+    log_params(policy_net_params, "policy_net_params")
+  if value_net_params:
+    log_params(value_net_params, "value_net_params")
+  if policy_and_value_net_params:
+    log_params(policy_and_value_net_params, "policy_and_value_net_params")
+
+  if value_losses:
+    logging.vlog(1, "value_losses: %s", np.stack(value_losses))
+  if ppo_objective:
+    logging.vlog(1, "ppo_objective: %s", np.stack(ppo_objective))
+  if average_rewards:
+    logging.vlog(1, "average_rewards: %s", average_rewards)
 
   return ((policy_net_params, value_net_params), average_rewards,
           np.stack(value_losses), np.stack(ppo_objective))
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 1e8968c39..edc4f43c1 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -43,7 +43,8 @@
                      "trajectory.")
 flags.DEFINE_float("learning_rate", 1e-3, "Learning rate.")
 flags.DEFINE_boolean("jax_debug_nans", False,
-                     "Setting to true will help to debug nans.")
+                     "Setting to true will help to debug nans and disable jit.")
+flags.DEFINE_boolean("disable_jit", False, "Setting to true will disable jit.")
 
 
 def common_layers():
@@ -67,19 +68,16 @@ def run_training_loop():
     ppo.training_loop(
         env_name=FLAGS.env_name,
         epochs=FLAGS.epochs,
-        policy_net_fun=functools.partial(
-            ppo.policy_net, bottom_layers=common_layers()),
-        value_net_fun=functools.partial(
-            ppo.value_net, bottom_layers=common_layers()),
-        policy_optimizer_fun=optimizer_fun,
-        value_optimizer_fun=optimizer_fun,
+        policy_and_value_net_fun=functools.partial(
+            ppo.policy_and_value_net, bottom_layers=common_layers()),
+        policy_and_value_optimizer_fun=optimizer_fun,
         batch_size=FLAGS.batch_size,
         num_optimizer_steps=FLAGS.num_optimizer_steps,
         boundary=FLAGS.boundary,
         max_timestep=FLAGS.max_timestep,
         random_seed=FLAGS.random_seed)
 
-  if FLAGS.jax_debug_nans:
+  if FLAGS.jax_debug_nans or FLAGS.disable_jit:
     with jax.disable_jit():
       run_training_loop()
   else:
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index cbd9b1d6d..7fe3ae898 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -114,9 +114,30 @@ def test_collect_trajectories(self):
     num_trajectories = 5
     trajectories = ppo.collect_trajectories(
         env,
-        policy_apply,
-        policy_params,
-        num_trajectories,
+        policy_fun=lambda obs: policy_apply(obs, policy_params),
+        num_trajectories=num_trajectories,
+        policy="categorical-sampling")
+
+    # Number of trajectories is as expected.
+    self.assertEqual(num_trajectories, len(trajectories))
+
+    # Shapes of observations, actions and rewards are as expected.
+    for observations, actions, rewards in trajectories:
+      # observations are one more in number than rewards or actions.
+      self.assertEqual((done_time_step + 2,) + observation_shape,
+                       observations.shape)
+      self.assertEqual((done_time_step + 1,), actions.shape)
+      self.assertEqual((done_time_step + 1,), rewards.shape)
+
+    # Test collect using a Policy and Value function.
+    pnv_params, pnv_apply = ppo.policy_and_value_net(
+        self.rng_key, (-1, -1) + observation_shape, num_actions,
+        [layers.Flatten(num_axis_to_keep=2)])
+
+    trajectories = ppo.collect_trajectories(
+        env,
+        policy_fun=lambda obs: pnv_apply(obs, pnv_params)[0],
+        num_trajectories=num_trajectories,
         policy="categorical-sampling")
 
     # Number of trajectories is as expected.
@@ -133,12 +154,8 @@ def test_collect_trajectories(self):
   def test_collect_trajectories_max_timestep(self):
     observation_shape = (2, 3, 4)
     num_actions = 2
-    policy_params, policy_apply = ppo.policy_net(
-        self.rng_key,
-        (-1, -1) + observation_shape,
-        num_actions,
-        # flatten except batch and time
-        # step dimensions.
+    pnv_params, pnv_apply = ppo.policy_and_value_net(
+        self.rng_key, (-1, -1) + observation_shape, num_actions,
         [layers.Flatten(num_axis_to_keep=2)])
 
     # We'll get done at time-step #5, starting from 0, therefore in 6 steps.
@@ -156,9 +173,8 @@ def test_collect_trajectories_max_timestep(self):
 
     trajectories = ppo.collect_trajectories(
         env,
-        policy_apply,
-        policy_params,
-        num_trajectories,
+        policy_fun=lambda obs: pnv_apply(obs, pnv_params)[0],
+        num_trajectories=num_trajectories,
         policy="categorical-sampling",
         max_timestep=max_timestep)
 
@@ -168,8 +184,7 @@ def test_collect_trajectories_max_timestep(self):
     # Shapes of observations, actions and rewards are as expected.
     for observations, actions, rewards in trajectories:
       # observations are one more in number than rewards or actions.
-      self.assertEqual((max_timestep,) + observation_shape,
-                       observations.shape)
+      self.assertEqual((max_timestep,) + observation_shape, observations.shape)
       self.assertEqual((max_timestep - 1,), actions.shape)
       self.assertEqual((max_timestep - 1,), rewards.shape)
 
@@ -271,7 +286,7 @@ def test_rewards_to_go_really_long_sequences(self):
     expected_r2g = np.zeros_like(masked_rewards)
     for t in range(T):
       for j in range(t, T):
-        expected_r2g[t] += (gamma**(j-t)) * masked_rewards[j]
+        expected_r2g[t] += (gamma**(j - t)) * masked_rewards[j]
 
     self.assertAllClose(expected_r2g, actual_r2g)
 
@@ -424,7 +439,7 @@ def test_compute_probab_ratios(self):
 
     mask = np.array([[1, 1, 0, 0], [1, 1, 1, 0]])
 
-    probab_ratios = ppo.compute_probab_ratios(p_old, p_new, actions, mask)
+    probab_ratios = ppo.compute_probab_ratios(p_new, p_old, actions, mask)
 
     self.assertAllClose(
         np.array([
@@ -516,6 +531,69 @@ def test_ppo_loss(self):
                      value_apply, value_params, observations, actions, rewards,
                      mask)
 
+  def test_combined_loss(self):
+    self.rng_key, key1, key2 = jax_random.split(self.rng_key, num=3)
+
+    B, T, A, OBS = 2, 10, 2, (28, 28, 3)  # pylint: disable=invalid-name
+    batch_observation_shape = (-1, -1) + OBS
+
+    old_params, _ = ppo.policy_and_value_net(
+        key1, batch_observation_shape, A, [layers.Flatten(num_axis_to_keep=2)])
+
+    new_params, net_apply = ppo.policy_and_value_net(
+        key2, batch_observation_shape, A, [layers.Flatten(num_axis_to_keep=2)])
+
+    # Generate a batch of observations.
+
+    observations = np.random.uniform(size=(B, T + 1) + OBS)
+    actions = np.random.randint(0, A, size=(B, T))
+    rewards = np.random.uniform(0, 1, size=(B, T))
+    mask = np.ones_like(rewards)
+
+    # Just test that this computes at all.
+    new_log_probabs, value_predictions = net_apply(observations, new_params)
+    old_log_probabs, _ = net_apply(observations, old_params)
+
+    gamma = 0.99
+    lambda_ = 0.95
+    epsilon = 0.2
+    c1 = 1.0
+    c2 = 0.01
+
+    value_loss_1 = ppo.value_loss_given_predictions(
+        value_predictions, rewards, mask, gamma=gamma)
+    ppo_loss_1 = ppo.ppo_loss_given_predictions(
+        new_log_probabs,
+        old_log_probabs,
+        value_predictions,
+        actions,
+        rewards,
+        mask,
+        gamma=gamma,
+        lambda_=lambda_,
+        epsilon=epsilon)
+
+    (combined_loss, ppo_loss_2, value_loss_2, entropy_bonus) = (
+        ppo.combined_loss(new_params,
+                          old_params,
+                          net_apply,
+                          observations,
+                          actions,
+                          rewards,
+                          mask,
+                          gamma=gamma,
+                          lambda_=lambda_,
+                          epsilon=epsilon,
+                          c1=c1,
+                          c2=c2)
+    )
+
+    # Test that these compute at all and are self consistent.
+    self.assertEqual(0.0, entropy_bonus)
+    self.assertNear(value_loss_1, value_loss_2, 1e-6)
+    self.assertNear(ppo_loss_1, ppo_loss_2, 1e-6)
+    self.assertNear(combined_loss, ppo_loss_2 + (c1 * value_loss_2), 1e-6)
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index a0aabc0da..8b80ada8f 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -29,12 +29,15 @@
 
 class PpoTrainingLoopTest(test.TestCase):
 
-  def test_training_loop(self):
-    env = gym.make("CartPole-v0")
+  def get_wrapped_env(self, name="CartPole-v0", max_episode_steps=2):
+    env = gym.make(name)
     # Usually gym envs are wrapped in TimeLimit wrapper.
     env = gym_utils.remove_time_limit_wrapper(env)
     # Limit this to a small number for tests.
-    env = gym.wrappers.TimeLimit(env, max_episode_steps=2)
+    return gym.wrappers.TimeLimit(env, max_episode_steps=max_episode_steps)
+
+  def test_training_loop(self):
+    env = self.get_wrapped_env("CartPole-v0", 2)
     num_epochs = 2
     batch_size = 2
     # Run the training loop.
@@ -45,6 +48,26 @@ def test_training_loop(self):
             ppo.policy_net, bottom_layers=[layers.Dense(1)]),
         value_net_fun=functools.partial(
             ppo.value_net, bottom_layers=[layers.Dense(1)]),
+        policy_optimizer_fun=ppo.optimizer_fun,
+        value_optimizer_fun=ppo.optimizer_fun,
+        batch_size=batch_size,
+        num_optimizer_steps=1,
+        random_seed=0)
+    self.assertLen(rewards, num_epochs)
+    self.assertLen(val_losses, num_epochs)
+    self.assertLen(ppo_objectives, num_epochs)
+
+  def test_training_loop_policy_and_value_function(self):
+    env = self.get_wrapped_env("CartPole-v0", 2)
+    num_epochs = 2
+    batch_size = 2
+    # Run the training loop.
+    _, rewards, val_losses, ppo_objectives = ppo.training_loop(
+        env=env,
+        epochs=num_epochs,
+        policy_and_value_net_fun=functools.partial(
+            ppo.policy_and_value_net, bottom_layers=[layers.Dense(1)]),
+        policy_and_value_optimizer_fun=ppo.optimizer_fun,
         batch_size=batch_size,
         num_optimizer_steps=1,
         random_seed=0)

From adf76adc8f45e1899e41808c860f509f633b18b3 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 29 Apr 2019 14:37:04 -0700
Subject: [PATCH 1969/2720] Fix Evolved Transformer TPU decoding and shrink
 tests to make them faster.

Evolved Transformer TPU fast decoding was broken due to an incorrect inplace update offset for the cache.

PiperOrigin-RevId: 245825698
---
 tensor2tensor/models/evolved_transformer.py   | 26 +++++++++++++++----
 .../models/evolved_transformer_test.py        | 17 +++++++-----
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index cf6468277..57adebad4 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -46,6 +46,22 @@
 _DECODER_FINAL_CONV_PADDING = 6
 
 
+def _capped_double_heads(num_heads, cap=16):
+  """Calculate the number of heads for the attention layers with more heads.
+
+  The number of heads will be twice the normal amount (num_heads), until it
+  reaches |cap| heads.
+
+  Args:
+    num_heads: the num_heads hparam for the model.
+    cap: the maximum number of heads |num_heads| will be doubled to.
+
+  Returns:
+    The number of heads for the attention layers that have more heads.
+  """
+  return max(min(num_heads * 2, cap), num_heads)
+
+
 @registry.register_model
 class EvolvedTransformer(transformer.Transformer):
   """The Evolved Transformer from arxiv.org/abs/1901.11117 ."""
@@ -289,7 +305,6 @@ def evolved_transformer_decoder(decoder_input,
           residual_state = hidden_state
           hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
 
-          # Attention with at least 16 heads.
           attention_cache = layer_cache[
               _SIXTEEN_HEAD_ATTENTION_NAME] if layer_cache is not None else None
           left_state = common_attention.multihead_attention(
@@ -299,7 +314,7 @@ def evolved_transformer_decoder(decoder_input,
               hparams.attention_key_channels or hparams.hidden_size,
               hparams.attention_value_channels or hparams.hidden_size,
               hparams.hidden_size,
-              max(16, hparams.num_heads),
+              _capped_double_heads(hparams.num_heads),
               hparams.attention_dropout,
               attention_type=hparams.self_attention_type,
               max_relative_position=hparams.max_relative_position,
@@ -385,7 +400,8 @@ def evolved_transformer_decoder(decoder_input,
               tmp = tf.expand_dims(tmp, axis=1)
               tmp = inplace_ops.alias_inplace_update(
                   tmp,
-                  decode_loop_step * tf.shape(hidden_state)[1],
+                  decode_loop_step * tf.shape(hidden_state)[1] +
+                  _DECODER_LEFT_CONV_PADDING,
                   tf.transpose(hidden_state, perm=[1, 0, 2]))
               tmp = tf.squeeze(tmp, axis=1)
               hidden_state = layer_cache[
@@ -622,12 +638,12 @@ def _init_evolved_transformer_cache(cache, hparams, batch_size,
                   common_attention.split_heads(
                       tf.zeros(
                           [batch_size, attention_init_length, key_channels]),
-                      max(16, hparams.num_heads)),
+                      _capped_double_heads(hparams.num_heads)),
               "v":
                   common_attention.split_heads(
                       tf.zeros(
                           [batch_size, attention_init_length, value_channels]),
-                      max(16, hparams.num_heads)),
+                      _capped_double_heads(hparams.num_heads)),
           },
           _VANILLA_ATTENTION_NAME: {
               "k":
diff --git a/tensor2tensor/models/evolved_transformer_test.py b/tensor2tensor/models/evolved_transformer_test.py
index 8691d162c..0761a06b4 100644
--- a/tensor2tensor/models/evolved_transformer_test.py
+++ b/tensor2tensor/models/evolved_transformer_test.py
@@ -30,12 +30,15 @@
 INPUT_LENGTH = 5
 TARGET_LENGTH = 7
 VOCAB_SIZE = 10
+DECODE_LENGTH = 3
 
 
 def get_model(hparams, has_input=True):
   hparams.layer_prepostprocess_dropout = 0.0
-  hparams.hidden_size = 16
+  hparams.hidden_size = 4
   hparams.num_heads = 1
+  hparams.num_encoder_layers = 1
+  hparams.num_decoder_layers = 1
 
   p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE, VOCAB_SIZE,
                                                    hparams)
@@ -70,7 +73,7 @@ def testEvolvedTransformer(self):
   def testSlowVsFast(self):
     model, features = get_model(transformer.transformer_tiny())
 
-    decode_length = 30
+    decode_length = DECODE_LENGTH
 
     out_logits, _ = model(features)
     out_logits = tf.squeeze(out_logits, axis=[2, 3])
@@ -105,7 +108,7 @@ def testSlowVsFastNoInput(self):
     model, features = get_model(
         transformer.transformer_tiny(), has_input=False)
 
-    decode_length = 30
+    decode_length = DECODE_LENGTH
 
     out_logits, _ = model(features)
     out_logits = tf.squeeze(out_logits, axis=[2, 3])
@@ -139,7 +142,7 @@ def testSlowVsFastNoInput(self):
   def testBeamVsFast(self):
     model, features = get_model(transformer.transformer_tiny())
 
-    decode_length = 30
+    decode_length = DECODE_LENGTH
 
     out_logits, _ = model(features)
     out_logits = tf.squeeze(out_logits, axis=[2, 3])
@@ -204,7 +207,7 @@ def _create_greedy_infer_model(self):
     return model, features
 
   def testGreedySlowTPUVsNonTPU(self):
-    decode_length = 30
+    decode_length = DECODE_LENGTH
 
     model, features = self._create_greedy_infer_model()
 
@@ -226,7 +229,7 @@ def testGreedySlowTPUVsNonTPU(self):
     self.assertAllClose(slow_tpu_res, slow_non_tpu_res)
 
   def testGreedyFastTPUVsNonTPU(self):
-    decode_length = 30
+    decode_length = DECODE_LENGTH
 
     model, features = self._create_greedy_infer_model()
 
@@ -246,7 +249,7 @@ def testGreedyFastTPUVsNonTPU(self):
     self.assertAllClose(fast_tpu_res, fast_non_tpu_res)
 
   def testGreedyTPUSlowVsFast(self):
-    decode_length = 30
+    decode_length = DECODE_LENGTH
 
     model, features = self._create_greedy_infer_model()
 

From b9b3ef675de940e976bee829763e5b5e97d8455d Mon Sep 17 00:00:00 2001
From: David Dohan <ddohan@google.com>
Date: Mon, 29 Apr 2019 14:49:03 -0700
Subject: [PATCH 1970/2720] {rnn,conv,parametrized} GRU implementation +
 NeuralGPU

PiperOrigin-RevId: 245828236
---
 tensor2tensor/trax/layers/__init__.py        |   1 +
 tensor2tensor/trax/layers/combinators.py     |  42 +++++-
 tensor2tensor/trax/layers/core.py            |  23 +++
 tensor2tensor/trax/layers/rnn.py             | 140 +++++++++++++++++++
 tensor2tensor/trax/layers/rnn_test.py        |  54 +++++++
 tensor2tensor/trax/models/__init__.py        |   6 +-
 tensor2tensor/trax/models/neural_gpu.py      |  82 +++++++++++
 tensor2tensor/trax/models/neural_gpu_test.py |  50 +++++++
 8 files changed, 393 insertions(+), 5 deletions(-)
 create mode 100644 tensor2tensor/trax/layers/rnn.py
 create mode 100644 tensor2tensor/trax/layers/rnn_test.py
 create mode 100644 tensor2tensor/trax/models/neural_gpu.py
 create mode 100644 tensor2tensor/trax/models/neural_gpu_test.py

diff --git a/tensor2tensor/trax/layers/__init__.py b/tensor2tensor/trax/layers/__init__.py
index b359e2a1a..25772a95f 100644
--- a/tensor2tensor/trax/layers/__init__.py
+++ b/tensor2tensor/trax/layers/__init__.py
@@ -25,3 +25,4 @@
 from tensor2tensor.trax.layers.base import *
 from tensor2tensor.trax.layers.combinators import *
 from tensor2tensor.trax.layers.core import *
+from tensor2tensor.trax.layers.rnn import *
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 67d9458f2..a291133c8 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import operator
 from tensor2tensor.trax import backend
 from tensor2tensor.trax.layers import base
 
@@ -130,26 +131,61 @@ def SecondBranch(x, **unused_kwargs):
   return x[1]  # Here x is a list of tensors, we select the second.
 
 
-def _nested_sum(inputs):  # pylint: disable=invalid-name
+def _nested_op(inputs, op):  # pylint: disable=invalid-name
   """Helper: sum a list of arrays or nested arrays."""
   # First the simple non-nested case.
   if not isinstance(inputs[0], (list, tuple)):
-    return sum(inputs)
+    return op(inputs)
   # In the nested case, sum on each axis separately.
   result_list = []
   for i in range(len(inputs[0])):
-    result_list.append(_nested_sum([x[i] for x in inputs]))
+    result_list.append(_nested_op([x[i] for x in inputs], op=op))
   if isinstance(inputs[0], list):
     return result_list
   return tuple(result_list)
 
 
+def _nested_sum(inputs):  # pylint: disable=invalid-name
+  return _nested_op(inputs=inputs, op=sum)
+
+
+def _nested_product(inputs):  # pylint: disable=invalid-name
+  return _nested_op(inputs=inputs, op=lambda xs: reduce(operator.mul, xs))
+
+
 @base.layer(output_shape=lambda input_shape_list: input_shape_list[0])
 def SumBranches(x, **unused_kwargs):
+  """Sum branches elementwise."""
   # Here x is a list of tensors of the same shape, or nested structures.
   return _nested_sum(x)
 
 
+@base.layer(output_shape=lambda input_shape_list: input_shape_list[0])
+def MultiplyBranches(x, **unused_kwargs):
+  """Multiply branches elementwise."""
+  return _nested_product(x)
+
+
+@base.layer(output_shape=lambda input_shape_list: input_shape_list[0])
+def GateBranches(x, **unused_kwargs):
+  """Implements a gating function on a (memory, gate, candidate) tuple.
+
+  Final update is memory * gate + (1-gate) * candidate
+
+  This gating equation may also be referred to as Highway Network.
+  Highway Networks: https://arxiv.org/abs/1505.00387
+
+  Args:
+    x: A tuple of (memory, gate, candidate)
+
+  Returns:
+    The result of applying gating.
+  """
+  assert len(x) == 3, x
+  state, gate, candidate = x
+  return gate * state + (1.0 - gate) * candidate
+
+
 def _concatenate_shape(input_shape, axis=-1):  # pylint: disable=invalid-name
   """Helper to determine the shape of Concatenate output."""
   ax = axis % len(input_shape[0])
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 31999fdce..3e6307223 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -74,16 +74,39 @@ def one_hot(x, size, dtype=np.float32):
 # Layers.
 
 
+@base.layer()
+def AddConstant(x, params, constant=0.0, **unused_kwargs):
+  del params
+  return x + constant
+
+
 @base.layer()
 def Relu(x, **unused_kwargs):
   return np.maximum(x, 0.)
 
 
+@base.layer()
+def Sigmoid(x, **unused_kwargs):
+  return 1. / (1. + np.exp(-x))
+
+
 @base.layer()
 def Tanh(x, **unused_kwargs):
   return np.tanh(x)
 
 
+@base.layer()
+def HardSigmoid(x, **unused_kwargs):
+  """Linear approximation to sigmoid."""
+  return np.maximum(0, np.minimum(1, (1 + x)))
+
+
+@base.layer()
+def HardTanh(x, **unused_kwargs):
+  """Linear approximation to tanh."""
+  return np.maximum(-1, np.minimum(1, x))
+
+
 @base.layer()
 def Exp(x, **unused_kwargs):
   return np.exp(x)
diff --git a/tensor2tensor/trax/layers/rnn.py b/tensor2tensor/trax/layers/rnn.py
new file mode 100644
index 000000000..91b3b69cd
--- /dev/null
+++ b/tensor2tensor/trax/layers/rnn.py
@@ -0,0 +1,140 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Implementations of common recurrent neural network cells (RNNs)."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import google_type_annotations
+from __future__ import print_function
+
+import google3
+from tensor2tensor.trax.layers import combinators
+from tensor2tensor.trax.layers import core
+
+
+def GRUCell(units):
+  """Builds a traditional GRU cell with dense internal transformations.
+
+  Gated Recurrent Unit paper: https://arxiv.org/abs/1412.3555
+
+
+  Args:
+    units: Number of hidden units.
+
+  Returns:
+    A Stax model representing a traditional GRU RNN cell.
+  """
+  return GeneralGRUCell(
+      candidate_transform=lambda: core.Dense(units=units),
+      memory_transform=combinators.Identity,
+      gate_nonlinearity=core.Sigmoid,
+      candidate_nonlinearity=core.Tanh)
+
+
+def ConvGRUCell(units, kernel_size=(3, 3)):
+  """Builds a convolutional GRU.
+
+  Paper: https://arxiv.org/abs/1511.06432.
+
+  Args:
+    units: Number of hidden units
+    kernel_size: Kernel size for convolution
+
+  Returns:
+    A Stax model representing a GRU cell with convolution transforms.
+  """
+
+  def BuildConv():
+    return core.Conv(filters=units, kernel_size=kernel_size, padding='SAME')
+
+  return GeneralGRUCell(
+      candidate_transform=BuildConv,
+      memory_transform=combinators.Identity,
+      gate_nonlinearity=core.Sigmoid,
+      candidate_nonlinearity=core.Tanh)
+
+
+def GeneralGRUCell(candidate_transform,
+                   memory_transform=combinators.Identity,
+                   gate_nonlinearity=core.Sigmoid,
+                   candidate_nonlinearity=core.Tanh,
+                   dropout_rate_c=0.1,
+                   sigmoid_bias=0.5):
+  r"""Parametrized Gated Recurrent Unit (GRU) cell construction.
+
+  GRU update equations:
+  $$ Update gate: u_t = \sigmoid(U' * s_{t-1} + B') $$
+  $$ Reset gate: r_t = \sigmoid(U'' * s_{t-1} + B'') $$
+  $$ Candidate memory: c_t = \tanh(U * (r_t \odot s_{t-1}) + B) $$
+  $$ New State: s_t = u_t \odot s_{t-1} + (1 - u_t) \odot c_t $$
+
+  See combinators.GateBranches for details on the gating function.
+
+
+  Args:
+    candidate_transform: Transform to apply inside the Candidate branch. Applied
+      before nonlinearities.
+    memory_transform: Optional transformation on the memory before gating.
+    gate_nonlinearity: Function to use as gate activation. Allows trying
+      alternatives to Sigmoid, such as HardSigmoid.
+    candidate_nonlinearity: Nonlinearity to apply after candidate branch. Allows
+      trying alternatives to traditional Tanh, such as HardTanh
+    dropout_rate_c: Amount of dropout on the transform (c) gate. Dropout works
+      best in a GRU when applied exclusively to this branch.
+    sigmoid_bias: Constant to add before sigmoid gates. Generally want to start
+      off with a positive bias.
+
+  Returns:
+    A model representing a GRU cell with specified transforms.
+  """
+  return combinators.Serial(
+      combinators.Branch(num_branches=3),
+      combinators.Parallel(
+          # s_{t-1} branch - optionally transform
+          # Typically is an identity.
+          memory_transform(),
+
+          # u_t (Update gate) branch
+          combinators.Serial(
+              candidate_transform(),
+              # Want bias to start out positive before sigmoids.
+              core.AddConstant(constant=sigmoid_bias),
+              gate_nonlinearity()),
+
+          # c_t (Candidate) branch
+          combinators.Serial(
+              combinators.Branch(num_branches=2),
+              combinators.Parallel(
+                  combinators.Identity(),
+                  # r_t (Reset) Branch
+                  combinators.Serial(
+                      candidate_transform(),
+                      # Want bias to start out positive before sigmoids.
+                      core.AddConstant(constant=sigmoid_bias),
+                      gate_nonlinearity())),
+              ## Gate S{t-1} with sigmoid(candidate_transform(S{t-1}))
+              combinators.MultiplyBranches(),
+
+              # Final projection + tanh to get Ct
+              candidate_transform(),
+              candidate_nonlinearity()),  # Candidate gate
+
+          # Only apply dropout on the C gate.
+          # Paper reports that 0.1 is a good default.
+          core.Dropout(rate=dropout_rate_c)),
+
+      # Gate memory and candidate
+      combinators.GateBranches())
diff --git a/tensor2tensor/trax/layers/rnn_test.py b/tensor2tensor/trax/layers/rnn_test.py
new file mode 100644
index 000000000..7a7ac23c7
--- /dev/null
+++ b/tensor2tensor/trax/layers/rnn_test.py
@@ -0,0 +1,54 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 google3.third_party.py.tensor2tensor.trax.layers.rnn."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.trax.backend import random as jax_random
+from tensor2tensor.trax.layers import rnn
+from google3.testing.pybase import googletest
+
+
+class RnnModelTest(googletest.TestCase):
+
+  def _test_cell_runs(self, model, input_shape, output_shape):
+    source = np.ones(input_shape, dtype=np.float32)
+
+    # Build params
+    rng = jax_random.get_prng(0)
+    model.initialize(input_shape, rng)
+
+    # Run network
+    output = model(source)
+
+    self.assertEqual(output_shape, output.shape)
+
+  def test_conv_gru_cell(self):
+    self._test_cell_runs(
+        rnn.ConvGRUCell(units=9, kernel_size=(3, 3)),
+        input_shape=(8, 1, 7, 9),
+        output_shape=(8, 1, 7, 9))
+
+  def test_gru_cell(self):
+    self._test_cell_runs(
+        rnn.GRUCell(units=9), input_shape=(8, 7, 9), output_shape=(8, 7, 9))
+
+
+if __name__ == '__main__':
+  googletest.main()
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index 8dd46c0de..8d2c8ef59 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -21,6 +21,7 @@
 import gin
 
 from tensor2tensor.trax.models import mlp
+from tensor2tensor.trax.models import neural_gpu
 from tensor2tensor.trax.models import resnet
 from tensor2tensor.trax.models import transformer
 
@@ -32,9 +33,10 @@ def model_configure(*args, **kwargs):
 
 
 # pylint: disable=invalid-name
+ChunkedTransformerLM = model_configure(transformer.ChunkedTransformerLM)
 MLP = model_configure(mlp.MLP)
+NeuralGPU = model_configure(neural_gpu.NeuralGPU)
 Resnet50 = model_configure(resnet.Resnet50)
-WideResnet = model_configure(resnet.WideResnet)
 TransformerEncoder = model_configure(transformer.TransformerEncoder)
 TransformerLM = model_configure(transformer.TransformerLM)
-ChunkedTransformerLM = model_configure(transformer.ChunkedTransformerLM)
+WideResnet = model_configure(resnet.WideResnet)
diff --git a/tensor2tensor/trax/models/neural_gpu.py b/tensor2tensor/trax/models/neural_gpu.py
new file mode 100644
index 000000000..46e9cc9fa
--- /dev/null
+++ b/tensor2tensor/trax/models/neural_gpu.py
@@ -0,0 +1,82 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Implementation of the improved Neural GPU (NGPU)."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import google_type_annotations
+from __future__ import print_function
+
+from tensor2tensor.trax import layers
+from tensor2tensor.trax.backend import numpy as np
+
+
+# TODO(ddohan): Combinator to add saturation costs to loss
+def SaturationCost(x, limit=0.9):
+  return np.minimum(0, np.abs(x) - limit)
+
+
+@layers.layer(output_shape=lambda input_shape_list: input_shape_list)
+def DiagonalGate(x, params, **kwargs):
+  """Split channels in 3 parts. Shifts 1st and 3rd sections to left/right."""
+  del params
+  del kwargs
+  # x : [batch, 1, length, depth]
+  x = np.pad(
+      x, [(0, 0), (0, 0), (1, 1), (0, 0)], mode='constant', constant_values=0.0)
+  depth = x.shape[-1] // 3
+  assert 3 * depth == x.shape[-1], ('Depth must be divisible by 3', depth,
+                                    x.shape)
+  xs = [
+      x[:, :, :-2, :depth], x[:, :, 1:-1, depth:2 * depth],
+      x[:, :, 2:, 2 * depth:3 * depth]
+  ]
+  return np.concatenate(xs, axis=3)
+
+
+def ConvDiagonalGRU(units, kernel_size=(3, 3)):
+  """Build convolutional GRU with diagonal gating as in ImprovedNGPU."""
+
+  def BuildConv():
+    return layers.Conv(filters=units, kernel_size=kernel_size, padding='SAME')
+
+  return layers.GeneralGRUCell(
+      candidate_transform=BuildConv,
+      memory_transform=DiagonalGate,
+      gate_nonlinearity=layers.HardSigmoid,
+      candidate_nonlinearity=layers.HardTanh)
+
+
+def NeuralGPU(feature_depth=96, steps=16, vocab_size=2):
+  """Implementation of Neural GPU: https://arxiv.org/abs/1702.08727.
+
+  Args:
+    feature_depth: Number of memory channels
+    steps: Number of times depthwise recurrence steps.
+    vocab_size: Vocabulary size.
+
+  Returns:
+    A NeuralGPU Stax model.
+  """
+  xs = []
+  xs.append(
+      layers.Embedding(feature_depth=feature_depth, vocab_size=vocab_size))
+  core = ConvDiagonalGRU(units=feature_depth)
+  xs.extend([core] * steps)
+  xs.append(layers.Dense(vocab_size))
+  xs.append(layers.LogSoftmax())
+
+  return layers.Serial(*xs)
diff --git a/tensor2tensor/trax/models/neural_gpu_test.py b/tensor2tensor/trax/models/neural_gpu_test.py
new file mode 100644
index 000000000..8b5ed38b9
--- /dev/null
+++ b/tensor2tensor/trax/models/neural_gpu_test.py
@@ -0,0 +1,50 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 google3.third_party.py.tensor2tensor.trax.models.neural_gpu."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl import logging
+import numpy as np
+from tensor2tensor.trax.backend import random as jax_random
+from tensor2tensor.trax.models import neural_gpu
+from google3.testing.pybase import googletest
+
+
+class NeuralGPUTest(googletest.TestCase):
+
+  def test_ngpu(self):
+    vocab_size = 2
+    in_shape = [3, 5, 7]
+    source = np.ones(in_shape, dtype=np.int32)
+
+    model = neural_gpu.NeuralGPU(
+        feature_depth=30, steps=4, vocab_size=vocab_size)
+    # Build params
+    rng = jax_random.get_prng(0)
+    logging.info(model)
+    model.initialize(in_shape, rng)
+
+    # Run network
+    output = model(source)
+
+    self.assertEqual(tuple(in_shape + [vocab_size]), output.shape)
+
+
+if __name__ == '__main__':
+  googletest.main()

From c4b4e7f728ef5ac51d87ca12657bf66deaed863f Mon Sep 17 00:00:00 2001
From: Peng Wang <wangpeng@google.com>
Date: Mon, 29 Apr 2019 16:32:03 -0700
Subject: [PATCH 1971/2720] Internal changes.

PiperOrigin-RevId: 245847810
---
 tensor2tensor/trax/backend.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index 872352c0f..17526c473 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -27,8 +27,6 @@
 import jax.scipy.special as jax_special
 import numpy as onp
 
-import tensorflow as tf
-
 
 
 _JAX_BACKEND = {
@@ -112,4 +110,3 @@ def backend(name="jax"):
   if name == "numpy":
     return _NUMPY_BACKEND
   return _JAX_BACKEND
-

From 539acda4f6501e83f389e86d7cb16d4f919d7e48 Mon Sep 17 00:00:00 2001
From: David Dohan <ddohan@google.com>
Date: Mon, 29 Apr 2019 17:24:47 -0700
Subject: [PATCH 1972/2720] Use tf.test.TestCase in tests

PiperOrigin-RevId: 245856606
---
 tensor2tensor/trax/layers/rnn_test.py        | 6 +++---
 tensor2tensor/trax/models/neural_gpu_test.py | 8 +++-----
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/trax/layers/rnn_test.py b/tensor2tensor/trax/layers/rnn_test.py
index 7a7ac23c7..b70934742 100644
--- a/tensor2tensor/trax/layers/rnn_test.py
+++ b/tensor2tensor/trax/layers/rnn_test.py
@@ -22,10 +22,10 @@
 import numpy as np
 from tensor2tensor.trax.backend import random as jax_random
 from tensor2tensor.trax.layers import rnn
-from google3.testing.pybase import googletest
+import tensorflow as tf
 
 
-class RnnModelTest(googletest.TestCase):
+class RnnModelTest(tf.test.TestCase):
 
   def _test_cell_runs(self, model, input_shape, output_shape):
     source = np.ones(input_shape, dtype=np.float32)
@@ -51,4 +51,4 @@ def test_gru_cell(self):
 
 
 if __name__ == '__main__':
-  googletest.main()
+  tf.test.main()
diff --git a/tensor2tensor/trax/models/neural_gpu_test.py b/tensor2tensor/trax/models/neural_gpu_test.py
index 8b5ed38b9..4f09a71c6 100644
--- a/tensor2tensor/trax/models/neural_gpu_test.py
+++ b/tensor2tensor/trax/models/neural_gpu_test.py
@@ -19,14 +19,13 @@
 from __future__ import division
 from __future__ import print_function
 
-from absl import logging
 import numpy as np
 from tensor2tensor.trax.backend import random as jax_random
 from tensor2tensor.trax.models import neural_gpu
-from google3.testing.pybase import googletest
+import tensorflow as tf
 
 
-class NeuralGPUTest(googletest.TestCase):
+class NeuralGPUTest(tf.test.TestCase):
 
   def test_ngpu(self):
     vocab_size = 2
@@ -37,7 +36,6 @@ def test_ngpu(self):
         feature_depth=30, steps=4, vocab_size=vocab_size)
     # Build params
     rng = jax_random.get_prng(0)
-    logging.info(model)
     model.initialize(in_shape, rng)
 
     # Run network
@@ -47,4 +45,4 @@ def test_ngpu(self):
 
 
 if __name__ == '__main__':
-  googletest.main()
+  tf.test.main()

From 6c3b5d4100f9175784a10e9637631ba4d12271a9 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 29 Apr 2019 18:38:47 -0700
Subject: [PATCH 1973/2720] Minor inheritance clean-up among En-De translate
 problems.

PiperOrigin-RevId: 245866222
---
 .../data_generators/translate_ende.py         | 26 ++++----
 .../data_generators/translate_ende_test.py    | 61 +++++++++++++++++++
 2 files changed, 72 insertions(+), 15 deletions(-)
 create mode 100644 tensor2tensor/data_generators/translate_ende_test.py

diff --git a/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index 37607d6b5..88b8ded8e 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -56,12 +56,8 @@
 
 
 @registry.register_problem
-class TranslateEndeWmt8k(translate.TranslateProblem):
-  """Problem spec for WMT En-De translation."""
-
-  @property
-  def approx_vocab_size(self):
-    return 2**13  # 8192
+class TranslateEndeWmt32k(translate.TranslateProblem):
+  """En-de translation trained on WMT corpus."""
 
   @property
   def additional_training_datasets(self):
@@ -74,15 +70,6 @@ def source_data_files(self, dataset_split):
     return train_datasets if train else _ENDE_EVAL_DATASETS
 
 
-@registry.register_problem
-class TranslateEndeWmt32k(TranslateEndeWmt8k):
-  """En-de translation trained on WMT corpus."""
-
-  @property
-  def approx_vocab_size(self):
-    return 2**15  # 32768
-
-
 @registry.register_problem
 class TranslateEndeWmtClean32k(TranslateEndeWmt32k):
   """En-de translation trained on WMT with further cleaning."""
@@ -177,6 +164,15 @@ def use_vocab_from_other_problem(self):
     return TranslateEndeWmt32k()
 
 
+@registry.register_problem
+class TranslateEndeWmt8k(TranslateEndeWmt32k):
+  """Problem spec for WMT En-De translation."""
+
+  @property
+  def approx_vocab_size(self):
+    return 2**13  # 8192
+
+
 @registry.register_problem
 class TranslateEndeWmt8kPacked(TranslateEndeWmt8k):
 
diff --git a/tensor2tensor/data_generators/translate_ende_test.py b/tensor2tensor/data_generators/translate_ende_test.py
new file mode 100644
index 000000000..b8b5e4550
--- /dev/null
+++ b/tensor2tensor/data_generators/translate_ende_test.py
@@ -0,0 +1,61 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.data_generators.translate_ende."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import translate_ende
+
+import tensorflow as tf
+
+
+class TranslateEndeTest(tf.test.TestCase):
+  """Tests that some TranslateEnde subclasses inherit information correctly."""
+
+  def test_vocab_size(self):
+    wmt_8k = translate_ende.TranslateEndeWmt8k()
+    wmt_32k = translate_ende.TranslateEndeWmt32k()
+    self.assertEqual(wmt_8k.approx_vocab_size, 8192)
+    self.assertEqual(wmt_32k.approx_vocab_size, 32768)
+
+  def test_additional_datasets(self):
+    wmt_8k = translate_ende.TranslateEndeWmt8k()
+    wmt_32k = translate_ende.TranslateEndeWmt32k()
+    self.assertListEqual(wmt_8k.additional_training_datasets, [])
+    self.assertListEqual(wmt_32k.additional_training_datasets, [])
+
+  def test_source_data_files(self):
+    wmt_8k = translate_ende.TranslateEndeWmt8k()
+    wmt_32k = translate_ende.TranslateEndeWmt32k()
+    eval_split = problem.DatasetSplit.EVAL
+    train_split = problem.DatasetSplit.TRAIN
+
+    wmt_8k_eval_files = wmt_8k.source_data_files(eval_split)
+    wmt_32k_eval_files = wmt_32k.source_data_files(eval_split)
+    self.assertListEqual(wmt_8k_eval_files, wmt_32k_eval_files)
+    self.assertGreater(len(wmt_8k_eval_files), 0)
+
+    wmt_8k_train_files = wmt_8k.source_data_files(train_split)
+    wmt_32k_train_files = wmt_32k.source_data_files(train_split)
+    self.assertListEqual(wmt_8k_train_files, wmt_32k_train_files)
+    self.assertGreater(len(wmt_8k_train_files), 0)
+
+
+if __name__ == '__main__':
+  tf.test.main()

From 272500b6efe353aeb638d2745ed56e519462ca31 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 29 Apr 2019 18:50:46 -0700
Subject: [PATCH 1974/2720] Adding capability in Trajectory and BatchTrajectory
 to pad observations. This padded observation array can be used as an input to
 a policy net.

PiperOrigin-RevId: 245867571
---
 tensor2tensor/envs/trajectory.py      |  40 ++++++++++
 tensor2tensor/envs/trajectory_test.py | 103 ++++++++++++++++++++++++--
 2 files changed, 138 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index 557f2db3d..29677b4e4 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -93,6 +93,10 @@ def reward(self):
         processed_rewards += ts.processed_reward
     return raw_rewards, processed_rewards
 
+  @property
+  def observations_np(self):
+    return np.stack([ts.observation for ts in self.time_steps])
+
 
 class BatchTrajectory(object):
   """Basically a batch of active trajectories and a list of completed ones."""
@@ -273,3 +277,39 @@ def num_time_steps(self):
 
     num_time_steps = sum(t.num_time_steps for t in self.trajectories)
     return num_time_steps + self.num_completed_time_steps
+
+  @property
+  def num_completed_trajectories(self):
+    """Returns the number of completed trajectories."""
+    return len(self.completed_trajectories)
+
+  def observations_np(self, boundary=20):
+    """Pads the observations in all the trajectories and returns them.
+
+    Args:
+      boundary: integer, Observations will be padded to (n * boundary) + 1 where
+          n is an integer.
+
+    Returns:
+      a tuple(padded_observations, time_steps), with shapes:
+      padded_observations: (self.batch_size, n * boundary + 1) + OBS
+      time_steps: integer list of length = self.batch_size
+    """
+    list_observations_np_ts = [t.observations_np for t in self.trajectories]
+    # Every element in `list_observations_np_ts` is shaped (t,) + OBS
+    OBS = list_observations_np_ts[0].shape[1:]  # pylint: disable=invalid-name
+
+    num_time_steps = [t.num_time_steps for t in self.trajectories]
+    t_max = max(num_time_steps)
+    # t_max is rounded to the next multiple of `boundary`
+    boundary = int(boundary)
+    bucket_length = boundary * int(np.ceil(float(t_max) / boundary))
+
+    def padding_config(obs):
+      # We're padding the first axis only, since that is the time-step.
+      num_to_pad = bucket_length + 1 - obs.shape[0]
+      return [(0, num_to_pad)] + [(0, 0)] * len(OBS)
+
+    return np.stack([
+        np.pad(obs, padding_config(obs), "constant")
+        for obs in list_observations_np_ts]), num_time_steps
diff --git a/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index 92105c220..acc50c1bb 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -90,6 +90,15 @@ def test_reward(self):
     self.assertEqual(5, raw_reward)
     self.assertEqual(500, processed_reward)
 
+  def test_observation_np(self):
+    t = trajectory.Trajectory()
+    ts = 5
+    shape = (3, 4)
+    for _ in range(ts):
+      t.add_time_step(observation=np.random.uniform(size=shape), done=False)
+
+    self.assertEqual((ts,) + shape, t.observations_np.shape)
+
 
 class BatchTrajectoryTest(tf.test.TestCase):
 
@@ -111,7 +120,7 @@ def test_creation(self):
     bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
 
     self.assertEqual(self.BATCH_SIZE, len(bt.trajectories))
-    self.assertEqual(0, len(bt.completed_trajectories))
+    self.assertEqual(0, bt.num_completed_trajectories)
 
   def test_reset_all(self):
     bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
@@ -125,7 +134,7 @@ def test_reset_all(self):
     # Assert that all trajectories are active and not done (reset never marks
     # anything as done).
     self.assertTrue(all(t.is_active for t in bt.trajectories))
-    self.assertEqual(0, len(bt.completed_trajectories))
+    self.assertEqual(0, bt.num_completed_trajectories)
 
   def test_num_time_steps(self):
     bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
@@ -150,7 +159,7 @@ def test_reset_some(self):
         all(not t.is_active for t in bt.trajectories[self.BATCH_SIZE // 2:]))
 
     # Nothing is done anyways.
-    self.assertEqual(0, len(bt.completed_trajectories))
+    self.assertEqual(0, bt.num_completed_trajectories)
 
   def test_step(self):
     bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
@@ -179,7 +188,7 @@ def test_step(self):
     bt.step(new_observations, raw_rewards, processed_rewards, dones, actions)
 
     # Expect to see `num_done` number of completed trajectories.
-    self.assertEqual(num_done, len(bt.completed_trajectories))
+    self.assertEqual(num_done, bt.num_completed_trajectories)
 
     # Expect to see that the rest are marked as active.
     num_active = sum(t.is_active for t in bt.trajectories)
@@ -207,7 +216,7 @@ def test_desired_placement_of_rewards_and_actions(self):
     bt.step(new_observations, raw_rewards, processed_rewards, dones, actions)
 
     # Assert that nothing is done, since dones is False
-    self.assertEqual(0, len(bt.completed_trajectories))
+    self.assertEqual(0, bt.num_completed_trajectories)
 
     # The only trajectory is active.
     self.assertEqual(batch_size, len(bt.trajectories))
@@ -238,6 +247,90 @@ def test_desired_placement_of_rewards_and_actions(self):
     self.assertEqual(processed_rewards[0], ts[1].processed_reward)
     self.assertIsNone(ts[0].processed_reward)
 
+  def test_observations_np(self):
+    bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
+    indices = np.arange(self.BATCH_SIZE)
+    observations, _, _, _ = self.get_random_observations_rewards_actions_dones()
+
+    # Have to call reset first.
+    bt.reset(indices, observations)
+
+    # Number of time-steps now looks like the following:
+    # (1, 1, 1, 1, 1, 1, 1, 1, 1, 1)
+    lengths = np.full((self.BATCH_SIZE,), 1)
+
+    ts = 5
+    for _ in range(ts):
+      (observations, rewards, actions, dones
+      ) = self.get_random_observations_rewards_actions_dones()
+      dones[...] = False
+      bt.step(observations, rewards, rewards, dones, actions)
+
+    # Number of time-steps now looks like the following:
+    # (6, 6, 6, 6, 6, 6, 6, 6, 6, 6)
+    lengths = lengths + ts
+
+    # Now let's mark the first two as done.
+    observations, _, _, _ = self.get_random_observations_rewards_actions_dones(
+        batch_size=2)
+    bt.reset(np.array([0, 1]), observations)
+
+    # Number of time-steps now looks like the following:
+    # (1, 1, 6, 6, 6, 6, 6, 6, 6, 6)
+    lengths[0] = lengths[1] = 1
+
+    for _ in range(ts):
+      (observations, rewards, actions, dones
+      ) = self.get_random_observations_rewards_actions_dones()
+      dones[...] = False
+      bt.step(observations, rewards, rewards, dones, actions)
+
+    # Number of time-steps now looks like the following:
+    # (6, 6, 11, 11, 11, 11, 11, 11, 11, 11)
+    lengths = lengths + ts
+
+    boundary = 20
+    padded_obs_np, padded_lengths = bt.observations_np(boundary=boundary)
+
+    # The lengths are what we expect them to be.
+    self.assertAllEqual(lengths, padded_lengths)
+
+    # The padded_observations are the shape we expect them to be.
+    self.assertEqual((self.BATCH_SIZE, boundary + 1) + self.OBSERVATION_SHAPE,
+                     padded_obs_np.shape)
+
+    # Let's do 10 more steps (to go on the other side of the boundary.
+    ts = 10
+    for _ in range(ts):
+      (observations, rewards, actions, dones
+      ) = self.get_random_observations_rewards_actions_dones()
+      dones[...] = False
+      bt.step(observations, rewards, rewards, dones, actions)
+
+    # Number of time-steps now looks like the following:
+    # (16, 16, 21, 21, 21, 21, 21, 21, 21, 21)
+    lengths = lengths + ts
+
+    padded_obs_np, padded_lengths = bt.observations_np(boundary=boundary)
+
+    # The lengths are what we expect them to be.
+    self.assertAllEqual(lengths, padded_lengths)
+
+    # The padded_observations are the shape we expect them to be.
+    self.assertEqual((self.BATCH_SIZE,
+                      (2 * boundary) + 1) + self.OBSERVATION_SHAPE,
+                     padded_obs_np.shape)
+
+    # Test that the padding is the only part that is all 0s.
+    # NOTE: There is almost 0 probability that the random observation is all 0s.
+    zero_obs = np.full(self.OBSERVATION_SHAPE, 0.)
+    for b in range(self.BATCH_SIZE):
+      # The first lengths[b] will be actual data, rest is 0s.
+      for ts in range(lengths[b]):
+        self.assertFalse(np.all(zero_obs == padded_obs_np[b][ts]))
+
+      for ts in range(lengths[b], len(padded_obs_np[b])):
+        self.assertAllEqual(zero_obs, padded_obs_np[b][ts])
 
 if __name__ == '__main__':
   tf.test.main()

From c126e6969a493aff6a7889018833884ab2d011d4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 29 Apr 2019 23:04:07 -0700
Subject: [PATCH 1975/2720] Two bug fixes.

PiperOrigin-RevId: 245889159
---
 tensor2tensor/trax/trax.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 7908aabba..ea60997a4 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -29,6 +29,7 @@
 import time
 
 from absl import logging
+import cloudpickle
 
 import gin
 
@@ -151,13 +152,20 @@ def save_gin(output_dir, sw=None):
 
 def save_state(state, output_dir, keep=False):
   """Save State and optionally gin config."""
+  # TODO(gilmer, lukaszkaiser): figure out how to use cloudpickle in python3.
+  # Currently the code throws an error when run in python3.
+  if sys.version_info[0] < 3:
+    pkl_module = cloudpickle
+  else:
+    pkl_module = pickle
   params_file = os.path.join(output_dir, "model.pkl")
+  params = jax.unreplicate(state.params)
   with gfile.GFile(params_file, "wb") as f:
-    pickle.dump((state.params, state.step, state.history), f)
+    pkl_module.dump((params, state.step, state.history), f)
   if keep:
     params_file = os.path.join(output_dir, "model_{}.pkl".format(state.step))
     with gfile.GFile(params_file, "wb") as f:
-      pickle.dump((state.params, state.step, state.history), f)
+      pkl_module.dump((params, state.step, state.history), f)
   log("Model saved to %s" % params_file, stdout=False)
 
 
@@ -466,6 +474,7 @@ def train(output_dir,
       step += 1
 
       if step in save_steps:
+        params = trax_opt.get_params(opt_state)
         save_state(State(params=params, step=step, history=history),
                    output_dir,
                    keep=True)

From cfb7afacd172e98f2e029600cfb56d82f13514dc Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Tue, 30 Apr 2019 04:18:37 -0700
Subject: [PATCH 1976/2720] Add Dense layer for deterministic variational
 inference.

References
+ https://github.com/Microsoft/deterministic-variational-inference

PiperOrigin-RevId: 245922349
---
 tensor2tensor/layers/bayes.py      | 167 +++++++++++++++++++++++++++++
 tensor2tensor/layers/bayes_test.py |  76 ++++++++++++-
 2 files changed, 241 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 2752baa00..18c50560d 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import functools
+import math
 from tensor2tensor.keras import constraints
 from tensor2tensor.keras import initializers
 from tensor2tensor.keras import regularizers
@@ -225,6 +226,172 @@ def dropped_inputs():
                      lambda: super(Conv2DVariationalDropout, self).call(inputs))
 
 
+@add_weight
+class DenseDVI(tf.keras.layers.Dense):
+  """Densely-connected layer with deterministic VI (Wu et al., 2018).
+
+  This layer computes a variational inference approximation via first and second
+  moments. It is accurate if the kernel and bias initializers return factorized
+  normal random variables and the number of units is sufficiently large. The
+  advantage is that the forward pass is deterministic, reducing variance of
+  gradients during training. The disadvantage is an O(features^2*units) compute
+  and O(features^2 + features*units) memory complexity. In comparison,
+  DenseReparameterization has O(features*units) compute and memory complexity.
+
+  #### Examples
+
+  Below implements deterministic variational inference for Bayesian
+  feedforward network regression. We use the exact expected log-likelihood from
+  Wu et al. (2018), Eq. 8. Assume 2-D real-valued tensors of `features` and
+  `labels` of shapes `[batch_size, num_features]` and `[batch_size, 1]`
+  respectively.
+
+  ```python
+  from tensor2tensor.layers import bayes
+
+  model = tf.keras.Sequential([
+      bayes.DenseDVI(256, activation=tf.nn.relu),
+      bayes.DenseDVI(256, activation=tf.nn.relu),
+      bayes.DenseDVI(1, activation=None),
+  ])
+  locs = model(features)
+  nll = 0.5 * tf.reduce_mean(locs.distribution.variance() +
+                             (labels - locs.distribution.mean())**2)
+  kl = sum(model.losses) / total_dataset_size
+  loss = nll + kl
+  train_op = tf.train.AdamOptimizer(0.1).minimize(loss)
+  ```
+
+  For evaluation, feed in data and use, e.g., `predictions.distribution.mean()`
+  to make predictions via the posterior predictive distribution.
+
+  ```python
+  predictions = ed.Normal(loc=locs.distribution.mean(),
+                          scale=locs.distribution.variance() + 1.)
+  ```
+  """
+
+  def __init__(self,
+               units,
+               activation=None,
+               use_bias=True,
+               kernel_initializer='trainable_normal',
+               bias_initializer='zero',
+               kernel_regularizer='normal_kl_divergence',
+               bias_regularizer=None,
+               activity_regularizer=None,
+               **kwargs):
+    super(DenseDVI, self).__init__(
+        units=units,
+        activation=activation,
+        use_bias=use_bias,
+        kernel_initializer=initializers.get(kernel_initializer),
+        bias_initializer=initializers.get(bias_initializer),
+        kernel_regularizer=regularizers.get(kernel_regularizer),
+        bias_regularizer=regularizers.get(bias_regularizer),
+        activity_regularizer=regularizers.get(activity_regularizer),
+        **kwargs)
+
+  def call_weights(self):
+    """Calls any weights if the initializer is itself a layer."""
+    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
+      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
+    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
+      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
+
+  def call(self, inputs):
+    self.call_weights()
+    if (not isinstance(inputs, ed.RandomVariable) and
+        not isinstance(self.kernel, ed.RandomVariable) and
+        not isinstance(self.bias, ed.RandomVariable)):
+      return super(DenseDVI, self).call(inputs)
+    inputs_mean, inputs_variance, inputs_covariance = get_moments(inputs)
+    kernel_mean, kernel_variance, _ = get_moments(self.kernel)
+    if self.use_bias:
+      bias_mean, _, bias_covariance = get_moments(self.bias)
+
+    # E[outputs] = E[inputs] * E[kernel] + E[bias]
+    mean = tf.tensordot(inputs_mean, kernel_mean, [[-1], [0]])
+    if self.use_bias:
+      mean = tf.nn.bias_add(mean, bias_mean)
+
+    # Cov = E[inputs**2] Cov(kernel) + E[W]^T Cov(inputs) E[W] + Cov(bias)
+    # For first term, assume Cov(kernel) = 0 on off-diagonals so we only
+    # compute diagonal term.
+    covariance_diag = tf.tensordot(inputs_variance + inputs_mean**2,
+                                   kernel_variance, [[-1], [0]])
+    # Compute quadratic form E[W]^T Cov E[W] from right-to-left. First is
+    #  [..., features, features], [features, units] -> [..., features, units].
+    cov_w = tf.tensordot(inputs_covariance, kernel_mean, [[-1], [0]])
+    # Next is [..., features, units], [features, units] -> [..., units, units].
+    w_cov_w = tf.tensordot(cov_w, kernel_mean, [[-2], [0]])
+    covariance = w_cov_w
+    if self.use_bias:
+      covariance += bias_covariance
+    covariance = tf.matrix_set_diag(
+        covariance, tf.matrix_diag_part(covariance) + covariance_diag)
+
+    if self.activation in (tf.keras.activations.relu, tf.nn.relu):
+      # Compute activation's moments with variable names from Wu et al. (2018).
+      variance = tf.matrix_diag_part(covariance)
+      scale = tf.sqrt(variance)
+      mu = mean / (scale + tf.keras.backend.epsilon())
+      mean = scale * soft_relu(mu)
+
+      pairwise_variances = (tf.expand_dims(variance, -1) *
+                            tf.expand_dims(variance, -2))  # [..., units, units]
+      rho = covariance / tf.sqrt(pairwise_variances +
+                                 tf.keras.backend.epsilon())
+      rho = tf.clip_by_value(rho,
+                             -1. / (1. + tf.keras.backend.epsilon()),
+                             1. / (1. + tf.keras.backend.epsilon()))
+      s = covariance / (rho + tf.keras.backend.epsilon())
+      mu1 = tf.expand_dims(mu, -1)  # [..., units, 1]
+      mu2 = tf.matrix_transpose(mu1)  # [..., 1, units]
+      a = (soft_relu(mu1) * soft_relu(mu2) +
+           rho * tfp.distributions.Normal(0., 1.).cdf(mu1) *
+           tfp.distributions.Normal(0., 1.).cdf(mu2))
+      gh = tf.asinh(rho)
+      bar_rho = tf.sqrt(1. - rho**2)
+      gr = gh + rho / (1. + bar_rho)
+      # Include numerically stable versions of gr and rho when multiplying or
+      # dividing them. The sign of gr*rho and rho/gr is always positive.
+      safe_gr = tf.abs(gr) + 0.5 * tf.keras.backend.epsilon()
+      safe_rho = tf.abs(rho) + tf.keras.backend.epsilon()
+      exp_negative_q = gr / (2. * math.pi) * tf.exp(
+          -safe_rho / (2. * safe_gr * (1 + bar_rho)) +
+          (gh - rho) / (safe_gr * safe_rho) * mu1 * mu2)
+      covariance = s * (a + exp_negative_q)
+    elif self.activation not in (tf.keras.activations.linear, None):
+      raise NotImplementedError('Activation is {}. Deterministic variational '
+                                'inference is only available if activation is '
+                                'ReLU or None.'.format(self.activation))
+
+    return ed.MultivariateNormalFullCovariance(mean, covariance)
+
+
+def get_moments(x):
+  """Gets first and second moments of input."""
+  if isinstance(x, ed.RandomVariable):
+    mean = x.distribution.mean()
+    variance = x.distribution.variance()
+    try:
+      covariance = x.distribution.covariance()
+    except NotImplementedError:
+      covariance = tf.zeros(x.shape.concatenate(x.shape[-1]), dtype=x.dtype)
+      covariance = tf.matrix_set_diag(covariance, variance)
+  else:
+    mean = x
+    variance = tf.zeros_like(x)
+    covariance = tf.zeros(x.shape.concatenate(x.shape[-1]), dtype=x.dtype)
+  return mean, variance, covariance
+
+
+def soft_relu(x):
+  return (tfp.distributions.Normal(0., 1.).prob(x) +
+          x * tfp.distributions.Normal(0., 1.).cdf(x))
+
+
 @add_weight
 class DenseReparameterization(tf.keras.layers.Dense):
   """Bayesian densely-connected layer estimated via reparameterization.
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index a53b0ebb4..1737f558b 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -113,6 +113,18 @@ def testTrainableNormalStddevConstraint(self):
     self.assertAllGreater(res, 0.)
 
   @parameterized.parameters(
+      {"layer": bayes.DenseDVI,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "zeros",
+       "all_close": True},
+      {"layer": bayes.DenseDVI,
+       "kernel_initializer": "trainable_normal",
+       "bias_initializer": "zeros",
+       "all_close": False},
+      {"layer": bayes.DenseDVI,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "trainable_normal",
+       "all_close": False},
       {"layer": bayes.DenseReparameterization,
        "kernel_initializer": "zeros",
        "bias_initializer": "zeros",
@@ -155,7 +167,8 @@ def testDenseKernel(self,
     self.evaluate(tf.global_variables_initializer())
     res1, res2 = self.evaluate([outputs1, outputs2])
     self.assertEqual(res1.shape, (5, 3, 4))
-    self.assertAllGreaterEqual(res1, 0.)
+    if layer != bayes.DenseDVI:
+      self.assertAllGreaterEqual(res1, 0.)
     if all_close:
       self.assertAllClose(res1, res2)
     else:
@@ -163,6 +176,7 @@ def testDenseKernel(self,
     model.get_config()
 
   @parameterized.parameters(
+      {"layer": bayes.DenseDVI},
       {"layer": bayes.DenseReparameterization},
       {"layer": bayes.DenseVariationalDropout},
   )
@@ -184,9 +198,11 @@ def take_mean(f, *args, **kwargs):
     res1, res2 = self.evaluate([outputs1, outputs2])
     self.assertEqual(res1.shape, (5, 3, 4))
     self.assertNotAllClose(res1, res2)
-    self.assertAllClose(res2, np.zeros((5, 3, 4)), atol=1e-4)
+    if layer != bayes.DenseDVI:
+      self.assertAllClose(res2, np.zeros((5, 3, 4)), atol=1e-4)
 
   @parameterized.parameters(
+      {"layer": bayes.DenseDVI},
       {"layer": bayes.DenseReparameterization},
       {"layer": bayes.DenseVariationalDropout},
   )
@@ -238,6 +254,7 @@ def testDenseLoss(self, layer):
       self.assertIsNotNone(grad)
 
   @parameterized.parameters(
+      {"layer": bayes.DenseDVI},
       {"layer": bayes.DenseReparameterization},
       {"layer": bayes.DenseVariationalDropout},
   )
@@ -259,6 +276,7 @@ def testDenseModel(self, layer):
     self.assertLen(model.losses, 1)
 
   @parameterized.parameters(
+      {"layer": bayes.DenseDVI},
       {"layer": bayes.DenseReparameterization},
       {"layer": bayes.DenseVariationalDropout},
   )
@@ -282,6 +300,60 @@ class DenseSubclass(layer):
     self.assertEqual(res.shape, (3, 2))
     self.assertLen(model.losses, 1)
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testDenseDVIIsDeterministic(self):
+    """Tests that DenseDVI network has a deterministic loss function."""
+    features = tf.to_float(np.random.rand(3, 2))
+    labels = tf.to_float(np.random.rand(3, 1))
+    model = tf.keras.Sequential([
+        bayes.DenseDVI(5, activation=tf.nn.relu),
+        bayes.DenseDVI(1, activation=None),
+    ])
+    outputs = model(features, training=True)
+    nll = -tf.reduce_sum(outputs.distribution.log_prob(labels))
+    kl = sum(model.losses)
+    loss = nll + kl
+    self.evaluate(tf.global_variables_initializer())
+    res1 = self.evaluate(loss)
+    res2 = self.evaluate(loss)
+    self.assertEqual(res1, res2)
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testDenseDVIMoments(self):
+    """Verifies DenseDVI's moments empirically with samples."""
+    tf.set_random_seed(377269)
+    batch_size = 3
+    num_features = 5
+    units = 128
+    num_samples = 50000
+    inputs = tf.to_float(np.random.rand(batch_size, num_features))
+    layer = bayes.DenseDVI(units, activation=tf.nn.relu)
+
+    outputs1 = layer(inputs)
+    mean1 = outputs1.distribution.mean()
+    covariance1 = outputs1.distribution.covariance()
+
+    kernel_samples = layer.kernel.distribution.sample(num_samples)
+    outputs2 = layer.activation(
+        tf.einsum("bd,sdu->sbu", inputs, kernel_samples) +
+        tf.reshape(layer.bias, [1, 1, units]))
+    mean2 = tf.reduce_mean(outputs2, axis=0)
+    centered_outputs2 = tf.transpose(outputs2 - mean2, [1, 2, 0])
+    covariance2 = tf.matmul(centered_outputs2,
+                            centered_outputs2,
+                            transpose_b=True) / float(num_samples)
+
+    self.evaluate(tf.global_variables_initializer())
+    mean1_val, covariance1_val, mean2_val, covariance2_val = self.evaluate(
+        [mean1, covariance1, mean2, covariance2])
+    # Check % of mismatches is not too high according to heuristic thresholds.
+    num_mismatches = np.sum(np.abs(mean1_val - mean2_val) > 5e-3)
+    percent_mismatches = num_mismatches / float(batch_size * units)
+    self.assertLessEqual(percent_mismatches, 0.05)
+    num_mismatches = np.sum(np.abs(covariance1_val - covariance2_val) > 5e-3)
+    percent_mismatches = num_mismatches / float(batch_size * units * units)
+    self.assertLessEqual(percent_mismatches, 0.05)
+
   @test_utils.run_in_graph_and_eager_modes()
   def testGaussianProcessPosterior(self):
     train_batch_size = 3

From f25c089ea29a627033b6ba974a3d5bd09c577e63 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 30 Apr 2019 10:09:48 -0700
Subject: [PATCH 1977/2720] typos.

PiperOrigin-RevId: 245970313
---
 tensor2tensor/trax/layers/rnn_test.py        | 2 +-
 tensor2tensor/trax/models/neural_gpu_test.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/layers/rnn_test.py b/tensor2tensor/trax/layers/rnn_test.py
index b70934742..80794c131 100644
--- a/tensor2tensor/trax/layers/rnn_test.py
+++ b/tensor2tensor/trax/layers/rnn_test.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for google3.third_party.py.tensor2tensor.trax.layers.rnn."""
+"""Tests for trax.layers.rnn."""
 
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensor2tensor/trax/models/neural_gpu_test.py b/tensor2tensor/trax/models/neural_gpu_test.py
index 4f09a71c6..9e49d0a65 100644
--- a/tensor2tensor/trax/models/neural_gpu_test.py
+++ b/tensor2tensor/trax/models/neural_gpu_test.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for google3.third_party.py.tensor2tensor.trax.models.neural_gpu."""
+"""Tests for trax.models.neural_gpu."""
 
 from __future__ import absolute_import
 from __future__ import division

From 72e9512590dcd208f3754e46afba213b4157e139 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 30 Apr 2019 10:27:25 -0700
Subject: [PATCH 1978/2720] RenderedEnv image resize expects [width, height]
 instead of [height, width].

PiperOrigin-RevId: 245973890
---
 tensor2tensor/rl/gym_utils.py      | 25 ++++++++++++-------------
 tensor2tensor/rl/gym_utils_test.py |  4 ++--
 2 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index c4cd1da1d..af31020d9 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -99,30 +99,29 @@ def __init__(self, env, mode="rgb_array", low=0, high=255, resize_to=None):
     else:
       assert len(resize_to) == 2
       self.should_resize = True
+      num_channels = sample_frame.shape[-1]
       self.observation_space = gym.spaces.Box(
           low=low,
           high=high,
-          shape=list(resize_to) + list(sample_frame.shape[-1:]),
+          shape=list(resize_to) + [num_channels],
           dtype=sample_frame.dtype)
 
+  def _maybe_resize(self, obs):
+    if not self.should_resize:
+      return obs
+    height, width = self.observation_space.shape[:2]
+    img = Image.fromarray(obs)
+    img = img.resize([width, height], resample=Image.ANTIALIAS)
+    return np.array(img)
+
   def step(self, action):
     _, reward, done, info = self.env.step(action)
-    obs = self.env.render(mode=self.mode)
-    if self.should_resize:
-      img = Image.fromarray(obs)
-      img = img.resize(
-          self.observation_space.shape[:-1], resample=Image.ANTIALIAS)
-      obs = np.array(img)
+    obs = self._maybe_resize(self.env.render(mode=self.mode))
     return obs, reward, done, info
 
   def reset(self, **kwargs):
     self.env.reset(**kwargs)
-    obs = self.env.render(mode=self.mode)
-    if self.should_resize:
-      img = Image.fromarray(obs)
-      img = img.resize(self.observation_space.shape[:-1],
-                       resample=Image.ANTIALIAS)
-      obs = np.asarray(img)
+    obs = self._maybe_resize(self.env.render(mode=self.mode))
     return obs
 
 
diff --git a/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index 267f846b0..47d602922 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -84,9 +84,9 @@ def test_unlimited_env(self):
     self.assertTrue(env._max_episode_steps is None)
 
   def test_rendered_env(self):
-    env = gym_utils.RenderedEnv(SimpleEnv(), resize_to=(64, 64))
+    env = gym_utils.RenderedEnv(SimpleEnv(), resize_to=(64, 12))
     obs, _, _, _ = env.step(1)
-    self.assertTrue(np.allclose(np.zeros([64, 64, 3], np.uint8), obs))
+    self.assertTrue(np.allclose(np.zeros([64, 12, 3], np.uint8), obs))
 
   def test_gym_registration(self):
     reg_id, env = gym_utils.register_gym_env(

From a64b8fd3970bf41bcbd30545c9da55a873065138 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 30 Apr 2019 10:42:42 -0700
Subject: [PATCH 1979/2720] Add optional variables argument to optimize(). The
 default behavior remains unchanged

PiperOrigin-RevId: 245977065
---
 tensor2tensor/utils/t2t_model.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 41c2252d9..7d5943b44 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -707,14 +707,15 @@ def loss(self, logits, features):
           features["targets"],
           weights=features.get("targets_mask"))
 
-  def optimize(self, loss, num_async_replicas=1, use_tpu=False):
+  def optimize(self, loss, num_async_replicas=1, use_tpu=False, variables=None):
     """Return a training op minimizing loss."""
     lr = learning_rate.learning_rate_schedule(self.hparams)
     if num_async_replicas > 1:
       log_info("Dividing learning rate by num_async_replicas: %d",
                num_async_replicas)
     lr /= math.sqrt(float(num_async_replicas))
-    train_op = optimize.optimize(loss, lr, self.hparams, use_tpu=use_tpu)
+    train_op = optimize.optimize(
+        loss, lr, self.hparams, use_tpu=use_tpu, variables=variables)
     return train_op
 
   def set_mode(self, mode):

From 6347886b21207d02dd80fe811ae4e3255470ed0b Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 30 Apr 2019 11:14:22 -0700
Subject: [PATCH 1980/2720] internal

PiperOrigin-RevId: 245984411
---
 tensor2tensor/trax/layers/rnn.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/trax/layers/rnn.py b/tensor2tensor/trax/layers/rnn.py
index 91b3b69cd..82d8619a3 100644
--- a/tensor2tensor/trax/layers/rnn.py
+++ b/tensor2tensor/trax/layers/rnn.py
@@ -20,7 +20,6 @@
 from __future__ import google_type_annotations
 from __future__ import print_function
 
-import google3
 from tensor2tensor.trax.layers import combinators
 from tensor2tensor.trax.layers import core
 

From 2a362fd36c0167440bf013cb26de46086d4020f5 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 30 Apr 2019 16:48:51 -0700
Subject: [PATCH 1981/2720] A working Transformer in JAX and updates to tests.

PiperOrigin-RevId: 246048738
---
 .../trax/configs/transformer_wmt_ende_8gb.gin |  53 ++++
 tensor2tensor/trax/inputs.py                  |  35 ++-
 tensor2tensor/trax/layers/attention.py        |  49 ++--
 tensor2tensor/trax/layers/base.py             |  31 ++-
 tensor2tensor/trax/layers/base_test.py        |  86 +------
 tensor2tensor/trax/layers/combinators.py      |  57 ++++-
 tensor2tensor/trax/layers/combinators_test.py |  37 +++
 tensor2tensor/trax/layers/core_test.py        |  83 +++++++
 tensor2tensor/trax/layers/rnn.py              |  11 +-
 tensor2tensor/trax/layers/rnn_test.py         |  25 +-
 tensor2tensor/trax/models/__init__.py         |   1 +
 tensor2tensor/trax/models/neural_gpu_test.py  |  24 +-
 tensor2tensor/trax/models/transformer.py      | 232 +++++++-----------
 13 files changed, 430 insertions(+), 294 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
 create mode 100644 tensor2tensor/trax/layers/combinators_test.py
 create mode 100644 tensor2tensor/trax/layers/core_test.py

diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
new file mode 100644
index 000000000..f2a9a1e59
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
@@ -0,0 +1,53 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 64
+batch_fun.eval_batch_size = 64
+batch_fun.max_eval_length = 1024
+batch_fun.buckets_include_inputs_in_length=True
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_translate_ende_wmt32k'
+inputs.append_targets = True
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 0.1
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_preprocess
+wmt_preprocess.max_length = 512
+wmt_preprocess.max_eval_length = 1024
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 1000
+train.eval_steps = 10
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.Transformer
+train.run_debug_step = False
+train.train_steps = 500000
+
+# Parameters for Transformer:
+# ==============================================================================
+Transformer.dropout = 0.1
+Transformer.feature_depth = 512
+Transformer.feedforward_depth = 2048
+Transformer.max_len = 2048
+Transformer.mode = 'train'
+Transformer.num_heads = 8
+Transformer.num_layers = 6
+Transformer.vocab_size = 33300
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 2142f88a0..cab5cfd7d 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -88,6 +88,10 @@ def numpy_stream(dataset):
     length = input_shape[0]
     input_shape = tuple(
         [tuple([length // num_chunks] + list(input_shape)[1:])] * num_chunks)
+  if append_targets:
+    # TODO(lukaszkaiser): remove the assumption that input and target
+    # shapes are the same, which is used below for now.
+    input_shape = (input_shape, input_shape)
 
   return Inputs(train_stream=lambda: numpy_stream(train_batches),
                 train_eval_stream=lambda: numpy_stream(train_eval_batches),
@@ -262,6 +266,7 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
 def batch_fun(dataset, training, shapes, target_names, num_devices,
               batch_size_per_device=32, batch_size=None, eval_batch_size=32,
               bucket_length=32, buckets=None,
+              buckets_include_inputs_in_length=False,
               batch_shuffle_size=128, max_eval_length=None):
   """Batching function."""
   del target_names
@@ -303,8 +308,12 @@ def batch_fun(dataset, training, shapes, target_names, num_devices,
 
   if buckets:
     tf.logging.info("Bucketing with buckets %s." % str(buckets))
-    def example_length(_, target):
-      return tf.shape(target)[0]
+    def example_length(example_inputs, target):
+      """The length function used by bucket_by_sequence_length to bucket."""
+      other_length = 0
+      if buckets_include_inputs_in_length:
+        other_length = tf.shape(example_inputs["inputs"])[0]
+      return tf.maximum(tf.shape(target)[0], other_length)
     boundaries, batch_sizes = buckets
     dataset = dataset.apply(tf.data.experimental.bucket_by_sequence_length(
         example_length, boundaries, batch_sizes,
@@ -353,6 +362,28 @@ def eval_target_right_length(_, target):
   return dataset
 
 
+# TODO(lukaszkaiser): find a single more abstract way of text pre-processing.
+@gin.configurable(blacklist=["dataset", "training"])
+def wmt_preprocess(dataset, training, max_length=-1, max_eval_length=-1):
+  """Preprocessing for LM1B: filter out targets exceeding maximum length."""
+
+  def train_right_length(example, target):
+    l = tf.maximum(tf.shape(example["inputs"])[0], tf.shape(target)[0])
+    return tf.less(l, max_length + 1)
+
+  def eval_right_length(example, target):
+    l = tf.maximum(tf.shape(example["inputs"])[0], tf.shape(target)[0])
+    return tf.less(l, max_eval_length + 1)
+
+  if max_length > 0 and training:
+    dataset = dataset.filter(train_right_length)
+
+  if max_eval_length > 0 and not training:
+    dataset = dataset.filter(eval_right_length)
+
+  return dataset
+
+
 @gin.configurable(whitelist=["preprocess_fun", "shuffle_buffer_size"])
 def shuffle_and_batch_data(dataset,
                            target_names,
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index b2fe70814..dcaef56c5 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -40,38 +40,23 @@ def PaddingMask(x, params, pad=0, **kwargs):
   return np.reshape(x != pad, (x.shape[0], 1, 1, x.shape[-1]))
 
 
-def MakeTargetMask(target, pad=0):
-  """Create an attention mask to hide padding and future words."""
-  target_mask = (target != pad)[ :, np.newaxis, :]
-  target_dtype = target_mask.dtype
-  causal_mask = onp.tril(onp.ones((1, target.shape[-1], target.shape[-1]),
-                                  dtype=target_dtype), k=0)
-  target_mask = target_mask & causal_mask
-  return np.expand_dims(target_mask, axis=1)
-
-
-def PreparePairedSequenceBatch(source, target_in, pad=0):
-  """Build masks for this batch.
-
-  Args:
-    source: (batch, source_len) array of integer-coded symbols for inputs
-    target_in: (batch, batch_len) array of integer-coded symbols for targets
-    pad: int: the padding symbol used to pad the above
-
-  Returns:
-    Prepared batch of tuple of arrays: source, input-target, shifted-target,
-    source mask, target mask, source-target "memory" mask, minibatch token count
-  """
-  target = target_in[:, :-1]
-  target_y = target_in[:, 1:]
-  source_mask = np.reshape(source != pad,
-                           (source.shape[0], 1, 1, source.shape[-1]))
-  target_mask = MakeTargetMask(target, pad)
-  memory_mask = (
-      np.reshape(np.arange(target.shape[-1]) < source.shape[-1], [-1, 1]))
-  ntokens = np.sum(target_y != pad)
-  return (source, target, target_y,
-          source_mask, target_mask, memory_mask, ntokens)
+def EncoderDecoderMaskShape(inputs):
+  """Helper: shape for encoder-decoder mask."""
+  (padding_mask_shape, decoder_input_shape) = inputs
+  batch_size = padding_mask_shape[0]
+  input_length = padding_mask_shape[-1]
+  target_length = decoder_input_shape[1]
+  return (batch_size, 1, target_length, input_length)
+
+
+@base.layer(output_shape=EncoderDecoderMaskShape)
+def EncoderDecoderMask(x, **unused_kwargs):
+  """Make encoder-decoder mask from a padding mask and decoder input."""
+  (padding_mask, decoder_input) = x
+  padding_mask = np.reshape(
+      padding_mask, (padding_mask.shape[0], 1, 1, padding_mask.shape[-1]))
+  # Final mask shape is [batch, 1 for heads, decoder-len, encoder-len].
+  return padding_mask + np.ones((1, 1, decoder_input.shape[1], 1))
 
 
 # Layer normalization.
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 5f9285c9a..a0ac29d0a 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -22,6 +22,9 @@
 import inspect
 import traceback
 
+import numpy as onp
+from tensor2tensor.trax.backend import random
+
 
 class Layer(object):
   """Layer object, base class. Handles parameter sharing."""
@@ -170,7 +173,7 @@ def shapes(x):
   """Get a structure of shapes for a structure of nested arrays."""
   def shape(x):
     try:
-      return x.shape
+      return tuple([int(i) for i in x.shape])
     except Exception:  # pylint: disable=broad-except
       return []
   return nested_map(x, shape)
@@ -274,3 +277,29 @@ def call_fun(self, x, params=(), **kwargs):
 
     return cls
   return layer_decorator
+
+
+def _random_inputs(input_shape, rng, integer_inputs=False):
+  """Create random floats of the given shape."""
+  if isinstance(input_shape[0], int):  # Non-nested shape.
+    if not integer_inputs:
+      return random.uniform(rng, input_shape, minval=-1.0, maxval=1.0)
+    return random.bernoulli(rng, 0.5, input_shape).astype(onp.int32)
+  elif isinstance(input_shape, (list, tuple)):  # Nested shape.
+    return [_random_inputs(shape, rng, integer_inputs) for shape in input_shape]
+  else:
+    raise TypeError(type(input_shape))
+
+
+def check_shape_agreement(layer_instance, input_shape, integer_inputs=False):
+  """Check if layer.output_shape agrees with the actual output shape."""
+  rng1, rng2, rng3 = random.split(random.get_prng(0), 3)
+  output_shape = layer_instance.output_shape(input_shape)
+  output_shape = nested_map(output_shape, int)  # Make non-numpy.
+  params = layer_instance.initialize(input_shape, rng1)
+  inputs = _random_inputs(input_shape, rng2, integer_inputs=integer_inputs)
+  result = layer_instance(inputs, params, rng=rng3)
+  result_shape = shapes(result)
+  msg = 'output shape %s != real result shape %s' % (output_shape, result_shape)
+  assert output_shape == result_shape, msg
+  return output_shape
diff --git a/tensor2tensor/trax/layers/base_test.py b/tensor2tensor/trax/layers/base_test.py
index ed0baf65d..b8b9bc2c9 100644
--- a/tensor2tensor/trax/layers/base_test.py
+++ b/tensor2tensor/trax/layers/base_test.py
@@ -13,91 +13,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for base layers."""
+"""Tests for base layer."""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
 from absl.testing import absltest
-import numpy as onp
-from tensor2tensor.trax import layers
-from tensor2tensor.trax.backend import random
+from tensor2tensor.trax.layers import base
 
 
-def random_inputs(rng, input_shape):
-  if isinstance(input_shape, tuple):
-    return rng.randn(*input_shape).astype(onp.float32)
-  elif isinstance(input_shape, list):
-    return [random_inputs(rng, shape) for shape in input_shape]
-  else:
-    raise TypeError(type(input_shape))
+class BaseLayerTest(absltest.TestCase):
 
+  def test_layer_decorator_and_shape_agreement(self):
+    @base.layer()
+    def add_one(x, **unused_kwargs):
+      return x + 1
 
-def check_shape_agreement(test_case, layer, input_shape):
-  rng_key1, rng_key2 = random.split(random.get_prng(0))
-  result_shape = layer.output_shape(input_shape)
-  params = layer.initialize(input_shape, rng_key1)
-  inputs = random_inputs(onp.random.RandomState(0), input_shape)
-  result = layer(inputs, params, rng=rng_key2)
-  test_case.assertEqual(result.shape, result_shape)
-  return result_shape
-
-
-def check_layer(test_case, layer, input_shape):
-  return check_shape_agreement(test_case, layer, input_shape)
-
-
-class SlaxTest(absltest.TestCase):
-
-  def test_flatten_n(self):
-    input_shape = (29, 87, 10, 20, 30)
-
-    actual_shape = check_layer(self, layers.Flatten(), input_shape)
-    self.assertEqual(actual_shape, (29, 87 * 10 * 20 * 30))
-
-    actual_shape = check_layer(self, layers.Flatten(num_axis_to_keep=2),
-                               input_shape)
-    self.assertEqual(actual_shape, (29, 87, 10 * 20 * 30))
-
-    actual_shape = check_layer(self, layers.Flatten(num_axis_to_keep=3),
-                               input_shape)
-    self.assertEqual(actual_shape, (29, 87, 10, 20 * 30))
-
-    actual_shape = check_layer(self, layers.Flatten(num_axis_to_keep=4),
-                               input_shape)
-    self.assertEqual(actual_shape, (29, 87, 10, 20, 30))
-
-    # Not enough dimensions.
-    with self.assertRaises(ValueError):
-      check_layer(self, layers.Flatten(num_axis_to_keep=5), input_shape)
-
-    with self.assertRaises(ValueError):
-      check_layer(self, layers.Flatten(num_axis_to_keep=6), input_shape)
-
-  def test_div(self):
-    layer = layers.Div(divisor=2.0)
-    input_np = onp.array([[1, 2, 3], [4, 5, 6]], dtype=onp.float32)
-    output_np = layer(input_np)
-    # absltest doesn't have ndarray equalities.
-    expected_output_np = input_np / 2.0
-    self.assertAlmostEqual(
-        0.0,
-        onp.sum((output_np - expected_output_np) ** 2),
-        delta=1e-6)
-
-  def test_dense_param_sharing(self):
-    model1 = layers.Serial(layers.Dense(32), layers.Dense(32))
-    layer = layers.Dense(32)
-    model2 = layers.Serial(layer, layer)
-    rng = random.get_prng(0)
-    params1 = model1.initialize((-1, 32), rng)
-    params2 = model2.initialize((-1, 32), rng)
-    # The first parameters have 2 kernels of size (32, 32).
-    self.assertEqual((32, 32), params1[0][0].shape)
-    self.assertEqual((32, 32), params1[1][0].shape)
-    # The second parameters have 1 kernel of size (32, 32) and an empty dict.
-    self.assertEqual((32, 32), params2[0][0].shape)
-    self.assertEqual((), params2[1])
+    output_shape = base.check_shape_agreement(
+        add_one(), (12, 17))  # pylint: disable=no-value-for-parameter
+    self.assertEqual(output_shape, (12, 17))
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index a291133c8..5889d6d98 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -20,6 +20,8 @@
 from __future__ import print_function
 
 import operator
+import six
+
 from tensor2tensor.trax import backend
 from tensor2tensor.trax.layers import base
 
@@ -73,6 +75,37 @@ def Identity(x, **unused_kwargs):
   return x
 
 
+def Unnest(x):
+  """Helper: remove nesting in x, return a flat tuple."""
+  if not isinstance(x, (list, tuple)):
+    return (x,)
+  return tuple([z for y in x for z in Unnest(y)])  # pylint: disable=g-complex-comprehension
+
+
+def UnnestShape(shape):
+  """Unnest a nested structure of shapes."""
+
+  class Shape(object):
+    """Since shapes are tuples, make them a class to not unnest too far."""
+
+    def __init__(self, shape):
+      self.shape = shape
+
+  def MakeShape(nested_shape):
+    """Make all shape-tuples in the nested object shape-classes."""
+    if isinstance(nested_shape[0], int):  # Not nested.
+      return Shape(nested_shape)
+    return [MakeShape(shape) for shape in nested_shape]
+
+  # Unnest on the level of shape-classes and bring back shape-tuples.
+  return tuple([y.shape for y in Unnest(MakeShape(shape))])
+
+
+@base.layer(output_shape=UnnestShape)
+def UnnestBranches(x, **unused_kwargs):
+  return Unnest(x)
+
+
 # Re-ordering layer.
 def _reorder_shape(input_shape, output=None):  # pylint: disable=invalid-name
   """Helper to determine the shape of reorder output."""
@@ -131,6 +164,11 @@ def SecondBranch(x, **unused_kwargs):
   return x[1]  # Here x is a list of tensors, we select the second.
 
 
+@base.layer(output_shape=lambda input_shape_list: input_shape_list[2])
+def ThirdBranch(x, **unused_kwargs):
+  return x[2]  # Here x is a list of tensors, we select the third.
+
+
 def _nested_op(inputs, op):  # pylint: disable=invalid-name
   """Helper: sum a list of arrays or nested arrays."""
   # First the simple non-nested case.
@@ -150,7 +188,8 @@ def _nested_sum(inputs):  # pylint: disable=invalid-name
 
 
 def _nested_product(inputs):  # pylint: disable=invalid-name
-  return _nested_op(inputs=inputs, op=lambda xs: reduce(operator.mul, xs))
+  return _nested_op(
+      inputs=inputs, op=lambda xs: six.moves.reduce(operator.mul, xs))
 
 
 @base.layer(output_shape=lambda input_shape_list: input_shape_list[0])
@@ -227,9 +266,19 @@ def call(self, inputs, params=(), **kwargs):
     return [layer(x, params=p, rng=r, **kwargs)
             for layer, x, p, r in zip(self._layers, inputs, params, rngs)]
 
-  def output_shape(self, input_shapes):
-    return tuple([layer.output_shape(shape)
-                  for layer, shape in zip(self._layers, input_shapes)])
+  def output_shape(self, input_shape):
+    output_shapes = []
+    for i, layer in enumerate(self._layers):
+      try:
+        output_shapes.append(layer.output_shape(input_shape[i]))
+      except Exception:
+        # Since this is a widely used combinator, we improve errors here.
+        # Private methods are accessed as an exception for that reason.
+        name, trace = layer.__class__.__name__, base._short_traceback()  # pylint: disable=protected-access
+        raise base.LayerError(
+            name, 'output_shape',
+            layer._caller, input_shape[i], trace)  # pylint: disable=protected-access
+    return tuple(output_shapes)
 
   def new_parameters(self, input_shape, rng):
     rngs = backend.random.split(rng, self._nlayers)
diff --git a/tensor2tensor/trax/layers/combinators_test.py b/tensor2tensor/trax/layers/combinators_test.py
new file mode 100644
index 000000000..7cf4c0ce0
--- /dev/null
+++ b/tensor2tensor/trax/layers/combinators_test.py
@@ -0,0 +1,37 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 combinator layers."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+from tensor2tensor.trax.layers import base
+from tensor2tensor.trax.layers import combinators
+
+
+class CombinatorLayerTest(absltest.TestCase):
+
+  def test_unnest_branches(self):
+    input_shape = ((2, 3), [(4, 5), (6, 7)], (8, 9, 10))
+    expected_shape = ((2, 3), (4, 5), (6, 7), (8, 9, 10))
+    output_shape = base.check_shape_agreement(
+        combinators.UnnestBranches(), input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tensor2tensor/trax/layers/core_test.py b/tensor2tensor/trax/layers/core_test.py
new file mode 100644
index 000000000..ff4c4f32b
--- /dev/null
+++ b/tensor2tensor/trax/layers/core_test.py
@@ -0,0 +1,83 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 core layers."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+import numpy as onp
+from tensor2tensor.trax import backend
+from tensor2tensor.trax.layers import base
+from tensor2tensor.trax.layers import combinators
+from tensor2tensor.trax.layers import core
+
+
+class CoreLayerTest(absltest.TestCase):
+
+  def test_flatten_n(self):
+    input_shape = (29, 87, 10, 20, 30)
+
+    actual_shape = base.check_shape_agreement(core.Flatten(), input_shape)
+    self.assertEqual(actual_shape, (29, 87 * 10 * 20 * 30))
+
+    actual_shape = base.check_shape_agreement(
+        core.Flatten(num_axis_to_keep=2), input_shape)
+    self.assertEqual(actual_shape, (29, 87, 10 * 20 * 30))
+
+    actual_shape = base.check_shape_agreement(
+        core.Flatten(num_axis_to_keep=3), input_shape)
+    self.assertEqual(actual_shape, (29, 87, 10, 20 * 30))
+
+    actual_shape = base.check_shape_agreement(
+        core.Flatten(num_axis_to_keep=4), input_shape)
+    self.assertEqual(actual_shape, (29, 87, 10, 20, 30))
+
+    # Not enough dimensions.
+    with self.assertRaises(ValueError):
+      base.check_shape_agreement(core.Flatten(num_axis_to_keep=5), input_shape)
+
+    with self.assertRaises(ValueError):
+      base.check_shape_agreement(core.Flatten(num_axis_to_keep=6), input_shape)
+
+  def test_div(self):
+    layer = core.Div(divisor=2.0)
+    input_np = onp.array([[1, 2, 3], [4, 5, 6]], dtype=onp.float32)
+    output_np = layer(input_np)
+    # absltest doesn't have ndarray equalities.
+    expected_output_np = input_np / 2.0
+    self.assertAlmostEqual(
+        0.0,
+        onp.sum((output_np - expected_output_np) ** 2),
+        delta=1e-6)
+
+  def test_dense_param_sharing(self):
+    model1 = combinators.Serial(core.Dense(32), core.Dense(32))
+    layer = core.Dense(32)
+    model2 = combinators.Serial(layer, layer)
+    rng = backend.random.get_prng(0)
+    params1 = model1.initialize((-1, 32), rng)
+    params2 = model2.initialize((-1, 32), rng)
+    # The first parameters have 2 kernels of size (32, 32).
+    self.assertEqual((32, 32), params1[0][0].shape)
+    self.assertEqual((32, 32), params1[1][0].shape)
+    # The second parameters have 1 kernel of size (32, 32) and an empty dict.
+    self.assertEqual((32, 32), params2[0][0].shape)
+    self.assertEqual((), params2[1])
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tensor2tensor/trax/layers/rnn.py b/tensor2tensor/trax/layers/rnn.py
index 82d8619a3..13df5ae75 100644
--- a/tensor2tensor/trax/layers/rnn.py
+++ b/tensor2tensor/trax/layers/rnn.py
@@ -17,7 +17,6 @@
 
 from __future__ import absolute_import
 from __future__ import division
-from __future__ import google_type_annotations
 from __future__ import print_function
 
 from tensor2tensor.trax.layers import combinators
@@ -129,11 +128,11 @@ def GeneralGRUCell(candidate_transform,
 
               # Final projection + tanh to get Ct
               candidate_transform(),
-              candidate_nonlinearity()),  # Candidate gate
-
-          # Only apply dropout on the C gate.
-          # Paper reports that 0.1 is a good default.
-          core.Dropout(rate=dropout_rate_c)),
+              candidate_nonlinearity(),  # Candidate gate
 
+              # Only apply dropout on the C gate.
+              # Paper reports that 0.1 is a good default.
+              core.Dropout(rate=dropout_rate_c)),
+      ),
       # Gate memory and candidate
       combinators.GateBranches())
diff --git a/tensor2tensor/trax/layers/rnn_test.py b/tensor2tensor/trax/layers/rnn_test.py
index 80794c131..805961f2e 100644
--- a/tensor2tensor/trax/layers/rnn_test.py
+++ b/tensor2tensor/trax/layers/rnn_test.py
@@ -13,31 +13,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for trax.layers.rnn."""
+"""Tests for rnn layers."""
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import numpy as np
-from tensor2tensor.trax.backend import random as jax_random
+from absl.testing import absltest
+from tensor2tensor.trax.layers import base
 from tensor2tensor.trax.layers import rnn
-import tensorflow as tf
 
 
-class RnnModelTest(tf.test.TestCase):
+class RnnLayerTest(absltest.TestCase):
 
-  def _test_cell_runs(self, model, input_shape, output_shape):
-    source = np.ones(input_shape, dtype=np.float32)
-
-    # Build params
-    rng = jax_random.get_prng(0)
-    model.initialize(input_shape, rng)
-
-    # Run network
-    output = model(source)
-
-    self.assertEqual(output_shape, output.shape)
+  def _test_cell_runs(self, layer, input_shape, output_shape):
+    final_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(output_shape, final_shape)
 
   def test_conv_gru_cell(self):
     self._test_cell_runs(
@@ -51,4 +42,4 @@ def test_gru_cell(self):
 
 
 if __name__ == '__main__':
-  tf.test.main()
+  absltest.main()
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index 8d2c8ef59..e90ab053b 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -37,6 +37,7 @@ def model_configure(*args, **kwargs):
 MLP = model_configure(mlp.MLP)
 NeuralGPU = model_configure(neural_gpu.NeuralGPU)
 Resnet50 = model_configure(resnet.Resnet50)
+Transformer = model_configure(transformer.Transformer)
 TransformerEncoder = model_configure(transformer.TransformerEncoder)
 TransformerLM = model_configure(transformer.TransformerLM)
 WideResnet = model_configure(resnet.WideResnet)
diff --git a/tensor2tensor/trax/models/neural_gpu_test.py b/tensor2tensor/trax/models/neural_gpu_test.py
index 9e49d0a65..ab3cbecfa 100644
--- a/tensor2tensor/trax/models/neural_gpu_test.py
+++ b/tensor2tensor/trax/models/neural_gpu_test.py
@@ -19,30 +19,22 @@
 from __future__ import division
 from __future__ import print_function
 
-import numpy as np
-from tensor2tensor.trax.backend import random as jax_random
+from absl.testing import absltest
+from tensor2tensor.trax.layers import base
 from tensor2tensor.trax.models import neural_gpu
-import tensorflow as tf
 
 
-class NeuralGPUTest(tf.test.TestCase):
+class NeuralGPUTest(absltest.TestCase):
 
   def test_ngpu(self):
     vocab_size = 2
-    in_shape = [3, 5, 7]
-    source = np.ones(in_shape, dtype=np.int32)
-
+    input_shape = [3, 5, 7]
     model = neural_gpu.NeuralGPU(
         feature_depth=30, steps=4, vocab_size=vocab_size)
-    # Build params
-    rng = jax_random.get_prng(0)
-    model.initialize(in_shape, rng)
-
-    # Run network
-    output = model(source)
-
-    self.assertEqual(tuple(in_shape + [vocab_size]), output.shape)
+    final_shape = base.check_shape_agreement(
+        model, tuple(input_shape), integer_inputs=True)
+    self.assertEqual(tuple(input_shape + [vocab_size]), final_shape)
 
 
 if __name__ == '__main__':
-  tf.test.main()
+  absltest.main()
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index a2b082a1c..98c5d3242 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Transformer Model."""
+"""Transformer Models."""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
@@ -273,159 +273,111 @@ def ChunkedTransformerLM(vocab_size,
   )
 
 
-# TODO(lukaszkaiser): rewrite the model below.
+def EncoderDecoderLayer(feature_depth,
+                        feedforward_depth,
+                        num_heads,
+                        dropout,
+                        mode):
+  """Transformer encoder-decoder layer.
 
+  The input is a triple pair (encoder, mask, decoder_input) where
+  the mask is created from the original source to prevent attending
+  to the padding part of the encoder.
 
-def Transformer(source_vocab_size,
-                target_vocab_size,
-                mode='train',
-                num_layers=6,
+  Args:
+    feature_depth: int:  depth of embedding
+    feedforward_depth: int: depth of feed-forward layer
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate (how much to drop out)
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    the layer, returning a triple (encoder, mask, decoder_activations).
+  """
+  # Decoder self-attending to decoder.
+  self_attention = layers.Residual(
+      layers.LayerNorm(),
+      layers.Branch(),
+      layers.Parallel(layers.Identity(),  # activation for (q, k, v)
+                      layers.CausalMask(axis=-2)),  # attention mask
+      layers.MultiHeadedAttention(feature_depth, num_heads=num_heads,
+                                  dropout=dropout, mode=mode),
+      layers.Dropout(rate=dropout, mode=mode)
+  )
+  # Decoder attending to encoder.
+  encoder_decoder_attention = layers.Serial(
+      layers.Reorder(output=((2, 0, 0), 1)),  # ((dec, enc, enc), mask)
+      layers.MultiHeadedAttentionQKV(  # ((q, k, v), mask) --> new v
+          feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),
+      layers.Dropout(rate=dropout, mode=mode),
+  )
+  return layers.Serial(
+      layers.Parallel(layers.Identity(), layers.Identity(), self_attention),
+      layers.Branch(),
+      layers.Parallel(layers.Identity(), encoder_decoder_attention),
+      layers.UnnestBranches(),   # (encoder, mask, old_act, new_act)
+      layers.Reorder(output=(0, 1, (2, 3))),
+      layers.Parallel(  # Residual after encoder-decoder attention.
+          layers.Identity(), layers.Identity(), layers.SumBranches()),
+      layers.Parallel(  # Feed-forward on the third component (decoder).
+          layers.Identity(), layers.Identity(), ResidualFeedForward(
+              feature_depth, feedforward_depth, dropout, mode=mode)
+      )
+  )
+
+
+# TODO(lukaszkaiser): allow different source and target vocabularies.
+def Transformer(vocab_size,
                 feature_depth=512,
                 feedforward_depth=2048,
+                num_layers=6,
                 num_heads=8,
                 dropout=0.1,
-                shared_embedding=True,
-                max_len=200,
-                return_evals=False):
-  """Transformer model.
+                max_len=2048,
+                mode='train'):
+  """Transformer.
+
+  This model expects on input a pair (source, target).
 
   Args:
-    source_vocab_size: int: source vocab size
-    target_vocab_size: int: target vocab size
-    mode: str: 'train' or 'eval'
-    num_layers: int: number of encoder/decoder layers
+    vocab_size: int: vocab size (shared source and target).
     feature_depth: int:  depth of embedding
     feedforward_depth: int: depth of feed-forward layer
+    num_layers: int: number of encoder/decoder layers
     num_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
-    shared_embedding: bool: specify whether source/target embeddings are tied.
     max_len: int: maximum symbol length for positional encoding
-    return_evals: bool: whether to generate decode-time evaluation functions
+    mode: str: 'train' or 'eval'
 
   Returns:
-    A namedtuple containing model 'init' and 'apply' functions for training and
-  the 'evals' functions that itself returns a namedtuple containing evaluation
-  functions for the trained encoder, decoder, and generator substax.
+    the Transformer model.
   """
-  # Input embedding and positional encoding
-  inject_position = layers.Serial(
-      layers.Dropout(dropout, mode=mode),
-      layers.PositionalEncoding(feature_depth, max_len=max_len)
+  embedding = layers.Serial(
+      layers.Embedding(feature_depth, vocab_size),
+      layers.Dropout(rate=dropout, mode=mode),
+      layers.PositionalEncoding(max_len=max_len)
+  )
+  encoder = layers.Serial(
+      layers.Branch(),  # Branch input to create embedding and mask.
+      layers.Parallel(embedding, layers.PaddingMask()),
+      layers.Serial(*[EncoderLayer(feature_depth, feedforward_depth, num_heads,
+                                   dropout, mode)
+                      for _ in range(num_layers)]),
+      layers.Parallel(layers.LayerNorm(), layers.Identity())
+  )
+  stack = [EncoderDecoderLayer(feature_depth, feedforward_depth, num_heads,
+                               dropout, mode)
+           for _ in range(num_layers)]
+  return layers.Serial(
+      layers.Parallel(layers.Identity(), layers.ShiftRight()),
+      layers.Parallel(encoder, embedding),
+      layers.UnnestBranches(),  # (encoder, encoder_mask, decoder_input)
+      layers.Reorder(output=(0, (1, 2), 2)),
+      layers.Parallel(  # (encoder_mask, decoder_input) -> encoder-decoder mask
+          layers.Identity(), layers.EncoderDecoderMask(), layers.Identity()),
+      layers.Serial(*stack),
+      layers.ThirdBranch(),
+      layers.LayerNorm(),
+      layers.Dense(vocab_size),
+      layers.LogSoftmax()
   )
-  if shared_embedding:
-    assert source_vocab_size == target_vocab_size
-    # Weight-shared Embedding
-    embedding = layers.Share(layers.Embedding(feature_depth, source_vocab_size))
-    source_embedding_layer = layers.Serial(embedding, inject_position)
-    target_embedding_layer = source_embedding_layer
-  else:
-    source_embedding = layers.Embedding(feature_depth, source_vocab_size)
-    target_embedding = layers.Embedding(feature_depth, target_vocab_size)
-    source_embedding_layer = layers.Serial(source_embedding, inject_position)
-    target_embedding_layer = layers.Serial(target_embedding, inject_position)
-
-  # Multi-headed Attention and Feed-forward layers
-  multi_attention = layers.MultiHeadedAttention(
-      feature_depth, num_heads=num_heads, dropout=dropout, mode=mode)
-
-  # Encoder
-  @layers.Lambda
-  def Encoder(source, source_mask):
-    """Transformer encoder stack.
-
-    Args:
-      source: layer variable: raw source sequences
-      source_mask: layer variable: self-attention mask
-
-    Returns:
-      Layer variable that outputs encoded source.
-    """
-    encoder_layer = layers.Serial(
-        # input attends to self
-        layers.Residual(layers.LayerNorm(),
-                        layers.Branch(size=4),
-                        layers.Parallel(layers.Identity(),  # query
-                                        layers.Identity(),  # key
-                                        layers.Identity(),  # value
-                                        source_mask),  # attention mask
-                        multi_attention,
-                        layers.Dropout(dropout, mode=mode)),
-        # feed-forward
-        ResidualFeedForward(
-            feature_depth, feedforward_depth, dropout, mode=mode),
-    )
-    return layers.Serial(
-        source,
-        source_embedding_layer,
-        layers.repeat(encoder_layer, num_layers),
-        layers.LayerNorm(),
-    )
-
-  # Decoder
-  @layers.Lambda
-  def Decoder(memory, target, target_mask, memory_mask):
-    """Transformer decoder stack.
-
-    Args:
-      memory: layer variable: encoded source sequences
-      target: layer variable: raw target sequences
-      target_mask: layer variable: self-attention mask
-      memory_mask: layer variable: memory attention mask
-
-    Returns:
-      Layer variable that outputs encoded source.
-    """
-    decoder_layer = layers.Serial(
-        # target attends to self
-        layers.Residual(layers.LayerNorm(),
-                        layers.Branch(size=4),
-                        layers.Parallel(layers.Identity(),  # query
-                                        layers.Identity(),  # key
-                                        layers.Identity(),  # value
-                                        target_mask),  # attention mask
-                        multi_attention,
-                        layers.Dropout(dropout, mode=mode)),
-        # target attends to encoded source
-        layers.Residual(layers.LayerNorm(),
-                        layers.Branch(size=4),
-                        layers.Parallel(layers.Identity(),  # query
-                                        memory,  # key
-                                        memory,  # value
-                                        memory_mask),  # attention mask
-                        multi_attention,
-                        layers.Dropout(dropout, mode=mode)),
-        # feed-forward
-        ResidualFeedForward(
-            feature_depth, feedforward_depth, dropout, mode=mode)
-    )
-    return layers.Serial(
-        target,
-        target_embedding_layer,
-        layers.repeat(decoder_layer, num_layers),
-        layers.LayerNorm(),
-    )
-
-  # The Transformer
-  @layers.Lambda
-  def transformer(source, target, source_mask, target_mask, memory_mask):  # pylint: disable=invalid-name
-    encoded_source = Encoder(source, source_mask)
-    return Decoder(encoded_source, target, target_mask, memory_mask)
-
-  # Finally, bind the generator transform to use later for inference.
-  @layers.Lambda
-  def Generator(encoded_target):
-    return layers.Serial(
-        encoded_target,
-        layers.Dense(target_vocab_size),
-        layers.LogSoftmax
-    )
-
-  # Model-Building and Evaluation Functions
-  # Get entire model's the layer pair
-  top_init, top_apply = Generator(transformer)
-
-  # By default act as a normal constructor and emit an (init, apply) pair.
-  if not return_evals:
-    return (top_init, top_apply)
-  else:
-    raise ValueError('inference in this model is still a work in progress')

From e8a5927f6bf7baa2bc946a97dbd244d76ada0aa8 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 30 Apr 2019 20:47:54 -0700
Subject: [PATCH 1982/2720] Changes to allow to run trax with eager-mode TF.

PiperOrigin-RevId: 246073975
---
 tensor2tensor/trax/inputs.py      | 44 ++++++++++++++++++++-----------
 tensor2tensor/trax/layers/core.py |  2 +-
 tensor2tensor/trax/trainer.py     |  7 +++++
 3 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index cab5cfd7d..86bf93d98 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -230,27 +230,39 @@ def _select_features(example, feature_list=None):
   return {f: example[f] for f in feature_list if f in example}
 
 
+def _eager_dataset_iterator(dataset):
+  for item in dataset:
+    flat = tf.nest.flatten(item)
+    flat = [el.numpy() for el in flat]
+    yield tf.nest.pack_sequence_as(item, flat)
+
+
 def _train_and_eval_dataset_v1(problem_name, data_dir):
   """Return train and evaluation datasets, feature info and supervised keys."""
-  assert not tf.executing_eagerly(), "tf.eager mode must be turned off."
-  problem = t2t_problems.problem(problem_name)
-  train_dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, data_dir)
-  train_dataset = train_dataset.map(_select_features)
-  eval_dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, data_dir)
-  eval_dataset = eval_dataset.map(_select_features)
-  hparams = problem.get_hparams()
-  # We take a few training examples to guess the shapes.
-  input_shapes, target_shapes = [], []
-  example_tensor = train_dataset.make_one_shot_iterator().get_next()
-  sess = tf.Session()
-  example1 = sess.run(example_tensor)
-  example2 = sess.run(example_tensor)
-  example3 = sess.run(example_tensor)
+  with tf.device("cpu:0"):
+    problem = t2t_problems.problem(problem_name)
+    train_dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, data_dir)
+    train_dataset = train_dataset.map(_select_features)
+    eval_dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, data_dir)
+    eval_dataset = eval_dataset.map(_select_features)
+    hparams = problem.get_hparams()
+    # We take a few training examples to guess the shapes.
+    input_shapes, target_shapes, examples = [], [], []
+    if tf.executing_eagerly():
+      for example in _eager_dataset_iterator(train_dataset.take(3)):
+        examples.append(example)
+    else:
+      example_tensor = train_dataset.make_one_shot_iterator().get_next()
+      sess = tf.Session()
+      example1 = sess.run(example_tensor)
+      example2 = sess.run(example_tensor)
+      example3 = sess.run(example_tensor)
+      examples = [example1, example2, example3]
   # We use "inputs" as input except for purely auto-regressive tasks like
   # language models where "targets" are used as input_key.
-  input_key = "inputs" if "inputs" in example1 else "targets"
+  input_key = "inputs" if "inputs" in examples[0] else "targets"
   supervised_keys = ([input_key], ["targets"])
-  for example in [example1, example2, example3]:
+  for example in examples:
     input_shapes.append(list(example[input_key].shape))
     target_shapes.append(list(example["targets"].shape))
   input_vocab_size = hparams.vocab_size[input_key]
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 3e6307223..257bf84b0 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -82,7 +82,7 @@ def AddConstant(x, params, constant=0.0, **unused_kwargs):
 
 @base.layer()
 def Relu(x, **unused_kwargs):
-  return np.maximum(x, 0.)
+  return np.maximum(x, np.array(0, dtype=x.dtype))
 
 
 @base.layer()
diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index b5f9f5ffe..efbd582da 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -30,6 +30,9 @@
 import jax
 from tensor2tensor.trax import trax
 
+import tensorflow as tf
+
+
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string("dataset", None, "Which dataset to use.")
@@ -43,6 +46,7 @@
                           "Configuration parameters (gin string).")
 flags.DEFINE_integer("log_level", logging.INFO, "Log level.")
 flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
+flags.DEFINE_bool("tf_eager", False, "Whether we're running TF in eager mode.")
 
 
 def _default_output_dir():
@@ -84,6 +88,9 @@ def _setup_gin():
 def main(_):
   logging.set_verbosity(FLAGS.log_level)
 
+  if FLAGS.tf_eager:
+    tf.enable_eager_execution()
+
   _setup_gin()
 
   # Setup output directory

From 46fc2da1ca7d5acf06016f4896a9bfe7d9915942 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 1 May 2019 12:41:54 -0700
Subject: [PATCH 1983/2720] Collect trajectories using an EnvProblem and a
 policy function.  - EnvProblem and Trajectory changes to output observations
 (and actions and rewards) as numpy arrays.  - EnvProblem.reset() should also
 clear the internal cache of completed and incomplete trajectories.    - This
 led to a nasty bug where we use env again to collect trajectories      after
 resetting, but it silently kept returning the earlier trajectories      since
 it had 'num_trajectories' worth of completed trajectories.

PiperOrigin-RevId: 246179894
---
 tensor2tensor/envs/env_problem.py            |  14 ++-
 tensor2tensor/envs/env_problem_test.py       |  27 +++++
 tensor2tensor/envs/env_problem_utils.py      | 113 ++++++++++++++++++-
 tensor2tensor/envs/env_problem_utils_test.py |  40 ++++++-
 tensor2tensor/envs/trajectory.py             |  35 ++++--
 tensor2tensor/envs/trajectory_test.py        |  39 ++++++-
 6 files changed, 250 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 0cddfa74e..9f0e09dc1 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -477,15 +477,21 @@ def reset(self, indices=None):
     Subclasses should override _reset to do the actual reset if something other
     than the default implementation is desired.
 
+    NOTE: With `indices` as None the recorded trajectories are also erased since
+        the expecation is that we want to re-use the whole env class from
+        scratch.
+
     Args:
-      indices: Indices of environments to reset. If None all envs are reset.
+      indices: Indices of environments to reset. If None all envs are reset as
+          well as trajectories are erased.
 
     Returns:
       Batch of initial observations of reset environments.
     """
 
     if indices is None:
-      indices = np.arange(self.trajectories.batch_size)
+      self.trajectories.reset_batch_trajectories()
+      indices = np.arange(self.batch_size)
 
     # If this is empty (not None) then don't do anything, no env was done.
     if indices.size == 0:
@@ -497,7 +503,7 @@ def reset(self, indices=None):
     processed_observations = self.process_observations(observations)
 
     # Record history.
-    self.trajectories.reset(indices, observations)
+    self.trajectories.reset(indices, processed_observations)
 
     return processed_observations
 
@@ -731,7 +737,7 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
     # Write the completed data into these files
 
-    num_completed_trajectories = len(self.trajectories.completed_trajectories)
+    num_completed_trajectories = self.trajectories.num_completed_trajectories
     num_shards = len(files_list)
     if num_completed_trajectories < num_shards:
       tf.logging.warning(
diff --git a/tensor2tensor/envs/env_problem_test.py b/tensor2tensor/envs/env_problem_test.py
index ba2fc3678..f5060fb99 100644
--- a/tensor2tensor/envs/env_problem_test.py
+++ b/tensor2tensor/envs/env_problem_test.py
@@ -350,6 +350,33 @@ def reward_modality(self):
     self.assertEqual(dev_trajectories, dev_trajectories_ds)
     self.assertEqual(dev_timesteps, dev_timesteps_ds)
 
+  def test_resets_properly(self):
+    base_env_name = "CartPole-v0"
+    batch_size = 5
+    reward_range = (-1, 1)
+    nsteps = 100
+
+    env = env_problem.EnvProblem(
+        base_env_name=base_env_name,
+        batch_size=batch_size,
+        reward_range=reward_range)
+    env.name = base_env_name
+
+    num_dones = 0
+    while num_dones == 0:
+      env, num_dones, _ = self.play_env(env=env,
+                                        nsteps=nsteps,
+                                        batch_size=batch_size,
+                                        reward_range=reward_range)
+
+    # Some completed trajectories have been generated.
+    self.assertGreater(env.trajectories.num_completed_trajectories, 0)
+
+    # This should clear the env completely of any state.
+    env.reset()
+
+    # Assert that there aren't any completed trajectories in the env now.
+    self.assertEqual(env.trajectories.num_completed_trajectories, 0)
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 24e845cd4..3a9378e58 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -27,8 +27,7 @@ def done_indices(dones):
   return np.argwhere(dones).squeeze(axis=1)
 
 
-def play_env_problem_randomly(env_problem,
-                              num_steps):
+def play_env_problem_randomly(env_problem, num_steps):
   """Plays the env problem by randomly sampling actions for `num_steps`."""
   # Reset all environments.
   env_problem.reset()
@@ -36,11 +35,117 @@ def play_env_problem_randomly(env_problem,
   # Play all environments, sampling random actions each time.
   for _ in range(num_steps):
     # Sample batch_size actions from the action space and stack them.
-    actions = np.stack([env_problem.action_space.sample() for _ in range(
-        env_problem.batch_size)])
+    actions = np.stack([
+        env_problem.action_space.sample() for _ in range(env_problem.batch_size)
+    ])
 
     # Execute actions, observations are stored in `env_problem`.
     _, _, dones, _ = env_problem.step(actions)
 
     # Get the indices where we are done and reset those.
     env_problem.reset(indices=done_indices(dones))
+
+
+def play_env_problem_with_policy(env,
+                                 policy_fun,
+                                 num_trajectories=1,
+                                 max_timestep=None,
+                                 boundary=20):
+  """Plays the given env with the policy function to collect trajectories.
+
+  Args:
+    env: environment object, should be a subclass of env_problem.EnvProblem.
+    policy_fun: callable, taking in observations((B, T) + OBS) and returning
+        back log-probabilities (B, T, A).
+    num_trajectories: int, number of trajectories to collect.
+    max_timestep: int or None, if not None or a negative number, we cut any
+        trajectory that exceeds this time and mark that as completed by
+        resetting that trajectory.
+    boundary: this is the bucket length, we pad the observations to integer
+        multiples of this + 1 and then feed the padded observations to the
+        policy_fun.
+
+  Returns:
+    Completed trajectories that is a list of triples of (observation, action,
+    reward) ndarrays.
+  """
+
+  def multinomial_sample(probs):
+    """Sample from this vector of probabilities.
+
+    Args:
+      probs: numpy array of shape (A,) where A is the number of actions, these
+        must sum up to 1.0
+
+    Returns:
+      an integer of which action to pick.
+    """
+    return int(np.argwhere(np.random.multinomial(1, probs) == 1))
+
+  # We need to reset all environments.
+  env.reset()
+
+  while True:
+    # Get all the observations for all the active trajectories.
+    # Shape is (B, T) + OBS
+    padded_observations = env.trajectories.observations_np(boundary=boundary)
+    lengths = env.trajectories.trajectory_lengths
+
+    B, T = padded_observations.shape[:2]  # pylint: disable=invalid-name
+
+    assert B == env.batch_size
+    assert (B,) == lengths.shape
+
+    log_prob_actions = policy_fun(padded_observations)
+    assert (B, T) == log_prob_actions.shape[:2]
+    A = log_prob_actions.shape[2]  # pylint: disable=invalid-name
+
+    # We need the log_probs of those actions that correspond to the last actual
+    # time-step.
+    index = lengths - 1  # Since we want to index using lengths.
+    log_probs = log_prob_actions[np.arange(B)[:, None],
+                                 index[:, None],
+                                 np.arange(A)]
+    assert (B, A) == log_probs.shape
+
+    # Convert to probs, since we need to do categorical sampling.
+    probs = np.exp(log_probs)
+
+    # Now pick actions from this probs array.
+    actions = np.apply_along_axis(multinomial_sample, 1, probs)
+
+    # Step through the env.
+    _, _, dones, _ = env.step(actions)
+
+    # Get the indices where we are done ...
+    done_idxs = done_indices(dones)
+
+    # ... and reset those.
+    if done_idxs.size:
+      env.reset(indices=done_idxs)
+
+    # Do we have enough trajectories right now?
+    if env.trajectories.num_completed_trajectories >= num_trajectories:
+      break
+
+    if max_timestep is None or max_timestep < 1:
+      continue
+
+    # Are there any trajectories that have exceeded the time-limit we want.
+    lengths = env.trajectories.trajectory_lengths
+    exceeded_time_limit_idxs = done_indices(lengths > max_timestep)
+
+    # If so, reset these as well.
+    if exceeded_time_limit_idxs.size:
+      env.reset(indices=exceeded_time_limit_idxs)
+    # Do we have enough trajectories right now?
+    if env.trajectories.num_completed_trajectories >= num_trajectories:
+      break
+
+  # We have the trajectories we need, return a list of triples:
+  # (observations, actions, rewards)
+  completed_trajectories = []
+  for trajectory in env.trajectories.completed_trajectories[:num_trajectories]:
+    completed_trajectories.append(trajectory.as_numpy)
+
+  return completed_trajectories
diff --git a/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index 8f1e0db51..8b335c9d5 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -19,6 +19,9 @@
 from __future__ import division
 from __future__ import print_function
 
+import numpy as np
+
+from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.envs import tic_tac_toe_env  # pylint: disable=unused-import
 from tensor2tensor.envs import tic_tac_toe_env_problem
@@ -43,6 +46,41 @@ def test_play_env_problem_randomly(self):
         num_steps * batch_size + len(ep.trajectories.completed_trajectories) +
         batch_size, ep.trajectories.num_time_steps)
 
+  def test_play_env_problem_with_policy(self):
+    env = env_problem.EnvProblem(
+        base_env_name="CartPole-v0",
+        batch_size=2,
+        reward_range=(-1, 1))
+
+    def policy_fun(observations):
+      b, t = observations.shape[:2]
+      a = env.action_space.n
+      p = np.random.uniform(size=(b, t, a))
+      p = np.exp(p)
+      p = p / np.sum(p, axis=-1, keepdims=True)
+      return np.log(p)
+
+    max_timestep = 15
+    num_trajectories = 2
+    trajectories = env_problem_utils.play_env_problem_with_policy(
+        env, policy_fun, num_trajectories=num_trajectories,
+        max_timestep=max_timestep, boundary=20)
+
+    self.assertEqual(num_trajectories, len(trajectories))
+
+    # Check shapes within trajectories.
+    traj = trajectories[0]
+    T = traj[1].shape[0]  # pylint: disable=invalid-name
+    self.assertEqual((T+1, 4), traj[0].shape)  # (4,) is OBS
+    self.assertEqual((T,), traj[2].shape)
+    self.assertLessEqual(T, max_timestep)
+
+    traj = trajectories[1]
+    T = traj[1].shape[0]  # pylint: disable=invalid-name
+    self.assertEqual((T+1, 4), traj[0].shape)
+    self.assertEqual((T,), traj[2].shape)
+    self.assertLessEqual(T, max_timestep)
+
 
-if __name__ == '__main__':
+if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index 29677b4e4..451001e18 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -97,6 +97,20 @@ def reward(self):
   def observations_np(self):
     return np.stack([ts.observation for ts in self.time_steps])
 
+  @property
+  def actions_np(self):
+    # The last action is None, so let's skip it.
+    return np.stack([ts.action for ts in self.time_steps[:-1]])
+
+  @property
+  def rewards_np(self):
+    # The first reward is None, so let's skip it.
+    return np.stack([ts.processed_reward for ts in self.time_steps[1:]])
+
+  @property
+  def as_numpy(self):
+    return self.observations_np, self.actions_np, self.rewards_np
+
 
 class BatchTrajectory(object):
   """Basically a batch of active trajectories and a list of completed ones."""
@@ -265,24 +279,34 @@ def step(self, observations, raw_rewards, processed_rewards, dones, actions):
         # `reset` should be called on it.
         assert not self._trajectories[index].is_active
 
+  @staticmethod
+  def _trajectory_lengths(trajectories):
+    return np.array([t.num_time_steps for t in trajectories])
+
   @property
   def num_completed_time_steps(self):
     """Returns the number of time-steps in completed trajectories."""
 
-    return sum(t.num_time_steps for t in self.completed_trajectories)
+    return sum(BatchTrajectory._trajectory_lengths(self.completed_trajectories))
 
   @property
   def num_time_steps(self):
     """Returns the number of time-steps in completed and incomplete trajectories."""
 
-    num_time_steps = sum(t.num_time_steps for t in self.trajectories)
+    num_time_steps = sum(BatchTrajectory._trajectory_lengths(self.trajectories))
     return num_time_steps + self.num_completed_time_steps
 
+  @property
+  def trajectory_lengths(self):
+    return BatchTrajectory._trajectory_lengths(self.trajectories)
+
   @property
   def num_completed_trajectories(self):
     """Returns the number of completed trajectories."""
     return len(self.completed_trajectories)
 
+  # TODO(afrozm): Take in an already padded observation ndarray and just append
+  # the last time-step and adding more padding if needed.
   def observations_np(self, boundary=20):
     """Pads the observations in all the trajectories and returns them.
 
@@ -291,16 +315,13 @@ def observations_np(self, boundary=20):
           n is an integer.
 
     Returns:
-      a tuple(padded_observations, time_steps), with shapes:
       padded_observations: (self.batch_size, n * boundary + 1) + OBS
-      time_steps: integer list of length = self.batch_size
     """
     list_observations_np_ts = [t.observations_np for t in self.trajectories]
     # Every element in `list_observations_np_ts` is shaped (t,) + OBS
     OBS = list_observations_np_ts[0].shape[1:]  # pylint: disable=invalid-name
 
-    num_time_steps = [t.num_time_steps for t in self.trajectories]
-    t_max = max(num_time_steps)
+    t_max = max(self.trajectory_lengths)
     # t_max is rounded to the next multiple of `boundary`
     boundary = int(boundary)
     bucket_length = boundary * int(np.ceil(float(t_max) / boundary))
@@ -312,4 +333,4 @@ def padding_config(obs):
 
     return np.stack([
         np.pad(obs, padding_config(obs), "constant")
-        for obs in list_observations_np_ts]), num_time_steps
+        for obs in list_observations_np_ts])
diff --git a/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index acc50c1bb..e583cf80d 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -99,6 +99,39 @@ def test_observation_np(self):
 
     self.assertEqual((ts,) + shape, t.observations_np.shape)
 
+  def test_as_numpy(self):
+    t = trajectory.Trajectory()
+    shape = (3, 4)
+
+    # We'll have `ts` observations and `ts-1` actions and rewards.
+    ts = 5
+    num_actions = 6
+    observations = np.random.uniform(size=(ts,) + shape)
+    actions = np.random.choice(range(num_actions), size=(ts-1,))
+    rewards = np.random.choice([-1, 0, 1], size=(ts-1,))
+
+    # First time-step has no reward.
+    t.add_time_step(observation=observations[0],
+                    done=False,
+                    action=actions[0])
+    for i in range(1, ts - 1):
+      t.add_time_step(observation=observations[i],
+                      done=False,
+                      raw_reward=rewards[i-1],
+                      processed_reward=rewards[i-1],
+                      action=actions[i])
+    # Last time-step has no action.
+    t.add_time_step(observation=observations[-1],
+                    done=False,
+                    raw_reward=rewards[-1],
+                    processed_reward=rewards[-1])
+
+    traj_np = t.as_numpy
+
+    self.assertAllEqual(observations, traj_np[0])
+    self.assertAllEqual(actions, traj_np[1])
+    self.assertAllEqual(rewards, traj_np[2])
+
 
 class BatchTrajectoryTest(tf.test.TestCase):
 
@@ -290,7 +323,8 @@ def test_observations_np(self):
     lengths = lengths + ts
 
     boundary = 20
-    padded_obs_np, padded_lengths = bt.observations_np(boundary=boundary)
+    padded_obs_np = bt.observations_np(boundary=boundary)
+    padded_lengths = bt.trajectory_lengths
 
     # The lengths are what we expect them to be.
     self.assertAllEqual(lengths, padded_lengths)
@@ -311,7 +345,8 @@ def test_observations_np(self):
     # (16, 16, 21, 21, 21, 21, 21, 21, 21, 21)
     lengths = lengths + ts
 
-    padded_obs_np, padded_lengths = bt.observations_np(boundary=boundary)
+    padded_obs_np = bt.observations_np(boundary=boundary)
+    padded_lengths = bt.trajectory_lengths
 
     # The lengths are what we expect them to be.
     self.assertAllEqual(lengths, padded_lengths)

From 69d15e78a6aac641ecf0a5dac7fda57643166dd4 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 1 May 2019 16:37:01 -0700
Subject: [PATCH 1984/2720] Correct multi-device batch size re-computation in
 trax.

PiperOrigin-RevId: 246223051
---
 tensor2tensor/layers/bayes_test.py             | 5 ++---
 tensor2tensor/layers/reversible_layers_test.py | 7 ++++---
 tensor2tensor/trax/inputs.py                   | 5 +++++
 tensor2tensor/trax/trax.py                     | 6 +++---
 4 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 1737f558b..bab361192 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -26,7 +26,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow as tf
-from tensorflow_probability import edward2 as ed
+import tensorflow_probability as tfp
 tf.compat.v1.enable_eager_execution()
 
 
@@ -192,7 +192,7 @@ def take_mean(f, *args, **kwargs):
     inputs = tf.to_float(np.random.rand(5, 3, 7))
     model = layer(4, activation=tf.nn.relu, use_bias=False)
     outputs1 = model(inputs)
-    with ed.interception(take_mean):
+    with tfp.edward2.interception(take_mean):
       outputs2 = model(inputs)
     self.evaluate(tf.global_variables_initializer())
     res1, res2 = self.evaluate([outputs1, outputs2])
@@ -581,4 +581,3 @@ def testMixtureLogistic(self):
 
 if __name__ == "__main__":
   tf.test.main()
-
diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
index 335623fb5..8917e9f50 100644
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -26,11 +26,12 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow as tf
-from tensorflow_probability import edward2 as ed
+import tensorflow_probability as tfp
 tf.compat.v1.enable_eager_execution()
 
 
 
+
 class ReversibleLayersTest(parameterized.TestCase, tf.test.TestCase):
 
   @test_utils.run_in_graph_and_eager_modes()
@@ -133,7 +134,7 @@ class Exp(tf.keras.layers.Layer):
       """Exponential activation function for reversible networks."""
 
       def __call__(self, inputs, *args, **kwargs):
-        if not isinstance(inputs, ed.RandomVariable):
+        if not isinstance(inputs, tfp.edward2.RandomVariable):
           return super(Exp, self).__call__(inputs, *args, **kwargs)
         return reversible.TransformedRandomVariable(inputs, self)
 
@@ -146,7 +147,7 @@ def reverse(self, inputs):
       def log_det_jacobian(self, inputs):
         return -tf.log(inputs)
 
-    x = ed.Normal(0., 1.)
+    x = tfp.edward2.Normal(0., 1.)
     y = Exp()(x)
     y_sample = self.evaluate(y.distribution.sample())
     y_log_prob = self.evaluate(y.distribution.log_prob(y_sample))
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 86bf93d98..9b15dc830 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -151,6 +151,11 @@ def dataset_to_stream(dataset, input_name, num_chunks=0, append_targets=False):
   """Takes a tf.Dataset and creates a numpy stream of ready batches."""
   for example in tfds.as_numpy(dataset):
     inp, out = example[0][input_name], example[1]
+    # Some accelerators don't handle uint8 well, cast to int.
+    if isinstance(inp, np.uint8):
+      inp = inp.astype(np.uint32)
+    if isinstance(out, np.uint8):
+      out = out.astype(np.uint32)
     if len(out.shape) > 1 and out.shape[-1] == 1:
       out = np.squeeze(out, axis=-1)
     if num_chunks > 0:
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index ea60997a4..6577d245b 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -303,10 +303,10 @@ def mapped_predict(x, params, rng):
         jax_random.split(rng, num_devices))
     # Need to reduce the [device, per-device-batch, ...] tensors back to
     # a [batch, ...] tensor. The tensors may be nested.
-    if not isinstance(x, (list, tuple)):  # Not nested.
-      batch_size = x.shape[0]
+    if not isinstance(pred, (list, tuple)):  # Not nested.
+      batch_size = pred.shape[0] * pred.shape[1]
       return np.reshape(pred, [batch_size] + list(pred.shape[2:]))
-    batch_size = x[0].shape[0]
+    batch_size = pred[0].shape[0] * pred[0].shape[1]
     return [np.reshape(p, [batch_size] + list(p.shape[2:])) for p in pred]
 
   return predict

From a3363515c5085eb483e9d8597fb34965f0481ad1 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 1 May 2019 22:31:59 -0700
Subject: [PATCH 1985/2720] Embarrassing Error in PPO implementation, do
 min(prob * advantage, clipped prob * advantage) rather then min(prob, clipped
 prob) * advantage, because advantage can be negative.

This doesn't seem to help too much anyways though, seems to actually be hurting us in Pong/CartPole, but I'm looking at other implementations and "comparing notes".

PiperOrigin-RevId: 246258751
---
 tensor2tensor/trax/rlax/ppo.py      | 11 ++++----
 tensor2tensor/trax/rlax/ppo_test.py | 41 ++++++++++++++++++-----------
 2 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 5968d6da6..e1d7fe212 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -588,14 +588,15 @@ def compute_probab_ratios(p_new, p_old, actions, reward_mask):
   return probab_ratios
 
 
-def clipped_probab_ratios(probab_ratios, reward_mask, epsilon=0.2):
-  return reward_mask * np.clip(probab_ratios, 1 - epsilon, 1 + epsilon)
+def clipped_probab_ratios(probab_ratios, epsilon=0.2):
+  return np.clip(probab_ratios, 1 - epsilon, 1 + epsilon)
 
 
 def clipped_objective(probab_ratios, advantages, reward_mask, epsilon=0.2):
-  c1 = probab_ratios * reward_mask
-  c2 = clipped_probab_ratios(probab_ratios, reward_mask, epsilon=epsilon)
-  return np.minimum(c1, c2) * advantages
+  return np.minimum(
+      probab_ratios * advantages,
+      clipped_probab_ratios(probab_ratios, epsilon=epsilon) * advantages
+      ) * reward_mask
 
 
 @functools.partial(jit, static_argnums=(0, 3))
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 7fe3ae898..547c60188 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -453,14 +453,12 @@ def test_clipped_probab_ratios(self):
         [2.5, 2.0, 0.1, 1.0],
     ])
 
-    mask = np.array([[1, 1, 0, 0], [1, 1, 1, 0]])
-
-    clipped_probab_ratios = ppo.clipped_probab_ratios(probab_ratios, mask, 0.1)
+    clipped_probab_ratios = ppo.clipped_probab_ratios(probab_ratios, 0.1)
 
     self.assertAllClose(
         np.array([
-            [1.1, 1.0, 0, 0],
-            [1.1, 1.1, 0.9, 0],
+            [1.1, 1.0, 0.9, 0.9],
+            [1.1, 1.1, 0.9, 1.0],
         ]), clipped_probab_ratios)
 
   def test_clipped_objective(self):
@@ -470,32 +468,45 @@ def test_clipped_objective(self):
     ])
 
     advantages = np.array([
-        [0.1, 0.1, 0.5, 0.7],
-        [2.0, 2.0, 2.0, 2.0],
+        [0.1, -0.1, 0.5, 0.7],
+        [2.0, -2.0, 2.0, 2.0],
     ])
 
     mask = np.array([[1, 1, 0, 0], [1, 1, 1, 0]])
 
     epsilon = 0.1
 
-    unused_clipped_probab_ratios = np.array([
+    clipped_probab_ratios = np.array([
         [1.1, 1.1, 0.9, 0.9],
         [1.1, 1.1, 0.9, 1.0],
     ])
 
-    minimums = np.array([
-        [1.1, 1.1, 0.5, 0.7],
-        [1.1, 1.1, 0.1, 1.0],
+    unused_advantages_x_probab_ratios = np.array([
+        [0.15, -0.2, 0.25, 0.49],
+        [5.00, -4.0, 0.20, 2.00]
+    ])
+
+    unused_advantages_x_clipped_probab_ratios = np.array([
+        [0.11, -0.11, 0.45, 0.63],
+        [2.20, -2.20, 1.80, 2.00]
     ])
 
-    # advantages * minimums * mask
+    unused_minimums = np.array([
+        [0.11, -0.2, 0.25, 0.49],
+        [2.20, -4.0, 0.20, 2.00]
+    ])
+
+    # minimums * mask
     objective = np.array([
-        [0.11, 0.11, 0.0, 0.0],
-        [2.2, 2.2, 0.2, 0.0],
+        [0.11, -0.2, 0.0, 0.],
+        [2.20, -4.0, 0.2, 0.]
     ])
 
     # Assert that we computed things correctly in this test.
-    self.assertAllClose(advantages * mask * minimums, objective)
+    self.assertAllClose(
+        np.minimum(probab_ratios * advantages,
+                   clipped_probab_ratios * advantages) * mask,
+        objective)
 
     self.assertAllClose(
         objective,

From f679aba4a254cb7f5c6cea11f3e431226a269957 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 2 May 2019 09:16:22 -0700
Subject: [PATCH 1986/2720] minor adjustments to logging messages.

PiperOrigin-RevId: 246330432
---
 tensor2tensor/envs/env_problem.py      | 4 ++--
 tensor2tensor/envs/gym_spaces_utils.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 9f0e09dc1..1f589535a 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -388,10 +388,10 @@ def num_rewards(self):
     # Pre-conditions: reward range is finite.
     #               : processed rewards are discrete.
     if not self.is_reward_range_finite:
-      tf.logging.error("Infinite reward range, `num_rewards returning None`")
+      tf.logging.warn("Infinite reward range, `num_rewards returning None`")
       return None
     if not self.is_processed_rewards_discrete:
-      tf.logging.error(
+      tf.logging.warn(
           "Processed rewards are not discrete, `num_rewards` returning None")
       return None
 
diff --git a/tensor2tensor/envs/gym_spaces_utils.py b/tensor2tensor/envs/gym_spaces_utils.py
index a4ca516ff..9b96948de 100644
--- a/tensor2tensor/envs/gym_spaces_utils.py
+++ b/tensor2tensor/envs/gym_spaces_utils.py
@@ -99,8 +99,8 @@ def cardinality(gym_space):
   """
 
   if (gym_space.dtype == np.float32) or (gym_space.dtype == np.float64):
-    tf.logging.error("Returning None for a float gym space's cardinality: ",
-                     gym_space)
+    tf.logging.warn("Returning None for a float gym space's cardinality: %s",
+                    gym_space)
     return None
 
   if isinstance(gym_space, Discrete):

From 5eb6e612eaabf3dd23f3e965e54d9f690c039fad Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 2 May 2019 13:05:39 -0700
Subject: [PATCH 1987/2720]  - Add option in ppo_main to toggle between
 separate or unified policy and value nets.   - Then have the option to have
 different learning rates for them.  - Add an option that will flatten
 observation dimensions at the bottom layers.

PiperOrigin-RevId: 246374983
---
 tensor2tensor/trax/rlax/ppo_main.py | 72 +++++++++++++++++++++--------
 1 file changed, 54 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index edc4f43c1..66fb3735c 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -41,42 +41,78 @@
 flags.DEFINE_integer("max_timestep", None,
                      "If set to an integer, maximum number of time-steps in a "
                      "trajectory.")
-flags.DEFINE_float("learning_rate", 1e-3, "Learning rate.")
+flags.DEFINE_float("policy_and_value_net_learning_rate", 1e-3, "Learning rate.")
+flags.DEFINE_float("policy_net_learning_rate", 3e-4,
+                   "Learning rate for the policy net only.")
+flags.DEFINE_float("value_net_learning_rate", 1e-3,
+                   "Learning rate for the value net only.")
 flags.DEFINE_boolean("jax_debug_nans", False,
                      "Setting to true will help to debug nans and disable jit.")
 flags.DEFINE_boolean("disable_jit", False, "Setting to true will disable jit.")
+flags.DEFINE_boolean("combined_policy_and_value_function", False,
+                     "If True there is a single network that determines policy"
+                     "and values.")
+flags.DEFINE_integer("flatten_non_batch_time_dims", False,
+                     "If true, we flatten except the first two dimensions.")
 
 
 def common_layers():
   cur_layers = []
-  if FLAGS.env_name == "Pong-v0":
+  if FLAGS.flatten_non_batch_time_dims:
     cur_layers = [layers.Div(divisor=255.0), layers.Flatten(num_axis_to_keep=2)]
   return cur_layers + [layers.Dense(16), layers.Relu(),
                        layers.Dense(4), layers.Relu()]
 
 
+def run_training_loop():
+  """Run the PPO training loop."""
+
+  policy_net_fun = None
+  value_net_fun = None
+  policy_and_value_net_fun = None
+  policy_optimizer_fun = None
+  value_optimizer_fun = None
+  policy_and_value_optimizer_fun = None
+
+  if FLAGS.combined_policy_and_value_function:
+    policy_and_value_net_fun = functools.partial(
+        ppo.policy_and_value_net, bottom_layers=common_layers())
+    policy_and_value_optimizer_fun = get_optimizer_fun(
+        FLAGS.policy_and_value_net_learning_rate)
+  else:
+    policy_net_fun = functools.partial(ppo.policy_net,
+                                       bottom_layers=common_layers())
+    value_net_fun = functools.partial(ppo.value_net,
+                                      bottom_layers=common_layers())
+    policy_optimizer_fun = get_optimizer_fun(FLAGS.policy_net_learning_rate)
+    value_optimizer_fun = get_optimizer_fun(FLAGS.value_net_learning_rate)
+
+  ppo.training_loop(
+      env_name=FLAGS.env_name,
+      epochs=FLAGS.epochs,
+      policy_net_fun=policy_net_fun,
+      value_net_fun=value_net_fun,
+      policy_and_value_net_fun=policy_and_value_net_fun,
+      policy_optimizer_fun=policy_optimizer_fun,
+      value_optimizer_fun=value_optimizer_fun,
+      policy_and_value_optimizer_fun=policy_and_value_optimizer_fun,
+      batch_size=FLAGS.batch_size,
+      num_optimizer_steps=FLAGS.num_optimizer_steps,
+      boundary=FLAGS.boundary,
+      max_timestep=FLAGS.max_timestep,
+      random_seed=FLAGS.random_seed)
+
+
+def get_optimizer_fun(learning_rate):
+  return functools.partial(ppo.optimizer_fun, step_size=learning_rate)
+
+
 def main(argv):
   del argv
 
   if FLAGS.jax_debug_nans:
     config.update("jax_debug_nans", True)
 
-  def run_training_loop():
-    optimizer_fun = functools.partial(
-        ppo.optimizer_fun, step_size=FLAGS.learning_rate)
-
-    ppo.training_loop(
-        env_name=FLAGS.env_name,
-        epochs=FLAGS.epochs,
-        policy_and_value_net_fun=functools.partial(
-            ppo.policy_and_value_net, bottom_layers=common_layers()),
-        policy_and_value_optimizer_fun=optimizer_fun,
-        batch_size=FLAGS.batch_size,
-        num_optimizer_steps=FLAGS.num_optimizer_steps,
-        boundary=FLAGS.boundary,
-        max_timestep=FLAGS.max_timestep,
-        random_seed=FLAGS.random_seed)
-
   if FLAGS.jax_debug_nans or FLAGS.disable_jit:
     with jax.disable_jit():
       run_training_loop()

From dabd95e13650dc996d956ced2299d96b6fa968d6 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 2 May 2019 16:14:45 -0700
Subject: [PATCH 1988/2720] Add 20k and 50k setups for mbrl.

PiperOrigin-RevId: 246410996
---
 .../rl/trainer_model_based_params.py          | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 65d8e6046..bd7d7fa55 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -332,6 +332,26 @@ def rlmb_base_stochastic_discrete_sticky_actions():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_base_stochastic_discrete_20k():
+  """Base setting with stochastic discrete model with 20k steps."""
+  hparams = rlmb_base_stochastic_discrete()
+  # Our num_real_env_frames should be divisible by real_ppo_epoch_length*epochs
+  # Here we decrease epochs to 6 and make this number 16*200*6.
+  hparams.num_real_env_frames = 19200
+  hparams.epochs = 6
+  hparams.ppo_epochs_num = 2000  # Increase PPO steps as we have less epochs.
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic_discrete_50k():
+  """Base setting with stochastic discrete model with 50k steps."""
+  hparams = rlmb_base_stochastic_discrete()
+  hparams.num_real_env_frames = 48000
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_base_stochastic_discrete_75k_model_steps():
   """Base setting with stochastic discrete model with 75k WM steps."""

From 0a4e9122c62ac4db970733c92ea2a9bea40eb8ea Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 2 May 2019 16:22:55 -0700
Subject: [PATCH 1989/2720] Correct a subtle but equally embarassing bug in PPO
 (for the combined policy and value network case).

The PPO Loss is parametrized only over probs of the current policy, everything
else: advantages (whose inputs are rewards and value predictions) have to be
carried out with the old parameters.
PiperOrigin-RevId: 246412333
---
 tensor2tensor/trax/rlax/ppo.py      | 4 ++--
 tensor2tensor/trax/rlax/ppo_test.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index e1d7fe212..db3caa0ea 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -742,10 +742,10 @@ def combined_loss(new_params,
                   c1=1.0,
                   c2=0.01):
   """Computes the combined (clipped loss + value loss) given observations."""
-  log_probab_actions_new, value_predictions = policy_and_value_net_apply(
+  log_probab_actions_new, _ = policy_and_value_net_apply(
       padded_observations, new_params)
 
-  log_probab_actions_old, _ = policy_and_value_net_apply(
+  log_probab_actions_old, value_predictions = policy_and_value_net_apply(
       padded_observations, old_params)
 
   # (combined_loss, ppo_loss, value_loss, entropy_bonus)
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 547c60188..65b40de74 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -562,8 +562,8 @@ def test_combined_loss(self):
     mask = np.ones_like(rewards)
 
     # Just test that this computes at all.
-    new_log_probabs, value_predictions = net_apply(observations, new_params)
-    old_log_probabs, _ = net_apply(observations, old_params)
+    new_log_probabs, _ = net_apply(observations, new_params)
+    old_log_probabs, value_predictions = net_apply(observations, old_params)
 
     gamma = 0.99
     lambda_ = 0.95

From b40d5da1b38e5ab100c20e99754179368cce622c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 2 May 2019 17:36:17 -0700
Subject: [PATCH 1990/2720] Add testing gin configs.

PiperOrigin-RevId: 246423640
---
 .../configs/resnet50_imagenet_8gb_testing.gin | 44 ++++++++++++++++
 .../configs/transformer_lm1b_8gb_testing.gin  | 52 +++++++++++++++++++
 2 files changed, 96 insertions(+)
 create mode 100644 tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
 create mode 100644 tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin

diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
new file mode 100644
index 000000000..cf14b3882
--- /dev/null
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
@@ -0,0 +1,44 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.learning_rate
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 32
+batch_fun.bucket_length = 32
+batch_fun.buckets = None
+batch_fun.eval_batch_size = 32
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_image_imagenet224'
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+EvalAdjustingSchedule.constant = 1.0
+MultifactorSchedule.factors = 'constant * linear_warmup'
+MultifactorSchedule.warmup_steps = 400
+
+# Parameters for momentum:
+# ==============================================================================
+momentum.mass = 0.9
+
+
+# Parameters for Resnet50:
+# ==============================================================================
+Resnet50.hidden_size = 64
+Resnet50.num_output_classes = 1001
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 2000
+train.eval_steps = 20
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.Resnet50
+train.optimizer = @trax.optimizers.momentum
+train.train_steps = 100000
+train.lr_schedule = @learning_rate.EvalAdjustingSchedule
+
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
new file mode 100644
index 000000000..56415758a
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
@@ -0,0 +1,52 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 128
+batch_fun.eval_batch_size = 128
+batch_fun.max_eval_length = 2048
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_languagemodel_lm1b32k'
+inputs.input_name = 'targets'
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 0.1
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
+lm1b_preprocess.max_target_length = 512
+lm1b_preprocess.max_eval_target_length = 2048
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 1000
+train.eval_steps = 10
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerLM
+train.run_debug_step = False
+train.train_steps = 100000
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.dropout = 0.1
+TransformerLM.feature_depth = 512
+TransformerLM.feedforward_depth = 2048
+TransformerLM.max_len = 2048
+TransformerLM.mode = 'train'
+TransformerLM.num_heads = 8
+TransformerLM.num_layers = 6
+TransformerLM.vocab_size = 32000

From bc05cd6e0c9000db2b98ef4e9b2785a1a9ce0118 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 3 May 2019 13:31:56 -0700
Subject: [PATCH 1991/2720] Add and subtract 1 trick on probabilities not to
 exceed 1.

PiperOrigin-RevId: 246564015
---
 tensor2tensor/envs/env_problem_utils.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 3a9378e58..cb76f7ea5 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -111,6 +111,15 @@ def multinomial_sample(probs):
     # Convert to probs, since we need to do categorical sampling.
     probs = np.exp(log_probs)
 
+    # Sometimes log_probs contains a 0, it shouldn't. This makes the
+    # probabilities sum up to more than 1, since the addition happens
+    # in float64, so just add and subtract 1.0 to zero those probabilites
+    # out. Real example encountered probs = [1e-8, 1.0, 1e-22]
+    #
+    # Also testing for this is brittle.
+    probs += 1
+    probs -= 1
+
     # Now pick actions from this probs array.
     actions = np.apply_along_axis(multinomial_sample, 1, probs)
 

From 1aa26eaf0dbb9841b915bf8a7820355e3d527436 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 3 May 2019 17:06:09 -0700
Subject: [PATCH 1992/2720] - Use the correct params for computing PPO
 objective. - Use EnvProblem for collecting trajectories.   - We still keep
 the current collect using a single env. - For the separate network case, use
 early stopping for policy optimization.   - TODO: Implement the same thing
 for the combined network. - Normalize advantages. - Use separate values of
 the following for the combined network and separate network:   - number of
 optimizer steps   - learning rate - Change epsilon (in PPO loss) from a
 schedule, to a constant.

With these changes, and separate policy and value networks Acrobot-v1 works well.

PiperOrigin-RevId: 246600046
---
 tensor2tensor/trax/rlax/ppo.py      | 281 ++++++++++++++++++----------
 tensor2tensor/trax/rlax/ppo_main.py | 211 +++++++++++++++------
 2 files changed, 333 insertions(+), 159 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index db3caa0ea..91294d300 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -60,6 +60,8 @@
 from jax import numpy as np
 from jax import random as jax_random
 import numpy as onp
+from tensor2tensor.envs import env_problem
+from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.trax import layers
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
@@ -70,6 +72,8 @@
 EPSILON = 0.1
 EPOCHS = 50  # 100
 NUM_OPTIMIZER_STEPS = 100
+POLICY_ONLY_NUM_OPTIMIZER_STEPS = 80
+VALUE_ONLY_NUM_OPTIMIZER_STEPS = 80
 PRINT_EVERY_OPTIMIZER_STEP = 20
 BATCH_TRAJECTORIES = 32
 POLICY = "categorical-sampling"
@@ -122,10 +126,12 @@ def policy_and_value_net(rng_key,
   # Now, with the current logits, one head computes action probabilities and the
   # other computes the value function.
   # NOTE: The LogSoftmax instead of the Softmax because of numerical stability.
-  cur_layers.extend([layers.Branch(), layers.Parallel(
-      layers.Serial(layers.Dense(num_actions), layers.LogSoftmax()),
-      layers.Dense(1)
-  )])
+  cur_layers.extend([
+      layers.Branch(),
+      layers.Parallel(
+          layers.Serial(layers.Dense(num_actions), layers.LogSoftmax()),
+          layers.Dense(1))
+  ])
   net = layers.Serial(*cur_layers)
   return net.initialize(batch_observations_shape, rng_key), net
 
@@ -144,12 +150,11 @@ def log_params(params, name="params"):
       # Empty tuple.
       continue
     if not isinstance(param, (list, tuple)):
-      logging.error(
-          "%s[%d] : (%s) = [%s]", name, i, param.shape, onp.array(param))
+      logging.error("%s[%d] : (%s) = [%s]", name, i, param.shape,
+                    onp.array(param))
     else:
       for j, p in enumerate(param):
-        logging.error(
-            "\t%s[%d, %d] = [%s]", name, i, j, onp.array(p))
+        logging.error("\t%s[%d, %d] = [%s]", name, i, j, onp.array(p))
 
 
 # Should this be collect 'n' trajectories, or
@@ -161,6 +166,7 @@ def collect_trajectories(env,
                          num_trajectories=1,
                          policy="greedy",
                          max_timestep=None,
+                         boundary=20,
                          epsilon=0.1):
   """Collect trajectories with the given policy net and behaviour.
 
@@ -169,10 +175,11 @@ def collect_trajectories(env,
     policy_fun: observations(B,T+1) -> log-probabs(B,T+1, A) callable.
     num_trajectories: int, number of trajectories.
     policy: string, "greedy", "epsilon-greedy", or "categorical-sampling" i.e.
-        how to use the policy_fun to return an action.
+      how to use the policy_fun to return an action.
     max_timestep: int or None, the index of the maximum time-step at which we
-        return the trajectory, None for ending a trajectory only when env
-        returns done.
+      return the trajectory, None for ending a trajectory only when env returns
+      done.
+    boundary: int, boundary for padding, used in EnvProblem envs.
     epsilon: float, the epsilon for `epsilon-greedy` policy.
 
   Returns:
@@ -182,6 +189,16 @@ def collect_trajectories(env,
     action[i] = (B, T_i)
     reward[i] = (B, T_i)
   """
+
+  if isinstance(env, env_problem.EnvProblem):
+    # This is an env_problem, run its collect function.
+    return env_problem_utils.play_env_problem_with_policy(
+        env,
+        policy_fun,
+        num_trajectories=num_trajectories,
+        max_timestep=max_timestep,
+        boundary=boundary)
+
   trajectories = []
 
   for t in range(num_trajectories):
@@ -261,12 +278,12 @@ def collect_trajectories(env,
       logging.vlog(
           2, "  Collected time-step[ %5d] of trajectory[ %5d] in [%0.2f] msec.",
           ts, t, get_time(ts_start))
-    logging.vlog(
-        2, " Collected trajectory[ %5d] in [%0.2f] msec.", t, get_time(t_start))
+    logging.vlog(2, " Collected trajectory[ %5d] in [%0.2f] msec.", t,
+                 get_time(t_start))
 
     # This means we are done we're been terminated early.
-    assert done or (
-        max_timestep and max_timestep >= observation_history.shape[1])
+    assert done or (max_timestep and
+                    max_timestep >= observation_history.shape[1])
     # observation_history is (1, T+1) + OBS, lets squeeze out the batch dim.
     observation_history = np.squeeze(observation_history, axis=0)
     trajectories.append(
@@ -287,7 +304,7 @@ def get_padding_value(dtype):
     padding_value = np.uint8(0)
   elif dtype == np.uint16:
     padding_value = np.uint16(0)
-  elif dtype == np.float32:
+  elif dtype == np.float32 or dtype == np.float64:
     padding_value = 0.0
   else:
     padding_value = 0
@@ -595,8 +612,8 @@ def clipped_probab_ratios(probab_ratios, epsilon=0.2):
 def clipped_objective(probab_ratios, advantages, reward_mask, epsilon=0.2):
   return np.minimum(
       probab_ratios * advantages,
-      clipped_probab_ratios(probab_ratios, epsilon=epsilon) * advantages
-      ) * reward_mask
+      clipped_probab_ratios(probab_ratios, epsilon=epsilon) *
+      advantages) * reward_mask
 
 
 @functools.partial(jit, static_argnums=(0, 3))
@@ -635,15 +652,16 @@ def ppo_loss(policy_net_apply,
   assert (B, T + 1) == log_probab_actions_new.shape[:2]
   assert log_probab_actions_old.shape[-1] == log_probab_actions_new.shape[-1]
 
-  return ppo_loss_given_predictions(log_probab_actions_new,
-                                    log_probab_actions_old,
-                                    predicted_values,
-                                    padded_actions,
-                                    padded_rewards,
-                                    reward_mask,
-                                    gamma=gamma,
-                                    lambda_=lambda_,
-                                    epsilon=epsilon)
+  return ppo_loss_given_predictions(
+      log_probab_actions_new,
+      log_probab_actions_old,
+      predicted_values,
+      padded_actions,
+      padded_rewards,
+      reward_mask,
+      gamma=gamma,
+      lambda_=lambda_,
+      epsilon=epsilon)
 
 
 @jit
@@ -677,11 +695,12 @@ def ppo_loss_given_predictions(log_probab_actions_new,
   advantages = gae_advantages(
       td_deltas, reward_mask, lambda_=lambda_, gamma=gamma)
 
+  # Normalize the advantages.
+  advantages = (advantages - np.mean(advantages)) / np.std(advantages)
+
   # (B, T)
-  ratios = compute_probab_ratios(log_probab_actions_new,
-                                 log_probab_actions_old,
-                                 padded_actions,
-                                 reward_mask)
+  ratios = compute_probab_ratios(log_probab_actions_new, log_probab_actions_old,
+                                 padded_actions, reward_mask)
   assert (B, T) == ratios.shape
 
   # (B, T)
@@ -711,15 +730,16 @@ def combined_loss_given_predictions(log_probab_actions_new,
   """Computes the combined (clipped loss + value loss) given predictions."""
   loss_value = value_loss_given_predictions(
       value_prediction, padded_rewards, reward_mask, gamma=gamma)
-  loss_ppo = ppo_loss_given_predictions(log_probab_actions_new,
-                                        log_probab_actions_old,
-                                        value_prediction,
-                                        padded_actions,
-                                        padded_rewards,
-                                        reward_mask,
-                                        gamma=gamma,
-                                        lambda_=lambda_,
-                                        epsilon=epsilon)
+  loss_ppo = ppo_loss_given_predictions(
+      log_probab_actions_new,
+      log_probab_actions_old,
+      value_prediction,
+      padded_actions,
+      padded_rewards,
+      reward_mask,
+      gamma=gamma,
+      lambda_=lambda_,
+      epsilon=epsilon)
   # TODO(afrozm): Add the entropy bonus, but since we don't do that in T2T
   # we'll skip if for now.
   entropy_bonus = 0.0
@@ -749,17 +769,18 @@ def combined_loss(new_params,
       padded_observations, old_params)
 
   # (combined_loss, ppo_loss, value_loss, entropy_bonus)
-  return combined_loss_given_predictions(log_probab_actions_new,
-                                         log_probab_actions_old,
-                                         value_predictions,
-                                         padded_actions,
-                                         padded_rewards,
-                                         reward_mask,
-                                         c1=c1,
-                                         c2=c2,
-                                         gamma=gamma,
-                                         lambda_=lambda_,
-                                         epsilon=epsilon)
+  return combined_loss_given_predictions(
+      log_probab_actions_new,
+      log_probab_actions_old,
+      value_predictions,
+      padded_actions,
+      padded_rewards,
+      reward_mask,
+      c1=c1,
+      c2=c2,
+      gamma=gamma,
+      lambda_=lambda_,
+      epsilon=epsilon)
 
 
 @functools.partial(jit, static_argnums=(2, 3, 5))
@@ -833,6 +854,7 @@ def policy_and_value_opt_step(i,
                               lambda_=0.95,
                               epsilon=0.1):
   """Policy and Value optimizer step."""
+
   # Combined loss function given the new params.
   def policy_and_value_loss(params):
     """Returns the combined loss given just parameters."""
@@ -862,27 +884,31 @@ def get_time(t1, t2=None):
   return round((t2 - t1) * 1000, 2)
 
 
-def training_loop(env=None,
-                  env_name="CartPole-v0",
-                  epochs=EPOCHS,
-                  policy_net_fun=None,
-                  value_net_fun=None,
-                  policy_and_value_net_fun=None,
-                  policy_optimizer_fun=None,
-                  value_optimizer_fun=None,
-                  policy_and_value_optimizer_fun=None,
-                  batch_size=BATCH_TRAJECTORIES,
-                  num_optimizer_steps=NUM_OPTIMIZER_STEPS,
-                  print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
-                  boundary=20,
-                  max_timestep=None,
-                  random_seed=None,
-                  gamma=GAMMA,
-                  lambda_=LAMBDA,
-                  epsilon=EPSILON,
-                  c1=1.0,
-                  c2=0.01):
+def training_loop(
+    env=None,
+    epochs=EPOCHS,
+    policy_net_fun=None,
+    value_net_fun=None,
+    policy_and_value_net_fun=None,
+    policy_optimizer_fun=None,
+    value_optimizer_fun=None,
+    policy_and_value_optimizer_fun=None,
+    batch_size=BATCH_TRAJECTORIES,
+    num_optimizer_steps=NUM_OPTIMIZER_STEPS,
+    policy_only_num_optimizer_steps=POLICY_ONLY_NUM_OPTIMIZER_STEPS,
+    value_only_num_optimizer_steps=VALUE_ONLY_NUM_OPTIMIZER_STEPS,
+    print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
+    target_kl=0.01,
+    boundary=20,
+    max_timestep=None,
+    random_seed=None,
+    gamma=GAMMA,
+    lambda_=LAMBDA,
+    epsilon=EPSILON,
+    c1=1.0,
+    c2=0.01):
   """Runs the training loop for PPO, with fixed policy and value nets."""
+  assert env
   jax_rng_key = trax.get_random_number_generator_and_set_seed(random_seed)
 
   value_losses = []
@@ -890,8 +916,6 @@ def training_loop(env=None,
   combined_losses = []
   average_rewards = []
 
-  env = env if env is not None else gym.make(env_name)
-
   # Batch Observations Shape = [-1, -1] + OBS, because we will eventually call
   # policy and value networks on shape [B, T] +_OBS
   batch_observations_shape = (-1, -1) + env.observation_space.shape
@@ -913,6 +937,8 @@ def training_loop(env=None,
     # Initialize the optimizers.
     policy_and_value_opt_state, policy_and_value_opt_update = (
         policy_and_value_optimizer_fun(policy_and_value_net_params))
+
+    policy_and_value_net_apply = jit(policy_and_value_net_apply)
   else:
     # Initialize the policy and value functions.
     assert policy_net_fun and value_net_fun
@@ -924,19 +950,23 @@ def training_loop(env=None,
                                                       batch_observations_shape,
                                                       num_actions)
 
+    policy_net_apply = jit(policy_net_apply)
+    value_net_apply = jit(value_net_apply)
+
     # Initialize the optimizers.
     ppo_opt_state, ppo_opt_update = policy_optimizer_fun(policy_net_params)
     value_opt_state, value_opt_update = value_optimizer_fun(value_net_params)
 
   # A function that will call the appropriate policy function with parameters.
   def get_policy_output(observations):
+    # Get the fresh params for collecting the policy.
     if policy_net_apply is not None:
-      assert policy_net_params
-      return policy_net_apply(observations, policy_net_params)
+      return policy_net_apply(observations, trax_opt.get_params(ppo_opt_state))
+
+    assert policy_and_value_net_apply
 
-    assert policy_and_value_net_apply and policy_and_value_net_params
     policy_predictions, unused_value_predictions = policy_and_value_net_apply(
-        observations, policy_and_value_net_params)
+        observations, trax_opt.get_params(policy_and_value_opt_state))
     return policy_predictions
 
   for i in range(epochs):
@@ -949,8 +979,19 @@ def get_policy_output(observations):
         num_trajectories=batch_size,
         policy=POLICY,
         max_timestep=max_timestep,
+        boundary=boundary,
         epsilon=(10.0 / (i + 10.0)))  # this is a different epsilon.
 
+    logging.vlog(1, "Collecting trajectories took %0.2f msec.", get_time(t))
+
+    # These were the params that were used to collect the trajectory.
+    if policy_and_value_net_apply:
+      policy_and_value_net_params = trax_opt.get_params(
+          policy_and_value_opt_state)
+    else:
+      policy_net_params = trax_opt.get_params(ppo_opt_state)
+      value_net_params = trax_opt.get_params(value_opt_state)
+
     avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
     max_reward = max(np.sum(traj[2]) for traj in trajs)
     min_reward = min(np.sum(traj[2]) for traj in trajs)
@@ -958,16 +999,20 @@ def get_policy_output(observations):
 
     logging.vlog(1, "Rewards average=[%0.2f], max=[%0.2f], min=[%0.2f]",
                  avg_reward, max_reward, min_reward)
-    logging.vlog(1, "Collecting trajectories took %0.2f msec.", get_time(t))
+    logging.vlog(2, "Rewards: %s", [float(np.sum(traj[2])) for traj in trajs])
+    logging.vlog(1, "Average Rewards: %s", average_rewards)
+
     logging.vlog(1,
                  "Trajectory Length average=[%0.2f], max=[%0.2f], min=[%0.2f]",
                  float(sum(len(traj[0]) for traj in trajs)) / len(trajs),
                  max(len(traj[0]) for traj in trajs),
                  min(len(traj[0]) for traj in trajs))
+    logging.vlog(2, "Trajectory Lengths: %s", [len(traj[0]) for traj in trajs])
 
     t = time.time()
     (_, reward_mask, padded_observations, padded_actions,
-     padded_rewards) = pad_trajectories(trajs, boundary=boundary)
+     padded_rewards) = pad_trajectories(
+         trajs, boundary=boundary)
 
     logging.vlog(1, "Padding trajectories took %0.2f msec.", get_time(t))
     logging.vlog(1, "Padded Observations' shape [%s]",
@@ -983,13 +1028,17 @@ def get_policy_output(observations):
     assert (B, T + 1) + env.observation_space.shape == padded_observations.shape
 
     # Linear annealing from 0.1 to 0.0
-    epsilon_schedule = epsilon if epochs == 1 else epsilon * (1.0 -
-                                                              (i /
-                                                               (epochs - 1)))
+    # epsilon_schedule = epsilon if epochs == 1 else epsilon * (1.0 -
+    #                                                           (i /
+    #                                                            (epochs - 1)))
+
+    # Constant epsilon.
+    epsilon_schedule = epsilon
 
     # Compute value and ppo losses.
     cur_value_loss, cur_ppo_loss, cur_combined_loss = None, None, None
     if policy_and_value_net_apply is not None:
+      logging.vlog(2, "Starting to compute P&V loss.")
       t = time.time()
       cur_combined_loss, cur_ppo_loss, cur_value_loss, _ = (
           combined_loss(
@@ -1038,7 +1087,8 @@ def get_policy_output(observations):
 
     value_losses.append(cur_value_loss)
     ppo_objective.append(-1.0 * cur_ppo_loss)
-    combined_losses.append(cur_combined_loss)
+    if cur_combined_loss:
+      combined_losses.append(cur_combined_loss)
 
     if policy_and_value_net_apply:
       logging.vlog(1, "Policy and Value Optimization")
@@ -1051,6 +1101,8 @@ def get_policy_output(observations):
             policy_and_value_opt_state,
             policy_and_value_opt_update,
             policy_and_value_net_apply,
+            # for the entirety of this loop, this should refer to params that
+            # were used to collect the trajectory.
             policy_and_value_net_params,
             padded_observations,
             padded_actions,
@@ -1062,16 +1114,17 @@ def get_policy_output(observations):
             lambda_=lambda_,
             epsilon=epsilon_schedule)
         t2 = time.time()
-        # Get the new params.
-        new_policy_and_value_net_params = trax_opt.get_params(
-            policy_and_value_opt_state)
         if ((j + 1) %
             print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
           # Compute and log the loss.
+          # Get the new params.
+          new_policy_and_value_net_params = trax_opt.get_params(
+              policy_and_value_opt_state)
           (loss_combined, loss_ppo, loss_value, unused_entropy_bonus) = (
               combined_loss(
                   new_policy_and_value_net_params,
-                  policy_and_value_net_params,  # old params
+                  # old params, that were used to collect the trajectory
+                  policy_and_value_net_params,
                   policy_and_value_net_apply,
                   padded_observations,
                   padded_actions,
@@ -1088,25 +1141,26 @@ def get_policy_output(observations):
               1,
               "Combined Loss(value, ppo) [%10.2f] -> [%10.2f(%10.2f,%10.2f)]",
               cur_combined_loss, loss_combined, loss_value, loss_ppo)
-        # Update the params.
-        policy_and_value_net_params = new_policy_and_value_net_params
+
+      # Update the params.
+      policy_and_value_net_params = new_policy_and_value_net_params
 
       logging.vlog(
-          1, "Total PPO loss reduction [%0.2f]%%",
+          1, "Total Combined Loss reduction [%0.2f]%%",
           (100 *
            (cur_combined_loss - loss_combined) / np.abs(cur_combined_loss)))
 
       logging.info(
           "Epoch [% 6d], Reward[min, max, avg] [%10.2f,%10.2f,%10.2f], Combined"
-          " Loss(value, ppo) [%10.2f(%10.2f,%10.2f)], took [%10.2f msec]",
-          i, min_reward, max_reward, avg_reward, loss_combined, loss_value,
+          " Loss(value, ppo) [%10.2f(%10.2f,%10.2f)], took [%10.2f msec]", i,
+          min_reward, max_reward, avg_reward, loss_combined, loss_value,
           loss_ppo, get_time(t1))
     else:
       # Run optimizers.
       logging.vlog(1, "PPO Optimization")
       t1 = time.time()
 
-      for j in range(num_optimizer_steps):
+      for j in range(policy_only_num_optimizer_steps):
         t = time.time()
         # Update the optimizer state.
         ppo_opt_state = ppo_opt_step(
@@ -1128,8 +1182,27 @@ def get_policy_output(observations):
         t2 = time.time()
         # Get the new params.
         new_policy_net_params = trax_opt.get_params(ppo_opt_state)
-        if ((j + 1) %
-            print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
+
+        # These are the "old" params - policy_net_params
+
+        # Compute the approx KL for early stopping.
+        log_probab_actions_old = policy_net_apply(padded_observations,
+                                                  policy_net_params)
+        log_probab_actions_new = policy_net_apply(padded_observations,
+                                                  new_policy_net_params)
+
+        approx_kl = np.mean(log_probab_actions_old - log_probab_actions_new)
+
+        early_stopping = approx_kl > 1.5 * target_kl
+        if early_stopping:
+          logging.vlog(
+              1, "Early stopping policy optimization at iter: %d, "
+              "with approx_kl: %0.2f", j, approx_kl)
+          # We don't return right-away, we want the below to execute on the last
+          # iteration.
+
+        if (((j + 1) % print_every_optimizer_steps == 0) or
+            (j == num_optimizer_steps - 1) or early_stopping):
           new_ppo_loss = ppo_loss(
               policy_net_apply,
               new_policy_net_params,
@@ -1147,15 +1220,21 @@ def get_policy_output(observations):
           logging.vlog(1, "One PPO grad desc took: %0.2f msec", get_time(t, t2))
           logging.vlog(1, "PPO loss [%10.2f] -> [%10.2f]", cur_ppo_loss,
                        new_ppo_loss)
-        # Update the params.
-        policy_net_params = new_policy_net_params
+
+        if early_stopping:
+          break
+
+      # Update the params ONLY AND ONLY AFTER we complete all the optimization
+      # iterations, till then `policy_net_params` should refer to the params
+      # that were used in collecting the policy.
+      # policy_net_params = trax_opt.get_params(ppo_opt_state)
 
       logging.vlog(1, "Total PPO loss reduction [%0.2f]%%",
                    (100 * (cur_ppo_loss - new_ppo_loss) / np.abs(cur_ppo_loss)))
 
       logging.vlog(1, "Value Optimization")
 
-      for j in range(num_optimizer_steps):
+      for j in range(value_only_num_optimizer_steps):
         t = time.time()
         value_opt_state = value_opt_step(
             j,
@@ -1208,9 +1287,11 @@ def get_policy_output(observations):
   if value_losses:
     logging.vlog(1, "value_losses: %s", np.stack(value_losses))
   if ppo_objective:
-    logging.vlog(1, "ppo_objective: %s", np.stack(ppo_objective))
+    logging.vlog(1, "ppo_objective:\n%s", np.stack(ppo_objective))
+  if combined_losses:
+    logging.vlog(1, "combined_losses:\n%s", np.stack(combined_losses))
   if average_rewards:
-    logging.vlog(1, "average_rewards: %s", average_rewards)
+    logging.vlog(1, "average_rewards:\n%s", average_rewards)
 
   return ((policy_net_params, value_net_params), average_rewards,
           np.stack(value_losses), np.stack(ppo_objective))
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 66fb3735c..ac80982e9 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -13,7 +13,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""PPO binary over a gym env."""
+r"""PPO binary over a gym env.
+
+Sample invocation:
+
+ENV_PROBLEM_NAME=Acrobot-v1
+COMBINED_NETWORK=false
+EPOCHS=100
+BATCH_SIZE=32
+RANDOM_SEED=0
+BOUNDARY=100
+
+
+blaze run --config=cuda third_party/py/tensor2tensor/trax/rlax:ppo_main -- \
+  --env_problem_name=${ENV_PROBLEM_NAME} \
+  --combined_policy_and_value_function=${COMBINED_NETWORK} \
+  --epochs=${EPOCHS} \
+  --batch_size=${BATCH_SIZE} \
+  --random_seed=${RANDOM_SEED} \
+  --boundary=${BOUNDARY} \
+  --vmodule=*/tensor2tensor/*=1 \
+  --alsologtostderr \
+
+"""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -23,84 +45,111 @@
 
 from absl import app
 from absl import flags
+import gym
 import jax
 from jax.config import config
+from tensor2tensor.envs import env_problem
+from tensor2tensor.envs import rendered_env_problem
+from tensor2tensor.rl import gym_utils
 from tensor2tensor.trax import layers
 from tensor2tensor.trax.rlax import ppo
 
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string("env_name", None, "Name of the environment to make.")
-flags.DEFINE_string("t2t_gym_env", None, "Name of the T2TGymEnv to make.")
+flags.DEFINE_string("env_problem_name", None, "Name of the EnvProblem to make.")
+
 flags.DEFINE_integer("epochs", 100, "Number of epochs to run for.")
 flags.DEFINE_integer("random_seed", 0, "Random seed.")
 flags.DEFINE_integer("batch_size", 32, "Batch of trajectories needed.")
-flags.DEFINE_integer("num_optimizer_steps", 100, "Number of optimizer steps.")
-flags.DEFINE_integer("boundary", 20,
-                     "We pad trajectories at integer multiples of this number.")
-flags.DEFINE_integer("max_timestep", None,
-                     "If set to an integer, maximum number of time-steps in a "
-                     "trajectory.")
-flags.DEFINE_float("policy_and_value_net_learning_rate", 1e-3, "Learning rate.")
-flags.DEFINE_float("policy_net_learning_rate", 3e-4,
-                   "Learning rate for the policy net only.")
-flags.DEFINE_float("value_net_learning_rate", 1e-3,
-                   "Learning rate for the value net only.")
-flags.DEFINE_boolean("jax_debug_nans", False,
-                     "Setting to true will help to debug nans and disable jit.")
+
+flags.DEFINE_integer(
+    "boundary", 20, "We pad trajectories at integer multiples of this number.")
+# -1: returns env as is.
+# None: unwraps and returns without TimeLimit wrapper.
+# Any other number: imposes this restriction.
+flags.DEFINE_integer(
+    "max_timestep", None,
+    "If set to an integer, maximum number of time-steps in a "
+    "trajectory.")
+
+flags.DEFINE_boolean(
+    "jax_debug_nans", False,
+    "Setting to true will help to debug nans and disable jit.")
 flags.DEFINE_boolean("disable_jit", False, "Setting to true will disable jit.")
-flags.DEFINE_boolean("combined_policy_and_value_function", False,
-                     "If True there is a single network that determines policy"
-                     "and values.")
-flags.DEFINE_integer("flatten_non_batch_time_dims", False,
+
+# If resize is True, then we create RenderedEnvProblem, so this has to be set to
+# False for something like CartPole.
+flags.DEFINE_boolean("resize", False, "If true, resize the game frame")
+flags.DEFINE_integer("resized_height", 105, "Resized height of the game frame.")
+flags.DEFINE_integer("resized_width", 80, "Resized width of the game frame.")
+
+flags.DEFINE_boolean(
+    "combined_policy_and_value_function", False,
+    "If True there is a single network that determines policy"
+    "and values.")
+
+flags.DEFINE_boolean("flatten_non_batch_time_dims", False,
                      "If true, we flatten except the first two dimensions.")
 
+# Number of optimizer steps of the combined net, policy net and value net.
+flags.DEFINE_integer("num_optimizer_steps", 100, "Number of optimizer steps.")
+flags.DEFINE_integer("policy_only_num_optimizer_steps", 80,
+                     "Number of optimizer steps policy only.")
+flags.DEFINE_integer("value_only_num_optimizer_steps", 80,
+                     "Number of optimizer steps value only.")
+
+# Learning rate of the combined net, policy net and value net.
+flags.DEFINE_float("learning_rate", 5e-4, "Learning rate.")
+flags.DEFINE_float("policy_only_learning_rate", 1e-3,
+                   "Learning rate for policy network only.")
+flags.DEFINE_float("value_only_learning_rate", 1e-3,
+                   "Learning rate for value network only.")
+
+# Target KL is used for doing early stopping in the
+flags.DEFINE_float("target_kl", 0.01, "Policy iteration early stopping")
+
 
 def common_layers():
   cur_layers = []
   if FLAGS.flatten_non_batch_time_dims:
     cur_layers = [layers.Div(divisor=255.0), layers.Flatten(num_axis_to_keep=2)]
-  return cur_layers + [layers.Dense(16), layers.Relu(),
-                       layers.Dense(4), layers.Relu()]
-
-
-def run_training_loop():
-  """Run the PPO training loop."""
-
-  policy_net_fun = None
-  value_net_fun = None
-  policy_and_value_net_fun = None
-  policy_optimizer_fun = None
-  value_optimizer_fun = None
-  policy_and_value_optimizer_fun = None
-
-  if FLAGS.combined_policy_and_value_function:
-    policy_and_value_net_fun = functools.partial(
-        ppo.policy_and_value_net, bottom_layers=common_layers())
-    policy_and_value_optimizer_fun = get_optimizer_fun(
-        FLAGS.policy_and_value_net_learning_rate)
-  else:
-    policy_net_fun = functools.partial(ppo.policy_net,
-                                       bottom_layers=common_layers())
-    value_net_fun = functools.partial(ppo.value_net,
-                                      bottom_layers=common_layers())
-    policy_optimizer_fun = get_optimizer_fun(FLAGS.policy_net_learning_rate)
-    value_optimizer_fun = get_optimizer_fun(FLAGS.value_net_learning_rate)
-
-  ppo.training_loop(
-      env_name=FLAGS.env_name,
-      epochs=FLAGS.epochs,
-      policy_net_fun=policy_net_fun,
-      value_net_fun=value_net_fun,
-      policy_and_value_net_fun=policy_and_value_net_fun,
-      policy_optimizer_fun=policy_optimizer_fun,
-      value_optimizer_fun=value_optimizer_fun,
-      policy_and_value_optimizer_fun=policy_and_value_optimizer_fun,
+  return cur_layers + [
+      layers.Dense(64),
+      layers.Tanh(),
+      layers.Dense(64),
+      layers.Tanh()
+  ]
+
+
+def make_env():
+  """Creates the env."""
+  if FLAGS.env_name:
+    return gym.make(FLAGS.env_name)
+
+  assert FLAGS.env_problem_name
+
+  # No resizing needed, so let's be on the normal EnvProblem.
+  if not FLAGS.resize:  # None or False
+    return env_problem.EnvProblem(
+        base_env_name=FLAGS.env_problem_name,
+        batch_size=FLAGS.batch_size,
+        reward_range=(-1, 1))
+
+  wrapper_fn = functools.partial(
+      gym_utils.gym_env_wrapper, **{
+          "rl_env_max_episode_steps": FLAGS.max_timestep,
+          "maxskip_env": True,
+          "rendered_env": True,
+          "rendered_env_resize_to": (FLAGS.resized_height, FLAGS.resized_width),
+          "sticky_actions": False
+      })
+
+  return rendered_env_problem.RenderedEnvProblem(
+      base_env_name=FLAGS.env_problem_name,
       batch_size=FLAGS.batch_size,
-      num_optimizer_steps=FLAGS.num_optimizer_steps,
-      boundary=FLAGS.boundary,
-      max_timestep=FLAGS.max_timestep,
-      random_seed=FLAGS.random_seed)
+      env_wrapper_fn=wrapper_fn,
+      reward_range=(-1, 1))
 
 
 def get_optimizer_fun(learning_rate):
@@ -113,11 +162,55 @@ def main(argv):
   if FLAGS.jax_debug_nans:
     config.update("jax_debug_nans", True)
 
+  # Make an env here.
+  env = make_env()
+  assert env
+
+  def run_training_loop():
+    """Runs the training loop."""
+    policy_net_fun = None
+    value_net_fun = None
+    policy_and_value_net_fun = None
+    policy_optimizer_fun = None
+    value_optimizer_fun = None
+    policy_and_value_optimizer_fun = None
+
+    if FLAGS.combined_policy_and_value_function:
+      policy_and_value_net_fun = functools.partial(
+          ppo.policy_and_value_net, bottom_layers=common_layers())
+      policy_and_value_optimizer_fun = get_optimizer_fun(FLAGS.learning_rate)
+    else:
+      policy_net_fun = functools.partial(
+          ppo.policy_net, bottom_layers=common_layers())
+      value_net_fun = functools.partial(
+          ppo.value_net, bottom_layers=common_layers())
+      policy_optimizer_fun = get_optimizer_fun(FLAGS.policy_only_learning_rate)
+      value_optimizer_fun = get_optimizer_fun(FLAGS.value_only_learning_rate)
+
+    ppo.training_loop(
+        env=env,
+        epochs=FLAGS.epochs,
+        policy_net_fun=policy_net_fun,
+        value_net_fun=value_net_fun,
+        policy_and_value_net_fun=policy_and_value_net_fun,
+        policy_optimizer_fun=policy_optimizer_fun,
+        value_optimizer_fun=value_optimizer_fun,
+        policy_and_value_optimizer_fun=policy_and_value_optimizer_fun,
+        batch_size=FLAGS.batch_size,
+        num_optimizer_steps=FLAGS.num_optimizer_steps,
+        policy_only_num_optimizer_steps=FLAGS.policy_only_num_optimizer_steps,
+        value_only_num_optimizer_steps=FLAGS.value_only_num_optimizer_steps,
+        target_kl=FLAGS.target_kl,
+        boundary=FLAGS.boundary,
+        max_timestep=FLAGS.max_timestep,
+        random_seed=FLAGS.random_seed)
+
   if FLAGS.jax_debug_nans or FLAGS.disable_jit:
     with jax.disable_jit():
       run_training_loop()
   else:
     run_training_loop()
 
+
 if __name__ == "__main__":
   app.run(main)

From cca7982b971565f927f2c6773d42637b3c455729 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 3 May 2019 22:00:11 -0700
Subject: [PATCH 1993/2720] Typo.

PiperOrigin-RevId: 246622180
---
 tensor2tensor/trax/rlax/ppo_main.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index ac80982e9..ce5afc042 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -24,8 +24,7 @@
 RANDOM_SEED=0
 BOUNDARY=100
 
-
-blaze run --config=cuda third_party/py/tensor2tensor/trax/rlax:ppo_main -- \
+python trax/rlax/ppo_main.py \
   --env_problem_name=${ENV_PROBLEM_NAME} \
   --combined_policy_and_value_function=${COMBINED_NETWORK} \
   --epochs=${EPOCHS} \
@@ -34,7 +33,6 @@
   --boundary=${BOUNDARY} \
   --vmodule=*/tensor2tensor/*=1 \
   --alsologtostderr \
-
 """
 
 from __future__ import absolute_import

From 53d0582af625c232a36dee7417eede9846180b96 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Sat, 4 May 2019 19:16:38 -0700
Subject: [PATCH 1994/2720] Trax layers: rename Identity->Copy and make Branch
 take layers as input. Allow named tuples and add tests.

PiperOrigin-RevId: 246683951
---
 tensor2tensor/trax/layers/attention.py        |  31 ++-
 tensor2tensor/trax/layers/base.py             |  35 ++-
 tensor2tensor/trax/layers/combinators.py      | 249 ++++++++++++------
 tensor2tensor/trax/layers/combinators_test.py |  46 +++-
 tensor2tensor/trax/layers/core.py             |   5 +-
 tensor2tensor/trax/layers/rnn.py              |  30 ++-
 tensor2tensor/trax/models/mlp.py              |  12 +-
 tensor2tensor/trax/models/neural_gpu.py       |  20 +-
 tensor2tensor/trax/models/resnet.py           | 108 ++++----
 tensor2tensor/trax/models/transformer.py      | 228 ++++++++--------
 tensor2tensor/trax/rlax/ppo.py                |   3 +-
 tensor2tensor/trax/trax.py                    |  28 +-
 12 files changed, 486 insertions(+), 309 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index dcaef56c5..4fcd78c94 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -56,7 +56,7 @@ def EncoderDecoderMask(x, **unused_kwargs):
   padding_mask = np.reshape(
       padding_mask, (padding_mask.shape[0], 1, 1, padding_mask.shape[-1]))
   # Final mask shape is [batch, 1 for heads, decoder-len, encoder-len].
-  return padding_mask + np.ones((1, 1, decoder_input.shape[1], 1))
+  return padding_mask + np.zeros((1, 1, decoder_input.shape[1], 1))
 
 
 # Layer normalization.
@@ -170,7 +170,8 @@ def _multihead_attention_output_shape(  # pylint: disable=invalid-name
     input_shapes, **unused_kwargs):
   """Helper: calculate multihead attention output shape."""
   q_shape = input_shapes[0][0]  # Inputs are ((q, k, v), mask).
-  return q_shape
+  mask_shape = input_shapes[1]
+  return q_shape, mask_shape
 
 
 @base.layer(output_shape=_multihead_attention_output_shape)
@@ -187,7 +188,7 @@ def PureMultiHeadedAttention(x, params, num_heads=8, dropout=0.0,
     **kwargs: other arguments including the rng
 
   Returns:
-    Pure Multi-headed attention layer (no Dense transforms on input).
+    Pure Multi-headed attention result, and the mask.
   """
   del params
   rng = kwargs.get('rng', None)
@@ -205,10 +206,11 @@ def JoinHeads(x):  # pylint: disable=invalid-name
     return np.reshape(
         np.transpose(x, (0, 2, 1, 3)), (nbatch, -1, num_heads*head_depth))
   # Split heads, dot-product attention, rejoin heads.
-  return JoinHeads(
+  res = JoinHeads(
       DotProductAttention(
           SplitHeads(q), SplitHeads(k), SplitHeads(v), mask,
           dropout=dropout, mode=mode, rng=rng))
+  return res, mask  # Keep the mask.
 
 
 def MultiHeadedAttentionQKV(
@@ -224,7 +226,7 @@ def MultiHeadedAttentionQKV(
     mode: str: 'train' or 'eval'
 
   Returns:
-    Multi-headed self-attention layer.
+    Multi-headed self-attention result and the mask.
   """
   return combinators.Serial(
       combinators.Parallel(
@@ -233,12 +235,12 @@ def MultiHeadedAttentionQKV(
               core.Dense(feature_depth),
               core.Dense(feature_depth),
           ),
-          combinators.Identity()
+          combinators.Copy()
       ),
       PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
           feature_depth=feature_depth, num_heads=num_heads,
           dropout=dropout, mode=mode),
-      core.Dense(feature_depth),
+      combinators.Parallel(core.Dense(feature_depth), combinators.Copy())
   )
 
 
@@ -259,8 +261,10 @@ def MultiHeadedAttention(
   """
   return combinators.Serial(
       combinators.Parallel(
-          combinators.Branch(num_branches=3),  # q = k = v = first input
-          combinators.Identity()  # pass the mask
+          # q = k = v = first input
+          combinators.Branch(
+              combinators.Copy(), combinators.Copy(), combinators.Copy()),
+          combinators.Copy()  # pass the mask
       ),
       MultiHeadedAttentionQKV(  # pylint: disable=no-value-for-parameter
           feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),
@@ -350,9 +354,9 @@ def ChunkedCausalMultiHeadedAttention(
     Multi-headed self-attention layer.
   """
   prepare_attention_input = combinators.Serial(
-      combinators.Branch(),
-      combinators.Parallel(
-          combinators.Branch(num_branches=3),  # q = k = v = first input
+      combinators.Branch(
+          combinators.Branch(  # q = k = v = first input
+              combinators.Copy(), combinators.Copy(), combinators.Copy()),
           CausalMask(axis=-2),  # pylint: disable=no-value-for-parameter
       ),
       combinators.Parallel(
@@ -361,7 +365,7 @@ def ChunkedCausalMultiHeadedAttention(
               core.Dense(feature_depth),
               core.Dense(feature_depth),
           ),
-          combinators.Identity()
+          combinators.Copy()
       )
   )
   return combinators.Serial(
@@ -370,6 +374,7 @@ def ChunkedCausalMultiHeadedAttention(
       combinators.Map(PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
           feature_depth=feature_depth, num_heads=num_heads,
           dropout=dropout, mode=mode), check_shapes=False),
+      combinators.Map(combinators.Select(0), check_shapes=False),  # drop masks
       combinators.Map(core.Dense(feature_depth))
   )
 
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index a0ac29d0a..bd56d554f 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -74,6 +74,14 @@ def new_parameters(self, input_shape, rng):
 
   # End of subclassing interface, all functions below are internal.
 
+  def output_shape_catch_errors(self, input_shape):
+    """Same as self.output_shape but with better error reporting."""
+    try:
+      return self.output_shape(input_shape)
+    except Exception:
+      name, trace = self.__class__.__name__, _short_traceback()
+      raise LayerError(name, 'output_shape', self._caller, input_shape, trace)
+
   def initialize(self, input_shape, rng):
     """Initialize the layer given an input shape and rng.
 
@@ -280,13 +288,34 @@ def call_fun(self, x, params=(), **kwargs):
 
 
 def _random_inputs(input_shape, rng, integer_inputs=False):
-  """Create random floats of the given shape."""
-  if isinstance(input_shape[0], int):  # Non-nested shape.
+  """Create random floats of the given shape.
+
+  Args:
+    input_shape: Could be either:
+        list/tuple of ints, ex: (210, 160, 3) or
+        list/tuple of nested shapes, ex: [(210, 160, 3), (105, 80, 3)] or
+        dictionary of nested shapes, ex: {"obs": [(28, 28, 1), (4,)],
+                                          "sensors": [(3,4), (4, 9)]} or
+        any other combination of these, ex: list of dictionaries of tuples etc.
+    rng: random number generator.
+    integer_inputs: boolean, True if we want arrays of integers, otherwise we
+        produce float32s.
+
+  Returns:
+    Random values of the type and shape specified.
+  """
+  if not isinstance(input_shape, dict) and isinstance(input_shape[0], int):
+    # Non-nested shape, create a random tuple.
     if not integer_inputs:
       return random.uniform(rng, input_shape, minval=-1.0, maxval=1.0)
     return random.bernoulli(rng, 0.5, input_shape).astype(onp.int32)
-  elif isinstance(input_shape, (list, tuple)):  # Nested shape.
+  elif isinstance(input_shape, list):  # Nested shape: list.
     return [_random_inputs(shape, rng, integer_inputs) for shape in input_shape]
+  elif isinstance(input_shape, tuple):  # Nested shape: tuple.
+    return tuple(_random_inputs(list(input_shape), rng, integer_inputs))
+  elif isinstance(input_shape, dict):  # Nested shape: dict.
+    return {k: _random_inputs(input_shape[k], rng, integer_inputs)
+            for k in input_shape}
   else:
     raise TypeError(type(input_shape))
 
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 5889d6d98..7b789276e 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -46,16 +46,7 @@ def call(self, x, params=(), **kwargs):
   def output_shape(self, input_shape):
     cur_shape = input_shape
     for layer in self._layers:
-      try:
-        cur_shape = layer.output_shape(cur_shape)
-      except Exception:
-        # Since this is a widely used combinator, we improve errors here.
-        # Private methods are accessed as an exception for that reason.
-        name, trace = layer.__class__.__name__, base._short_traceback()  # pylint: disable=protected-access
-        raise base.LayerError(
-            name, 'output_shape',
-            layer._caller, input_shape, trace)  # pylint: disable=protected-access
-
+      cur_shape = layer.output_shape_catch_errors(cur_shape)
     return cur_shape
 
   def new_parameters(self, input_shape, rng):
@@ -70,8 +61,8 @@ def new_parameters(self, input_shape, rng):
 
 
 @base.layer()
-def Identity(x, **unused_kwargs):
-  """Identity layer, return the inputs."""
+def Copy(x, **unused_kwargs):
+  """Copy layer, return the inputs."""
   return x
 
 
@@ -107,16 +98,8 @@ def UnnestBranches(x, **unused_kwargs):
 
 
 # Re-ordering layer.
-def _reorder_shape(input_shape, output=None):  # pylint: disable=invalid-name
-  """Helper to determine the shape of reorder output."""
-  if output is None:
-    return input_shape
-  return base.nested_map(output, lambda i: input_shape[i])
-
-
-@base.layer(output_shape=_reorder_shape)
-def Reorder(x, params, output=None, **kwargs):
-  """Reorder a tuple into another tuple.
+class Select(base.Layer):
+  """Select elements from a tuple or create another tuple from them.
 
   For example, we can re-order (x, y) into (y, x) or even (y, (x, y), y).
   The output argument specifies how to re-order, using integers that refer
@@ -126,12 +109,13 @@ def Reorder(x, params, output=None, **kwargs):
 
   then
 
-    Reorder(input, output=(1, 0, 2))   = (y, x, z)
-    Reorder(input, output=(0, 0))      = (x, x)
-    Reorder(input, output=(0, (1, 1))) = (x, (y, y))
-    Reorder(input, output=((2, 0), (1, 1))) = ((z, x), (y, y))
+    Select(0)                = x
+    Select((1, 0, 2))        = (y, x, z)
+    Select((0, 0))           = (x, x)
+    Select((0, (1, 1)))      = (x, (y, y))
+    Select(((2, 0), (1, 1))) = ((z, x), (y, y))
 
-  By default (if no output is given) Reorder does nothing (Identity).
+  By default (if no output is given) Select does nothing (Copy).
 
   Args:
     x: the input tuple to re-order.
@@ -142,35 +126,102 @@ def Reorder(x, params, output=None, **kwargs):
   Returns:
     The re-ordered tuple with the same shape as output.
   """
-  del params, kwargs
-  if output is None:
-    return x
-  return base.nested_map(output, lambda i: x[i])
 
+  def __init__(self, output=None):
+    super(Select, self).__init__()
+    self._output = output
 
-@base.layer(output_shape=lambda shape, num_branches=2: [shape] * num_branches)
-def Branch(x, params, num_branches=2, **kwargs):
-  del params, kwargs
-  return [x] * num_branches
+  def call(self, x, params=(), **kwargs):
+    del params, kwargs
+    if self._output is None:
+      return x
+    return base.nested_map(self._output, lambda i: x[i])
 
+  def output_shape(self, input_shape):
+    if self._output is None:
+      return input_shape
+    return base.nested_map(self._output, lambda i: input_shape[i])
 
-@base.layer(output_shape=lambda input_shape_list: input_shape_list[0])
-def FirstBranch(x, **unused_kwargs):
-  return x[0]  # Here x is a list of tensors, we select the first.
+  def new_parameters(self, input_shape, rng):
+    return ()
 
 
-@base.layer(output_shape=lambda input_shape_list: input_shape_list[1])
-def SecondBranch(x, **unused_kwargs):
-  return x[1]  # Here x is a list of tensors, we select the second.
+class Branch(base.Layer):
+  """Combinator for applying layers to copies of the input.
 
+  This layer is often used to create parallel towers in neural networks:
+  * Branch(Copy(), Copy()) -- creates a pair with copied input
+  * Branch(main, shortcut) -- start a residual tower (see Residual below)
 
-@base.layer(output_shape=lambda input_shape_list: input_shape_list[2])
-def ThirdBranch(x, **unused_kwargs):
-  return x[2]  # Here x is a list of tensors, we select the third.
+  Args:
+    *layers: a sequence of layers.
+    **kwlayers: a dictionary of layers.
+
+  Returns:
+    A new layer in which each of the given layers has been applied to
+    a copy of the input independently.
+  """
+
+  def __init__(self, *layers, **kwlayers):
+    super(Branch, self).__init__()
+    if layers and kwlayers:
+      raise ValueError('Cannot specify a Branch with both a list and dict.')
+    layers = layers or kwlayers
+    self._nlayers = len(layers)
+    self._layers = layers
+
+  def call(self, x, params=(), **kwargs):
+    # Split the random number generators.
+    rng = kwargs.pop('rng', None)
+    rngs = (None,) * self._nlayers
+    if rng is not None:
+      rngs = backend.random.split(rng, self._nlayers)
+    # If layers are a list or a tuple, just apply them.
+    if isinstance(self._layers, (list, tuple)):
+      res = [layer(x, params=p, rng=r, **kwargs)
+             for layer, p, r in zip(self._layers, params, rngs)]
+      return tuple(res)
+    # If layers are a dictionary, apply to matching keys.
+    assert isinstance(self._layers, dict)
+    result, counter = {}, 0
+    for k in self._layers:
+      result[k] = self._layers[k](
+          x, params=params[k], rng=rngs[counter], **kwargs)
+      counter += 1
+    return result
+
+  def output_shape(self, input_shape):
+    output_shapes = []
+    # If the argument layers are a sequence, apply each to calculate shape.
+    if not isinstance(self._layers, dict):
+      for layer in self._layers:
+        output_shapes.append(layer.output_shape_catch_errors(input_shape))
+      return tuple(output_shapes)
+    # If layers are a dictionary, apply to the input shape.
+    result = {}
+    for k in self._layers:
+      result[k] = self._layers[k].output_shape_catch_errors(input_shape)
+    return result
+
+  def new_parameters(self, input_shape, rng):
+    rngs = backend.random.split(rng, self._nlayers)
+    # If the argument layers are a sequence, create parameters for each one.
+    if not isinstance(self._layers, dict):
+      return [layer.initialize(input_shape, rng) for layer, rng
+              in zip(self._layers, rngs)]
+    # If the argument layers are a dictionary, create a dictionary too.
+    result, counter = {}, 0
+    for k in self._layers:
+      result[k] = self._layers[k].initialize(input_shape, rngs[counter])
+      counter += 1
+    return result
 
 
 def _nested_op(inputs, op):  # pylint: disable=invalid-name
-  """Helper: sum a list of arrays or nested arrays."""
+  """Helper: apply op over a list of arrays or nested arrays."""
+  # If input is a dictionary, apply to the values (ignore keys).
+  if isinstance(inputs, dict):
+    return _nested_op(list(inputs.values()), op)
   # First the simple non-nested case.
   if not isinstance(inputs[0], (list, tuple)):
     return op(inputs)
@@ -192,21 +243,27 @@ def _nested_product(inputs):  # pylint: disable=invalid-name
       inputs=inputs, op=lambda xs: six.moves.reduce(operator.mul, xs))
 
 
-@base.layer(output_shape=lambda input_shape_list: input_shape_list[0])
-def SumBranches(x, **unused_kwargs):
-  """Sum branches elementwise."""
+def _first_from_tuple_or_dict(tuple_or_dict):  # pylint: disable=invalid-name
+  """Helper: return the first element from a tuple or dict."""
+  for x in tuple_or_dict:
+    return x
+
+
+@base.layer(output_shape=_first_from_tuple_or_dict)
+def Add(x, **unused_kwargs):
+  """Add branches elementwise."""
   # Here x is a list of tensors of the same shape, or nested structures.
   return _nested_sum(x)
 
 
-@base.layer(output_shape=lambda input_shape_list: input_shape_list[0])
-def MultiplyBranches(x, **unused_kwargs):
+@base.layer(output_shape=_first_from_tuple_or_dict)
+def Multiply(x, **unused_kwargs):
   """Multiply branches elementwise."""
   return _nested_product(x)
 
 
-@base.layer(output_shape=lambda input_shape_list: input_shape_list[0])
-def GateBranches(x, **unused_kwargs):
+@base.layer(output_shape=_first_from_tuple_or_dict)
+def Gate(x, **unused_kwargs):
   """Implements a gating function on a (memory, gate, candidate) tuple.
 
   Final update is memory * gate + (1-gate) * candidate
@@ -227,6 +284,8 @@ def GateBranches(x, **unused_kwargs):
 
 def _concatenate_shape(input_shape, axis=-1):  # pylint: disable=invalid-name
   """Helper to determine the shape of Concatenate output."""
+  if isinstance(input_shape, dict):  # For named tuples, just use the values.
+    input_shape = list(input_shape.values())
   ax = axis % len(input_shape[0])
   concat_size = sum(shape[ax] for shape in input_shape)
   out_shape = input_shape[0][:ax] + (concat_size,) + input_shape[0][ax+1:]
@@ -236,70 +295,100 @@ def _concatenate_shape(input_shape, axis=-1):  # pylint: disable=invalid-name
 @base.layer(output_shape=_concatenate_shape)
 def Concatenate(x, params, axis=-1, **kwargs):
   del params, kwargs
+  if isinstance(x, dict):  # For dictionaries, just use the values.
+    x = list(x.values())
   return backend.numpy.concatenate(x, axis)
 
 
 class Parallel(base.Layer):
-  """Combinator for composing layers in parallel.
+  """Combinator for applying layers to parts of a tuple.
 
-  This layer is often used with the Branch and SumBranches layers.
+  This layer is often used with the Branch and Add layers.
 
   Args:
     *layers: a sequence of layers.
+    **kwlayers: a dictionary of layers.
 
   Returns:
-    A new layer representing parallel composition of the given layers.
-    The new layer takes a sequence of inputs and returns a sequence of outputs
-    with the same length as the argument `layers`.
+    A new layer in which each of the given layers has been applied to
+    its corresponding argument in the input tuple or dictionary.
   """
 
-  def __init__(self, *layers):
+  def __init__(self, *layers, **kwlayers):
     super(Parallel, self).__init__()
+    if layers and kwlayers:
+      raise ValueError('Cannot specify a Parallel with both a list and dict.')
+    layers = layers or kwlayers
     self._nlayers = len(layers)
     self._layers = layers
 
   def call(self, inputs, params=(), **kwargs):
+    # Split the random number generators.
     rng = kwargs.pop('rng', None)
     rngs = (None,) * self._nlayers
     if rng is not None:
       rngs = backend.random.split(rng, self._nlayers)
-    return [layer(x, params=p, rng=r, **kwargs)
-            for layer, x, p, r in zip(self._layers, inputs, params, rngs)]
+    # If layers are a list or a tuple, just apply them.
+    if not isinstance(self._layers, dict):
+      res = [layer(x, params=p, rng=r, **kwargs)
+             for layer, x, p, r in zip(self._layers, inputs, params, rngs)]
+      # Return a list if inputs are a list and a tuple if inputs are a tuple.
+      if isinstance(inputs, list):
+        return res
+      return tuple(res)
+    # If layers are a dictionary, apply to matching keys.
+    result, counter = {}, 0
+    for k in inputs:
+      if k in self._layers:
+        result[k] = self._layers[k](
+            inputs[k], params=params[k], rng=rngs[counter], **kwargs)
+        counter += 1
+      else:
+        result[k] = inputs[k]
+    return result
 
   def output_shape(self, input_shape):
     output_shapes = []
-    for i, layer in enumerate(self._layers):
-      try:
-        output_shapes.append(layer.output_shape(input_shape[i]))
-      except Exception:
-        # Since this is a widely used combinator, we improve errors here.
-        # Private methods are accessed as an exception for that reason.
-        name, trace = layer.__class__.__name__, base._short_traceback()  # pylint: disable=protected-access
-        raise base.LayerError(
-            name, 'output_shape',
-            layer._caller, input_shape[i], trace)  # pylint: disable=protected-access
-    return tuple(output_shapes)
+    # If the argument layers are a sequence, apply each to calculate shape.
+    if not isinstance(self._layers, dict):
+      for i, layer in enumerate(self._layers):
+        output_shapes.append(layer.output_shape_catch_errors(input_shape[i]))
+      return tuple(output_shapes)
+    # If layers are a dictionary, apply to matching keys in the input shape.
+    result = {}
+    for k in input_shape:
+      if k in self._layers:
+        result[k] = self._layers[k].output_shape_catch_errors(input_shape[k])
+      else:
+        result[k] = input_shape[k]
+    return result
 
   def new_parameters(self, input_shape, rng):
     rngs = backend.random.split(rng, self._nlayers)
-    return [layer.initialize(shape, rng) for layer, shape, rng
-            in zip(self._layers, input_shape, rngs)]
+    # If the argument layers are a sequence, create parameters for each one.
+    if not isinstance(self._layers, dict):
+      return [layer.initialize(shape, rng) for layer, shape, rng
+              in zip(self._layers, input_shape, rngs)]
+    # If the argument layers are a dictionary, create a dictionary too.
+    result, counter = {}, 0
+    for k in self._layers:
+      result[k] = self._layers[k].initialize(input_shape[k], rngs[counter])
+      counter += 1
+    return result
 
 
 def Residual(*layers, **kwargs):
   """Constructs a residual version of layers, summing input to layers output."""
-  shortcut = kwargs.get('shortcut', Identity())  # pylint: disable=no-value-for-parameter
+  shortcut = kwargs.get('shortcut', Copy())  # pylint: disable=no-value-for-parameter
   if len(layers) > 1:
     return Serial(
-        Branch(),  # pylint: disable=no-value-for-parameter
-        Parallel(Serial(*layers), shortcut),
-        SumBranches()  # pylint: disable=no-value-for-parameter
+        Branch(Serial(*layers), shortcut),
+        Add()  # pylint: disable=no-value-for-parameter
     )
   elif len(layers) == 1:
     return Serial(
-        Branch(),  # pylint: disable=no-value-for-parameter
-        Parallel(layers[0], shortcut),
-        SumBranches()  # pylint: disable=no-value-for-parameter
+        Branch(layers[0], shortcut),
+        Add()  # pylint: disable=no-value-for-parameter
     )
   else:
     raise ValueError('Empty residual combinator.')
diff --git a/tensor2tensor/trax/layers/combinators_test.py b/tensor2tensor/trax/layers/combinators_test.py
index 7cf4c0ce0..e8412304f 100644
--- a/tensor2tensor/trax/layers/combinators_test.py
+++ b/tensor2tensor/trax/layers/combinators_test.py
@@ -25,6 +25,50 @@
 
 class CombinatorLayerTest(absltest.TestCase):
 
+  def test_branch(self):
+    input_shape = (2, 3)
+    expected_shape = ((2, 3), (2, 3))
+    output_shape = base.check_shape_agreement(
+        combinators.Branch(combinators.Copy(), combinators.Copy()), input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+  def test_branch_named(self):
+    input_shape = (2, 3)
+    expected_shape = {'a': (2, 3), 'b': (2, 3)}
+    output_shape = base.check_shape_agreement(
+        combinators.Branch(a=combinators.Copy(), b=combinators.Copy()),
+        input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+  def test_parallel(self):
+    input_shape = ((2, 3), (2, 3))
+    expected_shape = ((2, 3), (2, 3))
+    output_shape = base.check_shape_agreement(
+        combinators.Parallel(combinators.Copy(), combinators.Copy()),
+        input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+  def test_parallel_named(self):
+    input_shape = {'a': (2, 3), 'b': (2, 3)}
+    expected_shape = {'a': (2, 3), 'b': (2, 3)}
+    output_shape = base.check_shape_agreement(
+        combinators.Parallel(a=combinators.Copy()), input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+  def test_select(self):
+    input_shape = ((2, 3), (3, 4))
+    expected_shape = (3, 4)
+    output_shape = base.check_shape_agreement(
+        combinators.Select(1), input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+  def test_select_named(self):
+    input_shape = {'a': (2, 3), 'b': (3, 4)}
+    expected_shape = (3, 4)
+    output_shape = base.check_shape_agreement(
+        combinators.Select('b'), input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
   def test_unnest_branches(self):
     input_shape = ((2, 3), [(4, 5), (6, 7)], (8, 9, 10))
     expected_shape = ((2, 3), (4, 5), (6, 7), (8, 9, 10))
@@ -33,5 +77,5 @@ def test_unnest_branches(self):
     self.assertEqual(output_shape, expected_shape)
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
   absltest.main()
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 257bf84b0..31f3890d1 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -151,8 +151,9 @@ def output_shape(self, input_shape):
     return tuple(input_shape[:-1]) + (self._units,)
 
   def new_parameters(self, input_shape, rng):
-    w = self._kernel_initializer((input_shape[-1], self._units), rng)
-    b = self._bias_initializer((self._units,), rng)
+    rng1, rng2 = backend.random.split(rng, 2)
+    w = self._kernel_initializer((input_shape[-1], self._units), rng1)
+    b = self._bias_initializer((self._units,), rng2)
     return (w, b)
 
 
diff --git a/tensor2tensor/trax/layers/rnn.py b/tensor2tensor/trax/layers/rnn.py
index 13df5ae75..9bfd60d83 100644
--- a/tensor2tensor/trax/layers/rnn.py
+++ b/tensor2tensor/trax/layers/rnn.py
@@ -37,7 +37,7 @@ def GRUCell(units):
   """
   return GeneralGRUCell(
       candidate_transform=lambda: core.Dense(units=units),
-      memory_transform=combinators.Identity,
+      memory_transform=combinators.Copy,
       gate_nonlinearity=core.Sigmoid,
       candidate_nonlinearity=core.Tanh)
 
@@ -60,13 +60,13 @@ def BuildConv():
 
   return GeneralGRUCell(
       candidate_transform=BuildConv,
-      memory_transform=combinators.Identity,
+      memory_transform=combinators.Copy,
       gate_nonlinearity=core.Sigmoid,
       candidate_nonlinearity=core.Tanh)
 
 
 def GeneralGRUCell(candidate_transform,
-                   memory_transform=combinators.Identity,
+                   memory_transform=combinators.Copy,
                    gate_nonlinearity=core.Sigmoid,
                    candidate_nonlinearity=core.Tanh,
                    dropout_rate_c=0.1,
@@ -79,7 +79,7 @@ def GeneralGRUCell(candidate_transform,
   $$ Candidate memory: c_t = \tanh(U * (r_t \odot s_{t-1}) + B) $$
   $$ New State: s_t = u_t \odot s_{t-1} + (1 - u_t) \odot c_t $$
 
-  See combinators.GateBranches for details on the gating function.
+  See combinators.Gate for details on the gating function.
 
 
   Args:
@@ -99,8 +99,7 @@ def GeneralGRUCell(candidate_transform,
     A model representing a GRU cell with specified transforms.
   """
   return combinators.Serial(
-      combinators.Branch(num_branches=3),
-      combinators.Parallel(
+      combinators.Branch(
           # s_{t-1} branch - optionally transform
           # Typically is an identity.
           memory_transform(),
@@ -110,21 +109,23 @@ def GeneralGRUCell(candidate_transform,
               candidate_transform(),
               # Want bias to start out positive before sigmoids.
               core.AddConstant(constant=sigmoid_bias),
-              gate_nonlinearity()),
+              gate_nonlinearity()
+          ),
 
           # c_t (Candidate) branch
           combinators.Serial(
-              combinators.Branch(num_branches=2),
-              combinators.Parallel(
-                  combinators.Identity(),
+              combinators.Branch(
+                  combinators.Copy(),
                   # r_t (Reset) Branch
                   combinators.Serial(
                       candidate_transform(),
                       # Want bias to start out positive before sigmoids.
                       core.AddConstant(constant=sigmoid_bias),
-                      gate_nonlinearity())),
+                      gate_nonlinearity()
+                  )
+              ),
               ## Gate S{t-1} with sigmoid(candidate_transform(S{t-1}))
-              combinators.MultiplyBranches(),
+              combinators.Multiply(),
 
               # Final projection + tanh to get Ct
               candidate_transform(),
@@ -132,7 +133,8 @@ def GeneralGRUCell(candidate_transform,
 
               # Only apply dropout on the C gate.
               # Paper reports that 0.1 is a good default.
-              core.Dropout(rate=dropout_rate_c)),
+              core.Dropout(rate=dropout_rate_c)
+          ),
       ),
       # Gate memory and candidate
-      combinators.GateBranches())
+      combinators.Gate())
diff --git a/tensor2tensor/trax/models/mlp.py b/tensor2tensor/trax/models/mlp.py
index f93536821..fa248e4c3 100644
--- a/tensor2tensor/trax/models/mlp.py
+++ b/tensor2tensor/trax/models/mlp.py
@@ -19,18 +19,18 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.trax import layers
+from tensor2tensor.trax import layers as tl
 
 
 def MLP(num_hidden_layers=2,
         hidden_size=512,
-        activation_fn=layers.Relu,
+        activation_fn=tl.Relu,
         num_output_classes=10,
         mode="train"):
   """Multi-layer feed-forward neural network with non-linear activations."""
   del mode
-  cur_layers = [layers.Flatten()]
+  cur_layers = [tl.Flatten()]
   for _ in range(num_hidden_layers):
-    cur_layers += [layers.Dense(hidden_size), activation_fn()]
-  cur_layers += [layers.Dense(num_output_classes), layers.LogSoftmax()]
-  return layers.Serial(*cur_layers)
+    cur_layers += [tl.Dense(hidden_size), activation_fn()]
+  cur_layers += [tl.Dense(num_output_classes), tl.LogSoftmax()]
+  return tl.Serial(*cur_layers)
diff --git a/tensor2tensor/trax/models/neural_gpu.py b/tensor2tensor/trax/models/neural_gpu.py
index 46e9cc9fa..721f35c79 100644
--- a/tensor2tensor/trax/models/neural_gpu.py
+++ b/tensor2tensor/trax/models/neural_gpu.py
@@ -20,7 +20,7 @@
 from __future__ import google_type_annotations
 from __future__ import print_function
 
-from tensor2tensor.trax import layers
+from tensor2tensor.trax import layers as tl
 from tensor2tensor.trax.backend import numpy as np
 
 
@@ -29,7 +29,7 @@ def SaturationCost(x, limit=0.9):
   return np.minimum(0, np.abs(x) - limit)
 
 
-@layers.layer(output_shape=lambda input_shape_list: input_shape_list)
+@tl.layer(output_shape=lambda input_shape_list: input_shape_list)
 def DiagonalGate(x, params, **kwargs):
   """Split channels in 3 parts. Shifts 1st and 3rd sections to left/right."""
   del params
@@ -51,13 +51,13 @@ def ConvDiagonalGRU(units, kernel_size=(3, 3)):
   """Build convolutional GRU with diagonal gating as in ImprovedNGPU."""
 
   def BuildConv():
-    return layers.Conv(filters=units, kernel_size=kernel_size, padding='SAME')
+    return tl.Conv(filters=units, kernel_size=kernel_size, padding='SAME')
 
-  return layers.GeneralGRUCell(
+  return tl.GeneralGRUCell(
       candidate_transform=BuildConv,
       memory_transform=DiagonalGate,
-      gate_nonlinearity=layers.HardSigmoid,
-      candidate_nonlinearity=layers.HardTanh)
+      gate_nonlinearity=tl.HardSigmoid,
+      candidate_nonlinearity=tl.HardTanh)
 
 
 def NeuralGPU(feature_depth=96, steps=16, vocab_size=2):
@@ -73,10 +73,10 @@ def NeuralGPU(feature_depth=96, steps=16, vocab_size=2):
   """
   xs = []
   xs.append(
-      layers.Embedding(feature_depth=feature_depth, vocab_size=vocab_size))
+      tl.Embedding(feature_depth=feature_depth, vocab_size=vocab_size))
   core = ConvDiagonalGRU(units=feature_depth)
   xs.extend([core] * steps)
-  xs.append(layers.Dense(vocab_size))
-  xs.append(layers.LogSoftmax())
+  xs.append(tl.Dense(vocab_size))
+  xs.append(tl.LogSoftmax())
 
-  return layers.Serial(*xs)
+  return tl.Serial(*xs)
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index aa6316792..2b7f801bc 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -19,32 +19,30 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.trax import layers
+from tensor2tensor.trax import layers as tl
 
 
 def ConvBlock(kernel_size, filters, strides):
   """ResNet convolutional striding block."""
   ks = kernel_size
   filters1, filters2, filters3 = filters
-  main = layers.Serial(
-      layers.Conv(filters1, (1, 1), strides),
-      layers.BatchNorm(),
-      layers.Relu(),
-      layers.Conv(filters2, (ks, ks), padding='SAME'),
-      layers.BatchNorm(),
-      layers.Relu(),
-      layers.Conv(filters3, (1, 1)),
-      layers.BatchNorm()
+  main = tl.Serial(
+      tl.Conv(filters1, (1, 1), strides),
+      tl.BatchNorm(),
+      tl.Relu(),
+      tl.Conv(filters2, (ks, ks), padding='SAME'),
+      tl.BatchNorm(),
+      tl.Relu(),
+      tl.Conv(filters3, (1, 1)),
+      tl.BatchNorm()
   )
-  shortcut = layers.Serial(
-      layers.Conv(filters3, (1, 1), strides),
-      layers.BatchNorm()
+  shortcut = tl.Serial(
+      tl.Conv(filters3, (1, 1), strides),
+      tl.BatchNorm()
   )
-  return layers.Serial(
-      layers.Branch(),
-      layers.Parallel(main, shortcut),
-      layers.SumBranches(),
-      layers.Relu()
+  return tl.Serial(
+      tl.Residual(main, shortcut=shortcut),
+      tl.Relu()
   )
 
 
@@ -52,21 +50,19 @@ def IdentityBlock(kernel_size, filters):
   """ResNet identical size block."""
   ks = kernel_size
   filters1, filters2, filters3 = filters
-  main = layers.Serial(
-      layers.Conv(filters1, (1, 1)),
-      layers.BatchNorm(),
-      layers.Relu(),
-      layers.Conv(filters2, (ks, ks), padding='SAME'),
-      layers.BatchNorm(),
-      layers.Relu(),
-      layers.Conv(filters3, (1, 1)),
-      layers.BatchNorm()
+  main = tl.Serial(
+      tl.Conv(filters1, (1, 1)),
+      tl.BatchNorm(),
+      tl.Relu(),
+      tl.Conv(filters2, (ks, ks), padding='SAME'),
+      tl.BatchNorm(),
+      tl.Relu(),
+      tl.Conv(filters3, (1, 1)),
+      tl.BatchNorm()
   )
-  return layers.Serial(
-      layers.Branch(),
-      layers.Parallel(main, layers.Identity()),
-      layers.SumBranches(),
-      layers.Relu()
+  return tl.Serial(
+      tl.Residual(main),
+      tl.Relu()
   )
 
 
@@ -82,10 +78,10 @@ def Resnet50(hidden_size=64, num_output_classes=1001, mode='train'):
     The ResNet model with the given layer and output sizes.
   """
   del mode
-  return layers.Serial(
-      layers.Conv(hidden_size, (7, 7), (2, 2), 'SAME'),
-      layers.BatchNorm(), layers.Relu(),
-      layers.MaxPool(pool_size=(3, 3), strides=(2, 2)),
+  return tl.Serial(
+      tl.Conv(hidden_size, (7, 7), (2, 2), 'SAME'),
+      tl.BatchNorm(), tl.Relu(),
+      tl.MaxPool(pool_size=(3, 3), strides=(2, 2)),
       ConvBlock(3, [hidden_size, hidden_size, 4 * hidden_size], (1, 1)),
       IdentityBlock(3, [hidden_size, hidden_size, 4 * hidden_size]),
       IdentityBlock(3, [hidden_size, hidden_size, 4 * hidden_size]),
@@ -102,20 +98,25 @@ def Resnet50(hidden_size=64, num_output_classes=1001, mode='train'):
       ConvBlock(3, [8 * hidden_size, 8 * hidden_size, 32*hidden_size], (2, 2)),
       IdentityBlock(3, [8 * hidden_size, 8 * hidden_size, 32 * hidden_size]),
       IdentityBlock(3, [8 * hidden_size, 8 * hidden_size, 32 * hidden_size]),
-      layers.AvgPool(pool_size=(7, 7)), layers.Flatten(),
-      layers.Dense(num_output_classes), layers.LogSoftmax())
+      tl.AvgPool(pool_size=(7, 7)),
+      tl.Flatten(),
+      tl.Dense(num_output_classes),
+      tl.LogSoftmax()
+  )
 
 
 def WideResnetBlock(channels, strides=(1, 1), channel_mismatch=False):
   """WideResnet convolutational block."""
-  main = layers.Serial(layers.BatchNorm(), layers.Relu(),
-                       layers.Conv(channels, (3, 3), strides, padding='SAME'),
-                       layers.BatchNorm(), layers.Relu(),
-                       layers.Conv(channels, (3, 3), padding='SAME'))
-  shortcut = layers.Identity() if not channel_mismatch else layers.Conv(
+  main = tl.Serial(
+      tl.BatchNorm(),
+      tl.Relu(),
+      tl.Conv(channels, (3, 3), strides, padding='SAME'),
+      tl.BatchNorm(),
+      tl.Relu(),
+      tl.Conv(channels, (3, 3), padding='SAME'))
+  shortcut = tl.Copy() if not channel_mismatch else tl.Conv(
       channels, (3, 3), strides, padding='SAME')
-  return layers.Serial(
-      layers.Branch(), layers.Parallel(main, shortcut), layers.SumBranches())
+  return tl.Residual(main, shortcut=shortcut)
 
 
 def WideResnetGroup(n, channels, strides=(1, 1)):
@@ -123,7 +124,7 @@ def WideResnetGroup(n, channels, strides=(1, 1)):
   blocks += [WideResnetBlock(channels, strides, channel_mismatch=True)]
   for _ in range(n - 1):
     blocks += [WideResnetBlock(channels, (1, 1))]
-  return layers.Serial(*blocks)
+  return tl.Serial(*blocks)
 
 
 def WideResnet(num_blocks=3, hidden_size=64, num_output_classes=10,
@@ -140,10 +141,15 @@ def WideResnet(num_blocks=3, hidden_size=64, num_output_classes=10,
     The WideResnet model with given layer and output sizes.
   """
   del mode
-  return layers.Serial(
-      layers.Conv(hidden_size, (3, 3), padding='SAME'),
+  return tl.Serial(
+      tl.Conv(hidden_size, (3, 3), padding='SAME'),
       WideResnetGroup(num_blocks, hidden_size),
       WideResnetGroup(num_blocks, hidden_size * 2, (2, 2)),
-      WideResnetGroup(num_blocks, hidden_size * 4, (2, 2)), layers.BatchNorm(),
-      layers.Relu(), layers.AvgPool(pool_size=(8, 8)), layers.Flatten(),
-      layers.Dense(num_output_classes), layers.LogSoftmax())
+      WideResnetGroup(num_blocks, hidden_size * 4, (2, 2)),
+      tl.BatchNorm(),
+      tl.Relu(),
+      tl.AvgPool(pool_size=(8, 8)),
+      tl.Flatten(),
+      tl.Dense(num_output_classes),
+      tl.LogSoftmax()
+  )
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 98c5d3242..85e5997bf 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -18,7 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.trax import layers
+from tensor2tensor.trax import layers as tl
 
 
 def ResidualFeedForward(feature_depth,
@@ -26,13 +26,13 @@ def ResidualFeedForward(feature_depth,
                         dropout,
                         mode):
   """Residual feed-forward layer with normalization at start."""
-  return layers.Residual(
-      layers.LayerNorm(),
-      layers.Dense(feedforward_depth),
-      layers.Relu(),
-      layers.Dropout(rate=dropout, mode=mode),
-      layers.Dense(feature_depth),
-      layers.Dropout(rate=dropout, mode=mode)
+  return tl.Residual(
+      tl.LayerNorm(),
+      tl.Dense(feedforward_depth),
+      tl.Relu(),
+      tl.Dropout(rate=dropout, mode=mode),
+      tl.Dense(feature_depth),
+      tl.Dropout(rate=dropout, mode=mode)
   )
 
 
@@ -57,22 +57,18 @@ def EncoderLayer(feature_depth,
   Returns:
     the layer, returning a pair (actiavtions, mask).
   """
-  # The encoder block expects (activation, mask) as input and returns
-  # the new activations only, we add the mask back to output next.
-  encoder_block = layers.Serial(
-      layers.Residual(  # Attention block here.
-          layers.Parallel(layers.LayerNorm(), layers.Identity()),
-          layers.MultiHeadedAttention(feature_depth, num_heads=num_heads,
-                                      dropout=dropout, mode=mode),
-          layers.Dropout(rate=dropout, mode=mode),
-          shortcut=layers.FirstBranch()
+  return tl.Serial(
+      tl.Residual(  # Attention block here.
+          tl.Parallel(tl.LayerNorm(), tl.Copy()),
+          tl.MultiHeadedAttention(feature_depth, num_heads=num_heads,
+                                  dropout=dropout, mode=mode),
+          tl.Parallel(tl.Dropout(rate=dropout, mode=mode), tl.Copy())
       ),
-      ResidualFeedForward(feature_depth, feedforward_depth, dropout, mode=mode)
-  )
-  # Now we add the mask back.
-  return layers.Serial(
-      layers.Reorder(output=((0, 1), 1)),  # (x, mask) --> ((x, mask), mask)
-      layers.Parallel(encoder_block, layers.Identity())
+      tl.Parallel(
+          ResidualFeedForward(
+              feature_depth, feedforward_depth, dropout, mode=mode),
+          tl.Div(divisor=2.0)  # Mask added to itself in the residual, divide.
+      )
   )
 
 
@@ -101,22 +97,21 @@ def TransformerEncoder(vocab_size,
   Returns:
     the Transformer encoder layer.
   """
-  input_embedding = layers.Serial(
-      layers.Embedding(feature_depth, vocab_size),
-      layers.Dropout(rate=dropout, mode=mode),
-      layers.PositionalEncoding(max_len=max_len)
+  input_embedding = tl.Serial(
+      tl.Embedding(feature_depth, vocab_size),
+      tl.Dropout(rate=dropout, mode=mode),
+      tl.PositionalEncoding(max_len=max_len)
   )
-  return layers.Serial(
-      layers.Branch(),  # Branch input to create embedding and mask.
-      layers.Parallel(input_embedding, layers.PaddingMask()),
-      layers.Serial(*[EncoderLayer(feature_depth, feedforward_depth, num_heads,
-                                   dropout, mode)
-                      for _ in range(num_layers)]),
-      layers.FirstBranch(),  # Drop the mask.
-      layers.LayerNorm(),
-      layers.Mean(axis=1),  # Average on length.
-      layers.Dense(num_classes),
-      layers.LogSoftmax()
+  return tl.Serial(
+      tl.Branch(input_embedding, tl.PaddingMask()),
+      tl.Serial(*[EncoderLayer(feature_depth, feedforward_depth, num_heads,
+                               dropout, mode)
+                  for _ in range(num_layers)]),
+      tl.Select(0),  # Drop the mask.
+      tl.LayerNorm(),
+      tl.Mean(axis=1),  # Average on length.
+      tl.Dense(num_classes),
+      tl.LogSoftmax()
   )
 
 
@@ -137,15 +132,14 @@ def DecoderLayer(feature_depth,
   Returns:
     the layer.
   """
-  return layers.Serial(
-      layers.Residual(  # Self-attention block.
-          layers.LayerNorm(),
-          layers.Branch(),
-          layers.Parallel(layers.Identity(),  # activation for (q, k, v)
-                          layers.CausalMask(axis=-2)),  # attention mask
-          layers.MultiHeadedAttention(feature_depth, num_heads=num_heads,
-                                      dropout=dropout, mode=mode),
-          layers.Dropout(rate=dropout, mode=mode)
+  return tl.Serial(
+      tl.Residual(  # Self-attention block.
+          tl.LayerNorm(),
+          tl.Branch(tl.Copy(), tl.CausalMask(axis=-2)),  # Create mask.
+          tl.MultiHeadedAttention(feature_depth, num_heads=num_heads,
+                                  dropout=dropout, mode=mode),
+          tl.Select(0),  # Drop the mask.
+          tl.Dropout(rate=dropout, mode=mode)
       ),
       ResidualFeedForward(feature_depth, feedforward_depth, dropout, mode=mode)
   )
@@ -174,17 +168,17 @@ def TransformerLM(vocab_size,
   Returns:
     the layer.
   """
-  return layers.Serial(
-      layers.ShiftRight(),
-      layers.Embedding(feature_depth, vocab_size),
-      layers.Dropout(rate=dropout, mode=mode),
-      layers.PositionalEncoding(max_len=max_len),
-      layers.Serial(*[DecoderLayer(feature_depth, feedforward_depth, num_heads,
-                                   dropout, mode)
-                      for _ in range(num_layers)]),
-      layers.LayerNorm(),
-      layers.Dense(vocab_size),
-      layers.LogSoftmax()
+  return tl.Serial(
+      tl.ShiftRight(),
+      tl.Embedding(feature_depth, vocab_size),
+      tl.Dropout(rate=dropout, mode=mode),
+      tl.PositionalEncoding(max_len=max_len),
+      tl.Serial(*[DecoderLayer(feature_depth, feedforward_depth, num_heads,
+                               dropout, mode)
+                  for _ in range(num_layers)]),
+      tl.LayerNorm(),
+      tl.Dense(vocab_size),
+      tl.LogSoftmax()
   )
 
 
@@ -207,15 +201,15 @@ def ChunkedDecoderLayer(feature_depth,
   Returns:
     the layer.
   """
-  return layers.Serial(
-      layers.Residual(  # Self-attention block.
-          layers.Map(layers.LayerNorm()),
-          layers.ChunkedCausalMultiHeadedAttention(
+  return tl.Serial(
+      tl.Residual(  # Self-attention block.
+          tl.Map(tl.LayerNorm()),
+          tl.ChunkedCausalMultiHeadedAttention(
               feature_depth, num_heads=num_heads, dropout=dropout,
               chunk_selector=chunk_selector, mode=mode),
-          layers.Map(layers.Dropout(rate=dropout, mode=mode)),
+          tl.Map(tl.Dropout(rate=dropout, mode=mode)),
       ),
-      layers.Map(ResidualFeedForward(
+      tl.Map(ResidualFeedForward(
           feature_depth, feedforward_depth, dropout, mode=mode))
   )
 
@@ -261,15 +255,15 @@ def ChunkedTransformerLM(vocab_size,
                                dropout, chunk_selector, mode)
            for _ in range(num_layers)]
   # Below each Map(L) applies the layer L to each chunk independently.
-  return layers.Serial(
-      layers.ShiftRight(),
-      layers.Map(layers.Embedding(feature_depth, vocab_size)),
-      layers.Map(layers.Dropout(rate=dropout, mode=mode)),
-      layers.PositionalEncoding(max_len=max_len),
-      layers.Serial(*stack),
-      layers.Map(layers.LayerNorm()),
-      layers.Map(layers.Dense(vocab_size)),
-      layers.Map(layers.LogSoftmax()),
+  return tl.Serial(
+      tl.ShiftRight(),
+      tl.Map(tl.Embedding(feature_depth, vocab_size)),
+      tl.Map(tl.Dropout(rate=dropout, mode=mode)),
+      tl.PositionalEncoding(max_len=max_len),
+      tl.Serial(*stack),
+      tl.Map(tl.LayerNorm()),
+      tl.Map(tl.Dense(vocab_size)),
+      tl.Map(tl.LogSoftmax()),
   )
 
 
@@ -295,32 +289,31 @@ def EncoderDecoderLayer(feature_depth,
     the layer, returning a triple (encoder, mask, decoder_activations).
   """
   # Decoder self-attending to decoder.
-  self_attention = layers.Residual(
-      layers.LayerNorm(),
-      layers.Branch(),
-      layers.Parallel(layers.Identity(),  # activation for (q, k, v)
-                      layers.CausalMask(axis=-2)),  # attention mask
-      layers.MultiHeadedAttention(feature_depth, num_heads=num_heads,
-                                  dropout=dropout, mode=mode),
-      layers.Dropout(rate=dropout, mode=mode)
+  self_attention = tl.Residual(
+      tl.LayerNorm(),
+      tl.Branch(tl.Copy(), tl.CausalMask(axis=-2)),  # create mask
+      tl.MultiHeadedAttention(feature_depth, num_heads=num_heads,
+                              dropout=dropout, mode=mode),
+      tl.Select(0),  # drop mask
+      tl.Dropout(rate=dropout, mode=mode)
   )
   # Decoder attending to encoder.
-  encoder_decoder_attention = layers.Serial(
-      layers.Reorder(output=((2, 0, 0), 1)),  # ((dec, enc, enc), mask)
-      layers.MultiHeadedAttentionQKV(  # ((q, k, v), mask) --> new v
+  encoder_decoder_attention = tl.Serial(
+      tl.Select(((2, 0, 0), 1)),  # ((dec, enc, enc), mask)
+      tl.MultiHeadedAttentionQKV(  # ((q, k, v), mask) --> new, mask
           feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),
-      layers.Dropout(rate=dropout, mode=mode),
+      tl.Select(0),  # drop the mask
+      tl.Dropout(rate=dropout, mode=mode),
   )
-  return layers.Serial(
-      layers.Parallel(layers.Identity(), layers.Identity(), self_attention),
-      layers.Branch(),
-      layers.Parallel(layers.Identity(), encoder_decoder_attention),
-      layers.UnnestBranches(),   # (encoder, mask, old_act, new_act)
-      layers.Reorder(output=(0, 1, (2, 3))),
-      layers.Parallel(  # Residual after encoder-decoder attention.
-          layers.Identity(), layers.Identity(), layers.SumBranches()),
-      layers.Parallel(  # Feed-forward on the third component (decoder).
-          layers.Identity(), layers.Identity(), ResidualFeedForward(
+  return tl.Serial(
+      tl.Parallel(tl.Copy(), tl.Copy(), self_attention),
+      tl.Branch(tl.Copy(), encoder_decoder_attention),
+      tl.UnnestBranches(),   # (encoder, mask, old_act, new_act)
+      tl.Select((0, 1, (2, 3))),
+      tl.Parallel(  # Residual after encoder-decoder attention.
+          tl.Copy(), tl.Copy(), tl.Add()),
+      tl.Parallel(  # Feed-forward on the third component (decoder).
+          tl.Copy(), tl.Copy(), ResidualFeedForward(
               feature_depth, feedforward_depth, dropout, mode=mode)
       )
   )
@@ -352,32 +345,31 @@ def Transformer(vocab_size,
   Returns:
     the Transformer model.
   """
-  embedding = layers.Serial(
-      layers.Embedding(feature_depth, vocab_size),
-      layers.Dropout(rate=dropout, mode=mode),
-      layers.PositionalEncoding(max_len=max_len)
+  embedding = tl.Serial(
+      tl.Embedding(feature_depth, vocab_size),
+      tl.Dropout(rate=dropout, mode=mode),
+      tl.PositionalEncoding(max_len=max_len)
   )
-  encoder = layers.Serial(
-      layers.Branch(),  # Branch input to create embedding and mask.
-      layers.Parallel(embedding, layers.PaddingMask()),
-      layers.Serial(*[EncoderLayer(feature_depth, feedforward_depth, num_heads,
-                                   dropout, mode)
-                      for _ in range(num_layers)]),
-      layers.Parallel(layers.LayerNorm(), layers.Identity())
+  encoder = tl.Serial(
+      tl.Branch(embedding, tl.PaddingMask()),
+      tl.Serial(*[EncoderLayer(feature_depth, feedforward_depth, num_heads,
+                               dropout, mode)
+                  for _ in range(num_layers)]),
+      tl.Parallel(tl.LayerNorm(), tl.Copy())
   )
   stack = [EncoderDecoderLayer(feature_depth, feedforward_depth, num_heads,
                                dropout, mode)
            for _ in range(num_layers)]
-  return layers.Serial(
-      layers.Parallel(layers.Identity(), layers.ShiftRight()),
-      layers.Parallel(encoder, embedding),
-      layers.UnnestBranches(),  # (encoder, encoder_mask, decoder_input)
-      layers.Reorder(output=(0, (1, 2), 2)),
-      layers.Parallel(  # (encoder_mask, decoder_input) -> encoder-decoder mask
-          layers.Identity(), layers.EncoderDecoderMask(), layers.Identity()),
-      layers.Serial(*stack),
-      layers.ThirdBranch(),
-      layers.LayerNorm(),
-      layers.Dense(vocab_size),
-      layers.LogSoftmax()
+  return tl.Serial(
+      tl.Parallel(tl.Copy(), tl.ShiftRight()),
+      tl.Parallel(encoder, embedding),
+      tl.UnnestBranches(),  # (encoder, encoder_mask, decoder_input)
+      tl.Select((0, (1, 2), 2)),
+      tl.Parallel(  # (encoder_mask, decoder_input) -> encoder-decoder mask
+          tl.Copy(), tl.EncoderDecoderMask(), tl.Copy()),
+      tl.Serial(*stack),
+      tl.Select(2),  # Drop encoder and mask.
+      tl.LayerNorm(),
+      tl.Dense(vocab_size),
+      tl.LogSoftmax()
   )
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 91294d300..df4f44795 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -127,8 +127,7 @@ def policy_and_value_net(rng_key,
   # other computes the value function.
   # NOTE: The LogSoftmax instead of the Softmax because of numerical stability.
   cur_layers.extend([
-      layers.Branch(),
-      layers.Parallel(
+      layers.Branch(
           layers.Serial(layers.Dense(num_actions), layers.LogSoftmax()),
           layers.Dense(1))
   ])
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 6577d245b..5834fe99c 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -376,7 +376,8 @@ def train(output_dir,
           num_devices=None,
           random_seed=None,
           run_debug_step=False,
-          save_forward_graph=False):
+          save_graphs=True,
+          save_backward_graph=False):
   """Train the model on the inputs.
 
   Args:
@@ -400,8 +401,8 @@ def train(output_dir,
     random_seed: the random seed to use; time/os dependent if None (default).
     run_debug_step: bool, if True, will run the model and loss without @jit for
       one step.
-    save_forward_graph: bool, if True, save forward computation graph to file.
-
+    save_graphs: bool, if True, save computation graph to file.
+    save_backward_graph: bool, if True, save backward graph to file too.
   Returns:
     trax.State
   """
@@ -510,13 +511,22 @@ def train(output_dir,
         eval_sw=eval_sw,
         history=history)
 
-    # Save computation graph
-    if save_forward_graph and step == 1:
-      # Dump forward computation graph to file.
-      computation = jax.xla_computation(model_predict_eval)(
+    # Save computation graph (single-device only for now).
+    if save_graphs and step == 1 and num_devices == 1:
+      # Dump computation graphs to files.
+      forward_computation = jax.xla_computation(model_predict_eval)(
           next_train_batch[0], params=params, rng=rng)
-      with gfile.GFile(os.path.join(output_dir, "forward_graph.dot"), "w") as f:
-        f.write(computation.GetHloDotGraph())
+      with gfile.GFile(os.path.join(output_dir, "forward.txt"), "w") as f:
+        f.write(forward_computation.GetHloText())
+      with gfile.GFile(os.path.join(output_dir, "forward.dot"), "w") as f:
+        f.write(forward_computation.GetHloDotGraph())
+      backward_computation = jax.xla_computation(jit_update_fun)(
+          step, opt_state, next_train_batch, rngs)
+      with gfile.GFile(os.path.join(output_dir, "backward.txt"), "w") as f:
+        f.write(backward_computation.GetHloText())
+      if save_backward_graph:  # Backward graphs can be large so we guard it.
+        with gfile.GFile(os.path.join(output_dir, "backward.dot"), "w") as f:
+          f.write(backward_computation.GetHloDotGraph())
 
     # Save state
     save_state(State(params=params, step=step, history=history), output_dir)

From ccbe1326885d2326a62b6397f3282a0164418098 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sun, 5 May 2019 00:21:47 -0700
Subject: [PATCH 1995/2720] [trax] update trax to use new optimizers api

PiperOrigin-RevId: 246698591
---
 tensor2tensor/trax/optimizers.py | 28 +++++++++++++++--
 tensor2tensor/trax/rlax/ppo.py   | 54 +++++++++++++++++++-------------
 tensor2tensor/trax/trax.py       | 30 +++++++++---------
 3 files changed, 73 insertions(+), 39 deletions(-)

diff --git a/tensor2tensor/trax/optimizers.py b/tensor2tensor/trax/optimizers.py
index e0c5c9bc9..5347de230 100644
--- a/tensor2tensor/trax/optimizers.py
+++ b/tensor2tensor/trax/optimizers.py
@@ -21,7 +21,9 @@
 
 import gin
 
+import jax
 from jax.experimental import optimizers as opt
+import numpy as onp
 
 
 def opt_configure(*args, **kwargs):
@@ -40,5 +42,27 @@ def opt_configure(*args, **kwargs):
 inverse_time_decay = opt_configure(opt.inverse_time_decay)
 piecewise_constant = opt_configure(opt.piecewise_constant)
 
-# Get params
-get_params = opt.get_params
+
+def parallelize(opt_maker):
+  """Transform an optimizer maker into a parallel one with replicated params."""
+
+  def parallel_opt_maker(*args, **kwargs):  # pylint:disable=missing-docstring
+    init_fun, update_fun, get_params = opt_maker(*args, **kwargs)
+
+    num_devices = jax.lib.xla_bridge.device_count()
+    replicate_array = lambda x: onp.broadcast_to(x, (num_devices,) + x.shape)
+    unreplicate_array = lambda x: x.mean(0)  # an alternative is x[0]
+
+    def init_replicated(params):
+      if num_devices > 1:
+        params = jax.tree_util.tree_map(replicate_array, params)
+      return init_fun(params)
+
+    def get_params_unreplicated(opt_state):
+      params = get_params(opt_state)
+      if num_devices > 1:
+        params = jax.tree_util.tree_map(unreplicate_array, params)
+      return params
+
+    return init_replicated, update_fun, get_params, get_params_unreplicated
+  return parallel_opt_maker
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index df4f44795..8b759b80d 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -136,10 +136,10 @@ def policy_and_value_net(rng_key,
 
 
 def optimizer_fun(net_params, step_size=1e-3):
-  opt_init, opt_update = trax_opt.adam(
+  opt_init, opt_update, get_params = trax_opt.adam(
       step_size=step_size, b1=0.9, b2=0.999, eps=1e-08)
   opt_state = opt_init(net_params)
-  return opt_state, opt_update
+  return opt_state, opt_update, get_params
 
 
 def log_params(params, name="params"):
@@ -782,10 +782,11 @@ def combined_loss(new_params,
       epsilon=epsilon)
 
 
-@functools.partial(jit, static_argnums=(2, 3, 5))
+@functools.partial(jit, static_argnums=(2, 3, 4, 6))
 def ppo_opt_step(i,
                  opt_state,
                  ppo_opt_update,
+                 ppo_get_params,
                  policy_net_apply,
                  old_policy_params,
                  value_net_apply,
@@ -798,7 +799,7 @@ def ppo_opt_step(i,
                  lambda_=0.95,
                  epsilon=0.1):
   """PPO optimizer step."""
-  new_policy_params = trax_opt.get_params(opt_state)
+  new_policy_params = ppo_get_params(opt_state)
   g = grad(
       ppo_loss, argnums=1)(
           policy_net_apply,
@@ -816,17 +817,18 @@ def ppo_opt_step(i,
   return ppo_opt_update(i, g, opt_state)
 
 
-@functools.partial(jit, static_argnums=(2, 3))
+@functools.partial(jit, static_argnums=(2, 3, 4))
 def value_opt_step(i,
                    opt_state,
                    opt_update,
+                   get_params,
                    value_net_apply,
                    padded_observations,
                    padded_rewards,
                    reward_mask,
                    gamma=0.99):
   """Value optimizer step."""
-  value_params = trax_opt.get_params(opt_state)
+  value_params = get_params(opt_state)
   # Note this partial application here and argnums above in ppo_opt_step.
   g = grad(functools.partial(value_loss, value_net_apply))(
       value_params,
@@ -837,10 +839,11 @@ def value_opt_step(i,
   return opt_update(i, g, opt_state)
 
 
-@functools.partial(jit, static_argnums=(2, 3))
+@functools.partial(jit, static_argnums=(2, 3, 4))
 def policy_and_value_opt_step(i,
                               opt_state,
                               opt_update,
+                              get_params,
                               policy_and_value_net_apply,
                               old_params,
                               padded_observations,
@@ -872,7 +875,7 @@ def policy_and_value_loss(params):
         epsilon=epsilon)
     return loss
 
-  new_params = trax_opt.get_params(opt_state)
+  new_params = get_params(opt_state)
   g = grad(policy_and_value_loss)(new_params)
   return opt_update(i, g, opt_state)
 
@@ -934,8 +937,10 @@ def training_loop(
         policy_and_value_net_fun(subkey, batch_observations_shape, num_actions))
 
     # Initialize the optimizers.
-    policy_and_value_opt_state, policy_and_value_opt_update = (
+    policy_and_value_optimizer = (
         policy_and_value_optimizer_fun(policy_and_value_net_params))
+    (policy_and_value_opt_state, policy_and_value_opt_update,
+     policy_and_value_get_params) = policy_and_value_optimizer
 
     policy_and_value_net_apply = jit(policy_and_value_net_apply)
   else:
@@ -953,19 +958,21 @@ def training_loop(
     value_net_apply = jit(value_net_apply)
 
     # Initialize the optimizers.
-    ppo_opt_state, ppo_opt_update = policy_optimizer_fun(policy_net_params)
-    value_opt_state, value_opt_update = value_optimizer_fun(value_net_params)
+    ppo_opt_state, ppo_opt_update, ppo_get_params = (
+        policy_optimizer_fun(policy_net_params))
+    value_opt_state, value_opt_update, value_get_params = (
+        value_optimizer_fun(value_net_params))
 
   # A function that will call the appropriate policy function with parameters.
   def get_policy_output(observations):
     # Get the fresh params for collecting the policy.
     if policy_net_apply is not None:
-      return policy_net_apply(observations, trax_opt.get_params(ppo_opt_state))
+      return policy_net_apply(observations, ppo_get_params(ppo_opt_state))
 
     assert policy_and_value_net_apply
 
     policy_predictions, unused_value_predictions = policy_and_value_net_apply(
-        observations, trax_opt.get_params(policy_and_value_opt_state))
+        observations, policy_and_value_get_params(policy_and_value_opt_state))
     return policy_predictions
 
   for i in range(epochs):
@@ -985,11 +992,11 @@ def get_policy_output(observations):
 
     # These were the params that were used to collect the trajectory.
     if policy_and_value_net_apply:
-      policy_and_value_net_params = trax_opt.get_params(
+      policy_and_value_net_params = policy_and_value_get_params(
           policy_and_value_opt_state)
     else:
-      policy_net_params = trax_opt.get_params(ppo_opt_state)
-      value_net_params = trax_opt.get_params(value_opt_state)
+      policy_net_params = ppo_get_params(ppo_opt_state)
+      value_net_params = value_get_params(value_opt_state)
 
     avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
     max_reward = max(np.sum(traj[2]) for traj in trajs)
@@ -1099,6 +1106,7 @@ def get_policy_output(observations):
             j,
             policy_and_value_opt_state,
             policy_and_value_opt_update,
+            policy_and_value_get_params,
             policy_and_value_net_apply,
             # for the entirety of this loop, this should refer to params that
             # were used to collect the trajectory.
@@ -1117,7 +1125,7 @@ def get_policy_output(observations):
             print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
           # Compute and log the loss.
           # Get the new params.
-          new_policy_and_value_net_params = trax_opt.get_params(
+          new_policy_and_value_net_params = policy_and_value_get_params(
               policy_and_value_opt_state)
           (loss_combined, loss_ppo, loss_value, unused_entropy_bonus) = (
               combined_loss(
@@ -1166,6 +1174,7 @@ def get_policy_output(observations):
             j,
             ppo_opt_state,
             ppo_opt_update,
+            ppo_get_params,
             policy_net_apply,
             policy_net_params,
             value_net_apply,
@@ -1180,7 +1189,7 @@ def get_policy_output(observations):
         )
         t2 = time.time()
         # Get the new params.
-        new_policy_net_params = trax_opt.get_params(ppo_opt_state)
+        new_policy_net_params = ppo_get_params(ppo_opt_state)
 
         # These are the "old" params - policy_net_params
 
@@ -1226,7 +1235,7 @@ def get_policy_output(observations):
       # Update the params ONLY AND ONLY AFTER we complete all the optimization
       # iterations, till then `policy_net_params` should refer to the params
       # that were used in collecting the policy.
-      # policy_net_params = trax_opt.get_params(ppo_opt_state)
+      # policy_net_params = ppo_get_params(ppo_opt_state)
 
       logging.vlog(1, "Total PPO loss reduction [%0.2f]%%",
                    (100 * (cur_ppo_loss - new_ppo_loss) / np.abs(cur_ppo_loss)))
@@ -1239,13 +1248,14 @@ def get_policy_output(observations):
             j,
             value_opt_state,
             value_opt_update,
+            value_get_params,
             value_net_apply,
             padded_observations,
             padded_rewards,
             reward_mask,
             gamma=gamma)
         t2 = time.time()
-        value_net_params = trax_opt.get_params(value_opt_state)
+        value_net_params = value_get_params(value_opt_state)
         if ((j + 1) %
             print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
           new_value_loss = value_loss(
@@ -1266,8 +1276,8 @@ def get_policy_output(observations):
       logging.vlog(1, "Grad desc took %0.2f msec", get_time(t1))
 
       # Set the optimized params to new params.
-      policy_net_params = trax_opt.get_params(ppo_opt_state)
-      value_net_params = trax_opt.get_params(value_opt_state)
+      policy_net_params = ppo_get_params(ppo_opt_state)
+      value_net_params = value_get_params(value_opt_state)
 
       logging.info(
           "Epoch [% 6d], Reward[min, max, avg] [%10.2f,%10.2f,%10.2f], "
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 5834fe99c..9810284ff 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -159,13 +159,12 @@ def save_state(state, output_dir, keep=False):
   else:
     pkl_module = pickle
   params_file = os.path.join(output_dir, "model.pkl")
-  params = jax.unreplicate(state.params)
   with gfile.GFile(params_file, "wb") as f:
-    pkl_module.dump((params, state.step, state.history), f)
+    pkl_module.dump((state.params, state.step, state.history), f)
   if keep:
     params_file = os.path.join(output_dir, "model_{}.pkl".format(state.step))
     with gfile.GFile(params_file, "wb") as f:
-      pkl_module.dump((params, state.step, state.history), f)
+      pkl_module.dump((state.params, state.step, state.history), f)
   log("Model saved to %s" % params_file, stdout=False)
 
 
@@ -317,8 +316,8 @@ def _jit_update_fun(predict_fun, loss_fun, optimizer, lr_fun, num_devices):
   if num_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
     def single_update(i, opt_state, batch, rng):
       rng, subrng = jax_random.split(rng[0])
-      _, opt_update = optimizer(lr_fun)
-      params = trax_opt.get_params(opt_state)
+      _, opt_update, get_params = optimizer(lr_fun)
+      params = get_params(opt_state)
       return opt_update(i, backend.grad(loss_fun)(
           params, batch, predict_fun, rng), opt_state), [subrng]
     return backend.jit(single_update)
@@ -328,15 +327,15 @@ def mapped_update(i, opt_state, batch, rng):
     """This is a multi-device version of the update function above."""
     # We assume all tensors have the first dimension = num_devices.
     rng, subrng = jax_random.split(rng)
-    _, opt_update = optimizer(lr_fun)
-    params = trax_opt.get_params(opt_state)
+    _, opt_update, get_params = optimizer(lr_fun)
+    params = get_params(opt_state)
     grads = backend.grad(loss_fun)(params, batch, predict_fun, rng)
     grads = jax.tree_util.tree_map(
         lambda g: lax.psum(g, "batch"), grads)
     return opt_update(i, grads, opt_state), subrng
 
   def update(i, opt_state, batch, rng):
-    return mapped_update(jax.replicate(i), opt_state, batch, rng)
+    return mapped_update(numpy.repeat(i, num_devices), opt_state, batch, rng)
 
   return update
 
@@ -421,7 +420,8 @@ def train(output_dir,
   state = restore_state(output_dir)
   history = state.history
   lr_fun = lr_schedule(history)
-  opt_init, _ = optimizer(lr_fun)
+  opt_init, _, get_rep_params, get_params = (
+      trax_opt.parallelize(optimizer)(lr_fun))
   model_train = model(mode="train")
   model_predict_eval = model(mode="eval")
 
@@ -438,8 +438,6 @@ def train(output_dir,
     model_input_shape = tuple([-1] + list(inputs.input_shape))
   params = state.params or model_train.initialize(model_input_shape, init_rng)
   opt_state = opt_init(params)
-  if num_devices > 1:  # TODO(lukaszkaiser): use everywhere when pmap is stable.
-    opt_state = jax.replicate(opt_state)
 
   # jit model_predict and update so they're fast
   jit_model_predict_eval = _jit_predict_fun(model_predict_eval, num_devices)
@@ -475,7 +473,7 @@ def train(output_dir,
       step += 1
 
       if step in save_steps:
-        params = trax_opt.get_params(opt_state)
+        params = get_params(opt_state)
         save_state(State(params=params, step=step, history=history),
                    output_dir,
                    keep=True)
@@ -494,17 +492,19 @@ def train(output_dir,
                       epoch_steps / epoch_time, step=step)
 
     # Print number of parameters
-    params = trax_opt.get_params(opt_state)
+    params = get_params(opt_state)
     if step == 1:
       sizes = layers.sizes(params)
       total_size = layers.nested_reduce(sizes, sum)
       step_log(step, "Total trainable parameters size: %d" % total_size)
 
-    # Evaluate
+    # Evaluate in parallel
+    replicated_params = get_rep_params(opt_state)
     evaluate_train_and_eval(
         step=step,
         inputs=inputs,
-        predict_fun=functools.partial(jit_model_predict_eval, params=params),
+        predict_fun=functools.partial(jit_model_predict_eval,
+                                      params=replicated_params),
         eval_steps=eval_steps,
         rng=rng,
         train_sw=train_sw,

From 67c26e4a8c6e114543228c1906d639360d9b4fe2 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 6 May 2019 17:56:13 -0700
Subject: [PATCH 1996/2720] Probabs while sampling actions from policy still
 aren't normalized sometimes, very close, but the numpy check fails, so
 normalize them. (Maybe do all that in JAX?)

PiperOrigin-RevId: 246932988
---
 tensor2tensor/envs/env_problem_utils.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index cb76f7ea5..f5b05ff17 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -106,7 +106,8 @@ def multinomial_sample(probs):
     log_probs = log_prob_actions[np.arange(B)[:, None],
                                  index[:, None],
                                  np.arange(A)]
-    assert (B, A) == log_probs.shape
+    assert (B, A) == log_probs.shape, \
+        "B=%d, A=%d, log_probs.shape=%s" % (B, A, log_probs.shape)
 
     # Convert to probs, since we need to do categorical sampling.
     probs = np.exp(log_probs)
@@ -114,12 +115,17 @@ def multinomial_sample(probs):
     # Sometimes log_probs contains a 0, it shouldn't. This makes the
     # probabilities sum up to more than 1, since the addition happens
     # in float64, so just add and subtract 1.0 to zero those probabilites
-    # out. Real example encountered probs = [1e-8, 1.0, 1e-22]
+    # out.
     #
     # Also testing for this is brittle.
     probs += 1
     probs -= 1
 
+    # For some reason, sometimes, this isn't the case.
+    probs_sum = np.sum(probs, axis=1, keepdims=True)
+    if not all(probs_sum == 1.0):
+      probs = probs / probs_sum
+
     # Now pick actions from this probs array.
     actions = np.apply_along_axis(multinomial_sample, 1, probs)
 

From 18b3f88acded6a1c98c3b8543dfa8575ea5c6090 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 6 May 2019 17:57:54 -0700
Subject: [PATCH 1997/2720] - Fix KL and Entropy computation in PPO, this gives
 much smoother results. - Reuse log-probabs and value-predictions from the
 collect-policy.     - In functions that need them (ppo_loss, combined_loss,
 etc).     - Rather than recompute them with params. - Stop passing in
 value_params and value_apply to ppo_loss/ppo_opt step.     - As a result of
 the above.

PiperOrigin-RevId: 246933187
---
 tensor2tensor/trax/rlax/ppo.py      | 255 +++++++++++++++++-----------
 tensor2tensor/trax/rlax/ppo_test.py |  93 +++++++++-
 2 files changed, 240 insertions(+), 108 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 8b759b80d..a3be4d62d 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -615,12 +615,11 @@ def clipped_objective(probab_ratios, advantages, reward_mask, epsilon=0.2):
       advantages) * reward_mask
 
 
-@functools.partial(jit, static_argnums=(0, 3))
+@functools.partial(jit, static_argnums=(0,))
 def ppo_loss(policy_net_apply,
              new_policy_params,
-             old_policy_params,
-             value_net_apply,
-             value_net_params,
+             log_probab_actions_old,
+             value_predictions_old,
              padded_observations,
              padded_actions,
              padded_rewards,
@@ -631,42 +630,35 @@ def ppo_loss(policy_net_apply,
   """PPO objective, with an eventual minus sign, given observations."""
   B, T = padded_rewards.shape  # pylint: disable=invalid-name
   assert (B, T + 1) == padded_observations.shape[:2]
+  assert (B, T + 1) == log_probab_actions_old.shape[:2]
+  assert (B, T + 1, 1) == value_predictions_old.shape
   assert (B, T) == padded_actions.shape
   assert (B, T) == padded_rewards.shape
   assert (B, T) == reward_mask.shape
 
-  # Compute predicted values and predicted log-probs and hand it over to
-  # `ppo_loss_given_predictions`.
-
-  # (B, T+1, 1)
-  predicted_values = value_net_apply(padded_observations, value_net_params)
-  assert (B, T + 1, 1) == predicted_values.shape
+  # Compute predicted log-probs and hand over to `ppo_loss_given_predictions`.
 
   # log_probab_actions_{old,new} are both (B, T+1, A)
-  log_probab_actions_old = policy_net_apply(padded_observations,
-                                            old_policy_params)
   log_probab_actions_new = policy_net_apply(padded_observations,
                                             new_policy_params)
-  assert (B, T + 1) == log_probab_actions_old.shape[:2]
   assert (B, T + 1) == log_probab_actions_new.shape[:2]
   assert log_probab_actions_old.shape[-1] == log_probab_actions_new.shape[-1]
 
-  return ppo_loss_given_predictions(
-      log_probab_actions_new,
-      log_probab_actions_old,
-      predicted_values,
-      padded_actions,
-      padded_rewards,
-      reward_mask,
-      gamma=gamma,
-      lambda_=lambda_,
-      epsilon=epsilon)
+  return ppo_loss_given_predictions(log_probab_actions_new,
+                                    log_probab_actions_old,
+                                    value_predictions_old,
+                                    padded_actions,
+                                    padded_rewards,
+                                    reward_mask,
+                                    gamma=gamma,
+                                    lambda_=lambda_,
+                                    epsilon=epsilon)
 
 
 @jit
 def ppo_loss_given_predictions(log_probab_actions_new,
                                log_probab_actions_old,
-                               predicted_values,
+                               value_predictions_old,
                                padded_actions,
                                padded_rewards,
                                reward_mask,
@@ -679,13 +671,13 @@ def ppo_loss_given_predictions(log_probab_actions_new,
   assert (B, T) == reward_mask.shape
 
   _, _, A = log_probab_actions_old.shape  # pylint: disable=invalid-name
-  assert (B, T + 1, 1) == predicted_values.shape
+  assert (B, T + 1, 1) == value_predictions_old.shape
   assert (B, T + 1, A) == log_probab_actions_old.shape
   assert (B, T + 1, A) == log_probab_actions_new.shape
 
   # (B, T)
   td_deltas = deltas(
-      np.squeeze(predicted_values, axis=2),  # (B, T+1)
+      np.squeeze(value_predictions_old, axis=2),  # (B, T+1)
       padded_rewards,
       reward_mask,
       gamma=gamma)
@@ -739,17 +731,15 @@ def combined_loss_given_predictions(log_probab_actions_new,
       gamma=gamma,
       lambda_=lambda_,
       epsilon=epsilon)
-  # TODO(afrozm): Add the entropy bonus, but since we don't do that in T2T
-  # we'll skip if for now.
-  entropy_bonus = 0.0
+  entropy_bonus = approximate_entropy(log_probab_actions_new, reward_mask)
   return (loss_ppo + (c1 * loss_value) - (c2 * entropy_bonus), loss_ppo,
           loss_value, entropy_bonus)
 
 
-# TODO(afrozm): Pass in `log_probab_actions_old` instead of re=computing it.
-@functools.partial(jit, static_argnums=(2,))
+@functools.partial(jit, static_argnums=(3,))
 def combined_loss(new_params,
-                  old_params,
+                  log_probab_actions_old,
+                  value_predictions_old,
                   policy_and_value_net_apply,
                   padded_observations,
                   padded_actions,
@@ -764,33 +754,28 @@ def combined_loss(new_params,
   log_probab_actions_new, _ = policy_and_value_net_apply(
       padded_observations, new_params)
 
-  log_probab_actions_old, value_predictions = policy_and_value_net_apply(
-      padded_observations, old_params)
-
   # (combined_loss, ppo_loss, value_loss, entropy_bonus)
-  return combined_loss_given_predictions(
-      log_probab_actions_new,
-      log_probab_actions_old,
-      value_predictions,
-      padded_actions,
-      padded_rewards,
-      reward_mask,
-      c1=c1,
-      c2=c2,
-      gamma=gamma,
-      lambda_=lambda_,
-      epsilon=epsilon)
+  return combined_loss_given_predictions(log_probab_actions_new,
+                                         log_probab_actions_old,
+                                         value_predictions_old,
+                                         padded_actions,
+                                         padded_rewards,
+                                         reward_mask,
+                                         c1=c1,
+                                         c2=c2,
+                                         gamma=gamma,
+                                         lambda_=lambda_,
+                                         epsilon=epsilon)
 
 
-@functools.partial(jit, static_argnums=(2, 3, 4, 6))
+@functools.partial(jit, static_argnums=(2, 3, 4))
 def ppo_opt_step(i,
                  opt_state,
                  ppo_opt_update,
                  ppo_get_params,
                  policy_net_apply,
-                 old_policy_params,
-                 value_net_apply,
-                 value_net_params,
+                 log_probab_actions_old,
+                 value_predictions_old,
                  padded_observations,
                  padded_actions,
                  padded_rewards,
@@ -804,9 +789,8 @@ def ppo_opt_step(i,
       ppo_loss, argnums=1)(
           policy_net_apply,
           new_policy_params,
-          old_policy_params,
-          value_net_apply,
-          value_net_params,
+          log_probab_actions_old,
+          value_predictions_old,
           padded_observations,
           padded_actions,
           padded_rewards,
@@ -845,7 +829,8 @@ def policy_and_value_opt_step(i,
                               opt_update,
                               get_params,
                               policy_and_value_net_apply,
-                              old_params,
+                              log_probab_actions_old,
+                              value_predictions_old,
                               padded_observations,
                               padded_actions,
                               padded_rewards,
@@ -862,7 +847,8 @@ def policy_and_value_loss(params):
     """Returns the combined loss given just parameters."""
     (loss, _, _, _) = combined_loss(
         params,
-        old_params,
+        log_probab_actions_old,
+        value_predictions_old,
         policy_and_value_net_apply,
         padded_observations,
         padded_actions,
@@ -886,6 +872,44 @@ def get_time(t1, t2=None):
   return round((t2 - t1) * 1000, 2)
 
 
+def approximate_kl(log_prob_new, log_prob_old, mask):
+  """Computes the approximate KL divergence between the old and new log-probs.
+
+  Args:
+    log_prob_new: (B, T+1, A) log probs new
+    log_prob_old: (B, T+1, A) log probs old
+    mask: (B, T)
+
+  Returns:
+    Approximate KL.
+  """
+  diff = log_prob_old - log_prob_new
+  # Cut the last time-step out.
+  diff = diff[:, :-1]
+  # Mask out the irrelevant part.
+  diff *= mask[:, :, np.newaxis]  # make mask (B, T, 1)
+  # Average on non-masked part.
+  return np.sum(diff) / np.sum(mask)
+
+
+def approximate_entropy(log_probs, mask):
+  """Computes the approximate entropy for the given log-probs.
+
+  Args:
+    log_probs: (B, T+1, A) log probs
+    mask: (B, T) mask.
+
+  Returns:
+    Approximate entropy.
+  """
+  # Cut the last time-step out.
+  lp = log_probs[:, :-1]
+  # Mask out the irrelevant part.
+  lp *= mask[:, :, np.newaxis]  # make mask (B, T, 1)
+  # Average on non-masked part and take negative.
+  return - (np.sum(lp) / np.sum(mask))
+
+
 def training_loop(
     env=None,
     epochs=EPOCHS,
@@ -963,25 +987,35 @@ def training_loop(
     value_opt_state, value_opt_update, value_get_params = (
         value_optimizer_fun(value_net_params))
 
-  # A function that will call the appropriate policy function with parameters.
-  def get_policy_output(observations):
-    # Get the fresh params for collecting the policy.
-    if policy_net_apply is not None:
-      return policy_net_apply(observations, ppo_get_params(ppo_opt_state))
+  for i in range(epochs):
 
-    assert policy_and_value_net_apply
+    # Params we'll use to collect the trajectories.
+    if policy_and_value_net_apply:
+      policy_and_value_net_params = policy_and_value_get_params(
+          policy_and_value_opt_state)
+    else:
+      policy_net_params = ppo_get_params(ppo_opt_state)
+      value_net_params = value_get_params(value_opt_state)
 
-    policy_predictions, unused_value_predictions = policy_and_value_net_apply(
-        observations, policy_and_value_get_params(policy_and_value_opt_state))
-    return policy_predictions
+    # A function to get the policy and value predictions.
+    def get_predictions(observations):
+      if policy_net_apply is not None:
+        # Get the fresh params for collecting the policy.
+        return (policy_net_apply(observations, policy_net_params),
+                value_net_apply(observations, value_net_params))
+
+      assert policy_and_value_net_apply
+
+      # Get the fresh params for collecting the policy.
+      return policy_and_value_net_apply(observations,
+                                        policy_and_value_net_params)
 
-  for i in range(epochs):
     t = time.time()
     t0 = t
     logging.vlog(1, "Epoch [% 6d] collecting trajectories.", i)
     trajs = collect_trajectories(
         env,
-        policy_fun=get_policy_output,
+        policy_fun=lambda observations: get_predictions(observations)[0],
         num_trajectories=batch_size,
         policy=POLICY,
         max_timestep=max_timestep,
@@ -1003,10 +1037,10 @@ def get_policy_output(observations):
     min_reward = min(np.sum(traj[2]) for traj in trajs)
     average_rewards.append(avg_reward)
 
-    logging.vlog(1, "Rewards average=[%0.2f], max=[%0.2f], min=[%0.2f]",
-                 avg_reward, max_reward, min_reward)
-    logging.vlog(2, "Rewards: %s", [float(np.sum(traj[2])) for traj in trajs])
-    logging.vlog(1, "Average Rewards: %s", average_rewards)
+    logging.vlog(1, "Rewards avg=[%0.2f], max=[%0.2f], min=[%0.2f], all=%s",
+                 avg_reward, max_reward, min_reward,
+                 [float(np.sum(traj[2])) for traj in trajs])
+    logging.vlog(1, "Average Rewards:\n%s", average_rewards)
 
     logging.vlog(1,
                  "Trajectory Length average=[%0.2f], max=[%0.2f], min=[%0.2f]",
@@ -1026,6 +1060,11 @@ def get_policy_output(observations):
     logging.vlog(1, "Padded Actions' shape [%s]", str(padded_actions.shape))
     logging.vlog(1, "Padded Rewards' shape [%s]", str(padded_rewards.shape))
 
+    # Calculate log-probabilities and value predictions of the trajectories.
+    # We'll pass these to the loss functions so as to not get recomputed.
+    log_probabs_traj, value_predictions_traj = get_predictions(
+        padded_observations)
+
     # Some assertions.
     B, T = padded_actions.shape  # pylint: disable=invalid-name
     assert (B, T) == padded_rewards.shape
@@ -1043,13 +1082,14 @@ def get_policy_output(observations):
 
     # Compute value and ppo losses.
     cur_value_loss, cur_ppo_loss, cur_combined_loss = None, None, None
-    if policy_and_value_net_apply is not None:
+    if policy_and_value_net_apply:
       logging.vlog(2, "Starting to compute P&V loss.")
       t = time.time()
-      cur_combined_loss, cur_ppo_loss, cur_value_loss, _ = (
+      cur_combined_loss, cur_ppo_loss, cur_value_loss, entropy_bonus = (
           combined_loss(
               policy_and_value_net_params,
-              policy_and_value_net_params,
+              log_probabs_traj,
+              value_predictions_traj,
               policy_and_value_net_apply,
               padded_observations,
               padded_actions,
@@ -1079,9 +1119,8 @@ def get_policy_output(observations):
       cur_ppo_loss = ppo_loss(
           policy_net_apply,
           policy_net_params,
-          policy_net_params,
-          value_net_apply,
-          value_net_params,
+          log_probabs_traj,
+          value_predictions_traj,
           padded_observations,
           padded_actions,
           padded_rewards,
@@ -1108,9 +1147,8 @@ def get_policy_output(observations):
             policy_and_value_opt_update,
             policy_and_value_get_params,
             policy_and_value_net_apply,
-            # for the entirety of this loop, this should refer to params that
-            # were used to collect the trajectory.
-            policy_and_value_net_params,
+            log_probabs_traj,
+            value_predictions_traj,
             padded_observations,
             padded_actions,
             padded_rewards,
@@ -1120,18 +1158,33 @@ def get_policy_output(observations):
             gamma=gamma,
             lambda_=lambda_,
             epsilon=epsilon_schedule)
+
+        # Compute the approx KL for early stopping.
+        new_policy_and_value_net_params = policy_and_value_get_params(
+            policy_and_value_opt_state)
+
+        log_probab_actions_new, _ = policy_and_value_net_apply(
+            padded_observations, new_policy_and_value_net_params)
+
+        approx_kl = approximate_kl(log_probab_actions_new,
+                                   log_probabs_traj,
+                                   reward_mask)
+
+        early_stopping = approx_kl > 1.5 * target_kl
+
         t2 = time.time()
-        if ((j + 1) %
-            print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
+        if (((j + 1) %
+             print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1)
+            or early_stopping):
           # Compute and log the loss.
           # Get the new params.
           new_policy_and_value_net_params = policy_and_value_get_params(
               policy_and_value_opt_state)
-          (loss_combined, loss_ppo, loss_value, unused_entropy_bonus) = (
+          (loss_combined, loss_ppo, loss_value, entropy_bonus) = (
               combined_loss(
                   new_policy_and_value_net_params,
-                  # old params, that were used to collect the trajectory
-                  policy_and_value_net_params,
+                  log_probabs_traj,
+                  value_predictions_traj,
                   policy_and_value_net_apply,
                   padded_observations,
                   padded_actions,
@@ -1144,10 +1197,12 @@ def get_policy_output(observations):
                   c2=c2))
           logging.vlog(1, "One Policy and Value grad desc took: %0.2f msec",
                        get_time(t, t2))
-          logging.vlog(
-              1,
-              "Combined Loss(value, ppo) [%10.2f] -> [%10.2f(%10.2f,%10.2f)]",
-              cur_combined_loss, loss_combined, loss_value, loss_ppo)
+          logging.vlog(1, "Combined Loss(value, ppo, entropy_bonus) [%10.2f] ->"
+                          " [%10.2f(%10.2f,%10.2f,%10.2f)]", cur_combined_loss,
+                       loss_combined, loss_value, loss_ppo, entropy_bonus)
+
+        if early_stopping:
+          break
 
       # Update the params.
       policy_and_value_net_params = new_policy_and_value_net_params
@@ -1176,9 +1231,8 @@ def get_policy_output(observations):
             ppo_opt_update,
             ppo_get_params,
             policy_net_apply,
-            policy_net_params,
-            value_net_apply,
-            value_net_params,
+            log_probabs_traj,
+            value_predictions_traj,
             padded_observations,
             padded_actions,
             padded_rewards,
@@ -1188,18 +1242,14 @@ def get_policy_output(observations):
             epsilon=epsilon_schedule,
         )
         t2 = time.time()
+        # Compute the approx KL for early stopping.
         # Get the new params.
         new_policy_net_params = ppo_get_params(ppo_opt_state)
-
-        # These are the "old" params - policy_net_params
-
-        # Compute the approx KL for early stopping.
-        log_probab_actions_old = policy_net_apply(padded_observations,
-                                                  policy_net_params)
         log_probab_actions_new = policy_net_apply(padded_observations,
                                                   new_policy_net_params)
-
-        approx_kl = np.mean(log_probab_actions_old - log_probab_actions_new)
+        approx_kl = approximate_kl(log_probab_actions_new,
+                                   log_probabs_traj,
+                                   reward_mask)
 
         early_stopping = approx_kl > 1.5 * target_kl
         if early_stopping:
@@ -1214,9 +1264,8 @@ def get_policy_output(observations):
           new_ppo_loss = ppo_loss(
               policy_net_apply,
               new_policy_net_params,
-              policy_net_params,
-              value_net_apply,
-              value_net_params,
+              log_probabs_traj,
+              value_predictions_traj,
               padded_observations,
               padded_actions,
               padded_rewards,
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 65b40de74..202766a88 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -537,9 +537,12 @@ def test_ppo_loss(self):
     rewards = np.random.uniform(0, 1, size=(B, T))
     mask = np.ones_like(rewards)
 
+    log_probs_old = policy_apply(observations, old_policy_params)
+    value_predictions_old = value_apply(observations, value_params)
+
     # Just test that this computes at all.
-    _ = ppo.ppo_loss(policy_apply, new_policy_params, old_policy_params,
-                     value_apply, value_params, observations, actions, rewards,
+    _ = ppo.ppo_loss(policy_apply, new_policy_params, log_probs_old,
+                     value_predictions_old, observations, actions, rewards,
                      mask)
 
   def test_combined_loss(self):
@@ -586,7 +589,8 @@ def test_combined_loss(self):
 
     (combined_loss, ppo_loss_2, value_loss_2, entropy_bonus) = (
         ppo.combined_loss(new_params,
-                          old_params,
+                          old_log_probabs,
+                          value_predictions,
                           net_apply,
                           observations,
                           actions,
@@ -600,11 +604,90 @@ def test_combined_loss(self):
     )
 
     # Test that these compute at all and are self consistent.
-    self.assertEqual(0.0, entropy_bonus)
+    self.assertGreater(entropy_bonus, 0.0)
     self.assertNear(value_loss_1, value_loss_2, 1e-6)
     self.assertNear(ppo_loss_1, ppo_loss_2, 1e-6)
-    self.assertNear(combined_loss, ppo_loss_2 + (c1 * value_loss_2), 1e-6)
+    self.assertNear(combined_loss,
+                    ppo_loss_2 + (c1 * value_loss_2) - (c2 * entropy_bonus),
+                    1e-6)
+
+  def test_approximate_kl(self):
+    # (2, 4+1, 4)
+    p_old = np.array([[
+        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
+        [np.log(0.4), np.log(0.1), np.log(0.4), np.log(0.1)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+    ], [
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [np.log(0.1), np.log(0.1), np.log(0.4), np.log(0.4)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+    ]])
+
+    # (2, 4+1, 4)
+    p_new = np.array([[
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [np.log(0.4), np.log(0.1), np.log(0.1), np.log(0.3)],
+        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
+    ], [
+        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
+        [np.log(0.1), np.log(0.1), np.log(0.2), np.log(0.6)],
+        [np.log(0.3), np.log(0.1), np.log(0.3), np.log(0.3)],
+        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
+        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
+    ]])
+
+    # (2, 4)
+    mask = np.array([
+        [1, 1, 0, 0],
+        [1, 1, 1, 0]
+    ])
+
+    self.assertNear(
+        ppo.approximate_kl(p_new, p_old, mask),
+        -ppo.approximate_entropy(p_old, mask) +
+        ppo.approximate_entropy(p_new, mask),
+        1e-6)
+
+  def test_get_approximate_entropy(self):
+    # (2, 4+1, 4)
+    log_probs = np.array([[
+        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
+        [np.log(0.4), np.log(0.1), np.log(0.4), np.log(0.1)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+    ], [
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [np.log(0.1), np.log(0.1), np.log(0.4), np.log(0.4)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+    ]])
+
+    # (2, 4)
+    mask = np.array([
+        [1, 1, 0, 0],
+        [1, 1, 1, 0]
+    ])
+
+    # Removing the last time-step and the masked stuff, gets us this.
+    filtered_log_probs = np.array([[
+        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
+        [np.log(0.4), np.log(0.1), np.log(0.4), np.log(0.1)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [np.log(0.1), np.log(0.1), np.log(0.4), np.log(0.4)],
+        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+    ]])
 
+    self.assertNear(ppo.approximate_entropy(log_probs, mask),
+                    -np.sum(filtered_log_probs) / 5.0,
+                    1e-6)
 
 if __name__ == "__main__":
   test.main()

From d765fa34a1843abea1b22ec1c9297cb44b09233d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 6 May 2019 20:39:31 -0700
Subject: [PATCH 1998/2720] LR defaults to the same one used by Spinup.

PiperOrigin-RevId: 246949173
---
 tensor2tensor/trax/rlax/ppo_main.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index ce5afc042..a0563b129 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -98,8 +98,8 @@
                      "Number of optimizer steps value only.")
 
 # Learning rate of the combined net, policy net and value net.
-flags.DEFINE_float("learning_rate", 5e-4, "Learning rate.")
-flags.DEFINE_float("policy_only_learning_rate", 1e-3,
+flags.DEFINE_float("learning_rate", 1e-3, "Learning rate.")
+flags.DEFINE_float("policy_only_learning_rate", 3e-4,
                    "Learning rate for policy network only.")
 flags.DEFINE_float("value_only_learning_rate", 1e-3,
                    "Learning rate for value network only.")

From 82b67b202181e59c9da83b75c335d5083273aad4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 6 May 2019 21:44:59 -0700
Subject: [PATCH 1999/2720] Internal

PiperOrigin-RevId: 246954330
---
 .../data_generators/text_problems.py          | 27 +++++++++++++++++++
 tensor2tensor/data_generators/translate.py    | 11 +++++---
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 0a5fd66ac..4c2d2cbf4 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -29,6 +29,7 @@
 from __future__ import print_function
 
 import os
+import re
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -611,6 +612,24 @@ def txt_line_iterator(txt_path):
       yield line.strip()
 
 
+def txt_and_label_iterator(txt_path):
+  """Iterate through lines of file."""
+  problem_pattern_without_vocab_size = re.compile("(.*)\tExtra_Label: (.*)")
+  with tf.gfile.Open(txt_path) as f:
+    for line in f:
+      results = problem_pattern_without_vocab_size.search(line.strip())
+      try:
+        line = results.group(1)
+        extra_label = int(results.group(2))
+      except AttributeError:
+        raise ValueError(
+            "Please provide the file in the right format, with each line having"
+            " the following format:\n<word_1 word_2 ... word_n>\\t"
+            "Extra_Label:\\s<int_label>"
+        )
+      yield [line, extra_label]
+
+
 def text2text_txt_iterator(source_txt_path, target_txt_path):
   """Yield dicts for Text2TextProblem.generate_samples from lines of files."""
   for inputs, targets in zip(
@@ -618,6 +637,14 @@ def text2text_txt_iterator(source_txt_path, target_txt_path):
     yield {"inputs": inputs, "targets": targets}
 
 
+def text2text_txt_iterator_with_label(source_txt_path, target_txt_path):
+  """Yield dicts for Text2TextProblem.generate_samples from lines of files."""
+  for inputs, (targets, extra_label) in zip(
+      txt_line_iterator(source_txt_path),
+      txt_and_label_iterator(target_txt_path)):
+    yield {"inputs": inputs, "targets": targets, "extra_label": [extra_label]}
+
+
 def text2text_distill_iterator(source_txt_path, target_txt_path,
                                distill_txt_path):
   """Yield dicts for Text2TextProblem.generate_samples from lines of files."""
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 1128ba63d..712fb5e08 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -58,7 +58,12 @@ def vocab_data_files(self):
     """Files to be passed to get_or_generate_vocab."""
     return self.source_data_files(problem.DatasetSplit.TRAIN)
 
-  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+  def generate_samples(
+      self,
+      data_dir,
+      tmp_dir,
+      dataset_split,
+      custom_iterator=text_problems.text2text_txt_iterator):
     datasets = self.source_data_files(dataset_split)
     tag = "dev"
     datatypes_to_clean = None
@@ -68,8 +73,8 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
     data_path = compile_data(
         tmp_dir, datasets, "%s-compiled-%s" % (self.name, tag),
         datatypes_to_clean=datatypes_to_clean)
-    return text_problems.text2text_txt_iterator(data_path + ".lang1",
-                                                data_path + ".lang2")
+
+    return custom_iterator(data_path + ".lang1", data_path + ".lang2")
 
   def generate_text_for_vocab(self, data_dir, tmp_dir):
     return generator_utils.generate_lines_for_vocab(tmp_dir,

From 8f9b0b2a522a53a2dde29bbd5eabf536bebfef5d Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 7 May 2019 14:06:26 -0700
Subject: [PATCH 2000/2720] Fail in trax if num_devices is not all of them (jax
 doesn't run in that config yet, let's fail for now to avoid surprises).

PiperOrigin-RevId: 247086543
---
 tensor2tensor/trax/trax.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 9810284ff..8ee28955b 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -407,7 +407,12 @@ def train(output_dir,
   """
   if save_steps is None:
     save_steps = []
-  num_devices = num_devices or jax.lib.xla_bridge.device_count()
+  device_count = jax.lib.xla_bridge.device_count()
+  num_devices = num_devices or device_count
+  # TODO(lukaszkaiser): remove this restriction when possible.
+  if num_devices != device_count:
+    raise ValueError("Jax cannot work yet with num_devices != all devices: "
+                     "%d != %d" % (num_devices, device_count))
   rng = get_random_number_generator_and_set_seed(random_seed)
   gfile.makedirs(output_dir)
   # Create summary writers and history.

From 84ab74010fedc1c0cdadb34f7ca35a1e6705d667 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 7 May 2019 14:12:20 -0700
Subject: [PATCH 2001/2720] trax.layers.attention.ShiftRight able to handle
 inputs of rank > 2. Will be useful for shifting (B,T) + OBS arrays on T.

PiperOrigin-RevId: 247087666
---
 tensor2tensor/trax/layers/attention.py      |  3 +-
 tensor2tensor/trax/layers/attention_test.py | 46 +++++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 tensor2tensor/trax/layers/attention_test.py

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 4fcd78c94..dd1d6461b 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -383,7 +383,8 @@ def ChunkedCausalMultiHeadedAttention(
 def ShiftRight(x, **unused_kwargs):
   """Layer to shift the tensor to the right by padding on axis 1."""
   if not isinstance(x, (list, tuple)):  # non-chunked inputs
-    pad_widths = [(0, 0), (1, 0)]
+    pad_widths = [(0, 0)] * len(x.shape)
+    pad_widths[1] = (1, 0)  # Padding on axis=1
     padded = np.pad(x, pad_widths, mode='constant')
     return padded[:, :-1]
   # Handling chunked inputs. Recall that the list of chunks represents a big
diff --git a/tensor2tensor/trax/layers/attention_test.py b/tensor2tensor/trax/layers/attention_test.py
new file mode 100644
index 000000000..cfa4a51bf
--- /dev/null
+++ b/tensor2tensor/trax/layers/attention_test.py
@@ -0,0 +1,46 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.layers.attention."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as onp
+from tensor2tensor.trax.layers import attention
+from tensorflow import test
+
+
+class AttentionTest(test.TestCase):
+
+  def test_shift_right(self):
+    # Test shifts right on axis=1
+    layer = attention.ShiftRight()
+    input_np = onp.arange(2*3*3).reshape(2, 3, 3)
+    output_np = layer(input_np)
+    self.assertEqual(input_np.shape, output_np.shape)
+    self.assertAllEqual(onp.array([[[0, 0, 0],
+                                    [0, 1, 2],
+                                    [3, 4, 5]],
+
+                                   [[0, 0, 0],
+                                    [9, 10, 11],
+                                    [12, 13, 14]]]),
+                        output_np)
+
+
+if __name__ == '__main__':
+  test.main()

From 3c3b3b1e38e185f56ebd4875584fa4923a2081de Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 7 May 2019 15:08:22 -0700
Subject: [PATCH 2002/2720] Disable evolved_transformer_test.py for the time
 being.

PiperOrigin-RevId: 247098725
---
 oss_scripts/oss_tests.sh | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 4d104a643..3fefdf5f5 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -41,6 +41,7 @@ set_status
 # test_utils.py is not a test, but pytest thinks it is.
 # subword_text_encoder_ops_test, pack_sequences_ops_test: interface with C++ ops
 # trax tests need C++
+# TODO(davidso): Re-enable EvolvedTransformer when possible.
 # others (see below) enable eager, so can't be tested along with the others in
 # pytest
 pytest --disable-warnings \
@@ -61,6 +62,7 @@ pytest --disable-warnings \
   --ignore=tensor2tensor/layers/modalities_test.py \
   --ignore=tensor2tensor/layers/ngram_test.py \
   --ignore=tensor2tensor/layers/reversible_layers_test.py \
+  --ignore=tensor2tensor/models/evolved_transformer_test.py \
   --ignore=tensor2tensor/models/research \
   --ignore=tensor2tensor/models/video/nfg_conv3d_test.py \
   --ignore=tensor2tensor/models/video/nfg_conv_lstm_test.py \

From 691abb19880664d8db5fb19cd34c0cb551465376 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 7 May 2019 15:56:44 -0700
Subject: [PATCH 2003/2720] trax - fix neural_gpu imports

PiperOrigin-RevId: 247107420
---
 tensor2tensor/trax/models/neural_gpu.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/trax/models/neural_gpu.py b/tensor2tensor/trax/models/neural_gpu.py
index 721f35c79..c578053df 100644
--- a/tensor2tensor/trax/models/neural_gpu.py
+++ b/tensor2tensor/trax/models/neural_gpu.py
@@ -17,7 +17,6 @@
 
 from __future__ import absolute_import
 from __future__ import division
-from __future__ import google_type_annotations
 from __future__ import print_function
 
 from tensor2tensor.trax import layers as tl

From ed6343ad565151a51e07a0d5df2d1644aad71e35 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 7 May 2019 16:07:20 -0700
Subject: [PATCH 2004/2720] Disable bayes_test for the time being.

PiperOrigin-RevId: 247109463
---
 oss_scripts/oss_tests.sh | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 3fefdf5f5..1e1920afc 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -83,10 +83,12 @@ set_status
 # Travis Error:
 # ImportError: /usr/lib/x86_64-linux-gnu/libstdc++.so.6: version `GLIBCXX_3.4.21' not found (required by /home/travis/virtualenv/python3.6.3/lib/python3.6/site-packages/jaxlib/_pywrap_xla.so)
 
+# TODO(trandustin): Re-enable bayes_test when we can.
+#  tensor2tensor/layers/bayes_test.py \
+
 # These tests enable eager, so are tested separately.
 pytest --disable-warnings \
   tensor2tensor/data_generators/problem_test.py \
-  tensor2tensor/layers/bayes_test.py \
   tensor2tensor/layers/common_attention_test.py \
   tensor2tensor/layers/common_layers_test.py \
   tensor2tensor/layers/common_video_test.py \

From 280d2f7360cedf1ea324cdbdf53cdc7164c837f8 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 7 May 2019 16:20:41 -0700
Subject: [PATCH 2005/2720] Trax: split layers into smaller files, move chunked
 Transformer to research, add tests.

PiperOrigin-RevId: 247111854
---
 tensor2tensor/trax/layers/__init__.py         |   4 +
 tensor2tensor/trax/layers/attention.py        | 127 --------
 tensor2tensor/trax/layers/convolution.py      | 152 +++++++++
 tensor2tensor/trax/layers/convolution_test.py |  36 +++
 tensor2tensor/trax/layers/core.py             | 292 +-----------------
 tensor2tensor/trax/layers/initializers.py     |  50 +++
 .../trax/layers/initializers_test.py          |  36 +++
 tensor2tensor/trax/layers/normalization.py    |  82 +++++
 .../trax/layers/normalization_test.py         |  42 +++
 tensor2tensor/trax/layers/pooling.py          |  80 +++++
 tensor2tensor/trax/layers/pooling_test.py     |  36 +++
 tensor2tensor/trax/layers/rnn.py              |   4 +-
 tensor2tensor/trax/models/__init__.py         |   3 +-
 .../models/research/chunked_transformer.py    | 230 ++++++++++++++
 tensor2tensor/trax/models/transformer.py      |  85 -----
 15 files changed, 770 insertions(+), 489 deletions(-)
 create mode 100644 tensor2tensor/trax/layers/convolution.py
 create mode 100644 tensor2tensor/trax/layers/convolution_test.py
 create mode 100644 tensor2tensor/trax/layers/initializers.py
 create mode 100644 tensor2tensor/trax/layers/initializers_test.py
 create mode 100644 tensor2tensor/trax/layers/normalization.py
 create mode 100644 tensor2tensor/trax/layers/normalization_test.py
 create mode 100644 tensor2tensor/trax/layers/pooling.py
 create mode 100644 tensor2tensor/trax/layers/pooling_test.py
 create mode 100644 tensor2tensor/trax/models/research/chunked_transformer.py

diff --git a/tensor2tensor/trax/layers/__init__.py b/tensor2tensor/trax/layers/__init__.py
index 25772a95f..87c3fe6b8 100644
--- a/tensor2tensor/trax/layers/__init__.py
+++ b/tensor2tensor/trax/layers/__init__.py
@@ -24,5 +24,9 @@
 from tensor2tensor.trax.layers.attention import *
 from tensor2tensor.trax.layers.base import *
 from tensor2tensor.trax.layers.combinators import *
+from tensor2tensor.trax.layers.convolution import *
 from tensor2tensor.trax.layers.core import *
+from tensor2tensor.trax.layers.initializers import *
+from tensor2tensor.trax.layers.normalization import *
+from tensor2tensor.trax.layers.pooling import *
 from tensor2tensor.trax.layers.rnn import *
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index dd1d6461b..7c5de1982 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -59,25 +59,6 @@ def EncoderDecoderMask(x, **unused_kwargs):
   return padding_mask + np.zeros((1, 1, decoder_input.shape[1], 1))
 
 
-# Layer normalization.
-def _layer_norm_new_params(input_shape, rng, epsilon=1e-6):  # pylint: disable=invalid-name
-  """Helper: create layer norm parameters."""
-  del rng, epsilon
-  features = input_shape[-1]
-  scale = np.ones(features)
-  bias = np.zeros(features)
-  return (scale, bias)
-
-
-@base.layer(new_parameters=_layer_norm_new_params)
-def LayerNorm(x, params, epsilon=1e-6, **unused_kwargs):
-  (scale, bias) = params
-  mean = np.mean(x, axis=-1, keepdims=True)
-  variance = np.mean((x - mean)**2, axis=-1, keepdims=True)
-  norm_inputs = (x - mean) / np.sqrt(variance + epsilon)
-  return norm_inputs * scale + bias
-
-
 # Positional encoding.
 def _positional_encoding_new_params(input_shape, rng, max_len=2048):  # pylint: disable=invalid-name
   """Helper: create positional encoding parameters."""
@@ -271,114 +252,6 @@ def MultiHeadedAttention(
   )
 
 
-# Chunked attention.
-def _chunked_selector_output_shape(  # pylint: disable=invalid-name
-    input_shapes, selector=None, **unused_kwargs):
-  """Helper: calculate output shape for chunked key selector (see below)."""
-  # Read the main function below first, the shape logic just follows the ops.
-  selector = selector or (lambda x: [] if x < 1 else [x-1])
-  triples, _ = zip(*input_shapes)
-  (query_shapes, key_shapes, value_shapes) = zip(*triples)
-  result = []
-  for i in range(len(input_shapes)):
-    selected = selector(i)
-    cur_key_shape, cur_value_shape = key_shapes[i], value_shapes[i]
-    # Since keys and values are [batch, length, depth] we concatenate on axis=1.
-    new_key_len = sum([key_shapes[j][1] for j in selected]) + cur_key_shape[1]
-    new_key_shape = (cur_key_shape[0], new_key_len, cur_key_shape[2])
-    new_value_len = sum(
-        [value_shapes[j][1] for j in selected]) + cur_value_shape[1]
-    new_value_shape = (cur_value_shape[0], new_value_len, cur_value_shape[2])
-    # Masks are (1, query-len, key-len).
-    new_mask_shape = (1, query_shapes[i][1], new_key_len)
-    new_shape = ((query_shapes[i], new_key_shape, new_value_shape),
-                 new_mask_shape)
-    result.append(new_shape)
-  return tuple(result)
-
-
-@base.layer(output_shape=_chunked_selector_output_shape)
-def ChunkedAttentionSelector(x, params, selector=None, **kwargs):
-  """Select which chunks to attend to in chunked attention.
-
-  Args:
-    x: inputs, a list of elements of the form (q, k, v), mask for each chunk.
-    params: parameters (unused).
-    selector: a function from chunk_number -> list of chunk numbers that says
-      which other chunks should be appended to the given one (previous if None).
-    **kwargs: unused other arguments.
-
-  Returns:
-    a list of elements of the form (q, k', v'), mask' where k', v' and mask' are
-    concatenations of k, v and identity-extended masks from selected chunks.
-  """
-  del params, kwargs
-  selector = selector or (lambda x: [] if x < 1 else [x-1])
-  triples, masks = zip(*x)
-  (queries, keys, values) = zip(*triples)
-  result = []
-  for i in range(len(x)):
-    selected = selector(i)
-    # Since keys and values are [batch, length, depth] we concatenate on axis=1.
-    # We also always include the current key or value at the end.
-    new_key_list = [keys[j] for j in selected]
-    new_key = np.concatenate(new_key_list + [keys[i]], axis=1)
-    new_value = np.concatenate(
-        [values[j] for j in selected] + [values[i]], axis=1)
-    # Masks are (1, query-len, key-len) so we concatenate on axis=2.
-    new_mask_shapes = [(1, queries[i].shape[1], key.shape[1])
-                       for key in new_key_list]
-    cur_mask = masks[i]
-    # Masks are all-1 for the added chunks (no masking).
-    new_mask_list = [np.ones(s, dtype=cur_mask.dtype) for s in new_mask_shapes]
-    # We still use the current (often causal) mask for the final chunk.
-    new_mask = np.concatenate(new_mask_list + [cur_mask], axis=2)
-    result.append(((queries[i], new_key, new_value), new_mask))
-  return tuple(result)
-
-
-def ChunkedCausalMultiHeadedAttention(
-    feature_depth, num_heads=8, dropout=0.0, chunk_selector=None, mode='train'):
-  """Transformer-style causal multi-headed attention operating on chunks.
-
-  Accepts inputs that are a list of chunks and applies causal attention.
-
-  Args:
-    feature_depth: int:  depth of embedding
-    num_heads: int: number of attention heads
-    dropout: float: dropout rate
-    chunk_selector: a function from chunk number to list of chunks to attend.
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    Multi-headed self-attention layer.
-  """
-  prepare_attention_input = combinators.Serial(
-      combinators.Branch(
-          combinators.Branch(  # q = k = v = first input
-              combinators.Copy(), combinators.Copy(), combinators.Copy()),
-          CausalMask(axis=-2),  # pylint: disable=no-value-for-parameter
-      ),
-      combinators.Parallel(
-          combinators.Parallel(
-              core.Dense(feature_depth),
-              core.Dense(feature_depth),
-              core.Dense(feature_depth),
-          ),
-          combinators.Copy()
-      )
-  )
-  return combinators.Serial(
-      combinators.Map(prepare_attention_input),
-      ChunkedAttentionSelector(selector=chunk_selector),  # pylint: disable=no-value-for-parameter
-      combinators.Map(PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
-          feature_depth=feature_depth, num_heads=num_heads,
-          dropout=dropout, mode=mode), check_shapes=False),
-      combinators.Map(combinators.Select(0), check_shapes=False),  # drop masks
-      combinators.Map(core.Dense(feature_depth))
-  )
-
-
 @base.layer()
 def ShiftRight(x, **unused_kwargs):
   """Layer to shift the tensor to the right by padding on axis 1."""
diff --git a/tensor2tensor/trax/layers/convolution.py b/tensor2tensor/trax/layers/convolution.py
new file mode 100644
index 000000000..4bf17f817
--- /dev/null
+++ b/tensor2tensor/trax/layers/convolution.py
@@ -0,0 +1,152 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Trax convolution layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import itertools
+
+from jax import lax
+
+import numpy as onp
+from tensor2tensor.trax.layers import base
+from tensor2tensor.trax.layers import initializers as init
+
+
+def PadtypeToPads(in_shape, window_shape, window_strides, padding):
+  """Convert padding string to list of pairs of pad values."""
+  padding = padding.upper()
+  if padding == 'SAME':
+    out_shape = onp.ceil(
+        onp.true_divide(in_shape, window_strides)).astype(int)
+    pad_sizes = [max((out_size - 1) * stride + window_shape - in_size, 0)
+                 for out_size, stride, window_shape, in_size
+                 in zip(out_shape, window_strides, window_shape, in_shape)]
+    return [(pad_size // 2, pad_size - pad_size // 2)
+            for pad_size in pad_sizes]
+  elif padding == 'VALID':
+    return [(0, 0)] * len(in_shape)
+  else:
+    msg = 'Unknown padding type: {}.'
+    raise TypeError(msg.format(padding))
+
+
+class Conv(base.Layer):
+  """Layer constructor function for a general convolution layer."""
+
+  def __init__(self, filters, kernel_size, strides=None, padding='VALID',
+               dimension_numbers=('NHWC', 'HWIO', 'NHWC'),
+               kernel_initializer=None,
+               bias_initializer=init.RandomNormalInitializer(1e-6)):
+    super(Conv, self).__init__()
+    self._filters = filters
+    self._kernel_size = kernel_size
+    self._padding = padding
+    self._dimension_numbers = dimension_numbers
+    self._lhs_spec, self._rhs_spec, self._out_spec = dimension_numbers
+    self._one = (1,) * len(kernel_size)
+    self._strides = strides or self._one
+    self._bias_initializer = bias_initializer
+    rhs_spec = self._rhs_spec
+    self._kernel_initializer = kernel_initializer
+    if kernel_initializer is None:
+      self._kernel_initializer = init.GlorotNormalInitializer(
+          rhs_spec.index('O'), rhs_spec.index('I'))
+
+  def call(self, x, params=(), **kwargs):
+    del kwargs
+    w, b = params
+    return lax.conv_general_dilated(
+        x, w, self._strides, self._padding, self._one, self._one,
+        self._dimension_numbers) + b
+
+  def _kernel_shape(self, input_shape):
+    """Helper to calculate the kernel shape."""
+    kernel_size_iter = iter(self._kernel_size)
+    return [self._filters if c == 'O' else
+            input_shape[self._lhs_spec.index('C')] if c == 'I' else
+            next(kernel_size_iter) for c in self._rhs_spec]
+
+  def _conv_shape_tuple(self, lhs_shape, rhs_shape, strides, pads):
+    """Compute the shape of a conv given input shapes in canonical order."""
+    if isinstance(pads, str):
+      pads = PadtypeToPads(lhs_shape[2:], rhs_shape[2:], strides, pads)
+    if len(pads) != len(lhs_shape) - 2:
+      msg = 'Wrong number of explicit pads for conv: expected {}, got {}.'
+      raise TypeError(msg.format(len(lhs_shape) - 2, len(pads)))
+    lhs_padded = onp.add(lhs_shape[2:], onp.add(*zip(*pads)))
+    out_space = onp.floor_divide(
+        onp.subtract(lhs_padded, rhs_shape[2:]), strides) + 1
+    out_space = onp.maximum(0, out_space)
+    out_shape = (lhs_shape[0], rhs_shape[0]) + tuple(out_space)
+    return tuple(out_shape)
+
+  def _conv_general_permutations(self, dimension_numbers):
+    """Utility for convolution dimension permutations relative to Conv HLO."""
+    lhs_spec, rhs_spec, out_spec = dimension_numbers
+    lhs_char, rhs_char, out_char = ('N', 'C'), ('O', 'I'), ('N', 'C')
+    charpairs = (lhs_char, rhs_char, out_char)
+    for i, (a, b) in enumerate(charpairs):
+      if not (dimension_numbers[i].count(a) == 1 and
+              dimension_numbers[i].count(b) == 1):
+        msg = ('convolution dimension_numbers[{}] must contain the characters '
+               '"{}" and "{}" exatly once, got {}.')
+        raise TypeError(msg.format(i, a, b, dimension_numbers[i]))
+      if len(dimension_numbers[i]) != len(set(dimension_numbers[i])):
+        msg = ('convolution dimension_numbers[{}] cannot have duplicate '
+               'characters, got {}.')
+        raise TypeError(msg.format(i, dimension_numbers[i]))
+    if not (set(lhs_spec) - set(lhs_char) == set(rhs_spec) - set(rhs_char) ==
+            set(out_spec) - set(out_char)):
+      msg = ('convolution dimension_numbers elements must each have the same '
+             'set of spatial characters, got {}.')
+      raise TypeError(msg.format(dimension_numbers))
+
+    def GetPerm(spec, charpair):
+      spatial = (i for i, c in enumerate(spec) if c not in charpair)
+      if spec is not rhs_spec:
+        spatial = sorted(spatial, key=lambda i: rhs_spec.index(spec[i]))
+      return (spec.index(charpair[0]), spec.index(charpair[1])) + tuple(spatial)
+
+    lhs_perm, rhs_perm, out_perm = map(GetPerm, dimension_numbers, charpairs)
+    return lhs_perm, rhs_perm, out_perm
+
+  def _conv_general_shape_tuple(self, lhs_shape, rhs_shape, window_strides,
+                                padding, dimension_numbers):
+    """Generalized computation of conv shape."""
+    lhs_perm, rhs_perm, out_perm = self._conv_general_permutations(
+        dimension_numbers)
+    lhs_trans = onp.take(lhs_shape, lhs_perm)
+    rhs_trans = onp.take(rhs_shape, rhs_perm)
+    out_trans = self._conv_shape_tuple(
+        lhs_trans, rhs_trans, window_strides, padding)
+    return tuple(onp.take(out_trans, onp.argsort(out_perm)))
+
+  def output_shape(self, input_shape):
+    kernel_shape = self._kernel_shape(input_shape)
+    return self._conv_general_shape_tuple(
+        input_shape, kernel_shape,
+        self._strides, self._padding, self._dimension_numbers)
+
+  def new_parameters(self, input_shape, rng):
+    kernel_shape = self._kernel_shape(input_shape)
+    bias_shape = [self._filters if c == 'C' else 1 for c in self._out_spec]
+    bias_shape = tuple(itertools.dropwhile(lambda x: x == 1, bias_shape))
+    w = self._kernel_initializer(kernel_shape, rng)
+    b = self._bias_initializer(bias_shape, rng)
+    return (w, b)
diff --git a/tensor2tensor/trax/layers/convolution_test.py b/tensor2tensor/trax/layers/convolution_test.py
new file mode 100644
index 000000000..d52db6289
--- /dev/null
+++ b/tensor2tensor/trax/layers/convolution_test.py
@@ -0,0 +1,36 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 convolution layers."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+from tensor2tensor.trax.layers import base
+from tensor2tensor.trax.layers import convolution
+
+
+class ConvolutionLayerTest(absltest.TestCase):
+
+  def test_conv(self):
+    input_shape = (29, 5, 5, 20)
+    result_shape = base.check_shape_agreement(
+        convolution.Conv(30, (3, 3)), input_shape)
+    self.assertEqual(result_shape, (29, 3, 3, 30))
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 31f3890d1..0872d63d6 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -19,65 +19,12 @@
 from __future__ import division
 from __future__ import print_function
 
-import itertools
 import operator as op
-
-from jax import lax
-
-import numpy as onp
 from six.moves import reduce
 from tensor2tensor.trax import backend
 from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers import base
-
-# Following the convention used in Keras and tf.layers, we use CamelCase for the
-# names of layer constructors, like Conv and Relu, while using snake_case for
-# other functions, like lax.conv and relu. To allow this, we disable below.
-# pylint: disable=invalid-name
-
-
-# Initializers.
-
-
-def RandomNormalInitializer(stddev=1e-2):
-  """An initializer function for random normal coefficients."""
-  def init(shape, rng):
-    return (stddev * backend.random.normal(rng, shape)).astype('float32')
-  return init
-
-
-def GlorotNormalInitializer(out_dim=0, in_dim=1, scale=onp.sqrt(2)):
-  """An initializer function for random Glorot-scaled coefficients."""
-  def init(shape, rng):
-    fan_in, fan_out = shape[in_dim], shape[out_dim]
-    size = onp.prod(onp.delete(shape, [in_dim, out_dim]))
-    std = scale / np.sqrt((fan_in + fan_out) / 2. * size)
-    return (std * backend.random.normal(rng, shape)).astype('float32')
-  return init
-
-
-def GlorotUniformInitializer(out_dim=0, in_dim=1):
-  """An initializer function for random uniform Glorot-scaled coefficients."""
-  def init(shape, rng):
-    fan_in, fan_out = shape[in_dim], shape[out_dim]
-    std = np.sqrt(2.0 / (fan_in + fan_out))
-    a = np.sqrt(3.0) * std
-    return backend.random.uniform(rng, shape, minval=-a, maxval=a)
-  return init
-
-
-def one_hot(x, size, dtype=np.float32):
-  """Make a n+1 dim one-hot array from n dim int-categorical array."""
-  return np.array(x[..., np.newaxis] == np.arange(size), dtype)
-
-
-# Layers.
-
-
-@base.layer()
-def AddConstant(x, params, constant=0.0, **unused_kwargs):
-  del params
-  return x + constant
+from tensor2tensor.trax.layers import initializers as init
 
 
 @base.layer()
@@ -135,8 +82,8 @@ class Dense(base.Layer):
   """Layer constructor function for a dense (fully-connected) layer."""
 
   def __init__(self, units,
-               kernel_initializer=GlorotUniformInitializer(),
-               bias_initializer=RandomNormalInitializer(1e-6)):
+               kernel_initializer=init.GlorotUniformInitializer(),
+               bias_initializer=init.RandomNormalInitializer(1e-6)):
     super(Dense, self).__init__()
     self._units = units
     self._kernel_initializer = kernel_initializer
@@ -161,7 +108,7 @@ class Embedding(base.Layer):
   """Layer constructor function for an embedding layer."""
 
   def __init__(self, feature_depth, vocab_size,
-               kernel_initializer=GlorotUniformInitializer()):
+               kernel_initializer=init.GlorotUniformInitializer()):
     super(Embedding, self).__init__()
     self._feature_depth = feature_depth
     self._vocab_size = vocab_size
@@ -179,130 +126,8 @@ def new_parameters(self, input_shape, rng):
         (self._vocab_size, self._feature_depth), rng)
 
 
-def padtype_to_pads(in_shape, window_shape, window_strides, padding):
-  """Convert padding string to list of pairs of pad values."""
-  padding = padding.upper()
-  if padding == 'SAME':
-    out_shape = onp.ceil(
-        onp.true_divide(in_shape, window_strides)).astype(int)
-    pad_sizes = [max((out_size - 1) * stride + window_shape - in_size, 0)
-                 for out_size, stride, window_shape, in_size
-                 in zip(out_shape, window_strides, window_shape, in_shape)]
-    return [(pad_size // 2, pad_size - pad_size // 2)
-            for pad_size in pad_sizes]
-  elif padding == 'VALID':
-    return [(0, 0)] * len(in_shape)
-  else:
-    msg = 'Unknown padding type: {}.'
-    raise TypeError(msg.format(padding))
-
-
-class Conv(base.Layer):
-  """Layer constructor function for a general convolution layer."""
-
-  def __init__(self, filters, kernel_size, strides=None, padding='VALID',
-               dimension_numbers=('NHWC', 'HWIO', 'NHWC'),
-               kernel_initializer=None,
-               bias_initializer=RandomNormalInitializer(1e-6)):
-    super(Conv, self).__init__()
-    self._filters = filters
-    self._kernel_size = kernel_size
-    self._padding = padding
-    self._dimension_numbers = dimension_numbers
-    self._lhs_spec, self._rhs_spec, self._out_spec = dimension_numbers
-    self._one = (1,) * len(kernel_size)
-    self._strides = strides or self._one
-    self._bias_initializer = bias_initializer
-    rhs_spec = self._rhs_spec
-    self._kernel_initializer = kernel_initializer or GlorotNormalInitializer(
-        rhs_spec.index('O'), rhs_spec.index('I'))
-
-  def call(self, x, params=(), **kwargs):
-    del kwargs
-    w, b = params
-    return lax.conv_general_dilated(
-        x, w, self._strides, self._padding, self._one, self._one,
-        self._dimension_numbers) + b
-
-  def _kernel_shape(self, input_shape):
-    """Helper to calculate the kernel shape."""
-    kernel_size_iter = iter(self._kernel_size)
-    return [self._filters if c == 'O' else
-            input_shape[self._lhs_spec.index('C')] if c == 'I' else
-            next(kernel_size_iter) for c in self._rhs_spec]
-
-  def _conv_shape_tuple(self, lhs_shape, rhs_shape, strides, pads):
-    """Compute the shape of a conv given input shapes in canonical order."""
-    if isinstance(pads, str):
-      pads = padtype_to_pads(lhs_shape[2:], rhs_shape[2:], strides, pads)
-    if len(pads) != len(lhs_shape) - 2:
-      msg = 'Wrong number of explicit pads for conv: expected {}, got {}.'
-      raise TypeError(msg.format(len(lhs_shape) - 2, len(pads)))
-    lhs_padded = onp.add(lhs_shape[2:], onp.add(*zip(*pads)))
-    out_space = onp.floor_divide(
-        onp.subtract(lhs_padded, rhs_shape[2:]), strides) + 1
-    out_space = onp.maximum(0, out_space)
-    out_shape = (lhs_shape[0], rhs_shape[0]) + tuple(out_space)
-    return tuple(out_shape)
-
-  def _conv_general_permutations(self, dimension_numbers):
-    """Utility for convolution dimension permutations relative to Conv HLO."""
-    lhs_spec, rhs_spec, out_spec = dimension_numbers
-    lhs_char, rhs_char, out_char = ('N', 'C'), ('O', 'I'), ('N', 'C')
-    charpairs = (lhs_char, rhs_char, out_char)
-    for i, (a, b) in enumerate(charpairs):
-      if not (dimension_numbers[i].count(a) == 1 and
-              dimension_numbers[i].count(b) == 1):
-        msg = ('convolution dimension_numbers[{}] must contain the characters '
-               '"{}" and "{}" exatly once, got {}.')
-        raise TypeError(msg.format(i, a, b, dimension_numbers[i]))
-      if len(dimension_numbers[i]) != len(set(dimension_numbers[i])):
-        msg = ('convolution dimension_numbers[{}] cannot have duplicate '
-               'characters, got {}.')
-        raise TypeError(msg.format(i, dimension_numbers[i]))
-    if not (set(lhs_spec) - set(lhs_char) == set(rhs_spec) - set(rhs_char) ==
-            set(out_spec) - set(out_char)):
-      msg = ('convolution dimension_numbers elements must each have the same '
-             'set of spatial characters, got {}.')
-      raise TypeError(msg.format(dimension_numbers))
-
-    def getperm(spec, charpair):
-      spatial = (i for i, c in enumerate(spec) if c not in charpair)
-      if spec is not rhs_spec:
-        spatial = sorted(spatial, key=lambda i: rhs_spec.index(spec[i]))
-      return (spec.index(charpair[0]), spec.index(charpair[1])) + tuple(spatial)
-
-    lhs_perm, rhs_perm, out_perm = map(getperm, dimension_numbers, charpairs)
-    return lhs_perm, rhs_perm, out_perm
-
-  def _conv_general_shape_tuple(self, lhs_shape, rhs_shape, window_strides,
-                                padding, dimension_numbers):
-    """Generalized computation of conv shape."""
-    lhs_perm, rhs_perm, out_perm = self._conv_general_permutations(
-        dimension_numbers)
-    lhs_trans = onp.take(lhs_shape, lhs_perm)
-    rhs_trans = onp.take(rhs_shape, rhs_perm)
-    out_trans = self._conv_shape_tuple(
-        lhs_trans, rhs_trans, window_strides, padding)
-    return tuple(onp.take(out_trans, onp.argsort(out_perm)))
-
-  def output_shape(self, input_shape):
-    kernel_shape = self._kernel_shape(input_shape)
-    return self._conv_general_shape_tuple(
-        input_shape, kernel_shape,
-        self._strides, self._padding, self._dimension_numbers)
-
-  def new_parameters(self, input_shape, rng):
-    kernel_shape = self._kernel_shape(input_shape)
-    bias_shape = [self._filters if c == 'C' else 1 for c in self._out_spec]
-    bias_shape = tuple(itertools.dropwhile(lambda x: x == 1, bias_shape))
-    w = self._kernel_initializer(kernel_shape, rng)
-    b = self._bias_initializer(bias_shape, rng)
-    return (w, b)
-
-
 # Flatten.
-def _flatten_output_shape(input_shape, num_axis_to_keep=1):
+def _flatten_output_shape(input_shape, num_axis_to_keep=1):  # pylint: disable=invalid-name
   """Output shape of a flatten layer."""
   if num_axis_to_keep >= len(input_shape):
     raise ValueError(
@@ -318,100 +143,6 @@ def Flatten(x, params, num_axis_to_keep=1, **kwargs):
   return np.reshape(x, (x.shape[:num_axis_to_keep] + (-1,)))
 
 
-# Batch normalization.
-def _batch_norm_new_params(input_shape, rng, axis=(0, 1, 2),
-                           center=True, scale=True, **kwargs):
-  """Helper to initialize batch norm params."""
-  del rng, kwargs
-  axis = (axis,) if np.isscalar(axis) else axis
-  shape = tuple(d for i, d in enumerate(input_shape) if i not in axis)
-  beta = np.zeros(shape, dtype='float32') if center else ()
-  gamma = np.ones(shape, dtype='float32') if scale else ()
-  return (beta, gamma)
-
-
-@base.layer(new_parameters=_batch_norm_new_params)
-def BatchNorm(x, params, axis=(0, 1, 2), epsilon=1e-5,
-              center=True, scale=True, **unused_kwargs):
-  """Layer construction function for a batch normalization layer."""
-  mean = np.mean(x, axis, keepdims=True)
-  # Fast but less numerically-stable variance calculation than np.var.
-  m1 = np.mean(x**2, axis, keepdims=True)
-  var = m1 - mean**2
-  z = (x - mean) / np.sqrt(var + epsilon)
-
-  # Expand the parameters to have the right axes.
-  beta, gamma = params
-  # TODO(phawkins): np.expand_dims should accept an axis tuple.
-  # (https://github.com/numpy/numpy/issues/12290)
-  ed = tuple(None if i in axis else slice(None) for i in range(np.ndim(x)))
-  beta = beta[ed]
-  gamma = gamma[ed]
-
-  # Return the z rescaled by the parameters if requested.
-  if center and scale:
-    return gamma * z + beta
-  if center:
-    return z + beta
-  if scale:
-    return gamma * z
-  return z
-
-
-# Pooling.
-def _pooling_output_shape(input_shape, pool_size=(2, 2),
-                          strides=None, padding='VALID'):
-  """Helper: compute the output shape for the pooling layer."""
-  dims = (1,) + pool_size + (1,)  # NHWC
-  spatial_strides = strides or (1,) * len(pool_size)
-  strides = (1,) + spatial_strides + (1,)
-  pads = padtype_to_pads(input_shape, dims, strides, padding)
-  operand_padded = onp.add(input_shape, onp.add(*zip(*pads)))
-  t = onp.floor_divide(onp.subtract(operand_padded, dims), strides) + 1
-  return tuple(t)
-
-
-def _pooling_general(inputs, reducer, init_val, rescaler=None,
-                     pool_size=(2, 2), strides=None, padding='VALID'):
-  """Helper: general pooling computation used in pooling layers later."""
-  spatial_strides = strides or (1,) * len(pool_size)
-  rescale = rescaler(pool_size, spatial_strides, padding) if rescaler else None
-  dims = (1,) + pool_size + (1,)  # NHWC
-  strides = (1,) + spatial_strides + (1,)
-  out = lax.reduce_window(inputs, init_val, reducer, dims, strides, padding)
-  return rescale(out, inputs) if rescale else out
-
-
-@base.layer(output_shape=_pooling_output_shape)
-def MaxPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
-  del params, kw
-  return _pooling_general(x, lax.max, -np.inf, pool_size=pool_size,
-                          strides=strides, padding=padding)
-
-
-@base.layer(output_shape=_pooling_output_shape)
-def SumPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
-  del params, kw
-  return _pooling_general(x, lax.add, 0., pool_size=pool_size,
-                          strides=strides, padding=padding)
-
-
-def _normalize_by_window_size(dims, spatial_strides, padding):
-  def rescale(outputs, inputs):
-    one = np.ones(inputs.shape[1:-1], dtype=inputs.dtype)
-    window_sizes = lax.reduce_window(
-        one, 0., lax.add, dims, spatial_strides, padding)
-    return outputs / window_sizes[..., np.newaxis]
-  return rescale
-
-
-@base.layer(output_shape=_pooling_output_shape)
-def AvgPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
-  del params, kw
-  return _pooling_general(x, lax.add, 0., _normalize_by_window_size,
-                          pool_size, strides=strides, padding=padding)
-
-
 @base.layer()
 def Dropout(x, params, rate=0.0, mode='train', rng=None, **kwargs):
   """Layer construction function for a dropout layer with given rate."""
@@ -436,8 +167,19 @@ def Div(x, params, divisor=1.0, **kwargs):
   return x / divisor
 
 
+@base.layer()
+def AddConstant(x, params, constant=0.0, **unused_kwargs):
+  del params
+  return x + constant
+
+
+def one_hot(x, size, dtype=np.float32):  # pylint: disable=invalid-name
+  """Make a n+1 dim one-hot array from n dim int-categorical array."""
+  return np.array(x[..., np.newaxis] == np.arange(size), dtype)
+
+
 # Mean.
-def _mean_output_shape(input_shape, axis=-1, keepdims=False):
+def _mean_output_shape(input_shape, axis=-1, keepdims=False):  # pylint: disable=invalid-name
   shape1 = list(input_shape)[:axis]  # Shape before axis.
   shape2 = list(input_shape)[axis:][1:]  # Shape after axis.
   mid_shape = [1] if keepdims else []
diff --git a/tensor2tensor/trax/layers/initializers.py b/tensor2tensor/trax/layers/initializers.py
new file mode 100644
index 000000000..9505d13f9
--- /dev/null
+++ b/tensor2tensor/trax/layers/initializers.py
@@ -0,0 +1,50 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Trax initializers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as onp
+from tensor2tensor.trax import backend
+
+
+def RandomNormalInitializer(stddev=1e-2):
+  """An initializer function for random normal coefficients."""
+  def Init(shape, rng):
+    return (stddev * backend.random.normal(rng, shape)).astype('float32')
+  return Init
+
+
+def GlorotNormalInitializer(out_dim=0, in_dim=1, scale=onp.sqrt(2)):
+  """An initializer function for random Glorot-scaled coefficients."""
+  def Init(shape, rng):
+    fan_in, fan_out = shape[in_dim], shape[out_dim]
+    size = onp.prod(onp.delete(shape, [in_dim, out_dim]))
+    std = scale / backend.numpy.sqrt((fan_in + fan_out) / 2. * size)
+    return (std * backend.random.normal(rng, shape)).astype('float32')
+  return Init
+
+
+def GlorotUniformInitializer(out_dim=0, in_dim=1):
+  """An initializer function for random uniform Glorot-scaled coefficients."""
+  def Init(shape, rng):
+    fan_in, fan_out = shape[in_dim], shape[out_dim]
+    std = backend.numpy.sqrt(2.0 / (fan_in + fan_out))
+    a = backend.numpy.sqrt(3.0) * std
+    return backend.random.uniform(rng, shape, minval=-a, maxval=a)
+  return Init
diff --git a/tensor2tensor/trax/layers/initializers_test.py b/tensor2tensor/trax/layers/initializers_test.py
new file mode 100644
index 000000000..bb5ad1b0c
--- /dev/null
+++ b/tensor2tensor/trax/layers/initializers_test.py
@@ -0,0 +1,36 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 initializers."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+from tensor2tensor.trax.backend import random
+from tensor2tensor.trax.layers import initializers
+
+
+class InitializersTest(absltest.TestCase):
+
+  def test_random_normal(self):
+    initializer = initializers.RandomNormalInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tensor2tensor/trax/layers/normalization.py b/tensor2tensor/trax/layers/normalization.py
new file mode 100644
index 000000000..deb8ef490
--- /dev/null
+++ b/tensor2tensor/trax/layers/normalization.py
@@ -0,0 +1,82 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Trax normalization layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.trax.backend import numpy as np
+from tensor2tensor.trax.layers import base
+
+
+# Batch normalization.
+def BatchNormParams(input_shape, rng, axis=(0, 1, 2),
+                    center=True, scale=True, **kwargs):
+  """Helper to initialize batch norm params."""
+  del rng, kwargs
+  axis = (axis,) if np.isscalar(axis) else axis
+  shape = tuple(d for i, d in enumerate(input_shape) if i not in axis)
+  beta = np.zeros(shape, dtype='float32') if center else ()
+  gamma = np.ones(shape, dtype='float32') if scale else ()
+  return (beta, gamma)
+
+
+@base.layer(new_parameters=BatchNormParams)
+def BatchNorm(x, params, axis=(0, 1, 2), epsilon=1e-5,
+              center=True, scale=True, **unused_kwargs):
+  """Layer construction function for a batch normalization layer."""
+  mean = np.mean(x, axis, keepdims=True)
+  # Fast but less numerically-stable variance calculation than np.var.
+  m1 = np.mean(x**2, axis, keepdims=True)
+  var = m1 - mean**2
+  z = (x - mean) / np.sqrt(var + epsilon)
+
+  # Expand the parameters to have the right axes.
+  beta, gamma = params
+  # TODO(phawkins): np.expand_dims should accept an axis tuple.
+  # (https://github.com/numpy/numpy/issues/12290)
+  ed = tuple(None if i in axis else slice(None) for i in range(np.ndim(x)))
+  beta = beta[ed]
+  gamma = gamma[ed]
+
+  # Return the z rescaled by the parameters if requested.
+  if center and scale:
+    return gamma * z + beta
+  if center:
+    return z + beta
+  if scale:
+    return gamma * z
+  return z
+
+
+# Layer normalization.
+def LayerNormParams(input_shape, rng, epsilon=1e-6):
+  """Helper: create layer norm parameters."""
+  del rng, epsilon
+  features = input_shape[-1]
+  scale = np.ones(features)
+  bias = np.zeros(features)
+  return (scale, bias)
+
+
+@base.layer(new_parameters=LayerNormParams)
+def LayerNorm(x, params, epsilon=1e-6, **unused_kwargs):
+  (scale, bias) = params
+  mean = np.mean(x, axis=-1, keepdims=True)
+  variance = np.mean((x - mean)**2, axis=-1, keepdims=True)
+  norm_inputs = (x - mean) / np.sqrt(variance + epsilon)
+  return norm_inputs * scale + bias
diff --git a/tensor2tensor/trax/layers/normalization_test.py b/tensor2tensor/trax/layers/normalization_test.py
new file mode 100644
index 000000000..894d31693
--- /dev/null
+++ b/tensor2tensor/trax/layers/normalization_test.py
@@ -0,0 +1,42 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 normalization layers."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+from tensor2tensor.trax.layers import base
+from tensor2tensor.trax.layers import normalization
+
+
+class NormalizationLayerTest(absltest.TestCase):
+
+  def test_batch_norm(self):
+    input_shape = (29, 5, 7, 20)
+    result_shape = base.check_shape_agreement(
+        normalization.BatchNorm(), input_shape)
+    self.assertEqual(result_shape, input_shape)
+
+  def test_layer_norm(self):
+    input_shape = (29, 5, 7, 20)
+    result_shape = base.check_shape_agreement(
+        normalization.LayerNorm(), input_shape)
+    self.assertEqual(result_shape, input_shape)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tensor2tensor/trax/layers/pooling.py b/tensor2tensor/trax/layers/pooling.py
new file mode 100644
index 000000000..c04b1aeb6
--- /dev/null
+++ b/tensor2tensor/trax/layers/pooling.py
@@ -0,0 +1,80 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Trax pooling layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from jax import lax
+
+import numpy as onp
+from tensor2tensor.trax.backend import numpy as np
+from tensor2tensor.trax.layers import base
+from tensor2tensor.trax.layers import convolution
+
+
+def PoolingOutputShape(input_shape, pool_size=(2, 2),
+                       strides=None, padding='VALID'):
+  """Helper: compute the output shape for the pooling layer."""
+  dims = (1,) + pool_size + (1,)  # NHWC
+  spatial_strides = strides or (1,) * len(pool_size)
+  strides = (1,) + spatial_strides + (1,)
+  pads = convolution.PadtypeToPads(input_shape, dims, strides, padding)
+  operand_padded = onp.add(input_shape, onp.add(*zip(*pads)))
+  t = onp.floor_divide(onp.subtract(operand_padded, dims), strides) + 1
+  return tuple(t)
+
+
+def PoolingGeneral(inputs, reducer, init_val, rescaler=None,
+                   pool_size=(2, 2), strides=None, padding='VALID'):
+  """Helper: general pooling computation used in pooling layers later."""
+  spatial_strides = strides or (1,) * len(pool_size)
+  rescale = rescaler(pool_size, spatial_strides, padding) if rescaler else None
+  dims = (1,) + pool_size + (1,)  # NHWC
+  strides = (1,) + spatial_strides + (1,)
+  out = lax.reduce_window(inputs, init_val, reducer, dims, strides, padding)
+  return rescale(out, inputs) if rescale else out
+
+
+@base.layer(output_shape=PoolingOutputShape)
+def MaxPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
+  del params, kw
+  return PoolingGeneral(x, lax.max, -np.inf, pool_size=pool_size,
+                        strides=strides, padding=padding)
+
+
+@base.layer(output_shape=PoolingOutputShape)
+def SumPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
+  del params, kw
+  return PoolingGeneral(x, lax.add, 0., pool_size=pool_size,
+                        strides=strides, padding=padding)
+
+
+def _normalize_by_window_size(dims, spatial_strides, padding):  # pylint: disable=invalid-name
+  def Rescale(outputs, inputs):
+    one = np.ones(inputs.shape[1:-1], dtype=inputs.dtype)
+    window_sizes = lax.reduce_window(
+        one, 0., lax.add, dims, spatial_strides, padding)
+    return outputs / window_sizes[..., np.newaxis]
+  return Rescale
+
+
+@base.layer(output_shape=PoolingOutputShape)
+def AvgPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
+  del params, kw
+  return PoolingGeneral(x, lax.add, 0., _normalize_by_window_size,
+                        pool_size, strides=strides, padding=padding)
diff --git a/tensor2tensor/trax/layers/pooling_test.py b/tensor2tensor/trax/layers/pooling_test.py
new file mode 100644
index 000000000..8924a8b5e
--- /dev/null
+++ b/tensor2tensor/trax/layers/pooling_test.py
@@ -0,0 +1,36 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 conv layers."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+from tensor2tensor.trax.layers import base
+from tensor2tensor.trax.layers import pooling
+
+
+class PoolingLayerTest(absltest.TestCase):
+
+  def test_avg_pool(self):
+    input_shape = (29, 4, 4, 20)
+    result_shape = base.check_shape_agreement(
+        pooling.AvgPool(pool_size=(2, 2), strides=(2, 2)), input_shape)
+    self.assertEqual(result_shape, (29, 2, 2, 20))
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tensor2tensor/trax/layers/rnn.py b/tensor2tensor/trax/layers/rnn.py
index 9bfd60d83..b0caa2bbc 100644
--- a/tensor2tensor/trax/layers/rnn.py
+++ b/tensor2tensor/trax/layers/rnn.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 from tensor2tensor.trax.layers import combinators
+from tensor2tensor.trax.layers import convolution
 from tensor2tensor.trax.layers import core
 
 
@@ -56,7 +57,8 @@ def ConvGRUCell(units, kernel_size=(3, 3)):
   """
 
   def BuildConv():
-    return core.Conv(filters=units, kernel_size=kernel_size, padding='SAME')
+    return convolution.Conv(
+        filters=units, kernel_size=kernel_size, padding='SAME')
 
   return GeneralGRUCell(
       candidate_transform=BuildConv,
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index e90ab053b..2a0117725 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -24,6 +24,7 @@
 from tensor2tensor.trax.models import neural_gpu
 from tensor2tensor.trax.models import resnet
 from tensor2tensor.trax.models import transformer
+from tensor2tensor.trax.models.research import chunked_transformer
 
 
 # Ginify
@@ -33,7 +34,7 @@ def model_configure(*args, **kwargs):
 
 
 # pylint: disable=invalid-name
-ChunkedTransformerLM = model_configure(transformer.ChunkedTransformerLM)
+ChunkedTransformerLM = model_configure(chunked_transformer.ChunkedTransformerLM)
 MLP = model_configure(mlp.MLP)
 NeuralGPU = model_configure(neural_gpu.NeuralGPU)
 Resnet50 = model_configure(resnet.Resnet50)
diff --git a/tensor2tensor/trax/models/research/chunked_transformer.py b/tensor2tensor/trax/models/research/chunked_transformer.py
new file mode 100644
index 000000000..cb200b0db
--- /dev/null
+++ b/tensor2tensor/trax/models/research/chunked_transformer.py
@@ -0,0 +1,230 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Chunked Transformer Models."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.trax import layers as tl
+from tensor2tensor.trax.backend import numpy as np
+
+
+# Chunked attention.
+def _chunked_selector_output_shape(  # pylint: disable=invalid-name
+    input_shapes, selector=None, **unused_kwargs):
+  """Helper: calculate output shape for chunked key selector (see below)."""
+  # Read the main function below first, the shape logic just follows the ops.
+  selector = selector or (lambda x: [] if x < 1 else [x-1])
+  triples, _ = zip(*input_shapes)
+  (query_shapes, key_shapes, value_shapes) = zip(*triples)
+  result = []
+  for i in range(len(input_shapes)):
+    selected = selector(i)
+    cur_key_shape, cur_value_shape = key_shapes[i], value_shapes[i]
+    # Since keys and values are [batch, length, depth] we concatenate on axis=1.
+    new_key_len = sum([key_shapes[j][1] for j in selected]) + cur_key_shape[1]
+    new_key_shape = (cur_key_shape[0], new_key_len, cur_key_shape[2])
+    new_value_len = sum(
+        [value_shapes[j][1] for j in selected]) + cur_value_shape[1]
+    new_value_shape = (cur_value_shape[0], new_value_len, cur_value_shape[2])
+    # Masks are (1, query-len, key-len).
+    new_mask_shape = (1, query_shapes[i][1], new_key_len)
+    new_shape = ((query_shapes[i], new_key_shape, new_value_shape),
+                 new_mask_shape)
+    result.append(new_shape)
+  return tuple(result)
+
+
+@tl.layer(output_shape=_chunked_selector_output_shape)
+def ChunkedAttentionSelector(x, params, selector=None, **kwargs):
+  """Select which chunks to attend to in chunked attention.
+
+  Args:
+    x: inputs, a list of elements of the form (q, k, v), mask for each chunk.
+    params: parameters (unused).
+    selector: a function from chunk_number -> list of chunk numbers that says
+      which other chunks should be appended to the given one (previous if None).
+    **kwargs: unused other arguments.
+
+  Returns:
+    a list of elements of the form (q, k', v'), mask' where k', v' and mask' are
+    concatenations of k, v and identity-extended masks from selected chunks.
+  """
+  del params, kwargs
+  selector = selector or (lambda x: [] if x < 1 else [x-1])
+  triples, masks = zip(*x)
+  (queries, keys, values) = zip(*triples)
+  result = []
+  for i in range(len(x)):
+    selected = selector(i)
+    # Since keys and values are [batch, length, depth] we concatenate on axis=1.
+    # We also always include the current key or value at the end.
+    new_key_list = [keys[j] for j in selected]
+    new_key = np.concatenate(new_key_list + [keys[i]], axis=1)
+    new_value = np.concatenate(
+        [values[j] for j in selected] + [values[i]], axis=1)
+    # Masks are (1, query-len, key-len) so we concatenate on axis=2.
+    new_mask_shapes = [(1, queries[i].shape[1], key.shape[1])
+                       for key in new_key_list]
+    cur_mask = masks[i]
+    # Masks are all-1 for the added chunks (no masking).
+    new_mask_list = [np.ones(s, dtype=cur_mask.dtype) for s in new_mask_shapes]
+    # We still use the current (often causal) mask for the final chunk.
+    new_mask = np.concatenate(new_mask_list + [cur_mask], axis=2)
+    result.append(((queries[i], new_key, new_value), new_mask))
+  return tuple(result)
+
+
+def ChunkedCausalMultiHeadedAttention(
+    feature_depth, num_heads=8, dropout=0.0, chunk_selector=None, mode='train'):
+  """Transformer-style causal multi-headed attention operating on chunks.
+
+  Accepts inputs that are a list of chunks and applies causal attention.
+
+  Args:
+    feature_depth: int:  depth of embedding
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate
+    chunk_selector: a function from chunk number to list of chunks to attend.
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    Multi-headed self-attention layer.
+  """
+  prepare_attention_input = tl.Serial(
+      tl.Branch(
+          tl.Branch(  # q = k = v = first input
+              tl.Copy(), tl.Copy(), tl.Copy()),
+          tl.CausalMask(axis=-2),
+      ),
+      tl.Parallel(
+          tl.Parallel(
+              tl.Dense(feature_depth),
+              tl.Dense(feature_depth),
+              tl.Dense(feature_depth),
+          ),
+          tl.Copy()
+      )
+  )
+  return tl.Serial(
+      tl.Map(prepare_attention_input),
+      ChunkedAttentionSelector(selector=chunk_selector),  # pylint: disable=no-value-for-parameter
+      tl.Map(tl.PureMultiHeadedAttention(
+          feature_depth=feature_depth, num_heads=num_heads,
+          dropout=dropout, mode=mode), check_shapes=False),
+      tl.Map(tl.Select(0), check_shapes=False),  # drop masks
+      tl.Map(tl.Dense(feature_depth))
+  )
+
+
+def ResidualFeedForward(feature_depth,
+                        feedforward_depth,
+                        dropout,
+                        mode):
+  """Residual feed-forward layer with normalization at start."""
+  return tl.Residual(
+      tl.LayerNorm(),
+      tl.Dense(feedforward_depth),
+      tl.Relu(),
+      tl.Dropout(rate=dropout, mode=mode),
+      tl.Dense(feature_depth),
+      tl.Dropout(rate=dropout, mode=mode)
+  )
+
+
+def ChunkedDecoderLayer(feature_depth,
+                        feedforward_depth,
+                        num_heads,
+                        dropout,
+                        chunk_selector,
+                        mode):
+  """Transformer decoder layer operating on chunks.
+
+  Args:
+    feature_depth: int:  depth of embedding
+    feedforward_depth: int: depth of feed-forward layer
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate (how much to drop out)
+    chunk_selector: a function from chunk number to list of chunks to attend.
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    the layer.
+  """
+  return tl.Serial(
+      tl.Residual(  # Self-attention block.
+          tl.Map(tl.LayerNorm()),
+          ChunkedCausalMultiHeadedAttention(
+              feature_depth, num_heads=num_heads, dropout=dropout,
+              chunk_selector=chunk_selector, mode=mode),
+          tl.Map(tl.Dropout(rate=dropout, mode=mode)),
+      ),
+      tl.Map(ResidualFeedForward(
+          feature_depth, feedforward_depth, dropout, mode=mode))
+  )
+
+
+def ChunkedTransformerLM(vocab_size,
+                         feature_depth=512,
+                         feedforward_depth=2048,
+                         num_layers=6,
+                         num_heads=8,
+                         dropout=0.1,
+                         chunk_selector=None,
+                         max_len=2048,
+                         mode='train'):
+  """Transformer language model operating on chunks.
+
+  The input to this  model is a sequence presented as a list or tuple of chunks:
+    (chunk1, chunk2, chunks3, ..., chunkN).
+  Each chunk should have the same shape (batch, chunk-length) and together they
+  represent a long sequence that's a concatenation chunk1,chunk2,...,chunkN.
+
+  Chunked Transformer emulates the operation of a Transformer on this long
+  sequence except for the chunked attention layer, which may attend to only
+  a subset of the chunks to reduce memory use.
+
+  Args:
+    vocab_size: int: vocab size
+    feature_depth: int:  depth of embedding
+    feedforward_depth: int: depth of feed-forward layer
+    num_layers: int: number of encoder/decoder layers
+    num_heads: int: number of attention heads
+    dropout: float: dropout rate (how much to drop out)
+    chunk_selector: a function from chunk number to list of chunks to attend
+      (if None, attends to the previous chunks which is equivalent to setting
+       chunk_selector(x) = [] if x < 1 else [x-1] (TransformerXL); we attend
+       to the current chunk with a causal mask too, selected chunks unmasked).
+    max_len: int: maximum symbol length for positional encoding
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    the layer.
+  """
+  stack = [ChunkedDecoderLayer(feature_depth, feedforward_depth, num_heads,
+                               dropout, chunk_selector, mode)
+           for _ in range(num_layers)]
+  # Below each Map(L) applies the layer L to each chunk independently.
+  return tl.Serial(
+      tl.ShiftRight(),
+      tl.Map(tl.Embedding(feature_depth, vocab_size)),
+      tl.Map(tl.Dropout(rate=dropout, mode=mode)),
+      tl.PositionalEncoding(max_len=max_len),
+      tl.Serial(*stack),
+      tl.Map(tl.LayerNorm()),
+      tl.Map(tl.Dense(vocab_size)),
+      tl.Map(tl.LogSoftmax()),
+  )
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 85e5997bf..7cf3486b3 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -182,91 +182,6 @@ def TransformerLM(vocab_size,
   )
 
 
-def ChunkedDecoderLayer(feature_depth,
-                        feedforward_depth,
-                        num_heads,
-                        dropout,
-                        chunk_selector,
-                        mode):
-  """Transformer decoder layer operating on chunks.
-
-  Args:
-    feature_depth: int:  depth of embedding
-    feedforward_depth: int: depth of feed-forward layer
-    num_heads: int: number of attention heads
-    dropout: float: dropout rate (how much to drop out)
-    chunk_selector: a function from chunk number to list of chunks to attend.
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    the layer.
-  """
-  return tl.Serial(
-      tl.Residual(  # Self-attention block.
-          tl.Map(tl.LayerNorm()),
-          tl.ChunkedCausalMultiHeadedAttention(
-              feature_depth, num_heads=num_heads, dropout=dropout,
-              chunk_selector=chunk_selector, mode=mode),
-          tl.Map(tl.Dropout(rate=dropout, mode=mode)),
-      ),
-      tl.Map(ResidualFeedForward(
-          feature_depth, feedforward_depth, dropout, mode=mode))
-  )
-
-
-def ChunkedTransformerLM(vocab_size,
-                         feature_depth=512,
-                         feedforward_depth=2048,
-                         num_layers=6,
-                         num_heads=8,
-                         dropout=0.1,
-                         chunk_selector=None,
-                         max_len=2048,
-                         mode='train'):
-  """Transformer language model operating on chunks.
-
-  The input to this  model is a sequence presented as a list or tuple of chunks:
-    (chunk1, chunk2, chunks3, ..., chunkN).
-  Each chunk should have the same shape (batch, chunk-length) and together they
-  represent a long sequence that's a concatenation chunk1,chunk2,...,chunkN.
-
-  Chunked Transformer emulates the operation of a Transformer on this long
-  sequence except for the chunked attention layer, which may attend to only
-  a subset of the chunks to reduce memory use.
-
-  Args:
-    vocab_size: int: vocab size
-    feature_depth: int:  depth of embedding
-    feedforward_depth: int: depth of feed-forward layer
-    num_layers: int: number of encoder/decoder layers
-    num_heads: int: number of attention heads
-    dropout: float: dropout rate (how much to drop out)
-    chunk_selector: a function from chunk number to list of chunks to attend
-      (if None, attends to the previous chunks which is equivalent to setting
-       chunk_selector(x) = [] if x < 1 else [x-1] (TransformerXL); we attend
-       to the current chunk with a causal mask too, selected chunks unmasked).
-    max_len: int: maximum symbol length for positional encoding
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    the layer.
-  """
-  stack = [ChunkedDecoderLayer(feature_depth, feedforward_depth, num_heads,
-                               dropout, chunk_selector, mode)
-           for _ in range(num_layers)]
-  # Below each Map(L) applies the layer L to each chunk independently.
-  return tl.Serial(
-      tl.ShiftRight(),
-      tl.Map(tl.Embedding(feature_depth, vocab_size)),
-      tl.Map(tl.Dropout(rate=dropout, mode=mode)),
-      tl.PositionalEncoding(max_len=max_len),
-      tl.Serial(*stack),
-      tl.Map(tl.LayerNorm()),
-      tl.Map(tl.Dense(vocab_size)),
-      tl.Map(tl.LogSoftmax()),
-  )
-
-
 def EncoderDecoderLayer(feature_depth,
                         feedforward_depth,
                         num_heads,

From ddb9665eec22be5136a6e076929e19667a6d1c98 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 7 May 2019 17:04:36 -0700
Subject: [PATCH 2006/2720] Internal change.

PiperOrigin-RevId: 247119337
---
 tensor2tensor/problems_colab.py | 36 +++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)
 create mode 100644 tensor2tensor/problems_colab.py

diff --git a/tensor2tensor/problems_colab.py b/tensor2tensor/problems_colab.py
new file mode 100644
index 000000000..12d679659
--- /dev/null
+++ b/tensor2tensor/problems_colab.py
@@ -0,0 +1,36 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Access T2T Problems."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.data_generators import all_problems
+from tensor2tensor.utils import registry
+
+
+def problem(name):
+  return registry.problem(name)
+
+
+def available():
+  return sorted(registry.list_problems())
+
+
+# Import problem modules
+_modules = list(all_problems.MODULES)
+
+all_problems.import_modules(_modules)

From a298052cc245cd9081fe507c5386ce04a91c90d0 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 7 May 2019 17:23:53 -0700
Subject: [PATCH 2007/2720] Splits out the transformer decoder layers so they
 may be called separately.

PiperOrigin-RevId: 247122386
---
 tensor2tensor/models/transformer.py | 255 ++++++++++++++++------------
 1 file changed, 149 insertions(+), 106 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 5b5a1b1a4..1eac82db5 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1336,6 +1336,133 @@ def transformer_prepare_decoder(targets, hparams, features=None):
   return (decoder_input, decoder_self_attention_bias)
 
 
+def transformer_decoder_layer(decoder_input,
+                              decoder_self_attention_bias,
+                              layer_idx,
+                              hparams,
+                              encoder_output=None,
+                              encoder_decoder_attention_bias=None,
+                              cache=None,
+                              decode_loop_step=None,
+                              nonpadding=None,
+                              save_weights_to=None,
+                              make_image_summary=False,
+                              losses=None,
+                              layer_collection=None,
+                              recurrent_memory_by_layer=None,
+                              chunk_number=None):
+  """A single transformer decoder layer."""
+  x = decoder_input
+  layer = layer_idx
+  layer_name = "layer_%d" % layer
+  layer_cache = cache[layer_name] if cache is not None else None
+
+  attention_dropout_broadcast_dims = (
+      common_layers.comma_separated_string_to_integer_list(
+          getattr(hparams, "attention_dropout_broadcast_dims", "")))
+
+  if recurrent_memory_by_layer is not None:
+    recurrent_memory = recurrent_memory_by_layer[layer_name]
+  else:
+    recurrent_memory = None
+
+  if layer < hparams.get("num_area_layers", 0):
+    max_area_width = hparams.get("max_area_width", 1)
+    max_area_height = hparams.get("max_area_height", 1)
+    memory_height = hparams.get("max_area_height", 1)
+  else:
+    max_area_width = 1
+    max_area_height = 1
+    memory_height = 1
+  with tf.variable_scope(layer_name):
+    with tf.variable_scope("self_attention"):
+      y = common_attention.multihead_attention(
+          common_layers.layer_preprocess(
+              x, hparams, layer_collection=layer_collection),
+          None,
+          decoder_self_attention_bias,
+          hparams.attention_key_channels or hparams.hidden_size,
+          hparams.attention_value_channels or hparams.hidden_size,
+          hparams.hidden_size,
+          hparams.num_heads,
+          hparams.attention_dropout,
+          attention_type=hparams.self_attention_type,
+          max_relative_position=hparams.max_relative_position,
+          heads_share_relative_embedding=(
+              hparams.heads_share_relative_embedding),
+          add_relative_to_values=hparams.add_relative_to_values,
+          save_weights_to=save_weights_to,
+          cache=layer_cache,
+          make_image_summary=make_image_summary,
+          dropout_broadcast_dims=attention_dropout_broadcast_dims,
+          max_length=hparams.get("max_length"),
+          decode_loop_step=decode_loop_step,
+          vars_3d=hparams.get("attention_variables_3d"),
+          activation_dtype=hparams.get("activation_dtype", "float32"),
+          weight_dtype=hparams.get("weight_dtype", "float32"),
+          layer_collection=layer_collection,
+          recurrent_memory=recurrent_memory,
+          chunk_number=chunk_number,
+          hard_attention_k=hparams.get("hard_attention_k", 0),
+          max_area_width=max_area_width,
+          max_area_height=max_area_height,
+          memory_height=memory_height,
+          area_key_mode=hparams.get("area_key_mode", "none"),
+          area_value_mode=hparams.get("area_value_mode", "none"),
+          training=(hparams.get(
+              "mode",
+              tf.estimator.ModeKeys.TRAIN) == tf.estimator.ModeKeys.TRAIN))
+      x = common_layers.layer_postprocess(x, y, hparams)
+    if encoder_output is not None:
+      with tf.variable_scope("encdec_attention"):
+        y = common_attention.multihead_attention(
+            common_layers.layer_preprocess(
+                x, hparams, layer_collection=layer_collection),
+            encoder_output,
+            encoder_decoder_attention_bias,
+            hparams.attention_key_channels or hparams.hidden_size,
+            hparams.attention_value_channels or hparams.hidden_size,
+            hparams.hidden_size,
+            hparams.num_heads,
+            hparams.attention_dropout,
+            max_relative_position=hparams.max_relative_position,
+            heads_share_relative_embedding=(
+                hparams.heads_share_relative_embedding),
+            add_relative_to_values=hparams.add_relative_to_values,
+            save_weights_to=save_weights_to,
+            cache=layer_cache,
+            make_image_summary=make_image_summary,
+            dropout_broadcast_dims=attention_dropout_broadcast_dims,
+            max_length=hparams.get("max_length"),
+            vars_3d=hparams.get("attention_variables_3d"),
+            activation_dtype=hparams.get("activation_dtype", "float32"),
+            weight_dtype=hparams.get("weight_dtype", "float32"),
+            layer_collection=layer_collection,
+            hard_attention_k=hparams.get("hard_attention_k", 0),
+            max_area_width=max_area_width,
+            max_area_height=max_area_height,
+            memory_height=memory_height,
+            area_key_mode=hparams.get("area_key_mode", "none"),
+            area_value_mode=hparams.get("area_value_mode", "none"),
+            training=(hparams.get(
+                "mode",
+                tf.estimator.ModeKeys.TRAIN) == tf.estimator.ModeKeys.TRAIN))
+        x = common_layers.layer_postprocess(x, y, hparams)
+    with tf.variable_scope("ffn"):
+      y = transformer_ffn_layer(
+          common_layers.layer_preprocess(
+              x, hparams, layer_collection=layer_collection),
+          hparams,
+          conv_padding="LEFT",
+          nonpadding_mask=nonpadding,
+          losses=losses,
+          cache=layer_cache,
+          decode_loop_step=decode_loop_step,
+          layer_collection=layer_collection)
+      x = common_layers.layer_postprocess(x, y, hparams)
+      return x
+
+
 def transformer_decoder(decoder_input,
                         encoder_output,
                         decoder_self_attention_bias,
@@ -1350,8 +1477,7 @@ def transformer_decoder(decoder_input,
                         losses=None,
                         layer_collection=None,
                         recurrent_memory_by_layer=None,
-                        chunk_number=None,
-                        ):
+                        chunk_number=None):
   """A stack of transformer layers.
 
   Args:
@@ -1377,8 +1503,8 @@ def transformer_decoder(decoder_input,
       key created from the variable scope (including name).
     make_image_summary: Whether to make an attention image summary.
     losses: optional list onto which to append extra training losses
-    layer_collection: A tensorflow_kfac.LayerCollection. Only used by the
-      KFAC optimizer. Default is None.
+    layer_collection: A tensorflow_kfac.LayerCollection. Only used by the KFAC
+      optimizer. Default is None.
     recurrent_memory_by_layer: Optional dict, mapping layer names to instances
       of transformer_memory.RecurrentMemory. Default is None.
     chunk_number: an optional integer Tensor with shape [batch] used to operate
@@ -1388,9 +1514,6 @@ def transformer_decoder(decoder_input,
     y: a Tensors
   """
   x = decoder_input
-  attention_dropout_broadcast_dims = (
-      common_layers.comma_separated_string_to_integer_list(
-          getattr(hparams, "attention_dropout_broadcast_dims", "")))
 
   mlperf_log.transformer_print(
       key=mlperf_log.MODEL_HP_NUM_HIDDEN_LAYERS,
@@ -1410,106 +1533,26 @@ def transformer_decoder(decoder_input,
       hparams=hparams)
 
   with tf.variable_scope(name):
-    for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
-      layer_name = "layer_%d" % layer
-      layer_cache = cache[layer_name] if cache is not None else None
-      if recurrent_memory_by_layer is not None:
-        recurrent_memory = recurrent_memory_by_layer[layer_name]
-      else:
-        recurrent_memory = None
+    for layer_idx in range(hparams.num_decoder_layers or
+                           hparams.num_hidden_layers):
+      x = transformer_decoder_layer(
+          x,
+          decoder_self_attention_bias,
+          layer_idx,
+          hparams,
+          encoder_decoder_attention_bias=encoder_decoder_attention_bias,
+          encoder_output=encoder_output,
+          cache=cache,
+          decode_loop_step=decode_loop_step,
+          nonpadding=nonpadding,
+          save_weights_to=save_weights_to,
+          make_image_summary=make_image_summary,
+          losses=losses,
+          layer_collection=layer_collection,
+          recurrent_memory_by_layer=recurrent_memory_by_layer,
+          chunk_number=chunk_number,
+      )
 
-      if layer < hparams.get("num_area_layers", 0):
-        max_area_width = hparams.get("max_area_width", 1)
-        max_area_height = hparams.get("max_area_height", 1)
-        memory_height = hparams.get("max_area_height", 1)
-      else:
-        max_area_width = 1
-        max_area_height = 1
-        memory_height = 1
-      with tf.variable_scope(layer_name):
-        with tf.variable_scope("self_attention"):
-          y = common_attention.multihead_attention(
-              common_layers.layer_preprocess(
-                  x, hparams, layer_collection=layer_collection),
-              None,
-              decoder_self_attention_bias,
-              hparams.attention_key_channels or hparams.hidden_size,
-              hparams.attention_value_channels or hparams.hidden_size,
-              hparams.hidden_size,
-              hparams.num_heads,
-              hparams.attention_dropout,
-              attention_type=hparams.self_attention_type,
-              max_relative_position=hparams.max_relative_position,
-              heads_share_relative_embedding=(
-                  hparams.heads_share_relative_embedding),
-              add_relative_to_values=hparams.add_relative_to_values,
-              save_weights_to=save_weights_to,
-              cache=layer_cache,
-              make_image_summary=make_image_summary,
-              dropout_broadcast_dims=attention_dropout_broadcast_dims,
-              max_length=hparams.get("max_length"),
-              decode_loop_step=decode_loop_step,
-              vars_3d=hparams.get("attention_variables_3d"),
-              activation_dtype=hparams.get("activation_dtype", "float32"),
-              weight_dtype=hparams.get("weight_dtype", "float32"),
-              layer_collection=layer_collection,
-              recurrent_memory=recurrent_memory,
-              chunk_number=chunk_number,
-              hard_attention_k=hparams.get("hard_attention_k", 0),
-              max_area_width=max_area_width,
-              max_area_height=max_area_height,
-              memory_height=memory_height,
-              area_key_mode=hparams.get("area_key_mode", "none"),
-              area_value_mode=hparams.get("area_value_mode", "none"),
-              training=(hparams.get("mode", tf.estimator.ModeKeys.TRAIN)
-                        == tf.estimator.ModeKeys.TRAIN))
-          x = common_layers.layer_postprocess(x, y, hparams)
-        if encoder_output is not None:
-          with tf.variable_scope("encdec_attention"):
-            y = common_attention.multihead_attention(
-                common_layers.layer_preprocess(
-                    x, hparams, layer_collection=layer_collection),
-                encoder_output,
-                encoder_decoder_attention_bias,
-                hparams.attention_key_channels or hparams.hidden_size,
-                hparams.attention_value_channels or hparams.hidden_size,
-                hparams.hidden_size,
-                hparams.num_heads,
-                hparams.attention_dropout,
-                max_relative_position=hparams.max_relative_position,
-                heads_share_relative_embedding=(
-                    hparams.heads_share_relative_embedding),
-                add_relative_to_values=hparams.add_relative_to_values,
-                save_weights_to=save_weights_to,
-                cache=layer_cache,
-                make_image_summary=make_image_summary,
-                dropout_broadcast_dims=attention_dropout_broadcast_dims,
-                max_length=hparams.get("max_length"),
-                vars_3d=hparams.get("attention_variables_3d"),
-                activation_dtype=hparams.get("activation_dtype", "float32"),
-                weight_dtype=hparams.get("weight_dtype", "float32"),
-                layer_collection=layer_collection,
-                hard_attention_k=hparams.get("hard_attention_k", 0),
-                max_area_width=max_area_width,
-                max_area_height=max_area_height,
-                memory_height=memory_height,
-                area_key_mode=hparams.get("area_key_mode", "none"),
-                area_value_mode=hparams.get("area_value_mode", "none"),
-                training=(hparams.get("mode", tf.estimator.ModeKeys.TRAIN)
-                          == tf.estimator.ModeKeys.TRAIN))
-            x = common_layers.layer_postprocess(x, y, hparams)
-        with tf.variable_scope("ffn"):
-          y = transformer_ffn_layer(
-              common_layers.layer_preprocess(
-                  x, hparams, layer_collection=layer_collection),
-              hparams,
-              conv_padding="LEFT",
-              nonpadding_mask=nonpadding,
-              losses=losses,
-              cache=layer_cache,
-              decode_loop_step=decode_loop_step,
-              layer_collection=layer_collection)
-          x = common_layers.layer_postprocess(x, y, hparams)
     # if normalization is done in layer_preprocess, then it should also be done
     # on the output, since the output can grow very large, being the sum of
     # a whole stack of unnormalized layer outputs.

From c0ee93fdbebb15376c97e4d9e19c16110e40e25d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 7 May 2019 17:45:16 -0700
Subject: [PATCH 2008/2720] Bump setup.py to 1.13.3 -- Travis is green.

PiperOrigin-RevId: 247125411
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 5b895b689..e8f49626a 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.13.2',
+    version='1.13.3',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',

From 27dfdd2c646dbcaaad80af7aed72e3053c42efad Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 8 May 2019 08:45:47 -0700
Subject: [PATCH 2009/2720] Add init Trax research models.

PiperOrigin-RevId: 247220362
---
 tensor2tensor/trax/models/research/__init__.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)
 create mode 100644 tensor2tensor/trax/models/research/__init__.py

diff --git a/tensor2tensor/trax/models/research/__init__.py b/tensor2tensor/trax/models/research/__init__.py
new file mode 100644
index 000000000..4872e5d5d
--- /dev/null
+++ b/tensor2tensor/trax/models/research/__init__.py
@@ -0,0 +1,15 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+

From a4b4326221e325d822457b45aa3a98c8e89850c0 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 8 May 2019 08:49:41 -0700
Subject: [PATCH 2010/2720] Bump version to 1.13.4

PiperOrigin-RevId: 247220994
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index e8f49626a..f305eab7c 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.13.3',
+    version='1.13.4',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',

From 5180e24458f4f937b925ddf27c587ae1dcf9423e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 8 May 2019 11:09:21 -0700
Subject: [PATCH 2011/2720] Fixed SliceNet for Image2TextProblems.

PiperOrigin-RevId: 247250287
---
 tensor2tensor/models/slicenet.py      | 17 +++++++++--------
 tensor2tensor/models/slicenet_test.py | 22 ++++++++++++++++++++++
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/models/slicenet.py b/tensor2tensor/models/slicenet.py
index 8006974b5..55254b3a1 100644
--- a/tensor2tensor/models/slicenet.py
+++ b/tensor2tensor/models/slicenet.py
@@ -30,6 +30,7 @@
 import tensorflow as tf
 
 
+# pylint: disable=unused-argument
 def attention(targets_shifted, inputs_encoded, norm_fn, hparams, bias=None):
   """Complete attention layer with preprocessing."""
   separabilities = [hparams.separability, hparams.separability]
@@ -46,8 +47,11 @@ def attention(targets_shifted, inputs_encoded, norm_fn, hparams, bias=None):
     targets_timed = tf.squeeze(targets_timed, 2)
     target_shape = tf.shape(targets_timed)
     targets_segment = tf.zeros([target_shape[0], target_shape[1]])
-    target_attention_bias = common_attention.attention_bias(
-        targets_segment, targets_segment, lower_triangular=True)
+    target_attention_bias = common_attention.attention_bias_lower_triangle(
+        target_shape[1])
+    inputs_encoded = common_layers.flatten4d3d(inputs_encoded)
+    # TODO(jbaccash): use input bias parameter. This code seems to assume fixed
+    # size inputs.
     inputs_attention_bias = tf.zeros([
         tf.shape(inputs_encoded)[0], hparams.num_heads,
         tf.shape(targets_segment)[1],
@@ -75,10 +79,8 @@ def attention(targets_shifted, inputs_encoded, norm_fn, hparams, bias=None):
         hparams.attention_dropout,
         name="encdec_attention")
     return tf.expand_dims(qv, 2)
-  elif hparams.attention_type == "simple":
-    targets_with_attention = common_layers.simple_attention(
-        targets_timed, inputs_encoded, bias=bias)
-    return norm_fn(targets_shifted + targets_with_attention, name="attn_norm")
+  else:
+    raise ValueError("Unsupported attention_type: %s" % hparams.attention_type)
 
 
 def multi_conv_res(x, padding, name, layers, hparams, mask=None, source=None):
@@ -324,7 +326,7 @@ def slicenet_params1():
   hparams.add_hparam("kernel_scheme", "3.7.15.31")
   hparams.add_hparam("audio_compression", 8)
   # attention-related flags
-  hparams.add_hparam("attention_type", "simple")
+  hparams.add_hparam("attention_type", "transformer")
   hparams.add_hparam("num_heads", 8)
   hparams.add_hparam("attention_key_channels", 0)
   hparams.add_hparam("attention_value_channels", 0)
@@ -352,7 +354,6 @@ def slicenet_params1_noam():
 def slicenet_params1_tiny():
   """Version for fast local runs."""
   hparams = slicenet_params1()
-  hparams.attention_type = "simple"
   hparams.separability = 0
   hparams.hidden_size = 128
   hparams.num_hidden_layers = 2
diff --git a/tensor2tensor/models/slicenet_test.py b/tensor2tensor/models/slicenet_test.py
index cf38c2b04..4faf71817 100644
--- a/tensor2tensor/models/slicenet_test.py
+++ b/tensor2tensor/models/slicenet_test.py
@@ -21,6 +21,7 @@
 import numpy as np
 
 from tensor2tensor.data_generators import cifar  # pylint: disable=unused-import
+from tensor2tensor.data_generators import mscoco  # pylint: disable=unused-import
 from tensor2tensor.layers import modalities  # pylint: disable=unused-import
 from tensor2tensor.models import slicenet
 from tensor2tensor.utils import registry
@@ -51,6 +52,27 @@ def testSliceNet(self):
       res = session.run(logits)
     self.assertEqual(res.shape, (3, 1, 1, 1, 10))
 
+  def testSliceNetImageToText(self):
+    x = np.random.randint(256, size=(3, 5, 5, 3))
+    y = np.random.randint(10, size=(3, 5, 1, 1))
+    hparams = slicenet.slicenet_params1_tiny()
+    hparams.add_hparam("data_dir", "")
+    problem = registry.problem("image_ms_coco_characters")
+    p_hparams = problem.get_hparams(hparams)
+    hparams.problem_hparams = p_hparams
+    with self.test_session() as session:
+      features = {
+          "inputs": tf.constant(x, dtype=tf.int32),
+          "targets": tf.constant(y, dtype=tf.int32),
+          "target_space_id": tf.constant(1, dtype=tf.int32),
+      }
+      model = slicenet.SliceNet(hparams, tf.estimator.ModeKeys.TRAIN,
+                                p_hparams)
+      logits, _ = model(features)
+      session.run(tf.global_variables_initializer())
+      res = session.run(logits)
+    self.assertEqual(res.shape, (3, 5, 1, 1, 258))
+
 
 if __name__ == "__main__":
   tf.test.main()

From 6ce1a6ce0b6c55d151e2594c7df1fe8644b6bb61 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 8 May 2019 11:13:38 -0700
Subject: [PATCH 2012/2720] Fix bayes_test for open-source.

I tested locally with a github repo. It failed before the commit and works with the commit. Hopefully that means Travis will also pass!

PiperOrigin-RevId: 247251235
---
 oss_scripts/oss_tests.sh      |  4 +---
 tensor2tensor/layers/bayes.py | 29 +++++++++++++++++++++++++++--
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 1e1920afc..3fefdf5f5 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -83,12 +83,10 @@ set_status
 # Travis Error:
 # ImportError: /usr/lib/x86_64-linux-gnu/libstdc++.so.6: version `GLIBCXX_3.4.21' not found (required by /home/travis/virtualenv/python3.6.3/lib/python3.6/site-packages/jaxlib/_pywrap_xla.so)
 
-# TODO(trandustin): Re-enable bayes_test when we can.
-#  tensor2tensor/layers/bayes_test.py \
-
 # These tests enable eager, so are tested separately.
 pytest --disable-warnings \
   tensor2tensor/data_generators/problem_test.py \
+  tensor2tensor/layers/bayes_test.py \
   tensor2tensor/layers/common_attention_test.py \
   tensor2tensor/layers/common_layers_test.py \
   tensor2tensor/layers/common_video_test.py \
diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 18c50560d..672225b10 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -213,7 +213,7 @@ def dropped_inputs():
 
     # Following tf.keras.Dropout, only apply variational dropout if training
     # flag is True. The kernel must also be a random variable.
-    training_value = tf.contrib.util.constant_value(training)
+    training_value = smart_constant_value(training)
     if training_value is not None:
       if training_value and isinstance(self.kernel, ed.RandomVariable):
         return dropped_inputs()
@@ -226,6 +226,31 @@ def dropped_inputs():
                      lambda: super(Conv2DVariationalDropout, self).call(inputs))
 
 
+# From `tensorflow/python/framework/smart_cond.py`
+def smart_constant_value(pred):
+  """Return the bool value for `pred`, or None if `pred` had a dynamic value.
+
+  Arguments:
+    pred: A scalar, either a Python bool or tensor.
+
+  Returns:
+    True or False if `pred` has a constant boolean value, None otherwise.
+
+  Raises:
+    TypeError: If `pred` is not a Tensor or bool.
+  """
+  if pred in {0, 1}:  # Accept 1/0 as valid boolean values
+    pred_value = bool(pred)
+  elif isinstance(pred, bool):
+    pred_value = pred
+  elif isinstance(pred, tf.Tensor):
+    pred_value = tf.contrib.util.constant_value(pred)
+  else:
+    raise TypeError('`pred` must be a Tensor, or a Python bool, or 1 or 0. '
+                    'Found instead: %s' % pred)
+  return pred_value
+
+
 @add_weight
 class DenseDVI(tf.keras.layers.Dense):
   """Densely-connected layer with deterministic VI (Wu et al., 2018).
@@ -518,7 +543,7 @@ def dropped_inputs():
 
     # Following tf.keras.Dropout, only apply variational dropout if training
     # flag is True. The kernel must also be a random variable.
-    training_value = tf.contrib.util.constant_value(training)
+    training_value = smart_constant_value(training)
     if training_value is not None:
       if training_value and isinstance(self.kernel, ed.RandomVariable):
         return dropped_inputs()

From de99648093456d9362699af405458a2845c2fe69 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 8 May 2019 12:07:21 -0700
Subject: [PATCH 2013/2720] re-enable multi-pass scheduled sampling. add
 sequential scheduled sampling.

PiperOrigin-RevId: 247262418
---
 tensor2tensor/layers/common_hparams.py    |   5 +-
 tensor2tensor/utils/scheduled_sampling.py | 276 ++++++++++++++++++++++
 tensor2tensor/utils/t2t_model.py          |  69 ++++--
 3 files changed, 321 insertions(+), 29 deletions(-)
 create mode 100644 tensor2tensor/utils/scheduled_sampling.py

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 26703e78e..d8bdc06b3 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -247,11 +247,10 @@ def basic_params1():
       # determined by scheduled_sampling_gold_mixin_prob. Control the number
       # of passes with scheduled_sampling_num_passes.
       scheduled_sampling_prob=0.0,
+      scheduled_sampling_method="parallel",  # parallel or sequential.
       scheduled_sampling_warmup_steps=50000,
       scheduled_sampling_gold_mixin_prob=0.5,
-      # TODO(duckworthd): Uncomment when we can ascertain why adding an
-      # extra field to HParam causes test failures.
-      # scheduled_sampling_num_passes=1,
+      scheduled_sampling_num_passes=1,
 
       # This setting controls whether to copy variables around in a daisy chain
       # (if true) or leave their placement to TensorFlow. It only affects multi
diff --git a/tensor2tensor/utils/scheduled_sampling.py b/tensor2tensor/utils/scheduled_sampling.py
new file mode 100644
index 000000000..46cbc0c38
--- /dev/null
+++ b/tensor2tensor/utils/scheduled_sampling.py
@@ -0,0 +1,276 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Scheduled Sampling.
+
+This module implemented scheduled sampling as described in (Bengio et al, 2015).
+The entry points are two functions,
+
+`sequential_scheduled_sampling_for_t2tmodel()`:
+  scheduled sampling adapted to instances of T2TModel.
+
+`sequential_scheduled_sampling()`:
+  raw implementation of scheduled sampling. May be used independent of T2T.
+
+**WARNING** This code is VERY slow. Its runtime is at least O(n^2) for
+sequences of length n. For models with self-attention, its runtime is O(n^3).
+
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import google_type_annotations
+from __future__ import print_function
+
+import copy
+
+from tensor2tensor.layers import common_layers
+import tensorflow as tf
+
+from tensorflow.python.ops import inplace_ops  # pylint: disable=g-direct-tensorflow-import
+
+
+def sequential_scheduled_sampling_for_t2tmodel(t2tmodel, features):
+  """Schedule Sampling for T2TModels.
+
+  Args:
+    t2tmodel: T2TModel instance.
+    features: {str: Tensor}. Input features.
+
+  Returns:
+    ss_logits: [batch_size, seq_len, 1, 1, vocab_size].
+    losses_dict: {str: scalar Tensor}. Losses to minimize.
+  """
+  targets = features["targets"]
+  targets_size = common_layers.shape_list(targets)
+  batch_size = targets_size[0]
+  seq_len = targets_size[1]
+  targets = tf.reshape(targets, [batch_size, seq_len])
+
+  adapter = ScheduledSamplingAdapter(t2tmodel, features)
+  ss_tokens, ss_logits, losses_dict = sequential_scheduled_sampling(
+      infer_fn=adapter.infer_fn,
+      mix_fn=adapter.mix_fn,
+      loss_fn=adapter.loss_fn,
+      targets=targets)
+
+  _ = ss_tokens  # unused.
+  targets_vocab_size = t2tmodel.problem_hparams.vocab_size["targets"]
+  ss_logits = tf.reshape(ss_logits,
+                         [batch_size, seq_len, 1, 1, targets_vocab_size])
+
+  return ss_logits, losses_dict
+
+
+def sequential_scheduled_sampling(infer_fn, mix_fn, loss_fn, targets):
+  """Scheduled Sampling.
+
+  Args:
+    infer_fn: Function. Computes logits for all timesteps.
+    mix_fn: Function. Mixes gold and sample tokens.
+    loss_fn: Function. Computes loss between gold tokens and logits.
+    targets: Tensor of shape [batch_size, seq_len]. Gold tokens.
+
+  Returns:
+    ss_tokens: Tensor of shape [batch_size, seq_len]. Scheduled sampling tokens.
+    ss_logits: Tensor of shape [batch_size, seq_len, vocab_size]. Logits for
+      next token when conditioning on ss_tokens.
+    losses_dict: {str: scalar Tensor}. Losses to optimize.
+  """
+  targets_shape = common_layers.shape_list(targets)
+  batch_size = targets_shape[0]
+  seq_len = targets_shape[1]
+
+  if not targets.shape.is_fully_defined():
+    # TODO(duckworthd): When running on GPU, I get the following error. Solve
+    # it to enable use on other devices.
+    #
+    #   Cannot use 'Identity_186' as input to
+    #   'transformer/parallel_0_7/transformer/transformer/symbol_modality_16282_512/shared/convert_gradient_to_tensor_HBc3xYw22Mw'
+    #   because 'Identity_186' is in a while loop.
+
+    raise ValueError(
+        "The following code only works on TPU. As targets.shape isn't fully "
+        "defined, I am assuming you are using a different device.")
+
+  def cond_fn(i, ss_tokens):
+    """True if i < seq_len."""
+    _ = ss_tokens
+    return i < seq_len
+
+  def body_fn(i, ss_tokens):
+    """Constructs conditioning tokens for scheduled sampling."""
+    # next_token_logits depends on timesteps 0...i-1.
+    #
+    # [batch_size, seq_len] -> [batch_size, seq_len, vocab_size]
+    ss_tokens_logits = infer_fn(ss_tokens)
+
+    # Same as 'next_token_logits = ss_tokens_logits[:, i, :]'.
+    vocab_size = common_layers.shape_list(ss_tokens_logits)[2]
+    next_token_logits = tf.slice(
+        ss_tokens_logits, begin=[0, i, 0], size=[batch_size, 1, vocab_size])
+    next_token_logits = tf.squeeze(next_token_logits, axis=[1])
+
+    # [batch_size, vocab_size] -> [batch_size]
+    sampled_next_tokens = _sample_next_tokens(next_token_logits)
+
+    # Same as 'gold_next_tokens = targets[:, i]'.
+    gold_next_tokens = tf.slice(targets, begin=[0, i], size=[batch_size, 1])
+    gold_next_tokens = tf.squeeze(gold_next_tokens, axis=[1])
+
+    next_tokens = mix_fn(gold_next_tokens, sampled_next_tokens)
+    ss_tokens = _update_timestep(ss_tokens, timestep=i, values=next_tokens)
+
+    return i+1, tf.stop_gradient(ss_tokens)
+
+  # tf.while_loop() over all timesteps. Generate scheduled sampling tokens.
+  i = 0
+  ss_tokens = tf.zeros([batch_size, seq_len], dtype=tf.int32)
+  i, ss_tokens = tf.while_loop(cond_fn, body_fn, [i, ss_tokens])
+
+  ss_logits = infer_fn(ss_tokens)
+  return ss_tokens, ss_logits, loss_fn(targets, ss_logits)
+
+
+def _mix_tokens(p_sample, gold_targets, sampled_targets):
+  """Interleave sampled and gold tokens randomly.
+
+  Args:
+    p_sample: float in [0, 1]. Probability a token will come from
+      'sampled_targets'. 0 means all-gold, 1 means all-sampled.
+    gold_targets: Tensor. Gold token IDs.
+    sampled_targets: Tensor. Sampled token IDs. Same shape as 'gold_targets'.
+
+  Returns:
+    Tensor of same shape as 'gold_targets' containing a mix of tokens from
+    'gold_targets' and 'sampled_targets'.
+  """
+  targets_shape = common_layers.shape_list(sampled_targets)
+  return tf.where(
+      tf.less(tf.random_uniform(targets_shape), p_sample),
+      sampled_targets, gold_targets)
+
+
+def _sample_next_tokens(logits):
+  """Sample tokens for next timestep."""
+  batch_size = common_layers.shape_list(logits)[0]
+  next_tokens = tf.random.categorical(logits, 1)
+  next_tokens = tf.cast(next_tokens, tf.int32)
+  next_tokens = tf.reshape(next_tokens, [batch_size])
+  return next_tokens
+
+
+def _update_timestep(x, timestep, values):
+  """Set x[:, timestep] = values.
+
+  This operation is **NOT** differentiable.
+
+  Args:
+    x: Tensor of shape [batch_size, seq_len, ...]
+    timestep: int or scalar Tensor. Index to update in x.
+    values: Tensor of shape [batch_size, ...]. New values for x[:, i].
+
+  Returns:
+    Copy of 'x' after setting x[:, timestep] = values.
+  """
+  perm = range(x.shape.ndims)
+  perm[0], perm[1] = perm[1], perm[0]
+  x = tf.transpose(x, perm)
+  x = inplace_ops.alias_inplace_update(x, timestep, values)
+  x = tf.transpose(x, perm)
+  return x
+
+
+def _inverse_exp_decay_mix_prob(p_max, num_warmup_steps):
+  """Exponentially increase to p_max over a warmup period."""
+  return (p_max *
+          common_layers.inverse_exp_decay(
+              num_warmup_steps,
+              min_value=0.001))
+
+
+class ScheduledSamplingAdapter(object):
+  """Adapts T2TModel for sequential_scheduled_sampling()."""
+
+  def __init__(self, t2tmodel, features):
+    self._t2tmodel = t2tmodel
+    self._features = features
+
+    hparams = self._t2tmodel.hparams
+    assert hparams.mode == tf.estimator.ModeKeys.TRAIN, hparams.mode
+
+  def infer_fn(self, partial_targets):
+    """Computes logits for all timesteps.
+
+    Args:
+      partial_targets: [batch_size, seq_len]. Targets to condition on.
+
+    Returns:
+      next_token_logits: [batch_size, seq_len, vocab_size]
+    """
+    batch_size, seq_len = common_layers.shape_list(partial_targets)
+    partial_targets = tf.reshape(partial_targets, [batch_size, seq_len, 1, 1])
+    features = copy.copy(self._features)
+    features["targets"] = partial_targets
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+      transformed_features = self._t2tmodel.bottom(features)
+
+      with tf.variable_scope("body"):
+        body_outputs, losses = self._t2tmodel._normalize_body_output(  # pylint: disable=protected-access
+            self._t2tmodel.body(transformed_features))
+        assert losses == {"extra": 0.0}, (
+            "Auxiliary losses are not propagated in this code. %s"
+            % (losses,))
+
+      logits = self._t2tmodel.top(body_outputs, features)
+
+    vocab_size = self._t2tmodel.problem_hparams.vocab_size["targets"]
+    logits = tf.reshape(logits, [batch_size, seq_len, vocab_size])
+    return logits
+
+  def mix_fn(self, gold_tokens, sampled_tokens):
+    """Mixes gold and sampled tokens randomly."""
+    hparams = self._t2tmodel.hparams
+    p_sample = _inverse_exp_decay_mix_prob(
+        hparams.scheduled_sampling_gold_mixin_prob,
+        hparams.scheduled_sampling_warmup_steps)
+    return _mix_tokens(
+        p_sample=p_sample,
+        gold_targets=gold_tokens,
+        sampled_targets=sampled_tokens)
+
+  def loss_fn(self, targets, logits):
+    """Constructs loss dict.
+
+    Args:
+      targets: [batch_size, seq_len]
+      logits: [batch_size, seq_len, vocab_size]
+
+    Returns:
+      {str: Tensor of shape []}. Losses.
+    """
+    batch_size, seq_len, vocab_size = common_layers.shape_list(logits)
+    targets = tf.reshape(targets, [batch_size, seq_len, 1, 1])
+    logits = tf.reshape(logits, [batch_size, seq_len, 1, 1, vocab_size])
+    features = copy.copy(self._features)
+    features["targets"] = targets
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+      losses = {
+          "training": self._t2tmodel.loss(logits, features),
+      }
+
+    return losses
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 7d5943b44..97e51ba30 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -43,6 +43,7 @@
 from tensor2tensor.utils import optimize
 from tensor2tensor.utils import quantization
 from tensor2tensor.utils import registry
+from tensor2tensor.utils import scheduled_sampling
 
 import tensorflow as tf
 
@@ -1785,10 +1786,11 @@ def maybe_scheduled_sampling(self, features, logits, losses):
     This is the identity unless self.hparams.scheduled_sampling_prob > 0
     (default).
 
-    **WARNING**: This is not a faithful implementation of scheduled sampling.
-    This implementation samples tokens for timestep t condtioned on gold tokens
-    1...t-1. A proper implementation must condition on a mix of gold and
-    sampled tokens. Doing so is not efficient for models such like Transformer.
+    **WARNING**: If hparams.scheduled_sampling_method == "parallel", this is
+    not a faithful implementation of scheduled sampling. This implementation
+    samples tokens for timestep t condtioned on gold tokens 1...t-1. A proper
+    implementation must condition on a mix of gold and sampled tokens. Doing
+    so is not efficient for models such like Transformer.
 
     Args:
       features: {str: Tensor}. Features sharded along batch dimension.
@@ -1828,12 +1830,14 @@ def maybe_scheduled_sampling(self, features, logits, losses):
     assert vocab_size is not None
     assert hparams.vocab_divisor == 1
 
+    # TODO(duckworthd): Move to scheduled_sampling.py.
     def sample(x):
       """Multinomial sampling from a n-dimensional tensor."""
       samples = tf.multinomial(tf.reshape(x, [-1, vocab_size]), 1)
       reshaped_samples = tf.reshape(samples, common_layers.shape_list(x)[:-1])
       return tf.to_int32(reshaped_samples)
 
+    # TODO(duckworthd): Move to scheduled_sampling.py.
     def mix_gold_sampled(gold_targets, sampled_targets, mixin_prob):
       """Interleave sampled and gold tokens randomly."""
       return tf.where(
@@ -1843,7 +1847,8 @@ def mix_gold_sampled(gold_targets, sampled_targets, mixin_prob):
           sampled_targets,
           gold_targets)
 
-    def sampled_results(features, logits, mixin_prob):
+    # TODO(duckworthd): Move to scheduled_sampling.py.
+    def parallel_scheduled_sampling_pass(features, logits, mixin_prob):
       """Generate scheduled sampling results."""
       sampled_targets = sample(logits)
       new_targets = mix_gold_sampled(features["targets"],
@@ -1879,27 +1884,39 @@ def sampled_results(features, logits, mixin_prob):
     tf.logging.info("Using scheduled sampling.")
     assert hparams.scheduled_sampling_prob == 1.0, (
         "hparams.scheduled_sampling_prob must be 0 or 1.")
-    # Gradually increase over a warmup period. Lower numbers mean more gold
-    # tokens.
-    mixin_prob = (
-        hparams.scheduled_sampling_gold_mixin_prob *
-        common_layers.inverse_exp_decay(
-            hparams.scheduled_sampling_warmup_steps,
-            min_value=0.001)
-    )
-
-    # Apply scheduled sampling over N passes. The logits from the (n-1)-th pass
-    # will be mixed with gold tokens for conditioning in the n-th pass.
-    scheduled_sampling_num_passes = getattr(
-        hparams, "scheduled_sampling_num_passes", 1)
-    assert scheduled_sampling_num_passes > 0, (
-        "hparams.scheduled_sampling_num_passes must be > 0 if "
-        "hparams.scheduled_sampling_prob > 0.0")
-    new_logits = logits
-    new_losses = losses
-    for _ in range(scheduled_sampling_num_passes):
-      new_logits, new_losses = sampled_results(features, new_logits, mixin_prob)
-    return new_logits, new_losses
+
+    if hparams.scheduled_sampling_method == "sequential":
+      tf.logging.info("Using SEQUENTIAL scheduled sampling.")
+      assert hparams.scheduled_sampling_num_passes == 1, (
+          "hparams.scheduled_sampling_num_passes must equal 1 if "
+          "doing sequential scheduled sampling.")
+      return scheduled_sampling.sequential_scheduled_sampling_for_t2tmodel(
+          self, features)
+    elif hparams.scheduled_sampling_method == "parallel":
+      tf.logging.info("Using PARALLEL scheduled sampling.")
+      # TODO(duckworthd): Move this block to scheduled_sampling.py.
+
+      # Gradually increase over a warmup period. Lower numbers mean more gold
+      # tokens.
+      mixin_prob = scheduled_sampling._inverse_exp_decay_mix_prob(  # pylint: disable=protected-access
+          hparams.scheduled_sampling_gold_mixin_prob,
+          hparams.scheduled_sampling_warmup_steps)
+
+      # Apply scheduled sampling over N passes. The logits from the (n-1)-th
+      # pass will be mixed with gold tokens for conditioning in the n-th pass.
+      assert hparams.scheduled_sampling_num_passes > 0, (
+          "hparams.scheduled_sampling_num_passes must be > 0 if "
+          "hparams.scheduled_sampling_prob > 0.0")
+      new_logits = logits
+      new_losses = losses
+      for _ in range(hparams.scheduled_sampling_num_passes):
+        new_logits, new_losses = parallel_scheduled_sampling_pass(
+            features, new_logits, mixin_prob)
+      return new_logits, new_losses
+    else:
+      raise ValueError(
+          "Unknown scheduled_sampling_method = %s" % (
+              hparams.scheduled_sampling_method,))
 
 
 def _with_timing(fn, msg, silent=False):

From 099cec7c23169145fa026472f7037a85fc4afd7f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 8 May 2019 14:45:09 -0700
Subject: [PATCH 2014/2720] Various warmup schedules for scheduled sampling.

PiperOrigin-RevId: 247295474
---
 tensor2tensor/layers/common_hparams.py    |  1 +
 tensor2tensor/layers/common_layers.py     | 42 +++++++++++++++++++++--
 tensor2tensor/utils/scheduled_sampling.py | 17 +++++----
 tensor2tensor/utils/t2t_model.py          |  5 ++-
 4 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index d8bdc06b3..2af2e56f3 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -251,6 +251,7 @@ def basic_params1():
       scheduled_sampling_warmup_steps=50000,
       scheduled_sampling_gold_mixin_prob=0.5,
       scheduled_sampling_num_passes=1,
+      scheduled_sampling_warmup_schedule="exp",  # exp, linear, or sigmoid.
 
       # This setting controls whether to copy variables around in a daisy chain
       # (if true) or leave their placement to TensorFlow. It only affects multi
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 4a6b16865..a45b87a11 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -153,7 +153,7 @@ def hard_tanh(x, saturation_limit=0.9):
 
 
 def inverse_exp_decay(max_step, min_value=0.01, step=None):
-  """Inverse-decay exponentially from 0.01 to 1.0 reached at max_step."""
+  """Inverse-decay exponentially from min_value to 1.0 reached at max_step."""
   inv_base = tf.exp(tf.log(min_value) / float(max_step))
   if step is None:
     step = tf.train.get_global_step()
@@ -164,7 +164,7 @@ def inverse_exp_decay(max_step, min_value=0.01, step=None):
 
 
 def inverse_lin_decay(max_step, min_value=0.01, step=None):
-  """Inverse-decay linearly from 0.01 to 1.0 reached at max_step."""
+  """Inverse-decay linearly from min_value to 1.0 reached at max_step."""
   if step is None:
     step = tf.train.get_global_step()
   if step is None:
@@ -174,6 +174,44 @@ def inverse_lin_decay(max_step, min_value=0.01, step=None):
   return progress * (1.0 - min_value) + min_value
 
 
+def inverse_sigmoid_decay(max_step, min_value=0.01, step=None):
+  """Inverse-decay linearly from min_value to 1.0 reached at max_step."""
+  if step is None:
+    step = tf.train.get_global_step()
+  if step is None:
+    return 1.0
+  step = to_float(step)
+
+  def sigmoid(x):
+    return 1 / (1 + tf.exp(-x))
+
+  def inv_sigmoid(y):
+    return tf.log(y / (1 - y))
+
+  assert min_value > 0, (
+      "sigmoid's output is always >0 and <1. min_value must respect "
+      "these bounds for interpolation to work.")
+  assert min_value < 0.5, "Must choose min_value on the left half of sigmoid."
+
+  # Find
+  #   x  s.t. sigmoid(x ) = y_min and
+  #   x' s.t. sigmoid(x') = y_max
+  # We will map [0, max_step] to [x_min, x_max].
+  y_min = min_value
+  y_max = 1.0 - min_value
+  x_min = inv_sigmoid(y_min)
+  x_max = inv_sigmoid(y_max)
+
+  x = tf.minimum(step / float(max_step), 1.0)  # [0, 1]
+  x = x_min + (x_max - x_min) * x  # [x_min, x_max]
+  y = sigmoid(x)  # [y_min, y_max]
+
+  y = (y - y_min) / (y_max - y_min)  # [0, 1]
+  y = y * (1.0 - y_min)  # [0, 1-y_min]
+  y += y_min  # [y_min, 1]
+  return y
+
+
 def shakeshake2_py(x, y, equal=False, individual=False):
   """The shake-shake sum of 2 tensors, python version."""
   if equal:
diff --git a/tensor2tensor/utils/scheduled_sampling.py b/tensor2tensor/utils/scheduled_sampling.py
index 46cbc0c38..21815a9d4 100644
--- a/tensor2tensor/utils/scheduled_sampling.py
+++ b/tensor2tensor/utils/scheduled_sampling.py
@@ -193,12 +193,14 @@ def _update_timestep(x, timestep, values):
   return x
 
 
-def _inverse_exp_decay_mix_prob(p_max, num_warmup_steps):
-  """Exponentially increase to p_max over a warmup period."""
-  return (p_max *
-          common_layers.inverse_exp_decay(
-              num_warmup_steps,
-              min_value=0.001))
+def inverse_decay_mix_prob(warmup_schedule_name, p_max, num_warmup_steps):
+  """Interpolate from 0.001 to 'p_max' over 'num_warmup_steps'."""
+  warmup_schedule_fn = {
+      "exp": common_layers.inverse_exp_decay,
+      "linear": common_layers.inverse_lin_decay,
+      "sigmoid": common_layers.inverse_sigmoid_decay,
+  }[warmup_schedule_name]
+  return p_max * warmup_schedule_fn(num_warmup_steps, min_value=0.001)
 
 
 class ScheduledSamplingAdapter(object):
@@ -244,7 +246,8 @@ def infer_fn(self, partial_targets):
   def mix_fn(self, gold_tokens, sampled_tokens):
     """Mixes gold and sampled tokens randomly."""
     hparams = self._t2tmodel.hparams
-    p_sample = _inverse_exp_decay_mix_prob(
+    p_sample = inverse_decay_mix_prob(
+        hparams.scheduled_sampling_warmup_schedule,
         hparams.scheduled_sampling_gold_mixin_prob,
         hparams.scheduled_sampling_warmup_steps)
     return _mix_tokens(
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 97e51ba30..5afdfa781 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1882,6 +1882,8 @@ def parallel_scheduled_sampling_pass(features, logits, mixin_prob):
       return new_logits, new_losses
 
     tf.logging.info("Using scheduled sampling.")
+    tf.logging.info("Warming scheduled sampling up with schedule: %s",
+                    hparams.scheduled_sampling_warmup_schedule)
     assert hparams.scheduled_sampling_prob == 1.0, (
         "hparams.scheduled_sampling_prob must be 0 or 1.")
 
@@ -1898,7 +1900,8 @@ def parallel_scheduled_sampling_pass(features, logits, mixin_prob):
 
       # Gradually increase over a warmup period. Lower numbers mean more gold
       # tokens.
-      mixin_prob = scheduled_sampling._inverse_exp_decay_mix_prob(  # pylint: disable=protected-access
+      mixin_prob = scheduled_sampling.inverse_decay_mix_prob(
+          hparams.scheduled_sampling_warmup_schedule,
           hparams.scheduled_sampling_gold_mixin_prob,
           hparams.scheduled_sampling_warmup_steps)
 

From b4c2d57263a5da9eafbf6a79555fc219a22e2ccb Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 8 May 2019 21:19:38 -0700
Subject: [PATCH 2015/2720] PPO changes for improvements in the combined
 network case. - Incorporate old value predictions while calculating value
 loss. This   technique is from the OpenAI Baselines implementation. Though I
 haven't done   a thorough test of whether this helps or not. We do this only
 in the case of   the combined network. See `value_loss_given_predictions`. -
 Combined Loss should compute loss wrt new value predictions. So this was an  
 error in the earlier code. See `combined_loss_given_predictions`. - This is
 the thing that helped the most, but probably indicates bad JAX random  
 initialization:

  While initializing the combined network, we pass the same random seed that
  was used to initialize the policy network in the two networks case.

  Maybe I should vary the random seed on re-runs, currently my random seed is
  fixed at 0.
- Minor cleanup: There were several redundant get_param calls in the
  training_loop.

PiperOrigin-RevId: 247352436
---
 tensor2tensor/trax/rlax/ppo.py      | 92 ++++++++++++++++-------------
 tensor2tensor/trax/rlax/ppo_main.py | 18 +++---
 tensor2tensor/trax/rlax/ppo_test.py | 11 ++--
 3 files changed, 67 insertions(+), 54 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index a3be4d62d..bcae0a839 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -443,7 +443,9 @@ def value_loss(value_net_apply,
                observations,
                rewards,
                reward_mask,
-               gamma=0.99):
+               gamma=0.99,
+               epsilon=0.2,
+               value_prediction_old=None):
   """Computes the value loss.
 
   Args:
@@ -454,6 +456,10 @@ def value_loss(value_net_apply,
     rewards: np.ndarray of shape (B, T) of rewards.
     reward_mask: np.ndarray of shape (B, T), the mask over rewards.
     gamma: float, discount factor.
+    epsilon: float, clip-fraction, used if value_value_prediction_old isn't None
+    value_prediction_old: np.ndarray of shape (B, T+1, 1) of value predictions
+        using the old parameters. If provided, we incorporate this in the loss
+        as well. This is from the OpenAI baselines implementation.
 
   Returns:
     The average L2 value loss, averaged over instances where reward_mask is 1.
@@ -467,14 +473,17 @@ def value_loss(value_net_apply,
   assert (B, T + 1, 1) == value_prediction.shape
 
   return value_loss_given_predictions(value_prediction, rewards, reward_mask,
-                                      gamma)
+                                      gamma, epsilon=epsilon,
+                                      value_prediction_old=value_prediction_old)
 
 
 @jit
 def value_loss_given_predictions(value_prediction,
                                  rewards,
                                  reward_mask,
-                                 gamma=0.99):
+                                 gamma=0.99,
+                                 epsilon=0.2,
+                                 value_prediction_old=None):
   """Computes the value loss given the prediction of the value function.
 
   Args:
@@ -482,6 +491,10 @@ def value_loss_given_predictions(value_prediction,
     rewards: np.ndarray of shape (B, T) of rewards.
     reward_mask: np.ndarray of shape (B, T), the mask over rewards.
     gamma: float, discount factor.
+    epsilon: float, clip-fraction, used if value_value_prediction_old isn't None
+    value_prediction_old: np.ndarray of shape (B, T+1, 1) of value predictions
+        using the old parameters. If provided, we incorporate this in the loss
+        as well. This is from the OpenAI baselines implementation.
 
   Returns:
     The average L2 value loss, averaged over instances where reward_mask is 1.
@@ -496,6 +509,16 @@ def value_loss_given_predictions(value_prediction,
   r2g = rewards_to_go(rewards, reward_mask, gamma=gamma)  # (B, T)
   loss = (value_prediction - r2g)**2
 
+  # From the baselines implementation.
+  if value_prediction_old is not None:
+    value_prediction_old = np.squeeze(value_prediction_old, axis=2)  # (B, T+1)
+    value_prediction_old = value_prediction_old[:, :-1] * reward_mask  # (B, T)
+
+    v_clipped = value_prediction_old + np.clip(
+        value_prediction - value_prediction_old, -epsilon, epsilon)
+    v_clipped_loss = (v_clipped - r2g)**2
+    loss = np.maximum(v_clipped_loss, loss)
+
   # Take an average on only the points where mask != 0.
   return np.sum(loss) / np.sum(reward_mask)
 
@@ -709,7 +732,8 @@ def ppo_loss_given_predictions(log_probab_actions_new,
 @jit
 def combined_loss_given_predictions(log_probab_actions_new,
                                     log_probab_actions_old,
-                                    value_prediction,
+                                    value_prediction_new,
+                                    value_prediction_old,
                                     padded_actions,
                                     padded_rewards,
                                     reward_mask,
@@ -720,11 +744,12 @@ def combined_loss_given_predictions(log_probab_actions_new,
                                     c2=0.01):
   """Computes the combined (clipped loss + value loss) given predictions."""
   loss_value = value_loss_given_predictions(
-      value_prediction, padded_rewards, reward_mask, gamma=gamma)
+      value_prediction_new, padded_rewards, reward_mask, gamma=gamma,
+      value_prediction_old=value_prediction_old, epsilon=epsilon)
   loss_ppo = ppo_loss_given_predictions(
       log_probab_actions_new,
       log_probab_actions_old,
-      value_prediction,
+      value_prediction_old,
       padded_actions,
       padded_rewards,
       reward_mask,
@@ -751,21 +776,22 @@ def combined_loss(new_params,
                   c1=1.0,
                   c2=0.01):
   """Computes the combined (clipped loss + value loss) given observations."""
-  log_probab_actions_new, _ = policy_and_value_net_apply(
+  log_probab_actions_new, value_predictions_new = policy_and_value_net_apply(
       padded_observations, new_params)
 
   # (combined_loss, ppo_loss, value_loss, entropy_bonus)
   return combined_loss_given_predictions(log_probab_actions_new,
                                          log_probab_actions_old,
+                                         value_predictions_new,
                                          value_predictions_old,
                                          padded_actions,
                                          padded_rewards,
                                          reward_mask,
-                                         c1=c1,
-                                         c2=c2,
                                          gamma=gamma,
                                          lambda_=lambda_,
-                                         epsilon=epsilon)
+                                         epsilon=epsilon,
+                                         c1=c1,
+                                         c2=c2)
 
 
 @functools.partial(jit, static_argnums=(2, 3, 4))
@@ -863,6 +889,7 @@ def policy_and_value_loss(params):
 
   new_params = get_params(opt_state)
   g = grad(policy_and_value_loss)(new_params)
+  # TODO(afrozm): Maybe clip gradients?
   return opt_update(i, g, opt_state)
 
 
@@ -954,19 +981,19 @@ def training_loop(
   policy_net_params, policy_net_apply = None, None
   value_net_params, value_net_apply = None, None
   if policy_and_value_net_fun is not None:
-    jax_rng_key, subkey = jax_random.split(jax_rng_key)
+    jax_rng_key, key1, key2 = jax_random.split(jax_rng_key, num=3)
 
     # Initialize the policy and value network.
     policy_and_value_net_params, policy_and_value_net_apply = (
-        policy_and_value_net_fun(subkey, batch_observations_shape, num_actions))
+        policy_and_value_net_fun(key1, batch_observations_shape, num_actions))
+
+    policy_and_value_net_apply = jit(policy_and_value_net_apply)
 
     # Initialize the optimizers.
     policy_and_value_optimizer = (
         policy_and_value_optimizer_fun(policy_and_value_net_params))
     (policy_and_value_opt_state, policy_and_value_opt_update,
      policy_and_value_get_params) = policy_and_value_optimizer
-
-    policy_and_value_net_apply = jit(policy_and_value_net_apply)
   else:
     # Initialize the policy and value functions.
     assert policy_net_fun and value_net_fun
@@ -1000,15 +1027,13 @@ def training_loop(
     # A function to get the policy and value predictions.
     def get_predictions(observations):
       if policy_net_apply is not None:
-        # Get the fresh params for collecting the policy.
         return (policy_net_apply(observations, policy_net_params),
                 value_net_apply(observations, value_net_params))
 
       assert policy_and_value_net_apply
 
-      # Get the fresh params for collecting the policy.
-      return policy_and_value_net_apply(observations,
-                                        policy_and_value_net_params)
+      return policy_and_value_net_apply(
+          observations, policy_and_value_net_params)
 
     t = time.time()
     t0 = t
@@ -1024,14 +1049,6 @@ def get_predictions(observations):
 
     logging.vlog(1, "Collecting trajectories took %0.2f msec.", get_time(t))
 
-    # These were the params that were used to collect the trajectory.
-    if policy_and_value_net_apply:
-      policy_and_value_net_params = policy_and_value_get_params(
-          policy_and_value_opt_state)
-    else:
-      policy_net_params = ppo_get_params(ppo_opt_state)
-      value_net_params = value_get_params(value_opt_state)
-
     avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
     max_reward = max(np.sum(traj[2]) for traj in trajs)
     min_reward = min(np.sum(traj[2]) for traj in trajs)
@@ -1104,6 +1121,7 @@ def get_predictions(observations):
           1, "Calculating P&V loss [%10.2f(%10.2f, %10.2f)] took %0.2f msec.",
           cur_combined_loss, cur_value_loss, cur_ppo_loss, get_time(t))
     else:
+      logging.vlog(2, "Starting to compute Value loss.")
       t = time.time()
       cur_value_loss = value_loss(
           value_net_apply,
@@ -1116,6 +1134,7 @@ def get_predictions(observations):
       logging.vlog(1, "Calculating value loss took %0.2f msec.", get_time(t))
 
       t = time.time()
+      logging.vlog(2, "Starting to compute PPO loss.")
       cur_ppo_loss = ppo_loss(
           policy_net_apply,
           policy_net_params,
@@ -1171,15 +1190,18 @@ def get_predictions(observations):
                                    reward_mask)
 
         early_stopping = approx_kl > 1.5 * target_kl
+        if early_stopping:
+          logging.vlog(
+              1, "Early stopping policy and value optimization at iter: %d, "
+              "with approx_kl: %0.2f", j, approx_kl)
+          # We don't return right-away, we want the below to execute on the last
+          # iteration.
 
         t2 = time.time()
         if (((j + 1) %
              print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1)
             or early_stopping):
           # Compute and log the loss.
-          # Get the new params.
-          new_policy_and_value_net_params = policy_and_value_get_params(
-              policy_and_value_opt_state)
           (loss_combined, loss_ppo, loss_value, entropy_bonus) = (
               combined_loss(
                   new_policy_and_value_net_params,
@@ -1204,9 +1226,6 @@ def get_predictions(observations):
         if early_stopping:
           break
 
-      # Update the params.
-      policy_and_value_net_params = new_policy_and_value_net_params
-
       logging.vlog(
           1, "Total Combined Loss reduction [%0.2f]%%",
           (100 *
@@ -1281,11 +1300,6 @@ def get_predictions(observations):
         if early_stopping:
           break
 
-      # Update the params ONLY AND ONLY AFTER we complete all the optimization
-      # iterations, till then `policy_net_params` should refer to the params
-      # that were used in collecting the policy.
-      # policy_net_params = ppo_get_params(ppo_opt_state)
-
       logging.vlog(1, "Total PPO loss reduction [%0.2f]%%",
                    (100 * (cur_ppo_loss - new_ppo_loss) / np.abs(cur_ppo_loss)))
 
@@ -1324,10 +1338,6 @@ def get_predictions(observations):
 
       logging.vlog(1, "Grad desc took %0.2f msec", get_time(t1))
 
-      # Set the optimized params to new params.
-      policy_net_params = ppo_get_params(ppo_opt_state)
-      value_net_params = value_get_params(value_opt_state)
-
       logging.info(
           "Epoch [% 6d], Reward[min, max, avg] [%10.2f,%10.2f,%10.2f], "
           "ppo loss [%10.2f], value loss [%10.2f], took [%10.2f msec]", i,
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index a0563b129..f36e89e1a 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -106,18 +106,18 @@
 
 # Target KL is used for doing early stopping in the
 flags.DEFINE_float("target_kl", 0.01, "Policy iteration early stopping")
+flags.DEFINE_float("value_coef", 1.0,
+                   "Coefficient of Value Loss term in combined loss.")
+flags.DEFINE_float("entropy_coef", 0.01,
+                   "Coefficient of the Entropy Bonus term in combined loss.")
 
 
 def common_layers():
   cur_layers = []
   if FLAGS.flatten_non_batch_time_dims:
     cur_layers = [layers.Div(divisor=255.0), layers.Flatten(num_axis_to_keep=2)]
-  return cur_layers + [
-      layers.Dense(64),
-      layers.Tanh(),
-      layers.Dense(64),
-      layers.Tanh()
-  ]
+  body = [layers.Dense(64), layers.Tanh(), layers.Dense(64), layers.Tanh()]
+  return cur_layers + body
 
 
 def make_env():
@@ -194,14 +194,16 @@ def run_training_loop():
         policy_optimizer_fun=policy_optimizer_fun,
         value_optimizer_fun=value_optimizer_fun,
         policy_and_value_optimizer_fun=policy_and_value_optimizer_fun,
-        batch_size=FLAGS.batch_size,
         num_optimizer_steps=FLAGS.num_optimizer_steps,
         policy_only_num_optimizer_steps=FLAGS.policy_only_num_optimizer_steps,
         value_only_num_optimizer_steps=FLAGS.value_only_num_optimizer_steps,
+        batch_size=FLAGS.batch_size,
         target_kl=FLAGS.target_kl,
         boundary=FLAGS.boundary,
         max_timestep=FLAGS.max_timestep,
-        random_seed=FLAGS.random_seed)
+        random_seed=FLAGS.random_seed,
+        c1=FLAGS.value_coef,
+        c2=FLAGS.entropy_coef)
 
   if FLAGS.jax_debug_nans or FLAGS.disable_jit:
     with jax.disable_jit():
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 202766a88..2be00ba5d 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -565,8 +565,8 @@ def test_combined_loss(self):
     mask = np.ones_like(rewards)
 
     # Just test that this computes at all.
-    new_log_probabs, _ = net_apply(observations, new_params)
-    old_log_probabs, value_predictions = net_apply(observations, old_params)
+    new_log_probabs, value_predictions_new = net_apply(observations, new_params)
+    old_log_probabs, value_predictions_old = net_apply(observations, old_params)
 
     gamma = 0.99
     lambda_ = 0.95
@@ -575,11 +575,12 @@ def test_combined_loss(self):
     c2 = 0.01
 
     value_loss_1 = ppo.value_loss_given_predictions(
-        value_predictions, rewards, mask, gamma=gamma)
+        value_predictions_new, rewards, mask, gamma=gamma,
+        value_prediction_old=value_predictions_old, epsilon=epsilon)
     ppo_loss_1 = ppo.ppo_loss_given_predictions(
         new_log_probabs,
         old_log_probabs,
-        value_predictions,
+        value_predictions_old,
         actions,
         rewards,
         mask,
@@ -590,7 +591,7 @@ def test_combined_loss(self):
     (combined_loss, ppo_loss_2, value_loss_2, entropy_bonus) = (
         ppo.combined_loss(new_params,
                           old_log_probabs,
-                          value_predictions,
+                          value_predictions_old,
                           net_apply,
                           observations,
                           actions,

From 12b771dd3253991b81c5f0ad1e1868bc90ba6949 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 9 May 2019 12:37:34 -0700
Subject: [PATCH 2016/2720] Made ShiftRight support other dtypes than int.

PiperOrigin-RevId: 247476656
---
 tensor2tensor/trax/layers/attention.py      |  3 ++-
 tensor2tensor/trax/layers/attention_test.py | 21 +++++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 7c5de1982..8763660f7 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -258,7 +258,8 @@ def ShiftRight(x, **unused_kwargs):
   if not isinstance(x, (list, tuple)):  # non-chunked inputs
     pad_widths = [(0, 0)] * len(x.shape)
     pad_widths[1] = (1, 0)  # Padding on axis=1
-    padded = np.pad(x, pad_widths, mode='constant')
+    padded = np.pad(x, pad_widths, mode='constant',
+                    constant_values=x.dtype.type(0))
     return padded[:, :-1]
   # Handling chunked inputs. Recall that the list of chunks represents a big
   # sequence (the concatenation of the chunks). We want to shift that sequence,
diff --git a/tensor2tensor/trax/layers/attention_test.py b/tensor2tensor/trax/layers/attention_test.py
index cfa4a51bf..3517b5986 100644
--- a/tensor2tensor/trax/layers/attention_test.py
+++ b/tensor2tensor/trax/layers/attention_test.py
@@ -41,6 +41,27 @@ def test_shift_right(self):
                                     [12, 13, 14]]]),
                         output_np)
 
+  def test_shift_right_float(self):
+    layer = attention.ShiftRight()
+    input_np = onp.arange(2*3*3).reshape(2, 3, 3).astype(onp.float32)
+    # Test on a float array.
+    input_np = input_np.astype(onp.float32)
+    input_np /= 2.0
+    self.assertEqual(input_np.dtype, onp.float32)
+
+    output_np = layer(input_np)
+    self.assertEqual(input_np.shape, output_np.shape)
+    self.assertEqual(output_np.dtype, onp.float32)
+
+    self.assertAllEqual(onp.array([[[0., 0., 0.],
+                                    [0., 0.5, 1.],
+                                    [1.5, 2., 2.5]],
+
+                                   [[0., 0., 0.],
+                                    [4.5, 5., 5.5],
+                                    [6., 6.5, 7.]]]),
+                        output_np)
+
 
 if __name__ == '__main__':
   test.main()

From 1b0a83184d4f61f810f4a081c5f028ab123b21ad Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 9 May 2019 12:37:59 -0700
Subject: [PATCH 2017/2720] Internal.

PiperOrigin-RevId: 247476735
---
 tensor2tensor/utils/scheduled_sampling.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensor2tensor/utils/scheduled_sampling.py b/tensor2tensor/utils/scheduled_sampling.py
index 21815a9d4..80def28bd 100644
--- a/tensor2tensor/utils/scheduled_sampling.py
+++ b/tensor2tensor/utils/scheduled_sampling.py
@@ -31,7 +31,6 @@
 
 from __future__ import absolute_import
 from __future__ import division
-from __future__ import google_type_annotations
 from __future__ import print_function
 
 import copy

From 5255d747f63e6990e53413fabb0115e83a2f62d4 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 9 May 2019 16:44:39 -0700
Subject: [PATCH 2018/2720] Use SM3 for big Transformer LM in Trax.

PiperOrigin-RevId: 247520619
---
 .../trax/configs/transformer_big_lm1b_8gb.gin       | 13 +++++++------
 tensor2tensor/trax/optimizers.py                    |  1 +
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
index 99ee6ea5d..760cf6b6b 100644
--- a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
@@ -5,9 +5,9 @@ import tensor2tensor.trax.trax
 
 # Parameters for batch_fun:
 # ==============================================================================
-batch_fun.batch_size_per_device = 16
-batch_fun.eval_batch_size = 128
-batch_fun.max_eval_length = 2048
+batch_fun.batch_size_per_device = 32
+batch_fun.eval_batch_size = 64
+batch_fun.max_eval_length = 512
 
 # Parameters for inputs:
 # ==============================================================================
@@ -22,14 +22,14 @@ masked_mean.mask_id = 0
 # Parameters for MultifactorSchedule:
 # ==============================================================================
 MultifactorSchedule.constant = 0.1
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
+MultifactorSchedule.factors = 'constant * linear_warmup'
+MultifactorSchedule.warmup_steps = 16000
 
 # Parameters for preprocess_fun:
 # ==============================================================================
 shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
 lm1b_preprocess.max_target_length = 512
-lm1b_preprocess.max_eval_target_length = 2048
+lm1b_preprocess.max_eval_target_length = 512
 
 # Parameters for train:
 # ==============================================================================
@@ -37,6 +37,7 @@ train.eval_frequency = 1000
 train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
+train.optimizer = @trax.optimizers.sm3
 train.run_debug_step = False
 train.train_steps = 500000
 
diff --git a/tensor2tensor/trax/optimizers.py b/tensor2tensor/trax/optimizers.py
index 5347de230..4d8e1fef6 100644
--- a/tensor2tensor/trax/optimizers.py
+++ b/tensor2tensor/trax/optimizers.py
@@ -32,6 +32,7 @@ def opt_configure(*args, **kwargs):
 
 # Optimizers
 sgd = opt_configure(opt.sgd)
+sm3 = opt_configure(opt.sm3)
 adam = opt_configure(opt.adam)
 momentum = opt_configure(opt.momentum)
 rmsprop = opt_configure(opt.rmsprop)

From 92250f7e3c9a8ede72e0e21a82a858b67e7accd2 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 9 May 2019 17:28:33 -0700
Subject: [PATCH 2019/2720] ppo.policy_and_value_net now makes a two tower
 architecture now.

PiperOrigin-RevId: 247527418
---
 tensor2tensor/trax/rlax/ppo.py                | 22 ++++++++++---------
 tensor2tensor/trax/rlax/ppo_main.py           |  2 +-
 tensor2tensor/trax/rlax/ppo_test.py           | 12 +++++-----
 .../trax/rlax/ppo_training_loop_test.py       |  3 ++-
 4 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index bcae0a839..2852ef474 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -115,23 +115,25 @@ def value_net(rng_key,
 def policy_and_value_net(rng_key,
                          batch_observations_shape,
                          num_actions,
-                         bottom_layers=None):
+                         bottom_layers_fn=None):
   """A policy and value net function."""
 
   # Layers.
-  cur_layers = []
-  if bottom_layers is not None:
-    cur_layers.extend(bottom_layers)
 
   # Now, with the current logits, one head computes action probabilities and the
   # other computes the value function.
   # NOTE: The LogSoftmax instead of the Softmax because of numerical stability.
-  cur_layers.extend([
-      layers.Branch(
-          layers.Serial(layers.Dense(num_actions), layers.LogSoftmax()),
-          layers.Dense(1))
-  ])
-  net = layers.Serial(*cur_layers)
+
+  tower1 = [] if bottom_layers_fn is None else bottom_layers_fn()
+  tower2 = [] if bottom_layers_fn is None else bottom_layers_fn()
+
+  tower1.extend([layers.Dense(num_actions), layers.LogSoftmax()])
+  tower2.extend([layers.Dense(1)])
+
+  net = layers.Branch(
+      layers.Serial(*tower1),
+      layers.Serial(*tower2),
+  )
   return net.initialize(batch_observations_shape, rng_key), net
 
 
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index f36e89e1a..bff2ea74c 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -175,7 +175,7 @@ def run_training_loop():
 
     if FLAGS.combined_policy_and_value_function:
       policy_and_value_net_fun = functools.partial(
-          ppo.policy_and_value_net, bottom_layers=common_layers())
+          ppo.policy_and_value_net, bottom_layers_fn=common_layers)
       policy_and_value_optimizer_fun = get_optimizer_fun(FLAGS.learning_rate)
     else:
       policy_net_fun = functools.partial(
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 2be00ba5d..fe7a2a905 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -83,7 +83,7 @@ def test_policy_and_value_net(self):
     num_actions = 2
     pnv_params, pnv_apply = ppo.policy_and_value_net(
         self.rng_key, batch_observation_shape, num_actions,
-        [layers.Flatten(num_axis_to_keep=2)])
+        lambda: [layers.Flatten(num_axis_to_keep=2)])
     batch = 2
     time_steps = 10
     batch_of_observations = np.random.uniform(
@@ -132,7 +132,7 @@ def test_collect_trajectories(self):
     # Test collect using a Policy and Value function.
     pnv_params, pnv_apply = ppo.policy_and_value_net(
         self.rng_key, (-1, -1) + observation_shape, num_actions,
-        [layers.Flatten(num_axis_to_keep=2)])
+        lambda: [layers.Flatten(num_axis_to_keep=2)])
 
     trajectories = ppo.collect_trajectories(
         env,
@@ -156,7 +156,7 @@ def test_collect_trajectories_max_timestep(self):
     num_actions = 2
     pnv_params, pnv_apply = ppo.policy_and_value_net(
         self.rng_key, (-1, -1) + observation_shape, num_actions,
-        [layers.Flatten(num_axis_to_keep=2)])
+        lambda: [layers.Flatten(num_axis_to_keep=2)])
 
     # We'll get done at time-step #5, starting from 0, therefore in 6 steps.
     done_time_step = 5
@@ -552,10 +552,12 @@ def test_combined_loss(self):
     batch_observation_shape = (-1, -1) + OBS
 
     old_params, _ = ppo.policy_and_value_net(
-        key1, batch_observation_shape, A, [layers.Flatten(num_axis_to_keep=2)])
+        key1, batch_observation_shape, A,
+        lambda: [layers.Flatten(num_axis_to_keep=2)])
 
     new_params, net_apply = ppo.policy_and_value_net(
-        key2, batch_observation_shape, A, [layers.Flatten(num_axis_to_keep=2)])
+        key2, batch_observation_shape, A,
+        lambda: [layers.Flatten(num_axis_to_keep=2)])
 
     # Generate a batch of observations.
 
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index 8b80ada8f..c4f4e7ffd 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -66,7 +66,8 @@ def test_training_loop_policy_and_value_function(self):
         env=env,
         epochs=num_epochs,
         policy_and_value_net_fun=functools.partial(
-            ppo.policy_and_value_net, bottom_layers=[layers.Dense(1)]),
+            ppo.policy_and_value_net,
+            bottom_layers_fn=lambda: [layers.Dense(1)]),
         policy_and_value_optimizer_fun=ppo.optimizer_fun,
         batch_size=batch_size,
         num_optimizer_steps=1,

From 88f6336d9269f2d92663be30442f86ca66d3369b Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 9 May 2019 20:14:46 -0700
Subject: [PATCH 2020/2720] Typo.

PiperOrigin-RevId: 247543806
---
 tensor2tensor/trax/layers/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index bd56d554f..a71a72f20 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -47,7 +47,7 @@ def output_shape(self, input_shape):
     """The shape of the output of this layer given the shape of the input.
 
     Note that all arguments and return values can be tuples or dictionaries
-    or arbitraty nested structures composed of tuples and dictionaries.
+    or arbitrary nested structures composed of tuples and dictionaries.
 
     Args:
       input_shape: a tuple representing the shape of the input.

From 2d4b8fce795f5815e0a7b33fd8b5854dbc8c2445 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 9 May 2019 21:37:45 -0700
Subject: [PATCH 2021/2720] Add a new combinator `Rebatch` that treats the
 specified starting number of dimensions as batch.

This is useful if we want to operate on input like (batch, time, actual_input)
by considering batch and time to be a virtual batch.

PiperOrigin-RevId: 247550396
---
 tensor2tensor/trax/layers/combinators.py      | 54 +++++++++++++++++++
 tensor2tensor/trax/layers/combinators_test.py | 20 +++++++
 2 files changed, 74 insertions(+)

diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 7b789276e..80a66ec5a 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -23,6 +23,7 @@
 import six
 
 from tensor2tensor.trax import backend
+from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers import base
 
 
@@ -436,3 +437,56 @@ def new_parameters(self, input_shape, rng):
           raise ValueError('Map layer can only be applied to list of elements '
                            'with the same shapes. Shapes: %s' % str(shape))
     return self._layer.initialize(first_shape, rng)
+
+
+class Rebatch(base.Layer):
+  """Combinator for treating the first `n` dims as batch.
+
+  Args:
+    layer: subclass of base.Layer, a layer to apply to the input.
+    num_batch_dims: int, the number of leading dimensions to consider as batch.
+
+  Returns:
+    A new layer that will reshape the input into a virtual batch, apply the
+    layer and unbatch the virtual batch.
+  """
+
+  def __init__(self, layer, num_batch_dims=1):
+    super(Rebatch, self).__init__()
+    self._layer = layer
+    self._num_batch_dims = num_batch_dims
+
+  def _modify_shape(self, input_shape):
+    input_shape = tuple(input_shape)
+    batch_dims, non_batch_dims = (input_shape[:self._num_batch_dims],
+                                  input_shape[self._num_batch_dims:])
+    new_batch_dim = six.moves.reduce(operator.mul, batch_dims)
+    return (new_batch_dim,) + non_batch_dims, batch_dims
+
+  def _unmodify_shape(self, input_shape, batch_dims):
+    return batch_dims + tuple(input_shape[1:])
+
+  def _modify(self, inp):
+    modified_shape, batch_dims = self._modify_shape(inp.shape)
+    return np.reshape(inp, modified_shape), batch_dims
+
+  def _unmodify(self, inp, batch_dims):
+    return np.reshape(inp, self._unmodify_shape(inp.shape, batch_dims))
+
+  def call(self, inp, params=(), **kwargs):
+    if isinstance(inp, (tuple, list)):
+      # TODO(afrozm): This should be easy to do though.
+      # Tip from Lukasz - base.nested_map(self._modify, inp)
+      raise ValueError("Rebatch doesn't support list/tuple inputs now.")
+    inp, batch_dims = self._modify(inp)
+    out = self._layer(inp, params=params, **kwargs)
+    return self._unmodify(out, batch_dims)
+
+  def output_shape(self, input_shape):
+    modified_shape, batch_dims = self._modify_shape(input_shape)
+    out = self._layer.output_shape(modified_shape)
+    return self._unmodify_shape(out, batch_dims)
+
+  def new_parameters(self, input_shape, rng):
+    modified_shape, _ = self._modify_shape(input_shape)
+    return self._layer.initialize(modified_shape, rng)
diff --git a/tensor2tensor/trax/layers/combinators_test.py b/tensor2tensor/trax/layers/combinators_test.py
index e8412304f..5092e3eb1 100644
--- a/tensor2tensor/trax/layers/combinators_test.py
+++ b/tensor2tensor/trax/layers/combinators_test.py
@@ -21,6 +21,7 @@
 from absl.testing import absltest
 from tensor2tensor.trax.layers import base
 from tensor2tensor.trax.layers import combinators
+from tensor2tensor.trax.layers import convolution
 
 
 class CombinatorLayerTest(absltest.TestCase):
@@ -76,6 +77,25 @@ def test_unnest_branches(self):
         combinators.UnnestBranches(), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
+  def test_rebatch(self):
+    input_shape = (29, 5, 5, 20)
+    result_shape = base.check_shape_agreement(
+        convolution.Conv(30, (3, 3)), input_shape)
+    self.assertEqual(result_shape, (29, 3, 3, 30))
+
+    input_shape = (29, 5, 5, 20)
+    result_shape = base.check_shape_agreement(
+        combinators.Rebatch(
+            convolution.Conv(30, (3, 3)), num_batch_dims=1),
+        input_shape)
+    self.assertEqual(result_shape, (29, 3, 3, 30))
+
+    input_shape = (19, 29, 5, 5, 20)
+    result_shape = base.check_shape_agreement(
+        combinators.Rebatch(
+            convolution.Conv(30, (3, 3)), num_batch_dims=2),
+        input_shape)
+    self.assertEqual(result_shape, (19, 29, 3, 3, 30))
 
 if __name__ == '__main__':
   absltest.main()

From 6bd3812ed8bc5a12dffc2b840b9bc6aa1aa27204 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Fri, 10 May 2019 20:41:27 -0700
Subject: [PATCH 2022/2720] Minor, allow continuous evaluation in t2t_eval

PiperOrigin-RevId: 247718824
---
 tensor2tensor/bin/t2t_eval.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/bin/t2t_eval.py b/tensor2tensor/bin/t2t_eval.py
index 3bdb4634d..4791972c2 100644
--- a/tensor2tensor/bin/t2t_eval.py
+++ b/tensor2tensor/bin/t2t_eval.py
@@ -51,8 +51,12 @@ def main(_):
 
   estimator = trainer_lib.create_estimator(
       FLAGS.model, hparams, config, use_tpu=FLAGS.use_tpu)
-  predictions = estimator.evaluate(eval_input_fn, steps=FLAGS.eval_steps)
-  tf.logging.info(predictions)
+  ckpt_iter = trainer_lib.next_checkpoint(
+      hparams.model_dir, FLAGS.eval_timeout_mins)
+  for ckpt_path in ckpt_iter:
+    predictions = estimator.evaluate(
+        eval_input_fn, steps=FLAGS.eval_steps, checkpoint_path=ckpt_path)
+    tf.logging.info(predictions)
 
 
 if __name__ == "__main__":

From dad03c9335884d9d9e7b9d1dc502a406e551b23d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sat, 11 May 2019 08:47:43 -0700
Subject: [PATCH 2023/2720] [trax] fix optimizers.parallelize
 replicate/unreplicate logic for sm3

PiperOrigin-RevId: 247759135
---
 tensor2tensor/trax/optimizers.py | 19 +++++++++---------
 tensor2tensor/trax/trax_test.py  | 33 ++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/trax/optimizers.py b/tensor2tensor/trax/optimizers.py
index 4d8e1fef6..77bbfbb58 100644
--- a/tensor2tensor/trax/optimizers.py
+++ b/tensor2tensor/trax/optimizers.py
@@ -44,25 +44,26 @@ def opt_configure(*args, **kwargs):
 piecewise_constant = opt_configure(opt.piecewise_constant)
 
 
+# TODO(mattjj): upstream this to jax.experimental.optimizers.
 def parallelize(opt_maker):
-  """Transform an optimizer maker into a parallel one with replicated params."""
+  """Transform an optimizer maker into a parallel one with replicated state."""
+  num_devices = jax.lib.xla_bridge.device_count()
+  replicate_array = lambda x: onp.broadcast_to(x, (num_devices,) + x.shape)
+  unreplicate_array = lambda x: x.mean(0)  # an alternative is just x[0]
 
   def parallel_opt_maker(*args, **kwargs):  # pylint:disable=missing-docstring
     init_fun, update_fun, get_params = opt_maker(*args, **kwargs)
 
-    num_devices = jax.lib.xla_bridge.device_count()
-    replicate_array = lambda x: onp.broadcast_to(x, (num_devices,) + x.shape)
-    unreplicate_array = lambda x: x.mean(0)  # an alternative is x[0]
-
     def init_replicated(params):
+      opt_state = init_fun(params)
       if num_devices > 1:
-        params = jax.tree_util.tree_map(replicate_array, params)
-      return init_fun(params)
+        opt_state = jax.tree_util.tree_map(replicate_array, opt_state)
+      return opt_state
 
     def get_params_unreplicated(opt_state):
-      params = get_params(opt_state)
       if num_devices > 1:
-        params = jax.tree_util.tree_map(unreplicate_array, params)
+        opt_state = jax.tree_util.tree_map(unreplicate_array, opt_state)
+      params = get_params(opt_state)
       return params
 
     return init_replicated, update_fun, get_params, get_params_unreplicated
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index 682635404..eb2ec5e9a 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -29,6 +29,7 @@
 
 from tensor2tensor.trax import inputs as inputs_lib
 from tensor2tensor.trax import models
+from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
 
 from tensorflow import test
@@ -91,6 +92,38 @@ def test_train_eval_predict(self):
       inputs = inputs(1).train_stream()
       model()(next(inputs)[0], state.params)
 
+  def test_train_eval_predict_sm3(self):
+    with self.tmp_dir() as output_dir:
+      # Prepare model and inputs
+      num_classes = 4
+      train_steps = 2
+      eval_steps = 2
+      model = functools.partial(models.MLP,
+                                hidden_size=16,
+                                num_output_classes=num_classes)
+      inputs = lambda _: test_inputs(num_classes)
+
+      # Train and evaluate
+      state = trax.train(output_dir,
+                         model=model,
+                         inputs=inputs,
+                         train_steps=train_steps,
+                         eval_steps=eval_steps,
+                         optimizer=trax_opt.sm3)
+
+      # Assert total train steps
+      self.assertEqual(train_steps, state.step)
+
+      # Assert 2 evaluations ran
+      train_acc = state.history.get("train", "metrics/accuracy")
+      eval_acc = state.history.get("eval", "metrics/accuracy")
+      self.assertEqual(len(train_acc), len(eval_acc))
+      self.assertEqual(2, len(eval_acc))
+
+      # Predict with final params
+      inputs = inputs(1).train_stream()
+      model()(next(inputs)[0], state.params)
+
 
 if __name__ == "__main__":
   config.config_with_absl()

From 838aca4960f851cd759307481ea904038c1a1ab5 Mon Sep 17 00:00:00 2001
From: Shawn Simister <simister@google.com>
Date: Sun, 12 May 2019 00:41:06 -0700
Subject: [PATCH 2024/2720] Adding NeuralStack and NeuralQueue models under
 /research

PiperOrigin-RevId: 247806737
---
 tensor2tensor/layers/common_layers.py         |  30 ++
 tensor2tensor/models/__init__.py              |   1 +
 tensor2tensor/models/research/neural_stack.py | 483 ++++++++++++++++++
 .../models/research/neural_stack_test.py      | 291 +++++++++++
 4 files changed, 805 insertions(+)
 create mode 100644 tensor2tensor/models/research/neural_stack.py
 create mode 100644 tensor2tensor/models/research/neural_stack_test.py

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index a45b87a11..c1cd291c6 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1251,6 +1251,21 @@ def length_from_embedding(emb):
   return tf.cast(tf.reduce_sum(mask_from_embedding(emb), [1, 2, 3]), tf.int32)
 
 
+def mask_pos_gt(source_length, target_length):
+  """A mask with 1.0 wherever source_pos > target_pos and 0.0 elsewhere.
+
+  Args:
+    source_length: an integer
+    target_length: an integer
+  Returns:
+    a Tensor with shape [1, target_length, source_length]
+  """
+  return tf.expand_dims(
+      tf.cast(tf.greater(tf.expand_dims(tf.range(target_length), axis=0),
+                         tf.expand_dims(tf.range(source_length), axis=1)),
+              dtype=tf.float32), axis=0)
+
+
 def mask_leq(target_length, source_length):
   """A mask with 1.0 wherever source_pos <= target_pos and 0.0 elsewhere.
 
@@ -1268,6 +1283,21 @@ def mask_leq(target_length, source_length):
       out_shape=[1, target_length, source_length])
 
 
+def mask_pos_lt(source_length, target_length):
+  """A mask with 1.0 wherever source_pos < target_pos and 0.0 elsewhere.
+
+  Args:
+    source_length: an integer
+    target_length: an integer
+  Returns:
+    a Tensor with shape [1, target_length, source_length]
+  """
+  return tf.expand_dims(
+      tf.cast(tf.less(tf.expand_dims(tf.range(target_length), axis=0),
+                      tf.expand_dims(tf.range(source_length), axis=1)),
+              dtype=tf.float32), axis=0)
+
+
 def relu_density_logit(x, reduce_dims):
   """logit(density(x)).
 
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index b9fd5bf47..be54a772c 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -55,6 +55,7 @@
 from tensor2tensor.models.research import lm_experiments
 from tensor2tensor.models.research import moe_experiments
 from tensor2tensor.models.research import multiquery_paper
+from tensor2tensor.models.research import neural_stack
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research import similarity_transformer
 from tensor2tensor.models.research import super_lm
diff --git a/tensor2tensor/models/research/neural_stack.py b/tensor2tensor/models/research/neural_stack.py
new file mode 100644
index 000000000..ac6af5995
--- /dev/null
+++ b/tensor2tensor/models/research/neural_stack.py
@@ -0,0 +1,483 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Stacks and Queues implemented as encoder-decoder models.
+
+Based off of the following research:
+
+Learning to Transduce with Unbounded Memory
+Edward Grefenstette, Karl Moritz Hermann, Mustafa Suleyman, Phil Blunsom
+https://arxiv.org/abs/1506.02516, 2015
+
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_hparams
+from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
+
+import tensorflow as tf
+
+
+class NeuralStackCell(tf.nn.rnn_cell.RNNCell):
+  """An RNN cell base class that can implement a stack or queue.
+  """
+
+  def __init__(self, num_units, memory_size, embedding_size,
+               num_read_heads=1, num_write_heads=1, reuse=None):
+    """Create a new NeuralStackCell.
+
+    Args:
+      num_units: The number of hidden units in the RNN cell.
+      memory_size: The maximum memory size allocated for the stack.
+      embedding_size:  The embedding width of the individual stack values.
+      num_read_heads: This should always be 1 for a regular stack.
+      num_write_heads: This should always be 1 for a regular stack.
+      reuse: Whether to reuse the weights.
+    """
+    super(NeuralStackCell, self).__init__(dtype=tf.float32, _reuse=reuse)
+    self._num_units = num_units
+    self._embedding_size = embedding_size
+    self._memory_size = memory_size
+    self._num_read_heads = num_read_heads
+    self._num_write_heads = num_write_heads
+
+  @property
+  def state_size(self):
+    """The NeuralStackCell maintains a tuple of state values.
+
+    Returns:
+      (controller_state.shape,
+       read_values.shape,
+       memory_values.shape,
+       read_strengths.shape,
+       write_strengths.shape)
+    """
+    return (tf.TensorShape([self._num_units]),
+            tf.TensorShape([self._num_read_heads, self._embedding_size]),
+            tf.TensorShape([self._memory_size, self._embedding_size]),
+            tf.TensorShape([self._num_read_heads, self._memory_size, 1]),
+            tf.TensorShape([self._num_write_heads, self._memory_size, 1]))
+
+  @property
+  def output_size(self):
+    return tf.TensorShape([self._num_read_heads, self._embedding_size])
+
+  def initialize_write_strengths(self, batch_size):
+    """Initialize write strengths to write to the first memory address.
+
+    This is exposed as it's own function so that it can be overridden to provide
+    alternate write adressing schemes.
+
+    Args:
+      batch_size: The size of the current batch.
+
+    Returns:
+      A tf.float32 tensor of shape [num_write_heads, memory_size, 1] where the
+      first element in the second dimension is set to 1.0.
+    """
+    return tf.expand_dims(
+        tf.one_hot([[0] * self._num_write_heads] * batch_size,
+                   depth=self._memory_size, dtype=tf.float32), axis=3)
+
+  def zero_state(self, batch_size, dtype):
+    """Initialize the tuple of state values to zeros except write strengths.
+
+    Args:
+      batch_size: The size of the current batch.
+      dtype: The default datatype to initialize to.
+
+    Returns:
+      (controller_state.shape,
+       read_values.shape,
+       memory_values.shape,
+       read_strengths.shape,
+       write_strengths.shape)
+    """
+    state = list(super(NeuralStackCell, self).zero_state(batch_size, dtype))
+    state[4] = self.initialize_write_strengths(batch_size)
+    return tuple(state)
+
+  def build_read_mask(self):
+    """Creates a mask which allows us to attenuate subsequent read strengths.
+
+    This is exposed as it's own function so that it can be overridden to provide
+    alternate read adressing schemes.
+
+    Returns:
+      A tf.float32 tensor of shape [1, memory_size, memory_size]
+    """
+    return common_layers.mask_pos_gt(self._memory_size, self._memory_size)
+
+  def add_scalar_projection(self, name, size):
+    """A helper function for mapping scalar controller outputs.
+
+    Args:
+      name: A prefix for the variable names.
+      size: The desired number of scalar outputs.
+
+    Returns:
+      A tuple of (weights, bias) where weights has shape [num_units, size] and
+      bias has shape [size].
+    """
+    weights = self.add_variable(
+        name + "_projection_weights",
+        shape=[self._num_units, size],
+        dtype=self.dtype)
+    bias = self.add_variable(
+        name + "_projection_bias",
+        shape=[size],
+        initializer=tf.zeros_initializer(dtype=self.dtype))
+    return weights, bias
+
+  def add_vector_projection(self, name, size):
+    """A helper function for mapping embedding controller outputs.
+
+    Args:
+      name: A prefix for the variable names.
+      size: The desired number of embedding outputs.
+
+    Returns:
+      A tuple of (weights, bias) where weights has shape
+      [num_units, size * embedding_size] and bias has shape
+      [size * embedding_size].
+    """
+    weights = self.add_variable(
+        name + "_projection_weights",
+        shape=[self._num_units, size * self._embedding_size],
+        dtype=self.dtype)
+    bias = self.add_variable(
+        name + "_projection_bias",
+        shape=[size * self._embedding_size],
+        initializer=tf.zeros_initializer(dtype=self.dtype))
+    return weights, bias
+
+  def build_controller(self):
+    """Create the RNN and output projections for controlling the stack.
+    """
+    with tf.name_scope("controller"):
+      self.rnn = tf.contrib.rnn.BasicRNNCell(self._num_units)
+      self._input_proj = self.add_variable(
+          "input_projection_weights",
+          shape=[(self._embedding_size * self._num_read_heads) +
+                 (self._embedding_size * self._num_write_heads),
+                 self._num_units],
+          dtype=self.dtype)
+      self._input_bias = self.add_variable(
+          "input_projection_bias",
+          shape=[self._num_units],
+          initializer=tf.zeros_initializer(dtype=self.dtype))
+      self._push_proj, self._push_bias = self.add_scalar_projection(
+          "push", self._num_write_heads)
+      self._pop_proj, self._pop_bias = self.add_scalar_projection(
+          "pop", self._num_write_heads)
+      self._value_proj, self._value_bias = self.add_vector_projection(
+          "value", self._num_write_heads)
+      self._output_proj, self._output_bias = self.add_vector_projection(
+          "output", self._num_read_heads)
+
+  def build(self, _):
+    """Build the controller, read mask and write shift convolutional filter.
+
+    The write shift convolutional filter is a simple 3x3 convolution which is
+    used to advance the read heads to the next memory address at each step. This
+    filter can be changed to move the read heads in other ways.
+    """
+    self.read_mask = self.build_read_mask()
+    self.write_shift_convolution = tf.reshape(tf.one_hot([[3]], depth=9),
+                                              shape=[3, 3, 1, 1])
+    self.build_controller()
+
+    self.built = True
+
+  def call_controller(self, inputs, state, batch_size):
+    """Make a call to the neural stack controller.
+
+    See Section 3.1 of Grefenstette et al., 2015.
+
+    Args:
+      inputs: The combined inputs to the controller consisting of the current
+         input value concatenated with the read values from the previous
+         timestep with shape [batch_size, (num_write_heads + num_read_heads)
+         * embedding_size].
+      state: The hidden state from the previous time step.
+      batch_size: The size of the current batch of input values.
+
+    Returns:
+      A tuple of outputs and the new hidden state value:
+      (push_strengths, pop_strengths, write_values, outputs, state)
+    """
+    with tf.name_scope("controller"):
+      rnn_input = tf.tanh(tf.nn.bias_add(tf.matmul(
+          inputs, self._input_proj), self._input_bias))
+
+      (rnn_output, state) = self.rnn(rnn_input, state)
+
+      push_strengths = tf.reshape(
+          tf.sigmoid(tf.nn.bias_add(tf.matmul(
+              rnn_output, self._push_proj), self._push_bias)),
+          shape=[batch_size, self._num_write_heads, 1, 1])
+
+      pop_strengths = tf.reshape(
+          tf.sigmoid(tf.nn.bias_add(tf.matmul(
+              rnn_output, self._pop_proj), self._pop_bias)),
+          shape=[batch_size, self._num_write_heads, 1, 1])
+
+      write_values = tf.reshape(
+          tf.tanh(tf.nn.bias_add(tf.matmul(
+              rnn_output, self._value_proj), self._value_bias)),
+          shape=[batch_size, self._num_read_heads, self._embedding_size])
+
+      outputs = tf.reshape(
+          tf.tanh(tf.nn.bias_add(tf.matmul(
+              rnn_output, self._output_proj), self._output_bias)),
+          shape=[batch_size, self._num_read_heads, self._embedding_size])
+
+    return push_strengths, pop_strengths, write_values, outputs, state
+
+  def call(self, inputs, state):
+    """Evaluates one timestep of the current neural stack cell.
+
+    See section 3.4 of Grefenstette et al., 2015.
+
+    Args:
+      inputs: The inputs to the neural stack cell should be a tf.float32 tensor
+        with shape [batch_size, max_timesteps, 1, embedding_size]
+      state: The tuple of state values from the previous timestep.
+
+    Returns:
+      The output value of the stack as well as the new tuple of state values.
+      (outputs, (controller_state, read_values, memory_values, read_strengths,
+                 write_strengths))
+    """
+    (controller_state,
+     read_values,
+     memory_values,
+     read_strengths,
+     write_strengths) = state
+
+    batch_size = tf.shape(inputs)[0]
+
+    # Concatenate the current input value with the read value from  the previous
+    # timestep before feeding them into the controller.
+    controller_inputs = tf.concat([
+        tf.reshape(
+            read_values,
+            shape=[batch_size, self._num_read_heads * self._embedding_size]),
+        tf.reshape(
+            inputs,
+            shape=[batch_size, self._num_write_heads * self._embedding_size])
+    ], axis=1)
+
+    # Call the controller and get controller interface values.
+    (push_strengths, pop_strengths,
+     write_values, outputs, controller_state) = self.call_controller(
+         controller_inputs, controller_state, batch_size)
+
+    # Always write input values to memory regardless of push strength.
+    # See Equation-1 in Grefenstette et al., 2015.
+    memory_values += tf.reduce_sum(
+        tf.expand_dims(write_values, axis=1) * write_strengths, axis=1)
+
+    # Attenuate the read strengths of existing memory values depending on the
+    # current pop strength.
+    # See Equation-2 in Grefenstette et al., 2015.
+    read_strengths = tf.nn.relu(
+        read_strengths - tf.nn.relu(pop_strengths - tf.reduce_sum(
+            tf.reshape(read_strengths,
+                       shape=[batch_size, 1, 1, self._memory_size]) *
+            self.read_mask, axis=3, keepdims=True)))
+
+    # Set read strength for the current timestep based on the push strength.
+    read_strengths = read_strengths + push_strengths * write_strengths
+
+    # Calculate the "top" value of the stack by looking at read strengths.
+    # See Equation-3 in Grefenstette et al., 2015.
+    read_values = tf.reduce_sum(
+        tf.minimum(
+            read_strengths,
+            tf.nn.relu(1 - tf.reshape(
+                tf.reduce_sum(read_strengths * self.read_mask,
+                              axis=2,
+                              keepdims=True),
+                shape=[
+                    batch_size, self._num_read_heads, self._memory_size, 1
+                ]))) * tf.expand_dims(memory_values, axis=1),
+        axis=2)
+
+    # Shift the write strengths forward by one memory address for the next step.
+    write_strengths = tf.nn.conv2d(
+        write_strengths, self.write_shift_convolution, [1, 1, 1, 1],
+        padding="SAME")
+
+    return (outputs, (controller_state,
+                      read_values,
+                      memory_values,
+                      read_strengths,
+                      write_strengths))
+
+
+class NeuralQueueCell(NeuralStackCell):
+  """An subclass of the NeuralStackCell which reads from the opposite direction.
+
+  See section 3.2 of Grefenstette et al., 2015.
+  """
+
+  def build_read_mask(self):
+    """Uses mask_pos_lt() instead of mask_pos_gt() to reverse read values.
+
+    Returns:
+      A tf.float32 tensor of shape [1, memory_size, memory_size].
+    """
+    return common_layers.mask_pos_lt(self._memory_size, self._memory_size)
+
+
+@registry.register_model
+class NeuralStackModel(t2t_model.T2TModel):
+  """An encoder-decoder T2TModel that uses NeuralStackCells.
+  """
+
+  def cell(self, hidden_size):
+    """Build an RNN cell.
+
+    This is exposed as it's own function so that it can be overridden to provide
+    different types of RNN cells.
+
+    Args:
+      hidden_size: The hidden size of the cell.
+
+    Returns:
+      A new RNNCell with the given hidden size.
+    """
+    return NeuralStackCell(hidden_size,
+                           self._hparams.memory_size,
+                           self._hparams.embedding_size)
+
+  def _rnn(self, inputs, name, initial_state=None, sequence_length=None):
+    """A helper method to build tf.nn.dynamic_rnn.
+
+    Args:
+      inputs: The inputs to the RNN. A tensor of shape
+              [batch_size, max_seq_length, embedding_size]
+      name: A namespace for the RNN.
+      initial_state: An optional initial state for the RNN.
+      sequence_length: An optional sequence length for the RNN.
+
+    Returns:
+      A tf.nn.dynamic_rnn operator.
+    """
+    layers = [self.cell(layer_size)
+              for layer_size in self._hparams.controller_layer_sizes]
+    with tf.variable_scope(name):
+      return tf.nn.dynamic_rnn(
+          tf.contrib.rnn.MultiRNNCell(layers),
+          inputs,
+          initial_state=initial_state,
+          sequence_length=sequence_length,
+          dtype=tf.float32,
+          time_major=False)
+
+  def body(self, features):
+    """Build the main body of the model.
+
+    Args:
+      features: A dict of "inputs" and "targets" which have already been passed
+        through an embedding layer. Inputs should have shape
+        [batch_size, max_seq_length, 1, embedding_size]. Targets should have
+        shape [batch_size, max_seq_length, 1, 1]
+
+    Returns:
+      The logits which get passed to the top of the model for inference.
+      A tensor of shape [batch_size, seq_length, 1, embedding_size]
+    """
+    inputs = features.get("inputs")
+    targets = features["targets"]
+
+    if inputs is not None:
+      inputs = common_layers.flatten4d3d(inputs)
+      _, final_encoder_state = self._rnn(tf.reverse(inputs, axis=[1]),
+                                         "encoder")
+    else:
+      final_encoder_state = None
+
+    shifted_targets = common_layers.shift_right(targets)
+    decoder_outputs, _ = self._rnn(
+        common_layers.flatten4d3d(shifted_targets),
+        "decoder",
+        initial_state=final_encoder_state)
+    return decoder_outputs
+
+
+@registry.register_model
+class NeuralQueueModel(NeuralStackModel):
+  """Subcalss of NeuralStackModel which implements a queue.
+  """
+
+  def cell(self, hidden_size):
+    """Build a NeuralQueueCell instead of a NeuralStackCell.
+
+    Args:
+      hidden_size: The hidden size of the cell.
+
+    Returns:
+      A new NeuralQueueCell with the given hidden size.
+    """
+    return NeuralQueueCell(hidden_size,
+                           self._hparams.memory_size,
+                           self._hparams.embedding_size)
+
+
+@registry.register_hparams
+def lstm_transduction():
+  """HParams for LSTM base on transduction tasks."""
+  hparams = common_hparams.basic_params1()
+  hparams.daisy_chain_variables = False
+  hparams.batch_size = 10
+  hparams.clip_grad_norm = 1.0
+  hparams.hidden_size = 128
+  hparams.num_hidden_layers = 4
+  hparams.initializer = "uniform_unit_scaling"
+  hparams.initializer_gain = 1.0
+  hparams.optimizer = "RMSProp"
+  hparams.learning_rate = 0.01
+  hparams.weight_decay = 0.0
+
+  hparams.add_hparam("memory_size", 128)
+  hparams.add_hparam("embedding_size", 32)
+  return hparams
+
+
+@registry.register_hparams
+def neural_stack():
+  """HParams for neural stacks and queues."""
+  hparams = common_hparams.basic_params1()
+  hparams.daisy_chain_variables = False
+  hparams.batch_size = 10
+  hparams.clip_grad_norm = 1.0
+  hparams.initializer = "uniform_unit_scaling"
+  hparams.initializer_gain = 1.0
+  hparams.optimizer = "RMSProp"
+  hparams.learning_rate = 0.0001
+  hparams.weight_decay = 0.0
+
+  hparams.add_hparam("controller_layer_sizes", [256, 512])
+  hparams.add_hparam("memory_size", 128)
+  hparams.add_hparam("embedding_size", 64)
+  hparams.hidden_size = hparams.embedding_size
+  return hparams
diff --git a/tensor2tensor/models/research/neural_stack_test.py b/tensor2tensor/models/research/neural_stack_test.py
new file mode 100644
index 000000000..79a31662b
--- /dev/null
+++ b/tensor2tensor/models/research/neural_stack_test.py
@@ -0,0 +1,291 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 NeuralStackCell, NeuralQueueCell and NeuralStackModel."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import mock
+import numpy as np
+
+from tensor2tensor.layers import modalities
+from tensor2tensor.models.research import neural_stack
+
+import tensorflow as tf
+
+
+def build_fake_controller(cell):
+  """Create a scalar variable to track the timestep.
+
+  Args:
+    cell: The NeuralStackCell to add the variable to.
+  """
+  cell.current_step = cell.add_variable(
+      "current_step", [],
+      initializer=tf.constant_initializer(-1),
+      dtype=tf.int32,
+      trainable=False)
+
+
+def call_fake_controller(push_values, pop_values, read_values, output_values):
+  """Mock a RNN controller from a set of expected outputs.
+
+  Args:
+    push_values: Expected controller push values.
+    pop_values: Expected controller pop values.
+    read_values: Expected controller read values.
+    output_values: Expected controller output values.
+
+  Returns:
+    A callable which behaves like the call method of an NeuralStackCell.
+  """
+  def call(cell, inputs, state, batch_size):
+    del inputs
+    del batch_size
+    next_step = tf.assign_add(cell.current_step, tf.constant(1))
+    return (
+        tf.slice(tf.constant(push_values), [next_step, 0], [1, -1]),
+        tf.slice(tf.constant(pop_values), [next_step, 0], [1, -1]),
+        tf.slice(tf.constant(read_values), [next_step, 0, 0], [1, -1, -1]),
+        tf.slice(tf.constant(output_values), [next_step, 0, 0], [1, -1, -1]),
+        state
+    )
+  return call
+
+
+class NeuralStackCellTest(tf.test.TestCase):
+
+  def test_controller_shapes(self):
+    """Check that all the NeuralStackCell tensor shapes are correct.
+    """
+
+    batch_size = 5
+    embedding_size = 3
+    memory_size = 6
+    num_units = 8
+
+    stack = neural_stack.NeuralStackCell(num_units, memory_size, embedding_size)
+
+    stack.build(None)
+
+    self.assertEqual([1, embedding_size], stack.output_size)
+    self.assertEqual([1, memory_size, memory_size], stack.read_mask.shape)
+    self.assertEqual([3, 3, 1, 1], stack.write_shift_convolution.shape)
+
+    stack_input = tf.zeros([batch_size, 1, embedding_size], dtype=tf.float32)
+
+    zero_state = stack.zero_state(batch_size, tf.float32)
+
+    (controller_state,
+     previous_values,
+     memory_values,
+     read_strengths,
+     write_strengths) = zero_state
+
+    self.assertEqual([batch_size, num_units], controller_state.shape)
+    self.assertEqual([batch_size, 1, embedding_size], previous_values.shape)
+    self.assertEqual([batch_size, memory_size, embedding_size],
+                     memory_values.shape)
+    self.assertEqual([batch_size, 1, memory_size, 1], read_strengths.shape)
+    self.assertEqual([batch_size, 1, memory_size, 1], write_strengths.shape)
+
+    rnn_input = tf.concat([
+        tf.reshape(
+            previous_values,
+            shape=[batch_size, embedding_size]),
+        tf.reshape(
+            stack_input,
+            shape=[batch_size, embedding_size])
+    ], axis=1)
+    self.assertEqual([batch_size, 2 * embedding_size], rnn_input.shape)
+
+    (push_strengths,
+     pop_strengths,
+     new_values,
+     outputs,
+     controller_next_state) = stack.call_controller(rnn_input,
+                                                    controller_state,
+                                                    batch_size)
+
+    self.assertEqual([batch_size, 1, 1, 1], push_strengths.shape)
+    self.assertEqual([batch_size, 1, 1, 1], pop_strengths.shape)
+    self.assertEqual([batch_size, 1, embedding_size], new_values.shape)
+    self.assertEqual([batch_size, 1, embedding_size], outputs.shape)
+    self.assertEqual([batch_size, num_units], controller_next_state.shape)
+
+    (outputs, (controller_next_state,
+               read_values,
+               next_memory_values,
+               next_read_strengths,
+               next_write_strengths)) = stack.call(stack_input, zero_state)
+
+    self.assertEqual([batch_size, 1, embedding_size], outputs.shape)
+    self.assertEqual([batch_size, num_units], controller_next_state.shape)
+    self.assertEqual([batch_size, 1, embedding_size], read_values.shape)
+    self.assertEqual([batch_size, memory_size, embedding_size],
+                     next_memory_values.shape)
+    self.assertEqual([batch_size, 1, memory_size, 1], next_read_strengths.shape)
+    self.assertEqual([batch_size, 1, memory_size, 1],
+                     next_write_strengths.shape)
+
+    # Make sure that stack output shapes match stack input shapes
+    self.assertEqual(controller_next_state.shape, controller_state.shape)
+    self.assertEqual(read_values.shape, previous_values.shape)
+    self.assertEqual(next_memory_values.shape, memory_values.shape)
+    self.assertEqual(next_read_strengths.shape, read_strengths.shape)
+    self.assertEqual(next_write_strengths.shape, write_strengths.shape)
+
+  @mock.patch.object(neural_stack.NeuralStackCell, "build_controller",
+                     build_fake_controller)
+  @mock.patch.object(neural_stack.NeuralStackCell, "call_controller",
+                     call_fake_controller(
+                         push_values=[[1.0], [1.0], [0.0]],
+                         pop_values=[[0.0], [0.0], [1.0]],
+                         read_values=[[[1.0, 0.0, 0.0]],
+                                      [[0.0, 1.0, 0.0]],
+                                      [[0.0, 0.0, 1.0]]],
+                         output_values=[[[0.0, 0.0, 0.0]],
+                                        [[0.0, 0.0, 0.0]],
+                                        [[0.0, 0.0, 0.0]]]))
+  def test_push_pop(self):
+    """Test pushing a popping from a NeuralStackCell.
+    """
+    input_values = np.array([[[[1.0, 0.0, 0.0]],
+                              [[0.0, 1.0, 0.0]],
+                              [[0.0, 0.0, 1.0]]]])
+
+    expected_values = np.array([[[1.0, 0.0, 0.0],
+                                 [0.0, 1.0, 0.0],
+                                 [0.0, 0.0, 1.0],
+                                 [0.0, 0.0, 0.0],
+                                 [0.0, 0.0, 0.0],
+                                 [0.0, 0.0, 0.0]]])
+    expected_read_strengths = np.array([
+        [[[1.0], [0.0], [0.0], [0.0], [0.0], [0.0]]]])
+    expected_write_strengths = np.array([
+        [[[0.0], [0.0], [0.], [1.0], [0.0], [0.0]]]])
+    expected_top = np.array([[[1.0, 0.0, 0.0]]])
+
+    stack = neural_stack.NeuralStackCell(8, 6, 3)
+    stack_input = tf.constant(input_values, dtype=tf.float32)
+    (outputs, state) = tf.nn.dynamic_rnn(cell=stack,
+                                         inputs=stack_input,
+                                         time_major=False,
+                                         dtype=tf.float32)
+
+    with self.test_session() as sess:
+      sess.run(tf.global_variables_initializer())
+      _, state_vals = sess.run([outputs, state])
+      (_, stack_top, values, read_strengths, write_strengths) = state_vals
+
+      self.assertAllClose(expected_top, stack_top)
+      self.assertAllClose(expected_values, values)
+      self.assertAllClose(expected_read_strengths, read_strengths)
+      self.assertAllClose(expected_write_strengths, write_strengths)
+
+
+class NeuralQueueCellTest(tf.test.TestCase):
+
+  @mock.patch.object(neural_stack.NeuralQueueCell, "build_controller",
+                     build_fake_controller)
+  @mock.patch.object(neural_stack.NeuralQueueCell, "call_controller",
+                     call_fake_controller(
+                         push_values=[[1.0], [1.0], [0.0]],
+                         pop_values=[[0.0], [0.0], [1.0]],
+                         read_values=[[[1.0, 0.0, 0.0]],
+                                      [[0.0, 1.0, 0.0]],
+                                      [[0.0, 0.0, 1.0]]],
+                         output_values=[[[0.0, 0.0, 0.0]],
+                                        [[0.0, 0.0, 0.0]],
+                                        [[0.0, 0.0, 0.0]]]))
+  def test_enqueue_dequeue(self):
+    """Test enqueueing a dequeueing from a NeuralQueueCell.
+    """
+    input_values = np.array([[[[1.0, 0.0, 0.0]],
+                              [[0.0, 1.0, 0.0]],
+                              [[0.0, 0.0, 1.0]]]])
+    expected_values = np.array([[[1.0, 0.0, 0.0],
+                                 [0.0, 1.0, 0.0],
+                                 [0.0, 0.0, 1.0],
+                                 [0.0, 0.0, 0.0],
+                                 [0.0, 0.0, 0.0],
+                                 [0.0, 0.0, 0.0]]])
+    expected_read_strengths = np.array([
+        [[[0.0], [1.0], [0.0], [0.0], [0.0], [0.0]]]])
+    expected_write_strengths = np.array([
+        [[[0.0], [0.0], [0.0], [1.0], [0.0], [0.0]]]])
+    expected_front = np.array([[[0.0, 1.0, 0.0]]])
+
+    queue = neural_stack.NeuralQueueCell(8, 6, 3)
+    rnn_input = tf.constant(input_values, dtype=tf.float32)
+    (outputs, state) = tf.nn.dynamic_rnn(cell=queue,
+                                         inputs=rnn_input,
+                                         time_major=False,
+                                         dtype=tf.float32)
+
+    with self.test_session() as sess:
+      sess.run(tf.global_variables_initializer())
+      _, state_vals = sess.run([outputs, state])
+      (_, queue_front, values, read_strengths, write_strengths) = state_vals
+
+      self.assertAllClose(expected_front, queue_front)
+      self.assertAllClose(expected_values, values)
+      self.assertAllClose(expected_read_strengths, read_strengths)
+      self.assertAllClose(expected_write_strengths, write_strengths)
+
+
+class NeuralStackModelTest(tf.test.TestCase):
+
+  def test_model_shapes(self):
+    """Test a few of the important output shapes for NeuralStackModel.
+    """
+    batch_size = 100
+    seq_length = 80
+    embedding_size = 64
+    vocab_size = 128
+
+    hparams = neural_stack.neural_stack()
+    problem_hparams = tf.contrib.training.HParams()
+
+    problem_hparams.add_hparam("modality", {
+        "inputs": modalities.ModalityType.SYMBOL,
+        "targets": modalities.ModalityType.SYMBOL,
+    })
+    problem_hparams.add_hparam("vocab_size", {
+        "inputs": vocab_size,
+        "targets": vocab_size,
+    })
+    model = neural_stack.NeuralStackModel(hparams,
+                                          problem_hparams=problem_hparams)
+
+    features = {
+        "inputs": tf.ones([batch_size, seq_length, 1, 1],
+                          dtype=tf.int32),
+        "targets": tf.ones([batch_size, seq_length, 1, 1], dtype=tf.int32)
+    }
+
+    transformed_features = model.bottom(features)
+
+    self.assertEqual([batch_size, seq_length, 1, embedding_size],
+                     transformed_features["inputs"].shape)
+
+    logits = model.body(transformed_features)
+
+    self.assertEqual([batch_size, seq_length, 1, embedding_size], logits.shape)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 21695e8d4fc4a20ba3ab4df643a69140239bab0a Mon Sep 17 00:00:00 2001
From: Shawn Simister <simister@google.com>
Date: Mon, 13 May 2019 10:55:00 -0700
Subject: [PATCH 2025/2720] Adding a new eval metric that tracks the average
 accuracy of a generated sequence by the measuring the longest prefix of
 matching symbols instead of matching the whole sequence.

PiperOrigin-RevId: 247968702
---
 tensor2tensor/utils/metrics.py      | 37 ++++++++++++++++++++++++++++-
 tensor2tensor/utils/metrics_test.py | 20 ++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 91aaad345..e9af1a22f 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -51,6 +51,7 @@ class Metrics(object):
   ROUGE_2_F = "rouge_2_fscore"
   ROUGE_L_F = "rouge_L_fscore"
   EDIT_DISTANCE = "edit_distance"
+  PREFIX_ACCURACY = "prefix_accuracy"
   WORD_ERROR_RATE = "word_error_rate"
   SET_PRECISION = "set_precision"
   SET_RECALL = "set_recall"
@@ -197,6 +198,41 @@ def padded_sequence_accuracy(predictions,
     return correct_seq, tf.constant(1.0)
 
 
+def prefix_accuracy(predictions,
+                    labels,
+                    weights_fn=common_layers.weights_nonzero):
+  """Average # of correct tokens at start of sequences, ignoring padding 0s.
+
+  See section 4.3 of Learning to Transduce with Unbounded Memory,
+  Grefenstette et al., 2015.
+
+  Args:
+    predictions: Tensor of shape [`batch_size`, `length`, 1, `num_classes`] and
+        type tf.float32 representing the logits, 0-padded.
+    labels: Tensor of shape [`batch_size`, `length`, 1, 1] and type tf.int32
+        representing the labels of same length as logits and 0-padded.
+    weights_fn: ignored. The weights returned are the total length of the ground
+        truth labels, excluding 0-paddings.
+
+  Returns:
+    (prefix accuracy, 1.0)
+
+  Raises:
+    ValueError: if weights_fn is not common_layers.weights_nonzero.
+  """
+  if weights_fn is not common_layers.weights_nonzero:
+    raise ValueError("Only weights_nonzero can be used for this metric.")
+
+  predictions = tf.to_int32(tf.squeeze(tf.argmax(predictions, axis=-1), axis=2))
+  labels = tf.squeeze(labels, axis=(2, 3))
+  seq_len = tf.reduce_sum(
+      tf.cast(tf.not_equal(labels, tf.constant(0)), dtype=tf.float32), axis=1)
+  matching_elements = tf.equal(labels, predictions)
+  prefix_len = tf.reduce_sum(
+      tf.cumprod(tf.cast(matching_elements, tf.float32), axis=1), axis=1)
+  return tf.reduce_mean(prefix_len / seq_len), tf.constant(1.0)
+
+
 def sequence_edit_distance(predictions,
                            labels,
                            weights_fn=common_layers.weights_nonzero):
@@ -798,7 +834,6 @@ def pearson_correlation_coefficient(predictions, labels, weights_fn=None):
     Metrics.ROUGE_2_F: rouge.rouge_2_fscore,
     Metrics.ROUGE_L_F: rouge.rouge_l_fscore,
     Metrics.EDIT_DISTANCE: sequence_edit_distance,
-    Metrics.WORD_ERROR_RATE: word_error_rate,
     Metrics.SOFTMAX_CROSS_ENTROPY_ONE_HOT: softmax_cross_entropy_one_hot,
     Metrics.SIGMOID_ACCURACY_ONE_HOT: sigmoid_accuracy_one_hot,
     Metrics.SIGMOID_RECALL_ONE_HOT: sigmoid_recall_one_hot,
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index ede480d8f..e259f5124 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -56,6 +56,26 @@ def testAccuracyTopKMetric(self):
     self.assertAlmostEqual(actual1, expected)
     self.assertAlmostEqual(actual2, 1.0)
 
+  def testPrefixAccuracy(self):
+    vocab_size = 10
+    predictions = tf.one_hot(
+        tf.constant([[[1], [2], [3], [4], [9], [6], [7], [8]],
+                     [[1], [2], [3], [4], [5], [9], [7], [8]],
+                     [[1], [2], [3], [4], [5], [9], [7], [0]]]),
+        vocab_size)
+    labels = tf.expand_dims(
+        tf.constant([[[1], [2], [3], [4], [5], [6], [7], [8]],
+                     [[1], [2], [3], [4], [5], [6], [7], [8]],
+                     [[1], [2], [3], [4], [5], [6], [7], [0]]]),
+        axis=-1)
+    expected_accuracy = np.average([4.0 / 8.0,
+                                    5.0 / 8.0,
+                                    5.0 / 7.0])
+    accuracy, _ = metrics.prefix_accuracy(predictions, labels)
+    with self.test_session() as session:
+      accuracy_value = session.run(accuracy)
+      self.assertAlmostEqual(expected_accuracy, accuracy_value)
+
   def testSequenceAccuracyMetric(self):
     predictions = np.random.randint(4, size=(12, 12, 12, 1))
     targets = np.random.randint(4, size=(12, 12, 12, 1))

From 3ca0d7718b521a3c1ca97f29d2b75de0825e03a4 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 13 May 2019 15:17:14 -0700
Subject: [PATCH 2026/2720] While sampling in `env_problem_utils.py` cast up to
 `np.float64` This seems to fix things.

PiperOrigin-RevId: 248021102
---
 tensor2tensor/envs/env_problem_utils.py | 17 +++++------------
 1 file changed, 5 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index f5b05ff17..04327e15a 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -112,19 +112,12 @@ def multinomial_sample(probs):
     # Convert to probs, since we need to do categorical sampling.
     probs = np.exp(log_probs)
 
-    # Sometimes log_probs contains a 0, it shouldn't. This makes the
-    # probabilities sum up to more than 1, since the addition happens
-    # in float64, so just add and subtract 1.0 to zero those probabilites
-    # out.
+    # Let's cast up to float64, because that's what numpy does when sampling
+    # and it leads to the sum(pvals[:-1]) > 1.0 error.
     #
-    # Also testing for this is brittle.
-    probs += 1
-    probs -= 1
-
-    # For some reason, sometimes, this isn't the case.
-    probs_sum = np.sum(probs, axis=1, keepdims=True)
-    if not all(probs_sum == 1.0):
-      probs = probs / probs_sum
+    # We also re-normalize when we do this.
+    probs = np.float64(probs)
+    probs /= np.sum(probs, axis=1, keepdims=True)
 
     # Now pick actions from this probs array.
     actions = np.apply_along_axis(multinomial_sample, 1, probs)

From ee38a70b41d05c2873198b21dc9e600ca7fcde6c Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 13 May 2019 17:47:23 -0700
Subject: [PATCH 2027/2720] Thread random number generator through PPO and
 policy collect.

PiperOrigin-RevId: 248047016
---
 tensor2tensor/envs/env_problem_utils.py      |   6 +-
 tensor2tensor/envs/env_problem_utils_test.py |   4 +-
 tensor2tensor/trax/rlax/ppo.py               | 106 +++++++++++++------
 tensor2tensor/trax/rlax/ppo_main.py          |   8 +-
 tensor2tensor/trax/rlax/ppo_test.py          |  46 +++++---
 5 files changed, 117 insertions(+), 53 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 04327e15a..44a4e5133 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -50,7 +50,8 @@ def play_env_problem_with_policy(env,
                                  policy_fun,
                                  num_trajectories=1,
                                  max_timestep=None,
-                                 boundary=20):
+                                 boundary=20,
+                                 rng=None):
   """Plays the given env with the policy function to collect trajectories.
 
   Args:
@@ -64,6 +65,7 @@ def play_env_problem_with_policy(env,
     boundary: this is the bucket length, we pad the observations to integer
         multiples of this + 1 and then feed the padded observations to the
         policy_fun.
+    rng: jax rng, splittable.
 
   Returns:
     Completed trajectories that is a list of triples of (observation, action,
@@ -96,7 +98,7 @@ def multinomial_sample(probs):
     assert B == env.batch_size
     assert (B,) == lengths.shape
 
-    log_prob_actions = policy_fun(padded_observations)
+    log_prob_actions, _, rng = policy_fun(padded_observations, rng=rng)
     assert (B, T) == log_prob_actions.shape[:2]
     A = log_prob_actions.shape[2]  # pylint: disable=invalid-name
 
diff --git a/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index 8b335c9d5..50dd8683b 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -52,13 +52,13 @@ def test_play_env_problem_with_policy(self):
         batch_size=2,
         reward_range=(-1, 1))
 
-    def policy_fun(observations):
+    def policy_fun(observations, rng=None):
       b, t = observations.shape[:2]
       a = env.action_space.n
       p = np.random.uniform(size=(b, t, a))
       p = np.exp(p)
       p = p / np.sum(p, axis=-1, keepdims=True)
-      return np.log(p)
+      return np.log(p), (), rng
 
     max_timestep = 15
     num_trajectories = 2
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 2852ef474..4a61a733b 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -168,7 +168,8 @@ def collect_trajectories(env,
                          policy="greedy",
                          max_timestep=None,
                          boundary=20,
-                         epsilon=0.1):
+                         epsilon=0.1,
+                         rng=None):
   """Collect trajectories with the given policy net and behaviour.
 
   Args:
@@ -182,6 +183,7 @@ def collect_trajectories(env,
       done.
     boundary: int, boundary for padding, used in EnvProblem envs.
     epsilon: float, the epsilon for `epsilon-greedy` policy.
+    rng: jax rng, splittable.
 
   Returns:
     trajectory: list of (observation, action, reward) tuples, where each element
@@ -198,7 +200,8 @@ def collect_trajectories(env,
         policy_fun,
         num_trajectories=num_trajectories,
         max_timestep=max_timestep,
-        boundary=boundary)
+        boundary=boundary,
+        rng=rng)
 
   trajectories = []
 
@@ -222,7 +225,7 @@ def collect_trajectories(env,
       ts_start = time.time()
       # Run the policy, to pick an action, shape is (1, t, A) because
       # observation_history is shaped (1, t) + OBS
-      predictions = policy_fun(observation_history)
+      predictions, _, rng = policy_fun(observation_history, rng=rng)
 
       # We need the predictions for the last time-step, so squeeze the batch
       # dimension and take the last time-step.
@@ -447,7 +450,8 @@ def value_loss(value_net_apply,
                reward_mask,
                gamma=0.99,
                epsilon=0.2,
-               value_prediction_old=None):
+               value_prediction_old=None,
+               rng=None):
   """Computes the value loss.
 
   Args:
@@ -462,6 +466,7 @@ def value_loss(value_net_apply,
     value_prediction_old: np.ndarray of shape (B, T+1, 1) of value predictions
         using the old parameters. If provided, we incorporate this in the loss
         as well. This is from the OpenAI baselines implementation.
+    rng: jax rng, splittable.
 
   Returns:
     The average L2 value loss, averaged over instances where reward_mask is 1.
@@ -471,7 +476,7 @@ def value_loss(value_net_apply,
   assert (B, T + 1) == observations.shape[:2]
 
   # NOTE: observations is (B, T+1) + OBS, value_prediction is (B, T+1, 1)
-  value_prediction = value_net_apply(observations, value_net_params)
+  value_prediction = value_net_apply(observations, value_net_params, rng=rng)
   assert (B, T + 1, 1) == value_prediction.shape
 
   return value_loss_given_predictions(value_prediction, rewards, reward_mask,
@@ -651,7 +656,8 @@ def ppo_loss(policy_net_apply,
              reward_mask,
              gamma=0.99,
              lambda_=0.95,
-             epsilon=0.2):
+             epsilon=0.2,
+             rng=None):
   """PPO objective, with an eventual minus sign, given observations."""
   B, T = padded_rewards.shape  # pylint: disable=invalid-name
   assert (B, T + 1) == padded_observations.shape[:2]
@@ -665,7 +671,8 @@ def ppo_loss(policy_net_apply,
 
   # log_probab_actions_{old,new} are both (B, T+1, A)
   log_probab_actions_new = policy_net_apply(padded_observations,
-                                            new_policy_params)
+                                            new_policy_params,
+                                            rng=rng)
   assert (B, T + 1) == log_probab_actions_new.shape[:2]
   assert log_probab_actions_old.shape[-1] == log_probab_actions_new.shape[-1]
 
@@ -776,10 +783,11 @@ def combined_loss(new_params,
                   lambda_=0.95,
                   epsilon=0.2,
                   c1=1.0,
-                  c2=0.01):
+                  c2=0.01,
+                  rng=None):
   """Computes the combined (clipped loss + value loss) given observations."""
   log_probab_actions_new, value_predictions_new = policy_and_value_net_apply(
-      padded_observations, new_params)
+      padded_observations, new_params, rng=rng)
 
   # (combined_loss, ppo_loss, value_loss, entropy_bonus)
   return combined_loss_given_predictions(log_probab_actions_new,
@@ -810,7 +818,8 @@ def ppo_opt_step(i,
                  reward_mask,
                  gamma=0.99,
                  lambda_=0.95,
-                 epsilon=0.1):
+                 epsilon=0.1,
+                 rng=None):
   """PPO optimizer step."""
   new_policy_params = ppo_get_params(opt_state)
   g = grad(
@@ -825,7 +834,8 @@ def ppo_opt_step(i,
           reward_mask,
           gamma=gamma,
           lambda_=lambda_,
-          epsilon=epsilon)
+          epsilon=epsilon,
+          rng=rng)
   return ppo_opt_update(i, g, opt_state)
 
 
@@ -838,7 +848,8 @@ def value_opt_step(i,
                    padded_observations,
                    padded_rewards,
                    reward_mask,
-                   gamma=0.99):
+                   gamma=0.99,
+                   rng=None):
   """Value optimizer step."""
   value_params = get_params(opt_state)
   # Note this partial application here and argnums above in ppo_opt_step.
@@ -847,7 +858,8 @@ def value_opt_step(i,
       padded_observations,
       padded_rewards,
       reward_mask,
-      gamma=gamma)
+      gamma=gamma,
+      rng=rng)
   return opt_update(i, g, opt_state)
 
 
@@ -867,7 +879,8 @@ def policy_and_value_opt_step(i,
                               c2=0.01,
                               gamma=0.99,
                               lambda_=0.95,
-                              epsilon=0.1):
+                              epsilon=0.1,
+                              rng=None):
   """Policy and Value optimizer step."""
 
   # Combined loss function given the new params.
@@ -886,7 +899,8 @@ def policy_and_value_loss(params):
         c2=c2,
         gamma=gamma,
         lambda_=lambda_,
-        epsilon=epsilon)
+        epsilon=epsilon,
+        rng=rng)
     return loss
 
   new_params = get_params(opt_state)
@@ -1027,26 +1041,34 @@ def training_loop(
       value_net_params = value_get_params(value_opt_state)
 
     # A function to get the policy and value predictions.
-    def get_predictions(observations):
+    def get_predictions(observations, rng=None):
+      """Returns log-probs, value predictions and key back."""
+      key, key1, key2 = jax_random.split(rng, num=3)
+
       if policy_net_apply is not None:
-        return (policy_net_apply(observations, policy_net_params),
-                value_net_apply(observations, value_net_params))
+        return (policy_net_apply(observations, policy_net_params, rng=key1),
+                value_net_apply(observations, value_net_params, rng=key2),
+                key)
 
       assert policy_and_value_net_apply
 
-      return policy_and_value_net_apply(
-          observations, policy_and_value_net_params)
+      log_probs, value_preds = policy_and_value_net_apply(
+          observations, policy_and_value_net_params, rng=key1)
+
+      return log_probs, value_preds, key
 
     t = time.time()
     t0 = t
     logging.vlog(1, "Epoch [% 6d] collecting trajectories.", i)
+    jax_rng_key, key = jax_random.split(jax_rng_key)
     trajs = collect_trajectories(
         env,
-        policy_fun=lambda observations: get_predictions(observations)[0],
+        policy_fun=get_predictions,
         num_trajectories=batch_size,
         policy=POLICY,
         max_timestep=max_timestep,
         boundary=boundary,
+        rng=key,
         epsilon=(10.0 / (i + 10.0)))  # this is a different epsilon.
 
     logging.vlog(1, "Collecting trajectories took %0.2f msec.", get_time(t))
@@ -1081,8 +1103,9 @@ def get_predictions(observations):
 
     # Calculate log-probabilities and value predictions of the trajectories.
     # We'll pass these to the loss functions so as to not get recomputed.
-    log_probabs_traj, value_predictions_traj = get_predictions(
-        padded_observations)
+    jax_rng_key, key = jax_random.split(jax_rng_key)
+    log_probabs_traj, value_predictions_traj, _ = get_predictions(
+        padded_observations, rng=key)
 
     # Some assertions.
     B, T = padded_actions.shape  # pylint: disable=invalid-name
@@ -1101,6 +1124,7 @@ def get_predictions(observations):
 
     # Compute value and ppo losses.
     cur_value_loss, cur_ppo_loss, cur_combined_loss = None, None, None
+    jax_rng_key, key1, key2 = jax_random.split(jax_rng_key, num=3)
     if policy_and_value_net_apply:
       logging.vlog(2, "Starting to compute P&V loss.")
       t = time.time()
@@ -1118,7 +1142,8 @@ def get_predictions(observations):
               lambda_=lambda_,
               epsilon=epsilon_schedule,
               c1=c1,
-              c2=c2))
+              c2=c2,
+              rng=key1))
       logging.vlog(
           1, "Calculating P&V loss [%10.2f(%10.2f, %10.2f)] took %0.2f msec.",
           cur_combined_loss, cur_value_loss, cur_ppo_loss, get_time(t))
@@ -1131,7 +1156,8 @@ def get_predictions(observations):
           padded_observations,
           padded_rewards,
           reward_mask,
-          gamma=gamma)
+          gamma=gamma,
+          rng=key1)
 
       logging.vlog(1, "Calculating value loss took %0.2f msec.", get_time(t))
 
@@ -1148,7 +1174,8 @@ def get_predictions(observations):
           reward_mask,
           gamma=gamma,
           lambda_=lambda_,
-          epsilon=epsilon_schedule)
+          epsilon=epsilon_schedule,
+          rng=key2)
       logging.vlog(1, "Calculating PPO loss took %0.2f msec.", get_time(t))
 
     value_losses.append(cur_value_loss)
@@ -1156,10 +1183,13 @@ def get_predictions(observations):
     if cur_combined_loss:
       combined_losses.append(cur_combined_loss)
 
+    jax_rng_key, key1, key2 = jax_random.split(jax_rng_key, num=3)
     if policy_and_value_net_apply:
       logging.vlog(1, "Policy and Value Optimization")
       t1 = time.time()
+      keys = jax_random.split(key1, num=num_optimizer_steps)
       for j in range(num_optimizer_steps):
+        k1, k2, k3 = jax_random.split(keys[j], num=3)
         t = time.time()
         # Update the optimizer state.
         policy_and_value_opt_state = policy_and_value_opt_step(
@@ -1178,14 +1208,15 @@ def get_predictions(observations):
             c2=c2,
             gamma=gamma,
             lambda_=lambda_,
-            epsilon=epsilon_schedule)
+            epsilon=epsilon_schedule,
+            rng=k1)
 
         # Compute the approx KL for early stopping.
         new_policy_and_value_net_params = policy_and_value_get_params(
             policy_and_value_opt_state)
 
         log_probab_actions_new, _ = policy_and_value_net_apply(
-            padded_observations, new_policy_and_value_net_params)
+            padded_observations, new_policy_and_value_net_params, rng=k2)
 
         approx_kl = approximate_kl(log_probab_actions_new,
                                    log_probabs_traj,
@@ -1218,7 +1249,8 @@ def get_predictions(observations):
                   lambda_=lambda_,
                   epsilon=epsilon_schedule,
                   c1=c1,
-                  c2=c2))
+                  c2=c2,
+                  rng=k3))
           logging.vlog(1, "One Policy and Value grad desc took: %0.2f msec",
                        get_time(t, t2))
           logging.vlog(1, "Combined Loss(value, ppo, entropy_bonus) [%10.2f] ->"
@@ -1242,8 +1274,9 @@ def get_predictions(observations):
       # Run optimizers.
       logging.vlog(1, "PPO Optimization")
       t1 = time.time()
-
+      keys1 = jax_random.split(key1, num=policy_only_num_optimizer_steps)
       for j in range(policy_only_num_optimizer_steps):
+        k1, k2, k3 = jax_random.split(keys1[j], num=3)
         t = time.time()
         # Update the optimizer state.
         ppo_opt_state = ppo_opt_step(
@@ -1261,13 +1294,15 @@ def get_predictions(observations):
             gamma=gamma,
             lambda_=lambda_,
             epsilon=epsilon_schedule,
+            rng=k1,
         )
         t2 = time.time()
         # Compute the approx KL for early stopping.
         # Get the new params.
         new_policy_net_params = ppo_get_params(ppo_opt_state)
         log_probab_actions_new = policy_net_apply(padded_observations,
-                                                  new_policy_net_params)
+                                                  new_policy_net_params,
+                                                  rng=k2)
         approx_kl = approximate_kl(log_probab_actions_new,
                                    log_probabs_traj,
                                    reward_mask)
@@ -1294,6 +1329,7 @@ def get_predictions(observations):
               gamma=gamma,
               lambda_=lambda_,
               epsilon=epsilon_schedule,
+              rng=k3,
           )
           logging.vlog(1, "One PPO grad desc took: %0.2f msec", get_time(t, t2))
           logging.vlog(1, "PPO loss [%10.2f] -> [%10.2f]", cur_ppo_loss,
@@ -1307,7 +1343,9 @@ def get_predictions(observations):
 
       logging.vlog(1, "Value Optimization")
 
+      keys2 = jax_random.split(key2, num=value_only_num_optimizer_steps)
       for j in range(value_only_num_optimizer_steps):
+        k1, k2, k3 = jax_random.split(keys2[j], num=3)
         t = time.time()
         value_opt_state = value_opt_step(
             j,
@@ -1318,7 +1356,8 @@ def get_predictions(observations):
             padded_observations,
             padded_rewards,
             reward_mask,
-            gamma=gamma)
+            gamma=gamma,
+            rng=k1)
         t2 = time.time()
         value_net_params = value_get_params(value_opt_state)
         if ((j + 1) %
@@ -1329,7 +1368,8 @@ def get_predictions(observations):
               padded_observations,
               padded_rewards,
               reward_mask,
-              gamma=gamma)
+              gamma=gamma,
+              rng=k2)
           logging.vlog(1, "One value grad desc took: %0.2f msec",
                        get_time(t, t2))
           logging.vlog(1, "Value loss [%10.2f] -> [%10.2f]", cur_value_loss,
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index bff2ea74c..e453210b3 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -83,11 +83,11 @@
 flags.DEFINE_integer("resized_width", 80, "Resized width of the game frame.")
 
 flags.DEFINE_boolean(
-    "combined_policy_and_value_function", False,
+    "combined_network", False,
     "If True there is a single network that determines policy"
     "and values.")
 
-flags.DEFINE_boolean("flatten_non_batch_time_dims", False,
+flags.DEFINE_boolean("flatten_dims", False,
                      "If true, we flatten except the first two dimensions.")
 
 # Number of optimizer steps of the combined net, policy net and value net.
@@ -114,7 +114,7 @@
 
 def common_layers():
   cur_layers = []
-  if FLAGS.flatten_non_batch_time_dims:
+  if FLAGS.flatten_dims:
     cur_layers = [layers.Div(divisor=255.0), layers.Flatten(num_axis_to_keep=2)]
   body = [layers.Dense(64), layers.Tanh(), layers.Dense(64), layers.Tanh()]
   return cur_layers + body
@@ -173,7 +173,7 @@ def run_training_loop():
     value_optimizer_fun = None
     policy_and_value_optimizer_fun = None
 
-    if FLAGS.combined_policy_and_value_function:
+    if FLAGS.combined_network:
       policy_and_value_net_fun = functools.partial(
           ppo.policy_and_value_net, bottom_layers_fn=common_layers)
       policy_and_value_optimizer_fun = get_optimizer_fun(FLAGS.learning_rate)
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index fe7a2a905..25243924a 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -96,10 +96,11 @@ def test_policy_and_value_net(self):
     self.assertEqual((batch, time_steps, 1), pnv_output[1].shape)
 
   def test_collect_trajectories(self):
+    self.rng_key, key1, key2, key3, key4 = jax_random.split(self.rng_key, num=5)
     observation_shape = (2, 3, 4)
     num_actions = 2
     policy_params, policy_apply = ppo.policy_net(
-        self.rng_key,
+        key1,
         (-1, -1) + observation_shape,
         num_actions,
         # flatten except batch and time
@@ -111,12 +112,17 @@ def test_collect_trajectories(self):
     env = fake_env.FakeEnv(
         observation_shape, num_actions, done_time_step=done_time_step)
 
+    def policy_fun(obs, rng=None):
+      rng, r = jax_random.split(rng)
+      return policy_apply(obs, policy_params, rng=r), (), rng
+
     num_trajectories = 5
     trajectories = ppo.collect_trajectories(
         env,
-        policy_fun=lambda obs: policy_apply(obs, policy_params),
+        policy_fun=policy_fun,
         num_trajectories=num_trajectories,
-        policy="categorical-sampling")
+        policy="categorical-sampling",
+        rng=key2)
 
     # Number of trajectories is as expected.
     self.assertEqual(num_trajectories, len(trajectories))
@@ -131,14 +137,20 @@ def test_collect_trajectories(self):
 
     # Test collect using a Policy and Value function.
     pnv_params, pnv_apply = ppo.policy_and_value_net(
-        self.rng_key, (-1, -1) + observation_shape, num_actions,
+        key3, (-1, -1) + observation_shape, num_actions,
         lambda: [layers.Flatten(num_axis_to_keep=2)])
 
+    def pnv_fun(obs, rng=None):
+      rng, r = jax_random.split(rng)
+      lp, v = pnv_apply(obs, pnv_params, rng=r)
+      return lp, v, rng
+
     trajectories = ppo.collect_trajectories(
         env,
-        policy_fun=lambda obs: pnv_apply(obs, pnv_params)[0],
+        policy_fun=pnv_fun,
         num_trajectories=num_trajectories,
-        policy="categorical-sampling")
+        policy="categorical-sampling",
+        rng=key4)
 
     # Number of trajectories is as expected.
     self.assertEqual(num_trajectories, len(trajectories))
@@ -152,12 +164,18 @@ def test_collect_trajectories(self):
       self.assertEqual((done_time_step + 1,), rewards.shape)
 
   def test_collect_trajectories_max_timestep(self):
+    self.rng_key, key1, key2 = jax_random.split(self.rng_key, num=3)
     observation_shape = (2, 3, 4)
     num_actions = 2
     pnv_params, pnv_apply = ppo.policy_and_value_net(
-        self.rng_key, (-1, -1) + observation_shape, num_actions,
+        key1, (-1, -1) + observation_shape, num_actions,
         lambda: [layers.Flatten(num_axis_to_keep=2)])
 
+    def pnv_fun(obs, rng=None):
+      rng, r = jax_random.split(rng)
+      lp, v = pnv_apply(obs, pnv_params, rng=r)
+      return lp, v, rng
+
     # We'll get done at time-step #5, starting from 0, therefore in 6 steps.
     done_time_step = 5
     env = fake_env.FakeEnv(
@@ -173,10 +191,11 @@ def test_collect_trajectories_max_timestep(self):
 
     trajectories = ppo.collect_trajectories(
         env,
-        policy_fun=lambda obs: pnv_apply(obs, pnv_params)[0],
+        policy_fun=pnv_fun,
         num_trajectories=num_trajectories,
         policy="categorical-sampling",
-        max_timestep=max_timestep)
+        max_timestep=max_timestep,
+        rng=key2)
 
     # Number of trajectories is as expected.
     self.assertEqual(num_trajectories, len(trajectories))
@@ -291,6 +310,8 @@ def test_rewards_to_go_really_long_sequences(self):
     self.assertAllClose(expected_r2g, actual_r2g)
 
   def test_value_loss(self):
+    self.rng_key, key = jax_random.split(self.rng_key, num=2)
+
     rewards = np.array([
         [1, 2, 4, 8, 16, 32, 64, 128],
         [1, 1, 1, 1, 1, 1, 1, 1],
@@ -309,8 +330,8 @@ def test_value_loss(self):
     observation_shape = (210, 160, 3)  # atari pong
     random_observations = np.random.uniform(size=(B, T + 1) + observation_shape)
 
-    def value_net_apply(observations, params):
-      del params
+    def value_net_apply(observations, params, rng=None):
+      del params, rng
       # pylint: disable=invalid-name
       B, T_p_1, OBS = (observations.shape[0], observations.shape[1],
                        observations.shape[2:])
@@ -324,7 +345,8 @@ def value_net_apply(observations, params):
           random_observations,
           rewards,
           rewards_mask,
-          gamma=gamma)
+          gamma=gamma,
+          rng=key)
 
     self.assertNear(53.3637084961, value_loss, 1e-6)
 

From 307ec34d9dfcb894f547ffe301ad48f5745c19d4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 14 May 2019 07:54:31 -0700
Subject: [PATCH 2028/2720] Small updates on Universal Transformer model:

PiperOrigin-RevId: 248135808
---
 .../models/research/universal_transformer.py  | 23 +++++-
 .../research/universal_transformer_util.py    | 70 ++++++++-----------
 2 files changed, 50 insertions(+), 43 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index ab93d54b6..185679cc0 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -41,7 +41,8 @@
 class UniversalTransformer(transformer.Transformer):
   """Universal Transformer: Depth-wise recurrent transformer model."""
 
-  def encode(self, inputs, target_space, hparams, features=None, losses=None):
+  def encode(self, inputs, target_space, hparams, features=None, losses=None,
+             **kwargs):
     """Encode Universal Transformer inputs.
 
     It is similar to "transformer.encode", but it uses
@@ -56,6 +57,7 @@ def encode(self, inputs, target_space, hparams, features=None, losses=None):
       features: optionally pass the entire features dictionary as well.
         This is needed now for "packed" datasets.
       losses: Unused.
+      **kwargs: additional arguments to pass to encoder_function
 
     Returns:
       Tuple of:
@@ -96,7 +98,8 @@ def decode(self,
              cache=None,
              decode_loop_step=None,
              nonpadding=None,
-             losses=None):
+             losses=None,
+             ** kwargs):
     """Decode Universal Transformer outputs from encoder representation.
 
     It is similar to "transformer.decode", but it uses
@@ -117,6 +120,7 @@ def decode(self,
       decode_loop_step: Unused.
       nonpadding: optional Tensor with shape [batch_size, decoder_length]
       losses: Unused.
+      **kwargs: additional arguments to pass to decoder_function
 
     Returns:
        Tuple of:
@@ -582,6 +586,14 @@ def adaptive_universal_transformer_tiny():
   return hparams
 
 
+@registry.register_hparams
+def adaptive_universal_transformer_sepconv_tiny():
+  hparams = universal_transformer_tiny()
+  hparams.recurrence_type = "act"
+  hparams.transformer_ffn_type = "sepconv"
+  return hparams
+
+
 @registry.register_hparams
 def adaptive_universal_transformer_global_base():
   hparams = universal_transformer_base()
@@ -784,6 +796,13 @@ def universal_transformer_sepconv_base():
   return hparams
 
 
+@registry.register_hparams
+def universal_transformer_sepconv_tiny():
+  hparams = universal_transformer_tiny()
+  hparams.transformer_ffn_type = "sepconv"
+  return hparams
+
+
 @registry.register_ranged_hparams
 def universal_transformer_base_range(rhp):
   """Range of hyperparameters."""
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 062e397a4..b7c6ae25c 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -37,7 +37,7 @@
 The recurrent transition function in fact controls how steps communicate with
 each other in depth. For instance, the recurrent transition, can be a simple
 identity function which passes the output of a step as the input to next step.
-Or it can be an LSTM (filliped vertically) next to the transformer which
+Or it can be an LSTM (flipped vertically) next to the transformer which
 controls how state of the model changes in depth.
 
 """
@@ -598,7 +598,7 @@ def universal_transformer_highway(layer_inputs,
   """Universal Transformer with highway connection.
 
 
-  It transforms the state using a block contaaining sel-attention and transition
+  It transforms the state using a block containing self-attention and transition
   function  and wrap the whole block with a highway connection.
   (the new state is a combination of the state and the transformed-state
   based on cary/transform gates.)
@@ -653,8 +653,7 @@ def universal_transformer_highway(layer_inputs,
       bias_initializer=tf.constant_initializer(hparams.transform_bias_init),
       activation=tf.sigmoid,
       pad_remover=pad_remover,
-      preprocess=True,
-      postprocess=True)
+      preprocess=True)
 
   if hparams.couple_carry_transform_gates:
     carry_gate = tf.subtract(1.0, transform_gate, name="carry")
@@ -668,8 +667,7 @@ def universal_transformer_highway(layer_inputs,
         bias_initializer=tf.constant_initializer(-hparams.transform_bias_init),
         activation=tf.sigmoid,
         pad_remover=pad_remover,
-        preprocess=True,
-        postprocess=True)
+        preprocess=True)
 
   new_state = state * carry_gate + transformed_state * transform_gate
 
@@ -747,8 +745,7 @@ def universal_transformer_skip(layer_inputs,
       bias_initializer=tf.constant_initializer(hparams.transform_bias_init),
       activation=tf.sigmoid,
       pad_remover=pad_remover,
-      preprocess=True,
-      postprocess=True)
+      preprocess=True)
 
   if hparams.couple_carry_transform_gates:
     carry_gate = tf.subtract(1.0, transform_gate, name="carry")
@@ -762,8 +759,7 @@ def universal_transformer_skip(layer_inputs,
         bias_initializer=tf.constant_initializer(-hparams.transform_bias_init),
         activation=tf.sigmoid,
         pad_remover=pad_remover,
-        preprocess=True,
-        postprocess=True)
+        preprocess=True)
 
   tf.contrib.summary.scalar("skip_transform_gate_layer",
                             tf.reduce_mean(transform_gate))
@@ -881,9 +877,7 @@ def universal_transformer_with_gru_as_transition_function(
         name="update",
         bias_initializer=tf.constant_initializer(1.0),
         activation=tf.sigmoid,
-        pad_remover=pad_remover,
-        preprocess=False,
-        postprocess=False)
+        pad_remover=pad_remover)
 
     tf.contrib.summary.scalar("gru_update_gate",
                               tf.reduce_mean(transition_function_update_gate))
@@ -895,9 +889,7 @@ def universal_transformer_with_gru_as_transition_function(
         name="reset",
         bias_initializer=tf.constant_initializer(1.0),
         activation=tf.sigmoid,
-        pad_remover=pad_remover,
-        preprocess=False,
-        postprocess=False)
+        pad_remover=pad_remover)
 
     tf.contrib.summary.scalar("gru_reset_gate",
                               tf.reduce_mean(transition_function_reset_gate))
@@ -910,9 +902,7 @@ def universal_transformer_with_gru_as_transition_function(
         name="candidate",
         bias_initializer=tf.zeros_initializer(),
         activation=tf.tanh,
-        pad_remover=pad_remover,
-        preprocess=False,
-        postprocess=False)
+        pad_remover=pad_remover)
 
     transition_function_output = (
         (1 - transition_function_update_gate) * transition_function_input +
@@ -975,9 +965,7 @@ def universal_transformer_with_lstm_as_transition_function(
         name="input",
         bias_initializer=tf.zeros_initializer(),
         activation=tf.sigmoid,
-        pad_remover=pad_remover,
-        preprocess=False,
-        postprocess=False)
+        pad_remover=pad_remover)
 
     tf.contrib.summary.scalar("lstm_input_gate",
                               tf.reduce_mean(transition_function_input_gate))
@@ -989,9 +977,7 @@ def universal_transformer_with_lstm_as_transition_function(
         name="forget",
         bias_initializer=tf.zeros_initializer(),
         activation=None,
-        pad_remover=pad_remover,
-        preprocess=False,
-        postprocess=False)
+        pad_remover=pad_remover)
     forget_bias_tensor = tf.constant(hparams.lstm_forget_bias)
     transition_function_forget_gate = tf.sigmoid(
         transition_function_forget_gate + forget_bias_tensor)
@@ -1006,9 +992,7 @@ def universal_transformer_with_lstm_as_transition_function(
         name="output",
         bias_initializer=tf.zeros_initializer(),
         activation=tf.sigmoid,
-        pad_remover=pad_remover,
-        preprocess=False,
-        postprocess=False)
+        pad_remover=pad_remover)
 
     tf.contrib.summary.scalar("lstm_output_gate",
                               tf.reduce_mean(transition_function_output_gate))
@@ -1020,9 +1004,7 @@ def universal_transformer_with_lstm_as_transition_function(
         name="input_modulation",
         bias_initializer=tf.zeros_initializer(),
         activation=tf.tanh,
-        pad_remover=pad_remover,
-        preprocess=False,
-        postprocess=False)
+        pad_remover=pad_remover)
 
     transition_function_memory = (
         memory * transition_function_forget_gate +
@@ -1214,6 +1196,7 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
 
 def _ffn_layer_multi_inputs(inputs_list,
                             hparams,
+                            output_size=None,
                             ffn_layer_type="dense",
                             name="ffn",
                             kernel_initializer=None,
@@ -1227,14 +1210,15 @@ def _ffn_layer_multi_inputs(inputs_list,
   Args:
     inputs_list: list of input tensors
     hparams: hyper-parameters
+    output_size: dimentionality of the output
     ffn_layer_type: dense / dense_dropconnect/ dense_relu_dense
     name: name
     kernel_initializer: kernel initializer
     bias_initializer: bias initializer
     activation: activation function
     pad_remover: pad remover
-    preprocess: if preprocess the input
-    postprocess: if postprocess the output
+    preprocess: if preprocess the input --> default: layer-norm
+    postprocess: if postprocess the output --> default: drop-out and residual
 
   Returns:
     a tensor
@@ -1247,10 +1231,14 @@ def _ffn_layer_multi_inputs(inputs_list,
   num_inputs = len(inputs_list)
   assert num_inputs > 0
 
-  if preprocess and num_inputs == 1:
-    inputs_list[0] = common_layers.layer_preprocess(inputs_list[0], hparams)
+  if preprocess:
+    # In case of having more than one input to the ffn,
+    # we just apply layer norm on them independently as preprocessing
+    for i, inputs in enumerate(inputs_list):
+      inputs_list[i] = common_layers.layer_preprocess(inputs_list[i], hparams)
 
-  if postprocess:
+  # for the residual connection
+  if postprocess and num_inputs == 1:
     original_inputs = inputs_list[0]
 
   # the output size is the hidden size of the main inputs
@@ -1280,7 +1268,7 @@ def remove_pads(x):
   if ffn_layer_type == "dense":
     output = common_layers.dense(
         ffn_inputs,
-        hparams.hidden_size,
+        hparams.hidden_size if output_size is None else output_size,
         name=name,
         activation=activation,
         use_bias=True,
@@ -1290,7 +1278,7 @@ def remove_pads(x):
   elif ffn_layer_type == "dense_dropconnect":
     output = common_layers.dense_dropconnect(
         ffn_inputs,
-        hparams.hidden_size,
+        hparams.hidden_size if output_size is None else output_size,
         name=name,
         dropconnect_dropout=hparams.dropconnect_dropout,
         output_activation=activation)
@@ -1300,7 +1288,7 @@ def remove_pads(x):
     output = common_layers.dense_relu_dense(
         ffn_inputs,
         hparams.filter_size,
-        hparams.hidden_size,
+        hparams.hidden_size if output_size is None else output_size,
         name=name,
         dropout=hparams.relu_dropout,
         output_activation=activation,
@@ -1317,11 +1305,11 @@ def remove_pads(x):
   if postprocess:
     if num_inputs == 1:
       output = common_layers.layer_postprocess(original_inputs, output, hparams)
-    else:  # only dropout (no residual)x
+    else:  # only dropout (no residual)
       hp = copy.copy(hparams)
       hp.layer_postprocess_sequence = hp.layer_postprocess_sequence.replace(
           "a", "")
-      output = common_layers.layer_postprocess(original_inputs, output, hp)
+      output = common_layers.layer_postprocess(None, output, hp)
 
   return output
 

From 1c7c10c60abc31308b40ae6c850e5c9e363dd4a9 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 14 May 2019 10:16:52 -0700
Subject: [PATCH 2029/2720] Fork JAX optimizers into Trax and make them
 objects.

PiperOrigin-RevId: 248160296
---
 .../trax/configs/resnet50_imagenet_8gb.gin    |   6 +-
 .../configs/resnet50_imagenet_8gb_testing.gin |   6 +-
 .../trax/configs/transformer_big_lm1b_8gb.gin |   2 +-
 .../trax/configs/wide_resnet_cifar10_8gb.gin  |   6 +-
 tensor2tensor/trax/optimizers.py              |  70 -----
 tensor2tensor/trax/optimizers/__init__.py     |  36 +++
 tensor2tensor/trax/optimizers/base.py         | 243 ++++++++++++++++++
 tensor2tensor/trax/rlax/ppo.py                |   6 +-
 tensor2tensor/trax/trax.py                    |  73 +++---
 tensor2tensor/trax/trax_test.py               |   6 +-
 10 files changed, 339 insertions(+), 115 deletions(-)
 delete mode 100644 tensor2tensor/trax/optimizers.py
 create mode 100644 tensor2tensor/trax/optimizers/__init__.py
 create mode 100644 tensor2tensor/trax/optimizers/base.py

diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
index 243ea5a39..22d99b9b6 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
@@ -22,9 +22,9 @@ EvalAdjustingSchedule.constant = 1.0
 MultifactorSchedule.factors = 'constant * linear_warmup'
 MultifactorSchedule.warmup_steps = 400
 
-# Parameters for momentum:
+# Parameters for Momentum:
 # ==============================================================================
-momentum.mass = 0.9
+Momentum.mass = 0.9
 
 
 # Parameters for Resnet50:
@@ -38,7 +38,7 @@ train.eval_frequency = 2000
 train.eval_steps = 20
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.Resnet50
-train.optimizer = @trax.optimizers.momentum
+train.optimizer = @trax.optimizers.Momentum
 train.train_steps = 1000000
 train.lr_schedule = @learning_rate.EvalAdjustingSchedule
 
diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
index cf14b3882..ad1c50cf0 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
@@ -22,9 +22,9 @@ EvalAdjustingSchedule.constant = 1.0
 MultifactorSchedule.factors = 'constant * linear_warmup'
 MultifactorSchedule.warmup_steps = 400
 
-# Parameters for momentum:
+# Parameters for Momentum:
 # ==============================================================================
-momentum.mass = 0.9
+Momentum.mass = 0.9
 
 
 # Parameters for Resnet50:
@@ -38,7 +38,7 @@ train.eval_frequency = 2000
 train.eval_steps = 20
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.Resnet50
-train.optimizer = @trax.optimizers.momentum
+train.optimizer = @trax.optimizers.Momentum
 train.train_steps = 100000
 train.lr_schedule = @learning_rate.EvalAdjustingSchedule
 
diff --git a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
index 760cf6b6b..6c88ca76a 100644
--- a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
@@ -37,7 +37,7 @@ train.eval_frequency = 1000
 train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
-train.optimizer = @trax.optimizers.sm3
+train.optimizer = @trax.optimizers.SM3
 train.run_debug_step = False
 train.train_steps = 500000
 
diff --git a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
index c2e8a605f..500931580 100644
--- a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
+++ b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
@@ -22,9 +22,9 @@ EvalAdjustingSchedule.constant = 1.0
 MultifactorSchedule.factors = 'constant * linear_warmup'
 MultifactorSchedule.warmup_steps = 400
 
-# Parameters for momentum:
+# Parameters for Momentum:
 # ==============================================================================
-momentum.mass = 0.9
+Momentum.mass = 0.9
 
 # Parameters for preprocess_fun:
 # ==============================================================================
@@ -42,6 +42,6 @@ train.eval_frequency = 2000
 train.eval_steps = 20
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.WideResnet
-train.optimizer = @trax.optimizers.momentum
+train.optimizer = @trax.optimizers.Momentum
 train.train_steps = 1000000
 train.lr_schedule = @learning_rate.EvalAdjustingSchedule
diff --git a/tensor2tensor/trax/optimizers.py b/tensor2tensor/trax/optimizers.py
deleted file mode 100644
index 77bbfbb58..000000000
--- a/tensor2tensor/trax/optimizers.py
+++ /dev/null
@@ -1,70 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""gin-configurable optimizers and learning rate functions."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gin
-
-import jax
-from jax.experimental import optimizers as opt
-import numpy as onp
-
-
-def opt_configure(*args, **kwargs):
-  kwargs["module"] = "trax.optimizers"
-  return gin.external_configurable(*args, **kwargs)
-
-# Optimizers
-sgd = opt_configure(opt.sgd)
-sm3 = opt_configure(opt.sm3)
-adam = opt_configure(opt.adam)
-momentum = opt_configure(opt.momentum)
-rmsprop = opt_configure(opt.rmsprop)
-
-# Learning rates
-constant = opt_configure(opt.constant)
-exponential_decay = opt_configure(opt.exponential_decay)
-inverse_time_decay = opt_configure(opt.inverse_time_decay)
-piecewise_constant = opt_configure(opt.piecewise_constant)
-
-
-# TODO(mattjj): upstream this to jax.experimental.optimizers.
-def parallelize(opt_maker):
-  """Transform an optimizer maker into a parallel one with replicated state."""
-  num_devices = jax.lib.xla_bridge.device_count()
-  replicate_array = lambda x: onp.broadcast_to(x, (num_devices,) + x.shape)
-  unreplicate_array = lambda x: x.mean(0)  # an alternative is just x[0]
-
-  def parallel_opt_maker(*args, **kwargs):  # pylint:disable=missing-docstring
-    init_fun, update_fun, get_params = opt_maker(*args, **kwargs)
-
-    def init_replicated(params):
-      opt_state = init_fun(params)
-      if num_devices > 1:
-        opt_state = jax.tree_util.tree_map(replicate_array, opt_state)
-      return opt_state
-
-    def get_params_unreplicated(opt_state):
-      if num_devices > 1:
-        opt_state = jax.tree_util.tree_map(unreplicate_array, opt_state)
-      params = get_params(opt_state)
-      return params
-
-    return init_replicated, update_fun, get_params, get_params_unreplicated
-  return parallel_opt_maker
diff --git a/tensor2tensor/trax/optimizers/__init__.py b/tensor2tensor/trax/optimizers/__init__.py
new file mode 100644
index 000000000..7e3fd7f8e
--- /dev/null
+++ b/tensor2tensor/trax/optimizers/__init__.py
@@ -0,0 +1,36 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Optimizers defined in trax."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gin
+
+from tensor2tensor.trax.optimizers import base
+
+
+def opt_configure(*args, **kwargs):
+  kwargs["module"] = "trax.optimizers"
+  return gin.external_configurable(*args, **kwargs)
+
+# Optimizers (using upper-case names).
+# pylint: disable=invalid-name
+SGD = opt_configure(base.SGD)
+Momentum = opt_configure(base.Momentum)
+RMSProp = opt_configure(base.RMSProp)
+Adam = opt_configure(base.Adam)
+SM3 = opt_configure(base.SM3)
diff --git a/tensor2tensor/trax/optimizers/base.py b/tensor2tensor/trax/optimizers/base.py
new file mode 100644
index 000000000..a1e9d6f0d
--- /dev/null
+++ b/tensor2tensor/trax/optimizers/base.py
@@ -0,0 +1,243 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Trax base optimizer class."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import jax.numpy as np
+from six.moves import reduce
+
+from tensor2tensor.trax.layers import base as layers
+
+
+def tree_flatten(tree):
+  """Flatten a tree into a list."""
+  if isinstance(tree, (list, tuple)):
+    # In python, sum of lists starting from [] is the concatenation.
+    return sum([tree_flatten(t) for t in tree], [])
+  if isinstance(tree, dict):
+    # Only use the values in case of a dictionary node.
+    return sum([tree_flatten(v) for v in tree.values()], [])
+  return [tree]
+
+
+def tree_unflatten(flat, tree):
+  """Unflatten a list into a tree given the tree shape as second argument.
+
+  Args:
+    flat: a flat list of elements to be assembled into a tree.
+    tree: a tree with the structure we want to have in the new tree.
+
+  Returns:
+    A pair (new_tree, rest_of_flat) where the new tree that has the structure
+    of tree but with leaves from flat, and the remaining elements of flat if
+    more were provided than the number of leaves of tree (useful for recursion).
+  """
+  if isinstance(tree, (list, tuple)):
+    new_tree, rest = [], flat
+    for t in tree:
+      new_t, rest = tree_unflatten(rest, t)
+      new_tree.append(new_t)
+    new_tree = tuple(new_tree) if isinstance(tree, tuple) else new_tree
+    return new_tree, rest
+  if isinstance(tree, dict):
+    new_tree, rest = {}, flat
+    for k in tree:
+      new_v, rest = tree_unflatten(rest, tree[k])
+      new_tree[k] = new_v
+    return new_tree, rest
+  return flat[0], flat[1:]
+
+
+class Optimizer(object):
+  """Optimizer object, base class. Maps per-parameter functions to trees."""
+
+  def __init__(self, step_size):
+    """Optimizers take the step size function (learning rate) as argument."""
+    if callable(step_size):
+      self._step_size = step_size
+    else:
+      self._step_size = lambda _: step_size
+
+  def init(self, x):
+    """Create optimizer slots for the parameter x."""
+    raise NotImplementedError
+
+  def update(self, i, g, x, s):
+    """Update the parameter x at step i with gradient g using state s."""
+    raise NotImplementedError
+
+  # End subclass interface.
+
+  def step_size(self, i):
+    return self._step_size(i)
+
+  def tree_init(self, x_tree):
+    return [self.init(x) for x in tree_flatten(x_tree)]
+
+  def tree_update(self, i, grad_tree, x_tree, opt_state):
+    grad_flat = tree_flatten(grad_tree)
+    x_flat = tree_flatten(x_tree)
+    updated_pairs = [self.update(i, g, x, s)
+                     for (g, x, s) in zip(grad_flat, x_flat, opt_state)]
+    new_x_flat, new_opt_state = zip(*updated_pairs)
+    new_x, _ = tree_unflatten(new_x_flat, x_tree)
+    return new_x, new_opt_state
+
+
+# Utilities.
+
+
+def l2_norm(tree):
+  """Compute the l2 norm of a pytree of arrays. Useful for weight decay."""
+  leaves = tree_flatten(tree)
+  return np.sqrt(sum(np.vdot(x, x) for x in leaves))
+
+
+def clip_grads(grad_tree, max_norm):
+  """Clip gradients stored as a pytree of arrays to maximum norm `max_norm`."""
+  norm = l2_norm(grad_tree)
+  normalize = lambda g: np.where(norm < max_norm, g, g * (max_norm / norm))
+  return layers.nested_map(normalize, grad_tree)
+
+
+# Optimizers.
+
+
+class SGD(Optimizer):
+  """Plain SGD optimizer."""
+
+  def init(self, x):
+    return None
+
+  def update(self, i, g, x, state):
+    del state
+    return x - self.step_size(i) * g, None
+
+
+class Momentum(Optimizer):
+  """Nestrov momentum optimizer."""
+
+  def __init__(self, step_size, mass=0.9):
+    """Initializer with a step size function and mass."""
+    super(Momentum, self).__init__(step_size)
+    self._mass = mass
+
+  def init(self, x):
+    return np.zeros_like(x)
+
+  def update(self, i, g, x, velocity):
+    new_velocity = self._mass * velocity - (1. - self._mass) * g
+    return x + self.step_size(i) * new_velocity, new_velocity
+
+
+class RMSProp(Optimizer):
+  """RMSProp optimizer."""
+
+  def __init__(self, step_size, gamma=0.9, eps=1e-8):
+    """Initializer with a step size function, gamma and epsilon."""
+    super(RMSProp, self).__init__(step_size)
+    self._gamma = gamma
+    self._epsilon = eps
+
+  def init(self, x):
+    return np.ones_like(x)
+
+  def update(self, i, g, x, avg_sq_grad):
+    avg_sq_grad = avg_sq_grad * self._gamma + g**2 * (1. - self._gamma)
+    x = x - self.step_size(i) * g / (np.sqrt(avg_sq_grad) + self._epsilon)
+    return x, avg_sq_grad
+
+
+class Adam(Optimizer):
+  """Adam optimizer."""
+
+  def __init__(self, step_size, b1=0.9, b2=0.999, eps=1e-8):
+    """Create the Adam optimizer.
+
+    Args:
+      step_size: a callable representing a step size schedule
+        that maps the iteration index to positive scalar.
+      b1: optional, a positive scalar value for beta_1, the exponential decay
+        rate for the first moment estimates (default 0.9).
+      b2: optional, a positive scalar value for beta_2, the exponential decay
+         rate for the second moment estimates (default 0.999).
+      eps: optional, a positive scalar value for epsilon, a small constant for
+        numerical stability (default 1e-8).
+    """
+    super(Adam, self).__init__(step_size)
+    self._b1 = b1
+    self._b2 = b2
+    self._eps = eps
+
+  def init(self, x):
+    m = np.zeros_like(x)
+    v = np.zeros_like(x)
+    return m, v
+
+  def update(self, i, g, x, state):
+    m, v = state
+    b1, b2, eps = self._b1, self._b2, self._eps
+    m = (1 - b1) * g + b1 * m  # First  moment estimate.
+    v = (1 - b2) * (g ** 2) + b2 * v  # Second moment estimate.
+    mhat = m / (1 - b1 ** (i + 1))  # Bias correction.
+    vhat = v / (1 - b2 ** (i + 1))
+    x = x - self.step_size(i) * mhat / (np.sqrt(vhat) + eps)
+    return x, (m, v)
+
+
+class SM3(Optimizer):
+  """SM3 optimizer."""
+
+  def __init__(self, step_size, momentum=0.9):
+    """Create the SM3 optimizer.
+
+    Memory-Efficient Adaptive Optimization for Large-Scale Learning.
+    https://arxiv.org/abs/1901.11150
+
+    Args:
+      step_size: a callable representing a step size schedule
+        that maps the iteration index to positive scalar.
+      momentum: optional, a positive scalar value for momentum
+    """
+    super(SM3, self).__init__(step_size)
+    self._momentum = momentum
+
+  def init(self, x):
+    vs = [np.zeros(sz, dtype=x.dtype) for sz in x.shape]
+    return (np.zeros_like(x), vs)
+
+  def update(self, i, g, x, state):
+    m, vs = state
+
+    def splice(seq, i, x):
+      lst = list(seq)
+      lst[i:i+1] = x
+      return lst
+
+    def broadcast_into(ndim, x, axis):
+      idx = splice([None] * ndim, axis, [slice(None)])
+      return x[tuple(idx)]
+
+    vs = [broadcast_into(g.ndim, v, i) for i, v in enumerate(vs)]
+    accum = reduce(np.minimum, vs) + g ** 2
+    accum_inv_sqrt = np.where(accum > 0, 1. / np.sqrt(accum), 0)
+    m = (1. - self._momentum) * (g * accum_inv_sqrt) + self._momentum * m
+    x = x - self._step_size(i) * m
+    vs = [accum.max(splice(range(x.ndim), j, [])) for j in range(x.ndim)]
+    return x, (m, vs)
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 4a61a733b..da1095597 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -138,8 +138,10 @@ def policy_and_value_net(rng_key,
 
 
 def optimizer_fun(net_params, step_size=1e-3):
-  opt_init, opt_update, get_params = trax_opt.adam(
-      step_size=step_size, b1=0.9, b2=0.999, eps=1e-08)
+  opt = trax_opt.Adam(step_size=step_size, b1=0.9, b2=0.999, eps=1e-08)
+  opt_init = lambda x: (x, opt.tree_init(x))
+  opt_update = lambda i, g, s: opt.tree_update(i, g, s[0], s[1])
+  get_params = lambda x: x[0]
   opt_state = opt_init(net_params)
   return opt_state, opt_update, get_params
 
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 8ee28955b..547611ab9 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""trax main training functions."""
+"""Trax main training functions."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -168,6 +168,15 @@ def save_state(state, output_dir, keep=False):
   log("Model saved to %s" % params_file, stdout=False)
 
 
+def _save_replicated(opt_state, step, history, num_devices, output_dir, keep):
+  """Save state but given a possibly replicated opt_state."""
+  if num_devices > 1:
+    unreplicate = lambda x: x.mean(0)
+    opt_state = layers.nested_map(opt_state, unreplicate)
+    save_state(State(params=opt_state, step=step, history=history),
+               output_dir, keep=keep)
+
+
 # Metrics to calculate and report.
 _METRICS = {
     "accuracy": accuracy,
@@ -311,15 +320,14 @@ def mapped_predict(x, params, rng):
   return predict
 
 
-def _jit_update_fun(predict_fun, loss_fun, optimizer, lr_fun, num_devices):
+def _jit_update_fun(predict_fun, loss_fun, optimizer, num_devices):
   """Get jit-ed update function for loss, optimizer, learning rate function."""
   if num_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
     def single_update(i, opt_state, batch, rng):
       rng, subrng = jax_random.split(rng[0])
-      _, opt_update, get_params = optimizer(lr_fun)
-      params = get_params(opt_state)
-      return opt_update(i, backend.grad(loss_fun)(
-          params, batch, predict_fun, rng), opt_state), [subrng]
+      params, opt_slots = opt_state
+      return optimizer.tree_update(i, backend.grad(loss_fun)(
+          params, batch, predict_fun, rng), params, opt_slots), [subrng]
     return backend.jit(single_update)
 
   @functools.partial(backend.pmap, axis_name="batch")
@@ -327,12 +335,11 @@ def mapped_update(i, opt_state, batch, rng):
     """This is a multi-device version of the update function above."""
     # We assume all tensors have the first dimension = num_devices.
     rng, subrng = jax_random.split(rng)
-    _, opt_update, get_params = optimizer(lr_fun)
-    params = get_params(opt_state)
+    params, opt_slots = opt_state
     grads = backend.grad(loss_fun)(params, batch, predict_fun, rng)
     grads = jax.tree_util.tree_map(
         lambda g: lax.psum(g, "batch"), grads)
-    return opt_update(i, grads, opt_state), subrng
+    return optimizer.tree_update(i, grads, params, opt_slots), subrng
 
   def update(i, opt_state, batch, rng):
     return mapped_update(numpy.repeat(i, num_devices), opt_state, batch, rng)
@@ -366,7 +373,7 @@ def train(output_dir,
           model=gin.REQUIRED,
           loss_fun=loss,
           inputs=trax_inputs.inputs,
-          optimizer=trax_opt.adam,
+          optimizer=trax_opt.Adam,
           lr_schedule=lr.MultifactorSchedule,
           train_steps=1000,
           save_steps=None,
@@ -386,8 +393,7 @@ def train(output_dir,
     loss_fun: callable with signature: params, trax.inputs.Inputs, model, rng
       -> loss.
     inputs: callable returning trax.inputs.Inputs.
-    optimizer: The optimizer as a callable taking a learning_rate callable and
-      returning 2 callables, opt_init and opt_update.
+    optimizer: The optimizer (see optimizers/base.py for signature).
     lr_schedule: A learning rate schedule as a function that takes history and
       returns a function from step to learning rate (a float).
     train_steps: int, total number of training steps.
@@ -425,8 +431,7 @@ def train(output_dir,
   state = restore_state(output_dir)
   history = state.history
   lr_fun = lr_schedule(history)
-  opt_init, _, get_rep_params, get_params = (
-      trax_opt.parallelize(optimizer)(lr_fun))
+  opt = optimizer(lr_fun)
   model_train = model(mode="train")
   model_predict_eval = model(mode="eval")
 
@@ -441,13 +446,19 @@ def train(output_dir,
         [tuple([-1] + list(shape)) for shape in inputs.input_shape])
   else:  # Otherwise just add [-1] to the input shape.
     model_input_shape = tuple([-1] + list(inputs.input_shape))
-  params = state.params or model_train.initialize(model_input_shape, init_rng)
-  opt_state = opt_init(params)
+  if state.params:
+    params = state.params[0]
+    opt_state = state.params
+  else:
+    params = model_train.initialize(model_input_shape, init_rng)
+    opt_state = (params, opt.tree_init(params))
+  if num_devices > 1:
+    replicate = lambda x: numpy.broadcast_to(x, (num_devices,) + x.shape)
+    opt_state = layers.nested_map(opt_state, replicate)
 
   # jit model_predict and update so they're fast
   jit_model_predict_eval = _jit_predict_fun(model_predict_eval, num_devices)
-  jit_update_fun = _jit_update_fun(
-      model_train, loss_fun, optimizer, lr_fun, num_devices)
+  jit_update_fun = _jit_update_fun(model_train, loss_fun, opt, num_devices)
 
   train_stream = inputs.train_stream()
   epoch_steps = [train_steps]  # Only training if eval_frequency is 0 or None.
@@ -478,10 +489,8 @@ def train(output_dir,
       step += 1
 
       if step in save_steps:
-        params = get_params(opt_state)
-        save_state(State(params=params, step=step, history=history),
-                   output_dir,
-                   keep=True)
+        _save_replicated(opt_state, step, history, num_devices,
+                         output_dir, True)
 
       # LR log
       if step == 1 or step % 10 == 0:
@@ -497,19 +506,21 @@ def train(output_dir,
                       epoch_steps / epoch_time, step=step)
 
     # Print number of parameters
-    params = get_params(opt_state)
     if step == 1:
-      sizes = layers.sizes(params)
+      sizes = layers.sizes(opt_state[0])
+      if num_devices > 1:
+        unreplicate = lambda x: x.mean(0)
+        single_params = layers.nested_map(opt_state[0], unreplicate)
+        sizes = layers.sizes(single_params)
       total_size = layers.nested_reduce(sizes, sum)
       step_log(step, "Total trainable parameters size: %d" % total_size)
 
     # Evaluate in parallel
-    replicated_params = get_rep_params(opt_state)
     evaluate_train_and_eval(
         step=step,
         inputs=inputs,
         predict_fun=functools.partial(jit_model_predict_eval,
-                                      params=replicated_params),
+                                      params=opt_state[0]),
         eval_steps=eval_steps,
         rng=rng,
         train_sw=train_sw,
@@ -518,6 +529,7 @@ def train(output_dir,
 
     # Save computation graph (single-device only for now).
     if save_graphs and step == 1 and num_devices == 1:
+      params = opt_state[0]
       # Dump computation graphs to files.
       forward_computation = jax.xla_computation(model_predict_eval)(
           next_train_batch[0], params=params, rng=rng)
@@ -534,7 +546,8 @@ def train(output_dir,
           f.write(backward_computation.GetHloDotGraph())
 
     # Save state
-    save_state(State(params=params, step=step, history=history), output_dir)
+    _save_replicated(opt_state, step, history, num_devices,
+                     output_dir, False)
 
     # Save Gin config
     # Gin only tracks the used parameters, so we save it after the first epoch.
@@ -545,12 +558,12 @@ def train(output_dir,
     old_lr_fun = lr_fun
     lr_fun = lr_schedule(history)
     if lr_fun != old_lr_fun:  # For performance, only jit if there is a change.
-      jit_update_fun = _jit_update_fun(
-          model_train, loss_fun, optimizer, lr_fun, num_devices)
+      opt = optimizer(lr_fun)
+      jit_update_fun = _jit_update_fun(model_train, loss_fun, opt, num_devices)
 
     # Flush summary writers
     train_sw.flush()
     eval_sw.flush()
 
   step_log(step, "Training done")
-  return State(params=params, step=step, history=history)
+  return State(params=opt_state, step=step, history=history)
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index eb2ec5e9a..88ed84354 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -90,7 +90,7 @@ def test_train_eval_predict(self):
 
       # Predict with final params
       inputs = inputs(1).train_stream()
-      model()(next(inputs)[0], state.params)
+      model()(next(inputs)[0], state.params[0])
 
   def test_train_eval_predict_sm3(self):
     with self.tmp_dir() as output_dir:
@@ -109,7 +109,7 @@ def test_train_eval_predict_sm3(self):
                          inputs=inputs,
                          train_steps=train_steps,
                          eval_steps=eval_steps,
-                         optimizer=trax_opt.sm3)
+                         optimizer=trax_opt.SM3)
 
       # Assert total train steps
       self.assertEqual(train_steps, state.step)
@@ -122,7 +122,7 @@ def test_train_eval_predict_sm3(self):
 
       # Predict with final params
       inputs = inputs(1).train_stream()
-      model()(next(inputs)[0], state.params)
+      model()(next(inputs)[0], state.params[0])
 
 
 if __name__ == "__main__":

From a792a4f38fac069646f48a4157569b24b2b6db94 Mon Sep 17 00:00:00 2001
From: Ben Goodrich <bgoodrich@google.com>
Date: Tue, 14 May 2019 14:12:23 -0700
Subject: [PATCH 2030/2720] Add placeholder directory for factual accuracy
 paper

PiperOrigin-RevId: 248208619
---
 tensor2tensor/data_generators/wikifact/README.md | 4 ++++
 1 file changed, 4 insertions(+)
 create mode 100644 tensor2tensor/data_generators/wikifact/README.md

diff --git a/tensor2tensor/data_generators/wikifact/README.md b/tensor2tensor/data_generators/wikifact/README.md
new file mode 100644
index 000000000..024ad72ac
--- /dev/null
+++ b/tensor2tensor/data_generators/wikifact/README.md
@@ -0,0 +1,4 @@
+# Assessing the Factual Accuracy of Generated Text
+
+This directory will contain the code and scripts to generate data and train
+models from the paper *Assessing the Factual Accuracy of Generated Text*.

From cec26dbd782ea7e4c07377e8d1f9391eb0c5a65c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 14 May 2019 18:06:53 -0700
Subject: [PATCH 2031/2720] Added Trax Demo colab notebook.

PiperOrigin-RevId: 248249635
---
 tensor2tensor/trax/README.md                  |   5 +
 .../trax/notebooks/trax_demo_iclr2019.ipynb   | 830 ++++++++++++++++++
 2 files changed, 835 insertions(+)
 create mode 100644 tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb

diff --git a/tensor2tensor/trax/README.md b/tensor2tensor/trax/README.md
index 71715aaf2..feff3f200 100644
--- a/tensor2tensor/trax/README.md
+++ b/tensor2tensor/trax/README.md
@@ -22,6 +22,11 @@ the best parts into core JAX.
 
 ### Examples
 
+#### Example Colab
+
+See our example constructing language models from scratch in a GPU-backed colab notebook at
+[Trax Demo](https://colab.sandbox.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb)
+
 #### MLP on MNIST
 
 
diff --git a/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb b/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb
new file mode 100644
index 000000000..f43396640
--- /dev/null
+++ b/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb
@@ -0,0 +1,830 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "name": "Trax Demo",
+      "version": "0.3.2",
+      "provenance": [],
+      "collapsed_sections": []
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "o4WGihMLneYq",
+        "colab_type": "text"
+      },
+      "source": [
+        "# Trax: Train Models in JAX\n",
+        "\n",
+        "[JAX](https://github.com/google/jax) allows you to write [numpy](https://www.numpy.org/) and run it fast on accelerators.\n",
+        "\n",
+        "This makes ML research more *fun* and *clear* so we made\n",
+        "* [Trax](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/trax): a library of models in JAX.\n",
+        "\n",
+        "In this demo we show how to:\n",
+        "* Train a Trax model on a toy copy problem.\n",
+        "* Decode from a pre-trained [Transformer](https://arxiv.org/abs/1706.03762) language model.\n",
+        "* Define [Transformer](https://arxiv.org/abs/1706.03762) from scratch in Trax.\n",
+        "* Do research in Trax: play with hard attention to see how it impacts training and results.\n",
+        "\n",
+        "We would like your feedback!\n",
+        "* What are the parts you like or dislike in JAX and Trax?\n",
+        "* Will you start doing your research in Trax? If not, why? What would change your mind?\n",
+        "* What should we focus on? Speed, cleanliness, memory use?\n",
+        "* If you cannot tell us in person, please add your feedback on [this github issue](https://github.com/tensorflow/tensor2tensor/issues/1478).\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8YQw0hySTVlK",
+        "colab_type": "text"
+      },
+      "source": [
+        "## Installs\n",
+        "\n",
+        "We install jax and trax and download a pretrained model and vocab file."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "vAWJVzYRnbDU",
+        "colab_type": "code",
+        "outputId": "6cdeff6f-3fc9-406f-feaf-fd1f8d9de775",
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 578
+        }
+      },
+      "source": [
+        "# Install JAX for GPU and Tensor2Tensor.\n",
+        "!pip install --upgrade -q https://storage.googleapis.com/jax-wheels/cuda100/jaxlib-0.1.14-cp36-none-linux_x86_64.whl\n",
+        "!pip install --upgrade -q jax==0.1.27\n",
+        "!pip install --upgrade -q tensor2tensor==1.13.4\n",
+        "# Grab language-model checkpoint and vocab file.\n",
+        "!rm -f model.pkl\n",
+        "!wget https://storage.googleapis.com/traxdemo/model.pkl\n",
+        "!wget https://storage.googleapis.com/traxdemo/vocab.lm1b.en.32768\n",
+        "# Show GPU type.\n",
+        "!nvidia-smi -L"
+      ],
+      "execution_count": 2,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "text": [
+            "\u001b[K     |████████████████████████████████| 44.6MB 1.2MB/s \n",
+            "\u001b[K     |████████████████████████████████| 174kB 3.5MB/s \n",
+            "\u001b[K     |████████████████████████████████| 61kB 24.4MB/s \n",
+            "\u001b[?25h  Building wheel for jax (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
+            "  Building wheel for opt-einsum (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
+            "\u001b[K     |████████████████████████████████| 1.4MB 3.4MB/s \n",
+            "\u001b[K     |████████████████████████████████| 686kB 45.8MB/s \n",
+            "\u001b[K     |████████████████████████████████| 143kB 40.2MB/s \n",
+            "\u001b[K     |████████████████████████████████| 296kB 32.6MB/s \n",
+            "\u001b[?25h  Building wheel for pypng (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
+            "--2019-05-14 22:57:21--  https://storage.googleapis.com/traxdemo/model.pkl\n",
+            "Resolving storage.googleapis.com (storage.googleapis.com)... 209.85.234.128, 2607:f8b0:4001:c12::80\n",
+            "Connecting to storage.googleapis.com (storage.googleapis.com)|209.85.234.128|:443... connected.\n",
+            "HTTP request sent, awaiting response... 200 OK\n",
+            "Length: 211170062 (201M) [application/octet-stream]\n",
+            "Saving to: ‘model.pkl’\n",
+            "\n",
+            "model.pkl           100%[===================>] 201.39M   101MB/s    in 2.0s    \n",
+            "\n",
+            "2019-05-14 22:57:23 (101 MB/s) - ‘model.pkl’ saved [211170062/211170062]\n",
+            "\n",
+            "--2019-05-14 22:57:23--  https://storage.googleapis.com/traxdemo/vocab.lm1b.en.32768\n",
+            "Resolving storage.googleapis.com (storage.googleapis.com)... 64.233.183.128, 2607:f8b0:4001:c07::80\n",
+            "Connecting to storage.googleapis.com (storage.googleapis.com)|64.233.183.128|:443... connected.\n",
+            "HTTP request sent, awaiting response... 200 OK\n",
+            "Length: 297760 (291K) [application/octet-stream]\n",
+            "Saving to: ‘vocab.lm1b.en.32768’\n",
+            "\n",
+            "vocab.lm1b.en.32768 100%[===================>] 290.78K  --.-KB/s    in 0.007s  \n",
+            "\n",
+            "2019-05-14 22:57:24 (40.8 MB/s) - ‘vocab.lm1b.en.32768’ saved [297760/297760]\n",
+            "\n",
+            "GPU 0: Tesla T4 (UUID: GPU-1959cc75-52ab-cf03-e5fa-36aee0d59bc5)\n"
+          ],
+          "name": "stdout"
+        }
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "vvFrqacVS6B6",
+        "colab_type": "text"
+      },
+      "source": [
+        "## Imports"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "dYq8J8uBn9ZC",
+        "colab_type": "code",
+        "outputId": "db8ca8de-164c-4355-8abb-a493e7f9f393",
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 136
+        }
+      },
+      "source": [
+        "from six.moves import cPickle\n",
+        "import os\n",
+        "import datetime\n",
+        "import random\n",
+        "\n",
+        "import numpy as onp\n",
+        "from matplotlib import pyplot as plt\n",
+        "\n",
+        "from jax.ops import index, index_update\n",
+        "\n",
+        "from tensor2tensor.trax import trax\n",
+        "from tensor2tensor.trax import layers as tl\n",
+        "from tensor2tensor.trax import inputs as trax_input\n",
+        "from tensor2tensor.trax import models as trax_models\n",
+        "from tensor2tensor.trax import optimizers as trax_optimizers\n",
+        "from tensor2tensor.trax import backend\n",
+        "from tensor2tensor.trax.backend import numpy as np\n",
+        "from tensor2tensor.trax.backend import random as trax_random"
+      ],
+      "execution_count": 3,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
+            "For more information, please see:\n",
+            "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
+            "  * https://github.com/tensorflow/addons\n",
+            "If you depend on functionality not listed there, please file an issue.\n",
+            "\n"
+          ],
+          "name": "stdout"
+        }
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zR6RVHx4lPzA",
+        "colab_type": "text"
+      },
+      "source": [
+        "# Toy Copy Problem\n",
+        "\n",
+        "Here we define batched random integer inputs for a trivial sequence-copy learning task."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "wGmWmpIslQYv",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "VOCAB_SIZE = 128\n",
+        "def toy_problem_inputs(num_devices, batch_size=64,\n",
+        "                       train_lengths=[10, 20], eval_lengths=[20]):\n",
+        "  \"\"\"Make Inputs for the toy problem of the language 0w0w for w in [1..127]*.\n",
+        "\n",
+        "  Args:\n",
+        "    num_devices: how many devices to build the inputs for (assert 1 for colab).\n",
+        "    batch_size: how large are the batches.\n",
+        "    train_lengths: lengths of w for training.\n",
+        "    eval_lengths: lengths of w for eval.\n",
+        "\n",
+        "  Returns:\n",
+        "    trax.inputs.Inputs\n",
+        "  \"\"\"\n",
+        "  assert num_devices == 1\n",
+        "  def random_minibatches(length_list):\n",
+        "    \"\"\"Generate a stream of random mini-batches.\"\"\"\n",
+        "    while True:\n",
+        "      length = random.choice(length_list)\n",
+        "      w = onp.random.randint(low=1, high=VOCAB_SIZE-1,\n",
+        "                            size=(batch_size, length // 2))\n",
+        "      zero = onp.zeros([batch_size, 1], onp.int32)\n",
+        "      x = onp.concatenate([zero, w, zero, w], axis=1)\n",
+        "      yield (x, x)  # In a language model input and output are the same.\n",
+        "\n",
+        "  return trax_input.Inputs(\n",
+        "      train_stream=lambda: random_minibatches(train_lengths),\n",
+        "      train_eval_stream=lambda: random_minibatches(train_lengths),\n",
+        "      eval_stream=lambda: random_minibatches(eval_lengths),\n",
+        "      input_shape=(None,))"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "eU0mpaf1lRky",
+        "colab_type": "code",
+        "outputId": "bf94086c-5d97-462b-b565-d4ba5f59b6c4",
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 51
+        }
+      },
+      "source": [
+        "inputs = toy_problem_inputs(1)\n",
+        "print(next(inputs.train_stream())[0][0])"
+      ],
+      "execution_count": 5,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "text": [
+            "[  0  68  91  99 107 115 113 111  17 102  48   0  68  91  99 107 115 113\n",
+            " 111  17 102  48]\n"
+          ],
+          "name": "stdout"
+        }
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KvNaSWu5g2Vm",
+        "colab_type": "text"
+      },
+      "source": [
+        "## Baseline Transformer on Toy Problem"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "AGDmtrgcl73M",
+        "colab_type": "code",
+        "outputId": "4c0f12e9-10ec-4e67-9f15-d2cc7084c083",
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 748
+        }
+      },
+      "source": [
+        "timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M\")\n",
+        "output_dir = os.path.expanduser(\"~/trax_lm_%s\" % timestamp)\n",
+        "def model(mode):\n",
+        "  return trax_models.TransformerLM(\n",
+        "      VOCAB_SIZE, feature_depth=128,\n",
+        "      feedforward_depth=256, num_layers=3,\n",
+        "      num_heads=4, mode=mode)\n",
+        "_ = trax.train(model=model,\n",
+        "               inputs=toy_problem_inputs,\n",
+        "               output_dir=output_dir,\n",
+        "               train_steps=3000,\n",
+        "               eval_steps=10,\n",
+        "               eval_frequency=1000)"
+      ],
+      "execution_count": 5,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "text": [
+            "Step      0: Starting training using 1 devices\n",
+            "\n",
+            "Step      1: Ran 1 train steps in 36.77 secs\n",
+            "Step      1: Total trainable parameters size: 692736\n",
+            "Step      1: Evaluation\n",
+            "Step      1: train           accuracy |  0.00616714\n",
+            "Step      1: train neg_log_perplexity | -5.06836748\n",
+            "Step      1: train               loss |  5.06836748\n",
+            "Step      1: eval            accuracy |  0.00610795\n",
+            "Step      1: eval  neg_log_perplexity | -5.20451212\n",
+            "Step      1: eval                loss |  5.20451212\n",
+            "Step      1: Finished evaluation\n",
+            "\n",
+            "Step   1000: Ran 999 train steps in 89.13 secs\n",
+            "Step   1000: Evaluation\n",
+            "Step   1000: train           accuracy |  0.45719695\n",
+            "Step   1000: train neg_log_perplexity | -2.71764731\n",
+            "Step   1000: train               loss |  2.71764731\n",
+            "Step   1000: eval            accuracy |  0.41278410\n",
+            "Step   1000: eval  neg_log_perplexity | -2.94052887\n",
+            "Step   1000: eval                loss |  2.94052887\n",
+            "Step   1000: Finished evaluation\n",
+            "\n",
+            "Step   2000: Ran 1000 train steps in 15.61 secs\n",
+            "Step   2000: Evaluation\n",
+            "Step   2000: train           accuracy |  0.43169984\n",
+            "Step   2000: train neg_log_perplexity | -2.82782769\n",
+            "Step   2000: train               loss |  2.82782769\n",
+            "Step   2000: eval            accuracy |  0.41278410\n",
+            "Step   2000: eval  neg_log_perplexity | -2.92255998\n",
+            "Step   2000: eval                loss |  2.92255998\n",
+            "Step   2000: Finished evaluation\n",
+            "\n",
+            "Step   3000: Ran 1000 train steps in 15.64 secs\n",
+            "Step   3000: Evaluation\n",
+            "Step   3000: train           accuracy |  0.45053267\n",
+            "Step   3000: train neg_log_perplexity | -2.73254609\n",
+            "Step   3000: train               loss |  2.73254609\n",
+            "Step   3000: eval            accuracy |  0.41249999\n",
+            "Step   3000: eval  neg_log_perplexity | -2.92720962\n",
+            "Step   3000: eval                loss |  2.92720962\n",
+            "Step   3000: Finished evaluation\n",
+            "Step   3000: Training done\n"
+          ],
+          "name": "stdout"
+        }
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "eapBBkRUuho7",
+        "colab_type": "text"
+      },
+      "source": [
+        "# Decoding from a Pre-Trained Transformer Language Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "H6hVQ3v5iC00",
+        "colab_type": "code",
+        "outputId": "812949cc-4294-4a42-f55a-c40f65e151f8",
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 187
+        }
+      },
+      "source": [
+        "# load model checkpoint\n",
+        "with open(\"model.pkl\", \"rb\") as f:\n",
+        "   (params, step, history) = cPickle.load(f, encoding=\"latin1\")\n",
+        "\n",
+        "# lm1b subword vocab\n",
+        "def clean(x):\n",
+        "  return x[1:-2]\n",
+        "with open(\"vocab.lm1b.en.32768\", \"r\") as fp:\n",
+        "  vocab = list(map(clean, fp.readlines()))\n",
+        "vocab_map = {v:idx for idx,v in enumerate(vocab)}\n",
+        "\n",
+        "list(enumerate(vocab))[:10]"
+      ],
+      "execution_count": 6,
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "[(0, '<pad>_'),\n",
+              " (1, '<EOS>_'),\n",
+              " (2, 'the_'),\n",
+              " (3, ' , _'),\n",
+              " (4, ' ._'),\n",
+              " (5, 'to_'),\n",
+              " (6, 'of_'),\n",
+              " (7, 'a_'),\n",
+              " (8, 'and_'),\n",
+              " (9, 'in_')]"
+            ]
+          },
+          "metadata": {
+            "tags": []
+          },
+          "execution_count": 6
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "W-7s9RXQNIru",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "tlm = trax_models.TransformerLM(\n",
+        "  dropout=0.1, \n",
+        "  feature_depth=512, \n",
+        "  feedforward_depth=2048, \n",
+        "  max_len=2048, \n",
+        "  mode='eval', \n",
+        "  num_heads=8, \n",
+        "  num_layers=6, \n",
+        "  vocab_size=32000)"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "iLdtplDpdTMr",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "def gumbel_sample(v, temperature=0.8):\n",
+        "  u = onp.random.uniform(low=1e-9, high=1.0, size=v.shape)\n",
+        "  g = -onp.log(-onp.log(u))\n",
+        "  return np.argmax(v + g * temperature)"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "IHSbtHzPjW6i",
+        "colab_type": "code",
+        "outputId": "7a8306b7-6c6b-41ba-c5aa-c1a76d9d8037",
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 102
+        }
+      },
+      "source": [
+        "prompt = \"Please_\"\n",
+        "num_samples = 5\n",
+        "max_length = 20\n",
+        "for _ in range(num_samples):\n",
+        "  enc = [vocab_map[w] for w in str.split(prompt)]\n",
+        "  pos = len(enc)\n",
+        "  rng = trax_random.get_prng(0)\n",
+        "  data = np.zeros((1, 50), dtype=np.int32)\n",
+        "  data = index_update(data, index[1, 0:pos], enc)\n",
+        "\n",
+        "  while pos < max_length:\n",
+        "    tmp = tlm(data, params=params, rng=rng)\n",
+        "    next_sym = gumbel_sample(tmp[0, pos])\n",
+        "    data = index_update(data, index[1, pos], next_sym)\n",
+        "    pos += 1\n",
+        "    if int(next_sym) == 1:\n",
+        "      break\n",
+        "\n",
+        "  print(\"\".join([vocab[idx] for idx in onp.array(data)[0, 0:pos]]))"
+      ],
+      "execution_count": 10,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "text": [
+            "Please_write_to_him_to_tell_him_about_the_Wallace_and_Gromit_films_. _and_to_give_him_this_\n",
+            "Please_do_not_turn_to_making_sure_your_children_are_already_in_school_or_that_you_have_school_ ._<EOS>_\n",
+            "Please_read_the_full_prospectus_to_see_if_the_proposed_transaction_may_be_accurate_ ._<EOS>_\n",
+            "Please_note_that_the_new_policy_has_been_strengthened_by_the_fact_that_Britney_Spears_ ' _mother_ , _Janet_Jackson_\n",
+            "Please_ , _please_aim_at_your_brother_ , _if_you_want_to_ ._<EOS>_\n"
+          ],
+          "name": "stdout"
+        }
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Ym8otS7HpUIO",
+        "colab_type": "text"
+      },
+      "source": [
+        "# Transformer from Scratch\n",
+        "\n",
+        "Here we re-implement multiheaded self-attention and a transformer language model from scratch using only a few simple linear primitives from trax.\n",
+        "\n",
+        "Note in particular the commented modifications in the core  __DotProductAttention__ function as an example of how easy it is to modify layers and models for research using Trax."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "uw-GIdm2p_4X",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "def DotProductAttention(query, key, value, mask, dropout, mode, rng, hard_k=4):\n",
+        "  \"\"\"Core dot product self-attention.\n",
+        "  Args:\n",
+        "    query: array of representations\n",
+        "    key: array of representations\n",
+        "    value: array of representations\n",
+        "    mask: attention-mask, gates attention\n",
+        "    dropout: float: dropout rate\n",
+        "    mode: 'eval' or 'train': whether to use dropout\n",
+        "    rng: JAX PRNGKey: subkey for disposable use\n",
+        "  Returns:\n",
+        "    Self attention for q, k, v arrays.\n",
+        "  \"\"\"\n",
+        "  depth = np.shape(query)[-1]\n",
+        "  dots = np.matmul(query, np.swapaxes(key, -1, -2)) / np.sqrt(depth)\n",
+        "  if mask is not None:\n",
+        "    dots = np.where(mask, dots, -1e9)\n",
+        "  # Softmax.\n",
+        "  dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))\n",
+        "  # ----------------------------------------------------------------------\n",
+        "  # As an example of a simple research modification, we modify the typical \n",
+        "  # dot-product attention mechanism with top-k \"hard attention\":\n",
+        "  # ----------------------------------------------------------------------\n",
+        "  if hard_k > 0:\n",
+        "    top_k = np.sort(dots)[..., -hard_k]  # Get the top-kth weight.\n",
+        "    dots -= top_k[..., np.newaxis]  # Subtract (be 0 for lower ones).\n",
+        "    dots = np.maximum(dots, 0)\n",
+        "    dots /= np.sum(dots, axis=-1, keepdims=True)  # Re-normalize.\n",
+        "  # ----------------------------------------------------------------------\n",
+        "  if dropout >= 1.0:\n",
+        "    raise ValueError('Dropout rates must be lower than 1.')\n",
+        "  if dropout is not None and dropout > 0.0 and mode == 'train':\n",
+        "    keep = backend.random.bernoulli(rng, 1.0 - dropout, dots.shape)\n",
+        "    dots = np.where(keep, dots / (1.0 - dropout), 0)\n",
+        "  out = np.matmul(dots, value)\n",
+        "  # Uncomment to see an example TRAX stack trace to this point:\n",
+        "  # ----------------------------------------------------------------------\n",
+        "  # raise ValueError(\"err\")\n",
+        "  # ----------------------------------------------------------------------\n",
+        "  return out\n",
+        "\n",
+        "\n",
+        "def _multihead_attention_output_shape(  # pylint: disable=invalid-name\n",
+        "    input_shapes, **unused_kwargs):\n",
+        "  \"\"\"Helper: calculate multihead attention output shape.\"\"\"\n",
+        "  q_shape = input_shapes[0][0]  # Inputs are ((q, k, v), mask).\n",
+        "  mask_shape = input_shapes[1]\n",
+        "  return q_shape, mask_shape\n",
+        "\n",
+        "\n",
+        "@tl.layer(output_shape=_multihead_attention_output_shape)\n",
+        "def PureMultiHeadedAttention(x, params, num_heads=8, dropout=0.0,\n",
+        "                             mode='train', **kwargs):\n",
+        "  \"\"\"Pure transformer-style multi-headed attention.\n",
+        "  Args:\n",
+        "    x: inputs ((q, k, v), mask)\n",
+        "    params: parameters (none)\n",
+        "    num_heads: int: number of attention heads\n",
+        "    dropout: float: dropout rate\n",
+        "    mode: str: 'train' or 'eval'\n",
+        "    **kwargs: other arguments including the rng\n",
+        "  Returns:\n",
+        "    Pure Multi-headed attention result, and the mask.\n",
+        "  \"\"\"\n",
+        "  del params\n",
+        "  rng = kwargs.get('rng', None)\n",
+        "  (q, k, v), mask = x\n",
+        "  feature_depth = q.shape[-1]\n",
+        "  assert feature_depth % num_heads == 0\n",
+        "  head_depth = feature_depth // num_heads\n",
+        "  nbatch = np.shape(q)[0]\n",
+        "  # nbatch, seqlen, feature_depth --> nbatch, num_heads, seqlen, head_depth\n",
+        "  def SplitHeads(x):\n",
+        "    return np.transpose(\n",
+        "        np.reshape(x, (nbatch, -1, num_heads, head_depth)), (0, 2, 1, 3))\n",
+        "  # nbatch, num_heads, seqlen, head_depth --> nbatch, seqlen, feature_depth\n",
+        "  def JoinHeads(x):  # pylint: disable=invalid-name\n",
+        "    return np.reshape(\n",
+        "        np.transpose(x, (0, 2, 1, 3)), (nbatch, -1, num_heads*head_depth))\n",
+        "  # Split heads, dot-product attention, rejoin heads.\n",
+        "  res = JoinHeads(\n",
+        "      DotProductAttention(\n",
+        "          SplitHeads(q), SplitHeads(k), SplitHeads(v), mask,\n",
+        "          dropout=dropout, mode=mode, rng=rng))\n",
+        "  return res, mask  # Keep the mask.\n",
+        "\n",
+        "\n",
+        "def MultiHeadedAttentionQKV(\n",
+        "    feature_depth, num_heads=8, dropout=0.0, mode='train'):\n",
+        "  \"\"\"Transformer-style multi-headed attention.\n",
+        "  Accepts inputs of the form (q, k, v), mask.\n",
+        "  Args:\n",
+        "    feature_depth: int:  depth of embedding\n",
+        "    num_heads: int: number of attention heads\n",
+        "    dropout: float: dropout rate\n",
+        "    mode: str: 'train' or 'eval'\n",
+        "  Returns:\n",
+        "    Multi-headed self-attention result and the mask.\n",
+        "  \"\"\"\n",
+        "  return tl.Serial(\n",
+        "      tl.Parallel(\n",
+        "          tl.Parallel(\n",
+        "              tl.Dense(feature_depth),\n",
+        "              tl.Dense(feature_depth),\n",
+        "              tl.Dense(feature_depth),\n",
+        "          ),\n",
+        "          tl.Copy()\n",
+        "      ),\n",
+        "      PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter\n",
+        "          feature_depth=feature_depth, num_heads=num_heads,\n",
+        "          dropout=dropout, mode=mode),\n",
+        "      tl.Parallel(tl.Dense(feature_depth), tl.Copy())\n",
+        "  )\n",
+        "\n",
+        "\n",
+        "def MultiHeadedAttention(\n",
+        "    feature_depth, num_heads=8, dropout=0.0, mode='train'):\n",
+        "  \"\"\"Transformer-style multi-headed attention.\n",
+        "  Accepts inputs of the form (x, mask) and constructs (q, k, v) from x.\n",
+        "  Args:\n",
+        "    feature_depth: int:  depth of embedding\n",
+        "    num_heads: int: number of attention heads\n",
+        "    dropout: float: dropout rate\n",
+        "    mode: str: 'train' or 'eval'\n",
+        "  Returns:\n",
+        "    Multi-headed self-attention layer.\n",
+        "  \"\"\"\n",
+        "  return tl.Serial(\n",
+        "      tl.Parallel(\n",
+        "          # q = k = v = first input\n",
+        "          tl.Branch(\n",
+        "              tl.Copy(), tl.Copy(), tl.Copy()),\n",
+        "          tl.Copy()  # pass the mask\n",
+        "      ),\n",
+        "      MultiHeadedAttentionQKV(  # pylint: disable=no-value-for-parameter\n",
+        "          feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),\n",
+        "  )"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "Ge42t7VZl-d2",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "def ResidualFeedForward(feature_depth,\n",
+        "                        feedforward_depth,\n",
+        "                        dropout,\n",
+        "                        mode):\n",
+        "  \"\"\"Residual feed-forward layer with normalization at start.\"\"\"\n",
+        "  return tl.Residual(\n",
+        "      tl.LayerNorm(),\n",
+        "      tl.Dense(feedforward_depth),\n",
+        "      tl.Relu(),\n",
+        "      tl.Dropout(rate=dropout, mode=mode),\n",
+        "      tl.Dense(feature_depth),\n",
+        "      tl.Dropout(rate=dropout, mode=mode)\n",
+        "  )\n",
+        "\n",
+        "\n",
+        "def DecoderLayer(feature_depth,\n",
+        "                 feedforward_depth,\n",
+        "                 num_heads,\n",
+        "                 dropout,\n",
+        "                 mode):\n",
+        "  \"\"\"Transformer decoder layer.\n",
+        "  Args:\n",
+        "    feature_depth: int:  depth of embedding\n",
+        "    feedforward_depth: int: depth of feed-forward layer\n",
+        "    num_heads: int: number of attention heads\n",
+        "    dropout: float: dropout rate (how much to drop out)\n",
+        "    mode: str: 'train' or 'eval'\n",
+        "  Returns:\n",
+        "    the layer.\n",
+        "  \"\"\"\n",
+        "  return tl.Serial(\n",
+        "      tl.Residual(  # Self-attention block.\n",
+        "          tl.LayerNorm(),\n",
+        "          tl.Branch(tl.Copy(), tl.CausalMask(axis=-2)),  # Create mask.\n",
+        "          # We replace the \"stock\" self-attention layer with the one defined\n",
+        "          # above:\n",
+        "          # tl.MultiHeadedAttention(feature_depth, num_heads=num_heads,\n",
+        "          #                         dropout=dropout, mode=mode),\n",
+        "          MultiHeadedAttention(feature_depth, num_heads=num_heads,\n",
+        "                                  dropout=dropout, mode=mode),\n",
+        "          tl.Select(0),  # Drop the mask.\n",
+        "          tl.Dropout(rate=dropout, mode=mode)\n",
+        "      ),\n",
+        "      ResidualFeedForward(feature_depth, feedforward_depth, dropout, mode=mode)\n",
+        "  )\n",
+        "\n",
+        "\n",
+        "def TransformerLM(vocab_size,\n",
+        "                  feature_depth=512,\n",
+        "                  feedforward_depth=2048,\n",
+        "                  num_layers=6,\n",
+        "                  num_heads=8,\n",
+        "                  dropout=0.1,\n",
+        "                  max_len=2048,\n",
+        "                  mode='train'):\n",
+        "  \"\"\"Transformer language model (only uses the decoder part of Transformer).\n",
+        "  Args:\n",
+        "    vocab_size: int: vocab size\n",
+        "    feature_depth: int:  depth of embedding\n",
+        "    feedforward_depth: int: depth of feed-forward layer\n",
+        "    num_layers: int: number of encoder/decoder layers\n",
+        "    num_heads: int: number of attention heads\n",
+        "    dropout: float: dropout rate (how much to drop out)\n",
+        "    max_len: int: maximum symbol length for positional encoding\n",
+        "    mode: str: 'train' or 'eval'\n",
+        "  Returns:\n",
+        "    the layer.\n",
+        "  \"\"\"\n",
+        "  return tl.Serial(\n",
+        "      tl.ShiftRight(),\n",
+        "      tl.Embedding(feature_depth, vocab_size),\n",
+        "      tl.Dropout(rate=dropout, mode=mode),\n",
+        "      tl.PositionalEncoding(max_len=max_len),\n",
+        "      tl.Serial(*[DecoderLayer(feature_depth, feedforward_depth, num_heads,\n",
+        "                               dropout, mode)\n",
+        "                  for _ in range(num_layers)]),\n",
+        "      tl.LayerNorm(),\n",
+        "      tl.Dense(vocab_size),\n",
+        "      tl.LogSoftmax()\n",
+        "  )"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "WZxnwjAEqYDh",
+        "colab_type": "code",
+        "outputId": "f90e965d-2625-4e56-9038-65c087639051",
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 748
+        }
+      },
+      "source": [
+        "timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M\")\n",
+        "output_dir = os.path.expanduser(\"~/trax_lm_%s\" % timestamp)\n",
+        "def new_model(mode):\n",
+        "  return TransformerLM(\n",
+        "      VOCAB_SIZE, feature_depth=128,\n",
+        "      feedforward_depth=256, num_layers=3,\n",
+        "      num_heads=4, mode=mode)\n",
+        "_ = trax.train(model=new_model,\n",
+        "           inputs=toy_problem_inputs,\n",
+        "           output_dir=output_dir,\n",
+        "           train_steps=3000,\n",
+        "           eval_steps=10,\n",
+        "           eval_frequency=1000)"
+      ],
+      "execution_count": 22,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "text": [
+            "Step      0: Starting training using 1 devices\n",
+            "\n",
+            "Step      1: Ran 1 train steps in 42.29 secs\n",
+            "Step      1: Total trainable parameters size: 692736\n",
+            "Step      1: Evaluation\n",
+            "Step      1: train           accuracy |  0.00686553\n",
+            "Step      1: train neg_log_perplexity | -5.42891455\n",
+            "Step      1: train               loss |  5.42891455\n",
+            "Step      1: eval            accuracy |  0.00809659\n",
+            "Step      1: eval  neg_log_perplexity | -5.39403439\n",
+            "Step      1: eval                loss |  5.39403439\n",
+            "Step      1: Finished evaluation\n",
+            "\n",
+            "Step   1000: Ran 999 train steps in 109.64 secs\n",
+            "Step   1000: Evaluation\n",
+            "Step   1000: train           accuracy |  0.12875238\n",
+            "Step   1000: train neg_log_perplexity | -4.29979420\n",
+            "Step   1000: train               loss |  4.29979420\n",
+            "Step   1000: eval            accuracy |  0.09928977\n",
+            "Step   1000: eval  neg_log_perplexity | -4.45948172\n",
+            "Step   1000: eval                loss |  4.45948172\n",
+            "Step   1000: Finished evaluation\n",
+            "\n",
+            "Step   2000: Ran 1000 train steps in 16.89 secs\n",
+            "Step   2000: Evaluation\n",
+            "Step   2000: train           accuracy |  0.53104877\n",
+            "Step   2000: train neg_log_perplexity | -2.33383632\n",
+            "Step   2000: train               loss |  2.33383632\n",
+            "Step   2000: eval            accuracy |  0.54900569\n",
+            "Step   2000: eval  neg_log_perplexity | -2.24813342\n",
+            "Step   2000: eval                loss |  2.24813342\n",
+            "Step   2000: Finished evaluation\n",
+            "\n",
+            "Step   3000: Ran 1000 train steps in 16.91 secs\n",
+            "Step   3000: Evaluation\n",
+            "Step   3000: train           accuracy |  0.56715208\n",
+            "Step   3000: train neg_log_perplexity | -2.15219927\n",
+            "Step   3000: train               loss |  2.15219927\n",
+            "Step   3000: eval            accuracy |  0.54928976\n",
+            "Step   3000: eval  neg_log_perplexity | -2.25436211\n",
+            "Step   3000: eval                loss |  2.25436211\n",
+            "Step   3000: Finished evaluation\n",
+            "Step   3000: Training done\n"
+          ],
+          "name": "stdout"
+        }
+      ]
+    }
+  ]
+}
\ No newline at end of file

From 2e88e8e92f01bc7b140a6ee67a451493e68e8226 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 14 May 2019 19:40:44 -0700
Subject: [PATCH 2032/2720] Fix Trax Demo Notebook

PiperOrigin-RevId: 248259172
---
 .../trax/notebooks/trax_demo_iclr2019.ipynb   | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb b/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb
index f43396640..ddb354a4e 100644
--- a/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb
+++ b/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb
@@ -15,6 +15,30 @@
     "accelerator": "GPU"
   },
   "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ySEmBgmqMSIJ",
+        "colab_type": "text"
+      },
+      "source": [
+        "##### Copyright 2019 Google LLC.\n",
+        "\n",
+        "Licensed under the Apache License, Version 2.0 (the \"License\");\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",
+        "https://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."
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {

From 96ae8a0c5337efe46aab9ad2a27b33bd81cc8f57 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 15 May 2019 09:52:17 -0700
Subject: [PATCH 2033/2720] Internal.

PiperOrigin-RevId: 248350800
---
 tensor2tensor/trax/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/README.md b/tensor2tensor/trax/README.md
index feff3f200..459b9b916 100644
--- a/tensor2tensor/trax/README.md
+++ b/tensor2tensor/trax/README.md
@@ -25,7 +25,7 @@ the best parts into core JAX.
 #### Example Colab
 
 See our example constructing language models from scratch in a GPU-backed colab notebook at
-[Trax Demo](https://colab.sandbox.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb)
+[Trax Demo](https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb)
 
 #### MLP on MNIST
 

From 9d4b8c79646cac3ccd468dd1dc0815d45b99b0ce Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 15 May 2019 10:17:11 -0700
Subject: [PATCH 2034/2720] Correct SM3 implementation and make it the default.

PiperOrigin-RevId: 248355543
---
 .../trax/configs/transformer_lm1b_8gb.gin     |  5 +-
 .../configs/transformer_lm1b_8gb_testing.gin  |  5 +-
 tensor2tensor/trax/learning_rate.py           |  4 +-
 tensor2tensor/trax/optimizers/base.py         | 67 +++++++++++++------
 tensor2tensor/trax/trax.py                    |  2 +-
 5 files changed, 57 insertions(+), 26 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index 89aa4bd48..3cf306276 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -21,8 +21,8 @@ masked_mean.mask_id = 0
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-MultifactorSchedule.constant = 0.1
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.constant = 0.3
+MultifactorSchedule.factors = 'constant * linear_warmup'
 MultifactorSchedule.warmup_steps = 8000
 
 # Parameters for preprocess_fun:
@@ -37,6 +37,7 @@ train.eval_frequency = 1000
 train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
+train.optimizer = @trax.optimizers.SM3
 train.run_debug_step = False
 train.train_steps = 500000
 
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
index 56415758a..4650208f2 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
@@ -21,8 +21,8 @@ masked_mean.mask_id = 0
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-MultifactorSchedule.constant = 0.1
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.constant = 0.3
+MultifactorSchedule.factors = 'constant * linear_warmup'
 MultifactorSchedule.warmup_steps = 8000
 
 # Parameters for preprocess_fun:
@@ -37,6 +37,7 @@ train.eval_frequency = 1000
 train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
+train.optimizer = @trax.optimizers.SM3
 train.run_debug_step = False
 train.train_steps = 100000
 
diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index 9146ec176..4363eddf8 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -40,9 +40,9 @@
 
 @gin.configurable(blacklist=["history"])
 def MultifactorSchedule(history=None,
-                        factors="constant * linear_warmup * rsqrt_decay",
+                        factors="constant * linear_warmup",
                         constant=0.1,
-                        warmup_steps=100,
+                        warmup_steps=400,
                         decay_factor=0.5,
                         steps_per_decay=20000):
   """Factor-based learning rate schedule.
diff --git a/tensor2tensor/trax/optimizers/base.py b/tensor2tensor/trax/optimizers/base.py
index a1e9d6f0d..8aea5f386 100644
--- a/tensor2tensor/trax/optimizers/base.py
+++ b/tensor2tensor/trax/optimizers/base.py
@@ -20,7 +20,6 @@
 from __future__ import print_function
 
 import jax.numpy as np
-from six.moves import reduce
 
 from tensor2tensor.trax.layers import base as layers
 
@@ -222,22 +221,52 @@ def init(self, x):
     vs = [np.zeros(sz, dtype=x.dtype) for sz in x.shape]
     return (np.zeros_like(x), vs)
 
+  def _update_diagonal(self, step, g, x, m, v):
+    v[0] += g * g
+    preconditioner = np.where(v[0] > 0, 1.0 / np.sqrt(v[0]),
+                              np.zeros_like(v[0]))
+    preconditioned_g = preconditioner * g
+    m = (1 - self._momentum) * preconditioned_g + self._momentum * m
+    x = x - self.step_size(step) * m
+    return x, (m, v)
+
+  def _expanded_shape(self, shape, axis):
+    # Replaces a `shape` of [M, N, K] with 1 in all dimensions except for i.
+    # For eg: i = 1 returns [1, N, 1].
+    rank = len(shape)
+    return [1] * axis + [shape[axis]] + [1] * (rank - axis - 1)
+
+  def _minimum(self, tensor_list):
+    minimum = tensor_list[0]
+    for i in range(1, len(tensor_list)):
+      minimum = np.minimum(minimum, tensor_list[i])
+    return minimum
+
+  def _update_sketched(self, step, g, x, m, v):
+    """Update for higher-rank parameters."""
+    shape = x.shape
+    rank = len(shape)
+    reshaped_accumulators = [np.reshape(v[i], self._expanded_shape(shape, i))
+                             for i in range(rank)]
+    current_accumulator = self._minimum(reshaped_accumulators)
+    current_accumulator += g * g
+    accumulator_inv_sqrt = np.where(current_accumulator > 0.0,
+                                    1.0 / np.sqrt(current_accumulator),
+                                    np.zeros_like(current_accumulator))
+    preconditioned_gradient = g * accumulator_inv_sqrt
+    m = (1.0 - self._momentum) * preconditioned_gradient + self._momentum * m
+    x = x - self.step_size(step) * m
+    for i in range(len(v)):
+      axes = list(range(int(i))) + list(range(int(i) + 1, rank))
+      dim_accumulator = np.amax(current_accumulator, axis=axes)
+      v[i] = dim_accumulator
+    return x, (m, v)
+
   def update(self, i, g, x, state):
-    m, vs = state
-
-    def splice(seq, i, x):
-      lst = list(seq)
-      lst[i:i+1] = x
-      return lst
-
-    def broadcast_into(ndim, x, axis):
-      idx = splice([None] * ndim, axis, [slice(None)])
-      return x[tuple(idx)]
-
-    vs = [broadcast_into(g.ndim, v, i) for i, v in enumerate(vs)]
-    accum = reduce(np.minimum, vs) + g ** 2
-    accum_inv_sqrt = np.where(accum > 0, 1. / np.sqrt(accum), 0)
-    m = (1. - self._momentum) * (g * accum_inv_sqrt) + self._momentum * m
-    x = x - self._step_size(i) * m
-    vs = [accum.max(splice(range(x.ndim), j, [])) for j in range(x.ndim)]
-    return x, (m, vs)
+    m, v = state
+    shape = x.shape
+    rank = len(shape)
+    if rank > 1:
+      return self._update_sketched(i, g, x, m, v)
+    else:
+      return self._update_diagonal(i, g, x, m, v)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 547611ab9..0b8c93dd6 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -373,7 +373,7 @@ def train(output_dir,
           model=gin.REQUIRED,
           loss_fun=loss,
           inputs=trax_inputs.inputs,
-          optimizer=trax_opt.Adam,
+          optimizer=trax_opt.SM3,
           lr_schedule=lr.MultifactorSchedule,
           train_steps=1000,
           save_steps=None,

From 4ce366131ce69d1005f035e14677609f7dfdb580 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 15 May 2019 13:39:22 -0700
Subject: [PATCH 2035/2720] Open sourcing the search space used in "The Evolved
 Transformer." arXiv preprint arXiv:1901.11117 (2019).

PiperOrigin-RevId: 248395650
---
 .../neural_architecture_search/__init__.py    |   16 +
 .../neural_architecture_search/nas_layers.py  |  686 +++++++++++
 .../nas_layers_test.py                        |  320 +++++
 .../neural_architecture_search/nas_model.py   | 1030 +++++++++++++++++
 .../nas_model_test.py                         |  468 ++++++++
 5 files changed, 2520 insertions(+)
 create mode 100644 tensor2tensor/models/neural_architecture_search/__init__.py
 create mode 100644 tensor2tensor/models/neural_architecture_search/nas_layers.py
 create mode 100644 tensor2tensor/models/neural_architecture_search/nas_layers_test.py
 create mode 100644 tensor2tensor/models/neural_architecture_search/nas_model.py
 create mode 100644 tensor2tensor/models/neural_architecture_search/nas_model_test.py

diff --git a/tensor2tensor/models/neural_architecture_search/__init__.py b/tensor2tensor/models/neural_architecture_search/__init__.py
new file mode 100644
index 000000000..b775a72bd
--- /dev/null
+++ b/tensor2tensor/models/neural_architecture_search/__init__.py
@@ -0,0 +1,16 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+
diff --git a/tensor2tensor/models/neural_architecture_search/nas_layers.py b/tensor2tensor/models/neural_architecture_search/nas_layers.py
new file mode 100644
index 000000000..c4c9839a1
--- /dev/null
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers.py
@@ -0,0 +1,686 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Bank of layers for Translation NAS searches.
+
+All encoder layers are registered in the global LayerRegistry ENCODER_LAYERS.
+All decoder layers are registered on the global LayerRegistry DECODER_LAYERS.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import abc
+import six
+
+from tensor2tensor.layers import common_attention
+
+import tensorflow as tf
+
+# Registry layer keys.
+ATTEND_TO_ENCODER_REGISTRY_KEY = "attend_to_encoder"
+ATTENTION_32_HEADS_REGISTRY_KEY = "attention_32_heads"
+ATTENTION_16_HEADS_REGISTRY_KEY = "attention_16_heads"
+ATTENTION_4_HEADS_REGISTRY_KEY = "attention_4_heads"
+DEPTHWISE_CONV_3X1_REGISTRY_KEY = "depthwise_conv_3x1"
+DEPTHWISE_CONV_5X1_REGISTRY_KEY = "depthwise_conv_5x1"
+DEPTHWISE_CONV_7X1_REGISTRY_KEY = "depthwise_conv_7x1"
+DILATED_CONV_3X1_REGISTRY_KEY = "dilated_conv_3x1"
+DILATED_CONV_5X1_REGISTRY_KEY = "dilated_conv_5x1"
+GATED_LINEAR_UNIT_REGISTRY_KEY = "gated_linear_unit"
+IDENTITY_REGISTRY_KEY = "identity"
+# Lightweight convolution naming convention uses "R_X" where X is the variable
+# reduction factor.
+LIGHTWEIGHT_CONV_3X1_R_1_REGISTRY_KEY = "lightweight_conv_3x1_r_1"
+LIGHTWEIGHT_CONV_3X1_R_4_REGISTRY_KEY = "lightweight_conv_3x1_r_4"
+LIGHTWEIGHT_CONV_3X1_R_16_REGISTRY_KEY = "lightweight_conv_3x1_r_16"
+LIGHTWEIGHT_CONV_5X1_R_1_REGISTRY_KEY = "lightweight_conv_5x1_r_1"
+LIGHTWEIGHT_CONV_5X1_R_4_REGISTRY_KEY = "lightweight_conv_5x1_r_4"
+LIGHTWEIGHT_CONV_5X1_R_16_REGISTRY_KEY = "lightweight_conv_5x1_r_16"
+LIGHTWEIGHT_CONV_7X1_R_1_REGISTRY_KEY = "lightweight_conv_7x1_r_1"
+LIGHTWEIGHT_CONV_7X1_R_4_REGISTRY_KEY = "lightweight_conv_7x1_r_4"
+LIGHTWEIGHT_CONV_7X1_R_16_REGISTRY_KEY = "lightweight_conv_7x1_r_16"
+LIGHTWEIGHT_CONV_15X1_R_1_REGISTRY_KEY = "lightweight_conv_15x1_r_1"
+LIGHTWEIGHT_CONV_15X1_R_4_REGISTRY_KEY = "lightweight_conv_15x1_r_4"
+LIGHTWEIGHT_CONV_15X1_R_16_REGISTRY_KEY = "lightweight_conv_15x1_r_16"
+SEPARABLE_CONV_3X1_REGISTRY_KEY = "separable_conv_3x1"
+SEPARABLE_CONV_5X1_REGISTRY_KEY = "separable_conv_5x1"
+SEPARABLE_CONV_7X1_REGISTRY_KEY = "separable_conv_7x1"
+SEPARABLE_CONV_9X1_REGISTRY_KEY = "separable_conv_9x1"
+SEPARABLE_CONV_11X1_REGISTRY_KEY = "separable_conv_11x1"
+SEPARABLE_CONV_13X1_REGISTRY_KEY = "separable_conv_13x1"
+SEPARABLE_CONV_15X1_REGISTRY_KEY = "separable_conv_15x1"
+STANDARD_CONV_1X1_REGISTRY_KEY = "standard_conv_1x1"
+STANDARD_CONV_3X1_REGISTRY_KEY = "standard_conv_3x1"
+STANDARD_CONV_5X1_REGISTRY_KEY = "standard_conv_5x1"
+STANDARD_ATTENTION_REGISTRY_KEY = "standard_attention"
+
+
+class TranslationLayer(object):
+  """Interface for the layers used in the Translation search space."""
+
+  __metaclass__ = abc.ABCMeta
+
+  @abc.abstractmethod
+  def _apply_logic(self, input_tensor, output_depth, hparams, var_scope_suffix,
+                   nonpadding, mask_future, **kwargs):
+    """Applies the layer specific logic to the `input_tensor`.
+
+    This is called by `apply_layer()` to apply the subclass specific logic to
+    the preprocessed `input_tensor`.
+
+    Args:
+      input_tensor: [batch_size, batch time_steps, embedding_depth] tensor.
+      output_depth: Depth of the output tensor.
+      hparams: Hyperparameters for the layer.
+      var_scope_suffix: Suffix appended to the end of the variable scope.
+      nonpadding: a [batch_size, batch time_steps] tensor with 1 where each
+        batch member has sequence information and 0 everywhere else. This is
+        used to mask out the irrelevant padded portions of the input.
+      mask_future: Boolean. If False, information moves across the
+        spatial/temporal dimension freely. If True, each timestep can only
+        process the information that has come before it.
+      **kwargs: Subclass-specific arguments.
+
+    Returns:
+      logic_output: [batch_size, batch time_steps, output_depth] tensor output
+                    of the logic.
+    """
+
+  def apply_layer(self,
+                  input_tensor,
+                  residual_tensor,
+                  output_depth,
+                  activation,
+                  hparams,
+                  var_scope_suffix,
+                  nonpadding,
+                  mask_future,
+                  layer_preprocess_fn=None,
+                  postprocess_dropout=True,
+                  **kwargs):
+    """Applies the layer to the input.
+
+    Also applies pad masking, preprocessing, postprocessing, and nonlinearity.
+
+    Args:
+      input_tensor: [batch_size, batch time_steps, embedding_depth] tensor.
+      residual_tensor: Tensor that gets added to the output residually if
+        `layer_postprocess` is True.
+      output_depth: Depth of the output tensor.
+      activation: Activation to be applied to the `layer_output`. If None, no
+        activation will be applied.
+      hparams: Hyperparameters for the layer.
+      var_scope_suffix: Suffix appended to the end of the variable scope.
+      nonpadding: a [batch_size, batch time_steps] tensor with 1 where each
+        batch member has sequence information and 0 everywhere else. This is
+        used to mask out the irrelevant padded portions of the input.
+      mask_future: Boolean. If False, information moves across the
+        spatial/temporal dimension freely. If True, each timestep can only
+        process the information that has come before it.
+      layer_preprocess_fn: Preprocess function applied to the input.
+      postprocess_dropout: Whether or not to apply dropout.
+      **kwargs: Arguments used by specific TranslationLayers.
+
+    Returns:
+      layer_output: The output of the layer.
+    """
+    input_depth = input_tensor.shape.as_list()[-1]
+    layer_output = input_tensor
+    if nonpadding is not None:
+      nonpadding_input_tiled = tf.tile(
+          tf.expand_dims(nonpadding, 2), [1, 1, input_depth])
+      layer_output *= nonpadding_input_tiled
+
+    if layer_preprocess_fn:
+      layer_output = layer_preprocess_fn(layer_output)
+      if nonpadding is not None:
+        layer_output *= nonpadding_input_tiled
+
+    layer_output = self._apply_logic(layer_output, output_depth, hparams,
+                                     var_scope_suffix, nonpadding, mask_future,
+                                     **kwargs)
+
+    if activation:
+      layer_output = activation(layer_output)
+
+    if postprocess_dropout:
+      layer_output = tf.nn.dropout(layer_output, 1 - hparams.relu_dropout)
+
+    if residual_tensor is not None:
+      layer_output += residual_tensor
+
+    # Remove the output padding items.
+    if nonpadding is not None:
+      nonpadding_output_tiled = tf.tile(
+          tf.expand_dims(nonpadding, 2), [1, 1, output_depth])
+      layer_output *= nonpadding_output_tiled
+
+    return layer_output
+
+  @abc.abstractmethod
+  def num_params(self, input_depth, output_depth, **kwargs):
+    """Returns num_params in the layer for the given input and output depths.
+
+    NOTE: This does not include layer norm parameters that appear in
+      layer_preprocess or layer_postprocess!
+
+    Args:
+      input_depth: The depth of the input.
+      output_depth: The depth of the output.
+      **kwargs: TranslationLayer specific arguments.
+    """
+
+
+class LayerRegisteredError(Exception):
+  """Layer name is already used in LayerRegistry."""
+
+
+class LayerRegistry(object):
+  """Registry of TranslationLayers.
+
+  The registry is a mapping of string names to TranslationLayers. Layers can be
+  added to the registry via `registry_layer()` and can be accessed via `get()`.
+  """
+
+  def __init__(self):
+    self._layers = {}
+
+  def register_layer(self, name, translation_layer):
+    """Register a TranslationLayer under the key `name`."""
+    if name in self._layers and self._layers[name] != translation_layer:
+      raise LayerRegisteredError(
+          "Already registered %s in layer registry with a different object!" %
+          name)
+
+    self._layers[name] = translation_layer
+
+  def get(self, name):
+    return self._layers[name]
+
+  def get_layer_names(self):
+    return sorted(six.iterkeys(self._layers))
+
+
+DECODER_LAYERS = LayerRegistry()
+ENCODER_LAYERS = LayerRegistry()
+
+
+class ConvLayerBase(TranslationLayer):
+  """Convolution TranslationLayer base class."""
+
+  def __init__(self, conv_type, conv_width, dilation_rate):
+    self._conv_type = conv_type
+    self._conv_width = conv_width
+    self._dilation_rate = dilation_rate
+
+  def _conv_function(self, input_tensor, output_depth, padding):
+    """Conv function that will be applied to the input tensor."""
+    raise NotImplementedError()
+
+  def _apply_logic(self, input_tensor, output_depth, hparams, var_scope_suffix,
+                   nonpadding, mask_future, **unused_kwargs):
+    """Applies conv logic to `input_tensor`."""
+    with tf.variable_scope("%s_conv_%s" % (self._conv_type, var_scope_suffix)):
+      if mask_future:
+        # Pad shift the inputs so that temporal information does not leak. This
+        # must be used in tandem with VALID padding.
+        pad_amount = int(self._conv_width - 1) * self._dilation_rate
+        logic_output = tf.pad(
+            input_tensor, paddings=[[0, 0], [pad_amount, 0], [0, 0]])
+        padding = "VALID"
+      else:
+        logic_output = input_tensor
+        padding = "SAME"
+
+      logic_output = tf.expand_dims(logic_output, 2)
+      logic_output = self._conv_function(logic_output, output_depth, padding)
+
+      logic_output = tf.squeeze(logic_output, 2)
+    return logic_output
+
+
+class SeparableConvLayer(ConvLayerBase):
+  """Separable convolution TranslationLayer base class."""
+
+  def __init__(self, conv_width):
+    super(SeparableConvLayer, self).__init__("separable", conv_width, 1)
+
+  def _conv_function(self, input_tensor, output_depth, padding):
+    conv_output = tf.squeeze(input_tensor, 2)
+    separable_conv_1d = tf.layers.SeparableConv1D(
+        output_depth,
+        self._conv_width,
+        padding=padding,
+        name="separable_conv_%sx1" % self._conv_width)
+    conv_output = separable_conv_1d.apply(conv_output)
+    return tf.expand_dims(conv_output, 2)
+
+  def num_params(self, input_depth, output_depth, **unused_kwargs):
+    return (self._conv_width * input_depth + input_depth * output_depth +
+            output_depth)
+
+
+class StandardConvLayer(ConvLayerBase):
+  """Standard convolutional TranslationLayer base class."""
+
+  def __init__(self, conv_width):
+    super(StandardConvLayer, self).__init__("standard", conv_width, 1)
+
+  def _conv_function(self, input_tensor, output_depth, padding):
+    return tf.layers.conv2d(
+        input_tensor,
+        output_depth, [self._conv_width, 1],
+        padding=padding,
+        name="conv_%sx1" % self._conv_width)
+
+  def num_params(self, input_depth, output_depth, **unused_kwargs):
+    return self._conv_width * input_depth * output_depth + output_depth
+
+
+def calculate_depthwise_channel_multiplier(input_depth, output_depth):
+  """Calculates channel multiplier for depthwise convolution."""
+  # Check to see if the output_depth >= input_depth
+  # and output_depth % input_depth == 0. If this is the case then we
+  # can satify the output_depth constraint, so the channel multiplier
+  # will be set accordingly.
+  if ((output_depth >= input_depth) and (output_depth % input_depth == 0)):
+    return output_depth / input_depth
+  return 1
+
+
+class DepthwiseConvLayer(ConvLayerBase):
+  """Depthwise convolution TranslationLayer base class."""
+
+  def __init__(self, conv_width):
+    super(DepthwiseConvLayer, self).__init__("depthwise", conv_width, 1)
+
+  def _conv_function(self, input_tensor, output_depth, padding):
+    input_depth = input_tensor.shape.as_list()[-1]
+    if not ((output_depth >= input_depth) and
+            (output_depth % input_depth == 0)):
+      raise ValueError(
+          "Depthwise layer output_depth (%s) must be greater or equal to and "
+          "a multiple of the depth of the "
+          "input tensor (%s)." % (output_depth, input_depth))
+    channel_multiplier = calculate_depthwise_channel_multiplier(
+        input_depth, output_depth)
+    kernel = tf.get_variable(
+        "kernel", [self._conv_width, 1, input_depth, channel_multiplier])
+    return tf.nn.depthwise_conv2d(
+        input_tensor,
+        kernel, [1, 1, 1, 1],
+        padding=padding,
+        name="depthwise_conv_%sx1" % str(self._conv_width))
+
+  def num_params(self, input_depth, output_depth, **unused_kwargs):
+    channel_multiplier = calculate_depthwise_channel_multiplier(
+        input_depth, output_depth)
+    return self._conv_width * input_depth * channel_multiplier
+
+
+class LightweightConvLayer(ConvLayerBase):
+  """Lightweight convolution TranslationLayer base class."""
+
+  def __init__(self, conv_width, num_repeat):
+    super(LightweightConvLayer, self).__init__("depthwise", conv_width, 1)
+    self._num_repeat = num_repeat
+
+  def _conv_function(self, input_tensor, output_depth, padding):
+    input_depth = input_tensor.shape.as_list()[-1]
+    if not ((output_depth >= input_depth) and
+            (output_depth % input_depth == 0)):
+      raise ValueError(
+          "Depthwise layer output_depth (%s) must be greater or equal to and "
+          "a multiple of the depth of the "
+          "input tensor (%s)." % (output_depth, input_depth))
+    channel_multiplier = calculate_depthwise_channel_multiplier(
+        input_depth, output_depth)
+
+    num_input_variables = input_depth / self._num_repeat
+    kernel_base = tf.get_variable(
+        "kernel_base",
+        [self._conv_width, 1, num_input_variables, channel_multiplier])
+    kernel = tf.concat([kernel_base] * self._num_repeat, axis=2)
+
+    num_nonrepeated_variables = input_depth % self._num_repeat
+    if num_nonrepeated_variables:
+      nonrepeated_variables = tf.get_variable(
+          "nonrepeated_kernel_variables",
+          [self._conv_width, 1, num_nonrepeated_variables, channel_multiplier])
+      kernel = tf.concat([kernel, nonrepeated_variables], axis=2)
+
+    kernel = tf.nn.softmax(kernel, axis=0)
+    return tf.nn.depthwise_conv2d(
+        input_tensor,
+        kernel, [1, 1, 1, 1],
+        padding=padding,
+        name="lightweight_conv_%sx1_r_%s" % (str(self._conv_width),
+                                             str(self._num_repeat)))
+
+  def num_params(self, input_depth, output_depth, **unused_kwargs):
+    channel_multiplier = calculate_depthwise_channel_multiplier(
+        input_depth, output_depth)
+    return self._conv_width * (int(input_depth / self._num_repeat) + (
+        input_depth % self._num_repeat)) * channel_multiplier
+
+
+class DilatedConvLayer(ConvLayerBase):
+  """Dilated convolution TranslationLayer base class."""
+
+  def __init__(self, conv_width):
+    super(DilatedConvLayer, self).__init__("dilated", conv_width, 2)
+
+  def _conv_function(self, input_tensor, output_depth, padding):
+    input_depth = input_tensor.shape.as_list()[-1]
+    kernel = tf.get_variable("kernel",
+                             [self._conv_width, 1, input_depth, output_depth])
+    return tf.nn.atrous_conv2d(
+        input_tensor,
+        kernel,
+        self._dilation_rate,
+        padding=padding,
+        name="dilated_conv_%sx1" % str(self._conv_width))
+
+  def num_params(self, input_depth, output_depth, **unused_kwargs):
+    return self._conv_width * input_depth * output_depth
+
+
+class AttentionLayer(TranslationLayer):
+  """Attention layer base class."""
+
+  def __init__(self,
+               hidden_dim_multiplier,
+               project_q,
+               project_k,
+               project_v,
+               num_heads=None):
+    self._hidden_dim_multiplier = hidden_dim_multiplier
+    self._project_q = project_q
+    self._project_k = project_k
+    self._project_v = project_v
+    self._num_heads = num_heads
+
+  def _apply_logic(self,
+                   input_tensor,
+                   output_depth,
+                   hparams,
+                   var_scope_suffix,
+                   nonpadding,
+                   mask_future,
+                   decoder_self_attention_bias=None,
+                   attention_dropout_broadcast_dims=None,
+                   **kwargs):
+    """Applies attention logic to `input_tensor`."""
+    with tf.variable_scope("standard_attention_layer_" + var_scope_suffix):
+      hidden_depth = int(
+          input_tensor.shape.as_list()[-1] * self._hidden_dim_multiplier)
+
+      attention_bias = decoder_self_attention_bias
+
+      # TODO(davidso): This dropout rate differs from the other layers. This
+      #                should be fixed so that they all use the same dropout
+      #                rate.
+      num_heads = self._num_heads
+      if num_heads is None:
+        num_heads = hparams.num_heads
+      logic_output = common_attention.multihead_attention(
+          input_tensor,
+          None,
+          attention_bias,
+          hidden_depth,
+          hidden_depth,
+          output_depth,
+          num_heads,
+          hparams.attention_dropout,
+          attention_type=hparams.self_attention_type,
+          max_relative_position=hparams.max_relative_position,
+          dropout_broadcast_dims=attention_dropout_broadcast_dims)
+    return logic_output
+
+  def num_params(self, input_depth, output_depth, **unused_kwargs):
+    # First account for the hidden to output projection params.
+    hidden_depth = input_depth * self._hidden_dim_multiplier
+    output_params = hidden_depth * output_depth
+
+    # Next account for all the hidden projections.
+    num_projections = sum([self._project_q, self._project_k, self._project_v])
+    return input_depth * hidden_depth * num_projections + output_params
+
+
+class AttendToEncoderLayerBase(TranslationLayer):
+  """Attend to encoder base, with configurable encoder attend points."""
+
+  def _determine_encoder_block_index(self, block_number, num_encoder_blocks):
+    """Determine the encoder block index to attend to."""
+    raise NotImplementedError()
+
+  def _apply_logic(self,
+                   input_tensor,
+                   output_depth,
+                   hparams,
+                   var_scope_suffix,
+                   nonpadding,
+                   mask_future,
+                   encoder_decoder_attention_bias,
+                   encoder_block_outputs,
+                   block_number,
+                   attention_dropout_broadcast_dims=None,
+                   **unused_kwargs):
+    """Applies attention logic to `input_tensor`."""
+    with tf.variable_scope("attend_to_encoder_layer_" + var_scope_suffix):
+      hidden_depth = int(input_tensor.shape.as_list()[-1])
+      num_encoder_blocks = len(encoder_block_outputs)
+      encoder_block_index = self._determine_encoder_block_index(
+          block_number, num_encoder_blocks)
+      encoder_layer = encoder_block_outputs[encoder_block_index]
+
+      # TODO(davidso): This dropout rate differs from the other layers. This
+      #                should be fixed so that they all use the same dropout
+      #                rate.
+      logic_output = common_attention.multihead_attention(
+          input_tensor,
+          encoder_layer,
+          encoder_decoder_attention_bias,
+          hidden_depth,
+          hidden_depth,
+          output_depth,
+          hparams.num_heads,
+          hparams.attention_dropout,
+          attention_type=hparams.self_attention_type,
+          max_relative_position=hparams.max_relative_position,
+          dropout_broadcast_dims=attention_dropout_broadcast_dims)
+
+    return logic_output
+
+  # Assumes uniform encoder output depths.
+  def num_params(self, input_depth, output_depth, **kwargs):
+    try:
+      encoder_depth = kwargs["encoder_depth"]
+    except KeyError:
+      raise ValueError("`encoder_depth` must be in kwargs passed to "
+                       "AttendToEncoder.num_params().")
+    hidden_depth = input_depth
+
+    # The number of params is comprised of the projection from the input tensor
+    # to its hidden tensor, the two encoder tensor projects to its hidden
+    # tensors, and the projection from the hidden concatenation to the output
+    # tensor.
+    return (input_depth * hidden_depth + 2 * encoder_depth * hidden_depth +
+            hidden_depth * output_depth)
+
+
+class AttendToEncoderTopDownLayer(AttendToEncoderLayerBase):
+  """Attend to the encoder starting with the highest layer, then moving down.
+
+    This allows the decoder to see higher level features first and then
+    eventually move on to incorporate lower level information.
+  """
+
+  def __init__(self, delay, increment_step):
+    self.delay = delay
+    self.increment_step = increment_step
+
+  def _determine_encoder_block_index(self, block_number, num_encoder_blocks):
+    """Attend to final encoder block output first, then move down."""
+    return max(
+        0, num_encoder_blocks - max(
+            0, (block_number - self.delay) * self.increment_step) - 1)
+
+
+class GatedLinearUnitLayer(TranslationLayer):
+  """Gated Linaer Unit Layer."""
+
+  def __init__(self):
+    pass
+
+  def _apply_logic(self, input_tensor, output_depth, hparams, var_scope_suffix,
+                   nonpadding, mask_future, **unused_kwargs):
+    values = tf.layers.dense(input_tensor, output_depth)
+    gates = tf.layers.dense(
+        input_tensor, output_depth, activation=tf.nn.sigmoid)
+    return values * gates
+
+  def num_params(self, input_depth, output_depth, **unused_kwargs):
+    return input_depth * output_depth * 2 + output_depth * 2
+
+
+class IdentityLayer(TranslationLayer):
+  """Identity TranslationLayer."""
+
+  def _apply_logic(self, input_tensor, output_depth, hparams, var_scope_suffix,
+                   nonpadding, mask_future, **unused_kwargs):
+    input_depth = input_tensor.shape.as_list()[-1]
+    if output_depth != input_depth:
+      raise ValueError(
+          "Identity layer output_depth (%s) must be equal to the depth of the "
+          "input tensor (%s)." % (output_depth, input_depth))
+    return input_tensor
+
+  def num_params(self, input_depth, output_depth, **unused_kwargs):
+    return 0
+
+
+def register_encoder_decoder_layer(name, translation_layer):
+  ENCODER_LAYERS.register_layer(name, translation_layer)
+  DECODER_LAYERS.register_layer(name, translation_layer)
+
+
+# Register all strictly decoder layers.
+DECODER_LAYERS.register_layer(
+    ATTEND_TO_ENCODER_REGISTRY_KEY,
+    AttendToEncoderTopDownLayer(delay=0, increment_step=0))
+
+# Register all encoder and decoder layers.
+register_encoder_decoder_layer(IDENTITY_REGISTRY_KEY, IdentityLayer())
+
+register_encoder_decoder_layer(SEPARABLE_CONV_3X1_REGISTRY_KEY,
+                               SeparableConvLayer(conv_width=3))
+register_encoder_decoder_layer(SEPARABLE_CONV_5X1_REGISTRY_KEY,
+                               SeparableConvLayer(conv_width=5))
+register_encoder_decoder_layer(SEPARABLE_CONV_7X1_REGISTRY_KEY,
+                               SeparableConvLayer(conv_width=7))
+register_encoder_decoder_layer(SEPARABLE_CONV_9X1_REGISTRY_KEY,
+                               SeparableConvLayer(conv_width=9))
+register_encoder_decoder_layer(SEPARABLE_CONV_11X1_REGISTRY_KEY,
+                               SeparableConvLayer(conv_width=11))
+register_encoder_decoder_layer(SEPARABLE_CONV_13X1_REGISTRY_KEY,
+                               SeparableConvLayer(conv_width=13))
+register_encoder_decoder_layer(SEPARABLE_CONV_15X1_REGISTRY_KEY,
+                               SeparableConvLayer(conv_width=15))
+register_encoder_decoder_layer(STANDARD_CONV_1X1_REGISTRY_KEY,
+                               StandardConvLayer(conv_width=1))
+register_encoder_decoder_layer(STANDARD_CONV_3X1_REGISTRY_KEY,
+                               StandardConvLayer(conv_width=3))
+register_encoder_decoder_layer(STANDARD_CONV_5X1_REGISTRY_KEY,
+                               StandardConvLayer(conv_width=5))
+register_encoder_decoder_layer(DEPTHWISE_CONV_3X1_REGISTRY_KEY,
+                               DepthwiseConvLayer(conv_width=3))
+register_encoder_decoder_layer(DEPTHWISE_CONV_5X1_REGISTRY_KEY,
+                               DepthwiseConvLayer(conv_width=5))
+register_encoder_decoder_layer(DEPTHWISE_CONV_7X1_REGISTRY_KEY,
+                               DepthwiseConvLayer(conv_width=7))
+register_encoder_decoder_layer(DILATED_CONV_3X1_REGISTRY_KEY,
+                               DilatedConvLayer(conv_width=3))
+register_encoder_decoder_layer(DILATED_CONV_5X1_REGISTRY_KEY,
+                               DilatedConvLayer(conv_width=5))
+
+
+register_encoder_decoder_layer(LIGHTWEIGHT_CONV_3X1_R_1_REGISTRY_KEY,
+                               LightweightConvLayer(conv_width=3, num_repeat=1))
+register_encoder_decoder_layer(LIGHTWEIGHT_CONV_3X1_R_4_REGISTRY_KEY,
+                               LightweightConvLayer(conv_width=3, num_repeat=4))
+register_encoder_decoder_layer(
+    LIGHTWEIGHT_CONV_3X1_R_16_REGISTRY_KEY,
+    LightweightConvLayer(conv_width=3, num_repeat=16))
+register_encoder_decoder_layer(LIGHTWEIGHT_CONV_5X1_R_1_REGISTRY_KEY,
+                               LightweightConvLayer(conv_width=5, num_repeat=1))
+register_encoder_decoder_layer(LIGHTWEIGHT_CONV_5X1_R_4_REGISTRY_KEY,
+                               LightweightConvLayer(conv_width=5, num_repeat=4))
+register_encoder_decoder_layer(
+    LIGHTWEIGHT_CONV_5X1_R_16_REGISTRY_KEY,
+    LightweightConvLayer(conv_width=5, num_repeat=16))
+register_encoder_decoder_layer(LIGHTWEIGHT_CONV_7X1_R_1_REGISTRY_KEY,
+                               LightweightConvLayer(conv_width=7, num_repeat=1))
+register_encoder_decoder_layer(LIGHTWEIGHT_CONV_7X1_R_4_REGISTRY_KEY,
+                               LightweightConvLayer(conv_width=7, num_repeat=4))
+register_encoder_decoder_layer(
+    LIGHTWEIGHT_CONV_7X1_R_16_REGISTRY_KEY,
+    LightweightConvLayer(conv_width=7, num_repeat=16))
+register_encoder_decoder_layer(
+    LIGHTWEIGHT_CONV_15X1_R_1_REGISTRY_KEY,
+    LightweightConvLayer(conv_width=15, num_repeat=1))
+register_encoder_decoder_layer(
+    LIGHTWEIGHT_CONV_15X1_R_4_REGISTRY_KEY,
+    LightweightConvLayer(conv_width=15, num_repeat=4))
+register_encoder_decoder_layer(
+    LIGHTWEIGHT_CONV_15X1_R_16_REGISTRY_KEY,
+    LightweightConvLayer(conv_width=15, num_repeat=16))
+
+register_encoder_decoder_layer(
+    GATED_LINEAR_UNIT_REGISTRY_KEY,
+    GatedLinearUnitLayer())
+
+
+register_encoder_decoder_layer(
+    STANDARD_ATTENTION_REGISTRY_KEY,
+    AttentionLayer(
+        hidden_dim_multiplier=1, project_q=True, project_k=True,
+        project_v=True))
+register_encoder_decoder_layer(
+    ATTENTION_16_HEADS_REGISTRY_KEY,
+    AttentionLayer(
+        hidden_dim_multiplier=1,
+        project_q=True,
+        project_k=True,
+        project_v=True,
+        num_heads=16))
+register_encoder_decoder_layer(
+    ATTENTION_32_HEADS_REGISTRY_KEY,
+    AttentionLayer(
+        hidden_dim_multiplier=1,
+        project_q=True,
+        project_k=True,
+        project_v=True,
+        num_heads=32))
+register_encoder_decoder_layer(
+    ATTENTION_4_HEADS_REGISTRY_KEY,
+    AttentionLayer(
+        hidden_dim_multiplier=1,
+        project_q=True,
+        project_k=True,
+        project_v=True,
+        num_heads=4))
diff --git a/tensor2tensor/models/neural_architecture_search/nas_layers_test.py b/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
new file mode 100644
index 000000000..07c088aab
--- /dev/null
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
@@ -0,0 +1,320 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Layers tests."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import copy
+import itertools
+from absl.testing import parameterized
+import numpy as np
+
+from tensor2tensor.layers import common_attention
+from tensor2tensor.models import transformer
+from tensor2tensor.models.neural_architecture_search import nas_layers as layers
+
+import tensorflow as tf
+
+_BATCH_SIZE = 32
+_TOTAL_SEQUENCE_LENGTH = 20
+_INPUT_DEPTH = 256
+_NUM_BLOCKS = 6
+_BLOCK_NUMBER = 3
+
+# The list of prefixes for layers that will not be tested for resizing outputs.
+_RESIZE_EXEMPT_LAYER_PREFIXES = [
+    "depthwise_conv", "squeeze_and_excitation", "identity", "lightweight_conv",
+]
+
+
+def _apply_encoder_layer(translation_layer, output_depth, nonpadding_list):
+  """Applies an encoder layer with basic arguments."""
+
+  input_tensor = tf.random_uniform(
+      [_BATCH_SIZE, _TOTAL_SEQUENCE_LENGTH, _INPUT_DEPTH]) / 4.0
+  nonpadding = tf.constant(nonpadding_list)
+  residual_tensor = tf.random_uniform(
+      [_BATCH_SIZE, _TOTAL_SEQUENCE_LENGTH, output_depth])
+  hparams = transformer.transformer_base()
+
+  return translation_layer.apply_layer(
+      input_tensor,
+      residual_tensor,
+      output_depth,
+      tf.nn.relu,
+      hparams,
+      "",
+      mask_future=False,
+      nonpadding=nonpadding,
+      layer_preprocess_fn=None,
+      postprocess_dropout=True)
+
+
+def _apply_decoder_layer(translation_layer, input_tensor, output_depth,
+                         encoder_depth):
+  """Applies an decoder layer with basic arguments."""
+
+  residual_tensor_values = np.random.rand(
+      *[_BATCH_SIZE, _TOTAL_SEQUENCE_LENGTH, output_depth]) - .5
+  residual_tensor = tf.constant(residual_tensor_values, dtype=tf.float32)
+  encoder_output_values = np.random.rand(
+      *[_BATCH_SIZE, _TOTAL_SEQUENCE_LENGTH, encoder_depth]) - .5
+  encoder_output = tf.constant(encoder_output_values, dtype=tf.float32)
+  encoder_block_outputs = [encoder_output] * _NUM_BLOCKS
+  hparams = transformer.transformer_base()
+  hparams.attention_dropout = 0
+  decoder_self_attention_bias = (
+      common_attention.attention_bias_lower_triangle(_TOTAL_SEQUENCE_LENGTH))
+
+  output_tensor = translation_layer.apply_layer(
+      input_tensor,
+      residual_tensor,
+      output_depth,
+      None,
+      hparams,
+      "",
+      nonpadding=None,
+      mask_future=True,
+      layer_preprocess_fn=None,
+      postprocess_dropout=False,
+      decoder_self_attention_bias=decoder_self_attention_bias,
+      encoder_decoder_attention_bias=None,
+      encoder_block_outputs=encoder_block_outputs,
+      block_number=_BLOCK_NUMBER)
+
+  return output_tensor
+
+
+def _zero_after_index_copy(feed_input, zero_after_index):
+  """Creates a copy of `feed_input` with zeros after `zero_after_index`."""
+  transformed_feed_input = copy.deepcopy(feed_input)
+  for i in range(_BATCH_SIZE):
+    for j in range(zero_after_index + 1, _TOTAL_SEQUENCE_LENGTH):
+      transformed_feed_input[i][j] = [0.0] * len(transformed_feed_input[i][j])
+  return transformed_feed_input
+
+
+def _get_empirical_parameters():
+  """Gets the number of parameters built into the current Tensorflow graph."""
+  trainable_variables_list = tf.trainable_variables()
+
+  empirical_num_params = 0
+  for variable_tensor in trainable_variables_list:
+    empirical_num_params += np.prod(variable_tensor.shape)
+
+  return empirical_num_params
+
+
+def _create_nonpadding_list():
+  """Creates the `nonpadding_list` for applying the encoder layers."""
+  nonpadding_list = []
+  for i in range(_BATCH_SIZE):
+    nonpadding_list.append([1.0] * min(i + 2, _TOTAL_SEQUENCE_LENGTH) +
+                           [0.0] * max((_TOTAL_SEQUENCE_LENGTH - i - 2), 0))
+  return nonpadding_list
+
+
+class LayersTest(parameterized.TestCase, tf.test.TestCase):
+  """Tests params, residual capabilities, padding leaks, and output shape."""
+
+  # Test that the encoder registry contains all the expected layers.
+  def test_encoder_registry(self):
+    encoder_layers = [
+        "separable_conv_3x1",
+        "separable_conv_5x1",
+        "separable_conv_7x1",
+        "separable_conv_9x1",
+        "separable_conv_11x1",
+        "separable_conv_13x1",
+        "separable_conv_15x1",
+        "standard_conv_1x1",
+        "standard_conv_3x1",
+        "standard_conv_5x1",
+        "depthwise_conv_3x1",
+        "depthwise_conv_5x1",
+        "depthwise_conv_7x1",
+        "dilated_conv_3x1",
+        "dilated_conv_5x1",
+        "standard_attention",
+        "identity",
+        "attention_4_heads",
+        "attention_16_heads",
+        "attention_32_heads",
+        "gated_linear_unit",
+        "lightweight_conv_3x1_r_1",
+        "lightweight_conv_3x1_r_4",
+        "lightweight_conv_3x1_r_16",
+        "lightweight_conv_5x1_r_1",
+        "lightweight_conv_5x1_r_4",
+        "lightweight_conv_5x1_r_16",
+        "lightweight_conv_7x1_r_1",
+        "lightweight_conv_7x1_r_4",
+        "lightweight_conv_7x1_r_16",
+        "lightweight_conv_15x1_r_1",
+        "lightweight_conv_15x1_r_4",
+        "lightweight_conv_15x1_r_16",
+    ]
+    self.assertSameElements(encoder_layers,
+                            layers.ENCODER_LAYERS.get_layer_names())
+
+  # Test that the decoder registry contains all the expected layers.
+  def test_decoder_registry(self):
+    decoder_layers = sorted([
+        "separable_conv_3x1",
+        "separable_conv_5x1",
+        "separable_conv_7x1",
+        "separable_conv_9x1",
+        "separable_conv_11x1",
+        "separable_conv_13x1",
+        "separable_conv_15x1",
+        "standard_conv_1x1",
+        "standard_conv_3x1",
+        "standard_conv_5x1",
+        "depthwise_conv_3x1",
+        "depthwise_conv_5x1",
+        "depthwise_conv_7x1",
+        "dilated_conv_3x1",
+        "dilated_conv_5x1",
+        "standard_attention",
+        "attend_to_encoder",
+        "identity",
+        "attention_4_heads",
+        "attention_16_heads",
+        "attention_32_heads",
+        "gated_linear_unit",
+        "lightweight_conv_3x1_r_1",
+        "lightweight_conv_3x1_r_4",
+        "lightweight_conv_3x1_r_16",
+        "lightweight_conv_5x1_r_1",
+        "lightweight_conv_5x1_r_4",
+        "lightweight_conv_5x1_r_16",
+        "lightweight_conv_7x1_r_1",
+        "lightweight_conv_7x1_r_4",
+        "lightweight_conv_7x1_r_16",
+        "lightweight_conv_15x1_r_1",
+        "lightweight_conv_15x1_r_4",
+        "lightweight_conv_15x1_r_16",
+    ])
+    self.assertSameElements(decoder_layers,
+                            layers.DECODER_LAYERS.get_layer_names())
+
+  # Test encoder layer. This includes checking that output dims are as
+  # expected, checking that num_params() agrees with the empirical number of
+  # variables produced, and that information does not leak from 0 padded
+  # areas of the input.
+  @parameterized.parameters(
+      itertools.product(layers.ENCODER_LAYERS.get_layer_names(),
+                        (256, 128, 512)))
+  def test_encoder_layer(self, translation_layer_name, output_depth):
+    with self.test_session(graph=tf.Graph()) as sess:
+      nonpadding_list = _create_nonpadding_list()
+      for prefix in _RESIZE_EXEMPT_LAYER_PREFIXES:
+        if prefix in translation_layer_name:
+          output_depth = _INPUT_DEPTH
+      translation_layer = layers.ENCODER_LAYERS.get(translation_layer_name)
+      output_tensor = _apply_encoder_layer(translation_layer, output_depth,
+                                           nonpadding_list)
+
+      # Check that the output shape is as expected.
+      self.assertEqual(output_tensor.shape.as_list(),
+                       [_BATCH_SIZE, _TOTAL_SEQUENCE_LENGTH, output_depth])
+
+      # Check that the number of parameters is as expected.
+      empirical_num_params = _get_empirical_parameters()
+      reported_num_params = translation_layer.num_params(
+          _INPUT_DEPTH, output_depth)
+      self.assertEqual(empirical_num_params, reported_num_params)
+
+      # Make sure padding is applied properly (no leaks).
+      sess.run(tf.global_variables_initializer())
+      output = sess.run(output_tensor)
+
+    for i, j in itertools.product(
+        range(_BATCH_SIZE), range(_TOTAL_SEQUENCE_LENGTH)):
+      if nonpadding_list[i][j] == 0:
+        self.assertAllEqual(output[i][j], np.array([0] * output_depth),
+                            "Output row %s, column %s not zeroed out." % (i, j))
+
+  # Test decoder layer. This includes checking that output dims are as
+  # expected, checking that num_params() agrees with the empirical number of
+  # variables produced, and that temporal information does not leak.
+  @parameterized.parameters(
+      itertools.product(layers.DECODER_LAYERS.get_layer_names(),
+                        (256, 128, 512)))
+  def test_decoder_layer(self, translation_layer_name, output_depth):
+    with self.test_session(graph=tf.Graph()) as sess:
+
+      # Check that the output shape is as expected.
+      input_tensor = tf.placeholder(
+          tf.float32, [_BATCH_SIZE, _TOTAL_SEQUENCE_LENGTH, _INPUT_DEPTH])
+      encoder_depth = int(_INPUT_DEPTH / 2)
+      for prefix in _RESIZE_EXEMPT_LAYER_PREFIXES:
+        if prefix in translation_layer_name:
+          output_depth = _INPUT_DEPTH
+      translation_layer = layers.DECODER_LAYERS.get(translation_layer_name)
+      output_tensor = _apply_decoder_layer(translation_layer, input_tensor,
+                                           output_depth, encoder_depth)
+      self.assertEqual(output_tensor.shape.as_list(),
+                       [_BATCH_SIZE, _TOTAL_SEQUENCE_LENGTH, output_depth])
+
+      # Check that the number of parameters is as expected.
+      empirical_num_params = _get_empirical_parameters()
+      reported_num_params = translation_layer.num_params(
+          _INPUT_DEPTH,
+          output_depth,
+          encoder_depth=encoder_depth)
+      self.assertEqual(empirical_num_params, reported_num_params)
+
+      # Check that there is no temporal information leak. Specifically, check
+      # that values before `test_index` remain unchanged, while the values
+      # after it have changed. Sums are used because two values could
+      # potentially be the same between the zero and non-zero portions, even
+      # if the masking is working correctly. Note: This assumes that the
+      # output at t is dependent on the input at t.
+      feed_input = np.random.random(
+          [_BATCH_SIZE, _TOTAL_SEQUENCE_LENGTH, _INPUT_DEPTH]) / 10.0
+      test_index = int(_TOTAL_SEQUENCE_LENGTH / 2)
+      transformed_feed_input = _zero_after_index_copy(feed_input, test_index)
+
+      # Produce the outputs for both types of input.
+      feed_dict = {
+          v: np.random.rand(*v.shape.as_list()) - .5
+          for v in tf.all_variables()
+      }
+      feed_dict[input_tensor] = feed_input
+      control_output = sess.run(output_tensor, feed_dict)
+
+      feed_dict[input_tensor] = transformed_feed_input
+      variable_output = sess.run(output_tensor, feed_dict)
+
+      self.assertAllClose(
+          control_output[:, :test_index + 1],
+          variable_output[:, :test_index + 1],
+          rtol=1)
+
+      with self.assertRaises(
+          AssertionError,
+          msg="Time-masked portion of output too close to control output."):
+        self.assertAllClose(
+            control_output[:, test_index + 1:],
+            variable_output[:, test_index + 1:],
+            rtol=1)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/neural_architecture_search/nas_model.py b/tensor2tensor/models/neural_architecture_search/nas_model.py
new file mode 100644
index 000000000..37d7d04e2
--- /dev/null
+++ b/tensor2tensor/models/neural_architecture_search/nas_model.py
@@ -0,0 +1,1030 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""TranslationNasNet class which can be modified and still used with t2t."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import abc
+import six
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_layers
+from tensor2tensor.models import transformer
+from tensor2tensor.models.neural_architecture_search import nas_layers as layers
+from tensor2tensor.utils import metrics
+from tensor2tensor.utils import t2t_model
+import tensorflow as tf
+
+
+# Keys for the activation map.
+LEAKY_RELU_ACTIVATION_KEY = "leaky_relu"
+NONE_ACTIVATION_KEY = "none"
+RELU_ACTIVATION_KEY = "relu"
+SIGMOID_ACTIVATION_KEY = "sigmoid"
+SWISH_ACTIVATION_KEY = "swish"
+SOFTMAX_ACTIVATION_KEY = "softmax"
+
+# Mapping from string names to activation function.
+ACTIVATION_MAP = {
+    SWISH_ACTIVATION_KEY: tf.nn.swish,
+    LEAKY_RELU_ACTIVATION_KEY: tf.nn.leaky_relu,
+    RELU_ACTIVATION_KEY: tf.nn.relu,
+    NONE_ACTIVATION_KEY: None,
+    SIGMOID_ACTIVATION_KEY: tf.nn.sigmoid,
+    SOFTMAX_ACTIVATION_KEY: tf.nn.softmax
+}
+
+# Norm strings.
+LAYER_NORM_KEY = "layer_norm"
+NO_NORM_KEY = "none"
+
+# Combiner function strings.
+ADD_COMBINER_FUNC_KEY = "add"
+MULTIPLY_COMBINER_FUNC_KEY = "multiply"
+CONCAT_COMBINER_FUNC_KEY = "concat"
+
+# Layers that force the output_dim to be equal to the input_dim if
+# enforce_fixed_output_sizes is True.
+LAYERS_TO_FIX_OUTPUT_SIZE = [
+    layers.IDENTITY_REGISTRY_KEY,
+]
+
+# Depthwise layers that the output dimension will need to be changed for
+# if channel multiplier cannot be changed to match output dimension.
+DEPTHWISE_LAYERS = [
+    layers.DEPTHWISE_CONV_3X1_REGISTRY_KEY,
+    layers.DEPTHWISE_CONV_5X1_REGISTRY_KEY,
+    layers.DEPTHWISE_CONV_7X1_REGISTRY_KEY
+]
+
+DEAD_BRANCH_KEY = "dead_branch"
+
+
+def should_alter_output_dim(layer_name, enforce_fixed_output_sizes, input_depth,
+                            output_depth):
+  """Check if the output_depth for the specified layer should be changed."""
+  # Check to see if output_depth should be changed if we are using
+  # a depthwise operation and the channel multiplier is returned as 1,
+  # which means that the depthwise multiplier could not be set to match
+  # output_depth.
+  change_dim_for_depthwise = ((layer_name in DEPTHWISE_LAYERS) and
+                              (layers.calculate_depthwise_channel_multiplier(
+                                  input_depth, output_depth) == 1))
+  # See if layer is in LAYERS_TO_FIX_OUTPUT_SIZE and if it is then we
+  # know that the output_dim must be input_dim.
+  change_dim_for_other = layer_name in LAYERS_TO_FIX_OUTPUT_SIZE
+  # Must be sure enforce_fixed_output_sizes is true.
+  return ((change_dim_for_depthwise or change_dim_for_other) and
+          enforce_fixed_output_sizes)
+
+
+def get_activation_names():
+  return ACTIVATION_MAP.keys()
+
+
+def _pad_shallow_tensors(tensors, pad_value):
+  """Pads the shorter tensors to be as long as the longest."""
+  max_dim = 0
+  for tensor in tensors:
+    dim = tensor.shape.as_list()[-1]
+    if dim > max_dim:
+      max_dim = dim
+
+  output_tensors = []
+  for tensor in tensors:
+    dim = tensor.shape.as_list()[-1]
+    if tensor.shape.as_list()[-1] < max_dim:
+      output_tensors.append(
+          tf.pad(
+              tensor, [[0, 0], [0, 0], [0, max_dim - dim]],
+              constant_values=pad_value))
+    else:
+      output_tensors.append(tensor)
+  print(output_tensors)
+
+  return output_tensors
+
+
+class CombinerFunction(object):
+  """Interface for combiner functions."""
+
+  __metaclass__ = abc.ABCMeta
+
+  @abc.abstractmethod
+  def combine_tensors(self, tensors):
+    """Combines `tensors`.
+
+    Args:
+      tensors: List of tensors to combine.
+
+    Returns:
+      Combined tensor.
+    """
+
+  @abc.abstractmethod
+  def combined_output_dim(self, output_dims):
+    """Determines the output dimension of the combined tensor.
+
+    Args:
+      output_dims: List of output dimensions of combined tensors.
+
+    Returns:
+      Output dimension of the combined tensor.
+    """
+
+
+class AddCombiner(CombinerFunction):
+  """Addition CombinerFunction."""
+
+  def combine_tensors(self, tensors):
+    assert tensors
+
+    if len(tensors) == 1:
+      return tensors[0]
+
+    tensors_to_combine = _pad_shallow_tensors(tensors, 0)
+
+    output_tensor = tensors_to_combine[0] + tensors_to_combine[1]
+    for tensor in tensors_to_combine[2:]:
+      output_tensor += tensor
+
+    return output_tensor
+
+  def combined_output_dim(self, output_dims):
+    return max(output_dims)
+
+
+class MultiplyCombiner(CombinerFunction):
+  """Multiply CombinerFunction."""
+
+  def combine_tensors(self, tensors):
+    assert tensors
+
+    if len(tensors) == 1:
+      return tensors[0]
+
+    tensors_to_combine = _pad_shallow_tensors(tensors, 1)
+
+    output_tensor = tensors_to_combine[0] * tensors_to_combine[1]
+    for tensor in tensors_to_combine[2:]:
+      output_tensor *= tensor
+
+    return output_tensor
+
+  def combined_output_dim(self, output_dims):
+    return max(output_dims)
+
+
+class ConcatCombiner(CombinerFunction):
+  """Concat CombinerFunction."""
+
+  def combine_tensors(self, tensors):
+    assert tensors
+
+    if len(tensors) == 1:
+      return tensors[0]
+
+    return tf.concat(tensors, 2)
+
+  def combined_output_dim(self, output_dims):
+    concat_tensor_dim = 0
+    for output_dim in output_dims:
+      concat_tensor_dim += output_dim
+
+    return concat_tensor_dim
+
+
+# Dict of combiner functions where each key is the function key string and each
+# value is a function that takes a list of tensors and outputs the tensors'
+# combination.
+COMBINER_FUNCTIONS = {
+    ADD_COMBINER_FUNC_KEY: AddCombiner,
+    MULTIPLY_COMBINER_FUNC_KEY: MultiplyCombiner,
+    CONCAT_COMBINER_FUNC_KEY: ConcatCombiner,
+}
+
+
+class NasSeq2Seq(transformer.Transformer):
+  """Configurable seq2seq model that uses NAS-like branching.
+
+  Builds a directed graph of operations with arbitrary branching.
+  """
+  __metaclass__ = abc.ABCMeta
+
+  def encode(self, inputs, target_space, hparams, features=None, losses=None):
+    """Encode inputs using _encoder().
+
+    This performs the same way as transformer.Transformer.encode with the
+    encoder portion replaced with _encoder().
+
+    Args:
+      inputs: Input [batch_size, input_length, input_height, hidden_dim] tensor
+        which will be flattened along the two spatial dimensions.
+      target_space: scalar, target space ID.
+      hparams: Hyperparmeters for model.
+      features: Optionally pass the entire features dictionary as well. This is
+        needed now for "packed" datasets.
+      losses: Unused list of losses.
+
+    Returns:
+      Tuple of:
+          encoder_output: Encoder representation.
+              [batch_size, input_length, hidden_dim]
+          encoder_decoder_attention_bias: Bias and mask weights for
+              encodre-decoder attention. [batch_size, input_length]
+
+    Raises:
+      ValueError: If encoder type not found.
+    """
+    inputs = common_layers.flatten4d3d(inputs)
+
+    encoder_input, self_attention_bias, encoder_decoder_attention_bias = (
+        transformer.transformer_prepare_encoder(
+            inputs, target_space, hparams, features=features))
+
+    encoder_input = tf.nn.dropout(encoder_input,
+                                  1.0 - hparams.layer_prepostprocess_dropout)
+
+    encoder_output = self._encoder(
+        encoder_input,
+        self_attention_bias,
+        hparams,
+        nonpadding=transformer.features_to_nonpadding(features, "inputs"),
+        save_weights_to=self.attention_weights)
+
+    return encoder_output, encoder_decoder_attention_bias
+
+  def decode(self,
+             decoder_input,
+             encoder_output,
+             encoder_decoder_attention_bias,
+             decoder_self_attention_bias,
+             hparams,
+             cache=None,
+             nonpadding=None,
+             losses=None):
+    """Decode inputs using _decoder().
+
+    This performs the same way as transformer.Transformer.decode with the
+    decoder portion replaced with _decoder().
+
+    Args:
+      decoder_input: Inputs to bottom of the model. [batch_size, decoder_length,
+        hidden_dim]
+      encoder_output: Encoder representation. [batch_size, input_length,
+        hidden_dim]
+      encoder_decoder_attention_bias: Bias and mask weights for encoder-decoder
+        attention. [batch_size, input_length]
+      decoder_self_attention_bias: Bias and mask weights for decoder
+        self-attention. [batch_size, decoder_length]
+      hparams: Hyperparmeters for model.
+      cache: Dict, containing tensors which are the results of previous
+        attentions, used for fast decoding.
+      nonpadding: Optional Tensor with shape [batch_size, decoder_length]
+      losses: Unused losses.
+
+    Returns:
+      Final decoder representation. [batch_size, decoder_length, hidden_dim]
+    """
+    decoder_input = tf.nn.dropout(decoder_input,
+                                  1.0 - hparams.layer_prepostprocess_dropout)
+
+    decoder_output = self._decoder(
+        decoder_input,
+        encoder_output,
+        decoder_self_attention_bias,
+        encoder_decoder_attention_bias,
+        hparams,
+        cache=cache,
+        nonpadding=nonpadding,
+        save_weights_to=self.attention_weights)
+
+    if (common_layers.is_xla_compiled() and
+        hparams.mode == tf.estimator.ModeKeys.TRAIN):
+      # TPU does not react kindly to extra dimensions.
+      return decoder_output
+
+    # Expand since t2t expects 4d tensors.
+    return tf.expand_dims(decoder_output, axis=2)
+
+  def _encoder(self,
+               encoder_input,
+               encoder_self_attention_bias,
+               hparams,
+               nonpadding=None,
+               save_weights_to=None):
+    encoder_output, encoder_block_outputs = nas_encoder(
+        encoder_input, encoder_self_attention_bias, hparams, nonpadding)
+    self._encoder_block_outputs = encoder_block_outputs
+    return encoder_output
+
+  def _decoder(self,
+               decoder_input,
+               encoder_output,
+               decoder_self_attention_bias,
+               encoder_decoder_attention_bias,
+               hparams,
+               cache=None,
+               nonpadding=None,
+               save_weights_to=None):
+    assert self._encoder_block_outputs
+    return nas_decoder(decoder_input, self._encoder_block_outputs,
+                       decoder_self_attention_bias,
+                       encoder_decoder_attention_bias, hparams)
+
+  def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
+    """Construct EstimatorSpec for EVAL mode."""
+    if self.hparams.use_tpu:
+      return self._tpu_estimator_spec_eval(features, logits, labels, loss,
+                                           losses_dict)
+    return self._gpu_estimator_spec_eval(features, logits, labels, loss,
+                                         losses_dict)
+
+  # This function is overridden because py_func is not supported on distributed
+  # training, which is necessary for NAS. This function works
+  # the exact same way as the original Transformer.estimator_spec_eval(),
+  # except only neg log perplexity is accepted as a metric.
+  def _gpu_estimator_spec_eval(self, features, logits, labels, loss,
+                               losses_dict):
+    """Construct EstimatorSpec for GPU EVAL mode."""
+    hparams = self.hparams
+
+    if not hasattr(hparams, "problem"):
+      raise NotImplementedError(
+          "hparams is missing attribute `problem`. Seq2SeqNasNet must "
+          "be used with a problem.")
+
+    # TPU is not supported.
+    eval_metrics_fns = metrics.create_evaluation_metrics([hparams.problem],
+                                                         hparams)
+    eval_metrics = {}
+    for metric_name, metric_fn in six.iteritems(eval_metrics_fns):
+      if "rouge" not in metric_name and "bleu" not in metric_name:
+        eval_metrics[metric_name] = metric_fn(logits, features,
+                                              features["targets"])
+
+    return tf.estimator.EstimatorSpec(
+        tf.estimator.ModeKeys.EVAL,
+        predictions={"predictions": logits},
+        eval_metric_ops=eval_metrics,
+        loss=loss)
+
+  def _tpu_estimator_spec_eval(self, features, logits, labels, loss,
+                               losses_dict):
+    """Construct EstimatorSpec for TPU EVAL mode."""
+    del losses_dict
+    hparams = self.hparams
+
+    if not hasattr(hparams, "problem"):
+      raise NotImplementedError(
+          "hparams is missing attribute `problem`. Seq2SeqNasNet must "
+          "be used with a problem.")
+
+    problem = hparams.problem
+    t2t_model.remove_summaries()
+    eval_metrics_fn = t2t_model.create_tpu_eval_metrics_fn(problem, hparams)
+    if isinstance(logits, dict):
+      # For TPU, logits dict will be passed as keyword arguments to
+      # eval_metrics_fn. Here we add the labels to those arguments.
+      logits.update({"labels": labels})
+      return tf.contrib.tpu.TPUEstimatorSpec(
+          tf.estimator.ModeKeys.EVAL,
+          eval_metrics=(eval_metrics_fn, logits),
+          loss=loss)
+    else:
+      return tf.contrib.tpu.TPUEstimatorSpec(
+          tf.estimator.ModeKeys.EVAL,
+          eval_metrics=(eval_metrics_fn, [logits, labels]),
+          loss=loss)
+
+  def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha,
+                   use_tpu):
+    """Forced slow beam decode.
+
+    Args:
+      features: an map of string to `Tensor`.
+      decode_length: an integer.  How many additional timesteps to decode.
+      beam_size: number of beams.
+      top_beams: an integer. How many of the beams to return.
+      alpha: Float that controls the length penalty. larger the alpha, stronger
+        the preference for longer translations.
+      use_tpu: Whether or not TPU is being used.
+
+    Returns:
+      A dict of decoding results {
+          "outputs": integer `Tensor` of decoded ids of shape
+              [batch_size, <= decode_length] if beam_size == 1 or
+              [batch_size, top_beams, <= decode_length].
+          "scores": decoding log probs from the beam search,
+              None if using greedy decoding (beam_size=1).
+      }
+    """
+    return self._beam_decode_slow(features, decode_length, beam_size, top_beams,
+                                  alpha, use_tpu)
+
+
+def _apply_layer_norm(input_tensor, nonpadding, hparams):
+  """Applies Tensor2Tensor layer_norm to |input_tensor|."""
+  input_depth = input_tensor.shape.as_list()[-1]
+  if nonpadding is not None:
+    nonpadding_input_tiled = tf.tile(
+        tf.expand_dims(nonpadding, 2), [1, 1, input_depth])
+    output_tensor = input_tensor * nonpadding_input_tiled
+
+  output_tensor = common_layers.layer_preprocess(input_tensor, hparams)
+  if nonpadding is not None:
+    output_tensor *= nonpadding_input_tiled
+
+  return output_tensor
+
+
+def _apply_nas_branch(
+    norm, layer_norm_dict, hidden_states, nonpadding, hparams, input_index,
+    layer_name, activation_name, layer_registry, output_dim, branch_scope_name,
+    mask_future, dropout_broadcast_dims, encoder_decoder_attention_bias,
+    encoder_block_outputs, decoder_self_attention_bias, block_number):
+  """Applies a single NAS branch."""
+  with tf.variable_scope(branch_scope_name):
+    # Apply layer norm to an individual layer at most one time.
+    if norm == LAYER_NORM_KEY:
+      try:
+        output_tensor = layer_norm_dict[input_index]
+      except KeyError:
+        output_tensor = _apply_layer_norm(hidden_states[input_index],
+                                          nonpadding, hparams)
+        layer_norm_dict[input_index] = output_tensor
+    elif norm == NO_NORM_KEY:
+      output_tensor = hidden_states[input_index]
+    else:
+      raise ValueError("norm must be either '%s' or '%s'. Got %s" %
+                       (LAYER_NORM_KEY, NO_NORM_KEY, norm))
+
+    layer_class = layer_registry.get(layer_name)
+    activation = ACTIVATION_MAP[activation_name]
+
+    postprocess_dropout = layer_name != layers.IDENTITY_REGISTRY_KEY
+    output_tensor = layer_class.apply_layer(
+        output_tensor,
+        None,
+        int(output_dim),
+        activation,
+        hparams,
+        branch_scope_name,
+        mask_future=mask_future,
+        layer_preprocess_fn=None,
+        postprocess_dropout=postprocess_dropout,
+        nonpadding=nonpadding,
+        attention_dropout_broadcast_dims=dropout_broadcast_dims,
+        encoder_decoder_attention_bias=encoder_decoder_attention_bias,
+        encoder_block_outputs=encoder_block_outputs,
+        block_number=block_number,
+        decoder_self_attention_bias=decoder_self_attention_bias)
+
+    return output_tensor
+
+
+def apply_nas_layers(input_tensor,
+                     left_inputs,
+                     left_layers,
+                     left_activations,
+                     left_output_dims,
+                     left_norms,
+                     right_inputs,
+                     right_layers,
+                     right_activations,
+                     right_output_dims,
+                     right_norms,
+                     combiner_functions,
+                     final_combiner_function,
+                     num_blocks,
+                     nonpadding,
+                     layer_registry,
+                     mask_future,
+                     hparams,
+                     var_scope,
+                     encoder_decoder_attention_bias=None,
+                     encoder_block_outputs=None,
+                     decoder_self_attention_bias=None,
+                     final_layer_norm=True,
+                     enforce_fixed_output_sizes=True):
+  """Applies layers with NAS-like branching.
+
+  Args:
+    input_tensor: Input [batch_size, input_length, hidden_dim] sequence tensor.
+    left_inputs: Int list of left branch hidden layer input indexes.
+    left_layers: String list of left branch layers.
+    left_activations: String list of left branch activations.
+    left_output_dims: String list of left branch output dimensions.
+    left_norms: String list of left branch norms.
+    right_inputs: Int list of right branch hidden layer input indexes.
+    right_layers: String list of right branch layers.
+    right_activations: String list of right branch activations.
+    right_output_dims: String list of right branch output dimensions.
+    right_norms: String list of right branch norms.
+    combiner_functions: String list of branch combining functions.
+    final_combiner_function: String. The final combiner function that combines
+      all the unused hidden layers in a block.
+    num_blocks: The number of blocks. This is the number of times the given
+      layers will be repeated.
+    nonpadding: Tensor with 1s at all nonpadding time step positions and 0s
+      everywhere else.
+    layer_registry: The LayerRegistry that holds all valid layers.
+    mask_future: Whether or not to mask future sequence values.
+    hparams: Hyperparameters for the model.
+    var_scope: The variable scope name.
+    encoder_decoder_attention_bias: The attention bias for decoder attending to
+      `encoder_output`.
+    encoder_block_outputs: List of tensors. The encoder block outputs, listed in
+      order.
+    decoder_self_attention_bias: The self attention bias for decoders. This
+      needs to be set for decoders.
+    final_layer_norm: Whether or not to apply a final layer_norm to the output
+      of the model.
+    enforce_fixed_output_sizes: Whether or not to automatically resize output
+      dimensions to match the input dimension if `should_alter_output_dim()`
+      returns True.
+
+  Raises:
+    ValueError: When branching inputs are not of the same length.
+    ValueError: If item in left_norms is not LAYER_NORM_KEY or NO_NORM_KEY.
+    ValueError: If item in right_norms is not LAYER_NORM_KEY or NO_NORM_KEY.
+
+  Returns:
+    Output of applied layers and list of each block's outputs in order.
+  """
+
+  if not (len(left_inputs) == len(left_layers) == len(left_activations) ==
+          len(left_output_dims) == len(left_norms) == len(right_inputs) ==
+          len(right_layers) == len(right_activations) == len(right_output_dims)
+          == len(right_norms) == len(combiner_functions)):
+    raise ValueError("All branching inputs must be of the same length.")
+
+  block_output = None
+  modified_left_inputs = [
+      left_inputs[i]
+      for i in range(len(left_inputs))
+      if left_layers[i] != DEAD_BRANCH_KEY
+  ]
+  modified_right_inputs = [
+      right_inputs[i]
+      for i in range(len(right_inputs))
+      if right_layers[i] != DEAD_BRANCH_KEY
+  ]
+  unused_block_hidden_states = [
+      i for i in range(len(left_inputs) + 1)
+      if i not in modified_left_inputs and i not in modified_right_inputs
+  ]
+  assert unused_block_hidden_states
+
+  block_outputs = []
+
+  with tf.variable_scope(var_scope):
+    dropout_broadcast_dims = (
+        common_layers.comma_separated_string_to_integer_list(
+            getattr(hparams, "attention_dropout_broadcast_dims", "")))
+
+    for block_num in range(num_blocks):
+      # h_0 is the input tensor.
+      # Keep a dict for layer norm states.
+      if block_output is not None:
+        block_hidden_states = [block_output]
+      else:
+        block_hidden_states = [input_tensor]
+      layer_norm_dict = {}
+
+      with tf.variable_scope("block_%d" % block_num):
+
+        for i, (left_input, left_layer_name, left_activation_name,
+                left_output_dim, left_norm, right_input, right_layer_name,
+                right_activation_name, right_output_dim, right_norm,
+                combiner) in enumerate(
+                    zip(left_inputs, left_layers, left_activations,
+                        left_output_dims, left_norms, right_inputs,
+                        right_layers, right_activations, right_output_dims,
+                        right_norms, combiner_functions)):
+          left_input = int(left_input)
+          right_input = int(right_input)
+
+          with tf.variable_scope("layer_%d" % i):
+
+            assert not (left_layer_name == DEAD_BRANCH_KEY and
+                        right_layer_name == DEAD_BRANCH_KEY)
+
+            if left_layer_name != DEAD_BRANCH_KEY:
+
+              left_raw_input_tensor = block_hidden_states[left_input]
+              left_input_dim = left_raw_input_tensor.shape.as_list()[-1]
+              if should_alter_output_dim(left_layer_name,
+                                         enforce_fixed_output_sizes,
+                                         left_input_dim, left_output_dim):
+                left_output_dim = left_input_dim
+
+              # First process the left branch.
+              left_tensor = _apply_nas_branch(
+                  norm=left_norm,
+                  layer_norm_dict=layer_norm_dict,
+                  hidden_states=block_hidden_states,
+                  nonpadding=nonpadding,
+                  hparams=hparams,
+                  input_index=left_input,
+                  layer_name=left_layer_name,
+                  activation_name=left_activation_name,
+                  layer_registry=layer_registry,
+                  output_dim=left_output_dim,
+                  branch_scope_name="left_%s" % str(i),
+                  mask_future=mask_future,
+                  dropout_broadcast_dims=dropout_broadcast_dims,
+                  encoder_decoder_attention_bias=encoder_decoder_attention_bias,
+                  encoder_block_outputs=encoder_block_outputs,
+                  decoder_self_attention_bias=decoder_self_attention_bias,
+                  block_number=block_num)
+
+            if right_layer_name != DEAD_BRANCH_KEY:
+              right_raw_input_tensor = block_hidden_states[right_input]
+              right_input_dim = right_raw_input_tensor.shape.as_list()[-1]
+              if should_alter_output_dim(right_layer_name,
+                                         enforce_fixed_output_sizes,
+                                         right_input_dim, right_output_dim):
+                right_output_dim = right_input_dim
+              # Next process the right branch.
+              right_tensor = _apply_nas_branch(
+                  norm=right_norm,
+                  layer_norm_dict=layer_norm_dict,
+                  hidden_states=block_hidden_states,
+                  nonpadding=nonpadding,
+                  hparams=hparams,
+                  input_index=right_input,
+                  layer_name=right_layer_name,
+                  activation_name=right_activation_name,
+                  layer_registry=layer_registry,
+                  output_dim=right_output_dim,
+                  branch_scope_name="right_%s" % str(i),
+                  mask_future=mask_future,
+                  dropout_broadcast_dims=dropout_broadcast_dims,
+                  encoder_decoder_attention_bias=encoder_decoder_attention_bias,
+                  encoder_block_outputs=encoder_block_outputs,
+                  decoder_self_attention_bias=decoder_self_attention_bias,
+                  block_number=block_num)
+
+            # Combine the branches.
+            if left_layer_name == DEAD_BRANCH_KEY:
+              hidden_tensor = right_tensor
+            elif right_layer_name == DEAD_BRANCH_KEY:
+              hidden_tensor = left_tensor
+            else:
+              hidden_tensor = COMBINER_FUNCTIONS[combiner]().combine_tensors(
+                  [left_tensor, right_tensor])
+            block_hidden_states.append(hidden_tensor)
+
+      states_to_combine = [
+          block_hidden_states[j] for j in unused_block_hidden_states
+      ]
+      block_output = COMBINER_FUNCTIONS[final_combiner_function](
+      ).combine_tensors(states_to_combine)
+      block_outputs.append(block_output)
+
+  if final_layer_norm:
+    final_output = common_layers.layer_preprocess(block_output, hparams)
+    block_outputs = [
+        common_layers.layer_preprocess(block_output, hparams)
+        for block_output in block_outputs
+    ]
+    return final_output, block_outputs
+  else:
+    return block_output, block_outputs
+
+
+def nas_encoder(encoder_input,
+                encoder_self_attention_bias,
+                hparams,
+                nonpadding=None,
+                final_layer_norm=True):
+  """Encoder for NAS-style model.
+
+  Args:
+    encoder_input: Input tensor.
+    encoder_self_attention_bias: Attention bias tensor with 0s for all valid
+      postions and large negative numbers for the padding positions.
+    hparams: transformer.Transformer hparams that must also contain:
+      + encoder_<left|right>_inputs: List of ints specifying the hidden layer
+        input indexes for the <left|right> branches.
+      + encoder_<left|right>_layers: String list of layers. Each string must be
+        the name of a TranslationLayer registered in layers.py's ENCODER_LAYERS.
+      + encoder_<left|right>_activations: String list of activations. Each
+        string in this list must have a corresponding activation in
+        ACTIVATION_MAP.
+      + encoder_<left|right>_output_dims: Int list of output dimensions for
+        <left|right> branch layers.
+      + encoder_<left|right>_norms: String list of norms to apply to the
+        <left|right> layer branches. Each item must be either LAYER_NORM_KEY or
+        NO_NORM_KEY.
+      + encoder_num_blocks: The number of blocks in the encoder. This determines
+        how many times the given layers will be repeated.
+      + encoder_combiner_functions: String list of functions used to combine
+        left and right branches. Must be a COMBINER_FUNCTION key.
+    nonpadding: Tensor with 1s at all nonpadding positions and 0s everywhere
+      else. If None (default), then nonpadding will be determined from
+      encoder_self_attention_bias.
+    final_layer_norm: Whether or not to apply a final layer_norm to the output
+      of the encoder.
+
+  Returns:
+    Encoder output and list of each encoder block's output in order.
+  """
+  if nonpadding is None:
+    padding = common_attention.attention_bias_to_padding(
+        encoder_self_attention_bias)
+    nonpadding = 1.0 - padding
+  return apply_nas_layers(
+      input_tensor=encoder_input,
+      left_inputs=hparams.encoder_left_inputs,
+      left_layers=hparams.encoder_left_layers,
+      left_activations=hparams.encoder_left_activations,
+      left_output_dims=hparams.encoder_left_output_dims,
+      left_norms=hparams.encoder_left_norms,
+      right_inputs=hparams.encoder_right_inputs,
+      right_layers=hparams.encoder_right_layers,
+      right_activations=hparams.encoder_right_activations,
+      right_output_dims=hparams.encoder_right_output_dims,
+      right_norms=hparams.encoder_right_norms,
+      num_blocks=hparams.encoder_num_blocks,
+      combiner_functions=hparams.encoder_combiner_functions,
+      final_combiner_function=hparams.encoder_final_combiner_function,
+      nonpadding=nonpadding,
+      layer_registry=layers.ENCODER_LAYERS,
+      mask_future=False,
+      hparams=hparams,
+      var_scope="encoder",
+      final_layer_norm=final_layer_norm)
+
+
+def nas_decoder(decoder_input,
+                encoder_block_outputs,
+                decoder_self_attention_bias,
+                encoder_decoder_attention_bias,
+                hparams,
+                final_layer_norm=True):
+  """Decoder for NAS-style model.
+
+  Args:
+    decoder_input: Input tensor.
+    encoder_block_outputs: List of tensors. The encoder block outputs, listed in
+      order.
+    decoder_self_attention_bias: Attention bias that the decoder uses when
+      attending to itself. This should have 0s for all valid positions and large
+      negative numbers for all hidden future positions.
+    encoder_decoder_attention_bias: Attention bias that the decoder uses when
+      attending to the encoder. This should be 0s at all valid positions and
+      large negative numbers for all padded positions.
+    hparams: transformer.Transformer hparams that must also contain:
+      + decoder_<left|right>_inputs: List of ints specifying the hidden layer
+        input indexes for the <left|right> branches.
+      + decoder_<left|right>_layers: String list of layers. Each string must be
+        the name of a TranslationLayer registered in layers.py's DECODER_LAYERS.
+      + decoder_<left|right>_activations: String list of activations. Each
+        string in this list must have a corresponding activation in
+        ACTIVATION_MAP.
+      + decoder_<left|right>_output_dims: Int list of output dimensions for
+        <left|right> branch layers.
+      + decoder_<left|right>_norms: String list of norms to apply to the
+        <left|right> layer branches. Each item must be either LAYER_NORM_KEY or
+        NO_NORM_KEY.
+      + decoder_num_blocks: The number of blocks in the decoder. This determines
+        how many times the given layers will be repeated.
+      + decoder_combiner_functions: String list of functions used to combine
+        left and right branches. Must be a COMBINER_FUNCTION key.
+      hparams may also optionally contain:
+      + enforce_output_size: Boolean that determines whether or not the decoder
+        output must be resized to hparams.hidden_size. If True, the output will
+        be resized if it not equal to hparams.hidden_size. If False, the output
+        will not be resized. If this field is not set, behavior defaults to
+        True.
+    final_layer_norm: Whether or not to apply a final layer norm to the output
+      of the decoder.
+
+  Returns:
+    Decoder output tensor.
+  """
+  # encoder_depth is wrong because it doesn't matter here.
+
+  # Enforce that the output tensor depth is equal to the depth of the encoding.
+  (_, output_depth, _, _) = calculate_branching_model_parameters(
+      encoding_depth=hparams.hidden_size,
+      left_inputs=hparams.decoder_left_inputs,
+      left_layers=hparams.decoder_left_layers,
+      left_output_dims=hparams.decoder_left_output_dims,
+      right_inputs=hparams.decoder_right_inputs,
+      right_layers=hparams.decoder_right_layers,
+      right_output_dims=hparams.decoder_right_output_dims,
+      combiner_functions=hparams.decoder_combiner_functions,
+      final_combiner_function=hparams.decoder_final_combiner_function,
+      layer_registry=layers.DECODER_LAYERS,
+      num_blocks=hparams.decoder_num_blocks,
+      encoder_depth=hparams.hidden_size)
+  improper_output_size = output_depth != hparams.hidden_size
+
+  try:
+    enforce_output_size = hparams.enforce_output_size
+  except AttributeError:
+    enforce_output_size = True
+  resize_output = enforce_output_size and improper_output_size
+
+  decoder_blocks_output, _ = apply_nas_layers(
+      input_tensor=decoder_input,
+      left_inputs=hparams.decoder_left_inputs,
+      left_layers=hparams.decoder_left_layers,
+      left_activations=hparams.decoder_left_activations,
+      left_output_dims=hparams.decoder_left_output_dims,
+      left_norms=hparams.decoder_left_norms,
+      right_inputs=hparams.decoder_right_inputs,
+      right_layers=hparams.decoder_right_layers,
+      right_activations=hparams.decoder_right_activations,
+      right_output_dims=hparams.decoder_right_output_dims,
+      right_norms=hparams.decoder_right_norms,
+      num_blocks=hparams.decoder_num_blocks,
+      combiner_functions=hparams.decoder_combiner_functions,
+      final_combiner_function=hparams.decoder_final_combiner_function,
+      nonpadding=None,
+      layer_registry=layers.DECODER_LAYERS,
+      mask_future=True,
+      hparams=hparams,
+      var_scope="decoder",
+      decoder_self_attention_bias=decoder_self_attention_bias,
+      encoder_decoder_attention_bias=encoder_decoder_attention_bias,
+      encoder_block_outputs=encoder_block_outputs,
+      final_layer_norm=final_layer_norm)
+
+  if not resize_output:
+    return decoder_blocks_output
+
+  # Resize output if necessary.
+  dense_layer = layers.DECODER_LAYERS.get(layers.STANDARD_CONV_1X1_REGISTRY_KEY)
+  output = dense_layer.apply_layer(
+      decoder_blocks_output,
+      None,
+      hparams.hidden_size,
+      None,
+      hparams,
+      "decoder_resize_dense",
+      mask_future=True,
+      layer_preprocess_fn=None,
+      postprocess_dropout=True,
+      nonpadding=None,
+      attention_dropout_broadcast_dims=None,
+      encoder_decoder_attention_bias=None,
+      encoder_block_outputs=None,
+      decoder_self_attention_bias=None,
+  )
+  if final_layer_norm:
+    output = common_layers.layer_preprocess(output, hparams)
+
+  return output
+
+
+def calculate_branching_model_parameters(encoding_depth,
+                                         left_inputs,
+                                         left_layers,
+                                         left_output_dims,
+                                         right_inputs,
+                                         right_layers,
+                                         right_output_dims,
+                                         combiner_functions,
+                                         layer_registry,
+                                         num_blocks,
+                                         final_combiner_function,
+                                         encoder_depth=None,
+                                         enforce_output_size=False,
+                                         enforce_fixed_output_sizes=True):
+  """Calculates the number of parameters in the given model portion.
+
+  Args:
+    encoding_depth: Integer. The depth of the initial input tensor.
+    left_inputs: Integer list. The indexes of the hidden layer inputs for the
+      left branch.
+    left_layers: String list. The names of the left branch layers.
+    left_output_dims: Integer list. The output dimensions for each of the left
+      branch layers.
+    right_inputs: Integer list. The indexes of the hidden layer inputs for the
+      right branch.
+    right_layers: String list. The names of the right branch layers.
+    right_output_dims: Integer list. The output dimensions of each of the right
+      branch layers.
+    combiner_functions: String list. The functions used to combine the left and
+      right branch tensors.
+    layer_registry: layers.LayerRegistry. The LayerRegistry that contains the
+      layers.TranslationLayers needed to construct the model.
+    num_blocks: Integer. The number of times the given layers are repeated to
+      produce the model.
+    final_combiner_function: String. The COMBINER_FUNCTIONS key for the combiner
+      used to combine the unused hidden dimensions.
+    encoder_depth: Integer. The depth of the final encoder layer.
+    enforce_output_size: Boolean. If True, include parameters for the addition
+      of a dense layer that projects the final output to the appropriate
+      `encoding_depth` if it is not already that size. If False, do not add any
+      additional parameters.
+    enforce_fixed_output_sizes: Whether or not to automatically resize output
+      dimensions to match the input dimension if `should_alter_output_dim()`
+      returns True.
+
+  Raises:
+    ValueError: When the layer config lists are not of equal length.
+
+  Returns:
+    total_parameters: The total number of parameters in the model, accounting
+      for repeated blocks.
+    output_depth: The depth of the block output tensor.
+    hidden_depths: The depths of the hidden layers.
+    unused_outputs: List of integer indexes of the hidden layers that are not
+      used as input, and therefore are concatenated to produce the block
+      output.
+  """
+  if not (len(left_inputs) == len(left_layers) == len(left_output_dims) ==
+          len(right_inputs) == len(right_layers) == len(right_output_dims) ==
+          len(combiner_functions)):
+    raise ValueError("Layer configs must be of equal length.")
+
+  total_parameters = 0
+  output_depth = encoding_depth
+  for _ in range(num_blocks):
+    hidden_depths = [output_depth]
+    unused_outputs = set(range(len(left_inputs) + 1))
+
+    for (left_input, left_layer, left_output_dim, right_input,
+         right_layer, right_output_dim, combiner_function) in zip(
+             left_inputs, left_layers, left_output_dims, right_inputs,
+             right_layers, right_output_dims, combiner_functions):
+
+      assert not (left_layer == DEAD_BRANCH_KEY and
+                  right_layer == DEAD_BRANCH_KEY)
+
+      if left_layer == DEAD_BRANCH_KEY:
+        left_parameters = 0
+
+      else:
+        left_input_dim = hidden_depths[left_input]
+        if should_alter_output_dim(left_layer, enforce_fixed_output_sizes,
+                                   left_input_dim, left_output_dim):
+          left_output_dim = left_input_dim
+
+        left_parameters = layer_registry.get(left_layer).num_params(
+            left_input_dim, left_output_dim, encoder_depth=encoder_depth)
+
+      if right_layer == DEAD_BRANCH_KEY:
+        right_parameters = 0
+
+      else:
+        right_input_dim = hidden_depths[right_input]
+        if should_alter_output_dim(right_layer, enforce_fixed_output_sizes,
+                                   right_input_dim, right_output_dim):
+          right_output_dim = right_input_dim
+
+        right_parameters = layer_registry.get(right_layer).num_params(
+            right_input_dim, right_output_dim, encoder_depth=encoder_depth)
+
+      total_parameters += left_parameters + right_parameters
+
+      if left_layer == DEAD_BRANCH_KEY:
+        hidden_dim = right_output_dim
+      elif right_layer == DEAD_BRANCH_KEY:
+        hidden_dim = left_output_dim
+      else:
+        hidden_dim = COMBINER_FUNCTIONS[combiner_function](
+        ).combined_output_dim([left_output_dim, right_output_dim])
+      hidden_depths.append(hidden_dim)
+
+      try:
+        if left_layer != DEAD_BRANCH_KEY:
+          unused_outputs.remove(left_input)
+      except KeyError:
+        pass
+      try:
+        if right_layer != DEAD_BRANCH_KEY:
+          unused_outputs.remove(right_input)
+      except KeyError:
+        pass
+
+    # All unused outputs combined_together.
+    unused_hidden_depths = [hidden_depths[index] for index in unused_outputs]
+    output_depth = COMBINER_FUNCTIONS[final_combiner_function](
+    ).combined_output_dim(unused_hidden_depths)
+
+  # Add the resizing layer if needed.
+  if output_depth != encoding_depth and enforce_output_size:
+    total_parameters += layer_registry.get(
+        layers.STANDARD_CONV_1X1_REGISTRY_KEY).num_params(
+            output_depth, encoding_depth, encoder_depth=encoder_depth)
+
+  return (total_parameters, output_depth, hidden_depths, unused_outputs)
diff --git a/tensor2tensor/models/neural_architecture_search/nas_model_test.py b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
new file mode 100644
index 000000000..c74415969
--- /dev/null
+++ b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
@@ -0,0 +1,468 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 Translation NasNets."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from absl.testing import parameterized
+import numpy as np
+from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.layers import common_attention
+from tensor2tensor.models import transformer
+from tensor2tensor.models.neural_architecture_search import nas_layers as layers
+from tensor2tensor.models.neural_architecture_search import nas_model as translation_nas_net
+import tensorflow as tf
+
+_BATCH_SIZE = 5
+_INPUT_LENGTH = 5
+_TARGET_LENGTH = 6
+_VOCAB_SIZE = 8
+_HIDDEN_SIZE = 512
+_EMBEDDING_DEPTH = _HIDDEN_SIZE
+
+
+def _list_product(num_list):
+  """Computes product of all elements in a list."""
+  product = 1
+  for num in num_list:
+    product *= num
+  return product
+
+
+def _get_transformer_branching_encoder_config():
+  """Returns config for the Transformer encoder."""
+  num_blocks = 2
+  left_inputs = [0, 1, 2, 3]
+  left_layers = [
+      layers.STANDARD_ATTENTION_REGISTRY_KEY,
+      layers.STANDARD_CONV_1X1_REGISTRY_KEY,
+      layers.STANDARD_CONV_1X1_REGISTRY_KEY, layers.IDENTITY_REGISTRY_KEY
+  ]
+  left_output_dims = [512, 2048, 512, 512]
+  right_inputs = [0, 1, 1, 3]
+  right_layers = [
+      layers.IDENTITY_REGISTRY_KEY, translation_nas_net.DEAD_BRANCH_KEY,
+      layers.IDENTITY_REGISTRY_KEY, translation_nas_net.DEAD_BRANCH_KEY
+  ]
+  right_output_dims = [512, 512, 512, 512]
+  combiner_functions = [
+      translation_nas_net.ADD_COMBINER_FUNC_KEY,
+      translation_nas_net.ADD_COMBINER_FUNC_KEY,
+      translation_nas_net.ADD_COMBINER_FUNC_KEY,
+      translation_nas_net.ADD_COMBINER_FUNC_KEY
+  ]
+  dummy_activations = [translation_nas_net.NONE_ACTIVATION_KEY] * 4
+  dummy_norms = [translation_nas_net.NO_NORM_KEY] * 4
+  layer_registry = layers.ENCODER_LAYERS
+  is_decoder = False
+  final_combiner_function = translation_nas_net.CONCAT_COMBINER_FUNC_KEY
+
+  return (num_blocks, left_inputs, left_layers, left_output_dims, right_inputs,
+          right_layers, right_output_dims, combiner_functions,
+          final_combiner_function, dummy_activations, dummy_norms,
+          layer_registry, is_decoder)
+
+
+def _get_transformer_branching_decoder_config():
+  """Returns config for the Transformer decoder."""
+  num_blocks = 2
+  left_inputs = [0, 1, 2, 3, 4]
+  left_layers = [
+      layers.STANDARD_ATTENTION_REGISTRY_KEY,
+      layers.ATTEND_TO_ENCODER_REGISTRY_KEY,
+      layers.STANDARD_CONV_1X1_REGISTRY_KEY,
+      layers.STANDARD_CONV_1X1_REGISTRY_KEY, layers.IDENTITY_REGISTRY_KEY
+  ]
+  left_output_dims = [512, 512, 1024, 256, 512]
+  right_inputs = [0, 1, 2, 3, 2]
+  right_layers = [
+      layers.IDENTITY_REGISTRY_KEY, layers.IDENTITY_REGISTRY_KEY,
+      layers.STANDARD_CONV_1X1_REGISTRY_KEY,
+      layers.STANDARD_CONV_1X1_REGISTRY_KEY, layers.IDENTITY_REGISTRY_KEY
+  ]
+  right_output_dims = [512, 512, 1024, 256, 512]
+  combiner_functions = [
+      translation_nas_net.ADD_COMBINER_FUNC_KEY,
+      translation_nas_net.ADD_COMBINER_FUNC_KEY,
+      translation_nas_net.CONCAT_COMBINER_FUNC_KEY,
+      translation_nas_net.CONCAT_COMBINER_FUNC_KEY,
+      translation_nas_net.ADD_COMBINER_FUNC_KEY
+  ]
+  dummy_activations = [translation_nas_net.NONE_ACTIVATION_KEY] * 5
+  dummy_norms = [translation_nas_net.NO_NORM_KEY] * 5
+  layer_registry = layers.DECODER_LAYERS
+  is_decoder = True
+  final_combiner_function = translation_nas_net.CONCAT_COMBINER_FUNC_KEY
+
+  return (num_blocks, left_inputs, left_layers, left_output_dims, right_inputs,
+          right_layers, right_output_dims, combiner_functions,
+          final_combiner_function, dummy_activations, dummy_norms,
+          layer_registry, is_decoder)
+
+
+def _add_transformer_branching_hparams(hparams):
+  (encoder_num_blocks, encoder_left_inputs, encoder_left_layers,
+   encoder_left_output_dims, encoder_right_inputs, encoder_right_layers,
+   encoder_right_output_dims, encoder_combiner_functions,
+   encoder_final_combiner_function, encoder_dummy_activations,
+   encoder_dummy_norms, _, _) = _get_transformer_branching_encoder_config()
+
+  # Transformer encoder.
+  hparams.add_hparam("encoder_left_inputs", encoder_left_inputs)
+  hparams.add_hparam("encoder_left_layers", encoder_left_layers)
+  hparams.add_hparam("encoder_left_activations", encoder_dummy_activations)
+  hparams.add_hparam("encoder_left_output_dims", encoder_left_output_dims)
+  hparams.add_hparam("encoder_left_norms", encoder_dummy_norms)
+  hparams.add_hparam("encoder_right_inputs", encoder_right_inputs)
+  hparams.add_hparam("encoder_right_layers", encoder_right_layers)
+  hparams.add_hparam("encoder_right_activations", encoder_dummy_activations)
+  hparams.add_hparam("encoder_right_output_dims", encoder_right_output_dims)
+  hparams.add_hparam("encoder_right_norms", encoder_dummy_norms)
+  hparams.add_hparam("encoder_combiner_functions", encoder_combiner_functions)
+  hparams.add_hparam("encoder_num_blocks", encoder_num_blocks)
+  hparams.add_hparam("encoder_final_combiner_function",
+                     encoder_final_combiner_function)
+
+  (decoder_num_blocks, decoder_left_inputs, decoder_left_layers,
+   decoder_left_output_dims, decoder_right_inputs, decoder_right_layers,
+   decoder_right_output_dims, decoder_combiner_functions,
+   decoder_final_combiner_function, decoder_dummy_activations,
+   decoder_dummy_norms, _, _) = _get_transformer_branching_decoder_config()
+
+  # Transformer decoder.
+  hparams.add_hparam("decoder_left_inputs", decoder_left_inputs)
+  hparams.add_hparam("decoder_left_layers", decoder_left_layers)
+  hparams.add_hparam("decoder_left_activations", decoder_dummy_activations)
+  hparams.add_hparam("decoder_left_output_dims", decoder_left_output_dims)
+  hparams.add_hparam("decoder_left_norms", decoder_dummy_norms)
+  hparams.add_hparam("decoder_right_inputs", decoder_right_inputs)
+  hparams.add_hparam("decoder_right_layers", decoder_right_layers)
+  hparams.add_hparam("decoder_right_activations", decoder_dummy_activations)
+  hparams.add_hparam("decoder_right_output_dims", decoder_right_output_dims)
+  hparams.add_hparam("decoder_right_norms", decoder_dummy_norms)
+  hparams.add_hparam("decoder_combiner_functions", decoder_combiner_functions)
+  hparams.add_hparam("decoder_num_blocks", decoder_num_blocks)
+  hparams.add_hparam("decoder_final_combiner_function",
+                     decoder_final_combiner_function)
+
+
+class TranslationNasNetTest(parameterized.TestCase, tf.test.TestCase):
+
+  def _test_model(self, model_cls, hparams):
+    """Test a Translation Nas Net model."""
+    tf.reset_default_graph()
+
+    hparams.filter_size = 32
+    hparams.num_heads = 1
+    hparams.layer_prepostprocess_dropout = 0.0
+    hparams.hidden_size = _HIDDEN_SIZE
+
+    p_hparams = problem_hparams.test_problem_hparams(_VOCAB_SIZE, _VOCAB_SIZE,
+                                                     hparams)
+    hparams.problems = [p_hparams]
+
+    inputs = -1 + np.random.random_integers(
+        _VOCAB_SIZE, size=(_BATCH_SIZE, _INPUT_LENGTH, 1, 1))
+    targets = -1 + np.random.random_integers(
+        _VOCAB_SIZE, size=(_BATCH_SIZE, _TARGET_LENGTH, 1, 1))
+    features = {
+        "inputs": tf.constant(inputs, dtype=tf.int32, name="inputs"),
+        "targets": tf.constant(targets, dtype=tf.int32, name="targets"),
+        "target_space_id": tf.constant(1, dtype=tf.int32)
+    }
+
+    model = model_cls(hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
+    logits, _ = model(features)
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      res = session.run(logits)
+    self.assertEqual(res.shape,
+                     (_BATCH_SIZE, _TARGET_LENGTH, 1, 1, _VOCAB_SIZE))
+
+  def _get_encoder_hparams(self):
+    hparams = transformer.transformer_small()
+    hparams.add_hparam("encoder_layer_list",
+                       layers.ENCODER_LAYERS.get_layer_names())
+    hparams.add_hparam("encoder_output_dim_list", [32] + [64] *
+                       (len(hparams.encoder_layer_list) - 2) + [32])
+    hparams.add_hparam("encoder_activation_list", ["none"] + ["relu"] *
+                       (len(hparams.encoder_layer_list) - 1))
+    hparams.add_hparam("encoder_norm_list", ["none"] + ["layer_norm"] *
+                       (len(hparams.encoder_layer_list) - 1))
+    return hparams
+
+  def test_nas_seq2seq(self):
+    hparams = self._get_encoder_hparams()
+    _add_transformer_branching_hparams(hparams)
+    self._test_model(translation_nas_net.NasSeq2Seq, hparams)
+
+  def _get_wrong_output_dim_decoder_hparams(self):
+    tf.reset_default_graph()
+
+    hparams = transformer.transformer_base()
+    _add_transformer_branching_hparams(hparams)
+    hparams.num_heads = 1
+    # Purposely scale up the final embedding depth.
+    wrong_output_size = _EMBEDDING_DEPTH + 1
+    hparams.decoder_left_output_dims[
+        -2] = hparams.decoder_left_output_dims[-2] + 1
+    hparams.decoder_left_output_dims[-1] = wrong_output_size
+
+    return hparams, wrong_output_size
+
+  def test_nas_decoder_resizing_output(self):
+    hparams, wrong_size = self._get_wrong_output_dim_decoder_hparams()
+    hparams.enforce_output_size = False
+    input_tensor = tf.zeros([_BATCH_SIZE, _INPUT_LENGTH, _EMBEDDING_DEPTH])
+    decoder_self_attention_bias = (
+        common_attention.attention_bias_lower_triangle(_INPUT_LENGTH))
+    with tf.variable_scope("wrong"):
+      wrong_size_decoder_output = translation_nas_net.nas_decoder(
+          decoder_input=input_tensor,
+          encoder_block_outputs=[input_tensor] * hparams.encoder_num_blocks,
+          decoder_self_attention_bias=decoder_self_attention_bias,
+          encoder_decoder_attention_bias=None,
+          hparams=hparams)
+
+    # Now add the correction.
+    hparams.enforce_output_size = True
+    with tf.variable_scope("correct"):
+      correct_size_decoder_output = translation_nas_net.nas_decoder(
+          decoder_input=input_tensor,
+          encoder_block_outputs=[input_tensor] * hparams.encoder_num_blocks,
+          decoder_self_attention_bias=decoder_self_attention_bias,
+          encoder_decoder_attention_bias=None,
+          hparams=hparams)
+
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      wrong_output, correct_output = session.run(
+          [wrong_size_decoder_output, correct_size_decoder_output])
+    self.assertEqual(wrong_output.shape,
+                     (_BATCH_SIZE, _INPUT_LENGTH, wrong_size))
+    self.assertEqual(correct_output.shape,
+                     (_BATCH_SIZE, _INPUT_LENGTH, _EMBEDDING_DEPTH))
+
+  @parameterized.parameters([(_get_transformer_branching_encoder_config,
+                              [512, 512, 2048, 512, 512]),
+                             (_get_transformer_branching_decoder_config,
+                              [512, 512, 512, 2048, 512, 512])])
+  def test_calculate_branching_model_parameters_transformer(
+      self, get_config, expected_hidden_depths):
+    tf.reset_default_graph()
+
+    (num_blocks, left_inputs, left_layers, left_output_dims, right_inputs,
+     right_layers, right_output_dims, combiner_functions,
+     final_combiner_function, dummy_activations, dummy_norms, layer_registry,
+     is_decoder) = get_config()
+
+    # Get predicted number of parameters.
+    (predicted_num_params, output_size, hidden_depths,
+     _) = translation_nas_net.calculate_branching_model_parameters(
+         encoding_depth=_EMBEDDING_DEPTH,
+         left_inputs=left_inputs,
+         left_layers=left_layers,
+         left_output_dims=left_output_dims,
+         right_inputs=right_inputs,
+         right_layers=right_layers,
+         right_output_dims=right_output_dims,
+         combiner_functions=combiner_functions,
+         final_combiner_function=final_combiner_function,
+         layer_registry=layer_registry,
+         num_blocks=num_blocks,
+         encoder_depth=_EMBEDDING_DEPTH)
+
+    # Create model graph.
+    input_tensor = tf.zeros([32, _INPUT_LENGTH, _EMBEDDING_DEPTH])
+    hparams = transformer.transformer_small()
+
+    if is_decoder:
+      nonpadding = None
+      mask_future = True
+      decoder_self_attention_bias = (
+          common_attention.attention_bias_lower_triangle(_INPUT_LENGTH))
+      encoder_block_outputs = [input_tensor] * 6
+    else:
+      nonpadding = tf.ones([32, _INPUT_LENGTH])
+      mask_future = False
+      decoder_self_attention_bias = None
+      encoder_block_outputs = None
+
+    translation_nas_net.apply_nas_layers(
+        input_tensor=input_tensor,
+        left_inputs=left_inputs,
+        left_layers=left_layers,
+        left_activations=dummy_activations,
+        left_output_dims=left_output_dims,
+        left_norms=dummy_norms,
+        right_inputs=right_inputs,
+        right_layers=right_layers,
+        right_activations=dummy_activations,
+        right_output_dims=right_output_dims,
+        right_norms=dummy_norms,
+        combiner_functions=combiner_functions,
+        final_combiner_function=final_combiner_function,
+        num_blocks=num_blocks,
+        nonpadding=nonpadding,
+        layer_registry=layer_registry,
+        mask_future=mask_future,
+        hparams=hparams,
+        var_scope="test",
+        encoder_decoder_attention_bias=None,
+        encoder_block_outputs=encoder_block_outputs,
+        decoder_self_attention_bias=decoder_self_attention_bias,
+        final_layer_norm=False)
+
+    # Count graph variables.
+    trainable_variables_list = tf.trainable_variables()
+    empirical_num_params = 0
+    for variable_tensor in trainable_variables_list:
+      empirical_num_params += _list_product(variable_tensor.shape.as_list())
+
+    # Compare.
+    self.assertEqual(empirical_num_params, predicted_num_params)
+    self.assertEqual(output_size, _EMBEDDING_DEPTH)
+    self.assertEqual(hidden_depths, expected_hidden_depths)
+
+  @parameterized.parameters([True, False])
+  def test_calculate_branching_model_parameters_decoder_resize(
+      self, enforce_output_size):
+    tf.reset_default_graph()
+
+    hparams, _ = self._get_wrong_output_dim_decoder_hparams()
+    hparams.enforce_output_size = enforce_output_size
+    hparams.decoder_left_norms = [translation_nas_net.NO_NORM_KEY] * 5
+    hparams.decoder_right_norms = [translation_nas_net.NO_NORM_KEY] * 5
+
+    # Get predicted number of parameters.
+    (predicted_num_params, _, _,
+     _) = translation_nas_net.calculate_branching_model_parameters(
+         encoding_depth=_EMBEDDING_DEPTH,
+         left_inputs=hparams.decoder_left_inputs,
+         left_layers=hparams.decoder_left_layers,
+         left_output_dims=hparams.decoder_left_output_dims,
+         right_inputs=hparams.decoder_right_inputs,
+         right_layers=hparams.decoder_right_layers,
+         right_output_dims=hparams.decoder_right_output_dims,
+         combiner_functions=hparams.decoder_combiner_functions,
+         final_combiner_function=hparams.decoder_final_combiner_function,
+         layer_registry=layers.DECODER_LAYERS,
+         num_blocks=hparams.decoder_num_blocks,
+         encoder_depth=_EMBEDDING_DEPTH,
+         enforce_output_size=enforce_output_size)
+
+    # Count graph variables.
+    input_tensor = tf.zeros([_BATCH_SIZE, _INPUT_LENGTH, _EMBEDDING_DEPTH])
+    decoder_self_attention_bias = (
+        common_attention.attention_bias_lower_triangle(_INPUT_LENGTH))
+    _ = translation_nas_net.nas_decoder(
+        decoder_input=input_tensor,
+        encoder_block_outputs=[input_tensor] * hparams.encoder_num_blocks,
+        decoder_self_attention_bias=decoder_self_attention_bias,
+        encoder_decoder_attention_bias=None,
+        hparams=hparams,
+        final_layer_norm=False)
+    trainable_variables_list = tf.trainable_variables()
+    empirical_num_params = 0
+    for variable_tensor in trainable_variables_list:
+      empirical_num_params += _list_product(variable_tensor.shape.as_list())
+
+    self.assertEqual(empirical_num_params, predicted_num_params)
+
+  def test_calculate_branching_model_parameters_output_size_only_final(self):
+    left_inputs = [0, 1, 2, 3]
+    right_inputs = [0, 1, 2, 3]
+    left_output_dims = [1, 10, 100, 1000]
+    right_output_dims = [10000, 100000, 1000000, 10000000]
+    right_layers = [
+        layers.IDENTITY_REGISTRY_KEY, layers.STANDARD_CONV_1X1_REGISTRY_KEY,
+        layers.STANDARD_CONV_1X1_REGISTRY_KEY, layers.IDENTITY_REGISTRY_KEY
+    ]
+    combiner_functions = [
+        translation_nas_net.ADD_COMBINER_FUNC_KEY,
+        translation_nas_net.ADD_COMBINER_FUNC_KEY,
+        translation_nas_net.MULTIPLY_COMBINER_FUNC_KEY,
+        translation_nas_net.CONCAT_COMBINER_FUNC_KEY
+    ]
+
+    (num_blocks, _, left_layers, _, _, _, _, _, final_combiner_function,
+     dummy_activations, dummy_norms, layer_registry,
+     _) = _get_transformer_branching_encoder_config()
+
+    # Get predicted number of parameters.
+    (_, output_size, _,
+     _) = translation_nas_net.calculate_branching_model_parameters(
+         encoding_depth=_EMBEDDING_DEPTH,
+         left_inputs=left_inputs,
+         left_layers=left_layers,
+         left_output_dims=left_output_dims,
+         right_inputs=right_inputs,
+         right_layers=right_layers,
+         right_output_dims=right_output_dims,
+         combiner_functions=combiner_functions,
+         final_combiner_function=final_combiner_function,
+         layer_registry=layer_registry,
+         num_blocks=num_blocks,
+         encoder_depth=_EMBEDDING_DEPTH,
+         enforce_output_size=False,
+         enforce_fixed_output_sizes=False)
+
+    self.assertEqual(output_size, 10001000)
+
+  def test_calculate_branching_model_parameters_output_size_last_two(self):
+    left_inputs = [0, 1, 2, 2]
+    right_inputs = [0, 1, 2, 2]
+    left_output_dims = [1, 10, 100, 1000]
+    right_output_dims = [10000, 100000, 1000000, 10000000]
+    right_layers = [
+        layers.IDENTITY_REGISTRY_KEY, layers.STANDARD_CONV_1X1_REGISTRY_KEY,
+        layers.STANDARD_CONV_1X1_REGISTRY_KEY, layers.IDENTITY_REGISTRY_KEY
+    ]
+    combiner_functions = [
+        translation_nas_net.ADD_COMBINER_FUNC_KEY,
+        translation_nas_net.ADD_COMBINER_FUNC_KEY,
+        translation_nas_net.MULTIPLY_COMBINER_FUNC_KEY,
+        translation_nas_net.CONCAT_COMBINER_FUNC_KEY
+    ]
+
+    (num_blocks, _, left_layers, _, _, _, _, _, final_combiner_function,
+     dummy_activations, dummy_norms, layer_registry,
+     _) = _get_transformer_branching_encoder_config()
+
+    # Get predicted number of parameters.
+    (_, output_size, _,
+     _) = translation_nas_net.calculate_branching_model_parameters(
+         encoding_depth=_EMBEDDING_DEPTH,
+         left_inputs=left_inputs,
+         left_layers=left_layers,
+         left_output_dims=left_output_dims,
+         right_inputs=right_inputs,
+         right_layers=right_layers,
+         right_output_dims=right_output_dims,
+         combiner_functions=combiner_functions,
+         final_combiner_function=final_combiner_function,
+         layer_registry=layer_registry,
+         num_blocks=num_blocks,
+         encoder_depth=_EMBEDDING_DEPTH,
+         enforce_output_size=False,
+         enforce_fixed_output_sizes=False)
+
+    self.assertEqual(output_size, 11001000)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From d49e9a0a2a30174de5225acf276d5ffc19cde100 Mon Sep 17 00:00:00 2001
From: Dumitru Erhan <dumitru@google.com>
Date: Wed, 15 May 2019 15:14:50 -0700
Subject: [PATCH 2036/2720] python3 conversion of glow ops.

PiperOrigin-RevId: 248415239
---
 tensor2tensor/models/research/glow_ops_test.py | 3 +++
 tensor2tensor/models/research/glow_test.py     | 2 ++
 2 files changed, 5 insertions(+)

diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 46914ec56..c77215753 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Lint as: python2, python3
 """Tests for tensor2tensor.models.research.glow_ops."""
 
 from __future__ import absolute_import
@@ -23,6 +24,8 @@
 import tempfile
 from absl.testing import parameterized
 import numpy as np
+from six.moves import range
+from six.moves import zip
 from tensor2tensor.models.research import glow
 from tensor2tensor.models.research import glow_ops
 from tensor2tensor.utils import hparam
diff --git a/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
index 653b23769..e157f53cc 100644
--- a/tensor2tensor/models/research/glow_test.py
+++ b/tensor2tensor/models/research/glow_test.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Lint as: python2, python3
 """Tests for tensor2tensor.models.research.glow_model."""
 
 from __future__ import absolute_import
@@ -22,6 +23,7 @@
 import os
 import tempfile
 import numpy as np
+from six.moves import range
 from tensor2tensor import problems
 from tensor2tensor.data_generators import cifar  # pylint: disable=unused-import
 from tensor2tensor.models.research import glow

From 61889b8d184c9bef4e5bbe93ad085483a5b80ab5 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 17 May 2019 14:34:00 -0700
Subject: [PATCH 2037/2720] Add an option to use distributional RL to the PPO
 code.

PiperOrigin-RevId: 248789235
---
 tensor2tensor/models/research/rl.py    | 62 ++++++++++++++++++++++--
 tensor2tensor/rl/ppo.py                | 67 +++++++++++++++++++++-----
 tensor2tensor/rl/ppo_learner.py        | 46 +++++++++++++-----
 tensor2tensor/rl/trainer_model_free.py |  5 +-
 4 files changed, 151 insertions(+), 29 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 350410a54..5e2f6ef48 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -21,6 +21,7 @@
 import gym
 import six
 
+from tensor2tensor.data_generators import gym_env
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.envs import tic_tac_toe_env
@@ -134,6 +135,14 @@ def ppo_original_params():
   return hparams
 
 
+@registry.register_hparams
+def ppo_dist_params():
+  """Parameters based on the original paper modified for distributional RL."""
+  hparams = ppo_original_params()
+  hparams.learning_rate_constant = 5e-4
+  return hparams
+
+
 @registry.register_hparams
 def ppo_original_tiny():
   """Parameters based on the original PPO paper, tiny version."""
@@ -277,13 +286,14 @@ def make_simulated_env_fn_from_hparams(real_env, hparams, **extra_kwargs):
   )
 
 
-def get_policy(observations, hparams, action_space):
+def get_policy(observations, hparams, action_space, distributional_size=1):
   """Get a policy network.
 
   Args:
     observations: observations
     hparams: parameters
     action_space: action space
+    distributional_size: optional number of buckets for distributional RL
 
   Returns:
     Tuple (action logits, value).
@@ -312,6 +322,9 @@ def get_policy(observations, hparams, action_space):
     num_target_frames = hparams.video_num_target_frames
   except AttributeError:
     num_target_frames = 1
+  target_value_shape_suffix = [num_target_frames]
+  if distributional_size > 1:
+    target_value_shape_suffix = [num_target_frames, distributional_size]
   features = {
       "inputs": observations,
       "input_action": tf.zeros(obs_shape[:2] + [1], dtype=tf.int32),
@@ -324,12 +337,16 @@ def get_policy(observations, hparams, action_space):
       "target_policy": tf.zeros(
           obs_shape[:1] + [num_target_frames] + [action_space.n]),
       "target_value": tf.zeros(
-          obs_shape[:1] + [num_target_frames])
+          obs_shape[:1] + target_value_shape_suffix)
   }
+  model.distributional_value_size = max(distributional_size, 1)
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
     t2t_model.create_dummy_vars()
     (targets, _) = model(features)
-  return (targets["target_policy"][:, 0, :], targets["target_value"][:, 0])
+  target_values = targets["target_value"][:, 0]
+  if distributional_size > 1:
+    target_values = targets["target_value"][:, :]
+  return (targets["target_policy"][:, 0, :], target_values)
 
 
 @registry.register_hparams
@@ -413,6 +430,9 @@ def rlmf_original():
       eval_rl_env_max_episode_steps=1000,
       resize_height_factor=2,
       resize_width_factor=2,
+      distributional_size=1,  # In distributional RL, number of buckets.
+      distributional_subscale=0.04,  # How to scale values to buckets.
+      distributional_threshold=0.0,  # Optimism threshold for experiments.
       grayscale=0,
       rl_env_max_episode_steps=-1,
       # If set, use this as the gym env name, instead of changing game mode etc.
@@ -420,6 +440,7 @@ def rlmf_original():
       # Controls whether we should derive observation space, do some
       # pre-processing etc. See T2TGymEnv._derive_observation_space.
       rl_should_derive_observation_space=True,
+      aunused=0,  # unused param for multi-run settings.
   )
 
 
@@ -452,6 +473,34 @@ def rlmf_base():
   return hparams
 
 
+@registry.register_ranged_hparams
+def rlmf_5runs(rhp):
+  rhp.set_discrete("aunused", list(range(5)))
+
+
+@registry.register_ranged_hparams
+def rlmf_5runs_atari(rhp):
+  rhp.set_categorical("game", gym_env.ATARI_GAMES_WITH_HUMAN_SCORE_NICE)
+  rhp.set_discrete("aunused", list(range(5)))
+
+
+@registry.register_hparams
+def rlmf_dist():
+  """Distributional set of hparams for model-free PPO."""
+  hparams = rlmf_original()
+  hparams.distributional_size = 1024
+  hparams.base_algo_params = "ppo_dist_params"
+  return hparams
+
+
+@registry.register_hparams
+def rlmf_dist_threshold():
+  """Distributional set of hparams for model-free PPO."""
+  hparams = rlmf_dist()
+  hparams.distributional_threshold = 0.5
+  return hparams
+
+
 @registry.register_hparams
 def rlmf_tiny():
   """Tiny set of hparams for model-free PPO."""
@@ -497,6 +546,10 @@ def rlmf_eval():
 
 class PolicyBase(t2t_model.T2TModel):
 
+  def __init__(self, *args, **kwargs):
+    super(PolicyBase, self).__init__(*args, **kwargs)
+    self.distributional_value_size = 1
+
   def loss(self, *args, **kwargs):
     return 0.0
 
@@ -658,8 +711,7 @@ def body(self, features):
       )
       logits = clip_logits(logits, self.hparams)
       logits = tf.expand_dims(logits, axis=1)
-
-      value = tf.layers.dense(x, 1)
+      value = tf.layers.dense(x, self.distributional_value_size)
     return {"target_policy": logits, "target_value": value}
 
 
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index 249be9e0f..21841da77 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -30,7 +30,8 @@
 import tensorflow_probability as tfp
 
 
-def define_ppo_step(data_points, hparams, action_space, lr):
+def define_ppo_step(data_points, hparams, action_space, lr,
+                    distributional_size=1, distributional_subscale=0.04):
   """Define ppo step."""
   observation, action, discounted_reward, norm_advantage, old_pdf = data_points
 
@@ -38,9 +39,9 @@ def define_ppo_step(data_points, hparams, action_space, lr):
   observation = tf.reshape(
       observation, [obs_shape[0] * obs_shape[1]] + obs_shape[2:]
   )
-  (logits, new_value) = get_policy(observation, hparams, action_space)
+  (logits, new_value) = get_policy(observation, hparams, action_space,
+                                   distributional_size=distributional_size)
   logits = tf.reshape(logits, obs_shape[:2] + [action_space.n])
-  new_value = tf.reshape(new_value, obs_shape[:2])
   new_policy_dist = tfp.distributions.Categorical(logits=logits)
 
   new_pdf = new_policy_dist.prob(action)
@@ -53,8 +54,22 @@ def define_ppo_step(data_points, hparams, action_space, lr):
                                    ratio * norm_advantage)
   policy_loss = -tf.reduce_mean(surrogate_objective)
 
-  value_error = new_value - discounted_reward
-  value_loss = hparams.value_loss_coef * tf.reduce_mean(value_error ** 2)
+  if distributional_size > 1:
+    new_value = tf.reshape(new_value, obs_shape[:2] + [distributional_size])
+    new_value = tf.nn.log_softmax(new_value, axis=-1)
+    # We assume the values range from (-half, half) -- set subscale accordingly.
+    half = (distributional_size // 2) * distributional_subscale
+    # To make values integers, we add half (to move range to (0, 2*half) and
+    # then multiply by subscale after which we floor to get nearest int.
+    quantized_dr = tf.floor(
+        (discounted_reward + half) / distributional_subscale)
+    hot_dr = tf.one_hot(tf.cast(quantized_dr, tf.int32), distributional_size)
+    value_loss = - tf.reduce_sum(new_value * hot_dr, axis=-1)
+    value_loss = hparams.value_loss_coef * tf.reduce_mean(value_loss)
+  else:
+    new_value = tf.reshape(new_value, obs_shape[:2])
+    value_error = new_value - discounted_reward
+    value_loss = hparams.value_loss_coef * tf.reduce_mean(value_error ** 2)
 
   entropy = new_policy_dist.entropy()
   entropy_loss = -hparams.entropy_loss_coef * tf.reduce_mean(entropy)
@@ -68,9 +83,29 @@ def define_ppo_step(data_points, hparams, action_space, lr):
     return [tf.identity(x) for x in losses]
 
 
-def define_ppo_epoch(memory, hparams, action_space, batch_size):
+def _distributional_to_value(value_d, size, subscale, threshold):
+  """Get a scalar value out of a value distribution in distributional RL."""
+  half = size // 2
+  value_range = (tf.to_float(tf.range(-half, half)) + 0.5) * subscale
+  probs = tf.nn.softmax(value_d)
+
+  if threshold == 0.0:
+    return tf.reduce_sum(probs * value_range, axis=-1)
+
+  # accumulated_probs[..., i] is the sum of probabilities in buckets upto i
+  # so it is the probability that value <= i'th bucket value
+  accumulated_probs = tf.cumsum(probs, axis=-1)
+  # New probs are 0 on all lower buckets, until the threshold
+  probs = tf.where(accumulated_probs < threshold, tf.zeros_like(probs), probs)
+  probs /= tf.reduce_sum(probs, axis=-1, keepdims=True)  # Re-normalize.
+  return tf.reduce_sum(probs * value_range, axis=-1)
+
+
+def define_ppo_epoch(memory, hparams, action_space, batch_size,
+                     distributional_size=1, distributional_subscale=0.04,
+                     distributional_threshold=0.0):
   """PPO epoch."""
-  observation, reward, done, action, old_pdf, value = memory
+  observation, reward, done, action, old_pdf, value_sm = memory
 
   # This is to avoid propagating gradients through simulated environment.
   observation = tf.stop_gradient(observation)
@@ -79,9 +114,15 @@ def define_ppo_epoch(memory, hparams, action_space, batch_size):
   if hasattr(hparams, "rewards_preprocessing_fun"):
     reward = hparams.rewards_preprocessing_fun(reward)
   done = tf.stop_gradient(done)
-  value = tf.stop_gradient(value)
+  value_sm = tf.stop_gradient(value_sm)
   old_pdf = tf.stop_gradient(old_pdf)
 
+  value = value_sm
+  if distributional_size > 1:
+    value = _distributional_to_value(
+        value_sm, distributional_size, distributional_subscale,
+        distributional_threshold)
+
   advantage = calculate_generalized_advantage_estimator(
       reward, value, done, hparams.gae_gamma, hparams.gae_lambda)
 
@@ -117,10 +158,12 @@ def define_ppo_epoch(memory, hparams, action_space, batch_size):
 
   ppo_step_rets = tf.scan(
       lambda a, i: add_lists_elementwise(  # pylint: disable=g-long-lambda
-          a, define_ppo_step([tf.gather(t, indices_of_batches[i, :])
-                              for t in input_tensors],
-                             hparams, action_space, lr
-                            )),
+          a, define_ppo_step(
+              [tf.gather(t, indices_of_batches[i, :]) for t in input_tensors],
+              hparams, action_space, lr,
+              distributional_size=distributional_size,
+              distributional_subscale=distributional_subscale
+          )),
       tf.range(number_of_batches),
       [0., 0., 0.],
       parallel_iterations=1)
diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 86b78ea9a..a7b6a166c 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -38,10 +38,15 @@
 class PPOLearner(PolicyLearner):
   """PPO for policy learning."""
 
-  def __init__(self, *args, **kwargs):
-    super(PPOLearner, self).__init__(*args, **kwargs)
+  def __init__(self, frame_stack_size, base_event_dir, agent_model_dir,
+               total_num_epochs, **kwargs):
+    super(PPOLearner, self).__init__(
+        frame_stack_size, base_event_dir, agent_model_dir, total_num_epochs)
     self._num_completed_iterations = 0
     self._lr_decay_start = None
+    self._distributional_size = kwargs.get("distributional_size", 1)
+    self._distributional_subscale = kwargs.get("distributional_subscale", 0.04)
+    self._distributional_threshold = kwargs.get("distributional_threshold", 0.0)
 
   def train(self,
             env_fn,
@@ -81,6 +86,9 @@ def train(self,
                   hparams,
                   eval_env_fn,
                   sampling_temp,
+                  distributional_size=self._distributional_size,
+                  distributional_subscale=self._distributional_subscale,
+                  distributional_threshold=self._distributional_threshold,
                   frame_stack_size=self.frame_stack_size,
                   force_beginning_resets=simulated))
 
@@ -134,6 +142,7 @@ def evaluate(self, env_fn, hparams, sampling_temp):
             frame_stack_size=self.frame_stack_size,
             force_beginning_resets=False,
             sampling_temp=sampling_temp,
+            distributional_size=self._distributional_size,
         )
         model_saver = tf.train.Saver(
             tf.global_variables(hparams.policy_network + "/.*")
@@ -153,6 +162,9 @@ def _define_train(
     ppo_hparams,
     eval_env_fn=None,
     sampling_temp=1.0,
+    distributional_size=1,
+    distributional_subscale=0.04,
+    distributional_threshold=0.0,
     **collect_kwargs
 ):
   """Define the training setup."""
@@ -163,9 +175,13 @@ def _define_train(
           "ppo_train",
           eval_phase=False,
           sampling_temp=sampling_temp,
+          distributional_size=distributional_size,
           **collect_kwargs))
   ppo_summary = ppo.define_ppo_epoch(
-      memory, ppo_hparams, train_env.action_space, train_env.batch_size)
+      memory, ppo_hparams, train_env.action_space, train_env.batch_size,
+      distributional_size=distributional_size,
+      distributional_subscale=distributional_subscale,
+      distributional_threshold=distributional_threshold)
   train_summary = tf.summary.merge([collect_summary, ppo_summary])
 
   if ppo_hparams.eval_every_epochs:
@@ -179,6 +195,7 @@ def _define_train(
             "ppo_eval",
             eval_phase=True,
             sampling_temp=0.0,
+            distributional_size=distributional_size,
             **collect_kwargs))
     return (train_summary, eval_collect_summary, (train_initialization,
                                                   eval_initialization))
@@ -251,10 +268,13 @@ def _run_train(ppo_hparams,
             model_save_fn(model_dir)
 
 
-def _rollout_metadata(batch_env):
+def _rollout_metadata(batch_env, distributional_size=1):
   """Metadata for rollouts."""
   batch_env_shape = batch_env.observ.get_shape().as_list()
   batch_size = [batch_env_shape[0]]
+  value_size = batch_size
+  if distributional_size > 1:
+    value_size = batch_size + [distributional_size]
   shapes_types_names = [
       # TODO(piotrmilos): possibly retrieve the observation type for batch_env
       (batch_size + batch_env_shape[1:], batch_env.observ_dtype, "observation"),
@@ -263,7 +283,7 @@ def _rollout_metadata(batch_env):
       (batch_size + list(batch_env.action_shape), batch_env.action_dtype,
        "action"),
       (batch_size, tf.float32, "pdf"),
-      (batch_size, tf.float32, "value_function"),
+      (value_size, tf.float32, "value_function"),
   ]
   return shapes_types_names
 
@@ -308,7 +328,8 @@ def simulate(self, action):
 
 
 def _define_collect(batch_env, ppo_hparams, scope, frame_stack_size, eval_phase,
-                    sampling_temp, force_beginning_resets):
+                    sampling_temp, force_beginning_resets,
+                    distributional_size=1):
   """Collect trajectories.
 
   Args:
@@ -319,6 +340,7 @@ def _define_collect(batch_env, ppo_hparams, scope, frame_stack_size, eval_phase,
     eval_phase: TODO(koz4k): Write docstring.
     sampling_temp: Sampling temperature for the policy.
     force_beginning_resets: Whether to reset at the beginning of each episode.
+    distributional_size: optional, number of buckets in distributional RL.
 
   Returns:
     Returns memory (observations, rewards, dones, actions,
@@ -343,7 +365,7 @@ def _define_collect(batch_env, ppo_hparams, scope, frame_stack_size, eval_phase,
       batch_env = w[0](batch_env, **w[1])
       to_initialize.append(batch_env)
 
-    rollout_metadata = _rollout_metadata(batch_env)
+    rollout_metadata = _rollout_metadata(batch_env, distributional_size)
     speculum = batch_env.speculum
 
     def initialization_lambda(sess):
@@ -389,12 +411,15 @@ def step(index, scores_sum, scores_num):
       # operation. We are waiting for tf.copy:
       # https://github.com/tensorflow/tensorflow/issues/11186
       obs_copy = batch_env.observ + 0
+      value_fun_shape = (num_agents,)
+      if distributional_size > 1:
+        value_fun_shape = (num_agents, distributional_size)
 
       def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
         """Step of the environment."""
 
         (logits, value_function) = get_policy(
-            obs_copy, ppo_hparams, batch_env.action_space
+            obs_copy, ppo_hparams, batch_env.action_space, distributional_size
         )
         action = common_layers.sample_with_temperature(logits, sampling_temp)
         action = tf.cast(action, tf.int32)
@@ -404,7 +429,7 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
 
         pdf = tfp.distributions.Categorical(logits=logits).prob(action)
         pdf = tf.reshape(pdf, shape=(num_agents,))
-        value_function = tf.reshape(value_function, shape=(num_agents,))
+        value_function = tf.reshape(value_function, shape=value_fun_shape)
         done = tf.reshape(done, shape=(num_agents,))
 
         with tf.control_dependencies([reward, done]):
@@ -418,7 +443,7 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
           env_step,
           [
               tf.constant(0.0, shape=(num_agents,)),
-              tf.constant(0.0, shape=(num_agents,)),
+              tf.constant(0.0, shape=value_fun_shape),
               tf.constant(False, shape=(num_agents,))
           ],
           parallel_iterations=1,
@@ -427,7 +452,6 @@ def env_step(arg1, arg2, arg3):  # pylint: disable=unused-argument
 
       with tf.control_dependencies([pdf, value_function]):
         obs, reward, done, action = speculum.dequeue()
-
         to_save = [obs, reward, done, action, pdf, value_function]
         save_ops = [
             tf.scatter_update(memory_slot, index, value)
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index fdcc6ec3b..a6cbafb70 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -90,7 +90,10 @@ def train(hparams, output_dir, env_problem_name, report_fn=None):
                   misc_utils.pprint_hparams(hparams))
   tf.logging.vlog(1, "Using hparams.base_algo: %s", hparams.base_algo)
   learner = rl_utils.LEARNERS[hparams.base_algo](
-      hparams.frame_stack_size, output_dir, output_dir, total_num_epochs=1
+      hparams.frame_stack_size, output_dir, output_dir, total_num_epochs=1,
+      distributional_size=hparams.get("distributional_size", 1),
+      distributional_subscale=hparams.get("distributional_subscale", 0.04),
+      distributional_threshold=hparams.get("distributional_threshold", 0.0),
   )
 
   policy_hparams = trainer_lib.create_hparams(hparams.base_algo_params)

From 019109d5984ae743b447c04583a78475eced10e5 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 20 May 2019 10:51:18 -0700
Subject: [PATCH 2038/2720] Add an output_dtype to RenderedEnv, the processed
 observation will be of this type. Done since TPUs don't support uint8.

PiperOrigin-RevId: 249081113
---
 tensor2tensor/envs/mujoco_problems.py |  3 ++-
 tensor2tensor/rl/gym_utils.py         | 27 +++++++++++++++++++++------
 tensor2tensor/rl/gym_utils_test.py    |  5 +++++
 tensor2tensor/trax/rlax/ppo_main.py   |  3 ++-
 4 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/envs/mujoco_problems.py b/tensor2tensor/envs/mujoco_problems.py
index 927321d1d..928b392e3 100644
--- a/tensor2tensor/envs/mujoco_problems.py
+++ b/tensor2tensor/envs/mujoco_problems.py
@@ -39,7 +39,8 @@ def __init__(self):
             "maxskip_env": False,
             "rendered_env": True,
             "rendered_env_resize_to": None,  # Do not resize frames
-            "sticky_actions": False
+            "sticky_actions": False,
+            "output_dtype": None,
         })
     super(ReacherEnvProblem, self).__init__(
         base_env_name=base_env_name, env_wrapper_fn=wrapper_fn)
diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index af31020d9..89ae3698f 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -83,13 +83,20 @@ def reset(self, **kwargs):
 class RenderedEnv(gym.Wrapper):
   """Simple Env wrapper to override observations with rendered rgb values."""
 
-  def __init__(self, env, mode="rgb_array", low=0, high=255, resize_to=None):
+  def __init__(self,
+               env,
+               mode="rgb_array",
+               low=0,
+               high=255,
+               resize_to=None,
+               output_dtype=None):
     gym.Wrapper.__init__(self, env)
     # Get a sample frame to correctly set observation space
     self.mode = mode
     sample_frame = self.render(mode=self.mode)
     assert sample_frame is not None
     self.should_resize = False
+    self.output_dtype = output_dtype
     if resize_to is None:
       self.observation_space = gym.spaces.Box(
           low=low,
@@ -112,7 +119,9 @@ def _maybe_resize(self, obs):
     height, width = self.observation_space.shape[:2]
     img = Image.fromarray(obs)
     img = img.resize([width, height], resample=Image.ANTIALIAS)
-    return np.array(img)
+    if self.output_dtype is None:
+      return np.array(img)
+    return np.array(img).astype(self.output_dtype)
 
   def step(self, action):
     _, reward, done, info = self.env.step(action)
@@ -148,7 +157,7 @@ def remove_time_limit_wrapper(env):
 
 
 def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env,
-                    rendered_env_resize_to, sticky_actions):
+                    rendered_env_resize_to, sticky_actions, output_dtype):
   """Wraps a gym environment. see make_gym_env for details."""
   # rl_env_max_episode_steps is None or int.
   assert ((not rl_env_max_episode_steps) or
@@ -167,7 +176,8 @@ def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env,
     env = MaxAndSkipEnv(env)  # pylint: disable=redefined-variable-type
 
   if rendered_env:
-    env = RenderedEnv(env, resize_to=rendered_env_resize_to)
+    env = RenderedEnv(
+        env, resize_to=rendered_env_resize_to, output_dtype=output_dtype)
 
   if wrap_with_time_limit:
     env = gym.wrappers.TimeLimit(
@@ -180,7 +190,8 @@ def make_gym_env(name,
                  maxskip_env=False,
                  rendered_env=False,
                  rendered_env_resize_to=None,
-                 sticky_actions=False):
+                 sticky_actions=False,
+                 output_dtype=None):
   """Create a gym env optionally with a time limit and maxskip wrapper.
 
   NOTE: The returned env may already be wrapped with TimeLimit!
@@ -196,13 +207,17 @@ def make_gym_env(name,
     rendered_env_resize_to: a list of [height, width] to change the original
       resolution of the native environment render.
     sticky_actions: whether to use sticky_actions before MaxAndSkip wrapper.
+    output_dtype: numpy datatype that we want the observation to be in, if None
+      this defaults to the env's observation dtype. Useful for TPUs since they
+      don't support uint8 which is a default observation type for a lot of envs.
 
   Returns:
     An instance of `gym.Env` or `gym.Wrapper`.
   """
   env = gym.make(name)
   return gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env,
-                         rendered_env, rendered_env_resize_to, sticky_actions)
+                         rendered_env, rendered_env_resize_to, sticky_actions,
+                         output_dtype)
 
 
 def register_gym_env(class_entry_point, version="v0", kwargs=None):
diff --git a/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index 47d602922..120ac0c30 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -88,6 +88,11 @@ def test_rendered_env(self):
     obs, _, _, _ = env.step(1)
     self.assertTrue(np.allclose(np.zeros([64, 12, 3], np.uint8), obs))
 
+    env = gym_utils.RenderedEnv(SimpleEnv(), resize_to=(64, 12),
+                                output_dtype=np.float32)
+    obs, _, _, _ = env.step(1)
+    self.assertTrue(np.allclose(np.zeros([64, 12, 3], np.float32), obs))
+
   def test_gym_registration(self):
     reg_id, env = gym_utils.register_gym_env(
         "tensor2tensor.rl.gym_utils_test:SimpleEnv")
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index e453210b3..617dd21de 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -140,7 +140,8 @@ def make_env():
           "maxskip_env": True,
           "rendered_env": True,
           "rendered_env_resize_to": (FLAGS.resized_height, FLAGS.resized_width),
-          "sticky_actions": False
+          "sticky_actions": False,
+          "output_dtype": None,
       })
 
   return rendered_env_problem.RenderedEnvProblem(

From b3ad0ad352ed1c28c0da20760aeab34513d5b5b2 Mon Sep 17 00:00:00 2001
From: Shawn Simister <simister@google.com>
Date: Mon, 20 May 2019 13:19:36 -0700
Subject: [PATCH 2039/2720] Fixing control dependency which caused flaky tests
 when RNN steps ran in parallel.

PiperOrigin-RevId: 249112369
---
 tensor2tensor/models/research/neural_stack.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/research/neural_stack.py b/tensor2tensor/models/research/neural_stack.py
index ac6af5995..b347e6032 100644
--- a/tensor2tensor/models/research/neural_stack.py
+++ b/tensor2tensor/models/research/neural_stack.py
@@ -286,9 +286,10 @@ def call(self, inputs, state):
     ], axis=1)
 
     # Call the controller and get controller interface values.
-    (push_strengths, pop_strengths,
-     write_values, outputs, controller_state) = self.call_controller(
-         controller_inputs, controller_state, batch_size)
+    with tf.control_dependencies([read_strengths]):
+      (push_strengths, pop_strengths,
+       write_values, outputs, controller_state) = self.call_controller(
+           controller_inputs, controller_state, batch_size)
 
     # Always write input values to memory regardless of push strength.
     # See Equation-1 in Grefenstette et al., 2015.

From a9161f0a1f1a9f1326fc667f35ef80e614102724 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 20 May 2019 13:28:39 -0700
Subject: [PATCH 2040/2720] Changes for handling trajectory truncation.

EnvProblem - Add a method, `truncate` that calls truncate on the Trajectory.

env_problem_utils - Now if a trajectory reaches `max_timestep`, we don't call `reset` and let it continue. `reset` is only called when env returns done.

BatchTrajectory - Add a method, `truncate_trajectories` that puts the truncated trajectories in the completed bin, and adds new trajectories with the same starting observation and the truncated trajectory's last observation.

PiperOrigin-RevId: 249114001
---
 tensor2tensor/envs/env_problem.py       |  8 +++++++
 tensor2tensor/envs/env_problem_utils.py | 25 ++++++++++++++-----
 tensor2tensor/envs/trajectory.py        | 32 +++++++++++++++++++++++++
 tensor2tensor/envs/trajectory_test.py   | 25 +++++++++++++++++++
 4 files changed, 84 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 1f589535a..7aacec777 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -471,6 +471,14 @@ def _reset(self, indices):
     # rest being the dimensionality of the observation.
     return np.stack([self._envs[index].reset() for index in indices])
 
+  def truncate(self, indices=None):
+    """Truncates trajectories at the specified indices."""
+
+    if indices is None:
+      indices = np.arange(self.batch_size)
+
+    self.trajectories.truncate_trajectories(indices)
+
   def reset(self, indices=None):
     """Resets environments at given indices.
 
diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 44a4e5133..0e61911cc 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -51,6 +51,7 @@ def play_env_problem_with_policy(env,
                                  num_trajectories=1,
                                  max_timestep=None,
                                  boundary=20,
+                                 idx=0,
                                  rng=None):
   """Plays the given env with the policy function to collect trajectories.
 
@@ -60,11 +61,13 @@ def play_env_problem_with_policy(env,
         back log-probabilities (B, T, A).
     num_trajectories: int, number of trajectories to collect.
     max_timestep: int or None, if not None or a negative number, we cut any
-        trajectory that exceeds this time and mark that as completed by
-        resetting that trajectory.
+        trajectory that exceeds this time put it in the completed bin, and
+        *dont* reset the env.
     boundary: this is the bucket length, we pad the observations to integer
         multiples of this + 1 and then feed the padded observations to the
         policy_fun.
+    idx: int, index on the number of times this function is being called, we may
+        want to reset only when idx == 0 for instance.
     rng: jax rng, splittable.
 
   Returns:
@@ -84,10 +87,14 @@ def multinomial_sample(probs):
     """
     return int(np.argwhere(np.random.multinomial(1, probs) == 1))
 
-  # We need to reset all environments.
-  env.reset()
+  # We need to reset all environments, if we're coming here the first time.
+  if idx == 0 or max_timestep is None or max_timestep <= 0:
+    env.reset()
+  else:
+    # Clear completed trajectories held internally.
+    env.trajectories.clear_completed_trajectories()
 
-  while True:
+  while env.trajectories.num_completed_trajectories < num_trajectories:
     # Get all the observations for all the active trajectories.
     # Shape is (B, T) + OBS
     padded_observations = env.trajectories.observations_np(boundary=boundary)
@@ -147,7 +154,10 @@ def multinomial_sample(probs):
 
     # If so, reset these as well.
     if exceeded_time_limit_idxs.size:
-      env.reset(indices=exceeded_time_limit_idxs)
+      # This just cuts the trajectory, doesn't reset the env, so it continues
+      # from where it left off.
+      env.truncate(indices=exceeded_time_limit_idxs)
+
     # Do we have enough trajectories right now?
     if env.trajectories.num_completed_trajectories >= num_trajectories:
       break
@@ -158,4 +168,7 @@ def multinomial_sample(probs):
   for trajectory in env.trajectories.completed_trajectories[:num_trajectories]:
     completed_trajectories.append(trajectory.as_numpy)
 
+  # Keep the rest of the trajectories, if any, in our kitty.
+  env.trajectories.clear_completed_trajectories(num=num_trajectories)
+
   return completed_trajectories
diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index 451001e18..412fca063 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -144,6 +144,13 @@ def trajectories(self):
   def completed_trajectories(self):
     return self._completed_trajectories
 
+  def clear_completed_trajectories(self, num=None):
+    """Clear the first `num` completed trajectories, or all if num is None."""
+    if num is None:
+      self._completed_trajectories = []
+    else:
+      self._completed_trajectories = self._completed_trajectories[num:]
+
   def _complete_trajectory(self, trajectory, index):
     """Completes the given trajectory at the given index."""
 
@@ -158,6 +165,31 @@ def _complete_trajectory(self, trajectory, index):
     # Make a new one to replace it.
     self._trajectories[index] = Trajectory()
 
+  def truncate_trajectories(self, indices):
+    """Truncate trajectories at specified indices.
+
+     This puts the truncated trajectories in the completed list and makes new
+     trajectories with the observation from the trajectory that was truncated at
+     the same index.
+
+    Args:
+        indices: iterable with the indices to truncate.
+    """
+    observations = []
+    for index in indices:
+      trajectory = self._trajectories[index]
+      assert trajectory.is_active, "Trajectory to truncate can't be inactive."
+
+      # NOTE: We don't mark the last time-step as done.
+
+      # Collect the observations.
+      observations.append(trajectory.last_time_step.observation)
+
+    # Call reset on these indices, this will make new trajectories with the same
+    # observation as the existing ones, but in new trajectories. The existing
+    # trajectories are marked as completed.
+    self.reset(indices, np.stack(observations))
+
   def reset(self, indices, observations):
     """Resets trajectories at given indices and populates observations.
 
diff --git a/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index e583cf80d..b4d2083d3 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -194,6 +194,31 @@ def test_reset_some(self):
     # Nothing is done anyways.
     self.assertEqual(0, bt.num_completed_trajectories)
 
+  def test_truncate(self):
+    batch_size = 1
+    bt = trajectory.BatchTrajectory(batch_size=batch_size)
+
+    indices = np.arange(batch_size)
+    observations, _, _, _ = (
+        self.get_random_observations_rewards_actions_dones(
+            batch_size=batch_size))
+
+    # Have to call reset first.
+    bt.reset(indices, observations)
+
+    self.assertEqual(0, bt.num_completed_trajectories)
+
+    bt.truncate_trajectories(indices)
+
+    self.assertEqual(batch_size, bt.num_completed_trajectories)
+
+    # Assert they are all active, since the last observation was duplicated.
+    self.assertTrue(all(t.is_active for t in bt.trajectories))
+
+    # Test that the observation is the same.
+    self.assertAllEqual(bt.trajectories[0].last_time_step.observation,
+                        bt.completed_trajectories[0].last_time_step.observation)
+
   def test_step(self):
     bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
 

From af2519fd3dbabc97b20465d6bff6565afaae3f39 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 20 May 2019 15:59:19 -0700
Subject: [PATCH 2041/2720] Small time-stacked CNN policy for Atari.

PiperOrigin-RevId: 249144422
---
 tensor2tensor/trax/models/atari_cnn.py      | 49 +++++++++++++++++++++
 tensor2tensor/trax/models/atari_cnn_test.py | 49 +++++++++++++++++++++
 2 files changed, 98 insertions(+)
 create mode 100644 tensor2tensor/trax/models/atari_cnn.py
 create mode 100644 tensor2tensor/trax/models/atari_cnn_test.py

diff --git a/tensor2tensor/trax/models/atari_cnn.py b/tensor2tensor/trax/models/atari_cnn.py
new file mode 100644
index 000000000..3d38c9584
--- /dev/null
+++ b/tensor2tensor/trax/models/atari_cnn.py
@@ -0,0 +1,49 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 net for playing Atari games using PPO."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.trax import layers as tl
+
+
+def AtariCnn(hidden_sizes=(32, 32), output_size=128):
+  # Input's shape = (B, T, H, W, C)
+  return tl.Serial(
+      tl.Div(divisor=255.0),
+      # Have 4 copies of the input, each one shifted to the right by one.
+      tl.Branch(tl.Copy(), tl.ShiftRight(),
+                tl.Serial(
+                    tl.ShiftRight(),
+                    tl.ShiftRight(),
+                ), tl.Serial(
+                    tl.ShiftRight(),
+                    tl.ShiftRight(),
+                    tl.ShiftRight(),
+                )),
+      # Concatenated on the last axis.
+      tl.Concatenate(axis=-1),  # (B, T, H, W, 4C)
+      tl.Rebatch(tl.Conv(hidden_sizes[0], (5, 5), (2, 2), 'SAME'), 2),
+      tl.Relu(),
+      tl.Rebatch(tl.Conv(hidden_sizes[1], (5, 5), (2, 2), 'SAME'), 2),
+      tl.Relu(),
+      tl.Flatten(num_axis_to_keep=2),  # B, T and rest.
+      tl.Dense(output_size),
+      tl.Relu(),
+      # Eventually this is shaped (B, T, output_size)
+  )
diff --git a/tensor2tensor/trax/models/atari_cnn_test.py b/tensor2tensor/trax/models/atari_cnn_test.py
new file mode 100644
index 000000000..7b36db92a
--- /dev/null
+++ b/tensor2tensor/trax/models/atari_cnn_test.py
@@ -0,0 +1,49 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.models.atari_cnn."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+import operator as op
+import numpy as onp
+from tensor2tensor.trax.backend import random as jax_random
+from tensor2tensor.trax.models import atari_cnn
+from tensorflow import test
+
+
+class AtariCnnTest(test.TestCase):
+
+  def test_computes(self):
+    rng_key = jax_random.get_prng(0)
+    hidden_size = (4, 4)
+    output_size = 6
+    policy = atari_cnn.AtariCnn(
+        hidden_sizes=hidden_size, output_size=output_size)
+    B, T, OBS = 2, 2, (28, 28, 3)  # pylint: disable=invalid-name
+    rng_key, key = jax_random.split(rng_key)
+    params = policy.initialize((-1, -1) + OBS, key)
+    x = onp.arange(B * (T + 1) * functools.reduce(op.mul, OBS)).reshape(
+        B, T + 1, *OBS)
+    rng_key, key = jax_random.split(rng_key)
+    y = policy(x, params, rng=key)
+    self.assertEqual((B, T + 1, output_size), y.shape)
+
+
+if __name__ == "__main__":
+  test.main()

From 7998502a152fe3725d2ebc92cfb9459023dca418 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 20 May 2019 16:04:08 -0700
Subject: [PATCH 2042/2720] Checkpointing, Eval and a few other changes in PPO.

Other changes being:
- Exact entropy instead of approximate entropy.
- Option added for a One vs Two Tower architecture in the combined network.
- Write average train and eval rewards in output_dir.
- Add flag to enable/disable early stopping.

PiperOrigin-RevId: 249145562
---
 tensor2tensor/trax/rlax/ppo.py      | 230 ++++++++++++++++++----------
 tensor2tensor/trax/rlax/ppo_main.py |  85 ++++++++--
 tensor2tensor/trax/rlax/ppo_test.py |  60 ++------
 3 files changed, 235 insertions(+), 140 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index da1095597..c1d71297a 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -50,6 +50,8 @@
 from __future__ import print_function
 
 import functools
+import os
+import pickle
 import time
 
 from absl import logging
@@ -65,6 +67,7 @@
 from tensor2tensor.trax import layers
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
+from tensorflow.io import gfile
 
 DEBUG_LOGGING = False
 GAMMA = 0.99
@@ -115,7 +118,8 @@ def value_net(rng_key,
 def policy_and_value_net(rng_key,
                          batch_observations_shape,
                          num_actions,
-                         bottom_layers_fn=None):
+                         bottom_layers_fn=None,
+                         two_towers=True):
   """A policy and value net function."""
 
   # Layers.
@@ -124,16 +128,27 @@ def policy_and_value_net(rng_key,
   # other computes the value function.
   # NOTE: The LogSoftmax instead of the Softmax because of numerical stability.
 
-  tower1 = [] if bottom_layers_fn is None else bottom_layers_fn()
-  tower2 = [] if bottom_layers_fn is None else bottom_layers_fn()
+  net = None
+  if not two_towers:
+    tower = [] if bottom_layers_fn is None else bottom_layers_fn()
+    tower.extend([
+        layers.Branch(
+            layers.Serial(layers.Dense(num_actions), layers.LogSoftmax()),
+            layers.Dense(1))
+    ])
+    net = layers.Serial(*tower)
+  else:
+    tower1 = [] if bottom_layers_fn is None else bottom_layers_fn()
+    tower2 = [] if bottom_layers_fn is None else bottom_layers_fn()
 
-  tower1.extend([layers.Dense(num_actions), layers.LogSoftmax()])
-  tower2.extend([layers.Dense(1)])
+    tower1.extend([layers.Dense(num_actions), layers.LogSoftmax()])
+    tower2.extend([layers.Dense(1)])
 
-  net = layers.Branch(
-      layers.Serial(*tower1),
-      layers.Serial(*tower2),
-  )
+    net = layers.Branch(
+        layers.Serial(*tower1),
+        layers.Serial(*tower2),
+    )
+  assert net
   return net.initialize(batch_observations_shape, rng_key), net
 
 
@@ -171,6 +186,7 @@ def collect_trajectories(env,
                          max_timestep=None,
                          boundary=20,
                          epsilon=0.1,
+                         idx=0,
                          rng=None):
   """Collect trajectories with the given policy net and behaviour.
 
@@ -185,6 +201,7 @@ def collect_trajectories(env,
       done.
     boundary: int, boundary for padding, used in EnvProblem envs.
     epsilon: float, the epsilon for `epsilon-greedy` policy.
+    idx: int, index on the number of times this function is being called
     rng: jax rng, splittable.
 
   Returns:
@@ -203,6 +220,7 @@ def collect_trajectories(env,
         num_trajectories=num_trajectories,
         max_timestep=max_timestep,
         boundary=boundary,
+        idx=idx,
         rng=rng)
 
   trajectories = []
@@ -466,8 +484,8 @@ def value_loss(value_net_apply,
     gamma: float, discount factor.
     epsilon: float, clip-fraction, used if value_value_prediction_old isn't None
     value_prediction_old: np.ndarray of shape (B, T+1, 1) of value predictions
-        using the old parameters. If provided, we incorporate this in the loss
-        as well. This is from the OpenAI baselines implementation.
+      using the old parameters. If provided, we incorporate this in the loss as
+      well. This is from the OpenAI baselines implementation.
     rng: jax rng, splittable.
 
   Returns:
@@ -481,9 +499,13 @@ def value_loss(value_net_apply,
   value_prediction = value_net_apply(observations, value_net_params, rng=rng)
   assert (B, T + 1, 1) == value_prediction.shape
 
-  return value_loss_given_predictions(value_prediction, rewards, reward_mask,
-                                      gamma, epsilon=epsilon,
-                                      value_prediction_old=value_prediction_old)
+  return value_loss_given_predictions(
+      value_prediction,
+      rewards,
+      reward_mask,
+      gamma,
+      epsilon=epsilon,
+      value_prediction_old=value_prediction_old)
 
 
 @jit
@@ -502,8 +524,8 @@ def value_loss_given_predictions(value_prediction,
     gamma: float, discount factor.
     epsilon: float, clip-fraction, used if value_value_prediction_old isn't None
     value_prediction_old: np.ndarray of shape (B, T+1, 1) of value predictions
-        using the old parameters. If provided, we incorporate this in the loss
-        as well. This is from the OpenAI baselines implementation.
+      using the old parameters. If provided, we incorporate this in the loss as
+      well. This is from the OpenAI baselines implementation.
 
   Returns:
     The average L2 value loss, averaged over instances where reward_mask is 1.
@@ -672,21 +694,21 @@ def ppo_loss(policy_net_apply,
   # Compute predicted log-probs and hand over to `ppo_loss_given_predictions`.
 
   # log_probab_actions_{old,new} are both (B, T+1, A)
-  log_probab_actions_new = policy_net_apply(padded_observations,
-                                            new_policy_params,
-                                            rng=rng)
+  log_probab_actions_new = policy_net_apply(
+      padded_observations, new_policy_params, rng=rng)
   assert (B, T + 1) == log_probab_actions_new.shape[:2]
   assert log_probab_actions_old.shape[-1] == log_probab_actions_new.shape[-1]
 
-  return ppo_loss_given_predictions(log_probab_actions_new,
-                                    log_probab_actions_old,
-                                    value_predictions_old,
-                                    padded_actions,
-                                    padded_rewards,
-                                    reward_mask,
-                                    gamma=gamma,
-                                    lambda_=lambda_,
-                                    epsilon=epsilon)
+  return ppo_loss_given_predictions(
+      log_probab_actions_new,
+      log_probab_actions_old,
+      value_predictions_old,
+      padded_actions,
+      padded_rewards,
+      reward_mask,
+      gamma=gamma,
+      lambda_=lambda_,
+      epsilon=epsilon)
 
 
 @jit
@@ -755,8 +777,12 @@ def combined_loss_given_predictions(log_probab_actions_new,
                                     c2=0.01):
   """Computes the combined (clipped loss + value loss) given predictions."""
   loss_value = value_loss_given_predictions(
-      value_prediction_new, padded_rewards, reward_mask, gamma=gamma,
-      value_prediction_old=value_prediction_old, epsilon=epsilon)
+      value_prediction_new,
+      padded_rewards,
+      reward_mask,
+      gamma=gamma,
+      value_prediction_old=value_prediction_old,
+      epsilon=epsilon)
   loss_ppo = ppo_loss_given_predictions(
       log_probab_actions_new,
       log_probab_actions_old,
@@ -767,7 +793,7 @@ def combined_loss_given_predictions(log_probab_actions_new,
       gamma=gamma,
       lambda_=lambda_,
       epsilon=epsilon)
-  entropy_bonus = approximate_entropy(log_probab_actions_new, reward_mask)
+  entropy_bonus = masked_entropy(log_probab_actions_new, reward_mask)
   return (loss_ppo + (c1 * loss_value) - (c2 * entropy_bonus), loss_ppo,
           loss_value, entropy_bonus)
 
@@ -792,18 +818,19 @@ def combined_loss(new_params,
       padded_observations, new_params, rng=rng)
 
   # (combined_loss, ppo_loss, value_loss, entropy_bonus)
-  return combined_loss_given_predictions(log_probab_actions_new,
-                                         log_probab_actions_old,
-                                         value_predictions_new,
-                                         value_predictions_old,
-                                         padded_actions,
-                                         padded_rewards,
-                                         reward_mask,
-                                         gamma=gamma,
-                                         lambda_=lambda_,
-                                         epsilon=epsilon,
-                                         c1=c1,
-                                         c2=c2)
+  return combined_loss_given_predictions(
+      log_probab_actions_new,
+      log_probab_actions_old,
+      value_predictions_new,
+      value_predictions_old,
+      padded_actions,
+      padded_rewards,
+      reward_mask,
+      gamma=gamma,
+      lambda_=lambda_,
+      epsilon=epsilon,
+      c1=c1,
+      c2=c2)
 
 
 @functools.partial(jit, static_argnums=(2, 3, 4))
@@ -937,22 +964,34 @@ def approximate_kl(log_prob_new, log_prob_old, mask):
   return np.sum(diff) / np.sum(mask)
 
 
-def approximate_entropy(log_probs, mask):
-  """Computes the approximate entropy for the given log-probs.
+def masked_entropy(log_probs, mask):
+  """Computes the entropy for the given log-probs.
 
   Args:
     log_probs: (B, T+1, A) log probs
     mask: (B, T) mask.
 
   Returns:
-    Approximate entropy.
+    Entropy.
   """
   # Cut the last time-step out.
   lp = log_probs[:, :-1]
   # Mask out the irrelevant part.
   lp *= mask[:, :, np.newaxis]  # make mask (B, T, 1)
+  p = np.exp(lp) * mask[:, :, np.newaxis]  # (B, T, 1)
   # Average on non-masked part and take negative.
-  return - (np.sum(lp) / np.sum(mask))
+  return -(np.sum(lp * p) / np.sum(mask))
+
+
+def evaluate_policy(eval_env, get_predictions, boundary, rng=None):
+  trajs = env_problem_utils.play_env_problem_with_policy(
+      eval_env,
+      get_predictions,
+      boundary=boundary,
+      idx=0,  # reset always
+      rng=rng)
+  avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
+  return avg_reward
 
 
 def training_loop(
@@ -977,15 +1016,24 @@ def training_loop(
     lambda_=LAMBDA,
     epsilon=EPSILON,
     c1=1.0,
-    c2=0.01):
+    c2=0.01,
+    output_dir=None,
+    eval_every_n=1000,
+    eval_env=None,
+    enable_early_stopping=True):
   """Runs the training loop for PPO, with fixed policy and value nets."""
   assert env
+
+  if output_dir:
+    gfile.makedirs(output_dir)
+
   jax_rng_key = trax.get_random_number_generator_and_set_seed(random_seed)
 
   value_losses = []
   ppo_objective = []
   combined_losses = []
   average_rewards = []
+  eval_average_rewards = []
 
   # Batch Observations Shape = [-1, -1] + OBS, because we will eventually call
   # policy and value networks on shape [B, T] +_OBS
@@ -1049,8 +1097,7 @@ def get_predictions(observations, rng=None):
 
       if policy_net_apply is not None:
         return (policy_net_apply(observations, policy_net_params, rng=key1),
-                value_net_apply(observations, value_net_params, rng=key2),
-                key)
+                value_net_apply(observations, value_net_params, rng=key2), key)
 
       assert policy_and_value_net_apply
 
@@ -1059,6 +1106,26 @@ def get_predictions(observations, rng=None):
 
       return log_probs, value_preds, key
 
+    # Save params and evaluate the policy.
+    if output_dir and (i % eval_every_n == 0):
+      jax_rng_key, key = jax_random.split(jax_rng_key, num=2)
+
+      logging.vlog(1, "Epoch [% 6d] saving model and evaluating policy.", i)
+      params_file = os.path.join(output_dir, "model-%06d.pkl" % i)
+      eval_rewards_file = os.path.join(output_dir, "eval_average_rewards")
+      with gfile.GFile(params_file, "wb") as f:
+        if policy_and_value_net_params:
+          pickle.dump(policy_and_value_net_params, f)
+        else:
+          pickle.dump((policy_net_params, value_net_params), f)
+
+      # TODO(afrozm): Dump in jaxboard or somewhere?
+      avg_reward = evaluate_policy(eval_env, get_predictions, boundary, rng=key)
+      eval_average_rewards.append(avg_reward)
+      logging.info("Epoch [% 6d] Policy Evaluation = %10.2f", i, avg_reward)
+      with gfile.GFile(eval_rewards_file, "w") as f:
+        f.write(", ".join([str(r) for r in eval_average_rewards]) + "\n")
+
     t = time.time()
     t0 = t
     logging.vlog(1, "Epoch [% 6d] collecting trajectories.", i)
@@ -1071,6 +1138,7 @@ def get_predictions(observations, rng=None):
         max_timestep=max_timestep,
         boundary=boundary,
         rng=key,
+        idx=i,
         epsilon=(10.0 / (i + 10.0)))  # this is a different epsilon.
 
     logging.vlog(1, "Collecting trajectories took %0.2f msec.", get_time(t))
@@ -1085,6 +1153,11 @@ def get_predictions(observations, rng=None):
                  [float(np.sum(traj[2])) for traj in trajs])
     logging.vlog(1, "Average Rewards:\n%s", average_rewards)
 
+    # TODO(afrozm): Dump in jaxboard or somewhere?
+    if output_dir:
+      with gfile.GFile(os.path.join(output_dir, "average_rewards"), "w") as f:
+        f.write(", ".join([str(r) for r in average_rewards]) + "\n")
+
     logging.vlog(1,
                  "Trajectory Length average=[%0.2f], max=[%0.2f], min=[%0.2f]",
                  float(sum(len(traj[0]) for traj in trajs)) / len(trajs),
@@ -1105,6 +1178,13 @@ def get_predictions(observations, rng=None):
 
     # Calculate log-probabilities and value predictions of the trajectories.
     # We'll pass these to the loss functions so as to not get recomputed.
+
+    # NOTE:
+    # There is a slight problem here, if the policy network contains
+    # stochasticity in the log-probabilities (ex: dropout), then calculating
+    # these again here is not going to be correct and should be done in the
+    # collect function.
+
     jax_rng_key, key = jax_random.split(jax_rng_key)
     log_probabs_traj, value_predictions_traj, _ = get_predictions(
         padded_observations, rng=key)
@@ -1147,8 +1227,10 @@ def get_predictions(observations, rng=None):
               c2=c2,
               rng=key1))
       logging.vlog(
-          1, "Calculating P&V loss [%10.2f(%10.2f, %10.2f)] took %0.2f msec.",
-          cur_combined_loss, cur_value_loss, cur_ppo_loss, get_time(t))
+          1,
+          "Calculating P&V loss [%10.2f(%10.2f, %10.2f, %10.2f)] took %0.2f msec.",
+          cur_combined_loss, cur_value_loss, cur_ppo_loss, entropy_bonus,
+          get_time(t))
     else:
       logging.vlog(2, "Starting to compute Value loss.")
       t = time.time()
@@ -1220,11 +1302,10 @@ def get_predictions(observations, rng=None):
         log_probab_actions_new, _ = policy_and_value_net_apply(
             padded_observations, new_policy_and_value_net_params, rng=k2)
 
-        approx_kl = approximate_kl(log_probab_actions_new,
-                                   log_probabs_traj,
+        approx_kl = approximate_kl(log_probab_actions_new, log_probabs_traj,
                                    reward_mask)
 
-        early_stopping = approx_kl > 1.5 * target_kl
+        early_stopping = enable_early_stopping and approx_kl > 1.5 * target_kl
         if early_stopping:
           logging.vlog(
               1, "Early stopping policy and value optimization at iter: %d, "
@@ -1233,9 +1314,8 @@ def get_predictions(observations, rng=None):
           # iteration.
 
         t2 = time.time()
-        if (((j + 1) %
-             print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1)
-            or early_stopping):
+        if (((j + 1) % print_every_optimizer_steps == 0) or
+            (j == num_optimizer_steps - 1) or early_stopping):
           # Compute and log the loss.
           (loss_combined, loss_ppo, loss_value, entropy_bonus) = (
               combined_loss(
@@ -1255,9 +1335,10 @@ def get_predictions(observations, rng=None):
                   rng=k3))
           logging.vlog(1, "One Policy and Value grad desc took: %0.2f msec",
                        get_time(t, t2))
-          logging.vlog(1, "Combined Loss(value, ppo, entropy_bonus) [%10.2f] ->"
-                          " [%10.2f(%10.2f,%10.2f,%10.2f)]", cur_combined_loss,
-                       loss_combined, loss_value, loss_ppo, entropy_bonus)
+          logging.vlog(
+              1, "Combined Loss(value, ppo, entropy_bonus) [%10.2f] ->"
+              " [%10.2f(%10.2f,%10.2f,%10.2f)]", cur_combined_loss,
+              loss_combined, loss_value, loss_ppo, entropy_bonus)
 
         if early_stopping:
           break
@@ -1268,10 +1349,11 @@ def get_predictions(observations, rng=None):
            (cur_combined_loss - loss_combined) / np.abs(cur_combined_loss)))
 
       logging.info(
-          "Epoch [% 6d], Reward[min, max, avg] [%10.2f,%10.2f,%10.2f], Combined"
-          " Loss(value, ppo) [%10.2f(%10.2f,%10.2f)], took [%10.2f msec]", i,
-          min_reward, max_reward, avg_reward, loss_combined, loss_value,
-          loss_ppo, get_time(t1))
+          "Epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
+          " Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)], took [%2.5f msec], Average Rewards(last 10):%s",
+          i, min_reward, max_reward,
+          avg_reward, loss_combined, loss_value, loss_ppo, entropy_bonus,
+          get_time(t1), ", ".join([str(a) for a in average_rewards[-10:]]))
     else:
       # Run optimizers.
       logging.vlog(1, "PPO Optimization")
@@ -1302,14 +1384,12 @@ def get_predictions(observations, rng=None):
         # Compute the approx KL for early stopping.
         # Get the new params.
         new_policy_net_params = ppo_get_params(ppo_opt_state)
-        log_probab_actions_new = policy_net_apply(padded_observations,
-                                                  new_policy_net_params,
-                                                  rng=k2)
-        approx_kl = approximate_kl(log_probab_actions_new,
-                                   log_probabs_traj,
+        log_probab_actions_new = policy_net_apply(
+            padded_observations, new_policy_net_params, rng=k2)
+        approx_kl = approximate_kl(log_probab_actions_new, log_probabs_traj,
                                    reward_mask)
 
-        early_stopping = approx_kl > 1.5 * target_kl
+        early_stopping = enable_early_stopping and approx_kl > 1.5 * target_kl
         if early_stopping:
           logging.vlog(
               1, "Early stopping policy optimization at iter: %d, "
@@ -1388,14 +1468,6 @@ def get_predictions(observations, rng=None):
           min_reward, max_reward, avg_reward, new_ppo_loss, new_value_loss,
           get_time(t0))
 
-  # Log the parameters, just for the sake of it.
-  if policy_net_params:
-    log_params(policy_net_params, "policy_net_params")
-  if value_net_params:
-    log_params(value_net_params, "value_net_params")
-  if policy_and_value_net_params:
-    log_params(policy_and_value_net_params, "policy_and_value_net_params")
-
   if value_losses:
     logging.vlog(1, "value_losses: %s", np.stack(value_losses))
   if ppo_objective:
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 617dd21de..1ea80f29d 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -46,10 +46,12 @@
 import gym
 import jax
 from jax.config import config
+import numpy as onp
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import rendered_env_problem
 from tensor2tensor.rl import gym_utils
 from tensor2tensor.trax import layers
+from tensor2tensor.trax.models import atari_cnn
 from tensor2tensor.trax.rlax import ppo
 
 FLAGS = flags.FLAGS
@@ -58,7 +60,7 @@
 flags.DEFINE_string("env_problem_name", None, "Name of the EnvProblem to make.")
 
 flags.DEFINE_integer("epochs", 100, "Number of epochs to run for.")
-flags.DEFINE_integer("random_seed", 0, "Random seed.")
+flags.DEFINE_string("random_seed", None, "Random seed.")
 flags.DEFINE_integer("batch_size", 32, "Batch of trajectories needed.")
 
 flags.DEFINE_integer(
@@ -69,7 +71,14 @@
 flags.DEFINE_integer(
     "max_timestep", None,
     "If set to an integer, maximum number of time-steps in a "
-    "trajectory.")
+    "trajectory. The bare env is wrapped with TimeLimit wrapper.")
+
+# This is different from max_timestep is that in the above, the env is wrapped
+# in a TimeLimit wrapper, vs here we use this in the collect function.
+flags.DEFINE_integer(
+    "truncation_timestep", None,
+    "If set to an integer, maximum number of time-steps in a "
+    "trajectory. Used in the collect procedure.")
 
 flags.DEFINE_boolean(
     "jax_debug_nans", False,
@@ -87,6 +96,11 @@
     "If True there is a single network that determines policy"
     "and values.")
 
+flags.DEFINE_bool(
+    "two_towers", True,
+    "In the combined network case should we make one tower or"
+    "two.")
+
 flags.DEFINE_boolean("flatten_dims", False,
                      "If true, we flatten except the first two dimensions.")
 
@@ -96,6 +110,9 @@
                      "Number of optimizer steps policy only.")
 flags.DEFINE_integer("value_only_num_optimizer_steps", 80,
                      "Number of optimizer steps value only.")
+flags.DEFINE_integer(
+    "print_every_optimizer_steps", 1,
+    "How often to log during the policy optimization process.")
 
 # Learning rate of the combined net, policy net and value net.
 flags.DEFINE_float("learning_rate", 1e-3, "Learning rate.")
@@ -110,9 +127,24 @@
                    "Coefficient of Value Loss term in combined loss.")
 flags.DEFINE_float("entropy_coef", 0.01,
                    "Coefficient of the Entropy Bonus term in combined loss.")
+flags.DEFINE_float("gamma", 0.99, "Policy iteration early stopping")
+flags.DEFINE_float("lambda_", 0.95, "Policy iteration early stopping")
+flags.DEFINE_float("epsilon", 0.1, "Policy iteration early stopping")
+
+flags.DEFINE_string("output_dir", "", "Output dir.")
+flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
+flags.DEFINE_bool("enable_early_stopping", True,
+                  "Whether to enable early stopping.")
+flags.DEFINE_bool("xm", False, "Are we running on borg?.")
+flags.DEFINE_integer("eval_every_n", 100, "How frequently to eval the policy.")
+flags.DEFINE_integer("eval_batch_size", 4, "Batch size for evaluation.")
 
 
 def common_layers():
+  # TODO(afrozm): Refactor.
+  if "Pong" in FLAGS.env_problem_name:
+    return atari_layers()
+
   cur_layers = []
   if FLAGS.flatten_dims:
     cur_layers = [layers.Div(divisor=255.0), layers.Flatten(num_axis_to_keep=2)]
@@ -120,7 +152,11 @@ def common_layers():
   return cur_layers + body
 
 
-def make_env():
+def atari_layers():
+  return [atari_cnn.AtariCnn()]
+
+
+def make_env(batch_size=8):
   """Creates the env."""
   if FLAGS.env_name:
     return gym.make(FLAGS.env_name)
@@ -131,7 +167,7 @@ def make_env():
   if not FLAGS.resize:  # None or False
     return env_problem.EnvProblem(
         base_env_name=FLAGS.env_problem_name,
-        batch_size=FLAGS.batch_size,
+        batch_size=batch_size,
         reward_range=(-1, 1))
 
   wrapper_fn = functools.partial(
@@ -141,12 +177,12 @@ def make_env():
           "rendered_env": True,
           "rendered_env_resize_to": (FLAGS.resized_height, FLAGS.resized_width),
           "sticky_actions": False,
-          "output_dtype": None,
+          "output_dtype": onp.int32 if FLAGS.use_tpu else None,
       })
 
   return rendered_env_problem.RenderedEnvProblem(
       base_env_name=FLAGS.env_problem_name,
-      batch_size=FLAGS.batch_size,
+      batch_size=batch_size,
       env_wrapper_fn=wrapper_fn,
       reward_range=(-1, 1))
 
@@ -160,11 +196,22 @@ def main(argv):
 
   if FLAGS.jax_debug_nans:
     config.update("jax_debug_nans", True)
+  if FLAGS.use_tpu:
+    config.update("jax_platform_name", "tpu")
+
+  # TODO(afrozm): Refactor.
+  if "Pong" in FLAGS.env_problem_name and FLAGS.xm:
+    from tensor2tensor.rl.google import atari_utils  # pylint: disable=g-import-not-at-top
+    FLAGS.atari_roms_path = "local_ram_fs_tmp"
+    atari_utils.copy_roms()
 
   # Make an env here.
-  env = make_env()
+  env = make_env(batch_size=FLAGS.batch_size)
   assert env
 
+  eval_env = make_env(batch_size=FLAGS.eval_batch_size)
+  assert eval_env
+
   def run_training_loop():
     """Runs the training loop."""
     policy_net_fun = None
@@ -176,7 +223,9 @@ def run_training_loop():
 
     if FLAGS.combined_network:
       policy_and_value_net_fun = functools.partial(
-          ppo.policy_and_value_net, bottom_layers_fn=common_layers)
+          ppo.policy_and_value_net,
+          bottom_layers_fn=common_layers,
+          two_towers=FLAGS.two_towers)
       policy_and_value_optimizer_fun = get_optimizer_fun(FLAGS.learning_rate)
     else:
       policy_net_fun = functools.partial(
@@ -186,6 +235,12 @@ def run_training_loop():
       policy_optimizer_fun = get_optimizer_fun(FLAGS.policy_only_learning_rate)
       value_optimizer_fun = get_optimizer_fun(FLAGS.value_only_learning_rate)
 
+    random_seed = None
+    try:
+      random_seed = int(FLAGS.random_seed)
+    except Exception:  # pylint: disable=broad-except
+      pass
+
     ppo.training_loop(
         env=env,
         epochs=FLAGS.epochs,
@@ -198,13 +253,21 @@ def run_training_loop():
         num_optimizer_steps=FLAGS.num_optimizer_steps,
         policy_only_num_optimizer_steps=FLAGS.policy_only_num_optimizer_steps,
         value_only_num_optimizer_steps=FLAGS.value_only_num_optimizer_steps,
+        print_every_optimizer_steps=FLAGS.print_every_optimizer_steps,
         batch_size=FLAGS.batch_size,
         target_kl=FLAGS.target_kl,
         boundary=FLAGS.boundary,
-        max_timestep=FLAGS.max_timestep,
-        random_seed=FLAGS.random_seed,
+        max_timestep=FLAGS.truncation_timestep,
+        random_seed=random_seed,
         c1=FLAGS.value_coef,
-        c2=FLAGS.entropy_coef)
+        c2=FLAGS.entropy_coef,
+        gamma=FLAGS.gamma,
+        lambda_=FLAGS.lambda_,
+        epsilon=FLAGS.epsilon,
+        enable_early_stopping=FLAGS.enable_early_stopping,
+        output_dir=FLAGS.output_dir,
+        eval_every_n=FLAGS.eval_every_n,
+        eval_env=eval_env)
 
   if FLAGS.jax_debug_nans or FLAGS.disable_jit:
     with jax.disable_jit():
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 25243924a..18bcbd524 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -636,50 +636,7 @@ def test_combined_loss(self):
                     ppo_loss_2 + (c1 * value_loss_2) - (c2 * entropy_bonus),
                     1e-6)
 
-  def test_approximate_kl(self):
-    # (2, 4+1, 4)
-    p_old = np.array([[
-        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
-        [np.log(0.4), np.log(0.1), np.log(0.4), np.log(0.1)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-    ], [
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-        [np.log(0.1), np.log(0.1), np.log(0.4), np.log(0.4)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-    ]])
-
-    # (2, 4+1, 4)
-    p_new = np.array([[
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-        [np.log(0.4), np.log(0.1), np.log(0.1), np.log(0.3)],
-        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
-    ], [
-        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
-        [np.log(0.1), np.log(0.1), np.log(0.2), np.log(0.6)],
-        [np.log(0.3), np.log(0.1), np.log(0.3), np.log(0.3)],
-        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
-        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
-    ]])
-
-    # (2, 4)
-    mask = np.array([
-        [1, 1, 0, 0],
-        [1, 1, 1, 0]
-    ])
-
-    self.assertNear(
-        ppo.approximate_kl(p_new, p_old, mask),
-        -ppo.approximate_entropy(p_old, mask) +
-        ppo.approximate_entropy(p_new, mask),
-        1e-6)
-
-  def test_get_approximate_entropy(self):
+  def test_masked_entropy(self):
     # (2, 4+1, 4)
     log_probs = np.array([[
         [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
@@ -701,16 +658,19 @@ def test_get_approximate_entropy(self):
         [1, 1, 1, 0]
     ])
 
+    def plp(p):
+      return p * np.log(p)
+
     # Removing the last time-step and the masked stuff, gets us this.
     filtered_log_probs = np.array([[
-        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
-        [np.log(0.4), np.log(0.1), np.log(0.4), np.log(0.1)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-        [np.log(0.1), np.log(0.1), np.log(0.4), np.log(0.4)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
+        [plp(0.1), plp(0.2), plp(0.6), plp(0.1)],
+        [plp(0.4), plp(0.1), plp(0.4), plp(0.1)],
+        [plp(0.3), plp(0.1), plp(0.5), plp(0.1)],
+        [plp(0.1), plp(0.1), plp(0.4), plp(0.4)],
+        [plp(0.3), plp(0.1), plp(0.5), plp(0.1)],
     ]])
 
-    self.assertNear(ppo.approximate_entropy(log_probs, mask),
+    self.assertNear(ppo.masked_entropy(log_probs, mask),
                     -np.sum(filtered_log_probs) / 5.0,
                     1e-6)
 

From 2a2dbfba119873e114d6c642ee76dbc69bb434a8 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 20 May 2019 18:02:37 -0700
Subject: [PATCH 2043/2720] Add `render` to EnvProblem, we can see the env
 locally with this function.

PiperOrigin-RevId: 249164639
---
 tensor2tensor/envs/env_problem.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 7aacec777..b3bbdaaa0 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -515,6 +515,24 @@ def reset(self, indices=None):
 
     return processed_observations
 
+  def render(self, mode="human", indices=None):
+    """Calls render with the given mode on the specified indices.
+
+    Args:
+      mode: rendering mode.
+      indices: array of indices, calls render on everything if indices is None.
+
+    Returns:
+      a list of return values from the environments rendered.
+    """
+
+    if indices is None:
+      indices = np.arange(self.batch_size)
+    ret_vals = []
+    for index in indices:
+      ret_vals.append(self._envs[index].render(mode=mode))
+    return ret_vals
+
   def _step(self, actions):
     """Takes a step in all environments, shouldn't pre-process or record.
 

From 777c6a8f477bcb36ffab5c9bdd42c54a2932501a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 20 May 2019 19:19:17 -0700
Subject: [PATCH 2044/2720] Rename Copy to NoOp and extend Select so the
 UnnestBranches can be removed.

PiperOrigin-RevId: 249173585
---
 tensor2tensor/trax/layers/attention.py        |  53 +++------
 tensor2tensor/trax/layers/base.py             |  63 +++++++---
 tensor2tensor/trax/layers/combinators.py      | 108 ++++++++++--------
 tensor2tensor/trax/layers/combinators_test.py |  15 +--
 tensor2tensor/trax/layers/convolution.py      |   5 +-
 tensor2tensor/trax/layers/core.py             |   6 +
 tensor2tensor/trax/layers/core_test.py        |   4 +-
 tensor2tensor/trax/layers/rnn.py              |   8 +-
 tensor2tensor/trax/models/atari_cnn.py        |   2 +-
 .../models/research/chunked_transformer.py    |  51 +++++++--
 tensor2tensor/trax/models/resnet.py           |   2 +-
 tensor2tensor/trax/models/transformer.py      |  34 +++---
 12 files changed, 211 insertions(+), 140 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 8763660f7..5aa7b2e9a 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -49,7 +49,7 @@ def EncoderDecoderMaskShape(inputs):
   return (batch_size, 1, target_length, input_length)
 
 
-@base.layer(output_shape=EncoderDecoderMaskShape)
+@base.layer(output_shape=EncoderDecoderMaskShape, stack_items_to_pass=0)
 def EncoderDecoderMask(x, **unused_kwargs):
   """Make encoder-decoder mask from a padding mask and decoder input."""
   (padding_mask, decoder_input) = x
@@ -63,10 +63,7 @@ def EncoderDecoderMask(x, **unused_kwargs):
 def _positional_encoding_new_params(input_shape, rng, max_len=2048):  # pylint: disable=invalid-name
   """Helper: create positional encoding parameters."""
   del rng
-  # Check if we are operating on chunked inputs by checking if the first
-  # shape is a list/tuple of shapes (otherwise it's an int or numpy array).
-  is_chunked = isinstance(input_shape[0], (list, tuple))
-  feature_depth = input_shape[0][-1] if is_chunked else input_shape[-1]
+  feature_depth = input_shape[-1]
   pe = onp.zeros((max_len, feature_depth), dtype=onp.float32)
   position = onp.arange(0, max_len)[:, onp.newaxis]
   div_term = onp.exp(
@@ -80,17 +77,8 @@ def _positional_encoding_new_params(input_shape, rng, max_len=2048):  # pylint:
 @base.layer(new_parameters=_positional_encoding_new_params)
 def PositionalEncoding(x, params, **unused_kwargs):
   """Implements bare positional encoding."""
-  if not isinstance(x, (list, tuple)):  # non-chunked inputs
-    symbol_size = np.shape(x)[1]
-    return x + params[:, :symbol_size, :]
-  # Chunked case: apply to all chunks selecting as much as needed.
-  offset = 0
-  results = []
-  for chunk in x:
-    symbol_size = np.shape(chunk)[1]
-    results.append(chunk + params[:, offset:offset + symbol_size, :])
-    offset += symbol_size
-  return results
+  symbol_size = np.shape(x)[1]
+  return x + params[:, :symbol_size, :]
 
 
 def DotProductAttention(query, key, value, mask, dropout, mode, rng):
@@ -150,12 +138,15 @@ def apply_fun(params, inputs, **kwargs):  # pylint: disable=invalid-name
 def _multihead_attention_output_shape(  # pylint: disable=invalid-name
     input_shapes, **unused_kwargs):
   """Helper: calculate multihead attention output shape."""
-  q_shape = input_shapes[0][0]  # Inputs are ((q, k, v), mask).
-  mask_shape = input_shapes[1]
-  return q_shape, mask_shape
+  q_shape = input_shapes[0]  # Inputs are (q, k, v, mask).
+  v_shape = input_shapes[2]  # Inputs are (q, k, v, mask).
+  mask_shape = input_shapes[3]
+  res_shape = list(q_shape[:-1]) + [v_shape[-1]]
+  return tuple(res_shape), mask_shape
 
 
-@base.layer(output_shape=_multihead_attention_output_shape)
+@base.layer(output_shape=_multihead_attention_output_shape,
+            stack_items_to_pass=4)
 def PureMultiHeadedAttention(x, params, num_heads=8, dropout=0.0,
                              mode='train', **kwargs):
   """Pure transformer-style multi-headed attention.
@@ -173,7 +164,7 @@ def PureMultiHeadedAttention(x, params, num_heads=8, dropout=0.0,
   """
   del params
   rng = kwargs.get('rng', None)
-  (q, k, v), mask = x
+  q, k, v, mask = x
   feature_depth = q.shape[-1]
   assert feature_depth % num_heads == 0
   head_depth = feature_depth // num_heads
@@ -211,17 +202,15 @@ def MultiHeadedAttentionQKV(
   """
   return combinators.Serial(
       combinators.Parallel(
-          combinators.Parallel(
-              core.Dense(feature_depth),
-              core.Dense(feature_depth),
-              core.Dense(feature_depth),
-          ),
-          combinators.Copy()
+          core.Dense(feature_depth),
+          core.Dense(feature_depth),
+          core.Dense(feature_depth),
+          combinators.NoOp()
       ),
       PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
           feature_depth=feature_depth, num_heads=num_heads,
           dropout=dropout, mode=mode),
-      combinators.Parallel(core.Dense(feature_depth), combinators.Copy())
+      combinators.Parallel(core.Dense(feature_depth), combinators.NoOp())
   )
 
 
@@ -241,12 +230,8 @@ def MultiHeadedAttention(
     Multi-headed self-attention layer.
   """
   return combinators.Serial(
-      combinators.Parallel(
-          # q = k = v = first input
-          combinators.Branch(
-              combinators.Copy(), combinators.Copy(), combinators.Copy()),
-          combinators.Copy()  # pass the mask
-      ),
+      combinators.Dup(),
+      combinators.Dup(),
       MultiHeadedAttentionQKV(  # pylint: disable=no-value-for-parameter
           feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),
   )
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index a71a72f20..1f9e167d5 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -43,7 +43,7 @@ def call(self, x, params=(), **kwargs):
     """Call this layer in input x using the given parameters."""
     raise NotImplementedError
 
-  def output_shape(self, input_shape):
+  def output_shape_fun(self, input_shape):
     """The shape of the output of this layer given the shape of the input.
 
     Note that all arguments and return values can be tuples or dictionaries
@@ -72,12 +72,19 @@ def new_parameters(self, input_shape, rng):
     """
     raise NotImplementedError
 
+  def stack_items_to_pass(self):
+    """How many of the top stack items do we process."""
+    return 0
+
   # End of subclassing interface, all functions below are internal.
 
-  def output_shape_catch_errors(self, input_shape):
+  def output_shape(self, input_shape):
     """Same as self.output_shape but with better error reporting."""
     try:
-      return self.output_shape(input_shape)
+      is_list = isinstance(input_shape, (list, tuple))
+      is_list = is_list and isinstance(input_shape[0], (list, tuple))
+      n = self.stack_items_to_pass() if is_list else 0
+      return _apply_to_first_n(self.output_shape_fun, input_shape, n)
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback()
       raise LayerError(name, 'output_shape', self._caller, input_shape, trace)
@@ -106,10 +113,16 @@ def initialize(self, input_shape, rng):
 
       # First call of this layer, create parameters.
       self._first_init = False
+      is_list = isinstance(input_shape, (list, tuple))
+      is_list = is_list and isinstance(input_shape[0], (list, tuple))
+      if is_list and self.stack_items_to_pass() > 0:
+        input_shape = input_shape[:self.stack_items_to_pass()]
+        if len(input_shape) == 1:
+          input_shape = input_shape[0]
       self._params = self.new_parameters(input_shape, rng)
       return self._params
     except Exception:
-      name, trace = self.__class__.__name__, _short_traceback()
+      name, trace = self.__class__.__name__, _short_traceback(skip=3)
       raise LayerError(name, 'initialize', self._caller, input_shape, trace)
 
   def __call__(self, x, params=(), **kwargs):
@@ -119,10 +132,12 @@ def __call__(self, x, params=(), **kwargs):
       # Note: to make sure jit tracers can decide this branch in python we
       #   use "params is ()" instead of, e.g., "not params" or "params == ()".
       if params is ():  # pylint: disable=literal-comparison
-        return self.call(x, params=self._params, **kwargs)
+        params = self._params
       # In this case, we're called for the first time: cache parameters.
       self._params = params
-      return self.call(x, params=params, **kwargs)
+      f = lambda y: self.call(y, params=params, **kwargs)
+      n = self.stack_items_to_pass() if isinstance(x, (list, tuple)) else 0
+      return _apply_to_first_n(f, x, n)
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback()
       raise LayerError(name, 'call', self._caller, shapes(x), trace)
@@ -155,6 +170,22 @@ def message(self):
     return prefix + caller + shapes_str + self._traceback
 
 
+def _apply_to_first_n(f, x, n):
+  """Helper: apply f to first n elements on the stack x if n > 0."""
+  if n < 1:
+    return f(x)
+  argument, rest = x[:n], x[n:]
+  if n == 1:
+    argument = argument[0]
+  result = f(argument)
+  if n == 1:
+    result = [result]
+  result = list(result) + list(rest)
+  if isinstance(x, tuple):
+    result = tuple(result)
+  return result
+
+
 def nested_map(x, f):
   """Map the function f to the nested structure x (dicts, tuples, lists)."""
   if isinstance(x, list):
@@ -171,9 +202,9 @@ def nested_reduce(x, f):
   if isinstance(x, list):
     return f([nested_reduce(y, f) for y in x])
   if isinstance(x, tuple):
-    return f(tuple([nested_reduce(y, f) for y in x]))
-  if isinstance(x, dict):
-    return f({k: nested_reduce(x[k], f) for k in x})
+    return f([nested_reduce(y, f) for y in x])
+  if isinstance(x, dict):  # We apply f only to values in the dicts.
+    return f([nested_reduce(v, f) for v in x.values()])
   return x
 
 
@@ -224,7 +255,7 @@ def _shorten_file_path(line):
   return line[:first_quote] + '[...]/' + new_path + line[second_quote + 1:]
 
 
-def _short_traceback(skip=3):
+def _short_traceback(skip=7):
   """Cleaned-up form of traceback."""
   counter, res = 0, []
   # Skipping 3 lines by default: the top (useless) and self-call.
@@ -246,10 +277,15 @@ def _short_traceback(skip=3):
 # Decorator for making layers from functions.
 
 
-def layer(output_shape=None, new_parameters=None):
+def layer(output_shape=None, new_parameters=None, stack_items_to_pass=1):
   """Create a layer class from a function."""
   def layer_decorator(call):
     """Decorating the call function."""
+
+    def stack_items_to_pass_fun(self):
+      del self
+      return stack_items_to_pass
+
     def output_shape_fun(self, input_shape):
       if output_shape is None:
         return input_shape
@@ -280,8 +316,9 @@ def call_fun(self, x, params=(), **kwargs):
     # Create the class.
     cls = type(call.__name__, (Layer,),
                {'call': call_fun,
-                'output_shape': output_shape_fun,
-                'new_parameters': new_parameters_fun})
+                'output_shape_fun': output_shape_fun,
+                'new_parameters': new_parameters_fun,
+                'stack_items_to_pass': stack_items_to_pass_fun})
 
     return cls
   return layer_decorator
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 80a66ec5a..2361e4bdc 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -44,10 +44,10 @@ def call(self, x, params=(), **kwargs):
       x = layer(x, p, rng=rng, **kwargs)
     return x
 
-  def output_shape(self, input_shape):
+  def output_shape_fun(self, input_shape):
     cur_shape = input_shape
     for layer in self._layers:
-      cur_shape = layer.output_shape_catch_errors(cur_shape)
+      cur_shape = layer.output_shape(cur_shape)
     return cur_shape
 
   def new_parameters(self, input_shape, rng):
@@ -62,40 +62,23 @@ def new_parameters(self, input_shape, rng):
 
 
 @base.layer()
-def Copy(x, **unused_kwargs):
-  """Copy layer, return the inputs."""
+def NoOp(x, **unused_kwargs):
+  """NoOp layer, return the inputs."""
   return x
 
 
-def Unnest(x):
-  """Helper: remove nesting in x, return a flat tuple."""
-  if not isinstance(x, (list, tuple)):
-    return (x,)
-  return tuple([z for y in x for z in Unnest(y)])  # pylint: disable=g-complex-comprehension
-
-
-def UnnestShape(shape):
-  """Unnest a nested structure of shapes."""
-
-  class Shape(object):
-    """Since shapes are tuples, make them a class to not unnest too far."""
-
-    def __init__(self, shape):
-      self.shape = shape
-
-  def MakeShape(nested_shape):
-    """Make all shape-tuples in the nested object shape-classes."""
-    if isinstance(nested_shape[0], int):  # Not nested.
-      return Shape(nested_shape)
-    return [MakeShape(shape) for shape in nested_shape]
+def _dup(x):  # pylint: disable=invalid-name
+  """Helper: copy the top element of a list or a tuple."""
+  if isinstance(x, list):
+    return [x[0]] + x
+  assert isinstance(x, tuple)
+  return tuple([x[0]] + list(x))
 
-  # Unnest on the level of shape-classes and bring back shape-tuples.
-  return tuple([y.shape for y in Unnest(MakeShape(shape))])
 
-
-@base.layer(output_shape=UnnestShape)
-def UnnestBranches(x, **unused_kwargs):
-  return Unnest(x)
+@base.layer(output_shape=_dup, stack_items_to_pass=0)
+def Dup(x, **unused_kwargs):
+  """Duplicate (copy) the first element on the stack."""
+  return _dup(x)
 
 
 # Re-ordering layer.
@@ -116,32 +99,56 @@ class Select(base.Layer):
     Select((0, (1, 1)))      = (x, (y, y))
     Select(((2, 0), (1, 1))) = ((z, x), (y, y))
 
-  By default (if no output is given) Select does nothing (Copy).
+  By default (if no output is given) Select does nothing (NoOp).
+  It is also possible to name the inputs to access tuple elements, e.g.:
+
+  Select(inputs=('encoder', ('decoder', 'mask')), output='decoder')
+
+  will transform a tuple (x, (y, x)) into y.
 
   Args:
     x: the input tuple to re-order.
     params: layer parameters (unused).
     output: the specification of the output tuple: a nested tuple of ints.
+    input: the specification of the input tuple if we need to disassemble it.
     **kwargs: other arguments (unused).
 
   Returns:
     The re-ordered tuple with the same shape as output.
   """
 
-  def __init__(self, output=None):
+  def __init__(self, output=None, inputs=None):
     super(Select, self).__init__()
     self._output = output
+    if inputs is None:
+      self._map = lambda x, i: x[i]
+    else:
+      self._input_map = {}
+      self._build_input_map(inputs, [])
+      def InputMapping(x, i):
+        cur = x
+        for idx in self._input_map[i]:
+          cur = cur[idx]
+        return cur
+      self._map = InputMapping
+
+  def _build_input_map(self, inputs, prefix):
+    for i, e in enumerate(inputs):
+      if isinstance(e, (list, tuple)):
+        self._build_input_map(e, prefix + [i])
+      else:
+        self._input_map[e] = prefix + [i]
 
   def call(self, x, params=(), **kwargs):
     del params, kwargs
     if self._output is None:
       return x
-    return base.nested_map(self._output, lambda i: x[i])
+    return base.nested_map(self._output, lambda i: self._map(x, i))
 
-  def output_shape(self, input_shape):
+  def output_shape_fun(self, input_shape):
     if self._output is None:
       return input_shape
-    return base.nested_map(self._output, lambda i: input_shape[i])
+    return base.nested_map(self._output, lambda i: self._map(input_shape, i))
 
   def new_parameters(self, input_shape, rng):
     return ()
@@ -191,17 +198,17 @@ def call(self, x, params=(), **kwargs):
       counter += 1
     return result
 
-  def output_shape(self, input_shape):
+  def output_shape_fun(self, input_shape):
     output_shapes = []
     # If the argument layers are a sequence, apply each to calculate shape.
     if not isinstance(self._layers, dict):
       for layer in self._layers:
-        output_shapes.append(layer.output_shape_catch_errors(input_shape))
+        output_shapes.append(layer.output_shape(input_shape))
       return tuple(output_shapes)
     # If layers are a dictionary, apply to the input shape.
     result = {}
     for k in self._layers:
-      result[k] = self._layers[k].output_shape_catch_errors(input_shape)
+      result[k] = self._layers[k].output_shape(input_shape)
     return result
 
   def new_parameters(self, input_shape, rng):
@@ -250,20 +257,20 @@ def _first_from_tuple_or_dict(tuple_or_dict):  # pylint: disable=invalid-name
     return x
 
 
-@base.layer(output_shape=_first_from_tuple_or_dict)
+@base.layer(output_shape=_first_from_tuple_or_dict, stack_items_to_pass=0)
 def Add(x, **unused_kwargs):
   """Add branches elementwise."""
   # Here x is a list of tensors of the same shape, or nested structures.
   return _nested_sum(x)
 
 
-@base.layer(output_shape=_first_from_tuple_or_dict)
+@base.layer(output_shape=_first_from_tuple_or_dict, stack_items_to_pass=0)
 def Multiply(x, **unused_kwargs):
   """Multiply branches elementwise."""
   return _nested_product(x)
 
 
-@base.layer(output_shape=_first_from_tuple_or_dict)
+@base.layer(output_shape=_first_from_tuple_or_dict, stack_items_to_pass=0)
 def Gate(x, **unused_kwargs):
   """Implements a gating function on a (memory, gate, candidate) tuple.
 
@@ -293,7 +300,7 @@ def _concatenate_shape(input_shape, axis=-1):  # pylint: disable=invalid-name
   return out_shape
 
 
-@base.layer(output_shape=_concatenate_shape)
+@base.layer(output_shape=_concatenate_shape, stack_items_to_pass=0)
 def Concatenate(x, params, axis=-1, **kwargs):
   del params, kwargs
   if isinstance(x, dict):  # For dictionaries, just use the values.
@@ -323,6 +330,9 @@ def __init__(self, *layers, **kwlayers):
     self._nlayers = len(layers)
     self._layers = layers
 
+  def stack_items_to_pass(self):
+    return self._nlayers
+
   def call(self, inputs, params=(), **kwargs):
     # Split the random number generators.
     rng = kwargs.pop('rng', None)
@@ -348,18 +358,18 @@ def call(self, inputs, params=(), **kwargs):
         result[k] = inputs[k]
     return result
 
-  def output_shape(self, input_shape):
+  def output_shape_fun(self, input_shape):
     output_shapes = []
     # If the argument layers are a sequence, apply each to calculate shape.
     if not isinstance(self._layers, dict):
       for i, layer in enumerate(self._layers):
-        output_shapes.append(layer.output_shape_catch_errors(input_shape[i]))
+        output_shapes.append(layer.output_shape(input_shape[i]))
       return tuple(output_shapes)
     # If layers are a dictionary, apply to matching keys in the input shape.
     result = {}
     for k in input_shape:
       if k in self._layers:
-        result[k] = self._layers[k].output_shape_catch_errors(input_shape[k])
+        result[k] = self._layers[k].output_shape(input_shape[k])
       else:
         result[k] = input_shape[k]
     return result
@@ -380,7 +390,7 @@ def new_parameters(self, input_shape, rng):
 
 def Residual(*layers, **kwargs):
   """Constructs a residual version of layers, summing input to layers output."""
-  shortcut = kwargs.get('shortcut', Copy())  # pylint: disable=no-value-for-parameter
+  shortcut = kwargs.get('shortcut', NoOp())  # pylint: disable=no-value-for-parameter
   if len(layers) > 1:
     return Serial(
         Branch(Serial(*layers), shortcut),
@@ -426,7 +436,7 @@ def call(self, inputs, params=(), **kwargs):
       return result
     return tuple(result)
 
-  def output_shape(self, input_shapes):
+  def output_shape_fun(self, input_shapes):
     return tuple([self._layer.output_shape(shape) for shape in input_shapes])
 
   def new_parameters(self, input_shape, rng):
@@ -482,7 +492,7 @@ def call(self, inp, params=(), **kwargs):
     out = self._layer(inp, params=params, **kwargs)
     return self._unmodify(out, batch_dims)
 
-  def output_shape(self, input_shape):
+  def output_shape_fun(self, input_shape):
     modified_shape, batch_dims = self._modify_shape(input_shape)
     out = self._layer.output_shape(modified_shape)
     return self._unmodify_shape(out, batch_dims)
diff --git a/tensor2tensor/trax/layers/combinators_test.py b/tensor2tensor/trax/layers/combinators_test.py
index 5092e3eb1..e0384a73e 100644
--- a/tensor2tensor/trax/layers/combinators_test.py
+++ b/tensor2tensor/trax/layers/combinators_test.py
@@ -30,14 +30,14 @@ def test_branch(self):
     input_shape = (2, 3)
     expected_shape = ((2, 3), (2, 3))
     output_shape = base.check_shape_agreement(
-        combinators.Branch(combinators.Copy(), combinators.Copy()), input_shape)
+        combinators.Branch(combinators.NoOp(), combinators.NoOp()), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
   def test_branch_named(self):
     input_shape = (2, 3)
     expected_shape = {'a': (2, 3), 'b': (2, 3)}
     output_shape = base.check_shape_agreement(
-        combinators.Branch(a=combinators.Copy(), b=combinators.Copy()),
+        combinators.Branch(a=combinators.NoOp(), b=combinators.NoOp()),
         input_shape)
     self.assertEqual(output_shape, expected_shape)
 
@@ -45,7 +45,7 @@ def test_parallel(self):
     input_shape = ((2, 3), (2, 3))
     expected_shape = ((2, 3), (2, 3))
     output_shape = base.check_shape_agreement(
-        combinators.Parallel(combinators.Copy(), combinators.Copy()),
+        combinators.Parallel(combinators.NoOp(), combinators.NoOp()),
         input_shape)
     self.assertEqual(output_shape, expected_shape)
 
@@ -53,7 +53,7 @@ def test_parallel_named(self):
     input_shape = {'a': (2, 3), 'b': (2, 3)}
     expected_shape = {'a': (2, 3), 'b': (2, 3)}
     output_shape = base.check_shape_agreement(
-        combinators.Parallel(a=combinators.Copy()), input_shape)
+        combinators.Parallel(a=combinators.NoOp()), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
   def test_select(self):
@@ -70,13 +70,6 @@ def test_select_named(self):
         combinators.Select('b'), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
-  def test_unnest_branches(self):
-    input_shape = ((2, 3), [(4, 5), (6, 7)], (8, 9, 10))
-    expected_shape = ((2, 3), (4, 5), (6, 7), (8, 9, 10))
-    output_shape = base.check_shape_agreement(
-        combinators.UnnestBranches(), input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
   def test_rebatch(self):
     input_shape = (29, 5, 5, 20)
     result_shape = base.check_shape_agreement(
diff --git a/tensor2tensor/trax/layers/convolution.py b/tensor2tensor/trax/layers/convolution.py
index 4bf17f817..769a93f6f 100644
--- a/tensor2tensor/trax/layers/convolution.py
+++ b/tensor2tensor/trax/layers/convolution.py
@@ -68,6 +68,9 @@ def __init__(self, filters, kernel_size, strides=None, padding='VALID',
       self._kernel_initializer = init.GlorotNormalInitializer(
           rhs_spec.index('O'), rhs_spec.index('I'))
 
+  def stack_items_to_pass(self):
+    return 1
+
   def call(self, x, params=(), **kwargs):
     del kwargs
     w, b = params
@@ -137,7 +140,7 @@ def _conv_general_shape_tuple(self, lhs_shape, rhs_shape, window_strides,
         lhs_trans, rhs_trans, window_strides, padding)
     return tuple(onp.take(out_trans, onp.argsort(out_perm)))
 
-  def output_shape(self, input_shape):
+  def output_shape_fun(self, input_shape):
     kernel_shape = self._kernel_shape(input_shape)
     return self._conv_general_shape_tuple(
         input_shape, kernel_shape,
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 0872d63d6..0b5a075a8 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -89,6 +89,9 @@ def __init__(self, units,
     self._kernel_initializer = kernel_initializer
     self._bias_initializer = bias_initializer
 
+  def stack_items_to_pass(self):
+    return 1
+
   def call(self, x, params, **kwargs):
     del kwargs
     w, b = params
@@ -114,6 +117,9 @@ def __init__(self, feature_depth, vocab_size,
     self._vocab_size = vocab_size
     self._kernel_initializer = kernel_initializer
 
+  def stack_items_to_pass(self):
+    return 1
+
   def call(self, x, params, **kwargs):
     del kwargs
     return np.take(params, x, axis=0)
diff --git a/tensor2tensor/trax/layers/core_test.py b/tensor2tensor/trax/layers/core_test.py
index ff4c4f32b..765fe92cf 100644
--- a/tensor2tensor/trax/layers/core_test.py
+++ b/tensor2tensor/trax/layers/core_test.py
@@ -47,10 +47,10 @@ def test_flatten_n(self):
     self.assertEqual(actual_shape, (29, 87, 10, 20, 30))
 
     # Not enough dimensions.
-    with self.assertRaises(ValueError):
+    with self.assertRaises(base.LayerError):
       base.check_shape_agreement(core.Flatten(num_axis_to_keep=5), input_shape)
 
-    with self.assertRaises(ValueError):
+    with self.assertRaises(base.LayerError):
       base.check_shape_agreement(core.Flatten(num_axis_to_keep=6), input_shape)
 
   def test_div(self):
diff --git a/tensor2tensor/trax/layers/rnn.py b/tensor2tensor/trax/layers/rnn.py
index b0caa2bbc..9e788c4b3 100644
--- a/tensor2tensor/trax/layers/rnn.py
+++ b/tensor2tensor/trax/layers/rnn.py
@@ -38,7 +38,7 @@ def GRUCell(units):
   """
   return GeneralGRUCell(
       candidate_transform=lambda: core.Dense(units=units),
-      memory_transform=combinators.Copy,
+      memory_transform=combinators.NoOp,
       gate_nonlinearity=core.Sigmoid,
       candidate_nonlinearity=core.Tanh)
 
@@ -62,13 +62,13 @@ def BuildConv():
 
   return GeneralGRUCell(
       candidate_transform=BuildConv,
-      memory_transform=combinators.Copy,
+      memory_transform=combinators.NoOp,
       gate_nonlinearity=core.Sigmoid,
       candidate_nonlinearity=core.Tanh)
 
 
 def GeneralGRUCell(candidate_transform,
-                   memory_transform=combinators.Copy,
+                   memory_transform=combinators.NoOp,
                    gate_nonlinearity=core.Sigmoid,
                    candidate_nonlinearity=core.Tanh,
                    dropout_rate_c=0.1,
@@ -117,7 +117,7 @@ def GeneralGRUCell(candidate_transform,
           # c_t (Candidate) branch
           combinators.Serial(
               combinators.Branch(
-                  combinators.Copy(),
+                  combinators.NoOp(),
                   # r_t (Reset) Branch
                   combinators.Serial(
                       candidate_transform(),
diff --git a/tensor2tensor/trax/models/atari_cnn.py b/tensor2tensor/trax/models/atari_cnn.py
index 3d38c9584..6f3beb29b 100644
--- a/tensor2tensor/trax/models/atari_cnn.py
+++ b/tensor2tensor/trax/models/atari_cnn.py
@@ -27,7 +27,7 @@ def AtariCnn(hidden_sizes=(32, 32), output_size=128):
   return tl.Serial(
       tl.Div(divisor=255.0),
       # Have 4 copies of the input, each one shifted to the right by one.
-      tl.Branch(tl.Copy(), tl.ShiftRight(),
+      tl.Branch(tl.NoOp(), tl.ShiftRight(),
                 tl.Serial(
                     tl.ShiftRight(),
                     tl.ShiftRight(),
diff --git a/tensor2tensor/trax/models/research/chunked_transformer.py b/tensor2tensor/trax/models/research/chunked_transformer.py
index cb200b0db..143ba7aa5 100644
--- a/tensor2tensor/trax/models/research/chunked_transformer.py
+++ b/tensor2tensor/trax/models/research/chunked_transformer.py
@@ -18,10 +18,47 @@
 from __future__ import division
 from __future__ import print_function
 
+import numpy as onp
+
 from tensor2tensor.trax import layers as tl
 from tensor2tensor.trax.backend import numpy as np
 
 
+# Chunked positional encoding.
+def _chunked_positional_encoding_new_params(input_shape, rng, max_len=2048):  # pylint: disable=invalid-name
+  """Helper: create positional encoding parameters."""
+  del rng
+  # Check if we are operating on chunked inputs by checking if the first
+  # shape is a list/tuple of shapes (otherwise it's an int or numpy array).
+  is_chunked = isinstance(input_shape[0], (list, tuple))
+  feature_depth = input_shape[0][-1] if is_chunked else input_shape[-1]
+  pe = onp.zeros((max_len, feature_depth), dtype=onp.float32)
+  position = onp.arange(0, max_len)[:, onp.newaxis]
+  div_term = onp.exp(
+      onp.arange(0, feature_depth, 2) * -(onp.log(10000.0) / feature_depth))
+  pe[:, 0::2] = onp.sin(position * div_term)
+  pe[:, 1::2] = onp.cos(position * div_term)
+  pe = pe[onp.newaxis, :, :]  # [1, max_len, feature_depth]
+  return np.array(pe)  # These are trainable parameters, initialized as above.
+
+
+@tl.layer(new_parameters=_chunked_positional_encoding_new_params,
+          stack_items_to_pass=0)
+def ChunkedPositionalEncoding(x, params, **unused_kwargs):
+  """Implements bare positional encoding."""
+  if not isinstance(x, (list, tuple)):  # non-chunked inputs
+    symbol_size = np.shape(x)[1]
+    return x + params[:, :symbol_size, :]
+  # Chunked case: apply to all chunks selecting as much as needed.
+  offset = 0
+  results = []
+  for chunk in x:
+    symbol_size = np.shape(chunk)[1]
+    results.append(chunk + params[:, offset:offset + symbol_size, :])
+    offset += symbol_size
+  return results
+
+
 # Chunked attention.
 def _chunked_selector_output_shape(  # pylint: disable=invalid-name
     input_shapes, selector=None, **unused_kwargs):
@@ -42,13 +79,13 @@ def _chunked_selector_output_shape(  # pylint: disable=invalid-name
     new_value_shape = (cur_value_shape[0], new_value_len, cur_value_shape[2])
     # Masks are (1, query-len, key-len).
     new_mask_shape = (1, query_shapes[i][1], new_key_len)
-    new_shape = ((query_shapes[i], new_key_shape, new_value_shape),
+    new_shape = (query_shapes[i], new_key_shape, new_value_shape,
                  new_mask_shape)
     result.append(new_shape)
   return tuple(result)
 
 
-@tl.layer(output_shape=_chunked_selector_output_shape)
+@tl.layer(output_shape=_chunked_selector_output_shape, stack_items_to_pass=0)
 def ChunkedAttentionSelector(x, params, selector=None, **kwargs):
   """Select which chunks to attend to in chunked attention.
 
@@ -60,7 +97,7 @@ def ChunkedAttentionSelector(x, params, selector=None, **kwargs):
     **kwargs: unused other arguments.
 
   Returns:
-    a list of elements of the form (q, k', v'), mask' where k', v' and mask' are
+    a list of elements of the form (q, k', v', mask') where k', v' and mask' are
     concatenations of k, v and identity-extended masks from selected chunks.
   """
   del params, kwargs
@@ -84,7 +121,7 @@ def ChunkedAttentionSelector(x, params, selector=None, **kwargs):
     new_mask_list = [np.ones(s, dtype=cur_mask.dtype) for s in new_mask_shapes]
     # We still use the current (often causal) mask for the final chunk.
     new_mask = np.concatenate(new_mask_list + [cur_mask], axis=2)
-    result.append(((queries[i], new_key, new_value), new_mask))
+    result.append((queries[i], new_key, new_value, new_mask))
   return tuple(result)
 
 
@@ -107,7 +144,7 @@ def ChunkedCausalMultiHeadedAttention(
   prepare_attention_input = tl.Serial(
       tl.Branch(
           tl.Branch(  # q = k = v = first input
-              tl.Copy(), tl.Copy(), tl.Copy()),
+              tl.NoOp(), tl.NoOp(), tl.NoOp()),
           tl.CausalMask(axis=-2),
       ),
       tl.Parallel(
@@ -116,7 +153,7 @@ def ChunkedCausalMultiHeadedAttention(
               tl.Dense(feature_depth),
               tl.Dense(feature_depth),
           ),
-          tl.Copy()
+          tl.NoOp()
       )
   )
   return tl.Serial(
@@ -222,7 +259,7 @@ def ChunkedTransformerLM(vocab_size,
       tl.ShiftRight(),
       tl.Map(tl.Embedding(feature_depth, vocab_size)),
       tl.Map(tl.Dropout(rate=dropout, mode=mode)),
-      tl.PositionalEncoding(max_len=max_len),
+      ChunkedPositionalEncoding(max_len=max_len),  # pylint: disable=no-value-for-parameter
       tl.Serial(*stack),
       tl.Map(tl.LayerNorm()),
       tl.Map(tl.Dense(vocab_size)),
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index 2b7f801bc..bf0e6aba1 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -114,7 +114,7 @@ def WideResnetBlock(channels, strides=(1, 1), channel_mismatch=False):
       tl.BatchNorm(),
       tl.Relu(),
       tl.Conv(channels, (3, 3), padding='SAME'))
-  shortcut = tl.Copy() if not channel_mismatch else tl.Conv(
+  shortcut = tl.NoOp() if not channel_mismatch else tl.Conv(
       channels, (3, 3), strides, padding='SAME')
   return tl.Residual(main, shortcut=shortcut)
 
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 7cf3486b3..7002b622f 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -59,10 +59,10 @@ def EncoderLayer(feature_depth,
   """
   return tl.Serial(
       tl.Residual(  # Attention block here.
-          tl.Parallel(tl.LayerNorm(), tl.Copy()),
+          tl.Parallel(tl.LayerNorm(), tl.NoOp()),
           tl.MultiHeadedAttention(feature_depth, num_heads=num_heads,
                                   dropout=dropout, mode=mode),
-          tl.Parallel(tl.Dropout(rate=dropout, mode=mode), tl.Copy())
+          tl.Parallel(tl.Dropout(rate=dropout, mode=mode), tl.NoOp())
       ),
       tl.Parallel(
           ResidualFeedForward(
@@ -135,7 +135,7 @@ def DecoderLayer(feature_depth,
   return tl.Serial(
       tl.Residual(  # Self-attention block.
           tl.LayerNorm(),
-          tl.Branch(tl.Copy(), tl.CausalMask(axis=-2)),  # Create mask.
+          tl.Branch(tl.NoOp(), tl.CausalMask(axis=-2)),  # Create mask.
           tl.MultiHeadedAttention(feature_depth, num_heads=num_heads,
                                   dropout=dropout, mode=mode),
           tl.Select(0),  # Drop the mask.
@@ -206,7 +206,7 @@ def EncoderDecoderLayer(feature_depth,
   # Decoder self-attending to decoder.
   self_attention = tl.Residual(
       tl.LayerNorm(),
-      tl.Branch(tl.Copy(), tl.CausalMask(axis=-2)),  # create mask
+      tl.Branch(tl.NoOp(), tl.CausalMask(axis=-2)),  # create mask
       tl.MultiHeadedAttention(feature_depth, num_heads=num_heads,
                               dropout=dropout, mode=mode),
       tl.Select(0),  # drop mask
@@ -214,21 +214,21 @@ def EncoderDecoderLayer(feature_depth,
   )
   # Decoder attending to encoder.
   encoder_decoder_attention = tl.Serial(
-      tl.Select(((2, 0, 0), 1)),  # ((dec, enc, enc), mask)
-      tl.MultiHeadedAttentionQKV(  # ((q, k, v), mask) --> new, mask
+      tl.Select((2, 0, 0, 1)),  # (dec, enc, enc, mask)
+      tl.MultiHeadedAttentionQKV(  # (q, k, v, mask) --> new, mask
           feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),
       tl.Select(0),  # drop the mask
       tl.Dropout(rate=dropout, mode=mode),
   )
   return tl.Serial(
-      tl.Parallel(tl.Copy(), tl.Copy(), self_attention),
-      tl.Branch(tl.Copy(), encoder_decoder_attention),
-      tl.UnnestBranches(),   # (encoder, mask, old_act, new_act)
-      tl.Select((0, 1, (2, 3))),
+      tl.Parallel(tl.NoOp(), tl.NoOp(), self_attention),
+      tl.Branch(tl.NoOp(), encoder_decoder_attention),
+      tl.Select(inputs=(('encoder', 'mask', 'old_act'), 'new_act'),
+                output=('encoder', 'mask', ('old_act', 'new_act'))),
       tl.Parallel(  # Residual after encoder-decoder attention.
-          tl.Copy(), tl.Copy(), tl.Add()),
+          tl.NoOp(), tl.NoOp(), tl.Add()),
       tl.Parallel(  # Feed-forward on the third component (decoder).
-          tl.Copy(), tl.Copy(), ResidualFeedForward(
+          tl.NoOp(), tl.NoOp(), ResidualFeedForward(
               feature_depth, feedforward_depth, dropout, mode=mode)
       )
   )
@@ -270,18 +270,18 @@ def Transformer(vocab_size,
       tl.Serial(*[EncoderLayer(feature_depth, feedforward_depth, num_heads,
                                dropout, mode)
                   for _ in range(num_layers)]),
-      tl.Parallel(tl.LayerNorm(), tl.Copy())
+      tl.Parallel(tl.LayerNorm(), tl.NoOp())
   )
   stack = [EncoderDecoderLayer(feature_depth, feedforward_depth, num_heads,
                                dropout, mode)
            for _ in range(num_layers)]
   return tl.Serial(
-      tl.Parallel(tl.Copy(), tl.ShiftRight()),
+      tl.Parallel(tl.NoOp(), tl.ShiftRight()),
       tl.Parallel(encoder, embedding),
-      tl.UnnestBranches(),  # (encoder, encoder_mask, decoder_input)
-      tl.Select((0, (1, 2), 2)),
+      tl.Select(inputs=(('encoder', 'mask'), 'decoder'),
+                output=('encoder', ('mask', 'decoder'), 'decoder')),
       tl.Parallel(  # (encoder_mask, decoder_input) -> encoder-decoder mask
-          tl.Copy(), tl.EncoderDecoderMask(), tl.Copy()),
+          tl.NoOp(), tl.EncoderDecoderMask(), tl.NoOp()),
       tl.Serial(*stack),
       tl.Select(2),  # Drop encoder and mask.
       tl.LayerNorm(),

From e61eba618c04a6c423f6e78979d089033a8fa7c7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 21 May 2019 01:33:07 -0700
Subject: [PATCH 2045/2720] carry early tokens in parallel scheduled sampling.

PiperOrigin-RevId: 249207421
---
 tensor2tensor/utils/t2t_model.py | 51 +++++++++++++++++++++++++-------
 1 file changed, 40 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 5afdfa781..547067692 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1838,22 +1838,50 @@ def sample(x):
       return tf.to_int32(reshaped_samples)
 
     # TODO(duckworthd): Move to scheduled_sampling.py.
-    def mix_gold_sampled(gold_targets, sampled_targets, mixin_prob):
+    def mix_gold_sampled(gold_targets,
+                         sampled_targets,
+                         mixin_prob,
+                         i,
+                         prev_new_targets):
       """Interleave sampled and gold tokens randomly."""
-      return tf.where(
-          tf.less(
-              tf.random_uniform(common_layers.shape_list(sampled_targets)),
-              mixin_prob),
+      # Resample each location iid.
+      should_use_sampled_targets = tf.less(
+          tf.random_uniform(common_layers.shape_list(sampled_targets)),
+          mixin_prob)
+      mixed_targets = tf.where(
+          should_use_sampled_targets,
           sampled_targets,
           gold_targets)
 
+      # Reuse sample tokens for earlier timesteps.
+      new_targets = tf.where(
+          is_later_timestep(gold_targets, i),
+          mixed_targets,
+          prev_new_targets)
+      return new_targets
+
+    # TODO(duckworthd): Move to scheduled_sampling.py.
+    def is_later_timestep(x, pass_idx):
+      """Constructs mask based on timestep."""
+      assert x.shape.ndims == 4, x.shape
+      x_shape = tf.shape(x)
+      batch_size = x_shape[0]
+      num_timesteps = x_shape[1]
+      timesteps = tf.range(num_timesteps)
+      timesteps = tf.reshape(timesteps, [1, num_timesteps, 1, 1])
+      timesteps = tf.tile(timesteps, [batch_size, 1, 1, 1])
+      return tf.greater_equal(timesteps, pass_idx)
+
     # TODO(duckworthd): Move to scheduled_sampling.py.
-    def parallel_scheduled_sampling_pass(features, logits, mixin_prob):
+    def parallel_scheduled_sampling_pass(
+        i, prev_new_targets, features, logits, mixin_prob):
       """Generate scheduled sampling results."""
       sampled_targets = sample(logits)
       new_targets = mix_gold_sampled(features["targets"],
                                      sampled_targets,
-                                     mixin_prob)
+                                     mixin_prob,
+                                     i,
+                                     prev_new_targets)
       new_targets = tf.stop_gradient(new_targets)  # Treat new_targets as given.
       new_features = copy.copy(features)
       new_features["targets"] = new_targets
@@ -1879,7 +1907,7 @@ def parallel_scheduled_sampling_pass(features, logits, mixin_prob):
         else:
           new_losses["training"] = 0.0
 
-      return new_logits, new_losses
+      return new_targets, new_logits, new_losses
 
     tf.logging.info("Using scheduled sampling.")
     tf.logging.info("Warming scheduled sampling up with schedule: %s",
@@ -1912,9 +1940,10 @@ def parallel_scheduled_sampling_pass(features, logits, mixin_prob):
           "hparams.scheduled_sampling_prob > 0.0")
       new_logits = logits
       new_losses = losses
-      for _ in range(hparams.scheduled_sampling_num_passes):
-        new_logits, new_losses = parallel_scheduled_sampling_pass(
-            features, new_logits, mixin_prob)
+      prev_new_targets = features["targets"]
+      for i in range(hparams.scheduled_sampling_num_passes):
+        prev_new_targets, new_logits, new_losses = parallel_scheduled_sampling_pass(
+            i, prev_new_targets, features, new_logits, mixin_prob)
       return new_logits, new_losses
     else:
       raise ValueError(

From ebf34e7ef5e5a6708442444c297b0c76d78a838b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 21 May 2019 06:36:03 -0700
Subject: [PATCH 2046/2720] Internal

PiperOrigin-RevId: 249239055
---
 tensor2tensor/layers/common_layers.py | 85 ++++++++++++++++++++++-----
 1 file changed, 69 insertions(+), 16 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index c1cd291c6..3670f7134 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1868,6 +1868,39 @@ def padded_cross_entropy(logits,
     return tf.reduce_sum(xent * weights), tf.reduce_sum(weights)
 
 
+def gather_tensor_by_mixture_index(value,
+                                   mixture_indices,
+                                   batch_size,
+                                   num_mixtures,
+                                   reshape=True):
+  """Gather the elements of a tensor, based on the mixture element id provided.
+
+  The tensor should be shaped as (num_mixtures * batch_size, dim2, dim3...),
+  and the mixture indices should be (batch_size), holding one mixture_id for
+  each element in the batch
+
+  Args:
+    value: a `Tensor` with shape `[num_mixtures * batch, dim2, dim3 ...]`. If
+    reshape is false, it should be [num_mixtures, batch, dim3, dim4 ..]
+    mixture_indices: `[batch_size]`.
+    batch_size: an int `Scalar`.
+    num_mixtures: an int `Scalar`.
+    reshape: bool
+
+  Returns:
+    selected_values: a `Tensor`.  Selected values from original tensor
+
+  """
+  original_shape = shape_list(value)
+  individual_element_indices = tf.range(batch_size)
+  stacked_mixture_element_indices = tf.stack(
+      (mixture_indices, individual_element_indices), -1)
+  if reshape:
+    value = tf.reshape(value, [num_mixtures, -1] + original_shape[1:])
+  selected_values = tf.gather_nd(value, stacked_mixture_element_indices)
+  return selected_values
+
+
 def padded_cross_entropy_mixture(logits,
                                  labels,
                                  label_smoothing,
@@ -1905,15 +1938,17 @@ def padded_cross_entropy_mixture(logits,
   Raises:
     ValueError: in case of unsupported argument types.
   """
+
+  (logits, mixture_labels, supervised_mode) = logits
+
   logit_shapes = shape_list(
       logits)  # batch_size * num_mixtures, timesteps, 1, 1, vocab_size
   batch_size = tf.cast(logit_shapes[0] / num_mixtures, dtype=tf.int32)
-  timesteps = logit_shapes[1]
-  vocab_size = logit_shapes[4]
 
   new_shape_for_xent = [num_mixtures] + shape_list(labels)
   labels = tf.tile(labels, [num_mixtures, 1, 1, 1])
 
+  # get xent loss for all mixtures
   xent, weights = padded_cross_entropy(logits, labels, label_smoothing,
                                        weights_fn, reduce_sum, cutoff, gaussian)
 
@@ -1926,38 +1961,56 @@ def padded_cross_entropy_mixture(logits,
 
   # if we need to compute the best logits
   if return_best_logits:
-    best_mixture_indices = tf.cast(tf.argmin(xent, 0), dtype=tf.int32)
-    individual_element_indices = tf.range(batch_size)
-    stacked_mixture_element_indices = tf.stack((tf.squeeze(
-        best_mixture_indices, axis=[1, 2]), individual_element_indices), -1)
-    best_logits = tf.reshape(logits,
-                             [num_mixtures, -1, timesteps, 1, 1, vocab_size])
-    best_logits = tf.gather_nd(best_logits, stacked_mixture_element_indices)
-    best_logits = tf.reshape(best_logits,
-                             [batch_size, timesteps, 1, 1, vocab_size])
+    if supervised_mode:
+      return_mixture_indices = tf.squeeze(
+          tf.cast(tf.argmin(xent, 0), dtype=tf.int32), axis=[1, 2])
+    else:
+      return_mixture_indices = mixture_labels
+    best_logits = gather_tensor_by_mixture_index(logits, return_mixture_indices,
+                                                 batch_size, num_mixtures)
 
   with tf.control_dependencies([
       tf.assert_equal(
           tf.shape(xent)[:3], [num_mixtures, batch_size, 1],
-          message="Each batch element should have a probability value for each mixture element"
+          message="Each batch element should have a probability value for "
+          "each mixture element"
       )
   ]):
-    xent_min = tf.reduce_min(xent, axis=0)
+    best_mixtures = tf.squeeze(
+        tf.cast(tf.argmin(xent, 0), dtype=tf.int32), axis=[1, 2])
+    if mixture_labels is not None:
+      mixture_accuracy = tf.metrics.accuracy(
+          mixture_labels, best_mixtures, name="mixture_accuracy")
+      tf.summary.scalar("mixture_acc_plot", mixture_accuracy[1])
+    if supervised_mode:
+      xent_min = gather_tensor_by_mixture_index(
+          xent, mixture_labels, batch_size, num_mixtures, reshape=False)
+    else:
+      xent_min = tf.reduce_min(xent, axis=0)
     xent_max = tf.reduce_max(xent, axis=0)
     weights = tf.reduce_mean(weights, axis=0)
 
   with tf.control_dependencies([
       tf.assert_equal(
           tf.shape(xent_min)[0], [batch_size],
-          message="There should be batch_size elements after selecting best mixture probabilities"
+          message="There should be batch_size elements after selecting best "
+          "mixture probabilities"
       )
   ]):
     summed_xent_min = tf.reduce_sum(xent_min)
     summed_xent_max = tf.reduce_sum(xent_max)
     summed_weights = tf.reduce_sum(weights)
 
-    tf.summary.scalar("mixture_xents_min", summed_xent_min / summed_weights)
-    tf.summary.scalar("mixture_xents_max", summed_xent_max / summed_weights)
+    for mixture in range(num_mixtures):
+      num_assigned_mixtures = tf.reduce_sum(
+          tf.cast(tf.equal(best_mixtures, mixture), tf.int32))
+      tf.summary.scalar("assigned_mixture_%d" % (mixture),
+                        num_assigned_mixtures / batch_size)
+
+    tf.summary.scalar("selected_mixture_xents_value",
+                      summed_xent_min / summed_weights)
+    tf.summary.scalar("max_mixture_xents_value",
+                      summed_xent_max / summed_weights)
 
   if return_best_logits:
     return summed_xent_min, summed_weights, best_logits

From b639d16e48ec1c078f5808798f94b40c870d897c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 21 May 2019 21:14:02 -0700
Subject: [PATCH 2047/2720] sampled top-k hard attention.

PiperOrigin-RevId: 249381613
---
 tensor2tensor/layers/common_attention.py      | 52 +++++++++++++------
 tensor2tensor/layers/common_attention_test.py | 20 +++++--
 tensor2tensor/layers/transformer_layers.py    |  1 +
 tensor2tensor/models/transformer.py           |  3 ++
 4 files changed, 56 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 07c66c73c..7778f0d3d 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1472,12 +1472,18 @@ def grouped_attention_multihead(query_antecedent,
     return o, extra_loss
 
 
-def harden_attention_weights(weights, hard_attention_k):
-  """Make attention weights non-0 only on the top-hard_attention_k ones."""
+def harden_attention_weights(weights, k, gumbel_noise_weight):
+  """Make attention weights non-0 only on the top k ones."""
+  if gumbel_noise_weight > 0.:
+    gumbel_noise = -tf.log(-tf.log(tf.random_uniform(tf.shape(weights),
+                                                     minval=1e-5,
+                                                     maxval=1 - 1e-5)))
+    weights += gumbel_noise * gumbel_noise_weight
+
   # Subtract the top-kth weight and zero-out all lower ones.
   # Note that currently in case of numerical ties it will retain more
   # than k elements. In the future, we may want to avoid this.
-  weights -= common_layers.top_kth_iterative(weights, hard_attention_k)
+  weights -= common_layers.top_kth_iterative(weights, k)
   weights = tf.nn.relu(weights)
   # Re-normalize the weights.
   weights_sum = tf.reduce_sum(weights, axis=-1, keep_dims=True)
@@ -1498,7 +1504,8 @@ def dot_product_attention(q,
                           dropout_broadcast_dims=None,
                           activation_dtype=None,
                           weight_dtype=None,
-                          hard_attention_k=0):
+                          hard_attention_k=0,
+                          gumbel_noise_weight=0.0):
   """Dot-product attention.
 
   Args:
@@ -1522,6 +1529,9 @@ def dot_product_attention(q,
       mixed precision.
     weight_dtype: The dtype weights are stored in when using mixed precision
     hard_attention_k: integer, if > 0 triggers hard attention (picking top-k)
+    gumbel_noise_weight: if > 0, apply Gumbel noise with weight
+      `gumbel_noise_weight` before picking top-k. This is a no op if
+      hard_attention_k <= 0.
 
   Returns:
     Tensor with shape [..., length_q, depth_v].
@@ -1536,7 +1546,8 @@ def dot_product_attention(q,
     logits = maybe_upcast(logits, activation_dtype, weight_dtype)
     weights = tf.nn.softmax(logits, name="attention_weights")
     if hard_attention_k > 0:
-      weights = harden_attention_weights(weights, hard_attention_k)
+      weights = harden_attention_weights(weights, hard_attention_k,
+                                         gumbel_noise_weight)
     weights = common_layers.cast_like(weights, q)
     if save_weights_to is not None:
       save_weights_to[scope.name] = weights
@@ -1630,7 +1641,8 @@ def dot_product_attention_relative(q,
                                    make_image_summary=True,
                                    cache=False,
                                    allow_memory=False,
-                                   hard_attention_k=0):
+                                   hard_attention_k=0,
+                                   gumbel_noise_weight=0.0):
   """Calculate relative position-aware dot-product self-attention.
 
   The attention calculation is augmented with learned representations for the
@@ -1655,6 +1667,9 @@ def dot_product_attention_relative(q,
       the length dimension of k/v/bias may be longer than the queries, and it is
       assumed that the extra memory entries precede the non-memory entries.
     hard_attention_k: integer, if > 0 triggers hard attention (picking top-k)
+    gumbel_noise_weight: if > 0, apply Gumbel noise with weight
+      `gumbel_noise_weight` before picking top-k. This is a no op if
+      hard_attention_k <= 0.
 
   Returns:
     A Tensor.
@@ -1692,7 +1707,8 @@ def dot_product_attention_relative(q,
       logits += bias
     weights = tf.nn.softmax(logits, name="attention_weights")
     if hard_attention_k > 0:
-      weights = harden_attention_weights(weights, hard_attention_k)
+      weights = harden_attention_weights(weights, hard_attention_k,
+                                         gumbel_noise_weight)
     if save_weights_to is not None:
       save_weights_to[scope.name] = weights
       save_weights_to[scope.name + "/logits"] = logits
@@ -3992,6 +4008,7 @@ def multihead_attention(query_antecedent,
                         recurrent_memory=None,
                         chunk_number=None,
                         hard_attention_k=0,
+                        gumbel_noise_weight=0.0,
                         max_area_width=1,
                         max_area_height=1,
                         memory_height=1,
@@ -4056,6 +4073,9 @@ def multihead_attention(query_antecedent,
     chunk_number: an optional integer Tensor with shape [batch] used to operate
       the recurrent_memory.
     hard_attention_k: integer, if > 0 triggers hard attention (picking top-k).
+    gumbel_noise_weight: if > 0, apply Gumbel noise with weight
+      `gumbel_noise_weight` before picking top-k. This is a no op if
+      hard_attention_k <= 0.
     max_area_width: the max width allowed for an area.
     max_area_height: the max height allowed for an area.
     memory_height: the height of the memory.
@@ -4198,13 +4218,14 @@ def multihead_attention(query_antecedent,
             area_value_mode=area_value_mode,
             training=training)
       else:
-        x = dot_product_attention(q, k, v, bias, dropout_rate, image_shapes,
-                                  save_weights_to=save_weights_to,
-                                  make_image_summary=make_image_summary,
-                                  dropout_broadcast_dims=dropout_broadcast_dims,
-                                  activation_dtype=kwargs.get(
-                                      "activation_dtype"),
-                                  hard_attention_k=hard_attention_k)
+        x = dot_product_attention(
+            q, k, v, bias, dropout_rate, image_shapes,
+            save_weights_to=save_weights_to,
+            make_image_summary=make_image_summary,
+            dropout_broadcast_dims=dropout_broadcast_dims,
+            activation_dtype=kwargs.get("activation_dtype"),
+            hard_attention_k=hard_attention_k,
+            gumbel_noise_weight=gumbel_noise_weight)
     elif attention_type == "dot_product_relative":
       x = dot_product_attention_relative(
           q,
@@ -4218,7 +4239,8 @@ def multihead_attention(query_antecedent,
           make_image_summary=make_image_summary,
           cache=cache is not None,
           allow_memory=recurrent_memory is not None,
-          hard_attention_k=hard_attention_k)
+          hard_attention_k=hard_attention_k,
+          gumbel_noise_weight=gumbel_noise_weight)
     elif attention_type == "dot_product_unmasked_relative_v2":
       x = dot_product_unmasked_self_attention_relative_v2(
           q,
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 174dceb20..fe31f4dd7 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -44,22 +44,32 @@ def testAddPositionalEmbedding(self):
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 3, 12))
 
+  @parameterized.named_parameters(
+      ("hard_top_k", 0.0),
+      ("sampled_top_k_default", 1.0),
+      ("sampled_top_k_2", 2.0),
+  )
   @test_utils.run_in_graph_and_eager_modes()
-  def testHardenAttentionWeights(self):
+  def testHardenAttentionWeights(self, gumbel_noise_weight):
     x = np.random.rand(5, 3, 12)
     y = common_attention.harden_attention_weights(
-        tf.nn.softmax(tf.constant(x, dtype=tf.float32)), 3)
+        tf.nn.softmax(tf.constant(x, dtype=tf.float32)), 3, gumbel_noise_weight)
     res = self.evaluate(y)
     self.assertEqual(res.shape, (5, 3, 12))
 
+  @parameterized.named_parameters(
+      ("hard_top_k", -0.5),
+      ("sampled_top_k", 0.5),
+  )
   @test_utils.run_in_graph_and_eager_modes()
-  def testHardenAttentionAllZeros(self):
+  def testHardenAttentionAllZeros(self, gumbel_noise_weight):
     """Check if the hardening code does not divide by zero for all zeros."""
     x = np.zeros((5, 3, 12), dtype=np.float32)
     y = common_attention.harden_attention_weights(
-        tf.constant(x, dtype=tf.float32), 3)
+        tf.constant(x, dtype=tf.float32), 3, gumbel_noise_weight)
     res = self.evaluate(y)
-    self.assertAllClose(res, x)
+    if gumbel_noise_weight <= 0.0:
+      self.assertAllClose(res, x)
 
   @parameterized.parameters(
       {"input_shape": (5, 3, 12)},
diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 9b23ecad5..59a0ced56 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -213,6 +213,7 @@ def transformer_encoder(encoder_input,
               activation_dtype=hparams.get("activation_dtype", "float32"),
               weight_dtype=hparams.get("weight_dtype", "float32"),
               hard_attention_k=hparams.get("hard_attention_k", 0),
+              gumbel_noise_weight=hparams.get("gumbel_noise_weight", 0.0),
               max_area_width=max_area_width,
               max_area_height=max_area_height,
               memory_height=memory_height,
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 1eac82db5..5ce060686 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1404,6 +1404,7 @@ def transformer_decoder_layer(decoder_input,
           recurrent_memory=recurrent_memory,
           chunk_number=chunk_number,
           hard_attention_k=hparams.get("hard_attention_k", 0),
+          gumbel_noise_weight=hparams.get("gumbel_noise_weight", 0.0),
           max_area_width=max_area_width,
           max_area_height=max_area_height,
           memory_height=memory_height,
@@ -1439,6 +1440,7 @@ def transformer_decoder_layer(decoder_input,
             weight_dtype=hparams.get("weight_dtype", "float32"),
             layer_collection=layer_collection,
             hard_attention_k=hparams.get("hard_attention_k", 0),
+            gumbel_noise_weight=hparams.get("gumbel_noise_weight", 0.0),
             max_area_width=max_area_width,
             max_area_height=max_area_height,
             memory_height=memory_height,
@@ -1673,6 +1675,7 @@ def transformer_base_v1():
   hparams.add_hparam("unidirectional_encoder", False)
   # For hard attention.
   hparams.add_hparam("hard_attention_k", 0)
+  hparams.add_hparam("gumbel_noise_weight", 0.0)
   return hparams
 
 

From a877a71af0f0ac63a61f9640ebde52fe219acc53 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Wed, 22 May 2019 10:32:15 -0700
Subject: [PATCH 2048/2720] [TRAX] Move @pmap outside of predict function to
 avoid unnecessary recompilation during eval.

If you write in JAX:
```
def f(y):

  @jit
  def g(x):
    return x + 1

  return g(y)
```
this means that *each* time you run `f` you create a new `g`. They are different Python functions, after all. JAX attaches the jit cache to the function `g`, so this means that each time you call `f` we build and `jit`-compile a new instance of `g`.

The fix is easy enough, just hoist the definition of `g` outside of `f`, i.e.:
```
@jit
def g(x):
  return x + 1

def f(y):
  return g(y)
```

PiperOrigin-RevId: 249474567
---
 tensor2tensor/trax/trax.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 0b8c93dd6..b1606b45c 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -295,16 +295,18 @@ def epochs(steps=None, epoch_steps=1):
 
 def _jit_predict_fun(model_predict, num_devices):
   """Use jit on model_predict if required."""
+
+  # Multi-devices, pmap and run.
+  @functools.partial(backend.pmap, axis_name="batch")
+  def mapped_predict(x, params, rng):
+    return model_predict(x, params, rng=rng)
+
   def predict(x, params=(), rng=None):
     """Predict function jited and parallelized as requested."""
     # On one device, jit and run.
     if num_devices == 1:
       return backend.jit(model_predict)(x, params, rng=rng)
 
-    # Multi-devices, pmap and run.
-    @functools.partial(backend.pmap, axis_name="batch")
-    def mapped_predict(x, params, rng):
-      return model_predict(x, params, rng=rng)
     pred = mapped_predict(
         reshape_by_device(x, num_devices),
         params,

From aa838f92d9f8df5ebd26423e5f0ea95ab22b0480 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 22 May 2019 10:35:19 -0700
Subject: [PATCH 2049/2720] Use Adam for single-gpu testing of Transformer LM.

PiperOrigin-RevId: 249475196
---
 tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
index 4650208f2..fe5b38ae1 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
@@ -21,8 +21,8 @@ masked_mean.mask_id = 0
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-MultifactorSchedule.constant = 0.3
-MultifactorSchedule.factors = 'constant * linear_warmup'
+MultifactorSchedule.constant = 0.1
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
 MultifactorSchedule.warmup_steps = 8000
 
 # Parameters for preprocess_fun:
@@ -37,7 +37,7 @@ train.eval_frequency = 1000
 train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
-train.optimizer = @trax.optimizers.SM3
+train.optimizer = @trax.optimizers.Adam
 train.run_debug_step = False
 train.train_steps = 100000
 

From 4cdecce421806b8cc6621f5ae9e556f02b8284c5 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 22 May 2019 12:01:11 -0700
Subject: [PATCH 2050/2720] Start transitioning Transformer to stack-based
 layer semantics. Add needed combinators and corrections to exising layers.
 Warning: as this progresses some combinators may still change or get removed.

PiperOrigin-RevId: 249492509
---
 tensor2tensor/trax/layers/attention.py        |   2 +-
 tensor2tensor/trax/layers/base.py             |   2 +
 tensor2tensor/trax/layers/combinators.py      | 134 ++++++++++++++++--
 tensor2tensor/trax/layers/core.py             |   4 +-
 .../models/research/chunked_transformer.py    |  13 +-
 tensor2tensor/trax/models/transformer.py      |  46 +++---
 6 files changed, 154 insertions(+), 47 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 5aa7b2e9a..022a0ff6a 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -210,7 +210,7 @@ def MultiHeadedAttentionQKV(
       PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
           feature_depth=feature_depth, num_heads=num_heads,
           dropout=dropout, mode=mode),
-      combinators.Parallel(core.Dense(feature_depth), combinators.NoOp())
+      core.Dense(feature_depth),
   )
 
 
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 1f9e167d5..de6a225ce 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -178,6 +178,8 @@ def _apply_to_first_n(f, x, n):
   if n == 1:
     argument = argument[0]
   result = f(argument)
+  if not rest:
+    return result
   if n == 1:
     result = [result]
   result = list(result) + list(rest)
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 2361e4bdc..de62e8e63 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -67,6 +67,18 @@ def NoOp(x, **unused_kwargs):
   return x
 
 
+def _print_shape(x, message='PrintShape'):  # pylint: disable=invalid-name
+  print(message + ' ; stack shape = ' + str(x))
+  return x
+
+
+@base.layer(output_shape=_print_shape, stack_items_to_pass=0)
+def PrintShape(x, message='PrintShape', **unused_kwargs):
+  """NoOp layer that prints the shape of the stack."""
+  _print_shape(base.shapes(x), message=message)
+  return x
+
+
 def _dup(x):  # pylint: disable=invalid-name
   """Helper: copy the top element of a list or a tuple."""
   if isinstance(x, list):
@@ -81,6 +93,72 @@ def Dup(x, **unused_kwargs):
   return _dup(x)
 
 
+def _swap(x):  # pylint: disable=invalid-name
+  """Helper: swap the top two elements of a list or a tuple."""
+  if isinstance(x, list):
+    return [x[1], x[0]] + x[2:]
+  assert isinstance(x, tuple)
+  return tuple([x[1], x[0]] + list(x[2:]))
+
+
+@base.layer(output_shape=_swap, stack_items_to_pass=0)
+def Swap(x, **unused_kwargs):
+  """Swap the first two element on the stack."""
+  return _swap(x)
+
+
+def _top_shape(x_shape):  # pylint: disable=invalid-name
+  """Helper: shape of top element of a stack."""
+  if isinstance(x_shape[0], (list, tuple)):
+    return x_shape[0]
+  return x_shape
+
+
+@base.layer(output_shape=_top_shape, stack_items_to_pass=0)
+def _Top(x, **unused_kwargs):
+  """Top element from the stack."""
+  if isinstance(x, (list, tuple)):
+    return x[0]
+  return x
+
+
+def _drop(x):  # pylint: disable=invalid-name
+  """Helper: pop top element of a stack (make it a non-list if length is 1)."""
+  result = x[1:]
+  if len(result) == 1:
+    return result[0]
+  return result
+
+
+@base.layer(output_shape=_drop, stack_items_to_pass=0)
+def Drop(x, **unused_kwargs):
+  """Drop first element from the stack."""
+  return _drop(x)
+
+
+def _flatten_shape(x_shape):  # pylint: disable=invalid-name
+  """Helper: shape of the flatten operation."""
+  shapes = []
+  for shape in x_shape:
+    if isinstance(shape[0], (list, tuple)):
+      shapes.extend(shape)
+    else:
+      shapes.append(shape)
+  return tuple(shapes)
+
+
+@base.layer(output_shape=_flatten_shape, stack_items_to_pass=0)
+def Flatten(xs, **unused_kwargs):
+  """Flatten lists."""
+  res = []
+  for x in xs:
+    if isinstance(x, (list, tuple)):
+      res.extend(list(x))
+    else:
+      res.append(x)
+  return tuple(res)
+
+
 # Re-ordering layer.
 class Select(base.Layer):
   """Select elements from a tuple or create another tuple from them.
@@ -242,13 +320,43 @@ def _nested_op(inputs, op):  # pylint: disable=invalid-name
   return tuple(result_list)
 
 
-def _nested_sum(inputs):  # pylint: disable=invalid-name
-  return _nested_op(inputs=inputs, op=sum)
+def _binary_op(inputs, op):  # pylint: disable=invalid-name
+  """Helper: apply op to the first 2 elements."""
+  xs, rest = inputs[:2], inputs[2:]
+  s = _nested_op(xs, op)
+  if not rest:
+    return s
+  if not isinstance(s, (list, tuple)):
+    s = [s]
+  res = list(s) + list(rest)
+  # TODO(lukaszkaiser): should we drop this tuple/list distinction?
+  if isinstance(s, tuple):
+    res = tuple(res)
+  return res
+
+
+def _binary_op_shape(stack_shape):  # pylint: disable=invalid-name
+  """Helper: shape for the top-two operation above (shape-preserving op)."""
+  if len(stack_shape) == 2:
+    return stack_shape[0]
+  return tuple([stack_shape[0]] + list(stack_shape[2:]))
+
+
+@base.layer(output_shape=_binary_op_shape, stack_items_to_pass=0)
+def Add(x, **unused_kwargs):
+  """Add first and second element on the stack."""
+  # Here x is a list of tensors of the same shape, or nested structures.
+  return _binary_op(x, op=sum)
+
+
+@base.layer(output_shape=_binary_op_shape, stack_items_to_pass=0)
+def Multiply(x, **unused_kwargs):
+  """Multiply first and second element on the stack."""
+  return _binary_op(x, op=lambda xs: six.moves.reduce(operator.mul, xs))
 
 
-def _nested_product(inputs):  # pylint: disable=invalid-name
-  return _nested_op(
-      inputs=inputs, op=lambda xs: six.moves.reduce(operator.mul, xs))
+def _nested_sum(inputs):  # pylint: disable=invalid-name
+  return _nested_op(inputs=inputs, op=sum)
 
 
 def _first_from_tuple_or_dict(tuple_or_dict):  # pylint: disable=invalid-name
@@ -258,18 +366,12 @@ def _first_from_tuple_or_dict(tuple_or_dict):  # pylint: disable=invalid-name
 
 
 @base.layer(output_shape=_first_from_tuple_or_dict, stack_items_to_pass=0)
-def Add(x, **unused_kwargs):
+def AddAll(x, **unused_kwargs):
   """Add branches elementwise."""
   # Here x is a list of tensors of the same shape, or nested structures.
   return _nested_sum(x)
 
 
-@base.layer(output_shape=_first_from_tuple_or_dict, stack_items_to_pass=0)
-def Multiply(x, **unused_kwargs):
-  """Multiply branches elementwise."""
-  return _nested_product(x)
-
-
 @base.layer(output_shape=_first_from_tuple_or_dict, stack_items_to_pass=0)
 def Gate(x, **unused_kwargs):
   """Implements a gating function on a (memory, gate, candidate) tuple.
@@ -390,15 +492,17 @@ def new_parameters(self, input_shape, rng):
 
 def Residual(*layers, **kwargs):
   """Constructs a residual version of layers, summing input to layers output."""
-  shortcut = kwargs.get('shortcut', NoOp())  # pylint: disable=no-value-for-parameter
+  shortcut = kwargs.get('shortcut', _Top())  # pylint: disable=no-value-for-parameter
   if len(layers) > 1:
     return Serial(
-        Branch(Serial(*layers), shortcut),
+        Branch(shortcut, Serial(*layers)),
+        Flatten(),  # pylint: disable=no-value-for-parameter
         Add()  # pylint: disable=no-value-for-parameter
     )
   elif len(layers) == 1:
     return Serial(
-        Branch(layers[0], shortcut),
+        Branch(shortcut, layers[0]),
+        Flatten(),  # pylint: disable=no-value-for-parameter
         Add()  # pylint: disable=no-value-for-parameter
     )
   else:
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 0b5a075a8..14fa7184f 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -97,7 +97,7 @@ def call(self, x, params, **kwargs):
     w, b = params
     return np.dot(x, w) + b
 
-  def output_shape(self, input_shape):
+  def output_shape_fun(self, input_shape):
     return tuple(input_shape[:-1]) + (self._units,)
 
   def new_parameters(self, input_shape, rng):
@@ -124,7 +124,7 @@ def call(self, x, params, **kwargs):
     del kwargs
     return np.take(params, x, axis=0)
 
-  def output_shape(self, input_shape):
+  def output_shape_fun(self, input_shape):
     return tuple(input_shape) + (self._feature_depth,)
 
   def new_parameters(self, input_shape, rng):
diff --git a/tensor2tensor/trax/models/research/chunked_transformer.py b/tensor2tensor/trax/models/research/chunked_transformer.py
index 143ba7aa5..aa3a3ee5c 100644
--- a/tensor2tensor/trax/models/research/chunked_transformer.py
+++ b/tensor2tensor/trax/models/research/chunked_transformer.py
@@ -167,12 +167,21 @@ def ChunkedCausalMultiHeadedAttention(
   )
 
 
+# Chunked residual.
+def Residual(*layers, **unused_kwargs):
+  """Constructs a residual version of layers, summing input to layers output."""
+  return tl.Serial(
+      tl.Branch(tl.Serial(*layers), tl.NoOp()),
+      tl.AddAll()
+  )
+
+
 def ResidualFeedForward(feature_depth,
                         feedforward_depth,
                         dropout,
                         mode):
   """Residual feed-forward layer with normalization at start."""
-  return tl.Residual(
+  return Residual(
       tl.LayerNorm(),
       tl.Dense(feedforward_depth),
       tl.Relu(),
@@ -202,7 +211,7 @@ def ChunkedDecoderLayer(feature_depth,
     the layer.
   """
   return tl.Serial(
-      tl.Residual(  # Self-attention block.
+      Residual(  # Self-attention block.
           tl.Map(tl.LayerNorm()),
           ChunkedCausalMultiHeadedAttention(
               feature_depth, num_heads=num_heads, dropout=dropout,
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 7002b622f..e0f07dacd 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -59,16 +59,13 @@ def EncoderLayer(feature_depth,
   """
   return tl.Serial(
       tl.Residual(  # Attention block here.
-          tl.Parallel(tl.LayerNorm(), tl.NoOp()),
+          tl.LayerNorm(),
           tl.MultiHeadedAttention(feature_depth, num_heads=num_heads,
                                   dropout=dropout, mode=mode),
-          tl.Parallel(tl.Dropout(rate=dropout, mode=mode), tl.NoOp())
+          tl.Dropout(rate=dropout, mode=mode)
       ),
-      tl.Parallel(
-          ResidualFeedForward(
-              feature_depth, feedforward_depth, dropout, mode=mode),
-          tl.Div(divisor=2.0)  # Mask added to itself in the residual, divide.
-      )
+      ResidualFeedForward(
+          feature_depth, feedforward_depth, dropout, mode=mode),
   )
 
 
@@ -189,7 +186,7 @@ def EncoderDecoderLayer(feature_depth,
                         mode):
   """Transformer encoder-decoder layer.
 
-  The input is a triple pair (encoder, mask, decoder_input) where
+  The input is a triple pair (decoder_input, mask, encoder) where
   the mask is created from the original source to prevent attending
   to the padding part of the encoder.
 
@@ -201,36 +198,31 @@ def EncoderDecoderLayer(feature_depth,
     mode: str: 'train' or 'eval'
 
   Returns:
-    the layer, returning a triple (encoder, mask, decoder_activations).
+    the layer, returning a triple (decoder_activations, mask, encoder).
   """
   # Decoder self-attending to decoder.
   self_attention = tl.Residual(
       tl.LayerNorm(),
-      tl.Branch(tl.NoOp(), tl.CausalMask(axis=-2)),  # create mask
+      tl.Dup(),
+      tl.CausalMask(axis=-2),  # Create the self-attention mask.
+      tl.Swap(),  # Put mask behind the activations.
       tl.MultiHeadedAttention(feature_depth, num_heads=num_heads,
                               dropout=dropout, mode=mode),
-      tl.Select(0),  # drop mask
+      tl.Swap(),  # Put self-attention mask on top.
+      tl.Drop(),   # Drop self-attention mask.
       tl.Dropout(rate=dropout, mode=mode)
   )
   # Decoder attending to encoder.
   encoder_decoder_attention = tl.Serial(
-      tl.Select((2, 0, 0, 1)),  # (dec, enc, enc, mask)
-      tl.MultiHeadedAttentionQKV(  # (q, k, v, mask) --> new, mask
+      tl.Select((0, 2, 2, 1, 2)),  # (dec, enc, enc, mask, enc-copy)
+      tl.MultiHeadedAttentionQKV(  # (q, k, v, mask, ...) --> (new, mask, ...)
           feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),
-      tl.Select(0),  # drop the mask
       tl.Dropout(rate=dropout, mode=mode),
   )
   return tl.Serial(
-      tl.Parallel(tl.NoOp(), tl.NoOp(), self_attention),
-      tl.Branch(tl.NoOp(), encoder_decoder_attention),
-      tl.Select(inputs=(('encoder', 'mask', 'old_act'), 'new_act'),
-                output=('encoder', 'mask', ('old_act', 'new_act'))),
-      tl.Parallel(  # Residual after encoder-decoder attention.
-          tl.NoOp(), tl.NoOp(), tl.Add()),
-      tl.Parallel(  # Feed-forward on the third component (decoder).
-          tl.NoOp(), tl.NoOp(), ResidualFeedForward(
-              feature_depth, feedforward_depth, dropout, mode=mode)
-      )
+      self_attention,
+      tl.Residual(encoder_decoder_attention),
+      ResidualFeedForward(feature_depth, feedforward_depth, dropout, mode=mode)
   )
 
 
@@ -270,7 +262,7 @@ def Transformer(vocab_size,
       tl.Serial(*[EncoderLayer(feature_depth, feedforward_depth, num_heads,
                                dropout, mode)
                   for _ in range(num_layers)]),
-      tl.Parallel(tl.LayerNorm(), tl.NoOp())
+      tl.LayerNorm()
   )
   stack = [EncoderDecoderLayer(feature_depth, feedforward_depth, num_heads,
                                dropout, mode)
@@ -279,11 +271,11 @@ def Transformer(vocab_size,
       tl.Parallel(tl.NoOp(), tl.ShiftRight()),
       tl.Parallel(encoder, embedding),
       tl.Select(inputs=(('encoder', 'mask'), 'decoder'),
-                output=('encoder', ('mask', 'decoder'), 'decoder')),
+                output=('decoder', ('mask', 'decoder'), 'encoder')),
       tl.Parallel(  # (encoder_mask, decoder_input) -> encoder-decoder mask
           tl.NoOp(), tl.EncoderDecoderMask(), tl.NoOp()),
       tl.Serial(*stack),
-      tl.Select(2),  # Drop encoder and mask.
+      tl.Select(0),  # Drop mask and encoder.
       tl.LayerNorm(),
       tl.Dense(vocab_size),
       tl.LogSoftmax()

From 7910b5712acbc2fd631d8acf280761d890ee5af0 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 22 May 2019 13:52:51 -0700
Subject: [PATCH 2051/2720] Corrections to distributed PPO in Atari, use
 simplified discounted reward.

PiperOrigin-RevId: 249513102
---
 tensor2tensor/models/research/rl.py | 22 ++++++++++++++++++++--
 tensor2tensor/rl/ppo.py             | 22 ++++++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 5e2f6ef48..18ebf5f8d 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -532,9 +532,9 @@ def rlmf_dqn_tiny():
 def rlmf_eval():
   """Eval set of hparams for model-free PPO."""
   hparams = rlmf_original()
-  hparams.batch_size = 8
+  hparams.batch_size = 16
   hparams.eval_sampling_temps = [0.0, 0.5, 1.0]
-  hparams.eval_rl_env_max_episode_steps = -1
+  hparams.eval_rl_env_max_episode_steps = 20000
   hparams.add_hparam("ppo_epoch_length", 128)
   hparams.add_hparam("ppo_optimization_batch_size", 32)
   hparams.add_hparam("ppo_epochs_num", 10000)
@@ -544,6 +544,24 @@ def rlmf_eval():
   return hparams
 
 
+@registry.register_hparams
+def rlmf_eval_dist():
+  """Distributional set of hparams for model-free PPO."""
+  hparams = rlmf_eval()
+  hparams.distributional_size = 4096
+  hparams.distributional_subscale = 0.08
+  hparams.base_algo_params = "ppo_dist_params"
+  return hparams
+
+
+@registry.register_hparams
+def rlmf_eval_dist_threshold():
+  """Distributional set of hparams for model-free PPO."""
+  hparams = rlmf_eval_dist()
+  hparams.distributional_threshold = 0.2
+  return hparams
+
+
 class PolicyBase(t2t_model.T2TModel):
 
   def __init__(self, *args, **kwargs):
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index 21841da77..2acada674 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -122,11 +122,19 @@ def define_ppo_epoch(memory, hparams, action_space, batch_size,
     value = _distributional_to_value(
         value_sm, distributional_size, distributional_subscale,
         distributional_threshold)
+  plain_value = value
+  if distributional_threshold > 1:
+    plain_value = _distributional_to_value(
+        value_sm, distributional_size, distributional_subscale, 0.0)
 
   advantage = calculate_generalized_advantage_estimator(
       reward, value, done, hparams.gae_gamma, hparams.gae_lambda)
 
   discounted_reward = tf.stop_gradient(advantage + value[:-1])
+  if distributional_size > 1:
+    end_values = plain_value[-1]
+    discounted_reward = tf.stop_gradient(discounted_rewards(
+        reward, done, hparams.gae_gamma, end_values))
 
   advantage_mean, advantage_variance = tf.nn.moments(advantage, axes=[0, 1],
                                                      keep_dims=True)
@@ -207,3 +215,17 @@ def calculate_generalized_advantage_estimator(
       tf.zeros_like(delta[0, :]),
       parallel_iterations=1), [0])
   return tf.check_numerics(return_, "return")
+
+
+def discounted_rewards(reward, done, gae_gamma, end_values):
+  """Discounted rewards."""
+  not_done = 1 - tf.cast(done, tf.float32)
+  end_values = end_values * not_done[-1, :]
+  return_ = tf.scan(
+      lambda agg, cur: cur + gae_gamma * agg,
+      reward * not_done,
+      initializer=end_values,
+      reverse=True,
+      back_prop=False,
+      parallel_iterations=2)
+  return tf.check_numerics(return_, "return")

From 4068a6246a469b8fd0ca988990f8934122b2a2cd Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 22 May 2019 16:28:18 -0700
Subject: [PATCH 2052/2720] PPO Eval changes - max timestep to eval for and
 other sampling strategies.

- Add a option to cut the eval once it exceeds 10,000 (default value) steps.
- We eval with the following three policies:
  - Categorical sampling.
  - Gumbel sampling with 0.5 temperature.
  - Epsilon-Greedy with 0.2 epsilon.

PiperOrigin-RevId: 249544685
---
 tensor2tensor/trax/rlax/ppo.py      | 38 ++++++++++++++++++++---------
 tensor2tensor/trax/rlax/ppo_main.py | 11 ++++++---
 2 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index c1d71297a..dff7c2bf2 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -983,15 +983,28 @@ def masked_entropy(log_probs, mask):
   return -(np.sum(lp * p) / np.sum(mask))
 
 
-def evaluate_policy(eval_env, get_predictions, boundary, rng=None):
-  trajs = env_problem_utils.play_env_problem_with_policy(
-      eval_env,
-      get_predictions,
-      boundary=boundary,
-      idx=0,  # reset always
-      rng=rng)
-  avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
-  return avg_reward
+def evaluate_policy(eval_env,
+                    get_predictions,
+                    boundary,
+                    max_timestep=10000,
+                    rng=None):
+  """Evaluate the policy."""
+
+  avg_rewards = []
+  for policy in [env_problem_utils.CATEGORICAL_SAMPLING,
+                 env_problem_utils.GUMBEL_SAMPLING,
+                 env_problem_utils.EPSILON_GREEDY]:
+    trajs = env_problem_utils.play_env_problem_with_policy(
+        eval_env,
+        get_predictions,
+        boundary=boundary,
+        max_timestep=max_timestep,
+        idx=0,  # reset always
+        policy_sampling=policy,
+        rng=rng)
+    avg_rewards.append(
+        float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs))
+  return tuple(avg_rewards)
 
 
 def training_loop(
@@ -1011,6 +1024,7 @@ def training_loop(
     target_kl=0.01,
     boundary=20,
     max_timestep=None,
+    max_timestep_eval=20000,
     random_seed=None,
     gamma=GAMMA,
     lambda_=LAMBDA,
@@ -1120,9 +1134,11 @@ def get_predictions(observations, rng=None):
           pickle.dump((policy_net_params, value_net_params), f)
 
       # TODO(afrozm): Dump in jaxboard or somewhere?
-      avg_reward = evaluate_policy(eval_env, get_predictions, boundary, rng=key)
+      avg_reward = evaluate_policy(eval_env, get_predictions, boundary,
+                                   max_timestep=max_timestep_eval, rng=key)
       eval_average_rewards.append(avg_reward)
-      logging.info("Epoch [% 6d] Policy Evaluation = %10.2f", i, avg_reward)
+      logging.info("Epoch [% 6d] Policy Evaluation = (%10.2f, %10.2f, %10.2f)",
+                   i, avg_reward[0], avg_reward[1], avg_reward[2])
       with gfile.GFile(eval_rewards_file, "w") as f:
         f.write(", ".join([str(r) for r in eval_average_rewards]) + "\n")
 
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 1ea80f29d..0519cb1f3 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -50,6 +50,7 @@
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import rendered_env_problem
 from tensor2tensor.rl import gym_utils
+from tensor2tensor.rl.google import atari_utils
 from tensor2tensor.trax import layers
 from tensor2tensor.trax.models import atari_cnn
 from tensor2tensor.trax.rlax import ppo
@@ -79,6 +80,10 @@
     "truncation_timestep", None,
     "If set to an integer, maximum number of time-steps in a "
     "trajectory. Used in the collect procedure.")
+flags.DEFINE_integer(
+    "truncation_timestep_eval", 20000,
+    "If set to an integer, maximum number of time-steps in an evaluation "
+    "trajectory. Used in the collect procedure.")
 
 flags.DEFINE_boolean(
     "jax_debug_nans", False,
@@ -142,7 +147,7 @@
 
 def common_layers():
   # TODO(afrozm): Refactor.
-  if "Pong" in FLAGS.env_problem_name:
+  if "NoFrameskip" in FLAGS.env_problem_name:
     return atari_layers()
 
   cur_layers = []
@@ -200,8 +205,7 @@ def main(argv):
     config.update("jax_platform_name", "tpu")
 
   # TODO(afrozm): Refactor.
-  if "Pong" in FLAGS.env_problem_name and FLAGS.xm:
-    from tensor2tensor.rl.google import atari_utils  # pylint: disable=g-import-not-at-top
+  if "NoFrameskip" in FLAGS.env_problem_name and FLAGS.xm:
     FLAGS.atari_roms_path = "local_ram_fs_tmp"
     atari_utils.copy_roms()
 
@@ -258,6 +262,7 @@ def run_training_loop():
         target_kl=FLAGS.target_kl,
         boundary=FLAGS.boundary,
         max_timestep=FLAGS.truncation_timestep,
+        max_timestep_eval=FLAGS.truncation_timestep_eval,
         random_seed=random_seed,
         c1=FLAGS.value_coef,
         c2=FLAGS.entropy_coef,

From 2f0df9fd7603732e31055495f1572feafb8d19f9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 22 May 2019 16:46:22 -0700
Subject: [PATCH 2053/2720] s/10k/20k in one more place.

PiperOrigin-RevId: 249547842
---
 tensor2tensor/trax/rlax/ppo.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index dff7c2bf2..2fb84f5c1 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -986,7 +986,7 @@ def masked_entropy(log_probs, mask):
 def evaluate_policy(eval_env,
                     get_predictions,
                     boundary,
-                    max_timestep=10000,
+                    max_timestep=20000,
                     rng=None):
   """Evaluate the policy."""
 

From afe714689acdcedf588d7f30df30f8983c6066e6 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 23 May 2019 10:03:17 -0700
Subject: [PATCH 2054/2720] Minor refactoring in PPO. Delete a bunch of code
 that will not get used. Actually implement categorical/gumbel/epsilon-greedy
 in env_problem_utils (forgot to include this file in an earlier change).

PiperOrigin-RevId: 249664412
---
 tensor2tensor/envs/env_problem_utils.py       |  89 ++-
 tensor2tensor/trax/rlax/ppo.py                | 626 +++---------------
 tensor2tensor/trax/rlax/ppo_main.py           |  45 +-
 tensor2tensor/trax/rlax/ppo_test.py           | 201 +-----
 .../trax/rlax/ppo_training_loop_test.py       |  90 +--
 5 files changed, 199 insertions(+), 852 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 0e61911cc..83cd156ae 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -21,6 +21,10 @@
 
 import numpy as np
 
+CATEGORICAL_SAMPLING = "categorical"
+EPSILON_GREEDY = "epsilon-greedy"
+GUMBEL_SAMPLING = "gumbel"
+
 
 def done_indices(dones):
   """Calculates the indices where dones has True."""
@@ -52,7 +56,10 @@ def play_env_problem_with_policy(env,
                                  max_timestep=None,
                                  boundary=20,
                                  idx=0,
-                                 rng=None):
+                                 rng=None,
+                                 policy_sampling=CATEGORICAL_SAMPLING,
+                                 temperature=0.5,
+                                 eps=0.1):
   """Plays the given env with the policy function to collect trajectories.
 
   Args:
@@ -69,23 +76,64 @@ def play_env_problem_with_policy(env,
     idx: int, index on the number of times this function is being called, we may
         want to reset only when idx == 0 for instance.
     rng: jax rng, splittable.
+    policy_sampling: string, how to select an action given a policy, one of:
+        CATEGORICAL_SAMPLING, GREEDY, GUMBEL_SAMPLING
+    temperature: float, temperature used in gumbel sampling.
+    eps: float, epsilon to use in epsilon greedy.
+
 
   Returns:
     Completed trajectories that is a list of triples of (observation, action,
     reward) ndarrays.
   """
 
-  def multinomial_sample(probs):
-    """Sample from this vector of probabilities.
+  def categorical_sample(log_probs):
+    """Categorical sampling."""
+
+    def multinomial_sample(probs):
+      """Sample from this vector of probabilities.
+
+      Args:
+        probs: numpy array of shape (A,) where A is the number of actions, these
+          must sum up to 1.0
 
-    Args:
-      probs: numpy array of shape (A,) where A is the number of actions, these
-        must sum up to 1.0
+      Returns:
+        an integer of which action to pick.
+      """
+
+      return int(np.argwhere(np.random.multinomial(1, probs) == 1))
+
+    # Convert to probs, since we need to do categorical sampling.
+    probs = np.exp(log_probs)
 
-    Returns:
-      an integer of which action to pick.
-    """
-    return int(np.argwhere(np.random.multinomial(1, probs) == 1))
+    # Let's cast up to float64, because that's what numpy does when sampling
+    # and it leads to the sum(pvals[:-1]) > 1.0 error.
+    #
+    # We also re-normalize when we do this.
+    probs = np.float64(probs)
+    probs /= np.sum(probs, axis=1, keepdims=True)
+
+    # Now pick actions from this probs array.
+    return np.apply_along_axis(multinomial_sample, 1, probs)
+
+  def gumbel_sample(log_probs):
+    """Gumbel sampling."""
+    u = np.random.uniform(low=1e-6, high=1.0 - 1e-6, size=log_probs.shape)
+    g = -np.log(-np.log(u))
+    return np.argmax((log_probs / temperature) + g, axis=1)
+
+  def epsilon_greedy(log_probs):
+    """Epsilon greedy sampling."""
+    _, A = log_probs.shape  # pylint: disable=invalid-name
+    actions = []
+    for log_prob in log_probs:
+      # Pick the argmax action.
+      action = np.argmax(log_prob)
+      if np.random.uniform() < eps:
+        # Pick an action at random.
+        action = np.random.choice(range(A))
+      actions.append(action)
+    return np.stack(actions)
 
   # We need to reset all environments, if we're coming here the first time.
   if idx == 0 or max_timestep is None or max_timestep <= 0:
@@ -118,18 +166,15 @@ def multinomial_sample(probs):
     assert (B, A) == log_probs.shape, \
         "B=%d, A=%d, log_probs.shape=%s" % (B, A, log_probs.shape)
 
-    # Convert to probs, since we need to do categorical sampling.
-    probs = np.exp(log_probs)
-
-    # Let's cast up to float64, because that's what numpy does when sampling
-    # and it leads to the sum(pvals[:-1]) > 1.0 error.
-    #
-    # We also re-normalize when we do this.
-    probs = np.float64(probs)
-    probs /= np.sum(probs, axis=1, keepdims=True)
-
-    # Now pick actions from this probs array.
-    actions = np.apply_along_axis(multinomial_sample, 1, probs)
+    actions = None
+    if policy_sampling == CATEGORICAL_SAMPLING:
+      actions = categorical_sample(log_probs)
+    elif policy_sampling == GUMBEL_SAMPLING:
+      actions = gumbel_sample(log_probs)
+    elif policy_sampling == EPSILON_GREEDY:
+      actions = epsilon_greedy(log_probs)
+    else:
+      raise ValueError("Unknown sampling policy [%s]" % policy_sampling)
 
     # Step through the env.
     _, _, dones, _ = env.step(actions)
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 2fb84f5c1..0d4390882 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -64,6 +64,7 @@
 import numpy as onp
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import env_problem_utils
+from tensor2tensor.trax import jaxboard
 from tensor2tensor.trax import layers
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
@@ -75,44 +76,8 @@
 EPSILON = 0.1
 EPOCHS = 50  # 100
 NUM_OPTIMIZER_STEPS = 100
-POLICY_ONLY_NUM_OPTIMIZER_STEPS = 80
-VALUE_ONLY_NUM_OPTIMIZER_STEPS = 80
 PRINT_EVERY_OPTIMIZER_STEP = 20
 BATCH_TRAJECTORIES = 32
-POLICY = "categorical-sampling"
-
-
-def policy_net(rng_key,
-               batch_observations_shape,
-               num_actions,
-               bottom_layers=None):
-  """A policy net function."""
-  # Use the bottom_layers as the bottom part of the network and just add the
-  # required layers on top of it.
-  if bottom_layers is None:
-    bottom_layers = []
-
-  # NOTE: The LogSoftmax instead of the Softmax.
-  bottom_layers.extend([layers.Dense(num_actions), layers.LogSoftmax()])
-  net = layers.Serial(*bottom_layers)
-
-  return net.initialize(batch_observations_shape, rng_key), net
-
-
-def value_net(rng_key,
-              batch_observations_shape,
-              num_actions,
-              bottom_layers=None):
-  """A value net function."""
-  del num_actions
-
-  if bottom_layers is None:
-    bottom_layers = []
-  bottom_layers.extend([
-      layers.Dense(1),
-  ])
-  net = layers.Serial(*bottom_layers)
-  return net.initialize(batch_observations_shape, rng_key), net
 
 
 def policy_and_value_net(rng_key,
@@ -161,20 +126,6 @@ def optimizer_fun(net_params, step_size=1e-3):
   return opt_state, opt_update, get_params
 
 
-def log_params(params, name="params"):
-  """Dumps the params with `logging.error`."""
-  for i, param in enumerate(params):
-    if not param:
-      # Empty tuple.
-      continue
-    if not isinstance(param, (list, tuple)):
-      logging.error("%s[%d] : (%s) = [%s]", name, i, param.shape,
-                    onp.array(param))
-    else:
-      for j, p in enumerate(param):
-        logging.error("\t%s[%d, %d] = [%s]", name, i, j, onp.array(p))
-
-
 # Should this be collect 'n' trajectories, or
 # Run the env for 'n' steps and take completed trajectories, or
 # Any other option?
@@ -182,7 +133,7 @@ def log_params(params, name="params"):
 def collect_trajectories(env,
                          policy_fun,
                          num_trajectories=1,
-                         policy="greedy",
+                         policy=env_problem_utils.CATEGORICAL_SAMPLING,
                          max_timestep=None,
                          boundary=20,
                          epsilon=0.1,
@@ -212,108 +163,18 @@ def collect_trajectories(env,
     reward[i] = (B, T_i)
   """
 
-  if isinstance(env, env_problem.EnvProblem):
-    # This is an env_problem, run its collect function.
-    return env_problem_utils.play_env_problem_with_policy(
-        env,
-        policy_fun,
-        num_trajectories=num_trajectories,
-        max_timestep=max_timestep,
-        boundary=boundary,
-        idx=idx,
-        rng=rng)
-
-  trajectories = []
-
-  for t in range(num_trajectories):
-    t_start = time.time()
-    rewards = []
-    actions = []
-    done = False
-
-    observation = env.reset()
-
-    # This is currently shaped (1, 1) + OBS, but new observations will keep
-    # getting added to it, making it eventually (1, T+1) + OBS
-    observation_history = observation[np.newaxis, np.newaxis, :]
-
-    # Run either till we're done OR if max_timestep is defined only till that
-    # timestep.
-    ts = 0
-    while ((not done) and
-           (not max_timestep or observation_history.shape[1] < max_timestep)):
-      ts_start = time.time()
-      # Run the policy, to pick an action, shape is (1, t, A) because
-      # observation_history is shaped (1, t) + OBS
-      predictions, _, rng = policy_fun(observation_history, rng=rng)
-
-      # We need the predictions for the last time-step, so squeeze the batch
-      # dimension and take the last time-step.
-      predictions = np.squeeze(predictions, axis=0)[-1]
-
-      # Policy can be run in one of the following ways:
-      #  - Greedy
-      #  - Epsilon-Greedy
-      #  - Categorical-Sampling
-      action = None
-      if policy == "greedy":
-        action = np.argmax(predictions)
-      elif policy == "epsilon-greedy":
-        # A schedule for epsilon is 1/k where k is the episode number sampled.
-        if onp.random.random() < epsilon:
-          # Choose an action at random.
-          action = onp.random.randint(0, high=len(predictions))
-        else:
-          # Return the best action.
-          action = np.argmax(predictions)
-      elif policy == "categorical-sampling":
-        # NOTE: The predictions aren't probabilities but log-probabilities
-        # instead, since they were computed with LogSoftmax.
-        # So just np.exp them to make them probabilities.
-        predictions = np.exp(predictions)
-        action = onp.argwhere(onp.random.multinomial(1, predictions) == 1)
-      else:
-        raise ValueError("Unknown policy: %s" % policy)
-
-      # NOTE: Assumption, single batch.
-      try:
-        action = int(action)
-      except TypeError as err:
-        # Let's dump some information before we die off.
-        logging.error("Cannot convert action into an integer: [%s]", err)
-        logging.error("action.shape: [%s]", action.shape)
-        logging.error("action: [%s]", action)
-        logging.error("predictions.shape: [%s]", predictions.shape)
-        logging.error("predictions: [%s]", predictions)
-        logging.error("observation_history: [%s]", observation_history)
-        raise err
-
-      observation, reward, done, _ = env.step(action)
-
-      # observation is of shape OBS, so add extra dims and concatenate on the
-      # time dimension.
-      observation_history = np.concatenate(
-          [observation_history, observation[np.newaxis, np.newaxis, :]], axis=1)
-
-      rewards.append(reward)
-      actions.append(action)
-
-      ts += 1
-      logging.vlog(
-          2, "  Collected time-step[ %5d] of trajectory[ %5d] in [%0.2f] msec.",
-          ts, t, get_time(ts_start))
-    logging.vlog(2, " Collected trajectory[ %5d] in [%0.2f] msec.", t,
-                 get_time(t_start))
-
-    # This means we are done we're been terminated early.
-    assert done or (max_timestep and
-                    max_timestep >= observation_history.shape[1])
-    # observation_history is (1, T+1) + OBS, lets squeeze out the batch dim.
-    observation_history = np.squeeze(observation_history, axis=0)
-    trajectories.append(
-        (observation_history, np.stack(actions), np.stack(rewards)))
-
-  return trajectories
+  assert isinstance(env, env_problem.EnvProblem)
+  # This is an env_problem, run its collect function.
+  return env_problem_utils.play_env_problem_with_policy(
+      env,
+      policy_fun,
+      num_trajectories=num_trajectories,
+      max_timestep=max_timestep,
+      boundary=boundary,
+      policy_sampling=policy,
+      eps=epsilon,
+      idx=idx,
+      rng=rng)
 
 
 # This function can probably be simplified, ask how?
@@ -462,52 +323,6 @@ def rewards_to_go(rewards, mask, gamma=0.99):
   return np.flip(np.stack(r2gs, axis=1), axis=1)
 
 
-@functools.partial(jit, static_argnums=(0,))
-def value_loss(value_net_apply,
-               value_net_params,
-               observations,
-               rewards,
-               reward_mask,
-               gamma=0.99,
-               epsilon=0.2,
-               value_prediction_old=None,
-               rng=None):
-  """Computes the value loss.
-
-  Args:
-    value_net_apply: value net apply function with signature (params, ndarray of
-      shape (B, T+1) + OBS) -> ndarray(B, T+1, 1)
-    value_net_params: params of value_net_apply.
-    observations: np.ndarray of shape (B, T+1) + OBS
-    rewards: np.ndarray of shape (B, T) of rewards.
-    reward_mask: np.ndarray of shape (B, T), the mask over rewards.
-    gamma: float, discount factor.
-    epsilon: float, clip-fraction, used if value_value_prediction_old isn't None
-    value_prediction_old: np.ndarray of shape (B, T+1, 1) of value predictions
-      using the old parameters. If provided, we incorporate this in the loss as
-      well. This is from the OpenAI baselines implementation.
-    rng: jax rng, splittable.
-
-  Returns:
-    The average L2 value loss, averaged over instances where reward_mask is 1.
-  """
-
-  B, T = rewards.shape  # pylint: disable=invalid-name
-  assert (B, T + 1) == observations.shape[:2]
-
-  # NOTE: observations is (B, T+1) + OBS, value_prediction is (B, T+1, 1)
-  value_prediction = value_net_apply(observations, value_net_params, rng=rng)
-  assert (B, T + 1, 1) == value_prediction.shape
-
-  return value_loss_given_predictions(
-      value_prediction,
-      rewards,
-      reward_mask,
-      gamma,
-      epsilon=epsilon,
-      value_prediction_old=value_prediction_old)
-
-
 @jit
 def value_loss_given_predictions(value_prediction,
                                  rewards,
@@ -669,48 +484,6 @@ def clipped_objective(probab_ratios, advantages, reward_mask, epsilon=0.2):
       advantages) * reward_mask
 
 
-@functools.partial(jit, static_argnums=(0,))
-def ppo_loss(policy_net_apply,
-             new_policy_params,
-             log_probab_actions_old,
-             value_predictions_old,
-             padded_observations,
-             padded_actions,
-             padded_rewards,
-             reward_mask,
-             gamma=0.99,
-             lambda_=0.95,
-             epsilon=0.2,
-             rng=None):
-  """PPO objective, with an eventual minus sign, given observations."""
-  B, T = padded_rewards.shape  # pylint: disable=invalid-name
-  assert (B, T + 1) == padded_observations.shape[:2]
-  assert (B, T + 1) == log_probab_actions_old.shape[:2]
-  assert (B, T + 1, 1) == value_predictions_old.shape
-  assert (B, T) == padded_actions.shape
-  assert (B, T) == padded_rewards.shape
-  assert (B, T) == reward_mask.shape
-
-  # Compute predicted log-probs and hand over to `ppo_loss_given_predictions`.
-
-  # log_probab_actions_{old,new} are both (B, T+1, A)
-  log_probab_actions_new = policy_net_apply(
-      padded_observations, new_policy_params, rng=rng)
-  assert (B, T + 1) == log_probab_actions_new.shape[:2]
-  assert log_probab_actions_old.shape[-1] == log_probab_actions_new.shape[-1]
-
-  return ppo_loss_given_predictions(
-      log_probab_actions_new,
-      log_probab_actions_old,
-      value_predictions_old,
-      padded_actions,
-      padded_rewards,
-      reward_mask,
-      gamma=gamma,
-      lambda_=lambda_,
-      epsilon=epsilon)
-
-
 @jit
 def ppo_loss_given_predictions(log_probab_actions_new,
                                log_probab_actions_old,
@@ -833,65 +606,6 @@ def combined_loss(new_params,
       c2=c2)
 
 
-@functools.partial(jit, static_argnums=(2, 3, 4))
-def ppo_opt_step(i,
-                 opt_state,
-                 ppo_opt_update,
-                 ppo_get_params,
-                 policy_net_apply,
-                 log_probab_actions_old,
-                 value_predictions_old,
-                 padded_observations,
-                 padded_actions,
-                 padded_rewards,
-                 reward_mask,
-                 gamma=0.99,
-                 lambda_=0.95,
-                 epsilon=0.1,
-                 rng=None):
-  """PPO optimizer step."""
-  new_policy_params = ppo_get_params(opt_state)
-  g = grad(
-      ppo_loss, argnums=1)(
-          policy_net_apply,
-          new_policy_params,
-          log_probab_actions_old,
-          value_predictions_old,
-          padded_observations,
-          padded_actions,
-          padded_rewards,
-          reward_mask,
-          gamma=gamma,
-          lambda_=lambda_,
-          epsilon=epsilon,
-          rng=rng)
-  return ppo_opt_update(i, g, opt_state)
-
-
-@functools.partial(jit, static_argnums=(2, 3, 4))
-def value_opt_step(i,
-                   opt_state,
-                   opt_update,
-                   get_params,
-                   value_net_apply,
-                   padded_observations,
-                   padded_rewards,
-                   reward_mask,
-                   gamma=0.99,
-                   rng=None):
-  """Value optimizer step."""
-  value_params = get_params(opt_state)
-  # Note this partial application here and argnums above in ppo_opt_step.
-  g = grad(functools.partial(value_loss, value_net_apply))(
-      value_params,
-      padded_observations,
-      padded_rewards,
-      reward_mask,
-      gamma=gamma,
-      rng=rng)
-  return opt_update(i, g, opt_state)
-
-
 @functools.partial(jit, static_argnums=(2, 3, 4))
 def policy_and_value_opt_step(i,
                               opt_state,
@@ -990,7 +704,7 @@ def evaluate_policy(eval_env,
                     rng=None):
   """Evaluate the policy."""
 
-  avg_rewards = []
+  avg_rewards = {}
   for policy in [env_problem_utils.CATEGORICAL_SAMPLING,
                  env_problem_utils.GUMBEL_SAMPLING,
                  env_problem_utils.EPSILON_GREEDY]:
@@ -1002,24 +716,18 @@ def evaluate_policy(eval_env,
         idx=0,  # reset always
         policy_sampling=policy,
         rng=rng)
-    avg_rewards.append(
-        float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs))
-  return tuple(avg_rewards)
+    avg_rewards[policy] = float(
+        sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
+  return avg_rewards
 
 
 def training_loop(
     env=None,
     epochs=EPOCHS,
-    policy_net_fun=None,
-    value_net_fun=None,
     policy_and_value_net_fun=None,
-    policy_optimizer_fun=None,
-    value_optimizer_fun=None,
     policy_and_value_optimizer_fun=None,
     batch_size=BATCH_TRAJECTORIES,
     num_optimizer_steps=NUM_OPTIMIZER_STEPS,
-    policy_only_num_optimizer_steps=POLICY_ONLY_NUM_OPTIMIZER_STEPS,
-    value_only_num_optimizer_steps=VALUE_ONLY_NUM_OPTIMIZER_STEPS,
     print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
     target_kl=0.01,
     boundary=20,
@@ -1041,13 +749,16 @@ def training_loop(
   if output_dir:
     gfile.makedirs(output_dir)
 
+  # Create summary writers and history.
+  train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
+  eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
+
   jax_rng_key = trax.get_random_number_generator_and_set_seed(random_seed)
 
   value_losses = []
   ppo_objective = []
   combined_losses = []
   average_rewards = []
-  eval_average_rewards = []
 
   # Batch Observations Shape = [-1, -1] + OBS, because we will eventually call
   # policy and value networks on shape [B, T] +_OBS
@@ -1058,62 +769,31 @@ def training_loop(
 
   policy_and_value_net_params, policy_and_value_net_apply = None, None
   policy_and_value_opt_state, policy_and_value_opt_update = None, None
-  policy_net_params, policy_net_apply = None, None
-  value_net_params, value_net_apply = None, None
-  if policy_and_value_net_fun is not None:
-    jax_rng_key, key1, key2 = jax_random.split(jax_rng_key, num=3)
-
-    # Initialize the policy and value network.
-    policy_and_value_net_params, policy_and_value_net_apply = (
-        policy_and_value_net_fun(key1, batch_observations_shape, num_actions))
-
-    policy_and_value_net_apply = jit(policy_and_value_net_apply)
-
-    # Initialize the optimizers.
-    policy_and_value_optimizer = (
-        policy_and_value_optimizer_fun(policy_and_value_net_params))
-    (policy_and_value_opt_state, policy_and_value_opt_update,
-     policy_and_value_get_params) = policy_and_value_optimizer
-  else:
-    # Initialize the policy and value functions.
-    assert policy_net_fun and value_net_fun
-    jax_rng_key, key1, key2 = jax_random.split(jax_rng_key, num=3)
 
-    policy_net_params, policy_net_apply = policy_net_fun(
-        key1, batch_observations_shape, num_actions)
-    value_net_params, value_net_apply = value_net_fun(key2,
-                                                      batch_observations_shape,
-                                                      num_actions)
+  jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)
+
+  # Initialize the policy and value network.
+  policy_and_value_net_params, policy_and_value_net_apply = (
+      policy_and_value_net_fun(key1, batch_observations_shape, num_actions))
 
-    policy_net_apply = jit(policy_net_apply)
-    value_net_apply = jit(value_net_apply)
+  policy_and_value_net_apply = jit(policy_and_value_net_apply)
 
-    # Initialize the optimizers.
-    ppo_opt_state, ppo_opt_update, ppo_get_params = (
-        policy_optimizer_fun(policy_net_params))
-    value_opt_state, value_opt_update, value_get_params = (
-        value_optimizer_fun(value_net_params))
+  # Initialize the optimizers.
+  policy_and_value_optimizer = (
+      policy_and_value_optimizer_fun(policy_and_value_net_params))
+  (policy_and_value_opt_state, policy_and_value_opt_update,
+   policy_and_value_get_params) = policy_and_value_optimizer
 
   for i in range(epochs):
 
     # Params we'll use to collect the trajectories.
-    if policy_and_value_net_apply:
-      policy_and_value_net_params = policy_and_value_get_params(
-          policy_and_value_opt_state)
-    else:
-      policy_net_params = ppo_get_params(ppo_opt_state)
-      value_net_params = value_get_params(value_opt_state)
+    policy_and_value_net_params = policy_and_value_get_params(
+        policy_and_value_opt_state)
 
     # A function to get the policy and value predictions.
     def get_predictions(observations, rng=None):
       """Returns log-probs, value predictions and key back."""
-      key, key1, key2 = jax_random.split(rng, num=3)
-
-      if policy_net_apply is not None:
-        return (policy_net_apply(observations, policy_net_params, rng=key1),
-                value_net_apply(observations, value_net_params, rng=key2), key)
-
-      assert policy_and_value_net_apply
+      key, key1 = jax_random.split(rng, num=2)
 
       log_probs, value_preds = policy_and_value_net_apply(
           observations, policy_and_value_net_params, rng=key1)
@@ -1126,31 +806,22 @@ def get_predictions(observations, rng=None):
 
       logging.vlog(1, "Epoch [% 6d] saving model and evaluating policy.", i)
       params_file = os.path.join(output_dir, "model-%06d.pkl" % i)
-      eval_rewards_file = os.path.join(output_dir, "eval_average_rewards")
       with gfile.GFile(params_file, "wb") as f:
-        if policy_and_value_net_params:
-          pickle.dump(policy_and_value_net_params, f)
-        else:
-          pickle.dump((policy_net_params, value_net_params), f)
+        pickle.dump(policy_and_value_net_params, f)
 
-      # TODO(afrozm): Dump in jaxboard or somewhere?
       avg_reward = evaluate_policy(eval_env, get_predictions, boundary,
                                    max_timestep=max_timestep_eval, rng=key)
-      eval_average_rewards.append(avg_reward)
-      logging.info("Epoch [% 6d] Policy Evaluation = (%10.2f, %10.2f, %10.2f)",
-                   i, avg_reward[0], avg_reward[1], avg_reward[2])
-      with gfile.GFile(eval_rewards_file, "w") as f:
-        f.write(", ".join([str(r) for r in eval_average_rewards]) + "\n")
+      for k, v in avg_reward.items():
+        eval_sw.scalar("eval/mean_reward/%s" % k, v, step=i)
+        logging.info("Epoch [% 6d] Policy Evaluation [%s] = %10.2f", i, k, v)
 
     t = time.time()
-    t0 = t
     logging.vlog(1, "Epoch [% 6d] collecting trajectories.", i)
     jax_rng_key, key = jax_random.split(jax_rng_key)
     trajs = collect_trajectories(
         env,
         policy_fun=get_predictions,
         num_trajectories=batch_size,
-        policy=POLICY,
         max_timestep=max_timestep,
         boundary=boundary,
         rng=key,
@@ -1164,6 +835,8 @@ def get_predictions(observations, rng=None):
     min_reward = min(np.sum(traj[2]) for traj in trajs)
     average_rewards.append(avg_reward)
 
+    train_sw.scalar("train/mean_reward", avg_reward, step=i)
+
     logging.vlog(1, "Rewards avg=[%0.2f], max=[%0.2f], min=[%0.2f], all=%s",
                  avg_reward, max_reward, min_reward,
                  [float(np.sum(traj[2])) for traj in trajs])
@@ -1222,68 +895,36 @@ def get_predictions(observations, rng=None):
 
     # Compute value and ppo losses.
     cur_value_loss, cur_ppo_loss, cur_combined_loss = None, None, None
-    jax_rng_key, key1, key2 = jax_random.split(jax_rng_key, num=3)
-    if policy_and_value_net_apply:
-      logging.vlog(2, "Starting to compute P&V loss.")
-      t = time.time()
-      cur_combined_loss, cur_ppo_loss, cur_value_loss, entropy_bonus = (
-          combined_loss(
-              policy_and_value_net_params,
-              log_probabs_traj,
-              value_predictions_traj,
-              policy_and_value_net_apply,
-              padded_observations,
-              padded_actions,
-              padded_rewards,
-              reward_mask,
-              gamma=gamma,
-              lambda_=lambda_,
-              epsilon=epsilon_schedule,
-              c1=c1,
-              c2=c2,
-              rng=key1))
-      logging.vlog(
-          1,
-          "Calculating P&V loss [%10.2f(%10.2f, %10.2f, %10.2f)] took %0.2f msec.",
-          cur_combined_loss, cur_value_loss, cur_ppo_loss, entropy_bonus,
-          get_time(t))
-    else:
-      logging.vlog(2, "Starting to compute Value loss.")
-      t = time.time()
-      cur_value_loss = value_loss(
-          value_net_apply,
-          value_net_params,
-          padded_observations,
-          padded_rewards,
-          reward_mask,
-          gamma=gamma,
-          rng=key1)
-
-      logging.vlog(1, "Calculating value loss took %0.2f msec.", get_time(t))
-
-      t = time.time()
-      logging.vlog(2, "Starting to compute PPO loss.")
-      cur_ppo_loss = ppo_loss(
-          policy_net_apply,
-          policy_net_params,
-          log_probabs_traj,
-          value_predictions_traj,
-          padded_observations,
-          padded_actions,
-          padded_rewards,
-          reward_mask,
-          gamma=gamma,
-          lambda_=lambda_,
-          epsilon=epsilon_schedule,
-          rng=key2)
-      logging.vlog(1, "Calculating PPO loss took %0.2f msec.", get_time(t))
+    jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)
+    logging.vlog(2, "Starting to compute P&V loss.")
+    t = time.time()
+    cur_combined_loss, cur_ppo_loss, cur_value_loss, entropy_bonus = (
+        combined_loss(
+            policy_and_value_net_params,
+            log_probabs_traj,
+            value_predictions_traj,
+            policy_and_value_net_apply,
+            padded_observations,
+            padded_actions,
+            padded_rewards,
+            reward_mask,
+            gamma=gamma,
+            lambda_=lambda_,
+            epsilon=epsilon_schedule,
+            c1=c1,
+            c2=c2,
+            rng=key1))
+    logging.vlog(
+        1,
+        "Calculating P&V loss [%10.2f(%10.2f, %10.2f, %10.2f)] took %0.2f msec.",
+        cur_combined_loss, cur_value_loss, cur_ppo_loss, entropy_bonus,
+        get_time(t))
 
     value_losses.append(cur_value_loss)
     ppo_objective.append(-1.0 * cur_ppo_loss)
-    if cur_combined_loss:
-      combined_losses.append(cur_combined_loss)
+    combined_losses.append(cur_combined_loss)
 
-    jax_rng_key, key1, key2 = jax_random.split(jax_rng_key, num=3)
+    jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)
     if policy_and_value_net_apply:
       logging.vlog(1, "Policy and Value Optimization")
       t1 = time.time()
@@ -1366,132 +1007,15 @@ def get_predictions(observations, rng=None):
 
       logging.info(
           "Epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
-          " Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)], took [%2.5f msec], Average Rewards(last 10):%s",
+          " Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)], took "
+          "[%2.5f msec], Average Rewards(last 10):%s",
           i, min_reward, max_reward,
           avg_reward, loss_combined, loss_value, loss_ppo, entropy_bonus,
           get_time(t1), ", ".join([str(a) for a in average_rewards[-10:]]))
-    else:
-      # Run optimizers.
-      logging.vlog(1, "PPO Optimization")
-      t1 = time.time()
-      keys1 = jax_random.split(key1, num=policy_only_num_optimizer_steps)
-      for j in range(policy_only_num_optimizer_steps):
-        k1, k2, k3 = jax_random.split(keys1[j], num=3)
-        t = time.time()
-        # Update the optimizer state.
-        ppo_opt_state = ppo_opt_step(
-            j,
-            ppo_opt_state,
-            ppo_opt_update,
-            ppo_get_params,
-            policy_net_apply,
-            log_probabs_traj,
-            value_predictions_traj,
-            padded_observations,
-            padded_actions,
-            padded_rewards,
-            reward_mask,
-            gamma=gamma,
-            lambda_=lambda_,
-            epsilon=epsilon_schedule,
-            rng=k1,
-        )
-        t2 = time.time()
-        # Compute the approx KL for early stopping.
-        # Get the new params.
-        new_policy_net_params = ppo_get_params(ppo_opt_state)
-        log_probab_actions_new = policy_net_apply(
-            padded_observations, new_policy_net_params, rng=k2)
-        approx_kl = approximate_kl(log_probab_actions_new, log_probabs_traj,
-                                   reward_mask)
-
-        early_stopping = enable_early_stopping and approx_kl > 1.5 * target_kl
-        if early_stopping:
-          logging.vlog(
-              1, "Early stopping policy optimization at iter: %d, "
-              "with approx_kl: %0.2f", j, approx_kl)
-          # We don't return right-away, we want the below to execute on the last
-          # iteration.
-
-        if (((j + 1) % print_every_optimizer_steps == 0) or
-            (j == num_optimizer_steps - 1) or early_stopping):
-          new_ppo_loss = ppo_loss(
-              policy_net_apply,
-              new_policy_net_params,
-              log_probabs_traj,
-              value_predictions_traj,
-              padded_observations,
-              padded_actions,
-              padded_rewards,
-              reward_mask,
-              gamma=gamma,
-              lambda_=lambda_,
-              epsilon=epsilon_schedule,
-              rng=k3,
-          )
-          logging.vlog(1, "One PPO grad desc took: %0.2f msec", get_time(t, t2))
-          logging.vlog(1, "PPO loss [%10.2f] -> [%10.2f]", cur_ppo_loss,
-                       new_ppo_loss)
-
-        if early_stopping:
-          break
-
-      logging.vlog(1, "Total PPO loss reduction [%0.2f]%%",
-                   (100 * (cur_ppo_loss - new_ppo_loss) / np.abs(cur_ppo_loss)))
 
-      logging.vlog(1, "Value Optimization")
+  logging.vlog(1, "value_losses: %s", np.stack(value_losses))
+  logging.vlog(1, "ppo_objective:\n%s", np.stack(ppo_objective))
+  logging.vlog(1, "combined_losses:\n%s", np.stack(combined_losses))
+  logging.vlog(1, "average_rewards:\n%s", average_rewards)
 
-      keys2 = jax_random.split(key2, num=value_only_num_optimizer_steps)
-      for j in range(value_only_num_optimizer_steps):
-        k1, k2, k3 = jax_random.split(keys2[j], num=3)
-        t = time.time()
-        value_opt_state = value_opt_step(
-            j,
-            value_opt_state,
-            value_opt_update,
-            value_get_params,
-            value_net_apply,
-            padded_observations,
-            padded_rewards,
-            reward_mask,
-            gamma=gamma,
-            rng=k1)
-        t2 = time.time()
-        value_net_params = value_get_params(value_opt_state)
-        if ((j + 1) %
-            print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
-          new_value_loss = value_loss(
-              value_net_apply,
-              value_net_params,
-              padded_observations,
-              padded_rewards,
-              reward_mask,
-              gamma=gamma,
-              rng=k2)
-          logging.vlog(1, "One value grad desc took: %0.2f msec",
-                       get_time(t, t2))
-          logging.vlog(1, "Value loss [%10.2f] -> [%10.2f]", cur_value_loss,
-                       new_value_loss)
-      logging.vlog(1, "Total value loss reduction [%0.2f]%%",
-                   (100 *
-                    (cur_value_loss - new_value_loss) / np.abs(cur_value_loss)))
-
-      logging.vlog(1, "Grad desc took %0.2f msec", get_time(t1))
-
-      logging.info(
-          "Epoch [% 6d], Reward[min, max, avg] [%10.2f,%10.2f,%10.2f], "
-          "ppo loss [%10.2f], value loss [%10.2f], took [%10.2f msec]", i,
-          min_reward, max_reward, avg_reward, new_ppo_loss, new_value_loss,
-          get_time(t0))
-
-  if value_losses:
-    logging.vlog(1, "value_losses: %s", np.stack(value_losses))
-  if ppo_objective:
-    logging.vlog(1, "ppo_objective:\n%s", np.stack(ppo_objective))
-  if combined_losses:
-    logging.vlog(1, "combined_losses:\n%s", np.stack(combined_losses))
-  if average_rewards:
-    logging.vlog(1, "average_rewards:\n%s", average_rewards)
-
-  return ((policy_net_params, value_net_params), average_rewards,
-          np.stack(value_losses), np.stack(ppo_objective))
+  return (average_rewards, np.stack(value_losses), np.stack(ppo_objective))
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 0519cb1f3..09cc4055a 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -96,11 +96,6 @@
 flags.DEFINE_integer("resized_height", 105, "Resized height of the game frame.")
 flags.DEFINE_integer("resized_width", 80, "Resized width of the game frame.")
 
-flags.DEFINE_boolean(
-    "combined_network", False,
-    "If True there is a single network that determines policy"
-    "and values.")
-
 flags.DEFINE_bool(
     "two_towers", True,
     "In the combined network case should we make one tower or"
@@ -111,20 +106,12 @@
 
 # Number of optimizer steps of the combined net, policy net and value net.
 flags.DEFINE_integer("num_optimizer_steps", 100, "Number of optimizer steps.")
-flags.DEFINE_integer("policy_only_num_optimizer_steps", 80,
-                     "Number of optimizer steps policy only.")
-flags.DEFINE_integer("value_only_num_optimizer_steps", 80,
-                     "Number of optimizer steps value only.")
 flags.DEFINE_integer(
     "print_every_optimizer_steps", 1,
     "How often to log during the policy optimization process.")
 
 # Learning rate of the combined net, policy net and value net.
 flags.DEFINE_float("learning_rate", 1e-3, "Learning rate.")
-flags.DEFINE_float("policy_only_learning_rate", 3e-4,
-                   "Learning rate for policy network only.")
-flags.DEFINE_float("value_only_learning_rate", 1e-3,
-                   "Learning rate for value network only.")
 
 # Target KL is used for doing early stopping in the
 flags.DEFINE_float("target_kl", 0.01, "Policy iteration early stopping")
@@ -218,26 +205,12 @@ def main(argv):
 
   def run_training_loop():
     """Runs the training loop."""
-    policy_net_fun = None
-    value_net_fun = None
-    policy_and_value_net_fun = None
-    policy_optimizer_fun = None
-    value_optimizer_fun = None
-    policy_and_value_optimizer_fun = None
-
-    if FLAGS.combined_network:
-      policy_and_value_net_fun = functools.partial(
-          ppo.policy_and_value_net,
-          bottom_layers_fn=common_layers,
-          two_towers=FLAGS.two_towers)
-      policy_and_value_optimizer_fun = get_optimizer_fun(FLAGS.learning_rate)
-    else:
-      policy_net_fun = functools.partial(
-          ppo.policy_net, bottom_layers=common_layers())
-      value_net_fun = functools.partial(
-          ppo.value_net, bottom_layers=common_layers())
-      policy_optimizer_fun = get_optimizer_fun(FLAGS.policy_only_learning_rate)
-      value_optimizer_fun = get_optimizer_fun(FLAGS.value_only_learning_rate)
+
+    policy_and_value_net_fun = functools.partial(
+        ppo.policy_and_value_net,
+        bottom_layers_fn=common_layers,
+        two_towers=FLAGS.two_towers)
+    policy_and_value_optimizer_fun = get_optimizer_fun(FLAGS.learning_rate)
 
     random_seed = None
     try:
@@ -248,15 +221,9 @@ def run_training_loop():
     ppo.training_loop(
         env=env,
         epochs=FLAGS.epochs,
-        policy_net_fun=policy_net_fun,
-        value_net_fun=value_net_fun,
         policy_and_value_net_fun=policy_and_value_net_fun,
-        policy_optimizer_fun=policy_optimizer_fun,
-        value_optimizer_fun=value_optimizer_fun,
         policy_and_value_optimizer_fun=policy_and_value_optimizer_fun,
         num_optimizer_steps=FLAGS.num_optimizer_steps,
-        policy_only_num_optimizer_steps=FLAGS.policy_only_num_optimizer_steps,
-        value_only_num_optimizer_steps=FLAGS.value_only_num_optimizer_steps,
         print_every_optimizer_steps=FLAGS.print_every_optimizer_steps,
         batch_size=FLAGS.batch_size,
         target_kl=FLAGS.target_kl,
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 18bcbd524..4658c7ad2 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -24,7 +24,6 @@
 import numpy as np
 from tensor2tensor.trax import layers
 from tensor2tensor.trax import trax
-from tensor2tensor.trax.rlax import fake_env
 from tensor2tensor.trax.rlax import ppo
 from tensorflow import test
 
@@ -34,49 +33,6 @@ class PpoTest(test.TestCase):
   def setUp(self):
     self.rng_key = trax.get_random_number_generator_and_set_seed(0)
 
-  def test_policy_net(self):
-    observation_shape = (3, 4)
-    num_actions = 2
-    policy_params, policy_apply = ppo.policy_net(
-        self.rng_key,
-        (-1, -1) + observation_shape,
-        num_actions,
-        # flatten except batch and time
-        # step dimensions.
-        [layers.Flatten(num_axis_to_keep=2)])
-
-    # Generate a batch of observations.
-    batch = 2
-    time_steps = 10
-    batch_of_observations = np.random.uniform(
-        size=(batch, time_steps) + observation_shape)
-
-    # Apply the policy net on observations
-    policy_output = policy_apply(batch_of_observations, policy_params)
-
-    # Verify certain expectations on the output.
-    self.assertEqual((batch, time_steps, num_actions), policy_output.shape)
-
-    # Also exp of last axis normalizes to 1, since these are log-probabilities.
-    sum_actions = np.sum(np.exp(policy_output), axis=-1)
-    self.assertAllClose(np.ones_like(sum_actions), sum_actions)
-
-  def test_value_net(self):
-    observation_shape = (3, 4, 5)
-    num_actions = 2
-    value_params, value_apply = ppo.value_net(
-        self.rng_key,
-        (-1, -1) + observation_shape,
-        num_actions, [layers.Flatten(num_axis_to_keep=2)])
-    batch = 2
-    time_steps = 10
-    batch_of_observations = np.random.uniform(
-        size=(batch, time_steps) + observation_shape)
-    value_output = value_apply(batch_of_observations, value_params)
-
-    # NOTE: The extra dimension at the end because of Dense(1).
-    self.assertEqual((batch, time_steps, 1), value_output.shape)
-
   def test_policy_and_value_net(self):
     observation_shape = (3, 4, 5)
     batch_observation_shape = (-1, -1) + observation_shape
@@ -95,118 +51,6 @@ def test_policy_and_value_net(self):
     self.assertEqual((batch, time_steps, num_actions), pnv_output[0].shape)
     self.assertEqual((batch, time_steps, 1), pnv_output[1].shape)
 
-  def test_collect_trajectories(self):
-    self.rng_key, key1, key2, key3, key4 = jax_random.split(self.rng_key, num=5)
-    observation_shape = (2, 3, 4)
-    num_actions = 2
-    policy_params, policy_apply = ppo.policy_net(
-        key1,
-        (-1, -1) + observation_shape,
-        num_actions,
-        # flatten except batch and time
-        # step dimensions.
-        [layers.Flatten(num_axis_to_keep=2)])
-
-    # We'll get done at time-step #5, starting from 0, therefore in 6 steps.
-    done_time_step = 5
-    env = fake_env.FakeEnv(
-        observation_shape, num_actions, done_time_step=done_time_step)
-
-    def policy_fun(obs, rng=None):
-      rng, r = jax_random.split(rng)
-      return policy_apply(obs, policy_params, rng=r), (), rng
-
-    num_trajectories = 5
-    trajectories = ppo.collect_trajectories(
-        env,
-        policy_fun=policy_fun,
-        num_trajectories=num_trajectories,
-        policy="categorical-sampling",
-        rng=key2)
-
-    # Number of trajectories is as expected.
-    self.assertEqual(num_trajectories, len(trajectories))
-
-    # Shapes of observations, actions and rewards are as expected.
-    for observations, actions, rewards in trajectories:
-      # observations are one more in number than rewards or actions.
-      self.assertEqual((done_time_step + 2,) + observation_shape,
-                       observations.shape)
-      self.assertEqual((done_time_step + 1,), actions.shape)
-      self.assertEqual((done_time_step + 1,), rewards.shape)
-
-    # Test collect using a Policy and Value function.
-    pnv_params, pnv_apply = ppo.policy_and_value_net(
-        key3, (-1, -1) + observation_shape, num_actions,
-        lambda: [layers.Flatten(num_axis_to_keep=2)])
-
-    def pnv_fun(obs, rng=None):
-      rng, r = jax_random.split(rng)
-      lp, v = pnv_apply(obs, pnv_params, rng=r)
-      return lp, v, rng
-
-    trajectories = ppo.collect_trajectories(
-        env,
-        policy_fun=pnv_fun,
-        num_trajectories=num_trajectories,
-        policy="categorical-sampling",
-        rng=key4)
-
-    # Number of trajectories is as expected.
-    self.assertEqual(num_trajectories, len(trajectories))
-
-    # Shapes of observations, actions and rewards are as expected.
-    for observations, actions, rewards in trajectories:
-      # observations are one more in number than rewards or actions.
-      self.assertEqual((done_time_step + 2,) + observation_shape,
-                       observations.shape)
-      self.assertEqual((done_time_step + 1,), actions.shape)
-      self.assertEqual((done_time_step + 1,), rewards.shape)
-
-  def test_collect_trajectories_max_timestep(self):
-    self.rng_key, key1, key2 = jax_random.split(self.rng_key, num=3)
-    observation_shape = (2, 3, 4)
-    num_actions = 2
-    pnv_params, pnv_apply = ppo.policy_and_value_net(
-        key1, (-1, -1) + observation_shape, num_actions,
-        lambda: [layers.Flatten(num_axis_to_keep=2)])
-
-    def pnv_fun(obs, rng=None):
-      rng, r = jax_random.split(rng)
-      lp, v = pnv_apply(obs, pnv_params, rng=r)
-      return lp, v, rng
-
-    # We'll get done at time-step #5, starting from 0, therefore in 6 steps.
-    done_time_step = 5
-    env = fake_env.FakeEnv(
-        observation_shape, num_actions, done_time_step=done_time_step)
-
-    num_trajectories = 5
-
-    # Let's collect trajectories only till `max_timestep`.
-    max_timestep = 3
-
-    # we're testing when we early stop the trajectory.
-    assert max_timestep < done_time_step
-
-    trajectories = ppo.collect_trajectories(
-        env,
-        policy_fun=pnv_fun,
-        num_trajectories=num_trajectories,
-        policy="categorical-sampling",
-        max_timestep=max_timestep,
-        rng=key2)
-
-    # Number of trajectories is as expected.
-    self.assertEqual(num_trajectories, len(trajectories))
-
-    # Shapes of observations, actions and rewards are as expected.
-    for observations, actions, rewards in trajectories:
-      # observations are one more in number than rewards or actions.
-      self.assertEqual((max_timestep,) + observation_shape, observations.shape)
-      self.assertEqual((max_timestep - 1,), actions.shape)
-      self.assertEqual((max_timestep - 1,), rewards.shape)
-
   def test_pad_trajectories(self):
     observation_shape = (2, 3, 4)
     trajectories = []
@@ -310,8 +154,6 @@ def test_rewards_to_go_really_long_sequences(self):
     self.assertAllClose(expected_r2g, actual_r2g)
 
   def test_value_loss(self):
-    self.rng_key, key = jax_random.split(self.rng_key, num=2)
-
     rewards = np.array([
         [1, 2, 4, 8, 16, 32, 64, 128],
         [1, 1, 1, 1, 1, 1, 1, 1],
@@ -339,14 +181,14 @@ def value_net_apply(observations, params, rng=None):
       return np.ones((B, T_p_1, 1))
       # pylint: enable=invalid-name
 
+    value_prediction = value_net_apply(random_observations, [])
+
     with jax.disable_jit():
-      value_loss = ppo.value_loss(
-          value_net_apply, [],
-          random_observations,
+      value_loss = ppo.value_loss_given_predictions(
+          value_prediction,
           rewards,
           rewards_mask,
-          gamma=gamma,
-          rng=key)
+          gamma)
 
     self.assertNear(53.3637084961, value_loss, 1e-6)
 
@@ -534,39 +376,6 @@ def test_clipped_objective(self):
         objective,
         ppo.clipped_objective(probab_ratios, advantages, mask, epsilon))
 
-  def test_ppo_loss(self):
-    self.rng_key, key1, key2, key3 = jax_random.split(self.rng_key, num=4)
-
-    B, T, A, OBS = 2, 10, 2, (28, 28, 3)  # pylint: disable=invalid-name
-    batch_observation_shape = (-1, -1) + OBS
-
-    old_policy_params, _ = ppo.policy_net(key1, batch_observation_shape, A,
-                                          [layers.Flatten(num_axis_to_keep=2)])
-
-    new_policy_params, policy_apply = ppo.policy_net(
-        key2,
-        batch_observation_shape, A,
-        [layers.Flatten(num_axis_to_keep=2)])
-
-    value_params, value_apply = ppo.value_net(
-        key3, batch_observation_shape, A,
-        [layers.Flatten(num_axis_to_keep=2)])
-
-    # Generate a batch of observations.
-
-    observations = np.random.uniform(size=(B, T + 1) + OBS)
-    actions = np.random.randint(0, A, size=(B, T))
-    rewards = np.random.uniform(0, 1, size=(B, T))
-    mask = np.ones_like(rewards)
-
-    log_probs_old = policy_apply(observations, old_policy_params)
-    value_predictions_old = value_apply(observations, value_params)
-
-    # Just test that this computes at all.
-    _ = ppo.ppo_loss(policy_apply, new_policy_params, log_probs_old,
-                     value_predictions_old, observations, actions, rewards,
-                     mask)
-
   def test_combined_loss(self):
     self.rng_key, key1, key2 = jax_random.split(self.rng_key, num=3)
 
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index c4f4e7ffd..02c3971e2 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -19,62 +19,64 @@
 from __future__ import division
 from __future__ import print_function
 
+import contextlib
 import functools
-import gym
+import tempfile
+from tensor2tensor.envs import env_problem
 from tensor2tensor.rl import gym_utils
 from tensor2tensor.trax import layers
 from tensor2tensor.trax.rlax import ppo
 from tensorflow import test
+from tensorflow.io import gfile
 
 
 class PpoTrainingLoopTest(test.TestCase):
 
   def get_wrapped_env(self, name="CartPole-v0", max_episode_steps=2):
-    env = gym.make(name)
-    # Usually gym envs are wrapped in TimeLimit wrapper.
-    env = gym_utils.remove_time_limit_wrapper(env)
-    # Limit this to a small number for tests.
-    return gym.wrappers.TimeLimit(env, max_episode_steps=max_episode_steps)
+    wrapper_fn = functools.partial(
+        gym_utils.gym_env_wrapper,
+        **{
+            "rl_env_max_episode_steps": max_episode_steps,
+            "maxskip_env": False,
+            "rendered_env": False,
+            "rendered_env_resize_to": None,  # Do not resize frames
+            "sticky_actions": False,
+            "output_dtype": None,
+        })
 
-  def test_training_loop(self):
-    env = self.get_wrapped_env("CartPole-v0", 2)
-    num_epochs = 2
-    batch_size = 2
-    # Run the training loop.
-    _, rewards, val_losses, ppo_objectives = ppo.training_loop(
-        env=env,
-        epochs=num_epochs,
-        policy_net_fun=functools.partial(
-            ppo.policy_net, bottom_layers=[layers.Dense(1)]),
-        value_net_fun=functools.partial(
-            ppo.value_net, bottom_layers=[layers.Dense(1)]),
-        policy_optimizer_fun=ppo.optimizer_fun,
-        value_optimizer_fun=ppo.optimizer_fun,
-        batch_size=batch_size,
-        num_optimizer_steps=1,
-        random_seed=0)
-    self.assertLen(rewards, num_epochs)
-    self.assertLen(val_losses, num_epochs)
-    self.assertLen(ppo_objectives, num_epochs)
+    return env_problem.EnvProblem(base_env_name=name,
+                                  batch_size=1,
+                                  env_wrapper_fn=wrapper_fn,
+                                  reward_range=(-1, 1))
+
+  @contextlib.contextmanager
+  def tmp_dir(self):
+    tmp = tempfile.mkdtemp(dir=self.get_temp_dir())
+    yield tmp
+    gfile.rmtree(tmp)
 
-  def test_training_loop_policy_and_value_function(self):
-    env = self.get_wrapped_env("CartPole-v0", 2)
-    num_epochs = 2
-    batch_size = 2
-    # Run the training loop.
-    _, rewards, val_losses, ppo_objectives = ppo.training_loop(
-        env=env,
-        epochs=num_epochs,
-        policy_and_value_net_fun=functools.partial(
-            ppo.policy_and_value_net,
-            bottom_layers_fn=lambda: [layers.Dense(1)]),
-        policy_and_value_optimizer_fun=ppo.optimizer_fun,
-        batch_size=batch_size,
-        num_optimizer_steps=1,
-        random_seed=0)
-    self.assertLen(rewards, num_epochs)
-    self.assertLen(val_losses, num_epochs)
-    self.assertLen(ppo_objectives, num_epochs)
+  def test_training_loop(self):
+    with self.tmp_dir() as output_dir:
+      env = self.get_wrapped_env("CartPole-v0", 2)
+      eval_env = self.get_wrapped_env("CartPole-v0", 2)
+      num_epochs = 2
+      batch_size = 2
+      # Run the training loop.
+      rewards, val_losses, ppo_objectives = ppo.training_loop(
+          env=env,
+          eval_env=eval_env,
+          epochs=num_epochs,
+          policy_and_value_net_fun=functools.partial(
+              ppo.policy_and_value_net,
+              bottom_layers_fn=lambda: [layers.Dense(1)]),
+          policy_and_value_optimizer_fun=ppo.optimizer_fun,
+          batch_size=batch_size,
+          num_optimizer_steps=1,
+          output_dir=output_dir,
+          random_seed=0)
+      self.assertLen(rewards, num_epochs)
+      self.assertLen(val_losses, num_epochs)
+      self.assertLen(ppo_objectives, num_epochs)
 
 
 if __name__ == "__main__":

From 7904dd1efcacd004661d9c3b8bf551e84ca51573 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 23 May 2019 12:10:51 -0700
Subject: [PATCH 2055/2720] More cleanup in PPO.

PiperOrigin-RevId: 249690233
---
 tensor2tensor/trax/rlax/ppo.py                | 27 ++-----------------
 tensor2tensor/trax/rlax/ppo_main.py           |  9 +------
 .../trax/rlax/ppo_training_loop_test.py       |  5 +---
 3 files changed, 4 insertions(+), 37 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 0d4390882..35331901e 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -755,11 +755,6 @@ def training_loop(
 
   jax_rng_key = trax.get_random_number_generator_and_set_seed(random_seed)
 
-  value_losses = []
-  ppo_objective = []
-  combined_losses = []
-  average_rewards = []
-
   # Batch Observations Shape = [-1, -1] + OBS, because we will eventually call
   # policy and value networks on shape [B, T] +_OBS
   batch_observations_shape = (-1, -1) + env.observation_space.shape
@@ -833,19 +828,12 @@ def get_predictions(observations, rng=None):
     avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
     max_reward = max(np.sum(traj[2]) for traj in trajs)
     min_reward = min(np.sum(traj[2]) for traj in trajs)
-    average_rewards.append(avg_reward)
 
     train_sw.scalar("train/mean_reward", avg_reward, step=i)
 
     logging.vlog(1, "Rewards avg=[%0.2f], max=[%0.2f], min=[%0.2f], all=%s",
                  avg_reward, max_reward, min_reward,
                  [float(np.sum(traj[2])) for traj in trajs])
-    logging.vlog(1, "Average Rewards:\n%s", average_rewards)
-
-    # TODO(afrozm): Dump in jaxboard or somewhere?
-    if output_dir:
-      with gfile.GFile(os.path.join(output_dir, "average_rewards"), "w") as f:
-        f.write(", ".join([str(r) for r in average_rewards]) + "\n")
 
     logging.vlog(1,
                  "Trajectory Length average=[%0.2f], max=[%0.2f], min=[%0.2f]",
@@ -920,10 +908,6 @@ def get_predictions(observations, rng=None):
         cur_combined_loss, cur_value_loss, cur_ppo_loss, entropy_bonus,
         get_time(t))
 
-    value_losses.append(cur_value_loss)
-    ppo_objective.append(-1.0 * cur_ppo_loss)
-    combined_losses.append(cur_combined_loss)
-
     jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)
     if policy_and_value_net_apply:
       logging.vlog(1, "Policy and Value Optimization")
@@ -1008,14 +992,7 @@ def get_predictions(observations, rng=None):
       logging.info(
           "Epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
           " Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)], took "
-          "[%2.5f msec], Average Rewards(last 10):%s",
+          "[%2.5f msec].",
           i, min_reward, max_reward,
           avg_reward, loss_combined, loss_value, loss_ppo, entropy_bonus,
-          get_time(t1), ", ".join([str(a) for a in average_rewards[-10:]]))
-
-  logging.vlog(1, "value_losses: %s", np.stack(value_losses))
-  logging.vlog(1, "ppo_objective:\n%s", np.stack(ppo_objective))
-  logging.vlog(1, "combined_losses:\n%s", np.stack(combined_losses))
-  logging.vlog(1, "average_rewards:\n%s", average_rewards)
-
-  return (average_rewards, np.stack(value_losses), np.stack(ppo_objective))
+          get_time(t1))
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 09cc4055a..0c17139f5 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -101,9 +101,6 @@
     "In the combined network case should we make one tower or"
     "two.")
 
-flags.DEFINE_boolean("flatten_dims", False,
-                     "If true, we flatten except the first two dimensions.")
-
 # Number of optimizer steps of the combined net, policy net and value net.
 flags.DEFINE_integer("num_optimizer_steps", 100, "Number of optimizer steps.")
 flags.DEFINE_integer(
@@ -137,11 +134,7 @@ def common_layers():
   if "NoFrameskip" in FLAGS.env_problem_name:
     return atari_layers()
 
-  cur_layers = []
-  if FLAGS.flatten_dims:
-    cur_layers = [layers.Div(divisor=255.0), layers.Flatten(num_axis_to_keep=2)]
-  body = [layers.Dense(64), layers.Tanh(), layers.Dense(64), layers.Tanh()]
-  return cur_layers + body
+  return [layers.Dense(64), layers.Tanh(), layers.Dense(64), layers.Tanh()]
 
 
 def atari_layers():
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index 02c3971e2..de34888a8 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -62,7 +62,7 @@ def test_training_loop(self):
       num_epochs = 2
       batch_size = 2
       # Run the training loop.
-      rewards, val_losses, ppo_objectives = ppo.training_loop(
+      ppo.training_loop(
           env=env,
           eval_env=eval_env,
           epochs=num_epochs,
@@ -74,9 +74,6 @@ def test_training_loop(self):
           num_optimizer_steps=1,
           output_dir=output_dir,
           random_seed=0)
-      self.assertLen(rewards, num_epochs)
-      self.assertLen(val_losses, num_epochs)
-      self.assertLen(ppo_objectives, num_epochs)
 
 
 if __name__ == "__main__":

From 1f9c68d830320a3884f6ef170741eaead3e71354 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 23 May 2019 14:55:41 -0700
Subject: [PATCH 2056/2720] Internal

PiperOrigin-RevId: 249721109
---
 tensor2tensor/models/transformer.py | 33 ++++++++++++++++++++---------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 5ce060686..d98c82cce 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -268,7 +268,6 @@ def body(self, features):
           recurrent_memory_by_layer=self.recurrent_memory_by_layer,
           chunk_number=chunk_number_each_example,
           )
-
     decoder_output = self.decode(
         decoder_input,
         encoder_output,
@@ -279,7 +278,6 @@ def body(self, features):
         losses=losses,
         **decode_kwargs
         )
-
     expected_attentions = features.get("expected_attentions")
     if expected_attentions is not None:
       attention_loss = common_attention.encoder_decoder_attention_loss(
@@ -608,6 +606,17 @@ def forced_logits():
         ret["outputs"] = ret["outputs"][:, :, partial_targets_length:]
     return ret
 
+  def get_decode_start_id(self):
+    """Returns the id of the first decoder input symbol.
+
+    The default case maps None to a vector of 0's for transformer. This method
+    can be overridden to return a different id by a model wanting to use a
+    different decoder start symbol. The id returned by this method is used to
+    index the embedding matrix, and retrieve the vector that will be used as the
+    first input to the decoder
+    """
+    return None
+
   def _fast_decode(self,
                    features,
                    decode_length,
@@ -751,8 +760,9 @@ def preprocess_targets(targets, i):
       # Shifts the targets along by one for the input which pads with zeros.
       # If the modality already maps GO to the zero embeddings this is not
       # needed.
-      targets = tf.cond(
-          tf.equal(i, 0), lambda: tf.zeros_like(targets), lambda: targets)
+      if not self.get_decode_start_id():
+        targets = tf.cond(
+            tf.equal(i, 0), lambda: tf.zeros_like(targets), lambda: targets)
 
       if positional_encoding is not None:
         targets += positional_encoding[:, i:i + 1]
@@ -771,7 +781,6 @@ def symbols_to_logits_fn(ids, i, cache):
       targets = preprocess_targets(targets, i)
 
       bias = decoder_self_attention_bias[:, :, i:i + 1, :i + 1]
-
       with tf.variable_scope("body"):
         body_outputs = dp(
             self.decode,
@@ -808,6 +817,8 @@ def forced_logits():
             tf.less(i, partial_targets_length), forced_logits, lambda: ret)
       return ret, cache
 
+    sos_id = self.get_decode_start_id() or 0
+
     ret = fast_decode(
         encoder_output=encoder_output,
         encoder_decoder_attention_bias=encoder_decoder_attention_bias,
@@ -820,7 +831,8 @@ def forced_logits():
         top_beams=top_beams,
         alpha=alpha,
         batch_size=batch_size,
-        force_decode_length=self._decode_hparams.force_decode_length)
+        force_decode_length=self._decode_hparams.force_decode_length,
+        sos_id=sos_id)
     if partial_targets is not None:
       if beam_size <= 1 or top_beams <= 1:
         ret["outputs"] = ret["outputs"][:, partial_targets_length:]
@@ -1278,7 +1290,7 @@ def features_to_nonpadding(features, inputs_or_targets="inputs"):
   return None
 
 
-def transformer_prepare_decoder(targets, hparams, features=None):
+def transformer_prepare_decoder(targets, hparams, features=None, pad=None):
   """Prepare one shard of the model for the decoder.
 
   Args:
@@ -1286,6 +1298,7 @@ def transformer_prepare_decoder(targets, hparams, features=None):
     hparams: run hyperparameters
     features: optionally pass the entire features dictionary as well. This is
       needed now for "packed" datasets.
+    pad: vector to use for padding when shifting targets right
 
   Returns:
     decoder_input: a Tensor, bottom of decoder stack
@@ -1318,7 +1331,7 @@ def transformer_prepare_decoder(targets, hparams, features=None):
   if hparams.proximity_bias:
     decoder_self_attention_bias += common_attention.attention_bias_proximal(
         common_layers.shape_list(targets)[1])
-  decoder_input = common_layers.shift_right_3d(targets)
+  decoder_input = common_layers.shift_right_3d(targets, pad)
   if hparams.pos == "timing":
     if targets_position is not None:
       decoder_input = common_attention.add_timing_signal_1d_given_position(
@@ -1552,8 +1565,8 @@ def transformer_decoder(decoder_input,
           losses=losses,
           layer_collection=layer_collection,
           recurrent_memory_by_layer=recurrent_memory_by_layer,
-          chunk_number=chunk_number,
-      )
+          chunk_number=chunk_number
+          )
 
     # if normalization is done in layer_preprocess, then it should also be done
     # on the output, since the output can grow very large, being the sum of

From d13662459a67777f64ab619afc3223cd1f36f631 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 23 May 2019 16:28:16 -0700
Subject: [PATCH 2057/2720] Implement restore from policy and `smart` save.

`Smart Save`: Save a policy when atleast a fraction (currently set to 0.5)
of the iterations are marked as `done` since the last time the policy was
saved.

PiperOrigin-RevId: 249737667
---
 tensor2tensor/envs/env_problem_utils.py      | 20 +++--
 tensor2tensor/envs/env_problem_utils_test.py |  2 +-
 tensor2tensor/trax/rlax/ppo.py               | 86 +++++++++++++++-----
 tensor2tensor/trax/rlax/ppo_main.py          | 14 +++-
 4 files changed, 93 insertions(+), 29 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 83cd156ae..8f2ffed47 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -55,7 +55,7 @@ def play_env_problem_with_policy(env,
                                  num_trajectories=1,
                                  max_timestep=None,
                                  boundary=20,
-                                 idx=0,
+                                 reset=True,
                                  rng=None,
                                  policy_sampling=CATEGORICAL_SAMPLING,
                                  temperature=0.5,
@@ -73,8 +73,8 @@ def play_env_problem_with_policy(env,
     boundary: this is the bucket length, we pad the observations to integer
         multiples of this + 1 and then feed the padded observations to the
         policy_fun.
-    idx: int, index on the number of times this function is being called, we may
-        want to reset only when idx == 0 for instance.
+    reset: bool, true if we want to reset the envs. The envs are also reset if
+        max_max_timestep is None or < 0
     rng: jax rng, splittable.
     policy_sampling: string, how to select an action given a policy, one of:
         CATEGORICAL_SAMPLING, GREEDY, GUMBEL_SAMPLING
@@ -83,8 +83,8 @@ def play_env_problem_with_policy(env,
 
 
   Returns:
-    Completed trajectories that is a list of triples of (observation, action,
-    reward) ndarrays.
+    A tuple, (trajectories, number of completed trajectories). Where
+    trajectories is a list of triples of (observation, action, reward) ndarrays.
   """
 
   def categorical_sample(log_probs):
@@ -136,12 +136,14 @@ def epsilon_greedy(log_probs):
     return np.stack(actions)
 
   # We need to reset all environments, if we're coming here the first time.
-  if idx == 0 or max_timestep is None or max_timestep <= 0:
+  if reset or max_timestep is None or max_timestep <= 0:
     env.reset()
   else:
     # Clear completed trajectories held internally.
     env.trajectories.clear_completed_trajectories()
 
+  num_done_trajectories = 0
+
   while env.trajectories.num_completed_trajectories < num_trajectories:
     # Get all the observations for all the active trajectories.
     # Shape is (B, T) + OBS
@@ -179,6 +181,10 @@ def epsilon_greedy(log_probs):
     # Step through the env.
     _, _, dones, _ = env.step(actions)
 
+    # Count the number of done trajectories, the others could just have been
+    # truncated.
+    num_done_trajectories += np.sum(dones)
+
     # Get the indices where we are done ...
     done_idxs = done_indices(dones)
 
@@ -216,4 +222,4 @@ def epsilon_greedy(log_probs):
   # Keep the rest of the trajectories, if any, in our kitty.
   env.trajectories.clear_completed_trajectories(num=num_trajectories)
 
-  return completed_trajectories
+  return completed_trajectories, num_done_trajectories
diff --git a/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index 50dd8683b..4885be940 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -62,7 +62,7 @@ def policy_fun(observations, rng=None):
 
     max_timestep = 15
     num_trajectories = 2
-    trajectories = env_problem_utils.play_env_problem_with_policy(
+    trajectories, _ = env_problem_utils.play_env_problem_with_policy(
         env, policy_fun, num_trajectories=num_trajectories,
         max_timestep=max_timestep, boundary=20)
 
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 35331901e..53498da92 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -137,7 +137,7 @@ def collect_trajectories(env,
                          max_timestep=None,
                          boundary=20,
                          epsilon=0.1,
-                         idx=0,
+                         reset=True,
                          rng=None):
   """Collect trajectories with the given policy net and behaviour.
 
@@ -152,10 +152,12 @@ def collect_trajectories(env,
       done.
     boundary: int, boundary for padding, used in EnvProblem envs.
     epsilon: float, the epsilon for `epsilon-greedy` policy.
-    idx: int, index on the number of times this function is being called
+    reset: bool, true if we want to reset the envs. The envs are also reset if
+      max_max_timestep is None or < 0
     rng: jax rng, splittable.
 
   Returns:
+    A tuple (trajectory, number of trajectories that are done)
     trajectory: list of (observation, action, reward) tuples, where each element
     `i` is a tuple of numpy arrays with shapes as follows:
     observation[i] = (B, T_i + 1)
@@ -173,7 +175,7 @@ def collect_trajectories(env,
       boundary=boundary,
       policy_sampling=policy,
       eps=epsilon,
-      idx=idx,
+      reset=reset,
       rng=rng)
 
 
@@ -708,12 +710,12 @@ def evaluate_policy(eval_env,
   for policy in [env_problem_utils.CATEGORICAL_SAMPLING,
                  env_problem_utils.GUMBEL_SAMPLING,
                  env_problem_utils.EPSILON_GREEDY]:
-    trajs = env_problem_utils.play_env_problem_with_policy(
+    trajs, _ = env_problem_utils.play_env_problem_with_policy(
         eval_env,
         get_predictions,
         boundary=boundary,
         max_timestep=max_timestep,
-        idx=0,  # reset always
+        reset=True,
         policy_sampling=policy,
         rng=rng)
     avg_rewards[policy] = float(
@@ -721,6 +723,29 @@ def evaluate_policy(eval_env,
   return avg_rewards
 
 
+def maybe_restore_params(output_dir, policy_and_value_net_params):
+  """Maybe restore the params from the checkpoint dir.
+
+  Args:
+    output_dir: Directory where saved model checkpoints are stored.
+    policy_and_value_net_params: Default params, returned if model is'nt found.
+
+  Returns:
+    triple (restore (bool), params, iter(int)) where iter is the epoch from
+    which we restored the params, 0 is restore = False.
+  """
+  model_files = gfile.glob(os.path.join(output_dir, "model-??????.pkl"))
+  if not model_files:
+    return False, policy_and_value_net_params, 0
+
+  model_file = sorted(model_files)[-1]
+  model_file_basename = os.path.basename(model_file)  # model-??????.pkl
+  i = int(filter(str.isdigit, model_file_basename))
+  with gfile.GFile(model_file, "rb") as f:
+    policy_and_value_net_params = pickle.load(f)
+  return True, policy_and_value_net_params, i
+
+
 def training_loop(
     env=None,
     epochs=EPOCHS,
@@ -742,12 +767,13 @@ def training_loop(
     output_dir=None,
     eval_every_n=1000,
     eval_env=None,
+    done_frac_for_policy_save=0.5,
     enable_early_stopping=True):
   """Runs the training loop for PPO, with fixed policy and value nets."""
   assert env
+  assert output_dir
 
-  if output_dir:
-    gfile.makedirs(output_dir)
+  gfile.makedirs(output_dir)
 
   # Create summary writers and history.
   train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
@@ -762,15 +788,22 @@ def training_loop(
   assert isinstance(env.action_space, gym.spaces.Discrete)
   num_actions = env.action_space.n
 
-  policy_and_value_net_params, policy_and_value_net_apply = None, None
-  policy_and_value_opt_state, policy_and_value_opt_update = None, None
-
   jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)
 
   # Initialize the policy and value network.
   policy_and_value_net_params, policy_and_value_net_apply = (
       policy_and_value_net_fun(key1, batch_observations_shape, num_actions))
 
+  # Maybe restore the policy params. If there is nothing to restore, then
+  # iteration = 0 and policy_and_value_net_params are returned as is.
+  restore, policy_and_value_net_params, iteration = (
+      maybe_restore_params(output_dir, policy_and_value_net_params))
+
+  if restore:
+    logging.info("Restored parameters from iteration [%d]", iteration)
+    # We should start from the next iteration.
+    iteration += 1
+
   policy_and_value_net_apply = jit(policy_and_value_net_apply)
 
   # Initialize the optimizers.
@@ -779,7 +812,9 @@ def training_loop(
   (policy_and_value_opt_state, policy_and_value_opt_update,
    policy_and_value_get_params) = policy_and_value_optimizer
 
-  for i in range(epochs):
+  num_trajectories_done = 0
+
+  for i in range(iteration, epochs):
 
     # Params we'll use to collect the trajectories.
     policy_and_value_net_params = policy_and_value_get_params(
@@ -795,14 +830,11 @@ def get_predictions(observations, rng=None):
 
       return log_probs, value_preds, key
 
-    # Save params and evaluate the policy.
-    if output_dir and (i % eval_every_n == 0):
+    # Evaluate the policy.
+    if (i % eval_every_n == 0) or (i == epochs - 1):
       jax_rng_key, key = jax_random.split(jax_rng_key, num=2)
 
-      logging.vlog(1, "Epoch [% 6d] saving model and evaluating policy.", i)
-      params_file = os.path.join(output_dir, "model-%06d.pkl" % i)
-      with gfile.GFile(params_file, "wb") as f:
-        pickle.dump(policy_and_value_net_params, f)
+      logging.vlog(1, "Epoch [% 6d] evaluating policy.", i)
 
       avg_reward = evaluate_policy(eval_env, get_predictions, boundary,
                                    max_timestep=max_timestep_eval, rng=key)
@@ -813,16 +845,30 @@ def get_predictions(observations, rng=None):
     t = time.time()
     logging.vlog(1, "Epoch [% 6d] collecting trajectories.", i)
     jax_rng_key, key = jax_random.split(jax_rng_key)
-    trajs = collect_trajectories(
+    trajs, num_done = collect_trajectories(
         env,
         policy_fun=get_predictions,
         num_trajectories=batch_size,
         max_timestep=max_timestep,
         boundary=boundary,
         rng=key,
-        idx=i,
+        reset=(i == 0) or restore,
         epsilon=(10.0 / (i + 10.0)))  # this is a different epsilon.
 
+    # Save parameters every time we see the end of atleast a fraction of batch
+    # number of trajectories that are done (not completed -- completed includes
+    # truncated and done).
+    # Or if this is the last iteration.
+    num_trajectories_done += num_done
+    if ((num_trajectories_done >= done_frac_for_policy_save * batch_size) or
+        (i == epochs - 1)):
+      logging.vlog(1, "Epoch [% 6d] saving model.", i)
+      params_file = os.path.join(output_dir, "model-%06d.pkl" % i)
+      with gfile.GFile(params_file, "wb") as f:
+        pickle.dump(policy_and_value_net_params, f)
+      # Reset this number.
+      num_trajectories_done = 0
+
     logging.vlog(1, "Collecting trajectories took %0.2f msec.", get_time(t))
 
     avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
@@ -996,3 +1042,5 @@ def get_predictions(observations, rng=None):
           i, min_reward, max_reward,
           avg_reward, loss_combined, loss_value, loss_ppo, entropy_bonus,
           get_time(t1))
+
+      restore = False
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 0c17139f5..1f222d60a 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -69,7 +69,7 @@
 # -1: returns env as is.
 # None: unwraps and returns without TimeLimit wrapper.
 # Any other number: imposes this restriction.
-flags.DEFINE_integer(
+flags.DEFINE_string(
     "max_timestep", None,
     "If set to an integer, maximum number of time-steps in a "
     "trajectory. The bare env is wrapped with TimeLimit wrapper.")
@@ -127,6 +127,9 @@
 flags.DEFINE_bool("xm", False, "Are we running on borg?.")
 flags.DEFINE_integer("eval_every_n", 100, "How frequently to eval the policy.")
 flags.DEFINE_integer("eval_batch_size", 4, "Batch size for evaluation.")
+flags.DEFINE_float("done_frac_for_policy_save", 0.5,
+                   "Fraction of the trajectories that should be done to "
+                   "checkpoint the policy.")
 
 
 def common_layers():
@@ -155,9 +158,15 @@ def make_env(batch_size=8):
         batch_size=batch_size,
         reward_range=(-1, 1))
 
+  max_timestep = None
+  try:
+    max_timestep = int(FLAGS.max_timestep)
+  except Exception:  # pylint: disable=broad-except
+    pass
+
   wrapper_fn = functools.partial(
       gym_utils.gym_env_wrapper, **{
-          "rl_env_max_episode_steps": FLAGS.max_timestep,
+          "rl_env_max_episode_steps": max_timestep,
           "maxskip_env": True,
           "rendered_env": True,
           "rendered_env_resize_to": (FLAGS.resized_height, FLAGS.resized_width),
@@ -232,6 +241,7 @@ def run_training_loop():
         enable_early_stopping=FLAGS.enable_early_stopping,
         output_dir=FLAGS.output_dir,
         eval_every_n=FLAGS.eval_every_n,
+        done_frac_for_policy_save=FLAGS.done_frac_for_policy_save,
         eval_env=eval_env)
 
   if FLAGS.jax_debug_nans or FLAGS.disable_jit:

From aca30d0213565d60bc8442ef8ab50f2a26e54981 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 24 May 2019 16:12:06 -0700
Subject: [PATCH 2058/2720] Add timing information to TensorBoard. Move policy
 saving towards end of the epoch.

PiperOrigin-RevId: 249916210
---
 tensor2tensor/trax/rlax/ppo.py                | 294 +++++++++++-------
 tensor2tensor/trax/rlax/ppo_main.py           |  24 +-
 .../trax/rlax/ppo_training_loop_test.py       |   1 +
 3 files changed, 189 insertions(+), 130 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 53498da92..bce25f2b0 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -707,9 +707,10 @@ def evaluate_policy(eval_env,
   """Evaluate the policy."""
 
   avg_rewards = {}
-  for policy in [env_problem_utils.CATEGORICAL_SAMPLING,
-                 env_problem_utils.GUMBEL_SAMPLING,
-                 env_problem_utils.EPSILON_GREEDY]:
+  for policy in [
+      env_problem_utils.CATEGORICAL_SAMPLING, env_problem_utils.GUMBEL_SAMPLING,
+      env_problem_utils.EPSILON_GREEDY
+  ]:
     trajs, _ = env_problem_utils.play_env_problem_with_policy(
         eval_env,
         get_predictions,
@@ -718,8 +719,8 @@ def evaluate_policy(eval_env,
         reset=True,
         policy_sampling=policy,
         rng=rng)
-    avg_rewards[policy] = float(
-        sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
+    avg_rewards[policy] = float(sum(
+        np.sum(traj[2]) for traj in trajs)) / len(trajs)
   return avg_rewards
 
 
@@ -768,17 +769,25 @@ def training_loop(
     eval_every_n=1000,
     eval_env=None,
     done_frac_for_policy_save=0.5,
-    enable_early_stopping=True):
+    enable_early_stopping=True,
+    env_name=None,
+):
   """Runs the training loop for PPO, with fixed policy and value nets."""
   assert env
   assert output_dir
+  assert env_name
 
   gfile.makedirs(output_dir)
 
   # Create summary writers and history.
   train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
+  timing_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "timing"))
   eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
 
+  train_sw.text("env_name", env_name)
+  timing_sw.text("env_name", env_name)
+  eval_sw.text("env_name", env_name)
+
   jax_rng_key = trax.get_random_number_generator_and_set_seed(random_seed)
 
   # Batch Observations Shape = [-1, -1] + OBS, because we will eventually call
@@ -813,8 +822,11 @@ def training_loop(
    policy_and_value_get_params) = policy_and_value_optimizer
 
   num_trajectories_done = 0
+  last_saved_at = 0
 
+  logging.info("Starting the PPO training loop.")
   for i in range(iteration, epochs):
+    epoch_start_time = time.time()
 
     # Params we'll use to collect the trajectories.
     policy_and_value_net_params = policy_and_value_get_params(
@@ -831,18 +843,24 @@ def get_predictions(observations, rng=None):
       return log_probs, value_preds, key
 
     # Evaluate the policy.
-    if (i % eval_every_n == 0) or (i == epochs - 1):
+    policy_eval_start_time = time.time()
+    if ((i + 1) % eval_every_n == 0) or (i == epochs - 1):
       jax_rng_key, key = jax_random.split(jax_rng_key, num=2)
 
       logging.vlog(1, "Epoch [% 6d] evaluating policy.", i)
 
-      avg_reward = evaluate_policy(eval_env, get_predictions, boundary,
-                                   max_timestep=max_timestep_eval, rng=key)
+      avg_reward = evaluate_policy(
+          eval_env,
+          get_predictions,
+          boundary,
+          max_timestep=max_timestep_eval,
+          rng=key)
       for k, v in avg_reward.items():
         eval_sw.scalar("eval/mean_reward/%s" % k, v, step=i)
         logging.info("Epoch [% 6d] Policy Evaluation [%s] = %10.2f", i, k, v)
+    policy_eval_time = get_time(policy_eval_start_time)
 
-    t = time.time()
+    trajectory_collection_start_time = time.time()
     logging.vlog(1, "Epoch [% 6d] collecting trajectories.", i)
     jax_rng_key, key = jax_random.split(jax_rng_key)
     trajs, num_done = collect_trajectories(
@@ -854,22 +872,10 @@ def get_predictions(observations, rng=None):
         rng=key,
         reset=(i == 0) or restore,
         epsilon=(10.0 / (i + 10.0)))  # this is a different epsilon.
+    trajectory_collection_time = get_time(trajectory_collection_start_time)
 
-    # Save parameters every time we see the end of atleast a fraction of batch
-    # number of trajectories that are done (not completed -- completed includes
-    # truncated and done).
-    # Or if this is the last iteration.
-    num_trajectories_done += num_done
-    if ((num_trajectories_done >= done_frac_for_policy_save * batch_size) or
-        (i == epochs - 1)):
-      logging.vlog(1, "Epoch [% 6d] saving model.", i)
-      params_file = os.path.join(output_dir, "model-%06d.pkl" % i)
-      with gfile.GFile(params_file, "wb") as f:
-        pickle.dump(policy_and_value_net_params, f)
-      # Reset this number.
-      num_trajectories_done = 0
-
-    logging.vlog(1, "Collecting trajectories took %0.2f msec.", get_time(t))
+    logging.vlog(1, "Collecting trajectories took %0.2f msec.",
+                 trajectory_collection_time)
 
     avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
     max_reward = max(np.sum(traj[2]) for traj in trajs)
@@ -888,12 +894,14 @@ def get_predictions(observations, rng=None):
                  min(len(traj[0]) for traj in trajs))
     logging.vlog(2, "Trajectory Lengths: %s", [len(traj[0]) for traj in trajs])
 
-    t = time.time()
+    padding_start_time = time.time()
     (_, reward_mask, padded_observations, padded_actions,
      padded_rewards) = pad_trajectories(
          trajs, boundary=boundary)
+    padding_time = get_time(padding_start_time)
 
-    logging.vlog(1, "Padding trajectories took %0.2f msec.", get_time(t))
+    logging.vlog(1, "Padding trajectories took %0.2f msec.",
+                 get_time(padding_start_time))
     logging.vlog(1, "Padded Observations' shape [%s]",
                  str(padded_observations.shape))
     logging.vlog(1, "Padded Actions' shape [%s]", str(padded_actions.shape))
@@ -908,9 +916,11 @@ def get_predictions(observations, rng=None):
     # these again here is not going to be correct and should be done in the
     # collect function.
 
+    log_prob_recompute_start_time = time.time()
     jax_rng_key, key = jax_random.split(jax_rng_key)
     log_probabs_traj, value_predictions_traj, _ = get_predictions(
         padded_observations, rng=key)
+    log_prob_recompute_time = get_time(log_prob_recompute_start_time)
 
     # Some assertions.
     B, T = padded_actions.shape  # pylint: disable=invalid-name
@@ -928,10 +938,9 @@ def get_predictions(observations, rng=None):
     epsilon_schedule = epsilon
 
     # Compute value and ppo losses.
-    cur_value_loss, cur_ppo_loss, cur_combined_loss = None, None, None
     jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)
     logging.vlog(2, "Starting to compute P&V loss.")
-    t = time.time()
+    loss_compute_start_time = time.time()
     cur_combined_loss, cur_ppo_loss, cur_value_loss, entropy_bonus = (
         combined_loss(
             policy_and_value_net_params,
@@ -948,99 +957,146 @@ def get_predictions(observations, rng=None):
             c1=c1,
             c2=c2,
             rng=key1))
+    loss_compute_time = get_time(loss_compute_start_time)
     logging.vlog(
         1,
         "Calculating P&V loss [%10.2f(%10.2f, %10.2f, %10.2f)] took %0.2f msec.",
         cur_combined_loss, cur_value_loss, cur_ppo_loss, entropy_bonus,
-        get_time(t))
+        get_time(loss_compute_start_time))
 
     jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)
-    if policy_and_value_net_apply:
-      logging.vlog(1, "Policy and Value Optimization")
-      t1 = time.time()
-      keys = jax_random.split(key1, num=num_optimizer_steps)
-      for j in range(num_optimizer_steps):
-        k1, k2, k3 = jax_random.split(keys[j], num=3)
-        t = time.time()
-        # Update the optimizer state.
-        policy_and_value_opt_state = policy_and_value_opt_step(
-            j,
-            policy_and_value_opt_state,
-            policy_and_value_opt_update,
-            policy_and_value_get_params,
-            policy_and_value_net_apply,
-            log_probabs_traj,
-            value_predictions_traj,
-            padded_observations,
-            padded_actions,
-            padded_rewards,
-            reward_mask,
-            c1=c1,
-            c2=c2,
-            gamma=gamma,
-            lambda_=lambda_,
-            epsilon=epsilon_schedule,
-            rng=k1)
-
-        # Compute the approx KL for early stopping.
-        new_policy_and_value_net_params = policy_and_value_get_params(
-            policy_and_value_opt_state)
-
-        log_probab_actions_new, _ = policy_and_value_net_apply(
-            padded_observations, new_policy_and_value_net_params, rng=k2)
-
-        approx_kl = approximate_kl(log_probab_actions_new, log_probabs_traj,
-                                   reward_mask)
-
-        early_stopping = enable_early_stopping and approx_kl > 1.5 * target_kl
-        if early_stopping:
-          logging.vlog(
-              1, "Early stopping policy and value optimization at iter: %d, "
-              "with approx_kl: %0.2f", j, approx_kl)
-          # We don't return right-away, we want the below to execute on the last
-          # iteration.
-
-        t2 = time.time()
-        if (((j + 1) % print_every_optimizer_steps == 0) or
-            (j == num_optimizer_steps - 1) or early_stopping):
-          # Compute and log the loss.
-          (loss_combined, loss_ppo, loss_value, entropy_bonus) = (
-              combined_loss(
-                  new_policy_and_value_net_params,
-                  log_probabs_traj,
-                  value_predictions_traj,
-                  policy_and_value_net_apply,
-                  padded_observations,
-                  padded_actions,
-                  padded_rewards,
-                  reward_mask,
-                  gamma=gamma,
-                  lambda_=lambda_,
-                  epsilon=epsilon_schedule,
-                  c1=c1,
-                  c2=c2,
-                  rng=k3))
-          logging.vlog(1, "One Policy and Value grad desc took: %0.2f msec",
-                       get_time(t, t2))
-          logging.vlog(
-              1, "Combined Loss(value, ppo, entropy_bonus) [%10.2f] ->"
-              " [%10.2f(%10.2f,%10.2f,%10.2f)]", cur_combined_loss,
-              loss_combined, loss_value, loss_ppo, entropy_bonus)
-
-        if early_stopping:
-          break
-
-      logging.vlog(
-          1, "Total Combined Loss reduction [%0.2f]%%",
-          (100 *
-           (cur_combined_loss - loss_combined) / np.abs(cur_combined_loss)))
-
-      logging.info(
-          "Epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
-          " Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)], took "
-          "[%2.5f msec].",
-          i, min_reward, max_reward,
-          avg_reward, loss_combined, loss_value, loss_ppo, entropy_bonus,
-          get_time(t1))
-
-      restore = False
+    logging.vlog(1, "Policy and Value Optimization")
+    optimization_start_time = time.time()
+    keys = jax_random.split(key1, num=num_optimizer_steps)
+    for j in range(num_optimizer_steps):
+      k1, k2, k3 = jax_random.split(keys[j], num=3)
+      t = time.time()
+      # Update the optimizer state.
+      policy_and_value_opt_state = policy_and_value_opt_step(
+          j,
+          policy_and_value_opt_state,
+          policy_and_value_opt_update,
+          policy_and_value_get_params,
+          policy_and_value_net_apply,
+          log_probabs_traj,
+          value_predictions_traj,
+          padded_observations,
+          padded_actions,
+          padded_rewards,
+          reward_mask,
+          c1=c1,
+          c2=c2,
+          gamma=gamma,
+          lambda_=lambda_,
+          epsilon=epsilon_schedule,
+          rng=k1)
+
+      # Compute the approx KL for early stopping.
+      new_policy_and_value_net_params = policy_and_value_get_params(
+          policy_and_value_opt_state)
+
+      log_probab_actions_new, _ = policy_and_value_net_apply(
+          padded_observations, new_policy_and_value_net_params, rng=k2)
+
+      approx_kl = approximate_kl(log_probab_actions_new, log_probabs_traj,
+                                 reward_mask)
+
+      early_stopping = enable_early_stopping and approx_kl > 1.5 * target_kl
+      if early_stopping:
+        logging.vlog(
+            1, "Early stopping policy and value optimization at iter: %d, "
+            "with approx_kl: %0.2f", j, approx_kl)
+        # We don't return right-away, we want the below to execute on the last
+        # iteration.
+
+      t2 = time.time()
+      if (((j + 1) % print_every_optimizer_steps == 0) or
+          (j == num_optimizer_steps - 1) or early_stopping):
+        # Compute and log the loss.
+        (loss_combined, loss_ppo, loss_value, entropy_bonus) = (
+            combined_loss(
+                new_policy_and_value_net_params,
+                log_probabs_traj,
+                value_predictions_traj,
+                policy_and_value_net_apply,
+                padded_observations,
+                padded_actions,
+                padded_rewards,
+                reward_mask,
+                gamma=gamma,
+                lambda_=lambda_,
+                epsilon=epsilon_schedule,
+                c1=c1,
+                c2=c2,
+                rng=k3))
+        logging.vlog(1, "One Policy and Value grad desc took: %0.2f msec",
+                     get_time(t, t2))
+        logging.vlog(
+            1, "Combined Loss(value, ppo, entropy_bonus) [%10.2f] ->"
+            " [%10.2f(%10.2f,%10.2f,%10.2f)]", cur_combined_loss, loss_combined,
+            loss_value, loss_ppo, entropy_bonus)
+
+      if early_stopping:
+        break
+
+    optimization_time = get_time(optimization_start_time)
+
+    logging.vlog(
+        1, "Total Combined Loss reduction [%0.2f]%%",
+        (100 * (cur_combined_loss - loss_combined) / np.abs(cur_combined_loss)))
+
+    # Save parameters every time we see the end of at least a fraction of batch
+    # number of trajectories that are done (not completed -- completed includes
+    # truncated and done).
+    # Also don't save too frequently, enforce a minimum gap.
+    # Or if this is the last iteration.
+    policy_save_start_time = time.time()
+    num_trajectories_done += num_done
+    if (((num_trajectories_done >= done_frac_for_policy_save * batch_size)
+         and (i - last_saved_at > eval_every_n)) or (i == epochs - 1)):
+      logging.vlog(1, "Epoch [% 6d] saving model.", i)
+      params_file = os.path.join(output_dir, "model-%06d.pkl" % i)
+      with gfile.GFile(params_file, "wb") as f:
+        pickle.dump(policy_and_value_net_params, f)
+      # Reset this number.
+      num_trajectories_done = 0
+      last_saved_at = i
+    policy_save_time = get_time(policy_save_start_time)
+
+    epoch_time = get_time(epoch_start_time)
+
+    logging.info(
+        "Epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
+        " Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)]", i, min_reward,
+        max_reward, avg_reward, loss_combined, loss_value, loss_ppo,
+        entropy_bonus)
+
+    timing_dict = {
+        "epoch": epoch_time,
+        "policy_eval": policy_eval_time,
+        "trajectory_collection": trajectory_collection_time,
+        "padding": padding_time,
+        "log_prob_recompute": log_prob_recompute_time,
+        "loss_compute": loss_compute_time,
+        "optimization": optimization_time,
+        "policy_save": policy_save_time,
+    }
+
+    for k, v in timing_dict.items():
+      timing_sw.scalar("timing/%s" % k, v, step=i)
+
+    max_key_len = max(len(k) for k in timing_dict)
+    timing_info_list = [
+        "%s : % 10.2f" % (k.rjust(max_key_len + 1), v)
+        for k, v in sorted(timing_dict.items())
+    ]
+    logging.info("Epoch [% 6d], Timings: \n%s", i, "\n".join(timing_info_list))
+
+    # Reset restore.
+    restore = False
+
+    # Flush summary writers once in a while.
+    if (i+1) % 1000 == 0 or i == epochs - 1:
+      train_sw.flush()
+      timing_sw.flush()
+      eval_sw.flush()
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 1f222d60a..080acc63a 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -43,7 +43,7 @@
 
 from absl import app
 from absl import flags
-import gym
+from absl import logging
 import jax
 from jax.config import config
 import numpy as onp
@@ -57,7 +57,6 @@
 
 FLAGS = flags.FLAGS
 
-flags.DEFINE_string("env_name", None, "Name of the environment to make.")
 flags.DEFINE_string("env_problem_name", None, "Name of the EnvProblem to make.")
 
 flags.DEFINE_integer("epochs", 100, "Number of epochs to run for.")
@@ -124,12 +123,13 @@
 flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
 flags.DEFINE_bool("enable_early_stopping", True,
                   "Whether to enable early stopping.")
-flags.DEFINE_bool("xm", False, "Are we running on borg?.")
+flags.DEFINE_bool("xm", False, "Copy atari roms?")
 flags.DEFINE_integer("eval_every_n", 100, "How frequently to eval the policy.")
 flags.DEFINE_integer("eval_batch_size", 4, "Batch size for evaluation.")
-flags.DEFINE_float("done_frac_for_policy_save", 0.5,
-                   "Fraction of the trajectories that should be done to "
-                   "checkpoint the policy.")
+flags.DEFINE_float(
+    "done_frac_for_policy_save", 0.5,
+    "Fraction of the trajectories that should be done to "
+    "checkpoint the policy.")
 
 
 def common_layers():
@@ -146,10 +146,6 @@ def atari_layers():
 
 def make_env(batch_size=8):
   """Creates the env."""
-  if FLAGS.env_name:
-    return gym.make(FLAGS.env_name)
-
-  assert FLAGS.env_problem_name
 
   # No resizing needed, so let's be on the normal EnvProblem.
   if not FLAGS.resize:  # None or False
@@ -190,8 +186,11 @@ def main(argv):
 
   if FLAGS.jax_debug_nans:
     config.update("jax_debug_nans", True)
+
   if FLAGS.use_tpu:
     config.update("jax_platform_name", "tpu")
+  else:
+    config.update("jax_platform_name", "gpu")
 
   # TODO(afrozm): Refactor.
   if "NoFrameskip" in FLAGS.env_problem_name and FLAGS.xm:
@@ -207,6 +206,7 @@ def main(argv):
 
   def run_training_loop():
     """Runs the training loop."""
+    logging.info("Starting the training loop.")
 
     policy_and_value_net_fun = functools.partial(
         ppo.policy_and_value_net,
@@ -242,7 +242,9 @@ def run_training_loop():
         output_dir=FLAGS.output_dir,
         eval_every_n=FLAGS.eval_every_n,
         done_frac_for_policy_save=FLAGS.done_frac_for_policy_save,
-        eval_env=eval_env)
+        eval_env=eval_env,
+        env_name=str(FLAGS.env_problem_name),
+    )
 
   if FLAGS.jax_debug_nans or FLAGS.disable_jit:
     with jax.disable_jit():
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index de34888a8..0727b51da 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -73,6 +73,7 @@ def test_training_loop(self):
           batch_size=batch_size,
           num_optimizer_steps=1,
           output_dir=output_dir,
+          env_name="CartPole-v0",
           random_seed=0)
 
 

From 8cae1775ff033208e4b80b6cbbdea22af86d65cc Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 24 May 2019 16:30:50 -0700
Subject: [PATCH 2059/2720] Internal

PiperOrigin-RevId: 249918669
---
 tensor2tensor/trax/rlax/ppo_main.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 080acc63a..e1d89d0c8 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -55,6 +55,7 @@
 from tensor2tensor.trax.models import atari_cnn
 from tensor2tensor.trax.rlax import ppo
 
+
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string("env_problem_name", None, "Name of the EnvProblem to make.")
@@ -183,6 +184,7 @@ def get_optimizer_fun(learning_rate):
 
 def main(argv):
   del argv
+  logging.info("Starting PPO Main.")
 
   if FLAGS.jax_debug_nans:
     config.update("jax_debug_nans", True)

From 86495653dbebefb8504e3fab02fbfa97a91f1b15 Mon Sep 17 00:00:00 2001
From: David So <davidso@google.com>
Date: Tue, 28 May 2019 10:59:47 -0700
Subject: [PATCH 2060/2720] Replace tf.gather with tf.slice, which is faster.

PiperOrigin-RevId: 250319612
---
 tensor2tensor/models/evolved_transformer.py | 34 ++++++++++-----------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 57adebad4..8d3d99e61 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -408,17 +408,18 @@ def evolved_transformer_decoder(decoder_input,
                   _CONV_BRANCHES_FIRST_LAYER_NAME] = tf.transpose(
                       tmp, perm=[1, 0, 2])
 
-              left_state_indexes = [
-                  decode_loop_step + i
-                  for i in range(_DECODER_LEFT_CONV_PADDING + 1)
-              ]
-              left_state = tf.gather(hidden_state, left_state_indexes, axis=1)
-              right_state_indexes = [
-                  decode_loop_step + i +
-                  (_DECODER_LEFT_CONV_PADDING - _DECODER_RIGHT_CONV_PADDING)
-                  for i in range(_DECODER_RIGHT_CONV_PADDING + 1)
-              ]
-              right_state = tf.gather(hidden_state, right_state_indexes, axis=1)
+              batch_size = hidden_state.shape.as_list()[0]
+              left_state = tf.slice(hidden_state, [0, decode_loop_step, 0], [
+                  batch_size, _DECODER_LEFT_CONV_PADDING + 1,
+                  hparams.hidden_size
+              ])
+              right_state = tf.slice(hidden_state, [
+                  0, decode_loop_step + _DECODER_LEFT_CONV_PADDING -
+                  _DECODER_RIGHT_CONV_PADDING, 0
+              ], [
+                  batch_size, _DECODER_RIGHT_CONV_PADDING + 1,
+                  hparams.hidden_size
+              ])
 
           else:  # No caching.
             left_state = tf.pad(
@@ -484,12 +485,11 @@ def evolved_transformer_decoder(decoder_input,
                   _CONV_BRANCHES_SECOND_LAYER_NAME] = tf.transpose(
                       tmp, perm=[1, 0, 2])
 
-              hidden_state_indexes = [
-                  decode_loop_step + i
-                  for i in range(_DECODER_FINAL_CONV_PADDING + 1)
-              ]
-              hidden_state = tf.gather(
-                  hidden_state, hidden_state_indexes, axis=1)
+              batch_size = hidden_state.shape.as_list()[0]
+              hidden_state = tf.slice(hidden_state, [0, decode_loop_step, 0], [
+                  batch_size, _DECODER_FINAL_CONV_PADDING + 1,
+                  hparams.hidden_size * 2
+              ])
           else:
             hidden_state = tf.pad(
                 hidden_state,

From 72a05e61e6a6ed9736a4f3c39f82c50a577e6df4 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 28 May 2019 13:50:49 -0700
Subject: [PATCH 2061/2720] Keep only the last checkpoint around. And use
 cPikcle (pickle didn't help on the cloud).

PiperOrigin-RevId: 250354246
---
 tensor2tensor/trax/rlax/ppo.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index bce25f2b0..b136ea26f 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -51,10 +51,10 @@
 
 import functools
 import os
-import pickle
 import time
 
 from absl import logging
+import cloudpickle as pickle
 import gym
 from jax import grad
 from jax import jit
@@ -1052,12 +1052,19 @@ def get_predictions(observations, rng=None):
     # Or if this is the last iteration.
     policy_save_start_time = time.time()
     num_trajectories_done += num_done
+    # TODO(afrozm): Refactor to trax.save_state.
     if (((num_trajectories_done >= done_frac_for_policy_save * batch_size)
-         and (i - last_saved_at > eval_every_n)) or (i == epochs - 1)):
+         and (i - last_saved_at > eval_every_n)
+         and (((i + 1) % eval_every_n == 0)))
+        or (i == epochs - 1)):
       logging.vlog(1, "Epoch [% 6d] saving model.", i)
+      old_model_files = gfile.glob(os.path.join(output_dir, "model-??????.pkl"))
       params_file = os.path.join(output_dir, "model-%06d.pkl" % i)
       with gfile.GFile(params_file, "wb") as f:
         pickle.dump(policy_and_value_net_params, f)
+      # Remove the old model files.
+      for path in old_model_files:
+        gfile.remove(path)
       # Reset this number.
       num_trajectories_done = 0
       last_saved_at = i

From 80f4e98da5511dd465abec89c6fa5a7aae03facc Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 28 May 2019 15:22:03 -0700
Subject: [PATCH 2062/2720] Allow access to epoch information in model-based
 RL.

PiperOrigin-RevId: 250372933
---
 tensor2tensor/models/research/rl.py            | 16 +++++++++++++++-
 tensor2tensor/rl/ppo.py                        |  6 ++++--
 tensor2tensor/rl/ppo_learner.py                |  5 ++++-
 tensor2tensor/rl/trainer_model_based.py        |  9 ++++-----
 tensor2tensor/rl/trainer_model_based_params.py | 16 ++++++++++++++++
 5 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 18ebf5f8d..9a668e8dd 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -72,6 +72,7 @@ def ppo_base_v1():
   hparams.add_hparam("logits_clip", 0.0)
   hparams.add_hparam("dropout_ppo", 0.1)
   hparams.add_hparam("effective_num_agents", None)
+  hparams.add_hparam("use_epochs", True)
   # TODO(afrozm): Clean this up, this is used in PPO learner to get modalities.
   hparams.add_hparam("policy_problem_name", "dummy_policy_problem")
   return hparams
@@ -286,7 +287,8 @@ def make_simulated_env_fn_from_hparams(real_env, hparams, **extra_kwargs):
   )
 
 
-def get_policy(observations, hparams, action_space, distributional_size=1):
+def get_policy(observations, hparams, action_space,
+               distributional_size=1, epoch=-1):
   """Get a policy network.
 
   Args:
@@ -294,6 +296,7 @@ def get_policy(observations, hparams, action_space, distributional_size=1):
     hparams: parameters
     action_space: action space
     distributional_size: optional number of buckets for distributional RL
+    epoch: optional epoch number
 
   Returns:
     Tuple (action logits, value).
@@ -327,6 +330,7 @@ def get_policy(observations, hparams, action_space, distributional_size=1):
     target_value_shape_suffix = [num_target_frames, distributional_size]
   features = {
       "inputs": observations,
+      "epoch": tf.constant(epoch + 1),
       "input_action": tf.zeros(obs_shape[:2] + [1], dtype=tf.int32),
       "input_reward": tf.zeros(obs_shape[:2] + [1], dtype=tf.int32),
       "targets": tf.zeros(obs_shape[:1] + [num_target_frames] + obs_shape[2:]),
@@ -340,6 +344,7 @@ def get_policy(observations, hparams, action_space, distributional_size=1):
           obs_shape[:1] + target_value_shape_suffix)
   }
   model.distributional_value_size = max(distributional_size, 1)
+  model.use_epochs = hparams.use_epochs
   with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
     t2t_model.create_dummy_vars()
     (targets, _) = model(features)
@@ -567,6 +572,7 @@ class PolicyBase(t2t_model.T2TModel):
   def __init__(self, *args, **kwargs):
     super(PolicyBase, self).__init__(*args, **kwargs)
     self.distributional_value_size = 1
+    self.use_epochs = False
 
   def loss(self, *args, **kwargs):
     return 0.0
@@ -721,6 +727,14 @@ def body(self, features):
                            activation=tf.nn.relu, padding="same")
 
       flat_x = tf.layers.flatten(x)
+      if self.use_epochs:
+        epoch = features["epoch"] + tf.zeros([x_shape[0]], dtype=tf.int32)
+        # Randomly set epoch to 0 in some cases as that's the inference value.
+        rand = tf.random.uniform([x_shape[0]])
+        epoch = tf.where(rand < 0.1, tf.zeros_like(epoch), epoch)
+        # Embed the epoch number.
+        emb_epoch = common_layers.embedding(epoch, 32, 32)  # [batch, 32]
+        flat_x = tf.concat([flat_x, emb_epoch], axis=1)
       flat_x = tf.layers.dropout(flat_x, rate=dropout)
       x = tf.layers.dense(flat_x, 128, activation=tf.nn.relu)
 
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index 2acada674..4baafc3bf 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -30,7 +30,7 @@
 import tensorflow_probability as tfp
 
 
-def define_ppo_step(data_points, hparams, action_space, lr,
+def define_ppo_step(data_points, hparams, action_space, lr, epoch=-1,
                     distributional_size=1, distributional_subscale=0.04):
   """Define ppo step."""
   observation, action, discounted_reward, norm_advantage, old_pdf = data_points
@@ -40,6 +40,7 @@ def define_ppo_step(data_points, hparams, action_space, lr,
       observation, [obs_shape[0] * obs_shape[1]] + obs_shape[2:]
   )
   (logits, new_value) = get_policy(observation, hparams, action_space,
+                                   epoch=epoch,
                                    distributional_size=distributional_size)
   logits = tf.reshape(logits, obs_shape[:2] + [action_space.n])
   new_policy_dist = tfp.distributions.Categorical(logits=logits)
@@ -103,7 +104,7 @@ def _distributional_to_value(value_d, size, subscale, threshold):
 
 def define_ppo_epoch(memory, hparams, action_space, batch_size,
                      distributional_size=1, distributional_subscale=0.04,
-                     distributional_threshold=0.0):
+                     distributional_threshold=0.0, epoch=-1):
   """PPO epoch."""
   observation, reward, done, action, old_pdf, value_sm = memory
 
@@ -169,6 +170,7 @@ def define_ppo_epoch(memory, hparams, action_space, batch_size,
           a, define_ppo_step(
               [tf.gather(t, indices_of_batches[i, :]) for t in input_tensors],
               hparams, action_space, lr,
+              epoch=epoch,
               distributional_size=distributional_size,
               distributional_subscale=distributional_subscale
           )),
diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index a7b6a166c..770f333c2 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -89,6 +89,7 @@ def train(self,
                   distributional_size=self._distributional_size,
                   distributional_subscale=self._distributional_subscale,
                   distributional_threshold=self._distributional_threshold,
+                  epoch=epoch if simulated else -1,
                   frame_stack_size=self.frame_stack_size,
                   force_beginning_resets=simulated))
 
@@ -165,6 +166,7 @@ def _define_train(
     distributional_size=1,
     distributional_subscale=0.04,
     distributional_threshold=0.0,
+    epoch=-1,
     **collect_kwargs
 ):
   """Define the training setup."""
@@ -181,7 +183,8 @@ def _define_train(
       memory, ppo_hparams, train_env.action_space, train_env.batch_size,
       distributional_size=distributional_size,
       distributional_subscale=distributional_subscale,
-      distributional_threshold=distributional_threshold)
+      distributional_threshold=distributional_threshold,
+      epoch=epoch)
   train_summary = tf.summary.merge([collect_summary, ppo_summary])
 
   if ppo_hparams.eval_every_epochs:
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 37b4504dc..5a57f875a 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -57,8 +57,8 @@ def real_env_step_increment(hparams):
   ))
 
 
-def world_model_step_increment(hparams, is_initial_epoch):
-  if is_initial_epoch:
+def world_model_step_increment(hparams, epoch):
+  if epoch in [0, 1, 4, 9, 14]:
     multiplier = hparams.initial_epoch_train_steps_multiplier
   else:
     multiplier = 1
@@ -162,6 +162,7 @@ def train_agent(real_env, learner, world_model_dir, hparams, epoch):
 
   final_epoch = hparams.epochs - 1
   is_special_epoch = (epoch + 3) == final_epoch or (epoch + 7) == final_epoch
+  is_special_epoch = is_special_epoch or (epoch == 1)  # Make 1 special too.
   is_final_epoch = epoch == final_epoch
   env_step_multiplier = 3 if is_final_epoch else 2 if is_special_epoch else 1
   learner.train(
@@ -200,9 +201,7 @@ def train_world_model(
     env, data_dir, output_dir, hparams, world_model_steps_num, epoch
 ):
   """Train the world model on problem_name."""
-  world_model_steps_num += world_model_step_increment(
-      hparams, is_initial_epoch=(epoch == 0)
-  )
+  world_model_steps_num += world_model_step_increment(hparams, epoch)
   model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
   model_hparams.learning_rate = model_hparams.learning_rate_constant
   if epoch > 0:
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index bd7d7fa55..fc50108d8 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -360,6 +360,22 @@ def rlmb_base_stochastic_discrete_75k_model_steps():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_base_stochastic_discrete_20k_model_steps():
+  """Base SD setting with 20k WM steps."""
+  hparams = rlmb_base_stochastic_discrete()
+  hparams.model_train_steps = 20000
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_base_stochastic_discrete_30k_model_steps():
+  """Base SD setting with 20k WM steps."""
+  hparams = rlmb_base_stochastic_discrete()
+  hparams.model_train_steps = 30000
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_base_stochastic_discrete_200k():
   """Base setting with stochastic discrete model with 200k steps."""

From 6667752e92d50edd078b93e7a3af695b1588f003 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 28 May 2019 17:02:00 -0700
Subject: [PATCH 2063/2720] Corrections to distributional PPO implementation
 wrt. end values.

PiperOrigin-RevId: 250390932
---
 tensor2tensor/models/research/rl.py |  4 +-
 tensor2tensor/rl/ppo.py             | 65 +++++++++++++++++++----------
 2 files changed, 46 insertions(+), 23 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 9a668e8dd..f0262d94a 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -140,7 +140,7 @@ def ppo_original_params():
 def ppo_dist_params():
   """Parameters based on the original paper modified for distributional RL."""
   hparams = ppo_original_params()
-  hparams.learning_rate_constant = 5e-4
+  hparams.learning_rate_constant = 1e-3
   return hparams
 
 
@@ -563,7 +563,7 @@ def rlmf_eval_dist():
 def rlmf_eval_dist_threshold():
   """Distributional set of hparams for model-free PPO."""
   hparams = rlmf_eval_dist()
-  hparams.distributional_threshold = 0.2
+  hparams.distributional_threshold = 0.5
   return hparams
 
 
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index 4baafc3bf..16bad15cf 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -33,7 +33,9 @@
 def define_ppo_step(data_points, hparams, action_space, lr, epoch=-1,
                     distributional_size=1, distributional_subscale=0.04):
   """Define ppo step."""
-  observation, action, discounted_reward, norm_advantage, old_pdf = data_points
+  del distributional_subscale
+  (observation, action, discounted_reward, discounted_reward_probs,
+   norm_advantage, old_pdf) = data_points
 
   obs_shape = common_layers.shape_list(observation)
   observation = tf.reshape(
@@ -58,14 +60,22 @@ def define_ppo_step(data_points, hparams, action_space, lr, epoch=-1,
   if distributional_size > 1:
     new_value = tf.reshape(new_value, obs_shape[:2] + [distributional_size])
     new_value = tf.nn.log_softmax(new_value, axis=-1)
-    # We assume the values range from (-half, half) -- set subscale accordingly.
-    half = (distributional_size // 2) * distributional_subscale
-    # To make values integers, we add half (to move range to (0, 2*half) and
-    # then multiply by subscale after which we floor to get nearest int.
-    quantized_dr = tf.floor(
-        (discounted_reward + half) / distributional_subscale)
-    hot_dr = tf.one_hot(tf.cast(quantized_dr, tf.int32), distributional_size)
-    value_loss = - tf.reduce_sum(new_value * hot_dr, axis=-1)
+    # The above is the new value distribution. We are also given as discounted
+    # reward the value distribution and the corresponding probabilities.
+    # The given discounted reward is already rounded to integers but in range
+    # increased by 2x for greater fidelity. Increase range of new_values here.
+    new_value_shifted = tf.concat([new_value[1:], new_value[-1:]], axis=0)
+    new_value_mean = (new_value + new_value_shifted) / 2
+    new_value = tf.concat([tf.expand_dims(new_value, axis=-1),
+                           tf.expand_dims(new_value_mean, axis=-1)], -1)
+    new_value = tf.reshape(new_value, tf.shape(new_value_mean))
+    # Cast discounted reward to integers and gather the new log-probs for them.
+    discounted_reward = tf.cast(discounted_reward, tf.int32)
+    value_loss = tf.batch_gather(new_value, discounted_reward)
+    # Weight the gathered (new) log-probs by the old probabilities.
+    discounted_reward_probs = tf.expand_dims(discounted_reward_probs, axis=1)
+    value_loss = - tf.reduce_sum(value_loss * discounted_reward_probs, axis=-1)
+    # Take the mean over batch and time as final loss, multiply by coefficient.
     value_loss = hparams.value_loss_coef * tf.reduce_mean(value_loss)
   else:
     new_value = tf.reshape(new_value, obs_shape[:2])
@@ -123,19 +133,32 @@ def define_ppo_epoch(memory, hparams, action_space, batch_size,
     value = _distributional_to_value(
         value_sm, distributional_size, distributional_subscale,
         distributional_threshold)
-  plain_value = value
-  if distributional_threshold > 1:
-    plain_value = _distributional_to_value(
-        value_sm, distributional_size, distributional_subscale, 0.0)
 
   advantage = calculate_generalized_advantage_estimator(
       reward, value, done, hparams.gae_gamma, hparams.gae_lambda)
 
-  discounted_reward = tf.stop_gradient(advantage + value[:-1])
   if distributional_size > 1:
-    end_values = plain_value[-1]
-    discounted_reward = tf.stop_gradient(discounted_rewards(
-        reward, done, hparams.gae_gamma, end_values))
+    # Create discounted reward values range.
+    half = distributional_size // 2
+    value_range = tf.to_float(tf.range(-half, half)) + 0.5  # Mid-bucket value.
+    value_range *= distributional_subscale
+    # Acquire new discounted rewards by using the above range as end-values.
+    end_values = tf.expand_dims(value_range, 0)
+    discounted_reward = discounted_rewards(
+        reward, done, hparams.gae_gamma, end_values)
+    # Re-normalize the discounted rewards to integers, in [0, dist_size] range.
+    discounted_reward /= distributional_subscale
+    discounted_reward += half
+    discounted_reward = tf.maximum(discounted_reward, 0.0)
+    discounted_reward = tf.minimum(discounted_reward, distributional_size)
+    # Multiply the rewards by 2 for greater fidelity and round to integers.
+    discounted_reward = tf.stop_gradient(tf.round(2 * discounted_reward))
+    # The probabilities corresponding to the end values from old predictions.
+    discounted_reward_prob = tf.stop_gradient(value_sm[-1])
+    discounted_reward_prob = tf.nn.softmax(discounted_reward_prob, axis=-1)
+  else:
+    discounted_reward = tf.stop_gradient(advantage + value[:-1])
+    discounted_reward_prob = discounted_reward  # Unused in this case.
 
   advantage_mean, advantage_variance = tf.nn.moments(advantage, axes=[0, 1],
                                                      keep_dims=True)
@@ -163,7 +186,7 @@ def define_ppo_epoch(memory, hparams, action_space, batch_size,
   indices_of_batches = tf.reshape(shuffled_indices,
                                   shape=(-1, hparams.optimization_batch_size))
   input_tensors = [observation, action, discounted_reward,
-                   advantage_normalized, old_pdf]
+                   discounted_reward_prob, advantage_normalized, old_pdf]
 
   ppo_step_rets = tf.scan(
       lambda a, i: add_lists_elementwise(  # pylint: disable=g-long-lambda
@@ -221,11 +244,11 @@ def calculate_generalized_advantage_estimator(
 
 def discounted_rewards(reward, done, gae_gamma, end_values):
   """Discounted rewards."""
-  not_done = 1 - tf.cast(done, tf.float32)
-  end_values = end_values * not_done[-1, :]
+  not_done = tf.expand_dims(1 - tf.cast(done, tf.float32), axis=2)
+  end_values = end_values * not_done[-1, :, :]
   return_ = tf.scan(
       lambda agg, cur: cur + gae_gamma * agg,
-      reward * not_done,
+      tf.expand_dims(reward, axis=2) * not_done,
       initializer=end_values,
       reverse=True,
       back_prop=False,

From c1918fb7a70ba175932df2b3feaeaa0ae7377b4e Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 29 May 2019 09:15:50 -0700
Subject: [PATCH 2064/2720] Internal

PiperOrigin-RevId: 250504295
---
 tensor2tensor/trax/rlax/ppo_main.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index e1d89d0c8..6dc40134e 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -50,7 +50,6 @@
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import rendered_env_problem
 from tensor2tensor.rl import gym_utils
-from tensor2tensor.rl.google import atari_utils
 from tensor2tensor.trax import layers
 from tensor2tensor.trax.models import atari_cnn
 from tensor2tensor.trax.rlax import ppo
@@ -194,10 +193,6 @@ def main(argv):
   else:
     config.update("jax_platform_name", "gpu")
 
-  # TODO(afrozm): Refactor.
-  if "NoFrameskip" in FLAGS.env_problem_name and FLAGS.xm:
-    FLAGS.atari_roms_path = "local_ram_fs_tmp"
-    atari_utils.copy_roms()
 
   # Make an env here.
   env = make_env(batch_size=FLAGS.batch_size)

From 0bf204cd4ac2f3cb57ad3f677662bbfc57b16891 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 29 May 2019 09:25:18 -0700
Subject: [PATCH 2065/2720] Internal

PiperOrigin-RevId: 250505862
---
 tensor2tensor/data_generators/translate.py      | 4 ++++
 tensor2tensor/data_generators/translate_enro.py | 4 ++++
 2 files changed, 8 insertions(+)

diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 712fb5e08..ff1014a41 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -22,6 +22,7 @@
 import gzip
 import os
 import tarfile
+import zipfile
 from tensor2tensor.data_generators import cleaner_en_xx
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -178,6 +179,9 @@ def compile_data(tmp_dir, datasets, filename, datatypes_to_clean=None):
         compressed_filepath = os.path.join(tmp_dir, compressed_filename)
         if url.startswith("http"):
           generator_utils.maybe_download(tmp_dir, compressed_filename, url)
+        if compressed_filename.endswith(".zip"):
+          zipfile.ZipFile(os.path.join(compressed_filepath),
+                          "r").extractall(tmp_dir)
 
         if dataset[1][0] == "tmx":
           cleaning_requested = "tmx" in datatypes_to_clean
diff --git a/tensor2tensor/data_generators/translate_enro.py b/tensor2tensor/data_generators/translate_enro.py
index 35115570f..f03e96729 100644
--- a/tensor2tensor/data_generators/translate_enro.py
+++ b/tensor2tensor/data_generators/translate_enro.py
@@ -32,6 +32,10 @@
     [
         "http://www.statmt.org/europarl/v7/ro-en.tgz",
         ("europarl-v7.ro-en.en", "europarl-v7.ro-en.ro")
+    ],
+    [
+        "http://opus.nlpl.eu/download.php?f=SETIMES/v2/moses/en-ro.txt.zip",
+        ("SETIMES.en-ro.en", "SETIMES.en-ro.ro")
     ]
 ]
 _ENRO_TEST_DATASETS = [

From 44c281b4bc5e58833de683e60fa4a313514bc927 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 29 May 2019 11:49:45 -0700
Subject: [PATCH 2066/2720] Correct shape mistake from previous commit.

PiperOrigin-RevId: 250536110
---
 tensor2tensor/rl/ppo.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index 16bad15cf..92d65b35b 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -60,6 +60,7 @@ def define_ppo_step(data_points, hparams, action_space, lr, epoch=-1,
   if distributional_size > 1:
     new_value = tf.reshape(new_value, obs_shape[:2] + [distributional_size])
     new_value = tf.nn.log_softmax(new_value, axis=-1)
+    value_shape = common_layers.shape_list(new_value)
     # The above is the new value distribution. We are also given as discounted
     # reward the value distribution and the corresponding probabilities.
     # The given discounted reward is already rounded to integers but in range
@@ -68,7 +69,7 @@ def define_ppo_step(data_points, hparams, action_space, lr, epoch=-1,
     new_value_mean = (new_value + new_value_shifted) / 2
     new_value = tf.concat([tf.expand_dims(new_value, axis=-1),
                            tf.expand_dims(new_value_mean, axis=-1)], -1)
-    new_value = tf.reshape(new_value, tf.shape(new_value_mean))
+    new_value = tf.reshape(new_value, value_shape[:-1] + [2 * value_shape[-1]])
     # Cast discounted reward to integers and gather the new log-probs for them.
     discounted_reward = tf.cast(discounted_reward, tf.int32)
     value_loss = tf.batch_gather(new_value, discounted_reward)

From 81cb61809641736eaea88a55706be4e985aefbd9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 29 May 2019 14:21:00 -0700
Subject: [PATCH 2067/2720] Compute and log raw unclipped rewards.

PiperOrigin-RevId: 250565529
---
 tensor2tensor/envs/trajectory.py |  7 ++++++-
 tensor2tensor/trax/rlax/ppo.py   | 18 ++++++++++++++----
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index 412fca063..3c412e499 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -107,9 +107,14 @@ def rewards_np(self):
     # The first reward is None, so let's skip it.
     return np.stack([ts.processed_reward for ts in self.time_steps[1:]])
 
+  @property
+  def raw_rewards_np(self):
+    return np.stack([ts.raw_reward for ts in self.time_steps[1:]])
+
   @property
   def as_numpy(self):
-    return self.observations_np, self.actions_np, self.rewards_np
+    return (self.observations_np, self.actions_np, self.rewards_np,
+            self.raw_rewards_np)
 
 
 class BatchTrajectory(object):
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index b136ea26f..b79de34b1 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -167,7 +167,7 @@ def collect_trajectories(env,
 
   assert isinstance(env, env_problem.EnvProblem)
   # This is an env_problem, run its collect function.
-  return env_problem_utils.play_env_problem_with_policy(
+  trajs, num_done = env_problem_utils.play_env_problem_with_policy(
       env,
       policy_fun,
       num_trajectories=num_trajectories,
@@ -177,6 +177,8 @@ def collect_trajectories(env,
       eps=epsilon,
       reset=reset,
       rng=rng)
+  # Skip returning raw_rewards here, since they aren't used.
+  return [(t[0], t[1], t[2]) for t in trajs], num_done
 
 
 # This function can probably be simplified, ask how?
@@ -707,6 +709,7 @@ def evaluate_policy(eval_env,
   """Evaluate the policy."""
 
   avg_rewards = {}
+  avg_rewards_unclipped = {}
   for policy in [
       env_problem_utils.CATEGORICAL_SAMPLING, env_problem_utils.GUMBEL_SAMPLING,
       env_problem_utils.EPSILON_GREEDY
@@ -721,7 +724,9 @@ def evaluate_policy(eval_env,
         rng=rng)
     avg_rewards[policy] = float(sum(
         np.sum(traj[2]) for traj in trajs)) / len(trajs)
-  return avg_rewards
+    avg_rewards_unclipped[policy] = float(sum(
+        np.sum(traj[3]) for traj in trajs)) / len(trajs)
+  return avg_rewards, avg_rewards_unclipped
 
 
 def maybe_restore_params(output_dir, policy_and_value_net_params):
@@ -849,7 +854,7 @@ def get_predictions(observations, rng=None):
 
       logging.vlog(1, "Epoch [% 6d] evaluating policy.", i)
 
-      avg_reward = evaluate_policy(
+      avg_reward, avg_reward_unclipped = evaluate_policy(
           eval_env,
           get_predictions,
           boundary,
@@ -857,7 +862,12 @@ def get_predictions(observations, rng=None):
           rng=key)
       for k, v in avg_reward.items():
         eval_sw.scalar("eval/mean_reward/%s" % k, v, step=i)
-        logging.info("Epoch [% 6d] Policy Evaluation [%s] = %10.2f", i, k, v)
+        logging.info("Epoch [% 6d] Policy Evaluation (clipped) [%s] = %10.2f",
+                     i, k, v)
+      for k, v in avg_reward_unclipped.items():
+        eval_sw.scalar("eval/mean_reward_unclipped/%s" % k, v, step=i)
+        logging.info("Epoch [% 6d] Policy Evaluation (unclipped) [%s] = %10.2f",
+                     i, k, v)
     policy_eval_time = get_time(policy_eval_start_time)
 
     trajectory_collection_start_time = time.time()

From dd9930607d87b57e1ae655d2c861d608b187b0e5 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 29 May 2019 15:58:15 -0700
Subject: [PATCH 2068/2720] Use lax.scan in `rewards_to_go` and `deltas`.

PiperOrigin-RevId: 250584446
---
 tensor2tensor/trax/rlax/ppo.py | 65 +++++++++++++---------------------
 1 file changed, 25 insertions(+), 40 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index b79de34b1..cd8b4b43d 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -129,7 +129,6 @@ def optimizer_fun(net_params, step_size=1e-3):
 # Should this be collect 'n' trajectories, or
 # Run the env for 'n' steps and take completed trajectories, or
 # Any other option?
-# TODO(afrozm): Replace this with EnvProblem?
 def collect_trajectories(env,
                          policy_fun,
                          num_trajectories=1,
@@ -275,7 +274,6 @@ def pad_trajectories(trajectories, boundary=20):
       padded_observations), np.stack(padded_actions), np.stack(padded_rewards)
 
 
-# TODO(afrozm): JAX-ify this, this is too slow for pong.
 def rewards_to_go(rewards, mask, gamma=0.99):
   r"""Computes rewards to go.
 
@@ -292,39 +290,23 @@ def rewards_to_go(rewards, mask, gamma=0.99):
   Returns:
     rewards to go, np.ndarray of shape (B, T).
   """
-  B, T = rewards.shape  # pylint: disable=invalid-name,unused-variable
+  # B, T = rewards.shape
 
   masked_rewards = rewards * mask  # (B, T)
 
-  # We use the following recurrence relation, derived from the equation above:
-  #
-  # r2g[t+1] = (r2g[t] - r[t]) / gamma
-  #
-  # This means we'll need to calculate r2g[0] first and then r2g[1] and so on ..
-  #
-  # **However** this leads to overflows for long sequences: r2g[t] - r[t] > 0
-  # and gamma < 1.0, so the division keeps increasing.
-  #
-  # So we just run the recurrence in reverse, i.e.
-  #
-  # r2g[t] = r[t] + (gamma*r2g[t+1])
-  #
-  # This is much better, but might have lost updates since the (small) rewards
-  # at earlier time-steps may get added to a (very?) large sum.
-
-  # Compute r2g_{T-1} at the start and then compute backwards in time.
-  r2gs = [masked_rewards[:, -1]]
-
-  # Go from T-2 down to 0.
-  for t in reversed(range(T - 1)):
-    r2gs.append(masked_rewards[:, t] + (gamma * r2gs[-1]))
-
-  # The list should have length T.
-  assert T == len(r2gs)
-
-  # First we stack them in the correct way to make it (B, T), but these are
-  # still from newest (T-1) to oldest (0), so then we flip it on time axis.
-  return np.flip(np.stack(r2gs, axis=1), axis=1)
+  reversed_rewards = np.flip(masked_rewards, axis=1)  # (B, T) flipped on time.
+  rrt = np.transpose(reversed_rewards)  # (T, B) transpose to scan over time.
+
+  def discounting_add(carry, reward):
+    x = reward + (gamma * carry)
+    return x, x
+
+  _, ys = lax.scan(discounting_add,
+                   np.zeros_like(rrt[0], dtype=np.float32),
+                   rrt.astype(np.float32))
+
+  # ys is (T, B) and T is in reverse order.
+  return np.flip(np.transpose(ys), axis=1)
 
 
 @jit
@@ -373,7 +355,6 @@ def value_loss_given_predictions(value_prediction,
   return np.sum(loss) / np.sum(reward_mask)
 
 
-# TODO(afrozm): JAX-ify this, this is too slow for pong.
 def deltas(predicted_values, rewards, mask, gamma=0.99):
   r"""Computes TD-residuals from V(s) and rewards.
 
@@ -392,14 +373,18 @@ def deltas(predicted_values, rewards, mask, gamma=0.99):
     ndarray of shape (B, T) of one-step TD-residuals.
   """
 
-  # `d`s are basically one-step TD residuals.
-  d = []
-  _, T = rewards.shape  # pylint: disable=invalid-name
-  for t in range(T):
-    d.append(rewards[:, t] + (gamma * predicted_values[:, t + 1]) -
-             predicted_values[:, t])
+  v0 = np.transpose(predicted_values)[0]        # (B,) just V_{b, 1}
+  v1 = np.transpose(predicted_values)[:-1]      # (T, B) without V_{b, T+1}
+  rt = rewards.astype(np.float32).T   # (T, B)
+
+  def td_residual(carry, inps):
+    r, v_next = inps
+    v = carry
+    return v_next, (r + gamma * v_next - v)
+
+  _, d = lax.scan(td_residual, v0, [rt, v1])
 
-  return np.array(d).T * mask
+  return np.transpose(d) * mask
 
 
 def gae_advantages(td_deltas, mask, lambda_=0.95, gamma=0.99):

From a03d78a1764a5d58912ca918889aefcfe5541f22 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 29 May 2019 17:23:21 -0700
Subject: [PATCH 2069/2720] Interpret Python list of layers as wrapped by a
 Serial layer.

Within Serial, perform recursive flattening so that lists of lists of
(lists of...) layers are treated as a single flat list of layers. These changes
enable cleaner/easier composition of layers (e.g., inlined list comprehensions).

Also modify some name patterns:
  - num_... --> n_...  # number of ...'s
  - ..._depth --> d_...  # dimensionality of ...
  - ..._fun --> ..._fn  # function (callable) as a parameter

PiperOrigin-RevId: 250599386
---
 .../chunked_transformer_imagenet64_8gb.gin    |  10 +-
 .../trax/configs/resnet50_imagenet_8gb.gin    |   4 +-
 .../configs/resnet50_imagenet_8gb_testing.gin |   4 +-
 .../trax/configs/transformer_big_lm1b_8gb.gin |   8 +-
 .../trax/configs/transformer_imdb_8gb.gin     |  10 +-
 .../trax/configs/transformer_lm1b_8gb.gin     |   8 +-
 .../configs/transformer_lm1b_8gb_testing.gin  |   8 +-
 .../trax/configs/transformer_wmt_ende_8gb.gin |   8 +-
 .../trax/configs/wide_resnet_cifar10_8gb.gin  |   6 +-
 tensor2tensor/trax/layers/attention.py        |  69 +++--
 tensor2tensor/trax/layers/base.py             |  26 +-
 tensor2tensor/trax/layers/combinators.py      |  75 ++---
 tensor2tensor/trax/layers/combinators_test.py |   6 +-
 tensor2tensor/trax/layers/convolution.py      |   2 +-
 tensor2tensor/trax/layers/core.py             |  14 +-
 tensor2tensor/trax/layers/rnn.py              |  71 ++---
 tensor2tensor/trax/models/atari_cnn.py        |  19 +-
 tensor2tensor/trax/models/mlp.py              |  18 +-
 tensor2tensor/trax/models/neural_gpu.py       |  20 +-
 tensor2tensor/trax/models/neural_gpu_test.py  |   3 +-
 .../models/research/chunked_transformer.py    |  71 +++--
 tensor2tensor/trax/models/resnet.py           | 140 +++++-----
 tensor2tensor/trax/models/transformer.py      | 264 +++++++++---------
 tensor2tensor/trax/rlax/ppo.py                |  58 ++--
 .../trax/rlax/ppo_training_loop_test.py       |   4 +-
 tensor2tensor/trax/trax.py                    | 141 +++++-----
 tensor2tensor/trax/trax_test.py               |  35 +--
 27 files changed, 536 insertions(+), 566 deletions(-)

diff --git a/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin b/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin
index 2c7fd1fb9..119b3f740 100644
--- a/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin
+++ b/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin
@@ -11,7 +11,7 @@ batch_fun.max_eval_length = 12288  # 64 * 64 * 3
 
 # Parameters for inputs:
 # ==============================================================================
-inputs.num_chunks = 64
+inputs.n_chunks = 64
 inputs.data_dir = None
 inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
 inputs.input_name = 'targets'
@@ -33,11 +33,11 @@ train.train_steps = 500000
 
 # Parameters for ChunkedTransformerLM:
 # ==============================================================================
+ChunkedTransformerLM.d_feature = 1024
+ChunkedTransformerLM.d_feedforward = 4096
 ChunkedTransformerLM.dropout = 0.1
-ChunkedTransformerLM.feature_depth = 1024
-ChunkedTransformerLM.feedforward_depth = 4096
 ChunkedTransformerLM.max_len = 12288  # 64 * 64 * 3
 ChunkedTransformerLM.mode = 'train'
-ChunkedTransformerLM.num_heads = 4
-ChunkedTransformerLM.num_layers = 3
+ChunkedTransformerLM.n_heads = 4
+ChunkedTransformerLM.n_layers = 3
 ChunkedTransformerLM.vocab_size = 256
diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
index 22d99b9b6..94e9f6b37 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
@@ -29,8 +29,8 @@ Momentum.mass = 0.9
 
 # Parameters for Resnet50:
 # ==============================================================================
-Resnet50.hidden_size = 64
-Resnet50.num_output_classes = 1001
+Resnet50.d_hidden = 64
+Resnet50.n_output_classes = 1001
 
 # Parameters for train:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
index ad1c50cf0..6ecd54a04 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
@@ -29,8 +29,8 @@ Momentum.mass = 0.9
 
 # Parameters for Resnet50:
 # ==============================================================================
-Resnet50.hidden_size = 64
-Resnet50.num_output_classes = 1001
+Resnet50.d_hidden = 64
+Resnet50.n_output_classes = 1001
 
 # Parameters for train:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
index 6c88ca76a..fe6ac3a25 100644
--- a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
@@ -43,11 +43,11 @@ train.train_steps = 500000
 
 # Parameters for TransformerLM:
 # ==============================================================================
+TransformerLM.d_feature = 1024
+TransformerLM.d_feedforward = 8192
 TransformerLM.dropout = 0.1
-TransformerLM.feature_depth = 1024
-TransformerLM.feedforward_depth = 8192
 TransformerLM.max_len = 2048
 TransformerLM.mode = 'train'
-TransformerLM.num_heads = 8
-TransformerLM.num_layers = 8
+TransformerLM.n_heads = 8
+TransformerLM.n_layers = 8
 TransformerLM.vocab_size = 32000
diff --git a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
index 3fdac10aa..784cb5bde 100644
--- a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
@@ -39,12 +39,12 @@ train.train_steps = 1000
 
 # Parameters for TransformerLM:
 # ==============================================================================
+TransformerEncoder.d_feature = 512
+TransformerEncoder.d_feedforward = 2048
 TransformerEncoder.dropout = 0.1
-TransformerEncoder.feature_depth = 512
-TransformerEncoder.feedforward_depth = 2048
 TransformerEncoder.max_len = 2048
 TransformerEncoder.mode = 'train'
-TransformerEncoder.num_classes = 10
-TransformerEncoder.num_heads = 8
-TransformerEncoder.num_layers = 6
+TransformerEncoder.n_classes = 10
+TransformerEncoder.n_heads = 8
+TransformerEncoder.n_layers = 6
 TransformerEncoder.vocab_size = 32000
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index 3cf306276..a6ac5102b 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -43,11 +43,11 @@ train.train_steps = 500000
 
 # Parameters for TransformerLM:
 # ==============================================================================
+TransformerLM.d_feature = 512
+TransformerLM.d_feedforward = 2048
 TransformerLM.dropout = 0.1
-TransformerLM.feature_depth = 512
-TransformerLM.feedforward_depth = 2048
 TransformerLM.max_len = 2048
 TransformerLM.mode = 'train'
-TransformerLM.num_heads = 8
-TransformerLM.num_layers = 6
+TransformerLM.n_heads = 8
+TransformerLM.n_layers = 6
 TransformerLM.vocab_size = 32000
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
index fe5b38ae1..e3c09a5de 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
@@ -43,11 +43,11 @@ train.train_steps = 100000
 
 # Parameters for TransformerLM:
 # ==============================================================================
+TransformerLM.d_feature = 512
+TransformerLM.d_feedforward = 2048
 TransformerLM.dropout = 0.1
-TransformerLM.feature_depth = 512
-TransformerLM.feedforward_depth = 2048
 TransformerLM.max_len = 2048
 TransformerLM.mode = 'train'
-TransformerLM.num_heads = 8
-TransformerLM.num_layers = 6
+TransformerLM.n_heads = 8
+TransformerLM.n_layers = 6
 TransformerLM.vocab_size = 32000
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
index f2a9a1e59..abcdfe7b6 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
@@ -43,11 +43,11 @@ train.train_steps = 500000
 
 # Parameters for Transformer:
 # ==============================================================================
+Transformer.d_feature= 512
+Transformer.d_feedforward = 2048
 Transformer.dropout = 0.1
-Transformer.feature_depth = 512
-Transformer.feedforward_depth = 2048
 Transformer.max_len = 2048
 Transformer.mode = 'train'
-Transformer.num_heads = 8
-Transformer.num_layers = 6
+Transformer.n_heads = 8
+Transformer.n_layers = 6
 Transformer.vocab_size = 33300
diff --git a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
index 500931580..ff3b2aa43 100644
--- a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
+++ b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
@@ -32,9 +32,9 @@ shuffle_and_batch_data.preprocess_fun=@trax.inputs.cifar10_no_augmentation_prepr
 
 # Parameters for WideResnet:
 # ==============================================================================
-WideResnet.num_blocks = 3
-WideResnet.hidden_size = 64
-WideResnet.num_output_classes = 10
+WideResnet.d_hidden = 64
+WideResnet.n_blocks = 3
+WideResnet.n_output_classes = 10
 
 # Parameters for train:
 # ==============================================================================
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 022a0ff6a..30aef4877 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -63,14 +63,14 @@ def EncoderDecoderMask(x, **unused_kwargs):
 def _positional_encoding_new_params(input_shape, rng, max_len=2048):  # pylint: disable=invalid-name
   """Helper: create positional encoding parameters."""
   del rng
-  feature_depth = input_shape[-1]
-  pe = onp.zeros((max_len, feature_depth), dtype=onp.float32)
+  d_feature = input_shape[-1]
+  pe = onp.zeros((max_len, d_feature), dtype=onp.float32)
   position = onp.arange(0, max_len)[:, onp.newaxis]
   div_term = onp.exp(
-      onp.arange(0, feature_depth, 2) * -(onp.log(10000.0) / feature_depth))
+      onp.arange(0, d_feature, 2) * -(onp.log(10000.0) / d_feature))
   pe[:, 0::2] = onp.sin(position * div_term)
   pe[:, 1::2] = onp.cos(position * div_term)
-  pe = pe[onp.newaxis, :, :]  # [1, max_len, feature_depth]
+  pe = pe[onp.newaxis, :, :]  # [1, max_len, d_feature]
   return np.array(pe)  # These are trainable parameters, initialized as above.
 
 
@@ -122,17 +122,17 @@ def PureDotProductAttention(dropout=0.0, mode='train'):
   Returns:
     Pure single-headed attention layer. (No Dense transforms on input.)
   """
-  def init_fun(_, input_shapes):  # pylint: disable=invalid-name
+  def init_fn(_, input_shapes):  # pylint: disable=invalid-name
     q_shape, _, v_shape, _ = input_shapes
     output_shape = q_shape[:-1] + (v_shape[-1],)
     return output_shape, ()
-  def apply_fun(params, inputs, **kwargs):  # pylint: disable=invalid-name
+  def apply_fn(params, inputs, **kwargs):  # pylint: disable=invalid-name
     del params
     q, k, v, mask = inputs
     rng = kwargs.get('rng', None)
     return DotProductAttention(q, k, v, mask,
                                dropout=dropout, mode=mode, rng=rng)
-  return init_fun, apply_fun
+  return init_fn, apply_fn
 
 
 def _multihead_attention_output_shape(  # pylint: disable=invalid-name
@@ -147,14 +147,14 @@ def _multihead_attention_output_shape(  # pylint: disable=invalid-name
 
 @base.layer(output_shape=_multihead_attention_output_shape,
             stack_items_to_pass=4)
-def PureMultiHeadedAttention(x, params, num_heads=8, dropout=0.0,
-                             mode='train', **kwargs):
+def PureMultiHeadedAttention(x, params, n_heads=8, dropout=0.0, mode='train',
+                             **kwargs):
   """Pure transformer-style multi-headed attention.
 
   Args:
     x: inputs ((q, k, v), mask)
     params: parameters (none)
-    num_heads: int: number of attention heads
+    n_heads: int: number of attention heads
     dropout: float: dropout rate
     mode: str: 'train' or 'eval'
     **kwargs: other arguments including the rng
@@ -165,18 +165,18 @@ def PureMultiHeadedAttention(x, params, num_heads=8, dropout=0.0,
   del params
   rng = kwargs.get('rng', None)
   q, k, v, mask = x
-  feature_depth = q.shape[-1]
-  assert feature_depth % num_heads == 0
-  head_depth = feature_depth // num_heads
+  d_feature = q.shape[-1]
+  assert d_feature % n_heads == 0
+  d_head = d_feature // n_heads
   nbatch = np.shape(q)[0]
-  # nbatch, seqlen, feature_depth --> nbatch, num_heads, seqlen, head_depth
+  # nbatch, seqlen, d_feature --> nbatch, n_heads, seqlen, d_head
   def SplitHeads(x):
     return np.transpose(
-        np.reshape(x, (nbatch, -1, num_heads, head_depth)), (0, 2, 1, 3))
-  # nbatch, num_heads, seqlen, head_depth --> nbatch, seqlen, feature_depth
+        np.reshape(x, (nbatch, -1, n_heads, d_head)), (0, 2, 1, 3))
+  # nbatch, n_heads, seqlen, d_head --> nbatch, seqlen, d_feature
   def JoinHeads(x):  # pylint: disable=invalid-name
     return np.reshape(
-        np.transpose(x, (0, 2, 1, 3)), (nbatch, -1, num_heads*head_depth))
+        np.transpose(x, (0, 2, 1, 3)), (nbatch, -1, n_heads * d_head))
   # Split heads, dot-product attention, rejoin heads.
   res = JoinHeads(
       DotProductAttention(
@@ -185,56 +185,53 @@ def JoinHeads(x):  # pylint: disable=invalid-name
   return res, mask  # Keep the mask.
 
 
-def MultiHeadedAttentionQKV(
-    feature_depth, num_heads=8, dropout=0.0, mode='train'):
+def MultiHeadedAttentionQKV(d_feature, n_heads=8, dropout=0.0, mode='train'):
   """Transformer-style multi-headed attention.
 
   Accepts inputs of the form (q, k, v), mask.
 
   Args:
-    feature_depth: int:  depth of embedding
-    num_heads: int: number of attention heads
+    d_feature: int:  dimensionality of feature embedding
+    n_heads: int: number of attention heads
     dropout: float: dropout rate
     mode: str: 'train' or 'eval'
 
   Returns:
     Multi-headed self-attention result and the mask.
   """
-  return combinators.Serial(
+  return [
       combinators.Parallel(
-          core.Dense(feature_depth),
-          core.Dense(feature_depth),
-          core.Dense(feature_depth),
-          combinators.NoOp()
+          core.Dense(d_feature),
+          core.Dense(d_feature),
+          core.Dense(d_feature),
       ),
       PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
-          feature_depth=feature_depth, num_heads=num_heads,
-          dropout=dropout, mode=mode),
-      core.Dense(feature_depth),
-  )
+          d_feature=d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+      core.Dense(d_feature),
+  ]
 
 
 def MultiHeadedAttention(
-    feature_depth, num_heads=8, dropout=0.0, mode='train'):
+    d_feature, n_heads=8, dropout=0.0, mode='train'):
   """Transformer-style multi-headed attention.
 
   Accepts inputs of the form (x, mask) and constructs (q, k, v) from x.
 
   Args:
-    feature_depth: int:  depth of embedding
-    num_heads: int: number of attention heads
+    d_feature: int:  dimensionality of feature embedding
+    n_heads: int: number of attention heads
     dropout: float: dropout rate
     mode: str: 'train' or 'eval'
 
   Returns:
     Multi-headed self-attention layer.
   """
-  return combinators.Serial(
+  return [
       combinators.Dup(),
       combinators.Dup(),
       MultiHeadedAttentionQKV(  # pylint: disable=no-value-for-parameter
-          feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),
-  )
+          d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+  ]
 
 
 @base.layer()
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index de6a225ce..978cee0d1 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -43,7 +43,7 @@ def call(self, x, params=(), **kwargs):
     """Call this layer in input x using the given parameters."""
     raise NotImplementedError
 
-  def output_shape_fun(self, input_shape):
+  def output_shape_fn(self, input_shape):
     """The shape of the output of this layer given the shape of the input.
 
     Note that all arguments and return values can be tuples or dictionaries
@@ -84,7 +84,7 @@ def output_shape(self, input_shape):
       is_list = isinstance(input_shape, (list, tuple))
       is_list = is_list and isinstance(input_shape[0], (list, tuple))
       n = self.stack_items_to_pass() if is_list else 0
-      return _apply_to_first_n(self.output_shape_fun, input_shape, n)
+      return _apply_to_first_n(self.output_shape_fn, input_shape, n)
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback()
       raise LayerError(name, 'output_shape', self._caller, input_shape, trace)
@@ -284,23 +284,23 @@ def layer(output_shape=None, new_parameters=None, stack_items_to_pass=1):
   def layer_decorator(call):
     """Decorating the call function."""
 
-    def stack_items_to_pass_fun(self):
+    def stack_items_to_pass_fn(self):
       del self
       return stack_items_to_pass
 
-    def output_shape_fun(self, input_shape):
+    def output_shape_fn(self, input_shape):
       if output_shape is None:
         return input_shape
       kwargs = self._init_kwargs  # pylint: disable=protected-access
       return output_shape(input_shape, **kwargs)
 
-    def new_parameters_fun(self, input_shape, rng):
+    def new_parameters_fn(self, input_shape, rng):
       if new_parameters is None:
         return ()
       kwargs = self._init_kwargs  # pylint: disable=protected-access
       return new_parameters(input_shape, rng, **kwargs)
 
-    def call_fun(self, x, params=(), **kwargs):
+    def call_fn(self, x, params=(), **kwargs):
       """The call function of the created class, derived from call."""
       # Merge on-call kwargs with class-kwargs.
       call_kwargs = kwargs.copy()
@@ -309,18 +309,18 @@ def call_fun(self, x, params=(), **kwargs):
       return call(x, params=params, **call_kwargs)
 
     # Set doc for python help.
-    call_fun.__doc__ = call.__doc__
+    call_fn.__doc__ = call.__doc__
     if output_shape is None:
-      output_shape_fun.__doc__ = output_shape.__doc__
+      output_shape_fn.__doc__ = output_shape.__doc__
     if new_parameters is None:
-      new_parameters_fun.__doc__ = new_parameters.__doc__
+      new_parameters_fn.__doc__ = new_parameters.__doc__
 
     # Create the class.
     cls = type(call.__name__, (Layer,),
-               {'call': call_fun,
-                'output_shape_fun': output_shape_fun,
-                'new_parameters': new_parameters_fun,
-                'stack_items_to_pass': stack_items_to_pass_fun})
+               {'call': call_fn,
+                'output_shape_fn': output_shape_fn,
+                'new_parameters': new_parameters_fn,
+                'stack_items_to_pass': stack_items_to_pass_fn})
 
     return cls
   return layer_decorator
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index de62e8e63..6e45bd55c 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -27,13 +27,35 @@
 from tensor2tensor.trax.layers import base
 
 
+def _DeepFlatten(xs):  # pylint: disable=invalid-name
+  for x in xs:
+    if isinstance(x, (list, tuple)):
+      for y in _DeepFlatten(x):
+        yield y
+    else:
+      yield x
+
+
+def _EnsureSublayers(layers):
+  # TODO(jonni): Implement for dict if dicts remain important.
+  if isinstance(layers, dict):
+    return layers
+  sublayers_not_lists = []
+  for layer in layers:
+    sublayers_not_lists.append(
+        Serial(layer) if isinstance(layer, list) else layer)
+  return sublayers_not_lists
+
+
 class Serial(base.Layer):
   """Layer composing a number of sub-layers in a serial way.."""
 
   def __init__(self, *layers):
     super(Serial, self).__init__()
-    self._nlayers = len(layers)
+    layers = list(_DeepFlatten(layers))
+    # TODO(jonni): Consider flattening (unpacking) also embedded Serial layers.
     self._layers = layers
+    self._nlayers = len(layers)
 
   def call(self, x, params=(), **kwargs):
     rng = kwargs.pop('rng', None)
@@ -44,7 +66,7 @@ def call(self, x, params=(), **kwargs):
       x = layer(x, p, rng=rng, **kwargs)
     return x
 
-  def output_shape_fun(self, input_shape):
+  def output_shape_fn(self, input_shape):
     cur_shape = input_shape
     for layer in self._layers:
       cur_shape = layer.output_shape(cur_shape)
@@ -150,13 +172,7 @@ def _flatten_shape(x_shape):  # pylint: disable=invalid-name
 @base.layer(output_shape=_flatten_shape, stack_items_to_pass=0)
 def Flatten(xs, **unused_kwargs):
   """Flatten lists."""
-  res = []
-  for x in xs:
-    if isinstance(x, (list, tuple)):
-      res.extend(list(x))
-    else:
-      res.append(x)
-  return tuple(res)
+  return tuple(_DeepFlatten(xs))
 
 
 # Re-ordering layer.
@@ -223,7 +239,7 @@ def call(self, x, params=(), **kwargs):
       return x
     return base.nested_map(self._output, lambda i: self._map(x, i))
 
-  def output_shape_fun(self, input_shape):
+  def output_shape_fn(self, input_shape):
     if self._output is None:
       return input_shape
     return base.nested_map(self._output, lambda i: self._map(input_shape, i))
@@ -253,6 +269,7 @@ def __init__(self, *layers, **kwlayers):
     if layers and kwlayers:
       raise ValueError('Cannot specify a Branch with both a list and dict.')
     layers = layers or kwlayers
+    layers = _EnsureSublayers(layers)
     self._nlayers = len(layers)
     self._layers = layers
 
@@ -276,7 +293,7 @@ def call(self, x, params=(), **kwargs):
       counter += 1
     return result
 
-  def output_shape_fun(self, input_shape):
+  def output_shape_fn(self, input_shape):
     output_shapes = []
     # If the argument layers are a sequence, apply each to calculate shape.
     if not isinstance(self._layers, dict):
@@ -429,6 +446,7 @@ def __init__(self, *layers, **kwlayers):
     if layers and kwlayers:
       raise ValueError('Cannot specify a Parallel with both a list and dict.')
     layers = layers or kwlayers
+    layers = _EnsureSublayers(layers)
     self._nlayers = len(layers)
     self._layers = layers
 
@@ -460,7 +478,7 @@ def call(self, inputs, params=(), **kwargs):
         result[k] = inputs[k]
     return result
 
-  def output_shape_fun(self, input_shape):
+  def output_shape_fn(self, input_shape):
     output_shapes = []
     # If the argument layers are a sequence, apply each to calculate shape.
     if not isinstance(self._layers, dict):
@@ -493,20 +511,11 @@ def new_parameters(self, input_shape, rng):
 def Residual(*layers, **kwargs):
   """Constructs a residual version of layers, summing input to layers output."""
   shortcut = kwargs.get('shortcut', _Top())  # pylint: disable=no-value-for-parameter
-  if len(layers) > 1:
-    return Serial(
-        Branch(shortcut, Serial(*layers)),
-        Flatten(),  # pylint: disable=no-value-for-parameter
-        Add()  # pylint: disable=no-value-for-parameter
-    )
-  elif len(layers) == 1:
-    return Serial(
-        Branch(shortcut, layers[0]),
-        Flatten(),  # pylint: disable=no-value-for-parameter
-        Add()  # pylint: disable=no-value-for-parameter
-    )
-  else:
-    raise ValueError('Empty residual combinator.')
+  return [
+      Branch(shortcut, Serial(layers)),  # Use Serial here to flatten layers.
+      Flatten(),  # pylint: disable=no-value-for-parameter
+      Add(),  # pylint: disable=no-value-for-parameter
+  ]
 
 
 class Map(base.Layer):
@@ -540,7 +549,7 @@ def call(self, inputs, params=(), **kwargs):
       return result
     return tuple(result)
 
-  def output_shape_fun(self, input_shapes):
+  def output_shape_fn(self, input_shapes):
     return tuple([self._layer.output_shape(shape) for shape in input_shapes])
 
   def new_parameters(self, input_shape, rng):
@@ -558,22 +567,22 @@ class Rebatch(base.Layer):
 
   Args:
     layer: subclass of base.Layer, a layer to apply to the input.
-    num_batch_dims: int, the number of leading dimensions to consider as batch.
+    n_batch_dims: int, the number of leading dimensions to consider as batch.
 
   Returns:
     A new layer that will reshape the input into a virtual batch, apply the
     layer and unbatch the virtual batch.
   """
 
-  def __init__(self, layer, num_batch_dims=1):
+  def __init__(self, layer, n_batch_dims=1):
     super(Rebatch, self).__init__()
     self._layer = layer
-    self._num_batch_dims = num_batch_dims
+    self._n_batch_dims = n_batch_dims
 
   def _modify_shape(self, input_shape):
     input_shape = tuple(input_shape)
-    batch_dims, non_batch_dims = (input_shape[:self._num_batch_dims],
-                                  input_shape[self._num_batch_dims:])
+    batch_dims, non_batch_dims = (input_shape[:self._n_batch_dims],
+                                  input_shape[self._n_batch_dims:])
     new_batch_dim = six.moves.reduce(operator.mul, batch_dims)
     return (new_batch_dim,) + non_batch_dims, batch_dims
 
@@ -596,7 +605,7 @@ def call(self, inp, params=(), **kwargs):
     out = self._layer(inp, params=params, **kwargs)
     return self._unmodify(out, batch_dims)
 
-  def output_shape_fun(self, input_shape):
+  def output_shape_fn(self, input_shape):
     modified_shape, batch_dims = self._modify_shape(input_shape)
     out = self._layer.output_shape(modified_shape)
     return self._unmodify_shape(out, batch_dims)
diff --git a/tensor2tensor/trax/layers/combinators_test.py b/tensor2tensor/trax/layers/combinators_test.py
index e0384a73e..ef91d08a8 100644
--- a/tensor2tensor/trax/layers/combinators_test.py
+++ b/tensor2tensor/trax/layers/combinators_test.py
@@ -78,15 +78,13 @@ def test_rebatch(self):
 
     input_shape = (29, 5, 5, 20)
     result_shape = base.check_shape_agreement(
-        combinators.Rebatch(
-            convolution.Conv(30, (3, 3)), num_batch_dims=1),
+        combinators.Rebatch(convolution.Conv(30, (3, 3)), n_batch_dims=1),
         input_shape)
     self.assertEqual(result_shape, (29, 3, 3, 30))
 
     input_shape = (19, 29, 5, 5, 20)
     result_shape = base.check_shape_agreement(
-        combinators.Rebatch(
-            convolution.Conv(30, (3, 3)), num_batch_dims=2),
+        combinators.Rebatch(convolution.Conv(30, (3, 3)), n_batch_dims=2),
         input_shape)
     self.assertEqual(result_shape, (19, 29, 3, 3, 30))
 
diff --git a/tensor2tensor/trax/layers/convolution.py b/tensor2tensor/trax/layers/convolution.py
index 769a93f6f..83267be09 100644
--- a/tensor2tensor/trax/layers/convolution.py
+++ b/tensor2tensor/trax/layers/convolution.py
@@ -140,7 +140,7 @@ def _conv_general_shape_tuple(self, lhs_shape, rhs_shape, window_strides,
         lhs_trans, rhs_trans, window_strides, padding)
     return tuple(onp.take(out_trans, onp.argsort(out_perm)))
 
-  def output_shape_fun(self, input_shape):
+  def output_shape_fn(self, input_shape):
     kernel_shape = self._kernel_shape(input_shape)
     return self._conv_general_shape_tuple(
         input_shape, kernel_shape,
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 14fa7184f..3467d0b9b 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -97,7 +97,7 @@ def call(self, x, params, **kwargs):
     w, b = params
     return np.dot(x, w) + b
 
-  def output_shape_fun(self, input_shape):
+  def output_shape_fn(self, input_shape):
     return tuple(input_shape[:-1]) + (self._units,)
 
   def new_parameters(self, input_shape, rng):
@@ -110,10 +110,10 @@ def new_parameters(self, input_shape, rng):
 class Embedding(base.Layer):
   """Layer constructor function for an embedding layer."""
 
-  def __init__(self, feature_depth, vocab_size,
+  def __init__(self, d_feature, vocab_size,
                kernel_initializer=init.GlorotUniformInitializer()):
     super(Embedding, self).__init__()
-    self._feature_depth = feature_depth
+    self._d_feature = d_feature  # feature dimensionality
     self._vocab_size = vocab_size
     self._kernel_initializer = kernel_initializer
 
@@ -124,12 +124,12 @@ def call(self, x, params, **kwargs):
     del kwargs
     return np.take(params, x, axis=0)
 
-  def output_shape_fun(self, input_shape):
-    return tuple(input_shape) + (self._feature_depth,)
+  def output_shape_fn(self, input_shape):
+    return tuple(input_shape) + (self._d_feature,)
 
   def new_parameters(self, input_shape, rng):
     return self._kernel_initializer(
-        (self._vocab_size, self._feature_depth), rng)
+        (self._vocab_size, self._d_feature), rng)
 
 
 # Flatten.
@@ -154,7 +154,7 @@ def Dropout(x, params, rate=0.0, mode='train', rng=None, **kwargs):
   """Layer construction function for a dropout layer with given rate."""
   del params, kwargs
   if rng is None:
-    msg = ('Dropout layer requires apply_fun to be called with a rng keyword '
+    msg = ('Dropout layer requires apply_fn to be called with a rng keyword '
            'argument. That is, instead of `Dropout(params, inputs)`, call '
            'it like `Dropout(params, inputs, rng=key)`.')
     raise ValueError(msg)
diff --git a/tensor2tensor/trax/layers/rnn.py b/tensor2tensor/trax/layers/rnn.py
index 9e788c4b3..8c51562b6 100644
--- a/tensor2tensor/trax/layers/rnn.py
+++ b/tensor2tensor/trax/layers/rnn.py
@@ -19,7 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.trax.layers import combinators
+from tensor2tensor.trax.layers import combinators as cb
 from tensor2tensor.trax.layers import convolution
 from tensor2tensor.trax.layers import core
 
@@ -38,7 +38,7 @@ def GRUCell(units):
   """
   return GeneralGRUCell(
       candidate_transform=lambda: core.Dense(units=units),
-      memory_transform=combinators.NoOp,
+      memory_transform=cb.NoOp,
       gate_nonlinearity=core.Sigmoid,
       candidate_nonlinearity=core.Tanh)
 
@@ -62,13 +62,13 @@ def BuildConv():
 
   return GeneralGRUCell(
       candidate_transform=BuildConv,
-      memory_transform=combinators.NoOp,
+      memory_transform=cb.NoOp,
       gate_nonlinearity=core.Sigmoid,
       candidate_nonlinearity=core.Tanh)
 
 
 def GeneralGRUCell(candidate_transform,
-                   memory_transform=combinators.NoOp,
+                   memory_transform=cb.NoOp,
                    gate_nonlinearity=core.Sigmoid,
                    candidate_nonlinearity=core.Tanh,
                    dropout_rate_c=0.1,
@@ -100,43 +100,26 @@ def GeneralGRUCell(candidate_transform,
   Returns:
     A model representing a GRU cell with specified transforms.
   """
-  return combinators.Serial(
-      combinators.Branch(
-          # s_{t-1} branch - optionally transform
-          # Typically is an identity.
-          memory_transform(),
-
-          # u_t (Update gate) branch
-          combinators.Serial(
-              candidate_transform(),
-              # Want bias to start out positive before sigmoids.
-              core.AddConstant(constant=sigmoid_bias),
-              gate_nonlinearity()
-          ),
-
-          # c_t (Candidate) branch
-          combinators.Serial(
-              combinators.Branch(
-                  combinators.NoOp(),
-                  # r_t (Reset) Branch
-                  combinators.Serial(
-                      candidate_transform(),
-                      # Want bias to start out positive before sigmoids.
-                      core.AddConstant(constant=sigmoid_bias),
-                      gate_nonlinearity()
-                  )
-              ),
-              ## Gate S{t-1} with sigmoid(candidate_transform(S{t-1}))
-              combinators.Multiply(),
-
-              # Final projection + tanh to get Ct
-              candidate_transform(),
-              candidate_nonlinearity(),  # Candidate gate
-
-              # Only apply dropout on the C gate.
-              # Paper reports that 0.1 is a good default.
-              core.Dropout(rate=dropout_rate_c)
-          ),
-      ),
-      # Gate memory and candidate
-      combinators.Gate())
+  gate_block = [  # u_t
+      candidate_transform(),
+      core.AddConstant(constant=sigmoid_bias),
+      gate_nonlinearity(),
+  ]
+  reset_block = [  # r_t
+      candidate_transform(),
+      core.AddConstant(constant=sigmoid_bias),  # Want bias to start positive.
+      gate_nonlinearity(),
+  ]
+  candidate_block = [
+      cb.Branch([], reset_block),
+      cb.Multiply(),  # Gate S{t-1} with sigmoid(candidate_transform(S{t-1}))
+      candidate_transform(),  # Final projection + tanh to get Ct
+      candidate_nonlinearity(),  # Candidate gate
+
+      # Only apply dropout on the C gate. Paper reports 0.1 as a good default.
+      core.Dropout(rate=dropout_rate_c)
+  ]
+  return cb.Serial([
+      cb.Branch(memory_transform(), gate_block, candidate_block),
+      cb.Gate(),
+  ])
diff --git a/tensor2tensor/trax/models/atari_cnn.py b/tensor2tensor/trax/models/atari_cnn.py
index 6f3beb29b..7b102c1de 100644
--- a/tensor2tensor/trax/models/atari_cnn.py
+++ b/tensor2tensor/trax/models/atari_cnn.py
@@ -24,18 +24,15 @@
 
 def AtariCnn(hidden_sizes=(32, 32), output_size=128):
   # Input's shape = (B, T, H, W, C)
-  return tl.Serial(
+  return tl.Serial([
       tl.Div(divisor=255.0),
       # Have 4 copies of the input, each one shifted to the right by one.
-      tl.Branch(tl.NoOp(), tl.ShiftRight(),
-                tl.Serial(
-                    tl.ShiftRight(),
-                    tl.ShiftRight(),
-                ), tl.Serial(
-                    tl.ShiftRight(),
-                    tl.ShiftRight(),
-                    tl.ShiftRight(),
-                )),
+      tl.Branch(
+          [],
+          [tl.ShiftRight()],
+          [tl.ShiftRight(), tl.ShiftRight()],
+          [tl.ShiftRight(), tl.ShiftRight(), tl.ShiftRight()]
+      ),
       # Concatenated on the last axis.
       tl.Concatenate(axis=-1),  # (B, T, H, W, 4C)
       tl.Rebatch(tl.Conv(hidden_sizes[0], (5, 5), (2, 2), 'SAME'), 2),
@@ -46,4 +43,4 @@ def AtariCnn(hidden_sizes=(32, 32), output_size=128):
       tl.Dense(output_size),
       tl.Relu(),
       # Eventually this is shaped (B, T, output_size)
-  )
+  ])
diff --git a/tensor2tensor/trax/models/mlp.py b/tensor2tensor/trax/models/mlp.py
index fa248e4c3..9283fad4e 100644
--- a/tensor2tensor/trax/models/mlp.py
+++ b/tensor2tensor/trax/models/mlp.py
@@ -22,15 +22,17 @@
 from tensor2tensor.trax import layers as tl
 
 
-def MLP(num_hidden_layers=2,
-        hidden_size=512,
+def MLP(n_hidden_layers=2,
+        d_hidden=512,
         activation_fn=tl.Relu,
-        num_output_classes=10,
+        n_output_classes=10,
         mode="train"):
   """Multi-layer feed-forward neural network with non-linear activations."""
   del mode
-  cur_layers = [tl.Flatten()]
-  for _ in range(num_hidden_layers):
-    cur_layers += [tl.Dense(hidden_size), activation_fn()]
-  cur_layers += [tl.Dense(num_output_classes), tl.LogSoftmax()]
-  return tl.Serial(*cur_layers)
+
+  return [
+      tl.Flatten(),
+      [[tl.Dense(d_hidden), activation_fn()] for _ in range(n_hidden_layers)],
+      tl.Dense(n_output_classes),
+      tl.LogSoftmax(),
+  ]
diff --git a/tensor2tensor/trax/models/neural_gpu.py b/tensor2tensor/trax/models/neural_gpu.py
index c578053df..e76e91c2c 100644
--- a/tensor2tensor/trax/models/neural_gpu.py
+++ b/tensor2tensor/trax/models/neural_gpu.py
@@ -59,23 +59,21 @@ def BuildConv():
       candidate_nonlinearity=tl.HardTanh)
 
 
-def NeuralGPU(feature_depth=96, steps=16, vocab_size=2):
+def NeuralGPU(d_feature=96, steps=16, vocab_size=2):
   """Implementation of Neural GPU: https://arxiv.org/abs/1702.08727.
 
   Args:
-    feature_depth: Number of memory channels
+    d_feature: Number of memory channels (dimensionality of feature embedding).
     steps: Number of times depthwise recurrence steps.
     vocab_size: Vocabulary size.
 
   Returns:
     A NeuralGPU Stax model.
   """
-  xs = []
-  xs.append(
-      tl.Embedding(feature_depth=feature_depth, vocab_size=vocab_size))
-  core = ConvDiagonalGRU(units=feature_depth)
-  xs.extend([core] * steps)
-  xs.append(tl.Dense(vocab_size))
-  xs.append(tl.LogSoftmax())
-
-  return tl.Serial(*xs)
+  core = ConvDiagonalGRU(units=d_feature)
+  return tl.Serial([
+      tl.Embedding(d_feature=d_feature, vocab_size=vocab_size),
+      [core] * steps,
+      tl.Dense(vocab_size),
+      tl.LogSoftmax(),
+  ])
diff --git a/tensor2tensor/trax/models/neural_gpu_test.py b/tensor2tensor/trax/models/neural_gpu_test.py
index ab3cbecfa..8583376d3 100644
--- a/tensor2tensor/trax/models/neural_gpu_test.py
+++ b/tensor2tensor/trax/models/neural_gpu_test.py
@@ -29,8 +29,7 @@ class NeuralGPUTest(absltest.TestCase):
   def test_ngpu(self):
     vocab_size = 2
     input_shape = [3, 5, 7]
-    model = neural_gpu.NeuralGPU(
-        feature_depth=30, steps=4, vocab_size=vocab_size)
+    model = neural_gpu.NeuralGPU(d_feature=30, steps=4, vocab_size=vocab_size)
     final_shape = base.check_shape_agreement(
         model, tuple(input_shape), integer_inputs=True)
     self.assertEqual(tuple(input_shape + [vocab_size]), final_shape)
diff --git a/tensor2tensor/trax/models/research/chunked_transformer.py b/tensor2tensor/trax/models/research/chunked_transformer.py
index aa3a3ee5c..2288445a0 100644
--- a/tensor2tensor/trax/models/research/chunked_transformer.py
+++ b/tensor2tensor/trax/models/research/chunked_transformer.py
@@ -31,14 +31,14 @@ def _chunked_positional_encoding_new_params(input_shape, rng, max_len=2048):  #
   # Check if we are operating on chunked inputs by checking if the first
   # shape is a list/tuple of shapes (otherwise it's an int or numpy array).
   is_chunked = isinstance(input_shape[0], (list, tuple))
-  feature_depth = input_shape[0][-1] if is_chunked else input_shape[-1]
-  pe = onp.zeros((max_len, feature_depth), dtype=onp.float32)
+  d_feature = input_shape[0][-1] if is_chunked else input_shape[-1]
+  pe = onp.zeros((max_len, d_feature), dtype=onp.float32)
   position = onp.arange(0, max_len)[:, onp.newaxis]
   div_term = onp.exp(
-      onp.arange(0, feature_depth, 2) * -(onp.log(10000.0) / feature_depth))
+      onp.arange(0, d_feature, 2) * -(onp.log(10000.0) / d_feature))
   pe[:, 0::2] = onp.sin(position * div_term)
   pe[:, 1::2] = onp.cos(position * div_term)
-  pe = pe[onp.newaxis, :, :]  # [1, max_len, feature_depth]
+  pe = pe[onp.newaxis, :, :]  # [1, max_len, d_feature]
   return np.array(pe)  # These are trainable parameters, initialized as above.
 
 
@@ -126,14 +126,14 @@ def ChunkedAttentionSelector(x, params, selector=None, **kwargs):
 
 
 def ChunkedCausalMultiHeadedAttention(
-    feature_depth, num_heads=8, dropout=0.0, chunk_selector=None, mode='train'):
+    d_feature, n_heads=8, dropout=0.0, chunk_selector=None, mode='train'):
   """Transformer-style causal multi-headed attention operating on chunks.
 
   Accepts inputs that are a list of chunks and applies causal attention.
 
   Args:
-    feature_depth: int:  depth of embedding
-    num_heads: int: number of attention heads
+    d_feature: int:  depth of embedding
+    n_heads: int: number of attention heads
     dropout: float: dropout rate
     chunk_selector: a function from chunk number to list of chunks to attend.
     mode: str: 'train' or 'eval'
@@ -149,9 +149,9 @@ def ChunkedCausalMultiHeadedAttention(
       ),
       tl.Parallel(
           tl.Parallel(
-              tl.Dense(feature_depth),
-              tl.Dense(feature_depth),
-              tl.Dense(feature_depth),
+              tl.Dense(d_feature),
+              tl.Dense(d_feature),
+              tl.Dense(d_feature),
           ),
           tl.NoOp()
       )
@@ -160,10 +160,10 @@ def ChunkedCausalMultiHeadedAttention(
       tl.Map(prepare_attention_input),
       ChunkedAttentionSelector(selector=chunk_selector),  # pylint: disable=no-value-for-parameter
       tl.Map(tl.PureMultiHeadedAttention(
-          feature_depth=feature_depth, num_heads=num_heads,
+          d_feature=d_feature, n_heads=n_heads,
           dropout=dropout, mode=mode), check_shapes=False),
       tl.Map(tl.Select(0), check_shapes=False),  # drop masks
-      tl.Map(tl.Dense(feature_depth))
+      tl.Map(tl.Dense(d_feature))
   )
 
 
@@ -176,33 +176,30 @@ def Residual(*layers, **unused_kwargs):
   )
 
 
-def ResidualFeedForward(feature_depth,
-                        feedforward_depth,
-                        dropout,
-                        mode):
+def ResidualFeedForward(d_feature, d_feedforward, dropout, mode):
   """Residual feed-forward layer with normalization at start."""
   return Residual(
       tl.LayerNorm(),
-      tl.Dense(feedforward_depth),
+      tl.Dense(d_feedforward),
       tl.Relu(),
       tl.Dropout(rate=dropout, mode=mode),
-      tl.Dense(feature_depth),
+      tl.Dense(d_feature),
       tl.Dropout(rate=dropout, mode=mode)
   )
 
 
-def ChunkedDecoderLayer(feature_depth,
-                        feedforward_depth,
-                        num_heads,
+def ChunkedDecoderLayer(d_feature,
+                        d_feedforward,
+                        n_heads,
                         dropout,
                         chunk_selector,
                         mode):
   """Transformer decoder layer operating on chunks.
 
   Args:
-    feature_depth: int:  depth of embedding
-    feedforward_depth: int: depth of feed-forward layer
-    num_heads: int: number of attention heads
+    d_feature: int:  depth of embedding
+    d_feedforward: int: depth of feed-forward layer
+    n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
     chunk_selector: a function from chunk number to list of chunks to attend.
     mode: str: 'train' or 'eval'
@@ -214,20 +211,20 @@ def ChunkedDecoderLayer(feature_depth,
       Residual(  # Self-attention block.
           tl.Map(tl.LayerNorm()),
           ChunkedCausalMultiHeadedAttention(
-              feature_depth, num_heads=num_heads, dropout=dropout,
+              d_feature, n_heads=n_heads, dropout=dropout,
               chunk_selector=chunk_selector, mode=mode),
           tl.Map(tl.Dropout(rate=dropout, mode=mode)),
       ),
       tl.Map(ResidualFeedForward(
-          feature_depth, feedforward_depth, dropout, mode=mode))
+          d_feature, d_feedforward, dropout, mode=mode))
   )
 
 
 def ChunkedTransformerLM(vocab_size,
-                         feature_depth=512,
-                         feedforward_depth=2048,
-                         num_layers=6,
-                         num_heads=8,
+                         d_feature=512,
+                         d_feedforward=2048,
+                         n_layers=6,
+                         n_heads=8,
                          dropout=0.1,
                          chunk_selector=None,
                          max_len=2048,
@@ -245,10 +242,10 @@ def ChunkedTransformerLM(vocab_size,
 
   Args:
     vocab_size: int: vocab size
-    feature_depth: int:  depth of embedding
-    feedforward_depth: int: depth of feed-forward layer
-    num_layers: int: number of encoder/decoder layers
-    num_heads: int: number of attention heads
+    d_feature: int:  depth of embedding
+    d_feedforward: int: depth of feed-forward layer
+    n_layers: int: number of encoder/decoder layers
+    n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
     chunk_selector: a function from chunk number to list of chunks to attend
       (if None, attends to the previous chunks which is equivalent to setting
@@ -260,13 +257,13 @@ def ChunkedTransformerLM(vocab_size,
   Returns:
     the layer.
   """
-  stack = [ChunkedDecoderLayer(feature_depth, feedforward_depth, num_heads,
+  stack = [ChunkedDecoderLayer(d_feature, d_feedforward, n_heads,
                                dropout, chunk_selector, mode)
-           for _ in range(num_layers)]
+           for _ in range(n_layers)]
   # Below each Map(L) applies the layer L to each chunk independently.
   return tl.Serial(
       tl.ShiftRight(),
-      tl.Map(tl.Embedding(feature_depth, vocab_size)),
+      tl.Map(tl.Embedding(d_feature, vocab_size)),
       tl.Map(tl.Dropout(rate=dropout, mode=mode)),
       ChunkedPositionalEncoding(max_len=max_len),  # pylint: disable=no-value-for-parameter
       tl.Serial(*stack),
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index bf0e6aba1..3221c660f 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -24,9 +24,10 @@
 
 def ConvBlock(kernel_size, filters, strides):
   """ResNet convolutional striding block."""
+  # TODO(jonni): Use good defaults so Resnet50 code is cleaner / less redundant.
   ks = kernel_size
   filters1, filters2, filters3 = filters
-  main = tl.Serial(
+  main = [
       tl.Conv(filters1, (1, 1), strides),
       tl.BatchNorm(),
       tl.Relu(),
@@ -34,23 +35,24 @@ def ConvBlock(kernel_size, filters, strides):
       tl.BatchNorm(),
       tl.Relu(),
       tl.Conv(filters3, (1, 1)),
-      tl.BatchNorm()
-  )
-  shortcut = tl.Serial(
+      tl.BatchNorm(),
+  ]
+  shortcut = [
       tl.Conv(filters3, (1, 1), strides),
-      tl.BatchNorm()
-  )
-  return tl.Serial(
+      tl.BatchNorm(),
+  ]
+  return [
       tl.Residual(main, shortcut=shortcut),
-      tl.Relu()
-  )
+      tl.Relu(),
+  ]
 
 
 def IdentityBlock(kernel_size, filters):
   """ResNet identical size block."""
+  # TODO(jonni): Use good defaults so Resnet50 code is cleaner / less redundant.
   ks = kernel_size
   filters1, filters2, filters3 = filters
-  main = tl.Serial(
+  main = [
       tl.Conv(filters1, (1, 1)),
       tl.BatchNorm(),
       tl.Relu(),
@@ -58,98 +60,100 @@ def IdentityBlock(kernel_size, filters):
       tl.BatchNorm(),
       tl.Relu(),
       tl.Conv(filters3, (1, 1)),
-      tl.BatchNorm()
-  )
-  return tl.Serial(
+      tl.BatchNorm(),
+  ]
+  return [
       tl.Residual(main),
-      tl.Relu()
-  )
+      tl.Relu(),
+  ]
 
 
-def Resnet50(hidden_size=64, num_output_classes=1001, mode='train'):
+def Resnet50(d_hidden=64, n_output_classes=1001, mode='train'):
   """ResNet.
 
   Args:
-    hidden_size: the size of the first hidden layer (multiplied later).
-    num_output_classes: how many classes to distinguish.
-    mode: whether we are training or evaluating or doing inference.
+    d_hidden: Dimensionality of the first hidden layer (multiplied later).
+    n_output_classes: Number of distinct output classes.
+    mode: Whether we are training or evaluating or doing inference.
 
   Returns:
-    The ResNet model with the given layer and output sizes.
+    The list of layers comprising a ResNet model with the given parameters.
   """
   del mode
-  return tl.Serial(
-      tl.Conv(hidden_size, (7, 7), (2, 2), 'SAME'),
-      tl.BatchNorm(), tl.Relu(),
+  return [
+      tl.Conv(d_hidden, (7, 7), (2, 2), 'SAME'),
+      tl.BatchNorm(),
+      tl.Relu(),
       tl.MaxPool(pool_size=(3, 3), strides=(2, 2)),
-      ConvBlock(3, [hidden_size, hidden_size, 4 * hidden_size], (1, 1)),
-      IdentityBlock(3, [hidden_size, hidden_size, 4 * hidden_size]),
-      IdentityBlock(3, [hidden_size, hidden_size, 4 * hidden_size]),
-      ConvBlock(3, [2 * hidden_size, 2 * hidden_size, 8 * hidden_size], (2, 2)),
-      IdentityBlock(3, [2 * hidden_size, 2 * hidden_size, 8 * hidden_size]),
-      IdentityBlock(3, [2 * hidden_size, 2 * hidden_size, 8 * hidden_size]),
-      IdentityBlock(3, [2 * hidden_size, 2 * hidden_size, 8 * hidden_size]),
-      ConvBlock(3, [4 * hidden_size, 4 * hidden_size, 16*hidden_size], (2, 2)),
-      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size, 16 * hidden_size]),
-      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size, 16 * hidden_size]),
-      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size, 16 * hidden_size]),
-      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size, 16 * hidden_size]),
-      IdentityBlock(3, [4 * hidden_size, 4 * hidden_size, 16 * hidden_size]),
-      ConvBlock(3, [8 * hidden_size, 8 * hidden_size, 32*hidden_size], (2, 2)),
-      IdentityBlock(3, [8 * hidden_size, 8 * hidden_size, 32 * hidden_size]),
-      IdentityBlock(3, [8 * hidden_size, 8 * hidden_size, 32 * hidden_size]),
+      ConvBlock(3, [d_hidden, d_hidden, 4 * d_hidden], (1, 1)),
+      IdentityBlock(3, [d_hidden, d_hidden, 4 * d_hidden]),
+      IdentityBlock(3, [d_hidden, d_hidden, 4 * d_hidden]),
+      ConvBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden], (2, 2)),
+      IdentityBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden]),
+      IdentityBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden]),
+      IdentityBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden]),
+      ConvBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden], (2, 2)),
+      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden]),
+      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden]),
+      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden]),
+      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden]),
+      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden]),
+      ConvBlock(3, [8 * d_hidden, 8 * d_hidden, 32 * d_hidden], (2, 2)),
+      IdentityBlock(3, [8 * d_hidden, 8 * d_hidden, 32 * d_hidden]),
+      IdentityBlock(3, [8 * d_hidden, 8 * d_hidden, 32 * d_hidden]),
       tl.AvgPool(pool_size=(7, 7)),
       tl.Flatten(),
-      tl.Dense(num_output_classes),
-      tl.LogSoftmax()
-  )
+      tl.Dense(n_output_classes),
+      tl.LogSoftmax(),
+  ]
 
 
-def WideResnetBlock(channels, strides=(1, 1), channel_mismatch=False):
-  """WideResnet convolutational block."""
-  main = tl.Serial(
+def WideResnetBlock(channels, strides=(1, 1)):
+  """WideResnet convolutional block."""
+  return [
       tl.BatchNorm(),
       tl.Relu(),
       tl.Conv(channels, (3, 3), strides, padding='SAME'),
       tl.BatchNorm(),
       tl.Relu(),
-      tl.Conv(channels, (3, 3), padding='SAME'))
-  shortcut = tl.NoOp() if not channel_mismatch else tl.Conv(
-      channels, (3, 3), strides, padding='SAME')
-  return tl.Residual(main, shortcut=shortcut)
+      tl.Conv(channels, (3, 3), padding='SAME'),
+  ]
 
 
 def WideResnetGroup(n, channels, strides=(1, 1)):
-  blocks = []
-  blocks += [WideResnetBlock(channels, strides, channel_mismatch=True)]
-  for _ in range(n - 1):
-    blocks += [WideResnetBlock(channels, (1, 1))]
-  return tl.Serial(*blocks)
+  shortcut = [
+      tl.Conv(channels, (3, 3), strides, padding='SAME'),
+  ]
+  return [
+      tl.Residual(WideResnetBlock(channels, strides), shortcut=shortcut),
+      tl.Residual([WideResnetBlock(channels, (1, 1))
+                   for _ in range(n - 1)]),
+  ]
 
 
-def WideResnet(num_blocks=3, hidden_size=64, num_output_classes=10,
+def WideResnet(n_blocks=3, d_hidden=64, n_output_classes=10,
                mode='train'):
   """WideResnet from https://arxiv.org/pdf/1605.07146.pdf.
 
   Args:
-    num_blocks: int, number of blocks in a group.
-    hidden_size: the size of the first hidden layer (multiplied later).
-    num_output_classes: int, number of classes to distinguish.
-    mode: is it training or eval.
+    n_blocks: int, number of blocks in a group.
+    d_hidden: Dimensionality of the first hidden layer (multiplied later).
+    n_output_classes: int, number of distinct output classes.
+    mode: Whether we are training or evaluating or doing inference.
 
   Returns:
-    The WideResnet model with given layer and output sizes.
+    The list of layers comprising a WideResnet model with the given parameters.
   """
   del mode
-  return tl.Serial(
-      tl.Conv(hidden_size, (3, 3), padding='SAME'),
-      WideResnetGroup(num_blocks, hidden_size),
-      WideResnetGroup(num_blocks, hidden_size * 2, (2, 2)),
-      WideResnetGroup(num_blocks, hidden_size * 4, (2, 2)),
+  return [
+      tl.Conv(d_hidden, (3, 3), padding='SAME'),
+      WideResnetGroup(n_blocks, d_hidden),
+      WideResnetGroup(n_blocks, d_hidden * 2, (2, 2)),
+      WideResnetGroup(n_blocks, d_hidden * 4, (2, 2)),
       tl.BatchNorm(),
       tl.Relu(),
       tl.AvgPool(pool_size=(8, 8)),
       tl.Flatten(),
-      tl.Dense(num_output_classes),
-      tl.LogSoftmax()
-  )
+      tl.Dense(n_output_classes),
+      tl.LogSoftmax(),
+  ]
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index e0f07dacd..0dd30d241 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -21,60 +21,56 @@
 from tensor2tensor.trax import layers as tl
 
 
-def ResidualFeedForward(feature_depth,
-                        feedforward_depth,
-                        dropout,
-                        mode):
-  """Residual feed-forward layer with normalization at start."""
-  return tl.Residual(
+def FeedForward(d_feature, d_feedforward, dropout, mode):
+  """Feed-forward block with layer normalization at start."""
+  return [
       tl.LayerNorm(),
-      tl.Dense(feedforward_depth),
+      tl.Dense(d_feedforward),
       tl.Relu(),
       tl.Dropout(rate=dropout, mode=mode),
-      tl.Dense(feature_depth),
-      tl.Dropout(rate=dropout, mode=mode)
-  )
+      tl.Dense(d_feature),
+      tl.Dropout(rate=dropout, mode=mode),
+  ]
 
 
-def EncoderLayer(feature_depth,
-                 feedforward_depth,
-                 num_heads,
-                 dropout,
-                 mode):
-  """Transformer encoder layer.
+def EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
+  """Transformer encoder block.
 
   The input to the encoder is a pair (embedded source, mask) where
   the mask is created from the original source to prevent attending
   to the padding part of the input.
 
   Args:
-    feature_depth: int:  depth of embedding
-    feedforward_depth: int: depth of feed-forward layer
-    num_heads: int: number of attention heads
+    d_feature: int:  depth of embedding
+    d_feedforward: int: depth of feed-forward layer
+    n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
 
   Returns:
-    the layer, returning a pair (actiavtions, mask).
+    the layer, returning a pair (activations, mask).
   """
-  return tl.Serial(
-      tl.Residual(  # Attention block here.
-          tl.LayerNorm(),
-          tl.MultiHeadedAttention(feature_depth, num_heads=num_heads,
-                                  dropout=dropout, mode=mode),
-          tl.Dropout(rate=dropout, mode=mode)
-      ),
-      ResidualFeedForward(
-          feature_depth, feedforward_depth, dropout, mode=mode),
-  )
+  attention = [
+      tl.LayerNorm(),
+      tl.MultiHeadedAttention(
+          d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+      tl.Dropout(rate=dropout, mode=mode),
+  ]
+  feed_forward = [
+      FeedForward(d_feature, d_feedforward, dropout, mode=mode),
+  ]
+  return [
+      tl.Residual(attention),
+      tl.Residual(feed_forward),
+  ]
 
 
 def TransformerEncoder(vocab_size,
                        num_classes=10,
-                       feature_depth=512,
-                       feedforward_depth=2048,
-                       num_layers=6,
-                       num_heads=8,
+                       d_feature=512,
+                       d_feedforward=2048,
+                       n_layers=6,
+                       n_heads=8,
                        dropout=0.1,
                        max_len=2048,
                        mode='train'):
@@ -83,10 +79,10 @@ def TransformerEncoder(vocab_size,
   Args:
     vocab_size: int: vocab size
     num_classes: how many classes on output
-    feature_depth: int:  depth of embedding
-    feedforward_depth: int: depth of feed-forward layer
-    num_layers: int: number of encoder/decoder layers
-    num_heads: int: number of attention heads
+    d_feature: int:  depth of embedding
+    d_feedforward: int: depth of feed-forward layer
+    n_layers: int: number of encoder/decoder layers
+    n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
     max_len: int: maximum symbol length for positional encoding
     mode: str: 'train' or 'eval'
@@ -94,59 +90,58 @@ def TransformerEncoder(vocab_size,
   Returns:
     the Transformer encoder layer.
   """
-  input_embedding = tl.Serial(
-      tl.Embedding(feature_depth, vocab_size),
+  positional_embedder = [
+      tl.Embedding(d_feature, vocab_size),
       tl.Dropout(rate=dropout, mode=mode),
-      tl.PositionalEncoding(max_len=max_len)
-  )
-  return tl.Serial(
-      tl.Branch(input_embedding, tl.PaddingMask()),
-      tl.Serial(*[EncoderLayer(feature_depth, feedforward_depth, num_heads,
-                               dropout, mode)
-                  for _ in range(num_layers)]),
-      tl.Select(0),  # Drop the mask.
+      tl.PositionalEncoding(max_len=max_len),
+  ]
+  return [
+      tl.Branch(positional_embedder, tl.PaddingMask()),  # Create mask.
+      [EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode)
+       for _ in range(n_layers)],
+      tl.Select(0),  # Drop mask.
       tl.LayerNorm(),
       tl.Mean(axis=1),  # Average on length.
       tl.Dense(num_classes),
-      tl.LogSoftmax()
-  )
+      tl.LogSoftmax(),
+  ]
 
 
-def DecoderLayer(feature_depth,
-                 feedforward_depth,
-                 num_heads,
-                 dropout,
-                 mode):
+def DecoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
   """Transformer decoder layer.
 
   Args:
-    feature_depth: int:  depth of embedding
-    feedforward_depth: int: depth of feed-forward layer
-    num_heads: int: number of attention heads
+    d_feature: int:  depth of embedding
+    d_feedforward: int: depth of feed-forward layer
+    n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
 
   Returns:
     the layer.
   """
-  return tl.Serial(
-      tl.Residual(  # Self-attention block.
-          tl.LayerNorm(),
-          tl.Branch(tl.NoOp(), tl.CausalMask(axis=-2)),  # Create mask.
-          tl.MultiHeadedAttention(feature_depth, num_heads=num_heads,
-                                  dropout=dropout, mode=mode),
-          tl.Select(0),  # Drop the mask.
-          tl.Dropout(rate=dropout, mode=mode)
-      ),
-      ResidualFeedForward(feature_depth, feedforward_depth, dropout, mode=mode)
-  )
+  self_attention = [
+      tl.LayerNorm(),
+      tl.Branch(tl.NoOp(), tl.CausalMask(axis=-2)),  # Create mask.
+      tl.MultiHeadedAttention(
+          d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+      tl.Select(0),  # Drop mask.
+      tl.Dropout(rate=dropout, mode=mode),
+  ]
+  feed_forward = [
+      FeedForward(d_feature, d_feedforward, dropout, mode=mode),
+  ]
+  return [
+      tl.Residual(self_attention),
+      tl.Residual(feed_forward),
+  ]
 
 
 def TransformerLM(vocab_size,
-                  feature_depth=512,
-                  feedforward_depth=2048,
-                  num_layers=6,
-                  num_heads=8,
+                  d_feature=512,
+                  d_feedforward=2048,
+                  n_layers=6,
+                  n_heads=8,
                   dropout=0.1,
                   max_len=2048,
                   mode='train'):
@@ -154,10 +149,10 @@ def TransformerLM(vocab_size,
 
   Args:
     vocab_size: int: vocab size
-    feature_depth: int:  depth of embedding
-    feedforward_depth: int: depth of feed-forward layer
-    num_layers: int: number of encoder/decoder layers
-    num_heads: int: number of attention heads
+    d_feature: int:  depth of embedding
+    d_feedforward: int: depth of feed-forward layer
+    n_layers: int: number of encoder/decoder layers
+    n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
     max_len: int: maximum symbol length for positional encoding
     mode: str: 'train' or 'eval'
@@ -165,25 +160,23 @@ def TransformerLM(vocab_size,
   Returns:
     the layer.
   """
-  return tl.Serial(
-      tl.ShiftRight(),
-      tl.Embedding(feature_depth, vocab_size),
+  positional_embedder = [
+      tl.Embedding(d_feature, vocab_size),
       tl.Dropout(rate=dropout, mode=mode),
       tl.PositionalEncoding(max_len=max_len),
-      tl.Serial(*[DecoderLayer(feature_depth, feedforward_depth, num_heads,
-                               dropout, mode)
-                  for _ in range(num_layers)]),
+  ]
+  return [
+      tl.ShiftRight(),
+      positional_embedder,
+      [DecoderBlock(d_feature, d_feedforward, n_heads, dropout, mode)
+       for _ in range(n_layers)],
       tl.LayerNorm(),
       tl.Dense(vocab_size),
-      tl.LogSoftmax()
-  )
+      tl.LogSoftmax(),
+  ]
 
 
-def EncoderDecoderLayer(feature_depth,
-                        feedforward_depth,
-                        num_heads,
-                        dropout,
-                        mode):
+def EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode):
   """Transformer encoder-decoder layer.
 
   The input is a triple pair (decoder_input, mask, encoder) where
@@ -191,47 +184,49 @@ def EncoderDecoderLayer(feature_depth,
   to the padding part of the encoder.
 
   Args:
-    feature_depth: int:  depth of embedding
-    feedforward_depth: int: depth of feed-forward layer
-    num_heads: int: number of attention heads
+    d_feature: int:  depth of embedding
+    d_feedforward: int: depth of feed-forward layer
+    n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
 
   Returns:
     the layer, returning a triple (decoder_activations, mask, encoder).
   """
-  # Decoder self-attending to decoder.
-  self_attention = tl.Residual(
+  decoder_self_attention = [
+      # TODO(jonni): Work on combinators so that this flow is cleaner/clearer.
       tl.LayerNorm(),
       tl.Dup(),
       tl.CausalMask(axis=-2),  # Create the self-attention mask.
       tl.Swap(),  # Put mask behind the activations.
-      tl.MultiHeadedAttention(feature_depth, num_heads=num_heads,
+      tl.MultiHeadedAttention(d_feature, n_heads=n_heads,
                               dropout=dropout, mode=mode),
       tl.Swap(),  # Put self-attention mask on top.
       tl.Drop(),   # Drop self-attention mask.
-      tl.Dropout(rate=dropout, mode=mode)
-  )
-  # Decoder attending to encoder.
-  encoder_decoder_attention = tl.Serial(
+      tl.Dropout(rate=dropout, mode=mode),
+  ]
+  decoder_to_encoder_attention = [
       tl.Select((0, 2, 2, 1, 2)),  # (dec, enc, enc, mask, enc-copy)
       tl.MultiHeadedAttentionQKV(  # (q, k, v, mask, ...) --> (new, mask, ...)
-          feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),
+          d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Dropout(rate=dropout, mode=mode),
-  )
-  return tl.Serial(
-      self_attention,
-      tl.Residual(encoder_decoder_attention),
-      ResidualFeedForward(feature_depth, feedforward_depth, dropout, mode=mode)
-  )
+  ]
+  feed_forward = [
+      FeedForward(d_feature, d_feedforward, dropout, mode=mode),
+  ]
+  return [
+      tl.Residual(decoder_self_attention),
+      tl.Residual(decoder_to_encoder_attention),
+      tl.Residual(feed_forward),
+  ]
 
 
 # TODO(lukaszkaiser): allow different source and target vocabularies.
 def Transformer(vocab_size,
-                feature_depth=512,
-                feedforward_depth=2048,
-                num_layers=6,
-                num_heads=8,
+                d_feature=512,
+                d_feedforward=2048,
+                n_layers=6,
+                n_heads=8,
                 dropout=0.1,
                 max_len=2048,
                 mode='train'):
@@ -241,10 +236,10 @@ def Transformer(vocab_size,
 
   Args:
     vocab_size: int: vocab size (shared source and target).
-    feature_depth: int:  depth of embedding
-    feedforward_depth: int: depth of feed-forward layer
-    num_layers: int: number of encoder/decoder layers
-    num_heads: int: number of attention heads
+    d_feature: int:  depth of embedding
+    d_feedforward: int: depth of feed-forward layer
+    n_layers: int: number of encoder/decoder layers
+    n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
     max_len: int: maximum symbol length for positional encoding
     mode: str: 'train' or 'eval'
@@ -252,31 +247,28 @@ def Transformer(vocab_size,
   Returns:
     the Transformer model.
   """
-  embedding = tl.Serial(
-      tl.Embedding(feature_depth, vocab_size),
+  positional_embedder = [
+      tl.Embedding(d_feature, vocab_size),
       tl.Dropout(rate=dropout, mode=mode),
-      tl.PositionalEncoding(max_len=max_len)
-  )
-  encoder = tl.Serial(
-      tl.Branch(embedding, tl.PaddingMask()),
-      tl.Serial(*[EncoderLayer(feature_depth, feedforward_depth, num_heads,
-                               dropout, mode)
-                  for _ in range(num_layers)]),
-      tl.LayerNorm()
-  )
-  stack = [EncoderDecoderLayer(feature_depth, feedforward_depth, num_heads,
-                               dropout, mode)
-           for _ in range(num_layers)]
-  return tl.Serial(
+      tl.PositionalEncoding(max_len=max_len),
+  ]
+  encoder = [
+      tl.Branch(positional_embedder, tl.PaddingMask()),
+      [EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode)
+       for _ in range(n_layers)],
+      tl.LayerNorm(),
+  ]
+  return [
       tl.Parallel(tl.NoOp(), tl.ShiftRight()),
-      tl.Parallel(encoder, embedding),
+      tl.Parallel(encoder, positional_embedder),
       tl.Select(inputs=(('encoder', 'mask'), 'decoder'),
                 output=('decoder', ('mask', 'decoder'), 'encoder')),
-      tl.Parallel(  # (encoder_mask, decoder_input) -> encoder-decoder mask
-          tl.NoOp(), tl.EncoderDecoderMask(), tl.NoOp()),
-      tl.Serial(*stack),
+      # (encoder_mask, decoder_input) -> encoder-decoder mask
+      tl.Parallel(tl.NoOp(), tl.EncoderDecoderMask(), tl.NoOp()),
+      [EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode)
+       for _ in range(n_layers)],
       tl.Select(0),  # Drop mask and encoder.
       tl.LayerNorm(),
       tl.Dense(vocab_size),
-      tl.LogSoftmax()
-  )
+      tl.LogSoftmax(),
+  ]
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index cd8b4b43d..3581f3f7a 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -65,7 +65,7 @@
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.trax import jaxboard
-from tensor2tensor.trax import layers
+from tensor2tensor.trax import layers as tl
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
 from tensorflow.io import gfile
@@ -82,8 +82,8 @@
 
 def policy_and_value_net(rng_key,
                          batch_observations_shape,
-                         num_actions,
-                         bottom_layers_fn=None,
+                         n_actions,
+                         bottom_layers_fn=(),
                          two_towers=True):
   """A policy and value net function."""
 
@@ -93,31 +93,23 @@ def policy_and_value_net(rng_key,
   # other computes the value function.
   # NOTE: The LogSoftmax instead of the Softmax because of numerical stability.
 
-  net = None
-  if not two_towers:
-    tower = [] if bottom_layers_fn is None else bottom_layers_fn()
-    tower.extend([
-        layers.Branch(
-            layers.Serial(layers.Dense(num_actions), layers.LogSoftmax()),
-            layers.Dense(1))
-    ])
-    net = layers.Serial(*tower)
+  if two_towers:
+    net = tl.Branch(
+        [bottom_layers_fn(), tl.Dense(n_actions), tl.LogSoftmax()],
+        [bottom_layers_fn(), tl.Dense(1)]
+    )
   else:
-    tower1 = [] if bottom_layers_fn is None else bottom_layers_fn()
-    tower2 = [] if bottom_layers_fn is None else bottom_layers_fn()
-
-    tower1.extend([layers.Dense(num_actions), layers.LogSoftmax()])
-    tower2.extend([layers.Dense(1)])
-
-    net = layers.Branch(
-        layers.Serial(*tower1),
-        layers.Serial(*tower2),
+    net = tl.Serial(
+        bottom_layers_fn(),
+        tl.Branch(
+            [tl.Dense(n_actions), tl.LogSoftmax()],
+            [tl.Dense(1)]
+        )
     )
-  assert net
   return net.initialize(batch_observations_shape, rng_key), net
 
 
-def optimizer_fun(net_params, step_size=1e-3):
+def optimizer_fn(net_params, step_size=1e-3):
   opt = trax_opt.Adam(step_size=step_size, b1=0.9, b2=0.999, eps=1e-08)
   opt_init = lambda x: (x, opt.tree_init(x))
   opt_update = lambda i, g, s: opt.tree_update(i, g, s[0], s[1])
@@ -130,7 +122,7 @@ def optimizer_fun(net_params, step_size=1e-3):
 # Run the env for 'n' steps and take completed trajectories, or
 # Any other option?
 def collect_trajectories(env,
-                         policy_fun,
+                         policy_fn,
                          num_trajectories=1,
                          policy=env_problem_utils.CATEGORICAL_SAMPLING,
                          max_timestep=None,
@@ -142,10 +134,10 @@ def collect_trajectories(env,
 
   Args:
     env: A gym env interface, for now this is not-batched.
-    policy_fun: observations(B,T+1) -> log-probabs(B,T+1, A) callable.
+    policy_fn: observations(B,T+1) -> log-probabs(B,T+1, A) callable.
     num_trajectories: int, number of trajectories.
     policy: string, "greedy", "epsilon-greedy", or "categorical-sampling" i.e.
-      how to use the policy_fun to return an action.
+      how to use the policy_fn to return an action.
     max_timestep: int or None, the index of the maximum time-step at which we
       return the trajectory, None for ending a trajectory only when env returns
       done.
@@ -168,7 +160,7 @@ def collect_trajectories(env,
   # This is an env_problem, run its collect function.
   trajs, num_done = env_problem_utils.play_env_problem_with_policy(
       env,
-      policy_fun,
+      policy_fn,
       num_trajectories=num_trajectories,
       max_timestep=max_timestep,
       boundary=boundary,
@@ -740,8 +732,8 @@ def maybe_restore_params(output_dir, policy_and_value_net_params):
 def training_loop(
     env=None,
     epochs=EPOCHS,
-    policy_and_value_net_fun=None,
-    policy_and_value_optimizer_fun=None,
+    policy_and_value_net_fn=None,
+    policy_and_value_optimizer_fn=None,
     batch_size=BATCH_TRAJECTORIES,
     num_optimizer_steps=NUM_OPTIMIZER_STEPS,
     print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
@@ -785,13 +777,13 @@ def training_loop(
   batch_observations_shape = (-1, -1) + env.observation_space.shape
 
   assert isinstance(env.action_space, gym.spaces.Discrete)
-  num_actions = env.action_space.n
+  n_actions = env.action_space.n
 
   jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)
 
   # Initialize the policy and value network.
   policy_and_value_net_params, policy_and_value_net_apply = (
-      policy_and_value_net_fun(key1, batch_observations_shape, num_actions))
+      policy_and_value_net_fn(key1, batch_observations_shape, n_actions))
 
   # Maybe restore the policy params. If there is nothing to restore, then
   # iteration = 0 and policy_and_value_net_params are returned as is.
@@ -807,7 +799,7 @@ def training_loop(
 
   # Initialize the optimizers.
   policy_and_value_optimizer = (
-      policy_and_value_optimizer_fun(policy_and_value_net_params))
+      policy_and_value_optimizer_fn(policy_and_value_net_params))
   (policy_and_value_opt_state, policy_and_value_opt_update,
    policy_and_value_get_params) = policy_and_value_optimizer
 
@@ -860,7 +852,7 @@ def get_predictions(observations, rng=None):
     jax_rng_key, key = jax_random.split(jax_rng_key)
     trajs, num_done = collect_trajectories(
         env,
-        policy_fun=get_predictions,
+        policy_fn=get_predictions,
         num_trajectories=batch_size,
         max_timestep=max_timestep,
         boundary=boundary,
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index 0727b51da..eae8dd3e2 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -66,10 +66,10 @@ def test_training_loop(self):
           env=env,
           eval_env=eval_env,
           epochs=num_epochs,
-          policy_and_value_net_fun=functools.partial(
+          policy_and_value_net_fn=functools.partial(
               ppo.policy_and_value_net,
               bottom_layers_fn=lambda: [layers.Dense(1)]),
-          policy_and_value_optimizer_fun=ppo.optimizer_fun,
+          policy_and_value_optimizer_fn=ppo.optimizer_fn,
           batch_size=batch_size,
           num_optimizer_steps=1,
           output_dir=output_dir,
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index b1606b45c..0c6bde87b 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -168,9 +168,9 @@ def save_state(state, output_dir, keep=False):
   log("Model saved to %s" % params_file, stdout=False)
 
 
-def _save_replicated(opt_state, step, history, num_devices, output_dir, keep):
+def _save_replicated(opt_state, step, history, n_devices, output_dir, keep):
   """Save state but given a possibly replicated opt_state."""
-  if num_devices > 1:
+  if n_devices > 1:
     unreplicate = lambda x: x.mean(0)
     opt_state = layers.nested_map(opt_state, unreplicate)
     save_state(State(params=opt_state, step=step, history=history),
@@ -185,14 +185,14 @@ def _save_replicated(opt_state, step, history, num_devices, output_dir, keep):
 }
 
 
-def evaluate_train_and_eval(step, inputs, predict_fun, eval_steps, rng,
+def evaluate_train_and_eval(step, inputs, predict_fn, eval_steps, rng,
                             train_sw=None, eval_sw=None, history=None):
   """Evalaute on train and eval data, and log metrics."""
   step_log(step, "Evaluation")
   train_metrics, eval_metrics = [
       evaluate(  # pylint: disable=g-complex-comprehension
           itertools.islice(input_stream(), eval_steps),
-          predict_fun,
+          predict_fn,
           _METRICS,
           rng)
       for input_stream in
@@ -205,14 +205,14 @@ def evaluate_train_and_eval(step, inputs, predict_fun, eval_steps, rng,
   return train_metrics, eval_metrics
 
 
-def evaluate(inputs_stream, predict_fun, metric_funs, rng):
+def evaluate(inputs_stream, predict_fn, metric_fns, rng):
   """Evaluate.
 
   Args:
     inputs_stream: iterable of inputs to evaluate on.
-    predict_fun: function from inputs to predictions. params should already be
+    predict_fn: function from inputs to predictions. params should already be
       partially applied.
-    metric_funs: dict from metric name to metric function, which takes inputs
+    metric_fns: dict from metric name to metric function, which takes inputs
       and predictions and returns a scalar metric value.
     rng: random number generator.
 
@@ -225,8 +225,8 @@ def evaluate(inputs_stream, predict_fun, metric_funs, rng):
   for inp in inputs_stream:
     count += 1
     rng, subrng = jax_random.split(rng)
-    preds = predict_fun(inp[0], rng=subrng)
-    for m, f in six.iteritems(metric_funs):
+    preds = predict_fn(inp[0], rng=subrng)
+    for m, f in six.iteritems(metric_fns):
       metrics[m] += f(inp, preds)
   return {m: v / count for (m, v) in six.iteritems(metrics)}
 
@@ -293,7 +293,7 @@ def epochs(steps=None, epoch_steps=1):
       break
 
 
-def _jit_predict_fun(model_predict, num_devices):
+def _jit_predict_fn(model_predict, n_devices):
   """Use jit on model_predict if required."""
 
   # Multi-devices, pmap and run.
@@ -304,13 +304,13 @@ def mapped_predict(x, params, rng):
   def predict(x, params=(), rng=None):
     """Predict function jited and parallelized as requested."""
     # On one device, jit and run.
-    if num_devices == 1:
+    if n_devices == 1:
       return backend.jit(model_predict)(x, params, rng=rng)
 
     pred = mapped_predict(
-        reshape_by_device(x, num_devices),
+        reshape_by_device(x, n_devices),
         params,
-        jax_random.split(rng, num_devices))
+        jax_random.split(rng, n_devices))
     # Need to reduce the [device, per-device-batch, ...] tensors back to
     # a [batch, ...] tensor. The tensors may be nested.
     if not isinstance(pred, (list, tuple)):  # Not nested.
@@ -322,58 +322,58 @@ def predict(x, params=(), rng=None):
   return predict
 
 
-def _jit_update_fun(predict_fun, loss_fun, optimizer, num_devices):
+def _jit_update_fn(predict_fn, loss_fn, optimizer, n_devices):
   """Get jit-ed update function for loss, optimizer, learning rate function."""
-  if num_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
+  if n_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
     def single_update(i, opt_state, batch, rng):
       rng, subrng = jax_random.split(rng[0])
       params, opt_slots = opt_state
-      return optimizer.tree_update(i, backend.grad(loss_fun)(
-          params, batch, predict_fun, rng), params, opt_slots), [subrng]
+      return optimizer.tree_update(i, backend.grad(loss_fn)(
+          params, batch, predict_fn, rng), params, opt_slots), [subrng]
     return backend.jit(single_update)
 
   @functools.partial(backend.pmap, axis_name="batch")
   def mapped_update(i, opt_state, batch, rng):
     """This is a multi-device version of the update function above."""
-    # We assume all tensors have the first dimension = num_devices.
+    # We assume all tensors have the first dimension = n_devices.
     rng, subrng = jax_random.split(rng)
     params, opt_slots = opt_state
-    grads = backend.grad(loss_fun)(params, batch, predict_fun, rng)
+    grads = backend.grad(loss_fn)(params, batch, predict_fn, rng)
     grads = jax.tree_util.tree_map(
         lambda g: lax.psum(g, "batch"), grads)
     return optimizer.tree_update(i, grads, params, opt_slots), subrng
 
   def update(i, opt_state, batch, rng):
-    return mapped_update(numpy.repeat(i, num_devices), opt_state, batch, rng)
+    return mapped_update(numpy.repeat(i, n_devices), opt_state, batch, rng)
 
   return update
 
 
-def _reshape_by_device_single(x, num_devices):
-  """Reshape x into a shape [num_devices, ...]."""
+def _reshape_by_device_single(x, n_devices):
+  """Reshape x into a shape [n_devices, ...]."""
   x_shape = list(x.shape)
   batch_size = x_shape[0]
-  batch_size_per_device = batch_size // num_devices
-  # We require that num_devices divides batch_size evenly.
-  if batch_size_per_device * num_devices != batch_size:
+  batch_size_per_device = batch_size // n_devices
+  # We require that n_devices divides batch_size evenly.
+  if batch_size_per_device * n_devices != batch_size:
     logging.fatal(
-        "We require that num_devices[%d] divides batch_size[%d] evenly.",
-        num_devices, batch_size)
+        "We require that n_devices[%d] divides batch_size[%d] evenly.",
+        n_devices, batch_size)
   # New shape.
-  new_shape_prefix = [num_devices, batch_size_per_device]
+  new_shape_prefix = [n_devices, batch_size_per_device]
   return np.reshape(x, new_shape_prefix + x_shape[1:])
 
 
-def reshape_by_device(x, num_devices):
-  """Reshape possibly nested x into a shape [num_devices, ...]."""
+def reshape_by_device(x, n_devices):
+  """Reshape possibly nested x into a shape [n_devices, ...]."""
   return layers.nested_map(
-      x, lambda x: _reshape_by_device_single(x, num_devices))
+      x, lambda x: _reshape_by_device_single(x, n_devices))
 
 
 @gin.configurable(blacklist=["output_dir"])
 def train(output_dir,
           model=gin.REQUIRED,
-          loss_fun=loss,
+          loss_fn=loss,
           inputs=trax_inputs.inputs,
           optimizer=trax_opt.SM3,
           lr_schedule=lr.MultifactorSchedule,
@@ -381,7 +381,7 @@ def train(output_dir,
           save_steps=None,
           eval_steps=10,
           eval_frequency=100,
-          num_devices=None,
+          n_devices=None,
           random_seed=None,
           run_debug_step=False,
           save_graphs=True,
@@ -390,9 +390,9 @@ def train(output_dir,
 
   Args:
     output_dir: Directory where to put the logs and checkpoints.
-    model: The model to train as a callable returning 2 callables, an init_fun
-      and apply_fun.
-    loss_fun: callable with signature: params, trax.inputs.Inputs, model, rng
+    model: The model to train as a callable returning 2 callables, an init_fn
+      and apply_fn.
+    loss_fn: callable with signature: params, trax.inputs.Inputs, model, rng
       -> loss.
     inputs: callable returning trax.inputs.Inputs.
     optimizer: The optimizer (see optimizers/base.py for signature).
@@ -404,7 +404,7 @@ def train(output_dir,
     eval_steps: int, num of steps per evaluation. If None or 0, eval disabled.
     eval_frequency: int, how often to run evaluation (every eval_frequency
       steps). If None or 0, eval disabled.
-    num_devices: how many devices to use (if None, default, use all available)
+    n_devices: how many devices to use (if None, default, use all available)
     random_seed: the random seed to use; time/os dependent if None (default).
     run_debug_step: bool, if True, will run the model and loss without @jit for
       one step.
@@ -416,31 +416,32 @@ def train(output_dir,
   if save_steps is None:
     save_steps = []
   device_count = jax.lib.xla_bridge.device_count()
-  num_devices = num_devices or device_count
+  n_devices = n_devices or device_count
   # TODO(lukaszkaiser): remove this restriction when possible.
-  if num_devices != device_count:
-    raise ValueError("Jax cannot work yet with num_devices != all devices: "
-                     "%d != %d" % (num_devices, device_count))
+  if n_devices != device_count:
+    raise ValueError("Jax cannot work yet with n_devices != all devices: "
+                     "%d != %d" % (n_devices, device_count))
   rng = get_random_number_generator_and_set_seed(random_seed)
   gfile.makedirs(output_dir)
   # Create summary writers and history.
   train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
   eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
 
-  inputs = inputs(num_devices)
+  inputs = inputs(n_devices)
 
   # Setup optimizer and model
   state = restore_state(output_dir)
   history = state.history
-  lr_fun = lr_schedule(history)
-  opt = optimizer(lr_fun)
-  model_train = model(mode="train")
-  model_predict_eval = model(mode="eval")
+  lr_fn = lr_schedule(history)
+  opt = optimizer(lr_fn)
+
+  model_train = layers.Serial(model(mode="train"))
+  model_predict_eval = layers.Serial(model(mode="eval"))
 
   # Setup state
   step = state.step or 0
   rng, init_rng = jax_random.split(rng)
-  rngs = jax_random.split(rng, num_devices)
+  rngs = jax_random.split(rng, n_devices)
   first_shape = inputs.input_shape[0]
   # If the inputs are a tuple/list, add [-1] (batch) to each element.
   if isinstance(first_shape, (list, tuple)):
@@ -454,13 +455,13 @@ def train(output_dir,
   else:
     params = model_train.initialize(model_input_shape, init_rng)
     opt_state = (params, opt.tree_init(params))
-  if num_devices > 1:
-    replicate = lambda x: numpy.broadcast_to(x, (num_devices,) + x.shape)
+  if n_devices > 1:
+    replicate = lambda x: numpy.broadcast_to(x, (n_devices,) + x.shape)
     opt_state = layers.nested_map(opt_state, replicate)
 
   # jit model_predict and update so they're fast
-  jit_model_predict_eval = _jit_predict_fun(model_predict_eval, num_devices)
-  jit_update_fun = _jit_update_fun(model_train, loss_fun, opt, num_devices)
+  jit_model_predict_eval = _jit_predict_fn(model_predict_eval, n_devices)
+  jit_update_fn = _jit_update_fn(model_train, loss_fn, opt, n_devices)
 
   train_stream = inputs.train_stream()
   epoch_steps = [train_steps]  # Only training if eval_frequency is 0 or None.
@@ -468,11 +469,11 @@ def train(output_dir,
     epoch_steps = itertools.chain([1,  # first epoch only 1 step
                                    eval_frequency - 1],
                                   itertools.repeat(eval_frequency))
-  step_log(step, "Starting training using %d devices" % num_devices)
+  step_log(step, "Starting training using %d devices" % n_devices)
 
   # Non-compiled debug step helps find problems in models easier.
   if run_debug_step:
-    debug_loss = loss_fun(params, next(train_stream), model_train, rng)
+    debug_loss = loss_fn(params, next(train_stream), model_train, rng)
     step_log(step, "Debug step loss %.8f" % debug_loss)
 
   for epoch, epoch_steps in epochs(train_steps, epoch_steps):
@@ -485,19 +486,18 @@ def train(output_dir,
     for _ in range(epoch_steps):
       # Train
       next_train_batch = next(train_stream)
-      if num_devices > 1:  # TODO(lukaszkaiser): use everywhere when possible.
-        next_train_batch = reshape_by_device(next_train_batch, num_devices)
-      opt_state, rngs = jit_update_fun(step, opt_state, next_train_batch, rngs)
+      if n_devices > 1:  # TODO(lukaszkaiser): use everywhere when possible.
+        next_train_batch = reshape_by_device(next_train_batch, n_devices)
+      opt_state, rngs = jit_update_fn(step, opt_state, next_train_batch, rngs)
       step += 1
 
       if step in save_steps:
-        _save_replicated(opt_state, step, history, num_devices,
-                         output_dir, True)
+        _save_replicated(opt_state, step, history, n_devices, output_dir, True)
 
       # LR log
       if step == 1 or step % 10 == 0:
         train_sw.scalar("training/learning rate",
-                        lr_fun(step), step=step)
+                        lr_fn(step), step=step)
 
     # Timer
     epoch_time = time.time() - start_time
@@ -510,7 +510,7 @@ def train(output_dir,
     # Print number of parameters
     if step == 1:
       sizes = layers.sizes(opt_state[0])
-      if num_devices > 1:
+      if n_devices > 1:
         unreplicate = lambda x: x.mean(0)
         single_params = layers.nested_map(opt_state[0], unreplicate)
         sizes = layers.sizes(single_params)
@@ -521,8 +521,8 @@ def train(output_dir,
     evaluate_train_and_eval(
         step=step,
         inputs=inputs,
-        predict_fun=functools.partial(jit_model_predict_eval,
-                                      params=opt_state[0]),
+        predict_fn=functools.partial(jit_model_predict_eval,
+                                     params=opt_state[0]),
         eval_steps=eval_steps,
         rng=rng,
         train_sw=train_sw,
@@ -530,7 +530,7 @@ def train(output_dir,
         history=history)
 
     # Save computation graph (single-device only for now).
-    if save_graphs and step == 1 and num_devices == 1:
+    if save_graphs and step == 1 and n_devices == 1:
       params = opt_state[0]
       # Dump computation graphs to files.
       forward_computation = jax.xla_computation(model_predict_eval)(
@@ -539,7 +539,7 @@ def train(output_dir,
         f.write(forward_computation.GetHloText())
       with gfile.GFile(os.path.join(output_dir, "forward.dot"), "w") as f:
         f.write(forward_computation.GetHloDotGraph())
-      backward_computation = jax.xla_computation(jit_update_fun)(
+      backward_computation = jax.xla_computation(jit_update_fn)(
           step, opt_state, next_train_batch, rngs)
       with gfile.GFile(os.path.join(output_dir, "backward.txt"), "w") as f:
         f.write(backward_computation.GetHloText())
@@ -548,8 +548,7 @@ def train(output_dir,
           f.write(backward_computation.GetHloDotGraph())
 
     # Save state
-    _save_replicated(opt_state, step, history, num_devices,
-                     output_dir, False)
+    _save_replicated(opt_state, step, history, n_devices, output_dir, False)
 
     # Save Gin config
     # Gin only tracks the used parameters, so we save it after the first epoch.
@@ -557,11 +556,11 @@ def train(output_dir,
       save_gin(output_dir, train_sw)
 
     # Update learning rate with new history
-    old_lr_fun = lr_fun
-    lr_fun = lr_schedule(history)
-    if lr_fun != old_lr_fun:  # For performance, only jit if there is a change.
-      opt = optimizer(lr_fun)
-      jit_update_fun = _jit_update_fun(model_train, loss_fun, opt, num_devices)
+    old_lr_fn = lr_fn
+    lr_fn = lr_schedule(history)
+    if lr_fn != old_lr_fn:  # For performance, only jit if there is a change.
+      opt = optimizer(lr_fn)
+      jit_update_fn = _jit_update_fn(model_train, loss_fn, opt, n_devices)
 
     # Flush summary writers
     train_sw.flush()
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index 88ed84354..71e24ea47 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -28,6 +28,7 @@
 import numpy as np
 
 from tensor2tensor.trax import inputs as inputs_lib
+from tensor2tensor.trax import layers
 from tensor2tensor.trax import models
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
@@ -36,7 +37,7 @@
 from tensorflow.io import gfile
 
 
-def test_inputs(num_classes):
+def test_inputs(n_classes):
   """Make trax.inputs.Inputs."""
   batch_size = 2
   input_shape = (6, 6, 3)
@@ -44,7 +45,7 @@ def test_inputs(num_classes):
   def input_stream():
     while True:
       yield (np.random.rand(*([batch_size] + list(input_shape))),
-             np.random.randint(num_classes, size=batch_size))
+             np.random.randint(n_classes, size=batch_size))
 
   return inputs_lib.Inputs(
       train_stream=input_stream,
@@ -64,17 +65,17 @@ def tmp_dir(self):
   def test_train_eval_predict(self):
     with self.tmp_dir() as output_dir:
       # Prepare model and inputs
-      num_classes = 4
+      n_classes = 4
       train_steps = 2
       eval_steps = 2
-      model = functools.partial(models.MLP,
-                                hidden_size=16,
-                                num_output_classes=num_classes)
-      inputs = lambda _: test_inputs(num_classes)
+      model_fn = functools.partial(models.MLP,
+                                   d_hidden=16,
+                                   n_output_classes=n_classes)
+      inputs = lambda _: test_inputs(n_classes)
 
       # Train and evaluate
       state = trax.train(output_dir,
-                         model=model,
+                         model=model_fn,
                          inputs=inputs,
                          train_steps=train_steps,
                          eval_steps=eval_steps)
@@ -90,22 +91,23 @@ def test_train_eval_predict(self):
 
       # Predict with final params
       inputs = inputs(1).train_stream()
-      model()(next(inputs)[0], state.params[0])
+      model = layers.Serial(model_fn())
+      model(next(inputs)[0], state.params[0])
 
   def test_train_eval_predict_sm3(self):
     with self.tmp_dir() as output_dir:
       # Prepare model and inputs
-      num_classes = 4
+      n_classes = 4
       train_steps = 2
       eval_steps = 2
-      model = functools.partial(models.MLP,
-                                hidden_size=16,
-                                num_output_classes=num_classes)
-      inputs = lambda _: test_inputs(num_classes)
+      model_fn = functools.partial(models.MLP,
+                                   d_hidden=16,
+                                   n_output_classes=n_classes)
+      inputs = lambda _: test_inputs(n_classes)
 
       # Train and evaluate
       state = trax.train(output_dir,
-                         model=model,
+                         model=model_fn,
                          inputs=inputs,
                          train_steps=train_steps,
                          eval_steps=eval_steps,
@@ -122,7 +124,8 @@ def test_train_eval_predict_sm3(self):
 
       # Predict with final params
       inputs = inputs(1).train_stream()
-      model()(next(inputs)[0], state.params[0])
+      model = layers.Serial(model_fn())
+      model(next(inputs)[0], state.params[0])
 
 
 if __name__ == "__main__":

From 70197062a4bd8d04ded87bb09347dc3454cbdccd Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 29 May 2019 17:39:16 -0700
Subject: [PATCH 2070/2720] Add ability to run eval multiple times and report
 the average.

PiperOrigin-RevId: 250601682
---
 tensor2tensor/trax/rlax/ppo.py      | 48 ++++++++++++++++++-----------
 tensor2tensor/trax/rlax/ppo_main.py |  2 ++
 2 files changed, 32 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 3581f3f7a..7241197d6 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -49,6 +49,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import functools
 import os
 import time
@@ -682,27 +683,36 @@ def evaluate_policy(eval_env,
                     get_predictions,
                     boundary,
                     max_timestep=20000,
+                    num_evals=1,
                     rng=None):
   """Evaluate the policy."""
 
-  avg_rewards = {}
-  avg_rewards_unclipped = {}
-  for policy in [
-      env_problem_utils.CATEGORICAL_SAMPLING, env_problem_utils.GUMBEL_SAMPLING,
-      env_problem_utils.EPSILON_GREEDY
-  ]:
-    trajs, _ = env_problem_utils.play_env_problem_with_policy(
-        eval_env,
-        get_predictions,
-        boundary=boundary,
-        max_timestep=max_timestep,
-        reset=True,
-        policy_sampling=policy,
-        rng=rng)
-    avg_rewards[policy] = float(sum(
-        np.sum(traj[2]) for traj in trajs)) / len(trajs)
-    avg_rewards_unclipped[policy] = float(sum(
-        np.sum(traj[3]) for traj in trajs)) / len(trajs)
+  avg_rewards = collections.defaultdict(float)
+  avg_rewards_unclipped = collections.defaultdict(float)
+  for _ in range(num_evals):
+    for policy in [
+        env_problem_utils.CATEGORICAL_SAMPLING,
+        env_problem_utils.GUMBEL_SAMPLING,
+        env_problem_utils.EPSILON_GREEDY
+    ]:
+      trajs, _ = env_problem_utils.play_env_problem_with_policy(
+          eval_env,
+          get_predictions,
+          boundary=boundary,
+          max_timestep=max_timestep,
+          reset=True,
+          policy_sampling=policy,
+          rng=rng)
+      avg_rewards[policy] += float(sum(
+          np.sum(traj[2]) for traj in trajs)) / len(trajs)
+      avg_rewards_unclipped[policy] += float(sum(
+          np.sum(traj[3]) for traj in trajs)) / len(trajs)
+
+  # Now average these out.
+  for k in avg_rewards:
+    avg_rewards[k] /= num_evals
+    avg_rewards_unclipped[k] /= num_evals
+
   return avg_rewards, avg_rewards_unclipped
 
 
@@ -753,6 +763,7 @@ def training_loop(
     done_frac_for_policy_save=0.5,
     enable_early_stopping=True,
     env_name=None,
+    num_evals=1,
 ):
   """Runs the training loop for PPO, with fixed policy and value nets."""
   assert env
@@ -836,6 +847,7 @@ def get_predictions(observations, rng=None):
           get_predictions,
           boundary,
           max_timestep=max_timestep_eval,
+          num_evals=num_evals,
           rng=key)
       for k, v in avg_reward.items():
         eval_sw.scalar("eval/mean_reward/%s" % k, v, step=i)
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 6dc40134e..882849850 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -126,6 +126,7 @@
 flags.DEFINE_bool("xm", False, "Copy atari roms?")
 flags.DEFINE_integer("eval_every_n", 100, "How frequently to eval the policy.")
 flags.DEFINE_integer("eval_batch_size", 4, "Batch size for evaluation.")
+flags.DEFINE_integer("num_evals", 1, "Number of times to evaluate.")
 flags.DEFINE_float(
     "done_frac_for_policy_save", 0.5,
     "Fraction of the trajectories that should be done to "
@@ -240,6 +241,7 @@ def run_training_loop():
         eval_every_n=FLAGS.eval_every_n,
         done_frac_for_policy_save=FLAGS.done_frac_for_policy_save,
         eval_env=eval_env,
+        num_evals=FLAGS.num_evals,
         env_name=str(FLAGS.env_problem_name),
     )
 

From 33b2c62dbd28477b6d6c0375646c0e8de68ca5f2 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 29 May 2019 18:57:49 -0700
Subject: [PATCH 2071/2720] Fix breakage from a recent change s/fun/fn in
 ppo_main.py

PiperOrigin-RevId: 250611310
---
 tensor2tensor/trax/rlax/ppo_main.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 882849850..0a66a08ae 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -178,8 +178,8 @@ def make_env(batch_size=8):
       reward_range=(-1, 1))
 
 
-def get_optimizer_fun(learning_rate):
-  return functools.partial(ppo.optimizer_fun, step_size=learning_rate)
+def get_optimizer_fn(learning_rate):
+  return functools.partial(ppo.optimizer_fn, step_size=learning_rate)
 
 
 def main(argv):
@@ -206,11 +206,11 @@ def run_training_loop():
     """Runs the training loop."""
     logging.info("Starting the training loop.")
 
-    policy_and_value_net_fun = functools.partial(
+    policy_and_value_net_fn = functools.partial(
         ppo.policy_and_value_net,
         bottom_layers_fn=common_layers,
         two_towers=FLAGS.two_towers)
-    policy_and_value_optimizer_fun = get_optimizer_fun(FLAGS.learning_rate)
+    policy_and_value_optimizer_fn = get_optimizer_fn(FLAGS.learning_rate)
 
     random_seed = None
     try:
@@ -221,8 +221,8 @@ def run_training_loop():
     ppo.training_loop(
         env=env,
         epochs=FLAGS.epochs,
-        policy_and_value_net_fun=policy_and_value_net_fun,
-        policy_and_value_optimizer_fun=policy_and_value_optimizer_fun,
+        policy_and_value_net_fn=policy_and_value_net_fn,
+        policy_and_value_optimizer_fn=policy_and_value_optimizer_fn,
         num_optimizer_steps=FLAGS.num_optimizer_steps,
         print_every_optimizer_steps=FLAGS.print_every_optimizer_steps,
         batch_size=FLAGS.batch_size,

From a5668962bcf9d84dc5b192db73f4b575e595172c Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 30 May 2019 12:41:59 -0700
Subject: [PATCH 2072/2720] Correct deltas, the lax version was wrong. Will
 come up with a better test.

PiperOrigin-RevId: 250741347
---
 tensor2tensor/trax/rlax/ppo.py | 19 +++++++------------
 1 file changed, 7 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 7241197d6..be1deb074 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -366,18 +366,13 @@ def deltas(predicted_values, rewards, mask, gamma=0.99):
     ndarray of shape (B, T) of one-step TD-residuals.
   """
 
-  v0 = np.transpose(predicted_values)[0]        # (B,) just V_{b, 1}
-  v1 = np.transpose(predicted_values)[:-1]      # (T, B) without V_{b, T+1}
-  rt = rewards.astype(np.float32).T   # (T, B)
-
-  def td_residual(carry, inps):
-    r, v_next = inps
-    v = carry
-    return v_next, (r + gamma * v_next - v)
-
-  _, d = lax.scan(td_residual, v0, [rt, v1])
-
-  return np.transpose(d) * mask
+  # Predicted values at time t, cutting off the last to have shape (B, T).
+  predicted_values_bt = predicted_values[:, :-1]
+  # Predicted values at time t+1, by cutting off the first to have shape (B, T)
+  predicted_values_btplus1 = predicted_values[:, 1:]
+  # Return the deltas as defined above.
+  return (
+      rewards + (gamma * predicted_values_btplus1) - predicted_values_bt) * mask
 
 
 def gae_advantages(td_deltas, mask, lambda_=0.95, gamma=0.99):

From ed518a1b3d01eea3168cc545e7d13d2dd833c7a9 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 30 May 2019 12:42:17 -0700
Subject: [PATCH 2073/2720] Set default eval batch size to 32 and double number
 of episode steps for higher scores in some games.

PiperOrigin-RevId: 250741406
---
 tensor2tensor/models/research/rl.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index f0262d94a..039d8809f 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -538,8 +538,10 @@ def rlmf_eval():
   """Eval set of hparams for model-free PPO."""
   hparams = rlmf_original()
   hparams.batch_size = 16
-  hparams.eval_sampling_temps = [0.0, 0.5, 1.0]
-  hparams.eval_rl_env_max_episode_steps = 20000
+  hparams.eval_batch_size = 32
+  hparams.eval_episodes_num = 2
+  hparams.eval_sampling_temps = [0.5, 0.0, 1.0]
+  hparams.eval_rl_env_max_episode_steps = 40000
   hparams.add_hparam("ppo_epoch_length", 128)
   hparams.add_hparam("ppo_optimization_batch_size", 32)
   hparams.add_hparam("ppo_epochs_num", 10000)

From 121ee60a3b57a092264aa5b5bf69ad194cafb118 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Thu, 30 May 2019 13:30:16 -0700
Subject: [PATCH 2074/2720] Adds MovingMnist to the T2T dataset generators
 using utilites from tensorflow-datasets wherever appropriate.

PiperOrigin-RevId: 250750540
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/moving_mnist.py | 154 ++++++++++++++++++
 2 files changed, 155 insertions(+)
 create mode 100644 tensor2tensor/data_generators/moving_mnist.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 0315d918f..79687e31b 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -54,6 +54,7 @@
     "tensor2tensor.data_generators.lm1b_imdb",
     "tensor2tensor.data_generators.lm1b_mnli",
     "tensor2tensor.data_generators.mnist",
+    "tensor2tensor.data_generators.moving_mnist",
     "tensor2tensor.data_generators.mrpc",
     "tensor2tensor.data_generators.mscoco",
     "tensor2tensor.data_generators.multinli",
diff --git a/tensor2tensor/data_generators/moving_mnist.py b/tensor2tensor/data_generators/moving_mnist.py
new file mode 100644
index 000000000..7d1502c67
--- /dev/null
+++ b/tensor2tensor/data_generators/moving_mnist.py
@@ -0,0 +1,154 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Moving MNIST dataset.
+
+Unsupervised Learning of Video Representations using LSTMs
+Nitish Srivastava, Elman Mansimov, Ruslan Salakhutdinov
+https://arxiv.org/abs/1502.04681
+
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import numpy as np
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import video_utils
+from tensor2tensor.layers import modalities
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+import tensorflow_datasets as tfds
+from tensorflow_datasets.video import moving_sequence
+
+
+DATA_URL = (
+    "http://www.cs.toronto.edu/~nitish/unsupervised_video/mnist_test_seq.npy")
+SPLIT_TO_SIZE = {
+    problem.DatasetSplit.TRAIN: 100000,
+    problem.DatasetSplit.EVAL: 10000,
+    problem.DatasetSplit.TEST: 10000}
+
+
+@registry.register_problem
+class VideoMovingMnist(video_utils.VideoProblem):
+  """MovingMnist Dataset."""
+
+  @property
+  def num_channels(self):
+    return 1
+
+  @property
+  def frame_height(self):
+    return 64
+
+  @property
+  def frame_width(self):
+    return 64
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  # num_videos * num_frames
+  @property
+  def total_number_of_frames(self):
+    return 100000 * 20
+
+  def max_frames_per_video(self, hparams):
+    return 20
+
+  @property
+  def random_skip(self):
+    return False
+
+  def eval_metrics(self):
+    return []
+
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each."""
+    return [
+        {"split": problem.DatasetSplit.TRAIN, "shards": 10},
+        {"split": problem.DatasetSplit.EVAL, "shards": 1},
+        {"split": problem.DatasetSplit.TEST, "shards": 1}]
+
+  @property
+  def extra_reading_spec(self):
+    """Additional data fields to store on disk and their decoders."""
+    data_fields = {
+        "frame_number": tf.FixedLenFeature([1], tf.int64),
+    }
+    decoders = {
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
+            tensor_key="frame_number"),
+    }
+    return data_fields, decoders
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    p.modality = {"inputs": modalities.ModalityType.VIDEO,
+                  "targets": modalities.ModalityType.VIDEO}
+    p.vocab_size = {"inputs": 256,
+                    "targets": 256}
+
+  def get_test_iterator(self, tmp_dir):
+    path = generator_utils.maybe_download(
+        tmp_dir, os.path.basename(DATA_URL), DATA_URL)
+    with tf.io.gfile.GFile(path, "rb") as fp:
+      mnist_test = np.load(fp)
+    mnist_test = np.transpose(mnist_test, (1, 0, 2, 3))
+    mnist_test = np.expand_dims(mnist_test, axis=-1)
+    mnist_test = tf.data.Dataset.from_tensor_slices(mnist_test)
+    return mnist_test.make_initializable_iterator()
+
+  def map_fn(self, image, label):
+    sequence = moving_sequence.image_as_moving_sequence(
+        image, sequence_length=20)
+    return sequence.image_sequence
+
+  def get_train_iterator(self):
+    mnist_ds = tfds.load("mnist", split=tfds.Split.TRAIN, as_supervised=True)
+    mnist_ds = mnist_ds.repeat()
+    moving_mnist_ds = mnist_ds.map(self.map_fn).batch(2)
+    moving_mnist_ds = moving_mnist_ds.map(lambda x: tf.reduce_max(x, axis=0))
+    return moving_mnist_ds.make_initializable_iterator()
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    with tf.Graph().as_default():
+      # train and eval set are generated on-the-fly.
+      # test set is the official test-set.
+      if dataset_split == problem.DatasetSplit.TEST:
+        moving_ds = self.get_test_iterator(tmp_dir)
+      else:
+        moving_ds = self.get_train_iterator()
+
+      next_video = moving_ds.get_next()
+      with tf.Session() as sess:
+        sess.run(moving_ds.initializer)
+
+        n_samples = SPLIT_TO_SIZE[dataset_split]
+        for _ in range(n_samples):
+          next_video_np = sess.run(next_video)
+          for frame_number, frame in enumerate(next_video_np):
+            yield {
+                "frame_number": [frame_number],
+                "frame": frame,
+            }

From c5a40dbd257c7f03ae1ad13dfbdc83a2264e1727 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 30 May 2019 14:05:19 -0700
Subject: [PATCH 2075/2720] Add basic shape tests for trax models.

PiperOrigin-RevId: 250757860
---
 tensor2tensor/trax/models/mlp_test.py         | 37 ++++++++++++++
 tensor2tensor/trax/models/resnet_test.py      | 43 ++++++++++++++++
 tensor2tensor/trax/models/transformer_test.py | 50 +++++++++++++++++++
 3 files changed, 130 insertions(+)
 create mode 100644 tensor2tensor/trax/models/mlp_test.py
 create mode 100644 tensor2tensor/trax/models/resnet_test.py
 create mode 100644 tensor2tensor/trax/models/transformer_test.py

diff --git a/tensor2tensor/trax/models/mlp_test.py b/tensor2tensor/trax/models/mlp_test.py
new file mode 100644
index 000000000..ae47e1cdf
--- /dev/null
+++ b/tensor2tensor/trax/models/mlp_test.py
@@ -0,0 +1,37 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 MLP."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+from tensor2tensor.trax import layers as tl
+from tensor2tensor.trax.models import mlp
+
+
+class MLPTest(absltest.TestCase):
+
+  def test_mlp(self):
+    input_shape = (3, 28, 28, 1)
+    model = mlp.MLP(d_hidden=32, n_output_classes=10)
+    final_shape = tl.check_shape_agreement(tl.Serial(model), input_shape)
+    self.assertEqual((3, 10), final_shape)
+
+
+if __name__ == '__main__':
+  absltest.main()
diff --git a/tensor2tensor/trax/models/resnet_test.py b/tensor2tensor/trax/models/resnet_test.py
new file mode 100644
index 000000000..81e17e869
--- /dev/null
+++ b/tensor2tensor/trax/models/resnet_test.py
@@ -0,0 +1,43 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 Resnet models."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+from tensor2tensor.trax import layers as tl
+from tensor2tensor.trax.models import resnet
+
+
+class ResnetTest(absltest.TestCase):
+
+  def test_resnet(self):
+    input_shape = (3, 256, 256, 3)
+    model = resnet.Resnet50(d_hidden=8, n_output_classes=10)
+    final_shape = tl.check_shape_agreement(tl.Serial(model), input_shape)
+    self.assertEqual((3, 10), final_shape)
+
+  def test_wide_resnet(self):
+    input_shape = (3, 32, 32, 3)
+    model = resnet.WideResnet(n_blocks=1, n_output_classes=10)
+    final_shape = tl.check_shape_agreement(tl.Serial(model), input_shape)
+    self.assertEqual((3, 10), final_shape)
+
+
+if __name__ == '__main__':
+  absltest.main()
diff --git a/tensor2tensor/trax/models/transformer_test.py b/tensor2tensor/trax/models/transformer_test.py
new file mode 100644
index 000000000..b8c33e450
--- /dev/null
+++ b/tensor2tensor/trax/models/transformer_test.py
@@ -0,0 +1,50 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 Transformer models."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+from tensor2tensor.trax import layers as tl
+from tensor2tensor.trax.models import transformer
+
+
+class TransformerTest(absltest.TestCase):
+
+  def test_transformer_lm(self):
+    vocab_size = 16
+    input_shape = [3, 5]
+    model = transformer.TransformerLM(
+        vocab_size, d_feature=32, d_feedforward=64, n_layers=2, n_heads=2)
+    final_shape = tl.check_shape_agreement(
+        tl.Serial(model), tuple(input_shape), integer_inputs=True)
+    self.assertEqual(tuple(input_shape + [vocab_size]), final_shape)
+
+  def test_transformer(self):
+    vocab_size = 16
+    single_input_shape = [3, 5]
+    input_shape = (tuple(single_input_shape), tuple(single_input_shape))
+    model = transformer.Transformer(
+        vocab_size, d_feature=32, d_feedforward=64, n_layers=2, n_heads=2)
+    final_shape = tl.check_shape_agreement(
+        tl.Serial(model), input_shape, integer_inputs=True)
+    self.assertEqual(tuple(single_input_shape + [vocab_size]), final_shape)
+
+
+if __name__ == '__main__':
+  absltest.main()

From 354995685843f0a5c65d690506059dbfb52846a6 Mon Sep 17 00:00:00 2001
From: Peng Wang <wangpeng@google.com>
Date: Thu, 30 May 2019 15:27:26 -0700
Subject: [PATCH 2076/2720] Type promotion, jit-argument canonicalization and
 various tf-numpy improvements that make trax/transformer work.

PiperOrigin-RevId: 250775245
---
 tensor2tensor/trax/backend.py          |  6 ++++++
 tensor2tensor/trax/layers/attention.py |  2 +-
 tensor2tensor/trax/optimizers/base.py  |  2 +-
 tensor2tensor/trax/trax.py             | 23 ++++++++++++++++-------
 4 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index 17526c473..df6a0d33e 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -30,6 +30,7 @@
 
 
 _JAX_BACKEND = {
+    "name": "jax",
     "np": jnp,
     "logsumexp": jax_special.logsumexp,
     "jit": jax.jit,
@@ -44,11 +45,16 @@
 
 
 _NUMPY_BACKEND = {
+    "name": "numpy",
     "np": onp,
     "jit": (lambda f: f),
 }
 
 
+def get_name():
+  return backend()["name"]
+
+
 def logsumexp(*args, **kwargs):
   return backend()["logsumexp"](*args, **kwargs)
 
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 30aef4877..4e2583e77 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -31,7 +31,7 @@
 def CausalMask(x, params, axis=-1, **kwargs):
   del params, kwargs
   size = x.shape[axis]
-  return onp.tril(onp.ones((1, size, size), dtype=x.dtype), k=0)
+  return onp.tril(onp.ones((1, size, size), dtype=onp.bool_), k=0)
 
 
 @base.layer(output_shape=lambda shape, pad=0: (shape[0], 1, 1, shape[-1]))
diff --git a/tensor2tensor/trax/optimizers/base.py b/tensor2tensor/trax/optimizers/base.py
index 8aea5f386..223a4e67a 100644
--- a/tensor2tensor/trax/optimizers/base.py
+++ b/tensor2tensor/trax/optimizers/base.py
@@ -19,7 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
-import jax.numpy as np
+from tensor2tensor.trax.backend import numpy as np
 
 from tensor2tensor.trax.layers import base as layers
 
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 0c6bde87b..1f8177408 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -293,9 +293,16 @@ def epochs(steps=None, epoch_steps=1):
       break
 
 
-def _jit_predict_fn(model_predict, n_devices):
+@gin.configurable
+def _jit_predict_fn(model_predict, n_devices, jit=True):
   """Use jit on model_predict if required."""
 
+  if n_devices == 1:
+    if jit:
+      return backend.jit(model_predict)
+    else:
+      return model_predict
+
   # Multi-devices, pmap and run.
   @functools.partial(backend.pmap, axis_name="batch")
   def mapped_predict(x, params, rng):
@@ -304,9 +311,6 @@ def mapped_predict(x, params, rng):
   def predict(x, params=(), rng=None):
     """Predict function jited and parallelized as requested."""
     # On one device, jit and run.
-    if n_devices == 1:
-      return backend.jit(model_predict)(x, params, rng=rng)
-
     pred = mapped_predict(
         reshape_by_device(x, n_devices),
         params,
@@ -322,7 +326,8 @@ def predict(x, params=(), rng=None):
   return predict
 
 
-def _jit_update_fn(predict_fn, loss_fn, optimizer, n_devices):
+@gin.configurable
+def _jit_update_fn(predict_fn, loss_fn, optimizer, n_devices, jit=True):
   """Get jit-ed update function for loss, optimizer, learning rate function."""
   if n_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
     def single_update(i, opt_state, batch, rng):
@@ -330,7 +335,10 @@ def single_update(i, opt_state, batch, rng):
       params, opt_slots = opt_state
       return optimizer.tree_update(i, backend.grad(loss_fn)(
           params, batch, predict_fn, rng), params, opt_slots), [subrng]
-    return backend.jit(single_update)
+    if jit:
+      return backend.jit(single_update)
+    else:
+      return single_update
 
   @functools.partial(backend.pmap, axis_name="batch")
   def mapped_update(i, opt_state, batch, rng):
@@ -530,7 +538,8 @@ def train(output_dir,
         history=history)
 
     # Save computation graph (single-device only for now).
-    if save_graphs and step == 1 and n_devices == 1:
+    if (save_graphs and backend.get_name() == "jax" and step == 1 and
+        n_devices == 1):
       params = opt_state[0]
       # Dump computation graphs to files.
       forward_computation = jax.xla_computation(model_predict_eval)(

From dfc84418d15ce4b04ef69ba3efa90273c2c438df Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 30 May 2019 16:47:32 -0700
Subject: [PATCH 2077/2720] Revert to the loop version of rewards_to_go, since
 it is faster.

PiperOrigin-RevId: 250789766
---
 tensor2tensor/trax/rlax/ppo.py | 60 ++++++++++++++++++++++++++--------
 1 file changed, 46 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index be1deb074..93a7d4180 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -283,23 +283,55 @@ def rewards_to_go(rewards, mask, gamma=0.99):
   Returns:
     rewards to go, np.ndarray of shape (B, T).
   """
-  # B, T = rewards.shape
+  B, T = rewards.shape  # pylint: disable=invalid-name,unused-variable
 
   masked_rewards = rewards * mask  # (B, T)
 
-  reversed_rewards = np.flip(masked_rewards, axis=1)  # (B, T) flipped on time.
-  rrt = np.transpose(reversed_rewards)  # (T, B) transpose to scan over time.
-
-  def discounting_add(carry, reward):
-    x = reward + (gamma * carry)
-    return x, x
-
-  _, ys = lax.scan(discounting_add,
-                   np.zeros_like(rrt[0], dtype=np.float32),
-                   rrt.astype(np.float32))
-
-  # ys is (T, B) and T is in reverse order.
-  return np.flip(np.transpose(ys), axis=1)
+  # The lax.scan version of this is slow, but we still show it here for
+  # completeness.
+  #   rewards_rev = np.flip(masked_rewards, axis=1)  # (B, T) flipped on time.
+  #   rrt = np.transpose(rewards_rev)  # (T, B) transpose to scan over time.
+  #
+  #   def discounting_add(carry, reward):
+  #     x = reward + (gamma * carry)
+  #     return x, x
+  #
+  #   _, ys = lax.scan(discounting_add,
+  #                    np.zeros_like(rrt[0], dtype=np.float32),
+  #                    rrt.astype(np.float32))
+  #
+  #   # ys is (T, B) and T is in reverse order.
+  #   return np.flip(np.transpose(ys), axis=1)
+
+  # We use the following recurrence relation, derived from the equation above:
+  #
+  # r2g[t+1] = (r2g[t] - r[t]) / gamma
+  #
+  # This means we'll need to calculate r2g[0] first and then r2g[1] and so on ..
+  #
+  # **However** this leads to overflows for long sequences: r2g[t] - r[t] > 0
+  # and gamma < 1.0, so the division keeps increasing.
+  #
+  # So we just run the recurrence in reverse, i.e.
+  #
+  # r2g[t] = r[t] + (gamma*r2g[t+1])
+  #
+  # This is much better, but might have lost updates since the (small) rewards
+  # at earlier time-steps may get added to a (very?) large sum.
+
+  # Compute r2g_{T-1} at the start and then compute backwards in time.
+  r2gs = [masked_rewards[:, -1]]
+
+  # Go from T-2 down to 0.
+  for t in reversed(range(T - 1)):
+    r2gs.append(masked_rewards[:, t] + (gamma * r2gs[-1]))
+
+  # The list should have length T.
+  assert T == len(r2gs)
+
+  # First we stack them in the correct way to make it (B, T), but these are
+  # still from newest (T-1) to oldest (0), so then we flip it on time axis.
+  return np.flip(np.stack(r2gs, axis=1), axis=1)
 
 
 @jit

From ebaef4732a2529a7f434fd31c0df0197364e6627 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 30 May 2019 21:30:28 -0700
Subject: [PATCH 2078/2720] Rename tests to make clear what is tested and allow
 backend changes without gin.

PiperOrigin-RevId: 250822120
---
 tensor2tensor/trax/backend.py                 | 6 +++++-
 tensor2tensor/trax/models/mlp_test.py         | 4 +++-
 tensor2tensor/trax/models/transformer_test.py | 7 +++++--
 3 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index df6a0d33e..4fc0431df 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -111,8 +111,12 @@ def __getattr__(self, attr):
 
 
 
+default_backend_name = "jax"
+
+
 @gin.configurable()
-def backend(name="jax"):
+def backend(name=None):
+  name = name or default_backend_name
   if name == "numpy":
     return _NUMPY_BACKEND
   return _JAX_BACKEND
diff --git a/tensor2tensor/trax/models/mlp_test.py b/tensor2tensor/trax/models/mlp_test.py
index ae47e1cdf..67e51bde4 100644
--- a/tensor2tensor/trax/models/mlp_test.py
+++ b/tensor2tensor/trax/models/mlp_test.py
@@ -20,13 +20,15 @@
 from __future__ import print_function
 
 from absl.testing import absltest
+from tensor2tensor.trax import backend
 from tensor2tensor.trax import layers as tl
 from tensor2tensor.trax.models import mlp
 
 
 class MLPTest(absltest.TestCase):
 
-  def test_mlp(self):
+  def test_mlp_forward_shape(self):
+    """Run the MLP model forward and check output shape."""
     input_shape = (3, 28, 28, 1)
     model = mlp.MLP(d_hidden=32, n_output_classes=10)
     final_shape = tl.check_shape_agreement(tl.Serial(model), input_shape)
diff --git a/tensor2tensor/trax/models/transformer_test.py b/tensor2tensor/trax/models/transformer_test.py
index b8c33e450..97574b2e4 100644
--- a/tensor2tensor/trax/models/transformer_test.py
+++ b/tensor2tensor/trax/models/transformer_test.py
@@ -20,13 +20,15 @@
 from __future__ import print_function
 
 from absl.testing import absltest
+from tensor2tensor.trax import backend
 from tensor2tensor.trax import layers as tl
 from tensor2tensor.trax.models import transformer
 
 
 class TransformerTest(absltest.TestCase):
 
-  def test_transformer_lm(self):
+  def test_transformer_lm_forward_shape(self):
+    """Run the Transformer LM forward and check output shape."""
     vocab_size = 16
     input_shape = [3, 5]
     model = transformer.TransformerLM(
@@ -35,7 +37,8 @@ def test_transformer_lm(self):
         tl.Serial(model), tuple(input_shape), integer_inputs=True)
     self.assertEqual(tuple(input_shape + [vocab_size]), final_shape)
 
-  def test_transformer(self):
+  def test_transformer_forward_shape(self):
+    """Run the Transformer forward and check output shape."""
     vocab_size = 16
     single_input_shape = [3, 5]
     input_shape = (tuple(single_input_shape), tuple(single_input_shape))

From 2da59d24eb9367cbed20c98df559beccd11b7582 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 30 May 2019 22:32:17 -0700
Subject: [PATCH 2079/2720] Add masked local n-D attention.

PiperOrigin-RevId: 250827437
---
 tensor2tensor/layers/common_attention.py      | 612 +++++++++++++++++-
 tensor2tensor/layers/common_attention_test.py | 286 +++++++-
 2 files changed, 896 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 7778f0d3d..7b38d6c87 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -20,6 +20,7 @@
 
 import collections
 import functools
+import itertools
 import math
 import operator
 
@@ -1146,6 +1147,22 @@ def split_heads_2d(x, num_heads):
   return tf.transpose(split_last_dimension(x, num_heads), [0, 3, 1, 2, 4])
 
 
+def split_heads_nd(x, num_heads):
+  """Split the depth dimension (last dimension) into multiple heads.
+
+  Args:
+    x: a [batch, d1, ..., dn, depth] tensor
+    num_heads: an integer
+
+  Returns:
+    a [batch, num_heads, d1, ..., dn, depth // num_heads]
+  """
+  num_dimensions = len(common_layers.shape_list(x)) - 2
+  return tf.transpose(
+      split_last_dimension(x, num_heads), [0, num_dimensions + 1] +
+      list(range(1, num_dimensions + 1)) + [num_dimensions + 2])
+
+
 @expert_utils.add_name_scope()
 def combine_heads(x):
   """Inverse of split_heads.
@@ -1173,6 +1190,21 @@ def combine_heads_2d(x):
   return combine_last_two_dimensions(tf.transpose(x, [0, 2, 3, 1, 4]))
 
 
+def combine_heads_nd(x):
+  """Inverse of split_heads_nd.
+
+  Args:
+    x: a [batch, num_heads, d1, ..., dn, depth // num_heads] tensor
+
+  Returns:
+    a [batch, d1, ...., dn, depth] tensor
+  """
+  num_dimensions = len(common_layers.shape_list(x)) - 3
+  return combine_last_two_dimensions(
+      tf.transpose(x, [0] + list(range(2, num_dimensions + 2)) +
+                   [1, num_dimensions + 2]))
+
+
 def attention_image_summary(attn, image_shapes=None):
   """Compute color image summary.
 
@@ -3773,6 +3805,27 @@ def right_shift_blockwise(x, query_shape, name=None):
     return output
 
 
+def right_shift_blockwise_nd(x, block_shape):
+  """Right shift once in every block.
+
+  Args:
+    x: a [batch, d1, d2, ..., dn, depth] tensor
+    block_shape: a tuple (q1, q2, ..., qn) representing the block shape
+
+  Returns:
+    a [batch, d1, d2, ..., dn, depth] tensor, right shifted.
+  """
+  blocked_x = break_into_blocks_nd(x, block_shape)
+  blocked_x_shape = common_layers.shape_list(blocked_x)
+  blocked_x = tf.reshape(blocked_x,
+                         [blocked_x_shape[0], -1, blocked_x_shape[-1]])
+  padded_x = tf.pad(blocked_x, [[0, 0], [1, 0], [0, 0]])
+  x = tf.slice(padded_x, [0, 0, 0],
+               [-1, np.prod(blocked_x_shape[1:-1], dtype=np.int32), -1])
+  x = tf.reshape(x, blocked_x_shape)
+  return put_back_blocks_nd(x, block_shape)
+
+
 def masked_local_attention_2d(q,
                               k,
                               v,
@@ -3866,6 +3919,363 @@ def masked_local_attention_2d(q,
     return output
 
 
+def masked_local_attention_nd(q,
+                              k,
+                              v,
+                              query_shape,
+                              memory_flange,
+                              decode_step=None,
+                              name=None):
+  """Masked local attention nd.
+
+  Each position in q can attend to positions in memory that are positioned less
+  than or equal to query position according to raster scan ordering and are in
+  the same memory block. A memory block is n-dimensional and each dimension 'i'
+  is of size q[i] + 2 * m[i] except for the first dimension which is of size
+  q[0] + m[0]. NOTE: This computation assumes memory_flange is divisible by
+  query_shape in every dimension.
+
+  Args:
+    q: a [batch, heads, d1, d2, ..., dn, depth_k] tensor or a [batch, heads, 1,
+      1, ..., 1, depth_k] tensor in decoding mode.
+    k: a [batch, heads, d1, d2, ..., dn, depth_k] tensor
+    v: a [batch, heads, d1, d2, ..., dn, depth_v] tensor
+    query_shape: a tuple (q1, q2, ..., qn) indicating the shape of query blocks.
+    memory_flange: a tuple (m1, m2, ..., mn) indicating the number of extra
+      positions in the attention memory. memory_shape=[q1 + m1, d2 + 2 * m2,
+      ..., dn + 2 * mn]
+    decode_step: an integer in fast decoding mode.
+    name: an optional string
+
+  Returns:
+    a [batch, head, d1, d2, ..., dn, depth_v] tensor or
+      [batch, head, 1, 1, ..., 1, depth_v] if decode_step is not None.
+  """
+  assert all([m % b == 0 for m, b in zip(memory_flange, query_shape)])
+  with tf.variable_scope(
+      name, default_name="masked_local_attention_nd", values=[q, k, v]):
+    # This computation only applies to self attention, so assert q, k and v have
+    # the same dimensions.
+    if decode_step is None:
+      q.get_shape().assert_is_compatible_with(k.get_shape())
+      q.get_shape()[:-1].assert_is_compatible_with(v.get_shape()[:-1])
+    else:
+      k.get_shape().assert_is_compatible_with(v.get_shape())
+
+    # move heads to batch dimension. This is needed to reduce number of
+    # dimensions as much as possible, since most ops support only up to 7
+    # dimensions.
+    q_shape = common_layers.shape_list(q)
+    k_shape = common_layers.shape_list(k)
+    v_shape = common_layers.shape_list(v)
+    q = tf.reshape(q, [-1] + q_shape[2:])
+    k = tf.reshape(k, [-1] + k_shape[2:])
+    v = tf.reshape(v, [-1] + v_shape[2:])
+
+    # Pad query, key, value to ensure multiple of corresponding lengths.
+    if decode_step is None:
+      # don't pad query in fast decoding mode. We only need to calculate self
+      # attention for one position.
+      q = pad_to_multiple_nd(q, query_shape)
+    k = pad_to_multiple_nd(k, query_shape)
+    v = pad_to_multiple_nd(v, query_shape)
+
+    # extract query and memory blocks
+    if decode_step is None:
+      q = break_into_blocks_nd(q, query_shape)
+    else:
+      # in fast decoding, q has 1 block with 1 item in it
+      # q shape will be [batch] + [1] * n + [1, depth] which is equivalent of
+      # [batch, b1, b2, ..., bn, items_in_block, depth] where there is 1 block
+      # and 1 item in that block
+      q = tf.reshape(q, [-1] + [1] * (len(q_shape) - 3) + [q_shape[-1]])
+    k = break_into_memory_blocks_nd(k, query_shape, memory_flange, masked=True)
+    v = break_into_memory_blocks_nd(v, query_shape, memory_flange, masked=True)
+
+    # extract just one block of k and v in fast decoding mode.
+    if decode_step is not None:
+      k = select_block_for_decode_step(k, decode_step, query_shape)
+      v = select_block_for_decode_step(v, decode_step, query_shape)
+
+    # flatten q, k and v to [batch, num_blocks, items_in_block, depth]
+    q, blocks_per_dim = flatten_blocks_nd(q)
+    k, _ = flatten_blocks_nd(k)
+    v, _ = flatten_blocks_nd(v)
+
+    # make attention bias for causal attention.
+    causal_attn_bias = causal_attention_bias_nd(
+        query_shape, memory_flange, decode_step=decode_step)
+    padding_attn_bias = tf.expand_dims(
+        embedding_to_padding(v[:1, :, :, :]) * -1e9, axis=-2)
+
+    if decode_step is None:
+      num_blocks = common_layers.shape_list(v)[1]
+      causal_attn_bias = tf.tile(causal_attn_bias, [1, num_blocks, 1, 1])
+      padding_attn_bias = tf.tile(
+          padding_attn_bias,
+          [1, 1, np.prod(query_shape, dtype=np.int32), 1])
+    attn_bias = tf.minimum(causal_attn_bias, padding_attn_bias)
+
+    # Calculate dot product attention
+    output = dot_product_attention(
+        q,
+        k,
+        v,
+        attn_bias,
+        dropout_rate=0.,
+        name=name or "masked_local_nd",
+        make_image_summary=False)
+
+    # restructure the output from blocks ordering to the original ordering
+    output = unflatten_blocks_nd(output, blocks_per_dim)
+    if decode_step is None:
+      # In fast decoding, output only contains one element, this is not needed.
+      output = put_back_blocks_nd(output, query_shape)
+
+    # bring back the heads dimension
+    output_shape = common_layers.shape_list(output)
+    output = tf.reshape(output, q_shape[:2] + output_shape[1:])
+    if decode_step is None:
+      # No padding is introduced in fast decoding, no need to do this.
+      output_shape = common_layers.shape_list(output)
+      output = tf.slice(output, [0] * len(output_shape),
+                        [-1, -1] + q_shape[2:-1] + [-1])
+    return output
+
+
+def select_block_for_decode_step(blocked_x, decode_step, query_shape):
+  """Selects one block from `x` that contains position `decode_step`.
+
+  NOTE: This method only works for blocked inputs. It selects one block around
+  `decode_step` position in blocked raster scan order.
+
+  Args:
+    blocked_x: a [batch, blocks_per_d1, ..., blocks_per_dn, b1 * ...* bn, depth]
+      tensor
+    decode_step: an integer
+    query_shape: a tuple (q1, q2, ..., qn) representing query shape
+
+  Returns:
+     a [batch, [1] * n, b1 * ... * bn, depth] tensor
+  """
+  blocked_x_shape = common_layers.shape_list(blocked_x)
+  # calculate the shape of the normal x
+  x_shape = [b * q for b, q in zip(blocked_x_shape[1:-2], query_shape)]
+  # Get the position of `decode_step` element in the unblocked x.
+  index = decode_step_to_index(decode_step, query_shape, x_shape)
+  # Convert it to the blocked positions.
+  blocked_index = [i // q for i, q in zip(index, query_shape)]
+  # TPU needs size to be non negative for the case when begin is not
+  # compile-time constants.
+  return tf.slice(blocked_x, [0] + blocked_index + [0, 0],
+                  [blocked_x_shape[0]] + [1] * len(blocked_index) +
+                  blocked_x_shape[-2:])
+
+
+def flatten_blocks_nd(x):
+  """Flattens blocks of the input tensor.
+
+  Args:
+    x: a [batch, b1, ..., bn, items_in_block, depth] tensor
+
+  Returns:
+    a flattened tensor of shape [batch, b1 * ...* bm, items_in_block, depth]
+    a list of [b1, ..., bn] which is used for unflattening.
+  """
+  x_shape = common_layers.shape_list(x)
+  num_blocks = np.prod(x_shape[1:-2], dtype=np.int32)
+  return tf.reshape(x, [-1, num_blocks] + x_shape[-2:]), x_shape[1:-2]
+
+
+def unflatten_blocks_nd(x, blocks_per_dimension):
+  """Converts a flattened tensor into a normal blocked tensor.
+
+  Args:
+    x: a [batch, d1 * ... dn, items_in_block, depth] tensor
+    blocks_per_dimension: a n-d list of integers for number of blocks in each
+      dimension.
+
+  Returns:
+    a [batch, d1, d2, ..., dn, items_in_block, depth] tensor
+  """
+  x_shape = common_layers.shape_list(x)
+  assert x_shape[1] == np.prod(blocks_per_dimension, dtype=np.int32)
+  return tf.reshape(x, [-1] + list(blocks_per_dimension) + x_shape[-2:])
+
+
+def break_into_memory_blocks_nd(x, query_shape, memory_flange, masked=False):
+  """Break a tensor into memory blocks around query blocks.
+
+  This requires memory_flange to be divisible by query_shape in every dimension.
+
+  Args:
+    x: a [batch, d1, d2, ..., dn, depth] tensor
+    query_shape: a n-d list of integers representing query shape
+    memory_flange: an n-d list of integers representing memory flange.
+    masked: a boolean for masked vs unmasked attention.
+
+  Returns:
+    a [batch, blocks_per_d1, ..., blocks_per_dn, b1 * ...* bn, depth] where bi
+      is the memory block size in dimension i which is equal to q[i] + 2m[i] or
+      q[i] + m[i] if masked attention and i = 1.
+  """
+  assert all([m % b == 0 for b, m in zip(query_shape, memory_flange)])
+
+  original_x_shape = common_layers.shape_list(x)
+  # calculate the total number of query blocks in each dimension
+  blocks_in_memory_flange = [m // b for b, m in zip(query_shape, memory_flange)]
+  num_query_blocks = [
+      l // q for l, q in zip(original_x_shape[1:-1], query_shape)
+  ]
+  # pad x to have enough items on the corners to form the  memory blocks.
+  if masked:
+    # Only pad the beginning of first dimension in masked mode.
+    x = tf.pad(x, [[0, 0], [memory_flange[0], 0]] +
+               [[p, p] for p in memory_flange[1:]] + [[0, 0]])
+  else:
+    x = tf.pad(x, [[0, 0]] + [[p, p] for p in memory_flange] + [[0, 0]])
+
+  query_blocks = break_into_blocks_nd(x, query_shape)
+  # stitch query blocks together to form memory blocks of the desired size.
+  start_indices_per_dimension = []
+  for dimension, blocks in enumerate(blocks_in_memory_flange):
+    if masked and dimension == 0:
+      # num blocks for first dimension in masked mode is blocks + 1
+      size = blocks + 1
+    else:
+      size = 2 * blocks + 1
+    start_indices_per_dimension.append(range(size))
+
+  slices = []
+  for start_indices in itertools.product(*start_indices_per_dimension):
+    start = [0] + list(start_indices) + [0, 0]
+    size = [-1] + num_query_blocks + [-1, -1]
+    s = tf.slice(query_blocks, start, size)
+    slices.append(s)
+  # concat slices in their query block dimension to form the full memory blocks
+  return tf.concat(slices, axis=-2)
+
+
+def break_into_blocks_nd(x, block_shape):
+  """Break input tensor into blocks of `block_shape`.
+
+  Args:
+    x: a [batch, d1, d2, ..., dn, depth] tensor
+    block_shape: a n-d list of integers representing block shape
+
+  Returns:
+    a [batch, d1//block1, ..., dn//blockn, block1 *... * blockn, depth] tensor
+  """
+  x_shape = common_layers.shape_list(x)
+  assert all([l % b == 0 for l, b in zip(x_shape[1:], block_shape)])
+  blocks_per_dimension = [l // b for l, b in zip(x_shape[1:], block_shape)]
+  # reshape to [-1, d1 // block1, block1, ..., dn // blockn, blockn, depth]
+  reshape_to = list(
+      itertools.chain.from_iterable(zip(blocks_per_dimension, block_shape)))
+  x = tf.reshape(x, [-1] + reshape_to + x_shape[-1:])
+  # transpose dimensions to bring the n-d blocks in consecutive dimensions.
+  block_dimensions_index = [2 * (i + 1) for i in range(len(block_shape))]
+  x = tf.transpose(x, [0] + [i - 1 for i in block_dimensions_index] +
+                   block_dimensions_index + [2 * len(block_shape) + 1])
+  return tf.reshape(x, [-1] + blocks_per_dimension +
+                    [np.prod(block_shape, dtype=np.int32)] + x_shape[-1:])
+
+
+def put_back_blocks_nd(x, block_shape):
+  """Restructure input tensor from blocks to normal ordering.
+
+  Args:
+    x: a [batch, b1, ..., bn, items_in_block, depth] tensor
+    block_shape: a n-d list of integers representing block shape.
+
+  Returns:
+    a [batch, d1, ..., dn, depth] where blocks are put back to form the
+      original tensor.
+  """
+  x_shape = common_layers.shape_list(x)
+  assert x_shape[-2] == np.prod(block_shape)
+  x = tf.reshape(x, x_shape[:-2] + list(block_shape) + x_shape[-1:])
+  block_dimension_index = [i + 1 for i in range(len(block_shape))]
+  block_shape_index = [b + len(block_shape) for b in block_dimension_index]
+  interleaved_dimensions = list(
+      itertools.chain.from_iterable(
+          zip(block_dimension_index, block_shape_index)))
+  x = tf.transpose(x, [0] + interleaved_dimensions + [2 * len(block_shape) + 1])
+  x_shape = common_layers.shape_list(x)
+  x = tf.reshape(x, [-1] + [
+      x_shape[2 * i + 1] * x_shape[2 * i + 2] for i in range(len(block_shape))
+  ] + x_shape[-1:])
+  return x
+
+
+def pad_to_multiple_nd(x, block_shape):
+  """Making sure x is a multiple of shape.
+
+  Args:
+    x: a [batch, d1, d2, ..., dn, depth] tensor
+    block_shape: a n-d list of integers representing block shape
+
+  Returns:
+    padded x where each dimension is a multiple of corresponding block length.
+  """
+  shape = common_layers.shape_list(x)
+  paddings = [-l % b for l, b in zip(shape[1:-1], block_shape)]
+  return tf.pad(x, [[0, 0]] + [[0, p] for p in paddings] + [[0, 0]])
+
+
+def causal_attention_bias_nd(query_shape, memory_flange, decode_step=None):
+  """Creates causal attention bias for local nd attention.
+
+  This assumes memory_flange is divisible by query_shape in every dimension.
+
+  Args:
+    query_shape: a n-d list of integers representing query shape
+    memory_flange: a n-d list of integers representing memory flange
+    decode_step: an integer
+
+  Returns:
+    a [1, 1, query_items, memory_items] tensor for masked attention bias or
+    a [1, 1, 1, memory_items] tensor if decode_step is not None.
+  """
+  assert all([m % q == 0 for q, m in zip(query_shape, memory_flange)])
+  blocks_per_memory_flange = [
+      m // q for q, m in zip(query_shape, memory_flange)
+  ]
+  # previous blocks will be half the number of all blocks if we select blocks
+  # to the left and right of center block in every dimension.
+  prev_blocks = np.prod([2 * b + 1 for b in blocks_per_memory_flange],
+                        dtype=np.int32) // 2
+  all_blocks = np.prod(
+      [blocks_per_memory_flange[0] + 1] +
+      [2 * b + 1 for b in blocks_per_memory_flange[1:]],
+      dtype=np.int32)
+  future_blocks = all_blocks - prev_blocks - 1
+  # add unmasked biases for all prev blocks and a lower triangle for the center
+  # block and all masked for future blocks.
+  items_in_block = np.prod(query_shape, dtype=np.int32)
+  items_in_query = items_in_block if decode_step is None else 1
+  prev_blocks_attn = tf.zeros(
+      [1, 1, items_in_query, prev_blocks * items_in_block])
+
+  # add mask for the center block
+  if decode_step is None:
+    center_block_attn = attention_bias_lower_triangle(items_in_block)
+  else:
+    step_in_block = decode_step % items_in_block
+    cond = tf.reshape(
+        tf.less_equal(tf.range(items_in_block, dtype=tf.int32), step_in_block),
+        [1, 1, items_in_query, items_in_block])
+    center_block_attn = tf.where(
+        cond, tf.zeros([1, 1, items_in_query, items_in_block]),
+        -1e9 * tf.ones([1, 1, items_in_query, items_in_block]))
+
+  # add mask for all future blocks
+  future_blocks_attn = -1e9 * tf.ones(
+      [1, 1, items_in_query, future_blocks * items_in_block])
+  return tf.concat([prev_blocks_attn, center_block_attn, future_blocks_attn],
+                   axis=3)
+
+
 def compute_attention_component(antecedent,
                                 total_depth,
                                 filter_width=1,
@@ -3873,7 +4283,7 @@ def compute_attention_component(antecedent,
                                 name="c",
                                 vars_3d_num_heads=0,
                                 layer_collection=None):
-  """Computes attention compoenent (query, key or value).
+  """Computes attention component (query, key or value).
 
   Args:
     antecedent: a Tensor with shape [batch, length, channels]
@@ -4387,6 +4797,206 @@ def multihead_attention_2d(query_antecedent,
     return x
 
 
+def multihead_attention_nd(query_antecedent,
+                           memory_antecedent,
+                           total_key_depth,
+                           total_value_depth,
+                           output_depth,
+                           num_heads,
+                           query_shape,
+                           memory_flange,
+                           masked=False,
+                           cache=None,
+                           decode_step=None,
+                           name=None):
+  """n-d Multihead scaled-dot-product attention with in/output transformations.
+
+  Args:
+    query_antecedent: a Tensor with shape [batch, d1, ..., dn, depth_q] or
+      [batch, 1, ..., 1, depth_q] if in fast decoding mode.
+    memory_antecedent: a Tensor with shape [batch, d1, ..., dn, depth_m] or None
+      for self attention.
+    total_key_depth: an integer
+    total_value_depth: an integer
+    output_depth: an integer
+    num_heads: an integer dividing total_key_depth and total_value_depth
+    query_shape: an tuple indicating the dimensions of each query block.
+    memory_flange: an integer indicating how much to look around a query block
+      in each dimension
+    masked: a boolean to specify whether to do masked or unmasked attention.
+    cache: a dict like: {
+      'k': [batch, num_heads, d1, ..., dn, depth_k // num_heads],
+      'v': [batch, num_heads, d1, ..., dn, depth_v // num_heads]} Caller should
+        initially pass zero tensors for `decode_step` == 0. This method will
+        update cache and caller should pass the same cache in consecutive calls.
+        This works for both GPU and TPU inference. Caller should pass the latest
+        query via `query_antecedent`. `memory_antecedent` should be None in this
+        case, since auto-regressive decoding only applies to self attention.
+    decode_step: integer to pass in decoding mode. `cache` and `decode_step`
+      should both be set in decoding mode. Caller can also pass an empty `cache`
+      without `decode_step`, for this method to initialize the cache for future
+      calls with `decode_step` > 0.
+    name: an optional string
+
+  Returns:
+    A Tensor of shape [batch, d1, ..., dn, output_depth] or
+    [batch, 1, ..., 1, output_depth] if decode_step is set.
+
+  Raises:
+    ValueError: if the key depth or value depth are not divisible by the
+      number of attention heads.
+  """
+  if total_key_depth % num_heads != 0:
+    raise ValueError("Key depth (%d) must be divisible by the number of "
+                     "attention heads (%d)." % (total_key_depth, num_heads))
+  if total_value_depth % num_heads != 0:
+    raise ValueError("Value depth (%d) must be divisible by the number of "
+                     "attention heads (%d)." % (total_value_depth, num_heads))
+  # Validate decoding input params are sensible.
+  if decode_step is not None:
+    assert "k" in cache and "v" in cache
+  if cache is not None:
+    assert memory_antecedent is None
+
+  with tf.variable_scope(
+      name,
+      default_name="multihead_attention_nd",
+      values=[query_antecedent, memory_antecedent]):
+    if decode_step is not None:
+      latest_antecedent = query_antecedent
+      q, latest_k, latest_v = compute_qkv(latest_antecedent, None,
+                                          total_key_depth, total_value_depth)
+      latest_k = split_heads_nd(latest_k, num_heads)
+      latest_v = split_heads_nd(latest_v, num_heads)
+      # put latest k and v into their correct position in cache.
+      k = cache["k"]
+      v = cache["v"]
+      k = put_item_in_decode_step(k, latest_k, decode_step, query_shape)
+      v = put_item_in_decode_step(v, latest_v, decode_step, query_shape)
+      cache["k"] = k
+      cache["v"] = v
+
+    else:
+      q, k, v = compute_qkv(query_antecedent, memory_antecedent,
+                            total_key_depth, total_value_depth)
+      k = split_heads_nd(k, num_heads)
+      v = split_heads_nd(v, num_heads)
+      if cache is not None:
+        cache["k"] = k
+        cache["v"] = v
+    # after splitting, shape is [batch, heads, d1, ..., dn, depth]
+    q = split_heads_nd(q, num_heads)
+    key_depth_per_head = total_key_depth // num_heads
+    q *= key_depth_per_head**-0.5
+    if masked:
+      x = masked_local_attention_nd(
+          q,
+          k,
+          v,
+          query_shape=query_shape,
+          memory_flange=memory_flange,
+          decode_step=decode_step)
+    else:
+      raise NotImplementedError(
+          "Unmaked multihead attention nd is not implemented")
+    x = combine_heads_nd(x)
+    x = common_layers.dense(
+        x, output_depth, use_bias=False, name="output_transform")
+    return x
+
+
+def decode_step_to_index(decode_step, query_shape, tensor_shape):
+  """Maps decode step to n-d index according to blocked raster scan order.
+
+  Args:
+    decode_step: an integer
+    query_shape: a tuple (q1, q2, ..., qn) representing the query shape
+    tensor_shape: a tuple (d1, d2, ..., dn) representing the tensor shape, minus
+      the batch and depth dimensions.
+
+  Returns:
+    a tuple (i1, i2, ..., in) representing the index of the element at
+    `decode_step` w.r.t. blocked raster scan order.
+  """
+  assert len(query_shape) == len(tensor_shape)
+  blocks_per_dimension = [t // q for t, q in zip(tensor_shape, query_shape)]
+  items_in_block = np.prod(query_shape, dtype=np.int32)
+  step_block = decode_step // items_in_block
+  step_within_block = decode_step % items_in_block
+
+  block_index = []
+  for q in blocks_per_dimension[::-1]:
+    block_index.insert(0, step_block % q)
+    step_block //= q
+
+  within_block_index = []
+  for q in query_shape[::-1]:
+    within_block_index.insert(0, step_within_block % q)
+    step_within_block //= q
+
+  final_index = [
+      w + b * q for w, b, q in zip(within_block_index, block_index, query_shape)
+  ]
+  return tuple(final_index)
+
+
+def get_item_at_decode_step(x, decode_step, query_shape):
+  """Extracts a single item from an n-d tensor at `decode_step` position.
+
+  Args:
+    x: a [batch, d1, d2, ..., dn, depth] tensor
+    decode_step: an integer
+    query_shape: a tuple (q1, q2, ..., qn) representing the query shape
+
+  Returns:
+    a [batch, 1, 1, ..., 1, depth] tensor that is a single element from `x` at
+    `decode_step` w.r.t. blocked raster scan order.
+  """
+  x_shape = common_layers.shape_list(x)
+  index = decode_step_to_index(decode_step, query_shape, x_shape[1:-1])
+  # TPU needs size to be non negative for the case when begins are not
+  # compile-time constants.
+  return tf.slice(x, [0] + list(index) + [0],
+                  [x_shape[0]] + [1] * len(index) + [x_shape[-1]])
+
+
+def put_item_in_decode_step(x, item, decode_step, query_shape):
+  """Puts a single item into an n-d tensor at `decode_step` position.
+
+  Args:
+    x: a [batch, heads, d1, d2, ..., dn, depth] tensor
+    item: a [batch, heads, 1, 1, ..., 1, depth] tensor
+    decode_step: an integer
+    query_shape: a tuple (q1, q2, ..., qn) representing the query shape
+
+  Returns:
+    a [batch, heads, d1, d2, ..., dn, depth] tensor with value at `decode_step`
+    w.r.t. blocked raster scan order is updated to be `item`.
+  """
+  x_shape = common_layers.shape_list(x)
+  index = decode_step_to_index(decode_step, query_shape, x_shape[2:-1])
+  # inplace_update only works on the first dimension, we need to flatten and
+  # move batch to be the second dimension.
+  flattened_x = tf.reshape(
+      x, [-1, x_shape[1], np.prod(x_shape[2:-1]), x_shape[-1]])
+  # transpose to [positions, batch, heads, depth]
+  flattened_x = tf.transpose(flattened_x, [2, 0, 1, 3])
+
+  flattened_index = 0
+  factor = 1
+  for d, idx in zip(x_shape[-2:1:-1], index[::-1]):
+    flattened_index += idx * factor
+    factor *= d
+
+  item_shape = common_layers.shape_list(item)
+  item = tf.reshape(item, item_shape[:2] + item_shape[-1:])
+  updated_x = inplace_ops.alias_inplace_update(flattened_x, flattened_index,
+                                               item)
+  # unflatten the results
+  updated_x = tf.transpose(updated_x, [1, 2, 0, 3])
+  return tf.reshape(updated_x, [-1, x_shape[1]] + x_shape[2:])
+
+
 def ffn_self_attention_layer(x,
                              filter_depth,
                              output_depth,
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index fe31f4dd7..f747059b8 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -98,6 +98,291 @@ def testDotProductAttention(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
+  @parameterized.parameters(
+      ([3, 10, 64], 4),
+      ([3, 10, 20, 64], 2),
+      ([3, 10, 20, 30, 64], 4),
+  )
+  def testSplitHeadsND(self, shape, num_heads):
+    t = tf.zeros(shape)
+    h = common_attention.split_heads_nd(t, num_heads)
+    res = self.evaluate(h)
+    self.assertEqual(
+        res.shape,
+        tuple(shape[:1] + [num_heads] + shape[1:-1] + [shape[-1] // num_heads]))
+
+  @parameterized.parameters(
+      ([3, 4, 10, 64],),
+      ([3, 2, 10, 20, 64],),
+      ([3, 4, 10, 20, 30, 64],),
+  )
+  def testCombineHeadsND(self, shape):
+    t = tf.zeros(shape)
+    h = common_attention.combine_heads_nd(t)
+    res = self.evaluate(h)
+    self.assertEqual(res.shape,
+                     tuple(shape[:1] + shape[2:-1] + [shape[-1] * shape[1]]))
+
+  @parameterized.parameters(
+      ([3, 4, 10, 64], (5,), (10,)),
+      ([3, 4, 10, 10, 64], (5, 5), (5, 5)),
+      ([3, 4, 10, 10, 10, 64], (5, 5, 5), (5, 5, 5)),
+  )
+  def testShapeMaskedLocalAttentionND(self, shape, query_shape, memory_flange):
+    q = k = v = tf.reshape(tf.range(np.prod(shape), dtype=tf.float32), shape)
+    val = common_attention.masked_local_attention_nd(q, k, v, query_shape,
+                                                     memory_flange)
+    res = self.evaluate(val)
+    self.assertEqual(res.shape, tuple(shape))
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testRightShiftBlockwiseND(self):
+    tensor = tf.convert_to_tensor(np.array([[
+        [[1], [2], [3], [4]],
+        [[5], [6], [7], [8]],
+        [[9], [10], [11], [12]],
+        [[13], [14], [15], [16]],
+    ]], dtype=np.float32))
+    val = common_attention.right_shift_blockwise_nd(tensor, (2, 2))
+    res = self.evaluate(val)
+    expected_val = np.array([[
+        [[0], [1], [6], [3]],
+        [[2], [5], [4], [7]],
+        [[8], [9], [14], [11]],
+        [[10], [13], [12], [15]],
+    ]], dtype=np.float32)
+    self.assertAllClose(expected_val, res)
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testContentMaskedLocalAttentionND(self):
+    def softmax(arr):
+      return np.exp(arr) / np.sum(np.exp(arr))
+
+    q = k = v = tf.convert_to_tensor(
+        np.array([[[
+            [[0.1], [0.1], [0.1], [0.1]],
+            [[0.1], [1.0], [1.0], [0.1]],
+            [[0.1], [1.0], [1.0], [0.1]],
+            [[0.1], [0.1], [0.1], [0.1]],
+        ]]], dtype=np.float32))
+    attn_weights = np.array([[[[softmax([-1e9, -1e9, -1e9, -1e9, 0.01]),
+                                softmax([-1e9, -1e9, -1e9, 0.01, 0.01]),
+                                softmax([-1e9, -1e9, -1e9, 0.01, 0.01]),
+                                softmax([-1e9, -1e9, -1e9, 0.01, 0.01])
+                               ],
+                               [softmax([-1e9, 0.01, 0.01, -1e9, 0.01]),
+                                softmax([0.1, 0.1, 0.1, 0.1, 1.0]),
+                                softmax([0.1, 0.1, 0.1, 1.0, 1.0]),
+                                softmax([0.01, 0.01, -1e9, 0.1, 0.01])
+                               ],
+                               [softmax([-1e9, 0.01, 0.1, -1e9, 0.01]),
+                                softmax([0.1, 1.0, 1.0, 0.1, 1.0]),
+                                softmax([1.0, 1.0, 0.1, 1.0, 1.0]),
+                                softmax([0.1, 0.01, -1e9, 0.1, 0.01])
+                               ],
+                               [softmax([-1e9, 0.01, 0.1, -1e9, 0.01]),
+                                softmax([0.01, 0.1, 0.1, 0.01, 0.01]),
+                                softmax([0.1, 0.1, 0.01, 0.01, 0.01]),
+                                softmax([0.1, 0.01, -1e9, 0.01, 0.01])
+                               ]]]])
+    blocked_v = np.array([[[[[0, 0, 0, 0, 0.1],
+                             [0, 0, 0, 0.1, 0.1],
+                             [0, 0, 0, 0.1, 0.1],
+                             [0, 0, 0, 0.1, 0.1]],
+                            [[0, 0.1, 0.1, 0, 0.1],
+                             [0.1, 0.1, 0.1, 0.1, 1],
+                             [0.1, 0.1, 0.1, 1, 1],
+                             [0.1, 0.1, 0, 1, 0.1]],
+                            [[0, 0.1, 1, 0, 0.1],
+                             [0.1, 1, 1, 0.1, 1],
+                             [1, 1, 0.1, 1, 1],
+                             [1, 0.1, 0, 1, 0.1]],
+                            [[0, 0.1, 1, 0, 0.1],
+                             [0.1, 1, 1, 0.1, 0.1],
+                             [1, 1, 0.1, 0.1, 0.1],
+                             [1, 0.1, 0, 0.1, 0.1]]]]])
+    expected_val = np.expand_dims(
+        np.sum(attn_weights * blocked_v, axis=4), axis=-1)
+    val = common_attention.masked_local_attention_nd(q, k, v, (1, 1), (1, 1))
+    res = self.evaluate(val)
+    self.assertAllClose(expected_val, res)
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testSelectBlockForDecodeStep(self):
+    tensor = tf.reshape(
+        tf.range(2 * 6 * 6 * 4, dtype=tf.float32), [2, 6, 6, 4, 1])
+    block = common_attention.select_block_for_decode_step(tensor, 20, (2, 2))
+    expected_tensor = tensor[:, 0:1, 5:6, :, :]
+    expected_value = self.evaluate(expected_tensor)
+    res = self.evaluate(block)
+    self.assertAllClose(expected_value, res)
+
+  @parameterized.parameters(
+      ((2, 6, 4, 10),),
+      ((2, 6, 6, 4, 10),),
+      ((2, 6, 6, 6, 4, 10),),
+  )
+  def testFlattenBlocksND(self, shape):
+    tensor = tf.zeros(shape, dtype=tf.float32)
+    value, _ = common_attention.flatten_blocks_nd(tensor)
+    res = self.evaluate(value)
+    self.assertAllClose(res.shape,
+                        (shape[0], np.prod(shape[1:-2]), shape[-2], shape[-1]))
+
+  @parameterized.parameters(
+      ((5,),),
+      ((5, 10),),
+      ((5, 10, 15),),
+  )
+  def testUnflattenBlocksND(self, blocks_per_dim):
+    tensor = tf.zeros([2, np.prod(blocks_per_dim), 6, 10])
+    value = common_attention.unflatten_blocks_nd(tensor, blocks_per_dim)
+    res = self.evaluate(value)
+    self.assertAllClose(res.shape, (2,) + blocks_per_dim + (6, 10))
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testBreakIntoMemoryBlocksND(self):
+    tensor = tf.convert_to_tensor(
+        np.array([[
+            [[1], [2], [3], [4]],
+            [[5], [6], [7], [8]],
+            [[9], [10], [11], [12]],
+            [[13], [14], [15], [16]],
+        ]]))
+    value = common_attention.break_into_memory_blocks_nd(tensor,
+                                                         (2, 2),
+                                                         (2, 2),
+                                                         masked=True)
+    res = self.evaluate(value)
+    expected_value = np.array([[
+        [
+            [
+                [0], [0], [0], [0], [0], [0], [0], [0], [0], [0], [0], [0],
+                [0], [0], [0], [0], [1], [2], [5], [6], [3], [4], [7], [8]
+            ],
+            [
+                [0], [0], [0], [0], [0], [0], [0], [0], [0], [0], [0], [0],
+                [1], [2], [5], [6], [3], [4], [7], [8], [0], [0], [0], [0]
+            ]
+        ],
+        [
+            [
+                [0], [0], [0], [0], [1], [2], [5], [6], [3], [4], [7], [8], [0],
+                [0], [0], [0], [9], [10], [13], [14], [11], [12], [15], [16]
+            ],
+            [
+                [1], [2], [5], [6], [3], [4], [7], [8], [0], [0], [0], [0], [9],
+                [10], [13], [14], [11], [12], [15], [16], [0], [0], [0], [0]
+            ]
+        ]]])
+    self.assertAllClose(expected_value, res)
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testBreakIntoBlocksND(self):
+    tensor = tf.convert_to_tensor(
+        np.array([[
+            [[1], [2], [3], [4]],
+            [[5], [6], [7], [8]],
+            [[9], [10], [11], [12]],
+            [[13], [14], [15], [16]],
+        ]]))
+    value = common_attention.break_into_blocks_nd(tensor, (2, 2))
+    res = self.evaluate(value)
+    expected_value = np.array([[
+        [[[1], [2], [5], [6]], [[3], [4], [7], [8]]],
+        [[[9], [10], [13], [14]], [[11], [12], [15], [16]]]
+    ]])
+    self.assertAllClose(expected_value, res)
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testPutBackBlocksND(self):
+    tensor = tf.convert_to_tensor(
+        np.array([[
+            [[[1], [2], [5], [6]], [[3], [4], [7], [8]]],
+            [[[9], [10], [13], [14]], [[11], [12], [15], [16]]]
+        ]]))
+    value = common_attention.put_back_blocks_nd(tensor, (2, 2))
+    res = self.evaluate(value)
+    expected_value = np.array([[
+        [[1], [2], [3], [4]],
+        [[5], [6], [7], [8]],
+        [[9], [10], [11], [12]],
+        [[13], [14], [15], [16]],
+    ]])
+    self.assertAllClose(expected_value, res)
+
+  @parameterized.parameters(
+      ((2, 100, 5), (7,), (2, 105, 5)),
+      ((2, 100, 100, 5), (5, 7), (2, 100, 105, 5)),
+      ((2, 100, 100, 100, 5), (10, 20, 30), (2, 100, 100, 120, 5))
+  )
+  def testPadToMultipleND(self, tensor_shape, block_shape, expected_shape):
+    tensor = tf.zeros(tensor_shape)
+    value = common_attention.pad_to_multiple_nd(tensor, block_shape)
+    res = self.evaluate(value)
+    self.assertAllClose(res.shape, expected_shape)
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testCausalAttentionBiasND(self):
+    bias = common_attention.causal_attention_bias_nd((2, 2), (2, 2))
+    res = self.evaluate(bias)
+    expected_val = np.array([[[
+        [0] * 17 + [-1e9] * 7,
+        [0] * 18 + [-1e9] * 6,
+        [0] * 19 + [-1e9] * 5,
+        [0] * 20 + [-1e9] * 4,
+    ]]])
+    self.assertAllClose(expected_val, res)
+
+  @parameterized.parameters(
+      ((1, 64, 10), (80,), (80,)),
+      ((1, 64, 64, 10), (8, 8), (16, 16)),
+      ((1, 5, 64, 64, 10), (1, 8, 8), (1, 8, 8))
+  )
+  def testMultiheadAttentionND(self, tensor_shape, query_shape, memory_flange):
+    query_antecedent = tf.zeros(tensor_shape)
+    value = common_attention.multihead_attention_nd(
+        query_antecedent=query_antecedent,
+        memory_antecedent=None,
+        total_key_depth=256,
+        total_value_depth=256,
+        output_depth=256,
+        num_heads=4,
+        query_shape=query_shape,
+        memory_flange=memory_flange,
+        masked=True)
+    res = self.evaluate(value)
+    self.assertAllClose(res.shape, tensor_shape[:-1] + (256,))
+
+  @parameterized.parameters(
+      (15, (5,), (100,), (15,)),
+      (10, (2, 2), (4, 4), (3, 0)),
+      (25, (2, 2, 3), (10, 10, 12), (0, 0, 7))
+  )
+  def testDecodeStepToIndex(self, decode_step, query_shape, tensor_shape,
+                            expected_index):
+    res = common_attention.decode_step_to_index(decode_step, query_shape,
+                                                tensor_shape)
+    self.assertAllClose(res, expected_index)
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testGetItemAtDecodeStep(self):
+    tensor = tf.reshape(tf.range(25 * 25 * 4), [1, 4, 25, 25, 1])
+    value = common_attention.get_item_at_decode_step(tensor, 100, (2, 5, 5))
+    res = self.evaluate(value)
+    expected_value = np.array([[[[[10]]]]])
+    self.assertAllClose(expected_value, res)
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testPutItemAtDecodeStep(self):
+    tensor = tf.zeros([1, 1, 10, 10, 1])
+    item = tf.ones([1, 1, 1, 1, 1])
+    value = common_attention.put_item_in_decode_step(tensor, item, 32, (2, 2))
+    res = self.evaluate(value)
+    expected_val = np.zeros([1, 1, 10, 10, 1])
+    expected_val[0, 0, 2, 6, 0] = 1
+    self.assertAllClose(expected_val, res)
+
   @parameterized.named_parameters(
       ("", 1, 1, 8, 4, 1, 2),
       ("dynamic_batch", None, 1, 8, 4, 1, 2),
@@ -1274,4 +1559,3 @@ def testMaskedDilatedAttention(self, batch, heads, length, depth_v,
 
 if __name__ == "__main__":
   tf.test.main()
-

From 96865f7e912d7faeb99756f5077f03a6cfa9b459 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Fri, 31 May 2019 08:51:18 -0700
Subject: [PATCH 2080/2720] [TRAX] Improve the efficiency of the multi-device
 _save_replicated function.

Profiling shows that save_replicated is responsible for triggering a large number of synchronous host to device transfers inside the call to .mean(0) on a ShardedDeviceArray.

Although we do also plan to improve the efficiency of the current code by improving JAX, we don't need to compute the mean at all since all replicas should be identical. Instead, use the value from the first replica.

PiperOrigin-RevId: 250893206
---
 tensor2tensor/trax/trax.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 1f8177408..52d53000d 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -171,10 +171,10 @@ def save_state(state, output_dir, keep=False):
 def _save_replicated(opt_state, step, history, n_devices, output_dir, keep):
   """Save state but given a possibly replicated opt_state."""
   if n_devices > 1:
-    unreplicate = lambda x: x.mean(0)
-    opt_state = layers.nested_map(opt_state, unreplicate)
-    save_state(State(params=opt_state, step=step, history=history),
-               output_dir, keep=keep)
+    first_replica = lambda x: x[0]
+    opt_state = layers.nested_map(opt_state, first_replica)
+  save_state(State(params=opt_state, step=step, history=history),
+             output_dir, keep=keep)
 
 
 # Metrics to calculate and report.

From c5ccb790d63edf24860544bdbc2921059b442e75 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 31 May 2019 09:40:46 -0700
Subject: [PATCH 2081/2720] Clarify relation between layer lists and layer
 objects. Remove combinators.NoOp(); can use empty list instead.

PiperOrigin-RevId: 250901543
---
 .../chunked_transformer_imagenet64_8gb.gin    |  2 +-
 tensor2tensor/trax/layers/combinators.py      | 73 +++++++++++++------
 tensor2tensor/trax/layers/combinators_test.py | 10 +--
 tensor2tensor/trax/layers/rnn.py              | 11 +--
 tensor2tensor/trax/models/atari_cnn.py        |  5 +-
 tensor2tensor/trax/models/atari_cnn_test.py   |  6 +-
 tensor2tensor/trax/models/mlp.py              |  6 +-
 tensor2tensor/trax/models/mlp_test.py         |  2 +-
 tensor2tensor/trax/models/neural_gpu.py       |  6 +-
 .../models/research/chunked_transformer.py    | 45 ++++++------
 tensor2tensor/trax/models/resnet.py           |  8 +-
 tensor2tensor/trax/models/resnet_test.py      |  4 +-
 tensor2tensor/trax/models/transformer.py      | 14 ++--
 tensor2tensor/trax/models/transformer_test.py |  4 +-
 tensor2tensor/trax/trax.py                    |  4 +-
 15 files changed, 112 insertions(+), 88 deletions(-)

diff --git a/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin b/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin
index 119b3f740..10b984fd1 100644
--- a/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin
+++ b/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin
@@ -11,7 +11,7 @@ batch_fun.max_eval_length = 12288  # 64 * 64 * 3
 
 # Parameters for inputs:
 # ==============================================================================
-inputs.n_chunks = 64
+inputs.num_chunks = 64
 inputs.data_dir = None
 inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
 inputs.input_name = 'targets'
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 6e45bd55c..8a2b80125 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -27,24 +27,53 @@
 from tensor2tensor.trax.layers import base
 
 
-def _DeepFlatten(xs):  # pylint: disable=invalid-name
+def Model(*layers):
+  """Ensures that a layer or list of layers can be treated as a model.
+
+  Currently, any subclass of base.Layer can be treated as a model.
+
+  Args:
+    *layers: One or more layer objects. In fuller detail, the list may contain
+        nested sublists, and the top-level list can also be a tuple.
+
+  Returns:
+    A single object that treated as a model, e.g., trained or evaluated.
+  """
+  return Serial(*layers)
+
+
+def _deep_flatten(xs):  # pylint: disable=invalid-name
   for x in xs:
     if isinstance(x, (list, tuple)):
-      for y in _DeepFlatten(x):
+      for y in _deep_flatten(x):
         yield y
     else:
       yield x
 
 
-def _EnsureSublayers(layers):
-  # TODO(jonni): Implement for dict if dicts remain important.
-  if isinstance(layers, dict):
-    return layers
-  sublayers_not_lists = []
-  for layer in layers:
-    sublayers_not_lists.append(
-        Serial(layer) if isinstance(layer, list) else layer)
-  return sublayers_not_lists
+def _ensure_sublayers(layers):  # pylint: disable=invalid-name
+  """Ensures that elements in a layer list (or dict) are layers.
+
+  Args:
+    layers: A list or dict whose elements/values can each be a layer, a list,
+        or a dict, and so on recursively.
+
+  Returns:
+    An analogous collection of layers in which embedded layer lists are
+    wrapped in Serial layer instances.
+  """
+  if not layers:  # None or an empty list can signal a no-op.
+    return Serial([])  # no-op, but still handles shapes and initialization
+  elif isinstance(layers, dict):
+    return {k: _ensure_sublayers(v) for k, v in layers.items()}
+  elif isinstance(layers, (list, tuple)):
+    sublayers_not_lists = []
+    for layer in layers:
+      sublayers_not_lists.append(
+          Serial(layer) if isinstance(layer, (list, tuple)) else layer)
+    return sublayers_not_lists
+  else:
+    raise TypeError(type(layers))
 
 
 class Serial(base.Layer):
@@ -52,7 +81,7 @@ class Serial(base.Layer):
 
   def __init__(self, *layers):
     super(Serial, self).__init__()
-    layers = list(_DeepFlatten(layers))
+    layers = list(_deep_flatten(layers))
     # TODO(jonni): Consider flattening (unpacking) also embedded Serial layers.
     self._layers = layers
     self._nlayers = len(layers)
@@ -83,12 +112,6 @@ def new_parameters(self, input_shape, rng):
     return params
 
 
-@base.layer()
-def NoOp(x, **unused_kwargs):
-  """NoOp layer, return the inputs."""
-  return x
-
-
 def _print_shape(x, message='PrintShape'):  # pylint: disable=invalid-name
   print(message + ' ; stack shape = ' + str(x))
   return x
@@ -96,7 +119,7 @@ def _print_shape(x, message='PrintShape'):  # pylint: disable=invalid-name
 
 @base.layer(output_shape=_print_shape, stack_items_to_pass=0)
 def PrintShape(x, message='PrintShape', **unused_kwargs):
-  """NoOp layer that prints the shape of the stack."""
+  """No-op layer that prints the shape of the stack."""
   _print_shape(base.shapes(x), message=message)
   return x
 
@@ -172,7 +195,7 @@ def _flatten_shape(x_shape):  # pylint: disable=invalid-name
 @base.layer(output_shape=_flatten_shape, stack_items_to_pass=0)
 def Flatten(xs, **unused_kwargs):
   """Flatten lists."""
-  return tuple(_DeepFlatten(xs))
+  return tuple(_deep_flatten(xs))
 
 
 # Re-ordering layer.
@@ -193,8 +216,8 @@ class Select(base.Layer):
     Select((0, (1, 1)))      = (x, (y, y))
     Select(((2, 0), (1, 1))) = ((z, x), (y, y))
 
-  By default (if no output is given) Select does nothing (NoOp).
-  It is also possible to name the inputs to access tuple elements, e.g.:
+  By default (if no output is given) Select does nothing. It is also possible
+  to name the inputs to access tuple elements, e.g.:
 
   Select(inputs=('encoder', ('decoder', 'mask')), output='decoder')
 
@@ -269,7 +292,7 @@ def __init__(self, *layers, **kwlayers):
     if layers and kwlayers:
       raise ValueError('Cannot specify a Branch with both a list and dict.')
     layers = layers or kwlayers
-    layers = _EnsureSublayers(layers)
+    layers = _ensure_sublayers(layers)
     self._nlayers = len(layers)
     self._layers = layers
 
@@ -446,7 +469,7 @@ def __init__(self, *layers, **kwlayers):
     if layers and kwlayers:
       raise ValueError('Cannot specify a Parallel with both a list and dict.')
     layers = layers or kwlayers
-    layers = _EnsureSublayers(layers)
+    layers = _ensure_sublayers(layers)
     self._nlayers = len(layers)
     self._layers = layers
 
@@ -530,6 +553,8 @@ class Map(base.Layer):
 
   def __init__(self, layer, check_shapes=True):
     super(Map, self).__init__()
+    if layer is None or isinstance(layer, (list, tuple)):
+      layer = Serial(layer)
     self._layer = layer
     # Generally a Map should be applied to lists where all elements have
     # the same shape -- because self._layer will only be initialized once
diff --git a/tensor2tensor/trax/layers/combinators_test.py b/tensor2tensor/trax/layers/combinators_test.py
index ef91d08a8..592735978 100644
--- a/tensor2tensor/trax/layers/combinators_test.py
+++ b/tensor2tensor/trax/layers/combinators_test.py
@@ -30,30 +30,28 @@ def test_branch(self):
     input_shape = (2, 3)
     expected_shape = ((2, 3), (2, 3))
     output_shape = base.check_shape_agreement(
-        combinators.Branch(combinators.NoOp(), combinators.NoOp()), input_shape)
+        combinators.Branch([], []), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
   def test_branch_named(self):
     input_shape = (2, 3)
     expected_shape = {'a': (2, 3), 'b': (2, 3)}
     output_shape = base.check_shape_agreement(
-        combinators.Branch(a=combinators.NoOp(), b=combinators.NoOp()),
-        input_shape)
+        combinators.Branch(a=[], b=[]), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
   def test_parallel(self):
     input_shape = ((2, 3), (2, 3))
     expected_shape = ((2, 3), (2, 3))
     output_shape = base.check_shape_agreement(
-        combinators.Parallel(combinators.NoOp(), combinators.NoOp()),
-        input_shape)
+        combinators.Parallel([], []), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
   def test_parallel_named(self):
     input_shape = {'a': (2, 3), 'b': (2, 3)}
     expected_shape = {'a': (2, 3), 'b': (2, 3)}
     output_shape = base.check_shape_agreement(
-        combinators.Parallel(a=combinators.NoOp()), input_shape)
+        combinators.Parallel(a=[]), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
   def test_select(self):
diff --git a/tensor2tensor/trax/layers/rnn.py b/tensor2tensor/trax/layers/rnn.py
index 8c51562b6..979dfeba8 100644
--- a/tensor2tensor/trax/layers/rnn.py
+++ b/tensor2tensor/trax/layers/rnn.py
@@ -38,7 +38,7 @@ def GRUCell(units):
   """
   return GeneralGRUCell(
       candidate_transform=lambda: core.Dense(units=units),
-      memory_transform=cb.NoOp,
+      memory_transform_fn=None,
       gate_nonlinearity=core.Sigmoid,
       candidate_nonlinearity=core.Tanh)
 
@@ -62,13 +62,13 @@ def BuildConv():
 
   return GeneralGRUCell(
       candidate_transform=BuildConv,
-      memory_transform=cb.NoOp,
+      memory_transform_fn=None,
       gate_nonlinearity=core.Sigmoid,
       candidate_nonlinearity=core.Tanh)
 
 
 def GeneralGRUCell(candidate_transform,
-                   memory_transform=cb.NoOp,
+                   memory_transform_fn=None,
                    gate_nonlinearity=core.Sigmoid,
                    candidate_nonlinearity=core.Tanh,
                    dropout_rate_c=0.1,
@@ -87,7 +87,7 @@ def GeneralGRUCell(candidate_transform,
   Args:
     candidate_transform: Transform to apply inside the Candidate branch. Applied
       before nonlinearities.
-    memory_transform: Optional transformation on the memory before gating.
+    memory_transform_fn: Optional transformation on the memory before gating.
     gate_nonlinearity: Function to use as gate activation. Allows trying
       alternatives to Sigmoid, such as HardSigmoid.
     candidate_nonlinearity: Nonlinearity to apply after candidate branch. Allows
@@ -119,7 +119,8 @@ def GeneralGRUCell(candidate_transform,
       # Only apply dropout on the C gate. Paper reports 0.1 as a good default.
       core.Dropout(rate=dropout_rate_c)
   ]
+  memory_transform = memory_transform_fn() if memory_transform_fn else []
   return cb.Serial([
-      cb.Branch(memory_transform(), gate_block, candidate_block),
+      cb.Branch(memory_transform, gate_block, candidate_block),
       cb.Gate(),
   ])
diff --git a/tensor2tensor/trax/models/atari_cnn.py b/tensor2tensor/trax/models/atari_cnn.py
index 7b102c1de..e9128b348 100644
--- a/tensor2tensor/trax/models/atari_cnn.py
+++ b/tensor2tensor/trax/models/atari_cnn.py
@@ -23,8 +23,9 @@
 
 
 def AtariCnn(hidden_sizes=(32, 32), output_size=128):
+  """An Atari CNN."""
   # Input's shape = (B, T, H, W, C)
-  return tl.Serial([
+  return tl.Model(
       tl.Div(divisor=255.0),
       # Have 4 copies of the input, each one shifted to the right by one.
       tl.Branch(
@@ -43,4 +44,4 @@ def AtariCnn(hidden_sizes=(32, 32), output_size=128):
       tl.Dense(output_size),
       tl.Relu(),
       # Eventually this is shaped (B, T, output_size)
-  ])
+  )
diff --git a/tensor2tensor/trax/models/atari_cnn_test.py b/tensor2tensor/trax/models/atari_cnn_test.py
index 7b36db92a..6d1ff70f3 100644
--- a/tensor2tensor/trax/models/atari_cnn_test.py
+++ b/tensor2tensor/trax/models/atari_cnn_test.py
@@ -33,15 +33,15 @@ def test_computes(self):
     rng_key = jax_random.get_prng(0)
     hidden_size = (4, 4)
     output_size = 6
-    policy = atari_cnn.AtariCnn(
+    model = atari_cnn.AtariCnn(
         hidden_sizes=hidden_size, output_size=output_size)
     B, T, OBS = 2, 2, (28, 28, 3)  # pylint: disable=invalid-name
     rng_key, key = jax_random.split(rng_key)
-    params = policy.initialize((-1, -1) + OBS, key)
+    params = model.initialize((-1, -1) + OBS, key)
     x = onp.arange(B * (T + 1) * functools.reduce(op.mul, OBS)).reshape(
         B, T + 1, *OBS)
     rng_key, key = jax_random.split(rng_key)
-    y = policy(x, params, rng=key)
+    y = model(x, params, rng=key)
     self.assertEqual((B, T + 1, output_size), y.shape)
 
 
diff --git a/tensor2tensor/trax/models/mlp.py b/tensor2tensor/trax/models/mlp.py
index 9283fad4e..f6f2a9575 100644
--- a/tensor2tensor/trax/models/mlp.py
+++ b/tensor2tensor/trax/models/mlp.py
@@ -27,12 +27,12 @@ def MLP(n_hidden_layers=2,
         activation_fn=tl.Relu,
         n_output_classes=10,
         mode="train"):
-  """Multi-layer feed-forward neural network with non-linear activations."""
+  """A multi-layer feedforward (perceptron) network."""
   del mode
 
-  return [
+  return tl.Model(
       tl.Flatten(),
       [[tl.Dense(d_hidden), activation_fn()] for _ in range(n_hidden_layers)],
       tl.Dense(n_output_classes),
       tl.LogSoftmax(),
-  ]
+  )
diff --git a/tensor2tensor/trax/models/mlp_test.py b/tensor2tensor/trax/models/mlp_test.py
index 67e51bde4..2ba7568d0 100644
--- a/tensor2tensor/trax/models/mlp_test.py
+++ b/tensor2tensor/trax/models/mlp_test.py
@@ -31,7 +31,7 @@ def test_mlp_forward_shape(self):
     """Run the MLP model forward and check output shape."""
     input_shape = (3, 28, 28, 1)
     model = mlp.MLP(d_hidden=32, n_output_classes=10)
-    final_shape = tl.check_shape_agreement(tl.Serial(model), input_shape)
+    final_shape = tl.check_shape_agreement(model, input_shape)
     self.assertEqual((3, 10), final_shape)
 
 
diff --git a/tensor2tensor/trax/models/neural_gpu.py b/tensor2tensor/trax/models/neural_gpu.py
index e76e91c2c..0d357a72c 100644
--- a/tensor2tensor/trax/models/neural_gpu.py
+++ b/tensor2tensor/trax/models/neural_gpu.py
@@ -54,7 +54,7 @@ def BuildConv():
 
   return tl.GeneralGRUCell(
       candidate_transform=BuildConv,
-      memory_transform=DiagonalGate,
+      memory_transform_fn=DiagonalGate,
       gate_nonlinearity=tl.HardSigmoid,
       candidate_nonlinearity=tl.HardTanh)
 
@@ -71,9 +71,9 @@ def NeuralGPU(d_feature=96, steps=16, vocab_size=2):
     A NeuralGPU Stax model.
   """
   core = ConvDiagonalGRU(units=d_feature)
-  return tl.Serial([
+  return tl.Model(
       tl.Embedding(d_feature=d_feature, vocab_size=vocab_size),
       [core] * steps,
       tl.Dense(vocab_size),
       tl.LogSoftmax(),
-  ])
+  )
diff --git a/tensor2tensor/trax/models/research/chunked_transformer.py b/tensor2tensor/trax/models/research/chunked_transformer.py
index 2288445a0..4223e98cd 100644
--- a/tensor2tensor/trax/models/research/chunked_transformer.py
+++ b/tensor2tensor/trax/models/research/chunked_transformer.py
@@ -141,39 +141,38 @@ def ChunkedCausalMultiHeadedAttention(
   Returns:
     Multi-headed self-attention layer.
   """
-  prepare_attention_input = tl.Serial(
+  prepare_attention_input = [
       tl.Branch(
-          tl.Branch(  # q = k = v = first input
-              tl.NoOp(), tl.NoOp(), tl.NoOp()),
-          tl.CausalMask(axis=-2),
+          tl.Branch([], [], []),  # q = k = v = first input
+          tl.CausalMask(axis=-2)
       ),
       tl.Parallel(
           tl.Parallel(
               tl.Dense(d_feature),
               tl.Dense(d_feature),
-              tl.Dense(d_feature),
+              tl.Dense(d_feature)
           ),
-          tl.NoOp()
+          []
       )
-  )
-  return tl.Serial(
+  ]
+  return [
       tl.Map(prepare_attention_input),
       ChunkedAttentionSelector(selector=chunk_selector),  # pylint: disable=no-value-for-parameter
-      tl.Map(tl.PureMultiHeadedAttention(
-          d_feature=d_feature, n_heads=n_heads,
-          dropout=dropout, mode=mode), check_shapes=False),
+      tl.Map(tl.PureMultiHeadedAttention(d_feature=d_feature, n_heads=n_heads,
+                                         dropout=dropout, mode=mode),
+             check_shapes=False),
       tl.Map(tl.Select(0), check_shapes=False),  # drop masks
       tl.Map(tl.Dense(d_feature))
-  )
+  ]
 
 
 # Chunked residual.
 def Residual(*layers, **unused_kwargs):
   """Constructs a residual version of layers, summing input to layers output."""
-  return tl.Serial(
-      tl.Branch(tl.Serial(*layers), tl.NoOp()),
+  return [
+      tl.Branch(layers, []),
       tl.AddAll()
-  )
+  ]
 
 
 def ResidualFeedForward(d_feature, d_feedforward, dropout, mode):
@@ -205,9 +204,9 @@ def ChunkedDecoderLayer(d_feature,
     mode: str: 'train' or 'eval'
 
   Returns:
-    the layer.
+    The layers comprising a chunked decoder.
   """
-  return tl.Serial(
+  return [
       Residual(  # Self-attention block.
           tl.Map(tl.LayerNorm()),
           ChunkedCausalMultiHeadedAttention(
@@ -217,7 +216,7 @@ def ChunkedDecoderLayer(d_feature,
       ),
       tl.Map(ResidualFeedForward(
           d_feature, d_feedforward, dropout, mode=mode))
-  )
+  ]
 
 
 def ChunkedTransformerLM(vocab_size,
@@ -257,16 +256,16 @@ def ChunkedTransformerLM(vocab_size,
   Returns:
     the layer.
   """
-  stack = [ChunkedDecoderLayer(d_feature, d_feedforward, n_heads,
-                               dropout, chunk_selector, mode)
-           for _ in range(n_layers)]
+  decoder_stack = [ChunkedDecoderLayer(d_feature, d_feedforward, n_heads,
+                                       dropout, chunk_selector, mode)
+                   for _ in range(n_layers)]
   # Below each Map(L) applies the layer L to each chunk independently.
-  return tl.Serial(
+  return tl.Model(
       tl.ShiftRight(),
       tl.Map(tl.Embedding(d_feature, vocab_size)),
       tl.Map(tl.Dropout(rate=dropout, mode=mode)),
       ChunkedPositionalEncoding(max_len=max_len),  # pylint: disable=no-value-for-parameter
-      tl.Serial(*stack),
+      decoder_stack,
       tl.Map(tl.LayerNorm()),
       tl.Map(tl.Dense(vocab_size)),
       tl.Map(tl.LogSoftmax()),
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index 3221c660f..8062275ac 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -80,7 +80,7 @@ def Resnet50(d_hidden=64, n_output_classes=1001, mode='train'):
     The list of layers comprising a ResNet model with the given parameters.
   """
   del mode
-  return [
+  return tl.Model(
       tl.Conv(d_hidden, (7, 7), (2, 2), 'SAME'),
       tl.BatchNorm(),
       tl.Relu(),
@@ -105,7 +105,7 @@ def Resnet50(d_hidden=64, n_output_classes=1001, mode='train'):
       tl.Flatten(),
       tl.Dense(n_output_classes),
       tl.LogSoftmax(),
-  ]
+  )
 
 
 def WideResnetBlock(channels, strides=(1, 1)):
@@ -145,7 +145,7 @@ def WideResnet(n_blocks=3, d_hidden=64, n_output_classes=10,
     The list of layers comprising a WideResnet model with the given parameters.
   """
   del mode
-  return [
+  return tl.Model(
       tl.Conv(d_hidden, (3, 3), padding='SAME'),
       WideResnetGroup(n_blocks, d_hidden),
       WideResnetGroup(n_blocks, d_hidden * 2, (2, 2)),
@@ -156,4 +156,4 @@ def WideResnet(n_blocks=3, d_hidden=64, n_output_classes=10,
       tl.Flatten(),
       tl.Dense(n_output_classes),
       tl.LogSoftmax(),
-  ]
+  )
diff --git a/tensor2tensor/trax/models/resnet_test.py b/tensor2tensor/trax/models/resnet_test.py
index 81e17e869..de36a6757 100644
--- a/tensor2tensor/trax/models/resnet_test.py
+++ b/tensor2tensor/trax/models/resnet_test.py
@@ -29,13 +29,13 @@ class ResnetTest(absltest.TestCase):
   def test_resnet(self):
     input_shape = (3, 256, 256, 3)
     model = resnet.Resnet50(d_hidden=8, n_output_classes=10)
-    final_shape = tl.check_shape_agreement(tl.Serial(model), input_shape)
+    final_shape = tl.check_shape_agreement(model, input_shape)
     self.assertEqual((3, 10), final_shape)
 
   def test_wide_resnet(self):
     input_shape = (3, 32, 32, 3)
     model = resnet.WideResnet(n_blocks=1, n_output_classes=10)
-    final_shape = tl.check_shape_agreement(tl.Serial(model), input_shape)
+    final_shape = tl.check_shape_agreement(model, input_shape)
     self.assertEqual((3, 10), final_shape)
 
 
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 0dd30d241..1bca46d78 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -122,7 +122,7 @@ def DecoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
   """
   self_attention = [
       tl.LayerNorm(),
-      tl.Branch(tl.NoOp(), tl.CausalMask(axis=-2)),  # Create mask.
+      tl.Branch([], tl.CausalMask(axis=-2)),  # Create mask.
       tl.MultiHeadedAttention(
           d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Select(0),  # Drop mask.
@@ -165,7 +165,7 @@ def TransformerLM(vocab_size,
       tl.Dropout(rate=dropout, mode=mode),
       tl.PositionalEncoding(max_len=max_len),
   ]
-  return [
+  return tl.Model(
       tl.ShiftRight(),
       positional_embedder,
       [DecoderBlock(d_feature, d_feedforward, n_heads, dropout, mode)
@@ -173,7 +173,7 @@ def TransformerLM(vocab_size,
       tl.LayerNorm(),
       tl.Dense(vocab_size),
       tl.LogSoftmax(),
-  ]
+  )
 
 
 def EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode):
@@ -258,17 +258,17 @@ def Transformer(vocab_size,
        for _ in range(n_layers)],
       tl.LayerNorm(),
   ]
-  return [
-      tl.Parallel(tl.NoOp(), tl.ShiftRight()),
+  return tl.Model(
+      tl.Parallel([], tl.ShiftRight()),
       tl.Parallel(encoder, positional_embedder),
       tl.Select(inputs=(('encoder', 'mask'), 'decoder'),
                 output=('decoder', ('mask', 'decoder'), 'encoder')),
       # (encoder_mask, decoder_input) -> encoder-decoder mask
-      tl.Parallel(tl.NoOp(), tl.EncoderDecoderMask(), tl.NoOp()),
+      tl.Parallel([], tl.EncoderDecoderMask(), []),
       [EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode)
        for _ in range(n_layers)],
       tl.Select(0),  # Drop mask and encoder.
       tl.LayerNorm(),
       tl.Dense(vocab_size),
       tl.LogSoftmax(),
-  ]
+  )
diff --git a/tensor2tensor/trax/models/transformer_test.py b/tensor2tensor/trax/models/transformer_test.py
index 97574b2e4..5d49a4a5f 100644
--- a/tensor2tensor/trax/models/transformer_test.py
+++ b/tensor2tensor/trax/models/transformer_test.py
@@ -34,7 +34,7 @@ def test_transformer_lm_forward_shape(self):
     model = transformer.TransformerLM(
         vocab_size, d_feature=32, d_feedforward=64, n_layers=2, n_heads=2)
     final_shape = tl.check_shape_agreement(
-        tl.Serial(model), tuple(input_shape), integer_inputs=True)
+        model, tuple(input_shape), integer_inputs=True)
     self.assertEqual(tuple(input_shape + [vocab_size]), final_shape)
 
   def test_transformer_forward_shape(self):
@@ -45,7 +45,7 @@ def test_transformer_forward_shape(self):
     model = transformer.Transformer(
         vocab_size, d_feature=32, d_feedforward=64, n_layers=2, n_heads=2)
     final_shape = tl.check_shape_agreement(
-        tl.Serial(model), input_shape, integer_inputs=True)
+        model, input_shape, integer_inputs=True)
     self.assertEqual(tuple(single_input_shape + [vocab_size]), final_shape)
 
 
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 52d53000d..eafd3dbbb 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -443,8 +443,8 @@ def train(output_dir,
   lr_fn = lr_schedule(history)
   opt = optimizer(lr_fn)
 
-  model_train = layers.Serial(model(mode="train"))
-  model_predict_eval = layers.Serial(model(mode="eval"))
+  model_train = model(mode="train")
+  model_predict_eval = model(mode="eval")
 
   # Setup state
   step = state.step or 0

From 53dff6b2bbadce691a0fc5e5cad50f38d2856f64 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 31 May 2019 12:30:21 -0700
Subject: [PATCH 2082/2720] Small comment clean-up.

PiperOrigin-RevId: 250935202
---
 tensor2tensor/trax/layers/attention.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 4e2583e77..e66a76fe3 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -152,7 +152,7 @@ def PureMultiHeadedAttention(x, params, n_heads=8, dropout=0.0, mode='train',
   """Pure transformer-style multi-headed attention.
 
   Args:
-    x: inputs ((q, k, v), mask)
+    x: inputs (q, k, v, mask)
     params: parameters (none)
     n_heads: int: number of attention heads
     dropout: float: dropout rate
@@ -188,7 +188,7 @@ def JoinHeads(x):  # pylint: disable=invalid-name
 def MultiHeadedAttentionQKV(d_feature, n_heads=8, dropout=0.0, mode='train'):
   """Transformer-style multi-headed attention.
 
-  Accepts inputs of the form (q, k, v), mask.
+  Accepts inputs of the form q, k, v, mask.
 
   Args:
     d_feature: int:  dimensionality of feature embedding

From 7eecac65aaeaab36e3e0756d0b87000ff020a83b Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 31 May 2019 12:58:16 -0700
Subject: [PATCH 2083/2720] Correct DQN learner API after recent changes.

PiperOrigin-RevId: 250940034
---
 tensor2tensor/rl/dopamine_connector.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 1686c86f4..45605c793 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -500,8 +500,10 @@ def _get_optimizer(params):
 class DQNLearner(PolicyLearner):
   """Interface for learning dqn implemented in dopamine."""
 
-  def __init__(self, *args, **kwargs):
-    super(DQNLearner, self).__init__(*args, **kwargs)
+  def __init__(self, frame_stack_size, base_event_dir, agent_model_dir,
+               total_num_epochs, **kwargs):
+    super(DQNLearner, self).__init__(
+        frame_stack_size, base_event_dir, agent_model_dir, total_num_epochs)
     self.completed_iterations = 0
 
   def _target_iteractions_and_steps(self, num_env_steps, save_continuously,

From 9f78c55b556eaa3e5128c9a07393780b9119c940 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 31 May 2019 13:26:26 -0700
Subject: [PATCH 2084/2720] Open-source discrete flows implementation.

PiperOrigin-RevId: 250944871
---
 tensor2tensor/layers/reversible_layers.py     | 643 ++++++++++++++++++
 .../layers/reversible_layers_test.py          | 382 +++++++++++
 2 files changed, 1025 insertions(+)

diff --git a/tensor2tensor/layers/reversible_layers.py b/tensor2tensor/layers/reversible_layers.py
index 4082d49e6..b45efba90 100644
--- a/tensor2tensor/layers/reversible_layers.py
+++ b/tensor2tensor/layers/reversible_layers.py
@@ -39,6 +39,649 @@ def __init__(self, reversible_layer, **kwargs):
     self.reverse = reversible_layer.call
 
 
+class DiscreteAutoregressiveFlow(tf.keras.layers.Layer):
+  """A discrete reversible layer.
+
+  The flow takes as input a one-hot Tensor of shape `[..., length, vocab_size]`.
+  The flow returns a Tensor of same shape and dtype. (To enable gradients, the
+  input must have float dtype.)
+
+  For the forward pass, the flow computes in serial:
+
+  ```none
+  outputs = []
+  for t in range(length):
+    new_inputs = [outputs, inputs[..., t, :]]
+    net = layer(new_inputs)
+    loc, scale = tf.split(net, 2, axis=-1)
+    loc = tf.argmax(loc, axis=-1)
+    scale = tf.argmax(scale, axis=-1)
+    new_outputs = (((inputs - loc) * inverse(scale)) % vocab_size)[..., -1, :]
+    outputs.append(new_outputs)
+  ```
+
+  For the reverse pass, the flow computes in parallel:
+
+  ```none
+  net = layer(inputs)
+  loc, scale = tf.split(net, 2, axis=-1)
+  loc = tf.argmax(loc, axis=-1)
+  scale = tf.argmax(scale, axis=-1)
+  outputs = (loc + scale * inputs) % vocab_size
+  ```
+
+  The modular arithmetic happens in one-hot space.
+
+  If `x` is a discrete random variable, the induced probability mass function on
+  the outputs `y = flow(x)` is
+
+  ```none
+  p(y) = p(flow.reverse(y)).
+  ```
+
+  The location-only transform is always invertible ([integers modulo
+  `vocab_size` form an additive group](
+  https://en.wikipedia.org/wiki/Modular_arithmetic)). The transform with a scale
+  is invertible if the scale and `vocab_size` are coprime (see
+  [prime fields](https://en.wikipedia.org/wiki/Finite_field)).
+  """
+
+  def __init__(self, layer, temperature, **kwargs):
+    """Constructs flow.
+
+    Args:
+      layer: Two-headed masked network taking the inputs and returning a
+        real-valued Tensor of shape `[..., length, 2*vocab_size]`.
+        Alternatively, `layer` may return a Tensor of shape
+        `[..., length, vocab_size]` to be used as the location transform; the
+        scale transform will be hard-coded to 1.
+      temperature: Positive value determining bias of gradient estimator.
+      **kwargs: kwargs of parent class.
+    """
+    super(DiscreteAutoregressiveFlow, self).__init__(**kwargs)
+    self.layer = layer
+    self.temperature = temperature
+
+  def build(self, input_shape):
+    input_shape = tf.TensorShape(input_shape)
+    self.vocab_size = input_shape[-1]
+    if isinstance(self.vocab_size, tf.Dimension):
+      self.vocab_size = self.vocab_size.value
+    if self.vocab_size is None:
+      raise ValueError('The last dimension of the inputs to '
+                       '`DiscreteAutoregressiveFlow` should be defined. Found '
+                       '`None`.')
+    self.built = True
+
+  def __call__(self, inputs, *args, **kwargs):
+    if not isinstance(inputs, ed.RandomVariable):
+      return super(DiscreteAutoregressiveFlow, self).__call__(
+          inputs, *args, **kwargs)
+    return TransformedRandomVariable(inputs, self)
+
+  def call(self, inputs, **kwargs):
+    """Forward pass for left-to-right autoregressive generation."""
+    inputs = tf.convert_to_tensor(inputs)
+    length = inputs.shape[-2].value
+    if length is None:
+      raise NotImplementedError('length dimension must be known.')
+    # Form initial sequence tensor of shape [..., 1, vocab_size]. In a loop, we
+    # incrementally build a Tensor of shape [..., t, vocab_size] as t grows.
+    outputs = self._initial_call(inputs[..., 0, :], length, **kwargs)
+    # TODO(trandustin): Use tf.while_loop. Unrolling is memory-expensive for big
+    # models and not valid for variable lengths.
+    for t in range(1, length):
+      outputs = self._per_timestep_call(outputs,
+                                        inputs[..., t, :],
+                                        length,
+                                        t,
+                                        **kwargs)
+    return outputs
+
+  def _initial_call(self, new_inputs, length, **kwargs):
+    """Returns Tensor of shape [..., 1, vocab_size].
+
+    Args:
+      new_inputs: Tensor of shape [..., vocab_size], the new input to generate
+        its output.
+      length: Length of final desired sequence.
+      **kwargs: Optional keyword arguments to layer.
+    """
+    inputs = new_inputs[..., tf.newaxis, :]
+    # TODO(trandustin): To handle variable lengths, extend MADE to subset its
+    # input and output layer weights rather than pad inputs.
+    batch_ndims = inputs.shape.ndims - 2
+    padded_inputs = tf.pad(
+        inputs, [[0, 0]] * batch_ndims + [[0, length - 1], [0, 0]])
+    net = self.layer(padded_inputs, **kwargs)
+    if net.shape[-1] == 2 * self.vocab_size:
+      raise NotImplementedError()
+      # TODO(trandustin): Enable scale.
+      # loc, scale = tf.split(net, 2, axis=-1)
+      # scale = scale[..., 0, :][..., tf.newaxis, :]
+      # scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
+      # inverse_scale = multiplicative_inverse(scale, self.vocab_size)
+    elif net.shape[-1] == self.vocab_size:
+      loc = net
+      # inverse_scale = tf.ones_like(inputs)
+    else:
+      raise ValueError('Output of layer does not have compatible dimensions.')
+    loc = loc[..., 0, :][..., tf.newaxis, :]
+    loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
+    # scaled_inputs = one_hot_multiply(inputs, inverse_scale)
+    scaled_inputs = inputs
+    outputs = one_hot_minus(scaled_inputs, loc)
+    return outputs
+
+  def _per_timestep_call(self,
+                         current_outputs,
+                         new_inputs,
+                         length,
+                         timestep,
+                         **kwargs):
+    """Returns Tensor of shape [..., timestep+1, vocab_size].
+
+    Args:
+      current_outputs: Tensor of shape [..., timestep, vocab_size], the so-far
+        generated sequence Tensor.
+      new_inputs: Tensor of shape [..., vocab_size], the new input to generate
+        its output given current_outputs.
+      length: Length of final desired sequence.
+      timestep: Current timestep.
+      **kwargs: Optional keyword arguments to layer.
+    """
+    inputs = tf.concat([current_outputs,
+                        new_inputs[..., tf.newaxis, :]], axis=-2)
+    # TODO(trandustin): To handle variable lengths, extend MADE to subset its
+    # input and output layer weights rather than pad inputs.
+    batch_ndims = inputs.shape.ndims - 2
+    padded_inputs = tf.pad(
+        inputs, [[0, 0]] * batch_ndims + [[0, length - timestep - 1], [0, 0]])
+    net = self.layer(padded_inputs, **kwargs)
+    if net.shape[-1] == 2 * self.vocab_size:
+      raise NotImplementedError()
+      # TODO(trandustin): Enable scale.
+      # loc, scale = tf.split(net, 2, axis=-1)
+      # scale = scale[..., :(timestep+1), :]
+      # scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
+      # inverse_scale = multiplicative_inverse(scale, self.vocab_size)
+    elif net.shape[-1] == self.vocab_size:
+      loc = net
+      # inverse_scale = tf.ones_like(inputs)
+    else:
+      raise ValueError('Output of layer does not have compatible dimensions.')
+    loc = loc[..., :(timestep+1), :]
+    loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
+    # scaled_inputs = one_hot_multiply(inputs, inverse_scale)
+    scaled_inputs = inputs
+    new_outputs = one_hot_minus(scaled_inputs, loc)
+    outputs = tf.concat([current_outputs,
+                         new_outputs[..., -1, :][..., tf.newaxis, :]], axis=-2)
+    if not tf.executing_eagerly():
+      outputs.set_shape([None] * batch_ndims + [timestep+1, self.vocab_size])
+    return outputs
+
+  def reverse(self, inputs, **kwargs):
+    """Reverse pass returning the inverse autoregressive transformation."""
+    if not self.built:
+      self._maybe_build(inputs)
+
+    net = self.layer(inputs, **kwargs)
+    if net.shape[-1] == 2 * self.vocab_size:
+      raise NotImplementedError()
+      # TODO(trandustin): Enable scale.
+      # loc, scale = tf.split(net, 2, axis=-2)
+      # scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
+    elif net.shape[-1] == self.vocab_size:
+      loc = net
+      # scale = tf.ones_like(inputs)
+    else:
+      raise ValueError('Output of layer does not have compatible dimensions.')
+    loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
+    # scaled_inputs = one_hot_multiply(inputs, scale)
+    scaled_inputs = inputs
+    outputs = one_hot_add(loc, scaled_inputs)
+    return outputs
+
+  def log_det_jacobian(self, inputs):
+    return tf.cast(0, inputs.dtype)
+
+
+class DiscreteBipartiteFlow(tf.keras.layers.Layer):
+  """A discrete reversible layer.
+
+  The flow takes as input a one-hot Tensor of shape `[..., length, vocab_size]`.
+  The flow returns a Tensor of same shape and dtype. (To enable gradients, the
+  input must have float dtype.)
+
+  For the forward pass, the flow computes:
+
+  ```none
+  net = layer(mask * inputs)
+  loc, scale = tf.split(net, 2, axis=-1)
+  loc = tf.argmax(loc, axis=-1)
+  scale = tf.argmax(scale, axis=-1)
+  outputs = ((inputs - (1-mask) * loc) * (1-mask) * inverse(scale)) % vocab_size
+  ```
+
+  For the reverse pass, the flow computes:
+
+  ```none
+  net = layer(mask * inputs)
+  loc, scale = tf.split(net, 2, axis=-1)
+  loc = tf.argmax(loc, axis=-1)
+  scale = tf.argmax(scale, axis=-1)
+  outputs = ((1-mask) * loc + (1-mask) * scale * inputs) % vocab_size
+  ```
+
+  The modular arithmetic happens in one-hot space.
+
+  If `x` is a discrete random variable, the induced probability mass function on
+  the outputs `y = flow(x)` is
+
+  ```none
+  p(y) = p(flow.reverse(y)).
+  ```
+
+  The location-only transform is always invertible ([integers modulo
+  `vocab_size` form an additive group](
+  https://en.wikipedia.org/wiki/Modular_arithmetic)). The transform with a scale
+  is invertible if the scale and `vocab_size` are coprime (see
+  [prime fields](https://en.wikipedia.org/wiki/Finite_field)).
+  """
+
+  def __init__(self, layer, mask, temperature, **kwargs):
+    """Constructs flow.
+
+    Args:
+      layer: Two-headed masked network taking the inputs and returning a
+        real-valued Tensor of shape `[..., length, 2*vocab_size]`.
+        Alternatively, `layer` may return a Tensor of shape
+        `[..., length, vocab_size]` to be used as the location transform; the
+        scale transform will be hard-coded to 1.
+      mask: binary Tensor of shape `[length]` forming the bipartite assignment.
+      temperature: Positive value determining bias of gradient estimator.
+      **kwargs: kwargs of parent class.
+    """
+    super(DiscreteBipartiteFlow, self).__init__(**kwargs)
+    self.layer = layer
+    self.mask = mask
+    self.temperature = temperature
+
+  def build(self, input_shape):
+    input_shape = tf.TensorShape(input_shape)
+    self.vocab_size = input_shape[-1]
+    if isinstance(self.vocab_size, tf.Dimension):
+      self.vocab_size = self.vocab_size.value
+    if self.vocab_size is None:
+      raise ValueError('The last dimension of the inputs to '
+                       '`DiscreteBipartiteFlow` should be defined. Found '
+                       '`None`.')
+    self.built = True
+
+  def __call__(self, inputs, *args, **kwargs):
+    if not isinstance(inputs, ed.RandomVariable):
+      return super(DiscreteBipartiteFlow, self).__call__(
+          inputs, *args, **kwargs)
+    return TransformedRandomVariable(inputs, self)
+
+  def call(self, inputs, **kwargs):
+    """Forward pass for bipartite generation."""
+    inputs = tf.convert_to_tensor(inputs)
+    batch_ndims = inputs.shape.ndims - 2
+    mask = tf.reshape(tf.cast(self.mask, inputs.dtype),
+                      [1] * batch_ndims + [-1, 1])
+    masked_inputs = mask * inputs
+    net = self.layer(masked_inputs, **kwargs)
+    if net.shape[-1] == 2 * self.vocab_size:
+      raise NotImplementedError()
+    elif net.shape[-1] == self.vocab_size:
+      loc = net
+    else:
+      raise ValueError('Output of layer does not have compatible dimensions.')
+    loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
+    masked_outputs = (1. - mask) * one_hot_minus(inputs, loc)
+    outputs = masked_inputs + masked_outputs
+    return outputs
+
+  def reverse(self, inputs, **kwargs):
+    """Reverse pass for the inverse bipartite transformation."""
+    if not self.built:
+      self._maybe_build(inputs)
+
+    inputs = tf.convert_to_tensor(inputs)
+    batch_ndims = inputs.shape.ndims - 2
+    mask = tf.reshape(tf.cast(self.mask, inputs.dtype),
+                      [1] * batch_ndims + [-1, 1])
+    masked_inputs = mask * inputs
+    net = self.layer(masked_inputs, **kwargs)
+    if net.shape[-1] == 2 * self.vocab_size:
+      raise NotImplementedError()
+    elif net.shape[-1] == self.vocab_size:
+      loc = net
+    else:
+      raise ValueError('Output of layer does not have compatible dimensions.')
+    loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
+    masked_outputs = (1. - mask) * one_hot_add(loc, inputs)
+    outputs = masked_inputs + masked_outputs
+    return outputs
+
+  def log_det_jacobian(self, inputs):
+    return tf.cast(0, inputs.dtype)
+
+
+class SinkhornAutoregressiveFlow(tf.keras.layers.Layer):
+  """A discrete reversible layer using Sinkhorn normalization for permutations.
+
+  The flow takes as input a one-hot Tensor of shape `[..., length, vocab_size]`.
+  The flow returns a Tensor of same shape and dtype. (To enable gradients, the
+  input must have float dtype.)
+  """
+
+  def __init__(self, layer, temperature, **kwargs):
+    """Constructs flow.
+
+    Args:
+      layer: Masked network taking inputs with shape `[..., length, vocab_size]`
+        and returning a real-valued Tensor of shape
+        `[..., length, vocab_size ** 2]`. Sinkhorn iterations are applied to
+        each `layer` output to produce permutation matrices.
+      temperature: Positive value determining bias of gradient estimator.
+      **kwargs: kwargs of parent class.
+    """
+    super(SinkhornAutoregressiveFlow, self).__init__(**kwargs)
+    self.layer = layer
+    self.temperature = temperature
+
+  def build(self, input_shape):
+    input_shape = tf.TensorShape(input_shape)
+    self.vocab_size = input_shape[-1]
+    if isinstance(self.vocab_size, tf.Dimension):
+      self.vocab_size = self.vocab_size.value
+    if self.vocab_size is None:
+      raise ValueError('The last dimension of the inputs to '
+                       '`DiscreteAutoregressiveFlow` should be defined. Found '
+                       '`None`.')
+    self.built = True
+
+  def __call__(self, inputs, *args, **kwargs):
+    if not isinstance(inputs, ed.RandomVariable):
+      return super(SinkhornAutoregressiveFlow, self).__call__(
+          inputs, *args, **kwargs)
+    return TransformedRandomVariable(inputs, self)
+
+  def call(self, inputs, **kwargs):
+    """Forward pass for left-to-right autoregressive generation."""
+    inputs = tf.convert_to_tensor(inputs)
+    length = inputs.shape[-2].value
+    if length is None:
+      raise NotImplementedError('length dimension must be known.')
+    # Form initial sequence tensor of shape [..., 1, vocab_size]. In a loop, we
+    # incrementally build a Tensor of shape [..., t, vocab_size] as t grows.
+    outputs = self._initial_call(inputs[..., 0, :], length, **kwargs)
+    for t in range(1, length):
+      outputs = self._per_timestep_call(outputs,
+                                        inputs[..., t, :],
+                                        length,
+                                        t,
+                                        **kwargs)
+    return outputs
+
+  def _initial_call(self, new_inputs, length, **kwargs):
+    """Returns Tensor of shape [..., 1, vocab_size].
+
+    Args:
+      new_inputs: Tensor of shape [..., vocab_size], the new input to generate
+        its output.
+      length: Length of final desired sequence.
+      **kwargs: Optional keyword arguments to layer.
+    """
+    inputs = new_inputs[..., tf.newaxis, :]
+    # TODO(trandustin): To handle variable lengths, extend MADE to subset its
+    # input and output layer weights rather than pad inputs.
+    batch_ndims = inputs.shape.ndims - 2
+    padded_inputs = tf.pad(
+        inputs, [[0, 0]] * batch_ndims + [[0, length - 1], [0, 0]])
+    temperature = 1.
+    logits = self.layer(padded_inputs / temperature, **kwargs)
+    logits = logits[..., 0, :][..., tf.newaxis, :]
+    logits = tf.reshape(
+        logits,
+        logits.shape[:-1].concatenate([self.vocab_size, self.vocab_size]))
+    soft = sinkhorn(logits)
+    hard = tf.cast(soft_to_hard_permutation(soft), inputs.dtype)
+    hard = tf.reshape(hard, logits.shape)
+    # Inverse of permutation matrix is its transpose.
+    # inputs is [batch_size, timestep + 1, vocab_size].
+    # hard is [batch_size, timestep + 1, vocab_size, vocab_size].
+    outputs = tf.matmul(inputs[..., tf.newaxis, :],
+                        hard,
+                        transpose_b=True)[..., 0, :]
+    return outputs
+
+  def _per_timestep_call(self,
+                         current_outputs,
+                         new_inputs,
+                         length,
+                         timestep,
+                         **kwargs):
+    """Returns Tensor of shape [..., timestep+1, vocab_size].
+
+    Args:
+      current_outputs: Tensor of shape [..., timestep, vocab_size], the so-far
+        generated sequence Tensor.
+      new_inputs: Tensor of shape [..., vocab_size], the new input to generate
+        its output given current_outputs.
+      length: Length of final desired sequence.
+      timestep: Current timestep.
+      **kwargs: Optional keyword arguments to layer.
+    """
+    inputs = tf.concat([current_outputs,
+                        new_inputs[..., tf.newaxis, :]], axis=-2)
+    # TODO(trandustin): To handle variable lengths, extend MADE to subset its
+    # input and output layer weights rather than pad inputs.
+    batch_ndims = inputs.shape.ndims - 2
+    padded_inputs = tf.pad(
+        inputs, [[0, 0]] * batch_ndims + [[0, length - timestep - 1], [0, 0]])
+    logits = self.layer(padded_inputs, **kwargs)
+    logits = logits[..., :(timestep+1), :]
+    logits = tf.reshape(
+        logits,
+        logits.shape[:-1].concatenate([self.vocab_size, self.vocab_size]))
+    soft = sinkhorn(logits / self.temperature)
+    hard = tf.cast(soft_to_hard_permutation(soft), inputs.dtype)
+    hard = tf.reshape(hard, logits.shape)
+    # Inverse of permutation matrix is its transpose.
+    # inputs is [batch_size, timestep + 1, vocab_size].
+    # hard is [batch_size, timestep + 1, vocab_size, vocab_size].
+    new_outputs = tf.matmul(inputs[..., tf.newaxis, :],
+                            hard,
+                            transpose_b=True)[..., 0, :]
+    outputs = tf.concat([current_outputs,
+                         new_outputs[..., -1, :][..., tf.newaxis, :]], axis=-2)
+    if not tf.executing_eagerly():
+      outputs.set_shape([None] * batch_ndims + [timestep+1, self.vocab_size])
+    return outputs
+
+  def reverse(self, inputs, **kwargs):
+    """Reverse pass returning the inverse autoregressive transformation."""
+    if not self.built:
+      self._maybe_build(inputs)
+
+    logits = self.layer(inputs, **kwargs)
+    logits = tf.reshape(
+        logits,
+        logits.shape[:-1].concatenate([self.vocab_size, self.vocab_size]))
+    soft = sinkhorn(logits / self.temperature, n_iters=20)
+    hard = soft_to_hard_permutation(soft)
+    hard = tf.reshape(hard, logits.shape)
+    # Recover the permutation by right-multiplying by the permutation matrix.
+    outputs = tf.matmul(inputs[..., tf.newaxis, :], hard)[..., 0, :]
+    return outputs
+
+  def log_det_jacobian(self, inputs):
+    return tf.cast(0, inputs.dtype)
+
+
+def soft_to_hard_permutation(inputs):
+  """Returns permutation matrices by solving a matching problem.
+
+  Solves linear sum assignment to convert doubly-stochastic matrices to
+  permutation matrices. It uses scipy.optimize.linear_sum_assignment to solve
+  the optimization problem max_P sum_i,j M_i,j P_i,j with P a permutation
+  matrix. Notice the negative sign; the reason, the original function solves a
+  minimization problem.
+
+  Code is adapted from Mena et al. [1].
+
+  [1] Gonzalo Mena, David Belanger, Scott Linderman, Jasper Snoek.
+  Learning latent permutations with Gumbel-Sinkhorn networks. International
+  Conference on Learning Representations, 2018.
+
+  Args:
+    inputs: A `Tensor` with shape `[:, vocab_size, vocab_size]` that is
+      doubly-stochastic in its last two dimensions.
+
+  Returns:
+    outputs: A hard permutation `Tensor` with the same shape as `inputs` (in
+      other words the last two dimensions are doubly-stochastic and each element
+      is 0 or 1).
+  """
+
+  def hungarian(x):
+    if x.ndim == 2:
+      x = np.reshape(x, [1, x.shape[0], x.shape[1]])
+    sol = np.zeros((x.shape[0], x.shape[1]), dtype=np.int32)
+    for i in range(x.shape[0]):
+      sol[i, :] = linear_sum_assignment(-x[i, :])[1].astype(np.int32)
+    return sol
+
+  vocab_size = inputs.shape[-1]
+  # Note: tf.py_func isn't currently supported on headless GPUs.
+  # TODO(vafa): Fix tf.py_func headless GPU bug.
+  permutation_lists = tf.py_func(hungarian, [inputs], tf.int32)
+  hard = tf.one_hot(permutation_lists, depth=vocab_size)
+  outputs = tf.stop_gradient(hard - inputs) + inputs
+  return outputs
+
+
+def one_hot_argmax(inputs, temperature, axis=-1):
+  """Returns one-hot of argmax with backward pass set to softmax-temperature."""
+  vocab_size = inputs.shape[-1]
+  hard = tf.one_hot(tf.argmax(inputs, axis=axis),
+                    depth=vocab_size,
+                    axis=axis,
+                    dtype=inputs.dtype)
+  soft = tf.nn.softmax(inputs / temperature, axis=axis)
+  outputs = soft + tf.stop_gradient(hard - soft)
+  return outputs
+
+
+def one_hot_add(inputs, shift):
+  """Performs (inputs + shift) % vocab_size in the one-hot space.
+
+  Args:
+    inputs: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
+      Tensor.
+    shift: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
+      Tensor specifying how much to shift the corresponding one-hot vector in
+      inputs. Soft values perform a "weighted shift": for example,
+      shift=[0.2, 0.3, 0.5] performs a linear combination of 0.2 * shifting by
+      zero; 0.3 * shifting by one; and 0.5 * shifting by two.
+
+  Returns:
+    Tensor of same shape and dtype as inputs.
+  """
+  # Compute circular 1-D convolution with shift as the kernel.
+  inputs = tf.cast(inputs, tf.complex64)
+  shift = tf.cast(shift, tf.complex64)
+  return tf.real(tf.signal.ifft(tf.signal.fft(inputs) * tf.signal.fft(shift)))
+
+
+def one_hot_minus(inputs, shift):
+  """Performs (inputs - shift) % vocab_size in the one-hot space.
+
+  Args:
+    inputs: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
+      Tensor.
+    shift: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
+      Tensor specifying how much to shift the corresponding one-hot vector in
+      inputs. Soft values perform a "weighted shift": for example,
+      shift=[0.2, 0.3, 0.5] performs a linear combination of 0.2 * shifting by
+      zero; 0.3 * shifting by one; and 0.5 * shifting by two.
+
+  Returns:
+    Tensor of same shape and dtype as inputs.
+  """
+  # TODO(trandustin): Implement with circular conv1d.
+  inputs = tf.convert_to_tensor(inputs)
+  shift = tf.cast(shift, inputs.dtype)
+  vocab_size = inputs.shape[-1].value
+  # Form a [..., vocab_size, vocab_size] matrix. Each batch element of
+  # inputs will vector-matrix multiply the vocab_size x vocab_size matrix. This
+  # "shifts" the inputs batch element by the corresponding shift batch element.
+  shift_matrix = tf.stack([tf.roll(shift, i, axis=-1)
+                           for i in range(vocab_size)], axis=-2)
+  outputs = tf.einsum('...v,...uv->...u', inputs, shift_matrix)
+  return outputs
+
+
+def py_multiplicative_inverse(a, n):
+  """Multiplicative inverse of a modulo n (in Python).
+
+  Implements extended Euclidean algorithm.
+
+  Args:
+    a: int-like np.ndarray.
+    n: int.
+
+  Returns:
+    Multiplicative inverse as an int32 np.ndarray with same shape as a.
+  """
+  batched_a = np.asarray(a, dtype=np.int32)
+  batched_inverse = []
+  for a in np.nditer(batched_a):
+    inverse = 0
+    new_inverse = 1
+    remainder = n
+    new_remainder = a
+    while new_remainder != 0:
+      quotient = remainder // new_remainder
+      (inverse, new_inverse) = (new_inverse, inverse - quotient * new_inverse)
+      (remainder, new_remainder) = (new_remainder,
+                                    remainder - quotient * new_remainder)
+    if remainder > 1:
+      return ValueError(
+          'Inverse for {} modulo {} does not exist.'.format(a, n))
+    if inverse < 0:
+      inverse += n
+    batched_inverse.append(inverse)
+  return np.asarray(batched_inverse, dtype=np.int32).reshape(batched_a.shape)
+
+
+def multiplicative_inverse(a, n):
+  """Multiplicative inverse of a modulo n.
+
+  Args:
+    a: Tensor of shape [..., vocab_size]. It denotes an integer in the one-hot
+      space.
+    n: int Tensor of shape [...].
+
+  Returns:
+    Tensor of same shape and dtype as a.
+  """
+  a = tf.convert_to_tensor(a)
+  n = tf.convert_to_tensor(n)
+  vocab_size = a.shape[-1].value
+  a_dtype = a.dtype
+  sparse_a = tf.argmax(a, axis=-1)
+  sparse_outputs = tf.py_func(
+      py_multiplicative_inverse, [sparse_a, n], tf.int32)
+  sparse_outputs.set_shape(sparse_a.shape)
+  outputs = tf.one_hot(sparse_outputs, depth=vocab_size, dtype=a_dtype)
+  return outputs
+
+
 class ActNorm(tf.keras.layers.Layer):
   """Actnorm, an affine reversible layer (Prafulla and Kingma, 2018).
 
diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
index 8917e9f50..68084f929 100644
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -30,10 +30,392 @@
 tf.compat.v1.enable_eager_execution()
 
 
+def _log_prob(self, x):
+  """Re-implementation of OneHotCategorical._log_prob for gradients wrt x."""
+  x = self._assert_valid_sample(x)
+  logits = self.logits
+  if (not x.shape.is_fully_defined() or
+      not logits.shape.is_fully_defined() or
+      x.shape != logits.shape):
+    logits = tf.ones_like(x, dtype=logits.dtype) * logits
+    x = tf.ones_like(logits, dtype=x.dtype) * x
+
+  logits_shape = tf.shape(tf.reduce_sum(logits, -1))
+  logits_2d = tf.reshape(logits, [-1, self.event_size])
+  x_2d = tf.reshape(x, [-1, self.event_size])
+  ret = -tf.nn.softmax_cross_entropy_with_logits_v2(
+      labels=x_2d, logits=logits_2d)
+  ret = tf.reshape(ret, logits_shape)
+  return ret
+
+tfp.distributions.OneHotCategorical._log_prob = _log_prob  # monkey patch
 
 
 class ReversibleLayersTest(parameterized.TestCase, tf.test.TestCase):
 
+  @parameterized.parameters(
+      # TODO(trandustin): Enable test.
+      # (False,),
+      (True,),
+  )
+  @test_utils.run_in_graph_and_eager_modes()
+  def testDiscreteAutoregressiveFlowCall(self, loc_only):
+    batch_size = 3
+    vocab_size = 79
+    length = 5
+    if loc_only:
+      units = vocab_size
+    else:
+      units = 2 * vocab_size
+    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
+    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
+    layer = reversible.DiscreteAutoregressiveFlow(
+        reversible.MADE(units, []), 1.)
+    outputs = layer(inputs)
+    self.evaluate(tf.global_variables_initializer())
+    outputs_val = self.evaluate(outputs)
+    self.assertEqual(outputs_val.shape, (batch_size, length, vocab_size))
+    self.assertAllGreaterEqual(outputs_val, 0)
+    self.assertAllLessEqual(outputs_val, vocab_size - 1)
+
+  @parameterized.parameters(
+      # TODO(trandustin): Enable test.
+      # (False,),
+      (True,),
+  )
+  @test_utils.run_in_graph_and_eager_modes()
+  def testDiscreteAutoregressiveFlowSample(self, loc_only):
+    batch_size = 5
+    length = 2
+    vocab_size = 2
+    if loc_only:
+      units = vocab_size
+    else:
+      units = 2 * vocab_size
+    layer = reversible.DiscreteAutoregressiveFlow(
+        reversible.MADE(units, []), 1.)
+    logits = tf.tile(tf.random_normal([length, vocab_size])[tf.newaxis],
+                     [batch_size, 1, 1])
+    base = tfp.edward2.OneHotCategorical(logits=logits, dtype=tf.float32)
+    outputs = layer(base)
+    _ = outputs.value  # need to do this to instantiate tf.variables
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(outputs)
+    self.assertEqual(res.shape, (batch_size, length, vocab_size))
+    self.assertAllGreaterEqual(res, 0)
+    self.assertAllLessEqual(res, vocab_size - 1)
+
+  @parameterized.parameters(
+      # TODO(trandustin): Enable test.
+      # (False,),
+      (True,),
+  )
+  @test_utils.run_in_graph_and_eager_modes()
+  def testDiscreteAutoregressiveFlowInverse(self, loc_only):
+    batch_size = 2
+    vocab_size = 79
+    length = 5
+    if loc_only:
+      units = vocab_size
+    else:
+      units = 2 * vocab_size
+    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
+    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
+    layer = reversible.DiscreteAutoregressiveFlow(
+        reversible.MADE(units, []), 1.)
+    rev_fwd_inputs = layer.reverse(layer(inputs))
+    fwd_rev_inputs = layer(layer.reverse(inputs))
+    self.evaluate(tf.global_variables_initializer())
+    inputs_val, rev_fwd_inputs_val, fwd_rev_inputs_val = self.evaluate(
+        [inputs, rev_fwd_inputs, fwd_rev_inputs])
+    self.assertAllClose(inputs_val, rev_fwd_inputs_val)
+    self.assertAllClose(inputs_val, fwd_rev_inputs_val)
+
+  @parameterized.parameters(
+      # TODO(trandustin): Enable test.
+      # (False,),
+      (True,),
+  )
+  @test_utils.run_in_graph_and_eager_modes()
+  def testDiscreteAutoregressiveFlowRandomVariable(self, loc_only):
+    batch_size = 2
+    length = 4
+    vocab_size = 5
+    if loc_only:
+      units = vocab_size
+    else:
+      units = 2 * vocab_size
+    base = tfp.edward2.OneHotCategorical(logits=tf.random_normal([batch_size,
+                                                                  length,
+                                                                  vocab_size]),
+                                         dtype=tf.float32)
+    flow = reversible.DiscreteAutoregressiveFlow(
+        reversible.MADE(units, [16, 16]), 1.)
+    flow_rv = flow(base)
+    self.assertEqual(flow_rv.dtype, tf.float32)
+
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(flow_rv)
+    self.assertEqual(res.shape, (batch_size, length, vocab_size))
+    self.assertAllGreaterEqual(res, 0)
+    self.assertAllLessEqual(res, vocab_size - 1)
+
+    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
+    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
+    outputs = flow(inputs)
+    rev_outputs = flow.reverse(outputs)
+    inputs_val, rev_outputs_val = self.evaluate([inputs, rev_outputs])
+    self.assertAllClose(inputs_val, rev_outputs_val)
+
+    inputs_log_prob = base.distribution.log_prob(inputs)
+    outputs_log_prob = flow_rv.distribution.log_prob(outputs)
+    res1, res2 = self.evaluate([inputs_log_prob, outputs_log_prob])
+    self.assertEqual(res1.shape, (batch_size, length))
+    self.assertAllClose(res1, res2)
+
+  @parameterized.parameters(
+      # TODO(trandustin): Enable test.
+      # (False,),
+      (True,),
+  )
+  @test_utils.run_in_graph_mode_only()
+  def testDiscreteAutoregressiveFlowReverseGradients(self, loc_only):
+    batch_size = 2
+    length = 4
+    vocab_size = 2
+    if loc_only:
+      units = vocab_size
+    else:
+      units = 2 * vocab_size
+    base = tfp.edward2.OneHotCategorical(
+        logits=tf.random_normal([batch_size, length, vocab_size]))
+    flow = reversible.DiscreteAutoregressiveFlow(
+        reversible.MADE(units, [16, 16]), 1.)
+    flow_rv = flow(base)
+    features = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
+    features = tf.one_hot(features, depth=vocab_size, dtype=tf.float32)
+    loss = -tf.reduce_sum(flow_rv.distribution.log_prob(features))
+    grads = tf.gradients(loss, flow.layer.weights)
+    self.evaluate(tf.global_variables_initializer())
+    _ = self.evaluate(grads)
+    for grad in grads:
+      self.assertIsNotNone(grad)
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testOneHotAddExactHard(self):
+    inputs = tf.constant([[0., 1., 0.],
+                          [0., 0., 1.]])
+    shift = tf.constant([[0., 1., 0.],
+                         [1., 0., 0.]])
+
+    outputs = reversible.one_hot_add(inputs, shift)
+    outputs_val = self.evaluate(outputs)
+    self.assertAllEqual(outputs_val, np.array([[0., 0., 1.],
+                                               [0., 0., 1.]]))
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testOneHotMinusExactHard(self):
+    inputs = tf.constant([[0., 1., 0.],
+                          [0., 0., 1.]])
+    shift = tf.constant([[0., 1., 0.],
+                         [1., 0., 0.]])
+
+    outputs = reversible.one_hot_minus(inputs, shift)
+    outputs_val = self.evaluate(outputs)
+    self.assertAllEqual(outputs_val, np.array([[1., 0., 0.],
+                                               [0., 0., 1.]]))
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testOneHotAddExactSoft(self):
+    inputs = tf.constant([[0., 1., 0.],
+                          [0., 0., 1.]])
+    shift = tf.constant([[0.1, 0.6, 0.3],
+                         [0.2, 0.4, 0.4]])
+
+    outputs = reversible.one_hot_add(inputs, shift)
+
+    shift_zero = inputs
+    shift_one = np.array([[0., 0., 1.],
+                          [1., 0., 0.]])
+    shift_two = np.array([[1., 0., 0.],
+                          [0., 1., 0.]])
+    expected_outputs = (shift[..., 0][..., tf.newaxis] * shift_zero +
+                        shift[..., 1][..., tf.newaxis] * shift_one +
+                        shift[..., 2][..., tf.newaxis] * shift_two)
+
+    actual_outputs_val, expected_outputs_val = self.evaluate([
+        outputs, expected_outputs])
+    self.assertAllClose(actual_outputs_val, expected_outputs_val)
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testOneHotMinusExactSoft(self):
+    inputs = tf.constant([[0., 1., 0.],
+                          [0., 0., 1.]])
+    shift = tf.constant([[0.1, 0.6, 0.3],
+                         [0.2, 0.4, 0.4]])
+
+    outputs = reversible.one_hot_minus(inputs, shift)
+
+    shift_zero = inputs
+    shift_one = np.array([[1., 0., 0.],
+                          [0., 1., 0.]])
+    shift_two = np.array([[0., 0., 1.],
+                          [1., 0., 0.]])
+    expected_outputs = (shift[..., 0][..., tf.newaxis] * shift_zero +
+                        shift[..., 1][..., tf.newaxis] * shift_one +
+                        shift[..., 2][..., tf.newaxis] * shift_two)
+
+    actual_outputs_val, expected_outputs_val = self.evaluate([
+        outputs, expected_outputs])
+    self.assertAllEqual(actual_outputs_val, expected_outputs_val)
+
+  @parameterized.parameters(
+      (reversible.one_hot_add,),
+      (reversible.one_hot_minus,),
+  )
+  @test_utils.run_in_graph_and_eager_modes()
+  def testOneHotAddShapeHard(self, one_hot_add_fn):
+    batch_size = 2
+    length = 4
+    vocab_size = 5
+    inputs = tf.random_uniform(
+        [batch_size, length], minval=0, maxval=vocab_size, dtype=tf.int32)
+    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
+    shift = tf.random_uniform(
+        [batch_size, length], minval=0, maxval=vocab_size, dtype=tf.int32)
+    shift = tf.one_hot(shift, depth=vocab_size)
+
+    outputs = one_hot_add_fn(inputs, shift)
+    outputs_val = self.evaluate(outputs)
+    self.assertEqual(outputs_val.shape, (batch_size, length, vocab_size))
+
+  @parameterized.parameters(
+      (reversible.one_hot_add,),
+      (reversible.one_hot_minus,),
+  )
+  @test_utils.run_in_graph_and_eager_modes()
+  def testOneHotAddShapeSoft(self, one_hot_add_fn):
+    batch_size = 2
+    length = 4
+    vocab_size = 5
+    inputs = tf.random_uniform([batch_size, length, vocab_size])
+    shift = tf.random_uniform([batch_size, length, vocab_size])
+
+    outputs = one_hot_add_fn(inputs, shift)
+    outputs_val = self.evaluate(outputs)
+    self.assertEqual(outputs_val.shape, (batch_size, length, vocab_size))
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testMultiplicativeInverse(self):
+    batch_size = 3
+    vocab_size = 79
+    length = 5
+    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
+    one_hot_inputs = tf.one_hot(inputs, depth=vocab_size)
+
+    one_hot_inv = reversible.multiplicative_inverse(one_hot_inputs, vocab_size)
+    inv_inputs = tf.argmax(one_hot_inv, axis=-1)
+    inputs_inv_inputs = tf.floormod(inputs * inv_inputs, vocab_size)
+    inputs_inv_inputs_val = self.evaluate(inputs_inv_inputs)
+    self.assertAllEqual(inputs_inv_inputs_val, np.ones((batch_size, length)))
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testDiscreteBipartiteFlowCall(self):
+    batch_size = 3
+    vocab_size = 79
+    length = 5
+    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
+    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
+    layer = reversible.DiscreteBipartiteFlow(
+        tf.identity,
+        mask=tf.random_uniform([length], minval=0, maxval=2, dtype=tf.int32),
+        temperature=1.)
+    outputs = layer(inputs)
+    self.evaluate(tf.global_variables_initializer())
+    outputs_val = self.evaluate(outputs)
+    self.assertEqual(outputs_val.shape, (batch_size, length, vocab_size))
+    self.assertAllGreaterEqual(outputs_val, 0)
+    self.assertAllLessEqual(outputs_val, vocab_size - 1)
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testDiscreteBipartiteFlowInverse(self):
+    batch_size = 2
+    vocab_size = 79
+    length = 5
+    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
+    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
+    layer = reversible.DiscreteBipartiteFlow(
+        tf.identity,
+        mask=tf.random_uniform([length], minval=0, maxval=2, dtype=tf.int32),
+        temperature=1.)
+    rev_fwd_inputs = layer.reverse(layer(inputs))
+    fwd_rev_inputs = layer(layer.reverse(inputs))
+    self.evaluate(tf.global_variables_initializer())
+    inputs_val, rev_fwd_inputs_val, fwd_rev_inputs_val = self.evaluate(
+        [inputs, rev_fwd_inputs, fwd_rev_inputs])
+    self.assertAllClose(inputs_val, rev_fwd_inputs_val)
+    self.assertAllClose(inputs_val, fwd_rev_inputs_val)
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testSinkhornAutoregressiveFlowCall(self):
+    batch_size = 3
+    vocab_size = 79
+    length = 5
+    units = vocab_size ** 2
+    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
+    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
+    layer = reversible.SinkhornAutoregressiveFlow(
+        reversible.MADE(units, []), 1.)
+    outputs = layer(inputs)
+    self.evaluate(tf.global_variables_initializer())
+    outputs_val = self.evaluate(outputs)
+    self.assertEqual(outputs_val.shape, (batch_size, length, vocab_size))
+    self.assertAllGreaterEqual(outputs_val, 0)
+    self.assertAllLessEqual(outputs_val, vocab_size - 1)
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testDiscreteSinkhornFlowInverse(self):
+    batch_size = 2
+    vocab_size = 79
+    length = 5
+    units = vocab_size ** 2
+    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
+    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
+    layer = reversible.SinkhornAutoregressiveFlow(
+        reversible.MADE(units, []), 1.)
+    rev_fwd_inputs = layer.reverse(layer(inputs))
+    fwd_rev_inputs = layer(layer.reverse(inputs))
+    self.evaluate(tf.global_variables_initializer())
+    inputs_val, rev_fwd_inputs_val, fwd_rev_inputs_val = self.evaluate(
+        [inputs, rev_fwd_inputs, fwd_rev_inputs])
+    self.assertAllEqual(inputs_val, rev_fwd_inputs_val)
+    self.assertAllEqual(inputs_val, fwd_rev_inputs_val)
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testApproximatelyStochastic(self):
+    rng = np.random.RandomState(0)
+    tf.set_random_seed(1)
+    for dims in [2, 5, 10]:
+      for batch_size in [1, 2, 10]:
+        log_alpha = rng.randn(batch_size, dims, dims)
+        result = reversible.sinkhorn(log_alpha)
+        result_val = self.evaluate(result)
+        self.assertAllClose(np.sum(result_val, 1),
+                            np.tile([1.0], (batch_size, dims)),
+                            atol=1e-3)
+        self.assertAllClose(np.sum(result_val, 2),
+                            np.tile([1.0], (batch_size, dims)),
+                            atol=1e-3)
+
+  @test_utils.run_in_graph_mode_only()
+  def test_soft_to_hard_permutation(self):
+    """The solution of the matching for the identity matrix is range(N)."""
+    dims = 10
+    identity = np.eye(dims)
+    result_matching = reversible.soft_to_hard_permutation(identity)
+    result_matching_val = self.evaluate(result_matching)
+    self.assertAllEqual(result_matching_val[0], np.eye(dims))
+
   @test_utils.run_in_graph_and_eager_modes()
   def testActNorm(self):
     np.random.seed(83243)

From d546695261d3feefa330110438005702d9be89eb Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Fri, 31 May 2019 13:35:20 -0700
Subject: [PATCH 2085/2720] Internal.

PiperOrigin-RevId: 250946632
---
 tensor2tensor/layers/reversible_layers.py     | 92 +++++++++++++------
 .../layers/reversible_layers_test.py          | 34 +++++++
 2 files changed, 97 insertions(+), 29 deletions(-)

diff --git a/tensor2tensor/layers/reversible_layers.py b/tensor2tensor/layers/reversible_layers.py
index b45efba90..997720951 100644
--- a/tensor2tensor/layers/reversible_layers.py
+++ b/tensor2tensor/layers/reversible_layers.py
@@ -155,21 +155,18 @@ def _initial_call(self, new_inputs, length, **kwargs):
         inputs, [[0, 0]] * batch_ndims + [[0, length - 1], [0, 0]])
     net = self.layer(padded_inputs, **kwargs)
     if net.shape[-1] == 2 * self.vocab_size:
-      raise NotImplementedError()
-      # TODO(trandustin): Enable scale.
-      # loc, scale = tf.split(net, 2, axis=-1)
-      # scale = scale[..., 0, :][..., tf.newaxis, :]
-      # scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
-      # inverse_scale = multiplicative_inverse(scale, self.vocab_size)
+      loc, scale = tf.split(net, 2, axis=-1)
+      scale = scale[..., 0:1, :]
+      scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
+      inverse_scale = multiplicative_inverse(scale, self.vocab_size)
+      scaled_inputs = one_hot_multiply(inputs, inverse_scale)
     elif net.shape[-1] == self.vocab_size:
       loc = net
-      # inverse_scale = tf.ones_like(inputs)
+      scaled_inputs = inputs
     else:
       raise ValueError('Output of layer does not have compatible dimensions.')
     loc = loc[..., 0, :][..., tf.newaxis, :]
     loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-    # scaled_inputs = one_hot_multiply(inputs, inverse_scale)
-    scaled_inputs = inputs
     outputs = one_hot_minus(scaled_inputs, loc)
     return outputs
 
@@ -199,21 +196,18 @@ def _per_timestep_call(self,
         inputs, [[0, 0]] * batch_ndims + [[0, length - timestep - 1], [0, 0]])
     net = self.layer(padded_inputs, **kwargs)
     if net.shape[-1] == 2 * self.vocab_size:
-      raise NotImplementedError()
-      # TODO(trandustin): Enable scale.
-      # loc, scale = tf.split(net, 2, axis=-1)
-      # scale = scale[..., :(timestep+1), :]
-      # scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
-      # inverse_scale = multiplicative_inverse(scale, self.vocab_size)
+      loc, scale = tf.split(net, 2, axis=-1)
+      scale = scale[..., :(timestep+1), :]
+      scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
+      inverse_scale = multiplicative_inverse(scale, self.vocab_size)
+      scaled_inputs = one_hot_multiply(inputs, inverse_scale)
     elif net.shape[-1] == self.vocab_size:
       loc = net
-      # inverse_scale = tf.ones_like(inputs)
+      scaled_inputs = inputs
     else:
       raise ValueError('Output of layer does not have compatible dimensions.')
     loc = loc[..., :(timestep+1), :]
     loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-    # scaled_inputs = one_hot_multiply(inputs, inverse_scale)
-    scaled_inputs = inputs
     new_outputs = one_hot_minus(scaled_inputs, loc)
     outputs = tf.concat([current_outputs,
                          new_outputs[..., -1, :][..., tf.newaxis, :]], axis=-2)
@@ -228,18 +222,15 @@ def reverse(self, inputs, **kwargs):
 
     net = self.layer(inputs, **kwargs)
     if net.shape[-1] == 2 * self.vocab_size:
-      raise NotImplementedError()
-      # TODO(trandustin): Enable scale.
-      # loc, scale = tf.split(net, 2, axis=-2)
-      # scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
+      loc, scale = tf.split(net, 2, axis=-2)
+      scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
+      scaled_inputs = one_hot_multiply(inputs, scale)
     elif net.shape[-1] == self.vocab_size:
       loc = net
-      # scale = tf.ones_like(inputs)
+      scaled_inputs = inputs
     else:
       raise ValueError('Output of layer does not have compatible dimensions.')
     loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-    # scaled_inputs = one_hot_multiply(inputs, scale)
-    scaled_inputs = inputs
     outputs = one_hot_add(loc, scaled_inputs)
     return outputs
 
@@ -334,13 +325,17 @@ def call(self, inputs, **kwargs):
     masked_inputs = mask * inputs
     net = self.layer(masked_inputs, **kwargs)
     if net.shape[-1] == 2 * self.vocab_size:
-      raise NotImplementedError()
+      loc, scale = tf.split(net, 2, axis=-1)
+      scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
+      inverse_scale = multiplicative_inverse(scale, self.vocab_size)
+      scaled_inputs = one_hot_multiply(inputs, inverse_scale)
     elif net.shape[-1] == self.vocab_size:
       loc = net
+      scaled_inputs = inputs
     else:
       raise ValueError('Output of layer does not have compatible dimensions.')
     loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-    masked_outputs = (1. - mask) * one_hot_minus(inputs, loc)
+    masked_outputs = (1. - mask) * one_hot_minus(scaled_inputs, loc)
     outputs = masked_inputs + masked_outputs
     return outputs
 
@@ -356,13 +351,16 @@ def reverse(self, inputs, **kwargs):
     masked_inputs = mask * inputs
     net = self.layer(masked_inputs, **kwargs)
     if net.shape[-1] == 2 * self.vocab_size:
-      raise NotImplementedError()
+      loc, scale = tf.split(net, 2, axis=-2)
+      scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
+      scaled_inputs = one_hot_multiply(inputs, scale)
     elif net.shape[-1] == self.vocab_size:
       loc = net
+      scaled_inputs = inputs
     else:
       raise ValueError('Output of layer does not have compatible dimensions.')
     loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-    masked_outputs = (1. - mask) * one_hot_add(loc, inputs)
+    masked_outputs = (1. - mask) * one_hot_add(loc, scaled_inputs)
     outputs = masked_inputs + masked_outputs
     return outputs
 
@@ -626,6 +624,42 @@ def one_hot_minus(inputs, shift):
   return outputs
 
 
+def one_hot_multiply(inputs, scale):
+  """Performs (inputs * scale) % vocab_size in the one-hot space.
+
+  Args:
+    inputs: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
+      Tensor.
+    scale: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
+      Tensor specifying how much to scale the corresponding one-hot vector in
+      inputs. Soft values perform a "weighted scale": for example,
+      scale=[0.2, 0.3, 0.5] performs a linear combination of
+      0.2 * scaling by zero; 0.3 * scaling by one; and 0.5 * scaling by two.
+
+  Returns:
+    Tensor of same shape and dtype as inputs.
+  """
+  # TODO(trandustin): Implement with circular conv1d.
+  inputs = tf.convert_to_tensor(inputs)
+  scale = tf.cast(scale, inputs.dtype)
+  batch_shape = inputs.shape[:-1].as_list()
+  vocab_size = inputs.shape[-1].value
+  # Form a [..., vocab_size, vocab_size] tensor. The ith row of the
+  # batched vocab_size x vocab_size matrix represents scaling inputs by i.
+  permutation_matrix = tf.floormod(
+      tf.tile(tf.range(vocab_size)[:, tf.newaxis], [1, vocab_size]) *
+      tf.range(vocab_size)[tf.newaxis], vocab_size)
+  permutation_matrix = tf.one_hot(permutation_matrix, depth=vocab_size, axis=-1)
+  # Scale the inputs according to the permutation matrix of all possible scales.
+  scaled_inputs = tf.einsum('...v,avu->...au', inputs, permutation_matrix)
+  scaled_inputs = tf.concat([tf.zeros(batch_shape + [1, vocab_size]),
+                             scaled_inputs[..., 1:, :]], axis=-2)
+  # Reduce rows of the scaled inputs by the scale values. This forms a
+  # weighted linear combination of scaling by zero, scaling by one, and so on.
+  outputs = tf.einsum('...v,...vu->...u', scale, scaled_inputs)
+  return outputs
+
+
 def py_multiplicative_inverse(a, n):
   """Multiplicative inverse of a modulo n (in Python).
 
diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
index 68084f929..83de8a798 100644
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -225,6 +225,18 @@ def testOneHotMinusExactHard(self):
     self.assertAllEqual(outputs_val, np.array([[1., 0., 0.],
                                                [0., 0., 1.]]))
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testOneHotMultiplyExactHard(self):
+    inputs = tf.constant([[0., 1., 0.],
+                          [0., 0., 1.]])
+    scale = tf.constant([[0., 1., 0.],
+                         [0., 0., 1.]])
+
+    outputs = reversible.one_hot_multiply(inputs, scale)
+    outputs_val = self.evaluate(outputs)
+    self.assertAllEqual(outputs_val, np.array([[0., 1., 0.],
+                                               [0., 1., 0.]]))
+
   @test_utils.run_in_graph_and_eager_modes()
   def testOneHotAddExactSoft(self):
     inputs = tf.constant([[0., 1., 0.],
@@ -269,6 +281,28 @@ def testOneHotMinusExactSoft(self):
         outputs, expected_outputs])
     self.assertAllEqual(actual_outputs_val, expected_outputs_val)
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testOneHotMultiplyExactSoft(self):
+    inputs = tf.constant([[0., 1., 0.],
+                          [0., 0., 1.]])
+    scale = tf.constant([[0.1, 0.6, 0.3],
+                         [0.2, 0.4, 0.4]])
+
+    outputs = reversible.one_hot_multiply(inputs, scale)
+
+    scale_zero = np.array([[0., 0., 0.],
+                           [0., 0., 0.]])
+    scale_one = inputs
+    scale_two = np.array([[0., 0., 1.],
+                          [0., 1., 0.]])
+    expected_outputs = (scale[..., 0][..., tf.newaxis] * scale_zero +
+                        scale[..., 1][..., tf.newaxis] * scale_one +
+                        scale[..., 2][..., tf.newaxis] * scale_two)
+
+    actual_outputs_val, expected_outputs_val = self.evaluate([
+        outputs, expected_outputs])
+    self.assertAllEqual(actual_outputs_val, expected_outputs_val)
+
   @parameterized.parameters(
       (reversible.one_hot_add,),
       (reversible.one_hot_minus,),

From 03e8ede7164fa4a24de6bb1d10655936f8cbbfcd Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 31 May 2019 14:38:15 -0700
Subject: [PATCH 2086/2720] Naming consistency pass: change remaining
 "num_...s" to "n_...s".

PiperOrigin-RevId: 250958916
---
 .../chunked_transformer_imagenet64_8gb.gin    |  2 +-
 tensor2tensor/trax/inputs.py                  | 80 +++++++++----------
 tensor2tensor/trax/inputs_test.py             |  4 +-
 tensor2tensor/trax/models/transformer.py      |  6 +-
 tensor2tensor/trax/rlax/fake_env.py           |  4 +-
 tensor2tensor/trax/rlax/fake_env_test.py      |  4 +-
 tensor2tensor/trax/rlax/ppo.py                | 44 +++++-----
 tensor2tensor/trax/rlax/ppo_main.py           |  8 +-
 tensor2tensor/trax/rlax/ppo_test.py           | 26 +++---
 .../trax/rlax/ppo_training_loop_test.py       |  6 +-
 10 files changed, 92 insertions(+), 92 deletions(-)

diff --git a/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin b/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin
index 10b984fd1..e72c76d61 100644
--- a/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin
+++ b/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin
@@ -11,10 +11,10 @@ batch_fun.max_eval_length = 12288  # 64 * 64 * 3
 
 # Parameters for inputs:
 # ==============================================================================
-inputs.num_chunks = 64
 inputs.data_dir = None
 inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
 inputs.input_name = 'targets'
+inputs.n_chunks = 64
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 9b15dc830..2c82e51f9 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -54,18 +54,18 @@
 _MAX_SKIP_EXAMPLES = 1e5
 
 
-@gin.configurable(blacklist=["num_devices"])
-def inputs(num_devices, dataset_name, data_dir=None, input_name=None,
-           num_chunks=0, append_targets=False):
+@gin.configurable(blacklist=["n_devices"])
+def inputs(n_devices, dataset_name, data_dir=None, input_name=None,
+           n_chunks=0, append_targets=False):
   """Make Inputs for built-in datasets.
 
   Args:
-    num_devices: how many devices to build the inputs for.
+    n_devices: how many devices to build the inputs for.
     dataset_name: a TFDS or T2T dataset name. If it's a T2T dataset name, prefix
       with "t2t_".
     data_dir: data directory.
     input_name: optional, name of the inputs from the dictionary.
-    num_chunks: optional, into how many pieces should we chunk (large inputs).
+    n_chunks: optional, into how many pieces should we chunk (large inputs).
     append_targets: optional, instead of inputs return a pair (inputs, targets)
       which is useful for autoregressive models.
 
@@ -77,17 +77,17 @@ def inputs(num_devices, dataset_name, data_dir=None, input_name=None,
 
   (train_batches, train_eval_batches, eval_batches,
    input_name, input_shape) = _train_and_eval_batches(
-       dataset_name, data_dir, input_name, num_devices)
+       dataset_name, data_dir, input_name, n_devices)
 
   def numpy_stream(dataset):
     return dataset_to_stream(
         dataset, input_name,
-        num_chunks=num_chunks, append_targets=append_targets)
+        n_chunks=n_chunks, append_targets=append_targets)
 
-  if num_chunks > 0:
+  if n_chunks > 0:
     length = input_shape[0]
     input_shape = tuple(
-        [tuple([length // num_chunks] + list(input_shape)[1:])] * num_chunks)
+        [tuple([length // n_chunks] + list(input_shape)[1:])] * n_chunks)
   if append_targets:
     # TODO(lukaszkaiser): remove the assumption that input and target
     # shapes are the same, which is used below for now.
@@ -99,15 +99,15 @@ def numpy_stream(dataset):
                 input_shape=input_shape)
 
 
-@gin.configurable(blacklist=["num_devices"])
+@gin.configurable(blacklist=["n_devices"])
 def random_inputs(
-    num_devices,
+    n_devices,
     input_shape=gin.REQUIRED, input_dtype=np.int32, input_range=(0, 255),
     output_shape=gin.REQUIRED, output_dtype=np.int32, output_range=(0, 9)):
   """Make random Inputs for debugging.
 
   Args:
-    num_devices: how many devices to build the inputs for.
+    n_devices: how many devices to build the inputs for.
     input_shape: the shape of inputs (including batch dimension).
     input_dtype: the type of the inputs (int32 by default).
     input_range: the range of inputs (defaults to (0, 255)).
@@ -118,14 +118,14 @@ def random_inputs(
   Returns:
     trax.inputs.Inputs
   """
-  if input_shape[0] % num_devices != 0:
+  if input_shape[0] % n_devices != 0:
     tf.logging.fatal(
-        "num_devices[%d] should divide the first dimension of input_shape[%s]",
-        num_devices, input_shape)
-  if output_shape[0] % num_devices != 0:
+        "n_devices[%d] should divide the first dimension of input_shape[%s]",
+        n_devices, input_shape)
+  if output_shape[0] % n_devices != 0:
     tf.logging.fatal(
-        "num_devices[%d] should divide the first dimension of output_shape[%s]",
-        num_devices, output_shape)
+        "n_devices[%d] should divide the first dimension of output_shape[%s]",
+        n_devices, output_shape)
 
   def random_minibatches():
     """Generate a stream of random mini-batches."""
@@ -147,7 +147,7 @@ def random_minibatches():
                 input_shape=input_shape_without_batch)
 
 
-def dataset_to_stream(dataset, input_name, num_chunks=0, append_targets=False):
+def dataset_to_stream(dataset, input_name, n_chunks=0, append_targets=False):
   """Takes a tf.Dataset and creates a numpy stream of ready batches."""
   for example in tfds.as_numpy(dataset):
     inp, out = example[0][input_name], example[1]
@@ -158,9 +158,9 @@ def dataset_to_stream(dataset, input_name, num_chunks=0, append_targets=False):
       out = out.astype(np.uint32)
     if len(out.shape) > 1 and out.shape[-1] == 1:
       out = np.squeeze(out, axis=-1)
-    if num_chunks > 0:
-      inp = np.split(inp, num_chunks, axis=1)
-      out = np.split(out, num_chunks, axis=1)
+    if n_chunks > 0:
+      inp = np.split(inp, n_chunks, axis=1)
+      out = np.split(out, n_chunks, axis=1)
     if append_targets:
       inp = (inp, out)
     yield inp, out
@@ -186,7 +186,7 @@ def train_and_eval_dataset(dataset_name, data_dir, train_shuffle_files=True,
      * the train tf.Dataset
      * the eval tf.Dataset
      * information about features: a python dictionary with feature names
-         as keys and an object as value that provides .shape and .num_classes.
+         as keys and an object as value that provides .shape and .n_classes.
      * supervised_keys: information what's the input and what's the target,
          ie., a pair of lists with input and target feature names.
   """
@@ -214,9 +214,9 @@ def train_and_eval_dataset(dataset_name, data_dir, train_shuffle_files=True,
   return train, valid, info.features, keys
 
 
-def _make_info(shape_list, num_classes):
+def _make_info(shape_list, n_classes):
   """Create an info-like tuple for feature given some shapes and vocab size."""
-  feature_info = collections.namedtuple("FeatureInfo", ["shape", "num_classes"])
+  feature_info = collections.namedtuple("FeatureInfo", ["shape", "n_classes"])
   cur_shape = list(shape_list[0])
   # We need to merge the provided shapes, put None where they disagree.
   for shape in shape_list:
@@ -226,7 +226,7 @@ def _make_info(shape_list, num_classes):
       if cur_shape[i] is not None:
         if shape[i] != cur_shape[i]:
           cur_shape[i] = None
-  return feature_info(cur_shape, num_classes)
+  return feature_info(cur_shape, n_classes)
 
 
 def _select_features(example, feature_list=None):
@@ -279,20 +279,20 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
 
 
 @gin.configurable(blacklist=["dataset", "training", "shapes",
-                             "target_names", "num_devices"])
-def batch_fun(dataset, training, shapes, target_names, num_devices,
+                             "target_names", "n_devices"])
+def batch_fun(dataset, training, shapes, target_names, n_devices,
               batch_size_per_device=32, batch_size=None, eval_batch_size=32,
               bucket_length=32, buckets=None,
               buckets_include_inputs_in_length=False,
               batch_shuffle_size=128, max_eval_length=None):
   """Batching function."""
   del target_names
-  # Batch size is batch_size_per_device * num_devices unless given directly.
-  batch_size = batch_size or batch_size_per_device * num_devices
+  # Batch size is batch_size_per_device * n_devices unless given directly.
+  batch_size = batch_size or batch_size_per_device * n_devices
   # If bucketing is not specified, check if target shapes are variable.
   cur_batch_size = batch_size if training else eval_batch_size
-  # Make cur_batch_size divisible by num_devices.
-  cur_batch_size = max(cur_batch_size // num_devices, 1) * num_devices
+  # Make cur_batch_size divisible by n_devices.
+  cur_batch_size = max(cur_batch_size // n_devices, 1) * n_devices
   # Create heuristic buckets is none are specified.
   if buckets is None:
     variable_target_shapes = False
@@ -318,8 +318,8 @@ def batch_fun(dataset, training, shapes, target_names, num_devices,
                             cur_batch_size // 16, 1]
       if not training:
         bucket_batch_sizes[-2] = cur_batch_size // max_eval_length
-      # Make batch sizes divisible by num_devices.
-      bucket_batch_sizes = [max(b // num_devices, 1) * num_devices
+      # Make batch sizes divisible by n_devices.
+      bucket_batch_sizes = [max(b // n_devices, 1) * n_devices
                             for b in bucket_batch_sizes]
       buckets = (bucket_boundaries, bucket_batch_sizes)
 
@@ -406,7 +406,7 @@ def shuffle_and_batch_data(dataset,
                            target_names,
                            features_info,
                            training,
-                           num_devices,
+                           n_devices,
                            shuffle_buffer_size=1024,
                            preprocess_fun=no_preprocess):
   """Shuffle and batch the given dataset."""
@@ -429,24 +429,24 @@ def append_targets(example):
   shapes = {k: features_info[k].shape for k in features_info}
   shapes = (shapes, shapes[target_names[0]])
   dataset = dataset.shuffle(shuffle_buffer_size)
-  dataset = batch_fun(dataset, training, shapes, target_names, num_devices)
+  dataset = batch_fun(dataset, training, shapes, target_names, n_devices)
   return dataset.prefetch(2)
 
 
-def _train_and_eval_batches(dataset, data_dir, input_name, num_devices):
+def _train_and_eval_batches(dataset, data_dir, input_name, n_devices):
   """Return train and eval batches with input name and shape."""
   (train_data, eval_data, features_info, keys) = train_and_eval_dataset(
       dataset, data_dir)
   input_names, target_names = keys[0], keys[1]
   train_batches = shuffle_and_batch_data(
       train_data, target_names, features_info, training=True,
-      num_devices=num_devices)
+      n_devices=n_devices)
   train_eval_batches = shuffle_and_batch_data(  # Data for eval-on-train.
       train_data, target_names, features_info, training=False,
-      num_devices=num_devices)
+      n_devices=n_devices)
   eval_batches = shuffle_and_batch_data(
       eval_data, target_names, features_info, training=False,
-      num_devices=num_devices)
+      n_devices=n_devices)
   input_name = input_name or input_names[0]
   input_shape = features_info[input_name].shape
   return (train_batches, train_eval_batches, eval_batches,
diff --git a/tensor2tensor/trax/inputs_test.py b/tensor2tensor/trax/inputs_test.py
index 1d9f353d2..80c76b01d 100644
--- a/tensor2tensor/trax/inputs_test.py
+++ b/tensor2tensor/trax/inputs_test.py
@@ -55,7 +55,7 @@ def test_batch_fun(self):
       self.assertEqual(example[0].shape[0], 10)  # Batch size = 10.
     self.assertEqual(count, 1)  # Just one batch here.
 
-  def test_batch_fun_num_devices(self):
+  def test_batch_fun_n_devices(self):
     dataset = test_dataset_ints([32])
     dataset = dataset.repeat(9)
     batches = inputs.batch_fun(
@@ -63,7 +63,7 @@ def test_batch_fun_num_devices(self):
     count = 0
     for example in tfds.as_numpy(batches):
       count += 1
-      # Batch size adjusted to be divisible by num_devices.
+      # Batch size adjusted to be divisible by n_devices.
       self.assertEqual(example[0].shape[0], 9)
     self.assertEqual(count, 1)  # Just one batch here.
 
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 1bca46d78..ae195b155 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -66,7 +66,7 @@ def EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
 
 
 def TransformerEncoder(vocab_size,
-                       num_classes=10,
+                       n_classes=10,
                        d_feature=512,
                        d_feedforward=2048,
                        n_layers=6,
@@ -78,7 +78,7 @@ def TransformerEncoder(vocab_size,
 
   Args:
     vocab_size: int: vocab size
-    num_classes: how many classes on output
+    n_classes: how many classes on output
     d_feature: int:  depth of embedding
     d_feedforward: int: depth of feed-forward layer
     n_layers: int: number of encoder/decoder layers
@@ -102,7 +102,7 @@ def TransformerEncoder(vocab_size,
       tl.Select(0),  # Drop mask.
       tl.LayerNorm(),
       tl.Mean(axis=1),  # Average on length.
-      tl.Dense(num_classes),
+      tl.Dense(n_classes),
       tl.LogSoftmax(),
   ]
 
diff --git a/tensor2tensor/trax/rlax/fake_env.py b/tensor2tensor/trax/rlax/fake_env.py
index cc4b6fd98..cff0fac7b 100644
--- a/tensor2tensor/trax/rlax/fake_env.py
+++ b/tensor2tensor/trax/rlax/fake_env.py
@@ -33,14 +33,14 @@ class FakeEnv(object):
 
   def __init__(self,
                input_shape=(4,),
-               num_actions=2,
+               n_actions=2,
                done_time_step=None,
                done_action=None):
     self._input_shape = input_shape
     self._done_time_step = done_time_step
     self._done_action = done_action
     self._t = 0
-    self.action_space = gym.spaces.Discrete(num_actions)
+    self.action_space = gym.spaces.Discrete(n_actions)
     self.observation_space = gym.spaces.Box(
         low=-1.0, high=1.0, shape=input_shape)
 
diff --git a/tensor2tensor/trax/rlax/fake_env_test.py b/tensor2tensor/trax/rlax/fake_env_test.py
index 34aba74d7..8c9ee9771 100644
--- a/tensor2tensor/trax/rlax/fake_env_test.py
+++ b/tensor2tensor/trax/rlax/fake_env_test.py
@@ -27,7 +27,7 @@ class FakeEnvTest(test.TestCase):
 
   def test_done_action(self):
     env = fake_env.FakeEnv(input_shape=(2, 3),
-                           num_actions=10,
+                           n_actions=10,
                            done_time_step=None,
                            done_action=9)
     env.reset()
@@ -44,7 +44,7 @@ def test_done_action(self):
 
   def test_done_time_step(self):
     env = fake_env.FakeEnv(input_shape=(2, 3),
-                           num_actions=10,
+                           n_actions=10,
                            done_time_step=10,
                            done_action=None)
     env.reset()
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 93a7d4180..2af68ddd7 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -76,7 +76,7 @@
 LAMBDA = 0.95
 EPSILON = 0.1
 EPOCHS = 50  # 100
-NUM_OPTIMIZER_STEPS = 100
+N_OPTIMIZER_STEPS = 100
 PRINT_EVERY_OPTIMIZER_STEP = 20
 BATCH_TRAJECTORIES = 32
 
@@ -124,7 +124,7 @@ def optimizer_fn(net_params, step_size=1e-3):
 # Any other option?
 def collect_trajectories(env,
                          policy_fn,
-                         num_trajectories=1,
+                         n_trajectories=1,
                          policy=env_problem_utils.CATEGORICAL_SAMPLING,
                          max_timestep=None,
                          boundary=20,
@@ -136,7 +136,7 @@ def collect_trajectories(env,
   Args:
     env: A gym env interface, for now this is not-batched.
     policy_fn: observations(B,T+1) -> log-probabs(B,T+1, A) callable.
-    num_trajectories: int, number of trajectories.
+    n_trajectories: int, number of trajectories.
     policy: string, "greedy", "epsilon-greedy", or "categorical-sampling" i.e.
       how to use the policy_fn to return an action.
     max_timestep: int or None, the index of the maximum time-step at which we
@@ -159,10 +159,10 @@ def collect_trajectories(env,
 
   assert isinstance(env, env_problem.EnvProblem)
   # This is an env_problem, run its collect function.
-  trajs, num_done = env_problem_utils.play_env_problem_with_policy(
+  trajs, n_done = env_problem_utils.play_env_problem_with_policy(
       env,
       policy_fn,
-      num_trajectories=num_trajectories,
+      num_trajectories=n_trajectories,
       max_timestep=max_timestep,
       boundary=boundary,
       policy_sampling=policy,
@@ -170,7 +170,7 @@ def collect_trajectories(env,
       reset=reset,
       rng=rng)
   # Skip returning raw_rewards here, since they aren't used.
-  return [(t[0], t[1], t[2]) for t in trajs], num_done
+  return [(t[0], t[1], t[2]) for t in trajs], n_done
 
 
 # This function can probably be simplified, ask how?
@@ -710,13 +710,13 @@ def evaluate_policy(eval_env,
                     get_predictions,
                     boundary,
                     max_timestep=20000,
-                    num_evals=1,
+                    n_evals=1,
                     rng=None):
   """Evaluate the policy."""
 
   avg_rewards = collections.defaultdict(float)
   avg_rewards_unclipped = collections.defaultdict(float)
-  for _ in range(num_evals):
+  for _ in range(n_evals):
     for policy in [
         env_problem_utils.CATEGORICAL_SAMPLING,
         env_problem_utils.GUMBEL_SAMPLING,
@@ -737,8 +737,8 @@ def evaluate_policy(eval_env,
 
   # Now average these out.
   for k in avg_rewards:
-    avg_rewards[k] /= num_evals
-    avg_rewards_unclipped[k] /= num_evals
+    avg_rewards[k] /= n_evals
+    avg_rewards_unclipped[k] /= n_evals
 
   return avg_rewards, avg_rewards_unclipped
 
@@ -772,7 +772,7 @@ def training_loop(
     policy_and_value_net_fn=None,
     policy_and_value_optimizer_fn=None,
     batch_size=BATCH_TRAJECTORIES,
-    num_optimizer_steps=NUM_OPTIMIZER_STEPS,
+    n_optimizer_steps=N_OPTIMIZER_STEPS,
     print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
     target_kl=0.01,
     boundary=20,
@@ -790,7 +790,7 @@ def training_loop(
     done_frac_for_policy_save=0.5,
     enable_early_stopping=True,
     env_name=None,
-    num_evals=1,
+    n_evals=1,
 ):
   """Runs the training loop for PPO, with fixed policy and value nets."""
   assert env
@@ -841,7 +841,7 @@ def training_loop(
   (policy_and_value_opt_state, policy_and_value_opt_update,
    policy_and_value_get_params) = policy_and_value_optimizer
 
-  num_trajectories_done = 0
+  n_trajectories_done = 0
   last_saved_at = 0
 
   logging.info("Starting the PPO training loop.")
@@ -874,7 +874,7 @@ def get_predictions(observations, rng=None):
           get_predictions,
           boundary,
           max_timestep=max_timestep_eval,
-          num_evals=num_evals,
+          n_evals=n_evals,
           rng=key)
       for k, v in avg_reward.items():
         eval_sw.scalar("eval/mean_reward/%s" % k, v, step=i)
@@ -889,10 +889,10 @@ def get_predictions(observations, rng=None):
     trajectory_collection_start_time = time.time()
     logging.vlog(1, "Epoch [% 6d] collecting trajectories.", i)
     jax_rng_key, key = jax_random.split(jax_rng_key)
-    trajs, num_done = collect_trajectories(
+    trajs, n_done = collect_trajectories(
         env,
         policy_fn=get_predictions,
-        num_trajectories=batch_size,
+        n_trajectories=batch_size,
         max_timestep=max_timestep,
         boundary=boundary,
         rng=key,
@@ -993,8 +993,8 @@ def get_predictions(observations, rng=None):
     jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)
     logging.vlog(1, "Policy and Value Optimization")
     optimization_start_time = time.time()
-    keys = jax_random.split(key1, num=num_optimizer_steps)
-    for j in range(num_optimizer_steps):
+    keys = jax_random.split(key1, num=n_optimizer_steps)
+    for j in range(n_optimizer_steps):
       k1, k2, k3 = jax_random.split(keys[j], num=3)
       t = time.time()
       # Update the optimizer state.
@@ -1037,7 +1037,7 @@ def get_predictions(observations, rng=None):
 
       t2 = time.time()
       if (((j + 1) % print_every_optimizer_steps == 0) or
-          (j == num_optimizer_steps - 1) or early_stopping):
+          (j == n_optimizer_steps - 1) or early_stopping):
         # Compute and log the loss.
         (loss_combined, loss_ppo, loss_value, entropy_bonus) = (
             combined_loss(
@@ -1077,9 +1077,9 @@ def get_predictions(observations, rng=None):
     # Also don't save too frequently, enforce a minimum gap.
     # Or if this is the last iteration.
     policy_save_start_time = time.time()
-    num_trajectories_done += num_done
+    n_trajectories_done += n_done
     # TODO(afrozm): Refactor to trax.save_state.
-    if (((num_trajectories_done >= done_frac_for_policy_save * batch_size)
+    if (((n_trajectories_done >= done_frac_for_policy_save * batch_size)
          and (i - last_saved_at > eval_every_n)
          and (((i + 1) % eval_every_n == 0)))
         or (i == epochs - 1)):
@@ -1092,7 +1092,7 @@ def get_predictions(observations, rng=None):
       for path in old_model_files:
         gfile.remove(path)
       # Reset this number.
-      num_trajectories_done = 0
+      n_trajectories_done = 0
       last_saved_at = i
     policy_save_time = get_time(policy_save_start_time)
 
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 0a66a08ae..478701b34 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -101,7 +101,7 @@
     "two.")
 
 # Number of optimizer steps of the combined net, policy net and value net.
-flags.DEFINE_integer("num_optimizer_steps", 100, "Number of optimizer steps.")
+flags.DEFINE_integer("n_optimizer_steps", 100, "Number of optimizer steps.")
 flags.DEFINE_integer(
     "print_every_optimizer_steps", 1,
     "How often to log during the policy optimization process.")
@@ -126,7 +126,7 @@
 flags.DEFINE_bool("xm", False, "Copy atari roms?")
 flags.DEFINE_integer("eval_every_n", 100, "How frequently to eval the policy.")
 flags.DEFINE_integer("eval_batch_size", 4, "Batch size for evaluation.")
-flags.DEFINE_integer("num_evals", 1, "Number of times to evaluate.")
+flags.DEFINE_integer("n_evals", 1, "Number of times to evaluate.")
 flags.DEFINE_float(
     "done_frac_for_policy_save", 0.5,
     "Fraction of the trajectories that should be done to "
@@ -223,7 +223,7 @@ def run_training_loop():
         epochs=FLAGS.epochs,
         policy_and_value_net_fn=policy_and_value_net_fn,
         policy_and_value_optimizer_fn=policy_and_value_optimizer_fn,
-        num_optimizer_steps=FLAGS.num_optimizer_steps,
+        n_optimizer_steps=FLAGS.n_optimizer_steps,
         print_every_optimizer_steps=FLAGS.print_every_optimizer_steps,
         batch_size=FLAGS.batch_size,
         target_kl=FLAGS.target_kl,
@@ -241,7 +241,7 @@ def run_training_loop():
         eval_every_n=FLAGS.eval_every_n,
         done_frac_for_policy_save=FLAGS.done_frac_for_policy_save,
         eval_env=eval_env,
-        num_evals=FLAGS.num_evals,
+        n_evals=FLAGS.n_evals,
         env_name=str(FLAGS.env_problem_name),
     )
 
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 4658c7ad2..2972a73fb 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -36,9 +36,9 @@ def setUp(self):
   def test_policy_and_value_net(self):
     observation_shape = (3, 4, 5)
     batch_observation_shape = (-1, -1) + observation_shape
-    num_actions = 2
+    n_actions = 2
     pnv_params, pnv_apply = ppo.policy_and_value_net(
-        self.rng_key, batch_observation_shape, num_actions,
+        self.rng_key, batch_observation_shape, n_actions,
         lambda: [layers.Flatten(num_axis_to_keep=2)])
     batch = 2
     time_steps = 10
@@ -48,14 +48,14 @@ def test_policy_and_value_net(self):
 
     # Output is a list, first is probab of actions and the next is value output.
     self.assertEqual(2, len(pnv_output))
-    self.assertEqual((batch, time_steps, num_actions), pnv_output[0].shape)
+    self.assertEqual((batch, time_steps, n_actions), pnv_output[0].shape)
     self.assertEqual((batch, time_steps, 1), pnv_output[1].shape)
 
   def test_pad_trajectories(self):
     observation_shape = (2, 3, 4)
     trajectories = []
-    num_trajectories = 7
-    num_actions = 10
+    n_trajectories = 7
+    n_actions = 10
 
     # Time-steps are between [min_allowable_time_step, max_allowable_time_step]
     max_allowable_time_step = 19
@@ -67,8 +67,8 @@ def test_pad_trajectories(self):
     # Bucket length.
     bucket_length = 15
 
-    # Make `num_trajectories` random trajectories.
-    for i in range(num_trajectories):
+    # Make `n_trajectories` random trajectories.
+    for i in range(n_trajectories):
       time_steps = np.random.randint(min_allowable_time_step,
                                      max_allowable_time_step + 1)
       if time_steps > max_time_step:
@@ -77,7 +77,7 @@ def test_pad_trajectories(self):
           0, 255, size=(time_steps + 1,) + observation_shape).astype(np.uint8)
       rewards = np.random.uniform(size=(time_steps,)).astype(np.float32)
       actions = np.random.randint(
-          0, num_actions, size=(time_steps,)).astype(np.int32)
+          0, n_actions, size=(time_steps,)).astype(np.int32)
       trajectories.append((observations, rewards, actions))
 
     # Now pad these trajectories.
@@ -96,16 +96,16 @@ def test_pad_trajectories(self):
 
     # Expectations on the padded shapes.
     self.assertEqual(padded_observations.shape, (
-        num_trajectories,
+        n_trajectories,
         expected_padding + 1,
     ) + observation_shape)
-    self.assertEqual(padded_actions.shape, (num_trajectories, expected_padding))
-    self.assertEqual(padded_rewards.shape, (num_trajectories, expected_padding))
-    self.assertEqual(reward_mask.shape, (num_trajectories, expected_padding))
+    self.assertEqual(padded_actions.shape, (n_trajectories, expected_padding))
+    self.assertEqual(padded_rewards.shape, (n_trajectories, expected_padding))
+    self.assertEqual(reward_mask.shape, (n_trajectories, expected_padding))
 
     # Assert that the padding lengths and reward mask are consistent.
     self.assertAllEqual(
-        np.full((num_trajectories,), expected_padding),
+        np.full((n_trajectories,), expected_padding),
         np.array(np.sum(reward_mask, axis=1)) + pad_lengths)
 
   def test_rewards_to_go(self):
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index eae8dd3e2..5e9a17a49 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -59,19 +59,19 @@ def test_training_loop(self):
     with self.tmp_dir() as output_dir:
       env = self.get_wrapped_env("CartPole-v0", 2)
       eval_env = self.get_wrapped_env("CartPole-v0", 2)
-      num_epochs = 2
+      n_epochs = 2
       batch_size = 2
       # Run the training loop.
       ppo.training_loop(
           env=env,
           eval_env=eval_env,
-          epochs=num_epochs,
+          epochs=n_epochs,
           policy_and_value_net_fn=functools.partial(
               ppo.policy_and_value_net,
               bottom_layers_fn=lambda: [layers.Dense(1)]),
           policy_and_value_optimizer_fn=ppo.optimizer_fn,
           batch_size=batch_size,
-          num_optimizer_steps=1,
+          n_optimizer_steps=1,
           output_dir=output_dir,
           env_name="CartPole-v0",
           random_seed=0)

From cefd39eb0af617af9a045090f97c88c501096b4e Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Sun, 2 Jun 2019 22:06:20 -0700
Subject: [PATCH 2087/2720] Use JAX abstact-eval to calculate shapes in Trax.

PiperOrigin-RevId: 251164989
---
 tensor2tensor/trax/backend.py                 |  16 +-
 tensor2tensor/trax/layers/attention.py        |  54 +----
 tensor2tensor/trax/layers/base.py             | 146 ++++++++----
 tensor2tensor/trax/layers/combinators.py      | 209 +++++-------------
 tensor2tensor/trax/layers/combinators_test.py |  21 --
 tensor2tensor/trax/layers/convolution.py      |  80 -------
 tensor2tensor/trax/layers/core.py             |  36 +--
 tensor2tensor/trax/layers/core_test.py        |   4 +-
 tensor2tensor/trax/layers/pooling.py          |  20 +-
 tensor2tensor/trax/models/atari_cnn_test.py   |   2 +-
 tensor2tensor/trax/models/neural_gpu.py       |   2 +-
 .../models/research/chunked_transformer.py    |  27 +--
 tensor2tensor/trax/models/transformer.py      |   4 +-
 tensor2tensor/trax/rlax/ppo.py                |   4 +-
 tensor2tensor/trax/rlax/ppo_test.py           |   4 +-
 tensor2tensor/trax/trax.py                    |  11 +-
 16 files changed, 208 insertions(+), 432 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index 4fc0431df..c3a0cee17 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import contextlib
 import gin
 
 import jax
@@ -111,12 +112,21 @@ def __getattr__(self, attr):
 
 
 
-default_backend_name = "jax"
+override_backend_name = None
 
 
 @gin.configurable()
-def backend(name=None):
-  name = name or default_backend_name
+def backend(name="jax"):
+  name = name if not override_backend_name else override_backend_name
   if name == "numpy":
     return _NUMPY_BACKEND
   return _JAX_BACKEND
+
+
+@contextlib.contextmanager
+def use_backend(name):
+  global override_backend_name
+  prev_name = override_backend_name
+  override_backend_name = name
+  yield
+  override_backend_name = prev_name
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index e66a76fe3..00763b551 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -27,29 +27,20 @@
 from tensor2tensor.trax.layers import core
 
 
-@base.layer(output_shape=lambda shape, axis=-1: (1, shape[axis], shape[axis]))
+@base.layer()
 def CausalMask(x, params, axis=-1, **kwargs):
   del params, kwargs
   size = x.shape[axis]
   return onp.tril(onp.ones((1, size, size), dtype=onp.bool_), k=0)
 
 
-@base.layer(output_shape=lambda shape, pad=0: (shape[0], 1, 1, shape[-1]))
+@base.layer()
 def PaddingMask(x, params, pad=0, **kwargs):
   del params, kwargs
   return np.reshape(x != pad, (x.shape[0], 1, 1, x.shape[-1]))
 
 
-def EncoderDecoderMaskShape(inputs):
-  """Helper: shape for encoder-decoder mask."""
-  (padding_mask_shape, decoder_input_shape) = inputs
-  batch_size = padding_mask_shape[0]
-  input_length = padding_mask_shape[-1]
-  target_length = decoder_input_shape[1]
-  return (batch_size, 1, target_length, input_length)
-
-
-@base.layer(output_shape=EncoderDecoderMaskShape, stack_items_to_pass=0)
+@base.layer(stack_items_to_pass=0)
 def EncoderDecoderMask(x, **unused_kwargs):
   """Make encoder-decoder mask from a padding mask and decoder input."""
   (padding_mask, decoder_input) = x
@@ -111,42 +102,7 @@ def DotProductAttention(query, key, value, mask, dropout, mode, rng):
   return out
 
 
-# TODO(lukaszkaiser): make this a layer.
-def PureDotProductAttention(dropout=0.0, mode='train'):
-  """Pure single-headed self-attention.
-
-  Args:
-    dropout: float: dropout rate
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    Pure single-headed attention layer. (No Dense transforms on input.)
-  """
-  def init_fn(_, input_shapes):  # pylint: disable=invalid-name
-    q_shape, _, v_shape, _ = input_shapes
-    output_shape = q_shape[:-1] + (v_shape[-1],)
-    return output_shape, ()
-  def apply_fn(params, inputs, **kwargs):  # pylint: disable=invalid-name
-    del params
-    q, k, v, mask = inputs
-    rng = kwargs.get('rng', None)
-    return DotProductAttention(q, k, v, mask,
-                               dropout=dropout, mode=mode, rng=rng)
-  return init_fn, apply_fn
-
-
-def _multihead_attention_output_shape(  # pylint: disable=invalid-name
-    input_shapes, **unused_kwargs):
-  """Helper: calculate multihead attention output shape."""
-  q_shape = input_shapes[0]  # Inputs are (q, k, v, mask).
-  v_shape = input_shapes[2]  # Inputs are (q, k, v, mask).
-  mask_shape = input_shapes[3]
-  res_shape = list(q_shape[:-1]) + [v_shape[-1]]
-  return tuple(res_shape), mask_shape
-
-
-@base.layer(output_shape=_multihead_attention_output_shape,
-            stack_items_to_pass=4)
+@base.layer(stack_items_to_pass=4)
 def PureMultiHeadedAttention(x, params, n_heads=8, dropout=0.0, mode='train',
                              **kwargs):
   """Pure transformer-style multi-headed attention.
@@ -234,7 +190,7 @@ def MultiHeadedAttention(
   ]
 
 
-@base.layer()
+@base.layer(input_is_int=True, stack_items_to_pass=0)
 def ShiftRight(x, **unused_kwargs):
   """Layer to shift the tensor to the right by padding on axis 1."""
   if not isinstance(x, (list, tuple)):  # non-chunked inputs
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 978cee0d1..3b58d156d 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -22,8 +22,11 @@
 import inspect
 import traceback
 
+import jax
+from jax.interpreters import partial_eval as pe
+
 import numpy as onp
-from tensor2tensor.trax.backend import random
+from tensor2tensor.trax import backend
 
 
 class Layer(object):
@@ -43,20 +46,6 @@ def call(self, x, params=(), **kwargs):
     """Call this layer in input x using the given parameters."""
     raise NotImplementedError
 
-  def output_shape_fn(self, input_shape):
-    """The shape of the output of this layer given the shape of the input.
-
-    Note that all arguments and return values can be tuples or dictionaries
-    or arbitrary nested structures composed of tuples and dictionaries.
-
-    Args:
-      input_shape: a tuple representing the shape of the input.
-
-    Returns:
-      The shape of the output.
-    """
-    raise NotImplementedError
-
   def new_parameters(self, input_shape, rng):
     """Create new parameters for the layer given an input shape and rng.
 
@@ -72,22 +61,45 @@ def new_parameters(self, input_shape, rng):
     """
     raise NotImplementedError
 
+  # TODO(lukaszkaiser): re-visit the 2 items below in the future.
   def stack_items_to_pass(self):
     """How many of the top stack items do we process."""
     return 0
 
+  def default_input_is_int(self):
+    """Whether the default inputs are ints or floats."""
+    return False
+
   # End of subclassing interface, all functions below are internal.
 
-  def output_shape(self, input_shape):
-    """Same as self.output_shape but with better error reporting."""
+  def output_shape(self, input_shape_and_type, params):
+    """Output shape and type for this layer given input shape and type.
+
+    Note that all arguments and return values can be tuples or dictionaries
+    or arbitrary nested structures composed of tuples and dictionaries.
+
+    Args:
+      input_shape_and_type: a ShapeType with shape and type of the input.
+      params: parameters for this layer.
+
+    Returns:
+      The shape and type of the output.
+    """
     try:
-      is_list = isinstance(input_shape, (list, tuple))
-      is_list = is_list and isinstance(input_shape[0], (list, tuple))
-      n = self.stack_items_to_pass() if is_list else 0
-      return _apply_to_first_n(self.output_shape_fn, input_shape, n)
+      with backend.use_backend('jax'):
+        rng = backend.random.get_prng(0)
+        def call_on_input(x, params):
+          f = lambda y: self.call(y, params=params, rng=rng)
+          n = self.stack_items_to_pass() if isinstance(x, (list, tuple)) else 0
+          return _apply_to_first_n(f, x, n)
+        params_shapes = nested_map(
+            params, lambda x: ShapeType(shape=x.shape, tp=x.dtype))
+        s = _eval_on_shapes(call_on_input, input_shape_and_type, params_shapes)
+      return s
     except Exception:
-      name, trace = self.__class__.__name__, _short_traceback()
-      raise LayerError(name, 'output_shape', self._caller, input_shape, trace)
+      name, trace = self.__class__.__name__, _short_traceback(skip=3)
+      raise LayerError(name, 'output_shape', self._caller,
+                       input_shape_and_type, trace)
 
   def initialize(self, input_shape, rng):
     """Initialize the layer given an input shape and rng.
@@ -143,6 +155,17 @@ def __call__(self, x, params=(), **kwargs):
       raise LayerError(name, 'call', self._caller, shapes(x), trace)
 
 
+class ShapeType(object):
+  """Store shape and type."""
+
+  def __init__(self, shape, tp):
+    self.shape = shape
+    self.tp = tp
+
+  def __repr__(self):
+    return '[shape:' + str(self.shape) + ', type:' + str(self.tp) + ']'
+
+
 class LayerError(Exception):
   """Exception raised in the layer stack.
 
@@ -170,6 +193,39 @@ def message(self):
     return prefix + caller + shapes_str + self._traceback
 
 
+# TODO(lukaszkaiser): remove this function once JAX has an analogue.
+def _eval_on_shapes(f, *args):
+  """Evaluate f given only shapes and types."""
+  def abstractify(x):
+    return jax.abstract_arrays.raise_to_shaped(jax.core.get_aval(x))
+
+  def make_array(arg):
+    return backend.numpy.zeros(shape=arg.shape, dtype=arg.tp)
+
+  def turn_back_into_pytree(x):
+    if isinstance(x, jax.core.JaxTuple):
+      return tuple([turn_back_into_pytree(y) for y in x])
+    return x
+
+  def get_shapes_and_types(x):
+    if isinstance(x, jax.core.AbstractTuple):
+      return tuple([get_shapes_and_types(y) for y in x])
+    return ShapeType(x.shape, x.dtype)
+
+  def f_jaxtuple(*jaxtuple_args):
+    args = map(turn_back_into_pytree, jaxtuple_args)
+    out = f(*args)
+    res, _ = jax.api_util.pytree_to_jaxtupletree(out)
+    return res
+
+  args_arrays = nested_map(args, make_array)
+  jaxtuple_args, _ = jax.util.unzip2(
+      map(jax.api_util.pytree_to_jaxtupletree, args_arrays))
+  res = pe.abstract_eval_fun(f_jaxtuple, *map(abstractify, jaxtuple_args))
+
+  return get_shapes_and_types(res)
+
+
 def _apply_to_first_n(f, x, n):
   """Helper: apply f to first n elements on the stack x if n > 0."""
   if n < 1:
@@ -279,7 +335,7 @@ def _short_traceback(skip=7):
 # Decorator for making layers from functions.
 
 
-def layer(output_shape=None, new_parameters=None, stack_items_to_pass=1):
+def layer(new_parameters=None, stack_items_to_pass=1, input_is_int=False):
   """Create a layer class from a function."""
   def layer_decorator(call):
     """Decorating the call function."""
@@ -288,11 +344,9 @@ def stack_items_to_pass_fn(self):
       del self
       return stack_items_to_pass
 
-    def output_shape_fn(self, input_shape):
-      if output_shape is None:
-        return input_shape
-      kwargs = self._init_kwargs  # pylint: disable=protected-access
-      return output_shape(input_shape, **kwargs)
+    def default_input_is_int_fn(self):
+      del self
+      return input_is_int
 
     def new_parameters_fn(self, input_shape, rng):
       if new_parameters is None:
@@ -310,15 +364,13 @@ def call_fn(self, x, params=(), **kwargs):
 
     # Set doc for python help.
     call_fn.__doc__ = call.__doc__
-    if output_shape is None:
-      output_shape_fn.__doc__ = output_shape.__doc__
     if new_parameters is None:
       new_parameters_fn.__doc__ = new_parameters.__doc__
 
     # Create the class.
     cls = type(call.__name__, (Layer,),
                {'call': call_fn,
-                'output_shape_fn': output_shape_fn,
+                'default_input_is_int': default_input_is_int_fn,
                 'new_parameters': new_parameters_fn,
                 'stack_items_to_pass': stack_items_to_pass_fn})
 
@@ -346,8 +398,8 @@ def _random_inputs(input_shape, rng, integer_inputs=False):
   if not isinstance(input_shape, dict) and isinstance(input_shape[0], int):
     # Non-nested shape, create a random tuple.
     if not integer_inputs:
-      return random.uniform(rng, input_shape, minval=-1.0, maxval=1.0)
-    return random.bernoulli(rng, 0.5, input_shape).astype(onp.int32)
+      return backend.random.uniform(rng, input_shape, minval=-1.0, maxval=1.0)
+    return backend.random.bernoulli(rng, 0.5, input_shape).astype(onp.int32)
   elif isinstance(input_shape, list):  # Nested shape: list.
     return [_random_inputs(shape, rng, integer_inputs) for shape in input_shape]
   elif isinstance(input_shape, tuple):  # Nested shape: tuple.
@@ -359,12 +411,32 @@ def _random_inputs(input_shape, rng, integer_inputs=False):
     raise TypeError(type(input_shape))
 
 
+def to_shape_and_type(x_shapes, integers):
+  """Make a shape-and-type tuple from shapes."""
+  if isinstance(x_shapes, dict):  # Nested shape: dict.
+    return {k: to_shape_and_type(x_shapes[k], integers) for k in x_shapes}
+  if isinstance(x_shapes, onp.ndarray):  # Numpy array shape
+    return ShapeType(shape=x_shapes.tolist(),
+                     tp=onp.int32 if integers else onp.float32)
+  if isinstance(x_shapes[0], (int, onp.int32, onp.int64)):
+    return ShapeType(shape=x_shapes,
+                     tp=onp.int32 if integers else onp.float32)
+  if isinstance(x_shapes, list):  # Nested shape: list.
+    return [to_shape_and_type(s, integers) for s in x_shapes]
+  if isinstance(x_shapes, tuple):  # Nested shape: tuple.
+    return tuple([to_shape_and_type(s, integers) for s in x_shapes])
+  assert False  # Should never get here.
+
+
 def check_shape_agreement(layer_instance, input_shape, integer_inputs=False):
   """Check if layer.output_shape agrees with the actual output shape."""
-  rng1, rng2, rng3 = random.split(random.get_prng(0), 3)
-  output_shape = layer_instance.output_shape(input_shape)
-  output_shape = nested_map(output_shape, int)  # Make non-numpy.
+  rng1, rng2, rng3 = backend.random.split(backend.random.get_prng(0), 3)
   params = layer_instance.initialize(input_shape, rng1)
+  input_shape_and_type = to_shape_and_type(input_shape, integer_inputs)
+  output_shape_and_type = layer_instance.output_shape(
+      input_shape_and_type, params)
+  output_shape = nested_map(output_shape_and_type, lambda x: x.shape)
+  output_shape = nested_map(output_shape, int)  # Make non-numpy.
   inputs = _random_inputs(input_shape, rng2, integer_inputs=integer_inputs)
   result = layer_instance(inputs, params, rng=rng3)
   result_shape = shapes(result)
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 8a2b80125..a24a0f093 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -95,71 +95,57 @@ def call(self, x, params=(), **kwargs):
       x = layer(x, p, rng=rng, **kwargs)
     return x
 
-  def output_shape_fn(self, input_shape):
-    cur_shape = input_shape
-    for layer in self._layers:
-      cur_shape = layer.output_shape(cur_shape)
-    return cur_shape
+  def output_shape(self, input_shape_and_type, params):
+    cur_shape_and_type = input_shape_and_type
+    for layer, param in zip(self._layers, params):
+      cur_shape_and_type = layer.output_shape(cur_shape_and_type, param)
+    return cur_shape_and_type
+
+  def default_input_is_int(self):
+    if self._nlayers == 0:
+      return False
+    return self._layers[0].default_input_is_int()
 
   def new_parameters(self, input_shape, rng):
     params = []
-    cur_shape = input_shape
+    cur_shape_and_type = base.to_shape_and_type(
+        input_shape, self.default_input_is_int())
     for layer in self._layers:
       rng, layer_rng = backend.random.split(rng)
+      cur_shape = base.nested_map(cur_shape_and_type, lambda x: x.shape)
       param = layer.initialize(cur_shape, layer_rng)
-      cur_shape = layer.output_shape(cur_shape)
+      pparam = layer._params   # pylint: disable=protected-access
+      cur_shape_and_type = layer.output_shape(cur_shape_and_type, pparam)
       params.append(param)
     return params
 
 
-def _print_shape(x, message='PrintShape'):  # pylint: disable=invalid-name
-  print(message + ' ; stack shape = ' + str(x))
-  return x
-
-
-@base.layer(output_shape=_print_shape, stack_items_to_pass=0)
+@base.layer(stack_items_to_pass=0)
 def PrintShape(x, message='PrintShape', **unused_kwargs):
   """No-op layer that prints the shape of the stack."""
-  _print_shape(base.shapes(x), message=message)
+  print(message + ' ; stack shape = ' + str(base.shapes(x)))
   return x
 
 
-def _dup(x):  # pylint: disable=invalid-name
-  """Helper: copy the top element of a list or a tuple."""
+@base.layer(stack_items_to_pass=0)
+def Dup(x, **unused_kwargs):
+  """Duplicate (copy) the first element on the stack."""
   if isinstance(x, list):
     return [x[0]] + x
   assert isinstance(x, tuple)
   return tuple([x[0]] + list(x))
 
 
-@base.layer(output_shape=_dup, stack_items_to_pass=0)
-def Dup(x, **unused_kwargs):
-  """Duplicate (copy) the first element on the stack."""
-  return _dup(x)
-
-
-def _swap(x):  # pylint: disable=invalid-name
-  """Helper: swap the top two elements of a list or a tuple."""
+@base.layer(stack_items_to_pass=0)
+def Swap(x, **unused_kwargs):
+  """Swap the first two element on the stack."""
   if isinstance(x, list):
     return [x[1], x[0]] + x[2:]
   assert isinstance(x, tuple)
   return tuple([x[1], x[0]] + list(x[2:]))
 
 
-@base.layer(output_shape=_swap, stack_items_to_pass=0)
-def Swap(x, **unused_kwargs):
-  """Swap the first two element on the stack."""
-  return _swap(x)
-
-
-def _top_shape(x_shape):  # pylint: disable=invalid-name
-  """Helper: shape of top element of a stack."""
-  if isinstance(x_shape[0], (list, tuple)):
-    return x_shape[0]
-  return x_shape
-
-
-@base.layer(output_shape=_top_shape, stack_items_to_pass=0)
+@base.layer(stack_items_to_pass=0)
 def _Top(x, **unused_kwargs):
   """Top element from the stack."""
   if isinstance(x, (list, tuple)):
@@ -167,33 +153,17 @@ def _Top(x, **unused_kwargs):
   return x
 
 
-def _drop(x):  # pylint: disable=invalid-name
-  """Helper: pop top element of a stack (make it a non-list if length is 1)."""
+@base.layer(stack_items_to_pass=0)
+def Drop(x, **unused_kwargs):
+  """Drop first element from the stack."""
   result = x[1:]
   if len(result) == 1:
     return result[0]
   return result
 
 
-@base.layer(output_shape=_drop, stack_items_to_pass=0)
-def Drop(x, **unused_kwargs):
-  """Drop first element from the stack."""
-  return _drop(x)
-
-
-def _flatten_shape(x_shape):  # pylint: disable=invalid-name
-  """Helper: shape of the flatten operation."""
-  shapes = []
-  for shape in x_shape:
-    if isinstance(shape[0], (list, tuple)):
-      shapes.extend(shape)
-    else:
-      shapes.append(shape)
-  return tuple(shapes)
-
-
-@base.layer(output_shape=_flatten_shape, stack_items_to_pass=0)
-def Flatten(xs, **unused_kwargs):
+@base.layer(stack_items_to_pass=0)
+def FlattenList(xs, **unused_kwargs):
   """Flatten lists."""
   return tuple(_deep_flatten(xs))
 
@@ -262,11 +232,6 @@ def call(self, x, params=(), **kwargs):
       return x
     return base.nested_map(self._output, lambda i: self._map(x, i))
 
-  def output_shape_fn(self, input_shape):
-    if self._output is None:
-      return input_shape
-    return base.nested_map(self._output, lambda i: self._map(input_shape, i))
-
   def new_parameters(self, input_shape, rng):
     return ()
 
@@ -275,23 +240,18 @@ class Branch(base.Layer):
   """Combinator for applying layers to copies of the input.
 
   This layer is often used to create parallel towers in neural networks:
-  * Branch(Copy(), Copy()) -- creates a pair with copied input
   * Branch(main, shortcut) -- start a residual tower (see Residual below)
 
   Args:
     *layers: a sequence of layers.
-    **kwlayers: a dictionary of layers.
 
   Returns:
     A new layer in which each of the given layers has been applied to
     a copy of the input independently.
   """
 
-  def __init__(self, *layers, **kwlayers):
+  def __init__(self, *layers):
     super(Branch, self).__init__()
-    if layers and kwlayers:
-      raise ValueError('Cannot specify a Branch with both a list and dict.')
-    layers = layers or kwlayers
     layers = _ensure_sublayers(layers)
     self._nlayers = len(layers)
     self._layers = layers
@@ -302,45 +262,26 @@ def call(self, x, params=(), **kwargs):
     rngs = (None,) * self._nlayers
     if rng is not None:
       rngs = backend.random.split(rng, self._nlayers)
-    # If layers are a list or a tuple, just apply them.
     if isinstance(self._layers, (list, tuple)):
       res = [layer(x, params=p, rng=r, **kwargs)
              for layer, p, r in zip(self._layers, params, rngs)]
       return tuple(res)
-    # If layers are a dictionary, apply to matching keys.
-    assert isinstance(self._layers, dict)
-    result, counter = {}, 0
-    for k in self._layers:
-      result[k] = self._layers[k](
-          x, params=params[k], rng=rngs[counter], **kwargs)
-      counter += 1
-    return result
 
-  def output_shape_fn(self, input_shape):
-    output_shapes = []
-    # If the argument layers are a sequence, apply each to calculate shape.
-    if not isinstance(self._layers, dict):
-      for layer in self._layers:
-        output_shapes.append(layer.output_shape(input_shape))
-      return tuple(output_shapes)
-    # If layers are a dictionary, apply to the input shape.
-    result = {}
-    for k in self._layers:
-      result[k] = self._layers[k].output_shape(input_shape)
-    return result
+  def default_input_is_int(self):
+    return self._layers[0].default_input_is_int()
 
   def new_parameters(self, input_shape, rng):
     rngs = backend.random.split(rng, self._nlayers)
-    # If the argument layers are a sequence, create parameters for each one.
     if not isinstance(self._layers, dict):
-      return [layer.initialize(input_shape, rng) for layer, rng
-              in zip(self._layers, rngs)]
-    # If the argument layers are a dictionary, create a dictionary too.
-    result, counter = {}, 0
-    for k in self._layers:
-      result[k] = self._layers[k].initialize(input_shape, rngs[counter])
-      counter += 1
-    return result
+      return [layer.initialize(input_shape, rng)
+              for layer, rng in zip(self._layers, rngs)]
+
+  def output_shape(self, input_shape, params):
+    output_shapes = []
+    if not isinstance(self._layers, dict):
+      for layer, param in zip(self._layers, params):
+        output_shapes.append(layer.output_shape(input_shape, param))
+      return tuple(output_shapes)
 
 
 def _nested_op(inputs, op):  # pylint: disable=invalid-name
@@ -375,44 +316,27 @@ def _binary_op(inputs, op):  # pylint: disable=invalid-name
   return res
 
 
-def _binary_op_shape(stack_shape):  # pylint: disable=invalid-name
-  """Helper: shape for the top-two operation above (shape-preserving op)."""
-  if len(stack_shape) == 2:
-    return stack_shape[0]
-  return tuple([stack_shape[0]] + list(stack_shape[2:]))
-
-
-@base.layer(output_shape=_binary_op_shape, stack_items_to_pass=0)
+@base.layer(stack_items_to_pass=0)
 def Add(x, **unused_kwargs):
   """Add first and second element on the stack."""
   # Here x is a list of tensors of the same shape, or nested structures.
   return _binary_op(x, op=sum)
 
 
-@base.layer(output_shape=_binary_op_shape, stack_items_to_pass=0)
+@base.layer(stack_items_to_pass=0)
 def Multiply(x, **unused_kwargs):
   """Multiply first and second element on the stack."""
   return _binary_op(x, op=lambda xs: six.moves.reduce(operator.mul, xs))
 
 
-def _nested_sum(inputs):  # pylint: disable=invalid-name
-  return _nested_op(inputs=inputs, op=sum)
-
-
-def _first_from_tuple_or_dict(tuple_or_dict):  # pylint: disable=invalid-name
-  """Helper: return the first element from a tuple or dict."""
-  for x in tuple_or_dict:
-    return x
-
-
-@base.layer(output_shape=_first_from_tuple_or_dict, stack_items_to_pass=0)
+@base.layer(stack_items_to_pass=0)
 def AddAll(x, **unused_kwargs):
   """Add branches elementwise."""
   # Here x is a list of tensors of the same shape, or nested structures.
-  return _nested_sum(x)
+  return _nested_op(x, op=sum)
 
 
-@base.layer(output_shape=_first_from_tuple_or_dict, stack_items_to_pass=0)
+@base.layer(stack_items_to_pass=0)
 def Gate(x, **unused_kwargs):
   """Implements a gating function on a (memory, gate, candidate) tuple.
 
@@ -432,17 +356,7 @@ def Gate(x, **unused_kwargs):
   return gate * state + (1.0 - gate) * candidate
 
 
-def _concatenate_shape(input_shape, axis=-1):  # pylint: disable=invalid-name
-  """Helper to determine the shape of Concatenate output."""
-  if isinstance(input_shape, dict):  # For named tuples, just use the values.
-    input_shape = list(input_shape.values())
-  ax = axis % len(input_shape[0])
-  concat_size = sum(shape[ax] for shape in input_shape)
-  out_shape = input_shape[0][:ax] + (concat_size,) + input_shape[0][ax+1:]
-  return out_shape
-
-
-@base.layer(output_shape=_concatenate_shape, stack_items_to_pass=0)
+@base.layer(stack_items_to_pass=0)
 def Concatenate(x, params, axis=-1, **kwargs):
   del params, kwargs
   if isinstance(x, dict):  # For dictionaries, just use the values.
@@ -476,6 +390,9 @@ def __init__(self, *layers, **kwlayers):
   def stack_items_to_pass(self):
     return self._nlayers
 
+  def default_input_is_int(self):
+    return any([layer.default_input_is_int() for layer in self._layers])
+
   def call(self, inputs, params=(), **kwargs):
     # Split the random number generators.
     rng = kwargs.pop('rng', None)
@@ -501,22 +418,6 @@ def call(self, inputs, params=(), **kwargs):
         result[k] = inputs[k]
     return result
 
-  def output_shape_fn(self, input_shape):
-    output_shapes = []
-    # If the argument layers are a sequence, apply each to calculate shape.
-    if not isinstance(self._layers, dict):
-      for i, layer in enumerate(self._layers):
-        output_shapes.append(layer.output_shape(input_shape[i]))
-      return tuple(output_shapes)
-    # If layers are a dictionary, apply to matching keys in the input shape.
-    result = {}
-    for k in input_shape:
-      if k in self._layers:
-        result[k] = self._layers[k].output_shape(input_shape[k])
-      else:
-        result[k] = input_shape[k]
-    return result
-
   def new_parameters(self, input_shape, rng):
     rngs = backend.random.split(rng, self._nlayers)
     # If the argument layers are a sequence, create parameters for each one.
@@ -536,7 +437,7 @@ def Residual(*layers, **kwargs):
   shortcut = kwargs.get('shortcut', _Top())  # pylint: disable=no-value-for-parameter
   return [
       Branch(shortcut, Serial(layers)),  # Use Serial here to flatten layers.
-      Flatten(),  # pylint: disable=no-value-for-parameter
+      FlattenList(),  # pylint: disable=no-value-for-parameter
       Add(),  # pylint: disable=no-value-for-parameter
   ]
 
@@ -574,9 +475,6 @@ def call(self, inputs, params=(), **kwargs):
       return result
     return tuple(result)
 
-  def output_shape_fn(self, input_shapes):
-    return tuple([self._layer.output_shape(shape) for shape in input_shapes])
-
   def new_parameters(self, input_shape, rng):
     first_shape = input_shape[0]
     if self._check_shapes:
@@ -630,11 +528,6 @@ def call(self, inp, params=(), **kwargs):
     out = self._layer(inp, params=params, **kwargs)
     return self._unmodify(out, batch_dims)
 
-  def output_shape_fn(self, input_shape):
-    modified_shape, batch_dims = self._modify_shape(input_shape)
-    out = self._layer.output_shape(modified_shape)
-    return self._unmodify_shape(out, batch_dims)
-
   def new_parameters(self, input_shape, rng):
     modified_shape, _ = self._modify_shape(input_shape)
     return self._layer.initialize(modified_shape, rng)
diff --git a/tensor2tensor/trax/layers/combinators_test.py b/tensor2tensor/trax/layers/combinators_test.py
index 592735978..bc5b8b593 100644
--- a/tensor2tensor/trax/layers/combinators_test.py
+++ b/tensor2tensor/trax/layers/combinators_test.py
@@ -33,13 +33,6 @@ def test_branch(self):
         combinators.Branch([], []), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
-  def test_branch_named(self):
-    input_shape = (2, 3)
-    expected_shape = {'a': (2, 3), 'b': (2, 3)}
-    output_shape = base.check_shape_agreement(
-        combinators.Branch(a=[], b=[]), input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
   def test_parallel(self):
     input_shape = ((2, 3), (2, 3))
     expected_shape = ((2, 3), (2, 3))
@@ -47,13 +40,6 @@ def test_parallel(self):
         combinators.Parallel([], []), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
-  def test_parallel_named(self):
-    input_shape = {'a': (2, 3), 'b': (2, 3)}
-    expected_shape = {'a': (2, 3), 'b': (2, 3)}
-    output_shape = base.check_shape_agreement(
-        combinators.Parallel(a=[]), input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
   def test_select(self):
     input_shape = ((2, 3), (3, 4))
     expected_shape = (3, 4)
@@ -61,13 +47,6 @@ def test_select(self):
         combinators.Select(1), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
-  def test_select_named(self):
-    input_shape = {'a': (2, 3), 'b': (3, 4)}
-    expected_shape = (3, 4)
-    output_shape = base.check_shape_agreement(
-        combinators.Select('b'), input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
   def test_rebatch(self):
     input_shape = (29, 5, 5, 20)
     result_shape = base.check_shape_agreement(
diff --git a/tensor2tensor/trax/layers/convolution.py b/tensor2tensor/trax/layers/convolution.py
index 83267be09..d721f2ec0 100644
--- a/tensor2tensor/trax/layers/convolution.py
+++ b/tensor2tensor/trax/layers/convolution.py
@@ -23,29 +23,10 @@
 
 from jax import lax
 
-import numpy as onp
 from tensor2tensor.trax.layers import base
 from tensor2tensor.trax.layers import initializers as init
 
 
-def PadtypeToPads(in_shape, window_shape, window_strides, padding):
-  """Convert padding string to list of pairs of pad values."""
-  padding = padding.upper()
-  if padding == 'SAME':
-    out_shape = onp.ceil(
-        onp.true_divide(in_shape, window_strides)).astype(int)
-    pad_sizes = [max((out_size - 1) * stride + window_shape - in_size, 0)
-                 for out_size, stride, window_shape, in_size
-                 in zip(out_shape, window_strides, window_shape, in_shape)]
-    return [(pad_size // 2, pad_size - pad_size // 2)
-            for pad_size in pad_sizes]
-  elif padding == 'VALID':
-    return [(0, 0)] * len(in_shape)
-  else:
-    msg = 'Unknown padding type: {}.'
-    raise TypeError(msg.format(padding))
-
-
 class Conv(base.Layer):
   """Layer constructor function for a general convolution layer."""
 
@@ -85,67 +66,6 @@ def _kernel_shape(self, input_shape):
             input_shape[self._lhs_spec.index('C')] if c == 'I' else
             next(kernel_size_iter) for c in self._rhs_spec]
 
-  def _conv_shape_tuple(self, lhs_shape, rhs_shape, strides, pads):
-    """Compute the shape of a conv given input shapes in canonical order."""
-    if isinstance(pads, str):
-      pads = PadtypeToPads(lhs_shape[2:], rhs_shape[2:], strides, pads)
-    if len(pads) != len(lhs_shape) - 2:
-      msg = 'Wrong number of explicit pads for conv: expected {}, got {}.'
-      raise TypeError(msg.format(len(lhs_shape) - 2, len(pads)))
-    lhs_padded = onp.add(lhs_shape[2:], onp.add(*zip(*pads)))
-    out_space = onp.floor_divide(
-        onp.subtract(lhs_padded, rhs_shape[2:]), strides) + 1
-    out_space = onp.maximum(0, out_space)
-    out_shape = (lhs_shape[0], rhs_shape[0]) + tuple(out_space)
-    return tuple(out_shape)
-
-  def _conv_general_permutations(self, dimension_numbers):
-    """Utility for convolution dimension permutations relative to Conv HLO."""
-    lhs_spec, rhs_spec, out_spec = dimension_numbers
-    lhs_char, rhs_char, out_char = ('N', 'C'), ('O', 'I'), ('N', 'C')
-    charpairs = (lhs_char, rhs_char, out_char)
-    for i, (a, b) in enumerate(charpairs):
-      if not (dimension_numbers[i].count(a) == 1 and
-              dimension_numbers[i].count(b) == 1):
-        msg = ('convolution dimension_numbers[{}] must contain the characters '
-               '"{}" and "{}" exatly once, got {}.')
-        raise TypeError(msg.format(i, a, b, dimension_numbers[i]))
-      if len(dimension_numbers[i]) != len(set(dimension_numbers[i])):
-        msg = ('convolution dimension_numbers[{}] cannot have duplicate '
-               'characters, got {}.')
-        raise TypeError(msg.format(i, dimension_numbers[i]))
-    if not (set(lhs_spec) - set(lhs_char) == set(rhs_spec) - set(rhs_char) ==
-            set(out_spec) - set(out_char)):
-      msg = ('convolution dimension_numbers elements must each have the same '
-             'set of spatial characters, got {}.')
-      raise TypeError(msg.format(dimension_numbers))
-
-    def GetPerm(spec, charpair):
-      spatial = (i for i, c in enumerate(spec) if c not in charpair)
-      if spec is not rhs_spec:
-        spatial = sorted(spatial, key=lambda i: rhs_spec.index(spec[i]))
-      return (spec.index(charpair[0]), spec.index(charpair[1])) + tuple(spatial)
-
-    lhs_perm, rhs_perm, out_perm = map(GetPerm, dimension_numbers, charpairs)
-    return lhs_perm, rhs_perm, out_perm
-
-  def _conv_general_shape_tuple(self, lhs_shape, rhs_shape, window_strides,
-                                padding, dimension_numbers):
-    """Generalized computation of conv shape."""
-    lhs_perm, rhs_perm, out_perm = self._conv_general_permutations(
-        dimension_numbers)
-    lhs_trans = onp.take(lhs_shape, lhs_perm)
-    rhs_trans = onp.take(rhs_shape, rhs_perm)
-    out_trans = self._conv_shape_tuple(
-        lhs_trans, rhs_trans, window_strides, padding)
-    return tuple(onp.take(out_trans, onp.argsort(out_perm)))
-
-  def output_shape_fn(self, input_shape):
-    kernel_shape = self._kernel_shape(input_shape)
-    return self._conv_general_shape_tuple(
-        input_shape, kernel_shape,
-        self._strides, self._padding, self._dimension_numbers)
-
   def new_parameters(self, input_shape, rng):
     kernel_shape = self._kernel_shape(input_shape)
     bias_shape = [self._filters if c == 'C' else 1 for c in self._out_spec]
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 3467d0b9b..c98694a41 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -19,8 +19,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import operator as op
-from six.moves import reduce
 from tensor2tensor.trax import backend
 from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers import base
@@ -97,9 +95,6 @@ def call(self, x, params, **kwargs):
     w, b = params
     return np.dot(x, w) + b
 
-  def output_shape_fn(self, input_shape):
-    return tuple(input_shape[:-1]) + (self._units,)
-
   def new_parameters(self, input_shape, rng):
     rng1, rng2 = backend.random.split(rng, 2)
     w = self._kernel_initializer((input_shape[-1], self._units), rng1)
@@ -120,32 +115,26 @@ def __init__(self, d_feature, vocab_size,
   def stack_items_to_pass(self):
     return 1
 
+  def default_input_is_int(self):
+    return True
+
   def call(self, x, params, **kwargs):
     del kwargs
     return np.take(params, x, axis=0)
 
-  def output_shape_fn(self, input_shape):
-    return tuple(input_shape) + (self._d_feature,)
-
   def new_parameters(self, input_shape, rng):
     return self._kernel_initializer(
         (self._vocab_size, self._d_feature), rng)
 
 
 # Flatten.
-def _flatten_output_shape(input_shape, num_axis_to_keep=1):  # pylint: disable=invalid-name
-  """Output shape of a flatten layer."""
-  if num_axis_to_keep >= len(input_shape):
-    raise ValueError(
-        "num_axis_to_keep[%d] should be less than input's rank[%d]" %
-        (num_axis_to_keep, len(input_shape)))
-  return tuple(input_shape[:num_axis_to_keep]) + (
-      reduce(op.mul, input_shape[num_axis_to_keep:], 1),)
-
-
-@base.layer(output_shape=_flatten_output_shape)
+@base.layer()
 def Flatten(x, params, num_axis_to_keep=1, **kwargs):
   del params, kwargs
+  if num_axis_to_keep >= len(x.shape):
+    raise ValueError(
+        "num_axis_to_keep[%d] should be less than input's rank[%d]" %
+        (num_axis_to_keep, len(x.shape)))
   return np.reshape(x, (x.shape[:num_axis_to_keep] + (-1,)))
 
 
@@ -185,14 +174,7 @@ def one_hot(x, size, dtype=np.float32):  # pylint: disable=invalid-name
 
 
 # Mean.
-def _mean_output_shape(input_shape, axis=-1, keepdims=False):  # pylint: disable=invalid-name
-  shape1 = list(input_shape)[:axis]  # Shape before axis.
-  shape2 = list(input_shape)[axis:][1:]  # Shape after axis.
-  mid_shape = [1] if keepdims else []
-  return tuple(shape1 + mid_shape + shape2)
-
-
-@base.layer(output_shape=_mean_output_shape)
+@base.layer()
 def Mean(x, params, axis=-1, keepdims=False, **kwargs):
   del params, kwargs
   return np.mean(x, axis=axis, keepdims=keepdims)
diff --git a/tensor2tensor/trax/layers/core_test.py b/tensor2tensor/trax/layers/core_test.py
index 765fe92cf..dd0775376 100644
--- a/tensor2tensor/trax/layers/core_test.py
+++ b/tensor2tensor/trax/layers/core_test.py
@@ -69,8 +69,8 @@ def test_dense_param_sharing(self):
     layer = core.Dense(32)
     model2 = combinators.Serial(layer, layer)
     rng = backend.random.get_prng(0)
-    params1 = model1.initialize((-1, 32), rng)
-    params2 = model2.initialize((-1, 32), rng)
+    params1 = model1.initialize((1, 32), rng)
+    params2 = model2.initialize((1, 32), rng)
     # The first parameters have 2 kernels of size (32, 32).
     self.assertEqual((32, 32), params1[0][0].shape)
     self.assertEqual((32, 32), params1[1][0].shape)
diff --git a/tensor2tensor/trax/layers/pooling.py b/tensor2tensor/trax/layers/pooling.py
index c04b1aeb6..3c07fd5fb 100644
--- a/tensor2tensor/trax/layers/pooling.py
+++ b/tensor2tensor/trax/layers/pooling.py
@@ -21,22 +21,8 @@
 
 from jax import lax
 
-import numpy as onp
 from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import convolution
-
-
-def PoolingOutputShape(input_shape, pool_size=(2, 2),
-                       strides=None, padding='VALID'):
-  """Helper: compute the output shape for the pooling layer."""
-  dims = (1,) + pool_size + (1,)  # NHWC
-  spatial_strides = strides or (1,) * len(pool_size)
-  strides = (1,) + spatial_strides + (1,)
-  pads = convolution.PadtypeToPads(input_shape, dims, strides, padding)
-  operand_padded = onp.add(input_shape, onp.add(*zip(*pads)))
-  t = onp.floor_divide(onp.subtract(operand_padded, dims), strides) + 1
-  return tuple(t)
 
 
 def PoolingGeneral(inputs, reducer, init_val, rescaler=None,
@@ -50,14 +36,14 @@ def PoolingGeneral(inputs, reducer, init_val, rescaler=None,
   return rescale(out, inputs) if rescale else out
 
 
-@base.layer(output_shape=PoolingOutputShape)
+@base.layer()
 def MaxPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
   del params, kw
   return PoolingGeneral(x, lax.max, -np.inf, pool_size=pool_size,
                         strides=strides, padding=padding)
 
 
-@base.layer(output_shape=PoolingOutputShape)
+@base.layer()
 def SumPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
   del params, kw
   return PoolingGeneral(x, lax.add, 0., pool_size=pool_size,
@@ -73,7 +59,7 @@ def Rescale(outputs, inputs):
   return Rescale
 
 
-@base.layer(output_shape=PoolingOutputShape)
+@base.layer()
 def AvgPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
   del params, kw
   return PoolingGeneral(x, lax.add, 0., _normalize_by_window_size,
diff --git a/tensor2tensor/trax/models/atari_cnn_test.py b/tensor2tensor/trax/models/atari_cnn_test.py
index 6d1ff70f3..0b1cb21d1 100644
--- a/tensor2tensor/trax/models/atari_cnn_test.py
+++ b/tensor2tensor/trax/models/atari_cnn_test.py
@@ -37,7 +37,7 @@ def test_computes(self):
         hidden_sizes=hidden_size, output_size=output_size)
     B, T, OBS = 2, 2, (28, 28, 3)  # pylint: disable=invalid-name
     rng_key, key = jax_random.split(rng_key)
-    params = model.initialize((-1, -1) + OBS, key)
+    params = model.initialize((1, 1) + OBS, key)
     x = onp.arange(B * (T + 1) * functools.reduce(op.mul, OBS)).reshape(
         B, T + 1, *OBS)
     rng_key, key = jax_random.split(rng_key)
diff --git a/tensor2tensor/trax/models/neural_gpu.py b/tensor2tensor/trax/models/neural_gpu.py
index 0d357a72c..dc19d52cb 100644
--- a/tensor2tensor/trax/models/neural_gpu.py
+++ b/tensor2tensor/trax/models/neural_gpu.py
@@ -28,7 +28,7 @@ def SaturationCost(x, limit=0.9):
   return np.minimum(0, np.abs(x) - limit)
 
 
-@tl.layer(output_shape=lambda input_shape_list: input_shape_list)
+@tl.layer()
 def DiagonalGate(x, params, **kwargs):
   """Split channels in 3 parts. Shifts 1st and 3rd sections to left/right."""
   del params
diff --git a/tensor2tensor/trax/models/research/chunked_transformer.py b/tensor2tensor/trax/models/research/chunked_transformer.py
index 4223e98cd..d819ae40a 100644
--- a/tensor2tensor/trax/models/research/chunked_transformer.py
+++ b/tensor2tensor/trax/models/research/chunked_transformer.py
@@ -60,32 +60,7 @@ def ChunkedPositionalEncoding(x, params, **unused_kwargs):
 
 
 # Chunked attention.
-def _chunked_selector_output_shape(  # pylint: disable=invalid-name
-    input_shapes, selector=None, **unused_kwargs):
-  """Helper: calculate output shape for chunked key selector (see below)."""
-  # Read the main function below first, the shape logic just follows the ops.
-  selector = selector or (lambda x: [] if x < 1 else [x-1])
-  triples, _ = zip(*input_shapes)
-  (query_shapes, key_shapes, value_shapes) = zip(*triples)
-  result = []
-  for i in range(len(input_shapes)):
-    selected = selector(i)
-    cur_key_shape, cur_value_shape = key_shapes[i], value_shapes[i]
-    # Since keys and values are [batch, length, depth] we concatenate on axis=1.
-    new_key_len = sum([key_shapes[j][1] for j in selected]) + cur_key_shape[1]
-    new_key_shape = (cur_key_shape[0], new_key_len, cur_key_shape[2])
-    new_value_len = sum(
-        [value_shapes[j][1] for j in selected]) + cur_value_shape[1]
-    new_value_shape = (cur_value_shape[0], new_value_len, cur_value_shape[2])
-    # Masks are (1, query-len, key-len).
-    new_mask_shape = (1, query_shapes[i][1], new_key_len)
-    new_shape = (query_shapes[i], new_key_shape, new_value_shape,
-                 new_mask_shape)
-    result.append(new_shape)
-  return tuple(result)
-
-
-@tl.layer(output_shape=_chunked_selector_output_shape, stack_items_to_pass=0)
+@tl.layer(stack_items_to_pass=0)
 def ChunkedAttentionSelector(x, params, selector=None, **kwargs):
   """Select which chunks to attend to in chunked attention.
 
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index ae195b155..95ed69ed0 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -95,7 +95,7 @@ def TransformerEncoder(vocab_size,
       tl.Dropout(rate=dropout, mode=mode),
       tl.PositionalEncoding(max_len=max_len),
   ]
-  return [
+  return tl.Model([
       tl.Branch(positional_embedder, tl.PaddingMask()),  # Create mask.
       [EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode)
        for _ in range(n_layers)],
@@ -104,7 +104,7 @@ def TransformerEncoder(vocab_size,
       tl.Mean(axis=1),  # Average on length.
       tl.Dense(n_classes),
       tl.LogSoftmax(),
-  ]
+  ])
 
 
 def DecoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 2af68ddd7..fc7be2670 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -810,9 +810,9 @@ def training_loop(
 
   jax_rng_key = trax.get_random_number_generator_and_set_seed(random_seed)
 
-  # Batch Observations Shape = [-1, -1] + OBS, because we will eventually call
+  # Batch Observations Shape = [1, 1] + OBS, because we will eventually call
   # policy and value networks on shape [B, T] +_OBS
-  batch_observations_shape = (-1, -1) + env.observation_space.shape
+  batch_observations_shape = (1, 1) + env.observation_space.shape
 
   assert isinstance(env.action_space, gym.spaces.Discrete)
   n_actions = env.action_space.n
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 2972a73fb..95533a617 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -35,7 +35,7 @@ def setUp(self):
 
   def test_policy_and_value_net(self):
     observation_shape = (3, 4, 5)
-    batch_observation_shape = (-1, -1) + observation_shape
+    batch_observation_shape = (1, 1) + observation_shape
     n_actions = 2
     pnv_params, pnv_apply = ppo.policy_and_value_net(
         self.rng_key, batch_observation_shape, n_actions,
@@ -380,7 +380,7 @@ def test_combined_loss(self):
     self.rng_key, key1, key2 = jax_random.split(self.rng_key, num=3)
 
     B, T, A, OBS = 2, 10, 2, (28, 28, 3)  # pylint: disable=invalid-name
-    batch_observation_shape = (-1, -1) + OBS
+    batch_observation_shape = (1, 1) + OBS
 
     old_params, _ = ppo.policy_and_value_net(
         key1, batch_observation_shape, A,
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index eafd3dbbb..604df1293 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -451,12 +451,15 @@ def train(output_dir,
   rng, init_rng = jax_random.split(rng)
   rngs = jax_random.split(rng, n_devices)
   first_shape = inputs.input_shape[0]
-  # If the inputs are a tuple/list, add [-1] (batch) to each element.
+  # If the inputs are a tuple/list, add [None] (batch) to each element.
   if isinstance(first_shape, (list, tuple)):
     model_input_shape = tuple(
-        [tuple([-1] + list(shape)) for shape in inputs.input_shape])
-  else:  # Otherwise just add [-1] to the input shape.
-    model_input_shape = tuple([-1] + list(inputs.input_shape))
+        tuple([None] + list(shape)) for shape in inputs.input_shape)
+  else:  # Otherwise just add [None] to the input shape.
+    model_input_shape = tuple([None] + list(inputs.input_shape))
+  # Change all None to 1 in input shape.
+  model_input_shape = layers.nested_map(
+      model_input_shape, lambda x: x if x else 1)
   if state.params:
     params = state.params[0]
     opt_state = state.params

From 5acf4a44cc2cbe91cd788734075376af0f8dd3f4 Mon Sep 17 00:00:00 2001
From: Chris Gorgolewski <chrisgo@google.com>
Date: Mon, 3 Jun 2019 11:08:18 -0700
Subject: [PATCH 2088/2720] Adding metadata to the wikisum dataset description.

PiperOrigin-RevId: 251269380
---
 .../data_generators/wikisum/README.md         | 99 +++++++++++++++++++
 1 file changed, 99 insertions(+)

diff --git a/tensor2tensor/data_generators/wikisum/README.md b/tensor2tensor/data_generators/wikisum/README.md
index f71d550da..0692a19eb 100644
--- a/tensor2tensor/data_generators/wikisum/README.md
+++ b/tensor2tensor/data_generators/wikisum/README.md
@@ -211,3 +211,102 @@ t2t-trainer \
   --data_dir=$DATA_DIR \
   --output_dir=$TRAIN_DIR
 ```
+
+
+## Dataset Metadata
+The following table is necessary for this dataset to be indexed by search
+engines such as <a href="https://g.co/datasetsearch">Google Dataset Search</a>.
+<div itemscope itemtype="http://schema.org/Dataset">
+<table>
+  <tr>
+    <th>property</th>
+    <th>value</th>
+  </tr>
+  <tr>
+    <td>name</td>
+    <td><code itemprop="name">wikisum</code></td>
+  </tr>
+  <tr>
+    <td>alternateName</td>
+    <td><code itemprop="alternateName">WikisumCommonCrawl</code></td>
+  </tr>
+  <tr>
+    <td>alternateName</td>
+    <td><code itemprop="alternateName">WikisumWeb</code></td>
+  </tr>
+  <tr>
+    <td>alternateName</td>
+    <td><code itemprop="alternateName">wkisum_commoncrawl</code></td>
+  </tr>
+  <tr>
+    <td>alternateName</td>
+    <td><code itemprop="alternateName">wikisum_web</code></td>
+  </tr>
+  <tr>
+    <td>url</td>
+    <td><code itemprop="url">https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/wikisum</code></td>
+  </tr>
+  <tr>
+    <td>sameAs</td>
+    <td><code itemprop="sameAs">https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/wikisum</code></td>
+  </tr>
+  <tr>
+    <td>description</td>
+    <td><code itemprop="description">The dataset from the
+paper [Generating Wikipedia by Summarizing Long
+Sequences](https://arxiv.org/abs/1801.10198). The task is to generate a
+Wikipedia article based on the contents of the cited references in that article
+and the top 10 Google search results for the article's title.\n
+\n
+There are 2 sources for the reference URLs used:
+\n
+1. [CommonCrawl](http://commoncrawl.org/), an open-source crawl of the web. The
+   advantage of using CommonCrawl is that the dataset is perfectly reproducible.
+   However, there is limited coverage of the reference URLs.
+1. Live web fetches. Coverage is considerably increased, but the content is
+   subject to change.\n
+\n
+The dataset includes:\n
+\n
+**URLs:** The dataset contains ~90M URLs total (~2.3M Wikipedia articles, each
+with ~40 reference URLs). The URLs in the dataset are available in sharded JSON
+files.\n
+\n
+**Wikipedia Articles:** We have processed the Wikipedia articles slightly to
+extract the title, section breaks, and section headings. The processed Wikipedia
+content is available in sharded `TFRecord` files containing serialized
+`tensorflow.Example` protocol buffers.\n
+\n
+**CommonCrawl References Index:** To enable efficiently extracting the reference
+URLs from CommonCrawl, we provide a JSON file per CommonCrawl file which maps a
+reference URL contained in that CommonCrawl file to a list of shard ids.
+These shards are the ones that contain one or more Wikipedia articles that cite
+this reference.</code></td>
+  </tr>
+  <tr>
+    <td>citation</td>
+    <td><code itemprop="citation">https://identifiers.org/arxiv:1801.10198</code></td>
+  </tr>
+  <tr>
+    <td>provider</td>
+    <td>
+      <div itemscope itemtype="http://schema.org/Organization" itemprop="provider">
+        <table>
+          <tr>
+            <th>property</th>
+            <th>value</th>
+          </tr>
+          <tr>
+            <td>name</td>
+            <td><code itemprop="name">Google</code></td>
+          </tr>
+          <tr>
+            <td>sameAs</td>
+            <td><code itemprop="sameAs">https://en.wikipedia.org/wiki/Google</code></td>
+          </tr>
+        </table>
+      </div>
+    </td>
+  </tr>
+</table>
+</div>

From 5fc18b815e1b573f207d278644e4a4937441984f Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 3 Jun 2019 13:20:45 -0700
Subject: [PATCH 2089/2720] Fix scale transform.

PiperOrigin-RevId: 251295286
---
 tensor2tensor/layers/reversible_layers.py     | 46 ++++++++-------
 .../layers/reversible_layers_test.py          | 57 ++++++++++++-------
 2 files changed, 61 insertions(+), 42 deletions(-)

diff --git a/tensor2tensor/layers/reversible_layers.py b/tensor2tensor/layers/reversible_layers.py
index 997720951..421026256 100644
--- a/tensor2tensor/layers/reversible_layers.py
+++ b/tensor2tensor/layers/reversible_layers.py
@@ -156,18 +156,20 @@ def _initial_call(self, new_inputs, length, **kwargs):
     net = self.layer(padded_inputs, **kwargs)
     if net.shape[-1] == 2 * self.vocab_size:
       loc, scale = tf.split(net, 2, axis=-1)
+      loc = loc[..., 0:1, :]
+      loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
       scale = scale[..., 0:1, :]
       scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
       inverse_scale = multiplicative_inverse(scale, self.vocab_size)
-      scaled_inputs = one_hot_multiply(inputs, inverse_scale)
+      shifted_inputs = one_hot_minus(inputs, loc)
+      outputs = one_hot_multiply(shifted_inputs, inverse_scale)
     elif net.shape[-1] == self.vocab_size:
       loc = net
-      scaled_inputs = inputs
+      loc = loc[..., 0:1, :]
+      loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
+      outputs = one_hot_minus(inputs, loc)
     else:
       raise ValueError('Output of layer does not have compatible dimensions.')
-    loc = loc[..., 0, :][..., tf.newaxis, :]
-    loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-    outputs = one_hot_minus(scaled_inputs, loc)
     return outputs
 
   def _per_timestep_call(self,
@@ -197,20 +199,21 @@ def _per_timestep_call(self,
     net = self.layer(padded_inputs, **kwargs)
     if net.shape[-1] == 2 * self.vocab_size:
       loc, scale = tf.split(net, 2, axis=-1)
+      loc = loc[..., :(timestep+1), :]
+      loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
       scale = scale[..., :(timestep+1), :]
       scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
       inverse_scale = multiplicative_inverse(scale, self.vocab_size)
-      scaled_inputs = one_hot_multiply(inputs, inverse_scale)
+      shifted_inputs = one_hot_minus(inputs, loc)
+      new_outputs = one_hot_multiply(shifted_inputs, inverse_scale)
     elif net.shape[-1] == self.vocab_size:
       loc = net
-      scaled_inputs = inputs
+      loc = loc[..., :(timestep+1), :]
+      loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
+      new_outputs = one_hot_minus(inputs, loc)
     else:
       raise ValueError('Output of layer does not have compatible dimensions.')
-    loc = loc[..., :(timestep+1), :]
-    loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-    new_outputs = one_hot_minus(scaled_inputs, loc)
-    outputs = tf.concat([current_outputs,
-                         new_outputs[..., -1, :][..., tf.newaxis, :]], axis=-2)
+    outputs = tf.concat([current_outputs, new_outputs[..., -1:, :]], axis=-2)
     if not tf.executing_eagerly():
       outputs.set_shape([None] * batch_ndims + [timestep+1, self.vocab_size])
     return outputs
@@ -222,7 +225,7 @@ def reverse(self, inputs, **kwargs):
 
     net = self.layer(inputs, **kwargs)
     if net.shape[-1] == 2 * self.vocab_size:
-      loc, scale = tf.split(net, 2, axis=-2)
+      loc, scale = tf.split(net, 2, axis=-1)
       scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
       scaled_inputs = one_hot_multiply(inputs, scale)
     elif net.shape[-1] == self.vocab_size:
@@ -326,16 +329,18 @@ def call(self, inputs, **kwargs):
     net = self.layer(masked_inputs, **kwargs)
     if net.shape[-1] == 2 * self.vocab_size:
       loc, scale = tf.split(net, 2, axis=-1)
+      loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
       scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
       inverse_scale = multiplicative_inverse(scale, self.vocab_size)
-      scaled_inputs = one_hot_multiply(inputs, inverse_scale)
+      shifted_inputs = one_hot_minus(inputs, loc)
+      masked_outputs = (1. - mask) * one_hot_multiply(shifted_inputs,
+                                                      inverse_scale)
     elif net.shape[-1] == self.vocab_size:
       loc = net
-      scaled_inputs = inputs
+      loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
+      masked_outputs = (1. - mask) * one_hot_minus(inputs, loc)
     else:
       raise ValueError('Output of layer does not have compatible dimensions.')
-    loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-    masked_outputs = (1. - mask) * one_hot_minus(scaled_inputs, loc)
     outputs = masked_inputs + masked_outputs
     return outputs
 
@@ -351,7 +356,7 @@ def reverse(self, inputs, **kwargs):
     masked_inputs = mask * inputs
     net = self.layer(masked_inputs, **kwargs)
     if net.shape[-1] == 2 * self.vocab_size:
-      loc, scale = tf.split(net, 2, axis=-2)
+      loc, scale = tf.split(net, 2, axis=-1)
       scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
       scaled_inputs = one_hot_multiply(inputs, scale)
     elif net.shape[-1] == self.vocab_size:
@@ -442,7 +447,7 @@ def _initial_call(self, new_inputs, length, **kwargs):
         inputs, [[0, 0]] * batch_ndims + [[0, length - 1], [0, 0]])
     temperature = 1.
     logits = self.layer(padded_inputs / temperature, **kwargs)
-    logits = logits[..., 0, :][..., tf.newaxis, :]
+    logits = logits[..., 0:1, :]
     logits = tf.reshape(
         logits,
         logits.shape[:-1].concatenate([self.vocab_size, self.vocab_size]))
@@ -495,8 +500,7 @@ def _per_timestep_call(self,
     new_outputs = tf.matmul(inputs[..., tf.newaxis, :],
                             hard,
                             transpose_b=True)[..., 0, :]
-    outputs = tf.concat([current_outputs,
-                         new_outputs[..., -1, :][..., tf.newaxis, :]], axis=-2)
+    outputs = tf.concat([current_outputs, new_outputs[..., -1:, :]], axis=-2)
     if not tf.executing_eagerly():
       outputs.set_shape([None] * batch_ndims + [timestep+1, self.vocab_size])
     return outputs
diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
index 83de8a798..fb1ac4644 100644
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -54,8 +54,7 @@ def _log_prob(self, x):
 class ReversibleLayersTest(parameterized.TestCase, tf.test.TestCase):
 
   @parameterized.parameters(
-      # TODO(trandustin): Enable test.
-      # (False,),
+      (False,),
       (True,),
   )
   @test_utils.run_in_graph_and_eager_modes()
@@ -65,12 +64,16 @@ def testDiscreteAutoregressiveFlowCall(self, loc_only):
     length = 5
     if loc_only:
       units = vocab_size
+      network = reversible.MADE(units, [])
     else:
       units = 2 * vocab_size
+      mask = tf.reshape([0] * vocab_size + [-1e10] + [0] * (vocab_size - 1),
+                        [1, 1, 2 * vocab_size])
+      network_ = reversible.MADE(units, [])
+      network = lambda inputs: mask + network_(inputs)
     inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
     inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
-    layer = reversible.DiscreteAutoregressiveFlow(
-        reversible.MADE(units, []), 1.)
+    layer = reversible.DiscreteAutoregressiveFlow(network, 1.)
     outputs = layer(inputs)
     self.evaluate(tf.global_variables_initializer())
     outputs_val = self.evaluate(outputs)
@@ -79,8 +82,7 @@ def testDiscreteAutoregressiveFlowCall(self, loc_only):
     self.assertAllLessEqual(outputs_val, vocab_size - 1)
 
   @parameterized.parameters(
-      # TODO(trandustin): Enable test.
-      # (False,),
+      (False,),
       (True,),
   )
   @test_utils.run_in_graph_and_eager_modes()
@@ -90,10 +92,14 @@ def testDiscreteAutoregressiveFlowSample(self, loc_only):
     vocab_size = 2
     if loc_only:
       units = vocab_size
+      network = reversible.MADE(units, [])
     else:
       units = 2 * vocab_size
-    layer = reversible.DiscreteAutoregressiveFlow(
-        reversible.MADE(units, []), 1.)
+      mask = tf.reshape([0] * vocab_size + [-1e10] + [0] * (vocab_size - 1),
+                        [1, 1, 2 * vocab_size])
+      network_ = reversible.MADE(units, [])
+      network = lambda inputs: mask + network_(inputs)
+    layer = reversible.DiscreteAutoregressiveFlow(network, 1.)
     logits = tf.tile(tf.random_normal([length, vocab_size])[tf.newaxis],
                      [batch_size, 1, 1])
     base = tfp.edward2.OneHotCategorical(logits=logits, dtype=tf.float32)
@@ -106,8 +112,7 @@ def testDiscreteAutoregressiveFlowSample(self, loc_only):
     self.assertAllLessEqual(res, vocab_size - 1)
 
   @parameterized.parameters(
-      # TODO(trandustin): Enable test.
-      # (False,),
+      (False,),
       (True,),
   )
   @test_utils.run_in_graph_and_eager_modes()
@@ -117,12 +122,16 @@ def testDiscreteAutoregressiveFlowInverse(self, loc_only):
     length = 5
     if loc_only:
       units = vocab_size
+      network = reversible.MADE(units, [])
     else:
       units = 2 * vocab_size
+      mask = tf.reshape([0] * vocab_size + [-1e10] + [0] * (vocab_size - 1),
+                        [1, 1, 2 * vocab_size])
+      network_ = reversible.MADE(units, [])
+      network = lambda inputs: mask + network_(inputs)
     inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
     inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
-    layer = reversible.DiscreteAutoregressiveFlow(
-        reversible.MADE(units, []), 1.)
+    layer = reversible.DiscreteAutoregressiveFlow(network, 1.)
     rev_fwd_inputs = layer.reverse(layer(inputs))
     fwd_rev_inputs = layer(layer.reverse(inputs))
     self.evaluate(tf.global_variables_initializer())
@@ -132,8 +141,7 @@ def testDiscreteAutoregressiveFlowInverse(self, loc_only):
     self.assertAllClose(inputs_val, fwd_rev_inputs_val)
 
   @parameterized.parameters(
-      # TODO(trandustin): Enable test.
-      # (False,),
+      (False,),
       (True,),
   )
   @test_utils.run_in_graph_and_eager_modes()
@@ -143,14 +151,18 @@ def testDiscreteAutoregressiveFlowRandomVariable(self, loc_only):
     vocab_size = 5
     if loc_only:
       units = vocab_size
+      network = reversible.MADE(units, [])
     else:
       units = 2 * vocab_size
+      mask = tf.reshape([0] * vocab_size + [-1e10] + [0] * (vocab_size - 1),
+                        [1, 1, 2 * vocab_size])
+      network_ = reversible.MADE(units, [])
+      network = lambda inputs: mask + network_(inputs)
     base = tfp.edward2.OneHotCategorical(logits=tf.random_normal([batch_size,
                                                                   length,
                                                                   vocab_size]),
                                          dtype=tf.float32)
-    flow = reversible.DiscreteAutoregressiveFlow(
-        reversible.MADE(units, [16, 16]), 1.)
+    flow = reversible.DiscreteAutoregressiveFlow(network, 1.)
     flow_rv = flow(base)
     self.assertEqual(flow_rv.dtype, tf.float32)
 
@@ -174,8 +186,7 @@ def testDiscreteAutoregressiveFlowRandomVariable(self, loc_only):
     self.assertAllClose(res1, res2)
 
   @parameterized.parameters(
-      # TODO(trandustin): Enable test.
-      # (False,),
+      (False,),
       (True,),
   )
   @test_utils.run_in_graph_mode_only()
@@ -185,17 +196,21 @@ def testDiscreteAutoregressiveFlowReverseGradients(self, loc_only):
     vocab_size = 2
     if loc_only:
       units = vocab_size
+      network = reversible.MADE(units, [16, 16])
     else:
       units = 2 * vocab_size
+      mask = tf.reshape([0] * vocab_size + [-1e10] + [0] * (vocab_size - 1),
+                        [1, 1, 2 * vocab_size])
+      network_ = reversible.MADE(units, [16, 16])
+      network = lambda inputs: mask + network_(inputs)
     base = tfp.edward2.OneHotCategorical(
         logits=tf.random_normal([batch_size, length, vocab_size]))
-    flow = reversible.DiscreteAutoregressiveFlow(
-        reversible.MADE(units, [16, 16]), 1.)
+    flow = reversible.DiscreteAutoregressiveFlow(network, 1.)
     flow_rv = flow(base)
     features = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
     features = tf.one_hot(features, depth=vocab_size, dtype=tf.float32)
     loss = -tf.reduce_sum(flow_rv.distribution.log_prob(features))
-    grads = tf.gradients(loss, flow.layer.weights)
+    grads = tf.gradients(loss, tf.trainable_variables())
     self.evaluate(tf.global_variables_initializer())
     _ = self.evaluate(grads)
     for grad in grads:

From 4bbac6c9598d6214056261c6f0753726b9c2db19 Mon Sep 17 00:00:00 2001
From: Zi Yang <ziy@google.com>
Date: Mon, 3 Jun 2019 13:45:27 -0700
Subject: [PATCH 2090/2720] [T2T] Added flags `teacher_dir` and `student_dir`
 and `skip_teacher_training`.

Flags `teacher_dir` and `student_dir` are used to explicitly specify teacher/student model paths. Flag `skip_teacher_training` is used if a teacher model has been trained.

PiperOrigin-RevId: 251300055
---
 tensor2tensor/bin/t2t_distill.py     | 52 ++++++++++++++++++++++------
 tensor2tensor/models/distillation.py |  8 +++--
 2 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/bin/t2t_distill.py b/tensor2tensor/bin/t2t_distill.py
index 2c05c9409..e58b3c218 100644
--- a/tensor2tensor/bin/t2t_distill.py
+++ b/tensor2tensor/bin/t2t_distill.py
@@ -17,6 +17,13 @@
 
 This script is intended to be used with --model=distillation. See the model for
 example hyperparameters and usage.
+
+If only output_dir is specified, then teacher_dir is `output_dir/teacher`, and
+the student_dir is `output_dir/student`. Logs are written inside `output_dir`.
+If teacher_dir is also specified explicitly, the student_dir is still
+`output_dir/student` and the logs are written into `output_dir`. If student_dir
+is further specified, the logs are written into student_dir unless output_dir is
+explicitly specified, which only contains the logs in this case.
 """
 from __future__ import absolute_import
 from __future__ import division
@@ -36,6 +43,18 @@
 flags = tf.flags
 FLAGS = flags.FLAGS
 
+flags.DEFINE_bool(
+    "skip_teacher_training", False,
+    "By default, we train teacher model. If set to True, skip the training.")
+flags.DEFINE_string(
+    "teacher_dir", None,
+    "Directory to teacher network. If not specified, `output_dir/teacher` is "
+    "used instead.")
+flags.DEFINE_string(
+    "student_dir", None,
+    "Directory to student network. If not specified, `output_dir/student` is "
+    "used instead.")
+
 
 def main(argv):
   tf.logging.set_verbosity(tf.logging.INFO)
@@ -58,24 +77,35 @@ def main(argv):
 
   root_output_dir = FLAGS.output_dir
 
+  if FLAGS.teacher_dir:
+    teacher_dir = FLAGS.teacher_dir
+  else:
+    teacher_dir = os.path.join(root_output_dir, "teacher")
+
   # Train Teacher ============
-  hparams = t2t_trainer.create_hparams()
-  hparams.distill_phase = "train"
-  teacher_dir = os.path.join(root_output_dir, "teacher")
-  FLAGS.output_dir = teacher_dir
+  if FLAGS.skip_teacher_training:
+    tf.logging.info("training teacher skipped")
+  else:
+    hparams = t2t_trainer.create_hparams()
+    hparams.distill_phase = "train"
+    FLAGS.output_dir = teacher_dir
+
+    exp_fn = t2t_trainer.create_experiment_fn()
+    run_config = t2t_trainer.create_run_config(hparams)
+    exp = exp_fn(run_config, hparams)
+    if t2t_trainer.is_chief():
+      t2t_trainer.save_metadata(hparams)
+    t2t_trainer.execute_schedule(exp)
 
-  exp_fn = t2t_trainer.create_experiment_fn()
-  run_config = t2t_trainer.create_run_config(hparams)
-  exp = exp_fn(run_config, hparams)
-  if t2t_trainer.is_chief():
-    t2t_trainer.save_metadata(hparams)
-  t2t_trainer.execute_schedule(exp)
   # ==========================
   # Train Student ============
   hparams = t2t_trainer.create_hparams()
   hparams.add_hparam("teacher_dir", teacher_dir)
   hparams.distill_phase = "distill"
-  student_dir = os.path.join(root_output_dir, "student")
+  if FLAGS.student_dir:
+    student_dir = FLAGS.student_dir
+  else:
+    student_dir = os.path.join(root_output_dir, "student")
   FLAGS.output_dir = student_dir
 
   exp_fn = t2t_trainer.create_experiment_fn()
diff --git a/tensor2tensor/models/distillation.py b/tensor2tensor/models/distillation.py
index 6600f48d6..0a7f69908 100644
--- a/tensor2tensor/models/distillation.py
+++ b/tensor2tensor/models/distillation.py
@@ -47,7 +47,8 @@ def __init__(self,
                mode=tf.estimator.ModeKeys.TRAIN,
                problem_hparams=None,
                data_parallelism=None,
-               decode_hparams=None):
+               decode_hparams=None,
+               **kwargs):
     assert hparams.distill_phase in ["train", "distill"]
 
     if hparams.distill_phase == "train" and hparams.teacher_learning_rate:
@@ -63,8 +64,9 @@ def __init__(self,
     self.student_model = registry.model(
         hparams.student_model)(self.student_hparams, mode, problem_hparams,
                                data_parallelism, decode_hparams)
-    super(Distillation, self).__init__(hparams, mode, problem_hparams,
-                                       data_parallelism, decode_hparams)
+    super(Distillation,
+          self).__init__(hparams, mode, problem_hparams, data_parallelism,
+                         decode_hparams, **kwargs)
 
   def body(self, features):
     hp = self.hparams

From 560c008f7d87502174765fac5ae3d822bbf6b243 Mon Sep 17 00:00:00 2001
From: Taylor Robie <taylorrobie@google.com>
Date: Mon, 3 Jun 2019 14:20:51 -0700
Subject: [PATCH 2091/2720] Implement sequence packing as a tf.data.Dataset
 transformation.

PiperOrigin-RevId: 251307162
---
 .../data_generators/generator_utils.py        | 499 ++++++++++++++----
 .../data_generators/generator_utils_test.py   |  89 ++++
 2 files changed, 472 insertions(+), 116 deletions(-)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index aae3e5572..33f3d3d2f 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -19,8 +19,10 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
 import gzip
 import math
+import multiprocessing
 import os
 import random
 import stat
@@ -695,10 +697,6 @@ def pack_dataset(dataset, length, keys=None, use_custom_ops=False):
   Sequences in the incoming examples are truncated to length "length", and the
   sequences in the output examples all have fixed (padded) length "length".
 
-  TODO(noam): This code is slow - the use_custom_ops option is faster, but
-  requiers a custom-built binary.  Resolve this so that it is easy to get
-  good perfomrance.
-
   Args:
     dataset: a tf.data.Dataset
     length: an integer
@@ -711,6 +709,7 @@ def pack_dataset(dataset, length, keys=None, use_custom_ops=False):
   shapes = dataset.output_shapes
   if keys is None:
     keys = shapes.keys()
+
   for k in keys:
     if k not in shapes:
       raise ValueError("Key %s not found in dataset.  Available keys are %s"
@@ -718,19 +717,11 @@ def pack_dataset(dataset, length, keys=None, use_custom_ops=False):
     if not shapes[k].is_compatible_with(tf.TensorShape([None])):
       raise ValueError("Tensors to be packed must be one-dimensional.")
 
-  # trim to length
-  dataset = dataset.map(lambda x: {k: x[k][:length] for k in keys})
-  # Setting batch_size=length ensures that the concatenated sequences (if they
-  # have length >=1) are sufficient to fill at least one packed example.
-  batch_size = length
-  dataset = dataset.padded_batch(
-      batch_size, padded_shapes={k: [-1] for k in keys})
-  if use_custom_ops and len(keys) == 2:
-    # custom op only handles 2 keys.
-    # TODO(noam): support other numbers of keys.
+  if use_custom_ops:
     return _pack_with_custom_ops(dataset, keys, length)
   else:
-    return _pack_with_tf_ops(dataset, keys, length)
+    packer = SequenceDatasetPacker(length, spacing=0, queue_size=10)
+    return packer(dataset, cycle_length=10, keys=keys)
 
 
 def _pack_with_custom_ops(dataset, keys, length):
@@ -750,7 +741,16 @@ def _pack_with_custom_ops(dataset, keys, length):
     a dataset.
   """
   from tensor2tensor.data_generators.ops import pack_sequences_ops  # pylint: disable=g-import-not-at-top
-  # faster and better packing but requires custom-built binary.
+
+  # trim to length
+  dataset = dataset.map(lambda x: {k: x[k][:length] for k in keys})
+  # Setting batch_size=length ensures that the concatenated sequences (if they
+  # have length >=1) are sufficient to fill at least one packed example.
+  batch_size = length
+  dataset = dataset.padded_batch(
+      batch_size, padded_shapes={k: [-1] for k in keys})
+
+  # better packing (may be faster) but requires custom-built binary.
   k1, k2 = keys
   def map_fn_custom(x):
     """Map-function."""
@@ -770,114 +770,381 @@ def map_fn_custom(x):
   return dataset
 
 
-def _pack_with_tf_ops(dataset, keys, length):
-  """Helper-function for packing a dataset which has already been batched.
+INDEX_DTYPE = tf.int32
 
-  See pack_dataset()
 
-  Uses tf.while_loop.  Slow.
+class SequenceDatasetPacker(object):
+  """Helper class for packing a dataset of sequences in an online fashon.
 
-  Args:
-    dataset: a dataset containing padded batches of examples.
-    keys: a list of strings
-    length: an integer
+  The input sequence is expected to be a tuple of 1D Tensors which will be
+  converted to a dataset which produces a dict of packed examples, example
+  positions, and segment ids.
 
-  Returns:
-    a dataset.
+  If `window_size` or `cycle_length` is specified multiple packing operations
+  will be performed in parallel to increase throughput. A value of None will
+  select default parallelism parameters. If this dataset will be run on a TPU,
+  specifying a cycle_length > 10 is recommended.
   """
-  empty_example = {}
-  for k in keys:
-    empty_example[k] = tf.zeros([0], dtype=tf.int64)
-    empty_example[k + "_position"] = tf.zeros([0], dtype=tf.int32)
-  keys_etc = empty_example.keys()
-
-  def write_packed_example(partial, outputs):
-    new_partial = empty_example.copy()
-    new_outputs = {}
-    for k in keys_etc:
-      new_outputs[k] = outputs[k].write(
-          outputs[k].size(),
-          tf.pad(partial[k], [[0, length - tf.size(partial[k])]]))
-    return new_partial, new_outputs
-
-  def map_fn(x):
-    """Internal function to flat_map over.
-
-    Consumes a batch of input examples and produces a variable number of output
-    examples.
-
-    Args:
-      x: a single example
+
+  def __init__(self, packed_length=256, spacing=0, queue_size=10,
+               chop_long_sequences=False):
+    self._packed_length = packed_length
+    self._spacing = spacing
+    self._queue_size = queue_size
+    self._chop_long_sequences = chop_long_sequences
+    self._num_sequences = None
+    self._token_dtype = None
+
+  def __call__(self, dataset, **kwargs):
+    if {"window_size", "cycle_length"}.intersection(kwargs):
+      return self._concurrent_pack(dataset, **kwargs)
+    return self._pack(dataset, **kwargs)
+
+  def _concurrent_pack(self, dataset, window_size=None, cycle_length=None,
+                       keys=None):
+    """Selects sensible default parallelism parameters based for a task."""
+
+    if window_size is None:
+      # This is a heuristic to fill all of the queues 10 times, and should do a
+      # reasonable job balancing parallelism (which benefits from lower window
+      # size) with packing efficiency (which suffers from edge effects when the
+      # window size is too low.)
+      window_size = int(self._packed_length / 8 * self._queue_size * 10)
+
+    if cycle_length is None:
+      # Typically binning one stream will saturate about 3 cores.
+
+      # Note on TPUs:
+      # cycle_length should still be explicitly set when training on TPUs,
+      # since the cpu count will be the local CPU count (which could be quite
+      # small), wereas the transforms will actually run on the TPU host
+      # controller which has a very robust CPU.
+      cycle_length = max([int(multiprocessing.cpu_count() / 3), 1])
+    return self._pack(dataset, window_size=window_size,
+                      cycle_length=cycle_length, keys=keys)
+
+  def _pack(self, dataset, window_size=None, cycle_length=None,
+            deterministic=False, keys=None):
+    """Main method for chaining together packing transformation steps."""
+    (dataset, self._num_sequences, self._token_dtype, keys
+    ) = self._standardize(dataset, keys)
+    if window_size is None:
+      dataset = self._scanning_pack(dataset)
+    else:
+      # Dataset.window splits nested Tensors.
+      re_zip = lambda *x: tf.data.Dataset.zip(x)
+      dataset = dataset.window(window_size).map(re_zip).interleave(
+          self._scanning_pack, cycle_length=cycle_length,
+          block_length=window_size,
+          num_parallel_calls=tf.data.experimental.AUTOTUNE)
+
+      if not deterministic:
+        # Sloppy interleave offers a marginal performance improvement.
+        options = tf.data.Options()
+        options.experimental_deterministic = False
+        dataset = dataset.with_options(options)
+
+    dataset = dataset.map(
+        self._finalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)
+    self._num_sequences, self._token_dtype = None, None
+
+    if keys:
+      def dict_pack(example):
+        output = {}
+        for i, key in enumerate(keys):
+          output[key] = example["contents"][:, i]
+          output[key + "_segmentation"] = example["segment"][:, i]
+          output[key + "_position"] = example["position"][:, i]
+        return output
+      dataset = dataset.map(dict_pack)
+    return dataset
+
+  def _standardize(self, dataset, keys):
+    """Force dataset structure into a tuple of Tensors."""
+    shapes = tf.compat.v1.data.get_output_shapes(dataset)
+
+    if isinstance(shapes, dict):
+      keys = tuple(shapes.keys())
+      dataset = dataset.map(lambda x: tuple(x[k] for k in keys))
+      shapes = tf.compat.v1.data.get_output_shapes(dataset)
+
+    if not all(isinstance(i, tf.TensorShape) for i in shapes):
+      # Internally this class expects tuples of Tensors, even for the degenerate
+      # case of a single sequence.
+      dataset = dataset.map(lambda x: (x,))
+      shapes = tf.compat.v1.data.get_output_shapes(dataset)
+
+    for s in shapes:
+      if not s.is_compatible_with(tf.TensorShape([None])):
+        raise ValueError("Tensors to be packed must be one-dimensional.")
+
+    if not shapes:
+      raise ValueError("Expected sequence dataset.")
+
+    if self._chop_long_sequences and len(shapes) != 1:
+      raise ValueError("chop_long_sequences expects a single sequence dataset.")
+
+    token_types = tf.compat.v1.data.get_output_types(dataset)
+    if len(set(token_types)) > 1:
+      raise ValueError("Inconsistent dtypes: {}".format(token_types))
+
+    return dataset, len(shapes), token_types[0], keys
+
+  def _eviction_fn(self, _):
+    return tuple(-tf.ones((self._packed_length,), dtype=self._token_dtype)
+                 for _ in range(self._num_sequences))
+
+  def _scan_initial_state(self):
+    """Create TensorArrays and indices to track bin assignment.
+
+    availability: TensorArray[queue_size, num_sequences]
+      This represents the number of tokens available in the ith bin.
+      See implementation note below.
+
+    contents: TensorArray[queue_size, num_sequences * 2]
+      This holds the actual contents of the packed strings as well as a bit
+      mask indicating where sequences begin. It is stored in a flat vector and
+      is accessed in offsets of packed_length.
+
+    top_index: scalar [0, queue_size)
+      Integer tensor indicating which index is the "top" bin. See implementation
+      note below.
+
+    IMPLEMENTATION_NOTE:
+      The FFD algorithm periodically pops the topmost queue and pushes a new
+      one to replace it. In order to replicate those semantics with a fixed size
+      TensorArray, indexing operations are shifted by top_index. For example,
+      instead of:
+        `queue_available.read(i)`
+
+      a read is instead performed as:
+        `queue_available.read((i - top_index) % queue_size)`
+
+      to account for the fact that the "ith" logical FFD queue is stored at
+      position j. This means that the pop / push update can be performed by
+      simply incrementing top_index. (And zeroing the old top_index position.)
+
     Returns:
-      a tf.data.Dataset
+      The state for the binning scan.
     """
-    partial = empty_example.copy()
-    i = tf.zeros([], dtype=tf.int32)
-    dynamic_batch_size = tf.shape(x[keys[0]])[0]
-    outputs = {}
-    for k in keys:
-      outputs[k] = tf.TensorArray(
-          tf.int64, size=0, dynamic_size=True, element_shape=[length])
-      outputs[k + "_position"] = tf.TensorArray(
-          tf.int32, size=0, dynamic_size=True, element_shape=[length])
-    def cond_fn(i, partial, outputs):
-      del partial, outputs
-      return i < dynamic_batch_size
-    def body_fn(i, partial, outputs):
-      """Body function for while_loop.
-
-      Args:
-        i: integer scalar
-        partial: dictionary of Tensor (partially-constructed example)
-        outputs: dictionary of TensorArray
-      Returns:
-        A triple containing the new values of the inputs.
-      """
-      can_append = True
-      one_example = {}
-      for k in keys:
-        val = x[k][i]
-        val = val[:tf.reduce_sum(tf.to_int32(tf.not_equal(val, 0)))]
-        one_example[k] = val
-      for k in keys:
-        can_append = tf.logical_and(
-            can_append,
-            tf.less_equal(
-                tf.size(partial[k]) + tf.size(one_example[k]), length))
-      def false_fn():
-        return write_packed_example(partial, outputs)
-      def true_fn():
-        return partial, outputs
-      partial, outputs = tf.cond(can_append, true_fn, false_fn)
-      new_partial = {}
-      for k in keys:
-        new_seq = one_example[k][:length]
-        new_seq_len = tf.size(new_seq)
-        new_partial[k] = tf.concat([partial[k], new_seq], 0)
-        new_partial[k + "_position"] = tf.concat(
-            [partial[k + "_position"],
-             tf.range(new_seq_len, dtype=tf.int32)], 0)
-      partial = new_partial
-      return i+1, partial, outputs
-
-    i, partial, outputs = tf.while_loop(
-        cond_fn, body_fn, (i, partial, outputs),
-        back_prop=False,
-        shape_invariants=(
-            tf.TensorShape([]),
-            {k: tf.TensorShape([None]) for k in keys_etc},
-            {k: tf.TensorShape(None) for k in keys_etc},
-            ))
-    partial, outputs = write_packed_example(partial, outputs)
-    packed = {k: outputs[k].stack() for k in keys_etc}
-    for k in keys:
-      packed[k + "_segmentation"] = (
-          tf.cumsum(tf.to_int32(tf.equal(packed[k + "_position"], 0)), axis=1) *
-          tf.to_int32(tf.not_equal(packed[k], 0)))
 
-    return tf.data.Dataset.from_tensor_slices(packed)
-  dataset = dataset.flat_map(map_fn)
-  return dataset
+    all_available = tf.ones((self._queue_size, self._num_sequences),
+                            dtype=INDEX_DTYPE) * self._packed_length
+    total_size = self._packed_length * self._queue_size
+    total_size_range = tf.range(total_size, dtype=INDEX_DTYPE)
+    empty = tf.zeros((total_size, self._num_sequences * 2),
+                     dtype=self._token_dtype)
+
+    availability = tf.TensorArray(
+        dtype=INDEX_DTYPE, size=self._queue_size, dynamic_size=False,
+        clear_after_read=False, element_shape=(self._num_sequences,)
+        ).scatter(tf.range(self._queue_size, dtype=INDEX_DTYPE), all_available)
+
+    contents = tf.TensorArray(
+        dtype=self._token_dtype, size=total_size, dynamic_size=False,
+        clear_after_read=False, element_shape=(self._num_sequences * 2,)
+        ).scatter(total_size_range, empty)
+
+    # Which index should be considered the "top" bucket for the purpose of
+    # the first-fit descending algorithm.
+    top_index = tf.zeros((), dtype=INDEX_DTYPE)
+
+    return availability, contents, top_index
+
+  def _scanning_pack(self, dataset):
+    """Apply scan based pack to a dataset."""
+    if self._chop_long_sequences:
+      dataset = dataset.map(lambda x: (x[:self._packed_length],))
+    else:
+      dataset = dataset.filter(lambda *x: tf.reduce_max(  # pylint: disable=g-long-lambda
+          tf.stack([tf.shape(i)[0] for i in x]), axis=0) <= self._packed_length)
+
+    # In order to retrieve the sequences which are still in the queue when the
+    # dataset is exhausted, we feed dummy sequences which are guaranteed to
+    # displace the remaining elements.
+    dataset = dataset.concatenate(
+        tf.data.Dataset.range(self._queue_size).map(self._eviction_fn))
+
+    initial_state = self._scan_initial_state()
+    step_fn = functools.partial(
+        _scan_step_fn, packed_length=self._packed_length,
+        queue_size=self._queue_size, spacing=self._spacing,
+        num_sequences=self._num_sequences, token_dtype=self._token_dtype)
+
+    dataset = dataset.apply(tf.data.experimental.scan(initial_state, step_fn))
+
+    is_real_sample = lambda valid_sample, _: valid_sample
+    return dataset.filter(is_real_sample)
+
+  def _compute_auxiliary_structure(self, contents_and_mask):
+    """Compute segment and position metadata."""
+    contents = contents_and_mask[:, :self._num_sequences]
+    start_mask = tf.cast(contents_and_mask[:, self._num_sequences:],
+                         dtype=INDEX_DTYPE)
+
+    segment = tf.cumsum(start_mask, axis=0)
+    uniform_count = tf.ones_like(segment[:, 0])
+    position = []
+    for i in range(self._num_sequences):
+      segment_slice = segment[:, i]
+      counts = tf.math.segment_sum(uniform_count, segment[:, i])
+      position.append(tf.range(self._packed_length) -  tf.cumsum(
+          tf.gather(counts, segment_slice - 1) * start_mask[:, i]))
+    position = tf.concat([i[:, tf.newaxis] for i in position], axis=1)
+
+    # Correct for padding tokens.
+    pad_mask = tf.cast(tf.not_equal(contents, 0), dtype=INDEX_DTYPE)
+    segment *= pad_mask
+    position *= pad_mask
+
+    return segment, position
+
+  def _finalize(self, _, contents):
+    """Structure output and compute segment and position metadata."""
+
+    # The output shape information is lost during the filter; however we can
+    # guarantee the shape. (That's the point of this exercise, after all!)
+    contents.set_shape((self._packed_length, self._num_sequences * 2))
+
+    # Both the dummy branch of the scan step function and the eviction dataset
+    # use vectors of minus one. The cost of this check is negligible and the
+    # leakage of such dummy sequences would be difficult to debug downstream.
+    check_leaks = tf.assert_none_equal(contents, -tf.ones_like(contents))
+    with tf.control_dependencies([check_leaks]):
+      contents = tf.identity(contents)
+
+    segment, position = self._compute_auxiliary_structure(contents)
+    return {"contents": contents[:, :self._num_sequences],
+            "segment": segment, "position": position}
+
+
+@tf.autograph.to_graph
+def _scan_step_fn(state, example, packed_length, queue_size, spacing,
+                  num_sequences, token_dtype):  # pylint: disable=g-doc-args
+  """Transform function used by tf.data.experimental.scan to process an example.
+
+  This is written as a stateless function rather than a class method because we
+  trace it with AutoGraph (in order to simplify the conditional), and this way
+  we don't have to worry about handling re-tracing semantics.
+
+  Args:
+    See the SequenceDatasetPacker class.
+
+  Returns:
+    The updated queue state, and either a packed example or a dummy sequence
+    which will be filtered out downstream.
+  """
+
+  # Convert TensorArray tuples to lists since we'll need to replace them.
+  availability, contents, top_index = state
+
+  lengths = tf.concat([tf.shape(i) for i in example], axis=0)
+  start_availability = availability.stack()
+  can_fit = tf.reduce_all(tf.greater_equal(start_availability, lengths), axis=1)
+  any_can_fit = tf.reduce_any(can_fit, axis=0)
+
+  # AutoGraph will convert this block to a tf.cond
+  if any_can_fit:
+    # This indicates where in the FFD queue rotation a given index sits
+    shifted_range = (
+        tf.range(queue_size, dtype=INDEX_DTYPE) - top_index) % queue_size
+
+    # Mark any indices which cannot accommodate the current example.
+    exclusion_mask = tf.cast(tf.logical_not(can_fit), INDEX_DTYPE) * queue_size
+
+    # Index in [0, queue_size) in which to place the sample. Note, this index
+    # is the position in the actual TensorArray, not the index of the FFD queue.
+    queue_index = (tf.reduce_min(shifted_range + exclusion_mask) +
+                   top_index) % queue_size
+
+    # NOTE(taylorrobie): We emit a non-empty Tensor for downstream checks.
+    output_contents = -tf.ones((1, num_sequences), dtype=token_dtype)
+
+  else:
+    index_range = top_index * packed_length + tf.range(packed_length)
+    output_contents = contents.gather(index_range)
+
+    # Reset the queue state.
+    availability = availability.write(
+        top_index, packed_length * tf.ones((num_sequences,), dtype=INDEX_DTYPE))
+    empty_contents = tf.zeros((packed_length, num_sequences * 2),
+                              dtype=token_dtype)
+    contents = contents.scatter(index_range, empty_contents)
+
+    queue_index = top_index
+    top_index = (top_index + 1) % queue_size
+
+  pre_assign_availability = availability.read(queue_index)
+  space_left = pre_assign_availability - lengths - spacing
+  availability = availability.write(queue_index, space_left)
+
+  # ============================================================================
+  # == Update contents =========================================================
+  # ============================================================================
+  # Consider the following case for a seq-to-seq packing:
+  #   (padding is represented as underscores)
+  #
+  #   Queue starting state:
+  #     [1, 3, 2, 4, 6, 1, _, _, _, _, _, ...]
+  #     [5, 9, _, _, _, _, _, _, _, _, _, ...]
+  #
+  #   Examples:
+  #     [4, 2, 4], [3]
+  #
+  #   Desired new queue state:
+  #     [1, 3, 2, 4, 6, 1, _, _, 4, 2, 4, _, _, ...]
+  #     [5, 9, _, _, 3, _, _, _, _, _, _, _, _, ...]
+  #
+  # This could be acomplished by creating a TensorArray for each of the two
+  # sequences, and scattering into the respective arrays. However TensorArray
+  # writes are extremely expensive relative to other operations. So instead we
+  # store the contents in a single TensorArray of shape (packed_length, 2), and
+  # we pad and concatenate the examples such that they can be added in a single
+  # assign:
+  #
+  #              [_, _, _, _, 4, 2, 4]
+  #              [3, _, _, _, _, _, _]
+  #                        +
+  #  [1, 3, 2, 4, 6, 1, _, _, _, _, _, ...]
+  #  [5, 9, _, _, _, _, _, _, _, _, _, ...]
+  #
+  # And in practice, the extra work of padding is neglidgable compared to
+  # the gain from vectorizing the TensorArray assign. We also store a bit mask
+  # denoting where sequences start which is used to compute segment and
+  # position metadata:
+  #
+  #              [_, _, _, _, 1, _, _]
+  #              [1, _, _, _, _, _, _]
+  #                        +
+  #  [1, _, _, _, _, _, _, _, _, _, _, ...]
+  #  [1, _, _, _, _, _, _, _, _, _, _, ...]
+  #
+  # Both the contents and the mask are concatenated in the same TensorArray
+  # for performance.
+
+  start_index = packed_length - pre_assign_availability
+  end_index = start_index + lengths
+  leftmost = tf.reduce_min(start_index, axis=0)
+  rightmost = tf.reduce_max(end_index, axis=0)
+  delta = rightmost - leftmost
+  pad_indices = [tf.stack((start_index[i] - leftmost, rightmost - end_index[i]))
+                 for i in range(num_sequences)]
+
+  padded_examples = [tf.pad(ex, padding[tf.newaxis, :])
+                     for ex, padding in zip(example, pad_indices)]
+  padded_examples = tf.transpose(tf.stack(padded_examples))
+  mask_update = tf.one_hot(start_index - leftmost, delta,
+                           dtype=contents.dtype, axis=0)
+
+  content_update = tf.concat([padded_examples, mask_update], axis=1)
+
+  index_range = (queue_index * packed_length +  # Offset into the right section.
+                 tf.range(delta, dtype=INDEX_DTYPE) + leftmost)
+  contents = contents.scatter(index_range, contents.gather(index_range) +
+                              content_update)
+
+  state = (availability, contents, top_index)
+  return state, (tf.logical_not(any_can_fit), output_contents)
 
 
 def make_tmp_dir(suffix="", prefix="tmp", dir=None):  # pylint: disable=redefined-builtin
diff --git a/tensor2tensor/data_generators/generator_utils_test.py b/tensor2tensor/data_generators/generator_utils_test.py
index 2b4fc1fe4..3b83a968d 100644
--- a/tensor2tensor/data_generators/generator_utils_test.py
+++ b/tensor2tensor/data_generators/generator_utils_test.py
@@ -30,6 +30,67 @@
 import tensorflow as tf
 
 
+INPUTS = (
+    (1, 2, 3),
+    (4, 5,),
+    (6,),
+)
+TARGETS = (
+    (10,),
+    (20, 30, 40),
+    (50, 60,),
+)
+INPUTS_PACKED = (
+    (1, 2, 3, 4, 5),
+    (6, 0, 0, 0, 0),
+)
+INPUTS_SEGMENTATION = (
+    (1, 1, 1, 2, 2),
+    (1, 0, 0, 0, 0),
+)
+INPUTS_POSITION = (
+    (0, 1, 2, 0, 1),
+    (0, 0, 0, 0, 0),
+)
+TARGETS_PACKED = (
+    (10, 20, 30, 40, 0),
+    (50, 60, 0, 0, 0),
+)
+TARGETS_SEGMENTATION = (
+    (1, 2, 2, 2, 0),
+    (1, 1, 0, 0, 0),
+)
+TARGETS_POSITION = (
+    (0, 0, 1, 2, 0),
+    (0, 1, 0, 0, 0),
+)
+
+
+def example_generator():
+  for i, t in zip(INPUTS, TARGETS):
+    yield {"inputs": list(i), "targets": list(t)}
+
+
+def trim_right(x):
+  x = {k: list(v) for k, v in x.items()}
+  while all(x.values()) and not any(i[-1] for i in x.values()):
+    _ = [i.pop() for i in x.values()]
+  return x
+
+
+def reference_packing(trim_fn=None):
+  no_trim = lambda x: {k: list(v) for k, v in x.items()}
+  trim_fn = trim_fn or no_trim
+  outputs = [INPUTS_PACKED, INPUTS_POSITION, INPUTS_SEGMENTATION,
+             TARGETS_PACKED, TARGETS_POSITION, TARGETS_SEGMENTATION]
+  for i, i_pos, i_seg, t, t_pos, t_seg in zip(*outputs):
+    output = trim_fn({"inputs": i, "inputs_position": i_pos,
+                      "inputs_segmentation": i_seg})
+    output.update(trim_fn({"targets": t, "targets_position": t_pos,
+                           "targets_segmentation": t_seg}))
+    yield output
+
+
 class GeneratorUtilsTest(tf.test.TestCase):
 
   def testGenerateFiles(self):
@@ -120,5 +181,33 @@ def testGetOrGenerateTxtVocab(self):
     self.assertIsNotNone(vocab2)
     self.assertEqual(vocab1.dump(), vocab2.dump())
 
+  def testPacking(self):
+    packed = generator_utils.pack_examples(
+        example_generator(), has_inputs=True, packed_length=5, queue_size=2,
+        spacing=0)
+    for example, reference in zip(packed, reference_packing(trim_right)):
+      self.assertAllEqual(set(example.keys()), set(reference.keys()))
+      for k in reference:
+        self.assertAllEqual(example[k], reference[k])
+
+  def testDatasetPacking(self):
+    dataset = tf.data.Dataset.from_generator(
+        example_generator,
+        output_types={"inputs": tf.int64, "targets": tf.int64},
+        output_shapes={"inputs": tf.TensorShape((None,)),
+                       "targets": tf.TensorShape((None,))}
+    )
+    dataset = generator_utils.pack_dataset(
+        dataset, length=5, keys=("inputs", "targets"), use_custom_ops=False)
+
+    with tf.Session().as_default() as sess:
+      batch = dataset.make_one_shot_iterator().get_next()
+      for reference in reference_packing():
+        example = sess.run(batch)
+        self.assertAllEqual(set(example.keys()), set(reference.keys()))
+        for k in reference:
+          self.assertAllEqual(example[k], reference[k])
+
+
 if __name__ == "__main__":
   tf.test.main()

From 6025e780943517d95388cd98dfd5001aaa3cbf26 Mon Sep 17 00:00:00 2001
From: Zi Yang <ziy@google.com>
Date: Mon, 3 Jun 2019 15:39:25 -0700
Subject: [PATCH 2092/2720] Added `as_text` flag for model exporting, defaulted
 to True for backward compatibility.

PiperOrigin-RevId: 251323621
---
 tensor2tensor/serving/export.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index ff582db38..8d103a39a 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -43,6 +43,11 @@
     "If None, we will use the latest checkpoint stored in the directory "
     "specified by --output_dir")
 
+tf.flags.DEFINE_bool(
+    "as_text", True,
+    "Whether to write the SavedModel proto in text format. Defaults to `False`."
+)
+
 
 def _get_hparams_path():
   """Get hyper-parameters file path."""
@@ -186,7 +191,7 @@ def main(_):
   exporter = tf.estimator.FinalExporter(
       "exporter",
       lambda: problem.serving_input_fn(hparams, decode_hparams, FLAGS.use_tpu),
-      as_text=True)
+      as_text=FLAGS.as_text)
 
   exporter.export(
       estimator,

From 20a01166e575441eef27572dec2fef4eb562fbf0 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Tue, 4 Jun 2019 05:29:33 -0700
Subject: [PATCH 2093/2720] [TRAX] Add a jax.device_get call to
 trax._save_replicated.

By using device_get, trax can tell JAX that it intends to fetch all of the weights from the device, rather than synchronously transferring them one by one.

PiperOrigin-RevId: 251420028
---
 tensor2tensor/trax/trax.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 604df1293..8e6ac36ad 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -173,6 +173,10 @@ def _save_replicated(opt_state, step, history, n_devices, output_dir, keep):
   if n_devices > 1:
     first_replica = lambda x: x[0]
     opt_state = layers.nested_map(opt_state, first_replica)
+  # This line, while optional, allows JAX to transfer arrays from the device to
+  # the host in parallel, which is particularly important for cloud TPU.
+  if backend.get_name() == "jax":
+    opt_state = jax.device_get(opt_state)
   save_state(State(params=opt_state, step=step, history=history),
              output_dir, keep=keep)
 

From 00ccea885299e821646ab3796570ee7bbc6852c5 Mon Sep 17 00:00:00 2001
From: Peng Wang <wangpeng@google.com>
Date: Tue, 4 Jun 2019 17:16:08 -0700
Subject: [PATCH 2094/2720] Adds switch `--tf_xla` to toggle XLA for tf-numpy

PiperOrigin-RevId: 251548636
---
 tensor2tensor/trax/trainer.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index efbd582da..edf3a1f43 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -47,6 +47,7 @@
 flags.DEFINE_integer("log_level", logging.INFO, "Log level.")
 flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
 flags.DEFINE_bool("tf_eager", False, "Whether we're running TF in eager mode.")
+flags.DEFINE_bool("tf_xla", False, "Whether to turn on XLA for TF.")
 
 
 def _default_output_dir():
@@ -91,6 +92,9 @@ def main(_):
   if FLAGS.tf_eager:
     tf.enable_eager_execution()
 
+  if FLAGS.tf_xla:
+    tf.config.optimizer.set_jit(True)
+
   _setup_gin()
 
   # Setup output directory

From 4ce5a1293ff4559bde55a9ea46031d55f28feae7 Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Tue, 4 Jun 2019 23:11:15 -0700
Subject: [PATCH 2095/2720] Add get_decode_end_id() to transformer.py so
 inherited models can use a different EOS symbol during decoding.

PiperOrigin-RevId: 251583883
---
 tensor2tensor/models/transformer.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index d98c82cce..ccf37db01 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -586,6 +586,7 @@ def forced_logits():
             tf.less(i, partial_targets_length), forced_logits, lambda: ret)
       return ret, cache
 
+    eos_id = self.get_decode_end_id() or beam_search.EOS_ID
     ret = fast_decode_tpu(
         encoder_output=encoder_output,
         encoder_decoder_attention_bias=encoder_decoder_attention_bias,
@@ -598,7 +599,8 @@ def forced_logits():
         top_beams=top_beams,
         alpha=alpha,
         batch_size=batch_size,
-        force_decode_length=self._decode_hparams.force_decode_length)
+        force_decode_length=self._decode_hparams.force_decode_length,
+        eos_id=eos_id)
     if partial_targets is not None:
       if beam_size <= 1 or top_beams <= 1:
         ret["outputs"] = ret["outputs"][:, partial_targets_length:]
@@ -617,6 +619,14 @@ def get_decode_start_id(self):
     """
     return None
 
+  def get_decode_end_id(self):
+    """Returns the id of the output symbol that terminates decoding.
+
+    This method can be overridden by a different model. The id returned by this
+    method is used to check if the generation is complete during decoding.
+    """
+    return None
+
   def _fast_decode(self,
                    features,
                    decode_length,
@@ -818,6 +828,7 @@ def forced_logits():
       return ret, cache
 
     sos_id = self.get_decode_start_id() or 0
+    eos_id = self.get_decode_end_id() or beam_search.EOS_ID
 
     ret = fast_decode(
         encoder_output=encoder_output,
@@ -832,7 +843,8 @@ def forced_logits():
         alpha=alpha,
         batch_size=batch_size,
         force_decode_length=self._decode_hparams.force_decode_length,
-        sos_id=sos_id)
+        sos_id=sos_id,
+        eos_id=eos_id)
     if partial_targets is not None:
       if beam_size <= 1 or top_beams <= 1:
         ret["outputs"] = ret["outputs"][:, partial_targets_length:]

From e7c4a8b5b85ff49943d55910bfc80356c3348f41 Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Wed, 5 Jun 2019 12:07:28 -0700
Subject: [PATCH 2096/2720] Fix log_prob accumulation during decoding: only
 include up to (and inclusive of) the first EOS generated.

PiperOrigin-RevId: 251690778
---
 tensor2tensor/models/transformer.py | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index ccf37db01..94f6dc9e6 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1023,11 +1023,14 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       next_id = common_layers.sample_with_temperature(
           logits, temperature, keep_top)
 
-      hit_eos |= tf.equal(next_id, eos_id)
-
       log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],
                                   axis=1)
-      log_prob += tf.gather_nd(log_probs, log_prob_indices)
+      log_prob += tf.gather_nd(
+          log_probs, log_prob_indices) * (1 - tf.to_float(hit_eos))
+      # Note(thangluong): we purposely update hit_eos after aggregating log_prob
+      # There is a subtle detail here that we want to include log_probs up to
+      # (and inclusive of) the first eos generated, but not subsequent tokens.
+      hit_eos |= tf.equal(next_id, eos_id)
 
       next_id = tf.expand_dims(next_id, axis=1)
       decoded_ids = tf.transpose(decoded_ids)
@@ -1167,14 +1170,19 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
         temperature = 0.0
       next_id = common_layers.sample_with_temperature(
           logits, temperature, keep_top)
-      hit_eos |= tf.equal(next_id, eos_id)
 
       log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],
                                   axis=1)
-      log_prob += tf.gather_nd(log_probs, log_prob_indices)
+      log_prob += tf.gather_nd(
+          log_probs, log_prob_indices) * (1 - tf.to_float(hit_eos))
+      # Note(thangluong): we purposely update hit_eos after aggregating log_prob
+      # There is a subtle detail here that we want to include log_probs up to
+      # (and inclusive of) the first eos generated, but not subsequent tokens.
+      hit_eos |= tf.equal(next_id, eos_id)
 
       next_id = tf.expand_dims(next_id, axis=1)
       decoded_ids = tf.concat([decoded_ids, next_id], axis=1)
+
       return i + 1, hit_eos, next_id, decoded_ids, cache, log_prob
 
     def is_not_finished(i, hit_eos, *_):

From dfacea61347315451eed40961e3e5d89e362c509 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 5 Jun 2019 13:40:38 -0700
Subject: [PATCH 2097/2720] Make convolutions re-batch in Trax, remove Rebatch
 layer.

PiperOrigin-RevId: 251708750
---
 tensor2tensor/trax/layers/combinators.py      | 49 -------------------
 tensor2tensor/trax/layers/combinators_test.py | 18 -------
 tensor2tensor/trax/layers/convolution.py      | 21 +++++++-
 tensor2tensor/trax/layers/convolution_test.py |  6 +++
 tensor2tensor/trax/models/atari_cnn.py        |  4 +-
 5 files changed, 28 insertions(+), 70 deletions(-)

diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index a24a0f093..9e6efb610 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -23,7 +23,6 @@
 import six
 
 from tensor2tensor.trax import backend
-from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers import base
 
 
@@ -483,51 +482,3 @@ def new_parameters(self, input_shape, rng):
           raise ValueError('Map layer can only be applied to list of elements '
                            'with the same shapes. Shapes: %s' % str(shape))
     return self._layer.initialize(first_shape, rng)
-
-
-class Rebatch(base.Layer):
-  """Combinator for treating the first `n` dims as batch.
-
-  Args:
-    layer: subclass of base.Layer, a layer to apply to the input.
-    n_batch_dims: int, the number of leading dimensions to consider as batch.
-
-  Returns:
-    A new layer that will reshape the input into a virtual batch, apply the
-    layer and unbatch the virtual batch.
-  """
-
-  def __init__(self, layer, n_batch_dims=1):
-    super(Rebatch, self).__init__()
-    self._layer = layer
-    self._n_batch_dims = n_batch_dims
-
-  def _modify_shape(self, input_shape):
-    input_shape = tuple(input_shape)
-    batch_dims, non_batch_dims = (input_shape[:self._n_batch_dims],
-                                  input_shape[self._n_batch_dims:])
-    new_batch_dim = six.moves.reduce(operator.mul, batch_dims)
-    return (new_batch_dim,) + non_batch_dims, batch_dims
-
-  def _unmodify_shape(self, input_shape, batch_dims):
-    return batch_dims + tuple(input_shape[1:])
-
-  def _modify(self, inp):
-    modified_shape, batch_dims = self._modify_shape(inp.shape)
-    return np.reshape(inp, modified_shape), batch_dims
-
-  def _unmodify(self, inp, batch_dims):
-    return np.reshape(inp, self._unmodify_shape(inp.shape, batch_dims))
-
-  def call(self, inp, params=(), **kwargs):
-    if isinstance(inp, (tuple, list)):
-      # TODO(afrozm): This should be easy to do though.
-      # Tip from Lukasz - base.nested_map(self._modify, inp)
-      raise ValueError("Rebatch doesn't support list/tuple inputs now.")
-    inp, batch_dims = self._modify(inp)
-    out = self._layer(inp, params=params, **kwargs)
-    return self._unmodify(out, batch_dims)
-
-  def new_parameters(self, input_shape, rng):
-    modified_shape, _ = self._modify_shape(input_shape)
-    return self._layer.initialize(modified_shape, rng)
diff --git a/tensor2tensor/trax/layers/combinators_test.py b/tensor2tensor/trax/layers/combinators_test.py
index bc5b8b593..4a2598db4 100644
--- a/tensor2tensor/trax/layers/combinators_test.py
+++ b/tensor2tensor/trax/layers/combinators_test.py
@@ -21,7 +21,6 @@
 from absl.testing import absltest
 from tensor2tensor.trax.layers import base
 from tensor2tensor.trax.layers import combinators
-from tensor2tensor.trax.layers import convolution
 
 
 class CombinatorLayerTest(absltest.TestCase):
@@ -47,23 +46,6 @@ def test_select(self):
         combinators.Select(1), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
-  def test_rebatch(self):
-    input_shape = (29, 5, 5, 20)
-    result_shape = base.check_shape_agreement(
-        convolution.Conv(30, (3, 3)), input_shape)
-    self.assertEqual(result_shape, (29, 3, 3, 30))
-
-    input_shape = (29, 5, 5, 20)
-    result_shape = base.check_shape_agreement(
-        combinators.Rebatch(convolution.Conv(30, (3, 3)), n_batch_dims=1),
-        input_shape)
-    self.assertEqual(result_shape, (29, 3, 3, 30))
-
-    input_shape = (19, 29, 5, 5, 20)
-    result_shape = base.check_shape_agreement(
-        combinators.Rebatch(convolution.Conv(30, (3, 3)), n_batch_dims=2),
-        input_shape)
-    self.assertEqual(result_shape, (19, 29, 3, 3, 30))
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/tensor2tensor/trax/layers/convolution.py b/tensor2tensor/trax/layers/convolution.py
index d721f2ec0..4e3b1866e 100644
--- a/tensor2tensor/trax/layers/convolution.py
+++ b/tensor2tensor/trax/layers/convolution.py
@@ -20,9 +20,12 @@
 from __future__ import print_function
 
 import itertools
+import operator
 
 from jax import lax
+import six
 
+from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers import base
 from tensor2tensor.trax.layers import initializers as init
 
@@ -52,12 +55,24 @@ def __init__(self, filters, kernel_size, strides=None, padding='VALID',
   def stack_items_to_pass(self):
     return 1
 
+  def _check_nhwc(self):
+    msg = 'Convolutions on more than 4 dimensions only supported in NHWC.'
+    assert self._lhs_spec == self._out_spec == 'NHWC', msg
+
   def call(self, x, params=(), **kwargs):
     del kwargs
     w, b = params
-    return lax.conv_general_dilated(
+    x_shape = list(x.shape)
+    if len(x_shape) > 4:
+      self._check_nhwc()
+      new_batch_dim = six.moves.reduce(operator.mul, x_shape[:-3])
+      x = np.reshape(x, [new_batch_dim] + x_shape[-3:])
+    res = lax.conv_general_dilated(
         x, w, self._strides, self._padding, self._one, self._one,
         self._dimension_numbers) + b
+    if len(x_shape) > 4:
+      res = np.reshape(res, x_shape[:-3] + list(res.shape[-3:]))
+    return res
 
   def _kernel_shape(self, input_shape):
     """Helper to calculate the kernel shape."""
@@ -67,6 +82,10 @@ def _kernel_shape(self, input_shape):
             next(kernel_size_iter) for c in self._rhs_spec]
 
   def new_parameters(self, input_shape, rng):
+    if len(input_shape) > 4:
+      self._check_nhwc()
+      new_batch_dim = six.moves.reduce(operator.mul, input_shape[:-3])
+      input_shape = [new_batch_dim] + list(input_shape[-3:])
     kernel_shape = self._kernel_shape(input_shape)
     bias_shape = [self._filters if c == 'C' else 1 for c in self._out_spec]
     bias_shape = tuple(itertools.dropwhile(lambda x: x == 1, bias_shape))
diff --git a/tensor2tensor/trax/layers/convolution_test.py b/tensor2tensor/trax/layers/convolution_test.py
index d52db6289..62b931fc0 100644
--- a/tensor2tensor/trax/layers/convolution_test.py
+++ b/tensor2tensor/trax/layers/convolution_test.py
@@ -31,6 +31,12 @@ def test_conv(self):
         convolution.Conv(30, (3, 3)), input_shape)
     self.assertEqual(result_shape, (29, 3, 3, 30))
 
+  def test_conv_rebatch(self):
+    input_shape = (3, 29, 5, 5, 20)
+    result_shape = base.check_shape_agreement(
+        convolution.Conv(30, (3, 3)), input_shape)
+    self.assertEqual(result_shape, (3, 29, 3, 3, 30))
+
 
 if __name__ == "__main__":
   absltest.main()
diff --git a/tensor2tensor/trax/models/atari_cnn.py b/tensor2tensor/trax/models/atari_cnn.py
index e9128b348..1c88b1e9e 100644
--- a/tensor2tensor/trax/models/atari_cnn.py
+++ b/tensor2tensor/trax/models/atari_cnn.py
@@ -36,9 +36,9 @@ def AtariCnn(hidden_sizes=(32, 32), output_size=128):
       ),
       # Concatenated on the last axis.
       tl.Concatenate(axis=-1),  # (B, T, H, W, 4C)
-      tl.Rebatch(tl.Conv(hidden_sizes[0], (5, 5), (2, 2), 'SAME'), 2),
+      tl.Conv(hidden_sizes[0], (5, 5), (2, 2), 'SAME'),
       tl.Relu(),
-      tl.Rebatch(tl.Conv(hidden_sizes[1], (5, 5), (2, 2), 'SAME'), 2),
+      tl.Conv(hidden_sizes[1], (5, 5), (2, 2), 'SAME'),
       tl.Relu(),
       tl.Flatten(num_axis_to_keep=2),  # B, T and rest.
       tl.Dense(output_size),

From 999c4123bcf012ae362650d732ba7f5a6c7becd0 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 5 Jun 2019 15:20:46 -0700
Subject: [PATCH 2098/2720] Pass dtype to layer.initialize which allows to
 remove layer.default_input_is_int.

PiperOrigin-RevId: 251730169
---
 tensor2tensor/trax/inputs.py                  | 38 +++++++++++-----
 tensor2tensor/trax/layers/README.md           |  5 ++-
 tensor2tensor/trax/layers/attention.py        |  7 +--
 tensor2tensor/trax/layers/base.py             | 44 ++++++++-----------
 tensor2tensor/trax/layers/combinators.py      | 39 +++++++---------
 tensor2tensor/trax/layers/convolution.py      |  3 +-
 tensor2tensor/trax/layers/core.py             | 16 ++++---
 tensor2tensor/trax/layers/core_test.py        |  4 +-
 tensor2tensor/trax/layers/normalization.py    |  8 ++--
 tensor2tensor/trax/models/atari_cnn.py        |  1 +
 tensor2tensor/trax/models/atari_cnn_test.py   |  2 +-
 .../models/research/chunked_transformer.py    |  5 ++-
 tensor2tensor/trax/models/resnet.py           |  2 +
 tensor2tensor/trax/rlax/ppo.py                |  8 +++-
 tensor2tensor/trax/rlax/ppo_test.py           |  6 +--
 tensor2tensor/trax/trax.py                    |  3 +-
 tensor2tensor/trax/trax_test.py               |  3 +-
 17 files changed, 107 insertions(+), 87 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 2c82e51f9..943d40978 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -43,10 +43,14 @@
 #     the shape of examples is [batch_fun.eval_batch_size, ...]
 # * input_shape: the shape of inputs
 #     the [...] above, without batch size
+# * input_dtype: the data type of inputs
+
 
 Inputs = collections.namedtuple(
     "_Inputs",
-    ["train_stream", "train_eval_stream", "eval_stream", "input_shape"])
+    ["train_stream", "train_eval_stream", "eval_stream",
+     "input_shape", "input_dtype"]
+)
 
 # How many examples from the stream to skip at random during training.
 # For now, we skip at most 100K examples for efficiency.
@@ -76,9 +80,12 @@ def inputs(n_devices, dataset_name, data_dir=None, input_name=None,
   data_dir = os.path.expanduser(data_dir)
 
   (train_batches, train_eval_batches, eval_batches,
-   input_name, input_shape) = _train_and_eval_batches(
+   input_name, input_shape, input_dtype) = _train_and_eval_batches(
        dataset_name, data_dir, input_name, n_devices)
 
+  if input_dtype == np.uint8:  # TPUs don't like uint8s, we cast to ints.
+    input_dtype = np.int32
+
   def numpy_stream(dataset):
     return dataset_to_stream(
         dataset, input_name,
@@ -88,15 +95,17 @@ def numpy_stream(dataset):
     length = input_shape[0]
     input_shape = tuple(
         [tuple([length // n_chunks] + list(input_shape)[1:])] * n_chunks)
+    input_dtype = tuple([input_dtype] * n_chunks)
   if append_targets:
     # TODO(lukaszkaiser): remove the assumption that input and target
     # shapes are the same, which is used below for now.
     input_shape = (input_shape, input_shape)
+    input_dtype = (input_dtype, input_dtype)
 
   return Inputs(train_stream=lambda: numpy_stream(train_batches),
                 train_eval_stream=lambda: numpy_stream(train_eval_batches),
                 eval_stream=lambda: numpy_stream(eval_batches),
-                input_shape=input_shape)
+                input_shape=input_shape, input_dtype=input_dtype)
 
 
 @gin.configurable(blacklist=["n_devices"])
@@ -144,7 +153,8 @@ def random_minibatches():
   return Inputs(train_stream=random_minibatches,
                 train_eval_stream=random_minibatches,
                 eval_stream=random_minibatches,
-                input_shape=input_shape_without_batch)
+                input_shape=input_shape_without_batch,
+                input_dtype=input_dtype)
 
 
 def dataset_to_stream(dataset, input_name, n_chunks=0, append_targets=False):
@@ -153,9 +163,9 @@ def dataset_to_stream(dataset, input_name, n_chunks=0, append_targets=False):
     inp, out = example[0][input_name], example[1]
     # Some accelerators don't handle uint8 well, cast to int.
     if isinstance(inp, np.uint8):
-      inp = inp.astype(np.uint32)
+      inp = inp.astype(np.int32)
     if isinstance(out, np.uint8):
-      out = out.astype(np.uint32)
+      out = out.astype(np.int32)
     if len(out.shape) > 1 and out.shape[-1] == 1:
       out = np.squeeze(out, axis=-1)
     if n_chunks > 0:
@@ -214,9 +224,10 @@ def train_and_eval_dataset(dataset_name, data_dir, train_shuffle_files=True,
   return train, valid, info.features, keys
 
 
-def _make_info(shape_list, n_classes):
+def _make_info(shape_list, n_classes, dtype):
   """Create an info-like tuple for feature given some shapes and vocab size."""
-  feature_info = collections.namedtuple("FeatureInfo", ["shape", "n_classes"])
+  feature_info = collections.namedtuple(
+      "FeatureInfo", ["shape", "n_classes", "dtype"])
   cur_shape = list(shape_list[0])
   # We need to merge the provided shapes, put None where they disagree.
   for shape in shape_list:
@@ -226,7 +237,7 @@ def _make_info(shape_list, n_classes):
       if cur_shape[i] is not None:
         if shape[i] != cur_shape[i]:
           cur_shape[i] = None
-  return feature_info(cur_shape, n_classes)
+  return feature_info(cur_shape, n_classes, dtype)
 
 
 def _select_features(example, feature_list=None):
@@ -272,8 +283,10 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
     target_shapes.append(list(example["targets"].shape))
   input_vocab_size = hparams.vocab_size[input_key]
   target_vocab_size = hparams.vocab_size["targets"]
-  input_info = _make_info(input_shapes, input_vocab_size)
-  target_info = _make_info(target_shapes, target_vocab_size)
+  input_dtype = examples[0][input_key].dtype
+  target_dtype = examples[0]["targets"].dtype
+  input_info = _make_info(input_shapes, input_vocab_size, input_dtype)
+  target_info = _make_info(target_shapes, target_vocab_size, target_dtype)
   info = {input_key: input_info, "targets": target_info}
   return train_dataset, eval_dataset, info, supervised_keys
 
@@ -449,5 +462,6 @@ def _train_and_eval_batches(dataset, data_dir, input_name, n_devices):
       n_devices=n_devices)
   input_name = input_name or input_names[0]
   input_shape = features_info[input_name].shape
+  input_dtype = features_info[input_name].dtype
   return (train_batches, train_eval_batches, eval_batches,
-          input_name, list(input_shape))
+          input_name, list(input_shape), input_dtype)
diff --git a/tensor2tensor/trax/layers/README.md b/tensor2tensor/trax/layers/README.md
index 679552bbc..fbdaedd50 100644
--- a/tensor2tensor/trax/layers/README.md
+++ b/tensor2tensor/trax/layers/README.md
@@ -23,8 +23,9 @@ and call functions to be used as follows.
 ```python
 layer = MyLayer()
 x = np.zeros(10)
-params = layer.initialize(x.shape)
-output = layer(x, params)
+rng = random.get_prng(0)
+params = layer.initialize(x.shape, x.dtype, rng)
+output = layer(x, params, rng=rng)
 ```
 
 ## Decorator
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 00763b551..1934e7da2 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -51,9 +51,10 @@ def EncoderDecoderMask(x, **unused_kwargs):
 
 
 # Positional encoding.
-def _positional_encoding_new_params(input_shape, rng, max_len=2048):  # pylint: disable=invalid-name
+def _positional_encoding_new_params(  # pylint: disable=invalid-name
+    input_shape, input_dtype, rng, max_len=2048):
   """Helper: create positional encoding parameters."""
-  del rng
+  del input_dtype, rng
   d_feature = input_shape[-1]
   pe = onp.zeros((max_len, d_feature), dtype=onp.float32)
   position = onp.arange(0, max_len)[:, onp.newaxis]
@@ -190,7 +191,7 @@ def MultiHeadedAttention(
   ]
 
 
-@base.layer(input_is_int=True, stack_items_to_pass=0)
+@base.layer(stack_items_to_pass=0)
 def ShiftRight(x, **unused_kwargs):
   """Layer to shift the tensor to the right by padding on axis 1."""
   if not isinstance(x, (list, tuple)):  # non-chunked inputs
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 3b58d156d..49d473f65 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -46,14 +46,15 @@ def call(self, x, params=(), **kwargs):
     """Call this layer in input x using the given parameters."""
     raise NotImplementedError
 
-  def new_parameters(self, input_shape, rng):
-    """Create new parameters for the layer given an input shape and rng.
+  def new_parameters(self, input_shape, input_dtype, rng):
+    """Create new parameters for the layer given an input shape, dtype and rng.
 
     Note that all arguments and return values can be tuples or dictionaries
     or arbitraty nested structures composed of tuples and dictionaries.
 
     Args:
       input_shape: a tuple representing the shape of the input.
+      input_dtype: numpy dtype of the input.
       rng: random number generator.
 
     Returns:
@@ -66,10 +67,6 @@ def stack_items_to_pass(self):
     """How many of the top stack items do we process."""
     return 0
 
-  def default_input_is_int(self):
-    """Whether the default inputs are ints or floats."""
-    return False
-
   # End of subclassing interface, all functions below are internal.
 
   def output_shape(self, input_shape_and_type, params):
@@ -93,7 +90,7 @@ def call_on_input(x, params):
           n = self.stack_items_to_pass() if isinstance(x, (list, tuple)) else 0
           return _apply_to_first_n(f, x, n)
         params_shapes = nested_map(
-            params, lambda x: ShapeType(shape=x.shape, tp=x.dtype))
+            params, lambda x: ShapeType(shape=x.shape, dtype=x.dtype))
         s = _eval_on_shapes(call_on_input, input_shape_and_type, params_shapes)
       return s
     except Exception:
@@ -101,8 +98,8 @@ def call_on_input(x, params):
       raise LayerError(name, 'output_shape', self._caller,
                        input_shape_and_type, trace)
 
-  def initialize(self, input_shape, rng):
-    """Initialize the layer given an input shape and rng.
+  def initialize(self, input_shape, input_dtype, rng):
+    """Initialize the layer given an input shape, dtype and rng.
 
     Returns new_parameters(input_shape, rng) on the first call and () on any
     subsequent call, as the layer is already initialized. This is used for
@@ -113,6 +110,7 @@ def initialize(self, input_shape, rng):
 
     Args:
       input_shape: a tuple representing the shape of the input.
+      input_dtype: numpy dtype of the input.
       rng: random number generator.
 
     Returns:
@@ -131,7 +129,7 @@ def initialize(self, input_shape, rng):
         input_shape = input_shape[:self.stack_items_to_pass()]
         if len(input_shape) == 1:
           input_shape = input_shape[0]
-      self._params = self.new_parameters(input_shape, rng)
+      self._params = self.new_parameters(input_shape, input_dtype, rng)
       return self._params
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback(skip=3)
@@ -158,12 +156,12 @@ def __call__(self, x, params=(), **kwargs):
 class ShapeType(object):
   """Store shape and type."""
 
-  def __init__(self, shape, tp):
+  def __init__(self, shape, dtype):
     self.shape = shape
-    self.tp = tp
+    self.dtype = dtype
 
   def __repr__(self):
-    return '[shape:' + str(self.shape) + ', type:' + str(self.tp) + ']'
+    return '[shape:' + str(self.shape) + ', dtype:' + str(self.dtype) + ']'
 
 
 class LayerError(Exception):
@@ -200,7 +198,7 @@ def abstractify(x):
     return jax.abstract_arrays.raise_to_shaped(jax.core.get_aval(x))
 
   def make_array(arg):
-    return backend.numpy.zeros(shape=arg.shape, dtype=arg.tp)
+    return backend.numpy.zeros(shape=arg.shape, dtype=arg.dtype)
 
   def turn_back_into_pytree(x):
     if isinstance(x, jax.core.JaxTuple):
@@ -335,7 +333,7 @@ def _short_traceback(skip=7):
 # Decorator for making layers from functions.
 
 
-def layer(new_parameters=None, stack_items_to_pass=1, input_is_int=False):
+def layer(new_parameters=None, stack_items_to_pass=1):
   """Create a layer class from a function."""
   def layer_decorator(call):
     """Decorating the call function."""
@@ -344,15 +342,11 @@ def stack_items_to_pass_fn(self):
       del self
       return stack_items_to_pass
 
-    def default_input_is_int_fn(self):
-      del self
-      return input_is_int
-
-    def new_parameters_fn(self, input_shape, rng):
+    def new_parameters_fn(self, input_shape, input_dtype, rng):
       if new_parameters is None:
         return ()
       kwargs = self._init_kwargs  # pylint: disable=protected-access
-      return new_parameters(input_shape, rng, **kwargs)
+      return new_parameters(input_shape, input_dtype, rng, **kwargs)
 
     def call_fn(self, x, params=(), **kwargs):
       """The call function of the created class, derived from call."""
@@ -370,7 +364,6 @@ def call_fn(self, x, params=(), **kwargs):
     # Create the class.
     cls = type(call.__name__, (Layer,),
                {'call': call_fn,
-                'default_input_is_int': default_input_is_int_fn,
                 'new_parameters': new_parameters_fn,
                 'stack_items_to_pass': stack_items_to_pass_fn})
 
@@ -417,10 +410,10 @@ def to_shape_and_type(x_shapes, integers):
     return {k: to_shape_and_type(x_shapes[k], integers) for k in x_shapes}
   if isinstance(x_shapes, onp.ndarray):  # Numpy array shape
     return ShapeType(shape=x_shapes.tolist(),
-                     tp=onp.int32 if integers else onp.float32)
+                     dtype=onp.int32 if integers else onp.float32)
   if isinstance(x_shapes[0], (int, onp.int32, onp.int64)):
     return ShapeType(shape=x_shapes,
-                     tp=onp.int32 if integers else onp.float32)
+                     dtype=onp.int32 if integers else onp.float32)
   if isinstance(x_shapes, list):  # Nested shape: list.
     return [to_shape_and_type(s, integers) for s in x_shapes]
   if isinstance(x_shapes, tuple):  # Nested shape: tuple.
@@ -431,8 +424,9 @@ def to_shape_and_type(x_shapes, integers):
 def check_shape_agreement(layer_instance, input_shape, integer_inputs=False):
   """Check if layer.output_shape agrees with the actual output shape."""
   rng1, rng2, rng3 = backend.random.split(backend.random.get_prng(0), 3)
-  params = layer_instance.initialize(input_shape, rng1)
   input_shape_and_type = to_shape_and_type(input_shape, integer_inputs)
+  input_dtype = nested_map(input_shape_and_type, lambda x: x.dtype)
+  params = layer_instance.initialize(input_shape, input_dtype, rng1)
   output_shape_and_type = layer_instance.output_shape(
       input_shape_and_type, params)
   output_shape = nested_map(output_shape_and_type, lambda x: x.shape)
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 9e6efb610..a00f86e8e 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -100,19 +100,18 @@ def output_shape(self, input_shape_and_type, params):
       cur_shape_and_type = layer.output_shape(cur_shape_and_type, param)
     return cur_shape_and_type
 
-  def default_input_is_int(self):
-    if self._nlayers == 0:
-      return False
-    return self._layers[0].default_input_is_int()
-
-  def new_parameters(self, input_shape, rng):
+  def new_parameters(self, input_shape, input_dtype, rng):
+    def MakeShapeType(shape, dtype):
+      if isinstance(dtype, (list, tuple)):
+        return tuple(MakeShapeType(s, t) for s, t in zip(shape, dtype))
+      return base.ShapeType(shape=shape, dtype=dtype)
     params = []
-    cur_shape_and_type = base.to_shape_and_type(
-        input_shape, self.default_input_is_int())
+    cur_shape_and_type = MakeShapeType(input_shape, input_dtype)
     for layer in self._layers:
       rng, layer_rng = backend.random.split(rng)
       cur_shape = base.nested_map(cur_shape_and_type, lambda x: x.shape)
-      param = layer.initialize(cur_shape, layer_rng)
+      cur_dtype = base.nested_map(cur_shape_and_type, lambda x: x.dtype)
+      param = layer.initialize(cur_shape, cur_dtype, layer_rng)
       pparam = layer._params   # pylint: disable=protected-access
       cur_shape_and_type = layer.output_shape(cur_shape_and_type, pparam)
       params.append(param)
@@ -231,7 +230,7 @@ def call(self, x, params=(), **kwargs):
       return x
     return base.nested_map(self._output, lambda i: self._map(x, i))
 
-  def new_parameters(self, input_shape, rng):
+  def new_parameters(self, input_shape, input_dtype, rng):
     return ()
 
 
@@ -266,13 +265,10 @@ def call(self, x, params=(), **kwargs):
              for layer, p, r in zip(self._layers, params, rngs)]
       return tuple(res)
 
-  def default_input_is_int(self):
-    return self._layers[0].default_input_is_int()
-
-  def new_parameters(self, input_shape, rng):
+  def new_parameters(self, input_shape, input_dtype, rng):
     rngs = backend.random.split(rng, self._nlayers)
     if not isinstance(self._layers, dict):
-      return [layer.initialize(input_shape, rng)
+      return [layer.initialize(input_shape, input_dtype, rng)
               for layer, rng in zip(self._layers, rngs)]
 
   def output_shape(self, input_shape, params):
@@ -389,9 +385,6 @@ def __init__(self, *layers, **kwlayers):
   def stack_items_to_pass(self):
     return self._nlayers
 
-  def default_input_is_int(self):
-    return any([layer.default_input_is_int() for layer in self._layers])
-
   def call(self, inputs, params=(), **kwargs):
     # Split the random number generators.
     rng = kwargs.pop('rng', None)
@@ -417,12 +410,12 @@ def call(self, inputs, params=(), **kwargs):
         result[k] = inputs[k]
     return result
 
-  def new_parameters(self, input_shape, rng):
+  def new_parameters(self, input_shape, input_dtype, rng):
     rngs = backend.random.split(rng, self._nlayers)
     # If the argument layers are a sequence, create parameters for each one.
     if not isinstance(self._layers, dict):
-      return [layer.initialize(shape, rng) for layer, shape, rng
-              in zip(self._layers, input_shape, rngs)]
+      return [layer.initialize(shape, dtype, rng) for layer, shape, dtype, rng
+              in zip(self._layers, input_shape, input_dtype, rngs)]
     # If the argument layers are a dictionary, create a dictionary too.
     result, counter = {}, 0
     for k in self._layers:
@@ -474,11 +467,11 @@ def call(self, inputs, params=(), **kwargs):
       return result
     return tuple(result)
 
-  def new_parameters(self, input_shape, rng):
+  def new_parameters(self, input_shape, input_dtype, rng):
     first_shape = input_shape[0]
     if self._check_shapes:
       for shape in input_shape:
         if shape != first_shape:
           raise ValueError('Map layer can only be applied to list of elements '
                            'with the same shapes. Shapes: %s' % str(shape))
-    return self._layer.initialize(first_shape, rng)
+    return self._layer.initialize(first_shape, input_dtype[0], rng)
diff --git a/tensor2tensor/trax/layers/convolution.py b/tensor2tensor/trax/layers/convolution.py
index 4e3b1866e..879f66ad7 100644
--- a/tensor2tensor/trax/layers/convolution.py
+++ b/tensor2tensor/trax/layers/convolution.py
@@ -81,7 +81,8 @@ def _kernel_shape(self, input_shape):
             input_shape[self._lhs_spec.index('C')] if c == 'I' else
             next(kernel_size_iter) for c in self._rhs_spec]
 
-  def new_parameters(self, input_shape, rng):
+  def new_parameters(self, input_shape, input_dtype, rng):
+    del input_dtype
     if len(input_shape) > 4:
       self._check_nhwc()
       new_batch_dim = six.moves.reduce(operator.mul, input_shape[:-3])
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index c98694a41..cb314a526 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import numpy as onp
+
 from tensor2tensor.trax import backend
 from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers import base
@@ -76,6 +78,11 @@ def Softplus(x, **unused_kwargs):
   return np.logaddexp(x, 0.)
 
 
+@base.layer()
+def ToFloat(x, **unused_kwargs):
+  return x.astype(onp.float32)
+
+
 class Dense(base.Layer):
   """Layer constructor function for a dense (fully-connected) layer."""
 
@@ -95,7 +102,8 @@ def call(self, x, params, **kwargs):
     w, b = params
     return np.dot(x, w) + b
 
-  def new_parameters(self, input_shape, rng):
+  def new_parameters(self, input_shape, input_dtype, rng):
+    del input_dtype
     rng1, rng2 = backend.random.split(rng, 2)
     w = self._kernel_initializer((input_shape[-1], self._units), rng1)
     b = self._bias_initializer((self._units,), rng2)
@@ -115,14 +123,12 @@ def __init__(self, d_feature, vocab_size,
   def stack_items_to_pass(self):
     return 1
 
-  def default_input_is_int(self):
-    return True
-
   def call(self, x, params, **kwargs):
     del kwargs
     return np.take(params, x, axis=0)
 
-  def new_parameters(self, input_shape, rng):
+  def new_parameters(self, input_shape, input_dtype, rng):
+    del input_dtype
     return self._kernel_initializer(
         (self._vocab_size, self._d_feature), rng)
 
diff --git a/tensor2tensor/trax/layers/core_test.py b/tensor2tensor/trax/layers/core_test.py
index dd0775376..0331e8440 100644
--- a/tensor2tensor/trax/layers/core_test.py
+++ b/tensor2tensor/trax/layers/core_test.py
@@ -69,8 +69,8 @@ def test_dense_param_sharing(self):
     layer = core.Dense(32)
     model2 = combinators.Serial(layer, layer)
     rng = backend.random.get_prng(0)
-    params1 = model1.initialize((1, 32), rng)
-    params2 = model2.initialize((1, 32), rng)
+    params1 = model1.initialize((1, 32), onp.float32, rng)
+    params2 = model2.initialize((1, 32), onp.float32, rng)
     # The first parameters have 2 kernels of size (32, 32).
     self.assertEqual((32, 32), params1[0][0].shape)
     self.assertEqual((32, 32), params1[1][0].shape)
diff --git a/tensor2tensor/trax/layers/normalization.py b/tensor2tensor/trax/layers/normalization.py
index deb8ef490..c42dd5f0f 100644
--- a/tensor2tensor/trax/layers/normalization.py
+++ b/tensor2tensor/trax/layers/normalization.py
@@ -24,10 +24,10 @@
 
 
 # Batch normalization.
-def BatchNormParams(input_shape, rng, axis=(0, 1, 2),
+def BatchNormParams(input_shape, input_dtype, rng, axis=(0, 1, 2),
                     center=True, scale=True, **kwargs):
   """Helper to initialize batch norm params."""
-  del rng, kwargs
+  del input_dtype, rng, kwargs
   axis = (axis,) if np.isscalar(axis) else axis
   shape = tuple(d for i, d in enumerate(input_shape) if i not in axis)
   beta = np.zeros(shape, dtype='float32') if center else ()
@@ -64,9 +64,9 @@ def BatchNorm(x, params, axis=(0, 1, 2), epsilon=1e-5,
 
 
 # Layer normalization.
-def LayerNormParams(input_shape, rng, epsilon=1e-6):
+def LayerNormParams(input_shape, input_dtype, rng, epsilon=1e-6):
   """Helper: create layer norm parameters."""
-  del rng, epsilon
+  del input_dtype, rng, epsilon
   features = input_shape[-1]
   scale = np.ones(features)
   bias = np.zeros(features)
diff --git a/tensor2tensor/trax/models/atari_cnn.py b/tensor2tensor/trax/models/atari_cnn.py
index 1c88b1e9e..29cb3aaa1 100644
--- a/tensor2tensor/trax/models/atari_cnn.py
+++ b/tensor2tensor/trax/models/atari_cnn.py
@@ -26,6 +26,7 @@ def AtariCnn(hidden_sizes=(32, 32), output_size=128):
   """An Atari CNN."""
   # Input's shape = (B, T, H, W, C)
   return tl.Model(
+      tl.ToFloat(),
       tl.Div(divisor=255.0),
       # Have 4 copies of the input, each one shifted to the right by one.
       tl.Branch(
diff --git a/tensor2tensor/trax/models/atari_cnn_test.py b/tensor2tensor/trax/models/atari_cnn_test.py
index 0b1cb21d1..e0bd2e730 100644
--- a/tensor2tensor/trax/models/atari_cnn_test.py
+++ b/tensor2tensor/trax/models/atari_cnn_test.py
@@ -37,7 +37,7 @@ def test_computes(self):
         hidden_sizes=hidden_size, output_size=output_size)
     B, T, OBS = 2, 2, (28, 28, 3)  # pylint: disable=invalid-name
     rng_key, key = jax_random.split(rng_key)
-    params = model.initialize((1, 1) + OBS, key)
+    params = model.initialize((1, 1) + OBS, onp.float32, key)
     x = onp.arange(B * (T + 1) * functools.reduce(op.mul, OBS)).reshape(
         B, T + 1, *OBS)
     rng_key, key = jax_random.split(rng_key)
diff --git a/tensor2tensor/trax/models/research/chunked_transformer.py b/tensor2tensor/trax/models/research/chunked_transformer.py
index d819ae40a..491728cd6 100644
--- a/tensor2tensor/trax/models/research/chunked_transformer.py
+++ b/tensor2tensor/trax/models/research/chunked_transformer.py
@@ -25,9 +25,10 @@
 
 
 # Chunked positional encoding.
-def _chunked_positional_encoding_new_params(input_shape, rng, max_len=2048):  # pylint: disable=invalid-name
+def _chunked_positional_encoding_new_params(  # pylint: disable=invalid-name
+    input_shape, input_dtype, rng, max_len=2048):
   """Helper: create positional encoding parameters."""
-  del rng
+  del input_dtype, rng
   # Check if we are operating on chunked inputs by checking if the first
   # shape is a list/tuple of shapes (otherwise it's an int or numpy array).
   is_chunked = isinstance(input_shape[0], (list, tuple))
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index 8062275ac..7ed36508d 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -81,6 +81,7 @@ def Resnet50(d_hidden=64, n_output_classes=1001, mode='train'):
   """
   del mode
   return tl.Model(
+      tl.ToFloat(),
       tl.Conv(d_hidden, (7, 7), (2, 2), 'SAME'),
       tl.BatchNorm(),
       tl.Relu(),
@@ -146,6 +147,7 @@ def WideResnet(n_blocks=3, d_hidden=64, n_output_classes=10,
   """
   del mode
   return tl.Model(
+      tl.ToFloat(),
       tl.Conv(d_hidden, (3, 3), padding='SAME'),
       WideResnetGroup(n_blocks, d_hidden),
       WideResnetGroup(n_blocks, d_hidden * 2, (2, 2)),
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index fc7be2670..6e4c0fc0d 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -83,6 +83,7 @@
 
 def policy_and_value_net(rng_key,
                          batch_observations_shape,
+                         observations_dtype,
                          n_actions,
                          bottom_layers_fn=(),
                          two_towers=True):
@@ -107,7 +108,8 @@ def policy_and_value_net(rng_key,
             [tl.Dense(1)]
         )
     )
-  return net.initialize(batch_observations_shape, rng_key), net
+  params = net.initialize(batch_observations_shape, observations_dtype, rng_key)
+  return params, net
 
 
 def optimizer_fn(net_params, step_size=1e-3):
@@ -813,6 +815,7 @@ def training_loop(
   # Batch Observations Shape = [1, 1] + OBS, because we will eventually call
   # policy and value networks on shape [B, T] +_OBS
   batch_observations_shape = (1, 1) + env.observation_space.shape
+  observations_dtype = env.observation_space.dtype
 
   assert isinstance(env.action_space, gym.spaces.Discrete)
   n_actions = env.action_space.n
@@ -821,7 +824,8 @@ def training_loop(
 
   # Initialize the policy and value network.
   policy_and_value_net_params, policy_and_value_net_apply = (
-      policy_and_value_net_fn(key1, batch_observations_shape, n_actions))
+      policy_and_value_net_fn(key1, batch_observations_shape,
+                              observations_dtype, n_actions))
 
   # Maybe restore the policy params. If there is nothing to restore, then
   # iteration = 0 and policy_and_value_net_params are returned as is.
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index 95533a617..c8a879d70 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -38,7 +38,7 @@ def test_policy_and_value_net(self):
     batch_observation_shape = (1, 1) + observation_shape
     n_actions = 2
     pnv_params, pnv_apply = ppo.policy_and_value_net(
-        self.rng_key, batch_observation_shape, n_actions,
+        self.rng_key, batch_observation_shape, np.float32, n_actions,
         lambda: [layers.Flatten(num_axis_to_keep=2)])
     batch = 2
     time_steps = 10
@@ -383,11 +383,11 @@ def test_combined_loss(self):
     batch_observation_shape = (1, 1) + OBS
 
     old_params, _ = ppo.policy_and_value_net(
-        key1, batch_observation_shape, A,
+        key1, batch_observation_shape, np.float32, A,
         lambda: [layers.Flatten(num_axis_to_keep=2)])
 
     new_params, net_apply = ppo.policy_and_value_net(
-        key2, batch_observation_shape, A,
+        key2, batch_observation_shape, np.float32, A,
         lambda: [layers.Flatten(num_axis_to_keep=2)])
 
     # Generate a batch of observations.
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 8e6ac36ad..10b14a777 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -468,7 +468,8 @@ def train(output_dir,
     params = state.params[0]
     opt_state = state.params
   else:
-    params = model_train.initialize(model_input_shape, init_rng)
+    params = model_train.initialize(
+        model_input_shape, inputs.input_dtype, init_rng)
     opt_state = (params, opt.tree_init(params))
   if n_devices > 1:
     replicate = lambda x: numpy.broadcast_to(x, (n_devices,) + x.shape)
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index 71e24ea47..1f8ee2504 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -51,7 +51,8 @@ def input_stream():
       train_stream=input_stream,
       train_eval_stream=input_stream,
       eval_stream=input_stream,
-      input_shape=input_shape)
+      input_shape=input_shape,
+      input_dtype=np.float32)
 
 
 class TraxTest(test.TestCase):

From df4a50be8db63d36bb2391e1adc04558c0e8e1c9 Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Wed, 5 Jun 2019 16:57:36 -0700
Subject: [PATCH 2099/2720] Infer hparams_path from checkpoint_path (if
 specified) in serving/export.py

PiperOrigin-RevId: 251748055
---
 tensor2tensor/serving/export.py | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 8d103a39a..55bcbeec5 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -54,11 +54,25 @@ def _get_hparams_path():
   hparams_path = None
   if FLAGS.output_dir:
     hparams_path = os.path.join(FLAGS.output_dir, "hparams.json")
-  else:
+  elif FLAGS.checkpoint_path:  # Infer hparams.json from checkpoint path
+    hparams_path = os.path.join(
+        os.path.dirname(FLAGS.checkpoint_path), "hparams.json")
+
+  # Check if hparams_path really exists
+  if hparams_path:
+    if tf.gfile.Exists(hparams_path):
+      tf.logging.info("hparams file %s exists", hparams_path)
+    else:
+      tf.logging.info("hparams file %s does not exist", hparams_path)
+      hparams_path = None
+
+  # Can't find hparams_path
+  if not hparams_path:
     tf.logging.warning(
-        "--output_dir not specified. Hyper-parameters will be infered from"
-        "--hparams_set and --hparams only. These may not match training time"
-        "hyper-parameters.")
+        "--output_dir not specified or file hparams.json does not exists. "
+        "Hyper-parameters will be infered from --hparams_set and "
+        "--hparams only. These may not match training time hyper-parameters.")
+
   return hparams_path
 
 

From 21f2494c69aa5d2ec151f1f9b14bd4cfc8cf8215 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Thu, 6 Jun 2019 21:05:04 +0200
Subject: [PATCH 2100/2720] Params for mbrl with dqn. (#1592)

---
 tensor2tensor/models/research/rl.py           | 28 +++++++++++++++++++
 .../rl/trainer_model_based_params.py          | 28 +++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 039d8809f..02f7dd2d1 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -408,6 +408,34 @@ def dqn_original_params():
   return hparams
 
 
+@registry.register_hparams
+def dqn_guess1_params():
+  """Guess 1 for DQN params."""
+  hparams = dqn_atari_base()
+  hparams.set_hparam("num_frames", int(1e6))
+  hparams.set_hparam("agent_update_period", 1)
+  hparams.set_hparam("agent_target_update_period", 400)
+  # Small replay buffer size was set for mistake, but it seems to work
+  hparams.set_hparam("replay_buffer_replay_capacity", 10000)
+  return hparams
+
+
+@registry.register_hparams
+def dqn_2m_replay_buffer_params():
+  """Guess 1 for DQN params, 2 milions transitions in replay buffer"""
+  hparams = dqn_guess1_params()
+  hparams.set_hparam("replay_buffer_replay_capacity", int(2e6) + int(1e5))
+  return hparams
+
+
+@registry.register_hparams
+def dqn_10m_replay_buffer_params():
+  """Guess 1 for DQN params, 10 milions transitions in replay buffer"""
+  hparams = dqn_guess1_params()
+  hparams.set_hparam("replay_buffer_replay_capacity", int(10e6))
+  return hparams
+
+
 def rlmf_tiny_overrides():
   """Parameters to override for tiny setting excluding agent-related hparams."""
   return dict(
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index fc50108d8..64529856c 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -210,6 +210,33 @@ def rlmb_dqn_base():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_dqn_guess1():
+  """rlmb_dqn guess1 params"""
+  hparams = rlmb_dqn_base()
+  hparams.set_hparam("base_algo_params", "dqn_guess1_params")
+  # At the moment no other option for evaluation, so we want long rollouts to
+  # not bias scores.
+  hparams.set_hparam("eval_rl_env_max_episode_steps", 5000)
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_dqn_guess1_2m_replay_buffer():
+  """rlmb_dqn guess1 params"""
+  hparams = rlmb_dqn_guess1()
+  hparams.set_hparam("base_algo_params", "dqn_2m_replay_buffer_params")
+  return hparams
+
+
+@registry.register_hparams
+def rlmb_dqn_guess1_10m_replay_buffer():
+  """rlmb_dqn guess1 params"""
+  hparams = rlmb_dqn_guess1()
+  hparams.set_hparam("base_algo_params", "dqn_10m_replay_buffer_params")
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_basetest():
   """Base setting but quicker with only 2 epochs."""
@@ -617,6 +644,7 @@ def rlmb_dqn_tiny():
   hparams = rlmb_dqn_base()
   hparams = hparams.override_from_dict(_rlmb_tiny_overrides())
   update_hparams(hparams, dict(
+      base_algo_params="dqn_guess1_params",
       simulated_rollout_length=2,
       dqn_time_limit=2,
       dqn_num_frames=128,

From fe77e4b8d6c5a522a4b365d5b66c0411b742db9d Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Thu, 6 Jun 2019 12:42:20 -0700
Subject: [PATCH 2101/2720] Merge of PR #1592

PiperOrigin-RevId: 251905735
---
 tensor2tensor/models/research/rl.py            | 4 ++--
 tensor2tensor/rl/trainer_model_based_params.py | 6 +++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 02f7dd2d1..4cb2bf161 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -422,7 +422,7 @@ def dqn_guess1_params():
 
 @registry.register_hparams
 def dqn_2m_replay_buffer_params():
-  """Guess 1 for DQN params, 2 milions transitions in replay buffer"""
+  """Guess 1 for DQN params, 2 milions transitions in replay buffer."""
   hparams = dqn_guess1_params()
   hparams.set_hparam("replay_buffer_replay_capacity", int(2e6) + int(1e5))
   return hparams
@@ -430,7 +430,7 @@ def dqn_2m_replay_buffer_params():
 
 @registry.register_hparams
 def dqn_10m_replay_buffer_params():
-  """Guess 1 for DQN params, 10 milions transitions in replay buffer"""
+  """Guess 1 for DQN params, 10 milions transitions in replay buffer."""
   hparams = dqn_guess1_params()
   hparams.set_hparam("replay_buffer_replay_capacity", int(10e6))
   return hparams
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 64529856c..4c5350fed 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -212,7 +212,7 @@ def rlmb_dqn_base():
 
 @registry.register_hparams
 def rlmb_dqn_guess1():
-  """rlmb_dqn guess1 params"""
+  """DQN guess1 params."""
   hparams = rlmb_dqn_base()
   hparams.set_hparam("base_algo_params", "dqn_guess1_params")
   # At the moment no other option for evaluation, so we want long rollouts to
@@ -223,7 +223,7 @@ def rlmb_dqn_guess1():
 
 @registry.register_hparams
 def rlmb_dqn_guess1_2m_replay_buffer():
-  """rlmb_dqn guess1 params"""
+  """DQN guess1 params, 2M replay buffer."""
   hparams = rlmb_dqn_guess1()
   hparams.set_hparam("base_algo_params", "dqn_2m_replay_buffer_params")
   return hparams
@@ -231,7 +231,7 @@ def rlmb_dqn_guess1_2m_replay_buffer():
 
 @registry.register_hparams
 def rlmb_dqn_guess1_10m_replay_buffer():
-  """rlmb_dqn guess1 params"""
+  """DQN guess1 params, 10M replay buffer."""
   hparams = rlmb_dqn_guess1()
   hparams.set_hparam("base_algo_params", "dqn_10m_replay_buffer_params")
   return hparams

From c668e4bc92a8a9600903cc9b0ee150f931ac43d2 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 6 Jun 2019 16:22:15 -0700
Subject: [PATCH 2102/2720] Remove experimental_export_device_assignment from
 TPUEstimator.export_savedmodel(), so as to remove rewrite_for_inference().

As a replacement, export_savedmodel() V2 API supports device_assignment where user call tpu.rewrite in model_fn and pass in device_assigment there.

PiperOrigin-RevId: 251950039
---
 tensor2tensor/utils/trainer_lib.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index f6016cfbb..7dc6f2448 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -314,8 +314,7 @@ def create_estimator(model_name,
         use_tpu=use_tpu,
         train_batch_size=batch_size,
         eval_batch_size=batch_size if "eval" in schedule else None,
-        predict_batch_size=predict_batch_size,
-        experimental_export_device_assignment=True)
+        predict_batch_size=predict_batch_size)
   else:
     estimator = tf.estimator.Estimator(
         model_fn=model_fn,

From 0487781a660815146d9ab0ec87ce35843b9c0db7 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 6 Jun 2019 22:25:23 -0700
Subject: [PATCH 2103/2720] Fix bug in policy evaluation, where we were only
 taking 1 trajectory (rather than batch size number of trajectories)  - This
 basically implies that we took the trajectory that completed the quickest
 (probably under-reporting how good the policy is, also probably not by much).
  - Also leads to a lot of wasted cycles by silently dropping the
 trajectories.

PiperOrigin-RevId: 251991424
---
 tensor2tensor/trax/rlax/ppo.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 6e4c0fc0d..f3d3068d1 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -727,6 +727,7 @@ def evaluate_policy(eval_env,
       trajs, _ = env_problem_utils.play_env_problem_with_policy(
           eval_env,
           get_predictions,
+          num_trajectories=eval_env.batch_size,
           boundary=boundary,
           max_timestep=max_timestep,
           reset=True,

From eb9b0f2096b8ce561d5d0e051839e201eefbf14e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 7 Jun 2019 09:17:17 -0700
Subject: [PATCH 2104/2720] Adding a new class wit the proper wmt18 training
 set including rapid data as defined at
 http://www.statmt.org/wmt18/translation-task.html

PiperOrigin-RevId: 252061140
---
 .../data_generators/translate_ende.py         | 25 +++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index 88b8ded8e..1247d0814 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -41,12 +41,23 @@
         ("training/europarl-v7.de-en.en", "training/europarl-v7.de-en.de")
     ],
 ]
+
 _ENDE_EVAL_DATASETS = [
     [
         "http://data.statmt.org/wmt17/translation-task/dev.tgz",
         ("dev/newstest2013.en", "dev/newstest2013.de")
     ],
 ]
+
+_ENDE_RAPID_TRAIN_DATASET = [
+    # additional training data available for WMT 18 news task training data
+    # as defined by http://www.statmt.org/wmt18/translation-task.html
+    [
+        "http://data.statmt.org/wmt18/translation-task/rapid2016.tgz",
+        ("rapid2016.de-en.en", "rapid2016.de-en.de"),
+    ],
+]
+
 _ENDE_PARACRAWL_DATASETS = [
     [
         "https://s3.amazonaws.com/web-language-models/paracrawl/release4/en-de.bicleaner07.tmx.gz",  # pylint: disable=line-too-long
@@ -70,6 +81,20 @@ def source_data_files(self, dataset_split):
     return train_datasets if train else _ENDE_EVAL_DATASETS
 
 
+@registry.register_problem
+class TranslateEnde2018Wmt32k(translate.TranslateProblem):
+  """En-de translation trained on WMT18 corpus."""
+
+  @property
+  def use_vocab_from_other_problem(self):
+    return TranslateEndeWmt32k()
+
+  @property
+  def additional_training_datasets(self):
+    """WMT18 adds rapid data."""
+    return _ENDE_RAPID_TRAIN_DATASET
+
+
 @registry.register_problem
 class TranslateEndeWmtClean32k(TranslateEndeWmt32k):
   """En-de translation trained on WMT with further cleaning."""

From b7d192425acad790e30b93eb1e5f06025f85f5e4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 7 Jun 2019 10:34:05 -0700
Subject: [PATCH 2105/2720] Internal

PiperOrigin-RevId: 252075275
---
 tensor2tensor/layers/common_layers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 3670f7134..5e95fed62 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1961,7 +1961,7 @@ def padded_cross_entropy_mixture(logits,
 
   # if we need to compute the best logits
   if return_best_logits:
-    if supervised_mode:
+    if not supervised_mode:
       return_mixture_indices = tf.squeeze(
           tf.cast(tf.argmin(xent, 0), dtype=tf.int32), axis=[1, 2])
     else:

From 6c7f2ab4b45ca0f2082a4b14b48cf7d4b45cbbea Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 7 Jun 2019 11:30:46 -0700
Subject: [PATCH 2106/2720] Construct observations that only go back `n`
 time-steps in time to give to the policy.  - This implicitly happens in
 training where n=128 because of truncation. But in eval n can be upto 20,000.
  - We instead introduce another parameter `len_history_for_policy`, i.e. how
 many last observations to take into account to give to the policy function
 and use that everywhere.

Also keep the last `n` observations when a trajectory is truncated
 - This currently doesn't happen even for training, so at trajectory boundaries we are currently missing context.
 - Right now n=1, so this is what is there currently, i.e. the default.

Keep evals limited to categorical and gumbel sampling.
 - Basically removes epsilon-greedy, this is 1/3rd of the eval time and categorical and gumbel seem the most useful anyways.

In policy application, bucket on `len_history_for_policy` rather than on boundary.
 - NOTE: In policy optimization, we still cut on boundary chunks. Since this happens downstream in ppo.py's pad_trajectories

PiperOrigin-RevId: 252086911
---
 tensor2tensor/envs/env_problem.py            |   4 +-
 tensor2tensor/envs/env_problem_utils.py      |  57 ++++----
 tensor2tensor/envs/env_problem_utils_test.py |  26 ++--
 tensor2tensor/envs/trajectory.py             |  66 +++++++---
 tensor2tensor/envs/trajectory_test.py        | 129 ++++++++++++++-----
 tensor2tensor/trax/rlax/ppo.py               |  60 +++++----
 tensor2tensor/trax/rlax/ppo_main.py          |   3 +
 7 files changed, 228 insertions(+), 117 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index b3bbdaaa0..2b8c63bff 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -471,13 +471,13 @@ def _reset(self, indices):
     # rest being the dimensionality of the observation.
     return np.stack([self._envs[index].reset() for index in indices])
 
-  def truncate(self, indices=None):
+  def truncate(self, indices=None, num_to_keep=1):
     """Truncates trajectories at the specified indices."""
 
     if indices is None:
       indices = np.arange(self.batch_size)
 
-    self.trajectories.truncate_trajectories(indices)
+    self.trajectories.truncate_trajectories(indices, num_to_keep=num_to_keep)
 
   def reset(self, indices=None):
     """Resets environments at given indices.
diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 8f2ffed47..3c7952480 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import time
 import numpy as np
 
 CATEGORICAL_SAMPLING = "categorical"
@@ -54,38 +55,39 @@ def play_env_problem_with_policy(env,
                                  policy_fun,
                                  num_trajectories=1,
                                  max_timestep=None,
-                                 boundary=20,
                                  reset=True,
                                  rng=None,
                                  policy_sampling=CATEGORICAL_SAMPLING,
                                  temperature=0.5,
-                                 eps=0.1):
+                                 eps=0.1,
+                                 len_history_for_policy=32,
+                                 num_to_keep=1):
   """Plays the given env with the policy function to collect trajectories.
 
   Args:
     env: environment object, should be a subclass of env_problem.EnvProblem.
     policy_fun: callable, taking in observations((B, T) + OBS) and returning
-        back log-probabilities (B, T, A).
+      back log-probabilities (B, T, A).
     num_trajectories: int, number of trajectories to collect.
     max_timestep: int or None, if not None or a negative number, we cut any
-        trajectory that exceeds this time put it in the completed bin, and
-        *dont* reset the env.
-    boundary: this is the bucket length, we pad the observations to integer
-        multiples of this + 1 and then feed the padded observations to the
-        policy_fun.
+      trajectory that exceeds this time put it in the completed bin, and *dont*
+      reset the env.
     reset: bool, true if we want to reset the envs. The envs are also reset if
-        max_max_timestep is None or < 0
+      max_max_timestep is None or < 0
     rng: jax rng, splittable.
     policy_sampling: string, how to select an action given a policy, one of:
-        CATEGORICAL_SAMPLING, GREEDY, GUMBEL_SAMPLING
+      CATEGORICAL_SAMPLING, GREEDY, GUMBEL_SAMPLING
     temperature: float, temperature used in gumbel sampling.
     eps: float, epsilon to use in epsilon greedy.
-
+    len_history_for_policy: int, the maximum history to keep for applying the
+      policy on. We also bucket observations on this number.
+    num_to_keep: int, while truncating trajectory how many time-steps to keep.
 
   Returns:
     A tuple, (trajectories, number of completed trajectories). Where
     trajectories is a list of triples of (observation, action, reward) ndarrays.
   """
+  t0 = time.time()
 
   def categorical_sample(log_probs):
     """Categorical sampling."""
@@ -144,26 +146,31 @@ def epsilon_greedy(log_probs):
 
   num_done_trajectories = 0
 
+  policy_application_total_time = 0
   while env.trajectories.num_completed_trajectories < num_trajectories:
     # Get all the observations for all the active trajectories.
     # Shape is (B, T) + OBS
-    padded_observations = env.trajectories.observations_np(boundary=boundary)
-    lengths = env.trajectories.trajectory_lengths
+    # Bucket on whatever length is needed.
+    padded_observations, lengths = env.trajectories.observations_np(
+        boundary=len_history_for_policy,
+        len_history_for_policy=len_history_for_policy)
 
     B, T = padded_observations.shape[:2]  # pylint: disable=invalid-name
 
     assert B == env.batch_size
     assert (B,) == lengths.shape
 
+    t1 = time.time()
     log_prob_actions, _, rng = policy_fun(padded_observations, rng=rng)
+    policy_application_total_time += (time.time() - t1)
+
     assert (B, T) == log_prob_actions.shape[:2]
     A = log_prob_actions.shape[2]  # pylint: disable=invalid-name
 
     # We need the log_probs of those actions that correspond to the last actual
     # time-step.
     index = lengths - 1  # Since we want to index using lengths.
-    log_probs = log_prob_actions[np.arange(B)[:, None],
-                                 index[:, None],
+    log_probs = log_prob_actions[np.arange(B)[:, None], index[:, None],
                                  np.arange(A)]
     assert (B, A) == log_probs.shape, \
         "B=%d, A=%d, log_probs.shape=%s" % (B, A, log_probs.shape)
@@ -192,10 +199,6 @@ def epsilon_greedy(log_probs):
     if done_idxs.size:
       env.reset(indices=done_idxs)
 
-    # Do we have enough trajectories right now?
-    if env.trajectories.num_completed_trajectories >= num_trajectories:
-      break
-
     if max_timestep is None or max_timestep < 1:
       continue
 
@@ -207,11 +210,7 @@ def epsilon_greedy(log_probs):
     if exceeded_time_limit_idxs.size:
       # This just cuts the trajectory, doesn't reset the env, so it continues
       # from where it left off.
-      env.truncate(indices=exceeded_time_limit_idxs)
-
-    # Do we have enough trajectories right now?
-    if env.trajectories.num_completed_trajectories >= num_trajectories:
-      break
+      env.truncate(indices=exceeded_time_limit_idxs, num_to_keep=num_to_keep)
 
   # We have the trajectories we need, return a list of triples:
   # (observations, actions, rewards)
@@ -219,7 +218,11 @@ def epsilon_greedy(log_probs):
   for trajectory in env.trajectories.completed_trajectories[:num_trajectories]:
     completed_trajectories.append(trajectory.as_numpy)
 
-  # Keep the rest of the trajectories, if any, in our kitty.
-  env.trajectories.clear_completed_trajectories(num=num_trajectories)
+  policy_application_time = round(1000 * policy_application_total_time, 2)
+  misc_time = round(1000 * (time.time() - t0), 2) - policy_application_time
+  timing_info = {
+      "trajectory_collection/policy_application": policy_application_time,
+      "trajectory_collection/misc": misc_time,
+  }
 
-  return completed_trajectories, num_done_trajectories
+  return completed_trajectories, num_done_trajectories, timing_info
diff --git a/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index 4885be940..8eb6d6501 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -48,12 +48,19 @@ def test_play_env_problem_randomly(self):
 
   def test_play_env_problem_with_policy(self):
     env = env_problem.EnvProblem(
-        base_env_name="CartPole-v0",
-        batch_size=2,
-        reward_range=(-1, 1))
+        base_env_name="CartPole-v0", batch_size=2, reward_range=(-1, 1))
+
+    # Let's make sure that at-most 4 observations come to the policy function.
+    len_history_for_policy = 4
 
     def policy_fun(observations, rng=None):
       b, t = observations.shape[:2]
+      # Assert that observations from time-step len_history_for_policy onwards
+      # are zeros.
+      self.assertTrue(
+          np.all(observations[:, len_history_for_policy:, ...] == 0))
+      self.assertFalse(
+          np.all(observations[:, :len_history_for_policy, ...] == 0))
       a = env.action_space.n
       p = np.random.uniform(size=(b, t, a))
       p = np.exp(p)
@@ -62,22 +69,25 @@ def policy_fun(observations, rng=None):
 
     max_timestep = 15
     num_trajectories = 2
-    trajectories, _ = env_problem_utils.play_env_problem_with_policy(
-        env, policy_fun, num_trajectories=num_trajectories,
-        max_timestep=max_timestep, boundary=20)
+    trajectories, _, _ = env_problem_utils.play_env_problem_with_policy(
+        env,
+        policy_fun,
+        num_trajectories=num_trajectories,
+        max_timestep=max_timestep,
+        len_history_for_policy=len_history_for_policy)
 
     self.assertEqual(num_trajectories, len(trajectories))
 
     # Check shapes within trajectories.
     traj = trajectories[0]
     T = traj[1].shape[0]  # pylint: disable=invalid-name
-    self.assertEqual((T+1, 4), traj[0].shape)  # (4,) is OBS
+    self.assertEqual((T + 1, 4), traj[0].shape)  # (4,) is OBS
     self.assertEqual((T,), traj[2].shape)
     self.assertLessEqual(T, max_timestep)
 
     traj = trajectories[1]
     T = traj[1].shape[0]  # pylint: disable=invalid-name
-    self.assertEqual((T+1, 4), traj[0].shape)
+    self.assertEqual((T + 1, 4), traj[0].shape)
     self.assertEqual((T,), traj[2].shape)
     self.assertLessEqual(T, max_timestep)
 
diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index 3c412e499..da04a8744 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -30,9 +30,12 @@
 class Trajectory(object):
   """Basically a list of TimeSteps with convenience methods."""
 
-  def __init__(self):
+  def __init__(self, time_steps=None):
     # Contains a list of time steps.
-    self._time_steps = []
+    if time_steps is None:
+      self._time_steps = []
+    else:
+      self._time_steps = time_steps
 
   def __str__(self):
     if not self.time_steps:
@@ -58,6 +61,22 @@ def change_last_time_step(self, **replace_time_step_kwargs):
     self._time_steps[-1] = self._time_steps[-1].replace(
         **replace_time_step_kwargs)
 
+  def truncate(self, num_to_keep=1):
+    """Truncate trajectories, keeping the last `num_to_keep` time-steps."""
+
+    # We return `ts_copy` back to the truncator.
+    ts_copy = self._time_steps[:]
+
+    # We keep the last few observations.
+    self._time_steps = self._time_steps[-num_to_keep:]
+
+    # NOTE: We will need to set the rewards to 0, to eliminate double counting.
+    for i in range(self.num_time_steps):
+      self._time_steps[i] = self._time_steps[i].replace(
+          raw_reward=0, processed_reward=0)
+
+    return Trajectory(time_steps=ts_copy)
+
   @property
   def last_time_step(self):
     # Pre-conditions: self._time_steps shouldn't be empty.
@@ -97,6 +116,11 @@ def reward(self):
   def observations_np(self):
     return np.stack([ts.observation for ts in self.time_steps])
 
+  def last_n_observations_np(self, n=None):
+    if n is not None:
+      n = -n  # pylint: disable=invalid-unary-operand-type
+    return np.stack([ts.observation for ts in self.time_steps[n:]])
+
   @property
   def actions_np(self):
     # The last action is None, so let's skip it.
@@ -170,7 +194,7 @@ def _complete_trajectory(self, trajectory, index):
     # Make a new one to replace it.
     self._trajectories[index] = Trajectory()
 
-  def truncate_trajectories(self, indices):
+  def truncate_trajectories(self, indices, num_to_keep=1):
     """Truncate trajectories at specified indices.
 
      This puts the truncated trajectories in the completed list and makes new
@@ -179,21 +203,19 @@ def truncate_trajectories(self, indices):
 
     Args:
         indices: iterable with the indices to truncate.
+        num_to_keep: int, number of last time-steps to keep while truncating.
     """
-    observations = []
     for index in indices:
       trajectory = self._trajectories[index]
       assert trajectory.is_active, "Trajectory to truncate can't be inactive."
 
-      # NOTE: We don't mark the last time-step as done.
+      # Now `trajectory` just consists of the last `num_to_keep` observations
+      # and actions. Rewards are zeroed out.
+      # The old data is placed in `old_trajectory`.
+      old_trajectory = trajectory.truncate(num_to_keep=num_to_keep)
 
-      # Collect the observations.
-      observations.append(trajectory.last_time_step.observation)
-
-    # Call reset on these indices, this will make new trajectories with the same
-    # observation as the existing ones, but in new trajectories. The existing
-    # trajectories are marked as completed.
-    self.reset(indices, np.stack(observations))
+      # We put the old data in _completed_trajectories.
+      self._completed_trajectories.append(old_trajectory)
 
   def reset(self, indices, observations):
     """Resets trajectories at given indices and populates observations.
@@ -344,21 +366,30 @@ def num_completed_trajectories(self):
 
   # TODO(afrozm): Take in an already padded observation ndarray and just append
   # the last time-step and adding more padding if needed.
-  def observations_np(self, boundary=20):
+  def observations_np(self, boundary=20, len_history_for_policy=20):
     """Pads the observations in all the trajectories and returns them.
 
     Args:
       boundary: integer, Observations will be padded to (n * boundary) + 1 where
-          n is an integer.
+        n is an integer.
+      len_history_for_policy: int, For each trajectory return only the last
+        `len_history_for_policy` observations. Set to None for all the
+        observations.
 
     Returns:
       padded_observations: (self.batch_size, n * boundary + 1) + OBS
     """
-    list_observations_np_ts = [t.observations_np for t in self.trajectories]
+    list_observations_np_ts = [
+        t.last_n_observations_np(n=len_history_for_policy)
+        for t in self.trajectories
+    ]
     # Every element in `list_observations_np_ts` is shaped (t,) + OBS
     OBS = list_observations_np_ts[0].shape[1:]  # pylint: disable=invalid-name
 
-    t_max = max(self.trajectory_lengths)
+    trajectory_lengths = np.stack(
+        [obs.shape[0] for obs in list_observations_np_ts])
+
+    t_max = max(trajectory_lengths)
     # t_max is rounded to the next multiple of `boundary`
     boundary = int(boundary)
     bucket_length = boundary * int(np.ceil(float(t_max) / boundary))
@@ -370,4 +401,5 @@ def padding_config(obs):
 
     return np.stack([
         np.pad(obs, padding_config(obs), "constant")
-        for obs in list_observations_np_ts])
+        for obs in list_observations_np_ts
+    ]), trajectory_lengths
diff --git a/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index b4d2083d3..b9a7f0bfc 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -99,6 +99,35 @@ def test_observation_np(self):
 
     self.assertEqual((ts,) + shape, t.observations_np.shape)
 
+  def test_truncate_and_last_n_observations_np(self):
+    t = trajectory.Trajectory()
+    ts = 5
+    shape = (3, 4)
+    for _ in range(ts):
+      t.add_time_step(observation=np.random.uniform(size=shape), done=False)
+
+    original_obs = np.copy(t.observations_np)
+    self.assertEqual((ts,) + shape, original_obs.shape)
+
+    # Now let's just get the observations from the last 2 steps.
+    num_to_keep = 2
+    truncated_original_obs = original_obs[-num_to_keep:, ...]
+
+    # Let's get the last `num_to_keep` observations
+    last_n_observations_np = np.copy(t.last_n_observations_np(n=num_to_keep))
+
+    # Now truncate the trajectory and get the same.
+    _ = t.truncate(num_to_keep=num_to_keep)
+    truncated_np = np.copy(t.observations_np)
+
+    # These should be the expected length.
+    self.assertEqual((2,) + shape, last_n_observations_np.shape)
+    self.assertEqual((2,) + shape, truncated_np.shape)
+
+    # Test the last `num_to_keep` are the same.
+    self.assertAllEqual(truncated_np, truncated_original_obs)
+    self.assertAllEqual(last_n_observations_np, truncated_original_obs)
+
   def test_as_numpy(self):
     t = trajectory.Trajectory()
     shape = (3, 4)
@@ -107,24 +136,24 @@ def test_as_numpy(self):
     ts = 5
     num_actions = 6
     observations = np.random.uniform(size=(ts,) + shape)
-    actions = np.random.choice(range(num_actions), size=(ts-1,))
-    rewards = np.random.choice([-1, 0, 1], size=(ts-1,))
+    actions = np.random.choice(range(num_actions), size=(ts - 1,))
+    rewards = np.random.choice([-1, 0, 1], size=(ts - 1,))
 
     # First time-step has no reward.
-    t.add_time_step(observation=observations[0],
-                    done=False,
-                    action=actions[0])
+    t.add_time_step(observation=observations[0], done=False, action=actions[0])
     for i in range(1, ts - 1):
-      t.add_time_step(observation=observations[i],
-                      done=False,
-                      raw_reward=rewards[i-1],
-                      processed_reward=rewards[i-1],
-                      action=actions[i])
+      t.add_time_step(
+          observation=observations[i],
+          done=False,
+          raw_reward=rewards[i - 1],
+          processed_reward=rewards[i - 1],
+          action=actions[i])
     # Last time-step has no action.
-    t.add_time_step(observation=observations[-1],
-                    done=False,
-                    raw_reward=rewards[-1],
-                    processed_reward=rewards[-1])
+    t.add_time_step(
+        observation=observations[-1],
+        done=False,
+        raw_reward=rewards[-1],
+        processed_reward=rewards[-1])
 
     traj_np = t.as_numpy
 
@@ -206,18 +235,36 @@ def test_truncate(self):
     # Have to call reset first.
     bt.reset(indices, observations)
 
+    # Take a few steps.
+    ts = 5
+    for _ in range(ts):
+      (observations, rewards, actions,
+       dones) = self.get_random_observations_rewards_actions_dones(
+           batch_size=batch_size)
+      dones[...] = False
+      bt.step(observations, rewards, rewards, dones, actions)
+
     self.assertEqual(0, bt.num_completed_trajectories)
 
-    bt.truncate_trajectories(indices)
+    num_to_keep = 2
+    bt.truncate_trajectories(indices, num_to_keep=num_to_keep)
 
     self.assertEqual(batch_size, bt.num_completed_trajectories)
 
-    # Assert they are all active, since the last observation was duplicated.
+    # Assert they are all active.
+    # Since the last `num_to_keep` observations were duplicated.
     self.assertTrue(all(t.is_active for t in bt.trajectories))
 
-    # Test that the observation is the same.
-    self.assertAllEqual(bt.trajectories[0].last_time_step.observation,
-                        bt.completed_trajectories[0].last_time_step.observation)
+    orig_obs = bt.completed_trajectories[0].observations_np
+    # + 1 because of the initial reset
+    self.assertEqual(ts + 1, orig_obs.shape[0])
+
+    trunc_obs = bt.trajectories[0].observations_np
+    self.assertEqual(num_to_keep, trunc_obs.shape[0])
+    self.assertEqual(num_to_keep, bt.trajectories[0].num_time_steps)
+
+    # Test that the observations are the same.
+    self.assertAllEqual(orig_obs[-num_to_keep:, ...], trunc_obs)
 
   def test_step(self):
     bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
@@ -319,8 +366,8 @@ def test_observations_np(self):
 
     ts = 5
     for _ in range(ts):
-      (observations, rewards, actions, dones
-      ) = self.get_random_observations_rewards_actions_dones()
+      (observations, rewards, actions,
+       dones) = self.get_random_observations_rewards_actions_dones()
       dones[...] = False
       bt.step(observations, rewards, rewards, dones, actions)
 
@@ -338,8 +385,8 @@ def test_observations_np(self):
     lengths[0] = lengths[1] = 1
 
     for _ in range(ts):
-      (observations, rewards, actions, dones
-      ) = self.get_random_observations_rewards_actions_dones()
+      (observations, rewards, actions,
+       dones) = self.get_random_observations_rewards_actions_dones()
       dones[...] = False
       bt.step(observations, rewards, rewards, dones, actions)
 
@@ -348,8 +395,10 @@ def test_observations_np(self):
     lengths = lengths + ts
 
     boundary = 20
-    padded_obs_np = bt.observations_np(boundary=boundary)
-    padded_lengths = bt.trajectory_lengths
+    len_history_for_policy = 40
+
+    padded_obs_np, padded_lengths = bt.observations_np(
+        boundary=boundary, len_history_for_policy=len_history_for_policy)
 
     # The lengths are what we expect them to be.
     self.assertAllEqual(lengths, padded_lengths)
@@ -358,11 +407,26 @@ def test_observations_np(self):
     self.assertEqual((self.BATCH_SIZE, boundary + 1) + self.OBSERVATION_SHAPE,
                      padded_obs_np.shape)
 
+    # Let's now request the last n = [1, 2 * boundary) steps for the history.
+    for len_history_for_policy in range(1, 2 * boundary):
+      # The expected lengths will now be:
+      truncated_lengths = [min(l, len_history_for_policy) for l in lengths]
+
+      padded_obs_np, padded_lengths = bt.observations_np(
+          boundary=boundary, len_history_for_policy=len_history_for_policy)
+
+      self.assertAllEqual(truncated_lengths, padded_lengths)
+
+      # This shouldn't change, since even if we request lengths > boundary + 1
+      # there are no trajectories that long.
+      self.assertEqual((self.BATCH_SIZE, boundary + 1) + self.OBSERVATION_SHAPE,
+                       padded_obs_np.shape)
+
     # Let's do 10 more steps (to go on the other side of the boundary.
     ts = 10
     for _ in range(ts):
-      (observations, rewards, actions, dones
-      ) = self.get_random_observations_rewards_actions_dones()
+      (observations, rewards, actions,
+       dones) = self.get_random_observations_rewards_actions_dones()
       dones[...] = False
       bt.step(observations, rewards, rewards, dones, actions)
 
@@ -370,16 +434,17 @@ def test_observations_np(self):
     # (16, 16, 21, 21, 21, 21, 21, 21, 21, 21)
     lengths = lengths + ts
 
-    padded_obs_np = bt.observations_np(boundary=boundary)
-    padded_lengths = bt.trajectory_lengths
+    len_history_for_policy = 40
+    padded_obs_np, padded_lengths = bt.observations_np(
+        boundary=boundary, len_history_for_policy=len_history_for_policy)
 
     # The lengths are what we expect them to be.
     self.assertAllEqual(lengths, padded_lengths)
 
     # The padded_observations are the shape we expect them to be.
-    self.assertEqual((self.BATCH_SIZE,
-                      (2 * boundary) + 1) + self.OBSERVATION_SHAPE,
-                     padded_obs_np.shape)
+    self.assertEqual(
+        (self.BATCH_SIZE, (2 * boundary) + 1) + self.OBSERVATION_SHAPE,
+        padded_obs_np.shape)
 
     # Test that the padding is the only part that is all 0s.
     # NOTE: There is almost 0 probability that the random observation is all 0s.
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index f3d3068d1..732cfe175 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -96,18 +96,14 @@ def policy_and_value_net(rng_key,
   # NOTE: The LogSoftmax instead of the Softmax because of numerical stability.
 
   if two_towers:
-    net = tl.Branch(
-        [bottom_layers_fn(), tl.Dense(n_actions), tl.LogSoftmax()],
-        [bottom_layers_fn(), tl.Dense(1)]
-    )
+    net = tl.Branch([bottom_layers_fn(),
+                     tl.Dense(n_actions),
+                     tl.LogSoftmax()],
+                    [bottom_layers_fn(), tl.Dense(1)])
   else:
     net = tl.Serial(
         bottom_layers_fn(),
-        tl.Branch(
-            [tl.Dense(n_actions), tl.LogSoftmax()],
-            [tl.Dense(1)]
-        )
-    )
+        tl.Branch([tl.Dense(n_actions), tl.LogSoftmax()], [tl.Dense(1)]))
   params = net.initialize(batch_observations_shape, observations_dtype, rng_key)
   return params, net
 
@@ -129,9 +125,9 @@ def collect_trajectories(env,
                          n_trajectories=1,
                          policy=env_problem_utils.CATEGORICAL_SAMPLING,
                          max_timestep=None,
-                         boundary=20,
                          epsilon=0.1,
                          reset=True,
+                         len_history_for_policy=32,
                          rng=None):
   """Collect trajectories with the given policy net and behaviour.
 
@@ -144,10 +140,11 @@ def collect_trajectories(env,
     max_timestep: int or None, the index of the maximum time-step at which we
       return the trajectory, None for ending a trajectory only when env returns
       done.
-    boundary: int, boundary for padding, used in EnvProblem envs.
     epsilon: float, the epsilon for `epsilon-greedy` policy.
     reset: bool, true if we want to reset the envs. The envs are also reset if
       max_max_timestep is None or < 0
+    len_history_for_policy: int, the maximum history to keep for applying the
+      policy on.
     rng: jax rng, splittable.
 
   Returns:
@@ -161,18 +158,18 @@ def collect_trajectories(env,
 
   assert isinstance(env, env_problem.EnvProblem)
   # This is an env_problem, run its collect function.
-  trajs, n_done = env_problem_utils.play_env_problem_with_policy(
+  trajs, n_done, timing_info = env_problem_utils.play_env_problem_with_policy(
       env,
       policy_fn,
       num_trajectories=n_trajectories,
       max_timestep=max_timestep,
-      boundary=boundary,
       policy_sampling=policy,
       eps=epsilon,
       reset=reset,
+      len_history_for_policy=len_history_for_policy,
       rng=rng)
   # Skip returning raw_rewards here, since they aren't used.
-  return [(t[0], t[1], t[2]) for t in trajs], n_done
+  return [(t[0], t[1], t[2]) for t in trajs], n_done, timing_info
 
 
 # This function can probably be simplified, ask how?
@@ -405,8 +402,8 @@ def deltas(predicted_values, rewards, mask, gamma=0.99):
   # Predicted values at time t+1, by cutting off the first to have shape (B, T)
   predicted_values_btplus1 = predicted_values[:, 1:]
   # Return the deltas as defined above.
-  return (
-      rewards + (gamma * predicted_values_btplus1) - predicted_values_bt) * mask
+  return (rewards +
+          (gamma * predicted_values_btplus1) - predicted_values_bt) * mask
 
 
 def gae_advantages(td_deltas, mask, lambda_=0.95, gamma=0.99):
@@ -710,9 +707,9 @@ def masked_entropy(log_probs, mask):
 
 def evaluate_policy(eval_env,
                     get_predictions,
-                    boundary,
                     max_timestep=20000,
                     n_evals=1,
+                    len_history_for_policy=32,
                     rng=None):
   """Evaluate the policy."""
 
@@ -722,21 +719,20 @@ def evaluate_policy(eval_env,
     for policy in [
         env_problem_utils.CATEGORICAL_SAMPLING,
         env_problem_utils.GUMBEL_SAMPLING,
-        env_problem_utils.EPSILON_GREEDY
     ]:
-      trajs, _ = env_problem_utils.play_env_problem_with_policy(
+      trajs, _, _ = env_problem_utils.play_env_problem_with_policy(
           eval_env,
           get_predictions,
           num_trajectories=eval_env.batch_size,
-          boundary=boundary,
           max_timestep=max_timestep,
           reset=True,
           policy_sampling=policy,
-          rng=rng)
+          rng=rng,
+          len_history_for_policy=len_history_for_policy)
       avg_rewards[policy] += float(sum(
           np.sum(traj[2]) for traj in trajs)) / len(trajs)
-      avg_rewards_unclipped[policy] += float(sum(
-          np.sum(traj[3]) for traj in trajs)) / len(trajs)
+      avg_rewards_unclipped[policy] += float(
+          sum(np.sum(traj[3]) for traj in trajs)) / len(trajs)
 
   # Now average these out.
   for k in avg_rewards:
@@ -794,6 +790,7 @@ def training_loop(
     enable_early_stopping=True,
     env_name=None,
     n_evals=1,
+    len_history_for_policy=4,
 ):
   """Runs the training loop for PPO, with fixed policy and value nets."""
   assert env
@@ -877,9 +874,9 @@ def get_predictions(observations, rng=None):
       avg_reward, avg_reward_unclipped = evaluate_policy(
           eval_env,
           get_predictions,
-          boundary,
           max_timestep=max_timestep_eval,
           n_evals=n_evals,
+          len_history_for_policy=len_history_for_policy,
           rng=key)
       for k, v in avg_reward.items():
         eval_sw.scalar("eval/mean_reward/%s" % k, v, step=i)
@@ -894,13 +891,13 @@ def get_predictions(observations, rng=None):
     trajectory_collection_start_time = time.time()
     logging.vlog(1, "Epoch [% 6d] collecting trajectories.", i)
     jax_rng_key, key = jax_random.split(jax_rng_key)
-    trajs, n_done = collect_trajectories(
+    trajs, n_done, timing_info = collect_trajectories(
         env,
         policy_fn=get_predictions,
         n_trajectories=batch_size,
         max_timestep=max_timestep,
-        boundary=boundary,
         rng=key,
+        len_history_for_policy=len_history_for_policy,
         reset=(i == 0) or restore,
         epsilon=(10.0 / (i + 10.0)))  # this is a different epsilon.
     trajectory_collection_time = get_time(trajectory_collection_start_time)
@@ -1084,10 +1081,9 @@ def get_predictions(observations, rng=None):
     policy_save_start_time = time.time()
     n_trajectories_done += n_done
     # TODO(afrozm): Refactor to trax.save_state.
-    if (((n_trajectories_done >= done_frac_for_policy_save * batch_size)
-         and (i - last_saved_at > eval_every_n)
-         and (((i + 1) % eval_every_n == 0)))
-        or (i == epochs - 1)):
+    if (((n_trajectories_done >= done_frac_for_policy_save * batch_size) and
+         (i - last_saved_at > eval_every_n) and
+         (((i + 1) % eval_every_n == 0))) or (i == epochs - 1)):
       logging.vlog(1, "Epoch [% 6d] saving model.", i)
       old_model_files = gfile.glob(os.path.join(output_dir, "model-??????.pkl"))
       params_file = os.path.join(output_dir, "model-%06d.pkl" % i)
@@ -1120,6 +1116,8 @@ def get_predictions(observations, rng=None):
         "policy_save": policy_save_time,
     }
 
+    timing_dict.update(timing_info)
+
     for k, v in timing_dict.items():
       timing_sw.scalar("timing/%s" % k, v, step=i)
 
@@ -1134,7 +1132,7 @@ def get_predictions(observations, rng=None):
     restore = False
 
     # Flush summary writers once in a while.
-    if (i+1) % 1000 == 0 or i == epochs - 1:
+    if (i + 1) % 1000 == 0 or i == epochs - 1:
       train_sw.flush()
       timing_sw.flush()
       eval_sw.flush()
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 478701b34..546781aaa 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -131,6 +131,8 @@
     "done_frac_for_policy_save", 0.5,
     "Fraction of the trajectories that should be done to "
     "checkpoint the policy.")
+flags.DEFINE_integer("len_history_for_policy", 4,
+                     "How much of history to give to the policy.")
 
 
 def common_layers():
@@ -243,6 +245,7 @@ def run_training_loop():
         eval_env=eval_env,
         n_evals=FLAGS.n_evals,
         env_name=str(FLAGS.env_problem_name),
+        len_history_for_policy=int(FLAGS.len_history_for_policy),
     )
 
   if FLAGS.jax_debug_nans or FLAGS.disable_jit:

From c394f62bebfb19cb59368425e63f1f82a15eba27 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 7 Jun 2019 23:59:06 +0200
Subject: [PATCH 2107/2720] Dump policy checkpoints every epoch in PPO, add an
 option to evaluate (#1597)

all epochs in the evaluator
---
 tensor2tensor/rl/evaluator.py   | 24 +++++++++++++++++++++++-
 tensor2tensor/rl/ppo_learner.py | 13 +++++++++----
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 2928665f3..3db1ec411 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -95,6 +95,10 @@
     "random_starts_step_limit", 10000,
     "Number of frames to choose from for random starts of the simulated env."
 )
+flags.DEFINE_bool(
+    "all_epochs", False,
+    "Whether to run the evaluator on policy checkpoints from all epochs."
+)
 
 # Unused flags needed to pass for multi-run infrastructure.
 flags.DEFINE_bool("autotune", False, "Unused here.")
@@ -477,6 +481,21 @@ def get_game_for_worker(map_name, directory_id):
   return games[game_id]
 
 
+def evaluate_all_epochs(
+    loop_hparams, planner_hparams, policy_dir, model_dir, eval_metrics_dir,
+    *args, **kwargs
+):
+  epoch_policy_dirs = tf.gfile.Glob(os.path.join(policy_dir, "epoch_*"))
+  for epoch_policy_dir in epoch_policy_dirs:
+    epoch_metrics_dir = os.path.join(eval_metrics_dir, "epoch_{}".format(
+        epoch_policy_dir.split("_")[-1]
+    ))
+    evaluate(
+        loop_hparams, planner_hparams, epoch_policy_dir, model_dir,
+        epoch_metrics_dir, *args, **kwargs
+    )
+
+
 def main(_):
   now = datetime.datetime.now()
   now_tag = now.strftime("%Y_%m_%d_%H_%M")
@@ -496,6 +515,7 @@ def main(_):
   model_dir = FLAGS.model_dir
   eval_metrics_dir = FLAGS.eval_metrics_dir
   debug_video_path = FLAGS.debug_video_path
+  evaluate_fn = evaluate
   if FLAGS.output_dir:
     cur_dir = FLAGS.output_dir
     if FLAGS.total_num_workers > 1:
@@ -510,7 +530,9 @@ def main(_):
     tf.logging.info("Writing metrics to %s." % eval_metrics_dir)
     if not tf.gfile.Exists(eval_metrics_dir):
       tf.gfile.MkDir(eval_metrics_dir)
-  evaluate(
+    if FLAGS.all_epochs:
+      evaluate_fn = evaluate_all_epochs
+  evaluate_fn(
       loop_hparams, planner_hparams, policy_dir, model_dir,
       eval_metrics_dir, FLAGS.agent, FLAGS.mode, FLAGS.eval_with_learner,
       FLAGS.log_every_steps if FLAGS.log_every_steps > 0 else None,
diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 770f333c2..9ca4ffccd 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -128,6 +128,7 @@ def train(self,
             train_summary_op,
             eval_summary_op,
             initializers,
+            epoch,
             report_fn=report_fn,
             model_save_fn=model_save_fn)
 
@@ -213,6 +214,7 @@ def _run_train(ppo_hparams,
                train_summary_op,
                eval_summary_op,
                initializers,
+               epoch,
                report_fn=None,
                model_save_fn=None):
   """Train."""
@@ -262,11 +264,14 @@ def _run_train(ppo_hparams,
         if (model_saver and ppo_hparams.save_models_every_epochs and
             (epoch_index % ppo_hparams.save_models_every_epochs == 0 or
              (epoch_index + 1) == num_target_iterations)):
-          ckpt_path = os.path.join(
-              model_dir,
-              "model.ckpt-{}".format(tf.train.global_step(sess, global_step))
+          ckpt_name = "model.ckpt-{}".format(
+              tf.train.global_step(sess, global_step)
           )
-          model_saver.save(sess, ckpt_path)
+          # Keep the last checkpoint from each epoch in a separate directory.
+          epoch_dir = os.path.join(model_dir, "epoch_{}".format(epoch))
+          tf.gfile.MakeDirs(epoch_dir)
+          for ckpt_dir in (model_dir, epoch_dir):
+            model_saver.save(sess, os.path.join(ckpt_dir, ckpt_name))
           if model_save_fn:
             model_save_fn(model_dir)
 

From cb76e97b4e956940e0a50f1ad80e21e815262ae2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Hubert=20Bry=C5=82kowski?= <hubert.brylkowski@brainly.com>
Date: Sat, 8 Jun 2019 00:21:15 +0200
Subject: [PATCH 2108/2720] Documentation for creating own model (#1589)

* Update mscoco.py

* docs for adding new model

* corrected contributing link
---
 docs/new_model.md | 97 +++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 94 insertions(+), 3 deletions(-)

diff --git a/docs/new_model.md b/docs/new_model.md
index 5968c8325..d56856213 100644
--- a/docs/new_model.md
+++ b/docs/new_model.md
@@ -5,12 +5,103 @@ version](https://badge.fury.io/py/tensor2tensor.svg)](https://badge.fury.io/py/t
 [![GitHub
 Issues](https://img.shields.io/github/issues/tensorflow/tensor2tensor.svg)](https://github.com/tensorflow/tensor2tensor/issues)
 [![Contributions
-welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CONTRIBUTING.md)
+welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](../CONTRIBUTING.md)
 [![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
 [![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
 
 Here we show how to create your own model in T2T.
 
-## The T2TModel class
+## The T2TModel class - abstract base class for models
 
-TODO: complete.
+  `T2TModel` has three typical usages:
+
+  1. Estimator: The method `make_estimator_model_fn` builds a `model_fn` for
+     the tf.Estimator workflow of training, evaluation, and prediction.
+     It performs the method `call`, which performs the core computation,
+     followed by `estimator_spec_train`, `estimator_spec_eval`, or
+     `estimator_spec_predict` depending on the tf.Estimator mode.
+  2. Layer: The method `call` enables `T2TModel` to be used a callable by
+     itself. It calls the following methods:
+
+     * `bottom`, which transforms features according to `problem_hparams`' input
+       and target `Modality`s;
+     * `body`, which takes features and performs the core model computation to
+        return output and any auxiliary loss terms;
+     * `top`, which takes features and the body output, and transforms them
+       according to `problem_hparams`' input and target `Modality`s to return
+       the final logits;
+     * `loss`, which takes the logits, forms any missing training loss, and sums
+       all loss terms.
+  3. Inference: The method `infer` enables `T2TModel` to make sequence
+     predictions by itself.
+
+
+## Creating your own model
+
+1. Create class that extends T2TModel 
+    in this example it will be a copy of existing basic fully connected network:
+    ```python
+    from tensor2tensor.utils import t2t_model
+ 
+    class MyFC(t2t_model.T2TModel):
+        pass
+    ```
+
+2. Implement body method:
+    ```python
+    class MyFC(t2t_model.T2TModel):
+      def body(self, features):
+        hparams = self.hparams
+        x = features["inputs"]
+        shape = common_layers.shape_list(x)
+        x = tf.reshape(x, [-1, shape[1] * shape[2] * shape[3]])  # Flatten input as in T2T they are all 4D vectors
+        for i in range(hparams.num_hidden_layers): # create layers
+          x = tf.layers.dense(x, hparams.hidden_size, name="layer_%d" % i)
+          x = tf.nn.dropout(x, keep_prob=1.0 - hparams.dropout)
+          x = tf.nn.relu(x)
+        return tf.expand_dims(tf.expand_dims(x, axis=1), axis=1)  # 4D For T2T.
+    ```
+    
+    method signature:
+    * Args:
+        * features: dict of str to Tensor, where each Tensor has shape [batch_size,
+       ..., hidden_size]. It typically contains keys `inputs` and `targets`.
+
+    * Returns one of:
+      * output: Tensor of pre-logit activations with shape [batch_size, ...,
+             hidden_size].
+      * losses: Either single loss as a scalar, a list, a Tensor (to be averaged),
+             or a dictionary of losses. If losses is a dictionary with the key
+             "training", losses["training"] is considered the final training
+             loss and output is considered logits; self.top and self.loss will
+             be skipped.
+
+3. Register your model
+    ```python
+    from tensor2tensor.utils import registry
+
+    @registry.register_model
+    class MyFC(t2t_model.T2TModel):
+       # ...
+    ```
+
+3. Use it with t2t tools as any other model
+
+    Have in mind that names are translated from camel case to snake_case `MyFC` -> `my_fc`
+    and that you need to point t2t to directory containing your model with `t2t_usr_dir` switch. 
+    For example if you want to train model on gcloud with 1 GPU worker on IMDB sentiment task you can run your model
+    by executing following command from your model class directory. 
+    
+    ```bash
+    t2t-trainer \
+      --model=my_fc \
+      --t2t_usr_dir=.
+      --cloud_mlengine --worker_gpu=1 \
+      --generate_data \
+      --data_dir='gs://data' \
+      --output_dir='gs://out' \
+      --problem=sentiment_imdb \
+      --hparams_set=basic_fc_small \
+      --train_steps=10000 \
+      --eval_steps=10 \
+    ```

From db02a3b7e7692ce2af58832c6ea3a77ad1266427 Mon Sep 17 00:00:00 2001
From: Kalpesh Krishna <kalpeshk2011@gmail.com>
Date: Fri, 7 Jun 2019 18:22:19 -0400
Subject: [PATCH 2109/2720] Adding extra linear to semantic hashing
 discretization bottleneck. (#1578)

---
 tensor2tensor/layers/discretization.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 970b705e4..f8f7c91e1 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -309,6 +309,8 @@ def embed(x,
       h1a = tf.layers.dense(c, filter_size, name="vch1a")
       h1b = tf.layers.dense(1.0 - c, filter_size, name="vch1b")
       h1 = h1a + h1b
+      h1 = tf.layers.dense(h1, hidden_size, name="vch_final_linear")
+
     elif bottleneck_kind == "gumbel-softmax":
       hot = tf.one_hot(x, 2**z_size)
       h1 = tf.layers.dense(hot, hidden_size, name="dae_dense")
@@ -773,6 +775,8 @@ def discrete_bottleneck(inputs,
       outputs_dense_a = tf.layers.dense(c, filter_size, name="vch1a")
       outputs_dense_b = tf.layers.dense(1.0 - c, filter_size, name="vch1b")
       outputs_dense = outputs_dense_a + outputs_dense_b
+      outputs_dense = tf.layers.dense(outputs_dense, hidden_size, name="vch_final_linear")
+
       dx = tf.to_int32(tf.stop_gradient(d))
       outputs_discrete = bit_to_int(dx, z_size)
       extra_loss = tf.constant(0.0)

From 7ad6c7fa38c9899f6cfc7834002662510597fde0 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <kozak000@gmail.com>
Date: Fri, 7 Jun 2019 14:59:35 -0700
Subject: [PATCH 2110/2720] Merge of PR #1597

PiperOrigin-RevId: 252125723
---
 docs/new_model.md                      | 97 +-------------------------
 tensor2tensor/layers/discretization.py |  4 --
 2 files changed, 3 insertions(+), 98 deletions(-)

diff --git a/docs/new_model.md b/docs/new_model.md
index d56856213..5968c8325 100644
--- a/docs/new_model.md
+++ b/docs/new_model.md
@@ -5,103 +5,12 @@ version](https://badge.fury.io/py/tensor2tensor.svg)](https://badge.fury.io/py/t
 [![GitHub
 Issues](https://img.shields.io/github/issues/tensorflow/tensor2tensor.svg)](https://github.com/tensorflow/tensor2tensor/issues)
 [![Contributions
-welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](../CONTRIBUTING.md)
+welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CONTRIBUTING.md)
 [![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
 [![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
 
 Here we show how to create your own model in T2T.
 
-## The T2TModel class - abstract base class for models
+## The T2TModel class
 
-  `T2TModel` has three typical usages:
-
-  1. Estimator: The method `make_estimator_model_fn` builds a `model_fn` for
-     the tf.Estimator workflow of training, evaluation, and prediction.
-     It performs the method `call`, which performs the core computation,
-     followed by `estimator_spec_train`, `estimator_spec_eval`, or
-     `estimator_spec_predict` depending on the tf.Estimator mode.
-  2. Layer: The method `call` enables `T2TModel` to be used a callable by
-     itself. It calls the following methods:
-
-     * `bottom`, which transforms features according to `problem_hparams`' input
-       and target `Modality`s;
-     * `body`, which takes features and performs the core model computation to
-        return output and any auxiliary loss terms;
-     * `top`, which takes features and the body output, and transforms them
-       according to `problem_hparams`' input and target `Modality`s to return
-       the final logits;
-     * `loss`, which takes the logits, forms any missing training loss, and sums
-       all loss terms.
-  3. Inference: The method `infer` enables `T2TModel` to make sequence
-     predictions by itself.
-
-
-## Creating your own model
-
-1. Create class that extends T2TModel 
-    in this example it will be a copy of existing basic fully connected network:
-    ```python
-    from tensor2tensor.utils import t2t_model
- 
-    class MyFC(t2t_model.T2TModel):
-        pass
-    ```
-
-2. Implement body method:
-    ```python
-    class MyFC(t2t_model.T2TModel):
-      def body(self, features):
-        hparams = self.hparams
-        x = features["inputs"]
-        shape = common_layers.shape_list(x)
-        x = tf.reshape(x, [-1, shape[1] * shape[2] * shape[3]])  # Flatten input as in T2T they are all 4D vectors
-        for i in range(hparams.num_hidden_layers): # create layers
-          x = tf.layers.dense(x, hparams.hidden_size, name="layer_%d" % i)
-          x = tf.nn.dropout(x, keep_prob=1.0 - hparams.dropout)
-          x = tf.nn.relu(x)
-        return tf.expand_dims(tf.expand_dims(x, axis=1), axis=1)  # 4D For T2T.
-    ```
-    
-    method signature:
-    * Args:
-        * features: dict of str to Tensor, where each Tensor has shape [batch_size,
-       ..., hidden_size]. It typically contains keys `inputs` and `targets`.
-
-    * Returns one of:
-      * output: Tensor of pre-logit activations with shape [batch_size, ...,
-             hidden_size].
-      * losses: Either single loss as a scalar, a list, a Tensor (to be averaged),
-             or a dictionary of losses. If losses is a dictionary with the key
-             "training", losses["training"] is considered the final training
-             loss and output is considered logits; self.top and self.loss will
-             be skipped.
-
-3. Register your model
-    ```python
-    from tensor2tensor.utils import registry
-
-    @registry.register_model
-    class MyFC(t2t_model.T2TModel):
-       # ...
-    ```
-
-3. Use it with t2t tools as any other model
-
-    Have in mind that names are translated from camel case to snake_case `MyFC` -> `my_fc`
-    and that you need to point t2t to directory containing your model with `t2t_usr_dir` switch. 
-    For example if you want to train model on gcloud with 1 GPU worker on IMDB sentiment task you can run your model
-    by executing following command from your model class directory. 
-    
-    ```bash
-    t2t-trainer \
-      --model=my_fc \
-      --t2t_usr_dir=.
-      --cloud_mlengine --worker_gpu=1 \
-      --generate_data \
-      --data_dir='gs://data' \
-      --output_dir='gs://out' \
-      --problem=sentiment_imdb \
-      --hparams_set=basic_fc_small \
-      --train_steps=10000 \
-      --eval_steps=10 \
-    ```
+TODO: complete.
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index f8f7c91e1..970b705e4 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -309,8 +309,6 @@ def embed(x,
       h1a = tf.layers.dense(c, filter_size, name="vch1a")
       h1b = tf.layers.dense(1.0 - c, filter_size, name="vch1b")
       h1 = h1a + h1b
-      h1 = tf.layers.dense(h1, hidden_size, name="vch_final_linear")
-
     elif bottleneck_kind == "gumbel-softmax":
       hot = tf.one_hot(x, 2**z_size)
       h1 = tf.layers.dense(hot, hidden_size, name="dae_dense")
@@ -775,8 +773,6 @@ def discrete_bottleneck(inputs,
       outputs_dense_a = tf.layers.dense(c, filter_size, name="vch1a")
       outputs_dense_b = tf.layers.dense(1.0 - c, filter_size, name="vch1b")
       outputs_dense = outputs_dense_a + outputs_dense_b
-      outputs_dense = tf.layers.dense(outputs_dense, hidden_size, name="vch_final_linear")
-
       dx = tf.to_int32(tf.stop_gradient(d))
       outputs_discrete = bit_to_int(dx, z_size)
       extra_loss = tf.constant(0.0)

From 0e0d4490765e7a485876c09d5028db58e9b30938 Mon Sep 17 00:00:00 2001
From: Hubert Bry?kowski <hubert+github@brylkowski.com>
Date: Fri, 7 Jun 2019 15:23:35 -0700
Subject: [PATCH 2111/2720] Merge of PR #1589

PiperOrigin-RevId: 252130487
---
 docs/new_model.md | 104 ++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 101 insertions(+), 3 deletions(-)

diff --git a/docs/new_model.md b/docs/new_model.md
index 5968c8325..861f83bb0 100644
--- a/docs/new_model.md
+++ b/docs/new_model.md
@@ -5,12 +5,110 @@ version](https://badge.fury.io/py/tensor2tensor.svg)](https://badge.fury.io/py/t
 [![GitHub
 Issues](https://img.shields.io/github/issues/tensorflow/tensor2tensor.svg)](https://github.com/tensorflow/tensor2tensor/issues)
 [![Contributions
-welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CONTRIBUTING.md)
+welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](../CONTRIBUTING.md)
 [![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
 [![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
 
 Here we show how to create your own model in T2T.
 
-## The T2TModel class
+## The T2TModel class - abstract base class for models
 
-TODO: complete.
+  `T2TModel` has three typical usages:
+
+  1. Estimator: The method `make_estimator_model_fn` builds a `model_fn` for
+     the tf.Estimator workflow of training, evaluation, and prediction.
+     It performs the method `call`, which performs the core computation,
+     followed by `estimator_spec_train`, `estimator_spec_eval`, or
+     `estimator_spec_predict` depending on the tf.Estimator mode.
+  2. Layer: The method `call` enables `T2TModel` to be used a callable by
+     itself. It calls the following methods:
+
+     * `bottom`, which transforms features according to `problem_hparams`' input
+       and target `Modality`s;
+     * `body`, which takes features and performs the core model computation to
+        return output and any auxiliary loss terms;
+     * `top`, which takes features and the body output, and transforms them
+       according to `problem_hparams`' input and target `Modality`s to return
+       the final logits;
+     * `loss`, which takes the logits, forms any missing training loss, and sums
+       all loss terms.
+  3. Inference: The method `infer` enables `T2TModel` to make sequence
+     predictions by itself.
+
+
+## Creating your own model
+
+1. Create class that extends T2TModel
+    in this example it will be a copy of existing basic fully connected network:
+
+```python
+    from tensor2tensor.utils import t2t_model
+
+    class MyFC(t2t_model.T2TModel):
+        pass
+```
+
+
+2. Implement body method:
+
+```python
+    class MyFC(t2t_model.T2TModel):
+      def body(self, features):
+        hparams = self.hparams
+        x = features["inputs"]
+        shape = common_layers.shape_list(x)
+        x = tf.reshape(x, [-1, shape[1] * shape[2] * shape[3]])  # Flatten input as in T2T they are all 4D vectors
+        for i in range(hparams.num_hidden_layers): # create layers
+          x = tf.layers.dense(x, hparams.hidden_size, name="layer_%d" % i)
+          x = tf.nn.dropout(x, keep_prob=1.0 - hparams.dropout)
+          x = tf.nn.relu(x)
+        return tf.expand_dims(tf.expand_dims(x, axis=1), axis=1)  # 4D For T2T.
+```
+
+
+Method Signature:
+
+  * Args:
+      * features: dict of str to Tensor, where each Tensor has shape [batch_size,
+     ..., hidden_size]. It typically contains keys `inputs` and `targets`.
+
+  * Returns one of:
+    * output: Tensor of pre-logit activations with shape [batch_size, ...,
+           hidden_size].
+    * losses: Either single loss as a scalar, a list, a Tensor (to be averaged),
+           or a dictionary of losses. If losses is a dictionary with the key
+           "training", losses["training"] is considered the final training
+           loss and output is considered logits; self.top and self.loss will
+           be skipped.
+
+3. Register your model
+
+```python
+    from tensor2tensor.utils import registry
+
+    @registry.register_model
+    class MyFC(t2t_model.T2TModel):
+       # ...
+```
+
+
+3. Use it with t2t tools as any other model
+
+    Have in mind that names are translated from camel case to snake_case `MyFC` -> `my_fc`
+    and that you need to point t2t to directory containing your model with `t2t_usr_dir` switch. 
+    For example if you want to train model on gcloud with 1 GPU worker on IMDB sentiment task you can run your model
+    by executing following command from your model class directory. 
+
+```bash
+    t2t-trainer \
+      --model=my_fc \
+      --t2t_usr_dir=.
+      --cloud_mlengine --worker_gpu=1 \
+      --generate_data \
+      --data_dir='gs://data' \
+      --output_dir='gs://out' \
+      --problem=sentiment_imdb \
+      --hparams_set=basic_fc_small \
+      --train_steps=10000 \
+      --eval_steps=10 \
+```

From e118269bcf2bbf25d76abff834bd93dfc29c3714 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 7 Jun 2019 15:49:00 -0700
Subject: [PATCH 2112/2720] Further break out collect timing.

PiperOrigin-RevId: 252134897
---
 tensor2tensor/envs/env_problem_utils.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 3c7952480..2ac197257 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -147,6 +147,7 @@ def epsilon_greedy(log_probs):
   num_done_trajectories = 0
 
   policy_application_total_time = 0
+  env_actions_total_time = 0
   while env.trajectories.num_completed_trajectories < num_trajectories:
     # Get all the observations for all the active trajectories.
     # Shape is (B, T) + OBS
@@ -186,7 +187,9 @@ def epsilon_greedy(log_probs):
       raise ValueError("Unknown sampling policy [%s]" % policy_sampling)
 
     # Step through the env.
+    t1 = time.time()
     _, _, dones, _ = env.step(actions)
+    env_actions_total_time += (time.time() - t1)
 
     # Count the number of done trajectories, the others could just have been
     # truncated.
@@ -196,8 +199,10 @@ def epsilon_greedy(log_probs):
     done_idxs = done_indices(dones)
 
     # ... and reset those.
+    t1 = time.time()
     if done_idxs.size:
       env.reset(indices=done_idxs)
+    env_actions_total_time += (time.time() - t1)
 
     if max_timestep is None or max_timestep < 1:
       continue
@@ -207,10 +212,12 @@ def epsilon_greedy(log_probs):
     exceeded_time_limit_idxs = done_indices(lengths > max_timestep)
 
     # If so, reset these as well.
+    t1 = time.time()
     if exceeded_time_limit_idxs.size:
       # This just cuts the trajectory, doesn't reset the env, so it continues
       # from where it left off.
       env.truncate(indices=exceeded_time_limit_idxs, num_to_keep=num_to_keep)
+    env_actions_total_time += (time.time() - t1)
 
   # We have the trajectories we need, return a list of triples:
   # (observations, actions, rewards)
@@ -218,11 +225,12 @@ def epsilon_greedy(log_probs):
   for trajectory in env.trajectories.completed_trajectories[:num_trajectories]:
     completed_trajectories.append(trajectory.as_numpy)
 
-  policy_application_time = round(1000 * policy_application_total_time, 2)
-  misc_time = round(1000 * (time.time() - t0), 2) - policy_application_time
+  misc_time = (time.time() - t0) - policy_application_total_time
   timing_info = {
-      "trajectory_collection/policy_application": policy_application_time,
+      "trajectory_collection/policy_application": policy_application_total_time,
       "trajectory_collection/misc": misc_time,
+      "trajectory_collection/env_actions": env_actions_total_time,
   }
+  timing_info = {k: round(1000 * v, 2) for k, v in timing_info.items()}
 
   return completed_trajectories, num_done_trajectories, timing_info

From e51b36d12aa7a04c86cd06df674a3cd565c43635 Mon Sep 17 00:00:00 2001
From: Kalpesh Krishna <kalpeshk2011@gmail.com>
Date: Fri, 7 Jun 2019 16:02:33 -0700
Subject: [PATCH 2113/2720] Merge of PR #1578

PiperOrigin-RevId: 252137294
---
 tensor2tensor/layers/discretization.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 970b705e4..667b6e5a4 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -309,6 +309,8 @@ def embed(x,
       h1a = tf.layers.dense(c, filter_size, name="vch1a")
       h1b = tf.layers.dense(1.0 - c, filter_size, name="vch1b")
       h1 = h1a + h1b
+      h1 = tf.layers.dense(h1, hidden_size, name="vch_final_linear")
+
     elif bottleneck_kind == "gumbel-softmax":
       hot = tf.one_hot(x, 2**z_size)
       h1 = tf.layers.dense(hot, hidden_size, name="dae_dense")
@@ -773,6 +775,9 @@ def discrete_bottleneck(inputs,
       outputs_dense_a = tf.layers.dense(c, filter_size, name="vch1a")
       outputs_dense_b = tf.layers.dense(1.0 - c, filter_size, name="vch1b")
       outputs_dense = outputs_dense_a + outputs_dense_b
+      outputs_dense = tf.layers.dense(outputs_dense, hidden_size,
+                                      name="vch_final_linear")
+
       dx = tf.to_int32(tf.stop_gradient(d))
       outputs_discrete = bit_to_int(dx, z_size)
       extra_loss = tf.constant(0.0)

From f45ca8491191fd52923737a7a5ed9ab55f2688ea Mon Sep 17 00:00:00 2001
From: Eugene Karaulov <EugKar@gmail.com>
Date: Sat, 8 Jun 2019 02:38:00 +0300
Subject: [PATCH 2114/2720] Using partial targets at inference time. (#1596)

---
 tensor2tensor/models/transformer.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 94f6dc9e6..41e6700a6 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -709,7 +709,10 @@ def _fast_decode(self,
             features=features)
       encoder_output = encoder_output[0]
       encoder_decoder_attention_bias = encoder_decoder_attention_bias[0]
-      partial_targets = None
+      if 'partial_targets' in features:
+        partial_targets = features['partial_targets']
+      else:
+        partial_targets = None
     else:
       # The problem has no inputs.
       encoder_output = None
@@ -722,6 +725,8 @@ def _fast_decode(self,
       if partial_targets is None:
         partial_targets = features["targets"]
       assert partial_targets is not None
+
+    if partial_targets is not None:
       partial_targets = common_layers.expand_squeeze_to_nd(partial_targets, 2)
       partial_targets = tf.to_int64(partial_targets)
       partial_targets_shape = common_layers.shape_list(partial_targets)

From 6577bd7eab70a9cbc9c201df6c4fa0be6d54c17c Mon Sep 17 00:00:00 2001
From: Max Sobol Mark <maxsobolmark@gmail.com>
Date: Fri, 7 Jun 2019 16:39:16 -0700
Subject: [PATCH 2115/2720] Updated link to DeepMind Math dataset (#1583)

---
 tensor2tensor/data_generators/algorithmic_math_deepmind.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/algorithmic_math_deepmind.py b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
index 18841b2f0..dce068130 100644
--- a/tensor2tensor/data_generators/algorithmic_math_deepmind.py
+++ b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
@@ -33,7 +33,7 @@
 import tensorflow as tf
 
 
-_URL = "https://storage.googleapis.com/mathematics-dataset/v1.0.tar.gz"
+_URL = "https://storage.cloud.google.com/mathematics-dataset/mathematics_dataset-v1.0.tar.gz"
 
 
 @registry.register_problem

From a002665bbcaffd8815b03f5eee4a23fbf472fabc Mon Sep 17 00:00:00 2001
From: Shufang Xie <xieshufang76@gmail.com>
Date: Sat, 8 Jun 2019 07:39:57 +0800
Subject: [PATCH 2116/2720] Only strip end of line (#1577)

---
 tensor2tensor/data_generators/text_encoder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 4fe071093..1d241164d 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -924,7 +924,7 @@ def _load_from_file_object(self, f):
     """
     subtoken_strings = []
     for line in f:
-      s = line.strip()
+      s = line.rstrip()
       # Some vocab files wrap words in single quotes, but others don't
       if ((s.startswith("'") and s.endswith("'")) or
           (s.startswith("\"") and s.endswith("\""))):

From 8b0cfed4e71adf625505979660ada268cd54f37a Mon Sep 17 00:00:00 2001
From: Eugene Karaulov <EugKar@gmail.com>
Date: Fri, 7 Jun 2019 16:38:42 -0700
Subject: [PATCH 2117/2720] Merge of PR #1596

PiperOrigin-RevId: 252143053
---
 tensor2tensor/data_generators/algorithmic_math_deepmind.py | 2 +-
 tensor2tensor/data_generators/text_encoder.py              | 2 +-
 tensor2tensor/models/transformer.py                        | 5 +----
 3 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/algorithmic_math_deepmind.py b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
index dce068130..18841b2f0 100644
--- a/tensor2tensor/data_generators/algorithmic_math_deepmind.py
+++ b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
@@ -33,7 +33,7 @@
 import tensorflow as tf
 
 
-_URL = "https://storage.cloud.google.com/mathematics-dataset/mathematics_dataset-v1.0.tar.gz"
+_URL = "https://storage.googleapis.com/mathematics-dataset/v1.0.tar.gz"
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 1d241164d..4fe071093 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -924,7 +924,7 @@ def _load_from_file_object(self, f):
     """
     subtoken_strings = []
     for line in f:
-      s = line.rstrip()
+      s = line.strip()
       # Some vocab files wrap words in single quotes, but others don't
       if ((s.startswith("'") and s.endswith("'")) or
           (s.startswith("\"") and s.endswith("\""))):
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 41e6700a6..c6a678250 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -709,10 +709,7 @@ def _fast_decode(self,
             features=features)
       encoder_output = encoder_output[0]
       encoder_decoder_attention_bias = encoder_decoder_attention_bias[0]
-      if 'partial_targets' in features:
-        partial_targets = features['partial_targets']
-      else:
-        partial_targets = None
+      partial_targets = features.get("partial_targets")
     else:
       # The problem has no inputs.
       encoder_output = None

From ac96bbb830d2843b3371cce35a7ce83cbc4ac100 Mon Sep 17 00:00:00 2001
From: Max Sobol Mark <maxsobolmark@gmail.com>
Date: Fri, 7 Jun 2019 16:39:31 -0700
Subject: [PATCH 2118/2720] Merge of PR #1583

PiperOrigin-RevId: 252143163
---
 tensor2tensor/data_generators/algorithmic_math_deepmind.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/algorithmic_math_deepmind.py b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
index 18841b2f0..dce068130 100644
--- a/tensor2tensor/data_generators/algorithmic_math_deepmind.py
+++ b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
@@ -33,7 +33,7 @@
 import tensorflow as tf
 
 
-_URL = "https://storage.googleapis.com/mathematics-dataset/v1.0.tar.gz"
+_URL = "https://storage.cloud.google.com/mathematics-dataset/mathematics_dataset-v1.0.tar.gz"
 
 
 @registry.register_problem

From a35b76c8d8785996c3291f2e5031da5c9b96b435 Mon Sep 17 00:00:00 2001
From: Shufang Xie <xieshufang76@gmail.com>
Date: Fri, 7 Jun 2019 16:40:21 -0700
Subject: [PATCH 2119/2720] Merge of PR #1577

PiperOrigin-RevId: 252143317
---
 tensor2tensor/data_generators/text_encoder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 4fe071093..1d241164d 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -924,7 +924,7 @@ def _load_from_file_object(self, f):
     """
     subtoken_strings = []
     for line in f:
-      s = line.strip()
+      s = line.rstrip()
       # Some vocab files wrap words in single quotes, but others don't
       if ((s.startswith("'") and s.endswith("'")) or
           (s.startswith("\"") and s.endswith("\""))):

From 54f3949c2bf09273e1e80482da150e1f5a05f6d9 Mon Sep 17 00:00:00 2001
From: Thomas O'Malley <omalleyt@google.com>
Date: Sat, 8 Jun 2019 12:03:44 -0700
Subject: [PATCH 2120/2720] Pop `training` arg before passing to sublayer.

PiperOrigin-RevId: 252221438
---
 tensor2tensor/layers/bayes.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 672225b10..d14eaa2d1 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -129,6 +129,7 @@ def call_weights(self):
 
   def call(self, *args, **kwargs):
     self.call_weights()
+    kwargs.pop('training', None)
     return super(Conv2DReparameterization, self).call(*args, **kwargs)
 
 
@@ -467,6 +468,7 @@ def call_weights(self):
 
   def call(self, *args, **kwargs):
     self.call_weights()
+    kwargs.pop('training', None)
     return super(DenseReparameterization, self).call(*args, **kwargs)
 
 

From 5222832036dd7ff1a369d3b4deb479af37c63998 Mon Sep 17 00:00:00 2001
From: David So <davidso@google.com>
Date: Mon, 10 Jun 2019 11:46:50 -0700
Subject: [PATCH 2121/2720] Update single_cycle_cos_decay to use
 hparams.train_steps and add deep setting for ET.

PiperOrigin-RevId: 252451273
---
 tensor2tensor/models/evolved_transformer.py | 16 ++++++++++------
 tensor2tensor/utils/learning_rate.py        |  7 ++++++-
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 8d3d99e61..2da51fbd2 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -724,12 +724,6 @@ def add_evolved_transformer_hparams(hparams):
   hparams.learning_rate_constant /= hparams.learning_rate_warmup_steps ** 0.5
   hparams.learning_rate_schedule = (
       "constant*linear_warmup*single_cycle_cos_decay*rsqrt_hidden_size")
-  # The current infrastructure does not support exposing
-  # `train_steps` to the decay functions, and so we are hard coding the decay
-  # steps here to match the default number of train steps used in `t2t_trainer`.
-  # TODO(davidso): Thread `train_steps` through to decay functions so we do not
-  # have to worry about a `learning_rate_decay_steps` mismatch.
-  hparams.learning_rate_decay_steps = 250000
   return hparams
 
 
@@ -745,6 +739,16 @@ def evolved_transformer_big():
   return add_evolved_transformer_hparams(transformer.transformer_big())
 
 
+@registry.register_hparams
+def evolved_transformer_deep():
+  """Deep parameters for Evolved Transformer model on WMT."""
+  hparams = add_evolved_transformer_hparams(transformer.transformer_big())
+  hparams.num_encoder_layers = 9
+  hparams.num_decoder_layers = 10
+  hparams.hidden_size = 640
+  return hparams
+
+
 @registry.register_hparams
 def evolved_transformer_base_tpu():
   """Base parameters for Evolved Transformer model on TPU."""
diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index 377280fa1..382742b0e 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -45,8 +45,13 @@ def learning_rate_factor(name, step_num, hparams):
     # "cosdecay" because it starts at 1 when the warmup steps end.
     x = tf.maximum(step_num, hparams.learning_rate_warmup_steps)
     step = x - hparams.learning_rate_warmup_steps
+    if hparams.train_steps <= hparams.learning_rate_warmup_steps:
+      raise ValueError("single_cycle_cos_decay cannot be used unless "
+                       "hparams.train_steps > "
+                       "hparams.learning_rate_warmup_steps")
     return tf.math.cos(
-        step * np.pi / hparams.learning_rate_decay_steps) / 2.0 + 0.5
+        step * np.pi /
+        (hparams.train_steps - hparams.learning_rate_warmup_steps)) / 2.0 + 0.5
   elif name == "rsqrt_decay":
     return tf.rsqrt(tf.maximum(step_num, hparams.learning_rate_warmup_steps))
   elif name == "rsqrt_normalized_decay":

From 4331f5e892f315a4c7217608c7f5e320edc04129 Mon Sep 17 00:00:00 2001
From: David So <davidso@google.com>
Date: Mon, 10 Jun 2019 17:19:19 -0700
Subject: [PATCH 2122/2720] Clean up NasSeq2Seq and integrate it with main run
 pipeline.

PiperOrigin-RevId: 252519066
---
 tensor2tensor/models/__init__.py              |   2 +
 .../neural_architecture_search/README.md      |   3 +
 .../neural_architecture_search/nas_layers.py  |  24 +-
 .../nas_layers_test.py                        |  10 +-
 .../neural_architecture_search/nas_model.py   | 249 ++++++++++++------
 .../nas_model_test.py                         |  48 ++--
 6 files changed, 221 insertions(+), 115 deletions(-)
 create mode 100644 tensor2tensor/models/neural_architecture_search/README.md

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index be54a772c..59ec38cb3 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -44,6 +44,8 @@
 from tensor2tensor.models import vanilla_gan
 from tensor2tensor.models import xception
 
+from tensor2tensor.models.neural_architecture_search import nas_model
+
 from tensor2tensor.models.research import adafactor_experiments
 from tensor2tensor.models.research import aligned
 from tensor2tensor.models.research import attention_lm
diff --git a/tensor2tensor/models/neural_architecture_search/README.md b/tensor2tensor/models/neural_architecture_search/README.md
new file mode 100644
index 000000000..c197c88ab
--- /dev/null
+++ b/tensor2tensor/models/neural_architecture_search/README.md
@@ -0,0 +1,3 @@
+This directory contains the configurable model code used in the Evolved
+Transformer paper (https://arxiv.org/abs/1901.11117). It can be used to train
+models in the search space as was done in the paper.
diff --git a/tensor2tensor/models/neural_architecture_search/nas_layers.py b/tensor2tensor/models/neural_architecture_search/nas_layers.py
index c4c9839a1..ae8bb5ee7 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_layers.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers.py
@@ -464,8 +464,8 @@ def num_params(self, input_depth, output_depth, **unused_kwargs):
 class AttendToEncoderLayerBase(TranslationLayer):
   """Attend to encoder base, with configurable encoder attend points."""
 
-  def _determine_encoder_block_index(self, block_number, num_encoder_blocks):
-    """Determine the encoder block index to attend to."""
+  def _determine_encoder_cell_index(self, cell_number, num_encoder_cells):
+    """Determine the encoder cell index to attend to."""
     raise NotImplementedError()
 
   def _apply_logic(self,
@@ -476,17 +476,17 @@ def _apply_logic(self,
                    nonpadding,
                    mask_future,
                    encoder_decoder_attention_bias,
-                   encoder_block_outputs,
-                   block_number,
+                   encoder_cell_outputs,
+                   cell_number,
                    attention_dropout_broadcast_dims=None,
                    **unused_kwargs):
     """Applies attention logic to `input_tensor`."""
     with tf.variable_scope("attend_to_encoder_layer_" + var_scope_suffix):
       hidden_depth = int(input_tensor.shape.as_list()[-1])
-      num_encoder_blocks = len(encoder_block_outputs)
-      encoder_block_index = self._determine_encoder_block_index(
-          block_number, num_encoder_blocks)
-      encoder_layer = encoder_block_outputs[encoder_block_index]
+      num_encoder_cells = len(encoder_cell_outputs)
+      encoder_cell_index = self._determine_encoder_cell_index(
+          cell_number, num_encoder_cells)
+      encoder_layer = encoder_cell_outputs[encoder_cell_index]
 
       # TODO(davidso): This dropout rate differs from the other layers. This
       #                should be fixed so that they all use the same dropout
@@ -534,11 +534,11 @@ def __init__(self, delay, increment_step):
     self.delay = delay
     self.increment_step = increment_step
 
-  def _determine_encoder_block_index(self, block_number, num_encoder_blocks):
-    """Attend to final encoder block output first, then move down."""
+  def _determine_encoder_cell_index(self, cell_number, num_encoder_cells):
+    """Attend to final encoder cell output first, then move down."""
     return max(
-        0, num_encoder_blocks - max(
-            0, (block_number - self.delay) * self.increment_step) - 1)
+        0, num_encoder_cells -
+        max(0, (cell_number - self.delay) * self.increment_step) - 1)
 
 
 class GatedLinearUnitLayer(TranslationLayer):
diff --git a/tensor2tensor/models/neural_architecture_search/nas_layers_test.py b/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
index 07c088aab..1d1fc705e 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
@@ -33,8 +33,8 @@
 _BATCH_SIZE = 32
 _TOTAL_SEQUENCE_LENGTH = 20
 _INPUT_DEPTH = 256
-_NUM_BLOCKS = 6
-_BLOCK_NUMBER = 3
+_NUM_CELLS = 6
+_CELL_NUMBER = 3
 
 # The list of prefixes for layers that will not be tested for resizing outputs.
 _RESIZE_EXEMPT_LAYER_PREFIXES = [
@@ -75,7 +75,7 @@ def _apply_decoder_layer(translation_layer, input_tensor, output_depth,
   encoder_output_values = np.random.rand(
       *[_BATCH_SIZE, _TOTAL_SEQUENCE_LENGTH, encoder_depth]) - .5
   encoder_output = tf.constant(encoder_output_values, dtype=tf.float32)
-  encoder_block_outputs = [encoder_output] * _NUM_BLOCKS
+  encoder_cell_outputs = [encoder_output] * _NUM_CELLS
   hparams = transformer.transformer_base()
   hparams.attention_dropout = 0
   decoder_self_attention_bias = (
@@ -94,8 +94,8 @@ def _apply_decoder_layer(translation_layer, input_tensor, output_depth,
       postprocess_dropout=False,
       decoder_self_attention_bias=decoder_self_attention_bias,
       encoder_decoder_attention_bias=None,
-      encoder_block_outputs=encoder_block_outputs,
-      block_number=_BLOCK_NUMBER)
+      encoder_cell_outputs=encoder_cell_outputs,
+      cell_number=_CELL_NUMBER)
 
   return output_tensor
 
diff --git a/tensor2tensor/models/neural_architecture_search/nas_model.py b/tensor2tensor/models/neural_architecture_search/nas_model.py
index 37d7d04e2..5c85390e1 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model.py
@@ -13,7 +13,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""TranslationNasNet class which can be modified and still used with t2t."""
+"""NasSeq2Seq class which can be configured to produce a variety of models.
+
+This was the class used in the Evolved Transformer paper
+(https://arxiv.org/abs/1901.11117) to create configurable models. It can be used
+to train models in the search space as was done in the paper.
+
+To use NasSeq2Seq:
+  - set model=nas_seq2_seq.
+  - set hparams_set=nas_seq2seq_base.
+  - use hparams to specify the configuration you want to run. See
+    nas_seq2seq_base() for an example.
+"""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -25,6 +36,7 @@
 from tensor2tensor.models import transformer
 from tensor2tensor.models.neural_architecture_search import nas_layers as layers
 from tensor2tensor.utils import metrics
+from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 import tensorflow as tf
 
@@ -217,10 +229,42 @@ def combined_output_dim(self, output_dims):
 }
 
 
+@registry.register_model
 class NasSeq2Seq(transformer.Transformer):
-  """Configurable seq2seq model that uses NAS-like branching.
-
-  Builds a directed graph of operations with arbitrary branching.
+  """Configurable seq2seq model used for Neural Architecture Search.
+
+  Models are defined by 26 hparam fields. They are:
+    - <encoder/decoder>_num_cells: The number of cells in the <encoder/decoder>.
+    - <encoder/decoder>_<left/right>_layers: List of layers used the
+                                             <encoder/decoder> <left/right>
+                                             branch. For available layers, see
+                                             the nas_layers.py file.
+    - <encoder/decoder>_<left/right_inputs>: List of inputs to the
+                                             <encoder/decoder> <left/right>
+                                             layers. Each index i specifies the
+                                             i_th layer's output with 0
+                                             representing the cell input
+                                             tensor.
+    - <encoder/decoder>_<left/right>_output_dims: List of absolute output
+                                                  dimensions for each layer.
+    - <encoder/decoder>_<left/right>_activation: List of activations applied
+                                                 after each layer.
+                                                 ACTIVATION_MAP holds the valid
+                                                 activations.
+    - <encoder/decoder>_<left/right>_norms: List of norms applied before each
+                                            layer. Must be either "layer_norm"
+                                            or "none".
+    - <encoder/decoder>_combiner_functions: List of functions used to combine
+                                            each left/right branch pair.
+                                            Options are listed in
+                                            COMBINER_FUNCTIONS.
+    - <encoder/decoder>_final_combiner_function: Function applied to combine
+                                                 all the block outputs that are
+                                                 not used as inputs to other
+                                                 blocks. Options are listed in
+                                                 COMBINER_FUNCTIONS.
+
+  For an example of how to set these hparams, please see nas_seq2seq_base().
   """
   __metaclass__ = abc.ABCMeta
 
@@ -326,9 +370,9 @@ def _encoder(self,
                hparams,
                nonpadding=None,
                save_weights_to=None):
-    encoder_output, encoder_block_outputs = nas_encoder(
+    encoder_output, encoder_cell_outputs = nas_encoder(
         encoder_input, encoder_self_attention_bias, hparams, nonpadding)
-    self._encoder_block_outputs = encoder_block_outputs
+    self._encoder_cell_outputs = encoder_cell_outputs
     return encoder_output
 
   def _decoder(self,
@@ -340,8 +384,8 @@ def _decoder(self,
                cache=None,
                nonpadding=None,
                save_weights_to=None):
-    assert self._encoder_block_outputs
-    return nas_decoder(decoder_input, self._encoder_block_outputs,
+    assert self._encoder_cell_outputs
+    return nas_decoder(decoder_input, self._encoder_cell_outputs,
                        decoder_self_attention_bias,
                        encoder_decoder_attention_bias, hparams)
 
@@ -364,7 +408,7 @@ def _gpu_estimator_spec_eval(self, features, logits, labels, loss,
 
     if not hasattr(hparams, "problem"):
       raise NotImplementedError(
-          "hparams is missing attribute `problem`. Seq2SeqNasNet must "
+          "hparams is missing attribute `problem`. NasSeq2Seq must "
           "be used with a problem.")
 
     # TPU is not supported.
@@ -390,7 +434,7 @@ def _tpu_estimator_spec_eval(self, features, logits, labels, loss,
 
     if not hasattr(hparams, "problem"):
       raise NotImplementedError(
-          "hparams is missing attribute `problem`. Seq2SeqNasNet must "
+          "hparams is missing attribute `problem`. NasSeq2Seq must "
           "be used with a problem.")
 
     problem = hparams.problem
@@ -451,11 +495,12 @@ def _apply_layer_norm(input_tensor, nonpadding, hparams):
   return output_tensor
 
 
-def _apply_nas_branch(
-    norm, layer_norm_dict, hidden_states, nonpadding, hparams, input_index,
-    layer_name, activation_name, layer_registry, output_dim, branch_scope_name,
-    mask_future, dropout_broadcast_dims, encoder_decoder_attention_bias,
-    encoder_block_outputs, decoder_self_attention_bias, block_number):
+def _apply_nas_branch(norm, layer_norm_dict, hidden_states, nonpadding, hparams,
+                      input_index, layer_name, activation_name, layer_registry,
+                      output_dim, branch_scope_name, mask_future,
+                      dropout_broadcast_dims, encoder_decoder_attention_bias,
+                      encoder_cell_outputs, decoder_self_attention_bias,
+                      cell_number):
   """Applies a single NAS branch."""
   with tf.variable_scope(branch_scope_name):
     # Apply layer norm to an individual layer at most one time.
@@ -489,8 +534,8 @@ def _apply_nas_branch(
         nonpadding=nonpadding,
         attention_dropout_broadcast_dims=dropout_broadcast_dims,
         encoder_decoder_attention_bias=encoder_decoder_attention_bias,
-        encoder_block_outputs=encoder_block_outputs,
-        block_number=block_number,
+        encoder_cell_outputs=encoder_cell_outputs,
+        cell_number=cell_number,
         decoder_self_attention_bias=decoder_self_attention_bias)
 
     return output_tensor
@@ -509,18 +554,18 @@ def apply_nas_layers(input_tensor,
                      right_norms,
                      combiner_functions,
                      final_combiner_function,
-                     num_blocks,
+                     num_cells,
                      nonpadding,
                      layer_registry,
                      mask_future,
                      hparams,
                      var_scope,
                      encoder_decoder_attention_bias=None,
-                     encoder_block_outputs=None,
+                     encoder_cell_outputs=None,
                      decoder_self_attention_bias=None,
                      final_layer_norm=True,
                      enforce_fixed_output_sizes=True):
-  """Applies layers with NAS-like branching.
+  """Applies layers with NasNet search space style branching.
 
   Args:
     input_tensor: Input [batch_size, input_length, hidden_dim] sequence tensor.
@@ -536,8 +581,8 @@ def apply_nas_layers(input_tensor,
     right_norms: String list of right branch norms.
     combiner_functions: String list of branch combining functions.
     final_combiner_function: String. The final combiner function that combines
-      all the unused hidden layers in a block.
-    num_blocks: The number of blocks. This is the number of times the given
+      all the unused hidden layers in a cell.
+    num_cells: The number of cells. This is the number of times the given
       layers will be repeated.
     nonpadding: Tensor with 1s at all nonpadding time step positions and 0s
       everywhere else.
@@ -547,7 +592,7 @@ def apply_nas_layers(input_tensor,
     var_scope: The variable scope name.
     encoder_decoder_attention_bias: The attention bias for decoder attending to
       `encoder_output`.
-    encoder_block_outputs: List of tensors. The encoder block outputs, listed in
+    encoder_cell_outputs: List of tensors. The encoder cell outputs, listed in
       order.
     decoder_self_attention_bias: The self attention bias for decoders. This
       needs to be set for decoders.
@@ -563,7 +608,7 @@ def apply_nas_layers(input_tensor,
     ValueError: If item in right_norms is not LAYER_NORM_KEY or NO_NORM_KEY.
 
   Returns:
-    Output of applied layers and list of each block's outputs in order.
+    Output of applied layers and list of each cell's outputs in order.
   """
 
   if not (len(left_inputs) == len(left_layers) == len(left_activations) ==
@@ -572,7 +617,7 @@ def apply_nas_layers(input_tensor,
           == len(right_norms) == len(combiner_functions)):
     raise ValueError("All branching inputs must be of the same length.")
 
-  block_output = None
+  cell_output = None
   modified_left_inputs = [
       left_inputs[i]
       for i in range(len(left_inputs))
@@ -583,29 +628,29 @@ def apply_nas_layers(input_tensor,
       for i in range(len(right_inputs))
       if right_layers[i] != DEAD_BRANCH_KEY
   ]
-  unused_block_hidden_states = [
+  unused_cell_hidden_states = [
       i for i in range(len(left_inputs) + 1)
       if i not in modified_left_inputs and i not in modified_right_inputs
   ]
-  assert unused_block_hidden_states
+  assert unused_cell_hidden_states
 
-  block_outputs = []
+  cell_outputs = []
 
   with tf.variable_scope(var_scope):
     dropout_broadcast_dims = (
         common_layers.comma_separated_string_to_integer_list(
             getattr(hparams, "attention_dropout_broadcast_dims", "")))
 
-    for block_num in range(num_blocks):
+    for cell_num in range(num_cells):
       # h_0 is the input tensor.
       # Keep a dict for layer norm states.
-      if block_output is not None:
-        block_hidden_states = [block_output]
+      if cell_output is not None:
+        cell_hidden_states = [cell_output]
       else:
-        block_hidden_states = [input_tensor]
+        cell_hidden_states = [input_tensor]
       layer_norm_dict = {}
 
-      with tf.variable_scope("block_%d" % block_num):
+      with tf.variable_scope("cell_%d" % cell_num):
 
         for i, (left_input, left_layer_name, left_activation_name,
                 left_output_dim, left_norm, right_input, right_layer_name,
@@ -625,7 +670,7 @@ def apply_nas_layers(input_tensor,
 
             if left_layer_name != DEAD_BRANCH_KEY:
 
-              left_raw_input_tensor = block_hidden_states[left_input]
+              left_raw_input_tensor = cell_hidden_states[left_input]
               left_input_dim = left_raw_input_tensor.shape.as_list()[-1]
               if should_alter_output_dim(left_layer_name,
                                          enforce_fixed_output_sizes,
@@ -636,7 +681,7 @@ def apply_nas_layers(input_tensor,
               left_tensor = _apply_nas_branch(
                   norm=left_norm,
                   layer_norm_dict=layer_norm_dict,
-                  hidden_states=block_hidden_states,
+                  hidden_states=cell_hidden_states,
                   nonpadding=nonpadding,
                   hparams=hparams,
                   input_index=left_input,
@@ -648,12 +693,12 @@ def apply_nas_layers(input_tensor,
                   mask_future=mask_future,
                   dropout_broadcast_dims=dropout_broadcast_dims,
                   encoder_decoder_attention_bias=encoder_decoder_attention_bias,
-                  encoder_block_outputs=encoder_block_outputs,
+                  encoder_cell_outputs=encoder_cell_outputs,
                   decoder_self_attention_bias=decoder_self_attention_bias,
-                  block_number=block_num)
+                  cell_number=cell_num)
 
             if right_layer_name != DEAD_BRANCH_KEY:
-              right_raw_input_tensor = block_hidden_states[right_input]
+              right_raw_input_tensor = cell_hidden_states[right_input]
               right_input_dim = right_raw_input_tensor.shape.as_list()[-1]
               if should_alter_output_dim(right_layer_name,
                                          enforce_fixed_output_sizes,
@@ -663,7 +708,7 @@ def apply_nas_layers(input_tensor,
               right_tensor = _apply_nas_branch(
                   norm=right_norm,
                   layer_norm_dict=layer_norm_dict,
-                  hidden_states=block_hidden_states,
+                  hidden_states=cell_hidden_states,
                   nonpadding=nonpadding,
                   hparams=hparams,
                   input_index=right_input,
@@ -675,9 +720,9 @@ def apply_nas_layers(input_tensor,
                   mask_future=mask_future,
                   dropout_broadcast_dims=dropout_broadcast_dims,
                   encoder_decoder_attention_bias=encoder_decoder_attention_bias,
-                  encoder_block_outputs=encoder_block_outputs,
+                  encoder_cell_outputs=encoder_cell_outputs,
                   decoder_self_attention_bias=decoder_self_attention_bias,
-                  block_number=block_num)
+                  cell_number=cell_num)
 
             # Combine the branches.
             if left_layer_name == DEAD_BRANCH_KEY:
@@ -687,24 +732,24 @@ def apply_nas_layers(input_tensor,
             else:
               hidden_tensor = COMBINER_FUNCTIONS[combiner]().combine_tensors(
                   [left_tensor, right_tensor])
-            block_hidden_states.append(hidden_tensor)
+            cell_hidden_states.append(hidden_tensor)
 
       states_to_combine = [
-          block_hidden_states[j] for j in unused_block_hidden_states
+          cell_hidden_states[j] for j in unused_cell_hidden_states
       ]
-      block_output = COMBINER_FUNCTIONS[final_combiner_function](
+      cell_output = COMBINER_FUNCTIONS[final_combiner_function](
       ).combine_tensors(states_to_combine)
-      block_outputs.append(block_output)
+      cell_outputs.append(cell_output)
 
   if final_layer_norm:
-    final_output = common_layers.layer_preprocess(block_output, hparams)
-    block_outputs = [
-        common_layers.layer_preprocess(block_output, hparams)
-        for block_output in block_outputs
+    final_output = common_layers.layer_preprocess(cell_output, hparams)
+    cell_outputs = [
+        common_layers.layer_preprocess(cell_output, hparams)
+        for cell_output in cell_outputs
     ]
-    return final_output, block_outputs
+    return final_output, cell_outputs
   else:
-    return block_output, block_outputs
+    return cell_output, cell_outputs
 
 
 def nas_encoder(encoder_input,
@@ -712,7 +757,7 @@ def nas_encoder(encoder_input,
                 hparams,
                 nonpadding=None,
                 final_layer_norm=True):
-  """Encoder for NAS-style model.
+  """Encoder for configurable NAS model.
 
   Args:
     encoder_input: Input tensor.
@@ -731,7 +776,7 @@ def nas_encoder(encoder_input,
       + encoder_<left|right>_norms: String list of norms to apply to the
         <left|right> layer branches. Each item must be either LAYER_NORM_KEY or
         NO_NORM_KEY.
-      + encoder_num_blocks: The number of blocks in the encoder. This determines
+      + encoder_num_cells: The number of cells in the encoder. This determines
         how many times the given layers will be repeated.
       + encoder_combiner_functions: String list of functions used to combine
         left and right branches. Must be a COMBINER_FUNCTION key.
@@ -742,7 +787,7 @@ def nas_encoder(encoder_input,
       of the encoder.
 
   Returns:
-    Encoder output and list of each encoder block's output in order.
+    Encoder output and list of each encoder cell's output in order.
   """
   if nonpadding is None:
     padding = common_attention.attention_bias_to_padding(
@@ -760,7 +805,7 @@ def nas_encoder(encoder_input,
       right_activations=hparams.encoder_right_activations,
       right_output_dims=hparams.encoder_right_output_dims,
       right_norms=hparams.encoder_right_norms,
-      num_blocks=hparams.encoder_num_blocks,
+      num_cells=hparams.encoder_num_cells,
       combiner_functions=hparams.encoder_combiner_functions,
       final_combiner_function=hparams.encoder_final_combiner_function,
       nonpadding=nonpadding,
@@ -772,16 +817,16 @@ def nas_encoder(encoder_input,
 
 
 def nas_decoder(decoder_input,
-                encoder_block_outputs,
+                encoder_cell_outputs,
                 decoder_self_attention_bias,
                 encoder_decoder_attention_bias,
                 hparams,
                 final_layer_norm=True):
-  """Decoder for NAS-style model.
+  """Decoder for configurable model.
 
   Args:
     decoder_input: Input tensor.
-    encoder_block_outputs: List of tensors. The encoder block outputs, listed in
+    encoder_cell_outputs: List of tensors. The encoder cell outputs, listed in
       order.
     decoder_self_attention_bias: Attention bias that the decoder uses when
       attending to itself. This should have 0s for all valid positions and large
@@ -802,7 +847,7 @@ def nas_decoder(decoder_input,
       + decoder_<left|right>_norms: String list of norms to apply to the
         <left|right> layer branches. Each item must be either LAYER_NORM_KEY or
         NO_NORM_KEY.
-      + decoder_num_blocks: The number of blocks in the decoder. This determines
+      + decoder_num_cells: The number of cells in the decoder. This determines
         how many times the given layers will be repeated.
       + decoder_combiner_functions: String list of functions used to combine
         left and right branches. Must be a COMBINER_FUNCTION key.
@@ -818,8 +863,6 @@ def nas_decoder(decoder_input,
   Returns:
     Decoder output tensor.
   """
-  # encoder_depth is wrong because it doesn't matter here.
-
   # Enforce that the output tensor depth is equal to the depth of the encoding.
   (_, output_depth, _, _) = calculate_branching_model_parameters(
       encoding_depth=hparams.hidden_size,
@@ -832,7 +875,7 @@ def nas_decoder(decoder_input,
       combiner_functions=hparams.decoder_combiner_functions,
       final_combiner_function=hparams.decoder_final_combiner_function,
       layer_registry=layers.DECODER_LAYERS,
-      num_blocks=hparams.decoder_num_blocks,
+      num_cells=hparams.decoder_num_cells,
       encoder_depth=hparams.hidden_size)
   improper_output_size = output_depth != hparams.hidden_size
 
@@ -842,7 +885,7 @@ def nas_decoder(decoder_input,
     enforce_output_size = True
   resize_output = enforce_output_size and improper_output_size
 
-  decoder_blocks_output, _ = apply_nas_layers(
+  decoder_cells_output, _ = apply_nas_layers(
       input_tensor=decoder_input,
       left_inputs=hparams.decoder_left_inputs,
       left_layers=hparams.decoder_left_layers,
@@ -854,7 +897,7 @@ def nas_decoder(decoder_input,
       right_activations=hparams.decoder_right_activations,
       right_output_dims=hparams.decoder_right_output_dims,
       right_norms=hparams.decoder_right_norms,
-      num_blocks=hparams.decoder_num_blocks,
+      num_cells=hparams.decoder_num_cells,
       combiner_functions=hparams.decoder_combiner_functions,
       final_combiner_function=hparams.decoder_final_combiner_function,
       nonpadding=None,
@@ -864,16 +907,16 @@ def nas_decoder(decoder_input,
       var_scope="decoder",
       decoder_self_attention_bias=decoder_self_attention_bias,
       encoder_decoder_attention_bias=encoder_decoder_attention_bias,
-      encoder_block_outputs=encoder_block_outputs,
+      encoder_cell_outputs=encoder_cell_outputs,
       final_layer_norm=final_layer_norm)
 
   if not resize_output:
-    return decoder_blocks_output
+    return decoder_cells_output
 
   # Resize output if necessary.
   dense_layer = layers.DECODER_LAYERS.get(layers.STANDARD_CONV_1X1_REGISTRY_KEY)
   output = dense_layer.apply_layer(
-      decoder_blocks_output,
+      decoder_cells_output,
       None,
       hparams.hidden_size,
       None,
@@ -885,7 +928,7 @@ def nas_decoder(decoder_input,
       nonpadding=None,
       attention_dropout_broadcast_dims=None,
       encoder_decoder_attention_bias=None,
-      encoder_block_outputs=None,
+      encoder_cell_outputs=None,
       decoder_self_attention_bias=None,
   )
   if final_layer_norm:
@@ -903,7 +946,7 @@ def calculate_branching_model_parameters(encoding_depth,
                                          right_output_dims,
                                          combiner_functions,
                                          layer_registry,
-                                         num_blocks,
+                                         num_cells,
                                          final_combiner_function,
                                          encoder_depth=None,
                                          enforce_output_size=False,
@@ -926,7 +969,7 @@ def calculate_branching_model_parameters(encoding_depth,
       right branch tensors.
     layer_registry: layers.LayerRegistry. The LayerRegistry that contains the
       layers.TranslationLayers needed to construct the model.
-    num_blocks: Integer. The number of times the given layers are repeated to
+    num_cells: Integer. The number of times the given layers are repeated to
       produce the model.
     final_combiner_function: String. The COMBINER_FUNCTIONS key for the combiner
       used to combine the unused hidden dimensions.
@@ -944,11 +987,11 @@ def calculate_branching_model_parameters(encoding_depth,
 
   Returns:
     total_parameters: The total number of parameters in the model, accounting
-      for repeated blocks.
-    output_depth: The depth of the block output tensor.
+      for repeated cells.
+    output_depth: The depth of the cell output tensor.
     hidden_depths: The depths of the hidden layers.
     unused_outputs: List of integer indexes of the hidden layers that are not
-      used as input, and therefore are concatenated to produce the block
+      used as input, and therefore are concatenated to produce the cell
       output.
   """
   if not (len(left_inputs) == len(left_layers) == len(left_output_dims) ==
@@ -958,7 +1001,7 @@ def calculate_branching_model_parameters(encoding_depth,
 
   total_parameters = 0
   output_depth = encoding_depth
-  for _ in range(num_blocks):
+  for _ in range(num_cells):
     hidden_depths = [output_depth]
     unused_outputs = set(range(len(left_inputs) + 1))
 
@@ -1028,3 +1071,61 @@ def calculate_branching_model_parameters(encoding_depth,
             output_depth, encoding_depth, encoder_depth=encoder_depth)
 
   return (total_parameters, output_depth, hidden_depths, unused_outputs)
+
+
+@registry.register_hparams
+def nas_seq2seq_base():
+  """Base parameters for Nas Seq2Seq model.
+
+  The default parameters are set to create the Transformer.
+
+  Returns:
+    Hyperparameters for Nas Seq2Seq model.
+  """
+  hparams = transformer.transformer_base()
+
+  hparams.add_hparam("encoder_num_cells", 6)
+  hparams.add_hparam("encoder_left_inputs", [0, 1, 2, 3])
+  hparams.add_hparam("encoder_left_layers", [
+      "standard_attention", "standard_conv_1x1", "standard_conv_1x1", "identity"
+  ])
+  hparams.add_hparam("encoder_left_output_dims", [512, 2048, 512, 512])
+  hparams.add_hparam("encoder_left_activations",
+                     ["none", "relu", "none", "none"])
+  hparams.add_hparam("encoder_left_norms",
+                     ["layer_norm", "layer_norm", "none", "none"])
+  hparams.add_hparam("encoder_right_inputs", [0, 1, 1, 1])
+  hparams.add_hparam("encoder_right_layers",
+                     ["identity", "dead_branch", "identity", "dead_branch"])
+  hparams.add_hparam("encoder_right_activations",
+                     ["none", "none", "none", "none"])
+  hparams.add_hparam("encoder_right_output_dims", [512, 512, 512, 512])
+  hparams.add_hparam("encoder_right_norms", ["none", "none", "none", "none"])
+  hparams.add_hparam("encoder_combiner_functions", ["add", "add", "add", "add"])
+  hparams.add_hparam("encoder_final_combiner_function", "add")
+
+  hparams.add_hparam("decoder_num_cells", 6)
+  hparams.add_hparam("decoder_left_inputs", [0, 1, 2, 3, 4])
+  hparams.add_hparam("decoder_left_layers", [
+      "standard_attention", "attend_to_encoder", "standard_conv_1x1",
+      "standard_conv_1x1", "identity"
+  ])
+  hparams.add_hparam("decoder_left_activations",
+                     ["none", "none", "relu", "none", "none"])
+  hparams.add_hparam("decoder_left_output_dims", [512, 512, 2048, 512, 512])
+  hparams.add_hparam("decoder_left_norms",
+                     ["layer_norm", "layer_norm", "layer_norm", "none", "none"])
+  hparams.add_hparam("decoder_right_inputs", [0, 1, 2, 2, 4])
+  hparams.add_hparam(
+      "decoder_right_layers",
+      ["identity", "identity", "dead_branch", "identity", "dead_branch"])
+  hparams.add_hparam("decoder_right_activations",
+                     ["none", "none", "none", "none", "none"])
+  hparams.add_hparam("decoder_right_output_dims", [512, 512, 512, 512, 512])
+  hparams.add_hparam("decoder_right_norms",
+                     ["none", "none", "none", "none", "none"])
+  hparams.add_hparam("decoder_combiner_functions",
+                     ["add", "add", "add", "add", "add"])
+  hparams.add_hparam("decoder_final_combiner_function", "add")
+
+  return hparams
diff --git a/tensor2tensor/models/neural_architecture_search/nas_model_test.py b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
index c74415969..c5fd3775c 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model_test.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for Translation NasNets."""
+"""Tests for NasSeq2Seq."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -45,7 +45,7 @@ def _list_product(num_list):
 
 def _get_transformer_branching_encoder_config():
   """Returns config for the Transformer encoder."""
-  num_blocks = 2
+  num_cells = 2
   left_inputs = [0, 1, 2, 3]
   left_layers = [
       layers.STANDARD_ATTENTION_REGISTRY_KEY,
@@ -71,7 +71,7 @@ def _get_transformer_branching_encoder_config():
   is_decoder = False
   final_combiner_function = translation_nas_net.CONCAT_COMBINER_FUNC_KEY
 
-  return (num_blocks, left_inputs, left_layers, left_output_dims, right_inputs,
+  return (num_cells, left_inputs, left_layers, left_output_dims, right_inputs,
           right_layers, right_output_dims, combiner_functions,
           final_combiner_function, dummy_activations, dummy_norms,
           layer_registry, is_decoder)
@@ -79,7 +79,7 @@ def _get_transformer_branching_encoder_config():
 
 def _get_transformer_branching_decoder_config():
   """Returns config for the Transformer decoder."""
-  num_blocks = 2
+  num_cells = 2
   left_inputs = [0, 1, 2, 3, 4]
   left_layers = [
       layers.STANDARD_ATTENTION_REGISTRY_KEY,
@@ -108,14 +108,14 @@ def _get_transformer_branching_decoder_config():
   is_decoder = True
   final_combiner_function = translation_nas_net.CONCAT_COMBINER_FUNC_KEY
 
-  return (num_blocks, left_inputs, left_layers, left_output_dims, right_inputs,
+  return (num_cells, left_inputs, left_layers, left_output_dims, right_inputs,
           right_layers, right_output_dims, combiner_functions,
           final_combiner_function, dummy_activations, dummy_norms,
           layer_registry, is_decoder)
 
 
 def _add_transformer_branching_hparams(hparams):
-  (encoder_num_blocks, encoder_left_inputs, encoder_left_layers,
+  (encoder_num_cells, encoder_left_inputs, encoder_left_layers,
    encoder_left_output_dims, encoder_right_inputs, encoder_right_layers,
    encoder_right_output_dims, encoder_combiner_functions,
    encoder_final_combiner_function, encoder_dummy_activations,
@@ -133,11 +133,11 @@ def _add_transformer_branching_hparams(hparams):
   hparams.add_hparam("encoder_right_output_dims", encoder_right_output_dims)
   hparams.add_hparam("encoder_right_norms", encoder_dummy_norms)
   hparams.add_hparam("encoder_combiner_functions", encoder_combiner_functions)
-  hparams.add_hparam("encoder_num_blocks", encoder_num_blocks)
+  hparams.add_hparam("encoder_num_cells", encoder_num_cells)
   hparams.add_hparam("encoder_final_combiner_function",
                      encoder_final_combiner_function)
 
-  (decoder_num_blocks, decoder_left_inputs, decoder_left_layers,
+  (decoder_num_cells, decoder_left_inputs, decoder_left_layers,
    decoder_left_output_dims, decoder_right_inputs, decoder_right_layers,
    decoder_right_output_dims, decoder_combiner_functions,
    decoder_final_combiner_function, decoder_dummy_activations,
@@ -155,12 +155,12 @@ def _add_transformer_branching_hparams(hparams):
   hparams.add_hparam("decoder_right_output_dims", decoder_right_output_dims)
   hparams.add_hparam("decoder_right_norms", decoder_dummy_norms)
   hparams.add_hparam("decoder_combiner_functions", decoder_combiner_functions)
-  hparams.add_hparam("decoder_num_blocks", decoder_num_blocks)
+  hparams.add_hparam("decoder_num_cells", decoder_num_cells)
   hparams.add_hparam("decoder_final_combiner_function",
                      decoder_final_combiner_function)
 
 
-class TranslationNasNetTest(parameterized.TestCase, tf.test.TestCase):
+class NasSeq2SeqTest(parameterized.TestCase, tf.test.TestCase):
 
   def _test_model(self, model_cls, hparams):
     """Test a Translation Nas Net model."""
@@ -233,7 +233,7 @@ def test_nas_decoder_resizing_output(self):
     with tf.variable_scope("wrong"):
       wrong_size_decoder_output = translation_nas_net.nas_decoder(
           decoder_input=input_tensor,
-          encoder_block_outputs=[input_tensor] * hparams.encoder_num_blocks,
+          encoder_cell_outputs=[input_tensor] * hparams.encoder_num_cells,
           decoder_self_attention_bias=decoder_self_attention_bias,
           encoder_decoder_attention_bias=None,
           hparams=hparams)
@@ -243,7 +243,7 @@ def test_nas_decoder_resizing_output(self):
     with tf.variable_scope("correct"):
       correct_size_decoder_output = translation_nas_net.nas_decoder(
           decoder_input=input_tensor,
-          encoder_block_outputs=[input_tensor] * hparams.encoder_num_blocks,
+          encoder_cell_outputs=[input_tensor] * hparams.encoder_num_cells,
           decoder_self_attention_bias=decoder_self_attention_bias,
           encoder_decoder_attention_bias=None,
           hparams=hparams)
@@ -265,7 +265,7 @@ def test_calculate_branching_model_parameters_transformer(
       self, get_config, expected_hidden_depths):
     tf.reset_default_graph()
 
-    (num_blocks, left_inputs, left_layers, left_output_dims, right_inputs,
+    (num_cells, left_inputs, left_layers, left_output_dims, right_inputs,
      right_layers, right_output_dims, combiner_functions,
      final_combiner_function, dummy_activations, dummy_norms, layer_registry,
      is_decoder) = get_config()
@@ -283,7 +283,7 @@ def test_calculate_branching_model_parameters_transformer(
          combiner_functions=combiner_functions,
          final_combiner_function=final_combiner_function,
          layer_registry=layer_registry,
-         num_blocks=num_blocks,
+         num_cells=num_cells,
          encoder_depth=_EMBEDDING_DEPTH)
 
     # Create model graph.
@@ -295,12 +295,12 @@ def test_calculate_branching_model_parameters_transformer(
       mask_future = True
       decoder_self_attention_bias = (
           common_attention.attention_bias_lower_triangle(_INPUT_LENGTH))
-      encoder_block_outputs = [input_tensor] * 6
+      encoder_cell_outputs = [input_tensor] * 6
     else:
       nonpadding = tf.ones([32, _INPUT_LENGTH])
       mask_future = False
       decoder_self_attention_bias = None
-      encoder_block_outputs = None
+      encoder_cell_outputs = None
 
     translation_nas_net.apply_nas_layers(
         input_tensor=input_tensor,
@@ -316,14 +316,14 @@ def test_calculate_branching_model_parameters_transformer(
         right_norms=dummy_norms,
         combiner_functions=combiner_functions,
         final_combiner_function=final_combiner_function,
-        num_blocks=num_blocks,
+        num_cells=num_cells,
         nonpadding=nonpadding,
         layer_registry=layer_registry,
         mask_future=mask_future,
         hparams=hparams,
         var_scope="test",
         encoder_decoder_attention_bias=None,
-        encoder_block_outputs=encoder_block_outputs,
+        encoder_cell_outputs=encoder_cell_outputs,
         decoder_self_attention_bias=decoder_self_attention_bias,
         final_layer_norm=False)
 
@@ -361,7 +361,7 @@ def test_calculate_branching_model_parameters_decoder_resize(
          combiner_functions=hparams.decoder_combiner_functions,
          final_combiner_function=hparams.decoder_final_combiner_function,
          layer_registry=layers.DECODER_LAYERS,
-         num_blocks=hparams.decoder_num_blocks,
+         num_cells=hparams.decoder_num_cells,
          encoder_depth=_EMBEDDING_DEPTH,
          enforce_output_size=enforce_output_size)
 
@@ -371,7 +371,7 @@ def test_calculate_branching_model_parameters_decoder_resize(
         common_attention.attention_bias_lower_triangle(_INPUT_LENGTH))
     _ = translation_nas_net.nas_decoder(
         decoder_input=input_tensor,
-        encoder_block_outputs=[input_tensor] * hparams.encoder_num_blocks,
+        encoder_cell_outputs=[input_tensor] * hparams.encoder_num_cells,
         decoder_self_attention_bias=decoder_self_attention_bias,
         encoder_decoder_attention_bias=None,
         hparams=hparams,
@@ -399,7 +399,7 @@ def test_calculate_branching_model_parameters_output_size_only_final(self):
         translation_nas_net.CONCAT_COMBINER_FUNC_KEY
     ]
 
-    (num_blocks, _, left_layers, _, _, _, _, _, final_combiner_function,
+    (num_cells, _, left_layers, _, _, _, _, _, final_combiner_function,
      dummy_activations, dummy_norms, layer_registry,
      _) = _get_transformer_branching_encoder_config()
 
@@ -416,7 +416,7 @@ def test_calculate_branching_model_parameters_output_size_only_final(self):
          combiner_functions=combiner_functions,
          final_combiner_function=final_combiner_function,
          layer_registry=layer_registry,
-         num_blocks=num_blocks,
+         num_cells=num_cells,
          encoder_depth=_EMBEDDING_DEPTH,
          enforce_output_size=False,
          enforce_fixed_output_sizes=False)
@@ -439,7 +439,7 @@ def test_calculate_branching_model_parameters_output_size_last_two(self):
         translation_nas_net.CONCAT_COMBINER_FUNC_KEY
     ]
 
-    (num_blocks, _, left_layers, _, _, _, _, _, final_combiner_function,
+    (num_cells, _, left_layers, _, _, _, _, _, final_combiner_function,
      dummy_activations, dummy_norms, layer_registry,
      _) = _get_transformer_branching_encoder_config()
 
@@ -456,7 +456,7 @@ def test_calculate_branching_model_parameters_output_size_last_two(self):
          combiner_functions=combiner_functions,
          final_combiner_function=final_combiner_function,
          layer_registry=layer_registry,
-         num_blocks=num_blocks,
+         num_cells=num_cells,
          encoder_depth=_EMBEDDING_DEPTH,
          enforce_output_size=False,
          enforce_fixed_output_sizes=False)

From d7f02686034e57b41f2460a08e6e489620243d26 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 11 Jun 2019 07:56:38 -0700
Subject: [PATCH 2123/2720] Add some combinator tests.

PiperOrigin-RevId: 252618826
---
 tensor2tensor/trax/layers/combinators_test.py | 78 +++++++++++++++++--
 1 file changed, 71 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/trax/layers/combinators_test.py b/tensor2tensor/trax/layers/combinators_test.py
index 4a2598db4..6f987d960 100644
--- a/tensor2tensor/trax/layers/combinators_test.py
+++ b/tensor2tensor/trax/layers/combinators_test.py
@@ -20,30 +20,94 @@
 
 from absl.testing import absltest
 from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import combinators
+from tensor2tensor.trax.layers import combinators as cb
+from tensor2tensor.trax.layers import core
+
+
+_EMPTY_STACK = ()
+_REST_OF_STACK = ((1, 5), (4,))
 
 
 class CombinatorLayerTest(absltest.TestCase):
 
+  def test_drop(self):
+    layer = cb.Drop()
+    input_shape = ((3, 2),)
+    expected_shape = _EMPTY_STACK
+    output_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+    input_shape = ((3, 2),) + _REST_OF_STACK
+    expected_shape = _REST_OF_STACK
+    output_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+  def test_dup(self):
+    layer = cb.Dup()
+    input_shape = ((3, 2),)
+    expected_shape = ((3, 2), (3, 2))
+    output_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+    input_shape = ((3, 2),) + _REST_OF_STACK
+    expected_shape = ((3, 2), (3, 2)) + _REST_OF_STACK
+    output_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+  def test_swap(self):
+    layer = cb.Swap()
+    input_shape = ((3, 2), (4, 7))
+    expected_shape = ((4, 7), (3, 2))
+    output_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+    input_shape = ((3, 2), (4, 7)) + _REST_OF_STACK
+    expected_shape = ((4, 7), (3, 2)) + _REST_OF_STACK
+    output_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+  def test_serial_no_op_list(self):
+    layer = cb.Serial([])
+    input_shape = ((3, 2), (4, 7))
+    expected_shape = ((3, 2), (4, 7))
+    output_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+    input_shape = ((3, 2), (4, 7)) + _REST_OF_STACK
+    expected_shape = ((3, 2), (4, 7)) + _REST_OF_STACK
+    output_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+  def test_serial_one_in_one_out(self):
+    layer = cb.Serial(core.Div(divisor=2.0))
+    input_shape = ((3, 2), (4, 7))
+    expected_shape = ((3, 2), (4, 7))
+    output_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+  def test_serial_div_div(self):
+    layer = cb.Serial(core.Div(divisor=2.0), core.Div(divisor=5.0))
+    input_shape = ((3, 2), (4, 7))
+    expected_shape = ((3, 2), (4, 7))
+    output_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
   def test_branch(self):
     input_shape = (2, 3)
     expected_shape = ((2, 3), (2, 3))
-    output_shape = base.check_shape_agreement(
-        combinators.Branch([], []), input_shape)
+    output_shape = base.check_shape_agreement(cb.Branch([], []), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
   def test_parallel(self):
     input_shape = ((2, 3), (2, 3))
     expected_shape = ((2, 3), (2, 3))
-    output_shape = base.check_shape_agreement(
-        combinators.Parallel([], []), input_shape)
+    output_shape = base.check_shape_agreement(cb.Parallel([], []), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
   def test_select(self):
     input_shape = ((2, 3), (3, 4))
     expected_shape = (3, 4)
-    output_shape = base.check_shape_agreement(
-        combinators.Select(1), input_shape)
+    output_shape = base.check_shape_agreement(cb.Select(1), input_shape)
     self.assertEqual(output_shape, expected_shape)
 
 

From 404335e94b3f70c1c07792d13b032925b96c59aa Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 12 Jun 2019 11:22:09 -0700
Subject: [PATCH 2124/2720] Factor out small rng function from combinator call
 methods.

PiperOrigin-RevId: 252863011
---
 tensor2tensor/trax/layers/combinators.py | 29 +++++++++---------------
 1 file changed, 11 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index a00f86e8e..9a499c1a9 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -75,6 +75,13 @@ def _ensure_sublayers(layers):  # pylint: disable=invalid-name
     raise TypeError(type(layers))
 
 
+def _pop_rng_and_split(args_dict, n_copies):  # pylint: disable=invalid-name
+  rng = args_dict.pop('rng', None)
+  if rng is None:
+    return (None,) * n_copies
+  return backend.random.split(rng, n_copies)
+
+
 class Serial(base.Layer):
   """Layer composing a number of sub-layers in a serial way.."""
 
@@ -86,10 +93,7 @@ def __init__(self, *layers):
     self._nlayers = len(layers)
 
   def call(self, x, params=(), **kwargs):
-    rng = kwargs.pop('rng', None)
-    rngs = (None,) * self._nlayers
-    if rng is not None:
-      rngs = backend.random.split(rng, self._nlayers)
+    rngs = _pop_rng_and_split(kwargs, self._nlayers)
     for layer, p, rng in zip(self._layers, params, rngs):
       x = layer(x, p, rng=rng, **kwargs)
     return x
@@ -255,11 +259,7 @@ def __init__(self, *layers):
     self._layers = layers
 
   def call(self, x, params=(), **kwargs):
-    # Split the random number generators.
-    rng = kwargs.pop('rng', None)
-    rngs = (None,) * self._nlayers
-    if rng is not None:
-      rngs = backend.random.split(rng, self._nlayers)
+    rngs = _pop_rng_and_split(kwargs, self._nlayers)
     if isinstance(self._layers, (list, tuple)):
       res = [layer(x, params=p, rng=r, **kwargs)
              for layer, p, r in zip(self._layers, params, rngs)]
@@ -386,11 +386,7 @@ def stack_items_to_pass(self):
     return self._nlayers
 
   def call(self, inputs, params=(), **kwargs):
-    # Split the random number generators.
-    rng = kwargs.pop('rng', None)
-    rngs = (None,) * self._nlayers
-    if rng is not None:
-      rngs = backend.random.split(rng, self._nlayers)
+    rngs = _pop_rng_and_split(kwargs, self._nlayers)
     # If layers are a list or a tuple, just apply them.
     if not isinstance(self._layers, dict):
       res = [layer(x, params=p, rng=r, **kwargs)
@@ -457,10 +453,7 @@ def __init__(self, layer, check_shapes=True):
     self._check_shapes = check_shapes
 
   def call(self, inputs, params=(), **kwargs):
-    rng = kwargs.pop('rng', None)
-    rngs = (None,) * len(inputs)
-    if rng is not None:
-      rngs = backend.random.split(rng, len(inputs))
+    rngs = _pop_rng_and_split(kwargs, len(inputs))
     result = [self._layer(x, params=params, rng=r, **kwargs)
               for x, r in zip(inputs, rngs)]
     if isinstance(inputs, list):

From ea130ba72831b9047da498696fa4022476ecbe81 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 12 Jun 2019 14:18:04 -0700
Subject: [PATCH 2125/2720] Rename "num_axis_to_keep" to "n_axes_to_keep".

PiperOrigin-RevId: 252897836
---
 tensor2tensor/trax/layers/core.py      | 10 +++++-----
 tensor2tensor/trax/layers/core_test.py | 10 +++++-----
 tensor2tensor/trax/models/atari_cnn.py |  2 +-
 tensor2tensor/trax/rlax/ppo_test.py    |  6 +++---
 4 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index cb314a526..97afa10cb 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -135,13 +135,13 @@ def new_parameters(self, input_shape, input_dtype, rng):
 
 # Flatten.
 @base.layer()
-def Flatten(x, params, num_axis_to_keep=1, **kwargs):
+def Flatten(x, params, n_axes_to_keep=1, **kwargs):
   del params, kwargs
-  if num_axis_to_keep >= len(x.shape):
+  if n_axes_to_keep >= len(x.shape):
     raise ValueError(
-        "num_axis_to_keep[%d] should be less than input's rank[%d]" %
-        (num_axis_to_keep, len(x.shape)))
-  return np.reshape(x, (x.shape[:num_axis_to_keep] + (-1,)))
+        "n_axes_to_keep[%d] should be less than input's rank[%d]" %
+        (n_axes_to_keep, len(x.shape)))
+  return np.reshape(x, (x.shape[:n_axes_to_keep] + (-1,)))
 
 
 @base.layer()
diff --git a/tensor2tensor/trax/layers/core_test.py b/tensor2tensor/trax/layers/core_test.py
index 0331e8440..4f076f5fd 100644
--- a/tensor2tensor/trax/layers/core_test.py
+++ b/tensor2tensor/trax/layers/core_test.py
@@ -35,23 +35,23 @@ def test_flatten_n(self):
     self.assertEqual(actual_shape, (29, 87 * 10 * 20 * 30))
 
     actual_shape = base.check_shape_agreement(
-        core.Flatten(num_axis_to_keep=2), input_shape)
+        core.Flatten(n_axes_to_keep=2), input_shape)
     self.assertEqual(actual_shape, (29, 87, 10 * 20 * 30))
 
     actual_shape = base.check_shape_agreement(
-        core.Flatten(num_axis_to_keep=3), input_shape)
+        core.Flatten(n_axes_to_keep=3), input_shape)
     self.assertEqual(actual_shape, (29, 87, 10, 20 * 30))
 
     actual_shape = base.check_shape_agreement(
-        core.Flatten(num_axis_to_keep=4), input_shape)
+        core.Flatten(n_axes_to_keep=4), input_shape)
     self.assertEqual(actual_shape, (29, 87, 10, 20, 30))
 
     # Not enough dimensions.
     with self.assertRaises(base.LayerError):
-      base.check_shape_agreement(core.Flatten(num_axis_to_keep=5), input_shape)
+      base.check_shape_agreement(core.Flatten(n_axes_to_keep=5), input_shape)
 
     with self.assertRaises(base.LayerError):
-      base.check_shape_agreement(core.Flatten(num_axis_to_keep=6), input_shape)
+      base.check_shape_agreement(core.Flatten(n_axes_to_keep=6), input_shape)
 
   def test_div(self):
     layer = core.Div(divisor=2.0)
diff --git a/tensor2tensor/trax/models/atari_cnn.py b/tensor2tensor/trax/models/atari_cnn.py
index 29cb3aaa1..f26dd6cc1 100644
--- a/tensor2tensor/trax/models/atari_cnn.py
+++ b/tensor2tensor/trax/models/atari_cnn.py
@@ -41,7 +41,7 @@ def AtariCnn(hidden_sizes=(32, 32), output_size=128):
       tl.Relu(),
       tl.Conv(hidden_sizes[1], (5, 5), (2, 2), 'SAME'),
       tl.Relu(),
-      tl.Flatten(num_axis_to_keep=2),  # B, T and rest.
+      tl.Flatten(n_axes_to_keep=2),  # B, T and rest.
       tl.Dense(output_size),
       tl.Relu(),
       # Eventually this is shaped (B, T, output_size)
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index c8a879d70..bf39fbba3 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -39,7 +39,7 @@ def test_policy_and_value_net(self):
     n_actions = 2
     pnv_params, pnv_apply = ppo.policy_and_value_net(
         self.rng_key, batch_observation_shape, np.float32, n_actions,
-        lambda: [layers.Flatten(num_axis_to_keep=2)])
+        lambda: [layers.Flatten(n_axes_to_keep=2)])
     batch = 2
     time_steps = 10
     batch_of_observations = np.random.uniform(
@@ -384,11 +384,11 @@ def test_combined_loss(self):
 
     old_params, _ = ppo.policy_and_value_net(
         key1, batch_observation_shape, np.float32, A,
-        lambda: [layers.Flatten(num_axis_to_keep=2)])
+        lambda: [layers.Flatten(n_axes_to_keep=2)])
 
     new_params, net_apply = ppo.policy_and_value_net(
         key2, batch_observation_shape, np.float32, A,
-        lambda: [layers.Flatten(num_axis_to_keep=2)])
+        lambda: [layers.Flatten(n_axes_to_keep=2)])
 
     # Generate a batch of observations.
 

From 95232123e3e5829732e101e37e8e3c017d02e680 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 12 Jun 2019 17:16:34 -0700
Subject: [PATCH 2126/2720] Transformer with position lookups and testing
 data-set.

PiperOrigin-RevId: 252932428
---
 tensor2tensor/data_generators/algorithmic.py  |  21 ++
 .../position_lookup_transformer_copy.gin      |  45 +++
 tensor2tensor/trax/inputs.py                  |  22 +-
 tensor2tensor/trax/layers/combinators.py      |   5 +-
 tensor2tensor/trax/models/__init__.py         |   3 +
 .../research/position_lookup_transformer.py   | 353 ++++++++++++++++++
 6 files changed, 445 insertions(+), 4 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
 create mode 100644 tensor2tensor/trax/models/research/position_lookup_transformer.py

diff --git a/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index ee59b486f..f5458f2ec 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -131,6 +131,27 @@ def num_symbols(self):
     return 10
 
 
+@registry.register_problem
+class AlgorithmicIdentityVocab95Train20Eval30(AlgorithmicIdentityBinary40):
+  """Problem spec for algorithmic decimal identity task."""
+
+  @property
+  def num_symbols(self):
+    return 95
+
+  @property
+  def train_length(self):
+    return 20
+
+  @property
+  def dev_length(self):
+    return 30
+
+  @property
+  def train_size(self):
+    return 1000000
+
+
 @registry.register_problem
 class AlgorithmicShiftDecimal40(AlgorithmicProblem):
   """Problem spec for algorithmic decimal shift task."""
diff --git a/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin b/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
new file mode 100644
index 000000000..d9ea5abd6
--- /dev/null
+++ b/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
@@ -0,0 +1,45 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 32
+batch_fun.eval_batch_size = 32
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.dataset_name = 't2t_algorithmic_identity_vocab95_train20_eval30'
+
+# Parameters for masked_mean:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.concat_preprocess
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 0.05
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 16000
+
+# Parameters for PositionLookupTransformerLM:
+# ==============================================================================
+PositionLookupTransformerLM.d_feature = 256
+PositionLookupTransformerLM.d_feedforward = 512
+PositionLookupTransformerLM.dropout = 0.01
+PositionLookupTransformerLM.max_len = 100
+PositionLookupTransformerLM.n_heads = 4
+PositionLookupTransformerLM.n_layers = 3
+PositionLookupTransformerLM.vocab_size = 128
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 1000
+train.eval_steps = 10
+train.model = @trax.models.PositionLookupTransformerLM
+train.optimizer = @trax.optimizers.Adam
+train.train_steps = 100000
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 943d40978..cc0d2b236 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -355,9 +355,9 @@ def example_length(example_inputs, target):
   return dataset
 
 
-# pylint: disable=unused-argument
 @gin.configurable(blacklist=["dataset", "training"])
 def cifar10_no_augmentation_preprocess(dataset, training):
+  del training
 
   def cast_image(features, targets):
     features["image"] = tf.cast(features["image"], tf.float32) / 255.0
@@ -367,8 +367,26 @@ def cast_image(features, targets):
   return dataset
 
 
-# pylint: disable=unused-argument
 def no_preprocess(dataset, training):
+  del training
+  return dataset
+
+
+@gin.configurable(blacklist=["dataset", "training"])
+def concat_preprocess(dataset, training, pad_symbol=0):
+  """Pre-processing function that concatenates input and target for LM."""
+  del training
+
+  def concat(features, targets):
+    inp = features["inputs"]
+    pad = tf.expand_dims(tf.zeros_like(inp[0]) + pad_symbol, axis=0)
+    concat = tf.concat([pad, inp, pad, targets], axis=0)
+    # Note: we're updating existing features dictionary here, so make sure
+    # it is not re-used in some other ways outside of this function.
+    features["inputs"] = concat
+    return features, concat
+
+  dataset = dataset.map(concat)
   return dataset
 
 
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 9a499c1a9..94d914b51 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -134,8 +134,9 @@ def Dup(x, **unused_kwargs):
   """Duplicate (copy) the first element on the stack."""
   if isinstance(x, list):
     return [x[0]] + x
-  assert isinstance(x, tuple)
-  return tuple([x[0]] + list(x))
+  if isinstance(x, tuple):
+    return tuple([x[0]] + list(x))
+  return [x, x]
 
 
 @base.layer(stack_items_to_pass=0)
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index 2a0117725..6dcdbe3b3 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -25,6 +25,7 @@
 from tensor2tensor.trax.models import resnet
 from tensor2tensor.trax.models import transformer
 from tensor2tensor.trax.models.research import chunked_transformer
+from tensor2tensor.trax.models.research import position_lookup_transformer
 
 
 # Ginify
@@ -37,6 +38,8 @@ def model_configure(*args, **kwargs):
 ChunkedTransformerLM = model_configure(chunked_transformer.ChunkedTransformerLM)
 MLP = model_configure(mlp.MLP)
 NeuralGPU = model_configure(neural_gpu.NeuralGPU)
+PositionLookupTransformerLM = model_configure(
+    position_lookup_transformer.PositionLookupTransformerLM)
 Resnet50 = model_configure(resnet.Resnet50)
 Transformer = model_configure(transformer.Transformer)
 TransformerEncoder = model_configure(transformer.TransformerEncoder)
diff --git a/tensor2tensor/trax/models/research/position_lookup_transformer.py b/tensor2tensor/trax/models/research/position_lookup_transformer.py
new file mode 100644
index 000000000..9d55aa8c7
--- /dev/null
+++ b/tensor2tensor/trax/models/research/position_lookup_transformer.py
@@ -0,0 +1,353 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Deep Lookups for Transformer Positions."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as onp
+
+from tensor2tensor.trax import layers as tl
+from tensor2tensor.trax.backend import numpy as np
+
+
+# pylint: disable=g-complex-comprehension
+# pylint: disable=no-value-for-parameter
+
+POS_VECTOR_SIZE = 32
+_ABSOLUTE_MAX_LEN = 10000
+_POSITIONS = onp.random.uniform(size=[_ABSOLUTE_MAX_LEN, POS_VECTOR_SIZE])
+
+
+@tl.layer()
+def NewPositionalEncoding(x, positions=None, **kwargs):
+  """Implements new positional encoding."""
+  del kwargs
+  x_length = np.shape(x)[1]
+  pos = np.array(positions)[np.newaxis, :x_length, :]
+  pos += np.zeros((np.shape(x)[0], 1, 1))  # Broadcast on batch.
+  res = np.concatenate([x, pos], axis=2)
+  return res
+
+
+@tl.layer(stack_items_to_pass=0)
+def CutPosition(xs, **unused_kwargs):
+  """Splits x into a pair (x[:position], position)."""
+  if not isinstance(xs, (list, tuple)):
+    xs = [xs]
+  x = xs[0]
+  res = [x[:, :, :-POS_VECTOR_SIZE], x[:, :, -POS_VECTOR_SIZE:]]
+  return tuple(res + list(xs[1:]))
+
+
+@tl.layer()
+def MixHeadsPos(x, h=8, **unused_kwargs):
+  """Mix x = (x0, p) into x0_h1, p, x0_h2, p, ...."""
+  head_size = (x.shape[2] - POS_VECTOR_SIZE) // h
+  p = x[:, :, -POS_VECTOR_SIZE:]
+  res, idx = [], 0
+  for _ in range(h):
+    res.append(x[:, :, idx:idx+head_size])
+    res.append(p)
+    idx += head_size
+  return np.concatenate(res, axis=-1)
+
+
+@tl.layer()
+def CombineHeadsPos(x, h=8, **unused_kwargs):
+  """Mix x = (x0, p0, ..., xH, pH) into x0, ...., xH, p_combined.
+
+  The positions are added as vectors.
+
+  Args:
+    x: input vector, concatenated (x0, p0, ..., xH, pH).
+    h: number of heads.
+
+  Returns:
+    the vector with combined positions.
+  """
+  head_size = int((x.shape[2] / h) - POS_VECTOR_SIZE)
+  res, positions, idx = [], [], 0
+  for _ in range(h):
+    res.append(x[:, :, idx:idx+head_size])
+    idx += head_size
+    positions.append(x[:, :, idx:idx+POS_VECTOR_SIZE])
+    idx += POS_VECTOR_SIZE
+  combined_position = sum(positions)
+  res.append(combined_position)
+  return np.concatenate(res, axis=-1)
+
+
+@tl.layer()
+def CopyHeadsPos(x, h=8, **unused_kwargs):
+  """Mix x = (x, p) into x_h1, p_h1, x_h2, p_h2, ...."""
+  head_size = (x.shape[2] - h*POS_VECTOR_SIZE) // h
+  p = x[:, :, -h*POS_VECTOR_SIZE:]
+  res, idx = [], 0
+  for i in range(h):
+    res.append(x[:, :, idx:idx+head_size])
+    res.append(p[:, :, i*POS_VECTOR_SIZE:(i+1)*POS_VECTOR_SIZE])
+    idx += head_size
+  return np.concatenate(res, axis=-1)
+
+
+def DeepFlatten(xs):
+  for x in xs:
+    if isinstance(x, (list, tuple)):
+      for y in DeepFlatten(x):
+        yield y
+    else:
+      yield x
+
+
+@tl.layer(stack_items_to_pass=0)
+def Unnest(xs, **unused_kwargs):
+  return [x for x in DeepFlatten(xs)]
+
+
+@tl.layer(stack_items_to_pass=0)
+def ConcatenateN(xs, params, n=2, axis=-1, **kwargs):
+  """Concatenate first N inputs (and output remainder as is if non-empty)."""
+  del params, kwargs
+  res = np.concatenate(xs[:n], axis)
+  rest = list(xs[n:])
+  if rest:
+    return tuple([res] + rest)
+  return res
+
+
+def PreservePosition(layer):
+  """Execute layer without position but preserve it in parallel."""
+  return tl.Serial(
+      CutPosition(),
+      layer,
+      ConcatenateN()
+  )
+
+
+def ApplyAndQueryPositions(layer, pos):
+  """Execute layer without position and pos-layers on positions.
+
+  This takes an embedding including position x = (emb, p), and
+  outputs layer(emb).pos1(x, p).....layer(emb).posn(x, p)
+  where pos=[pos1...posn].
+
+  Args:
+    layer: layer to be executed without position information.
+    pos: list of layers to be applied to positions.
+
+  Returns:
+    the result of this application.
+  """
+  n_heads = len(pos)
+  return tl.Serial(
+      tl.Dup(),
+      CutPosition(),
+      tl.Select(tuple([0] + [(2, 1)]*n_heads)),
+      tl.Parallel(*([layer] + pos)),
+      Unnest(),
+      ConcatenateN(n=n_heads + 1)
+  )
+
+
+@tl.layer()
+def QueryPositionKV(x, keys=None, values=None, binary=False, **unused_kwargs):
+  """Query a table with a position vector."""
+  if keys is None:
+    return x
+  k = np.array(keys)
+  v = np.array(values)
+  q = x
+  if binary:
+    q = np.concatenate([x, x], axis=-1)
+  return tl.DotProductAttention(q, k, v, None, None, None, None)
+
+
+def LearnedQP(keys=None, values=None, binary=False):
+  """Get (query, pos), make learned weight of qeury and return with pos."""
+  return tl.Parallel(
+      tl.Dense(1),
+      QueryPositionKV(keys=keys, values=values, binary=binary),
+  )
+
+
+@tl.layer(stack_items_to_pass=0)
+def SoftmaxBranches(x_list_in, n_branches=2, **unused_kwargs):
+  """Softmax xs.
+
+  The input xs is a list of embeddings and weights of the form
+  w_1 e_1 .... w_n e_n (followed by optional rest that is preserved).
+
+  Args:
+    x_list_in: the input weights and embeddings.
+    n_branches: what part of the list to use.
+
+  Returns:
+    softmax(w) * e for the joint weights w and embeddings e.
+  """
+  x_list, x_list_rest = x_list_in[:2*n_branches], x_list_in[2*n_branches:]
+  softmax_activations = [x_list[2*i] for i in range(n_branches)]
+  max_sa = softmax_activations[0]
+  for x in softmax_activations:
+    max_sa = np.maximum(max_sa, x)
+  softmax_activations = [x - max_sa for x in softmax_activations]
+  softmax_activations = [np.exp(x) for x in softmax_activations]
+  sum_sa = sum(softmax_activations)
+  softmax_activations = [x / sum_sa for x in softmax_activations]
+  res = sum([x_list[2*i+1] * softmax_activations[i] for i in range(n_branches)])
+  return tuple([res] + list(x_list_rest))
+
+
+def SumLearnedPick(positions):
+  """Get a pair (vec, pos) and pick new pos."""
+  succ_keys = positions[:-1, :]
+  succ_values = positions[1:, :]
+  subtract_1_keys = positions[1:, :]
+  subtract_1_values = positions[:-1, :]
+  l = int(positions.shape[0]) // 2
+  add_keys = np.array([np.concatenate([positions[i, :], positions[j, :]])
+                       for i in range(l) for j in range(l)])
+  add_values = np.array([positions[i + j, :]
+                         for i in range(l) for j in range(l)])
+  # TODO(lukaszkaiser): try this below: "for j in range(i) for i in range(2*l)"
+  sub_keys = np.array([np.concatenate([positions[i, :], positions[j, :]])
+                       for j in range(l) for i in range(l)])
+  sub_values = np.array([positions[max(i - j, 0), :]
+                         for j in range(l) for i in range(l)])
+  return tl.Serial(
+      tl.Branch(
+          LearnedQP(),
+          LearnedQP(keys=succ_keys, values=succ_values),
+          LearnedQP(keys=subtract_1_keys, values=subtract_1_values),
+          LearnedQP(keys=add_keys, values=add_values, binary=True),
+          LearnedQP(keys=sub_keys, values=sub_values, binary=True),
+      ),
+      Unnest(),
+      SoftmaxBranches(n_branches=5)
+  )
+
+
+def MultiHeadedAttentionPosition(
+    positions, d_feature, n_heads=8, dropout=0.0, mode='train'):
+  """Transformer-style multi-headed attention."""
+  return tl.Serial(
+      tl.Dup(),
+      tl.Dup(),
+      tl.Parallel(
+          ApplyAndQueryPositions(tl.Dense(d_feature),
+                                 pos=[SumLearnedPick(positions)
+                                      for _ in range(n_heads)]),
+          PreservePosition(tl.Dense(d_feature)),
+          PreservePosition(tl.Dense(d_feature)),
+      ),
+      tl.Parallel(
+          CopyHeadsPos(h=n_heads),
+          MixHeadsPos(h=n_heads),
+          MixHeadsPos(h=n_heads),
+      ),
+      tl.PureMultiHeadedAttention(
+          d_feature=d_feature, n_heads=n_heads,
+          dropout=dropout, mode=mode),
+      tl.Select(0),  # Drop the mask.
+      CombineHeadsPos(h=n_heads),
+      PreservePosition(tl.Dense(d_feature)),
+  )
+
+
+def ResidualFeedForward(d_feature,
+                        d_feedforward,
+                        dropout,
+                        mode):
+  """Residual feed-forward layer with normalization at start."""
+  stack = tl.Serial(
+      tl.LayerNorm(),
+      tl.Dense(d_feedforward),
+      tl.Relu(),
+      tl.Dropout(rate=dropout, mode=mode),
+      tl.Dense(d_feature),
+      tl.Dropout(rate=dropout, mode=mode)
+  )
+  return tl.Residual(PreservePosition(stack))
+
+
+def DecoderLayer(positions,
+                 d_feature,
+                 d_feedforward,
+                 n_heads,
+                 dropout,
+                 mode):
+  """Transformer decoder layer.
+
+  Args:
+    positions: random vectors for positions
+    d_feature: int:  depth of embedding
+    d_feedforward: int: depth of feed-forward layer
+    n_heads: int: number of attention heads
+    dropout: float: dropout rate (how much to drop out)
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    the layer.
+  """
+  return [
+      tl.Residual(  # Self-attention block.
+          PreservePosition(tl.LayerNorm()),
+          tl.Branch([],  # activation for (q, k, v)
+                    tl.CausalMask(axis=-2)),  # attention mask
+          MultiHeadedAttentionPosition(positions,
+                                       d_feature, n_heads=n_heads,
+                                       dropout=dropout, mode=mode),
+          PreservePosition(tl.Dropout(rate=dropout, mode=mode))
+      ),
+      ResidualFeedForward(d_feature, d_feedforward, dropout, mode=mode)
+  ]
+
+
+def PositionLookupTransformerLM(vocab_size=128,
+                                d_feature=256,
+                                d_feedforward=512,
+                                n_layers=3,
+                                n_heads=4,
+                                dropout=0.1,
+                                max_len=100,
+                                mode='train'):
+  """Transformer language model (only uses the decoder part of Transformer).
+
+  Args:
+    vocab_size: int: vocab size
+    d_feature: int:  depth of embedding
+    d_feedforward: int: depth of feed-forward layer
+    n_layers: int: number of layers
+    n_heads: int: number of attention heads
+    dropout: float: dropout rate (how much to drop out)
+    max_len: maximal length
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    the layer.
+  """
+  positions = _POSITIONS[:max_len, :]
+  return tl.Serial([
+      tl.ShiftRight(),
+      tl.Embedding(d_feature, vocab_size),
+      tl.Dropout(rate=dropout, mode=mode),
+      NewPositionalEncoding(positions=positions),
+      [DecoderLayer(positions, d_feature, d_feedforward, n_heads, dropout, mode)
+       for _ in range(n_layers)],
+      PreservePosition(tl.LayerNorm()),
+      tl.Dense(vocab_size),
+      tl.LogSoftmax()
+  ])

From ba1c763bf62d219ff7b756b71f051f046a9a4550 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 14 Jun 2019 09:37:20 -0700
Subject: [PATCH 2127/2720] Move trax training code into a class (first step of
 refactoring).

PiperOrigin-RevId: 253243119
---
 tensor2tensor/trax/trax.py | 370 ++++++++++++++++++++++---------------
 1 file changed, 217 insertions(+), 153 deletions(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 10b14a777..f904b5a60 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -140,7 +140,7 @@ def restore_state(output_dir):
   return State(step=step, params=params, history=history)
 
 
-def save_gin(output_dir, sw=None):
+def _save_gin(output_dir, sw=None):
   config_path = os.path.join(output_dir, "config.gin")
   config_str = gin.operative_config_str()
   with gfile.GFile(config_path, "w") as f:
@@ -181,6 +181,17 @@ def _save_replicated(opt_state, step, history, n_devices, output_dir, keep):
              output_dir, keep=keep)
 
 
+def _print_n_params(opt_state, n_devices, step):
+  """Print out the number of parameters."""
+  sizes = layers.sizes(opt_state[0])
+  if n_devices > 1:
+    unreplicate = lambda x: x.mean(0)
+    single_params = layers.nested_map(opt_state[0], unreplicate)
+    sizes = layers.sizes(single_params)
+  total_size = layers.nested_reduce(sizes, sum)
+  step_log(step, "Total trainable parameters size: %d" % total_size)
+
+
 # Metrics to calculate and report.
 _METRICS = {
     "accuracy": accuracy,
@@ -382,6 +393,190 @@ def reshape_by_device(x, n_devices):
       x, lambda x: _reshape_by_device_single(x, n_devices))
 
 
+class Trainer(object):
+  """Trax trainer.
+
+  A trainer allows to make training steps, train for full epochs,
+  save the training state and access evaluation data.
+  """
+
+  def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs, output_dir,
+               random_seed=None, n_devices=None, save_steps=None):
+    if save_steps is None:
+      save_steps = []
+    self._save_steps = save_steps
+    device_count = jax.lib.xla_bridge.device_count()
+    n_devices = n_devices or device_count
+    # TODO(lukaszkaiser): remove this restriction when possible.
+    if n_devices != device_count:
+      raise ValueError("Jax cannot work yet with n_devices != all devices: "
+                       "%d != %d" % (n_devices, device_count))
+    self._n_devices = n_devices
+    rng = get_random_number_generator_and_set_seed(random_seed)
+    self._output_dir = output_dir
+    gfile.makedirs(output_dir)
+    # Create summary writers and history.
+    self._train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
+    self._eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
+
+    # Create input streams.
+    inputs = inputs(n_devices)
+    self._inputs = inputs
+    self._train_stream = inputs.train_stream()
+
+    # Setup optimizer and model.
+    state = restore_state(output_dir)
+    history = state.history
+    self._lr_fn = lr_schedule(history)
+    opt = optimizer(self._lr_fn)
+
+    model_train = model(mode="train")
+    model_predict_eval = model(mode="eval")
+
+    # Setup state.
+    step = state.step or 0
+    rng, init_rng = jax_random.split(rng)
+    self._rngs = jax_random.split(rng, n_devices)
+    first_shape = inputs.input_shape[0]
+    # If the inputs are a tuple/list, add [None] (batch) to each element.
+    if isinstance(first_shape, (list, tuple)):
+      model_input_shape = tuple(
+          tuple([None] + list(shape)) for shape in inputs.input_shape)
+    else:  # Otherwise just add [None] to the input shape.
+      model_input_shape = tuple([None] + list(inputs.input_shape))
+    # Change all None to 1 in input shape.
+    model_input_shape = layers.nested_map(
+        model_input_shape, lambda x: x if x else 1)
+    if state.params:
+      params = state.params[0]
+      opt_state = state.params
+    else:
+      params = model_train.initialize(
+          model_input_shape, inputs.input_dtype, init_rng)
+      opt_state = (params, opt.tree_init(params))
+    if n_devices > 1:
+      replicate = lambda x: numpy.broadcast_to(x, (n_devices,) + x.shape)
+      opt_state = layers.nested_map(opt_state, replicate)
+
+    # jit model_predict and update so they're fast
+    self._jit_model_predict_eval = _jit_predict_fn(
+        model_predict_eval, n_devices)
+    self._jit_update_fn = _jit_update_fn(model_train, loss_fn, opt, n_devices)
+
+    self._step = step
+    self._model_train = model_train
+    self._model_predict_eval = model_predict_eval
+    self._loss_fn = loss_fn
+    self._optimizer = optimizer
+    self._opt_state = opt_state
+    self._history = history
+    self._lr_schedule = lr_schedule
+
+  @property
+  def step(self):
+    return self._step
+
+  @property
+  def n_devices(self):
+    return self._n_devices
+
+  @property
+  def state(self):
+    return State(params=self._opt_state, step=self._step, history=self._history)
+
+  def save_gin(self):
+    _save_gin(self._output_dir, self._train_sw)
+
+  def print_n_params(self):
+    _print_n_params(self._opt_state, self._n_devices, self._step)
+
+  def train_epoch(self, epoch_steps, eval_steps):
+    """Train for one epoch."""
+    # Log separator
+    print()
+
+    # Timer
+    start_time = time.time()
+
+    for _ in range(epoch_steps):
+      # Train
+      next_train_batch = next(self._train_stream)
+      if self._n_devices > 1:  # TODO(lukaszkaiser): use everywhere if possible.
+        next_train_batch = reshape_by_device(next_train_batch, self._n_devices)
+      self._opt_state, self._rngs = self._jit_update_fn(
+          self._step, self._opt_state, next_train_batch, self._rngs)
+      self._step += 1
+
+      if self._step in self._save_steps:
+        _save_replicated(self._opt_state, self._step, self._history,
+                         self._n_devices, self._output_dir, True)
+
+      # LR log
+      if self._step == 1 or self._step % 10 == 0:
+        self._train_sw.scalar("training/learning rate",
+                              self._lr_fn(self._step), step=self._step)
+
+    # Timer
+    epoch_time = time.time() - start_time
+    step_log(self._step, "Ran %d train steps in %0.2f secs" %
+             (epoch_steps, epoch_time))
+    if epoch_steps > 1:
+      self._train_sw.scalar("training/steps per second",
+                            epoch_steps / epoch_time, step=self._step)
+
+    # Evaluate in parallel
+    _, rng = jax_random.split(self._rngs[0])
+    evaluate_train_and_eval(
+        step=self._step,
+        inputs=self._inputs,
+        predict_fn=functools.partial(self._jit_model_predict_eval,
+                                     params=self._opt_state[0]),
+        eval_steps=eval_steps,
+        rng=rng,
+        train_sw=self._train_sw,
+        eval_sw=self._eval_sw,
+        history=self._history)
+
+    # Save state
+    _save_replicated(self._opt_state, self._step, self._history,
+                     self._n_devices, self._output_dir, False)
+
+    # Flush summary writers
+    self._train_sw.flush()
+    self._eval_sw.flush()
+
+  def update_learning_rate(self):
+    old_lr_fn = self._lr_fn
+    self._lr_fn = self._lr_schedule(self._history)
+    if self._lr_fn != old_lr_fn:  # For performance only jit if it's changed.
+      opt = self._optimizer(self._lr_fn)
+      self._jit_update_fn = _jit_update_fn(
+          self._model_train, self._loss_fn, opt, self._n_devices)
+
+  def save_computation_graphs(self, save_backward_graph):
+    """Dump computation graphs to files."""
+    if self._n_devices != 1:
+      return  # TODO(lukaszkaiser): make this work with more devices.
+    next_train_batch = next(self._train_stream)
+    output_dir = self._output_dir
+    if self._n_devices > 1:
+      next_train_batch = reshape_by_device(next_train_batch, self._n_devices)
+    params = self._opt_state[0]
+    forward_computation = jax.xla_computation(self._model_predict_eval)(
+        next_train_batch[0], params=params, rng=self._rngs[0])
+    with gfile.GFile(os.path.join(output_dir, "forward.txt"), "w") as f:
+      f.write(forward_computation.GetHloText())
+    with gfile.GFile(os.path.join(output_dir, "forward.dot"), "w") as f:
+      f.write(forward_computation.GetHloDotGraph())
+    backward_computation = jax.xla_computation(self._jit_update_fn)(
+        self._step, self._opt_state, next_train_batch, self._rngs)
+    with gfile.GFile(os.path.join(output_dir, "backward.txt"), "w") as f:
+      f.write(backward_computation.GetHloText())
+    if save_backward_graph:  # Backward graphs can be large so we guard it.
+      with gfile.GFile(os.path.join(output_dir, "backward.dot"), "w") as f:
+        f.write(backward_computation.GetHloDotGraph())
+
+
 @gin.configurable(blacklist=["output_dir"])
 def train(output_dir,
           model=gin.REQUIRED,
@@ -395,7 +590,6 @@ def train(output_dir,
           eval_frequency=100,
           n_devices=None,
           random_seed=None,
-          run_debug_step=False,
           save_graphs=True,
           save_backward_graph=False):
   """Train the model on the inputs.
@@ -418,170 +612,40 @@ def train(output_dir,
       steps). If None or 0, eval disabled.
     n_devices: how many devices to use (if None, default, use all available)
     random_seed: the random seed to use; time/os dependent if None (default).
-    run_debug_step: bool, if True, will run the model and loss without @jit for
-      one step.
     save_graphs: bool, if True, save computation graph to file.
     save_backward_graph: bool, if True, save backward graph to file too.
   Returns:
     trax.State
   """
-  if save_steps is None:
-    save_steps = []
-  device_count = jax.lib.xla_bridge.device_count()
-  n_devices = n_devices or device_count
-  # TODO(lukaszkaiser): remove this restriction when possible.
-  if n_devices != device_count:
-    raise ValueError("Jax cannot work yet with n_devices != all devices: "
-                     "%d != %d" % (n_devices, device_count))
-  rng = get_random_number_generator_and_set_seed(random_seed)
-  gfile.makedirs(output_dir)
-  # Create summary writers and history.
-  train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
-  eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
-
-  inputs = inputs(n_devices)
-
-  # Setup optimizer and model
-  state = restore_state(output_dir)
-  history = state.history
-  lr_fn = lr_schedule(history)
-  opt = optimizer(lr_fn)
-
-  model_train = model(mode="train")
-  model_predict_eval = model(mode="eval")
-
-  # Setup state
-  step = state.step or 0
-  rng, init_rng = jax_random.split(rng)
-  rngs = jax_random.split(rng, n_devices)
-  first_shape = inputs.input_shape[0]
-  # If the inputs are a tuple/list, add [None] (batch) to each element.
-  if isinstance(first_shape, (list, tuple)):
-    model_input_shape = tuple(
-        tuple([None] + list(shape)) for shape in inputs.input_shape)
-  else:  # Otherwise just add [None] to the input shape.
-    model_input_shape = tuple([None] + list(inputs.input_shape))
-  # Change all None to 1 in input shape.
-  model_input_shape = layers.nested_map(
-      model_input_shape, lambda x: x if x else 1)
-  if state.params:
-    params = state.params[0]
-    opt_state = state.params
-  else:
-    params = model_train.initialize(
-        model_input_shape, inputs.input_dtype, init_rng)
-    opt_state = (params, opt.tree_init(params))
-  if n_devices > 1:
-    replicate = lambda x: numpy.broadcast_to(x, (n_devices,) + x.shape)
-    opt_state = layers.nested_map(opt_state, replicate)
+  trainer = Trainer(model, loss_fn, optimizer, lr_schedule, inputs, output_dir,
+                    random_seed=random_seed, n_devices=n_devices,
+                    save_steps=save_steps)
 
-  # jit model_predict and update so they're fast
-  jit_model_predict_eval = _jit_predict_fn(model_predict_eval, n_devices)
-  jit_update_fn = _jit_update_fn(model_train, loss_fn, opt, n_devices)
-
-  train_stream = inputs.train_stream()
-  epoch_steps = [train_steps]  # Only training if eval_frequency is 0 or None.
+  epoch_steps = [train_steps]  # Only training if eval_frequency is 0 or None
   if eval_frequency and eval_steps > 0:
     epoch_steps = itertools.chain([1,  # first epoch only 1 step
                                    eval_frequency - 1],
                                   itertools.repeat(eval_frequency))
-  step_log(step, "Starting training using %d devices" % n_devices)
-
-  # Non-compiled debug step helps find problems in models easier.
-  if run_debug_step:
-    debug_loss = loss_fn(params, next(train_stream), model_train, rng)
-    step_log(step, "Debug step loss %.8f" % debug_loss)
-
-  for epoch, epoch_steps in epochs(train_steps, epoch_steps):
-    # Log separator
-    print()
-
-    # Timer
-    start_time = time.time()
+  step_log(trainer.step,
+           "Starting training using %d devices" % trainer.n_devices)
 
-    for _ in range(epoch_steps):
-      # Train
-      next_train_batch = next(train_stream)
-      if n_devices > 1:  # TODO(lukaszkaiser): use everywhere when possible.
-        next_train_batch = reshape_by_device(next_train_batch, n_devices)
-      opt_state, rngs = jit_update_fn(step, opt_state, next_train_batch, rngs)
-      step += 1
+  for _, epoch_steps in epochs(train_steps, epoch_steps):
+    trainer.train_epoch(epoch_steps, eval_steps)
 
-      if step in save_steps:
-        _save_replicated(opt_state, step, history, n_devices, output_dir, True)
-
-      # LR log
-      if step == 1 or step % 10 == 0:
-        train_sw.scalar("training/learning rate",
-                        lr_fn(step), step=step)
-
-    # Timer
-    epoch_time = time.time() - start_time
-    step_log(step, "Ran %d train steps in %0.2f secs" %
-             (epoch_steps, epoch_time))
-    if epoch_steps > 1:
-      train_sw.scalar("training/steps per second",
-                      epoch_steps / epoch_time, step=step)
-
-    # Print number of parameters
-    if step == 1:
-      sizes = layers.sizes(opt_state[0])
-      if n_devices > 1:
-        unreplicate = lambda x: x.mean(0)
-        single_params = layers.nested_map(opt_state[0], unreplicate)
-        sizes = layers.sizes(single_params)
-      total_size = layers.nested_reduce(sizes, sum)
-      step_log(step, "Total trainable parameters size: %d" % total_size)
-
-    # Evaluate in parallel
-    evaluate_train_and_eval(
-        step=step,
-        inputs=inputs,
-        predict_fn=functools.partial(jit_model_predict_eval,
-                                     params=opt_state[0]),
-        eval_steps=eval_steps,
-        rng=rng,
-        train_sw=train_sw,
-        eval_sw=eval_sw,
-        history=history)
-
-    # Save computation graph (single-device only for now).
-    if (save_graphs and backend.get_name() == "jax" and step == 1 and
-        n_devices == 1):
-      params = opt_state[0]
-      # Dump computation graphs to files.
-      forward_computation = jax.xla_computation(model_predict_eval)(
-          next_train_batch[0], params=params, rng=rng)
-      with gfile.GFile(os.path.join(output_dir, "forward.txt"), "w") as f:
-        f.write(forward_computation.GetHloText())
-      with gfile.GFile(os.path.join(output_dir, "forward.dot"), "w") as f:
-        f.write(forward_computation.GetHloDotGraph())
-      backward_computation = jax.xla_computation(jit_update_fn)(
-          step, opt_state, next_train_batch, rngs)
-      with gfile.GFile(os.path.join(output_dir, "backward.txt"), "w") as f:
-        f.write(backward_computation.GetHloText())
-      if save_backward_graph:  # Backward graphs can be large so we guard it.
-        with gfile.GFile(os.path.join(output_dir, "backward.dot"), "w") as f:
-          f.write(backward_computation.GetHloDotGraph())
-
-    # Save state
-    _save_replicated(opt_state, step, history, n_devices, output_dir, False)
+    # Update learning rate with new history
+    trainer.update_learning_rate()
 
-    # Save Gin config
-    # Gin only tracks the used parameters, so we save it after the first epoch.
-    if epoch == 1:
-      save_gin(output_dir, train_sw)
+    # Bookkeeping we do at the first step
+    if trainer.step == 1:
+      # Print number of parameters
+      trainer.print_n_params()
 
-    # Update learning rate with new history
-    old_lr_fn = lr_fn
-    lr_fn = lr_schedule(history)
-    if lr_fn != old_lr_fn:  # For performance, only jit if there is a change.
-      opt = optimizer(lr_fn)
-      jit_update_fn = _jit_update_fn(model_train, loss_fn, opt, n_devices)
+      # Save computation graph (single-device only for now)
+      if (save_graphs and backend.get_name() == "jax"):
+        trainer.save_computation_graphs(save_backward_graph)
 
-    # Flush summary writers
-    train_sw.flush()
-    eval_sw.flush()
+      # Save Gin config
+      trainer.save_gin()
 
-  step_log(step, "Training done")
-  return State(params=opt_state, step=step, history=history)
+  step_log(trainer.step, "Training done")
+  return trainer.state

From c4578f84b80723701029c0eb55921833225f64aa Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 14 Jun 2019 10:15:03 -0700
Subject: [PATCH 2128/2720] Remove chunked Transformer from research (as we
 took another route) and the related Map combinator.

PiperOrigin-RevId: 253249938
---
 tensor2tensor/trax/layers/combinators.py      |  40 ---
 tensor2tensor/trax/models/__init__.py         |   2 -
 .../models/research/chunked_transformer.py    | 248 ------------------
 3 files changed, 290 deletions(-)
 delete mode 100644 tensor2tensor/trax/models/research/chunked_transformer.py

diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 94d914b51..0f0bd72db 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -429,43 +429,3 @@ def Residual(*layers, **kwargs):
       FlattenList(),  # pylint: disable=no-value-for-parameter
       Add(),  # pylint: disable=no-value-for-parameter
   ]
-
-
-class Map(base.Layer):
-  """Combinator for applying a layer to a list or tuple.
-
-  Args:
-    layer: a layer to apply to each element.
-
-  Returns:
-    A new layer representing mapping layer to all elements of the input.
-  """
-
-  def __init__(self, layer, check_shapes=True):
-    super(Map, self).__init__()
-    if layer is None or isinstance(layer, (list, tuple)):
-      layer = Serial(layer)
-    self._layer = layer
-    # Generally a Map should be applied to lists where all elements have
-    # the same shape -- because self._layer will only be initialized once
-    # and it could have different parameters for different shapes. But there
-    # are valid cases -- e.g., when self._layer has no parameters -- where we
-    # can apply Map to different shapes -- set check_shapes=False in such cases.
-    self._check_shapes = check_shapes
-
-  def call(self, inputs, params=(), **kwargs):
-    rngs = _pop_rng_and_split(kwargs, len(inputs))
-    result = [self._layer(x, params=params, rng=r, **kwargs)
-              for x, r in zip(inputs, rngs)]
-    if isinstance(inputs, list):
-      return result
-    return tuple(result)
-
-  def new_parameters(self, input_shape, input_dtype, rng):
-    first_shape = input_shape[0]
-    if self._check_shapes:
-      for shape in input_shape:
-        if shape != first_shape:
-          raise ValueError('Map layer can only be applied to list of elements '
-                           'with the same shapes. Shapes: %s' % str(shape))
-    return self._layer.initialize(first_shape, input_dtype[0], rng)
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index 6dcdbe3b3..ad1d27b5c 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -24,7 +24,6 @@
 from tensor2tensor.trax.models import neural_gpu
 from tensor2tensor.trax.models import resnet
 from tensor2tensor.trax.models import transformer
-from tensor2tensor.trax.models.research import chunked_transformer
 from tensor2tensor.trax.models.research import position_lookup_transformer
 
 
@@ -35,7 +34,6 @@ def model_configure(*args, **kwargs):
 
 
 # pylint: disable=invalid-name
-ChunkedTransformerLM = model_configure(chunked_transformer.ChunkedTransformerLM)
 MLP = model_configure(mlp.MLP)
 NeuralGPU = model_configure(neural_gpu.NeuralGPU)
 PositionLookupTransformerLM = model_configure(
diff --git a/tensor2tensor/trax/models/research/chunked_transformer.py b/tensor2tensor/trax/models/research/chunked_transformer.py
deleted file mode 100644
index 491728cd6..000000000
--- a/tensor2tensor/trax/models/research/chunked_transformer.py
+++ /dev/null
@@ -1,248 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Chunked Transformer Models."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as onp
-
-from tensor2tensor.trax import layers as tl
-from tensor2tensor.trax.backend import numpy as np
-
-
-# Chunked positional encoding.
-def _chunked_positional_encoding_new_params(  # pylint: disable=invalid-name
-    input_shape, input_dtype, rng, max_len=2048):
-  """Helper: create positional encoding parameters."""
-  del input_dtype, rng
-  # Check if we are operating on chunked inputs by checking if the first
-  # shape is a list/tuple of shapes (otherwise it's an int or numpy array).
-  is_chunked = isinstance(input_shape[0], (list, tuple))
-  d_feature = input_shape[0][-1] if is_chunked else input_shape[-1]
-  pe = onp.zeros((max_len, d_feature), dtype=onp.float32)
-  position = onp.arange(0, max_len)[:, onp.newaxis]
-  div_term = onp.exp(
-      onp.arange(0, d_feature, 2) * -(onp.log(10000.0) / d_feature))
-  pe[:, 0::2] = onp.sin(position * div_term)
-  pe[:, 1::2] = onp.cos(position * div_term)
-  pe = pe[onp.newaxis, :, :]  # [1, max_len, d_feature]
-  return np.array(pe)  # These are trainable parameters, initialized as above.
-
-
-@tl.layer(new_parameters=_chunked_positional_encoding_new_params,
-          stack_items_to_pass=0)
-def ChunkedPositionalEncoding(x, params, **unused_kwargs):
-  """Implements bare positional encoding."""
-  if not isinstance(x, (list, tuple)):  # non-chunked inputs
-    symbol_size = np.shape(x)[1]
-    return x + params[:, :symbol_size, :]
-  # Chunked case: apply to all chunks selecting as much as needed.
-  offset = 0
-  results = []
-  for chunk in x:
-    symbol_size = np.shape(chunk)[1]
-    results.append(chunk + params[:, offset:offset + symbol_size, :])
-    offset += symbol_size
-  return results
-
-
-# Chunked attention.
-@tl.layer(stack_items_to_pass=0)
-def ChunkedAttentionSelector(x, params, selector=None, **kwargs):
-  """Select which chunks to attend to in chunked attention.
-
-  Args:
-    x: inputs, a list of elements of the form (q, k, v), mask for each chunk.
-    params: parameters (unused).
-    selector: a function from chunk_number -> list of chunk numbers that says
-      which other chunks should be appended to the given one (previous if None).
-    **kwargs: unused other arguments.
-
-  Returns:
-    a list of elements of the form (q, k', v', mask') where k', v' and mask' are
-    concatenations of k, v and identity-extended masks from selected chunks.
-  """
-  del params, kwargs
-  selector = selector or (lambda x: [] if x < 1 else [x-1])
-  triples, masks = zip(*x)
-  (queries, keys, values) = zip(*triples)
-  result = []
-  for i in range(len(x)):
-    selected = selector(i)
-    # Since keys and values are [batch, length, depth] we concatenate on axis=1.
-    # We also always include the current key or value at the end.
-    new_key_list = [keys[j] for j in selected]
-    new_key = np.concatenate(new_key_list + [keys[i]], axis=1)
-    new_value = np.concatenate(
-        [values[j] for j in selected] + [values[i]], axis=1)
-    # Masks are (1, query-len, key-len) so we concatenate on axis=2.
-    new_mask_shapes = [(1, queries[i].shape[1], key.shape[1])
-                       for key in new_key_list]
-    cur_mask = masks[i]
-    # Masks are all-1 for the added chunks (no masking).
-    new_mask_list = [np.ones(s, dtype=cur_mask.dtype) for s in new_mask_shapes]
-    # We still use the current (often causal) mask for the final chunk.
-    new_mask = np.concatenate(new_mask_list + [cur_mask], axis=2)
-    result.append((queries[i], new_key, new_value, new_mask))
-  return tuple(result)
-
-
-def ChunkedCausalMultiHeadedAttention(
-    d_feature, n_heads=8, dropout=0.0, chunk_selector=None, mode='train'):
-  """Transformer-style causal multi-headed attention operating on chunks.
-
-  Accepts inputs that are a list of chunks and applies causal attention.
-
-  Args:
-    d_feature: int:  depth of embedding
-    n_heads: int: number of attention heads
-    dropout: float: dropout rate
-    chunk_selector: a function from chunk number to list of chunks to attend.
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    Multi-headed self-attention layer.
-  """
-  prepare_attention_input = [
-      tl.Branch(
-          tl.Branch([], [], []),  # q = k = v = first input
-          tl.CausalMask(axis=-2)
-      ),
-      tl.Parallel(
-          tl.Parallel(
-              tl.Dense(d_feature),
-              tl.Dense(d_feature),
-              tl.Dense(d_feature)
-          ),
-          []
-      )
-  ]
-  return [
-      tl.Map(prepare_attention_input),
-      ChunkedAttentionSelector(selector=chunk_selector),  # pylint: disable=no-value-for-parameter
-      tl.Map(tl.PureMultiHeadedAttention(d_feature=d_feature, n_heads=n_heads,
-                                         dropout=dropout, mode=mode),
-             check_shapes=False),
-      tl.Map(tl.Select(0), check_shapes=False),  # drop masks
-      tl.Map(tl.Dense(d_feature))
-  ]
-
-
-# Chunked residual.
-def Residual(*layers, **unused_kwargs):
-  """Constructs a residual version of layers, summing input to layers output."""
-  return [
-      tl.Branch(layers, []),
-      tl.AddAll()
-  ]
-
-
-def ResidualFeedForward(d_feature, d_feedforward, dropout, mode):
-  """Residual feed-forward layer with normalization at start."""
-  return Residual(
-      tl.LayerNorm(),
-      tl.Dense(d_feedforward),
-      tl.Relu(),
-      tl.Dropout(rate=dropout, mode=mode),
-      tl.Dense(d_feature),
-      tl.Dropout(rate=dropout, mode=mode)
-  )
-
-
-def ChunkedDecoderLayer(d_feature,
-                        d_feedforward,
-                        n_heads,
-                        dropout,
-                        chunk_selector,
-                        mode):
-  """Transformer decoder layer operating on chunks.
-
-  Args:
-    d_feature: int:  depth of embedding
-    d_feedforward: int: depth of feed-forward layer
-    n_heads: int: number of attention heads
-    dropout: float: dropout rate (how much to drop out)
-    chunk_selector: a function from chunk number to list of chunks to attend.
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    The layers comprising a chunked decoder.
-  """
-  return [
-      Residual(  # Self-attention block.
-          tl.Map(tl.LayerNorm()),
-          ChunkedCausalMultiHeadedAttention(
-              d_feature, n_heads=n_heads, dropout=dropout,
-              chunk_selector=chunk_selector, mode=mode),
-          tl.Map(tl.Dropout(rate=dropout, mode=mode)),
-      ),
-      tl.Map(ResidualFeedForward(
-          d_feature, d_feedforward, dropout, mode=mode))
-  ]
-
-
-def ChunkedTransformerLM(vocab_size,
-                         d_feature=512,
-                         d_feedforward=2048,
-                         n_layers=6,
-                         n_heads=8,
-                         dropout=0.1,
-                         chunk_selector=None,
-                         max_len=2048,
-                         mode='train'):
-  """Transformer language model operating on chunks.
-
-  The input to this  model is a sequence presented as a list or tuple of chunks:
-    (chunk1, chunk2, chunks3, ..., chunkN).
-  Each chunk should have the same shape (batch, chunk-length) and together they
-  represent a long sequence that's a concatenation chunk1,chunk2,...,chunkN.
-
-  Chunked Transformer emulates the operation of a Transformer on this long
-  sequence except for the chunked attention layer, which may attend to only
-  a subset of the chunks to reduce memory use.
-
-  Args:
-    vocab_size: int: vocab size
-    d_feature: int:  depth of embedding
-    d_feedforward: int: depth of feed-forward layer
-    n_layers: int: number of encoder/decoder layers
-    n_heads: int: number of attention heads
-    dropout: float: dropout rate (how much to drop out)
-    chunk_selector: a function from chunk number to list of chunks to attend
-      (if None, attends to the previous chunks which is equivalent to setting
-       chunk_selector(x) = [] if x < 1 else [x-1] (TransformerXL); we attend
-       to the current chunk with a causal mask too, selected chunks unmasked).
-    max_len: int: maximum symbol length for positional encoding
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    the layer.
-  """
-  decoder_stack = [ChunkedDecoderLayer(d_feature, d_feedforward, n_heads,
-                                       dropout, chunk_selector, mode)
-                   for _ in range(n_layers)]
-  # Below each Map(L) applies the layer L to each chunk independently.
-  return tl.Model(
-      tl.ShiftRight(),
-      tl.Map(tl.Embedding(d_feature, vocab_size)),
-      tl.Map(tl.Dropout(rate=dropout, mode=mode)),
-      ChunkedPositionalEncoding(max_len=max_len),  # pylint: disable=no-value-for-parameter
-      decoder_stack,
-      tl.Map(tl.LayerNorm()),
-      tl.Map(tl.Dense(vocab_size)),
-      tl.Map(tl.LogSoftmax()),
-  )

From 91e7bcfe8df281b682b40b1488e1b9b8fdd455bf Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 14 Jun 2019 16:21:17 -0700
Subject: [PATCH 2129/2720] update the reacher env test.

PiperOrigin-RevId: 253316740
---
 tensor2tensor/envs/mujoco_problems_test.py | 19 ++-----------------
 1 file changed, 2 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/envs/mujoco_problems_test.py b/tensor2tensor/envs/mujoco_problems_test.py
index 8d90f616a..5a66f5cd1 100644
--- a/tensor2tensor/envs/mujoco_problems_test.py
+++ b/tensor2tensor/envs/mujoco_problems_test.py
@@ -33,7 +33,7 @@ def test_registration_and_interaction_with_env_problem(self):
     # This ensures that registration has occurred.
     ep = registry.env_problem("reacher_env_problem", batch_size=batch_size)
     ep.reset()
-    num_done, num_lost, num_won, num_draw = 0, 0, 0, 0
+    num_done = 0
     nsteps = 100
     for _ in range(nsteps):
       actions = np.stack([ep.action_space.sample() for _ in range(batch_size)])
@@ -48,25 +48,10 @@ def test_registration_and_interaction_with_env_problem(self):
       done_indices = env_problem_utils.done_indices(dones)
       ep.reset(done_indices)
       num_done += sum(dones)
-      for r, d in zip(rewards, dones):
-        if not d:
-          continue
-        if r == -1:
-          num_lost += 1
-        elif r == 0:
-          num_draw += 1
-        elif r == 1:
-          num_won += 1
-        else:
-          raise ValueError("reward should be -1, 0, 1 but is {}".format(r))
 
-    # Assert that something got done atleast, without that the next assert is
-    # meaningless.
+    # Assert that something got done atleast,
     self.assertGreater(num_done, 0)
 
-    # Assert that things are consistent.
-    self.assertEqual(num_done, num_won + num_lost + num_draw)
-
 
 if __name__ == "__main__":
   tf.test.main()

From a8b3ff16dd18070fdf9bea834d208f1631abed61 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sat, 15 Jun 2019 12:13:03 -0700
Subject: [PATCH 2130/2720] Remove dictionary support from layer inputs and
 outputs. Clean up naming and docstrings dealing with input and output shapes.

PiperOrigin-RevId: 253396261
---
 tensor2tensor/trax/layers/base.py        | 138 ++++++++++++-----------
 tensor2tensor/trax/layers/combinators.py |  66 +++--------
 2 files changed, 85 insertions(+), 119 deletions(-)

diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 49d473f65..a89927773 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -49,13 +49,12 @@ def call(self, x, params=(), **kwargs):
   def new_parameters(self, input_shape, input_dtype, rng):
     """Create new parameters for the layer given an input shape, dtype and rng.
 
-    Note that all arguments and return values can be tuples or dictionaries
-    or arbitraty nested structures composed of tuples and dictionaries.
-
     Args:
-      input_shape: a tuple representing the shape of the input.
+      input_shape: A tuple representing a shape (if this layer takes one input)
+          or a tuple of shapes (if this layer takes more than one input).
+          For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
       input_dtype: numpy dtype of the input.
-      rng: random number generator.
+      rng: A random number generator.
 
     Returns:
       The newly created parameters for this layer.
@@ -69,18 +68,18 @@ def stack_items_to_pass(self):
 
   # End of subclassing interface, all functions below are internal.
 
-  def output_shape(self, input_shape_and_type, params):
-    """Output shape and type for this layer given input shape and type.
-
-    Note that all arguments and return values can be tuples or dictionaries
-    or arbitrary nested structures composed of tuples and dictionaries.
+  def pseudo_call(self, pseudo_input, params):
+    """Computes what shapes and types this layer would produce for given input.
 
     Args:
-      input_shape_and_type: a ShapeType with shape and type of the input.
-      params: parameters for this layer.
+      pseudo_input: A ShapeType instance (input data minus the actual values)
+          or a tuple of ShapeType instances.
+      params: Parameters for this layer.
 
     Returns:
-      The shape and type of the output.
+      A ShapeType instance representing the shape and type of the output (if
+      this layer has one output) or a tuple of ShapeType instances (if this
+      layer has more than one output).
     """
     try:
       with backend.use_backend('jax'):
@@ -91,12 +90,11 @@ def call_on_input(x, params):
           return _apply_to_first_n(f, x, n)
         params_shapes = nested_map(
             params, lambda x: ShapeType(shape=x.shape, dtype=x.dtype))
-        s = _eval_on_shapes(call_on_input, input_shape_and_type, params_shapes)
+        s = _eval_on_shapes(call_on_input, pseudo_input, params_shapes)
       return s
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback(skip=3)
-      raise LayerError(name, 'output_shape', self._caller,
-                       input_shape_and_type, trace)
+      raise LayerError(name, 'pseudo_call', self._caller, pseudo_input, trace)
 
   def initialize(self, input_shape, input_dtype, rng):
     """Initialize the layer given an input shape, dtype and rng.
@@ -109,9 +107,11 @@ def initialize(self, input_shape, input_dtype, rng):
     or arbitraty nested structures composed of tuples and dictionaries.
 
     Args:
-      input_shape: a tuple representing the shape of the input.
+      input_shape: A tuple representing a shape (if this layer takes one input)
+          or a tuple of shapes (if this layer takes more than one input).
+          For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
       input_dtype: numpy dtype of the input.
-      rng: random number generator.
+      rng: A random number generator.
 
     Returns:
       Newly created parameters on the first call and () on all subsequent calls.
@@ -248,8 +248,6 @@ def nested_map(x, f):
     return [nested_map(y, f) for y in x]
   if isinstance(x, tuple):
     return tuple([nested_map(y, f) for y in x])
-  if isinstance(x, dict):
-    return {k: nested_map(x[k], f) for k in x}
   return f(x)
 
 
@@ -259,8 +257,6 @@ def nested_reduce(x, f):
     return f([nested_reduce(y, f) for y in x])
   if isinstance(x, tuple):
     return f([nested_reduce(y, f) for y in x])
-  if isinstance(x, dict):  # We apply f only to values in the dicts.
-    return f([nested_reduce(v, f) for v in x.values()])
   return x
 
 
@@ -372,23 +368,20 @@ def call_fn(self, x, params=(), **kwargs):
 
 
 def _random_inputs(input_shape, rng, integer_inputs=False):
-  """Create random floats of the given shape.
+  """Creates random floats or ints of the given shape.
 
   Args:
-    input_shape: Could be either:
-        list/tuple of ints, ex: (210, 160, 3) or
-        list/tuple of nested shapes, ex: [(210, 160, 3), (105, 80, 3)] or
-        dictionary of nested shapes, ex: {"obs": [(28, 28, 1), (4,)],
-                                          "sensors": [(3,4), (4, 9)]} or
-        any other combination of these, ex: list of dictionaries of tuples etc.
-    rng: random number generator.
-    integer_inputs: boolean, True if we want arrays of integers, otherwise we
-        produce float32s.
+    input_shape: A tuple representing a shape (if the layer takes one input)
+        or a tuple of shapes (if this layer takes more than one input).
+        For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
+    rng: A random number generator.
+    integer_inputs: If True, use numpy int32 to produce the random data, else
+        use float32.
 
   Returns:
-    Random values of the type and shape specified.
+    Random values with the shape and type specified.
   """
-  if not isinstance(input_shape, dict) and isinstance(input_shape[0], int):
+  if isinstance(input_shape[0], int):
     # Non-nested shape, create a random tuple.
     if not integer_inputs:
       return backend.random.uniform(rng, input_shape, minval=-1.0, maxval=1.0)
@@ -397,43 +390,56 @@ def _random_inputs(input_shape, rng, integer_inputs=False):
     return [_random_inputs(shape, rng, integer_inputs) for shape in input_shape]
   elif isinstance(input_shape, tuple):  # Nested shape: tuple.
     return tuple(_random_inputs(list(input_shape), rng, integer_inputs))
-  elif isinstance(input_shape, dict):  # Nested shape: dict.
-    return {k: _random_inputs(input_shape[k], rng, integer_inputs)
-            for k in input_shape}
   else:
     raise TypeError(type(input_shape))
 
 
-def to_shape_and_type(x_shapes, integers):
-  """Make a shape-and-type tuple from shapes."""
-  if isinstance(x_shapes, dict):  # Nested shape: dict.
-    return {k: to_shape_and_type(x_shapes[k], integers) for k in x_shapes}
-  if isinstance(x_shapes, onp.ndarray):  # Numpy array shape
-    return ShapeType(shape=x_shapes.tolist(),
-                     dtype=onp.int32 if integers else onp.float32)
-  if isinstance(x_shapes[0], (int, onp.int32, onp.int64)):
-    return ShapeType(shape=x_shapes,
-                     dtype=onp.int32 if integers else onp.float32)
-  if isinstance(x_shapes, list):  # Nested shape: list.
-    return [to_shape_and_type(s, integers) for s in x_shapes]
-  if isinstance(x_shapes, tuple):  # Nested shape: tuple.
-    return tuple([to_shape_and_type(s, integers) for s in x_shapes])
-  assert False  # Should never get here.
-
-
-def check_shape_agreement(layer_instance, input_shape, integer_inputs=False):
-  """Check if layer.output_shape agrees with the actual output shape."""
+def _is_tuple_of_shapes(shape):
+  # TODO(jonni): Find better way to distinguish a shape from a tuple of shapes.
+  if not isinstance(shape, tuple):
+    raise TypeError('shape must be a tuple or tuple of tuples, instead got:'
+                    ' {}'.format(shape))
+  return isinstance(shape, tuple) and isinstance(shape[0], tuple)
+
+
+def check_shape_agreement(layer_fn, input_shape, integer_inputs=False):
+  """Checks if the layer's call output agrees its pseudo_call predictions.
+
+  This function helps test layer mechanics and inter-layer connections that
+  aren't dependent on specific data values.
+
+  Args:
+    layer_fn: A Layer instance, viewed as a function from input shapes to
+        output shapes.
+    input_shape: A tuple representing a shape (if the layer takes one input)
+        or a tuple of shapes (if this layer takes more than one input).
+        For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
+    integer_inputs: If True, use numpy int32 as the type for the pseudo-data,
+        else use float32.
+
+  Returns:
+    A tuple representing either a single shape (if the layer has one output) or
+    a tuple of shape tuples (if the layer has more than one output).
+  """
   rng1, rng2, rng3 = backend.random.split(backend.random.get_prng(0), 3)
-  input_shape_and_type = to_shape_and_type(input_shape, integer_inputs)
-  input_dtype = nested_map(input_shape_and_type, lambda x: x.dtype)
-  params = layer_instance.initialize(input_shape, input_dtype, rng1)
-  output_shape_and_type = layer_instance.output_shape(
-      input_shape_and_type, params)
-  output_shape = nested_map(output_shape_and_type, lambda x: x.shape)
-  output_shape = nested_map(output_shape, int)  # Make non-numpy.
-  inputs = _random_inputs(input_shape, rng2, integer_inputs=integer_inputs)
-  result = layer_instance(inputs, params, rng=rng3)
-  result_shape = shapes(result)
+  input_dtype = onp.int32 if integer_inputs else onp.float32
+  if _is_tuple_of_shapes(input_shape):
+    pseudo_data = tuple(ShapeType(x, input_dtype) for x in input_shape)
+    input_dtype = tuple(input_dtype for _ in input_shape)
+  else:
+    pseudo_data = ShapeType(input_shape, input_dtype)
+  params = layer_fn.initialize(input_shape, input_dtype, rng1)
+  pseudo_output = layer_fn.pseudo_call(pseudo_data, params)
+  if isinstance(pseudo_output, tuple):
+    output_shape = tuple(x.shape for x in pseudo_output)
+  else:
+    output_shape = pseudo_output.shape
+
+  random_input = _random_inputs(input_shape, rng2, integer_inputs)
+  real_output = layer_fn(random_input, params, rng=rng3)
+  result_shape = shapes(real_output)
+
   msg = 'output shape %s != real result shape %s' % (output_shape, result_shape)
   assert output_shape == result_shape, msg
+  # TODO(jonni): Remove this assert? It makes test logs harder to read.
   return output_shape
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 0f0bd72db..e0b953aec 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -51,11 +51,11 @@ def _deep_flatten(xs):  # pylint: disable=invalid-name
 
 
 def _ensure_sublayers(layers):  # pylint: disable=invalid-name
-  """Ensures that elements in a layer list (or dict) are layers.
+  """Ensures that elements in a layer list are layers.
 
   Args:
-    layers: A list or dict whose elements/values can each be a layer, a list,
-        or a dict, and so on recursively.
+    layers: A tuple or list whose elements can each be a layer, tuple, or list,
+        and so on recursively.
 
   Returns:
     An analogous collection of layers in which embedded layer lists are
@@ -63,8 +63,6 @@ def _ensure_sublayers(layers):  # pylint: disable=invalid-name
   """
   if not layers:  # None or an empty list can signal a no-op.
     return Serial([])  # no-op, but still handles shapes and initialization
-  elif isinstance(layers, dict):
-    return {k: _ensure_sublayers(v) for k, v in layers.items()}
   elif isinstance(layers, (list, tuple)):
     sublayers_not_lists = []
     for layer in layers:
@@ -98,26 +96,20 @@ def call(self, x, params=(), **kwargs):
       x = layer(x, p, rng=rng, **kwargs)
     return x
 
-  def output_shape(self, input_shape_and_type, params):
-    cur_shape_and_type = input_shape_and_type
-    for layer, param in zip(self._layers, params):
-      cur_shape_and_type = layer.output_shape(cur_shape_and_type, param)
-    return cur_shape_and_type
-
   def new_parameters(self, input_shape, input_dtype, rng):
     def MakeShapeType(shape, dtype):
       if isinstance(dtype, (list, tuple)):
         return tuple(MakeShapeType(s, t) for s, t in zip(shape, dtype))
       return base.ShapeType(shape=shape, dtype=dtype)
     params = []
-    cur_shape_and_type = MakeShapeType(input_shape, input_dtype)
+    pseudo_data = MakeShapeType(input_shape, input_dtype)
     for layer in self._layers:
       rng, layer_rng = backend.random.split(rng)
-      cur_shape = base.nested_map(cur_shape_and_type, lambda x: x.shape)
-      cur_dtype = base.nested_map(cur_shape_and_type, lambda x: x.dtype)
+      cur_shape = base.nested_map(pseudo_data, lambda x: x.shape)
+      cur_dtype = base.nested_map(pseudo_data, lambda x: x.dtype)
       param = layer.initialize(cur_shape, cur_dtype, layer_rng)
       pparam = layer._params   # pylint: disable=protected-access
-      cur_shape_and_type = layer.output_shape(cur_shape_and_type, pparam)
+      pseudo_data = layer.pseudo_call(pseudo_data, pparam)
       params.append(param)
     return params
 
@@ -268,16 +260,8 @@ def call(self, x, params=(), **kwargs):
 
   def new_parameters(self, input_shape, input_dtype, rng):
     rngs = backend.random.split(rng, self._nlayers)
-    if not isinstance(self._layers, dict):
-      return [layer.initialize(input_shape, input_dtype, rng)
-              for layer, rng in zip(self._layers, rngs)]
-
-  def output_shape(self, input_shape, params):
-    output_shapes = []
-    if not isinstance(self._layers, dict):
-      for layer, param in zip(self._layers, params):
-        output_shapes.append(layer.output_shape(input_shape, param))
-      return tuple(output_shapes)
+    return [layer.initialize(input_shape, input_dtype, rng)
+            for layer, rng in zip(self._layers, rngs)]
 
 
 def _nested_op(inputs, op):  # pylint: disable=invalid-name
@@ -388,37 +372,13 @@ def stack_items_to_pass(self):
 
   def call(self, inputs, params=(), **kwargs):
     rngs = _pop_rng_and_split(kwargs, self._nlayers)
-    # If layers are a list or a tuple, just apply them.
-    if not isinstance(self._layers, dict):
-      res = [layer(x, params=p, rng=r, **kwargs)
-             for layer, x, p, r in zip(self._layers, inputs, params, rngs)]
-      # Return a list if inputs are a list and a tuple if inputs are a tuple.
-      if isinstance(inputs, list):
-        return res
-      return tuple(res)
-    # If layers are a dictionary, apply to matching keys.
-    result, counter = {}, 0
-    for k in inputs:
-      if k in self._layers:
-        result[k] = self._layers[k](
-            inputs[k], params=params[k], rng=rngs[counter], **kwargs)
-        counter += 1
-      else:
-        result[k] = inputs[k]
-    return result
+    return tuple(layer(x, params=p, rng=r, **kwargs)
+                 for layer, x, p, r in zip(self._layers, inputs, params, rngs))
 
   def new_parameters(self, input_shape, input_dtype, rng):
     rngs = backend.random.split(rng, self._nlayers)
-    # If the argument layers are a sequence, create parameters for each one.
-    if not isinstance(self._layers, dict):
-      return [layer.initialize(shape, dtype, rng) for layer, shape, dtype, rng
-              in zip(self._layers, input_shape, input_dtype, rngs)]
-    # If the argument layers are a dictionary, create a dictionary too.
-    result, counter = {}, 0
-    for k in self._layers:
-      result[k] = self._layers[k].initialize(input_shape[k], rngs[counter])
-      counter += 1
-    return result
+    return [layer.initialize(shape, dtype, rng) for layer, shape, dtype, rng
+            in zip(self._layers, input_shape, input_dtype, rngs)]
 
 
 def Residual(*layers, **kwargs):

From 4aaf287be92fbd2ede8ef0d6cfdb2cbe736dae89 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 17 Jun 2019 07:51:17 -0700
Subject: [PATCH 2131/2720] Clarify comments/names in heavily used base.layer
 decorator function.

PiperOrigin-RevId: 253579002
---
 tensor2tensor/trax/layers/base.py | 50 ++++++++++++++-----------------
 1 file changed, 22 insertions(+), 28 deletions(-)

diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index a89927773..50a3f280c 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -326,45 +326,39 @@ def _short_traceback(skip=7):
   return '\n'.join(res)
 
 
-# Decorator for making layers from functions.
-
-
 def layer(new_parameters=None, stack_items_to_pass=1):
-  """Create a layer class from a function."""
-  def layer_decorator(call):
-    """Decorating the call function."""
+  """Decorates a function to make it the call method of a new Layer class."""
+  # TODO(jonni): Consider renaming new_parameters to new_parameters_fn.
+
+  def _build_layer_class(raw_call_fn):
+    """Returns a Layer class built around the given call function."""
 
-    def stack_items_to_pass_fn(self):
+    def _stack_items_to_pass(self):
       del self
       return stack_items_to_pass
 
-    def new_parameters_fn(self, input_shape, input_dtype, rng):
+    def _new_parameters(self, input_shape, input_dtype, rng):
       if new_parameters is None:
         return ()
       kwargs = self._init_kwargs  # pylint: disable=protected-access
       return new_parameters(input_shape, input_dtype, rng, **kwargs)
 
-    def call_fn(self, x, params=(), **kwargs):
-      """The call function of the created class, derived from call."""
-      # Merge on-call kwargs with class-kwargs.
-      call_kwargs = kwargs.copy()
-      call_kwargs.update(self._init_kwargs)  # pylint: disable=protected-access
-      # Call with the merged kwargs.
-      return call(x, params=params, **call_kwargs)
-
-    # Set doc for python help.
-    call_fn.__doc__ = call.__doc__
-    if new_parameters is None:
-      new_parameters_fn.__doc__ = new_parameters.__doc__
-
-    # Create the class.
-    cls = type(call.__name__, (Layer,),
-               {'call': call_fn,
-                'new_parameters': new_parameters_fn,
-                'stack_items_to_pass': stack_items_to_pass_fn})
-
+    def _call_with_context(self, x, params=(), **kwargs):
+      """Calls raw_call_fn with extra keyword args from Layer.__init__."""
+      merged_kwargs = kwargs.copy()
+      merged_kwargs.update(self._init_kwargs)  # pylint: disable=protected-access
+      return raw_call_fn(x, params=params, **merged_kwargs)
+
+    # Set docstrings and create the class.
+    _call_with_context.__doc__ = raw_call_fn.__doc__
+    _new_parameters.__doc__ = new_parameters.__doc__  # None.__doc__ is None
+    cls = type(raw_call_fn.__name__, (Layer,),
+               {'call': _call_with_context,
+                'new_parameters': _new_parameters,
+                'stack_items_to_pass': _stack_items_to_pass})
     return cls
-  return layer_decorator
+
+  return _build_layer_class
 
 
 def _random_inputs(input_shape, rng, integer_inputs=False):

From 65c25121ad750127a9fbe021d63cd734087ca9c2 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 17 Jun 2019 10:30:06 -0700
Subject: [PATCH 2132/2720] Update Trax configs after recent changes.

PiperOrigin-RevId: 253608584
---
 .../chunked_transformer_imagenet64_8gb.gin    | 43 -------------------
 .../trax/configs/transformer_big_lm1b_8gb.gin |  1 -
 .../trax/configs/transformer_imdb_8gb.gin     |  1 -
 .../trax/configs/transformer_lm1b_8gb.gin     |  1 -
 .../configs/transformer_lm1b_8gb_testing.gin  |  1 -
 .../trax/configs/transformer_wmt_ende_8gb.gin |  1 -
 6 files changed, 48 deletions(-)
 delete mode 100644 tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin

diff --git a/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin b/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin
deleted file mode 100644
index e72c76d61..000000000
--- a/tensor2tensor/trax/configs/chunked_transformer_imagenet64_8gb.gin
+++ /dev/null
@@ -1,43 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fun:
-# ==============================================================================
-batch_fun.batch_size_per_device = 1
-batch_fun.eval_batch_size = 8
-batch_fun.max_eval_length = 12288  # 64 * 64 * 3
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
-inputs.input_name = 'targets'
-inputs.n_chunks = 64
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 0.1
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 1000
-train.eval_steps = 10
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.ChunkedTransformerLM
-train.run_debug_step = False
-train.train_steps = 500000
-
-# Parameters for ChunkedTransformerLM:
-# ==============================================================================
-ChunkedTransformerLM.d_feature = 1024
-ChunkedTransformerLM.d_feedforward = 4096
-ChunkedTransformerLM.dropout = 0.1
-ChunkedTransformerLM.max_len = 12288  # 64 * 64 * 3
-ChunkedTransformerLM.mode = 'train'
-ChunkedTransformerLM.n_heads = 4
-ChunkedTransformerLM.n_layers = 3
-ChunkedTransformerLM.vocab_size = 256
diff --git a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
index fe6ac3a25..4b3d0fc72 100644
--- a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
@@ -38,7 +38,6 @@ train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.SM3
-train.run_debug_step = False
 train.train_steps = 500000
 
 # Parameters for TransformerLM:
diff --git a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
index 784cb5bde..f25fcc2e4 100644
--- a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
@@ -34,7 +34,6 @@ train.eval_frequency = 100
 train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerEncoder
-train.run_debug_step = False
 train.train_steps = 1000
 
 # Parameters for TransformerLM:
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index a6ac5102b..9c0bfdd87 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -38,7 +38,6 @@ train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.SM3
-train.run_debug_step = False
 train.train_steps = 500000
 
 # Parameters for TransformerLM:
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
index e3c09a5de..e32ce968c 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
@@ -38,7 +38,6 @@ train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.Adam
-train.run_debug_step = False
 train.train_steps = 100000
 
 # Parameters for TransformerLM:
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
index abcdfe7b6..f5ceefbc6 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
@@ -38,7 +38,6 @@ train.eval_frequency = 1000
 train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.Transformer
-train.run_debug_step = False
 train.train_steps = 500000
 
 # Parameters for Transformer:

From 2bccf1cace7ea64ef476a61f2d05e4c7e6b70d7f Mon Sep 17 00:00:00 2001
From: Oscar Ramirez <oars@google.com>
Date: Mon, 17 Jun 2019 10:59:15 -0700
Subject: [PATCH 2133/2720] Update gym version.

PiperOrigin-RevId: 253615467
---
 tensor2tensor/rl/gym_utils.py      | 3 ++-
 tensor2tensor/rl/gym_utils_test.py | 5 ++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index 89ae3698f..d1c7bada3 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import gym
+import gym.wrappers
 import numpy as np
 from PIL import Image
 import tensorflow as tf
@@ -179,7 +180,7 @@ def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env,
     env = RenderedEnv(
         env, resize_to=rendered_env_resize_to, output_dtype=output_dtype)
 
-  if wrap_with_time_limit:
+  if wrap_with_time_limit and rl_env_max_episode_steps is not None:
     env = gym.wrappers.TimeLimit(
         env, max_episode_steps=rl_env_max_episode_steps)
   return env
diff --git a/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index 120ac0c30..007010e10 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -76,12 +76,11 @@ def test_making_timewrapped_env(self):
     self.assertTrue(isinstance(env, gym.wrappers.TimeLimit))
     self.assertEqual(1000, env._max_episode_steps)
 
-  # Make a time-wrapped environment with unlimited limit.
+  # Make an instance of the environment without a TimeLimit
   def test_unlimited_env(self):
     env = gym_utils.make_gym_env("CartPole-v0", rl_env_max_episode_steps=None)
     self.assertTrue(isinstance(env, gym.Env))
-    self.assertTrue(isinstance(env, gym.wrappers.TimeLimit))
-    self.assertTrue(env._max_episode_steps is None)
+    self.assertNotIsInstance(env, gym.wrappers.TimeLimit)
 
   def test_rendered_env(self):
     env = gym_utils.RenderedEnv(SimpleEnv(), resize_to=(64, 12))

From b2615aab938af99418ac0d1318338bf3030357fa Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 17 Jun 2019 15:04:19 -0700
Subject: [PATCH 2134/2720] Add MemoryEfficientTrainer

PiperOrigin-RevId: 253669173
---
 tensor2tensor/trax/trax.py | 146 ++++++++++++++++++++++++++++++++-----
 1 file changed, 128 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index f904b5a60..a6b94c892 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -185,7 +185,7 @@ def _print_n_params(opt_state, n_devices, step):
   """Print out the number of parameters."""
   sizes = layers.sizes(opt_state[0])
   if n_devices > 1:
-    unreplicate = lambda x: x.mean(0)
+    unreplicate = lambda x: x[0]
     single_params = layers.nested_map(opt_state[0], unreplicate)
     sizes = layers.sizes(single_params)
   total_size = layers.nested_reduce(sizes, sum)
@@ -246,6 +246,30 @@ def evaluate(inputs_stream, predict_fn, metric_fns, rng):
   return {m: v / count for (m, v) in six.iteritems(metrics)}
 
 
+def evaluate_loss_train_and_eval(step, inputs, compute_loss_fn, eval_steps,
+                                 rngs,
+                                 train_sw=None, eval_sw=None, history=None):
+  """More efficient evaluation that logs only the loss on train & eval data."""
+  step_log(step, "Evaluation")
+  train_eval_metrics = []
+  for input_stream in [inputs.train_eval_stream, inputs.eval_stream]:
+    total = 0.0
+    count = 0.0
+    for inp in itertools.islice(input_stream(), eval_steps):
+      loss_values, rngs = compute_loss_fn(inp, rngs)
+      total += float(numpy.mean(loss_values))
+      count += 1.0
+    metrics = {"loss": total / count}
+    train_eval_metrics.append(metrics)
+  train_metrics, eval_metrics = train_eval_metrics  # pylint: disable=unbalanced-tuple-unpacking
+  if train_sw:
+    log_metrics(train_metrics, train_sw, "train", step, history=history)
+  if eval_sw:
+    log_metrics(eval_metrics, eval_sw, "eval", step, history=history)
+  step_log(step, "Finished evaluation")
+  return train_metrics, eval_metrics
+
+
 def log_metrics(metrics, summ_writer, log_prefix, step, history=None):
   """Log metrics to summary writer and history."""
   rjust_len = max([len(name) for name in metrics])
@@ -372,6 +396,33 @@ def update(i, opt_state, batch, rng):
   return update
 
 
+@gin.configurable
+def _jit_compute_loss_fn(predict_fn, loss_fn, n_devices, jit=True):
+  """Get jit-ed function that computes the loss."""
+  if n_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
+    def single_compute_loss(opt_state, batch, rng):
+      rng, subrng = jax_random.split(rng[0])
+      return loss_fn(opt_state[0], batch, predict_fn, rng), [subrng]
+    if jit:
+      return backend.jit(single_compute_loss)
+    else:
+      return single_compute_loss
+
+  @functools.partial(backend.pmap, axis_name="batch")
+  def mapped_compute_loss(opt_state, batch, rng):
+    """This is a multi-device version of the update function above."""
+    # We assume all tensors have the first dimension = n_devices.
+    rng, subrng = jax_random.split(rng)
+    loss_val = loss_fn(opt_state[0], batch, predict_fn, rng)
+    return loss_val, subrng
+
+  def compute_loss(opt_state, batch, rng):
+    return mapped_compute_loss(
+        opt_state, reshape_by_device(batch, n_devices), rng)
+
+  return compute_loss
+
+
 def _reshape_by_device_single(x, n_devices):
   """Reshape x into a shape [n_devices, ...]."""
   x_shape = list(x.shape)
@@ -393,6 +444,7 @@ def reshape_by_device(x, n_devices):
       x, lambda x: _reshape_by_device_single(x, n_devices))
 
 
+@gin.configurable(whitelist=[])
 class Trainer(object):
   """Trax trainer.
 
@@ -448,12 +500,15 @@ def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs, output_dir,
     model_input_shape = layers.nested_map(
         model_input_shape, lambda x: x if x else 1)
     if state.params:
-      params = state.params[0]
       opt_state = state.params
     else:
-      params = model_train.initialize(
-          model_input_shape, inputs.input_dtype, init_rng)
-      opt_state = (params, opt.tree_init(params))
+      # JIT parameter initialization to avoid memory fragmentation
+      def initialize(input_shape, input_dtype, init_rng):
+        params = model_train.initialize(input_shape, input_dtype, init_rng)
+        opt_state = (params, opt.tree_init(params))
+        return opt_state
+      initialize = backend.jit(initialize, static_argnums=(0, 1))
+      opt_state = initialize(model_input_shape, inputs.input_dtype, init_rng)
     if n_devices > 1:
       replicate = lambda x: numpy.broadcast_to(x, (n_devices,) + x.shape)
       opt_state = layers.nested_map(opt_state, replicate)
@@ -513,8 +568,12 @@ def train_epoch(self, epoch_steps, eval_steps):
 
       # LR log
       if self._step == 1 or self._step % 10 == 0:
-        self._train_sw.scalar("training/learning rate",
-                              self._lr_fn(self._step), step=self._step)
+        # TODO(lukaszkaiser): it makes no sense to use an accelerator (e.g. TPU)
+        # in op-by-op mode just to compute the learning rate. However, there
+        # should be a cleaner approach that forceably swapping out the backend.
+        with backend.use_backend("numpy"):
+          self._train_sw.scalar("training/learning rate",
+                                self._lr_fn(self._step), step=self._step)
 
     # Timer
     epoch_time = time.time() - start_time
@@ -525,6 +584,17 @@ def train_epoch(self, epoch_steps, eval_steps):
                             epoch_steps / epoch_time, step=self._step)
 
     # Evaluate in parallel
+    self.evaluate(eval_steps)
+
+    # Save state
+    _save_replicated(self._opt_state, self._step, self._history,
+                     self._n_devices, self._output_dir, False)
+
+    # Flush summary writers
+    self._train_sw.flush()
+    self._eval_sw.flush()
+
+  def evaluate(self, eval_steps):
     _, rng = jax_random.split(self._rngs[0])
     evaluate_train_and_eval(
         step=self._step,
@@ -537,14 +607,6 @@ def train_epoch(self, epoch_steps, eval_steps):
         eval_sw=self._eval_sw,
         history=self._history)
 
-    # Save state
-    _save_replicated(self._opt_state, self._step, self._history,
-                     self._n_devices, self._output_dir, False)
-
-    # Flush summary writers
-    self._train_sw.flush()
-    self._eval_sw.flush()
-
   def update_learning_rate(self):
     old_lr_fn = self._lr_fn
     self._lr_fn = self._lr_schedule(self._history)
@@ -577,6 +639,51 @@ def save_computation_graphs(self, save_backward_graph):
         f.write(backward_computation.GetHloDotGraph())
 
 
+@gin.configurable(whitelist=[])
+class MemoryEfficientTrainer(Trainer):
+  """Trax trainer that aims to minimize memory usage.
+  """
+  # TODO(kitaev): memory efficiency should be a feature of the main Trainer
+  # class, but there's a separate class for now because this trainer only
+  # supports evaluating the loss (and not any other metrics).
+
+  def __init__(self, *args, **kwargs):
+    super(MemoryEfficientTrainer, self).__init__(*args, **kwargs)
+    # Model predictions can use large amounts of memory. The memory-efficient
+    # approach is to compute metrics on each replica and then aggregate. For now
+    # we only implement computing the loss, and not any other metrics.
+    self._jit_compute_loss = _jit_compute_loss_fn(
+        self._model_predict_eval, self._loss_fn, self._n_devices)
+
+  def evaluate(self, eval_steps):
+    # Evaluate only the loss function (a more efficient, jitted, implementation)
+    evaluate_loss_train_and_eval(
+        step=self._step,
+        inputs=self._inputs,
+        compute_loss_fn=functools.partial(self._jit_compute_loss,
+                                          self._opt_state),
+        eval_steps=eval_steps,
+        rngs=self._rngs,
+        train_sw=self._train_sw,
+        eval_sw=self._eval_sw,
+        history=self._history)
+
+  def update_learning_rate(self):
+    old_lr_fn = self._lr_fn
+    self._lr_fn = self._lr_schedule(self._history)
+    if self._lr_fn != old_lr_fn:
+      raise NotImplementedError(
+          "Loss function changed. Garbage collection for jitted functions is "
+          "not implemented in jax, so global accelerator memory allocated by "
+          "the jitted update function with the old loss cannot be reclaimed.")
+
+  def save_computation_graphs(self, save_backward_graph):
+    # TODO(kitaev): implement saving graphs while making sure that no op-by-op
+    # execution happens in the process.
+    del save_backward_graph
+    return
+
+
 @gin.configurable(blacklist=["output_dir"])
 def train(output_dir,
           model=gin.REQUIRED,
@@ -584,6 +691,7 @@ def train(output_dir,
           inputs=trax_inputs.inputs,
           optimizer=trax_opt.SM3,
           lr_schedule=lr.MultifactorSchedule,
+          trainer_class=Trainer,
           train_steps=1000,
           save_steps=None,
           eval_steps=10,
@@ -604,6 +712,7 @@ def train(output_dir,
     optimizer: The optimizer (see optimizers/base.py for signature).
     lr_schedule: A learning rate schedule as a function that takes history and
       returns a function from step to learning rate (a float).
+    trainer_class: The trainer class to use.
     train_steps: int, total number of training steps.
     save_steps: list of integers. Keep a model file at each of the supplied save
       steps.
@@ -617,9 +726,10 @@ def train(output_dir,
   Returns:
     trax.State
   """
-  trainer = Trainer(model, loss_fn, optimizer, lr_schedule, inputs, output_dir,
-                    random_seed=random_seed, n_devices=n_devices,
-                    save_steps=save_steps)
+  trainer = trainer_class(model, loss_fn, optimizer, lr_schedule, inputs,
+                          output_dir,
+                          random_seed=random_seed, n_devices=n_devices,
+                          save_steps=save_steps)
 
   epoch_steps = [train_steps]  # Only training if eval_frequency is 0 or None
   if eval_frequency and eval_steps > 0:

From b3a667e841953ad81cdd718c1b94e802df2891fe Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 18 Jun 2019 13:16:55 -0700
Subject: [PATCH 2135/2720] Optimize memory usage in trax

- Dropout and attention masking no longer store in global memory constants that have the same shape as the activations
- layers.one_hot no longer stores a large intermediate quantity in global memory

PiperOrigin-RevId: 253851993
---
 tensor2tensor/trax/backend.py          |  2 +-
 tensor2tensor/trax/layers/attention.py | 11 +++++++++--
 tensor2tensor/trax/layers/base.py      |  3 ++-
 tensor2tensor/trax/layers/core.py      | 11 ++++++++---
 4 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index c3a0cee17..6e2aa8967 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -40,7 +40,7 @@
     "random_uniform": jax_random.uniform,
     "random_normal": jax_random.normal,
     "random_bernoulli": jax_random.bernoulli,
-    "random_get_prng": jax_random.PRNGKey,
+    "random_get_prng": jax.jit(jax_random.PRNGKey),
     "random_split": jax_random.split,
 }
 
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 1934e7da2..725713ced 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import jax
 import numpy as onp
 
 from tensor2tensor.trax import backend
@@ -91,14 +92,20 @@ def DotProductAttention(query, key, value, mask, dropout, mode, rng):
   depth = np.shape(query)[-1]
   dots = np.matmul(query, np.swapaxes(key, -1, -2)) / np.sqrt(depth)
   if mask is not None:
-    dots = np.where(mask, dots, -1e9)
+    # TODO(kitaev): workaround for https://github.com/google/jax/issues/850
+    # We must ensure that both mask and the -1e9 constant have a data dependency
+    # on the input. Broadcasted copies of these use a lot of memory, so they
+    # should be computed at runtime (rather than being global constants).
+    if backend.get_name() == 'jax':
+      mask = jax.lax.tie_in(dots, mask)
+    dots = np.where(mask, dots, np.full_like(dots, -1e9))
   # Softmax.
   dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))
   if dropout >= 1.0:
     raise ValueError('Dropout rates must be lower than 1.')
   if dropout is not None and dropout > 0.0 and mode == 'train':
     keep = backend.random.bernoulli(rng, 1.0 - dropout, dots.shape)
-    dots = np.where(keep, dots / (1.0 - dropout), 0)
+    dots = np.where(keep, dots / (1.0 - dropout), np.zeros_like(dots))
   out = np.matmul(dots, value)
   return out
 
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 50a3f280c..58096f821 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -83,7 +83,8 @@ def pseudo_call(self, pseudo_input, params):
     """
     try:
       with backend.use_backend('jax'):
-        rng = backend.random.get_prng(0)
+        # Same as backend.random.get_prng(0), but no op-by-op execution.
+        rng = onp.zeros(2, onp.uint32)
         def call_on_input(x, params):
           f = lambda y: self.call(y, params=params, rng=rng)
           n = self.stack_items_to_pass() if isinstance(x, (list, tuple)) else 0
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 97afa10cb..8419dafbe 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import jax
 import numpy as onp
 
 from tensor2tensor.trax import backend
@@ -29,7 +30,7 @@
 
 @base.layer()
 def Relu(x, **unused_kwargs):
-  return np.maximum(x, np.array(0, dtype=x.dtype))
+  return np.maximum(x, np.zeros_like(x))
 
 
 @base.layer()
@@ -157,7 +158,7 @@ def Dropout(x, params, rate=0.0, mode='train', rng=None, **kwargs):
     raise ValueError('Dropout rate (%f) must be lower than 1.' % rate)
   if mode == 'train' and rate > 0.0:
     keep = backend.random.bernoulli(rng, 1.0 - rate, x.shape)
-    return np.where(keep, x / (1.0 - rate), 0)
+    return np.where(keep, x / (1.0 - rate), np.zeros_like(x))
   else:
     return x
 
@@ -176,7 +177,11 @@ def AddConstant(x, params, constant=0.0, **unused_kwargs):
 
 def one_hot(x, size, dtype=np.float32):  # pylint: disable=invalid-name
   """Make a n+1 dim one-hot array from n dim int-categorical array."""
-  return np.array(x[..., np.newaxis] == np.arange(size), dtype)
+  arange_size = np.arange(size)
+  if backend.get_name() == 'jax':
+    # Work around a jax broadcasting issue.
+    arange_size = jax.lax.tie_in(x, arange_size)
+  return np.array(x[..., np.newaxis] == arange_size, dtype)
 
 
 # Mean.

From eb6d8253ae45f0f80cdacea454cbca6f71f6e2f4 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 18 Jun 2019 14:57:45 -0700
Subject: [PATCH 2136/2720] Move towards py3 compatibility.

PiperOrigin-RevId: 253874773
---
 tensor2tensor/models/research/vqa_attention.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index c73f25ab3..801ef577f 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -51,11 +51,13 @@ class VqaAttentionBaseline(t2t_model.T2TModel):
 
   def body(self, features):
     hp = self.hparams
-    # pylint: disable=eval-used
+    model_fn = resnet_v1_152
+    if hp.image_model_fn != "resnet_v1_152":
+      model_fn = eval(hp.image_model_fn)  # pylint: disable=eval-used
     if hp.image_input_type == "image":
       image_feat = vqa_layers.image_embedding(
           features["inputs"],
-          model_fn=eval(hp.image_model_fn),
+          model_fn=model_fn,
           trainable=hp.train_resnet,
           is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
     else:

From 8c23cbb2f3634d7ba2d9ade1c88b935e07197218 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 18 Jun 2019 15:29:29 -0700
Subject: [PATCH 2137/2720] Reversible Transformer

PiperOrigin-RevId: 253881109
---
 .../transformer_revnet_imagenet64_8gb.gin     |  46 +++
 .../configs/transformer_revnet_lm1b_8gb.gin   |  56 +++
 tensor2tensor/trax/layers/combinators.py      |  11 +
 tensor2tensor/trax/models/__init__.py         |   2 +
 .../models/research/transformer_revnet.py     | 364 ++++++++++++++++++
 5 files changed, 479 insertions(+)
 create mode 100644 tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
 create mode 100644 tensor2tensor/trax/configs/transformer_revnet_lm1b_8gb.gin
 create mode 100644 tensor2tensor/trax/models/research/transformer_revnet.py

diff --git a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
new file mode 100644
index 000000000..9e62d3ff4
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
@@ -0,0 +1,46 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 32
+batch_fun.eval_batch_size = 8
+batch_fun.max_eval_length = 12288  # 64 * 64 * 3
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
+inputs.input_name = 'targets'
+inputs.n_chunks = 64
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 0.3
+MultifactorSchedule.factors = 'constant * linear_warmup'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 100
+train.eval_steps = 10
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerRevnetLM
+train.optimizer = @trax.optimizers.SM3
+train.train_steps = 500000
+train.trainer_class = @MemoryEfficientTrainer
+
+# Parameters for TransformerRevnetLM:
+# ==============================================================================
+TransformerRevnetLM.d_feature = 512
+TransformerRevnetLM.d_feedforward = 2048
+TransformerRevnetLM.dropout = 0.1
+TransformerRevnetLM.max_len = 12288  # 64 * 64 * 3
+TransformerRevnetLM.mode = 'train'
+TransformerRevnetLM.n_heads = 8
+TransformerRevnetLM.n_layers = 6
+TransformerRevnetLM.vocab_size = 256
+TransformerRevnetLM.n_chunks = 64
+TransformerRevnetLM.n_attention_chunks = 64
diff --git a/tensor2tensor/trax/configs/transformer_revnet_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_lm1b_8gb.gin
new file mode 100644
index 000000000..f1b7addec
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_revnet_lm1b_8gb.gin
@@ -0,0 +1,56 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 4096
+batch_fun.eval_batch_size = 128
+batch_fun.max_eval_length = 2048
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_languagemodel_lm1b32k_packed'
+inputs.input_name = 'targets'
+inputs.n_chunks = 32
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 0.3
+MultifactorSchedule.factors = 'constant * linear_warmup'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
+lm1b_preprocess.max_target_length = 512
+lm1b_preprocess.max_eval_target_length = 2048
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 100
+train.eval_steps = 10
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerRevnetLM
+train.optimizer = @trax.optimizers.SM3
+train.train_steps = 500000
+train.trainer_class = @MemoryEfficientTrainer
+
+# Parameters for TransformerRevnetLM:
+# ==============================================================================
+TransformerRevnetLM.d_feature = 512
+TransformerRevnetLM.d_feedforward = 2048
+TransformerRevnetLM.dropout = 0.1
+TransformerRevnetLM.max_len = 2048
+TransformerRevnetLM.mode = 'train'
+TransformerRevnetLM.n_heads = 8
+TransformerRevnetLM.n_layers = 6
+TransformerRevnetLM.vocab_size = 32000
+TransformerRevnetLM.n_chunks = 32
+TransformerRevnetLM.n_attention_chunks = 8
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index e0b953aec..41025f45a 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -303,6 +303,13 @@ def Add(x, **unused_kwargs):
   return _binary_op(x, op=sum)
 
 
+@base.layer(stack_items_to_pass=0)
+def SubtractTop(x, **unused_kwargs):
+  """Subtract the first element on the stack from the second element."""
+  # Here x is a list of tensors of the same shape, or nested structures.
+  return _binary_op(x, op=lambda xs: xs[1] - xs[0])
+
+
 @base.layer(stack_items_to_pass=0)
 def Multiply(x, **unused_kwargs):
   """Multiply first and second element on the stack."""
@@ -372,6 +379,10 @@ def stack_items_to_pass(self):
 
   def call(self, inputs, params=(), **kwargs):
     rngs = _pop_rng_and_split(kwargs, self._nlayers)
+    # Note that zip silently truncates its result if lengths don't match.
+    assert len(inputs) == self._nlayers
+    assert len(params) == self._nlayers
+    assert len(rngs) == self._nlayers
     return tuple(layer(x, params=p, rng=r, **kwargs)
                  for layer, x, p, r in zip(self._layers, inputs, params, rngs))
 
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index ad1d27b5c..22db0422e 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -25,6 +25,7 @@
 from tensor2tensor.trax.models import resnet
 from tensor2tensor.trax.models import transformer
 from tensor2tensor.trax.models.research import position_lookup_transformer
+from tensor2tensor.trax.models.research import transformer_revnet
 
 
 # Ginify
@@ -42,4 +43,5 @@ def model_configure(*args, **kwargs):
 Transformer = model_configure(transformer.Transformer)
 TransformerEncoder = model_configure(transformer.TransformerEncoder)
 TransformerLM = model_configure(transformer.TransformerLM)
+TransformerRevnetLM = model_configure(transformer_revnet.TransformerRevnetLM)
 WideResnet = model_configure(resnet.WideResnet)
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
new file mode 100644
index 000000000..dd17354a3
--- /dev/null
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -0,0 +1,364 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Transformer Models."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import jax
+
+from tensor2tensor.trax import backend
+from tensor2tensor.trax import layers as tl
+from tensor2tensor.trax.layers.combinators import _pop_rng_and_split
+
+
+# Layers are always CamelCase, but functions in general are snake_case
+# pylint: disable=invalid-name
+
+
+class Map(tl.Layer):
+  """Combinator for applying a layer to a list or tuple.
+
+  Args:
+    layer: a layer to apply to each element.
+
+  Returns:
+    A new layer representing mapping layer to all elements of the input.
+  """
+
+  def __init__(self, layer, check_shapes=True):
+    super(Map, self).__init__()
+    if layer is None or isinstance(layer, (list, tuple)):
+      layer = tl.Serial(layer)
+    self._layer = layer
+    # Generally a Map should be applied to lists where all elements have
+    # the same shape -- because self._layer will only be initialized once
+    # and it could have different parameters for different shapes. But there
+    # are valid cases -- e.g., when self._layer has no parameters -- where we
+    # can apply Map to different shapes -- set check_shapes=False in such cases.
+    self._check_shapes = check_shapes
+
+  def call(self, inputs, params=(), **kwargs):
+    rngs = _pop_rng_and_split(kwargs, len(inputs))
+    result = [self._layer(x, params=params, rng=r, **kwargs)
+              for x, r in zip(inputs, rngs)]
+    if isinstance(inputs, list):
+      return result
+    return tuple(result)
+
+  def new_parameters(self, input_shape, input_dtype, rng):
+    first_shape = input_shape[0]
+    if self._check_shapes:
+      for shape in input_shape:
+        if shape != first_shape:
+          raise ValueError('Map layer can only be applied to list of elements '
+                           'with the same shapes. Shapes: %s' % str(shape))
+    return self._layer.initialize(first_shape, input_dtype[0], rng)
+
+
+def FeedForward(d_feature, d_feedforward, dropout, mode):
+  """Feed-forward block with layer normalization at start."""
+  # TODO(kitaev): dropout is disabled to save memory
+  del dropout, mode
+  return [
+      tl.LayerNorm(),
+      tl.Dense(d_feedforward),
+      tl.Relu(),
+      # tl.Dropout(rate=dropout, mode=mode),
+      tl.Dense(d_feature),
+      # tl.Dropout(rate=dropout, mode=mode),
+  ]
+
+
+class ReversibleLayerMixin(object):
+  """Reversible Layer Mixin."""
+
+  def inverse_and_vjp(self, output, ct, params=(), **kwargs):
+    """Backward pass: computes the inverse of a layer and propagates gradients.
+
+    Args:
+      output: Output activations; can be a (possibly nested) tuple or list.
+      ct: gradient signal (cotangent) computed based on subsequent layers. If
+          None, no gradients are propagated. Otherwise the structure and shape
+          must match the output.
+      params: layer parameters
+      **kwargs: kwargs for the layer
+
+    Returns:
+      A tuple (x, x_ct), where x is the reconstructed input and x_ct is the
+      gradient signal for the input. If ct is None, x_ct will also be None.
+    """
+    if ct is None:
+      # Subclasses must override inverse_and_vjp, but in the case where ct is
+      # not None there is an unoptimized implementation below that they can
+      # delegate to.
+      raise NotImplementedError
+
+    # Note: jax.vjp does not allow us to use **kwargs in the signature here.
+    def _do_call(x, params, kwargs):
+      return super(ReversibleLayerMixin, self).__call__(x, params, **kwargs)
+
+    reconstructed_x, must_be_none = self.inverse_and_vjp(
+        output, None, params, **kwargs)
+    assert must_be_none is None
+    _, vjpfun = jax.vjp(_do_call, reconstructed_x, params, kwargs)
+    input_ct = vjpfun(ct)
+    return reconstructed_x, input_ct
+
+  def __call__(self, x, params=(), **kwargs):
+    assert backend.get_name() == 'jax', (
+        'Reversible layers are only supported in JAX')
+
+    # Retrieve shared parameters (cf. tl.Layer.__call__)
+    super(ReversibleLayerMixin, self).__call__(x, params, **kwargs)
+    if params is () and self._params:  # pylint: disable=literal-comparison
+      # TODO(kitaev): Figure out why parameter sharing doesn't work (if this
+      # explicit error isn't thrown, a jax tracer error occurs instead)
+      raise NotImplementedError(
+          'Parameter sharing between reversible layers is not implemented.')
+
+    @jax.custom_transforms
+    def do_call(x, params, kwargs):
+      return super(ReversibleLayerMixin, self).__call__(x, params, **kwargs)
+
+    def do_call_vjp(x, params, kwargs):
+      output = super(ReversibleLayerMixin, self).__call__(x, params, **kwargs)
+      def vjpfun(ct):
+        _, input_ct = self.inverse_and_vjp(output, ct, params, **kwargs)
+        return input_ct
+
+      return output, vjpfun
+
+    jax.defvjp_all(do_call, do_call_vjp)
+    return do_call(x, params, kwargs)
+
+
+@tl.layer(stack_items_to_pass=1)
+def Split(x, params, sections=2, axis=-1, **kwargs):
+  del params, kwargs
+  return list(backend.numpy.split(x, sections, axis))
+
+
+@tl.layer(stack_items_to_pass=1)
+def Duplicate(x, params, sections=2, **kwargs):
+  del params, kwargs
+  return [x for _ in range(sections)]
+
+
+class ReversibleHalfResidual(ReversibleLayerMixin, tl.Serial):
+  """Half of a RevNet-style residual (only updates part of the hidden state)."""
+
+  def __init__(self, residual_layers):
+    self.compute_residual = tl.Serial([
+        tl.Select(inputs=('x1_or_y1', 'x2'), output=('x2', 'x1_or_y1', 'x2')),
+        tl.Parallel(residual_layers, [], []),
+    ])
+
+    layers = [self.compute_residual, tl.Add()]
+    super(ReversibleHalfResidual, self).__init__(layers)
+
+    self.subtract_top = tl.SubtractTop()
+    self.reverse_layers = [self.compute_residual, self.subtract_top]
+
+  def inverse_and_vjp(self, output, ct, params=(), **kwargs):
+    rng = kwargs.pop('rng', None)
+    rngs = (None,) * self._nlayers
+    if rng is not None:
+      rngs = backend.random.split(rng, self._nlayers)
+
+    if ct is None:
+      reconstructed_x = output
+      # Note that self._layers aligns exactly with self.reverse_layers in terms
+      # of parameter and rng usage, so no re-ordering is required.
+      for layer, p, rng in zip(self.reverse_layers, params, rngs):
+        reconstructed_x = layer(reconstructed_x, p, rng=rng, **kwargs)
+      return reconstructed_x, None
+    else:
+      # Note: jax.vjp does not allow us to use **kwargs in the signature here.
+      def call_compute_residual(x, params, kwargs):
+        return self.compute_residual(x, params, **kwargs)
+
+      assert len(ct) == 2
+      ct = ((ct[0], ct[0], ct[1]))
+
+      compute_residual_kwargs = kwargs.copy()
+      compute_residual_kwargs['rng'] = rngs[0]
+      stack_with_residual, vjpfun = jax.vjp(
+          call_compute_residual, output, params[0], compute_residual_kwargs)
+      reconstructed_x = self.subtract_top(
+          stack_with_residual, params[-1], rng=rngs[-1], **kwargs)
+
+      x_ct, residual_params_ct, kwargs_ct = vjpfun(ct)
+      return reconstructed_x, (x_ct, (residual_params_ct, ()), kwargs_ct)
+
+
+class ReversibleSwap(ReversibleLayerMixin, tl.Swap):
+  """Swap the first two element on the stack."""
+
+  def inverse_and_vjp(self, output, ct, params=(), **kwargs):
+    if ct is None:
+      # Swap is its own inverse
+      return self.call(output, params, **kwargs), None
+    else:
+      return super(ReversibleSwap, self).inverse_and_vjp(
+          output, ct, params, **kwargs)
+
+
+def ReversibleResidual(layers_a, layers_b):
+  """RevNet-style reversible residual layer."""
+  return [
+      ReversibleHalfResidual(layers_a),  # (x1, x2) -> (z1, x2)
+      ReversibleSwap(),  # (z1, x2) -> (x2, z1)
+      ReversibleHalfResidual(layers_b),  # (x2, z1) -> (y2, z1)
+      ReversibleSwap(),  # (y2, z1) -> (z1, y2); where y1 := z1
+  ]
+
+
+class ReversibleSerial(ReversibleLayerMixin, tl.Serial):
+  """A reversible version of tl.Serial (requires reversible sub-layers)."""
+
+  def __init__(self, *layers):
+    super(ReversibleSerial, self).__init__(*layers)
+
+    # Note that self._layers has already been flattened to remove nested lists.
+    for i, layer in enumerate(self._layers):
+      if not isinstance(layer, ReversibleLayerMixin):
+        raise ValueError(
+            'Sub-layer {} of ReversibleSerial is not reversible: {}'.format(
+                i, layer))
+
+  def inverse_and_vjp(self, output, ct, params=(), **kwargs):
+    rng = kwargs.pop('rng', None)
+    rngs = (None,) * self._nlayers
+    if rng is not None:
+      rngs = backend.random.split(rng, self._nlayers)
+
+    layer_val = output
+    if ct is not None:
+      layer_ct = ct
+      params_ct = []
+    for layer, p, rng in reversed(zip(self._layers, params, rngs)):
+      layer_val, layer_ct = layer.inverse_and_vjp(
+          layer_val, layer_ct, p, rng=rng, **kwargs)
+      if ct is not None:
+        layer_ct, p_ct, kwargs_ct = layer_ct
+        params_ct.insert(0, p_ct)
+
+    # TODO(kitaev): Handle kwargs_ct properly. However, kwargs generally only
+    # contains the rng, which is non-differentiable.
+    for k in kwargs:
+      if k != 'rng':
+        raise NotImplementedError(
+            'ReversibleSerial does not support differentiation wrt kwargs,'
+            'and the key {} is not known to be non-differentiable.'.format(k))
+
+    if ct is not None:
+      return layer_val, (layer_ct, params_ct, kwargs_ct)
+    else:
+      return layer_val, None
+
+
+def DecoderBlock(d_feature, d_feedforward, n_heads, n_attention_chunks,
+                 dropout, mode):
+  """Reversible transformer decoder layer.
+
+  Args:
+    d_feature: int:  depth of embedding
+    d_feedforward: int: depth of feed-forward layer
+    n_heads: int: number of attention heads
+    n_attention_chunks: int: number of chunks for memory-efficient attention
+    dropout: float: dropout rate (how much to drop out)
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    the layer.
+  """
+  self_attention = [
+      tl.LayerNorm(),
+      tl.Branch([], tl.CausalMask(axis=-2)),  # Create mask.
+      tl.MultiHeadedAttention(
+          d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+      tl.Select(0),  # Drop mask.
+      tl.Dropout(rate=dropout, mode=mode),
+  ]
+
+  # TODO(kitaev): Memory-efficient attention. This chunking is temporary.
+  self_attention = [
+      Split(sections=n_attention_chunks, axis=-2),  # pylint: disable=no-value-for-parameter
+      Map(self_attention),
+      tl.Concatenate(axis=-2),
+  ]
+
+  feed_forward = [
+      FeedForward(d_feature, d_feedforward, dropout, mode=mode),
+  ]
+  return [
+      ReversibleResidual([self_attention], [feed_forward]),
+  ]
+
+
+def TransformerRevnetLM(vocab_size,
+                        d_feature=512,
+                        d_feedforward=2048,
+                        n_layers=6,
+                        n_heads=8,
+                        dropout=0.1,
+                        max_len=2048,
+                        n_chunks=32,
+                        n_attention_chunks=8,
+                        mode='train'):
+  """Reversible transformer language model (only uses a decoder, no encoder).
+
+  Args:
+    vocab_size: int: vocab size
+    d_feature: int:  depth of *each half* of the two-part features
+    d_feedforward: int: depth of feed-forward layer
+    n_layers: int: number of decoder layers
+    n_heads: int: number of attention heads
+    dropout: float: dropout rate (how much to drop out)
+    max_len: int: maximum symbol length for positional encoding
+    n_chunks: int: number of chunks (must match input pipeline)
+    n_attention_chunks: int: number of chunks for memory-efficient attention
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    the layer.
+  """
+  positional_embedder = [
+      tl.Embedding(d_feature, vocab_size),
+      # TODO(kitaev): dropout is disabled to save memory
+      # tl.Dropout(rate=dropout, mode=mode),
+      tl.PositionalEncoding(max_len=max_len),
+  ]
+  return tl.Model(
+      tl.Concatenate(),
+      tl.ShiftRight(),
+      positional_embedder,
+      Duplicate(),  # pylint: disable=no-value-for-parameter
+      ReversibleSerial([
+          DecoderBlock(d_feature, d_feedforward, n_heads, n_attention_chunks,
+                       dropout, mode)
+          for _ in range(n_layers)
+      ]),
+      tl.Parallel(tl.LayerNorm(), tl.LayerNorm()),
+      tl.Concatenate(),
+      Split(sections=n_chunks, axis=-2),  # pylint: disable=no-value-for-parameter
+      Map([
+          tl.Dense(vocab_size),
+          tl.LogSoftmax(),
+      ]),
+  )
+

From 1af4cb3be5d45e0e65493080592ee38cdfbbea1f Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Tue, 18 Jun 2019 17:51:01 -0700
Subject: [PATCH 2138/2720] Allow for overwriting
 transformer_prepare_encoder/decoder and introduce optionally type_ids in
 transformer_layers.transformer_prepare_encoder.

PiperOrigin-RevId: 253905794
---
 tensor2tensor/layers/transformer_layers.py | 11 ++++++++++-
 tensor2tensor/models/transformer.py        | 14 ++++++++++----
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 59a0ced56..318d6077d 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -32,7 +32,8 @@ def layers():
   return common_layers.layers()
 
 
-def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
+def transformer_prepare_encoder(inputs, target_space, hparams, features=None,
+                                type_ids=None):
   """Prepare one shard of the model for the encoder.
 
   Args:
@@ -41,6 +42,8 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
     hparams: run hyperparameters
     features: optionally pass the entire features dictionary as well.
       This is needed now for "packed" datasets.
+    type_ids: optional, an int64 Tensor of shape [batch, length] that allows
+      for adding type embeddings, similar to positional embeddings.
 
   Returns:
     encoder_input: a Tensor, bottom of encoder stack
@@ -108,6 +111,12 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
         encoder_input, hparams.max_length, "inputs_positional_embedding",
         inputs_position)
 
+  # Add type embeddings
+  if type_ids is not None:
+    encoder_input = common_attention.add_positional_embedding(
+        encoder_input, hparams.max_length, "inputs_type_embedding",
+        type_ids)
+
   encoder_self_attention_bias = common_layers.cast_like(
       encoder_self_attention_bias, encoder_input)
   encoder_decoder_attention_bias = common_layers.cast_like(
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index c6a678250..368fbcb67 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -56,7 +56,7 @@
 
 def transformer_encode(encoder_function, inputs, target_space, hparams,
                        attention_weights=None, features=None, losses=None,
-                       **kwargs):
+                       prepare_encoder_fn=None, **kwargs):
   """Encode transformer inputs.
 
   Args:
@@ -69,6 +69,7 @@ def transformer_encode(encoder_function, inputs, target_space, hparams,
     features: optionally pass the entire features dictionary as well. This is
       needed now for "packed" datasets.
     losses: optional list onto which to append extra training losses
+    prepare_encoder_fn: optional, alternative to transformer_prepare_encoder.
     **kwargs: additional arguments to pass to encoder_function
 
   Returns:
@@ -80,8 +81,10 @@ def transformer_encode(encoder_function, inputs, target_space, hparams,
   """
   inputs = common_layers.flatten4d3d(inputs)
 
+  if not prepare_encoder_fn:
+    prepare_encoder_fn = transformer_prepare_encoder
   encoder_input, self_attention_bias, encoder_decoder_attention_bias = (
-      transformer_prepare_encoder(
+      prepare_encoder_fn(
           inputs, target_space, hparams, features=features))
 
   mlperf_log.transformer_print(
@@ -189,13 +192,16 @@ def __init__(self, *args, **kwargs):
     self._encoder_function = transformer_encoder
     self._decoder_function = transformer_decoder
     self._init_cache_fn = _init_transformer_cache
+    self._prepare_encoder_fn = transformer_prepare_encoder
+    self._prepare_decoder_fn = transformer_prepare_decoder
 
   def encode(self, inputs, target_space, hparams, features=None, losses=None):
     """Encode transformer inputs, see transformer_encode."""
     return transformer_encode(
         self._encoder_function, inputs, target_space, hparams,
         attention_weights=self.attention_weights,
-        features=features, losses=losses)
+        features=features, losses=losses,
+        prepare_encoder_fn=self._prepare_encoder_fn)
 
   def decode(self,
              decoder_input,
@@ -245,7 +251,7 @@ def body(self, features):
     targets = features["targets"]
     targets_shape = common_layers.shape_list(targets)
     targets = common_layers.flatten4d3d(targets)
-    decoder_input, decoder_self_attention_bias = transformer_prepare_decoder(
+    decoder_input, decoder_self_attention_bias = self._prepare_decoder_fn(
         targets, hparams, features=features)
 
     # Not all subclasses of Transformer support keyword arguments related to

From ddf0ef2f1f5e7653887b7ed9997d52622888957c Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Wed, 19 Jun 2019 21:07:05 +0200
Subject: [PATCH 2139/2720] Introduce Rainbow. (#1607)

---
 tensor2tensor/models/research/rl.py           |   9 +
 tensor2tensor/rl/dopamine_connector.py        | 279 ++++++++++++++++--
 .../rl/trainer_model_based_params.py          |   8 +
 3 files changed, 275 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 4cb2bf161..a3ed18a6d 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -377,6 +377,7 @@ def dqn_atari_base():
       agent_epsilon_eval=0.001,
       agent_epsilon_decay_period=250000,  # agent steps
       agent_generates_trainable_dones=True,
+      agent_type="VanillaDQN",  # one of ["Rainbow", "VanillaDQN"]
 
       optimizer_class="RMSProp",
       optimizer_learning_rate=0.00025,
@@ -420,6 +421,14 @@ def dqn_guess1_params():
   return hparams
 
 
+@registry.register_hparams
+def dqn_guess1_rainbow_params():
+  """Guess 1 for DQN params."""
+  hparams = dqn_guess1_params()
+  hparams.set_hparam("agent_type", "Rainbow")
+  return hparams
+
+
 @registry.register_hparams
 def dqn_2m_replay_buffer_params():
   """Guess 1 for DQN params, 2 milions transitions in replay buffer."""
diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 45605c793..a1ab21bc1 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -24,14 +24,19 @@
 import sys
 
 from dopamine.agents.dqn import dqn_agent
+from dopamine.agents.rainbow import rainbow_agent
 from dopamine.replay_memory import circular_replay_buffer
-from dopamine.replay_memory.circular_replay_buffer import OutOfGraphReplayBuffer
-from dopamine.replay_memory.circular_replay_buffer import ReplayElement
+from dopamine.replay_memory.circular_replay_buffer import \
+    OutOfGraphReplayBuffer, ReplayElement
+from dopamine.replay_memory.prioritized_replay_buffer import \
+  OutOfGraphPrioritizedReplayBuffer, WrappedPrioritizedReplayBuffer
 import numpy as np
+
 from tensor2tensor.rl.policy_learner import PolicyLearner
 import tensorflow as tf
 
 # pylint: disable=g-import-not-at-top
+# pylint: disable=ungrouped-imports
 try:
   import cv2
 except ImportError:
@@ -41,7 +46,18 @@
 except ImportError:
   run_experiment = None
 # pylint: enable=g-import-not-at-top
-
+# pylint: enable=ungrouped-imports
+
+# TODO: Vanilla DQN and Rainbow have a lot of common code. Most likely we want
+#  to remove Vanilla DQN and only have Rainbow. To do so one needs to remove
+#  following:
+#    * _DQNAgent
+#    * BatchDQNAgent
+#    * _OutOfGraphReplayBuffer
+#    * "if" clause in create_agent()
+#    * parameter "agent_type" from dqn_atari_base() hparams and possibly other
+#      rlmb dqn hparams sets
+#  If we want to keep both Vanilla DQN and Rainbow, larger refactor is required.
 
 class _DQNAgent(dqn_agent.DQNAgent):
   """Modify dopamine DQNAgent to match our needs.
@@ -178,6 +194,201 @@ def choose_action(ix):
     return np.array([choose_action(ix) for ix in range(self.env_batch_size)])
 
 
+class _OutOfGraphReplayBuffer(OutOfGraphReplayBuffer):
+  """Replay not sampling artificial_terminal transition.
+
+  Adds to stored tuples "artificial_done" field (as last ReplayElement).
+  When sampling, ignores tuples for which artificial_done is True.
+
+  When adding new attributes check if there are loaded from disk, when using
+  load() method.
+
+  Attributes:
+      are_terminal_valid: A boolean indicating if newly added terminal
+        transitions should be marked as artificially done. Replay data loaded
+        from disk will not be overridden.
+  """
+
+  def __init__(self, artificial_done, **kwargs):
+    extra_storage_types = kwargs.pop("extra_storage_types", None) or []
+    extra_storage_types.append(ReplayElement("artificial_done", (), np.uint8))
+    super(_OutOfGraphReplayBuffer, self).__init__(
+        extra_storage_types=extra_storage_types, **kwargs)
+    self._artificial_done = artificial_done
+
+  def is_valid_transition(self, index):
+    valid = super(_OutOfGraphReplayBuffer, self).is_valid_transition(index)
+    valid &= not self.get_artificial_done_stack(index).any()
+    return valid
+
+  def get_artificial_done_stack(self, index):
+    return self.get_range(self._store["artificial_done"],
+                          index - self._stack_size + 1, index + 1)
+
+  def add(self, observation, action, reward, terminal, *args):
+    """Append artificial_done to *args and run parent method."""
+    # If this will be a problem for maintenance, we could probably override
+    # DQNAgent.add() method instead.
+    artificial_done = self._artificial_done and terminal
+    args = list(args)
+    args.append(artificial_done)
+    return super(_OutOfGraphReplayBuffer, self).add(observation, action, reward,
+                                                    terminal, *args)
+
+  def load(self, *args, **kwargs):
+    # Check that appropriate attributes are not overridden
+    are_terminal_valid = self._artificial_done
+    super(_OutOfGraphReplayBuffer, self).load(*args, **kwargs)
+    assert self._artificial_done == are_terminal_valid
+
+
+class _WrappedPrioritizedReplayBuffer(WrappedPrioritizedReplayBuffer):
+  """
+
+  Allows to pass out-of-graph-replay-buffer via wrapped_memory.
+  """
+  def __init__(self, wrapped_memory, batch_size, use_staging):
+    self.batch_size = batch_size
+    self.memory = wrapped_memory
+    self.create_sampling_ops(use_staging)
+
+
+class _RainbowAgent(rainbow_agent.RainbowAgent):
+  """Modify dopamine DQNAgent to match our needs.
+
+  Allow passing batch_size and replay_capacity to ReplayBuffer, allow not using
+  (some of) terminal episode transitions in training.
+  """
+
+  def __init__(self, replay_capacity, buffer_batch_size,
+               generates_trainable_dones, **kwargs):
+    self._replay_capacity = replay_capacity
+    self._buffer_batch_size = buffer_batch_size
+    self._generates_trainable_dones = generates_trainable_dones
+    super(_RainbowAgent, self).__init__(**kwargs)
+
+  def _build_replay_buffer(self, use_staging):
+    """Build WrappedReplayBuffer with custom OutOfGraphReplayBuffer."""
+    replay_buffer_kwargs = dict(
+        observation_shape=dqn_agent.NATURE_DQN_OBSERVATION_SHAPE,
+        stack_size=dqn_agent.NATURE_DQN_STACK_SIZE,
+        replay_capacity=self._replay_capacity,
+        batch_size=self._buffer_batch_size,
+        update_horizon=self.update_horizon,
+        gamma=self.gamma,
+        extra_storage_types=None,
+        observation_dtype=np.uint8,
+    )
+
+    replay_memory = _OutOfGraphPrioritizedReplayBuffer(
+        artificial_done=not self._generates_trainable_dones,
+        **replay_buffer_kwargs)
+
+    return _WrappedPrioritizedReplayBuffer(
+        wrapped_memory=replay_memory,
+        use_staging=use_staging, batch_size=self._buffer_batch_size)
+    # **replay_buffer_kwargs)
+
+
+class BatchRainbowAgent(_RainbowAgent):
+  """Batch agent for DQN.
+
+  Episodes are stored on done.
+
+  Assumes that all rollouts in batch would end at the same moment.
+  """
+
+  def __init__(self, env_batch_size, *args, **kwargs):
+    super(BatchRainbowAgent, self).__init__(*args, **kwargs)
+    self.env_batch_size = env_batch_size
+    obs_size = dqn_agent.NATURE_DQN_OBSERVATION_SHAPE
+    state_shape = [self.env_batch_size, obs_size[0], obs_size[1],
+                   dqn_agent.NATURE_DQN_STACK_SIZE]
+    self.state_batch = np.zeros(state_shape)
+    self.state = None  # assure it will be not used
+    self._observation = None  # assure it will be not used
+    self.reset_current_rollouts()
+
+  def reset_current_rollouts(self):
+    self._current_rollouts = [[] for _ in range(self.env_batch_size)]
+
+  def _record_observation(self, observation_batch):
+    # Set current observation. Represents an (batch_size x 84 x 84 x 1) image
+    # frame.
+    observation_batch = np.array(observation_batch)
+    self._observation_batch = observation_batch[:, :, :, 0]
+    # Swap out the oldest frames with the current frames.
+    self.state_batch = np.roll(self.state_batch, -1, axis=3)
+    self.state_batch[:, :, :, -1] = self._observation_batch
+
+  def _reset_state(self):
+    self.state_batch.fill(0)
+
+  def begin_episode(self, observation):
+    self._reset_state()
+    self._record_observation(observation)
+
+    if not self.eval_mode:
+      self._train_step()
+
+    self.action = self._select_action()
+    return self.action
+
+  def _update_current_rollouts(self, last_observation, action, reward,
+                               are_terminal):
+    transitions = zip(last_observation, action, reward, are_terminal)
+    for transition, rollout in zip(transitions, self._current_rollouts):
+      rollout.append(transition)
+
+  def _store_current_rollouts(self):
+    for rollout in self._current_rollouts:
+      for transition in rollout:
+        self._store_transition(*transition)
+    self.reset_current_rollouts()
+
+  def step(self, reward, observation):
+    self._last_observation = self._observation_batch
+    self._record_observation(observation)
+
+    if not self.eval_mode:
+      self._update_current_rollouts(self._last_observation, self.action, reward,
+                                    [False] * self.env_batch_size)
+      # We want to have the same train_step:env_step ratio not depending on
+      # batch size.
+      for _ in range(self.env_batch_size):
+        self._train_step()
+
+    self.action = self._select_action()
+    return self.action
+
+  def end_episode(self, reward):
+    if not self.eval_mode:
+      self._update_current_rollouts(
+          self._observation_batch, self.action, reward,
+          [True] * self.env_batch_size)
+      self._store_current_rollouts()
+
+  def _select_action(self):
+    epsilon = self.epsilon_eval
+    if not self.eval_mode:
+      epsilon = self.epsilon_fn(
+          self.epsilon_decay_period,
+          self.training_steps,
+          self.min_replay_history,
+          self.epsilon_train)
+
+    def choose_action(ix):
+      if random.random() <= epsilon:
+        # Choose a random action with probability epsilon.
+        return random.randint(0, self.num_actions - 1)
+      else:
+        # Choose the action with highest Q-value at the current state.
+        return self._sess.run(self._q_argmax,
+                              {self.state_ph: self.state_batch[ix:ix+1]})
+
+    return np.array([choose_action(ix) for ix in range(self.env_batch_size)])
+
+
 class BatchRunner(run_experiment.Runner):
   """Run a batch of environments.
 
@@ -223,7 +434,7 @@ def close(self):
     self._environment.close()
 
 
-class _OutOfGraphReplayBuffer(OutOfGraphReplayBuffer):
+class _OutOfGraphPrioritizedReplayBuffer(OutOfGraphPrioritizedReplayBuffer):
   """Replay not sampling artificial_terminal transition.
 
   Adds to stored tuples "artificial_done" field (as last ReplayElement).
@@ -240,34 +451,47 @@ class _OutOfGraphReplayBuffer(OutOfGraphReplayBuffer):
 
   def __init__(self, artificial_done, **kwargs):
     extra_storage_types = kwargs.pop("extra_storage_types", None) or []
+    assert not extra_storage_types, "Other extra_storage_types are " \
+                                    "currently not supported for this " \
+                                    "class."
     extra_storage_types.append(ReplayElement("artificial_done", (), np.uint8))
-    super(_OutOfGraphReplayBuffer, self).__init__(
+    super(_OutOfGraphPrioritizedReplayBuffer, self).__init__(
         extra_storage_types=extra_storage_types, **kwargs)
     self._artificial_done = artificial_done
 
   def is_valid_transition(self, index):
-    valid = super(_OutOfGraphReplayBuffer, self).is_valid_transition(index)
-    valid &= not self.get_artificial_done_stack(index).any()
+    valid = super(_OutOfGraphPrioritizedReplayBuffer, self).\
+        is_valid_transition(index)
+    if valid:
+      valid = not self.get_artificial_done_stack(index).any()
     return valid
 
   def get_artificial_done_stack(self, index):
     return self.get_range(self._store["artificial_done"],
                           index - self._stack_size + 1, index + 1)
 
-  def add(self, observation, action, reward, terminal, *args):
-    """Append artificial_done to *args and run parent method."""
+  def add(self, observation, action, reward, terminal, priority):
+    """Infer artificial_done and call parent method.
+
+    Note that OutOfGraphPrioritizedReplayBuffer (implicitly) assumes that
+    priority would be last argument in add. Here we write it explicitly.
+    Passing *args to this method is disabled on purpose, code start to gets to
+    convoluted with it.
+    """
     # If this will be a problem for maintenance, we could probably override
     # DQNAgent.add() method instead.
+    if not isinstance(priority, (float, np.floating)):
+      raise ValueError("priority should be float, got type {}"
+                       .format(type(priority)))
     artificial_done = self._artificial_done and terminal
-    args = list(args)
-    args.append(artificial_done)
-    return super(_OutOfGraphReplayBuffer, self).add(observation, action, reward,
-                                                    terminal, *args)
+    return super(_OutOfGraphPrioritizedReplayBuffer, self).add(
+        observation, action, reward, terminal, artificial_done, priority
+    )
 
   def load(self, *args, **kwargs):
     # Check that appropriate attributes are not overridden
     are_terminal_valid = self._artificial_done
-    super(_OutOfGraphReplayBuffer, self).load(*args, **kwargs)
+    super(_OutOfGraphPrioritizedReplayBuffer, self).load(*args, **kwargs)
     assert self._artificial_done == are_terminal_valid
 
 
@@ -280,6 +504,8 @@ def get_create_agent(agent_kwargs):
   Returns:
     Function(sess, environment, summary_writer) -> BatchDQNAgent instance.
   """
+  agent_kwargs = copy.deepcopy(agent_kwargs)
+  agent_type = agent_kwargs.pop("type")
 
   def create_agent(sess, environment, summary_writer=None):
     """Creates a DQN agent.
@@ -294,13 +520,24 @@ def create_agent(sess, environment, summary_writer=None):
     Returns:
       a DQN agent.
     """
-    return BatchDQNAgent(
-        env_batch_size=environment.batch_size,
-        sess=sess,
-        num_actions=environment.action_space.n,
-        summary_writer=summary_writer,
-        tf_device="/gpu:*",
-        **agent_kwargs)
+    if agent_type == "Rainbow":
+      return BatchRainbowAgent(
+          env_batch_size=environment.batch_size,
+          sess=sess,
+          num_actions=environment.action_space.n,
+          summary_writer=summary_writer,
+          tf_device="/gpu:*",
+          **agent_kwargs)
+    elif agent_type == "VanillaDQN":
+      return BatchDQNAgent(
+          env_batch_size=environment.batch_size,
+          sess=sess,
+          num_actions=environment.action_space.n,
+          summary_writer=summary_writer,
+          tf_device="/gpu:*",
+          **agent_kwargs)
+    else:
+      raise ValueError("Unknown agent_type {}".format(agent_type))
 
   return create_agent
 
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 4c5350fed..185d2b522 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -221,6 +221,14 @@ def rlmb_dqn_guess1():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_dqn_guess1_rainbow():
+  """rlmb_dqn guess1 params"""
+  hparams = rlmb_dqn_guess1()
+  hparams.set_hparam("base_algo_params", "dqn_guess1_rainbow_params")
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_dqn_guess1_2m_replay_buffer():
   """DQN guess1 params, 2M replay buffer."""

From f3c91c6440f8348343a9b904636f7f2aa7d2ed6e Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Wed, 19 Jun 2019 12:07:28 -0700
Subject: [PATCH 2140/2720] Merge of PR #1607

PiperOrigin-RevId: 254044123
---
 tensor2tensor/rl/dopamine_connector.py        | 34 ++++++++-----------
 .../rl/trainer_model_based_params.py          |  2 +-
 2 files changed, 15 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index a1ab21bc1..c3933ac8f 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -26,10 +26,10 @@
 from dopamine.agents.dqn import dqn_agent
 from dopamine.agents.rainbow import rainbow_agent
 from dopamine.replay_memory import circular_replay_buffer
-from dopamine.replay_memory.circular_replay_buffer import \
-    OutOfGraphReplayBuffer, ReplayElement
-from dopamine.replay_memory.prioritized_replay_buffer import \
-  OutOfGraphPrioritizedReplayBuffer, WrappedPrioritizedReplayBuffer
+from dopamine.replay_memory.circular_replay_buffer import OutOfGraphReplayBuffer
+from dopamine.replay_memory.circular_replay_buffer import ReplayElement
+from dopamine.replay_memory.prioritized_replay_buffer import OutOfGraphPrioritizedReplayBuffer
+from dopamine.replay_memory.prioritized_replay_buffer import WrappedPrioritizedReplayBuffer
 import numpy as np
 
 from tensor2tensor.rl.policy_learner import PolicyLearner
@@ -41,14 +41,16 @@
   import cv2
 except ImportError:
   cv2 = None
+
 try:
   from dopamine.discrete_domains import run_experiment
 except ImportError:
   run_experiment = None
+
 # pylint: enable=g-import-not-at-top
 # pylint: enable=ungrouped-imports
 
-# TODO: Vanilla DQN and Rainbow have a lot of common code. Most likely we want
+# TODO(rlmb): Vanilla DQN and Rainbow have a lot of common code. We will want
 #  to remove Vanilla DQN and only have Rainbow. To do so one needs to remove
 #  following:
 #    * _DQNAgent
@@ -59,6 +61,7 @@
 #      rlmb dqn hparams sets
 #  If we want to keep both Vanilla DQN and Rainbow, larger refactor is required.
 
+
 class _DQNAgent(dqn_agent.DQNAgent):
   """Modify dopamine DQNAgent to match our needs.
 
@@ -243,10 +246,8 @@ def load(self, *args, **kwargs):
 
 
 class _WrappedPrioritizedReplayBuffer(WrappedPrioritizedReplayBuffer):
-  """
+  """Allows to pass out-of-graph-replay-buffer via wrapped_memory."""
 
-  Allows to pass out-of-graph-replay-buffer via wrapped_memory.
-  """
   def __init__(self, wrapped_memory, batch_size, use_staging):
     self.batch_size = batch_size
     self.memory = wrapped_memory
@@ -451,17 +452,16 @@ class _OutOfGraphPrioritizedReplayBuffer(OutOfGraphPrioritizedReplayBuffer):
 
   def __init__(self, artificial_done, **kwargs):
     extra_storage_types = kwargs.pop("extra_storage_types", None) or []
-    assert not extra_storage_types, "Other extra_storage_types are " \
-                                    "currently not supported for this " \
-                                    "class."
+    msg = "Other extra_storage_types aren't currently supported for this class."
+    assert not extra_storage_types, msg
     extra_storage_types.append(ReplayElement("artificial_done", (), np.uint8))
     super(_OutOfGraphPrioritizedReplayBuffer, self).__init__(
         extra_storage_types=extra_storage_types, **kwargs)
     self._artificial_done = artificial_done
 
   def is_valid_transition(self, index):
-    valid = super(_OutOfGraphPrioritizedReplayBuffer, self).\
-        is_valid_transition(index)
+    valid = super(_OutOfGraphPrioritizedReplayBuffer,
+                  self).is_valid_transition(index)
     if valid:
       valid = not self.get_artificial_done_stack(index).any()
     return valid
@@ -471,13 +471,7 @@ def get_artificial_done_stack(self, index):
                           index - self._stack_size + 1, index + 1)
 
   def add(self, observation, action, reward, terminal, priority):
-    """Infer artificial_done and call parent method.
-
-    Note that OutOfGraphPrioritizedReplayBuffer (implicitly) assumes that
-    priority would be last argument in add. Here we write it explicitly.
-    Passing *args to this method is disabled on purpose, code start to gets to
-    convoluted with it.
-    """
+    """Infer artificial_done and call parent method."""
     # If this will be a problem for maintenance, we could probably override
     # DQNAgent.add() method instead.
     if not isinstance(priority, (float, np.floating)):
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 185d2b522..fb0fe1926 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -223,7 +223,7 @@ def rlmb_dqn_guess1():
 
 @registry.register_hparams
 def rlmb_dqn_guess1_rainbow():
-  """rlmb_dqn guess1 params"""
+  """Rainbow rlmb_dqn guess1 params."""
   hparams = rlmb_dqn_guess1()
   hparams.set_hparam("base_algo_params", "dqn_guess1_rainbow_params")
   return hparams

From 09bf4108186e2803b3307c8957a3be60fdaf16d0 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 19 Jun 2019 16:45:19 -0700
Subject: [PATCH 2141/2720] never increase batch size in update_hparams_for_tpu

PiperOrigin-RevId: 254098007
---
 tensor2tensor/models/transformer.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 368fbcb67..b397c9ebb 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -2428,8 +2428,9 @@ def update_hparams_for_tpu(hparams):
   # this hyperparameter is ignored.
   hparams.max_length = 64
 
-  # TPUs have less memory than GPUs, so decrease the batch size
-  hparams.batch_size = 2048
+  # TPUs have less memory than GPUs, so decrease the batch size if it's too high
+  if hparams.batch_size > 2048:
+    hparams.batch_size = 2048
 
   # Using noise broadcast in the dropout layers saves memory during training.
   hparams.attention_dropout_broadcast_dims = "0,1"  # batch, heads

From 6036fd86be3dd3a1e90d208694368e6a42aab3df Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 19 Jun 2019 17:31:45 -0700
Subject: [PATCH 2142/2720] Expand an error message with a known workaround;
 add flag to enable said workaround.

PiperOrigin-RevId: 254105522
---
 tensor2tensor/bin/t2t_decoder.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 72bda7b2b..88b45c230 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -58,6 +58,8 @@
 flags.DEFINE_string("score_file", "", "File to score. Each line in the file "
                     "must be in the format input \t target.")
 flags.DEFINE_bool("decode_in_memory", False, "Decode in memory.")
+flags.DEFINE_bool("disable_grappler_optimizations", False,
+                  "Disable Grappler if need be to avoid tensor format errors.")
 
 
 def create_hparams():
@@ -185,11 +187,14 @@ def main(_):
 
   hp = create_hparams()
   decode_hp = create_decode_hparams()
+  run_config = t2t_trainer.create_run_config(hp)
+  if FLAGS.disable_grappler_optimizations:
+    run_config.session_config.graph_options.rewrite_options.disable_meta_optimizer = True
 
   estimator = trainer_lib.create_estimator(
       FLAGS.model,
       hp,
-      t2t_trainer.create_run_config(hp),
+      run_config,
       decode_hparams=decode_hp,
       use_tpu=FLAGS.use_tpu)
 

From 0964f5c7b1d5f52f3c19e2a61a7fa4207635e71a Mon Sep 17 00:00:00 2001
From: Eugene Karaulov <EugKar@gmail.com>
Date: Thu, 20 Jun 2019 05:28:01 +0300
Subject: [PATCH 2143/2720] Storing encoder-decoder attention history at
 Transformer's cache during fast decoding. (#1602)

* Using partial targets at inference time.

* Saving attention history to Transformer's cache during fast decoding.
---
 tensor2tensor/models/transformer.py | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index b397c9ebb..fe3fcd4bc 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -27,6 +27,7 @@
 from __future__ import division
 from __future__ import print_function
 from six.moves import range  # pylint: disable=redefined-builtin
+import re
 
 from tensor2tensor.data_generators import librispeech
 from tensor2tensor.layers import common_attention
@@ -792,6 +793,23 @@ def preprocess_targets(targets, i):
       decoder_self_attention_bias += common_attention.attention_bias_proximal(
           decode_length)
 
+    # Create tensors for encoder-decoder attention history
+    att_cache = {"attention_history": {}}
+    num_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
+    att_batch_size, enc_seq_length = common_layers.shape_list(encoder_output)[0:2]
+    for layer in range(num_layers):
+      att_cache["attention_history"]["layer_%d" % layer] = tf.zeros(
+        [att_batch_size, hparams.num_heads, 0, enc_seq_length])
+
+    def update_decoder_attention_history(cache):
+      for k in filter(lambda x: "decoder" in x and not "self" in x and not "logits" in x,
+        self.attention_weights.keys()):
+        m = re.search(r"(layer_\d+)", k)
+        if m is None:
+          continue
+        cache["attention_history"][m[0]] = tf.concat(
+            [cache["attention_history"][m[0]], self.attention_weights[k]], axis=2)
+
     def symbols_to_logits_fn(ids, i, cache):
       """Go from ids to logits for next symbol."""
       ids = ids[:, -1:]
@@ -810,6 +828,8 @@ def symbols_to_logits_fn(ids, i, cache):
             cache,
             nonpadding=features_to_nonpadding(features, "targets"))
 
+      update_decoder_attention_history(cache)
+
       modality_name = hparams.name.get(
           "targets",
           modalities.get_name(target_modality))(hparams, target_vocab_size)
@@ -852,7 +872,8 @@ def forced_logits():
         batch_size=batch_size,
         force_decode_length=self._decode_hparams.force_decode_length,
         sos_id=sos_id,
-        eos_id=eos_id)
+        eos_id=eos_id,
+        cache=att_cache)
     if partial_targets is not None:
       if beam_size <= 1 or top_beams <= 1:
         ret["outputs"] = ret["outputs"][:, partial_targets_length:]

From 336d72f967895c5e0be583ca9fc2084b83f85004 Mon Sep 17 00:00:00 2001
From: Eugene Karaulov <EugKar@gmail.com>
Date: Wed, 19 Jun 2019 19:28:37 -0700
Subject: [PATCH 2144/2720] Merge of PR #1602

PiperOrigin-RevId: 254119136
---
 tensor2tensor/models/transformer.py | 33 +++++++++++++++++++----------
 1 file changed, 22 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index fe3fcd4bc..a33aed536 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -27,7 +27,6 @@
 from __future__ import division
 from __future__ import print_function
 from six.moves import range  # pylint: disable=redefined-builtin
-import re
 
 from tensor2tensor.data_generators import librispeech
 from tensor2tensor.layers import common_attention
@@ -796,19 +795,31 @@ def preprocess_targets(targets, i):
     # Create tensors for encoder-decoder attention history
     att_cache = {"attention_history": {}}
     num_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
-    att_batch_size, enc_seq_length = common_layers.shape_list(encoder_output)[0:2]
-    for layer in range(num_layers):
-      att_cache["attention_history"]["layer_%d" % layer] = tf.zeros(
-        [att_batch_size, hparams.num_heads, 0, enc_seq_length])
+    if encoder_output is not None:
+      att_batch_size, enc_seq_length = common_layers.shape_list(
+          encoder_output)[0:2]
+      for layer in range(num_layers):
+        att_cache["attention_history"]["layer_%d" % layer] = tf.zeros(
+            [att_batch_size, hparams.num_heads, 0, enc_seq_length])
 
     def update_decoder_attention_history(cache):
-      for k in filter(lambda x: "decoder" in x and not "self" in x and not "logits" in x,
-        self.attention_weights.keys()):
-        m = re.search(r"(layer_\d+)", k)
-        if m is None:
+      """Save attention weights in cache, e.g., for vizualization."""
+      for k in [x for x in self.attention_weights
+                if "decoder" in x and "self" not in x and "logits" not in x]:
+        idx = k.find("layer_")
+        if idx < 0:
           continue
-        cache["attention_history"][m[0]] = tf.concat(
-            [cache["attention_history"][m[0]], self.attention_weights[k]], axis=2)
+        # Get layer number from the string name.
+        layer_nbr = k[idx + 6:]
+        idx = 0
+        while idx + 1 < len(layer_nbr) and layer_nbr[:idx + 1].isdigit():
+          idx += 1
+        layer_nbr = "layer_%d" % int(layer_nbr[:idx])
+        if layer_nbr in cache["attention_history"]:
+          cache["attention_history"][layer_nbr] = tf.concat(
+              [cache["attention_history"][layer_nbr],
+               self.attention_weights[k]],
+              axis=2)
 
     def symbols_to_logits_fn(ids, i, cache):
       """Go from ids to logits for next symbol."""

From e6fbef39f0ca983034f1d4a17fdfbf0ccb3bfc02 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 19 Jun 2019 20:41:33 -0700
Subject: [PATCH 2145/2720] Reduce dimensionality of tensors in eval metrics on
 TPUs.

PiperOrigin-RevId: 254125422
---
 tensor2tensor/utils/t2t_model.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 547067692..284a7a081 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1989,6 +1989,18 @@ def create_dummy_vars():
 def create_tpu_eval_metrics_fn(problem, model_hparams):
   """Create the metrics_fn that TPUEstimatorSpec expects."""
 
+  def reduce_dimensions(predictions, labels):
+    """Reduce dimensions for high-dimensional predictions and labels."""
+    if len(predictions.get_shape()) > 5:
+      predictions_shape = common_layers.shape_list(predictions)
+      predictions = tf.reshape(
+          predictions, [predictions_shape[0], predictions_shape[1], -1,
+                        predictions_shape[-1]])
+      labels_shape = common_layers.shape_list(labels)
+      labels = tf.reshape(
+          labels, [labels_shape[0], labels_shape[1], -1])
+    return predictions, labels
+
   metric_fns = []
   eval_metrics = problem.eval_metric_fns(model_hparams)
 
@@ -1998,11 +2010,14 @@ def create_tpu_eval_metrics_fn(problem, model_hparams):
       weights_fn = modalities.get_weights_fn(v)
 
       def make_metric_fn(metric_fn):
+        """returns a metric_fn."""
         def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
           kwargs = {}
           args, _, keywords, _ = inspect.getargspec(metric_fn)
           if ("features" in args) or keywords:
             kwargs["features"] = features
+
+          logits, labels = reduce_dimensions(logits, labels)
           num, den = metric_fn(logits, labels, weights_fn=weights_fn, **kwargs)
           return tf.metrics.mean(num, den)
 
@@ -2018,11 +2033,14 @@ def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
     weights_fn = modalities.get_weights_fn(tm)
 
     def make_metric_fn(metric_fn):
+      """returns a metric fn."""
       def wrapped_metric_fn(logits, labels, features):
         kwargs = {}
         args, _, keywords, _ = inspect.getargspec(metric_fn)
         if ("features" in args) or keywords:
           kwargs["features"] = features
+
+        logits, labels = reduce_dimensions(logits, labels)
         num, den = metric_fn(logits, labels, weights_fn=weights_fn, **kwargs)
         return tf.metrics.mean(num, den)
 

From a32cf719fa62d8220fae0584eb9565a26feffba5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 19 Jun 2019 22:36:52 -0700
Subject: [PATCH 2146/2720] Bring back eval metrics for video problems.

PiperOrigin-RevId: 254135712
---
 tensor2tensor/data_generators/bair_robot_pushing.py | 3 ---
 tensor2tensor/data_generators/moving_mnist.py       | 3 ---
 2 files changed, 6 deletions(-)

diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index c9bcba37f..eadd3275f 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -79,9 +79,6 @@ def max_frames_per_video(self, hparams):
   def random_skip(self):
     return False
 
-  def eval_metrics(self):
-    return []
-
   @property
   def only_keep_videos_from_0th_frame(self):
     return True
diff --git a/tensor2tensor/data_generators/moving_mnist.py b/tensor2tensor/data_generators/moving_mnist.py
index 7d1502c67..f2140089f 100644
--- a/tensor2tensor/data_generators/moving_mnist.py
+++ b/tensor2tensor/data_generators/moving_mnist.py
@@ -79,9 +79,6 @@ def max_frames_per_video(self, hparams):
   def random_skip(self):
     return False
 
-  def eval_metrics(self):
-    return []
-
   @property
   def dataset_splits(self):
     """Splits of data to produce and number of output shards for each."""

From 848b14ba75ce4a2dd34229c92b16ae94ef32dacd Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 20 Jun 2019 09:57:19 -0700
Subject: [PATCH 2147/2720] Move closer to fully stack-based approach to inputs
 and outputs.   - Explicitly track n_inputs and n_outputs for each layer.   -
 Implement stack sementics via Serial combinator.   - Define/implement
 variable-width sublayer semantics for Parallel op.   - Remove lingering
 dependencies on non-sequence structure of args.   - Remove Select and Branch
 ops.   - Update all models to use the modified/restricted set of ops.

PiperOrigin-RevId: 254217127
---
 tensor2tensor/trax/layers/attention.py        |  39 +-
 tensor2tensor/trax/layers/base.py             | 233 +++++---
 tensor2tensor/trax/layers/combinators.py      | 545 ++++++++++--------
 tensor2tensor/trax/layers/combinators_test.py |  81 +--
 tensor2tensor/trax/layers/convolution.py      |   3 -
 tensor2tensor/trax/layers/core.py             |  14 +-
 tensor2tensor/trax/layers/core_test.py        |  34 +-
 tensor2tensor/trax/layers/rnn.py              |  22 +-
 tensor2tensor/trax/layers/rnn_test.py         |   4 +-
 tensor2tensor/trax/models/atari_cnn.py        |  24 +-
 .../research/position_lookup_transformer.py   |  23 +-
 .../models/research/transformer_revnet.py     |  10 +-
 tensor2tensor/trax/models/transformer.py      | 176 +++---
 tensor2tensor/trax/rlax/ppo.py                |  21 +-
 14 files changed, 706 insertions(+), 523 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 725713ced..f4d0d1eb1 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -24,7 +24,7 @@
 from tensor2tensor.trax import backend
 from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import combinators
+from tensor2tensor.trax.layers import combinators as cb
 from tensor2tensor.trax.layers import core
 
 
@@ -41,10 +41,10 @@ def PaddingMask(x, params, pad=0, **kwargs):
   return np.reshape(x != pad, (x.shape[0], 1, 1, x.shape[-1]))
 
 
-@base.layer(stack_items_to_pass=0)
+@base.layer(n_inputs=2)
 def EncoderDecoderMask(x, **unused_kwargs):
-  """Make encoder-decoder mask from a padding mask and decoder input."""
-  (padding_mask, decoder_input) = x
+  """Makes encoder-decoder mask from decoder input and a padding mask."""
+  decoder_input, padding_mask = x
   padding_mask = np.reshape(
       padding_mask, (padding_mask.shape[0], 1, 1, padding_mask.shape[-1]))
   # Final mask shape is [batch, 1 for heads, decoder-len, encoder-len].
@@ -110,7 +110,7 @@ def DotProductAttention(query, key, value, mask, dropout, mode, rng):
   return out
 
 
-@base.layer(stack_items_to_pass=4)
+@base.layer(n_inputs=4, n_outputs=2)
 def PureMultiHeadedAttention(x, params, n_heads=8, dropout=0.0, mode='train',
                              **kwargs):
   """Pure transformer-style multi-headed attention.
@@ -164,7 +164,7 @@ def MultiHeadedAttentionQKV(d_feature, n_heads=8, dropout=0.0, mode='train'):
     Multi-headed self-attention result and the mask.
   """
   return [
-      combinators.Parallel(
+      cb.Parallel(
           core.Dense(d_feature),
           core.Dense(d_feature),
           core.Dense(d_feature),
@@ -191,30 +191,17 @@ def MultiHeadedAttention(
     Multi-headed self-attention layer.
   """
   return [
-      combinators.Dup(),
-      combinators.Dup(),
+      cb.Dup(), cb.Dup(),
       MultiHeadedAttentionQKV(  # pylint: disable=no-value-for-parameter
           d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
   ]
 
 
-@base.layer(stack_items_to_pass=0)
+@base.layer()
 def ShiftRight(x, **unused_kwargs):
   """Layer to shift the tensor to the right by padding on axis 1."""
-  if not isinstance(x, (list, tuple)):  # non-chunked inputs
-    pad_widths = [(0, 0)] * len(x.shape)
-    pad_widths[1] = (1, 0)  # Padding on axis=1
-    padded = np.pad(x, pad_widths, mode='constant',
-                    constant_values=x.dtype.type(0))
-    return padded[:, :-1]
-  # Handling chunked inputs. Recall that the list of chunks represents a big
-  # sequence (the concatenation of the chunks). We want to shift that sequence,
-  # so we put a 0 in the beginning of the first chunk and the last element of
-  # that chunk is used as the new first element of the next chunk, and so on.
-  padded = []
-  last_value = np.zeros_like(x[0][:, -1])
-  for chunk in x:
-    padded_chunk = np.concatenate([last_value[:, np.newaxis], chunk], axis=1)
-    last_value = chunk[:, -1]
-    padded.append(padded_chunk[:, :-1])
-  return padded
+  pad_widths = [(0, 0)] * len(x.shape)
+  pad_widths[1] = (1, 0)  # Padding on axis=1
+  padded = np.pad(x, pad_widths, mode='constant',
+                  constant_values=x.dtype.type(0))
+  return padded[:, :-1]
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 58096f821..7c1eb9a46 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -30,29 +30,94 @@
 
 
 class Layer(object):
-  """Layer object, base class. Handles parameter sharing."""
+  """Base class for composable layers in a deep learning network.
+
+  A layer is a function from zero or more inputs to zero or more outputs,
+  possibly with trainable parameters. A layer is either atomic or composed
+  of sublayers. All layers provide accessor methods for these aspects:
+
+    - n_inputs()
+    - n_outputs()
+    - sublayers()
+
+  The inputs to a layer are activation tensors, packaged according to how many
+  there are:
+
+    - n_inputs = 0: an empty tuple ()
+    _ n_inputs = 1: the activation tensor (NOT wrapped in a tuple)
+    _ n_inputs > 1: a tuple of activation tensors
+
+  (The special treatment for the single-input case is intended as a
+  simplification for layer writers; this design choice may be revisited in the
+  future.)
+
+  The outputs from a layer are also activations tensors, packaged the same as
+  layer inputs:
+
+    - n_outputs = 0: an empty tuple ()
+    _ n_outputs = 1: the activation tensor (NOT wrapped in a tuple)
+    _ n_outputs > 1: a tuple of activation tensors
+
+  The runtime maintains a data stack with which layer calls are composed. One
+  can therefore view each layer as a function from stack state to stack state,
+  where the function's inputs are a slice from the stack, and the function's
+  outputs are spliced back into the stack.
+  """
 
   def __init__(self, **kwargs):
-    # We store kwargs by default, used below in creating a generic decorator.
-    self._init_kwargs = kwargs
-    # This field says if this layer's init has already been called or not.
-    self._first_init = True
-    # Cache parameters here, defaults empty params (we use () for that).
+    self._init_kwargs = kwargs  # can be used in creating a generic decorator
+    self._needs_init = True
     self._params = ()  # cached parameters
-    # Caller field storing info on where the caller class was created.
-    self._caller = _find_frame(inspect.stack())
+    self._caller = _find_frame(inspect.stack())  # for custom error messages
+
+  def __repr__(self):
+    class_str = self.__class__.__name__
+    fields_str = 'in={},out={}'.format(self.n_inputs(), self.n_outputs())
+    objs = self.sublayers()
+    if objs:
+      objs_str = ', '.join(str(x) for x in objs)
+      return '{}[{},layers=[{}]]'.format(class_str, fields_str, objs_str)
+    else:
+      return '{}[{}]'.format(class_str, fields_str)
+
+  def call(self, inputs, params=(), **kwargs):
+    """Applies this layer to given activation tensors, using trainable params.
 
-  def call(self, x, params=(), **kwargs):
-    """Call this layer in input x using the given parameters."""
+    Args:
+      inputs: Data tensors, matching the number (n_inputs) expected by this
+          layer. Specifically:
+            - n_inputs = 0: an empty tuple ()
+            - n_inputs = 1: a data tensor (NOT wrapped in a tuple)
+            - n_inputs > 1: a tuple of data tensors, with n_inputs items
+      params: A tuple of trainable parameters, with one element for this layer
+          and one for each of this layer's sublayers. If a layer (or sublayer)
+          has no trainable parameters, the corresponding params element is an
+          empty tuple.
+      **kwargs: Layer-specific keyword args.
+
+    Returns:
+      Data tensors, matching the number (n_outputs) promised by this layer.
+      Specifically:
+        - n_outputs = 0: an empty tuple
+        - n_outputs = 1: a data tensor (NOT wrapped in a tuple)
+        - n_outputs > 1: a tuple of data tensors, with n_outputs items
+      A tuple of activation tensors, one for each output.
+    """
     raise NotImplementedError
 
-  def new_parameters(self, input_shape, input_dtype, rng):
-    """Create new parameters for the layer given an input shape, dtype and rng.
+  def new_parameters(self, input_shapes, input_dtype, rng):
+    """Creates layer-specific parameters based on data shape, dtype and rng.
 
     Args:
-      input_shape: A tuple representing a shape (if this layer takes one input)
-          or a tuple of shapes (if this layer takes more than one input).
-          For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
+      input_shapes: A tuple, depending on the number of inputs (n_inputs)
+          expected by this layer:
+            - n_inputs = 0: an empty tuple ()
+            - n_inputs = 1: a tuple representing the shape of the input
+            - n_inputs > 1: a tuple of shape tuples, one for each input
+          For example:
+            - 0 inputs: ()
+            - 1 input: (210, 160, 3) [NOTE: no tuple wrapping the shape]
+            - 2 inputs: ((210, 160, 3), (105, 80, 3))
       input_dtype: numpy dtype of the input.
       rng: A random number generator.
 
@@ -61,19 +126,27 @@ def new_parameters(self, input_shape, input_dtype, rng):
     """
     raise NotImplementedError
 
-  # TODO(lukaszkaiser): re-visit the 2 items below in the future.
-  def stack_items_to_pass(self):
-    """How many of the top stack items do we process."""
-    return 0
+  def n_inputs(self):
+    """Specifies how many data tensors this layer expects as input."""
+    return 1  # Default is one input; subclasses can override.
+
+  def n_outputs(self):
+    """Specifies how many data tensors this layer promises as output."""
+    return 1  # Default is one output: subclasses can override.
+
+  def sublayers(self):
+    """Returns the sublayers contained in / managed by this layer."""
+    return ()  # Default is no sublayers; subclasses can override.
 
   # End of subclassing interface, all functions below are internal.
 
-  def pseudo_call(self, pseudo_input, params):
-    """Computes what shapes and types this layer would produce for given input.
+  def pseudo_call(self, pseudo_inputs, params):
+    """Computes shapes and types this layer would produce for the given inputs.
 
     Args:
-      pseudo_input: A ShapeType instance (input data minus the actual values)
-          or a tuple of ShapeType instances.
+      pseudo_inputs: A ShapeType instance (input data minus the actual values)
+          or a tuple of ShapeType instances, following the same conventions as
+          Layer.call's input arg.
       params: Parameters for this layer.
 
     Returns:
@@ -86,29 +159,24 @@ def pseudo_call(self, pseudo_input, params):
         # Same as backend.random.get_prng(0), but no op-by-op execution.
         rng = onp.zeros(2, onp.uint32)
         def call_on_input(x, params):
-          f = lambda y: self.call(y, params=params, rng=rng)
-          n = self.stack_items_to_pass() if isinstance(x, (list, tuple)) else 0
-          return _apply_to_first_n(f, x, n)
+          return self.call(x, params=params, rng=rng)
         params_shapes = nested_map(
             params, lambda x: ShapeType(shape=x.shape, dtype=x.dtype))
-        s = _eval_on_shapes(call_on_input, pseudo_input, params_shapes)
+        s = _eval_on_shapes(call_on_input, pseudo_inputs, params_shapes)
       return s
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback(skip=3)
-      raise LayerError(name, 'pseudo_call', self._caller, pseudo_input, trace)
+      raise LayerError(name, 'pseudo_call', self._caller, pseudo_inputs, trace)
 
-  def initialize(self, input_shape, input_dtype, rng):
+  def initialize(self, input_shapes, input_dtype, rng):
     """Initialize the layer given an input shape, dtype and rng.
 
-    Returns new_parameters(input_shape, rng) on the first call and () on any
+    Returns new_parameters(input_shapes, rng) on the first call and () on any
     subsequent call, as the layer is already initialized. This is used for
     networks that share parameters, so the layer only produces them once.
 
-    Note that all arguments and return values can be tuples or dictionaries
-    or arbitraty nested structures composed of tuples and dictionaries.
-
     Args:
-      input_shape: A tuple representing a shape (if this layer takes one input)
+      input_shapes: A tuple representing a shape (if this layer takes one input)
           or a tuple of shapes (if this layer takes more than one input).
           For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
       input_dtype: numpy dtype of the input.
@@ -118,23 +186,16 @@ def initialize(self, input_shape, input_dtype, rng):
       Newly created parameters on the first call and () on all subsequent calls.
     """
     try:
-      # Re-using this layer, no new parameters.
-      if not self._first_init:
+      # Initialize params once; store them for use when this layer is called.
+      if self._needs_init:
+        self._params = self.new_parameters(input_shapes, input_dtype, rng)
+        self._needs_init = False
+        return self._params
+      else:
         return ()
-
-      # First call of this layer, create parameters.
-      self._first_init = False
-      is_list = isinstance(input_shape, (list, tuple))
-      is_list = is_list and isinstance(input_shape[0], (list, tuple))
-      if is_list and self.stack_items_to_pass() > 0:
-        input_shape = input_shape[:self.stack_items_to_pass()]
-        if len(input_shape) == 1:
-          input_shape = input_shape[0]
-      self._params = self.new_parameters(input_shape, input_dtype, rng)
-      return self._params
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback(skip=3)
-      raise LayerError(name, 'initialize', self._caller, input_shape, trace)
+      raise LayerError(name, 'initialize', self._caller, input_shapes, trace)
 
   def __call__(self, x, params=(), **kwargs):
     try:
@@ -147,8 +208,7 @@ def __call__(self, x, params=(), **kwargs):
       # In this case, we're called for the first time: cache parameters.
       self._params = params
       f = lambda y: self.call(y, params=params, **kwargs)
-      n = self.stack_items_to_pass() if isinstance(x, (list, tuple)) else 0
-      return _apply_to_first_n(f, x, n)
+      return f(x)
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback()
       raise LayerError(name, 'call', self._caller, shapes(x), trace)
@@ -194,7 +254,7 @@ def message(self):
 
 # TODO(lukaszkaiser): remove this function once JAX has an analogue.
 def _eval_on_shapes(f, *args):
-  """Evaluate f given only shapes and types."""
+  """Evaluates f given only shapes and types."""
   def abstractify(x):
     return jax.abstract_arrays.raise_to_shaped(jax.core.get_aval(x))
 
@@ -327,28 +387,50 @@ def _short_traceback(skip=7):
   return '\n'.join(res)
 
 
-def layer(new_parameters=None, stack_items_to_pass=1):
+def _validate_call_input(x, n_inputs):
+  if n_inputs != 1:
+    if not isinstance(x, tuple):
+      raise TypeError(
+          'expected input to be a tuple; instead received {}'.format(type(x)))
+    if len(x) != n_inputs:
+      raise ValueError(
+          'input tuple length ({}) does not equal required number of inputs'
+          ' ({})'.format(len(x), n_inputs))
+
+
+def layer(new_parameters=None, n_inputs=1, n_outputs=1):
   """Decorates a function to make it the call method of a new Layer class."""
   # TODO(jonni): Consider renaming new_parameters to new_parameters_fn.
 
   def _build_layer_class(raw_call_fn):
     """Returns a Layer class built around the given call function."""
 
-    def _stack_items_to_pass(self):
+    def _n_inputs(self):
+      del self
+      return n_inputs
+
+    def _n_outputs(self):
       del self
-      return stack_items_to_pass
+      return n_outputs
 
-    def _new_parameters(self, input_shape, input_dtype, rng):
+    def _new_parameters(self, input_shapes, input_dtype, rng):
       if new_parameters is None:
         return ()
       kwargs = self._init_kwargs  # pylint: disable=protected-access
-      return new_parameters(input_shape, input_dtype, rng, **kwargs)
+      return new_parameters(input_shapes, input_dtype, rng, **kwargs)
+
+    def _is_empty(raw_output):
+      return raw_output is None or (isinstance(raw_output, (list, tuple))
+                                    and len(raw_output) == 0)  # pylint: disable=g-explicit-length-test
 
     def _call_with_context(self, x, params=(), **kwargs):
       """Calls raw_call_fn with extra keyword args from Layer.__init__."""
       merged_kwargs = kwargs.copy()
       merged_kwargs.update(self._init_kwargs)  # pylint: disable=protected-access
-      return raw_call_fn(x, params=params, **merged_kwargs)
+
+      _validate_call_input(x, n_inputs)
+      raw_output = raw_call_fn(x, params=params, **merged_kwargs)
+      return () if _is_empty(raw_output) else raw_output
 
     # Set docstrings and create the class.
     _call_with_context.__doc__ = raw_call_fn.__doc__
@@ -356,17 +438,18 @@ def _call_with_context(self, x, params=(), **kwargs):
     cls = type(raw_call_fn.__name__, (Layer,),
                {'call': _call_with_context,
                 'new_parameters': _new_parameters,
-                'stack_items_to_pass': _stack_items_to_pass})
+                'n_inputs': _n_inputs,
+                'n_outputs': _n_outputs})
     return cls
 
   return _build_layer_class
 
 
-def _random_inputs(input_shape, rng, integer_inputs=False):
+def _random_values(input_shapes, rng, integer_inputs=False):
   """Creates random floats or ints of the given shape.
 
   Args:
-    input_shape: A tuple representing a shape (if the layer takes one input)
+    input_shapes: A tuple representing a shape (if the layer takes one input)
         or a tuple of shapes (if this layer takes more than one input).
         For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
     rng: A random number generator.
@@ -376,17 +459,15 @@ def _random_inputs(input_shape, rng, integer_inputs=False):
   Returns:
     Random values with the shape and type specified.
   """
-  if isinstance(input_shape[0], int):
+  if isinstance(input_shapes[0], int):
     # Non-nested shape, create a random tuple.
     if not integer_inputs:
-      return backend.random.uniform(rng, input_shape, minval=-1.0, maxval=1.0)
-    return backend.random.bernoulli(rng, 0.5, input_shape).astype(onp.int32)
-  elif isinstance(input_shape, list):  # Nested shape: list.
-    return [_random_inputs(shape, rng, integer_inputs) for shape in input_shape]
-  elif isinstance(input_shape, tuple):  # Nested shape: tuple.
-    return tuple(_random_inputs(list(input_shape), rng, integer_inputs))
+      return backend.random.uniform(rng, input_shapes, minval=-1.0, maxval=1.0)
+    return backend.random.bernoulli(rng, 0.5, input_shapes).astype(onp.int32)
+  elif isinstance(input_shapes, tuple):  # Nested shape: tuple.
+    return tuple(_random_values(x, rng, integer_inputs) for x in input_shapes)
   else:
-    raise TypeError(type(input_shape))
+    raise TypeError(type(input_shapes))
 
 
 def _is_tuple_of_shapes(shape):
@@ -397,7 +478,7 @@ def _is_tuple_of_shapes(shape):
   return isinstance(shape, tuple) and isinstance(shape[0], tuple)
 
 
-def check_shape_agreement(layer_fn, input_shape, integer_inputs=False):
+def check_shape_agreement(layer_fn, input_shapes, integer_inputs=False):
   """Checks if the layer's call output agrees its pseudo_call predictions.
 
   This function helps test layer mechanics and inter-layer connections that
@@ -406,7 +487,7 @@ def check_shape_agreement(layer_fn, input_shape, integer_inputs=False):
   Args:
     layer_fn: A Layer instance, viewed as a function from input shapes to
         output shapes.
-    input_shape: A tuple representing a shape (if the layer takes one input)
+    input_shapes: A tuple representing a shape (if the layer takes one input)
         or a tuple of shapes (if this layer takes more than one input).
         For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
     integer_inputs: If True, use numpy int32 as the type for the pseudo-data,
@@ -418,19 +499,19 @@ def check_shape_agreement(layer_fn, input_shape, integer_inputs=False):
   """
   rng1, rng2, rng3 = backend.random.split(backend.random.get_prng(0), 3)
   input_dtype = onp.int32 if integer_inputs else onp.float32
-  if _is_tuple_of_shapes(input_shape):
-    pseudo_data = tuple(ShapeType(x, input_dtype) for x in input_shape)
-    input_dtype = tuple(input_dtype for _ in input_shape)
+  if _is_tuple_of_shapes(input_shapes):
+    pseudo_data = tuple(ShapeType(x, input_dtype) for x in input_shapes)
+    input_dtype = tuple(input_dtype for _ in input_shapes)
   else:
-    pseudo_data = ShapeType(input_shape, input_dtype)
-  params = layer_fn.initialize(input_shape, input_dtype, rng1)
+    pseudo_data = ShapeType(input_shapes, input_dtype)
+  params = layer_fn.initialize(input_shapes, input_dtype, rng1)
   pseudo_output = layer_fn.pseudo_call(pseudo_data, params)
   if isinstance(pseudo_output, tuple):
     output_shape = tuple(x.shape for x in pseudo_output)
   else:
     output_shape = pseudo_output.shape
 
-  random_input = _random_inputs(input_shape, rng2, integer_inputs)
+  random_input = _random_values(input_shapes, rng2, integer_inputs)
   real_output = layer_fn(random_input, params, rng=rng3)
   result_shape = shapes(real_output)
 
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 41025f45a..d2d83268e 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -19,9 +19,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import operator
-import six
-
 from tensor2tensor.trax import backend
 from tensor2tensor.trax.layers import base
 
@@ -41,13 +38,28 @@ def Model(*layers):
   return Serial(*layers)
 
 
-def _deep_flatten(xs):  # pylint: disable=invalid-name
-  for x in xs:
-    if isinstance(x, (list, tuple)):
-      for y in _deep_flatten(x):
-        yield y
-    else:
-      yield x
+def _deep_flatten(items):  # pylint: disable=invalid-name
+  """Returns a list of objects, flattening sublists/subtuples along the way.
+
+  Example: _deep_flatten([1, (2, 3, (4, 5), [6, 7]), [[[8]]]]) would return
+  the list [1, 2, 3, 4, 5, 6, 7, 8].
+
+  Args:
+    items: An iterable. If elements of this iterable are lists or tuples, they
+        will be (recursively) flattened until non-list non-tuple objects are
+        reached.
+
+  Returns:
+    A list of non-list, non-tuple objects.
+  """
+  def _flat_gen(xs):  # pylint: disable=invalid-name
+    for x in xs:
+      if isinstance(x, (list, tuple)):
+        for y in _flat_gen(x):
+          yield y
+      else:
+        yield x
+  return list(_flat_gen(items))
 
 
 def _ensure_sublayers(layers):  # pylint: disable=invalid-name
@@ -62,7 +74,7 @@ def _ensure_sublayers(layers):  # pylint: disable=invalid-name
     wrapped in Serial layer instances.
   """
   if not layers:  # None or an empty list can signal a no-op.
-    return Serial([])  # no-op, but still handles shapes and initialization
+    return Serial(None)  # no-op, but still handles shapes and initialization
   elif isinstance(layers, (list, tuple)):
     sublayers_not_lists = []
     for layer in layers:
@@ -80,190 +92,176 @@ def _pop_rng_and_split(args_dict, n_copies):  # pylint: disable=invalid-name
   return backend.random.split(rng, n_copies)
 
 
+def _count_items(xs):  # pylint: disable=invalid-name
+  return len(xs) if isinstance(xs, (list, tuple)) else 1
+
+
 class Serial(base.Layer):
-  """Layer composing a number of sub-layers in a serial way.."""
+  """Combinator that applies layers serially (by function composition).
+
+  A Serial combinator uses stack semantics to manage data for its sublayers.
+  Each sublayer sees only the inputs it needs and returns only the outputs it
+  has generated. The sublayers interact via the data stack. For instance, a
+  sublayer k, following sublayer j, gets called with the data stack in the
+  state left after layer j has applied. The Serial combinator then:
+
+    - takes N_in items off the top of the stack (N_in = k.n_inputs()) and calls
+      layer k, passing those items as arguments; and
+
+    - takes layer k's N_out return values (N_out = k.n_outputs()) and pushes
+      them onto the data stack.
+  """
 
   def __init__(self, *layers):
     super(Serial, self).__init__()
-    layers = list(_deep_flatten(layers))
-    # TODO(jonni): Consider flattening (unpacking) also embedded Serial layers.
-    self._layers = layers
-    self._nlayers = len(layers)
 
-  def call(self, x, params=(), **kwargs):
-    rngs = _pop_rng_and_split(kwargs, self._nlayers)
-    for layer, p, rng in zip(self._layers, params, rngs):
-      x = layer(x, p, rng=rng, **kwargs)
-    return x
+    layers = self._ensure_flat(layers)
+    self._sublayers = layers
+    self._n_layers = len(layers)
+
+    if not layers:
+      self._n_inputs = 1
+      self._n_outputs = 1
+    else:
+      self._n_inputs, self._n_outputs = self._n_inputs_n_outputs(layers)
+
+  def _ensure_flat(self, layers):
+    """Ensures that layers is a single flat list of Layer instances."""
+    del self
+    if len(layers) == 1 and layers[0] is None:
+      layers = []
+    else:
+      layers = _deep_flatten(layers)
+    for obj in layers:
+      if not isinstance(obj, base.Layer):
+        raise ValueError(
+            'Found nonlayer object ({}) in layers: {}.'.format(obj, layers))
+    return layers
+
+  def _n_inputs_n_outputs(self, layers):
+    del self
+    running_max = 0
+    running_total = 0
+    for layer in layers:
+      running_total += layer.n_inputs()
+      running_max = max(running_max, running_total)
+      running_total -= layer.n_outputs()
+    return running_max, (running_max - running_total)
+
+  def n_inputs(self):
+    return self._n_inputs
+
+  def n_outputs(self):
+    return self._n_outputs
+
+  def sublayers(self):
+    return self._sublayers
+
+  def _validate_call_inputs(self, xs):
+    if not isinstance(xs, tuple) and self._n_inputs != 1:
+      raise TypeError(
+          'Serial.call input must be a tuple; instead got {}'.format(xs))
+    if len(xs) < self.n_inputs():
+      raise ValueError(
+          'number of inputs ({}) to Serial.call less than n_inputs'
+          ' ({})'.format(len(xs), self.n_inputs()))
+
+  def call(self, xs, params=(), **kwargs):
+    self._validate_call_inputs(xs)
+    rngs = _pop_rng_and_split(kwargs, self._n_layers)
+    if not self._sublayers:  # No-op: leave args unchanged.
+      return xs
+
+    stack = xs
+    for layer, p, rng in zip(self._sublayers, params, rngs):
+      is_stack_just_one_item = (_count_items(stack) == 1)
+
+      # Give layer its args from the stack; treat 1-arg layer specially.
+      n_in = layer.n_inputs()
+      if n_in == 1 and is_stack_just_one_item:
+        inputs = stack
+      elif n_in == 1:
+        inputs = stack[0]
+      else:
+        inputs = stack[:n_in]
+      outputs = layer(inputs, p, rng=rng, **kwargs)
+
+      # Push outputs onto remaining stack (if any).
+      if n_in < _count_items(stack):
+        if layer.n_outputs() == 1:
+          outputs = (outputs,)
+        stack = outputs + stack[n_in:]
+      else:
+        stack = outputs  # NOTE: can be single value or tuple.
+
+    return stack
 
   def new_parameters(self, input_shape, input_dtype, rng):
     def MakeShapeType(shape, dtype):
       if isinstance(dtype, (list, tuple)):
         return tuple(MakeShapeType(s, t) for s, t in zip(shape, dtype))
       return base.ShapeType(shape=shape, dtype=dtype)
+
     params = []
-    pseudo_data = MakeShapeType(input_shape, input_dtype)
-    for layer in self._layers:
+    pseudo_xs = MakeShapeType(input_shape, input_dtype)
+    for layer in self._sublayers:
       rng, layer_rng = backend.random.split(rng)
-      cur_shape = base.nested_map(pseudo_data, lambda x: x.shape)
-      cur_dtype = base.nested_map(pseudo_data, lambda x: x.dtype)
-      param = layer.initialize(cur_shape, cur_dtype, layer_rng)
-      pparam = layer._params   # pylint: disable=protected-access
-      pseudo_data = layer.pseudo_call(pseudo_data, pparam)
-      params.append(param)
-    return params
 
+      # Give layer its args from pseudo_xs; treat 1-arg layer specially.
+      is_stack_just_one_item = (_count_items(pseudo_xs) == 1)
+      n_in = layer.n_inputs()
+      if n_in == 1 and is_stack_just_one_item:
+        inputs = pseudo_xs
+      elif n_in == 1:
+        inputs = pseudo_xs[0]
+      else:
+        inputs = pseudo_xs[:n_in]
 
-@base.layer(stack_items_to_pass=0)
-def PrintShape(x, message='PrintShape', **unused_kwargs):
-  """No-op layer that prints the shape of the stack."""
-  print(message + ' ; stack shape = ' + str(base.shapes(x)))
-  return x
+      in_shape = base.nested_map(inputs, lambda x: x.shape)
+      in_dtype = base.nested_map(inputs, lambda x: x.dtype)
+      param = layer.initialize(in_shape, in_dtype, layer_rng)
+      pparam = layer._params   # pylint: disable=protected-access
 
+      outputs = layer.pseudo_call(inputs, pparam)
 
-@base.layer(stack_items_to_pass=0)
-def Dup(x, **unused_kwargs):
-  """Duplicate (copy) the first element on the stack."""
-  if isinstance(x, list):
-    return [x[0]] + x
-  if isinstance(x, tuple):
-    return tuple([x[0]] + list(x))
-  return [x, x]
+      # Push outputs onto remaining pseudo_xs (if any).
+      if n_in < _count_items(pseudo_xs):
+        if layer.n_outputs() == 1:
+          outputs = (outputs,)
+        pseudo_xs = outputs + pseudo_xs[n_in:]
+      else:
+        pseudo_xs = outputs  # NOTE: can be single value or tuple.
+
+      params.append(param)
+    return params
 
 
-@base.layer(stack_items_to_pass=0)
-def Swap(x, **unused_kwargs):
-  """Swap the first two element on the stack."""
-  if isinstance(x, list):
-    return [x[1], x[0]] + x[2:]
-  assert isinstance(x, tuple)
-  return tuple([x[1], x[0]] + list(x[2:]))
+@base.layer(n_outputs=2)
+def Dup(x, **unused_kwargs):
+  """Duplicates (copies) an element."""
+  return (x, x)
 
 
-@base.layer(stack_items_to_pass=0)
-def _Top(x, **unused_kwargs):
-  """Top element from the stack."""
-  if isinstance(x, (list, tuple)):
-    return x[0]
-  return x
+@base.layer(n_inputs=2, n_outputs=2)
+def Swap(xs, **unused_kwargs):
+  """Swaps two elements."""
+  return (xs[1], xs[0])
 
 
-@base.layer(stack_items_to_pass=0)
+@base.layer(n_outputs=0)
 def Drop(x, **unused_kwargs):
-  """Drop first element from the stack."""
-  result = x[1:]
-  if len(result) == 1:
-    return result[0]
-  return result
+  """Drops one element."""
+  del x  # Just for the compiler.
+  return ()
 
 
-@base.layer(stack_items_to_pass=0)
+@base.layer(n_inputs=0)
 def FlattenList(xs, **unused_kwargs):
   """Flatten lists."""
+  # TODO(jonni): Consider renaming layer to DeepFlatten.
   return tuple(_deep_flatten(xs))
 
 
-# Re-ordering layer.
-class Select(base.Layer):
-  """Select elements from a tuple or create another tuple from them.
-
-  For example, we can re-order (x, y) into (y, x) or even (y, (x, y), y).
-  The output argument specifies how to re-order, using integers that refer
-  to indices in the input tuple. For example, if
-
-    input = (x, y, z)
-
-  then
-
-    Select(0)                = x
-    Select((1, 0, 2))        = (y, x, z)
-    Select((0, 0))           = (x, x)
-    Select((0, (1, 1)))      = (x, (y, y))
-    Select(((2, 0), (1, 1))) = ((z, x), (y, y))
-
-  By default (if no output is given) Select does nothing. It is also possible
-  to name the inputs to access tuple elements, e.g.:
-
-  Select(inputs=('encoder', ('decoder', 'mask')), output='decoder')
-
-  will transform a tuple (x, (y, x)) into y.
-
-  Args:
-    x: the input tuple to re-order.
-    params: layer parameters (unused).
-    output: the specification of the output tuple: a nested tuple of ints.
-    input: the specification of the input tuple if we need to disassemble it.
-    **kwargs: other arguments (unused).
-
-  Returns:
-    The re-ordered tuple with the same shape as output.
-  """
-
-  def __init__(self, output=None, inputs=None):
-    super(Select, self).__init__()
-    self._output = output
-    if inputs is None:
-      self._map = lambda x, i: x[i]
-    else:
-      self._input_map = {}
-      self._build_input_map(inputs, [])
-      def InputMapping(x, i):
-        cur = x
-        for idx in self._input_map[i]:
-          cur = cur[idx]
-        return cur
-      self._map = InputMapping
-
-  def _build_input_map(self, inputs, prefix):
-    for i, e in enumerate(inputs):
-      if isinstance(e, (list, tuple)):
-        self._build_input_map(e, prefix + [i])
-      else:
-        self._input_map[e] = prefix + [i]
-
-  def call(self, x, params=(), **kwargs):
-    del params, kwargs
-    if self._output is None:
-      return x
-    return base.nested_map(self._output, lambda i: self._map(x, i))
-
-  def new_parameters(self, input_shape, input_dtype, rng):
-    return ()
-
-
-class Branch(base.Layer):
-  """Combinator for applying layers to copies of the input.
-
-  This layer is often used to create parallel towers in neural networks:
-  * Branch(main, shortcut) -- start a residual tower (see Residual below)
-
-  Args:
-    *layers: a sequence of layers.
-
-  Returns:
-    A new layer in which each of the given layers has been applied to
-    a copy of the input independently.
-  """
-
-  def __init__(self, *layers):
-    super(Branch, self).__init__()
-    layers = _ensure_sublayers(layers)
-    self._nlayers = len(layers)
-    self._layers = layers
-
-  def call(self, x, params=(), **kwargs):
-    rngs = _pop_rng_and_split(kwargs, self._nlayers)
-    if isinstance(self._layers, (list, tuple)):
-      res = [layer(x, params=p, rng=r, **kwargs)
-             for layer, p, r in zip(self._layers, params, rngs)]
-      return tuple(res)
-
-  def new_parameters(self, input_shape, input_dtype, rng):
-    rngs = backend.random.split(rng, self._nlayers)
-    return [layer.initialize(input_shape, input_dtype, rng)
-            for layer, rng in zip(self._layers, rngs)]
-
-
 def _nested_op(inputs, op):  # pylint: disable=invalid-name
   """Helper: apply op over a list of arrays or nested arrays."""
   # If input is a dictionary, apply to the values (ignore keys).
@@ -281,50 +279,26 @@ def _nested_op(inputs, op):  # pylint: disable=invalid-name
   return tuple(result_list)
 
 
-def _binary_op(inputs, op):  # pylint: disable=invalid-name
-  """Helper: apply op to the first 2 elements."""
-  xs, rest = inputs[:2], inputs[2:]
-  s = _nested_op(xs, op)
-  if not rest:
-    return s
-  if not isinstance(s, (list, tuple)):
-    s = [s]
-  res = list(s) + list(rest)
-  # TODO(lukaszkaiser): should we drop this tuple/list distinction?
-  if isinstance(s, tuple):
-    res = tuple(res)
-  return res
-
+@base.layer(n_inputs=2)
+def Add(xs, **unused_kwargs):
+  """Adds two tensors."""
+  return xs[0] + xs[1]
 
-@base.layer(stack_items_to_pass=0)
-def Add(x, **unused_kwargs):
-  """Add first and second element on the stack."""
-  # Here x is a list of tensors of the same shape, or nested structures.
-  return _binary_op(x, op=sum)
 
+@base.layer(n_inputs=2)
+def SubtractTop(xs, **unused_kwargs):
+  """Subtracts the first tensor from the second."""
+  return xs[1] - xs[0]
 
-@base.layer(stack_items_to_pass=0)
-def SubtractTop(x, **unused_kwargs):
-  """Subtract the first element on the stack from the second element."""
-  # Here x is a list of tensors of the same shape, or nested structures.
-  return _binary_op(x, op=lambda xs: xs[1] - xs[0])
 
+@base.layer(n_inputs=2)
+def Multiply(xs, **unused_kwargs):
+  """Multiplies two tensors."""
+  return xs[0] * xs[1]
 
-@base.layer(stack_items_to_pass=0)
-def Multiply(x, **unused_kwargs):
-  """Multiply first and second element on the stack."""
-  return _binary_op(x, op=lambda xs: six.moves.reduce(operator.mul, xs))
 
-
-@base.layer(stack_items_to_pass=0)
-def AddAll(x, **unused_kwargs):
-  """Add branches elementwise."""
-  # Here x is a list of tensors of the same shape, or nested structures.
-  return _nested_op(x, op=sum)
-
-
-@base.layer(stack_items_to_pass=0)
-def Gate(x, **unused_kwargs):
+@base.layer(n_inputs=3)
+def Gate(xs, **unused_kwargs):
   """Implements a gating function on a (memory, gate, candidate) tuple.
 
   Final update is memory * gate + (1-gate) * candidate
@@ -333,70 +307,161 @@ def Gate(x, **unused_kwargs):
   Highway Networks: https://arxiv.org/abs/1505.00387
 
   Args:
-    x: A tuple of (memory, gate, candidate)
+    xs: A tuple of memory, gate, candidate
 
   Returns:
     The result of applying gating.
   """
-  assert len(x) == 3, x
-  state, gate, candidate = x
+  state, gate, candidate = xs
   return gate * state + (1.0 - gate) * candidate
 
 
-@base.layer(stack_items_to_pass=0)
-def Concatenate(x, params, axis=-1, **kwargs):
-  del params, kwargs
-  if isinstance(x, dict):  # For dictionaries, just use the values.
-    x = list(x.values())
-  return backend.numpy.concatenate(x, axis)
+class Concatenate(base.Layer):
+  """Concatenates n tensors into a single tensor."""
+
+  def __init__(self, n_items=2, axis=-1):
+    super(Concatenate, self).__init__()
+    self._n_items = n_items
+    self._axis = axis
+
+  def n_inputs(self):
+    return self._n_items
+
+  def new_parameters(self, input_shape, input_dtype, rng):
+    return ()
+
+  def call(self, xs, params=(), **kwargs):
+    del params, kwargs
+    return backend.numpy.concatenate(xs, self._axis)
 
 
 class Parallel(base.Layer):
-  """Combinator for applying layers to parts of a tuple.
+  """Combinator that applies a list of layers in parallel to its inputs.
+
+  Layers in the list apply to successive spans of inputs, where the spans are
+  determined how many inputs each layer takes. The resulting output is the
+  (flattened) concatenation of the resepective layer outputs.
+
+  For example, suppose one has three layers:
 
-  This layer is often used with the Branch and Add layers.
+    - F: 1 input, 1 output
+    - G: 3 inputs, 1 output
+    - H: 2 inputs, 2 outputs (h1, h2)
+
+  Then Parallel(F, G, H) will take 6 inputs and give 4 outputs:
+
+    - inputs: a, b, c, d, e, f
+    - outputs: F(a), G(b, c, d), h1, h2
+
+  As an important special case, a None argument to Parallel acts as if it takes
+  one argument, which it leaves unchanged. (It acts as a one-arg no-op.) For
+  example:
+
+    Parallel(None, F)
+
+  creates a layer that passes its first input unchanged and applies F to the
+  following input(s).
 
   Args:
-    *layers: a sequence of layers.
-    **kwlayers: a dictionary of layers.
+    *layers: A list of layers.
 
   Returns:
-    A new layer in which each of the given layers has been applied to
-    its corresponding argument in the input tuple or dictionary.
+    A new layer in which each of the given layers applies to its corresponding
+    span of elements in the dataflow stack.
   """
 
-  def __init__(self, *layers, **kwlayers):
+  def __init__(self, *layers):
     super(Parallel, self).__init__()
-    if layers and kwlayers:
-      raise ValueError('Cannot specify a Parallel with both a list and dict.')
-    layers = layers or kwlayers
-    layers = _ensure_sublayers(layers)
-    self._nlayers = len(layers)
-    self._layers = layers
-
-  def stack_items_to_pass(self):
-    return self._nlayers
+    layers = self._validate(layers)
+    self._n_layers = len(layers)
+    self._sublayers = layers
+    self._n_inputs = sum(x.n_inputs() for x in layers)
+    self._n_outputs = sum(x.n_outputs() for x in layers)
+
+  def _validate(self, layers):
+    if not layers or len(layers) < 2:
+      raise ValueError(
+          'layers ({}) must be a list with at least two elements'.format(
+              layers))
+    layers = list(layers)  # Ensure we can modify layers.
+    for i, obj in enumerate(layers):
+      if obj is None or obj == []:  # pylint: disable=g-explicit-bool-comparison
+        layers[i] = Serial(None)
+      elif isinstance(obj, (list, tuple)):
+        layers[i] = Serial(obj)
+      else:
+        if not isinstance(obj, base.Layer):
+          raise ValueError(
+              'Found nonlayer object ({}) in layers list: [{}].'.format(
+                  obj, layers))
+      if layers[i].n_inputs() == 0:
+        raise ValueError(
+            'Sublayer with n_inputs = 0 not allowed in Parallel:'
+            ' {}'.format(layers[i]))
+    return layers
+
+  def n_inputs(self):
+    return self._n_inputs
+
+  def n_outputs(self):
+    return self._n_outputs
+
+  def sublayers(self):
+    return self._sublayers
+
+  def _allot_to_sublayers(self, inputs):
+    """Divides Parallel's inputs for use by the sublayers.
+
+    Args:
+      inputs: Tuple of elements.
+
+    Returns:
+      A tuple that partitions this layer's inputs among its sublayers.
+      Sublayers that take one argument get that argument directly. All other
+      sublayers get a tuple of items.
+    """
+    start, end = 0, 0
+    sub_inputs = []
+    for layer in self._sublayers:
+      n_in = layer.n_inputs()
+      end = start + n_in
+      if n_in == 1:
+        sub_inputs.append(inputs[start])
+      else:
+        sub_inputs.append(inputs[start:end])
+      start = end
+    return tuple(sub_inputs)
 
   def call(self, inputs, params=(), **kwargs):
-    rngs = _pop_rng_and_split(kwargs, self._nlayers)
-    # Note that zip silently truncates its result if lengths don't match.
-    assert len(inputs) == self._nlayers
-    assert len(params) == self._nlayers
-    assert len(rngs) == self._nlayers
-    return tuple(layer(x, params=p, rng=r, **kwargs)
-                 for layer, x, p, r in zip(self._layers, inputs, params, rngs))
+    n_layers, layers = self._n_layers, self._sublayers
+    sublayer_inputs = self._allot_to_sublayers(inputs)
+    rngs = _pop_rng_and_split(kwargs, n_layers)
+    assert len(sublayer_inputs) == n_layers
+    assert len(params) == n_layers
+    assert len(rngs) == n_layers
+    outputs = []
+    for layer, x, p, r in zip(layers, sublayer_inputs, params, rngs):
+      # Note that zip silently truncates its result if lengths don't match.
+      sub_outputs = layer(x, params=p, rng=r, **kwargs)
+      if layer.n_outputs() == 1:
+        outputs.append(sub_outputs)
+      else:
+        outputs.extend(sub_outputs)
+    return outputs[0] if self.n_outputs() == 1 else tuple(outputs)
 
-  def new_parameters(self, input_shape, input_dtype, rng):
-    rngs = backend.random.split(rng, self._nlayers)
+  def new_parameters(self, input_shapes, input_dtypes, rng):
+    sublayer_shapes = self._allot_to_sublayers(input_shapes)
+    sublayer_dtypes = self._allot_to_sublayers(input_dtypes)
+    rngs = backend.random.split(rng, self._n_layers)
     return [layer.initialize(shape, dtype, rng) for layer, shape, dtype, rng
-            in zip(self._layers, input_shape, input_dtype, rngs)]
+            in zip(self._sublayers, sublayer_shapes, sublayer_dtypes, rngs)]
 
 
 def Residual(*layers, **kwargs):
   """Constructs a residual version of layers, summing input to layers output."""
-  shortcut = kwargs.get('shortcut', _Top())  # pylint: disable=no-value-for-parameter
+  shortcut = kwargs.get('shortcut')  # default None signals no-op
   return [
-      Branch(shortcut, Serial(layers)),  # Use Serial here to flatten layers.
-      FlattenList(),  # pylint: disable=no-value-for-parameter
+      Dup(),  # pylint: disable=no-value-for-parameter
+      Parallel(shortcut, layers),
       Add(),  # pylint: disable=no-value-for-parameter
   ]
diff --git a/tensor2tensor/trax/layers/combinators_test.py b/tensor2tensor/trax/layers/combinators_test.py
index 6f987d960..cf4d7907b 100644
--- a/tensor2tensor/trax/layers/combinators_test.py
+++ b/tensor2tensor/trax/layers/combinators_test.py
@@ -24,36 +24,22 @@
 from tensor2tensor.trax.layers import core
 
 
-_EMPTY_STACK = ()
-_REST_OF_STACK = ((1, 5), (4,))
-
-
 class CombinatorLayerTest(absltest.TestCase):
 
   def test_drop(self):
     layer = cb.Drop()
-    input_shape = ((3, 2),)
-    expected_shape = _EMPTY_STACK
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
-    input_shape = ((3, 2),) + _REST_OF_STACK
-    expected_shape = _REST_OF_STACK
+    input_shape = (3, 2)
+    expected_shape = ()
     output_shape = base.check_shape_agreement(layer, input_shape)
     self.assertEqual(output_shape, expected_shape)
 
   def test_dup(self):
     layer = cb.Dup()
-    input_shape = ((3, 2),)
+    input_shape = (3, 2)
     expected_shape = ((3, 2), (3, 2))
     output_shape = base.check_shape_agreement(layer, input_shape)
     self.assertEqual(output_shape, expected_shape)
 
-    input_shape = ((3, 2),) + _REST_OF_STACK
-    expected_shape = ((3, 2), (3, 2)) + _REST_OF_STACK
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
   def test_swap(self):
     layer = cb.Swap()
     input_shape = ((3, 2), (4, 7))
@@ -61,8 +47,10 @@ def test_swap(self):
     output_shape = base.check_shape_agreement(layer, input_shape)
     self.assertEqual(output_shape, expected_shape)
 
-    input_shape = ((3, 2), (4, 7)) + _REST_OF_STACK
-    expected_shape = ((4, 7), (3, 2)) + _REST_OF_STACK
+  def test_serial_no_op(self):
+    layer = cb.Serial(None)
+    input_shape = ((3, 2), (4, 7))
+    expected_shape = ((3, 2), (4, 7))
     output_shape = base.check_shape_agreement(layer, input_shape)
     self.assertEqual(output_shape, expected_shape)
 
@@ -73,43 +61,56 @@ def test_serial_no_op_list(self):
     output_shape = base.check_shape_agreement(layer, input_shape)
     self.assertEqual(output_shape, expected_shape)
 
-    input_shape = ((3, 2), (4, 7)) + _REST_OF_STACK
-    expected_shape = ((3, 2), (4, 7)) + _REST_OF_STACK
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
   def test_serial_one_in_one_out(self):
     layer = cb.Serial(core.Div(divisor=2.0))
-    input_shape = ((3, 2), (4, 7))
-    expected_shape = ((3, 2), (4, 7))
+    input_shape = (3, 2)
+    expected_shape = (3, 2)
     output_shape = base.check_shape_agreement(layer, input_shape)
     self.assertEqual(output_shape, expected_shape)
 
   def test_serial_div_div(self):
     layer = cb.Serial(core.Div(divisor=2.0), core.Div(divisor=5.0))
-    input_shape = ((3, 2), (4, 7))
-    expected_shape = ((3, 2), (4, 7))
+    input_shape = (3, 2)
+    expected_shape = (3, 2)
     output_shape = base.check_shape_agreement(layer, input_shape)
     self.assertEqual(output_shape, expected_shape)
 
-  def test_branch(self):
-    input_shape = (2, 3)
-    expected_shape = ((2, 3), (2, 3))
-    output_shape = base.check_shape_agreement(cb.Branch([], []), input_shape)
+  def test_serial_dup_dup(self):
+    layer = cb.Serial(cb.Dup(), cb.Dup())
+    input_shape = (3, 2)
+    expected_shape = ((3, 2), (3, 2), (3, 2))
+    output_shape = base.check_shape_agreement(layer, input_shape)
     self.assertEqual(output_shape, expected_shape)
 
-  def test_parallel(self):
-    input_shape = ((2, 3), (2, 3))
-    expected_shape = ((2, 3), (2, 3))
-    output_shape = base.check_shape_agreement(cb.Parallel([], []), input_shape)
+  def test_parallel_dup_dup(self):
+    layer = cb.Parallel(cb.Dup(), cb.Dup())
+    input_shape = ((3, 2), (4, 7))
+    expected_shape = ((3, 2), (3, 2), (4, 7), (4, 7))
+    output_shape = base.check_shape_agreement(layer, input_shape)
     self.assertEqual(output_shape, expected_shape)
 
-  def test_select(self):
-    input_shape = ((2, 3), (3, 4))
-    expected_shape = (3, 4)
-    output_shape = base.check_shape_agreement(cb.Select(1), input_shape)
+  def test_parallel_div_div(self):
+    layer = cb.Parallel(core.Div(divisor=0.5), core.Div(divisor=3.0))
+    input_shape = ((3, 2), (4, 7))
+    expected_shape = ((3, 2), (4, 7))
+    output_shape = base.check_shape_agreement(layer, input_shape)
     self.assertEqual(output_shape, expected_shape)
 
+  def test_parallel_no_ops(self):
+    layer = cb.Parallel([], None)
+    input_shape = ((3, 2), (4, 7))
+    expected_shape = ((3, 2), (4, 7))
+    output_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
+  def test_branch_op_not_defined(self):
+    with self.assertRaises(AttributeError):
+      cb.Branch([], [])
+
+  def test_select_op_not_defined(self):
+    input_shape = ((3, 2), (4, 7))
+    with self.assertRaises(AttributeError):
+      cb.Select(1, input_shape)
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/tensor2tensor/trax/layers/convolution.py b/tensor2tensor/trax/layers/convolution.py
index 879f66ad7..bd388a2ec 100644
--- a/tensor2tensor/trax/layers/convolution.py
+++ b/tensor2tensor/trax/layers/convolution.py
@@ -52,9 +52,6 @@ def __init__(self, filters, kernel_size, strides=None, padding='VALID',
       self._kernel_initializer = init.GlorotNormalInitializer(
           rhs_spec.index('O'), rhs_spec.index('I'))
 
-  def stack_items_to_pass(self):
-    return 1
-
   def _check_nhwc(self):
     msg = 'Convolutions on more than 4 dimensions only supported in NHWC.'
     assert self._lhs_spec == self._out_spec == 'NHWC', msg
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 8419dafbe..264423633 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -87,17 +87,14 @@ def ToFloat(x, **unused_kwargs):
 class Dense(base.Layer):
   """Layer constructor function for a dense (fully-connected) layer."""
 
-  def __init__(self, units,
+  def __init__(self, n_units,
                kernel_initializer=init.GlorotUniformInitializer(),
                bias_initializer=init.RandomNormalInitializer(1e-6)):
     super(Dense, self).__init__()
-    self._units = units
+    self._n_units = n_units
     self._kernel_initializer = kernel_initializer
     self._bias_initializer = bias_initializer
 
-  def stack_items_to_pass(self):
-    return 1
-
   def call(self, x, params, **kwargs):
     del kwargs
     w, b = params
@@ -106,8 +103,8 @@ def call(self, x, params, **kwargs):
   def new_parameters(self, input_shape, input_dtype, rng):
     del input_dtype
     rng1, rng2 = backend.random.split(rng, 2)
-    w = self._kernel_initializer((input_shape[-1], self._units), rng1)
-    b = self._bias_initializer((self._units,), rng2)
+    w = self._kernel_initializer((input_shape[-1], self._n_units), rng1)
+    b = self._bias_initializer((self._n_units,), rng2)
     return (w, b)
 
 
@@ -121,9 +118,6 @@ def __init__(self, d_feature, vocab_size,
     self._vocab_size = vocab_size
     self._kernel_initializer = kernel_initializer
 
-  def stack_items_to_pass(self):
-    return 1
-
   def call(self, x, params, **kwargs):
     del kwargs
     return np.take(params, x, axis=0)
diff --git a/tensor2tensor/trax/layers/core_test.py b/tensor2tensor/trax/layers/core_test.py
index 4f076f5fd..4409652db 100644
--- a/tensor2tensor/trax/layers/core_test.py
+++ b/tensor2tensor/trax/layers/core_test.py
@@ -31,20 +31,25 @@ class CoreLayerTest(absltest.TestCase):
   def test_flatten_n(self):
     input_shape = (29, 87, 10, 20, 30)
 
-    actual_shape = base.check_shape_agreement(core.Flatten(), input_shape)
-    self.assertEqual(actual_shape, (29, 87 * 10 * 20 * 30))
+    layer = core.Flatten()
+    expected_shape = (29, 87 * 10 * 20 * 30)
+    actual_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(actual_shape, expected_shape)
 
-    actual_shape = base.check_shape_agreement(
-        core.Flatten(n_axes_to_keep=2), input_shape)
-    self.assertEqual(actual_shape, (29, 87, 10 * 20 * 30))
+    layer = core.Flatten(n_axes_to_keep=2)
+    expected_shape = (29, 87, 10 * 20 * 30)
+    actual_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(actual_shape, expected_shape)
 
-    actual_shape = base.check_shape_agreement(
-        core.Flatten(n_axes_to_keep=3), input_shape)
-    self.assertEqual(actual_shape, (29, 87, 10, 20 * 30))
+    layer = core.Flatten(n_axes_to_keep=3)
+    expected_shape = (29, 87, 10, 20 * 30)
+    actual_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(actual_shape, expected_shape)
 
-    actual_shape = base.check_shape_agreement(
-        core.Flatten(n_axes_to_keep=4), input_shape)
-    self.assertEqual(actual_shape, (29, 87, 10, 20, 30))
+    layer = core.Flatten(n_axes_to_keep=4)
+    expected_shape = (29, 87, 10, 20, 30)
+    actual_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(actual_shape, expected_shape)
 
     # Not enough dimensions.
     with self.assertRaises(base.LayerError):
@@ -64,6 +69,13 @@ def test_div(self):
         onp.sum((output_np - expected_output_np) ** 2),
         delta=1e-6)
 
+  def test_div_shapes(self):
+    layer = core.Div(divisor=2.0)
+    input_shape = (3, 2)
+    expected_shape = (3, 2)
+    output_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(output_shape, expected_shape)
+
   def test_dense_param_sharing(self):
     model1 = combinators.Serial(core.Dense(32), core.Dense(32))
     layer = core.Dense(32)
diff --git a/tensor2tensor/trax/layers/rnn.py b/tensor2tensor/trax/layers/rnn.py
index 979dfeba8..f42ebfe35 100644
--- a/tensor2tensor/trax/layers/rnn.py
+++ b/tensor2tensor/trax/layers/rnn.py
@@ -24,32 +24,32 @@
 from tensor2tensor.trax.layers import core
 
 
-def GRUCell(units):
+def GRUCell(n_units):
   """Builds a traditional GRU cell with dense internal transformations.
 
   Gated Recurrent Unit paper: https://arxiv.org/abs/1412.3555
 
 
   Args:
-    units: Number of hidden units.
+    n_units: Number of hidden units.
 
   Returns:
     A Stax model representing a traditional GRU RNN cell.
   """
   return GeneralGRUCell(
-      candidate_transform=lambda: core.Dense(units=units),
+      candidate_transform=lambda: core.Dense(n_units),
       memory_transform_fn=None,
       gate_nonlinearity=core.Sigmoid,
       candidate_nonlinearity=core.Tanh)
 
 
-def ConvGRUCell(units, kernel_size=(3, 3)):
+def ConvGRUCell(n_units, kernel_size=(3, 3)):
   """Builds a convolutional GRU.
 
   Paper: https://arxiv.org/abs/1511.06432.
 
   Args:
-    units: Number of hidden units
+    n_units: Number of hidden units
     kernel_size: Kernel size for convolution
 
   Returns:
@@ -58,7 +58,7 @@ def ConvGRUCell(units, kernel_size=(3, 3)):
 
   def BuildConv():
     return convolution.Conv(
-        filters=units, kernel_size=kernel_size, padding='SAME')
+        filters=n_units, kernel_size=kernel_size, padding='SAME')
 
   return GeneralGRUCell(
       candidate_transform=BuildConv,
@@ -111,7 +111,8 @@ def GeneralGRUCell(candidate_transform,
       gate_nonlinearity(),
   ]
   candidate_block = [
-      cb.Branch([], reset_block),
+      cb.Dup(),
+      reset_block,
       cb.Multiply(),  # Gate S{t-1} with sigmoid(candidate_transform(S{t-1}))
       candidate_transform(),  # Final projection + tanh to get Ct
       candidate_nonlinearity(),  # Candidate gate
@@ -120,7 +121,8 @@ def GeneralGRUCell(candidate_transform,
       core.Dropout(rate=dropout_rate_c)
   ]
   memory_transform = memory_transform_fn() if memory_transform_fn else []
-  return cb.Serial([
-      cb.Branch(memory_transform, gate_block, candidate_block),
+  return cb.Model(
+      cb.Dup(), cb.Dup(),
+      cb.Parallel(memory_transform, gate_block, candidate_block),
       cb.Gate(),
-  ])
+  )
diff --git a/tensor2tensor/trax/layers/rnn_test.py b/tensor2tensor/trax/layers/rnn_test.py
index 805961f2e..c88ab4b39 100644
--- a/tensor2tensor/trax/layers/rnn_test.py
+++ b/tensor2tensor/trax/layers/rnn_test.py
@@ -32,13 +32,13 @@ def _test_cell_runs(self, layer, input_shape, output_shape):
 
   def test_conv_gru_cell(self):
     self._test_cell_runs(
-        rnn.ConvGRUCell(units=9, kernel_size=(3, 3)),
+        rnn.ConvGRUCell(9, kernel_size=(3, 3)),
         input_shape=(8, 1, 7, 9),
         output_shape=(8, 1, 7, 9))
 
   def test_gru_cell(self):
     self._test_cell_runs(
-        rnn.GRUCell(units=9), input_shape=(8, 7, 9), output_shape=(8, 7, 9))
+        rnn.GRUCell(9), input_shape=(8, 7, 9), output_shape=(8, 7, 9))
 
 
 if __name__ == '__main__':
diff --git a/tensor2tensor/trax/models/atari_cnn.py b/tensor2tensor/trax/models/atari_cnn.py
index f26dd6cc1..de9029504 100644
--- a/tensor2tensor/trax/models/atari_cnn.py
+++ b/tensor2tensor/trax/models/atari_cnn.py
@@ -24,19 +24,18 @@
 
 def AtariCnn(hidden_sizes=(32, 32), output_size=128):
   """An Atari CNN."""
-  # Input's shape = (B, T, H, W, C)
+  # TODO(jonni): Include link to paper?
+  # Input shape: (B, T, H, W, C)
+  # Output shape: (B, T, output_size)
   return tl.Model(
       tl.ToFloat(),
       tl.Div(divisor=255.0),
-      # Have 4 copies of the input, each one shifted to the right by one.
-      tl.Branch(
-          [],
-          [tl.ShiftRight()],
-          [tl.ShiftRight(), tl.ShiftRight()],
-          [tl.ShiftRight(), tl.ShiftRight(), tl.ShiftRight()]
-      ),
-      # Concatenated on the last axis.
-      tl.Concatenate(axis=-1),  # (B, T, H, W, 4C)
+
+      # Set up 4 successive game frames, concatenated on the last axis.
+      tl.Dup(), tl.Dup(), tl.Dup(),
+      tl.Parallel(None, _shift_right(1), _shift_right(2), _shift_right(3)),
+      tl.Concatenate(n_items=4, axis=-1),  # (B, T, H, W, 4C)
+
       tl.Conv(hidden_sizes[0], (5, 5), (2, 2), 'SAME'),
       tl.Relu(),
       tl.Conv(hidden_sizes[1], (5, 5), (2, 2), 'SAME'),
@@ -44,5 +43,8 @@ def AtariCnn(hidden_sizes=(32, 32), output_size=128):
       tl.Flatten(n_axes_to_keep=2),  # B, T and rest.
       tl.Dense(output_size),
       tl.Relu(),
-      # Eventually this is shaped (B, T, output_size)
   )
+
+
+def _shift_right(n):  # pylint: disable=invalid-name
+  return [tl.ShiftRight()] * n
diff --git a/tensor2tensor/trax/models/research/position_lookup_transformer.py b/tensor2tensor/trax/models/research/position_lookup_transformer.py
index 9d55aa8c7..3f637dd92 100644
--- a/tensor2tensor/trax/models/research/position_lookup_transformer.py
+++ b/tensor2tensor/trax/models/research/position_lookup_transformer.py
@@ -43,7 +43,8 @@ def NewPositionalEncoding(x, positions=None, **kwargs):
   return res
 
 
-@tl.layer(stack_items_to_pass=0)
+# TODO(lukaszkaiser): This used to have stack_items_to_pass=0; fix as needed.
+@tl.layer()
 def CutPosition(xs, **unused_kwargs):
   """Splits x into a pair (x[:position], position)."""
   if not isinstance(xs, (list, tuple)):
@@ -113,12 +114,14 @@ def DeepFlatten(xs):
       yield x
 
 
-@tl.layer(stack_items_to_pass=0)
+# TODO(lukaszkaiser): This used to have stack_items_to_pass=0; fix as needed.
+@tl.layer()
 def Unnest(xs, **unused_kwargs):
   return [x for x in DeepFlatten(xs)]
 
 
-@tl.layer(stack_items_to_pass=0)
+# TODO(lukaszkaiser): This used to have stack_items_to_pass=0; fix as needed.
+@tl.layer()
 def ConcatenateN(xs, params, n=2, axis=-1, **kwargs):
   """Concatenate first N inputs (and output remainder as is if non-empty)."""
   del params, kwargs
@@ -156,6 +159,7 @@ def ApplyAndQueryPositions(layer, pos):
   return tl.Serial(
       tl.Dup(),
       CutPosition(),
+      # TODO(lukaszkaiser): Rewrite without using Select.
       tl.Select(tuple([0] + [(2, 1)]*n_heads)),
       tl.Parallel(*([layer] + pos)),
       Unnest(),
@@ -184,7 +188,8 @@ def LearnedQP(keys=None, values=None, binary=False):
   )
 
 
-@tl.layer(stack_items_to_pass=0)
+# TODO(lukaszkaiser): This used to have stack_items_to_pass=0; fix as needed.
+@tl.layer()
 def SoftmaxBranches(x_list_in, n_branches=2, **unused_kwargs):
   """Softmax xs.
 
@@ -228,7 +233,8 @@ def SumLearnedPick(positions):
   sub_values = np.array([positions[max(i - j, 0), :]
                          for j in range(l) for i in range(l)])
   return tl.Serial(
-      tl.Branch(
+      tl.Dup(), tl.Dup(), tl.Dup(), tl.Dup(),
+      tl.Parallel(
           LearnedQP(),
           LearnedQP(keys=succ_keys, values=succ_values),
           LearnedQP(keys=subtract_1_keys, values=subtract_1_values),
@@ -261,7 +267,7 @@ def MultiHeadedAttentionPosition(
       tl.PureMultiHeadedAttention(
           d_feature=d_feature, n_heads=n_heads,
           dropout=dropout, mode=mode),
-      tl.Select(0),  # Drop the mask.
+      tl.Parallel([], tl.Drop()),  # Drop the mask.
       CombineHeadsPos(h=n_heads),
       PreservePosition(tl.Dense(d_feature)),
   )
@@ -305,8 +311,9 @@ def DecoderLayer(positions,
   return [
       tl.Residual(  # Self-attention block.
           PreservePosition(tl.LayerNorm()),
-          tl.Branch([],  # activation for (q, k, v)
-                    tl.CausalMask(axis=-2)),  # attention mask
+          tl.Dup(),
+          tl.Parallel([],  # activation for (q, k, v)
+                      tl.CausalMask(axis=-2)),  # attention mask
           MultiHeadedAttentionPosition(positions,
                                        d_feature, n_heads=n_heads,
                                        dropout=dropout, mode=mode),
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index dd17354a3..4acc29a5e 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -146,13 +146,13 @@ def vjpfun(ct):
     return do_call(x, params, kwargs)
 
 
-@tl.layer(stack_items_to_pass=1)
+@tl.layer()
 def Split(x, params, sections=2, axis=-1, **kwargs):
   del params, kwargs
   return list(backend.numpy.split(x, sections, axis))
 
 
-@tl.layer(stack_items_to_pass=1)
+@tl.layer()
 def Duplicate(x, params, sections=2, **kwargs):
   del params, kwargs
   return [x for _ in range(sections)]
@@ -163,6 +163,7 @@ class ReversibleHalfResidual(ReversibleLayerMixin, tl.Serial):
 
   def __init__(self, residual_layers):
     self.compute_residual = tl.Serial([
+        # TODO(jonni): Rewrite without using Select.
         tl.Select(inputs=('x1_or_y1', 'x2'), output=('x2', 'x1_or_y1', 'x2')),
         tl.Parallel(residual_layers, [], []),
     ])
@@ -288,10 +289,11 @@ def DecoderBlock(d_feature, d_feedforward, n_heads, n_attention_chunks,
   """
   self_attention = [
       tl.LayerNorm(),
-      tl.Branch([], tl.CausalMask(axis=-2)),  # Create mask.
+      tl.Dup(),
+      tl.Parallel([], tl.CausalMask(axis=-2)),  # Create mask.
       tl.MultiHeadedAttention(
           d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
-      tl.Select(0),  # Drop mask.
+      tl.Parallel([], tl.Drop()),  # Drop mask.
       tl.Dropout(rate=dropout, mode=mode),
   ]
 
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 95ed69ed0..8865ed6c0 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -34,11 +34,11 @@ def FeedForward(d_feature, d_feedforward, dropout, mode):
 
 
 def EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
-  """Transformer encoder block.
+  """Returns a layer sequence that implements a Transformer encoder block.
 
-  The input to the encoder is a pair (embedded source, mask) where
-  the mask is created from the original source to prevent attending
-  to the padding part of the input.
+  The input to the layer sequence is a pair, (activations, mask), where the
+  mask was created from the original source tokens to prevent attending to the
+  padding part of the input.
 
   Args:
     d_feature: int:  depth of embedding
@@ -48,7 +48,8 @@ def EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
     mode: str: 'train' or 'eval'
 
   Returns:
-    the layer, returning a pair (activations, mask).
+    A sequence of layers that maps an (activations, mask) pair to an
+    (activations, mask) pair.
   """
   attention = [
       tl.LayerNorm(),
@@ -74,7 +75,9 @@ def TransformerEncoder(vocab_size,
                        dropout=0.1,
                        max_len=2048,
                        mode='train'):
-  """Transformer encoder.
+  """Returns a Transformer encoder model.
+
+  The input to the model is a tensor of tokens.
 
   Args:
     vocab_size: int: vocab size
@@ -88,27 +91,31 @@ def TransformerEncoder(vocab_size,
     mode: str: 'train' or 'eval'
 
   Returns:
-    the Transformer encoder layer.
+    A Transformer model as a layer that maps from a tensor of tokens to
+    activations over a set of output classes.
   """
-  positional_embedder = [
+  embedder = [
       tl.Embedding(d_feature, vocab_size),
       tl.Dropout(rate=dropout, mode=mode),
       tl.PositionalEncoding(max_len=max_len),
   ]
-  return tl.Model([
-      tl.Branch(positional_embedder, tl.PaddingMask()),  # Create mask.
+  return tl.Model([                             #      tokens
+      tl.Dup(),                                 # toks toks
+      tl.Parallel(embedder, tl.PaddingMask()),  # vecs mask
       [EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode)
-       for _ in range(n_layers)],
-      tl.Select(0),  # Drop mask.
-      tl.LayerNorm(),
-      tl.Mean(axis=1),  # Average on length.
-      tl.Dense(n_classes),
-      tl.LogSoftmax(),
+       for _ in range(n_layers)],               # vecs mask
+      tl.Parallel([], tl.Drop()),               # ____  0
+      tl.LayerNorm(),                           # vecs
+      tl.Mean(axis=1),  # Average on length.    # vecs
+      tl.Dense(n_classes),                      # vecs
+      tl.LogSoftmax(),                          # vecs
   ])
 
 
 def DecoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
-  """Transformer decoder layer.
+  """Returns a layer sequence that implements a Transformer decoder block.
+
+  The input to the layer sequence is an activation tensor.
 
   Args:
     d_feature: int:  depth of embedding
@@ -118,15 +125,16 @@ def DecoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
     mode: str: 'train' or 'eval'
 
   Returns:
-    the layer.
+    A sequence of layers that maps an activation tensor to an activation tensor.
   """
   self_attention = [
-      tl.LayerNorm(),
-      tl.Branch([], tl.CausalMask(axis=-2)),  # Create mask.
+      tl.LayerNorm(),  # vec
+      tl.Dup(),  # vec vec
+      tl.Parallel([], tl.CausalMask(axis=-2)),  # vec mask
       tl.MultiHeadedAttention(
           d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
-      tl.Select(0),  # Drop mask.
-      tl.Dropout(rate=dropout, mode=mode),
+      tl.Parallel([], tl.Drop()),  # vec
+      tl.Dropout(rate=dropout, mode=mode),  # vec
   ]
   feed_forward = [
       FeedForward(d_feature, d_feedforward, dropout, mode=mode),
@@ -145,7 +153,10 @@ def TransformerLM(vocab_size,
                   dropout=0.1,
                   max_len=2048,
                   mode='train'):
-  """Transformer language model (only uses the decoder part of Transformer).
+  """Returns a Transformer language model.
+
+  The input to the model is a tensor of tokens. (This model uses only the
+  decoder part of the overall Transformer.)
 
   Args:
     vocab_size: int: vocab size
@@ -158,30 +169,31 @@ def TransformerLM(vocab_size,
     mode: str: 'train' or 'eval'
 
   Returns:
-    the layer.
+    A Transformer language model as a layer that maps from a tensor of tokens
+    to activations over a vocab set.
   """
-  positional_embedder = [
+  embedder = [
       tl.Embedding(d_feature, vocab_size),
       tl.Dropout(rate=dropout, mode=mode),
       tl.PositionalEncoding(max_len=max_len),
   ]
-  return tl.Model(
-      tl.ShiftRight(),
-      positional_embedder,
+  return tl.Model(                  # tokens
+      tl.ShiftRight(),              # toks
+      embedder,                     # vecs
       [DecoderBlock(d_feature, d_feedforward, n_heads, dropout, mode)
-       for _ in range(n_layers)],
-      tl.LayerNorm(),
-      tl.Dense(vocab_size),
-      tl.LogSoftmax(),
+       for _ in range(n_layers)],   # vecs
+      tl.LayerNorm(),               # vecs
+      tl.Dense(vocab_size),         # vecs
+      tl.LogSoftmax(),              # vecs
   )
 
 
 def EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode):
   """Transformer encoder-decoder layer.
 
-  The input is a triple pair (decoder_input, mask, encoder) where
-  the mask is created from the original source to prevent attending
-  to the padding part of the encoder.
+  The input is a triple (decoder_input, mask, encoder) where the mask is
+  created from the original source to prevent attending to the padding part
+  of the encoder.
 
   Args:
     d_feature: int:  depth of embedding
@@ -193,31 +205,30 @@ def EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode):
   Returns:
     the layer, returning a triple (decoder_activations, mask, encoder).
   """
-  decoder_self_attention = [
-      # TODO(jonni): Work on combinators so that this flow is cleaner/clearer.
-      tl.LayerNorm(),
-      tl.Dup(),
-      tl.CausalMask(axis=-2),  # Create the self-attention mask.
-      tl.Swap(),  # Put mask behind the activations.
+  decoder_self_attention = [                    #        vecs_d   pmask vecs_e
+      tl.LayerNorm(),                           #        vecs_d   ..... ......
+      tl.Dup(),                                 # vecs_d vecs_d   ..... ......
+      tl.Parallel([], tl.CausalMask(axis=-2)),  # ______ masks    ..... ......
       tl.MultiHeadedAttention(d_feature, n_heads=n_heads,
                               dropout=dropout, mode=mode),
-      tl.Swap(),  # Put self-attention mask on top.
-      tl.Drop(),   # Drop self-attention mask.
-      tl.Dropout(rate=dropout, mode=mode),
+      tl.Parallel([], tl.Drop()),               # ______   0      ..... ......
+      tl.Dropout(rate=dropout, mode=mode),      # vecs_d          ..... ......
   ]
-  decoder_to_encoder_attention = [
-      tl.Select((0, 2, 2, 1, 2)),  # (dec, enc, enc, mask, enc-copy)
-      tl.MultiHeadedAttentionQKV(  # (q, k, v, mask, ...) --> (new, mask, ...)
+  decoder_to_encoder_attention = [        # vecs_d        masks         vecs_e
+      tl.Parallel([], [], tl.Dup()),      # ______        _____  vecs_e vecs_e
+      tl.Parallel([], tl.Swap()),         # ______        vecs_e masks  ......
+      tl.Parallel([], tl.Dup()),          # ______ vecs_e vecs_e .....  ......
+      tl.MultiHeadedAttentionQKV(  # (q k v masks ... --> vecs_d masks ...)
           d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
-      tl.Dropout(rate=dropout, mode=mode),
+      tl.Dropout(rate=dropout, mode=mode),  # vecs_d mask vecs_e
   ]
   feed_forward = [
       FeedForward(d_feature, d_feedforward, dropout, mode=mode),
   ]
-  return [
-      tl.Residual(decoder_self_attention),
-      tl.Residual(decoder_to_encoder_attention),
-      tl.Residual(feed_forward),
+  return [                                        # vecs_d masks vecs_e
+      tl.Residual(decoder_self_attention),        # vecs_d masks vecs_e
+      tl.Residual(decoder_to_encoder_attention),  # vecs_d masks vecs_e
+      tl.Residual(feed_forward),                  # vecs_d masks vecs_e
   ]
 
 
@@ -230,9 +241,9 @@ def Transformer(vocab_size,
                 dropout=0.1,
                 max_len=2048,
                 mode='train'):
-  """Transformer.
+  """Returns a Transformer model.
 
-  This model expects on input a pair (source, target).
+  This model expects an input pair: target, source.
 
   Args:
     vocab_size: int: vocab size (shared source and target).
@@ -245,30 +256,43 @@ def Transformer(vocab_size,
     mode: str: 'train' or 'eval'
 
   Returns:
-    the Transformer model.
+    A Transformer model as a layer that maps from a target, source pair to
+    activations over a vocab set.
   """
-  positional_embedder = [
-      tl.Embedding(d_feature, vocab_size),
-      tl.Dropout(rate=dropout, mode=mode),
-      tl.PositionalEncoding(max_len=max_len),
+  embed = [                                    # tokens
+      tl.Embedding(d_feature, vocab_size),     # vecs
+      tl.Dropout(rate=dropout, mode=mode),     # vecs
+      tl.PositionalEncoding(max_len=max_len),  # vecs
   ]
-  encoder = [
-      tl.Branch(positional_embedder, tl.PaddingMask()),
+
+  encoder_stack = (  # masks vectors --> masks vectors
       [EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode)
-       for _ in range(n_layers)],
-      tl.LayerNorm(),
-  ]
-  return tl.Model(
-      tl.Parallel([], tl.ShiftRight()),
-      tl.Parallel(encoder, positional_embedder),
-      tl.Select(inputs=(('encoder', 'mask'), 'decoder'),
-                output=('decoder', ('mask', 'decoder'), 'encoder')),
-      # (encoder_mask, decoder_input) -> encoder-decoder mask
-      tl.Parallel([], tl.EncoderDecoderMask(), []),
+       for _ in range(n_layers)])
+
+  encoder_decoder_stack = (  # vecs_d masks vecs_e --> vecs_d masks vecs_e
       [EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode)
-       for _ in range(n_layers)],
-      tl.Select(0),  # Drop mask and encoder.
-      tl.LayerNorm(),
-      tl.Dense(vocab_size),
-      tl.LogSoftmax(),
+       for _ in range(n_layers)])
+
+  # Input: encoder_side_tokens, decoder_side_tokens
+  return tl.Model(  # tokens_e tokens_d
+      tl.Swap(),    # toks_d toks_e
+
+      # Encode.
+      tl.Parallel(                                    # toks_d        toks_e
+          [], [tl.Dup(),                              # ______ toks_e toks_e
+               tl.Parallel(embed, tl.PaddingMask()),  # ______ vecs_e masks
+               encoder_stack,                         # ______ vecs_e masks
+               tl.LayerNorm(),                        # ______ vecs_e .....
+               tl.Swap()]),                           # ______ masks vecs_e
+
+      # Decode.                                  #        toks_d masks vecs_e
+      tl.ShiftRight(),                           #        toks_d ..... ......
+      embed,                                     #        vecs_d ..... ......
+      tl.Dup(),                                  # vecs_d vecs_d ..... ......
+      tl.Parallel([], tl.EncoderDecoderMask()),  # ______    masks     ......
+      encoder_decoder_stack,                     # vecs_d    masks     vecs_e
+      tl.Parallel([], tl.Drop(), tl.Drop()),     # vecs_d
+      tl.LayerNorm(),                            # vecs_d
+      tl.Dense(vocab_size),                      # vecs_d
+      tl.LogSoftmax(),                           # vecs_d
   )
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 732cfe175..2f4008392 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -96,14 +96,23 @@ def policy_and_value_net(rng_key,
   # NOTE: The LogSoftmax instead of the Softmax because of numerical stability.
 
   if two_towers:
-    net = tl.Branch([bottom_layers_fn(),
-                     tl.Dense(n_actions),
-                     tl.LogSoftmax()],
-                    [bottom_layers_fn(), tl.Dense(1)])
+    layers = [
+        tl.Dup(),
+        tl.Parallel(
+            [bottom_layers_fn(), tl.Dense(n_actions), tl.LogSoftmax()],
+            [bottom_layers_fn(), tl.Dense(1)],
+        )
+    ]
   else:
-    net = tl.Serial(
+    layers = [
         bottom_layers_fn(),
-        tl.Branch([tl.Dense(n_actions), tl.LogSoftmax()], [tl.Dense(1)]))
+        tl.Dup(),
+        tl.Parallel(
+            [tl.Dense(n_actions), tl.LogSoftmax()],
+            [tl.Dense(1)],
+        )
+    ]
+  net = tl.Model(layers)
   params = net.initialize(batch_observations_shape, observations_dtype, rng_key)
   return params, net
 

From e6dacfa38ee65b545b9fdba9896a7231d544088f Mon Sep 17 00:00:00 2001
From: "Joshua V. Dillon" <jvdillon@google.com>
Date: Thu, 20 Jun 2019 19:22:05 -0700
Subject: [PATCH 2148/2720] Update use of TFP distributions' `probs`, `logits`
 properties to to use `probs_parameter`, `logits_parameter`. In the future
 properties `probs` `logits` will return `None` if that's how the distribution
 was parameterized.

PiperOrigin-RevId: 254318793
---
 tensor2tensor/layers/reversible_layers_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
index fb1ac4644..7db21c65c 100644
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -33,7 +33,7 @@
 def _log_prob(self, x):
   """Re-implementation of OneHotCategorical._log_prob for gradients wrt x."""
   x = self._assert_valid_sample(x)
-  logits = self.logits
+  logits = self.logits_parameter()
   if (not x.shape.is_fully_defined() or
       not logits.shape.is_fully_defined() or
       x.shape != logits.shape):

From f778c62ed37c00359ea4048c764289a146262b44 Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Thu, 20 Jun 2019 20:06:18 -0700
Subject: [PATCH 2149/2720] Fix a minor typo in problem.py

PiperOrigin-RevId: 254322989
---
 tensor2tensor/data_generators/problem.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index a8f12ee77..e1cb5d943 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -338,7 +338,7 @@ def example_reading_spec(self):
   def preprocess_example(self, example, mode, hparams):
     """Runtime preprocessing.
 
-    Return a dict or a tf.Data.Datset.from_tensor_slices (if you want each
+    Return a dict or a tf.data.Dataset.from_tensor_slices (if you want each
     example to turn into multiple).
 
     Args:

From 2a754e5bb56df55023e01b8c7e7ade807bf3a287 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 21 Jun 2019 22:20:36 -0700
Subject: [PATCH 2150/2720] * Modify similarity transformer model to predict
 embedding for both docstring and code. * Modify github function docstring
 problem to generate samples with "embed_code" feature, where 0 indicates that
 the input is docstring, and 1 for code.

PiperOrigin-RevId: 254522155
---
 .../data_generators/function_docstring.py     | 23 +++++++-----
 .../models/research/similarity_transformer.py | 36 ++++++++++++++++---
 2 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
index d2ce005ea..4fceebe8e 100644
--- a/tensor2tensor/data_generators/function_docstring.py
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -34,19 +34,22 @@ class GithubFunctionDocstring(text_problems.Text2TextProblem):
   ",".
   """
 
+  NUM_SHARDS = 100
+
   @property
   def base_url(self):
     return "gs://kubeflow-examples/t2t-code-search/raw_data"
 
   @property
   def pair_files_list(self):
-    return [
-        [
-            "{}/func-doc-pairs-000{:02}-of-00100.csv".format(self.base_url, i),
-            ("func-doc-pairs-000{:02}-of-00100.csv".format(i),)
-        ]
-        for i in range(100)
-    ]
+    files = []
+    for i in range(self.NUM_SHARDS):
+      files.append([
+          "{}/func-doc-pairs-{:05}-of-{:05}.csv".format(self.base_url, i,
+                                                        self.NUM_SHARDS),
+          ("func-doc-pairs-{:05}-of-{:05}.csv".format(i, self.NUM_SHARDS),)
+      ])
+    return files
 
   @property
   def is_generate_per_split(self):
@@ -89,7 +92,11 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
         for line in csv_file:
           reader = csv.reader(StringIO(line))
           for docstring_tokens, function_tokens in reader:
-            yield {"inputs": docstring_tokens, "targets": function_tokens}
+            yield {
+                "inputs": docstring_tokens,
+                "targets": function_tokens,
+                "embed_code": [0],
+            }
 
   def eval_metrics(self):
     return [
diff --git a/tensor2tensor/models/research/similarity_transformer.py b/tensor2tensor/models/research/similarity_transformer.py
index 4d316a961..0da345b19 100644
--- a/tensor2tensor/models/research/similarity_transformer.py
+++ b/tensor2tensor/models/research/similarity_transformer.py
@@ -37,10 +37,12 @@ def top(self, body_output, _):
     return body_output
 
   def body(self, features):
-    with tf.variable_scope('string_embedding'):
-      string_embedding = self.encode(features, 'inputs')
+    if self.hparams.mode != tf.estimator.ModeKeys.PREDICT:
+      # In training mode we need to embed both the queries and the code
+      # using the inputs and targets respectively.
+      with tf.variable_scope('string_embedding'):
+        string_embedding = self.encode(features, 'inputs')
 
-    if 'targets' in features:
       with tf.variable_scope('code_embedding'):
         code_embedding = self.encode(features, 'targets')
 
@@ -61,10 +63,34 @@ def body(self, features):
 
       loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels,
                                                      logits=logits)
+      return string_embedding_norm, {'training': loss}
+
+    # In predict mode we conditionally embed either the string query
+    # or the code based on the embed_code feature. In both cases the
+    # input will be in the inputs feature but the variable scope will
+    # be different
+    # Define predicates to be used with tf.cond
+    def embed_string():
+      with tf.variable_scope('string_embedding'):
+        string_embedding = self.encode(features, 'inputs')
+      return string_embedding
+
+    def embed_code():
+      with tf.variable_scope('code_embedding'):
+        code_embedding = self.encode(features, 'inputs')
+      return code_embedding
+
+    embed_code_feature = features.get('embed_code')
 
-      return string_embedding, {'training': loss}
+    # embed_code_feature will be a tensor because inputs will be a batch
+    # of inputs. We need to reduce that down to a single value for use
+    # with tf.cond; so we simply take the max of all the elements.
+    # This implicitly assume all inputs have the same value.
+    is_embed_code = tf.reduce_max(embed_code_feature)
+    result = tf.cond(is_embed_code > 0, embed_code, embed_string)
 
-    return string_embedding
+    result = tf.nn.l2_normalize(result)
+    return result
 
   def encode(self, features, input_key):
     hparams = self._hparams

From 902d49d04cfda135f718052297abafc674f7ff50 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Sat, 22 Jun 2019 11:33:53 -0700
Subject: [PATCH 2151/2720] Add Conv2DFlipout, DenseFlipout, LSTMCellFlipout.

conv2dflipout/denseflipout follow original source code. lstmcellflipout is mostly a cp+paste of tf.keras.layers.lstmcell, but with flipout perturbations.

PiperOrigin-RevId: 254569516
---
 tensor2tensor/layers/bayes.py      | 304 ++++++++++++++++++++++++++++-
 tensor2tensor/layers/bayes_test.py |  92 +++++++--
 2 files changed, 376 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index d14eaa2d1..aee45bfad 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -133,6 +133,71 @@ def call(self, *args, **kwargs):
     return super(Conv2DReparameterization, self).call(*args, **kwargs)
 
 
+class Conv2DFlipout(Conv2DReparameterization):
+  """2D convolution layer (e.g. spatial convolution over images).
+
+  The layer computes a variational Bayesian approximation to the distribution
+  over convolutional layers,
+
+  ```
+  p(outputs | inputs) = int conv2d(inputs; weights, bias) p(weights, bias)
+    dweights dbias.
+  ```
+
+  It does this with a stochastic forward pass, sampling from learnable
+  distributions on the kernel and bias. Gradients with respect to the
+  distributions' learnable parameters backpropagate via reparameterization.
+  Minimizing cross-entropy plus the layer's losses performs variational
+  minimum description length, i.e., it minimizes an upper bound to the negative
+  marginal likelihood.
+
+  This layer uses the Flipout estimator (Wen et al., 2018) for integrating with
+  respect to the `kernel`. Namely, it applies
+  pseudo-independent weight perturbations via independent sign flips for each
+  example, enabling variance reduction over independent weight perturbations.
+  For this estimator to work, the `kernel` random variable must be able
+  to decompose as a sum of its mean and a perturbation distribution; the
+  perturbation distribution must be independent across weight elements and
+  symmetric around zero (for example, a fully factorized Gaussian).
+  """
+
+  def call(self, inputs):
+    self.call_weights()
+    if not isinstance(self.kernel, ed.RandomVariable):
+      return super(Conv2DFlipout, self).call(inputs)
+    input_shape = tf.shape(inputs)
+    batch_dim = input_shape[0]
+    if self.data_format == 'channels_first':
+      channels = input_shape[1]
+      sign_input_shape = [batch_dim, channels, 1, 1]
+      sign_output_shape = [batch_dim, self.filters, 1, 1]
+    else:
+      channels = input_shape[-1]
+      sign_input_shape = [batch_dim, 1, 1, channels]
+      sign_output_shape = [batch_dim, 1, 1, self.filters]
+    sign_input = 2 * tf.random.uniform(sign_input_shape,
+                                       minval=0,
+                                       maxval=2,
+                                       dtype=inputs.dtype) - 1
+    sign_output = 2 * tf.random.uniform(sign_output_shape,
+                                        minval=0,
+                                        maxval=2,
+                                        dtype=inputs.dtype) - 1
+    kernel_mean = self.kernel.distribution.mean()
+    perturbation = self.kernel - kernel_mean
+    outputs = self._convolution_op(inputs, kernel_mean)
+    outputs += self._convolution_op(inputs * sign_input,
+                                    perturbation) * sign_output
+    if self.use_bias:
+      if self.data_format == 'channels_first':
+        outputs = tf.nn.bias_add(outputs, self.bias, data_format='NCHW')
+      else:
+        outputs = tf.nn.bias_add(outputs, self.bias, data_format='NHWC')
+    if self.activation is not None:
+      outputs = self.activation(outputs)
+    return outputs
+
+
 @add_weight
 class Conv2DVariationalDropout(tf.keras.layers.Conv2D):
   """2D convolution layer with variational dropout (Kingma et al., 2015).
@@ -207,7 +272,10 @@ def dropped_inputs():
           tf.keras.backend.epsilon())
       outputs = means + stddevs * tf.random_normal(tf.shape(stddevs))
       if self.use_bias:
-        outputs = tf.nn.bias_add(outputs, self.bias)
+        if self.data_format == 'channels_first':
+          outputs = tf.nn.bias_add(outputs, self.bias, data_format='NCHW')
+        else:
+          outputs = tf.nn.bias_add(outputs, self.bias, data_format='NHWC')
       if self.activation is not None:
         outputs = self.activation(outputs)
       return outputs
@@ -472,6 +540,65 @@ def call(self, *args, **kwargs):
     return super(DenseReparameterization, self).call(*args, **kwargs)
 
 
+class DenseFlipout(DenseReparameterization):
+  """Bayesian densely-connected layer estimated via Flipout (Wen et al., 2018).
+
+  The layer computes a variational Bayesian approximation to the distribution
+  over densely-connected layers,
+
+  ```
+  p(outputs | inputs) = int dense(inputs; weights, bias) p(weights, bias)
+    dweights dbias.
+  ```
+
+  It does this with a stochastic forward pass, sampling from learnable
+  distributions on the kernel and bias. Gradients with respect to the
+  distributions' learnable parameters backpropagate via reparameterization.
+  Minimizing cross-entropy plus the layer's losses performs variational
+  minimum description length, i.e., it minimizes an upper bound to the negative
+  marginal likelihood.
+
+  This layer uses the Flipout estimator (Wen et al., 2018) for integrating with
+  respect to the `kernel`. Namely, it applies
+  pseudo-independent weight perturbations via independent sign flips for each
+  example, enabling variance reduction over independent weight perturbations.
+  For this estimator to work, the `kernel` random variable must be able
+  to decompose as a sum of its mean and a perturbation distribution; the
+  perturbation distribution must be independent across weight elements and
+  symmetric around zero (for example, a fully factorized Gaussian).
+  """
+
+  def call(self, inputs):
+    self.call_weights()
+    if not isinstance(self.kernel, ed.RandomVariable):
+      return super(DenseFlipout, self).call(inputs)
+    input_shape = tf.shape(inputs)
+    sign_input = 2 * tf.random.uniform(input_shape,
+                                       minval=0,
+                                       maxval=2,
+                                       dtype=inputs.dtype) - 1
+    sign_output = 2 * tf.random.uniform(tf.concat([input_shape[:-1],
+                                                   [self.units]], 0),
+                                        minval=0,
+                                        maxval=2,
+                                        dtype=inputs.dtype) - 1
+    kernel_mean = self.kernel.distribution.mean()
+    perturbation = self.kernel - kernel_mean
+    if inputs.shape.ndims <= 2:
+      outputs = tf.matmul(inputs, kernel_mean)
+      outputs += tf.matmul(inputs * sign_input, perturbation) * sign_output
+    else:
+      outputs = tf.tensordot(inputs, kernel_mean, [[-1], [0]])
+      outputs += tf.tensordot(inputs * sign_input,
+                              perturbation,
+                              [[-1], [0]]) * sign_output
+    if self.use_bias:
+      outputs = tf.nn.bias_add(outputs, self.bias)
+    if self.activation is not None:
+      outputs = self.activation(outputs)
+    return outputs
+
+
 @add_weight
 class DenseVariationalDropout(tf.keras.layers.Dense):
   """Densely-connected layer with variational dropout (Kingma et al., 2015).
@@ -680,7 +807,6 @@ def call_weights(self):
     if isinstance(self.bias_initializer, tf.keras.layers.Layer):
       self.bias = self.bias_initializer(self.bias.shape, self.dtype)
 
-  # NOTE: This will not be called in TF < 1.11.
   def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
     """Get the initial state and side-effect sampling of stochastic weights."""
     if self.built:
@@ -689,6 +815,180 @@ def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
         inputs=inputs, batch_size=batch_size, dtype=dtype)
 
 
+class LSTMCellFlipout(LSTMCellReparameterization):
+  """Bayesian LSTM cell class estimated via Flipout (Wen et al., 2018).
+
+  The layer computes a variational Bayesian approximation to the distribution
+  over LSTM cell functions,
+
+  ```
+  p(outputs | inputs) = int lstm_cell(inputs; weights, bias) p(weights, bias)
+    dweights dbias,
+  ```
+
+  where the weights consist of both input and recurrent weights.
+
+  It does this with a stochastic forward pass, sampling from learnable
+  distributions on the kernel, recurrent kernel, and bias. Gradients with
+  respect to the distributions' learnable parameters backpropagate via
+  reparameterization.  Minimizing cross-entropy plus the layer's losses performs
+  variational minimum description length, i.e., it minimizes an upper bound to
+  the negative marginal likelihood.
+
+  This layer uses the Flipout estimator (Wen et al., 2018) for integrating with
+  respect to the `kernel` and `recurrent_kernel`. Namely, it applies
+  pseudo-independent weight perturbations via independent sign flips for each
+  example, enabling variance reduction over independent weight perturbations.
+  For this estimator to work, the `kernel` and `recurrent_kernel` random
+  variable must be able to decompose as a sum of its mean and a perturbation
+  distribution; the perturbation distribution must be independent across weight
+  elements and symmetric around zero (for example, a fully factorized Gaussian).
+  """
+
+  def _call_sign_flips(self, inputs=None, batch_size=None, dtype=None):
+    """Builds per-example sign flips for pseudo-independent perturbations."""
+    # TODO(trandustin): We add and call this method separately from build().
+    # This is because build() operates on a static input_shape. We need dynamic
+    # input shapes as we operate on the batch size which is often dynamic.
+    if inputs is not None:
+      batch_size = tf.shape(inputs)[0]
+      dtype = inputs.dtype
+    input_dim = tf.shape(self.kernel)[0]
+    self.sign_input = 2 * tf.random.uniform(
+        [batch_size, 4 * input_dim], minval=0, maxval=2, dtype=dtype) - 1
+    self.sign_output = 2 * tf.random.uniform(
+        [batch_size, 4 * self.units], minval=0, maxval=2, dtype=dtype) - 1
+    self.recurrent_sign_input = 2 * tf.random.uniform(
+        [batch_size, 4 * self.units], minval=0, maxval=2, dtype=dtype) - 1
+    self.recurrent_sign_output = 2 * tf.random.uniform(
+        [batch_size, 4 * self.units], minval=0, maxval=2, dtype=dtype) - 1
+
+  def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
+    """Get the initial state and side-effect sampling of stochastic weights."""
+    if self.built:
+      self._call_sign_flips(inputs, batch_size, dtype)
+    return super(LSTMCellFlipout, self).get_initial_state(
+        inputs=inputs, batch_size=batch_size, dtype=dtype)
+
+  def _compute_carry_and_output(self, x, h_tm1, c_tm1):
+    """Computes carry and output using split kernels."""
+    if not isinstance(self.recurrent_kernel, ed.RandomVariable):
+      return super(LSTMCellFlipout, self)._compute_carry_and_output(x,
+                                                                    h_tm1,
+                                                                    c_tm1)
+    x_i, x_f, x_c, x_o = x
+    h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o = h_tm1
+    kernel_mean = self.recurrent_kernel.distribution.mean()
+    perturbation = self.recurrent_kernel - kernel_mean
+    k_i, k_f, k_c, k_o = tf.split(kernel_mean, num_or_size_splits=4, axis=1)
+    p_i, p_f, p_c, p_o = tf.split(perturbation, num_or_size_splits=4, axis=1)
+    si_i, si_f, si_c, si_o = tf.split(self.recurrent_sign_input,
+                                      num_or_size_splits=4, axis=1)
+    so_i, so_f, so_c, so_o = tf.split(self.recurrent_sign_output,
+                                      num_or_size_splits=4, axis=1)
+    z0 = (x_i + tf.keras.backend.dot(h_tm1_i, k_i) +
+          tf.keras.backend.dot(h_tm1_i * si_i, p_i) * so_i)
+    z1 = (x_f + tf.keras.backend.dot(h_tm1_f, k_f) +
+          tf.keras.backend.dot(h_tm1_f * si_f, p_f) * so_f)
+    z2 = (x_c + tf.keras.backend.dot(h_tm1_c, k_c) +
+          tf.keras.backend.dot(h_tm1_c * si_c, p_c) * so_c)
+    z3 = (x_o + tf.keras.backend.dot(h_tm1_o, k_o) +
+          tf.keras.backend.dot(h_tm1_o * si_o, p_o) * so_o)
+    i = self.recurrent_activation(z0)
+    f = self.recurrent_activation(z1)
+    c = f * c_tm1 + i * self.activation(z2)
+    o = self.recurrent_activation(z3)
+    return c, o
+
+  def call(self, inputs, states, training=None):
+    # TODO(trandustin): Enable option for Flipout on only the kernel or
+    # recurrent_kernel. If only one is a random variable, we currently default
+    # to weight reparameterization.
+    if (not isinstance(self.kernel, ed.RandomVariable) or
+        not isinstance(self.recurrent_kernel, ed.RandomVariable)):
+      return super(LSTMCellFlipout, self).call(inputs, states, training)
+    if not hasattr(self, 'sign_input'):
+      self._call_sign_flips(inputs)
+    h_tm1 = states[0]  # previous memory state
+    c_tm1 = states[1]  # previous carry state
+
+    dp_mask = self.get_dropout_mask_for_cell(inputs, training, count=4)
+    rec_dp_mask = self.get_recurrent_dropout_mask_for_cell(
+        h_tm1, training, count=4)
+
+    if self.implementation == 1:
+      if 0 < self.dropout < 1.:
+        inputs_i = inputs * dp_mask[0]
+        inputs_f = inputs * dp_mask[1]
+        inputs_c = inputs * dp_mask[2]
+        inputs_o = inputs * dp_mask[3]
+      else:
+        inputs_i = inputs
+        inputs_f = inputs
+        inputs_c = inputs
+        inputs_o = inputs
+      kernel_mean = self.kernel.distribution.mean()
+      perturbation = self.kernel - kernel_mean
+      k_i, k_f, k_c, k_o = tf.split(kernel_mean, num_or_size_splits=4, axis=1)
+      p_i, p_f, p_c, p_o = tf.split(perturbation, num_or_size_splits=4, axis=1)
+      si_i, si_f, si_c, si_o = tf.split(self.sign_input,
+                                        num_or_size_splits=4, axis=1)
+      so_i, so_f, so_c, so_o = tf.split(self.sign_output,
+                                        num_or_size_splits=4, axis=1)
+      x_i = (tf.keras.backend.dot(inputs_i, k_i) +
+             tf.keras.backend.dot(inputs_i * si_i, p_i) * so_i)
+      x_f = (tf.keras.backend.dot(inputs_f, k_f) +
+             tf.keras.backend.dot(inputs_f * si_f, p_f) * so_f)
+      x_c = (tf.keras.backend.dot(inputs_c, k_c) +
+             tf.keras.backend.dot(inputs_c * si_c, p_c) * so_c)
+      x_o = (tf.keras.backend.dot(inputs_o, k_o) +
+             tf.keras.backend.dot(inputs_o * si_o, p_o) * so_o)
+      if self.use_bias:
+        b_i, b_f, b_c, b_o = tf.split(
+            self.bias, num_or_size_splits=4, axis=0)
+        x_i = tf.keras.backend.bias_add(x_i, b_i)
+        x_f = tf.keras.backend.bias_add(x_f, b_f)
+        x_c = tf.keras.backend.bias_add(x_c, b_c)
+        x_o = tf.keras.backend.bias_add(x_o, b_o)
+
+      if 0 < self.recurrent_dropout < 1.:
+        h_tm1_i = h_tm1 * rec_dp_mask[0]
+        h_tm1_f = h_tm1 * rec_dp_mask[1]
+        h_tm1_c = h_tm1 * rec_dp_mask[2]
+        h_tm1_o = h_tm1 * rec_dp_mask[3]
+      else:
+        h_tm1_i = h_tm1
+        h_tm1_f = h_tm1
+        h_tm1_c = h_tm1
+        h_tm1_o = h_tm1
+      x = (x_i, x_f, x_c, x_o)
+      h_tm1 = (h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o)
+      c, o = self._compute_carry_and_output(x, h_tm1, c_tm1)
+    else:
+      if 0. < self.dropout < 1.:
+        inputs = inputs * dp_mask[0]
+      kernel_mean = self.kernel.distribution.mean()
+      perturbation = self.kernel - kernel_mean
+      z = tf.keras.backend.dot(inputs, kernel_mean)
+      z += tf.keras.backend.dot(inputs * self.sign_input,
+                                perturbation) * self.sign_output
+      if 0. < self.recurrent_dropout < 1.:
+        h_tm1 = h_tm1 * rec_dp_mask[0]
+      recurrent_kernel_mean = self.recurrent_kernel.distribution.mean()
+      perturbation = self.recurrent_kernel - recurrent_kernel_mean
+      z += tf.keras.backend.dot(h_tm1, recurrent_kernel_mean)
+      z += tf.keras.backend.dot(h_tm1 * self.recurrent_sign_input,
+                                perturbation) * self.recurrent_sign_output
+      if self.use_bias:
+        z = tf.keras.backend.bias_add(z, self.bias)
+
+      z = tf.split(z, num_or_size_splits=4, axis=1)
+      c, o = self._compute_carry_and_output_fused(z, c_tm1)
+
+    h = o * self.activation(c)
+    return h, [h, c]
+
+
 class Zeros(object):
   """Function returning zeros tensor of same shape excluding the last dim."""
 
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index bab361192..473216318 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -33,6 +33,18 @@
 class BayesTest(parameterized.TestCase, tf.test.TestCase):
 
   @parameterized.parameters(
+      {"layer": bayes.Conv2DFlipout,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "zeros",
+       "all_close": True},
+      {"layer": bayes.Conv2DFlipout,
+       "kernel_initializer": "trainable_normal",
+       "bias_initializer": "zeros",
+       "all_close": False},
+      {"layer": bayes.Conv2DFlipout,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "trainable_normal",
+       "all_close": False},
       {"layer": bayes.Conv2DReparameterization,
        "kernel_initializer": "zeros",
        "bias_initializer": "zeros",
@@ -84,6 +96,7 @@ def testConv2DKernel(self,
     model.get_config()
 
   @parameterized.parameters(
+      {"layer": bayes.Conv2DFlipout},
       {"layer": bayes.Conv2DReparameterization},
       {"layer": bayes.Conv2DVariationalDropout},
   )
@@ -125,6 +138,18 @@ def testTrainableNormalStddevConstraint(self):
        "kernel_initializer": "zeros",
        "bias_initializer": "trainable_normal",
        "all_close": False},
+      {"layer": bayes.DenseFlipout,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "zeros",
+       "all_close": True},
+      {"layer": bayes.DenseFlipout,
+       "kernel_initializer": "trainable_normal",
+       "bias_initializer": "zeros",
+       "all_close": False},
+      {"layer": bayes.DenseFlipout,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "trainable_normal",
+       "all_close": False},
       {"layer": bayes.DenseReparameterization,
        "kernel_initializer": "zeros",
        "bias_initializer": "zeros",
@@ -177,6 +202,7 @@ def testDenseKernel(self,
 
   @parameterized.parameters(
       {"layer": bayes.DenseDVI},
+      {"layer": bayes.DenseFlipout},
       {"layer": bayes.DenseReparameterization},
       {"layer": bayes.DenseVariationalDropout},
   )
@@ -203,6 +229,7 @@ def take_mean(f, *args, **kwargs):
 
   @parameterized.parameters(
       {"layer": bayes.DenseDVI},
+      {"layer": bayes.DenseFlipout},
       {"layer": bayes.DenseReparameterization},
       {"layer": bayes.DenseVariationalDropout},
   )
@@ -255,6 +282,7 @@ def testDenseLoss(self, layer):
 
   @parameterized.parameters(
       {"layer": bayes.DenseDVI},
+      {"layer": bayes.DenseFlipout},
       {"layer": bayes.DenseReparameterization},
       {"layer": bayes.DenseVariationalDropout},
   )
@@ -277,6 +305,7 @@ def testDenseModel(self, layer):
 
   @parameterized.parameters(
       {"layer": bayes.DenseDVI},
+      {"layer": bayes.DenseFlipout},
       {"layer": bayes.DenseReparameterization},
       {"layer": bayes.DenseVariationalDropout},
   )
@@ -394,39 +423,57 @@ def testGaussianProcessPrior(self):
     self.assertLessEqual(log_prob_val, 0.)
     self.assertEqual(outputs_val.shape, (batch_size, output_dim))
 
-  @parameterized.named_parameters(
-      {"testcase_name": "_no_uncertainty",
+  @parameterized.parameters(
+      {"lstm_cell": bayes.LSTMCellFlipout,
        "kernel_initializer": "zeros",
        "recurrent_initializer": "orthogonal",
        "bias_initializer": "zeros",
        "all_close": True},
-      {"testcase_name": "_kernel_uncertainty",
+      {"lstm_cell": bayes.LSTMCellFlipout,
        "kernel_initializer": "trainable_normal",
        "recurrent_initializer": "orthogonal",
        "bias_initializer": "zeros",
        "all_close": False},
-      {"testcase_name": "_recurrent_uncertainty",
+      {"lstm_cell": bayes.LSTMCellFlipout,
+       "kernel_initializer": "zeros",
+       "recurrent_initializer": "orthogonal",
+       "bias_initializer": "trainable_normal",
+       "all_close": False},
+      {"lstm_cell": bayes.LSTMCellReparameterization,
+       "kernel_initializer": "zeros",
+       "recurrent_initializer": "orthogonal",
+       "bias_initializer": "zeros",
+       "all_close": True},
+      {"lstm_cell": bayes.LSTMCellReparameterization,
+       "kernel_initializer": "trainable_normal",
+       "recurrent_initializer": "orthogonal",
+       "bias_initializer": "zeros",
+       "all_close": False},
+      {"lstm_cell": bayes.LSTMCellReparameterization,
        "kernel_initializer": "zeros",
        "recurrent_initializer": "trainable_normal",
        "bias_initializer": "zeros",
        "all_close": False},
-      {"testcase_name": "_bias_uncertainty",
+      {"lstm_cell": bayes.LSTMCellReparameterization,
        "kernel_initializer": "zeros",
        "recurrent_initializer": "orthogonal",
        "bias_initializer": "trainable_normal",
        "all_close": False},
   )
   @test_utils.run_in_graph_and_eager_modes
-  def testLSTMCellReparameterization(
-      self, kernel_initializer, recurrent_initializer, bias_initializer,
-      all_close):
+  def testLSTMCell(self,
+                   lstm_cell,
+                   kernel_initializer,
+                   recurrent_initializer,
+                   bias_initializer,
+                   all_close):
     batch_size, timesteps, dim = 5, 3, 12
     hidden_size = 10
     inputs = tf.to_float(np.random.rand(batch_size, timesteps, dim))
-    cell = bayes.LSTMCellReparameterization(
-        hidden_size, kernel_initializer=kernel_initializer,
-        recurrent_initializer=recurrent_initializer,
-        bias_initializer=bias_initializer)
+    cell = lstm_cell(hidden_size,
+                     kernel_initializer=kernel_initializer,
+                     recurrent_initializer=recurrent_initializer,
+                     bias_initializer=bias_initializer)
     noise = tf.to_float(np.random.rand(1, hidden_size))
     h0, c0 = cell.get_initial_state(inputs)
     state = (h0 + noise, c0)
@@ -444,11 +491,15 @@ def testLSTMCellReparameterization(
       self.assertNotAllClose(res1, res3)
     cell.get_config()
 
+  @parameterized.parameters(
+      {"lstm_cell": bayes.LSTMCellFlipout},
+      {"lstm_cell": bayes.LSTMCellReparameterization},
+  )
   @test_utils.run_in_graph_and_eager_modes()
-  def testLSTMCellReparameterizationLoss(self):
+  def testLSTMCellLoss(self, lstm_cell):
     features = tf.to_float(np.random.rand(5, 1, 12))
     labels = tf.to_float(np.random.rand(5, 10))
-    cell = bayes.LSTMCellReparameterization(10)
+    cell = lstm_cell(10)
     state = (tf.zeros([1, 10]), tf.zeros([1, 10]))
 
     # Imagine this is the 1st epoch.
@@ -500,12 +551,16 @@ def testLSTMCellReparameterizationLoss(self):
     for grad in grads:
       self.assertIsNotNone(grad)
 
+  @parameterized.parameters(
+      {"lstm_cell": bayes.LSTMCellFlipout},
+      {"lstm_cell": bayes.LSTMCellReparameterization},
+  )
   @test_utils.run_in_graph_and_eager_modes()
-  def testLSTMCellReparameterizationModel(self):
+  def testLSTMCellModel(self, lstm_cell):
     batch_size, timesteps, dim = 5, 3, 12
     hidden_size = 10
     inputs = tf.to_float(np.random.rand(batch_size, timesteps, dim))
-    cell = bayes.LSTMCellReparameterization(hidden_size)
+    cell = lstm_cell(hidden_size)
     model = tf.keras.Sequential([
         tf.keras.layers.RNN(cell, return_sequences=True)
     ])
@@ -525,8 +580,9 @@ def testLSTMCellReparameterizationModel(self):
     # each call, so these should be different.
     self.assertNotAllClose(res1, res2)
     # NOTE: We didn't call `cell.call_weights` again before computing
-    # `outputs3`, so the cell should have had the same weights as it did during
-    # computation of `outputs2`, and thus yielded the same output tensor.
+    # `outputs3`, so the cell should have had the same weights as it did
+    # during computation of `outputs2`, and thus yielded the same output
+    # tensor.
     self.assertAllClose(res2, res3)
     self.assertLen(model.losses, 2)
 

From d9f8074fc391594656176dfbc77bfd3faa81892d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 24 Jun 2019 13:15:48 -0700
Subject: [PATCH 2152/2720] Change layer names according to
 s/MultiHeadedAttention/Attention/. Also, for Attention layers, make n_heads
 have default value 1.

PiperOrigin-RevId: 254820673
---
 tensor2tensor/trax/layers/attention.py             | 13 +++++--------
 .../models/research/position_lookup_transformer.py | 14 ++++++--------
 .../trax/models/research/transformer_revnet.py     |  3 +--
 tensor2tensor/trax/models/transformer.py           | 11 ++++-------
 4 files changed, 16 insertions(+), 25 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index f4d0d1eb1..793684f0b 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -111,8 +111,7 @@ def DotProductAttention(query, key, value, mask, dropout, mode, rng):
 
 
 @base.layer(n_inputs=4, n_outputs=2)
-def PureMultiHeadedAttention(x, params, n_heads=8, dropout=0.0, mode='train',
-                             **kwargs):
+def PureAttention(x, params, n_heads=1, dropout=0.0, mode='train', **kwargs):
   """Pure transformer-style multi-headed attention.
 
   Args:
@@ -149,7 +148,7 @@ def JoinHeads(x):  # pylint: disable=invalid-name
   return res, mask  # Keep the mask.
 
 
-def MultiHeadedAttentionQKV(d_feature, n_heads=8, dropout=0.0, mode='train'):
+def AttentionQKV(d_feature, n_heads=1, dropout=0.0, mode='train'):
   """Transformer-style multi-headed attention.
 
   Accepts inputs of the form q, k, v, mask.
@@ -169,14 +168,13 @@ def MultiHeadedAttentionQKV(d_feature, n_heads=8, dropout=0.0, mode='train'):
           core.Dense(d_feature),
           core.Dense(d_feature),
       ),
-      PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
+      PureAttention(  # pylint: disable=no-value-for-parameter
           d_feature=d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
       core.Dense(d_feature),
   ]
 
 
-def MultiHeadedAttention(
-    d_feature, n_heads=8, dropout=0.0, mode='train'):
+def Attention(d_feature, n_heads=1, dropout=0.0, mode='train'):
   """Transformer-style multi-headed attention.
 
   Accepts inputs of the form (x, mask) and constructs (q, k, v) from x.
@@ -192,8 +190,7 @@ def MultiHeadedAttention(
   """
   return [
       cb.Dup(), cb.Dup(),
-      MultiHeadedAttentionQKV(  # pylint: disable=no-value-for-parameter
-          d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+      AttentionQKV(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
   ]
 
 
diff --git a/tensor2tensor/trax/models/research/position_lookup_transformer.py b/tensor2tensor/trax/models/research/position_lookup_transformer.py
index 3f637dd92..596adedc8 100644
--- a/tensor2tensor/trax/models/research/position_lookup_transformer.py
+++ b/tensor2tensor/trax/models/research/position_lookup_transformer.py
@@ -246,8 +246,8 @@ def SumLearnedPick(positions):
   )
 
 
-def MultiHeadedAttentionPosition(
-    positions, d_feature, n_heads=8, dropout=0.0, mode='train'):
+def AttentionPosition(positions, d_feature, n_heads=8, dropout=0.0,
+                      mode='train'):
   """Transformer-style multi-headed attention."""
   return tl.Serial(
       tl.Dup(),
@@ -264,9 +264,8 @@ def MultiHeadedAttentionPosition(
           MixHeadsPos(h=n_heads),
           MixHeadsPos(h=n_heads),
       ),
-      tl.PureMultiHeadedAttention(
-          d_feature=d_feature, n_heads=n_heads,
-          dropout=dropout, mode=mode),
+      tl.PureAttention(d_feature=d_feature, n_heads=n_heads, dropout=dropout,
+                       mode=mode),
       tl.Parallel([], tl.Drop()),  # Drop the mask.
       CombineHeadsPos(h=n_heads),
       PreservePosition(tl.Dense(d_feature)),
@@ -314,9 +313,8 @@ def DecoderLayer(positions,
           tl.Dup(),
           tl.Parallel([],  # activation for (q, k, v)
                       tl.CausalMask(axis=-2)),  # attention mask
-          MultiHeadedAttentionPosition(positions,
-                                       d_feature, n_heads=n_heads,
-                                       dropout=dropout, mode=mode),
+          AttentionPosition(positions, d_feature, n_heads=n_heads,
+                            dropout=dropout, mode=mode),
           PreservePosition(tl.Dropout(rate=dropout, mode=mode))
       ),
       ResidualFeedForward(d_feature, d_feedforward, dropout, mode=mode)
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 4acc29a5e..5d069f8cd 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -291,8 +291,7 @@ def DecoderBlock(d_feature, d_feedforward, n_heads, n_attention_chunks,
       tl.LayerNorm(),
       tl.Dup(),
       tl.Parallel([], tl.CausalMask(axis=-2)),  # Create mask.
-      tl.MultiHeadedAttention(
-          d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+      tl.Attention(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Parallel([], tl.Drop()),  # Drop mask.
       tl.Dropout(rate=dropout, mode=mode),
   ]
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 8865ed6c0..b9f62ced9 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -53,8 +53,7 @@ def EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
   """
   attention = [
       tl.LayerNorm(),
-      tl.MultiHeadedAttention(
-          d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+      tl.Attention(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Dropout(rate=dropout, mode=mode),
   ]
   feed_forward = [
@@ -131,8 +130,7 @@ def DecoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
       tl.LayerNorm(),  # vec
       tl.Dup(),  # vec vec
       tl.Parallel([], tl.CausalMask(axis=-2)),  # vec mask
-      tl.MultiHeadedAttention(
-          d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+      tl.Attention(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Parallel([], tl.Drop()),  # vec
       tl.Dropout(rate=dropout, mode=mode),  # vec
   ]
@@ -209,8 +207,7 @@ def EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode):
       tl.LayerNorm(),                           #        vecs_d   ..... ......
       tl.Dup(),                                 # vecs_d vecs_d   ..... ......
       tl.Parallel([], tl.CausalMask(axis=-2)),  # ______ masks    ..... ......
-      tl.MultiHeadedAttention(d_feature, n_heads=n_heads,
-                              dropout=dropout, mode=mode),
+      tl.Attention(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Parallel([], tl.Drop()),               # ______   0      ..... ......
       tl.Dropout(rate=dropout, mode=mode),      # vecs_d          ..... ......
   ]
@@ -218,7 +215,7 @@ def EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode):
       tl.Parallel([], [], tl.Dup()),      # ______        _____  vecs_e vecs_e
       tl.Parallel([], tl.Swap()),         # ______        vecs_e masks  ......
       tl.Parallel([], tl.Dup()),          # ______ vecs_e vecs_e .....  ......
-      tl.MultiHeadedAttentionQKV(  # (q k v masks ... --> vecs_d masks ...)
+      tl.AttentionQKV(  # (q k v masks ... --> vecs_d masks ...)
           d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Dropout(rate=dropout, mode=mode),  # vecs_d mask vecs_e
   ]

From 20e1c458198e92a57d4aaaaebedb081da12515d4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 24 Jun 2019 17:14:29 -0700
Subject: [PATCH 2153/2720] Add bits and nats test to latent_layers.

PiperOrigin-RevId: 254869334
---
 tensor2tensor/layers/latent_layers_test.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
index 78a83d008..94ef58801 100644
--- a/tensor2tensor/layers/latent_layers_test.py
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -91,6 +91,22 @@ def imagetransformer_latent_tiny():
 
 class LatentLayersTest(tf.test.TestCase):
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testComputeBitsAndNats(self):
+    reconstruction_loss = tf.random_uniform(())
+    prior_loss = tf.random_uniform(())
+    data_dim = tf.random_uniform((), maxval=1000, dtype=tf.int32)
+    latent_dim = tf.random_uniform((), maxval=1000, dtype=tf.int32)
+    nats_per_dim, bits_per_dim = latent_layers.compute_nats_and_bits_per_dim(
+        data_dim,
+        latent_dim,
+        reconstruction_loss,
+        prior_loss)
+
+    nats_per_dim_py, bits_per_dim_conv_py = self.evaluate(
+        [nats_per_dim, bits_per_dim * tf.log(2.)])
+    self.assertAllClose(nats_per_dim_py, bits_per_dim_conv_py)
+
   @test_utils.run_in_graph_and_eager_modes()
   def testTransformerAutoencoder(self):
     hparams = imagetransformer_latent_tiny()

From eb5fe490b47cf10514f2ef4cb7fd5b5a5f17669e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 25 Jun 2019 11:49:06 -0700
Subject: [PATCH 2154/2720] Add SIGMOID_ACCURACY metric

PiperOrigin-RevId: 255015841
---
 tensor2tensor/utils/metrics.py      | 20 ++++++++++++++++++++
 tensor2tensor/utils/metrics_test.py | 16 ++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index e9af1a22f..b85b60d9c 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -57,6 +57,7 @@ class Metrics(object):
   SET_RECALL = "set_recall"
   SOFTMAX_CROSS_ENTROPY_ONE_HOT = "softmax_cross_entropy_one_hot"
   SIGMOID_ACCURACY_ONE_HOT = "sigmoid_accuracy_one_hot"
+  SIGMOID_ACCURACY = "sigmoid_accuracy"
   SIGMOID_RECALL_ONE_HOT = "sigmoid_recall_one_hot"
   SIGMOID_PRECISION_ONE_HOT = "sigmoid_precision_one_hot"
   SIGMOID_CROSS_ENTROPY_ONE_HOT = "sigmoid_cross_entropy_one_hot"
@@ -487,6 +488,24 @@ def sigmoid_accuracy_one_hot(logits, labels, weights_fn=None):
     return accuracy, tf.constant(1.0)
 
 
+def sigmoid_accuracy(logits, labels, weights_fn=None):
+  """Calculate accuracy for a set, given integer labels and logits.
+
+  Args:
+    logits: Tensor of size [batch-size, o=1, p=1, num-classes]
+    labels: Tensor of size [batch-size, o=1, p=1]
+    weights_fn: Function that takes in labels and weighs examples (unused)
+  Returns:
+    accuracy (scalar), weights
+  """
+  with tf.variable_scope("sigmoid_accuracy", values=[logits, labels]):
+    del weights_fn
+    predictions = tf.nn.sigmoid(logits)
+    predictions = tf.argmax(predictions, -1)
+    _, accuracy = tf.metrics.accuracy(labels=labels, predictions=predictions)
+    return accuracy, tf.constant(1.0)
+
+
 def sigmoid_precision_one_hot(logits, labels, weights_fn=None):
   """Calculate precision for a set, given one-hot labels and logits.
 
@@ -835,6 +854,7 @@ def pearson_correlation_coefficient(predictions, labels, weights_fn=None):
     Metrics.ROUGE_L_F: rouge.rouge_l_fscore,
     Metrics.EDIT_DISTANCE: sequence_edit_distance,
     Metrics.SOFTMAX_CROSS_ENTROPY_ONE_HOT: softmax_cross_entropy_one_hot,
+    Metrics.SIGMOID_ACCURACY: sigmoid_accuracy,
     Metrics.SIGMOID_ACCURACY_ONE_HOT: sigmoid_accuracy_one_hot,
     Metrics.SIGMOID_RECALL_ONE_HOT: sigmoid_recall_one_hot,
     Metrics.SIGMOID_PRECISION_ONE_HOT: sigmoid_precision_one_hot,
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index e259f5124..3172ad19d 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -225,6 +225,22 @@ def testSigmoidAccuracyOneHot(self):
       s = session.run(score)
     self.assertEqual(s, 0.5)
 
+  def testSigmoidAccuracy(self):
+    logits = np.array([
+        [-1., 1.],
+        [1., -1.],
+        [-1., 1.],
+        [1., -1.]
+    ])
+    labels = np.array([1, 0, 0, 1])
+
+    with self.test_session() as session:
+      score, _ = metrics.sigmoid_accuracy(logits, labels)
+      session.run(tf.global_variables_initializer())
+      session.run(tf.local_variables_initializer())
+      s = session.run(score)
+    self.assertEqual(s, 0.5)
+
   def testSigmoidPrecisionOneHot(self):
     logits = np.array([
         [-1., 1.],

From beb26959f1292ccf8360351c2299325754321419 Mon Sep 17 00:00:00 2001
From: Benjamin Kramer <kramerb@google.com>
Date: Tue, 25 Jun 2019 22:18:04 -0700
Subject: [PATCH 2155/2720] Use tf.linalg.inv on TPU

PiperOrigin-RevId: 255117361
---
 tensor2tensor/models/research/glow_ops.py | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 888055040..84c44f6a3 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -389,12 +389,7 @@ def invertible_1x1_conv(name, x, reverse=False):
       w = tf.reshape(w, [1, 1] + w_shape)
       x = tf.nn.conv2d(x, w, [1, 1, 1, 1], "SAME", data_format="NHWC")
     else:
-      # TODO(b/111271662): Remove when supported.
-      def tpu_inv(m):
-        """tf.linalg.inv workaround until it is supported on TPU."""
-        q, r = tf.linalg.qr(m)
-        return tf.linalg.triangular_solve(r, tf.transpose(q), lower=False)
-      w_inv = tf.reshape(tpu_inv(w), [1, 1]+w_shape)
+      w_inv = tf.reshape(tf.linalg.inv(w), [1, 1]+w_shape)
       x = tf.nn.conv2d(
           x, w_inv, [1, 1, 1, 1], "SAME", data_format="NHWC")
       objective *= -1

From 61469eea3330a37551c1e5824a950a7395b8461a Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 26 Jun 2019 09:54:44 -0700
Subject: [PATCH 2156/2720] Fix Reversible Transformer

PiperOrigin-RevId: 255208944
---
 tensor2tensor/trax/inputs.py                  |  4 +-
 .../models/research/transformer_revnet.py     | 97 +++++++++++++------
 2 files changed, 72 insertions(+), 29 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index cc0d2b236..93793bcba 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -169,8 +169,8 @@ def dataset_to_stream(dataset, input_name, n_chunks=0, append_targets=False):
     if len(out.shape) > 1 and out.shape[-1] == 1:
       out = np.squeeze(out, axis=-1)
     if n_chunks > 0:
-      inp = np.split(inp, n_chunks, axis=1)
-      out = np.split(out, n_chunks, axis=1)
+      inp = tuple(np.split(inp, n_chunks, axis=1))
+      out = tuple(np.split(out, n_chunks, axis=1))
     if append_targets:
       inp = (inp, out)
     yield inp, out
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 5d069f8cd..d2e646801 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -39,7 +39,7 @@ class Map(tl.Layer):
     A new layer representing mapping layer to all elements of the input.
   """
 
-  def __init__(self, layer, check_shapes=True):
+  def __init__(self, layer, sections=1, check_shapes=True):
     super(Map, self).__init__()
     if layer is None or isinstance(layer, (list, tuple)):
       layer = tl.Serial(layer)
@@ -50,13 +50,20 @@ def __init__(self, layer, check_shapes=True):
     # are valid cases -- e.g., when self._layer has no parameters -- where we
     # can apply Map to different shapes -- set check_shapes=False in such cases.
     self._check_shapes = check_shapes
+    self._sections = sections
+
+  def n_inputs(self):
+    """Specifies how many data tensors this layer expects as input."""
+    return self._sections
+
+  def n_outputs(self):
+    """Specifies how many data tensors this layer promises as output."""
+    return self._sections
 
   def call(self, inputs, params=(), **kwargs):
     rngs = _pop_rng_and_split(kwargs, len(inputs))
     result = [self._layer(x, params=params, rng=r, **kwargs)
               for x, r in zip(inputs, rngs)]
-    if isinstance(inputs, list):
-      return result
     return tuple(result)
 
   def new_parameters(self, input_shape, input_dtype, rng):
@@ -90,7 +97,7 @@ def inverse_and_vjp(self, output, ct, params=(), **kwargs):
     """Backward pass: computes the inverse of a layer and propagates gradients.
 
     Args:
-      output: Output activations; can be a (possibly nested) tuple or list.
+      output: Output activations; can be a (possibly nested) tuple.
       ct: gradient signal (cotangent) computed based on subsequent layers. If
           None, no gradients are propagated. Otherwise the structure and shape
           must match the output.
@@ -146,16 +153,48 @@ def vjpfun(ct):
     return do_call(x, params, kwargs)
 
 
+class Split(tl.Layer):
+  """Splits the input into sections along an axis."""
+
+  def __init__(self, sections=2, axis=-1):
+    super(Split, self).__init__()
+    self._sections = sections
+    self._axis = axis
+
+  def call(self, inputs, params=(), **kwargs):
+    del params, kwargs
+    return tuple(backend.numpy.split(inputs, self._sections, self._axis))
+
+  def new_parameters(self, input_shapes, input_dtype, rng):
+    return ()
+
+  def n_inputs(self):
+    """Specifies how many data tensors this layer expects as input."""
+    return 1
+
+  def n_outputs(self):
+    """Specifies how many data tensors this layer promises as output."""
+    return self._sections
+
+
 @tl.layer()
-def Split(x, params, sections=2, axis=-1, **kwargs):
+def Chunk(x, params, sections=2, **kwargs):
   del params, kwargs
-  return list(backend.numpy.split(x, sections, axis))
+  assert x.shape[1] % sections == 0
+  return backend.numpy.reshape(x, (
+      x.shape[0] * sections,
+      x.shape[1] // sections,
+      ) + x.shape[2:])
 
 
 @tl.layer()
-def Duplicate(x, params, sections=2, **kwargs):
+def Unchunk(x, params, sections=2, **kwargs):
   del params, kwargs
-  return [x for _ in range(sections)]
+  assert x.shape[0] % sections == 0
+  return backend.numpy.reshape(x, (
+      x.shape[0] // sections,
+      x.shape[1] * sections,
+      ) + x.shape[2:])
 
 
 class ReversibleHalfResidual(ReversibleLayerMixin, tl.Serial):
@@ -163,27 +202,31 @@ class ReversibleHalfResidual(ReversibleLayerMixin, tl.Serial):
 
   def __init__(self, residual_layers):
     self.compute_residual = tl.Serial([
-        # TODO(jonni): Rewrite without using Select.
-        tl.Select(inputs=('x1_or_y1', 'x2'), output=('x2', 'x1_or_y1', 'x2')),
+        # (x1_or_y1, x2) -> (x2, x1_or_y1, x2)
+        tl.Parallel([], tl.Dup()),
+        tl.Swap(),
         tl.Parallel(residual_layers, [], []),
     ])
 
-    layers = [self.compute_residual, tl.Add()]
+    layers = [
+        self.compute_residual,
+        tl.Parallel(tl.Add(), [])
+    ]
     super(ReversibleHalfResidual, self).__init__(layers)
 
-    self.subtract_top = tl.SubtractTop()
+    self.subtract_top = tl.Parallel(tl.SubtractTop(), [])
     self.reverse_layers = [self.compute_residual, self.subtract_top]
 
   def inverse_and_vjp(self, output, ct, params=(), **kwargs):
     rng = kwargs.pop('rng', None)
-    rngs = (None,) * self._nlayers
+    rngs = (None,) * self._n_layers
     if rng is not None:
-      rngs = backend.random.split(rng, self._nlayers)
+      rngs = backend.random.split(rng, self._n_layers)
 
     if ct is None:
       reconstructed_x = output
-      # Note that self._layers aligns exactly with self.reverse_layers in terms
-      # of parameter and rng usage, so no re-ordering is required.
+      # Note that self.sublayers() aligns exactly with self.reverse_layers in
+      # terms of parameter and rng usage, so no re-ordering is required.
       for layer, p, rng in zip(self.reverse_layers, params, rngs):
         reconstructed_x = layer(reconstructed_x, p, rng=rng, **kwargs)
       return reconstructed_x, None
@@ -234,8 +277,8 @@ class ReversibleSerial(ReversibleLayerMixin, tl.Serial):
   def __init__(self, *layers):
     super(ReversibleSerial, self).__init__(*layers)
 
-    # Note that self._layers has already been flattened to remove nested lists.
-    for i, layer in enumerate(self._layers):
+    # Note that sublayers has already been flattened to remove nested lists.
+    for i, layer in enumerate(self.sublayers()):
       if not isinstance(layer, ReversibleLayerMixin):
         raise ValueError(
             'Sub-layer {} of ReversibleSerial is not reversible: {}'.format(
@@ -243,15 +286,15 @@ def __init__(self, *layers):
 
   def inverse_and_vjp(self, output, ct, params=(), **kwargs):
     rng = kwargs.pop('rng', None)
-    rngs = (None,) * self._nlayers
+    rngs = (None,) * self._n_layers
     if rng is not None:
-      rngs = backend.random.split(rng, self._nlayers)
+      rngs = backend.random.split(rng, self._n_layers)
 
     layer_val = output
     if ct is not None:
       layer_ct = ct
       params_ct = []
-    for layer, p, rng in reversed(zip(self._layers, params, rngs)):
+    for layer, p, rng in reversed(zip(self.sublayers(), params, rngs)):
       layer_val, layer_ct = layer.inverse_and_vjp(
           layer_val, layer_ct, p, rng=rng, **kwargs)
       if ct is not None:
@@ -298,9 +341,9 @@ def DecoderBlock(d_feature, d_feedforward, n_heads, n_attention_chunks,
 
   # TODO(kitaev): Memory-efficient attention. This chunking is temporary.
   self_attention = [
-      Split(sections=n_attention_chunks, axis=-2),  # pylint: disable=no-value-for-parameter
-      Map(self_attention),
-      tl.Concatenate(axis=-2),
+      Chunk(sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
+      self_attention,
+      Unchunk(sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
   ]
 
   feed_forward = [
@@ -345,10 +388,10 @@ def TransformerRevnetLM(vocab_size,
       tl.PositionalEncoding(max_len=max_len),
   ]
   return tl.Model(
-      tl.Concatenate(),
+      tl.Concatenate(n_items=n_chunks),
       tl.ShiftRight(),
       positional_embedder,
-      Duplicate(),  # pylint: disable=no-value-for-parameter
+      tl.Dup(),
       ReversibleSerial([
           DecoderBlock(d_feature, d_feedforward, n_heads, n_attention_chunks,
                        dropout, mode)
@@ -360,6 +403,6 @@ def TransformerRevnetLM(vocab_size,
       Map([
           tl.Dense(vocab_size),
           tl.LogSoftmax(),
-      ]),
+      ], sections=n_chunks),
   )
 

From 54c1d3dff4641d62529ab34523c337c4467c1866 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 26 Jun 2019 10:09:07 -0700
Subject: [PATCH 2157/2720] Fix wasteful memory allocations in dropout shape
 calculation

PiperOrigin-RevId: 255212058
---
 tensor2tensor/trax/layers/base.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 7c1eb9a46..80b583f08 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -156,13 +156,15 @@ def pseudo_call(self, pseudo_inputs, params):
     """
     try:
       with backend.use_backend('jax'):
-        # Same as backend.random.get_prng(0), but no op-by-op execution.
-        rng = onp.zeros(2, onp.uint32)
-        def call_on_input(x, params):
+        # Beware: using an actual RNG (as opposed to this ShapeType stub) would
+        # cause a large number of dropout masks to be computed and permanently
+        # stored in global memory.
+        rng = ShapeType(shape=(2,), dtype=onp.uint32)
+        def call_on_input(x, params, rng):
           return self.call(x, params=params, rng=rng)
         params_shapes = nested_map(
             params, lambda x: ShapeType(shape=x.shape, dtype=x.dtype))
-        s = _eval_on_shapes(call_on_input, pseudo_inputs, params_shapes)
+        s = _eval_on_shapes(call_on_input, pseudo_inputs, params_shapes, rng)
       return s
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback(skip=3)

From da6fd084b6263aeb6871dbbf2af6dec8b7912557 Mon Sep 17 00:00:00 2001
From: Etienne Pot <epot@google.com>
Date: Wed, 26 Jun 2019 10:10:59 -0700
Subject: [PATCH 2158/2720] Fix mode kwargs for AtariCNN and NeuralGPU

PiperOrigin-RevId: 255212377
---
 tensor2tensor/trax/models/atari_cnn.py  | 4 +++-
 tensor2tensor/trax/models/neural_gpu.py | 5 ++++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/models/atari_cnn.py b/tensor2tensor/trax/models/atari_cnn.py
index de9029504..0eb4494d9 100644
--- a/tensor2tensor/trax/models/atari_cnn.py
+++ b/tensor2tensor/trax/models/atari_cnn.py
@@ -22,8 +22,10 @@
 from tensor2tensor.trax import layers as tl
 
 
-def AtariCnn(hidden_sizes=(32, 32), output_size=128):
+def AtariCnn(hidden_sizes=(32, 32), output_size=128, mode='train'):
   """An Atari CNN."""
+  del mode
+
   # TODO(jonni): Include link to paper?
   # Input shape: (B, T, H, W, C)
   # Output shape: (B, T, output_size)
diff --git a/tensor2tensor/trax/models/neural_gpu.py b/tensor2tensor/trax/models/neural_gpu.py
index dc19d52cb..b89c1f9b6 100644
--- a/tensor2tensor/trax/models/neural_gpu.py
+++ b/tensor2tensor/trax/models/neural_gpu.py
@@ -59,17 +59,20 @@ def BuildConv():
       candidate_nonlinearity=tl.HardTanh)
 
 
-def NeuralGPU(d_feature=96, steps=16, vocab_size=2):
+def NeuralGPU(d_feature=96, steps=16, vocab_size=2, mode='train'):
   """Implementation of Neural GPU: https://arxiv.org/abs/1702.08727.
 
   Args:
     d_feature: Number of memory channels (dimensionality of feature embedding).
     steps: Number of times depthwise recurrence steps.
     vocab_size: Vocabulary size.
+    mode: Whether we are training or evaluating or doing inference.
 
   Returns:
     A NeuralGPU Stax model.
   """
+  del mode
+
   core = ConvDiagonalGRU(units=d_feature)
   return tl.Model(
       tl.Embedding(d_feature=d_feature, vocab_size=vocab_size),

From e1cf771cef3072967908484f1e3f66574c6fb6f5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 26 Jun 2019 10:21:00 -0700
Subject: [PATCH 2159/2720] Remove all dropout in reversible transformer

PiperOrigin-RevId: 255214277
---
 .../trax/models/research/transformer_revnet.py      | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index d2e646801..ae7dd0da8 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -78,15 +78,15 @@ def new_parameters(self, input_shape, input_dtype, rng):
 
 def FeedForward(d_feature, d_feedforward, dropout, mode):
   """Feed-forward block with layer normalization at start."""
-  # TODO(kitaev): dropout is disabled to save memory
+  # TODO(kitaev): add dropout. Dropout is typically performed by adding noise to
+  # the activations, but when the size of the activations is very large it is
+  # more efficient to add noise to the *parameters* instead.
   del dropout, mode
   return [
       tl.LayerNorm(),
       tl.Dense(d_feedforward),
       tl.Relu(),
-      # tl.Dropout(rate=dropout, mode=mode),
       tl.Dense(d_feature),
-      # tl.Dropout(rate=dropout, mode=mode),
   ]
 
 
@@ -334,9 +334,9 @@ def DecoderBlock(d_feature, d_feedforward, n_heads, n_attention_chunks,
       tl.LayerNorm(),
       tl.Dup(),
       tl.Parallel([], tl.CausalMask(axis=-2)),  # Create mask.
-      tl.Attention(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+      # TODO(kitaev): add dropout
+      tl.Attention(d_feature, n_heads=n_heads, dropout=None, mode=mode),
       tl.Parallel([], tl.Drop()),  # Drop mask.
-      tl.Dropout(rate=dropout, mode=mode),
   ]
 
   # TODO(kitaev): Memory-efficient attention. This chunking is temporary.
@@ -383,8 +383,7 @@ def TransformerRevnetLM(vocab_size,
   """
   positional_embedder = [
       tl.Embedding(d_feature, vocab_size),
-      # TODO(kitaev): dropout is disabled to save memory
-      # tl.Dropout(rate=dropout, mode=mode),
+      # TODO(kitaev): add dropout
       tl.PositionalEncoding(max_len=max_len),
   ]
   return tl.Model(

From eb048f69c7ea860324122b87cb9caf59c52a27f3 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 26 Jun 2019 15:41:28 -0700
Subject: [PATCH 2160/2720] Internal

PiperOrigin-RevId: 255279570
---
 tensor2tensor/layers/common_layers.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 5e95fed62..d25f5e7e9 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1938,7 +1938,7 @@ def padded_cross_entropy_mixture(logits,
   Raises:
     ValueError: in case of unsupported argument types.
   """
-
+  # TODO(karishmamalkan): Fix documentation and refactor name
   (logits, mixture_labels, supervised_mode) = logits
 
   logit_shapes = shape_list(
@@ -1982,6 +1982,15 @@ def padded_cross_entropy_mixture(logits,
       mixture_accuracy = tf.metrics.accuracy(
           mixture_labels, best_mixtures, name="mixture_accuracy")
       tf.summary.scalar("mixture_acc_plot", mixture_accuracy[1])
+
+    # plot a summary for the difference between the top 2 losses
+    if num_mixtures > 1:
+      xent_reshaped = tf.transpose(tf.squeeze(xent), perm=[1, 0])
+      top_2_mixtures = tf.reduce_mean(
+          -tf.math.top_k(-xent_reshaped, k=2)[0], axis=0)
+      tf.summary.scalar("difference_top_2",
+                        top_2_mixtures[0] - top_2_mixtures[1])
+
     if supervised_mode:
       xent_min = gather_tensor_by_mixture_index(
           xent, mixture_labels, batch_size, num_mixtures, reshape=False)
@@ -2013,7 +2022,7 @@ def padded_cross_entropy_mixture(logits,
                       summed_xent_max / summed_weights)
 
   if return_best_logits:
-    return summed_xent_min, summed_weights, best_logits
+    return summed_xent_min, summed_weights, best_logits, return_mixture_indices
   else:
     return summed_xent_min, summed_weights
 

From 7860d2670868a1747c462e23847ad5ec4395cfe4 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 26 Jun 2019 23:03:47 -0700
Subject: [PATCH 2161/2720] Try restoring the model in reverse chronological
 order in PPO, we may hit an error if we were unable to complete writing the
 model file for any reason, so fall back on the previous model file if one
 exists and pick up from there.

PiperOrigin-RevId: 255335440
---
 tensor2tensor/trax/rlax/ppo.py | 23 ++++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 2f4008392..832660774 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -763,15 +763,20 @@ def maybe_restore_params(output_dir, policy_and_value_net_params):
     which we restored the params, 0 is restore = False.
   """
   model_files = gfile.glob(os.path.join(output_dir, "model-??????.pkl"))
-  if not model_files:
-    return False, policy_and_value_net_params, 0
-
-  model_file = sorted(model_files)[-1]
-  model_file_basename = os.path.basename(model_file)  # model-??????.pkl
-  i = int(filter(str.isdigit, model_file_basename))
-  with gfile.GFile(model_file, "rb") as f:
-    policy_and_value_net_params = pickle.load(f)
-  return True, policy_and_value_net_params, i
+  for model_file in reversed(sorted(model_files)):
+    logging.info("Trying to restore model from %s", model_file)
+    try:
+      with gfile.GFile(model_file, "rb") as f:
+        loaded_policy_and_value_net_params = pickle.load(f)
+        policy_and_value_net_params = loaded_policy_and_value_net_params
+      model_file_basename = os.path.basename(model_file)  # model-??????.pkl
+      i = int(filter(str.isdigit, model_file_basename))
+      return True, policy_and_value_net_params, i
+    except EOFError as e:
+      logging.error("Unable to load model from: %s with %s", model_file, e)
+      # Try an older version.
+      continue
+  return False, policy_and_value_net_params, 0
 
 
 def training_loop(

From a26fe8adc9a333c08149b32ebe048f7ac21b18fe Mon Sep 17 00:00:00 2001
From: David So <davidso@google.com>
Date: Thu, 27 Jun 2019 01:12:45 -0700
Subject: [PATCH 2162/2720] Cosine learning rate decay with multiple cycles.

PiperOrigin-RevId: 255350125
---
 tensor2tensor/utils/learning_rate.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index 382742b0e..26b397f8b 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -52,6 +52,14 @@ def learning_rate_factor(name, step_num, hparams):
     return tf.math.cos(
         step * np.pi /
         (hparams.train_steps - hparams.learning_rate_warmup_steps)) / 2.0 + 0.5
+  elif name == "multi_cycle_cos_decay":
+    # Cosine decay with a variable number of cycles. This is different from
+    # "cosdecay" because it starts at 1 when the warmup steps end. Use
+    # hparams.learning_rate_decay_steps to determine the number of cycles.
+    x = tf.maximum(step_num, hparams.learning_rate_warmup_steps)
+    step = x - hparams.learning_rate_warmup_steps
+    return tf.math.cos(
+        step * np.pi / hparams.learning_rate_decay_steps) / 2.0 + 0.5
   elif name == "rsqrt_decay":
     return tf.rsqrt(tf.maximum(step_num, hparams.learning_rate_warmup_steps))
   elif name == "rsqrt_normalized_decay":

From cb5b80c47c6e16f331403c2f4700bb024d3944ac Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 27 Jun 2019 15:22:53 -0700
Subject: [PATCH 2163/2720] Memory-efficient attention for reversible
 transformer

PiperOrigin-RevId: 255489639
---
 .../transformer_revnet_imagenet64_8gb.gin     |   7 +-
 .../models/research/transformer_revnet.py     | 390 ++++++++++++++++--
 2 files changed, 367 insertions(+), 30 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
index 9e62d3ff4..cfafd4b42 100644
--- a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
@@ -5,7 +5,7 @@ import tensor2tensor.trax.trax
 
 # Parameters for batch_fun:
 # ==============================================================================
-batch_fun.batch_size_per_device = 32
+batch_fun.batch_size_per_device = 16
 batch_fun.eval_batch_size = 8
 batch_fun.max_eval_length = 12288  # 64 * 64 * 3
 
@@ -24,7 +24,7 @@ MultifactorSchedule.warmup_steps = 8000
 
 # Parameters for train:
 # ==============================================================================
-train.eval_frequency = 100
+train.eval_frequency = 25
 train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerRevnetLM
@@ -43,4 +43,5 @@ TransformerRevnetLM.n_heads = 8
 TransformerRevnetLM.n_layers = 6
 TransformerRevnetLM.vocab_size = 256
 TransformerRevnetLM.n_chunks = 64
-TransformerRevnetLM.n_attention_chunks = 64
+TransformerRevnetLM.n_attention_chunks = 1
+TransformerRevnetLM.attention_loop_stride = 512
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index ae7dd0da8..cb8f03e48 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -19,9 +19,11 @@
 from __future__ import print_function
 
 import jax
+import numpy as onp
 
 from tensor2tensor.trax import backend
 from tensor2tensor.trax import layers as tl
+from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers.combinators import _pop_rng_and_split
 
 
@@ -129,8 +131,6 @@ def __call__(self, x, params=(), **kwargs):
     assert backend.get_name() == 'jax', (
         'Reversible layers are only supported in JAX')
 
-    # Retrieve shared parameters (cf. tl.Layer.__call__)
-    super(ReversibleLayerMixin, self).__call__(x, params, **kwargs)
     if params is () and self._params:  # pylint: disable=literal-comparison
       # TODO(kitaev): Figure out why parameter sharing doesn't work (if this
       # explicit error isn't thrown, a jax tracer error occurs instead)
@@ -249,6 +249,332 @@ def call_compute_residual(x, params, kwargs):
       return reconstructed_x, (x_ct, (residual_params_ct, ()), kwargs_ct)
 
 
+@tl.layer(n_inputs=1, n_outputs=1)
+def SplitHeads(x, params, n_heads=1, **kwargs):
+  del params, kwargs
+  d_feature = x.shape[-1]
+  assert d_feature % n_heads == 0
+  d_head = d_feature // n_heads
+  n_batch = np.shape(x)[0]
+  # n_batch, seqlen, d_feature --> n_batch, n_heads, seqlen, d_head
+  return np.transpose(
+      np.reshape(x, (n_batch, -1, n_heads, d_head)), (0, 2, 1, 3))
+
+
+@tl.layer(n_inputs=1, n_outputs=1)
+def JoinHeads(x, params, **kwargs):
+  del params, kwargs
+  n_batch = np.shape(x)[0]
+  seqlen = np.shape(x)[2]
+  # n_batch, n_heads, seqlen, d_head --> n_batch, seqlen, d_feature
+  return np.reshape(np.transpose(x, (0, 2, 1, 3)), (n_batch, seqlen, -1))
+
+
+class ApplyAttentionWrapper(tl.Parallel):
+  """Same as tl.Parallel(attention, [], []), but implements forward_and_vjp.
+
+  See MemoryEfficientDotProductAttention for why this is needed.
+  """
+
+  def __init__(self, attention):
+    assert hasattr(attention, 'forward_and_vjp')
+    super(ApplyAttentionWrapper, self).__init__(attention, [], [])
+    self.attention = attention
+
+  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
+    # Simultaneous forward pass and backprop through the attention mechanism.
+    qkv = inputs[:3]
+    passthrough = inputs[3:]
+    out_ct = ct[0]
+    passthrough_ct = ct[1:]
+
+    out, qkv_ct = self.attention.forward_and_vjp(
+        qkv, out_ct, params=(), **kwargs)
+    return (out,) + passthrough, qkv_ct + passthrough_ct
+
+
+class DotProductAttention(tl.Layer):
+  """A standard (non-memory-efficient) dot product attention implementation.
+
+  This class sets up the API that is required to implement
+  MemoryEfficientDotProductAttention.
+  """
+
+  def __init__(self, dropout, mode):
+    super(DotProductAttention, self).__init__()
+    self._dropout = dropout
+    self._mode = mode
+
+  def call(self, inputs, params=(), rng=None, **kwargs):
+    del params
+    q, k, v = inputs
+    mask_size = q.shape[-2]
+    mask = np.tril(np.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
+    res = tl.DotProductAttention(
+        q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=rng)
+    return res
+
+  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
+    # Simultaneous forward pass and backprop through the attention mechanism.
+    def do_call(x):
+      return self.call(x, params, **kwargs)
+    output, vjpfun = jax.vjp(do_call, inputs)
+    return output, vjpfun(ct)[0]
+
+  def new_parameters(self, input_shapes, input_dtype, rng):
+    return ()
+
+  def n_inputs(self):
+    return 3
+
+  def n_outputs(self):
+    return 1
+
+
+class MemoryEfficientDotProductAttention(DotProductAttention):
+  """Memory-efficient dot product attention."""
+
+  def __init__(self, loop_stride, dropout, mode):
+    super(MemoryEfficientDotProductAttention, self).__init__(dropout, mode)
+    self._loop_stride = loop_stride
+
+  def call(self, inputs, params=(), **kwargs):
+    output, _ = self.forward_and_vjp(inputs, None, params=params, **kwargs)
+    return output
+
+  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
+    # This is the core of the memory-efficient attention implementation, where
+    # we use the jax.lax.while_loop primitive to compute attention for a small
+    # set of query positions at a time. Note how in the backwards pass, we
+    # compute both the forward direction (to recover the previous layer's
+    # activations) and the backward direction simultaneously. This allows us to
+    # only use a single loop, where the inner portion of the loop does a slice
+    # of the forward+backward joint computation. Unfortunately we have had to
+    # introduce a large number of wrapper classes (including
+    # ReversibleAttentionHalfResidual and ApplyAttentionWrapper) for the sole
+    # purpose of connecting this implementation of forward_and_vjp with the core
+    # backprop implementation.
+
+    query, key, value = inputs
+    depth = np.shape(query)[-1]
+    do_backprop = ct is not None
+
+    def make_mask(N, M, k):
+      x = np.arange(N, dtype=np.int32)
+      y = np.arange(M, dtype=np.int32)
+      mask = jax.lax.lt(
+          (jax.lax.broadcast_in_dim(
+              x, shape=(N, M), broadcast_dimensions=(0,)) + k),
+          jax.lax.broadcast(y, [N]))
+      mask = jax.lax.convert_element_type(mask, np.float32)
+      return mask
+
+    def forward_slice(query_slice, q_loop_idx, key, value):
+      """Forward pass for a subset of the query vectors."""
+      dots = np.matmul(
+          query_slice, np.swapaxes(key, -1, -2)) / np.sqrt(depth)
+
+      # Causal masking
+      mask = make_mask(dots.shape[-2], dots.shape[-1], q_loop_idx)
+      dots = dots - 1e9 * mask
+
+      # Softmax.
+      dots = np.exp(dots - dots.max(axis=-1, keepdims=True))
+      dots = dots / dots.sum(axis=-1, keepdims=True)
+      out_slice = np.matmul(dots, value)
+      return out_slice
+
+    def forward_and_vjp_slice(query_slice, q_loop_idx, key, value, ct_slice):
+      output_slice, vjpfun = jax.vjp(
+          forward_slice, query_slice, q_loop_idx, key, value)
+      return output_slice, vjpfun(ct_slice)
+
+    q_loop_idx = np.zeros((), dtype=np.int32)
+    q_loop_max = query.shape[2]
+    q_loop_stride = self._loop_stride
+    assert q_loop_max % q_loop_stride == 0, (
+        'Stride must evenly divide the number of query elements.')
+
+    out_accum = np.zeros_like(query)
+    if do_backprop:
+      query_ct_accum = np.zeros_like(query)
+      key_ct_accum = np.zeros_like(key)
+      value_ct_accum = np.zeros_like(value)
+      init_vals = (
+          q_loop_idx, out_accum,
+          query_ct_accum, key_ct_accum, value_ct_accum)
+    else:
+      init_vals = (q_loop_idx, out_accum)
+
+    def cond_fun(vals):
+      q_loop_idx = vals[0]
+      return jax.lax.lt(q_loop_idx, q_loop_max)
+
+    def body_fun(vals):
+      """Compute a slice of the attention mechanism."""
+      if do_backprop:
+        (q_loop_idx, out_accum,
+         query_ct_accum, key_ct_accum, value_ct_accum) = vals
+      else:
+        q_loop_idx, out_accum = vals
+
+      query_slice = jax.lax.dynamic_slice_in_dim(
+          query, q_loop_idx, q_loop_stride, axis=2)
+
+      if do_backprop:
+        ct_slice = jax.lax.dynamic_slice_in_dim(
+            ct, q_loop_idx, q_loop_stride, axis=2)
+        out_slice, partial_ct = forward_and_vjp_slice(
+            query_slice, q_loop_idx, key, value, ct_slice)
+        query_ct_accum = jax.lax.dynamic_update_slice_in_dim(
+            query_ct_accum, partial_ct[0], q_loop_idx, axis=2)
+        # ignore partial_ct[1], which is wrt the loop idx
+        key_ct_accum = key_ct_accum + partial_ct[2]
+        value_ct_accum = value_ct_accum + partial_ct[3]
+      else:
+        out_slice = forward_slice(query_slice, q_loop_idx, key, value)
+
+      out_accum = jax.lax.dynamic_update_slice_in_dim(
+          out_accum, out_slice, q_loop_idx, axis=2)
+      q_loop_idx = q_loop_idx + q_loop_stride
+
+      if do_backprop:
+        return (q_loop_idx, out_accum,
+                query_ct_accum, key_ct_accum, value_ct_accum)
+      else:
+        return (q_loop_idx, out_accum)
+
+    final_vals = jax.lax.while_loop(cond_fun, body_fun, init_vals)
+
+    if not do_backprop:
+      return final_vals[1], None
+    else:
+      return final_vals[1], final_vals[2:]
+
+
+class ReversibleAttentionHalfResidual(ReversibleLayerMixin, tl.Serial):
+  """Half of a RevNet-style residual that performs attention.
+
+  If inputs are (x1, x2), then outputs are (x1 + z, x2) where:
+  z = post_attention(attention(pre_attention(x1)))
+
+  The post_attention layers must be linear in their input (typically they will
+  consists of reshaping and dense linear layers). This allows back-propagating
+  the gradient signal from the output of ReversibleAttentionHalfResidual to the
+  output of the "attention" portion based only on the network parameters.
+
+  The forward pass is equivalent to using
+  ReversibleHalfResidual([pre_attention, attention, post_attention]), but the
+  backward pass uses attention.forward_and_vjp. See
+  MemoryEfficientDotProductAttention for why forward_and_vjp is helpful.
+  """
+
+  def __init__(self, pre_attention, attention, post_attention):
+    self.pre_attention = tl.Serial([
+        # (x1_or_y1, x2) -> (x2, x1_or_y1, x2)
+        tl.Parallel([], tl.Dup()),
+        tl.Swap(),
+        tl.Parallel(pre_attention, [], []),
+    ])
+    assert hasattr(attention, 'forward_and_vjp')
+    self.attention = ApplyAttentionWrapper(attention)
+    self.post_attention = tl.Parallel(post_attention, [], [])
+
+    layers = [
+        self.pre_attention,
+        self.attention,
+        self.post_attention,
+        tl.Parallel(tl.Add(), []),
+    ]
+    super(ReversibleAttentionHalfResidual, self).__init__(layers)
+
+    self.subtract_top = tl.Parallel(tl.SubtractTop(), [])
+    self.reverse_layers = [
+        self.pre_attention,
+        self.attention,
+        self.post_attention,
+        self.subtract_top,
+    ]
+
+  def inverse_and_vjp(self, output, ct, params=(), **kwargs):
+    rng = kwargs.pop('rng', None)
+    rngs = (None,) * self._n_layers
+    if rng is not None:
+      rngs = backend.random.split(rng, self._n_layers)
+
+    if ct is None:
+      reconstructed_x = output
+      # Note that self.sublayers() aligns exactly with self.reverse_layers in
+      # terms of parameter and rng usage, so no re-ordering is required.
+      for layer, p, rng in zip(self.reverse_layers, params, rngs):
+        reconstructed_x = layer(reconstructed_x, p, rng=rng, **kwargs)
+      return reconstructed_x, None
+    else:
+      # Forward pass through self.pre_attention, while preparing for
+      # later backprop.
+      # Note: jax.vjp does not allow us to use **kwargs in the signature here.
+      def call_pre_attention(x, params, kwargs):
+        return self.pre_attention(x, params, **kwargs)
+      pre_attention_kwargs = kwargs.copy()
+      pre_attention_kwargs['rng'] = rngs[0]
+      stack, pre_attention_vjpfun = jax.vjp(
+          call_pre_attention, output, params[0], pre_attention_kwargs)
+
+      # Backprop through adding the residual
+      assert len(ct) == 2
+      ct = saved_ct = (ct[0], ct[0], ct[1])
+
+      # Backprop through self.post_attention with respect to the inputs only
+      call_post_attention_kwargs = kwargs.copy()
+      call_post_attention_kwargs['rng'] = rngs[2]
+      def call_post_attention(x):
+        return self.post_attention(x, params[2], **call_post_attention_kwargs)
+      # Note: these are *not* the actual inputs to self.post_attention.
+      # If self.post_attention is not linear, we will get incorrect gradients.
+      dummy_inputs = (stack[-3], stack[-2], stack[-1])
+      _, post_attention_vjpfun = jax.vjp(call_post_attention, dummy_inputs)
+      (ct,) = post_attention_vjpfun(ct)
+
+      # Simultaneous forward pass and backprop through the attention mechanism
+      attention_kwargs = kwargs.copy()
+      attention_kwargs['rng'] = rngs[1]
+      stack, ct = self.attention.forward_and_vjp(
+          stack, ct, **attention_kwargs)
+      attention_params_ct = ()
+
+      # Backprop through self.pre_attention
+      (x_ct,
+       pre_attention_params_ct,
+       pre_attention_kwargs_ct) = pre_attention_vjpfun(ct)
+
+      # Forward pass for self.post_attention, and backprop with respect to the
+      # parameters only
+      def call_post_attention2(params, kwargs):
+        return self.post_attention(stack, params, **kwargs)
+      stack, post_attention_vjpfun = jax.vjp(
+          call_post_attention2, params[2], call_post_attention_kwargs)
+      (post_attention_params_ct,
+       post_attention_kwargs_ct) = post_attention_vjpfun(saved_ct)
+
+      # Forward pass through subtracting the residual
+      reconstructed_x = self.subtract_top(
+          stack, params[-1], rng=rngs[-1], **kwargs)
+
+      params_ct = (
+          pre_attention_params_ct,
+          attention_params_ct,
+          post_attention_params_ct,
+          (),
+          )
+
+      # We don't actually backprop through the kwargs, but the API requires that
+      # we provide a value for kwargs_ct.
+      kwargs_ct = pre_attention_kwargs_ct
+      del post_attention_kwargs_ct
+
+      return reconstructed_x, (x_ct, params_ct, kwargs_ct)
+
+
 class ReversibleSwap(ReversibleLayerMixin, tl.Swap):
   """Swap the first two element on the stack."""
 
@@ -261,16 +587,6 @@ def inverse_and_vjp(self, output, ct, params=(), **kwargs):
           output, ct, params, **kwargs)
 
 
-def ReversibleResidual(layers_a, layers_b):
-  """RevNet-style reversible residual layer."""
-  return [
-      ReversibleHalfResidual(layers_a),  # (x1, x2) -> (z1, x2)
-      ReversibleSwap(),  # (z1, x2) -> (x2, z1)
-      ReversibleHalfResidual(layers_b),  # (x2, z1) -> (y2, z1)
-      ReversibleSwap(),  # (y2, z1) -> (z1, y2); where y1 := z1
-  ]
-
-
 class ReversibleSerial(ReversibleLayerMixin, tl.Serial):
   """A reversible version of tl.Serial (requires reversible sub-layers)."""
 
@@ -316,33 +632,47 @@ def inverse_and_vjp(self, output, ct, params=(), **kwargs):
 
 
 def DecoderBlock(d_feature, d_feedforward, n_heads, n_attention_chunks,
-                 dropout, mode):
+                 attention_loop_stride, dropout, mode):
   """Reversible transformer decoder layer.
 
   Args:
     d_feature: int:  depth of embedding
     d_feedforward: int: depth of feed-forward layer
     n_heads: int: number of attention heads
-    n_attention_chunks: int: number of chunks for memory-efficient attention
+    n_attention_chunks: int: number of chunks for attention
+    attention_loop_stride: int: number of query elements to compute attention
+      for in parallel. Set to 0 to disable memory-efficient attention.
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
 
   Returns:
     the layer.
   """
-  self_attention = [
+
+  pre_attention = [
+      Chunk(sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
       tl.LayerNorm(),
-      tl.Dup(),
-      tl.Parallel([], tl.CausalMask(axis=-2)),  # Create mask.
-      # TODO(kitaev): add dropout
-      tl.Attention(d_feature, n_heads=n_heads, dropout=None, mode=mode),
-      tl.Parallel([], tl.Drop()),  # Drop mask.
+      tl.Dup(), tl.Dup(),
+      tl.Parallel(
+          [tl.Dense(d_feature), SplitHeads(n_heads=n_heads)],  # pylint: disable=no-value-for-parameter
+          [tl.Dense(d_feature), SplitHeads(n_heads=n_heads)],  # pylint: disable=no-value-for-parameter
+          [tl.Dense(d_feature), SplitHeads(n_heads=n_heads)],  # pylint: disable=no-value-for-parameter
+      ),
   ]
 
-  # TODO(kitaev): Memory-efficient attention. This chunking is temporary.
-  self_attention = [
-      Chunk(sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
-      self_attention,
+  # TODO(kitaev): add dropout
+  if attention_loop_stride < 1:
+    # Use the standard implementation if no loop_stride is provided.
+    attention = DotProductAttention(dropout=None, mode=mode)
+  else:
+    attention = MemoryEfficientDotProductAttention(
+        loop_stride=attention_loop_stride, dropout=None, mode=mode)
+
+  # ReversibleAttentionHalfResidual requires that post_attention be linear in
+  # its input (so the backward pass can be computed without knowing the input)
+  post_attention = [
+      JoinHeads(),  # pylint: disable=no-value-for-parameter
+      tl.Dense(d_feature),
       Unchunk(sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
   ]
 
@@ -350,7 +680,10 @@ def DecoderBlock(d_feature, d_feedforward, n_heads, n_attention_chunks,
       FeedForward(d_feature, d_feedforward, dropout, mode=mode),
   ]
   return [
-      ReversibleResidual([self_attention], [feed_forward]),
+      ReversibleAttentionHalfResidual(pre_attention, attention, post_attention),
+      ReversibleSwap(),
+      ReversibleHalfResidual(feed_forward),
+      ReversibleSwap(),
   ]
 
 
@@ -363,6 +696,7 @@ def TransformerRevnetLM(vocab_size,
                         max_len=2048,
                         n_chunks=32,
                         n_attention_chunks=8,
+                        attention_loop_stride=0,
                         mode='train'):
   """Reversible transformer language model (only uses a decoder, no encoder).
 
@@ -375,7 +709,9 @@ def TransformerRevnetLM(vocab_size,
     dropout: float: dropout rate (how much to drop out)
     max_len: int: maximum symbol length for positional encoding
     n_chunks: int: number of chunks (must match input pipeline)
-    n_attention_chunks: int: number of chunks for memory-efficient attention
+    n_attention_chunks: int: number of chunks for attention
+    attention_loop_stride: int: number of query elements to compute attention
+      for in parallel. Set to 0 to disable memory-efficient attention.
     mode: str: 'train' or 'eval'
 
   Returns:
@@ -393,7 +729,7 @@ def TransformerRevnetLM(vocab_size,
       tl.Dup(),
       ReversibleSerial([
           DecoderBlock(d_feature, d_feedforward, n_heads, n_attention_chunks,
-                       dropout, mode)
+                       attention_loop_stride, dropout, mode)
           for _ in range(n_layers)
       ]),
       tl.Parallel(tl.LayerNorm(), tl.LayerNorm()),

From 58435743965dfc3bed28fda62664397fc18fff76 Mon Sep 17 00:00:00 2001
From: Zi Yang <ziy@google.com>
Date: Thu, 27 Jun 2019 15:44:19 -0700
Subject: [PATCH 2164/2720] [T2T] Fixed miscellaneous issues during TPU model
 exporting and serving.

1. Fixed "AttributeError: '_DefinedFunction' object has no attribute 'get_operations'"
  - Added function graph isinstance check in prune_unconnected_ops_from_xla

2. Fixed "ValueError: prediction values with "inputs: DatasetToSingleElement:0" must be from the default graph."
  - Added tf.identity to all values in the predictions as a workaround.

PiperOrigin-RevId: 255493712
---
 tensor2tensor/utils/t2t_model.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 284a7a081..27405470e 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1698,6 +1698,12 @@ def estimator_spec_predict(self, features, use_tpu=False):
       outputs = infer_out
       scores = None
 
+    # Workaround for "ValueError: prediction values must be from the default
+    # graph" during TPU model exporting.
+    # TODO(b/130501786): remove tf.identity once default graph mismatch is fixed
+    for name, feature in features.items():
+      features[name] = tf.identity(feature)
+
     inputs = features.get("inputs")
     if inputs is None:
       inputs = features["targets"]

From f98cc3035f58559015858807cfd0679d105ba056 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Fri, 28 Jun 2019 07:53:24 +0200
Subject: [PATCH 2165/2720] Add hiperparameters for DQN evaluation. (#1614)

---
 tensor2tensor/models/research/rl.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index a3ed18a6d..67b2ae719 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -421,6 +421,14 @@ def dqn_guess1_params():
   return hparams
 
 
+@registry.register_hparams
+def dqn_guess1_params_eval():
+  """Params for dqn_guess1 evaluation (with evaluator.py)."""
+  hparams = dqn_guess1_params()
+  hparams.set_hparam("eval_episodes_num", 64)
+  return hparams
+
+
 @registry.register_hparams
 def dqn_guess1_rainbow_params():
   """Guess 1 for DQN params."""

From 6fcb16ab091436122b34f2481b8103aa273e58a2 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 28 Jun 2019 14:19:26 -0700
Subject: [PATCH 2166/2720] Make activation functions configurable.

PiperOrigin-RevId: 255667138
---
 tensor2tensor/trax/layers/__init__.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/tensor2tensor/trax/layers/__init__.py b/tensor2tensor/trax/layers/__init__.py
index 87c3fe6b8..5eed08920 100644
--- a/tensor2tensor/trax/layers/__init__.py
+++ b/tensor2tensor/trax/layers/__init__.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import gin
 # We create a flat layers.* namespace for uniform calling conventions as we
 # upstream changes.
 # pylint: disable=wildcard-import
@@ -30,3 +31,21 @@
 from tensor2tensor.trax.layers.normalization import *
 from tensor2tensor.trax.layers.pooling import *
 from tensor2tensor.trax.layers.rnn import *
+
+
+# Ginify
+def layer_configure(*args, **kwargs):
+  kwargs["module"] = "trax.layers"
+  return gin.external_configurable(*args, **kwargs)
+
+# pylint: disable=used-before-assignment
+# pylint: disable=invalid-name
+Relu = layer_configure(Relu)
+Sigmoid = layer_configure(Sigmoid)
+Tanh = layer_configure(Tanh)
+HardSigmoid = layer_configure(HardSigmoid)
+HardTanh = layer_configure(HardTanh)
+Exp = layer_configure(Exp)
+LogSoftmax = layer_configure(LogSoftmax)
+Softmax = layer_configure(Softmax)
+Softplus = layer_configure(Softplus)

From a65aeb062d9e2bc181c0caf714493df8c2a6a564 Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Fri, 28 Jun 2019 15:16:51 -0700
Subject: [PATCH 2167/2720] Adding num_types to accompany type_ids in
 transformer_layers.transformer_prepare_encoder()

PiperOrigin-RevId: 255677664
---
 tensor2tensor/layers/transformer_layers.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 318d6077d..f61043fc8 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -33,7 +33,7 @@ def layers():
 
 
 def transformer_prepare_encoder(inputs, target_space, hparams, features=None,
-                                type_ids=None):
+                                type_ids=None, num_types=None):
   """Prepare one shard of the model for the encoder.
 
   Args:
@@ -44,6 +44,7 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None,
       This is needed now for "packed" datasets.
     type_ids: optional, an int64 Tensor of shape [batch, length] that allows
       for adding type embeddings, similar to positional embeddings.
+    num_types: optional, an int that decides the number of types in type_ids.
 
   Returns:
     encoder_input: a Tensor, bottom of encoder stack
@@ -113,9 +114,10 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None,
 
   # Add type embeddings
   if type_ids is not None:
+    if not num_types:
+      raise ValueError("Need to set num_types as well.")
     encoder_input = common_attention.add_positional_embedding(
-        encoder_input, hparams.max_length, "inputs_type_embedding",
-        type_ids)
+        encoder_input, num_types, "inputs_type_embedding", type_ids)
 
   encoder_self_attention_bias = common_layers.cast_like(
       encoder_self_attention_bias, encoder_input)

From 0d6b2ffe75ff9b2edc72ec90a661f7553f0f63ee Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Fri, 28 Jun 2019 17:48:25 -0700
Subject: [PATCH 2168/2720] Allow for specifying different bottleneck ratios in
 ResNet model

PiperOrigin-RevId: 255700177
---
 tensor2tensor/models/resnet.py | 50 ++++++++++++++++++++++++----------
 1 file changed, 36 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index ee0c3bc72..3e33e4c02 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -197,7 +197,8 @@ def residual_block(inputs,
                    data_format="channels_first",
                    use_td=False,
                    targeting_rate=None,
-                   keep_prob=None):
+                   keep_prob=None,
+                   bottleneck_ratio=None):
   """Standard building block for residual networks with BN before convolutions.
 
   Args:
@@ -220,11 +221,14 @@ def residual_block(inputs,
     targeting_rate: `float` proportion of weights to target with targeted
       dropout.
     keep_prob: `float` keep probability for targeted dropout.
+    bottleneck_ratio: unused parameter to keep the same function signature as
+        `bottleneck_block`.
 
   Returns:
     The output `Tensor` of the block.
   """
   del final_block
+  del bottleneck_ratio
   shortcut = inputs
   inputs = batch_norm_relu(inputs, is_training, data_format=data_format)
 
@@ -266,7 +270,8 @@ def bottleneck_block(inputs,
                      data_format="channels_first",
                      use_td=False,
                      targeting_rate=None,
-                     keep_prob=None):
+                     keep_prob=None,
+                     bottleneck_ratio=4):
   """Bottleneck block variant for residual networks with BN after convolutions.
 
   Args:
@@ -290,6 +295,8 @@ def bottleneck_block(inputs,
     targeting_rate: `float` proportion of weights to target with targeted
       dropout.
     keep_prob: `float` keep probability for targeted dropout.
+    bottleneck_ratio: `int`, how much we scale up filters.
+
 
   Returns:
     The output `Tensor` of the block.
@@ -327,7 +334,7 @@ def bottleneck_block(inputs,
   inputs = batch_norm_relu(inputs, is_training, data_format=data_format)
   inputs = conv2d_fixed_padding(
       inputs=inputs,
-      filters=4 * filters,
+      filters=bottleneck_ratio * filters,
       kernel_size=1,
       strides=1,
       data_format=data_format,
@@ -355,7 +362,8 @@ def block_layer(inputs,
                 data_format="channels_first",
                 use_td=False,
                 targeting_rate=None,
-                keep_prob=None):
+                keep_prob=None,
+                bottleneck_ratio=4):
   """Creates one layer of blocks for the ResNet model.
 
   Args:
@@ -374,12 +382,15 @@ def block_layer(inputs,
     targeting_rate: `float` proportion of weights to target with targeted
       dropout.
     keep_prob: `float` keep probability for targeted dropout.
+    bottleneck_ratio: `int`, how much we scale up filters in bottleneck block.
 
   Returns:
     The output `Tensor` of the block layer.
   """
-  # Bottleneck blocks end with 4x the number of filters as they start with
-  filters_out = 4 * filters if block_fn is bottleneck_block else filters
+  # Bottleneck blocks end with bottleneck_ratio x the number of filters
+  filters_out = filters
+  if block_fn is bottleneck_block:
+    filters_out = bottleneck_ratio * filters
 
   def projection_shortcut(inputs):
     """Project identity branch."""
@@ -407,7 +418,8 @@ def projection_shortcut(inputs):
       data_format,
       use_td=use_td,
       targeting_rate=targeting_rate,
-      keep_prob=keep_prob)
+      keep_prob=keep_prob,
+      bottleneck_ratio=bottleneck_ratio)
 
   for i in range(1, blocks):
     inputs = block_fn(
@@ -419,7 +431,8 @@ def projection_shortcut(inputs):
         data_format,
         use_td=use_td,
         targeting_rate=targeting_rate,
-        keep_prob=keep_prob)
+        keep_prob=keep_prob,
+        bottleneck_ratio=bottleneck_ratio)
 
   return tf.identity(inputs, name)
 
@@ -433,7 +446,8 @@ def resnet_v2(inputs,
               is_cifar=False,
               use_td=False,
               targeting_rate=None,
-              keep_prob=None):
+              keep_prob=None,
+              bottleneck_ratios=None):
   """Resnet model.
 
   Args:
@@ -454,6 +468,8 @@ def resnet_v2(inputs,
     targeting_rate: `float` proportion of weights to target with targeted
       dropout.
     keep_prob: `float` keep probability for targeted dropout.
+    bottleneck_ratios: list of `int`s, how much we scale up filters in
+      bottleneck blocks.
 
   Returns:
     Pre-logit activations.
@@ -469,7 +485,8 @@ def resnet_v2(inputs,
       data_format=data_format,
       use_td=use_td,
       targeting_rate=targeting_rate,
-      keep_prob=keep_prob)
+      keep_prob=keep_prob,
+      bottleneck_ratio=bottleneck_ratios[0])
   inputs = block_layer(
       inputs=inputs,
       filters=filters[2],
@@ -481,7 +498,8 @@ def resnet_v2(inputs,
       data_format=data_format,
       use_td=use_td,
       targeting_rate=targeting_rate,
-      keep_prob=keep_prob)
+      keep_prob=keep_prob,
+      bottleneck_ratio=bottleneck_ratios[1])
   inputs = block_layer(
       inputs=inputs,
       filters=filters[3],
@@ -493,7 +511,8 @@ def resnet_v2(inputs,
       data_format=data_format,
       use_td=use_td,
       targeting_rate=targeting_rate,
-      keep_prob=keep_prob)
+      keep_prob=keep_prob,
+      bottleneck_ratio=bottleneck_ratios[2])
   if not is_cifar:
     inputs = block_layer(
         inputs=inputs,
@@ -506,7 +525,8 @@ def resnet_v2(inputs,
         data_format=data_format,
         use_td=use_td,
         targeting_rate=targeting_rate,
-        keep_prob=keep_prob)
+        keep_prob=keep_prob,
+        bottleneck_ratio=bottleneck_ratios[3])
 
   return inputs
 
@@ -562,7 +582,8 @@ def body(self, features):
         is_cifar=hp.is_cifar,
         use_td=hp.use_td,
         targeting_rate=hp.targeting_rate,
-        keep_prob=hp.keep_prob)
+        keep_prob=hp.keep_prob,
+        bottleneck_ratios=hp.bottleneck_ratios)
 
     if hp.use_nchw:
       out = tf.transpose(out, [0, 2, 3, 1])
@@ -620,6 +641,7 @@ def resnet_base():
 
   # Model-specific parameters
   hparams.add_hparam("layer_sizes", [3, 4, 6, 3])
+  hparams.add_hparam("bottleneck_ratios", [4, 4, 4, 4])
   hparams.add_hparam("filter_sizes", [64, 64, 128, 256, 512])
   hparams.add_hparam("block_fn", "bottleneck")
   hparams.add_hparam("use_nchw", True)

From c1c2b84705d4d2df92609e6ec0e6cd02cad4df44 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 28 Jun 2019 17:57:08 -0700
Subject: [PATCH 2169/2720] Adjust PositionLookupTransformer to recent stack
 semantics change.

PiperOrigin-RevId: 255701020
---
 .../research/position_lookup_transformer.py   | 73 +++++++------------
 1 file changed, 28 insertions(+), 45 deletions(-)

diff --git a/tensor2tensor/trax/models/research/position_lookup_transformer.py b/tensor2tensor/trax/models/research/position_lookup_transformer.py
index 596adedc8..3b84fa265 100644
--- a/tensor2tensor/trax/models/research/position_lookup_transformer.py
+++ b/tensor2tensor/trax/models/research/position_lookup_transformer.py
@@ -32,6 +32,14 @@
 _POSITIONS = onp.random.uniform(size=[_ABSOLUTE_MAX_LEN, POS_VECTOR_SIZE])
 
 
+def Dup2():
+  """Copy first 2 elements of the stack: (a, b, ...) -> (a, b, a, b, ...)."""
+  return [                              # Stack is (a, b, ...)
+      tl.Parallel(tl.Dup(), tl.Dup()),  # Stack is (a, a, b, b, ...)
+      tl.Parallel([], tl.Swap())        # Stack is (a, b, a, b, ...)
+  ]
+
+
 @tl.layer()
 def NewPositionalEncoding(x, positions=None, **kwargs):
   """Implements new positional encoding."""
@@ -43,15 +51,10 @@ def NewPositionalEncoding(x, positions=None, **kwargs):
   return res
 
 
-# TODO(lukaszkaiser): This used to have stack_items_to_pass=0; fix as needed.
-@tl.layer()
-def CutPosition(xs, **unused_kwargs):
+@tl.layer(n_inputs=1, n_outputs=2)
+def CutAtPosition(x, **unused_kwargs):
   """Splits x into a pair (x[:position], position)."""
-  if not isinstance(xs, (list, tuple)):
-    xs = [xs]
-  x = xs[0]
-  res = [x[:, :, :-POS_VECTOR_SIZE], x[:, :, -POS_VECTOR_SIZE:]]
-  return tuple(res + list(xs[1:]))
+  return tuple([x[:, :, :-POS_VECTOR_SIZE], x[:, :, -POS_VECTOR_SIZE:]])
 
 
 @tl.layer()
@@ -114,30 +117,12 @@ def DeepFlatten(xs):
       yield x
 
 
-# TODO(lukaszkaiser): This used to have stack_items_to_pass=0; fix as needed.
-@tl.layer()
-def Unnest(xs, **unused_kwargs):
-  return [x for x in DeepFlatten(xs)]
-
-
-# TODO(lukaszkaiser): This used to have stack_items_to_pass=0; fix as needed.
-@tl.layer()
-def ConcatenateN(xs, params, n=2, axis=-1, **kwargs):
-  """Concatenate first N inputs (and output remainder as is if non-empty)."""
-  del params, kwargs
-  res = np.concatenate(xs[:n], axis)
-  rest = list(xs[n:])
-  if rest:
-    return tuple([res] + rest)
-  return res
-
-
 def PreservePosition(layer):
   """Execute layer without position but preserve it in parallel."""
   return tl.Serial(
-      CutPosition(),
+      CutAtPosition(),
       layer,
-      ConcatenateN()
+      tl.Concatenate(n_items=2)
   )
 
 
@@ -157,13 +142,13 @@ def ApplyAndQueryPositions(layer, pos):
   """
   n_heads = len(pos)
   return tl.Serial(
-      tl.Dup(),
-      CutPosition(),
-      # TODO(lukaszkaiser): Rewrite without using Select.
-      tl.Select(tuple([0] + [(2, 1)]*n_heads)),
+      tl.Dup(),                    # (x, x)
+      CutAtPosition(),          # (x_content, x_position, x)
+      tl.Parallel([], tl.Swap()),  # (x_content, x, x_position)
+      [tl.Parallel([], Dup2()) for _ in range(n_heads - 1)],
+      # Now the stack is x_content, (x, x_position) * n_heads.
       tl.Parallel(*([layer] + pos)),
-      Unnest(),
-      ConcatenateN(n=n_heads + 1)
+      tl.Concatenate(n_items=n_heads + 1)
   )
 
 
@@ -188,22 +173,21 @@ def LearnedQP(keys=None, values=None, binary=False):
   )
 
 
-# TODO(lukaszkaiser): This used to have stack_items_to_pass=0; fix as needed.
-@tl.layer()
-def SoftmaxBranches(x_list_in, n_branches=2, **unused_kwargs):
+@tl.layer(n_inputs=10, n_outputs=1)
+def Softmax5Branches(x_list, n_branches=2, **unused_kwargs):
   """Softmax xs.
 
   The input xs is a list of embeddings and weights of the form
   w_1 e_1 .... w_n e_n (followed by optional rest that is preserved).
 
   Args:
-    x_list_in: the input weights and embeddings.
+    x_list: the input weights and embeddings.
     n_branches: what part of the list to use.
 
   Returns:
     softmax(w) * e for the joint weights w and embeddings e.
   """
-  x_list, x_list_rest = x_list_in[:2*n_branches], x_list_in[2*n_branches:]
+  assert n_branches == 5
   softmax_activations = [x_list[2*i] for i in range(n_branches)]
   max_sa = softmax_activations[0]
   for x in softmax_activations:
@@ -213,7 +197,7 @@ def SoftmaxBranches(x_list_in, n_branches=2, **unused_kwargs):
   sum_sa = sum(softmax_activations)
   softmax_activations = [x / sum_sa for x in softmax_activations]
   res = sum([x_list[2*i+1] * softmax_activations[i] for i in range(n_branches)])
-  return tuple([res] + list(x_list_rest))
+  return res
 
 
 def SumLearnedPick(positions):
@@ -233,7 +217,7 @@ def SumLearnedPick(positions):
   sub_values = np.array([positions[max(i - j, 0), :]
                          for j in range(l) for i in range(l)])
   return tl.Serial(
-      tl.Dup(), tl.Dup(), tl.Dup(), tl.Dup(),
+      Dup2(), Dup2(), Dup2(), Dup2(),
       tl.Parallel(
           LearnedQP(),
           LearnedQP(keys=succ_keys, values=succ_values),
@@ -241,8 +225,7 @@ def SumLearnedPick(positions):
           LearnedQP(keys=add_keys, values=add_values, binary=True),
           LearnedQP(keys=sub_keys, values=sub_values, binary=True),
       ),
-      Unnest(),
-      SoftmaxBranches(n_branches=5)
+      Softmax5Branches(n_branches=5)
   )
 
 
@@ -345,7 +328,7 @@ def PositionLookupTransformerLM(vocab_size=128,
     the layer.
   """
   positions = _POSITIONS[:max_len, :]
-  return tl.Serial([
+  return tl.Serial(
       tl.ShiftRight(),
       tl.Embedding(d_feature, vocab_size),
       tl.Dropout(rate=dropout, mode=mode),
@@ -355,4 +338,4 @@ def PositionLookupTransformerLM(vocab_size=128,
       PreservePosition(tl.LayerNorm()),
       tl.Dense(vocab_size),
       tl.LogSoftmax()
-  ])
+  )

From 1c77b7279f55401f7f7c00280b1bc414cc2d4f4d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 28 Jun 2019 20:14:08 -0700
Subject: [PATCH 2170/2720] Skip the trajectory_collection/misc time, this is
 just confusing.

Right now we do:
                    trajectory_collection :     877.29
        trajectory_collection/env_actions :     261.57
               trajectory_collection/misc :     682.60
 trajectory_collection/policy_application :     193.81

And ideally we'd like sum of trajectory_collection/* to be <= the value of trajectory_collection, with the rest as misc.

PiperOrigin-RevId: 255711981
---
 tensor2tensor/envs/env_problem_utils.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 2ac197257..715d22260 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -87,7 +87,6 @@ def play_env_problem_with_policy(env,
     A tuple, (trajectories, number of completed trajectories). Where
     trajectories is a list of triples of (observation, action, reward) ndarrays.
   """
-  t0 = time.time()
 
   def categorical_sample(log_probs):
     """Categorical sampling."""
@@ -225,10 +224,8 @@ def epsilon_greedy(log_probs):
   for trajectory in env.trajectories.completed_trajectories[:num_trajectories]:
     completed_trajectories.append(trajectory.as_numpy)
 
-  misc_time = (time.time() - t0) - policy_application_total_time
   timing_info = {
       "trajectory_collection/policy_application": policy_application_total_time,
-      "trajectory_collection/misc": misc_time,
       "trajectory_collection/env_actions": env_actions_total_time,
   }
   timing_info = {k: round(1000 * v, 2) for k, v in timing_info.items()}

From 0c8e40834b2ffb8c95217ce6abb4c2e40eda23e8 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 28 Jun 2019 22:20:11 -0700
Subject: [PATCH 2171/2720] Expose the actual run-time of an env in the stats.
 This is to see which part in trajectory collection is slow.

Sample run:
                      trajectory_collection :    1308.29
          trajectory_collection/env_actions :     516.65
 trajectory_collection/env_actions/bare_env :     505.12
   trajectory_collection/policy_application :     346.97

PiperOrigin-RevId: 255720590
---
 tensor2tensor/envs/env_problem.py       | 3 +++
 tensor2tensor/envs/env_problem_utils.py | 5 ++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 2b8c63bff..930259311 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -23,6 +23,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import time
 import gym
 from gym.core import Env
 import numpy as np
@@ -558,7 +559,9 @@ def _step(self, actions):
 
     # Take steps in all environments.
     for env, action in zip(self._envs, actions):
+      t1 = time.time()
       observation, reward, done, info = env.step(action)
+      info["__bare_env_run_time__"] = time.time() - t1
 
       observations.append(observation)
       rewards.append(reward)
diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 715d22260..11aebd68d 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -147,6 +147,7 @@ def epsilon_greedy(log_probs):
 
   policy_application_total_time = 0
   env_actions_total_time = 0
+  bare_env_run_time = 0
   while env.trajectories.num_completed_trajectories < num_trajectories:
     # Get all the observations for all the active trajectories.
     # Shape is (B, T) + OBS
@@ -187,8 +188,9 @@ def epsilon_greedy(log_probs):
 
     # Step through the env.
     t1 = time.time()
-    _, _, dones, _ = env.step(actions)
+    _, _, dones, infos = env.step(actions)
     env_actions_total_time += (time.time() - t1)
+    bare_env_run_time += sum(info["__bare_env_run_time__"] for info in infos)
 
     # Count the number of done trajectories, the others could just have been
     # truncated.
@@ -227,6 +229,7 @@ def epsilon_greedy(log_probs):
   timing_info = {
       "trajectory_collection/policy_application": policy_application_total_time,
       "trajectory_collection/env_actions": env_actions_total_time,
+      "trajectory_collection/env_actions/bare_env": bare_env_run_time,
   }
   timing_info = {k: round(1000 * v, 2) for k, v in timing_info.items()}
 

From 1fb49ea481925128e4fe8a11652f959f327b7bc9 Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Sat, 29 Jun 2019 00:50:06 -0700
Subject: [PATCH 2172/2720] Modify t2t_distill to add student_dir to hparams
 (beside teacher_dir)

PiperOrigin-RevId: 255730407
---
 tensor2tensor/bin/t2t_distill.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/bin/t2t_distill.py b/tensor2tensor/bin/t2t_distill.py
index e58b3c218..1815d0b10 100644
--- a/tensor2tensor/bin/t2t_distill.py
+++ b/tensor2tensor/bin/t2t_distill.py
@@ -107,6 +107,7 @@ def main(argv):
   else:
     student_dir = os.path.join(root_output_dir, "student")
   FLAGS.output_dir = student_dir
+  hparams.add_hparam("student_dir", student_dir)
 
   exp_fn = t2t_trainer.create_experiment_fn()
   run_config = t2t_trainer.create_run_config(hparams)
@@ -165,6 +166,7 @@ def create_student_experiment(run_config, hparams, argv):
     t2t_trainer.set_hparams_from_args(argv[1:])
 
   hparams.add_hparam("teacher_dir", FLAGS.teacher_dir)
+  hparams.add_hparam("student_dir", FLAGS.student_dir)
   hparams.distill_phase = "distill"
   exp_fn = t2t_trainer.create_experiment_fn()
   exp = exp_fn(run_config, hparams)

From d38f3435ded822e585d1fc7136f3ece857a41c8d Mon Sep 17 00:00:00 2001
From: Zi Yang <ziy@google.com>
Date: Mon, 1 Jul 2019 13:57:21 -0700
Subject: [PATCH 2173/2720] [T2T] Fixed high usage of TPU HBM "Arguments"
 during serving   - Added flag for export_saved_model_api_version (default to
 1)   - Added maybe_use_guarantee_const_getter_model_fn and
 use_guarantee_const_getter flag. It marks all weights as constant, which may
 improves TPU inference performance because it prevents the weights being
 transferred to the TPU. It will increase HBM "program" usage and reduce HBM
 "arguments" usage during TPU model serving.

PiperOrigin-RevId: 256026810
---
 tensor2tensor/bin/t2t_trainer.py   | 11 +++++
 tensor2tensor/serving/export.py    |  4 +-
 tensor2tensor/utils/trainer_lib.py | 68 ++++++++++++++++++++++++++++--
 3 files changed, 78 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 0c76819c6..d146f84a7 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -58,6 +58,15 @@
 flags.DEFINE_bool("use_tpu", False, "Whether to use TPU.")
 flags.DEFINE_bool("use_tpu_estimator", False, "Whether to use TPUEstimator. "
                   "This is always enabled when use_tpu is True.")
+flags.DEFINE_integer("export_saved_model_api_version", 1,
+                     "ExportSavedModelApiVersion, 1 (V1, default) or 2 (V2). "
+                     "Default V2 uses model_fn_inference_on_tpu for rewrite."
+                     "Flag use_guarantee_const is only enabled in V2.")
+flags.DEFINE_bool("use_guarantee_const_getter", False,
+                  "Whether to use GuaranteeConst Ops to mark all weights as "
+                  "constant. It may improve TPU inference performance and "
+                  "reduce HBM arguments usage. Only available when "
+                  "export_saved_model_api_version=2 and use_tpu=True.")
 flags.DEFINE_bool("xla_compile", False,
                   "Whether to use XLA to compile model_fn.")
 flags.DEFINE_integer("xla_jit_level", -1,
@@ -197,6 +206,8 @@ def create_experiment_fn():
       use_tpu=FLAGS.use_tpu,
       use_tpu_estimator=FLAGS.use_tpu_estimator,
       use_xla=FLAGS.xla_compile,
+      export_saved_model_api_version=FLAGS.export_saved_model_api_version,
+      use_guarantee_const_getter=FLAGS.use_guarantee_const_getter,
       warm_start_from=FLAGS.warm_start_from,
       decode_from_file=FLAGS.decode_from_file,
       decode_to_file=FLAGS.decode_to_file,
diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 55bcbeec5..2737f6961 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -82,7 +82,9 @@ def create_estimator(run_config, hparams):
       hparams,
       run_config,
       decode_hparams=decoding.decode_hparams(FLAGS.decode_hparams),
-      use_tpu=FLAGS.use_tpu)
+      use_tpu=FLAGS.use_tpu,
+      export_saved_model_api_version=FLAGS.export_saved_model_api_version,
+      use_guarantee_const_getter=FLAGS.use_guarantee_const_getter)
 
 
 def create_hparams():
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 7dc6f2448..42c8daff8 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import collections
+import contextlib
 import json
 import os
 import random
@@ -35,6 +36,7 @@
 
 import tensorflow as tf
 
+from tensorflow.contrib.tpu.python.tpu import tpu_estimator
 from tensorflow.core.protobuf import rewriter_config_pb2
 from tensorflow.python import debug
 
@@ -285,7 +287,9 @@ def create_estimator(model_name,
                      decode_hparams=None,
                      use_tpu=False,
                      use_tpu_estimator=False,
-                     use_xla=False):
+                     use_xla=False,
+                     export_saved_model_api_version=1,
+                     use_guarantee_const_getter=False):
   """Create a T2T Estimator."""
   model_fn = t2t_model.T2TModel.make_estimator_model_fn(
       model_name, hparams, decode_hparams=decode_hparams, use_tpu=use_tpu)
@@ -307,14 +311,66 @@ def create_estimator(model_name,
     if decode_hparams and run_config.tpu_config:
       decode_hparams.add_hparam("iterations_per_loop",
                                 run_config.tpu_config.iterations_per_loop)
+    if export_saved_model_api_version == 1:
+      api_version_enum_name = tpu_estimator.ExportSavedModelApiVersion.V1
+      estimator_model_fn = model_fn
+    elif export_saved_model_api_version == 2:
+      api_version_enum_name = tpu_estimator.ExportSavedModelApiVersion.V2
+
+      def maybe_use_guarantee_const_getter_model_fn(features, labels, mode,
+                                                    params):
+        """Wrapper model_fn with guarantee_const getter."""
+        if not use_guarantee_const_getter:
+          return model_fn(features, labels, mode, params)
+
+        # It marks all weights as constant, which may improves TPU inference
+        # performance because it prevents the weights being transferred to the
+        # TPU. It will increase HBM "program" usage and reduce HBM "arguments"
+        # usage during TPU model serving.
+        def guarantee_const_getter(getter, name, *args, **kwargs):
+          with tf.control_dependencies(None):
+            return tf.guarantee_const(
+                getter(name, *args, **kwargs), name=name + "/GuaranteeConst")
+
+        @contextlib.contextmanager
+        def guarantee_const_scope():
+          var_scope = tf.get_variable_scope()
+          prev_custom_getter = var_scope.custom_getter
+          prev_caching_device = var_scope.caching_device
+          var_scope.set_custom_getter(guarantee_const_getter)
+          var_scope.set_caching_device(lambda op: op.device)
+          yield
+          var_scope.set_custom_getter(prev_custom_getter)
+          var_scope.set_caching_device(prev_caching_device)
+
+        with guarantee_const_scope():
+          return model_fn(features, labels, mode, params)
+
+      def tpu_model_fn(features, labels, mode, params):
+        """Wrapper model_fn with tpu.rewrite / TPUPartitionedCall."""
+        if mode == tf.estimator.ModeKeys.PREDICT and params["use_tpu"]:
+          return tpu_estimator.model_fn_inference_on_tpu(
+              maybe_use_guarantee_const_getter_model_fn,
+              features=features,
+              labels=labels,
+              config=None,
+              params=params,
+              batch_config=None)
+        else:
+          return model_fn(features, labels, mode, params)
+
+      estimator_model_fn = tpu_model_fn
+    else:
+      raise ValueError("Flag export_saved_model_api_version must be 1 or 2.")
     estimator = tf.contrib.tpu.TPUEstimator(
-        model_fn=model_fn,
+        model_fn=estimator_model_fn,
         model_dir=run_config.model_dir,
         config=run_config,
         use_tpu=use_tpu,
         train_batch_size=batch_size,
         eval_batch_size=batch_size if "eval" in schedule else None,
-        predict_batch_size=predict_batch_size)
+        predict_batch_size=predict_batch_size,
+        export_saved_model_api_version=api_version_enum_name)
   else:
     estimator = tf.estimator.Estimator(
         model_fn=model_fn,
@@ -633,6 +689,8 @@ def create_experiment(
     use_tpu=False,
     use_tpu_estimator=False,
     use_xla=False,
+    export_saved_model_api_version=1,
+    use_guarantee_const_getter=False,
     additional_train_hooks=None,
     additional_eval_hooks=None,
     warm_start_from=None,
@@ -668,7 +726,9 @@ def create_experiment(
       decode_hparams=decode_hparams,
       use_tpu=use_tpu,
       use_tpu_estimator=use_tpu_estimator,
-      use_xla=use_xla)
+      use_xla=use_xla,
+      export_saved_model_api_version=export_saved_model_api_version,
+      use_guarantee_const_getter=use_guarantee_const_getter)
 
   # Input fns from Problem
   problem = hparams.problem

From 28935cd38f3b8347f8d78f45bb6a426d56206c5f Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Mon, 1 Jul 2019 17:33:55 -0700
Subject: [PATCH 2174/2720] When loading hparams from json, only modify the
 hparams if "hparams" arg is specified.

PiperOrigin-RevId: 256068614
---
 tensor2tensor/utils/hparams_lib.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/utils/hparams_lib.py b/tensor2tensor/utils/hparams_lib.py
index ff7f8756e..07be4e4f6 100644
--- a/tensor2tensor/utils/hparams_lib.py
+++ b/tensor2tensor/utils/hparams_lib.py
@@ -67,11 +67,12 @@ def create_hparams_from_json(json_path, hparams=None):
     # Prevent certain keys from overwriting the passed-in hparams.
     # TODO(trandustin): Remove this hack after registries are available to avoid
     # saving them as functions.
-    hparams_values.pop("bottom", None)
-    hparams_values.pop("loss", None)
-    hparams_values.pop("name", None)
-    hparams_values.pop("top", None)
-    hparams_values.pop("weights_fn", None)
+    if hparams:
+      hparams_values.pop("bottom", None)
+      hparams_values.pop("loss", None)
+      hparams_values.pop("name", None)
+      hparams_values.pop("top", None)
+      hparams_values.pop("weights_fn", None)
     new_hparams = hparam.HParams(**hparams_values)
     # Some keys are in new_hparams but not hparams, so we need to be more
     #   careful than simply using parse_json() from HParams

From 07db31dfb8510f5c283bdb769505ee48402b4194 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 1 Jul 2019 18:12:12 -0700
Subject: [PATCH 2175/2720] Move rlax envs to a new package

PiperOrigin-RevId: 256074302
---
 tensor2tensor/trax/rlax/envs/__init__.py          | 15 +++++++++++++++
 tensor2tensor/trax/rlax/{ => envs}/fake_env.py    |  0
 .../trax/rlax/{ => envs}/fake_env_test.py         |  2 +-
 3 files changed, 16 insertions(+), 1 deletion(-)
 create mode 100644 tensor2tensor/trax/rlax/envs/__init__.py
 rename tensor2tensor/trax/rlax/{ => envs}/fake_env.py (100%)
 rename tensor2tensor/trax/rlax/{ => envs}/fake_env_test.py (97%)

diff --git a/tensor2tensor/trax/rlax/envs/__init__.py b/tensor2tensor/trax/rlax/envs/__init__.py
new file mode 100644
index 000000000..4872e5d5d
--- /dev/null
+++ b/tensor2tensor/trax/rlax/envs/__init__.py
@@ -0,0 +1,15 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
diff --git a/tensor2tensor/trax/rlax/fake_env.py b/tensor2tensor/trax/rlax/envs/fake_env.py
similarity index 100%
rename from tensor2tensor/trax/rlax/fake_env.py
rename to tensor2tensor/trax/rlax/envs/fake_env.py
diff --git a/tensor2tensor/trax/rlax/fake_env_test.py b/tensor2tensor/trax/rlax/envs/fake_env_test.py
similarity index 97%
rename from tensor2tensor/trax/rlax/fake_env_test.py
rename to tensor2tensor/trax/rlax/envs/fake_env_test.py
index 8c9ee9771..e6dabf471 100644
--- a/tensor2tensor/trax/rlax/fake_env_test.py
+++ b/tensor2tensor/trax/rlax/envs/fake_env_test.py
@@ -19,7 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.trax.rlax import fake_env
+from tensor2tensor.trax.rlax.envs import fake_env
 from tensorflow import test
 
 

From 95866303475da15ca30366a9266041504a63fc62 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 3 Jul 2019 03:20:42 -0700
Subject: [PATCH 2176/2720] Move input preparation to separate methods in
 transformer.

PiperOrigin-RevId: 256335110
---
 tensor2tensor/models/transformer.py | 105 +++++++++++++++-------------
 1 file changed, 56 insertions(+), 49 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index a33aed536..1968d1ca7 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -241,7 +241,7 @@ def body(self, features):
     losses = []
 
     if self.has_input:
-      inputs = features["inputs"]
+      inputs = self._prepare_inputs_for_body(features)
       target_space = features["target_space_id"]
       encoder_output, encoder_decoder_attention_bias = self.encode(
           inputs, target_space, hparams, features=features, losses=losses)
@@ -298,6 +298,20 @@ def body(self, features):
     else:
       return ret
 
+  def _prepare_inputs_for_body(self, features):
+    """Prepare inputs for body.
+
+    Args:
+      features: Map of string to model features. Should contain
+          "inputs": Transformer inputs. [batch_size, input_length, 1,
+            hidden_dim].
+
+    Returns:
+      Inputs which will be passed to the model. [batch_size, input_length, 1,
+          hidden_dim]
+    """
+    return features["inputs"]
+
   def _greedy_infer(self, features, decode_length, use_tpu=False):
     """Fast version of greedy decoding.
 
@@ -370,6 +384,39 @@ def _beam_decode(self,
       return self._fast_decode(features, decode_length, beam_size, top_beams,
                                alpha)
 
+  def _prepare_inputs_for_decode(self, features):
+    """Prepare inputs for decoding.
+
+    Args:
+      features: A map of string to model features.
+
+    Returns:
+      Inputs after fixing shape and applying modality.
+    """
+    dp = self._data_parallelism
+    hparams = self._hparams
+    inputs = features["inputs"]
+    # TODO(llion): Clean up this reshaping logic.
+    inputs = tf.expand_dims(inputs, axis=1)
+    if len(inputs.shape) < 5:
+      inputs = tf.expand_dims(inputs, axis=4)
+    s = common_layers.shape_list(inputs)
+    inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
+    # _shard_features called to ensure that the variable names match
+    inputs = self._shard_features({"inputs": inputs})["inputs"]
+    input_modality = self._problem_hparams.modality["inputs"]
+    input_vocab_size = self._problem_hparams.vocab_size["inputs"]
+    if input_vocab_size is not None and hasattr(hparams, "vocab_divisor"):
+      input_vocab_size += (-input_vocab_size) % hparams.vocab_divisor
+    modality_name = hparams.name.get("inputs",
+                                     modalities.get_name(input_modality))(
+                                         hparams, input_vocab_size)
+    with tf.variable_scope(modality_name):
+      bottom = hparams.bottom.get("inputs",
+                                  modalities.get_bottom(input_modality))
+      inputs = dp(bottom, inputs, hparams, input_vocab_size)
+    return inputs
+
   def _fast_decode_tpu(self,
                        features,
                        decode_length,
@@ -416,34 +463,14 @@ def _fast_decode_tpu(self,
       target_vocab_size += (-target_vocab_size) % hparams.vocab_divisor
 
     if self.has_input:
-      inputs = features["inputs"]
+      inputs_shape = common_layers.shape_list(features["inputs"])
       if target_modality == modalities.ModalityType.CLASS_LABEL:
         decode_length = 1
       else:
         decode_length = (
-            common_layers.shape_list(inputs)[1] + features.get(
-                "decode_length", decode_length))
-
-      # TODO(llion): Clean up this reshaping logic.
-      inputs = tf.expand_dims(inputs, axis=1)
-      if len(inputs.shape) < 5:
-        inputs = tf.expand_dims(inputs, axis=4)
-      s = common_layers.shape_list(inputs)
-      batch_size = s[0]
-      inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
-      # _shard_features called to ensure that the variable names match
-      inputs = self._shard_features({"inputs": inputs})["inputs"]
-      input_modality = self._problem_hparams.modality["inputs"]
-      input_vocab_size = self._problem_hparams.vocab_size["inputs"]
-      if input_vocab_size is not None and hasattr(hparams, "vocab_divisor"):
-        input_vocab_size += (-input_vocab_size) % hparams.vocab_divisor
-      modality_name = hparams.name.get(
-          "inputs",
-          modalities.get_name(input_modality))(hparams, input_vocab_size)
-      with tf.variable_scope(modality_name):
-        bottom = hparams.bottom.get("inputs",
-                                    modalities.get_bottom(input_modality))
-        inputs = dp(bottom, inputs, hparams, input_vocab_size)
+            inputs_shape[1] + features.get("decode_length", decode_length))
+      batch_size = inputs_shape[0]
+      inputs = self._prepare_inputs_for_decode(features)
       with tf.variable_scope("body"):
         encoder_output, encoder_decoder_attention_bias = dp(
             self.encode,
@@ -678,34 +705,14 @@ def _fast_decode(self,
           " If you want to decode from a dataset, use the non-packed version"
           " of the dataset when decoding.")
     if self.has_input:
-      inputs = features["inputs"]
+      inputs_shape = common_layers.shape_list(features["inputs"])
       if target_modality == modalities.ModalityType.CLASS_LABEL:
         decode_length = 1
       else:
         decode_length = (
-            common_layers.shape_list(inputs)[1] + features.get(
-                "decode_length", decode_length))
-
-      # TODO(llion): Clean up this reshaping logic.
-      inputs = tf.expand_dims(inputs, axis=1)
-      if len(inputs.shape) < 5:
-        inputs = tf.expand_dims(inputs, axis=4)
-      s = common_layers.shape_list(inputs)
-      batch_size = s[0]
-      inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
-      # _shard_features called to ensure that the variable names match
-      inputs = self._shard_features({"inputs": inputs})["inputs"]
-      input_modality = self._problem_hparams.modality["inputs"]
-      input_vocab_size = self._problem_hparams.vocab_size["inputs"]
-      if input_vocab_size is not None and hasattr(hparams, "vocab_divisor"):
-        input_vocab_size += (-input_vocab_size) % hparams.vocab_divisor
-      modality_name = hparams.name.get(
-          "inputs",
-          modalities.get_name(input_modality))(hparams, input_vocab_size)
-      with tf.variable_scope(modality_name):
-        bottom = hparams.bottom.get("inputs",
-                                    modalities.get_bottom(input_modality))
-        inputs = dp(bottom, inputs, hparams, input_vocab_size)
+            inputs_shape[1] + features.get("decode_length", decode_length))
+      batch_size = inputs_shape[0]
+      inputs = self._prepare_inputs_for_decode(features)
       with tf.variable_scope("body"):
         encoder_output, encoder_decoder_attention_bias = dp(
             self.encode,

From 8fd75ea4c6d2dd54d47ac0ac0da88e838c41fcfb Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 3 Jul 2019 11:27:17 -0700
Subject: [PATCH 2177/2720] Internal

PiperOrigin-RevId: 256405806
---
 tensor2tensor/layers/common_layers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index d25f5e7e9..09c3a5d34 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1985,7 +1985,7 @@ def padded_cross_entropy_mixture(logits,
 
     # plot a summary for the difference between the top 2 losses
     if num_mixtures > 1:
-      xent_reshaped = tf.transpose(tf.squeeze(xent), perm=[1, 0])
+      xent_reshaped = tf.transpose(tf.squeeze(xent, axis=[2, 3]), perm=[1, 0])
       top_2_mixtures = tf.reduce_mean(
           -tf.math.top_k(-xent_reshaped, k=2)[0], axis=0)
       tf.summary.scalar("difference_top_2",

From e9e72ba346b44662c98f4e091af3d3a7d5abbdc9 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 8 Jul 2019 14:07:18 -0700
Subject: [PATCH 2178/2720] Fix T2T breaking at HEAD.

PiperOrigin-RevId: 257052138
---
 tensor2tensor/layers/bayes_test.py | 163 +++++++++++++++--------------
 1 file changed, 82 insertions(+), 81 deletions(-)

diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 473216318..e528abbdd 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -32,87 +32,88 @@
 
 class BayesTest(parameterized.TestCase, tf.test.TestCase):
 
-  @parameterized.parameters(
-      {"layer": bayes.Conv2DFlipout,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "zeros",
-       "all_close": True},
-      {"layer": bayes.Conv2DFlipout,
-       "kernel_initializer": "trainable_normal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"layer": bayes.Conv2DFlipout,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "trainable_normal",
-       "all_close": False},
-      {"layer": bayes.Conv2DReparameterization,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "zeros",
-       "all_close": True},
-      {"layer": bayes.Conv2DReparameterization,
-       "kernel_initializer": "trainable_normal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"layer": bayes.Conv2DReparameterization,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "trainable_normal",
-       "all_close": False},
-      {"layer": bayes.Conv2DVariationalDropout,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "zeros",
-       "all_close": True},
-      {"layer": bayes.Conv2DVariationalDropout,
-       "kernel_initializer": "trainable_normal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"layer": bayes.Conv2DVariationalDropout,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "trainable_normal",
-       "all_close": False},
-  )
-  @test_utils.run_in_graph_and_eager_modes
-  def testConv2DKernel(self,
-                       layer,
-                       kernel_initializer,
-                       bias_initializer,
-                       all_close):
-    tf.keras.backend.set_learning_phase(1)  # training time
-    inputs = tf.to_float(np.random.rand(5, 4, 4, 12))
-    model = layer(4,
-                  kernel_size=2,
-                  kernel_initializer=kernel_initializer,
-                  bias_initializer=bias_initializer,
-                  activation=tf.nn.relu)
-    outputs1 = model(inputs)
-    outputs2 = model(inputs)
-    self.evaluate(tf.global_variables_initializer())
-    res1, res2 = self.evaluate([outputs1, outputs2])
-    self.assertEqual(res1.shape, (5, 3, 3, 4))
-    self.assertAllGreaterEqual(res1, 0.)
-    if all_close:
-      self.assertAllClose(res1, res2)
-    else:
-      self.assertNotAllClose(res1, res2)
-    model.get_config()
-
-  @parameterized.parameters(
-      {"layer": bayes.Conv2DFlipout},
-      {"layer": bayes.Conv2DReparameterization},
-      {"layer": bayes.Conv2DVariationalDropout},
-  )
-  @test_utils.run_in_graph_and_eager_modes()
-  def testConv2DModel(self, layer):
-    inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
-    model = tf.keras.Sequential([
-        layer(3, kernel_size=2, padding="SAME", activation=tf.nn.relu),
-        tf.keras.layers.Flatten(),
-        tf.keras.layers.Dense(2, activation=None),
-    ])
-    outputs = model(inputs, training=True)
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(outputs)
-    self.assertEqual(res.shape, (3, 2))
-    self.assertLen(model.losses, 1)
+    # TODO(trandustin): Re-enable tests.
+#   @parameterized.parameters(
+#       {"layer": bayes.Conv2DFlipout,
+#        "kernel_initializer": "zeros",
+#        "bias_initializer": "zeros",
+#        "all_close": True},
+#       {"layer": bayes.Conv2DFlipout,
+#        "kernel_initializer": "trainable_normal",
+#        "bias_initializer": "zeros",
+#        "all_close": False},
+#       {"layer": bayes.Conv2DFlipout,
+#        "kernel_initializer": "zeros",
+#        "bias_initializer": "trainable_normal",
+#        "all_close": False},
+#       {"layer": bayes.Conv2DReparameterization,
+#        "kernel_initializer": "zeros",
+#        "bias_initializer": "zeros",
+#        "all_close": True},
+#       {"layer": bayes.Conv2DReparameterization,
+#        "kernel_initializer": "trainable_normal",
+#        "bias_initializer": "zeros",
+#        "all_close": False},
+#       {"layer": bayes.Conv2DReparameterization,
+#        "kernel_initializer": "zeros",
+#        "bias_initializer": "trainable_normal",
+#        "all_close": False},
+#       {"layer": bayes.Conv2DVariationalDropout,
+#        "kernel_initializer": "zeros",
+#        "bias_initializer": "zeros",
+#        "all_close": True},
+#       {"layer": bayes.Conv2DVariationalDropout,
+#        "kernel_initializer": "trainable_normal",
+#        "bias_initializer": "zeros",
+#        "all_close": False},
+#       {"layer": bayes.Conv2DVariationalDropout,
+#        "kernel_initializer": "zeros",
+#        "bias_initializer": "trainable_normal",
+#        "all_close": False},
+#   )
+#   @test_utils.run_in_graph_and_eager_modes
+#   def testConv2DKernel(self,
+#                        layer,
+#                        kernel_initializer,
+#                        bias_initializer,
+#                        all_close):
+#     tf.keras.backend.set_learning_phase(1)  # training time
+#     inputs = tf.to_float(np.random.rand(5, 4, 4, 12))
+#     model = layer(4,
+#                   kernel_size=2,
+#                   kernel_initializer=kernel_initializer,
+#                   bias_initializer=bias_initializer,
+#                   activation=tf.nn.relu)
+#     outputs1 = model(inputs)
+#     outputs2 = model(inputs)
+#     self.evaluate(tf.global_variables_initializer())
+#     res1, res2 = self.evaluate([outputs1, outputs2])
+#     self.assertEqual(res1.shape, (5, 3, 3, 4))
+#     self.assertAllGreaterEqual(res1, 0.)
+#     if all_close:
+#       self.assertAllClose(res1, res2)
+#     else:
+#       self.assertNotAllClose(res1, res2)
+#     model.get_config()
+#
+#   @parameterized.parameters(
+#       {"layer": bayes.Conv2DFlipout},
+#       {"layer": bayes.Conv2DReparameterization},
+#       {"layer": bayes.Conv2DVariationalDropout},
+#   )
+#   @test_utils.run_in_graph_and_eager_modes()
+#   def testConv2DModel(self, layer):
+#     inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
+#     model = tf.keras.Sequential([
+#         layer(3, kernel_size=2, padding="SAME", activation=tf.nn.relu),
+#         tf.keras.layers.Flatten(),
+#         tf.keras.layers.Dense(2, activation=None),
+#     ])
+#     outputs = model(inputs, training=True)
+#     self.evaluate(tf.global_variables_initializer())
+#     res = self.evaluate(outputs)
+#     self.assertEqual(res.shape, (3, 2))
+#     self.assertLen(model.losses, 1)
 
   @test_utils.run_in_graph_and_eager_modes
   def testTrainableNormalStddevConstraint(self):

From e72023412d80a5e61dc2cafc4d56f4839239918f Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 8 Jul 2019 17:23:49 -0700
Subject: [PATCH 2179/2720] Add SparseGaussianProcess layer.

requested by frederick.j.hoffman@gmail.com (e-mail). Like the math, we implement a sparse GP by re-using all of the GP layer's computation. :-)

PiperOrigin-RevId: 257090199
---
 tensor2tensor/layers/bayes.py      | 140 +++++++++++++++++++++++++++++
 tensor2tensor/layers/bayes_test.py |  25 ++++++
 2 files changed, 165 insertions(+)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index aee45bfad..0a1066e56 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -1225,6 +1225,146 @@ def get_config(self):
     return dict(list(base_config.items()) + list(config.items()))
 
 
+@add_weight
+class SparseGaussianProcess(GaussianProcess):
+  r"""Gaussian process layer with inducing input and output variables.
+
+  The layer represents a distribution over functions, where a
+  stochastic forward pass appears as
+
+  ```none
+  f ~ GP(f | inducing_inputs, inducing_outputs; mean_fn, covariance_fn)
+  outputs = f(inputs)
+  ```
+
+  The arguments `inducing_inputs` and `inducing_outputs`
+  capture data that the GP "memorizes", i.e., it forms a posterior predictive
+  distribution. Typically in a variational inference scheme (and by default),
+  the inducing outputs are normally distributed with learnable location and
+  scale parameters, and the inducing inputs are learnable parameters.
+
+  Given a call to `inputs` with these defaults, an equivalent formulation in
+  terms of function outputs is
+
+  ```none
+  inducing_outputs ~ Normal(inducing_outputs | mean, stddev)
+  outputs ~ \prod_{unit=1}^{units} MultivariateNormal(output[:, unit] |
+      mean = mean_fn(inputs) + Knm Kmm^{-1} (inducing_outputs[:, unit]-mean),
+      covariance = Knn - Knm Kmm^{-1} Kmn)
+  ```
+
+  where Knm is the covariance function evaluated between all `inputs` and
+  `inducing_inputs`; Knn is between all `inputs`; Kmm is between all
+  `inducing_inputs`; and mean is the mean function evaluated on
+  `inducing_inputs`. The multivariate normal is correlated across input
+  dimensions and is independent across output dimensions.
+
+  #### Examples
+
+  We demonstrate a three-layer deep GP with variational inference (Salimbeni and
+  Deisenroth, 2017; Damianou and Lawrence, 2013). The code snippet mirrors
+  Figure 5 of Bayesian Layers. We apply it for regression given batches of
+  spatial inputs and vector-valued outputs. We flatten inputs to use the
+  default squared exponential kernel; this naturally extends to pass in a
+  more sophisticated kernel function.
+
+  ```python
+  from tensor2tensor.layers import bayes
+
+  batch_size = 256
+  dataset_size = 10000
+  features, labels = load_spatial_data(batch_size)
+
+  model = tf.keras.Sequential([
+    tf.keras.layers.Flatten(),
+    layers.SparseGaussianProcess(256, num_inducing=512),
+    layers.SparseGaussianProcess(256, num_inducing=512),
+    layers.SparseGaussianProcess(10, num_inducing=512),
+  ])
+  predictions = model(features)
+  nll = tf.losses.mean_squared_error(labels=labels, predictions=predictions)
+  kl = sum(model.losses) / dataset_size
+  loss = nll + kl
+  train_op = tf.train.AdamOptimizer().minimize(loss)
+  ```
+  """
+
+  def __init__(
+      self,
+      units,
+      num_inducing,
+      mean_fn=Zeros(),
+      covariance_fn=ExponentiatedQuadratic(variance=1., lengthscale=1.),
+      inducing_inputs_initializer='random_normal',
+      inducing_outputs_initializer='trainable_normal',
+      inducing_inputs_regularizer=None,
+      inducing_outputs_regularizer='normal_kl_divergence',
+      inducing_inputs_constraint=None,
+      inducing_outputs_constraint=None,
+      **kwargs):
+    """Constructs layer.
+
+    Args:
+      units: integer, dimensionality of layer.
+      num_inducing: integer, number of inducing points for the approximation.
+      mean_fn: Mean function, a callable taking an inputs Tensor of shape
+        [batch, ...] and returning a Tensor of shape [batch].
+      covariance_fn: Covariance function, a callable taking two input Tensors
+        of shape [batch_x1, ...] and [batch_x2, ...] respectively, and returning
+        a positive semi-definite matrix of shape [batch_x1, batch_x2].
+      inducing_inputs_initializer: Initializer for the inducing inputs.
+      inducing_outputs_initializer: Initializer for the inducing outputs.
+      inducing_inputs_regularizer: Regularizer function applied to the inducing
+        inputs.
+      inducing_outputs_regularizer: Regularizer function applied to the inducing
+        outputs.
+      inducing_inputs_constraint: Constraint function applied to the inducing
+        inputs.
+      inducing_outputs_constraint: Constraint function applied to the inducing
+        outputs.
+      **kwargs: kwargs passed to parent class.
+    """
+    super(SparseGaussianProcess, self).__init__(
+        units=units,
+        mean_fn=mean_fn,
+        covariance_fn=covariance_fn,
+        conditional_inputs=None,
+        conditional_outputs=None,
+        **kwargs)
+    self.num_inducing = num_inducing
+    self.inducing_inputs_initializer = initializers.get(
+        inducing_inputs_initializer)
+    self.inducing_outputs_initializer = initializers.get(
+        inducing_outputs_initializer)
+    self.inducing_inputs_regularizer = regularizers.get(
+        inducing_inputs_regularizer)
+    self.inducing_outputs_regularizer = regularizers.get(
+        inducing_outputs_regularizer)
+    self.inducing_inputs_constraint = constraints.get(
+        inducing_inputs_constraint)
+    self.inducing_outputs_constraint = constraints.get(
+        inducing_outputs_constraint)
+
+  def build(self, input_shape=None):
+    input_shape = tf.TensorShape(input_shape)
+    input_dim = input_shape[-1]
+    if isinstance(input_dim, tf.Dimension):
+      input_dim = input_dim.value
+    self.conditional_inputs = self.add_weight(
+        shape=(self.num_inducing, input_dim),
+        name='inducing_inputs',
+        initializer=self.inducing_inputs_initializer,
+        regularizer=self.inducing_inputs_regularizer,
+        constraint=self.inducing_inputs_constraint)
+    self.conditional_outputs = self.add_weight(
+        shape=(self.num_inducing, self.units),
+        name='inducing_outputs',
+        initializer=self.inducing_outputs_initializer,
+        regularizer=self.inducing_outputs_regularizer,
+        constraint=self.inducing_outputs_constraint)
+    super(SparseGaussianProcess, self).build(input_shape)
+
+
 class BayesianLinearModel(tf.keras.Model):
   r"""Bayesian linear model with standard normal prior over its coefficients.
 
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index e528abbdd..5c2b56e00 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -424,6 +424,31 @@ def testGaussianProcessPrior(self):
     self.assertLessEqual(log_prob_val, 0.)
     self.assertEqual(outputs_val.shape, (batch_size, output_dim))
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testSparseGaussianProcess(self):
+    dataset_size = 10
+    batch_size = 3
+    input_dim = 4
+    output_dim = 5
+    features = tf.to_float(np.random.rand(batch_size, input_dim))
+    labels = tf.to_float(np.random.rand(batch_size, output_dim))
+    model = bayes.SparseGaussianProcess(output_dim, num_inducing=2)
+    with tf.GradientTape() as tape:
+      predictions = model(features)
+      nll = -tf.reduce_mean(predictions.distribution.log_prob(labels))
+      kl = sum(model.losses) / dataset_size
+      loss = nll + kl
+
+    self.evaluate(tf.global_variables_initializer())
+    grads = tape.gradient(nll, model.variables)
+    for grad in grads:
+      self.assertIsNotNone(grad)
+
+    loss_val, predictions_val = self.evaluate([loss, predictions])
+    self.assertEqual(loss_val.shape, ())
+    self.assertGreaterEqual(loss_val, 0.)
+    self.assertEqual(predictions_val.shape, (batch_size, output_dim))
+
   @parameterized.parameters(
       {"lstm_cell": bayes.LSTMCellFlipout,
        "kernel_initializer": "zeros",

From 33d8e96f815d19f4f516e6677c4183f91769f3c3 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 9 Jul 2019 10:47:52 -0700
Subject: [PATCH 2180/2720] Implement an env for hyperparameter tuning during
 training

PiperOrigin-RevId: 257226759
---
 tensor2tensor/trax/rlax/envs/online_tune.py   | 142 ++++++++++++++++++
 .../trax/rlax/envs/online_tune_test.py        | 105 +++++++++++++
 tensor2tensor/trax/trax.py                    |  27 ++--
 3 files changed, 265 insertions(+), 9 deletions(-)
 create mode 100644 tensor2tensor/trax/rlax/envs/online_tune.py
 create mode 100644 tensor2tensor/trax/rlax/envs/online_tune_test.py

diff --git a/tensor2tensor/trax/rlax/envs/online_tune.py b/tensor2tensor/trax/rlax/envs/online_tune.py
new file mode 100644
index 000000000..e4944bcbd
--- /dev/null
+++ b/tensor2tensor/trax/rlax/envs/online_tune.py
@@ -0,0 +1,142 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""An environment for tuning model hyperparameters during training."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+
+import gym
+
+from tensor2tensor.trax import inputs as trax_inputs
+from tensor2tensor.trax import optimizers as trax_opt
+from tensor2tensor.trax import trax
+from tensorflow.io import gfile
+
+
+class OnlineTune(object):
+  """An environment for tuning model hyperparameters during training.
+
+  A rollout is one instance of training a specific model on a specific problem.
+  Observations are the values of some evaluation metric. Actions control
+  hyperparameter changes during training. Reward is the change of the evaluation
+  metric. One environment step corresponds to a fixed number of training steps.
+
+  For now we only support tuning the learning rate.
+  """
+
+  # Chosen so that the opposite actions cancel each other out, so random walk
+  # has a median of 1.
+  DEFAULT_ACTION_MULTIPLIERS = [1.0 / 1.5, 1.0 / 1.25, 1.0, 1.25, 1.5]
+
+  def __init__(self,
+               model,
+               output_dir,
+               trainer_class=trax.Trainer,
+               loss_fn=trax.loss,
+               optimizer=trax_opt.SM3,
+               inputs=trax_inputs.inputs,
+               action_multipliers=None,
+               history_mode="eval",
+               metric="metrics/accuracy",
+               train_steps=100,
+               eval_steps=10,
+               env_steps=100,
+               start_lr=0.001):
+    if action_multipliers is None:
+      action_multipliers = self.DEFAULT_ACTION_MULTIPLIERS
+    self._model = model
+    self._trainer_fn = functools.partial(
+        trainer_class,
+        model=model,
+        loss_fn=loss_fn,
+        optimizer=optimizer,
+        lr_schedule=(lambda history: lambda step: self._current_lr),
+        inputs=inputs)
+    self._action_multipliers = action_multipliers
+    self._history_mode = history_mode
+    self._metric = metric
+    self._train_steps = train_steps
+    self._eval_steps = eval_steps
+    self._env_steps = env_steps
+    self._start_lr = start_lr
+    self._trainer = None
+
+    self.output_dir = output_dir
+    # Action is an index in self._action_multipliers.
+    self.action_space = gym.spaces.Discrete(len(self._action_multipliers))
+    # Observation is the value of the metric specified in self._metric.
+    self.observation_space = gym.spaces.Box(
+        low=float("-inf"), high=float("+inf"), shape=())
+
+  def _remove_output_dir(self):
+    if gfile.exists(self.output_dir):
+      gfile.rmtree(self.output_dir)
+
+  @property
+  def _current_metric_value(self):
+    metric_sequence = self._trainer.state.history.get(self._history_mode,
+                                                      self._metric)
+    assert metric_sequence
+    (_, metric_value) = metric_sequence[-1]
+    return metric_value
+
+  @property
+  def trainer(self):
+    if self._trainer is None:
+      raise ValueError("The environment has to be reset first.")
+    return self._trainer
+
+  def reset(self):
+    # TODO(pkozakowski): Don't erase the data. Will be done in the next CL.
+    self._remove_output_dir()
+    gfile.makedirs(self.output_dir)
+
+    self._current_lr = self._start_lr
+    self._step = 0
+    self._trainer = self._trainer_fn(output_dir=self.output_dir)
+    self._trainer.evaluate(self._eval_steps)
+    return self._current_metric_value
+
+  def step(self, action):
+    """Step the environment.
+
+    One environment step corresponds to self.train_steps training steps.
+
+    Args:
+      action: (int) Action to take. An index in self.action_multipliers.
+
+    Returns:
+      Tuple (observation, reward, done, info). observation is a singleton vector
+        with the current value of the metric. reward is the difference in the
+        metric since the last step. done is set after reaching self.env_steps
+        environment steps. info is an empty dict.
+    """
+    self._current_lr *= self._action_multipliers[action]
+    self._trainer.update_learning_rate(force_jit=True)
+    last_metric_value = self._current_metric_value
+    self._trainer.train_epoch(self._train_steps, self._eval_steps)
+    self._step += 1
+    current_metric_value = self._current_metric_value
+    observation = current_metric_value
+    reward = current_metric_value - last_metric_value
+    done = self._step == self._env_steps
+    return (observation, reward, done, {})
+
+  def close(self):
+    self._remove_output_dir()
diff --git a/tensor2tensor/trax/rlax/envs/online_tune_test.py b/tensor2tensor/trax/rlax/envs/online_tune_test.py
new file mode 100644
index 000000000..70eb507f7
--- /dev/null
+++ b/tensor2tensor/trax/rlax/envs/online_tune_test.py
@@ -0,0 +1,105 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.rlax.online_tune."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+
+import numpy as np
+
+from tensor2tensor.trax import inputs as trax_inputs
+from tensor2tensor.trax import models
+from tensor2tensor.trax import trax
+from tensor2tensor.trax.rlax.envs import online_tune
+from tensorflow import test
+
+HISTORY_MODE = "eval"
+METRIC = "metrics/accuracy"
+
+
+class MockTrainer(trax.Trainer):
+
+  def __init__(self, metrics_to_report, *args, **kwargs):
+    super(MockTrainer, self).__init__(*args, **kwargs)
+    self.learning_rates = []
+    # Copy the list so we can modify it later.
+    self.metrics_to_report = metrics_to_report[:]
+
+  def train_epoch(self, epoch_steps, eval_steps):
+    del epoch_steps
+    self.learning_rates.append(self.learning_rate)
+    self.evaluate(eval_steps)
+
+  def evaluate(self, eval_steps):
+    del eval_steps
+    self.state.history.append(
+        mode=HISTORY_MODE,
+        metric=METRIC,
+        step=self.step,
+        value=self.metrics_to_report.pop(0))
+
+
+class OnlineTuneTest(test.TestCase):
+
+  def test_communicates_with_trainer(self):
+    action_multipliers = [0.8, 1.0, 1.25]
+    metrics_to_report = [0.1, 0.5, 0.8, 0.9]
+    actions_to_take = [0, 1, 2]
+    expected_observations = metrics_to_report
+    # Metric difference in consecutive timesteps.
+    expected_rewards = [0.4, 0.3, 0.1]
+    expected_dones = [False, False, True]
+    expected_learning_rates = [0.0008, 0.0008, 0.001]
+
+    env = online_tune.OnlineTune(
+        trainer_class=functools.partial(MockTrainer, metrics_to_report),
+        model=functools.partial(
+            models.MLP, n_hidden_layers=0, n_output_classes=1),
+        inputs=functools.partial(
+            trax_inputs.random_inputs,
+            input_shape=(1, 1),
+            input_dtype=np.float32,
+            output_shape=(1, 1),
+            output_dtype=np.float32),
+        output_dir=self.get_temp_dir(),
+        action_multipliers=action_multipliers,
+        history_mode=HISTORY_MODE,
+        metric=METRIC,
+        train_steps=1,
+        eval_steps=1,
+        env_steps=len(actions_to_take))
+
+    actual_observations = [env.reset()]
+    actual_rewards = []
+    actual_dones = []
+    for action in actions_to_take:
+      (observation, reward, done, _) = env.step(action)
+      actual_observations.append(observation)
+      actual_rewards.append(reward)
+      actual_dones.append(done)
+
+    np.testing.assert_allclose(actual_observations, expected_observations)
+    np.testing.assert_allclose(actual_rewards, expected_rewards)
+    self.assertEqual(actual_dones, expected_dones)
+    np.testing.assert_allclose(env.trainer.learning_rates,
+                               expected_learning_rates)
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index a6b94c892..d3adb3e76 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -539,6 +539,14 @@ def n_devices(self):
   def state(self):
     return State(params=self._opt_state, step=self._step, history=self._history)
 
+  @property
+  def learning_rate(self):
+    # TODO(lukaszkaiser): it makes no sense to use an accelerator (e.g. TPU)
+    # in op-by-op mode just to compute the learning rate. However, there
+    # should be a cleaner approach that forceably swapping out the backend.
+    with backend.use_backend("numpy"):
+      return self._lr_fn(self._step)
+
   def save_gin(self):
     _save_gin(self._output_dir, self._train_sw)
 
@@ -572,8 +580,7 @@ def train_epoch(self, epoch_steps, eval_steps):
         # in op-by-op mode just to compute the learning rate. However, there
         # should be a cleaner approach that forceably swapping out the backend.
         with backend.use_backend("numpy"):
-          self._train_sw.scalar("training/learning rate",
-                                self._lr_fn(self._step), step=self._step)
+          self._train_sw.scalar("training/learning rate", self.learning_rate)
 
     # Timer
     epoch_time = time.time() - start_time
@@ -607,10 +614,11 @@ def evaluate(self, eval_steps):
         eval_sw=self._eval_sw,
         history=self._history)
 
-  def update_learning_rate(self):
+  def update_learning_rate(self, force_jit=False):
     old_lr_fn = self._lr_fn
     self._lr_fn = self._lr_schedule(self._history)
-    if self._lr_fn != old_lr_fn:  # For performance only jit if it's changed.
+    # For performance only jit if it's changed or we force it.
+    if self._lr_fn != old_lr_fn or force_jit:
       opt = self._optimizer(self._lr_fn)
       self._jit_update_fn = _jit_update_fn(
           self._model_train, self._loss_fn, opt, self._n_devices)
@@ -668,14 +676,15 @@ def evaluate(self, eval_steps):
         eval_sw=self._eval_sw,
         history=self._history)
 
-  def update_learning_rate(self):
+  def update_learning_rate(self, force_jit=False):
     old_lr_fn = self._lr_fn
     self._lr_fn = self._lr_schedule(self._history)
-    if self._lr_fn != old_lr_fn:
+    if self._lr_fn != old_lr_fn or force_jit:
       raise NotImplementedError(
-          "Loss function changed. Garbage collection for jitted functions is "
-          "not implemented in jax, so global accelerator memory allocated by "
-          "the jitted update function with the old loss cannot be reclaimed.")
+          "Loss function changed or jitting was requested. Garbage collection "
+          "for jitted functions is not implemented in jax, so global "
+          "accelerator memory allocated by the jitted update function with the "
+          "old loss cannot be reclaimed.")
 
   def save_computation_graphs(self, save_backward_graph):
     # TODO(kitaev): implement saving graphs while making sure that no op-by-op

From dd4cf87e140cb4cdaaf859f31b81791828d115bf Mon Sep 17 00:00:00 2001
From: Norman Mu <normanmu@google.com>
Date: Tue, 9 Jul 2019 11:31:41 -0700
Subject: [PATCH 2181/2720] Fix WideResNet definition in trax/models

PiperOrigin-RevId: 257236969
---
 .../trax/configs/wide_resnet_cifar10_8gb.gin      |  2 +-
 tensor2tensor/trax/models/resnet.py               | 15 +++++++--------
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
index ff3b2aa43..74eabe53f 100644
--- a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
+++ b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
@@ -32,7 +32,7 @@ shuffle_and_batch_data.preprocess_fun=@trax.inputs.cifar10_no_augmentation_prepr
 
 # Parameters for WideResnet:
 # ==============================================================================
-WideResnet.d_hidden = 64
+WideResnet.widen_factor = 10
 WideResnet.n_blocks = 3
 WideResnet.n_output_classes = 10
 
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index 7ed36508d..beeaa5780 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -132,13 +132,12 @@ def WideResnetGroup(n, channels, strides=(1, 1)):
   ]
 
 
-def WideResnet(n_blocks=3, d_hidden=64, n_output_classes=10,
-               mode='train'):
+def WideResnet(n_blocks=3, widen_factor=1, n_output_classes=10, mode='train'):
   """WideResnet from https://arxiv.org/pdf/1605.07146.pdf.
 
   Args:
-    n_blocks: int, number of blocks in a group.
-    d_hidden: Dimensionality of the first hidden layer (multiplied later).
+    n_blocks: int, number of blocks in a group. total layers = 6n + 4.
+    widen_factor: int, widening factor of each group. k=1 is vanilla resnet.
     n_output_classes: int, number of distinct output classes.
     mode: Whether we are training or evaluating or doing inference.
 
@@ -148,10 +147,10 @@ def WideResnet(n_blocks=3, d_hidden=64, n_output_classes=10,
   del mode
   return tl.Model(
       tl.ToFloat(),
-      tl.Conv(d_hidden, (3, 3), padding='SAME'),
-      WideResnetGroup(n_blocks, d_hidden),
-      WideResnetGroup(n_blocks, d_hidden * 2, (2, 2)),
-      WideResnetGroup(n_blocks, d_hidden * 4, (2, 2)),
+      tl.Conv(16, (3, 3), padding='SAME'),
+      WideResnetGroup(n_blocks, 16 * widen_factor),
+      WideResnetGroup(n_blocks, 32 * widen_factor, (2, 2)),
+      WideResnetGroup(n_blocks, 64 * widen_factor, (2, 2)),
       tl.BatchNorm(),
       tl.Relu(),
       tl.AvgPool(pool_size=(8, 8)),

From 2ff384a84f389d7384373441226e27d343b36118 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 9 Jul 2019 17:55:18 -0700
Subject: [PATCH 2182/2720] Make OnlineTune work with PPO

PiperOrigin-RevId: 257309870
---
 tensor2tensor/envs/env_problem.py             | 27 ++++++--
 tensor2tensor/envs/rendered_env_problem.py    |  9 +--
 tensor2tensor/trax/models/__init__.py         |  2 +
 tensor2tensor/trax/rlax/__init__.py           | 15 ++++
 .../online_tune_wide_resnet_cifar10.gin       | 40 +++++++++++
 tensor2tensor/trax/rlax/envs/__init__.py      | 19 ++++++
 .../{online_tune.py => online_tune_env.py}    | 53 ++++++++++-----
 ...e_tune_test.py => online_tune_env_test.py} | 50 +++++++++-----
 tensor2tensor/trax/rlax/ppo_main.py           | 38 +++++++++--
 .../trax/rlax/ppo_training_loop_test.py       | 68 +++++++++++++------
 10 files changed, 246 insertions(+), 75 deletions(-)
 create mode 100644 tensor2tensor/trax/rlax/__init__.py
 create mode 100644 tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
 rename tensor2tensor/trax/rlax/envs/{online_tune.py => online_tune_env.py} (77%)
 rename tensor2tensor/trax/rlax/envs/{online_tune_test.py => online_tune_env_test.py} (75%)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 930259311..d608b4fc8 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -129,7 +129,9 @@ def __init__(self,
                base_env_name=None,
                batch_size=None,
                env_wrapper_fn=None,
-               reward_range=(-np.inf, np.inf)):
+               reward_range=(-np.inf, np.inf),
+               discrete_rewards=True,
+               **env_kwargs):
     """Initializes this class by creating the envs and managing trajectories.
 
     Args:
@@ -142,6 +144,9 @@ def __init__(self,
       reward_range: (tuple(number, number)) the first element is the minimum
         reward and the second is the maximum reward, used to clip and process
         the raw reward in `process_rewards`.
+      discrete_rewards: (bool) whether to round the rewards to the nearest
+        integer.
+      **env_kwargs: (dict) Additional kwargs to pass to the environments.
     """
 
     # Call the super's ctor.
@@ -162,6 +167,9 @@ def __init__(self,
     # in `process_rewards`.
     self._reward_range = reward_range
 
+    # If set, we discretize the rewards and treat them as integers.
+    self._discrete_rewards = discrete_rewards
+
     # Initialize the environment(s).
 
     # This can either be a list of environments of len `batch_size` or this can
@@ -181,7 +189,7 @@ def __init__(self,
     self._env_wrapper_fn = env_wrapper_fn
 
     if batch_size is not None:
-      self.initialize(batch_size=batch_size)
+      self.initialize(batch_size=batch_size, **env_kwargs)
 
   @property
   def batch_size(self):
@@ -246,7 +254,7 @@ def initialize(self, **kwargs):
     assert self._reward_range is not None
     assert self._trajectories is not None
 
-  def initialize_environments(self, batch_size=1):
+  def initialize_environments(self, batch_size=1, **env_kwargs):
     """Initializes the environments and trajectories.
 
     Subclasses can override this if they don't want a default implementation
@@ -255,11 +263,14 @@ def initialize_environments(self, batch_size=1):
 
     Args:
       batch_size: (int) Number of `self.base_env_name` envs to initialize.
+      **env_kwargs: (dict) Kwargs to pass to gym.make.
     """
     assert batch_size >= 1
     self._batch_size = batch_size
 
-    self._envs = [gym.make(self.base_env_name) for _ in range(batch_size)]
+    self._envs = [
+        gym.make(self.base_env_name, **env_kwargs) for _ in range(batch_size)
+    ]
     if self._env_wrapper_fn is not None:
       self._envs = list(map(self._env_wrapper_fn, self._envs))
 
@@ -351,7 +362,7 @@ def is_reward_range_finite(self):
     return (min_reward != -np.inf) and (max_reward != np.inf)
 
   def process_rewards(self, rewards):
-    """Clips, rounds, and changes to integer type.
+    """Clips the rewards, optionally rounds them and casts to integer.
 
     Args:
       rewards: numpy array of raw (float) rewards.
@@ -364,8 +375,10 @@ def process_rewards(self, rewards):
 
     # Clips at min and max reward.
     rewards = np.clip(rewards, min_reward, max_reward)
-    # Round to (nearest) int and convert to integral type.
-    rewards = np.around(rewards, decimals=0).astype(np.int64)
+
+    if self._discrete_rewards:
+      # Round to (nearest) int and convert to integral type.
+      rewards = np.around(rewards, decimals=0).astype(np.int64)
     return rewards
 
   @property
diff --git a/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
index 5760ae47a..dba768495 100644
--- a/tensor2tensor/envs/rendered_env_problem.py
+++ b/tensor2tensor/envs/rendered_env_problem.py
@@ -48,14 +48,9 @@ class RenderedEnvProblem(env_problem.EnvProblem, video_utils.VideoProblem):
   `RenderedEnvProblem`, `EnvProblem`, `Env`, `VideoProblem`, `Problem`
   """
 
-  def __init__(self,
-               base_env_name=None,
-               batch_size=None,
-               env_wrapper_fn=None,
-               reward_range=(-np.inf, np.inf)):
+  def __init__(self, *args, **kwargs):
     """Initialize by calling both parents' constructors."""
-    env_problem.EnvProblem.__init__(self, base_env_name, batch_size,
-                                    env_wrapper_fn, reward_range)
+    env_problem.EnvProblem.__init__(self, *args, **kwargs)
     video_utils.VideoProblem.__init__(self)
 
   def initialize_environments(self, batch_size=1):
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index 22db0422e..8b363f27b 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -20,6 +20,7 @@
 
 import gin
 
+from tensor2tensor.trax.models import atari_cnn
 from tensor2tensor.trax.models import mlp
 from tensor2tensor.trax.models import neural_gpu
 from tensor2tensor.trax.models import resnet
@@ -35,6 +36,7 @@ def model_configure(*args, **kwargs):
 
 
 # pylint: disable=invalid-name
+AtariCnn = model_configure(atari_cnn.AtariCnn)
 MLP = model_configure(mlp.MLP)
 NeuralGPU = model_configure(neural_gpu.NeuralGPU)
 PositionLookupTransformerLM = model_configure(
diff --git a/tensor2tensor/trax/rlax/__init__.py b/tensor2tensor/trax/rlax/__init__.py
new file mode 100644
index 000000000..4872e5d5d
--- /dev/null
+++ b/tensor2tensor/trax/rlax/__init__.py
@@ -0,0 +1,15 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
diff --git a/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
new file mode 100644
index 000000000..1d84a1b5f
--- /dev/null
+++ b/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
@@ -0,0 +1,40 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.rlax.envs
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size = 32
+batch_fun.bucket_length = 32
+batch_fun.buckets = None
+batch_fun.eval_batch_size = 32
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 'cifar10'
+
+# Parameters for Momentum:
+# ==============================================================================
+Momentum.mass = 0.9
+
+# Parameters for shuffle_and_batch_data:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun = @trax.inputs.cifar10_no_augmentation_preprocess
+
+# Parameters for WideResnet:
+# ==============================================================================
+WideResnet.d_hidden = 64
+WideResnet.n_blocks = 3
+WideResnet.n_output_classes = 10
+
+# Parameters for OnlineTune:
+# ==============================================================================
+OnlineTuneEnv.inputs = @trax.inputs.inputs
+OnlineTuneEnv.model = @trax.models.WideResnet
+OnlineTuneEnv.optimizer = @trax.optimizers.Momentum
+OnlineTuneEnv.start_lr = 0.01
+OnlineTuneEnv.train_steps = 500
+OnlineTuneEnv.eval_steps = 100
+OnlineTuneEnv.env_steps = 100
diff --git a/tensor2tensor/trax/rlax/envs/__init__.py b/tensor2tensor/trax/rlax/envs/__init__.py
index 4872e5d5d..fa029618a 100644
--- a/tensor2tensor/trax/rlax/envs/__init__.py
+++ b/tensor2tensor/trax/rlax/envs/__init__.py
@@ -13,3 +13,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Environments defined in RLAX."""
+
+import gin
+from gym.envs.registration import register
+
+from tensor2tensor.trax.rlax.envs import online_tune_env
+
+
+# Ginify and register in gym.
+def configure_and_register_env(env_class):
+  register(
+      id="{}-v0".format(env_class.__name__),
+      entry_point="tensor2tensor.trax.rlax.envs:{}".format(env_class.__name__),
+  )
+  return gin.external_configurable(env_class, module="trax.rlax.envs")
+
+
+# pylint: disable=invalid-name
+OnlineTuneEnv = configure_and_register_env(online_tune_env.OnlineTuneEnv)
diff --git a/tensor2tensor/trax/rlax/envs/online_tune.py b/tensor2tensor/trax/rlax/envs/online_tune_env.py
similarity index 77%
rename from tensor2tensor/trax/rlax/envs/online_tune.py
rename to tensor2tensor/trax/rlax/envs/online_tune_env.py
index e4944bcbd..b29e9f46c 100644
--- a/tensor2tensor/trax/rlax/envs/online_tune.py
+++ b/tensor2tensor/trax/rlax/envs/online_tune_env.py
@@ -20,8 +20,10 @@
 from __future__ import print_function
 
 import functools
+import os
 
 import gym
+import numpy as np
 
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import optimizers as trax_opt
@@ -29,7 +31,7 @@
 from tensorflow.io import gfile
 
 
-class OnlineTune(object):
+class OnlineTuneEnv(gym.Env):
   """An environment for tuning model hyperparameters during training.
 
   A rollout is one instance of training a specific model on a specific problem.
@@ -77,16 +79,40 @@ def __init__(self,
     self._start_lr = start_lr
     self._trainer = None
 
-    self.output_dir = output_dir
+    self._output_dir = output_dir
+    gfile.makedirs(self._output_dir)
     # Action is an index in self._action_multipliers.
     self.action_space = gym.spaces.Discrete(len(self._action_multipliers))
-    # Observation is the value of the metric specified in self._metric.
+    # Observation is a singleton vector with the value of the metric specified
+    # in self._metric.
     self.observation_space = gym.spaces.Box(
-        low=float("-inf"), high=float("+inf"), shape=())
+        low=float("-inf"), high=float("+inf"), shape=(1,))
 
-  def _remove_output_dir(self):
-    if gfile.exists(self.output_dir):
-      gfile.rmtree(self.output_dir)
+  @property
+  def _next_trajectory_dir(self):
+    """Assigns a new output dir for a trajectory under self._output_dir.
+
+    Directory names are consecutive integers starting from zero. New directory
+    index is assigned as the maximum of past indices plus one. Directories that
+    are not integers are ignored.
+
+    Returns:
+      A path of the new directory.
+    """
+    trajectory_dirs = gfile.listdir(self._output_dir)
+
+    def int_or_none(s):
+      try:
+        return int(s)
+      except TypeError:
+        return None
+
+    past_trajectory_ids = [
+        trajectory_id for trajectory_id in map(int_or_none, trajectory_dirs)
+        if trajectory_id is not None]
+    next_trajectory_id = max([-1] + past_trajectory_ids) + 1
+
+    return os.path.join(self._output_dir, str(next_trajectory_id))
 
   @property
   def _current_metric_value(self):
@@ -103,15 +129,11 @@ def trainer(self):
     return self._trainer
 
   def reset(self):
-    # TODO(pkozakowski): Don't erase the data. Will be done in the next CL.
-    self._remove_output_dir()
-    gfile.makedirs(self.output_dir)
-
     self._current_lr = self._start_lr
     self._step = 0
-    self._trainer = self._trainer_fn(output_dir=self.output_dir)
+    self._trainer = self._trainer_fn(output_dir=self._next_trajectory_dir)
     self._trainer.evaluate(self._eval_steps)
-    return self._current_metric_value
+    return np.array([self._current_metric_value])
 
   def step(self, action):
     """Step the environment.
@@ -133,10 +155,7 @@ def step(self, action):
     self._trainer.train_epoch(self._train_steps, self._eval_steps)
     self._step += 1
     current_metric_value = self._current_metric_value
-    observation = current_metric_value
+    observation = np.array([current_metric_value])
     reward = current_metric_value - last_metric_value
     done = self._step == self._env_steps
     return (observation, reward, done, {})
-
-  def close(self):
-    self._remove_output_dir()
diff --git a/tensor2tensor/trax/rlax/envs/online_tune_test.py b/tensor2tensor/trax/rlax/envs/online_tune_env_test.py
similarity index 75%
rename from tensor2tensor/trax/rlax/envs/online_tune_test.py
rename to tensor2tensor/trax/rlax/envs/online_tune_env_test.py
index 70eb507f7..20afa96f1 100644
--- a/tensor2tensor/trax/rlax/envs/online_tune_test.py
+++ b/tensor2tensor/trax/rlax/envs/online_tune_env_test.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for tensor2tensor.trax.rlax.online_tune."""
+"""Tests for tensor2tensor.trax.rlax.online_tune_env."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -26,8 +26,9 @@
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import models
 from tensor2tensor.trax import trax
-from tensor2tensor.trax.rlax.envs import online_tune
+from tensor2tensor.trax.rlax.envs import online_tune_env
 from tensorflow import test
+from tensorflow.io import gfile
 
 HISTORY_MODE = "eval"
 METRIC = "metrics/accuracy"
@@ -38,8 +39,8 @@ class MockTrainer(trax.Trainer):
   def __init__(self, metrics_to_report, *args, **kwargs):
     super(MockTrainer, self).__init__(*args, **kwargs)
     self.learning_rates = []
-    # Copy the list so we can modify it later.
-    self.metrics_to_report = metrics_to_report[:]
+    # Copy the sequence to a list so we can modify it later.
+    self.metrics_to_report = list(metrics_to_report)
 
   def train_epoch(self, epoch_steps, eval_steps):
     del epoch_steps
@@ -57,17 +58,9 @@ def evaluate(self, eval_steps):
 
 class OnlineTuneTest(test.TestCase):
 
-  def test_communicates_with_trainer(self):
-    action_multipliers = [0.8, 1.0, 1.25]
-    metrics_to_report = [0.1, 0.5, 0.8, 0.9]
-    actions_to_take = [0, 1, 2]
-    expected_observations = metrics_to_report
-    # Metric difference in consecutive timesteps.
-    expected_rewards = [0.4, 0.3, 0.1]
-    expected_dones = [False, False, True]
-    expected_learning_rates = [0.0008, 0.0008, 0.001]
-
-    env = online_tune.OnlineTune(
+  @staticmethod
+  def _create_env(output_dir, metrics_to_report=(0.0,), action_multipliers=()):
+    return online_tune_env.OnlineTuneEnv(
         trainer_class=functools.partial(MockTrainer, metrics_to_report),
         model=functools.partial(
             models.MLP, n_hidden_layers=0, n_output_classes=1),
@@ -77,14 +70,28 @@ def test_communicates_with_trainer(self):
             input_dtype=np.float32,
             output_shape=(1, 1),
             output_dtype=np.float32),
-        output_dir=self.get_temp_dir(),
+        output_dir=output_dir,
         action_multipliers=action_multipliers,
         history_mode=HISTORY_MODE,
         metric=METRIC,
         train_steps=1,
         eval_steps=1,
-        env_steps=len(actions_to_take))
+        env_steps=(len(metrics_to_report) - 1))
 
+  def test_communicates_with_trainer(self):
+    action_multipliers = [0.8, 1.0, 1.25]
+    metrics_to_report = [0.1, 0.5, 0.8, 0.9]
+    actions_to_take = [0, 1, 2]
+    expected_observations = np.expand_dims(metrics_to_report, axis=1)
+    # Metric difference in consecutive timesteps.
+    expected_rewards = [0.4, 0.3, 0.1]
+    expected_dones = [False, False, True]
+    expected_learning_rates = [0.0008, 0.0008, 0.001]
+
+    env = self._create_env(
+        output_dir=self.get_temp_dir(),
+        metrics_to_report=metrics_to_report,
+        action_multipliers=action_multipliers)
     actual_observations = [env.reset()]
     actual_rewards = []
     actual_dones = []
@@ -100,6 +107,15 @@ def test_communicates_with_trainer(self):
     np.testing.assert_allclose(env.trainer.learning_rates,
                                expected_learning_rates)
 
+  def test_creates_new_trajectory_dirs(self):
+    output_dir = self.get_temp_dir()
+    env = self._create_env(output_dir=output_dir)
+    self.assertEqual(set(gfile.listdir(output_dir)), set())
+    env.reset()
+    self.assertEqual(set(gfile.listdir(output_dir)), {"0"})
+    env.reset()
+    self.assertEqual(set(gfile.listdir(output_dir)), {"0", "1"})
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 546781aaa..fc4f69c5d 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -40,10 +40,12 @@
 from __future__ import print_function
 
 import functools
+import os
 
 from absl import app
 from absl import flags
 from absl import logging
+import gin
 import jax
 from jax.config import config
 import numpy as onp
@@ -51,7 +53,8 @@
 from tensor2tensor.envs import rendered_env_problem
 from tensor2tensor.rl import gym_utils
 from tensor2tensor.trax import layers
-from tensor2tensor.trax.models import atari_cnn
+from tensor2tensor.trax import models
+from tensor2tensor.trax.rlax import envs  # pylint: disable=unused-import
 from tensor2tensor.trax.rlax import ppo
 
 
@@ -120,6 +123,10 @@
 flags.DEFINE_float("epsilon", 0.1, "Policy iteration early stopping")
 
 flags.DEFINE_string("output_dir", "", "Output dir.")
+flags.DEFINE_multi_string("config_file", None,
+                          "Configuration file with parameters (.gin).")
+flags.DEFINE_multi_string("config", None,
+                          "Configuration parameters (gin string).")
 flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
 flags.DEFINE_bool("enable_early_stopping", True,
                   "Whether to enable early stopping.")
@@ -133,6 +140,8 @@
     "checkpoint the policy.")
 flags.DEFINE_integer("len_history_for_policy", 4,
                      "How much of history to give to the policy.")
+flags.DEFINE_bool("clip_rewards", True,
+                  "Whether to clip and discretize the rewards.")
 
 
 def common_layers():
@@ -144,18 +153,23 @@ def common_layers():
 
 
 def atari_layers():
-  return [atari_cnn.AtariCnn()]
+  return [models.AtariCnn()]
 
 
-def make_env(batch_size=8):
+def make_env(batch_size=8, **env_kwargs):
   """Creates the env."""
 
+  if FLAGS.clip_rewards:
+    env_kwargs.update({"reward_range": (-1, 1), "discrete_rewards": True})
+  else:
+    env_kwargs.update({"discrete_rewards": False})
+
   # No resizing needed, so let's be on the normal EnvProblem.
   if not FLAGS.resize:  # None or False
     return env_problem.EnvProblem(
         base_env_name=FLAGS.env_problem_name,
         batch_size=batch_size,
-        reward_range=(-1, 1))
+        **env_kwargs)
 
   max_timestep = None
   try:
@@ -177,7 +191,7 @@ def make_env(batch_size=8):
       base_env_name=FLAGS.env_problem_name,
       batch_size=batch_size,
       env_wrapper_fn=wrapper_fn,
-      reward_range=(-1, 1))
+      **env_kwargs)
 
 
 def get_optimizer_fn(learning_rate):
@@ -197,11 +211,21 @@ def main(argv):
     config.update("jax_platform_name", "gpu")
 
 
+  gin_configs = FLAGS.config or []
+  gin.parse_config_files_and_bindings(FLAGS.config_file, gin_configs)
+
+  # TODO(pkozakowski): Find a better way to determine this.
+  if "OnlineTuneEnv" in FLAGS.env_problem_name:
+    # TODO(pkozakowski): Separate env output dirs by train/eval and epoch.
+    env_kwargs = {"output_dir": os.path.join(FLAGS.output_dir, "envs")}
+  else:
+    env_kwargs = {}
+
   # Make an env here.
-  env = make_env(batch_size=FLAGS.batch_size)
+  env = make_env(batch_size=FLAGS.batch_size, **env_kwargs)
   assert env
 
-  eval_env = make_env(batch_size=FLAGS.eval_batch_size)
+  eval_env = make_env(batch_size=FLAGS.eval_batch_size, **env_kwargs)
   assert eval_env
 
   def run_training_loop():
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index 5e9a17a49..05ef3a797 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -21,10 +21,18 @@
 
 import contextlib
 import functools
+import os
 import tempfile
+
+import gin
+import numpy as np
+
 from tensor2tensor.envs import env_problem
 from tensor2tensor.rl import gym_utils
+from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import layers
+from tensor2tensor.trax import models
+from tensor2tensor.trax.rlax import envs  # pylint: disable=unused-import
 from tensor2tensor.trax.rlax import ppo
 from tensorflow import test
 from tensorflow.io import gfile
@@ -47,7 +55,8 @@ def get_wrapped_env(self, name="CartPole-v0", max_episode_steps=2):
     return env_problem.EnvProblem(base_env_name=name,
                                   batch_size=1,
                                   env_wrapper_fn=wrapper_fn,
-                                  reward_range=(-1, 1))
+                                  reward_range=(-1, 1),
+                                  discrete_rewards=False)
 
   @contextlib.contextmanager
   def tmp_dir(self):
@@ -55,26 +64,45 @@ def tmp_dir(self):
     yield tmp
     gfile.rmtree(tmp)
 
-  def test_training_loop(self):
+  def _run_training_loop(self, env_name, output_dir):
+    env = self.get_wrapped_env(env_name, 2)
+    eval_env = self.get_wrapped_env(env_name, 2)
+    n_epochs = 2
+    batch_size = 2
+    # Run the training loop.
+    ppo.training_loop(
+        env=env,
+        eval_env=eval_env,
+        epochs=n_epochs,
+        policy_and_value_net_fn=functools.partial(
+            ppo.policy_and_value_net,
+            bottom_layers_fn=lambda: [layers.Dense(1)]),
+        policy_and_value_optimizer_fn=ppo.optimizer_fn,
+        batch_size=batch_size,
+        n_optimizer_steps=1,
+        output_dir=output_dir,
+        env_name=env_name,
+        random_seed=0)
+
+  def test_training_loop_cartpole(self):
+    with self.tmp_dir() as output_dir:
+      self._run_training_loop("CartPole-v0", output_dir)
+
+  def test_training_loop_onlinetune(self):
     with self.tmp_dir() as output_dir:
-      env = self.get_wrapped_env("CartPole-v0", 2)
-      eval_env = self.get_wrapped_env("CartPole-v0", 2)
-      n_epochs = 2
-      batch_size = 2
-      # Run the training loop.
-      ppo.training_loop(
-          env=env,
-          eval_env=eval_env,
-          epochs=n_epochs,
-          policy_and_value_net_fn=functools.partial(
-              ppo.policy_and_value_net,
-              bottom_layers_fn=lambda: [layers.Dense(1)]),
-          policy_and_value_optimizer_fn=ppo.optimizer_fn,
-          batch_size=batch_size,
-          n_optimizer_steps=1,
-          output_dir=output_dir,
-          env_name="CartPole-v0",
-          random_seed=0)
+      gin.bind_parameter("OnlineTuneEnv.model", functools.partial(
+          models.MLP, n_hidden_layers=0, n_output_classes=1))
+      gin.bind_parameter("OnlineTuneEnv.inputs", functools.partial(
+          trax_inputs.random_inputs,
+          input_shape=(1, 1),
+          input_dtype=np.float32,
+          output_shape=(1, 1),
+          output_dtype=np.float32))
+      gin.bind_parameter("OnlineTuneEnv.train_steps", 2)
+      gin.bind_parameter("OnlineTuneEnv.eval_steps", 2)
+      gin.bind_parameter(
+          "OnlineTuneEnv.output_dir", os.path.join(output_dir, "envs"))
+      self._run_training_loop("OnlineTuneEnv-v0", output_dir)
 
 
 if __name__ == "__main__":

From 5551789ebfe90e9a471714a5785c134dc9dd000a Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Tue, 9 Jul 2019 20:09:52 -0700
Subject: [PATCH 2183/2720] Move Gaussian process layers into another module.

The bayes.py module is becoming a bit too unwieldy to navigate. Reorganizing also allows for faster and more modular imports.

PiperOrigin-RevId: 257324288
---
 tensor2tensor/layers/bayes.py                 | 437 +----------------
 tensor2tensor/layers/bayes_test.py            |  98 +---
 tensor2tensor/layers/gaussian_process.py      | 464 ++++++++++++++++++
 tensor2tensor/layers/gaussian_process_test.py | 131 +++++
 4 files changed, 597 insertions(+), 533 deletions(-)
 create mode 100644 tensor2tensor/layers/gaussian_process.py
 create mode 100644 tensor2tensor/layers/gaussian_process_test.py

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 0a1066e56..49dfa0f0f 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Bayesian layers."""
+"""Bayesian neural network layers."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -989,441 +989,6 @@ def call(self, inputs, states, training=None):
     return h, [h, c]
 
 
-class Zeros(object):
-  """Function returning zeros tensor of same shape excluding the last dim."""
-
-  def __call__(self, inputs):
-    return tf.zeros(tf.shape(inputs)[:-1], inputs.dtype)
-
-  def get_config(self):
-    return {}
-
-
-class ExponentiatedQuadratic(object):
-  """Exponentiated quadratic kernel."""
-
-  def __init__(self, variance, lengthscale):
-    self.variance = variance
-    self.lengthscale = lengthscale
-
-  def __call__(self, x1, x2):
-    """Computes exponentiated quadratic over all pairs of inputs.
-
-    Args:
-      x1: Tensor of shape [batch_x1, ...]. Slices along the batch axis denote an
-        individual input to be passed to the kernel. It is computed pairwise
-        with each input sliced from x2.
-      x2: Tensor of shape [batch_x2, ...]. Slices along the batch axis denote an
-        individual input passed to the kernel function. It is computed pairwise
-        with each input sliced from x1.
-
-    Returns:
-      Tensor of shape [batch_x1, batch_x2].
-    """
-    size = tf.convert_to_tensor(x1).shape.ndims
-    if size > 2:
-      raise NotImplementedError('Multiple feature dimensions is not yet '
-                                'supported.')
-    x1 = x1 / self.lengthscale
-    x2 = x2 / self.lengthscale
-    x1_squared = tf.reduce_sum(tf.square(x1), list(range(1, len(x1.shape))))
-    x2_squared = tf.reduce_sum(tf.square(x2), list(range(1, len(x2.shape))))
-    square = (x1_squared[:, tf.newaxis] +
-              x2_squared[tf.newaxis, :] -
-              2 * tf.matmul(x1, x2, transpose_b=True))
-    return self.variance * tf.exp(-square / 2)
-
-  def get_config(self):
-    return {'variance': self.variance, 'lengthscale': self.lengthscale}
-
-
-class LinearKernel(object):
-  """Linear kernel, optionally on top of a feature extractor (e.g., encoder)."""
-
-  def __init__(self, variance, bias, encoder=tf.identity):
-    self.variance = variance
-    self.bias = bias
-    self.encoder = encoder
-
-  def __call__(self, x1, x2):
-    """Computes scaled dot product of over all pairs of encoded inputs.
-
-    Args:
-      x1: Tensor of shape [batch_x1] + encoder domain. Slices along the batch
-        axis denote an individual input to be passed to the kernel. It is
-        computed pairwise with each input sliced from x2.
-      x2: Tensor of shape [batch_x2] + encoder domain. Slices along the batch
-        axis denote an individual input to be passed to the kernel. It is
-        computed pairwise with each input sliced from x1.
-
-    Returns:
-      Tensor of shape [batch_x1, batch_x2].
-    """
-    encoded_x1 = self.encoder(x1)
-    encoded_x2 = self.encoder(x2)
-    dot_product = tf.matmul(encoded_x1, encoded_x2, transpose_b=True)
-    return self.variance * dot_product + self.bias
-
-  def get_config(self):
-    return {
-        'variance': self.variance,
-        'bias': self.bias,
-        'encoder': tf.keras.utils.serialize_keras_object(self.encoder),
-    }
-
-
-class GaussianProcess(tf.keras.layers.Layer):
-  r"""Gaussian process layer.
-
-  The layer represents a distribution over functions, where a
-  stochastic forward pass appears as
-
-  ```none
-  f ~ GP(f | conditional_inputs, conditional_outputs; mean_fn, covariance_fn)
-  outputs = f(inputs)
-  ```
-
-  The optional arguments `conditional_inputs` and `conditional_outputs`
-  capture data that the GP "memorizes", i.e., it forms a posterior predictive
-  distribution. If left unspecified, the GP posits a prior predictive.
-
-  Given a call to `inputs`, an equivalent formulation in terms of function
-  outputs is
-
-  ```none
-  outputs ~ \prod_{unit=1}^{units} MultivariateNormal(output[:, unit] |
-      mean = mean_fn(inputs) + Knm Kmm^{-1} (conditional_outputs[:, unit]-mean),
-      covariance = Knn - Knm Kmm^{-1} Kmn)
-  ```
-
-  where Knm is the covariance function evaluated between all `inputs` and
-  `conditional_inputs`; Knn is between all `inputs`; Kmm is between all
-  `conditional_inputs`; and mean is the mean function evaluated on
-  `conditional_inputs`. The multivariate normal is correlated across input
-  dimensions and is independent across output dimensions.
-  """
-
-  def __init__(
-      self,
-      units,
-      mean_fn=Zeros(),
-      covariance_fn=ExponentiatedQuadratic(variance=1., lengthscale=1.),
-      conditional_inputs=None,
-      conditional_outputs=None,
-      **kwargs):
-    """Constructs layer.
-
-    Args:
-      units: integer, dimensionality of layer.
-      mean_fn: Mean function, a callable taking an inputs Tensor of shape
-        [batch, ...] and returning a Tensor of shape [batch].
-      covariance_fn: Covariance function, a callable taking two input Tensors
-        of shape [batch_x1, ...] and [batch_x2, ...] respectively, and returning
-        a positive semi-definite matrix of shape [batch_x1, batch_x2].
-      conditional_inputs: Tensor of shape [batch, ...], where batch must be the
-        same as conditional_outputs', and ellipses must match layer inputs.
-      conditional_outputs: Tensor of shape [batch, units], where batch must be
-        the same as conditional_inputs' and units is the layer's units size.
-      **kwargs: kwargs passed to parent class.
-    """
-    super(GaussianProcess, self).__init__(**kwargs)
-    self.units = int(units)
-    self.mean_fn = mean_fn
-    self.covariance_fn = covariance_fn
-    self.conditional_inputs = conditional_inputs
-    self.conditional_outputs = conditional_outputs
-
-    self.supports_masking = True
-    self.input_spec = tf.keras.layers.InputSpec(min_ndim=2)
-
-  def build(self, input_shape=None):
-    # Don't track trainable variables such as in the kernel. The user should
-    # refer to any via, e.g., self.covariance_fn or the user environment.
-    self.built = True
-
-  def call(self, inputs):
-    if self.conditional_inputs is None and self.conditional_outputs is None:
-      covariance_matrix = self.covariance_fn(inputs, inputs)
-      # Tile locations so output has shape [units, batch_size]. Covariance will
-      # broadcast to [units, batch_size, batch_size], and we perform
-      # shape manipulations to get a random variable over [batch_size, units].
-      loc = self.mean_fn(inputs)
-      loc = tf.tile(loc[tf.newaxis], [self.units] + [1] * len(loc.shape))
-    else:
-      knn = self.covariance_fn(inputs, inputs)
-      knm = self.covariance_fn(inputs, self.conditional_inputs)
-      kmm = self.covariance_fn(self.conditional_inputs, self.conditional_inputs)
-      kmm = tf.matrix_set_diag(
-          kmm, tf.matrix_diag_part(kmm) + tf.keras.backend.epsilon())
-      kmm_tril = tf.linalg.cholesky(kmm)
-      kmm_tril_operator = tf.linalg.LinearOperatorLowerTriangular(kmm_tril)
-      knm_operator = tf.linalg.LinearOperatorFullMatrix(knm)
-
-      # TODO(trandustin): Vectorize linear algebra for multiple outputs. For
-      # now, we do each separately and stack to obtain a locations Tensor of
-      # shape [units, batch_size].
-      loc = []
-      for conditional_outputs_unit in tf.unstack(self.conditional_outputs,
-                                                 axis=-1):
-        center = conditional_outputs_unit - self.mean_fn(
-            self.conditional_inputs)
-        loc_unit = knm_operator.matvec(
-            kmm_tril_operator.solvevec(kmm_tril_operator.solvevec(center),
-                                       adjoint=True))
-        loc.append(loc_unit)
-      loc = tf.stack(loc) + self.mean_fn(inputs)[tf.newaxis]
-
-      covariance_matrix = knn
-      covariance_matrix -= knm_operator.matmul(
-          kmm_tril_operator.solve(
-              kmm_tril_operator.solve(knm, adjoint_arg=True), adjoint=True))
-
-    covariance_matrix = tf.matrix_set_diag(
-        covariance_matrix,
-        tf.matrix_diag_part(covariance_matrix) + tf.keras.backend.epsilon())
-
-    # Form a multivariate normal random variable with batch_shape units and
-    # event_shape batch_size. Then make it be independent across the units
-    # dimension. Then transpose its dimensions so it is [batch_size, units].
-    random_variable = ed.MultivariateNormalFullCovariance(
-        loc=loc, covariance_matrix=covariance_matrix)
-    random_variable = ed.Independent(random_variable.distribution,
-                                     reinterpreted_batch_ndims=1)
-    bijector = tfp.bijectors.Inline(
-        forward_fn=lambda x: tf.transpose(x, [1, 0]),
-        inverse_fn=lambda y: tf.transpose(y, [1, 0]),
-        forward_event_shape_fn=lambda input_shape: input_shape[::-1],
-        forward_event_shape_tensor_fn=lambda input_shape: input_shape[::-1],
-        inverse_log_det_jacobian_fn=lambda y: tf.cast(0, y.dtype),
-        forward_min_event_ndims=2)
-    random_variable = ed.TransformedDistribution(random_variable.distribution,
-                                                 bijector=bijector)
-    return random_variable
-
-  def compute_output_shape(self, input_shape):
-    input_shape = tf.TensorShape(input_shape)
-    input_shape = input_shape.with_rank_at_least(2)
-    input_dim = input_shape[-1]
-    if isinstance(input_dim, tf.Dimension):
-      input_dim = input_dim.value
-    if input_dim is None:
-      raise ValueError(
-          'The innermost dimension of input_shape must be defined, but saw: %s'
-          % input_shape)
-    return input_shape[:-1].concatenate(self.units)
-
-  def get_config(self):
-    config = {
-        'units': self.units,
-        'mean_fn': tf.keras.utils.serialize_keras_object(self.mean_fn),
-        'covariance_fn': tf.keras.utils.serialize_keras_object(
-            self.covariance_fn),
-        'conditional_inputs': None,  # don't serialize as it can be large
-        'conditional_outputs': None,  # don't serialize as it can be large
-    }
-    base_config = super(GaussianProcess, self).get_config()
-    return dict(list(base_config.items()) + list(config.items()))
-
-
-@add_weight
-class SparseGaussianProcess(GaussianProcess):
-  r"""Gaussian process layer with inducing input and output variables.
-
-  The layer represents a distribution over functions, where a
-  stochastic forward pass appears as
-
-  ```none
-  f ~ GP(f | inducing_inputs, inducing_outputs; mean_fn, covariance_fn)
-  outputs = f(inputs)
-  ```
-
-  The arguments `inducing_inputs` and `inducing_outputs`
-  capture data that the GP "memorizes", i.e., it forms a posterior predictive
-  distribution. Typically in a variational inference scheme (and by default),
-  the inducing outputs are normally distributed with learnable location and
-  scale parameters, and the inducing inputs are learnable parameters.
-
-  Given a call to `inputs` with these defaults, an equivalent formulation in
-  terms of function outputs is
-
-  ```none
-  inducing_outputs ~ Normal(inducing_outputs | mean, stddev)
-  outputs ~ \prod_{unit=1}^{units} MultivariateNormal(output[:, unit] |
-      mean = mean_fn(inputs) + Knm Kmm^{-1} (inducing_outputs[:, unit]-mean),
-      covariance = Knn - Knm Kmm^{-1} Kmn)
-  ```
-
-  where Knm is the covariance function evaluated between all `inputs` and
-  `inducing_inputs`; Knn is between all `inputs`; Kmm is between all
-  `inducing_inputs`; and mean is the mean function evaluated on
-  `inducing_inputs`. The multivariate normal is correlated across input
-  dimensions and is independent across output dimensions.
-
-  #### Examples
-
-  We demonstrate a three-layer deep GP with variational inference (Salimbeni and
-  Deisenroth, 2017; Damianou and Lawrence, 2013). The code snippet mirrors
-  Figure 5 of Bayesian Layers. We apply it for regression given batches of
-  spatial inputs and vector-valued outputs. We flatten inputs to use the
-  default squared exponential kernel; this naturally extends to pass in a
-  more sophisticated kernel function.
-
-  ```python
-  from tensor2tensor.layers import bayes
-
-  batch_size = 256
-  dataset_size = 10000
-  features, labels = load_spatial_data(batch_size)
-
-  model = tf.keras.Sequential([
-    tf.keras.layers.Flatten(),
-    layers.SparseGaussianProcess(256, num_inducing=512),
-    layers.SparseGaussianProcess(256, num_inducing=512),
-    layers.SparseGaussianProcess(10, num_inducing=512),
-  ])
-  predictions = model(features)
-  nll = tf.losses.mean_squared_error(labels=labels, predictions=predictions)
-  kl = sum(model.losses) / dataset_size
-  loss = nll + kl
-  train_op = tf.train.AdamOptimizer().minimize(loss)
-  ```
-  """
-
-  def __init__(
-      self,
-      units,
-      num_inducing,
-      mean_fn=Zeros(),
-      covariance_fn=ExponentiatedQuadratic(variance=1., lengthscale=1.),
-      inducing_inputs_initializer='random_normal',
-      inducing_outputs_initializer='trainable_normal',
-      inducing_inputs_regularizer=None,
-      inducing_outputs_regularizer='normal_kl_divergence',
-      inducing_inputs_constraint=None,
-      inducing_outputs_constraint=None,
-      **kwargs):
-    """Constructs layer.
-
-    Args:
-      units: integer, dimensionality of layer.
-      num_inducing: integer, number of inducing points for the approximation.
-      mean_fn: Mean function, a callable taking an inputs Tensor of shape
-        [batch, ...] and returning a Tensor of shape [batch].
-      covariance_fn: Covariance function, a callable taking two input Tensors
-        of shape [batch_x1, ...] and [batch_x2, ...] respectively, and returning
-        a positive semi-definite matrix of shape [batch_x1, batch_x2].
-      inducing_inputs_initializer: Initializer for the inducing inputs.
-      inducing_outputs_initializer: Initializer for the inducing outputs.
-      inducing_inputs_regularizer: Regularizer function applied to the inducing
-        inputs.
-      inducing_outputs_regularizer: Regularizer function applied to the inducing
-        outputs.
-      inducing_inputs_constraint: Constraint function applied to the inducing
-        inputs.
-      inducing_outputs_constraint: Constraint function applied to the inducing
-        outputs.
-      **kwargs: kwargs passed to parent class.
-    """
-    super(SparseGaussianProcess, self).__init__(
-        units=units,
-        mean_fn=mean_fn,
-        covariance_fn=covariance_fn,
-        conditional_inputs=None,
-        conditional_outputs=None,
-        **kwargs)
-    self.num_inducing = num_inducing
-    self.inducing_inputs_initializer = initializers.get(
-        inducing_inputs_initializer)
-    self.inducing_outputs_initializer = initializers.get(
-        inducing_outputs_initializer)
-    self.inducing_inputs_regularizer = regularizers.get(
-        inducing_inputs_regularizer)
-    self.inducing_outputs_regularizer = regularizers.get(
-        inducing_outputs_regularizer)
-    self.inducing_inputs_constraint = constraints.get(
-        inducing_inputs_constraint)
-    self.inducing_outputs_constraint = constraints.get(
-        inducing_outputs_constraint)
-
-  def build(self, input_shape=None):
-    input_shape = tf.TensorShape(input_shape)
-    input_dim = input_shape[-1]
-    if isinstance(input_dim, tf.Dimension):
-      input_dim = input_dim.value
-    self.conditional_inputs = self.add_weight(
-        shape=(self.num_inducing, input_dim),
-        name='inducing_inputs',
-        initializer=self.inducing_inputs_initializer,
-        regularizer=self.inducing_inputs_regularizer,
-        constraint=self.inducing_inputs_constraint)
-    self.conditional_outputs = self.add_weight(
-        shape=(self.num_inducing, self.units),
-        name='inducing_outputs',
-        initializer=self.inducing_outputs_initializer,
-        regularizer=self.inducing_outputs_regularizer,
-        constraint=self.inducing_outputs_constraint)
-    super(SparseGaussianProcess, self).build(input_shape)
-
-
-class BayesianLinearModel(tf.keras.Model):
-  r"""Bayesian linear model with standard normal prior over its coefficients.
-
-  A forward pass computes the mean of the exact predictive distribution
-
-  ```none
-  p(outputs | inputs) = \int Normal(outputs | coeffs * inputs, noise_variance)
-                             Normal(coeffs | 0, 1) dweights dbias.
-  ```
-
-  It takes a Tensor of shape [batch_size, input_dim] as input and returns a
-  Normal random variable of shape [batch_size] representing its outputs.
-  After `fit()`, the forward pass computes the exact posterior predictive
-  distribution.
-  """
-
-  def __init__(self, noise_variance, **kwargs):
-    super(BayesianLinearModel, self).__init__(**kwargs)
-    self.noise_variance = noise_variance
-    self.coeffs_precision_tril_op = None
-    self.coeffs_mean = None
-
-  def call(self, inputs):
-    if self.coeffs_mean is None and self.coeffs_precision_tril_op is None:
-      # p(mean(ynew) | xnew) = Normal(ynew | mean = 0, variance = xnew xnew^T)
-      predictive_mean = 0.
-      predictive_variance = tf.reduce_sum(tf.square(inputs), -1)
-    else:
-      # p(mean(ynew) | xnew, x, y) = Normal(ynew |
-      #   mean = xnew (1/noise_variance) (1/noise_variance x^T x + I)^{-1}x^T y,
-      #   variance = xnew (1/noise_variance x^T x + I)^{-1} xnew^T)
-      predictive_mean = tf.einsum('nm,m->n', inputs, self.coeffs_mean)
-      predictive_covariance = tf.matmul(
-          inputs,
-          self.coeffs_precision_tril_op.solve(
-              self.coeffs_precision_tril_op.solve(inputs, adjoint_arg=True),
-              adjoint=True))
-      predictive_variance = tf.diag_part(predictive_covariance)
-    return ed.Normal(loc=predictive_mean, scale=tf.sqrt(predictive_variance))
-
-  def fit(self, x=None, y=None):
-    # p(coeffs | x, y) = Normal(coeffs |
-    #   mean = (1/noise_variance) (1/noise_variance x^T x + I)^{-1} x^T y,
-    #   covariance = (1/noise_variance x^T x + I)^{-1})
-    # TODO(trandustin): We newly fit the data at each call. Extend to do
-    # Bayesian updating.
-    kernel_matrix = tf.matmul(x, x, transpose_a=True) / self.noise_variance
-    coeffs_precision = tf.matrix_set_diag(
-        kernel_matrix, tf.matrix_diag_part(kernel_matrix) + 1.)
-    coeffs_precision_tril = tf.linalg.cholesky(coeffs_precision)
-    self.coeffs_precision_tril_op = tf.linalg.LinearOperatorLowerTriangular(
-        coeffs_precision_tril)
-    self.coeffs_mean = self.coeffs_precision_tril_op.solvevec(
-        self.coeffs_precision_tril_op.solvevec(tf.einsum('nm,n->m', x, y)),
-        adjoint=True) / self.noise_variance
-    # TODO(trandustin): To be fully Keras-compatible, return History object.
-    return
-
-
 class MixtureLogistic(tf.keras.layers.Layer):
   """Stochastic output layer, distributed as a mixture of logistics."""
 
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 5c2b56e00..25a4b3b91 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for common Bayes."""
+"""Tests for Bayesian neural network layers."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -384,71 +384,6 @@ def testDenseDVIMoments(self):
     percent_mismatches = num_mismatches / float(batch_size * units * units)
     self.assertLessEqual(percent_mismatches, 0.05)
 
-  @test_utils.run_in_graph_and_eager_modes()
-  def testGaussianProcessPosterior(self):
-    train_batch_size = 3
-    test_batch_size = 2
-    input_dim = 4
-    output_dim = 5
-    features = tf.to_float(np.random.rand(train_batch_size, input_dim))
-    labels = tf.to_float(np.random.rand(train_batch_size, output_dim))
-    layer = bayes.GaussianProcess(output_dim,
-                                  conditional_inputs=features,
-                                  conditional_outputs=labels)
-    test_features = tf.to_float(np.random.rand(test_batch_size, input_dim))
-    test_labels = tf.to_float(np.random.rand(test_batch_size, output_dim))
-    test_outputs = layer(test_features)
-    test_nats = -test_outputs.distribution.log_prob(test_labels)
-    self.evaluate(tf.global_variables_initializer())
-    test_nats_val, outputs_val = self.evaluate([test_nats, test_outputs])
-    self.assertEqual(test_nats_val.shape, ())
-    self.assertGreaterEqual(test_nats_val, 0.)
-    self.assertEqual(outputs_val.shape, (test_batch_size, output_dim))
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testGaussianProcessPrior(self):
-    batch_size = 3
-    input_dim = 4
-    output_dim = 5
-    features = tf.to_float(np.random.rand(batch_size, input_dim))
-    labels = tf.to_float(np.random.rand(batch_size, output_dim))
-    model = tf.keras.Sequential([
-        tf.keras.layers.Dense(2, activation=None),
-        bayes.GaussianProcess(output_dim),
-    ])
-    outputs = model(features)
-    log_prob = outputs.distribution.log_prob(labels)
-    self.evaluate(tf.global_variables_initializer())
-    log_prob_val, outputs_val = self.evaluate([log_prob, outputs])
-    self.assertEqual(log_prob_val.shape, ())
-    self.assertLessEqual(log_prob_val, 0.)
-    self.assertEqual(outputs_val.shape, (batch_size, output_dim))
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testSparseGaussianProcess(self):
-    dataset_size = 10
-    batch_size = 3
-    input_dim = 4
-    output_dim = 5
-    features = tf.to_float(np.random.rand(batch_size, input_dim))
-    labels = tf.to_float(np.random.rand(batch_size, output_dim))
-    model = bayes.SparseGaussianProcess(output_dim, num_inducing=2)
-    with tf.GradientTape() as tape:
-      predictions = model(features)
-      nll = -tf.reduce_mean(predictions.distribution.log_prob(labels))
-      kl = sum(model.losses) / dataset_size
-      loss = nll + kl
-
-    self.evaluate(tf.global_variables_initializer())
-    grads = tape.gradient(nll, model.variables)
-    for grad in grads:
-      self.assertIsNotNone(grad)
-
-    loss_val, predictions_val = self.evaluate([loss, predictions])
-    self.assertEqual(loss_val.shape, ())
-    self.assertGreaterEqual(loss_val, 0.)
-    self.assertEqual(predictions_val.shape, (batch_size, output_dim))
-
   @parameterized.parameters(
       {"lstm_cell": bayes.LSTMCellFlipout,
        "kernel_initializer": "zeros",
@@ -612,37 +547,6 @@ def testLSTMCellModel(self, lstm_cell):
     self.assertAllClose(res2, res3)
     self.assertLen(model.losses, 2)
 
-  @test_utils.run_in_graph_and_eager_modes()
-  def testBayesianLinearModel(self):
-    """Tests that model makes reasonable predictions."""
-    np.random.seed(42)
-    train_batch_size = 5
-    test_batch_size = 2
-    num_features = 3
-    noise_variance = 0.01
-    coeffs = tf.range(num_features, dtype=tf.float32)
-    features = tf.to_float(np.random.randn(train_batch_size, num_features))
-    labels = (tf.tensordot(features, coeffs, [[-1], [0]])
-              + noise_variance * tf.to_float(np.random.randn(train_batch_size)))
-
-    model = bayes.BayesianLinearModel(noise_variance=noise_variance)
-    model.fit(features, labels)
-
-    test_features = tf.to_float(np.random.randn(test_batch_size, num_features))
-    test_labels = tf.tensordot(test_features, coeffs, [[-1], [0]])
-    outputs = model(test_features)
-    test_predictions = outputs.distribution.mean()
-    test_predictions_variance = outputs.distribution.variance()
-
-    [
-        test_labels_val, test_predictions_val, test_predictions_variance_val,
-    ] = self.evaluate(
-        [test_labels, test_predictions, test_predictions_variance])
-    self.assertEqual(test_predictions_val.shape, (test_batch_size,))
-    self.assertEqual(test_predictions_variance_val.shape, (test_batch_size,))
-    self.assertAllClose(test_predictions_val, test_labels_val, atol=0.1)
-    self.assertAllLessEqual(test_predictions_variance_val, noise_variance)
-
   @test_utils.run_in_graph_and_eager_modes()
   def testMixtureLogistic(self):
     batch_size = 3
diff --git a/tensor2tensor/layers/gaussian_process.py b/tensor2tensor/layers/gaussian_process.py
new file mode 100644
index 000000000..c0714457b
--- /dev/null
+++ b/tensor2tensor/layers/gaussian_process.py
@@ -0,0 +1,464 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Gaussian process layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.keras import constraints
+from tensor2tensor.keras import initializers
+from tensor2tensor.keras import regularizers
+from tensor2tensor.layers import bayes
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tensorflow_probability import edward2 as ed
+
+
+class Zeros(object):
+  """Function returning zeros tensor of same shape excluding the last dim."""
+
+  def __call__(self, inputs):
+    return tf.zeros(tf.shape(inputs)[:-1], inputs.dtype)
+
+  def get_config(self):
+    return {}
+
+
+class ExponentiatedQuadratic(object):
+  """Exponentiated quadratic kernel."""
+
+  def __init__(self, variance, lengthscale):
+    self.variance = variance
+    self.lengthscale = lengthscale
+
+  def __call__(self, x1, x2):
+    """Computes exponentiated quadratic over all pairs of inputs.
+
+    Args:
+      x1: Tensor of shape [batch_x1, ...]. Slices along the batch axis denote an
+        individual input to be passed to the kernel. It is computed pairwise
+        with each input sliced from x2.
+      x2: Tensor of shape [batch_x2, ...]. Slices along the batch axis denote an
+        individual input passed to the kernel function. It is computed pairwise
+        with each input sliced from x1.
+
+    Returns:
+      Tensor of shape [batch_x1, batch_x2].
+    """
+    size = tf.convert_to_tensor(x1).shape.ndims
+    if size > 2:
+      raise NotImplementedError('Multiple feature dimensions is not yet '
+                                'supported.')
+    x1 = x1 / self.lengthscale
+    x2 = x2 / self.lengthscale
+    x1_squared = tf.reduce_sum(tf.square(x1), list(range(1, len(x1.shape))))
+    x2_squared = tf.reduce_sum(tf.square(x2), list(range(1, len(x2.shape))))
+    square = (x1_squared[:, tf.newaxis] +
+              x2_squared[tf.newaxis, :] -
+              2 * tf.matmul(x1, x2, transpose_b=True))
+    return self.variance * tf.exp(-square / 2)
+
+  def get_config(self):
+    return {'variance': self.variance, 'lengthscale': self.lengthscale}
+
+
+class LinearKernel(object):
+  """Linear kernel, optionally on top of a feature extractor (e.g., encoder)."""
+
+  def __init__(self, variance, bias, encoder=tf.identity):
+    self.variance = variance
+    self.bias = bias
+    self.encoder = encoder
+
+  def __call__(self, x1, x2):
+    """Computes scaled dot product of over all pairs of encoded inputs.
+
+    Args:
+      x1: Tensor of shape [batch_x1] + encoder domain. Slices along the batch
+        axis denote an individual input to be passed to the kernel. It is
+        computed pairwise with each input sliced from x2.
+      x2: Tensor of shape [batch_x2] + encoder domain. Slices along the batch
+        axis denote an individual input to be passed to the kernel. It is
+        computed pairwise with each input sliced from x1.
+
+    Returns:
+      Tensor of shape [batch_x1, batch_x2].
+    """
+    encoded_x1 = self.encoder(x1)
+    encoded_x2 = self.encoder(x2)
+    dot_product = tf.matmul(encoded_x1, encoded_x2, transpose_b=True)
+    return self.variance * dot_product + self.bias
+
+  def get_config(self):
+    return {
+        'variance': self.variance,
+        'bias': self.bias,
+        'encoder': tf.keras.utils.serialize_keras_object(self.encoder),
+    }
+
+
+class GaussianProcess(tf.keras.layers.Layer):
+  r"""Gaussian process layer.
+
+  The layer represents a distribution over functions, where a
+  stochastic forward pass appears as
+
+  ```none
+  f ~ GP(f | conditional_inputs, conditional_outputs; mean_fn, covariance_fn)
+  outputs = f(inputs)
+  ```
+
+  The optional arguments `conditional_inputs` and `conditional_outputs`
+  capture data that the GP "memorizes", i.e., it forms a posterior predictive
+  distribution. If left unspecified, the GP posits a prior predictive.
+
+  Given a call to `inputs`, an equivalent formulation in terms of function
+  outputs is
+
+  ```none
+  outputs ~ \prod_{unit=1}^{units} MultivariateNormal(output[:, unit] |
+      mean = mean_fn(inputs) + Knm Kmm^{-1} (conditional_outputs[:, unit]-mean),
+      covariance = Knn - Knm Kmm^{-1} Kmn)
+  ```
+
+  where Knm is the covariance function evaluated between all `inputs` and
+  `conditional_inputs`; Knn is between all `inputs`; Kmm is between all
+  `conditional_inputs`; and mean is the mean function evaluated on
+  `conditional_inputs`. The multivariate normal is correlated across input
+  dimensions and is independent across output dimensions.
+  """
+
+  def __init__(
+      self,
+      units,
+      mean_fn=Zeros(),
+      covariance_fn=ExponentiatedQuadratic(variance=1., lengthscale=1.),
+      conditional_inputs=None,
+      conditional_outputs=None,
+      **kwargs):
+    """Constructs layer.
+
+    Args:
+      units: integer, dimensionality of layer.
+      mean_fn: Mean function, a callable taking an inputs Tensor of shape
+        [batch, ...] and returning a Tensor of shape [batch].
+      covariance_fn: Covariance function, a callable taking two input Tensors
+        of shape [batch_x1, ...] and [batch_x2, ...] respectively, and returning
+        a positive semi-definite matrix of shape [batch_x1, batch_x2].
+      conditional_inputs: Tensor of shape [batch, ...], where batch must be the
+        same as conditional_outputs', and ellipses must match layer inputs.
+      conditional_outputs: Tensor of shape [batch, units], where batch must be
+        the same as conditional_inputs' and units is the layer's units size.
+      **kwargs: kwargs passed to parent class.
+    """
+    super(GaussianProcess, self).__init__(**kwargs)
+    self.units = int(units)
+    self.mean_fn = mean_fn
+    self.covariance_fn = covariance_fn
+    self.conditional_inputs = conditional_inputs
+    self.conditional_outputs = conditional_outputs
+
+    self.supports_masking = True
+    self.input_spec = tf.keras.layers.InputSpec(min_ndim=2)
+
+  def build(self, input_shape=None):
+    # Don't track trainable variables such as in the kernel. The user should
+    # refer to any via, e.g., self.covariance_fn or the user environment.
+    self.built = True
+
+  def call(self, inputs):
+    if self.conditional_inputs is None and self.conditional_outputs is None:
+      covariance_matrix = self.covariance_fn(inputs, inputs)
+      # Tile locations so output has shape [units, batch_size]. Covariance will
+      # broadcast to [units, batch_size, batch_size], and we perform
+      # shape manipulations to get a random variable over [batch_size, units].
+      loc = self.mean_fn(inputs)
+      loc = tf.tile(loc[tf.newaxis], [self.units] + [1] * len(loc.shape))
+    else:
+      knn = self.covariance_fn(inputs, inputs)
+      knm = self.covariance_fn(inputs, self.conditional_inputs)
+      kmm = self.covariance_fn(self.conditional_inputs, self.conditional_inputs)
+      kmm = tf.matrix_set_diag(
+          kmm, tf.matrix_diag_part(kmm) + tf.keras.backend.epsilon())
+      kmm_tril = tf.linalg.cholesky(kmm)
+      kmm_tril_operator = tf.linalg.LinearOperatorLowerTriangular(kmm_tril)
+      knm_operator = tf.linalg.LinearOperatorFullMatrix(knm)
+
+      # TODO(trandustin): Vectorize linear algebra for multiple outputs. For
+      # now, we do each separately and stack to obtain a locations Tensor of
+      # shape [units, batch_size].
+      loc = []
+      for conditional_outputs_unit in tf.unstack(self.conditional_outputs,
+                                                 axis=-1):
+        center = conditional_outputs_unit - self.mean_fn(
+            self.conditional_inputs)
+        loc_unit = knm_operator.matvec(
+            kmm_tril_operator.solvevec(kmm_tril_operator.solvevec(center),
+                                       adjoint=True))
+        loc.append(loc_unit)
+      loc = tf.stack(loc) + self.mean_fn(inputs)[tf.newaxis]
+
+      covariance_matrix = knn
+      covariance_matrix -= knm_operator.matmul(
+          kmm_tril_operator.solve(
+              kmm_tril_operator.solve(knm, adjoint_arg=True), adjoint=True))
+
+    covariance_matrix = tf.matrix_set_diag(
+        covariance_matrix,
+        tf.matrix_diag_part(covariance_matrix) + tf.keras.backend.epsilon())
+
+    # Form a multivariate normal random variable with batch_shape units and
+    # event_shape batch_size. Then make it be independent across the units
+    # dimension. Then transpose its dimensions so it is [batch_size, units].
+    random_variable = ed.MultivariateNormalFullCovariance(
+        loc=loc, covariance_matrix=covariance_matrix)
+    random_variable = ed.Independent(random_variable.distribution,
+                                     reinterpreted_batch_ndims=1)
+    bijector = tfp.bijectors.Inline(
+        forward_fn=lambda x: tf.transpose(x, [1, 0]),
+        inverse_fn=lambda y: tf.transpose(y, [1, 0]),
+        forward_event_shape_fn=lambda input_shape: input_shape[::-1],
+        forward_event_shape_tensor_fn=lambda input_shape: input_shape[::-1],
+        inverse_log_det_jacobian_fn=lambda y: tf.cast(0, y.dtype),
+        forward_min_event_ndims=2)
+    random_variable = ed.TransformedDistribution(random_variable.distribution,
+                                                 bijector=bijector)
+    return random_variable
+
+  def compute_output_shape(self, input_shape):
+    input_shape = tf.TensorShape(input_shape)
+    input_shape = input_shape.with_rank_at_least(2)
+    input_dim = input_shape[-1]
+    if isinstance(input_dim, tf.Dimension):
+      input_dim = input_dim.value
+    if input_dim is None:
+      raise ValueError(
+          'The innermost dimension of input_shape must be defined, but saw: %s'
+          % input_shape)
+    return input_shape[:-1].concatenate(self.units)
+
+  def get_config(self):
+    config = {
+        'units': self.units,
+        'mean_fn': tf.keras.utils.serialize_keras_object(self.mean_fn),
+        'covariance_fn': tf.keras.utils.serialize_keras_object(
+            self.covariance_fn),
+        'conditional_inputs': None,  # don't serialize as it can be large
+        'conditional_outputs': None,  # don't serialize as it can be large
+    }
+    base_config = super(GaussianProcess, self).get_config()
+    return dict(list(base_config.items()) + list(config.items()))
+
+
+@bayes.add_weight
+class SparseGaussianProcess(GaussianProcess):
+  r"""Gaussian process layer with inducing input and output variables.
+
+  The layer represents a distribution over functions, where a
+  stochastic forward pass appears as
+
+  ```none
+  f ~ GP(f | inducing_inputs, inducing_outputs; mean_fn, covariance_fn)
+  outputs = f(inputs)
+  ```
+
+  The arguments `inducing_inputs` and `inducing_outputs`
+  capture data that the GP "memorizes", i.e., it forms a posterior predictive
+  distribution. Typically in a variational inference scheme (and by default),
+  the inducing outputs are normally distributed with learnable location and
+  scale parameters, and the inducing inputs are learnable parameters.
+
+  Given a call to `inputs` with these defaults, an equivalent formulation in
+  terms of function outputs is
+
+  ```none
+  inducing_outputs ~ Normal(inducing_outputs | mean, stddev)
+  outputs ~ \prod_{unit=1}^{units} MultivariateNormal(output[:, unit] |
+      mean = mean_fn(inputs) + Knm Kmm^{-1} (inducing_outputs[:, unit]-mean),
+      covariance = Knn - Knm Kmm^{-1} Kmn)
+  ```
+
+  where Knm is the covariance function evaluated between all `inputs` and
+  `inducing_inputs`; Knn is between all `inputs`; Kmm is between all
+  `inducing_inputs`; and mean is the mean function evaluated on
+  `inducing_inputs`. The multivariate normal is correlated across input
+  dimensions and is independent across output dimensions.
+
+  #### Examples
+
+  We demonstrate a three-layer deep GP with variational inference (Salimbeni and
+  Deisenroth, 2017; Damianou and Lawrence, 2013). The code snippet mirrors
+  Figure 5 of Bayesian Layers. We apply it for regression given batches of
+  spatial inputs and vector-valued outputs. We flatten inputs to use the
+  default squared exponential kernel; this naturally extends to pass in a
+  more sophisticated kernel function.
+
+  ```python
+  from tensor2tensor.layers import bayes
+
+  batch_size = 256
+  dataset_size = 10000
+  features, labels = load_spatial_data(batch_size)
+
+  model = tf.keras.Sequential([
+    tf.keras.layers.Flatten(),
+    layers.SparseGaussianProcess(256, num_inducing=512),
+    layers.SparseGaussianProcess(256, num_inducing=512),
+    layers.SparseGaussianProcess(10, num_inducing=512),
+  ])
+  predictions = model(features)
+  nll = tf.losses.mean_squared_error(labels=labels, predictions=predictions)
+  kl = sum(model.losses) / dataset_size
+  loss = nll + kl
+  train_op = tf.train.AdamOptimizer().minimize(loss)
+  ```
+  """
+
+  def __init__(
+      self,
+      units,
+      num_inducing,
+      mean_fn=Zeros(),
+      covariance_fn=ExponentiatedQuadratic(variance=1., lengthscale=1.),
+      inducing_inputs_initializer='random_normal',
+      inducing_outputs_initializer='trainable_normal',
+      inducing_inputs_regularizer=None,
+      inducing_outputs_regularizer='normal_kl_divergence',
+      inducing_inputs_constraint=None,
+      inducing_outputs_constraint=None,
+      **kwargs):
+    """Constructs layer.
+
+    Args:
+      units: integer, dimensionality of layer.
+      num_inducing: integer, number of inducing points for the approximation.
+      mean_fn: Mean function, a callable taking an inputs Tensor of shape
+        [batch, ...] and returning a Tensor of shape [batch].
+      covariance_fn: Covariance function, a callable taking two input Tensors
+        of shape [batch_x1, ...] and [batch_x2, ...] respectively, and returning
+        a positive semi-definite matrix of shape [batch_x1, batch_x2].
+      inducing_inputs_initializer: Initializer for the inducing inputs.
+      inducing_outputs_initializer: Initializer for the inducing outputs.
+      inducing_inputs_regularizer: Regularizer function applied to the inducing
+        inputs.
+      inducing_outputs_regularizer: Regularizer function applied to the inducing
+        outputs.
+      inducing_inputs_constraint: Constraint function applied to the inducing
+        inputs.
+      inducing_outputs_constraint: Constraint function applied to the inducing
+        outputs.
+      **kwargs: kwargs passed to parent class.
+    """
+    super(SparseGaussianProcess, self).__init__(
+        units=units,
+        mean_fn=mean_fn,
+        covariance_fn=covariance_fn,
+        conditional_inputs=None,
+        conditional_outputs=None,
+        **kwargs)
+    self.num_inducing = num_inducing
+    self.inducing_inputs_initializer = initializers.get(
+        inducing_inputs_initializer)
+    self.inducing_outputs_initializer = initializers.get(
+        inducing_outputs_initializer)
+    self.inducing_inputs_regularizer = regularizers.get(
+        inducing_inputs_regularizer)
+    self.inducing_outputs_regularizer = regularizers.get(
+        inducing_outputs_regularizer)
+    self.inducing_inputs_constraint = constraints.get(
+        inducing_inputs_constraint)
+    self.inducing_outputs_constraint = constraints.get(
+        inducing_outputs_constraint)
+
+  def build(self, input_shape=None):
+    input_shape = tf.TensorShape(input_shape)
+    input_dim = input_shape[-1]
+    if isinstance(input_dim, tf.Dimension):
+      input_dim = input_dim.value
+    self.conditional_inputs = self.add_weight(
+        shape=(self.num_inducing, input_dim),
+        name='inducing_inputs',
+        initializer=self.inducing_inputs_initializer,
+        regularizer=self.inducing_inputs_regularizer,
+        constraint=self.inducing_inputs_constraint)
+    self.conditional_outputs = self.add_weight(
+        shape=(self.num_inducing, self.units),
+        name='inducing_outputs',
+        initializer=self.inducing_outputs_initializer,
+        regularizer=self.inducing_outputs_regularizer,
+        constraint=self.inducing_outputs_constraint)
+    super(SparseGaussianProcess, self).build(input_shape)
+
+
+class BayesianLinearModel(tf.keras.Model):
+  r"""Bayesian linear model with standard normal prior over its coefficients.
+
+  A forward pass computes the mean of the exact predictive distribution
+
+  ```none
+  p(outputs | inputs) = \int Normal(outputs | coeffs * inputs, noise_variance)
+                             Normal(coeffs | 0, 1) dweights dbias.
+  ```
+
+  It takes a Tensor of shape [batch_size, input_dim] as input and returns a
+  Normal random variable of shape [batch_size] representing its outputs.
+  After `fit()`, the forward pass computes the exact posterior predictive
+  distribution.
+  """
+
+  def __init__(self, noise_variance, **kwargs):
+    super(BayesianLinearModel, self).__init__(**kwargs)
+    self.noise_variance = noise_variance
+    self.coeffs_precision_tril_op = None
+    self.coeffs_mean = None
+
+  def call(self, inputs):
+    if self.coeffs_mean is None and self.coeffs_precision_tril_op is None:
+      # p(mean(ynew) | xnew) = Normal(ynew | mean = 0, variance = xnew xnew^T)
+      predictive_mean = 0.
+      predictive_variance = tf.reduce_sum(tf.square(inputs), -1)
+    else:
+      # p(mean(ynew) | xnew, x, y) = Normal(ynew |
+      #   mean = xnew (1/noise_variance) (1/noise_variance x^T x + I)^{-1}x^T y,
+      #   variance = xnew (1/noise_variance x^T x + I)^{-1} xnew^T)
+      predictive_mean = tf.einsum('nm,m->n', inputs, self.coeffs_mean)
+      predictive_covariance = tf.matmul(
+          inputs,
+          self.coeffs_precision_tril_op.solve(
+              self.coeffs_precision_tril_op.solve(inputs, adjoint_arg=True),
+              adjoint=True))
+      predictive_variance = tf.diag_part(predictive_covariance)
+    return ed.Normal(loc=predictive_mean, scale=tf.sqrt(predictive_variance))
+
+  def fit(self, x=None, y=None):
+    # p(coeffs | x, y) = Normal(coeffs |
+    #   mean = (1/noise_variance) (1/noise_variance x^T x + I)^{-1} x^T y,
+    #   covariance = (1/noise_variance x^T x + I)^{-1})
+    # TODO(trandustin): We newly fit the data at each call. Extend to do
+    # Bayesian updating.
+    kernel_matrix = tf.matmul(x, x, transpose_a=True) / self.noise_variance
+    coeffs_precision = tf.matrix_set_diag(
+        kernel_matrix, tf.matrix_diag_part(kernel_matrix) + 1.)
+    coeffs_precision_tril = tf.linalg.cholesky(coeffs_precision)
+    self.coeffs_precision_tril_op = tf.linalg.LinearOperatorLowerTriangular(
+        coeffs_precision_tril)
+    self.coeffs_mean = self.coeffs_precision_tril_op.solvevec(
+        self.coeffs_precision_tril_op.solvevec(tf.einsum('nm,n->m', x, y)),
+        adjoint=True) / self.noise_variance
+    # TODO(trandustin): To be fully Keras-compatible, return History object.
+    return
diff --git a/tensor2tensor/layers/gaussian_process_test.py b/tensor2tensor/layers/gaussian_process_test.py
new file mode 100644
index 000000000..a4d33d9e3
--- /dev/null
+++ b/tensor2tensor/layers/gaussian_process_test.py
@@ -0,0 +1,131 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 Gaussian process layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensor2tensor.layers import gaussian_process
+from tensor2tensor.utils import test_utils
+
+import tensorflow as tf
+tf.compat.v1.enable_eager_execution()
+
+
+class GaussianProcessTest(tf.test.TestCase):
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testGaussianProcessPosterior(self):
+    train_batch_size = 3
+    test_batch_size = 2
+    input_dim = 4
+    output_dim = 5
+    features = tf.to_float(np.random.rand(train_batch_size, input_dim))
+    labels = tf.to_float(np.random.rand(train_batch_size, output_dim))
+    layer = gaussian_process.GaussianProcess(output_dim,
+                                             conditional_inputs=features,
+                                             conditional_outputs=labels)
+    test_features = tf.to_float(np.random.rand(test_batch_size, input_dim))
+    test_labels = tf.to_float(np.random.rand(test_batch_size, output_dim))
+    test_outputs = layer(test_features)
+    test_nats = -test_outputs.distribution.log_prob(test_labels)
+    self.evaluate(tf.global_variables_initializer())
+    test_nats_val, outputs_val = self.evaluate([test_nats, test_outputs])
+    self.assertEqual(test_nats_val.shape, ())
+    self.assertGreaterEqual(test_nats_val, 0.)
+    self.assertEqual(outputs_val.shape, (test_batch_size, output_dim))
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testGaussianProcessPrior(self):
+    batch_size = 3
+    input_dim = 4
+    output_dim = 5
+    features = tf.to_float(np.random.rand(batch_size, input_dim))
+    labels = tf.to_float(np.random.rand(batch_size, output_dim))
+    model = tf.keras.Sequential([
+        tf.keras.layers.Dense(2, activation=None),
+        gaussian_process.GaussianProcess(output_dim),
+    ])
+    outputs = model(features)
+    log_prob = outputs.distribution.log_prob(labels)
+    self.evaluate(tf.global_variables_initializer())
+    log_prob_val, outputs_val = self.evaluate([log_prob, outputs])
+    self.assertEqual(log_prob_val.shape, ())
+    self.assertLessEqual(log_prob_val, 0.)
+    self.assertEqual(outputs_val.shape, (batch_size, output_dim))
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testSparseGaussianProcess(self):
+    dataset_size = 10
+    batch_size = 3
+    input_dim = 4
+    output_dim = 5
+    features = tf.to_float(np.random.rand(batch_size, input_dim))
+    labels = tf.to_float(np.random.rand(batch_size, output_dim))
+    model = gaussian_process.SparseGaussianProcess(output_dim, num_inducing=2)
+    with tf.GradientTape() as tape:
+      predictions = model(features)
+      nll = -tf.reduce_mean(predictions.distribution.log_prob(labels))
+      kl = sum(model.losses) / dataset_size
+      loss = nll + kl
+
+    self.evaluate(tf.global_variables_initializer())
+    grads = tape.gradient(nll, model.variables)
+    for grad in grads:
+      self.assertIsNotNone(grad)
+
+    loss_val, predictions_val = self.evaluate([loss, predictions])
+    self.assertEqual(loss_val.shape, ())
+    self.assertGreaterEqual(loss_val, 0.)
+    self.assertEqual(predictions_val.shape, (batch_size, output_dim))
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testBayesianLinearModel(self):
+    """Tests that model makes reasonable predictions."""
+    np.random.seed(42)
+    train_batch_size = 5
+    test_batch_size = 2
+    num_features = 3
+    noise_variance = 0.01
+    coeffs = tf.range(num_features, dtype=tf.float32)
+    features = tf.to_float(np.random.randn(train_batch_size, num_features))
+    labels = (tf.tensordot(features, coeffs, [[-1], [0]])
+              + noise_variance * tf.to_float(np.random.randn(train_batch_size)))
+
+    model = gaussian_process.BayesianLinearModel(noise_variance=noise_variance)
+    model.fit(features, labels)
+
+    test_features = tf.to_float(np.random.randn(test_batch_size, num_features))
+    test_labels = tf.tensordot(test_features, coeffs, [[-1], [0]])
+    outputs = model(test_features)
+    test_predictions = outputs.distribution.mean()
+    test_predictions_variance = outputs.distribution.variance()
+
+    [
+        test_labels_val, test_predictions_val, test_predictions_variance_val,
+    ] = self.evaluate(
+        [test_labels, test_predictions, test_predictions_variance])
+    self.assertEqual(test_predictions_val.shape, (test_batch_size,))
+    self.assertEqual(test_predictions_variance_val.shape, (test_batch_size,))
+    self.assertAllClose(test_predictions_val, test_labels_val, atol=0.1)
+    self.assertAllLessEqual(test_predictions_variance_val, noise_variance)
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 7c7ccea9d622ed3e85bd56ec0e5a41f573d4e424 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 10 Jul 2019 11:44:41 -0700
Subject: [PATCH 2184/2720] Remove categorical sampling (equivalent to Gumbel)
 and allow various sampling temperatures

PiperOrigin-RevId: 257449707
---
 tensor2tensor/envs/env_problem_utils.py | 42 +++----------------------
 tensor2tensor/trax/rlax/ppo.py          | 23 ++++++++------
 2 files changed, 18 insertions(+), 47 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 11aebd68d..8064c4670 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -22,7 +22,6 @@
 import time
 import numpy as np
 
-CATEGORICAL_SAMPLING = "categorical"
 EPSILON_GREEDY = "epsilon-greedy"
 GUMBEL_SAMPLING = "gumbel"
 
@@ -57,8 +56,8 @@ def play_env_problem_with_policy(env,
                                  max_timestep=None,
                                  reset=True,
                                  rng=None,
-                                 policy_sampling=CATEGORICAL_SAMPLING,
-                                 temperature=0.5,
+                                 policy_sampling=GUMBEL_SAMPLING,
+                                 temperature=1.0,
                                  eps=0.1,
                                  len_history_for_policy=32,
                                  num_to_keep=1):
@@ -76,8 +75,8 @@ def play_env_problem_with_policy(env,
       max_max_timestep is None or < 0
     rng: jax rng, splittable.
     policy_sampling: string, how to select an action given a policy, one of:
-      CATEGORICAL_SAMPLING, GREEDY, GUMBEL_SAMPLING
-    temperature: float, temperature used in gumbel sampling.
+      EPSILON_GREEDY, GUMBEL_SAMPLING
+    temperature: float, temperature used in Gumbel sampling.
     eps: float, epsilon to use in epsilon greedy.
     len_history_for_policy: int, the maximum history to keep for applying the
       policy on. We also bucket observations on this number.
@@ -88,35 +87,6 @@ def play_env_problem_with_policy(env,
     trajectories is a list of triples of (observation, action, reward) ndarrays.
   """
 
-  def categorical_sample(log_probs):
-    """Categorical sampling."""
-
-    def multinomial_sample(probs):
-      """Sample from this vector of probabilities.
-
-      Args:
-        probs: numpy array of shape (A,) where A is the number of actions, these
-          must sum up to 1.0
-
-      Returns:
-        an integer of which action to pick.
-      """
-
-      return int(np.argwhere(np.random.multinomial(1, probs) == 1))
-
-    # Convert to probs, since we need to do categorical sampling.
-    probs = np.exp(log_probs)
-
-    # Let's cast up to float64, because that's what numpy does when sampling
-    # and it leads to the sum(pvals[:-1]) > 1.0 error.
-    #
-    # We also re-normalize when we do this.
-    probs = np.float64(probs)
-    probs /= np.sum(probs, axis=1, keepdims=True)
-
-    # Now pick actions from this probs array.
-    return np.apply_along_axis(multinomial_sample, 1, probs)
-
   def gumbel_sample(log_probs):
     """Gumbel sampling."""
     u = np.random.uniform(low=1e-6, high=1.0 - 1e-6, size=log_probs.shape)
@@ -177,9 +147,7 @@ def epsilon_greedy(log_probs):
         "B=%d, A=%d, log_probs.shape=%s" % (B, A, log_probs.shape)
 
     actions = None
-    if policy_sampling == CATEGORICAL_SAMPLING:
-      actions = categorical_sample(log_probs)
-    elif policy_sampling == GUMBEL_SAMPLING:
+    if policy_sampling == GUMBEL_SAMPLING:
       actions = gumbel_sample(log_probs)
     elif policy_sampling == EPSILON_GREEDY:
       actions = epsilon_greedy(log_probs)
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 832660774..de0cde9e1 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -132,7 +132,7 @@ def optimizer_fn(net_params, step_size=1e-3):
 def collect_trajectories(env,
                          policy_fn,
                          n_trajectories=1,
-                         policy=env_problem_utils.CATEGORICAL_SAMPLING,
+                         policy=env_problem_utils.GUMBEL_SAMPLING,
                          max_timestep=None,
                          epsilon=0.1,
                          reset=True,
@@ -716,6 +716,7 @@ def masked_entropy(log_probs, mask):
 
 def evaluate_policy(eval_env,
                     get_predictions,
+                    temperatures,
                     max_timestep=20000,
                     n_evals=1,
                     len_history_for_policy=32,
@@ -725,22 +726,20 @@ def evaluate_policy(eval_env,
   avg_rewards = collections.defaultdict(float)
   avg_rewards_unclipped = collections.defaultdict(float)
   for _ in range(n_evals):
-    for policy in [
-        env_problem_utils.CATEGORICAL_SAMPLING,
-        env_problem_utils.GUMBEL_SAMPLING,
-    ]:
+    for temperature in temperatures:
       trajs, _, _ = env_problem_utils.play_env_problem_with_policy(
           eval_env,
           get_predictions,
           num_trajectories=eval_env.batch_size,
           max_timestep=max_timestep,
           reset=True,
-          policy_sampling=policy,
+          policy_sampling=env_problem_utils.GUMBEL_SAMPLING,
+          temperature=temperature,
           rng=rng,
           len_history_for_policy=len_history_for_policy)
-      avg_rewards[policy] += float(sum(
+      avg_rewards[temperature] += float(sum(
           np.sum(traj[2]) for traj in trajs)) / len(trajs)
-      avg_rewards_unclipped[policy] += float(
+      avg_rewards_unclipped[temperature] += float(
           sum(np.sum(traj[3]) for traj in trajs)) / len(trajs)
 
   # Now average these out.
@@ -805,6 +804,7 @@ def training_loop(
     env_name=None,
     n_evals=1,
     len_history_for_policy=4,
+    eval_temperatures=(1.0, 0.5),
 ):
   """Runs the training loop for PPO, with fixed policy and value nets."""
   assert env
@@ -888,17 +888,20 @@ def get_predictions(observations, rng=None):
       avg_reward, avg_reward_unclipped = evaluate_policy(
           eval_env,
           get_predictions,
+          temperatures=eval_temperatures,
           max_timestep=max_timestep_eval,
           n_evals=n_evals,
           len_history_for_policy=len_history_for_policy,
           rng=key)
       for k, v in avg_reward.items():
         eval_sw.scalar("eval/mean_reward/%s" % k, v, step=i)
-        logging.info("Epoch [% 6d] Policy Evaluation (clipped) [%s] = %10.2f",
+        logging.info("Epoch [% 6d] Policy Evaluation (clipped) "
+                     "[temperature %s] = %10.2f",
                      i, k, v)
       for k, v in avg_reward_unclipped.items():
         eval_sw.scalar("eval/mean_reward_unclipped/%s" % k, v, step=i)
-        logging.info("Epoch [% 6d] Policy Evaluation (unclipped) [%s] = %10.2f",
+        logging.info("Epoch [% 6d] Policy Evaluation (unclipped) "
+                     "[temperature %s] = %10.2f",
                      i, k, v)
     policy_eval_time = get_time(policy_eval_start_time)
 

From 06144dfdd20978bacdd1cbf5cbedd5b128cbdf45 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 10 Jul 2019 12:52:15 -0700
Subject: [PATCH 2185/2720] Remove boilerplate by enforcing subclasses in
 Bayesian neural net layers.

A previous commit removed boilerplate in several classes by making the layers be subclasses of another. I'm still not sure I like the tradeoff of adding complexity via a dependence hierarchy. For now, I'm doing it here to enforce consistency across all layers.

PiperOrigin-RevId: 257462496
---
 tensor2tensor/layers/bayes.py | 157 ++++++++++++----------------------
 1 file changed, 56 insertions(+), 101 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 49dfa0f0f..32dcb25a8 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -198,8 +198,7 @@ def call(self, inputs):
     return outputs
 
 
-@add_weight
-class Conv2DVariationalDropout(tf.keras.layers.Conv2D):
+class Conv2DVariationalDropout(Conv2DReparameterization):
   """2D convolution layer with variational dropout (Kingma et al., 2015).
 
   Implementation follows the additive parameterization of
@@ -241,13 +240,6 @@ def __init__(self,
         bias_constraint=constraints.get(bias_constraint),
         **kwargs)
 
-  def call_weights(self):
-    """Calls any weights if the initializer is itself a layer."""
-    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
-      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
-    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
-
   def call(self, inputs, training=None):
     self.call_weights()
     if training is None:
@@ -321,7 +313,60 @@ def smart_constant_value(pred):
 
 
 @add_weight
-class DenseDVI(tf.keras.layers.Dense):
+class DenseReparameterization(tf.keras.layers.Dense):
+  """Bayesian densely-connected layer estimated via reparameterization.
+
+  The layer computes a variational Bayesian approximation to the distribution
+  over densely-connected layers,
+
+  ```
+  p(outputs | inputs) = int dense(inputs; weights, bias) p(weights, bias)
+    dweights dbias.
+  ```
+
+  It does this with a stochastic forward pass, sampling from learnable
+  distributions on the kernel and bias. Gradients with respect to the
+  distributions' learnable parameters backpropagate via reparameterization.
+  Minimizing cross-entropy plus the layer's losses performs variational
+  minimum description length, i.e., it minimizes an upper bound to the negative
+  marginal likelihood.
+  """
+
+  def __init__(self,
+               units,
+               activation=None,
+               use_bias=True,
+               kernel_initializer='trainable_normal',
+               bias_initializer='zero',
+               kernel_regularizer='normal_kl_divergence',
+               bias_regularizer=None,
+               activity_regularizer=None,
+               **kwargs):
+    super(DenseReparameterization, self).__init__(
+        units=units,
+        activation=activation,
+        use_bias=use_bias,
+        kernel_initializer=initializers.get(kernel_initializer),
+        bias_initializer=initializers.get(bias_initializer),
+        kernel_regularizer=regularizers.get(kernel_regularizer),
+        bias_regularizer=regularizers.get(bias_regularizer),
+        activity_regularizer=regularizers.get(activity_regularizer),
+        **kwargs)
+
+  def call_weights(self):
+    """Calls any weights if the initializer is itself a layer."""
+    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
+      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
+    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
+      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
+
+  def call(self, *args, **kwargs):
+    self.call_weights()
+    kwargs.pop('training', None)
+    return super(DenseReparameterization, self).call(*args, **kwargs)
+
+
+class DenseDVI(DenseReparameterization):
   """Densely-connected layer with deterministic VI (Wu et al., 2018).
 
   This layer computes a variational inference approximation via first and second
@@ -365,34 +410,6 @@ class DenseDVI(tf.keras.layers.Dense):
   ```
   """
 
-  def __init__(self,
-               units,
-               activation=None,
-               use_bias=True,
-               kernel_initializer='trainable_normal',
-               bias_initializer='zero',
-               kernel_regularizer='normal_kl_divergence',
-               bias_regularizer=None,
-               activity_regularizer=None,
-               **kwargs):
-    super(DenseDVI, self).__init__(
-        units=units,
-        activation=activation,
-        use_bias=use_bias,
-        kernel_initializer=initializers.get(kernel_initializer),
-        bias_initializer=initializers.get(bias_initializer),
-        kernel_regularizer=regularizers.get(kernel_regularizer),
-        bias_regularizer=regularizers.get(bias_regularizer),
-        activity_regularizer=regularizers.get(activity_regularizer),
-        **kwargs)
-
-  def call_weights(self):
-    """Calls any weights if the initializer is itself a layer."""
-    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
-      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
-    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
-
   def call(self, inputs):
     self.call_weights()
     if (not isinstance(inputs, ed.RandomVariable) and
@@ -486,60 +503,6 @@ def soft_relu(x):
           x * tfp.distributions.Normal(0., 1.).cdf(x))
 
 
-@add_weight
-class DenseReparameterization(tf.keras.layers.Dense):
-  """Bayesian densely-connected layer estimated via reparameterization.
-
-  The layer computes a variational Bayesian approximation to the distribution
-  over densely-connected layers,
-
-  ```
-  p(outputs | inputs) = int dense(inputs; weights, bias) p(weights, bias)
-    dweights dbias.
-  ```
-
-  It does this with a stochastic forward pass, sampling from learnable
-  distributions on the kernel and bias. Gradients with respect to the
-  distributions' learnable parameters backpropagate via reparameterization.
-  Minimizing cross-entropy plus the layer's losses performs variational
-  minimum description length, i.e., it minimizes an upper bound to the negative
-  marginal likelihood.
-  """
-
-  def __init__(self,
-               units,
-               activation=None,
-               use_bias=True,
-               kernel_initializer='trainable_normal',
-               bias_initializer='zero',
-               kernel_regularizer='normal_kl_divergence',
-               bias_regularizer=None,
-               activity_regularizer=None,
-               **kwargs):
-    super(DenseReparameterization, self).__init__(
-        units=units,
-        activation=activation,
-        use_bias=use_bias,
-        kernel_initializer=initializers.get(kernel_initializer),
-        bias_initializer=initializers.get(bias_initializer),
-        kernel_regularizer=regularizers.get(kernel_regularizer),
-        bias_regularizer=regularizers.get(bias_regularizer),
-        activity_regularizer=regularizers.get(activity_regularizer),
-        **kwargs)
-
-  def call_weights(self):
-    """Calls any weights if the initializer is itself a layer."""
-    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
-      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
-    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
-
-  def call(self, *args, **kwargs):
-    self.call_weights()
-    kwargs.pop('training', None)
-    return super(DenseReparameterization, self).call(*args, **kwargs)
-
-
 class DenseFlipout(DenseReparameterization):
   """Bayesian densely-connected layer estimated via Flipout (Wen et al., 2018).
 
@@ -599,8 +562,7 @@ def call(self, inputs):
     return outputs
 
 
-@add_weight
-class DenseVariationalDropout(tf.keras.layers.Dense):
+class DenseVariationalDropout(DenseReparameterization):
   """Densely-connected layer with variational dropout (Kingma et al., 2015).
 
   Implementation follows the additive parameterization of
@@ -628,13 +590,6 @@ def __init__(self,
         activity_regularizer=regularizers.get(activity_regularizer),
         **kwargs)
 
-  def call_weights(self):
-    """Calls any weights if the initializer is itself a layer."""
-    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
-      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
-    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
-
   def call(self, inputs, training=None):
     self.call_weights()
     if training is None:

From 3856200e65c32ed8da9c7d7745399abdf3a7e7f0 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 10 Jul 2019 17:46:26 -0700
Subject: [PATCH 2186/2720] Add noise contrastive priors.

Future changes will refactor the NCP experiments code.

PiperOrigin-RevId: 257519280
---
 tensor2tensor/layers/bayes.py      | 223 ++++++++++++++++++++++++++++-
 tensor2tensor/layers/bayes_test.py |  33 +++++
 2 files changed, 251 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 32dcb25a8..db65bbc31 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -262,12 +262,12 @@ def dropped_inputs():
       stddevs = tf.sqrt(
           self._convolution_op(tf.square(inputs), tf.exp(log_variance)) +
           tf.keras.backend.epsilon())
-      outputs = means + stddevs * tf.random_normal(tf.shape(stddevs))
       if self.use_bias:
         if self.data_format == 'channels_first':
-          outputs = tf.nn.bias_add(outputs, self.bias, data_format='NCHW')
+          means = tf.nn.bias_add(means, self.bias, data_format='NCHW')
         else:
-          outputs = tf.nn.bias_add(outputs, self.bias, data_format='NHWC')
+          means = tf.nn.bias_add(means, self.bias, data_format='NHWC')
+      outputs = ed.Normal(loc=means, scale=stddevs)
       if self.activation is not None:
         outputs = self.activation(outputs)
       return outputs
@@ -618,9 +618,9 @@ def dropped_inputs():
         stddevs = tf.sqrt(
             tf.tensordot(tf.square(inputs), tf.exp(log_variance), [[-1], [0]]) +
             tf.keras.backend.epsilon())
-      outputs = means + stddevs * tf.random_normal(tf.shape(stddevs))
       if self.use_bias:
-        outputs = tf.nn.bias_add(outputs, self.bias)
+        means = tf.nn.bias_add(means, self.bias)
+      outputs = ed.Normal(loc=means, scale=stddevs)
       if self.activation is not None:
         outputs = self.activation(outputs)
       return outputs
@@ -944,6 +944,219 @@ def call(self, inputs, states, training=None):
     return h, [h, c]
 
 
+class NCPNormalPerturb(tf.keras.layers.Layer):
+  """Noise contrastive prior for continuous inputs (Hafner et al., 2018).
+
+  The layer doubles the inputs' batch size and adds a random normal perturbation
+  to the concatenated second batch. This acts an input prior to be used in
+  combination with an output prior. The output prior reduces the second batch
+  (reverting to the inputs' original shape) and computes a regularizer that
+  matches the second batch towards some output (e.g., uniform distribution).
+  This layer implementation is inspired by the Aboleth library.
+
+  #### Examples
+
+  Below implements neural network regression with heteroskedastic noise,
+  noise contrastive priors, and being Bayesian only at the mean's output layer.
+
+  ```python
+  from tensor2tensor.layers import bayes
+
+  batch_size, dataset_size = 128, 1000
+  features, labels = get_some_dataset()
+
+  inputs = keras.Input(shape=(25,))
+  x = bayes.NCPNormalPerturb()(inputs)  # double input batch
+  x = layers.Dense(64, activation='relu')(x)
+  x = layers.Dense(64, activation='relu')(x)
+  means = bayes.DenseVariationalDropout(1, activation=None)(x)  # get mean dist.
+  means = bayes.NCPNormalOutput(labels)(means)  # halve input batch
+  stddevs = tf.keras.layers.Dense(1, activation='softplus')(x[:batch_size])
+  outputs = tf.keras.layers.Lambda(lambda x: ed.Normal(x[0], x[1]))([means,
+                                                                     stddevs])
+  model = tf.keras.Model(inputs=inputs, outputs=outputs)
+
+  predictions = model(features)
+  loss = tf.reduce_mean(predictions.distribution.log_prob(labels))
+  loss += model.losses[0] / dataset_size  # KL regularizer for output layer
+  loss += model.losses[-1]
+  train_op = tf.train.AdamOptimizer(0.1).minimize(loss)
+  ```
+
+  The network applies `bayes.NCPNormalPerturb()` to double the input batch
+  size and add Gaussian noise to the second half; then feedforward layers; then
+  `bayes.DenseVariational` to be Bayesian about the output density's mean; then
+  `bayes.NCPNormalOutput` centered at the labels to revert to the batch size
+  and compute a loss on the second half; then parameterize the output density's
+  standard deviations; then compute the total loss function as the sum of the
+  model's negative log-likelihood, KL divergence for the Bayesian mean layer,
+  and NCP loss.
+  """
+
+  def __init__(self, mean=0., stddev=1., seed=None, **kwargs):
+    self.mean = mean
+    self.stddev = stddev
+    self.seed = seed
+    super(NCPNormalPerturb, self).__init__(**kwargs)
+
+  def call(self, inputs):
+    noise = tf.random.normal(tf.shape(inputs),
+                             mean=self.mean,
+                             stddev=self.stddev,
+                             dtype=inputs.dtype,
+                             seed=self.seed)
+    perturbed_inputs = inputs + noise
+    return tf.concat([inputs, perturbed_inputs], 0)
+
+
+class NCPCategoricalPerturb(tf.keras.layers.Layer):
+  """Noise contrastive prior for discrete inputs (Hafner et al., 2018).
+
+  The layer doubles the inputs' batch size and randomly flips categories
+  for the concatenated second batch (all features must be integer-valued). This
+  acts an input prior to be used in combination with an output prior. The output
+  prior reduces the second batch (reverting to the inputs' original shape) and
+  computes a regularizer that matches the second batch towards some output
+  (e.g., uniform distribution). This layer implementation is inspired by the
+  Aboleth library.
+
+  #### Examples
+
+  Below implements neural network regression with heteroskedastic noise,
+  noise contrastive priors, and being Bayesian only at the mean's output layer.
+
+  ```python
+  from tensor2tensor.layers import bayes
+
+  batch_size, dataset_size = 128, 1000
+  features, labels = get_some_dataset()
+
+  inputs = keras.Input(shape=(25,))
+  x = bayes.NCPCategoricalPerturb(10)(inputs)  # double input batch
+  x = layers.Dense(64, activation='relu')(x)
+  x = layers.Dense(64, activation='relu')(x)
+  means = bayes.DenseVariationalDropout(1, activation=None)(x)  # get mean dist.
+  means = bayes.NCPNormalOutput(labels)(means)  # halve input batch
+  stddevs = tf.keras.layers.Dense(1, activation='softplus')(x[:batch_size])
+  outputs = tf.keras.layers.Lambda(lambda x: ed.Normal(x[0], x[1]))([means,
+                                                                     stddevs])
+  model = tf.keras.Model(inputs=inputs, outputs=outputs)
+
+  predictions = model(features)
+  loss = tf.reduce_mean(predictions.distribution.log_prob(labels))
+  loss += model.losses[0] / dataset_size  # KL regularizer for output layer
+  loss += model.losses[-1]
+  train_op = tf.train.AdamOptimizer(0.1).minimize(loss)
+  ```
+
+  The network applies `bayes.NCPCategoricalPerturb()` to double the input batch
+  size and flip categories for the second half; then feedforward layers; then
+  `bayes.DenseVariational` to be Bayesian about the output density's mean; then
+  `bayes.NCPNormalOutput` centered at the labels to revert to the batch size
+  and compute a loss on the second half; then parameterize the output density's
+  standard deviations; then compute the total loss function as the sum of the
+  model's negative log-likelihood, KL divergence for the Bayesian mean layer,
+  and NCP loss.
+  """
+
+  def __init__(self, input_dim, probs=0.1, **kwargs):
+    """Creates layer.
+
+    Args:
+      input_dim: int > 0. Size of the category, i.e. maximum integer index + 1.
+      probs: Probability that a category is randomly flipped.
+      **kwargs: kwargs to parent class.
+    """
+    self.input_dim = input_dim
+    self.probs = probs
+    super(NCPCategoricalPerturb, self).__init__(**kwargs)
+
+  def call(self, inputs):
+    mask = tf.cast(tf.random.uniform(tf.shape(inputs)) <= self.probs,
+                   inputs.dtype)
+    flips = tf.random.uniform(
+        tf.shape(inputs), minval=0, maxval=self.input_dim, dtype=inputs.dtype)
+    flipped_inputs = mask * flips + (1 - mask) * inputs
+    return tf.concat([inputs, flipped_inputs], 0)
+
+
+class NCPNormalOutput(tf.keras.layers.Layer):
+  """Noise contrastive prior for continuous outputs (Hafner et al., 2018).
+
+  The layer returns the first half of the inputs' batch. It computes a KL
+  regularizer as a side-effect, which matches the inputs' second half towards a
+  normal distribution (the output prior), and averaged over the number of inputs
+  in the second half. This layer is typically in combination with an input prior
+  which doubles the batch. This layer implementation is inspired by the Aboleth
+  library.
+
+  The layer computes the exact KL divergence from a normal distribution to
+  the input RandomVariable. It is an unbiased estimate if the input
+  RandomVariable has random parameters. If the input is a Tensor, then it
+  assumes its density is `ed.Normal(input, 1.)`, i.e., mean squared error loss.
+
+  #### Examples
+
+  Below implements neural network regression with heteroskedastic noise,
+  noise contrastive priors, and being Bayesian only at the mean's output layer.
+
+  ```python
+  from tensor2tensor.layers import bayes
+
+  batch_size, dataset_size = 128, 1000
+  features, labels = get_some_dataset()
+
+  inputs = keras.Input(shape=(25,))
+  x = bayes.NCPNormalPerturb()(inputs)  # double input batch
+  x = layers.Dense(64, activation='relu')(x)
+  x = layers.Dense(64, activation='relu')(x)
+  means = bayes.DenseVariationalDropout(1, activation=None)(x)  # get mean dist.
+  means = bayes.NCPNormalOutput(labels)(means)  # halve input batch
+  stddevs = tf.keras.layers.Dense(1, activation='softplus')(x[:batch_size])
+  outputs = tf.keras.layers.Lambda(lambda x: ed.Normal(x[0], x[1]))([means,
+                                                                     stddevs])
+  model = tf.keras.Model(inputs=inputs, outputs=outputs)
+
+  predictions = model(features)
+  loss = tf.reduce_mean(predictions.distribution.log_prob(labels))
+  loss += model.losses[0] / dataset_size  # KL regularizer for output layer
+  loss += model.losses[-1]
+  train_op = tf.train.AdamOptimizer(0.1).minimize(loss)
+  ```
+
+  The network applies `bayes.NCPNormalPerturb()` to double the input batch
+  size and add Gaussian noise to the second half; then feedforward layers; then
+  `bayes.DenseVariational` to be Bayesian about the output density's mean; then
+  `bayes.NCPNormalOutput` centered at the labels to revert to the batch size
+  and compute a loss on the second half; then parameterize the output density's
+  standard deviations; then compute the total loss function as the sum of the
+  model's negative log-likelihood, KL divergence for the Bayesian mean layer,
+  and NCP loss.
+  """
+
+  def __init__(self, mean=0., stddev=1., **kwargs):
+    self.mean = mean
+    self.stddev = stddev
+    super(NCPNormalOutput, self).__init__(**kwargs)
+
+  def call(self, inputs):
+    if not isinstance(inputs, ed.RandomVariable):
+      # Default to a unit normal, i.e., derived from mean squared error loss.
+      inputs = ed.Normal(loc=inputs, scale=1.)
+    batch_size = tf.shape(inputs)[0] // 2
+    # TODO(trandustin): Depend on github's ed2 for indexing RVs. This is a hack.
+    # _, _ = inputs[:batch_size], inputs[batch_size:]
+    original_inputs = ed.RandomVariable(inputs.distribution[:batch_size],
+                                        value=inputs.value[:batch_size])
+    perturbed_inputs = ed.RandomVariable(inputs.distribution[batch_size:],
+                                         value=inputs.value[batch_size:])
+    loss = tf.reduce_sum(
+        tfp.distributions.Normal(self.mean, self.stddev).kl_divergence(
+            perturbed_inputs.distribution)) / tf.to_float(batch_size)
+    self.add_loss(loss)
+    return original_inputs
+
+
 class MixtureLogistic(tf.keras.layers.Layer):
   """Stochastic output layer, distributed as a mixture of logistics."""
 
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 25a4b3b91..cf243326b 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -27,6 +27,7 @@
 
 import tensorflow as tf
 import tensorflow_probability as tfp
+ed = tfp.edward2
 tf.compat.v1.enable_eager_execution()
 
 
@@ -547,6 +548,38 @@ def testLSTMCellModel(self, lstm_cell):
     self.assertAllClose(res2, res3)
     self.assertLen(model.losses, 2)
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testNCPNormalPerturb(self):
+    batch_size = 3
+    inputs = tf.to_float(np.random.rand(batch_size, 4))
+    model = bayes.NCPNormalPerturb()
+    outputs = model(inputs)
+    inputs_val, outputs_val = self.evaluate([inputs, outputs])
+    self.assertEqual(outputs_val.shape, (2 * batch_size, 4))
+    self.assertAllEqual(inputs_val, outputs_val[:batch_size])
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testNCPCategoricalPerturb(self):
+    input_dim = 5
+    batch_size = 3
+    inputs = tf.to_float(np.random.choice(input_dim, size=(batch_size, 4)))
+    model = bayes.NCPCategoricalPerturb(input_dim)
+    outputs = model(inputs)
+    inputs_val, outputs_val = self.evaluate([inputs, outputs])
+    self.assertEqual(outputs_val.shape, (2 * batch_size, 4))
+    self.assertAllEqual(inputs_val, outputs_val[:batch_size])
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testNCPNormalOutput(self):
+    batch_size = 3
+    features = ed.Normal(loc=tf.random.normal([2 * batch_size, 1]), scale=1.)
+    labels = tf.to_float(np.random.rand(batch_size))
+    model = bayes.NCPNormalOutput(mean=labels)
+    predictions = model(features)
+    features_val, predictions_val = self.evaluate([features, predictions])
+    self.assertLen(model.losses, 1)
+    self.assertAllEqual(features_val[:batch_size], predictions_val)
+
   @test_utils.run_in_graph_and_eager_modes()
   def testMixtureLogistic(self):
     batch_size = 3

From 35fac6d51d0c24b286814662286bd7b126391769 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 10 Jul 2019 18:24:17 -0700
Subject: [PATCH 2187/2720] Report the standard deviation of rewards

PiperOrigin-RevId: 257523916
---
 tensor2tensor/trax/rlax/ppo.py | 76 +++++++++++++++++++++++-----------
 1 file changed, 51 insertions(+), 25 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index de0cde9e1..0877394c7 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -723,8 +723,8 @@ def evaluate_policy(eval_env,
                     rng=None):
   """Evaluate the policy."""
 
-  avg_rewards = collections.defaultdict(float)
-  avg_rewards_unclipped = collections.defaultdict(float)
+  processed_reward_sums = collections.defaultdict(list)
+  raw_reward_sums = collections.defaultdict(list)
   for _ in range(n_evals):
     for temperature in temperatures:
       trajs, _, _ = env_problem_utils.play_env_problem_with_policy(
@@ -737,17 +737,19 @@ def evaluate_policy(eval_env,
           temperature=temperature,
           rng=rng,
           len_history_for_policy=len_history_for_policy)
-      avg_rewards[temperature] += float(sum(
-          np.sum(traj[2]) for traj in trajs)) / len(trajs)
-      avg_rewards_unclipped[temperature] += float(
-          sum(np.sum(traj[3]) for traj in trajs)) / len(trajs)
-
-  # Now average these out.
-  for k in avg_rewards:
-    avg_rewards[k] /= n_evals
-    avg_rewards_unclipped[k] /= n_evals
-
-  return avg_rewards, avg_rewards_unclipped
+      processed_reward_sums[temperature].extend(sum(traj[2]) for traj in trajs)
+      raw_reward_sums[temperature].extend(sum(traj[3]) for traj in trajs)
+
+  # Return the mean and standard deviation for each temperature.
+  def compute_stats(reward_dict):
+    return {
+        temperature: {"mean": onp.mean(rewards), "std": onp.std(rewards)}
+        for (temperature, rewards) in reward_dict.items()
+    }
+  return {
+      "processed": compute_stats(processed_reward_sums),
+      "raw": compute_stats(raw_reward_sums),
+  }
 
 
 def maybe_restore_params(output_dir, policy_and_value_net_params):
@@ -778,6 +780,39 @@ def maybe_restore_params(output_dir, policy_and_value_net_params):
   return False, policy_and_value_net_params, 0
 
 
+def write_eval_reward_summaries(reward_stats_by_mode, summary_writer, epoch):
+  """Writes evaluation reward statistics to summary and logs them.
+
+  Args:
+    reward_stats_by_mode: Nested dict of structure:
+      {
+          "raw": {
+              <temperature 1>: {
+                  "mean": <reward mean>,
+                  "std": <reward std>,
+              },
+              <temperature 2>: ...
+          },
+          "processed": ...
+      }
+    summary_writer: jaxboard.SummaryWriter.
+    epoch: Current epoch number.
+  """
+  for (reward_mode, reward_stats_by_temp) in reward_stats_by_mode.items():
+    for (temperature, reward_stats) in reward_stats_by_temp.items():
+      for (stat_name, stat) in reward_stats.items():
+        summary_writer.scalar(
+            "eval/{reward_mode}_reward_{stat_name}/"
+            "temperature_{temperature}".format(reward_mode=reward_mode,
+                                               stat_name=stat_name,
+                                               temperature=temperature),
+            stat, step=epoch)
+      logging.info("Epoch [% 6d] Policy Evaluation (%s reward) "
+                   "[temperature %.2f] = %10.2f (+/- %.2f)",
+                   epoch, reward_mode, temperature,
+                   reward_stats["mean"], reward_stats["std"])
+
+
 def training_loop(
     env=None,
     epochs=EPOCHS,
@@ -885,7 +920,7 @@ def get_predictions(observations, rng=None):
 
       logging.vlog(1, "Epoch [% 6d] evaluating policy.", i)
 
-      avg_reward, avg_reward_unclipped = evaluate_policy(
+      reward_stats = evaluate_policy(
           eval_env,
           get_predictions,
           temperatures=eval_temperatures,
@@ -893,16 +928,7 @@ def get_predictions(observations, rng=None):
           n_evals=n_evals,
           len_history_for_policy=len_history_for_policy,
           rng=key)
-      for k, v in avg_reward.items():
-        eval_sw.scalar("eval/mean_reward/%s" % k, v, step=i)
-        logging.info("Epoch [% 6d] Policy Evaluation (clipped) "
-                     "[temperature %s] = %10.2f",
-                     i, k, v)
-      for k, v in avg_reward_unclipped.items():
-        eval_sw.scalar("eval/mean_reward_unclipped/%s" % k, v, step=i)
-        logging.info("Epoch [% 6d] Policy Evaluation (unclipped) "
-                     "[temperature %s] = %10.2f",
-                     i, k, v)
+      write_eval_reward_summaries(reward_stats, eval_sw, epoch=i)
     policy_eval_time = get_time(policy_eval_start_time)
 
     trajectory_collection_start_time = time.time()
@@ -926,7 +952,7 @@ def get_predictions(observations, rng=None):
     max_reward = max(np.sum(traj[2]) for traj in trajs)
     min_reward = min(np.sum(traj[2]) for traj in trajs)
 
-    train_sw.scalar("train/mean_reward", avg_reward, step=i)
+    train_sw.scalar("train/reward_mean_truncated", avg_reward, step=i)
 
     logging.vlog(1, "Rewards avg=[%0.2f], max=[%0.2f], min=[%0.2f], all=%s",
                  avg_reward, max_reward, min_reward,

From f8a6647a1257afa0b89c8fbac0e90898945a72c0 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 10 Jul 2019 18:46:31 -0700
Subject: [PATCH 2188/2720] Split the rng during policy evaluation

PiperOrigin-RevId: 257526359
---
 tensor2tensor/trax/rlax/ppo.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 0877394c7..301c3f322 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -725,7 +725,7 @@ def evaluate_policy(eval_env,
 
   processed_reward_sums = collections.defaultdict(list)
   raw_reward_sums = collections.defaultdict(list)
-  for _ in range(n_evals):
+  for eval_rng in jax_random.split(rng, num=n_evals):
     for temperature in temperatures:
       trajs, _, _ = env_problem_utils.play_env_problem_with_policy(
           eval_env,
@@ -735,7 +735,7 @@ def evaluate_policy(eval_env,
           reset=True,
           policy_sampling=env_problem_utils.GUMBEL_SAMPLING,
           temperature=temperature,
-          rng=rng,
+          rng=eval_rng,
           len_history_for_policy=len_history_for_policy)
       processed_reward_sums[temperature].extend(sum(traj[2]) for traj in trajs)
       raw_reward_sums[temperature].extend(sum(traj[3]) for traj in trajs)

From 5a4d76af1ed3dee7e6b4fcfd7e1440b19f17f5b9 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 11 Jul 2019 11:41:18 -0700
Subject: [PATCH 2189/2720] Enable different source/target vocabs in Trax
 Transformer.

PiperOrigin-RevId: 257650056
---
 tensor2tensor/trax/models/transformer.py      | 42 ++++++++++++-------
 tensor2tensor/trax/models/transformer_test.py | 23 +++++++---
 2 files changed, 45 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index b9f62ced9..8b821be53 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -229,8 +229,8 @@ def EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode):
   ]
 
 
-# TODO(lukaszkaiser): allow different source and target vocabularies.
-def Transformer(vocab_size,
+def Transformer(input_vocab_size,
+                output_vocab_size=None,
                 d_feature=512,
                 d_feedforward=2048,
                 n_layers=6,
@@ -243,7 +243,9 @@ def Transformer(vocab_size,
   This model expects an input pair: target, source.
 
   Args:
-    vocab_size: int: vocab size (shared source and target).
+    input_vocab_size: int: vocab size of the source.
+    output_vocab_size: int (optional): vocab size of the target. If None, the
+      source and target are assumed to have the same vocab.
     d_feature: int:  depth of embedding
     d_feedforward: int: depth of feed-forward layer
     n_layers: int: number of encoder/decoder layers
@@ -256,12 +258,22 @@ def Transformer(vocab_size,
     A Transformer model as a layer that maps from a target, source pair to
     activations over a vocab set.
   """
-  embed = [                                    # tokens
-      tl.Embedding(d_feature, vocab_size),     # vecs
-      tl.Dropout(rate=dropout, mode=mode),     # vecs
-      tl.PositionalEncoding(max_len=max_len),  # vecs
+  in_embed = [                                    # tokens
+      tl.Embedding(d_feature, input_vocab_size),  # vecs
+      tl.Dropout(rate=dropout, mode=mode),        # vecs
+      tl.PositionalEncoding(max_len=max_len),     # vecs
   ]
 
+  if output_vocab_size is None:
+    output_vocab_size = input_vocab_size
+    out_embed = in_embed
+  else:
+    out_embed = [                                    # tokens
+        tl.Embedding(d_feature, output_vocab_size),  # vecs
+        tl.Dropout(rate=dropout, mode=mode),         # vecs
+        tl.PositionalEncoding(max_len=max_len),      # vecs
+    ]
+
   encoder_stack = (  # masks vectors --> masks vectors
       [EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode)
        for _ in range(n_layers)])
@@ -275,21 +287,21 @@ def Transformer(vocab_size,
       tl.Swap(),    # toks_d toks_e
 
       # Encode.
-      tl.Parallel(                                    # toks_d        toks_e
-          [], [tl.Dup(),                              # ______ toks_e toks_e
-               tl.Parallel(embed, tl.PaddingMask()),  # ______ vecs_e masks
-               encoder_stack,                         # ______ vecs_e masks
-               tl.LayerNorm(),                        # ______ vecs_e .....
-               tl.Swap()]),                           # ______ masks vecs_e
+      tl.Parallel(                                       # toks_d        toks_e
+          [], [tl.Dup(),                                 # ______ toks_e toks_e
+               tl.Parallel(in_embed, tl.PaddingMask()),  # ______ vecs_e masks
+               encoder_stack,                            # ______ vecs_e masks
+               tl.LayerNorm(),                           # ______ vecs_e .....
+               tl.Swap()]),                              # ______ masks  vecs_e
 
       # Decode.                                  #        toks_d masks vecs_e
       tl.ShiftRight(),                           #        toks_d ..... ......
-      embed,                                     #        vecs_d ..... ......
+      out_embed,                                 #        vecs_d ..... ......
       tl.Dup(),                                  # vecs_d vecs_d ..... ......
       tl.Parallel([], tl.EncoderDecoderMask()),  # ______    masks     ......
       encoder_decoder_stack,                     # vecs_d    masks     vecs_e
       tl.Parallel([], tl.Drop(), tl.Drop()),     # vecs_d
       tl.LayerNorm(),                            # vecs_d
-      tl.Dense(vocab_size),                      # vecs_d
+      tl.Dense(output_vocab_size),               # vecs_d
       tl.LogSoftmax(),                           # vecs_d
   )
diff --git a/tensor2tensor/trax/models/transformer_test.py b/tensor2tensor/trax/models/transformer_test.py
index 5d49a4a5f..1ec6a3832 100644
--- a/tensor2tensor/trax/models/transformer_test.py
+++ b/tensor2tensor/trax/models/transformer_test.py
@@ -20,12 +20,13 @@
 from __future__ import print_function
 
 from absl.testing import absltest
+from absl.testing import parameterized
 from tensor2tensor.trax import backend
 from tensor2tensor.trax import layers as tl
 from tensor2tensor.trax.models import transformer
 
 
-class TransformerTest(absltest.TestCase):
+class TransformerTest(parameterized.TestCase):
 
   def test_transformer_lm_forward_shape(self):
     """Run the Transformer LM forward and check output shape."""
@@ -37,16 +38,28 @@ def test_transformer_lm_forward_shape(self):
         model, tuple(input_shape), integer_inputs=True)
     self.assertEqual(tuple(input_shape + [vocab_size]), final_shape)
 
-  def test_transformer_forward_shape(self):
+  def _test_transformer_forward_shape(self, input_vocab_size,
+                                      output_vocab_size):
     """Run the Transformer forward and check output shape."""
-    vocab_size = 16
     single_input_shape = [3, 5]
     input_shape = (tuple(single_input_shape), tuple(single_input_shape))
     model = transformer.Transformer(
-        vocab_size, d_feature=32, d_feedforward=64, n_layers=2, n_heads=2)
+        input_vocab_size, output_vocab_size,
+        d_feature=32, d_feedforward=64, n_layers=2, n_heads=2)
     final_shape = tl.check_shape_agreement(
         model, input_shape, integer_inputs=True)
-    self.assertEqual(tuple(single_input_shape + [vocab_size]), final_shape)
+    expected_shape = (tuple(single_input_shape +
+                            [output_vocab_size if output_vocab_size is not None
+                             else input_vocab_size]))
+    self.assertEqual(expected_shape, final_shape)
+
+  @parameterized.named_parameters(
+      ('same_vocab', 16, None),
+      ('same_size', 16, 16),
+      ('different_size', 16, 50))
+  def test_transformer_forward_shape(self, input_vocab_size, output_vocab_size):
+    """Run the Transformer forward and check output shape."""
+    self._test_transformer_forward_shape(input_vocab_size, output_vocab_size)
 
 
 if __name__ == '__main__':

From 7e99358a5dd44ddf5531dc8795fbc8a86a644ea2 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 11 Jul 2019 16:29:26 -0700
Subject: [PATCH 2190/2720] Add a flag --parallelize_envs that runs the envs in
 the EnvProblem in parallel.

If set to True, this uses the number of cpu cores, else runs the envs sequentially as in the current implementation.

Whether this actually speeds things up is indeterminate, rather for Atari this actually *slows* things down, probably because of the GIL.

Hopefully though for the `RPCEnv` and in processes where the step can take a long time, this should hopefully help a lot -- but this remains to be seen.

Let's check this in for now and try to fix it if the RPC code is slower.

PiperOrigin-RevId: 257704752
---
 tensor2tensor/envs/env_problem.py   | 40 +++++++++++++++++++----------
 tensor2tensor/trax/rlax/ppo_main.py |  8 ++++++
 2 files changed, 35 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index d608b4fc8..ebbf10d27 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -23,7 +23,9 @@
 from __future__ import division
 from __future__ import print_function
 
+import multiprocessing.pool
 import time
+
 import gym
 from gym.core import Env
 import numpy as np
@@ -131,6 +133,7 @@ def __init__(self,
                env_wrapper_fn=None,
                reward_range=(-np.inf, np.inf),
                discrete_rewards=True,
+               parallelism=1,
                **env_kwargs):
     """Initializes this class by creating the envs and managing trajectories.
 
@@ -146,6 +149,8 @@ def __init__(self,
         the raw reward in `process_rewards`.
       discrete_rewards: (bool) whether to round the rewards to the nearest
         integer.
+      parallelism: (int) If this is greater than one then we run the envs in
+        parallel using multi-threading.
       **env_kwargs: (dict) Additional kwargs to pass to the environments.
     """
 
@@ -176,6 +181,8 @@ def __init__(self,
     # be a Neural Network, in which case it will be fed input with first
     # dimension = `batch_size`.
     self._envs = None
+    self._pool = None
+    self._parallelism = parallelism
 
     self._observation_space = None
     self._action_space = None
@@ -254,7 +261,7 @@ def initialize(self, **kwargs):
     assert self._reward_range is not None
     assert self._trajectories is not None
 
-  def initialize_environments(self, batch_size=1, **env_kwargs):
+  def initialize_environments(self, batch_size=1, parallelism=1, **env_kwargs):
     """Initializes the environments and trajectories.
 
     Subclasses can override this if they don't want a default implementation
@@ -263,6 +270,8 @@ def initialize_environments(self, batch_size=1, **env_kwargs):
 
     Args:
       batch_size: (int) Number of `self.base_env_name` envs to initialize.
+      parallelism: (int) If this is greater than one then we run the envs in
+        parallel using multi-threading.
       **env_kwargs: (dict) Kwargs to pass to gym.make.
     """
     assert batch_size >= 1
@@ -271,6 +280,8 @@ def initialize_environments(self, batch_size=1, **env_kwargs):
     self._envs = [
         gym.make(self.base_env_name, **env_kwargs) for _ in range(batch_size)
     ]
+    self._parallelism = parallelism
+    self._pool = multiprocessing.pool.ThreadPool(self._parallelism)
     if self._env_wrapper_fn is not None:
       self._envs = list(map(self._env_wrapper_fn, self._envs))
 
@@ -564,22 +575,25 @@ def _step(self, actions):
     #               : len(actions) == len(self._envs)
     self.assert_common_preconditions()
     assert len(actions) == len(self._envs)
+    assert self.batch_size == len(actions)
 
-    observations = []
-    rewards = []
-    dones = []
-    infos = []
+    observations = [None] * self.batch_size
+    rewards = [None] * self.batch_size
+    dones = [None] * self.batch_size
+    infos = [{} for _ in range(self.batch_size)]
 
-    # Take steps in all environments.
-    for env, action in zip(self._envs, actions):
+    def apply_step(i):
       t1 = time.time()
-      observation, reward, done, info = env.step(action)
-      info["__bare_env_run_time__"] = time.time() - t1
+      observations[i], rewards[i], dones[i], infos[i] = self._envs[i].step(
+          actions[i])
+      t2 = time.time()
+      infos[i]["__bare_env_run_time__"] = t2 - t1
 
-      observations.append(observation)
-      rewards.append(reward)
-      dones.append(done)
-      infos.append(info)
+    if self._parallelism > 1:
+      self._pool.map(apply_step, range(self.batch_size))
+    else:
+      for i in range(self.batch_size):
+        apply_step(i)
 
     # Convert each list (observations, rewards, ...) into np.array and return a
     # tuple.
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index fc4f69c5d..998429d4c 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -40,6 +40,7 @@
 from __future__ import print_function
 
 import functools
+import multiprocessing
 import os
 
 from absl import app
@@ -142,6 +143,8 @@
                      "How much of history to give to the policy.")
 flags.DEFINE_bool("clip_rewards", True,
                   "Whether to clip and discretize the rewards.")
+flags.DEFINE_boolean("parallelize_envs", False,
+                     "If true, sets parallelism to number of cpu cores.")
 
 
 def common_layers():
@@ -164,11 +167,15 @@ def make_env(batch_size=8, **env_kwargs):
   else:
     env_kwargs.update({"discrete_rewards": False})
 
+  # TODO(afrozm): Should we leave out some cores?
+  parallelism = multiprocessing.cpu_count() if FLAGS.parallelize_envs else 1
+
   # No resizing needed, so let's be on the normal EnvProblem.
   if not FLAGS.resize:  # None or False
     return env_problem.EnvProblem(
         base_env_name=FLAGS.env_problem_name,
         batch_size=batch_size,
+        parallelism=parallelism,
         **env_kwargs)
 
   max_timestep = None
@@ -190,6 +197,7 @@ def make_env(batch_size=8, **env_kwargs):
   return rendered_env_problem.RenderedEnvProblem(
       base_env_name=FLAGS.env_problem_name,
       batch_size=batch_size,
+      parallelism=parallelism,
       env_wrapper_fn=wrapper_fn,
       **env_kwargs)
 

From d83ecb78205767a65f85e1f4f6a55689e30f6699 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 11 Jul 2019 18:53:52 -0700
Subject: [PATCH 2191/2720] Make attention key/value dimensions configurable

PiperOrigin-RevId: 257724784
---
 .../models/research/transformer_revnet.py     | 24 +++++++++++++------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index cb8f03e48..800d18f9b 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -631,13 +631,16 @@ def inverse_and_vjp(self, output, ct, params=(), **kwargs):
       return layer_val, None
 
 
-def DecoderBlock(d_feature, d_feedforward, n_heads, n_attention_chunks,
-                 attention_loop_stride, dropout, mode):
+def DecoderBlock(d_feature, d_feedforward, d_attention_key, d_attention_value,
+                 n_heads, n_attention_chunks, attention_loop_stride,
+                 dropout, mode):
   """Reversible transformer decoder layer.
 
   Args:
     d_feature: int:  depth of embedding
     d_feedforward: int: depth of feed-forward layer
+    d_attention_key: int: depth of key vector for each attention head
+    d_attention_value: int: depth of value vector for each attention head
     n_heads: int: number of attention heads
     n_attention_chunks: int: number of chunks for attention
     attention_loop_stride: int: number of query elements to compute attention
@@ -654,9 +657,9 @@ def DecoderBlock(d_feature, d_feedforward, n_heads, n_attention_chunks,
       tl.LayerNorm(),
       tl.Dup(), tl.Dup(),
       tl.Parallel(
-          [tl.Dense(d_feature), SplitHeads(n_heads=n_heads)],  # pylint: disable=no-value-for-parameter
-          [tl.Dense(d_feature), SplitHeads(n_heads=n_heads)],  # pylint: disable=no-value-for-parameter
-          [tl.Dense(d_feature), SplitHeads(n_heads=n_heads)],  # pylint: disable=no-value-for-parameter
+          [tl.Dense(d_attention_key * n_heads), SplitHeads(n_heads=n_heads)],  # pylint: disable=no-value-for-parameter
+          [tl.Dense(d_attention_key * n_heads), SplitHeads(n_heads=n_heads)],  # pylint: disable=no-value-for-parameter
+          [tl.Dense(d_attention_value * n_heads), SplitHeads(n_heads=n_heads)],  # pylint: disable=no-value-for-parameter
       ),
   ]
 
@@ -690,6 +693,8 @@ def DecoderBlock(d_feature, d_feedforward, n_heads, n_attention_chunks,
 def TransformerRevnetLM(vocab_size,
                         d_feature=512,
                         d_feedforward=2048,
+                        d_attention_key=64,
+                        d_attention_value=64,
                         n_layers=6,
                         n_heads=8,
                         dropout=0.1,
@@ -704,6 +709,8 @@ def TransformerRevnetLM(vocab_size,
     vocab_size: int: vocab size
     d_feature: int:  depth of *each half* of the two-part features
     d_feedforward: int: depth of feed-forward layer
+    d_attention_key: int: depth of key vector for each attention head
+    d_attention_value: int: depth of value vector for each attention head
     n_layers: int: number of decoder layers
     n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
@@ -728,8 +735,11 @@ def TransformerRevnetLM(vocab_size,
       positional_embedder,
       tl.Dup(),
       ReversibleSerial([
-          DecoderBlock(d_feature, d_feedforward, n_heads, n_attention_chunks,
-                       attention_loop_stride, dropout, mode)
+          # pylint: disable=g-complex-comprehension
+          DecoderBlock(d_feature, d_feedforward,
+                       d_attention_key, d_attention_value, n_heads,
+                       n_attention_chunks, attention_loop_stride,
+                       dropout, mode)
           for _ in range(n_layers)
       ]),
       tl.Parallel(tl.LayerNorm(), tl.LayerNorm()),

From eefb4b1987647f2ced9707453b5cb1d6a4196f47 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 11 Jul 2019 19:08:18 -0700
Subject: [PATCH 2192/2720] Update transformer_revnet_imagenet64_8gb.gin

PiperOrigin-RevId: 257726403
---
 .../transformer_revnet_imagenet64_8gb.gin      | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
index cfafd4b42..f8cb6b4d6 100644
--- a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
@@ -5,7 +5,7 @@ import tensor2tensor.trax.trax
 
 # Parameters for batch_fun:
 # ==============================================================================
-batch_fun.batch_size_per_device = 16
+batch_fun.batch_size_per_device = 8
 batch_fun.eval_batch_size = 8
 batch_fun.max_eval_length = 12288  # 64 * 64 * 3
 
@@ -14,18 +14,18 @@ batch_fun.max_eval_length = 12288  # 64 * 64 * 3
 inputs.data_dir = None
 inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
 inputs.input_name = 'targets'
-inputs.n_chunks = 64
+inputs.n_chunks = 16
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-MultifactorSchedule.constant = 0.3
+MultifactorSchedule.constant = 0.2
 MultifactorSchedule.factors = 'constant * linear_warmup'
-MultifactorSchedule.warmup_steps = 8000
+MultifactorSchedule.warmup_steps = 2000
 
 # Parameters for train:
 # ==============================================================================
-train.eval_frequency = 25
-train.eval_steps = 10
+train.eval_frequency = 100
+train.eval_steps = 8
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerRevnetLM
 train.optimizer = @trax.optimizers.SM3
@@ -36,12 +36,14 @@ train.trainer_class = @MemoryEfficientTrainer
 # ==============================================================================
 TransformerRevnetLM.d_feature = 512
 TransformerRevnetLM.d_feedforward = 2048
+TransformerRevnetLM.d_attention_key = 32
+TransformerRevnetLM.d_attention_value = 32
 TransformerRevnetLM.dropout = 0.1
 TransformerRevnetLM.max_len = 12288  # 64 * 64 * 3
 TransformerRevnetLM.mode = 'train'
-TransformerRevnetLM.n_heads = 8
+TransformerRevnetLM.n_heads = 4
 TransformerRevnetLM.n_layers = 6
 TransformerRevnetLM.vocab_size = 256
-TransformerRevnetLM.n_chunks = 64
+TransformerRevnetLM.n_chunks = 16
 TransformerRevnetLM.n_attention_chunks = 1
 TransformerRevnetLM.attention_loop_stride = 512

From d0ddbd4c38f05c6413da8108b663f8da5b08b41a Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 12 Jul 2019 08:14:01 -0700
Subject: [PATCH 2193/2720] add model_dir to hparams in decoder binary.

PiperOrigin-RevId: 257805253
---
 tensor2tensor/bin/t2t_decoder.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 88b45c230..3f9a56000 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -191,6 +191,10 @@ def main(_):
   if FLAGS.disable_grappler_optimizations:
     run_config.session_config.graph_options.rewrite_options.disable_meta_optimizer = True
 
+  # summary-hook in tf.estimator.EstimatorSpec requires
+  # hparams.model_dir to be set.
+  hp.add_hparam("model_dir", run_config.model_dir)
+
   estimator = trainer_lib.create_estimator(
       FLAGS.model,
       hp,

From 1a1e5bb3b173f6b45461e28c27103dfe750fea65 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 12 Jul 2019 10:15:52 -0700
Subject: [PATCH 2194/2720] Ginify ppo.training_loop

PiperOrigin-RevId: 257824097
---
 tensor2tensor/trax/rlax/configs/atari.gin     | 22 +++++
 .../online_tune_wide_resnet_cifar10.gin       | 24 ++++-
 tensor2tensor/trax/rlax/ppo.py                | 56 +++++++++---
 tensor2tensor/trax/rlax/ppo_main.py           | 87 +++----------------
 .../trax/rlax/ppo_training_loop_test.py       |  2 -
 5 files changed, 97 insertions(+), 94 deletions(-)
 create mode 100644 tensor2tensor/trax/rlax/configs/atari.gin

diff --git a/tensor2tensor/trax/rlax/configs/atari.gin b/tensor2tensor/trax/rlax/configs/atari.gin
new file mode 100644
index 000000000..d2f98901d
--- /dev/null
+++ b/tensor2tensor/trax/rlax/configs/atari.gin
@@ -0,0 +1,22 @@
+import tensor2tensor.trax.rlax
+
+# Parameters for ppo.training_loop:
+# ==============================================================================
+ppo.training_loop.epochs = 40000
+ppo.training_loop.n_optimizer_steps = 4
+ppo.training_loop.target_kl = 0.01
+ppo.training_loop.boundary = 20
+ppo.training_loop.max_timestep = 128
+ppo.training_loop.max_timestep_eval = 20000
+ppo.training_loop.random_seed = 0
+ppo.training_loop.gamma = 0.99
+ppo.training_loop.lambda_ = 0.95
+ppo.training_loop.epsilon = 0.1
+ppo.training_loop.c1 = 1.0
+ppo.training_loop.c2 = 0.01
+ppo.training_loop.eval_every_n = 500
+ppo.training_loop.done_frac_for_policy_save = 0.9
+ppo.training_loop.enable_early_stopping = False
+ppo.training_loop.n_evals = 16
+ppo.training_loop.len_history_for_policy = 4
+ppo.training_loop.eval_temperatures = (1.0, 0.5)
diff --git a/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
index 1d84a1b5f..688083821 100644
--- a/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
@@ -1,6 +1,7 @@
 import tensor2tensor.trax.inputs
 import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.rlax
 import tensor2tensor.trax.rlax.envs
 
 # Parameters for batch_fun:
@@ -25,11 +26,11 @@ shuffle_and_batch_data.preprocess_fun = @trax.inputs.cifar10_no_augmentation_pre
 
 # Parameters for WideResnet:
 # ==============================================================================
-WideResnet.d_hidden = 64
+WideResnet.widen_factor = 2
 WideResnet.n_blocks = 3
 WideResnet.n_output_classes = 10
 
-# Parameters for OnlineTune:
+# Parameters for OnlineTuneEnv:
 # ==============================================================================
 OnlineTuneEnv.inputs = @trax.inputs.inputs
 OnlineTuneEnv.model = @trax.models.WideResnet
@@ -38,3 +39,22 @@ OnlineTuneEnv.start_lr = 0.01
 OnlineTuneEnv.train_steps = 500
 OnlineTuneEnv.eval_steps = 100
 OnlineTuneEnv.env_steps = 100
+
+# Parameters for ppo.training_loop:
+# ==============================================================================
+ppo.training_loop.n_optimizer_steps = 30
+ppo.training_loop.boundary = 20
+ppo.training_loop.max_timestep = 128
+ppo.training_loop.max_timestep_eval = 20000
+ppo.training_loop.random_seed = 0
+ppo.training_loop.gamma = 0.99
+ppo.training_loop.lambda_ = 0.95
+ppo.training_loop.epsilon = 0.1
+ppo.training_loop.c1 = 1.0
+ppo.training_loop.c2 = 0.01
+ppo.training_loop.eval_every_n = 10
+ppo.training_loop.done_frac_for_policy_save = 0
+ppo.training_loop.enable_early_stopping = True
+ppo.training_loop.n_evals = 1
+ppo.training_loop.len_history_for_policy = 1
+ppo.training_loop.eval_temperatures = (1.0,)
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 301c3f322..076fb4b60 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -56,6 +56,7 @@
 
 from absl import logging
 import cloudpickle as pickle
+import gin
 import gym
 from jax import grad
 from jax import jit
@@ -813,12 +814,15 @@ def write_eval_reward_summaries(reward_stats_by_mode, summary_writer, epoch):
                    reward_stats["mean"], reward_stats["std"])
 
 
+@gin.configurable(blacklist=["output_dir"])
 def training_loop(
-    env=None,
+    env,
+    eval_env,
+    env_name,
+    policy_and_value_net_fn,
+    policy_and_value_optimizer_fn,
+    output_dir,
     epochs=EPOCHS,
-    policy_and_value_net_fn=None,
-    policy_and_value_optimizer_fn=None,
-    batch_size=BATCH_TRAJECTORIES,
     n_optimizer_steps=N_OPTIMIZER_STEPS,
     print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
     target_kl=0.01,
@@ -831,21 +835,47 @@ def training_loop(
     epsilon=EPSILON,
     c1=1.0,
     c2=0.01,
-    output_dir=None,
     eval_every_n=1000,
-    eval_env=None,
     done_frac_for_policy_save=0.5,
     enable_early_stopping=True,
-    env_name=None,
     n_evals=1,
     len_history_for_policy=4,
     eval_temperatures=(1.0, 0.5),
 ):
-  """Runs the training loop for PPO, with fixed policy and value nets."""
-  assert env
-  assert output_dir
-  assert env_name
+  """Runs the training loop for PPO, with fixed policy and value nets.
 
+  Args:
+    env: gym.Env to use for training.
+    eval_env: gym.Env to use for evaluation.
+    env_name: Name of the environment.
+    policy_and_value_net_fn: Function defining the policy and value network.
+    policy_and_value_optimizer_fn: Function defining the optimizer.
+    output_dir: Output dir.
+    epochs: Number of epochs to run for.
+    n_optimizer_steps: Number of optimizer steps.
+    print_every_optimizer_steps: How often to log during the policy optimization
+      process.
+    target_kl: Policy iteration early stopping.
+    boundary: We pad trajectories at integer multiples of this number.
+    max_timestep: If set to an integer, maximum number of time-steps in
+      a trajectory. Used in the collect procedure.
+    max_timestep_eval: If set to an integer, maximum number of time-steps in an
+      evaluation trajectory. Used in the collect procedure.
+    random_seed: Random seed.
+    gamma: Reward discount factor.
+    lambda_: N-step TD-error discount factor in GAE.
+    epsilon: Random action probability in epsilon-greedy sampling.
+    c1: Value loss coefficient.
+    c2: Entropy loss coefficient.
+    eval_every_n: How frequently to eval the policy.
+    done_frac_for_policy_save: Fraction of the trajectories that should be done
+      to checkpoint the policy.
+    enable_early_stopping: Whether to enable early stopping.
+    n_evals: Number of times to evaluate.
+    len_history_for_policy: How much of history to give to the policy.
+    eval_temperatures: Sequence of temperatures to try for categorical sampling
+      during evaluation.
+  """
   gfile.makedirs(output_dir)
 
   # Create summary writers and history.
@@ -937,7 +967,7 @@ def get_predictions(observations, rng=None):
     trajs, n_done, timing_info = collect_trajectories(
         env,
         policy_fn=get_predictions,
-        n_trajectories=batch_size,
+        n_trajectories=env.batch_size,
         max_timestep=max_timestep,
         rng=key,
         len_history_for_policy=len_history_for_policy,
@@ -1124,7 +1154,7 @@ def get_predictions(observations, rng=None):
     policy_save_start_time = time.time()
     n_trajectories_done += n_done
     # TODO(afrozm): Refactor to trax.save_state.
-    if (((n_trajectories_done >= done_frac_for_policy_save * batch_size) and
+    if (((n_trajectories_done >= done_frac_for_policy_save * env.batch_size) and
          (i - last_saved_at > eval_every_n) and
          (((i + 1) % eval_every_n == 0))) or (i == epochs - 1)):
       logging.vlog(1, "Epoch [% 6d] saving model.", i)
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 998429d4c..167c34464 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -18,19 +18,18 @@
 Sample invocation:
 
 ENV_PROBLEM_NAME=Acrobot-v1
-COMBINED_NETWORK=false
-EPOCHS=100
 BATCH_SIZE=32
+EPOCHS=100
 RANDOM_SEED=0
 BOUNDARY=100
 
 python trax/rlax/ppo_main.py \
   --env_problem_name=${ENV_PROBLEM_NAME} \
-  --combined_policy_and_value_function=${COMBINED_NETWORK} \
-  --epochs=${EPOCHS} \
   --batch_size=${BATCH_SIZE} \
-  --random_seed=${RANDOM_SEED} \
-  --boundary=${BOUNDARY} \
+  --config=ppo.training_loop.epochs=${EPOCHS} \
+  --config=ppo.training_loop.random_seed=${RANDOM_SEED} \
+  --config=ppo.training_loop.boundary=${BOUNDARY} \
+  --output_dir=${HOME}/ppo_acrobot \
   --vmodule=*/tensor2tensor/*=1 \
   --alsologtostderr \
 """
@@ -63,12 +62,6 @@
 
 flags.DEFINE_string("env_problem_name", None, "Name of the EnvProblem to make.")
 
-flags.DEFINE_integer("epochs", 100, "Number of epochs to run for.")
-flags.DEFINE_string("random_seed", None, "Random seed.")
-flags.DEFINE_integer("batch_size", 32, "Batch of trajectories needed.")
-
-flags.DEFINE_integer(
-    "boundary", 20, "We pad trajectories at integer multiples of this number.")
 # -1: returns env as is.
 # None: unwraps and returns without TimeLimit wrapper.
 # Any other number: imposes this restriction.
@@ -77,17 +70,6 @@
     "If set to an integer, maximum number of time-steps in a "
     "trajectory. The bare env is wrapped with TimeLimit wrapper.")
 
-# This is different from max_timestep is that in the above, the env is wrapped
-# in a TimeLimit wrapper, vs here we use this in the collect function.
-flags.DEFINE_integer(
-    "truncation_timestep", None,
-    "If set to an integer, maximum number of time-steps in a "
-    "trajectory. Used in the collect procedure.")
-flags.DEFINE_integer(
-    "truncation_timestep_eval", 20000,
-    "If set to an integer, maximum number of time-steps in an evaluation "
-    "trajectory. Used in the collect procedure.")
-
 flags.DEFINE_boolean(
     "jax_debug_nans", False,
     "Setting to true will help to debug nans and disable jit.")
@@ -104,43 +86,19 @@
     "In the combined network case should we make one tower or"
     "two.")
 
-# Number of optimizer steps of the combined net, policy net and value net.
-flags.DEFINE_integer("n_optimizer_steps", 100, "Number of optimizer steps.")
-flags.DEFINE_integer(
-    "print_every_optimizer_steps", 1,
-    "How often to log during the policy optimization process.")
-
 # Learning rate of the combined net, policy net and value net.
 flags.DEFINE_float("learning_rate", 1e-3, "Learning rate.")
 
-# Target KL is used for doing early stopping in the
-flags.DEFINE_float("target_kl", 0.01, "Policy iteration early stopping")
-flags.DEFINE_float("value_coef", 1.0,
-                   "Coefficient of Value Loss term in combined loss.")
-flags.DEFINE_float("entropy_coef", 0.01,
-                   "Coefficient of the Entropy Bonus term in combined loss.")
-flags.DEFINE_float("gamma", 0.99, "Policy iteration early stopping")
-flags.DEFINE_float("lambda_", 0.95, "Policy iteration early stopping")
-flags.DEFINE_float("epsilon", 0.1, "Policy iteration early stopping")
-
 flags.DEFINE_string("output_dir", "", "Output dir.")
 flags.DEFINE_multi_string("config_file", None,
                           "Configuration file with parameters (.gin).")
 flags.DEFINE_multi_string("config", None,
                           "Configuration parameters (gin string).")
 flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
-flags.DEFINE_bool("enable_early_stopping", True,
-                  "Whether to enable early stopping.")
 flags.DEFINE_bool("xm", False, "Copy atari roms?")
-flags.DEFINE_integer("eval_every_n", 100, "How frequently to eval the policy.")
+flags.DEFINE_integer("batch_size", 32,
+                     "Number of parallel environments during training.")
 flags.DEFINE_integer("eval_batch_size", 4, "Batch size for evaluation.")
-flags.DEFINE_integer("n_evals", 1, "Number of times to evaluate.")
-flags.DEFINE_float(
-    "done_frac_for_policy_save", 0.5,
-    "Fraction of the trajectories that should be done to "
-    "checkpoint the policy.")
-flags.DEFINE_integer("len_history_for_policy", 4,
-                     "How much of history to give to the policy.")
 flags.DEFINE_bool("clip_rewards", True,
                   "Whether to clip and discretize the rewards.")
 flags.DEFINE_boolean("parallelize_envs", False,
@@ -246,38 +204,13 @@ def run_training_loop():
         two_towers=FLAGS.two_towers)
     policy_and_value_optimizer_fn = get_optimizer_fn(FLAGS.learning_rate)
 
-    random_seed = None
-    try:
-      random_seed = int(FLAGS.random_seed)
-    except Exception:  # pylint: disable=broad-except
-      pass
-
     ppo.training_loop(
-        env=env,
-        epochs=FLAGS.epochs,
-        policy_and_value_net_fn=policy_and_value_net_fn,
-        policy_and_value_optimizer_fn=policy_and_value_optimizer_fn,
-        n_optimizer_steps=FLAGS.n_optimizer_steps,
-        print_every_optimizer_steps=FLAGS.print_every_optimizer_steps,
-        batch_size=FLAGS.batch_size,
-        target_kl=FLAGS.target_kl,
-        boundary=FLAGS.boundary,
-        max_timestep=FLAGS.truncation_timestep,
-        max_timestep_eval=FLAGS.truncation_timestep_eval,
-        random_seed=random_seed,
-        c1=FLAGS.value_coef,
-        c2=FLAGS.entropy_coef,
-        gamma=FLAGS.gamma,
-        lambda_=FLAGS.lambda_,
-        epsilon=FLAGS.epsilon,
-        enable_early_stopping=FLAGS.enable_early_stopping,
         output_dir=FLAGS.output_dir,
-        eval_every_n=FLAGS.eval_every_n,
-        done_frac_for_policy_save=FLAGS.done_frac_for_policy_save,
+        env=env,
         eval_env=eval_env,
-        n_evals=FLAGS.n_evals,
         env_name=str(FLAGS.env_problem_name),
-        len_history_for_policy=int(FLAGS.len_history_for_policy),
+        policy_and_value_net_fn=policy_and_value_net_fn,
+        policy_and_value_optimizer_fn=policy_and_value_optimizer_fn,
     )
 
   if FLAGS.jax_debug_nans or FLAGS.disable_jit:
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index 05ef3a797..580da77b0 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -68,7 +68,6 @@ def _run_training_loop(self, env_name, output_dir):
     env = self.get_wrapped_env(env_name, 2)
     eval_env = self.get_wrapped_env(env_name, 2)
     n_epochs = 2
-    batch_size = 2
     # Run the training loop.
     ppo.training_loop(
         env=env,
@@ -78,7 +77,6 @@ def _run_training_loop(self, env_name, output_dir):
             ppo.policy_and_value_net,
             bottom_layers_fn=lambda: [layers.Dense(1)]),
         policy_and_value_optimizer_fn=ppo.optimizer_fn,
-        batch_size=batch_size,
         n_optimizer_steps=1,
         output_dir=output_dir,
         env_name=env_name,

From 80f7da73cdcc685d22f9fcefb7d9227f29cf9cd2 Mon Sep 17 00:00:00 2001
From: Peng Wang <wangpeng@google.com>
Date: Fri, 12 Jul 2019 15:14:21 -0700
Subject: [PATCH 2195/2720] tf-numpy improvements to enable trax/resnet

PiperOrigin-RevId: 257877879
---
 tensor2tensor/trax/backend.py              | 89 ++++++++++++++++++++++
 tensor2tensor/trax/inputs.py               | 11 +--
 tensor2tensor/trax/layers/combinators.py   | 20 +++--
 tensor2tensor/trax/layers/convolution.py   |  8 +-
 tensor2tensor/trax/layers/normalization.py | 22 ++++--
 tensor2tensor/trax/layers/pooling.py       | 36 ++-------
 tensor2tensor/trax/models/resnet_test.py   |  2 +
 tensor2tensor/trax/optimizers/base.py      | 27 +++++--
 tensor2tensor/trax/trax.py                 | 10 ++-
 9 files changed, 162 insertions(+), 63 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index 6e2aa8967..dbb251bb6 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -23,17 +23,85 @@
 import gin
 
 import jax
+from jax import lax
 from jax import random as jax_random
 import jax.numpy as jnp
 import jax.scipy.special as jax_special
 import numpy as onp
+import tensorflow_datasets as tfds
 
 
 
+def jax_conv(inp, fltr, window_strides, padding, dimension_numbers,
+             filter_dilation=None):
+  """A wrapper around `lax.conv_general_dilated`.
+
+  It requires `dimension_numbers` and disallows `inp_dilation`.
+
+  Args:
+    inp: an (N+2)-D array. The input of the convolution.
+    fltr: an (N+2)-D array. The filter (i.e. kernel) of the convolution.
+    window_strides: the strides for moving the convolution window.
+    padding: a string, either "VALID" or "SAME". The padding algorithm.
+    dimension_numbers: a tuple of three strings encoding the data format of
+      input, filter and output. "I" means input; "O" means output; "C" means
+      channel; other characters such as "W", "H" and "D" means spatial
+      dimensions.
+    filter_dilation: the dilation rates for the filter. Dilating the filter
+      means adding "holes" to the filter.
+
+  Returns:
+    An (N+2)-D array. The convolution result.
+  """
+  return lax.conv_general_dilated(inp, fltr, window_strides, padding,
+                                  lhs_dilation=None,
+                                  rhs_dilation=filter_dilation,
+                                  dimension_numbers=dimension_numbers)
+
+
+def _pooling_general(inputs, reducer, init_val, rescaler=None,
+                     pool_size=(2, 2), strides=None, padding="VALID"):
+  """Helper: general pooling computation used in pooling layers later."""
+  spatial_strides = strides or (1,) * len(pool_size)
+  rescale = rescaler(pool_size, spatial_strides, padding) if rescaler else None
+  dims = (1,) + pool_size + (1,)  # NHWC
+  strides = (1,) + spatial_strides + (1,)
+  out = lax.reduce_window(inputs, init_val, reducer, dims, strides, padding)
+  return rescale(out, inputs) if rescale else out
+
+
+def jax_max_pool(x, pool_size, strides, padding):
+  return _pooling_general(x, lax.max, -jnp.inf, pool_size=pool_size,
+                          strides=strides, padding=padding)
+
+
+def jax_sum_pool(x, pool_size, strides, padding):
+  return _pooling_general(x, lax.add, 0., pool_size=pool_size,
+                          strides=strides, padding=padding)
+
+
+def _normalize_by_window_size(dims, spatial_strides, padding):  # pylint: disable=invalid-name
+  def rescale(outputs, inputs):
+    one = jnp.ones(inputs.shape[1:-1], dtype=inputs.dtype)
+    window_sizes = lax.reduce_window(
+        one, 0., lax.add, dims, spatial_strides, padding)
+    return outputs / window_sizes[..., jnp.newaxis]
+  return rescale
+
+
+def jax_avg_pool(x, pool_size, strides, padding):
+  return _pooling_general(x, lax.add, 0., _normalize_by_window_size,
+                          pool_size, strides=strides, padding=padding)
+
+
 _JAX_BACKEND = {
     "name": "jax",
     "np": jnp,
     "logsumexp": jax_special.logsumexp,
+    "conv": jax_conv,
+    "avg_pool": jax_avg_pool,
+    "max_pool": jax_max_pool,
+    "sum_pool": jax_sum_pool,
     "jit": jax.jit,
     "grad": jax.grad,
     "pmap": jax.pmap,
@@ -42,6 +110,7 @@
     "random_bernoulli": jax_random.bernoulli,
     "random_get_prng": jax.jit(jax_random.PRNGKey),
     "random_split": jax_random.split,
+    "dataset_as_numpy": tfds.as_numpy,
 }
 
 
@@ -60,6 +129,22 @@ def logsumexp(*args, **kwargs):
   return backend()["logsumexp"](*args, **kwargs)
 
 
+def conv(*args, **kwargs):
+  return backend()["conv"](*args, **kwargs)
+
+
+def avg_pool(*args, **kwargs):
+  return backend()["avg_pool"](*args, **kwargs)
+
+
+def max_pool(*args, **kwargs):
+  return backend()["max_pool"](*args, **kwargs)
+
+
+def sum_pool(*args, **kwargs):
+  return backend()["sum_pool"](*args, **kwargs)
+
+
 def jit(*args, **kwargs):
   return backend()["jit"](*args, **kwargs)
 
@@ -72,6 +157,10 @@ def pmap(*args, **kwargs):
   return backend()["pmap"](*args, **kwargs)
 
 
+def dataset_as_numpy(*args, **kwargs):
+  return backend()["dataset_as_numpy"](*args, **kwargs)
+
+
 # For numpy and random modules, we need to call "backend()" lazily, only when
 # the function is called -- so that it can be set by gin configs.
 # (Otherwise, backend() is called on import before gin-config is parsed.)
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 93793bcba..381954de7 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -24,10 +24,11 @@
 import random
 
 import gin
-
-import numpy as np
+import numpy as onp
 
 from tensor2tensor import problems_colab as t2t_problems
+from tensor2tensor.trax import backend
+from tensor2tensor.trax.backend import numpy as np
 import tensorflow as tf
 import tensorflow_datasets as tfds
 
@@ -139,9 +140,9 @@ def random_inputs(
   def random_minibatches():
     """Generate a stream of random mini-batches."""
     if input_dtype in [np.float16, np.float32, np.float64]:
-      rand = np.random.uniform
+      rand = onp.random.uniform
     else:
-      rand = np.random.random_integers
+      rand = onp.random.random_integers
     while True:
       inp = rand(input_range[0], input_range[1], input_shape)
       inp = inp.astype(input_dtype)
@@ -159,7 +160,7 @@ def random_minibatches():
 
 def dataset_to_stream(dataset, input_name, n_chunks=0, append_targets=False):
   """Takes a tf.Dataset and creates a numpy stream of ready batches."""
-  for example in tfds.as_numpy(dataset):
+  for example in backend.dataset_as_numpy(dataset):
     inp, out = example[0][input_name], example[1]
     # Some accelerators don't handle uint8 well, cast to int.
     if isinstance(inp, np.uint8):
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index d2d83268e..c3da0e3cb 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 from tensor2tensor.trax import backend
+from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers import base
 
 
@@ -161,7 +162,8 @@ def _validate_call_inputs(self, xs):
     if not isinstance(xs, tuple) and self._n_inputs != 1:
       raise TypeError(
           'Serial.call input must be a tuple; instead got {}'.format(xs))
-    if len(xs) < self.n_inputs():
+    len_xs = 1 if isinstance(xs, np.ndarray) else len(xs)
+    if len_xs < self.n_inputs():
       raise ValueError(
           'number of inputs ({}) to Serial.call less than n_inputs'
           ' ({})'.format(len(xs), self.n_inputs()))
@@ -361,16 +363,18 @@ class Parallel(base.Layer):
 
   creates a layer that passes its first input unchanged and applies F to the
   following input(s).
-
-  Args:
-    *layers: A list of layers.
-
-  Returns:
-    A new layer in which each of the given layers applies to its corresponding
-    span of elements in the dataflow stack.
   """
 
   def __init__(self, *layers):
+    """The constructor.
+
+    Args:
+      *layers: A list of layers.
+
+    Returns:
+      A new layer in which each of the given layers applies to its corresponding
+      span of elements in the dataflow stack.
+    """
     super(Parallel, self).__init__()
     layers = self._validate(layers)
     self._n_layers = len(layers)
diff --git a/tensor2tensor/trax/layers/convolution.py b/tensor2tensor/trax/layers/convolution.py
index bd388a2ec..c30669af9 100644
--- a/tensor2tensor/trax/layers/convolution.py
+++ b/tensor2tensor/trax/layers/convolution.py
@@ -22,9 +22,9 @@
 import itertools
 import operator
 
-from jax import lax
 import six
 
+from tensor2tensor.trax import backend
 from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers import base
 from tensor2tensor.trax.layers import initializers as init
@@ -64,9 +64,9 @@ def call(self, x, params=(), **kwargs):
       self._check_nhwc()
       new_batch_dim = six.moves.reduce(operator.mul, x_shape[:-3])
       x = np.reshape(x, [new_batch_dim] + x_shape[-3:])
-    res = lax.conv_general_dilated(
-        x, w, self._strides, self._padding, self._one, self._one,
-        self._dimension_numbers) + b
+    res = backend.conv(
+        x, w, self._strides, self._padding, self._dimension_numbers,
+        self._one) + b
     if len(x_shape) > 4:
       res = np.reshape(res, x_shape[:-3] + list(res.shape[-3:]))
     return res
diff --git a/tensor2tensor/trax/layers/normalization.py b/tensor2tensor/trax/layers/normalization.py
index c42dd5f0f..d6b07cc23 100644
--- a/tensor2tensor/trax/layers/normalization.py
+++ b/tensor2tensor/trax/layers/normalization.py
@@ -43,7 +43,9 @@ def BatchNorm(x, params, axis=(0, 1, 2), epsilon=1e-5,
   # Fast but less numerically-stable variance calculation than np.var.
   m1 = np.mean(x**2, axis, keepdims=True)
   var = m1 - mean**2
-  z = (x - mean) / np.sqrt(var + epsilon)
+  # x mustn't be onp.ndarray here; otherwise `x-mean` will call mean.__rsub__
+  # with each element of x, resulting in an onp.ndarray with dtype `object`.
+  z = (x - mean) / np.sqrt(var + epsilon).astype(x.dtype)
 
   # Expand the parameters to have the right axes.
   beta, gamma = params
@@ -55,12 +57,18 @@ def BatchNorm(x, params, axis=(0, 1, 2), epsilon=1e-5,
 
   # Return the z rescaled by the parameters if requested.
   if center and scale:
-    return gamma * z + beta
-  if center:
-    return z + beta
-  if scale:
-    return gamma * z
-  return z
+    ret = gamma * z + beta
+  elif center:
+    ret = z + beta
+  elif scale:
+    ret = gamma * z
+  else:
+    ret = z
+  assert ret.dtype == x.dtype, ('The dtype of the output (%s) of batch norm is '
+                                'not the same as the input (%s). Batch norm '
+                                'should not change the dtype' %
+                                (ret.dtype, x.dtype))
+  return ret
 
 
 # Layer normalization.
diff --git a/tensor2tensor/trax/layers/pooling.py b/tensor2tensor/trax/layers/pooling.py
index 3c07fd5fb..3782491cf 100644
--- a/tensor2tensor/trax/layers/pooling.py
+++ b/tensor2tensor/trax/layers/pooling.py
@@ -19,48 +19,26 @@
 from __future__ import division
 from __future__ import print_function
 
-from jax import lax
-
-from tensor2tensor.trax.backend import numpy as np
+from tensor2tensor.trax import backend
 from tensor2tensor.trax.layers import base
 
 
-def PoolingGeneral(inputs, reducer, init_val, rescaler=None,
-                   pool_size=(2, 2), strides=None, padding='VALID'):
-  """Helper: general pooling computation used in pooling layers later."""
-  spatial_strides = strides or (1,) * len(pool_size)
-  rescale = rescaler(pool_size, spatial_strides, padding) if rescaler else None
-  dims = (1,) + pool_size + (1,)  # NHWC
-  strides = (1,) + spatial_strides + (1,)
-  out = lax.reduce_window(inputs, init_val, reducer, dims, strides, padding)
-  return rescale(out, inputs) if rescale else out
-
-
 @base.layer()
 def MaxPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
   del params, kw
-  return PoolingGeneral(x, lax.max, -np.inf, pool_size=pool_size,
-                        strides=strides, padding=padding)
+  return backend.max_pool(x, pool_size=pool_size, strides=strides,
+                          padding=padding)
 
 
 @base.layer()
 def SumPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
   del params, kw
-  return PoolingGeneral(x, lax.add, 0., pool_size=pool_size,
-                        strides=strides, padding=padding)
-
-
-def _normalize_by_window_size(dims, spatial_strides, padding):  # pylint: disable=invalid-name
-  def Rescale(outputs, inputs):
-    one = np.ones(inputs.shape[1:-1], dtype=inputs.dtype)
-    window_sizes = lax.reduce_window(
-        one, 0., lax.add, dims, spatial_strides, padding)
-    return outputs / window_sizes[..., np.newaxis]
-  return Rescale
+  return backend.sum_pool(x, pool_size=pool_size, strides=strides,
+                          padding=padding)
 
 
 @base.layer()
 def AvgPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
   del params, kw
-  return PoolingGeneral(x, lax.add, 0., _normalize_by_window_size,
-                        pool_size, strides=strides, padding=padding)
+  return backend.avg_pool(x, pool_size=pool_size, strides=strides,
+                          padding=padding)
diff --git a/tensor2tensor/trax/models/resnet_test.py b/tensor2tensor/trax/models/resnet_test.py
index de36a6757..4751ac9a8 100644
--- a/tensor2tensor/trax/models/resnet_test.py
+++ b/tensor2tensor/trax/models/resnet_test.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 from absl.testing import absltest
+from tensor2tensor.trax import backend
 from tensor2tensor.trax import layers as tl
 from tensor2tensor.trax.models import resnet
 
@@ -39,5 +40,6 @@ def test_wide_resnet(self):
     self.assertEqual((3, 10), final_shape)
 
 
+
 if __name__ == '__main__':
   absltest.main()
diff --git a/tensor2tensor/trax/optimizers/base.py b/tensor2tensor/trax/optimizers/base.py
index 223a4e67a..cc3505412 100644
--- a/tensor2tensor/trax/optimizers/base.py
+++ b/tensor2tensor/trax/optimizers/base.py
@@ -20,7 +20,6 @@
 from __future__ import print_function
 
 from tensor2tensor.trax.backend import numpy as np
-
 from tensor2tensor.trax.layers import base as layers
 
 
@@ -89,10 +88,21 @@ def step_size(self, i):
   def tree_init(self, x_tree):
     return [self.init(x) for x in tree_flatten(x_tree)]
 
+  def _update_and_check(self, i, g, x, s):
+    new_x, new_s = self.update(i, g, x, s)
+    if isinstance(x, np.ndarray):
+      assert isinstance(new_x, np.ndarray), ("The type of the new parameter "
+                                             "values should be np.ndarray; "
+                                             "got %s" % type(new_x))
+      assert new_x.dtype == x.dtype, ("The dtype of the new parameter values "
+                                      "(%s) is not the same as the old one (%s)"
+                                      % (new_x.dtype, x.dtype))
+    return new_x, new_s
+
   def tree_update(self, i, grad_tree, x_tree, opt_state):
     grad_flat = tree_flatten(grad_tree)
     x_flat = tree_flatten(x_tree)
-    updated_pairs = [self.update(i, g, x, s)
+    updated_pairs = [self._update_and_check(i, g, x, s)
                      for (g, x, s) in zip(grad_flat, x_flat, opt_state)]
     new_x_flat, new_opt_state = zip(*updated_pairs)
     new_x, _ = tree_unflatten(new_x_flat, x_tree)
@@ -126,7 +136,7 @@ def init(self, x):
 
   def update(self, i, g, x, state):
     del state
-    return x - self.step_size(i) * g, None
+    return x - (self.step_size(i) * g).astype(x.dtype), None
 
 
 class Momentum(Optimizer):
@@ -142,7 +152,7 @@ def init(self, x):
 
   def update(self, i, g, x, velocity):
     new_velocity = self._mass * velocity - (1. - self._mass) * g
-    return x + self.step_size(i) * new_velocity, new_velocity
+    return x + (self.step_size(i) * new_velocity).astype(x.dtype), new_velocity
 
 
 class RMSProp(Optimizer):
@@ -159,7 +169,8 @@ def init(self, x):
 
   def update(self, i, g, x, avg_sq_grad):
     avg_sq_grad = avg_sq_grad * self._gamma + g**2 * (1. - self._gamma)
-    x = x - self.step_size(i) * g / (np.sqrt(avg_sq_grad) + self._epsilon)
+    x = x - (self.step_size(i) * g /
+             (np.sqrt(avg_sq_grad) + self._epsilon)).astype(x.dtype)
     return x, avg_sq_grad
 
 
@@ -196,7 +207,7 @@ def update(self, i, g, x, state):
     v = (1 - b2) * (g ** 2) + b2 * v  # Second moment estimate.
     mhat = m / (1 - b1 ** (i + 1))  # Bias correction.
     vhat = v / (1 - b2 ** (i + 1))
-    x = x - self.step_size(i) * mhat / (np.sqrt(vhat) + eps)
+    x = x - (self.step_size(i) * mhat / (np.sqrt(vhat) + eps)).astype(x.dtype)
     return x, (m, v)
 
 
@@ -227,7 +238,7 @@ def _update_diagonal(self, step, g, x, m, v):
                               np.zeros_like(v[0]))
     preconditioned_g = preconditioner * g
     m = (1 - self._momentum) * preconditioned_g + self._momentum * m
-    x = x - self.step_size(step) * m
+    x = x - (self.step_size(step) * m).astype(x.dtype)
     return x, (m, v)
 
   def _expanded_shape(self, shape, axis):
@@ -255,7 +266,7 @@ def _update_sketched(self, step, g, x, m, v):
                                     np.zeros_like(current_accumulator))
     preconditioned_gradient = g * accumulator_inv_sqrt
     m = (1.0 - self._momentum) * preconditioned_gradient + self._momentum * m
-    x = x - self.step_size(step) * m
+    x = x - (self.step_size(step) * m).astype(x.dtype)
     for i in range(len(v)):
       axes = list(range(int(i))) + list(range(int(i) + 1, rank))
       dim_accumulator = np.amax(current_accumulator, axis=axes)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index d3adb3e76..076dd5826 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -423,6 +423,11 @@ def compute_loss(opt_state, batch, rng):
   return compute_loss
 
 
+@gin.configurable
+def _is_jit_init(value=True):
+  return value
+
+
 def _reshape_by_device_single(x, n_devices):
   """Reshape x into a shape [n_devices, ...]."""
   x_shape = list(x.shape)
@@ -502,12 +507,13 @@ def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs, output_dir,
     if state.params:
       opt_state = state.params
     else:
-      # JIT parameter initialization to avoid memory fragmentation
       def initialize(input_shape, input_dtype, init_rng):
         params = model_train.initialize(input_shape, input_dtype, init_rng)
         opt_state = (params, opt.tree_init(params))
         return opt_state
-      initialize = backend.jit(initialize, static_argnums=(0, 1))
+      if _is_jit_init():
+        # JIT parameter initialization to avoid memory fragmentation
+        initialize = backend.jit(initialize, static_argnums=(0, 1))
       opt_state = initialize(model_input_shape, inputs.input_dtype, init_rng)
     if n_devices > 1:
       replicate = lambda x: numpy.broadcast_to(x, (n_devices,) + x.shape)

From ca1606891eb3a4d1368ba8850289013939befbcc Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 12 Jul 2019 15:18:51 -0700
Subject: [PATCH 2196/2720] Add ability to stuff arbitrary data in a time-step
 and therefore a trajectory.

This helps for example to add log-probabilities of a selected action.

In addition to log-probabilities we also use it to stuff the value prediction
at each time-step as well.

Not only does this save us re-computing these things but is the right thing to
do in the face of randomness in the policy/value network (ex: dropout).

PiperOrigin-RevId: 257878599
---
 tensor2tensor/envs/env_problem.py            | 13 +--
 tensor2tensor/envs/env_problem_utils.py      | 14 +++-
 tensor2tensor/envs/env_problem_utils_test.py |  6 +-
 tensor2tensor/envs/time_step.py              | 21 +++--
 tensor2tensor/envs/time_step_test.py         |  7 +-
 tensor2tensor/envs/trajectory.py             | 33 +++++++-
 tensor2tensor/envs/trajectory_test.py        | 16 +++-
 tensor2tensor/trax/rlax/ppo.py               | 85 ++++++++++++++------
 tensor2tensor/trax/rlax/ppo_test.py          | 14 +++-
 9 files changed, 160 insertions(+), 49 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index ebbf10d27..516a05a86 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -599,7 +599,7 @@ def apply_step(i):
     # tuple.
     return tuple(map(np.stack, [observations, rewards, dones, infos]))
 
-  def step(self, actions):
+  def step(self, actions, infos=None):
     """Takes a step in all environments.
 
     Subclasses should override _step to do the actual reset if something other
@@ -607,12 +607,14 @@ def step(self, actions):
 
     Args:
       actions: Batch of actions.
+      infos: (optional) a dictionary of keys and values, where all the values
+        have the first dimension as batch_size.
 
     Returns:
-      (preprocessed_observations, processed_rewards, dones, infos).
+      (preprocessed_observations, processed_rewards, dones, env_infos).
     """
 
-    observations, raw_rewards, dones, infos = self._step(actions)
+    observations, raw_rewards, dones, env_infos = self._step(actions)
 
     # Process rewards.
     raw_rewards = raw_rewards.astype(np.float32)
@@ -623,9 +625,10 @@ def step(self, actions):
 
     # Record history.
     self.trajectories.step(processed_observations, raw_rewards,
-                           processed_rewards, dones, actions)
+                           processed_rewards, dones, actions,
+                           infos=infos)
 
-    return processed_observations, processed_rewards, dones, infos
+    return processed_observations, processed_rewards, dones, env_infos
 
   def example_reading_spec(self):
     """Data fields to store on disk and their decoders."""
diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 8064c4670..69ec1d945 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -132,7 +132,8 @@ def epsilon_greedy(log_probs):
     assert (B,) == lengths.shape
 
     t1 = time.time()
-    log_prob_actions, _, rng = policy_fun(padded_observations, rng=rng)
+    log_prob_actions, value_predictions, rng = policy_fun(
+        padded_observations, rng=rng)
     policy_application_total_time += (time.time() - t1)
 
     assert (B, T) == log_prob_actions.shape[:2]
@@ -143,8 +144,12 @@ def epsilon_greedy(log_probs):
     index = lengths - 1  # Since we want to index using lengths.
     log_probs = log_prob_actions[np.arange(B)[:, None], index[:, None],
                                  np.arange(A)]
+    value_preds = value_predictions[np.arange(B)[:, None], index[:, None],
+                                    np.arange(1)]
     assert (B, A) == log_probs.shape, \
         "B=%d, A=%d, log_probs.shape=%s" % (B, A, log_probs.shape)
+    assert (B, 1) == value_preds.shape, \
+        "B=%d, value_preds.shape=%s" % (B, value_preds.shape)
 
     actions = None
     if policy_sampling == GUMBEL_SAMPLING:
@@ -156,9 +161,12 @@ def epsilon_greedy(log_probs):
 
     # Step through the env.
     t1 = time.time()
-    _, _, dones, infos = env.step(actions)
+    _, _, dones, env_infos = env.step(
+        actions, infos={"log_prob_actions": log_probs,
+                        "value_predictions": value_preds})
     env_actions_total_time += (time.time() - t1)
-    bare_env_run_time += sum(info["__bare_env_run_time__"] for info in infos)
+    bare_env_run_time += sum(
+        info["__bare_env_run_time__"] for info in env_infos)
 
     # Count the number of done trajectories, the others could just have been
     # truncated.
diff --git a/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index 8eb6d6501..6a0288cde 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -65,7 +65,7 @@ def policy_fun(observations, rng=None):
       p = np.random.uniform(size=(b, t, a))
       p = np.exp(p)
       p = p / np.sum(p, axis=-1, keepdims=True)
-      return np.log(p), (), rng
+      return np.log(p), np.log(p), rng
 
     max_timestep = 15
     num_trajectories = 2
@@ -83,12 +83,16 @@ def policy_fun(observations, rng=None):
     T = traj[1].shape[0]  # pylint: disable=invalid-name
     self.assertEqual((T + 1, 4), traj[0].shape)  # (4,) is OBS
     self.assertEqual((T,), traj[2].shape)
+    self.assertEqual(T, len(traj[4]["log_prob_actions"]))
+    self.assertEqual(T, len(traj[4]["value_predictions"]))
     self.assertLessEqual(T, max_timestep)
 
     traj = trajectories[1]
     T = traj[1].shape[0]  # pylint: disable=invalid-name
     self.assertEqual((T + 1, 4), traj[0].shape)
     self.assertEqual((T,), traj[2].shape)
+    self.assertEqual(T, len(traj[4]["log_prob_actions"]))
+    self.assertEqual(T, len(traj[4]["value_predictions"]))
     self.assertLessEqual(T, max_timestep)
 
 
diff --git a/tensor2tensor/envs/time_step.py b/tensor2tensor/envs/time_step.py
index f02567d9e..7f6bcc233 100644
--- a/tensor2tensor/envs/time_step.py
+++ b/tensor2tensor/envs/time_step.py
@@ -18,18 +18,24 @@
 Let:
 r_t = Reward(s_{t-1}, a_{t-1}, s_t)  - reward for getting into a state.
 d_t = Done(s_t)                      - is this state terminal.
+a_t = Action performed at state s_t
+i_t = (optional) Dictionary of key, value pairs of miscellaneous data.
 
 Then the sequence of states, actions and rewards looks like the following:
 
-s0, a0 s1/r1/d1, a1 s2/r2/d2, a2 s3/r3/d3, ...
+s0, a0/i0 s1/r1/d1, a1/i1 s2/r2/d2, a2/i2 s3/r3/d3, ...
 
-TimeStep holds (s_t, d_t, r_t, a_t).
+TimeStep holds (s_t, d_t, r_t, a_t, i_t).
 
 NOTE: When we call step on an environment at time-step t, we supply a_t and in
 return the env gives us s_{t+1}, d_{t+1}, r_{t+1}
 
-So, we'd have to add the actions a_t to the current time-step, but add the
+So, we'd have to add the actions a_t/i_t to the current time-step, but add the
 observations, rewards and dones to a new time-step.
+
+NOTE: wrt `info` - A good solution could be to have two additional fields in
+TimeStep - structured algo_info (a namedtuple, possibly different for every
+algorithm, or None if we don't use any) and unstructured env_info (a dict).))
 """
 
 from __future__ import absolute_import
@@ -42,7 +48,8 @@
 class TimeStep(
     collections.namedtuple(
         "TimeStep",
-        ["observation", "done", "raw_reward", "processed_reward", "action"])):
+        ["observation", "done", "raw_reward", "processed_reward", "action",
+         "info"])):
   """This class represents the time-step as mentioned above."""
 
   def replace(self, **kwargs):
@@ -61,7 +68,9 @@ def create_time_step(cls,
                        done=False,
                        raw_reward=None,
                        processed_reward=None,
-                       action=None):
+                       action=None,
+                       info=None):
     """Creates a TimeStep with both rewards and actions as optional."""
 
-    return cls(observation, done, raw_reward, processed_reward, action)
+    return cls(observation, done, raw_reward, processed_reward, action,
+               info)
diff --git a/tensor2tensor/envs/time_step_test.py b/tensor2tensor/envs/time_step_test.py
index 02eeb9612..1ab8f26b7 100644
--- a/tensor2tensor/envs/time_step_test.py
+++ b/tensor2tensor/envs/time_step_test.py
@@ -28,19 +28,21 @@ class TimeStepTest(tf.test.TestCase):
 
   def test_create_time_step(self):
     ts = time_step.TimeStep.create_time_step(
-        observation=1, done=True, raw_reward=1.0, processed_reward=1, action=1)
+        observation=1, done=True, raw_reward=1.0, processed_reward=1, action=1,
+        info={1: 1, 2: 4})
 
     self.assertEqual(1, ts.observation)
     self.assertTrue(ts.done)
     self.assertNear(1.0, ts.raw_reward, 1e-6)
     self.assertEqual(1, ts.processed_reward)
     self.assertEqual(1, ts.action)
+    self.assertEqual({1: 1, 2: 4}, ts.info)
 
   def test_replace(self):
     ts = time_step.TimeStep.create_time_step(observation=1, action=1)
     self.assertFalse(ts.done)
 
-    tsr = ts.replace(action=2, done=True)
+    tsr = ts.replace(action=2, done=True, info={1: 1, 2: 4})
 
     # Asert that ts didn't change.
     self.assertFalse(ts.done)
@@ -51,6 +53,7 @@ def test_replace(self):
     self.assertTrue(tsr.done)
     self.assertEqual(1, tsr.observation)  # unchanged
     self.assertEqual(2, tsr.action)  # changed
+    self.assertEqual({1: 1, 2: 4}, tsr.info)
 
 
 if __name__ == '__main__':
diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index da04a8744..1e6787588 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -126,6 +126,17 @@ def actions_np(self):
     # The last action is None, so let's skip it.
     return np.stack([ts.action for ts in self.time_steps[:-1]])
 
+  @property
+  def info_np(self):
+    if not self.time_steps or not self.time_steps[0].info:
+      return None
+    info_np_dict = {}
+    for info_key in self.time_steps[0].info:
+      # Same as actions, the last info is missing, so we skip it.
+      info_np_dict[info_key] = np.stack(
+          [ts.info[info_key] for ts in self.time_steps[:-1]])
+    return info_np_dict
+
   @property
   def rewards_np(self):
     # The first reward is None, so let's skip it.
@@ -137,8 +148,9 @@ def raw_rewards_np(self):
 
   @property
   def as_numpy(self):
+    # TODO(afrozm): Return a named tuple here, ex: TrajectoryArrays
     return (self.observations_np, self.actions_np, self.rewards_np,
-            self.raw_rewards_np)
+            self.raw_rewards_np, self.info_np)
 
 
 class BatchTrajectory(object):
@@ -271,7 +283,8 @@ def complete_all_trajectories(self):
       assert trajectory.is_active
       self._complete_trajectory(trajectory, index)
 
-  def step(self, observations, raw_rewards, processed_rewards, dones, actions):
+  def step(self, observations, raw_rewards, processed_rewards, dones, actions,
+           infos=None):
     """Record the information obtained from taking a step in all envs.
 
     Records (observation, rewards, done) in a new time-step and actions in the
@@ -293,6 +306,8 @@ def step(self, observations, raw_rewards, processed_rewards, dones, actions):
       actions: ndarray of first dimension self.batch_size, containing actions
         applied at the current time-step, which leads to the observations
         rewards and done at the next time-step, i.e. a_t
+      infos: (optional) a dictionary of keys and values, where all the values
+        have the first dimension as self.batch_size.
     """
     # Pre-conditions
     assert isinstance(observations, np.ndarray)
@@ -300,6 +315,8 @@ def step(self, observations, raw_rewards, processed_rewards, dones, actions):
     assert isinstance(processed_rewards, np.ndarray)
     assert isinstance(dones, np.ndarray)
     assert isinstance(actions, np.ndarray)
+    if infos:
+      assert isinstance(infos, dict)
 
     # We assume that we step in all envs, i.e. not like reset where we can reset
     # some envs and not others.
@@ -308,6 +325,14 @@ def step(self, observations, raw_rewards, processed_rewards, dones, actions):
     assert self.batch_size == processed_rewards.shape[0]
     assert self.batch_size == dones.shape[0]
     assert self.batch_size == actions.shape[0]
+    if infos:
+      for _, v in infos.items():
+        assert self.batch_size == len(v)
+
+    def extract_info_at_index(infos, index):
+      if not infos:
+        return None
+      return {k: v[index] for k, v in infos.items()}
 
     for index in range(self.batch_size):
       trajectory = self._trajectories[index]
@@ -320,7 +345,9 @@ def step(self, observations, raw_rewards, processed_rewards, dones, actions):
       assert trajectory.is_active
 
       # To this trajectory's last time-step, set actions.
-      trajectory.change_last_time_step(action=actions[index])
+      trajectory.change_last_time_step(
+          action=actions[index],
+          info=extract_info_at_index(infos, index))
 
       # Create a new time-step to add observation, done & rewards (no actions).
       trajectory.add_time_step(
diff --git a/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index b9a7f0bfc..8658aee21 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -138,16 +138,23 @@ def test_as_numpy(self):
     observations = np.random.uniform(size=(ts,) + shape)
     actions = np.random.choice(range(num_actions), size=(ts - 1,))
     rewards = np.random.choice([-1, 0, 1], size=(ts - 1,))
+    squares = np.arange(ts - 1) ** 2
+    cubes = np.arange(ts - 1) ** 3
+
+    def get_info(i):
+      return {"sq": squares[i], "cu": cubes[i]}
 
     # First time-step has no reward.
-    t.add_time_step(observation=observations[0], done=False, action=actions[0])
+    t.add_time_step(observation=observations[0], done=False, action=actions[0],
+                    info=get_info(0))
     for i in range(1, ts - 1):
       t.add_time_step(
           observation=observations[i],
           done=False,
           raw_reward=rewards[i - 1],
           processed_reward=rewards[i - 1],
-          action=actions[i])
+          action=actions[i],
+          info=get_info(i))
     # Last time-step has no action.
     t.add_time_step(
         observation=observations[-1],
@@ -161,6 +168,9 @@ def test_as_numpy(self):
     self.assertAllEqual(actions, traj_np[1])
     self.assertAllEqual(rewards, traj_np[2])
 
+    self.assertAllEqual(squares, traj_np[4]["sq"])
+    self.assertAllEqual(cubes, traj_np[4]["cu"])
+
 
 class BatchTrajectoryTest(tf.test.TestCase):
 
@@ -457,5 +467,5 @@ def test_observations_np(self):
       for ts in range(lengths[b], len(padded_obs_np[b])):
         self.assertAllEqual(zero_obs, padded_obs_np[b][ts])
 
-if __name__ == '__main__':
+if __name__ == "__main__":
   tf.test.main()
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 076fb4b60..f727d812d 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -179,7 +179,10 @@ def collect_trajectories(env,
       len_history_for_policy=len_history_for_policy,
       rng=rng)
   # Skip returning raw_rewards here, since they aren't used.
-  return [(t[0], t[1], t[2]) for t in trajs], n_done, timing_info
+
+  # t is the return value of Trajectory.as_numpy, so:
+  # (observation, action, processed_reward, raw_reward, infos)
+  return [(t[0], t[1], t[2], t[4]) for t in trajs], n_done, timing_info
 
 
 # This function can probably be simplified, ask how?
@@ -223,7 +226,7 @@ def pad_trajectories(trajectories, boundary=20):
   """
 
   # Let's compute max(t) over all trajectories.
-  t_max = max(r.shape[0] for (_, _, r) in trajectories)
+  t_max = max(r.shape[0] for (_, _, r, _) in trajectories)
 
   # t_max is rounded to the next multiple of `boundary`
   boundary = int(boundary)
@@ -233,9 +236,11 @@ def pad_trajectories(trajectories, boundary=20):
   padded_observations = []
   padded_actions = []
   padded_rewards = []
+  padded_infos = collections.defaultdict(list)
   padded_lengths = []
   reward_masks = []
-  for (o, a, r) in trajectories:
+
+  for (o, a, r, i) in trajectories:
     # Determine the amount to pad, this holds true for obs, actions and rewards.
     num_to_pad = bucket_length + 1 - o.shape[0]
     padded_lengths.append(num_to_pad)
@@ -244,13 +249,13 @@ def pad_trajectories(trajectories, boundary=20):
       padded_actions.append(a)
       padded_rewards.append(r)
       reward_masks.append(onp.ones_like(r, dtype=np.int32))
+      if i:
+        for k, v in i.items():
+          padded_infos[k].append(v)
       continue
 
     # First pad observations.
-    padding_config = [(0, num_to_pad, 0)]
-    for _ in range(o.ndim - 1):
-      padding_config.append((0, 0, 0))
-    padding_config = tuple(padding_config)
+    padding_config = tuple([(0, num_to_pad, 0)] + [(0, 0, 0)] * (o.ndim - 1))
 
     padding_value = get_padding_value(o.dtype)
     action_padding_value = get_padding_value(a.dtype)
@@ -272,8 +277,20 @@ def pad_trajectories(trajectories, boundary=20):
     reward_mask = onp.ones_like(r, dtype=np.int32)
     reward_masks.append(lax.pad(reward_mask, 0, padding_config))
 
+    if i:
+      for k, v in i.items():
+        # Create a padding configuration for this value.
+        padding_config = [(0, num_to_pad, 0)] + [(0, 0, 0)] * (v.ndim - 1)
+        padded_infos[k].append(lax.pad(v, 0.0, tuple(padding_config)))
+
+  # Now stack these padded_infos if they exist.
+  stacked_padded_infos = None
+  if padded_infos:
+    stacked_padded_infos = {k: np.stack(v) for k, v in padded_infos.items()}
+
   return padded_lengths, np.stack(reward_masks), np.stack(
-      padded_observations), np.stack(padded_actions), np.stack(padded_rewards)
+      padded_observations), np.stack(padded_actions), np.stack(
+          padded_rewards), stacked_padded_infos
 
 
 def rewards_to_go(rewards, mask, gamma=0.99):
@@ -997,7 +1014,7 @@ def get_predictions(observations, rng=None):
 
     padding_start_time = time.time()
     (_, reward_mask, padded_observations, padded_actions,
-     padded_rewards) = pad_trajectories(
+     padded_rewards, padded_infos) = pad_trajectories(
          trajs, boundary=boundary)
     padding_time = get_time(padding_start_time)
 
@@ -1008,21 +1025,6 @@ def get_predictions(observations, rng=None):
     logging.vlog(1, "Padded Actions' shape [%s]", str(padded_actions.shape))
     logging.vlog(1, "Padded Rewards' shape [%s]", str(padded_rewards.shape))
 
-    # Calculate log-probabilities and value predictions of the trajectories.
-    # We'll pass these to the loss functions so as to not get recomputed.
-
-    # NOTE:
-    # There is a slight problem here, if the policy network contains
-    # stochasticity in the log-probabilities (ex: dropout), then calculating
-    # these again here is not going to be correct and should be done in the
-    # collect function.
-
-    log_prob_recompute_start_time = time.time()
-    jax_rng_key, key = jax_random.split(jax_rng_key)
-    log_probabs_traj, value_predictions_traj, _ = get_predictions(
-        padded_observations, rng=key)
-    log_prob_recompute_time = get_time(log_prob_recompute_start_time)
-
     # Some assertions.
     B, T = padded_actions.shape  # pylint: disable=invalid-name
     assert (B, T) == padded_rewards.shape
@@ -1030,6 +1032,41 @@ def get_predictions(observations, rng=None):
     assert (B, T + 1) == padded_observations.shape[:2]
     assert (B, T + 1) + env.observation_space.shape == padded_observations.shape
 
+    log_prob_recompute_start_time = time.time()
+    assert ("log_prob_actions" in padded_infos and
+            "value_predictions" in padded_infos)
+    # These are the actual log-probabs and value predictions seen while picking
+    # the actions.
+    actual_log_probabs_traj = padded_infos["log_prob_actions"]
+    actual_value_predictions_traj = padded_infos["value_predictions"]
+
+    assert (B, T) == actual_log_probabs_traj.shape[:2]
+    A = actual_log_probabs_traj.shape[2]  # pylint: disable=invalid-name
+    assert (B, T, 1) == actual_value_predictions_traj.shape
+
+    # TODO(afrozm): log-probabs doesn't need to be (B, T+1, A) it can do with
+    # (B, T, A), so make that change throughout.
+
+    # NOTE: We don't have the log-probabs and value-predictions for the last
+    # observation, so we re-calculate for everything, but use the original ones
+    # for all but the last time-step.
+    jax_rng_key, key = jax_random.split(jax_rng_key)
+    log_probabs_traj, value_predictions_traj, _ = get_predictions(
+        padded_observations, rng=key)
+
+    assert (B, T + 1, A) == log_probabs_traj.shape
+    assert (B, T + 1, 1) == value_predictions_traj.shape
+
+    # Concatenate the last time-step's log-probabs and value predictions to the
+    # actual log-probabs and value predictions and use those going forward.
+    log_probabs_traj = np.concatenate(
+        (actual_log_probabs_traj, log_probabs_traj[:, -1:, :]), axis=1)
+    value_predictions_traj = np.concatenate(
+        (actual_value_predictions_traj, value_predictions_traj[:, -1:, :]),
+        axis=1)
+
+    log_prob_recompute_time = get_time(log_prob_recompute_start_time)
+
     # Linear annealing from 0.1 to 0.0
     # epsilon_schedule = epsilon if epochs == 1 else epsilon * (1.0 -
     #                                                           (i /
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rlax/ppo_test.py
index bf39fbba3..0fd4337c8 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rlax/ppo_test.py
@@ -31,6 +31,7 @@
 class PpoTest(test.TestCase):
 
   def setUp(self):
+    super(PpoTest, self).setUp()
     self.rng_key = trax.get_random_number_generator_and_set_seed(0)
 
   def test_policy_and_value_net(self):
@@ -78,7 +79,11 @@ def test_pad_trajectories(self):
       rewards = np.random.uniform(size=(time_steps,)).astype(np.float32)
       actions = np.random.randint(
           0, n_actions, size=(time_steps,)).astype(np.int32)
-      trajectories.append((observations, rewards, actions))
+      infos = {
+          "a": np.random.uniform(size=(time_steps,)).astype(np.float32),
+          "b": np.random.uniform(size=(time_steps,)).astype(np.float32)
+      }
+      trajectories.append((observations, rewards, actions, infos))
 
     # Now pad these trajectories.
     padded_trajectories = ppo.pad_trajectories(
@@ -92,7 +97,7 @@ def test_pad_trajectories(self):
 
     # Get the padded objects.
     (pad_lengths, reward_mask, padded_observations, padded_actions,
-     padded_rewards) = padded_trajectories
+     padded_rewards, padded_infos) = padded_trajectories
 
     # Expectations on the padded shapes.
     self.assertEqual(padded_observations.shape, (
@@ -103,6 +108,11 @@ def test_pad_trajectories(self):
     self.assertEqual(padded_rewards.shape, (n_trajectories, expected_padding))
     self.assertEqual(reward_mask.shape, (n_trajectories, expected_padding))
 
+    self.assertEqual(padded_infos["a"].shape,
+                     (n_trajectories, expected_padding))
+    self.assertEqual(padded_infos["b"].shape,
+                     (n_trajectories, expected_padding))
+
     # Assert that the padding lengths and reward mask are consistent.
     self.assertAllEqual(
         np.full((n_trajectories,), expected_padding),

From abbd929558dd29115acc9d0f035f1efddb45566d Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 12 Jul 2019 18:54:01 -0700
Subject: [PATCH 2197/2720] Extract a base class from EnvProblem

PiperOrigin-RevId: 257907213
---
 tensor2tensor/envs/__init__.py                |   3 +-
 tensor2tensor/envs/env_problem.py             | 291 ++++--------------
 tensor2tensor/envs/env_problem_utils_test.py  |   4 +-
 tensor2tensor/envs/gym_env_problem.py         | 279 +++++++++++++++++
 ...roblem_test.py => gym_env_problem_test.py} |  32 +-
 tensor2tensor/envs/rendered_env_problem.py    |  13 +-
 tensor2tensor/envs/tic_tac_toe_env_problem.py |   4 +-
 tensor2tensor/trax/rlax/ppo_main.py           |   4 +-
 .../trax/rlax/ppo_training_loop_test.py       |  12 +-
 9 files changed, 371 insertions(+), 271 deletions(-)
 create mode 100644 tensor2tensor/envs/gym_env_problem.py
 rename tensor2tensor/envs/{env_problem_test.py => gym_env_problem_test.py} (93%)

diff --git a/tensor2tensor/envs/__init__.py b/tensor2tensor/envs/__init__.py
index 675172fb9..8c036a176 100644
--- a/tensor2tensor/envs/__init__.py
+++ b/tensor2tensor/envs/__init__.py
@@ -19,7 +19,6 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.envs import env_problem
+from tensor2tensor.envs import gym_env_problem
 from tensor2tensor.envs import tic_tac_toe_env
 from tensor2tensor.envs import tic_tac_toe_env_problem
-
diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 516a05a86..890f74374 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -23,10 +23,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import multiprocessing.pool
-import time
-
-import gym
 from gym.core import Env
 import numpy as np
 import six
@@ -47,7 +43,7 @@
 
 
 class EnvProblem(Env, problem.Problem):
-  """An env which generates data like a problem class.
+  """Base class of an env which generates data like a problem class.
 
   EnvProblem is both a gym Env and a Problem, since it subclasses both.
 
@@ -56,7 +52,13 @@ class EnvProblem(Env, problem.Problem):
   reset is stored within this class and is persisted on disk when we call
   `generate_data` on it.
 
-  Subclasses *should* override the following functions, since they are used in
+  Subclasses *should* override the following functions:
+  - initialize_environments
+  - _reset
+  - _step
+  - _render
+
+  In addition, they should ovveride the following functions, which are used in
   the `hparams` function to return modalities and vocab_sizes.
   - input_modality
   - input_vocab_size
@@ -67,44 +69,8 @@ class EnvProblem(Env, problem.Problem):
 
   NON NATIVELY BATCHED ENVS:
 
-  The default implementations of the other major functions, should work well for
-  cases where the env is not batched by default ex: any gym env. In this case we
-  create `batch_size` number of envs and store them in a list. Any function then
-  that interacts with the envs, like reset, step or close goes over the env list
-  to do the needful, ex: when reset is called with specific indices we reset
-  only those indices, etc.
-
-  The usage of this class will look like the following:
-
-  # 1. Creates and initializes the env_problem.
-  ep = env_problem.EnvProblem(...)
-
-  # 2. One needs to call reset() at the start, this resets all envs.
-  ep.reset()
-
-  # 3. Call step with actions for all envs, i.e. len(action) = batch_size
-  obs, rewards, dones, infos = ep.step(actions)
-
-  # 4. Figure out which envs got done and reset only those.
-  ep.reset(indices=env_problem_utils.done_indices(dones))
-
-  # 5. Go back to Step #3 to further interact with the env or just dump the
-  # generated data to disk by calling:
-  ep.generate_data(...)
-
-  # 6. If we now need to use this object again to play a few more iterations
-  # perhaps with a different batch size or maybe not recording the data, then
-  # we need to re-initialize environments and do some book-keeping, call:
-  ep.initialize_environments(batch_size)
-
-  # 7. Go back to Step #2, i.e. reset all envs.
-
-  NOTE: Look at `EnvProblemTest.test_interaction_with_env` and/or
-  `EnvProblemTest.test_generate_data`
-
-  NOTE: We rely heavily that the underlying environments expose a gym style
-  interface, i.e. in addition to reset(), step() and close() we have access to
-  the following properties: observation_space, action_space, reward_range.
+  The implementation for cases where the env is not batched by default is
+  `gym_env_problem.GymEnvProblem`.
 
   NATIVELY BATCHED ENVS:
 
@@ -128,9 +94,7 @@ class EnvProblem(Env, problem.Problem):
   """
 
   def __init__(self,
-               base_env_name=None,
                batch_size=None,
-               env_wrapper_fn=None,
                reward_range=(-np.inf, np.inf),
                discrete_rewards=True,
                parallelism=1,
@@ -138,12 +102,8 @@ def __init__(self,
     """Initializes this class by creating the envs and managing trajectories.
 
     Args:
-      base_env_name: (string) passed to `gym.make` to make the underlying
-        environment.
       batch_size: (int or None) How many envs to make in the non natively
         batched mode.
-      env_wrapper_fn: (callable(env): env) Applies gym wrappers to the base
-        environment.
       reward_range: (tuple(number, number)) the first element is the minimum
         reward and the second is the maximum reward, used to clip and process
         the raw reward in `process_rewards`.
@@ -157,10 +117,6 @@ def __init__(self,
     # Call the super's ctor.
     problem.Problem.__init__(self, was_reversed=False, was_copy=False)
 
-    # Name for the base environment, will be used in `gym.make` in
-    # the default implementation of `initialize_environments`.
-    self._base_env_name = base_env_name
-
     # An env generates data when it is given actions by an agent which is either
     # a policy or a human -- this is supposed to be the `id` of the agent.
     #
@@ -175,14 +131,7 @@ def __init__(self,
     # If set, we discretize the rewards and treat them as integers.
     self._discrete_rewards = discrete_rewards
 
-    # Initialize the environment(s).
-
-    # This can either be a list of environments of len `batch_size` or this can
-    # be a Neural Network, in which case it will be fed input with first
-    # dimension = `batch_size`.
-    self._envs = None
-    self._pool = None
-    self._parallelism = parallelism
+    self._parallelism = None
 
     self._observation_space = None
     self._action_space = None
@@ -193,8 +142,6 @@ def __init__(self,
 
     self._batch_size = None
 
-    self._env_wrapper_fn = env_wrapper_fn
-
     if batch_size is not None:
       self.initialize(batch_size=batch_size, **env_kwargs)
 
@@ -204,54 +151,20 @@ def batch_size(self):
     # places in ppo_learner.py -- re-evaluate if needed.
     return self._batch_size
 
-  @property
-  def base_env_name(self):
-    return self._base_env_name
-
   @property
   def trajectories(self):
     return self._trajectories
 
-  def _verify_same_spaces(self):
-    """Verifies that all the envs have the same observation and action space."""
-
-    # Pre-conditions: self._envs is initialized.
-
-    if self._envs is None:
-      raise ValueError("Environments not initialized.")
-
-    if not isinstance(self._envs, list):
-      tf.logging.warning("Not checking observation and action space "
-                         "compatibility across envs, since there is just one.")
-      return
-
-    # NOTE: We compare string representations of observation_space and
-    # action_space because compositional classes like space.Tuple don't return
-    # true on object comparison.
-
-    if not all(
-        str(env.observation_space) == str(self.observation_space)
-        for env in self._envs):
-      err_str = ("All environments should have the same observation space, but "
-                 "don't.")
-      tf.logging.error(err_str)
-      # Log all observation spaces.
-      for i, env in enumerate(self._envs):
-        tf.logging.error("Env[%d] has observation space [%s]", i,
-                         env.observation_space)
-      raise ValueError(err_str)
-
-    if not all(
-        str(env.action_space) == str(self.action_space) for env in self._envs):
-      err_str = "All environments should have the same action space, but don't."
-      tf.logging.error(err_str)
-      # Log all action spaces.
-      for i, env in enumerate(self._envs):
-        tf.logging.error("Env[%d] has action space [%s]", i, env.action_space)
-      raise ValueError(err_str)
-
-  def initialize(self, **kwargs):
-    self.initialize_environments(**kwargs)
+  def initialize(self, batch_size=1, **kwargs):
+    self.initialize_environments(batch_size=batch_size, **kwargs)
+
+    self._batch_size = batch_size
+
+    # This data structure stores the history of each env.
+    #
+    # NOTE: Even if the env is a NN and can step in all batches concurrently, it
+    # is still valuable to store the trajectories separately.
+    self._trajectories = trajectory.BatchTrajectory(batch_size=batch_size)
 
     # Assert that *all* the above are now set, we should do this since
     # subclasses can override `initialize_environments`.
@@ -259,70 +172,20 @@ def initialize(self, **kwargs):
     assert self._observation_space is not None
     assert self._action_space is not None
     assert self._reward_range is not None
-    assert self._trajectories is not None
 
-  def initialize_environments(self, batch_size=1, parallelism=1, **env_kwargs):
-    """Initializes the environments and trajectories.
-
-    Subclasses can override this if they don't want a default implementation
-    which initializes `batch_size` environments, but must take care to
-    initialize self._trajectories (this is checked in __init__ anyways).
+  def initialize_environments(self, batch_size=1, parallelism=1, **kwargs):
+    """Initializes the environments.
 
     Args:
-      batch_size: (int) Number of `self.base_env_name` envs to initialize.
-      parallelism: (int) If this is greater than one then we run the envs in
-        parallel using multi-threading.
-      **env_kwargs: (dict) Kwargs to pass to gym.make.
+      batch_size: (int) Number of envs to initialize.
+      parallelism: (int) If this is greater than one then we allow the
+        implementation to use multi-threading to step the envs.
+      **kwargs: (dict) Any additional args needed to initialize the envs.
     """
-    assert batch_size >= 1
-    self._batch_size = batch_size
-
-    self._envs = [
-        gym.make(self.base_env_name, **env_kwargs) for _ in range(batch_size)
-    ]
-    self._parallelism = parallelism
-    self._pool = multiprocessing.pool.ThreadPool(self._parallelism)
-    if self._env_wrapper_fn is not None:
-      self._envs = list(map(self._env_wrapper_fn, self._envs))
-
-    # If self.observation_space and self.action_space aren't None, then it means
-    # that this is a re-initialization of this class, in that case make sure
-    # that this matches our previous behaviour.
-    if self._observation_space:
-      assert str(self._observation_space) == str(
-          self._envs[0].observation_space)
-    else:
-      # This means that we are initializing this class for the first time.
-      #
-      # We set this equal to the first env's observation space, later on we'll
-      # verify that all envs have the same observation space.
-      self._observation_space = self._envs[0].observation_space
-
-    # Similarly for action_space
-    if self._action_space:
-      assert str(self._action_space) == str(self._envs[0].action_space)
-    else:
-      self._action_space = self._envs[0].action_space
-
-    self._verify_same_spaces()
-
-    # If self.reward_range is None, i.e. this means that we should take the
-    # reward range of the env.
-    if self.reward_range is None:
-      self._reward_range = self._envs[0].reward_range
-
-    # This data structure stores the history of each env.
-    #
-    # NOTE: Even if the env is a NN and can step in all batches concurrently, it
-    # is still valuable to store the trajectories separately.
-    self._trajectories = trajectory.BatchTrajectory(batch_size=batch_size)
+    raise NotImplementedError
 
   def assert_common_preconditions(self):
-    # Asserts on the common pre-conditions of:
-    #  - self._envs is initialized.
-    #  - self._envs is a list.
-    assert self._envs
-    assert isinstance(self._envs, list)
+    pass
 
   @property
   def observation_space(self):
@@ -448,53 +311,21 @@ def unwrapped(self):
     return self
 
   def seed(self, seed=None):
-    if not self._envs:
-      tf.logging.info("`seed` called on non-existent envs, doing nothing.")
-      return None
-
-    if not isinstance(self._envs, list):
-      tf.logging.warning("`seed` called on non-list envs, doing nothing.")
-      return None
-
-    tf.logging.warning(
-        "Called `seed` on EnvProblem, calling seed on the underlying envs.")
-    for env in self._envs:
-      env.seed(seed)
-
     return [seed]
 
   def close(self):
-    if not self._envs:
-      tf.logging.info("`close` called on non-existent envs, doing nothing.")
-      return
-
-    if not isinstance(self._envs, list):
-      tf.logging.warning("`close` called on non-list envs, doing nothing.")
-      return
-
-    # Call close on all the envs one by one.
-    for env in self._envs:
-      env.close()
+    pass
 
   def _reset(self, indices):
     """Resets environments at indices shouldn't pre-process or record.
 
-    Subclasses should override this to do the actual reset if something other
-    than the default implementation is desired.
-
     Args:
       indices: list of indices of underlying envs to call reset on.
 
     Returns:
       np.ndarray of stacked observations from the reset-ed envs.
     """
-
-    # Pre-conditions: common_preconditions, see `assert_common_preconditions`.
-    self.assert_common_preconditions()
-
-    # This returns a numpy array with first dimension `len(indices)` and the
-    # rest being the dimensionality of the observation.
-    return np.stack([self._envs[index].reset() for index in indices])
+    raise NotImplementedError
 
   def truncate(self, indices=None, num_to_keep=1):
     """Truncates trajectories at the specified indices."""
@@ -532,6 +363,9 @@ def reset(self, indices=None):
           "`reset` called with empty indices array, this is a no-op.")
       return None
 
+    # Pre-conditions: common_preconditions, see `assert_common_preconditions`.
+    self.assert_common_preconditions()
+
     observations = self._reset(indices)
     processed_observations = self.process_observations(observations)
 
@@ -540,12 +374,24 @@ def reset(self, indices=None):
 
     return processed_observations
 
-  def render(self, mode="human", indices=None):
-    """Calls render with the given mode on the specified indices.
+  def _render(self, indices, mode="human"):
+    """Renders the environments with the given mode on the specified indices.
 
     Args:
+      indices: array of indices.
       mode: rendering mode.
+
+    Returns:
+      a list of return values from the environments rendered.
+    """
+    raise NotImplementedError
+
+  def render(self, indices=None, mode="human"):
+    """Renders the environments with the given mode on the specified indices.
+
+    Args:
       indices: array of indices, calls render on everything if indices is None.
+      mode: rendering mode.
 
     Returns:
       a list of return values from the environments rendered.
@@ -553,51 +399,18 @@ def render(self, mode="human", indices=None):
 
     if indices is None:
       indices = np.arange(self.batch_size)
-    ret_vals = []
-    for index in indices:
-      ret_vals.append(self._envs[index].render(mode=mode))
-    return ret_vals
+    return self._render(indices, mode)
 
   def _step(self, actions):
     """Takes a step in all environments, shouldn't pre-process or record.
 
-    Subclasses should override this to do the actual step if something other
-    than the default implementation is desired.
-
     Args:
       actions: (np.ndarray) with first dimension equal to the batch size.
 
     Returns:
       a tuple of stacked raw observations, raw rewards, dones and infos.
     """
-
-    # Pre-conditions: common_preconditions, see `assert_common_preconditions`.
-    #               : len(actions) == len(self._envs)
-    self.assert_common_preconditions()
-    assert len(actions) == len(self._envs)
-    assert self.batch_size == len(actions)
-
-    observations = [None] * self.batch_size
-    rewards = [None] * self.batch_size
-    dones = [None] * self.batch_size
-    infos = [{} for _ in range(self.batch_size)]
-
-    def apply_step(i):
-      t1 = time.time()
-      observations[i], rewards[i], dones[i], infos[i] = self._envs[i].step(
-          actions[i])
-      t2 = time.time()
-      infos[i]["__bare_env_run_time__"] = t2 - t1
-
-    if self._parallelism > 1:
-      self._pool.map(apply_step, range(self.batch_size))
-    else:
-      for i in range(self.batch_size):
-        apply_step(i)
-
-    # Convert each list (observations, rewards, ...) into np.array and return a
-    # tuple.
-    return tuple(map(np.stack, [observations, rewards, dones, infos]))
+    raise NotImplementedError
 
   def step(self, actions, infos=None):
     """Takes a step in all environments.
@@ -613,6 +426,10 @@ def step(self, actions, infos=None):
     Returns:
       (preprocessed_observations, processed_rewards, dones, env_infos).
     """
+    # Pre-conditions: common_preconditions, see `assert_common_preconditions`.
+    #               : len(actions) == len(self._envs)
+    self.assert_common_preconditions()
+    assert self.batch_size == len(actions)
 
     observations, raw_rewards, dones, env_infos = self._step(actions)
 
diff --git a/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index 6a0288cde..e7416bc1c 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -21,8 +21,8 @@
 
 import numpy as np
 
-from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import env_problem_utils
+from tensor2tensor.envs import gym_env_problem
 from tensor2tensor.envs import tic_tac_toe_env  # pylint: disable=unused-import
 from tensor2tensor.envs import tic_tac_toe_env_problem
 
@@ -47,7 +47,7 @@ def test_play_env_problem_randomly(self):
         batch_size, ep.trajectories.num_time_steps)
 
   def test_play_env_problem_with_policy(self):
-    env = env_problem.EnvProblem(
+    env = gym_env_problem.GymEnvProblem(
         base_env_name="CartPole-v0", batch_size=2, reward_range=(-1, 1))
 
     # Let's make sure that at-most 4 observations come to the policy function.
diff --git a/tensor2tensor/envs/gym_env_problem.py b/tensor2tensor/envs/gym_env_problem.py
new file mode 100644
index 000000000..475c7fc55
--- /dev/null
+++ b/tensor2tensor/envs/gym_env_problem.py
@@ -0,0 +1,279 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 class for envs that store their history.
+
+EnvProblem subclasses Problem and also implements the Gym interface (step,
+reset, render, close, seed)
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import multiprocessing.pool
+import time
+
+import gym
+import numpy as np
+from tensor2tensor.envs import env_problem
+from tensor2tensor.envs import trajectory
+import tensorflow as tf
+
+
+class GymEnvProblem(env_problem.EnvProblem):
+  """An EnvProblem implemented as a batch of gym envs.
+
+  This implementation should work well for cases where the env is not batched by
+  default ex: any gym env. In this case we create `batch_size` number of envs
+  and store them in a list. Any function then that interacts with the envs, like
+  reset, step or close goes over the env list to do the needful, ex: when reset
+  is called with specific indices we reset only those indices, etc.
+
+  The usage of this class will look like the following:
+
+  # 1. Creates and initializes the env_problem.
+  ep = env_problem.EnvProblem(...)
+
+  # 2. One needs to call reset() at the start, this resets all envs.
+  ep.reset()
+
+  # 3. Call step with actions for all envs, i.e. len(action) = batch_size
+  obs, rewards, dones, infos = ep.step(actions)
+
+  # 4. Figure out which envs got done and reset only those.
+  ep.reset(indices=env_problem_utils.done_indices(dones))
+
+  # 5. Go back to Step #3 to further interact with the env or just dump the
+  # generated data to disk by calling:
+  ep.generate_data(...)
+
+  # 6. If we now need to use this object again to play a few more iterations
+  # perhaps with a different batch size or maybe not recording the data, then
+  # we need to re-initialize environments and do some book-keeping, call:
+  ep.initialize_environments(batch_size)
+
+  # 7. Go back to Step #2, i.e. reset all envs.
+
+  NOTE: Look at `EnvProblemTest.test_interaction_with_env` and/or
+  `EnvProblemTest.test_generate_data`
+
+  NOTE: We rely heavily that the underlying environments expose a gym style
+  interface, i.e. in addition to reset(), step() and close() we have access to
+  the following properties: observation_space, action_space, reward_range.
+  """
+
+  def __init__(self, base_env_name=None, env_wrapper_fn=None, **kwargs):
+    """Initializes this class by creating the envs and managing trajectories.
+
+    Args:
+      base_env_name: (string) passed to `gym.make` to make the underlying
+        environment.
+      env_wrapper_fn: (callable(env): env) Applies gym wrappers to the base
+        environment.
+      **kwargs: (dict) Arguments passed to the base class.
+    """
+    # Name for the base environment, will be used in `gym.make` in
+    # the default implementation of `initialize_environments`.
+    self._base_env_name = base_env_name
+
+    # An env generates data when it is given actions by an agent which is either
+    # a policy or a human -- this is supposed to be the `id` of the agent.
+    #
+    # In practice, this is used only to store (and possibly retrieve) history
+    # to an appropriate directory.
+    self._agent_id = "default"
+
+    # Initialize the environment(s).
+
+    # This can either be a list of environments of len `batch_size` or this can
+    # be a Neural Network, in which case it will be fed input with first
+    # dimension = `batch_size`.
+    self._envs = None
+    self._pool = None
+
+    self._env_wrapper_fn = env_wrapper_fn
+
+    # Call the super's ctor. It will use some of the member fields, so we call
+    # it in the end.
+    super(GymEnvProblem, self).__init__(**kwargs)
+
+  @property
+  def base_env_name(self):
+    return self._base_env_name
+
+  def _verify_same_spaces(self):
+    """Verifies that all the envs have the same observation and action space."""
+
+    # Pre-conditions: self._envs is initialized.
+
+    if self._envs is None:
+      raise ValueError("Environments not initialized.")
+
+    if not isinstance(self._envs, list):
+      tf.logging.warning("Not checking observation and action space "
+                         "compatibility across envs, since there is just one.")
+      return
+
+    # NOTE: We compare string representations of observation_space and
+    # action_space because compositional classes like space.Tuple don't return
+    # true on object comparison.
+
+    if not all(
+        str(env.observation_space) == str(self.observation_space)
+        for env in self._envs):
+      err_str = ("All environments should have the same observation space, but "
+                 "don't.")
+      tf.logging.error(err_str)
+      # Log all observation spaces.
+      for i, env in enumerate(self._envs):
+        tf.logging.error("Env[%d] has observation space [%s]", i,
+                         env.observation_space)
+      raise ValueError(err_str)
+
+    if not all(
+        str(env.action_space) == str(self.action_space) for env in self._envs):
+      err_str = "All environments should have the same action space, but don't."
+      tf.logging.error(err_str)
+      # Log all action spaces.
+      for i, env in enumerate(self._envs):
+        tf.logging.error("Env[%d] has action space [%s]", i, env.action_space)
+      raise ValueError(err_str)
+
+  def initialize_environments(self, batch_size=1, parallelism=1, **kwargs):
+    """Initializes the environments.
+
+    Args:
+      batch_size: (int) Number of `self.base_env_name` envs to initialize.
+      parallelism: (int) If this is greater than one then we run the envs in
+        parallel using multi-threading.
+      **kwargs: (dict) Kwargs to pass to gym.make.
+    """
+    assert batch_size >= 1
+
+    self._envs = [
+        gym.make(self.base_env_name, **kwargs) for _ in range(batch_size)
+    ]
+    self._parallelism = parallelism
+    self._pool = multiprocessing.pool.ThreadPool(self._parallelism)
+    if self._env_wrapper_fn is not None:
+      self._envs = list(map(self._env_wrapper_fn, self._envs))
+
+    # If self.observation_space and self.action_space aren't None, then it means
+    # that this is a re-initialization of this class, in that case make sure
+    # that this matches our previous behaviour.
+    if self._observation_space:
+      assert str(self._observation_space) == str(
+          self._envs[0].observation_space)
+    else:
+      # This means that we are initializing this class for the first time.
+      #
+      # We set this equal to the first env's observation space, later on we'll
+      # verify that all envs have the same observation space.
+      self._observation_space = self._envs[0].observation_space
+
+    # Similarly for action_space
+    if self._action_space:
+      assert str(self._action_space) == str(self._envs[0].action_space)
+    else:
+      self._action_space = self._envs[0].action_space
+
+    self._verify_same_spaces()
+
+    # If self.reward_range is None, i.e. this means that we should take the
+    # reward range of the env.
+    if self.reward_range is None:
+      self._reward_range = self._envs[0].reward_range
+
+    # This data structure stores the history of each env.
+    #
+    # NOTE: Even if the env is a NN and can step in all batches concurrently, it
+    # is still valuable to store the trajectories separately.
+    self._trajectories = trajectory.BatchTrajectory(batch_size=batch_size)
+
+  def seed(self, seed=None):
+    if not self._envs:
+      tf.logging.info("`seed` called on non-existent envs, doing nothing.")
+      return None
+
+    if not isinstance(self._envs, list):
+      tf.logging.warning("`seed` called on non-list envs, doing nothing.")
+      return None
+
+    tf.logging.warning(
+        "Called `seed` on EnvProblem, calling seed on the underlying envs.")
+    for env in self._envs:
+      env.seed(seed)
+
+    return super(GymEnvProblem, self).seed(seed=seed)
+
+  def close(self):
+    if not self._envs:
+      tf.logging.info("`close` called on non-existent envs, doing nothing.")
+      return
+
+    if not isinstance(self._envs, list):
+      tf.logging.warning("`close` called on non-list envs, doing nothing.")
+      return
+
+    # Call close on all the envs one by one.
+    for env in self._envs:
+      env.close()
+
+  def _reset(self, indices):
+    """Resets environments at indices shouldn't pre-process or record.
+
+    Args:
+      indices: list of indices of underlying envs to call reset on.
+
+    Returns:
+      np.ndarray of stacked observations from the reset-ed envs.
+    """
+    # This returns a numpy array with first dimension `len(indices)` and the
+    # rest being the dimensionality of the observation.
+    return np.stack([self._envs[index].reset() for index in indices])
+
+  def _step(self, actions):
+    """Takes a step in all environments, shouldn't pre-process or record.
+
+    Args:
+      actions: (np.ndarray) with first dimension equal to the batch size.
+
+    Returns:
+      a tuple of stacked raw observations, raw rewards, dones and infos.
+    """
+    assert len(actions) == len(self._envs)
+
+    observations = [None] * self.batch_size
+    rewards = [None] * self.batch_size
+    dones = [None] * self.batch_size
+    infos = [{} for _ in range(self.batch_size)]
+
+    def apply_step(i):
+      t1 = time.time()
+      observations[i], rewards[i], dones[i], infos[i] = self._envs[i].step(
+          actions[i])
+      t2 = time.time()
+      infos[i]["__bare_env_run_time__"] = t2 - t1
+
+    if self._parallelism > 1:
+      self._pool.map(apply_step, range(self.batch_size))
+    else:
+      for i in range(self.batch_size):
+        apply_step(i)
+
+    # Convert each list (observations, rewards, ...) into np.array and return a
+    # tuple.
+    return tuple(map(np.stack, [observations, rewards, dones, infos]))
diff --git a/tensor2tensor/envs/env_problem_test.py b/tensor2tensor/envs/gym_env_problem_test.py
similarity index 93%
rename from tensor2tensor/envs/env_problem_test.py
rename to tensor2tensor/envs/gym_env_problem_test.py
index f5060fb99..7df54d938 100644
--- a/tensor2tensor/envs/env_problem_test.py
+++ b/tensor2tensor/envs/gym_env_problem_test.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for tensor2tensor.envs.env_problem."""
+"""Tests for tensor2tensor.envs.gym_env_problem."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -27,11 +27,12 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import env_problem_utils
+from tensor2tensor.envs import gym_env_problem
 from tensor2tensor.layers import modalities
 import tensorflow as tf
 
 
-class EnvProblemTest(tf.test.TestCase):
+class GymEnvProblemTest(tf.test.TestCase):
 
   def setUp(self):
     self.tmp_dir = os.path.join(tf.test.get_temp_dir(), "tmp_dir")
@@ -41,7 +42,8 @@ def tearDown(self):
     tf.gfile.DeleteRecursively(self.tmp_dir)
 
   def test_setup(self):
-    ep = env_problem.EnvProblem(base_env_name="CartPole-v0", batch_size=5)
+    ep = gym_env_problem.GymEnvProblem(
+        base_env_name="CartPole-v0", batch_size=5)
     # Checks that environments were created and they are `batch_size` in number.
     ep.assert_common_preconditions()
 
@@ -64,7 +66,7 @@ def test_setup(self):
   def test_reward_range(self):
     # Passing reward_range=None means take the reward range of the underlying
     # environment as the reward range.
-    ep = env_problem.EnvProblem(
+    ep = gym_env_problem.GymEnvProblem(
         base_env_name="FrozenLake-v0", batch_size=5, reward_range=None)
     ep.assert_common_preconditions()
 
@@ -77,7 +79,7 @@ def test_reward_range(self):
 
   def test_default_processed_rewards_discrete(self):
     # This differs in the above because it has a Tuple observation space.
-    ep = env_problem.EnvProblem(
+    ep = gym_env_problem.GymEnvProblem(
         base_env_name="KellyCoinflip-v0", batch_size=5, reward_range=None)
     ep.assert_common_preconditions()
 
@@ -103,7 +105,7 @@ def test_default_processed_rewards_discrete(self):
   def test_interaction_with_env(self):
     batch_size = 5
     reward_range = (-1, 1)
-    ep = env_problem.EnvProblem(
+    ep = gym_env_problem.GymEnvProblem(
         base_env_name="KellyCoinflip-v0",
         batch_size=batch_size,
         reward_range=reward_range)
@@ -185,22 +187,22 @@ def play_env(self,
                base_env_name=None,
                batch_size=5,
                reward_range=None):
-    """Creates `EnvProblem` with the given arguments and plays it randomly.
+    """Creates `GymEnvProblem` with the given arguments and plays it randomly.
 
     Args:
       env: optional env.
       nsteps: plays the env randomly for nsteps.
-      base_env_name: passed to EnvProblem's init.
-      batch_size: passed to EnvProblem's init.
-      reward_range: passed to EnvProblem's init.
+      base_env_name: passed to GymEnvProblem's init.
+      batch_size: passed to GymEnvProblem's init.
+      reward_range: passed to GymEnvProblem's init.
 
     Returns:
-      tuple of env_problem, number of trajectories done, number of trajectories
-      done in the last step.
+      tuple of gym_env_problem, number of trajectories done,
+      number of trajectories done in the last step.
     """
 
     if env is None:
-      env = env_problem.EnvProblem(
+      env = gym_env_problem.GymEnvProblem(
           base_env_name=base_env_name,
           batch_size=batch_size,
           reward_range=reward_range)
@@ -272,7 +274,7 @@ def test_generate_data(self):
   def test_problem_dataset_works(self):
 
     # We need to derive this class to set the required methods.
-    class TestEnv(env_problem.EnvProblem):
+    class TestEnv(gym_env_problem.GymEnvProblem):
       name = "TestEnv"
 
       @property
@@ -356,7 +358,7 @@ def test_resets_properly(self):
     reward_range = (-1, 1)
     nsteps = 100
 
-    env = env_problem.EnvProblem(
+    env = gym_env_problem.GymEnvProblem(
         base_env_name=base_env_name,
         batch_size=batch_size,
         reward_range=reward_range)
diff --git a/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
index dba768495..a26c3cd97 100644
--- a/tensor2tensor/envs/rendered_env_problem.py
+++ b/tensor2tensor/envs/rendered_env_problem.py
@@ -23,6 +23,7 @@
 import six
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.envs import env_problem
+from tensor2tensor.envs import gym_env_problem
 import tensorflow as tf
 
 _IMAGE_ENCODED_FIELD = "image/encoded"
@@ -34,7 +35,8 @@
 _FORMAT = "png"
 
 
-class RenderedEnvProblem(env_problem.EnvProblem, video_utils.VideoProblem):
+class RenderedEnvProblem(gym_env_problem.GymEnvProblem,
+                         video_utils.VideoProblem):
   """An `EnvProblem` when observations are RGB arrays.
 
   This takes care of wrapping a rendered gym environment to behave like a
@@ -50,11 +52,11 @@ class RenderedEnvProblem(env_problem.EnvProblem, video_utils.VideoProblem):
 
   def __init__(self, *args, **kwargs):
     """Initialize by calling both parents' constructors."""
-    env_problem.EnvProblem.__init__(self, *args, **kwargs)
+    gym_env_problem.GymEnvProblem.__init__(self, *args, **kwargs)
     video_utils.VideoProblem.__init__(self)
 
   def initialize_environments(self, batch_size=1):
-    env_problem.EnvProblem.initialize_environments(self, batch_size)
+    gym_env_problem.GymEnvProblem.initialize_environments(self, batch_size)
     # Assert the underlying gym environment has correct observation space
     assert len(self.observation_spec.shape) == 3
 
@@ -62,7 +64,8 @@ def example_reading_spec(self):
     """Return a mix of env and video data fields and decoders."""
     video_fields, video_decoders = (
         video_utils.VideoProblem.example_reading_spec(self))
-    env_fields, env_decoders = env_problem.EnvProblem.example_reading_spec(self)
+    env_fields, env_decoders = (
+        gym_env_problem.GymEnvProblem.example_reading_spec(self))
 
     # Remove raw observations field since we want to capture them as videos.
     env_fields.pop(env_problem.OBSERVATION_FIELD)
@@ -81,7 +84,7 @@ def example_reading_spec(self):
 
   def _generate_time_steps(self, trajectory_list):
     """Transforms time step observations to frames of a video."""
-    for time_step in env_problem.EnvProblem._generate_time_steps(
+    for time_step in gym_env_problem.GymEnvProblem._generate_time_steps(
         self, trajectory_list):
       # Convert the rendered observations from numpy to png format.
       frame_np = np.array(time_step.pop(env_problem.OBSERVATION_FIELD))
diff --git a/tensor2tensor/envs/tic_tac_toe_env_problem.py b/tensor2tensor/envs/tic_tac_toe_env_problem.py
index ae5b877d8..c609245f8 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem.py
@@ -19,13 +19,13 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.envs import env_problem
+from tensor2tensor.envs import gym_env_problem
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
 
 @registry.register_env_problem
-class TicTacToeEnvProblem(env_problem.EnvProblem):
+class TicTacToeEnvProblem(gym_env_problem.GymEnvProblem):
   """Plays `batch_size` games of tic-tac-toe."""
 
   def __init__(self):
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 167c34464..998739d46 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -49,7 +49,7 @@
 import jax
 from jax.config import config
 import numpy as onp
-from tensor2tensor.envs import env_problem
+from tensor2tensor.envs import gym_env_problem
 from tensor2tensor.envs import rendered_env_problem
 from tensor2tensor.rl import gym_utils
 from tensor2tensor.trax import layers
@@ -130,7 +130,7 @@ def make_env(batch_size=8, **env_kwargs):
 
   # No resizing needed, so let's be on the normal EnvProblem.
   if not FLAGS.resize:  # None or False
-    return env_problem.EnvProblem(
+    return gym_env_problem.GymEnvProblem(
         base_env_name=FLAGS.env_problem_name,
         batch_size=batch_size,
         parallelism=parallelism,
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index 580da77b0..fa65a22ac 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -27,7 +27,7 @@
 import gin
 import numpy as np
 
-from tensor2tensor.envs import env_problem
+from tensor2tensor.envs import gym_env_problem
 from tensor2tensor.rl import gym_utils
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import layers
@@ -52,11 +52,11 @@ def get_wrapped_env(self, name="CartPole-v0", max_episode_steps=2):
             "output_dtype": None,
         })
 
-    return env_problem.EnvProblem(base_env_name=name,
-                                  batch_size=1,
-                                  env_wrapper_fn=wrapper_fn,
-                                  reward_range=(-1, 1),
-                                  discrete_rewards=False)
+    return gym_env_problem.GymEnvProblem(base_env_name=name,
+                                         batch_size=1,
+                                         env_wrapper_fn=wrapper_fn,
+                                         reward_range=(-1, 1),
+                                         discrete_rewards=False)
 
   @contextlib.contextmanager
   def tmp_dir(self):

From 9a7f09ef0b66676c50c39799b29e9a8bd5378716 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 15 Jul 2019 15:03:53 -0700
Subject: [PATCH 2198/2720] fix full transformer and transformer wmt ende
 training hparams

PiperOrigin-RevId: 258245126
---
 .../trax/configs/transformer_wmt_ende_8gb.gin          | 10 +++++++++-
 tensor2tensor/trax/models/transformer.py               |  1 +
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
index f5ceefbc6..edabb4b50 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
@@ -22,7 +22,8 @@ masked_mean.mask_id = 0
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-MultifactorSchedule.constant = 0.1
+# 0.044 ~= 512^-0.5 = feature_depth^-0.5
+MultifactorSchedule.constant = 0.044
 MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
 MultifactorSchedule.warmup_steps = 8000
 
@@ -39,6 +40,13 @@ train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.Transformer
 train.train_steps = 500000
+train.optimizer = @trax.optimizers.Adam
+
+# Parameters for Adam:
+# ==============================================================================
+Adam.b1 = 0.9
+Adam.b2 = 0.98
+Adam.eps = 1e-9
 
 # Parameters for Transformer:
 # ==============================================================================
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 8b821be53..3279ea002 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -212,6 +212,7 @@ def EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode):
       tl.Dropout(rate=dropout, mode=mode),      # vecs_d          ..... ......
   ]
   decoder_to_encoder_attention = [        # vecs_d        masks         vecs_e
+      tl.LayerNorm(),                     # vecs_d        masks         vecs_e
       tl.Parallel([], [], tl.Dup()),      # ______        _____  vecs_e vecs_e
       tl.Parallel([], tl.Swap()),         # ______        vecs_e masks  ......
       tl.Parallel([], tl.Dup()),          # ______ vecs_e vecs_e .....  ......

From d124912c80fa40f2a6b4afca95ba3d54b6b7a7c6 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 15 Jul 2019 18:13:16 -0700
Subject: [PATCH 2199/2720] Internal

PiperOrigin-RevId: 258277479
---
 tensor2tensor/layers/common_layers.py | 159 --------------------------
 1 file changed, 159 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 09c3a5d34..a59d508fd 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1868,165 +1868,6 @@ def padded_cross_entropy(logits,
     return tf.reduce_sum(xent * weights), tf.reduce_sum(weights)
 
 
-def gather_tensor_by_mixture_index(value,
-                                   mixture_indices,
-                                   batch_size,
-                                   num_mixtures,
-                                   reshape=True):
-  """Gather the elements of a tensor, based on the mixture element id provided.
-
-  The tensor should be shaped as (num_mixtures * batch_size, dim2, dim3...),
-  and the mixture indices should be (batch_size), holding one mixture_id for
-  each element in the batch
-
-  Args:
-    value: a `Tensor` with shape `[num_mixtures * batch, dim2, dim3 ...]`. If
-    reshape is false, it should be [num_mixtures, batch, dim3, dim4 ..]
-    mixture_indices: `[batch_size]`.
-    batch_size: an int `Scalar`.
-    num_mixtures: an int `Scalar`.
-    reshape: bool
-
-  Returns:
-    selected_values: a `Tensor`.  Selected values from original tensor
-
-  """
-  original_shape = shape_list(value)
-  individual_element_indices = tf.range(batch_size)
-  stacked_mixture_element_indices = tf.stack(
-      (mixture_indices, individual_element_indices), -1)
-  if reshape:
-    value = tf.reshape(value, [num_mixtures, -1] + original_shape[1:])
-  selected_values = tf.gather_nd(value, stacked_mixture_element_indices)
-  return selected_values
-
-
-def padded_cross_entropy_mixture(logits,
-                                 labels,
-                                 label_smoothing,
-                                 num_mixtures,
-                                 weights_fn=weights_nonzero,
-                                 reduce_sum=False,
-                                 cutoff=0.0,
-                                 gaussian=False,
-                                 return_best_logits=False):
-  """Compute cross-entropy assuming 0s are padding.
-
-  Computes a loss numerator (the sum of losses), and loss denominator
-  (the number of non-padding tokens).
-
-  Computes cross-entropy for each mixture, and returns the corresponding values
-  for the mixture with the highest probability
-
-  Args:
-    logits: `Tensor` with shape `[batch * num_mixtures, timesteps, vocab_size]`.
-      optionally a FactoredTensor.
-    labels: an integer `Tensor` with shape `[batch, timesteps]`.
-    label_smoothing: a floating point `Scalar`.
-    num_mixtures: an integer.
-    weights_fn: A function from labels to weights.
-    reduce_sum: a Boolean, whether to sum at the end or not.
-    cutoff: a float, at which point to have no loss.
-    gaussian: If true, use a Gaussian distribution for label smoothing
-    return_best_logits: If true, return the logits of the mixture with highest
-    probabilities for an example
-
-  Returns:
-    loss_numerator: a `Scalar`.  Sum of losses.
-    loss_denominator: a `Scalar.  The number of non-padding target tokens.
-
-  Raises:
-    ValueError: in case of unsupported argument types.
-  """
-  # TODO(karishmamalkan): Fix documentation and refactor name
-  (logits, mixture_labels, supervised_mode) = logits
-
-  logit_shapes = shape_list(
-      logits)  # batch_size * num_mixtures, timesteps, 1, 1, vocab_size
-  batch_size = tf.cast(logit_shapes[0] / num_mixtures, dtype=tf.int32)
-
-  new_shape_for_xent = [num_mixtures] + shape_list(labels)
-  labels = tf.tile(labels, [num_mixtures, 1, 1, 1])
-
-  # get xent loss for all mixtures
-  xent, weights = padded_cross_entropy(logits, labels, label_smoothing,
-                                       weights_fn, reduce_sum, cutoff, gaussian)
-
-  # reshape xent and weights to have the num_mixtures as first dimension
-  xent = tf.reshape(xent, new_shape_for_xent)
-  weights = tf.reshape(weights, new_shape_for_xent[:-1])
-
-  # sum up sentence neg log probs
-  xent = tf.reduce_sum(xent, axis=2)
-
-  # if we need to compute the best logits
-  if return_best_logits:
-    if not supervised_mode:
-      return_mixture_indices = tf.squeeze(
-          tf.cast(tf.argmin(xent, 0), dtype=tf.int32), axis=[1, 2])
-    else:
-      return_mixture_indices = mixture_labels
-    best_logits = gather_tensor_by_mixture_index(logits, return_mixture_indices,
-                                                 batch_size, num_mixtures)
-
-  with tf.control_dependencies([
-      tf.assert_equal(
-          tf.shape(xent)[:3], [num_mixtures, batch_size, 1],
-          message="Each batch element should have a probability value for "
-          "each mixture element"
-      )
-  ]):
-    best_mixtures = tf.squeeze(
-        tf.cast(tf.argmin(xent, 0), dtype=tf.int32), axis=[1, 2])
-    if mixture_labels is not None:
-      mixture_accuracy = tf.metrics.accuracy(
-          mixture_labels, best_mixtures, name="mixture_accuracy")
-      tf.summary.scalar("mixture_acc_plot", mixture_accuracy[1])
-
-    # plot a summary for the difference between the top 2 losses
-    if num_mixtures > 1:
-      xent_reshaped = tf.transpose(tf.squeeze(xent, axis=[2, 3]), perm=[1, 0])
-      top_2_mixtures = tf.reduce_mean(
-          -tf.math.top_k(-xent_reshaped, k=2)[0], axis=0)
-      tf.summary.scalar("difference_top_2",
-                        top_2_mixtures[0] - top_2_mixtures[1])
-
-    if supervised_mode:
-      xent_min = gather_tensor_by_mixture_index(
-          xent, mixture_labels, batch_size, num_mixtures, reshape=False)
-    else:
-      xent_min = tf.reduce_min(xent, axis=0)
-    xent_max = tf.reduce_max(xent, axis=0)
-    weights = tf.reduce_mean(weights, axis=0)
-
-  with tf.control_dependencies([
-      tf.assert_equal(
-          tf.shape(xent_min)[0], [batch_size],
-          message="There should be batch_size elements after selecting best "
-          "mixture probabilities"
-      )
-  ]):
-    summed_xent_min = tf.reduce_sum(xent_min)
-    summed_xent_max = tf.reduce_sum(xent_max)
-    summed_weights = tf.reduce_sum(weights)
-
-    for mixture in range(num_mixtures):
-      num_assigned_mixtures = tf.reduce_sum(
-          tf.cast(tf.equal(best_mixtures, mixture), tf.int32))
-      tf.summary.scalar("assigned_mixture_%d" % (mixture),
-                        num_assigned_mixtures / batch_size)
-
-    tf.summary.scalar("selected_mixture_xents_value",
-                      summed_xent_min / summed_weights)
-    tf.summary.scalar("max_mixture_xents_value",
-                      summed_xent_max / summed_weights)
-
-  if return_best_logits:
-    return summed_xent_min, summed_weights, best_logits, return_mixture_indices
-  else:
-    return summed_xent_min, summed_weights
-
-
 def _weights_one_third(labels):
   """Returns Tensor of shape [batch, height, width]. Each element is 1/3."""
   return tf.ones(tf.shape(labels)[:-1]) / 3.

From f48bda1165d26f63421fe9f3f01fc5974209c9a5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 15 Jul 2019 18:31:42 -0700
Subject: [PATCH 2200/2720] Define a Neural Process Model.

Classes for training and inference with (attentive) neural processes.

PiperOrigin-RevId: 258279491
---
 tensor2tensor/layers/gaussian_process.py      | 351 ++++++++++++++++++
 tensor2tensor/layers/gaussian_process_test.py | 163 +++++++-
 2 files changed, 513 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/gaussian_process.py b/tensor2tensor/layers/gaussian_process.py
index c0714457b..dc42cf6ae 100644
--- a/tensor2tensor/layers/gaussian_process.py
+++ b/tensor2tensor/layers/gaussian_process.py
@@ -462,3 +462,354 @@ def fit(self, x=None, y=None):
         adjoint=True) / self.noise_variance
     # TODO(trandustin): To be fully Keras-compatible, return History object.
     return
+
+
+def batch_mlp(inputs, hidden_sizes):
+  """Apply MLP to the final axis of a 3D tensor.
+
+  Args:
+    inputs: input Tensor of shape [batch_size, n, d_in].
+    hidden_sizes: An iterable containing the hidden layer sizes of the MLP.
+
+  Returns:
+    Tensor of shape [batch_size, n, d_out] where d_out = output_sizes[-1].
+  """
+  batch_size, _, filter_size = inputs.shape.as_list()
+  hidden = tf.reshape(inputs, (-1, filter_size))
+
+  for size in hidden_sizes[:-1]:
+    hidden = tf.keras.layers.Dense(size, activation=tf.nn.relu)(hidden)
+
+  output = tf.keras.layers.Dense(hidden_sizes[-1], activation=None)(hidden)
+  output = tf.reshape(output, (batch_size, -1, hidden_sizes[-1]))
+  return output
+
+
+# TODO(adityagrover): Reimplement using preexisting attention routines in T2T
+def uniform_attention(q, v):
+  """Computes uniform attention. Equivalent to neural process.
+
+  Args:
+    q: queries. Tensor of shape [batch_size, m, d_k].
+    v: values. Tensor of shape [batch_size, n, d_v].
+
+  Returns:
+    Tensor of shape [batch_size, m, d_v].
+  """
+  total_points = tf.shape(q)[1]
+  rep = tf.reduce_mean(v, axis=1, keepdims=True)  # [batch_size, 1, d_v]
+  rep = tf.tile(rep, [1, total_points, 1])
+  return rep
+
+
+def laplace_attention(q, k, v, scale, normalise):
+  """Computes laplace exponential attention.
+
+  Args:
+    q: queries. Tensor of shape [batch_size, m, d_k].
+    k: keys. Tensor of shape [batch_size, n, d_k].
+    v: values. Tensor of shape [batch_size, n, d_v].
+    scale: float that scales the L1 distance.
+    normalise: Boolean that determines whether weights sum to 1.
+
+  Returns:
+    Tensor of shape [batch_size, m, d_v].
+  """
+  k = tf.expand_dims(k, axis=1)  # [batch_size, 1, n, d_k]
+  q = tf.expand_dims(q, axis=2)  # [batch_size, m, 1, d_k]
+  unnorm_weights = - tf.abs((k - q) / scale)  # [batch_size, m, n, d_k]
+  unnorm_weights = tf.reduce_sum(unnorm_weights, axis=-1)  # [batch_size, m, n]
+  if normalise:
+    weight_fn = tf.nn.softmax
+  else:
+    weight_fn = lambda x: 1 + tf.tanh(x)
+  weights = weight_fn(unnorm_weights)  # [batch_size, m, n]
+  rep = tf.einsum('bik,bkj->bij', weights, v)  # [batch_size, m, d_v]
+  return rep
+
+
+def dot_product_attention(q, k, v, normalise):
+  """Computes dot product attention.
+
+  Args:
+    q: queries. Tensor of  shape [batch_size, m, d_k].
+    k: keys. Tensor of shape [batch_size, n, d_k].
+    v: values. Tensor of shape [batch_size, n, d_v].
+    normalise: Boolean that determines whether weights sum to 1.
+
+  Returns:
+    Tensor of shape [batch_size, m, d_v].
+  """
+  d_k = tf.shape(q)[-1]
+  scale = tf.sqrt(tf.cast(d_k, tf.float32))
+  unnorm_weights = tf.einsum('bjk,bik->bij', k, q) / scale  # [batch_size,m,n]
+  if normalise:
+    weight_fn = tf.nn.softmax
+  else:
+    weight_fn = tf.sigmoid
+  weights = weight_fn(unnorm_weights)  # [batch_size,m,n]
+  rep = tf.einsum('bik,bkj->bij', weights, v)  # [batch_size,m,d_v]
+  return rep
+
+
+def multihead_attention(q, k, v, num_heads=8):
+  """Computes multi-head attention.
+
+  Args:
+    q: queries. Tensor of  shape [batch_size, m, d_k].
+    k: keys. Tensor of shape [batch_size, n, d_k].
+    v: values. Tensor of shape [batch_size, n, d_v].
+    num_heads: number of heads. Should divide d_v.
+
+  Returns:
+    Tensor of shape [batch_size, m, d_v].
+  """
+  d_k = q.get_shape().as_list()[-1]
+  d_v = v.get_shape().as_list()[-1]
+  head_size = int(d_v / num_heads)
+  key_initializer = tf.random_normal_initializer(stddev=d_k**-0.5)
+  value_initializer = tf.random_normal_initializer(stddev=d_v**-0.5)
+  rep = tf.constant(0.0)
+  for h in range(num_heads):
+    o = dot_product_attention(
+        tf.keras.layers.Conv1D(
+            head_size, 1, kernel_initializer=key_initializer,
+            name='wq%d' % h, use_bias=False, padding='VALID')(q),
+        tf.keras.layers.Conv1D(
+            head_size, 1, kernel_initializer=key_initializer,
+            name='wk%d' % h, use_bias=False, padding='VALID')(k),
+        tf.keras.layers.Conv1D(
+            head_size, 1, kernel_initializer=key_initializer,
+            name='wv%d' % h, use_bias=False, padding='VALID')(v),
+        normalise=True)
+    rep += tf.keras.layers.Conv1D(d_v, 1, kernel_initializer=value_initializer,
+                                  name='wo%d' % h, use_bias=False,
+                                  padding='VALID')(o)
+  return rep
+
+
+# TODO(adityagrover): Implement via T2T.
+class Attention(object):
+  """The Attention module."""
+
+  def __init__(self, rep, output_sizes, att_type, scale=1., normalise=True,
+               num_heads=8):
+    """Creates a attention module.
+
+    Takes in context inputs, target inputs and
+    representations of each context input/output pair
+    to output an aggregated representation of the context data.
+
+    Args:
+      rep: transformation to apply to contexts before computing attention.
+          One of: ['identity', 'mlp'].
+      output_sizes: list of number of hidden units per layer of mlp.
+          Used only if rep == 'mlp'.
+      att_type: type of attention. One of the following:
+          ['uniform', 'laplace', 'dot_product', 'multihead']
+      scale: scale of attention.
+      normalise: Boolean determining whether to:
+          1. apply softmax to weights so they sum to 1 across context pts or
+          2. apply custom transformation to have weights in [0, 1].
+      num_heads: number of heads for multihead.
+    """
+    self._rep = rep
+    self._output_sizes = output_sizes
+    self._type = att_type
+    self._scale = scale
+    self._normalise = normalise
+    if self._type == 'multihead':
+      self._num_heads = num_heads
+
+  def __call__(self, x1, x2, r):
+    """Applies attention to create aggregated representation of r.
+
+    Args:
+      x1: Tensor of shape [B ,n1, d_x].
+      x2: Tensor of shape [batch_size, n2, d_x].
+      r: Tensor of shape [batch_size, n1, d].
+
+    Returns:
+      Tensor of shape [batch_size, n2, d]
+
+    Raises:
+      NameError: The argument for rep/type was invalid.
+    """
+    if self._rep == 'identity':
+      k, q = (x1, x2)
+    elif self._rep == 'mlp':
+      k = batch_mlp(x1, self._output_sizes)
+      q = batch_mlp(x2, self._output_sizes)
+    else:
+      raise NameError("'rep' not among ['identity', 'mlp']")
+
+    if self._type == 'uniform':
+      rep = uniform_attention(q, r)
+    elif self._type == 'laplace':
+      rep = laplace_attention(q, k, r, self._scale, self._normalise)
+    elif self._type == 'dot_product':
+      rep = dot_product_attention(q, k, r, self._normalise)
+    elif self._type == 'multihead':
+      rep = multihead_attention(q, k, r, self._num_heads)
+    else:
+      raise NameError(("'att_type' not among ['uniform', 'laplace', "
+                       "'dot_product', 'multihead']"))
+
+    return rep
+
+
+# TODO(adityagrover): Make the encoder and decoder configurable.
+class NeuralProcess(tf.keras.Model):
+  """Attentive Neural Process (Kim et al., 2019; Garnelo et al., 2018)."""
+
+  def __init__(self,
+               latent_encoder_sizes,
+               num_latents,
+               decoder_sizes,
+               use_deterministic_path=True,
+               deterministic_encoder_sizes=None,
+               attention_wrapper=None):
+    """Initializes the Neural Process model.
+
+    Args:
+      latent_encoder_sizes: (list of ints) Hidden layer sizes for latent
+          encoder.
+      num_latents: (int) Dimensionality of global latent variable.
+      decoder_sizes: (list of ints) Hidden layer sizes for decoder
+      use_deterministic_path: (bool) Uses deterministic encoder as well if True.
+      deterministic_encoder_sizes: (list of ints) Hidden layer sizes for
+          deterministic encoder.
+      attention_wrapper: Instance of Attention class to apply for
+          determinitic encoder embedding.
+    """
+    super(NeuralProcess, self).__init__()
+    self._num_latents = num_latents
+    self._latent_encoder_sizes = latent_encoder_sizes
+    self._deterministic_encoder_sizes = deterministic_encoder_sizes
+    self._decoder_sizes = decoder_sizes
+    self._use_deterministic_path = use_deterministic_path
+    self._attention = attention_wrapper
+
+  def latent_encoder(self, x, y):
+    """Encodes the inputs into one representation.
+
+    Args:
+      x: Tensor of shape [batch_size, observations, d_x]. For the prior, these
+         are context x-values. For the posterior, these are target x-values.
+      y: Tensor of shape [batch_size, observations, d_y]. For the prior, these
+         are context y-values. For the posterior, these are target y-values.
+
+    Returns:
+      A normal distribution over tensors of shape [batch_size, num_latents].
+    """
+    encoder_input = tf.concat([x, y], axis=-1)
+    per_example_embedding = batch_mlp(
+        encoder_input, self._latent_encoder_sizes)
+    dataset_embedding = tf.reduce_mean(per_example_embedding, axis=1)
+    hidden = tf.keras.layers.Dense(
+        (self._latent_encoder_sizes[-1] + self._num_latents)//2,
+        activation=tf.nn.relu)(dataset_embedding)
+    loc = tf.keras.layers.Dense(self._num_latents, activation=None)(hidden)
+    untransformed_scale = tf.keras.layers.Dense(self._num_latents,
+                                                activation=None)(hidden)
+    # Constraint scale following Garnelo et al. (2018).
+    scale_diag = 0.1 + 0.9 * tf.sigmoid(untransformed_scale)
+    return ed.MultivariateNormalDiag(loc=loc,
+                                     scale_diag=scale_diag)
+
+  def deterministic_encoder(self, context_x, context_y, target_x):
+    """Encodes the inputs into one representation.
+
+    Args:
+      context_x: Tensor of shape [batch_size, observations, d_x].
+        Observed x-values.
+      context_y: Tensor of shape [batch_size, observations, d_y].
+        Observed y-values.
+      target_x: Tensor of shape [batch_size, target_observations, d_x].
+        Target x-values.
+
+    Returns:
+      Encodings. Tensor of shape [batch_size, target_observations, d].
+    """
+    encoder_input = tf.concat([context_x, context_y], axis=-1)
+    per_example_embedding = batch_mlp(encoder_input,
+                                      self._deterministic_encoder_sizes)
+    per_target_embedding = self._attention(context_x,
+                                           target_x,
+                                           per_example_embedding)
+    return per_target_embedding
+
+  def decoder(self, representation, target_x):
+    """Decodes the individual targets.
+
+    Args:
+      representation: The representation of the context for target predictions.
+          Tensor of shape [batch_size, target_observations, ?].
+      target_x: The x locations for the target query.
+          Tensor of shape [batch_size, target_observations, d_x].
+
+    Returns:
+      dist: A multivariate Gaussian over the target points. A distribution over
+          tensors of shape [batch_size, target_observations, d_y].
+    """
+    decoder_input = tf.concat([representation, target_x], axis=-1)
+    hidden = batch_mlp(decoder_input, self._decoder_sizes)
+    loc, untransformed_scale = tf.split(hidden, 2, axis=-1)
+    scale_diag = 0.1 + 0.9 * tf.nn.softplus(untransformed_scale)
+    return tfp.distributions.MultivariateNormalDiag(loc=loc,
+                                                    scale_diag=scale_diag)
+
+  def __call__(self, query, target_y=None):
+    """Returns the predicted mean and variance at the target points.
+
+    Args:
+      query: Nested tuple containing ((context_x, context_y), target_x) where:
+              context_x is Tensor of shape [batch_size, num_contexts, d_x].
+                  Contains the x values of the context points.
+              context_y is Tensor of shape [batch_size, num_contexts, d_y].
+                  Contains the y values of the context points.
+              target_x is Tensor of shape [batch_size, num_targets, d_x].
+                  Contains the x values of the target points.
+      target_y: The ground truth y values of the target y.
+          Tensor of shape [batch_size, num_targets, d_y].
+
+    Returns:
+      predictive_dist: Predictive posterior distribution over the predicted y.
+    """
+
+    (context_x, context_y), target_x = query
+    num_targets = tf.shape(target_x)[1]
+    prior = self.latent_encoder(context_x, context_y)
+
+    # For training, when target_y is available, use targets for latent encoder.
+    # Note that targets contain contexts by design.
+    # For testing, when target_y unavailable, use contexts for latent encoder.
+    if target_y is None:
+      latent_rep = prior
+    else:
+      posterior = self.latent_encoder(target_x, target_y)
+      latent_rep = posterior
+    latent_rep = tf.tile(tf.expand_dims(latent_rep, axis=1),
+                         [1, num_targets, 1])
+    if self._use_deterministic_path:
+      deterministic_rep = self.deterministic_encoder(context_x,
+                                                     context_y,
+                                                     target_x)
+      representation = tf.concat([deterministic_rep, latent_rep], axis=-1)
+    else:
+      representation = latent_rep
+
+    predictive_dist = self.decoder(representation, target_x)
+
+    if target_y is not None:
+      kl = tf.expand_dims(
+          posterior.distribution.kl_divergence(prior.distribution),
+          -1)
+      self.add_loss(lambda: kl)
+
+    return predictive_dist
+
+  call = __call__
+
+
+
diff --git a/tensor2tensor/layers/gaussian_process_test.py b/tensor2tensor/layers/gaussian_process_test.py
index a4d33d9e3..c89d3666c 100644
--- a/tensor2tensor/layers/gaussian_process_test.py
+++ b/tensor2tensor/layers/gaussian_process_test.py
@@ -127,5 +127,166 @@ def testBayesianLinearModel(self):
     self.assertAllLessEqual(test_predictions_variance_val, noise_variance)
 
 
-if __name__ == "__main__":
+def train_neural_process(model,
+                         train_data,
+                         valid_data,
+                         num_epochs,
+                         batch_size,
+                         learning_rate=1e-4):
+  """Trains the NeuralProcess model.
+
+  Validation data is used for early stopping,
+
+  Args:
+    model: A NeuralProcess Model subclassing Keras model.
+    train_data: (4-tuple of tensors) Values of x and y for contexts and targets.
+    valid_data: 4-tuple of tensors) Values of x and y for contexts and targets.
+    num_epochs: (int) Number of epochs to train the model for.
+    batch_size: (int) Size of batch.
+    learning_rate: (float) Learning rate for Adam optimizer.
+
+  Returns:
+    best_loss: (float) Average validation loss of best early-stopped model.
+  """
+  optimizer = tf.keras.optimizers.Adam(learning_rate)
+  context_x, context_y, target_x, target_y = train_data
+  valid_context_x, valid_context_y, valid_target_x, valid_target_y = valid_data
+  train_data_size = target_x.shape[0]
+  num_updates_per_epoch = train_data_size//batch_size
+  best_loss = np.inf
+  valid_query = (valid_context_x, valid_context_y), valid_target_x
+
+  for _ in range(num_epochs):
+    for i in range(num_updates_per_epoch):
+      start_idx, end_idx = batch_size*i, batch_size*(i+1)
+      batch_query = ((context_x[start_idx:end_idx],
+                      context_y[start_idx:end_idx]),
+                     target_x[start_idx:end_idx])
+      batch_target_y = target_y[start_idx:end_idx]
+      num_targets = tf.shape(batch_target_y)[1]
+      with tf.GradientTape() as tape:
+        predictive_dist = model(batch_query, batch_target_y)
+        log_p = predictive_dist.log_prob(batch_target_y)
+        kl = tf.tile(model.losses[-1], [1, num_targets])
+        loss = -tf.reduce_mean(log_p - kl/tf.cast(num_targets, tf.float32))
+      gradients = tape.gradient(loss, model.trainable_variables)
+      optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+    predictive_dist = model(valid_query, valid_target_y)
+    log_p = predictive_dist.log_prob(valid_target_y)
+    kl = tf.tile(model.losses[-1], [1, tf.shape(valid_target_y)[1]])
+    valid_loss = -tf.reduce_mean(log_p - kl/tf.cast(num_targets, tf.float32))
+    if valid_loss < best_loss:
+      best_loss = valid_loss
+
+  return best_loss
+
+
+class NeuralProcessTest(tf.test.TestCase):
+
+  def setUp(self):
+    # Create a dummy multi-task fake dataset
+    num_train_problems = 32
+    num_valid_problems = 32
+    num_targets = 50
+    num_contexts = 10
+    input_dim = 5
+
+    def _create_fake_dataset(num_problems):
+      target_x = tf.cast(np.random.rand(num_problems,
+                                        num_targets,
+                                        input_dim),
+                         tf.float32)
+      target_y = tf.cast(np.random.rand(num_problems, num_targets, 1),
+                         tf.float32)
+      context_x, context_y = (target_x[:, :num_contexts, :],
+                              target_y[:, :num_contexts, :])
+      return (context_x, context_y, target_x, target_y)
+
+    self.train_data = _create_fake_dataset(num_train_problems)
+    self.valid_data = _create_fake_dataset(num_valid_problems)
+
+    hidden_size = 128
+    num_latents = 16
+
+    np_attention_wrapper = gaussian_process.Attention(
+        rep='identity', output_sizes=None, att_type='uniform')
+    self.np_model = gaussian_process.NeuralProcess(
+        latent_encoder_sizes=[hidden_size]*4,
+        num_latents=num_latents,
+        decoder_sizes=[hidden_size]*2 + [2],
+        use_deterministic_path=True,
+        deterministic_encoder_sizes=[hidden_size]*4,
+        attention_wrapper=np_attention_wrapper)
+
+    anp_attention_wrapper = gaussian_process.Attention(
+        rep='mlp', output_sizes=[hidden_size]*2, att_type='multihead')
+    self.anp_model = gaussian_process.NeuralProcess(
+        latent_encoder_sizes=[hidden_size]*4,
+        num_latents=num_latents,
+        decoder_sizes=[hidden_size]*2 + [2],
+        use_deterministic_path=True,
+        deterministic_encoder_sizes=[hidden_size]*4,
+        attention_wrapper=anp_attention_wrapper)
+
+    self.models = [self.np_model, self.anp_model]
+    self.num_latents, self.hidden_size, self.num_targets = (num_latents,
+                                                            hidden_size,
+                                                            num_targets)
+    super(NeuralProcessTest, self).setUp()
+
+  def test_termination(self):
+    for model in self.models:
+      validation_loss = train_neural_process(
+          model,
+          self.train_data,
+          self.valid_data,
+          num_epochs=2,
+          batch_size=16,
+          learning_rate=1e-4)
+
+      self.assertGreaterEqual(validation_loss, 0.)
+
+  def test_latent_encoder(self):
+    valid_context_x, valid_context_y, _, _ = self.valid_data
+    batch_size = valid_context_x.shape[0]
+
+    for model in self.models:
+      dist = model.latent_encoder(valid_context_x, valid_context_y).distribution
+      self.assertEqual(dist.loc.shape, (batch_size, self.num_latents))
+      self.assertEqual(dist.scale.shape,
+                       (batch_size, self.num_latents, self.num_latents))
+
+  def test_deterministic_encoder(self):
+    valid_context_x, valid_context_y, valid_target_x, _ = self.valid_data
+    batch_size = valid_context_x.shape[0]
+
+    for model in self.models:
+      embedding = model.deterministic_encoder(
+          valid_context_x, valid_context_y, valid_target_x)
+      self.assertEqual(embedding.shape, (batch_size, self.num_targets,
+                                         self.hidden_size))
+
+  def test_call(self):
+    valid_context_x, valid_context_y, valid_target_x, valid_target_y = self.valid_data
+    batch_size = valid_context_x.shape[0]
+
+    for model in self.models:
+      query = (valid_context_x, valid_context_y), valid_target_x
+      # test 'training' when target_y is available
+      predictive_dist = model(query, valid_target_y)
+      self.assertEqual(predictive_dist.loc.shape, (batch_size, self.num_targets,
+                                                   1))
+      self.assertEqual(predictive_dist.scale.shape,
+                       (batch_size, self.num_targets, 1, 1))
+      self.assertAllGreaterEqual(model.losses, 0.)
+
+      # test 'testing' when target_y is unavailable
+      predictive_dist = model(query)
+      self.assertEqual(predictive_dist.loc.shape, (batch_size, self.num_targets,
+                                                   1))
+      self.assertEqual(predictive_dist.scale.shape,
+                       (batch_size, self.num_targets, 1, 1))
+
+
+if __name__ == '__main__':
   tf.test.main()

From 99d7a69081eea10696729c09199260b1578b6217 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Tue, 16 Jul 2019 13:31:57 -0700
Subject: [PATCH 2201/2720] Add layers for hierarchical priors/posteriors.

Notes:

1. I implemented hierarchical distributions as a separate layer. This allows one to apply Flipout or local reparameterization on the independent weights matrix in the noncentered parameterization. We wouldn't be able to do so if we wrote the hierarchical distribution as a weights initializer.

2. We directly use half-Cauchy distributions whereas Louizos et al. (2017) use a Gamma-Inverse Gamma decomposition. When used with log-normal variational distributions, this allows them to get analytic KLs. I don't think analytic vs stochastic KL estimates matter, particularly for the scale parameters. In contrast, the default I chose are half-Cauchy variational distributions which are spikier (more sparse) than log-normal.

3. Conv2DHierarchical uses Flipout instead of local reparameterization which is biased.

Orthogonal to this CL, I noticed our default weight initializers are nonstandard. Namely, the posterior mean initializations are not input/output dependent. I imagine making them dependent is more important than making the posterior scale initializations input/output dependent. In a future CL, I'll play with them and the half-Cauchy initializations.

PiperOrigin-RevId: 258429639
---
 tensor2tensor/keras/initializers.py      |  75 ++++++++
 tensor2tensor/keras/initializers_test.py |  51 ++++++
 tensor2tensor/keras/regularizers.py      |  28 +++
 tensor2tensor/keras/regularizers_test.py |  46 +++++
 tensor2tensor/layers/bayes.py            | 224 +++++++++++++++++++++++
 tensor2tensor/layers/bayes_test.py       |  31 +++-
 6 files changed, 452 insertions(+), 3 deletions(-)
 create mode 100644 tensor2tensor/keras/initializers_test.py
 create mode 100644 tensor2tensor/keras/regularizers_test.py

diff --git a/tensor2tensor/keras/initializers.py b/tensor2tensor/keras/initializers.py
index 92c1bf450..07ba0e822 100644
--- a/tensor2tensor/keras/initializers.py
+++ b/tensor2tensor/keras/initializers.py
@@ -125,6 +125,80 @@ def __call__(self, shape, dtype=None, partition_info=None):
                                       dtype=dtype)
 
 
+class TrainableHalfCauchy(tf.keras.layers.Layer):
+  """Half-Cauchy distribution initializer with trainable parameters."""
+
+  def __init__(self,
+               loc_initializer=tf.keras.initializers.truncated_normal(
+                   stddev=1e-5),
+               scale_initializer=tf.keras.initializers.truncated_normal(
+                   mean=1., stddev=1e-5),
+               loc_regularizer=None,
+               scale_regularizer=None,
+               loc_constraint=None,
+               scale_constraint='positive',
+               seed=None,
+               dtype=tf.float32,
+               **kwargs):
+    """Constructs the initializer."""
+    super(TrainableHalfCauchy, self).__init__(dtype=dtype, **kwargs)
+    self.loc_initializer = get(loc_initializer)
+    self.scale_initializer = get(scale_initializer)
+    self.loc_regularizer = regularizers.get(loc_regularizer)
+    self.scale_regularizer = regularizers.get(scale_regularizer)
+    self.loc_constraint = constraints.get(loc_constraint)
+    self.scale_constraint = constraints.get(scale_constraint)
+    self.seed = seed
+
+  def build(self, shape, dtype=None):
+    if dtype is None:
+      dtype = self.dtype
+
+    self.loc = self.add_weight(
+        'loc',
+        shape=shape,
+        initializer=self.loc_initializer,
+        regularizer=self.loc_regularizer,
+        constraint=self.loc_constraint,
+        dtype=dtype,
+        trainable=True)
+    self.scale = self.add_weight(
+        'scale',
+        shape=shape,
+        initializer=self.scale_initializer,
+        regularizer=self.scale_regularizer,
+        constraint=self.scale_constraint,
+        dtype=dtype,
+        trainable=True)
+    self.built = True
+
+  def __call__(self, shape, dtype=None, partition_info=None):
+    del partition_info  # unused arg
+    if not self.built:
+      self.build(shape, dtype)
+    return ed.Independent(
+        ed.HalfCauchy(loc=self.loc, scale=self.scale).distribution,
+        reinterpreted_batch_ndims=len(shape))
+
+  def get_config(self):
+    return {
+        'loc_initializer':
+            tf.keras.initializers.serialize(self.loc_initializer),
+        'scale_initializer':
+            tf.keras.initializers.serialize(self.scale_initializer),
+        'loc_regularizer':
+            tf.keras.regularizers.serialize(self.loc_regularizer),
+        'scale_regularizer':
+            tf.keras.regularizers.serialize(self.scale_regularizer),
+        'loc_constraint':
+            tf.keras.constraints.serialize(self.loc_constraint),
+        'scale_constraint':
+            tf.keras.constraints.serialize(self.scale_constraint),
+        'seed': self.seed,
+        'dtype': self.dtype,
+    }
+
+
 class TrainableNormal(tf.keras.layers.Layer):
   """Random normal op as an initializer with trainable mean and stddev."""
 
@@ -256,6 +330,7 @@ def get_config(self):
 
 # pylint: disable=invalid-name
 scaled_normal_std_dev = ScaledNormalStdDev
+trainable_half_cauchy = TrainableHalfCauchy
 trainable_normal = TrainableNormal
 trainable_he_normal = TrainableHeNormal
 trainable_glorot_normal = TrainableGlorotNormal
diff --git a/tensor2tensor/keras/initializers_test.py b/tensor2tensor/keras/initializers_test.py
new file mode 100644
index 000000000..145ee86bf
--- /dev/null
+++ b/tensor2tensor/keras/initializers_test.py
@@ -0,0 +1,51 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 Keras-style initializers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.keras import initializers
+from tensor2tensor.utils import test_utils
+
+import tensorflow as tf
+tf.compat.v1.enable_eager_execution()
+
+
+class InitializersTest(tf.test.TestCase):
+
+  @test_utils.run_in_graph_and_eager_modes
+  def testTrainableHalfCauchy(self):
+    shape = (3,)
+    initializer = initializers.get('trainable_half_cauchy')
+    half_cauchy = initializer(shape)
+    self.evaluate(tf.global_variables_initializer())
+    loc_value, scale_value = self.evaluate([
+        # Get distribution of rv -> get distribution of Independent.
+        half_cauchy.distribution.distribution.loc,
+        half_cauchy.distribution.distribution.scale])
+    self.assertAllClose(loc_value, np.zeros(shape), atol=1e-4)
+    self.assertAllClose(scale_value, np.ones(shape), atol=1e-4)
+
+    half_cauchy_value = self.evaluate(half_cauchy)
+    self.assertAllEqual(half_cauchy_value.shape, shape)
+    self.assertAllGreaterEqual(half_cauchy_value, 0.)
+
+
+if __name__ == '__main__':
+  tf.test.main()
diff --git a/tensor2tensor/keras/regularizers.py b/tensor2tensor/keras/regularizers.py
index 4220ababb..5247d32b8 100644
--- a/tensor2tensor/keras/regularizers.py
+++ b/tensor2tensor/keras/regularizers.py
@@ -25,6 +25,33 @@
 from tensorflow_probability import edward2 as ed
 
 
+class HalfCauchyKLDivergence(tf.keras.regularizers.Regularizer):
+  """KL divergence regularizer from an input to the half-Cauchy distribution."""
+
+  def __init__(self, loc=0., scale=1.):
+    """Constructs regularizer where default uses the standard half-Cauchy."""
+    self.loc = loc
+    self.scale = scale
+
+  def __call__(self, x):
+    """Computes regularization using an unbiased Monte Carlo estimate."""
+    prior = ed.Independent(
+        ed.HalfCauchy(
+            loc=tf.broadcast_to(self.loc, x.distribution.event_shape),
+            scale=tf.broadcast_to(self.scale, x.distribution.event_shape)
+        ).distribution,
+        reinterpreted_batch_ndims=len(x.distribution.event_shape))
+    negative_entropy = x.distribution.log_prob(x)
+    cross_entropy = -prior.distribution.log_prob(x)
+    return negative_entropy + cross_entropy
+
+  def get_config(self):
+    return {
+        'loc': self.loc,
+        'scale': self.scale,
+    }
+
+
 class LogUniformKLDivergence(tf.keras.regularizers.Regularizer):
   """KL divergence regularizer from an input to the log-uniform distribution."""
 
@@ -85,6 +112,7 @@ def get_config(self):
 # Compatibility aliases, following tf.keras
 
 # pylint: disable=invalid-name
+half_cauchy_kl_divergence = HalfCauchyKLDivergence
 log_uniform_kl_divergence = LogUniformKLDivergence
 normal_kl_divergence = NormalKLDivergence
 # pylint: enable=invalid-name
diff --git a/tensor2tensor/keras/regularizers_test.py b/tensor2tensor/keras/regularizers_test.py
new file mode 100644
index 000000000..45536cff1
--- /dev/null
+++ b/tensor2tensor/keras/regularizers_test.py
@@ -0,0 +1,46 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 Keras-style regularizers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.keras import regularizers
+from tensor2tensor.utils import test_utils
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+ed = tfp.edward2
+tf.compat.v1.enable_eager_execution()
+
+
+class RegularizersTest(tf.test.TestCase):
+
+  @test_utils.run_in_graph_and_eager_modes
+  def testHalfCauchyKLDivergence(self):
+    shape = (3,)
+    regularizer = regularizers.get('half_cauchy_kl_divergence')
+    variational_posterior = ed.Independent(
+        ed.LogNormal(loc=tf.zeros(shape), scale=1.).distribution,
+        reinterpreted_batch_ndims=1)
+    kl = regularizer(variational_posterior)
+    kl_value = self.evaluate(kl)
+    self.assertGreaterEqual(kl_value, 0.)
+
+
+if __name__ == '__main__':
+  tf.test.main()
diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index db65bbc31..dc61003df 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -198,6 +198,125 @@ def call(self, inputs):
     return outputs
 
 
+class Conv2DHierarchical(Conv2DFlipout):
+  """2D convolution layer with hierarchical distributions.
+
+  The layer computes a variational Bayesian approximation to the distribution
+  over convolutional layers, and where the distribution over weights
+  involves a hierarchical distribution with hidden unit noise coupling vectors
+  of the kernel weight matrix (Louizos et al., 2017),
+
+  ```
+  p(outputs | inputs) = int conv2d(inputs; new_kernel, bias) p(kernel,
+    local_scales, global_scale, bias) dkernel dlocal_scales dglobal_scale dbias.
+  ```
+
+  It does this with a stochastic forward pass, sampling from learnable
+  distributions on the kernel and bias. The kernel is written in non-centered
+  parameterization where
+
+  ```
+  new_kernel[i, j] = kernel[i, j] * local_scale[j] * global_scale.
+  ```
+
+  That is, there is "local" multiplicative noise which couples weights for each
+  output filter. There is also a "global" multiplicative noise which couples the
+  entire weight matrix. By default, the weights are normally distributed and the
+  local and global noises are half-Cauchy distributed; this makes the kernel a
+  horseshoe distribution (Carvalho et al., 2009; Polson and Scott, 2012).
+
+  The estimation uses Flipout for variance reduction with respect to sampling
+  the full weights. Gradients with respect to the distributions' learnable
+  parameters backpropagate via reparameterization. Minimizing cross-entropy
+  plus the layer's losses performs variational minimum description length,
+  i.e., it minimizes an upper bound to the negative marginal likelihood.
+  """
+
+  def __init__(self,
+               filters,
+               kernel_size,
+               strides=(1, 1),
+               padding='valid',
+               data_format=None,
+               dilation_rate=(1, 1),
+               activation=None,
+               use_bias=True,
+               kernel_initializer='trainable_normal',
+               bias_initializer='zeros',
+               local_scale_initializer='trainable_half_cauchy',
+               global_scale_initializer='trainable_half_cauchy',
+               kernel_regularizer='normal_kl_divergence',
+               bias_regularizer=None,
+               local_scale_regularizer='half_cauchy_kl_divergence',
+               global_scale_regularizer=regularizers.HalfCauchyKLDivergence(
+                   scale=1e-5),
+               activity_regularizer=None,
+               kernel_constraint=None,
+               bias_constraint=None,
+               local_scale_constraint='positive',
+               global_scale_constraint='positive',
+               **kwargs):
+    self.local_scale_initializer = initializers.get(local_scale_initializer)
+    self.global_scale_initializer = initializers.get(global_scale_initializer)
+    self.local_scale_regularizer = regularizers.get(local_scale_regularizer)
+    self.global_scale_regularizer = regularizers.get(global_scale_regularizer)
+    self.local_scale_constraint = constraints.get(local_scale_constraint)
+    self.global_scale_constraint = constraints.get(global_scale_constraint)
+    super(Conv2DHierarchical, self).__init__(
+        filters=filters,
+        kernel_size=kernel_size,
+        strides=strides,
+        padding=padding,
+        data_format=data_format,
+        dilation_rate=dilation_rate,
+        activation=activation,
+        use_bias=use_bias,
+        kernel_initializer=initializers.get(kernel_initializer),
+        bias_initializer=initializers.get(bias_initializer),
+        kernel_regularizer=regularizers.get(kernel_regularizer),
+        bias_regularizer=regularizers.get(bias_regularizer),
+        activity_regularizer=regularizers.get(activity_regularizer),
+        kernel_constraint=constraints.get(kernel_constraint),
+        bias_constraint=constraints.get(bias_constraint),
+        **kwargs)
+
+  def build(self, input_shape):
+    self.local_scale = self.add_weight(
+        shape=(self.filters,),
+        name='local_scale',
+        initializer=self.local_scale_initializer,
+        regularizer=self.local_scale_regularizer,
+        constraint=self.local_scale_constraint)
+    self.global_scale = self.add_weight(
+        shape=(),
+        name='global_scale',
+        initializer=self.global_scale_initializer,
+        regularizer=self.global_scale_regularizer,
+        constraint=self.global_scale_constraint)
+    super(Conv2DHierarchical, self).build(input_shape)
+
+  def call_weights(self):
+    """Calls any weights if the initializer is itself a layer."""
+    if isinstance(self.local_scale_initializer, tf.keras.layers.Layer):
+      self.local_scale = self.local_scale_initializer(self.local_scale.shape,
+                                                      self.dtype)
+    if isinstance(self.global_scale_initializer, tf.keras.layers.Layer):
+      self.global_scale = self.global_scale_initializer(self.global_scale.shape,
+                                                        self.dtype)
+    super(Conv2DHierarchical, self).call_weights()
+
+  def call(self, inputs, training=None):
+    self.call_weights()
+    if self.data_format == 'channels_first':
+      local_scale = tf.reshape(self.local_scale, [1, -1, 1, 1])
+    else:
+      local_scale = tf.reshape(self.local_scale, [1, 1, 1, -1])
+    # TODO(trandustin): Figure out what to set local/global scales to at test
+    # time. Means don't exist for Half-Cauchy approximate posteriors.
+    inputs *= local_scale * self.global_scale
+    return super(Conv2DHierarchical, self).call(inputs, training=training)
+
+
 class Conv2DVariationalDropout(Conv2DReparameterization):
   """2D convolution layer with variational dropout (Kingma et al., 2015).
 
@@ -640,6 +759,111 @@ def dropped_inputs():
                      lambda: super(DenseVariationalDropout, self).call(inputs))
 
 
+class DenseHierarchical(DenseVariationalDropout):
+  """Bayesian densely-connected layer with hierarchical distributions.
+
+  The layer computes a variational Bayesian approximation to the distribution
+  over densely-connected layers, and where the distribution over weights
+  involves a hierarchical distribution with hidden unit noise coupling vectors
+  of the kernel weight matrix (Louizos et al., 2017),
+
+  ```
+  p(outputs | inputs) = int dense(inputs; new_kernel, bias) p(kernel,
+    local_scales, global_scale, bias) dkernel dlocal_scales dglobal_scale dbias.
+  ```
+
+  It does this with a stochastic forward pass, sampling from learnable
+  distributions on the kernel and bias. The kernel is written in non-centered
+  parameterization where
+
+  ```
+  new_kernel[i, j] = kernel[i, j] * local_scale[i] * global_scale.
+  ```
+
+  That is, there is "local" multiplicative noise which couples weights for each
+  input neuron. There is also a "global" multiplicative noise which couples the
+  entire weight matrix. By default, the weights are normally distributed and the
+  local and global noises are half-Cauchy distributed; this makes the kernel a
+  horseshoe distribution (Carvalho et al., 2009; Polson and Scott, 2012).
+
+  The estimation uses local reparameterization to avoid sampling the full
+  weights. Gradients with respect to the distributions' learnable parameters
+  backpropagate via reparameterization. Minimizing cross-entropy plus the
+  layer's losses performs variational minimum description length, i.e., it
+  minimizes an upper bound to the negative marginal likelihood.
+  """
+
+  def __init__(self,
+               units,
+               activation=None,
+               use_bias=True,
+               kernel_initializer='trainable_normal',
+               bias_initializer='zero',
+               local_scale_initializer='trainable_half_cauchy',
+               global_scale_initializer='trainable_half_cauchy',
+               kernel_regularizer='normal_kl_divergence',
+               bias_regularizer=None,
+               local_scale_regularizer='half_cauchy_kl_divergence',
+               global_scale_regularizer=regularizers.HalfCauchyKLDivergence(
+                   scale=1e-5),
+               activity_regularizer=None,
+               local_scale_constraint='positive',
+               global_scale_constraint='positive',
+               **kwargs):
+    self.local_scale_initializer = initializers.get(local_scale_initializer)
+    self.global_scale_initializer = initializers.get(global_scale_initializer)
+    self.local_scale_regularizer = regularizers.get(local_scale_regularizer)
+    self.global_scale_regularizer = regularizers.get(global_scale_regularizer)
+    self.local_scale_constraint = constraints.get(local_scale_constraint)
+    self.global_scale_constraint = constraints.get(global_scale_constraint)
+    super(DenseHierarchical, self).__init__(
+        units=units,
+        activation=activation,
+        use_bias=use_bias,
+        kernel_initializer=initializers.get(kernel_initializer),
+        bias_initializer=initializers.get(bias_initializer),
+        kernel_regularizer=regularizers.get(kernel_regularizer),
+        bias_regularizer=regularizers.get(bias_regularizer),
+        activity_regularizer=regularizers.get(activity_regularizer),
+        **kwargs)
+
+  def build(self, input_shape):
+    input_shape = tf.TensorShape(input_shape)
+    input_dim = input_shape[-1]
+    if isinstance(input_dim, tf.Dimension):
+      input_dim = input_dim.value
+    self.local_scale = self.add_weight(
+        shape=(input_dim,),
+        name='local_scale',
+        initializer=self.local_scale_initializer,
+        regularizer=self.local_scale_regularizer,
+        constraint=self.local_scale_constraint)
+    self.global_scale = self.add_weight(
+        shape=(),
+        name='global_scale',
+        initializer=self.global_scale_initializer,
+        regularizer=self.global_scale_regularizer,
+        constraint=self.global_scale_constraint)
+    super(DenseHierarchical, self).build(input_shape)
+
+  def call_weights(self):
+    """Calls any weights if the initializer is itself a layer."""
+    if isinstance(self.local_scale_initializer, tf.keras.layers.Layer):
+      self.local_scale = self.local_scale_initializer(self.local_scale.shape,
+                                                      self.dtype)
+    if isinstance(self.global_scale_initializer, tf.keras.layers.Layer):
+      self.global_scale = self.global_scale_initializer(self.global_scale.shape,
+                                                        self.dtype)
+    super(DenseHierarchical, self).call_weights()
+
+  def call(self, inputs, training=None):
+    self.call_weights()
+    # TODO(trandustin): Figure out what to set local/global scales to at test
+    # time. Means don't exist for Half-Cauchy approximate posteriors.
+    inputs *= self.local_scale[tf.newaxis, :] * self.global_scale
+    return super(DenseHierarchical, self).call(inputs, training=training)
+
+
 @add_weight
 class LSTMCellReparameterization(tf.keras.layers.LSTMCell):
   """Bayesian LSTM cell class estimated via reparameterization.
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index cf243326b..3319063cf 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -47,6 +47,18 @@ class BayesTest(parameterized.TestCase, tf.test.TestCase):
 #        "kernel_initializer": "zeros",
 #        "bias_initializer": "trainable_normal",
 #        "all_close": False},
+#       {"layer": bayes.Conv2DHierarchical,
+#        "kernel_initializer": "zeros",
+#        "bias_initializer": "zeros",
+#        "all_close": True},
+#       {"layer": bayes.Conv2DHierarchical,
+#        "kernel_initializer": "trainable_normal",
+#        "bias_initializer": "zeros",
+#        "all_close": False},
+#       {"layer": bayes.Conv2DHierarchical,
+#        "kernel_initializer": "zeros",
+#        "bias_initializer": "trainable_normal",
+#        "all_close": False},
 #       {"layer": bayes.Conv2DReparameterization,
 #        "kernel_initializer": "zeros",
 #        "bias_initializer": "zeros",
@@ -99,6 +111,7 @@ class BayesTest(parameterized.TestCase, tf.test.TestCase):
 #
 #   @parameterized.parameters(
 #       {"layer": bayes.Conv2DFlipout},
+#       {"layer": bayes.Conv2DHierarchical},
 #       {"layer": bayes.Conv2DReparameterization},
 #       {"layer": bayes.Conv2DVariationalDropout},
 #   )
@@ -114,7 +127,10 @@ class BayesTest(parameterized.TestCase, tf.test.TestCase):
 #     self.evaluate(tf.global_variables_initializer())
 #     res = self.evaluate(outputs)
 #     self.assertEqual(res.shape, (3, 2))
-#     self.assertLen(model.losses, 1)
+#     if layer == bayes.Conv2DHierarchical:
+#       self.assertLen(model.losses, 3)
+#     else:
+#       self.assertLen(model.losses, 1)
 
   @test_utils.run_in_graph_and_eager_modes
   def testTrainableNormalStddevConstraint(self):
@@ -234,6 +250,7 @@ def take_mean(f, *args, **kwargs):
       {"layer": bayes.DenseFlipout},
       {"layer": bayes.DenseReparameterization},
       {"layer": bayes.DenseVariationalDropout},
+      {"layer": bayes.DenseHierarchical},
   )
   @test_utils.run_in_graph_and_eager_modes()
   def testDenseLoss(self, layer):
@@ -287,6 +304,7 @@ def testDenseLoss(self, layer):
       {"layer": bayes.DenseFlipout},
       {"layer": bayes.DenseReparameterization},
       {"layer": bayes.DenseVariationalDropout},
+      {"layer": bayes.DenseHierarchical},
   )
   @test_utils.run_in_graph_and_eager_modes()
   def testDenseModel(self, layer):
@@ -303,13 +321,17 @@ def testDenseModel(self, layer):
     self.evaluate(tf.global_variables_initializer())
     res = self.evaluate(outputs)
     self.assertEqual(res.shape, (3, 2))
-    self.assertLen(model.losses, 1)
+    if layer == bayes.DenseHierarchical:
+      self.assertLen(model.losses, 3)
+    else:
+      self.assertLen(model.losses, 1)
 
   @parameterized.parameters(
       {"layer": bayes.DenseDVI},
       {"layer": bayes.DenseFlipout},
       {"layer": bayes.DenseReparameterization},
       {"layer": bayes.DenseVariationalDropout},
+      {"layer": bayes.DenseHierarchical},
   )
   @test_utils.run_in_graph_and_eager_modes()
   def testDenseSubclass(self, layer):
@@ -329,7 +351,10 @@ class DenseSubclass(layer):
     self.evaluate(tf.global_variables_initializer())
     res = self.evaluate(outputs)
     self.assertEqual(res.shape, (3, 2))
-    self.assertLen(model.losses, 1)
+    if layer == bayes.DenseHierarchical:
+      self.assertLen(model.losses, 3)
+    else:
+      self.assertLen(model.losses, 1)
 
   @test_utils.run_in_graph_and_eager_modes()
   def testDenseDVIIsDeterministic(self):

From 09ce63a22e218180e3c5e66bbc45922749b2033f Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Tue, 16 Jul 2019 15:09:41 -0700
Subject: [PATCH 2202/2720] Fix variational dropout layers if kernel is
 deterministic.

This also removes extraneous calls to call_weights() if kernel is deterministic. Ex: Conv2DVariationalDropout(..., kernel_initializer='he_normal'). call calls call_weights(). Then passes to super which is Conv2DReparameterization. call calls call_weights(). Then passes to super which is Conv2D.

PiperOrigin-RevId: 258449129
---
 tensor2tensor/layers/bayes.py | 34 +++++++++++++++++-----------------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index dc61003df..0d08341fe 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -162,9 +162,9 @@ class Conv2DFlipout(Conv2DReparameterization):
   """
 
   def call(self, inputs):
-    self.call_weights()
     if not isinstance(self.kernel, ed.RandomVariable):
       return super(Conv2DFlipout, self).call(inputs)
+    self.call_weights()
     input_shape = tf.shape(inputs)
     batch_dim = input_shape[0]
     if self.data_format == 'channels_first':
@@ -360,6 +360,8 @@ def __init__(self,
         **kwargs)
 
   def call(self, inputs, training=None):
+    if not isinstance(self.kernel, ed.RandomVariable):
+      return super(Conv2DVariationalDropout, self).call(inputs)
     self.call_weights()
     if training is None:
       training = tf.keras.backend.learning_phase()
@@ -392,18 +394,16 @@ def dropped_inputs():
       return outputs
 
     # Following tf.keras.Dropout, only apply variational dropout if training
-    # flag is True. The kernel must also be a random variable.
+    # flag is True.
     training_value = smart_constant_value(training)
     if training_value is not None:
-      if training_value and isinstance(self.kernel, ed.RandomVariable):
+      if training_value:
         return dropped_inputs()
       else:
         return super(Conv2DVariationalDropout, self).call(inputs)
-    else:
-      return tf.cond(tf.logical_and(training,
-                                    isinstance(self.kernel, ed.RandomVariable)),
-                     dropped_inputs,
-                     lambda: super(Conv2DVariationalDropout, self).call(inputs))
+    return tf.cond(training,
+                   dropped_inputs,
+                   lambda: super(Conv2DVariationalDropout, self).call(inputs))
 
 
 # From `tensorflow/python/framework/smart_cond.py`
@@ -530,11 +530,11 @@ class DenseDVI(DenseReparameterization):
   """
 
   def call(self, inputs):
-    self.call_weights()
     if (not isinstance(inputs, ed.RandomVariable) and
         not isinstance(self.kernel, ed.RandomVariable) and
         not isinstance(self.bias, ed.RandomVariable)):
       return super(DenseDVI, self).call(inputs)
+    self.call_weights()
     inputs_mean, inputs_variance, inputs_covariance = get_moments(inputs)
     kernel_mean, kernel_variance, _ = get_moments(self.kernel)
     if self.use_bias:
@@ -651,9 +651,9 @@ class DenseFlipout(DenseReparameterization):
   """
 
   def call(self, inputs):
-    self.call_weights()
     if not isinstance(self.kernel, ed.RandomVariable):
       return super(DenseFlipout, self).call(inputs)
+    self.call_weights()
     input_shape = tf.shape(inputs)
     sign_input = 2 * tf.random.uniform(input_shape,
                                        minval=0,
@@ -710,6 +710,8 @@ def __init__(self,
         **kwargs)
 
   def call(self, inputs, training=None):
+    if not isinstance(self.kernel, ed.RandomVariable):
+      return super(DenseVariationalDropout, self).call(inputs)
     self.call_weights()
     if training is None:
       training = tf.keras.backend.learning_phase()
@@ -745,18 +747,16 @@ def dropped_inputs():
       return outputs
 
     # Following tf.keras.Dropout, only apply variational dropout if training
-    # flag is True. The kernel must also be a random variable.
+    # flag is True.
     training_value = smart_constant_value(training)
     if training_value is not None:
-      if training_value and isinstance(self.kernel, ed.RandomVariable):
+      if training_value:
         return dropped_inputs()
       else:
         return super(DenseVariationalDropout, self).call(inputs)
-    else:
-      return tf.cond(tf.logical_and(training,
-                                    isinstance(self.kernel, ed.RandomVariable)),
-                     dropped_inputs,
-                     lambda: super(DenseVariationalDropout, self).call(inputs))
+    return tf.cond(training,
+                   dropped_inputs,
+                   lambda: super(DenseVariationalDropout, self).call(inputs))
 
 
 class DenseHierarchical(DenseVariationalDropout):

From 1511949be2ba00f71297ac3f3f2d732ec265da88 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Wed, 17 Jul 2019 10:18:55 -0700
Subject: [PATCH 2203/2720] Fix Bayesian convolutional layers.

TensorFlow Keras made self._convolution_op created dynamically (https://github.com/tensorflow/tensorflow/commit/546308e322a6b95542ba9f3cbb14136128aaad1e). I made the same change within the Bayesian conv layers that don't rely on calls to super.

PiperOrigin-RevId: 258592332
---
 tensor2tensor/layers/bayes.py      |  51 ++++++--
 tensor2tensor/layers/bayes_test.py | 195 ++++++++++++++---------------
 2 files changed, 137 insertions(+), 109 deletions(-)

diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
index 0d08341fe..eff347dc5 100644
--- a/tensor2tensor/layers/bayes.py
+++ b/tensor2tensor/layers/bayes.py
@@ -165,8 +165,32 @@ def call(self, inputs):
     if not isinstance(self.kernel, ed.RandomVariable):
       return super(Conv2DFlipout, self).call(inputs)
     self.call_weights()
+    outputs = self._apply_kernel(inputs)
+    if self.use_bias:
+      if self.data_format == 'channels_first':
+        outputs = tf.nn.bias_add(outputs, self.bias, data_format='NCHW')
+      else:
+        outputs = tf.nn.bias_add(outputs, self.bias, data_format='NHWC')
+    if self.activation is not None:
+      outputs = self.activation(outputs)
+    return outputs
+
+  def _apply_kernel(self, inputs):
     input_shape = tf.shape(inputs)
     batch_dim = input_shape[0]
+    if self._convolution_op is None:
+      padding = self.padding
+      if self.padding == 'causal':
+        padding = 'valid'
+      if not isinstance(padding, (list, tuple)):
+        padding = padding.upper()
+      self._convolution_op = functools.partial(
+          tf.nn.convolution,
+          strides=self.strides,
+          padding=padding,
+          data_format='NHWC' if self.data_format == 'channels_last' else 'NCHW',
+          dilations=self.dilation_rate)
+
     if self.data_format == 'channels_first':
       channels = input_shape[1]
       sign_input_shape = [batch_dim, channels, 1, 1]
@@ -188,13 +212,6 @@ def call(self, inputs):
     outputs = self._convolution_op(inputs, kernel_mean)
     outputs += self._convolution_op(inputs * sign_input,
                                     perturbation) * sign_output
-    if self.use_bias:
-      if self.data_format == 'channels_first':
-        outputs = tf.nn.bias_add(outputs, self.bias, data_format='NCHW')
-      else:
-        outputs = tf.nn.bias_add(outputs, self.bias, data_format='NHWC')
-    if self.activation is not None:
-      outputs = self.activation(outputs)
     return outputs
 
 
@@ -305,16 +322,16 @@ def call_weights(self):
                                                         self.dtype)
     super(Conv2DHierarchical, self).call_weights()
 
-  def call(self, inputs, training=None):
-    self.call_weights()
+  def _apply_kernel(self, inputs):
+    outputs = super(Conv2DHierarchical, self)._apply_kernel(inputs)
     if self.data_format == 'channels_first':
       local_scale = tf.reshape(self.local_scale, [1, -1, 1, 1])
     else:
       local_scale = tf.reshape(self.local_scale, [1, 1, 1, -1])
     # TODO(trandustin): Figure out what to set local/global scales to at test
     # time. Means don't exist for Half-Cauchy approximate posteriors.
-    inputs *= local_scale * self.global_scale
-    return super(Conv2DHierarchical, self).call(inputs, training=training)
+    outputs *= local_scale * self.global_scale
+    return outputs
 
 
 class Conv2DVariationalDropout(Conv2DReparameterization):
@@ -365,6 +382,18 @@ def call(self, inputs, training=None):
     self.call_weights()
     if training is None:
       training = tf.keras.backend.learning_phase()
+    if self._convolution_op is None:
+      padding = self.padding
+      if self.padding == 'causal':
+        padding = 'valid'
+      if not isinstance(padding, (list, tuple)):
+        padding = padding.upper()
+      self._convolution_op = functools.partial(
+          tf.nn.convolution,
+          strides=self.strides,
+          padding=padding,
+          data_format='NHWC' if self.data_format == 'channels_last' else 'NCHW',
+          dilations=self.dilation_rate)
 
     def dropped_inputs():
       """Forward pass with dropout."""
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
index 3319063cf..c123a0b19 100644
--- a/tensor2tensor/layers/bayes_test.py
+++ b/tensor2tensor/layers/bayes_test.py
@@ -33,104 +33,103 @@
 
 class BayesTest(parameterized.TestCase, tf.test.TestCase):
 
-    # TODO(trandustin): Re-enable tests.
-#   @parameterized.parameters(
-#       {"layer": bayes.Conv2DFlipout,
-#        "kernel_initializer": "zeros",
-#        "bias_initializer": "zeros",
-#        "all_close": True},
-#       {"layer": bayes.Conv2DFlipout,
-#        "kernel_initializer": "trainable_normal",
-#        "bias_initializer": "zeros",
-#        "all_close": False},
-#       {"layer": bayes.Conv2DFlipout,
-#        "kernel_initializer": "zeros",
-#        "bias_initializer": "trainable_normal",
-#        "all_close": False},
-#       {"layer": bayes.Conv2DHierarchical,
-#        "kernel_initializer": "zeros",
-#        "bias_initializer": "zeros",
-#        "all_close": True},
-#       {"layer": bayes.Conv2DHierarchical,
-#        "kernel_initializer": "trainable_normal",
-#        "bias_initializer": "zeros",
-#        "all_close": False},
-#       {"layer": bayes.Conv2DHierarchical,
-#        "kernel_initializer": "zeros",
-#        "bias_initializer": "trainable_normal",
-#        "all_close": False},
-#       {"layer": bayes.Conv2DReparameterization,
-#        "kernel_initializer": "zeros",
-#        "bias_initializer": "zeros",
-#        "all_close": True},
-#       {"layer": bayes.Conv2DReparameterization,
-#        "kernel_initializer": "trainable_normal",
-#        "bias_initializer": "zeros",
-#        "all_close": False},
-#       {"layer": bayes.Conv2DReparameterization,
-#        "kernel_initializer": "zeros",
-#        "bias_initializer": "trainable_normal",
-#        "all_close": False},
-#       {"layer": bayes.Conv2DVariationalDropout,
-#        "kernel_initializer": "zeros",
-#        "bias_initializer": "zeros",
-#        "all_close": True},
-#       {"layer": bayes.Conv2DVariationalDropout,
-#        "kernel_initializer": "trainable_normal",
-#        "bias_initializer": "zeros",
-#        "all_close": False},
-#       {"layer": bayes.Conv2DVariationalDropout,
-#        "kernel_initializer": "zeros",
-#        "bias_initializer": "trainable_normal",
-#        "all_close": False},
-#   )
-#   @test_utils.run_in_graph_and_eager_modes
-#   def testConv2DKernel(self,
-#                        layer,
-#                        kernel_initializer,
-#                        bias_initializer,
-#                        all_close):
-#     tf.keras.backend.set_learning_phase(1)  # training time
-#     inputs = tf.to_float(np.random.rand(5, 4, 4, 12))
-#     model = layer(4,
-#                   kernel_size=2,
-#                   kernel_initializer=kernel_initializer,
-#                   bias_initializer=bias_initializer,
-#                   activation=tf.nn.relu)
-#     outputs1 = model(inputs)
-#     outputs2 = model(inputs)
-#     self.evaluate(tf.global_variables_initializer())
-#     res1, res2 = self.evaluate([outputs1, outputs2])
-#     self.assertEqual(res1.shape, (5, 3, 3, 4))
-#     self.assertAllGreaterEqual(res1, 0.)
-#     if all_close:
-#       self.assertAllClose(res1, res2)
-#     else:
-#       self.assertNotAllClose(res1, res2)
-#     model.get_config()
-#
-#   @parameterized.parameters(
-#       {"layer": bayes.Conv2DFlipout},
-#       {"layer": bayes.Conv2DHierarchical},
-#       {"layer": bayes.Conv2DReparameterization},
-#       {"layer": bayes.Conv2DVariationalDropout},
-#   )
-#   @test_utils.run_in_graph_and_eager_modes()
-#   def testConv2DModel(self, layer):
-#     inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
-#     model = tf.keras.Sequential([
-#         layer(3, kernel_size=2, padding="SAME", activation=tf.nn.relu),
-#         tf.keras.layers.Flatten(),
-#         tf.keras.layers.Dense(2, activation=None),
-#     ])
-#     outputs = model(inputs, training=True)
-#     self.evaluate(tf.global_variables_initializer())
-#     res = self.evaluate(outputs)
-#     self.assertEqual(res.shape, (3, 2))
-#     if layer == bayes.Conv2DHierarchical:
-#       self.assertLen(model.losses, 3)
-#     else:
-#       self.assertLen(model.losses, 1)
+  @parameterized.parameters(
+      {"layer": bayes.Conv2DFlipout,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "zeros",
+       "all_close": True},
+      {"layer": bayes.Conv2DFlipout,
+       "kernel_initializer": "trainable_normal",
+       "bias_initializer": "zeros",
+       "all_close": False},
+      {"layer": bayes.Conv2DFlipout,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "trainable_normal",
+       "all_close": False},
+      {"layer": bayes.Conv2DHierarchical,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "zeros",
+       "all_close": True},
+      {"layer": bayes.Conv2DHierarchical,
+       "kernel_initializer": "trainable_normal",
+       "bias_initializer": "zeros",
+       "all_close": False},
+      {"layer": bayes.Conv2DHierarchical,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "trainable_normal",
+       "all_close": False},
+      {"layer": bayes.Conv2DReparameterization,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "zeros",
+       "all_close": True},
+      {"layer": bayes.Conv2DReparameterization,
+       "kernel_initializer": "trainable_normal",
+       "bias_initializer": "zeros",
+       "all_close": False},
+      {"layer": bayes.Conv2DReparameterization,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "trainable_normal",
+       "all_close": False},
+      {"layer": bayes.Conv2DVariationalDropout,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "zeros",
+       "all_close": True},
+      {"layer": bayes.Conv2DVariationalDropout,
+       "kernel_initializer": "trainable_normal",
+       "bias_initializer": "zeros",
+       "all_close": False},
+      {"layer": bayes.Conv2DVariationalDropout,
+       "kernel_initializer": "zeros",
+       "bias_initializer": "trainable_normal",
+       "all_close": False},
+  )
+  @test_utils.run_in_graph_and_eager_modes
+  def testConv2DKernel(self,
+                       layer,
+                       kernel_initializer,
+                       bias_initializer,
+                       all_close):
+    tf.keras.backend.set_learning_phase(1)  # training time
+    inputs = tf.to_float(np.random.rand(5, 4, 4, 12))
+    model = layer(4,
+                  kernel_size=2,
+                  kernel_initializer=kernel_initializer,
+                  bias_initializer=bias_initializer,
+                  activation=tf.nn.relu)
+    outputs1 = model(inputs)
+    outputs2 = model(inputs)
+    self.evaluate(tf.global_variables_initializer())
+    res1, res2 = self.evaluate([outputs1, outputs2])
+    self.assertEqual(res1.shape, (5, 3, 3, 4))
+    self.assertAllGreaterEqual(res1, 0.)
+    if all_close:
+      self.assertAllClose(res1, res2)
+    else:
+      self.assertNotAllClose(res1, res2)
+    model.get_config()
+
+  @parameterized.parameters(
+      {"layer": bayes.Conv2DFlipout},
+      {"layer": bayes.Conv2DHierarchical},
+      {"layer": bayes.Conv2DReparameterization},
+      {"layer": bayes.Conv2DVariationalDropout},
+  )
+  @test_utils.run_in_graph_and_eager_modes()
+  def testConv2DModel(self, layer):
+    inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
+    model = tf.keras.Sequential([
+        layer(3, kernel_size=2, padding="SAME", activation=tf.nn.relu),
+        tf.keras.layers.Flatten(),
+        tf.keras.layers.Dense(2, activation=None),
+    ])
+    outputs = model(inputs, training=True)
+    self.evaluate(tf.global_variables_initializer())
+    res = self.evaluate(outputs)
+    self.assertEqual(res.shape, (3, 2))
+    if layer == bayes.Conv2DHierarchical:
+      self.assertLen(model.losses, 3)
+    else:
+      self.assertLen(model.losses, 1)
 
   @test_utils.run_in_graph_and_eager_modes
   def testTrainableNormalStddevConstraint(self):

From 6e8476f2f04edde42931f9426062ee1265de879d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 17 Jul 2019 10:59:35 -0700
Subject: [PATCH 2204/2720] fix vocab parameter name in wmt ende config

PiperOrigin-RevId: 258601458
---
 tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
index edabb4b50..656144f57 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
@@ -57,4 +57,4 @@ Transformer.max_len = 2048
 Transformer.mode = 'train'
 Transformer.n_heads = 8
 Transformer.n_layers = 6
-Transformer.vocab_size = 33300
+Transformer.input_vocab_size = 33300

From 2c761783a7aacd6800d445d10ad3676a56365514 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 17 Jul 2019 11:59:34 -0700
Subject: [PATCH 2205/2720] Implement SimulatedEnvProblem

PiperOrigin-RevId: 258613974
---
 tensor2tensor/envs/env_problem.py             |  30 ++-
 tensor2tensor/envs/gym_env_problem.py         |  49 +++--
 tensor2tensor/trax/backend.py                 |   2 +
 .../trax/rlax/ppo_training_loop_test.py       | 110 +++++++++-
 .../trax/rlax/simulated_env_problem.py        | 192 ++++++++++++++++++
 .../trax/rlax/simulated_env_problem_test.py   | 120 +++++++++++
 6 files changed, 456 insertions(+), 47 deletions(-)
 create mode 100644 tensor2tensor/trax/rlax/simulated_env_problem.py
 create mode 100644 tensor2tensor/trax/rlax/simulated_env_problem_test.py

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 890f74374..e65bf6405 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -54,6 +54,9 @@ class EnvProblem(Env, problem.Problem):
 
   Subclasses *should* override the following functions:
   - initialize_environments
+  - observation_space
+  - action_space
+  - reward_range
   - _reset
   - _step
   - _render
@@ -95,7 +98,6 @@ class EnvProblem(Env, problem.Problem):
 
   def __init__(self,
                batch_size=None,
-               reward_range=(-np.inf, np.inf),
                discrete_rewards=True,
                parallelism=1,
                **env_kwargs):
@@ -104,9 +106,6 @@ def __init__(self,
     Args:
       batch_size: (int or None) How many envs to make in the non natively
         batched mode.
-      reward_range: (tuple(number, number)) the first element is the minimum
-        reward and the second is the maximum reward, used to clip and process
-        the raw reward in `process_rewards`.
       discrete_rewards: (bool) whether to round the rewards to the nearest
         integer.
       parallelism: (int) If this is greater than one then we run the envs in
@@ -124,18 +123,11 @@ def __init__(self,
     # to an appropriate directory.
     self._agent_id = "default"
 
-    # We clip rewards to this range before processing them further, as described
-    # in `process_rewards`.
-    self._reward_range = reward_range
-
     # If set, we discretize the rewards and treat them as integers.
     self._discrete_rewards = discrete_rewards
 
     self._parallelism = None
 
-    self._observation_space = None
-    self._action_space = None
-
     # A data structure to hold the `batch_size` currently active trajectories
     # and also the ones that are completed, i.e. done.
     self._trajectories = None
@@ -168,10 +160,10 @@ def initialize(self, batch_size=1, **kwargs):
 
     # Assert that *all* the above are now set, we should do this since
     # subclasses can override `initialize_environments`.
-    assert self._envs is not None
-    assert self._observation_space is not None
-    assert self._action_space is not None
-    assert self._reward_range is not None
+    self.assert_common_preconditions()
+    assert self.observation_space is not None
+    assert self.action_space is not None
+    assert self.reward_range is not None
 
   def initialize_environments(self, batch_size=1, parallelism=1, **kwargs):
     """Initializes the environments.
@@ -189,7 +181,7 @@ def assert_common_preconditions(self):
 
   @property
   def observation_space(self):
-    return self._observation_space
+    raise NotImplementedError
 
   @property
   def observation_spec(self):
@@ -210,7 +202,7 @@ def process_observations(self, observations):
 
   @property
   def action_space(self):
-    return self._action_space
+    raise NotImplementedError
 
   @property
   def action_spec(self):
@@ -228,7 +220,9 @@ def num_actions(self):
 
   @property
   def reward_range(self):
-    return self._reward_range
+    # We clip rewards to this range before processing them further, as described
+    # in `process_rewards`.
+    raise NotImplementedError
 
   @property
   def is_reward_range_finite(self):
diff --git a/tensor2tensor/envs/gym_env_problem.py b/tensor2tensor/envs/gym_env_problem.py
index 475c7fc55..fa628f7ad 100644
--- a/tensor2tensor/envs/gym_env_problem.py
+++ b/tensor2tensor/envs/gym_env_problem.py
@@ -75,7 +75,8 @@ class GymEnvProblem(env_problem.EnvProblem):
   the following properties: observation_space, action_space, reward_range.
   """
 
-  def __init__(self, base_env_name=None, env_wrapper_fn=None, **kwargs):
+  def __init__(self, base_env_name=None, env_wrapper_fn=None, reward_range=None,
+               **kwargs):
     """Initializes this class by creating the envs and managing trajectories.
 
     Args:
@@ -83,6 +84,10 @@ def __init__(self, base_env_name=None, env_wrapper_fn=None, **kwargs):
         environment.
       env_wrapper_fn: (callable(env): env) Applies gym wrappers to the base
         environment.
+      reward_range: (tuple(number, number) or None) the first element is the
+        minimum reward and the second is the maximum reward, used to clip and
+        process the raw reward in `process_rewards`. If None, this is inferred
+        from the inner environments.
       **kwargs: (dict) Arguments passed to the base class.
     """
     # Name for the base environment, will be used in `gym.make` in
@@ -96,6 +101,10 @@ def __init__(self, base_env_name=None, env_wrapper_fn=None, **kwargs):
     # to an appropriate directory.
     self._agent_id = "default"
 
+    # We clip rewards to this range before processing them further, as described
+    # in `process_rewards`.
+    self._reward_range = reward_range
+
     # Initialize the environment(s).
 
     # This can either be a list of environments of len `batch_size` or this can
@@ -171,25 +180,6 @@ def initialize_environments(self, batch_size=1, parallelism=1, **kwargs):
     if self._env_wrapper_fn is not None:
       self._envs = list(map(self._env_wrapper_fn, self._envs))
 
-    # If self.observation_space and self.action_space aren't None, then it means
-    # that this is a re-initialization of this class, in that case make sure
-    # that this matches our previous behaviour.
-    if self._observation_space:
-      assert str(self._observation_space) == str(
-          self._envs[0].observation_space)
-    else:
-      # This means that we are initializing this class for the first time.
-      #
-      # We set this equal to the first env's observation space, later on we'll
-      # verify that all envs have the same observation space.
-      self._observation_space = self._envs[0].observation_space
-
-    # Similarly for action_space
-    if self._action_space:
-      assert str(self._action_space) == str(self._envs[0].action_space)
-    else:
-      self._action_space = self._envs[0].action_space
-
     self._verify_same_spaces()
 
     # If self.reward_range is None, i.e. this means that we should take the
@@ -203,6 +193,25 @@ def initialize_environments(self, batch_size=1, parallelism=1, **kwargs):
     # is still valuable to store the trajectories separately.
     self._trajectories = trajectory.BatchTrajectory(batch_size=batch_size)
 
+  def assert_common_preconditions(self):
+    # Asserts on the common pre-conditions of:
+    #  - self._envs is initialized.
+    #  - self._envs is a list.
+    assert self._envs
+    assert isinstance(self._envs, list)
+
+  @property
+  def observation_space(self):
+    return self._envs[0].observation_space
+
+  @property
+  def action_space(self):
+    return self._envs[0].action_space
+
+  @property
+  def reward_range(self):
+    return self._reward_range
+
   def seed(self, seed=None):
     if not self._envs:
       tf.logging.info("`seed` called on non-existent envs, doing nothing.")
diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index dbb251bb6..07cba1542 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -118,6 +118,8 @@ def jax_avg_pool(x, pool_size, strides, padding):
     "name": "numpy",
     "np": onp,
     "jit": (lambda f: f),
+    "random_get_prng": lambda seed: None,
+    "random_split": lambda prng, num=2: (None,) * num,
 }
 
 
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index fa65a22ac..25d5f49cc 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -21,19 +21,25 @@
 
 import contextlib
 import functools
+import itertools
 import os
 import tempfile
 
 import gin
+import gym
 import numpy as np
 
 from tensor2tensor.envs import gym_env_problem
 from tensor2tensor.rl import gym_utils
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import layers
+from tensor2tensor.trax import learning_rate as lr
 from tensor2tensor.trax import models
+from tensor2tensor.trax import optimizers as trax_opt
+from tensor2tensor.trax import trax
 from tensor2tensor.trax.rlax import envs  # pylint: disable=unused-import
 from tensor2tensor.trax.rlax import ppo
+from tensor2tensor.trax.rlax import simulated_env_problem
 from tensorflow import test
 from tensorflow.io import gfile
 
@@ -55,7 +61,6 @@ def get_wrapped_env(self, name="CartPole-v0", max_episode_steps=2):
     return gym_env_problem.GymEnvProblem(base_env_name=name,
                                          batch_size=1,
                                          env_wrapper_fn=wrapper_fn,
-                                         reward_range=(-1, 1),
                                          discrete_rewards=False)
 
   @contextlib.contextmanager
@@ -64,9 +69,7 @@ def tmp_dir(self):
     yield tmp
     gfile.rmtree(tmp)
 
-  def _run_training_loop(self, env_name, output_dir):
-    env = self.get_wrapped_env(env_name, 2)
-    eval_env = self.get_wrapped_env(env_name, 2)
+  def _run_training_loop(self, env, eval_env, output_dir):
     n_epochs = 2
     # Run the training loop.
     ppo.training_loop(
@@ -79,28 +82,117 @@ def _run_training_loop(self, env_name, output_dir):
         policy_and_value_optimizer_fn=ppo.optimizer_fn,
         n_optimizer_steps=1,
         output_dir=output_dir,
-        env_name=env_name,
+        env_name="SomeEnv",
         random_seed=0)
 
   def test_training_loop_cartpole(self):
     with self.tmp_dir() as output_dir:
-      self._run_training_loop("CartPole-v0", output_dir)
+      self._run_training_loop(
+          env=self.get_wrapped_env("CartPole-v0", 2),
+          eval_env=self.get_wrapped_env("CartPole-v0", 2),
+          output_dir=output_dir,
+      )
 
   def test_training_loop_onlinetune(self):
     with self.tmp_dir() as output_dir:
       gin.bind_parameter("OnlineTuneEnv.model", functools.partial(
-          models.MLP, n_hidden_layers=0, n_output_classes=1))
+          models.MLP,
+          n_hidden_layers=0,
+          n_output_classes=1,
+      ))
       gin.bind_parameter("OnlineTuneEnv.inputs", functools.partial(
           trax_inputs.random_inputs,
           input_shape=(1, 1),
           input_dtype=np.float32,
           output_shape=(1, 1),
-          output_dtype=np.float32))
+          output_dtype=np.float32,
+      ))
       gin.bind_parameter("OnlineTuneEnv.train_steps", 2)
       gin.bind_parameter("OnlineTuneEnv.eval_steps", 2)
       gin.bind_parameter(
           "OnlineTuneEnv.output_dir", os.path.join(output_dir, "envs"))
-      self._run_training_loop("OnlineTuneEnv-v0", output_dir)
+      self._run_training_loop(
+          env=self.get_wrapped_env("OnlineTuneEnv-v0", 2),
+          eval_env=self.get_wrapped_env("OnlineTuneEnv-v0", 2),
+          output_dir=output_dir,
+      )
+
+  def test_training_loop_simulated(self):
+    n_actions = 5
+    history_shape = (3, 2, 3)
+    action_shape = (3,)
+    obs_shape = (3, 3)
+    reward_shape = (3, 1)
+
+    def model(mode):
+      del mode
+      return layers.Serial(
+          layers.Parallel(
+              layers.Flatten(),  # Observation stack.
+              layers.Embedding(d_feature=1, vocab_size=n_actions),  # Action.
+          ),
+          layers.Concatenate(),
+          layers.Dense(n_units=1),
+          layers.Dup(),
+          layers.Parallel(
+              layers.Dense(n_units=obs_shape[1]),  # New observation.
+              None,  # Reward.
+          )
+      )
+
+    def inputs(n_devices):
+      del n_devices
+      stream = itertools.repeat((
+          (np.zeros(history_shape), np.zeros(action_shape, dtype=np.int32)),
+          (np.zeros(obs_shape), np.zeros(reward_shape)),
+      ))
+      return trax_inputs.Inputs(
+          train_stream=lambda: stream,
+          train_eval_stream=lambda: stream,
+          eval_stream=lambda: stream,
+          input_shape=(history_shape[1:], action_shape[1:]),
+          input_dtype=(np.float32, np.int32),
+      )
+
+    def loss(*args, **kwargs):
+      del args
+      del kwargs
+      return 0.0
+
+    with self.tmp_dir() as output_dir:
+      # Run fake training just to save the parameters.
+      trainer = trax.Trainer(
+          model=model,
+          loss_fn=loss,
+          inputs=inputs,
+          optimizer=trax_opt.SM3,
+          lr_schedule=lr.MultifactorSchedule,
+          output_dir=output_dir,
+      )
+      trainer.train_epoch(epoch_steps=1, eval_steps=1)
+
+      # Repeat the initial observations over and over again.
+      stream = itertools.repeat(np.zeros(history_shape))
+      env_fn = functools.partial(
+          simulated_env_problem.SimulatedEnvProblem,
+          model=model,
+          history_length=history_shape[1],
+          trajectory_length=3,
+          batch_size=history_shape[0],
+          observation_space=gym.spaces.Box(
+              low=-np.inf, high=np.inf, shape=(obs_shape[1],)),
+          action_space=gym.spaces.Discrete(n=n_actions),
+          reward_range=(-1, 1),
+          discrete_rewards=False,
+          initial_observation_stream=stream,
+          output_dir=output_dir,
+      )
+
+      self._run_training_loop(
+          env=env_fn(),
+          eval_env=env_fn(),
+          output_dir=output_dir,
+      )
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/trax/rlax/simulated_env_problem.py b/tensor2tensor/trax/rlax/simulated_env_problem.py
new file mode 100644
index 000000000..fd3204a41
--- /dev/null
+++ b/tensor2tensor/trax/rlax/simulated_env_problem.py
@@ -0,0 +1,192 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""EnvProblem for environments simulated by a TRAX model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import random
+
+import numpy as np
+
+from tensor2tensor.envs import env_problem
+from tensor2tensor.trax import backend
+from tensor2tensor.trax import trax
+from tensor2tensor.trax.backend import random as jax_random
+
+
+class SimulatedEnvProblem(env_problem.EnvProblem):
+  """EnvProblem for environments simulated by TRAX models.
+
+  Wraps an autoregressive TRAX model of signature
+  (observation_history, action) -> (observation, reward) in an EnvProblem.
+  The model is assumed to take a fixed number of last observations as input
+  and produce a single observation, which is fed back into the model in the
+  next environment step.
+
+  Shape requirements (without the batch dimension):
+    observation: Consistent with observation_space.
+    observation_history: (history_length,) + observation.shape.
+    action: Consistent with action_space.
+    reward: (1,). The singleton dimension is removed in step().
+
+  The initial observations to start the model are taken from
+  initial_observation_stream. This iterator in incremented in every reset().
+
+  A checkpoint saved by the TRAX trainer should be available in output_dir.
+  """
+
+  def __init__(self, model, history_length, trajectory_length, batch_size,
+               observation_space, action_space, reward_range, discrete_rewards,
+               initial_observation_stream, output_dir):
+    """Initializes the env.
+
+    Args:
+      model: TRAX model.
+      history_length: (int) Number of last observations fed into the model.
+      trajectory_length: (int) Length of each trajectory unrolled from the
+        model.
+      batch_size: (int) Number of simulated environments run in parallel.
+      observation_space: (gym.Space) Observation space.
+      action_space: (gym.Space) Action space.
+      reward_range: (tuple) Pair (min_reward, max_reward).
+      discrete_rewards: (bool) Whether to discretize the rewards.
+      initial_observation_stream: Iterator yielding batches of initial
+        observations for the model.
+      output_dir: (str) Output dir.
+    """
+    # TODO(pkozakowski): At some point we will have a "predict" mode which we
+    # should use here. When this happens, change the mode.
+    self._model_predict = backend.jit(model(mode="eval"))
+    self._history_length = history_length
+    self._trajectory_length = trajectory_length
+    self._observation_space = observation_space
+    self._action_space = action_space
+    self._reward_range = reward_range
+    self._output_dir = output_dir
+
+    self._model_params = None
+    self._rng = None
+    self._initial_observation_stream = None
+    self._history = None
+    self._steps = None
+
+    # Call the super's ctor. It will use some of the member fields, so we call
+    # it in the end.
+    super(SimulatedEnvProblem, self).__init__(
+        batch_size=batch_size,
+        discrete_rewards=discrete_rewards,
+        initial_observation_stream=initial_observation_stream,
+    )
+
+    self.seed()
+
+  def initialize_environments(self,
+                              initial_observation_stream,
+                              batch_size=1,
+                              parallelism=1):
+    """Initializes the environments.
+
+    Args:
+      initial_observation_stream: Iterator yielding batches of initial
+        observations for the model.
+      batch_size: (int) Number of environments in a batch.
+      parallelism: (int) Unused.
+    """
+    del parallelism
+
+    model_state = trax.restore_state(self._output_dir)
+    # model_state.params is a pair (model_params, optimizer_state).
+    (self._model_params, _) = model_state.params
+    self._initial_observation_stream = initial_observation_stream
+
+    self._history = None
+    self._steps = np.zeros(batch_size)
+
+  @property
+  def observation_space(self):
+    return self._observation_space
+
+  @property
+  def action_space(self):
+    return self._action_space
+
+  @property
+  def reward_range(self):
+    return self._reward_range
+
+  def seed(self, seed=None):
+    if seed is None:
+      seed = random.randint(0, 2**31 - 1)
+    self._rng = jax_random.get_prng(seed)
+    return super(SimulatedEnvProblem, self).seed(seed=seed)
+
+  def _reset(self, indices):
+    """Resets environments at the given indices.
+
+    Args:
+      indices: list of indices of underlying envs to call reset on.
+
+    Returns:
+      np.ndarray of batched observations from the reset envs.
+    """
+    history = next(self._initial_observation_stream)
+    assert history.shape == ((self._batch_size, self._history_length) +
+                             self.observation_space.shape)
+
+    if self._history is None:
+      # At the first reset, all indices should be triggered.
+      assert set(indices) == set(range(self._batch_size))
+      self._history = np.array(history)
+    else:
+      history = history[indices, ...]
+      self._history[indices, ...] = history
+
+    # Reset the step counters.
+    self._steps[indices] = 0
+
+    # Return just the last timestep at the given indices.
+    return history[:, -1, ...]
+
+  def _step(self, actions):
+    """Takes a step in all environments.
+
+    Args:
+      actions: (np.ndarray) with first dimension equal to the batch size.
+
+    Returns:
+      a tuple of batched raw observations, raw rewards, dones and infos.
+    """
+    # Predict the next observation.
+    (subrng, self._rng) = jax_random.split(self._rng)
+    (observation, reward) = self._model_predict((self._history, actions),
+                                                params=self._model_params,
+                                                rng=subrng)
+
+    # Roll the history one timestep back and append the new observation.
+    self._history = np.roll(self._history, shift=-1, axis=1)
+    self._history[:, -1, ...] = observation
+
+    # Increment the step counters and determine which envs are done.
+    self._steps += 1
+    done = self._steps == self._trajectory_length
+
+    # Call copy() to get the data as numpy arrays.
+    observation = observation.copy()
+    # Reshape the rewards to get rid of the extra dimension.
+    reward = np.squeeze(reward.copy(), axis=1)
+    return (observation, reward, done, {})
diff --git a/tensor2tensor/trax/rlax/simulated_env_problem_test.py b/tensor2tensor/trax/rlax/simulated_env_problem_test.py
new file mode 100644
index 000000000..db2c6ceb9
--- /dev/null
+++ b/tensor2tensor/trax/rlax/simulated_env_problem_test.py
@@ -0,0 +1,120 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.rlax.simulated_env_problem."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gym
+import mock
+import numpy as np
+
+from tensor2tensor.trax import backend
+from tensor2tensor.trax import trax
+from tensor2tensor.trax.rlax import simulated_env_problem
+from tensorflow import test
+
+
+class SimulatedEnvProblemTest(test.TestCase):
+
+  @staticmethod
+  @mock.patch.object(trax, "restore_state", autospec=True)
+  def _create_env(mock_restore_state, model, initial_observations,
+                  trajectory_length):
+    # (model_params, opt_state)
+    mock_restore_state.return_value.params = (None, None)
+    space = gym.spaces.Discrete(100)
+    return simulated_env_problem.SimulatedEnvProblem(
+        model=model,
+        history_length=initial_observations.shape[2],
+        trajectory_length=trajectory_length,
+        batch_size=1,
+        observation_space=space,
+        action_space=space,
+        reward_range=(-1, 1),
+        discrete_rewards=True,
+        initial_observation_stream=iter(initial_observations),
+        output_dir=None,
+    )
+
+  def test_communicates_with_model(self):
+    # Mock model increasing the observation by action, reward is the parity of
+    # the new observation.
+    def mock_transition(inputs, *args, **kwargs):
+      del args
+      del kwargs
+      (observations, actions) = inputs
+      new_observations = observations[:, -1] + actions
+      rewards = np.array([[int(new_observations % 2 == 0)]])
+      return (new_observations, rewards)
+
+    mock_model_fn = mock.MagicMock()
+    mock_model_fn.return_value.side_effect = mock_transition
+    mock_model = mock_model_fn.return_value
+
+    actions_to_take = np.array([[1], [3]])
+    initial_observations = np.array([[[0, 1, 2, 3]]])
+    expected_observations = np.array([[3], [4], [7]])
+    expected_rewards = np.array([[1], [0]])
+    expected_dones = np.array([[False], [True]])
+    expected_histories = np.array([[[0, 1, 2, 3]], [[1, 2, 3, 4]]])
+    expected_actions = actions_to_take
+
+    with backend.use_backend("numpy"):
+      env = self._create_env(  # pylint: disable=no-value-for-parameter
+          model=mock_model_fn,
+          initial_observations=initial_observations,
+          trajectory_length=len(actions_to_take),
+      )
+      actual_observations = [env.reset()]
+      actual_rewards = []
+      actual_dones = []
+      actual_histories = []
+      actual_actions = []
+      for action in actions_to_take:
+        (observation, reward, done, _) = env.step(action)
+        actual_observations.append(observation)
+        actual_rewards.append(reward)
+        actual_dones.append(done)
+        # Mock call is a tuple (args, kwargs). There is one positional argument,
+        # which is a tuple (history, action).
+        (((history, action),), _) = mock_model.call_args
+        actual_actions.append(action)
+        actual_histories.append(history)
+
+    np.testing.assert_array_equal(actual_observations, expected_observations)
+    np.testing.assert_array_equal(actual_rewards, expected_rewards)
+    np.testing.assert_array_equal(actual_dones, expected_dones)
+    np.testing.assert_array_equal(actual_histories, expected_histories)
+    np.testing.assert_array_equal(actual_actions, expected_actions)
+
+  def test_takes_new_initial_frames(self):
+    initial_observations = np.array([[[0, 1, 2]], [[3, 4, 5]]])
+
+    with backend.use_backend("numpy"):
+      env = self._create_env(  # pylint: disable=no-value-for-parameter
+          model=mock.MagicMock(),
+          initial_observations=initial_observations,
+          trajectory_length=2,
+      )
+      env.reset()
+      observation = env.reset()
+      np.testing.assert_array_equal(observation, [5])
+
+
+if __name__ == "__main__":
+  test.main()

From 8ecb757aa98d24be1d434fe28299d5790dafd771 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Thu, 18 Jul 2019 20:14:37 +0200
Subject: [PATCH 2206/2720] Fix setting agent to eval_mode during evaluation.
 (#1634)

---
 tensor2tensor/rl/dopamine_connector.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index c3933ac8f..48fd92c92 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -826,7 +826,7 @@ def evaluate(self, env_fn, hparams, sampling_temp):
       agent = runner._agent  # pylint: disable=protected-access
       runner.close()
       del runner
-      agent.eval = True
+      agent.eval_mode = True
 
       for _ in range(hparams.eval_episodes_num):
         # Run single episode

From e7c23fca1b794cf40cb86ff21cd95d2cfb07d963 Mon Sep 17 00:00:00 2001
From: David <david.rau@hotmail.de>
Date: Thu, 18 Jul 2019 21:36:15 +0200
Subject: [PATCH 2207/2720] Loss twice multiplied with loss_coef (#1627)

In the case of topk gating, the loss is falsely multiplied twice by the loss coefficient.
---
 tensor2tensor/utils/expert_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 8a97c764c..c5b7f1875 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -1044,7 +1044,7 @@ def local_moe(x,
           noisy_gating=True,
           noise_epsilon=1e-2)
       importance = tf.reduce_sum(gates, 0)
-      loss = loss_coef * (cv_squared(importance) + cv_squared(load))
+      loss = (cv_squared(importance) + cv_squared(load))
     else:
       assert hparams.gating_type == "vq"
       tf.logging.info("Using VQ gating")

From 41e35ef0c4c2c2d7d65b2c96c55ab9f8895dc273 Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Thu, 18 Jul 2019 11:14:59 -0700
Subject: [PATCH 2208/2720] Merge of PR #1634

PiperOrigin-RevId: 258808073
---
 tensor2tensor/utils/expert_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index c5b7f1875..8a97c764c 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -1044,7 +1044,7 @@ def local_moe(x,
           noisy_gating=True,
           noise_epsilon=1e-2)
       importance = tf.reduce_sum(gates, 0)
-      loss = (cv_squared(importance) + cv_squared(load))
+      loss = loss_coef * (cv_squared(importance) + cv_squared(load))
     else:
       assert hparams.gating_type == "vq"
       tf.logging.info("Using VQ gating")

From 2bc21898f6335072ee5fb6e73a968436fa58db18 Mon Sep 17 00:00:00 2001
From: David <david.rau@hotmail.de>
Date: Thu, 18 Jul 2019 12:36:47 -0700
Subject: [PATCH 2209/2720] Merge of PR #1627

PiperOrigin-RevId: 258823663
---
 tensor2tensor/utils/expert_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 8a97c764c..c5b7f1875 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -1044,7 +1044,7 @@ def local_moe(x,
           noisy_gating=True,
           noise_epsilon=1e-2)
       importance = tf.reduce_sum(gates, 0)
-      loss = loss_coef * (cv_squared(importance) + cv_squared(load))
+      loss = (cv_squared(importance) + cv_squared(load))
     else:
       assert hparams.gating_type == "vq"
       tf.logging.info("Using VQ gating")

From 21394f1a9c3a4cb745f3716d2ea83135bba99b62 Mon Sep 17 00:00:00 2001
From: Thuong-Hai Pham <phamthuonghai@users.noreply.github.com>
Date: Fri, 19 Jul 2019 04:07:32 +0800
Subject: [PATCH 2210/2720] Should not generate summary during decoding in
 dot_product_relative_atention (#1618)

---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 7b38d6c87..b8c029cfd 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1745,7 +1745,7 @@ def dot_product_attention_relative(q,
       save_weights_to[scope.name] = weights
       save_weights_to[scope.name + "/logits"] = logits
     weights = tf.nn.dropout(weights, 1.0 - dropout_rate)
-    if not tf.get_variable_scope().reuse and make_image_summary:
+    if not tf.get_variable_scope().reuse and common_layers.should_generate_summaries() and make_image_summary:
       attention_image_summary(weights, image_shapes)
     return _relative_attention_inner(weights, v, relations_values, False)
 

From 7e77fa90079c73d5af91d0316e4f4d8a039b04e6 Mon Sep 17 00:00:00 2001
From: Thuong-Hai Pham <phamthuonghai@users.noreply.github.com>
Date: Thu, 18 Jul 2019 13:08:03 -0700
Subject: [PATCH 2211/2720] Merge of PR #1618

PiperOrigin-RevId: 258829641
---
 tensor2tensor/layers/common_attention.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index b8c029cfd..aba173123 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1745,7 +1745,9 @@ def dot_product_attention_relative(q,
       save_weights_to[scope.name] = weights
       save_weights_to[scope.name + "/logits"] = logits
     weights = tf.nn.dropout(weights, 1.0 - dropout_rate)
-    if not tf.get_variable_scope().reuse and common_layers.should_generate_summaries() and make_image_summary:
+    if (not tf.get_variable_scope().reuse and
+        common_layers.should_generate_summaries() and
+        make_image_summary):
       attention_image_summary(weights, image_shapes)
     return _relative_attention_inner(weights, v, relations_values, False)
 

From 69a81e81a912ba3090490b2c16f7be3ab75404e6 Mon Sep 17 00:00:00 2001
From: Guillermo Peralta Scura <gperaltascura@gmail.com>
Date: Thu, 18 Jul 2019 16:53:27 -0400
Subject: [PATCH 2212/2720] Add English-Spanish translation problem (#1626)

---
 README.md                                     |   1 +
 docs/index.md                                 |   1 +
 .../data_generators/translate_enes.py         | 121 ++++++++++++++++++
 3 files changed, 123 insertions(+)
 create mode 100644 tensor2tensor/data_generators/translate_enes.py

diff --git a/README.md b/README.md
index 27c992491..d51eff0a9 100644
--- a/README.md
+++ b/README.md
@@ -210,6 +210,7 @@ There are a number of translation data-sets in T2T:
 * English-Czech: `--problem=translate_encs_wmt32k`
 * English-Chinese: `--problem=translate_enzh_wmt32k`
 * English-Vietnamese: `--problem=translate_envi_iwslt32k`
+* English-Spanish: `--problem=translate_enes_wmt32k`
 
 You can get translations in the other direction by appending `_rev` to
 the problem name, e.g., for German-English use
diff --git a/docs/index.md b/docs/index.md
index 38c6120eb..26298a9d2 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -111,6 +111,7 @@ There are a number of translation data-sets in T2T:
 * English-Czech: `--problem=translate_encs_wmt32k`
 * English-Chinese: `--problem=translate_enzh_wmt32k`
 * English-Vietnamese: `--problem=translate_envi_iwslt32k`
+* English-Spanish: `--problem=translate_enes_wmt32k`
 
 You can get translations in the other direction by appending `_rev` to
 the problem name, e.g., for German-English use
diff --git a/tensor2tensor/data_generators/translate_enes.py b/tensor2tensor/data_generators/translate_enes.py
new file mode 100644
index 000000000..90b396134
--- /dev/null
+++ b/tensor2tensor/data_generators/translate_enes.py
@@ -0,0 +1,121 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Data generators for translation data-sets."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators import translate
+from tensor2tensor.utils import registry
+
+# End-of-sentence marker.
+EOS = text_encoder.EOS_ID
+
+_ENES_TRAIN_DATASETS = [
+    [
+        "http://www.statmt.org/wmt13/training-parallel-commoncrawl.tgz",
+        ("commoncrawl.es-en.en", "commoncrawl.es-en.es")
+    ],
+    [
+        "http://www.statmt.org/wmt13/training-parallel-europarl-v7.tgz",
+        ("training/europarl-v7.es-en.en", "training/europarl-v7.es-en.es")
+    ],
+    [
+        "http://www.statmt.org/wmt13/training-parallel-un.tgz",
+        ("un/undoc.2000.es-en.en", "un/undoc.2000.es-en.es")
+    ],
+    [
+        "https://s3.amazonaws.com/web-language-models/paracrawl/release1/paracrawl-release1.en-es.zipporah0-dedup-clean.tgz",
+        ("paracrawl-release1.en-es.zipporah0-dedup-clean.en",
+         "paracrawl-release1.en-es.zipporah0-dedup-clean.es")
+    ]
+]
+_ENES_TEST_DATASETS = [
+    [
+        "http://data.statmt.org/wmt17/translation-task/dev.tgz",
+        ("dev/newstest2013.en", "dev/newstest2013.es")
+    ],
+]
+
+@registry.register_problem
+class TranslateEnesWmt32k(translate.TranslateProblem):
+  """En-es translation trained on WMT corpus."""
+
+  @property
+  def additional_training_datasets(self):
+    """Allow subclasses to add training datasets."""
+    return []
+
+  def source_data_files(self, dataset_split):
+    train = dataset_split == problem.DatasetSplit.TRAIN
+    train_datasets = _ENES_TRAIN_DATASETS + self.additional_training_datasets
+    return train_datasets if train else _ENES_TEST_DATASETS
+
+  def vocab_data_files(self):
+    return _ENES_TRAIN_DATASETS
+
+@registry.register_problem
+class TranslateEnesWmtClean32k(TranslateEnesWmt32k):
+  """En-es translation trained on WMT with further cleaning."""
+
+  @property
+  def use_vocab_from_other_problem(self):
+    return TranslateEnesWmt32k()
+
+  @property
+  def datatypes_to_clean(self):
+    return ["txt"]
+
+@registry.register_problem
+class TranslateEnesWmt32kPacked(TranslateEnesWmt32k):
+
+  @property
+  def packed_length(self):
+    return 256
+
+  @property
+  def use_vocab_from_other_problem(self):
+    return TranslateEnesWmt32k()
+
+@registry.register_problem
+class TranslateEnesWmt8k(TranslateEnesWmt32k):
+  """Problem spec for WMT En-Es translation."""
+
+  @property
+  def approx_vocab_size(self):
+    return 2**13  # 8192
+
+@registry.register_problem
+class TranslateEnesWmt8kPacked(TranslateEnesWmt8k):
+
+  @property
+  def packed_length(self):
+    return 256
+
+  @property
+  def use_vocab_from_other_problem(self):
+    return TranslateEnesWmt8k()
+
+@registry.register_problem
+class TranslateEnesWmtCharacters(TranslateEnesWmt8k):
+  """Problem spec for WMT En-Es translation."""
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
\ No newline at end of file

From 77146aa87bcfa6269b6f7c0332cff3d404890f8d Mon Sep 17 00:00:00 2001
From: Guillermo Peralta Scura <gperaltascura@gmail.com>
Date: Thu, 18 Jul 2019 13:53:51 -0700
Subject: [PATCH 2213/2720] Merge of PR #1626

PiperOrigin-RevId: 258838381
---
 docs/walkthrough.md                             | 1 +
 tensor2tensor/data_generators/all_problems.py   | 1 +
 tensor2tensor/data_generators/translate_enes.py | 8 +++++++-
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 27c992491..d51eff0a9 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -210,6 +210,7 @@ There are a number of translation data-sets in T2T:
 * English-Czech: `--problem=translate_encs_wmt32k`
 * English-Chinese: `--problem=translate_enzh_wmt32k`
 * English-Vietnamese: `--problem=translate_envi_iwslt32k`
+* English-Spanish: `--problem=translate_enes_wmt32k`
 
 You can get translations in the other direction by appending `_rev` to
 the problem name, e.g., for German-English use
diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 79687e31b..19ccffc78 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -78,6 +78,7 @@
     "tensor2tensor.data_generators.transduction_problems",
     "tensor2tensor.data_generators.translate_encs",
     "tensor2tensor.data_generators.translate_ende",
+    "tensor2tensor.data_generators.translate_enes",
     "tensor2tensor.data_generators.translate_enet",
     "tensor2tensor.data_generators.translate_enfr",
     "tensor2tensor.data_generators.translate_enid",
diff --git a/tensor2tensor/data_generators/translate_enes.py b/tensor2tensor/data_generators/translate_enes.py
index 90b396134..1619a5adc 100644
--- a/tensor2tensor/data_generators/translate_enes.py
+++ b/tensor2tensor/data_generators/translate_enes.py
@@ -53,6 +53,7 @@
     ],
 ]
 
+
 @registry.register_problem
 class TranslateEnesWmt32k(translate.TranslateProblem):
   """En-es translation trained on WMT corpus."""
@@ -70,6 +71,7 @@ def source_data_files(self, dataset_split):
   def vocab_data_files(self):
     return _ENES_TRAIN_DATASETS
 
+
 @registry.register_problem
 class TranslateEnesWmtClean32k(TranslateEnesWmt32k):
   """En-es translation trained on WMT with further cleaning."""
@@ -82,6 +84,7 @@ def use_vocab_from_other_problem(self):
   def datatypes_to_clean(self):
     return ["txt"]
 
+
 @registry.register_problem
 class TranslateEnesWmt32kPacked(TranslateEnesWmt32k):
 
@@ -93,6 +96,7 @@ def packed_length(self):
   def use_vocab_from_other_problem(self):
     return TranslateEnesWmt32k()
 
+
 @registry.register_problem
 class TranslateEnesWmt8k(TranslateEnesWmt32k):
   """Problem spec for WMT En-Es translation."""
@@ -101,6 +105,7 @@ class TranslateEnesWmt8k(TranslateEnesWmt32k):
   def approx_vocab_size(self):
     return 2**13  # 8192
 
+
 @registry.register_problem
 class TranslateEnesWmt8kPacked(TranslateEnesWmt8k):
 
@@ -112,10 +117,11 @@ def packed_length(self):
   def use_vocab_from_other_problem(self):
     return TranslateEnesWmt8k()
 
+
 @registry.register_problem
 class TranslateEnesWmtCharacters(TranslateEnesWmt8k):
   """Problem spec for WMT En-Es translation."""
 
   @property
   def vocab_type(self):
-    return text_problems.VocabType.CHARACTER
\ No newline at end of file
+    return text_problems.VocabType.CHARACTER

From 19a96660424becf75b2265d4f01c868f2f6cc254 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 19 Jul 2019 15:01:09 -0700
Subject: [PATCH 2214/2720] Change to shorter names that match the source
 paper: d_ff and d_model.

PiperOrigin-RevId: 259045383
---
 .../position_lookup_transformer_copy.gin      |  4 +-
 .../trax/configs/transformer_big_lm1b_8gb.gin |  4 +-
 .../trax/configs/transformer_imdb_8gb.gin     |  4 +-
 .../trax/configs/transformer_lm1b_8gb.gin     |  4 +-
 .../configs/transformer_lm1b_8gb_testing.gin  |  4 +-
 .../transformer_revnet_imagenet64_8gb.gin     |  4 +-
 .../configs/transformer_revnet_lm1b_8gb.gin   |  4 +-
 .../trax/configs/transformer_wmt_ende_8gb.gin |  4 +-
 .../research/position_lookup_transformer.py   | 44 +++++------
 .../models/research/transformer_revnet.py     | 38 ++++-----
 tensor2tensor/trax/models/transformer.py      | 78 +++++++++----------
 tensor2tensor/trax/models/transformer_test.py |  4 +-
 12 files changed, 98 insertions(+), 98 deletions(-)

diff --git a/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin b/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
index d9ea5abd6..f6aef3874 100644
--- a/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
+++ b/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
@@ -28,8 +28,8 @@ MultifactorSchedule.warmup_steps = 16000
 
 # Parameters for PositionLookupTransformerLM:
 # ==============================================================================
-PositionLookupTransformerLM.d_feature = 256
-PositionLookupTransformerLM.d_feedforward = 512
+PositionLookupTransformerLM.d_model = 256
+PositionLookupTransformerLM.d_ff = 512
 PositionLookupTransformerLM.dropout = 0.01
 PositionLookupTransformerLM.max_len = 100
 PositionLookupTransformerLM.n_heads = 4
diff --git a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
index 4b3d0fc72..63b693f7b 100644
--- a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
@@ -42,8 +42,8 @@ train.train_steps = 500000
 
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.d_feature = 1024
-TransformerLM.d_feedforward = 8192
+TransformerLM.d_model = 1024
+TransformerLM.d_ff = 8192
 TransformerLM.dropout = 0.1
 TransformerLM.max_len = 2048
 TransformerLM.mode = 'train'
diff --git a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
index f25fcc2e4..58391c63d 100644
--- a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
@@ -38,8 +38,8 @@ train.train_steps = 1000
 
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerEncoder.d_feature = 512
-TransformerEncoder.d_feedforward = 2048
+TransformerEncoder.d_model = 512
+TransformerEncoder.d_ff = 2048
 TransformerEncoder.dropout = 0.1
 TransformerEncoder.max_len = 2048
 TransformerEncoder.mode = 'train'
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index 9c0bfdd87..5549858d7 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -42,8 +42,8 @@ train.train_steps = 500000
 
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.d_feature = 512
-TransformerLM.d_feedforward = 2048
+TransformerLM.d_model = 512
+TransformerLM.d_ff = 2048
 TransformerLM.dropout = 0.1
 TransformerLM.max_len = 2048
 TransformerLM.mode = 'train'
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
index e32ce968c..56f779b52 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
@@ -42,8 +42,8 @@ train.train_steps = 100000
 
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.d_feature = 512
-TransformerLM.d_feedforward = 2048
+TransformerLM.d_model = 512
+TransformerLM.d_ff = 2048
 TransformerLM.dropout = 0.1
 TransformerLM.max_len = 2048
 TransformerLM.mode = 'train'
diff --git a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
index f8cb6b4d6..7716aed65 100644
--- a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
@@ -34,8 +34,8 @@ train.trainer_class = @MemoryEfficientTrainer
 
 # Parameters for TransformerRevnetLM:
 # ==============================================================================
-TransformerRevnetLM.d_feature = 512
-TransformerRevnetLM.d_feedforward = 2048
+TransformerRevnetLM.d_model = 512
+TransformerRevnetLM.d_ff = 2048
 TransformerRevnetLM.d_attention_key = 32
 TransformerRevnetLM.d_attention_value = 32
 TransformerRevnetLM.dropout = 0.1
diff --git a/tensor2tensor/trax/configs/transformer_revnet_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_lm1b_8gb.gin
index f1b7addec..3f4cd14aa 100644
--- a/tensor2tensor/trax/configs/transformer_revnet_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_revnet_lm1b_8gb.gin
@@ -44,8 +44,8 @@ train.trainer_class = @MemoryEfficientTrainer
 
 # Parameters for TransformerRevnetLM:
 # ==============================================================================
-TransformerRevnetLM.d_feature = 512
-TransformerRevnetLM.d_feedforward = 2048
+TransformerRevnetLM.d_model = 512
+TransformerRevnetLM.d_ff = 2048
 TransformerRevnetLM.dropout = 0.1
 TransformerRevnetLM.max_len = 2048
 TransformerRevnetLM.mode = 'train'
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
index 656144f57..01972aeec 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
@@ -50,8 +50,8 @@ Adam.eps = 1e-9
 
 # Parameters for Transformer:
 # ==============================================================================
-Transformer.d_feature= 512
-Transformer.d_feedforward = 2048
+Transformer.d_model= 512
+Transformer.d_ff = 2048
 Transformer.dropout = 0.1
 Transformer.max_len = 2048
 Transformer.mode = 'train'
diff --git a/tensor2tensor/trax/models/research/position_lookup_transformer.py b/tensor2tensor/trax/models/research/position_lookup_transformer.py
index 3b84fa265..969812745 100644
--- a/tensor2tensor/trax/models/research/position_lookup_transformer.py
+++ b/tensor2tensor/trax/models/research/position_lookup_transformer.py
@@ -229,51 +229,51 @@ def SumLearnedPick(positions):
   )
 
 
-def AttentionPosition(positions, d_feature, n_heads=8, dropout=0.0,
+def AttentionPosition(positions, d_model, n_heads=8, dropout=0.0,
                       mode='train'):
   """Transformer-style multi-headed attention."""
   return tl.Serial(
       tl.Dup(),
       tl.Dup(),
       tl.Parallel(
-          ApplyAndQueryPositions(tl.Dense(d_feature),
+          ApplyAndQueryPositions(tl.Dense(d_model),
                                  pos=[SumLearnedPick(positions)
                                       for _ in range(n_heads)]),
-          PreservePosition(tl.Dense(d_feature)),
-          PreservePosition(tl.Dense(d_feature)),
+          PreservePosition(tl.Dense(d_model)),
+          PreservePosition(tl.Dense(d_model)),
       ),
       tl.Parallel(
           CopyHeadsPos(h=n_heads),
           MixHeadsPos(h=n_heads),
           MixHeadsPos(h=n_heads),
       ),
-      tl.PureAttention(d_feature=d_feature, n_heads=n_heads, dropout=dropout,
+      tl.PureAttention(d_model=d_model, n_heads=n_heads, dropout=dropout,
                        mode=mode),
       tl.Parallel([], tl.Drop()),  # Drop the mask.
       CombineHeadsPos(h=n_heads),
-      PreservePosition(tl.Dense(d_feature)),
+      PreservePosition(tl.Dense(d_model)),
   )
 
 
-def ResidualFeedForward(d_feature,
-                        d_feedforward,
+def ResidualFeedForward(d_model,
+                        d_ff,
                         dropout,
                         mode):
   """Residual feed-forward layer with normalization at start."""
   stack = tl.Serial(
       tl.LayerNorm(),
-      tl.Dense(d_feedforward),
+      tl.Dense(d_ff),
       tl.Relu(),
       tl.Dropout(rate=dropout, mode=mode),
-      tl.Dense(d_feature),
+      tl.Dense(d_model),
       tl.Dropout(rate=dropout, mode=mode)
   )
   return tl.Residual(PreservePosition(stack))
 
 
 def DecoderLayer(positions,
-                 d_feature,
-                 d_feedforward,
+                 d_model,
+                 d_ff,
                  n_heads,
                  dropout,
                  mode):
@@ -281,8 +281,8 @@ def DecoderLayer(positions,
 
   Args:
     positions: random vectors for positions
-    d_feature: int:  depth of embedding
-    d_feedforward: int: depth of feed-forward layer
+    d_model: int:  depth of embedding
+    d_ff: int: depth of feed-forward layer
     n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
@@ -296,17 +296,17 @@ def DecoderLayer(positions,
           tl.Dup(),
           tl.Parallel([],  # activation for (q, k, v)
                       tl.CausalMask(axis=-2)),  # attention mask
-          AttentionPosition(positions, d_feature, n_heads=n_heads,
+          AttentionPosition(positions, d_model, n_heads=n_heads,
                             dropout=dropout, mode=mode),
           PreservePosition(tl.Dropout(rate=dropout, mode=mode))
       ),
-      ResidualFeedForward(d_feature, d_feedforward, dropout, mode=mode)
+      ResidualFeedForward(d_model, d_ff, dropout, mode=mode)
   ]
 
 
 def PositionLookupTransformerLM(vocab_size=128,
-                                d_feature=256,
-                                d_feedforward=512,
+                                d_model=256,
+                                d_ff=512,
                                 n_layers=3,
                                 n_heads=4,
                                 dropout=0.1,
@@ -316,8 +316,8 @@ def PositionLookupTransformerLM(vocab_size=128,
 
   Args:
     vocab_size: int: vocab size
-    d_feature: int:  depth of embedding
-    d_feedforward: int: depth of feed-forward layer
+    d_model: int:  depth of embedding
+    d_ff: int: depth of feed-forward layer
     n_layers: int: number of layers
     n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
@@ -330,10 +330,10 @@ def PositionLookupTransformerLM(vocab_size=128,
   positions = _POSITIONS[:max_len, :]
   return tl.Serial(
       tl.ShiftRight(),
-      tl.Embedding(d_feature, vocab_size),
+      tl.Embedding(d_model, vocab_size),
       tl.Dropout(rate=dropout, mode=mode),
       NewPositionalEncoding(positions=positions),
-      [DecoderLayer(positions, d_feature, d_feedforward, n_heads, dropout, mode)
+      [DecoderLayer(positions, d_model, d_ff, n_heads, dropout, mode)
        for _ in range(n_layers)],
       PreservePosition(tl.LayerNorm()),
       tl.Dense(vocab_size),
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 800d18f9b..7a5fe0a29 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -78,7 +78,7 @@ def new_parameters(self, input_shape, input_dtype, rng):
     return self._layer.initialize(first_shape, input_dtype[0], rng)
 
 
-def FeedForward(d_feature, d_feedforward, dropout, mode):
+def FeedForward(d_model, d_ff, dropout, mode):
   """Feed-forward block with layer normalization at start."""
   # TODO(kitaev): add dropout. Dropout is typically performed by adding noise to
   # the activations, but when the size of the activations is very large it is
@@ -86,9 +86,9 @@ def FeedForward(d_feature, d_feedforward, dropout, mode):
   del dropout, mode
   return [
       tl.LayerNorm(),
-      tl.Dense(d_feedforward),
+      tl.Dense(d_ff),
       tl.Relu(),
-      tl.Dense(d_feature),
+      tl.Dense(d_model),
   ]
 
 
@@ -252,11 +252,11 @@ def call_compute_residual(x, params, kwargs):
 @tl.layer(n_inputs=1, n_outputs=1)
 def SplitHeads(x, params, n_heads=1, **kwargs):
   del params, kwargs
-  d_feature = x.shape[-1]
-  assert d_feature % n_heads == 0
-  d_head = d_feature // n_heads
+  d_model = x.shape[-1]
+  assert d_model % n_heads == 0
+  d_head = d_model // n_heads
   n_batch = np.shape(x)[0]
-  # n_batch, seqlen, d_feature --> n_batch, n_heads, seqlen, d_head
+  # n_batch, seqlen, d_model --> n_batch, n_heads, seqlen, d_head
   return np.transpose(
       np.reshape(x, (n_batch, -1, n_heads, d_head)), (0, 2, 1, 3))
 
@@ -266,7 +266,7 @@ def JoinHeads(x, params, **kwargs):
   del params, kwargs
   n_batch = np.shape(x)[0]
   seqlen = np.shape(x)[2]
-  # n_batch, n_heads, seqlen, d_head --> n_batch, seqlen, d_feature
+  # n_batch, n_heads, seqlen, d_head --> n_batch, seqlen, d_model
   return np.reshape(np.transpose(x, (0, 2, 1, 3)), (n_batch, seqlen, -1))
 
 
@@ -631,14 +631,14 @@ def inverse_and_vjp(self, output, ct, params=(), **kwargs):
       return layer_val, None
 
 
-def DecoderBlock(d_feature, d_feedforward, d_attention_key, d_attention_value,
+def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
                  n_heads, n_attention_chunks, attention_loop_stride,
                  dropout, mode):
   """Reversible transformer decoder layer.
 
   Args:
-    d_feature: int:  depth of embedding
-    d_feedforward: int: depth of feed-forward layer
+    d_model: int:  depth of embedding
+    d_ff: int: depth of feed-forward layer
     d_attention_key: int: depth of key vector for each attention head
     d_attention_value: int: depth of value vector for each attention head
     n_heads: int: number of attention heads
@@ -675,12 +675,12 @@ def DecoderBlock(d_feature, d_feedforward, d_attention_key, d_attention_value,
   # its input (so the backward pass can be computed without knowing the input)
   post_attention = [
       JoinHeads(),  # pylint: disable=no-value-for-parameter
-      tl.Dense(d_feature),
+      tl.Dense(d_model),
       Unchunk(sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
   ]
 
   feed_forward = [
-      FeedForward(d_feature, d_feedforward, dropout, mode=mode),
+      FeedForward(d_model, d_ff, dropout, mode=mode),
   ]
   return [
       ReversibleAttentionHalfResidual(pre_attention, attention, post_attention),
@@ -691,8 +691,8 @@ def DecoderBlock(d_feature, d_feedforward, d_attention_key, d_attention_value,
 
 
 def TransformerRevnetLM(vocab_size,
-                        d_feature=512,
-                        d_feedforward=2048,
+                        d_model=512,
+                        d_ff=2048,
                         d_attention_key=64,
                         d_attention_value=64,
                         n_layers=6,
@@ -707,8 +707,8 @@ def TransformerRevnetLM(vocab_size,
 
   Args:
     vocab_size: int: vocab size
-    d_feature: int:  depth of *each half* of the two-part features
-    d_feedforward: int: depth of feed-forward layer
+    d_model: int:  depth of *each half* of the two-part features
+    d_ff: int: depth of feed-forward layer
     d_attention_key: int: depth of key vector for each attention head
     d_attention_value: int: depth of value vector for each attention head
     n_layers: int: number of decoder layers
@@ -725,7 +725,7 @@ def TransformerRevnetLM(vocab_size,
     the layer.
   """
   positional_embedder = [
-      tl.Embedding(d_feature, vocab_size),
+      tl.Embedding(d_model, vocab_size),
       # TODO(kitaev): add dropout
       tl.PositionalEncoding(max_len=max_len),
   ]
@@ -736,7 +736,7 @@ def TransformerRevnetLM(vocab_size,
       tl.Dup(),
       ReversibleSerial([
           # pylint: disable=g-complex-comprehension
-          DecoderBlock(d_feature, d_feedforward,
+          DecoderBlock(d_model, d_ff,
                        d_attention_key, d_attention_value, n_heads,
                        n_attention_chunks, attention_loop_stride,
                        dropout, mode)
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 3279ea002..0839a6c1f 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -21,19 +21,19 @@
 from tensor2tensor.trax import layers as tl
 
 
-def FeedForward(d_feature, d_feedforward, dropout, mode):
+def FeedForward(d_model, d_ff, dropout, mode):
   """Feed-forward block with layer normalization at start."""
   return [
       tl.LayerNorm(),
-      tl.Dense(d_feedforward),
+      tl.Dense(d_ff),
       tl.Relu(),
       tl.Dropout(rate=dropout, mode=mode),
-      tl.Dense(d_feature),
+      tl.Dense(d_model),
       tl.Dropout(rate=dropout, mode=mode),
   ]
 
 
-def EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
+def EncoderBlock(d_model, d_ff, n_heads, dropout, mode):
   """Returns a layer sequence that implements a Transformer encoder block.
 
   The input to the layer sequence is a pair, (activations, mask), where the
@@ -41,8 +41,8 @@ def EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
   padding part of the input.
 
   Args:
-    d_feature: int:  depth of embedding
-    d_feedforward: int: depth of feed-forward layer
+    d_model: int:  depth of embedding
+    d_ff: int: depth of feed-forward layer
     n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
@@ -53,11 +53,11 @@ def EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
   """
   attention = [
       tl.LayerNorm(),
-      tl.Attention(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+      tl.Attention(d_model, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Dropout(rate=dropout, mode=mode),
   ]
   feed_forward = [
-      FeedForward(d_feature, d_feedforward, dropout, mode=mode),
+      FeedForward(d_model, d_ff, dropout, mode=mode),
   ]
   return [
       tl.Residual(attention),
@@ -67,8 +67,8 @@ def EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
 
 def TransformerEncoder(vocab_size,
                        n_classes=10,
-                       d_feature=512,
-                       d_feedforward=2048,
+                       d_model=512,
+                       d_ff=2048,
                        n_layers=6,
                        n_heads=8,
                        dropout=0.1,
@@ -81,8 +81,8 @@ def TransformerEncoder(vocab_size,
   Args:
     vocab_size: int: vocab size
     n_classes: how many classes on output
-    d_feature: int:  depth of embedding
-    d_feedforward: int: depth of feed-forward layer
+    d_model: int:  depth of embedding
+    d_ff: int: depth of feed-forward layer
     n_layers: int: number of encoder/decoder layers
     n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
@@ -94,14 +94,14 @@ def TransformerEncoder(vocab_size,
     activations over a set of output classes.
   """
   embedder = [
-      tl.Embedding(d_feature, vocab_size),
+      tl.Embedding(d_model, vocab_size),
       tl.Dropout(rate=dropout, mode=mode),
       tl.PositionalEncoding(max_len=max_len),
   ]
   return tl.Model([                             #      tokens
       tl.Dup(),                                 # toks toks
       tl.Parallel(embedder, tl.PaddingMask()),  # vecs mask
-      [EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode)
+      [EncoderBlock(d_model, d_ff, n_heads, dropout, mode)
        for _ in range(n_layers)],               # vecs mask
       tl.Parallel([], tl.Drop()),               # ____  0
       tl.LayerNorm(),                           # vecs
@@ -111,14 +111,14 @@ def TransformerEncoder(vocab_size,
   ])
 
 
-def DecoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
+def DecoderBlock(d_model, d_ff, n_heads, dropout, mode):
   """Returns a layer sequence that implements a Transformer decoder block.
 
   The input to the layer sequence is an activation tensor.
 
   Args:
-    d_feature: int:  depth of embedding
-    d_feedforward: int: depth of feed-forward layer
+    d_model: int:  depth of embedding
+    d_ff: int: depth of feed-forward layer
     n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
@@ -130,12 +130,12 @@ def DecoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
       tl.LayerNorm(),  # vec
       tl.Dup(),  # vec vec
       tl.Parallel([], tl.CausalMask(axis=-2)),  # vec mask
-      tl.Attention(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+      tl.Attention(d_model, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Parallel([], tl.Drop()),  # vec
       tl.Dropout(rate=dropout, mode=mode),  # vec
   ]
   feed_forward = [
-      FeedForward(d_feature, d_feedforward, dropout, mode=mode),
+      FeedForward(d_model, d_ff, dropout, mode=mode),
   ]
   return [
       tl.Residual(self_attention),
@@ -144,8 +144,8 @@ def DecoderBlock(d_feature, d_feedforward, n_heads, dropout, mode):
 
 
 def TransformerLM(vocab_size,
-                  d_feature=512,
-                  d_feedforward=2048,
+                  d_model=512,
+                  d_ff=2048,
                   n_layers=6,
                   n_heads=8,
                   dropout=0.1,
@@ -158,8 +158,8 @@ def TransformerLM(vocab_size,
 
   Args:
     vocab_size: int: vocab size
-    d_feature: int:  depth of embedding
-    d_feedforward: int: depth of feed-forward layer
+    d_model: int:  depth of embedding
+    d_ff: int: depth of feed-forward layer
     n_layers: int: number of encoder/decoder layers
     n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
@@ -171,14 +171,14 @@ def TransformerLM(vocab_size,
     to activations over a vocab set.
   """
   embedder = [
-      tl.Embedding(d_feature, vocab_size),
+      tl.Embedding(d_model, vocab_size),
       tl.Dropout(rate=dropout, mode=mode),
       tl.PositionalEncoding(max_len=max_len),
   ]
   return tl.Model(                  # tokens
       tl.ShiftRight(),              # toks
       embedder,                     # vecs
-      [DecoderBlock(d_feature, d_feedforward, n_heads, dropout, mode)
+      [DecoderBlock(d_model, d_ff, n_heads, dropout, mode)
        for _ in range(n_layers)],   # vecs
       tl.LayerNorm(),               # vecs
       tl.Dense(vocab_size),         # vecs
@@ -186,7 +186,7 @@ def TransformerLM(vocab_size,
   )
 
 
-def EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode):
+def EncoderDecoder(d_model, d_ff, n_heads, dropout, mode):
   """Transformer encoder-decoder layer.
 
   The input is a triple (decoder_input, mask, encoder) where the mask is
@@ -194,8 +194,8 @@ def EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode):
   of the encoder.
 
   Args:
-    d_feature: int:  depth of embedding
-    d_feedforward: int: depth of feed-forward layer
+    d_model: int:  depth of embedding
+    d_ff: int: depth of feed-forward layer
     n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
@@ -207,7 +207,7 @@ def EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode):
       tl.LayerNorm(),                           #        vecs_d   ..... ......
       tl.Dup(),                                 # vecs_d vecs_d   ..... ......
       tl.Parallel([], tl.CausalMask(axis=-2)),  # ______ masks    ..... ......
-      tl.Attention(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+      tl.Attention(d_model, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Parallel([], tl.Drop()),               # ______   0      ..... ......
       tl.Dropout(rate=dropout, mode=mode),      # vecs_d          ..... ......
   ]
@@ -217,11 +217,11 @@ def EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode):
       tl.Parallel([], tl.Swap()),         # ______        vecs_e masks  ......
       tl.Parallel([], tl.Dup()),          # ______ vecs_e vecs_e .....  ......
       tl.AttentionQKV(  # (q k v masks ... --> vecs_d masks ...)
-          d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+          d_model, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Dropout(rate=dropout, mode=mode),  # vecs_d mask vecs_e
   ]
   feed_forward = [
-      FeedForward(d_feature, d_feedforward, dropout, mode=mode),
+      FeedForward(d_model, d_ff, dropout, mode=mode),
   ]
   return [                                        # vecs_d masks vecs_e
       tl.Residual(decoder_self_attention),        # vecs_d masks vecs_e
@@ -232,8 +232,8 @@ def EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode):
 
 def Transformer(input_vocab_size,
                 output_vocab_size=None,
-                d_feature=512,
-                d_feedforward=2048,
+                d_model=512,
+                d_ff=2048,
                 n_layers=6,
                 n_heads=8,
                 dropout=0.1,
@@ -247,8 +247,8 @@ def Transformer(input_vocab_size,
     input_vocab_size: int: vocab size of the source.
     output_vocab_size: int (optional): vocab size of the target. If None, the
       source and target are assumed to have the same vocab.
-    d_feature: int:  depth of embedding
-    d_feedforward: int: depth of feed-forward layer
+    d_model: int:  depth of embedding
+    d_ff: int: depth of feed-forward layer
     n_layers: int: number of encoder/decoder layers
     n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
@@ -260,7 +260,7 @@ def Transformer(input_vocab_size,
     activations over a vocab set.
   """
   in_embed = [                                    # tokens
-      tl.Embedding(d_feature, input_vocab_size),  # vecs
+      tl.Embedding(d_model, input_vocab_size),  # vecs
       tl.Dropout(rate=dropout, mode=mode),        # vecs
       tl.PositionalEncoding(max_len=max_len),     # vecs
   ]
@@ -270,17 +270,17 @@ def Transformer(input_vocab_size,
     out_embed = in_embed
   else:
     out_embed = [                                    # tokens
-        tl.Embedding(d_feature, output_vocab_size),  # vecs
+        tl.Embedding(d_model, output_vocab_size),  # vecs
         tl.Dropout(rate=dropout, mode=mode),         # vecs
         tl.PositionalEncoding(max_len=max_len),      # vecs
     ]
 
   encoder_stack = (  # masks vectors --> masks vectors
-      [EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode)
+      [EncoderBlock(d_model, d_ff, n_heads, dropout, mode)
        for _ in range(n_layers)])
 
   encoder_decoder_stack = (  # vecs_d masks vecs_e --> vecs_d masks vecs_e
-      [EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode)
+      [EncoderDecoder(d_model, d_ff, n_heads, dropout, mode)
        for _ in range(n_layers)])
 
   # Input: encoder_side_tokens, decoder_side_tokens
diff --git a/tensor2tensor/trax/models/transformer_test.py b/tensor2tensor/trax/models/transformer_test.py
index 1ec6a3832..8b3a89d9c 100644
--- a/tensor2tensor/trax/models/transformer_test.py
+++ b/tensor2tensor/trax/models/transformer_test.py
@@ -33,7 +33,7 @@ def test_transformer_lm_forward_shape(self):
     vocab_size = 16
     input_shape = [3, 5]
     model = transformer.TransformerLM(
-        vocab_size, d_feature=32, d_feedforward=64, n_layers=2, n_heads=2)
+        vocab_size, d_model=32, d_ff=64, n_layers=2, n_heads=2)
     final_shape = tl.check_shape_agreement(
         model, tuple(input_shape), integer_inputs=True)
     self.assertEqual(tuple(input_shape + [vocab_size]), final_shape)
@@ -45,7 +45,7 @@ def _test_transformer_forward_shape(self, input_vocab_size,
     input_shape = (tuple(single_input_shape), tuple(single_input_shape))
     model = transformer.Transformer(
         input_vocab_size, output_vocab_size,
-        d_feature=32, d_feedforward=64, n_layers=2, n_heads=2)
+        d_model=32, d_ff=64, n_layers=2, n_heads=2)
     final_shape = tl.check_shape_agreement(
         model, input_shape, integer_inputs=True)
     expected_shape = (tuple(single_input_shape +

From 18e85d6f1e7e9e138c9e2613cfccb48aea301d7d Mon Sep 17 00:00:00 2001
From: Jacob Burnim <jburnim@google.com>
Date: Mon, 22 Jul 2019 14:05:59 -0700
Subject: [PATCH 2215/2720] fix monkey-patched OneHotCategorical._log_prob.

PiperOrigin-RevId: 259402217
---
 tensor2tensor/layers/reversible_layers_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
index 7db21c65c..d523e81b2 100644
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ b/tensor2tensor/layers/reversible_layers_test.py
@@ -32,8 +32,8 @@
 
 def _log_prob(self, x):
   """Re-implementation of OneHotCategorical._log_prob for gradients wrt x."""
-  x = self._assert_valid_sample(x)
   logits = self.logits_parameter()
+  x = self._maybe_assert_valid_sample(x, dtype=logits.dtype)
   if (not x.shape.is_fully_defined() or
       not logits.shape.is_fully_defined() or
       x.shape != logits.shape):

From acc1886914e0da18d03d597c1b3f18598e21beec Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 22 Jul 2019 16:32:53 -0700
Subject: [PATCH 2216/2720] Remove parallelism from EnvProblem __init__, now it
 is just used in GymEnvProblem.

PiperOrigin-RevId: 259431034
---
 tensor2tensor/envs/env_problem.py | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index e65bf6405..5e260aa5b 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -99,8 +99,7 @@ class EnvProblem(Env, problem.Problem):
   def __init__(self,
                batch_size=None,
                discrete_rewards=True,
-               parallelism=1,
-               **env_kwargs):
+               **kwargs):
     """Initializes this class by creating the envs and managing trajectories.
 
     Args:
@@ -108,9 +107,7 @@ def __init__(self,
         batched mode.
       discrete_rewards: (bool) whether to round the rewards to the nearest
         integer.
-      parallelism: (int) If this is greater than one then we run the envs in
-        parallel using multi-threading.
-      **env_kwargs: (dict) Additional kwargs to pass to the environments.
+      **kwargs: (dict) Additional kwargs to pass to `self.initialize`.
     """
 
     # Call the super's ctor.
@@ -126,8 +123,6 @@ def __init__(self,
     # If set, we discretize the rewards and treat them as integers.
     self._discrete_rewards = discrete_rewards
 
-    self._parallelism = None
-
     # A data structure to hold the `batch_size` currently active trajectories
     # and also the ones that are completed, i.e. done.
     self._trajectories = None
@@ -135,7 +130,7 @@ def __init__(self,
     self._batch_size = None
 
     if batch_size is not None:
-      self.initialize(batch_size=batch_size, **env_kwargs)
+      self.initialize(batch_size=batch_size, **kwargs)
 
   @property
   def batch_size(self):

From 2a0d6a303d1e6d46874bc12308b78aea74d062b9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 22 Jul 2019 16:40:58 -0700
Subject: [PATCH 2217/2720] Introduce env specific arguments in GymEnvProblem.

Sometimes we need to initialize the `batch_size` envs in GymEnvProblem with
different parameters.

Ex: Different output_directories, different servers to talk to etc.

This CL introduces support for that.

PiperOrigin-RevId: 259432521
---
 tensor2tensor/envs/gym_env_problem.py      | 25 +++++++++-
 tensor2tensor/envs/gym_env_problem_test.py | 53 ++++++++++++++++++++++
 2 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/envs/gym_env_problem.py b/tensor2tensor/envs/gym_env_problem.py
index fa628f7ad..c9bcdca53 100644
--- a/tensor2tensor/envs/gym_env_problem.py
+++ b/tensor2tensor/envs/gym_env_problem.py
@@ -23,6 +23,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import copy
 import multiprocessing.pool
 import time
 
@@ -161,19 +162,39 @@ def _verify_same_spaces(self):
         tf.logging.error("Env[%d] has action space [%s]", i, env.action_space)
       raise ValueError(err_str)
 
-  def initialize_environments(self, batch_size=1, parallelism=1, **kwargs):
+  def initialize_environments(self,
+                              batch_size=1,
+                              parallelism=1,
+                              per_env_kwargs=None,
+                              **kwargs):
     """Initializes the environments.
 
     Args:
       batch_size: (int) Number of `self.base_env_name` envs to initialize.
       parallelism: (int) If this is greater than one then we run the envs in
         parallel using multi-threading.
+      per_env_kwargs: (list or None) An optional list of dictionaries to pass to
+        gym.make. If not None, length should match `batch_size`.
       **kwargs: (dict) Kwargs to pass to gym.make.
     """
     assert batch_size >= 1
+    if per_env_kwargs is not None:
+      assert batch_size == len(per_env_kwargs)
+    else:
+      per_env_kwargs = [{} for _ in range(batch_size)]
+
+    # By now `per_env_kwargs` is a list of dictionaries of size batch_size.
+    # The individual dictionaries maybe empty.
+
+    def union_dicts(dict1, dict2):
+      """Union `dict1` and `dict2`."""
+      copy_dict1 = copy.copy(dict1)
+      copy_dict1.update(dict2)
+      return copy_dict1
 
     self._envs = [
-        gym.make(self.base_env_name, **kwargs) for _ in range(batch_size)
+        gym.make(self.base_env_name, **union_dicts(kwargs, env_kwarg))
+        for env_kwarg in per_env_kwargs
     ]
     self._parallelism = parallelism
     self._pool = multiprocessing.pool.ThreadPool(self._parallelism)
diff --git a/tensor2tensor/envs/gym_env_problem_test.py b/tensor2tensor/envs/gym_env_problem_test.py
index 7df54d938..fe88d9b04 100644
--- a/tensor2tensor/envs/gym_env_problem_test.py
+++ b/tensor2tensor/envs/gym_env_problem_test.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import os
+import gym
 from gym.spaces import Box
 from gym.spaces import Discrete
 import numpy as np
@@ -380,5 +381,57 @@ def test_resets_properly(self):
     # Assert that there aren't any completed trajectories in the env now.
     self.assertEqual(env.trajectories.num_completed_trajectories, 0)
 
+  def test_per_env_kwargs(self):
+
+    # Creating a dummy class where we specify the action at which the env
+    # returns done.
+    class TestPerEnvKwargsEnv(gym.Env):
+      """Test environment with the `done action` specified."""
+
+      action_space = Discrete(3)
+      observation_space = Box(low=-1.0, high=1.0, shape=())
+
+      def __init__(self, done_action=0):
+        self._done_action = done_action
+
+      def _generate_ob(self):
+        return self.observation_space.sample()
+
+      def step(self, action):
+        done = self._done_action == action
+        reward = 1 if done else 0
+        return (self._generate_ob(), reward, done, {})
+
+      def reset(self):
+        return self._generate_ob()
+
+    # Registering it with gym.
+    test_env_name = "TestPerEnvKwargsEnv-v0"
+    gym.envs.register(id=test_env_name, entry_point=TestPerEnvKwargsEnv)
+
+    # Creating a batch of those with different done actions.
+    base_env_name = test_env_name
+    batch_size = 2
+    reward_range = (-1, 1)
+    per_env_kwargs = [{"done_action": 1}, {"done_action": 2}]
+
+    env = gym_env_problem.GymEnvProblem(
+        base_env_name=base_env_name,
+        batch_size=batch_size,
+        reward_range=reward_range,
+        per_env_kwargs=per_env_kwargs)
+
+    _ = env.reset()
+
+    # Finally querying the done actions.
+
+    _, _, d, _ = env.step(np.array([0, 0]))
+    self.assertFalse(d[0])
+    self.assertFalse(d[1])
+
+    _, _, d, _ = env.step(np.array([1, 2]))
+    self.assertTrue(d[0])
+    self.assertTrue(d[1])
+
 if __name__ == "__main__":
   tf.test.main()

From 83b38f1391a43ed415c085cefb1feb547b9a3a5c Mon Sep 17 00:00:00 2001
From: Blake Hechtman <blakehechtman@google.com>
Date: Mon, 22 Jul 2019 22:31:54 -0700
Subject: [PATCH 2218/2720] [TENSOR2TENSOR] Use an einsum instead of matmul.

PiperOrigin-RevId: 259471932
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index aba173123..4d35a1498 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1570,7 +1570,7 @@ def dot_product_attention(q,
   """
   with tf.variable_scope(
       name, default_name="dot_product_attention", values=[q, k, v]) as scope:
-    logits = tf.matmul(q, k, transpose_b=True)  # [..., length_q, length_kv]
+    logits = tf.einsum("...kd,...qd->...qk", k, q)
     if bias is not None:
       bias = common_layers.cast_like(bias, logits)
       logits += bias

From 3c72d2a93910a1fc64eba89432fad09c2216ae2d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 23 Jul 2019 09:36:43 -0700
Subject: [PATCH 2219/2720] Reverting the removing parallelism change.

Reverting now, will figure out later.

PiperOrigin-RevId: 259552773
---
 tensor2tensor/envs/env_problem.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 5e260aa5b..e65bf6405 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -99,7 +99,8 @@ class EnvProblem(Env, problem.Problem):
   def __init__(self,
                batch_size=None,
                discrete_rewards=True,
-               **kwargs):
+               parallelism=1,
+               **env_kwargs):
     """Initializes this class by creating the envs and managing trajectories.
 
     Args:
@@ -107,7 +108,9 @@ def __init__(self,
         batched mode.
       discrete_rewards: (bool) whether to round the rewards to the nearest
         integer.
-      **kwargs: (dict) Additional kwargs to pass to `self.initialize`.
+      parallelism: (int) If this is greater than one then we run the envs in
+        parallel using multi-threading.
+      **env_kwargs: (dict) Additional kwargs to pass to the environments.
     """
 
     # Call the super's ctor.
@@ -123,6 +126,8 @@ def __init__(self,
     # If set, we discretize the rewards and treat them as integers.
     self._discrete_rewards = discrete_rewards
 
+    self._parallelism = None
+
     # A data structure to hold the `batch_size` currently active trajectories
     # and also the ones that are completed, i.e. done.
     self._trajectories = None
@@ -130,7 +135,7 @@ def __init__(self,
     self._batch_size = None
 
     if batch_size is not None:
-      self.initialize(batch_size=batch_size, **kwargs)
+      self.initialize(batch_size=batch_size, **env_kwargs)
 
   @property
   def batch_size(self):

From 8eb6c27f2ad7e71e131fef540622764ac8bbd839 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 23 Jul 2019 11:59:28 -0700
Subject: [PATCH 2220/2720] Add unit test for ASR input, and Area Attention.

PiperOrigin-RevId: 259582396
---
 tensor2tensor/models/transformer_test.py | 38 ++++++++++++++++++++++--
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 20f1ee726..6f55d44cb 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 import numpy as np
 
+from tensor2tensor.data_generators import librispeech
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.models import transformer
 
@@ -41,9 +42,10 @@ def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
   hparams.num_heads = 1
   hparams.layer_prepostprocess_dropout = 0.0
 
-  p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE,
-                                                   VOCAB_SIZE,
-                                                   hparams)
+  if hparams.get("problem_hparams", None) is None:
+    p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE,
+                                                     VOCAB_SIZE,
+                                                     hparams)
   if not has_input:
     del p_hparams.modality["inputs"]
   hparams.problem_hparams = p_hparams
@@ -72,6 +74,36 @@ def testTransformer(self):
       res = session.run(logits)
     self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, 1, 1, VOCAB_SIZE))
 
+  def testTransformerLibrispeech(self, params=None):
+    model_hparams = transformer.transformer_small()
+    if params is not None:  # Add or Set any provided HParams
+      assert isinstance(params, dict)
+      for param_name in params:
+        if hasattr(model_hparams, param_name):
+          model_hparams.set_hparam(param_name, params[param_name])
+        else:
+          model_hparams.add_hparam(param_name, params[param_name])
+    problem = librispeech.Librispeech()
+    model_hparams.problem_hparams = problem.get_hparams(model_hparams)
+    model_hparams._problem_hparams = model_hparams.problem_hparams
+    model, features = get_model(model_hparams)
+    model._problem_hparams.modality = {"inputs": "speech_recognition",
+                                       "targets": "symbol"}
+    features["inputs"] = np.random.rand(
+        BATCH_SIZE, INPUT_LENGTH, 80, 3).astype("float32")  # modify for speech
+
+    logits, _ = model(features)
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      res = session.run(logits)
+    self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, 1, 1, VOCAB_SIZE))
+
+  def testTransformerLibrispeechWithAreaAttention(self):
+    self.testTransformerLibrispeech({"max_area_width": 2,
+                                     "num_area_layers": 1,
+                                     "area_key_mode": "mean",
+                                     "area_value_mode": "sum"})
+
   def testTransformerRelative(self):
     model, features = get_model(transformer.transformer_relative_tiny())
     logits, _ = model(features)

From c9036a69d39695ee2dfd58ebe77fd0cc8d708177 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 23 Jul 2019 12:04:01 -0700
Subject: [PATCH 2221/2720] add WER metric to the dictionary of metric
 functions, so T2T acknowledges its existence.

PiperOrigin-RevId: 259583596
---
 tensor2tensor/utils/metrics.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index b85b60d9c..661566de1 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -866,4 +866,5 @@ def pearson_correlation_coefficient(predictions, labels, weights_fn=None):
     Metrics.DMOL_PERPLEXITY: dmol_neg_log_perplexity,
     Metrics.ABS_ERR: abs_error,
     Metrics.IMAGE_RMSE: image_rmse,
+    Metrics.WORD_ERROR_RATE: word_error_rate,
 }

From 7f1dec4da1c863f8d0b10dabaeda6a4e8a09ed5e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 23 Jul 2019 12:11:01 -0700
Subject: [PATCH 2222/2720] add hparams only if they're not already set.

PiperOrigin-RevId: 259584798
---
 .../data_generators/speech_recognition.py     | 28 +++++++++++--------
 1 file changed, 16 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 5426988b3..3201a92fd 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -43,21 +43,25 @@ class SpeechRecognitionProblem(problem.Problem):
   """Base class for speech recognition problems."""
 
   def hparams(self, defaults, model_hparams):
+    def add_if_absent(p, attr, value):
+      if not hasattr(p, attr):
+        p.add_hparam(attr, value)
+
     p = model_hparams
     # Filterbank extraction in bottom instead of preprocess_example is faster.
-    p.add_hparam("audio_preproc_in_bottom", False)
+    add_if_absent(p, "audio_preproc_in_bottom", False)
     # The trainer seems to reserve memory for all members of the input dict
-    p.add_hparam("audio_keep_example_waveforms", False)
-    p.add_hparam("audio_sample_rate", 16000)
-    p.add_hparam("audio_preemphasis", 0.97)
-    p.add_hparam("audio_dither", 1.0 / np.iinfo(np.int16).max)
-    p.add_hparam("audio_frame_length", 25.0)
-    p.add_hparam("audio_frame_step", 10.0)
-    p.add_hparam("audio_lower_edge_hertz", 20.0)
-    p.add_hparam("audio_upper_edge_hertz", 8000.0)
-    p.add_hparam("audio_num_mel_bins", 80)
-    p.add_hparam("audio_add_delta_deltas", True)
-    p.add_hparam("num_zeropad_frames", 250)
+    add_if_absent(p, "audio_keep_example_waveforms", False)
+    add_if_absent(p, "audio_sample_rate", 16000)
+    add_if_absent(p, "audio_preemphasis", 0.97)
+    add_if_absent(p, "audio_dither", 1.0 / np.iinfo(np.int16).max)
+    add_if_absent(p, "audio_frame_length", 25.0)
+    add_if_absent(p, "audio_frame_step", 10.0)
+    add_if_absent(p, "audio_lower_edge_hertz", 20.0)
+    add_if_absent(p, "audio_upper_edge_hertz", 8000.0)
+    add_if_absent(p, "audio_num_mel_bins", 80)
+    add_if_absent(p, "audio_add_delta_deltas", True)
+    add_if_absent(p, "num_zeropad_frames", 250)
 
     p = defaults
     p.modality = {"inputs": modalities.ModalityType.SPEECH_RECOGNITION,

From b196b45b8ba8cdcc0086aca0a7acf1bc67c0ac1d Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 23 Jul 2019 14:00:09 -0700
Subject: [PATCH 2223/2720] Allow custom gradients in Trax layers, start
 simplifying reversible Transformer.

PiperOrigin-RevId: 259605806
---
 tensor2tensor/trax/layers/base.py             | 54 +++++++++++-
 tensor2tensor/trax/layers/base_test.py        | 60 +++++++++++++
 .../models/research/transformer_revnet.py     | 88 ++++++++-----------
 3 files changed, 148 insertions(+), 54 deletions(-)

diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 80b583f08..503230588 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -138,6 +138,29 @@ def sublayers(self):
     """Returns the sublayers contained in / managed by this layer."""
     return ()  # Default is no sublayers; subclasses can override.
 
+  @property
+  def has_custom_grad(self):
+    """Whether to use custom gradients (in which case, see below)."""
+    return False
+
+  def custom_grad(self, inputs, output, grad, params, **kwargs):
+    """Custom backward pass to propagate gradients in a custom way.
+
+    Args:
+      inputs: Input activations; can be a (possibly nested) tuple.
+      output: The result of running this layer on inputs.
+      grad: gradient signal (called cotangent in jax) computed based on
+        subsequent layers. The structure and shape must match output.
+      params: layer parameters
+      **kwargs: kwargs for the layer
+
+    Returns:
+      The custom gradient signal for the input. Note that we need to return
+      a gradient for each argument of call, so it will usually be a triple
+      of signals: the gradient for inputs, parameters, and kwargs.
+    """
+    raise NotImplementedError
+
   # End of subclassing interface, all functions below are internal.
 
   def pseudo_call(self, pseudo_inputs, params):
@@ -209,8 +232,33 @@ def __call__(self, x, params=(), **kwargs):
         params = self._params
       # In this case, we're called for the first time: cache parameters.
       self._params = params
-      f = lambda y: self.call(y, params=params, **kwargs)
-      return f(x)
+
+      if not self.has_custom_grad:
+        return self.call(x, params=params, **kwargs)
+
+      # Custom gradients part.
+      assert backend.get_name() == 'jax', (
+          'Custom gradients are only supported in JAX for now.')
+
+      # See this link for how custom transformations are defined in JAX:
+      # https://jax.readthedocs.io/en/latest/jax.html#jax.custom_transforms
+      @jax.custom_transforms
+      def do_call(y, params, kwargs):
+        return self.call(y, params=params, **kwargs)
+
+      # This is the custom gradient (vector-jacobian product in JAX) function.
+      # For the exact specification of this custom transformation see this link:
+      # https://jax.readthedocs.io/en/latest/jax.html#jax.defjvp_all
+      # Note that we make arguments positional to allow gradients wrt. them.
+      def do_call_vjp(y, params, kwargs):
+        output = self.call(y, params=params, **kwargs)
+        def vjpfun(grad):
+          return self.custom_grad(y, output, grad, params, **kwargs)
+        return output, vjpfun
+
+      jax.defvjp_all(do_call, do_call_vjp)
+      return do_call(x, params, kwargs)
+
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback()
       raise LayerError(name, 'call', self._caller, shapes(x), trace)
@@ -370,7 +418,7 @@ def _shorten_file_path(line):
   return line[:first_quote] + '[...]/' + new_path + line[second_quote + 1:]
 
 
-def _short_traceback(skip=7):
+def _short_traceback(skip=3):
   """Cleaned-up form of traceback."""
   counter, res = 0, []
   # Skipping 3 lines by default: the top (useless) and self-call.
diff --git a/tensor2tensor/trax/layers/base_test.py b/tensor2tensor/trax/layers/base_test.py
index b8b9bc2c9..7b15484ec 100644
--- a/tensor2tensor/trax/layers/base_test.py
+++ b/tensor2tensor/trax/layers/base_test.py
@@ -19,6 +19,7 @@
 from __future__ import print_function
 
 from absl.testing import absltest
+from tensor2tensor.trax import backend
 from tensor2tensor.trax.layers import base
 
 
@@ -33,6 +34,65 @@ def add_one(x, **unused_kwargs):
         add_one(), (12, 17))  # pylint: disable=no-value-for-parameter
     self.assertEqual(output_shape, (12, 17))
 
+  def test_custom_zero_grad(self):
+
+    class IdWithZeroGrad(base.Layer):
+
+      def call(self, x, params, **kwargs):
+        del params, kwargs
+        return x
+
+      def new_parameters(self, input_shapes, input_dtype, rng):
+        del input_shapes, input_dtype, rng
+        return ()
+
+      @property
+      def has_custom_grad(self):
+        return True
+
+      def custom_grad(self, inputs, output, ct, params, **kwargs):
+        return (backend.numpy.zeros_like(ct), None, None)
+
+    layer = IdWithZeroGrad()
+    rng = backend.random.get_prng(0)
+    params = ()
+    input_shape = (9, 17)
+    random_input = backend.random.uniform(rng, input_shape, minval=-1.0,
+                                          maxval=1.0)
+    f = lambda x: backend.numpy.mean(layer(x, params, rng=rng))
+    grad = backend.grad(f)(random_input)
+    self.assertEqual(grad.shape, input_shape)  # Gradient for each input.
+    self.assertEqual(sum(sum(grad * grad)), 0.0)  # Each one is 0.
+
+  def test_custom_id_grad(self):
+
+    class IdWithIdGrad(base.Layer):
+
+      def call(self, x, params, **kwargs):
+        del params, kwargs
+        return x
+
+      def new_parameters(self, input_shapes, input_dtype, rng):
+        del input_shapes, input_dtype, rng
+        return ()
+
+      @property
+      def has_custom_grad(self):
+        return True
+
+      def custom_grad(self, inputs, output, ct, params, **kwargs):
+        return (inputs, None, None)
+
+    layer = IdWithIdGrad()
+    rng = backend.random.get_prng(0)
+    params = ()
+    input_shape = (9, 17)
+    random_input = backend.random.uniform(rng, input_shape, minval=-1.0,
+                                          maxval=1.0)
+    f = lambda x: backend.numpy.mean(layer(x, params, rng=rng))
+    grad = backend.grad(f)(random_input)
+    self.assertEqual(grad.shape, input_shape)  # Gradient for each input.
+    self.assertEqual(sum(sum(grad)), sum(sum(random_input)))  # Same as input.
 
 if __name__ == "__main__":
   absltest.main()
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 7a5fe0a29..1dfd3f368 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -32,16 +32,19 @@
 
 
 class Map(tl.Layer):
-  """Combinator for applying a layer to a list or tuple.
+  """Combinator for applying a layer to a list or tuple."""
 
-  Args:
-    layer: a layer to apply to each element.
+  def __init__(self, layer, n_sections=1, check_shapes=True):
+    """Initialize the combinator.
 
-  Returns:
-    A new layer representing mapping layer to all elements of the input.
-  """
+    Args:
+      layer: a layer to apply to each element.
+      n_sections: how many sections to map to (default: 1).
+      check_shapes: whether to check that shapes are identical (default: true).
 
-  def __init__(self, layer, sections=1, check_shapes=True):
+    Returns:
+      A new layer representing mapping layer to all elements of the input.
+    """
     super(Map, self).__init__()
     if layer is None or isinstance(layer, (list, tuple)):
       layer = tl.Serial(layer)
@@ -52,15 +55,15 @@ def __init__(self, layer, sections=1, check_shapes=True):
     # are valid cases -- e.g., when self._layer has no parameters -- where we
     # can apply Map to different shapes -- set check_shapes=False in such cases.
     self._check_shapes = check_shapes
-    self._sections = sections
+    self._n_sections = n_sections
 
   def n_inputs(self):
     """Specifies how many data tensors this layer expects as input."""
-    return self._sections
+    return self._n_sections
 
   def n_outputs(self):
     """Specifies how many data tensors this layer promises as output."""
-    return self._sections
+    return self._n_sections
 
   def call(self, inputs, params=(), **kwargs):
     rngs = _pop_rng_and_split(kwargs, len(inputs))
@@ -118,7 +121,7 @@ def inverse_and_vjp(self, output, ct, params=(), **kwargs):
 
     # Note: jax.vjp does not allow us to use **kwargs in the signature here.
     def _do_call(x, params, kwargs):
-      return super(ReversibleLayerMixin, self).__call__(x, params, **kwargs)
+      return super(ReversibleLayerMixin, self).call(x, params=params, **kwargs)
 
     reconstructed_x, must_be_none = self.inverse_and_vjp(
         output, None, params, **kwargs)
@@ -127,43 +130,27 @@ def _do_call(x, params, kwargs):
     input_ct = vjpfun(ct)
     return reconstructed_x, input_ct
 
-  def __call__(self, x, params=(), **kwargs):
-    assert backend.get_name() == 'jax', (
-        'Reversible layers are only supported in JAX')
+  @property
+  def has_custom_grad(self):
+    return True
 
-    if params is () and self._params:  # pylint: disable=literal-comparison
-      # TODO(kitaev): Figure out why parameter sharing doesn't work (if this
-      # explicit error isn't thrown, a jax tracer error occurs instead)
-      raise NotImplementedError(
-          'Parameter sharing between reversible layers is not implemented.')
-
-    @jax.custom_transforms
-    def do_call(x, params, kwargs):
-      return super(ReversibleLayerMixin, self).__call__(x, params, **kwargs)
-
-    def do_call_vjp(x, params, kwargs):
-      output = super(ReversibleLayerMixin, self).__call__(x, params, **kwargs)
-      def vjpfun(ct):
-        _, input_ct = self.inverse_and_vjp(output, ct, params, **kwargs)
-        return input_ct
-
-      return output, vjpfun
-
-    jax.defvjp_all(do_call, do_call_vjp)
-    return do_call(x, params, kwargs)
+  def custom_grad(self, inputs, output, ct, params, **kwargs):
+    del inputs
+    _, input_ct = self.inverse_and_vjp(output, ct, params, **kwargs)
+    return input_ct
 
 
 class Split(tl.Layer):
   """Splits the input into sections along an axis."""
 
-  def __init__(self, sections=2, axis=-1):
+  def __init__(self, n_sections=2, axis=-1):
     super(Split, self).__init__()
-    self._sections = sections
+    self._n_sections = n_sections
     self._axis = axis
 
   def call(self, inputs, params=(), **kwargs):
     del params, kwargs
-    return tuple(backend.numpy.split(inputs, self._sections, self._axis))
+    return tuple(backend.numpy.split(inputs, self._n_sections, self._axis))
 
   def new_parameters(self, input_shapes, input_dtype, rng):
     return ()
@@ -174,26 +161,26 @@ def n_inputs(self):
 
   def n_outputs(self):
     """Specifies how many data tensors this layer promises as output."""
-    return self._sections
+    return self._n_sections
 
 
 @tl.layer()
-def Chunk(x, params, sections=2, **kwargs):
+def Chunk(x, params, n_sections=2, **kwargs):
   del params, kwargs
-  assert x.shape[1] % sections == 0
+  assert x.shape[1] % n_sections == 0
   return backend.numpy.reshape(x, (
-      x.shape[0] * sections,
-      x.shape[1] // sections,
+      x.shape[0] * n_sections,
+      x.shape[1] // n_sections,
       ) + x.shape[2:])
 
 
 @tl.layer()
-def Unchunk(x, params, sections=2, **kwargs):
+def Unchunk(x, params, n_sections=2, **kwargs):
   del params, kwargs
-  assert x.shape[0] % sections == 0
+  assert x.shape[0] % n_sections == 0
   return backend.numpy.reshape(x, (
-      x.shape[0] // sections,
-      x.shape[1] * sections,
+      x.shape[0] // n_sections,
+      x.shape[1] * n_sections,
       ) + x.shape[2:])
 
 
@@ -653,7 +640,7 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
   """
 
   pre_attention = [
-      Chunk(sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
+      Chunk(n_sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
       tl.LayerNorm(),
       tl.Dup(), tl.Dup(),
       tl.Parallel(
@@ -676,7 +663,7 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
   post_attention = [
       JoinHeads(),  # pylint: disable=no-value-for-parameter
       tl.Dense(d_model),
-      Unchunk(sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
+      Unchunk(n_sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
   ]
 
   feed_forward = [
@@ -744,10 +731,9 @@ def TransformerRevnetLM(vocab_size,
       ]),
       tl.Parallel(tl.LayerNorm(), tl.LayerNorm()),
       tl.Concatenate(),
-      Split(sections=n_chunks, axis=-2),  # pylint: disable=no-value-for-parameter
+      Split(n_sections=n_chunks, axis=-2),  # pylint: disable=no-value-for-parameter
       Map([
           tl.Dense(vocab_size),
           tl.LogSoftmax(),
-      ], sections=n_chunks),
+      ], n_sections=n_chunks),
   )
-

From c7cfa8b6edbc0db6adb02f8a61e371fe9d9d40e0 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 23 Jul 2019 14:14:08 -0700
Subject: [PATCH 2224/2720] add test to make sure WER metric is in the
 dictionary of metric functions

PiperOrigin-RevId: 259608794
---
 tensor2tensor/utils/metrics_test.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 3172ad19d..5e0037385 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -127,6 +127,10 @@ def testSequenceEditDistanceMetric(self):
 
   def testWordErrorRateMetric(self):
 
+    # Ensure availability of the WER metric function in the dictionary.
+    assert metrics.Metrics.WORD_ERROR_RATE in metrics.METRICS_FNS
+
+    # Test if WER is computed correctly.
     ref = np.asarray([
         # a b c
         [97, 34, 98, 34, 99],

From 63c015f964c1166d181d8efd232abd856574fd83 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 23 Jul 2019 19:01:17 -0700
Subject: [PATCH 2225/2720] Add Adafactor optimizer.

Add Adafactor to optimizers, provide example base-transformer wmt ende training configs for adam, adafactor, sm3 optimizers.

PiperOrigin-RevId: 259658045
---
 .../transformer_wmt_ende_8gb_adafactor.gin    |  63 ++++++++++
 ....gin => transformer_wmt_ende_8gb_adam.gin} |  14 +--
 .../configs/transformer_wmt_ende_8gb_sm3.gin  |  57 +++++++++
 tensor2tensor/trax/optimizers/__init__.py     |   1 +
 tensor2tensor/trax/optimizers/base.py         | 108 ++++++++++++++++++
 5 files changed, 236 insertions(+), 7 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
 rename tensor2tensor/trax/configs/{transformer_wmt_ende_8gb.gin => transformer_wmt_ende_8gb_adam.gin} (98%)
 create mode 100644 tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin

diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
new file mode 100644
index 000000000..e2b9ce3b6
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
@@ -0,0 +1,63 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 64
+batch_fun.eval_batch_size = 64
+batch_fun.max_eval_length = 1024
+batch_fun.buckets_include_inputs_in_length=True
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_translate_ende_wmt32k'
+inputs.append_targets = True
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 1.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 10000
+
+# Parameters for Adafactor:
+# ==============================================================================
+Adafactor.beta1 = 0.0
+Adafactor.decay_rate = 0.8
+Adafactor.clipping_threshold = 1.0
+Adafactor.epsilon1 = 1e-30
+Adafactor.epsilon2 = 0.001
+Adafactor.factored = True
+Adafactor.multiply_by_parameter_scale = True
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_preprocess
+wmt_preprocess.max_length = 512
+wmt_preprocess.max_eval_length = 1024
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 1000
+train.eval_steps = 10
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.Transformer
+train.train_steps = 500000
+train.optimizer = @trax.optimizers.Adafactor
+
+# Parameters for Transformer:
+# ==============================================================================
+Transformer.d_model = 512
+Transformer.d_ff = 2048
+Transformer.dropout = 0.1
+Transformer.max_len = 2048
+Transformer.mode = 'train'
+Transformer.n_heads = 8
+Transformer.n_layers = 6
+Transformer.input_vocab_size = 33300
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
similarity index 98%
rename from tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
rename to tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
index 01972aeec..e167181e9 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
@@ -22,11 +22,17 @@ masked_mean.mask_id = 0
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-# 0.044 ~= 512^-0.5 = feature_depth^-0.5
+# 0.044 ~= 512^-0.5 = d_model^-0.5
 MultifactorSchedule.constant = 0.044
 MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
 MultifactorSchedule.warmup_steps = 8000
 
+# Parameters for Adam:
+# ==============================================================================
+Adam.b1 = 0.9
+Adam.b2 = 0.98
+Adam.eps = 1e-9
+
 # Parameters for preprocess_fun:
 # ==============================================================================
 shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_preprocess
@@ -42,12 +48,6 @@ train.model = @trax.models.Transformer
 train.train_steps = 500000
 train.optimizer = @trax.optimizers.Adam
 
-# Parameters for Adam:
-# ==============================================================================
-Adam.b1 = 0.9
-Adam.b2 = 0.98
-Adam.eps = 1e-9
-
 # Parameters for Transformer:
 # ==============================================================================
 Transformer.d_model= 512
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin
new file mode 100644
index 000000000..a5ae925e9
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin
@@ -0,0 +1,57 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 64
+batch_fun.eval_batch_size = 64
+batch_fun.max_eval_length = 1024
+batch_fun.buckets_include_inputs_in_length=True
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_translate_ende_wmt32k'
+inputs.append_targets = True
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 0.1
+MultifactorSchedule.factors = 'constant * linear_warmup'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for SM3:
+# ==============================================================================
+SM3.momentum = 0.9
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_preprocess
+wmt_preprocess.max_length = 512
+wmt_preprocess.max_eval_length = 1024
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 1000
+train.eval_steps = 10
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.Transformer
+train.train_steps = 500000
+train.optimizer = @trax.optimizers.SM3
+
+# Parameters for Transformer:
+# ==============================================================================
+Transformer.d_model= 512
+Transformer.d_ff = 2048
+Transformer.dropout = 0.1
+Transformer.max_len = 2048
+Transformer.mode = 'train'
+Transformer.n_heads = 8
+Transformer.n_layers = 6
+Transformer.input_vocab_size = 33300
diff --git a/tensor2tensor/trax/optimizers/__init__.py b/tensor2tensor/trax/optimizers/__init__.py
index 7e3fd7f8e..5974e76b9 100644
--- a/tensor2tensor/trax/optimizers/__init__.py
+++ b/tensor2tensor/trax/optimizers/__init__.py
@@ -33,4 +33,5 @@ def opt_configure(*args, **kwargs):
 Momentum = opt_configure(base.Momentum)
 RMSProp = opt_configure(base.RMSProp)
 Adam = opt_configure(base.Adam)
+Adafactor = opt_configure(base.Adafactor)
 SM3 = opt_configure(base.SM3)
diff --git a/tensor2tensor/trax/optimizers/base.py b/tensor2tensor/trax/optimizers/base.py
index cc3505412..b9f9fac3c 100644
--- a/tensor2tensor/trax/optimizers/base.py
+++ b/tensor2tensor/trax/optimizers/base.py
@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
+
 from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers import base as layers
 
@@ -211,6 +213,112 @@ def update(self, i, g, x, state):
     return x, (m, v)
 
 
+class Adafactor(Optimizer):
+  """Adafactor optimizer."""
+
+  def __init__(self,
+               step_size,
+               decay_rate=0.8,
+               beta1=0.0,
+               clipping_threshold=1.0,
+               factored=True,
+               multiply_by_parameter_scale=True,
+               epsilon1=1e-30,
+               epsilon2=1e-3):
+    """Create the Adafactor optimizer.
+
+    Adafactor is described in https://arxiv.org/abs/1804.04235.
+
+    Args:
+      step_size: function i -> float, trax-provided learning rate schedule.
+      decay_rate: float: controls second-moment exponential decay schedule.
+      beta1: a float value between 0 and 1, enables momentum and uses extra
+        memory if nonzero!  Off by default.
+      clipping_threshold: an optional float >= 1, if None no update clipping.
+      factored: boolean: whether to use factored second-moment estimator for 2d
+        variables.
+      multiply_by_parameter_scale: boolean: if True, then scale provided
+        step_size by parameter norm. if False, provided step_size is absolute
+        step size.
+      epsilon1: Regularization constant for squared gradient.
+      epsilon2: Regularization constant for parameter scale.
+    """
+    super(Adafactor, self).__init__(step_size)
+    self._multiply_by_parameter_scale = multiply_by_parameter_scale
+    self._beta1 = beta1
+    self._clipping_threshold = clipping_threshold
+    self._factored = factored
+    self._epsilon1 = epsilon1
+    self._epsilon2 = epsilon2
+    self._step_size = step_size
+    self._decay_rate = functools.partial(self._decay_rate_pow,
+                                         exponent=decay_rate)
+
+  @staticmethod
+  def _decay_rate_pow(i, exponent=0.8):
+    """Default Adafactor second-moment decay schedule."""
+    t = np.array(i, np.float32) + 1.0
+    return 1.0 - t**(-exponent)
+
+  def init(self, x):
+    shape = x.shape
+    state = []
+    if self._factored and len(shape) >= 2:
+      v_row = np.zeros(shape[:-1], dtype=np.float32)
+      v_col = np.zeros(shape[:-2] + shape[-1:], dtype=np.float32)
+      state.extend([v_row, v_col])
+    else:
+      v = np.zeros_like(x)
+      state.append(v)
+    if self._beta1:
+      m = np.zeros_like(x)
+      state.append(m)
+    return state
+
+  def update(self, i, g, x, state):
+    updates = []
+    decay_rate = self._decay_rate(i)
+    update_scale = self._step_size(i)
+    if self._multiply_by_parameter_scale:
+      update_scale *= np.maximum(np.sqrt(np.mean(x * x)), self._epsilon2)
+    mixing_rate = 1.0 - decay_rate
+
+    g_sqr = g * g + self._epsilon1
+    if self._factored and len(x.shape) >= 2:
+      v_row = state.pop(0)
+      v_col = state.pop(0)
+      new_v_row = decay_rate * v_row + mixing_rate * np.mean(g_sqr, axis=-1)
+      new_v_col = decay_rate * v_col + mixing_rate * np.mean(g_sqr, axis=-2)
+      updates.extend([new_v_row, new_v_col])
+      row_col_mean = np.mean(new_v_row, axis=-1, keepdims=True)
+      row_factor = (new_v_row / row_col_mean)**-0.5
+      col_factor = (new_v_col)**-0.5
+      y = (
+          g * np.expand_dims(row_factor, axis=-1) *
+          np.expand_dims(col_factor, axis=-2))
+    else:
+      v = state.pop(0)
+      new_v = decay_rate * v + mixing_rate * g_sqr
+      updates.append(new_v)
+      y = g * (new_v)**-0.5
+
+    if self._clipping_threshold is not None:
+      clipping_denom = (
+          np.maximum(1.0,
+                     np.sqrt(np.mean(y * y)) / self._clipping_threshold))
+      y /= clipping_denom
+
+    subtrahend = update_scale * y
+    if self._beta1:
+      m = state.pop(0)
+      new_m = self._beta1 * m + (1.0 - self._beta1) * subtrahend
+      subtrahend = new_m
+      updates.append(new_m)
+
+    new_x = x - subtrahend
+    return new_x, updates
+
+
 class SM3(Optimizer):
   """SM3 optimizer."""
 

From f05ec67e0446f069fad84324d10ffc13bbf95fb5 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 24 Jul 2019 09:59:54 -0700
Subject: [PATCH 2226/2720] Pass learning rate etc. to optimizers explicitly as
 opt_params to avoid rejitting

PiperOrigin-RevId: 259761692
---
 tensor2tensor/trax/optimizers/base.py         | 235 ++++++++++--------
 .../trax/rlax/envs/online_tune_env.py         |   2 +-
 tensor2tensor/trax/rlax/ppo.py                |  37 ++-
 tensor2tensor/trax/rlax/ppo_main.py           |   2 +-
 .../trax/rlax/simulated_env_problem.py        |   3 +-
 tensor2tensor/trax/trax.py                    | 114 +++++----
 tensor2tensor/trax/trax_test.py               |  32 ++-
 7 files changed, 259 insertions(+), 166 deletions(-)

diff --git a/tensor2tensor/trax/optimizers/base.py b/tensor2tensor/trax/optimizers/base.py
index b9f9fac3c..41c113534 100644
--- a/tensor2tensor/trax/optimizers/base.py
+++ b/tensor2tensor/trax/optimizers/base.py
@@ -67,48 +67,73 @@ def tree_unflatten(flat, tree):
 class Optimizer(object):
   """Optimizer object, base class. Maps per-parameter functions to trees."""
 
-  def __init__(self, step_size):
-    """Optimizers take the step size function (learning rate) as argument."""
-    if callable(step_size):
-      self._step_size = step_size
-    else:
-      self._step_size = lambda _: step_size
+  def __init__(self, learning_rate, *init_opt_params):
+    """Initialize the optimizer.
 
-  def init(self, x):
-    """Create optimizer slots for the parameter x."""
+    Takes the initial optimizer parameters as positional arguments. They are fed
+    back to the optimizer in tree_update, in the same order. They can be changed
+    between updates, e.g. for learning rate schedules.
+
+    The constructor should be overridden in derived classes to give names to the
+    optimizer parameters, so the gin configuration can set them.
+
+    Args:
+      learning_rate: The initial learning rate.
+      *init_opt_params: Initial values of any additional optimizer parameters.
+    """
+    self._init_opt_params = tuple(
+        map(np.array, (learning_rate,) + init_opt_params))
+
+  def init(self, params):
+    """Create optimizer slots for the given parameters."""
     raise NotImplementedError
 
-  def update(self, i, g, x, s):
-    """Update the parameter x at step i with gradient g using state s."""
+  def update(self, step, grads, params, slots, opt_params):
+    """Update a single parameter array.
+
+    Args:
+      step: Current step.
+      grads: Gradients.
+      params: Parameters.
+      slots: Optimizer slots (e.g. gradient moments).
+      opt_params: Optimizer (hyper)parameters (e.g. learning rate, momentum).
+
+    Returns:
+      (new_params, new_slots)
+    """
     raise NotImplementedError
 
   # End subclass interface.
 
-  def step_size(self, i):
-    return self._step_size(i)
-
-  def tree_init(self, x_tree):
-    return [self.init(x) for x in tree_flatten(x_tree)]
-
-  def _update_and_check(self, i, g, x, s):
-    new_x, new_s = self.update(i, g, x, s)
-    if isinstance(x, np.ndarray):
-      assert isinstance(new_x, np.ndarray), ("The type of the new parameter "
-                                             "values should be np.ndarray; "
-                                             "got %s" % type(new_x))
-      assert new_x.dtype == x.dtype, ("The dtype of the new parameter values "
-                                      "(%s) is not the same as the old one (%s)"
-                                      % (new_x.dtype, x.dtype))
-    return new_x, new_s
-
-  def tree_update(self, i, grad_tree, x_tree, opt_state):
-    grad_flat = tree_flatten(grad_tree)
-    x_flat = tree_flatten(x_tree)
-    updated_pairs = [self._update_and_check(i, g, x, s)
-                     for (g, x, s) in zip(grad_flat, x_flat, opt_state)]
-    new_x_flat, new_opt_state = zip(*updated_pairs)
-    new_x, _ = tree_unflatten(new_x_flat, x_tree)
-    return new_x, new_opt_state
+  def tree_init(self, param_tree):
+    return (
+        [self.init(param) for param in tree_flatten(param_tree)],
+        self._init_opt_params,
+    )
+
+  def _update_and_check(self, step, grads, params, slots, opt_params):
+    """Update a single parameter array and check types."""
+    new_params, new_slots = self.update(
+        step, grads, params, slots, opt_params)
+    if isinstance(params, np.ndarray):
+      assert isinstance(new_params, np.ndarray), (
+          "The type of the new parameter values should be np.ndarray; got %s" %
+          type(new_params))
+      assert new_params.dtype == params.dtype, (
+          "The dtype of the new parameter values (%s) is not the same as the "
+          "old one (%s)" % (new_params.dtype, params.dtype))
+    return new_params, new_slots
+
+  def tree_update(self, step, grad_tree, param_tree, slots, opt_params):
+    grads_flat = tree_flatten(grad_tree)
+    params_flat = tree_flatten(param_tree)
+    updated_pairs = [
+        self._update_and_check(step, grad, param, slot, opt_params)
+        for (grad, param, slot) in zip(grads_flat, params_flat, slots)
+    ]
+    new_params_flat, new_slots = zip(*updated_pairs)
+    new_params, _ = tree_unflatten(new_params_flat, param_tree)
+    return new_params, new_slots
 
 
 # Utilities.
@@ -133,58 +158,59 @@ def clip_grads(grad_tree, max_norm):
 class SGD(Optimizer):
   """Plain SGD optimizer."""
 
-  def init(self, x):
+  def init(self, params):
     return None
 
-  def update(self, i, g, x, state):
-    del state
-    return x - (self.step_size(i) * g).astype(x.dtype), None
+  def update(self, step, grads, params, slots, opt_params):
+    del step
+    del slots
+    (learning_rate,) = opt_params
+    return params - (learning_rate * grads).astype(params.dtype), None
 
 
 class Momentum(Optimizer):
   """Nestrov momentum optimizer."""
 
-  def __init__(self, step_size, mass=0.9):
-    """Initializer with a step size function and mass."""
-    super(Momentum, self).__init__(step_size)
-    self._mass = mass
+  def __init__(self, learning_rate, mass=0.9):  # pylint: disable=useless-super-delegation
+    super(Momentum, self).__init__(learning_rate, mass)
 
-  def init(self, x):
-    return np.zeros_like(x)
+  def init(self, params):
+    return np.zeros_like(params)
 
-  def update(self, i, g, x, velocity):
-    new_velocity = self._mass * velocity - (1. - self._mass) * g
-    return x + (self.step_size(i) * new_velocity).astype(x.dtype), new_velocity
+  def update(self, step, grads, params, velocity, opt_params):
+    del step
+    (learning_rate, mass) = opt_params
+    new_velocity = mass * velocity - (1. - mass) * grads
+    new_params = params + (learning_rate * new_velocity).astype(params.dtype)
+    return (new_params, new_velocity)
 
 
 class RMSProp(Optimizer):
   """RMSProp optimizer."""
 
-  def __init__(self, step_size, gamma=0.9, eps=1e-8):
-    """Initializer with a step size function, gamma and epsilon."""
-    super(RMSProp, self).__init__(step_size)
-    self._gamma = gamma
-    self._epsilon = eps
+  def __init__(self, learning_rate, gamma=0.9, eps=1e-8):  # pylint: disable=useless-super-delegation
+    super(RMSProp, self).__init__(learning_rate, gamma, eps)
 
-  def init(self, x):
-    return np.ones_like(x)
+  def init(self, params):
+    return np.ones_like(params)
 
-  def update(self, i, g, x, avg_sq_grad):
-    avg_sq_grad = avg_sq_grad * self._gamma + g**2 * (1. - self._gamma)
-    x = x - (self.step_size(i) * g /
-             (np.sqrt(avg_sq_grad) + self._epsilon)).astype(x.dtype)
-    return x, avg_sq_grad
+  def update(self, step, grads, params, avg_sq_grad, opt_params):
+    del step
+    (learning_rate, gamma, eps) = opt_params
+    avg_sq_grad = avg_sq_grad * gamma + grads**2 * (1. - gamma)
+    params = params - (learning_rate * grads /
+                       (np.sqrt(avg_sq_grad) + eps)).astype(params.dtype)
+    return params, avg_sq_grad
 
 
 class Adam(Optimizer):
   """Adam optimizer."""
 
-  def __init__(self, step_size, b1=0.9, b2=0.999, eps=1e-8):
+  def __init__(self, learning_rate, b1=0.9, b2=0.999, eps=1e-8):  # pylint: disable=useless-super-delegation
     """Create the Adam optimizer.
 
     Args:
-      step_size: a callable representing a step size schedule
-        that maps the iteration index to positive scalar.
+      learning_rate: a postitive scalar value for the initial learning rate.
       b1: optional, a positive scalar value for beta_1, the exponential decay
         rate for the first moment estimates (default 0.9).
       b2: optional, a positive scalar value for beta_2, the exponential decay
@@ -192,25 +218,23 @@ def __init__(self, step_size, b1=0.9, b2=0.999, eps=1e-8):
       eps: optional, a positive scalar value for epsilon, a small constant for
         numerical stability (default 1e-8).
     """
-    super(Adam, self).__init__(step_size)
-    self._b1 = b1
-    self._b2 = b2
-    self._eps = eps
+    super(Adam, self).__init__(learning_rate, b1, b2, eps)
 
-  def init(self, x):
-    m = np.zeros_like(x)
-    v = np.zeros_like(x)
+  def init(self, params):
+    m = np.zeros_like(params)
+    v = np.zeros_like(params)
     return m, v
 
-  def update(self, i, g, x, state):
-    m, v = state
-    b1, b2, eps = self._b1, self._b2, self._eps
-    m = (1 - b1) * g + b1 * m  # First  moment estimate.
-    v = (1 - b2) * (g ** 2) + b2 * v  # Second moment estimate.
-    mhat = m / (1 - b1 ** (i + 1))  # Bias correction.
-    vhat = v / (1 - b2 ** (i + 1))
-    x = x - (self.step_size(i) * mhat / (np.sqrt(vhat) + eps)).astype(x.dtype)
-    return x, (m, v)
+  def update(self, step, grads, params, slots, opt_params):
+    m, v = slots
+    learning_rate, b1, b2, eps = opt_params
+    m = (1 - b1) * grads + b1 * m  # First  moment estimate.
+    v = (1 - b2) * (grads ** 2) + b2 * v  # Second moment estimate.
+    mhat = m / (1 - b1 ** (step + 1))  # Bias correction.
+    vhat = v / (1 - b2 ** (step + 1))
+    params = params - (
+        learning_rate * mhat / (np.sqrt(vhat) + eps)).astype(params.dtype)
+    return params, (m, v)
 
 
 class Adafactor(Optimizer):
@@ -322,32 +346,31 @@ def update(self, i, g, x, state):
 class SM3(Optimizer):
   """SM3 optimizer."""
 
-  def __init__(self, step_size, momentum=0.9):
+  def __init__(self, learning_rate, momentum=0.9):  # pylint: disable=useless-super-delegation
     """Create the SM3 optimizer.
 
     Memory-Efficient Adaptive Optimization for Large-Scale Learning.
     https://arxiv.org/abs/1901.11150
 
     Args:
-      step_size: a callable representing a step size schedule
-        that maps the iteration index to positive scalar.
+      learning_rate: a postitive scalar value for the initial learning rate.
       momentum: optional, a positive scalar value for momentum
     """
-    super(SM3, self).__init__(step_size)
-    self._momentum = momentum
+    super(SM3, self).__init__(learning_rate, momentum)
 
-  def init(self, x):
-    vs = [np.zeros(sz, dtype=x.dtype) for sz in x.shape]
-    return (np.zeros_like(x), vs)
+  def init(self, params):
+    vs = [np.zeros(sz, dtype=params.dtype) for sz in params.shape]
+    return (np.zeros_like(params), vs)
 
-  def _update_diagonal(self, step, g, x, m, v):
-    v[0] += g * g
+  def _update_diagonal(self, grads, params, m, v, opt_params):
+    (learning_rate, momentum) = opt_params
+    v[0] += grads * grads
     preconditioner = np.where(v[0] > 0, 1.0 / np.sqrt(v[0]),
                               np.zeros_like(v[0]))
-    preconditioned_g = preconditioner * g
-    m = (1 - self._momentum) * preconditioned_g + self._momentum * m
-    x = x - (self.step_size(step) * m).astype(x.dtype)
-    return x, (m, v)
+    preconditioned_grads = preconditioner * grads
+    m = (1 - momentum) * preconditioned_grads + momentum * m
+    params = params - (learning_rate * m).astype(params.dtype)
+    return params, (m, v)
 
   def _expanded_shape(self, shape, axis):
     # Replaces a `shape` of [M, N, K] with 1 in all dimensions except for i.
@@ -361,31 +384,33 @@ def _minimum(self, tensor_list):
       minimum = np.minimum(minimum, tensor_list[i])
     return minimum
 
-  def _update_sketched(self, step, g, x, m, v):
+  def _update_sketched(self, grads, params, m, v, opt_params):
     """Update for higher-rank parameters."""
-    shape = x.shape
+    (learning_rate, momentum) = opt_params
+    shape = params.shape
     rank = len(shape)
     reshaped_accumulators = [np.reshape(v[i], self._expanded_shape(shape, i))
                              for i in range(rank)]
     current_accumulator = self._minimum(reshaped_accumulators)
-    current_accumulator += g * g
+    current_accumulator += grads * grads
     accumulator_inv_sqrt = np.where(current_accumulator > 0.0,
                                     1.0 / np.sqrt(current_accumulator),
                                     np.zeros_like(current_accumulator))
-    preconditioned_gradient = g * accumulator_inv_sqrt
-    m = (1.0 - self._momentum) * preconditioned_gradient + self._momentum * m
-    x = x - (self.step_size(step) * m).astype(x.dtype)
+    preconditioned_gradient = grads * accumulator_inv_sqrt
+    m = (1.0 - momentum) * preconditioned_gradient + momentum * m
+    params = params - (learning_rate * m).astype(params.dtype)
     for i in range(len(v)):
       axes = list(range(int(i))) + list(range(int(i) + 1, rank))
       dim_accumulator = np.amax(current_accumulator, axis=axes)
       v[i] = dim_accumulator
-    return x, (m, v)
+    return params, (m, v)
 
-  def update(self, i, g, x, state):
-    m, v = state
-    shape = x.shape
+  def update(self, step, grads, params, slots, opt_params):
+    del step
+    m, v = slots
+    shape = params.shape
     rank = len(shape)
     if rank > 1:
-      return self._update_sketched(i, g, x, m, v)
+      return self._update_sketched(grads, params, m, v, opt_params)
     else:
-      return self._update_diagonal(i, g, x, m, v)
+      return self._update_diagonal(grads, params, m, v, opt_params)
diff --git a/tensor2tensor/trax/rlax/envs/online_tune_env.py b/tensor2tensor/trax/rlax/envs/online_tune_env.py
index b29e9f46c..352956df6 100644
--- a/tensor2tensor/trax/rlax/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rlax/envs/online_tune_env.py
@@ -150,7 +150,7 @@ def step(self, action):
         environment steps. info is an empty dict.
     """
     self._current_lr *= self._action_multipliers[action]
-    self._trainer.update_learning_rate(force_jit=True)
+    self._trainer.update_learning_rate()
     last_metric_value = self._current_metric_value
     self._trainer.train_epoch(self._train_steps, self._eval_steps)
     self._step += 1
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index f727d812d..4e276813a 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -118,13 +118,36 @@ def policy_and_value_net(rng_key,
   return params, net
 
 
-def optimizer_fn(net_params, step_size=1e-3):
-  opt = trax_opt.Adam(step_size=step_size, b1=0.9, b2=0.999, eps=1e-08)
-  opt_init = lambda x: (x, opt.tree_init(x))
-  opt_update = lambda i, g, s: opt.tree_update(i, g, s[0], s[1])
-  get_params = lambda x: x[0]
-  opt_state = opt_init(net_params)
-  return opt_state, opt_update, get_params
+def optimizer_fn(net_params, learning_rate=1e-3):
+  """Exposes a convenient interface for the optimizer.
+
+  Args:
+    net_params: A nested structure of network parameters.
+    learning_rate: Learning rate.
+
+  Returns:
+    A tuple (opt_state, opt_update, get_params), where:
+      opt_state: Pair (net_params, opt_slots) - initial optimization state.
+      opt_update: Function (step, grads, opt_state) -> opt_state doing one
+        optimization step.
+      get_params: Function opt_state -> net_params for extracting the network
+        parameters from the optimization state.
+  """
+  opt = trax_opt.Adam(learning_rate=learning_rate, b1=0.9, b2=0.999, eps=1e-08)
+  (init_slots, init_nontrainable_slots) = opt.tree_init(net_params)
+  init_state = (net_params, init_slots)
+
+  def opt_update(step, grads, opt_state):
+    (params, slots) = opt_state
+    # Pass the initial nontrainable_slots as we don't tune them during training.
+    # (yet!)
+    return opt.tree_update(step, grads, params, slots, init_nontrainable_slots)
+
+  def get_params(opt_state):
+    (params, _) = opt_state
+    return params
+
+  return init_state, opt_update, get_params
 
 
 # Should this be collect 'n' trajectories, or
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 998739d46..d4980f821 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -161,7 +161,7 @@ def make_env(batch_size=8, **env_kwargs):
 
 
 def get_optimizer_fn(learning_rate):
-  return functools.partial(ppo.optimizer_fn, step_size=learning_rate)
+  return functools.partial(ppo.optimizer_fn, learning_rate=learning_rate)
 
 
 def main(argv):
diff --git a/tensor2tensor/trax/rlax/simulated_env_problem.py b/tensor2tensor/trax/rlax/simulated_env_problem.py
index fd3204a41..0834527ac 100644
--- a/tensor2tensor/trax/rlax/simulated_env_problem.py
+++ b/tensor2tensor/trax/rlax/simulated_env_problem.py
@@ -110,8 +110,7 @@ def initialize_environments(self,
     del parallelism
 
     model_state = trax.restore_state(self._output_dir)
-    # model_state.params is a pair (model_params, optimizer_state).
-    (self._model_params, _) = model_state.params
+    self._model_params = model_state.opt_state.params
     self._initial_observation_stream = initial_observation_stream
 
     self._history = None
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 076dd5826..c0c28ae1a 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -124,20 +124,31 @@ def step_log(step, s):
   log("Step % 6d: %s" % (step, s))
 
 
-State = collections.namedtuple("_State", ["step", "params", "history"])
+State = collections.namedtuple("_State", [
+    "step",       # Current training step number.
+    "opt_state",  # OptState.
+    "history",    # trax.history.History.
+])
+
+
+OptState = collections.namedtuple("_OptState", [
+    "params",      # Model parameters.
+    "slots",       # Per-parameter optimizer state, e.g. gradient moments.
+    "opt_params",  # Optimizer (hyper)parameters, e.g. learning rate, momentum.
+])
 
 
 def restore_state(output_dir):
   """Restore State."""
   params_file = os.path.join(output_dir, "model.pkl")
   if not gfile.exists(params_file):
-    return State(step=None, params=None, history=trax_history.History())
+    return State(step=None, opt_state=None, history=trax_history.History())
 
   with gfile.GFile(params_file, "rb") as f:
-    (params, step, history) = pickle.load(f)
+    (opt_state, step, history) = pickle.load(f)
   log("Model loaded from %s at step %d" % (params_file, step))
   logging.debug("From loaded model : history = %s", history)
-  return State(step=step, params=params, history=history)
+  return State(step=step, opt_state=OptState(*opt_state), history=history)
 
 
 def _save_gin(output_dir, sw=None):
@@ -160,11 +171,11 @@ def save_state(state, output_dir, keep=False):
     pkl_module = pickle
   params_file = os.path.join(output_dir, "model.pkl")
   with gfile.GFile(params_file, "wb") as f:
-    pkl_module.dump((state.params, state.step, state.history), f)
+    pkl_module.dump((tuple(state.opt_state), state.step, state.history), f)
   if keep:
     params_file = os.path.join(output_dir, "model_{}.pkl".format(state.step))
     with gfile.GFile(params_file, "wb") as f:
-      pkl_module.dump((state.params, state.step, state.history), f)
+      pkl_module.dump((tuple(state.opt_state), state.step, state.history), f)
   log("Model saved to %s" % params_file, stdout=False)
 
 
@@ -172,21 +183,21 @@ def _save_replicated(opt_state, step, history, n_devices, output_dir, keep):
   """Save state but given a possibly replicated opt_state."""
   if n_devices > 1:
     first_replica = lambda x: x[0]
-    opt_state = layers.nested_map(opt_state, first_replica)
+    opt_state = OptState(*layers.nested_map(opt_state, first_replica))
   # This line, while optional, allows JAX to transfer arrays from the device to
   # the host in parallel, which is particularly important for cloud TPU.
   if backend.get_name() == "jax":
     opt_state = jax.device_get(opt_state)
-  save_state(State(params=opt_state, step=step, history=history),
+  save_state(State(opt_state=opt_state, step=step, history=history),
              output_dir, keep=keep)
 
 
 def _print_n_params(opt_state, n_devices, step):
   """Print out the number of parameters."""
-  sizes = layers.sizes(opt_state[0])
+  sizes = layers.sizes(opt_state.params)
   if n_devices > 1:
     unreplicate = lambda x: x[0]
-    single_params = layers.nested_map(opt_state[0], unreplicate)
+    single_params = layers.nested_map(opt_state.params, unreplicate)
     sizes = layers.sizes(single_params)
   total_size = layers.nested_reduce(sizes, sum)
   step_log(step, "Total trainable parameters size: %d" % total_size)
@@ -370,10 +381,11 @@ def _jit_update_fn(predict_fn, loss_fn, optimizer, n_devices, jit=True):
   """Get jit-ed update function for loss, optimizer, learning rate function."""
   if n_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
     def single_update(i, opt_state, batch, rng):
+      params, slots, opt_params = opt_state
       rng, subrng = jax_random.split(rng[0])
-      params, opt_slots = opt_state
-      return optimizer.tree_update(i, backend.grad(loss_fn)(
-          params, batch, predict_fn, rng), params, opt_slots), [subrng]
+      grads = backend.grad(loss_fn)(opt_state.params, batch, predict_fn, rng)
+      return optimizer.tree_update(
+          i, grads, params, slots, opt_params), [subrng]
     if jit:
       return backend.jit(single_update)
     else:
@@ -383,12 +395,13 @@ def single_update(i, opt_state, batch, rng):
   def mapped_update(i, opt_state, batch, rng):
     """This is a multi-device version of the update function above."""
     # We assume all tensors have the first dimension = n_devices.
+    params, slots, opt_params = opt_state
     rng, subrng = jax_random.split(rng)
-    params, opt_slots = opt_state
-    grads = backend.grad(loss_fn)(params, batch, predict_fn, rng)
+    grads = backend.grad(loss_fn)(opt_state.params, batch, predict_fn, rng)
     grads = jax.tree_util.tree_map(
         lambda g: lax.psum(g, "batch"), grads)
-    return optimizer.tree_update(i, grads, params, opt_slots), subrng
+    return optimizer.tree_update(
+        i, grads, params, slots, opt_params), subrng
 
   def update(i, opt_state, batch, rng):
     return mapped_update(numpy.repeat(i, n_devices), opt_state, batch, rng)
@@ -483,15 +496,15 @@ def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs, output_dir,
 
     # Setup optimizer and model.
     state = restore_state(output_dir)
+    step = state.step or 0
     history = state.history
     self._lr_fn = lr_schedule(history)
-    opt = optimizer(self._lr_fn)
+    opt = optimizer(learning_rate=self._lr_fn(step))
 
     model_train = model(mode="train")
     model_predict_eval = model(mode="eval")
 
     # Setup state.
-    step = state.step or 0
     rng, init_rng = jax_random.split(rng)
     self._rngs = jax_random.split(rng, n_devices)
     first_shape = inputs.input_shape[0]
@@ -504,20 +517,19 @@ def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs, output_dir,
     # Change all None to 1 in input shape.
     model_input_shape = layers.nested_map(
         model_input_shape, lambda x: x if x else 1)
-    if state.params:
-      opt_state = state.params
+    if state.opt_state:
+      opt_state = state.opt_state
     else:
       def initialize(input_shape, input_dtype, init_rng):
         params = model_train.initialize(input_shape, input_dtype, init_rng)
-        opt_state = (params, opt.tree_init(params))
-        return opt_state
+        (slots, opt_params) = opt.tree_init(params)
+        return OptState(params, slots, opt_params)
       if _is_jit_init():
         # JIT parameter initialization to avoid memory fragmentation
         initialize = backend.jit(initialize, static_argnums=(0, 1))
-      opt_state = initialize(model_input_shape, inputs.input_dtype, init_rng)
-    if n_devices > 1:
-      replicate = lambda x: numpy.broadcast_to(x, (n_devices,) + x.shape)
-      opt_state = layers.nested_map(opt_state, replicate)
+      opt_state = initialize(
+          model_input_shape, inputs.input_dtype, init_rng)
+    opt_state = OptState(*layers.nested_map(opt_state, self._maybe_replicate))
 
     # jit model_predict and update so they're fast
     self._jit_model_predict_eval = _jit_predict_fn(
@@ -528,11 +540,12 @@ def initialize(input_shape, input_dtype, init_rng):
     self._model_train = model_train
     self._model_predict_eval = model_predict_eval
     self._loss_fn = loss_fn
-    self._optimizer = optimizer
     self._opt_state = opt_state
     self._history = history
     self._lr_schedule = lr_schedule
 
+    self.update_learning_rate()
+
   @property
   def step(self):
     return self._step
@@ -543,7 +556,8 @@ def n_devices(self):
 
   @property
   def state(self):
-    return State(params=self._opt_state, step=self._step, history=self._history)
+    return State(
+        opt_state=self._opt_state, step=self._step, history=self._history)
 
   @property
   def learning_rate(self):
@@ -553,12 +567,33 @@ def learning_rate(self):
     with backend.use_backend("numpy"):
       return self._lr_fn(self._step)
 
+  def _maybe_replicate(self, x):
+    if self._n_devices > 1:
+      return numpy.broadcast_to(x, (self._n_devices,) + x.shape)
+    else:
+      return x
+
   def save_gin(self):
     _save_gin(self._output_dir, self._train_sw)
 
   def print_n_params(self):
     _print_n_params(self._opt_state, self._n_devices, self._step)
 
+  def _train_step(self, next_train_batch):
+    """Run one training step and update self._opt_state."""
+    # Calculate the current learning rate.
+    learning_rate = self._maybe_replicate(np.array(self._lr_fn(self._step)))
+    opt_state = self._opt_state
+    opt_params = opt_state.opt_params
+    opt_params = (learning_rate,) + opt_params[1:]
+    opt_state = opt_state._replace(opt_params=opt_params)
+
+    # Run the update.
+    (params, slots), self._rngs = self._jit_update_fn(
+        self._step, opt_state, next_train_batch, self._rngs)
+    self._opt_state = opt_state._replace(params=params, slots=slots)
+    self._step += 1
+
   def train_epoch(self, epoch_steps, eval_steps):
     """Train for one epoch."""
     # Log separator
@@ -572,9 +607,8 @@ def train_epoch(self, epoch_steps, eval_steps):
       next_train_batch = next(self._train_stream)
       if self._n_devices > 1:  # TODO(lukaszkaiser): use everywhere if possible.
         next_train_batch = reshape_by_device(next_train_batch, self._n_devices)
-      self._opt_state, self._rngs = self._jit_update_fn(
-          self._step, self._opt_state, next_train_batch, self._rngs)
-      self._step += 1
+
+      self._train_step(next_train_batch)
 
       if self._step in self._save_steps:
         _save_replicated(self._opt_state, self._step, self._history,
@@ -620,14 +654,8 @@ def evaluate(self, eval_steps):
         eval_sw=self._eval_sw,
         history=self._history)
 
-  def update_learning_rate(self, force_jit=False):
-    old_lr_fn = self._lr_fn
+  def update_learning_rate(self):
     self._lr_fn = self._lr_schedule(self._history)
-    # For performance only jit if it's changed or we force it.
-    if self._lr_fn != old_lr_fn or force_jit:
-      opt = self._optimizer(self._lr_fn)
-      self._jit_update_fn = _jit_update_fn(
-          self._model_train, self._loss_fn, opt, self._n_devices)
 
   def save_computation_graphs(self, save_backward_graph):
     """Dump computation graphs to files."""
@@ -682,16 +710,6 @@ def evaluate(self, eval_steps):
         eval_sw=self._eval_sw,
         history=self._history)
 
-  def update_learning_rate(self, force_jit=False):
-    old_lr_fn = self._lr_fn
-    self._lr_fn = self._lr_schedule(self._history)
-    if self._lr_fn != old_lr_fn or force_jit:
-      raise NotImplementedError(
-          "Loss function changed or jitting was requested. Garbage collection "
-          "for jitted functions is not implemented in jax, so global "
-          "accelerator memory allocated by the jitted update function with the "
-          "old loss cannot be reclaimed.")
-
   def save_computation_graphs(self, save_backward_graph):
     # TODO(kitaev): implement saving graphs while making sure that no op-by-op
     # execution happens in the process.
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index 1f8ee2504..956b64f10 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -93,7 +93,7 @@ def test_train_eval_predict(self):
       # Predict with final params
       inputs = inputs(1).train_stream()
       model = layers.Serial(model_fn())
-      model(next(inputs)[0], state.params[0])
+      model(next(inputs)[0], state.opt_state.params)
 
   def test_train_eval_predict_sm3(self):
     with self.tmp_dir() as output_dir:
@@ -126,7 +126,35 @@ def test_train_eval_predict_sm3(self):
       # Predict with final params
       inputs = inputs(1).train_stream()
       model = layers.Serial(model_fn())
-      model(next(inputs)[0], state.params[0])
+      model(next(inputs)[0], state.opt_state.params)
+
+  def test_train_restart(self):
+    with self.tmp_dir() as output_dir:
+      # Prepare model and inputs
+      n_classes = 4
+      train_steps = 2
+      eval_steps = 2
+      model_fn = functools.partial(models.MLP,
+                                   d_hidden=16,
+                                   n_output_classes=n_classes)
+      inputs = lambda _: test_inputs(n_classes)
+
+      # Train and evaluate
+      trax.train(output_dir,
+                 model=model_fn,
+                 inputs=inputs,
+                 train_steps=train_steps,
+                 eval_steps=eval_steps)
+
+      # Restart training
+      state = trax.train(output_dir,
+                         model=model_fn,
+                         inputs=inputs,
+                         train_steps=train_steps,
+                         eval_steps=eval_steps)
+
+      # Assert total train steps
+      self.assertEqual(state.step, 2 * train_steps)
 
 
 if __name__ == "__main__":

From a9830d08a9d7a3f9e889600cd8eb9da69601bc2c Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 24 Jul 2019 10:10:02 -0700
Subject: [PATCH 2227/2720] Add a reset() method to Trainer for resetting the
 model parameters without rejitting.

PiperOrigin-RevId: 259764087
---
 .../trax/rlax/envs/online_tune_env.py         |  8 +-
 .../trax/rlax/envs/online_tune_env_test.py    |  7 +-
 tensor2tensor/trax/trax.py                    | 96 ++++++++++++-------
 3 files changed, 70 insertions(+), 41 deletions(-)

diff --git a/tensor2tensor/trax/rlax/envs/online_tune_env.py b/tensor2tensor/trax/rlax/envs/online_tune_env.py
index 352956df6..b5cba2427 100644
--- a/tensor2tensor/trax/rlax/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rlax/envs/online_tune_env.py
@@ -19,7 +19,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import functools
 import os
 
 import gym
@@ -63,8 +62,7 @@ def __init__(self,
     if action_multipliers is None:
       action_multipliers = self.DEFAULT_ACTION_MULTIPLIERS
     self._model = model
-    self._trainer_fn = functools.partial(
-        trainer_class,
+    self._trainer = trainer_class(
         model=model,
         loss_fn=loss_fn,
         optimizer=optimizer,
@@ -77,7 +75,6 @@ def __init__(self,
     self._eval_steps = eval_steps
     self._env_steps = env_steps
     self._start_lr = start_lr
-    self._trainer = None
 
     self._output_dir = output_dir
     gfile.makedirs(self._output_dir)
@@ -131,7 +128,7 @@ def trainer(self):
   def reset(self):
     self._current_lr = self._start_lr
     self._step = 0
-    self._trainer = self._trainer_fn(output_dir=self._next_trajectory_dir)
+    self._trainer.reset(output_dir=self._next_trajectory_dir)
     self._trainer.evaluate(self._eval_steps)
     return np.array([self._current_metric_value])
 
@@ -150,7 +147,6 @@ def step(self, action):
         environment steps. info is an empty dict.
     """
     self._current_lr *= self._action_multipliers[action]
-    self._trainer.update_learning_rate()
     last_metric_value = self._current_metric_value
     self._trainer.train_epoch(self._train_steps, self._eval_steps)
     self._step += 1
diff --git a/tensor2tensor/trax/rlax/envs/online_tune_env_test.py b/tensor2tensor/trax/rlax/envs/online_tune_env_test.py
index 20afa96f1..02f13b054 100644
--- a/tensor2tensor/trax/rlax/envs/online_tune_env_test.py
+++ b/tensor2tensor/trax/rlax/envs/online_tune_env_test.py
@@ -39,8 +39,13 @@ class MockTrainer(trax.Trainer):
   def __init__(self, metrics_to_report, *args, **kwargs):
     super(MockTrainer, self).__init__(*args, **kwargs)
     self.learning_rates = []
+    self.init_metrics_to_report = metrics_to_report
+    self.metrics_to_report = None
+
+  def reset(self, output_dir):
+    super(MockTrainer, self).reset(output_dir)
     # Copy the sequence to a list so we can modify it later.
-    self.metrics_to_report = list(metrics_to_report)
+    self.metrics_to_report = list(self.init_metrics_to_report)
 
   def train_epoch(self, epoch_steps, eval_steps):
     del epoch_steps
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index c0c28ae1a..e58ee558f 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -383,7 +383,7 @@ def _jit_update_fn(predict_fn, loss_fn, optimizer, n_devices, jit=True):
     def single_update(i, opt_state, batch, rng):
       params, slots, opt_params = opt_state
       rng, subrng = jax_random.split(rng[0])
-      grads = backend.grad(loss_fn)(opt_state.params, batch, predict_fn, rng)
+      grads = backend.grad(loss_fn)(params, batch, predict_fn, rng)
       return optimizer.tree_update(
           i, grads, params, slots, opt_params), [subrng]
     if jit:
@@ -397,7 +397,7 @@ def mapped_update(i, opt_state, batch, rng):
     # We assume all tensors have the first dimension = n_devices.
     params, slots, opt_params = opt_state
     rng, subrng = jax_random.split(rng)
-    grads = backend.grad(loss_fn)(opt_state.params, batch, predict_fn, rng)
+    grads = backend.grad(loss_fn)(params, batch, predict_fn, rng)
     grads = jax.tree_util.tree_map(
         lambda g: lax.psum(g, "batch"), grads)
     return optimizer.tree_update(
@@ -470,8 +470,9 @@ class Trainer(object):
   save the training state and access evaluation data.
   """
 
-  def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs, output_dir,
-               random_seed=None, n_devices=None, save_steps=None):
+  def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs,
+               output_dir=None, random_seed=None, n_devices=None,
+               save_steps=None):
     if save_steps is None:
       save_steps = []
     self._save_steps = save_steps
@@ -483,24 +484,13 @@ def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs, output_dir,
                        "%d != %d" % (n_devices, device_count))
     self._n_devices = n_devices
     rng = get_random_number_generator_and_set_seed(random_seed)
-    self._output_dir = output_dir
-    gfile.makedirs(output_dir)
-    # Create summary writers and history.
-    self._train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
-    self._eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
-
-    # Create input streams.
     inputs = inputs(n_devices)
     self._inputs = inputs
-    self._train_stream = inputs.train_stream()
 
-    # Setup optimizer and model.
-    state = restore_state(output_dir)
-    step = state.step or 0
-    history = state.history
-    self._lr_fn = lr_schedule(history)
-    opt = optimizer(learning_rate=self._lr_fn(step))
+    # Initialize the learning rate to a dummy value. It will be set in reset().
+    opt = optimizer(learning_rate=0.0)
 
+    # Setup the model.
     model_train = model(mode="train")
     model_predict_eval = model(mode="eval")
 
@@ -517,33 +507,71 @@ def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs, output_dir,
     # Change all None to 1 in input shape.
     model_input_shape = layers.nested_map(
         model_input_shape, lambda x: x if x else 1)
-    if state.opt_state:
-      opt_state = state.opt_state
-    else:
-      def initialize(input_shape, input_dtype, init_rng):
-        params = model_train.initialize(input_shape, input_dtype, init_rng)
-        (slots, opt_params) = opt.tree_init(params)
-        return OptState(params, slots, opt_params)
-      if _is_jit_init():
-        # JIT parameter initialization to avoid memory fragmentation
-        initialize = backend.jit(initialize, static_argnums=(0, 1))
-      opt_state = initialize(
-          model_input_shape, inputs.input_dtype, init_rng)
-    opt_state = OptState(*layers.nested_map(opt_state, self._maybe_replicate))
+    def initialize(input_shape, input_dtype, init_rng):
+      params = model_train.initialize(input_shape, input_dtype, init_rng)
+      (slots, opt_params) = opt.tree_init(params)
+      return OptState(params, slots, opt_params)
+    if _is_jit_init():
+      # JIT parameter initialization to avoid memory fragmentation
+      initialize = backend.jit(initialize, static_argnums=(0, 1))
+    self._initialize = lambda: initialize(  # pylint: disable=g-long-lambda
+        model_input_shape, self._inputs.input_dtype, init_rng)
 
     # jit model_predict and update so they're fast
     self._jit_model_predict_eval = _jit_predict_fn(
         model_predict_eval, n_devices)
     self._jit_update_fn = _jit_update_fn(model_train, loss_fn, opt, n_devices)
 
-    self._step = step
     self._model_train = model_train
     self._model_predict_eval = model_predict_eval
     self._loss_fn = loss_fn
-    self._opt_state = opt_state
-    self._history = history
     self._lr_schedule = lr_schedule
 
+    # Those fields will be set in reset().
+    self._output_dir = None
+    self._train_sw = None
+    self._eval_sw = None
+    self._history = None
+    self._lr_fn = None
+    self._opt_state = None
+    self._step = None
+
+    if output_dir is not None:
+      self.reset(output_dir)
+
+  def reset(self, output_dir):
+    """Reset the model parameters.
+
+    Restores the parameters from the given output_dir if a checkpoint exists,
+    otherwise randomly initializes them.
+
+    Does not re-jit the model.
+
+    Args:
+      output_dir: Output directory.
+    """
+    self._output_dir = output_dir
+    gfile.makedirs(output_dir)
+    # Create summary writers and history.
+    self._train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
+    self._eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
+
+    # Reset the training stream.
+    self._train_stream = self._inputs.train_stream()
+
+    # Restore the training state.
+    state = restore_state(output_dir)
+    self._step = state.step or 0
+    history = state.history
+    self._lr_fn = self._lr_schedule(history)
+    self._history = history
+    if state.opt_state:
+      opt_state = state.opt_state
+    else:
+      opt_state = self._initialize()
+    self._opt_state = OptState(*layers.nested_map(
+        opt_state, self._maybe_replicate))
+
     self.update_learning_rate()
 
   @property

From cd80370b4a7737066bad13b49f043a1011641268 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 24 Jul 2019 11:15:28 -0700
Subject: [PATCH 2228/2720] add librispeech_train_full_test_other problem def,
 make minor doc fixes.

PiperOrigin-RevId: 259778368
---
 tensor2tensor/data_generators/librispeech.py | 52 +++++++++++++++++++-
 1 file changed, 50 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index 1bab1d79a..9d2841f39 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -227,6 +227,54 @@ def filepattern(self, data_dir, mode, shard=None):
     return "%s-%s%s*" % (path, suffix, shard_str)
 
 
+@registry.register_problem()
+class LibrispeechTrainFullTestOther(Librispeech):
+  """Problem to train on full 960h, but evaluate on clean data only."""
+
+  def training_filepaths(self, data_dir, num_shards, shuffled):
+    return Librispeech.training_filepaths(self, data_dir, num_shards, shuffled)
+
+  def dev_filepaths(self, data_dir, num_shards, shuffled):
+    return LibrispeechNoisy.dev_filepaths(self, data_dir, num_shards, shuffled)
+
+  def test_filepaths(self, data_dir, num_shards, shuffled):
+    return LibrispeechNoisy.test_filepaths(self, data_dir, num_shards, shuffled)
+
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    raise Exception("Generate librispeech and librispeech_noisy data.")
+
+  def filepattern(self, data_dir, mode, shard=None):
+    """Get filepattern for data files for mode.
+
+    Matches mode to a suffix.
+    * DatasetSplit.TRAIN: train
+    * DatasetSplit.EVAL: dev
+    * DatasetSplit.TEST: test
+    * tf.estimator.ModeKeys.PREDICT: dev
+
+    Args:
+      data_dir: str, data directory.
+      mode: DatasetSplit
+      shard: int, if provided, will only read data from the specified shard.
+
+    Returns:
+      filepattern str
+    """
+    shard_str = "-%05d" % shard if shard is not None else ""
+    if mode == problem.DatasetSplit.TRAIN:
+      path = os.path.join(data_dir, "librispeech")
+      suffix = "train"
+    elif mode in [problem.DatasetSplit.EVAL, tf.estimator.ModeKeys.PREDICT]:
+      path = os.path.join(data_dir, "librispeech_noisy")
+      suffix = "dev"
+    else:
+      assert mode == problem.DatasetSplit.TEST
+      path = os.path.join(data_dir, "librispeech_noisy")
+      suffix = "test"
+
+    return "%s-%s%s*" % (path, suffix, shard_str)
+
+
 @registry.register_problem()
 class LibrispeechCleanSmall(Librispeech):
   """Problem spec for Librispeech using 100h clean train and clean eval data."""
@@ -249,9 +297,9 @@ class LibrispeechClean(Librispeech):
 
 @registry.register_problem()
 class LibrispeechNoisy(Librispeech):
-  """Problem spec for Librispeech using 400h noisy train and noisy eval data."""
+  """Problem spec for Librispeech using 500h noisy train and noisy eval data."""
 
-  # Select only the clean data
+  # Select only the noisy data
   TRAIN_DATASETS = _LIBRISPEECH_TRAIN_DATASETS[2:]
   DEV_DATASETS = _LIBRISPEECH_DEV_DATASETS[1:]
   TEST_DATASETS = _LIBRISPEECH_TEST_DATASETS[1:]

From 7632ed01e739cd124c8bac85f121f0f49ddd86cf Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 24 Jul 2019 13:15:02 -0700
Subject: [PATCH 2229/2720] Starting a colab to demonstrate Trax layers.

PiperOrigin-RevId: 259802336
---
 tensor2tensor/trax/layers/demo.ipynb | 189 +++++++++++++++++++++++++++
 1 file changed, 189 insertions(+)
 create mode 100644 tensor2tensor/trax/layers/demo.ipynb

diff --git a/tensor2tensor/trax/layers/demo.ipynb b/tensor2tensor/trax/layers/demo.ipynb
new file mode 100644
index 000000000..1e3d0303f
--- /dev/null
+++ b/tensor2tensor/trax/layers/demo.ipynb
@@ -0,0 +1,189 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "7yuytuIllsv1"
+      },
+      "source": [
+        "# Trax Layers\n",
+        "\n",
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "colab": {
+          "height": 119
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 88481,
+          "status": "ok",
+          "timestamp": 1563927238895,
+          "user": {
+            "displayName": "Lukasz Kaiser",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mC8pChl87HbK_eOtVhtNPwUVx8btvfyYzH9UHn3=s64",
+            "userId": "13267693649565518272"
+          },
+          "user_tz": 420
+        },
+        "id": "oILRLCWN_16u",
+        "outputId": "3f750014-c633-4162-ad07-f3c56c273304"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "T2T: skipped importing 1 data_generators modules. OK if no other errors. Depend on _heavy or problem-specific py_binary targets if trying to use a module that was skipped.\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "matplotlib.use() must be called *before* pylab, matplotlib.pyplot,\n",
+            "or matplotlib.backends is imported for the first time.\n",
+            "\n"
+          ]
+        }
+      ],
+      "source": [
+        "#@title Licence and python imports.\n",
+        "# Copyright 2018 Google LLC.\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",
+        "# https://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",
+        "import datetime\n",
+        "import numpy as onp\n",
+        "\n",
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 5,
+      "metadata": {
+        "cellView": "both",
+        "colab": {
+          "height": 51
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 2992,
+          "status": "ok",
+          "timestamp": 1563927313403,
+          "user": {
+            "displayName": "Lukasz Kaiser",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mC8pChl87HbK_eOtVhtNPwUVx8btvfyYzH9UHn3=s64",
+            "userId": "13267693649565518272"
+          },
+          "user_tz": 420
+        },
+        "id": "vlGjGoGMTt-D",
+        "outputId": "7a7b5a1e-c01e-4a5e-eeb1-a88d9500aad4"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "/bin/sh: pip: command not found\n",
+            "/bin/sh: pip: command not found\n"
+          ]
+        }
+      ],
+      "source": [
+        "#@title Install TensorFlow, Tensor2Tensor and Trax.\n",
+        "\n",
+        "! pip install -q tensorflow\n",
+        "! pip install -q -U tensor2tensor\n",
+        "\n",
+        "from tensor2tensor.trax import trax\n",
+        "from tensor2tensor.trax import backend\n",
+        "from tensor2tensor.trax import layers as tl"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 6,
+      "metadata": {
+        "colab": {
+          "height": 68
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 434,
+          "status": "ok",
+          "timestamp": 1563927323194,
+          "user": {
+            "displayName": "Lukasz Kaiser",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mC8pChl87HbK_eOtVhtNPwUVx8btvfyYzH9UHn3=s64",
+            "userId": "13267693649565518272"
+          },
+          "user_tz": 420
+        },
+        "id": "V09viOSEQvQe",
+        "outputId": "822ec44b-1e2e-4b6e-9fc7-9fc29f5d3783"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "[[-7 -6 -5 -4 -3]\n",
+            " [-2 -1  0  1  2]\n",
+            " [ 3  4  5  6  7]]\n",
+            "[[0 0 0 0 0]\n",
+            " [0 0 0 1 2]\n",
+            " [3 4 5 6 7]]\n"
+          ]
+        }
+      ],
+      "source": [
+        "x  = onp.arange(-7, 8).reshape(3, -1)\n",
+        "rng = backend.random.get_prng(0)\n",
+        "layer = tl.Relu()\n",
+        "params = layer.initialize(x.shape, x.dtype, rng)\n",
+        "output = layer(x, params, rng=rng)\n",
+        "print(x)\n",
+        "print(output)"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [],
+      "last_runtime": {
+        "build_target": "//learning/deepmind/dm_python:dm_notebook",
+        "kind": "private"
+      },
+      "name": "Trax Layers",
+      "provenance": [
+        {
+          "file_id": "1EH76AWQ_pvT4i8ZXfkv-SCV4MrmllEl5",
+          "timestamp": 1563927451951
+        }
+      ],
+      "version": "0.3.2"
+    },
+    "kernelspec": {
+      "display_name": "Python 2",
+      "name": "python2"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

From d470434f8ee584fb131170e2cb730a273d1c5a20 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 25 Jul 2019 10:08:32 -0700
Subject: [PATCH 2230/2720] No accelerators in lr calculation in Trax step (to
 avoid allocations).

PiperOrigin-RevId: 259968721
---
 tensor2tensor/trax/trax.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index e58ee558f..f97a9570e 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -610,7 +610,7 @@ def print_n_params(self):
   def _train_step(self, next_train_batch):
     """Run one training step and update self._opt_state."""
     # Calculate the current learning rate.
-    learning_rate = self._maybe_replicate(np.array(self._lr_fn(self._step)))
+    learning_rate = self._maybe_replicate(np.array(self.learning_rate))
     opt_state = self._opt_state
     opt_params = opt_state.opt_params
     opt_params = (learning_rate,) + opt_params[1:]

From 210282baf89580a081ef6952264fb1ea50d22300 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 25 Jul 2019 13:35:11 -0700
Subject: [PATCH 2231/2720] Update Adafactor optimizer to new signature.

PiperOrigin-RevId: 260010526
---
 tensor2tensor/trax/optimizers/base.py | 70 ++++++++++++++-------------
 1 file changed, 36 insertions(+), 34 deletions(-)

diff --git a/tensor2tensor/trax/optimizers/base.py b/tensor2tensor/trax/optimizers/base.py
index 41c113534..ec2b15a43 100644
--- a/tensor2tensor/trax/optimizers/base.py
+++ b/tensor2tensor/trax/optimizers/base.py
@@ -240,8 +240,9 @@ def update(self, step, grads, params, slots, opt_params):
 class Adafactor(Optimizer):
   """Adafactor optimizer."""
 
+  # TODO(levskaya): refactor to use newer RL friendly parameter passing.
   def __init__(self,
-               step_size,
+               learning_rate,
                decay_rate=0.8,
                beta1=0.0,
                clipping_threshold=1.0,
@@ -254,7 +255,7 @@ def __init__(self,
     Adafactor is described in https://arxiv.org/abs/1804.04235.
 
     Args:
-      step_size: function i -> float, trax-provided learning rate schedule.
+      learning_rate: float: trax-provided learning rate.
       decay_rate: float: controls second-moment exponential decay schedule.
       beta1: a float value between 0 and 1, enables momentum and uses extra
         memory if nonzero!  Off by default.
@@ -262,19 +263,18 @@ def __init__(self,
       factored: boolean: whether to use factored second-moment estimator for 2d
         variables.
       multiply_by_parameter_scale: boolean: if True, then scale provided
-        step_size by parameter norm. if False, provided step_size is absolute
-        step size.
+        learning_rate by parameter norm. if False, provided learning_rate is
+        absolute step size.
       epsilon1: Regularization constant for squared gradient.
       epsilon2: Regularization constant for parameter scale.
     """
-    super(Adafactor, self).__init__(step_size)
+    super(Adafactor, self).__init__(learning_rate)
     self._multiply_by_parameter_scale = multiply_by_parameter_scale
+    self._factored = factored
     self._beta1 = beta1
     self._clipping_threshold = clipping_threshold
-    self._factored = factored
     self._epsilon1 = epsilon1
     self._epsilon2 = epsilon2
-    self._step_size = step_size
     self._decay_rate = functools.partial(self._decay_rate_pow,
                                          exponent=decay_rate)
 
@@ -284,49 +284,51 @@ def _decay_rate_pow(i, exponent=0.8):
     t = np.array(i, np.float32) + 1.0
     return 1.0 - t**(-exponent)
 
-  def init(self, x):
-    shape = x.shape
-    state = []
+  def init(self, params):
+    shape = params.shape
+    slots = []
     if self._factored and len(shape) >= 2:
       v_row = np.zeros(shape[:-1], dtype=np.float32)
       v_col = np.zeros(shape[:-2] + shape[-1:], dtype=np.float32)
-      state.extend([v_row, v_col])
+      slots.extend([v_row, v_col])
     else:
-      v = np.zeros_like(x)
-      state.append(v)
+      v = np.zeros_like(params)
+      slots.append(v)
     if self._beta1:
-      m = np.zeros_like(x)
-      state.append(m)
-    return state
+      m = np.zeros_like(params)
+      slots.append(m)
+    return slots
 
-  def update(self, i, g, x, state):
+  def update(self, step, grads, params, slots, opt_params):
     updates = []
-    decay_rate = self._decay_rate(i)
-    update_scale = self._step_size(i)
+    (learning_rate,) = opt_params
+    decay_rate = self._decay_rate(step)
+    update_scale = learning_rate
     if self._multiply_by_parameter_scale:
-      update_scale *= np.maximum(np.sqrt(np.mean(x * x)), self._epsilon2)
+      update_scale *= np.maximum(
+          np.sqrt(np.mean(params * params)), self._epsilon2)
     mixing_rate = 1.0 - decay_rate
 
-    g_sqr = g * g + self._epsilon1
-    if self._factored and len(x.shape) >= 2:
-      v_row = state.pop(0)
-      v_col = state.pop(0)
-      new_v_row = decay_rate * v_row + mixing_rate * np.mean(g_sqr, axis=-1)
-      new_v_col = decay_rate * v_col + mixing_rate * np.mean(g_sqr, axis=-2)
+    grads_sqr = grads * grads + self._epsilon1
+    if self._factored and len(params.shape) >= 2:
+      v_row = slots.pop(0)
+      v_col = slots.pop(0)
+      new_v_row = decay_rate * v_row + mixing_rate * np.mean(grads_sqr, axis=-1)
+      new_v_col = decay_rate * v_col + mixing_rate * np.mean(grads_sqr, axis=-2)
       updates.extend([new_v_row, new_v_col])
       row_col_mean = np.mean(new_v_row, axis=-1, keepdims=True)
       row_factor = (new_v_row / row_col_mean)**-0.5
       col_factor = (new_v_col)**-0.5
       y = (
-          g * np.expand_dims(row_factor, axis=-1) *
+          grads * np.expand_dims(row_factor, axis=-1) *
           np.expand_dims(col_factor, axis=-2))
     else:
-      v = state.pop(0)
-      new_v = decay_rate * v + mixing_rate * g_sqr
+      v = slots.pop(0)
+      new_v = decay_rate * v + mixing_rate * grads_sqr
       updates.append(new_v)
-      y = g * (new_v)**-0.5
+      y = grads * (new_v)**-0.5
 
-    if self._clipping_threshold is not None:
+    if self._clipping_threshold:
       clipping_denom = (
           np.maximum(1.0,
                      np.sqrt(np.mean(y * y)) / self._clipping_threshold))
@@ -334,13 +336,13 @@ def update(self, i, g, x, state):
 
     subtrahend = update_scale * y
     if self._beta1:
-      m = state.pop(0)
+      m = slots.pop(0)
       new_m = self._beta1 * m + (1.0 - self._beta1) * subtrahend
       subtrahend = new_m
       updates.append(new_m)
 
-    new_x = x - subtrahend
-    return new_x, updates
+    new_params = params - subtrahend
+    return new_params, updates
 
 
 class SM3(Optimizer):

From 223c78434d3155a3824546532219cfa30947519e Mon Sep 17 00:00:00 2001
From: Peng Wang <wangpeng@google.com>
Date: Thu, 25 Jul 2019 16:35:50 -0700
Subject: [PATCH 2232/2720] Some performance improvements for tf-numpy

PiperOrigin-RevId: 260045417
---
 tensor2tensor/trax/inputs.py  | 22 ++++++++++++----------
 tensor2tensor/trax/trainer.py | 12 ++++++++++++
 tensor2tensor/trax/trax.py    |  3 +++
 3 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 381954de7..86ed644ee 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -162,16 +162,18 @@ def dataset_to_stream(dataset, input_name, n_chunks=0, append_targets=False):
   """Takes a tf.Dataset and creates a numpy stream of ready batches."""
   for example in backend.dataset_as_numpy(dataset):
     inp, out = example[0][input_name], example[1]
-    # Some accelerators don't handle uint8 well, cast to int.
-    if isinstance(inp, np.uint8):
-      inp = inp.astype(np.int32)
-    if isinstance(out, np.uint8):
-      out = out.astype(np.int32)
-    if len(out.shape) > 1 and out.shape[-1] == 1:
-      out = np.squeeze(out, axis=-1)
-    if n_chunks > 0:
-      inp = tuple(np.split(inp, n_chunks, axis=1))
-      out = tuple(np.split(out, n_chunks, axis=1))
+    # All input-pipeline processing should be on CPU.
+    with tf.device("cpu:0"):
+      # Some accelerators don't handle uint8 well, cast to int.
+      if isinstance(inp, np.uint8):
+        inp = inp.astype(np.int32)
+      if isinstance(out, np.uint8):
+        out = out.astype(np.int32)
+      if len(out.shape) > 1 and out.shape[-1] == 1:
+        out = np.squeeze(out, axis=-1)
+      if n_chunks > 0:
+        inp = tuple(np.split(inp, n_chunks, axis=1))
+        out = tuple(np.split(out, n_chunks, axis=1))
     if append_targets:
       inp = (inp, out)
     yield inp, out
diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index edf3a1f43..ed29ef542 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -48,6 +48,10 @@
 flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
 flags.DEFINE_bool("tf_eager", False, "Whether we're running TF in eager mode.")
 flags.DEFINE_bool("tf_xla", False, "Whether to turn on XLA for TF.")
+flags.DEFINE_bool("tf_opt_pin_to_host", False, "Whether to turn on TF "
+                  "pin-to-host optimization.")
+flags.DEFINE_bool("tf_opt_layout", False, "Whether to turn on TF layout "
+                  "optimization.")
 
 
 def _default_output_dir():
@@ -95,6 +99,14 @@ def main(_):
   if FLAGS.tf_xla:
     tf.config.optimizer.set_jit(True)
 
+  tf.config.optimizer.set_experimental_options(
+      {"pin_to_host_optimization": FLAGS.tf_opt_pin_to_host}
+  )
+
+  tf.config.optimizer.set_experimental_options(
+      {"layout_optimizer": FLAGS.tf_opt_layout}
+  )
+
   _setup_gin()
 
   # Setup output directory
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index f97a9570e..f5574df05 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -571,6 +571,9 @@ def reset(self, output_dir):
       opt_state = self._initialize()
     self._opt_state = OptState(*layers.nested_map(
         opt_state, self._maybe_replicate))
+    if not state.opt_state:
+      _save_replicated(self._opt_state, self._step, self._history,
+                       self._n_devices, self._output_dir, False)
 
     self.update_learning_rate()
 

From bc6a8cf43c2ad7e290da12646eed75ba5b061b6a Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 26 Jul 2019 13:55:17 -0700
Subject: [PATCH 2233/2720] [T2T] Fixes a bug that resulted in incorrect TPU
 summaries on tensorboard.

Description: At the moment, training and eval summaries are written out to the same directory. This can manifest itself in a number of ways:

1) Tensorboard incorrectly truncates training summaries if the corresponding summary is not present in eval.
2) Tensorboard mixes the losses from train and eval and labels them both as 'train'. This can cause the training loss to appear jagged.

PiperOrigin-RevId: 260207502
---
 tensor2tensor/utils/t2t_model.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 27405470e..402116cdf 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1546,7 +1546,10 @@ def create_train_host_call(self):
     return create_host_call(self.hparams.model_dir)
 
   def create_eval_host_call(self):
-    return self.create_train_host_call()
+    eval_dir = os.path.join(
+        self.hparams.model_dir,
+        self.hparams.get("eval_dir_name", "eval"))
+    return create_host_call(eval_dir)
 
   def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
     """Constructs `tf.estimator.EstimatorSpec` for TRAIN (training) mode."""

From 80559082f0eed6ef1a9bcd5792243350e476455a Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 27 Jul 2019 07:47:42 -0700
Subject: [PATCH 2234/2720] gRPC server env and client for T2T

- env_service.proto : Contains the service (EnvService) and proto definitions.
- env_service_servicer.py : Contains the implementation of the service.
- env_service_server.py : Contains the server that serves these requests.
- client_env.py : A gym Env that talks to a server that implements EnvService.
- env_service_client.py : A debug utility that connects to a remote cluster of envs.
- utils.py : Utilities for proto conversions, gym spaces, envs.

* Needs to be done:
- Put the client on a cluster hosted GPU too.
- Server to run OnlineTunerEnv and run PPO on it.
- More tests being added async.
- For OSS to work *without* proto compiler, we'll need to add the python files generated from the .proto file to the source code.

PiperOrigin-RevId: 260298971
---
 tensor2tensor/envs/__init__.py                |  16 ++
 tensor2tensor/envs/client_env.py              | 137 ++++++++++++++++++
 tensor2tensor/envs/env_problem.py             |   7 +-
 tensor2tensor/envs/env_problem_utils.py       |  60 +++++++-
 tensor2tensor/envs/env_service.proto          | 118 +++++++++++++++
 tensor2tensor/envs/env_service_client.py      |  63 ++++++++
 .../envs/env_service_serialization.py         |  89 ++++++++++++
 .../envs/env_service_serialization_test.py    |  95 ++++++++++++
 tensor2tensor/envs/env_service_server.py      |  95 ++++++++++++
 tensor2tensor/envs/env_service_servicer.py    | 111 ++++++++++++++
 .../envs/env_service_servicer_test.py         |  95 ++++++++++++
 tensor2tensor/envs/gym_env_problem.py         |  15 +-
 tensor2tensor/envs/rendered_env_problem.py    |  10 +-
 13 files changed, 904 insertions(+), 7 deletions(-)
 create mode 100644 tensor2tensor/envs/client_env.py
 create mode 100644 tensor2tensor/envs/env_service.proto
 create mode 100644 tensor2tensor/envs/env_service_client.py
 create mode 100644 tensor2tensor/envs/env_service_serialization.py
 create mode 100644 tensor2tensor/envs/env_service_serialization_test.py
 create mode 100644 tensor2tensor/envs/env_service_server.py
 create mode 100644 tensor2tensor/envs/env_service_servicer.py
 create mode 100644 tensor2tensor/envs/env_service_servicer_test.py

diff --git a/tensor2tensor/envs/__init__.py b/tensor2tensor/envs/__init__.py
index 8c036a176..9b78dafac 100644
--- a/tensor2tensor/envs/__init__.py
+++ b/tensor2tensor/envs/__init__.py
@@ -19,6 +19,22 @@
 from __future__ import division
 from __future__ import print_function
 
+from gym.envs.registration import register
+
+from tensor2tensor.envs import client_env
 from tensor2tensor.envs import gym_env_problem
 from tensor2tensor.envs import tic_tac_toe_env
 from tensor2tensor.envs import tic_tac_toe_env_problem
+
+
+def register_env(env_class):
+  register(
+      id="{}-v0".format(env_class.__name__),
+      entry_point="tensor2tensor.envs:{}".format(env_class.__name__),
+  )
+  return env_class
+
+
+# TODO(afrozm): Register TicTacToeEnv the same way.
+# register_env(tic_tac_toe_env.TicTacToeEnv)
+ClientEnv = register_env(client_env.ClientEnv)  # pylint: disable=invalid-name
diff --git a/tensor2tensor/envs/client_env.py b/tensor2tensor/envs/client_env.py
new file mode 100644
index 000000000..f4a2a1b0b
--- /dev/null
+++ b/tensor2tensor/envs/client_env.py
@@ -0,0 +1,137 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Client Env that connects to a distributed env."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl import logging
+import grpc
+from grpc import loas2
+import gym
+import numpy as np
+from tensor2tensor.envs import env_service_pb2
+from tensor2tensor.envs import env_service_pb2_grpc
+from tensor2tensor.envs import env_service_serialization as serialization
+
+
+class ClientEnv(gym.Env):
+  """Creates a connection to a remote env, and calls RPC methods on it."""
+
+  @staticmethod
+  def create_channel(remote_env_address):
+    return grpc.secure_channel(remote_env_address,
+                               loas2.loas2_channel_credentials())
+
+  @staticmethod
+  def run_step(stub, discrete_action):
+    action_proto = env_service_pb2.Action(discrete_action=discrete_action)
+    step_request = env_service_pb2.StepRequest()
+    step_request.action.CopyFrom(action_proto)
+    return stub.Step(step_request)
+
+  @staticmethod
+  def run_reset(stub):
+    return stub.Reset(env_service_pb2.ResetRequest())
+
+  @staticmethod
+  def run_close(stub, channel):
+    close_response = stub.Close(env_service_pb2.CloseRequest())
+    channel.close()
+    return close_response
+
+  @staticmethod
+  def run_render(stub, mode="rgb_array"):
+    return stub.Render(env_service_pb2.RenderRequest(mode=mode))
+
+  @staticmethod
+  def run_get_env_info(stub):
+    env_info_response = stub.GetEnvInfo(env_service_pb2.EnvInfoRequest())
+    gym_observation_space = serialization.proto_to_gym_space(
+        env_info_response.observation_space)
+    gym_action_space = serialization.proto_to_gym_space(
+        env_info_response.action_space)
+    reward_range = (env_info_response.reward_range.low,
+                    env_info_response.reward_range.high)
+    return (gym_action_space, gym_observation_space, reward_range,
+            env_info_response.batch_size)
+
+  def __init__(self, remote_env_address=None):
+    logging.vlog(1, "Making a ClientEnv with remote address: [%s]",
+                 remote_env_address)
+    assert remote_env_address is not None
+    # Make a channel and stub on the remote env address.
+    self._remote_env_address = remote_env_address
+
+    self._channel = None
+    self._stub = None
+    self.initialize_stub()
+    assert self._stub is not None
+
+    # We now have to do an RPC to determine spaces and reward range.
+    #
+    # NOTE: If all these are same across replicas, then we technically only need
+    # to do this once on the 'master' replica (say 0), but `GymEnvProblem`
+    # checks that they are all the same.
+    (self.action_space, self.observation_space, self.reward_range,
+     self._server_env_batch_size) = (
+         ClientEnv.run_get_env_info(self._stub))
+
+  def initialize_stub(self):
+    self._channel = ClientEnv.create_channel(self._remote_env_address)
+    # TODO(afrozm): Why is this done?
+    grpc.channel_ready_future(self._channel).result()
+    self._stub = env_service_pb2_grpc.EnvServiceStub(self._channel)
+
+  def _maybe_squeeze_array(self, np_array):
+    # Usually this client is talking to a server env that is running a single
+    # element batch, if so, this client should strip out the batch dimension
+    # before reporting the observation upstream (since this is a plain gym env,
+    # not an EnvProblem), the upstream EnvProblem will then batch across
+    # multiple ClientEnvs.
+    if isinstance(
+        np_array, np.ndarray
+    ) and self._server_env_batch_size == 1 and np_array.shape[0] == 1:
+      np_array = np.squeeze(np_array, axis=0)
+    return np_array
+
+  def reset(self):
+    # Run the RPC.
+    reset_response_proto = ClientEnv.run_reset(self._stub)
+    # Convert the TensorProto to numpy.
+    obs_np = serialization.tensor_proto_to_numpy_array(
+        reset_response_proto.observation.observation)
+    return self._maybe_squeeze_array(obs_np)
+
+  def close(self):
+    ClientEnv.run_close(self._stub, self._channel)
+
+  def render(self, mode="rgb_array"):
+    render_response = ClientEnv.run_render(self._stub, mode=mode)
+    if not render_response:
+      return
+    # Parse out the numpy array.
+    return serialization.tensor_proto_to_numpy_array(
+        render_response.observation.observation)
+
+  def step(self, action):
+    step_response = ClientEnv.run_step(self._stub, action)
+    observation = self._maybe_squeeze_array(
+        serialization.tensor_proto_to_numpy_array(
+            step_response.observation.observation))
+    info = {k: v for k, v in step_response.info.info_map.items()}
+    return observation, step_response.reward, step_response.done, info
diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index e65bf6405..909307c0f 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -126,14 +126,17 @@ def __init__(self,
     # If set, we discretize the rewards and treat them as integers.
     self._discrete_rewards = discrete_rewards
 
-    self._parallelism = None
-
     # A data structure to hold the `batch_size` currently active trajectories
     # and also the ones that are completed, i.e. done.
     self._trajectories = None
 
     self._batch_size = None
 
+    self._parallelism = None
+    # The parallelism is passes in via env_kwargs because it will be used by
+    # `GymEnvProblem` to paralellize env actions across a batch.
+    env_kwargs["parallelism"] = parallelism
+
     if batch_size is not None:
       self.initialize(batch_size=batch_size, **env_kwargs)
 
diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 69ec1d945..65d074c56 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -19,9 +19,14 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
 import time
 import numpy as np
 
+from tensor2tensor.envs import gym_env_problem
+from tensor2tensor.envs import rendered_env_problem
+from tensor2tensor.rl import gym_utils
+
 EPSILON_GREEDY = "epsilon-greedy"
 GUMBEL_SAMPLING = "gumbel"
 
@@ -162,8 +167,11 @@ def epsilon_greedy(log_probs):
     # Step through the env.
     t1 = time.time()
     _, _, dones, env_infos = env.step(
-        actions, infos={"log_prob_actions": log_probs,
-                        "value_predictions": value_preds})
+        actions,
+        infos={
+            "log_prob_actions": log_probs,
+            "value_predictions": value_preds
+        })
     env_actions_total_time += (time.time() - t1)
     bare_env_run_time += sum(
         info["__bare_env_run_time__"] for info in env_infos)
@@ -210,3 +218,51 @@ def epsilon_greedy(log_probs):
   timing_info = {k: round(1000 * v, 2) for k, v in timing_info.items()}
 
   return completed_trajectories, num_done_trajectories, timing_info
+
+
+def make_env(batch_size=1,
+             env_problem_name="",
+             resize=True,
+             resized_height=105,
+             resized_width=80,
+             max_timestep="None",
+             clip_rewards=True,
+             parallelism=1,
+             use_tpu=False,
+             **env_kwargs):
+  """Creates the env."""
+
+  if clip_rewards:
+    env_kwargs.update({"reward_range": (-1, 1), "discrete_rewards": True})
+  else:
+    env_kwargs.update({"discrete_rewards": False})
+
+  # No resizing needed, so let's be on the normal EnvProblem.
+  if not resize:  # None or False
+    return gym_env_problem.GymEnvProblem(
+        base_env_name=env_problem_name,
+        batch_size=batch_size,
+        parallelism=parallelism,
+        **env_kwargs)
+
+  try:
+    max_timestep = int(max_timestep)
+  except Exception:  # pylint: disable=broad-except
+    max_timestep = None
+
+  wrapper_fn = functools.partial(
+      gym_utils.gym_env_wrapper, **{
+          "rl_env_max_episode_steps": max_timestep,
+          "maxskip_env": True,
+          "rendered_env": True,
+          "rendered_env_resize_to": (resized_height, resized_width),
+          "sticky_actions": False,
+          "output_dtype": np.int32 if use_tpu else None,
+      })
+
+  return rendered_env_problem.RenderedEnvProblem(
+      base_env_name=env_problem_name,
+      batch_size=batch_size,
+      parallelism=parallelism,
+      env_wrapper_fn=wrapper_fn,
+      **env_kwargs)
diff --git a/tensor2tensor/envs/env_service.proto b/tensor2tensor/envs/env_service.proto
new file mode 100644
index 000000000..e433204dc
--- /dev/null
+++ b/tensor2tensor/envs/env_service.proto
@@ -0,0 +1,118 @@
+syntax = "proto3";
+option cc_enable_arenas = true;
+
+package third_party.py.tensor2tensor.trax.rlax.envs;
+
+import "third_party/tensorflow/core/framework/tensor.proto";
+import "third_party/tensorflow/core/framework/tensor_shape.proto";
+import "third_party/tensorflow/core/framework/types.proto";
+
+// We use tensorflow.TensorProto to represent numpy arrays.
+
+message Action {
+  oneof payload {
+    int64 discrete_action = 1;
+    tensorflow.TensorProto continuous_action = 2;
+  }
+}
+
+message Observation {
+  tensorflow.TensorProto observation = 1;
+}
+
+message Info {
+  map<string, double> info_map = 1;
+}
+
+message StepRequest {
+  Action action = 1;
+}
+
+message StepResponse {
+  Observation observation = 1;
+  double reward = 2;
+  bool done = 3;
+  Info info = 4;
+}
+
+message ResetRequest {}
+
+message ResetResponse {
+  Observation observation = 1;
+}
+
+message CloseRequest {}
+message CloseResponse {}
+
+message RenderRequest {
+  string mode = 1;
+}
+
+message RenderResponse {
+  Observation observation = 1;
+}
+
+message EnvInfoRequest {}
+
+message SpaceBox {
+  tensorflow.DataType dtype = 1;
+  tensorflow.TensorShapeProto shape = 2;
+  tensorflow.TensorProto low = 3;
+  tensorflow.TensorProto high = 4;
+}
+
+message SpaceDiscrete {
+  int32 num_actions = 1;
+}
+
+message GymSpace {
+  // TODO(afrozm): Move spaces can be added as needed, or composite spaces.
+  oneof gym_space {
+    bool unimplemented_space = 1;
+    SpaceBox box = 2;
+    SpaceDiscrete discrete = 3;
+  }
+}
+
+message RewardRange {
+  double low = 1;
+  double high = 2;
+}
+
+message EnvInfoResponse {
+  GymSpace observation_space = 1;
+  GymSpace action_space = 2;
+  RewardRange reward_range = 3;
+  int64 batch_size = 4;
+}
+
+service EnvService {
+  // Reset
+  rpc Reset(ResetRequest) returns (ResetResponse) {
+    option fail_fast = true;
+  }
+
+  // Step
+  rpc Step(StepRequest) returns (StepResponse) {
+    option fail_fast = true;
+  }
+
+  // Close
+  rpc Close(CloseRequest) returns (CloseResponse) {
+    option fail_fast = true;
+    option deadline = 10;
+  }
+
+  // Render
+  rpc Render(RenderRequest) returns (RenderResponse) {
+    option fail_fast = true;
+    option deadline = 10;
+  }
+
+  // Observation and Action Space.
+  rpc GetEnvInfo(EnvInfoRequest) returns (EnvInfoResponse) {
+    option fail_fast = true;
+    option deadline = 10;
+  }
+}
+
diff --git a/tensor2tensor/envs/env_service_client.py b/tensor2tensor/envs/env_service_client.py
new file mode 100644
index 000000000..b49e86db2
--- /dev/null
+++ b/tensor2tensor/envs/env_service_client.py
@@ -0,0 +1,63 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 client binary that talks to remote envs, for debugging."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import pdb
+from absl import app
+from absl import flags
+import numpy as np  # pylint: disable=unused-import
+from tensor2tensor import envs  # pylint: disable=unused-import
+from tensor2tensor.envs import client_env
+from tensor2tensor.envs import env_problem_utils
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("server_bns", "", "Server's BNS.")
+flags.DEFINE_integer("replicas", 0, "Number of replicas in the server.")
+
+
+def main(argv):
+  del argv
+
+  if FLAGS.replicas == 0:
+    env = client_env.ClientEnv(FLAGS.server_bns)
+    pdb.set_trace()
+    env.close()
+    return
+
+  # Replicated server.
+  per_env_kwargs = [{
+      "remote_env_address": os.path.join(FLAGS.server_bns, str(replica))
+  } for replica in range(FLAGS.replicas)]
+  env = env_problem_utils.make_env(
+      batch_size=FLAGS.replicas,
+      env_problem_name="ClientEnv-v0",
+      resize=False,
+      parallelism=FLAGS.replicas,
+      per_env_kwargs=per_env_kwargs)
+
+  pdb.set_trace()
+
+  env.close()
+
+
+if __name__ == "__main__":
+  app.run(main)
diff --git a/tensor2tensor/envs/env_service_serialization.py b/tensor2tensor/envs/env_service_serialization.py
new file mode 100644
index 000000000..51c3f8226
--- /dev/null
+++ b/tensor2tensor/envs/env_service_serialization.py
@@ -0,0 +1,89 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Utilities for serializing numpy arrays, gym spaces and envs."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gym
+import numpy as np
+from tensor2tensor.envs import env_service_pb2
+from tensorflow.python.framework import tensor_util  # pylint: disable=g-direct-tensorflow-import
+
+
+def numpy_array_to_observation(array):
+  obs = env_service_pb2.Observation()
+  obs.observation.CopyFrom(tensor_util.make_tensor_proto(array))
+  return obs
+
+
+def tensor_proto_to_numpy_array(tensor_proto):
+  return tensor_util.MakeNdarray(tensor_proto)
+
+
+def step_request_from_discrete_action(action):
+  action_proto = env_service_pb2.Action(discrete_action=action)
+  step_request = env_service_pb2.StepRequest()
+  step_request.action.CopyFrom(action_proto)
+  return step_request
+
+
+def gym_space_to_proto(gym_space):
+  """Converts a gym space to `env_service_pb2.GymSpace`."""
+
+  if isinstance(gym_space, gym.spaces.Discrete):
+    return env_service_pb2.GymSpace(
+        discrete=env_service_pb2.SpaceDiscrete(num_actions=gym_space.n))
+  elif isinstance(gym_space, gym.spaces.Box):
+    space_proto = env_service_pb2.GymSpace()
+    box_proto = space_proto.box
+
+    # Set low & high first, we can set shape and type from it later.
+    box_proto.low.CopyFrom(tensor_util.make_tensor_proto(gym_space.low))
+    box_proto.high.CopyFrom(tensor_util.make_tensor_proto(gym_space.high))
+
+    # dtype and shape.
+    box_proto.dtype = box_proto.low.dtype
+    box_proto.shape.CopyFrom(box_proto.low.tensor_shape)
+
+    return space_proto
+
+  # A space that we haven't implemented.
+  return env_service_pb2.GymSpace(unimplemented_space=True)
+
+
+def proto_to_gym_space(gym_space_proto):
+  """Converts a `env_service_pb2.GymSpace` to a `gym.spaces`."""
+
+  if gym_space_proto.unimplemented_space:
+    return None
+
+  if gym_space_proto.HasField("discrete"):
+    return gym.spaces.Discrete(gym_space_proto.discrete.num_actions)
+
+  assert gym_space_proto.HasField("box")
+
+  low_np = tensor_proto_to_numpy_array(gym_space_proto.box.low)
+  high_np = tensor_proto_to_numpy_array(gym_space_proto.box.high)
+
+  return gym.spaces.Box(low=low_np, high=high_np, dtype=low_np.dtype)
+
+
+def reward_range_to_proto(reward_range=None):
+  if reward_range is None:
+    reward_range = (-np.inf, np.inf)
+  return env_service_pb2.RewardRange(low=reward_range[0], high=reward_range[1])
diff --git a/tensor2tensor/envs/env_service_serialization_test.py b/tensor2tensor/envs/env_service_serialization_test.py
new file mode 100644
index 000000000..92b3d6feb
--- /dev/null
+++ b/tensor2tensor/envs/env_service_serialization_test.py
@@ -0,0 +1,95 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.envs.env_service_serialization."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gym
+import numpy as np
+
+from tensor2tensor.envs import env_service_serialization as utils
+
+from tensorflow import test
+from tensorflow.core.framework import types_pb2  # pylint: disable=g-direct-tensorflow-import
+
+
+class UtilsTest(test.TestCase):
+
+  def test_conversion(self):
+    np_a = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
+    obs = utils.numpy_array_to_observation(np_a)
+
+    tp_a = obs.observation
+    np_tp_a = utils.tensor_proto_to_numpy_array(tp_a)
+
+    np.testing.assert_array_equal(np_a, np_tp_a)
+
+  def test_step_request_from_discrete_action(self):
+    discrete_action = 6
+    step_request = utils.step_request_from_discrete_action(discrete_action)
+    action_request = step_request.action
+    self.assertTrue(action_request.HasField("discrete_action"))
+    self.assertEqual("discrete_action", action_request.WhichOneof("payload"))
+    self.assertEqual(discrete_action, action_request.discrete_action)
+
+  def test_gym_space_to_proto_discrete(self):
+    num_actions = 77
+    space = gym.spaces.Discrete(num_actions)
+    space_proto = utils.gym_space_to_proto(space)
+
+    self.assertFalse(space_proto.HasField("box"))
+    self.assertTrue(space_proto.HasField("discrete"))
+    self.assertEqual(num_actions, space_proto.discrete.num_actions)
+
+  def test_gym_space_to_proto_box(self):
+    space = gym.spaces.Box(low=0, high=255, shape=(28, 29, 3), dtype=np.uint8)
+    space_proto = utils.gym_space_to_proto(space)
+
+    self.assertTrue(space_proto.HasField("box"))
+    self.assertEqual(types_pb2.DT_UINT8, space_proto.box.dtype)
+
+    self.assertEqual(28, space_proto.box.shape.dim[0].size)
+    self.assertEqual(29, space_proto.box.shape.dim[1].size)
+    self.assertEqual(3, space_proto.box.shape.dim[2].size)
+
+  def test_proto_to_gym_space_discrete(self):
+    num_actions = 77
+    space = gym.spaces.Discrete(num_actions)
+    space_proto = utils.gym_space_to_proto(space)
+    space_gym = utils.proto_to_gym_space(space_proto)
+    space_gym.n = num_actions
+
+  def test_proto_to_gym_space_box(self):
+    space = gym.spaces.Box(low=-1.0, high=1.0, shape=(28, 29), dtype=np.float32)
+    space_proto = utils.gym_space_to_proto(space)
+    space_gym = utils.proto_to_gym_space(space_proto)
+    self.assertEqual(np.float32, space_gym.dtype)
+    self.assertAllEqual(space.shape, space_gym.shape)
+
+  def test_reward_range_to_proto(self):
+    reward_proto = utils.reward_range_to_proto((-12, +13))
+    self.assertEqual(-12, reward_proto.low)
+    self.assertEqual(+13, reward_proto.high)
+
+    reward_proto = utils.reward_range_to_proto()
+    self.assertEqual(-np.inf, reward_proto.low)
+    self.assertEqual(np.inf, reward_proto.high)
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensor2tensor/envs/env_service_server.py b/tensor2tensor/envs/env_service_server.py
new file mode 100644
index 000000000..dd41c02c5
--- /dev/null
+++ b/tensor2tensor/envs/env_service_server.py
@@ -0,0 +1,95 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Server that acts as a remote env."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import time
+from absl import app
+from absl import flags
+from absl import logging
+from concurrent import futures
+import grpc
+from grpc import loas2
+from tensor2tensor.envs import env_problem_utils
+from tensor2tensor.envs import env_service_pb2_grpc
+from tensor2tensor.envs import env_service_servicer
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_bool("xm", False, "Copy atari roms?")
+flags.DEFINE_integer("env_service_port", 7777, "Port on which to run.")
+flags.DEFINE_string("env_problem_name", None, "Name of the EnvProblem to make.")
+flags.DEFINE_string(
+    "max_timestep", None,
+    "If set to an integer, maximum number of time-steps in a "
+    "trajectory. The bare env is wrapped with TimeLimit wrapper.")
+flags.DEFINE_boolean("resize", False, "If true, resize the game frame")
+flags.DEFINE_integer("resized_height", 105, "Resized height of the game frame.")
+flags.DEFINE_integer("resized_width", 80, "Resized width of the game frame.")
+flags.DEFINE_string("output_dir", "", "Output dir.")
+flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
+flags.DEFINE_integer("replica", 0, "Basically to append to output_dir")
+flags.DEFINE_bool("clip_rewards", True,
+                  "Whether to clip and discretize the rewards.")
+
+# Since we're only dealing with 1 GPU machines here.
+_MAX_CONCURRENCY = 1
+_ADDRESS_FORMAT = "[::]:{}"
+
+
+def add_port(server):
+  server_credentials = loas2.loas2_server_credentials()
+  return server.add_secure_port(_ADDRESS_FORMAT.format(FLAGS.env_service_port),
+                                server_credentials)
+
+
+def serve(output_dir, env):
+  del output_dir  # may use later.
+  server = grpc.server(futures.ThreadPoolExecutor(max_workers=_MAX_CONCURRENCY))
+  servicer = env_service_servicer.EnvServiceServicer(env)
+  env_service_pb2_grpc.add_EnvServiceServicer_to_server(servicer, server)
+  port = add_port(server)
+  server.start()
+  logging.info("Starting server on port %s", port)
+  while True:
+    time.sleep(60 * 60 * 24)  # sleep for a day only to sleep again.
+
+
+def main(argv):
+  del argv
+  output_dir = FLAGS.output_dir
+
+  output_dir = os.path.join(output_dir, str(FLAGS.replica))
+
+  env = env_problem_utils.make_env(
+      batch_size=1,
+      env_problem_name=FLAGS.env_problem_name,
+      resize=FLAGS.resize,
+      resized_height=FLAGS.resized_height,
+      resized_width=FLAGS.resized_width,
+      max_timestep=FLAGS.max_timestep,
+      clip_rewards=FLAGS.clip_rewards)
+
+  logging.info("Replica[%s] is ready to serve requests.", FLAGS.replica)
+  serve(output_dir, env)
+
+
+if __name__ == "__main__":
+  app.run(main)
diff --git a/tensor2tensor/envs/env_service_servicer.py b/tensor2tensor/envs/env_service_servicer.py
new file mode 100644
index 000000000..c7bc72c7d
--- /dev/null
+++ b/tensor2tensor/envs/env_service_servicer.py
@@ -0,0 +1,111 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Implementation of the EnvService RPC."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl import logging
+import grpc
+import numpy as np
+from tensor2tensor.envs import env_service_pb2
+from tensor2tensor.envs import env_service_pb2_grpc
+from tensor2tensor.envs import env_service_serialization as serialization
+
+
+class EnvServiceServicer(env_service_pb2_grpc.EnvServiceServicer):
+  """Implementation of the EnvService service defined in env_service.proto."""
+
+  def __init__(self, env):
+    self._env = env
+
+  def Reset(self, request, context):
+    """Reset."""
+    logging.vlog(1, "EnvServiceServicer is being reset.")
+
+    obs = self._env.reset()
+    reset_response = env_service_pb2.ResetResponse()
+    # Anything more efficient?
+    reset_response.observation.CopyFrom(
+        serialization.numpy_array_to_observation(obs))
+
+    return reset_response
+
+  def Step(self, step_request, context):
+    """Step."""
+    logging.vlog(1, "EnvServiceServicer is being stepped.")
+
+    step_response = env_service_pb2.StepResponse()
+    action = step_request.action
+
+    if "discrete_action" != action.WhichOneof("payload"):
+      context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+      context.set_details("Method not implemented for non-discrete actions!")
+      return step_response
+
+    obs, rewards, dones, infos = self._env.step(
+        np.array([action.discrete_action]))
+
+    step_response.observation.CopyFrom(
+        serialization.numpy_array_to_observation(obs))
+    step_response.reward = rewards
+    step_response.done = dones
+
+    # TODO(afrozm): Take care of this later. `info` is an np array of dicts.
+    if len(infos) > 1:
+      logging.error("Skipping adding the info for other elements in batch.")
+
+    for k, v in infos[0].items():
+      step_response.info.info_map[k] = v
+
+    return step_response
+
+  def Close(self, request, context):
+    """Close."""
+
+    self._env.close()
+    return env_service_pb2.CloseResponse()
+
+  def Render(self, request, context):
+    """Render."""
+
+    mode = request.mode or "rgb_array"
+    rendered_value = self._env.render(mode=mode)
+    response = env_service_pb2.RenderResponse()
+    if (rendered_value is not None) and isinstance(rendered_value, np.ndarray):
+      response.observation = serialization.numpy_array_to_observation(
+          rendered_value)
+
+    return response
+
+  def GetEnvInfo(self, request, context):
+    # Request is empty.
+    del request
+    del context
+
+    response = env_service_pb2.EnvInfoResponse()
+
+    response.observation_space.CopyFrom(
+        serialization.gym_space_to_proto(self._env.observation_space))
+    response.action_space.CopyFrom(
+        serialization.gym_space_to_proto(self._env.action_space))
+    response.reward_range.CopyFrom(
+        serialization.reward_range_to_proto(self._env.reward_range))
+    # Usually these envs aren't batched envs, in that case batch size = 1.
+    response.batch_size = getattr(self._env, "batch_size", 1)
+
+    return response
diff --git a/tensor2tensor/envs/env_service_servicer_test.py b/tensor2tensor/envs/env_service_servicer_test.py
new file mode 100644
index 000000000..11285088a
--- /dev/null
+++ b/tensor2tensor/envs/env_service_servicer_test.py
@@ -0,0 +1,95 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.envs.env_service_servicer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gym
+import mock
+import numpy as np
+from tensor2tensor.envs import env_service_pb2
+from tensor2tensor.envs import env_service_serialization
+from tensor2tensor.envs import env_service_servicer
+from tensorflow import test
+
+
+class EnvServiceServicerTest(test.TestCase):
+
+  def test_get_env_info(self):
+    env = gym.make("CartPole-v0")
+    env_ss = env_service_servicer.EnvServiceServicer(env)
+    env_info = env_ss.GetEnvInfo(None, None)
+
+    self.assertIsInstance(env_info, env_service_pb2.EnvInfoResponse)
+
+    self.assertTrue(env_info.observation_space.HasField("box"))
+    self.assertTrue(env_info.action_space.HasField("discrete"))
+
+    self.assertEqual(1, len(env_info.observation_space.box.shape.dim))
+    self.assertEqual(4, env_info.observation_space.box.shape.dim[0].size)
+    self.assertEqual(2, env_info.action_space.discrete.num_actions)
+
+    self.assertEqual(-np.inf, env_info.reward_range.low)
+    self.assertEqual(np.inf, env_info.reward_range.high)
+
+    self.assertEqual(1, env_info.batch_size)
+
+  def test_reset(self):
+    # Set expectation on a mock.
+    reset_obs = np.array([0.1, 0.2, 0.3, 0.4])
+    env = mock.Mock()
+    env.reset.return_value = reset_obs
+
+    # Call reset.
+    env_ss = env_service_servicer.EnvServiceServicer(env)
+    reset_response = env_ss.Reset(None, None)
+
+    # Assert the set expectation.
+    self.assertIsInstance(reset_response, env_service_pb2.ResetResponse)
+    self.assertAllEqual(
+        reset_obs,
+        env_service_serialization.tensor_proto_to_numpy_array(
+            reset_response.observation.observation))
+
+  def test_step(self):
+    action = 3
+    step_obs = np.array([1.1, 1.2, 1.3, 1.4])
+    reward = 1.2
+    done = False
+    info = {"k1": 1, "k2": 2}
+
+    env = mock.Mock()
+    env.step.return_value = (step_obs, reward, done, [info])
+
+    env_ss = env_service_servicer.EnvServiceServicer(env)
+    step_request = env_service_pb2.StepRequest(
+        action=env_service_pb2.Action(discrete_action=action))
+    step_response = env_ss.Step(step_request, None)
+
+    self.assertAllEqual(
+        step_obs,
+        env_service_serialization.tensor_proto_to_numpy_array(
+            step_response.observation.observation))
+    self.assertEqual(reward, step_response.reward)
+    self.assertEqual(done, step_response.done)
+    self.assertEqual(info["k1"], step_response.info.info_map["k1"])
+    self.assertEqual(info["k2"], step_response.info.info_map["k2"])
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensor2tensor/envs/gym_env_problem.py b/tensor2tensor/envs/gym_env_problem.py
index c9bcdca53..3c8437091 100644
--- a/tensor2tensor/envs/gym_env_problem.py
+++ b/tensor2tensor/envs/gym_env_problem.py
@@ -273,7 +273,20 @@ def _reset(self, indices):
     """
     # This returns a numpy array with first dimension `len(indices)` and the
     # rest being the dimensionality of the observation.
-    return np.stack([self._envs[index].reset() for index in indices])
+
+    num_envs_to_reset = len(indices)
+    observations = [None] * num_envs_to_reset
+
+    def reset_at(idx):
+      observations[idx] = self._envs[indices[idx]].reset()
+
+    if self._parallelism > 1:
+      self._pool.map(reset_at, range(num_envs_to_reset))
+    else:
+      for i in range(num_envs_to_reset):
+        reset_at(i)
+
+    return np.stack(observations)
 
   def _step(self, actions):
     """Takes a step in all environments, shouldn't pre-process or record.
diff --git a/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
index a26c3cd97..92a7a7e54 100644
--- a/tensor2tensor/envs/rendered_env_problem.py
+++ b/tensor2tensor/envs/rendered_env_problem.py
@@ -55,8 +55,14 @@ def __init__(self, *args, **kwargs):
     gym_env_problem.GymEnvProblem.__init__(self, *args, **kwargs)
     video_utils.VideoProblem.__init__(self)
 
-  def initialize_environments(self, batch_size=1):
-    gym_env_problem.GymEnvProblem.initialize_environments(self, batch_size)
+  def initialize_environments(self,
+                              batch_size=1,
+                              parallelism=1,
+                              per_env_kwargs=None,
+                              **kwargs):
+    gym_env_problem.GymEnvProblem.initialize_environments(
+        self, batch_size=batch_size, parallelism=parallelism,
+        per_env_kwargs=per_env_kwargs, **kwargs)
     # Assert the underlying gym environment has correct observation space
     assert len(self.observation_spec.shape) == 3
 

From 50df96cc42e05c481aced3019ef9ebf6fb4c08b8 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sun, 28 Jul 2019 06:22:20 -0700
Subject: [PATCH 2235/2720] Add test for client_env.py, now everything is
 reasonably tested.

PiperOrigin-RevId: 260375198
---
 tensor2tensor/envs/client_env.py      | 20 +++---
 tensor2tensor/envs/client_env_test.py | 99 +++++++++++++++++++++++++++
 2 files changed, 111 insertions(+), 8 deletions(-)
 create mode 100644 tensor2tensor/envs/client_env_test.py

diff --git a/tensor2tensor/envs/client_env.py b/tensor2tensor/envs/client_env.py
index f4a2a1b0b..ea7e3f61a 100644
--- a/tensor2tensor/envs/client_env.py
+++ b/tensor2tensor/envs/client_env.py
@@ -70,16 +70,20 @@ def run_get_env_info(stub):
     return (gym_action_space, gym_observation_space, reward_range,
             env_info_response.batch_size)
 
-  def __init__(self, remote_env_address=None):
-    logging.vlog(1, "Making a ClientEnv with remote address: [%s]",
-                 remote_env_address)
-    assert remote_env_address is not None
-    # Make a channel and stub on the remote env address.
-    self._remote_env_address = remote_env_address
-
+  def __init__(self, remote_env_address=None, stub=None):
     self._channel = None
     self._stub = None
-    self.initialize_stub()
+    self._remote_env_address = None
+
+    if stub is not None:
+      self._stub = stub
+    else:
+      assert remote_env_address is not None
+      logging.vlog(1, "Making a ClientEnv with remote address: [%s]",
+                   remote_env_address)
+      self._remote_env_address = remote_env_address
+      self.initialize_stub()
+
     assert self._stub is not None
 
     # We now have to do an RPC to determine spaces and reward range.
diff --git a/tensor2tensor/envs/client_env_test.py b/tensor2tensor/envs/client_env_test.py
new file mode 100644
index 000000000..aa06b9661
--- /dev/null
+++ b/tensor2tensor/envs/client_env_test.py
@@ -0,0 +1,99 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.envs.client_env."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gym
+import mock
+import numpy as np
+from tensor2tensor.envs import client_env
+from tensor2tensor.envs import env_service_pb2
+from tensor2tensor.envs import env_service_serialization
+from tensorflow import test
+
+
+class ClientEnvTest(test.TestCase):
+
+  def configure_env_info_on_mock(self, mock_obj):
+    env_info_response = env_service_pb2.EnvInfoResponse()
+    env_info_response.observation_space.box.CopyFrom(
+        env_service_serialization.gym_space_to_proto(
+            gym.spaces.Box(low=0, high=255, shape=(28, 28, 3))).box)
+    env_info_response.action_space.discrete.num_actions = 6
+    env_info_response.reward_range.low = -1
+    env_info_response.reward_range.high = 1
+    env_info_response.batch_size = 1
+    mock_obj.GetEnvInfo.return_value = env_info_response
+
+  def test_get_env_info(self):
+    mock_stub = mock.Mock()
+    self.configure_env_info_on_mock(mock_stub)
+
+    env = client_env.ClientEnv(stub=mock_stub)
+
+    self.assertIsInstance(env.action_space, gym.spaces.Discrete)
+    self.assertIsInstance(env.observation_space, gym.spaces.Box)
+
+    self.assertEqual(6, env.action_space.n)
+    self.assertEqual((28, 28, 3), env.observation_space.shape)
+    self.assertEqual((-1, 1), env.reward_range)
+
+  def test_reset(self):
+    mock_stub = mock.Mock()
+    self.configure_env_info_on_mock(mock_stub)
+    obs_np = np.random.uniform(size=(1, 28, 28, 3))
+    reset_response = env_service_pb2.ResetResponse()
+    reset_response.observation.CopyFrom(
+        env_service_serialization.numpy_array_to_observation(obs_np))
+    mock_stub.Reset.return_value = reset_response
+
+    env = client_env.ClientEnv(stub=mock_stub)
+
+    self.assertAllEqual(np.squeeze(obs_np, axis=0), env.reset())
+
+  def test_step(self):
+    mock_stub = mock.Mock()
+    self.configure_env_info_on_mock(mock_stub)
+    obs_np = np.random.uniform(size=(1, 28, 28, 3))
+    reward = 0.5
+    done = True
+    step_response = env_service_pb2.StepResponse(reward=reward, done=done)
+    step_response.observation.CopyFrom(
+        env_service_serialization.numpy_array_to_observation(obs_np))
+    step_response.info.info_map["k1"] = 1
+    step_response.info.info_map["k2"] = 2
+    mock_stub.Step.return_value = step_response
+
+    action = 4
+    step_request = env_service_pb2.StepRequest(
+        action=env_service_pb2.Action(discrete_action=action))
+
+    env = client_env.ClientEnv(stub=mock_stub)
+    step_retval = env.step(action)
+
+    mock_stub.Step.assert_called_with(step_request)
+    self.assertAllEqual(np.squeeze(obs_np, axis=0), step_retval[0])
+    self.assertEqual(reward, step_retval[1])
+    self.assertEqual(done, step_retval[2])
+    self.assertEqual(1, step_retval[3]["k1"])
+    self.assertEqual(2, step_retval[3]["k2"])
+
+
+if __name__ == "__main__":
+  test.main()

From acbb94b5226ada8d91cf11c6837fe2c160f0e3fe Mon Sep 17 00:00:00 2001
From: Adam Roberts <adarob@google.com>
Date: Mon, 29 Jul 2019 12:15:40 -0700
Subject: [PATCH 2236/2720] Make SequenceDatasetPacker pay attention to `keys`
 argument.

PiperOrigin-RevId: 260552329
---
 tensor2tensor/data_generators/generator_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 33f3d3d2f..43d4ff14d 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -864,7 +864,7 @@ def _standardize(self, dataset, keys):
     shapes = tf.compat.v1.data.get_output_shapes(dataset)
 
     if isinstance(shapes, dict):
-      keys = tuple(shapes.keys())
+      keys = keys or tuple(shapes.keys())
       dataset = dataset.map(lambda x: tuple(x[k] for k in keys))
       shapes = tf.compat.v1.data.get_output_shapes(dataset)
 

From 9bc78862300ee6c3f5f8b2986e1be89a06718077 Mon Sep 17 00:00:00 2001
From: David Dohan <ddohan@google.com>
Date: Mon, 29 Jul 2019 16:44:35 -0700
Subject: [PATCH 2237/2720] Add causal convolution layer.

PiperOrigin-RevId: 260607887
---
 tensor2tensor/trax/layers/convolution.py      | 34 +++++++++++++++++++
 tensor2tensor/trax/layers/convolution_test.py | 11 ++++++
 2 files changed, 45 insertions(+)

diff --git a/tensor2tensor/trax/layers/convolution.py b/tensor2tensor/trax/layers/convolution.py
index c30669af9..9b556ef8d 100644
--- a/tensor2tensor/trax/layers/convolution.py
+++ b/tensor2tensor/trax/layers/convolution.py
@@ -90,3 +90,37 @@ def new_parameters(self, input_shape, input_dtype, rng):
     w = self._kernel_initializer(kernel_shape, rng)
     b = self._bias_initializer(bias_shape, rng)
     return (w, b)
+
+
+class CausalConv(Conv):
+  """Causal (masked) convolution for [batch x time x depth] sequences.
+
+  Maintains causality along time axis. Used in language modeling tasks.
+  """
+
+  def __init__(self,
+               filters,
+               kernel_width=3,
+               kernel_initializer=None,
+               bias_initializer=init.RandomNormalInitializer(1e-6)):
+    super(CausalConv, self).__init__(
+        filters=filters,
+        kernel_size=(kernel_width,),
+        strides=None,
+        padding='VALID',
+        dimension_numbers=('NWC', 'WIO', 'NWC'),
+        kernel_initializer=kernel_initializer,
+        bias_initializer=bias_initializer)
+
+  def call(self, x, params=(), **kwargs):
+    assert self._padding == 'VALID'
+    # Left pad with 0s. Applying an unmasked valid convolution on top of this
+    # yields a causal convolution.
+    # TODO(ddohan): Support strided and dilated convolutions.
+    rate = 1
+    effective_kernel_size = int((self._kernel_size[0] - 1) * rate + 1)
+    pad = effective_kernel_size - 1
+    x_leftpad = np.pad(x, pad_width=[[0, 0], [pad, 0], [0, 0]], mode='constant')
+
+    res = super(CausalConv, self).call(x_leftpad, params)
+    return res
diff --git a/tensor2tensor/trax/layers/convolution_test.py b/tensor2tensor/trax/layers/convolution_test.py
index 62b931fc0..b17f4721f 100644
--- a/tensor2tensor/trax/layers/convolution_test.py
+++ b/tensor2tensor/trax/layers/convolution_test.py
@@ -38,5 +38,16 @@ def test_conv_rebatch(self):
     self.assertEqual(result_shape, (3, 29, 3, 3, 30))
 
 
+class CausalConvolutionTest(absltest.TestCase):
+
+  def test_causal_conv(self):
+    input_shape = (29, 5, 20)
+    conv = convolution.CausalConv(filters=30, kernel_width=3)
+    result_shape = base.check_shape_agreement(conv, input_shape)
+    self.assertEqual(result_shape, (29, 5, 30))
+
+    # TODO(ddohan): How to test for causality? Gradient check between positions?
+
+
 if __name__ == "__main__":
   absltest.main()

From 937a05164b1056819bb2f8603e6026f23832f8c6 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 29 Jul 2019 17:30:17 -0700
Subject: [PATCH 2238/2720] Make Adafactor configurable and use in
 transformer-revnet (improved memory use and so results).

PiperOrigin-RevId: 260615704
---
 .../transformer_revnet_imagenet64_8gb.gin     | 10 ++--
 .../configs/transformer_revnet_lm1b_8gb.gin   | 56 -------------------
 tensor2tensor/trax/optimizers/base.py         | 56 ++++++++++---------
 3 files changed, 34 insertions(+), 88 deletions(-)
 delete mode 100644 tensor2tensor/trax/configs/transformer_revnet_lm1b_8gb.gin

diff --git a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
index 7716aed65..dd4e1f47e 100644
--- a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
@@ -18,9 +18,9 @@ inputs.n_chunks = 16
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-MultifactorSchedule.constant = 0.2
-MultifactorSchedule.factors = 'constant * linear_warmup'
-MultifactorSchedule.warmup_steps = 2000
+MultifactorSchedule.constant = 2.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
 
 # Parameters for train:
 # ==============================================================================
@@ -28,13 +28,13 @@ train.eval_frequency = 100
 train.eval_steps = 8
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerRevnetLM
-train.optimizer = @trax.optimizers.SM3
+train.optimizer = @trax.optimizers.Adafactor
 train.train_steps = 500000
 train.trainer_class = @MemoryEfficientTrainer
 
 # Parameters for TransformerRevnetLM:
 # ==============================================================================
-TransformerRevnetLM.d_model = 512
+TransformerRevnetLM.d_model = 1024
 TransformerRevnetLM.d_ff = 2048
 TransformerRevnetLM.d_attention_key = 32
 TransformerRevnetLM.d_attention_value = 32
diff --git a/tensor2tensor/trax/configs/transformer_revnet_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_lm1b_8gb.gin
deleted file mode 100644
index 3f4cd14aa..000000000
--- a/tensor2tensor/trax/configs/transformer_revnet_lm1b_8gb.gin
+++ /dev/null
@@ -1,56 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fun:
-# ==============================================================================
-batch_fun.batch_size_per_device = 4096
-batch_fun.eval_batch_size = 128
-batch_fun.max_eval_length = 2048
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_languagemodel_lm1b32k_packed'
-inputs.input_name = 'targets'
-inputs.n_chunks = 32
-
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 0.3
-MultifactorSchedule.factors = 'constant * linear_warmup'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
-lm1b_preprocess.max_target_length = 512
-lm1b_preprocess.max_eval_target_length = 2048
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 100
-train.eval_steps = 10
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerRevnetLM
-train.optimizer = @trax.optimizers.SM3
-train.train_steps = 500000
-train.trainer_class = @MemoryEfficientTrainer
-
-# Parameters for TransformerRevnetLM:
-# ==============================================================================
-TransformerRevnetLM.d_model = 512
-TransformerRevnetLM.d_ff = 2048
-TransformerRevnetLM.dropout = 0.1
-TransformerRevnetLM.max_len = 2048
-TransformerRevnetLM.mode = 'train'
-TransformerRevnetLM.n_heads = 8
-TransformerRevnetLM.n_layers = 6
-TransformerRevnetLM.vocab_size = 32000
-TransformerRevnetLM.n_chunks = 32
-TransformerRevnetLM.n_attention_chunks = 8
diff --git a/tensor2tensor/trax/optimizers/base.py b/tensor2tensor/trax/optimizers/base.py
index ec2b15a43..725b7c3c3 100644
--- a/tensor2tensor/trax/optimizers/base.py
+++ b/tensor2tensor/trax/optimizers/base.py
@@ -19,8 +19,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import functools
-
 from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers import base as layers
 
@@ -243,11 +241,13 @@ class Adafactor(Optimizer):
   # TODO(levskaya): refactor to use newer RL friendly parameter passing.
   def __init__(self,
                learning_rate,
-               decay_rate=0.8,
-               beta1=0.0,
-               clipping_threshold=1.0,
                factored=True,
                multiply_by_parameter_scale=True,
+               do_clipping=True,
+               do_momentum=False,
+               beta1=0.0,
+               decay_rate=0.8,
+               clipping_threshold=1.0,
                epsilon1=1e-30,
                epsilon2=1e-3):
     """Create the Adafactor optimizer.
@@ -256,27 +256,29 @@ def __init__(self,
 
     Args:
       learning_rate: float: trax-provided learning rate.
-      decay_rate: float: controls second-moment exponential decay schedule.
-      beta1: a float value between 0 and 1, enables momentum and uses extra
-        memory if nonzero!  Off by default.
-      clipping_threshold: an optional float >= 1, if None no update clipping.
       factored: boolean: whether to use factored second-moment estimator for 2d
         variables.
       multiply_by_parameter_scale: boolean: if True, then scale provided
         learning_rate by parameter norm. if False, provided learning_rate is
         absolute step size.
+      do_clipping: whether to clip gradients; if True, set clipping_theshold.
+      do_momentum: whether to use momentum; if True, set beta1.
+      beta1: a float value between 0 and 1, enables momentum and uses extra
+        memory if nonzero!  Off by default.
+      decay_rate: float: controls second-moment exponential decay schedule.
+      clipping_threshold: an optional float >= 1, if None no update clipping.
       epsilon1: Regularization constant for squared gradient.
       epsilon2: Regularization constant for parameter scale.
     """
-    super(Adafactor, self).__init__(learning_rate)
-    self._multiply_by_parameter_scale = multiply_by_parameter_scale
+    # These 4 parameters are not configurable once the class is created.
     self._factored = factored
-    self._beta1 = beta1
-    self._clipping_threshold = clipping_threshold
-    self._epsilon1 = epsilon1
-    self._epsilon2 = epsilon2
-    self._decay_rate = functools.partial(self._decay_rate_pow,
-                                         exponent=decay_rate)
+    self._multiply_by_parameter_scale = multiply_by_parameter_scale
+    self._do_clipping = do_clipping
+    self._do_momentum = do_momentum
+    # Dynamically configurable parameters will be passed to the update function.
+    super(Adafactor, self).__init__(
+        learning_rate, beta1, decay_rate, clipping_threshold,
+        epsilon1, epsilon2)
 
   @staticmethod
   def _decay_rate_pow(i, exponent=0.8):
@@ -294,22 +296,23 @@ def init(self, params):
     else:
       v = np.zeros_like(params)
       slots.append(v)
-    if self._beta1:
+    if self._do_momentum:
       m = np.zeros_like(params)
       slots.append(m)
     return slots
 
   def update(self, step, grads, params, slots, opt_params):
     updates = []
-    (learning_rate,) = opt_params
-    decay_rate = self._decay_rate(step)
+    (learning_rate, beta1, decay_rate, clipping_threshold,
+     epsilon1, epsilon2) = opt_params
+    decay_rate = self._decay_rate_pow(step, exponent=decay_rate)
     update_scale = learning_rate
     if self._multiply_by_parameter_scale:
       update_scale *= np.maximum(
-          np.sqrt(np.mean(params * params)), self._epsilon2)
+          np.sqrt(np.mean(params * params)), epsilon2)
     mixing_rate = 1.0 - decay_rate
 
-    grads_sqr = grads * grads + self._epsilon1
+    grads_sqr = grads * grads + epsilon1
     if self._factored and len(params.shape) >= 2:
       v_row = slots.pop(0)
       v_col = slots.pop(0)
@@ -328,16 +331,15 @@ def update(self, step, grads, params, slots, opt_params):
       updates.append(new_v)
       y = grads * (new_v)**-0.5
 
-    if self._clipping_threshold:
+    if self._do_clipping:
       clipping_denom = (
-          np.maximum(1.0,
-                     np.sqrt(np.mean(y * y)) / self._clipping_threshold))
+          np.maximum(1.0, np.sqrt(np.mean(y * y)) / clipping_threshold))
       y /= clipping_denom
 
     subtrahend = update_scale * y
-    if self._beta1:
+    if self._do_momentum:
       m = slots.pop(0)
-      new_m = self._beta1 * m + (1.0 - self._beta1) * subtrahend
+      new_m = beta1 * m + (1.0 - beta1) * subtrahend
       subtrahend = new_m
       updates.append(new_m)
 

From 200ce96f835d34cca6817303ea2fdcf2e3cb021b Mon Sep 17 00:00:00 2001
From: Sepehr Sameni <Sepehr.Sameni@gmail.com>
Date: Tue, 30 Jul 2019 21:45:13 +0430
Subject: [PATCH 2239/2720] correct typo in add_timing_signal_nd (#1651)

---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 4d35a1498..33116776b 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -619,7 +619,7 @@ def add_timing_signal_nd(x, min_timescale=1.0, max_timescale=1.0e4):
   memory inputs to attention.
 
   The use of relative position is possible because sin(a+b) and cos(a+b) can be
-  experessed in terms of b, sin(a) and cos(a).
+  expressed in terms of b, sin(a) and cos(a).
 
   x is a Tensor with n "positional" dimensions, e.g. one dimension for a
   sequence or two dimensions for an image

From ce56fd9703d4f4fc89c47e81b10f8c275a3da178 Mon Sep 17 00:00:00 2001
From: dong-s <dongxin1024@yahoo.com>
Date: Wed, 31 Jul 2019 01:16:15 +0800
Subject: [PATCH 2240/2720] fix decode bug (#1645)

---
 tensor2tensor/visualization/visualization.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/visualization/visualization.py b/tensor2tensor/visualization/visualization.py
index 0fae8722e..d411e75ca 100644
--- a/tensor2tensor/visualization/visualization.py
+++ b/tensor2tensor/visualization/visualization.py
@@ -58,12 +58,17 @@ def encode(self, input_str):
   def decode(self, integers):
     """List of ints to str."""
     integers = list(np.squeeze(integers))
-    return self.encoders["inputs"].decode(integers)
+    return self.encoders['targets'].decode(integers)
+
+  def encode_list(self, integers):
+    """List of ints to list of str."""
+    integers = list(np.squeeze(integers))
+    return self.encoders['inputs'].decode_list(integers)
 
   def decode_list(self, integers):
     """List of ints to list of str."""
     integers = list(np.squeeze(integers))
-    return self.encoders["inputs"].decode_list(integers)
+    return self.encoders['targets'].decode_list(integers)
 
   def get_vis_data_from_string(self, sess, input_string):
     """Constructs the data needed for visualizing attentions.
@@ -104,7 +109,7 @@ def get_vis_data_from_string(self, sess, input_string):
     })
 
     output_string = self.decode(out)
-    input_list = self.decode_list(encoded_inputs)
+    input_list = self.encode_list(encoded_inputs)
     output_list = self.decode_list(out)
 
     return output_string, input_list, output_list, att_mats

From bba231fe9c945bc9f1a77f83a1dee1bab02fd172 Mon Sep 17 00:00:00 2001
From: Vinh Nguyen <vinh.nguyenx@gmail.com>
Date: Wed, 31 Jul 2019 03:21:44 +1000
Subject: [PATCH 2241/2720] Adding automatic mixed precision support (#1637)

* adding GPU auto mixed precision training

* add warning: manual vs. amp

* fix os environment check

* add _use_locking attribute to optimizer

* add _name attribute to optimizer

* change OS flag to TF_ENABLE_AUTO_MIXED_PRECISION
---
 tensor2tensor/utils/optimize.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index ae962c32f..fa620da3b 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 import numpy as np
+import os
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import adafactor as adafactor_lib
@@ -40,7 +41,7 @@ def _mixed_precision_is_enabled(hparams):
   return activation_dtype == tf.float16 and weight_dtype == tf.float32
 
 
-def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
+def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None, gpu_auto_mixed_precision=False):
   """Minimize loss."""
   loss = weight_decay_and_noise(loss, hparams, learning_rate)
   loss = tf.identity(loss, name="total_loss")
@@ -65,6 +66,16 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
   opt = ConditionalOptimizer(hparams.optimizer, learning_rate, hparams, use_tpu)
   if use_tpu:
     opt = tf.contrib.tpu.CrossShardOptimizer(opt)
+  if os.environ.get('TF_ENABLE_AUTO_MIXED_PRECISION', default='0') == '1' or gpu_auto_mixed_precision:
+      if use_tpu:
+          raise(RuntimeError("GPU auto mixed precision cannot be used with TPU"))
+      elif _mixed_precision_is_enabled(hparams):
+          raise(RuntimeError("GPU auto mixed precision cannot be used with manual mixed precision"))
+      else:
+          setattr(opt, '_use_locking', 'True')
+          setattr(opt, '_name', 'ConditionalOptimizer')
+          opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)
+
   opt_summaries = []
   if common_layers.should_generate_summaries():
     tf.summary.scalar("learning_rate", learning_rate)

From 4e0daf5ab4b8184beac85bd087b88e22d7b01cb0 Mon Sep 17 00:00:00 2001
From: dong-s <dongxin1024@yahoo.com>
Date: Tue, 30 Jul 2019 10:26:04 -0700
Subject: [PATCH 2242/2720] Merge of PR #1645

PiperOrigin-RevId: 260743903
---
 tensor2tensor/layers/common_attention.py     |  2 +-
 tensor2tensor/utils/optimize.py              | 13 +------------
 tensor2tensor/visualization/visualization.py |  6 +++---
 3 files changed, 5 insertions(+), 16 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 33116776b..4d35a1498 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -619,7 +619,7 @@ def add_timing_signal_nd(x, min_timescale=1.0, max_timescale=1.0e4):
   memory inputs to attention.
 
   The use of relative position is possible because sin(a+b) and cos(a+b) can be
-  expressed in terms of b, sin(a) and cos(a).
+  experessed in terms of b, sin(a) and cos(a).
 
   x is a Tensor with n "positional" dimensions, e.g. one dimension for a
   sequence or two dimensions for an image
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index fa620da3b..ae962c32f 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -18,7 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 import numpy as np
-import os
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import adafactor as adafactor_lib
@@ -41,7 +40,7 @@ def _mixed_precision_is_enabled(hparams):
   return activation_dtype == tf.float16 and weight_dtype == tf.float32
 
 
-def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None, gpu_auto_mixed_precision=False):
+def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
   """Minimize loss."""
   loss = weight_decay_and_noise(loss, hparams, learning_rate)
   loss = tf.identity(loss, name="total_loss")
@@ -66,16 +65,6 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None, gpu_au
   opt = ConditionalOptimizer(hparams.optimizer, learning_rate, hparams, use_tpu)
   if use_tpu:
     opt = tf.contrib.tpu.CrossShardOptimizer(opt)
-  if os.environ.get('TF_ENABLE_AUTO_MIXED_PRECISION', default='0') == '1' or gpu_auto_mixed_precision:
-      if use_tpu:
-          raise(RuntimeError("GPU auto mixed precision cannot be used with TPU"))
-      elif _mixed_precision_is_enabled(hparams):
-          raise(RuntimeError("GPU auto mixed precision cannot be used with manual mixed precision"))
-      else:
-          setattr(opt, '_use_locking', 'True')
-          setattr(opt, '_name', 'ConditionalOptimizer')
-          opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)
-
   opt_summaries = []
   if common_layers.should_generate_summaries():
     tf.summary.scalar("learning_rate", learning_rate)
diff --git a/tensor2tensor/visualization/visualization.py b/tensor2tensor/visualization/visualization.py
index d411e75ca..3519751b9 100644
--- a/tensor2tensor/visualization/visualization.py
+++ b/tensor2tensor/visualization/visualization.py
@@ -58,17 +58,17 @@ def encode(self, input_str):
   def decode(self, integers):
     """List of ints to str."""
     integers = list(np.squeeze(integers))
-    return self.encoders['targets'].decode(integers)
+    return self.encoders["targets"].decode(integers)
 
   def encode_list(self, integers):
     """List of ints to list of str."""
     integers = list(np.squeeze(integers))
-    return self.encoders['inputs'].decode_list(integers)
+    return self.encoders["inputs"].decode_list(integers)
 
   def decode_list(self, integers):
     """List of ints to list of str."""
     integers = list(np.squeeze(integers))
-    return self.encoders['targets'].decode_list(integers)
+    return self.encoders["targets"].decode_list(integers)
 
   def get_vis_data_from_string(self, sess, input_string):
     """Constructs the data needed for visualizing attentions.

From d13c33183998f8de31c49d83a3d4048eade2e293 Mon Sep 17 00:00:00 2001
From: Sepehr Sameni <Sepehr.Sameni@gmail.com>
Date: Tue, 30 Jul 2019 10:26:18 -0700
Subject: [PATCH 2243/2720] Merge of PR #1651

PiperOrigin-RevId: 260743960
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 4d35a1498..33116776b 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -619,7 +619,7 @@ def add_timing_signal_nd(x, min_timescale=1.0, max_timescale=1.0e4):
   memory inputs to attention.
 
   The use of relative position is possible because sin(a+b) and cos(a+b) can be
-  experessed in terms of b, sin(a) and cos(a).
+  expressed in terms of b, sin(a) and cos(a).
 
   x is a Tensor with n "positional" dimensions, e.g. one dimension for a
   sequence or two dimensions for an image

From 5bfe69a7d68b7d61d51fac36c6088f94b9d6fdc6 Mon Sep 17 00:00:00 2001
From: Vinh Nguyen <vinh.nguyenx@gmail.com>
Date: Tue, 30 Jul 2019 11:09:28 -0700
Subject: [PATCH 2244/2720] Merge of PR #1637

PiperOrigin-RevId: 260754631
---
 tensor2tensor/utils/optimize.py | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index ae962c32f..fd8bf46c3 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -17,8 +17,9 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-import numpy as np
 
+import os
+import numpy as np
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import adafactor as adafactor_lib
 from tensor2tensor.utils import misc_utils
@@ -40,7 +41,12 @@ def _mixed_precision_is_enabled(hparams):
   return activation_dtype == tf.float16 and weight_dtype == tf.float32
 
 
-def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
+def optimize(loss,
+             learning_rate,
+             hparams,
+             use_tpu=False,
+             variables=None,
+             gpu_auto_mixed_precision=False):
   """Minimize loss."""
   loss = weight_decay_and_noise(loss, hparams, learning_rate)
   loss = tf.identity(loss, name="total_loss")
@@ -65,6 +71,18 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None):
   opt = ConditionalOptimizer(hparams.optimizer, learning_rate, hparams, use_tpu)
   if use_tpu:
     opt = tf.contrib.tpu.CrossShardOptimizer(opt)
+  if gpu_auto_mixed_precision or os.environ.get(
+      "TF_ENABLE_AUTO_MIXED_PRECISION", "0") == "1":
+    if use_tpu:
+      raise RuntimeError("GPU auto mixed precision cannot be used with TPU")
+    elif _mixed_precision_is_enabled(hparams):
+      raise RuntimeError(
+          "GPU auto mixed precision cannot be used with manual mixed precision")
+    else:
+      setattr(opt, "_use_locking", "True")
+      setattr(opt, "_name", "ConditionalOptimizer")
+      opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)
+
   opt_summaries = []
   if common_layers.should_generate_summaries():
     tf.summary.scalar("learning_rate", learning_rate)

From e940b5cd6436463b8163d2072229b52277bff220 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 31 Jul 2019 07:46:16 -0700
Subject: [PATCH 2245/2720] Transformer Librispeech test case cleanup and
 addition

PiperOrigin-RevId: 260924863
---
 tensor2tensor/models/transformer_test.py | 72 +++++++++++++++---------
 1 file changed, 46 insertions(+), 26 deletions(-)

diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 6f55d44cb..678a25341 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -64,10 +64,42 @@ def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
   return model_cls(hparams, mode, p_hparams), features
 
 
+def small_librispeech_model(param_overrides=None):
+  hparams = transformer.transformer_small()
+  hparams.hidden_size = 8
+  hparams.filter_size = 32
+  hparams.num_heads = 1
+  hparams.layer_prepostprocess_dropout = 0.0
+  p_hparams = librispeech.Librispeech().get_hparams(hparams)
+  p_hparams.vocab_size["targets"] = VOCAB_SIZE
+  hparams.problem_hparams = p_hparams
+  model = transformer.Transformer(hparams, problem_hparams=p_hparams)
+  if param_overrides is not None:  # Add or Set any provided HParams
+    assert isinstance(param_overrides, dict)
+    for param_name in param_overrides:
+      if hasattr(hparams, param_name):
+        hparams.set_hparam(param_name, param_overrides[param_name])
+      else:
+        hparams.add_hparam(param_name, param_overrides[param_name])
+  inputs = np.random.rand(
+      BATCH_SIZE, INPUT_LENGTH, 80, 3).astype("float32")  # modify for speech
+  targets = np.random.randint(
+      VOCAB_SIZE, size=(BATCH_SIZE, TARGET_LENGTH, 1, 1))
+  features = {
+      "inputs": tf.constant(inputs, dtype=tf.float32, name="inputs"),
+      "targets": tf.constant(targets, dtype=tf.int32, name="targets"),
+      "target_space_id": tf.constant(1, dtype=tf.int32)
+  }
+  return model, features
+
+
 class TransformerTest(tf.test.TestCase):
 
-  def testTransformer(self):
-    model, features = get_model(transformer.transformer_small())
+  def testTransformer(self, get_model_fn=None, p=None):
+    if get_model_fn:
+      model, features = get_model_fn(param_overrides=p)
+    else:
+      model, features = get_model(transformer.transformer_small())
     logits, _ = model(features)
     with self.test_session() as session:
       session.run(tf.global_variables_initializer())
@@ -75,30 +107,15 @@ def testTransformer(self):
     self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, 1, 1, VOCAB_SIZE))
 
   def testTransformerLibrispeech(self, params=None):
-    model_hparams = transformer.transformer_small()
-    if params is not None:  # Add or Set any provided HParams
-      assert isinstance(params, dict)
-      for param_name in params:
-        if hasattr(model_hparams, param_name):
-          model_hparams.set_hparam(param_name, params[param_name])
-        else:
-          model_hparams.add_hparam(param_name, params[param_name])
-    problem = librispeech.Librispeech()
-    model_hparams.problem_hparams = problem.get_hparams(model_hparams)
-    model_hparams._problem_hparams = model_hparams.problem_hparams
-    model, features = get_model(model_hparams)
-    model._problem_hparams.modality = {"inputs": "speech_recognition",
-                                       "targets": "symbol"}
-    features["inputs"] = np.random.rand(
-        BATCH_SIZE, INPUT_LENGTH, 80, 3).astype("float32")  # modify for speech
+    self.testTransformer(get_model_fn=small_librispeech_model, p=params)
 
-    logits, _ = model(features)
-    with self.test_session() as session:
-      session.run(tf.global_variables_initializer())
-      res = session.run(logits)
-    self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, 1, 1, VOCAB_SIZE))
+  def testLibrispeechSlowVsFast(self, params=None):
+    self.testSlowVsFast(get_model_fn=small_librispeech_model, p=params)
 
-  def testTransformerLibrispeechWithAreaAttention(self):
+  def testLibrispeechMultihead(self, params=None):
+    self.testTransformerLibrispeech({"num_heads": 2})
+
+  def testLibrispeechWithAreaAttention(self):
     self.testTransformerLibrispeech({"max_area_width": 2,
                                      "num_area_layers": 1,
                                      "area_key_mode": "mean",
@@ -112,8 +129,11 @@ def testTransformerRelative(self):
       res = session.run(logits)
     self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, 1, 1, VOCAB_SIZE))
 
-  def testSlowVsFast(self):
-    model, features = get_model(transformer.transformer_small())
+  def testSlowVsFast(self, get_model_fn=None, p=None):
+    if get_model_fn:
+      model, features = get_model_fn(param_overrides=p)
+    else:
+      model, features = get_model(transformer.transformer_small())
 
     decode_length = 3
 

From d425b813ad3acb71ca0705bf429ea235c30a69c5 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 31 Jul 2019 08:47:16 -0700
Subject: [PATCH 2246/2720] Introduce per-element weights in trax.train

PiperOrigin-RevId: 260934106
---
 tensor2tensor/trax/backend.py   |  7 +++-
 tensor2tensor/trax/trax.py      | 67 +++++++++++++++++++++----------
 tensor2tensor/trax/trax_test.py | 71 +++++++++++++++++++++++++++++++--
 3 files changed, 118 insertions(+), 27 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index 07cba1542..a78381abd 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -219,5 +219,8 @@ def use_backend(name):
   global override_backend_name
   prev_name = override_backend_name
   override_backend_name = name
-  yield
-  override_backend_name = prev_name
+  # Run the decorated function in try-finally in case it throws, e.g. for tests.
+  try:
+    yield
+  finally:
+    override_backend_name = prev_name
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index f5574df05..219f95299 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -52,65 +52,88 @@
 from tensorflow.io import gfile
 
 
-def _make_list(predictions, targets):
-  """Helper: make predictions and targets lists, check they match on length."""
+@gin.configurable
+def unpack_batch(batch, has_weights=False):
+  """Unpacks a training batch into inputs, targets and weights."""
+  if has_weights:
+    assert len(batch) == 3  # (inputs, targets, weights)
+    return batch
+  else:
+    inputs, targets = batch
+    if isinstance(inputs, (list, tuple)):
+      # If weights are not provided, use scalar 1s and rely on broadcasting.
+      weights = [1.0] * len(inputs)
+    else:
+      weights = 1.0
+    return inputs, targets, weights
+
+
+def _make_list(predictions, targets, weights):
+  """Make predictions, targets and weights lists, check they match on length."""
   #  Our models sometimes return predictions in lists, make it a list always.
   # TODO(lukaszkaiser): make abstractions for nested structures and refactor.
   if not isinstance(predictions, (list, tuple)):
     if isinstance(targets, (list, tuple)):
       raise ValueError("Targets are a list or tuple but predictions are not.")
-    predictions, targets = [predictions], [targets]
+    if isinstance(weights, (list, tuple)):
+      raise ValueError("Weights are a list or tuple but predictions are not.")
+    predictions, targets, weights = [predictions], [targets], [weights]
   if len(predictions) != len(targets):
     raise ValueError("Predictions and targets have different lengths.")
-  return list(predictions), list(targets)
+  if len(predictions) != len(weights):
+    raise ValueError("Predictions and weights have different lengths.")
+  return list(predictions), list(targets), list(weights)
 
 
-@gin.configurable(blacklist=["inputs", "targets"])
-def masked_mean(inputs, targets, mask_id=None):
-  """Mean of the inputs but counting only those where targets != mask_id."""
+@gin.configurable(blacklist=["inputs", "targets", "weights"])
+def masked_mean(inputs, targets, weights, mask_id=None):
+  """Weighted mean of the inputs, excluding where targets == mask_id."""
   inputs = [x.astype(np.float32) for x in inputs]
   # We assume all elements in the list contribute equally.
   # TODO(lukaszkaiser): remove this assumption (e.g., when masks differ).
   length = len(inputs)
-  if mask_id is None:
-    # TODO(lukaszkaiser): can we just divide the sum by length? XLA optimizes?
-    return sum([np.mean(x) / length for x in inputs])
-  unmask = [1.0 - np.equal(t, mask_id).astype(np.float32) for t in targets]
-  return sum([np.sum(x * m) / (length * np.sum(m))
-              for x, m in zip(inputs, unmask)])
+  if mask_id is not None:
+    weights = [w * (1.0 - np.equal(t, mask_id).astype(np.float32))
+               for t, w in zip(targets, weights)]
+  weight_sums = [t.size if np.isscalar(w) else np.sum(w)
+                 for w, t in zip(weights, targets)]
+  return sum([np.sum(x * w) / (length * s)
+              for x, w, s in zip(inputs, weights, weight_sums)])
 
 
 def accuracy(batch, model_predictions):
   """Calculate accuracy."""
-  _, targets = batch
-  model_predictions, targets = _make_list(model_predictions, targets)
+  _, targets, weights = unpack_batch(batch)
+  model_predictions, targets, weights = _make_list(
+      model_predictions, targets, weights)
   correct = []
   for (prediction, target) in zip(model_predictions, targets):
     predicted_class = np.argmax(prediction, axis=-1)
     correct.append(np.equal(predicted_class, target))
-  return masked_mean(correct, targets)
+  return masked_mean(correct, targets, weights)
 
 
 def neg_log_perplexity(batch, model_predictions):
   """Calculate negative log perplexity."""
-  _, targets = batch
-  model_predictions, targets = _make_list(model_predictions, targets)
+  _, targets, weights = unpack_batch(batch)
+  model_predictions, targets, weights = _make_list(
+      model_predictions, targets, weights)
   xent = []
   for (prediction, target) in zip(model_predictions, targets):
     hot_target = layers.one_hot(target, prediction.shape[-1])
     xent.append(np.sum(prediction * hot_target, axis=-1))
-  return masked_mean(xent, targets)
+  return masked_mean(xent, targets, weights)
 
 
 def loss(params, batch, model_predict, rng):
   """Calculate loss."""
-  inputs, targets = batch
+  inputs, targets, weights = unpack_batch(batch)
   predictions = model_predict(inputs, params, rng=rng)
-  predictions, targets = _make_list(predictions, targets)
+  predictions, targets, weights = _make_list(predictions, targets, weights)
   xent = []
   for (pred, target) in zip(predictions, targets):
     xent.append(np.sum(pred * layers.one_hot(target, pred.shape[-1]), axis=-1))
-  return - masked_mean(xent, targets)
+  return - masked_mean(xent, targets, weights)
 
 
 def log(s, stdout=True):
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index 956b64f10..28d2d425b 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -23,10 +23,12 @@
 import functools
 import tempfile
 
+import gin
 from jax import test_util  # pylint: disable=unused-import
 from jax.config import config
 import numpy as np
 
+from tensor2tensor.trax import backend
 from tensor2tensor.trax import inputs as inputs_lib
 from tensor2tensor.trax import layers
 from tensor2tensor.trax import models
@@ -37,15 +39,20 @@
 from tensorflow.io import gfile
 
 
-def test_inputs(n_classes):
+def test_inputs(n_classes, with_weights=False):
   """Make trax.inputs.Inputs."""
   batch_size = 2
   input_shape = (6, 6, 3)
 
   def input_stream():
     while True:
-      yield (np.random.rand(*([batch_size] + list(input_shape))),
-             np.random.randint(n_classes, size=batch_size))
+      inputs = np.random.rand(*([batch_size] + list(input_shape)))
+      targets = np.random.randint(n_classes, size=batch_size)
+      weights = np.random.rand(batch_size)
+      if with_weights:
+        yield inputs, targets, weights
+      else:
+        yield inputs, targets
 
   return inputs_lib.Inputs(
       train_stream=input_stream,
@@ -156,6 +163,64 @@ def test_train_restart(self):
       # Assert total train steps
       self.assertEqual(state.step, 2 * train_steps)
 
+  def test_train_with_weights(self):
+    with self.tmp_dir() as output_dir:
+      gin.bind_parameter("unpack_batch.has_weights", True)
+
+      # Prepare model and inputs
+      n_classes = 4
+      train_steps = 2
+      eval_steps = 2
+      model_fn = functools.partial(models.MLP,
+                                   d_hidden=16,
+                                   n_output_classes=n_classes)
+      inputs = lambda _: test_inputs(n_classes, with_weights=True)
+
+      # Train and evaluate
+      state = trax.train(output_dir,
+                         model=model_fn,
+                         inputs=inputs,
+                         train_steps=train_steps,
+                         eval_steps=eval_steps)
+
+      # Assert total train steps
+      self.assertEqual(state.step, train_steps)
+
+
+class MaskedMeanTest(test.TestCase):
+
+  def test_computes_basic_mean(self):
+    inputs = [np.array([1, 2, 3])]
+    targets = [np.zeros(3)]
+    weights = [1]
+    with backend.use_backend("numpy"):
+      mean = trax.masked_mean(inputs, targets, weights)
+      np.testing.assert_allclose(mean, 2)
+
+  def test_computes_mean_with_weights(self):
+    inputs = [np.array([1, 2, 3])]
+    targets = [np.zeros(3)]
+    weights = [np.array([3, 1, 0])]
+    with backend.use_backend("numpy"):
+      mean = trax.masked_mean(inputs, targets, weights)
+      np.testing.assert_allclose(mean, 1.25)
+
+  def test_computes_mean_with_mask(self):
+    inputs = [np.array([1, 2, 3])]
+    targets = [np.array([1, 0, 0])]
+    weights = [1]
+    with backend.use_backend("numpy"):
+      mean = trax.masked_mean(inputs, targets, weights, mask_id=1)
+      np.testing.assert_allclose(mean, 2.5)
+
+  def test_computes_mean_with_weights_and_mask(self):
+    inputs = [np.array([1, 2, 4])]
+    targets = [np.array([1, 0, 0])]
+    weights = [np.array([10, 4, 1])]
+    with backend.use_backend("numpy"):
+      mean = trax.masked_mean(inputs, targets, weights, mask_id=1)
+      np.testing.assert_allclose(mean, 2.4)
+
 
 if __name__ == "__main__":
   config.config_with_absl()

From 36d03439394c90beeb6c77b9bcf9589d1fe22a53 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 31 Jul 2019 21:56:12 -0700
Subject: [PATCH 2247/2720] Fix silent but confusing bug in TRAX example
 notebook.

PiperOrigin-RevId: 261058969
---
 tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb b/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb
index ddb354a4e..86663d97f 100644
--- a/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb
+++ b/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb
@@ -489,12 +489,12 @@
         "  pos = len(enc)\n",
         "  rng = trax_random.get_prng(0)\n",
         "  data = np.zeros((1, 50), dtype=np.int32)\n",
-        "  data = index_update(data, index[1, 0:pos], enc)\n",
+        "  data = index_update(data, index[0, 0:pos], enc)\n",
         "\n",
         "  while pos < max_length:\n",
         "    tmp = tlm(data, params=params, rng=rng)\n",
         "    next_sym = gumbel_sample(tmp[0, pos])\n",
-        "    data = index_update(data, index[1, pos], next_sym)\n",
+        "    data = index_update(data, index[0, pos], next_sym)\n",
         "    pos += 1\n",
         "    if int(next_sym) == 1:\n",
         "      break\n",
@@ -851,4 +851,4 @@
       ]
     }
   ]
-}
\ No newline at end of file
+}

From 0de4615f802ebb7e3369dae8824cefde83cebe17 Mon Sep 17 00:00:00 2001
From: Peng Wang <wangpeng@google.com>
Date: Wed, 31 Jul 2019 21:57:13 -0700
Subject: [PATCH 2248/2720] Turns on --tf_eager and --tf_xla by default

PiperOrigin-RevId: 261059061
---
 tensor2tensor/trax/trainer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index ed29ef542..dfa82bce6 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -46,8 +46,8 @@
                           "Configuration parameters (gin string).")
 flags.DEFINE_integer("log_level", logging.INFO, "Log level.")
 flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
-flags.DEFINE_bool("tf_eager", False, "Whether we're running TF in eager mode.")
-flags.DEFINE_bool("tf_xla", False, "Whether to turn on XLA for TF.")
+flags.DEFINE_bool("tf_eager", True, "Whether we're running TF in eager mode.")
+flags.DEFINE_bool("tf_xla", True, "Whether to turn on XLA for TF.")
 flags.DEFINE_bool("tf_opt_pin_to_host", False, "Whether to turn on TF "
                   "pin-to-host optimization.")
 flags.DEFINE_bool("tf_opt_layout", False, "Whether to turn on TF layout "

From 892de556638c7be095b9cd7311bcae21c577530f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 1 Aug 2019 07:40:47 -0700
Subject: [PATCH 2249/2720] fix documentation. fast decoding for
 dot_product_relative works.

PiperOrigin-RevId: 261119443
---
 tensor2tensor/models/transformer.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 1968d1ca7..429702cd3 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -372,9 +372,7 @@ def _beam_decode(self,
         "dot_product", "dot_product_relative"
     ]):
       # Caching is not guaranteed to work with attention types other than
-      # dot_product.
-      # TODO(petershaw): Support fast decoding when using relative
-      # position representations, i.e. "dot_product_relative" attention.
+      # dot_product and dot_product_relative.
       return self._beam_decode_slow(features, decode_length, beam_size,
                                     top_beams, alpha, use_tpu)
     with tf.variable_scope(self.name):

From 84d111d71f63e5e3be9269783afa5212fe31f0ff Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 1 Aug 2019 10:54:19 -0700
Subject: [PATCH 2250/2720] Support conv3d_transpose and
 conv3d_transpose_with_blocks in MeshTensorflow.

PiperOrigin-RevId: 261154181
---
 tensor2tensor/models/mtf_resnet.py | 85 ++++++++++++------------------
 1 file changed, 35 insertions(+), 50 deletions(-)

diff --git a/tensor2tensor/models/mtf_resnet.py b/tensor2tensor/models/mtf_resnet.py
index 3ffac2d19..794efc1b8 100644
--- a/tensor2tensor/models/mtf_resnet.py
+++ b/tensor2tensor/models/mtf_resnet.py
@@ -79,58 +79,44 @@ def bottleneck_block(inputs,
   """
   shortcut = inputs
 
-  filter_h_dim = mtf.Dimension("filter_height", 3)
-  filter_w_dim = mtf.Dimension("filter_width", 3)
-  one_h_dim = mtf.Dimension("filter_height", 1)
-  one_w_dim = mtf.Dimension("filter_width", 1)
-
   if projection_shortcut is not None:
     filters_dim = mtf.Dimension("filtersp", filters)
-    kernel = mtf.get_variable(
-        inputs.mesh, "kernel", mtf.Shape(
-            [one_h_dim, one_w_dim, inputs.shape.dims[-1], filters_dim]))
-    shortcut = projection_shortcut(inputs, kernel)
+    shortcut = projection_shortcut(inputs, filters_dim)
 
   # First conv block
-  filters1_dim = mtf.Dimension("filters1", filters)
-  kernel1 = mtf.get_variable(
-      inputs.mesh, "kernel1", mtf.Shape(
-          [one_h_dim, one_w_dim, inputs.shape.dims[-1], filters1_dim]))
-  inputs = mtf.conv2d_with_blocks(
+  inputs = mtf.layers.conv2d_with_blocks(
       inputs,
-      kernel1,
-      strides=[1, 1, 1, 1],
+      mtf.Dimension("filters1", filters),
+      filter_size=[1, 1],
+      strides=[1, 1],
       padding="SAME",
-      h_blocks_dim=None, w_blocks_dim=col_blocks_dim)
+      h_blocks_dim=None, w_blocks_dim=col_blocks_dim,
+      name="conv0")
 
   # TODO(nikip): Add Dropout?
   inputs = batch_norm_relu(inputs, is_training)
 
   # Second conv block
-  filters2_dim = mtf.Dimension("filters2", 4*filters)
-  kernel2 = mtf.get_variable(
-      inputs.mesh, "kernel2", mtf.Shape(
-          [filter_h_dim, filter_w_dim, filters1_dim, filters2_dim]))
-  inputs = mtf.conv2d_with_blocks(
+  inputs = mtf.layers.conv2d_with_blocks(
       inputs,
-      kernel2,
-      strides=[1, 1, 1, 1],
+      mtf.Dimension("filters2", 4 * filters),
+      filter_size=[3, 3],
+      strides=[1, 1],
       padding="SAME",
-      h_blocks_dim=row_blocks_dim, w_blocks_dim=col_blocks_dim)
+      h_blocks_dim=row_blocks_dim, w_blocks_dim=col_blocks_dim,
+      name="conv1")
 
   inputs = batch_norm_relu(inputs, is_training)
 
   # Third wide conv filter block
-  filters3_dim = mtf.Dimension("filters3", filters)
-  filters3_kernel = mtf.get_variable(
-      inputs.mesh, "wide_kernel", mtf.Shape(
-          [one_h_dim, one_w_dim, filters2_dim, filters3_dim]))
-  inputs = mtf.conv2d_with_blocks(
+  inputs = mtf.layers.conv2d_with_blocks(
       inputs,
-      filters3_kernel,
-      strides,
+      mtf.Dimension("filters3", filters),
+      filter_size=[1, 1],
+      strides=strides,
       padding="SAME",
-      h_blocks_dim=None, w_blocks_dim=col_blocks_dim)
+      h_blocks_dim=None, w_blocks_dim=col_blocks_dim,
+      name="conv2")
 
   # TODO(nikip): Althought the original resnet code has this batch norm, in our
   # setup this is causing no gradients to be passed. Investigate further.
@@ -170,14 +156,16 @@ def block_layer(inputs,
   """
   with tf.variable_scope(name, default_name="block_layer"):
     # Only the first block per block_layer uses projection_shortcut and strides
-    def projection_shortcut(inputs, kernel):
+    def projection_shortcut(inputs, output_dim):
       """Project identity branch."""
-      inputs = mtf.conv2d_with_blocks(
+      inputs = mtf.layers.conv2d_with_blocks(
           inputs,
-          kernel,
+          output_dim,
+          filter_size=[1, 1],
           strides=strides,
           padding="SAME",
-          h_blocks_dim=None, w_blocks_dim=col_blocks_dim)
+          h_blocks_dim=None, w_blocks_dim=col_blocks_dim,
+          name="shortcut0")
       return batch_norm_relu(
           inputs, is_training, relu=False)
 
@@ -231,9 +219,7 @@ def mtf_model_fn(self, features, mesh):
     # Declare all the dimensions
     batch_dim = mtf.Dimension("batch", hparams.batch_size)
     hidden_dim = mtf.Dimension("hidden", hparams.hidden_size)
-    filter_h_dim = mtf.Dimension("filter_height", 7)
-    filter_w_dim = mtf.Dimension("filter_width", 7)
-    filters = mtf.Dimension("filters", hparams.filter_sizes[0])
+    filter_dim = mtf.Dimension("filters", hparams.filter_sizes[0])
     rows_dim = mtf.Dimension("rows_size", hparams.rows_size)
     cols_dim = mtf.Dimension("cols_size", hparams.cols_size)
     row_blocks_dim = mtf.Dimension("row_blocks", hparams.row_blocks)
@@ -258,15 +244,14 @@ def mtf_model_fn(self, features, mesh):
                           rows_dim, cols_dim, channels_dim])
 
     x = mtf.to_float(x)
-    initial_filters = mtf.get_variable(
-        mesh, "init_filters",
-        mtf.Shape([filter_h_dim, filter_w_dim, channels_dim, filters]))
-    x = mtf.conv2d_with_blocks(
+    x = mtf.layers.conv2d_with_blocks(
         x,
-        initial_filters,
-        strides=[1, 1, 1, 1],
+        filter_dim,
+        filter_size=[3, 3],
+        strides=[1, 1],
         padding="SAME",
-        h_blocks_dim=None, w_blocks_dim=col_blocks_dim)
+        h_blocks_dim=None, w_blocks_dim=col_blocks_dim,
+        name="initial_filter")
 
     x = batch_norm_relu(x, is_training)
 
@@ -280,7 +265,7 @@ def mtf_model_fn(self, features, mesh):
             inputs=x,
             filters=hparams.filter_sizes[0],
             blocks=hparams.layer_sizes[0],
-            strides=[1, 1, 1, 1],
+            strides=[1, 1],
             is_training=is_training,
             name="block_layer1",
             row_blocks_dim=None,
@@ -289,7 +274,7 @@ def mtf_model_fn(self, features, mesh):
             inputs=x,
             filters=hparams.filter_sizes[1],
             blocks=hparams.layer_sizes[1],
-            strides=[1, 1, 1, 1],
+            strides=[1, 1],
             is_training=is_training,
             name="block_layer2",
             row_blocks_dim=None,
@@ -298,7 +283,7 @@ def mtf_model_fn(self, features, mesh):
             inputs=x,
             filters=hparams.filter_sizes[2],
             blocks=hparams.layer_sizes[2],
-            strides=[1, 1, 1, 1],
+            strides=[1, 1],
             is_training=is_training,
             name="block_layer3",
             row_blocks_dim=None,

From 01000fbe1ee82cbb9fee0de4be3bcb4608000ea7 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 1 Aug 2019 13:05:36 -0700
Subject: [PATCH 2251/2720] Add grpcio to install_requires on T2T. This is for
 gRPC.

PiperOrigin-RevId: 261181316
---
 setup.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/setup.py b/setup.py
index f305eab7c..dc1de00a0 100644
--- a/setup.py
+++ b/setup.py
@@ -40,6 +40,7 @@
         'gevent',
         'gin-config',
         'google-api-python-client',
+        'grpcio',
         'gunicorn',
         'gym',
         'h5py',

From ee08f6e48592f8e6b05791405d9875b5a64f861e Mon Sep 17 00:00:00 2001
From: Peng Wang <wangpeng@google.com>
Date: Thu, 1 Aug 2019 13:07:40 -0700
Subject: [PATCH 2252/2720] Internal change

PiperOrigin-RevId: 261181702
---
 tensor2tensor/trax/inputs.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 86ed644ee..67f893d6b 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -323,11 +323,11 @@ def batch_fun(dataset, training, shapes, target_names, n_devices,
                            bucket_length, bucket_length * 2,
                            bucket_length * 4, bucket_length * 8,
                            bucket_length * 16]
-      # We will pad to boundaries which pads to bucket_boundary - 1: add 1 here.
-      bucket_boundaries = [b + 1 for b in bucket_boundaries]
       if not training:
         max_eval_length = max_eval_length or bucket_length * 32
         bucket_boundaries[-1] = max_eval_length
+      # We will pad to boundaries which pads to bucket_boundary - 1: add 1 here.
+      bucket_boundaries = [b + 1 for b in bucket_boundaries]
       bucket_batch_sizes = [cur_batch_size * 4, cur_batch_size * 2,
                             cur_batch_size, cur_batch_size // 2,
                             cur_batch_size // 4, cur_batch_size // 8,

From 0a3da0fbb8da5f6b7d0d4642dc0507637387c26e Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 1 Aug 2019 13:32:14 -0700
Subject: [PATCH 2253/2720] Update TF_LATEST in .travis.yml to 1.14.*

PiperOrigin-RevId: 261186361
---
 .travis.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index dea932d82..f979b8146 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -14,13 +14,13 @@ env:
     - T2T_PROBLEM=algorithmic_reverse_binary40_test
     - T2T_DATA_DIR=/tmp/t2t-data
     - T2T_TRAIN_DIR=/tmp/t2t-train
-    - TF_LATEST="1.13.*"
+    - TF_LATEST="1.14.*"
     # This is necessary to have gsutil work with Python 2.7
     - BOTO_CONFIG=/dev/null
   matrix:
     # We test against the latest stable TensorFlow and tf-nightly.
     # If updating, also update TF_LATEST above
-    - TF_VERSION="1.13.*"
+    - TF_VERSION="1.14.*"
     - TF_VERSION="tf-nightly"
 install:
   - ./oss_scripts/oss_pip_install.sh

From 1df125e1014776b7c5777c6dbc5452cf80b6ad9e Mon Sep 17 00:00:00 2001
From: konradczechowski <konrad.czechowski@gmail.com>
Date: Thu, 1 Aug 2019 22:48:03 +0200
Subject: [PATCH 2254/2720] Rainbow params with larger epsilon. (#1635)

---
 tensor2tensor/models/research/rl.py            |  9 +++++++++
 tensor2tensor/rl/trainer_model_based_params.py | 13 +++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 67b2ae719..68ba22a10 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -437,6 +437,15 @@ def dqn_guess1_rainbow_params():
   return hparams
 
 
+@registry.register_hparams
+def dqn_rainbow_params():
+  """Rainbow params."""
+  hparams = dqn_guess1_params()
+  hparams.set_hparam("agent_type", "Rainbow")
+  hparams.set_hparam("replay_buffer_replay_capacity", int(2e6) + int(1e5))
+  return hparams
+
+
 @registry.register_hparams
 def dqn_2m_replay_buffer_params():
   """Guess 1 for DQN params, 2 milions transitions in replay buffer."""
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index fb0fe1926..32553bb10 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -229,6 +229,19 @@ def rlmb_dqn_guess1_rainbow():
   return hparams
 
 
+@registry.register_hparams
+def rlmb_dqn_rainbow_large_epsilon():
+  """Rainbow rlmb_dqn params."""
+  hparams = rlmb_dqn_guess1()
+  hparams.set_hparam("base_algo_params", "dqn_rainbow_params")
+  hparams.set_hparam("dqn_agent_epsilon_train", 0.1)
+  hparams.add_hparam("real_dqn_agent_epsilon_train", 0.02)
+  simulated_rollout_length = 10
+  hparams.set_hparam("simulated_rollout_length", simulated_rollout_length)
+  hparams.set_hparam("dqn_time_limit", simulated_rollout_length)
+  return hparams
+
+
 @registry.register_hparams
 def rlmb_dqn_guess1_2m_replay_buffer():
   """DQN guess1 params, 2M replay buffer."""

From 0c8a559bf9dc667aa9066f277053e553cfcf9d52 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 1 Aug 2019 13:54:47 -0700
Subject: [PATCH 2255/2720] Give different output directories for train env and
 eval env. We may not want to mix up these trajectories.

PiperOrigin-RevId: 261190975
---
 tensor2tensor/trax/rlax/ppo_main.py | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index d4980f821..7694161ca 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -181,17 +181,26 @@ def main(argv):
   gin.parse_config_files_and_bindings(FLAGS.config_file, gin_configs)
 
   # TODO(pkozakowski): Find a better way to determine this.
+  env_kwargs = {}
+  train_env_kwargs = {}
+  eval_env_kwargs = {}
   if "OnlineTuneEnv" in FLAGS.env_problem_name:
     # TODO(pkozakowski): Separate env output dirs by train/eval and epoch.
-    env_kwargs = {"output_dir": os.path.join(FLAGS.output_dir, "envs")}
-  else:
-    env_kwargs = {}
+    train_env_kwargs = {}
+    train_env_kwargs.update(env_kwargs)
+    train_env_kwargs["output_dir"] = os.path.join(FLAGS.output_dir,
+                                                  "envs/train")
+
+    eval_env_kwargs = {}
+    eval_env_kwargs.update(env_kwargs)
+    eval_env_kwargs["output_dir"] = os.path.join(FLAGS.output_dir,
+                                                 "envs/eval")
 
   # Make an env here.
-  env = make_env(batch_size=FLAGS.batch_size, **env_kwargs)
+  env = make_env(batch_size=FLAGS.batch_size, **train_env_kwargs)
   assert env
 
-  eval_env = make_env(batch_size=FLAGS.eval_batch_size, **env_kwargs)
+  eval_env = make_env(batch_size=FLAGS.eval_batch_size, **eval_env_kwargs)
   assert eval_env
 
   def run_training_loop():

From 50a04804eb4cee4798691edf4a67e9fd72b482f7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 1 Aug 2019 15:47:18 -0700
Subject: [PATCH 2256/2720] Split out the self_attention from the decoder_layer
 in Tensor2Tensor

PiperOrigin-RevId: 261213641
---
 tensor2tensor/models/transformer.py | 67 ++++++++++++++++++++++-------
 1 file changed, 51 insertions(+), 16 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 429702cd3..24e44167a 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1414,22 +1414,20 @@ def transformer_prepare_decoder(targets, hparams, features=None, pad=None):
   return (decoder_input, decoder_self_attention_bias)
 
 
-def transformer_decoder_layer(decoder_input,
-                              decoder_self_attention_bias,
-                              layer_idx,
-                              hparams,
-                              encoder_output=None,
-                              encoder_decoder_attention_bias=None,
-                              cache=None,
-                              decode_loop_step=None,
-                              nonpadding=None,
-                              save_weights_to=None,
-                              make_image_summary=False,
-                              losses=None,
-                              layer_collection=None,
-                              recurrent_memory_by_layer=None,
-                              chunk_number=None):
-  """A single transformer decoder layer."""
+def transformer_self_attention_layer(decoder_input,
+                                     decoder_self_attention_bias,
+                                     layer_idx,
+                                     hparams,
+                                     encoder_output=None,
+                                     encoder_decoder_attention_bias=None,
+                                     cache=None,
+                                     decode_loop_step=None,
+                                     save_weights_to=None,
+                                     make_image_summary=False,
+                                     layer_collection=None,
+                                     recurrent_memory_by_layer=None,
+                                     chunk_number=None):
+  """A single transformer self-attention layer."""
   x = decoder_input
   layer = layer_idx
   layer_name = "layer_%d" % layer
@@ -1528,6 +1526,43 @@ def transformer_decoder_layer(decoder_input,
                 "mode",
                 tf.estimator.ModeKeys.TRAIN) == tf.estimator.ModeKeys.TRAIN))
         x = common_layers.layer_postprocess(x, y, hparams)
+    return x, layer_cache
+
+
+def transformer_decoder_layer(decoder_input,
+                              decoder_self_attention_bias,
+                              layer_idx,
+                              hparams,
+                              encoder_output=None,
+                              encoder_decoder_attention_bias=None,
+                              cache=None,
+                              decode_loop_step=None,
+                              nonpadding=None,
+                              save_weights_to=None,
+                              make_image_summary=False,
+                              losses=None,
+                              layer_collection=None,
+                              recurrent_memory_by_layer=None,
+                              chunk_number=None):
+  """A single transformer decoder layer."""
+  x, layer_cache = transformer_self_attention_layer(
+      decoder_input=decoder_input,
+      decoder_self_attention_bias=decoder_self_attention_bias,
+      layer_idx=layer_idx,
+      hparams=hparams,
+      encoder_output=encoder_output,
+      encoder_decoder_attention_bias=encoder_decoder_attention_bias,
+      cache=cache,
+      decode_loop_step=decode_loop_step,
+      save_weights_to=save_weights_to,
+      make_image_summary=make_image_summary,
+      layer_collection=layer_collection,
+      recurrent_memory_by_layer=recurrent_memory_by_layer,
+      chunk_number=chunk_number)
+
+  layer = layer_idx
+  layer_name = "layer_%d" % layer
+  with tf.variable_scope(layer_name):
     with tf.variable_scope("ffn"):
       y = transformer_ffn_layer(
           common_layers.layer_preprocess(

From c5de299dc238553b7387a25c01357755bd30cf0f Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 2 Aug 2019 11:29:47 -0700
Subject: [PATCH 2257/2720] Separate out utils from env_service_server.

PiperOrigin-RevId: 261358619
---
 tensor2tensor/envs/env_service_server.py | 38 ++++--------------
 tensor2tensor/envs/server_utils.py       | 51 ++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 30 deletions(-)
 create mode 100644 tensor2tensor/envs/server_utils.py

diff --git a/tensor2tensor/envs/env_service_server.py b/tensor2tensor/envs/env_service_server.py
index dd41c02c5..24b268ce9 100644
--- a/tensor2tensor/envs/env_service_server.py
+++ b/tensor2tensor/envs/env_service_server.py
@@ -20,33 +20,29 @@
 from __future__ import print_function
 
 import os
-import time
 from absl import app
 from absl import flags
 from absl import logging
-from concurrent import futures
-import grpc
-from grpc import loas2
 from tensor2tensor.envs import env_problem_utils
-from tensor2tensor.envs import env_service_pb2_grpc
-from tensor2tensor.envs import env_service_servicer
+from tensor2tensor.envs import server_utils
 
 FLAGS = flags.FLAGS
 
 flags.DEFINE_bool("xm", False, "Copy atari roms?")
 flags.DEFINE_integer("env_service_port", 7777, "Port on which to run.")
 flags.DEFINE_string("env_problem_name", None, "Name of the EnvProblem to make.")
-flags.DEFINE_string(
-    "max_timestep", None,
-    "If set to an integer, maximum number of time-steps in a "
-    "trajectory. The bare env is wrapped with TimeLimit wrapper.")
+flags.DEFINE_string("max_timestep",
+                    None,
+                    "If set to an integer, maximum number of time-steps in a "
+                    "trajectory. The bare env is TimeLimit wrapped.")
 flags.DEFINE_boolean("resize", False, "If true, resize the game frame")
 flags.DEFINE_integer("resized_height", 105, "Resized height of the game frame.")
 flags.DEFINE_integer("resized_width", 80, "Resized width of the game frame.")
 flags.DEFINE_string("output_dir", "", "Output dir.")
 flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
 flags.DEFINE_integer("replica", 0, "Basically to append to output_dir")
-flags.DEFINE_bool("clip_rewards", True,
+flags.DEFINE_bool("clip_rewards",
+                  True,
                   "Whether to clip and discretize the rewards.")
 
 # Since we're only dealing with 1 GPU machines here.
@@ -54,24 +50,6 @@
 _ADDRESS_FORMAT = "[::]:{}"
 
 
-def add_port(server):
-  server_credentials = loas2.loas2_server_credentials()
-  return server.add_secure_port(_ADDRESS_FORMAT.format(FLAGS.env_service_port),
-                                server_credentials)
-
-
-def serve(output_dir, env):
-  del output_dir  # may use later.
-  server = grpc.server(futures.ThreadPoolExecutor(max_workers=_MAX_CONCURRENCY))
-  servicer = env_service_servicer.EnvServiceServicer(env)
-  env_service_pb2_grpc.add_EnvServiceServicer_to_server(servicer, server)
-  port = add_port(server)
-  server.start()
-  logging.info("Starting server on port %s", port)
-  while True:
-    time.sleep(60 * 60 * 24)  # sleep for a day only to sleep again.
-
-
 def main(argv):
   del argv
   output_dir = FLAGS.output_dir
@@ -88,7 +66,7 @@ def main(argv):
       clip_rewards=FLAGS.clip_rewards)
 
   logging.info("Replica[%s] is ready to serve requests.", FLAGS.replica)
-  serve(output_dir, env)
+  server_utils.serve(output_dir, env, FLAGS.env_service_port)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/envs/server_utils.py b/tensor2tensor/envs/server_utils.py
new file mode 100644
index 000000000..3df60b1d7
--- /dev/null
+++ b/tensor2tensor/envs/server_utils.py
@@ -0,0 +1,51 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Utilities for env_service_server.py."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import time
+from absl import logging
+from concurrent import futures
+import grpc
+from grpc import loas2
+
+from tensor2tensor.envs import env_service_pb2_grpc
+from tensor2tensor.envs import env_service_servicer
+
+# Since we're only dealing with 1 GPU machines here.
+_MAX_CONCURRENCY = 1
+_ADDRESS_FORMAT = "[::]:{}"
+
+
+def add_port(server, port):
+  server_credentials = loas2.loas2_server_credentials()
+  return server.add_secure_port(
+      _ADDRESS_FORMAT.format(port), server_credentials)
+
+
+def serve(output_dir, env, port):
+  del output_dir  # may use later.
+  server = grpc.server(futures.ThreadPoolExecutor(max_workers=_MAX_CONCURRENCY))
+  servicer = env_service_servicer.EnvServiceServicer(env)
+  env_service_pb2_grpc.add_EnvServiceServicer_to_server(servicer, server)
+  serving_port = add_port(server, port)
+  server.start()
+  logging.info("Starting server on port %s", serving_port)
+  while True:
+    time.sleep(60 * 60 * 24)  # sleep for a day only to sleep again.

From c8ef0d65d40c8e58a8ee055b5a51684b8e48bf5e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 2 Aug 2019 13:19:05 -0700
Subject: [PATCH 2258/2720] Convert TF dtypes from TFDS to Numpy dtypes.

PiperOrigin-RevId: 261378650
---
 tensor2tensor/trax/inputs.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 67f893d6b..5f49f0046 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -84,6 +84,9 @@ def inputs(n_devices, dataset_name, data_dir=None, input_name=None,
    input_name, input_shape, input_dtype) = _train_and_eval_batches(
        dataset_name, data_dir, input_name, n_devices)
 
+  if isinstance(input_dtype, tf.DType):
+    input_dtype = input_dtype.as_numpy_dtype
+
   if input_dtype == np.uint8:  # TPUs don't like uint8s, we cast to ints.
     input_dtype = np.int32
 

From 0bea793e381cc7c2fe72681902a358e40ee0928c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 2 Aug 2019 16:05:14 -0700
Subject: [PATCH 2259/2720] Pass data_dir when loading a tfds dataset.

PiperOrigin-RevId: 261410051
---
 tensor2tensor/trax/inputs.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 5f49f0046..8a60b17a5 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -219,10 +219,10 @@ def train_and_eval_dataset(dataset_name, data_dir, train_shuffle_files=True,
   if tfds.Split.VALIDATION not in splits:
     eval_split = tfds.Split.TEST
   train = tfds.load(
-      name=dataset_name, split=tfds.Split.TRAIN,
+      name=dataset_name, split=tfds.Split.TRAIN, data_dir=data_dir,
       as_dataset_kwargs={"shuffle_files": train_shuffle_files})
   valid = tfds.load(
-      name=dataset_name, split=eval_split,
+      name=dataset_name, split=eval_split, data_dir=data_dir,
       as_dataset_kwargs={"shuffle_files": test_shuffle_files})
   keys = None
   if info.supervised_keys:

From 4b0353fdb203cf9100e3a66dfc283bd707eec1bb Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sun, 4 Aug 2019 14:27:55 -0700
Subject: [PATCH 2260/2720] TRAX's env_service_server.py based on T2T's. This
 is needed to link in and initialize TRAX specific envs and gin config files.

PiperOrigin-RevId: 261588197
---
 .../trax/rlax/envs/env_service_server.py      | 114 ++++++++++++++++++
 1 file changed, 114 insertions(+)
 create mode 100644 tensor2tensor/trax/rlax/envs/env_service_server.py

diff --git a/tensor2tensor/trax/rlax/envs/env_service_server.py b/tensor2tensor/trax/rlax/envs/env_service_server.py
new file mode 100644
index 000000000..4531c317a
--- /dev/null
+++ b/tensor2tensor/trax/rlax/envs/env_service_server.py
@@ -0,0 +1,114 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Server that acts as a remote env.
+
+NOTE: This is a fork from T2T's `env_service_server.py` since we need to
+link in some TRAX specific envs and gin configuration. This also enables
+eager execution.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from absl import app
+from absl import flags
+from absl import logging
+import gin
+from tensor2tensor.envs import env_problem_utils
+from tensor2tensor.envs import server_utils
+from tensor2tensor.rl.google import atari_utils
+from tensor2tensor.trax.rlax import envs  # pylint: disable=unused-import
+import tensorflow as tf
+
+
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_bool("xm", False, "Copy atari roms?")
+flags.DEFINE_integer("env_service_port", 7777, "Port on which to run.")
+flags.DEFINE_string("env_problem_name", None, "Name of the EnvProblem to make.")
+flags.DEFINE_string("max_timestep",
+                    None,
+                    "If set to an integer, maximum number of time-steps in a "
+                    "trajectory. The bare env is TimeLimit wrapped.")
+flags.DEFINE_boolean("resize", False, "If true, resize the game frame")
+flags.DEFINE_integer("resized_height", 105, "Resized height of the game frame.")
+flags.DEFINE_integer("resized_width", 80, "Resized width of the game frame.")
+flags.DEFINE_string("output_dir", "", "Output dir.")
+flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
+flags.DEFINE_integer("replica", 0, "Basically to append to output_dir")
+flags.DEFINE_bool("clip_rewards",
+                  True,
+                  "Whether to clip and discretize the rewards.")
+
+# Gin related flags.
+flags.DEFINE_multi_string("gin_config_file",
+                          None,
+                          "Configuration file with parameters (.gin).")
+flags.DEFINE_multi_string("gin_config_string",
+                          [],
+                          "Configuration parameters (gin string).")
+
+
+# TODO(afrozm): Check this.
+flags.DEFINE_bool("enable_eager_execution", False, "")
+
+# Since we're only dealing with 1 GPU machines here.
+_MAX_CONCURRENCY = 1
+_ADDRESS_FORMAT = "[::]:{}"
+
+
+def initialize_gin():
+  gin_bindings = FLAGS.gin_config_string
+  if not (FLAGS.gin_config_file or gin_bindings):
+    return
+  gin.parse_config_files_and_bindings(FLAGS.gin_config_file, gin_bindings)
+
+
+def main(argv):
+  del argv
+
+  if FLAGS.enable_eager_execution:
+    tf.enable_eager_execution()
+
+  output_dir = FLAGS.output_dir
+
+  # Initialize Gin.
+  initialize_gin()
+
+  output_dir = os.path.join(output_dir, str(FLAGS.replica))
+
+  env_kwargs = {"output_dir": output_dir}
+
+  env = env_problem_utils.make_env(
+      batch_size=1,
+      env_problem_name=FLAGS.env_problem_name,
+      resize=FLAGS.resize,
+      resized_height=FLAGS.resized_height,
+      resized_width=FLAGS.resized_width,
+      max_timestep=FLAGS.max_timestep,
+      clip_rewards=FLAGS.clip_rewards,
+      **env_kwargs)
+
+  logging.info("Replica[%s] is ready to serve requests.", FLAGS.replica)
+  server_utils.serve(output_dir, env, FLAGS.env_service_port)
+
+
+if __name__ == "__main__":
+  app.run(main)

From 62cf66fda3494e815895973dbcd94bb983ae5bc8 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sun, 4 Aug 2019 14:29:34 -0700
Subject: [PATCH 2261/2720] Import T2T envs into PPO. Add flags/code in PPO to
 support construction of ClientEnv.

PiperOrigin-RevId: 261588264
---
 tensor2tensor/trax/rlax/ppo_main.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 7694161ca..4b1249c47 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -49,12 +49,14 @@
 import jax
 from jax.config import config
 import numpy as onp
+from tensor2tensor import envs  # pylint: disable=unused-import
 from tensor2tensor.envs import gym_env_problem
 from tensor2tensor.envs import rendered_env_problem
 from tensor2tensor.rl import gym_utils
+from tensor2tensor.rl.google import atari_utils  # GOOGLE-INTERNAL:
 from tensor2tensor.trax import layers
 from tensor2tensor.trax import models
-from tensor2tensor.trax.rlax import envs  # pylint: disable=unused-import
+from tensor2tensor.trax.rlax import envs as rlax_envs  # pylint: disable=unused-import
 from tensor2tensor.trax.rlax import ppo
 
 
@@ -105,6 +107,11 @@
                      "If true, sets parallelism to number of cpu cores.")
 
 
+# TODO(afrozm): Find a better way to do these configurations.
+flags.DEFINE_string("train_server_bns", "", "Train Server's BNS.")
+flags.DEFINE_string("eval_server_bns", "", "Eval Server's BNS.")
+
+
 def common_layers():
   # TODO(afrozm): Refactor.
   if "NoFrameskip" in FLAGS.env_problem_name:
@@ -196,6 +203,15 @@ def main(argv):
     eval_env_kwargs["output_dir"] = os.path.join(FLAGS.output_dir,
                                                  "envs/eval")
 
+  if "ClientEnv" in FLAGS.env_problem_name:
+    train_env_kwargs["per_env_kwargs"] = [{
+        "remote_env_address": os.path.join(FLAGS.train_server_bns, str(replica))
+    } for replica in range(FLAGS.batch_size)]
+
+    eval_env_kwargs["per_env_kwargs"] = [{
+        "remote_env_address": os.path.join(FLAGS.eval_server_bns, str(replica))
+    } for replica in range(FLAGS.eval_batch_size)]
+
   # Make an env here.
   env = make_env(batch_size=FLAGS.batch_size, **train_env_kwargs)
   assert env

From ea52c76ab411287d4f4867bfdc9e4e55390d58dc Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sun, 4 Aug 2019 14:33:12 -0700
Subject: [PATCH 2262/2720] Gin files for env_server and ppo binaries for
 learning rate tuning.

PiperOrigin-RevId: 261588529
---
 ...params_online_tune_wide_resnet_cifar10.gin | 41 +++++++++++++++++++
 ...params_online_tune_wide_resnet_cifar10.gin | 21 ++++++++++
 2 files changed, 62 insertions(+)
 create mode 100644 tensor2tensor/trax/rlax/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin
 create mode 100644 tensor2tensor/trax/rlax/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin

diff --git a/tensor2tensor/trax/rlax/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rlax/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin
new file mode 100644
index 000000000..4ec0423de
--- /dev/null
+++ b/tensor2tensor/trax/rlax/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin
@@ -0,0 +1,41 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.rlax
+import tensor2tensor.trax.rlax.envs
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size = 32
+batch_fun.bucket_length = 32
+batch_fun.buckets = None
+batch_fun.eval_batch_size = 32
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 'cifar10'
+
+# Parameters for Momentum:
+# ==============================================================================
+Momentum.mass = 0.9
+
+# Parameters for shuffle_and_batch_data:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun = @trax.inputs.cifar10_no_augmentation_preprocess
+
+# Parameters for WideResnet:
+# ==============================================================================
+WideResnet.widen_factor = 2
+WideResnet.n_blocks = 3
+WideResnet.n_output_classes = 10
+
+# Parameters for OnlineTuneEnv:
+# ==============================================================================
+OnlineTuneEnv.inputs = @trax.inputs.inputs
+OnlineTuneEnv.model = @trax.models.WideResnet
+OnlineTuneEnv.optimizer = @trax.optimizers.Momentum
+OnlineTuneEnv.start_lr = 0.01
+OnlineTuneEnv.train_steps = 500
+OnlineTuneEnv.eval_steps = 50
+OnlineTuneEnv.env_steps = 128
diff --git a/tensor2tensor/trax/rlax/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rlax/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
new file mode 100644
index 000000000..b67345b8c
--- /dev/null
+++ b/tensor2tensor/trax/rlax/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
@@ -0,0 +1,21 @@
+import tensor2tensor.trax.rlax
+
+# Parameters for ppo.training_loop:
+# ==============================================================================
+ppo.training_loop.n_optimizer_steps = 30
+ppo.training_loop.boundary = 128
+ppo.training_loop.max_timestep = 128
+ppo.training_loop.max_timestep_eval = 128
+ppo.training_loop.random_seed = 0
+ppo.training_loop.gamma = 0.99
+ppo.training_loop.lambda_ = 0.95
+ppo.training_loop.epsilon = 0.1
+ppo.training_loop.c1 = 1.0
+ppo.training_loop.c2 = 0.01
+ppo.training_loop.eval_every_n = 10
+ppo.training_loop.done_frac_for_policy_save = 0
+ppo.training_loop.enable_early_stopping = True
+ppo.training_loop.n_evals = 1
+ppo.training_loop.len_history_for_policy = 1  # this needs to be bumped up.
+ppo.training_loop.eval_temperatures = (1.0,)
+ppo.training_loop.epochs = 1000

From 25d0f3aecb122e1e240539241fb58f5b4de32399 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 5 Aug 2019 15:43:45 -0700
Subject: [PATCH 2263/2720] Refactor most basic reversible layers from
 transformer_revnet into Trax layers, make names more descriptive.

PiperOrigin-RevId: 261786022
---
 tensor2tensor/trax/layers/__init__.py         |   1 +
 tensor2tensor/trax/layers/reversible.py       | 128 ++++++++
 tensor2tensor/trax/layers/reversible_test.py  |  37 +++
 .../models/research/transformer_revnet.py     | 296 ++++++------------
 4 files changed, 268 insertions(+), 194 deletions(-)
 create mode 100644 tensor2tensor/trax/layers/reversible.py
 create mode 100644 tensor2tensor/trax/layers/reversible_test.py

diff --git a/tensor2tensor/trax/layers/__init__.py b/tensor2tensor/trax/layers/__init__.py
index 5eed08920..f6a956b02 100644
--- a/tensor2tensor/trax/layers/__init__.py
+++ b/tensor2tensor/trax/layers/__init__.py
@@ -30,6 +30,7 @@
 from tensor2tensor.trax.layers.initializers import *
 from tensor2tensor.trax.layers.normalization import *
 from tensor2tensor.trax.layers.pooling import *
+from tensor2tensor.trax.layers.reversible import *
 from tensor2tensor.trax.layers.rnn import *
 
 
diff --git a/tensor2tensor/trax/layers/reversible.py b/tensor2tensor/trax/layers/reversible.py
new file mode 100644
index 000000000..7d25e93b6
--- /dev/null
+++ b/tensor2tensor/trax/layers/reversible.py
@@ -0,0 +1,128 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Implementations of reversible layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import jax
+from tensor2tensor.trax import backend
+from tensor2tensor.trax.layers import base
+from tensor2tensor.trax.layers import combinators as cb
+
+
+class ReversibleLayer(base.Layer):
+  """Reversible Layer."""
+
+  def reverse(self, output, params=(), **kwargs):
+    """Reverse this layer: compute input given output."""
+    raise NotImplementedError
+
+  def reverse_and_grad(self, output, grad, params=(), **kwargs):
+    """Backward pass: computes the inverse of a layer and propagates gradients.
+
+    While you may choose to only implement reverse, some layers implement this
+    function directly as computation may be shared between reversing and
+    computing gradients.
+
+    Args:
+      output: Output activations; can be a (possibly nested) tuple.
+      grad: gradient signal (cotangent) computed based on subsequent layers.
+        The structure and shape must match the output.
+      params: layer parameters
+      **kwargs: kwargs for the layer
+
+    Returns:
+      A tuple (x, x_grad), where x is the reconstructed input and x_grad is the
+      gradient signal for the input.
+    """
+    # Note: jax.vjp does not allow us to use **kwargs in the signature here.
+    def _do_call(x, params, kwargs):
+      return super(ReversibleLayer, self).call(x, params=params, **kwargs)
+
+    reconstructed_x = self.reverse(output, params, **kwargs)
+    _, vjpfun = jax.vjp(_do_call, reconstructed_x, params, kwargs)
+    x_grad = vjpfun(grad)
+    return reconstructed_x, x_grad
+
+  @property
+  def has_custom_grad(self):
+    return True
+
+  def custom_grad(self, inputs, output, ct, params, **kwargs):
+    del inputs
+    _, input_ct = self.reverse_and_grad(output, ct, params, **kwargs)
+    return input_ct
+
+
+class ReversibleSwap(ReversibleLayer, cb.Swap):
+  """Swap the first two element on the stack."""
+
+  def reverse(self, output, params=(), **kwargs):
+    # Swap is its own inverse
+    return self.call(output, params, **kwargs)
+
+
+class ReversibleSerial(ReversibleLayer, cb.Serial):
+  """A reversible version of tl.Serial (requires reversible sub-layers)."""
+
+  def __init__(self, *layers):
+    super(ReversibleSerial, self).__init__(*layers)
+
+    # Note that sublayers has already been flattened to remove nested lists.
+    for i, layer in enumerate(self.sublayers()):
+      if not isinstance(layer, ReversibleLayer):
+        raise ValueError(
+            'Sub-layer {} of ReversibleSerial is not reversible: {}'.format(
+                i, layer))
+
+  def reverse(self, output, params=(), **kwargs):
+    rng = kwargs.pop('rng', None)
+    rngs = (None,) * self._n_layers
+    if rng is not None:
+      rngs = backend.random.split(rng, self._n_layers)
+
+    layer_val = output
+    for layer, p, rng in reversed(zip(self.sublayers(), params, rngs)):
+      layer_val = layer.reverse(layer_val, p, rng=rng, **kwargs)
+
+    return layer_val
+
+  def reverse_and_grad(self, output, ct, params=(), **kwargs):
+    rng = kwargs.pop('rng', None)
+    rngs = (None,) * self._n_layers
+    if rng is not None:
+      rngs = backend.random.split(rng, self._n_layers)
+
+    layer_val = output
+    layer_ct = ct
+    params_ct = []
+    for layer, p, rng in reversed(zip(self.sublayers(), params, rngs)):
+      layer_val, layer_ct = layer.reverse_and_grad(
+          layer_val, layer_ct, p, rng=rng, **kwargs)
+      layer_ct, p_ct, kwargs_ct = layer_ct
+      params_ct.insert(0, p_ct)
+
+    # TODO(kitaev): Handle kwargs_ct properly. However, kwargs generally only
+    # contains the rng, which is non-differentiable.
+    for k in kwargs:
+      if k != 'rng':
+        raise NotImplementedError(
+            'ReversibleSerial does not support differentiation wrt kwargs,'
+            'and the key {} is not known to be non-differentiable.'.format(k))
+
+    return layer_val, (layer_ct, params_ct, kwargs_ct)
diff --git a/tensor2tensor/trax/layers/reversible_test.py b/tensor2tensor/trax/layers/reversible_test.py
new file mode 100644
index 000000000..6f63963b4
--- /dev/null
+++ b/tensor2tensor/trax/layers/reversible_test.py
@@ -0,0 +1,37 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 reversible layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+from tensor2tensor.trax.layers import base
+from tensor2tensor.trax.layers import reversible
+
+
+class ReversibleLayerTest(absltest.TestCase):
+
+  def test_reversible_swap(self):
+    layer = reversible.ReversibleSwap()
+    input_shape = ((2, 3), (3, 3))
+    final_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual(final_shape, input_shape[::-1])
+
+
+if __name__ == '__main__':
+  absltest.main()
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 1dfd3f368..1e2154774 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -95,51 +95,6 @@ def FeedForward(d_model, d_ff, dropout, mode):
   ]
 
 
-class ReversibleLayerMixin(object):
-  """Reversible Layer Mixin."""
-
-  def inverse_and_vjp(self, output, ct, params=(), **kwargs):
-    """Backward pass: computes the inverse of a layer and propagates gradients.
-
-    Args:
-      output: Output activations; can be a (possibly nested) tuple.
-      ct: gradient signal (cotangent) computed based on subsequent layers. If
-          None, no gradients are propagated. Otherwise the structure and shape
-          must match the output.
-      params: layer parameters
-      **kwargs: kwargs for the layer
-
-    Returns:
-      A tuple (x, x_ct), where x is the reconstructed input and x_ct is the
-      gradient signal for the input. If ct is None, x_ct will also be None.
-    """
-    if ct is None:
-      # Subclasses must override inverse_and_vjp, but in the case where ct is
-      # not None there is an unoptimized implementation below that they can
-      # delegate to.
-      raise NotImplementedError
-
-    # Note: jax.vjp does not allow us to use **kwargs in the signature here.
-    def _do_call(x, params, kwargs):
-      return super(ReversibleLayerMixin, self).call(x, params=params, **kwargs)
-
-    reconstructed_x, must_be_none = self.inverse_and_vjp(
-        output, None, params, **kwargs)
-    assert must_be_none is None
-    _, vjpfun = jax.vjp(_do_call, reconstructed_x, params, kwargs)
-    input_ct = vjpfun(ct)
-    return reconstructed_x, input_ct
-
-  @property
-  def has_custom_grad(self):
-    return True
-
-  def custom_grad(self, inputs, output, ct, params, **kwargs):
-    del inputs
-    _, input_ct = self.inverse_and_vjp(output, ct, params, **kwargs)
-    return input_ct
-
-
 class Split(tl.Layer):
   """Splits the input into sections along an axis."""
 
@@ -184,7 +139,7 @@ def Unchunk(x, params, n_sections=2, **kwargs):
       ) + x.shape[2:])
 
 
-class ReversibleHalfResidual(ReversibleLayerMixin, tl.Serial):
+class ReversibleHalfResidual(tl.ReversibleLayer, tl.Serial):
   """Half of a RevNet-style residual (only updates part of the hidden state)."""
 
   def __init__(self, residual_layers):
@@ -204,36 +159,40 @@ def __init__(self, residual_layers):
     self.subtract_top = tl.Parallel(tl.SubtractTop(), [])
     self.reverse_layers = [self.compute_residual, self.subtract_top]
 
-  def inverse_and_vjp(self, output, ct, params=(), **kwargs):
+  def reverse(self, output, params=(), **kwargs):
+    reconstructed_x = output
     rng = kwargs.pop('rng', None)
     rngs = (None,) * self._n_layers
     if rng is not None:
       rngs = backend.random.split(rng, self._n_layers)
+    # Note that self.sublayers() aligns exactly with self.reverse_layers in
+    # terms of parameter and rng usage, so no re-ordering is required.
+    for layer, p, rng in zip(self.reverse_layers, params, rngs):
+      reconstructed_x = layer(reconstructed_x, p, rng=rng, **kwargs)
+    return reconstructed_x
 
-    if ct is None:
-      reconstructed_x = output
-      # Note that self.sublayers() aligns exactly with self.reverse_layers in
-      # terms of parameter and rng usage, so no re-ordering is required.
-      for layer, p, rng in zip(self.reverse_layers, params, rngs):
-        reconstructed_x = layer(reconstructed_x, p, rng=rng, **kwargs)
-      return reconstructed_x, None
-    else:
-      # Note: jax.vjp does not allow us to use **kwargs in the signature here.
-      def call_compute_residual(x, params, kwargs):
-        return self.compute_residual(x, params, **kwargs)
+  def reverse_and_grad(self, output, ct, params=(), **kwargs):
+    rng = kwargs.pop('rng', None)
+    rngs = (None,) * self._n_layers
+    if rng is not None:
+      rngs = backend.random.split(rng, self._n_layers)
 
-      assert len(ct) == 2
-      ct = ((ct[0], ct[0], ct[1]))
+    # Note: jax.vjp does not allow us to use **kwargs in the signature here.
+    def call_compute_residual(x, params, kwargs):
+      return self.compute_residual(x, params, **kwargs)
+
+    assert len(ct) == 2
+    ct = ((ct[0], ct[0], ct[1]))
 
-      compute_residual_kwargs = kwargs.copy()
-      compute_residual_kwargs['rng'] = rngs[0]
-      stack_with_residual, vjpfun = jax.vjp(
-          call_compute_residual, output, params[0], compute_residual_kwargs)
-      reconstructed_x = self.subtract_top(
-          stack_with_residual, params[-1], rng=rngs[-1], **kwargs)
+    compute_residual_kwargs = kwargs.copy()
+    compute_residual_kwargs['rng'] = rngs[0]
+    stack_with_residual, vjpfun = jax.vjp(
+        call_compute_residual, output, params[0], compute_residual_kwargs)
+    reconstructed_x = self.subtract_top(
+        stack_with_residual, params[-1], rng=rngs[-1], **kwargs)
 
-      x_ct, residual_params_ct, kwargs_ct = vjpfun(ct)
-      return reconstructed_x, (x_ct, (residual_params_ct, ()), kwargs_ct)
+    x_ct, residual_params_ct, kwargs_ct = vjpfun(ct)
+    return reconstructed_x, (x_ct, (residual_params_ct, ()), kwargs_ct)
 
 
 @tl.layer(n_inputs=1, n_outputs=1)
@@ -439,7 +398,7 @@ def body_fun(vals):
       return final_vals[1], final_vals[2:]
 
 
-class ReversibleAttentionHalfResidual(ReversibleLayerMixin, tl.Serial):
+class ReversibleAttentionHalfResidual(tl.ReversibleLayer, tl.Serial):
   """Half of a RevNet-style residual that performs attention.
 
   If inputs are (x1, x2), then outputs are (x1 + z, x2) where:
@@ -483,139 +442,88 @@ def __init__(self, pre_attention, attention, post_attention):
         self.subtract_top,
     ]
 
-  def inverse_and_vjp(self, output, ct, params=(), **kwargs):
+  def reverse(self, output, params=(), **kwargs):
     rng = kwargs.pop('rng', None)
     rngs = (None,) * self._n_layers
     if rng is not None:
       rngs = backend.random.split(rng, self._n_layers)
 
-    if ct is None:
-      reconstructed_x = output
-      # Note that self.sublayers() aligns exactly with self.reverse_layers in
-      # terms of parameter and rng usage, so no re-ordering is required.
-      for layer, p, rng in zip(self.reverse_layers, params, rngs):
-        reconstructed_x = layer(reconstructed_x, p, rng=rng, **kwargs)
-      return reconstructed_x, None
-    else:
-      # Forward pass through self.pre_attention, while preparing for
-      # later backprop.
-      # Note: jax.vjp does not allow us to use **kwargs in the signature here.
-      def call_pre_attention(x, params, kwargs):
-        return self.pre_attention(x, params, **kwargs)
-      pre_attention_kwargs = kwargs.copy()
-      pre_attention_kwargs['rng'] = rngs[0]
-      stack, pre_attention_vjpfun = jax.vjp(
-          call_pre_attention, output, params[0], pre_attention_kwargs)
-
-      # Backprop through adding the residual
-      assert len(ct) == 2
-      ct = saved_ct = (ct[0], ct[0], ct[1])
-
-      # Backprop through self.post_attention with respect to the inputs only
-      call_post_attention_kwargs = kwargs.copy()
-      call_post_attention_kwargs['rng'] = rngs[2]
-      def call_post_attention(x):
-        return self.post_attention(x, params[2], **call_post_attention_kwargs)
-      # Note: these are *not* the actual inputs to self.post_attention.
-      # If self.post_attention is not linear, we will get incorrect gradients.
-      dummy_inputs = (stack[-3], stack[-2], stack[-1])
-      _, post_attention_vjpfun = jax.vjp(call_post_attention, dummy_inputs)
-      (ct,) = post_attention_vjpfun(ct)
-
-      # Simultaneous forward pass and backprop through the attention mechanism
-      attention_kwargs = kwargs.copy()
-      attention_kwargs['rng'] = rngs[1]
-      stack, ct = self.attention.forward_and_vjp(
-          stack, ct, **attention_kwargs)
-      attention_params_ct = ()
-
-      # Backprop through self.pre_attention
-      (x_ct,
-       pre_attention_params_ct,
-       pre_attention_kwargs_ct) = pre_attention_vjpfun(ct)
-
-      # Forward pass for self.post_attention, and backprop with respect to the
-      # parameters only
-      def call_post_attention2(params, kwargs):
-        return self.post_attention(stack, params, **kwargs)
-      stack, post_attention_vjpfun = jax.vjp(
-          call_post_attention2, params[2], call_post_attention_kwargs)
-      (post_attention_params_ct,
-       post_attention_kwargs_ct) = post_attention_vjpfun(saved_ct)
-
-      # Forward pass through subtracting the residual
-      reconstructed_x = self.subtract_top(
-          stack, params[-1], rng=rngs[-1], **kwargs)
-
-      params_ct = (
-          pre_attention_params_ct,
-          attention_params_ct,
-          post_attention_params_ct,
-          (),
-          )
-
-      # We don't actually backprop through the kwargs, but the API requires that
-      # we provide a value for kwargs_ct.
-      kwargs_ct = pre_attention_kwargs_ct
-      del post_attention_kwargs_ct
-
-      return reconstructed_x, (x_ct, params_ct, kwargs_ct)
-
-
-class ReversibleSwap(ReversibleLayerMixin, tl.Swap):
-  """Swap the first two element on the stack."""
-
-  def inverse_and_vjp(self, output, ct, params=(), **kwargs):
-    if ct is None:
-      # Swap is its own inverse
-      return self.call(output, params, **kwargs), None
-    else:
-      return super(ReversibleSwap, self).inverse_and_vjp(
-          output, ct, params, **kwargs)
-
-
-class ReversibleSerial(ReversibleLayerMixin, tl.Serial):
-  """A reversible version of tl.Serial (requires reversible sub-layers)."""
-
-  def __init__(self, *layers):
-    super(ReversibleSerial, self).__init__(*layers)
+    reconstructed_x = output
+    # Note that self.sublayers() aligns exactly with self.reverse_layers in
+    # terms of parameter and rng usage, so no re-ordering is required.
+    for layer, p, rng in zip(self.reverse_layers, params, rngs):
+      reconstructed_x = layer.reverse(reconstructed_x, p, rng=rng, **kwargs)
+    return reconstructed_x
 
-    # Note that sublayers has already been flattened to remove nested lists.
-    for i, layer in enumerate(self.sublayers()):
-      if not isinstance(layer, ReversibleLayerMixin):
-        raise ValueError(
-            'Sub-layer {} of ReversibleSerial is not reversible: {}'.format(
-                i, layer))
-
-  def inverse_and_vjp(self, output, ct, params=(), **kwargs):
+  def reverse_and_grad(self, output, ct, params=(), **kwargs):
     rng = kwargs.pop('rng', None)
     rngs = (None,) * self._n_layers
     if rng is not None:
       rngs = backend.random.split(rng, self._n_layers)
 
-    layer_val = output
-    if ct is not None:
-      layer_ct = ct
-      params_ct = []
-    for layer, p, rng in reversed(zip(self.sublayers(), params, rngs)):
-      layer_val, layer_ct = layer.inverse_and_vjp(
-          layer_val, layer_ct, p, rng=rng, **kwargs)
-      if ct is not None:
-        layer_ct, p_ct, kwargs_ct = layer_ct
-        params_ct.insert(0, p_ct)
-
-    # TODO(kitaev): Handle kwargs_ct properly. However, kwargs generally only
-    # contains the rng, which is non-differentiable.
-    for k in kwargs:
-      if k != 'rng':
-        raise NotImplementedError(
-            'ReversibleSerial does not support differentiation wrt kwargs,'
-            'and the key {} is not known to be non-differentiable.'.format(k))
-
-    if ct is not None:
-      return layer_val, (layer_ct, params_ct, kwargs_ct)
-    else:
-      return layer_val, None
+    # Forward pass through self.pre_attention, while preparing for
+    # later backprop.
+    # Note: jax.vjp does not allow us to use **kwargs in the signature here.
+    def call_pre_attention(x, params, kwargs):
+      return self.pre_attention(x, params, **kwargs)
+    pre_attention_kwargs = kwargs.copy()
+    pre_attention_kwargs['rng'] = rngs[0]
+    stack, pre_attention_vjpfun = jax.vjp(
+        call_pre_attention, output, params[0], pre_attention_kwargs)
+
+    # Backprop through adding the residual
+    assert len(ct) == 2
+    ct = saved_ct = (ct[0], ct[0], ct[1])
+
+    # Backprop through self.post_attention with respect to the inputs only
+    call_post_attention_kwargs = kwargs.copy()
+    call_post_attention_kwargs['rng'] = rngs[2]
+    def call_post_attention(x):
+      return self.post_attention(x, params[2], **call_post_attention_kwargs)
+    # Note: these are *not* the actual inputs to self.post_attention.
+    # If self.post_attention is not linear, we will get incorrect gradients.
+    dummy_inputs = (stack[-3], stack[-2], stack[-1])
+    _, post_attention_vjpfun = jax.vjp(call_post_attention, dummy_inputs)
+    (ct,) = post_attention_vjpfun(ct)
+
+    # Simultaneous forward pass and backprop through the attention mechanism
+    attention_kwargs = kwargs.copy()
+    attention_kwargs['rng'] = rngs[1]
+    stack, ct = self.attention.forward_and_vjp(
+        stack, ct, **attention_kwargs)
+    attention_params_ct = ()
+
+    # Backprop through self.pre_attention
+    (x_ct,
+     pre_attention_params_ct,
+     pre_attention_kwargs_ct) = pre_attention_vjpfun(ct)
+
+    # Forward pass for self.post_attention, and backprop with respect to the
+    # parameters only
+    def call_post_attention2(params, kwargs):
+      return self.post_attention(stack, params, **kwargs)
+    stack, post_attention_vjpfun = jax.vjp(
+        call_post_attention2, params[2], call_post_attention_kwargs)
+    (post_attention_params_ct,
+     post_attention_kwargs_ct) = post_attention_vjpfun(saved_ct)
+
+    # Forward pass through subtracting the residual
+    reconstructed_x = self.subtract_top(
+        stack, params[-1], rng=rngs[-1], **kwargs)
+
+    params_ct = (
+        pre_attention_params_ct,
+        attention_params_ct,
+        post_attention_params_ct,
+        (),
+        )
+
+    # We don't actually backprop through the kwargs, but the API requires that
+    # we provide a value for kwargs_ct.
+    kwargs_ct = pre_attention_kwargs_ct
+    del post_attention_kwargs_ct
+
+    return reconstructed_x, (x_ct, params_ct, kwargs_ct)
 
 
 def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
@@ -671,9 +579,9 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
   ]
   return [
       ReversibleAttentionHalfResidual(pre_attention, attention, post_attention),
-      ReversibleSwap(),
+      tl.ReversibleSwap(),
       ReversibleHalfResidual(feed_forward),
-      ReversibleSwap(),
+      tl.ReversibleSwap(),
   ]
 
 
@@ -721,7 +629,7 @@ def TransformerRevnetLM(vocab_size,
       tl.ShiftRight(),
       positional_embedder,
       tl.Dup(),
-      ReversibleSerial([
+      tl.ReversibleSerial([
           # pylint: disable=g-complex-comprehension
           DecoderBlock(d_model, d_ff,
                        d_attention_key, d_attention_value, n_heads,

From 3931562581df1cccb919e4ffb8ae4da1edc54e4e Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 5 Aug 2019 15:54:37 -0700
Subject: [PATCH 2264/2720] Factor out causal attention.

PiperOrigin-RevId: 261788035
---
 tensor2tensor/trax/layers/attention.py   | 25 +++++++++++++++++++++++-
 tensor2tensor/trax/models/transformer.py | 10 ++--------
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 793684f0b..4828939c8 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -186,7 +186,7 @@ def Attention(d_feature, n_heads=1, dropout=0.0, mode='train'):
     mode: str: 'train' or 'eval'
 
   Returns:
-    Multi-headed self-attention layer.
+    Multi-headed self-attention result and the mask.
   """
   return [
       cb.Dup(), cb.Dup(),
@@ -194,6 +194,29 @@ def Attention(d_feature, n_heads=1, dropout=0.0, mode='train'):
   ]
 
 
+def CausalAttention(d_feature, n_heads=1, dropout=0.0, mode='train'):
+  """Transformer-style multi-headed causal attention.
+
+  # TODO(jonni,lukaszkaiser): standardize and improve layer comments.
+  Accepts inputs of the form x and constructs (q, k, v) and causal mask from x.
+
+  Args:
+    d_feature: int:  dimensionality of feature embedding
+    n_heads: int: number of attention heads
+    dropout: float: dropout rate
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    Multi-headed self-attention result.
+  """
+  return [
+      cb.Dup(),
+      cb.Parallel([], CausalMask(axis=-2)),  # pylint: disable=no-value-for-parameter
+      Attention(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+      cb.Parallel([], cb.Drop()),  # x
+  ]
+
+
 @base.layer()
 def ShiftRight(x, **unused_kwargs):
   """Layer to shift the tensor to the right by padding on axis 1."""
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 0839a6c1f..63b673c34 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -128,10 +128,7 @@ def DecoderBlock(d_model, d_ff, n_heads, dropout, mode):
   """
   self_attention = [
       tl.LayerNorm(),  # vec
-      tl.Dup(),  # vec vec
-      tl.Parallel([], tl.CausalMask(axis=-2)),  # vec mask
-      tl.Attention(d_model, n_heads=n_heads, dropout=dropout, mode=mode),
-      tl.Parallel([], tl.Drop()),  # vec
+      tl.CausalAttention(d_model, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Dropout(rate=dropout, mode=mode),  # vec
   ]
   feed_forward = [
@@ -205,10 +202,7 @@ def EncoderDecoder(d_model, d_ff, n_heads, dropout, mode):
   """
   decoder_self_attention = [                    #        vecs_d   pmask vecs_e
       tl.LayerNorm(),                           #        vecs_d   ..... ......
-      tl.Dup(),                                 # vecs_d vecs_d   ..... ......
-      tl.Parallel([], tl.CausalMask(axis=-2)),  # ______ masks    ..... ......
-      tl.Attention(d_model, n_heads=n_heads, dropout=dropout, mode=mode),
-      tl.Parallel([], tl.Drop()),               # ______   0      ..... ......
+      tl.CausalAttention(d_model, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Dropout(rate=dropout, mode=mode),      # vecs_d          ..... ......
   ]
   decoder_to_encoder_attention = [        # vecs_d        masks         vecs_e

From 3ffd612f189f17d8beb22089409e245be5bcf129 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 6 Aug 2019 09:57:37 -0700
Subject: [PATCH 2265/2720] Override preprocess_example method of
 function_docstring problem to add 'embed_code' attribute to a sample. The
 attribute indicates if the sample is code or docstring during inference.

PiperOrigin-RevId: 261936204
---
 tensor2tensor/data_generators/function_docstring.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
index 4fceebe8e..04ddbc8ad 100644
--- a/tensor2tensor/data_generators/function_docstring.py
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -98,6 +98,11 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
                 "embed_code": [0],
             }
 
+  def preprocess_example(self, example, mode, unused_hparams):
+    if mode != tf.estimator.ModeKeys.TRAIN:
+      example["embed_code"] = [0]
+    return example
+
   def eval_metrics(self):
     return [
         metrics.Metrics.ACC

From 43be271c8a3fa06cb06b5147f044cbdc8bb77535 Mon Sep 17 00:00:00 2001
From: Peng Wang <wangpeng@google.com>
Date: Tue, 6 Aug 2019 11:29:21 -0700
Subject: [PATCH 2266/2720] Provides `backend.eval_on_shapes`, a function
 transformer which runs the function with only shapes and dtypes.

PiperOrigin-RevId: 261958288
---
 tensor2tensor/trax/backend.py     | 77 +++++++++++++++++++++++++++++++
 tensor2tensor/trax/layers/base.py | 76 +++++-------------------------
 2 files changed, 89 insertions(+), 64 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index a78381abd..a9bdb0cc9 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -20,11 +20,13 @@
 from __future__ import print_function
 
 import contextlib
+import functools
 import gin
 
 import jax
 from jax import lax
 from jax import random as jax_random
+from jax.interpreters import partial_eval as pe
 import jax.numpy as jnp
 import jax.scipy.special as jax_special
 import numpy as onp
@@ -94,6 +96,76 @@ def jax_avg_pool(x, pool_size, strides, padding):
                           pool_size, strides=strides, padding=padding)
 
 
+def nested_map(x, f):
+  """Map the function f to the nested structure x (dicts, tuples, lists)."""
+  if isinstance(x, list):
+    return [nested_map(y, f) for y in x]
+  if isinstance(x, tuple):
+    return tuple([nested_map(y, f) for y in x])
+  return f(x)
+
+
+class ShapeType(object):
+  """Store shape and type."""
+
+  def __init__(self, shape, dtype):
+    self.shape = shape
+    self.dtype = dtype
+
+  def __repr__(self):
+    return "[shape:" + str(self.shape) + ", dtype:" + str(self.dtype) + "]"
+
+
+# TODO(lukaszkaiser): remove this function once JAX has an analogue.
+# pylint: disable=missing-docstring
+def _jax_eval_on_shapes(f, *args):
+  def abstractify(x):
+    return jax.abstract_arrays.raise_to_shaped(jax.core.get_aval(x))
+
+  def make_array(arg):
+    return jnp.zeros(shape=arg.shape, dtype=arg.dtype)
+
+  def turn_back_into_pytree(x):
+    if isinstance(x, jax.core.JaxTuple):
+      return tuple([turn_back_into_pytree(y) for y in x])
+    return x
+
+  def get_shapes_and_types(x):
+    if isinstance(x, jax.core.AbstractTuple):
+      return tuple([get_shapes_and_types(y) for y in x])
+    return ShapeType(x.shape, x.dtype)
+
+  def f_jaxtuple(*jaxtuple_args):
+    args = map(turn_back_into_pytree, jaxtuple_args)
+    out = f(*args)
+    res, _ = jax.api_util.pytree_to_jaxtupletree(out)
+    return res
+
+  args_arrays = nested_map(args, make_array)
+  jaxtuple_args, _ = jax.util.unzip2(
+      map(jax.api_util.pytree_to_jaxtupletree, args_arrays))
+  res = pe.abstract_eval_fun(f_jaxtuple, *map(abstractify, jaxtuple_args))
+
+  return get_shapes_and_types(res)
+
+
+def jax_eval_on_shapes(f):
+  """Returns a function that evaluates `f` given input shapes and dtypes.
+
+  It transforms function `f` to a function that performs the same computation as
+  `f` but only on shapes and dtypes (a.k.a. shape inference).
+
+  Args:
+    f: the function to be transformed.
+
+  Returns:
+    A function whose input arguments can be either the same as `f`'s or only
+    their shapes/dtypes represented by `ShapeType`, and whose return values are
+    `ShapeType`s with the same nested structure as `f`'s return values.
+  """
+  return functools.partial(_jax_eval_on_shapes, f)
+
+
 _JAX_BACKEND = {
     "name": "jax",
     "np": jnp,
@@ -105,6 +177,7 @@ def jax_avg_pool(x, pool_size, strides, padding):
     "jit": jax.jit,
     "grad": jax.grad,
     "pmap": jax.pmap,
+    "eval_on_shapes": jax_eval_on_shapes,
     "random_uniform": jax_random.uniform,
     "random_normal": jax_random.normal,
     "random_bernoulli": jax_random.bernoulli,
@@ -159,6 +232,10 @@ def pmap(*args, **kwargs):
   return backend()["pmap"](*args, **kwargs)
 
 
+def eval_on_shapes(*args, **kwargs):
+  return backend()["eval_on_shapes"](*args, **kwargs)
+
+
 def dataset_as_numpy(*args, **kwargs):
   return backend()["dataset_as_numpy"](*args, **kwargs)
 
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 503230588..0d76ef642 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -23,10 +23,11 @@
 import traceback
 
 import jax
-from jax.interpreters import partial_eval as pe
 
 import numpy as onp
 from tensor2tensor.trax import backend
+from tensor2tensor.trax.backend import nested_map
+from tensor2tensor.trax.backend import ShapeType
 
 
 class Layer(object):
@@ -178,16 +179,16 @@ def pseudo_call(self, pseudo_inputs, params):
       layer has more than one output).
     """
     try:
-      with backend.use_backend('jax'):
-        # Beware: using an actual RNG (as opposed to this ShapeType stub) would
-        # cause a large number of dropout masks to be computed and permanently
-        # stored in global memory.
-        rng = ShapeType(shape=(2,), dtype=onp.uint32)
-        def call_on_input(x, params, rng):
-          return self.call(x, params=params, rng=rng)
-        params_shapes = nested_map(
-            params, lambda x: ShapeType(shape=x.shape, dtype=x.dtype))
-        s = _eval_on_shapes(call_on_input, pseudo_inputs, params_shapes, rng)
+      # Beware: using an actual RNG (as opposed to this ShapeType stub) would
+      # cause a large number of dropout masks to be computed and permanently
+      # stored in global memory.
+      rng = ShapeType(shape=(2,), dtype=onp.uint32)
+      def call_on_input(x, params, rng):
+        return self.call(x, params=params, rng=rng)
+      params_shapes = nested_map(
+          params, lambda x: ShapeType(shape=x.shape, dtype=x.dtype))
+      s = backend.eval_on_shapes(call_on_input)(pseudo_inputs,
+                                                params_shapes, rng)
       return s
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback(skip=3)
@@ -264,17 +265,6 @@ def vjpfun(grad):
       raise LayerError(name, 'call', self._caller, shapes(x), trace)
 
 
-class ShapeType(object):
-  """Store shape and type."""
-
-  def __init__(self, shape, dtype):
-    self.shape = shape
-    self.dtype = dtype
-
-  def __repr__(self):
-    return '[shape:' + str(self.shape) + ', dtype:' + str(self.dtype) + ']'
-
-
 class LayerError(Exception):
   """Exception raised in the layer stack.
 
@@ -302,39 +292,6 @@ def message(self):
     return prefix + caller + shapes_str + self._traceback
 
 
-# TODO(lukaszkaiser): remove this function once JAX has an analogue.
-def _eval_on_shapes(f, *args):
-  """Evaluates f given only shapes and types."""
-  def abstractify(x):
-    return jax.abstract_arrays.raise_to_shaped(jax.core.get_aval(x))
-
-  def make_array(arg):
-    return backend.numpy.zeros(shape=arg.shape, dtype=arg.dtype)
-
-  def turn_back_into_pytree(x):
-    if isinstance(x, jax.core.JaxTuple):
-      return tuple([turn_back_into_pytree(y) for y in x])
-    return x
-
-  def get_shapes_and_types(x):
-    if isinstance(x, jax.core.AbstractTuple):
-      return tuple([get_shapes_and_types(y) for y in x])
-    return ShapeType(x.shape, x.dtype)
-
-  def f_jaxtuple(*jaxtuple_args):
-    args = map(turn_back_into_pytree, jaxtuple_args)
-    out = f(*args)
-    res, _ = jax.api_util.pytree_to_jaxtupletree(out)
-    return res
-
-  args_arrays = nested_map(args, make_array)
-  jaxtuple_args, _ = jax.util.unzip2(
-      map(jax.api_util.pytree_to_jaxtupletree, args_arrays))
-  res = pe.abstract_eval_fun(f_jaxtuple, *map(abstractify, jaxtuple_args))
-
-  return get_shapes_and_types(res)
-
-
 def _apply_to_first_n(f, x, n):
   """Helper: apply f to first n elements on the stack x if n > 0."""
   if n < 1:
@@ -353,15 +310,6 @@ def _apply_to_first_n(f, x, n):
   return result
 
 
-def nested_map(x, f):
-  """Map the function f to the nested structure x (dicts, tuples, lists)."""
-  if isinstance(x, list):
-    return [nested_map(y, f) for y in x]
-  if isinstance(x, tuple):
-    return tuple([nested_map(y, f) for y in x])
-  return f(x)
-
-
 def nested_reduce(x, f):
   """Fold the function f to the nested structure x (dicts, tuples, lists)."""
   if isinstance(x, list):

From 492691ec4fe93c637dc31f8343f928b29fd3ff7c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 7 Aug 2019 09:54:36 -0700
Subject: [PATCH 2267/2720] Fix GPU memory contention between jax and enabling
 tf_eager

PiperOrigin-RevId: 262156649
---
 tensor2tensor/trax/trainer.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index dfa82bce6..1eb71b950 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -28,6 +28,7 @@
 
 import gin
 import jax
+from tensor2tensor.trax import backend
 from tensor2tensor.trax import trax
 
 import tensorflow as tf
@@ -109,6 +110,12 @@ def main(_):
 
   _setup_gin()
 
+  if FLAGS.tf_eager and backend.get_name() in ("numpy", "jax"):
+    # Numpy backend doesn't benefit from having the input pipeline run on GPU,
+    # and jax backend has GPU memory contention if TF uses the GPU. Gin must be
+    # set up first before determining the backend.
+    tf.config.experimental.set_visible_devices([], "GPU")
+
   # Setup output directory
   output_dir = FLAGS.output_dir or _default_output_dir()
   trax.log("Using --output_dir %s" % output_dir)

From 673ca5096df85963f1d858218dd426a7290a6925 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 7 Aug 2019 13:30:50 -0700
Subject: [PATCH 2268/2720] Make attention type configurable via gin

PiperOrigin-RevId: 262202863
---
 .../transformer_revnet_imagenet64_8gb.gin     | 13 ++++++++++-
 tensor2tensor/trax/models/__init__.py         |  5 ++++
 .../models/research/transformer_revnet.py     | 23 ++++++++-----------
 3 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
index dd4e1f47e..0e158baf9 100644
--- a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
@@ -32,6 +32,16 @@ train.optimizer = @trax.optimizers.Adafactor
 train.train_steps = 500000
 train.trainer_class = @MemoryEfficientTrainer
 
+# Parameters for DotProductAttention:
+# ==============================================================================
+DotProductAttention.dropout = 0.0
+MemoryEfficientDotProductAttention.loop_stride = 512
+
+# Parameters for MemoryEfficientDotProductAttention:
+# ==============================================================================
+MemoryEfficientDotProductAttention.dropout = 0.0
+MemoryEfficientDotProductAttention.loop_stride = 512
+
 # Parameters for TransformerRevnetLM:
 # ==============================================================================
 TransformerRevnetLM.d_model = 1024
@@ -46,4 +56,5 @@ TransformerRevnetLM.n_layers = 6
 TransformerRevnetLM.vocab_size = 256
 TransformerRevnetLM.n_chunks = 16
 TransformerRevnetLM.n_attention_chunks = 1
-TransformerRevnetLM.attention_loop_stride = 512
+TransformerRevnetLM.attention_type = \
+    @trax.models.MemoryEfficientDotProductAttention
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index 8b363f27b..1b6b281b4 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -47,3 +47,8 @@ def model_configure(*args, **kwargs):
 TransformerLM = model_configure(transformer.TransformerLM)
 TransformerRevnetLM = model_configure(transformer_revnet.TransformerRevnetLM)
 WideResnet = model_configure(resnet.WideResnet)
+
+DotProductAttention = model_configure(
+    transformer_revnet.DotProductAttention, blacklist=["mode"])
+MemoryEfficientDotProductAttention = model_configure(
+    transformer_revnet.MemoryEfficientDotProductAttention, blacklist=["mode"])
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 1e2154774..4357663a3 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -283,6 +283,9 @@ class MemoryEfficientDotProductAttention(DotProductAttention):
   def __init__(self, loop_stride, dropout, mode):
     super(MemoryEfficientDotProductAttention, self).__init__(dropout, mode)
     self._loop_stride = loop_stride
+    # TODO(kitaev): implement attention dropout
+    assert dropout is None or dropout == 0.0, (
+        'Dropout is not implemented in MemoryEfficientDotProductAttention.')
 
   def call(self, inputs, params=(), **kwargs):
     output, _ = self.forward_and_vjp(inputs, None, params=params, **kwargs)
@@ -527,7 +530,7 @@ def call_post_attention2(params, kwargs):
 
 
 def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
-                 n_heads, n_attention_chunks, attention_loop_stride,
+                 n_heads, n_attention_chunks, attention_type,
                  dropout, mode):
   """Reversible transformer decoder layer.
 
@@ -538,8 +541,7 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
     d_attention_value: int: depth of value vector for each attention head
     n_heads: int: number of attention heads
     n_attention_chunks: int: number of chunks for attention
-    attention_loop_stride: int: number of query elements to compute attention
-      for in parallel. Set to 0 to disable memory-efficient attention.
+    attention_type: class: attention class to use, such as DotProductAttention.
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
 
@@ -558,13 +560,7 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
       ),
   ]
 
-  # TODO(kitaev): add dropout
-  if attention_loop_stride < 1:
-    # Use the standard implementation if no loop_stride is provided.
-    attention = DotProductAttention(dropout=None, mode=mode)
-  else:
-    attention = MemoryEfficientDotProductAttention(
-        loop_stride=attention_loop_stride, dropout=None, mode=mode)
+  attention = attention_type(mode=mode)
 
   # ReversibleAttentionHalfResidual requires that post_attention be linear in
   # its input (so the backward pass can be computed without knowing the input)
@@ -596,7 +592,7 @@ def TransformerRevnetLM(vocab_size,
                         max_len=2048,
                         n_chunks=32,
                         n_attention_chunks=8,
-                        attention_loop_stride=0,
+                        attention_type=DotProductAttention,
                         mode='train'):
   """Reversible transformer language model (only uses a decoder, no encoder).
 
@@ -612,8 +608,7 @@ def TransformerRevnetLM(vocab_size,
     max_len: int: maximum symbol length for positional encoding
     n_chunks: int: number of chunks (must match input pipeline)
     n_attention_chunks: int: number of chunks for attention
-    attention_loop_stride: int: number of query elements to compute attention
-      for in parallel. Set to 0 to disable memory-efficient attention.
+    attention_type: class: attention class to use, such as DotProductAttention.
     mode: str: 'train' or 'eval'
 
   Returns:
@@ -633,7 +628,7 @@ def TransformerRevnetLM(vocab_size,
           # pylint: disable=g-complex-comprehension
           DecoderBlock(d_model, d_ff,
                        d_attention_key, d_attention_value, n_heads,
-                       n_attention_chunks, attention_loop_stride,
+                       n_attention_chunks, attention_type,
                        dropout, mode)
           for _ in range(n_layers)
       ]),

From bf60fdc8d76b29948f362beb7796663e3e7081e2 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 7 Aug 2019 15:28:07 -0700
Subject: [PATCH 2269/2720] add area attention config for librispeech with
 transformer

PiperOrigin-RevId: 262227683
---
 tensor2tensor/models/transformer.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 24e44167a..7a86ce6b7 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -2718,6 +2718,17 @@ def transformer_librispeech_tpu_v2():
   return hparams
 
 
+@registry.register_hparams
+def transformer_librispeech_with_area_attention():
+  """HParams for training ASR model on Librispeech on TPU v2."""
+  hparams = transformer_librispeech_tpu_v2()
+  hparams.num_area_layers = 3  # area attn on first 3 encoder and decoder layers
+  hparams.max_area_width = 5
+  hparams.area_key_mode = "concat"
+  hparams.area_value_mode = "sum"
+  return hparams
+
+
 @registry.register_hparams
 def transformer_librispeech():
   """HParams for training ASR model on Librispeech."""

From 9133c3e71a3945cda7d0b6c9ea45bd858adc5b29 Mon Sep 17 00:00:00 2001
From: Peng Wang <wangpeng@google.com>
Date: Thu, 8 Aug 2019 11:14:37 -0700
Subject: [PATCH 2270/2720] Augments trax_test.py to cover the tf-numpy
 backend.

PiperOrigin-RevId: 262393742
---
 tensor2tensor/trax/backend.py   |  24 +++++++
 tensor2tensor/trax/trax.py      |   2 +-
 tensor2tensor/trax/trax_test.py | 118 ++++++++++++++++++++------------
 3 files changed, 100 insertions(+), 44 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index a9bdb0cc9..d73bb1af7 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -166,6 +166,26 @@ def jax_eval_on_shapes(f):
   return functools.partial(_jax_eval_on_shapes, f)
 
 
+# The default value of dtype is different from jax_random.randint
+def jax_randint(key, shape, minval, maxval, dtype=onp.int32):
+  """Sample uniform random values in [minval, maxval) with given shape/dtype.
+
+  Args:
+    key: a PRNGKey used as the random key.
+    shape: a tuple of nonnegative integers representing the shape.
+    minval: int or array of ints broadcast-compatible with ``shape``, a minimum
+      (inclusive) value for the range.
+    maxval: int or array of ints broadcast-compatible with  ``shape``, a maximum
+      (exclusive) value for the range.
+    dtype: optional, an int dtype for the returned values (default int32).
+
+  Returns:
+    A random array with the specified shape and dtype.
+  """
+  return jax_random.randint(key, shape, minval=minval, maxval=maxval,
+                            dtype=dtype)
+
+
 _JAX_BACKEND = {
     "name": "jax",
     "np": jnp,
@@ -179,6 +199,7 @@ def jax_eval_on_shapes(f):
     "pmap": jax.pmap,
     "eval_on_shapes": jax_eval_on_shapes,
     "random_uniform": jax_random.uniform,
+    "random_randint": jax_randint,
     "random_normal": jax_random.normal,
     "random_bernoulli": jax_random.bernoulli,
     "random_get_prng": jax.jit(jax_random.PRNGKey),
@@ -258,6 +279,9 @@ def split(self, prng, num=2):
   def uniform(self, *args, **kwargs):
     return backend()["random_uniform"](*args, **kwargs)
 
+  def randint(self, *args, **kwargs):
+    return backend()["random_randint"](*args, **kwargs)
+
   def normal(self, *args, **kwargs):
     return backend()["random_normal"](*args, **kwargs)
 
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 219f95299..200784a0e 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -95,7 +95,7 @@ def masked_mean(inputs, targets, weights, mask_id=None):
   if mask_id is not None:
     weights = [w * (1.0 - np.equal(t, mask_id).astype(np.float32))
                for t, w in zip(targets, weights)]
-  weight_sums = [t.size if np.isscalar(w) else np.sum(w)
+  weight_sums = [np.float32(t.size) if np.isscalar(w) else np.sum(w)
                  for w, t in zip(weights, targets)]
   return sum([np.sum(x * w) / (length * s)
               for x, w, s in zip(inputs, weights, weight_sums)])
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index 28d2d425b..22b7b7cf1 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -22,11 +22,13 @@
 import contextlib
 import functools
 import tempfile
+from absl.testing import parameterized
 
 import gin
+import jax
 from jax import test_util  # pylint: disable=unused-import
 from jax.config import config
-import numpy as np
+import numpy as onp
 
 from tensor2tensor.trax import backend
 from tensor2tensor.trax import inputs as inputs_lib
@@ -34,7 +36,9 @@
 from tensor2tensor.trax import models
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
+from tensor2tensor.trax.backend import numpy as np
 
+import tensorflow as tf
 from tensorflow import test
 from tensorflow.io import gfile
 
@@ -45,10 +49,14 @@ def test_inputs(n_classes, with_weights=False):
   input_shape = (6, 6, 3)
 
   def input_stream():
+    key = backend.random.get_prng(0)
     while True:
-      inputs = np.random.rand(*([batch_size] + list(input_shape)))
-      targets = np.random.randint(n_classes, size=batch_size)
-      weights = np.random.rand(batch_size)
+      keys = backend.random.split(key, 4)
+      key = keys[0]
+      inputs = backend.random.uniform(keys[1], [batch_size] + list(input_shape))
+      targets = backend.random.randint(keys[2], [batch_size], dtype=np.int32,
+                                       minval=0, maxval=n_classes)
+      weights = backend.random.uniform(keys[3], [batch_size])
       if with_weights:
         yield inputs, targets, weights
       else:
@@ -62,7 +70,10 @@ def input_stream():
       input_dtype=np.float32)
 
 
-class TraxTest(test.TestCase):
+BACKENDS = ["jax"]
+
+
+class TraxTest(test.TestCase, parameterized.TestCase):
 
   @contextlib.contextmanager
   def tmp_dir(self):
@@ -70,8 +81,13 @@ def tmp_dir(self):
     yield tmp
     gfile.rmtree(tmp)
 
-  def test_train_eval_predict(self):
-    with self.tmp_dir() as output_dir:
+  # TODO(wangpeng): Remove `skipTest`'s when tf-numpy's `pmap` is in place
+
+  @parameterized.parameters(BACKENDS)
+  def test_train_eval_predict(self, backend_name):
+    if jax.lib.xla_bridge.device_count() > 1 and backend_name == "tf":
+      self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
+    with backend.use_backend(backend_name), self.tmp_dir() as output_dir:
       # Prepare model and inputs
       n_classes = 4
       train_steps = 2
@@ -95,15 +111,18 @@ def test_train_eval_predict(self):
       train_acc = state.history.get("train", "metrics/accuracy")
       eval_acc = state.history.get("eval", "metrics/accuracy")
       self.assertEqual(len(train_acc), len(eval_acc))
-      self.assertEqual(2, len(eval_acc))
+      self.assertLen(eval_acc, 2)
 
       # Predict with final params
       inputs = inputs(1).train_stream()
       model = layers.Serial(model_fn())
       model(next(inputs)[0], state.opt_state.params)
 
-  def test_train_eval_predict_sm3(self):
-    with self.tmp_dir() as output_dir:
+  @parameterized.parameters(BACKENDS)
+  def test_train_eval_predict_sm3(self, backend_name):
+    if jax.lib.xla_bridge.device_count() > 1 and backend_name == "tf":
+      self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
+    with backend.use_backend(backend_name), self.tmp_dir() as output_dir:
       # Prepare model and inputs
       n_classes = 4
       train_steps = 2
@@ -128,15 +147,18 @@ def test_train_eval_predict_sm3(self):
       train_acc = state.history.get("train", "metrics/accuracy")
       eval_acc = state.history.get("eval", "metrics/accuracy")
       self.assertEqual(len(train_acc), len(eval_acc))
-      self.assertEqual(2, len(eval_acc))
+      self.assertLen(eval_acc, 2)
 
       # Predict with final params
       inputs = inputs(1).train_stream()
       model = layers.Serial(model_fn())
       model(next(inputs)[0], state.opt_state.params)
 
-  def test_train_restart(self):
-    with self.tmp_dir() as output_dir:
+  @parameterized.parameters(BACKENDS)
+  def test_train_restart(self, backend_name):
+    if jax.lib.xla_bridge.device_count() > 1 and backend_name == "tf":
+      self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
+    with backend.use_backend(backend_name), self.tmp_dir() as output_dir:
       # Prepare model and inputs
       n_classes = 4
       train_steps = 2
@@ -163,8 +185,11 @@ def test_train_restart(self):
       # Assert total train steps
       self.assertEqual(state.step, 2 * train_steps)
 
-  def test_train_with_weights(self):
-    with self.tmp_dir() as output_dir:
+  @parameterized.parameters(BACKENDS)
+  def test_train_with_weights(self, backend_name):
+    if jax.lib.xla_bridge.device_count() > 1 and backend_name == "tf":
+      self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
+    with backend.use_backend(backend_name), self.tmp_dir() as output_dir:
       gin.bind_parameter("unpack_batch.has_weights", True)
 
       # Prepare model and inputs
@@ -187,39 +212,46 @@ def test_train_with_weights(self):
       self.assertEqual(state.step, train_steps)
 
 
-class MaskedMeanTest(test.TestCase):
+MASKED_MEAN_TEST_BACKENDS = ["numpy"]
 
-  def test_computes_basic_mean(self):
-    inputs = [np.array([1, 2, 3])]
-    targets = [np.zeros(3)]
-    weights = [1]
-    with backend.use_backend("numpy"):
-      mean = trax.masked_mean(inputs, targets, weights)
-      np.testing.assert_allclose(mean, 2)
 
-  def test_computes_mean_with_weights(self):
-    inputs = [np.array([1, 2, 3])]
-    targets = [np.zeros(3)]
-    weights = [np.array([3, 1, 0])]
-    with backend.use_backend("numpy"):
-      mean = trax.masked_mean(inputs, targets, weights)
-      np.testing.assert_allclose(mean, 1.25)
+class MaskedMeanTest(test.TestCase, parameterized.TestCase):
 
-  def test_computes_mean_with_mask(self):
-    inputs = [np.array([1, 2, 3])]
-    targets = [np.array([1, 0, 0])]
-    weights = [1]
-    with backend.use_backend("numpy"):
+  @parameterized.parameters(MASKED_MEAN_TEST_BACKENDS)
+  def test_computes_basic_mean(self, backend_name):
+    with backend.use_backend(backend_name):
+      inputs = [np.array([1, 2, 3])]
+      targets = [np.zeros(3)]
+      weights = [1]
+      mean = trax.masked_mean(inputs, targets, weights)
+      onp.testing.assert_allclose(mean, 2)
+
+  @parameterized.parameters(MASKED_MEAN_TEST_BACKENDS)
+  def test_computes_mean_with_weights(self, backend_name):
+    with backend.use_backend(backend_name):
+      inputs = [np.array([1, 2, 3])]
+      targets = [np.zeros(3)]
+      weights = [np.array([3, 1, 0])]
+      mean = trax.masked_mean(inputs, targets, weights)
+      onp.testing.assert_allclose(mean, 1.25)
+
+  @parameterized.parameters(MASKED_MEAN_TEST_BACKENDS)
+  def test_computes_mean_with_mask(self, backend_name):
+    with backend.use_backend(backend_name):
+      inputs = [np.array([1, 2, 3])]
+      targets = [np.array([1, 0, 0])]
+      weights = [1]
       mean = trax.masked_mean(inputs, targets, weights, mask_id=1)
-      np.testing.assert_allclose(mean, 2.5)
-
-  def test_computes_mean_with_weights_and_mask(self):
-    inputs = [np.array([1, 2, 4])]
-    targets = [np.array([1, 0, 0])]
-    weights = [np.array([10, 4, 1])]
-    with backend.use_backend("numpy"):
+      onp.testing.assert_allclose(mean, 2.5)
+
+  @parameterized.parameters(MASKED_MEAN_TEST_BACKENDS)
+  def test_computes_mean_with_weights_and_mask(self, backend_name):
+    with backend.use_backend(backend_name):
+      inputs = [np.array([1, 2, 4])]
+      targets = [np.array([1, 0, 0])]
+      weights = [np.array([10, 4, 1])]
       mean = trax.masked_mean(inputs, targets, weights, mask_id=1)
-      np.testing.assert_allclose(mean, 2.4)
+      onp.testing.assert_allclose(mean, 2.4)
 
 
 if __name__ == "__main__":

From 4a4d7fa1d4ec4372a95d9e969c8d8b7b3b5a60aa Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 8 Aug 2019 16:25:02 -0700
Subject: [PATCH 2271/2720] Add DummyHashedAttention

PiperOrigin-RevId: 262456348
---
 .../transformer_revnet_imagenet64_8gb.gin     |   6 +-
 tensor2tensor/trax/models/__init__.py         |   2 +
 .../models/research/transformer_revnet.py     | 111 ++++++++++++++++++
 3 files changed, 118 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
index 0e158baf9..54efae800 100644
--- a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
@@ -35,13 +35,17 @@ train.trainer_class = @MemoryEfficientTrainer
 # Parameters for DotProductAttention:
 # ==============================================================================
 DotProductAttention.dropout = 0.0
-MemoryEfficientDotProductAttention.loop_stride = 512
 
 # Parameters for MemoryEfficientDotProductAttention:
 # ==============================================================================
 MemoryEfficientDotProductAttention.dropout = 0.0
 MemoryEfficientDotProductAttention.loop_stride = 512
 
+# Parameters for DummyHashedAttention:
+# ==============================================================================
+DummyHashedAttention.dropout = 0.0
+DummyHashedAttention.n_bins = 64
+
 # Parameters for TransformerRevnetLM:
 # ==============================================================================
 TransformerRevnetLM.d_model = 1024
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index 1b6b281b4..09d1b1b04 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -52,3 +52,5 @@ def model_configure(*args, **kwargs):
     transformer_revnet.DotProductAttention, blacklist=["mode"])
 MemoryEfficientDotProductAttention = model_configure(
     transformer_revnet.MemoryEfficientDotProductAttention, blacklist=["mode"])
+DummyHashedAttention = model_configure(
+    transformer_revnet.DummyHashedAttention, blacklist=["mode"])
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 4357663a3..b77001121 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -401,6 +401,117 @@ def body_fun(vals):
       return final_vals[1], final_vals[2:]
 
 
+class DummyHashedAttention(DotProductAttention):
+  """A stand-in for hash-based attention, but without a real hash function."""
+
+  def __init__(self, dropout, mode, n_bins=64):
+    super(DummyHashedAttention, self).__init__(dropout, mode)
+    self.n_bins = n_bins
+
+  def call(self, inputs, params=(), **kwargs):
+    output, _ = self.forward_and_vjp(inputs, None, params=params, **kwargs)
+    return output
+
+  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
+    del params, kwargs
+    q, k, v = inputs
+    # q/k/v are n_batch, n_heads, seqlen, d_head
+
+    assert k.shape[2] % self.n_bins == 0
+    bin_size = int(k.shape[2] // self.n_bins)
+
+    # q_bins/kv_bins are n_batch, n_heads, seqlen
+    # They specify which hash bucket the query/key/value vectors fall in. For
+    # now, instead of hashing we just put consecutive items in the same bucket.
+    q_bins = np.arange(q.shape[2], dtype=np.int32) // bin_size
+    q_bins = jax.lax.tie_in(q, q_bins)
+    q_bins = q_bins[None, None, :]
+    q_bins = np.broadcast_to(q_bins, q.shape[:-1])
+    q_bins = -q_bins
+    kv_bins = q_bins * 2
+
+    # q_t/kv_t are n_batch, n_heads, seqlen
+    q_t = jax.lax.tie_in(q, np.arange(q.shape[2]))
+    q_t = np.reshape(q_t, (1, 1, q_t.shape[0]))
+    q_t = np.broadcast_to(q_t, q.shape[:-1])
+    kv_t = q_t
+
+    def chunk_rank3(x):
+      return np.reshape(x, (x.shape[0], x.shape[1], self.n_bins, -1))
+
+    def chunk_rank4(x):
+      return np.reshape(
+          x, (x.shape[0], x.shape[1], self.n_bins, -1, x.shape[-1]))
+
+    def unchunk_rank4(x):
+      return np.reshape(x, (x.shape[0], x.shape[1], -1, x.shape[-1]))
+
+   # Sort everything by bin number (variables starting with "s" are sorted)
+    _, sq_t = jax.lax.sort_key_val(q_bins, q_t, dimension=2)
+
+    sq = np.take_along_axis(q, sq_t[:, :, :, None], axis=2)
+    if ct is not None:
+      so_ct = np.take_along_axis(ct, sq_t[:, :, :, None], axis=2)
+
+    _, skv_t = jax.lax.sort_key_val(kv_bins, kv_t, dimension=2)
+    sk = np.take_along_axis(k, skv_t[:, :, :, None], axis=2)
+    sv = np.take_along_axis(v, skv_t[:, :, :, None], axis=2)
+
+    @jax.jit
+    def binned_attn(sq, sk, sv):
+      """Performs attention on sorted queries/keys/values."""
+      # Split off a "bin" axis so that attention only occurs whithin chunks.
+      bq_t = chunk_rank3(sq_t)
+      bkv_t = chunk_rank3(skv_t)
+      bq = chunk_rank4(sq)
+      bk = chunk_rank4(sk)
+      bv = chunk_rank4(sv)
+
+      dots = np.matmul(bq, np.swapaxes(bk, -1, -2)) / np.sqrt(bq.shape[-1])
+
+      # Causal masking
+      mask = jax.lax.convert_element_type(
+          jax.lax.lt(bq_t[:, :, :, :, None], bkv_t[:, :, :, None, :]),
+          np.float32)
+      dots = dots - 1e9 * mask
+
+      # Softmax.
+      dots = np.exp(dots - dots.max(axis=-1, keepdims=True))
+      dots = dots / dots.sum(axis=-1, keepdims=True)
+      bo = np.matmul(dots, bv)
+
+      so = unchunk_rank4(bo)
+      return so
+
+    @jax.jit
+    def binned_attn_vjp(sq, sk, sv, so_ct):
+      so, vjpfun = jax.vjp(binned_attn, sq, sk, sv)
+      sqkv_ct = vjpfun(so_ct)
+      return so, sqkv_ct
+
+    if ct is None:
+      so = binned_attn(sq, sk, sv)
+      _, undo_q_sort = jax.lax.sort_key_val(sq_t, q_t, dimension=2)
+      out = np.take_along_axis(so, undo_q_sort[:, :, :, None], axis=2)
+      return out, None
+    else:
+      # Jax can construct a backward pass automatically, but it's about 2x
+      # slower than writing our own. The main reason is that the backward pass
+      # of gather is in general a scatter operation, but we know we're dealing
+      # with permutations so we use gather for the backward pass too.
+      so, (sq_ct, sk_ct, sv_ct) = binned_attn_vjp(sq, sk, sv, so_ct)
+
+      _, undo_q_sort = jax.lax.sort_key_val(sq_t, q_t, dimension=2)
+      out = np.take_along_axis(so, undo_q_sort[:, :, :, None], axis=2)
+      q_ct = np.take_along_axis(sq_ct, undo_q_sort[:, :, :, None], axis=2)
+
+      _, undo_kv_sort = jax.lax.sort_key_val(skv_t, kv_t, dimension=2)
+      k_ct = np.take_along_axis(sk_ct, undo_kv_sort[:, :, :, None], axis=2)
+      v_ct = np.take_along_axis(sv_ct, undo_kv_sort[:, :, :, None], axis=2)
+
+      return out, (q_ct, k_ct, v_ct)
+
+
 class ReversibleAttentionHalfResidual(tl.ReversibleLayer, tl.Serial):
   """Half of a RevNet-style residual that performs attention.
 

From 1e896cc67c692bd8d2e327dec82c81ab8ed6b24a Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 8 Aug 2019 18:03:35 -0700
Subject: [PATCH 2272/2720] Include all metrics in an OnlineTuneEnv observation

PiperOrigin-RevId: 262472256
---
 .../online_tune_wide_resnet_cifar10.gin       |  8 +++-
 .../trax/rlax/envs/online_tune_env.py         | 42 ++++++++++++-------
 .../trax/rlax/envs/online_tune_env_test.py    |  4 +-
 3 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
index 688083821..3743d52eb 100644
--- a/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
@@ -26,7 +26,7 @@ shuffle_and_batch_data.preprocess_fun = @trax.inputs.cifar10_no_augmentation_pre
 
 # Parameters for WideResnet:
 # ==============================================================================
-WideResnet.widen_factor = 2
+WideResnet.widen_factor = 10
 WideResnet.n_blocks = 3
 WideResnet.n_output_classes = 10
 
@@ -39,6 +39,12 @@ OnlineTuneEnv.start_lr = 0.01
 OnlineTuneEnv.train_steps = 500
 OnlineTuneEnv.eval_steps = 100
 OnlineTuneEnv.env_steps = 100
+OnlineTuneEnv.observation_metrics = [
+    ("train", "metrics/accuracy"),
+    ("train", "metrics/loss"),
+    ("eval", "metrics/accuracy"),
+    ("eval", "metrics/loss"),
+]
 
 # Parameters for ppo.training_loop:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rlax/envs/online_tune_env.py b/tensor2tensor/trax/rlax/envs/online_tune_env.py
index b5cba2427..ba05a1af1 100644
--- a/tensor2tensor/trax/rlax/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rlax/envs/online_tune_env.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import math
 import os
 
 import gym
@@ -53,8 +54,9 @@ def __init__(self,
                optimizer=trax_opt.SM3,
                inputs=trax_inputs.inputs,
                action_multipliers=None,
-               history_mode="eval",
-               metric="metrics/accuracy",
+               observation_metrics=(("eval", "metrics/accuracy"),),
+               include_lr_in_observation=False,
+               reward_metric=("eval", "metrics/accuracy"),
                train_steps=100,
                eval_steps=10,
                env_steps=100,
@@ -69,8 +71,9 @@ def __init__(self,
         lr_schedule=(lambda history: lambda step: self._current_lr),
         inputs=inputs)
     self._action_multipliers = action_multipliers
-    self._history_mode = history_mode
-    self._metric = metric
+    self._observation_metrics = observation_metrics
+    self._include_lr_in_observation = include_lr_in_observation
+    self._reward_metric = reward_metric
     self._train_steps = train_steps
     self._eval_steps = eval_steps
     self._env_steps = env_steps
@@ -80,10 +83,12 @@ def __init__(self,
     gfile.makedirs(self._output_dir)
     # Action is an index in self._action_multipliers.
     self.action_space = gym.spaces.Discrete(len(self._action_multipliers))
-    # Observation is a singleton vector with the value of the metric specified
-    # in self._metric.
+    # Observation is a vector with the values of the metrics specified in
+    # observation_metrics plus optionally the learning rate.
+    observation_dim = (
+        len(self._observation_metrics) + int(self._include_lr_in_observation))
     self.observation_space = gym.spaces.Box(
-        low=float("-inf"), high=float("+inf"), shape=(1,))
+        low=float("-inf"), high=float("+inf"), shape=(observation_dim,))
 
   @property
   def _next_trajectory_dir(self):
@@ -111,14 +116,21 @@ def int_or_none(s):
 
     return os.path.join(self._output_dir, str(next_trajectory_id))
 
-  @property
-  def _current_metric_value(self):
-    metric_sequence = self._trainer.state.history.get(self._history_mode,
-                                                      self._metric)
+  def _current_metric_value(self, metric):
+    metric_sequence = self._trainer.state.history.get(*metric)
     assert metric_sequence
     (_, metric_value) = metric_sequence[-1]
     return metric_value
 
+  @property
+  def _current_observation(self):
+    observation = list(
+        map(self._current_metric_value, self._observation_metrics))
+    if self._include_lr_in_observation:
+      # Logartihm of the learning rate.
+      observation.append(math.log(self._current_lr))
+    return np.array(observation)
+
   @property
   def trainer(self):
     if self._trainer is None:
@@ -130,7 +142,7 @@ def reset(self):
     self._step = 0
     self._trainer.reset(output_dir=self._next_trajectory_dir)
     self._trainer.evaluate(self._eval_steps)
-    return np.array([self._current_metric_value])
+    return self._current_observation
 
   def step(self, action):
     """Step the environment.
@@ -147,11 +159,11 @@ def step(self, action):
         environment steps. info is an empty dict.
     """
     self._current_lr *= self._action_multipliers[action]
-    last_metric_value = self._current_metric_value
+    last_metric_value = self._current_metric_value(self._reward_metric)
     self._trainer.train_epoch(self._train_steps, self._eval_steps)
     self._step += 1
-    current_metric_value = self._current_metric_value
-    observation = np.array([current_metric_value])
+    current_metric_value = self._current_metric_value(self._reward_metric)
+    observation = self._current_observation
     reward = current_metric_value - last_metric_value
     done = self._step == self._env_steps
     return (observation, reward, done, {})
diff --git a/tensor2tensor/trax/rlax/envs/online_tune_env_test.py b/tensor2tensor/trax/rlax/envs/online_tune_env_test.py
index 02f13b054..fd8e8438d 100644
--- a/tensor2tensor/trax/rlax/envs/online_tune_env_test.py
+++ b/tensor2tensor/trax/rlax/envs/online_tune_env_test.py
@@ -77,8 +77,8 @@ def _create_env(output_dir, metrics_to_report=(0.0,), action_multipliers=()):
             output_dtype=np.float32),
         output_dir=output_dir,
         action_multipliers=action_multipliers,
-        history_mode=HISTORY_MODE,
-        metric=METRIC,
+        observation_metrics=[(HISTORY_MODE, METRIC)],
+        reward_metric=(HISTORY_MODE, METRIC),
         train_steps=1,
         eval_steps=1,
         env_steps=(len(metrics_to_report) - 1))

From b690b5d2ec4ea9927a507ebe72523c8864ba8cad Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 8 Aug 2019 18:28:47 -0700
Subject: [PATCH 2273/2720] Don't save model checkpoints in OnlineTuneEnv

PiperOrigin-RevId: 262475193
---
 tensor2tensor/trax/rlax/envs/online_tune_env.py |  9 +++++++--
 tensor2tensor/trax/trax.py                      | 17 ++++++++++-------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/trax/rlax/envs/online_tune_env.py b/tensor2tensor/trax/rlax/envs/online_tune_env.py
index ba05a1af1..bfe16218a 100644
--- a/tensor2tensor/trax/rlax/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rlax/envs/online_tune_env.py
@@ -60,7 +60,10 @@ def __init__(self,
                train_steps=100,
                eval_steps=10,
                env_steps=100,
-               start_lr=0.001):
+               start_lr=0.001,
+               # Don't save checkpoints by default, as they tend to use a lot of
+               # space.
+               should_save_checkpoints=False):
     if action_multipliers is None:
       action_multipliers = self.DEFAULT_ACTION_MULTIPLIERS
     self._model = model
@@ -69,7 +72,9 @@ def __init__(self,
         loss_fn=loss_fn,
         optimizer=optimizer,
         lr_schedule=(lambda history: lambda step: self._current_lr),
-        inputs=inputs)
+        inputs=inputs,
+        should_save=should_save_checkpoints,
+    )
     self._action_multipliers = action_multipliers
     self._observation_metrics = observation_metrics
     self._include_lr_in_observation = include_lr_in_observation
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 200784a0e..ede9ca1e2 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -495,10 +495,11 @@ class Trainer(object):
 
   def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs,
                output_dir=None, random_seed=None, n_devices=None,
-               save_steps=None):
+               save_steps=None, should_save=True):
     if save_steps is None:
       save_steps = []
     self._save_steps = save_steps
+    self._should_save = should_save
     device_count = jax.lib.xla_bridge.device_count()
     n_devices = n_devices or device_count
     # TODO(lukaszkaiser): remove this restriction when possible.
@@ -595,8 +596,7 @@ def reset(self, output_dir):
     self._opt_state = OptState(*layers.nested_map(
         opt_state, self._maybe_replicate))
     if not state.opt_state:
-      _save_replicated(self._opt_state, self._step, self._history,
-                       self._n_devices, self._output_dir, False)
+      self._maybe_save_state(keep=False)
 
     self.update_learning_rate()
 
@@ -627,6 +627,11 @@ def _maybe_replicate(self, x):
     else:
       return x
 
+  def _maybe_save_state(self, keep):
+    if self._should_save:
+      _save_replicated(self._opt_state, self._step, self._history,
+                       self._n_devices, self._output_dir, keep)
+
   def save_gin(self):
     _save_gin(self._output_dir, self._train_sw)
 
@@ -665,8 +670,7 @@ def train_epoch(self, epoch_steps, eval_steps):
       self._train_step(next_train_batch)
 
       if self._step in self._save_steps:
-        _save_replicated(self._opt_state, self._step, self._history,
-                         self._n_devices, self._output_dir, True)
+        self._maybe_save_state(keep=True)
 
       # LR log
       if self._step == 1 or self._step % 10 == 0:
@@ -688,8 +692,7 @@ def train_epoch(self, epoch_steps, eval_steps):
     self.evaluate(eval_steps)
 
     # Save state
-    _save_replicated(self._opt_state, self._step, self._history,
-                     self._n_devices, self._output_dir, False)
+    self._maybe_save_state(keep=False)
 
     # Flush summary writers
     self._train_sw.flush()

From 8cf834aa375f8a0e613a1412d63b9ec7aedbe12b Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Fri, 9 Aug 2019 14:52:59 -0700
Subject: [PATCH 2274/2720] Weight decay not being used in AdamW in T2T?

PiperOrigin-RevId: 262639162
---
 tensor2tensor/models/transformer.py |  2 ++
 tensor2tensor/utils/optimize.py     | 11 ++---------
 2 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 7a86ce6b7..37140c9d8 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -2053,6 +2053,7 @@ def transformer_tall_pretrain_lm():
   hparams.learning_rate_constant = 2e-4
   hparams.learning_rate_schedule = ("linear_warmup*constant*cosdecay")
   hparams.optimizer = "adam_w"
+  hparams.weight_decay = 0.01 * hparams.learning_rate_constant
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.999
   hparams.optimizer_adam_epsilon = 1e-8
@@ -2098,6 +2099,7 @@ def transformer_tall_pretrain_lm_tpu():
   hparams.learning_rate_constant = 2e-4
   hparams.learning_rate_schedule = ("linear_warmup * constant * cosdecay")
   hparams.optimizer = "adam_w"
+  hparams.weight_decay = 0.01 * hparams.learning_rate_constant
   return hparams
 
 
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index fd8bf46c3..1df3fbbdd 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -159,16 +159,9 @@ def true_adam(learning_rate, hparams):
 
 @registry.register_optimizer
 def adam_w(learning_rate, hparams):
-  # Openai gpt used weight decay.
-  # Given the internals of AdamW, weight decay dependent on the
-  # learning rate is chosen to match the openai implementation.
-  # The weight decay update to each parameter is applied before the adam
-  # gradients computation, which is different from that described
-  # in the paper and in the openai implementation:
-  # https://arxiv.org/pdf/1711.05101.pdf
   return tf.contrib.opt.AdamWOptimizer(
-      0.01*learning_rate,
-      learning_rate,
+      weight_decay=hparams.weight_decay,
+      learning_rate=learning_rate,
       beta1=hparams.optimizer_adam_beta1,
       beta2=hparams.optimizer_adam_beta2,
       epsilon=hparams.optimizer_adam_epsilon)

From 0935c5248743668b03e40125e6f6503e6ba775ed Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 9 Aug 2019 15:26:00 -0700
Subject: [PATCH 2275/2720] Implement framestack for MLP policies

PiperOrigin-RevId: 262645693
---
 tensor2tensor/trax/models/__init__.py         |  1 +
 tensor2tensor/trax/models/atari_cnn.py        | 35 ++++++++++++++++---
 tensor2tensor/trax/models/atari_cnn_test.py   | 18 ++++++++++
 .../online_tune_wide_resnet_cifar10.gin       |  6 +++-
 tensor2tensor/trax/rlax/ppo_main.py           | 21 ++++-------
 5 files changed, 61 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index 09d1b1b04..30876c202 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -37,6 +37,7 @@ def model_configure(*args, **kwargs):
 
 # pylint: disable=invalid-name
 AtariCnn = model_configure(atari_cnn.AtariCnn)
+FrameStackMLP = model_configure(atari_cnn.FrameStackMLP)
 MLP = model_configure(mlp.MLP)
 NeuralGPU = model_configure(neural_gpu.NeuralGPU)
 PositionLookupTransformerLM = model_configure(
diff --git a/tensor2tensor/trax/models/atari_cnn.py b/tensor2tensor/trax/models/atari_cnn.py
index 0eb4494d9..7b1e8e251 100644
--- a/tensor2tensor/trax/models/atari_cnn.py
+++ b/tensor2tensor/trax/models/atari_cnn.py
@@ -22,7 +22,22 @@
 from tensor2tensor.trax import layers as tl
 
 
-def AtariCnn(hidden_sizes=(32, 32), output_size=128, mode='train'):
+def FrameStack(n_frames):
+  """Stacks a fixed number of frames along the dimension 1."""
+  # Input shape: (B, T, ..., C).
+  # Output shape: (B, T, ..., C * n_frames).
+  assert n_frames >= 1
+  return (
+      # Make n_frames copies of the input sequence.
+      [tl.Dup()] * (n_frames - 1),
+      # Shift copies to the right by [0, .., n_frames - 1] frames.
+      tl.Parallel(*map(_shift_right, range(n_frames))),
+      # Concatenate along the channel dimension.
+      tl.Concatenate(n_items=n_frames, axis=-1),
+  )
+
+
+def AtariCnn(n_frames=4, hidden_sizes=(32, 32), output_size=128, mode='train'):
   """An Atari CNN."""
   del mode
 
@@ -33,10 +48,8 @@ def AtariCnn(hidden_sizes=(32, 32), output_size=128, mode='train'):
       tl.ToFloat(),
       tl.Div(divisor=255.0),
 
-      # Set up 4 successive game frames, concatenated on the last axis.
-      tl.Dup(), tl.Dup(), tl.Dup(),
-      tl.Parallel(None, _shift_right(1), _shift_right(2), _shift_right(3)),
-      tl.Concatenate(n_items=4, axis=-1),  # (B, T, H, W, 4C)
+      # Set up n_frames successive game frames, concatenated on the last axis.
+      FrameStack(n_frames=n_frames),  # (B, T, H, W, 4C)
 
       tl.Conv(hidden_sizes[0], (5, 5), (2, 2), 'SAME'),
       tl.Relu(),
@@ -48,5 +61,17 @@ def AtariCnn(hidden_sizes=(32, 32), output_size=128, mode='train'):
   )
 
 
+def FrameStackMLP(n_frames=4, hidden_sizes=(64,), output_size=64,
+                  mode='train'):
+  """MLP operating on a fixed number of last frames."""
+  del mode
+
+  return tl.Model(
+      FrameStack(n_frames=n_frames),
+      [[tl.Dense(d_hidden), tl.Relu()] for d_hidden in hidden_sizes],
+      tl.Dense(output_size),
+  )
+
+
 def _shift_right(n):  # pylint: disable=invalid-name
   return [tl.ShiftRight()] * n
diff --git a/tensor2tensor/trax/models/atari_cnn_test.py b/tensor2tensor/trax/models/atari_cnn_test.py
index e0bd2e730..abb2a200d 100644
--- a/tensor2tensor/trax/models/atari_cnn_test.py
+++ b/tensor2tensor/trax/models/atari_cnn_test.py
@@ -45,5 +45,23 @@ def test_computes(self):
     self.assertEqual((B, T + 1, output_size), y.shape)
 
 
+class FrameStackMLPTest(test.TestCase):
+
+  def test_computes(self):
+    rng_key = jax_random.get_prng(0)
+    hidden_size = (4, 4)
+    output_size = 6
+    model = atari_cnn.FrameStackMLP(
+        hidden_sizes=hidden_size, output_size=output_size)
+    B, T, OBS = 2, 2, 3  # pylint: disable=invalid-name
+    rng_key, key = jax_random.split(rng_key)
+    params = model.initialize((1, 1, OBS), onp.float32, key)
+    x = onp.arange(B * (T + 1) * OBS).reshape(
+        B, T + 1, OBS)
+    rng_key, key = jax_random.split(rng_key)
+    y = model(x, params, rng=key)
+    self.assertEqual((B, T + 1, output_size), y.shape)
+
+
 if __name__ == "__main__":
   test.main()
diff --git a/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
index 3743d52eb..0fb7382cd 100644
--- a/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
@@ -24,6 +24,10 @@ Momentum.mass = 0.9
 # ==============================================================================
 shuffle_and_batch_data.preprocess_fun = @trax.inputs.cifar10_no_augmentation_preprocess
 
+# Parameters for FrameStackMLP:
+# ==============================================================================
+FrameStackMLP.n_frames = 4
+
 # Parameters for WideResnet:
 # ==============================================================================
 WideResnet.widen_factor = 10
@@ -62,5 +66,5 @@ ppo.training_loop.eval_every_n = 10
 ppo.training_loop.done_frac_for_policy_save = 0
 ppo.training_loop.enable_early_stopping = True
 ppo.training_loop.n_evals = 1
-ppo.training_loop.len_history_for_policy = 1
 ppo.training_loop.eval_temperatures = (1.0,)
+ppo.training_loop.len_history_for_policy = 4
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 4b1249c47..7fe717e68 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -54,7 +54,6 @@
 from tensor2tensor.envs import rendered_env_problem
 from tensor2tensor.rl import gym_utils
 from tensor2tensor.rl.google import atari_utils  # GOOGLE-INTERNAL:
-from tensor2tensor.trax import layers
 from tensor2tensor.trax import models
 from tensor2tensor.trax.rlax import envs as rlax_envs  # pylint: disable=unused-import
 from tensor2tensor.trax.rlax import ppo
@@ -112,18 +111,6 @@
 flags.DEFINE_string("eval_server_bns", "", "Eval Server's BNS.")
 
 
-def common_layers():
-  # TODO(afrozm): Refactor.
-  if "NoFrameskip" in FLAGS.env_problem_name:
-    return atari_layers()
-
-  return [layers.Dense(64), layers.Tanh(), layers.Dense(64), layers.Tanh()]
-
-
-def atari_layers():
-  return [models.AtariCnn()]
-
-
 def make_env(batch_size=8, **env_kwargs):
   """Creates the env."""
 
@@ -219,13 +206,19 @@ def main(argv):
   eval_env = make_env(batch_size=FLAGS.eval_batch_size, **eval_env_kwargs)
   assert eval_env
 
+  # TODO(afrozm): Refactor.
+  if "NoFrameskip" in FLAGS.env_problem_name:
+    bottom_layers_fn = models.AtariCnn
+  else:
+    bottom_layers_fn = models.FrameStackMLP
+
   def run_training_loop():
     """Runs the training loop."""
     logging.info("Starting the training loop.")
 
     policy_and_value_net_fn = functools.partial(
         ppo.policy_and_value_net,
-        bottom_layers_fn=common_layers,
+        bottom_layers_fn=bottom_layers_fn,
         two_towers=FLAGS.two_towers)
     policy_and_value_optimizer_fn = get_optimizer_fn(FLAGS.learning_rate)
 

From 41726d4c16c55fe3fb8b4ad2c7030910c84a8b08 Mon Sep 17 00:00:00 2001
From: Daniel De Freitas Adiwardana <adiwardana@google.com>
Date: Fri, 9 Aug 2019 17:12:02 -0700
Subject: [PATCH 2276/2720] Parameterizing name of labels feature name.

PiperOrigin-RevId: 262663532
---
 tensor2tensor/utils/data_reader.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index 7f29944e6..7f99eb585 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -558,7 +558,8 @@ def prepare_for_output(example):
       example["infer_targets"] = example.pop("targets")
       return example
     else:
-      return example, example["targets"]
+      return example, example[hparams.get(
+          key="labels_feature_name", default="targets")]
 
   dataset = dataset.map(prepare_for_output, num_parallel_calls=num_threads)
   dataset = dataset.prefetch(2)

From f7f8549a6421723154b366996b2c6559048ac3fb Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Mon, 12 Aug 2019 15:40:23 -0700
Subject: [PATCH 2277/2720] Implement SerializedSequenceSimulatedEnvProblem

PiperOrigin-RevId: 263020747
---
 .../trax/rlax/ppo_training_loop_test.py       |   6 +-
 .../trax/rlax/simulated_env_problem.py        | 307 +++++++++++++++---
 .../trax/rlax/simulated_env_problem_test.py   |  85 ++++-
 tensor2tensor/trax/rlax/space_serializer.py   | 151 +++++++++
 .../trax/rlax/space_serializer_test.py        |  94 ++++++
 5 files changed, 579 insertions(+), 64 deletions(-)
 create mode 100644 tensor2tensor/trax/rlax/space_serializer.py
 create mode 100644 tensor2tensor/trax/rlax/space_serializer_test.py

diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index 25d5f49cc..641f10155 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -171,10 +171,10 @@ def loss(*args, **kwargs):
       )
       trainer.train_epoch(epoch_steps=1, eval_steps=1)
 
-      # Repeat the initial observations over and over again.
+      # Repeat the history over and over again.
       stream = itertools.repeat(np.zeros(history_shape))
       env_fn = functools.partial(
-          simulated_env_problem.SimulatedEnvProblem,
+          simulated_env_problem.RawSimulatedEnvProblem,
           model=model,
           history_length=history_shape[1],
           trajectory_length=3,
@@ -184,7 +184,7 @@ def loss(*args, **kwargs):
           action_space=gym.spaces.Discrete(n=n_actions),
           reward_range=(-1, 1),
           discrete_rewards=False,
-          initial_observation_stream=stream,
+          history_stream=stream,
           output_dir=output_dir,
       )
 
diff --git a/tensor2tensor/trax/rlax/simulated_env_problem.py b/tensor2tensor/trax/rlax/simulated_env_problem.py
index 0834527ac..52798d06c 100644
--- a/tensor2tensor/trax/rlax/simulated_env_problem.py
+++ b/tensor2tensor/trax/rlax/simulated_env_problem.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
 import random
 
 import numpy as np
@@ -27,22 +28,11 @@
 from tensor2tensor.trax import backend
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.backend import random as jax_random
+from tensor2tensor.trax.rlax import space_serializer
 
 
 class SimulatedEnvProblem(env_problem.EnvProblem):
-  """EnvProblem for environments simulated by TRAX models.
-
-  Wraps an autoregressive TRAX model of signature
-  (observation_history, action) -> (observation, reward) in an EnvProblem.
-  The model is assumed to take a fixed number of last observations as input
-  and produce a single observation, which is fed back into the model in the
-  next environment step.
-
-  Shape requirements (without the batch dimension):
-    observation: Consistent with observation_space.
-    observation_history: (history_length,) + observation.shape.
-    action: Consistent with action_space.
-    reward: (1,). The singleton dimension is removed in step().
+  """EnvProblem base class for environments simulated by TRAX models.
 
   The initial observations to start the model are taken from
   initial_observation_stream. This iterator in incremented in every reset().
@@ -50,71 +40,66 @@ class SimulatedEnvProblem(env_problem.EnvProblem):
   A checkpoint saved by the TRAX trainer should be available in output_dir.
   """
 
-  def __init__(self, model, history_length, trajectory_length, batch_size,
-               observation_space, action_space, reward_range, discrete_rewards,
-               initial_observation_stream, output_dir):
+  def __init__(self, model, batch_size, observation_space, action_space,
+               reward_range, discrete_rewards, history_stream, output_dir):
     """Initializes the env.
 
     Args:
       model: TRAX model.
-      history_length: (int) Number of last observations fed into the model.
-      trajectory_length: (int) Length of each trajectory unrolled from the
-        model.
       batch_size: (int) Number of simulated environments run in parallel.
       observation_space: (gym.Space) Observation space.
       action_space: (gym.Space) Action space.
       reward_range: (tuple) Pair (min_reward, max_reward).
       discrete_rewards: (bool) Whether to discretize the rewards.
-      initial_observation_stream: Iterator yielding batches of initial
-        observations for the model.
+      history_stream: Iterator yielding batches of initial input data for the
+        model. The format is implementation-specific.
       output_dir: (str) Output dir.
     """
     # TODO(pkozakowski): At some point we will have a "predict" mode which we
     # should use here. When this happens, change the mode.
     self._model_predict = backend.jit(model(mode="eval"))
-    self._history_length = history_length
-    self._trajectory_length = trajectory_length
     self._observation_space = observation_space
     self._action_space = action_space
     self._reward_range = reward_range
     self._output_dir = output_dir
 
-    self._model_params = None
+    self._predict_fn = None
     self._rng = None
-    self._initial_observation_stream = None
-    self._history = None
-    self._steps = None
+    self._history_stream = None
 
     # Call the super's ctor. It will use some of the member fields, so we call
     # it in the end.
     super(SimulatedEnvProblem, self).__init__(
         batch_size=batch_size,
         discrete_rewards=discrete_rewards,
-        initial_observation_stream=initial_observation_stream,
+        history_stream=history_stream,
     )
 
     self.seed()
 
   def initialize_environments(self,
-                              initial_observation_stream,
+                              history_stream,
                               batch_size=1,
                               parallelism=1):
     """Initializes the environments.
 
     Args:
-      initial_observation_stream: Iterator yielding batches of initial
-        observations for the model.
+      history_stream: Iterator yielding batches of initial input data for the
+        model. The format is implementation-specific.
       batch_size: (int) Number of environments in a batch.
       parallelism: (int) Unused.
     """
     del parallelism
 
     model_state = trax.restore_state(self._output_dir)
-    self._model_params = model_state.opt_state.params
-    self._initial_observation_stream = initial_observation_stream
+    model_params = model_state.opt_state.params
+    self._predict_fn = functools.partial(
+        self._model_predict,
+        params=model_params,
+    )
+    self._history_stream = history_stream
 
-    self._history = None
-    self._steps = np.zeros(batch_size)
+    self._steps = np.zeros(batch_size, dtype=np.int32)
 
   @property
   def observation_space(self):
@@ -134,6 +119,37 @@ def seed(self, seed=None):
     self._rng = jax_random.get_prng(seed)
     return super(SimulatedEnvProblem, self).seed(seed=seed)
 
+  def _reset_model(self, predict_fn, indices, history, rng):
+    """Resets the environments at the given indices.
+
+    Should be implemented in subclasses.
+
+    Args:
+      predict_fn: Function running prediction with the model.
+      indices: List of indices of underlying envs to call reset on.
+      history: Initial input data for the model.
+      rng: Jax RNG.
+
+    Returns:
+      np.ndarray of batched observations from the reset envs.
+    """
+    raise NotImplementedError
+
+  def _step_model(self, predict_fn, actions, rng):
+    """Takes a step in all environments.
+
+    Should be implemented in subclasses.
+
+    Args:
+      predict_fn: Function running prediction with the model.
+      actions: (np.ndarray) with first dimension equal to the batch size.
+      rng: Jax RNG.
+
+    Returns:
+      a tuple of batched raw observations, rewards and dones.
+    """
+    raise NotImplementedError
+
   def _reset(self, indices):
     """Resets environments at the given indices.
 
@@ -143,7 +159,69 @@ def _reset(self, indices):
     Returns:
       np.ndarray of batched observations from the reset envs.
     """
-    history = next(self._initial_observation_stream)
+    history = next(self._history_stream)
+    (subrng, self._rng) = jax_random.split(self._rng)
+    return self._reset_model(self._predict_fn, indices, history, subrng)
+
+  def _step(self, actions):
+    """Takes a step in all environments.
+
+    Args:
+      actions: (np.ndarray) with first dimension equal to the batch size.
+
+    Returns:
+      a tuple of batched raw observations, raw rewards, dones and infos.
+    """
+    # Predict the next observation.
+    (subrng, self._rng) = jax_random.split(self._rng)
+    (observation, reward, done) = self._step_model(
+        self._predict_fn, actions, subrng)
+    return (observation, reward, done, {})
+
+
+class RawSimulatedEnvProblem(SimulatedEnvProblem):
+  """SimulatedEnvProblem running a model operating on raw tensors.
+
+  Wraps an autoregressive TRAX model of signature
+  (observation_history, action) -> (observation, reward) in an EnvProblem.
+  The model is assumed to take a fixed number of last observations as input
+  and produce a single observation, which is fed back into the model in the
+  next environment step.
+
+  Shape requirements (without the batch dimension):
+    observation: Consistent with observation_space.
+    observation_history: (history_length,) + observation.shape.
+    action: Consistent with action_space.
+    reward: (1,). The singleton dimension is removed in step().
+  """
+
+  def __init__(self, history_length, trajectory_length, *args, **kwargs):
+    """Initializes the env.
+
+    Args:
+      history_length: (int) Number of last observations fed into the model.
+      trajectory_length: (int) Length of each trajectory unrolled from the
+        model.
+      *args: (tuple) Positional arguments passed to the base class.
+      **kwargs: (dict) Keyword arguments passed to the base class.
+    """
+    self._history_length = history_length
+    self._trajectory_length = trajectory_length
+    self._history = None
+    self._steps = None
+
+    super(RawSimulatedEnvProblem, self).__init__(*args, **kwargs)
+
+  def initialize_environments(self, batch_size=1, **kwargs):
+    """Initializes the environments."""
+    self._history = None
+    self._steps = np.zeros(batch_size)
+    return super(RawSimulatedEnvProblem, self).initialize_environments(
+        batch_size=batch_size, **kwargs)
+
+  def _reset_model(self, predict_fn, indices, history, rng):
+    del predict_fn
+    del rng
     assert history.shape == ((self._batch_size, self._history_length) +
                              self.observation_space.shape)
 
@@ -161,20 +239,8 @@ def _reset(self, indices):
     # Return just the last timestep at the given indices.
     return history[:, -1, ...]
 
-  def _step(self, actions):
-    """Takes a step in all environments.
-
-    Args:
-      actions: (np.ndarray) with first dimension equal to the batch size.
-
-    Returns:
-      a tuple of batched raw observations, raw rewards, dones and infos.
-    """
-    # Predict the next observation.
-    (subrng, self._rng) = jax_random.split(self._rng)
-    (observation, reward) = self._model_predict((self._history, actions),
-                                                params=self._model_params,
-                                                rng=subrng)
+  def _step_model(self, predict_fn, actions, rng):
+    (observation, reward) = predict_fn((self._history, actions), rng=rng)
 
     # Roll the history one timestep back and append the new observation.
     self._history = np.roll(self._history, shift=-1, axis=1)
@@ -188,4 +254,143 @@ def _step(self, actions):
     observation = observation.copy()
     # Reshape the rewards to get rid of the extra dimension.
     reward = np.squeeze(reward.copy(), axis=1)
-    return (observation, reward, done, {})
+    return (observation, reward, done)
+
+
+def index_range_2d(begin_indices, length):
+  # Take all indices along the first dimension. Add another axis that'll
+  # broadcast along the second one.
+  first_dim = np.arange(len(begin_indices))[:, None]
+  # Take a range of indices along the second dimension. Offset it by
+  # begin_indices.
+  # TODO(pkozakowski): This materializes all indices of elements along the
+  # second dimension. Do it more efficiently if needed.
+  second_dim = np.arange(length)[None, :] + begin_indices[:, None]
+  return (first_dim, second_dim)
+
+
+def index_slice(indices):
+  first_dim = np.arange(len(indices))[:, None]
+  second_dim = indices[:, None]
+  return (first_dim, second_dim)
+
+
+class SerializedSequenceSimulatedEnvProblem(SimulatedEnvProblem):
+  """SimulatedEnvProblem running a model operating on sequences of symbols.
+
+  Wraps an autoregressive TRAX model of signature past_symbols -> symbol_probs
+  in an EnvProblem. The model is assumed to take a sequence of symbols as input
+  and produce distributions over all symbols in the sequence. The next symbol
+  is sampled and fed back to the model in the next decoding step.
+
+  Shape requirements (without the batch dimension):
+    past_symbols: (max_trajectory_length * L,)
+    symbol_probs: (max_trajectory_length * L, vocab_size)
+  where L is the representation length of one environment step.
+
+  Observations, actions, rewards and done flags are (de)serialized from/to
+  sequences of symbols using an EnvSerializer passed to the constructor.
+  """
+
+  def __init__(self, reward_fn, done_fn, vocab_size, max_trajectory_length,
+               *args, **kwargs):
+    """Initializes the env.
+
+    Args:
+      reward_fn: Function (previous_observation, current_observation) -> reward.
+      done_fn: Function (previous_observation, current_observation) -> done.
+      vocab_size: (int) Number of symbols in the vocabulary.
+      max_trajectory_length: (int) Maximum length of a trajectory unrolled from
+        the model.
+      *args: (tuple) Positional arguments passed to the base class.
+      **kwargs: (dict) Keyword arguments passed to the base class.
+    """
+    self._reward_fn = reward_fn
+    self._done_fn = done_fn
+    self._vocab_size = vocab_size
+    self._max_trajectory_length = max_trajectory_length
+    self._history = None
+    self._steps = None
+    self._observation_space = None
+    self._action_space = None
+    self._last_observations = None
+
+    super(SerializedSequenceSimulatedEnvProblem, self).__init__(*args, **kwargs)
+
+  def initialize_environments(self, batch_size=1, **kwargs):
+    """Initializes the environments."""
+    self._obs_serializer = space_serializer.create(
+        self.observation_space, self._vocab_size)
+    self._action_serializer = space_serializer.create(
+        self.action_space, self._vocab_size)
+    self._obs_repr_length = self._obs_serializer.representation_length
+    self._action_repr_length = self._action_serializer.representation_length
+    self._step_repr_length = self._obs_repr_length + self._action_repr_length
+    self._history = np.zeros((
+        batch_size,
+        self._max_trajectory_length * self._step_repr_length
+    ), dtype=np.int32)
+    self._steps = np.zeros(batch_size, dtype=np.int32)
+    self._last_observations = np.full(
+        (batch_size,) + self._observation_space.shape, np.nan)
+    return super(
+        SerializedSequenceSimulatedEnvProblem, self
+    ).initialize_environments(batch_size=batch_size, **kwargs)
+
+  @property
+  def _obs_repr_indices(self):
+    begin_indices = self._step_repr_length * self._steps
+    return index_range_2d(begin_indices, self._obs_repr_length)
+
+  @property
+  def _action_repr_indices(self):
+    begin_indices = self._step_repr_length * self._steps + self._obs_repr_length
+    return index_range_2d(begin_indices, self._action_repr_length)
+
+  def _predict_obs(self, predict_fn, rng):
+    def gumbel_sample(log_probs):
+      u = np.random.uniform(low=1e-6, high=1.0 - 1e-6, size=log_probs.shape)
+      g = -np.log(-np.log(u))
+      return np.argmax(log_probs + g, axis=-1)
+
+    for (i, subrng) in enumerate(jax_random.split(rng, self._obs_repr_length)):
+      symbol_index = self._steps * self._step_repr_length + i
+      log_probs = predict_fn(self._history, rng=subrng)[:, symbol_index, :]
+      self._history[:, symbol_index] = gumbel_sample(log_probs)
+
+    obs_repr = self._history[self._obs_repr_indices]
+    return self._obs_serializer.deserialize(obs_repr)
+
+  def _reset_model(self, predict_fn, indices, history, rng):
+    # TODO(pkozakowski): Random starts.
+    del history
+
+    self._steps[indices] = 0
+    observation = self._predict_obs(predict_fn, rng)[indices]
+    self._last_observations[indices] = observation
+    return observation
+
+  def _step_model(self, predict_fn, actions, rng):
+    action_repr = self._action_serializer.serialize(actions)
+    self._history[self._action_repr_indices] = action_repr
+    self._steps += 1
+    observation = self._predict_obs(predict_fn, rng)
+    reward = self._reward_fn(self._last_observations, observation)
+    done = self._done_fn(self._last_observations, observation)
+    self._last_observations = observation
+    done = np.logical_or(done, self._steps == self._max_trajectory_length)
+    return (observation, reward, done)
+
+
+def cartpole_done_fn(previous_observation, current_observation):
+  del previous_observation
+  x_threshold = 2.4
+  theta_threshold = 12 * 2 * np.pi / 360
+  x = current_observation[:, 0]
+  theta = current_observation[:, 2]
+  return np.logical_or(np.abs(x) > x_threshold, np.abs(theta) > theta_threshold)
+
+
+def cartpole_reward_fn(previous_observation, current_observation):
+  done = cartpole_done_fn(previous_observation, current_observation)
+  return 1.0 - done  # Unit reward for every timestep until the end.
diff --git a/tensor2tensor/trax/rlax/simulated_env_problem_test.py b/tensor2tensor/trax/rlax/simulated_env_problem_test.py
index db2c6ceb9..f4c895889 100644
--- a/tensor2tensor/trax/rlax/simulated_env_problem_test.py
+++ b/tensor2tensor/trax/rlax/simulated_env_problem_test.py
@@ -19,6 +19,9 @@
 from __future__ import division
 from __future__ import print_function
 
+import itertools
+
+import gin
 import gym
 import mock
 import numpy as np
@@ -29,25 +32,25 @@
 from tensorflow import test
 
 
-class SimulatedEnvProblemTest(test.TestCase):
+class RawSimulatedEnvProblemTest(test.TestCase):
 
   @staticmethod
   @mock.patch.object(trax, "restore_state", autospec=True)
-  def _create_env(mock_restore_state, model, initial_observations,
+  def _create_env(mock_restore_state, model, histories,
                   trajectory_length):
     # (model_params, opt_state)
     mock_restore_state.return_value.params = (None, None)
     space = gym.spaces.Discrete(100)
-    return simulated_env_problem.SimulatedEnvProblem(
+    return simulated_env_problem.RawSimulatedEnvProblem(
         model=model,
-        history_length=initial_observations.shape[2],
+        history_length=histories.shape[2],
         trajectory_length=trajectory_length,
         batch_size=1,
         observation_space=space,
         action_space=space,
         reward_range=(-1, 1),
         discrete_rewards=True,
-        initial_observation_stream=iter(initial_observations),
+        history_stream=iter(histories),
         output_dir=None,
     )
 
@@ -67,7 +70,7 @@ def mock_transition(inputs, *args, **kwargs):
     mock_model = mock_model_fn.return_value
 
     actions_to_take = np.array([[1], [3]])
-    initial_observations = np.array([[[0, 1, 2, 3]]])
+    histories = np.array([[[0, 1, 2, 3]]])
     expected_observations = np.array([[3], [4], [7]])
     expected_rewards = np.array([[1], [0]])
     expected_dones = np.array([[False], [True]])
@@ -77,7 +80,7 @@ def mock_transition(inputs, *args, **kwargs):
     with backend.use_backend("numpy"):
       env = self._create_env(  # pylint: disable=no-value-for-parameter
           model=mock_model_fn,
-          initial_observations=initial_observations,
+          histories=histories,
           trajectory_length=len(actions_to_take),
       )
       actual_observations = [env.reset()]
@@ -102,13 +105,13 @@ def mock_transition(inputs, *args, **kwargs):
     np.testing.assert_array_equal(actual_histories, expected_histories)
     np.testing.assert_array_equal(actual_actions, expected_actions)
 
-  def test_takes_new_initial_frames(self):
-    initial_observations = np.array([[[0, 1, 2]], [[3, 4, 5]]])
+  def test_takes_new_history(self):
+    histories = np.array([[[0, 1, 2]], [[3, 4, 5]]])
 
     with backend.use_backend("numpy"):
       env = self._create_env(  # pylint: disable=no-value-for-parameter
           model=mock.MagicMock(),
-          initial_observations=initial_observations,
+          histories=histories,
           trajectory_length=2,
       )
       env.reset()
@@ -116,5 +119,67 @@ def test_takes_new_initial_frames(self):
       np.testing.assert_array_equal(observation, [5])
 
 
+class SerializedSequenceSimulatedEnvProblemTest(test.TestCase):
+
+  @mock.patch.object(trax, "restore_state", autospec=True)
+  def test_communicates_with_model(self, mock_restore_state):
+    gin.bind_parameter("BoxSpaceSerializer.precision", 1)
+    vocab_size = 16
+    # Mock model predicting a fixed sequence of symbols. It is made such that
+    # the first two observations are equal and the last one is different.
+    symbols = [
+        1, 1, 2, 2,  # obs1
+        1, 1, 2, 2,  # obs2
+        1, 2, 2, 1,  # obs3
+    ]
+    def make_prediction(symbol):
+      one_hot = np.eye(vocab_size)[symbol]
+      log_probs = (1 - one_hot) * -100.0  # Virtually deterministic.
+      # (4 obs symbols + 1 action symbol) * 3 timesteps = 15.
+      return np.array([[log_probs] * 15])
+
+    mock_model_fn = mock.MagicMock()
+    mock_model = mock_model_fn.return_value
+    mock_model.side_effect = map(make_prediction, symbols)
+
+    with backend.use_backend("numpy"):
+      # (model_params, opt_state)
+      mock_restore_state.return_value.params = (None, None)
+      env = simulated_env_problem.SerializedSequenceSimulatedEnvProblem(
+          model=mock_model_fn,
+          reward_fn=(lambda _1, _2: np.array([0.5])),
+          done_fn=(lambda _1, _2: np.array([False])),
+          vocab_size=vocab_size,
+          max_trajectory_length=3,
+          batch_size=1,
+          observation_space=gym.spaces.Box(low=0, high=5, shape=(4,)),
+          action_space=gym.spaces.Discrete(2),
+          reward_range=(-1, 1),
+          discrete_rewards=False,
+          history_stream=itertools.repeat(None),
+          output_dir=None,
+      )
+      obs1 = env.reset()
+      ((inputs,), _) = mock_model.call_args
+
+      act1 = 0
+      (obs2, reward, done, _) = env.step(np.array([act1]))
+      ((inputs,), _) = mock_model.call_args
+      self.assertEqual(inputs[0, 4], act1)
+      np.testing.assert_array_equal(inputs[0, :4], symbols[:4])
+      np.testing.assert_array_equal(obs1, obs2)
+      np.testing.assert_array_equal(reward, [0.5])
+      np.testing.assert_array_equal(done, [False])
+
+      act2 = 1
+      (obs3, reward, done, _) = env.step(np.array([act2]))
+      ((inputs,), _) = mock_model.call_args
+      self.assertEqual(inputs[0, 9], act2)
+      np.testing.assert_array_equal(inputs[0, 5:9], symbols[4:8])
+      self.assertFalse(np.array_equal(obs2, obs3))
+      np.testing.assert_array_equal(reward, [0.5])
+      np.testing.assert_array_equal(done, [False])
+
+
 if __name__ == "__main__":
   test.main()
diff --git a/tensor2tensor/trax/rlax/space_serializer.py b/tensor2tensor/trax/rlax/space_serializer.py
new file mode 100644
index 000000000..220fb7d5b
--- /dev/null
+++ b/tensor2tensor/trax/rlax/space_serializer.py
@@ -0,0 +1,151 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Serialization of elements of Gym spaces into discrete sequences."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gin
+import gym
+import numpy as np
+
+
+class SpaceSerializer(object):
+  """Base class for Gym space serializers.
+
+  Attrs:
+    space_type: (type) Gym space class that this SpaceSerializer corresponds
+      to. Should be defined in subclasses.
+    representation_length: (int) Number of symbols in the representation of
+      every element of the space.
+  """
+
+  space_type = None
+  representation_length = None
+
+  def __init__(self, space, vocab_size):
+    """Creates a SpaceSerializer.
+
+    Subclasses should retain the signature.
+
+    Args:
+      space: (gym.Space) Gym space of type self.space_type.
+      vocab_size: (int) Number of symbols in the vocabulary.
+    """
+    assert isinstance(space, self.space_type)
+    self._space = space
+    self._vocab_size = vocab_size
+
+  def serialize(self, data):
+    """Serializes a batch of space elements into a discrete sequences.
+
+    Should be defined in subclasses.
+
+    Args:
+      data: A batch of batch_size elements of the Gym space to be serialized.
+
+    Returns:
+      int32 array of shape (batch_size, self.representation_length).
+    """
+    raise NotImplementedError
+
+  def deserialize(self, representation):
+    """Deserializes a batch of discrete sequences into space elements.
+
+    Should be defined in subclasses.
+
+    Args:
+      representation: int32 Numpy array of shape
+        (batch_size, self.representation_length) to be deserialized.
+
+    Returns:
+      A batch of batch_size deserialized elements of the Gym space.
+    """
+    raise NotImplementedError
+
+
+def create(space, vocab_size):
+  """Creates a SpaceSerializer for the given Gym space."""
+  return {
+      gym.spaces.Box: BoxSpaceSerializer,
+      gym.spaces.Discrete: DiscreteSpaceSerializer,
+  }[type(space)](space, vocab_size)
+
+
+@gin.configurable(whitelist=["precision"])
+class BoxSpaceSerializer(SpaceSerializer):
+  """Serializer for gym.spaces.Box.
+
+  Assumes that the space is bounded. Internally rescales it to the [0, 1]
+  interval and uses a fixed-precision encoding.
+  """
+
+  space_type = gym.spaces.Box
+
+  def __init__(self, space, vocab_size, precision):
+    super(BoxSpaceSerializer, self).__init__(space, vocab_size)
+    assert space.is_bounded(), "Only bounded spaces are supported."
+    self._precision = precision
+
+  def serialize(self, data):
+    array = data
+    batch_size = array.shape[0]
+    array = (array - self._space.low) / (self._space.high - self._space.low)
+    digits = []
+    for digit_index in range(-1, -self._precision - 1, -1):
+      threshold = self._vocab_size ** digit_index
+      digit = np.array(array / threshold).astype(np.int32) % self._vocab_size
+      digits.append(digit)
+      array -= digit * threshold
+    digits = np.stack(digits, axis=-1)
+    return np.reshape(digits, (batch_size, -1))
+
+  def deserialize(self, representation):
+    digits = representation
+    batch_size = digits.shape[0]
+    digits = np.reshape(digits, (batch_size, -1, self._precision))
+    array = np.zeros(digits.shape[:-1])
+    for digit_index_in_seq in range(self._precision):
+      digit_index = -digit_index_in_seq - 1
+      array += self._vocab_size ** digit_index * digits[..., digit_index_in_seq]
+    array = np.reshape(array, (batch_size,) + self._space.shape)
+    return array * (self._space.high - self._space.low) + self._space.low
+
+  @property
+  def representation_length(self):
+    return self._precision * self._space.low.size
+
+
+class DiscreteSpaceSerializer(SpaceSerializer):
+  """Serializer for gym.spaces.Discrete.
+
+  Assumes that the size of the space fits in the number of symbols.
+  """
+
+  space_type = gym.spaces.Discrete
+  representation_length = 1
+
+  def __init__(self, space, vocab_size):
+    super(DiscreteSpaceSerializer, self).__init__(space, vocab_size)
+    assert space.n <= vocab_size, (
+        "Discrete space size should fit in the number of symbols.")
+
+  def serialize(self, data):
+    return np.reshape(data, (-1, 1))
+
+  def deserialize(self, representation):
+    return np.reshape(representation, -1)
diff --git a/tensor2tensor/trax/rlax/space_serializer_test.py b/tensor2tensor/trax/rlax/space_serializer_test.py
new file mode 100644
index 000000000..0c62b3bbe
--- /dev/null
+++ b/tensor2tensor/trax/rlax/space_serializer_test.py
@@ -0,0 +1,94 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.rlax.space_serializer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gin
+import gym
+import numpy as np
+
+from tensor2tensor.trax.rlax import space_serializer
+from tensorflow import test
+
+
+class BoxSpaceSerializerTest(test.TestCase):
+
+  def setUp(self):
+    super(BoxSpaceSerializerTest, self).setUp()
+    # Enough precision to represent float32s accurately.
+    gin.bind_parameter("BoxSpaceSerializer.precision", 4)
+    self._space = gym.spaces.Box(low=-10, high=10, shape=(2,))
+    self._serializer = space_serializer.create(
+        self._space,
+        # Weird vocab_size to test that it doesn't only work with powers of 2.
+        vocab_size=257)
+
+  def _sample_batch(self):
+    return np.reshape(self._space.sample(), (1,) + self._space.shape)
+
+  def test_representation_length(self):
+    input_array = self._sample_batch()
+    representation = self._serializer.serialize(input_array)
+    self.assertEqual(
+        representation.shape, (1, self._serializer.representation_length))
+
+  def test_commutes(self):
+    input_array = self._sample_batch()
+    representation = self._serializer.serialize(input_array)
+    output_array = self._serializer.deserialize(representation)
+    np.testing.assert_array_almost_equal(input_array, output_array)
+
+  def test_representation_changes(self):
+    array1 = self._sample_batch()
+    array2 = -array1
+    (repr1, repr2) = tuple(map(self._serializer.serialize, (array1, array2)))
+    self.assertFalse(np.array_equal(repr1, repr2))
+
+
+class DiscreteSpaceSerializerTest(test.TestCase):
+
+  def setUp(self):
+    super(DiscreteSpaceSerializerTest, self).setUp()
+    self._space = gym.spaces.Discrete(n=2)
+    self._serializer = space_serializer.create(self._space, vocab_size=2)
+
+  def _sample_batch(self):
+    return np.reshape(self._space.sample(), (1,) + self._space.shape)
+
+  def test_representation_length(self):
+    input_array = self._sample_batch()
+    representation = self._serializer.serialize(input_array)
+    self.assertEqual(
+        representation.shape, (1, self._serializer.representation_length))
+
+  def test_commutes(self):
+    input_array = self._sample_batch()
+    representation = self._serializer.serialize(input_array)
+    output_array = self._serializer.deserialize(representation)
+    np.testing.assert_array_almost_equal(input_array, output_array)
+
+  def test_representation_changes(self):
+    array1 = self._sample_batch()
+    array2 = 1 - array1
+    (repr1, repr2) = tuple(map(self._serializer.serialize, (array1, array2)))
+    self.assertFalse(np.array_equal(repr1, repr2))
+
+
+if __name__ == "__main__":
+  test.main()

From 5882ccd2e67ad5571052a1bd422ccc6f794a59d6 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 12 Aug 2019 15:47:50 -0700
Subject: [PATCH 2278/2720] Add an infer() implementation to the aligned model.

PiperOrigin-RevId: 263022175
---
 tensor2tensor/models/research/aligned.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/tensor2tensor/models/research/aligned.py b/tensor2tensor/models/research/aligned.py
index 26a857647..edea1218a 100644
--- a/tensor2tensor/models/research/aligned.py
+++ b/tensor2tensor/models/research/aligned.py
@@ -224,6 +224,23 @@ def _pseudolocal_bias(x):
     decoder_output = dp(tf.expand_dims, x, 2)
     return decoder_output, extra_loss
 
+  def infer(self,
+            features=None,
+            decode_length=1,
+            beam_size=1,
+            top_beams=1,
+            alpha=0.0,
+            use_tpu=False):
+    """Predict."""
+    features["targets"] = tf.identity(features["inputs"])
+    logits, _ = self(features)
+    log_probs = common_layers.log_prob_from_logits(logits)
+    predictions, scores = common_layers.argmax_with_score(log_probs)
+    return {
+        "outputs": predictions,
+        "scores": scores,
+    }
+
 
 def get_batch_coordinate(x):
   """Return a flat int32 tensor of shape [1, batch_size*length, 1]."""
@@ -247,6 +264,7 @@ def aligned_base():
     a hparams object
   """
   hparams = common_hparams.basic_params1()
+  hparams.force_full_predict = True
   hparams.hidden_size = 512
   hparams.batch_size = 5000
   hparams.max_length = 0

From 89f12fe9b45e9d5f30e97407407f1e277e37a031 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 12 Aug 2019 16:35:01 -0700
Subject: [PATCH 2279/2720] Use dropout in reformer

PiperOrigin-RevId: 263030896
---
 .../transformer_revnet_imagenet64_8gb.gin     |  2 +-
 .../models/research/transformer_revnet.py     | 51 +++++++++++++++----
 2 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
index 54efae800..26ba18308 100644
--- a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
@@ -52,7 +52,7 @@ TransformerRevnetLM.d_model = 1024
 TransformerRevnetLM.d_ff = 2048
 TransformerRevnetLM.d_attention_key = 32
 TransformerRevnetLM.d_attention_value = 32
-TransformerRevnetLM.dropout = 0.1
+TransformerRevnetLM.dropout = 0.0
 TransformerRevnetLM.max_len = 12288  # 64 * 64 * 3
 TransformerRevnetLM.mode = 'train'
 TransformerRevnetLM.n_heads = 4
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index b77001121..2cfec5fea 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -81,17 +81,36 @@ def new_parameters(self, input_shape, input_dtype, rng):
     return self._layer.initialize(first_shape, input_dtype[0], rng)
 
 
+@tl.layer()
+def BroadcastedDropout(x, params, rate=0.0, mode='train', broadcast_dims=(-2,),
+                       rng=None, **kwargs):
+  """Dropout, with broadcasting to save memory."""
+  del params, kwargs
+  if rng is None:
+    raise ValueError('BroadcastedDropout requires rng kwarg.')
+  if rate >= 1.0:
+    raise ValueError('Dropout rate (%f) must be lower than 1.' % rate)
+  if mode == 'train' and rate > 0.0:
+    noise_shape = list(x.shape)
+    for dim in broadcast_dims:
+      noise_shape[dim] = 1
+    keep_prob = jax.lax.tie_in(rng, 1.0 - rate)
+    keep = backend.random.bernoulli(rng, keep_prob, tuple(noise_shape))
+    multiplier = keep.astype(x.dtype) / jax.lax.tie_in(keep, keep_prob)
+    return x * multiplier
+  else:
+    return x
+
+
 def FeedForward(d_model, d_ff, dropout, mode):
   """Feed-forward block with layer normalization at start."""
-  # TODO(kitaev): add dropout. Dropout is typically performed by adding noise to
-  # the activations, but when the size of the activations is very large it is
-  # more efficient to add noise to the *parameters* instead.
-  del dropout, mode
   return [
       tl.LayerNorm(),
       tl.Dense(d_ff),
+      BroadcastedDropout(rate=dropout, mode=mode),  # pylint: disable=no-value-for-parameter
       tl.Relu(),
       tl.Dense(d_model),
+      BroadcastedDropout(rate=dropout, mode=mode),  # pylint: disable=no-value-for-parameter
   ]
 
 
@@ -283,15 +302,18 @@ class MemoryEfficientDotProductAttention(DotProductAttention):
   def __init__(self, loop_stride, dropout, mode):
     super(MemoryEfficientDotProductAttention, self).__init__(dropout, mode)
     self._loop_stride = loop_stride
-    # TODO(kitaev): implement attention dropout
-    assert dropout is None or dropout == 0.0, (
-        'Dropout is not implemented in MemoryEfficientDotProductAttention.')
+    if dropout >= 1.0:
+      raise ValueError('Dropout rates must be lower than 1.')
+    if mode == 'train':
+      self.dropout = dropout
+    else:
+      self.dropout = None
 
   def call(self, inputs, params=(), **kwargs):
     output, _ = self.forward_and_vjp(inputs, None, params=params, **kwargs)
     return output
 
-  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
+  def forward_and_vjp(self, inputs, ct, params=(), rng=None, **kwargs):
     # This is the core of the memory-efficient attention implementation, where
     # we use the jax.lax.while_loop primitive to compute attention for a small
     # set of query positions at a time. Note how in the backwards pass, we
@@ -330,6 +352,16 @@ def forward_slice(query_slice, q_loop_idx, key, value):
       # Softmax.
       dots = np.exp(dots - dots.max(axis=-1, keepdims=True))
       dots = dots / dots.sum(axis=-1, keepdims=True)
+
+      if self.dropout is not None and self.dropout > 0.0:
+        # Dropout is broadcast across the batch and head dimensions
+        dropout_shape = (1, 1, dots.shape[-2], dots.shape[-1])
+        slice_rng = jax.random.fold_in(rng, q_loop_idx)
+        keep_prob = jax.lax.tie_in(dots, 1.0 - self.dropout)
+        keep = backend.random.bernoulli(slice_rng, keep_prob, dropout_shape)
+        multiplier = keep.astype(dots.dtype) / jax.lax.tie_in(keep, keep_prob)
+        dots = dots * multiplier
+
       out_slice = np.matmul(dots, value)
       return out_slice
 
@@ -679,6 +711,7 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
       JoinHeads(),  # pylint: disable=no-value-for-parameter
       tl.Dense(d_model),
       Unchunk(n_sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
+      BroadcastedDropout(rate=dropout, mode=mode),  # pylint: disable=no-value-for-parameter
   ]
 
   feed_forward = [
@@ -727,7 +760,7 @@ def TransformerRevnetLM(vocab_size,
   """
   positional_embedder = [
       tl.Embedding(d_model, vocab_size),
-      # TODO(kitaev): add dropout
+      BroadcastedDropout(rate=dropout, mode=mode),  # pylint: disable=no-value-for-parameter
       tl.PositionalEncoding(max_len=max_len),
   ]
   return tl.Model(

From 4a32adced5e2bea10c48193c49b56db7ef85fa33 Mon Sep 17 00:00:00 2001
From: Joel Shor <joelshor@google.com>
Date: Tue, 13 Aug 2019 11:24:44 -0700
Subject: [PATCH 2280/2720] Internal change

PiperOrigin-RevId: 263177751
---
 docs/cloud_mlengine.md                 |  90 ---------
 docs/cloud_tpu.md                      |  50 -----
 docs/distributed_training.md           | 220 ----------------------
 docs/index.md                          | 127 -------------
 docs/multi_problem.md                  | 188 -------------------
 docs/new_model.md                      | 114 ------------
 docs/new_problem.md                    | 243 -------------------------
 docs/overview.md                       | 175 ------------------
 docs/tutorials/asr_with_transformer.md |   4 -
 tensor2tensor/models/video/savp.py     |   3 +-
 10 files changed, 2 insertions(+), 1212 deletions(-)
 delete mode 100644 docs/cloud_mlengine.md
 delete mode 100644 docs/cloud_tpu.md
 delete mode 100644 docs/distributed_training.md
 delete mode 100644 docs/index.md
 delete mode 100644 docs/multi_problem.md
 delete mode 100644 docs/new_model.md
 delete mode 100644 docs/new_problem.md
 delete mode 100644 docs/overview.md
 delete mode 100644 docs/tutorials/asr_with_transformer.md

diff --git a/docs/cloud_mlengine.md b/docs/cloud_mlengine.md
deleted file mode 100644
index 83ebe7e57..000000000
--- a/docs/cloud_mlengine.md
+++ /dev/null
@@ -1,90 +0,0 @@
-# Running on Cloud ML Engine
-
-Google Cloud Platform offers a managed training environment for TensorFlow
-models called [Cloud ML Engine](https://cloud.google.com/ml-engine/) and
-you can easily launch Tensor2Tensor on it, including for hyperparameter tuning.
-
-# Launch
-
-It's the same `t2t-trainer` you know and love with the addition of the
-`--cloud_mlengine` flag, which by default will launch on a 1-GPU machine
-in the default compute region. See the
-[docs for `gcloud compute`](https://cloud.google.com/compute/docs/gcloud-compute/#set_default_zone_and_region_in_your_local_client)
-to learn how to set the default compute region.
-
-```
-# Note that both the data dir and output dir have to be on GCS
-DATA_DIR=gs://my-bucket/data
-OUTPUT_DIR=gs://my-bucket/train
-t2t-trainer \
-  --problem=translate_ende_wmt32k \
-  --model=transformer \
-  --hparams_set=transformer_base \
-  --data_dir=$DATA_DIR \
-  --output_dir=$OUTPUT_DIR \
-  --cloud_mlengine
-```
-
-By passing `--worker_gpu=4` or `--worker_gpu=8` it will automatically launch on
-machines with 4 or 8 GPUs.
-
-You can additionally pass the `--cloud_mlengine_master_type` to select another
-kind of machine (see the [docs for
-`masterType`](https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#traininginput)
-for options, including
-[ML Engine machine
-types](https://cloud.google.com/ml-engine/docs/training-overview)
-and their
-[specs](https://cloud.google.com/compute/docs/machine-types)).
-If you provide this flag yourself, make sure you pass the
-correct value for `--worker_gpu` (for non-GPU machines, you should pass
-`--worker_gpu=0`).
-
-**Note**: `t2t-trainer` only currently supports launching with single machines,
-possibly with multiple GPUs. Multi-machine setups are not yet supported out of
-the box with the `--cloud_mlengine` flag, though multi-machine should in
-principle work just fine. Contributions/testers welcome.
-
-
-## `--t2t_usr_dir`
-
-Launching on Cloud ML Engine works with `--t2t_usr_dir` as well as long as the
-directory is fully self-contained (i.e. the imports only refer to other modules
-in the directory). If there are additional PyPI dependencies that you need, you
-can include a `requirements.txt` file in the directory specified by
-`t2t_usr_dir`.
-
-# Hyperparameter Tuning
-
-Hyperparameter tuning with `t2t-trainer` and Cloud ML Engine is also a breeze
-with `--hparams_range` and the `--autotune_*` flags:
-
-```
-t2t-trainer \
-  --problem=translate_ende_wmt32k \
-  --model=transformer \
-  --hparams_set=transformer_base \
-  --data_dir=$DATA_DIR \
-  --output_dir=$OUTPUT_DIR \
-  --cloud_mlengine \
-  --hparams_range=transformer_base_range \
-  --autotune_objective='metrics-translate_ende_wmt32k/neg_log_perplexity' \
-  --autotune_maximize \
-  --autotune_max_trials=100 \
-  --autotune_parallel_trials=3
-```
-
-The `--hparams_range` specifies the search space and should be registered with
-`@register_ranged_hparams`. It defines a `RangedHParams` object that sets
-search ranges and scales for various parameters. See `transformer_base_range`
-in
-[`transformer.py`](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/models/transformer.py)
-for an example.
-
-The metric name passed as `--autotune_objective` should be exactly what you'd
-see in TensorBoard. To minimize a metric, set `--autotune_maximize=False`.
-
-You control how many total trials to run with `--autotune_max_trials` and the
-number of jobs to launch in parallel with `--autotune_parallel_trials`.
-
-Happy tuning!
diff --git a/docs/cloud_tpu.md b/docs/cloud_tpu.md
deleted file mode 100644
index c0625e132..000000000
--- a/docs/cloud_tpu.md
+++ /dev/null
@@ -1,50 +0,0 @@
-# Running on Cloud TPUs
-
-Tensor2Tensor supports running on Google Cloud Platforms TPUs, chips
-specialized for ML training. See the official tutorials for [running the
-T2T Transformer for text on Cloud TPUs](https://cloud.google.com/tpu/docs/tutorials/transformer) and
-[Transformer for Speech Recognition](https://cloud.google.com/tpu/docs/tutorials/automated-speech-recognition).
-
-## Other models on TPU
-
-Many of Tensor2Tensor's models work on TPU.
-
-You can provision a VM and TPU with `ctpu up`. Use the `t2t-trainer` command
-on the VM as usual with the additional flags `--use_tpu` and
-`--cloud_tpu_name=$TPU_NAME`.
-
-Note that because the `TPUEstimator` does not catch the `OutOfRangeError`
-during evaluation, you should ensure that `--eval_steps` is small enough to
-not exhaust the evaluation data.
-
-A non-exhaustive list of T2T models that work on TPU:
-
-* Image generation: `imagetransformer` with `imagetransformer_base_tpu` (or
-  `imagetransformer_tiny_tpu`)
-* Super-resolution: `img2img_transformer` with `img2img_transformer_base_tpu`
-  (or `img2img_transformer_tiny_tpu`)
-* `resnet` with `resnet_50` (or `resnet_18` or `resnet_34`)
-* `revnet` with `revnet_104` (or `revnet_38_cifar`)
-* `shake_shake` with `shakeshake_tpu` (or `shakeshake_small`)
-
-## Example invocation
-
-Use `ctpu up` to bring up the VM and TPU machines; once the machines are ready
-it will SSH you into the VM and you can run the following:
-
-```
-# DATA_DIR and OUT_DIR should be GCS buckets
-# TPU_NAME should have been set automatically by the ctpu tool
-
-t2t-trainer \
-  --model=shake_shake \
-  --hparams_set=shakeshake_tpu \
-  --problem=image_cifar10 \
-  --train_steps=180000 \
-  --eval_steps=9 \
-  --local_eval_frequency=100 \
-  --data_dir=$DATA_DIR \
-  --output_dir=$OUT_DIR \
-  --use_tpu \
-  --cloud_tpu_name=$TPU_NAME
-```
diff --git a/docs/distributed_training.md b/docs/distributed_training.md
deleted file mode 100644
index f59974623..000000000
--- a/docs/distributed_training.md
+++ /dev/null
@@ -1,220 +0,0 @@
-# Distributed Training
-
-The `t2t-trainer` supports both synchronous and asynchronous distributed
-training.
-
-Note that it is almost always more efficient to train on a single machine with
-multiple GPUs/TPUs. Async training is less stable than sync training, and sync
-training is much faster on 1 machine than on multiple. For these reasons, we
-almost always train on single machines with multiple GPUs/TPUs.
-
-T2T uses TensorFlow Estimators and so distributed training is configured with
-the `TF_CONFIG` environment variable that is read by the
-[RunConfig](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/estimator/run_config.py)
-along with a set of flags that T2T uses to distribute the computation.
-
-## Shared output directory
-
-When using multiple machines, it is necessary that all nodes use the same
-`--output_dir`, which means that it should be set to a Google Cloud Storage
-bucket (`gs://...`) or a directory on a shared network filesystem.
-
-## Utility to produce `TF_CONFIG` and flags
-
-[`t2t-make-tf-configs`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/bin/t2t-make-tf-configs)
-generates the `TF_CONFIG` json strings and the necessary command-line flags for
-the jobs.
-
-Given a set of master and parameter server addresses, the script outputs, for
-each job, a line with the `TF_CONFIG` environment variable and the command-line
-flags necessary for distributed training. For each job, you should invoke the
-`t2t-trainer` with the `TF_CONFIG` value and flags that are output.
-
-## Eval jobs
-
-Eval jobs should set the following flags and do not need the `TF_CONFIG`
-environment variable to be set as the eval jobs run locally and do not
-communicate to the other jobs (the eval jobs read the model checkpoints that the
-trainer writes out):
-
-- `--schedule=continuous_eval_on_train_data` or
-  `--schedule=continuous_eval` (for dev data)
-- `--worker_job='/job:localhost'`
-- `--output_dir=$TRAIN_DIR`
-
-**Note that evaluation does not work distributed.** That is, distributed jobs
-should always use `--schedule=train`.
-
-## Examples
-
-### Sync training across multiple workers
-
-In this scenario, you wish to do synchronous training across multiple workers.
-Note that it is easier to simply use 1 worker with multiple GPUs and set
-`--worker_gpu=8`, but there may be cases where you may want to have multiple
-machines.
-
-You will need 1 `ip:port` for the master and then 1 `ip:port` for each worker.
-
-For this example we'll use 2 workers and these addresses:
-
-```
-# Master
-10.0.0.1:5555
-
-# Worker 1
-10.0.0.2:5555
-
-# Worker 2
-10.0.0.3:5555
-```
-
-Next we generate the `TF_CONFIG` and command-line-flags for each job.
-
-```
-$ t2t-make-tf-configs --masters='10.0.0.1:5555' --ps='10.0.0.2:5555,10.0.0.3:5555'
-Assuming SYNC distributed training with a single master and 2 workers
-'{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 0, "type": "master"}}'      --master=grpc://10.0.0.1:5555 --ps_replicas=2 --worker_replicas=1 --worker_gpu=0 --worker_id=0 --ps_gpu=1 --sync --schedule=train --worker_job='/job:master'
-'{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 0, "type": "ps"}}'  --schedule=run_std_server
-'{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 1, "type": "ps"}}'  --schedule=run_std_server
-```
-
-The output here is 1 line per job. Each line contains the `TF_CONFIG` to set
-for that job as well as the command-line flags to set for that job.
-
-It is a bit confusing that the workers are being passed to the `--ps` flag, but
-this is correct. When running in `--sync` mode, the `ps` are actually the
-workers. You can see in the next example below that when `--sync=False`, i.e.
-async mode, that the `ps` are in fact being used as parameter servers.
-
-Here's how we would start each job on their respective machines (the
-commands below assume that you're ssh'd into that job's machine):
-
-**Master**:
-
-```
-$ export TF_CONFIG='{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 0, "type": "master"}}'
-$ t2t-trainer \
-    --master=grpc://10.0.0.1:5555 \
-    --ps_replicas=2 \
-    --worker_replicas=1 \
-    --worker_gpu=0 \
-    --worker_id=0 \
-    --ps_gpu=1 \
-    --sync \
-    --schedule=train \
-    --worker_job='/job:master' \
-    --model=transformer \
-    --hparams_set=transformer_base \
-    --problem=translate_ende_wmt32k
-```
-
-**Worker 1**:
-
-```
-$ export TF_CONFIG='{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 0, "type": "ps"}}'
-$ t2t-trainer --schedule=run_std_server
-```
-
-**Worker 2**:
-
-```
-$ export TF_CONFIG='{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 1, "type": "ps"}}'
-$ t2t-trainer --schedule=run_std_server
-```
-
-Note that if you have more than 1 GPU on each worker machine, make sure to
-modify the `--ps_gpu` passed to the master.
-
-### Async training across multiple workers
-
-In this scenario, you wish to do asynchronous training across multiple workers
-with 1+ shared parameter servers.
-
-Note that async training is usually less stable than sync training and for that
-reason we almost always prefer sync training, but there may be cases where you
-want to do async distributed training.
-
-For this example we'll use 2 workers and 2 parameter servers:
-
-```
-# Worker 1
-10.0.0.1:5555
-
-# Worker 2
-10.0.0.2:5555
-
-# PS 1
-10.0.0.3:5555
-
-# PS 2
-10.0.0.4:5555
-```
-
-Next we generate the `TF_CONFIG` and command-line-flags for each job.
-
-```
-$ t2t-make-tf-configs --masters='10.0.0.1:5555,10.0.0.2:5555' --ps='10.0.0.3:5555,10.0.0.4:5555'
-Assuming ASYNC distributed training with 2 workers and 2 parameter servers
-'{"task": {"index": 0, "type": "chief"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}' --master=grpc://10.0.0.1:5555 --ps_replicas=2 --worker_replicas=2 --worker_gpu=1 --worker_id=0 --ps_gpu=0  --schedule=train --worker_job='/job:chief'
-'{"task": {"index": 0, "type": "worker"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'        --master=grpc://10.0.0.2:5555 --ps_replicas=2 --worker_replicas=2 --worker_gpu=1 --worker_id=1 --ps_gpu=0 --schedule=train --worker_job='/job:worker'
-'{"task": {"index": 0, "type": "ps"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'    --schedule=run_std_server
-'{"task": {"index": 1, "type": "ps"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'    --schedule=run_std_server
-```
-
-Here's how we would start each job on their respective machines (the
-commands below assume that you're ssh'd into that job's machine):
-
-**Worker 1**:
-
-```
-$ export TF_CONFIG='{"task": {"index": 0, "type": "chief"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'
-$ t2t-trainer \
-    --master=grpc://10.0.0.1:5555 \
-    --ps_replicas=2 \
-    --worker_replicas=2 \
-    --worker_gpu=1 \
-    --worker_id=0 \
-    --ps_gpu=0 \
-    --schedule=train \
-    --worker_job='/job:chief' \
-    --model=transformer \
-    --hparams_set=transformer_base \
-    --problem=translate_ende_wmt32k
-```
-
-**Worker 2**:
-
-```
-$ export TF_CONFIG='{"task": {"index": 0, "type": "worker"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'
-$ t2t-trainer \
-    --master=grpc://10.0.0.2:5555 \
-    --ps_replicas=2 \
-    --worker_replicas=2 \
-    --worker_gpu=1 \
-    --worker_id=1 \
-    --ps_gpu=0 \
-    --schedule=train \
-    --worker_job='/job:worker' \
-    --model=transformer \
-    --hparams_set=transformer_base \
-    --problem=translate_ende_wmt32k
-```
-
-**PS 1**:
-
-```
-$ export TF_CONFIG='{"task": {"index": 0, "type": "ps"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'
-$ t2t-trainer --schedule=run_std_server
-```
-
-**PS 2**:
-
-```
-$ export TF_CONFIG='{"task": {"index": 1, "type": "ps"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'
-$ t2t-trainer --schedule=run_std_server
-```
-
-Increase `--worker_gpu` on each of the workers if you have multiple GPUs. If the
-parameter servers are also using GPUs, set `--ps_gpu` to the number of GPUs on
-the parameter servers.
diff --git a/docs/index.md b/docs/index.md
deleted file mode 100644
index 26298a9d2..000000000
--- a/docs/index.md
+++ /dev/null
@@ -1,127 +0,0 @@
-# Tensor2Tensor Documentation
-
-[![PyPI
-version](https://badge.fury.io/py/tensor2tensor.svg)](https://badge.fury.io/py/tensor2tensor)
-[![GitHub
-Issues](https://img.shields.io/github/issues/tensorflow/tensor2tensor.svg)](https://github.com/tensorflow/tensor2tensor/issues)
-[![Contributions
-welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CONTRIBUTING.md)
-[![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
-[![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
-
-[Tensor2Tensor](https://github.com/tensorflow/tensor2tensor), or
-[T2T](https://github.com/tensorflow/tensor2tensor) for short, is a library
-of deep learning models and datasets designed to make deep learning more
-accessible and [accelerate ML
-research](https://research.googleblog.com/2017/06/accelerating-deep-learning-research.html).
-
-
-## Introduction
-
-* [Walkthrough](walkthrough.md): Install and run.
-* [IPython notebook](https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/hello_t2t.ipynb): Get a hands-on experience.
-
-## Basics
-
-* [Overview](overview.md): How all parts of T2T code are connected.
-* [New Problem](new_problem.md): Train T2T models on your data.
-* [New Model](new_model.md): Create your own T2T model.
-
-## Training in the cloud
-
-* [Training on Google Cloud ML](cloud_mlengine.md)
-* [Training on Google Cloud TPUs](cloud_tpu.md)
-* [Distributed Training](distributed_training.md)
-
-## Solving your task
-
-Below we list a number of tasks that can be solved with T2T when
-you train the appropriate model on the appropriate problem.
-We give the problem and model below and we suggest a setting of
-hyperparameters that we know works well in our setup. We usually
-run either on Cloud TPUs or on 8-GPU machines; you might need
-to modify the hyperparameters if you run on a different setup.
-
-### Image Classification
-
-For image classification, we have a number of standard data-sets:
-* ImageNet (a large data-set): `--problem=image_imagenet`, or one
-   of the re-scaled versions (`image_imagenet224`, `image_imagenet64`,
-   `image_imagenet32`)
-* CIFAR-10: `--problem=image_cifar10` (or
-    `--problem=image_cifar10_plain` to turn off data augmentation)
-* CIFAR-100: `--problem=image_cifar100`
-* MNIST: `--problem=image_mnist`
-
-For ImageNet, we suggest to use the ResNet or Xception, i.e.,
-use `--model=resnet --hparams_set=resnet_50` or
-`--model=xception --hparams_set=xception_base`.
-Resnet should get to above 76% top-1 accuracy on ImageNet.
-
-For CIFAR and MNIST, we suggest to try the shake-shake model:
-`--model=shake_shake --hparams_set=shakeshake_big`.
-This setting trained for `--train_steps=700000` should yield
-close to 97% accuracy on CIFAR-10.
-
-### Language Modeling
-
-For language modeling, we have these data-sets in T2T:
-* PTB (a small data-set): `--problem=languagemodel_ptb10k` for
-    word-level modeling and `--problem=languagemodel_ptb_characters`
-    for character-level modeling.
-* LM1B (a billion-word corpus): `--problem=languagemodel_lm1b32k` for
-    subword-level modeling and `--problem=languagemodel_lm1b_characters`
-    for character-level modeling.
-
-We suggest to start with `--model=transformer` on this task and use
-`--hparams_set=transformer_small` for PTB and
-`--hparams_set=transformer_base` for LM1B.
-
-### Sentiment Analysis
-
-For the task of recognizing the sentiment of a sentence, use
-* the IMDB data-set: `--problem=sentiment_imdb`
-
-We suggest to use `--model=transformer_encoder` here and since it is
-a small data-set, try `--hparams_set=transformer_tiny` and train for
-few steps (e.g., `--train_steps=2000`).
-
-### Speech Recognition
-
-For speech-to-text, we have these data-sets in T2T:
-* Librispeech (English speech to text): `--problem=librispeech` for
-    the whole set and `--problem=librispeech_clean` for a smaller
-    but nicely filtered part.
-
-### Summarization
-
-For summarizing longer text into shorter one we have these data-sets:
-* CNN/DailyMail articles summarized into a few sentences:
-  `--problem=summarize_cnn_dailymail32k`
-
-We suggest to use `--model=transformer` and
-`--hparams_set=transformer_prepend` for this task.
-This yields good ROUGE scores.
-
-### Translation
-
-There are a number of translation data-sets in T2T:
-* English-German: `--problem=translate_ende_wmt32k`
-* English-French: `--problem=translate_enfr_wmt32k`
-* English-Czech: `--problem=translate_encs_wmt32k`
-* English-Chinese: `--problem=translate_enzh_wmt32k`
-* English-Vietnamese: `--problem=translate_envi_iwslt32k`
-* English-Spanish: `--problem=translate_enes_wmt32k`
-
-You can get translations in the other direction by appending `_rev` to
-the problem name, e.g., for German-English use
-`--problem=translate_ende_wmt32k_rev`.
-
-For all translation problems, we suggest to try the Transformer model:
-`--model=transformer`. At first it is best to try the base setting,
-`--hparams_set=transformer_base`. When trained on 8 GPUs for 300K steps
-this should reach a BLEU score of about 28 on the English-German data-set,
-which is close to state-of-the art. If training on a single GPU, try the
-`--hparams_set=transformer_base_single_gpu` setting. For very good results
-or larger data-sets (e.g., for English-French), try the big model
-with `--hparams_set=transformer_big`.
diff --git a/docs/multi_problem.md b/docs/multi_problem.md
deleted file mode 100644
index d4e37d09d..000000000
--- a/docs/multi_problem.md
+++ /dev/null
@@ -1,188 +0,0 @@
-# Multi-problem training
-
-Multi-problem training is possible by defining [MultiProblem](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/multi_problem.py) sub-classes that specify a list of [Problem](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/problem.py) objects to include in training. In some cases, multi-problem training can be used to improve performance compared to training on individual problems.
-
-In the following sections we'll discuss MultiProblem from a usage perspective followed by that of someone wishing to build upon it.
-
-Please note the [T2T Walkthrough](https://github.com/tensorflow/tensor2tensor/blob/master/docs/walkthrough.md) documentation is a good place to start to understand the variety of component concepts we'll build on here.
-
-## Usage
-
-### Problem definition and datagen
-
-In this discussion we'll consider the following (large) multi-problem that includes ten different sub-problems. These include:
-
-1. A [language modeling](https://en.wikipedia.org/wiki/Language_model) [problem](https://github.com/tensorflow/tensor2tensor/blob/0dff89d64c3406d42717280cb9135a5ce7af793c/tensor2tensor/data_generators/wiki_lm.py#L223) operating on a corpus of German, English, French, and Romanian language wikipedia articles.
-2. Multiple compatible pairwise language translation problems (En -> De, En -> Fr, En -> Ro, De -> En, Fr -> En, Ro -> En)
-3. A compatible [version](https://github.com/tensorflow/tensor2tensor/blob/ef12bee72270b322165d073c39a650a189de39aa/tensor2tensor/data_generators/cnn_dailymail.py#L267) of the combined CNN/DailyMail news article summarization problem.
-4. A compatible [version](https://github.com/tensorflow/tensor2tensor/blob/ef12bee72270b322165d073c39a650a189de39aa/tensor2tensor/data_generators/multinli.py#L155) of the [MultiNLI](https://www.nyu.edu/projects/bowman/multinli/) textual entailment classification problem.
-5. A compatible [version](https://github.com/tensorflow/tensor2tensor/blob/1de13dbebccb415d89b0658e18a57e9607bafd32/tensor2tensor/data_generators/squad.py#L126) of the [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/) question/answer problem.
-
-```python
-
-@registry.register_problem
-class LanguagemodelMultiWikiTranslate(multi_problem.MultiProblem):
-  """Wiki multi-lingual LM and multiple translations."""
-
-  def __init__(self, was_reversed=False, was_copy=False):
-    super(LanguagemodelMultiWikiTranslate, self).__init__(
-        was_reversed, was_copy)
-    self.task_list.append(wiki_lm.LanguagemodelDeEnFrRoWiki64k())
-    self.task_list.append(translate_ende.TranslateEndeWmtMulti64k())
-    self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k())
-    self.task_list.append(translate_enro.TranslateEnroWmtMultiTiny64k())
-    self.task_list.append(translate_ende.TranslateEndeWmtMulti64k(
-        was_reversed=True))
-    self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k(
-        was_reversed=True))
-    self.task_list.append(translate_enro.TranslateEnroWmtMultiTiny64k(
-        was_reversed=True))
-    self.task_list.append(
-        cnn_dailymail.SummarizeCnnDailymailWikiLMMultiVocab64k())
-    self.task_list.append(multinli.MultiNLIWikiLMMultiVocab64k())
-    self.task_list.append(squad.SquadConcatMulti64k())
-
-  @property
-  def vocab_type(self):
-    return text_problems.VocabType.SUBWORD
-
-```
-
-The word "compatible" was used a lot above! That's because each of these problems have been modified to use the vocabulary produced by the Wikipedia-based language modeling problem, e.g. the following
-
-```python
-@registry.register_problem
-class SummarizeCnnDailymailWikiLMMultiVocab64k(SummarizeCnnDailymail32k):
-  """Summarize CNN and Daily Mail articles using multi-lingual 64k vocab."""
-
-  @property
-  def vocab_filename(self):
-    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
-```
-
-**Important note:** It's easy to miss the key point that, as implemented currently, the first task in the task list must be a language modelling problem and each included task must be modified to use the resulting vocabulary.
-
-With a properly defined and registered multi-problem we can now run datagen as follows:
-
-```bash
-
-t2t-datagen --problem=languagemodel_multi_wiki_translate
-
-```
-
-This will take approximately the following amount of space (and several hours):
-
-```bash
-(t2t) username@instance-2:~$ du -sh /tmp
-99G     /tmp
-(t2t) username@instance-2:~$ du -sh /tmp/t2t_datagen
-81G     /tmp/t2t_datagen
-```
-
-### Training
-
-Next we're ready to try training a model on this MultiProblem. Note that by not specifying `--data_dir` above TFExample's were by default generated into /tmp so that's what we'll explicitly provide here.
-
-```bash
-
-t2t-trainer --problem=languagemodel_multi_wiki_translate \
-    --model=transformer \
-    --hparams_set=transformer_tall_pretrain_lm_tpu_adafactor_large \
-    --output_dir ~/t2t_train/transformer_multi_2jan19 \
-    --data_dir=/tmp \
-    --train_steps=1 \
-    --eval_steps=1
-
-```
-
-The `hparams_set` parameter we provided above was [transformer_tall_pretrain_lm_tpu_adafactor_large](https://github.com/tensorflow/tensor2tensor/blob/08e83030acf3ef13d15ad6eaefaa0a67fb20b59d/tensor2tensor/models/transformer.py#L1721), also provided below:
-
-```python
-
-@registry.register_hparams
-def transformer_tall_pretrain_lm_tpu_adafactor_large():
-  """Hparams for transformer on LM pretraining on TPU, large model."""
-  hparams = transformer_tall_pretrain_lm_tpu_adafactor()
-  hparams.hidden_size = 1024
-  hparams.num_heads = 16
-  hparams.filter_size = 32768  # max fitting in 16G memory is 49152, batch 2
-  hparams.batch_size = 4
-  hparams.multiproblem_mixing_schedule = "constant"
-  # Task order: lm/en-de/en-fr/en-ro/de-en/fr-en/ro-en/cnndm/mnli/squad.
-  hparams.multiproblem_per_task_threshold = "320,80,160,2,80,160,2,20,5,5"
-  return hparams
-
-```
-
-Here it's worth noting a couple things, one that we have specified a `multi_problem_mixing_schedule` (which is required), consumed by [MultiProblem.mix_data](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/multi_problem.py#L280). When set to "constant" the strategy for sampling examples is not a function of step and is proportional only to the per-task "thresholds" which are by default equal (sample examples from each problem with equal probability).
-
-But notice we have also specified the (non-required) `multiproblem_per_task_threshold` parameter, also consumed by mix_data, and specifically used by [sample_task](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/multi_problem.py#L340) which defines non-uniform thresholds to inform a weighted random sampling. E.g. for two problems with weights 1 and 9 the first would be sampled 1/10 of the time and the other 9/10.
-
-### Inference
-
-You can try translating from English to German using a model previously trained on `LanguagemodelMultiWikiTranslate` (the one shown above) ([gs://tensor2tensor-checkpoints/transformer_multi_2jan19/](https://console.cloud.google.com/storage/browser/tensor2tensor-checkpoints/transformer_multi_2jan19/)). Just copy the checkpoint down to a local directory such as the one given via `--output_dir` below:
-
-```bash
-
-t2t-decoder --problem=languagemodel_multi_wiki_translate \
-    --model=transformer \
-    --hparams_set=transformer_tall_pretrain_lm_tpu_adafactor_large \
-    --decode_hparams='batch_size=1,multiproblem_task_id=64510' \
-    --hparams="" \
-    --output_dir=~/t2t_train/transformer_multi_2jan19 \
-    --decode_from_file ~/newstest2014.en \
-    --data_dir=~/t2t_train/transformer_multi_2jan19
-
-```
-
-Here we'll point `--data_dir` to the checkpoint directory which includes the vocab file `vocab.languagemodel_de_en_fr_ro_wiki64k.64000.subwords`; typically data_dir would point to the directory containing your TFRecord example dataset(s).
-
-The file passed to `--decode_from_file` is simply a file with one sentence to translate on each line (in its original form, not post-vocabulary-encoded).
-
-A key requirement for multi-problem inference is that we specify the ID of the problem for which we want to perform inference. But wait, why is the task ID 64510? We can see from the code for [`MultiProblem.update_task_ids`](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/multi_problem.py#L386) that TID's have a place at the end of the vocabulary.
-
-```python
-
-class MultiProblem(problem.Problem):
-  """MultiProblem base class."""
-
-  ...
-
-  def update_task_ids(self, encoder_vocab_size):
-    """Generate task_ids for each problem.
-    These ids correspond to the index of the task in the task_list.
-    Args:
-      encoder_vocab_size: the size of the vocab which is used to compute
-        the index offset.
-    """
-    for idx, task in enumerate(self.task_list):
-      task.set_task_id(idx + encoder_vocab_size)
-      tf.logging.info("Task %d (%s) has id %d." %
-                      (idx, task.name, task.task_id))
-
-```
-
-We can look up the task_id that is assigned to each task we may want to use for inference by instantiating the MultiProblem subclass and obtaining the value, in this case via the following:
-
-```python
-
-task_index = 1 # The second task in the list is En -> De
-LanguagemodelMultiWikiTranslate().task_list[task_index].task_id
-
-```
-
-For me running the `t2t-decode` command provided above gave the following output:
-
-```bash
-...
-
-INFO:tensorflow:Running local_init_op.
-INFO:tensorflow:Done running local_init_op.
-INFO:tensorflow:Inference results INPUT: hello world was the news of the day
-INFO:tensorflow:Inference results OUTPUT: Hallo Welt war die Nachricht des Tages
-INFO:tensorflow:Elapsed Time: 37.15079
-INFO:tensorflow:Averaged Single Token Generation Time: 3.3009222 (time 36.3101439 count 11)
-
-...
-
-```
diff --git a/docs/new_model.md b/docs/new_model.md
deleted file mode 100644
index 861f83bb0..000000000
--- a/docs/new_model.md
+++ /dev/null
@@ -1,114 +0,0 @@
-# T2T: Create Your Own Model
-
-[![PyPI
-version](https://badge.fury.io/py/tensor2tensor.svg)](https://badge.fury.io/py/tensor2tensor)
-[![GitHub
-Issues](https://img.shields.io/github/issues/tensorflow/tensor2tensor.svg)](https://github.com/tensorflow/tensor2tensor/issues)
-[![Contributions
-welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](../CONTRIBUTING.md)
-[![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
-[![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
-
-Here we show how to create your own model in T2T.
-
-## The T2TModel class - abstract base class for models
-
-  `T2TModel` has three typical usages:
-
-  1. Estimator: The method `make_estimator_model_fn` builds a `model_fn` for
-     the tf.Estimator workflow of training, evaluation, and prediction.
-     It performs the method `call`, which performs the core computation,
-     followed by `estimator_spec_train`, `estimator_spec_eval`, or
-     `estimator_spec_predict` depending on the tf.Estimator mode.
-  2. Layer: The method `call` enables `T2TModel` to be used a callable by
-     itself. It calls the following methods:
-
-     * `bottom`, which transforms features according to `problem_hparams`' input
-       and target `Modality`s;
-     * `body`, which takes features and performs the core model computation to
-        return output and any auxiliary loss terms;
-     * `top`, which takes features and the body output, and transforms them
-       according to `problem_hparams`' input and target `Modality`s to return
-       the final logits;
-     * `loss`, which takes the logits, forms any missing training loss, and sums
-       all loss terms.
-  3. Inference: The method `infer` enables `T2TModel` to make sequence
-     predictions by itself.
-
-
-## Creating your own model
-
-1. Create class that extends T2TModel
-    in this example it will be a copy of existing basic fully connected network:
-
-```python
-    from tensor2tensor.utils import t2t_model
-
-    class MyFC(t2t_model.T2TModel):
-        pass
-```
-
-
-2. Implement body method:
-
-```python
-    class MyFC(t2t_model.T2TModel):
-      def body(self, features):
-        hparams = self.hparams
-        x = features["inputs"]
-        shape = common_layers.shape_list(x)
-        x = tf.reshape(x, [-1, shape[1] * shape[2] * shape[3]])  # Flatten input as in T2T they are all 4D vectors
-        for i in range(hparams.num_hidden_layers): # create layers
-          x = tf.layers.dense(x, hparams.hidden_size, name="layer_%d" % i)
-          x = tf.nn.dropout(x, keep_prob=1.0 - hparams.dropout)
-          x = tf.nn.relu(x)
-        return tf.expand_dims(tf.expand_dims(x, axis=1), axis=1)  # 4D For T2T.
-```
-
-
-Method Signature:
-
-  * Args:
-      * features: dict of str to Tensor, where each Tensor has shape [batch_size,
-     ..., hidden_size]. It typically contains keys `inputs` and `targets`.
-
-  * Returns one of:
-    * output: Tensor of pre-logit activations with shape [batch_size, ...,
-           hidden_size].
-    * losses: Either single loss as a scalar, a list, a Tensor (to be averaged),
-           or a dictionary of losses. If losses is a dictionary with the key
-           "training", losses["training"] is considered the final training
-           loss and output is considered logits; self.top and self.loss will
-           be skipped.
-
-3. Register your model
-
-```python
-    from tensor2tensor.utils import registry
-
-    @registry.register_model
-    class MyFC(t2t_model.T2TModel):
-       # ...
-```
-
-
-3. Use it with t2t tools as any other model
-
-    Have in mind that names are translated from camel case to snake_case `MyFC` -> `my_fc`
-    and that you need to point t2t to directory containing your model with `t2t_usr_dir` switch. 
-    For example if you want to train model on gcloud with 1 GPU worker on IMDB sentiment task you can run your model
-    by executing following command from your model class directory. 
-
-```bash
-    t2t-trainer \
-      --model=my_fc \
-      --t2t_usr_dir=.
-      --cloud_mlengine --worker_gpu=1 \
-      --generate_data \
-      --data_dir='gs://data' \
-      --output_dir='gs://out' \
-      --problem=sentiment_imdb \
-      --hparams_set=basic_fc_small \
-      --train_steps=10000 \
-      --eval_steps=10 \
-```
diff --git a/docs/new_problem.md b/docs/new_problem.md
deleted file mode 100644
index 13f012b79..000000000
--- a/docs/new_problem.md
+++ /dev/null
@@ -1,243 +0,0 @@
-# T2T: Train on Your Own Data
-
-[![PyPI
-version](https://badge.fury.io/py/tensor2tensor.svg)](https://badge.fury.io/py/tensor2tensor)
-[![GitHub
-Issues](https://img.shields.io/github/issues/tensorflow/tensor2tensor.svg)](https://github.com/tensorflow/tensor2tensor/issues)
-[![Contributions
-welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CONTRIBUTING.md)
-[![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
-[![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
-
-Another good overview of this part together with training is given in
-[The Cloud ML Poetry Blog
-Post](https://cloud.google.com/blog/big-data/2018/02/cloud-poetry-training-and-hyperparameter-tuning-custom-text-models-on-cloud-ml-engine)
-
-Let's add a new dataset together and train the
-[Transformer](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/models/transformer.py)
-model on it. We'll give the model a line of poetry, and it will learn to
-generate the next line.
-
-# Defining the `Problem`
-
-For each problem we want to tackle we create a new subclass of `Problem` and
-register it. Let's call our problem `PoetryLines`.
-
-Since many text-to-text problems share similar methods, there's already a class
-called
-[`Text2TextProblem`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/text_problems.py)
-that extends the base problem class
-[`Problem`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/problem.py)
-and makes it easy to add text-to-text problems.
-
-In that same file, there are other base classes that make it easy to add text
-classification tasks (`Text2ClassProblem`) and language modeling tasks
-(`Text2SelfProblem`).
-
-For our problem, let's create the file `poetry_lines.py` and add our new
-problem, `PoetryLines`, which extends `Text2TextProblem` and register it so that
-it is accessible by command-line flag.
-
-Here's the Problem in full. We'll go step by step through it.
-
-```python
-import re
-
-from gutenberg import acquire
-from gutenberg import cleanup
-
-from tensor2tensor.data_generators import problem
-from tensor2tensor.data_generators import text_problems
-from tensor2tensor.utils import registry
-
-@registry.register_problem
-class PoetryLines(text_problems.Text2TextProblem):
-  """Predict next line of poetry from the last line. From Gutenberg texts."""
-
-  @property
-  def approx_vocab_size(self):
-    return 2**13  # ~8k
-
-  @property
-  def is_generate_per_split(self):
-    # generate_data will shard the data into TRAIN and EVAL for us.
-    return False
-
-  @property
-  def dataset_splits(self):
-    """Splits of data to produce and number of output shards for each."""
-    # 10% evaluation data
-    return [{
-        "split": problem.DatasetSplit.TRAIN,
-        "shards": 9,
-    }, {
-        "split": problem.DatasetSplit.EVAL,
-        "shards": 1,
-    }]
-
-  def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    del data_dir
-    del tmp_dir
-    del dataset_split
-
-
-    books = [
-        # bookid, skip N lines
-        (19221, 223),
-        (15553, 522),
-    ]
-
-    for (book_id, toskip) in books:
-      text = cleanup.strip_headers(acquire.load_etext(book_id)).strip()
-      lines = text.split("\n")[toskip:]
-      prev_line = None
-      ex_count = 0
-      for line in lines:
-        # Any line that is all upper case is a title or author name
-        if not line or line.upper() == line:
-          prev_line = None
-          continue
-
-        line = re.sub("[^a-z]+", " ", line.strip().lower())
-        if prev_line and line:
-          yield {
-              "inputs": prev_line,
-              "targets": line,
-          }
-          ex_count += 1
-        prev_line = line
-```
-
-## Vocabulary specification
-
-The text generated is encoded with a vocabulary for training. By default, it is
-a `SubwordTextEncoder` that is built with an approximate vocab size specified by
-the user. It's fully invertible (no out-of-vocab tokens) with a fixed-size vocab
-which makes it ideal for text problems.
-
-You can also choose to use a character-level encoder or a token encoder where
-you provide the vocab file yourself. See `Text2TextProblem.vocab_type`.
-
-Here we specify that we're going to have a vocabulary with approximately 8,000
-subwords.
-
-```python
-  @property
-  def approx_vocab_size(self):
-    return 2**13  # ~8k
-```
-
-## Splitting data between Train and Eval
-
-By setting `is_generate_per_split=False`, the `generate_samples` method will
-only be called once and the data will automatically be split across training and
-evaluation data for us. This is useful because for our dataset we don't have
-pre-existing "training" and "evaluation" sets. If we did, we'd set
-`is_generate_per_split=True` so that `generate_samples` was called once per data
-split.
-
-The `dataset_splits` method determines the fraction that goes to each split. The
-training data will be generated into 9 files and the evaluation data into 1.
-90% of the data will be for training. 10% of the data will be for evaluation.
-
-```python
-  @property
-  def is_generate_per_split(self):
-    # generate_data will shard the data into TRAIN and EVAL for us.
-    return False
-
-  @property
-  def dataset_splits(self):
-    """Splits of data to produce and number of output shards for each."""
-    # 10% evaluation data
-    return [{
-        "split": problem.DatasetSplit.TRAIN,
-        "shards": 9,
-    }, {
-        "split": problem.DatasetSplit.EVAL,
-        "shards": 1,
-    }]
-```
-
-## Generating samples
-
-`generate_samples` is the bulk of the code where we actually produce
-dictionaries of poetry line pairs ("inputs" and "targets").
-
-Some problems might require downloading, which can be done into `tmp_dir`. Some
-problems may use their own token vocabulary file, in which case it can be copied
-into `data_dir` before yielding samples.
-
-Here we iterate through the lines of a couple books of poetry and produce pairs
-of lines for the model to train against.
-
-```python
-  def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    del data_dir
-    del tmp_dir
-    del dataset_split
-
-    books = [
-        # bookid, skip N lines
-        (19221, 223),
-        (15553, 522),
-    ]
-
-    for (book_id, toskip) in books:
-      text = cleanup.strip_headers(acquire.load_etext(book_id)).strip()
-      lines = text.split("\n")[toskip:]
-      prev_line = None
-      ex_count = 0
-      for line in lines:
-        # Any line that is all upper case is a title or author name
-        if not line or line.upper() == line:
-          prev_line = None
-          continue
-
-        line = re.sub("[^a-z]+", " ", line.strip().lower())
-        if prev_line and line:
-          yield {
-              "inputs": prev_line,
-              "targets": line,
-          }
-          ex_count += 1
-        prev_line = line
-```
-
-That's all for the problem specification! We're ready to generate the data.
-
-# Run data generation
-
-You can generate data for your problem with `t2t-datagen` and the
-`--t2t_usr_dir` flag, which points to the directory containing an `__init__.py`
-file that imports the `poetry_lines` file we just wrote. See setup below.
-
-```bash
-USR_DIR=...
-PROBLEM=poetry_lines
-DATA_DIR=$HOME/t2t_data
-TMP_DIR=/tmp/t2t_datagen
-mkdir -p $DATA_DIR $TMP_DIR
-
-t2t-datagen \
-  --t2t_usr_dir=$USR_DIR \
-  --data_dir=$DATA_DIR \
-  --tmp_dir=$TMP_DIR \
-  --problem=$PROBLEM
-```
-
-`PROBLEM` is the name of the class that was registered with
-`@registry.register_problem`, but converted from `CamelCase` to `snake_case`.
-
-`USR_DIR` is a directory with the `poetry_lines.py` file and an
-`__init__.py` file that imports it (`from . import poetry_lines`).
-
-If you plan to contribute problems to the tensor2tensor repository, you can
-clone the repository and install it in developer mode with `pip install -e .`.
-
-# Train!
-
-You can train exactly as you do in the [walkthrough](walkthrough.md) with flags
-`--problem=poetry_lines` and `--t2t_usr_dir=$USR_DIR`.
-
-All done. Let us know what amazing poetry your model writes!
diff --git a/docs/overview.md b/docs/overview.md
deleted file mode 100644
index 9ea87bc50..000000000
--- a/docs/overview.md
+++ /dev/null
@@ -1,175 +0,0 @@
-# T2T: Life of an Example
-
-[![PyPI
-version](https://badge.fury.io/py/tensor2tensor.svg)](https://badge.fury.io/py/tensor2tensor)
-[![GitHub
-Issues](https://img.shields.io/github/issues/tensorflow/tensor2tensor.svg)](https://github.com/tensorflow/tensor2tensor/issues)
-[![Contributions
-welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CONTRIBUTING.md)
-[![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
-[![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
-
-This doc explains how a training example flows through T2T, from data generation
-to training, evaluation, and decoding.
-
-Some key files and their functions:
-
-*   [`t2t_trainer.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/bin/t2t_trainer.py) and [`trainer_lib.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/trainer_lib.py):
-    Main entrypoint for training and evaluation.  Constructs and runs all the
-    main components of the system (the `Problem`, the `HParams`, the
-    `Estimator`, the `Experiment`, the `input_fn`s and `model_fn`).
-*   [`common_hparams.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/layers/common_hparams.py):
-    `basic_params1` serves as the base for all model hyperparameters. Registered
-    model hparams functions always start with this default set of
-    hyperparameters.
-*   [`problem.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/problem.py):
-    Every dataset in T2T subclasses `Problem`. `Problem.input_fn` is the
-    Estimator input function.
-*   [`t2t_model.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/t2t_model.py):
-    Every model in T2T subclasses `T2TModel`. `T2TModel.estimator_model_fn` is
-    the Estimator model function.
-
-## Data Generation
-
-The `t2t-datagen` binary is the entrypoint for data generation. It simply looks
-up the `Problem` specified by `--problem` and calls
-`Problem.generate_data(data_dir, tmp_dir)`.
-
-All `Problem`s are expected to generate 2 sharded `TFRecords` files - 1 for
-training and 1 for evaluation - with `tensorflow.Example` protocol buffers. The
-expected names of the files are given by `Problem.{training, dev}_filepaths`.
-Typically, the features in the `Example` will be `"inputs"` and `"targets"`;
-however, some tasks have a different on-disk representation that is converted to
-`"inputs"` and `"targets"` online in the input pipeline (e.g. image features are
-typically stored with features `"image/encoded"` and `"image/format"` and the
-decoding happens in the input pipeline).
-
-For tasks that require a vocabulary, this is also the point at which the
-vocabulary is generated and all examples are encoded.
-
-There are several utility functions in
-[`generator_utils`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/generator_utils.py)
-that are commonly used by `Problem`s to generate data. Several are highlighted
-below:
-
-*   `generate_dataset_and_shuffle`: given 2 generators, 1 for training and 1 for
-    eval, yielding dictionaries of `<feature name, list< int or float or
-    string >>`, will produce sharded and shuffled `TFRecords` files with
-    `tensorflow.Example` protos.
-*   `maybe_download`: downloads a file at a URL to the given directory and
-    filename (see `maybe_download_from_drive` if the URL points to Google
-    Drive).
-*   `get_or_generate_vocab_inner`: given a target vocabulary size and a
-    generator that yields lines or tokens from the dataset, will build a
-    `SubwordTextEncoder` along with a backing vocabulary file that can be used
-    to map input strings to lists of ids.
-    [`SubwordTextEncoder`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/text_encoder.py)
-    uses word pieces and its encoding is fully invertible.
-
-## Data Input Pipeline
-
-Once the data is produced on disk, training, evaluation, and inference (if
-decoding from the dataset) consume it by way of the T2T input pipeline, defined
-by `Problem.input_fn`.
-
-The entire input pipeline is implemented with the new `tf.data.Dataset` API.
-
-The input function has 2 main parts: first, reading and processing individual
-examples, which is done is `Problem.dataset`, and second, batching, which is
-done in `Problem.input_fn` after the call to `Problem.dataset`.
-
-`Problem` subclasses may override the entire `input_fn` or portions of it (e.g.
-`example_reading_spec` to indicate the names, types, and shapes of features on
-disk). Typically they only override portions.
-
-### Batching
-
-Problems that have fixed size features (e.g. image problems) can use
-`hp.batch_size` to set the batch size.
-
-Variable length Problems are bucketed by sequence length and then batched out of
-those buckets.  This significantly improves performance over a naive batching
-scheme for variable length sequences because each example in a batch must be
-padded to match the example with the maximum length in the batch.
-
-Controlling hparams:
-
-* `hp.batch_size`: the approximate total number of tokens in
-  the batch (i.e. long sequences will have smaller actual batch size and short
-  sequences will have a larger actual batch size in order to generally have an
-  equal number of tokens in the batch).
-* `hp.max_length`: For variable length features, sequences with length longer
-  than this will be dropped during training (and also during eval if
-  `hp.eval_drop_long_sequences` is `True`). If not set, the maximum length of
-  examples is set to `hp.batch_size`.
-* `hp.batch_size_multiplier`: multiplier for the maximum length
-* `hp.min_length_bucket`: example length for the smallest bucket (i.e. the
-  smallest bucket will bucket examples up to this length).
-* `hp.length_bucket_step`: controls how spaced out the length buckets are.
-
-## Building the Model
-
-At this point, the input features typically have `"inputs"` and `"targets"`,
-each of which is a batched 4-D Tensor (e.g. of shape `[batch_size,
-sequence_length, 1, 1]` for text input or `[batch_size, height, width, 3]` for
-image input).
-
-The Estimator model function is created by `T2TModel.estimator_model_fn`, which
-may be overridden in its entirety by subclasses if desired. Typically,
-subclasses only override `T2TModel.body`.
-
-The model function constructs a `T2TModel`, calls it, and then calls
-`T2TModel.{estimator_spec_train, estimator_spec_eval, estimator_spec_predict}`
-depending on the mode.
-
-A call of a `T2TModel` internally calls `bottom`, `body`, `top`, and `loss`, all
-of which can be overridden by subclasses (typically only `body` is).
-
-The default implementations of `bottom`, `top`, and `loss` depend on the
-`Modality` specified for the input and target features (e.g.
-`SymbolModality.bottom` embeds integer tokens and `SymbolModality.loss` is
-`softmax_cross_entropy`).
-
-## `Estimator` and `Experiment`
-
-The actual training loop and related services (checkpointing, summaries,
-continuous evaluation, etc.) are all handled by `Estimator` and `Experiment`
-objects. `t2t_trainer.py` is the main entrypoint and uses `trainer_lib.py`
-to construct the various components.
-
-## Decoding
-
-* [`t2t_decoder.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/bin/t2t-decoder)
-* [`decoding.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/decoding.py)
-
-## System Overview for Train/Eval
-
-See `t2t_trainer.py` and `trainer_lib.py`.
-
-* Create HParams
-* Create `RunConfig`, including `Parallelism` object (i.e. `data_parallelism`)
-* Create `Experiment`, including hooks
-* Create `Estimator`
-  * `T2TModel.estimator_model_fn`
-    * `model(features)`
-      * `model.model_fn`
-        * `model.bottom`
-        * `model.body`
-        * `model.top`
-        * `model.loss`
-    * [TRAIN] `model.estimator_spec_train`
-      * `train_op = model.optimize`
-    * [EVAL] `model.estimator_spec_eval`
-      * Create metrics
-* Create input functions
-  * `Problem.input_fn`
-    * `Problem.dataset`
-    * Batching
-* Create hooks
-* Run Experiment --schedule (e.g. `exp.continuous_train_and_eval()`)
-  * `estimator.train`
-    * `train_op = model_fn(input_fn(mode=TRAIN))`
-    * Run train op
-  * `estimator.evaluate`
-    * `metrics = model_fn(input_fn(mode=EVAL))`
-    * Accumulate metrics
diff --git a/docs/tutorials/asr_with_transformer.md b/docs/tutorials/asr_with_transformer.md
deleted file mode 100644
index ee8aa0f53..000000000
--- a/docs/tutorials/asr_with_transformer.md
+++ /dev/null
@@ -1,4 +0,0 @@
-# Automated Speech Recognition with the Transformer model
-
-See the
-[official tutorial](https://cloud.google.com/tpu/docs/tutorials/automated-speech-recognition).
diff --git a/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index bfda04240..6e44e50e3 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -32,8 +32,9 @@
 from tensor2tensor.utils import update_ops_hook
 
 import tensorflow as tf
+import tensorflow_gan as tfgan
 
-gan_losses = tf.contrib.gan.losses.wargs
+gan_losses = tfgan.losses.wargs
 
 
 class NextFrameSavpBase(object):

From 1790cccb40d124e9726ee2d2c457b4cedac17d8e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 13 Aug 2019 11:34:12 -0700
Subject: [PATCH 2281/2720] Adjust Reformer output layers

This provides a clearer split between the portion of the network that is mathematically reversible, and the portion that is not.
Also adds dropout before the output projection and softmax.

PiperOrigin-RevId: 263179765
---
 .../models/research/transformer_revnet.py     | 70 ++++++++++++++++++-
 1 file changed, 67 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 2cfec5fea..7ff50fc18 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -138,6 +138,67 @@ def n_outputs(self):
     return self._n_sections
 
 
+class SplitForOutput(tl.ReversibleLayer):
+  """Splits activations into sections (for use right before the output layer).
+
+  After the reversible portion of the network, there is a final output portion
+  that's non-reversible (which at minimum includes normalization, output
+  projection, and log-softmax). The output portion needs to operate on chunks
+  of the sequence to avoid running out of memory for large vocabulary sizes.
+
+  This layer concatenates the two subparts of the activations along the feature
+  dimension, and then splits into chunks along the time dimension. We implement
+  it is a subclass of tl.ReversibleLayer because we want to ensure that multiple
+  copies of the activations don't exist simultaneously except in the middle of a
+  memory copy operation.
+  """
+
+  def __init__(self, n_sections=2, axis=-2):
+    super(SplitForOutput, self).__init__()
+    self._n_sections = n_sections
+    self._axis = axis
+
+  def n_inputs(self):
+    """Specifies how many data tensors this layer expects as input."""
+    return 2
+
+  def n_outputs(self):
+    """Specifies how many data tensors this layer promises as output."""
+    return self._n_sections
+
+  def new_parameters(self, input_shape, input_dtype, rng):
+    return ()
+
+  def call(self, inputs, params=(), **kwargs):
+    del params, kwargs
+    x1, x2 = inputs
+
+    x1_split = backend.numpy.split(x1, self._n_sections, self._axis)
+    x2_split = backend.numpy.split(x2, self._n_sections, self._axis)
+
+    res = [backend.numpy.concatenate(ys, -1) for ys in zip(x1_split, x2_split)]
+    return tuple(res)
+
+  def reverse(self, output, params=(), **kwargs):
+    del params, kwargs
+
+    x1_split = []
+    x2_split = []
+    for y in output:
+      y1, y2 = backend.numpy.split(y, 2, -1)
+      x1_split.append(y1)
+      x2_split.append(y2)
+
+    x1 = backend.numpy.concatenate(x1_split, self._axis)
+    x2 = backend.numpy.concatenate(x2_split, self._axis)
+
+    return (x1, x2)
+
+  def reverse_and_grad(self, output, ct, params=(), **kwargs):
+    del params, kwargs
+    return self.reverse(output), (self.reverse(ct), (), ())
+
+
 @tl.layer()
 def Chunk(x, params, n_sections=2, **kwargs):
   del params, kwargs
@@ -775,11 +836,14 @@ def TransformerRevnetLM(vocab_size,
                        n_attention_chunks, attention_type,
                        dropout, mode)
           for _ in range(n_layers)
+      ] + [
+          SplitForOutput(n_sections=n_chunks, axis=-2),  # pylint: disable=no-value-for-parameter
       ]),
-      tl.Parallel(tl.LayerNorm(), tl.LayerNorm()),
-      tl.Concatenate(),
-      Split(n_sections=n_chunks, axis=-2),  # pylint: disable=no-value-for-parameter
       Map([
+          # TODO(kitaev): Test whether dropout should go before or after the
+          # LayerNorm, and whether dropout broadcasting is needed here.
+          tl.LayerNorm(),
+          BroadcastedDropout(rate=dropout, mode=mode),  # pylint: disable=no-value-for-parameter
           tl.Dense(vocab_size),
           tl.LogSoftmax(),
       ], n_sections=n_chunks),

From b4e5758c69adccea3d393cf419e166369ea8c608 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 13 Aug 2019 11:45:26 -0700
Subject: [PATCH 2282/2720] Fuse batch and head dimensions in attention

PiperOrigin-RevId: 263182145
---
 .../models/research/transformer_revnet.py     | 175 +++++++++++-------
 1 file changed, 110 insertions(+), 65 deletions(-)

diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 7ff50fc18..c4be23e6a 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -275,25 +275,71 @@ def call_compute_residual(x, params, kwargs):
     return reconstructed_x, (x_ct, (residual_params_ct, ()), kwargs_ct)
 
 
-@tl.layer(n_inputs=1, n_outputs=1)
-def SplitHeads(x, params, n_heads=1, **kwargs):
-  del params, kwargs
-  d_model = x.shape[-1]
-  assert d_model % n_heads == 0
-  d_head = d_model // n_heads
-  n_batch = np.shape(x)[0]
-  # n_batch, seqlen, d_model --> n_batch, n_heads, seqlen, d_head
-  return np.transpose(
-      np.reshape(x, (n_batch, -1, n_heads, d_head)), (0, 2, 1, 3))
+class ComputeAttentionHeads(tl.Layer):
+  """Computes queries/keys/values via linear projection.
+
+  The output shape is (n_batch * n_heads, seqlen, d_head); the batch and head
+  dimensions are fused to allow for more efficient memory layouts.
+  """
 
+  def __init__(self, n_heads=1, d_head=64,
+               kernel_initializer=tl.initializers.GlorotUniformInitializer()):
+    super(ComputeAttentionHeads, self).__init__()
+    self._n_heads = n_heads
+    self._d_head = d_head
+    self._kernel_initializer = kernel_initializer
+    # The lack of a bias term here is consistent with the tensor2tensor
+    # implementation, and shouldn't have an effect on modeling quality.
+
+  def call(self, x, params, **kwargs):
+    del kwargs
+    seqlen = x.shape[1]
+    res = np.dot(x, params)
+
+    # n_batch, seqlen, n_heads*d_head -> n_batch, seqlen, n_heads, d_head
+    res = np.reshape(res, (x.shape[0], seqlen, self._n_heads, self._d_head))
+    # n_batch, seqlen, n_heads, d_head -> n_batch, n_heads, seqlen, d_head
+    res = np.transpose(res, (0, 2, 1, 3))
+    # n_batch, n_heads, seqlen, d_head -> n_batch*n_heads, seqlen, d_head
+    res = np.reshape(res, (-1, seqlen, self._d_head))
+
+    return res
+
+  def new_parameters(self, input_shape, input_dtype, rng):
+    del input_dtype
+    w = self._kernel_initializer(
+        (input_shape[-1], self._n_heads * self._d_head), rng)
+    return w
 
-@tl.layer(n_inputs=1, n_outputs=1)
-def JoinHeads(x, params, **kwargs):
-  del params, kwargs
-  n_batch = np.shape(x)[0]
-  seqlen = np.shape(x)[2]
-  # n_batch, n_heads, seqlen, d_head --> n_batch, seqlen, d_model
-  return np.reshape(np.transpose(x, (0, 2, 1, 3)), (n_batch, seqlen, -1))
+
+class ComputeAttentionOutput(tl.Layer):
+  """Joins outputs from different heads via linear projection."""
+
+  def __init__(self, n_heads=1, d_model=1024,
+               kernel_initializer=tl.initializers.GlorotUniformInitializer()):
+    super(ComputeAttentionOutput, self).__init__()
+    self._n_heads = n_heads
+    self._d_model = d_model
+    self._kernel_initializer = kernel_initializer
+    # The lack of a bias term here is consistent with the tensor2tensor
+    # implementation, and shouldn't have an effect on modeling quality.
+
+  def call(self, x, params, **kwargs):
+    del kwargs
+    seqlen = x.shape[1]
+    d_head = x.shape[2]
+
+    x = np.reshape(x, (-1, self._n_heads, seqlen, d_head))
+    x = np.transpose(x, (0, 2, 1, 3))  # -> n_batch, seqlen, n_heads, d_head
+    x = np.reshape(x, (-1, seqlen, self._n_heads * d_head))
+
+    return np.dot(x, params)
+
+  def new_parameters(self, input_shape, input_dtype, rng):
+    del input_dtype
+    w = self._kernel_initializer(
+        (input_shape[-1] * self._n_heads, self._d_model), rng)
+    return w
 
 
 class ApplyAttentionWrapper(tl.Parallel):
@@ -415,8 +461,8 @@ def forward_slice(query_slice, q_loop_idx, key, value):
       dots = dots / dots.sum(axis=-1, keepdims=True)
 
       if self.dropout is not None and self.dropout > 0.0:
-        # Dropout is broadcast across the batch and head dimensions
-        dropout_shape = (1, 1, dots.shape[-2], dots.shape[-1])
+        # Dropout is broadcast across the batch+head dimension
+        dropout_shape = (1, dots.shape[-2], dots.shape[-1])
         slice_rng = jax.random.fold_in(rng, q_loop_idx)
         keep_prob = jax.lax.tie_in(dots, 1.0 - self.dropout)
         keep = backend.random.bernoulli(slice_rng, keep_prob, dropout_shape)
@@ -432,7 +478,7 @@ def forward_and_vjp_slice(query_slice, q_loop_idx, key, value, ct_slice):
       return output_slice, vjpfun(ct_slice)
 
     q_loop_idx = np.zeros((), dtype=np.int32)
-    q_loop_max = query.shape[2]
+    q_loop_max = query.shape[-2]
     q_loop_stride = self._loop_stride
     assert q_loop_max % q_loop_stride == 0, (
         'Stride must evenly divide the number of query elements.')
@@ -461,15 +507,15 @@ def body_fun(vals):
         q_loop_idx, out_accum = vals
 
       query_slice = jax.lax.dynamic_slice_in_dim(
-          query, q_loop_idx, q_loop_stride, axis=2)
+          query, q_loop_idx, q_loop_stride, axis=-2)
 
       if do_backprop:
         ct_slice = jax.lax.dynamic_slice_in_dim(
-            ct, q_loop_idx, q_loop_stride, axis=2)
+            ct, q_loop_idx, q_loop_stride, axis=-2)
         out_slice, partial_ct = forward_and_vjp_slice(
             query_slice, q_loop_idx, key, value, ct_slice)
         query_ct_accum = jax.lax.dynamic_update_slice_in_dim(
-            query_ct_accum, partial_ct[0], q_loop_idx, axis=2)
+            query_ct_accum, partial_ct[0], q_loop_idx, axis=-2)
         # ignore partial_ct[1], which is wrt the loop idx
         key_ct_accum = key_ct_accum + partial_ct[2]
         value_ct_accum = value_ct_accum + partial_ct[3]
@@ -477,7 +523,7 @@ def body_fun(vals):
         out_slice = forward_slice(query_slice, q_loop_idx, key, value)
 
       out_accum = jax.lax.dynamic_update_slice_in_dim(
-          out_accum, out_slice, q_loop_idx, axis=2)
+          out_accum, out_slice, q_loop_idx, axis=-2)
       q_loop_idx = q_loop_idx + q_loop_stride
 
       if do_backprop:
@@ -508,63 +554,63 @@ def call(self, inputs, params=(), **kwargs):
   def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
     del params, kwargs
     q, k, v = inputs
-    # q/k/v are n_batch, n_heads, seqlen, d_head
+    # q/k/v are n_batch*n_heads, seqlen, d_head
 
-    assert k.shape[2] % self.n_bins == 0
-    bin_size = int(k.shape[2] // self.n_bins)
+    assert k.shape[-2] % self.n_bins == 0
+    bin_size = int(k.shape[-2] // self.n_bins)
 
-    # q_bins/kv_bins are n_batch, n_heads, seqlen
+    # q_bins/kv_bins are n_batch*n_heads, seqlen
     # They specify which hash bucket the query/key/value vectors fall in. For
     # now, instead of hashing we just put consecutive items in the same bucket.
-    q_bins = np.arange(q.shape[2], dtype=np.int32) // bin_size
+    q_bins = np.arange(q.shape[-2], dtype=np.int32) // bin_size
     q_bins = jax.lax.tie_in(q, q_bins)
-    q_bins = q_bins[None, None, :]
+    q_bins = q_bins[None, :]
     q_bins = np.broadcast_to(q_bins, q.shape[:-1])
     q_bins = -q_bins
     kv_bins = q_bins * 2
 
-    # q_t/kv_t are n_batch, n_heads, seqlen
-    q_t = jax.lax.tie_in(q, np.arange(q.shape[2]))
-    q_t = np.reshape(q_t, (1, 1, q_t.shape[0]))
+    # q_t/kv_t are n_batch*n_heads, seqlen
+    q_t = jax.lax.tie_in(q, np.arange(q.shape[-2]))
+    q_t = np.reshape(q_t, (1, q_t.shape[0]))
     q_t = np.broadcast_to(q_t, q.shape[:-1])
     kv_t = q_t
 
-    def chunk_rank3(x):
-      return np.reshape(x, (x.shape[0], x.shape[1], self.n_bins, -1))
+    def chunk_scalars(x):
+      return np.reshape(x, (x.shape[0], self.n_bins, -1))
 
-    def chunk_rank4(x):
+    def chunk_vectors(x):
       return np.reshape(
-          x, (x.shape[0], x.shape[1], self.n_bins, -1, x.shape[-1]))
+          x, (x.shape[0], self.n_bins, -1, x.shape[-1]))
 
-    def unchunk_rank4(x):
-      return np.reshape(x, (x.shape[0], x.shape[1], -1, x.shape[-1]))
+    def unchunk_vectors(x):
+      return np.reshape(x, (x.shape[0], -1, x.shape[-1]))
 
    # Sort everything by bin number (variables starting with "s" are sorted)
-    _, sq_t = jax.lax.sort_key_val(q_bins, q_t, dimension=2)
+    _, sq_t = jax.lax.sort_key_val(q_bins, q_t, dimension=-1)
 
-    sq = np.take_along_axis(q, sq_t[:, :, :, None], axis=2)
+    sq = np.take_along_axis(q, sq_t[:, :, None], axis=-2)
     if ct is not None:
-      so_ct = np.take_along_axis(ct, sq_t[:, :, :, None], axis=2)
+      so_ct = np.take_along_axis(ct, sq_t[:, :, None], axis=-2)
 
-    _, skv_t = jax.lax.sort_key_val(kv_bins, kv_t, dimension=2)
-    sk = np.take_along_axis(k, skv_t[:, :, :, None], axis=2)
-    sv = np.take_along_axis(v, skv_t[:, :, :, None], axis=2)
+    _, skv_t = jax.lax.sort_key_val(kv_bins, kv_t, dimension=-1)
+    sk = np.take_along_axis(k, skv_t[:, :, None], axis=-2)
+    sv = np.take_along_axis(v, skv_t[:, :, None], axis=-2)
 
     @jax.jit
     def binned_attn(sq, sk, sv):
       """Performs attention on sorted queries/keys/values."""
       # Split off a "bin" axis so that attention only occurs whithin chunks.
-      bq_t = chunk_rank3(sq_t)
-      bkv_t = chunk_rank3(skv_t)
-      bq = chunk_rank4(sq)
-      bk = chunk_rank4(sk)
-      bv = chunk_rank4(sv)
+      bq_t = chunk_scalars(sq_t)
+      bkv_t = chunk_scalars(skv_t)
+      bq = chunk_vectors(sq)
+      bk = chunk_vectors(sk)
+      bv = chunk_vectors(sv)
 
       dots = np.matmul(bq, np.swapaxes(bk, -1, -2)) / np.sqrt(bq.shape[-1])
 
       # Causal masking
       mask = jax.lax.convert_element_type(
-          jax.lax.lt(bq_t[:, :, :, :, None], bkv_t[:, :, :, None, :]),
+          jax.lax.lt(bq_t[:, :, :, None], bkv_t[:, :, None, :]),
           np.float32)
       dots = dots - 1e9 * mask
 
@@ -573,7 +619,7 @@ def binned_attn(sq, sk, sv):
       dots = dots / dots.sum(axis=-1, keepdims=True)
       bo = np.matmul(dots, bv)
 
-      so = unchunk_rank4(bo)
+      so = unchunk_vectors(bo)
       return so
 
     @jax.jit
@@ -584,8 +630,8 @@ def binned_attn_vjp(sq, sk, sv, so_ct):
 
     if ct is None:
       so = binned_attn(sq, sk, sv)
-      _, undo_q_sort = jax.lax.sort_key_val(sq_t, q_t, dimension=2)
-      out = np.take_along_axis(so, undo_q_sort[:, :, :, None], axis=2)
+      _, undo_q_sort = jax.lax.sort_key_val(sq_t, q_t, dimension=-1)
+      out = np.take_along_axis(so, undo_q_sort[:, :, None], axis=-2)
       return out, None
     else:
       # Jax can construct a backward pass automatically, but it's about 2x
@@ -594,13 +640,13 @@ def binned_attn_vjp(sq, sk, sv, so_ct):
       # with permutations so we use gather for the backward pass too.
       so, (sq_ct, sk_ct, sv_ct) = binned_attn_vjp(sq, sk, sv, so_ct)
 
-      _, undo_q_sort = jax.lax.sort_key_val(sq_t, q_t, dimension=2)
-      out = np.take_along_axis(so, undo_q_sort[:, :, :, None], axis=2)
-      q_ct = np.take_along_axis(sq_ct, undo_q_sort[:, :, :, None], axis=2)
+      _, undo_q_sort = jax.lax.sort_key_val(sq_t, q_t, dimension=-1)
+      out = np.take_along_axis(so, undo_q_sort[:, :, None], axis=-2)
+      q_ct = np.take_along_axis(sq_ct, undo_q_sort[:, :, None], axis=-2)
 
-      _, undo_kv_sort = jax.lax.sort_key_val(skv_t, kv_t, dimension=2)
-      k_ct = np.take_along_axis(sk_ct, undo_kv_sort[:, :, :, None], axis=2)
-      v_ct = np.take_along_axis(sv_ct, undo_kv_sort[:, :, :, None], axis=2)
+      _, undo_kv_sort = jax.lax.sort_key_val(skv_t, kv_t, dimension=-1)
+      k_ct = np.take_along_axis(sk_ct, undo_kv_sort[:, :, None], axis=-2)
+      v_ct = np.take_along_axis(sv_ct, undo_kv_sort[:, :, None], axis=-2)
 
       return out, (q_ct, k_ct, v_ct)
 
@@ -758,9 +804,9 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
       tl.LayerNorm(),
       tl.Dup(), tl.Dup(),
       tl.Parallel(
-          [tl.Dense(d_attention_key * n_heads), SplitHeads(n_heads=n_heads)],  # pylint: disable=no-value-for-parameter
-          [tl.Dense(d_attention_key * n_heads), SplitHeads(n_heads=n_heads)],  # pylint: disable=no-value-for-parameter
-          [tl.Dense(d_attention_value * n_heads), SplitHeads(n_heads=n_heads)],  # pylint: disable=no-value-for-parameter
+          [ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key)],
+          [ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key)],
+          [ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value)],
       ),
   ]
 
@@ -769,8 +815,7 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
   # ReversibleAttentionHalfResidual requires that post_attention be linear in
   # its input (so the backward pass can be computed without knowing the input)
   post_attention = [
-      JoinHeads(),  # pylint: disable=no-value-for-parameter
-      tl.Dense(d_model),
+      ComputeAttentionOutput(n_heads=n_heads, d_model=d_model),
       Unchunk(n_sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
       BroadcastedDropout(rate=dropout, mode=mode),  # pylint: disable=no-value-for-parameter
   ]

From 71966b90ca4d40ae80368df529c1895bb4f1d4fd Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 13 Aug 2019 22:16:43 -0700
Subject: [PATCH 2283/2720] Move PPO code to PPOTrainer

PiperOrigin-RevId: 263282309
---
 tensor2tensor/envs/env_problem_utils.py       |  29 +-
 tensor2tensor/trax/rlax/configs/atari.gin     |   2 -
 ...params_online_tune_wide_resnet_cifar10.gin |   2 -
 .../online_tune_wide_resnet_cifar10.gin       |   2 -
 tensor2tensor/trax/rlax/ppo.py                | 534 ++++++++++--------
 tensor2tensor/trax/rlax/ppo_main.py           |  23 +-
 .../trax/rlax/ppo_training_loop_test.py       |  11 +-
 7 files changed, 317 insertions(+), 286 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 65d074c56..23d2d5f9c 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -27,9 +27,6 @@
 from tensor2tensor.envs import rendered_env_problem
 from tensor2tensor.rl import gym_utils
 
-EPSILON_GREEDY = "epsilon-greedy"
-GUMBEL_SAMPLING = "gumbel"
-
 
 def done_indices(dones):
   """Calculates the indices where dones has True."""
@@ -61,9 +58,7 @@ def play_env_problem_with_policy(env,
                                  max_timestep=None,
                                  reset=True,
                                  rng=None,
-                                 policy_sampling=GUMBEL_SAMPLING,
                                  temperature=1.0,
-                                 eps=0.1,
                                  len_history_for_policy=32,
                                  num_to_keep=1):
   """Plays the given env with the policy function to collect trajectories.
@@ -79,10 +74,7 @@ def play_env_problem_with_policy(env,
     reset: bool, true if we want to reset the envs. The envs are also reset if
       max_max_timestep is None or < 0
     rng: jax rng, splittable.
-    policy_sampling: string, how to select an action given a policy, one of:
-      EPSILON_GREEDY, GUMBEL_SAMPLING
     temperature: float, temperature used in Gumbel sampling.
-    eps: float, epsilon to use in epsilon greedy.
     len_history_for_policy: int, the maximum history to keep for applying the
       policy on. We also bucket observations on this number.
     num_to_keep: int, while truncating trajectory how many time-steps to keep.
@@ -98,19 +90,6 @@ def gumbel_sample(log_probs):
     g = -np.log(-np.log(u))
     return np.argmax((log_probs / temperature) + g, axis=1)
 
-  def epsilon_greedy(log_probs):
-    """Epsilon greedy sampling."""
-    _, A = log_probs.shape  # pylint: disable=invalid-name
-    actions = []
-    for log_prob in log_probs:
-      # Pick the argmax action.
-      action = np.argmax(log_prob)
-      if np.random.uniform() < eps:
-        # Pick an action at random.
-        action = np.random.choice(range(A))
-      actions.append(action)
-    return np.stack(actions)
-
   # We need to reset all environments, if we're coming here the first time.
   if reset or max_timestep is None or max_timestep <= 0:
     env.reset()
@@ -156,13 +135,7 @@ def epsilon_greedy(log_probs):
     assert (B, 1) == value_preds.shape, \
         "B=%d, value_preds.shape=%s" % (B, value_preds.shape)
 
-    actions = None
-    if policy_sampling == GUMBEL_SAMPLING:
-      actions = gumbel_sample(log_probs)
-    elif policy_sampling == EPSILON_GREEDY:
-      actions = epsilon_greedy(log_probs)
-    else:
-      raise ValueError("Unknown sampling policy [%s]" % policy_sampling)
+    actions = gumbel_sample(log_probs)
 
     # Step through the env.
     t1 = time.time()
diff --git a/tensor2tensor/trax/rlax/configs/atari.gin b/tensor2tensor/trax/rlax/configs/atari.gin
index d2f98901d..7734fd0d9 100644
--- a/tensor2tensor/trax/rlax/configs/atari.gin
+++ b/tensor2tensor/trax/rlax/configs/atari.gin
@@ -11,12 +11,10 @@ ppo.training_loop.max_timestep_eval = 20000
 ppo.training_loop.random_seed = 0
 ppo.training_loop.gamma = 0.99
 ppo.training_loop.lambda_ = 0.95
-ppo.training_loop.epsilon = 0.1
 ppo.training_loop.c1 = 1.0
 ppo.training_loop.c2 = 0.01
 ppo.training_loop.eval_every_n = 500
 ppo.training_loop.done_frac_for_policy_save = 0.9
-ppo.training_loop.enable_early_stopping = False
 ppo.training_loop.n_evals = 16
 ppo.training_loop.len_history_for_policy = 4
 ppo.training_loop.eval_temperatures = (1.0, 0.5)
diff --git a/tensor2tensor/trax/rlax/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rlax/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
index b67345b8c..af365de63 100644
--- a/tensor2tensor/trax/rlax/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rlax/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
@@ -9,12 +9,10 @@ ppo.training_loop.max_timestep_eval = 128
 ppo.training_loop.random_seed = 0
 ppo.training_loop.gamma = 0.99
 ppo.training_loop.lambda_ = 0.95
-ppo.training_loop.epsilon = 0.1
 ppo.training_loop.c1 = 1.0
 ppo.training_loop.c2 = 0.01
 ppo.training_loop.eval_every_n = 10
 ppo.training_loop.done_frac_for_policy_save = 0
-ppo.training_loop.enable_early_stopping = True
 ppo.training_loop.n_evals = 1
 ppo.training_loop.len_history_for_policy = 1  # this needs to be bumped up.
 ppo.training_loop.eval_temperatures = (1.0,)
diff --git a/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
index 0fb7382cd..e6ffbea13 100644
--- a/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
@@ -59,12 +59,10 @@ ppo.training_loop.max_timestep_eval = 20000
 ppo.training_loop.random_seed = 0
 ppo.training_loop.gamma = 0.99
 ppo.training_loop.lambda_ = 0.95
-ppo.training_loop.epsilon = 0.1
 ppo.training_loop.c1 = 1.0
 ppo.training_loop.c2 = 0.01
 ppo.training_loop.eval_every_n = 10
 ppo.training_loop.done_frac_for_policy_save = 0
-ppo.training_loop.enable_early_stopping = True
 ppo.training_loop.n_evals = 1
 ppo.training_loop.eval_temperatures = (1.0,)
 ppo.training_loop.len_history_for_policy = 4
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rlax/ppo.py
index 4e276813a..c517ab55a 100644
--- a/tensor2tensor/trax/rlax/ppo.py
+++ b/tensor2tensor/trax/rlax/ppo.py
@@ -156,9 +156,7 @@ def get_params(opt_state):
 def collect_trajectories(env,
                          policy_fn,
                          n_trajectories=1,
-                         policy=env_problem_utils.GUMBEL_SAMPLING,
                          max_timestep=None,
-                         epsilon=0.1,
                          reset=True,
                          len_history_for_policy=32,
                          rng=None):
@@ -168,12 +166,9 @@ def collect_trajectories(env,
     env: A gym env interface, for now this is not-batched.
     policy_fn: observations(B,T+1) -> log-probabs(B,T+1, A) callable.
     n_trajectories: int, number of trajectories.
-    policy: string, "greedy", "epsilon-greedy", or "categorical-sampling" i.e.
-      how to use the policy_fn to return an action.
     max_timestep: int or None, the index of the maximum time-step at which we
       return the trajectory, None for ending a trajectory only when env returns
       done.
-    epsilon: float, the epsilon for `epsilon-greedy` policy.
     reset: bool, true if we want to reset the envs. The envs are also reset if
       max_max_timestep is None or < 0
     len_history_for_policy: int, the maximum history to keep for applying the
@@ -196,8 +191,6 @@ def collect_trajectories(env,
       policy_fn,
       num_trajectories=n_trajectories,
       max_timestep=max_timestep,
-      policy_sampling=policy,
-      eps=epsilon,
       reset=reset,
       len_history_for_policy=len_history_for_policy,
       rng=rng)
@@ -774,7 +767,6 @@ def evaluate_policy(eval_env,
           num_trajectories=eval_env.batch_size,
           max_timestep=max_timestep,
           reset=True,
-          policy_sampling=env_problem_utils.GUMBEL_SAMPLING,
           temperature=temperature,
           rng=eval_rng,
           len_history_for_policy=len_history_for_policy)
@@ -854,165 +846,130 @@ def write_eval_reward_summaries(reward_stats_by_mode, summary_writer, epoch):
                    reward_stats["mean"], reward_stats["std"])
 
 
-@gin.configurable(blacklist=["output_dir"])
-def training_loop(
-    env,
-    eval_env,
-    env_name,
-    policy_and_value_net_fn,
-    policy_and_value_optimizer_fn,
-    output_dir,
-    epochs=EPOCHS,
-    n_optimizer_steps=N_OPTIMIZER_STEPS,
-    print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
-    target_kl=0.01,
-    boundary=20,
-    max_timestep=None,
-    max_timestep_eval=20000,
-    random_seed=None,
-    gamma=GAMMA,
-    lambda_=LAMBDA,
-    epsilon=EPSILON,
-    c1=1.0,
-    c2=0.01,
-    eval_every_n=1000,
-    done_frac_for_policy_save=0.5,
-    enable_early_stopping=True,
-    n_evals=1,
-    len_history_for_policy=4,
-    eval_temperatures=(1.0, 0.5),
-):
-  """Runs the training loop for PPO, with fixed policy and value nets.
-
-  Args:
-    env: gym.Env to use for training.
-    eval_env: gym.Env to use for evaluation.
-    env_name: Name of the environment.
-    policy_and_value_net_fn: Function defining the policy and value network.
-    policy_and_value_optimizer_fn: Function defining the optimizer.
-    output_dir: Output dir.
-    epochs: Number of epochs to run for.
-    n_optimizer_steps: Number of optimizer steps.
-    print_every_optimizer_steps: How often to log during the policy optimization
-      process.
-    target_kl: Policy iteration early stopping.
-    boundary: We pad trajectories at integer multiples of this number.
-    max_timestep: If set to an integer, maximum number of time-steps in
-      a trajectory. Used in the collect procedure.
-    max_timestep_eval: If set to an integer, maximum number of time-steps in an
-      evaluation trajectory. Used in the collect procedure.
-    random_seed: Random seed.
-    gamma: Reward discount factor.
-    lambda_: N-step TD-error discount factor in GAE.
-    epsilon: Random action probability in epsilon-greedy sampling.
-    c1: Value loss coefficient.
-    c2: Entropy loss coefficient.
-    eval_every_n: How frequently to eval the policy.
-    done_frac_for_policy_save: Fraction of the trajectories that should be done
-      to checkpoint the policy.
-    enable_early_stopping: Whether to enable early stopping.
-    n_evals: Number of times to evaluate.
-    len_history_for_policy: How much of history to give to the policy.
-    eval_temperatures: Sequence of temperatures to try for categorical sampling
-      during evaluation.
-  """
-  gfile.makedirs(output_dir)
-
-  # Create summary writers and history.
-  train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
-  timing_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "timing"))
-  eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
-
-  train_sw.text("env_name", env_name)
-  timing_sw.text("env_name", env_name)
-  eval_sw.text("env_name", env_name)
-
-  jax_rng_key = trax.get_random_number_generator_and_set_seed(random_seed)
-
-  # Batch Observations Shape = [1, 1] + OBS, because we will eventually call
-  # policy and value networks on shape [B, T] +_OBS
-  batch_observations_shape = (1, 1) + env.observation_space.shape
-  observations_dtype = env.observation_space.dtype
-
-  assert isinstance(env.action_space, gym.spaces.Discrete)
-  n_actions = env.action_space.n
-
-  jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)
-
-  # Initialize the policy and value network.
-  policy_and_value_net_params, policy_and_value_net_apply = (
-      policy_and_value_net_fn(key1, batch_observations_shape,
-                              observations_dtype, n_actions))
-
-  # Maybe restore the policy params. If there is nothing to restore, then
-  # iteration = 0 and policy_and_value_net_params are returned as is.
-  restore, policy_and_value_net_params, iteration = (
-      maybe_restore_params(output_dir, policy_and_value_net_params))
-
-  if restore:
-    logging.info("Restored parameters from iteration [%d]", iteration)
-    # We should start from the next iteration.
-    iteration += 1
-
-  policy_and_value_net_apply = jit(policy_and_value_net_apply)
-
-  # Initialize the optimizers.
-  policy_and_value_optimizer = (
-      policy_and_value_optimizer_fn(policy_and_value_net_params))
-  (policy_and_value_opt_state, policy_and_value_opt_update,
-   policy_and_value_get_params) = policy_and_value_optimizer
-
-  n_trajectories_done = 0
-  last_saved_at = 0
-
-  logging.info("Starting the PPO training loop.")
-  for i in range(iteration, epochs):
+class PPOTrainer(object):
+  """PPO trainer."""
+
+  def __init__(
+      self,
+      train_env,
+      eval_env,
+      policy_and_value_net_fn,
+      policy_and_value_optimizer_fn,
+      output_dir,
+      n_optimizer_steps,
+      print_every_optimizer_steps,
+      target_kl,
+      boundary,
+      max_timestep,
+      max_timestep_eval,
+      random_seed,
+      gamma,
+      lambda_,
+      c1,
+      c2,
+      eval_every_n,
+      done_frac_for_policy_save,
+      n_evals,
+      len_history_for_policy,
+      eval_temperatures,
+  ):
+    self._train_env = train_env
+    self._eval_env = eval_env
+    self._n_optimizer_steps = n_optimizer_steps
+    self._print_every_optimizer_steps = print_every_optimizer_steps
+    self._target_kl = target_kl
+    self._boundary = boundary
+    self._max_timestep = max_timestep
+    self._max_timestep_eval = max_timestep_eval
+    self._gamma = gamma
+    self._lambda_ = lambda_
+    self._c1 = c1
+    self._c2 = c2
+    self._eval_every_n = eval_every_n
+    self._done_frac_for_policy_save = done_frac_for_policy_save
+    self._n_evals = n_evals
+    self._len_history_for_policy = len_history_for_policy
+    self._eval_temperatures = eval_temperatures
+
+    assert isinstance(self._train_env.action_space, gym.spaces.Discrete)
+    n_actions = self._train_env.action_space.n
+
+    # Batch Observations Shape = [1, 1] + OBS, because we will eventually call
+    # policy and value networks on shape [B, T] +_OBS
+    batch_observations_shape = (1, 1) + self._train_env.observation_space.shape
+    observations_dtype = self._train_env.observation_space.dtype
+
+    self._rng = trax.get_random_number_generator_and_set_seed(random_seed)
+    self._rng, key1 = jax_random.split(self._rng, num=2)
+
+    # Initialize the policy and value network.
+    policy_and_value_net_params, policy_and_value_net_apply = (
+        policy_and_value_net_fn(key1, batch_observations_shape,
+                                observations_dtype, n_actions))
+
+    self._policy_and_value_net_apply = jit(policy_and_value_net_apply)
+
+    # Maybe restore the policy params. If there is nothing to restore, then
+    # iteration = 0 and policy_and_value_net_params are returned as is.
+    restored, policy_and_value_net_params, self._epoch = (
+        maybe_restore_params(output_dir, policy_and_value_net_params))
+
+    if restored:
+      logging.info("Restored parameters from iteration [%d]", self._epoch)
+      # We should start from the next iteration.
+      self._epoch += 1
+
+    # Initialize the optimizers.
+    policy_and_value_optimizer = (
+        policy_and_value_optimizer_fn(policy_and_value_net_params))
+    (self._policy_and_value_opt_state, self._policy_and_value_opt_update,
+     self._policy_and_value_get_params) = policy_and_value_optimizer
+
+    self._output_dir = output_dir
+    gfile.makedirs(self._output_dir)
+
+    # Create summary writers and history.
+    self._train_sw = jaxboard.SummaryWriter(
+        os.path.join(self._output_dir, "train"))
+    self._timing_sw = jaxboard.SummaryWriter(
+        os.path.join(self._output_dir, "timing"))
+    self._eval_sw = jaxboard.SummaryWriter(
+        os.path.join(self._output_dir, "eval"))
+
+    self._should_reset = True
+    self._n_trajectories_done = 0
+
+    self._last_saved_at = 0
+
+  @property
+  def epoch(self):
+    return self._epoch
+
+  def train_epoch(self):
+    """Train one PPO epoch."""
     epoch_start_time = time.time()
 
-    # Params we'll use to collect the trajectories.
-    policy_and_value_net_params = policy_and_value_get_params(
-        policy_and_value_opt_state)
-
-    # A function to get the policy and value predictions.
-    def get_predictions(observations, rng=None):
-      """Returns log-probs, value predictions and key back."""
-      key, key1 = jax_random.split(rng, num=2)
-
-      log_probs, value_preds = policy_and_value_net_apply(
-          observations, policy_and_value_net_params, rng=key1)
-
-      return log_probs, value_preds, key
-
     # Evaluate the policy.
     policy_eval_start_time = time.time()
-    if ((i + 1) % eval_every_n == 0) or (i == epochs - 1):
-      jax_rng_key, key = jax_random.split(jax_rng_key, num=2)
-
-      logging.vlog(1, "Epoch [% 6d] evaluating policy.", i)
+    if (self._epoch + 1) % self._eval_every_n == 0:
+      self._rng, key = jax_random.split(self._rng, num=2)
+      self.evaluate()
 
-      reward_stats = evaluate_policy(
-          eval_env,
-          get_predictions,
-          temperatures=eval_temperatures,
-          max_timestep=max_timestep_eval,
-          n_evals=n_evals,
-          len_history_for_policy=len_history_for_policy,
-          rng=key)
-      write_eval_reward_summaries(reward_stats, eval_sw, epoch=i)
     policy_eval_time = get_time(policy_eval_start_time)
 
     trajectory_collection_start_time = time.time()
-    logging.vlog(1, "Epoch [% 6d] collecting trajectories.", i)
-    jax_rng_key, key = jax_random.split(jax_rng_key)
+    logging.vlog(1, "Epoch [% 6d] collecting trajectories.", self._epoch)
+    self._rng, key = jax_random.split(self._rng)
     trajs, n_done, timing_info = collect_trajectories(
-        env,
-        policy_fn=get_predictions,
-        n_trajectories=env.batch_size,
-        max_timestep=max_timestep,
+        self._train_env,
+        policy_fn=self._get_predictions,
+        n_trajectories=self._train_env.batch_size,
+        max_timestep=self._max_timestep,
         rng=key,
-        len_history_for_policy=len_history_for_policy,
-        reset=(i == 0) or restore,
-        epsilon=(10.0 / (i + 10.0)))  # this is a different epsilon.
+        len_history_for_policy=self._len_history_for_policy,
+        reset=self._should_reset,
+    )
+    self._should_reset = False
     trajectory_collection_time = get_time(trajectory_collection_start_time)
 
     logging.vlog(1, "Collecting trajectories took %0.2f msec.",
@@ -1022,7 +979,8 @@ def get_predictions(observations, rng=None):
     max_reward = max(np.sum(traj[2]) for traj in trajs)
     min_reward = min(np.sum(traj[2]) for traj in trajs)
 
-    train_sw.scalar("train/reward_mean_truncated", avg_reward, step=i)
+    self._train_sw.scalar(
+        "train/reward_mean_truncated", avg_reward, step=self._epoch)
 
     logging.vlog(1, "Rewards avg=[%0.2f], max=[%0.2f], min=[%0.2f], all=%s",
                  avg_reward, max_reward, min_reward,
@@ -1038,7 +996,7 @@ def get_predictions(observations, rng=None):
     padding_start_time = time.time()
     (_, reward_mask, padded_observations, padded_actions,
      padded_rewards, padded_infos) = pad_trajectories(
-         trajs, boundary=boundary)
+         trajs, boundary=self._boundary)
     padding_time = get_time(padding_start_time)
 
     logging.vlog(1, "Padding trajectories took %0.2f msec.",
@@ -1053,7 +1011,8 @@ def get_predictions(observations, rng=None):
     assert (B, T) == padded_rewards.shape
     assert (B, T) == reward_mask.shape
     assert (B, T + 1) == padded_observations.shape[:2]
-    assert (B, T + 1) + env.observation_space.shape == padded_observations.shape
+    assert ((B, T + 1) + self._train_env.observation_space.shape ==
+            padded_observations.shape)
 
     log_prob_recompute_start_time = time.time()
     assert ("log_prob_actions" in padded_infos and
@@ -1073,8 +1032,8 @@ def get_predictions(observations, rng=None):
     # NOTE: We don't have the log-probabs and value-predictions for the last
     # observation, so we re-calculate for everything, but use the original ones
     # for all but the last time-step.
-    jax_rng_key, key = jax_random.split(jax_rng_key)
-    log_probabs_traj, value_predictions_traj, _ = get_predictions(
+    self._rng, key = jax_random.split(self._rng)
+    log_probabs_traj, value_predictions_traj, _ = self._get_predictions(
         padded_observations, rng=key)
 
     assert (B, T + 1, A) == log_probabs_traj.shape
@@ -1090,33 +1049,24 @@ def get_predictions(observations, rng=None):
 
     log_prob_recompute_time = get_time(log_prob_recompute_start_time)
 
-    # Linear annealing from 0.1 to 0.0
-    # epsilon_schedule = epsilon if epochs == 1 else epsilon * (1.0 -
-    #                                                           (i /
-    #                                                            (epochs - 1)))
-
-    # Constant epsilon.
-    epsilon_schedule = epsilon
-
     # Compute value and ppo losses.
-    jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)
+    self._rng, key1 = jax_random.split(self._rng, num=2)
     logging.vlog(2, "Starting to compute P&V loss.")
     loss_compute_start_time = time.time()
     cur_combined_loss, cur_ppo_loss, cur_value_loss, entropy_bonus = (
         combined_loss(
-            policy_and_value_net_params,
+            self._policy_and_value_net_params,
             log_probabs_traj,
             value_predictions_traj,
-            policy_and_value_net_apply,
+            self._policy_and_value_net_apply,
             padded_observations,
             padded_actions,
             padded_rewards,
             reward_mask,
-            gamma=gamma,
-            lambda_=lambda_,
-            epsilon=epsilon_schedule,
-            c1=c1,
-            c2=c2,
+            gamma=self._gamma,
+            lambda_=self._lambda_,
+            c1=self._c1,
+            c2=self._c2,
             rng=key1))
     loss_compute_time = get_time(loss_compute_start_time)
     logging.vlog(
@@ -1125,44 +1075,40 @@ def get_predictions(observations, rng=None):
         cur_combined_loss, cur_value_loss, cur_ppo_loss, entropy_bonus,
         get_time(loss_compute_start_time))
 
-    jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)
+    self._rng, key1 = jax_random.split(self._rng, num=2)
     logging.vlog(1, "Policy and Value Optimization")
     optimization_start_time = time.time()
-    keys = jax_random.split(key1, num=n_optimizer_steps)
-    for j in range(n_optimizer_steps):
-      k1, k2, k3 = jax_random.split(keys[j], num=3)
+    keys = jax_random.split(key1, num=self._n_optimizer_steps)
+    for (j, key) in enumerate(keys):
+      k1, k2, k3 = jax_random.split(key, num=3)
       t = time.time()
       # Update the optimizer state.
-      policy_and_value_opt_state = policy_and_value_opt_step(
+      self._policy_and_value_opt_state = policy_and_value_opt_step(
           j,
-          policy_and_value_opt_state,
-          policy_and_value_opt_update,
-          policy_and_value_get_params,
-          policy_and_value_net_apply,
+          self._policy_and_value_opt_state,
+          self._policy_and_value_opt_update,
+          self._policy_and_value_get_params,
+          self._policy_and_value_net_apply,
           log_probabs_traj,
           value_predictions_traj,
           padded_observations,
           padded_actions,
           padded_rewards,
           reward_mask,
-          c1=c1,
-          c2=c2,
-          gamma=gamma,
-          lambda_=lambda_,
-          epsilon=epsilon_schedule,
+          c1=self._c1,
+          c2=self._c2,
+          gamma=self._gamma,
+          lambda_=self._lambda_,
           rng=k1)
 
       # Compute the approx KL for early stopping.
-      new_policy_and_value_net_params = policy_and_value_get_params(
-          policy_and_value_opt_state)
-
-      log_probab_actions_new, _ = policy_and_value_net_apply(
-          padded_observations, new_policy_and_value_net_params, rng=k2)
+      log_probab_actions_new, _ = self._policy_and_value_net_apply(
+          padded_observations, self._policy_and_value_net_params, rng=k2)
 
       approx_kl = approximate_kl(log_probab_actions_new, log_probabs_traj,
                                  reward_mask)
 
-      early_stopping = enable_early_stopping and approx_kl > 1.5 * target_kl
+      early_stopping = approx_kl > 1.5 * self._target_kl
       if early_stopping:
         logging.vlog(
             1, "Early stopping policy and value optimization at iter: %d, "
@@ -1171,24 +1117,23 @@ def get_predictions(observations, rng=None):
         # iteration.
 
       t2 = time.time()
-      if (((j + 1) % print_every_optimizer_steps == 0) or
-          (j == n_optimizer_steps - 1) or early_stopping):
+      if (((j + 1) % self._print_every_optimizer_steps == 0) or
+          (j == self._n_optimizer_steps - 1) or early_stopping):
         # Compute and log the loss.
         (loss_combined, loss_ppo, loss_value, entropy_bonus) = (
             combined_loss(
-                new_policy_and_value_net_params,
+                self._policy_and_value_net_params,
                 log_probabs_traj,
                 value_predictions_traj,
-                policy_and_value_net_apply,
+                self._policy_and_value_net_apply,
                 padded_observations,
                 padded_actions,
                 padded_rewards,
                 reward_mask,
-                gamma=gamma,
-                lambda_=lambda_,
-                epsilon=epsilon_schedule,
-                c1=c1,
-                c2=c2,
+                gamma=self._gamma,
+                lambda_=self._lambda_,
+                c1=self._c1,
+                c2=self._c2,
                 rng=k3))
         logging.vlog(1, "One Policy and Value grad desc took: %0.2f msec",
                      get_time(t, t2))
@@ -1212,30 +1157,21 @@ def get_predictions(observations, rng=None):
     # Also don't save too frequently, enforce a minimum gap.
     # Or if this is the last iteration.
     policy_save_start_time = time.time()
-    n_trajectories_done += n_done
+    self._n_trajectories_done += n_done
     # TODO(afrozm): Refactor to trax.save_state.
-    if (((n_trajectories_done >= done_frac_for_policy_save * env.batch_size) and
-         (i - last_saved_at > eval_every_n) and
-         (((i + 1) % eval_every_n == 0))) or (i == epochs - 1)):
-      logging.vlog(1, "Epoch [% 6d] saving model.", i)
-      old_model_files = gfile.glob(os.path.join(output_dir, "model-??????.pkl"))
-      params_file = os.path.join(output_dir, "model-%06d.pkl" % i)
-      with gfile.GFile(params_file, "wb") as f:
-        pickle.dump(policy_and_value_net_params, f)
-      # Remove the old model files.
-      for path in old_model_files:
-        gfile.remove(path)
-      # Reset this number.
-      n_trajectories_done = 0
-      last_saved_at = i
+    if ((self._n_trajectories_done >=
+         self._done_frac_for_policy_save * self._train_env.batch_size) and
+        (self._epoch - self._last_saved_at > self._eval_every_n) and
+        (((self._epoch + 1) % self._eval_every_n == 0))):
+      self.save()
     policy_save_time = get_time(policy_save_start_time)
 
     epoch_time = get_time(epoch_start_time)
 
     logging.info(
         "Epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
-        " Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)]", i, min_reward,
-        max_reward, avg_reward, loss_combined, loss_value, loss_ppo,
+        " Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)]", self._epoch,
+        min_reward, max_reward, avg_reward, loss_combined, loss_value, loss_ppo,
         entropy_bonus)
 
     timing_dict = {
@@ -1252,20 +1188,154 @@ def get_predictions(observations, rng=None):
     timing_dict.update(timing_info)
 
     for k, v in timing_dict.items():
-      timing_sw.scalar("timing/%s" % k, v, step=i)
+      self._timing_sw.scalar("timing/%s" % k, v, step=self._epoch)
 
     max_key_len = max(len(k) for k in timing_dict)
     timing_info_list = [
         "%s : % 10.2f" % (k.rjust(max_key_len + 1), v)
         for k, v in sorted(timing_dict.items())
     ]
-    logging.info("Epoch [% 6d], Timings: \n%s", i, "\n".join(timing_info_list))
+    logging.info(
+        "Epoch [% 6d], Timings: \n%s", self._epoch, "\n".join(timing_info_list))
 
-    # Reset restore.
-    restore = False
+    self._epoch += 1
 
     # Flush summary writers once in a while.
-    if (i + 1) % 1000 == 0 or i == epochs - 1:
-      train_sw.flush()
-      timing_sw.flush()
-      eval_sw.flush()
+    if (self._epoch + 1) % 1000 == 0:
+      self.flush_summaries()
+
+  def evaluate(self):
+    """Evaluate the agent."""
+    logging.vlog(1, "Epoch [% 6d] evaluating policy.", self._epoch)
+    self._rng, key = jax_random.split(self._rng, num=2)
+    reward_stats = evaluate_policy(
+        self._eval_env,
+        self._get_predictions,
+        temperatures=self._eval_temperatures,
+        max_timestep=self._max_timestep_eval,
+        n_evals=self._n_evals,
+        len_history_for_policy=self._len_history_for_policy,
+        rng=key)
+    write_eval_reward_summaries(reward_stats, self._eval_sw, epoch=self._epoch)
+
+  def save(self):
+    """Save the agent parameters."""
+    logging.vlog(1, "Epoch [% 6d] saving model.", self._epoch)
+    old_model_files = gfile.glob(
+        os.path.join(self._output_dir, "model-??????.pkl"))
+    params_file = os.path.join(self._output_dir, "model-%06d.pkl" % self._epoch)
+    with gfile.GFile(params_file, "wb") as f:
+      pickle.dump(self._policy_and_value_net_params, f)
+    # Remove the old model files.
+    for path in old_model_files:
+      gfile.remove(path)
+    # Reset this number.
+    self._n_trajectories_done = 0
+    self._last_saved_at = self._epoch
+
+  def flush_summaries(self):
+    self._train_sw.flush()
+    self._timing_sw.flush()
+    self._eval_sw.flush()
+
+  @property
+  def _policy_and_value_net_params(self):
+    return self._policy_and_value_get_params(self._policy_and_value_opt_state)
+
+  # A function to get the policy and value predictions.
+  def _get_predictions(self, observations, rng=None):
+    """Returns log-probs, value predictions and key back."""
+    key, key1 = jax_random.split(rng, num=2)
+
+    log_probs, value_preds = self._policy_and_value_net_apply(
+        observations, self._policy_and_value_net_params, rng=key1)
+
+    return log_probs, value_preds, key
+
+
+@gin.configurable(blacklist=["output_dir"])
+def training_loop(
+    train_env,
+    eval_env,
+    policy_and_value_net_fn,
+    policy_and_value_optimizer_fn,
+    output_dir,
+    epochs=EPOCHS,
+    n_optimizer_steps=N_OPTIMIZER_STEPS,
+    print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
+    target_kl=0.01,
+    boundary=20,
+    max_timestep=None,
+    max_timestep_eval=20000,
+    random_seed=None,
+    gamma=GAMMA,
+    lambda_=LAMBDA,
+    c1=1.0,
+    c2=0.01,
+    eval_every_n=1000,
+    done_frac_for_policy_save=0.5,
+    n_evals=1,
+    len_history_for_policy=4,
+    eval_temperatures=(1.0, 0.5),
+):
+  """Runs the training loop for PPO, with fixed policy and value nets.
+
+  Args:
+    train_env: gym.Env to use for training.
+    eval_env: gym.Env to use for evaluation.
+    policy_and_value_net_fn: Function defining the policy and value network.
+    policy_and_value_optimizer_fn: Function defining the optimizer.
+    output_dir: Output dir.
+    epochs: Number of epochs to run for.
+    n_optimizer_steps: Number of optimizer steps.
+    print_every_optimizer_steps: How often to log during the policy optimization
+      process.
+    target_kl: Policy iteration early stopping. Set to infinity to disable early
+      stopping.
+    boundary: We pad trajectories at integer multiples of this number.
+    max_timestep: If set to an integer, maximum number of time-steps in
+      a trajectory. Used in the collect procedure.
+    max_timestep_eval: If set to an integer, maximum number of time-steps in an
+      evaluation trajectory. Used in the collect procedure.
+    random_seed: Random seed.
+    gamma: Reward discount factor.
+    lambda_: N-step TD-error discount factor in GAE.
+    c1: Value loss coefficient.
+    c2: Entropy loss coefficient.
+    eval_every_n: How frequently to eval the policy.
+    done_frac_for_policy_save: Fraction of the trajectories that should be done
+      to checkpoint the policy.
+    n_evals: Number of times to evaluate.
+    len_history_for_policy: How much of history to give to the policy.
+    eval_temperatures: Sequence of temperatures to try for categorical sampling
+      during evaluation.
+  """
+  trainer = PPOTrainer(
+      train_env=train_env,
+      eval_env=eval_env,
+      policy_and_value_net_fn=policy_and_value_net_fn,
+      policy_and_value_optimizer_fn=policy_and_value_optimizer_fn,
+      output_dir=output_dir,
+      n_optimizer_steps=n_optimizer_steps,
+      print_every_optimizer_steps=print_every_optimizer_steps,
+      target_kl=target_kl,
+      boundary=boundary,
+      max_timestep=max_timestep,
+      max_timestep_eval=max_timestep_eval,
+      random_seed=random_seed,
+      gamma=gamma,
+      lambda_=lambda_,
+      c1=c1,
+      c2=c2,
+      eval_every_n=eval_every_n,
+      done_frac_for_policy_save=done_frac_for_policy_save,
+      n_evals=n_evals,
+      len_history_for_policy=len_history_for_policy,
+      eval_temperatures=eval_temperatures,
+  )
+  logging.info("Starting the PPO training loop.")
+  for _ in range(trainer.epoch, epochs):
+    trainer.train_epoch()
+  trainer.save()
+  trainer.evaluate()
+  trainer.flush_summaries()
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rlax/ppo_main.py
index 7fe717e68..a5990cef4 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rlax/ppo_main.py
@@ -175,20 +175,16 @@ def main(argv):
   gin.parse_config_files_and_bindings(FLAGS.config_file, gin_configs)
 
   # TODO(pkozakowski): Find a better way to determine this.
-  env_kwargs = {}
   train_env_kwargs = {}
   eval_env_kwargs = {}
   if "OnlineTuneEnv" in FLAGS.env_problem_name:
     # TODO(pkozakowski): Separate env output dirs by train/eval and epoch.
-    train_env_kwargs = {}
-    train_env_kwargs.update(env_kwargs)
-    train_env_kwargs["output_dir"] = os.path.join(FLAGS.output_dir,
-                                                  "envs/train")
-
-    eval_env_kwargs = {}
-    eval_env_kwargs.update(env_kwargs)
-    eval_env_kwargs["output_dir"] = os.path.join(FLAGS.output_dir,
-                                                 "envs/eval")
+    train_env_kwargs = {
+        "output_dir": os.path.join(FLAGS.output_dir, "envs/train")
+    }
+    eval_env_kwargs = {
+        "output_dir": os.path.join(FLAGS.output_dir, "envs/eval")
+    }
 
   if "ClientEnv" in FLAGS.env_problem_name:
     train_env_kwargs["per_env_kwargs"] = [{
@@ -200,8 +196,8 @@ def main(argv):
     } for replica in range(FLAGS.eval_batch_size)]
 
   # Make an env here.
-  env = make_env(batch_size=FLAGS.batch_size, **train_env_kwargs)
-  assert env
+  train_env = make_env(batch_size=FLAGS.batch_size, **train_env_kwargs)
+  assert train_env
 
   eval_env = make_env(batch_size=FLAGS.eval_batch_size, **eval_env_kwargs)
   assert eval_env
@@ -224,9 +220,8 @@ def run_training_loop():
 
     ppo.training_loop(
         output_dir=FLAGS.output_dir,
-        env=env,
+        train_env=train_env,
         eval_env=eval_env,
-        env_name=str(FLAGS.env_problem_name),
         policy_and_value_net_fn=policy_and_value_net_fn,
         policy_and_value_optimizer_fn=policy_and_value_optimizer_fn,
     )
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
index 641f10155..269a438f7 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rlax/ppo_training_loop_test.py
@@ -69,11 +69,11 @@ def tmp_dir(self):
     yield tmp
     gfile.rmtree(tmp)
 
-  def _run_training_loop(self, env, eval_env, output_dir):
+  def _run_training_loop(self, train_env, eval_env, output_dir):
     n_epochs = 2
     # Run the training loop.
     ppo.training_loop(
-        env=env,
+        train_env=train_env,
         eval_env=eval_env,
         epochs=n_epochs,
         policy_and_value_net_fn=functools.partial(
@@ -82,13 +82,12 @@ def _run_training_loop(self, env, eval_env, output_dir):
         policy_and_value_optimizer_fn=ppo.optimizer_fn,
         n_optimizer_steps=1,
         output_dir=output_dir,
-        env_name="SomeEnv",
         random_seed=0)
 
   def test_training_loop_cartpole(self):
     with self.tmp_dir() as output_dir:
       self._run_training_loop(
-          env=self.get_wrapped_env("CartPole-v0", 2),
+          train_env=self.get_wrapped_env("CartPole-v0", 2),
           eval_env=self.get_wrapped_env("CartPole-v0", 2),
           output_dir=output_dir,
       )
@@ -112,7 +111,7 @@ def test_training_loop_onlinetune(self):
       gin.bind_parameter(
           "OnlineTuneEnv.output_dir", os.path.join(output_dir, "envs"))
       self._run_training_loop(
-          env=self.get_wrapped_env("OnlineTuneEnv-v0", 2),
+          train_env=self.get_wrapped_env("OnlineTuneEnv-v0", 2),
           eval_env=self.get_wrapped_env("OnlineTuneEnv-v0", 2),
           output_dir=output_dir,
       )
@@ -189,7 +188,7 @@ def loss(*args, **kwargs):
       )
 
       self._run_training_loop(
-          env=env_fn(),
+          train_env=env_fn(),
           eval_env=env_fn(),
           output_dir=output_dir,
       )

From ef8ac56be4acd83af9a836812694aae1530f0793 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 14 Aug 2019 13:05:35 -0700
Subject: [PATCH 2284/2720] Trax: rename rlax->rl and make rl/ppo_main a
 top-level rl_trainer binary.

PiperOrigin-RevId: 263411743
---
 docs/cloud_mlengine.md                        |  90 +++++++
 docs/cloud_tpu.md                             |  50 ++++
 docs/distributed_training.md                  | 220 ++++++++++++++++
 docs/index.md                                 | 127 +++++++++
 docs/multi_problem.md                         | 188 ++++++++++++++
 docs/new_model.md                             | 114 ++++++++
 docs/new_problem.md                           | 243 ++++++++++++++++++
 docs/overview.md                              | 175 +++++++++++++
 docs/tutorials/asr_with_transformer.md        |   4 +
 tensor2tensor/trax/{rlax => rl}/__init__.py   |   0
 .../trax/{rlax => rl}/configs/atari.gin       |   2 +-
 ...params_online_tune_wide_resnet_cifar10.gin |   4 +-
 ...params_online_tune_wide_resnet_cifar10.gin |   2 +-
 .../online_tune_wide_resnet_cifar10.gin       |   4 +-
 .../trax/{rlax => rl}/envs/__init__.py        |   8 +-
 .../{rlax => rl}/envs/env_service_server.py   |   2 +-
 .../trax/{rlax => rl}/envs/fake_env.py        |   0
 .../trax/{rlax => rl}/envs/fake_env_test.py   |   4 +-
 .../trax/{rlax => rl}/envs/online_tune_env.py |   0
 .../{rlax => rl}/envs/online_tune_env_test.py |   4 +-
 tensor2tensor/trax/{rlax => rl}/ppo.py        |   0
 tensor2tensor/trax/{rlax => rl}/ppo_test.py   |   4 +-
 .../{rlax => rl}/ppo_training_loop_test.py    |   8 +-
 .../{rlax => rl}/simulated_env_problem.py     |   2 +-
 .../simulated_env_problem_test.py             |   4 +-
 .../trax/{rlax => rl}/space_serializer.py     |   0
 .../{rlax => rl}/space_serializer_test.py     |   4 +-
 .../trax/{rlax/ppo_main.py => rl_trainer.py}  |  10 +-
 28 files changed, 1243 insertions(+), 30 deletions(-)
 create mode 100644 docs/cloud_mlengine.md
 create mode 100644 docs/cloud_tpu.md
 create mode 100644 docs/distributed_training.md
 create mode 100644 docs/index.md
 create mode 100644 docs/multi_problem.md
 create mode 100644 docs/new_model.md
 create mode 100644 docs/new_problem.md
 create mode 100644 docs/overview.md
 create mode 100644 docs/tutorials/asr_with_transformer.md
 rename tensor2tensor/trax/{rlax => rl}/__init__.py (100%)
 rename tensor2tensor/trax/{rlax => rl}/configs/atari.gin (95%)
 rename tensor2tensor/trax/{rlax => rl}/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin (95%)
 rename tensor2tensor/trax/{rlax => rl}/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin (95%)
 rename tensor2tensor/trax/{rlax => rl}/configs/online_tune_wide_resnet_cifar10.gin (97%)
 rename tensor2tensor/trax/{rlax => rl}/envs/__init__.py (78%)
 rename tensor2tensor/trax/{rlax => rl}/envs/env_service_server.py (98%)
 rename tensor2tensor/trax/{rlax => rl}/envs/fake_env.py (100%)
 rename tensor2tensor/trax/{rlax => rl}/envs/fake_env_test.py (94%)
 rename tensor2tensor/trax/{rlax => rl}/envs/online_tune_env.py (100%)
 rename tensor2tensor/trax/{rlax => rl}/envs/online_tune_env_test.py (97%)
 rename tensor2tensor/trax/{rlax => rl}/ppo.py (100%)
 rename tensor2tensor/trax/{rlax => rl}/ppo_test.py (99%)
 rename tensor2tensor/trax/{rlax => rl}/ppo_training_loop_test.py (96%)
 rename tensor2tensor/trax/{rlax => rl}/simulated_env_problem.py (99%)
 rename tensor2tensor/trax/{rlax => rl}/simulated_env_problem_test.py (98%)
 rename tensor2tensor/trax/{rlax => rl}/space_serializer.py (100%)
 rename tensor2tensor/trax/{rlax => rl}/space_serializer_test.py (96%)
 rename tensor2tensor/trax/{rlax/ppo_main.py => rl_trainer.py} (97%)

diff --git a/docs/cloud_mlengine.md b/docs/cloud_mlengine.md
new file mode 100644
index 000000000..83ebe7e57
--- /dev/null
+++ b/docs/cloud_mlengine.md
@@ -0,0 +1,90 @@
+# Running on Cloud ML Engine
+
+Google Cloud Platform offers a managed training environment for TensorFlow
+models called [Cloud ML Engine](https://cloud.google.com/ml-engine/) and
+you can easily launch Tensor2Tensor on it, including for hyperparameter tuning.
+
+# Launch
+
+It's the same `t2t-trainer` you know and love with the addition of the
+`--cloud_mlengine` flag, which by default will launch on a 1-GPU machine
+in the default compute region. See the
+[docs for `gcloud compute`](https://cloud.google.com/compute/docs/gcloud-compute/#set_default_zone_and_region_in_your_local_client)
+to learn how to set the default compute region.
+
+```
+# Note that both the data dir and output dir have to be on GCS
+DATA_DIR=gs://my-bucket/data
+OUTPUT_DIR=gs://my-bucket/train
+t2t-trainer \
+  --problem=translate_ende_wmt32k \
+  --model=transformer \
+  --hparams_set=transformer_base \
+  --data_dir=$DATA_DIR \
+  --output_dir=$OUTPUT_DIR \
+  --cloud_mlengine
+```
+
+By passing `--worker_gpu=4` or `--worker_gpu=8` it will automatically launch on
+machines with 4 or 8 GPUs.
+
+You can additionally pass the `--cloud_mlengine_master_type` to select another
+kind of machine (see the [docs for
+`masterType`](https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#traininginput)
+for options, including
+[ML Engine machine
+types](https://cloud.google.com/ml-engine/docs/training-overview)
+and their
+[specs](https://cloud.google.com/compute/docs/machine-types)).
+If you provide this flag yourself, make sure you pass the
+correct value for `--worker_gpu` (for non-GPU machines, you should pass
+`--worker_gpu=0`).
+
+**Note**: `t2t-trainer` only currently supports launching with single machines,
+possibly with multiple GPUs. Multi-machine setups are not yet supported out of
+the box with the `--cloud_mlengine` flag, though multi-machine should in
+principle work just fine. Contributions/testers welcome.
+
+
+## `--t2t_usr_dir`
+
+Launching on Cloud ML Engine works with `--t2t_usr_dir` as well as long as the
+directory is fully self-contained (i.e. the imports only refer to other modules
+in the directory). If there are additional PyPI dependencies that you need, you
+can include a `requirements.txt` file in the directory specified by
+`t2t_usr_dir`.
+
+# Hyperparameter Tuning
+
+Hyperparameter tuning with `t2t-trainer` and Cloud ML Engine is also a breeze
+with `--hparams_range` and the `--autotune_*` flags:
+
+```
+t2t-trainer \
+  --problem=translate_ende_wmt32k \
+  --model=transformer \
+  --hparams_set=transformer_base \
+  --data_dir=$DATA_DIR \
+  --output_dir=$OUTPUT_DIR \
+  --cloud_mlengine \
+  --hparams_range=transformer_base_range \
+  --autotune_objective='metrics-translate_ende_wmt32k/neg_log_perplexity' \
+  --autotune_maximize \
+  --autotune_max_trials=100 \
+  --autotune_parallel_trials=3
+```
+
+The `--hparams_range` specifies the search space and should be registered with
+`@register_ranged_hparams`. It defines a `RangedHParams` object that sets
+search ranges and scales for various parameters. See `transformer_base_range`
+in
+[`transformer.py`](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/models/transformer.py)
+for an example.
+
+The metric name passed as `--autotune_objective` should be exactly what you'd
+see in TensorBoard. To minimize a metric, set `--autotune_maximize=False`.
+
+You control how many total trials to run with `--autotune_max_trials` and the
+number of jobs to launch in parallel with `--autotune_parallel_trials`.
+
+Happy tuning!
diff --git a/docs/cloud_tpu.md b/docs/cloud_tpu.md
new file mode 100644
index 000000000..c0625e132
--- /dev/null
+++ b/docs/cloud_tpu.md
@@ -0,0 +1,50 @@
+# Running on Cloud TPUs
+
+Tensor2Tensor supports running on Google Cloud Platforms TPUs, chips
+specialized for ML training. See the official tutorials for [running the
+T2T Transformer for text on Cloud TPUs](https://cloud.google.com/tpu/docs/tutorials/transformer) and
+[Transformer for Speech Recognition](https://cloud.google.com/tpu/docs/tutorials/automated-speech-recognition).
+
+## Other models on TPU
+
+Many of Tensor2Tensor's models work on TPU.
+
+You can provision a VM and TPU with `ctpu up`. Use the `t2t-trainer` command
+on the VM as usual with the additional flags `--use_tpu` and
+`--cloud_tpu_name=$TPU_NAME`.
+
+Note that because the `TPUEstimator` does not catch the `OutOfRangeError`
+during evaluation, you should ensure that `--eval_steps` is small enough to
+not exhaust the evaluation data.
+
+A non-exhaustive list of T2T models that work on TPU:
+
+* Image generation: `imagetransformer` with `imagetransformer_base_tpu` (or
+  `imagetransformer_tiny_tpu`)
+* Super-resolution: `img2img_transformer` with `img2img_transformer_base_tpu`
+  (or `img2img_transformer_tiny_tpu`)
+* `resnet` with `resnet_50` (or `resnet_18` or `resnet_34`)
+* `revnet` with `revnet_104` (or `revnet_38_cifar`)
+* `shake_shake` with `shakeshake_tpu` (or `shakeshake_small`)
+
+## Example invocation
+
+Use `ctpu up` to bring up the VM and TPU machines; once the machines are ready
+it will SSH you into the VM and you can run the following:
+
+```
+# DATA_DIR and OUT_DIR should be GCS buckets
+# TPU_NAME should have been set automatically by the ctpu tool
+
+t2t-trainer \
+  --model=shake_shake \
+  --hparams_set=shakeshake_tpu \
+  --problem=image_cifar10 \
+  --train_steps=180000 \
+  --eval_steps=9 \
+  --local_eval_frequency=100 \
+  --data_dir=$DATA_DIR \
+  --output_dir=$OUT_DIR \
+  --use_tpu \
+  --cloud_tpu_name=$TPU_NAME
+```
diff --git a/docs/distributed_training.md b/docs/distributed_training.md
new file mode 100644
index 000000000..f59974623
--- /dev/null
+++ b/docs/distributed_training.md
@@ -0,0 +1,220 @@
+# Distributed Training
+
+The `t2t-trainer` supports both synchronous and asynchronous distributed
+training.
+
+Note that it is almost always more efficient to train on a single machine with
+multiple GPUs/TPUs. Async training is less stable than sync training, and sync
+training is much faster on 1 machine than on multiple. For these reasons, we
+almost always train on single machines with multiple GPUs/TPUs.
+
+T2T uses TensorFlow Estimators and so distributed training is configured with
+the `TF_CONFIG` environment variable that is read by the
+[RunConfig](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/estimator/run_config.py)
+along with a set of flags that T2T uses to distribute the computation.
+
+## Shared output directory
+
+When using multiple machines, it is necessary that all nodes use the same
+`--output_dir`, which means that it should be set to a Google Cloud Storage
+bucket (`gs://...`) or a directory on a shared network filesystem.
+
+## Utility to produce `TF_CONFIG` and flags
+
+[`t2t-make-tf-configs`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/bin/t2t-make-tf-configs)
+generates the `TF_CONFIG` json strings and the necessary command-line flags for
+the jobs.
+
+Given a set of master and parameter server addresses, the script outputs, for
+each job, a line with the `TF_CONFIG` environment variable and the command-line
+flags necessary for distributed training. For each job, you should invoke the
+`t2t-trainer` with the `TF_CONFIG` value and flags that are output.
+
+## Eval jobs
+
+Eval jobs should set the following flags and do not need the `TF_CONFIG`
+environment variable to be set as the eval jobs run locally and do not
+communicate to the other jobs (the eval jobs read the model checkpoints that the
+trainer writes out):
+
+- `--schedule=continuous_eval_on_train_data` or
+  `--schedule=continuous_eval` (for dev data)
+- `--worker_job='/job:localhost'`
+- `--output_dir=$TRAIN_DIR`
+
+**Note that evaluation does not work distributed.** That is, distributed jobs
+should always use `--schedule=train`.
+
+## Examples
+
+### Sync training across multiple workers
+
+In this scenario, you wish to do synchronous training across multiple workers.
+Note that it is easier to simply use 1 worker with multiple GPUs and set
+`--worker_gpu=8`, but there may be cases where you may want to have multiple
+machines.
+
+You will need 1 `ip:port` for the master and then 1 `ip:port` for each worker.
+
+For this example we'll use 2 workers and these addresses:
+
+```
+# Master
+10.0.0.1:5555
+
+# Worker 1
+10.0.0.2:5555
+
+# Worker 2
+10.0.0.3:5555
+```
+
+Next we generate the `TF_CONFIG` and command-line-flags for each job.
+
+```
+$ t2t-make-tf-configs --masters='10.0.0.1:5555' --ps='10.0.0.2:5555,10.0.0.3:5555'
+Assuming SYNC distributed training with a single master and 2 workers
+'{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 0, "type": "master"}}'      --master=grpc://10.0.0.1:5555 --ps_replicas=2 --worker_replicas=1 --worker_gpu=0 --worker_id=0 --ps_gpu=1 --sync --schedule=train --worker_job='/job:master'
+'{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 0, "type": "ps"}}'  --schedule=run_std_server
+'{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 1, "type": "ps"}}'  --schedule=run_std_server
+```
+
+The output here is 1 line per job. Each line contains the `TF_CONFIG` to set
+for that job as well as the command-line flags to set for that job.
+
+It is a bit confusing that the workers are being passed to the `--ps` flag, but
+this is correct. When running in `--sync` mode, the `ps` are actually the
+workers. You can see in the next example below that when `--sync=False`, i.e.
+async mode, that the `ps` are in fact being used as parameter servers.
+
+Here's how we would start each job on their respective machines (the
+commands below assume that you're ssh'd into that job's machine):
+
+**Master**:
+
+```
+$ export TF_CONFIG='{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 0, "type": "master"}}'
+$ t2t-trainer \
+    --master=grpc://10.0.0.1:5555 \
+    --ps_replicas=2 \
+    --worker_replicas=1 \
+    --worker_gpu=0 \
+    --worker_id=0 \
+    --ps_gpu=1 \
+    --sync \
+    --schedule=train \
+    --worker_job='/job:master' \
+    --model=transformer \
+    --hparams_set=transformer_base \
+    --problem=translate_ende_wmt32k
+```
+
+**Worker 1**:
+
+```
+$ export TF_CONFIG='{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 0, "type": "ps"}}'
+$ t2t-trainer --schedule=run_std_server
+```
+
+**Worker 2**:
+
+```
+$ export TF_CONFIG='{"cluster": {"master": ["10.0.0.1:5555"], "ps": ["10.0.0.2:5555", "10.0.0.3:5555"]}, "environment": "cloud", "task": {"index": 1, "type": "ps"}}'
+$ t2t-trainer --schedule=run_std_server
+```
+
+Note that if you have more than 1 GPU on each worker machine, make sure to
+modify the `--ps_gpu` passed to the master.
+
+### Async training across multiple workers
+
+In this scenario, you wish to do asynchronous training across multiple workers
+with 1+ shared parameter servers.
+
+Note that async training is usually less stable than sync training and for that
+reason we almost always prefer sync training, but there may be cases where you
+want to do async distributed training.
+
+For this example we'll use 2 workers and 2 parameter servers:
+
+```
+# Worker 1
+10.0.0.1:5555
+
+# Worker 2
+10.0.0.2:5555
+
+# PS 1
+10.0.0.3:5555
+
+# PS 2
+10.0.0.4:5555
+```
+
+Next we generate the `TF_CONFIG` and command-line-flags for each job.
+
+```
+$ t2t-make-tf-configs --masters='10.0.0.1:5555,10.0.0.2:5555' --ps='10.0.0.3:5555,10.0.0.4:5555'
+Assuming ASYNC distributed training with 2 workers and 2 parameter servers
+'{"task": {"index": 0, "type": "chief"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}' --master=grpc://10.0.0.1:5555 --ps_replicas=2 --worker_replicas=2 --worker_gpu=1 --worker_id=0 --ps_gpu=0  --schedule=train --worker_job='/job:chief'
+'{"task": {"index": 0, "type": "worker"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'        --master=grpc://10.0.0.2:5555 --ps_replicas=2 --worker_replicas=2 --worker_gpu=1 --worker_id=1 --ps_gpu=0 --schedule=train --worker_job='/job:worker'
+'{"task": {"index": 0, "type": "ps"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'    --schedule=run_std_server
+'{"task": {"index": 1, "type": "ps"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'    --schedule=run_std_server
+```
+
+Here's how we would start each job on their respective machines (the
+commands below assume that you're ssh'd into that job's machine):
+
+**Worker 1**:
+
+```
+$ export TF_CONFIG='{"task": {"index": 0, "type": "chief"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'
+$ t2t-trainer \
+    --master=grpc://10.0.0.1:5555 \
+    --ps_replicas=2 \
+    --worker_replicas=2 \
+    --worker_gpu=1 \
+    --worker_id=0 \
+    --ps_gpu=0 \
+    --schedule=train \
+    --worker_job='/job:chief' \
+    --model=transformer \
+    --hparams_set=transformer_base \
+    --problem=translate_ende_wmt32k
+```
+
+**Worker 2**:
+
+```
+$ export TF_CONFIG='{"task": {"index": 0, "type": "worker"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'
+$ t2t-trainer \
+    --master=grpc://10.0.0.2:5555 \
+    --ps_replicas=2 \
+    --worker_replicas=2 \
+    --worker_gpu=1 \
+    --worker_id=1 \
+    --ps_gpu=0 \
+    --schedule=train \
+    --worker_job='/job:worker' \
+    --model=transformer \
+    --hparams_set=transformer_base \
+    --problem=translate_ende_wmt32k
+```
+
+**PS 1**:
+
+```
+$ export TF_CONFIG='{"task": {"index": 0, "type": "ps"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'
+$ t2t-trainer --schedule=run_std_server
+```
+
+**PS 2**:
+
+```
+$ export TF_CONFIG='{"task": {"index": 1, "type": "ps"}, "cluster": {"chief": ["10.0.0.1:5555"], "ps": ["10.0.0.3:5555", "10.0.0.4:5555"], "worker": ["10.0.0.2:5555"]}, "environment": "cloud"}'
+$ t2t-trainer --schedule=run_std_server
+```
+
+Increase `--worker_gpu` on each of the workers if you have multiple GPUs. If the
+parameter servers are also using GPUs, set `--ps_gpu` to the number of GPUs on
+the parameter servers.
diff --git a/docs/index.md b/docs/index.md
new file mode 100644
index 000000000..26298a9d2
--- /dev/null
+++ b/docs/index.md
@@ -0,0 +1,127 @@
+# Tensor2Tensor Documentation
+
+[![PyPI
+version](https://badge.fury.io/py/tensor2tensor.svg)](https://badge.fury.io/py/tensor2tensor)
+[![GitHub
+Issues](https://img.shields.io/github/issues/tensorflow/tensor2tensor.svg)](https://github.com/tensorflow/tensor2tensor/issues)
+[![Contributions
+welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CONTRIBUTING.md)
+[![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
+[![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
+
+[Tensor2Tensor](https://github.com/tensorflow/tensor2tensor), or
+[T2T](https://github.com/tensorflow/tensor2tensor) for short, is a library
+of deep learning models and datasets designed to make deep learning more
+accessible and [accelerate ML
+research](https://research.googleblog.com/2017/06/accelerating-deep-learning-research.html).
+
+
+## Introduction
+
+* [Walkthrough](walkthrough.md): Install and run.
+* [IPython notebook](https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/hello_t2t.ipynb): Get a hands-on experience.
+
+## Basics
+
+* [Overview](overview.md): How all parts of T2T code are connected.
+* [New Problem](new_problem.md): Train T2T models on your data.
+* [New Model](new_model.md): Create your own T2T model.
+
+## Training in the cloud
+
+* [Training on Google Cloud ML](cloud_mlengine.md)
+* [Training on Google Cloud TPUs](cloud_tpu.md)
+* [Distributed Training](distributed_training.md)
+
+## Solving your task
+
+Below we list a number of tasks that can be solved with T2T when
+you train the appropriate model on the appropriate problem.
+We give the problem and model below and we suggest a setting of
+hyperparameters that we know works well in our setup. We usually
+run either on Cloud TPUs or on 8-GPU machines; you might need
+to modify the hyperparameters if you run on a different setup.
+
+### Image Classification
+
+For image classification, we have a number of standard data-sets:
+* ImageNet (a large data-set): `--problem=image_imagenet`, or one
+   of the re-scaled versions (`image_imagenet224`, `image_imagenet64`,
+   `image_imagenet32`)
+* CIFAR-10: `--problem=image_cifar10` (or
+    `--problem=image_cifar10_plain` to turn off data augmentation)
+* CIFAR-100: `--problem=image_cifar100`
+* MNIST: `--problem=image_mnist`
+
+For ImageNet, we suggest to use the ResNet or Xception, i.e.,
+use `--model=resnet --hparams_set=resnet_50` or
+`--model=xception --hparams_set=xception_base`.
+Resnet should get to above 76% top-1 accuracy on ImageNet.
+
+For CIFAR and MNIST, we suggest to try the shake-shake model:
+`--model=shake_shake --hparams_set=shakeshake_big`.
+This setting trained for `--train_steps=700000` should yield
+close to 97% accuracy on CIFAR-10.
+
+### Language Modeling
+
+For language modeling, we have these data-sets in T2T:
+* PTB (a small data-set): `--problem=languagemodel_ptb10k` for
+    word-level modeling and `--problem=languagemodel_ptb_characters`
+    for character-level modeling.
+* LM1B (a billion-word corpus): `--problem=languagemodel_lm1b32k` for
+    subword-level modeling and `--problem=languagemodel_lm1b_characters`
+    for character-level modeling.
+
+We suggest to start with `--model=transformer` on this task and use
+`--hparams_set=transformer_small` for PTB and
+`--hparams_set=transformer_base` for LM1B.
+
+### Sentiment Analysis
+
+For the task of recognizing the sentiment of a sentence, use
+* the IMDB data-set: `--problem=sentiment_imdb`
+
+We suggest to use `--model=transformer_encoder` here and since it is
+a small data-set, try `--hparams_set=transformer_tiny` and train for
+few steps (e.g., `--train_steps=2000`).
+
+### Speech Recognition
+
+For speech-to-text, we have these data-sets in T2T:
+* Librispeech (English speech to text): `--problem=librispeech` for
+    the whole set and `--problem=librispeech_clean` for a smaller
+    but nicely filtered part.
+
+### Summarization
+
+For summarizing longer text into shorter one we have these data-sets:
+* CNN/DailyMail articles summarized into a few sentences:
+  `--problem=summarize_cnn_dailymail32k`
+
+We suggest to use `--model=transformer` and
+`--hparams_set=transformer_prepend` for this task.
+This yields good ROUGE scores.
+
+### Translation
+
+There are a number of translation data-sets in T2T:
+* English-German: `--problem=translate_ende_wmt32k`
+* English-French: `--problem=translate_enfr_wmt32k`
+* English-Czech: `--problem=translate_encs_wmt32k`
+* English-Chinese: `--problem=translate_enzh_wmt32k`
+* English-Vietnamese: `--problem=translate_envi_iwslt32k`
+* English-Spanish: `--problem=translate_enes_wmt32k`
+
+You can get translations in the other direction by appending `_rev` to
+the problem name, e.g., for German-English use
+`--problem=translate_ende_wmt32k_rev`.
+
+For all translation problems, we suggest to try the Transformer model:
+`--model=transformer`. At first it is best to try the base setting,
+`--hparams_set=transformer_base`. When trained on 8 GPUs for 300K steps
+this should reach a BLEU score of about 28 on the English-German data-set,
+which is close to state-of-the art. If training on a single GPU, try the
+`--hparams_set=transformer_base_single_gpu` setting. For very good results
+or larger data-sets (e.g., for English-French), try the big model
+with `--hparams_set=transformer_big`.
diff --git a/docs/multi_problem.md b/docs/multi_problem.md
new file mode 100644
index 000000000..d4e37d09d
--- /dev/null
+++ b/docs/multi_problem.md
@@ -0,0 +1,188 @@
+# Multi-problem training
+
+Multi-problem training is possible by defining [MultiProblem](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/multi_problem.py) sub-classes that specify a list of [Problem](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/problem.py) objects to include in training. In some cases, multi-problem training can be used to improve performance compared to training on individual problems.
+
+In the following sections we'll discuss MultiProblem from a usage perspective followed by that of someone wishing to build upon it.
+
+Please note the [T2T Walkthrough](https://github.com/tensorflow/tensor2tensor/blob/master/docs/walkthrough.md) documentation is a good place to start to understand the variety of component concepts we'll build on here.
+
+## Usage
+
+### Problem definition and datagen
+
+In this discussion we'll consider the following (large) multi-problem that includes ten different sub-problems. These include:
+
+1. A [language modeling](https://en.wikipedia.org/wiki/Language_model) [problem](https://github.com/tensorflow/tensor2tensor/blob/0dff89d64c3406d42717280cb9135a5ce7af793c/tensor2tensor/data_generators/wiki_lm.py#L223) operating on a corpus of German, English, French, and Romanian language wikipedia articles.
+2. Multiple compatible pairwise language translation problems (En -> De, En -> Fr, En -> Ro, De -> En, Fr -> En, Ro -> En)
+3. A compatible [version](https://github.com/tensorflow/tensor2tensor/blob/ef12bee72270b322165d073c39a650a189de39aa/tensor2tensor/data_generators/cnn_dailymail.py#L267) of the combined CNN/DailyMail news article summarization problem.
+4. A compatible [version](https://github.com/tensorflow/tensor2tensor/blob/ef12bee72270b322165d073c39a650a189de39aa/tensor2tensor/data_generators/multinli.py#L155) of the [MultiNLI](https://www.nyu.edu/projects/bowman/multinli/) textual entailment classification problem.
+5. A compatible [version](https://github.com/tensorflow/tensor2tensor/blob/1de13dbebccb415d89b0658e18a57e9607bafd32/tensor2tensor/data_generators/squad.py#L126) of the [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/) question/answer problem.
+
+```python
+
+@registry.register_problem
+class LanguagemodelMultiWikiTranslate(multi_problem.MultiProblem):
+  """Wiki multi-lingual LM and multiple translations."""
+
+  def __init__(self, was_reversed=False, was_copy=False):
+    super(LanguagemodelMultiWikiTranslate, self).__init__(
+        was_reversed, was_copy)
+    self.task_list.append(wiki_lm.LanguagemodelDeEnFrRoWiki64k())
+    self.task_list.append(translate_ende.TranslateEndeWmtMulti64k())
+    self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k())
+    self.task_list.append(translate_enro.TranslateEnroWmtMultiTiny64k())
+    self.task_list.append(translate_ende.TranslateEndeWmtMulti64k(
+        was_reversed=True))
+    self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k(
+        was_reversed=True))
+    self.task_list.append(translate_enro.TranslateEnroWmtMultiTiny64k(
+        was_reversed=True))
+    self.task_list.append(
+        cnn_dailymail.SummarizeCnnDailymailWikiLMMultiVocab64k())
+    self.task_list.append(multinli.MultiNLIWikiLMMultiVocab64k())
+    self.task_list.append(squad.SquadConcatMulti64k())
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.SUBWORD
+
+```
+
+The word "compatible" was used a lot above! That's because each of these problems have been modified to use the vocabulary produced by the Wikipedia-based language modeling problem, e.g. the following
+
+```python
+@registry.register_problem
+class SummarizeCnnDailymailWikiLMMultiVocab64k(SummarizeCnnDailymail32k):
+  """Summarize CNN and Daily Mail articles using multi-lingual 64k vocab."""
+
+  @property
+  def vocab_filename(self):
+    return wiki_lm.LanguagemodelDeEnFrRoWiki64k().vocab_filename
+```
+
+**Important note:** It's easy to miss the key point that, as implemented currently, the first task in the task list must be a language modelling problem and each included task must be modified to use the resulting vocabulary.
+
+With a properly defined and registered multi-problem we can now run datagen as follows:
+
+```bash
+
+t2t-datagen --problem=languagemodel_multi_wiki_translate
+
+```
+
+This will take approximately the following amount of space (and several hours):
+
+```bash
+(t2t) username@instance-2:~$ du -sh /tmp
+99G     /tmp
+(t2t) username@instance-2:~$ du -sh /tmp/t2t_datagen
+81G     /tmp/t2t_datagen
+```
+
+### Training
+
+Next we're ready to try training a model on this MultiProblem. Note that by not specifying `--data_dir` above TFExample's were by default generated into /tmp so that's what we'll explicitly provide here.
+
+```bash
+
+t2t-trainer --problem=languagemodel_multi_wiki_translate \
+    --model=transformer \
+    --hparams_set=transformer_tall_pretrain_lm_tpu_adafactor_large \
+    --output_dir ~/t2t_train/transformer_multi_2jan19 \
+    --data_dir=/tmp \
+    --train_steps=1 \
+    --eval_steps=1
+
+```
+
+The `hparams_set` parameter we provided above was [transformer_tall_pretrain_lm_tpu_adafactor_large](https://github.com/tensorflow/tensor2tensor/blob/08e83030acf3ef13d15ad6eaefaa0a67fb20b59d/tensor2tensor/models/transformer.py#L1721), also provided below:
+
+```python
+
+@registry.register_hparams
+def transformer_tall_pretrain_lm_tpu_adafactor_large():
+  """Hparams for transformer on LM pretraining on TPU, large model."""
+  hparams = transformer_tall_pretrain_lm_tpu_adafactor()
+  hparams.hidden_size = 1024
+  hparams.num_heads = 16
+  hparams.filter_size = 32768  # max fitting in 16G memory is 49152, batch 2
+  hparams.batch_size = 4
+  hparams.multiproblem_mixing_schedule = "constant"
+  # Task order: lm/en-de/en-fr/en-ro/de-en/fr-en/ro-en/cnndm/mnli/squad.
+  hparams.multiproblem_per_task_threshold = "320,80,160,2,80,160,2,20,5,5"
+  return hparams
+
+```
+
+Here it's worth noting a couple things, one that we have specified a `multi_problem_mixing_schedule` (which is required), consumed by [MultiProblem.mix_data](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/multi_problem.py#L280). When set to "constant" the strategy for sampling examples is not a function of step and is proportional only to the per-task "thresholds" which are by default equal (sample examples from each problem with equal probability).
+
+But notice we have also specified the (non-required) `multiproblem_per_task_threshold` parameter, also consumed by mix_data, and specifically used by [sample_task](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/multi_problem.py#L340) which defines non-uniform thresholds to inform a weighted random sampling. E.g. for two problems with weights 1 and 9 the first would be sampled 1/10 of the time and the other 9/10.
+
+### Inference
+
+You can try translating from English to German using a model previously trained on `LanguagemodelMultiWikiTranslate` (the one shown above) ([gs://tensor2tensor-checkpoints/transformer_multi_2jan19/](https://console.cloud.google.com/storage/browser/tensor2tensor-checkpoints/transformer_multi_2jan19/)). Just copy the checkpoint down to a local directory such as the one given via `--output_dir` below:
+
+```bash
+
+t2t-decoder --problem=languagemodel_multi_wiki_translate \
+    --model=transformer \
+    --hparams_set=transformer_tall_pretrain_lm_tpu_adafactor_large \
+    --decode_hparams='batch_size=1,multiproblem_task_id=64510' \
+    --hparams="" \
+    --output_dir=~/t2t_train/transformer_multi_2jan19 \
+    --decode_from_file ~/newstest2014.en \
+    --data_dir=~/t2t_train/transformer_multi_2jan19
+
+```
+
+Here we'll point `--data_dir` to the checkpoint directory which includes the vocab file `vocab.languagemodel_de_en_fr_ro_wiki64k.64000.subwords`; typically data_dir would point to the directory containing your TFRecord example dataset(s).
+
+The file passed to `--decode_from_file` is simply a file with one sentence to translate on each line (in its original form, not post-vocabulary-encoded).
+
+A key requirement for multi-problem inference is that we specify the ID of the problem for which we want to perform inference. But wait, why is the task ID 64510? We can see from the code for [`MultiProblem.update_task_ids`](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/multi_problem.py#L386) that TID's have a place at the end of the vocabulary.
+
+```python
+
+class MultiProblem(problem.Problem):
+  """MultiProblem base class."""
+
+  ...
+
+  def update_task_ids(self, encoder_vocab_size):
+    """Generate task_ids for each problem.
+    These ids correspond to the index of the task in the task_list.
+    Args:
+      encoder_vocab_size: the size of the vocab which is used to compute
+        the index offset.
+    """
+    for idx, task in enumerate(self.task_list):
+      task.set_task_id(idx + encoder_vocab_size)
+      tf.logging.info("Task %d (%s) has id %d." %
+                      (idx, task.name, task.task_id))
+
+```
+
+We can look up the task_id that is assigned to each task we may want to use for inference by instantiating the MultiProblem subclass and obtaining the value, in this case via the following:
+
+```python
+
+task_index = 1 # The second task in the list is En -> De
+LanguagemodelMultiWikiTranslate().task_list[task_index].task_id
+
+```
+
+For me running the `t2t-decode` command provided above gave the following output:
+
+```bash
+...
+
+INFO:tensorflow:Running local_init_op.
+INFO:tensorflow:Done running local_init_op.
+INFO:tensorflow:Inference results INPUT: hello world was the news of the day
+INFO:tensorflow:Inference results OUTPUT: Hallo Welt war die Nachricht des Tages
+INFO:tensorflow:Elapsed Time: 37.15079
+INFO:tensorflow:Averaged Single Token Generation Time: 3.3009222 (time 36.3101439 count 11)
+
+...
+
+```
diff --git a/docs/new_model.md b/docs/new_model.md
new file mode 100644
index 000000000..861f83bb0
--- /dev/null
+++ b/docs/new_model.md
@@ -0,0 +1,114 @@
+# T2T: Create Your Own Model
+
+[![PyPI
+version](https://badge.fury.io/py/tensor2tensor.svg)](https://badge.fury.io/py/tensor2tensor)
+[![GitHub
+Issues](https://img.shields.io/github/issues/tensorflow/tensor2tensor.svg)](https://github.com/tensorflow/tensor2tensor/issues)
+[![Contributions
+welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](../CONTRIBUTING.md)
+[![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
+[![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
+
+Here we show how to create your own model in T2T.
+
+## The T2TModel class - abstract base class for models
+
+  `T2TModel` has three typical usages:
+
+  1. Estimator: The method `make_estimator_model_fn` builds a `model_fn` for
+     the tf.Estimator workflow of training, evaluation, and prediction.
+     It performs the method `call`, which performs the core computation,
+     followed by `estimator_spec_train`, `estimator_spec_eval`, or
+     `estimator_spec_predict` depending on the tf.Estimator mode.
+  2. Layer: The method `call` enables `T2TModel` to be used a callable by
+     itself. It calls the following methods:
+
+     * `bottom`, which transforms features according to `problem_hparams`' input
+       and target `Modality`s;
+     * `body`, which takes features and performs the core model computation to
+        return output and any auxiliary loss terms;
+     * `top`, which takes features and the body output, and transforms them
+       according to `problem_hparams`' input and target `Modality`s to return
+       the final logits;
+     * `loss`, which takes the logits, forms any missing training loss, and sums
+       all loss terms.
+  3. Inference: The method `infer` enables `T2TModel` to make sequence
+     predictions by itself.
+
+
+## Creating your own model
+
+1. Create class that extends T2TModel
+    in this example it will be a copy of existing basic fully connected network:
+
+```python
+    from tensor2tensor.utils import t2t_model
+
+    class MyFC(t2t_model.T2TModel):
+        pass
+```
+
+
+2. Implement body method:
+
+```python
+    class MyFC(t2t_model.T2TModel):
+      def body(self, features):
+        hparams = self.hparams
+        x = features["inputs"]
+        shape = common_layers.shape_list(x)
+        x = tf.reshape(x, [-1, shape[1] * shape[2] * shape[3]])  # Flatten input as in T2T they are all 4D vectors
+        for i in range(hparams.num_hidden_layers): # create layers
+          x = tf.layers.dense(x, hparams.hidden_size, name="layer_%d" % i)
+          x = tf.nn.dropout(x, keep_prob=1.0 - hparams.dropout)
+          x = tf.nn.relu(x)
+        return tf.expand_dims(tf.expand_dims(x, axis=1), axis=1)  # 4D For T2T.
+```
+
+
+Method Signature:
+
+  * Args:
+      * features: dict of str to Tensor, where each Tensor has shape [batch_size,
+     ..., hidden_size]. It typically contains keys `inputs` and `targets`.
+
+  * Returns one of:
+    * output: Tensor of pre-logit activations with shape [batch_size, ...,
+           hidden_size].
+    * losses: Either single loss as a scalar, a list, a Tensor (to be averaged),
+           or a dictionary of losses. If losses is a dictionary with the key
+           "training", losses["training"] is considered the final training
+           loss and output is considered logits; self.top and self.loss will
+           be skipped.
+
+3. Register your model
+
+```python
+    from tensor2tensor.utils import registry
+
+    @registry.register_model
+    class MyFC(t2t_model.T2TModel):
+       # ...
+```
+
+
+3. Use it with t2t tools as any other model
+
+    Have in mind that names are translated from camel case to snake_case `MyFC` -> `my_fc`
+    and that you need to point t2t to directory containing your model with `t2t_usr_dir` switch. 
+    For example if you want to train model on gcloud with 1 GPU worker on IMDB sentiment task you can run your model
+    by executing following command from your model class directory. 
+
+```bash
+    t2t-trainer \
+      --model=my_fc \
+      --t2t_usr_dir=.
+      --cloud_mlengine --worker_gpu=1 \
+      --generate_data \
+      --data_dir='gs://data' \
+      --output_dir='gs://out' \
+      --problem=sentiment_imdb \
+      --hparams_set=basic_fc_small \
+      --train_steps=10000 \
+      --eval_steps=10 \
+```
diff --git a/docs/new_problem.md b/docs/new_problem.md
new file mode 100644
index 000000000..13f012b79
--- /dev/null
+++ b/docs/new_problem.md
@@ -0,0 +1,243 @@
+# T2T: Train on Your Own Data
+
+[![PyPI
+version](https://badge.fury.io/py/tensor2tensor.svg)](https://badge.fury.io/py/tensor2tensor)
+[![GitHub
+Issues](https://img.shields.io/github/issues/tensorflow/tensor2tensor.svg)](https://github.com/tensorflow/tensor2tensor/issues)
+[![Contributions
+welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CONTRIBUTING.md)
+[![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
+[![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
+
+Another good overview of this part together with training is given in
+[The Cloud ML Poetry Blog
+Post](https://cloud.google.com/blog/big-data/2018/02/cloud-poetry-training-and-hyperparameter-tuning-custom-text-models-on-cloud-ml-engine)
+
+Let's add a new dataset together and train the
+[Transformer](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/models/transformer.py)
+model on it. We'll give the model a line of poetry, and it will learn to
+generate the next line.
+
+# Defining the `Problem`
+
+For each problem we want to tackle we create a new subclass of `Problem` and
+register it. Let's call our problem `PoetryLines`.
+
+Since many text-to-text problems share similar methods, there's already a class
+called
+[`Text2TextProblem`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/text_problems.py)
+that extends the base problem class
+[`Problem`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/problem.py)
+and makes it easy to add text-to-text problems.
+
+In that same file, there are other base classes that make it easy to add text
+classification tasks (`Text2ClassProblem`) and language modeling tasks
+(`Text2SelfProblem`).
+
+For our problem, let's create the file `poetry_lines.py` and add our new
+problem, `PoetryLines`, which extends `Text2TextProblem` and register it so that
+it is accessible by command-line flag.
+
+Here's the Problem in full. We'll go step by step through it.
+
+```python
+import re
+
+from gutenberg import acquire
+from gutenberg import cleanup
+
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+
+@registry.register_problem
+class PoetryLines(text_problems.Text2TextProblem):
+  """Predict next line of poetry from the last line. From Gutenberg texts."""
+
+  @property
+  def approx_vocab_size(self):
+    return 2**13  # ~8k
+
+  @property
+  def is_generate_per_split(self):
+    # generate_data will shard the data into TRAIN and EVAL for us.
+    return False
+
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each."""
+    # 10% evaluation data
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 9,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    del data_dir
+    del tmp_dir
+    del dataset_split
+
+
+    books = [
+        # bookid, skip N lines
+        (19221, 223),
+        (15553, 522),
+    ]
+
+    for (book_id, toskip) in books:
+      text = cleanup.strip_headers(acquire.load_etext(book_id)).strip()
+      lines = text.split("\n")[toskip:]
+      prev_line = None
+      ex_count = 0
+      for line in lines:
+        # Any line that is all upper case is a title or author name
+        if not line or line.upper() == line:
+          prev_line = None
+          continue
+
+        line = re.sub("[^a-z]+", " ", line.strip().lower())
+        if prev_line and line:
+          yield {
+              "inputs": prev_line,
+              "targets": line,
+          }
+          ex_count += 1
+        prev_line = line
+```
+
+## Vocabulary specification
+
+The text generated is encoded with a vocabulary for training. By default, it is
+a `SubwordTextEncoder` that is built with an approximate vocab size specified by
+the user. It's fully invertible (no out-of-vocab tokens) with a fixed-size vocab
+which makes it ideal for text problems.
+
+You can also choose to use a character-level encoder or a token encoder where
+you provide the vocab file yourself. See `Text2TextProblem.vocab_type`.
+
+Here we specify that we're going to have a vocabulary with approximately 8,000
+subwords.
+
+```python
+  @property
+  def approx_vocab_size(self):
+    return 2**13  # ~8k
+```
+
+## Splitting data between Train and Eval
+
+By setting `is_generate_per_split=False`, the `generate_samples` method will
+only be called once and the data will automatically be split across training and
+evaluation data for us. This is useful because for our dataset we don't have
+pre-existing "training" and "evaluation" sets. If we did, we'd set
+`is_generate_per_split=True` so that `generate_samples` was called once per data
+split.
+
+The `dataset_splits` method determines the fraction that goes to each split. The
+training data will be generated into 9 files and the evaluation data into 1.
+90% of the data will be for training. 10% of the data will be for evaluation.
+
+```python
+  @property
+  def is_generate_per_split(self):
+    # generate_data will shard the data into TRAIN and EVAL for us.
+    return False
+
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each."""
+    # 10% evaluation data
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 9,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+```
+
+## Generating samples
+
+`generate_samples` is the bulk of the code where we actually produce
+dictionaries of poetry line pairs ("inputs" and "targets").
+
+Some problems might require downloading, which can be done into `tmp_dir`. Some
+problems may use their own token vocabulary file, in which case it can be copied
+into `data_dir` before yielding samples.
+
+Here we iterate through the lines of a couple books of poetry and produce pairs
+of lines for the model to train against.
+
+```python
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    del data_dir
+    del tmp_dir
+    del dataset_split
+
+    books = [
+        # bookid, skip N lines
+        (19221, 223),
+        (15553, 522),
+    ]
+
+    for (book_id, toskip) in books:
+      text = cleanup.strip_headers(acquire.load_etext(book_id)).strip()
+      lines = text.split("\n")[toskip:]
+      prev_line = None
+      ex_count = 0
+      for line in lines:
+        # Any line that is all upper case is a title or author name
+        if not line or line.upper() == line:
+          prev_line = None
+          continue
+
+        line = re.sub("[^a-z]+", " ", line.strip().lower())
+        if prev_line and line:
+          yield {
+              "inputs": prev_line,
+              "targets": line,
+          }
+          ex_count += 1
+        prev_line = line
+```
+
+That's all for the problem specification! We're ready to generate the data.
+
+# Run data generation
+
+You can generate data for your problem with `t2t-datagen` and the
+`--t2t_usr_dir` flag, which points to the directory containing an `__init__.py`
+file that imports the `poetry_lines` file we just wrote. See setup below.
+
+```bash
+USR_DIR=...
+PROBLEM=poetry_lines
+DATA_DIR=$HOME/t2t_data
+TMP_DIR=/tmp/t2t_datagen
+mkdir -p $DATA_DIR $TMP_DIR
+
+t2t-datagen \
+  --t2t_usr_dir=$USR_DIR \
+  --data_dir=$DATA_DIR \
+  --tmp_dir=$TMP_DIR \
+  --problem=$PROBLEM
+```
+
+`PROBLEM` is the name of the class that was registered with
+`@registry.register_problem`, but converted from `CamelCase` to `snake_case`.
+
+`USR_DIR` is a directory with the `poetry_lines.py` file and an
+`__init__.py` file that imports it (`from . import poetry_lines`).
+
+If you plan to contribute problems to the tensor2tensor repository, you can
+clone the repository and install it in developer mode with `pip install -e .`.
+
+# Train!
+
+You can train exactly as you do in the [walkthrough](walkthrough.md) with flags
+`--problem=poetry_lines` and `--t2t_usr_dir=$USR_DIR`.
+
+All done. Let us know what amazing poetry your model writes!
diff --git a/docs/overview.md b/docs/overview.md
new file mode 100644
index 000000000..9ea87bc50
--- /dev/null
+++ b/docs/overview.md
@@ -0,0 +1,175 @@
+# T2T: Life of an Example
+
+[![PyPI
+version](https://badge.fury.io/py/tensor2tensor.svg)](https://badge.fury.io/py/tensor2tensor)
+[![GitHub
+Issues](https://img.shields.io/github/issues/tensorflow/tensor2tensor.svg)](https://github.com/tensorflow/tensor2tensor/issues)
+[![Contributions
+welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CONTRIBUTING.md)
+[![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/tensor2tensor/Lobby)
+[![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
+
+This doc explains how a training example flows through T2T, from data generation
+to training, evaluation, and decoding.
+
+Some key files and their functions:
+
+*   [`t2t_trainer.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/bin/t2t_trainer.py) and [`trainer_lib.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/trainer_lib.py):
+    Main entrypoint for training and evaluation.  Constructs and runs all the
+    main components of the system (the `Problem`, the `HParams`, the
+    `Estimator`, the `Experiment`, the `input_fn`s and `model_fn`).
+*   [`common_hparams.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/layers/common_hparams.py):
+    `basic_params1` serves as the base for all model hyperparameters. Registered
+    model hparams functions always start with this default set of
+    hyperparameters.
+*   [`problem.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/problem.py):
+    Every dataset in T2T subclasses `Problem`. `Problem.input_fn` is the
+    Estimator input function.
+*   [`t2t_model.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/t2t_model.py):
+    Every model in T2T subclasses `T2TModel`. `T2TModel.estimator_model_fn` is
+    the Estimator model function.
+
+## Data Generation
+
+The `t2t-datagen` binary is the entrypoint for data generation. It simply looks
+up the `Problem` specified by `--problem` and calls
+`Problem.generate_data(data_dir, tmp_dir)`.
+
+All `Problem`s are expected to generate 2 sharded `TFRecords` files - 1 for
+training and 1 for evaluation - with `tensorflow.Example` protocol buffers. The
+expected names of the files are given by `Problem.{training, dev}_filepaths`.
+Typically, the features in the `Example` will be `"inputs"` and `"targets"`;
+however, some tasks have a different on-disk representation that is converted to
+`"inputs"` and `"targets"` online in the input pipeline (e.g. image features are
+typically stored with features `"image/encoded"` and `"image/format"` and the
+decoding happens in the input pipeline).
+
+For tasks that require a vocabulary, this is also the point at which the
+vocabulary is generated and all examples are encoded.
+
+There are several utility functions in
+[`generator_utils`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/generator_utils.py)
+that are commonly used by `Problem`s to generate data. Several are highlighted
+below:
+
+*   `generate_dataset_and_shuffle`: given 2 generators, 1 for training and 1 for
+    eval, yielding dictionaries of `<feature name, list< int or float or
+    string >>`, will produce sharded and shuffled `TFRecords` files with
+    `tensorflow.Example` protos.
+*   `maybe_download`: downloads a file at a URL to the given directory and
+    filename (see `maybe_download_from_drive` if the URL points to Google
+    Drive).
+*   `get_or_generate_vocab_inner`: given a target vocabulary size and a
+    generator that yields lines or tokens from the dataset, will build a
+    `SubwordTextEncoder` along with a backing vocabulary file that can be used
+    to map input strings to lists of ids.
+    [`SubwordTextEncoder`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/text_encoder.py)
+    uses word pieces and its encoding is fully invertible.
+
+## Data Input Pipeline
+
+Once the data is produced on disk, training, evaluation, and inference (if
+decoding from the dataset) consume it by way of the T2T input pipeline, defined
+by `Problem.input_fn`.
+
+The entire input pipeline is implemented with the new `tf.data.Dataset` API.
+
+The input function has 2 main parts: first, reading and processing individual
+examples, which is done is `Problem.dataset`, and second, batching, which is
+done in `Problem.input_fn` after the call to `Problem.dataset`.
+
+`Problem` subclasses may override the entire `input_fn` or portions of it (e.g.
+`example_reading_spec` to indicate the names, types, and shapes of features on
+disk). Typically they only override portions.
+
+### Batching
+
+Problems that have fixed size features (e.g. image problems) can use
+`hp.batch_size` to set the batch size.
+
+Variable length Problems are bucketed by sequence length and then batched out of
+those buckets.  This significantly improves performance over a naive batching
+scheme for variable length sequences because each example in a batch must be
+padded to match the example with the maximum length in the batch.
+
+Controlling hparams:
+
+* `hp.batch_size`: the approximate total number of tokens in
+  the batch (i.e. long sequences will have smaller actual batch size and short
+  sequences will have a larger actual batch size in order to generally have an
+  equal number of tokens in the batch).
+* `hp.max_length`: For variable length features, sequences with length longer
+  than this will be dropped during training (and also during eval if
+  `hp.eval_drop_long_sequences` is `True`). If not set, the maximum length of
+  examples is set to `hp.batch_size`.
+* `hp.batch_size_multiplier`: multiplier for the maximum length
+* `hp.min_length_bucket`: example length for the smallest bucket (i.e. the
+  smallest bucket will bucket examples up to this length).
+* `hp.length_bucket_step`: controls how spaced out the length buckets are.
+
+## Building the Model
+
+At this point, the input features typically have `"inputs"` and `"targets"`,
+each of which is a batched 4-D Tensor (e.g. of shape `[batch_size,
+sequence_length, 1, 1]` for text input or `[batch_size, height, width, 3]` for
+image input).
+
+The Estimator model function is created by `T2TModel.estimator_model_fn`, which
+may be overridden in its entirety by subclasses if desired. Typically,
+subclasses only override `T2TModel.body`.
+
+The model function constructs a `T2TModel`, calls it, and then calls
+`T2TModel.{estimator_spec_train, estimator_spec_eval, estimator_spec_predict}`
+depending on the mode.
+
+A call of a `T2TModel` internally calls `bottom`, `body`, `top`, and `loss`, all
+of which can be overridden by subclasses (typically only `body` is).
+
+The default implementations of `bottom`, `top`, and `loss` depend on the
+`Modality` specified for the input and target features (e.g.
+`SymbolModality.bottom` embeds integer tokens and `SymbolModality.loss` is
+`softmax_cross_entropy`).
+
+## `Estimator` and `Experiment`
+
+The actual training loop and related services (checkpointing, summaries,
+continuous evaluation, etc.) are all handled by `Estimator` and `Experiment`
+objects. `t2t_trainer.py` is the main entrypoint and uses `trainer_lib.py`
+to construct the various components.
+
+## Decoding
+
+* [`t2t_decoder.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/bin/t2t-decoder)
+* [`decoding.py`](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/utils/decoding.py)
+
+## System Overview for Train/Eval
+
+See `t2t_trainer.py` and `trainer_lib.py`.
+
+* Create HParams
+* Create `RunConfig`, including `Parallelism` object (i.e. `data_parallelism`)
+* Create `Experiment`, including hooks
+* Create `Estimator`
+  * `T2TModel.estimator_model_fn`
+    * `model(features)`
+      * `model.model_fn`
+        * `model.bottom`
+        * `model.body`
+        * `model.top`
+        * `model.loss`
+    * [TRAIN] `model.estimator_spec_train`
+      * `train_op = model.optimize`
+    * [EVAL] `model.estimator_spec_eval`
+      * Create metrics
+* Create input functions
+  * `Problem.input_fn`
+    * `Problem.dataset`
+    * Batching
+* Create hooks
+* Run Experiment --schedule (e.g. `exp.continuous_train_and_eval()`)
+  * `estimator.train`
+    * `train_op = model_fn(input_fn(mode=TRAIN))`
+    * Run train op
+  * `estimator.evaluate`
+    * `metrics = model_fn(input_fn(mode=EVAL))`
+    * Accumulate metrics
diff --git a/docs/tutorials/asr_with_transformer.md b/docs/tutorials/asr_with_transformer.md
new file mode 100644
index 000000000..ee8aa0f53
--- /dev/null
+++ b/docs/tutorials/asr_with_transformer.md
@@ -0,0 +1,4 @@
+# Automated Speech Recognition with the Transformer model
+
+See the
+[official tutorial](https://cloud.google.com/tpu/docs/tutorials/automated-speech-recognition).
diff --git a/tensor2tensor/trax/rlax/__init__.py b/tensor2tensor/trax/rl/__init__.py
similarity index 100%
rename from tensor2tensor/trax/rlax/__init__.py
rename to tensor2tensor/trax/rl/__init__.py
diff --git a/tensor2tensor/trax/rlax/configs/atari.gin b/tensor2tensor/trax/rl/configs/atari.gin
similarity index 95%
rename from tensor2tensor/trax/rlax/configs/atari.gin
rename to tensor2tensor/trax/rl/configs/atari.gin
index 7734fd0d9..fc6fa7fec 100644
--- a/tensor2tensor/trax/rlax/configs/atari.gin
+++ b/tensor2tensor/trax/rl/configs/atari.gin
@@ -1,4 +1,4 @@
-import tensor2tensor.trax.rlax
+import tensor2tensor.trax.rl
 
 # Parameters for ppo.training_loop:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rlax/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin
similarity index 95%
rename from tensor2tensor/trax/rlax/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin
rename to tensor2tensor/trax/rl/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin
index 4ec0423de..78195a528 100644
--- a/tensor2tensor/trax/rlax/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin
@@ -1,8 +1,8 @@
 import tensor2tensor.trax.inputs
 import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.rlax
-import tensor2tensor.trax.rlax.envs
+import tensor2tensor.trax.rl
+import tensor2tensor.trax.rl.envs
 
 # Parameters for batch_fun:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rlax/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
similarity index 95%
rename from tensor2tensor/trax/rlax/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
rename to tensor2tensor/trax/rl/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
index af365de63..dc1155525 100644
--- a/tensor2tensor/trax/rlax/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
@@ -1,4 +1,4 @@
-import tensor2tensor.trax.rlax
+import tensor2tensor.trax.rl
 
 # Parameters for ppo.training_loop:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/online_tune_wide_resnet_cifar10.gin
similarity index 97%
rename from tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
rename to tensor2tensor/trax/rl/configs/online_tune_wide_resnet_cifar10.gin
index e6ffbea13..eb0ebd98c 100644
--- a/tensor2tensor/trax/rlax/configs/online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/online_tune_wide_resnet_cifar10.gin
@@ -1,8 +1,8 @@
 import tensor2tensor.trax.inputs
 import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.rlax
-import tensor2tensor.trax.rlax.envs
+import tensor2tensor.trax.rl
+import tensor2tensor.trax.rl.envs
 
 # Parameters for batch_fun:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rlax/envs/__init__.py b/tensor2tensor/trax/rl/envs/__init__.py
similarity index 78%
rename from tensor2tensor/trax/rlax/envs/__init__.py
rename to tensor2tensor/trax/rl/envs/__init__.py
index fa029618a..7fae2be45 100644
--- a/tensor2tensor/trax/rlax/envs/__init__.py
+++ b/tensor2tensor/trax/rl/envs/__init__.py
@@ -13,21 +13,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Environments defined in RLAX."""
+"""Environments defined in RL."""
 
 import gin
 from gym.envs.registration import register
 
-from tensor2tensor.trax.rlax.envs import online_tune_env
+from tensor2tensor.trax.rl.envs import online_tune_env
 
 
 # Ginify and register in gym.
 def configure_and_register_env(env_class):
   register(
       id="{}-v0".format(env_class.__name__),
-      entry_point="tensor2tensor.trax.rlax.envs:{}".format(env_class.__name__),
+      entry_point="tensor2tensor.trax.rl.envs:{}".format(env_class.__name__),
   )
-  return gin.external_configurable(env_class, module="trax.rlax.envs")
+  return gin.external_configurable(env_class, module="trax.rl.envs")
 
 
 # pylint: disable=invalid-name
diff --git a/tensor2tensor/trax/rlax/envs/env_service_server.py b/tensor2tensor/trax/rl/envs/env_service_server.py
similarity index 98%
rename from tensor2tensor/trax/rlax/envs/env_service_server.py
rename to tensor2tensor/trax/rl/envs/env_service_server.py
index 4531c317a..15a5bbd4b 100644
--- a/tensor2tensor/trax/rlax/envs/env_service_server.py
+++ b/tensor2tensor/trax/rl/envs/env_service_server.py
@@ -33,7 +33,7 @@
 from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.envs import server_utils
 from tensor2tensor.rl.google import atari_utils
-from tensor2tensor.trax.rlax import envs  # pylint: disable=unused-import
+from tensor2tensor.trax.rl import envs  # pylint: disable=unused-import
 import tensorflow as tf
 
 
diff --git a/tensor2tensor/trax/rlax/envs/fake_env.py b/tensor2tensor/trax/rl/envs/fake_env.py
similarity index 100%
rename from tensor2tensor/trax/rlax/envs/fake_env.py
rename to tensor2tensor/trax/rl/envs/fake_env.py
diff --git a/tensor2tensor/trax/rlax/envs/fake_env_test.py b/tensor2tensor/trax/rl/envs/fake_env_test.py
similarity index 94%
rename from tensor2tensor/trax/rlax/envs/fake_env_test.py
rename to tensor2tensor/trax/rl/envs/fake_env_test.py
index e6dabf471..cff4d481f 100644
--- a/tensor2tensor/trax/rlax/envs/fake_env_test.py
+++ b/tensor2tensor/trax/rl/envs/fake_env_test.py
@@ -13,13 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for tensor2tensor.trax.rlax.fake_env."""
+"""Tests for tensor2tensor.trax.rl.fake_env."""
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.trax.rlax.envs import fake_env
+from tensor2tensor.trax.rl.envs import fake_env
 from tensorflow import test
 
 
diff --git a/tensor2tensor/trax/rlax/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
similarity index 100%
rename from tensor2tensor/trax/rlax/envs/online_tune_env.py
rename to tensor2tensor/trax/rl/envs/online_tune_env.py
diff --git a/tensor2tensor/trax/rlax/envs/online_tune_env_test.py b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
similarity index 97%
rename from tensor2tensor/trax/rlax/envs/online_tune_env_test.py
rename to tensor2tensor/trax/rl/envs/online_tune_env_test.py
index fd8e8438d..df9e1ab69 100644
--- a/tensor2tensor/trax/rlax/envs/online_tune_env_test.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for tensor2tensor.trax.rlax.online_tune_env."""
+"""Tests for tensor2tensor.trax.rl.online_tune_env."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -26,7 +26,7 @@
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import models
 from tensor2tensor.trax import trax
-from tensor2tensor.trax.rlax.envs import online_tune_env
+from tensor2tensor.trax.rl.envs import online_tune_env
 from tensorflow import test
 from tensorflow.io import gfile
 
diff --git a/tensor2tensor/trax/rlax/ppo.py b/tensor2tensor/trax/rl/ppo.py
similarity index 100%
rename from tensor2tensor/trax/rlax/ppo.py
rename to tensor2tensor/trax/rl/ppo.py
diff --git a/tensor2tensor/trax/rlax/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
similarity index 99%
rename from tensor2tensor/trax/rlax/ppo_test.py
rename to tensor2tensor/trax/rl/ppo_test.py
index 0fd4337c8..b91bf02c9 100644
--- a/tensor2tensor/trax/rlax/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for tensor2tensor.trax.rlax.ppo."""
+"""Tests for tensor2tensor.trax.rl.ppo."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -24,7 +24,7 @@
 import numpy as np
 from tensor2tensor.trax import layers
 from tensor2tensor.trax import trax
-from tensor2tensor.trax.rlax import ppo
+from tensor2tensor.trax.rl import ppo
 from tensorflow import test
 
 
diff --git a/tensor2tensor/trax/rlax/ppo_training_loop_test.py b/tensor2tensor/trax/rl/ppo_training_loop_test.py
similarity index 96%
rename from tensor2tensor/trax/rlax/ppo_training_loop_test.py
rename to tensor2tensor/trax/rl/ppo_training_loop_test.py
index 269a438f7..fb40fd86c 100644
--- a/tensor2tensor/trax/rlax/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rl/ppo_training_loop_test.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for tensor2tensor.trax.rlax.ppo's training_loop."""
+"""Tests for tensor2tensor.trax.rl.ppo's training_loop."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -37,9 +37,9 @@
 from tensor2tensor.trax import models
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
-from tensor2tensor.trax.rlax import envs  # pylint: disable=unused-import
-from tensor2tensor.trax.rlax import ppo
-from tensor2tensor.trax.rlax import simulated_env_problem
+from tensor2tensor.trax.rl import envs  # pylint: disable=unused-import
+from tensor2tensor.trax.rl import ppo
+from tensor2tensor.trax.rl import simulated_env_problem
 from tensorflow import test
 from tensorflow.io import gfile
 
diff --git a/tensor2tensor/trax/rlax/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
similarity index 99%
rename from tensor2tensor/trax/rlax/simulated_env_problem.py
rename to tensor2tensor/trax/rl/simulated_env_problem.py
index 52798d06c..f2dbccbc0 100644
--- a/tensor2tensor/trax/rlax/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -28,7 +28,7 @@
 from tensor2tensor.trax import backend
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.backend import random as jax_random
-from tensor2tensor.trax.rlax import space_serializer
+from tensor2tensor.trax.rl import space_serializer
 
 
 class SimulatedEnvProblem(env_problem.EnvProblem):
diff --git a/tensor2tensor/trax/rlax/simulated_env_problem_test.py b/tensor2tensor/trax/rl/simulated_env_problem_test.py
similarity index 98%
rename from tensor2tensor/trax/rlax/simulated_env_problem_test.py
rename to tensor2tensor/trax/rl/simulated_env_problem_test.py
index f4c895889..72149ec1e 100644
--- a/tensor2tensor/trax/rlax/simulated_env_problem_test.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem_test.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for tensor2tensor.trax.rlax.simulated_env_problem."""
+"""Tests for tensor2tensor.trax.rl.simulated_env_problem."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -28,7 +28,7 @@
 
 from tensor2tensor.trax import backend
 from tensor2tensor.trax import trax
-from tensor2tensor.trax.rlax import simulated_env_problem
+from tensor2tensor.trax.rl import simulated_env_problem
 from tensorflow import test
 
 
diff --git a/tensor2tensor/trax/rlax/space_serializer.py b/tensor2tensor/trax/rl/space_serializer.py
similarity index 100%
rename from tensor2tensor/trax/rlax/space_serializer.py
rename to tensor2tensor/trax/rl/space_serializer.py
diff --git a/tensor2tensor/trax/rlax/space_serializer_test.py b/tensor2tensor/trax/rl/space_serializer_test.py
similarity index 96%
rename from tensor2tensor/trax/rlax/space_serializer_test.py
rename to tensor2tensor/trax/rl/space_serializer_test.py
index 0c62b3bbe..288747c24 100644
--- a/tensor2tensor/trax/rlax/space_serializer_test.py
+++ b/tensor2tensor/trax/rl/space_serializer_test.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for tensor2tensor.trax.rlax.space_serializer."""
+"""Tests for tensor2tensor.trax.rl.space_serializer."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -23,7 +23,7 @@
 import gym
 import numpy as np
 
-from tensor2tensor.trax.rlax import space_serializer
+from tensor2tensor.trax.rl import space_serializer
 from tensorflow import test
 
 
diff --git a/tensor2tensor/trax/rlax/ppo_main.py b/tensor2tensor/trax/rl_trainer.py
similarity index 97%
rename from tensor2tensor/trax/rlax/ppo_main.py
rename to tensor2tensor/trax/rl_trainer.py
index a5990cef4..3f43f687d 100644
--- a/tensor2tensor/trax/rlax/ppo_main.py
+++ b/tensor2tensor/trax/rl_trainer.py
@@ -13,7 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-r"""PPO binary over a gym env.
+r"""Trainer for RL environments.
+
+For now we only support PPO as RL algorithm.
 
 Sample invocation:
 
@@ -23,7 +25,7 @@
 RANDOM_SEED=0
 BOUNDARY=100
 
-python trax/rlax/ppo_main.py \
+python trax/rl_trainer.py \
   --env_problem_name=${ENV_PROBLEM_NAME} \
   --batch_size=${BATCH_SIZE} \
   --config=ppo.training_loop.epochs=${EPOCHS} \
@@ -55,8 +57,8 @@
 from tensor2tensor.rl import gym_utils
 from tensor2tensor.rl.google import atari_utils  # GOOGLE-INTERNAL:
 from tensor2tensor.trax import models
-from tensor2tensor.trax.rlax import envs as rlax_envs  # pylint: disable=unused-import
-from tensor2tensor.trax.rlax import ppo
+from tensor2tensor.trax.rl import envs as rl_envs  # pylint: disable=unused-import
+from tensor2tensor.trax.rl import ppo
 
 
 FLAGS = flags.FLAGS

From ebd32ee2de9394c9f58bb5f27f53d787b1be681c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 14 Aug 2019 17:43:21 -0700
Subject: [PATCH 2285/2720] Name and polarity change ("self._needs_init" -->
 self._init_finished").

PiperOrigin-RevId: 263468129
---
 tensor2tensor/trax/layers/base.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 0d76ef642..3e2448a63 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -67,9 +67,9 @@ class Layer(object):
 
   def __init__(self, **kwargs):
     self._init_kwargs = kwargs  # can be used in creating a generic decorator
-    self._needs_init = True
     self._params = ()  # cached parameters
     self._caller = _find_frame(inspect.stack())  # for custom error messages
+    self._init_finished = False
 
   def __repr__(self):
     class_str = self.__class__.__name__
@@ -213,9 +213,9 @@ def initialize(self, input_shapes, input_dtype, rng):
     """
     try:
       # Initialize params once; store them for use when this layer is called.
-      if self._needs_init:
+      if not self._init_finished:
         self._params = self.new_parameters(input_shapes, input_dtype, rng)
-        self._needs_init = False
+        self._init_finished = True
         return self._params
       else:
         return ()

From 33783fd63bd0debe2138c5569698b31d9af350f6 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 15 Aug 2019 15:24:42 -0700
Subject: [PATCH 2286/2720] Use Transformer as policy

PiperOrigin-RevId: 263654116
---
 tensor2tensor/envs/env_problem_utils.py       |  8 +++--
 tensor2tensor/trax/models/__init__.py         |  1 +
 tensor2tensor/trax/models/atari_cnn.py        |  2 ++
 tensor2tensor/trax/models/transformer.py      | 35 +++++++++++++++++++
 tensor2tensor/trax/rl/configs/acrobot.gin     | 26 ++++++++++++++
 .../trax/rl/configs/acrobot_transformer.gin   | 30 ++++++++++++++++
 tensor2tensor/trax/rl/configs/atari.gin       |  2 ++
 ...params_online_tune_wide_resnet_cifar10.gin |  2 ++
 .../online_tune_wide_resnet_cifar10.gin       |  1 +
 tensor2tensor/trax/rl/ppo.py                  | 34 +++++++++++++-----
 .../trax/rl/ppo_training_loop_test.py         | 29 ++++++++++++---
 tensor2tensor/trax/rl_trainer.py              | 12 -------
 12 files changed, 153 insertions(+), 29 deletions(-)
 create mode 100644 tensor2tensor/trax/rl/configs/acrobot.gin
 create mode 100644 tensor2tensor/trax/rl/configs/acrobot_transformer.gin

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 23d2d5f9c..5a686ce83 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -59,6 +59,7 @@ def play_env_problem_with_policy(env,
                                  reset=True,
                                  rng=None,
                                  temperature=1.0,
+                                 boundary=32,
                                  len_history_for_policy=32,
                                  num_to_keep=1):
   """Plays the given env with the policy function to collect trajectories.
@@ -75,8 +76,9 @@ def play_env_problem_with_policy(env,
       max_max_timestep is None or < 0
     rng: jax rng, splittable.
     temperature: float, temperature used in Gumbel sampling.
-    len_history_for_policy: int, the maximum history to keep for applying the
-      policy on. We also bucket observations on this number.
+    boundary: int, pad the sequences to the multiples of this number.
+    len_history_for_policy: int or None, the maximum history to keep for
+      applying the policy on. If None, use the whole history.
     num_to_keep: int, while truncating trajectory how many time-steps to keep.
 
   Returns:
@@ -107,7 +109,7 @@ def gumbel_sample(log_probs):
     # Shape is (B, T) + OBS
     # Bucket on whatever length is needed.
     padded_observations, lengths = env.trajectories.observations_np(
-        boundary=len_history_for_policy,
+        boundary=boundary,
         len_history_for_policy=len_history_for_policy)
 
     B, T = padded_observations.shape[:2]  # pylint: disable=invalid-name
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index 30876c202..d1ee95066 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -44,6 +44,7 @@ def model_configure(*args, **kwargs):
     position_lookup_transformer.PositionLookupTransformerLM)
 Resnet50 = model_configure(resnet.Resnet50)
 Transformer = model_configure(transformer.Transformer)
+TransformerDecoder = model_configure(transformer.TransformerDecoder)
 TransformerEncoder = model_configure(transformer.TransformerEncoder)
 TransformerLM = model_configure(transformer.TransformerLM)
 TransformerRevnetLM = model_configure(transformer_revnet.TransformerRevnetLM)
diff --git a/tensor2tensor/trax/models/atari_cnn.py b/tensor2tensor/trax/models/atari_cnn.py
index 7b1e8e251..d50a295d9 100644
--- a/tensor2tensor/trax/models/atari_cnn.py
+++ b/tensor2tensor/trax/models/atari_cnn.py
@@ -27,6 +27,8 @@ def FrameStack(n_frames):
   # Input shape: (B, T, ..., C).
   # Output shape: (B, T, ..., C * n_frames).
   assert n_frames >= 1
+  if n_frames == 1:
+    return ()
   return (
       # Make n_frames copies of the input sequence.
       [tl.Dup()] * (n_frames - 1),
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 63b673c34..3d2604dca 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -140,6 +140,41 @@ def DecoderBlock(d_model, d_ff, n_heads, dropout, mode):
   ]
 
 
+def TransformerDecoder(d_model=512,
+                       d_ff=2048,
+                       n_layers=6,
+                       n_heads=8,
+                       dropout=0.1,
+                       max_len=2048,
+                       mode='train'):
+  """Returns a Transformer decoder model.
+
+  The input to the model is a continuous tensor. Does not shift the input to the
+  right, i.e. the output for timestep t is based on inputs up to timestep t
+  inclusively.
+
+  Args:
+    d_model: int:  depth of embedding
+    d_ff: int: depth of feed-forward layer
+    n_layers: int: number of encoder/decoder layers
+    n_heads: int: number of attention heads
+    dropout: float: dropout rate (how much to drop out)
+    max_len: int: maximum symbol length for positional encoding
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    A Transformer decoder as a layer that maps from a continuous tensor to
+    a continuous tensor.
+  """
+  return tl.Model(                  # vecs
+      tl.PositionalEncoding(max_len=max_len),
+      tl.Dense(d_model),            # vecs
+      [DecoderBlock(d_model, d_ff, n_heads, dropout, mode)
+       for _ in range(n_layers)],   # vecs
+      tl.LayerNorm(),               # vecs
+  )
+
+
 def TransformerLM(vocab_size,
                   d_model=512,
                   d_ff=2048,
diff --git a/tensor2tensor/trax/rl/configs/acrobot.gin b/tensor2tensor/trax/rl/configs/acrobot.gin
new file mode 100644
index 000000000..9e4b1a3f9
--- /dev/null
+++ b/tensor2tensor/trax/rl/configs/acrobot.gin
@@ -0,0 +1,26 @@
+import tensor2tensor.trax.models
+import tensor2tensor.trax.rl
+
+# Parameters for FrameStackMLP:
+# ==============================================================================
+FrameStackMLP.n_frames = 1
+
+# Parameters for ppo.training_loop:
+# ==============================================================================
+ppo.training_loop.epochs = 40000
+ppo.training_loop.n_optimizer_steps = 32
+ppo.training_loop.target_kl = 1000  # Virtually infinite.
+ppo.training_loop.boundary = 512
+ppo.training_loop.max_timestep = 512
+ppo.training_loop.max_timestep_eval = 20000
+ppo.training_loop.random_seed = None
+ppo.training_loop.gamma = 0.99
+ppo.training_loop.lambda_ = 0.95
+ppo.training_loop.c1 = 1.0
+ppo.training_loop.c2 = 0.0
+ppo.training_loop.eval_every_n = 500
+ppo.training_loop.done_frac_for_policy_save = 0.9
+ppo.training_loop.n_evals = 16
+ppo.training_loop.len_history_for_policy = 1
+ppo.training_loop.eval_temperatures = (1.0, 0.5)
+ppo.training_loop.policy_and_value_model = @trax.models.FrameStackMLP
diff --git a/tensor2tensor/trax/rl/configs/acrobot_transformer.gin b/tensor2tensor/trax/rl/configs/acrobot_transformer.gin
new file mode 100644
index 000000000..fabaa5715
--- /dev/null
+++ b/tensor2tensor/trax/rl/configs/acrobot_transformer.gin
@@ -0,0 +1,30 @@
+import tensor2tensor.trax.models
+import tensor2tensor.trax.rl
+
+# Parameters for TransformerDecoder:
+# ==============================================================================
+TransformerDecoder.d_model = 64
+TransformerDecoder.d_ff = 128
+TransformerDecoder.dropout = 0.1
+TransformerDecoder.n_heads = 2
+TransformerDecoder.n_layers = 1
+
+# Parameters for ppo.training_loop:
+# ==============================================================================
+ppo.training_loop.epochs = 40000
+ppo.training_loop.n_optimizer_steps = 32
+ppo.training_loop.target_kl = 1000  # Virtually infinite.
+ppo.training_loop.boundary = 512
+ppo.training_loop.max_timestep = 512
+ppo.training_loop.max_timestep_eval = 20000
+ppo.training_loop.random_seed = None
+ppo.training_loop.gamma = 0.99
+ppo.training_loop.lambda_ = 0.95
+ppo.training_loop.c1 = 1.0
+ppo.training_loop.c2 = 0.0
+ppo.training_loop.eval_every_n = 500
+ppo.training_loop.done_frac_for_policy_save = 0.9
+ppo.training_loop.n_evals = 16
+ppo.training_loop.len_history_for_policy = None
+ppo.training_loop.eval_temperatures = (1.0, 0.5)
+ppo.training_loop.policy_and_value_model = @trax.models.TransformerDecoder
diff --git a/tensor2tensor/trax/rl/configs/atari.gin b/tensor2tensor/trax/rl/configs/atari.gin
index fc6fa7fec..f0155b5f3 100644
--- a/tensor2tensor/trax/rl/configs/atari.gin
+++ b/tensor2tensor/trax/rl/configs/atari.gin
@@ -1,3 +1,4 @@
+import tensor2tensor.trax.models
 import tensor2tensor.trax.rl
 
 # Parameters for ppo.training_loop:
@@ -18,3 +19,4 @@ ppo.training_loop.done_frac_for_policy_save = 0.9
 ppo.training_loop.n_evals = 16
 ppo.training_loop.len_history_for_policy = 4
 ppo.training_loop.eval_temperatures = (1.0, 0.5)
+ppo.training_loop.policy_and_value_model = @trax.models.AtariCnn
diff --git a/tensor2tensor/trax/rl/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
index dc1155525..4691ad9d8 100644
--- a/tensor2tensor/trax/rl/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
@@ -1,3 +1,4 @@
+import tensor2tensor.trax.models
 import tensor2tensor.trax.rl
 
 # Parameters for ppo.training_loop:
@@ -17,3 +18,4 @@ ppo.training_loop.n_evals = 1
 ppo.training_loop.len_history_for_policy = 1  # this needs to be bumped up.
 ppo.training_loop.eval_temperatures = (1.0,)
 ppo.training_loop.epochs = 1000
+ppo.training_loop.policy_and_value_model = @trax.models.TransformerDecoder
diff --git a/tensor2tensor/trax/rl/configs/online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/online_tune_wide_resnet_cifar10.gin
index eb0ebd98c..4ffe114ec 100644
--- a/tensor2tensor/trax/rl/configs/online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/online_tune_wide_resnet_cifar10.gin
@@ -66,3 +66,4 @@ ppo.training_loop.done_frac_for_policy_save = 0
 ppo.training_loop.n_evals = 1
 ppo.training_loop.eval_temperatures = (1.0,)
 ppo.training_loop.len_history_for_policy = 4
+ppo.training_loop.policy_and_value_model = @trax.models.TransformerDecoder
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index c517ab55a..c398c2d1d 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -159,6 +159,7 @@ def collect_trajectories(env,
                          max_timestep=None,
                          reset=True,
                          len_history_for_policy=32,
+                         boundary=32,
                          rng=None):
   """Collect trajectories with the given policy net and behaviour.
 
@@ -171,8 +172,9 @@ def collect_trajectories(env,
       done.
     reset: bool, true if we want to reset the envs. The envs are also reset if
       max_max_timestep is None or < 0
-    len_history_for_policy: int, the maximum history to keep for applying the
-      policy on.
+    len_history_for_policy: int or None, the maximum history to keep for
+      applying the policy on. If None, use the full history.
+    boundary: int, pad the sequences to the multiples of this number.
     rng: jax rng, splittable.
 
   Returns:
@@ -193,6 +195,7 @@ def collect_trajectories(env,
       max_timestep=max_timestep,
       reset=reset,
       len_history_for_policy=len_history_for_policy,
+      boundary=boundary,
       rng=rng)
   # Skip returning raw_rewards here, since they aren't used.
 
@@ -853,8 +856,9 @@ def __init__(
       self,
       train_env,
       eval_env,
-      policy_and_value_net_fn,
+      policy_and_value_model,
       policy_and_value_optimizer_fn,
+      policy_and_value_two_towers,
       output_dir,
       n_optimizer_steps,
       print_every_optimizer_steps,
@@ -904,9 +908,15 @@ def __init__(
 
     # Initialize the policy and value network.
     policy_and_value_net_params, policy_and_value_net_apply = (
-        policy_and_value_net_fn(key1, batch_observations_shape,
-                                observations_dtype, n_actions))
-
+        policy_and_value_net(
+            rng_key=key1,
+            batch_observations_shape=batch_observations_shape,
+            observations_dtype=observations_dtype,
+            n_actions=n_actions,
+            bottom_layers_fn=policy_and_value_model,
+            two_towers=policy_and_value_two_towers,
+        )
+    )
     self._policy_and_value_net_apply = jit(policy_and_value_net_apply)
 
     # Maybe restore the policy params. If there is nothing to restore, then
@@ -967,6 +977,7 @@ def train_epoch(self):
         max_timestep=self._max_timestep,
         rng=key,
         len_history_for_policy=self._len_history_for_policy,
+        boundary=self._boundary,
         reset=self._should_reset,
     )
     self._should_reset = False
@@ -1257,9 +1268,10 @@ def _get_predictions(self, observations, rng=None):
 def training_loop(
     train_env,
     eval_env,
-    policy_and_value_net_fn,
+    policy_and_value_model,
     policy_and_value_optimizer_fn,
     output_dir,
+    policy_and_value_two_towers=False,
     epochs=EPOCHS,
     n_optimizer_steps=N_OPTIMIZER_STEPS,
     print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
@@ -1283,9 +1295,12 @@ def training_loop(
   Args:
     train_env: gym.Env to use for training.
     eval_env: gym.Env to use for evaluation.
-    policy_and_value_net_fn: Function defining the policy and value network.
+    policy_and_value_model: Function defining the policy and value network,
+      without the policy and value heads.
     policy_and_value_optimizer_fn: Function defining the optimizer.
     output_dir: Output dir.
+    policy_and_value_two_towers: Whether to use two separate models as the
+      policy and value networks. If False, share their parameters.
     epochs: Number of epochs to run for.
     n_optimizer_steps: Number of optimizer steps.
     print_every_optimizer_steps: How often to log during the policy optimization
@@ -1313,8 +1328,9 @@ def training_loop(
   trainer = PPOTrainer(
       train_env=train_env,
       eval_env=eval_env,
-      policy_and_value_net_fn=policy_and_value_net_fn,
+      policy_and_value_model=policy_and_value_model,
       policy_and_value_optimizer_fn=policy_and_value_optimizer_fn,
+      policy_and_value_two_towers=policy_and_value_two_towers,
       output_dir=output_dir,
       n_optimizer_steps=n_optimizer_steps,
       print_every_optimizer_steps=print_every_optimizer_steps,
diff --git a/tensor2tensor/trax/rl/ppo_training_loop_test.py b/tensor2tensor/trax/rl/ppo_training_loop_test.py
index fb40fd86c..7bfea65fc 100644
--- a/tensor2tensor/trax/rl/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rl/ppo_training_loop_test.py
@@ -69,20 +69,22 @@ def tmp_dir(self):
     yield tmp
     gfile.rmtree(tmp)
 
-  def _run_training_loop(self, train_env, eval_env, output_dir):
+  def _run_training_loop(self, train_env, eval_env, output_dir, model=None):
+    if model is None:
+      model = lambda: [layers.Dense(1)]
     n_epochs = 2
     # Run the training loop.
     ppo.training_loop(
         train_env=train_env,
         eval_env=eval_env,
         epochs=n_epochs,
-        policy_and_value_net_fn=functools.partial(
-            ppo.policy_and_value_net,
-            bottom_layers_fn=lambda: [layers.Dense(1)]),
+        policy_and_value_model=model,
         policy_and_value_optimizer_fn=ppo.optimizer_fn,
         n_optimizer_steps=1,
         output_dir=output_dir,
-        random_seed=0)
+        random_seed=0,
+        boundary=2,
+    )
 
   def test_training_loop_cartpole(self):
     with self.tmp_dir() as output_dir:
@@ -92,6 +94,23 @@ def test_training_loop_cartpole(self):
           output_dir=output_dir,
       )
 
+  def test_training_loop_cartpole_transformer(self):
+    with self.tmp_dir() as output_dir:
+      self._run_training_loop(
+          train_env=self.get_wrapped_env("CartPole-v0", 2),
+          eval_env=self.get_wrapped_env("CartPole-v0", 2),
+          output_dir=output_dir,
+          model=functools.partial(
+              models.TransformerDecoder,
+              d_model=1,
+              d_ff=1,
+              n_layers=1,
+              n_heads=1,
+              max_len=64,
+              mode="train",
+          ),
+      )
+
   def test_training_loop_onlinetune(self):
     with self.tmp_dir() as output_dir:
       gin.bind_parameter("OnlineTuneEnv.model", functools.partial(
diff --git a/tensor2tensor/trax/rl_trainer.py b/tensor2tensor/trax/rl_trainer.py
index 3f43f687d..f49fe178c 100644
--- a/tensor2tensor/trax/rl_trainer.py
+++ b/tensor2tensor/trax/rl_trainer.py
@@ -56,7 +56,6 @@
 from tensor2tensor.envs import rendered_env_problem
 from tensor2tensor.rl import gym_utils
 from tensor2tensor.rl.google import atari_utils  # GOOGLE-INTERNAL:
-from tensor2tensor.trax import models
 from tensor2tensor.trax.rl import envs as rl_envs  # pylint: disable=unused-import
 from tensor2tensor.trax.rl import ppo
 
@@ -204,27 +203,16 @@ def main(argv):
   eval_env = make_env(batch_size=FLAGS.eval_batch_size, **eval_env_kwargs)
   assert eval_env
 
-  # TODO(afrozm): Refactor.
-  if "NoFrameskip" in FLAGS.env_problem_name:
-    bottom_layers_fn = models.AtariCnn
-  else:
-    bottom_layers_fn = models.FrameStackMLP
-
   def run_training_loop():
     """Runs the training loop."""
     logging.info("Starting the training loop.")
 
-    policy_and_value_net_fn = functools.partial(
-        ppo.policy_and_value_net,
-        bottom_layers_fn=bottom_layers_fn,
-        two_towers=FLAGS.two_towers)
     policy_and_value_optimizer_fn = get_optimizer_fn(FLAGS.learning_rate)
 
     ppo.training_loop(
         output_dir=FLAGS.output_dir,
         train_env=train_env,
         eval_env=eval_env,
-        policy_and_value_net_fn=policy_and_value_net_fn,
         policy_and_value_optimizer_fn=policy_and_value_optimizer_fn,
     )
 

From 6a1b3a8c6d6961551e647c3cae058fb60d1cd191 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 16 Aug 2019 09:48:03 -0700
Subject: [PATCH 2287/2720] Correct bug with weights for translation model, add
 regression test.

PiperOrigin-RevId: 263789213
---
 ...former_wmt_ende_16gb_adafactor_testing.gin | 63 +++++++++++++++++++
 tensor2tensor/trax/trax.py                    |  4 +-
 2 files changed, 65 insertions(+), 2 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin

diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin
new file mode 100644
index 000000000..9577bb91b
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin
@@ -0,0 +1,63 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 128
+batch_fun.eval_batch_size = 64
+batch_fun.max_eval_length = 1024
+batch_fun.buckets_include_inputs_in_length=True
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_translate_ende_wmt32k'
+inputs.append_targets = True
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 1.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 20000
+
+# Parameters for Adafactor:
+# ==============================================================================
+Adafactor.beta1 = 0.0
+Adafactor.decay_rate = 0.8
+Adafactor.clipping_threshold = 1.0
+Adafactor.epsilon1 = 1e-30
+Adafactor.epsilon2 = 0.001
+Adafactor.factored = True
+Adafactor.multiply_by_parameter_scale = True
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_preprocess
+wmt_preprocess.max_length = 512
+wmt_preprocess.max_eval_length = 1024
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 1000
+train.eval_steps = 10
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.Transformer
+train.train_steps = 100000
+train.optimizer = @trax.optimizers.Adafactor
+
+# Parameters for Transformer:
+# ==============================================================================
+Transformer.d_model = 512
+Transformer.d_ff = 2048
+Transformer.dropout = 0.1
+Transformer.max_len = 2048
+Transformer.mode = 'train'
+Transformer.n_heads = 8
+Transformer.n_layers = 6
+Transformer.input_vocab_size = 33300
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index ede9ca1e2..4a9f258a1 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -60,9 +60,9 @@ def unpack_batch(batch, has_weights=False):
     return batch
   else:
     inputs, targets = batch
-    if isinstance(inputs, (list, tuple)):
+    if isinstance(targets, (list, tuple)):
       # If weights are not provided, use scalar 1s and rely on broadcasting.
-      weights = [1.0] * len(inputs)
+      weights = [1.0] * len(targets)
     else:
       weights = 1.0
     return inputs, targets, weights

From 669b05c20fe7515cb890e4ed9de1df37cdfc5de5 Mon Sep 17 00:00:00 2001
From: Yeming Wen <ywenxu@google.com>
Date: Sat, 17 Aug 2019 12:31:18 -0700
Subject: [PATCH 2288/2720] Add TrainableNormal test in the InitializersTest.

Implementation is based on https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/init_ops_v2_test.py#L256

PiperOrigin-RevId: 263958356
---
 tensor2tensor/keras/initializers_test.py | 29 ++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/tensor2tensor/keras/initializers_test.py b/tensor2tensor/keras/initializers_test.py
index 145ee86bf..b65f9aae3 100644
--- a/tensor2tensor/keras/initializers_test.py
+++ b/tensor2tensor/keras/initializers_test.py
@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import math
+
 import numpy as np
 from tensor2tensor.keras import initializers
 from tensor2tensor.utils import test_utils
@@ -46,6 +48,33 @@ def testTrainableHalfCauchy(self):
     self.assertAllEqual(half_cauchy_value.shape, shape)
     self.assertAllGreaterEqual(half_cauchy_value, 0.)
 
+  @test_utils.run_in_graph_and_eager_modes
+  def testTrainableNormal(self):
+    shape = (100,)
+    # TrainableNormal is expected to have var 1/shape[0]
+    # because it by default has the fan_in mode scale normal std initializer.
+    initializer = initializers.get('trainable_normal')
+    normal = initializer(shape)
+    self.evaluate(tf.global_variables_initializer())
+    loc_value, scale_value = self.evaluate([
+        # Get distribution of rv -> get distribution of Independent.
+        normal.distribution.distribution.loc,
+        normal.distribution.distribution.scale])
+    fan_in = shape[0]
+    target_scale = 1.
+    target_scale /= max(1., fan_in)
+    target_scale = math.sqrt(target_scale)
+
+    self.assertAllClose(loc_value, np.zeros(shape), atol=1e-4)
+    # Tolerance is larger because of the scale normal std initializer.
+    # In this case it has std around 0.01 (0.1*target_scale).
+    self.assertAllClose(
+        scale_value, target_scale * np.ones(shape), atol=5e-2)
+
+    # Test the TrainableNormal initializer has the specified shape.
+    normal_value = self.evaluate(normal)
+    self.assertAllEqual(normal_value.shape, shape)
+
 
 if __name__ == '__main__':
   tf.test.main()

From 897ad7968c8227bc5751634d70748b67abc76c3e Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Sun, 18 Aug 2019 20:41:51 -0700
Subject: [PATCH 2289/2720] Create a module for RL trainers

PiperOrigin-RevId: 264076721
---
 tensor2tensor/trax/rl/base_trainer.py         |  57 ++
 tensor2tensor/trax/rl/configs/acrobot.gin     |  42 +-
 .../trax/rl/configs/acrobot_transformer.gin   |  42 +-
 tensor2tensor/trax/rl/configs/atari.gin       |  46 +-
 ...params_online_tune_wide_resnet_cifar10.gin |  10 +-
 ...params_online_tune_wide_resnet_cifar10.gin |  45 +-
 .../online_tune_wide_resnet_cifar10.gin       |  42 +-
 tensor2tensor/trax/rl/ppo.py                  | 528 +-----------------
 tensor2tensor/trax/rl/ppo_trainer.py          | 495 ++++++++++++++++
 ...ining_loop_test.py => ppo_trainer_test.py} |   9 +-
 tensor2tensor/trax/rl/trainers.py             |  35 ++
 tensor2tensor/trax/rl_trainer.py              | 153 ++---
 12 files changed, 805 insertions(+), 699 deletions(-)
 create mode 100644 tensor2tensor/trax/rl/base_trainer.py
 create mode 100644 tensor2tensor/trax/rl/ppo_trainer.py
 rename tensor2tensor/trax/rl/{ppo_training_loop_test.py => ppo_trainer_test.py} (97%)
 create mode 100644 tensor2tensor/trax/rl/trainers.py

diff --git a/tensor2tensor/trax/rl/base_trainer.py b/tensor2tensor/trax/rl/base_trainer.py
new file mode 100644
index 000000000..68d7bc01c
--- /dev/null
+++ b/tensor2tensor/trax/rl/base_trainer.py
@@ -0,0 +1,57 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 class for RL trainers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl import logging
+from tensorflow.io import gfile
+
+
+class BaseTrainer(object):
+  """Base class for RL trainers."""
+
+  def __init__(self, train_env, eval_env, output_dir):
+    self._train_env = train_env
+    self._eval_env = eval_env
+    self._output_dir = output_dir
+    gfile.makedirs(self._output_dir)
+
+  @property
+  def epoch(self):
+    raise NotImplementedError
+
+  def train_epoch(self):
+    raise NotImplementedError
+
+  def evaluate(self):
+    raise NotImplementedError
+
+  def save(self):
+    raise NotImplementedError
+
+  def flush_summaries(self):
+    raise NotImplementedError
+
+  def training_loop(self, n_epochs):
+    logging.info("Starting the RL training loop.")
+    for _ in range(self.epoch, n_epochs):
+      self.train_epoch()
+    self.save()
+    self.evaluate()
+    self.flush_summaries()
diff --git a/tensor2tensor/trax/rl/configs/acrobot.gin b/tensor2tensor/trax/rl/configs/acrobot.gin
index 9e4b1a3f9..e122270ae 100644
--- a/tensor2tensor/trax/rl/configs/acrobot.gin
+++ b/tensor2tensor/trax/rl/configs/acrobot.gin
@@ -1,26 +1,30 @@
 import tensor2tensor.trax.models
-import tensor2tensor.trax.rl
+import tensor2tensor.trax.rl.trainers
 
 # Parameters for FrameStackMLP:
 # ==============================================================================
 FrameStackMLP.n_frames = 1
 
-# Parameters for ppo.training_loop:
+# Parameters for PPO:
 # ==============================================================================
-ppo.training_loop.epochs = 40000
-ppo.training_loop.n_optimizer_steps = 32
-ppo.training_loop.target_kl = 1000  # Virtually infinite.
-ppo.training_loop.boundary = 512
-ppo.training_loop.max_timestep = 512
-ppo.training_loop.max_timestep_eval = 20000
-ppo.training_loop.random_seed = None
-ppo.training_loop.gamma = 0.99
-ppo.training_loop.lambda_ = 0.95
-ppo.training_loop.c1 = 1.0
-ppo.training_loop.c2 = 0.0
-ppo.training_loop.eval_every_n = 500
-ppo.training_loop.done_frac_for_policy_save = 0.9
-ppo.training_loop.n_evals = 16
-ppo.training_loop.len_history_for_policy = 1
-ppo.training_loop.eval_temperatures = (1.0, 0.5)
-ppo.training_loop.policy_and_value_model = @trax.models.FrameStackMLP
+PPO.n_optimizer_steps = 32
+PPO.target_kl = 1000  # Virtually infinite.
+PPO.boundary = 512
+PPO.max_timestep = 512
+PPO.max_timestep_eval = 20000
+PPO.random_seed = None
+PPO.gamma = 0.99
+PPO.lambda_ = 0.95
+PPO.c1 = 1.0
+PPO.c2 = 0.0
+PPO.eval_every_n = 500
+PPO.done_frac_for_policy_save = 0.9
+PPO.n_evals = 16
+PPO.len_history_for_policy = 1
+PPO.eval_temperatures = (1.0, 0.5)
+PPO.policy_and_value_model = @trax.models.FrameStackMLP
+
+# Parameters for train_rl:
+# ==============================================================================
+train_rl.env_name = "Acrobot-v1"
+train_rl.n_epochs = 40000
diff --git a/tensor2tensor/trax/rl/configs/acrobot_transformer.gin b/tensor2tensor/trax/rl/configs/acrobot_transformer.gin
index fabaa5715..16343023d 100644
--- a/tensor2tensor/trax/rl/configs/acrobot_transformer.gin
+++ b/tensor2tensor/trax/rl/configs/acrobot_transformer.gin
@@ -1,5 +1,5 @@
 import tensor2tensor.trax.models
-import tensor2tensor.trax.rl
+import tensor2tensor.trax.rl.trainers
 
 # Parameters for TransformerDecoder:
 # ==============================================================================
@@ -9,22 +9,26 @@ TransformerDecoder.dropout = 0.1
 TransformerDecoder.n_heads = 2
 TransformerDecoder.n_layers = 1
 
-# Parameters for ppo.training_loop:
+# Parameters for PPO:
 # ==============================================================================
-ppo.training_loop.epochs = 40000
-ppo.training_loop.n_optimizer_steps = 32
-ppo.training_loop.target_kl = 1000  # Virtually infinite.
-ppo.training_loop.boundary = 512
-ppo.training_loop.max_timestep = 512
-ppo.training_loop.max_timestep_eval = 20000
-ppo.training_loop.random_seed = None
-ppo.training_loop.gamma = 0.99
-ppo.training_loop.lambda_ = 0.95
-ppo.training_loop.c1 = 1.0
-ppo.training_loop.c2 = 0.0
-ppo.training_loop.eval_every_n = 500
-ppo.training_loop.done_frac_for_policy_save = 0.9
-ppo.training_loop.n_evals = 16
-ppo.training_loop.len_history_for_policy = None
-ppo.training_loop.eval_temperatures = (1.0, 0.5)
-ppo.training_loop.policy_and_value_model = @trax.models.TransformerDecoder
+PPO.n_optimizer_steps = 32
+PPO.target_kl = 1000  # Virtually infinite.
+PPO.boundary = 512
+PPO.max_timestep = 512
+PPO.max_timestep_eval = 20000
+PPO.random_seed = None
+PPO.gamma = 0.99
+PPO.lambda_ = 0.95
+PPO.c1 = 1.0
+PPO.c2 = 0.0
+PPO.eval_every_n = 500
+PPO.done_frac_for_policy_save = 0.9
+PPO.n_evals = 16
+PPO.len_history_for_policy = None
+PPO.eval_temperatures = (1.0, 0.5)
+PPO.policy_and_value_model = @trax.models.TransformerDecoder
+
+# Parameters for train_rl:
+# ==============================================================================
+train_rl.env_name = "Acrobot-v1"
+train_rl.n_epochs = 40000
diff --git a/tensor2tensor/trax/rl/configs/atari.gin b/tensor2tensor/trax/rl/configs/atari.gin
index f0155b5f3..e4c0ab2c1 100644
--- a/tensor2tensor/trax/rl/configs/atari.gin
+++ b/tensor2tensor/trax/rl/configs/atari.gin
@@ -1,22 +1,30 @@
 import tensor2tensor.trax.models
-import tensor2tensor.trax.rl
+import tensor2tensor.trax.rl.trainers
 
-# Parameters for ppo.training_loop:
+# Parameters for PPO:
 # ==============================================================================
-ppo.training_loop.epochs = 40000
-ppo.training_loop.n_optimizer_steps = 4
-ppo.training_loop.target_kl = 0.01
-ppo.training_loop.boundary = 20
-ppo.training_loop.max_timestep = 128
-ppo.training_loop.max_timestep_eval = 20000
-ppo.training_loop.random_seed = 0
-ppo.training_loop.gamma = 0.99
-ppo.training_loop.lambda_ = 0.95
-ppo.training_loop.c1 = 1.0
-ppo.training_loop.c2 = 0.01
-ppo.training_loop.eval_every_n = 500
-ppo.training_loop.done_frac_for_policy_save = 0.9
-ppo.training_loop.n_evals = 16
-ppo.training_loop.len_history_for_policy = 4
-ppo.training_loop.eval_temperatures = (1.0, 0.5)
-ppo.training_loop.policy_and_value_model = @trax.models.AtariCnn
+PPO.n_optimizer_steps = 4
+PPO.target_kl = 0.01
+PPO.boundary = 20
+PPO.max_timestep = 128
+PPO.max_timestep_eval = 20000
+PPO.random_seed = None
+PPO.gamma = 0.99
+PPO.lambda_ = 0.95
+PPO.c1 = 1.0
+PPO.c2 = 0.01
+PPO.eval_every_n = 500
+PPO.done_frac_for_policy_save = 0.9
+PPO.n_evals = 16
+PPO.len_history_for_policy = 4
+PPO.eval_temperatures = (1.0, 0.5)
+PPO.policy_and_value_model = @trax.models.AtariCnn
+
+# Parameters for train_rl:
+# ==============================================================================
+train_rl.env_name = "PongNoFrameskip-v4"
+train_rl.n_epochs = 40000
+train_rl.clip_rewards = True
+train_rl.max_timestep = 10000
+train_rl.rendered_env = True
+train_rl.resize_dims = (105, 80)
diff --git a/tensor2tensor/trax/rl/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin
index 78195a528..42260dd64 100644
--- a/tensor2tensor/trax/rl/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin
@@ -26,7 +26,7 @@ shuffle_and_batch_data.preprocess_fun = @trax.inputs.cifar10_no_augmentation_pre
 
 # Parameters for WideResnet:
 # ==============================================================================
-WideResnet.widen_factor = 2
+WideResnet.widen_factor = 10
 WideResnet.n_blocks = 3
 WideResnet.n_output_classes = 10
 
@@ -38,4 +38,10 @@ OnlineTuneEnv.optimizer = @trax.optimizers.Momentum
 OnlineTuneEnv.start_lr = 0.01
 OnlineTuneEnv.train_steps = 500
 OnlineTuneEnv.eval_steps = 50
-OnlineTuneEnv.env_steps = 128
+OnlineTuneEnv.env_steps = 100
+OnlineTuneEnv.observation_metrics = [
+    ("train", "metrics/accuracy"),
+    ("train", "metrics/loss"),
+    ("eval", "metrics/accuracy"),
+    ("eval", "metrics/loss"),
+]
diff --git a/tensor2tensor/trax/rl/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
index 4691ad9d8..86b3692e6 100644
--- a/tensor2tensor/trax/rl/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
@@ -1,21 +1,30 @@
 import tensor2tensor.trax.models
-import tensor2tensor.trax.rl
+import tensor2tensor.trax.rl.trainers
 
-# Parameters for ppo.training_loop:
+# Parameters for FrameStackMLP:
 # ==============================================================================
-ppo.training_loop.n_optimizer_steps = 30
-ppo.training_loop.boundary = 128
-ppo.training_loop.max_timestep = 128
-ppo.training_loop.max_timestep_eval = 128
-ppo.training_loop.random_seed = 0
-ppo.training_loop.gamma = 0.99
-ppo.training_loop.lambda_ = 0.95
-ppo.training_loop.c1 = 1.0
-ppo.training_loop.c2 = 0.01
-ppo.training_loop.eval_every_n = 10
-ppo.training_loop.done_frac_for_policy_save = 0
-ppo.training_loop.n_evals = 1
-ppo.training_loop.len_history_for_policy = 1  # this needs to be bumped up.
-ppo.training_loop.eval_temperatures = (1.0,)
-ppo.training_loop.epochs = 1000
-ppo.training_loop.policy_and_value_model = @trax.models.TransformerDecoder
+FrameStackMLP.n_frames = 4
+
+# Parameters for PPO:
+# ==============================================================================
+PPO.n_optimizer_steps = 30
+PPO.target_kl = 0.1
+PPO.boundary = 128
+PPO.max_timestep = 128
+PPO.max_timestep_eval = 128
+PPO.random_seed = None
+PPO.gamma = 0.99
+PPO.lambda_ = 0.95
+PPO.c1 = 1.0
+PPO.c2 = 0.01
+PPO.eval_every_n = 10
+PPO.done_frac_for_policy_save = 0
+PPO.n_evals = 1
+PPO.len_history_for_policy = 4
+PPO.eval_temperatures = (1.0, 0.5)
+PPO.policy_and_value_model = @trax.models.FrameStackMLP
+
+# Parameters for train_rl:
+# ==============================================================================
+train_rl.env_name = "ClientEnv-v0"
+train_rl.n_epochs = 1000
diff --git a/tensor2tensor/trax/rl/configs/online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/online_tune_wide_resnet_cifar10.gin
index 4ffe114ec..f89769344 100644
--- a/tensor2tensor/trax/rl/configs/online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/online_tune_wide_resnet_cifar10.gin
@@ -1,8 +1,8 @@
 import tensor2tensor.trax.inputs
 import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.rl
 import tensor2tensor.trax.rl.envs
+import tensor2tensor.trax.rl.trainers
 
 # Parameters for batch_fun:
 # ==============================================================================
@@ -41,7 +41,7 @@ OnlineTuneEnv.model = @trax.models.WideResnet
 OnlineTuneEnv.optimizer = @trax.optimizers.Momentum
 OnlineTuneEnv.start_lr = 0.01
 OnlineTuneEnv.train_steps = 500
-OnlineTuneEnv.eval_steps = 100
+OnlineTuneEnv.eval_steps = 50
 OnlineTuneEnv.env_steps = 100
 OnlineTuneEnv.observation_metrics = [
     ("train", "metrics/accuracy"),
@@ -50,20 +50,26 @@ OnlineTuneEnv.observation_metrics = [
     ("eval", "metrics/loss"),
 ]
 
-# Parameters for ppo.training_loop:
+# Parameters for PPO:
 # ==============================================================================
-ppo.training_loop.n_optimizer_steps = 30
-ppo.training_loop.boundary = 20
-ppo.training_loop.max_timestep = 128
-ppo.training_loop.max_timestep_eval = 20000
-ppo.training_loop.random_seed = 0
-ppo.training_loop.gamma = 0.99
-ppo.training_loop.lambda_ = 0.95
-ppo.training_loop.c1 = 1.0
-ppo.training_loop.c2 = 0.01
-ppo.training_loop.eval_every_n = 10
-ppo.training_loop.done_frac_for_policy_save = 0
-ppo.training_loop.n_evals = 1
-ppo.training_loop.eval_temperatures = (1.0,)
-ppo.training_loop.len_history_for_policy = 4
-ppo.training_loop.policy_and_value_model = @trax.models.TransformerDecoder
+PPO.n_optimizer_steps = 30
+PPO.target_kl = 0.1
+PPO.boundary = 128
+PPO.max_timestep = 128
+PPO.max_timestep_eval = 128
+PPO.random_seed = None
+PPO.gamma = 0.99
+PPO.lambda_ = 0.95
+PPO.c1 = 1.0
+PPO.c2 = 0.01
+PPO.eval_every_n = 10
+PPO.done_frac_for_policy_save = 0
+PPO.n_evals = 1
+PPO.len_history_for_policy = 4
+PPO.eval_temperatures = (1.0, 0.5)
+PPO.policy_and_value_model = @trax.models.FrameStackMLP
+
+# Parameters for train_rl:
+# ==============================================================================
+train_rl.env_name = "OnlineTuneEnv-v0"
+train_rl.n_epochs = 1000
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index c398c2d1d..15e0891ad 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -56,8 +56,6 @@
 
 from absl import logging
 import cloudpickle as pickle
-import gin
-import gym
 from jax import grad
 from jax import jit
 from jax import lax
@@ -66,21 +64,9 @@
 import numpy as onp
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import env_problem_utils
-from tensor2tensor.trax import jaxboard
 from tensor2tensor.trax import layers as tl
-from tensor2tensor.trax import optimizers as trax_opt
-from tensor2tensor.trax import trax
 from tensorflow.io import gfile
 
-DEBUG_LOGGING = False
-GAMMA = 0.99
-LAMBDA = 0.95
-EPSILON = 0.1
-EPOCHS = 50  # 100
-N_OPTIMIZER_STEPS = 100
-PRINT_EVERY_OPTIMIZER_STEP = 20
-BATCH_TRAJECTORIES = 32
-
 
 def policy_and_value_net(rng_key,
                          batch_observations_shape,
@@ -118,12 +104,12 @@ def policy_and_value_net(rng_key,
   return params, net
 
 
-def optimizer_fn(net_params, learning_rate=1e-3):
+def optimizer_fn(optimizer, net_params):
   """Exposes a convenient interface for the optimizer.
 
   Args:
+    optimizer: Optimizer class to use.
     net_params: A nested structure of network parameters.
-    learning_rate: Learning rate.
 
   Returns:
     A tuple (opt_state, opt_update, get_params), where:
@@ -133,7 +119,7 @@ def optimizer_fn(net_params, learning_rate=1e-3):
       get_params: Function opt_state -> net_params for extracting the network
         parameters from the optimization state.
   """
-  opt = trax_opt.Adam(learning_rate=learning_rate, b1=0.9, b2=0.999, eps=1e-08)
+  opt = optimizer()
   (init_slots, init_nontrainable_slots) = opt.tree_init(net_params)
   init_state = (net_params, init_slots)
 
@@ -847,511 +833,3 @@ def write_eval_reward_summaries(reward_stats_by_mode, summary_writer, epoch):
                    "[temperature %.2f] = %10.2f (+/- %.2f)",
                    epoch, reward_mode, temperature,
                    reward_stats["mean"], reward_stats["std"])
-
-
-class PPOTrainer(object):
-  """PPO trainer."""
-
-  def __init__(
-      self,
-      train_env,
-      eval_env,
-      policy_and_value_model,
-      policy_and_value_optimizer_fn,
-      policy_and_value_two_towers,
-      output_dir,
-      n_optimizer_steps,
-      print_every_optimizer_steps,
-      target_kl,
-      boundary,
-      max_timestep,
-      max_timestep_eval,
-      random_seed,
-      gamma,
-      lambda_,
-      c1,
-      c2,
-      eval_every_n,
-      done_frac_for_policy_save,
-      n_evals,
-      len_history_for_policy,
-      eval_temperatures,
-  ):
-    self._train_env = train_env
-    self._eval_env = eval_env
-    self._n_optimizer_steps = n_optimizer_steps
-    self._print_every_optimizer_steps = print_every_optimizer_steps
-    self._target_kl = target_kl
-    self._boundary = boundary
-    self._max_timestep = max_timestep
-    self._max_timestep_eval = max_timestep_eval
-    self._gamma = gamma
-    self._lambda_ = lambda_
-    self._c1 = c1
-    self._c2 = c2
-    self._eval_every_n = eval_every_n
-    self._done_frac_for_policy_save = done_frac_for_policy_save
-    self._n_evals = n_evals
-    self._len_history_for_policy = len_history_for_policy
-    self._eval_temperatures = eval_temperatures
-
-    assert isinstance(self._train_env.action_space, gym.spaces.Discrete)
-    n_actions = self._train_env.action_space.n
-
-    # Batch Observations Shape = [1, 1] + OBS, because we will eventually call
-    # policy and value networks on shape [B, T] +_OBS
-    batch_observations_shape = (1, 1) + self._train_env.observation_space.shape
-    observations_dtype = self._train_env.observation_space.dtype
-
-    self._rng = trax.get_random_number_generator_and_set_seed(random_seed)
-    self._rng, key1 = jax_random.split(self._rng, num=2)
-
-    # Initialize the policy and value network.
-    policy_and_value_net_params, policy_and_value_net_apply = (
-        policy_and_value_net(
-            rng_key=key1,
-            batch_observations_shape=batch_observations_shape,
-            observations_dtype=observations_dtype,
-            n_actions=n_actions,
-            bottom_layers_fn=policy_and_value_model,
-            two_towers=policy_and_value_two_towers,
-        )
-    )
-    self._policy_and_value_net_apply = jit(policy_and_value_net_apply)
-
-    # Maybe restore the policy params. If there is nothing to restore, then
-    # iteration = 0 and policy_and_value_net_params are returned as is.
-    restored, policy_and_value_net_params, self._epoch = (
-        maybe_restore_params(output_dir, policy_and_value_net_params))
-
-    if restored:
-      logging.info("Restored parameters from iteration [%d]", self._epoch)
-      # We should start from the next iteration.
-      self._epoch += 1
-
-    # Initialize the optimizers.
-    policy_and_value_optimizer = (
-        policy_and_value_optimizer_fn(policy_and_value_net_params))
-    (self._policy_and_value_opt_state, self._policy_and_value_opt_update,
-     self._policy_and_value_get_params) = policy_and_value_optimizer
-
-    self._output_dir = output_dir
-    gfile.makedirs(self._output_dir)
-
-    # Create summary writers and history.
-    self._train_sw = jaxboard.SummaryWriter(
-        os.path.join(self._output_dir, "train"))
-    self._timing_sw = jaxboard.SummaryWriter(
-        os.path.join(self._output_dir, "timing"))
-    self._eval_sw = jaxboard.SummaryWriter(
-        os.path.join(self._output_dir, "eval"))
-
-    self._should_reset = True
-    self._n_trajectories_done = 0
-
-    self._last_saved_at = 0
-
-  @property
-  def epoch(self):
-    return self._epoch
-
-  def train_epoch(self):
-    """Train one PPO epoch."""
-    epoch_start_time = time.time()
-
-    # Evaluate the policy.
-    policy_eval_start_time = time.time()
-    if (self._epoch + 1) % self._eval_every_n == 0:
-      self._rng, key = jax_random.split(self._rng, num=2)
-      self.evaluate()
-
-    policy_eval_time = get_time(policy_eval_start_time)
-
-    trajectory_collection_start_time = time.time()
-    logging.vlog(1, "Epoch [% 6d] collecting trajectories.", self._epoch)
-    self._rng, key = jax_random.split(self._rng)
-    trajs, n_done, timing_info = collect_trajectories(
-        self._train_env,
-        policy_fn=self._get_predictions,
-        n_trajectories=self._train_env.batch_size,
-        max_timestep=self._max_timestep,
-        rng=key,
-        len_history_for_policy=self._len_history_for_policy,
-        boundary=self._boundary,
-        reset=self._should_reset,
-    )
-    self._should_reset = False
-    trajectory_collection_time = get_time(trajectory_collection_start_time)
-
-    logging.vlog(1, "Collecting trajectories took %0.2f msec.",
-                 trajectory_collection_time)
-
-    avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
-    max_reward = max(np.sum(traj[2]) for traj in trajs)
-    min_reward = min(np.sum(traj[2]) for traj in trajs)
-
-    self._train_sw.scalar(
-        "train/reward_mean_truncated", avg_reward, step=self._epoch)
-
-    logging.vlog(1, "Rewards avg=[%0.2f], max=[%0.2f], min=[%0.2f], all=%s",
-                 avg_reward, max_reward, min_reward,
-                 [float(np.sum(traj[2])) for traj in trajs])
-
-    logging.vlog(1,
-                 "Trajectory Length average=[%0.2f], max=[%0.2f], min=[%0.2f]",
-                 float(sum(len(traj[0]) for traj in trajs)) / len(trajs),
-                 max(len(traj[0]) for traj in trajs),
-                 min(len(traj[0]) for traj in trajs))
-    logging.vlog(2, "Trajectory Lengths: %s", [len(traj[0]) for traj in trajs])
-
-    padding_start_time = time.time()
-    (_, reward_mask, padded_observations, padded_actions,
-     padded_rewards, padded_infos) = pad_trajectories(
-         trajs, boundary=self._boundary)
-    padding_time = get_time(padding_start_time)
-
-    logging.vlog(1, "Padding trajectories took %0.2f msec.",
-                 get_time(padding_start_time))
-    logging.vlog(1, "Padded Observations' shape [%s]",
-                 str(padded_observations.shape))
-    logging.vlog(1, "Padded Actions' shape [%s]", str(padded_actions.shape))
-    logging.vlog(1, "Padded Rewards' shape [%s]", str(padded_rewards.shape))
-
-    # Some assertions.
-    B, T = padded_actions.shape  # pylint: disable=invalid-name
-    assert (B, T) == padded_rewards.shape
-    assert (B, T) == reward_mask.shape
-    assert (B, T + 1) == padded_observations.shape[:2]
-    assert ((B, T + 1) + self._train_env.observation_space.shape ==
-            padded_observations.shape)
-
-    log_prob_recompute_start_time = time.time()
-    assert ("log_prob_actions" in padded_infos and
-            "value_predictions" in padded_infos)
-    # These are the actual log-probabs and value predictions seen while picking
-    # the actions.
-    actual_log_probabs_traj = padded_infos["log_prob_actions"]
-    actual_value_predictions_traj = padded_infos["value_predictions"]
-
-    assert (B, T) == actual_log_probabs_traj.shape[:2]
-    A = actual_log_probabs_traj.shape[2]  # pylint: disable=invalid-name
-    assert (B, T, 1) == actual_value_predictions_traj.shape
-
-    # TODO(afrozm): log-probabs doesn't need to be (B, T+1, A) it can do with
-    # (B, T, A), so make that change throughout.
-
-    # NOTE: We don't have the log-probabs and value-predictions for the last
-    # observation, so we re-calculate for everything, but use the original ones
-    # for all but the last time-step.
-    self._rng, key = jax_random.split(self._rng)
-    log_probabs_traj, value_predictions_traj, _ = self._get_predictions(
-        padded_observations, rng=key)
-
-    assert (B, T + 1, A) == log_probabs_traj.shape
-    assert (B, T + 1, 1) == value_predictions_traj.shape
-
-    # Concatenate the last time-step's log-probabs and value predictions to the
-    # actual log-probabs and value predictions and use those going forward.
-    log_probabs_traj = np.concatenate(
-        (actual_log_probabs_traj, log_probabs_traj[:, -1:, :]), axis=1)
-    value_predictions_traj = np.concatenate(
-        (actual_value_predictions_traj, value_predictions_traj[:, -1:, :]),
-        axis=1)
-
-    log_prob_recompute_time = get_time(log_prob_recompute_start_time)
-
-    # Compute value and ppo losses.
-    self._rng, key1 = jax_random.split(self._rng, num=2)
-    logging.vlog(2, "Starting to compute P&V loss.")
-    loss_compute_start_time = time.time()
-    cur_combined_loss, cur_ppo_loss, cur_value_loss, entropy_bonus = (
-        combined_loss(
-            self._policy_and_value_net_params,
-            log_probabs_traj,
-            value_predictions_traj,
-            self._policy_and_value_net_apply,
-            padded_observations,
-            padded_actions,
-            padded_rewards,
-            reward_mask,
-            gamma=self._gamma,
-            lambda_=self._lambda_,
-            c1=self._c1,
-            c2=self._c2,
-            rng=key1))
-    loss_compute_time = get_time(loss_compute_start_time)
-    logging.vlog(
-        1,
-        "Calculating P&V loss [%10.2f(%10.2f, %10.2f, %10.2f)] took %0.2f msec.",
-        cur_combined_loss, cur_value_loss, cur_ppo_loss, entropy_bonus,
-        get_time(loss_compute_start_time))
-
-    self._rng, key1 = jax_random.split(self._rng, num=2)
-    logging.vlog(1, "Policy and Value Optimization")
-    optimization_start_time = time.time()
-    keys = jax_random.split(key1, num=self._n_optimizer_steps)
-    for (j, key) in enumerate(keys):
-      k1, k2, k3 = jax_random.split(key, num=3)
-      t = time.time()
-      # Update the optimizer state.
-      self._policy_and_value_opt_state = policy_and_value_opt_step(
-          j,
-          self._policy_and_value_opt_state,
-          self._policy_and_value_opt_update,
-          self._policy_and_value_get_params,
-          self._policy_and_value_net_apply,
-          log_probabs_traj,
-          value_predictions_traj,
-          padded_observations,
-          padded_actions,
-          padded_rewards,
-          reward_mask,
-          c1=self._c1,
-          c2=self._c2,
-          gamma=self._gamma,
-          lambda_=self._lambda_,
-          rng=k1)
-
-      # Compute the approx KL for early stopping.
-      log_probab_actions_new, _ = self._policy_and_value_net_apply(
-          padded_observations, self._policy_and_value_net_params, rng=k2)
-
-      approx_kl = approximate_kl(log_probab_actions_new, log_probabs_traj,
-                                 reward_mask)
-
-      early_stopping = approx_kl > 1.5 * self._target_kl
-      if early_stopping:
-        logging.vlog(
-            1, "Early stopping policy and value optimization at iter: %d, "
-            "with approx_kl: %0.2f", j, approx_kl)
-        # We don't return right-away, we want the below to execute on the last
-        # iteration.
-
-      t2 = time.time()
-      if (((j + 1) % self._print_every_optimizer_steps == 0) or
-          (j == self._n_optimizer_steps - 1) or early_stopping):
-        # Compute and log the loss.
-        (loss_combined, loss_ppo, loss_value, entropy_bonus) = (
-            combined_loss(
-                self._policy_and_value_net_params,
-                log_probabs_traj,
-                value_predictions_traj,
-                self._policy_and_value_net_apply,
-                padded_observations,
-                padded_actions,
-                padded_rewards,
-                reward_mask,
-                gamma=self._gamma,
-                lambda_=self._lambda_,
-                c1=self._c1,
-                c2=self._c2,
-                rng=k3))
-        logging.vlog(1, "One Policy and Value grad desc took: %0.2f msec",
-                     get_time(t, t2))
-        logging.vlog(
-            1, "Combined Loss(value, ppo, entropy_bonus) [%10.2f] ->"
-            " [%10.2f(%10.2f,%10.2f,%10.2f)]", cur_combined_loss, loss_combined,
-            loss_value, loss_ppo, entropy_bonus)
-
-      if early_stopping:
-        break
-
-    optimization_time = get_time(optimization_start_time)
-
-    logging.vlog(
-        1, "Total Combined Loss reduction [%0.2f]%%",
-        (100 * (cur_combined_loss - loss_combined) / np.abs(cur_combined_loss)))
-
-    # Save parameters every time we see the end of at least a fraction of batch
-    # number of trajectories that are done (not completed -- completed includes
-    # truncated and done).
-    # Also don't save too frequently, enforce a minimum gap.
-    # Or if this is the last iteration.
-    policy_save_start_time = time.time()
-    self._n_trajectories_done += n_done
-    # TODO(afrozm): Refactor to trax.save_state.
-    if ((self._n_trajectories_done >=
-         self._done_frac_for_policy_save * self._train_env.batch_size) and
-        (self._epoch - self._last_saved_at > self._eval_every_n) and
-        (((self._epoch + 1) % self._eval_every_n == 0))):
-      self.save()
-    policy_save_time = get_time(policy_save_start_time)
-
-    epoch_time = get_time(epoch_start_time)
-
-    logging.info(
-        "Epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
-        " Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)]", self._epoch,
-        min_reward, max_reward, avg_reward, loss_combined, loss_value, loss_ppo,
-        entropy_bonus)
-
-    timing_dict = {
-        "epoch": epoch_time,
-        "policy_eval": policy_eval_time,
-        "trajectory_collection": trajectory_collection_time,
-        "padding": padding_time,
-        "log_prob_recompute": log_prob_recompute_time,
-        "loss_compute": loss_compute_time,
-        "optimization": optimization_time,
-        "policy_save": policy_save_time,
-    }
-
-    timing_dict.update(timing_info)
-
-    for k, v in timing_dict.items():
-      self._timing_sw.scalar("timing/%s" % k, v, step=self._epoch)
-
-    max_key_len = max(len(k) for k in timing_dict)
-    timing_info_list = [
-        "%s : % 10.2f" % (k.rjust(max_key_len + 1), v)
-        for k, v in sorted(timing_dict.items())
-    ]
-    logging.info(
-        "Epoch [% 6d], Timings: \n%s", self._epoch, "\n".join(timing_info_list))
-
-    self._epoch += 1
-
-    # Flush summary writers once in a while.
-    if (self._epoch + 1) % 1000 == 0:
-      self.flush_summaries()
-
-  def evaluate(self):
-    """Evaluate the agent."""
-    logging.vlog(1, "Epoch [% 6d] evaluating policy.", self._epoch)
-    self._rng, key = jax_random.split(self._rng, num=2)
-    reward_stats = evaluate_policy(
-        self._eval_env,
-        self._get_predictions,
-        temperatures=self._eval_temperatures,
-        max_timestep=self._max_timestep_eval,
-        n_evals=self._n_evals,
-        len_history_for_policy=self._len_history_for_policy,
-        rng=key)
-    write_eval_reward_summaries(reward_stats, self._eval_sw, epoch=self._epoch)
-
-  def save(self):
-    """Save the agent parameters."""
-    logging.vlog(1, "Epoch [% 6d] saving model.", self._epoch)
-    old_model_files = gfile.glob(
-        os.path.join(self._output_dir, "model-??????.pkl"))
-    params_file = os.path.join(self._output_dir, "model-%06d.pkl" % self._epoch)
-    with gfile.GFile(params_file, "wb") as f:
-      pickle.dump(self._policy_and_value_net_params, f)
-    # Remove the old model files.
-    for path in old_model_files:
-      gfile.remove(path)
-    # Reset this number.
-    self._n_trajectories_done = 0
-    self._last_saved_at = self._epoch
-
-  def flush_summaries(self):
-    self._train_sw.flush()
-    self._timing_sw.flush()
-    self._eval_sw.flush()
-
-  @property
-  def _policy_and_value_net_params(self):
-    return self._policy_and_value_get_params(self._policy_and_value_opt_state)
-
-  # A function to get the policy and value predictions.
-  def _get_predictions(self, observations, rng=None):
-    """Returns log-probs, value predictions and key back."""
-    key, key1 = jax_random.split(rng, num=2)
-
-    log_probs, value_preds = self._policy_and_value_net_apply(
-        observations, self._policy_and_value_net_params, rng=key1)
-
-    return log_probs, value_preds, key
-
-
-@gin.configurable(blacklist=["output_dir"])
-def training_loop(
-    train_env,
-    eval_env,
-    policy_and_value_model,
-    policy_and_value_optimizer_fn,
-    output_dir,
-    policy_and_value_two_towers=False,
-    epochs=EPOCHS,
-    n_optimizer_steps=N_OPTIMIZER_STEPS,
-    print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
-    target_kl=0.01,
-    boundary=20,
-    max_timestep=None,
-    max_timestep_eval=20000,
-    random_seed=None,
-    gamma=GAMMA,
-    lambda_=LAMBDA,
-    c1=1.0,
-    c2=0.01,
-    eval_every_n=1000,
-    done_frac_for_policy_save=0.5,
-    n_evals=1,
-    len_history_for_policy=4,
-    eval_temperatures=(1.0, 0.5),
-):
-  """Runs the training loop for PPO, with fixed policy and value nets.
-
-  Args:
-    train_env: gym.Env to use for training.
-    eval_env: gym.Env to use for evaluation.
-    policy_and_value_model: Function defining the policy and value network,
-      without the policy and value heads.
-    policy_and_value_optimizer_fn: Function defining the optimizer.
-    output_dir: Output dir.
-    policy_and_value_two_towers: Whether to use two separate models as the
-      policy and value networks. If False, share their parameters.
-    epochs: Number of epochs to run for.
-    n_optimizer_steps: Number of optimizer steps.
-    print_every_optimizer_steps: How often to log during the policy optimization
-      process.
-    target_kl: Policy iteration early stopping. Set to infinity to disable early
-      stopping.
-    boundary: We pad trajectories at integer multiples of this number.
-    max_timestep: If set to an integer, maximum number of time-steps in
-      a trajectory. Used in the collect procedure.
-    max_timestep_eval: If set to an integer, maximum number of time-steps in an
-      evaluation trajectory. Used in the collect procedure.
-    random_seed: Random seed.
-    gamma: Reward discount factor.
-    lambda_: N-step TD-error discount factor in GAE.
-    c1: Value loss coefficient.
-    c2: Entropy loss coefficient.
-    eval_every_n: How frequently to eval the policy.
-    done_frac_for_policy_save: Fraction of the trajectories that should be done
-      to checkpoint the policy.
-    n_evals: Number of times to evaluate.
-    len_history_for_policy: How much of history to give to the policy.
-    eval_temperatures: Sequence of temperatures to try for categorical sampling
-      during evaluation.
-  """
-  trainer = PPOTrainer(
-      train_env=train_env,
-      eval_env=eval_env,
-      policy_and_value_model=policy_and_value_model,
-      policy_and_value_optimizer_fn=policy_and_value_optimizer_fn,
-      policy_and_value_two_towers=policy_and_value_two_towers,
-      output_dir=output_dir,
-      n_optimizer_steps=n_optimizer_steps,
-      print_every_optimizer_steps=print_every_optimizer_steps,
-      target_kl=target_kl,
-      boundary=boundary,
-      max_timestep=max_timestep,
-      max_timestep_eval=max_timestep_eval,
-      random_seed=random_seed,
-      gamma=gamma,
-      lambda_=lambda_,
-      c1=c1,
-      c2=c2,
-      eval_every_n=eval_every_n,
-      done_frac_for_policy_save=done_frac_for_policy_save,
-      n_evals=n_evals,
-      len_history_for_policy=len_history_for_policy,
-      eval_temperatures=eval_temperatures,
-  )
-  logging.info("Starting the PPO training loop.")
-  for _ in range(trainer.epoch, epochs):
-    trainer.train_epoch()
-  trainer.save()
-  trainer.evaluate()
-  trainer.flush_summaries()
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
new file mode 100644
index 000000000..bbd9b2baa
--- /dev/null
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -0,0 +1,495 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""PPO trainer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+import os
+import time
+
+from absl import logging
+import cloudpickle as pickle
+import gym
+from jax import jit
+from jax import numpy as np
+from jax import random as jax_random
+from tensor2tensor.trax import jaxboard
+from tensor2tensor.trax import models as trax_models
+from tensor2tensor.trax import optimizers as trax_opt
+from tensor2tensor.trax import trax
+from tensor2tensor.trax.rl import base_trainer
+from tensor2tensor.trax.rl import ppo
+from tensorflow.io import gfile
+
+
+DEBUG_LOGGING = False
+GAMMA = 0.99
+LAMBDA = 0.95
+EPSILON = 0.1
+EPOCHS = 50  # 100
+N_OPTIMIZER_STEPS = 100
+PRINT_EVERY_OPTIMIZER_STEP = 20
+BATCH_TRAJECTORIES = 32
+
+
+class PPO(base_trainer.BaseTrainer):
+  """PPO trainer."""
+
+  def __init__(
+      self,
+      train_env,
+      eval_env,
+      output_dir,
+      policy_and_value_model=trax_models.FrameStackMLP,
+      policy_and_value_optimizer=functools.partial(
+          trax_opt.Adam, learning_rate=1e-3),
+      policy_and_value_two_towers=False,
+      n_optimizer_steps=N_OPTIMIZER_STEPS,
+      print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
+      target_kl=0.01,
+      boundary=20,
+      max_timestep=None,
+      max_timestep_eval=20000,
+      random_seed=None,
+      gamma=GAMMA,
+      lambda_=LAMBDA,
+      c1=1.0,
+      c2=0.01,
+      eval_every_n=1000,
+      done_frac_for_policy_save=0.5,
+      n_evals=1,
+      len_history_for_policy=4,
+      eval_temperatures=(1.0, 0.5),
+  ):
+    """Creates the PPO trainer.
+
+    Args:
+      train_env: gym.Env to use for training.
+      eval_env: gym.Env to use for evaluation.
+      output_dir: Output dir.
+      policy_and_value_model: Function defining the policy and value network,
+        without the policy and value heads.
+      policy_and_value_optimizer: Function defining the optimizer.
+      policy_and_value_two_towers: Whether to use two separate models as the
+        policy and value networks. If False, share their parameters.
+      n_optimizer_steps: Number of optimizer steps.
+      print_every_optimizer_steps: How often to log during the policy
+        optimization process.
+      target_kl: Policy iteration early stopping. Set to infinity to disable
+        early stopping.
+      boundary: We pad trajectories at integer multiples of this number.
+      max_timestep: If set to an integer, maximum number of time-steps in
+        a trajectory. Used in the collect procedure.
+      max_timestep_eval: If set to an integer, maximum number of time-steps in
+        an evaluation trajectory. Used in the collect procedure.
+      random_seed: Random seed.
+      gamma: Reward discount factor.
+      lambda_: N-step TD-error discount factor in GAE.
+      c1: Value loss coefficient.
+      c2: Entropy loss coefficient.
+      eval_every_n: How frequently to eval the policy.
+      done_frac_for_policy_save: Fraction of the trajectories that should be
+        done to checkpoint the policy.
+      n_evals: Number of times to evaluate.
+      len_history_for_policy: How much of history to give to the policy.
+      eval_temperatures: Sequence of temperatures to try for categorical
+        sampling during evaluation.
+    """
+    super(PPO, self).__init__(train_env, eval_env, output_dir)
+
+    self._n_optimizer_steps = n_optimizer_steps
+    self._print_every_optimizer_steps = print_every_optimizer_steps
+    self._target_kl = target_kl
+    self._boundary = boundary
+    self._max_timestep = max_timestep
+    self._max_timestep_eval = max_timestep_eval
+    self._gamma = gamma
+    self._lambda_ = lambda_
+    self._c1 = c1
+    self._c2 = c2
+    self._eval_every_n = eval_every_n
+    self._done_frac_for_policy_save = done_frac_for_policy_save
+    self._n_evals = n_evals
+    self._len_history_for_policy = len_history_for_policy
+    self._eval_temperatures = eval_temperatures
+
+    assert isinstance(self._train_env.action_space, gym.spaces.Discrete)
+    n_actions = self._train_env.action_space.n
+
+    # Batch Observations Shape = [1, 1] + OBS, because we will eventually call
+    # policy and value networks on shape [B, T] +_OBS
+    batch_observations_shape = (1, 1) + self._train_env.observation_space.shape
+    observations_dtype = self._train_env.observation_space.dtype
+
+    self._rng = trax.get_random_number_generator_and_set_seed(random_seed)
+    self._rng, key1 = jax_random.split(self._rng, num=2)
+
+    # Initialize the policy and value network.
+    policy_and_value_net_params, policy_and_value_net_apply = (
+        ppo.policy_and_value_net(
+            rng_key=key1,
+            batch_observations_shape=batch_observations_shape,
+            observations_dtype=observations_dtype,
+            n_actions=n_actions,
+            bottom_layers_fn=policy_and_value_model,
+            two_towers=policy_and_value_two_towers,
+        )
+    )
+    self._policy_and_value_net_apply = jit(policy_and_value_net_apply)
+
+    # Maybe restore the policy params. If there is nothing to restore, then
+    # iteration = 0 and policy_and_value_net_params are returned as is.
+    restored, policy_and_value_net_params, self._epoch = (
+        ppo.maybe_restore_params(output_dir, policy_and_value_net_params))
+
+    if restored:
+      logging.info("Restored parameters from iteration [%d]", self._epoch)
+      # We should start from the next iteration.
+      self._epoch += 1
+
+    # Initialize the optimizer.
+    (self._policy_and_value_opt_state, self._policy_and_value_opt_update,
+     self._policy_and_value_get_params) = ppo.optimizer_fn(
+         policy_and_value_optimizer, policy_and_value_net_params)
+
+    # Create summary writers and history.
+    self._train_sw = jaxboard.SummaryWriter(
+        os.path.join(self._output_dir, "train"))
+    self._timing_sw = jaxboard.SummaryWriter(
+        os.path.join(self._output_dir, "timing"))
+    self._eval_sw = jaxboard.SummaryWriter(
+        os.path.join(self._output_dir, "eval"))
+
+    self._should_reset = True
+    self._n_trajectories_done = 0
+
+    self._last_saved_at = 0
+
+  @property
+  def epoch(self):
+    return self._epoch
+
+  def train_epoch(self):
+    """Train one PPO epoch."""
+    epoch_start_time = time.time()
+
+    # Evaluate the policy.
+    policy_eval_start_time = time.time()
+    if (self._epoch + 1) % self._eval_every_n == 0:
+      self._rng, key = jax_random.split(self._rng, num=2)
+      self.evaluate()
+
+    policy_eval_time = ppo.get_time(policy_eval_start_time)
+
+    trajectory_collection_start_time = time.time()
+    logging.vlog(1, "Epoch [% 6d] collecting trajectories.", self._epoch)
+    self._rng, key = jax_random.split(self._rng)
+    trajs, n_done, timing_info = ppo.collect_trajectories(
+        self._train_env,
+        policy_fn=self._get_predictions,
+        n_trajectories=self._train_env.batch_size,
+        max_timestep=self._max_timestep,
+        rng=key,
+        len_history_for_policy=self._len_history_for_policy,
+        boundary=self._boundary,
+        reset=self._should_reset,
+    )
+    self._should_reset = False
+    trajectory_collection_time = ppo.get_time(trajectory_collection_start_time)
+
+    logging.vlog(1, "Collecting trajectories took %0.2f msec.",
+                 trajectory_collection_time)
+
+    avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
+    max_reward = max(np.sum(traj[2]) for traj in trajs)
+    min_reward = min(np.sum(traj[2]) for traj in trajs)
+
+    self._train_sw.scalar(
+        "train/reward_mean_truncated", avg_reward, step=self._epoch)
+
+    logging.vlog(1, "Rewards avg=[%0.2f], max=[%0.2f], min=[%0.2f], all=%s",
+                 avg_reward, max_reward, min_reward,
+                 [float(np.sum(traj[2])) for traj in trajs])
+
+    logging.vlog(1,
+                 "Trajectory Length average=[%0.2f], max=[%0.2f], min=[%0.2f]",
+                 float(sum(len(traj[0]) for traj in trajs)) / len(trajs),
+                 max(len(traj[0]) for traj in trajs),
+                 min(len(traj[0]) for traj in trajs))
+    logging.vlog(2, "Trajectory Lengths: %s", [len(traj[0]) for traj in trajs])
+
+    padding_start_time = time.time()
+    (_, reward_mask, padded_observations, padded_actions,
+     padded_rewards, padded_infos) = ppo.pad_trajectories(
+         trajs, boundary=self._boundary)
+    padding_time = ppo.get_time(padding_start_time)
+
+    logging.vlog(1, "Padding trajectories took %0.2f msec.",
+                 ppo.get_time(padding_start_time))
+    logging.vlog(1, "Padded Observations' shape [%s]",
+                 str(padded_observations.shape))
+    logging.vlog(1, "Padded Actions' shape [%s]", str(padded_actions.shape))
+    logging.vlog(1, "Padded Rewards' shape [%s]", str(padded_rewards.shape))
+
+    # Some assertions.
+    B, T = padded_actions.shape  # pylint: disable=invalid-name
+    assert (B, T) == padded_rewards.shape
+    assert (B, T) == reward_mask.shape
+    assert (B, T + 1) == padded_observations.shape[:2]
+    assert ((B, T + 1) + self._train_env.observation_space.shape ==
+            padded_observations.shape)
+
+    log_prob_recompute_start_time = time.time()
+    assert ("log_prob_actions" in padded_infos and
+            "value_predictions" in padded_infos)
+    # These are the actual log-probabs and value predictions seen while picking
+    # the actions.
+    actual_log_probabs_traj = padded_infos["log_prob_actions"]
+    actual_value_predictions_traj = padded_infos["value_predictions"]
+
+    assert (B, T) == actual_log_probabs_traj.shape[:2]
+    A = actual_log_probabs_traj.shape[2]  # pylint: disable=invalid-name
+    assert (B, T, 1) == actual_value_predictions_traj.shape
+
+    # TODO(afrozm): log-probabs doesn't need to be (B, T+1, A) it can do with
+    # (B, T, A), so make that change throughout.
+
+    # NOTE: We don't have the log-probabs and value-predictions for the last
+    # observation, so we re-calculate for everything, but use the original ones
+    # for all but the last time-step.
+    self._rng, key = jax_random.split(self._rng)
+    log_probabs_traj, value_predictions_traj, _ = self._get_predictions(
+        padded_observations, rng=key)
+
+    assert (B, T + 1, A) == log_probabs_traj.shape
+    assert (B, T + 1, 1) == value_predictions_traj.shape
+
+    # Concatenate the last time-step's log-probabs and value predictions to the
+    # actual log-probabs and value predictions and use those going forward.
+    log_probabs_traj = np.concatenate(
+        (actual_log_probabs_traj, log_probabs_traj[:, -1:, :]), axis=1)
+    value_predictions_traj = np.concatenate(
+        (actual_value_predictions_traj, value_predictions_traj[:, -1:, :]),
+        axis=1)
+
+    log_prob_recompute_time = ppo.get_time(log_prob_recompute_start_time)
+
+    # Compute value and ppo losses.
+    self._rng, key1 = jax_random.split(self._rng, num=2)
+    logging.vlog(2, "Starting to compute P&V loss.")
+    loss_compute_start_time = time.time()
+    cur_combined_loss, cur_ppo_loss, cur_value_loss, entropy_bonus = (
+        ppo.combined_loss(
+            self._policy_and_value_net_params,
+            log_probabs_traj,
+            value_predictions_traj,
+            self._policy_and_value_net_apply,
+            padded_observations,
+            padded_actions,
+            padded_rewards,
+            reward_mask,
+            gamma=self._gamma,
+            lambda_=self._lambda_,
+            c1=self._c1,
+            c2=self._c2,
+            rng=key1))
+    loss_compute_time = ppo.get_time(loss_compute_start_time)
+    logging.vlog(
+        1,
+        "Calculating P&V loss [%10.2f(%10.2f, %10.2f, %10.2f)] took %0.2f msec.",
+        cur_combined_loss, cur_value_loss, cur_ppo_loss, entropy_bonus,
+        ppo.get_time(loss_compute_start_time))
+
+    self._rng, key1 = jax_random.split(self._rng, num=2)
+    logging.vlog(1, "Policy and Value Optimization")
+    optimization_start_time = time.time()
+    keys = jax_random.split(key1, num=self._n_optimizer_steps)
+    for (j, key) in enumerate(keys):
+      k1, k2, k3 = jax_random.split(key, num=3)
+      t = time.time()
+      # Update the optimizer state.
+      self._policy_and_value_opt_state = ppo.policy_and_value_opt_step(
+          j,
+          self._policy_and_value_opt_state,
+          self._policy_and_value_opt_update,
+          self._policy_and_value_get_params,
+          self._policy_and_value_net_apply,
+          log_probabs_traj,
+          value_predictions_traj,
+          padded_observations,
+          padded_actions,
+          padded_rewards,
+          reward_mask,
+          c1=self._c1,
+          c2=self._c2,
+          gamma=self._gamma,
+          lambda_=self._lambda_,
+          rng=k1)
+
+      # Compute the approx KL for early stopping.
+      log_probab_actions_new, _ = self._policy_and_value_net_apply(
+          padded_observations, self._policy_and_value_net_params, rng=k2)
+
+      approx_kl = ppo.approximate_kl(log_probab_actions_new, log_probabs_traj,
+                                     reward_mask)
+
+      early_stopping = approx_kl > 1.5 * self._target_kl
+      if early_stopping:
+        logging.vlog(
+            1, "Early stopping policy and value optimization at iter: %d, "
+            "with approx_kl: %0.2f", j, approx_kl)
+        # We don't return right-away, we want the below to execute on the last
+        # iteration.
+
+      t2 = time.time()
+      if (((j + 1) % self._print_every_optimizer_steps == 0) or
+          (j == self._n_optimizer_steps - 1) or early_stopping):
+        # Compute and log the loss.
+        (loss_combined, loss_ppo, loss_value, entropy_bonus) = (
+            ppo.combined_loss(
+                self._policy_and_value_net_params,
+                log_probabs_traj,
+                value_predictions_traj,
+                self._policy_and_value_net_apply,
+                padded_observations,
+                padded_actions,
+                padded_rewards,
+                reward_mask,
+                gamma=self._gamma,
+                lambda_=self._lambda_,
+                c1=self._c1,
+                c2=self._c2,
+                rng=k3))
+        logging.vlog(1, "One Policy and Value grad desc took: %0.2f msec",
+                     ppo.get_time(t, t2))
+        logging.vlog(
+            1, "Combined Loss(value, ppo, entropy_bonus) [%10.2f] ->"
+            " [%10.2f(%10.2f,%10.2f,%10.2f)]", cur_combined_loss, loss_combined,
+            loss_value, loss_ppo, entropy_bonus)
+
+      if early_stopping:
+        break
+
+    optimization_time = ppo.get_time(optimization_start_time)
+
+    logging.vlog(
+        1, "Total Combined Loss reduction [%0.2f]%%",
+        (100 * (cur_combined_loss - loss_combined) / np.abs(cur_combined_loss)))
+
+    # Save parameters every time we see the end of at least a fraction of batch
+    # number of trajectories that are done (not completed -- completed includes
+    # truncated and done).
+    # Also don't save too frequently, enforce a minimum gap.
+    # Or if this is the last iteration.
+    policy_save_start_time = time.time()
+    self._n_trajectories_done += n_done
+    # TODO(afrozm): Refactor to trax.save_state.
+    if ((self._n_trajectories_done >=
+         self._done_frac_for_policy_save * self._train_env.batch_size) and
+        (self._epoch - self._last_saved_at > self._eval_every_n) and
+        (((self._epoch + 1) % self._eval_every_n == 0))):
+      self.save()
+    policy_save_time = ppo.get_time(policy_save_start_time)
+
+    epoch_time = ppo.get_time(epoch_start_time)
+
+    logging.info(
+        "Epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
+        " Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)]", self._epoch,
+        min_reward, max_reward, avg_reward, loss_combined, loss_value, loss_ppo,
+        entropy_bonus)
+
+    timing_dict = {
+        "epoch": epoch_time,
+        "policy_eval": policy_eval_time,
+        "trajectory_collection": trajectory_collection_time,
+        "padding": padding_time,
+        "log_prob_recompute": log_prob_recompute_time,
+        "loss_compute": loss_compute_time,
+        "optimization": optimization_time,
+        "policy_save": policy_save_time,
+    }
+
+    timing_dict.update(timing_info)
+
+    for k, v in timing_dict.items():
+      self._timing_sw.scalar("timing/%s" % k, v, step=self._epoch)
+
+    max_key_len = max(len(k) for k in timing_dict)
+    timing_info_list = [
+        "%s : % 10.2f" % (k.rjust(max_key_len + 1), v)
+        for k, v in sorted(timing_dict.items())
+    ]
+    logging.info(
+        "Epoch [% 6d], Timings: \n%s", self._epoch, "\n".join(timing_info_list))
+
+    self._epoch += 1
+
+    # Flush summary writers once in a while.
+    if (self._epoch + 1) % 1000 == 0:
+      self.flush_summaries()
+
+  def evaluate(self):
+    """Evaluate the agent."""
+    logging.vlog(1, "Epoch [% 6d] evaluating policy.", self._epoch)
+    self._rng, key = jax_random.split(self._rng, num=2)
+    reward_stats = ppo.evaluate_policy(
+        self._eval_env,
+        self._get_predictions,
+        temperatures=self._eval_temperatures,
+        max_timestep=self._max_timestep_eval,
+        n_evals=self._n_evals,
+        len_history_for_policy=self._len_history_for_policy,
+        rng=key)
+    ppo.write_eval_reward_summaries(
+        reward_stats, self._eval_sw, epoch=self._epoch)
+
+  def save(self):
+    """Save the agent parameters."""
+    logging.vlog(1, "Epoch [% 6d] saving model.", self._epoch)
+    old_model_files = gfile.glob(
+        os.path.join(self._output_dir, "model-??????.pkl"))
+    params_file = os.path.join(self._output_dir, "model-%06d.pkl" % self._epoch)
+    with gfile.GFile(params_file, "wb") as f:
+      pickle.dump(self._policy_and_value_net_params, f)
+    # Remove the old model files.
+    for path in old_model_files:
+      gfile.remove(path)
+    # Reset this number.
+    self._n_trajectories_done = 0
+    self._last_saved_at = self._epoch
+
+  def flush_summaries(self):
+    self._train_sw.flush()
+    self._timing_sw.flush()
+    self._eval_sw.flush()
+
+  @property
+  def _policy_and_value_net_params(self):
+    return self._policy_and_value_get_params(self._policy_and_value_opt_state)
+
+  # A function to get the policy and value predictions.
+  def _get_predictions(self, observations, rng=None):
+    """Returns log-probs, value predictions and key back."""
+    key, key1 = jax_random.split(rng, num=2)
+
+    log_probs, value_preds = self._policy_and_value_net_apply(
+        observations, self._policy_and_value_net_params, rng=key1)
+
+    return log_probs, value_preds, key
diff --git a/tensor2tensor/trax/rl/ppo_training_loop_test.py b/tensor2tensor/trax/rl/ppo_trainer_test.py
similarity index 97%
rename from tensor2tensor/trax/rl/ppo_training_loop_test.py
rename to tensor2tensor/trax/rl/ppo_trainer_test.py
index 7bfea65fc..1d8e2c470 100644
--- a/tensor2tensor/trax/rl/ppo_training_loop_test.py
+++ b/tensor2tensor/trax/rl/ppo_trainer_test.py
@@ -38,13 +38,13 @@
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.rl import envs  # pylint: disable=unused-import
-from tensor2tensor.trax.rl import ppo
+from tensor2tensor.trax.rl import ppo_trainer
 from tensor2tensor.trax.rl import simulated_env_problem
 from tensorflow import test
 from tensorflow.io import gfile
 
 
-class PpoTrainingLoopTest(test.TestCase):
+class PpoTrainerTest(test.TestCase):
 
   def get_wrapped_env(self, name="CartPole-v0", max_episode_steps=2):
     wrapper_fn = functools.partial(
@@ -74,17 +74,16 @@ def _run_training_loop(self, train_env, eval_env, output_dir, model=None):
       model = lambda: [layers.Dense(1)]
     n_epochs = 2
     # Run the training loop.
-    ppo.training_loop(
+    trainer = ppo_trainer.PPO(
         train_env=train_env,
         eval_env=eval_env,
-        epochs=n_epochs,
         policy_and_value_model=model,
-        policy_and_value_optimizer_fn=ppo.optimizer_fn,
         n_optimizer_steps=1,
         output_dir=output_dir,
         random_seed=0,
         boundary=2,
     )
+    trainer.training_loop(n_epochs=n_epochs)
 
   def test_training_loop_cartpole(self):
     with self.tmp_dir() as output_dir:
diff --git a/tensor2tensor/trax/rl/trainers.py b/tensor2tensor/trax/rl/trainers.py
new file mode 100644
index 000000000..4be7e7318
--- /dev/null
+++ b/tensor2tensor/trax/rl/trainers.py
@@ -0,0 +1,35 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Trainers defined in trax.rl."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gin
+
+from tensor2tensor.trax.rl import ppo_trainer
+
+
+# Ginify
+def trainer_configure(*args, **kwargs):
+  kwargs["module"] = "trax.rl.trainers"
+  kwargs["blacklist"] = ["train_env", "eval_env", "output_dir"]
+  return gin.external_configurable(*args, **kwargs)
+
+
+# pylint: disable=invalid-name
+PPO = trainer_configure(ppo_trainer.PPO)
diff --git a/tensor2tensor/trax/rl_trainer.py b/tensor2tensor/trax/rl_trainer.py
index f49fe178c..39833c5e0 100644
--- a/tensor2tensor/trax/rl_trainer.py
+++ b/tensor2tensor/trax/rl_trainer.py
@@ -19,21 +19,13 @@
 
 Sample invocation:
 
-ENV_PROBLEM_NAME=Acrobot-v1
-BATCH_SIZE=32
-EPOCHS=100
-RANDOM_SEED=0
-BOUNDARY=100
-
+TRAIN_BATCH_SIZE=32
 python trax/rl_trainer.py \
-  --env_problem_name=${ENV_PROBLEM_NAME} \
-  --batch_size=${BATCH_SIZE} \
-  --config=ppo.training_loop.epochs=${EPOCHS} \
-  --config=ppo.training_loop.random_seed=${RANDOM_SEED} \
-  --config=ppo.training_loop.boundary=${BOUNDARY} \
+  --config_file=trax/rl/configs/acrobot.gin \
+  --train_batch_size=${TRAIN_BATCH_SIZE} \
   --output_dir=${HOME}/ppo_acrobot \
   --vmodule=*/tensor2tensor/*=1 \
-  --alsologtostderr \
+  --alsologtostderr
 """
 
 from __future__ import absolute_import
@@ -57,40 +49,16 @@
 from tensor2tensor.rl import gym_utils
 from tensor2tensor.rl.google import atari_utils  # GOOGLE-INTERNAL:
 from tensor2tensor.trax.rl import envs as rl_envs  # pylint: disable=unused-import
-from tensor2tensor.trax.rl import ppo
+from tensor2tensor.trax.rl import trainers as rl_trainers
 
 
 FLAGS = flags.FLAGS
 
-flags.DEFINE_string("env_problem_name", None, "Name of the EnvProblem to make.")
-
-# -1: returns env as is.
-# None: unwraps and returns without TimeLimit wrapper.
-# Any other number: imposes this restriction.
-flags.DEFINE_string(
-    "max_timestep", None,
-    "If set to an integer, maximum number of time-steps in a "
-    "trajectory. The bare env is wrapped with TimeLimit wrapper.")
-
 flags.DEFINE_boolean(
     "jax_debug_nans", False,
     "Setting to true will help to debug nans and disable jit.")
 flags.DEFINE_boolean("disable_jit", False, "Setting to true will disable jit.")
 
-# If resize is True, then we create RenderedEnvProblem, so this has to be set to
-# False for something like CartPole.
-flags.DEFINE_boolean("resize", False, "If true, resize the game frame")
-flags.DEFINE_integer("resized_height", 105, "Resized height of the game frame.")
-flags.DEFINE_integer("resized_width", 80, "Resized width of the game frame.")
-
-flags.DEFINE_bool(
-    "two_towers", True,
-    "In the combined network case should we make one tower or"
-    "two.")
-
-# Learning rate of the combined net, policy net and value net.
-flags.DEFINE_float("learning_rate", 1e-3, "Learning rate.")
-
 flags.DEFINE_string("output_dir", "", "Output dir.")
 flags.DEFINE_multi_string("config_file", None,
                           "Configuration file with parameters (.gin).")
@@ -98,11 +66,9 @@
                           "Configuration parameters (gin string).")
 flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
 flags.DEFINE_bool("xm", False, "Copy atari roms?")
-flags.DEFINE_integer("batch_size", 32,
+flags.DEFINE_integer("train_batch_size", 32,
                      "Number of parallel environments during training.")
 flags.DEFINE_integer("eval_batch_size", 4, "Batch size for evaluation.")
-flags.DEFINE_bool("clip_rewards", True,
-                  "Whether to clip and discretize the rewards.")
 flags.DEFINE_boolean("parallelize_envs", False,
                      "If true, sets parallelism to number of cpu cores.")
 
@@ -112,10 +78,11 @@
 flags.DEFINE_string("eval_server_bns", "", "Eval Server's BNS.")
 
 
-def make_env(batch_size=8, **env_kwargs):
+def make_env(name, batch_size, max_timestep, clip_rewards, rendered_env,
+             resize_dims, **env_kwargs):
   """Creates the env."""
 
-  if FLAGS.clip_rewards:
+  if clip_rewards:
     env_kwargs.update({"reward_range": (-1, 1), "discrete_rewards": True})
   else:
     env_kwargs.update({"discrete_rewards": False})
@@ -124,44 +91,63 @@ def make_env(batch_size=8, **env_kwargs):
   parallelism = multiprocessing.cpu_count() if FLAGS.parallelize_envs else 1
 
   # No resizing needed, so let's be on the normal EnvProblem.
-  if not FLAGS.resize:  # None or False
+  if not rendered_env:
     return gym_env_problem.GymEnvProblem(
-        base_env_name=FLAGS.env_problem_name,
+        base_env_name=name,
         batch_size=batch_size,
         parallelism=parallelism,
         **env_kwargs)
 
-  max_timestep = None
-  try:
-    max_timestep = int(FLAGS.max_timestep)
-  except Exception:  # pylint: disable=broad-except
-    pass
-
   wrapper_fn = functools.partial(
       gym_utils.gym_env_wrapper, **{
           "rl_env_max_episode_steps": max_timestep,
           "maxskip_env": True,
           "rendered_env": True,
-          "rendered_env_resize_to": (FLAGS.resized_height, FLAGS.resized_width),
+          "rendered_env_resize_to": resize_dims,
           "sticky_actions": False,
           "output_dtype": onp.int32 if FLAGS.use_tpu else None,
       })
 
   return rendered_env_problem.RenderedEnvProblem(
-      base_env_name=FLAGS.env_problem_name,
+      base_env_name=name,
       batch_size=batch_size,
       parallelism=parallelism,
       env_wrapper_fn=wrapper_fn,
       **env_kwargs)
 
 
-def get_optimizer_fn(learning_rate):
-  return functools.partial(ppo.optimizer_fn, learning_rate=learning_rate)
-
-
-def main(argv):
-  del argv
-  logging.info("Starting PPO Main.")
+# Not just "train" to avoid a conflict with trax.train in GIN files.
+@gin.configurable(blacklist=[
+    "output_dir", "train_batch_size", "eval_batch_size"])
+def train_rl(
+    output_dir,
+    train_batch_size,
+    eval_batch_size,
+    env_name="Acrobot-v1",
+    max_timestep=None,
+    clip_rewards=False,
+    rendered_env=False,
+    resize_dims=(105, 80),
+    trainer_class=rl_trainers.PPO,
+    n_epochs=10000,
+):
+  """Train the RL agent.
+
+  Args:
+    output_dir: Output directory.
+    train_batch_size: Number of parallel environments to use for training.
+    eval_batch_size: Number of parallel environments to use for evaluation.
+    env_name: Name of the environment.
+    max_timestep: Int or None, the maximum number of timesteps in a trajectory.
+      The environment is wrapped in a TimeLimit wrapper.
+    clip_rewards: Whether to clip and discretize the rewards.
+    rendered_env: Whether the environment has visual input. If so,
+      a RenderedEnvProblem will be used.
+    resize_dims: Pair (height, width), dimensions to resize the visual
+      observations to.
+    trainer_class: RLTrainer class to use.
+    n_epochs: Number epochs to run the training for.
+  """
 
   if FLAGS.jax_debug_nans:
     config.update("jax_debug_nans", True)
@@ -172,49 +158,54 @@ def main(argv):
     config.update("jax_platform_name", "gpu")
 
 
-  gin_configs = FLAGS.config or []
-  gin.parse_config_files_and_bindings(FLAGS.config_file, gin_configs)
-
   # TODO(pkozakowski): Find a better way to determine this.
   train_env_kwargs = {}
   eval_env_kwargs = {}
-  if "OnlineTuneEnv" in FLAGS.env_problem_name:
+  if "OnlineTuneEnv" in env_name:
     # TODO(pkozakowski): Separate env output dirs by train/eval and epoch.
     train_env_kwargs = {
-        "output_dir": os.path.join(FLAGS.output_dir, "envs/train")
+        "output_dir": os.path.join(output_dir, "envs/train")
     }
     eval_env_kwargs = {
-        "output_dir": os.path.join(FLAGS.output_dir, "envs/eval")
+        "output_dir": os.path.join(output_dir, "envs/eval")
     }
 
-  if "ClientEnv" in FLAGS.env_problem_name:
+  if "ClientEnv" in env_name:
     train_env_kwargs["per_env_kwargs"] = [{
         "remote_env_address": os.path.join(FLAGS.train_server_bns, str(replica))
-    } for replica in range(FLAGS.batch_size)]
+    } for replica in range(train_batch_size)]
 
     eval_env_kwargs["per_env_kwargs"] = [{
         "remote_env_address": os.path.join(FLAGS.eval_server_bns, str(replica))
-    } for replica in range(FLAGS.eval_batch_size)]
+    } for replica in range(eval_batch_size)]
+
+  common_env_kwargs = {
+      "name": env_name,
+      "max_timestep": max_timestep,
+      "clip_rewards": clip_rewards,
+      "rendered_env": rendered_env,
+      "resize_dims": resize_dims,
+  }
+  train_env_kwargs.update(common_env_kwargs)
+  eval_env_kwargs.update(common_env_kwargs)
 
   # Make an env here.
-  train_env = make_env(batch_size=FLAGS.batch_size, **train_env_kwargs)
+  train_env = make_env(batch_size=train_batch_size, **train_env_kwargs)
   assert train_env
 
-  eval_env = make_env(batch_size=FLAGS.eval_batch_size, **eval_env_kwargs)
+  eval_env = make_env(batch_size=eval_batch_size, **eval_env_kwargs)
   assert eval_env
 
   def run_training_loop():
     """Runs the training loop."""
     logging.info("Starting the training loop.")
 
-    policy_and_value_optimizer_fn = get_optimizer_fn(FLAGS.learning_rate)
-
-    ppo.training_loop(
+    trainer = trainer_class(
         output_dir=FLAGS.output_dir,
         train_env=train_env,
         eval_env=eval_env,
-        policy_and_value_optimizer_fn=policy_and_value_optimizer_fn,
     )
+    trainer.training_loop(n_epochs=n_epochs)
 
   if FLAGS.jax_debug_nans or FLAGS.disable_jit:
     with jax.disable_jit():
@@ -223,5 +214,19 @@ def run_training_loop():
     run_training_loop()
 
 
+def main(argv):
+  del argv
+  logging.info("Starting RL training.")
+
+  gin_configs = FLAGS.config or []
+  gin.parse_config_files_and_bindings(FLAGS.config_file, gin_configs)
+
+  train_rl(
+      output_dir=FLAGS.output_dir,
+      train_batch_size=FLAGS.train_batch_size,
+      eval_batch_size=FLAGS.eval_batch_size,
+  )
+
+
 if __name__ == "__main__":
   app.run(main)

From 8f1471865038f2a2226193c7211609b580d77453 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 19 Aug 2019 14:30:31 -0700
Subject: [PATCH 2290/2720] gRPC/loas2 isn't supported externally, switch to
 insecure channel/port.

PiperOrigin-RevId: 264242961
---
 tensor2tensor/envs/client_env.py   | 4 +---
 tensor2tensor/envs/server_utils.py | 5 +----
 2 files changed, 2 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/envs/client_env.py b/tensor2tensor/envs/client_env.py
index ea7e3f61a..46e57e136 100644
--- a/tensor2tensor/envs/client_env.py
+++ b/tensor2tensor/envs/client_env.py
@@ -21,7 +21,6 @@
 
 from absl import logging
 import grpc
-from grpc import loas2
 import gym
 import numpy as np
 from tensor2tensor.envs import env_service_pb2
@@ -34,8 +33,7 @@ class ClientEnv(gym.Env):
 
   @staticmethod
   def create_channel(remote_env_address):
-    return grpc.secure_channel(remote_env_address,
-                               loas2.loas2_channel_credentials())
+    return grpc.insecure_channel(remote_env_address)  # pylint: disable=unreachable
 
   @staticmethod
   def run_step(stub, discrete_action):
diff --git a/tensor2tensor/envs/server_utils.py b/tensor2tensor/envs/server_utils.py
index 3df60b1d7..825cb52a3 100644
--- a/tensor2tensor/envs/server_utils.py
+++ b/tensor2tensor/envs/server_utils.py
@@ -23,7 +23,6 @@
 from absl import logging
 from concurrent import futures
 import grpc
-from grpc import loas2
 
 from tensor2tensor.envs import env_service_pb2_grpc
 from tensor2tensor.envs import env_service_servicer
@@ -34,9 +33,7 @@
 
 
 def add_port(server, port):
-  server_credentials = loas2.loas2_server_credentials()
-  return server.add_secure_port(
-      _ADDRESS_FORMAT.format(port), server_credentials)
+  return server.add_insecure_port(_ADDRESS_FORMAT.format(port))  # pylint: disable=unreachable
 
 
 def serve(output_dir, env, port):

From 474e536945c6ed8939f3a4ebf269c43fa96ffef0 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 19 Aug 2019 17:10:22 -0700
Subject: [PATCH 2291/2720] Don't test on tf-nightly anymore, since it is going
 to be switching to Tf 2.0

PiperOrigin-RevId: 264277600
---
 .travis.yml                    |  5 +----
 oss_scripts/oss_pip_install.sh |  8 +-------
 oss_scripts/oss_tests.sh       | 21 +++++++++++----------
 3 files changed, 13 insertions(+), 21 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index f979b8146..45676b315 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -18,10 +18,7 @@ env:
     # This is necessary to have gsutil work with Python 2.7
     - BOTO_CONFIG=/dev/null
   matrix:
-    # We test against the latest stable TensorFlow and tf-nightly.
-    # If updating, also update TF_LATEST above
     - TF_VERSION="1.14.*"
-    - TF_VERSION="tf-nightly"
 install:
   - ./oss_scripts/oss_pip_install.sh
 script:
@@ -32,6 +29,6 @@ script:
   # errors on Travis.
   #
   # TODO(afrozm): Re-enable if this becomes an issue.
-  # - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "tf-nightly"  ]]; then
+  # - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]]; then
   #       pylint -j 2 tensor2tensor;
   #   fi
diff --git a/oss_scripts/oss_pip_install.sh b/oss_scripts/oss_pip_install.sh
index 9db9854a4..5246a8363 100755
--- a/oss_scripts/oss_pip_install.sh
+++ b/oss_scripts/oss_pip_install.sh
@@ -8,13 +8,7 @@ set -e  # fail and exit on any command erroring
 # Make sure we have the latest version of numpy - avoid problems we were
 # seeing with Python 3
 pip install -q -U numpy
-
-if [[ "$TF_VERSION" == "tf-nightly"  ]]
-then
-  pip install tf-nightly;
-else
-  pip install -q "tensorflow==$TF_VERSION"
-fi
+pip install -q "tensorflow==$TF_VERSION"
 
 # First ensure that the base dependencies are sufficient for a full import
 pip install -q -e .
diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 3fefdf5f5..d27f0d604 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -112,16 +112,17 @@ set_status
 pytest --disable-warnings tensor2tensor/data_generators/allen_brain_test.py
 set_status
 
-
-# Test models/research only against tf-nightly
-if [[ "$TRAVIS_PYTHON_VERSION" == "2.7"  ]] && [[ "$TF_VERSION" == "tf-nightly"  ]]
-then
-  # Ignores:
-  # * Glow requires the CIFAR-10 dataset to be generated
-  pytest --disable-warnings tensor2tensor/models/research \
-    --ignore=tensor2tensor/models/research/glow_test.py
-  set_status
-fi
+# TODO(afrozm): Enable this unconditionally?
+
+## Test models/research only against tf-nightly
+#if [[ "$TRAVIS_PYTHON_VERSION" == "2.7"  ]]
+#then
+#  # Ignores:
+#  # * Glow requires the CIFAR-10 dataset to be generated
+#  pytest --disable-warnings tensor2tensor/models/research \
+#    --ignore=tensor2tensor/models/research/glow_test.py
+#  set_status
+#fi
 
 if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]
 then

From 25fa8b7de9174a0bc8953253694cd8f03f071401 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 19 Aug 2019 17:41:08 -0700
Subject: [PATCH 2292/2720] Don't compute gradients wrt kwargs

PiperOrigin-RevId: 264283626
---
 tensor2tensor/trax/layers/base.py             |  8 +--
 tensor2tensor/trax/layers/base_test.py        |  4 +-
 tensor2tensor/trax/layers/reversible.py       | 29 ++++------
 .../models/research/transformer_revnet.py     | 54 ++++++-------------
 4 files changed, 34 insertions(+), 61 deletions(-)

diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 3e2448a63..52496b67f 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -243,22 +243,22 @@ def __call__(self, x, params=(), **kwargs):
 
       # See this link for how custom transformations are defined in JAX:
       # https://jax.readthedocs.io/en/latest/jax.html#jax.custom_transforms
+      # Note that we capture the kwargs and don't calculate gradients wrt. them.
       @jax.custom_transforms
-      def do_call(y, params, kwargs):
+      def do_call(y, params):
         return self.call(y, params=params, **kwargs)
 
       # This is the custom gradient (vector-jacobian product in JAX) function.
       # For the exact specification of this custom transformation see this link:
       # https://jax.readthedocs.io/en/latest/jax.html#jax.defjvp_all
-      # Note that we make arguments positional to allow gradients wrt. them.
-      def do_call_vjp(y, params, kwargs):
+      def do_call_vjp(y, params):
         output = self.call(y, params=params, **kwargs)
         def vjpfun(grad):
           return self.custom_grad(y, output, grad, params, **kwargs)
         return output, vjpfun
 
       jax.defvjp_all(do_call, do_call_vjp)
-      return do_call(x, params, kwargs)
+      return do_call(x, params)
 
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback()
diff --git a/tensor2tensor/trax/layers/base_test.py b/tensor2tensor/trax/layers/base_test.py
index 7b15484ec..a7dd593df 100644
--- a/tensor2tensor/trax/layers/base_test.py
+++ b/tensor2tensor/trax/layers/base_test.py
@@ -51,7 +51,7 @@ def has_custom_grad(self):
         return True
 
       def custom_grad(self, inputs, output, ct, params, **kwargs):
-        return (backend.numpy.zeros_like(ct), None, None)
+        return (backend.numpy.zeros_like(ct), None)
 
     layer = IdWithZeroGrad()
     rng = backend.random.get_prng(0)
@@ -81,7 +81,7 @@ def has_custom_grad(self):
         return True
 
       def custom_grad(self, inputs, output, ct, params, **kwargs):
-        return (inputs, None, None)
+        return (inputs, None)
 
     layer = IdWithIdGrad()
     rng = backend.random.get_prng(0)
diff --git a/tensor2tensor/trax/layers/reversible.py b/tensor2tensor/trax/layers/reversible.py
index 7d25e93b6..63dd3183e 100644
--- a/tensor2tensor/trax/layers/reversible.py
+++ b/tensor2tensor/trax/layers/reversible.py
@@ -47,17 +47,18 @@ def reverse_and_grad(self, output, grad, params=(), **kwargs):
       **kwargs: kwargs for the layer
 
     Returns:
-      A tuple (x, x_grad), where x is the reconstructed input and x_grad is the
-      gradient signal for the input.
+      A tuple (x, (x_grad, params_grad)), where x is the reconstructed input,
+      x_grad is the gradient signal for the input, and params_grad is the
+      gradient signal for the parameters.
     """
     # Note: jax.vjp does not allow us to use **kwargs in the signature here.
-    def _do_call(x, params, kwargs):
+    def _do_call(x, params):
       return super(ReversibleLayer, self).call(x, params=params, **kwargs)
 
     reconstructed_x = self.reverse(output, params, **kwargs)
-    _, vjpfun = jax.vjp(_do_call, reconstructed_x, params, kwargs)
-    x_grad = vjpfun(grad)
-    return reconstructed_x, x_grad
+    _, vjpfun = jax.vjp(_do_call, reconstructed_x, params)
+    x_params_grad = vjpfun(grad)
+    return reconstructed_x, x_params_grad
 
   @property
   def has_custom_grad(self):
@@ -65,8 +66,8 @@ def has_custom_grad(self):
 
   def custom_grad(self, inputs, output, ct, params, **kwargs):
     del inputs
-    _, input_ct = self.reverse_and_grad(output, ct, params, **kwargs)
-    return input_ct
+    _, inputs_params_ct = self.reverse_and_grad(output, ct, params, **kwargs)
+    return inputs_params_ct
 
 
 class ReversibleSwap(ReversibleLayer, cb.Swap):
@@ -114,15 +115,7 @@ def reverse_and_grad(self, output, ct, params=(), **kwargs):
     for layer, p, rng in reversed(zip(self.sublayers(), params, rngs)):
       layer_val, layer_ct = layer.reverse_and_grad(
           layer_val, layer_ct, p, rng=rng, **kwargs)
-      layer_ct, p_ct, kwargs_ct = layer_ct
+      layer_ct, p_ct = layer_ct
       params_ct.insert(0, p_ct)
 
-    # TODO(kitaev): Handle kwargs_ct properly. However, kwargs generally only
-    # contains the rng, which is non-differentiable.
-    for k in kwargs:
-      if k != 'rng':
-        raise NotImplementedError(
-            'ReversibleSerial does not support differentiation wrt kwargs,'
-            'and the key {} is not known to be non-differentiable.'.format(k))
-
-    return layer_val, (layer_ct, params_ct, kwargs_ct)
+    return layer_val, (layer_ct, params_ct)
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index c4be23e6a..75e357fe6 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -196,7 +196,7 @@ def reverse(self, output, params=(), **kwargs):
 
   def reverse_and_grad(self, output, ct, params=(), **kwargs):
     del params, kwargs
-    return self.reverse(output), (self.reverse(ct), (), ())
+    return self.reverse(output), (self.reverse(ct), ())
 
 
 @tl.layer()
@@ -257,22 +257,19 @@ def reverse_and_grad(self, output, ct, params=(), **kwargs):
     if rng is not None:
       rngs = backend.random.split(rng, self._n_layers)
 
-    # Note: jax.vjp does not allow us to use **kwargs in the signature here.
-    def call_compute_residual(x, params, kwargs):
-      return self.compute_residual(x, params, **kwargs)
+    def call_compute_residual(x, params):
+      return self.compute_residual(x, params, rng=rngs[0], **kwargs)
 
     assert len(ct) == 2
     ct = ((ct[0], ct[0], ct[1]))
 
-    compute_residual_kwargs = kwargs.copy()
-    compute_residual_kwargs['rng'] = rngs[0]
     stack_with_residual, vjpfun = jax.vjp(
-        call_compute_residual, output, params[0], compute_residual_kwargs)
+        call_compute_residual, output, params[0])
     reconstructed_x = self.subtract_top(
         stack_with_residual, params[-1], rng=rngs[-1], **kwargs)
 
-    x_ct, residual_params_ct, kwargs_ct = vjpfun(ct)
-    return reconstructed_x, (x_ct, (residual_params_ct, ()), kwargs_ct)
+    x_ct, residual_params_ct = vjpfun(ct)
+    return reconstructed_x, (x_ct, (residual_params_ct, ()))
 
 
 class ComputeAttentionHeads(tl.Layer):
@@ -716,23 +713,17 @@ def reverse_and_grad(self, output, ct, params=(), **kwargs):
 
     # Forward pass through self.pre_attention, while preparing for
     # later backprop.
-    # Note: jax.vjp does not allow us to use **kwargs in the signature here.
-    def call_pre_attention(x, params, kwargs):
-      return self.pre_attention(x, params, **kwargs)
-    pre_attention_kwargs = kwargs.copy()
-    pre_attention_kwargs['rng'] = rngs[0]
-    stack, pre_attention_vjpfun = jax.vjp(
-        call_pre_attention, output, params[0], pre_attention_kwargs)
+    def call_pre_attention(x, params):
+      return self.pre_attention(x, params, rng=rngs[0], **kwargs)
+    stack, pre_attention_vjpfun = jax.vjp(call_pre_attention, output, params[0])
 
     # Backprop through adding the residual
     assert len(ct) == 2
     ct = saved_ct = (ct[0], ct[0], ct[1])
 
     # Backprop through self.post_attention with respect to the inputs only
-    call_post_attention_kwargs = kwargs.copy()
-    call_post_attention_kwargs['rng'] = rngs[2]
     def call_post_attention(x):
-      return self.post_attention(x, params[2], **call_post_attention_kwargs)
+      return self.post_attention(x, params[2], rng=rngs[2], **kwargs)
     # Note: these are *not* the actual inputs to self.post_attention.
     # If self.post_attention is not linear, we will get incorrect gradients.
     dummy_inputs = (stack[-3], stack[-2], stack[-1])
@@ -740,25 +731,19 @@ def call_post_attention(x):
     (ct,) = post_attention_vjpfun(ct)
 
     # Simultaneous forward pass and backprop through the attention mechanism
-    attention_kwargs = kwargs.copy()
-    attention_kwargs['rng'] = rngs[1]
     stack, ct = self.attention.forward_and_vjp(
-        stack, ct, **attention_kwargs)
+        stack, ct, rng=rngs[1], **kwargs)
     attention_params_ct = ()
 
     # Backprop through self.pre_attention
-    (x_ct,
-     pre_attention_params_ct,
-     pre_attention_kwargs_ct) = pre_attention_vjpfun(ct)
+    x_ct, pre_attention_params_ct = pre_attention_vjpfun(ct)
 
     # Forward pass for self.post_attention, and backprop with respect to the
     # parameters only
-    def call_post_attention2(params, kwargs):
-      return self.post_attention(stack, params, **kwargs)
-    stack, post_attention_vjpfun = jax.vjp(
-        call_post_attention2, params[2], call_post_attention_kwargs)
-    (post_attention_params_ct,
-     post_attention_kwargs_ct) = post_attention_vjpfun(saved_ct)
+    def call_post_attention2(params):
+      return self.post_attention(stack, params, rng=rngs[2], **kwargs)
+    stack, post_attention_vjpfun = jax.vjp(call_post_attention2, params[2])
+    (post_attention_params_ct,) = post_attention_vjpfun(saved_ct)
 
     # Forward pass through subtracting the residual
     reconstructed_x = self.subtract_top(
@@ -771,12 +756,7 @@ def call_post_attention2(params, kwargs):
         (),
         )
 
-    # We don't actually backprop through the kwargs, but the API requires that
-    # we provide a value for kwargs_ct.
-    kwargs_ct = pre_attention_kwargs_ct
-    del post_attention_kwargs_ct
-
-    return reconstructed_x, (x_ct, params_ct, kwargs_ct)
+    return reconstructed_x, (x_ct, params_ct)
 
 
 def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,

From 61f698d81d8f0b9785c25bd740313909ccf50446 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 19 Aug 2019 17:51:07 -0700
Subject: [PATCH 2293/2720] Don't compute gradient wrt loop index

PiperOrigin-RevId: 264285023
---
 .../trax/models/research/transformer_revnet.py        | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 75e357fe6..0c92b7e68 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -470,8 +470,12 @@ def forward_slice(query_slice, q_loop_idx, key, value):
       return out_slice
 
     def forward_and_vjp_slice(query_slice, q_loop_idx, key, value, ct_slice):
+      # Capture q_loop_idx to avoid calculated gradients wrt. it.
+      def forward_slice_with_q_loop_idx(query_slice, key, value):
+        return forward_slice(query_slice, q_loop_idx, key, value)
+
       output_slice, vjpfun = jax.vjp(
-          forward_slice, query_slice, q_loop_idx, key, value)
+          forward_slice_with_q_loop_idx, query_slice, key, value)
       return output_slice, vjpfun(ct_slice)
 
     q_loop_idx = np.zeros((), dtype=np.int32)
@@ -513,9 +517,8 @@ def body_fun(vals):
             query_slice, q_loop_idx, key, value, ct_slice)
         query_ct_accum = jax.lax.dynamic_update_slice_in_dim(
             query_ct_accum, partial_ct[0], q_loop_idx, axis=-2)
-        # ignore partial_ct[1], which is wrt the loop idx
-        key_ct_accum = key_ct_accum + partial_ct[2]
-        value_ct_accum = value_ct_accum + partial_ct[3]
+        key_ct_accum = key_ct_accum + partial_ct[1]
+        value_ct_accum = value_ct_accum + partial_ct[2]
       else:
         out_slice = forward_slice(query_slice, q_loop_idx, key, value)
 

From 1133b36b075fd3d44c05535fc57a536c329dcb7e Mon Sep 17 00:00:00 2001
From: Pierre Ruyssen <pierrot@google.com>
Date: Tue, 20 Aug 2019 01:50:25 -0700
Subject: [PATCH 2294/2720] Update `tfds.load()` callers to specify
 `shuffle_files=True` when necessary.

PiperOrigin-RevId: 264338129
---
 tensor2tensor/v2/t2t.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/v2/t2t.py b/tensor2tensor/v2/t2t.py
index ffa4db73e..9e4363f28 100644
--- a/tensor2tensor/v2/t2t.py
+++ b/tensor2tensor/v2/t2t.py
@@ -75,8 +75,10 @@ def train_and_eval_dataset(dataset_name, data_dir):
   eval_split = tfds.Split.VALIDATION
   if tfds.Split.VALIDATION not in splits:
     eval_split = tfds.Split.TEST
-  train, valid = tfds.load(
-      name=dataset_name, split=[tfds.Split.TRAIN, eval_split])
+  train = tfds.load(
+      name=dataset_name, split=tfds.Split.TRAIN, shuffle_files=True)
+  valid = tfds.load(
+      name=dataset_name, split=eval_split, shuffle_files=False)
   keys = None
   if info.supervised_keys:
     keys = ([info.supervised_keys[0]], [info.supervised_keys[1]])

From 43cd6e4d11a29bede268f3b0414ab3cf901ea1c6 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Tue, 20 Aug 2019 11:21:21 -0700
Subject: [PATCH 2295/2720] If do_mask is false, targets should be d

PiperOrigin-RevId: 264426064
---
 tensor2tensor/models/research/transformer_vae.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 9cfe80933..eca25e1a3 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -507,6 +507,8 @@ def bn_inputs():
       # reshape back to 4d here
       if hparams.task == "image":
         targets = tf.reshape(targets, original_targets_shape)
+    else:
+      targets = d
 
   res = decode_transformer(inputs, ed, targets, hparams, "decoder",
                            causal=hparams.causal)

From 3ec84b14bf8550adc875ecca5d2643c50147cb6a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 20 Aug 2019 11:34:36 -0700
Subject: [PATCH 2296/2720] Remove obsolete v2 directory and add all_files for
 trax.

PiperOrigin-RevId: 264429293
---
 tensor2tensor/v2/keras_utils.py        |  50 ----
 tensor2tensor/v2/models/basic.py       |  92 -------
 tensor2tensor/v2/models/resnet.py      |  62 -----
 tensor2tensor/v2/models/transformer.py | 130 ---------
 tensor2tensor/v2/t2t.py                | 349 -------------------------
 tensor2tensor/v2/t2t_trainer.py        |  64 -----
 6 files changed, 747 deletions(-)
 delete mode 100644 tensor2tensor/v2/keras_utils.py
 delete mode 100644 tensor2tensor/v2/models/basic.py
 delete mode 100644 tensor2tensor/v2/models/resnet.py
 delete mode 100644 tensor2tensor/v2/models/transformer.py
 delete mode 100644 tensor2tensor/v2/t2t.py
 delete mode 100644 tensor2tensor/v2/t2t_trainer.py

diff --git a/tensor2tensor/v2/keras_utils.py b/tensor2tensor/v2/keras_utils.py
deleted file mode 100644
index 251b03796..000000000
--- a/tensor2tensor/v2/keras_utils.py
+++ /dev/null
@@ -1,50 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Utilities to use TF v1 layers with Keras and TF v2 easily."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import tensorflow as tf
-
-
-class FunctionLayer(tf.compat.v2.keras.layers.Layer):
-  """Layer made of a function. Stores all variables."""
-
-  def __init__(self, function, name=None):
-    if name is None:
-      name = function.__name__
-    super(FunctionLayer, self).__init__(name=name)
-    self._template = tf.compat.v1.make_template(name, function)
-    self._was_called = False
-
-  @property
-  def losses(self):
-    return []
-
-  def compute_mask(self, inputs, previous_mask):
-    return previous_mask
-
-  @tf.function
-  def _template_call(self, *args, **kwargs):
-    """Call to template but made in graph mode for better speed."""
-    return self._template(*args, **kwargs)
-
-  def call(self, *args, **kwargs):
-    if not self._was_called:  # Create variables at first call.
-      return self._template(*args, **kwargs)
-    return self._template_call(*args, **kwargs)
diff --git a/tensor2tensor/v2/models/basic.py b/tensor2tensor/v2/models/basic.py
deleted file mode 100644
index 113dcafc2..000000000
--- a/tensor2tensor/v2/models/basic.py
+++ /dev/null
@@ -1,92 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Basic models for testing simple tasks."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import tensorflow as tf
-import gin.tf
-
-
-@gin.configurable(whitelist=["num_hidden_layers", "hidden_size", "dropout"])
-class BasicFcRelu(tf.keras.Model):
-  """Basic fully-connected + ReLU model."""
-
-  def __init__(self, features_info=None, input_names=None, target_names=None,
-               num_hidden_layers=2, hidden_size=64, dropout=0.1):
-    super(BasicFcRelu, self).__init__()
-    self._input_name = input_names[0]
-    input_shape = features_info[self._input_name].shape
-    num_output_classes = features_info[target_names[0]].num_classes
-    self._num_hidden_layers = num_hidden_layers
-    self._dense_layers = []
-    self._dropout_layers = []
-
-    # Now the model.
-    self._flatten_layer = tf.keras.layers.Flatten(input_shape=input_shape)
-    for i in range(num_hidden_layers):
-      self._dense_layers.append(tf.keras.layers.Dense(
-          hidden_size, activation="relu", name="layer_%d" % i))
-      self._dropout_layers.append(tf.keras.layers.Dropout(
-          rate=dropout))
-    self._logits = tf.keras.layers.Dense(
-        num_output_classes, activation=None)
-
-  def call(self, inputs, training=False):
-    x = tf.cast(inputs[self._input_name], tf.float32) / 255.0
-    x = self._flatten_layer(x)
-    for i in range(self._num_hidden_layers):
-      x = self._dense_layers[i](x)
-      x = self._dropout_layers[i](x, training=training)
-    return self._logits(x)
-
-
-def basic_fc_large():
-  """Large set of parameters for this model."""
-  gin.bind_parameter("BasicFcRelu.num_hidden_layers", 3)
-  gin.bind_parameter("BasicFcRelu.hidden_size", 128)
-  gin.bind_parameter("BasicFcRelu.dropout", 0.3)
-  return BasicFcRelu
-
-
-# TODO(lukaszkaiser): could we allow coding like this? it's much easier!
-# This will run fine, but not train as new layers are made in each step!
-@gin.configurable(whitelist=["num_hidden_layers", "hidden_size", "dropout"])
-class BasicFcReluV2(tf.keras.Model):
-  """Basic fully-connected + ReLU model, nicer code version."""
-
-  def __init__(self, features_info=None, input_names=None, target_names=None,
-               num_hidden_layers=2, hidden_size=64, dropout=0.1):
-    super(BasicFcReluV2, self).__init__()
-    self._input_name = input_names[0]
-    self._input_shape = features_info[self._input_name].shape
-    self._num_output_classes = features_info[target_names[0]].num_classes
-    self._num_hidden_layers = num_hidden_layers
-    self._dropout = dropout
-    self._hidden_size = hidden_size
-
-  def call(self, inputs, training=False):
-    x = tf.cast(inputs[self._input_name], tf.float32) / 255.0
-    x = tf.keras.layers.Flatten(
-        input_shape=self._input_shape)(x)
-    for i in range(self._num_hidden_layers):
-      x = tf.keras.layers.Dense(
-          self._hidden_size, activation="relu", name="layer_%d" % i)(x)
-      x = tf.keras.layers.Dropout(rate=self._dropout)(x, training=training)
-    return tf.keras.layers.Dense(
-        self._num_output_classes, activation=None)(x)
diff --git a/tensor2tensor/v2/models/resnet.py b/tensor2tensor/v2/models/resnet.py
deleted file mode 100644
index db584e3fc..000000000
--- a/tensor2tensor/v2/models/resnet.py
+++ /dev/null
@@ -1,62 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Basic models for testing simple tasks."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.models import resnet
-import tensorflow as tf
-import gin.tf
-
-
-@gin.configurable(whitelist=["layer_sizes", "filter_sizes"])
-class Resnet(tf.keras.Model):
-  """Resnet."""
-
-  def __init__(self, features_info=None, input_names=None, target_names=None,
-               layer_sizes=None, filter_sizes=None):
-    super(Resnet, self).__init__()
-    # Base config for resnet-50.
-    if layer_sizes is None:
-      layer_sizes = [3, 4, 6, 3]
-    if filter_sizes is None:
-      filter_sizes = [64, 64, 128, 256, 512]
-    self._input_name = input_names[0]
-    num_output_classes = features_info[target_names[0]].num_classes
-
-    # Now the model.
-    def resnet_model(inputs, training=None):
-      return resnet.resnet_v2(
-          inputs,
-          resnet.bottleneck_block,
-          layer_sizes,
-          filter_sizes,
-          is_training=training,
-          is_cifar=True)
-
-    inputs = tf.keras.Input(shape=(32, 32, 3))
-    outputs = resnet_model(inputs)
-    self._resnet = tf.keras.Model(inputs=inputs, outputs=outputs)
-    self._logits = tf.keras.layers.Dense(
-        num_output_classes, activation=None)
-
-  def call(self, inputs, training=False):
-    x = tf.cast(inputs[self._input_name], tf.float32) / 255.0
-    x = self._resnet(x, training)
-    x = tf.reduce_mean(x, axis=[1, 2])
-    return self._logits(x)
diff --git a/tensor2tensor/v2/models/transformer.py b/tensor2tensor/v2/models/transformer.py
deleted file mode 100644
index edb27dcf2..000000000
--- a/tensor2tensor/v2/models/transformer.py
+++ /dev/null
@@ -1,130 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Transformer model from "Attention Is All You Need".
-
-The Transformer model consists of an encoder and a decoder. Both are stacks
-of self-attention layers followed by feed-forward layers. This model yields
-good results on a number of problems, especially in NLP and machine translation.
-
-See "Attention Is All You Need" (https://arxiv.org/abs/1706.03762) for the full
-description of the model and the results obtained with its early version.
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.layers import transformer_layers
-from tensor2tensor.models import transformer
-from tensor2tensor.v2 import keras_utils
-import tensorflow as tf
-import gin.tf
-
-
-@gin.configurable(whitelist=["hidden_size", "filter_size"])
-class Transformer(tf.keras.Model):
-  """Transformer."""
-
-  def __init__(self, features_info=None, input_names=None, target_names=None,
-               hidden_size=512, filter_size=2048):
-    super(Transformer, self).__init__()
-    # TODO(lukaszkaiser): gin'ify and split into encoder/decoder classes.
-    self._has_input = True if input_names else False
-    self._input_name = input_names[0]
-    self._target_name = target_names[0]
-    try:
-      target_vocab_size = features_info[self._target_name].num_classes
-    except AttributeError:
-      target_vocab_size = features_info[self._target_name].encoder.vocab_size
-    hparams = transformer.transformer_base()
-    hparams.hidden_size = hidden_size
-    hparams.filter_size = filter_size
-
-    # Now the model.
-    self._embedding = tf.keras.layers.Embedding(
-        target_vocab_size, hidden_size, mask_zero=True)
-    def transformer_encoder(inputs, features):
-      return transformer.transformer_encode(
-          transformer_layers.transformer_encoder, inputs, None,
-          hparams, features=features)
-
-    def transformer_prepare_decoder(targets, features):
-      return transformer.transformer_prepare_decoder(targets, hparams, features)
-
-    def transformer_decoder(decoder_input, encoder_output,
-                            encoder_decoder_attention_bias,
-                            decoder_self_attention_bias,
-                            features):
-      return transformer.transformer_decode(
-          transformer.transformer_decoder,
-          decoder_input,
-          encoder_output,
-          encoder_decoder_attention_bias,
-          decoder_self_attention_bias,
-          hparams,
-          nonpadding=transformer.features_to_nonpadding(features, "targets"))
-
-    if self._has_input:
-      self._encoder = keras_utils.FunctionLayer(transformer_encoder)
-    self._prepare_decoder = keras_utils.FunctionLayer(
-        transformer_prepare_decoder)
-    self._decoder = keras_utils.FunctionLayer(transformer_decoder)
-    self._logits = tf.keras.layers.Dense(
-        target_vocab_size, activation=None)
-
-  def call(self, features, training=False):
-    """Transformer main model_fn.
-
-    Args:
-      features: Map of features to the model. Should contain the following:
-          "inputs": Transformer inputs. [batch_size, input_length, 1,
-            hidden_dim].
-          "targets": Target decoder outputs. [batch_size, decoder_length, 1,
-            hidden_dim]
-          "target_space_id": A scalar int from data_generators.problem.SpaceID.
-      training: Whether we are training or not.
-
-    Returns:
-      Final decoder representation. [batch_size, decoder_length, hidden_dim]
-    """
-    if self._has_input:
-      inputs = features[self._input_name]
-      inputs = tf.expand_dims(self._embedding(inputs), 2)
-      encoder_output, encoder_decoder_attention_bias = self._encoder(
-          inputs, features)
-    else:
-      encoder_output, encoder_decoder_attention_bias = (None, None)
-
-    targets = features[self._target_name]
-    targets = self._embedding(targets)
-    decoder_input, decoder_self_attention_bias = self._prepare_decoder(
-        targets, features)
-    decoder_output = self._decoder(
-        decoder_input,
-        encoder_output,
-        encoder_decoder_attention_bias,
-        decoder_self_attention_bias,
-        features)
-
-    return self._logits(tf.squeeze(decoder_output, axis=2))
-
-
-def transformer_base_single_gpu():
-  """Single-gpu set of parameters for Transformer."""
-  gin.bind_parameter("T2TLearningRateSchedule.warmup_steps", 16000)
-  gin.bind_parameter("preprocess_fn.max_target_length", 256)
-  gin.bind_parameter("batch_fn.eval_batch_size", 8)
-  return Transformer
diff --git a/tensor2tensor/v2/t2t.py b/tensor2tensor/v2/t2t.py
deleted file mode 100644
index 9e4363f28..000000000
--- a/tensor2tensor/v2/t2t.py
+++ /dev/null
@@ -1,349 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""T2T models, configs and main training functions."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import collections
-import os
-
-import gin
-
-from tensor2tensor import problems
-from tensor2tensor.utils import data_reader
-from tensor2tensor.v2.models import basic
-from tensor2tensor.v2.models import resnet
-from tensor2tensor.v2.models import transformer
-
-import tensorflow as tf
-import tensorflow_datasets as tfds
-
-
-# Since there are few models and configs for now, we use this simple registry.
-# TODO(lukaszkaiser): find a better way to do this or remove altogether.
-_MODEL_REGISTRY = {
-    "basic_fc_relu": lambda: basic.BasicFcRelu,
-    "basic_fc_large": basic.basic_fc_large,
-    "basic_fc_relu_v2": lambda: basic.BasicFcReluV2,
-    "resnet": lambda: resnet.Resnet,
-    "transformer": transformer.transformer_base_single_gpu,
-}
-
-
-def train_and_eval_dataset(dataset_name, data_dir):
-  """Return train and evaluation datasets, feature info and supervised keys.
-
-  Args:
-    dataset_name: a string, the name of the dataset; if it starts with "v1_"
-      then we'll search T2T Problem registry for it, otherwise we assume it
-      is a dataset from TFDS and load it from there.
-    data_dir: directory where the data is located.
-
-  Returns:
-    a 4-tuple consisting of:
-     * the train tf.data.Dataset
-     * the eval tf.data.Dataset
-     * information about features: a python dictionary with feature names
-         as keys and an object as value that provides .shape and .num_classes.
-     * supervised_keys: information what's the input and what's the target,
-         ie., a pair of lists with input and target feature names.
-  """
-  if dataset_name.startswith("v1_"):
-    return _train_and_eval_dataset_v1(dataset_name[3:], data_dir)
-  dataset_builder = tfds.builder(dataset_name, data_dir=data_dir)
-  info = dataset_builder.info
-  splits = dataset_builder.info.splits
-  if tfds.Split.TRAIN not in splits:
-    raise ValueError("To train we require a train split in the dataset.")
-  if tfds.Split.VALIDATION not in splits and "test" not in splits:
-    raise ValueError("We require a validation or test split in the dataset.")
-  eval_split = tfds.Split.VALIDATION
-  if tfds.Split.VALIDATION not in splits:
-    eval_split = tfds.Split.TEST
-  train = tfds.load(
-      name=dataset_name, split=tfds.Split.TRAIN, shuffle_files=True)
-  valid = tfds.load(
-      name=dataset_name, split=eval_split, shuffle_files=False)
-  keys = None
-  if info.supervised_keys:
-    keys = ([info.supervised_keys[0]], [info.supervised_keys[1]])
-  return train, valid, info.features, keys
-
-
-def _make_info(shape_list, num_classes):
-  """Create an info-like tuple for feature given some shapes and vocab size."""
-  feature_info = collections.namedtuple("FeatureInfo", ["shape", "num_classes"])
-  cur_shape = list(shape_list[0])
-  # We need to merge the provided shapes, put None where they disagree.
-  for shape in shape_list:
-    if len(shape) != len(cur_shape):
-      raise ValueError("Shapes need to have the same number of dimensions.")
-    for i in range(len(shape)):
-      if cur_shape[i] is not None:
-        if shape[i] != cur_shape[i]:
-          cur_shape[i] = None
-  return feature_info(cur_shape, num_classes)
-
-
-def _select_features(example, feature_list=None):
-  """Select a subset of features from the example dict."""
-  feature_list = feature_list or ["inputs", "targets"]
-  return {f: example[f] for f in feature_list}
-
-
-def _train_and_eval_dataset_v1(problem_name, data_dir):
-  """Return train and evaluation datasets, feature info and supervised keys."""
-  problem = problems.problem(problem_name)
-  train_dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, data_dir)
-  train_dataset = train_dataset.map(_select_features)
-  eval_dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, data_dir)
-  eval_dataset = eval_dataset.map(_select_features)
-  supervised_keys = (["inputs"], ["targets"])
-  hparams = problem.get_hparams()
-  # We take a few training examples to guess the shapes.
-  input_shapes, target_shapes = [], []
-  for example in train_dataset.take(3):
-    input_shapes.append(example["inputs"].shape.as_list())
-    target_shapes.append(example["targets"].shape.as_list())
-  input_vocab_size = hparams.vocab_size["inputs"]
-  target_vocab_size = hparams.vocab_size["targets"]
-  input_info = _make_info(input_shapes, input_vocab_size)
-  target_info = _make_info(target_shapes, target_vocab_size)
-  info = {"inputs": input_info, "targets": target_info}
-  return train_dataset, eval_dataset, info, supervised_keys
-
-
-@gin.configurable(blacklist=["dataset", "training"])
-def preprocess_fn(dataset, training, max_target_length=-1):
-  def target_right_length(_, target):
-    if max_target_length < 1 or not training:
-      return tf.constant(True)
-    return tf.less(tf.shape(target)[0], max_target_length + 1)
-  dataset = dataset.filter(target_right_length)
-  return dataset
-
-
-@gin.configurable(blacklist=["dataset", "training", "shapes", "target_names"])
-def batch_fn(dataset, training, shapes, target_names,
-             batch_size=32, eval_batch_size=32, bucket_batch_length=32,
-             bucket_max_length=256, bucket_min_length=8,
-             bucket_length_step=1.1, buckets=None):
-  """Batching function."""
-  del target_names
-  # If bucketing is not specified, check if target shapes are variable.
-  cur_batch_size = batch_size if training else eval_batch_size
-  if buckets is None:
-    variable_target_shapes = False
-    target_shape = shapes[1]
-    for dim in target_shape:
-      if dim is None:
-        variable_target_shapes = True
-    tf.logging.info("Heuristically setting bucketing to %s based on shapes "
-                    "of target tensors." % variable_target_shapes)
-    if variable_target_shapes:
-      batch_size_per_token = cur_batch_size * bucket_batch_length
-      scheme = data_reader.batching_scheme(batch_size_per_token,
-                                           bucket_max_length,
-                                           bucket_min_length,
-                                           bucket_length_step,
-                                           drop_long_sequences=training)
-      buckets = (scheme["boundaries"], scheme["batch_sizes"])
-
-  if buckets:
-    tf.logging.info("Bucketing with buckets %s." % str(buckets))
-    def example_length(_, target):
-      return tf.shape(target)[0]
-    boundaries, batch_sizes = buckets
-    dataset = dataset.apply(tf.data.experimental.bucket_by_sequence_length(
-        example_length, boundaries, batch_sizes))
-  else:
-    dataset = dataset.padded_batch(cur_batch_size, shapes)
-  return dataset
-
-
-def shuffle_and_batch_data(dataset, target_names, features_info, training):
-  """Shuffle and batch the given dataset."""
-  def append_targets(example):
-    """Append targets to the example dictionary. Needed for Keras."""
-    if len(target_names) == 1:
-      return (example, example[target_names[0]])
-    targets = {}
-    for name in target_names:
-      targets[name] = example[name]
-    return (example, targets)
-  dataset = dataset.map(append_targets)
-  if training:
-    dataset = dataset.repeat()
-  shapes = {k: features_info[k].shape for k in features_info}
-  shapes = (shapes, shapes[target_names[0]])
-  dataset = dataset.shuffle(128)
-  dataset = preprocess_fn(dataset, training)
-  dataset = batch_fn(dataset, training, shapes, target_names)
-  return dataset.prefetch(8)
-
-
-@gin.configurable()
-class T2TLearningRateSchedule(
-    tf.keras.optimizers.schedules.LearningRateSchedule):
-  """A LearningRateSchedule that uses a T2T config."""
-
-  def __init__(self, schedule=None, constant=0.1, warmup_steps=200):
-    """Applies the give T2T schedule string with the given parameters."""
-    super(T2TLearningRateSchedule, self).__init__()
-    self.schedule = schedule or "constant * linear_warmup * rsqrt_decay"
-    self.constant = constant
-    self.warmup_steps = warmup_steps
-
-  def __call__(self, step):
-    ret = tf.constant(1.0)
-    for name in [n.strip() for n in self.schedule.split("*")]:
-      if name == "constant":
-        ret *= self.constant
-      elif name == "linear_warmup":
-        ret *= tf.minimum(1.0, step / self.warmup_steps)
-      elif name == "rsqrt_decay":
-        ret *= tf.rsqrt(tf.maximum(step, self.warmup_steps))
-      else:
-        raise ValueError("Unknown factor %s." % name)
-    tf.contrib.summary.scalar("learning_rate", ret)
-    return ret
-
-  def get_config(self):
-    return {
-        "schedule": self.schedule,
-        "constant": self.constant,
-        "warmup_steps": self.warmup_steps,
-    }
-
-
-@gin.configurable(blacklist=["model"])
-def optimize_fn(model,
-                optimizer=None,
-                learning_rate_schedule=None,
-                loss=None,
-                metrics=None):
-  """Compile the model in Keras."""
-  learning_rate_schedule = learning_rate_schedule or T2TLearningRateSchedule()
-  if optimizer:
-    optimizer = optimizer(learning_rate=learning_rate_schedule)
-  else:  # We use Adam by default with adjusted parameters.
-    optimizer = tf.keras.optimizers.Adam(
-        learning_rate=learning_rate_schedule,
-        beta_1=0.9, beta_2=0.997, epsilon=1e-9)
-  metrics = metrics or [tf.keras.metrics.sparse_categorical_accuracy]
-  def xent_loss(y, x):
-    return tf.keras.backend.sparse_categorical_crossentropy(
-        y, x, from_logits=True)
-  loss = loss or xent_loss
-  return model.compile(optimizer=optimizer,
-                       loss=loss,
-                       metrics=metrics)
-
-
-# We include in gin config everything that could be useful to share between
-# users, so when it gets saved in a .gin file it can be re-ran with few flags.
-@gin.configurable(blacklist=["data_dir", "output_dir"])
-def train_fn(data_dir=None, output_dir=None,
-             model_class=gin.REQUIRED, dataset=gin.REQUIRED,
-             input_names=None, target_names=None,
-             train_steps=1000, eval_steps=1, eval_frequency=100):
-  """Train the given model on the given dataset.
-
-  Args:
-    data_dir: Directory where the data is located.
-    output_dir: Directory where to put the logs and checkpoints.
-    model_class: The model class to train.
-    dataset: The name of the dataset to train on.
-    input_names: List of strings with the names of the features on input.
-    target_names: List of strings with the names of the target features.
-    train_steps: for how many steps to train.
-    eval_steps: for how many steps to do evaluation.
-    eval_frequency: how often (every this many steps) to run evaluation.
-  """
-  train_data, eval_data, features_info, keys = train_and_eval_dataset(
-      dataset, data_dir)
-  if input_names is None:
-    input_names = keys[0]
-  if target_names is None:
-    target_names = keys[1]
-  # TODO(lukaszkaiser): The use of distribution strategy below fails like this:
-  #   .../keras/models.py", line 93, in _clone_functional_model
-  #      for layer in model._input_layers:
-  #   AttributeError: 'BasicFcRelu' object has no attribute '_input_layers'
-  # strategy = tf.distribute.MirroredStrategy()
-  # with strategy.scope():
-  model = model_class(features_info=features_info,
-                      input_names=input_names, target_names=target_names)
-  optimize_fn(model)
-  train_batches = shuffle_and_batch_data(
-      train_data, target_names, features_info, training=True)
-  eval_batches = shuffle_and_batch_data(
-      eval_data, target_names, features_info, training=False)
-  # Need to run one training step just to get optimizer variables to load.
-  model.fit(train_batches, epochs=1, steps_per_epoch=1)
-
-  # Training loop.
-  callbacks = []
-  callbacks.append(tf.keras.callbacks.History())
-  callbacks.append(tf.keras.callbacks.BaseLogger())
-  last_epoch = 0
-  if output_dir is not None:
-    callbacks.append(tf.keras.callbacks.TensorBoard(log_dir=output_dir))
-    output_format = os.path.join(output_dir, "model-{epoch:05d}")
-    callbacks.append(tf.keras.callbacks.ModelCheckpoint(
-        filepath=output_format, save_weights_only=True))
-    checkpoints = tf.gfile.Glob(os.path.join(output_dir, "model-*"))
-    # Take basenames and strip the "model-" prefix.
-    checkpoints = [os.path.basename(ckpt)[6:] for ckpt in checkpoints]
-    # Get epoch numbers from the filenames and sort to obtain last epoch.
-    epoch_numbers = [int(ckpt[:5]) for ckpt in checkpoints if len(ckpt) > 4]
-    epoch_numbers.sort()
-    if epoch_numbers:
-      last_epoch = epoch_numbers[-1]
-      saved_path = os.path.join(output_dir, "model-%05d" % last_epoch)
-      model.load_weights(saved_path)
-  model.fit(train_batches,
-            epochs=train_steps // eval_frequency,
-            steps_per_epoch=eval_frequency,
-            validation_data=eval_batches,
-            validation_steps=eval_steps,
-            initial_epoch=last_epoch,
-            callbacks=callbacks)
-
-
-def t2t_train(model_name, dataset_name,
-              data_dir=None, output_dir=None, config_file=None, config=None):
-  """Main function to train the given model on the given dataset.
-
-  Args:
-    model_name: The name of the model to train.
-    dataset_name: The name of the dataset to train on.
-    data_dir: Directory where the data is located.
-    output_dir: Directory where to put the logs and checkpoints.
-    config_file: the gin configuration file to use.
-    config: string (in gin format) to override gin parameters.
-  """
-  if model_name not in _MODEL_REGISTRY:
-    raise ValueError("Model %s not in registry. Available models:\n * %s." %
-                     (model_name, "\n * ".join(_MODEL_REGISTRY.keys())))
-  model_class = _MODEL_REGISTRY[model_name]()
-  gin.bind_parameter("train_fn.model_class", model_class)
-  gin.bind_parameter("train_fn.dataset", dataset_name)
-  gin.parse_config_files_and_bindings(config_file, config)
-  # TODO(lukaszkaiser): save gin config in output_dir if provided?
-  train_fn(data_dir, output_dir=output_dir)
diff --git a/tensor2tensor/v2/t2t_trainer.py b/tensor2tensor/v2/t2t_trainer.py
deleted file mode 100644
index 9153b0978..000000000
--- a/tensor2tensor/v2/t2t_trainer.py
+++ /dev/null
@@ -1,64 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-r"""T2T trainer for TF 2.0.
-
-This trainer only supports a subset of models and features for now.
-
-Examples:
-
-- train a basic model on mnist:
-    v2/t2t_trainer.py --dataset=mnist --model=basic_fc_relu
-      --config="train_fn.train_steps=4000"
-"""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-from absl import app
-from absl import flags
-from tensor2tensor.v2 import t2t
-import tensorflow as tf
-
-tf.enable_v2_behavior()
-
-FLAGS = flags.FLAGS
-
-flags.DEFINE_string("dataset", None, "Which dataset to use.")
-flags.DEFINE_string("model", None, "Which model to train.")
-flags.DEFINE_string("data_dir", None, "Path to the directory with data.")
-flags.DEFINE_string("output_dir", None,
-                    "Path to the directory to save logs and checkpoints.")
-flags.DEFINE_multi_string("config_file", None,
-                          "Configuration file with parameters (.gin).")
-flags.DEFINE_multi_string("config", None,
-                          "Configuration parameters (gin string).")
-
-
-def main(argv):
-  del argv
-  data_dir, output_dir = FLAGS.data_dir, FLAGS.output_dir
-  if data_dir is not None:
-    data_dir = os.path.expanduser(data_dir)
-  if output_dir is not None:
-    output_dir = os.path.expanduser(output_dir)
-  t2t.t2t_train(FLAGS.model, FLAGS.dataset,
-                data_dir=data_dir, output_dir=output_dir,
-                config_file=FLAGS.config_file, config=FLAGS.config)
-
-
-if __name__ == "__main__":
-  app.run(main)

From 212b59909a22071418520327ec57929dc172a2a5 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 20 Aug 2019 16:27:32 -0700
Subject: [PATCH 2297/2720] Add tensorflow_gan repo in setup.py since it will
 be removed from tf.contrib.

PiperOrigin-RevId: 264492843
---
 setup.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/setup.py b/setup.py
index dc1de00a0..a4b6bad35 100644
--- a/setup.py
+++ b/setup.py
@@ -56,6 +56,7 @@
         'six',
         'sympy',
         'tensorflow-datasets',
+        'tensorflow-gan',
         'tensorflow-probability',
         'tqdm',
     ],

From 46cf96bfa4afba6afcdf3b72a76789af8be1381f Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 20 Aug 2019 18:19:29 -0700
Subject: [PATCH 2298/2720] Changes to be able to use and generate proto files
 externally.

- Service options : deadline and fail_fast don't seem to be supported
  externally by gRPC, remove them.
- Script added to generate the proto and service files (generate_py_proto.sh)
- The generated files are added to the source code in order for Travis to work.
  env_service_generated_pb2.py and env_service_generated_pb2_grpc.py

PiperOrigin-RevId: 264512062
---
 oss_scripts/generate_py_proto.sh              |  78 ++
 tensor2tensor/envs/__init__.py                |  23 +-
 tensor2tensor/envs/env_service.proto          |  16 +-
 .../envs/env_service_generated_pb2.py         | 924 ++++++++++++++++++
 .../envs/env_service_generated_pb2_grpc.py    | 131 +++
 5 files changed, 1157 insertions(+), 15 deletions(-)
 create mode 100755 oss_scripts/generate_py_proto.sh
 create mode 100644 tensor2tensor/envs/env_service_generated_pb2.py
 create mode 100644 tensor2tensor/envs/env_service_generated_pb2_grpc.py

diff --git a/oss_scripts/generate_py_proto.sh b/oss_scripts/generate_py_proto.sh
new file mode 100755
index 000000000..b72c0b75e
--- /dev/null
+++ b/oss_scripts/generate_py_proto.sh
@@ -0,0 +1,78 @@
+#!/bin/bash
+
+# This script use the protoc compiler to generate the python code of the
+# all of our proto files.
+
+
+# Function to prepend a pylint directive to skip the generated python file.
+function pylint_skip_file() {
+  local file_name=$1
+  printf "%s\n%s" "# pylint: skip-file" "$(cat ${file_name})" > ${file_name}
+}
+
+
+# Setup tmp directories
+TMP_DIR=$(mktemp -d)
+TMP_TF_DIR=${TMP_DIR}/tensorflow
+TMP_T2T_DIR="$PWD"
+
+echo "Temporary directory created: "
+echo ${TMP_DIR}
+
+
+TMP_T2T_PROTO_DIR="${TMP_T2T_DIR}/tensor2tensor/envs"
+ENV_SERVICE_PROTO="${TMP_T2T_PROTO_DIR}/env_service.proto"
+if [ ! -f ${ENV_SERVICE_PROTO} ]; then
+    echo "${ENV_SERVICE_PROTO} not found."
+    echo "Please run this script from the appropriate root directory."
+fi
+
+# Clone tensorflow repository.
+git clone https://github.com/tensorflow/tensorflow.git ${TMP_TF_DIR}
+
+# Install gRPC tools.
+pip install grpcio-tools
+
+# Invoke the grpc protoc compiler on env_service.proto
+python -m grpc_tools.protoc \
+  --proto_path=${TMP_TF_DIR}/ \
+  --proto_path=${TMP_T2T_DIR}/ \
+  --python_out=${TMP_T2T_DIR}/ \
+  --grpc_python_out=${TMP_T2T_DIR}/ \
+  ${ENV_SERVICE_PROTO}
+
+# Add pylint ignore and name the file as generated.
+GENERATED_ENV_SERVICE_PY="${TMP_T2T_PROTO_DIR}/env_service_generated_pb2.py"
+GENERATED_ENV_SERVICE_GRPC_PY="${TMP_T2T_PROTO_DIR}/env_service_generated_pb2_grpc.py"
+mv ${TMP_T2T_PROTO_DIR}/env_service_pb2.py ${GENERATED_ENV_SERVICE_PY}
+mv ${TMP_T2T_PROTO_DIR}/env_service_pb2_grpc.py ${GENERATED_ENV_SERVICE_GRPC_PY}
+pylint_skip_file "${GENERATED_ENV_SERVICE_PY}"
+pylint_skip_file "${GENERATED_ENV_SERVICE_GRPC_PY}"
+
+
+LICENSING_TEXT=$(cat <<-END
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+END
+)
+
+function add_licensing_text() {
+  local file_name=$1
+  printf "%s\n%s" "${LICENSING_TEXT}" "$(cat ${file_name})" > ${file_name}
+}
+
+add_licensing_text "${GENERATED_ENV_SERVICE_PY}"
+add_licensing_text "${GENERATED_ENV_SERVICE_GRPC_PY}"
+
diff --git a/tensor2tensor/envs/__init__.py b/tensor2tensor/envs/__init__.py
index 9b78dafac..3c335c2f2 100644
--- a/tensor2tensor/envs/__init__.py
+++ b/tensor2tensor/envs/__init__.py
@@ -15,9 +15,25 @@
 
 """Environments defined in T2T. Imports here force registration."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
+# Proto imports.
+
+
+# pylint: disable=g-import-not-at-top,g-statement-before-imports
+def _get_env_service():
+  from tensor2tensor.envs import env_service_generated_pb2 as env_service_pb2_
+  return env_service_pb2_
+
+
+def _get_env_service_grpc():
+  from tensor2tensor.envs import env_service_generated_pb2_grpc as env_service_pb2_grpc_
+  return env_service_pb2_grpc_
+# pylint: enable=g-import-not-at-top
+
+
+env_service_pb2 = _get_env_service()  # pylint: disable=invalid-name
+env_service_pb2_grpc = _get_env_service_grpc()  # pylint: disable=invalid-name
+del _get_env_service, _get_env_service_grpc
+# pylint: enable=g-statement-before-imports
 
 from gym.envs.registration import register
 
@@ -38,3 +54,4 @@ def register_env(env_class):
 # TODO(afrozm): Register TicTacToeEnv the same way.
 # register_env(tic_tac_toe_env.TicTacToeEnv)
 ClientEnv = register_env(client_env.ClientEnv)  # pylint: disable=invalid-name
+
diff --git a/tensor2tensor/envs/env_service.proto b/tensor2tensor/envs/env_service.proto
index e433204dc..46556b2e0 100644
--- a/tensor2tensor/envs/env_service.proto
+++ b/tensor2tensor/envs/env_service.proto
@@ -1,11 +1,11 @@
 syntax = "proto3";
 option cc_enable_arenas = true;
 
-package third_party.py.tensor2tensor.trax.rlax.envs;
+package tensor2tensor.trax.rlax.envs;
 
-import "third_party/tensorflow/core/framework/tensor.proto";
-import "third_party/tensorflow/core/framework/tensor_shape.proto";
-import "third_party/tensorflow/core/framework/types.proto";
+import "tensorflow/core/framework/tensor.proto";
+import "tensorflow/core/framework/tensor_shape.proto";
+import "tensorflow/core/framework/types.proto";
 
 // We use tensorflow.TensorProto to represent numpy arrays.
 
@@ -89,30 +89,22 @@ message EnvInfoResponse {
 service EnvService {
   // Reset
   rpc Reset(ResetRequest) returns (ResetResponse) {
-    option fail_fast = true;
   }
 
   // Step
   rpc Step(StepRequest) returns (StepResponse) {
-    option fail_fast = true;
   }
 
   // Close
   rpc Close(CloseRequest) returns (CloseResponse) {
-    option fail_fast = true;
-    option deadline = 10;
   }
 
   // Render
   rpc Render(RenderRequest) returns (RenderResponse) {
-    option fail_fast = true;
-    option deadline = 10;
   }
 
   // Observation and Action Space.
   rpc GetEnvInfo(EnvInfoRequest) returns (EnvInfoResponse) {
-    option fail_fast = true;
-    option deadline = 10;
   }
 }
 
diff --git a/tensor2tensor/envs/env_service_generated_pb2.py b/tensor2tensor/envs/env_service_generated_pb2.py
new file mode 100644
index 000000000..4efd96d89
--- /dev/null
+++ b/tensor2tensor/envs/env_service_generated_pb2.py
@@ -0,0 +1,924 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+
+# pylint: skip-file
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: tensor2tensor/envs/env_service.proto
+
+import sys
+_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from tensorflow.core.framework import tensor_pb2 as tensorflow_dot_core_dot_framework_dot_tensor__pb2
+from tensorflow.core.framework import tensor_shape_pb2 as tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2
+from tensorflow.core.framework import types_pb2 as tensorflow_dot_core_dot_framework_dot_types__pb2
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='tensor2tensor/envs/env_service.proto',
+  package='tensor2tensor.trax.rlax.envs',
+  syntax='proto3',
+  serialized_options=_b('\370\001\001'),
+  serialized_pb=_b('\n$tensor2tensor/envs/env_service.proto\x12\x1ctensor2tensor.trax.rlax.envs\x1a&tensorflow/core/framework/tensor.proto\x1a,tensorflow/core/framework/tensor_shape.proto\x1a%tensorflow/core/framework/types.proto\"d\n\x06\x41\x63tion\x12\x19\n\x0f\x64iscrete_action\x18\x01 \x01(\x03H\x00\x12\x34\n\x11\x63ontinuous_action\x18\x02 \x01(\x0b\x32\x17.tensorflow.TensorProtoH\x00\x42\t\n\x07payload\";\n\x0bObservation\x12,\n\x0bobservation\x18\x01 \x01(\x0b\x32\x17.tensorflow.TensorProto\"y\n\x04Info\x12\x41\n\x08info_map\x18\x01 \x03(\x0b\x32/.tensor2tensor.trax.rlax.envs.Info.InfoMapEntry\x1a.\n\x0cInfoMapEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x01:\x02\x38\x01\"C\n\x0bStepRequest\x12\x34\n\x06\x61\x63tion\x18\x01 \x01(\x0b\x32$.tensor2tensor.trax.rlax.envs.Action\"\x9e\x01\n\x0cStepResponse\x12>\n\x0bobservation\x18\x01 \x01(\x0b\x32).tensor2tensor.trax.rlax.envs.Observation\x12\x0e\n\x06reward\x18\x02 \x01(\x01\x12\x0c\n\x04\x64one\x18\x03 \x01(\x08\x12\x30\n\x04info\x18\x04 \x01(\x0b\x32\".tensor2tensor.trax.rlax.envs.Info\"\x0e\n\x0cResetRequest\"O\n\rResetResponse\x12>\n\x0bobservation\x18\x01 \x01(\x0b\x32).tensor2tensor.trax.rlax.envs.Observation\"\x0e\n\x0c\x43loseRequest\"\x0f\n\rCloseResponse\"\x1d\n\rRenderRequest\x12\x0c\n\x04mode\x18\x01 \x01(\t\"P\n\x0eRenderResponse\x12>\n\x0bobservation\x18\x01 \x01(\x0b\x32).tensor2tensor.trax.rlax.envs.Observation\"\x10\n\x0e\x45nvInfoRequest\"\xa9\x01\n\x08SpaceBox\x12#\n\x05\x64type\x18\x01 \x01(\x0e\x32\x14.tensorflow.DataType\x12+\n\x05shape\x18\x02 \x01(\x0b\x32\x1c.tensorflow.TensorShapeProto\x12$\n\x03low\x18\x03 \x01(\x0b\x32\x17.tensorflow.TensorProto\x12%\n\x04high\x18\x04 \x01(\x0b\x32\x17.tensorflow.TensorProto\"$\n\rSpaceDiscrete\x12\x13\n\x0bnum_actions\x18\x01 \x01(\x05\"\xae\x01\n\x08GymSpace\x12\x1d\n\x13unimplemented_space\x18\x01 \x01(\x08H\x00\x12\x35\n\x03\x62ox\x18\x02 \x01(\x0b\x32&.tensor2tensor.trax.rlax.envs.SpaceBoxH\x00\x12?\n\x08\x64iscrete\x18\x03 \x01(\x0b\x32+.tensor2tensor.trax.rlax.envs.SpaceDiscreteH\x00\x42\x0b\n\tgym_space\"(\n\x0bRewardRange\x12\x0b\n\x03low\x18\x01 \x01(\x01\x12\x0c\n\x04high\x18\x02 \x01(\x01\"\xe7\x01\n\x0f\x45nvInfoResponse\x12\x41\n\x11observation_space\x18\x01 \x01(\x0b\x32&.tensor2tensor.trax.rlax.envs.GymSpace\x12<\n\x0c\x61\x63tion_space\x18\x02 \x01(\x0b\x32&.tensor2tensor.trax.rlax.envs.GymSpace\x12?\n\x0creward_range\x18\x03 \x01(\x0b\x32).tensor2tensor.trax.rlax.envs.RewardRange\x12\x12\n\nbatch_size\x18\x04 \x01(\x03\x32\x89\x04\n\nEnvService\x12\x62\n\x05Reset\x12*.tensor2tensor.trax.rlax.envs.ResetRequest\x1a+.tensor2tensor.trax.rlax.envs.ResetResponse\"\x00\x12_\n\x04Step\x12).tensor2tensor.trax.rlax.envs.StepRequest\x1a*.tensor2tensor.trax.rlax.envs.StepResponse\"\x00\x12\x62\n\x05\x43lose\x12*.tensor2tensor.trax.rlax.envs.CloseRequest\x1a+.tensor2tensor.trax.rlax.envs.CloseResponse\"\x00\x12\x65\n\x06Render\x12+.tensor2tensor.trax.rlax.envs.RenderRequest\x1a,.tensor2tensor.trax.rlax.envs.RenderResponse\"\x00\x12k\n\nGetEnvInfo\x12,.tensor2tensor.trax.rlax.envs.EnvInfoRequest\x1a-.tensor2tensor.trax.rlax.envs.EnvInfoResponse\"\x00\x42\x03\xf8\x01\x01\x62\x06proto3')
+  ,
+  dependencies=[tensorflow_dot_core_dot_framework_dot_tensor__pb2.DESCRIPTOR,tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2.DESCRIPTOR,tensorflow_dot_core_dot_framework_dot_types__pb2.DESCRIPTOR,])
+
+
+
+
+_ACTION = _descriptor.Descriptor(
+  name='Action',
+  full_name='tensor2tensor.trax.rlax.envs.Action',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='discrete_action', full_name='tensor2tensor.trax.rlax.envs.Action.discrete_action', index=0,
+      number=1, type=3, cpp_type=2, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+    _descriptor.FieldDescriptor(
+      name='continuous_action', full_name='tensor2tensor.trax.rlax.envs.Action.continuous_action', index=1,
+      number=2, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+    _descriptor.OneofDescriptor(
+      name='payload', full_name='tensor2tensor.trax.rlax.envs.Action.payload',
+      index=0, containing_type=None, fields=[]),
+  ],
+  serialized_start=195,
+  serialized_end=295,
+)
+
+
+_OBSERVATION = _descriptor.Descriptor(
+  name='Observation',
+  full_name='tensor2tensor.trax.rlax.envs.Observation',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='observation', full_name='tensor2tensor.trax.rlax.envs.Observation.observation', index=0,
+      number=1, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=297,
+  serialized_end=356,
+)
+
+
+_INFO_INFOMAPENTRY = _descriptor.Descriptor(
+  name='InfoMapEntry',
+  full_name='tensor2tensor.trax.rlax.envs.Info.InfoMapEntry',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='key', full_name='tensor2tensor.trax.rlax.envs.Info.InfoMapEntry.key', index=0,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+    _descriptor.FieldDescriptor(
+      name='value', full_name='tensor2tensor.trax.rlax.envs.Info.InfoMapEntry.value', index=1,
+      number=2, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=_b('8\001'),
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=433,
+  serialized_end=479,
+)
+
+_INFO = _descriptor.Descriptor(
+  name='Info',
+  full_name='tensor2tensor.trax.rlax.envs.Info',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='info_map', full_name='tensor2tensor.trax.rlax.envs.Info.info_map', index=0,
+      number=1, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+  ],
+  extensions=[
+  ],
+  nested_types=[_INFO_INFOMAPENTRY, ],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=358,
+  serialized_end=479,
+)
+
+
+_STEPREQUEST = _descriptor.Descriptor(
+  name='StepRequest',
+  full_name='tensor2tensor.trax.rlax.envs.StepRequest',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='action', full_name='tensor2tensor.trax.rlax.envs.StepRequest.action', index=0,
+      number=1, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=481,
+  serialized_end=548,
+)
+
+
+_STEPRESPONSE = _descriptor.Descriptor(
+  name='StepResponse',
+  full_name='tensor2tensor.trax.rlax.envs.StepResponse',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='observation', full_name='tensor2tensor.trax.rlax.envs.StepResponse.observation', index=0,
+      number=1, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+    _descriptor.FieldDescriptor(
+      name='reward', full_name='tensor2tensor.trax.rlax.envs.StepResponse.reward', index=1,
+      number=2, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+    _descriptor.FieldDescriptor(
+      name='done', full_name='tensor2tensor.trax.rlax.envs.StepResponse.done', index=2,
+      number=3, type=8, cpp_type=7, label=1,
+      has_default_value=False, default_value=False,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+    _descriptor.FieldDescriptor(
+      name='info', full_name='tensor2tensor.trax.rlax.envs.StepResponse.info', index=3,
+      number=4, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=551,
+  serialized_end=709,
+)
+
+
+_RESETREQUEST = _descriptor.Descriptor(
+  name='ResetRequest',
+  full_name='tensor2tensor.trax.rlax.envs.ResetRequest',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=711,
+  serialized_end=725,
+)
+
+
+_RESETRESPONSE = _descriptor.Descriptor(
+  name='ResetResponse',
+  full_name='tensor2tensor.trax.rlax.envs.ResetResponse',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='observation', full_name='tensor2tensor.trax.rlax.envs.ResetResponse.observation', index=0,
+      number=1, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=727,
+  serialized_end=806,
+)
+
+
+_CLOSEREQUEST = _descriptor.Descriptor(
+  name='CloseRequest',
+  full_name='tensor2tensor.trax.rlax.envs.CloseRequest',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=808,
+  serialized_end=822,
+)
+
+
+_CLOSERESPONSE = _descriptor.Descriptor(
+  name='CloseResponse',
+  full_name='tensor2tensor.trax.rlax.envs.CloseResponse',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=824,
+  serialized_end=839,
+)
+
+
+_RENDERREQUEST = _descriptor.Descriptor(
+  name='RenderRequest',
+  full_name='tensor2tensor.trax.rlax.envs.RenderRequest',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='mode', full_name='tensor2tensor.trax.rlax.envs.RenderRequest.mode', index=0,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=841,
+  serialized_end=870,
+)
+
+
+_RENDERRESPONSE = _descriptor.Descriptor(
+  name='RenderResponse',
+  full_name='tensor2tensor.trax.rlax.envs.RenderResponse',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='observation', full_name='tensor2tensor.trax.rlax.envs.RenderResponse.observation', index=0,
+      number=1, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=872,
+  serialized_end=952,
+)
+
+
+_ENVINFOREQUEST = _descriptor.Descriptor(
+  name='EnvInfoRequest',
+  full_name='tensor2tensor.trax.rlax.envs.EnvInfoRequest',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=954,
+  serialized_end=970,
+)
+
+
+_SPACEBOX = _descriptor.Descriptor(
+  name='SpaceBox',
+  full_name='tensor2tensor.trax.rlax.envs.SpaceBox',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='dtype', full_name='tensor2tensor.trax.rlax.envs.SpaceBox.dtype', index=0,
+      number=1, type=14, cpp_type=8, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+    _descriptor.FieldDescriptor(
+      name='shape', full_name='tensor2tensor.trax.rlax.envs.SpaceBox.shape', index=1,
+      number=2, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+    _descriptor.FieldDescriptor(
+      name='low', full_name='tensor2tensor.trax.rlax.envs.SpaceBox.low', index=2,
+      number=3, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+    _descriptor.FieldDescriptor(
+      name='high', full_name='tensor2tensor.trax.rlax.envs.SpaceBox.high', index=3,
+      number=4, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=973,
+  serialized_end=1142,
+)
+
+
+_SPACEDISCRETE = _descriptor.Descriptor(
+  name='SpaceDiscrete',
+  full_name='tensor2tensor.trax.rlax.envs.SpaceDiscrete',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='num_actions', full_name='tensor2tensor.trax.rlax.envs.SpaceDiscrete.num_actions', index=0,
+      number=1, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=1144,
+  serialized_end=1180,
+)
+
+
+_GYMSPACE = _descriptor.Descriptor(
+  name='GymSpace',
+  full_name='tensor2tensor.trax.rlax.envs.GymSpace',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='unimplemented_space', full_name='tensor2tensor.trax.rlax.envs.GymSpace.unimplemented_space', index=0,
+      number=1, type=8, cpp_type=7, label=1,
+      has_default_value=False, default_value=False,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+    _descriptor.FieldDescriptor(
+      name='box', full_name='tensor2tensor.trax.rlax.envs.GymSpace.box', index=1,
+      number=2, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+    _descriptor.FieldDescriptor(
+      name='discrete', full_name='tensor2tensor.trax.rlax.envs.GymSpace.discrete', index=2,
+      number=3, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+    _descriptor.OneofDescriptor(
+      name='gym_space', full_name='tensor2tensor.trax.rlax.envs.GymSpace.gym_space',
+      index=0, containing_type=None, fields=[]),
+  ],
+  serialized_start=1183,
+  serialized_end=1357,
+)
+
+
+_REWARDRANGE = _descriptor.Descriptor(
+  name='RewardRange',
+  full_name='tensor2tensor.trax.rlax.envs.RewardRange',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='low', full_name='tensor2tensor.trax.rlax.envs.RewardRange.low', index=0,
+      number=1, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+    _descriptor.FieldDescriptor(
+      name='high', full_name='tensor2tensor.trax.rlax.envs.RewardRange.high', index=1,
+      number=2, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=1359,
+  serialized_end=1399,
+)
+
+
+_ENVINFORESPONSE = _descriptor.Descriptor(
+  name='EnvInfoResponse',
+  full_name='tensor2tensor.trax.rlax.envs.EnvInfoResponse',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='observation_space', full_name='tensor2tensor.trax.rlax.envs.EnvInfoResponse.observation_space', index=0,
+      number=1, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+    _descriptor.FieldDescriptor(
+      name='action_space', full_name='tensor2tensor.trax.rlax.envs.EnvInfoResponse.action_space', index=1,
+      number=2, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+    _descriptor.FieldDescriptor(
+      name='reward_range', full_name='tensor2tensor.trax.rlax.envs.EnvInfoResponse.reward_range', index=2,
+      number=3, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+    _descriptor.FieldDescriptor(
+      name='batch_size', full_name='tensor2tensor.trax.rlax.envs.EnvInfoResponse.batch_size', index=3,
+      number=4, type=3, cpp_type=2, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=1402,
+  serialized_end=1633,
+)
+
+_ACTION.fields_by_name['continuous_action'].message_type = tensorflow_dot_core_dot_framework_dot_tensor__pb2._TENSORPROTO
+_ACTION.oneofs_by_name['payload'].fields.append(
+  _ACTION.fields_by_name['discrete_action'])
+_ACTION.fields_by_name['discrete_action'].containing_oneof = _ACTION.oneofs_by_name['payload']
+_ACTION.oneofs_by_name['payload'].fields.append(
+  _ACTION.fields_by_name['continuous_action'])
+_ACTION.fields_by_name['continuous_action'].containing_oneof = _ACTION.oneofs_by_name['payload']
+_OBSERVATION.fields_by_name['observation'].message_type = tensorflow_dot_core_dot_framework_dot_tensor__pb2._TENSORPROTO
+_INFO_INFOMAPENTRY.containing_type = _INFO
+_INFO.fields_by_name['info_map'].message_type = _INFO_INFOMAPENTRY
+_STEPREQUEST.fields_by_name['action'].message_type = _ACTION
+_STEPRESPONSE.fields_by_name['observation'].message_type = _OBSERVATION
+_STEPRESPONSE.fields_by_name['info'].message_type = _INFO
+_RESETRESPONSE.fields_by_name['observation'].message_type = _OBSERVATION
+_RENDERRESPONSE.fields_by_name['observation'].message_type = _OBSERVATION
+_SPACEBOX.fields_by_name['dtype'].enum_type = tensorflow_dot_core_dot_framework_dot_types__pb2._DATATYPE
+_SPACEBOX.fields_by_name['shape'].message_type = tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2._TENSORSHAPEPROTO
+_SPACEBOX.fields_by_name['low'].message_type = tensorflow_dot_core_dot_framework_dot_tensor__pb2._TENSORPROTO
+_SPACEBOX.fields_by_name['high'].message_type = tensorflow_dot_core_dot_framework_dot_tensor__pb2._TENSORPROTO
+_GYMSPACE.fields_by_name['box'].message_type = _SPACEBOX
+_GYMSPACE.fields_by_name['discrete'].message_type = _SPACEDISCRETE
+_GYMSPACE.oneofs_by_name['gym_space'].fields.append(
+  _GYMSPACE.fields_by_name['unimplemented_space'])
+_GYMSPACE.fields_by_name['unimplemented_space'].containing_oneof = _GYMSPACE.oneofs_by_name['gym_space']
+_GYMSPACE.oneofs_by_name['gym_space'].fields.append(
+  _GYMSPACE.fields_by_name['box'])
+_GYMSPACE.fields_by_name['box'].containing_oneof = _GYMSPACE.oneofs_by_name['gym_space']
+_GYMSPACE.oneofs_by_name['gym_space'].fields.append(
+  _GYMSPACE.fields_by_name['discrete'])
+_GYMSPACE.fields_by_name['discrete'].containing_oneof = _GYMSPACE.oneofs_by_name['gym_space']
+_ENVINFORESPONSE.fields_by_name['observation_space'].message_type = _GYMSPACE
+_ENVINFORESPONSE.fields_by_name['action_space'].message_type = _GYMSPACE
+_ENVINFORESPONSE.fields_by_name['reward_range'].message_type = _REWARDRANGE
+DESCRIPTOR.message_types_by_name['Action'] = _ACTION
+DESCRIPTOR.message_types_by_name['Observation'] = _OBSERVATION
+DESCRIPTOR.message_types_by_name['Info'] = _INFO
+DESCRIPTOR.message_types_by_name['StepRequest'] = _STEPREQUEST
+DESCRIPTOR.message_types_by_name['StepResponse'] = _STEPRESPONSE
+DESCRIPTOR.message_types_by_name['ResetRequest'] = _RESETREQUEST
+DESCRIPTOR.message_types_by_name['ResetResponse'] = _RESETRESPONSE
+DESCRIPTOR.message_types_by_name['CloseRequest'] = _CLOSEREQUEST
+DESCRIPTOR.message_types_by_name['CloseResponse'] = _CLOSERESPONSE
+DESCRIPTOR.message_types_by_name['RenderRequest'] = _RENDERREQUEST
+DESCRIPTOR.message_types_by_name['RenderResponse'] = _RENDERRESPONSE
+DESCRIPTOR.message_types_by_name['EnvInfoRequest'] = _ENVINFOREQUEST
+DESCRIPTOR.message_types_by_name['SpaceBox'] = _SPACEBOX
+DESCRIPTOR.message_types_by_name['SpaceDiscrete'] = _SPACEDISCRETE
+DESCRIPTOR.message_types_by_name['GymSpace'] = _GYMSPACE
+DESCRIPTOR.message_types_by_name['RewardRange'] = _REWARDRANGE
+DESCRIPTOR.message_types_by_name['EnvInfoResponse'] = _ENVINFORESPONSE
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+Action = _reflection.GeneratedProtocolMessageType('Action', (_message.Message,), {
+  'DESCRIPTOR' : _ACTION,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.Action)
+  })
+_sym_db.RegisterMessage(Action)
+
+Observation = _reflection.GeneratedProtocolMessageType('Observation', (_message.Message,), {
+  'DESCRIPTOR' : _OBSERVATION,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.Observation)
+  })
+_sym_db.RegisterMessage(Observation)
+
+Info = _reflection.GeneratedProtocolMessageType('Info', (_message.Message,), {
+
+  'InfoMapEntry' : _reflection.GeneratedProtocolMessageType('InfoMapEntry', (_message.Message,), {
+    'DESCRIPTOR' : _INFO_INFOMAPENTRY,
+    '__module__' : 'tensor2tensor.envs.env_service_pb2'
+    # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.Info.InfoMapEntry)
+    })
+  ,
+  'DESCRIPTOR' : _INFO,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.Info)
+  })
+_sym_db.RegisterMessage(Info)
+_sym_db.RegisterMessage(Info.InfoMapEntry)
+
+StepRequest = _reflection.GeneratedProtocolMessageType('StepRequest', (_message.Message,), {
+  'DESCRIPTOR' : _STEPREQUEST,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.StepRequest)
+  })
+_sym_db.RegisterMessage(StepRequest)
+
+StepResponse = _reflection.GeneratedProtocolMessageType('StepResponse', (_message.Message,), {
+  'DESCRIPTOR' : _STEPRESPONSE,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.StepResponse)
+  })
+_sym_db.RegisterMessage(StepResponse)
+
+ResetRequest = _reflection.GeneratedProtocolMessageType('ResetRequest', (_message.Message,), {
+  'DESCRIPTOR' : _RESETREQUEST,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.ResetRequest)
+  })
+_sym_db.RegisterMessage(ResetRequest)
+
+ResetResponse = _reflection.GeneratedProtocolMessageType('ResetResponse', (_message.Message,), {
+  'DESCRIPTOR' : _RESETRESPONSE,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.ResetResponse)
+  })
+_sym_db.RegisterMessage(ResetResponse)
+
+CloseRequest = _reflection.GeneratedProtocolMessageType('CloseRequest', (_message.Message,), {
+  'DESCRIPTOR' : _CLOSEREQUEST,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.CloseRequest)
+  })
+_sym_db.RegisterMessage(CloseRequest)
+
+CloseResponse = _reflection.GeneratedProtocolMessageType('CloseResponse', (_message.Message,), {
+  'DESCRIPTOR' : _CLOSERESPONSE,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.CloseResponse)
+  })
+_sym_db.RegisterMessage(CloseResponse)
+
+RenderRequest = _reflection.GeneratedProtocolMessageType('RenderRequest', (_message.Message,), {
+  'DESCRIPTOR' : _RENDERREQUEST,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.RenderRequest)
+  })
+_sym_db.RegisterMessage(RenderRequest)
+
+RenderResponse = _reflection.GeneratedProtocolMessageType('RenderResponse', (_message.Message,), {
+  'DESCRIPTOR' : _RENDERRESPONSE,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.RenderResponse)
+  })
+_sym_db.RegisterMessage(RenderResponse)
+
+EnvInfoRequest = _reflection.GeneratedProtocolMessageType('EnvInfoRequest', (_message.Message,), {
+  'DESCRIPTOR' : _ENVINFOREQUEST,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.EnvInfoRequest)
+  })
+_sym_db.RegisterMessage(EnvInfoRequest)
+
+SpaceBox = _reflection.GeneratedProtocolMessageType('SpaceBox', (_message.Message,), {
+  'DESCRIPTOR' : _SPACEBOX,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.SpaceBox)
+  })
+_sym_db.RegisterMessage(SpaceBox)
+
+SpaceDiscrete = _reflection.GeneratedProtocolMessageType('SpaceDiscrete', (_message.Message,), {
+  'DESCRIPTOR' : _SPACEDISCRETE,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.SpaceDiscrete)
+  })
+_sym_db.RegisterMessage(SpaceDiscrete)
+
+GymSpace = _reflection.GeneratedProtocolMessageType('GymSpace', (_message.Message,), {
+  'DESCRIPTOR' : _GYMSPACE,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.GymSpace)
+  })
+_sym_db.RegisterMessage(GymSpace)
+
+RewardRange = _reflection.GeneratedProtocolMessageType('RewardRange', (_message.Message,), {
+  'DESCRIPTOR' : _REWARDRANGE,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.RewardRange)
+  })
+_sym_db.RegisterMessage(RewardRange)
+
+EnvInfoResponse = _reflection.GeneratedProtocolMessageType('EnvInfoResponse', (_message.Message,), {
+  'DESCRIPTOR' : _ENVINFORESPONSE,
+  '__module__' : 'tensor2tensor.envs.env_service_pb2'
+  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.EnvInfoResponse)
+  })
+_sym_db.RegisterMessage(EnvInfoResponse)
+
+
+DESCRIPTOR._options = None
+_INFO_INFOMAPENTRY._options = None
+
+_ENVSERVICE = _descriptor.ServiceDescriptor(
+  name='EnvService',
+  full_name='tensor2tensor.trax.rlax.envs.EnvService',
+  file=DESCRIPTOR,
+  index=0,
+  serialized_options=None,
+  serialized_start=1636,
+  serialized_end=2157,
+  methods=[
+  _descriptor.MethodDescriptor(
+    name='Reset',
+    full_name='tensor2tensor.trax.rlax.envs.EnvService.Reset',
+    index=0,
+    containing_service=None,
+    input_type=_RESETREQUEST,
+    output_type=_RESETRESPONSE,
+    serialized_options=None,
+  ),
+  _descriptor.MethodDescriptor(
+    name='Step',
+    full_name='tensor2tensor.trax.rlax.envs.EnvService.Step',
+    index=1,
+    containing_service=None,
+    input_type=_STEPREQUEST,
+    output_type=_STEPRESPONSE,
+    serialized_options=None,
+  ),
+  _descriptor.MethodDescriptor(
+    name='Close',
+    full_name='tensor2tensor.trax.rlax.envs.EnvService.Close',
+    index=2,
+    containing_service=None,
+    input_type=_CLOSEREQUEST,
+    output_type=_CLOSERESPONSE,
+    serialized_options=None,
+  ),
+  _descriptor.MethodDescriptor(
+    name='Render',
+    full_name='tensor2tensor.trax.rlax.envs.EnvService.Render',
+    index=3,
+    containing_service=None,
+    input_type=_RENDERREQUEST,
+    output_type=_RENDERRESPONSE,
+    serialized_options=None,
+  ),
+  _descriptor.MethodDescriptor(
+    name='GetEnvInfo',
+    full_name='tensor2tensor.trax.rlax.envs.EnvService.GetEnvInfo',
+    index=4,
+    containing_service=None,
+    input_type=_ENVINFOREQUEST,
+    output_type=_ENVINFORESPONSE,
+    serialized_options=None,
+  ),
+])
+_sym_db.RegisterServiceDescriptor(_ENVSERVICE)
+
+DESCRIPTOR.services_by_name['EnvService'] = _ENVSERVICE
+
+# @@protoc_insertion_point(module_scope)
\ No newline at end of file
diff --git a/tensor2tensor/envs/env_service_generated_pb2_grpc.py b/tensor2tensor/envs/env_service_generated_pb2_grpc.py
new file mode 100644
index 000000000..378939f00
--- /dev/null
+++ b/tensor2tensor/envs/env_service_generated_pb2_grpc.py
@@ -0,0 +1,131 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+
+# pylint: skip-file
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+import grpc
+
+from tensor2tensor.envs import env_service_pb2 as tensor2tensor_dot_envs_dot_env__service__pb2
+
+
+class EnvServiceStub(object):
+  # missing associated documentation comment in .proto file
+  pass
+
+  def __init__(self, channel):
+    """Constructor.
+
+    Args:
+      channel: A grpc.Channel.
+    """
+    self.Reset = channel.unary_unary(
+        '/tensor2tensor.trax.rlax.envs.EnvService/Reset',
+        request_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.ResetRequest.SerializeToString,
+        response_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.ResetResponse.FromString,
+        )
+    self.Step = channel.unary_unary(
+        '/tensor2tensor.trax.rlax.envs.EnvService/Step',
+        request_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.StepRequest.SerializeToString,
+        response_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.StepResponse.FromString,
+        )
+    self.Close = channel.unary_unary(
+        '/tensor2tensor.trax.rlax.envs.EnvService/Close',
+        request_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.CloseRequest.SerializeToString,
+        response_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.CloseResponse.FromString,
+        )
+    self.Render = channel.unary_unary(
+        '/tensor2tensor.trax.rlax.envs.EnvService/Render',
+        request_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.RenderRequest.SerializeToString,
+        response_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.RenderResponse.FromString,
+        )
+    self.GetEnvInfo = channel.unary_unary(
+        '/tensor2tensor.trax.rlax.envs.EnvService/GetEnvInfo',
+        request_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.EnvInfoRequest.SerializeToString,
+        response_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.EnvInfoResponse.FromString,
+        )
+
+
+class EnvServiceServicer(object):
+  # missing associated documentation comment in .proto file
+  pass
+
+  def Reset(self, request, context):
+    """Reset
+    """
+    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+    context.set_details('Method not implemented!')
+    raise NotImplementedError('Method not implemented!')
+
+  def Step(self, request, context):
+    """Step
+    """
+    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+    context.set_details('Method not implemented!')
+    raise NotImplementedError('Method not implemented!')
+
+  def Close(self, request, context):
+    """Close
+    """
+    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+    context.set_details('Method not implemented!')
+    raise NotImplementedError('Method not implemented!')
+
+  def Render(self, request, context):
+    """Render
+    """
+    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+    context.set_details('Method not implemented!')
+    raise NotImplementedError('Method not implemented!')
+
+  def GetEnvInfo(self, request, context):
+    """Observation and Action Space.
+    """
+    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+    context.set_details('Method not implemented!')
+    raise NotImplementedError('Method not implemented!')
+
+
+def add_EnvServiceServicer_to_server(servicer, server):
+  rpc_method_handlers = {
+      'Reset': grpc.unary_unary_rpc_method_handler(
+          servicer.Reset,
+          request_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.ResetRequest.FromString,
+          response_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.ResetResponse.SerializeToString,
+      ),
+      'Step': grpc.unary_unary_rpc_method_handler(
+          servicer.Step,
+          request_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.StepRequest.FromString,
+          response_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.StepResponse.SerializeToString,
+      ),
+      'Close': grpc.unary_unary_rpc_method_handler(
+          servicer.Close,
+          request_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.CloseRequest.FromString,
+          response_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.CloseResponse.SerializeToString,
+      ),
+      'Render': grpc.unary_unary_rpc_method_handler(
+          servicer.Render,
+          request_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.RenderRequest.FromString,
+          response_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.RenderResponse.SerializeToString,
+      ),
+      'GetEnvInfo': grpc.unary_unary_rpc_method_handler(
+          servicer.GetEnvInfo,
+          request_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.EnvInfoRequest.FromString,
+          response_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.EnvInfoResponse.SerializeToString,
+      ),
+  }
+  generic_handler = grpc.method_handlers_generic_handler(
+      'tensor2tensor.trax.rlax.envs.EnvService', rpc_method_handlers)
+  server.add_generic_rpc_handlers((generic_handler,))
\ No newline at end of file

From 4de97082ae159086973e590d284925ebfc69a944 Mon Sep 17 00:00:00 2001
From: Peng Wang <wangpeng@google.com>
Date: Wed, 21 Aug 2019 10:53:09 -0700
Subject: [PATCH 2299/2720] Add state to layers

PiperOrigin-RevId: 264644849
---
 tensor2tensor/envs/env_problem_utils.py       |  10 +-
 tensor2tensor/envs/env_problem_utils_test.py  |   6 +-
 tensor2tensor/trax/layers/attention.py        |   2 +-
 tensor2tensor/trax/layers/attention_test.py   |   4 +-
 tensor2tensor/trax/layers/base.py             |  73 +++++---
 tensor2tensor/trax/layers/base_test.py        |  12 +-
 tensor2tensor/trax/layers/combinators.py      |  53 ++++--
 tensor2tensor/trax/layers/convolution.py      |   6 +-
 tensor2tensor/trax/layers/core.py             |  12 +-
 tensor2tensor/trax/layers/core_test.py        |  16 +-
 tensor2tensor/trax/layers/normalization.py    | 135 +++++++++------
 .../trax/layers/normalization_test.py         |  47 ++++-
 tensor2tensor/trax/layers/reversible.py       |   2 +-
 tensor2tensor/trax/models/atari_cnn_test.py   |   8 +-
 tensor2tensor/trax/models/resnet.py           |  78 ++++-----
 tensor2tensor/trax/rl/ppo.py                  |  48 ++++--
 tensor2tensor/trax/rl/ppo_test.py             |  19 ++-
 tensor2tensor/trax/rl/ppo_trainer.py          | 101 ++++++-----
 tensor2tensor/trax/rl/ppo_trainer_test.py     |   7 +-
 .../trax/rl/simulated_env_problem.py          |  10 +-
 .../trax/rl/simulated_env_problem_test.py     |   4 +-
 tensor2tensor/trax/trax.py                    | 160 ++++++++++--------
 tensor2tensor/trax/trax_test.py               |  50 +++++-
 23 files changed, 546 insertions(+), 317 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 5a686ce83..bfe232f7b 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -57,6 +57,7 @@ def play_env_problem_with_policy(env,
                                  num_trajectories=1,
                                  max_timestep=None,
                                  reset=True,
+                                 state=None,
                                  rng=None,
                                  temperature=1.0,
                                  boundary=32,
@@ -73,7 +74,8 @@ def play_env_problem_with_policy(env,
       trajectory that exceeds this time put it in the completed bin, and *dont*
       reset the env.
     reset: bool, true if we want to reset the envs. The envs are also reset if
-      max_max_timestep is None or < 0
+      max_max_timestep is None or < 0.
+    state: the state for `policy_fn`.
     rng: jax rng, splittable.
     temperature: float, temperature used in Gumbel sampling.
     boundary: int, pad the sequences to the multiples of this number.
@@ -118,8 +120,8 @@ def gumbel_sample(log_probs):
     assert (B,) == lengths.shape
 
     t1 = time.time()
-    log_prob_actions, value_predictions, rng = policy_fun(
-        padded_observations, rng=rng)
+    log_prob_actions, value_predictions, state, rng = policy_fun(
+        padded_observations, state=state, rng=rng)
     policy_application_total_time += (time.time() - t1)
 
     assert (B, T) == log_prob_actions.shape[:2]
@@ -192,7 +194,7 @@ def gumbel_sample(log_probs):
   }
   timing_info = {k: round(1000 * v, 2) for k, v in timing_info.items()}
 
-  return completed_trajectories, num_done_trajectories, timing_info
+  return completed_trajectories, num_done_trajectories, timing_info, state
 
 
 def make_env(batch_size=1,
diff --git a/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index e7416bc1c..0bda965ad 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -53,7 +53,7 @@ def test_play_env_problem_with_policy(self):
     # Let's make sure that at-most 4 observations come to the policy function.
     len_history_for_policy = 4
 
-    def policy_fun(observations, rng=None):
+    def policy_fun(observations, state=None, rng=None):
       b, t = observations.shape[:2]
       # Assert that observations from time-step len_history_for_policy onwards
       # are zeros.
@@ -65,11 +65,11 @@ def policy_fun(observations, rng=None):
       p = np.random.uniform(size=(b, t, a))
       p = np.exp(p)
       p = p / np.sum(p, axis=-1, keepdims=True)
-      return np.log(p), np.log(p), rng
+      return np.log(p), np.log(p), state, rng
 
     max_timestep = 15
     num_trajectories = 2
-    trajectories, _, _ = env_problem_utils.play_env_problem_with_policy(
+    trajectories, _, _, _ = env_problem_utils.play_env_problem_with_policy(
         env,
         policy_fun,
         num_trajectories=num_trajectories,
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 4828939c8..f1ecd26b7 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -169,7 +169,7 @@ def AttentionQKV(d_feature, n_heads=1, dropout=0.0, mode='train'):
           core.Dense(d_feature),
       ),
       PureAttention(  # pylint: disable=no-value-for-parameter
-          d_feature=d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
+          n_heads=n_heads, dropout=dropout, mode=mode),
       core.Dense(d_feature),
   ]
 
diff --git a/tensor2tensor/trax/layers/attention_test.py b/tensor2tensor/trax/layers/attention_test.py
index 3517b5986..cb42732ed 100644
--- a/tensor2tensor/trax/layers/attention_test.py
+++ b/tensor2tensor/trax/layers/attention_test.py
@@ -30,7 +30,7 @@ def test_shift_right(self):
     # Test shifts right on axis=1
     layer = attention.ShiftRight()
     input_np = onp.arange(2*3*3).reshape(2, 3, 3)
-    output_np = layer(input_np)
+    output_np, _ = layer(input_np)
     self.assertEqual(input_np.shape, output_np.shape)
     self.assertAllEqual(onp.array([[[0, 0, 0],
                                     [0, 1, 2],
@@ -49,7 +49,7 @@ def test_shift_right_float(self):
     input_np /= 2.0
     self.assertEqual(input_np.dtype, onp.float32)
 
-    output_np = layer(input_np)
+    output_np, _ = layer(input_np)
     self.assertEqual(input_np.shape, output_np.shape)
     self.assertEqual(output_np.dtype, onp.float32)
 
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 52496b67f..79568de0e 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -81,7 +81,7 @@ def __repr__(self):
     else:
       return '{}[{}]'.format(class_str, fields_str)
 
-  def call(self, inputs, params=(), **kwargs):
+  def call(self, inputs, params=(), state=(), **kwargs):
     """Applies this layer to given activation tensors, using trainable params.
 
     Args:
@@ -94,6 +94,7 @@ def call(self, inputs, params=(), **kwargs):
           and one for each of this layer's sublayers. If a layer (or sublayer)
           has no trainable parameters, the corresponding params element is an
           empty tuple.
+      state: start state.
       **kwargs: Layer-specific keyword args.
 
     Returns:
@@ -106,6 +107,7 @@ def call(self, inputs, params=(), **kwargs):
     """
     raise NotImplementedError
 
+  # TODO(wangpeng): Should be called `new_parameters_and_state`.
   def new_parameters(self, input_shapes, input_dtype, rng):
     """Creates layer-specific parameters based on data shape, dtype and rng.
 
@@ -144,7 +146,7 @@ def has_custom_grad(self):
     """Whether to use custom gradients (in which case, see below)."""
     return False
 
-  def custom_grad(self, inputs, output, grad, params, **kwargs):
+  def custom_grad(self, inputs, output, grad, params, state, **kwargs):
     """Custom backward pass to propagate gradients in a custom way.
 
     Args:
@@ -153,6 +155,7 @@ def custom_grad(self, inputs, output, grad, params, **kwargs):
       grad: gradient signal (called cotangent in jax) computed based on
         subsequent layers. The structure and shape must match output.
       params: layer parameters
+      state: start state.
       **kwargs: kwargs for the layer
 
     Returns:
@@ -164,7 +167,7 @@ def custom_grad(self, inputs, output, grad, params, **kwargs):
 
   # End of subclassing interface, all functions below are internal.
 
-  def pseudo_call(self, pseudo_inputs, params):
+  def pseudo_call(self, pseudo_inputs, params, state):
     """Computes shapes and types this layer would produce for the given inputs.
 
     Args:
@@ -172,6 +175,7 @@ def pseudo_call(self, pseudo_inputs, params):
           or a tuple of ShapeType instances, following the same conventions as
           Layer.call's input arg.
       params: Parameters for this layer.
+      state: start state.
 
     Returns:
       A ShapeType instance representing the shape and type of the output (if
@@ -183,12 +187,12 @@ def pseudo_call(self, pseudo_inputs, params):
       # cause a large number of dropout masks to be computed and permanently
       # stored in global memory.
       rng = ShapeType(shape=(2,), dtype=onp.uint32)
-      def call_on_input(x, params, rng):
-        return self.call(x, params=params, rng=rng)
+      def call_on_input(x, params, state, rng):
+        return self.call(x, params=params, state=state, rng=rng)
       params_shapes = nested_map(
           params, lambda x: ShapeType(shape=x.shape, dtype=x.dtype))
       s = backend.eval_on_shapes(call_on_input)(pseudo_inputs,
-                                                params_shapes, rng)
+                                                params_shapes, state, rng)
       return s
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback(skip=3)
@@ -213,17 +217,23 @@ def initialize(self, input_shapes, input_dtype, rng):
     """
     try:
       # Initialize params once; store them for use when this layer is called.
+      # Needs to call new_parameters regardless of _init_finished because state
+      # also needs to be initialized. After jitting, graph pruning should be
+      # able to remove unnecessary computation.
+      # TODO(lukaszkaiser): Revisit this decision and see whether layers sharing
+      #   params should also share states.
+      params, state = self.new_parameters(input_shapes, input_dtype, rng)
       if not self._init_finished:
-        self._params = self.new_parameters(input_shapes, input_dtype, rng)
         self._init_finished = True
-        return self._params
+        self._params = params
       else:
-        return ()
+        params = ()
+      return (params, state)
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback(skip=3)
       raise LayerError(name, 'initialize', self._caller, input_shapes, trace)
 
-  def __call__(self, x, params=(), **kwargs):
+  def __call__(self, x, params=(), state=(), **kwargs):
     try:
       # If params are nothing, we may be reusing this layer.
       # Use the cached parameters to calculate the value.
@@ -231,34 +241,50 @@ def __call__(self, x, params=(), **kwargs):
       #   use "params is ()" instead of, e.g., "not params" or "params == ()".
       if params is ():  # pylint: disable=literal-comparison
         params = self._params
-      # In this case, we're called for the first time: cache parameters.
-      self._params = params
+      else:
+        # In this case, we're called for the first time: cache parameters.
+        self._params = params
 
       if not self.has_custom_grad:
-        return self.call(x, params=params, **kwargs)
+        return self.call(x, params=params, state=state, **kwargs)
 
       # Custom gradients part.
       assert backend.get_name() == 'jax', (
           'Custom gradients are only supported in JAX for now.')
 
+      # TODO(wangpeng): JAX doesn't support custom grads for functions with
+      #   auxiliary output yet (https://github.com/google/jax/issues/844). Will
+      #   remove the constraints on state below when this feature is added to
+      #   JAX.
+
+      assert state is (), (  # pylint: disable=literal-comparison
+          'Custom gradients do not allow non-trivial start state.')
+
+      def check_end_state(output_state):
+        output, state = output_state
+        assert state is (), (  # pylint: disable=literal-comparison
+            'Custom gradients do not allow non-trivial end state.')
+        return output
+
       # See this link for how custom transformations are defined in JAX:
       # https://jax.readthedocs.io/en/latest/jax.html#jax.custom_transforms
       # Note that we capture the kwargs and don't calculate gradients wrt. them.
       @jax.custom_transforms
       def do_call(y, params):
-        return self.call(y, params=params, **kwargs)
+        return check_end_state(self.call(y, params=params, state=(), **kwargs))
 
       # This is the custom gradient (vector-jacobian product in JAX) function.
       # For the exact specification of this custom transformation see this link:
       # https://jax.readthedocs.io/en/latest/jax.html#jax.defjvp_all
       def do_call_vjp(y, params):
-        output = self.call(y, params=params, **kwargs)
+        output = check_end_state(self.call(y, params=params, state=(),
+                                           **kwargs))
         def vjpfun(grad):
           return self.custom_grad(y, output, grad, params, **kwargs)
         return output, vjpfun
 
       jax.defvjp_all(do_call, do_call_vjp)
-      return do_call(x, params)
+      return do_call(x, params), ()
 
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback()
@@ -413,22 +439,23 @@ def _n_outputs(self):
 
     def _new_parameters(self, input_shapes, input_dtype, rng):
       if new_parameters is None:
-        return ()
+        return (), ()
       kwargs = self._init_kwargs  # pylint: disable=protected-access
-      return new_parameters(input_shapes, input_dtype, rng, **kwargs)
+      return new_parameters(input_shapes, input_dtype, rng, **kwargs), ()
 
     def _is_empty(raw_output):
       return raw_output is None or (isinstance(raw_output, (list, tuple))
                                     and len(raw_output) == 0)  # pylint: disable=g-explicit-length-test
 
-    def _call_with_context(self, x, params=(), **kwargs):
+    def _call_with_context(self, x, params=(), state=(), **kwargs):
       """Calls raw_call_fn with extra keyword args from Layer.__init__."""
       merged_kwargs = kwargs.copy()
       merged_kwargs.update(self._init_kwargs)  # pylint: disable=protected-access
 
       _validate_call_input(x, n_inputs)
       raw_output = raw_call_fn(x, params=params, **merged_kwargs)
-      return () if _is_empty(raw_output) else raw_output
+      output = () if _is_empty(raw_output) else raw_output
+      return (output, state)
 
     # Set docstrings and create the class.
     _call_with_context.__doc__ = raw_call_fn.__doc__
@@ -502,15 +529,15 @@ def check_shape_agreement(layer_fn, input_shapes, integer_inputs=False):
     input_dtype = tuple(input_dtype for _ in input_shapes)
   else:
     pseudo_data = ShapeType(input_shapes, input_dtype)
-  params = layer_fn.initialize(input_shapes, input_dtype, rng1)
-  pseudo_output = layer_fn.pseudo_call(pseudo_data, params)
+  params, state = layer_fn.initialize(input_shapes, input_dtype, rng1)
+  pseudo_output, _ = layer_fn.pseudo_call(pseudo_data, params, state)
   if isinstance(pseudo_output, tuple):
     output_shape = tuple(x.shape for x in pseudo_output)
   else:
     output_shape = pseudo_output.shape
 
   random_input = _random_values(input_shapes, rng2, integer_inputs)
-  real_output = layer_fn(random_input, params, rng=rng3)
+  real_output, _ = layer_fn(random_input, params, state=state, rng=rng3)
   result_shape = shapes(real_output)
 
   msg = 'output shape %s != real result shape %s' % (output_shape, result_shape)
diff --git a/tensor2tensor/trax/layers/base_test.py b/tensor2tensor/trax/layers/base_test.py
index a7dd593df..20b8020d6 100644
--- a/tensor2tensor/trax/layers/base_test.py
+++ b/tensor2tensor/trax/layers/base_test.py
@@ -40,11 +40,11 @@ class IdWithZeroGrad(base.Layer):
 
       def call(self, x, params, **kwargs):
         del params, kwargs
-        return x
+        return x, ()
 
       def new_parameters(self, input_shapes, input_dtype, rng):
         del input_shapes, input_dtype, rng
-        return ()
+        return (), ()
 
       @property
       def has_custom_grad(self):
@@ -59,7 +59,7 @@ def custom_grad(self, inputs, output, ct, params, **kwargs):
     input_shape = (9, 17)
     random_input = backend.random.uniform(rng, input_shape, minval=-1.0,
                                           maxval=1.0)
-    f = lambda x: backend.numpy.mean(layer(x, params, rng=rng))
+    f = lambda x: backend.numpy.mean(layer(x, params, rng=rng)[0])
     grad = backend.grad(f)(random_input)
     self.assertEqual(grad.shape, input_shape)  # Gradient for each input.
     self.assertEqual(sum(sum(grad * grad)), 0.0)  # Each one is 0.
@@ -70,11 +70,11 @@ class IdWithIdGrad(base.Layer):
 
       def call(self, x, params, **kwargs):
         del params, kwargs
-        return x
+        return x, ()
 
       def new_parameters(self, input_shapes, input_dtype, rng):
         del input_shapes, input_dtype, rng
-        return ()
+        return (), ()
 
       @property
       def has_custom_grad(self):
@@ -89,7 +89,7 @@ def custom_grad(self, inputs, output, ct, params, **kwargs):
     input_shape = (9, 17)
     random_input = backend.random.uniform(rng, input_shape, minval=-1.0,
                                           maxval=1.0)
-    f = lambda x: backend.numpy.mean(layer(x, params, rng=rng))
+    f = lambda x: backend.numpy.mean(layer(x, params, rng=rng)[0])
     grad = backend.grad(f)(random_input)
     self.assertEqual(grad.shape, input_shape)  # Gradient for each input.
     self.assertEqual(sum(sum(grad)), sum(sum(random_input)))  # Same as input.
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index c3da0e3cb..f51fb5ba8 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -168,14 +168,22 @@ def _validate_call_inputs(self, xs):
           'number of inputs ({}) to Serial.call less than n_inputs'
           ' ({})'.format(len(xs), self.n_inputs()))
 
-  def call(self, xs, params=(), **kwargs):
+  def call(self, xs, params=(), state=(), **kwargs):
     self._validate_call_inputs(xs)
     rngs = _pop_rng_and_split(kwargs, self._n_layers)
     if not self._sublayers:  # No-op: leave args unchanged.
-      return xs
+      return (xs, state)
 
     stack = xs
-    for layer, p, rng in zip(self._sublayers, params, rngs):
+    new_state = []
+    n_layers = self._n_layers
+    if n_layers != 1 and len(params) != n_layers:
+      raise ValueError('number of params ({}) not equal to number of layers '
+                       '({})'.format(len(params), n_layers))
+    if n_layers != 1 and len(state) != n_layers:
+      raise ValueError('number of params ({}) not equal to number of layers '
+                       '({})'.format(len(state), n_layers))
+    for layer, p, s, rng in zip(self._sublayers, params, state, rngs):
       is_stack_just_one_item = (_count_items(stack) == 1)
 
       # Give layer its args from the stack; treat 1-arg layer specially.
@@ -186,7 +194,8 @@ def call(self, xs, params=(), **kwargs):
         inputs = stack[0]
       else:
         inputs = stack[:n_in]
-      outputs = layer(inputs, p, rng=rng, **kwargs)
+      outputs, s = layer(inputs, p, state=s, rng=rng, **kwargs)
+      new_state.append(s)
 
       # Push outputs onto remaining stack (if any).
       if n_in < _count_items(stack):
@@ -196,7 +205,7 @@ def call(self, xs, params=(), **kwargs):
       else:
         stack = outputs  # NOTE: can be single value or tuple.
 
-    return stack
+    return stack, new_state
 
   def new_parameters(self, input_shape, input_dtype, rng):
     def MakeShapeType(shape, dtype):
@@ -205,6 +214,7 @@ def MakeShapeType(shape, dtype):
       return base.ShapeType(shape=shape, dtype=dtype)
 
     params = []
+    states = []
     pseudo_xs = MakeShapeType(input_shape, input_dtype)
     for layer in self._sublayers:
       rng, layer_rng = backend.random.split(rng)
@@ -221,10 +231,10 @@ def MakeShapeType(shape, dtype):
 
       in_shape = base.nested_map(inputs, lambda x: x.shape)
       in_dtype = base.nested_map(inputs, lambda x: x.dtype)
-      param = layer.initialize(in_shape, in_dtype, layer_rng)
+      param, state = layer.initialize(in_shape, in_dtype, layer_rng)
       pparam = layer._params   # pylint: disable=protected-access
 
-      outputs = layer.pseudo_call(inputs, pparam)
+      outputs, _ = layer.pseudo_call(inputs, pparam, state)
 
       # Push outputs onto remaining pseudo_xs (if any).
       if n_in < _count_items(pseudo_xs):
@@ -235,7 +245,8 @@ def MakeShapeType(shape, dtype):
         pseudo_xs = outputs  # NOTE: can be single value or tuple.
 
       params.append(param)
-    return params
+      states.append(state)
+    return params, states
 
 
 @base.layer(n_outputs=2)
@@ -330,11 +341,11 @@ def n_inputs(self):
     return self._n_items
 
   def new_parameters(self, input_shape, input_dtype, rng):
-    return ()
+    return (), ()
 
-  def call(self, xs, params=(), **kwargs):
+  def call(self, xs, params=(), state=(), **kwargs):
     del params, kwargs
-    return backend.numpy.concatenate(xs, self._axis)
+    return backend.numpy.concatenate(xs, self._axis), state
 
 
 class Parallel(base.Layer):
@@ -436,29 +447,37 @@ def _allot_to_sublayers(self, inputs):
       start = end
     return tuple(sub_inputs)
 
-  def call(self, inputs, params=(), **kwargs):
+  def call(self, inputs, params=(), state=(), **kwargs):
     n_layers, layers = self._n_layers, self._sublayers
     sublayer_inputs = self._allot_to_sublayers(inputs)
     rngs = _pop_rng_and_split(kwargs, n_layers)
     assert len(sublayer_inputs) == n_layers
     assert len(params) == n_layers
+    assert len(state) == n_layers
     assert len(rngs) == n_layers
     outputs = []
-    for layer, x, p, r in zip(layers, sublayer_inputs, params, rngs):
+    new_state = []
+    for layer, x, p, s, r in zip(layers, sublayer_inputs, params, state, rngs):
       # Note that zip silently truncates its result if lengths don't match.
-      sub_outputs = layer(x, params=p, rng=r, **kwargs)
+      sub_outputs, s = layer(x, params=p, state=s, rng=r, **kwargs)
       if layer.n_outputs() == 1:
         outputs.append(sub_outputs)
       else:
         outputs.extend(sub_outputs)
-    return outputs[0] if self.n_outputs() == 1 else tuple(outputs)
+      new_state.append(s)
+    output = outputs[0] if self.n_outputs() == 1 else tuple(outputs)
+    return output, new_state
 
   def new_parameters(self, input_shapes, input_dtypes, rng):
     sublayer_shapes = self._allot_to_sublayers(input_shapes)
     sublayer_dtypes = self._allot_to_sublayers(input_dtypes)
     rngs = backend.random.split(rng, self._n_layers)
-    return [layer.initialize(shape, dtype, rng) for layer, shape, dtype, rng
-            in zip(self._sublayers, sublayer_shapes, sublayer_dtypes, rngs)]
+    inits = [layer.initialize(shape, dtype, rng) for layer, shape, dtype, rng
+             in zip(self._sublayers, sublayer_shapes, sublayer_dtypes, rngs)]
+    if not inits:
+      return (), ()
+    else:
+      return tuple(zip(*inits))
 
 
 def Residual(*layers, **kwargs):
diff --git a/tensor2tensor/trax/layers/convolution.py b/tensor2tensor/trax/layers/convolution.py
index 9b556ef8d..aa9fd4efc 100644
--- a/tensor2tensor/trax/layers/convolution.py
+++ b/tensor2tensor/trax/layers/convolution.py
@@ -56,7 +56,7 @@ def _check_nhwc(self):
     msg = 'Convolutions on more than 4 dimensions only supported in NHWC.'
     assert self._lhs_spec == self._out_spec == 'NHWC', msg
 
-  def call(self, x, params=(), **kwargs):
+  def call(self, x, params=(), state=(), **kwargs):
     del kwargs
     w, b = params
     x_shape = list(x.shape)
@@ -69,7 +69,7 @@ def call(self, x, params=(), **kwargs):
         self._one) + b
     if len(x_shape) > 4:
       res = np.reshape(res, x_shape[:-3] + list(res.shape[-3:]))
-    return res
+    return res, state
 
   def _kernel_shape(self, input_shape):
     """Helper to calculate the kernel shape."""
@@ -89,7 +89,7 @@ def new_parameters(self, input_shape, input_dtype, rng):
     bias_shape = tuple(itertools.dropwhile(lambda x: x == 1, bias_shape))
     w = self._kernel_initializer(kernel_shape, rng)
     b = self._bias_initializer(bias_shape, rng)
-    return (w, b)
+    return (w, b), ()
 
 
 class CausalConv(Conv):
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 264423633..3e39a0f85 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -95,17 +95,17 @@ def __init__(self, n_units,
     self._kernel_initializer = kernel_initializer
     self._bias_initializer = bias_initializer
 
-  def call(self, x, params, **kwargs):
+  def call(self, x, params, state, **kwargs):
     del kwargs
     w, b = params
-    return np.dot(x, w) + b
+    return np.dot(x, w) + b, state
 
   def new_parameters(self, input_shape, input_dtype, rng):
     del input_dtype
     rng1, rng2 = backend.random.split(rng, 2)
     w = self._kernel_initializer((input_shape[-1], self._n_units), rng1)
     b = self._bias_initializer((self._n_units,), rng2)
-    return (w, b)
+    return (w, b), ()
 
 
 class Embedding(base.Layer):
@@ -118,14 +118,14 @@ def __init__(self, d_feature, vocab_size,
     self._vocab_size = vocab_size
     self._kernel_initializer = kernel_initializer
 
-  def call(self, x, params, **kwargs):
+  def call(self, x, params, state, **kwargs):
     del kwargs
-    return np.take(params, x, axis=0)
+    return np.take(params, x, axis=0), state
 
   def new_parameters(self, input_shape, input_dtype, rng):
     del input_dtype
     return self._kernel_initializer(
-        (self._vocab_size, self._d_feature), rng)
+        (self._vocab_size, self._d_feature), rng), ()
 
 
 # Flatten.
diff --git a/tensor2tensor/trax/layers/core_test.py b/tensor2tensor/trax/layers/core_test.py
index 4409652db..83356a5af 100644
--- a/tensor2tensor/trax/layers/core_test.py
+++ b/tensor2tensor/trax/layers/core_test.py
@@ -61,7 +61,7 @@ def test_flatten_n(self):
   def test_div(self):
     layer = core.Div(divisor=2.0)
     input_np = onp.array([[1, 2, 3], [4, 5, 6]], dtype=onp.float32)
-    output_np = layer(input_np)
+    output_np, _ = layer(input_np)
     # absltest doesn't have ndarray equalities.
     expected_output_np = input_np / 2.0
     self.assertAlmostEqual(
@@ -81,8 +81,8 @@ def test_dense_param_sharing(self):
     layer = core.Dense(32)
     model2 = combinators.Serial(layer, layer)
     rng = backend.random.get_prng(0)
-    params1 = model1.initialize((1, 32), onp.float32, rng)
-    params2 = model2.initialize((1, 32), onp.float32, rng)
+    params1, _ = model1.initialize((1, 32), onp.float32, rng)
+    params2, _ = model2.initialize((1, 32), onp.float32, rng)
     # The first parameters have 2 kernels of size (32, 32).
     self.assertEqual((32, 32), params1[0][0].shape)
     self.assertEqual((32, 32), params1[1][0].shape)
@@ -90,6 +90,16 @@ def test_dense_param_sharing(self):
     self.assertEqual((32, 32), params2[0][0].shape)
     self.assertEqual((), params2[1])
 
+  def test_dropout(self):
+    input_shape = (8, 7, 9)
+    output_shape = (8, 7, 9)
+    final_shape = base.check_shape_agreement(
+        core.Dropout(rate=0.1, mode="train"), input_shape)
+    self.assertEqual(final_shape, output_shape)
+    final_shape = base.check_shape_agreement(
+        core.Dropout(rate=0.1, mode="eval"), input_shape)
+    self.assertEqual(final_shape, output_shape)
+
 
 if __name__ == "__main__":
   absltest.main()
diff --git a/tensor2tensor/trax/layers/normalization.py b/tensor2tensor/trax/layers/normalization.py
index d6b07cc23..c5f0cca9a 100644
--- a/tensor2tensor/trax/layers/normalization.py
+++ b/tensor2tensor/trax/layers/normalization.py
@@ -23,56 +23,93 @@
 from tensor2tensor.trax.layers import base
 
 
-# Batch normalization.
-def BatchNormParams(input_shape, input_dtype, rng, axis=(0, 1, 2),
-                    center=True, scale=True, **kwargs):
-  """Helper to initialize batch norm params."""
-  del input_dtype, rng, kwargs
-  axis = (axis,) if np.isscalar(axis) else axis
-  shape = tuple(d for i, d in enumerate(input_shape) if i not in axis)
-  beta = np.zeros(shape, dtype='float32') if center else ()
-  gamma = np.ones(shape, dtype='float32') if scale else ()
-  return (beta, gamma)
-
-
-@base.layer(new_parameters=BatchNormParams)
-def BatchNorm(x, params, axis=(0, 1, 2), epsilon=1e-5,
-              center=True, scale=True, **unused_kwargs):
-  """Layer construction function for a batch normalization layer."""
-  mean = np.mean(x, axis, keepdims=True)
-  # Fast but less numerically-stable variance calculation than np.var.
-  m1 = np.mean(x**2, axis, keepdims=True)
-  var = m1 - mean**2
-  # x mustn't be onp.ndarray here; otherwise `x-mean` will call mean.__rsub__
-  # with each element of x, resulting in an onp.ndarray with dtype `object`.
-  z = (x - mean) / np.sqrt(var + epsilon).astype(x.dtype)
-
-  # Expand the parameters to have the right axes.
-  beta, gamma = params
-  # TODO(phawkins): np.expand_dims should accept an axis tuple.
-  # (https://github.com/numpy/numpy/issues/12290)
-  ed = tuple(None if i in axis else slice(None) for i in range(np.ndim(x)))
-  beta = beta[ed]
-  gamma = gamma[ed]
-
-  # Return the z rescaled by the parameters if requested.
-  if center and scale:
-    ret = gamma * z + beta
-  elif center:
-    ret = z + beta
-  elif scale:
-    ret = gamma * z
-  else:
-    ret = z
-  assert ret.dtype == x.dtype, ('The dtype of the output (%s) of batch norm is '
-                                'not the same as the input (%s). Batch norm '
-                                'should not change the dtype' %
-                                (ret.dtype, x.dtype))
-  return ret
+class BatchNorm(base.Layer):
+  """Batch normalization."""
+
+  def __init__(self, axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True,
+               momentum=None, mode='train'):
+    super(BatchNorm, self).__init__()
+    self._axis = axis
+    self._epsilon = epsilon
+    self._center = center
+    self._scale = scale
+    self._momentum = momentum
+    self._mode = mode
+
+  def new_parameters(self, input_shape, input_dtype, rng):
+    """Helper to initialize batch norm params."""
+    del input_dtype, rng
+    axis = self._axis
+    axis = (axis,) if np.isscalar(axis) else axis
+    shape = tuple(d for i, d in enumerate(input_shape) if i not in axis)
+    beta = np.zeros(shape, dtype='float32') if self._center else ()
+    gamma = np.ones(shape, dtype='float32') if self._scale else ()
+    def get_stats_axis(i, d):
+      if i in axis:
+        return 1
+      else:
+        return d
+    stats_shape = tuple(get_stats_axis(i, d) for i, d in enumerate(input_shape))
+    running_mean = np.zeros(stats_shape, dtype=np.float32)
+    running_var = np.zeros(stats_shape, dtype=np.float32)
+    num_batches = np.zeros((), dtype=np.int32)
+    return (beta, gamma), (running_mean, running_var, num_batches)
+
+  def call(self, x, params, state, **unused_kwargs):
+    """Layer construction function for a batch normalization layer."""
+
+    running_mean, running_var, num_batches = state
+
+    if self._mode == 'train':
+      mean = np.mean(x, self._axis, keepdims=True)
+      # Fast but less numerically-stable variance calculation than np.var.
+      m1 = np.mean(x**2, self._axis, keepdims=True)
+      var = m1 - mean**2
+      num_batches = num_batches + 1
+      if self._momentum is None:
+        # A simple average over all batches seen so far
+        exponential_average_factor = 1.0 / num_batches
+      else:
+        exponential_average_factor = self._momentum
+      def average(factor, new, old):
+        return (factor * new + (1 - factor) * old).astype(old.dtype)
+      running_mean = average(exponential_average_factor, mean, running_mean)
+      running_var = average(exponential_average_factor, var, running_var)
+      state = (running_mean, running_var, num_batches)
+    else:
+      mean = running_mean
+      var = running_var
+
+    z = (x - mean.astype(x.dtype)) / np.sqrt(var +
+                                             self._epsilon).astype(x.dtype)
+
+    # Expand the parameters to have the right axes.
+    beta, gamma = params
+    # TODO(phawkins): np.expand_dims should accept an axis tuple.
+    # (https://github.com/numpy/numpy/issues/12290)
+    ed = tuple(None if i in self._axis else slice(None)
+               for i in range(np.ndim(x)))
+    beta = beta[ed]
+    gamma = gamma[ed]
+
+    # Return the z rescaled by the parameters if requested.
+    if self._center and self._scale:
+      output = gamma * z + beta
+    elif self._center:
+      output = z + beta
+    elif self._scale:
+      output = gamma * z
+    else:
+      output = z
+    assert output.dtype == x.dtype, ('The dtype of the output (%s) of batch '
+                                     'norm is not the same as the input (%s). '
+                                     'Batch norm should not change the dtype' %
+                                     (output.dtype, x.dtype))
+    return output, state
 
 
 # Layer normalization.
-def LayerNormParams(input_shape, input_dtype, rng, epsilon=1e-6):
+def _layer_norm_params(input_shape, input_dtype, rng, epsilon=1e-6):
   """Helper: create layer norm parameters."""
   del input_dtype, rng, epsilon
   features = input_shape[-1]
@@ -81,8 +118,8 @@ def LayerNormParams(input_shape, input_dtype, rng, epsilon=1e-6):
   return (scale, bias)
 
 
-@base.layer(new_parameters=LayerNormParams)
-def LayerNorm(x, params, epsilon=1e-6, **unused_kwargs):
+@base.layer(new_parameters=_layer_norm_params)
+def LayerNorm(x, params, epsilon=1e-6, **unused_kwargs):  # pylint: disable=invalid-name
   (scale, bias) = params
   mean = np.mean(x, axis=-1, keepdims=True)
   variance = np.mean((x - mean)**2, axis=-1, keepdims=True)
diff --git a/tensor2tensor/trax/layers/normalization_test.py b/tensor2tensor/trax/layers/normalization_test.py
index 894d31693..d5a8067ee 100644
--- a/tensor2tensor/trax/layers/normalization_test.py
+++ b/tensor2tensor/trax/layers/normalization_test.py
@@ -19,19 +19,62 @@
 from __future__ import print_function
 
 from absl.testing import absltest
+import numpy as onp
+
+from tensor2tensor.trax import backend
+from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers import base
 from tensor2tensor.trax.layers import normalization
 
 
 class NormalizationLayerTest(absltest.TestCase):
 
-  def test_batch_norm(self):
+  def test_batch_norm_shape(self):
     input_shape = (29, 5, 7, 20)
     result_shape = base.check_shape_agreement(
         normalization.BatchNorm(), input_shape)
     self.assertEqual(result_shape, input_shape)
 
-  def test_layer_norm(self):
+  def test_batch_norm(self):
+    input_shape = (2, 3, 4)
+    input_dtype = np.float32
+    eps = 1e-5
+    rng = backend.random.get_prng(0)
+    inp1 = np.reshape(np.arange(np.prod(input_shape), dtype=input_dtype),
+                      input_shape)
+    m1 = 11.5
+    v1 = 47.9167
+    layer = normalization.BatchNorm(axis=(0, 1, 2))
+    params, state = layer.initialize(input_shape, input_dtype, rng)
+    onp.testing.assert_allclose(state[0], 0)
+    onp.testing.assert_allclose(state[1], 0)
+    self.assertEqual(state[2], 0)
+    out, state = layer(inp1, params, state)
+    onp.testing.assert_allclose(state[0], m1)
+    onp.testing.assert_allclose(state[1], v1, rtol=1e-6)
+    self.assertEqual(state[2], 1)
+    onp.testing.assert_allclose(out, (inp1 - m1) / np.sqrt(v1 + eps),
+                                rtol=1e-6)
+    inp2 = inp1 * 2 + 3
+    m2 = m1 * 2 + 3
+    v2 = v1 * 4
+    m12 = (m1 + m2) / 2
+    v12 = (v1 + v2) / 2
+    out, state = layer(inp2, params, state)
+    onp.testing.assert_allclose(state[0], m12)
+    onp.testing.assert_allclose(state[1], v12, rtol=1e-6)
+    self.assertEqual(state[2], 2)
+    onp.testing.assert_allclose(out, (inp2 - m2) / np.sqrt(v2 + eps),
+                                rtol=1e-6)
+    layer = normalization.BatchNorm(axis=(0, 1, 2), mode="eval")
+    inp3 = inp1 * 5 + 7
+    out, state_unchanged = layer(inp3, params, state)
+    for i in range(3):
+      onp.testing.assert_allclose(state_unchanged[i], state[i])
+    onp.testing.assert_allclose(out, (inp3 - m12) / np.sqrt(v12 + eps),
+                                rtol=1e-6)
+
+  def test_layer_norm_shape(self):
     input_shape = (29, 5, 7, 20)
     result_shape = base.check_shape_agreement(
         normalization.LayerNorm(), input_shape)
diff --git a/tensor2tensor/trax/layers/reversible.py b/tensor2tensor/trax/layers/reversible.py
index 63dd3183e..f10e45a50 100644
--- a/tensor2tensor/trax/layers/reversible.py
+++ b/tensor2tensor/trax/layers/reversible.py
@@ -53,7 +53,7 @@ def reverse_and_grad(self, output, grad, params=(), **kwargs):
     """
     # Note: jax.vjp does not allow us to use **kwargs in the signature here.
     def _do_call(x, params):
-      return super(ReversibleLayer, self).call(x, params=params, **kwargs)
+      return super(ReversibleLayer, self).call(x, params=params, **kwargs)[0]
 
     reconstructed_x = self.reverse(output, params, **kwargs)
     _, vjpfun = jax.vjp(_do_call, reconstructed_x, params)
diff --git a/tensor2tensor/trax/models/atari_cnn_test.py b/tensor2tensor/trax/models/atari_cnn_test.py
index abb2a200d..007ed6557 100644
--- a/tensor2tensor/trax/models/atari_cnn_test.py
+++ b/tensor2tensor/trax/models/atari_cnn_test.py
@@ -37,11 +37,11 @@ def test_computes(self):
         hidden_sizes=hidden_size, output_size=output_size)
     B, T, OBS = 2, 2, (28, 28, 3)  # pylint: disable=invalid-name
     rng_key, key = jax_random.split(rng_key)
-    params = model.initialize((1, 1) + OBS, onp.float32, key)
+    params, state = model.initialize((1, 1) + OBS, onp.float32, key)
     x = onp.arange(B * (T + 1) * functools.reduce(op.mul, OBS)).reshape(
         B, T + 1, *OBS)
     rng_key, key = jax_random.split(rng_key)
-    y = model(x, params, rng=key)
+    y, _ = model(x, params, state=state, rng=key)
     self.assertEqual((B, T + 1, output_size), y.shape)
 
 
@@ -55,11 +55,11 @@ def test_computes(self):
         hidden_sizes=hidden_size, output_size=output_size)
     B, T, OBS = 2, 2, 3  # pylint: disable=invalid-name
     rng_key, key = jax_random.split(rng_key)
-    params = model.initialize((1, 1, OBS), onp.float32, key)
+    params, state = model.initialize((1, 1, OBS), onp.float32, key)
     x = onp.arange(B * (T + 1) * OBS).reshape(
         B, T + 1, OBS)
     rng_key, key = jax_random.split(rng_key)
-    y = model(x, params, rng=key)
+    y, _ = model(x, params, state=state, rng=key)
     self.assertEqual((B, T + 1, output_size), y.shape)
 
 
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index beeaa5780..3551e10a5 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -22,24 +22,24 @@
 from tensor2tensor.trax import layers as tl
 
 
-def ConvBlock(kernel_size, filters, strides):
+def ConvBlock(kernel_size, filters, strides, mode='train'):
   """ResNet convolutional striding block."""
   # TODO(jonni): Use good defaults so Resnet50 code is cleaner / less redundant.
   ks = kernel_size
   filters1, filters2, filters3 = filters
   main = [
       tl.Conv(filters1, (1, 1), strides),
-      tl.BatchNorm(),
+      tl.BatchNorm(mode=mode),
       tl.Relu(),
       tl.Conv(filters2, (ks, ks), padding='SAME'),
-      tl.BatchNorm(),
+      tl.BatchNorm(mode=mode),
       tl.Relu(),
       tl.Conv(filters3, (1, 1)),
-      tl.BatchNorm(),
+      tl.BatchNorm(mode=mode),
   ]
   shortcut = [
       tl.Conv(filters3, (1, 1), strides),
-      tl.BatchNorm(),
+      tl.BatchNorm(mode=mode),
   ]
   return [
       tl.Residual(main, shortcut=shortcut),
@@ -47,20 +47,20 @@ def ConvBlock(kernel_size, filters, strides):
   ]
 
 
-def IdentityBlock(kernel_size, filters):
+def IdentityBlock(kernel_size, filters, mode='train'):
   """ResNet identical size block."""
   # TODO(jonni): Use good defaults so Resnet50 code is cleaner / less redundant.
   ks = kernel_size
   filters1, filters2, filters3 = filters
   main = [
       tl.Conv(filters1, (1, 1)),
-      tl.BatchNorm(),
+      tl.BatchNorm(mode=mode),
       tl.Relu(),
       tl.Conv(filters2, (ks, ks), padding='SAME'),
-      tl.BatchNorm(),
+      tl.BatchNorm(mode=mode),
       tl.Relu(),
       tl.Conv(filters3, (1, 1)),
-      tl.BatchNorm(),
+      tl.BatchNorm(mode=mode),
   ]
   return [
       tl.Residual(main),
@@ -79,29 +79,31 @@ def Resnet50(d_hidden=64, n_output_classes=1001, mode='train'):
   Returns:
     The list of layers comprising a ResNet model with the given parameters.
   """
-  del mode
   return tl.Model(
       tl.ToFloat(),
       tl.Conv(d_hidden, (7, 7), (2, 2), 'SAME'),
-      tl.BatchNorm(),
+      tl.BatchNorm(mode=mode),
       tl.Relu(),
       tl.MaxPool(pool_size=(3, 3), strides=(2, 2)),
-      ConvBlock(3, [d_hidden, d_hidden, 4 * d_hidden], (1, 1)),
-      IdentityBlock(3, [d_hidden, d_hidden, 4 * d_hidden]),
-      IdentityBlock(3, [d_hidden, d_hidden, 4 * d_hidden]),
-      ConvBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden], (2, 2)),
-      IdentityBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden]),
-      IdentityBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden]),
-      IdentityBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden]),
-      ConvBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden], (2, 2)),
-      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden]),
-      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden]),
-      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden]),
-      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden]),
-      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden]),
-      ConvBlock(3, [8 * d_hidden, 8 * d_hidden, 32 * d_hidden], (2, 2)),
-      IdentityBlock(3, [8 * d_hidden, 8 * d_hidden, 32 * d_hidden]),
-      IdentityBlock(3, [8 * d_hidden, 8 * d_hidden, 32 * d_hidden]),
+      ConvBlock(3, [d_hidden, d_hidden, 4 * d_hidden], (1, 1), mode=mode),
+      IdentityBlock(3, [d_hidden, d_hidden, 4 * d_hidden], mode=mode),
+      IdentityBlock(3, [d_hidden, d_hidden, 4 * d_hidden], mode=mode),
+      ConvBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden], (2, 2),
+                mode=mode),
+      IdentityBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden], mode=mode),
+      IdentityBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden], mode=mode),
+      IdentityBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden], mode=mode),
+      ConvBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden], (2, 2),
+                mode=mode),
+      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden], mode=mode),
+      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden], mode=mode),
+      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden], mode=mode),
+      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden], mode=mode),
+      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden], mode=mode),
+      ConvBlock(3, [8 * d_hidden, 8 * d_hidden, 32 * d_hidden], (2, 2),
+                mode=mode),
+      IdentityBlock(3, [8 * d_hidden, 8 * d_hidden, 32 * d_hidden], mode=mode),
+      IdentityBlock(3, [8 * d_hidden, 8 * d_hidden, 32 * d_hidden], mode=mode),
       tl.AvgPool(pool_size=(7, 7)),
       tl.Flatten(),
       tl.Dense(n_output_classes),
@@ -109,25 +111,26 @@ def Resnet50(d_hidden=64, n_output_classes=1001, mode='train'):
   )
 
 
-def WideResnetBlock(channels, strides=(1, 1)):
+def WideResnetBlock(channels, strides=(1, 1), mode='train'):
   """WideResnet convolutional block."""
   return [
-      tl.BatchNorm(),
+      tl.BatchNorm(mode=mode),
       tl.Relu(),
       tl.Conv(channels, (3, 3), strides, padding='SAME'),
-      tl.BatchNorm(),
+      tl.BatchNorm(mode=mode),
       tl.Relu(),
       tl.Conv(channels, (3, 3), padding='SAME'),
   ]
 
 
-def WideResnetGroup(n, channels, strides=(1, 1)):
+def WideResnetGroup(n, channels, strides=(1, 1), mode='train'):
   shortcut = [
       tl.Conv(channels, (3, 3), strides, padding='SAME'),
   ]
   return [
-      tl.Residual(WideResnetBlock(channels, strides), shortcut=shortcut),
-      tl.Residual([WideResnetBlock(channels, (1, 1))
+      tl.Residual(WideResnetBlock(channels, strides, mode=mode),
+                  shortcut=shortcut),
+      tl.Residual([WideResnetBlock(channels, (1, 1), mode=mode)
                    for _ in range(n - 1)]),
   ]
 
@@ -144,14 +147,13 @@ def WideResnet(n_blocks=3, widen_factor=1, n_output_classes=10, mode='train'):
   Returns:
     The list of layers comprising a WideResnet model with the given parameters.
   """
-  del mode
   return tl.Model(
       tl.ToFloat(),
       tl.Conv(16, (3, 3), padding='SAME'),
-      WideResnetGroup(n_blocks, 16 * widen_factor),
-      WideResnetGroup(n_blocks, 32 * widen_factor, (2, 2)),
-      WideResnetGroup(n_blocks, 64 * widen_factor, (2, 2)),
-      tl.BatchNorm(),
+      WideResnetGroup(n_blocks, 16 * widen_factor, mode=mode),
+      WideResnetGroup(n_blocks, 32 * widen_factor, (2, 2), mode=mode),
+      WideResnetGroup(n_blocks, 64 * widen_factor, (2, 2), mode=mode),
+      tl.BatchNorm(mode=mode),
       tl.Relu(),
       tl.AvgPool(pool_size=(8, 8)),
       tl.Flatten(),
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index 15e0891ad..cfe8edee0 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -100,8 +100,9 @@ def policy_and_value_net(rng_key,
         )
     ]
   net = tl.Model(layers)
-  params = net.initialize(batch_observations_shape, observations_dtype, rng_key)
-  return params, net
+  params, state = net.initialize(batch_observations_shape, observations_dtype,
+                                 rng_key)
+  return params, state, net
 
 
 def optimizer_fn(optimizer, net_params):
@@ -146,6 +147,7 @@ def collect_trajectories(env,
                          reset=True,
                          len_history_for_policy=32,
                          boundary=32,
+                         state=None,
                          rng=None):
   """Collect trajectories with the given policy net and behaviour.
 
@@ -161,6 +163,7 @@ def collect_trajectories(env,
     len_history_for_policy: int or None, the maximum history to keep for
       applying the policy on. If None, use the full history.
     boundary: int, pad the sequences to the multiples of this number.
+    state: state for `policy_fn`.
     rng: jax rng, splittable.
 
   Returns:
@@ -174,7 +177,7 @@ def collect_trajectories(env,
 
   assert isinstance(env, env_problem.EnvProblem)
   # This is an env_problem, run its collect function.
-  trajs, n_done, timing_info = env_problem_utils.play_env_problem_with_policy(
+  trajs, n_done, timing_info, state = env_problem_utils.play_env_problem_with_policy(
       env,
       policy_fn,
       num_trajectories=n_trajectories,
@@ -182,12 +185,13 @@ def collect_trajectories(env,
       reset=reset,
       len_history_for_policy=len_history_for_policy,
       boundary=boundary,
+      state=state,
       rng=rng)
   # Skip returning raw_rewards here, since they aren't used.
 
   # t is the return value of Trajectory.as_numpy, so:
   # (observation, action, processed_reward, raw_reward, infos)
-  return [(t[0], t[1], t[2], t[4]) for t in trajs], n_done, timing_info
+  return [(t[0], t[1], t[2], t[4]) for t in trajs], n_done, timing_info, state
 
 
 # This function can probably be simplified, ask how?
@@ -625,10 +629,12 @@ def combined_loss(new_params,
                   epsilon=0.2,
                   c1=1.0,
                   c2=0.01,
+                  state=None,
                   rng=None):
   """Computes the combined (clipped loss + value loss) given observations."""
-  log_probab_actions_new, value_predictions_new = policy_and_value_net_apply(
-      padded_observations, new_params, rng=rng)
+  (log_probab_actions_new, value_predictions_new), state = (
+      policy_and_value_net_apply(padded_observations, new_params, state,
+                                 rng=rng))
 
   # (combined_loss, ppo_loss, value_loss, entropy_bonus)
   return combined_loss_given_predictions(
@@ -643,7 +649,7 @@ def combined_loss(new_params,
       lambda_=lambda_,
       epsilon=epsilon,
       c1=c1,
-      c2=c2)
+      c2=c2), state
 
 
 @functools.partial(jit, static_argnums=(2, 3, 4))
@@ -663,13 +669,14 @@ def policy_and_value_opt_step(i,
                               gamma=0.99,
                               lambda_=0.95,
                               epsilon=0.1,
+                              state=None,
                               rng=None):
   """Policy and Value optimizer step."""
 
   # Combined loss function given the new params.
-  def policy_and_value_loss(params):
+  def policy_and_value_loss(params, state):
     """Returns the combined loss given just parameters."""
-    (loss, _, _, _) = combined_loss(
+    (loss, _, _, _), state = combined_loss(
         params,
         log_probab_actions_old,
         value_predictions_old,
@@ -683,13 +690,14 @@ def policy_and_value_loss(params):
         gamma=gamma,
         lambda_=lambda_,
         epsilon=epsilon,
+        state=state,
         rng=rng)
-    return loss
+    return loss, state
 
   new_params = get_params(opt_state)
-  g = grad(policy_and_value_loss)(new_params)
+  g, state = grad(policy_and_value_loss, has_aux=True)(new_params, state)
   # TODO(afrozm): Maybe clip gradients?
-  return opt_update(i, g, opt_state)
+  return opt_update(i, g, opt_state), state
 
 
 def get_time(t1, t2=None):
@@ -743,6 +751,7 @@ def evaluate_policy(eval_env,
                     max_timestep=20000,
                     n_evals=1,
                     len_history_for_policy=32,
+                    state=None,
                     rng=None):
   """Evaluate the policy."""
 
@@ -750,13 +759,14 @@ def evaluate_policy(eval_env,
   raw_reward_sums = collections.defaultdict(list)
   for eval_rng in jax_random.split(rng, num=n_evals):
     for temperature in temperatures:
-      trajs, _, _ = env_problem_utils.play_env_problem_with_policy(
+      trajs, _, _, state = env_problem_utils.play_env_problem_with_policy(
           eval_env,
           get_predictions,
           num_trajectories=eval_env.batch_size,
           max_timestep=max_timestep,
           reset=True,
           temperature=temperature,
+          state=state,
           rng=eval_rng,
           len_history_for_policy=len_history_for_policy)
       processed_reward_sums[temperature].extend(sum(traj[2]) for traj in trajs)
@@ -771,15 +781,16 @@ def compute_stats(reward_dict):
   return {
       "processed": compute_stats(processed_reward_sums),
       "raw": compute_stats(raw_reward_sums),
-  }
+  }, state
 
 
-def maybe_restore_params(output_dir, policy_and_value_net_params):
+def maybe_restore_params(output_dir, policy_and_value_net_params, state):
   """Maybe restore the params from the checkpoint dir.
 
   Args:
     output_dir: Directory where saved model checkpoints are stored.
     policy_and_value_net_params: Default params, returned if model is'nt found.
+    state: policy state.
 
   Returns:
     triple (restore (bool), params, iter(int)) where iter is the epoch from
@@ -790,16 +801,17 @@ def maybe_restore_params(output_dir, policy_and_value_net_params):
     logging.info("Trying to restore model from %s", model_file)
     try:
       with gfile.GFile(model_file, "rb") as f:
-        loaded_policy_and_value_net_params = pickle.load(f)
+        loaded_policy_and_value_net_params, loaded_state = pickle.load(f)
         policy_and_value_net_params = loaded_policy_and_value_net_params
+        state = loaded_state
       model_file_basename = os.path.basename(model_file)  # model-??????.pkl
       i = int(filter(str.isdigit, model_file_basename))
-      return True, policy_and_value_net_params, i
+      return True, policy_and_value_net_params, state, i
     except EOFError as e:
       logging.error("Unable to load model from: %s with %s", model_file, e)
       # Try an older version.
       continue
-  return False, policy_and_value_net_params, 0
+  return False, policy_and_value_net_params, state, 0
 
 
 def write_eval_reward_summaries(reward_stats_by_mode, summary_writer, epoch):
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
index b91bf02c9..c7ddc314e 100644
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -38,14 +38,14 @@ def test_policy_and_value_net(self):
     observation_shape = (3, 4, 5)
     batch_observation_shape = (1, 1) + observation_shape
     n_actions = 2
-    pnv_params, pnv_apply = ppo.policy_and_value_net(
+    pnv_params, pnv_state, pnv_apply = ppo.policy_and_value_net(
         self.rng_key, batch_observation_shape, np.float32, n_actions,
         lambda: [layers.Flatten(n_axes_to_keep=2)])
     batch = 2
     time_steps = 10
     batch_of_observations = np.random.uniform(
         size=(batch, time_steps) + observation_shape)
-    pnv_output = pnv_apply(batch_of_observations, pnv_params)
+    pnv_output, _ = pnv_apply(batch_of_observations, pnv_params, pnv_state)
 
     # Output is a list, first is probab of actions and the next is value output.
     self.assertEqual(2, len(pnv_output))
@@ -392,11 +392,11 @@ def test_combined_loss(self):
     B, T, A, OBS = 2, 10, 2, (28, 28, 3)  # pylint: disable=invalid-name
     batch_observation_shape = (1, 1) + OBS
 
-    old_params, _ = ppo.policy_and_value_net(
+    old_params, _, _ = ppo.policy_and_value_net(
         key1, batch_observation_shape, np.float32, A,
         lambda: [layers.Flatten(n_axes_to_keep=2)])
 
-    new_params, net_apply = ppo.policy_and_value_net(
+    new_params, state, net_apply = ppo.policy_and_value_net(
         key2, batch_observation_shape, np.float32, A,
         lambda: [layers.Flatten(n_axes_to_keep=2)])
 
@@ -408,8 +408,10 @@ def test_combined_loss(self):
     mask = np.ones_like(rewards)
 
     # Just test that this computes at all.
-    new_log_probabs, value_predictions_new = net_apply(observations, new_params)
-    old_log_probabs, value_predictions_old = net_apply(observations, old_params)
+    (new_log_probabs, value_predictions_new), _ = net_apply(observations,
+                                                            new_params, state)
+    (old_log_probabs, value_predictions_old), _ = net_apply(observations,
+                                                            old_params, state)
 
     gamma = 0.99
     lambda_ = 0.95
@@ -431,7 +433,7 @@ def test_combined_loss(self):
         lambda_=lambda_,
         epsilon=epsilon)
 
-    (combined_loss, ppo_loss_2, value_loss_2, entropy_bonus) = (
+    (combined_loss, ppo_loss_2, value_loss_2, entropy_bonus), state = (
         ppo.combined_loss(new_params,
                           old_log_probabs,
                           value_predictions_old,
@@ -444,7 +446,8 @@ def test_combined_loss(self):
                           lambda_=lambda_,
                           epsilon=epsilon,
                           c1=c1,
-                          c2=c2)
+                          c2=c2,
+                          state=state)
     )
 
     # Test that these compute at all and are self consistent.
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index bbd9b2baa..b5479d4ed 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -141,7 +141,7 @@ def __init__(
     self._rng, key1 = jax_random.split(self._rng, num=2)
 
     # Initialize the policy and value network.
-    policy_and_value_net_params, policy_and_value_net_apply = (
+    policy_and_value_net_params, self._model_state, policy_and_value_net_apply = (
         ppo.policy_and_value_net(
             rng_key=key1,
             batch_observations_shape=batch_observations_shape,
@@ -155,8 +155,9 @@ def __init__(
 
     # Maybe restore the policy params. If there is nothing to restore, then
     # iteration = 0 and policy_and_value_net_params are returned as is.
-    restored, policy_and_value_net_params, self._epoch = (
-        ppo.maybe_restore_params(output_dir, policy_and_value_net_params))
+    restored, policy_and_value_net_params, self._model_state, self._epoch = (
+        ppo.maybe_restore_params(output_dir, policy_and_value_net_params,
+                                 self._model_state))
 
     if restored:
       logging.info("Restored parameters from iteration [%d]", self._epoch)
@@ -200,11 +201,12 @@ def train_epoch(self):
     trajectory_collection_start_time = time.time()
     logging.vlog(1, "Epoch [% 6d] collecting trajectories.", self._epoch)
     self._rng, key = jax_random.split(self._rng)
-    trajs, n_done, timing_info = ppo.collect_trajectories(
+    trajs, n_done, timing_info, self._model_state = ppo.collect_trajectories(
         self._train_env,
         policy_fn=self._get_predictions,
         n_trajectories=self._train_env.batch_size,
         max_timestep=self._max_timestep,
+        state=self._model_state,
         rng=key,
         len_history_for_policy=self._len_history_for_policy,
         boundary=self._boundary,
@@ -274,8 +276,9 @@ def train_epoch(self):
     # observation, so we re-calculate for everything, but use the original ones
     # for all but the last time-step.
     self._rng, key = jax_random.split(self._rng)
-    log_probabs_traj, value_predictions_traj, _ = self._get_predictions(
-        padded_observations, rng=key)
+
+    log_probabs_traj, value_predictions_traj, self._model_state, _ = (
+        self._get_predictions(padded_observations, self._model_state, rng=key))
 
     assert (B, T + 1, A) == log_probabs_traj.shape
     assert (B, T + 1, 1) == value_predictions_traj.shape
@@ -294,21 +297,23 @@ def train_epoch(self):
     self._rng, key1 = jax_random.split(self._rng, num=2)
     logging.vlog(2, "Starting to compute P&V loss.")
     loss_compute_start_time = time.time()
-    cur_combined_loss, cur_ppo_loss, cur_value_loss, entropy_bonus = (
-        ppo.combined_loss(
-            self._policy_and_value_net_params,
-            log_probabs_traj,
-            value_predictions_traj,
-            self._policy_and_value_net_apply,
-            padded_observations,
-            padded_actions,
-            padded_rewards,
-            reward_mask,
-            gamma=self._gamma,
-            lambda_=self._lambda_,
-            c1=self._c1,
-            c2=self._c2,
-            rng=key1))
+    (cur_combined_loss, cur_ppo_loss, cur_value_loss,
+     entropy_bonus), self._model_state = (
+         ppo.combined_loss(
+             self._policy_and_value_net_params,
+             log_probabs_traj,
+             value_predictions_traj,
+             self._policy_and_value_net_apply,
+             padded_observations,
+             padded_actions,
+             padded_rewards,
+             reward_mask,
+             gamma=self._gamma,
+             lambda_=self._lambda_,
+             c1=self._c1,
+             c2=self._c2,
+             state=self._model_state,
+             rng=key1))
     loss_compute_time = ppo.get_time(loss_compute_start_time)
     logging.vlog(
         1,
@@ -324,7 +329,7 @@ def train_epoch(self):
       k1, k2, k3 = jax_random.split(key, num=3)
       t = time.time()
       # Update the optimizer state.
-      self._policy_and_value_opt_state = ppo.policy_and_value_opt_step(
+      self._policy_and_value_opt_state, self._model_state = ppo.policy_and_value_opt_step(
           j,
           self._policy_and_value_opt_state,
           self._policy_and_value_opt_update,
@@ -340,11 +345,14 @@ def train_epoch(self):
           c2=self._c2,
           gamma=self._gamma,
           lambda_=self._lambda_,
+          state=self._model_state,
           rng=k1)
 
       # Compute the approx KL for early stopping.
-      log_probab_actions_new, _ = self._policy_and_value_net_apply(
-          padded_observations, self._policy_and_value_net_params, rng=k2)
+      (log_probab_actions_new, _), self._model_state = (
+          self._policy_and_value_net_apply(padded_observations,
+                                           self._policy_and_value_net_params,
+                                           self._model_state, rng=k2))
 
       approx_kl = ppo.approximate_kl(log_probab_actions_new, log_probabs_traj,
                                      reward_mask)
@@ -361,21 +369,23 @@ def train_epoch(self):
       if (((j + 1) % self._print_every_optimizer_steps == 0) or
           (j == self._n_optimizer_steps - 1) or early_stopping):
         # Compute and log the loss.
-        (loss_combined, loss_ppo, loss_value, entropy_bonus) = (
-            ppo.combined_loss(
-                self._policy_and_value_net_params,
-                log_probabs_traj,
-                value_predictions_traj,
-                self._policy_and_value_net_apply,
-                padded_observations,
-                padded_actions,
-                padded_rewards,
-                reward_mask,
-                gamma=self._gamma,
-                lambda_=self._lambda_,
-                c1=self._c1,
-                c2=self._c2,
-                rng=k3))
+        (loss_combined, loss_ppo, loss_value,
+         entropy_bonus), self._model_state = (
+             ppo.combined_loss(
+                 self._policy_and_value_net_params,
+                 log_probabs_traj,
+                 value_predictions_traj,
+                 self._policy_and_value_net_apply,
+                 padded_observations,
+                 padded_actions,
+                 padded_rewards,
+                 reward_mask,
+                 gamma=self._gamma,
+                 lambda_=self._lambda_,
+                 c1=self._c1,
+                 c2=self._c2,
+                 state=self._model_state,
+                 rng=k3))
         logging.vlog(1, "One Policy and Value grad desc took: %0.2f msec",
                      ppo.get_time(t, t2))
         logging.vlog(
@@ -449,13 +459,14 @@ def evaluate(self):
     """Evaluate the agent."""
     logging.vlog(1, "Epoch [% 6d] evaluating policy.", self._epoch)
     self._rng, key = jax_random.split(self._rng, num=2)
-    reward_stats = ppo.evaluate_policy(
+    reward_stats, self._model_state = ppo.evaluate_policy(
         self._eval_env,
         self._get_predictions,
         temperatures=self._eval_temperatures,
         max_timestep=self._max_timestep_eval,
         n_evals=self._n_evals,
         len_history_for_policy=self._len_history_for_policy,
+        state=self._model_state,
         rng=key)
     ppo.write_eval_reward_summaries(
         reward_stats, self._eval_sw, epoch=self._epoch)
@@ -467,7 +478,7 @@ def save(self):
         os.path.join(self._output_dir, "model-??????.pkl"))
     params_file = os.path.join(self._output_dir, "model-%06d.pkl" % self._epoch)
     with gfile.GFile(params_file, "wb") as f:
-      pickle.dump(self._policy_and_value_net_params, f)
+      pickle.dump((self._policy_and_value_net_params, self._model_state), f)
     # Remove the old model files.
     for path in old_model_files:
       gfile.remove(path)
@@ -485,11 +496,11 @@ def _policy_and_value_net_params(self):
     return self._policy_and_value_get_params(self._policy_and_value_opt_state)
 
   # A function to get the policy and value predictions.
-  def _get_predictions(self, observations, rng=None):
+  def _get_predictions(self, observations, state, rng=None):
     """Returns log-probs, value predictions and key back."""
     key, key1 = jax_random.split(rng, num=2)
 
-    log_probs, value_preds = self._policy_and_value_net_apply(
-        observations, self._policy_and_value_net_params, rng=key1)
+    (log_probs, value_preds), state = self._policy_and_value_net_apply(
+        observations, self._policy_and_value_net_params, state, rng=key1)
 
-    return log_probs, value_preds, key
+    return log_probs, value_preds, state, key
diff --git a/tensor2tensor/trax/rl/ppo_trainer_test.py b/tensor2tensor/trax/rl/ppo_trainer_test.py
index 1d8e2c470..a8b48ec47 100644
--- a/tensor2tensor/trax/rl/ppo_trainer_test.py
+++ b/tensor2tensor/trax/rl/ppo_trainer_test.py
@@ -171,10 +171,9 @@ def inputs(n_devices):
           input_dtype=(np.float32, np.int32),
       )
 
-    def loss(*args, **kwargs):
-      del args
-      del kwargs
-      return 0.0
+    def loss(params, batch, model_predict, state, rng, **kwargs):
+      del params, batch, model_predict, rng, kwargs
+      return 0.0, state
 
     with self.tmp_dir() as output_dir:
       # Run fake training just to save the parameters.
diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index f2dbccbc0..5de39f353 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -65,6 +65,7 @@ def __init__(self, model, batch_size, observation_space, action_space,
 
     self._predict_fn = None
     self._rng = None
+    self._model_state = None
     self._history_stream = None
 
     # Call the super's ctor. It will use some of the member fields, so we call
@@ -93,6 +94,7 @@ def initialize_environments(self,
 
     model_state = trax.restore_state(self._output_dir)
     model_params = model_state.opt_state.params
+    self._model_state = model_state.model_state
     self._predict_fn = functools.partial(
         self._model_predict,
         params=model_params,
@@ -240,7 +242,8 @@ def _reset_model(self, predict_fn, indices, history, rng):
     return history[:, -1, ...]
 
   def _step_model(self, predict_fn, actions, rng):
-    (observation, reward) = predict_fn((self._history, actions), rng=rng)
+    (observation, reward), self._model_state = predict_fn(
+        (self._history, actions), state=self._model_state, rng=rng)
 
     # Roll the history one timestep back and append the new observation.
     self._history = np.roll(self._history, shift=-1, axis=1)
@@ -355,7 +358,10 @@ def gumbel_sample(log_probs):
 
     for (i, subrng) in enumerate(jax_random.split(rng, self._obs_repr_length)):
       symbol_index = self._steps * self._step_repr_length + i
-      log_probs = predict_fn(self._history, rng=subrng)[:, symbol_index, :]
+      log_probs, self._model_state = predict_fn(self._history,
+                                                state=self._model_state,
+                                                rng=subrng)
+      log_probs = log_probs[:, symbol_index, :]
       self._history[:, symbol_index] = gumbel_sample(log_probs)
 
     obs_repr = self._history[self._obs_repr_indices]
diff --git a/tensor2tensor/trax/rl/simulated_env_problem_test.py b/tensor2tensor/trax/rl/simulated_env_problem_test.py
index 72149ec1e..6f1dcd605 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem_test.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem_test.py
@@ -63,7 +63,7 @@ def mock_transition(inputs, *args, **kwargs):
       (observations, actions) = inputs
       new_observations = observations[:, -1] + actions
       rewards = np.array([[int(new_observations % 2 == 0)]])
-      return (new_observations, rewards)
+      return (new_observations, rewards), ()
 
     mock_model_fn = mock.MagicMock()
     mock_model_fn.return_value.side_effect = mock_transition
@@ -136,7 +136,7 @@ def make_prediction(symbol):
       one_hot = np.eye(vocab_size)[symbol]
       log_probs = (1 - one_hot) * -100.0  # Virtually deterministic.
       # (4 obs symbols + 1 action symbol) * 3 timesteps = 15.
-      return np.array([[log_probs] * 15])
+      return np.array([[log_probs] * 15]), ()
 
     mock_model_fn = mock.MagicMock()
     mock_model = mock_model_fn.return_value
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 4a9f258a1..9ecdd62e0 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -125,15 +125,15 @@ def neg_log_perplexity(batch, model_predictions):
   return masked_mean(xent, targets, weights)
 
 
-def loss(params, batch, model_predict, rng):
+def loss(params, batch, model_predict, state, rng):
   """Calculate loss."""
   inputs, targets, weights = unpack_batch(batch)
-  predictions = model_predict(inputs, params, rng=rng)
+  predictions, state = model_predict(inputs, params, state, rng=rng)
   predictions, targets, weights = _make_list(predictions, targets, weights)
   xent = []
   for (pred, target) in zip(predictions, targets):
     xent.append(np.sum(pred * layers.one_hot(target, pred.shape[-1]), axis=-1))
-  return - masked_mean(xent, targets, weights)
+  return - masked_mean(xent, targets, weights), state
 
 
 def log(s, stdout=True):
@@ -151,6 +151,7 @@ def step_log(step, s):
     "step",       # Current training step number.
     "opt_state",  # OptState.
     "history",    # trax.history.History.
+    "model_state",
 ])
 
 
@@ -165,13 +166,15 @@ def restore_state(output_dir):
   """Restore State."""
   params_file = os.path.join(output_dir, "model.pkl")
   if not gfile.exists(params_file):
-    return State(step=None, opt_state=None, history=trax_history.History())
+    return State(step=None, opt_state=None, history=trax_history.History(),
+                 model_state=None)
 
   with gfile.GFile(params_file, "rb") as f:
-    (opt_state, step, history) = pickle.load(f)
+    (opt_state, step, history, model_state) = pickle.load(f)
   log("Model loaded from %s at step %d" % (params_file, step))
   logging.debug("From loaded model : history = %s", history)
-  return State(step=step, opt_state=OptState(*opt_state), history=history)
+  return State(step=step, opt_state=OptState(*opt_state), history=history,
+               model_state=model_state)
 
 
 def _save_gin(output_dir, sw=None):
@@ -194,15 +197,18 @@ def save_state(state, output_dir, keep=False):
     pkl_module = pickle
   params_file = os.path.join(output_dir, "model.pkl")
   with gfile.GFile(params_file, "wb") as f:
-    pkl_module.dump((tuple(state.opt_state), state.step, state.history), f)
+    pkl_module.dump((tuple(state.opt_state), state.step, state.history,
+                     state.model_state), f)
   if keep:
     params_file = os.path.join(output_dir, "model_{}.pkl".format(state.step))
     with gfile.GFile(params_file, "wb") as f:
-      pkl_module.dump((tuple(state.opt_state), state.step, state.history), f)
+      pkl_module.dump((tuple(state.opt_state), state.step, state.history,
+                       state.model_state), f)
   log("Model saved to %s" % params_file, stdout=False)
 
 
-def _save_replicated(opt_state, step, history, n_devices, output_dir, keep):
+def _save_replicated(opt_state, step, history, model_state, n_devices,
+                     output_dir, keep):
   """Save state but given a possibly replicated opt_state."""
   if n_devices > 1:
     first_replica = lambda x: x[0]
@@ -211,8 +217,8 @@ def _save_replicated(opt_state, step, history, n_devices, output_dir, keep):
   # the host in parallel, which is particularly important for cloud TPU.
   if backend.get_name() == "jax":
     opt_state = jax.device_get(opt_state)
-  save_state(State(opt_state=opt_state, step=step, history=history),
-             output_dir, keep=keep)
+  save_state(State(opt_state=opt_state, step=step, history=history,
+                   model_state=model_state), output_dir, keep=keep)
 
 
 def _print_n_params(opt_state, n_devices, step):
@@ -234,27 +240,28 @@ def _print_n_params(opt_state, n_devices, step):
 }
 
 
-def evaluate_train_and_eval(step, inputs, predict_fn, eval_steps, rng,
+def evaluate_train_and_eval(step, inputs, predict_fn, eval_steps, state, rng,
                             train_sw=None, eval_sw=None, history=None):
   """Evalaute on train and eval data, and log metrics."""
   step_log(step, "Evaluation")
-  train_metrics, eval_metrics = [
-      evaluate(  # pylint: disable=g-complex-comprehension
-          itertools.islice(input_stream(), eval_steps),
-          predict_fn,
-          _METRICS,
-          rng)
-      for input_stream in
-      [inputs.train_eval_stream, inputs.eval_stream]]
+  metrics_list = []
+  for input_stream in [inputs.train_eval_stream, inputs.eval_stream]:
+    metrics, state = evaluate(  # pylint: disable=g-complex-comprehension
+        itertools.islice(input_stream(), eval_steps),
+        predict_fn,
+        _METRICS,
+        state, rng)
+    metrics_list.append(metrics)
+  train_metrics, eval_metrics = metrics_list  # pylint: disable=unbalanced-tuple-unpacking
   if train_sw:
     log_metrics(train_metrics, train_sw, "train", step, history=history)
   if eval_sw:
     log_metrics(eval_metrics, eval_sw, "eval", step, history=history)
   step_log(step, "Finished evaluation")
-  return train_metrics, eval_metrics
+  return train_metrics, eval_metrics, state
 
 
-def evaluate(inputs_stream, predict_fn, metric_fns, rng):
+def evaluate(inputs_stream, predict_fn, metric_fns, state, rng):
   """Evaluate.
 
   Args:
@@ -263,25 +270,27 @@ def evaluate(inputs_stream, predict_fn, metric_fns, rng):
       partially applied.
     metric_fns: dict from metric name to metric function, which takes inputs
       and predictions and returns a scalar metric value.
+    state: start state for `predict_fn`.
     rng: random number generator.
 
   Returns:
     metrics: dict from metric name to metric value averaged over the number of
       inputs.
+    state: end state for `predict_fn`.
   """
   metrics = collections.defaultdict(float)
   count = 0
   for inp in inputs_stream:
     count += 1
     rng, subrng = jax_random.split(rng)
-    preds = predict_fn(inp[0], rng=subrng)
+    preds, state = predict_fn(inp[0], state=state, rng=subrng)
     for m, f in six.iteritems(metric_fns):
       metrics[m] += f(inp, preds)
-  return {m: v / count for (m, v) in six.iteritems(metrics)}
+  return {m: v / count for (m, v) in six.iteritems(metrics)}, state
 
 
 def evaluate_loss_train_and_eval(step, inputs, compute_loss_fn, eval_steps,
-                                 rngs,
+                                 state, rngs,
                                  train_sw=None, eval_sw=None, history=None):
   """More efficient evaluation that logs only the loss on train & eval data."""
   step_log(step, "Evaluation")
@@ -290,7 +299,7 @@ def evaluate_loss_train_and_eval(step, inputs, compute_loss_fn, eval_steps,
     total = 0.0
     count = 0.0
     for inp in itertools.islice(input_stream(), eval_steps):
-      loss_values, rngs = compute_loss_fn(inp, rngs)
+      loss_values, state, rngs = compute_loss_fn(inp, state, rngs)
       total += float(numpy.mean(loss_values))
       count += 1.0
     metrics = {"loss": total / count}
@@ -301,12 +310,12 @@ def evaluate_loss_train_and_eval(step, inputs, compute_loss_fn, eval_steps,
   if eval_sw:
     log_metrics(eval_metrics, eval_sw, "eval", step, history=history)
   step_log(step, "Finished evaluation")
-  return train_metrics, eval_metrics
+  return train_metrics, eval_metrics, state
 
 
 def log_metrics(metrics, summ_writer, log_prefix, step, history=None):
   """Log metrics to summary writer and history."""
-  rjust_len = max([len(name) for name in metrics])
+  rjust_len = max([0] + [len(name) for name in metrics])
   for name, value in six.iteritems(metrics):
     step_log(step, "%s %s | % .8f" % (
         log_prefix.ljust(5), name.rjust(rjust_len), value))
@@ -378,23 +387,22 @@ def _jit_predict_fn(model_predict, n_devices, jit=True):
 
   # Multi-devices, pmap and run.
   @functools.partial(backend.pmap, axis_name="batch")
-  def mapped_predict(x, params, rng):
-    return model_predict(x, params, rng=rng)
+  def mapped_predict(x, params, state, rng):
+    return model_predict(x, params, state, rng=rng)
 
-  def predict(x, params=(), rng=None):
+  def predict(x, params=(), state=(), rng=None):
     """Predict function jited and parallelized as requested."""
-    # On one device, jit and run.
-    pred = mapped_predict(
+    pred, state = mapped_predict(
         reshape_by_device(x, n_devices),
         params,
+        state,
         jax_random.split(rng, n_devices))
     # Need to reduce the [device, per-device-batch, ...] tensors back to
     # a [batch, ...] tensor. The tensors may be nested.
-    if not isinstance(pred, (list, tuple)):  # Not nested.
-      batch_size = pred.shape[0] * pred.shape[1]
-      return np.reshape(pred, [batch_size] + list(pred.shape[2:]))
-    batch_size = pred[0].shape[0] * pred[0].shape[1]
-    return [np.reshape(p, [batch_size] + list(p.shape[2:])) for p in pred]
+    def combine(x):
+      batch_size = x.shape[0] * x.shape[1]
+      return np.reshape(x, [batch_size] + list(x.shape[2:]))
+    return layers.nested_map(pred, combine), state
 
   return predict
 
@@ -403,31 +411,34 @@ def predict(x, params=(), rng=None):
 def _jit_update_fn(predict_fn, loss_fn, optimizer, n_devices, jit=True):
   """Get jit-ed update function for loss, optimizer, learning rate function."""
   if n_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
-    def single_update(i, opt_state, batch, rng):
+    def single_update(i, opt_state, batch, state, rng):
       params, slots, opt_params = opt_state
       rng, subrng = jax_random.split(rng[0])
-      grads = backend.grad(loss_fn)(params, batch, predict_fn, rng)
+      grads, state = backend.grad(loss_fn, has_aux=True)(params, batch,
+                                                         predict_fn, state, rng)
       return optimizer.tree_update(
-          i, grads, params, slots, opt_params), [subrng]
+          i, grads, params, slots, opt_params), state, [subrng]
     if jit:
       return backend.jit(single_update)
     else:
       return single_update
 
   @functools.partial(backend.pmap, axis_name="batch")
-  def mapped_update(i, opt_state, batch, rng):
+  def mapped_update(i, opt_state, batch, state, rng):
     """This is a multi-device version of the update function above."""
     # We assume all tensors have the first dimension = n_devices.
     params, slots, opt_params = opt_state
     rng, subrng = jax_random.split(rng)
-    grads = backend.grad(loss_fn)(params, batch, predict_fn, rng)
+    grads, state = backend.grad(loss_fn, has_aux=True)(params, batch,
+                                                       predict_fn, state, rng)
     grads = jax.tree_util.tree_map(
         lambda g: lax.psum(g, "batch"), grads)
     return optimizer.tree_update(
-        i, grads, params, slots, opt_params), subrng
+        i, grads, params, slots, opt_params), state, subrng
 
-  def update(i, opt_state, batch, rng):
-    return mapped_update(numpy.repeat(i, n_devices), opt_state, batch, rng)
+  def update(i, opt_state, batch, state, rng):
+    return mapped_update(numpy.repeat(i, n_devices), opt_state, batch, state,
+                         rng)
 
   return update
 
@@ -436,25 +447,25 @@ def update(i, opt_state, batch, rng):
 def _jit_compute_loss_fn(predict_fn, loss_fn, n_devices, jit=True):
   """Get jit-ed function that computes the loss."""
   if n_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
-    def single_compute_loss(opt_state, batch, rng):
+    def single_compute_loss(opt_state, batch, state, rng):
       rng, subrng = jax_random.split(rng[0])
-      return loss_fn(opt_state[0], batch, predict_fn, rng), [subrng]
+      return loss_fn(opt_state[0], batch, predict_fn, state, rng), [subrng]
     if jit:
       return backend.jit(single_compute_loss)
     else:
       return single_compute_loss
 
   @functools.partial(backend.pmap, axis_name="batch")
-  def mapped_compute_loss(opt_state, batch, rng):
+  def mapped_compute_loss(opt_state, batch, state, rng):
     """This is a multi-device version of the update function above."""
     # We assume all tensors have the first dimension = n_devices.
     rng, subrng = jax_random.split(rng)
-    loss_val = loss_fn(opt_state[0], batch, predict_fn, rng)
-    return loss_val, subrng
+    loss_val, state = loss_fn(opt_state[0], batch, predict_fn, state, rng)
+    return loss_val, state, subrng
 
-  def compute_loss(opt_state, batch, rng):
+  def compute_loss(opt_state, batch, state, rng):
     return mapped_compute_loss(
-        opt_state, reshape_by_device(batch, n_devices), rng)
+        opt_state, reshape_by_device(batch, n_devices), state, rng)
 
   return compute_loss
 
@@ -532,9 +543,13 @@ def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs,
     model_input_shape = layers.nested_map(
         model_input_shape, lambda x: x if x else 1)
     def initialize(input_shape, input_dtype, init_rng):
-      params = model_train.initialize(input_shape, input_dtype, init_rng)
+      # We need to create a new model instance and not reuse `model_train` here,
+      # because `m.initialize` puts cached parameter values in `m` and hence the
+      # next call of `m.initialize` will give wrong results.
+      params, state = model(mode="train").initialize(input_shape, input_dtype,
+                                                     init_rng)
       (slots, opt_params) = opt.tree_init(params)
-      return OptState(params, slots, opt_params)
+      return (OptState(params, slots, opt_params), state)
     if _is_jit_init():
       # JIT parameter initialization to avoid memory fragmentation
       initialize = backend.jit(initialize, static_argnums=(0, 1))
@@ -559,6 +574,7 @@ def initialize(input_shape, input_dtype, init_rng):
     self._lr_fn = None
     self._opt_state = None
     self._step = None
+    self._model_state = None
 
     if output_dir is not None:
       self.reset(output_dir)
@@ -591,10 +607,14 @@ def reset(self, output_dir):
     self._history = history
     if state.opt_state:
       opt_state = state.opt_state
+      model_state = state.model_state
     else:
-      opt_state = self._initialize()
+      opt_state, model_state = self._initialize()
+      model_state = layers.nested_map(
+          model_state, self._maybe_replicate)
     self._opt_state = OptState(*layers.nested_map(
         opt_state, self._maybe_replicate))
+    self._model_state = model_state
     if not state.opt_state:
       self._maybe_save_state(keep=False)
 
@@ -611,7 +631,8 @@ def n_devices(self):
   @property
   def state(self):
     return State(
-        opt_state=self._opt_state, step=self._step, history=self._history)
+        opt_state=self._opt_state, step=self._step, history=self._history,
+        model_state=self._model_state)
 
   @property
   def learning_rate(self):
@@ -623,14 +644,15 @@ def learning_rate(self):
 
   def _maybe_replicate(self, x):
     if self._n_devices > 1:
-      return numpy.broadcast_to(x, (self._n_devices,) + x.shape)
+      return np.broadcast_to(x, (self._n_devices,) + x.shape)
     else:
       return x
 
   def _maybe_save_state(self, keep):
     if self._should_save:
       _save_replicated(self._opt_state, self._step, self._history,
-                       self._n_devices, self._output_dir, keep)
+                       self._model_state, self._n_devices, self._output_dir,
+                       keep)
 
   def save_gin(self):
     _save_gin(self._output_dir, self._train_sw)
@@ -648,8 +670,8 @@ def _train_step(self, next_train_batch):
     opt_state = opt_state._replace(opt_params=opt_params)
 
     # Run the update.
-    (params, slots), self._rngs = self._jit_update_fn(
-        self._step, opt_state, next_train_batch, self._rngs)
+    (params, slots), self._model_state, self._rngs = self._jit_update_fn(
+        self._step, opt_state, next_train_batch, self._model_state, self._rngs)
     self._opt_state = opt_state._replace(params=params, slots=slots)
     self._step += 1
 
@@ -700,12 +722,13 @@ def train_epoch(self, epoch_steps, eval_steps):
 
   def evaluate(self, eval_steps):
     _, rng = jax_random.split(self._rngs[0])
-    evaluate_train_and_eval(
+    _, _, self._model_state = evaluate_train_and_eval(
         step=self._step,
         inputs=self._inputs,
         predict_fn=functools.partial(self._jit_model_predict_eval,
                                      params=self._opt_state[0]),
         eval_steps=eval_steps,
+        state=self._model_state,
         rng=rng,
         train_sw=self._train_sw,
         eval_sw=self._eval_sw,
@@ -724,13 +747,15 @@ def save_computation_graphs(self, save_backward_graph):
       next_train_batch = reshape_by_device(next_train_batch, self._n_devices)
     params = self._opt_state[0]
     forward_computation = jax.xla_computation(self._model_predict_eval)(
-        next_train_batch[0], params=params, rng=self._rngs[0])
+        next_train_batch[0], params=params, state=self._model_state,
+        rng=self._rngs[0])
     with gfile.GFile(os.path.join(output_dir, "forward.txt"), "w") as f:
       f.write(forward_computation.GetHloText())
     with gfile.GFile(os.path.join(output_dir, "forward.dot"), "w") as f:
       f.write(forward_computation.GetHloDotGraph())
     backward_computation = jax.xla_computation(self._jit_update_fn)(
-        self._step, self._opt_state, next_train_batch, self._rngs)
+        self._step, self._opt_state, next_train_batch, self._model_state,
+        self._rngs)
     with gfile.GFile(os.path.join(output_dir, "backward.txt"), "w") as f:
       f.write(backward_computation.GetHloText())
     if save_backward_graph:  # Backward graphs can be large so we guard it.
@@ -756,12 +781,13 @@ def __init__(self, *args, **kwargs):
 
   def evaluate(self, eval_steps):
     # Evaluate only the loss function (a more efficient, jitted, implementation)
-    evaluate_loss_train_and_eval(
+    self._model_state = evaluate_loss_train_and_eval(
         step=self._step,
         inputs=self._inputs,
         compute_loss_fn=functools.partial(self._jit_compute_loss,
                                           self._opt_state),
         eval_steps=eval_steps,
+        state=self._model_state,
         rngs=self._rngs,
         train_sw=self._train_sw,
         eval_sw=self._eval_sw,
@@ -796,8 +822,8 @@ def train(output_dir,
     output_dir: Directory where to put the logs and checkpoints.
     model: The model to train as a callable returning 2 callables, an init_fn
       and apply_fn.
-    loss_fn: callable with signature: params, trax.inputs.Inputs, model, rng
-      -> loss.
+    loss_fn: callable with signature: params, trax.inputs.Inputs, model, state,
+      rng -> loss.
     inputs: callable returning trax.inputs.Inputs.
     optimizer: The optimizer (see optimizers/base.py for signature).
     lr_schedule: A learning rate schedule as a function that takes history and
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index 22b7b7cf1..bd07b0ba9 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -25,14 +25,15 @@
 from absl.testing import parameterized
 
 import gin
-import jax
 from jax import test_util  # pylint: disable=unused-import
 from jax.config import config
+from jax.lib import xla_bridge
 import numpy as onp
 
 from tensor2tensor.trax import backend
 from tensor2tensor.trax import inputs as inputs_lib
 from tensor2tensor.trax import layers
+from tensor2tensor.trax import learning_rate as lr
 from tensor2tensor.trax import models
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
@@ -45,7 +46,7 @@
 
 def test_inputs(n_classes, with_weights=False):
   """Make trax.inputs.Inputs."""
-  batch_size = 2
+  batch_size = 2 * xla_bridge.device_count()
   input_shape = (6, 6, 3)
 
   def input_stream():
@@ -85,16 +86,22 @@ def tmp_dir(self):
 
   @parameterized.parameters(BACKENDS)
   def test_train_eval_predict(self, backend_name):
-    if jax.lib.xla_bridge.device_count() > 1 and backend_name == "tf":
+    if xla_bridge.device_count() > 1 and backend_name == "tf":
       self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
     with backend.use_backend(backend_name), self.tmp_dir() as output_dir:
       # Prepare model and inputs
       n_classes = 4
       train_steps = 2
       eval_steps = 2
-      model_fn = functools.partial(models.MLP,
-                                   d_hidden=16,
-                                   n_output_classes=n_classes)
+      # Adds Dropout and BatchNorm to test state handling.
+      mlp = functools.partial(models.MLP,
+                              d_hidden=16,
+                              n_output_classes=n_classes)
+      def model_fn(mode="train"):
+        return layers.Model(layers.Dropout(mode=mode, rate=0.1),
+                            layers.BatchNorm(mode=mode),
+                            mlp(mode=mode))
+
       inputs = lambda _: test_inputs(n_classes)
 
       # Train and evaluate
@@ -120,7 +127,7 @@ def test_train_eval_predict(self, backend_name):
 
   @parameterized.parameters(BACKENDS)
   def test_train_eval_predict_sm3(self, backend_name):
-    if jax.lib.xla_bridge.device_count() > 1 and backend_name == "tf":
+    if xla_bridge.device_count() > 1 and backend_name == "tf":
       self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
     with backend.use_backend(backend_name), self.tmp_dir() as output_dir:
       # Prepare model and inputs
@@ -156,7 +163,7 @@ def test_train_eval_predict_sm3(self, backend_name):
 
   @parameterized.parameters(BACKENDS)
   def test_train_restart(self, backend_name):
-    if jax.lib.xla_bridge.device_count() > 1 and backend_name == "tf":
+    if xla_bridge.device_count() > 1 and backend_name == "tf":
       self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
     with backend.use_backend(backend_name), self.tmp_dir() as output_dir:
       # Prepare model and inputs
@@ -187,7 +194,7 @@ def test_train_restart(self, backend_name):
 
   @parameterized.parameters(BACKENDS)
   def test_train_with_weights(self, backend_name):
-    if jax.lib.xla_bridge.device_count() > 1 and backend_name == "tf":
+    if xla_bridge.device_count() > 1 and backend_name == "tf":
       self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
     with backend.use_backend(backend_name), self.tmp_dir() as output_dir:
       gin.bind_parameter("unpack_batch.has_weights", True)
@@ -211,6 +218,31 @@ def test_train_with_weights(self, backend_name):
       # Assert total train steps
       self.assertEqual(state.step, train_steps)
 
+  @parameterized.parameters(BACKENDS)
+  def test_reset_twice(self, backend_name):
+    if xla_bridge.device_count() > 1 and backend_name == "tf":
+      self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
+    with backend.use_backend(backend_name), self.tmp_dir() as output_dir1, \
+          self.tmp_dir() as output_dir2:
+      n_classes = 4
+      model_fn = functools.partial(models.MLP,
+                                   d_hidden=16,
+                                   n_output_classes=n_classes)
+      inputs = lambda _: test_inputs(n_classes)
+
+      trainer = trax.Trainer(
+          model=model_fn,
+          loss_fn=trax.loss,
+          optimizer=trax_opt.SM3,
+          lr_schedule=lr.MultifactorSchedule,
+          inputs=inputs,
+      )
+
+      trainer.reset(output_dir1)
+      trainer.evaluate(1)
+      trainer.reset(output_dir2)
+      trainer.evaluate(1)
+
 
 MASKED_MEAN_TEST_BACKENDS = ["numpy"]
 

From 4289bcfa12b75f434bfa2f73ca5adc088ea35731 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 21 Aug 2019 11:15:49 -0700
Subject: [PATCH 2300/2720] Add attr in test deps to fix a travis issue and
 remove gsutil since it isn't being used in tests.

PiperOrigin-RevId: 264650849
---
 setup.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index a4b6bad35..beb7607df 100644
--- a/setup.py
+++ b/setup.py
@@ -66,11 +66,13 @@
         'tensorflow-hub': ['tensorflow-hub>=0.1.1'],
         'tests': [
             'absl-py',
+            # Needed to fix a Travis pytest error.
+            # https://github.com/Julian/jsonschema/issues/449#issuecomment-411406525
+            'attrs>=17.4.0',
             'pytest>=3.8.0',
             'mock',
             'pylint',
             'jupyter',
-            'gsutil',
             'matplotlib',
             # Need atari extras for Travis tests, but because gym is already in
             # install_requires, pip skips the atari extras, so we instead do an

From 0591d94088b3e192107f6e85a8042ba50dcd4852 Mon Sep 17 00:00:00 2001
From: lazylife7157 <lazylife7157@gmail.com>
Date: Thu, 22 Aug 2019 04:28:05 +0900
Subject: [PATCH 2301/2720] Change confusing function name (#1669)

---
 tensor2tensor/utils/beam_search.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index 0632d3b12..85a32f8a1 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -724,8 +724,9 @@ def inner_loop(i, alive_seq, alive_log_probs, finished_seq, finished_scores,
     return (i + 1, alive_seq, alive_log_probs, finished_seq, finished_scores,
             finished_flags, states)
 
-  def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
-                   finished_scores, unused_finished_in_finished, unused_states):
+  def _is_not_finished(i, unused_alive_seq, alive_log_probs,
+                       unused_finished_seq, finished_scores,
+                       unused_finished_in_finished, unused_states):
     """Checking termination condition.
 
     We terminate when we decoded up to decode_length or the lowest scoring item
@@ -781,7 +782,7 @@ def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
     state_struc = nest.map_structure(get_state_shape_invariants, states)
   (_, alive_seq, alive_log_probs, finished_seq, finished_scores,
    finished_flags, states) = tf.while_loop(
-       _is_finished,
+       _is_not_finished,
        inner_loop, [
            tf.constant(0), alive_seq, alive_log_probs, finished_seq,
            finished_scores, finished_flags, states

From 67ca6058d4c26a7ec60d86f03eb2b8cc4881c322 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 21 Aug 2019 12:31:48 -0700
Subject: [PATCH 2302/2720] * Some more tests started enabling eager, so
 separate them out in pytest. * Disable reversible_layers_test.py since that
 would either need to be tested   against a new TFP release or against TFP
 nightly (this is a bigger change). * While debugging the above I found
 separating tests into top-level directories   and having a catch-all for all
 other tests makes the errors on travis *much   more readable* -- so doing
 that refactoring.

PiperOrigin-RevId: 264667040
---
 oss_scripts/oss_tests.sh | 74 +++++++++++++++++++++++++++++++---------
 1 file changed, 58 insertions(+), 16 deletions(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index d27f0d604..62343e253 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -28,8 +28,10 @@ fi
 python -c "from tensor2tensor.models import transformer; print(transformer.Transformer.__name__)"
 set_status
 
-# Run tests
-# Ignores:
+# We need to run some tests separately (because they enable eager or due to
+# other reasons). We also test the tests in the top-level-directories separately
+# to get more readable error messages.
+
 # Tested separately:
 #   * registry_test
 #   * trainer_lib_test
@@ -37,51 +39,64 @@ set_status
 #   * trainer_model_based_test
 #   * allen_brain_test
 #   * models/research
+
+
 # algorithmic_math_test: flaky
-# test_utils.py is not a test, but pytest thinks it is.
 # subword_text_encoder_ops_test, pack_sequences_ops_test: interface with C++ ops
-# trax tests need C++
-# TODO(davidso): Re-enable EvolvedTransformer when possible.
-# others (see below) enable eager, so can't be tested along with the others in
-# pytest
 pytest --disable-warnings \
-  --ignore=tensor2tensor/bin/t2t_trainer_test.py \
   --ignore=tensor2tensor/data_generators/algorithmic_math_test.py \
   --ignore=tensor2tensor/data_generators/allen_brain_test.py \
   --ignore=tensor2tensor/data_generators/ops/pack_sequences_ops_test.py \
   --ignore=tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py \
   --ignore=tensor2tensor/data_generators/problem_test.py \
+  tensor2tensor/data_generators
+set_status
+
+
+pytest --disable-warnings \
   --ignore=tensor2tensor/envs/mujoco_problems_test.py \
   --ignore=tensor2tensor/envs/rendered_env_problem_test.py \
+  tensor2tensor/envs/
+set_status
+
+
+pytest --disable-warnings \
   --ignore=tensor2tensor/layers/bayes_test.py \
   --ignore=tensor2tensor/layers/common_attention_test.py \
   --ignore=tensor2tensor/layers/common_layers_test.py \
   --ignore=tensor2tensor/layers/common_video_test.py \
   --ignore=tensor2tensor/layers/discretization_test.py \
+  --ignore=tensor2tensor/layers/gaussian_process_test.py \
   --ignore=tensor2tensor/layers/latent_layers_test.py \
   --ignore=tensor2tensor/layers/modalities_test.py \
   --ignore=tensor2tensor/layers/ngram_test.py \
   --ignore=tensor2tensor/layers/reversible_layers_test.py \
+  tensor2tensor/layers/
+set_status
+
+
+# TODO(davidso): Re-enable EvolvedTransformer when possible.
+pytest --disable-warnings \
   --ignore=tensor2tensor/models/evolved_transformer_test.py \
   --ignore=tensor2tensor/models/research \
   --ignore=tensor2tensor/models/video/nfg_conv3d_test.py \
   --ignore=tensor2tensor/models/video/nfg_conv_lstm_test.py \
   --ignore=tensor2tensor/models/video/nfg_conv_test.py \
   --ignore=tensor2tensor/models/video/nfg_uncond_test.py \
-  --ignore=tensor2tensor/rl \
-  --ignore=tensor2tensor/trax \
+  tensor2tensor/models/
+set_status
+
+
+# test_utils.py is not a test, but pytest thinks it is.
+pytest --disable-warnings \
   --ignore=tensor2tensor/utils/registry_test.py \
   --ignore=tensor2tensor/utils/t2t_model_test.py \
   --ignore=tensor2tensor/utils/test_utils.py \
   --ignore=tensor2tensor/utils/test_utils_test.py \
   --ignore=tensor2tensor/utils/trainer_lib_test.py \
-  --ignore=tensor2tensor/visualization/visualization_test.py \
-  --deselect=tensor2tensor/utils/beam_search_test.py::BeamSearchTest::testTPUBeam
+  tensor2tensor/utils/
 set_status
 
-# TODO(afrozm): Enable trax tests they currently need GLIBCXX_3.4.21
-# Travis Error:
-# ImportError: /usr/lib/x86_64-linux-gnu/libstdc++.so.6: version `GLIBCXX_3.4.21' not found (required by /home/travis/virtualenv/python3.6.3/lib/python3.6/site-packages/jaxlib/_pywrap_xla.so)
 
 # These tests enable eager, so are tested separately.
 pytest --disable-warnings \
@@ -91,15 +106,21 @@ pytest --disable-warnings \
   tensor2tensor/layers/common_layers_test.py \
   tensor2tensor/layers/common_video_test.py \
   tensor2tensor/layers/discretization_test.py \
+  tensor2tensor/layers/gaussian_process_test.py \
   tensor2tensor/layers/latent_layers_test.py \
   tensor2tensor/layers/modalities_test.py \
   tensor2tensor/layers/ngram_test.py \
-  tensor2tensor/layers/reversible_layers_test.py \
+  tensor2tensor/keras \
   tensor2tensor/utils/t2t_model_test.py \
   tensor2tensor/utils/test_utils_test.py \
   --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary
 set_status
 
+# TODO(afrozm): Re-enable once TFP's new release comes out (0.8.0) or modify
+# stuff so that we test against tfp-nightly.
+#pytest --disable-warnings tensor2tensor/layers/reversible_layers_test.py
+#set_status
+
 pytest --disable-warnings tensor2tensor/utils/registry_test.py
 set_status
 
@@ -112,6 +133,27 @@ set_status
 pytest --disable-warnings tensor2tensor/data_generators/allen_brain_test.py
 set_status
 
+# All other tests not tested above.
+
+# trax tests need C++
+# TODO(afrozm): Enable trax tests they currently need GLIBCXX_3.4.21
+# Travis Error:
+# ImportError: /usr/lib/x86_64-linux-gnu/libstdc++.so.6: version `GLIBCXX_3.4.21' not found (required by /home/travis/virtualenv/python3.6.3/lib/python3.6/site-packages/jaxlib/_pywrap_xla.so)
+pytest --disable-warnings \
+  --ignore=tensor2tensor/bin/t2t_trainer_test.py \
+  --ignore=tensor2tensor/data_generators \
+  --ignore=tensor2tensor/envs \
+  --ignore=tensor2tensor/keras \
+  --ignore=tensor2tensor/layers \
+  --ignore=tensor2tensor/models \
+  --ignore=tensor2tensor/rl \
+  --ignore=tensor2tensor/trax \
+  --ignore=tensor2tensor/utils \
+  --ignore=tensor2tensor/visualization \
+  --deselect=tensor2tensor/utils/beam_search_test.py::BeamSearchTest::testTPUBeam
+set_status
+
+
 # TODO(afrozm): Enable this unconditionally?
 
 ## Test models/research only against tf-nightly

From 5386fb120a6ba2f1c635f08f0b22413349cc6caa Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 21 Aug 2019 14:36:40 -0700
Subject: [PATCH 2303/2720] Small corrections.

PiperOrigin-RevId: 264694817
---
 tensor2tensor/trax/layers/base.py             |  4 +-
 tensor2tensor/trax/layers/combinators.py      |  2 +-
 .../models/research/transformer_revnet.py     | 48 ++++++++++---------
 3 files changed, 28 insertions(+), 26 deletions(-)

diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 79568de0e..842ec3025 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -258,12 +258,12 @@ def __call__(self, x, params=(), state=(), **kwargs):
       #   JAX.
 
       assert state is (), (  # pylint: disable=literal-comparison
-          'Custom gradients do not allow non-trivial start state.')
+          'Custom gradients require trivial start state. Got %s' % str(state))
 
       def check_end_state(output_state):
         output, state = output_state
         assert state is (), (  # pylint: disable=literal-comparison
-            'Custom gradients do not allow non-trivial end state.')
+            'Custom gradients require trivial end state. Got %s' % str(state))
         return output
 
       # See this link for how custom transformations are defined in JAX:
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index f51fb5ba8..800f48a7e 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -181,7 +181,7 @@ def call(self, xs, params=(), state=(), **kwargs):
       raise ValueError('number of params ({}) not equal to number of layers '
                        '({})'.format(len(params), n_layers))
     if n_layers != 1 and len(state) != n_layers:
-      raise ValueError('number of params ({}) not equal to number of layers '
+      raise ValueError('length of state ({}) not equal to number of layers '
                        '({})'.format(len(state), n_layers))
     for layer, p, s, rng in zip(self._sublayers, params, state, rngs):
       is_stack_just_one_item = (_count_items(stack) == 1)
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 0c92b7e68..ff5880b0f 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -65,11 +65,12 @@ def n_outputs(self):
     """Specifies how many data tensors this layer promises as output."""
     return self._n_sections
 
-  def call(self, inputs, params=(), **kwargs):
+  def call(self, inputs, params=(), state=(), **kwargs):
     rngs = _pop_rng_and_split(kwargs, len(inputs))
-    result = [self._layer(x, params=params, rng=r, **kwargs)
-              for x, r in zip(inputs, rngs)]
-    return tuple(result)
+    results = [self._layer(x, params=params, state=state, rng=r, **kwargs)
+               for x, r in zip(inputs, rngs)]
+    result_outputs, result_states = zip(*results)
+    return tuple(result_outputs), tuple(result_states)
 
   def new_parameters(self, input_shape, input_dtype, rng):
     first_shape = input_shape[0]
@@ -122,12 +123,13 @@ def __init__(self, n_sections=2, axis=-1):
     self._n_sections = n_sections
     self._axis = axis
 
-  def call(self, inputs, params=(), **kwargs):
+  def call(self, inputs, params=(), state=(), **kwargs):
     del params, kwargs
-    return tuple(backend.numpy.split(inputs, self._n_sections, self._axis))
+    res = tuple(backend.numpy.split(inputs, self._n_sections, self._axis))
+    return res, state
 
   def new_parameters(self, input_shapes, input_dtype, rng):
-    return ()
+    return (), ()
 
   def n_inputs(self):
     """Specifies how many data tensors this layer expects as input."""
@@ -167,9 +169,9 @@ def n_outputs(self):
     return self._n_sections
 
   def new_parameters(self, input_shape, input_dtype, rng):
-    return ()
+    return (), ()
 
-  def call(self, inputs, params=(), **kwargs):
+  def call(self, inputs, params=(), state=(), **kwargs):
     del params, kwargs
     x1, x2 = inputs
 
@@ -177,7 +179,7 @@ def call(self, inputs, params=(), **kwargs):
     x2_split = backend.numpy.split(x2, self._n_sections, self._axis)
 
     res = [backend.numpy.concatenate(ys, -1) for ys in zip(x1_split, x2_split)]
-    return tuple(res)
+    return tuple(res), state
 
   def reverse(self, output, params=(), **kwargs):
     del params, kwargs
@@ -288,7 +290,7 @@ def __init__(self, n_heads=1, d_head=64,
     # The lack of a bias term here is consistent with the tensor2tensor
     # implementation, and shouldn't have an effect on modeling quality.
 
-  def call(self, x, params, **kwargs):
+  def call(self, x, params, state, **kwargs):
     del kwargs
     seqlen = x.shape[1]
     res = np.dot(x, params)
@@ -300,13 +302,13 @@ def call(self, x, params, **kwargs):
     # n_batch, n_heads, seqlen, d_head -> n_batch*n_heads, seqlen, d_head
     res = np.reshape(res, (-1, seqlen, self._d_head))
 
-    return res
+    return res, state
 
   def new_parameters(self, input_shape, input_dtype, rng):
     del input_dtype
     w = self._kernel_initializer(
         (input_shape[-1], self._n_heads * self._d_head), rng)
-    return w
+    return w, ()
 
 
 class ComputeAttentionOutput(tl.Layer):
@@ -321,7 +323,7 @@ def __init__(self, n_heads=1, d_model=1024,
     # The lack of a bias term here is consistent with the tensor2tensor
     # implementation, and shouldn't have an effect on modeling quality.
 
-  def call(self, x, params, **kwargs):
+  def call(self, x, params, state, **kwargs):
     del kwargs
     seqlen = x.shape[1]
     d_head = x.shape[2]
@@ -330,13 +332,13 @@ def call(self, x, params, **kwargs):
     x = np.transpose(x, (0, 2, 1, 3))  # -> n_batch, seqlen, n_heads, d_head
     x = np.reshape(x, (-1, seqlen, self._n_heads * d_head))
 
-    return np.dot(x, params)
+    return np.dot(x, params), state
 
   def new_parameters(self, input_shape, input_dtype, rng):
     del input_dtype
     w = self._kernel_initializer(
         (input_shape[-1] * self._n_heads, self._d_model), rng)
-    return w
+    return w, ()
 
 
 class ApplyAttentionWrapper(tl.Parallel):
@@ -374,14 +376,14 @@ def __init__(self, dropout, mode):
     self._dropout = dropout
     self._mode = mode
 
-  def call(self, inputs, params=(), rng=None, **kwargs):
+  def call(self, inputs, params=(), state=(), rng=None, **kwargs):
     del params
     q, k, v = inputs
     mask_size = q.shape[-2]
     mask = np.tril(np.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
     res = tl.DotProductAttention(
         q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=rng)
-    return res
+    return res, state
 
   def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
     # Simultaneous forward pass and backprop through the attention mechanism.
@@ -391,7 +393,7 @@ def do_call(x):
     return output, vjpfun(ct)[0]
 
   def new_parameters(self, input_shapes, input_dtype, rng):
-    return ()
+    return (), ()
 
   def n_inputs(self):
     return 3
@@ -413,9 +415,9 @@ def __init__(self, loop_stride, dropout, mode):
     else:
       self.dropout = None
 
-  def call(self, inputs, params=(), **kwargs):
+  def call(self, inputs, params=(), state=(), **kwargs):
     output, _ = self.forward_and_vjp(inputs, None, params=params, **kwargs)
-    return output
+    return output, state
 
   def forward_and_vjp(self, inputs, ct, params=(), rng=None, **kwargs):
     # This is the core of the memory-efficient attention implementation, where
@@ -547,9 +549,9 @@ def __init__(self, dropout, mode, n_bins=64):
     super(DummyHashedAttention, self).__init__(dropout, mode)
     self.n_bins = n_bins
 
-  def call(self, inputs, params=(), **kwargs):
+  def call(self, inputs, params=(), state=(), **kwargs):
     output, _ = self.forward_and_vjp(inputs, None, params=params, **kwargs)
-    return output
+    return output, state
 
   def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
     del params, kwargs

From 8e19adfdefb0ca17b67b3b896cb3774a94c656b5 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 21 Aug 2019 14:58:27 -0700
Subject: [PATCH 2304/2720] Bump version to 1.14.0

PiperOrigin-RevId: 264699622
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index beb7607df..cfc678d98 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.13.4',
+    version='1.14.0',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',

From aa1a32d64c1c6ddc4f47445f0afd4f007ffc148f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 22 Aug 2019 10:12:46 -0700
Subject: [PATCH 2305/2720] [trax] use jax.eval_shape, reduce dependence on
 internal JAX APIs [jax_md] update to new jax.eval_shape api (includes dtypes)

PiperOrigin-RevId: 264861681
---
 tensor2tensor/trax/backend.py          | 40 +++-----------------------
 tensor2tensor/trax/layers/base_test.py |  4 +--
 2 files changed, 6 insertions(+), 38 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index d73bb1af7..866850f9b 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -20,13 +20,11 @@
 from __future__ import print_function
 
 import contextlib
-import functools
 import gin
 
 import jax
 from jax import lax
 from jax import random as jax_random
-from jax.interpreters import partial_eval as pe
 import jax.numpy as jnp
 import jax.scipy.special as jax_special
 import numpy as onp
@@ -116,39 +114,6 @@ def __repr__(self):
     return "[shape:" + str(self.shape) + ", dtype:" + str(self.dtype) + "]"
 
 
-# TODO(lukaszkaiser): remove this function once JAX has an analogue.
-# pylint: disable=missing-docstring
-def _jax_eval_on_shapes(f, *args):
-  def abstractify(x):
-    return jax.abstract_arrays.raise_to_shaped(jax.core.get_aval(x))
-
-  def make_array(arg):
-    return jnp.zeros(shape=arg.shape, dtype=arg.dtype)
-
-  def turn_back_into_pytree(x):
-    if isinstance(x, jax.core.JaxTuple):
-      return tuple([turn_back_into_pytree(y) for y in x])
-    return x
-
-  def get_shapes_and_types(x):
-    if isinstance(x, jax.core.AbstractTuple):
-      return tuple([get_shapes_and_types(y) for y in x])
-    return ShapeType(x.shape, x.dtype)
-
-  def f_jaxtuple(*jaxtuple_args):
-    args = map(turn_back_into_pytree, jaxtuple_args)
-    out = f(*args)
-    res, _ = jax.api_util.pytree_to_jaxtupletree(out)
-    return res
-
-  args_arrays = nested_map(args, make_array)
-  jaxtuple_args, _ = jax.util.unzip2(
-      map(jax.api_util.pytree_to_jaxtupletree, args_arrays))
-  res = pe.abstract_eval_fun(f_jaxtuple, *map(abstractify, jaxtuple_args))
-
-  return get_shapes_and_types(res)
-
-
 def jax_eval_on_shapes(f):
   """Returns a function that evaluates `f` given input shapes and dtypes.
 
@@ -163,7 +128,10 @@ def jax_eval_on_shapes(f):
     their shapes/dtypes represented by `ShapeType`, and whose return values are
     `ShapeType`s with the same nested structure as `f`'s return values.
   """
-  return functools.partial(_jax_eval_on_shapes, f)
+  def shape_fun(*args, **kwargs):
+    jax_shapes = jax.eval_shape(f, *args, **kwargs)
+    return nested_map(jax_shapes, lambda x: ShapeType(x.shape, x.dtype))
+  return shape_fun
 
 
 # The default value of dtype is different from jax_random.randint
diff --git a/tensor2tensor/trax/layers/base_test.py b/tensor2tensor/trax/layers/base_test.py
index 20b8020d6..103040c7e 100644
--- a/tensor2tensor/trax/layers/base_test.py
+++ b/tensor2tensor/trax/layers/base_test.py
@@ -51,7 +51,7 @@ def has_custom_grad(self):
         return True
 
       def custom_grad(self, inputs, output, ct, params, **kwargs):
-        return (backend.numpy.zeros_like(ct), None)
+        return (backend.numpy.zeros_like(ct), ())
 
     layer = IdWithZeroGrad()
     rng = backend.random.get_prng(0)
@@ -81,7 +81,7 @@ def has_custom_grad(self):
         return True
 
       def custom_grad(self, inputs, output, ct, params, **kwargs):
-        return (inputs, None)
+        return (inputs, ())
 
     layer = IdWithIdGrad()
     rng = backend.random.get_prng(0)

From 2c601178b82917a8bb7b650b36ba8caab4cf21f2 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 22 Aug 2019 11:24:29 -0700
Subject: [PATCH 2306/2720] Fix a bug in computing the step number for the
 optimizer

In the current version of the code, we reset the step number every epoch, but don't reset the optimizer slots, which means that in Adam the bias correction is computed incorrectly and the moment moving averages are overestimated: m_hat = sum[g_i*beta^{n-i} for i = 1..n] / sum[beta^{n-i} for i = 1..m], where m << n. Second moment has higher decay so is overestimated more, which means that the effective learning rate is decreased.

So we need to either 1. not reset the optimizer every epoch, or 2. reset everything, including the slots. I'm pretty sure we should do 1., because otherwise we don't get any advantages from using adaptive optimizers - the moving averages need time to warm up and ~10 optimization steps which we typically do in PPO is not enough. That's also what they do in OpenAI baselines: https://github.com/openai/baselines/blob/master/baselines/ppo2/model.py (that's not meant as any proof, just some extra evidence ;)).

PiperOrigin-RevId: 264879827
---
 tensor2tensor/trax/rl/ppo.py         | 13 ++++---
 tensor2tensor/trax/rl/ppo_trainer.py | 56 +++++++++++++++++-----------
 2 files changed, 42 insertions(+), 27 deletions(-)

diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index cfe8edee0..530dda7d5 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -793,25 +793,28 @@ def maybe_restore_params(output_dir, policy_and_value_net_params, state):
     state: policy state.
 
   Returns:
-    triple (restore (bool), params, iter(int)) where iter is the epoch from
-    which we restored the params, 0 is restore = False.
+    tuple (restore (bool), params, state, iter (int), opt_step (int)) where iter
+    is the epoch from which we restored the params, 0 is restore = False, and
+    opt_step is the total optimization step (sum of all optimization steps made
+    up to the current epoch).
   """
   model_files = gfile.glob(os.path.join(output_dir, "model-??????.pkl"))
   for model_file in reversed(sorted(model_files)):
     logging.info("Trying to restore model from %s", model_file)
     try:
       with gfile.GFile(model_file, "rb") as f:
-        loaded_policy_and_value_net_params, loaded_state = pickle.load(f)
+        loaded_policy_and_value_net_params, loaded_state, total_opt_step = (
+            pickle.load(f))
         policy_and_value_net_params = loaded_policy_and_value_net_params
         state = loaded_state
       model_file_basename = os.path.basename(model_file)  # model-??????.pkl
       i = int(filter(str.isdigit, model_file_basename))
-      return True, policy_and_value_net_params, state, i
+      return True, policy_and_value_net_params, state, i, total_opt_step
     except EOFError as e:
       logging.error("Unable to load model from: %s with %s", model_file, e)
       # Try an older version.
       continue
-  return False, policy_and_value_net_params, state, 0
+  return False, policy_and_value_net_params, state, 0, 0
 
 
 def write_eval_reward_summaries(reward_stats_by_mode, summary_writer, epoch):
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index b5479d4ed..1ff5fc6bd 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -155,9 +155,9 @@ def __init__(
 
     # Maybe restore the policy params. If there is nothing to restore, then
     # iteration = 0 and policy_and_value_net_params are returned as is.
-    restored, policy_and_value_net_params, self._model_state, self._epoch = (
-        ppo.maybe_restore_params(output_dir, policy_and_value_net_params,
-                                 self._model_state))
+    (restored, policy_and_value_net_params, self._model_state, self._epoch,
+     self._total_opt_step) = ppo.maybe_restore_params(
+         output_dir, policy_and_value_net_params, self._model_state)
 
     if restored:
       logging.info("Restored parameters from iteration [%d]", self._epoch)
@@ -329,24 +329,34 @@ def train_epoch(self):
       k1, k2, k3 = jax_random.split(key, num=3)
       t = time.time()
       # Update the optimizer state.
-      self._policy_and_value_opt_state, self._model_state = ppo.policy_and_value_opt_step(
-          j,
-          self._policy_and_value_opt_state,
-          self._policy_and_value_opt_update,
-          self._policy_and_value_get_params,
-          self._policy_and_value_net_apply,
-          log_probabs_traj,
-          value_predictions_traj,
-          padded_observations,
-          padded_actions,
-          padded_rewards,
-          reward_mask,
-          c1=self._c1,
-          c2=self._c2,
-          gamma=self._gamma,
-          lambda_=self._lambda_,
-          state=self._model_state,
-          rng=k1)
+      self._policy_and_value_opt_state, self._model_state = (
+          ppo.policy_and_value_opt_step(
+              # We pass the optimizer slots between PPO epochs, so we need to
+              # pass the optimization step as well, so for example the
+              # bias-correction in Adam is calculated properly. Alternatively we
+              # could reset the slots and the step in every PPO epoch, but then
+              # the moment estimates in adaptive optimizers would never have
+              # enough time to warm up. So it makes sense to reuse the slots,
+              # even though we're optimizing a different loss in every new
+              # epoch.
+              self._total_opt_step,
+              self._policy_and_value_opt_state,
+              self._policy_and_value_opt_update,
+              self._policy_and_value_get_params,
+              self._policy_and_value_net_apply,
+              log_probabs_traj,
+              value_predictions_traj,
+              padded_observations,
+              padded_actions,
+              padded_rewards,
+              reward_mask,
+              c1=self._c1,
+              c2=self._c2,
+              gamma=self._gamma,
+              lambda_=self._lambda_,
+              state=self._model_state,
+              rng=k1))
+      self._total_opt_step += 1
 
       # Compute the approx KL for early stopping.
       (log_probab_actions_new, _), self._model_state = (
@@ -478,7 +488,9 @@ def save(self):
         os.path.join(self._output_dir, "model-??????.pkl"))
     params_file = os.path.join(self._output_dir, "model-%06d.pkl" % self._epoch)
     with gfile.GFile(params_file, "wb") as f:
-      pickle.dump((self._policy_and_value_net_params, self._model_state), f)
+      pickle.dump(
+          (self._policy_and_value_net_params, self._model_state,
+           self._total_opt_step), f)
     # Remove the old model files.
     for path in old_model_files:
       gfile.remove(path)

From f082dae25c9e4087f72f3fdf2bb486cd1cb13027 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 22 Aug 2019 15:02:36 -0700
Subject: [PATCH 2307/2720] Add a testing config for ppo on pong.

PiperOrigin-RevId: 264927560
---
 .../trax/rl/configs/atari_regression_test.gin | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)
 create mode 100644 tensor2tensor/trax/rl/configs/atari_regression_test.gin

diff --git a/tensor2tensor/trax/rl/configs/atari_regression_test.gin b/tensor2tensor/trax/rl/configs/atari_regression_test.gin
new file mode 100644
index 000000000..98a2c896d
--- /dev/null
+++ b/tensor2tensor/trax/rl/configs/atari_regression_test.gin
@@ -0,0 +1,30 @@
+import tensor2tensor.trax.models
+import tensor2tensor.trax.rl.trainers
+
+# Parameters for PPO:
+# ==============================================================================
+PPO.n_optimizer_steps = 30
+PPO.target_kl = 0.01
+PPO.boundary = 20
+PPO.max_timestep = 128
+PPO.max_timestep_eval = 20000
+PPO.random_seed = None
+PPO.gamma = 0.99
+PPO.lambda_ = 0.95
+PPO.c1 = 1.0
+PPO.c2 = 0.01
+PPO.eval_every_n = 500
+PPO.done_frac_for_policy_save = 0.9
+PPO.n_evals = 16
+PPO.len_history_for_policy = 4
+PPO.eval_temperatures = (1.0, 0.5)
+PPO.policy_and_value_model = @trax.models.AtariCnn
+
+# Parameters for train_rl:
+# ==============================================================================
+train_rl.env_name = "PongNoFrameskip-v4"
+train_rl.n_epochs = 10000
+train_rl.clip_rewards = True
+train_rl.max_timestep = 10000
+train_rl.rendered_env = True
+train_rl.resize_dims = (105, 80)

From 286d7617b5d75f92be24fa8432aa7785cb42dd92 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 22 Aug 2019 16:05:20 -0700
Subject: [PATCH 2308/2720] Add weight decay to Adam and Adafactor.

PiperOrigin-RevId: 264940583
---
 tensor2tensor/trax/optimizers/base.py | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/trax/optimizers/base.py b/tensor2tensor/trax/optimizers/base.py
index 725b7c3c3..af592d0f5 100644
--- a/tensor2tensor/trax/optimizers/base.py
+++ b/tensor2tensor/trax/optimizers/base.py
@@ -204,11 +204,13 @@ def update(self, step, grads, params, avg_sq_grad, opt_params):
 class Adam(Optimizer):
   """Adam optimizer."""
 
-  def __init__(self, learning_rate, b1=0.9, b2=0.999, eps=1e-8):  # pylint: disable=useless-super-delegation
+  def __init__(self, learning_rate, weight_decay_rate=1e-5,  # pylint: disable=useless-super-delegation
+               b1=0.9, b2=0.999, eps=1e-5):
     """Create the Adam optimizer.
 
     Args:
       learning_rate: a postitive scalar value for the initial learning rate.
+      weight_decay_rate: rate at which to decay weights.
       b1: optional, a positive scalar value for beta_1, the exponential decay
         rate for the first moment estimates (default 0.9).
       b2: optional, a positive scalar value for beta_2, the exponential decay
@@ -216,7 +218,7 @@ def __init__(self, learning_rate, b1=0.9, b2=0.999, eps=1e-8):  # pylint: disabl
       eps: optional, a positive scalar value for epsilon, a small constant for
         numerical stability (default 1e-8).
     """
-    super(Adam, self).__init__(learning_rate, b1, b2, eps)
+    super(Adam, self).__init__(learning_rate, weight_decay_rate, b1, b2, eps)
 
   def init(self, params):
     m = np.zeros_like(params)
@@ -225,12 +227,12 @@ def init(self, params):
 
   def update(self, step, grads, params, slots, opt_params):
     m, v = slots
-    learning_rate, b1, b2, eps = opt_params
+    learning_rate, weight_decay_rate, b1, b2, eps = opt_params
     m = (1 - b1) * grads + b1 * m  # First  moment estimate.
     v = (1 - b2) * (grads ** 2) + b2 * v  # Second moment estimate.
     mhat = m / (1 - b1 ** (step + 1))  # Bias correction.
     vhat = v / (1 - b2 ** (step + 1))
-    params = params - (
+    params = (1 - weight_decay_rate) * params - (
         learning_rate * mhat / (np.sqrt(vhat) + eps)).astype(params.dtype)
     return params, (m, v)
 
@@ -238,7 +240,6 @@ def update(self, step, grads, params, slots, opt_params):
 class Adafactor(Optimizer):
   """Adafactor optimizer."""
 
-  # TODO(levskaya): refactor to use newer RL friendly parameter passing.
   def __init__(self,
                learning_rate,
                factored=True,
@@ -248,6 +249,7 @@ def __init__(self,
                beta1=0.0,
                decay_rate=0.8,
                clipping_threshold=1.0,
+               weight_decay_rate=1e-5,
                epsilon1=1e-30,
                epsilon2=1e-3):
     """Create the Adafactor optimizer.
@@ -267,6 +269,7 @@ def __init__(self,
         memory if nonzero!  Off by default.
       decay_rate: float: controls second-moment exponential decay schedule.
       clipping_threshold: an optional float >= 1, if None no update clipping.
+      weight_decay_rate: rate at which to decay weights.
       epsilon1: Regularization constant for squared gradient.
       epsilon2: Regularization constant for parameter scale.
     """
@@ -278,7 +281,7 @@ def __init__(self,
     # Dynamically configurable parameters will be passed to the update function.
     super(Adafactor, self).__init__(
         learning_rate, beta1, decay_rate, clipping_threshold,
-        epsilon1, epsilon2)
+        weight_decay_rate, epsilon1, epsilon2)
 
   @staticmethod
   def _decay_rate_pow(i, exponent=0.8):
@@ -304,7 +307,7 @@ def init(self, params):
   def update(self, step, grads, params, slots, opt_params):
     updates = []
     (learning_rate, beta1, decay_rate, clipping_threshold,
-     epsilon1, epsilon2) = opt_params
+     weight_decay_rate, epsilon1, epsilon2) = opt_params
     decay_rate = self._decay_rate_pow(step, exponent=decay_rate)
     update_scale = learning_rate
     if self._multiply_by_parameter_scale:
@@ -343,7 +346,7 @@ def update(self, step, grads, params, slots, opt_params):
       subtrahend = new_m
       updates.append(new_m)
 
-    new_params = params - subtrahend
+    new_params = (1 - weight_decay_rate) * params - subtrahend
     return new_params, updates
 
 

From 83f8c47813f0674c7588570c722a04f0e0c279e5 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 22 Aug 2019 17:08:48 -0700
Subject: [PATCH 2309/2720] Implement SimPLe

PiperOrigin-RevId: 264952474
---
 tensor2tensor/envs/env_problem.py             |   4 +
 tensor2tensor/trax/rl/__init__.py             |  35 ++++
 tensor2tensor/trax/rl/base_trainer.py         |   5 +-
 tensor2tensor/trax/rl/ppo_trainer.py          |  18 +-
 tensor2tensor/trax/rl/simple_trainer.py       | 171 ++++++++++++++++++
 tensor2tensor/trax/rl/simple_trainer_test.py  |  94 ++++++++++
 .../trax/rl/simulated_env_problem.py          | 105 +++++++++--
 .../trax/rl/simulated_env_problem_test.py     |   2 +-
 tensor2tensor/trax/rl/trainers.py             |   2 +
 tensor2tensor/trax/trax.py                    |  45 +++--
 tensor2tensor/trax/trax_test.py               |   6 +-
 11 files changed, 441 insertions(+), 46 deletions(-)
 create mode 100644 tensor2tensor/trax/rl/simple_trainer.py
 create mode 100644 tensor2tensor/trax/rl/simple_trainer_test.py

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 909307c0f..0fc27d9ce 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -232,6 +232,10 @@ def is_reward_range_finite(self):
     min_reward, max_reward = self.reward_range
     return (min_reward != -np.inf) and (max_reward != np.inf)
 
+  @property
+  def discrete_rewards(self):
+    return self._discrete_rewards
+
   def process_rewards(self, rewards):
     """Clips the rewards, optionally rounds them and casts to integer.
 
diff --git a/tensor2tensor/trax/rl/__init__.py b/tensor2tensor/trax/rl/__init__.py
index 4872e5d5d..b86f23d3b 100644
--- a/tensor2tensor/trax/rl/__init__.py
+++ b/tensor2tensor/trax/rl/__init__.py
@@ -13,3 +13,38 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Trax RL library."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gin
+
+from tensor2tensor.trax.rl import simulated_env_problem
+
+
+def configure_rl(*args, **kwargs):
+  kwargs["module"] = "trax.rl"
+  return gin.external_configurable(*args, **kwargs)
+
+
+def configure_simulated_env_problem(*args, **kwargs):
+  kwargs["blacklist"] = [
+      "batch_size", "observation_space", "action_space", "reward_range",
+      "discrete_rewards", "history_stream", "output_dir"]
+  return configure_rl(*args, **kwargs)
+
+
+# pylint: disable=invalid-name
+RawSimulatedEnvProblem = configure_simulated_env_problem(
+    simulated_env_problem.RawSimulatedEnvProblem)
+SerializedSequenceSimulatedEnvProblem = configure_simulated_env_problem(
+    simulated_env_problem.SerializedSequenceSimulatedEnvProblem)
+
+
+# pylint: disable=invalid-name
+cartpole_done_fn = configure_rl(simulated_env_problem.cartpole_done_fn)
+cartpole_reward_fn = configure_rl(simulated_env_problem.cartpole_reward_fn)
+acrobot_done_fn = configure_rl(simulated_env_problem.acrobot_done_fn)
+acrobot_reward_fn = configure_rl(simulated_env_problem.acrobot_reward_fn)
diff --git a/tensor2tensor/trax/rl/base_trainer.py b/tensor2tensor/trax/rl/base_trainer.py
index 68d7bc01c..dbdb652e2 100644
--- a/tensor2tensor/trax/rl/base_trainer.py
+++ b/tensor2tensor/trax/rl/base_trainer.py
@@ -27,8 +27,9 @@ class BaseTrainer(object):
   """Base class for RL trainers."""
 
   def __init__(self, train_env, eval_env, output_dir):
-    self._train_env = train_env
-    self._eval_env = eval_env
+    # Train and eval envs are settable.
+    self.train_env = train_env
+    self.eval_env = eval_env
     self._output_dir = output_dir
     gfile.makedirs(self._output_dir)
 
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index 1ff5fc6bd..d84cddb6e 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -129,13 +129,13 @@ def __init__(
     self._len_history_for_policy = len_history_for_policy
     self._eval_temperatures = eval_temperatures
 
-    assert isinstance(self._train_env.action_space, gym.spaces.Discrete)
-    n_actions = self._train_env.action_space.n
+    assert isinstance(self.train_env.action_space, gym.spaces.Discrete)
+    n_actions = self.train_env.action_space.n
 
     # Batch Observations Shape = [1, 1] + OBS, because we will eventually call
     # policy and value networks on shape [B, T] +_OBS
-    batch_observations_shape = (1, 1) + self._train_env.observation_space.shape
-    observations_dtype = self._train_env.observation_space.dtype
+    batch_observations_shape = (1, 1) + self.train_env.observation_space.shape
+    observations_dtype = self.train_env.observation_space.dtype
 
     self._rng = trax.get_random_number_generator_and_set_seed(random_seed)
     self._rng, key1 = jax_random.split(self._rng, num=2)
@@ -202,9 +202,9 @@ def train_epoch(self):
     logging.vlog(1, "Epoch [% 6d] collecting trajectories.", self._epoch)
     self._rng, key = jax_random.split(self._rng)
     trajs, n_done, timing_info, self._model_state = ppo.collect_trajectories(
-        self._train_env,
+        self.train_env,
         policy_fn=self._get_predictions,
-        n_trajectories=self._train_env.batch_size,
+        n_trajectories=self.train_env.batch_size,
         max_timestep=self._max_timestep,
         state=self._model_state,
         rng=key,
@@ -254,7 +254,7 @@ def train_epoch(self):
     assert (B, T) == padded_rewards.shape
     assert (B, T) == reward_mask.shape
     assert (B, T + 1) == padded_observations.shape[:2]
-    assert ((B, T + 1) + self._train_env.observation_space.shape ==
+    assert ((B, T + 1) + self.train_env.observation_space.shape ==
             padded_observations.shape)
 
     log_prob_recompute_start_time = time.time()
@@ -421,7 +421,7 @@ def train_epoch(self):
     self._n_trajectories_done += n_done
     # TODO(afrozm): Refactor to trax.save_state.
     if ((self._n_trajectories_done >=
-         self._done_frac_for_policy_save * self._train_env.batch_size) and
+         self._done_frac_for_policy_save * self.train_env.batch_size) and
         (self._epoch - self._last_saved_at > self._eval_every_n) and
         (((self._epoch + 1) % self._eval_every_n == 0))):
       self.save()
@@ -470,7 +470,7 @@ def evaluate(self):
     logging.vlog(1, "Epoch [% 6d] evaluating policy.", self._epoch)
     self._rng, key = jax_random.split(self._rng, num=2)
     reward_stats, self._model_state = ppo.evaluate_policy(
-        self._eval_env,
+        self.eval_env,
         self._get_predictions,
         temperatures=self._eval_temperatures,
         max_timestep=self._max_timestep_eval,
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
new file mode 100644
index 000000000..07b447188
--- /dev/null
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -0,0 +1,171 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 trainer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import itertools
+import os
+import random
+
+from absl import logging
+from jax import numpy as np
+from tensor2tensor.trax import inputs as trax_inputs
+from tensor2tensor.trax import trax
+from tensor2tensor.trax.rl import base_trainer
+from tensor2tensor.trax.rl import simulated_env_problem
+
+
+class SimPLe(base_trainer.BaseTrainer):
+  """SimPLe trainer."""
+
+  def __init__(
+      self,
+      train_env,
+      eval_env,
+      output_dir,
+      policy_trainer_class,
+      n_real_epochs=10,
+      data_eval_frac=0.05,
+      model_train_batch_size=64,
+      simulated_env_problem_class=(
+          simulated_env_problem.SerializedSequenceSimulatedEnvProblem),
+      simulated_batch_size=16,
+      n_simulated_epochs=1000,
+  ):
+    super(SimPLe, self).__init__(train_env, eval_env, output_dir)
+    self._policy_dir = os.path.join(output_dir, "policy")
+    self._policy_trainer = policy_trainer_class(
+        train_env=train_env,
+        eval_env=eval_env,
+        output_dir=self._policy_dir,
+    )
+    self._n_real_epochs = n_real_epochs
+    self._model_train_batch_size = model_train_batch_size
+    self._data_eval_frac = data_eval_frac
+    self._train_trajectories = []
+    self._eval_trajectories = []
+    self._model_dir = os.path.join(output_dir, "model")
+    self._sim_env = simulated_env_problem_class(
+        batch_size=None,
+        observation_space=train_env.observation_space,
+        action_space=train_env.action_space,
+        reward_range=train_env.reward_range,
+        discrete_rewards=train_env.discrete_rewards,
+        history_stream=None,  # TODO(pkozakowski): Support this.
+        output_dir=self._model_dir,
+    )
+    self._simulated_batch_size = simulated_batch_size
+    self._n_simulated_epochs = n_simulated_epochs
+    self._epoch = 0
+
+  @property
+  def epoch(self):
+    return self._epoch
+
+  def train_epoch(self):
+    self.collect_trajectories()
+    self.train_model()
+    self.train_policy()
+    self._epoch += 1
+
+  def evaluate(self):
+    self._policy_trainer.evaluate()
+
+  def save(self):
+    # Nothing to do, as we save stuff continuously.
+    pass
+
+  def flush_summaries(self):
+    # TODO(pkozakowski): Report some metrics, timing?
+    pass
+
+  def collect_trajectories(self):
+    logging.info("Epoch %d: collecting data", self._epoch)
+
+    self._policy_trainer.train_env = self.train_env
+    self._policy_trainer.training_loop(self._n_real_epochs)
+    self.train_env.trajectories.complete_all_trajectories()
+    trajectories = self.train_env.trajectories.completed_trajectories
+    pivot = int(len(trajectories) * (1 - self._data_eval_frac))
+    self._train_trajectories.extend(trajectories[:pivot])
+    self._eval_trajectories.extend(trajectories[pivot:])
+    # TODO(pkozakowski): Save trajectories to disk. Support restoring.
+
+  def _data_stream(self, trajectories, batch_size):
+    def make_batch(examples):
+      """Stack a structure of np arrays nested in lists/tuples."""
+      assert examples
+      if isinstance(examples[0], (list, tuple)):
+        return type(examples[0])(
+            make_batch([example[i] for example in examples])
+            for i in range(len(examples[0]))
+        )
+      else:
+        batch = np.stack(examples, axis=0)
+        pad_width = (
+            [(0, batch_size - len(examples))] +
+            [(0, 0)] * (len(batch.shape) - 1)
+        )
+        # Pad with zeros. This doesn't change anything, because we have weights
+        # in the examples.
+        return np.pad(batch, pad_width, mode="constant")
+
+    examples = [
+        example  # pylint: disable=g-complex-comprehension
+        for trajectory_examples in map(
+            self._sim_env.trajectory_to_training_examples, trajectories)
+        for example in trajectory_examples
+    ]
+    while True:
+      random.shuffle(examples)
+      for from_index in range(0, len(examples), batch_size):
+        example_list = examples[from_index:(from_index + batch_size)]
+        yield make_batch(example_list)
+
+  def train_model(self):
+    logging.info("Epoch %d: training model", self._epoch)
+
+    train_stream = lambda: self._data_stream(  # pylint: disable=g-long-lambda
+        self._train_trajectories, self._model_train_batch_size)
+    eval_stream = lambda: self._data_stream(  # pylint: disable=g-long-lambda
+        self._eval_trajectories, self._model_train_batch_size)
+    # Ignore n_devices for now.
+    inputs = lambda _: trax_inputs.Inputs(  # pylint: disable=g-long-lambda
+        train_stream=train_stream,
+        train_eval_stream=train_stream,
+        eval_stream=eval_stream,
+        input_shape=self._sim_env.model_input_shape,
+        input_dtype=self._sim_env.model_input_dtype,
+    )
+    trax.train(
+        model=self._sim_env.model,
+        inputs=inputs,
+        output_dir=self._model_dir,
+        has_weights=True,
+    )
+
+  def train_policy(self):
+    logging.info("Epoch %d: training policy", self._epoch)
+
+    self._sim_env.initialize(
+        batch_size=self._simulated_batch_size,
+        history_stream=itertools.repeat(None),
+    )
+    self._policy_trainer.train_env = self._sim_env
+    self._policy_trainer.training_loop(self._n_simulated_epochs)
diff --git a/tensor2tensor/trax/rl/simple_trainer_test.py b/tensor2tensor/trax/rl/simple_trainer_test.py
new file mode 100644
index 000000000..052f27000
--- /dev/null
+++ b/tensor2tensor/trax/rl/simple_trainer_test.py
@@ -0,0 +1,94 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.rl.simple_trainer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+
+import gin
+
+from tensor2tensor.envs import gym_env_problem
+from tensor2tensor.rl import gym_utils
+from tensor2tensor.trax import models
+from tensor2tensor.trax.rl import envs  # pylint: disable=unused-import
+from tensor2tensor.trax.rl import simulated_env_problem
+from tensor2tensor.trax.rl import trainers
+from tensorflow import test
+
+
+class SimpleTrainerTest(test.TestCase):
+
+  def _make_wrapped_env(self, name, max_episode_steps=2):
+    wrapper_fn = functools.partial(
+        gym_utils.gym_env_wrapper,
+        **{
+            "rl_env_max_episode_steps": max_episode_steps,
+            "maxskip_env": False,
+            "rendered_env": False,
+            "rendered_env_resize_to": None,  # Do not resize frames
+            "sticky_actions": False,
+            "output_dtype": None,
+        })
+
+    return gym_env_problem.GymEnvProblem(base_env_name=name,
+                                         batch_size=2,
+                                         env_wrapper_fn=wrapper_fn,
+                                         discrete_rewards=False)
+
+  def test_training_loop_acrobot(self):
+    gin.bind_parameter("BoxSpaceSerializer.precision", 2)
+    gin.bind_parameter("trax.train.train_steps", 1)
+    gin.bind_parameter("trax.train.eval_steps", 1)
+    trainer = trainers.SimPLe(
+        train_env=self._make_wrapped_env("Acrobot-v1"),
+        eval_env=self._make_wrapped_env("Acrobot-v1"),
+        output_dir=self.get_temp_dir(),
+        policy_trainer_class=functools.partial(
+            trainers.PPO,
+            policy_and_value_model=functools.partial(
+                models.FrameStackMLP,
+                n_frames=1,
+                hidden_sizes=(),
+                output_size=1,
+            ),
+            n_optimizer_steps=1,
+        ),
+        n_real_epochs=1,
+        model_train_batch_size=2,
+        simulated_env_problem_class=functools.partial(
+            simulated_env_problem.SerializedSequenceSimulatedEnvProblem,
+            model=functools.partial(
+                models.TransformerLM,
+                d_model=2,
+                n_layers=0,
+                max_len=64,
+            ),
+            reward_fn=simulated_env_problem.acrobot_reward_fn,
+            done_fn=simulated_env_problem.acrobot_done_fn,
+            vocab_size=4,
+            max_trajectory_length=4,
+        ),
+        simulated_batch_size=2,
+        n_simulated_epochs=1,
+    )
+    trainer.training_loop(n_epochs=1)
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index 5de39f353..07703ee9c 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -57,7 +57,8 @@ def __init__(self, model, batch_size, observation_space, action_space,
     """
     # TODO(pkozakowski): At some point we will have a "predict" mode which we
     # should use here. When this happens, change the mode.
-    self._model_predict = backend.jit(model(mode="eval"))
+    self._model = model
+    self._model_predict = backend.jit(self._model(mode="eval"))
     self._observation_space = observation_space
     self._action_space = action_space
     self._reward_range = reward_range
@@ -152,6 +153,17 @@ def _step_model(self, predict_fn, actions, rng):
     """
     raise NotImplementedError
 
+  def trajectory_to_training_examples(self, trajectory):
+    raise NotImplementedError
+
+  @property
+  def model_input_shape(self):
+    raise NotImplementedError
+
+  @property
+  def model_input_dtype(self):
+    raise NotImplementedError
+
   def _reset(self, indices):
     """Resets environments at the given indices.
 
@@ -180,6 +192,10 @@ def _step(self, actions):
         self._predict_fn, actions, subrng)
     return (observation, reward, done, {})
 
+  @property
+  def model(self):
+    return self._model
+
 
 class RawSimulatedEnvProblem(SimulatedEnvProblem):
   """SimulatedEnvProblem running a model operating on raw tensors.
@@ -295,16 +311,23 @@ class SerializedSequenceSimulatedEnvProblem(SimulatedEnvProblem):
   sequences of symbols using an EnvSerializer passed to the constructor.
   """
 
-  def __init__(self, reward_fn, done_fn, vocab_size, max_trajectory_length,
+  def __init__(self, model, reward_fn, done_fn, vocab_size,
+               max_trajectory_length, observation_space, action_space,
                *args, **kwargs):
     """Initializes the env.
 
     Args:
+      model: TRAX model to use for simulation. It's assumed to take keyword
+        arguments vocab_size and mode, where vocab_size is the number of symbols
+        in the vocabulary and mode is either "train" or "eval".
+
       reward_fn: Function (previous_observation, current_observation) -> reward.
       done_fn: Function (previous_observation, current_observation) -> done.
       vocab_size: (int) Number of symbols in the vocabulary.
       max_trajectory_length: (int) Maximum length of a trajectory unrolled from
         the model.
+      observation_space: (gym.Space) Observation space.
+      action_space: (gym.Space) Action space.
       *args: (tuple) Positional arguments passed to the base class.
       **kwargs: (dict) Keyword arguments passed to the base class.
     """
@@ -318,17 +341,27 @@ def __init__(self, reward_fn, done_fn, vocab_size, max_trajectory_length,
     self._action_space = None
     self._last_observations = None
 
-    super(SerializedSequenceSimulatedEnvProblem, self).__init__(*args, **kwargs)
-
-  def initialize_environments(self, batch_size=1, **kwargs):
-    """Initializes the environments."""
     self._obs_serializer = space_serializer.create(
-        self.observation_space, self._vocab_size)
+        observation_space, self._vocab_size)
     self._action_serializer = space_serializer.create(
-        self.action_space, self._vocab_size)
+        action_space, self._vocab_size)
     self._obs_repr_length = self._obs_serializer.representation_length
     self._action_repr_length = self._action_serializer.representation_length
     self._step_repr_length = self._obs_repr_length + self._action_repr_length
+
+    # We assume that the model takes vocab_size as an argument (e.g.
+    # TransformerLM).
+    model = functools.partial(model, vocab_size=vocab_size)
+    super(SerializedSequenceSimulatedEnvProblem, self).__init__(
+        *args,
+        model=model,
+        observation_space=observation_space,
+        action_space=action_space,
+        **kwargs
+    )
+
+  def initialize_environments(self, batch_size=1, **kwargs):
+    """Initializes the environments."""
     self._history = np.zeros((
         batch_size,
         self._max_trajectory_length * self._step_repr_length
@@ -336,9 +369,8 @@ def initialize_environments(self, batch_size=1, **kwargs):
     self._steps = np.zeros(batch_size, dtype=np.int32)
     self._last_observations = np.full(
         (batch_size,) + self._observation_space.shape, np.nan)
-    return super(
-        SerializedSequenceSimulatedEnvProblem, self
-    ).initialize_environments(batch_size=batch_size, **kwargs)
+    super(SerializedSequenceSimulatedEnvProblem, self).initialize_environments(
+        batch_size=batch_size, **kwargs)
 
   @property
   def _obs_repr_indices(self):
@@ -384,9 +416,46 @@ def _step_model(self, predict_fn, actions, rng):
     reward = self._reward_fn(self._last_observations, observation)
     done = self._done_fn(self._last_observations, observation)
     self._last_observations = observation
-    done = np.logical_or(done, self._steps == self._max_trajectory_length)
+    done = np.logical_or(done, self._steps == self._max_trajectory_length - 1)
     return (observation, reward, done)
 
+  def trajectory_to_training_examples(self, trajectory):
+    reprs = []
+    weights = []
+    for time_step in trajectory.time_steps:
+      # Serializers work on batches.
+      obs_repr = self._obs_serializer.serialize(
+          np.array([time_step.observation]))[0]
+      reprs.append(obs_repr)
+      # TODO(pkozakowski): Digit weighting.
+      weights.append(np.ones_like(obs_repr))
+      if time_step.action is not None:
+        action_repr = self._action_serializer.serialize(
+            np.array([time_step.action]))[0]
+        reprs.append(action_repr)
+        weights.append(np.zeros_like(action_repr))
+
+    def concat_and_pad(arrays):
+      (desired_length,) = self.model_input_shape
+      flat_array = np.concatenate(arrays, axis=0)
+      (actual_length,) = flat_array.shape
+      assert actual_length <= desired_length
+      return np.pad(
+          flat_array,
+          pad_width=((0, desired_length - actual_length),),
+          mode="constant",
+      )
+    (reprs, weights) = map(concat_and_pad, (reprs, weights))
+    return [(reprs, reprs, weights)]  # (inputs, targets, weights)
+
+  @property
+  def model_input_shape(self):
+    return (self._max_trajectory_length * self._step_repr_length,)
+
+  @property
+  def model_input_dtype(self):
+    return np.int32
+
 
 def cartpole_done_fn(previous_observation, current_observation):
   del previous_observation
@@ -400,3 +469,15 @@ def cartpole_done_fn(previous_observation, current_observation):
 def cartpole_reward_fn(previous_observation, current_observation):
   done = cartpole_done_fn(previous_observation, current_observation)
   return 1.0 - done  # Unit reward for every timestep until the end.
+
+
+def acrobot_done_fn(previous_observation, current_observation):
+  del previous_observation
+  theta1 = current_observation[:, 0]
+  theta2 = current_observation[:, 1]
+  return -np.cos(theta1) - np.cos(theta2 + theta1) > 1.0
+
+
+def acrobot_reward_fn(previous_observation, current_observation):
+  done = acrobot_done_fn(previous_observation, current_observation)
+  return -1.0 + done  # -1 reward for every timestep until the end.
diff --git a/tensor2tensor/trax/rl/simulated_env_problem_test.py b/tensor2tensor/trax/rl/simulated_env_problem_test.py
index 6f1dcd605..f824bf741 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem_test.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem_test.py
@@ -178,7 +178,7 @@ def make_prediction(symbol):
       np.testing.assert_array_equal(inputs[0, 5:9], symbols[4:8])
       self.assertFalse(np.array_equal(obs2, obs3))
       np.testing.assert_array_equal(reward, [0.5])
-      np.testing.assert_array_equal(done, [False])
+      np.testing.assert_array_equal(done, [True])
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/trax/rl/trainers.py b/tensor2tensor/trax/rl/trainers.py
index 4be7e7318..9368e3e61 100644
--- a/tensor2tensor/trax/rl/trainers.py
+++ b/tensor2tensor/trax/rl/trainers.py
@@ -22,6 +22,7 @@
 import gin
 
 from tensor2tensor.trax.rl import ppo_trainer
+from tensor2tensor.trax.rl import simple_trainer
 
 
 # Ginify
@@ -33,3 +34,4 @@ def trainer_configure(*args, **kwargs):
 
 # pylint: disable=invalid-name
 PPO = trainer_configure(ppo_trainer.PPO)
+SimPLe = trainer_configure(simple_trainer.SimPLe)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 9ecdd62e0..8b970a1a7 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -52,8 +52,7 @@
 from tensorflow.io import gfile
 
 
-@gin.configurable
-def unpack_batch(batch, has_weights=False):
+def unpack_batch(batch, has_weights):
   """Unpacks a training batch into inputs, targets and weights."""
   if has_weights:
     assert len(batch) == 3  # (inputs, targets, weights)
@@ -101,9 +100,9 @@ def masked_mean(inputs, targets, weights, mask_id=None):
               for x, w, s in zip(inputs, weights, weight_sums)])
 
 
-def accuracy(batch, model_predictions):
+def accuracy(batch, model_predictions, has_weights):
   """Calculate accuracy."""
-  _, targets, weights = unpack_batch(batch)
+  _, targets, weights = unpack_batch(batch, has_weights)
   model_predictions, targets, weights = _make_list(
       model_predictions, targets, weights)
   correct = []
@@ -113,9 +112,9 @@ def accuracy(batch, model_predictions):
   return masked_mean(correct, targets, weights)
 
 
-def neg_log_perplexity(batch, model_predictions):
+def neg_log_perplexity(batch, model_predictions, has_weights):
   """Calculate negative log perplexity."""
-  _, targets, weights = unpack_batch(batch)
+  _, targets, weights = unpack_batch(batch, has_weights)
   model_predictions, targets, weights = _make_list(
       model_predictions, targets, weights)
   xent = []
@@ -125,9 +124,9 @@ def neg_log_perplexity(batch, model_predictions):
   return masked_mean(xent, targets, weights)
 
 
-def loss(params, batch, model_predict, state, rng):
+def loss(params, batch, model_predict, state, rng, has_weights):
   """Calculate loss."""
-  inputs, targets, weights = unpack_batch(batch)
+  inputs, targets, weights = unpack_batch(batch, has_weights)
   predictions, state = model_predict(inputs, params, state, rng=rng)
   predictions, targets, weights = _make_list(predictions, targets, weights)
   xent = []
@@ -236,11 +235,12 @@ def _print_n_params(opt_state, n_devices, step):
 _METRICS = {
     "accuracy": accuracy,
     "neg_log_perplexity": neg_log_perplexity,
-    "loss": lambda x, y: - neg_log_perplexity(x, y),
+    "loss": lambda *args, **kwargs: - neg_log_perplexity(*args, **kwargs),
 }
 
 
 def evaluate_train_and_eval(step, inputs, predict_fn, eval_steps, state, rng,
+                            has_weights,
                             train_sw=None, eval_sw=None, history=None):
   """Evalaute on train and eval data, and log metrics."""
   step_log(step, "Evaluation")
@@ -250,7 +250,9 @@ def evaluate_train_and_eval(step, inputs, predict_fn, eval_steps, state, rng,
         itertools.islice(input_stream(), eval_steps),
         predict_fn,
         _METRICS,
-        state, rng)
+        state,
+        rng,
+        has_weights)
     metrics_list.append(metrics)
   train_metrics, eval_metrics = metrics_list  # pylint: disable=unbalanced-tuple-unpacking
   if train_sw:
@@ -261,7 +263,7 @@ def evaluate_train_and_eval(step, inputs, predict_fn, eval_steps, state, rng,
   return train_metrics, eval_metrics, state
 
 
-def evaluate(inputs_stream, predict_fn, metric_fns, state, rng):
+def evaluate(inputs_stream, predict_fn, metric_fns, state, rng, has_weights):
   """Evaluate.
 
   Args:
@@ -272,6 +274,7 @@ def evaluate(inputs_stream, predict_fn, metric_fns, state, rng):
       and predictions and returns a scalar metric value.
     state: start state for `predict_fn`.
     rng: random number generator.
+    has_weights: bool, whether weights are included in the inputs.
 
   Returns:
     metrics: dict from metric name to metric value averaged over the number of
@@ -285,12 +288,12 @@ def evaluate(inputs_stream, predict_fn, metric_fns, state, rng):
     rng, subrng = jax_random.split(rng)
     preds, state = predict_fn(inp[0], state=state, rng=subrng)
     for m, f in six.iteritems(metric_fns):
-      metrics[m] += f(inp, preds)
+      metrics[m] += f(inp, preds, has_weights=has_weights)
   return {m: v / count for (m, v) in six.iteritems(metrics)}, state
 
 
 def evaluate_loss_train_and_eval(step, inputs, compute_loss_fn, eval_steps,
-                                 state, rngs,
+                                 state, rngs, has_weights,
                                  train_sw=None, eval_sw=None, history=None):
   """More efficient evaluation that logs only the loss on train & eval data."""
   step_log(step, "Evaluation")
@@ -299,7 +302,7 @@ def evaluate_loss_train_and_eval(step, inputs, compute_loss_fn, eval_steps,
     total = 0.0
     count = 0.0
     for inp in itertools.islice(input_stream(), eval_steps):
-      loss_values, state, rngs = compute_loss_fn(inp, state, rngs)
+      loss_values, state, rngs = compute_loss_fn(inp, state, rngs, has_weights)
       total += float(numpy.mean(loss_values))
       count += 1.0
     metrics = {"loss": total / count}
@@ -506,11 +509,13 @@ class Trainer(object):
 
   def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs,
                output_dir=None, random_seed=None, n_devices=None,
-               save_steps=None, should_save=True):
+               save_steps=None, should_save=True, has_weights=False):
     if save_steps is None:
       save_steps = []
     self._save_steps = save_steps
     self._should_save = should_save
+    self._has_weights = has_weights
+    loss_fn = functools.partial(loss_fn, has_weights=self._has_weights)
     device_count = jax.lib.xla_bridge.device_count()
     n_devices = n_devices or device_count
     # TODO(lukaszkaiser): remove this restriction when possible.
@@ -732,7 +737,8 @@ def evaluate(self, eval_steps):
         rng=rng,
         train_sw=self._train_sw,
         eval_sw=self._eval_sw,
-        history=self._history)
+        history=self._history,
+        has_weights=self._has_weights)
 
   def update_learning_rate(self):
     self._lr_fn = self._lr_schedule(self._history)
@@ -789,6 +795,7 @@ def evaluate(self, eval_steps):
         eval_steps=eval_steps,
         state=self._model_state,
         rngs=self._rngs,
+        has_weights=self._has_weights,
         train_sw=self._train_sw,
         eval_sw=self._eval_sw,
         history=self._history)
@@ -815,7 +822,8 @@ def train(output_dir,
           n_devices=None,
           random_seed=None,
           save_graphs=True,
-          save_backward_graph=False):
+          save_backward_graph=False,
+          has_weights=False):
   """Train the model on the inputs.
 
   Args:
@@ -839,13 +847,14 @@ def train(output_dir,
     random_seed: the random seed to use; time/os dependent if None (default).
     save_graphs: bool, if True, save computation graph to file.
     save_backward_graph: bool, if True, save backward graph to file too.
+    has_weights: bool, whether weights are included in the inputs.
   Returns:
     trax.State
   """
   trainer = trainer_class(model, loss_fn, optimizer, lr_schedule, inputs,
                           output_dir,
                           random_seed=random_seed, n_devices=n_devices,
-                          save_steps=save_steps)
+                          save_steps=save_steps, has_weights=has_weights)
 
   epoch_steps = [train_steps]  # Only training if eval_frequency is 0 or None
   if eval_frequency and eval_steps > 0:
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index bd07b0ba9..60f83ab90 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -24,7 +24,6 @@
 import tempfile
 from absl.testing import parameterized
 
-import gin
 from jax import test_util  # pylint: disable=unused-import
 from jax.config import config
 from jax.lib import xla_bridge
@@ -197,8 +196,6 @@ def test_train_with_weights(self, backend_name):
     if xla_bridge.device_count() > 1 and backend_name == "tf":
       self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
     with backend.use_backend(backend_name), self.tmp_dir() as output_dir:
-      gin.bind_parameter("unpack_batch.has_weights", True)
-
       # Prepare model and inputs
       n_classes = 4
       train_steps = 2
@@ -213,7 +210,8 @@ def test_train_with_weights(self, backend_name):
                          model=model_fn,
                          inputs=inputs,
                          train_steps=train_steps,
-                         eval_steps=eval_steps)
+                         eval_steps=eval_steps,
+                         has_weights=True)
 
       # Assert total train steps
       self.assertEqual(state.step, train_steps)

From fe484ec9a1f2f1650b6ee63b0b1bcc38834ccf98 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 22 Aug 2019 17:46:56 -0700
Subject: [PATCH 2310/2720] Enable SYMBOL_WEIGHTS_ALL modality for Scheduled
 Sampling

PiperOrigin-RevId: 264958707
---
 tensor2tensor/utils/t2t_model.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 402116cdf..e7b361ca3 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1821,10 +1821,11 @@ def maybe_scheduled_sampling(self, features, logits, losses):
 
     # Only do scheduled sampling on language tasks.
     modality = problem_hparams.modality["targets"]
-    if modality != modalities.ModalityType.SYMBOL:
+    if modality not in [modalities.ModalityType.SYMBOL,
+                        modalities.ModalityType.SYMBOL_WEIGHTS_ALL]:
       assert hparams.scheduled_sampling_prob == 0, (
-          "Scheduled sampling only applies to ModalityType.SYMBOL. Set "
-          "hparams.scheduled_sampling_prob == 0.0.")
+          "Scheduled sampling only applies to ModalityType.{SYMBOL, "
+          "SYMBOL_WEIGHTS_ALL}. Set hparams.scheduled_sampling_prob == 0.0.")
       return (logits, losses)
 
     # Only do scheduled sampling when training.

From 96702413f4a2246dd11f1154cddfc04ea029476c Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Thu, 22 Aug 2019 18:49:11 -0700
Subject: [PATCH 2311/2720] Move Bayesian Layers from tensorflow/tensor2tensor
 to google/edward2.

Separate commits will clean up the namespace and add documentation to Edward2.

PiperOrigin-RevId: 264967108
---
 oss_scripts/oss_tests.sh                      |    9 -
 tensor2tensor/keras/__init__.py               |   16 -
 tensor2tensor/keras/constraints.py            |   78 -
 tensor2tensor/keras/initializers.py           |  373 -----
 tensor2tensor/keras/initializers_test.py      |   80 -
 tensor2tensor/keras/regularizers.py           |  154 --
 tensor2tensor/keras/regularizers_test.py      |   46 -
 tensor2tensor/layers/bayes.py                 | 1439 -----------------
 tensor2tensor/layers/bayes_test.py            |  626 -------
 tensor2tensor/layers/gaussian_process.py      |  815 ----------
 tensor2tensor/layers/gaussian_process_test.py |  292 ----
 tensor2tensor/layers/reversible_layers.py     | 1139 -------------
 .../layers/reversible_layers_test.py          |  590 -------
 13 files changed, 5657 deletions(-)
 delete mode 100644 tensor2tensor/keras/__init__.py
 delete mode 100644 tensor2tensor/keras/constraints.py
 delete mode 100644 tensor2tensor/keras/initializers.py
 delete mode 100644 tensor2tensor/keras/initializers_test.py
 delete mode 100644 tensor2tensor/keras/regularizers.py
 delete mode 100644 tensor2tensor/keras/regularizers_test.py
 delete mode 100644 tensor2tensor/layers/bayes.py
 delete mode 100644 tensor2tensor/layers/bayes_test.py
 delete mode 100644 tensor2tensor/layers/gaussian_process.py
 delete mode 100644 tensor2tensor/layers/gaussian_process_test.py
 delete mode 100644 tensor2tensor/layers/reversible_layers.py
 delete mode 100644 tensor2tensor/layers/reversible_layers_test.py

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 62343e253..889236f08 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -61,16 +61,13 @@ set_status
 
 
 pytest --disable-warnings \
-  --ignore=tensor2tensor/layers/bayes_test.py \
   --ignore=tensor2tensor/layers/common_attention_test.py \
   --ignore=tensor2tensor/layers/common_layers_test.py \
   --ignore=tensor2tensor/layers/common_video_test.py \
   --ignore=tensor2tensor/layers/discretization_test.py \
-  --ignore=tensor2tensor/layers/gaussian_process_test.py \
   --ignore=tensor2tensor/layers/latent_layers_test.py \
   --ignore=tensor2tensor/layers/modalities_test.py \
   --ignore=tensor2tensor/layers/ngram_test.py \
-  --ignore=tensor2tensor/layers/reversible_layers_test.py \
   tensor2tensor/layers/
 set_status
 
@@ -110,16 +107,11 @@ pytest --disable-warnings \
   tensor2tensor/layers/latent_layers_test.py \
   tensor2tensor/layers/modalities_test.py \
   tensor2tensor/layers/ngram_test.py \
-  tensor2tensor/keras \
   tensor2tensor/utils/t2t_model_test.py \
   tensor2tensor/utils/test_utils_test.py \
   --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary
 set_status
 
-# TODO(afrozm): Re-enable once TFP's new release comes out (0.8.0) or modify
-# stuff so that we test against tfp-nightly.
-#pytest --disable-warnings tensor2tensor/layers/reversible_layers_test.py
-#set_status
 
 pytest --disable-warnings tensor2tensor/utils/registry_test.py
 set_status
@@ -143,7 +135,6 @@ pytest --disable-warnings \
   --ignore=tensor2tensor/bin/t2t_trainer_test.py \
   --ignore=tensor2tensor/data_generators \
   --ignore=tensor2tensor/envs \
-  --ignore=tensor2tensor/keras \
   --ignore=tensor2tensor/layers \
   --ignore=tensor2tensor/models \
   --ignore=tensor2tensor/rl \
diff --git a/tensor2tensor/keras/__init__.py b/tensor2tensor/keras/__init__.py
deleted file mode 100644
index b775a72bd..000000000
--- a/tensor2tensor/keras/__init__.py
+++ /dev/null
@@ -1,16 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-
diff --git a/tensor2tensor/keras/constraints.py b/tensor2tensor/keras/constraints.py
deleted file mode 100644
index 95283c694..000000000
--- a/tensor2tensor/keras/constraints.py
+++ /dev/null
@@ -1,78 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Constraints."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import six
-import tensorflow as tf
-
-
-class Positive(tf.keras.constraints.Constraint):
-  """Positive constraint."""
-
-  def __init__(self, epsilon=tf.keras.backend.epsilon()):
-    self.epsilon = epsilon
-
-  def __call__(self, w):
-    return tf.maximum(w, self.epsilon)
-
-  def get_config(self):
-    return {'epsilon': self.epsilon}
-
-
-# Compatibility aliases, following tf.keras
-
-
-positive = Positive  # pylint: disable=invalid-name
-
-# Utility functions, following tf.keras
-
-
-def serialize(initializer):
-  return tf.keras.utils.serialize_keras_object(initializer)
-
-
-def deserialize(config, custom_objects=None):
-  return tf.keras.utils.deserialize_keras_object(
-      config,
-      module_objects=globals(),
-      custom_objects=custom_objects,
-      printable_module_name='constraints')
-
-
-def get(identifier, value=None):
-  """Getter for loading from strings; returns value if can't load."""
-  if value is None:
-    value = identifier
-  if identifier is None:
-    return None
-  elif isinstance(identifier, dict):
-    try:
-      return deserialize(identifier)
-    except ValueError:
-      return value
-  elif isinstance(identifier, six.string_types):
-    config = {'class_name': str(identifier), 'config': {}}
-    try:
-      return deserialize(config)
-    except ValueError:
-      return value
-  elif callable(identifier):
-    return identifier
-  return value
diff --git a/tensor2tensor/keras/initializers.py b/tensor2tensor/keras/initializers.py
deleted file mode 100644
index 07ba0e822..000000000
--- a/tensor2tensor/keras/initializers.py
+++ /dev/null
@@ -1,373 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Initializers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import math
-import six
-
-from tensor2tensor.keras import constraints
-from tensor2tensor.keras import regularizers
-import tensorflow as tf
-from tensorflow_probability import edward2 as ed
-
-
-# From `tensorflow/python/ops/init_ops.py`
-def _compute_fans(shape):
-  """Computes the number of input and output units for a weight shape.
-
-  Args:
-    shape: Integer shape tuple or TF tensor shape.
-
-  Returns:
-    A tuple of scalars (fan_in, fan_out).
-  """
-  if len(shape) < 1:  # Just to avoid errors for constants.
-    fan_in = fan_out = 1
-  elif len(shape) == 1:
-    fan_in = fan_out = shape[0]
-  elif len(shape) == 2:
-    fan_in = shape[0]
-    fan_out = shape[1]
-  else:
-    # Assuming convolution kernels (2D, 3D, or more).
-    # kernel shape: (..., input_depth, depth)
-    receptive_field_size = 1.
-    for dim in shape[:-2]:
-      receptive_field_size *= dim
-    fan_in = shape[-2] * receptive_field_size
-    fan_out = shape[-1] * receptive_field_size
-  if isinstance(fan_in, tf.Dimension):
-    fan_in = fan_in.value
-  if isinstance(fan_out, tf.Dimension):
-    fan_out = fan_out.value
-  return fan_in, fan_out
-
-
-class ScaledNormalStdDev(tf.keras.initializers.VarianceScaling):
-  """Initializer capable of adapting its scale to the shape of weights tensors.
-
-  This initializes the standard deviation parameter of a Trainable Normal
-  distribution with a scale based on the shape of the weights tensor.
-  Additionally, A small amount of noise will be added to break weigh symmetry.
-
-  With `distribution="truncated_normal" or "untruncated_normal"`, the standard
-  deviation (after truncation, if used) is `stddev = sqrt(scale / n)`, where n
-  is:
-    - number of input units in the weight tensor, if mode = "fan_in"
-    - number of output units, if mode = "fan_out"
-    - average of the numbers of input and output units, if mode = "fan_avg"
-
-  Args:
-    scale: Scaling factor (positive float).
-    mode: One of "fan_in", "fan_out", "fan_avg".
-    distribution: Random distribution to use. One of "truncated_normal", or
-      "untruncated_normal".
-    seed: A Python integer. Used to create random seeds. See
-      `tf.set_random_seed`
-      for behavior.
-    dtype: The data type. Only floating point types are supported.
-
-  Raises:
-    ValueError: In case of an invalid value for the "scale", mode" or
-      "distribution" arguments.
-  """
-
-  def __init__(self,
-               scale=1.0,
-               mode='fan_in',
-               distribution='untruncated_normal',
-               seed=None,
-               dtype=tf.float32):
-    distribution = distribution.lower()
-    if distribution not in {'truncated_normal', 'untruncated_normal'}:
-      raise ValueError('Invalid `distribution` argument:', distribution)
-    super(ScaledNormalStdDev, self).__init__(scale=scale, mode=mode,
-                                             distribution=distribution,
-                                             seed=seed, dtype=dtype)
-
-  def __call__(self, shape, dtype=None, partition_info=None):
-    if dtype is None:
-      dtype = self.dtype
-    scale = self.scale
-    scale_shape = shape
-    if partition_info is not None:
-      scale_shape = partition_info.full_shape
-    fan_in, fan_out = _compute_fans(scale_shape)
-    if self.mode == 'fan_in':
-      scale /= max(1., fan_in)
-    elif self.mode == 'fan_out':
-      scale /= max(1., fan_out)
-    else:
-      scale /= max(1., (fan_in + fan_out) / 2.)
-    if self.distribution == 'truncated_normal':
-      # constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.)
-      stddev = math.sqrt(scale) / .87962566103423978
-    else:  # self.distribution == 'untruncated_normal':
-      stddev = math.sqrt(scale)
-    return tf.random.truncated_normal(shape, mean=stddev, stddev=stddev*0.1,
-                                      dtype=dtype)
-
-
-class TrainableHalfCauchy(tf.keras.layers.Layer):
-  """Half-Cauchy distribution initializer with trainable parameters."""
-
-  def __init__(self,
-               loc_initializer=tf.keras.initializers.truncated_normal(
-                   stddev=1e-5),
-               scale_initializer=tf.keras.initializers.truncated_normal(
-                   mean=1., stddev=1e-5),
-               loc_regularizer=None,
-               scale_regularizer=None,
-               loc_constraint=None,
-               scale_constraint='positive',
-               seed=None,
-               dtype=tf.float32,
-               **kwargs):
-    """Constructs the initializer."""
-    super(TrainableHalfCauchy, self).__init__(dtype=dtype, **kwargs)
-    self.loc_initializer = get(loc_initializer)
-    self.scale_initializer = get(scale_initializer)
-    self.loc_regularizer = regularizers.get(loc_regularizer)
-    self.scale_regularizer = regularizers.get(scale_regularizer)
-    self.loc_constraint = constraints.get(loc_constraint)
-    self.scale_constraint = constraints.get(scale_constraint)
-    self.seed = seed
-
-  def build(self, shape, dtype=None):
-    if dtype is None:
-      dtype = self.dtype
-
-    self.loc = self.add_weight(
-        'loc',
-        shape=shape,
-        initializer=self.loc_initializer,
-        regularizer=self.loc_regularizer,
-        constraint=self.loc_constraint,
-        dtype=dtype,
-        trainable=True)
-    self.scale = self.add_weight(
-        'scale',
-        shape=shape,
-        initializer=self.scale_initializer,
-        regularizer=self.scale_regularizer,
-        constraint=self.scale_constraint,
-        dtype=dtype,
-        trainable=True)
-    self.built = True
-
-  def __call__(self, shape, dtype=None, partition_info=None):
-    del partition_info  # unused arg
-    if not self.built:
-      self.build(shape, dtype)
-    return ed.Independent(
-        ed.HalfCauchy(loc=self.loc, scale=self.scale).distribution,
-        reinterpreted_batch_ndims=len(shape))
-
-  def get_config(self):
-    return {
-        'loc_initializer':
-            tf.keras.initializers.serialize(self.loc_initializer),
-        'scale_initializer':
-            tf.keras.initializers.serialize(self.scale_initializer),
-        'loc_regularizer':
-            tf.keras.regularizers.serialize(self.loc_regularizer),
-        'scale_regularizer':
-            tf.keras.regularizers.serialize(self.scale_regularizer),
-        'loc_constraint':
-            tf.keras.constraints.serialize(self.loc_constraint),
-        'scale_constraint':
-            tf.keras.constraints.serialize(self.scale_constraint),
-        'seed': self.seed,
-        'dtype': self.dtype,
-    }
-
-
-class TrainableNormal(tf.keras.layers.Layer):
-  """Random normal op as an initializer with trainable mean and stddev."""
-
-  def __init__(self,
-               mean_initializer=tf.keras.initializers.truncated_normal(
-                   stddev=1e-5),
-               stddev_initializer='scaled_normal_std_dev',
-               mean_regularizer=None,
-               stddev_regularizer=None,
-               mean_constraint=None,
-               stddev_constraint='positive',
-               seed=None,
-               dtype=tf.float32,
-               **kwargs):
-    """Constructs the initializer."""
-    super(TrainableNormal, self).__init__(dtype=dtype, **kwargs)
-    self.mean_initializer = get(mean_initializer)
-    self.stddev_initializer = get(stddev_initializer)
-    self.mean_regularizer = regularizers.get(mean_regularizer)
-    self.stddev_regularizer = regularizers.get(stddev_regularizer)
-    self.mean_constraint = constraints.get(mean_constraint)
-    self.stddev_constraint = constraints.get(stddev_constraint)
-    self.seed = seed
-
-  def build(self, shape, dtype=None):
-    if dtype is None:
-      dtype = self.dtype
-
-    self.mean = self.add_weight(
-        'mean',
-        shape=shape,
-        initializer=self.mean_initializer,
-        regularizer=self.mean_regularizer,
-        constraint=self.mean_constraint,
-        dtype=dtype,
-        trainable=True)
-    self.stddev = self.add_weight(
-        'stddev',
-        shape=shape,
-        initializer=self.stddev_initializer,
-        regularizer=self.stddev_regularizer,
-        constraint=self.stddev_constraint,
-        dtype=dtype,
-        trainable=True)
-    self.built = True
-
-  def __call__(self, shape, dtype=None, partition_info=None):
-    del partition_info  # unused arg
-    if not self.built:
-      self.build(shape, dtype)
-    return ed.Independent(
-        ed.Normal(loc=self.mean, scale=self.stddev).distribution,
-        reinterpreted_batch_ndims=len(shape))
-
-  def get_config(self):
-    return {
-        'mean_initializer':
-            tf.keras.initializers.serialize(self.mean_initializer),
-        'stddev_initializer':
-            tf.keras.initializers.serialize(self.stddev_initializer),
-        'mean_regularizer':
-            tf.keras.regularizers.serialize(self.mean_regularizer),
-        'stddev_regularizer':
-            tf.keras.regularizers.serialize(self.stddev_regularizer),
-        'mean_constraint':
-            tf.keras.constraints.serialize(self.mean_constraint),
-        'stddev_constraint':
-            tf.keras.constraints.serialize(self.stddev_constraint),
-        'seed': self.seed,
-        'dtype': self.dtype,
-    }
-
-
-class TrainableHeNormal(TrainableNormal):
-  """Trainable normal initialized per He et al. 2015, given a ReLU nonlinearity.
-
-  The distribution is initialized to a Normal scaled by `sqrt(2 / fan_in)`,
-  where `fan_in` is the number of input units. A ReLU nonlinearity is assumed
-  for this initialization scheme.
-
-  References:
-    He K, Zhang X, Ren S, Sun J. Delving deep into rectifiers: Surpassing
-    human-level performance on imagenet classification. In Proceedings of the
-    IEEE international conference on computer vision 2015 (pp. 1026-1034).
-    https://arxiv.org/abs/1502.01852
-  """
-
-  def __init__(self, seed=None, dtype=tf.float32):
-    super(TrainableHeNormal, self).__init__(
-        stddev_initializer=ScaledNormalStdDev(scale=2.0, seed=seed,
-                                              dtype=dtype),
-        seed=seed, dtype=dtype)
-
-  def get_config(self):
-    return {
-        'seed': self.seed,
-        'dtype': self.dtype,
-    }
-
-
-class TrainableGlorotNormal(TrainableNormal):
-  """Trainable normal initialized per Glorot and Bengio, 2010.
-
-  The distribution is initialized to a Normal scaled by `sqrt(2 / fan_in +
-  fan_out)`, where `fan_in` is the number of input units and `fan_out` is the
-  number of output units.
-
-  References:
-    Glorot X, Bengio Y. Understanding the difficulty of training deep
-    feedforward neural networks. In Proceedings of the thirteenth international
-    conference on artificial intelligence and statistics 2010 Mar 31 (pp.
-    249-256). http://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf
-  """
-
-  def __init__(self, seed=None, dtype=tf.float32):
-    super(TrainableGlorotNormal, self).__init__(
-        stddev_initializer=ScaledNormalStdDev(mode='fan_avg', seed=seed,
-                                              dtype=dtype),
-        seed=seed, dtype=dtype)
-
-  def get_config(self):
-    return {
-        'seed': self.seed,
-        'dtype': self.dtype
-    }
-
-
-# Compatibility aliases, following tf.keras
-
-# pylint: disable=invalid-name
-scaled_normal_std_dev = ScaledNormalStdDev
-trainable_half_cauchy = TrainableHalfCauchy
-trainable_normal = TrainableNormal
-trainable_he_normal = TrainableHeNormal
-trainable_glorot_normal = TrainableGlorotNormal
-# pylint: enable=invalid-name
-
-# Utility functions, following tf.keras
-
-
-def serialize(initializer):
-  return tf.keras.utils.serialize_keras_object(initializer)
-
-
-def deserialize(config, custom_objects=None):
-  return tf.keras.utils.deserialize_keras_object(
-      config,
-      module_objects=globals(),
-      custom_objects=custom_objects,
-      printable_module_name='initializers')
-
-
-def get(identifier, value=None):
-  """Getter for loading from strings; returns value if can't load."""
-  if value is None:
-    value = identifier
-  if identifier is None:
-    return None
-  elif isinstance(identifier, dict):
-    try:
-      return deserialize(identifier)
-    except ValueError:
-      return value
-  elif isinstance(identifier, six.string_types):
-    config = {'class_name': str(identifier), 'config': {}}
-    try:
-      return deserialize(config)
-    except ValueError:
-      return value
-  elif callable(identifier):
-    return identifier
-  return value
diff --git a/tensor2tensor/keras/initializers_test.py b/tensor2tensor/keras/initializers_test.py
deleted file mode 100644
index b65f9aae3..000000000
--- a/tensor2tensor/keras/initializers_test.py
+++ /dev/null
@@ -1,80 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 Keras-style initializers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import math
-
-import numpy as np
-from tensor2tensor.keras import initializers
-from tensor2tensor.utils import test_utils
-
-import tensorflow as tf
-tf.compat.v1.enable_eager_execution()
-
-
-class InitializersTest(tf.test.TestCase):
-
-  @test_utils.run_in_graph_and_eager_modes
-  def testTrainableHalfCauchy(self):
-    shape = (3,)
-    initializer = initializers.get('trainable_half_cauchy')
-    half_cauchy = initializer(shape)
-    self.evaluate(tf.global_variables_initializer())
-    loc_value, scale_value = self.evaluate([
-        # Get distribution of rv -> get distribution of Independent.
-        half_cauchy.distribution.distribution.loc,
-        half_cauchy.distribution.distribution.scale])
-    self.assertAllClose(loc_value, np.zeros(shape), atol=1e-4)
-    self.assertAllClose(scale_value, np.ones(shape), atol=1e-4)
-
-    half_cauchy_value = self.evaluate(half_cauchy)
-    self.assertAllEqual(half_cauchy_value.shape, shape)
-    self.assertAllGreaterEqual(half_cauchy_value, 0.)
-
-  @test_utils.run_in_graph_and_eager_modes
-  def testTrainableNormal(self):
-    shape = (100,)
-    # TrainableNormal is expected to have var 1/shape[0]
-    # because it by default has the fan_in mode scale normal std initializer.
-    initializer = initializers.get('trainable_normal')
-    normal = initializer(shape)
-    self.evaluate(tf.global_variables_initializer())
-    loc_value, scale_value = self.evaluate([
-        # Get distribution of rv -> get distribution of Independent.
-        normal.distribution.distribution.loc,
-        normal.distribution.distribution.scale])
-    fan_in = shape[0]
-    target_scale = 1.
-    target_scale /= max(1., fan_in)
-    target_scale = math.sqrt(target_scale)
-
-    self.assertAllClose(loc_value, np.zeros(shape), atol=1e-4)
-    # Tolerance is larger because of the scale normal std initializer.
-    # In this case it has std around 0.01 (0.1*target_scale).
-    self.assertAllClose(
-        scale_value, target_scale * np.ones(shape), atol=5e-2)
-
-    # Test the TrainableNormal initializer has the specified shape.
-    normal_value = self.evaluate(normal)
-    self.assertAllEqual(normal_value.shape, shape)
-
-
-if __name__ == '__main__':
-  tf.test.main()
diff --git a/tensor2tensor/keras/regularizers.py b/tensor2tensor/keras/regularizers.py
deleted file mode 100644
index 5247d32b8..000000000
--- a/tensor2tensor/keras/regularizers.py
+++ /dev/null
@@ -1,154 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Regularizers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import six
-import tensorflow as tf
-
-from tensorflow_probability import edward2 as ed
-
-
-class HalfCauchyKLDivergence(tf.keras.regularizers.Regularizer):
-  """KL divergence regularizer from an input to the half-Cauchy distribution."""
-
-  def __init__(self, loc=0., scale=1.):
-    """Constructs regularizer where default uses the standard half-Cauchy."""
-    self.loc = loc
-    self.scale = scale
-
-  def __call__(self, x):
-    """Computes regularization using an unbiased Monte Carlo estimate."""
-    prior = ed.Independent(
-        ed.HalfCauchy(
-            loc=tf.broadcast_to(self.loc, x.distribution.event_shape),
-            scale=tf.broadcast_to(self.scale, x.distribution.event_shape)
-        ).distribution,
-        reinterpreted_batch_ndims=len(x.distribution.event_shape))
-    negative_entropy = x.distribution.log_prob(x)
-    cross_entropy = -prior.distribution.log_prob(x)
-    return negative_entropy + cross_entropy
-
-  def get_config(self):
-    return {
-        'loc': self.loc,
-        'scale': self.scale,
-    }
-
-
-class LogUniformKLDivergence(tf.keras.regularizers.Regularizer):
-  """KL divergence regularizer from an input to the log-uniform distribution."""
-
-  def __call__(self, x):
-    """Computes regularization given an ed.Normal random variable as input."""
-    if not isinstance(x, ed.RandomVariable):
-      raise ValueError('Input must be an ed.RandomVariable (for correct math, '
-                       'an ed.Normal random variable).')
-    # Clip magnitude of dropout rate, where we get the dropout rate alpha from
-    # the additive parameterization (Molchanov et al., 2017): for weight ~
-    # Normal(mu, sigma**2), the variance `sigma**2 = alpha * mu**2`.
-    mean = x.distribution.mean()
-    log_variance = tf.log(x.distribution.variance())
-    log_alpha = log_variance - tf.log(tf.square(mean) +
-                                      tf.keras.backend.epsilon())
-    log_alpha = tf.clip_by_value(log_alpha, -8., 8.)
-
-    # Set magic numbers for cubic polynomial approx. (Molchanov et al., 2017).
-    k1 = 0.63576
-    k2 = 1.8732
-    k3 = 1.48695
-    c = -k1
-    output = tf.reduce_sum(k1 * tf.nn.sigmoid(k2 + k3 * log_alpha) +
-                           -0.5 * tf.log1p(tf.exp(-log_alpha)) + c)
-    return output
-
-  def get_config(self):
-    return {}
-
-
-class NormalKLDivergence(tf.keras.regularizers.Regularizer):
-  """KL divergence regularizer from an input to the normal distribution."""
-
-  def __init__(self, mean=0., stddev=1.):
-    """Constructs regularizer where default is a KL towards the std normal."""
-    self.mean = mean
-    self.stddev = stddev
-
-  def __call__(self, x):
-    """Computes regularization given an ed.Normal random variable as input."""
-    if not isinstance(x, ed.RandomVariable):
-      raise ValueError('Input must be an ed.RandomVariable.')
-    random_variable = ed.Independent(
-        ed.Normal(
-            loc=tf.broadcast_to(self.mean, x.distribution.event_shape),
-            scale=tf.broadcast_to(self.stddev, x.distribution.event_shape)
-        ).distribution,
-        reinterpreted_batch_ndims=len(x.distribution.event_shape))
-    return random_variable.distribution.kl_divergence(x.distribution)
-
-  def get_config(self):
-    return {
-        'mean': self.mean,
-        'stddev': self.stddev,
-    }
-
-
-# Compatibility aliases, following tf.keras
-
-# pylint: disable=invalid-name
-half_cauchy_kl_divergence = HalfCauchyKLDivergence
-log_uniform_kl_divergence = LogUniformKLDivergence
-normal_kl_divergence = NormalKLDivergence
-# pylint: enable=invalid-name
-
-# Utility functions, following tf.keras
-
-
-def serialize(initializer):
-  return tf.keras.utils.serialize_keras_object(initializer)
-
-
-def deserialize(config, custom_objects=None):
-  return tf.keras.utils.deserialize_keras_object(
-      config,
-      module_objects=globals(),
-      custom_objects=custom_objects,
-      printable_module_name='regularizers')
-
-
-def get(identifier, value=None):
-  """Getter for loading from strings; returns value if can't load."""
-  if value is None:
-    value = identifier
-  if identifier is None:
-    return None
-  elif isinstance(identifier, dict):
-    try:
-      return deserialize(identifier)
-    except ValueError:
-      return value
-  elif isinstance(identifier, six.string_types):
-    config = {'class_name': str(identifier), 'config': {}}
-    try:
-      return deserialize(config)
-    except ValueError:
-      return value
-  elif callable(identifier):
-    return identifier
-  return value
diff --git a/tensor2tensor/keras/regularizers_test.py b/tensor2tensor/keras/regularizers_test.py
deleted file mode 100644
index 45536cff1..000000000
--- a/tensor2tensor/keras/regularizers_test.py
+++ /dev/null
@@ -1,46 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 Keras-style regularizers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.keras import regularizers
-from tensor2tensor.utils import test_utils
-
-import tensorflow as tf
-import tensorflow_probability as tfp
-ed = tfp.edward2
-tf.compat.v1.enable_eager_execution()
-
-
-class RegularizersTest(tf.test.TestCase):
-
-  @test_utils.run_in_graph_and_eager_modes
-  def testHalfCauchyKLDivergence(self):
-    shape = (3,)
-    regularizer = regularizers.get('half_cauchy_kl_divergence')
-    variational_posterior = ed.Independent(
-        ed.LogNormal(loc=tf.zeros(shape), scale=1.).distribution,
-        reinterpreted_batch_ndims=1)
-    kl = regularizer(variational_posterior)
-    kl_value = self.evaluate(kl)
-    self.assertGreaterEqual(kl_value, 0.)
-
-
-if __name__ == '__main__':
-  tf.test.main()
diff --git a/tensor2tensor/layers/bayes.py b/tensor2tensor/layers/bayes.py
deleted file mode 100644
index eff347dc5..000000000
--- a/tensor2tensor/layers/bayes.py
+++ /dev/null
@@ -1,1439 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Bayesian neural network layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import functools
-import math
-from tensor2tensor.keras import constraints
-from tensor2tensor.keras import initializers
-from tensor2tensor.keras import regularizers
-
-import tensorflow as tf
-import tensorflow_probability as tfp
-from tensorflow_probability import edward2 as ed
-
-
-def add_weight(cls):
-  """Decorator for Layers, overriding add_weight for trainable initializers."""
-  @functools.wraps(cls.add_weight)
-  def _add_weight(self,
-                  name=None,
-                  shape=None,
-                  dtype=None,
-                  initializer=None,
-                  regularizer=None,
-                  **kwargs):
-    """Adds weight."""
-    if isinstance(initializer, tf.keras.layers.Layer):
-      weight = initializer(shape, dtype)
-      self._trainable_weights.extend(initializer.trainable_weights)  # pylint: disable=protected-access
-      self._non_trainable_weights.extend(initializer.non_trainable_weights)  # pylint: disable=protected-access
-      if regularizer is not None:
-        # TODO(trandustin): Replace need for this with
-        # Layer._handle_weight_regularization. For Eager compatibility, random
-        # variable __init__s cannot apply TF ops (cl/220898007).
-        def loss_fn():
-          """Creates a regularization loss `Tensor`."""
-          with tf.name_scope(name + '/Regularizer'):
-            return regularizer(initializer(shape, dtype))
-        self.add_loss(loss_fn)
-      return weight
-    return super(cls, self).add_weight(name=name,
-                                       shape=shape,
-                                       dtype=dtype,
-                                       initializer=initializer,
-                                       regularizer=regularizer,
-                                       **kwargs)
-  cls.add_weight = _add_weight
-  return cls
-
-
-@add_weight
-class Conv2DReparameterization(tf.keras.layers.Conv2D):
-  """2D convolution layer (e.g. spatial convolution over images).
-
-  The layer computes a variational Bayesian approximation to the distribution
-  over convolutional layers,
-
-  ```
-  p(outputs | inputs) = int conv2d(inputs; weights, bias) p(weights, bias)
-    dweights dbias.
-  ```
-
-  It does this with a stochastic forward pass, sampling from learnable
-  distributions on the kernel and bias. Gradients with respect to the
-  distributions' learnable parameters backpropagate via reparameterization.
-  Minimizing cross-entropy plus the layer's losses performs variational
-  minimum description length, i.e., it minimizes an upper bound to the negative
-  marginal likelihood.
-  """
-
-  def __init__(self,
-               filters,
-               kernel_size,
-               strides=(1, 1),
-               padding='valid',
-               data_format=None,
-               dilation_rate=(1, 1),
-               activation=None,
-               use_bias=True,
-               kernel_initializer='trainable_normal',
-               bias_initializer='zeros',
-               kernel_regularizer='normal_kl_divergence',
-               bias_regularizer=None,
-               activity_regularizer=None,
-               kernel_constraint=None,
-               bias_constraint=None,
-               **kwargs):
-    super(Conv2DReparameterization, self).__init__(
-        filters=filters,
-        kernel_size=kernel_size,
-        strides=strides,
-        padding=padding,
-        data_format=data_format,
-        dilation_rate=dilation_rate,
-        activation=activation,
-        use_bias=use_bias,
-        kernel_initializer=initializers.get(kernel_initializer),
-        bias_initializer=initializers.get(bias_initializer),
-        kernel_regularizer=regularizers.get(kernel_regularizer),
-        bias_regularizer=regularizers.get(bias_regularizer),
-        activity_regularizer=regularizers.get(activity_regularizer),
-        kernel_constraint=constraints.get(kernel_constraint),
-        bias_constraint=constraints.get(bias_constraint),
-        **kwargs)
-
-  def call_weights(self):
-    """Calls any weights if the initializer is itself a layer."""
-    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
-      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
-    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
-
-  def call(self, *args, **kwargs):
-    self.call_weights()
-    kwargs.pop('training', None)
-    return super(Conv2DReparameterization, self).call(*args, **kwargs)
-
-
-class Conv2DFlipout(Conv2DReparameterization):
-  """2D convolution layer (e.g. spatial convolution over images).
-
-  The layer computes a variational Bayesian approximation to the distribution
-  over convolutional layers,
-
-  ```
-  p(outputs | inputs) = int conv2d(inputs; weights, bias) p(weights, bias)
-    dweights dbias.
-  ```
-
-  It does this with a stochastic forward pass, sampling from learnable
-  distributions on the kernel and bias. Gradients with respect to the
-  distributions' learnable parameters backpropagate via reparameterization.
-  Minimizing cross-entropy plus the layer's losses performs variational
-  minimum description length, i.e., it minimizes an upper bound to the negative
-  marginal likelihood.
-
-  This layer uses the Flipout estimator (Wen et al., 2018) for integrating with
-  respect to the `kernel`. Namely, it applies
-  pseudo-independent weight perturbations via independent sign flips for each
-  example, enabling variance reduction over independent weight perturbations.
-  For this estimator to work, the `kernel` random variable must be able
-  to decompose as a sum of its mean and a perturbation distribution; the
-  perturbation distribution must be independent across weight elements and
-  symmetric around zero (for example, a fully factorized Gaussian).
-  """
-
-  def call(self, inputs):
-    if not isinstance(self.kernel, ed.RandomVariable):
-      return super(Conv2DFlipout, self).call(inputs)
-    self.call_weights()
-    outputs = self._apply_kernel(inputs)
-    if self.use_bias:
-      if self.data_format == 'channels_first':
-        outputs = tf.nn.bias_add(outputs, self.bias, data_format='NCHW')
-      else:
-        outputs = tf.nn.bias_add(outputs, self.bias, data_format='NHWC')
-    if self.activation is not None:
-      outputs = self.activation(outputs)
-    return outputs
-
-  def _apply_kernel(self, inputs):
-    input_shape = tf.shape(inputs)
-    batch_dim = input_shape[0]
-    if self._convolution_op is None:
-      padding = self.padding
-      if self.padding == 'causal':
-        padding = 'valid'
-      if not isinstance(padding, (list, tuple)):
-        padding = padding.upper()
-      self._convolution_op = functools.partial(
-          tf.nn.convolution,
-          strides=self.strides,
-          padding=padding,
-          data_format='NHWC' if self.data_format == 'channels_last' else 'NCHW',
-          dilations=self.dilation_rate)
-
-    if self.data_format == 'channels_first':
-      channels = input_shape[1]
-      sign_input_shape = [batch_dim, channels, 1, 1]
-      sign_output_shape = [batch_dim, self.filters, 1, 1]
-    else:
-      channels = input_shape[-1]
-      sign_input_shape = [batch_dim, 1, 1, channels]
-      sign_output_shape = [batch_dim, 1, 1, self.filters]
-    sign_input = 2 * tf.random.uniform(sign_input_shape,
-                                       minval=0,
-                                       maxval=2,
-                                       dtype=inputs.dtype) - 1
-    sign_output = 2 * tf.random.uniform(sign_output_shape,
-                                        minval=0,
-                                        maxval=2,
-                                        dtype=inputs.dtype) - 1
-    kernel_mean = self.kernel.distribution.mean()
-    perturbation = self.kernel - kernel_mean
-    outputs = self._convolution_op(inputs, kernel_mean)
-    outputs += self._convolution_op(inputs * sign_input,
-                                    perturbation) * sign_output
-    return outputs
-
-
-class Conv2DHierarchical(Conv2DFlipout):
-  """2D convolution layer with hierarchical distributions.
-
-  The layer computes a variational Bayesian approximation to the distribution
-  over convolutional layers, and where the distribution over weights
-  involves a hierarchical distribution with hidden unit noise coupling vectors
-  of the kernel weight matrix (Louizos et al., 2017),
-
-  ```
-  p(outputs | inputs) = int conv2d(inputs; new_kernel, bias) p(kernel,
-    local_scales, global_scale, bias) dkernel dlocal_scales dglobal_scale dbias.
-  ```
-
-  It does this with a stochastic forward pass, sampling from learnable
-  distributions on the kernel and bias. The kernel is written in non-centered
-  parameterization where
-
-  ```
-  new_kernel[i, j] = kernel[i, j] * local_scale[j] * global_scale.
-  ```
-
-  That is, there is "local" multiplicative noise which couples weights for each
-  output filter. There is also a "global" multiplicative noise which couples the
-  entire weight matrix. By default, the weights are normally distributed and the
-  local and global noises are half-Cauchy distributed; this makes the kernel a
-  horseshoe distribution (Carvalho et al., 2009; Polson and Scott, 2012).
-
-  The estimation uses Flipout for variance reduction with respect to sampling
-  the full weights. Gradients with respect to the distributions' learnable
-  parameters backpropagate via reparameterization. Minimizing cross-entropy
-  plus the layer's losses performs variational minimum description length,
-  i.e., it minimizes an upper bound to the negative marginal likelihood.
-  """
-
-  def __init__(self,
-               filters,
-               kernel_size,
-               strides=(1, 1),
-               padding='valid',
-               data_format=None,
-               dilation_rate=(1, 1),
-               activation=None,
-               use_bias=True,
-               kernel_initializer='trainable_normal',
-               bias_initializer='zeros',
-               local_scale_initializer='trainable_half_cauchy',
-               global_scale_initializer='trainable_half_cauchy',
-               kernel_regularizer='normal_kl_divergence',
-               bias_regularizer=None,
-               local_scale_regularizer='half_cauchy_kl_divergence',
-               global_scale_regularizer=regularizers.HalfCauchyKLDivergence(
-                   scale=1e-5),
-               activity_regularizer=None,
-               kernel_constraint=None,
-               bias_constraint=None,
-               local_scale_constraint='positive',
-               global_scale_constraint='positive',
-               **kwargs):
-    self.local_scale_initializer = initializers.get(local_scale_initializer)
-    self.global_scale_initializer = initializers.get(global_scale_initializer)
-    self.local_scale_regularizer = regularizers.get(local_scale_regularizer)
-    self.global_scale_regularizer = regularizers.get(global_scale_regularizer)
-    self.local_scale_constraint = constraints.get(local_scale_constraint)
-    self.global_scale_constraint = constraints.get(global_scale_constraint)
-    super(Conv2DHierarchical, self).__init__(
-        filters=filters,
-        kernel_size=kernel_size,
-        strides=strides,
-        padding=padding,
-        data_format=data_format,
-        dilation_rate=dilation_rate,
-        activation=activation,
-        use_bias=use_bias,
-        kernel_initializer=initializers.get(kernel_initializer),
-        bias_initializer=initializers.get(bias_initializer),
-        kernel_regularizer=regularizers.get(kernel_regularizer),
-        bias_regularizer=regularizers.get(bias_regularizer),
-        activity_regularizer=regularizers.get(activity_regularizer),
-        kernel_constraint=constraints.get(kernel_constraint),
-        bias_constraint=constraints.get(bias_constraint),
-        **kwargs)
-
-  def build(self, input_shape):
-    self.local_scale = self.add_weight(
-        shape=(self.filters,),
-        name='local_scale',
-        initializer=self.local_scale_initializer,
-        regularizer=self.local_scale_regularizer,
-        constraint=self.local_scale_constraint)
-    self.global_scale = self.add_weight(
-        shape=(),
-        name='global_scale',
-        initializer=self.global_scale_initializer,
-        regularizer=self.global_scale_regularizer,
-        constraint=self.global_scale_constraint)
-    super(Conv2DHierarchical, self).build(input_shape)
-
-  def call_weights(self):
-    """Calls any weights if the initializer is itself a layer."""
-    if isinstance(self.local_scale_initializer, tf.keras.layers.Layer):
-      self.local_scale = self.local_scale_initializer(self.local_scale.shape,
-                                                      self.dtype)
-    if isinstance(self.global_scale_initializer, tf.keras.layers.Layer):
-      self.global_scale = self.global_scale_initializer(self.global_scale.shape,
-                                                        self.dtype)
-    super(Conv2DHierarchical, self).call_weights()
-
-  def _apply_kernel(self, inputs):
-    outputs = super(Conv2DHierarchical, self)._apply_kernel(inputs)
-    if self.data_format == 'channels_first':
-      local_scale = tf.reshape(self.local_scale, [1, -1, 1, 1])
-    else:
-      local_scale = tf.reshape(self.local_scale, [1, 1, 1, -1])
-    # TODO(trandustin): Figure out what to set local/global scales to at test
-    # time. Means don't exist for Half-Cauchy approximate posteriors.
-    outputs *= local_scale * self.global_scale
-    return outputs
-
-
-class Conv2DVariationalDropout(Conv2DReparameterization):
-  """2D convolution layer with variational dropout (Kingma et al., 2015).
-
-  Implementation follows the additive parameterization of
-  Molchanov et al. (2017).
-  """
-
-  def __init__(self,
-               filters,
-               kernel_size,
-               strides=(1, 1),
-               padding='valid',
-               data_format=None,
-               dilation_rate=(1, 1),
-               activation=None,
-               use_bias=True,
-               kernel_initializer='trainable_normal',
-               bias_initializer='zeros',
-               kernel_regularizer='log_uniform_kl_divergence',
-               bias_regularizer=None,
-               activity_regularizer=None,
-               kernel_constraint=None,
-               bias_constraint=None,
-               **kwargs):
-    super(Conv2DVariationalDropout, self).__init__(
-        filters=filters,
-        kernel_size=kernel_size,
-        strides=strides,
-        padding=padding,
-        data_format=data_format,
-        dilation_rate=dilation_rate,
-        activation=activation,
-        use_bias=use_bias,
-        kernel_initializer=initializers.get(kernel_initializer),
-        bias_initializer=initializers.get(bias_initializer),
-        kernel_regularizer=regularizers.get(kernel_regularizer),
-        bias_regularizer=regularizers.get(bias_regularizer),
-        activity_regularizer=regularizers.get(activity_regularizer),
-        kernel_constraint=constraints.get(kernel_constraint),
-        bias_constraint=constraints.get(bias_constraint),
-        **kwargs)
-
-  def call(self, inputs, training=None):
-    if not isinstance(self.kernel, ed.RandomVariable):
-      return super(Conv2DVariationalDropout, self).call(inputs)
-    self.call_weights()
-    if training is None:
-      training = tf.keras.backend.learning_phase()
-    if self._convolution_op is None:
-      padding = self.padding
-      if self.padding == 'causal':
-        padding = 'valid'
-      if not isinstance(padding, (list, tuple)):
-        padding = padding.upper()
-      self._convolution_op = functools.partial(
-          tf.nn.convolution,
-          strides=self.strides,
-          padding=padding,
-          data_format='NHWC' if self.data_format == 'channels_last' else 'NCHW',
-          dilations=self.dilation_rate)
-
-    def dropped_inputs():
-      """Forward pass with dropout."""
-      # Clip magnitude of dropout rate, where we get the dropout rate alpha from
-      # the additive parameterization (Molchanov et al., 2017): for weight ~
-      # Normal(mu, sigma**2), the variance `sigma**2 = alpha * mu**2`.
-      mean = self.kernel.distribution.mean()
-      log_variance = tf.log(self.kernel.distribution.variance())
-      log_alpha = log_variance - tf.log(tf.square(mean) +
-                                        tf.keras.backend.epsilon())
-      log_alpha = tf.clip_by_value(log_alpha, -8., 8.)
-      log_variance = log_alpha + tf.log(tf.square(mean) +
-                                        tf.keras.backend.epsilon())
-
-      means = self._convolution_op(inputs, mean)
-      stddevs = tf.sqrt(
-          self._convolution_op(tf.square(inputs), tf.exp(log_variance)) +
-          tf.keras.backend.epsilon())
-      if self.use_bias:
-        if self.data_format == 'channels_first':
-          means = tf.nn.bias_add(means, self.bias, data_format='NCHW')
-        else:
-          means = tf.nn.bias_add(means, self.bias, data_format='NHWC')
-      outputs = ed.Normal(loc=means, scale=stddevs)
-      if self.activation is not None:
-        outputs = self.activation(outputs)
-      return outputs
-
-    # Following tf.keras.Dropout, only apply variational dropout if training
-    # flag is True.
-    training_value = smart_constant_value(training)
-    if training_value is not None:
-      if training_value:
-        return dropped_inputs()
-      else:
-        return super(Conv2DVariationalDropout, self).call(inputs)
-    return tf.cond(training,
-                   dropped_inputs,
-                   lambda: super(Conv2DVariationalDropout, self).call(inputs))
-
-
-# From `tensorflow/python/framework/smart_cond.py`
-def smart_constant_value(pred):
-  """Return the bool value for `pred`, or None if `pred` had a dynamic value.
-
-  Arguments:
-    pred: A scalar, either a Python bool or tensor.
-
-  Returns:
-    True or False if `pred` has a constant boolean value, None otherwise.
-
-  Raises:
-    TypeError: If `pred` is not a Tensor or bool.
-  """
-  if pred in {0, 1}:  # Accept 1/0 as valid boolean values
-    pred_value = bool(pred)
-  elif isinstance(pred, bool):
-    pred_value = pred
-  elif isinstance(pred, tf.Tensor):
-    pred_value = tf.contrib.util.constant_value(pred)
-  else:
-    raise TypeError('`pred` must be a Tensor, or a Python bool, or 1 or 0. '
-                    'Found instead: %s' % pred)
-  return pred_value
-
-
-@add_weight
-class DenseReparameterization(tf.keras.layers.Dense):
-  """Bayesian densely-connected layer estimated via reparameterization.
-
-  The layer computes a variational Bayesian approximation to the distribution
-  over densely-connected layers,
-
-  ```
-  p(outputs | inputs) = int dense(inputs; weights, bias) p(weights, bias)
-    dweights dbias.
-  ```
-
-  It does this with a stochastic forward pass, sampling from learnable
-  distributions on the kernel and bias. Gradients with respect to the
-  distributions' learnable parameters backpropagate via reparameterization.
-  Minimizing cross-entropy plus the layer's losses performs variational
-  minimum description length, i.e., it minimizes an upper bound to the negative
-  marginal likelihood.
-  """
-
-  def __init__(self,
-               units,
-               activation=None,
-               use_bias=True,
-               kernel_initializer='trainable_normal',
-               bias_initializer='zero',
-               kernel_regularizer='normal_kl_divergence',
-               bias_regularizer=None,
-               activity_regularizer=None,
-               **kwargs):
-    super(DenseReparameterization, self).__init__(
-        units=units,
-        activation=activation,
-        use_bias=use_bias,
-        kernel_initializer=initializers.get(kernel_initializer),
-        bias_initializer=initializers.get(bias_initializer),
-        kernel_regularizer=regularizers.get(kernel_regularizer),
-        bias_regularizer=regularizers.get(bias_regularizer),
-        activity_regularizer=regularizers.get(activity_regularizer),
-        **kwargs)
-
-  def call_weights(self):
-    """Calls any weights if the initializer is itself a layer."""
-    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
-      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
-    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
-
-  def call(self, *args, **kwargs):
-    self.call_weights()
-    kwargs.pop('training', None)
-    return super(DenseReparameterization, self).call(*args, **kwargs)
-
-
-class DenseDVI(DenseReparameterization):
-  """Densely-connected layer with deterministic VI (Wu et al., 2018).
-
-  This layer computes a variational inference approximation via first and second
-  moments. It is accurate if the kernel and bias initializers return factorized
-  normal random variables and the number of units is sufficiently large. The
-  advantage is that the forward pass is deterministic, reducing variance of
-  gradients during training. The disadvantage is an O(features^2*units) compute
-  and O(features^2 + features*units) memory complexity. In comparison,
-  DenseReparameterization has O(features*units) compute and memory complexity.
-
-  #### Examples
-
-  Below implements deterministic variational inference for Bayesian
-  feedforward network regression. We use the exact expected log-likelihood from
-  Wu et al. (2018), Eq. 8. Assume 2-D real-valued tensors of `features` and
-  `labels` of shapes `[batch_size, num_features]` and `[batch_size, 1]`
-  respectively.
-
-  ```python
-  from tensor2tensor.layers import bayes
-
-  model = tf.keras.Sequential([
-      bayes.DenseDVI(256, activation=tf.nn.relu),
-      bayes.DenseDVI(256, activation=tf.nn.relu),
-      bayes.DenseDVI(1, activation=None),
-  ])
-  locs = model(features)
-  nll = 0.5 * tf.reduce_mean(locs.distribution.variance() +
-                             (labels - locs.distribution.mean())**2)
-  kl = sum(model.losses) / total_dataset_size
-  loss = nll + kl
-  train_op = tf.train.AdamOptimizer(0.1).minimize(loss)
-  ```
-
-  For evaluation, feed in data and use, e.g., `predictions.distribution.mean()`
-  to make predictions via the posterior predictive distribution.
-
-  ```python
-  predictions = ed.Normal(loc=locs.distribution.mean(),
-                          scale=locs.distribution.variance() + 1.)
-  ```
-  """
-
-  def call(self, inputs):
-    if (not isinstance(inputs, ed.RandomVariable) and
-        not isinstance(self.kernel, ed.RandomVariable) and
-        not isinstance(self.bias, ed.RandomVariable)):
-      return super(DenseDVI, self).call(inputs)
-    self.call_weights()
-    inputs_mean, inputs_variance, inputs_covariance = get_moments(inputs)
-    kernel_mean, kernel_variance, _ = get_moments(self.kernel)
-    if self.use_bias:
-      bias_mean, _, bias_covariance = get_moments(self.bias)
-
-    # E[outputs] = E[inputs] * E[kernel] + E[bias]
-    mean = tf.tensordot(inputs_mean, kernel_mean, [[-1], [0]])
-    if self.use_bias:
-      mean = tf.nn.bias_add(mean, bias_mean)
-
-    # Cov = E[inputs**2] Cov(kernel) + E[W]^T Cov(inputs) E[W] + Cov(bias)
-    # For first term, assume Cov(kernel) = 0 on off-diagonals so we only
-    # compute diagonal term.
-    covariance_diag = tf.tensordot(inputs_variance + inputs_mean**2,
-                                   kernel_variance, [[-1], [0]])
-    # Compute quadratic form E[W]^T Cov E[W] from right-to-left. First is
-    #  [..., features, features], [features, units] -> [..., features, units].
-    cov_w = tf.tensordot(inputs_covariance, kernel_mean, [[-1], [0]])
-    # Next is [..., features, units], [features, units] -> [..., units, units].
-    w_cov_w = tf.tensordot(cov_w, kernel_mean, [[-2], [0]])
-    covariance = w_cov_w
-    if self.use_bias:
-      covariance += bias_covariance
-    covariance = tf.matrix_set_diag(
-        covariance, tf.matrix_diag_part(covariance) + covariance_diag)
-
-    if self.activation in (tf.keras.activations.relu, tf.nn.relu):
-      # Compute activation's moments with variable names from Wu et al. (2018).
-      variance = tf.matrix_diag_part(covariance)
-      scale = tf.sqrt(variance)
-      mu = mean / (scale + tf.keras.backend.epsilon())
-      mean = scale * soft_relu(mu)
-
-      pairwise_variances = (tf.expand_dims(variance, -1) *
-                            tf.expand_dims(variance, -2))  # [..., units, units]
-      rho = covariance / tf.sqrt(pairwise_variances +
-                                 tf.keras.backend.epsilon())
-      rho = tf.clip_by_value(rho,
-                             -1. / (1. + tf.keras.backend.epsilon()),
-                             1. / (1. + tf.keras.backend.epsilon()))
-      s = covariance / (rho + tf.keras.backend.epsilon())
-      mu1 = tf.expand_dims(mu, -1)  # [..., units, 1]
-      mu2 = tf.matrix_transpose(mu1)  # [..., 1, units]
-      a = (soft_relu(mu1) * soft_relu(mu2) +
-           rho * tfp.distributions.Normal(0., 1.).cdf(mu1) *
-           tfp.distributions.Normal(0., 1.).cdf(mu2))
-      gh = tf.asinh(rho)
-      bar_rho = tf.sqrt(1. - rho**2)
-      gr = gh + rho / (1. + bar_rho)
-      # Include numerically stable versions of gr and rho when multiplying or
-      # dividing them. The sign of gr*rho and rho/gr is always positive.
-      safe_gr = tf.abs(gr) + 0.5 * tf.keras.backend.epsilon()
-      safe_rho = tf.abs(rho) + tf.keras.backend.epsilon()
-      exp_negative_q = gr / (2. * math.pi) * tf.exp(
-          -safe_rho / (2. * safe_gr * (1 + bar_rho)) +
-          (gh - rho) / (safe_gr * safe_rho) * mu1 * mu2)
-      covariance = s * (a + exp_negative_q)
-    elif self.activation not in (tf.keras.activations.linear, None):
-      raise NotImplementedError('Activation is {}. Deterministic variational '
-                                'inference is only available if activation is '
-                                'ReLU or None.'.format(self.activation))
-
-    return ed.MultivariateNormalFullCovariance(mean, covariance)
-
-
-def get_moments(x):
-  """Gets first and second moments of input."""
-  if isinstance(x, ed.RandomVariable):
-    mean = x.distribution.mean()
-    variance = x.distribution.variance()
-    try:
-      covariance = x.distribution.covariance()
-    except NotImplementedError:
-      covariance = tf.zeros(x.shape.concatenate(x.shape[-1]), dtype=x.dtype)
-      covariance = tf.matrix_set_diag(covariance, variance)
-  else:
-    mean = x
-    variance = tf.zeros_like(x)
-    covariance = tf.zeros(x.shape.concatenate(x.shape[-1]), dtype=x.dtype)
-  return mean, variance, covariance
-
-
-def soft_relu(x):
-  return (tfp.distributions.Normal(0., 1.).prob(x) +
-          x * tfp.distributions.Normal(0., 1.).cdf(x))
-
-
-class DenseFlipout(DenseReparameterization):
-  """Bayesian densely-connected layer estimated via Flipout (Wen et al., 2018).
-
-  The layer computes a variational Bayesian approximation to the distribution
-  over densely-connected layers,
-
-  ```
-  p(outputs | inputs) = int dense(inputs; weights, bias) p(weights, bias)
-    dweights dbias.
-  ```
-
-  It does this with a stochastic forward pass, sampling from learnable
-  distributions on the kernel and bias. Gradients with respect to the
-  distributions' learnable parameters backpropagate via reparameterization.
-  Minimizing cross-entropy plus the layer's losses performs variational
-  minimum description length, i.e., it minimizes an upper bound to the negative
-  marginal likelihood.
-
-  This layer uses the Flipout estimator (Wen et al., 2018) for integrating with
-  respect to the `kernel`. Namely, it applies
-  pseudo-independent weight perturbations via independent sign flips for each
-  example, enabling variance reduction over independent weight perturbations.
-  For this estimator to work, the `kernel` random variable must be able
-  to decompose as a sum of its mean and a perturbation distribution; the
-  perturbation distribution must be independent across weight elements and
-  symmetric around zero (for example, a fully factorized Gaussian).
-  """
-
-  def call(self, inputs):
-    if not isinstance(self.kernel, ed.RandomVariable):
-      return super(DenseFlipout, self).call(inputs)
-    self.call_weights()
-    input_shape = tf.shape(inputs)
-    sign_input = 2 * tf.random.uniform(input_shape,
-                                       minval=0,
-                                       maxval=2,
-                                       dtype=inputs.dtype) - 1
-    sign_output = 2 * tf.random.uniform(tf.concat([input_shape[:-1],
-                                                   [self.units]], 0),
-                                        minval=0,
-                                        maxval=2,
-                                        dtype=inputs.dtype) - 1
-    kernel_mean = self.kernel.distribution.mean()
-    perturbation = self.kernel - kernel_mean
-    if inputs.shape.ndims <= 2:
-      outputs = tf.matmul(inputs, kernel_mean)
-      outputs += tf.matmul(inputs * sign_input, perturbation) * sign_output
-    else:
-      outputs = tf.tensordot(inputs, kernel_mean, [[-1], [0]])
-      outputs += tf.tensordot(inputs * sign_input,
-                              perturbation,
-                              [[-1], [0]]) * sign_output
-    if self.use_bias:
-      outputs = tf.nn.bias_add(outputs, self.bias)
-    if self.activation is not None:
-      outputs = self.activation(outputs)
-    return outputs
-
-
-class DenseVariationalDropout(DenseReparameterization):
-  """Densely-connected layer with variational dropout (Kingma et al., 2015).
-
-  Implementation follows the additive parameterization of
-  Molchanov et al. (2017).
-  """
-
-  def __init__(self,
-               units,
-               activation=None,
-               use_bias=True,
-               kernel_initializer='trainable_normal',
-               bias_initializer='zero',
-               kernel_regularizer='log_uniform_kl_divergence',
-               bias_regularizer=None,
-               activity_regularizer=None,
-               **kwargs):
-    super(DenseVariationalDropout, self).__init__(
-        units=units,
-        activation=activation,
-        use_bias=use_bias,
-        kernel_initializer=initializers.get(kernel_initializer),
-        bias_initializer=initializers.get(bias_initializer),
-        kernel_regularizer=regularizers.get(kernel_regularizer),
-        bias_regularizer=regularizers.get(bias_regularizer),
-        activity_regularizer=regularizers.get(activity_regularizer),
-        **kwargs)
-
-  def call(self, inputs, training=None):
-    if not isinstance(self.kernel, ed.RandomVariable):
-      return super(DenseVariationalDropout, self).call(inputs)
-    self.call_weights()
-    if training is None:
-      training = tf.keras.backend.learning_phase()
-
-    def dropped_inputs():
-      """Forward pass with dropout."""
-      # Clip magnitude of dropout rate, where we get the dropout rate alpha from
-      # the additive parameterization (Molchanov et al., 2017): for weight ~
-      # Normal(mu, sigma**2), the variance `sigma**2 = alpha * mu**2`.
-      mean = self.kernel.distribution.mean()
-      log_variance = tf.log(self.kernel.distribution.variance())
-      log_alpha = log_variance - tf.log(tf.square(mean) +
-                                        tf.keras.backend.epsilon())
-      log_alpha = tf.clip_by_value(log_alpha, -8., 8.)
-      log_variance = log_alpha + tf.log(tf.square(mean) +
-                                        tf.keras.backend.epsilon())
-
-      if inputs.shape.ndims <= 2:
-        means = tf.matmul(inputs, mean)
-        stddevs = tf.sqrt(
-            tf.matmul(tf.square(inputs), tf.exp(log_variance)) +
-            tf.keras.backend.epsilon())
-      else:
-        means = tf.tensordot(inputs, mean, [[-1], [0]])
-        stddevs = tf.sqrt(
-            tf.tensordot(tf.square(inputs), tf.exp(log_variance), [[-1], [0]]) +
-            tf.keras.backend.epsilon())
-      if self.use_bias:
-        means = tf.nn.bias_add(means, self.bias)
-      outputs = ed.Normal(loc=means, scale=stddevs)
-      if self.activation is not None:
-        outputs = self.activation(outputs)
-      return outputs
-
-    # Following tf.keras.Dropout, only apply variational dropout if training
-    # flag is True.
-    training_value = smart_constant_value(training)
-    if training_value is not None:
-      if training_value:
-        return dropped_inputs()
-      else:
-        return super(DenseVariationalDropout, self).call(inputs)
-    return tf.cond(training,
-                   dropped_inputs,
-                   lambda: super(DenseVariationalDropout, self).call(inputs))
-
-
-class DenseHierarchical(DenseVariationalDropout):
-  """Bayesian densely-connected layer with hierarchical distributions.
-
-  The layer computes a variational Bayesian approximation to the distribution
-  over densely-connected layers, and where the distribution over weights
-  involves a hierarchical distribution with hidden unit noise coupling vectors
-  of the kernel weight matrix (Louizos et al., 2017),
-
-  ```
-  p(outputs | inputs) = int dense(inputs; new_kernel, bias) p(kernel,
-    local_scales, global_scale, bias) dkernel dlocal_scales dglobal_scale dbias.
-  ```
-
-  It does this with a stochastic forward pass, sampling from learnable
-  distributions on the kernel and bias. The kernel is written in non-centered
-  parameterization where
-
-  ```
-  new_kernel[i, j] = kernel[i, j] * local_scale[i] * global_scale.
-  ```
-
-  That is, there is "local" multiplicative noise which couples weights for each
-  input neuron. There is also a "global" multiplicative noise which couples the
-  entire weight matrix. By default, the weights are normally distributed and the
-  local and global noises are half-Cauchy distributed; this makes the kernel a
-  horseshoe distribution (Carvalho et al., 2009; Polson and Scott, 2012).
-
-  The estimation uses local reparameterization to avoid sampling the full
-  weights. Gradients with respect to the distributions' learnable parameters
-  backpropagate via reparameterization. Minimizing cross-entropy plus the
-  layer's losses performs variational minimum description length, i.e., it
-  minimizes an upper bound to the negative marginal likelihood.
-  """
-
-  def __init__(self,
-               units,
-               activation=None,
-               use_bias=True,
-               kernel_initializer='trainable_normal',
-               bias_initializer='zero',
-               local_scale_initializer='trainable_half_cauchy',
-               global_scale_initializer='trainable_half_cauchy',
-               kernel_regularizer='normal_kl_divergence',
-               bias_regularizer=None,
-               local_scale_regularizer='half_cauchy_kl_divergence',
-               global_scale_regularizer=regularizers.HalfCauchyKLDivergence(
-                   scale=1e-5),
-               activity_regularizer=None,
-               local_scale_constraint='positive',
-               global_scale_constraint='positive',
-               **kwargs):
-    self.local_scale_initializer = initializers.get(local_scale_initializer)
-    self.global_scale_initializer = initializers.get(global_scale_initializer)
-    self.local_scale_regularizer = regularizers.get(local_scale_regularizer)
-    self.global_scale_regularizer = regularizers.get(global_scale_regularizer)
-    self.local_scale_constraint = constraints.get(local_scale_constraint)
-    self.global_scale_constraint = constraints.get(global_scale_constraint)
-    super(DenseHierarchical, self).__init__(
-        units=units,
-        activation=activation,
-        use_bias=use_bias,
-        kernel_initializer=initializers.get(kernel_initializer),
-        bias_initializer=initializers.get(bias_initializer),
-        kernel_regularizer=regularizers.get(kernel_regularizer),
-        bias_regularizer=regularizers.get(bias_regularizer),
-        activity_regularizer=regularizers.get(activity_regularizer),
-        **kwargs)
-
-  def build(self, input_shape):
-    input_shape = tf.TensorShape(input_shape)
-    input_dim = input_shape[-1]
-    if isinstance(input_dim, tf.Dimension):
-      input_dim = input_dim.value
-    self.local_scale = self.add_weight(
-        shape=(input_dim,),
-        name='local_scale',
-        initializer=self.local_scale_initializer,
-        regularizer=self.local_scale_regularizer,
-        constraint=self.local_scale_constraint)
-    self.global_scale = self.add_weight(
-        shape=(),
-        name='global_scale',
-        initializer=self.global_scale_initializer,
-        regularizer=self.global_scale_regularizer,
-        constraint=self.global_scale_constraint)
-    super(DenseHierarchical, self).build(input_shape)
-
-  def call_weights(self):
-    """Calls any weights if the initializer is itself a layer."""
-    if isinstance(self.local_scale_initializer, tf.keras.layers.Layer):
-      self.local_scale = self.local_scale_initializer(self.local_scale.shape,
-                                                      self.dtype)
-    if isinstance(self.global_scale_initializer, tf.keras.layers.Layer):
-      self.global_scale = self.global_scale_initializer(self.global_scale.shape,
-                                                        self.dtype)
-    super(DenseHierarchical, self).call_weights()
-
-  def call(self, inputs, training=None):
-    self.call_weights()
-    # TODO(trandustin): Figure out what to set local/global scales to at test
-    # time. Means don't exist for Half-Cauchy approximate posteriors.
-    inputs *= self.local_scale[tf.newaxis, :] * self.global_scale
-    return super(DenseHierarchical, self).call(inputs, training=training)
-
-
-@add_weight
-class LSTMCellReparameterization(tf.keras.layers.LSTMCell):
-  """Bayesian LSTM cell class estimated via reparameterization.
-
-  The layer computes a variational Bayesian approximation to the distribution
-  over LSTM cell functions,
-
-  ```
-  p(outputs | inputs) = int lstm_cell(inputs; weights, bias) p(weights, bias)
-    dweights dbias,
-  ```
-
-  where the weights consist of both input and recurrent weights.
-
-  It does this with a stochastic forward pass, sampling from learnable
-  distributions on the kernel, recurrent kernel, and bias. Gradients with
-  respect to the distributions' learnable parameters backpropagate via
-  reparameterization.  Minimizing cross-entropy plus the layer's losses performs
-  variational minimum description length, i.e., it minimizes an upper bound to
-  the negative marginal likelihood.
-  """
-
-  def __init__(self,
-               units,
-               activation='tanh',
-               recurrent_activation='hard_sigmoid',
-               use_bias=True,
-               kernel_initializer='trainable_normal',
-               recurrent_initializer='trainable_normal',
-               bias_initializer='zeros',
-               unit_forget_bias=True,
-               kernel_regularizer='normal_kl_divergence',
-               recurrent_regularizer='normal_kl_divergence',
-               bias_regularizer=None,
-               kernel_constraint=None,
-               recurrent_constraint=None,
-               bias_constraint=None,
-               dropout=0.,
-               recurrent_dropout=0.,
-               implementation=1,
-               **kwargs):
-    super(LSTMCellReparameterization, self).__init__(
-        units=units,
-        activation=activation,
-        recurrent_activation=recurrent_activation,
-        use_bias=use_bias,
-        kernel_initializer=initializers.get(kernel_initializer),
-        recurrent_initializer=initializers.get(recurrent_initializer),
-        bias_initializer=initializers.get(bias_initializer),
-        unit_forget_bias=unit_forget_bias,
-        kernel_regularizer=regularizers.get(kernel_regularizer),
-        recurrent_regularizer=regularizers.get(recurrent_regularizer),
-        bias_regularizer=regularizers.get(bias_regularizer),
-        kernel_constraint=constraints.get(kernel_constraint),
-        recurrent_constraint=constraints.get(recurrent_constraint),
-        bias_constraint=constraints.get(bias_constraint),
-        dropout=dropout,
-        recurrent_dropout=recurrent_dropout,
-        implementation=implementation,
-        **kwargs)
-
-  def build(self, input_shape):
-    input_shape = tf.TensorShape(input_shape)
-    input_dim = input_shape[-1]
-    if isinstance(input_dim, tf.Dimension):
-      input_dim = input_dim.value
-    self.kernel = self.add_weight(
-        shape=(input_dim, self.units * 4),
-        name='kernel',
-        initializer=self.kernel_initializer,
-        regularizer=self.kernel_regularizer,
-        constraint=self.kernel_constraint)
-    self.recurrent_kernel = self.add_weight(
-        shape=(self.units, self.units * 4),
-        name='recurrent_kernel',
-        initializer=self.recurrent_initializer,
-        regularizer=self.recurrent_regularizer,
-        constraint=self.recurrent_constraint)
-
-    if self.use_bias:
-      if self.unit_forget_bias:
-        if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-          def bias_mean_initializer(_, *args, **kwargs):
-            return tf.concat([
-                tf.keras.initializers.truncated_normal(
-                    stddev=1e-5)((self.units,), *args, **kwargs),
-                tf.keras.initializers.truncated_normal(
-                    mean=1., stddev=1e-5)((self.units,), *args, **kwargs),
-                tf.keras.initializers.truncated_normal(
-                    stddev=1e-5)((self.units * 2,), *args, **kwargs),
-            ], axis=0)
-          bias_initializer = initializers.TrainableNormal(
-              mean_initializer=bias_mean_initializer)
-        else:
-          def bias_initializer(_, *args, **kwargs):
-            return tf.keras.backend.concatenate([
-                self.bias_initializer((self.units,), *args, **kwargs),
-                tf.keras.initializers.Ones()((self.units,), *args, **kwargs),
-                self.bias_initializer((self.units * 2,), *args, **kwargs),
-            ])
-      else:
-        bias_initializer = self.bias_initializer
-      self.bias = self.add_weight(
-          shape=(self.units * 4,),
-          name='bias',
-          initializer=bias_initializer,
-          regularizer=self.bias_regularizer,
-          constraint=self.bias_constraint)
-    else:
-      self.bias = None
-    self.built = True
-
-  def call_weights(self):
-    """Calls any weights if the initializer is itself a layer."""
-    if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
-      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
-    if isinstance(self.recurrent_initializer, tf.keras.layers.Layer):
-      self.recurrent_kernel = self.recurrent_initializer(
-          self.recurrent_kernel.shape, self.dtype)
-    if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
-
-  def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
-    """Get the initial state and side-effect sampling of stochastic weights."""
-    if self.built:
-      self.call_weights()
-    return super(LSTMCellReparameterization, self).get_initial_state(
-        inputs=inputs, batch_size=batch_size, dtype=dtype)
-
-
-class LSTMCellFlipout(LSTMCellReparameterization):
-  """Bayesian LSTM cell class estimated via Flipout (Wen et al., 2018).
-
-  The layer computes a variational Bayesian approximation to the distribution
-  over LSTM cell functions,
-
-  ```
-  p(outputs | inputs) = int lstm_cell(inputs; weights, bias) p(weights, bias)
-    dweights dbias,
-  ```
-
-  where the weights consist of both input and recurrent weights.
-
-  It does this with a stochastic forward pass, sampling from learnable
-  distributions on the kernel, recurrent kernel, and bias. Gradients with
-  respect to the distributions' learnable parameters backpropagate via
-  reparameterization.  Minimizing cross-entropy plus the layer's losses performs
-  variational minimum description length, i.e., it minimizes an upper bound to
-  the negative marginal likelihood.
-
-  This layer uses the Flipout estimator (Wen et al., 2018) for integrating with
-  respect to the `kernel` and `recurrent_kernel`. Namely, it applies
-  pseudo-independent weight perturbations via independent sign flips for each
-  example, enabling variance reduction over independent weight perturbations.
-  For this estimator to work, the `kernel` and `recurrent_kernel` random
-  variable must be able to decompose as a sum of its mean and a perturbation
-  distribution; the perturbation distribution must be independent across weight
-  elements and symmetric around zero (for example, a fully factorized Gaussian).
-  """
-
-  def _call_sign_flips(self, inputs=None, batch_size=None, dtype=None):
-    """Builds per-example sign flips for pseudo-independent perturbations."""
-    # TODO(trandustin): We add and call this method separately from build().
-    # This is because build() operates on a static input_shape. We need dynamic
-    # input shapes as we operate on the batch size which is often dynamic.
-    if inputs is not None:
-      batch_size = tf.shape(inputs)[0]
-      dtype = inputs.dtype
-    input_dim = tf.shape(self.kernel)[0]
-    self.sign_input = 2 * tf.random.uniform(
-        [batch_size, 4 * input_dim], minval=0, maxval=2, dtype=dtype) - 1
-    self.sign_output = 2 * tf.random.uniform(
-        [batch_size, 4 * self.units], minval=0, maxval=2, dtype=dtype) - 1
-    self.recurrent_sign_input = 2 * tf.random.uniform(
-        [batch_size, 4 * self.units], minval=0, maxval=2, dtype=dtype) - 1
-    self.recurrent_sign_output = 2 * tf.random.uniform(
-        [batch_size, 4 * self.units], minval=0, maxval=2, dtype=dtype) - 1
-
-  def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
-    """Get the initial state and side-effect sampling of stochastic weights."""
-    if self.built:
-      self._call_sign_flips(inputs, batch_size, dtype)
-    return super(LSTMCellFlipout, self).get_initial_state(
-        inputs=inputs, batch_size=batch_size, dtype=dtype)
-
-  def _compute_carry_and_output(self, x, h_tm1, c_tm1):
-    """Computes carry and output using split kernels."""
-    if not isinstance(self.recurrent_kernel, ed.RandomVariable):
-      return super(LSTMCellFlipout, self)._compute_carry_and_output(x,
-                                                                    h_tm1,
-                                                                    c_tm1)
-    x_i, x_f, x_c, x_o = x
-    h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o = h_tm1
-    kernel_mean = self.recurrent_kernel.distribution.mean()
-    perturbation = self.recurrent_kernel - kernel_mean
-    k_i, k_f, k_c, k_o = tf.split(kernel_mean, num_or_size_splits=4, axis=1)
-    p_i, p_f, p_c, p_o = tf.split(perturbation, num_or_size_splits=4, axis=1)
-    si_i, si_f, si_c, si_o = tf.split(self.recurrent_sign_input,
-                                      num_or_size_splits=4, axis=1)
-    so_i, so_f, so_c, so_o = tf.split(self.recurrent_sign_output,
-                                      num_or_size_splits=4, axis=1)
-    z0 = (x_i + tf.keras.backend.dot(h_tm1_i, k_i) +
-          tf.keras.backend.dot(h_tm1_i * si_i, p_i) * so_i)
-    z1 = (x_f + tf.keras.backend.dot(h_tm1_f, k_f) +
-          tf.keras.backend.dot(h_tm1_f * si_f, p_f) * so_f)
-    z2 = (x_c + tf.keras.backend.dot(h_tm1_c, k_c) +
-          tf.keras.backend.dot(h_tm1_c * si_c, p_c) * so_c)
-    z3 = (x_o + tf.keras.backend.dot(h_tm1_o, k_o) +
-          tf.keras.backend.dot(h_tm1_o * si_o, p_o) * so_o)
-    i = self.recurrent_activation(z0)
-    f = self.recurrent_activation(z1)
-    c = f * c_tm1 + i * self.activation(z2)
-    o = self.recurrent_activation(z3)
-    return c, o
-
-  def call(self, inputs, states, training=None):
-    # TODO(trandustin): Enable option for Flipout on only the kernel or
-    # recurrent_kernel. If only one is a random variable, we currently default
-    # to weight reparameterization.
-    if (not isinstance(self.kernel, ed.RandomVariable) or
-        not isinstance(self.recurrent_kernel, ed.RandomVariable)):
-      return super(LSTMCellFlipout, self).call(inputs, states, training)
-    if not hasattr(self, 'sign_input'):
-      self._call_sign_flips(inputs)
-    h_tm1 = states[0]  # previous memory state
-    c_tm1 = states[1]  # previous carry state
-
-    dp_mask = self.get_dropout_mask_for_cell(inputs, training, count=4)
-    rec_dp_mask = self.get_recurrent_dropout_mask_for_cell(
-        h_tm1, training, count=4)
-
-    if self.implementation == 1:
-      if 0 < self.dropout < 1.:
-        inputs_i = inputs * dp_mask[0]
-        inputs_f = inputs * dp_mask[1]
-        inputs_c = inputs * dp_mask[2]
-        inputs_o = inputs * dp_mask[3]
-      else:
-        inputs_i = inputs
-        inputs_f = inputs
-        inputs_c = inputs
-        inputs_o = inputs
-      kernel_mean = self.kernel.distribution.mean()
-      perturbation = self.kernel - kernel_mean
-      k_i, k_f, k_c, k_o = tf.split(kernel_mean, num_or_size_splits=4, axis=1)
-      p_i, p_f, p_c, p_o = tf.split(perturbation, num_or_size_splits=4, axis=1)
-      si_i, si_f, si_c, si_o = tf.split(self.sign_input,
-                                        num_or_size_splits=4, axis=1)
-      so_i, so_f, so_c, so_o = tf.split(self.sign_output,
-                                        num_or_size_splits=4, axis=1)
-      x_i = (tf.keras.backend.dot(inputs_i, k_i) +
-             tf.keras.backend.dot(inputs_i * si_i, p_i) * so_i)
-      x_f = (tf.keras.backend.dot(inputs_f, k_f) +
-             tf.keras.backend.dot(inputs_f * si_f, p_f) * so_f)
-      x_c = (tf.keras.backend.dot(inputs_c, k_c) +
-             tf.keras.backend.dot(inputs_c * si_c, p_c) * so_c)
-      x_o = (tf.keras.backend.dot(inputs_o, k_o) +
-             tf.keras.backend.dot(inputs_o * si_o, p_o) * so_o)
-      if self.use_bias:
-        b_i, b_f, b_c, b_o = tf.split(
-            self.bias, num_or_size_splits=4, axis=0)
-        x_i = tf.keras.backend.bias_add(x_i, b_i)
-        x_f = tf.keras.backend.bias_add(x_f, b_f)
-        x_c = tf.keras.backend.bias_add(x_c, b_c)
-        x_o = tf.keras.backend.bias_add(x_o, b_o)
-
-      if 0 < self.recurrent_dropout < 1.:
-        h_tm1_i = h_tm1 * rec_dp_mask[0]
-        h_tm1_f = h_tm1 * rec_dp_mask[1]
-        h_tm1_c = h_tm1 * rec_dp_mask[2]
-        h_tm1_o = h_tm1 * rec_dp_mask[3]
-      else:
-        h_tm1_i = h_tm1
-        h_tm1_f = h_tm1
-        h_tm1_c = h_tm1
-        h_tm1_o = h_tm1
-      x = (x_i, x_f, x_c, x_o)
-      h_tm1 = (h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o)
-      c, o = self._compute_carry_and_output(x, h_tm1, c_tm1)
-    else:
-      if 0. < self.dropout < 1.:
-        inputs = inputs * dp_mask[0]
-      kernel_mean = self.kernel.distribution.mean()
-      perturbation = self.kernel - kernel_mean
-      z = tf.keras.backend.dot(inputs, kernel_mean)
-      z += tf.keras.backend.dot(inputs * self.sign_input,
-                                perturbation) * self.sign_output
-      if 0. < self.recurrent_dropout < 1.:
-        h_tm1 = h_tm1 * rec_dp_mask[0]
-      recurrent_kernel_mean = self.recurrent_kernel.distribution.mean()
-      perturbation = self.recurrent_kernel - recurrent_kernel_mean
-      z += tf.keras.backend.dot(h_tm1, recurrent_kernel_mean)
-      z += tf.keras.backend.dot(h_tm1 * self.recurrent_sign_input,
-                                perturbation) * self.recurrent_sign_output
-      if self.use_bias:
-        z = tf.keras.backend.bias_add(z, self.bias)
-
-      z = tf.split(z, num_or_size_splits=4, axis=1)
-      c, o = self._compute_carry_and_output_fused(z, c_tm1)
-
-    h = o * self.activation(c)
-    return h, [h, c]
-
-
-class NCPNormalPerturb(tf.keras.layers.Layer):
-  """Noise contrastive prior for continuous inputs (Hafner et al., 2018).
-
-  The layer doubles the inputs' batch size and adds a random normal perturbation
-  to the concatenated second batch. This acts an input prior to be used in
-  combination with an output prior. The output prior reduces the second batch
-  (reverting to the inputs' original shape) and computes a regularizer that
-  matches the second batch towards some output (e.g., uniform distribution).
-  This layer implementation is inspired by the Aboleth library.
-
-  #### Examples
-
-  Below implements neural network regression with heteroskedastic noise,
-  noise contrastive priors, and being Bayesian only at the mean's output layer.
-
-  ```python
-  from tensor2tensor.layers import bayes
-
-  batch_size, dataset_size = 128, 1000
-  features, labels = get_some_dataset()
-
-  inputs = keras.Input(shape=(25,))
-  x = bayes.NCPNormalPerturb()(inputs)  # double input batch
-  x = layers.Dense(64, activation='relu')(x)
-  x = layers.Dense(64, activation='relu')(x)
-  means = bayes.DenseVariationalDropout(1, activation=None)(x)  # get mean dist.
-  means = bayes.NCPNormalOutput(labels)(means)  # halve input batch
-  stddevs = tf.keras.layers.Dense(1, activation='softplus')(x[:batch_size])
-  outputs = tf.keras.layers.Lambda(lambda x: ed.Normal(x[0], x[1]))([means,
-                                                                     stddevs])
-  model = tf.keras.Model(inputs=inputs, outputs=outputs)
-
-  predictions = model(features)
-  loss = tf.reduce_mean(predictions.distribution.log_prob(labels))
-  loss += model.losses[0] / dataset_size  # KL regularizer for output layer
-  loss += model.losses[-1]
-  train_op = tf.train.AdamOptimizer(0.1).minimize(loss)
-  ```
-
-  The network applies `bayes.NCPNormalPerturb()` to double the input batch
-  size and add Gaussian noise to the second half; then feedforward layers; then
-  `bayes.DenseVariational` to be Bayesian about the output density's mean; then
-  `bayes.NCPNormalOutput` centered at the labels to revert to the batch size
-  and compute a loss on the second half; then parameterize the output density's
-  standard deviations; then compute the total loss function as the sum of the
-  model's negative log-likelihood, KL divergence for the Bayesian mean layer,
-  and NCP loss.
-  """
-
-  def __init__(self, mean=0., stddev=1., seed=None, **kwargs):
-    self.mean = mean
-    self.stddev = stddev
-    self.seed = seed
-    super(NCPNormalPerturb, self).__init__(**kwargs)
-
-  def call(self, inputs):
-    noise = tf.random.normal(tf.shape(inputs),
-                             mean=self.mean,
-                             stddev=self.stddev,
-                             dtype=inputs.dtype,
-                             seed=self.seed)
-    perturbed_inputs = inputs + noise
-    return tf.concat([inputs, perturbed_inputs], 0)
-
-
-class NCPCategoricalPerturb(tf.keras.layers.Layer):
-  """Noise contrastive prior for discrete inputs (Hafner et al., 2018).
-
-  The layer doubles the inputs' batch size and randomly flips categories
-  for the concatenated second batch (all features must be integer-valued). This
-  acts an input prior to be used in combination with an output prior. The output
-  prior reduces the second batch (reverting to the inputs' original shape) and
-  computes a regularizer that matches the second batch towards some output
-  (e.g., uniform distribution). This layer implementation is inspired by the
-  Aboleth library.
-
-  #### Examples
-
-  Below implements neural network regression with heteroskedastic noise,
-  noise contrastive priors, and being Bayesian only at the mean's output layer.
-
-  ```python
-  from tensor2tensor.layers import bayes
-
-  batch_size, dataset_size = 128, 1000
-  features, labels = get_some_dataset()
-
-  inputs = keras.Input(shape=(25,))
-  x = bayes.NCPCategoricalPerturb(10)(inputs)  # double input batch
-  x = layers.Dense(64, activation='relu')(x)
-  x = layers.Dense(64, activation='relu')(x)
-  means = bayes.DenseVariationalDropout(1, activation=None)(x)  # get mean dist.
-  means = bayes.NCPNormalOutput(labels)(means)  # halve input batch
-  stddevs = tf.keras.layers.Dense(1, activation='softplus')(x[:batch_size])
-  outputs = tf.keras.layers.Lambda(lambda x: ed.Normal(x[0], x[1]))([means,
-                                                                     stddevs])
-  model = tf.keras.Model(inputs=inputs, outputs=outputs)
-
-  predictions = model(features)
-  loss = tf.reduce_mean(predictions.distribution.log_prob(labels))
-  loss += model.losses[0] / dataset_size  # KL regularizer for output layer
-  loss += model.losses[-1]
-  train_op = tf.train.AdamOptimizer(0.1).minimize(loss)
-  ```
-
-  The network applies `bayes.NCPCategoricalPerturb()` to double the input batch
-  size and flip categories for the second half; then feedforward layers; then
-  `bayes.DenseVariational` to be Bayesian about the output density's mean; then
-  `bayes.NCPNormalOutput` centered at the labels to revert to the batch size
-  and compute a loss on the second half; then parameterize the output density's
-  standard deviations; then compute the total loss function as the sum of the
-  model's negative log-likelihood, KL divergence for the Bayesian mean layer,
-  and NCP loss.
-  """
-
-  def __init__(self, input_dim, probs=0.1, **kwargs):
-    """Creates layer.
-
-    Args:
-      input_dim: int > 0. Size of the category, i.e. maximum integer index + 1.
-      probs: Probability that a category is randomly flipped.
-      **kwargs: kwargs to parent class.
-    """
-    self.input_dim = input_dim
-    self.probs = probs
-    super(NCPCategoricalPerturb, self).__init__(**kwargs)
-
-  def call(self, inputs):
-    mask = tf.cast(tf.random.uniform(tf.shape(inputs)) <= self.probs,
-                   inputs.dtype)
-    flips = tf.random.uniform(
-        tf.shape(inputs), minval=0, maxval=self.input_dim, dtype=inputs.dtype)
-    flipped_inputs = mask * flips + (1 - mask) * inputs
-    return tf.concat([inputs, flipped_inputs], 0)
-
-
-class NCPNormalOutput(tf.keras.layers.Layer):
-  """Noise contrastive prior for continuous outputs (Hafner et al., 2018).
-
-  The layer returns the first half of the inputs' batch. It computes a KL
-  regularizer as a side-effect, which matches the inputs' second half towards a
-  normal distribution (the output prior), and averaged over the number of inputs
-  in the second half. This layer is typically in combination with an input prior
-  which doubles the batch. This layer implementation is inspired by the Aboleth
-  library.
-
-  The layer computes the exact KL divergence from a normal distribution to
-  the input RandomVariable. It is an unbiased estimate if the input
-  RandomVariable has random parameters. If the input is a Tensor, then it
-  assumes its density is `ed.Normal(input, 1.)`, i.e., mean squared error loss.
-
-  #### Examples
-
-  Below implements neural network regression with heteroskedastic noise,
-  noise contrastive priors, and being Bayesian only at the mean's output layer.
-
-  ```python
-  from tensor2tensor.layers import bayes
-
-  batch_size, dataset_size = 128, 1000
-  features, labels = get_some_dataset()
-
-  inputs = keras.Input(shape=(25,))
-  x = bayes.NCPNormalPerturb()(inputs)  # double input batch
-  x = layers.Dense(64, activation='relu')(x)
-  x = layers.Dense(64, activation='relu')(x)
-  means = bayes.DenseVariationalDropout(1, activation=None)(x)  # get mean dist.
-  means = bayes.NCPNormalOutput(labels)(means)  # halve input batch
-  stddevs = tf.keras.layers.Dense(1, activation='softplus')(x[:batch_size])
-  outputs = tf.keras.layers.Lambda(lambda x: ed.Normal(x[0], x[1]))([means,
-                                                                     stddevs])
-  model = tf.keras.Model(inputs=inputs, outputs=outputs)
-
-  predictions = model(features)
-  loss = tf.reduce_mean(predictions.distribution.log_prob(labels))
-  loss += model.losses[0] / dataset_size  # KL regularizer for output layer
-  loss += model.losses[-1]
-  train_op = tf.train.AdamOptimizer(0.1).minimize(loss)
-  ```
-
-  The network applies `bayes.NCPNormalPerturb()` to double the input batch
-  size and add Gaussian noise to the second half; then feedforward layers; then
-  `bayes.DenseVariational` to be Bayesian about the output density's mean; then
-  `bayes.NCPNormalOutput` centered at the labels to revert to the batch size
-  and compute a loss on the second half; then parameterize the output density's
-  standard deviations; then compute the total loss function as the sum of the
-  model's negative log-likelihood, KL divergence for the Bayesian mean layer,
-  and NCP loss.
-  """
-
-  def __init__(self, mean=0., stddev=1., **kwargs):
-    self.mean = mean
-    self.stddev = stddev
-    super(NCPNormalOutput, self).__init__(**kwargs)
-
-  def call(self, inputs):
-    if not isinstance(inputs, ed.RandomVariable):
-      # Default to a unit normal, i.e., derived from mean squared error loss.
-      inputs = ed.Normal(loc=inputs, scale=1.)
-    batch_size = tf.shape(inputs)[0] // 2
-    # TODO(trandustin): Depend on github's ed2 for indexing RVs. This is a hack.
-    # _, _ = inputs[:batch_size], inputs[batch_size:]
-    original_inputs = ed.RandomVariable(inputs.distribution[:batch_size],
-                                        value=inputs.value[:batch_size])
-    perturbed_inputs = ed.RandomVariable(inputs.distribution[batch_size:],
-                                         value=inputs.value[batch_size:])
-    loss = tf.reduce_sum(
-        tfp.distributions.Normal(self.mean, self.stddev).kl_divergence(
-            perturbed_inputs.distribution)) / tf.to_float(batch_size)
-    self.add_loss(loss)
-    return original_inputs
-
-
-class MixtureLogistic(tf.keras.layers.Layer):
-  """Stochastic output layer, distributed as a mixture of logistics."""
-
-  def __init__(self, num_components, **kwargs):
-    super(MixtureLogistic, self).__init__(**kwargs)
-    self.num_components = num_components
-    self.layer = tf.keras.layers.Dense(num_components * 3)
-
-  def build(self, input_shape=None):
-    self.layer.build(input_shape)
-    self.built = True
-
-  def call(self, inputs):
-    net = self.layer(inputs)
-    logits, loc, unconstrained_scale = tf.split(net, 3, axis=-1)
-    scale = tf.nn.softplus(unconstrained_scale) + tf.keras.backend.epsilon()
-    return ed.MixtureSameFamily(
-        mixture_distribution=ed.Categorical(logits=logits).distribution,
-        components_distribution=ed.Logistic(loc=loc, scale=scale).distribution)
-
-  def compute_output_shape(self, input_shape):
-    return tf.TensorShape(input_shape)[:-1]
-
-  def get_config(self):
-    config = {'num_components': self.num_components}
-    base_config = super(MixtureLogistic, self).get_config()
-    return dict(list(base_config.items()) + list(config.items()))
diff --git a/tensor2tensor/layers/bayes_test.py b/tensor2tensor/layers/bayes_test.py
deleted file mode 100644
index c123a0b19..000000000
--- a/tensor2tensor/layers/bayes_test.py
+++ /dev/null
@@ -1,626 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 Bayesian neural network layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import parameterized
-import numpy as np
-
-from tensor2tensor.layers import bayes
-from tensor2tensor.utils import test_utils
-
-import tensorflow as tf
-import tensorflow_probability as tfp
-ed = tfp.edward2
-tf.compat.v1.enable_eager_execution()
-
-
-class BayesTest(parameterized.TestCase, tf.test.TestCase):
-
-  @parameterized.parameters(
-      {"layer": bayes.Conv2DFlipout,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "zeros",
-       "all_close": True},
-      {"layer": bayes.Conv2DFlipout,
-       "kernel_initializer": "trainable_normal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"layer": bayes.Conv2DFlipout,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "trainable_normal",
-       "all_close": False},
-      {"layer": bayes.Conv2DHierarchical,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "zeros",
-       "all_close": True},
-      {"layer": bayes.Conv2DHierarchical,
-       "kernel_initializer": "trainable_normal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"layer": bayes.Conv2DHierarchical,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "trainable_normal",
-       "all_close": False},
-      {"layer": bayes.Conv2DReparameterization,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "zeros",
-       "all_close": True},
-      {"layer": bayes.Conv2DReparameterization,
-       "kernel_initializer": "trainable_normal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"layer": bayes.Conv2DReparameterization,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "trainable_normal",
-       "all_close": False},
-      {"layer": bayes.Conv2DVariationalDropout,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "zeros",
-       "all_close": True},
-      {"layer": bayes.Conv2DVariationalDropout,
-       "kernel_initializer": "trainable_normal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"layer": bayes.Conv2DVariationalDropout,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "trainable_normal",
-       "all_close": False},
-  )
-  @test_utils.run_in_graph_and_eager_modes
-  def testConv2DKernel(self,
-                       layer,
-                       kernel_initializer,
-                       bias_initializer,
-                       all_close):
-    tf.keras.backend.set_learning_phase(1)  # training time
-    inputs = tf.to_float(np.random.rand(5, 4, 4, 12))
-    model = layer(4,
-                  kernel_size=2,
-                  kernel_initializer=kernel_initializer,
-                  bias_initializer=bias_initializer,
-                  activation=tf.nn.relu)
-    outputs1 = model(inputs)
-    outputs2 = model(inputs)
-    self.evaluate(tf.global_variables_initializer())
-    res1, res2 = self.evaluate([outputs1, outputs2])
-    self.assertEqual(res1.shape, (5, 3, 3, 4))
-    self.assertAllGreaterEqual(res1, 0.)
-    if all_close:
-      self.assertAllClose(res1, res2)
-    else:
-      self.assertNotAllClose(res1, res2)
-    model.get_config()
-
-  @parameterized.parameters(
-      {"layer": bayes.Conv2DFlipout},
-      {"layer": bayes.Conv2DHierarchical},
-      {"layer": bayes.Conv2DReparameterization},
-      {"layer": bayes.Conv2DVariationalDropout},
-  )
-  @test_utils.run_in_graph_and_eager_modes()
-  def testConv2DModel(self, layer):
-    inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
-    model = tf.keras.Sequential([
-        layer(3, kernel_size=2, padding="SAME", activation=tf.nn.relu),
-        tf.keras.layers.Flatten(),
-        tf.keras.layers.Dense(2, activation=None),
-    ])
-    outputs = model(inputs, training=True)
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(outputs)
-    self.assertEqual(res.shape, (3, 2))
-    if layer == bayes.Conv2DHierarchical:
-      self.assertLen(model.losses, 3)
-    else:
-      self.assertLen(model.losses, 1)
-
-  @test_utils.run_in_graph_and_eager_modes
-  def testTrainableNormalStddevConstraint(self):
-    layer = bayes.DenseReparameterization(
-        100, kernel_initializer="trainable_normal")
-    inputs = tf.random_normal([1, 1])
-    out = layer(inputs)
-    stddev = layer.kernel.distribution.stddev()
-    self.evaluate(tf.global_variables_initializer())
-    res, _ = self.evaluate([stddev, out])
-    self.assertAllGreater(res, 0.)
-
-  @parameterized.parameters(
-      {"layer": bayes.DenseDVI,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "zeros",
-       "all_close": True},
-      {"layer": bayes.DenseDVI,
-       "kernel_initializer": "trainable_normal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"layer": bayes.DenseDVI,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "trainable_normal",
-       "all_close": False},
-      {"layer": bayes.DenseFlipout,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "zeros",
-       "all_close": True},
-      {"layer": bayes.DenseFlipout,
-       "kernel_initializer": "trainable_normal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"layer": bayes.DenseFlipout,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "trainable_normal",
-       "all_close": False},
-      {"layer": bayes.DenseReparameterization,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "zeros",
-       "all_close": True},
-      {"layer": bayes.DenseReparameterization,
-       "kernel_initializer": "trainable_normal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"layer": bayes.DenseReparameterization,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "trainable_normal",
-       "all_close": False},
-      {"layer": bayes.DenseVariationalDropout,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "zeros",
-       "all_close": True},
-      {"layer": bayes.DenseVariationalDropout,
-       "kernel_initializer": "trainable_normal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"layer": bayes.DenseVariationalDropout,
-       "kernel_initializer": "zeros",
-       "bias_initializer": "trainable_normal",
-       "all_close": False},
-  )
-  @test_utils.run_in_graph_and_eager_modes
-  def testDenseKernel(self,
-                      layer,
-                      kernel_initializer,
-                      bias_initializer,
-                      all_close):
-    tf.keras.backend.set_learning_phase(1)  # training time
-    inputs = tf.to_float(np.random.rand(5, 3, 12))
-    model = layer(4,
-                  kernel_initializer=kernel_initializer,
-                  bias_initializer=bias_initializer,
-                  activation=tf.nn.relu)
-    outputs1 = model(inputs)
-    outputs2 = model(inputs)
-    self.evaluate(tf.global_variables_initializer())
-    res1, res2 = self.evaluate([outputs1, outputs2])
-    self.assertEqual(res1.shape, (5, 3, 4))
-    if layer != bayes.DenseDVI:
-      self.assertAllGreaterEqual(res1, 0.)
-    if all_close:
-      self.assertAllClose(res1, res2)
-    else:
-      self.assertNotAllClose(res1, res2)
-    model.get_config()
-
-  @parameterized.parameters(
-      {"layer": bayes.DenseDVI},
-      {"layer": bayes.DenseFlipout},
-      {"layer": bayes.DenseReparameterization},
-      {"layer": bayes.DenseVariationalDropout},
-  )
-  @test_utils.run_in_graph_and_eager_modes
-  def testDenseMean(self, layer):
-    """Tests that forward pass can use other values, e.g., posterior mean."""
-    tf.keras.backend.set_learning_phase(0)  # test time
-    def take_mean(f, *args, **kwargs):
-      """Sets random variable value to its mean."""
-      rv = f(*args, **kwargs)
-      rv._value = rv.distribution.mean()
-      return rv
-    inputs = tf.to_float(np.random.rand(5, 3, 7))
-    model = layer(4, activation=tf.nn.relu, use_bias=False)
-    outputs1 = model(inputs)
-    with tfp.edward2.interception(take_mean):
-      outputs2 = model(inputs)
-    self.evaluate(tf.global_variables_initializer())
-    res1, res2 = self.evaluate([outputs1, outputs2])
-    self.assertEqual(res1.shape, (5, 3, 4))
-    self.assertNotAllClose(res1, res2)
-    if layer != bayes.DenseDVI:
-      self.assertAllClose(res2, np.zeros((5, 3, 4)), atol=1e-4)
-
-  @parameterized.parameters(
-      {"layer": bayes.DenseDVI},
-      {"layer": bayes.DenseFlipout},
-      {"layer": bayes.DenseReparameterization},
-      {"layer": bayes.DenseVariationalDropout},
-      {"layer": bayes.DenseHierarchical},
-  )
-  @test_utils.run_in_graph_and_eager_modes()
-  def testDenseLoss(self, layer):
-    tf.keras.backend.set_learning_phase(1)  # training time
-    features = tf.to_float(np.random.rand(5, 12))
-    labels = tf.to_float(np.random.rand(5, 10))
-    model = layer(10)
-
-    # Imagine this is the 1st epoch.
-    with tf.GradientTape(persistent=True) as tape:
-      predictions = model(features)  # first call forces build
-      model(features)  # ensure robustness after multiple calls
-      nll = tf.losses.mean_squared_error(labels, predictions)
-      kl = sum(model.losses)
-
-    variables = [model.kernel_initializer.mean, model.kernel_initializer.stddev]
-    for v in variables:
-      self.assertIn(v, model.variables)
-
-    # This will be fine, since the layer was built inside this tape, and thus
-    # the distribution init ops were inside this tape.
-    grads = tape.gradient(nll, variables)
-    for grad in grads:
-      self.assertIsNotNone(grad)
-    grads = tape.gradient(kl, variables)
-    for grad in grads:
-      self.assertIsNotNone(grad)
-
-    # Imagine this is the 2nd epoch.
-    with tf.GradientTape(persistent=True) as tape:
-      predictions = model(features)  # build is not called
-      nll = tf.losses.mean_squared_error(labels, predictions)
-      kl = sum(model.losses)
-
-    variables = [model.kernel_initializer.mean, model.kernel_initializer.stddev]
-    for v in variables:
-      self.assertIn(v, model.variables)
-
-    # This would fail, since the layer was built inside the tape from the 1st
-    # epoch, and thus the distribution init ops were inside that tape instead of
-    # this tape. By using a callable for the variable, this will no longer fail.
-    grads = tape.gradient(nll, variables)
-    for grad in grads:
-      self.assertIsNotNone(grad)
-    grads = tape.gradient(kl, variables)
-    for grad in grads:
-      self.assertIsNotNone(grad)
-
-  @parameterized.parameters(
-      {"layer": bayes.DenseDVI},
-      {"layer": bayes.DenseFlipout},
-      {"layer": bayes.DenseReparameterization},
-      {"layer": bayes.DenseVariationalDropout},
-      {"layer": bayes.DenseHierarchical},
-  )
-  @test_utils.run_in_graph_and_eager_modes()
-  def testDenseModel(self, layer):
-    inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
-    model = tf.keras.Sequential([
-        tf.keras.layers.Conv2D(3,
-                               kernel_size=2,
-                               padding="SAME",
-                               activation=tf.nn.relu),
-        tf.keras.layers.Flatten(),
-        layer(2, activation=None),
-    ])
-    outputs = model(inputs, training=True)
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(outputs)
-    self.assertEqual(res.shape, (3, 2))
-    if layer == bayes.DenseHierarchical:
-      self.assertLen(model.losses, 3)
-    else:
-      self.assertLen(model.losses, 1)
-
-  @parameterized.parameters(
-      {"layer": bayes.DenseDVI},
-      {"layer": bayes.DenseFlipout},
-      {"layer": bayes.DenseReparameterization},
-      {"layer": bayes.DenseVariationalDropout},
-      {"layer": bayes.DenseHierarchical},
-  )
-  @test_utils.run_in_graph_and_eager_modes()
-  def testDenseSubclass(self, layer):
-    class DenseSubclass(layer):
-      pass
-
-    inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
-    model = tf.keras.Sequential([
-        tf.keras.layers.Conv2D(3,
-                               kernel_size=2,
-                               padding="SAME",
-                               activation=tf.nn.relu),
-        tf.keras.layers.Flatten(),
-        DenseSubclass(2, activation=None),
-    ])
-    outputs = model(inputs, training=True)
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(outputs)
-    self.assertEqual(res.shape, (3, 2))
-    if layer == bayes.DenseHierarchical:
-      self.assertLen(model.losses, 3)
-    else:
-      self.assertLen(model.losses, 1)
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testDenseDVIIsDeterministic(self):
-    """Tests that DenseDVI network has a deterministic loss function."""
-    features = tf.to_float(np.random.rand(3, 2))
-    labels = tf.to_float(np.random.rand(3, 1))
-    model = tf.keras.Sequential([
-        bayes.DenseDVI(5, activation=tf.nn.relu),
-        bayes.DenseDVI(1, activation=None),
-    ])
-    outputs = model(features, training=True)
-    nll = -tf.reduce_sum(outputs.distribution.log_prob(labels))
-    kl = sum(model.losses)
-    loss = nll + kl
-    self.evaluate(tf.global_variables_initializer())
-    res1 = self.evaluate(loss)
-    res2 = self.evaluate(loss)
-    self.assertEqual(res1, res2)
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testDenseDVIMoments(self):
-    """Verifies DenseDVI's moments empirically with samples."""
-    tf.set_random_seed(377269)
-    batch_size = 3
-    num_features = 5
-    units = 128
-    num_samples = 50000
-    inputs = tf.to_float(np.random.rand(batch_size, num_features))
-    layer = bayes.DenseDVI(units, activation=tf.nn.relu)
-
-    outputs1 = layer(inputs)
-    mean1 = outputs1.distribution.mean()
-    covariance1 = outputs1.distribution.covariance()
-
-    kernel_samples = layer.kernel.distribution.sample(num_samples)
-    outputs2 = layer.activation(
-        tf.einsum("bd,sdu->sbu", inputs, kernel_samples) +
-        tf.reshape(layer.bias, [1, 1, units]))
-    mean2 = tf.reduce_mean(outputs2, axis=0)
-    centered_outputs2 = tf.transpose(outputs2 - mean2, [1, 2, 0])
-    covariance2 = tf.matmul(centered_outputs2,
-                            centered_outputs2,
-                            transpose_b=True) / float(num_samples)
-
-    self.evaluate(tf.global_variables_initializer())
-    mean1_val, covariance1_val, mean2_val, covariance2_val = self.evaluate(
-        [mean1, covariance1, mean2, covariance2])
-    # Check % of mismatches is not too high according to heuristic thresholds.
-    num_mismatches = np.sum(np.abs(mean1_val - mean2_val) > 5e-3)
-    percent_mismatches = num_mismatches / float(batch_size * units)
-    self.assertLessEqual(percent_mismatches, 0.05)
-    num_mismatches = np.sum(np.abs(covariance1_val - covariance2_val) > 5e-3)
-    percent_mismatches = num_mismatches / float(batch_size * units * units)
-    self.assertLessEqual(percent_mismatches, 0.05)
-
-  @parameterized.parameters(
-      {"lstm_cell": bayes.LSTMCellFlipout,
-       "kernel_initializer": "zeros",
-       "recurrent_initializer": "orthogonal",
-       "bias_initializer": "zeros",
-       "all_close": True},
-      {"lstm_cell": bayes.LSTMCellFlipout,
-       "kernel_initializer": "trainable_normal",
-       "recurrent_initializer": "orthogonal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"lstm_cell": bayes.LSTMCellFlipout,
-       "kernel_initializer": "zeros",
-       "recurrent_initializer": "orthogonal",
-       "bias_initializer": "trainable_normal",
-       "all_close": False},
-      {"lstm_cell": bayes.LSTMCellReparameterization,
-       "kernel_initializer": "zeros",
-       "recurrent_initializer": "orthogonal",
-       "bias_initializer": "zeros",
-       "all_close": True},
-      {"lstm_cell": bayes.LSTMCellReparameterization,
-       "kernel_initializer": "trainable_normal",
-       "recurrent_initializer": "orthogonal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"lstm_cell": bayes.LSTMCellReparameterization,
-       "kernel_initializer": "zeros",
-       "recurrent_initializer": "trainable_normal",
-       "bias_initializer": "zeros",
-       "all_close": False},
-      {"lstm_cell": bayes.LSTMCellReparameterization,
-       "kernel_initializer": "zeros",
-       "recurrent_initializer": "orthogonal",
-       "bias_initializer": "trainable_normal",
-       "all_close": False},
-  )
-  @test_utils.run_in_graph_and_eager_modes
-  def testLSTMCell(self,
-                   lstm_cell,
-                   kernel_initializer,
-                   recurrent_initializer,
-                   bias_initializer,
-                   all_close):
-    batch_size, timesteps, dim = 5, 3, 12
-    hidden_size = 10
-    inputs = tf.to_float(np.random.rand(batch_size, timesteps, dim))
-    cell = lstm_cell(hidden_size,
-                     kernel_initializer=kernel_initializer,
-                     recurrent_initializer=recurrent_initializer,
-                     bias_initializer=bias_initializer)
-    noise = tf.to_float(np.random.rand(1, hidden_size))
-    h0, c0 = cell.get_initial_state(inputs)
-    state = (h0 + noise, c0)
-    outputs1, _ = cell(inputs[:, 0, :], state)
-    outputs2, _ = cell(inputs[:, 0, :], state)
-    cell.call_weights()
-    outputs3, _ = cell(inputs[:, 0, :], state)
-    self.evaluate(tf.global_variables_initializer())
-    res1, res2, res3 = self.evaluate([outputs1, outputs2, outputs3])
-    self.assertEqual(res1.shape, (batch_size, hidden_size))
-    self.assertAllClose(res1, res2)
-    if all_close:
-      self.assertAllClose(res1, res3)
-    else:
-      self.assertNotAllClose(res1, res3)
-    cell.get_config()
-
-  @parameterized.parameters(
-      {"lstm_cell": bayes.LSTMCellFlipout},
-      {"lstm_cell": bayes.LSTMCellReparameterization},
-  )
-  @test_utils.run_in_graph_and_eager_modes()
-  def testLSTMCellLoss(self, lstm_cell):
-    features = tf.to_float(np.random.rand(5, 1, 12))
-    labels = tf.to_float(np.random.rand(5, 10))
-    cell = lstm_cell(10)
-    state = (tf.zeros([1, 10]), tf.zeros([1, 10]))
-
-    # Imagine this is the 1st epoch.
-    with tf.GradientTape(persistent=True) as tape:
-      predictions, _ = cell(features[:, 0, :], state)  # first call forces build
-      cell(features[:, 0, :], state)  # ensure robustness after multiple calls
-      cell.get_initial_state(features[:, 0, :])
-      cell(features[:, 0, :], state)  # ensure robustness after multiple calls
-      nll = tf.losses.mean_squared_error(labels, predictions)
-      kl = sum(cell.losses)
-
-    variables = [
-        cell.kernel_initializer.mean, cell.kernel_initializer.stddev,
-        cell.recurrent_initializer.mean, cell.recurrent_initializer.stddev,
-    ]
-    for v in variables:
-      self.assertIn(v, cell.variables)
-
-    # This will be fine, since the layer was built inside this tape, and thus
-    # the distribution init ops were inside this tape.
-    grads = tape.gradient(nll, variables)
-    for grad in grads:
-      self.assertIsNotNone(grad)
-    grads = tape.gradient(kl, variables)
-    for grad in grads:
-      self.assertIsNotNone(grad)
-
-    # Imagine this is the 2nd epoch.
-    with tf.GradientTape(persistent=True) as tape:
-      cell.get_initial_state(features[:, 0, :])
-      predictions, _ = cell(features[:, 0, :], state)  # build is not called
-      nll = tf.losses.mean_squared_error(labels, predictions)
-      kl = sum(cell.losses)
-
-    variables = [
-        cell.kernel_initializer.mean, cell.kernel_initializer.stddev,
-        cell.recurrent_initializer.mean, cell.recurrent_initializer.stddev,
-    ]
-    for v in variables:
-      self.assertIn(v, cell.variables)
-
-    # This would fail, since the layer was built inside the tape from the 1st
-    # epoch, and thus the distribution init ops were inside that tape instead of
-    # this tape. By using a callable for the variable, this will no longer fail.
-    grads = tape.gradient(nll, variables)
-    for grad in grads:
-      self.assertIsNotNone(grad)
-    grads = tape.gradient(kl, variables)
-    for grad in grads:
-      self.assertIsNotNone(grad)
-
-  @parameterized.parameters(
-      {"lstm_cell": bayes.LSTMCellFlipout},
-      {"lstm_cell": bayes.LSTMCellReparameterization},
-  )
-  @test_utils.run_in_graph_and_eager_modes()
-  def testLSTMCellModel(self, lstm_cell):
-    batch_size, timesteps, dim = 5, 3, 12
-    hidden_size = 10
-    inputs = tf.to_float(np.random.rand(batch_size, timesteps, dim))
-    cell = lstm_cell(hidden_size)
-    model = tf.keras.Sequential([
-        tf.keras.layers.RNN(cell, return_sequences=True)
-    ])
-    outputs1 = model(inputs)
-    outputs2 = model(inputs)
-    state = (tf.zeros([1, hidden_size]), tf.zeros([1, hidden_size]))
-    outputs3 = []
-    for t in range(timesteps):
-      out, state = cell(inputs[:, t, :], state)
-      outputs3.append(out)
-    outputs3 = tf.stack(outputs3, axis=1)
-    self.evaluate(tf.global_variables_initializer())
-    res1, res2, res3 = self.evaluate([outputs1, outputs2, outputs3])
-    self.assertEqual(res1.shape, (batch_size, timesteps, hidden_size))
-    self.assertEqual(res3.shape, (batch_size, timesteps, hidden_size))
-    # NOTE: `cell.call_weights` should have been called at the beginning of
-    # each call, so these should be different.
-    self.assertNotAllClose(res1, res2)
-    # NOTE: We didn't call `cell.call_weights` again before computing
-    # `outputs3`, so the cell should have had the same weights as it did
-    # during computation of `outputs2`, and thus yielded the same output
-    # tensor.
-    self.assertAllClose(res2, res3)
-    self.assertLen(model.losses, 2)
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testNCPNormalPerturb(self):
-    batch_size = 3
-    inputs = tf.to_float(np.random.rand(batch_size, 4))
-    model = bayes.NCPNormalPerturb()
-    outputs = model(inputs)
-    inputs_val, outputs_val = self.evaluate([inputs, outputs])
-    self.assertEqual(outputs_val.shape, (2 * batch_size, 4))
-    self.assertAllEqual(inputs_val, outputs_val[:batch_size])
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testNCPCategoricalPerturb(self):
-    input_dim = 5
-    batch_size = 3
-    inputs = tf.to_float(np.random.choice(input_dim, size=(batch_size, 4)))
-    model = bayes.NCPCategoricalPerturb(input_dim)
-    outputs = model(inputs)
-    inputs_val, outputs_val = self.evaluate([inputs, outputs])
-    self.assertEqual(outputs_val.shape, (2 * batch_size, 4))
-    self.assertAllEqual(inputs_val, outputs_val[:batch_size])
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testNCPNormalOutput(self):
-    batch_size = 3
-    features = ed.Normal(loc=tf.random.normal([2 * batch_size, 1]), scale=1.)
-    labels = tf.to_float(np.random.rand(batch_size))
-    model = bayes.NCPNormalOutput(mean=labels)
-    predictions = model(features)
-    features_val, predictions_val = self.evaluate([features, predictions])
-    self.assertLen(model.losses, 1)
-    self.assertAllEqual(features_val[:batch_size], predictions_val)
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testMixtureLogistic(self):
-    batch_size = 3
-    features = tf.to_float(np.random.rand(batch_size, 4))
-    labels = tf.to_float(np.random.rand(batch_size))
-    model = tf.keras.Sequential([
-        tf.keras.layers.Dense(2, activation=None),
-        bayes.MixtureLogistic(5),
-    ])
-    outputs = model(features)
-    log_likelihood = tf.reduce_sum(outputs.distribution.log_prob(labels))
-    self.evaluate(tf.global_variables_initializer())
-    log_likelihood_val, outputs_val = self.evaluate([log_likelihood, outputs])
-    self.assertEqual(log_likelihood_val.shape, ())
-    self.assertLessEqual(log_likelihood_val, 0.)
-    self.assertEqual(outputs_val.shape, (batch_size,))
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/layers/gaussian_process.py b/tensor2tensor/layers/gaussian_process.py
deleted file mode 100644
index dc42cf6ae..000000000
--- a/tensor2tensor/layers/gaussian_process.py
+++ /dev/null
@@ -1,815 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Gaussian process layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.keras import constraints
-from tensor2tensor.keras import initializers
-from tensor2tensor.keras import regularizers
-from tensor2tensor.layers import bayes
-
-import tensorflow as tf
-import tensorflow_probability as tfp
-from tensorflow_probability import edward2 as ed
-
-
-class Zeros(object):
-  """Function returning zeros tensor of same shape excluding the last dim."""
-
-  def __call__(self, inputs):
-    return tf.zeros(tf.shape(inputs)[:-1], inputs.dtype)
-
-  def get_config(self):
-    return {}
-
-
-class ExponentiatedQuadratic(object):
-  """Exponentiated quadratic kernel."""
-
-  def __init__(self, variance, lengthscale):
-    self.variance = variance
-    self.lengthscale = lengthscale
-
-  def __call__(self, x1, x2):
-    """Computes exponentiated quadratic over all pairs of inputs.
-
-    Args:
-      x1: Tensor of shape [batch_x1, ...]. Slices along the batch axis denote an
-        individual input to be passed to the kernel. It is computed pairwise
-        with each input sliced from x2.
-      x2: Tensor of shape [batch_x2, ...]. Slices along the batch axis denote an
-        individual input passed to the kernel function. It is computed pairwise
-        with each input sliced from x1.
-
-    Returns:
-      Tensor of shape [batch_x1, batch_x2].
-    """
-    size = tf.convert_to_tensor(x1).shape.ndims
-    if size > 2:
-      raise NotImplementedError('Multiple feature dimensions is not yet '
-                                'supported.')
-    x1 = x1 / self.lengthscale
-    x2 = x2 / self.lengthscale
-    x1_squared = tf.reduce_sum(tf.square(x1), list(range(1, len(x1.shape))))
-    x2_squared = tf.reduce_sum(tf.square(x2), list(range(1, len(x2.shape))))
-    square = (x1_squared[:, tf.newaxis] +
-              x2_squared[tf.newaxis, :] -
-              2 * tf.matmul(x1, x2, transpose_b=True))
-    return self.variance * tf.exp(-square / 2)
-
-  def get_config(self):
-    return {'variance': self.variance, 'lengthscale': self.lengthscale}
-
-
-class LinearKernel(object):
-  """Linear kernel, optionally on top of a feature extractor (e.g., encoder)."""
-
-  def __init__(self, variance, bias, encoder=tf.identity):
-    self.variance = variance
-    self.bias = bias
-    self.encoder = encoder
-
-  def __call__(self, x1, x2):
-    """Computes scaled dot product of over all pairs of encoded inputs.
-
-    Args:
-      x1: Tensor of shape [batch_x1] + encoder domain. Slices along the batch
-        axis denote an individual input to be passed to the kernel. It is
-        computed pairwise with each input sliced from x2.
-      x2: Tensor of shape [batch_x2] + encoder domain. Slices along the batch
-        axis denote an individual input to be passed to the kernel. It is
-        computed pairwise with each input sliced from x1.
-
-    Returns:
-      Tensor of shape [batch_x1, batch_x2].
-    """
-    encoded_x1 = self.encoder(x1)
-    encoded_x2 = self.encoder(x2)
-    dot_product = tf.matmul(encoded_x1, encoded_x2, transpose_b=True)
-    return self.variance * dot_product + self.bias
-
-  def get_config(self):
-    return {
-        'variance': self.variance,
-        'bias': self.bias,
-        'encoder': tf.keras.utils.serialize_keras_object(self.encoder),
-    }
-
-
-class GaussianProcess(tf.keras.layers.Layer):
-  r"""Gaussian process layer.
-
-  The layer represents a distribution over functions, where a
-  stochastic forward pass appears as
-
-  ```none
-  f ~ GP(f | conditional_inputs, conditional_outputs; mean_fn, covariance_fn)
-  outputs = f(inputs)
-  ```
-
-  The optional arguments `conditional_inputs` and `conditional_outputs`
-  capture data that the GP "memorizes", i.e., it forms a posterior predictive
-  distribution. If left unspecified, the GP posits a prior predictive.
-
-  Given a call to `inputs`, an equivalent formulation in terms of function
-  outputs is
-
-  ```none
-  outputs ~ \prod_{unit=1}^{units} MultivariateNormal(output[:, unit] |
-      mean = mean_fn(inputs) + Knm Kmm^{-1} (conditional_outputs[:, unit]-mean),
-      covariance = Knn - Knm Kmm^{-1} Kmn)
-  ```
-
-  where Knm is the covariance function evaluated between all `inputs` and
-  `conditional_inputs`; Knn is between all `inputs`; Kmm is between all
-  `conditional_inputs`; and mean is the mean function evaluated on
-  `conditional_inputs`. The multivariate normal is correlated across input
-  dimensions and is independent across output dimensions.
-  """
-
-  def __init__(
-      self,
-      units,
-      mean_fn=Zeros(),
-      covariance_fn=ExponentiatedQuadratic(variance=1., lengthscale=1.),
-      conditional_inputs=None,
-      conditional_outputs=None,
-      **kwargs):
-    """Constructs layer.
-
-    Args:
-      units: integer, dimensionality of layer.
-      mean_fn: Mean function, a callable taking an inputs Tensor of shape
-        [batch, ...] and returning a Tensor of shape [batch].
-      covariance_fn: Covariance function, a callable taking two input Tensors
-        of shape [batch_x1, ...] and [batch_x2, ...] respectively, and returning
-        a positive semi-definite matrix of shape [batch_x1, batch_x2].
-      conditional_inputs: Tensor of shape [batch, ...], where batch must be the
-        same as conditional_outputs', and ellipses must match layer inputs.
-      conditional_outputs: Tensor of shape [batch, units], where batch must be
-        the same as conditional_inputs' and units is the layer's units size.
-      **kwargs: kwargs passed to parent class.
-    """
-    super(GaussianProcess, self).__init__(**kwargs)
-    self.units = int(units)
-    self.mean_fn = mean_fn
-    self.covariance_fn = covariance_fn
-    self.conditional_inputs = conditional_inputs
-    self.conditional_outputs = conditional_outputs
-
-    self.supports_masking = True
-    self.input_spec = tf.keras.layers.InputSpec(min_ndim=2)
-
-  def build(self, input_shape=None):
-    # Don't track trainable variables such as in the kernel. The user should
-    # refer to any via, e.g., self.covariance_fn or the user environment.
-    self.built = True
-
-  def call(self, inputs):
-    if self.conditional_inputs is None and self.conditional_outputs is None:
-      covariance_matrix = self.covariance_fn(inputs, inputs)
-      # Tile locations so output has shape [units, batch_size]. Covariance will
-      # broadcast to [units, batch_size, batch_size], and we perform
-      # shape manipulations to get a random variable over [batch_size, units].
-      loc = self.mean_fn(inputs)
-      loc = tf.tile(loc[tf.newaxis], [self.units] + [1] * len(loc.shape))
-    else:
-      knn = self.covariance_fn(inputs, inputs)
-      knm = self.covariance_fn(inputs, self.conditional_inputs)
-      kmm = self.covariance_fn(self.conditional_inputs, self.conditional_inputs)
-      kmm = tf.matrix_set_diag(
-          kmm, tf.matrix_diag_part(kmm) + tf.keras.backend.epsilon())
-      kmm_tril = tf.linalg.cholesky(kmm)
-      kmm_tril_operator = tf.linalg.LinearOperatorLowerTriangular(kmm_tril)
-      knm_operator = tf.linalg.LinearOperatorFullMatrix(knm)
-
-      # TODO(trandustin): Vectorize linear algebra for multiple outputs. For
-      # now, we do each separately and stack to obtain a locations Tensor of
-      # shape [units, batch_size].
-      loc = []
-      for conditional_outputs_unit in tf.unstack(self.conditional_outputs,
-                                                 axis=-1):
-        center = conditional_outputs_unit - self.mean_fn(
-            self.conditional_inputs)
-        loc_unit = knm_operator.matvec(
-            kmm_tril_operator.solvevec(kmm_tril_operator.solvevec(center),
-                                       adjoint=True))
-        loc.append(loc_unit)
-      loc = tf.stack(loc) + self.mean_fn(inputs)[tf.newaxis]
-
-      covariance_matrix = knn
-      covariance_matrix -= knm_operator.matmul(
-          kmm_tril_operator.solve(
-              kmm_tril_operator.solve(knm, adjoint_arg=True), adjoint=True))
-
-    covariance_matrix = tf.matrix_set_diag(
-        covariance_matrix,
-        tf.matrix_diag_part(covariance_matrix) + tf.keras.backend.epsilon())
-
-    # Form a multivariate normal random variable with batch_shape units and
-    # event_shape batch_size. Then make it be independent across the units
-    # dimension. Then transpose its dimensions so it is [batch_size, units].
-    random_variable = ed.MultivariateNormalFullCovariance(
-        loc=loc, covariance_matrix=covariance_matrix)
-    random_variable = ed.Independent(random_variable.distribution,
-                                     reinterpreted_batch_ndims=1)
-    bijector = tfp.bijectors.Inline(
-        forward_fn=lambda x: tf.transpose(x, [1, 0]),
-        inverse_fn=lambda y: tf.transpose(y, [1, 0]),
-        forward_event_shape_fn=lambda input_shape: input_shape[::-1],
-        forward_event_shape_tensor_fn=lambda input_shape: input_shape[::-1],
-        inverse_log_det_jacobian_fn=lambda y: tf.cast(0, y.dtype),
-        forward_min_event_ndims=2)
-    random_variable = ed.TransformedDistribution(random_variable.distribution,
-                                                 bijector=bijector)
-    return random_variable
-
-  def compute_output_shape(self, input_shape):
-    input_shape = tf.TensorShape(input_shape)
-    input_shape = input_shape.with_rank_at_least(2)
-    input_dim = input_shape[-1]
-    if isinstance(input_dim, tf.Dimension):
-      input_dim = input_dim.value
-    if input_dim is None:
-      raise ValueError(
-          'The innermost dimension of input_shape must be defined, but saw: %s'
-          % input_shape)
-    return input_shape[:-1].concatenate(self.units)
-
-  def get_config(self):
-    config = {
-        'units': self.units,
-        'mean_fn': tf.keras.utils.serialize_keras_object(self.mean_fn),
-        'covariance_fn': tf.keras.utils.serialize_keras_object(
-            self.covariance_fn),
-        'conditional_inputs': None,  # don't serialize as it can be large
-        'conditional_outputs': None,  # don't serialize as it can be large
-    }
-    base_config = super(GaussianProcess, self).get_config()
-    return dict(list(base_config.items()) + list(config.items()))
-
-
-@bayes.add_weight
-class SparseGaussianProcess(GaussianProcess):
-  r"""Gaussian process layer with inducing input and output variables.
-
-  The layer represents a distribution over functions, where a
-  stochastic forward pass appears as
-
-  ```none
-  f ~ GP(f | inducing_inputs, inducing_outputs; mean_fn, covariance_fn)
-  outputs = f(inputs)
-  ```
-
-  The arguments `inducing_inputs` and `inducing_outputs`
-  capture data that the GP "memorizes", i.e., it forms a posterior predictive
-  distribution. Typically in a variational inference scheme (and by default),
-  the inducing outputs are normally distributed with learnable location and
-  scale parameters, and the inducing inputs are learnable parameters.
-
-  Given a call to `inputs` with these defaults, an equivalent formulation in
-  terms of function outputs is
-
-  ```none
-  inducing_outputs ~ Normal(inducing_outputs | mean, stddev)
-  outputs ~ \prod_{unit=1}^{units} MultivariateNormal(output[:, unit] |
-      mean = mean_fn(inputs) + Knm Kmm^{-1} (inducing_outputs[:, unit]-mean),
-      covariance = Knn - Knm Kmm^{-1} Kmn)
-  ```
-
-  where Knm is the covariance function evaluated between all `inputs` and
-  `inducing_inputs`; Knn is between all `inputs`; Kmm is between all
-  `inducing_inputs`; and mean is the mean function evaluated on
-  `inducing_inputs`. The multivariate normal is correlated across input
-  dimensions and is independent across output dimensions.
-
-  #### Examples
-
-  We demonstrate a three-layer deep GP with variational inference (Salimbeni and
-  Deisenroth, 2017; Damianou and Lawrence, 2013). The code snippet mirrors
-  Figure 5 of Bayesian Layers. We apply it for regression given batches of
-  spatial inputs and vector-valued outputs. We flatten inputs to use the
-  default squared exponential kernel; this naturally extends to pass in a
-  more sophisticated kernel function.
-
-  ```python
-  from tensor2tensor.layers import bayes
-
-  batch_size = 256
-  dataset_size = 10000
-  features, labels = load_spatial_data(batch_size)
-
-  model = tf.keras.Sequential([
-    tf.keras.layers.Flatten(),
-    layers.SparseGaussianProcess(256, num_inducing=512),
-    layers.SparseGaussianProcess(256, num_inducing=512),
-    layers.SparseGaussianProcess(10, num_inducing=512),
-  ])
-  predictions = model(features)
-  nll = tf.losses.mean_squared_error(labels=labels, predictions=predictions)
-  kl = sum(model.losses) / dataset_size
-  loss = nll + kl
-  train_op = tf.train.AdamOptimizer().minimize(loss)
-  ```
-  """
-
-  def __init__(
-      self,
-      units,
-      num_inducing,
-      mean_fn=Zeros(),
-      covariance_fn=ExponentiatedQuadratic(variance=1., lengthscale=1.),
-      inducing_inputs_initializer='random_normal',
-      inducing_outputs_initializer='trainable_normal',
-      inducing_inputs_regularizer=None,
-      inducing_outputs_regularizer='normal_kl_divergence',
-      inducing_inputs_constraint=None,
-      inducing_outputs_constraint=None,
-      **kwargs):
-    """Constructs layer.
-
-    Args:
-      units: integer, dimensionality of layer.
-      num_inducing: integer, number of inducing points for the approximation.
-      mean_fn: Mean function, a callable taking an inputs Tensor of shape
-        [batch, ...] and returning a Tensor of shape [batch].
-      covariance_fn: Covariance function, a callable taking two input Tensors
-        of shape [batch_x1, ...] and [batch_x2, ...] respectively, and returning
-        a positive semi-definite matrix of shape [batch_x1, batch_x2].
-      inducing_inputs_initializer: Initializer for the inducing inputs.
-      inducing_outputs_initializer: Initializer for the inducing outputs.
-      inducing_inputs_regularizer: Regularizer function applied to the inducing
-        inputs.
-      inducing_outputs_regularizer: Regularizer function applied to the inducing
-        outputs.
-      inducing_inputs_constraint: Constraint function applied to the inducing
-        inputs.
-      inducing_outputs_constraint: Constraint function applied to the inducing
-        outputs.
-      **kwargs: kwargs passed to parent class.
-    """
-    super(SparseGaussianProcess, self).__init__(
-        units=units,
-        mean_fn=mean_fn,
-        covariance_fn=covariance_fn,
-        conditional_inputs=None,
-        conditional_outputs=None,
-        **kwargs)
-    self.num_inducing = num_inducing
-    self.inducing_inputs_initializer = initializers.get(
-        inducing_inputs_initializer)
-    self.inducing_outputs_initializer = initializers.get(
-        inducing_outputs_initializer)
-    self.inducing_inputs_regularizer = regularizers.get(
-        inducing_inputs_regularizer)
-    self.inducing_outputs_regularizer = regularizers.get(
-        inducing_outputs_regularizer)
-    self.inducing_inputs_constraint = constraints.get(
-        inducing_inputs_constraint)
-    self.inducing_outputs_constraint = constraints.get(
-        inducing_outputs_constraint)
-
-  def build(self, input_shape=None):
-    input_shape = tf.TensorShape(input_shape)
-    input_dim = input_shape[-1]
-    if isinstance(input_dim, tf.Dimension):
-      input_dim = input_dim.value
-    self.conditional_inputs = self.add_weight(
-        shape=(self.num_inducing, input_dim),
-        name='inducing_inputs',
-        initializer=self.inducing_inputs_initializer,
-        regularizer=self.inducing_inputs_regularizer,
-        constraint=self.inducing_inputs_constraint)
-    self.conditional_outputs = self.add_weight(
-        shape=(self.num_inducing, self.units),
-        name='inducing_outputs',
-        initializer=self.inducing_outputs_initializer,
-        regularizer=self.inducing_outputs_regularizer,
-        constraint=self.inducing_outputs_constraint)
-    super(SparseGaussianProcess, self).build(input_shape)
-
-
-class BayesianLinearModel(tf.keras.Model):
-  r"""Bayesian linear model with standard normal prior over its coefficients.
-
-  A forward pass computes the mean of the exact predictive distribution
-
-  ```none
-  p(outputs | inputs) = \int Normal(outputs | coeffs * inputs, noise_variance)
-                             Normal(coeffs | 0, 1) dweights dbias.
-  ```
-
-  It takes a Tensor of shape [batch_size, input_dim] as input and returns a
-  Normal random variable of shape [batch_size] representing its outputs.
-  After `fit()`, the forward pass computes the exact posterior predictive
-  distribution.
-  """
-
-  def __init__(self, noise_variance, **kwargs):
-    super(BayesianLinearModel, self).__init__(**kwargs)
-    self.noise_variance = noise_variance
-    self.coeffs_precision_tril_op = None
-    self.coeffs_mean = None
-
-  def call(self, inputs):
-    if self.coeffs_mean is None and self.coeffs_precision_tril_op is None:
-      # p(mean(ynew) | xnew) = Normal(ynew | mean = 0, variance = xnew xnew^T)
-      predictive_mean = 0.
-      predictive_variance = tf.reduce_sum(tf.square(inputs), -1)
-    else:
-      # p(mean(ynew) | xnew, x, y) = Normal(ynew |
-      #   mean = xnew (1/noise_variance) (1/noise_variance x^T x + I)^{-1}x^T y,
-      #   variance = xnew (1/noise_variance x^T x + I)^{-1} xnew^T)
-      predictive_mean = tf.einsum('nm,m->n', inputs, self.coeffs_mean)
-      predictive_covariance = tf.matmul(
-          inputs,
-          self.coeffs_precision_tril_op.solve(
-              self.coeffs_precision_tril_op.solve(inputs, adjoint_arg=True),
-              adjoint=True))
-      predictive_variance = tf.diag_part(predictive_covariance)
-    return ed.Normal(loc=predictive_mean, scale=tf.sqrt(predictive_variance))
-
-  def fit(self, x=None, y=None):
-    # p(coeffs | x, y) = Normal(coeffs |
-    #   mean = (1/noise_variance) (1/noise_variance x^T x + I)^{-1} x^T y,
-    #   covariance = (1/noise_variance x^T x + I)^{-1})
-    # TODO(trandustin): We newly fit the data at each call. Extend to do
-    # Bayesian updating.
-    kernel_matrix = tf.matmul(x, x, transpose_a=True) / self.noise_variance
-    coeffs_precision = tf.matrix_set_diag(
-        kernel_matrix, tf.matrix_diag_part(kernel_matrix) + 1.)
-    coeffs_precision_tril = tf.linalg.cholesky(coeffs_precision)
-    self.coeffs_precision_tril_op = tf.linalg.LinearOperatorLowerTriangular(
-        coeffs_precision_tril)
-    self.coeffs_mean = self.coeffs_precision_tril_op.solvevec(
-        self.coeffs_precision_tril_op.solvevec(tf.einsum('nm,n->m', x, y)),
-        adjoint=True) / self.noise_variance
-    # TODO(trandustin): To be fully Keras-compatible, return History object.
-    return
-
-
-def batch_mlp(inputs, hidden_sizes):
-  """Apply MLP to the final axis of a 3D tensor.
-
-  Args:
-    inputs: input Tensor of shape [batch_size, n, d_in].
-    hidden_sizes: An iterable containing the hidden layer sizes of the MLP.
-
-  Returns:
-    Tensor of shape [batch_size, n, d_out] where d_out = output_sizes[-1].
-  """
-  batch_size, _, filter_size = inputs.shape.as_list()
-  hidden = tf.reshape(inputs, (-1, filter_size))
-
-  for size in hidden_sizes[:-1]:
-    hidden = tf.keras.layers.Dense(size, activation=tf.nn.relu)(hidden)
-
-  output = tf.keras.layers.Dense(hidden_sizes[-1], activation=None)(hidden)
-  output = tf.reshape(output, (batch_size, -1, hidden_sizes[-1]))
-  return output
-
-
-# TODO(adityagrover): Reimplement using preexisting attention routines in T2T
-def uniform_attention(q, v):
-  """Computes uniform attention. Equivalent to neural process.
-
-  Args:
-    q: queries. Tensor of shape [batch_size, m, d_k].
-    v: values. Tensor of shape [batch_size, n, d_v].
-
-  Returns:
-    Tensor of shape [batch_size, m, d_v].
-  """
-  total_points = tf.shape(q)[1]
-  rep = tf.reduce_mean(v, axis=1, keepdims=True)  # [batch_size, 1, d_v]
-  rep = tf.tile(rep, [1, total_points, 1])
-  return rep
-
-
-def laplace_attention(q, k, v, scale, normalise):
-  """Computes laplace exponential attention.
-
-  Args:
-    q: queries. Tensor of shape [batch_size, m, d_k].
-    k: keys. Tensor of shape [batch_size, n, d_k].
-    v: values. Tensor of shape [batch_size, n, d_v].
-    scale: float that scales the L1 distance.
-    normalise: Boolean that determines whether weights sum to 1.
-
-  Returns:
-    Tensor of shape [batch_size, m, d_v].
-  """
-  k = tf.expand_dims(k, axis=1)  # [batch_size, 1, n, d_k]
-  q = tf.expand_dims(q, axis=2)  # [batch_size, m, 1, d_k]
-  unnorm_weights = - tf.abs((k - q) / scale)  # [batch_size, m, n, d_k]
-  unnorm_weights = tf.reduce_sum(unnorm_weights, axis=-1)  # [batch_size, m, n]
-  if normalise:
-    weight_fn = tf.nn.softmax
-  else:
-    weight_fn = lambda x: 1 + tf.tanh(x)
-  weights = weight_fn(unnorm_weights)  # [batch_size, m, n]
-  rep = tf.einsum('bik,bkj->bij', weights, v)  # [batch_size, m, d_v]
-  return rep
-
-
-def dot_product_attention(q, k, v, normalise):
-  """Computes dot product attention.
-
-  Args:
-    q: queries. Tensor of  shape [batch_size, m, d_k].
-    k: keys. Tensor of shape [batch_size, n, d_k].
-    v: values. Tensor of shape [batch_size, n, d_v].
-    normalise: Boolean that determines whether weights sum to 1.
-
-  Returns:
-    Tensor of shape [batch_size, m, d_v].
-  """
-  d_k = tf.shape(q)[-1]
-  scale = tf.sqrt(tf.cast(d_k, tf.float32))
-  unnorm_weights = tf.einsum('bjk,bik->bij', k, q) / scale  # [batch_size,m,n]
-  if normalise:
-    weight_fn = tf.nn.softmax
-  else:
-    weight_fn = tf.sigmoid
-  weights = weight_fn(unnorm_weights)  # [batch_size,m,n]
-  rep = tf.einsum('bik,bkj->bij', weights, v)  # [batch_size,m,d_v]
-  return rep
-
-
-def multihead_attention(q, k, v, num_heads=8):
-  """Computes multi-head attention.
-
-  Args:
-    q: queries. Tensor of  shape [batch_size, m, d_k].
-    k: keys. Tensor of shape [batch_size, n, d_k].
-    v: values. Tensor of shape [batch_size, n, d_v].
-    num_heads: number of heads. Should divide d_v.
-
-  Returns:
-    Tensor of shape [batch_size, m, d_v].
-  """
-  d_k = q.get_shape().as_list()[-1]
-  d_v = v.get_shape().as_list()[-1]
-  head_size = int(d_v / num_heads)
-  key_initializer = tf.random_normal_initializer(stddev=d_k**-0.5)
-  value_initializer = tf.random_normal_initializer(stddev=d_v**-0.5)
-  rep = tf.constant(0.0)
-  for h in range(num_heads):
-    o = dot_product_attention(
-        tf.keras.layers.Conv1D(
-            head_size, 1, kernel_initializer=key_initializer,
-            name='wq%d' % h, use_bias=False, padding='VALID')(q),
-        tf.keras.layers.Conv1D(
-            head_size, 1, kernel_initializer=key_initializer,
-            name='wk%d' % h, use_bias=False, padding='VALID')(k),
-        tf.keras.layers.Conv1D(
-            head_size, 1, kernel_initializer=key_initializer,
-            name='wv%d' % h, use_bias=False, padding='VALID')(v),
-        normalise=True)
-    rep += tf.keras.layers.Conv1D(d_v, 1, kernel_initializer=value_initializer,
-                                  name='wo%d' % h, use_bias=False,
-                                  padding='VALID')(o)
-  return rep
-
-
-# TODO(adityagrover): Implement via T2T.
-class Attention(object):
-  """The Attention module."""
-
-  def __init__(self, rep, output_sizes, att_type, scale=1., normalise=True,
-               num_heads=8):
-    """Creates a attention module.
-
-    Takes in context inputs, target inputs and
-    representations of each context input/output pair
-    to output an aggregated representation of the context data.
-
-    Args:
-      rep: transformation to apply to contexts before computing attention.
-          One of: ['identity', 'mlp'].
-      output_sizes: list of number of hidden units per layer of mlp.
-          Used only if rep == 'mlp'.
-      att_type: type of attention. One of the following:
-          ['uniform', 'laplace', 'dot_product', 'multihead']
-      scale: scale of attention.
-      normalise: Boolean determining whether to:
-          1. apply softmax to weights so they sum to 1 across context pts or
-          2. apply custom transformation to have weights in [0, 1].
-      num_heads: number of heads for multihead.
-    """
-    self._rep = rep
-    self._output_sizes = output_sizes
-    self._type = att_type
-    self._scale = scale
-    self._normalise = normalise
-    if self._type == 'multihead':
-      self._num_heads = num_heads
-
-  def __call__(self, x1, x2, r):
-    """Applies attention to create aggregated representation of r.
-
-    Args:
-      x1: Tensor of shape [B ,n1, d_x].
-      x2: Tensor of shape [batch_size, n2, d_x].
-      r: Tensor of shape [batch_size, n1, d].
-
-    Returns:
-      Tensor of shape [batch_size, n2, d]
-
-    Raises:
-      NameError: The argument for rep/type was invalid.
-    """
-    if self._rep == 'identity':
-      k, q = (x1, x2)
-    elif self._rep == 'mlp':
-      k = batch_mlp(x1, self._output_sizes)
-      q = batch_mlp(x2, self._output_sizes)
-    else:
-      raise NameError("'rep' not among ['identity', 'mlp']")
-
-    if self._type == 'uniform':
-      rep = uniform_attention(q, r)
-    elif self._type == 'laplace':
-      rep = laplace_attention(q, k, r, self._scale, self._normalise)
-    elif self._type == 'dot_product':
-      rep = dot_product_attention(q, k, r, self._normalise)
-    elif self._type == 'multihead':
-      rep = multihead_attention(q, k, r, self._num_heads)
-    else:
-      raise NameError(("'att_type' not among ['uniform', 'laplace', "
-                       "'dot_product', 'multihead']"))
-
-    return rep
-
-
-# TODO(adityagrover): Make the encoder and decoder configurable.
-class NeuralProcess(tf.keras.Model):
-  """Attentive Neural Process (Kim et al., 2019; Garnelo et al., 2018)."""
-
-  def __init__(self,
-               latent_encoder_sizes,
-               num_latents,
-               decoder_sizes,
-               use_deterministic_path=True,
-               deterministic_encoder_sizes=None,
-               attention_wrapper=None):
-    """Initializes the Neural Process model.
-
-    Args:
-      latent_encoder_sizes: (list of ints) Hidden layer sizes for latent
-          encoder.
-      num_latents: (int) Dimensionality of global latent variable.
-      decoder_sizes: (list of ints) Hidden layer sizes for decoder
-      use_deterministic_path: (bool) Uses deterministic encoder as well if True.
-      deterministic_encoder_sizes: (list of ints) Hidden layer sizes for
-          deterministic encoder.
-      attention_wrapper: Instance of Attention class to apply for
-          determinitic encoder embedding.
-    """
-    super(NeuralProcess, self).__init__()
-    self._num_latents = num_latents
-    self._latent_encoder_sizes = latent_encoder_sizes
-    self._deterministic_encoder_sizes = deterministic_encoder_sizes
-    self._decoder_sizes = decoder_sizes
-    self._use_deterministic_path = use_deterministic_path
-    self._attention = attention_wrapper
-
-  def latent_encoder(self, x, y):
-    """Encodes the inputs into one representation.
-
-    Args:
-      x: Tensor of shape [batch_size, observations, d_x]. For the prior, these
-         are context x-values. For the posterior, these are target x-values.
-      y: Tensor of shape [batch_size, observations, d_y]. For the prior, these
-         are context y-values. For the posterior, these are target y-values.
-
-    Returns:
-      A normal distribution over tensors of shape [batch_size, num_latents].
-    """
-    encoder_input = tf.concat([x, y], axis=-1)
-    per_example_embedding = batch_mlp(
-        encoder_input, self._latent_encoder_sizes)
-    dataset_embedding = tf.reduce_mean(per_example_embedding, axis=1)
-    hidden = tf.keras.layers.Dense(
-        (self._latent_encoder_sizes[-1] + self._num_latents)//2,
-        activation=tf.nn.relu)(dataset_embedding)
-    loc = tf.keras.layers.Dense(self._num_latents, activation=None)(hidden)
-    untransformed_scale = tf.keras.layers.Dense(self._num_latents,
-                                                activation=None)(hidden)
-    # Constraint scale following Garnelo et al. (2018).
-    scale_diag = 0.1 + 0.9 * tf.sigmoid(untransformed_scale)
-    return ed.MultivariateNormalDiag(loc=loc,
-                                     scale_diag=scale_diag)
-
-  def deterministic_encoder(self, context_x, context_y, target_x):
-    """Encodes the inputs into one representation.
-
-    Args:
-      context_x: Tensor of shape [batch_size, observations, d_x].
-        Observed x-values.
-      context_y: Tensor of shape [batch_size, observations, d_y].
-        Observed y-values.
-      target_x: Tensor of shape [batch_size, target_observations, d_x].
-        Target x-values.
-
-    Returns:
-      Encodings. Tensor of shape [batch_size, target_observations, d].
-    """
-    encoder_input = tf.concat([context_x, context_y], axis=-1)
-    per_example_embedding = batch_mlp(encoder_input,
-                                      self._deterministic_encoder_sizes)
-    per_target_embedding = self._attention(context_x,
-                                           target_x,
-                                           per_example_embedding)
-    return per_target_embedding
-
-  def decoder(self, representation, target_x):
-    """Decodes the individual targets.
-
-    Args:
-      representation: The representation of the context for target predictions.
-          Tensor of shape [batch_size, target_observations, ?].
-      target_x: The x locations for the target query.
-          Tensor of shape [batch_size, target_observations, d_x].
-
-    Returns:
-      dist: A multivariate Gaussian over the target points. A distribution over
-          tensors of shape [batch_size, target_observations, d_y].
-    """
-    decoder_input = tf.concat([representation, target_x], axis=-1)
-    hidden = batch_mlp(decoder_input, self._decoder_sizes)
-    loc, untransformed_scale = tf.split(hidden, 2, axis=-1)
-    scale_diag = 0.1 + 0.9 * tf.nn.softplus(untransformed_scale)
-    return tfp.distributions.MultivariateNormalDiag(loc=loc,
-                                                    scale_diag=scale_diag)
-
-  def __call__(self, query, target_y=None):
-    """Returns the predicted mean and variance at the target points.
-
-    Args:
-      query: Nested tuple containing ((context_x, context_y), target_x) where:
-              context_x is Tensor of shape [batch_size, num_contexts, d_x].
-                  Contains the x values of the context points.
-              context_y is Tensor of shape [batch_size, num_contexts, d_y].
-                  Contains the y values of the context points.
-              target_x is Tensor of shape [batch_size, num_targets, d_x].
-                  Contains the x values of the target points.
-      target_y: The ground truth y values of the target y.
-          Tensor of shape [batch_size, num_targets, d_y].
-
-    Returns:
-      predictive_dist: Predictive posterior distribution over the predicted y.
-    """
-
-    (context_x, context_y), target_x = query
-    num_targets = tf.shape(target_x)[1]
-    prior = self.latent_encoder(context_x, context_y)
-
-    # For training, when target_y is available, use targets for latent encoder.
-    # Note that targets contain contexts by design.
-    # For testing, when target_y unavailable, use contexts for latent encoder.
-    if target_y is None:
-      latent_rep = prior
-    else:
-      posterior = self.latent_encoder(target_x, target_y)
-      latent_rep = posterior
-    latent_rep = tf.tile(tf.expand_dims(latent_rep, axis=1),
-                         [1, num_targets, 1])
-    if self._use_deterministic_path:
-      deterministic_rep = self.deterministic_encoder(context_x,
-                                                     context_y,
-                                                     target_x)
-      representation = tf.concat([deterministic_rep, latent_rep], axis=-1)
-    else:
-      representation = latent_rep
-
-    predictive_dist = self.decoder(representation, target_x)
-
-    if target_y is not None:
-      kl = tf.expand_dims(
-          posterior.distribution.kl_divergence(prior.distribution),
-          -1)
-      self.add_loss(lambda: kl)
-
-    return predictive_dist
-
-  call = __call__
-
-
-
diff --git a/tensor2tensor/layers/gaussian_process_test.py b/tensor2tensor/layers/gaussian_process_test.py
deleted file mode 100644
index c89d3666c..000000000
--- a/tensor2tensor/layers/gaussian_process_test.py
+++ /dev/null
@@ -1,292 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 Gaussian process layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-
-from tensor2tensor.layers import gaussian_process
-from tensor2tensor.utils import test_utils
-
-import tensorflow as tf
-tf.compat.v1.enable_eager_execution()
-
-
-class GaussianProcessTest(tf.test.TestCase):
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testGaussianProcessPosterior(self):
-    train_batch_size = 3
-    test_batch_size = 2
-    input_dim = 4
-    output_dim = 5
-    features = tf.to_float(np.random.rand(train_batch_size, input_dim))
-    labels = tf.to_float(np.random.rand(train_batch_size, output_dim))
-    layer = gaussian_process.GaussianProcess(output_dim,
-                                             conditional_inputs=features,
-                                             conditional_outputs=labels)
-    test_features = tf.to_float(np.random.rand(test_batch_size, input_dim))
-    test_labels = tf.to_float(np.random.rand(test_batch_size, output_dim))
-    test_outputs = layer(test_features)
-    test_nats = -test_outputs.distribution.log_prob(test_labels)
-    self.evaluate(tf.global_variables_initializer())
-    test_nats_val, outputs_val = self.evaluate([test_nats, test_outputs])
-    self.assertEqual(test_nats_val.shape, ())
-    self.assertGreaterEqual(test_nats_val, 0.)
-    self.assertEqual(outputs_val.shape, (test_batch_size, output_dim))
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testGaussianProcessPrior(self):
-    batch_size = 3
-    input_dim = 4
-    output_dim = 5
-    features = tf.to_float(np.random.rand(batch_size, input_dim))
-    labels = tf.to_float(np.random.rand(batch_size, output_dim))
-    model = tf.keras.Sequential([
-        tf.keras.layers.Dense(2, activation=None),
-        gaussian_process.GaussianProcess(output_dim),
-    ])
-    outputs = model(features)
-    log_prob = outputs.distribution.log_prob(labels)
-    self.evaluate(tf.global_variables_initializer())
-    log_prob_val, outputs_val = self.evaluate([log_prob, outputs])
-    self.assertEqual(log_prob_val.shape, ())
-    self.assertLessEqual(log_prob_val, 0.)
-    self.assertEqual(outputs_val.shape, (batch_size, output_dim))
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testSparseGaussianProcess(self):
-    dataset_size = 10
-    batch_size = 3
-    input_dim = 4
-    output_dim = 5
-    features = tf.to_float(np.random.rand(batch_size, input_dim))
-    labels = tf.to_float(np.random.rand(batch_size, output_dim))
-    model = gaussian_process.SparseGaussianProcess(output_dim, num_inducing=2)
-    with tf.GradientTape() as tape:
-      predictions = model(features)
-      nll = -tf.reduce_mean(predictions.distribution.log_prob(labels))
-      kl = sum(model.losses) / dataset_size
-      loss = nll + kl
-
-    self.evaluate(tf.global_variables_initializer())
-    grads = tape.gradient(nll, model.variables)
-    for grad in grads:
-      self.assertIsNotNone(grad)
-
-    loss_val, predictions_val = self.evaluate([loss, predictions])
-    self.assertEqual(loss_val.shape, ())
-    self.assertGreaterEqual(loss_val, 0.)
-    self.assertEqual(predictions_val.shape, (batch_size, output_dim))
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testBayesianLinearModel(self):
-    """Tests that model makes reasonable predictions."""
-    np.random.seed(42)
-    train_batch_size = 5
-    test_batch_size = 2
-    num_features = 3
-    noise_variance = 0.01
-    coeffs = tf.range(num_features, dtype=tf.float32)
-    features = tf.to_float(np.random.randn(train_batch_size, num_features))
-    labels = (tf.tensordot(features, coeffs, [[-1], [0]])
-              + noise_variance * tf.to_float(np.random.randn(train_batch_size)))
-
-    model = gaussian_process.BayesianLinearModel(noise_variance=noise_variance)
-    model.fit(features, labels)
-
-    test_features = tf.to_float(np.random.randn(test_batch_size, num_features))
-    test_labels = tf.tensordot(test_features, coeffs, [[-1], [0]])
-    outputs = model(test_features)
-    test_predictions = outputs.distribution.mean()
-    test_predictions_variance = outputs.distribution.variance()
-
-    [
-        test_labels_val, test_predictions_val, test_predictions_variance_val,
-    ] = self.evaluate(
-        [test_labels, test_predictions, test_predictions_variance])
-    self.assertEqual(test_predictions_val.shape, (test_batch_size,))
-    self.assertEqual(test_predictions_variance_val.shape, (test_batch_size,))
-    self.assertAllClose(test_predictions_val, test_labels_val, atol=0.1)
-    self.assertAllLessEqual(test_predictions_variance_val, noise_variance)
-
-
-def train_neural_process(model,
-                         train_data,
-                         valid_data,
-                         num_epochs,
-                         batch_size,
-                         learning_rate=1e-4):
-  """Trains the NeuralProcess model.
-
-  Validation data is used for early stopping,
-
-  Args:
-    model: A NeuralProcess Model subclassing Keras model.
-    train_data: (4-tuple of tensors) Values of x and y for contexts and targets.
-    valid_data: 4-tuple of tensors) Values of x and y for contexts and targets.
-    num_epochs: (int) Number of epochs to train the model for.
-    batch_size: (int) Size of batch.
-    learning_rate: (float) Learning rate for Adam optimizer.
-
-  Returns:
-    best_loss: (float) Average validation loss of best early-stopped model.
-  """
-  optimizer = tf.keras.optimizers.Adam(learning_rate)
-  context_x, context_y, target_x, target_y = train_data
-  valid_context_x, valid_context_y, valid_target_x, valid_target_y = valid_data
-  train_data_size = target_x.shape[0]
-  num_updates_per_epoch = train_data_size//batch_size
-  best_loss = np.inf
-  valid_query = (valid_context_x, valid_context_y), valid_target_x
-
-  for _ in range(num_epochs):
-    for i in range(num_updates_per_epoch):
-      start_idx, end_idx = batch_size*i, batch_size*(i+1)
-      batch_query = ((context_x[start_idx:end_idx],
-                      context_y[start_idx:end_idx]),
-                     target_x[start_idx:end_idx])
-      batch_target_y = target_y[start_idx:end_idx]
-      num_targets = tf.shape(batch_target_y)[1]
-      with tf.GradientTape() as tape:
-        predictive_dist = model(batch_query, batch_target_y)
-        log_p = predictive_dist.log_prob(batch_target_y)
-        kl = tf.tile(model.losses[-1], [1, num_targets])
-        loss = -tf.reduce_mean(log_p - kl/tf.cast(num_targets, tf.float32))
-      gradients = tape.gradient(loss, model.trainable_variables)
-      optimizer.apply_gradients(zip(gradients, model.trainable_variables))
-    predictive_dist = model(valid_query, valid_target_y)
-    log_p = predictive_dist.log_prob(valid_target_y)
-    kl = tf.tile(model.losses[-1], [1, tf.shape(valid_target_y)[1]])
-    valid_loss = -tf.reduce_mean(log_p - kl/tf.cast(num_targets, tf.float32))
-    if valid_loss < best_loss:
-      best_loss = valid_loss
-
-  return best_loss
-
-
-class NeuralProcessTest(tf.test.TestCase):
-
-  def setUp(self):
-    # Create a dummy multi-task fake dataset
-    num_train_problems = 32
-    num_valid_problems = 32
-    num_targets = 50
-    num_contexts = 10
-    input_dim = 5
-
-    def _create_fake_dataset(num_problems):
-      target_x = tf.cast(np.random.rand(num_problems,
-                                        num_targets,
-                                        input_dim),
-                         tf.float32)
-      target_y = tf.cast(np.random.rand(num_problems, num_targets, 1),
-                         tf.float32)
-      context_x, context_y = (target_x[:, :num_contexts, :],
-                              target_y[:, :num_contexts, :])
-      return (context_x, context_y, target_x, target_y)
-
-    self.train_data = _create_fake_dataset(num_train_problems)
-    self.valid_data = _create_fake_dataset(num_valid_problems)
-
-    hidden_size = 128
-    num_latents = 16
-
-    np_attention_wrapper = gaussian_process.Attention(
-        rep='identity', output_sizes=None, att_type='uniform')
-    self.np_model = gaussian_process.NeuralProcess(
-        latent_encoder_sizes=[hidden_size]*4,
-        num_latents=num_latents,
-        decoder_sizes=[hidden_size]*2 + [2],
-        use_deterministic_path=True,
-        deterministic_encoder_sizes=[hidden_size]*4,
-        attention_wrapper=np_attention_wrapper)
-
-    anp_attention_wrapper = gaussian_process.Attention(
-        rep='mlp', output_sizes=[hidden_size]*2, att_type='multihead')
-    self.anp_model = gaussian_process.NeuralProcess(
-        latent_encoder_sizes=[hidden_size]*4,
-        num_latents=num_latents,
-        decoder_sizes=[hidden_size]*2 + [2],
-        use_deterministic_path=True,
-        deterministic_encoder_sizes=[hidden_size]*4,
-        attention_wrapper=anp_attention_wrapper)
-
-    self.models = [self.np_model, self.anp_model]
-    self.num_latents, self.hidden_size, self.num_targets = (num_latents,
-                                                            hidden_size,
-                                                            num_targets)
-    super(NeuralProcessTest, self).setUp()
-
-  def test_termination(self):
-    for model in self.models:
-      validation_loss = train_neural_process(
-          model,
-          self.train_data,
-          self.valid_data,
-          num_epochs=2,
-          batch_size=16,
-          learning_rate=1e-4)
-
-      self.assertGreaterEqual(validation_loss, 0.)
-
-  def test_latent_encoder(self):
-    valid_context_x, valid_context_y, _, _ = self.valid_data
-    batch_size = valid_context_x.shape[0]
-
-    for model in self.models:
-      dist = model.latent_encoder(valid_context_x, valid_context_y).distribution
-      self.assertEqual(dist.loc.shape, (batch_size, self.num_latents))
-      self.assertEqual(dist.scale.shape,
-                       (batch_size, self.num_latents, self.num_latents))
-
-  def test_deterministic_encoder(self):
-    valid_context_x, valid_context_y, valid_target_x, _ = self.valid_data
-    batch_size = valid_context_x.shape[0]
-
-    for model in self.models:
-      embedding = model.deterministic_encoder(
-          valid_context_x, valid_context_y, valid_target_x)
-      self.assertEqual(embedding.shape, (batch_size, self.num_targets,
-                                         self.hidden_size))
-
-  def test_call(self):
-    valid_context_x, valid_context_y, valid_target_x, valid_target_y = self.valid_data
-    batch_size = valid_context_x.shape[0]
-
-    for model in self.models:
-      query = (valid_context_x, valid_context_y), valid_target_x
-      # test 'training' when target_y is available
-      predictive_dist = model(query, valid_target_y)
-      self.assertEqual(predictive_dist.loc.shape, (batch_size, self.num_targets,
-                                                   1))
-      self.assertEqual(predictive_dist.scale.shape,
-                       (batch_size, self.num_targets, 1, 1))
-      self.assertAllGreaterEqual(model.losses, 0.)
-
-      # test 'testing' when target_y is unavailable
-      predictive_dist = model(query)
-      self.assertEqual(predictive_dist.loc.shape, (batch_size, self.num_targets,
-                                                   1))
-      self.assertEqual(predictive_dist.scale.shape,
-                       (batch_size, self.num_targets, 1, 1))
-
-
-if __name__ == '__main__':
-  tf.test.main()
diff --git a/tensor2tensor/layers/reversible_layers.py b/tensor2tensor/layers/reversible_layers.py
deleted file mode 100644
index 421026256..000000000
--- a/tensor2tensor/layers/reversible_layers.py
+++ /dev/null
@@ -1,1139 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Reversible layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-from scipy.optimize import linear_sum_assignment
-import tensorflow as tf
-import tensorflow_probability as tfp
-
-from tensorflow_probability import edward2 as ed
-
-
-class Reverse(tf.keras.layers.Layer):
-  """Swaps the forward and reverse transformations of a layer."""
-
-  def __init__(self, reversible_layer, **kwargs):
-    super(Reverse, self).__init__(**kwargs)
-    if not hasattr(reversible_layer, 'reverse'):
-      raise ValueError('Layer passed-in has not implemented "reverse" method: '
-                       '{}'.format(reversible_layer))
-    self.call = reversible_layer.reverse
-    self.reverse = reversible_layer.call
-
-
-class DiscreteAutoregressiveFlow(tf.keras.layers.Layer):
-  """A discrete reversible layer.
-
-  The flow takes as input a one-hot Tensor of shape `[..., length, vocab_size]`.
-  The flow returns a Tensor of same shape and dtype. (To enable gradients, the
-  input must have float dtype.)
-
-  For the forward pass, the flow computes in serial:
-
-  ```none
-  outputs = []
-  for t in range(length):
-    new_inputs = [outputs, inputs[..., t, :]]
-    net = layer(new_inputs)
-    loc, scale = tf.split(net, 2, axis=-1)
-    loc = tf.argmax(loc, axis=-1)
-    scale = tf.argmax(scale, axis=-1)
-    new_outputs = (((inputs - loc) * inverse(scale)) % vocab_size)[..., -1, :]
-    outputs.append(new_outputs)
-  ```
-
-  For the reverse pass, the flow computes in parallel:
-
-  ```none
-  net = layer(inputs)
-  loc, scale = tf.split(net, 2, axis=-1)
-  loc = tf.argmax(loc, axis=-1)
-  scale = tf.argmax(scale, axis=-1)
-  outputs = (loc + scale * inputs) % vocab_size
-  ```
-
-  The modular arithmetic happens in one-hot space.
-
-  If `x` is a discrete random variable, the induced probability mass function on
-  the outputs `y = flow(x)` is
-
-  ```none
-  p(y) = p(flow.reverse(y)).
-  ```
-
-  The location-only transform is always invertible ([integers modulo
-  `vocab_size` form an additive group](
-  https://en.wikipedia.org/wiki/Modular_arithmetic)). The transform with a scale
-  is invertible if the scale and `vocab_size` are coprime (see
-  [prime fields](https://en.wikipedia.org/wiki/Finite_field)).
-  """
-
-  def __init__(self, layer, temperature, **kwargs):
-    """Constructs flow.
-
-    Args:
-      layer: Two-headed masked network taking the inputs and returning a
-        real-valued Tensor of shape `[..., length, 2*vocab_size]`.
-        Alternatively, `layer` may return a Tensor of shape
-        `[..., length, vocab_size]` to be used as the location transform; the
-        scale transform will be hard-coded to 1.
-      temperature: Positive value determining bias of gradient estimator.
-      **kwargs: kwargs of parent class.
-    """
-    super(DiscreteAutoregressiveFlow, self).__init__(**kwargs)
-    self.layer = layer
-    self.temperature = temperature
-
-  def build(self, input_shape):
-    input_shape = tf.TensorShape(input_shape)
-    self.vocab_size = input_shape[-1]
-    if isinstance(self.vocab_size, tf.Dimension):
-      self.vocab_size = self.vocab_size.value
-    if self.vocab_size is None:
-      raise ValueError('The last dimension of the inputs to '
-                       '`DiscreteAutoregressiveFlow` should be defined. Found '
-                       '`None`.')
-    self.built = True
-
-  def __call__(self, inputs, *args, **kwargs):
-    if not isinstance(inputs, ed.RandomVariable):
-      return super(DiscreteAutoregressiveFlow, self).__call__(
-          inputs, *args, **kwargs)
-    return TransformedRandomVariable(inputs, self)
-
-  def call(self, inputs, **kwargs):
-    """Forward pass for left-to-right autoregressive generation."""
-    inputs = tf.convert_to_tensor(inputs)
-    length = inputs.shape[-2].value
-    if length is None:
-      raise NotImplementedError('length dimension must be known.')
-    # Form initial sequence tensor of shape [..., 1, vocab_size]. In a loop, we
-    # incrementally build a Tensor of shape [..., t, vocab_size] as t grows.
-    outputs = self._initial_call(inputs[..., 0, :], length, **kwargs)
-    # TODO(trandustin): Use tf.while_loop. Unrolling is memory-expensive for big
-    # models and not valid for variable lengths.
-    for t in range(1, length):
-      outputs = self._per_timestep_call(outputs,
-                                        inputs[..., t, :],
-                                        length,
-                                        t,
-                                        **kwargs)
-    return outputs
-
-  def _initial_call(self, new_inputs, length, **kwargs):
-    """Returns Tensor of shape [..., 1, vocab_size].
-
-    Args:
-      new_inputs: Tensor of shape [..., vocab_size], the new input to generate
-        its output.
-      length: Length of final desired sequence.
-      **kwargs: Optional keyword arguments to layer.
-    """
-    inputs = new_inputs[..., tf.newaxis, :]
-    # TODO(trandustin): To handle variable lengths, extend MADE to subset its
-    # input and output layer weights rather than pad inputs.
-    batch_ndims = inputs.shape.ndims - 2
-    padded_inputs = tf.pad(
-        inputs, [[0, 0]] * batch_ndims + [[0, length - 1], [0, 0]])
-    net = self.layer(padded_inputs, **kwargs)
-    if net.shape[-1] == 2 * self.vocab_size:
-      loc, scale = tf.split(net, 2, axis=-1)
-      loc = loc[..., 0:1, :]
-      loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-      scale = scale[..., 0:1, :]
-      scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
-      inverse_scale = multiplicative_inverse(scale, self.vocab_size)
-      shifted_inputs = one_hot_minus(inputs, loc)
-      outputs = one_hot_multiply(shifted_inputs, inverse_scale)
-    elif net.shape[-1] == self.vocab_size:
-      loc = net
-      loc = loc[..., 0:1, :]
-      loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-      outputs = one_hot_minus(inputs, loc)
-    else:
-      raise ValueError('Output of layer does not have compatible dimensions.')
-    return outputs
-
-  def _per_timestep_call(self,
-                         current_outputs,
-                         new_inputs,
-                         length,
-                         timestep,
-                         **kwargs):
-    """Returns Tensor of shape [..., timestep+1, vocab_size].
-
-    Args:
-      current_outputs: Tensor of shape [..., timestep, vocab_size], the so-far
-        generated sequence Tensor.
-      new_inputs: Tensor of shape [..., vocab_size], the new input to generate
-        its output given current_outputs.
-      length: Length of final desired sequence.
-      timestep: Current timestep.
-      **kwargs: Optional keyword arguments to layer.
-    """
-    inputs = tf.concat([current_outputs,
-                        new_inputs[..., tf.newaxis, :]], axis=-2)
-    # TODO(trandustin): To handle variable lengths, extend MADE to subset its
-    # input and output layer weights rather than pad inputs.
-    batch_ndims = inputs.shape.ndims - 2
-    padded_inputs = tf.pad(
-        inputs, [[0, 0]] * batch_ndims + [[0, length - timestep - 1], [0, 0]])
-    net = self.layer(padded_inputs, **kwargs)
-    if net.shape[-1] == 2 * self.vocab_size:
-      loc, scale = tf.split(net, 2, axis=-1)
-      loc = loc[..., :(timestep+1), :]
-      loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-      scale = scale[..., :(timestep+1), :]
-      scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
-      inverse_scale = multiplicative_inverse(scale, self.vocab_size)
-      shifted_inputs = one_hot_minus(inputs, loc)
-      new_outputs = one_hot_multiply(shifted_inputs, inverse_scale)
-    elif net.shape[-1] == self.vocab_size:
-      loc = net
-      loc = loc[..., :(timestep+1), :]
-      loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-      new_outputs = one_hot_minus(inputs, loc)
-    else:
-      raise ValueError('Output of layer does not have compatible dimensions.')
-    outputs = tf.concat([current_outputs, new_outputs[..., -1:, :]], axis=-2)
-    if not tf.executing_eagerly():
-      outputs.set_shape([None] * batch_ndims + [timestep+1, self.vocab_size])
-    return outputs
-
-  def reverse(self, inputs, **kwargs):
-    """Reverse pass returning the inverse autoregressive transformation."""
-    if not self.built:
-      self._maybe_build(inputs)
-
-    net = self.layer(inputs, **kwargs)
-    if net.shape[-1] == 2 * self.vocab_size:
-      loc, scale = tf.split(net, 2, axis=-1)
-      scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
-      scaled_inputs = one_hot_multiply(inputs, scale)
-    elif net.shape[-1] == self.vocab_size:
-      loc = net
-      scaled_inputs = inputs
-    else:
-      raise ValueError('Output of layer does not have compatible dimensions.')
-    loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-    outputs = one_hot_add(loc, scaled_inputs)
-    return outputs
-
-  def log_det_jacobian(self, inputs):
-    return tf.cast(0, inputs.dtype)
-
-
-class DiscreteBipartiteFlow(tf.keras.layers.Layer):
-  """A discrete reversible layer.
-
-  The flow takes as input a one-hot Tensor of shape `[..., length, vocab_size]`.
-  The flow returns a Tensor of same shape and dtype. (To enable gradients, the
-  input must have float dtype.)
-
-  For the forward pass, the flow computes:
-
-  ```none
-  net = layer(mask * inputs)
-  loc, scale = tf.split(net, 2, axis=-1)
-  loc = tf.argmax(loc, axis=-1)
-  scale = tf.argmax(scale, axis=-1)
-  outputs = ((inputs - (1-mask) * loc) * (1-mask) * inverse(scale)) % vocab_size
-  ```
-
-  For the reverse pass, the flow computes:
-
-  ```none
-  net = layer(mask * inputs)
-  loc, scale = tf.split(net, 2, axis=-1)
-  loc = tf.argmax(loc, axis=-1)
-  scale = tf.argmax(scale, axis=-1)
-  outputs = ((1-mask) * loc + (1-mask) * scale * inputs) % vocab_size
-  ```
-
-  The modular arithmetic happens in one-hot space.
-
-  If `x` is a discrete random variable, the induced probability mass function on
-  the outputs `y = flow(x)` is
-
-  ```none
-  p(y) = p(flow.reverse(y)).
-  ```
-
-  The location-only transform is always invertible ([integers modulo
-  `vocab_size` form an additive group](
-  https://en.wikipedia.org/wiki/Modular_arithmetic)). The transform with a scale
-  is invertible if the scale and `vocab_size` are coprime (see
-  [prime fields](https://en.wikipedia.org/wiki/Finite_field)).
-  """
-
-  def __init__(self, layer, mask, temperature, **kwargs):
-    """Constructs flow.
-
-    Args:
-      layer: Two-headed masked network taking the inputs and returning a
-        real-valued Tensor of shape `[..., length, 2*vocab_size]`.
-        Alternatively, `layer` may return a Tensor of shape
-        `[..., length, vocab_size]` to be used as the location transform; the
-        scale transform will be hard-coded to 1.
-      mask: binary Tensor of shape `[length]` forming the bipartite assignment.
-      temperature: Positive value determining bias of gradient estimator.
-      **kwargs: kwargs of parent class.
-    """
-    super(DiscreteBipartiteFlow, self).__init__(**kwargs)
-    self.layer = layer
-    self.mask = mask
-    self.temperature = temperature
-
-  def build(self, input_shape):
-    input_shape = tf.TensorShape(input_shape)
-    self.vocab_size = input_shape[-1]
-    if isinstance(self.vocab_size, tf.Dimension):
-      self.vocab_size = self.vocab_size.value
-    if self.vocab_size is None:
-      raise ValueError('The last dimension of the inputs to '
-                       '`DiscreteBipartiteFlow` should be defined. Found '
-                       '`None`.')
-    self.built = True
-
-  def __call__(self, inputs, *args, **kwargs):
-    if not isinstance(inputs, ed.RandomVariable):
-      return super(DiscreteBipartiteFlow, self).__call__(
-          inputs, *args, **kwargs)
-    return TransformedRandomVariable(inputs, self)
-
-  def call(self, inputs, **kwargs):
-    """Forward pass for bipartite generation."""
-    inputs = tf.convert_to_tensor(inputs)
-    batch_ndims = inputs.shape.ndims - 2
-    mask = tf.reshape(tf.cast(self.mask, inputs.dtype),
-                      [1] * batch_ndims + [-1, 1])
-    masked_inputs = mask * inputs
-    net = self.layer(masked_inputs, **kwargs)
-    if net.shape[-1] == 2 * self.vocab_size:
-      loc, scale = tf.split(net, 2, axis=-1)
-      loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-      scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
-      inverse_scale = multiplicative_inverse(scale, self.vocab_size)
-      shifted_inputs = one_hot_minus(inputs, loc)
-      masked_outputs = (1. - mask) * one_hot_multiply(shifted_inputs,
-                                                      inverse_scale)
-    elif net.shape[-1] == self.vocab_size:
-      loc = net
-      loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-      masked_outputs = (1. - mask) * one_hot_minus(inputs, loc)
-    else:
-      raise ValueError('Output of layer does not have compatible dimensions.')
-    outputs = masked_inputs + masked_outputs
-    return outputs
-
-  def reverse(self, inputs, **kwargs):
-    """Reverse pass for the inverse bipartite transformation."""
-    if not self.built:
-      self._maybe_build(inputs)
-
-    inputs = tf.convert_to_tensor(inputs)
-    batch_ndims = inputs.shape.ndims - 2
-    mask = tf.reshape(tf.cast(self.mask, inputs.dtype),
-                      [1] * batch_ndims + [-1, 1])
-    masked_inputs = mask * inputs
-    net = self.layer(masked_inputs, **kwargs)
-    if net.shape[-1] == 2 * self.vocab_size:
-      loc, scale = tf.split(net, 2, axis=-1)
-      scale = tf.cast(one_hot_argmax(scale, self.temperature), inputs.dtype)
-      scaled_inputs = one_hot_multiply(inputs, scale)
-    elif net.shape[-1] == self.vocab_size:
-      loc = net
-      scaled_inputs = inputs
-    else:
-      raise ValueError('Output of layer does not have compatible dimensions.')
-    loc = tf.cast(one_hot_argmax(loc, self.temperature), inputs.dtype)
-    masked_outputs = (1. - mask) * one_hot_add(loc, scaled_inputs)
-    outputs = masked_inputs + masked_outputs
-    return outputs
-
-  def log_det_jacobian(self, inputs):
-    return tf.cast(0, inputs.dtype)
-
-
-class SinkhornAutoregressiveFlow(tf.keras.layers.Layer):
-  """A discrete reversible layer using Sinkhorn normalization for permutations.
-
-  The flow takes as input a one-hot Tensor of shape `[..., length, vocab_size]`.
-  The flow returns a Tensor of same shape and dtype. (To enable gradients, the
-  input must have float dtype.)
-  """
-
-  def __init__(self, layer, temperature, **kwargs):
-    """Constructs flow.
-
-    Args:
-      layer: Masked network taking inputs with shape `[..., length, vocab_size]`
-        and returning a real-valued Tensor of shape
-        `[..., length, vocab_size ** 2]`. Sinkhorn iterations are applied to
-        each `layer` output to produce permutation matrices.
-      temperature: Positive value determining bias of gradient estimator.
-      **kwargs: kwargs of parent class.
-    """
-    super(SinkhornAutoregressiveFlow, self).__init__(**kwargs)
-    self.layer = layer
-    self.temperature = temperature
-
-  def build(self, input_shape):
-    input_shape = tf.TensorShape(input_shape)
-    self.vocab_size = input_shape[-1]
-    if isinstance(self.vocab_size, tf.Dimension):
-      self.vocab_size = self.vocab_size.value
-    if self.vocab_size is None:
-      raise ValueError('The last dimension of the inputs to '
-                       '`DiscreteAutoregressiveFlow` should be defined. Found '
-                       '`None`.')
-    self.built = True
-
-  def __call__(self, inputs, *args, **kwargs):
-    if not isinstance(inputs, ed.RandomVariable):
-      return super(SinkhornAutoregressiveFlow, self).__call__(
-          inputs, *args, **kwargs)
-    return TransformedRandomVariable(inputs, self)
-
-  def call(self, inputs, **kwargs):
-    """Forward pass for left-to-right autoregressive generation."""
-    inputs = tf.convert_to_tensor(inputs)
-    length = inputs.shape[-2].value
-    if length is None:
-      raise NotImplementedError('length dimension must be known.')
-    # Form initial sequence tensor of shape [..., 1, vocab_size]. In a loop, we
-    # incrementally build a Tensor of shape [..., t, vocab_size] as t grows.
-    outputs = self._initial_call(inputs[..., 0, :], length, **kwargs)
-    for t in range(1, length):
-      outputs = self._per_timestep_call(outputs,
-                                        inputs[..., t, :],
-                                        length,
-                                        t,
-                                        **kwargs)
-    return outputs
-
-  def _initial_call(self, new_inputs, length, **kwargs):
-    """Returns Tensor of shape [..., 1, vocab_size].
-
-    Args:
-      new_inputs: Tensor of shape [..., vocab_size], the new input to generate
-        its output.
-      length: Length of final desired sequence.
-      **kwargs: Optional keyword arguments to layer.
-    """
-    inputs = new_inputs[..., tf.newaxis, :]
-    # TODO(trandustin): To handle variable lengths, extend MADE to subset its
-    # input and output layer weights rather than pad inputs.
-    batch_ndims = inputs.shape.ndims - 2
-    padded_inputs = tf.pad(
-        inputs, [[0, 0]] * batch_ndims + [[0, length - 1], [0, 0]])
-    temperature = 1.
-    logits = self.layer(padded_inputs / temperature, **kwargs)
-    logits = logits[..., 0:1, :]
-    logits = tf.reshape(
-        logits,
-        logits.shape[:-1].concatenate([self.vocab_size, self.vocab_size]))
-    soft = sinkhorn(logits)
-    hard = tf.cast(soft_to_hard_permutation(soft), inputs.dtype)
-    hard = tf.reshape(hard, logits.shape)
-    # Inverse of permutation matrix is its transpose.
-    # inputs is [batch_size, timestep + 1, vocab_size].
-    # hard is [batch_size, timestep + 1, vocab_size, vocab_size].
-    outputs = tf.matmul(inputs[..., tf.newaxis, :],
-                        hard,
-                        transpose_b=True)[..., 0, :]
-    return outputs
-
-  def _per_timestep_call(self,
-                         current_outputs,
-                         new_inputs,
-                         length,
-                         timestep,
-                         **kwargs):
-    """Returns Tensor of shape [..., timestep+1, vocab_size].
-
-    Args:
-      current_outputs: Tensor of shape [..., timestep, vocab_size], the so-far
-        generated sequence Tensor.
-      new_inputs: Tensor of shape [..., vocab_size], the new input to generate
-        its output given current_outputs.
-      length: Length of final desired sequence.
-      timestep: Current timestep.
-      **kwargs: Optional keyword arguments to layer.
-    """
-    inputs = tf.concat([current_outputs,
-                        new_inputs[..., tf.newaxis, :]], axis=-2)
-    # TODO(trandustin): To handle variable lengths, extend MADE to subset its
-    # input and output layer weights rather than pad inputs.
-    batch_ndims = inputs.shape.ndims - 2
-    padded_inputs = tf.pad(
-        inputs, [[0, 0]] * batch_ndims + [[0, length - timestep - 1], [0, 0]])
-    logits = self.layer(padded_inputs, **kwargs)
-    logits = logits[..., :(timestep+1), :]
-    logits = tf.reshape(
-        logits,
-        logits.shape[:-1].concatenate([self.vocab_size, self.vocab_size]))
-    soft = sinkhorn(logits / self.temperature)
-    hard = tf.cast(soft_to_hard_permutation(soft), inputs.dtype)
-    hard = tf.reshape(hard, logits.shape)
-    # Inverse of permutation matrix is its transpose.
-    # inputs is [batch_size, timestep + 1, vocab_size].
-    # hard is [batch_size, timestep + 1, vocab_size, vocab_size].
-    new_outputs = tf.matmul(inputs[..., tf.newaxis, :],
-                            hard,
-                            transpose_b=True)[..., 0, :]
-    outputs = tf.concat([current_outputs, new_outputs[..., -1:, :]], axis=-2)
-    if not tf.executing_eagerly():
-      outputs.set_shape([None] * batch_ndims + [timestep+1, self.vocab_size])
-    return outputs
-
-  def reverse(self, inputs, **kwargs):
-    """Reverse pass returning the inverse autoregressive transformation."""
-    if not self.built:
-      self._maybe_build(inputs)
-
-    logits = self.layer(inputs, **kwargs)
-    logits = tf.reshape(
-        logits,
-        logits.shape[:-1].concatenate([self.vocab_size, self.vocab_size]))
-    soft = sinkhorn(logits / self.temperature, n_iters=20)
-    hard = soft_to_hard_permutation(soft)
-    hard = tf.reshape(hard, logits.shape)
-    # Recover the permutation by right-multiplying by the permutation matrix.
-    outputs = tf.matmul(inputs[..., tf.newaxis, :], hard)[..., 0, :]
-    return outputs
-
-  def log_det_jacobian(self, inputs):
-    return tf.cast(0, inputs.dtype)
-
-
-def soft_to_hard_permutation(inputs):
-  """Returns permutation matrices by solving a matching problem.
-
-  Solves linear sum assignment to convert doubly-stochastic matrices to
-  permutation matrices. It uses scipy.optimize.linear_sum_assignment to solve
-  the optimization problem max_P sum_i,j M_i,j P_i,j with P a permutation
-  matrix. Notice the negative sign; the reason, the original function solves a
-  minimization problem.
-
-  Code is adapted from Mena et al. [1].
-
-  [1] Gonzalo Mena, David Belanger, Scott Linderman, Jasper Snoek.
-  Learning latent permutations with Gumbel-Sinkhorn networks. International
-  Conference on Learning Representations, 2018.
-
-  Args:
-    inputs: A `Tensor` with shape `[:, vocab_size, vocab_size]` that is
-      doubly-stochastic in its last two dimensions.
-
-  Returns:
-    outputs: A hard permutation `Tensor` with the same shape as `inputs` (in
-      other words the last two dimensions are doubly-stochastic and each element
-      is 0 or 1).
-  """
-
-  def hungarian(x):
-    if x.ndim == 2:
-      x = np.reshape(x, [1, x.shape[0], x.shape[1]])
-    sol = np.zeros((x.shape[0], x.shape[1]), dtype=np.int32)
-    for i in range(x.shape[0]):
-      sol[i, :] = linear_sum_assignment(-x[i, :])[1].astype(np.int32)
-    return sol
-
-  vocab_size = inputs.shape[-1]
-  # Note: tf.py_func isn't currently supported on headless GPUs.
-  # TODO(vafa): Fix tf.py_func headless GPU bug.
-  permutation_lists = tf.py_func(hungarian, [inputs], tf.int32)
-  hard = tf.one_hot(permutation_lists, depth=vocab_size)
-  outputs = tf.stop_gradient(hard - inputs) + inputs
-  return outputs
-
-
-def one_hot_argmax(inputs, temperature, axis=-1):
-  """Returns one-hot of argmax with backward pass set to softmax-temperature."""
-  vocab_size = inputs.shape[-1]
-  hard = tf.one_hot(tf.argmax(inputs, axis=axis),
-                    depth=vocab_size,
-                    axis=axis,
-                    dtype=inputs.dtype)
-  soft = tf.nn.softmax(inputs / temperature, axis=axis)
-  outputs = soft + tf.stop_gradient(hard - soft)
-  return outputs
-
-
-def one_hot_add(inputs, shift):
-  """Performs (inputs + shift) % vocab_size in the one-hot space.
-
-  Args:
-    inputs: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
-      Tensor.
-    shift: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
-      Tensor specifying how much to shift the corresponding one-hot vector in
-      inputs. Soft values perform a "weighted shift": for example,
-      shift=[0.2, 0.3, 0.5] performs a linear combination of 0.2 * shifting by
-      zero; 0.3 * shifting by one; and 0.5 * shifting by two.
-
-  Returns:
-    Tensor of same shape and dtype as inputs.
-  """
-  # Compute circular 1-D convolution with shift as the kernel.
-  inputs = tf.cast(inputs, tf.complex64)
-  shift = tf.cast(shift, tf.complex64)
-  return tf.real(tf.signal.ifft(tf.signal.fft(inputs) * tf.signal.fft(shift)))
-
-
-def one_hot_minus(inputs, shift):
-  """Performs (inputs - shift) % vocab_size in the one-hot space.
-
-  Args:
-    inputs: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
-      Tensor.
-    shift: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
-      Tensor specifying how much to shift the corresponding one-hot vector in
-      inputs. Soft values perform a "weighted shift": for example,
-      shift=[0.2, 0.3, 0.5] performs a linear combination of 0.2 * shifting by
-      zero; 0.3 * shifting by one; and 0.5 * shifting by two.
-
-  Returns:
-    Tensor of same shape and dtype as inputs.
-  """
-  # TODO(trandustin): Implement with circular conv1d.
-  inputs = tf.convert_to_tensor(inputs)
-  shift = tf.cast(shift, inputs.dtype)
-  vocab_size = inputs.shape[-1].value
-  # Form a [..., vocab_size, vocab_size] matrix. Each batch element of
-  # inputs will vector-matrix multiply the vocab_size x vocab_size matrix. This
-  # "shifts" the inputs batch element by the corresponding shift batch element.
-  shift_matrix = tf.stack([tf.roll(shift, i, axis=-1)
-                           for i in range(vocab_size)], axis=-2)
-  outputs = tf.einsum('...v,...uv->...u', inputs, shift_matrix)
-  return outputs
-
-
-def one_hot_multiply(inputs, scale):
-  """Performs (inputs * scale) % vocab_size in the one-hot space.
-
-  Args:
-    inputs: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
-      Tensor.
-    scale: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
-      Tensor specifying how much to scale the corresponding one-hot vector in
-      inputs. Soft values perform a "weighted scale": for example,
-      scale=[0.2, 0.3, 0.5] performs a linear combination of
-      0.2 * scaling by zero; 0.3 * scaling by one; and 0.5 * scaling by two.
-
-  Returns:
-    Tensor of same shape and dtype as inputs.
-  """
-  # TODO(trandustin): Implement with circular conv1d.
-  inputs = tf.convert_to_tensor(inputs)
-  scale = tf.cast(scale, inputs.dtype)
-  batch_shape = inputs.shape[:-1].as_list()
-  vocab_size = inputs.shape[-1].value
-  # Form a [..., vocab_size, vocab_size] tensor. The ith row of the
-  # batched vocab_size x vocab_size matrix represents scaling inputs by i.
-  permutation_matrix = tf.floormod(
-      tf.tile(tf.range(vocab_size)[:, tf.newaxis], [1, vocab_size]) *
-      tf.range(vocab_size)[tf.newaxis], vocab_size)
-  permutation_matrix = tf.one_hot(permutation_matrix, depth=vocab_size, axis=-1)
-  # Scale the inputs according to the permutation matrix of all possible scales.
-  scaled_inputs = tf.einsum('...v,avu->...au', inputs, permutation_matrix)
-  scaled_inputs = tf.concat([tf.zeros(batch_shape + [1, vocab_size]),
-                             scaled_inputs[..., 1:, :]], axis=-2)
-  # Reduce rows of the scaled inputs by the scale values. This forms a
-  # weighted linear combination of scaling by zero, scaling by one, and so on.
-  outputs = tf.einsum('...v,...vu->...u', scale, scaled_inputs)
-  return outputs
-
-
-def py_multiplicative_inverse(a, n):
-  """Multiplicative inverse of a modulo n (in Python).
-
-  Implements extended Euclidean algorithm.
-
-  Args:
-    a: int-like np.ndarray.
-    n: int.
-
-  Returns:
-    Multiplicative inverse as an int32 np.ndarray with same shape as a.
-  """
-  batched_a = np.asarray(a, dtype=np.int32)
-  batched_inverse = []
-  for a in np.nditer(batched_a):
-    inverse = 0
-    new_inverse = 1
-    remainder = n
-    new_remainder = a
-    while new_remainder != 0:
-      quotient = remainder // new_remainder
-      (inverse, new_inverse) = (new_inverse, inverse - quotient * new_inverse)
-      (remainder, new_remainder) = (new_remainder,
-                                    remainder - quotient * new_remainder)
-    if remainder > 1:
-      return ValueError(
-          'Inverse for {} modulo {} does not exist.'.format(a, n))
-    if inverse < 0:
-      inverse += n
-    batched_inverse.append(inverse)
-  return np.asarray(batched_inverse, dtype=np.int32).reshape(batched_a.shape)
-
-
-def multiplicative_inverse(a, n):
-  """Multiplicative inverse of a modulo n.
-
-  Args:
-    a: Tensor of shape [..., vocab_size]. It denotes an integer in the one-hot
-      space.
-    n: int Tensor of shape [...].
-
-  Returns:
-    Tensor of same shape and dtype as a.
-  """
-  a = tf.convert_to_tensor(a)
-  n = tf.convert_to_tensor(n)
-  vocab_size = a.shape[-1].value
-  a_dtype = a.dtype
-  sparse_a = tf.argmax(a, axis=-1)
-  sparse_outputs = tf.py_func(
-      py_multiplicative_inverse, [sparse_a, n], tf.int32)
-  sparse_outputs.set_shape(sparse_a.shape)
-  outputs = tf.one_hot(sparse_outputs, depth=vocab_size, dtype=a_dtype)
-  return outputs
-
-
-class ActNorm(tf.keras.layers.Layer):
-  """Actnorm, an affine reversible layer (Prafulla and Kingma, 2018).
-
-  Weights use data-dependent initialization in which outputs have zero mean
-  and unit variance per channel (last dimension). The mean/variance statistics
-  are computed from the first batch of inputs.
-  """
-
-  def __init__(self, epsilon=tf.keras.backend.epsilon(), **kwargs):
-    super(ActNorm, self).__init__(**kwargs)
-    self.epsilon = epsilon
-
-  def build(self, input_shape):
-    input_shape = tf.TensorShape(input_shape)
-    last_dim = input_shape[-1]
-    if isinstance(last_dim, tf.Dimension):
-      last_dim = last_dim.value
-    if last_dim is None:
-      raise ValueError('The last dimension of the inputs to `ActNorm` '
-                       'should be defined. Found `None`.')
-    bias = self.add_weight('bias', [last_dim], dtype=self.dtype)
-    log_scale = self.add_weight('log_scale', [last_dim], dtype=self.dtype)
-    # Set data-dependent initializers.
-    bias = bias.assign(self.bias_initial_value)
-    with tf.control_dependencies([bias]):
-      self.bias = bias
-    log_scale = log_scale.assign(self.log_scale_initial_value)
-    with tf.control_dependencies([log_scale]):
-      self.log_scale = log_scale
-    self.built = True
-
-  def __call__(self, inputs, *args, **kwargs):
-    if not self.built:
-      mean, variance = tf.nn.moments(
-          inputs, axes=[i for i in range(inputs.shape.ndims - 1)])
-      self.bias_initial_value = -mean
-      # TODO(trandustin): Optionally, actnorm multiplies log_scale by a fixed
-      # log_scale factor (e.g., 3.) and initializes by
-      # initial_value / log_scale_factor.
-      self.log_scale_initial_value = tf.log(
-          1. / (tf.sqrt(variance) + self.epsilon))
-
-    if not isinstance(inputs, ed.RandomVariable):
-      return super(ActNorm, self).__call__(inputs, *args, **kwargs)
-    return TransformedRandomVariable(inputs, self)
-
-  def call(self, inputs):
-    return (inputs + self.bias) * tf.exp(self.log_scale)
-
-  def reverse(self, inputs):
-    return inputs * tf.exp(-self.log_scale) - self.bias
-
-  def log_det_jacobian(self, inputs):
-    """Returns log det | dx / dy | = num_events * sum log | scale |."""
-    del inputs  # unused
-    # Number of events is number of all elements excluding the batch and
-    # channel dimensions.
-    num_events = tf.reduce_prod(tf.shape(inputs)[1:-1])
-    log_det_jacobian = num_events * tf.reduce_sum(self.log_scale)
-    return log_det_jacobian
-
-
-class MADE(tf.keras.Model):
-  """Masked autoencoder for distribution estimation (Germain et al., 2015).
-
-  MADE takes as input a real Tensor of shape [..., length, channels] and returns
-  a Tensor of shape [..., length, units] and same dtype. It masks layer weights
-  to satisfy autoregressive constraints with respect to the length dimension. In
-  particular, for a given ordering, each input dimension of length can be
-  reconstructed from previous dimensions.
-
-  The output's units dimension captures per-time-step representations. For
-  example, setting units to 2 can parameterize the location and log-scale of an
-  autoregressive Gaussian distribution.
-  """
-
-  def __init__(self,
-               units,
-               hidden_dims,
-               input_order='left-to-right',
-               hidden_order='left-to-right',
-               activation=None,
-               use_bias=True,
-               **kwargs):
-    """Constructs network.
-
-    Args:
-      units: Positive integer, dimensionality of the output space.
-      hidden_dims: list with the number of hidden units per layer. It does not
-        include the output layer; those number of units will always be set to
-        the input dimension multiplied by `num_heads`. Each hidden unit size
-        must be at least the size of length (otherwise autoregressivity is not
-        possible).
-      input_order: Order of degrees to the input units: 'random',
-        'left-to-right', 'right-to-left', or an array of an explicit order.
-        For example, 'left-to-right' builds an autoregressive model
-        p(x) = p(x1) p(x2 | x1) ... p(xD | x<D).
-      hidden_order: Order of degrees to the hidden units: 'random',
-        'left-to-right'. If 'left-to-right', hidden units are allocated equally
-        (up to a remainder term) to each degree.
-      activation: Activation function.
-      use_bias: Whether to use a bias.
-      **kwargs: Keyword arguments of parent class.
-    """
-    super(MADE, self).__init__(**kwargs)
-    self.units = int(units)
-    self.hidden_dims = hidden_dims
-    self.input_order = input_order
-    self.hidden_order = hidden_order
-    self.activation = tf.keras.activations.get(activation)
-    self.use_bias = use_bias
-    self.network = tf.keras.Sequential([])
-
-  def build(self, input_shape):
-    input_shape = tf.TensorShape(input_shape)
-    length = input_shape[-2]
-    channels = input_shape[-1]
-    if isinstance(length, tf.Dimension):
-      length = length.value
-    if isinstance(channels, tf.Dimension):
-      channels = channels.value
-    if length is None or channels is None:
-      raise ValueError('The two last dimensions of the inputs to '
-                       '`MADE` should be defined. Found `None`.')
-    masks = create_masks(input_dim=length,
-                         hidden_dims=self.hidden_dims,
-                         input_order=self.input_order,
-                         hidden_order=self.hidden_order)
-
-    # Input-to-hidden layer: [..., length, channels] -> [..., hidden_dims[0]].
-    self.network.add(tf.keras.layers.Reshape([length * channels]))
-    # Tile the mask so each element repeats contiguously; this is compatible
-    # with the autoregressive contraints unlike naive tiling.
-    mask = masks[0]
-    mask = tf.tile(mask[:, tf.newaxis, :], [1, channels, 1])
-    mask = tf.reshape(mask, [mask.shape[0] * channels, mask.shape[-1]])
-    if self.hidden_dims:
-      layer = tf.keras.layers.Dense(
-          self.hidden_dims[0],
-          kernel_initializer=make_masked_initializer(mask),
-          kernel_constraint=make_masked_constraint(mask),
-          activation=self.activation,
-          use_bias=self.use_bias)
-      self.network.add(layer)
-
-    # Hidden-to-hidden layers: [..., hidden_dims[l-1]] -> [..., hidden_dims[l]].
-    for l in range(1, len(self.hidden_dims)):
-      layer = tf.keras.layers.Dense(
-          self.hidden_dims[l],
-          kernel_initializer=make_masked_initializer(masks[l]),
-          kernel_constraint=make_masked_constraint(masks[l]),
-          activation=self.activation,
-          use_bias=self.use_bias)
-      self.network.add(layer)
-
-    # Hidden-to-output layer: [..., hidden_dims[-1]] -> [..., length, units].
-    # Tile the mask so each element repeats contiguously; this is compatible
-    # with the autoregressive contraints unlike naive tiling.
-    if self.hidden_dims:
-      mask = masks[-1]
-    mask = tf.tile(mask[..., tf.newaxis], [1, 1, self.units])
-    mask = tf.reshape(mask, [mask.shape[0], mask.shape[1] * self.units])
-    layer = tf.keras.layers.Dense(
-        length * self.units,
-        kernel_initializer=make_masked_initializer(mask),
-        kernel_constraint=make_masked_constraint(mask),
-        activation=None,
-        use_bias=self.use_bias)
-    self.network.add(layer)
-    self.network.add(tf.keras.layers.Reshape([length, self.units]))
-    self.built = True
-
-  def call(self, inputs):
-    return self.network(inputs)
-
-
-def create_degrees(input_dim,
-                   hidden_dims,
-                   input_order='left-to-right',
-                   hidden_order='left-to-right'):
-  """Returns a list of degree vectors, one for each input and hidden layer.
-
-  A unit with degree d can only receive input from units with degree < d. Output
-  units always have the same degree as their associated input unit.
-
-  Args:
-    input_dim: Number of inputs.
-    hidden_dims: list with the number of hidden units per layer. It does not
-      include the output layer. Each hidden unit size must be at least the size
-      of length (otherwise autoregressivity is not possible).
-    input_order: Order of degrees to the input units: 'random', 'left-to-right',
-      'right-to-left', or an array of an explicit order. For example,
-      'left-to-right' builds an autoregressive model
-      p(x) = p(x1) p(x2 | x1) ... p(xD | x<D).
-    hidden_order: Order of degrees to the hidden units: 'random',
-      'left-to-right'. If 'left-to-right', hidden units are allocated equally
-      (up to a remainder term) to each degree.
-  """
-  if (isinstance(input_order, str) and
-      input_order not in ('random', 'left-to-right', 'right-to-left')):
-    raise ValueError('Input order is not valid.')
-  if hidden_order not in ('random', 'left-to-right'):
-    raise ValueError('Hidden order is not valid.')
-
-  degrees = []
-  if isinstance(input_order, str):
-    input_degrees = np.arange(1, input_dim + 1)
-    if input_order == 'right-to-left':
-      input_degrees = np.flip(input_degrees, 0)
-    elif input_order == 'random':
-      np.random.shuffle(input_degrees)
-  else:
-    input_order = np.array(input_order)
-    if np.all(np.sort(input_order) != np.arange(1, input_dim + 1)):
-      raise ValueError('invalid input order')
-    input_degrees = input_order
-  degrees.append(input_degrees)
-
-  for units in hidden_dims:
-    if hidden_order == 'random':
-      min_prev_degree = min(np.min(degrees[-1]), input_dim - 1)
-      hidden_degrees = np.random.randint(
-          low=min_prev_degree, high=input_dim, size=units)
-    elif hidden_order == 'left-to-right':
-      hidden_degrees = (np.arange(units) % max(1, input_dim - 1) +
-                        min(1, input_dim - 1))
-    degrees.append(hidden_degrees)
-  return degrees
-
-
-def create_masks(input_dim,
-                 hidden_dims,
-                 input_order='left-to-right',
-                 hidden_order='left-to-right'):
-  """Returns a list of binary mask matrices respecting autoregressive ordering.
-
-  Args:
-    input_dim: Number of inputs.
-    hidden_dims: list with the number of hidden units per layer. It does not
-      include the output layer; those number of units will always be set to
-      input_dim downstream. Each hidden unit size must be at least the size of
-      length (otherwise autoregressivity is not possible).
-    input_order: Order of degrees to the input units: 'random', 'left-to-right',
-      'right-to-left', or an array of an explicit order. For example,
-      'left-to-right' builds an autoregressive model
-      p(x) = p(x1) p(x2 | x1) ... p(xD | x<D).
-    hidden_order: Order of degrees to the hidden units: 'random',
-      'left-to-right'. If 'left-to-right', hidden units are allocated equally
-      (up to a remainder term) to each degree.
-  """
-  degrees = create_degrees(input_dim, hidden_dims, input_order, hidden_order)
-  masks = []
-  # Create input-to-hidden and hidden-to-hidden masks.
-  for input_degrees, output_degrees in zip(degrees[:-1], degrees[1:]):
-    mask = tf.cast(input_degrees[:, np.newaxis] <= output_degrees, tf.float32)
-    masks.append(mask)
-
-  # Create hidden-to-output mask.
-  mask = tf.cast(degrees[-1][:, np.newaxis] < degrees[0], tf.float32)
-  masks.append(mask)
-  return masks
-
-
-def make_masked_initializer(mask):
-  initializer = tf.keras.initializers.glorot_uniform()
-  def masked_initializer(shape, dtype=None, partition_info=None):
-    return mask * initializer(shape, dtype, partition_info)
-  return masked_initializer
-
-
-def make_masked_constraint(mask):
-  constraint = tf.identity
-  def masked_constraint(x):
-    return mask * constraint(x)
-  return masked_constraint
-
-
-def sinkhorn(inputs, n_iters=20):
-  """Performs incomplete Sinkhorn normalization to inputs.
-
-  By a theorem by Sinkhorn and Knopp [1], a sufficiently well-behaved  matrix
-  with positive entries can be turned into a doubly-stochastic matrix
-  (i.e. its rows and columns add up to one) via the succesive row and column
-  normalization.
-  -To ensure positivity, the effective input to sinkhorn has to be
-  exp(inputs) (elementwise).
-  -However, for stability, sinkhorn works in the log-space. It is only at
-   return time that entries are exponentiated.
-
-  Code is adapted from Mena et al. [2].
-
-  [1] Richard Sinkhorn and Paul Knopp. Concerning nonnegative matrices and
-  doubly stochastic matrices. Pacific Journal of Mathematics, 1967.
-
-  [2] Gonzalo Mena, David Belanger, Scott Linderman, Jasper Snoek.
-  Learning latent permutations with Gumbel-Sinkhorn networks. International
-  Conference on Learning Representations, 2018.
-
-  Args:
-    inputs: A `Tensor` with shape `[..., vocab_size, vocab_size]`.
-    n_iters: Number of sinkhorn iterations (in practice, as little as 20
-      iterations are needed to achieve decent convergence for `vocab_size` ~100)
-
-  Returns:
-    outputs: A `Tensor` of close-to-doubly-stochastic matrices with shape
-      `[:, vocab_size, vocab_size]`.
-  """
-  vocab_size = tf.shape(inputs)[-1]
-  log_alpha = tf.reshape(inputs, [-1, vocab_size, vocab_size])
-
-  for _ in range(n_iters):
-    log_alpha -= tf.reshape(tf.reduce_logsumexp(log_alpha, axis=2),
-                            [-1, vocab_size, 1])
-    log_alpha -= tf.reshape(tf.reduce_logsumexp(log_alpha, axis=1),
-                            [-1, 1, vocab_size])
-  outputs = tf.exp(log_alpha)
-  return outputs
-
-
-class TransformedDistribution(tfp.distributions.Distribution):
-  """Distribution of f(x), where x ~ p(x) and f is reversible."""
-
-  def __init__(self, base, reversible_layer, name=None):
-    """Constructs a transformed distribution.
-
-    Args:
-      base: Base distribution.
-      reversible_layer: Callable with methods `reverse` and `log_det_jacobian`.
-      name: Name for scoping operations in the class.
-    """
-    self.base = base
-    self.reversible_layer = reversible_layer
-    if name is None:
-      name = reversible_layer.name + base.name
-    super(TransformedDistribution, self).__init__(
-        base.dtype,
-        base.reparameterization_type,
-        base.validate_args,
-        base.allow_nan_stats,
-        parameters=dict(locals()),
-        name=name)
-
-  def _event_shape_tensor(self):
-    return self.base.event_shape_tensor()
-
-  def _event_shape(self):
-    return self.base.event_shape
-
-  def _batch_shape_tensor(self):
-    return self.base.batch_shape_tensor()
-
-  def _batch_shape(self):
-    return self.base.batch_shape
-
-  def __getitem__(self, slices):
-    overrides = {'base': self.base[slices]}
-    return self.copy(**overrides)
-
-  def _call_sample_n(self, sample_shape, seed, name, **kwargs):
-    x = self.base.sample(sample_shape, seed, **kwargs)
-    y = self.reversible_layer(x)
-    return y
-
-  def _log_prob(self, value):
-    x = self.reversible_layer.reverse(value)
-    log_det_jacobian = self.reversible_layer.log_det_jacobian(value)
-    return self.base.log_prob(x) + log_det_jacobian
-
-  def _prob(self, value):
-    if not hasattr(self.base, '_prob'):
-      return tf.exp(self.log_prob(value))
-    x = self.reversible_layer.reverse(value)
-    log_det_jacobian = self.reversible_layer.log_det_jacobian(value)
-    return self.base.prob(x) * tf.exp(log_det_jacobian)
-
-  def _log_cdf(self, value):
-    x = self.reversible_layer.reverse(value)
-    return self.base.log_cdf(x)
-
-  def _cdf(self, value):
-    x = self.reversible_layer.reverse(value)
-    return self.base.cdf(x)
-
-  def _log_survival_function(self, value):
-    x = self.reversible_layer.reverse(value)
-    return self.base.log_survival_function(x)
-
-  def _survival_function(self, value):
-    x = self.reversible_layer.reverse(value)
-    return self.base.survival_function(x)
-
-  def _quantile(self, value):
-    inverse_cdf = self.base.quantile(value)
-    return self.reversible_layer(inverse_cdf)
-
-  def _entropy(self):
-    dummy = tf.zeros(
-        tf.concat([self.batch_shape_tensor(), self.event_shape_tensor()], 0),
-        dtype=self.dtype)
-    log_det_jacobian = self.reversible_layer.log_det_jacobian(dummy)
-    entropy = self.base.entropy() - log_det_jacobian
-    return entropy
-
-
-@ed.interceptable
-def TransformedRandomVariable(random_variable,  # pylint: disable=invalid-name
-                              reversible_layer,
-                              name=None,
-                              sample_shape=(),
-                              value=None):
-  """Random variable for f(x), where x ~ p(x) and f is reversible."""
-  return ed.RandomVariable(
-      distribution=TransformedDistribution(random_variable.distribution,
-                                           reversible_layer,
-                                           name=name),
-      sample_shape=sample_shape,
-      value=value)
diff --git a/tensor2tensor/layers/reversible_layers_test.py b/tensor2tensor/layers/reversible_layers_test.py
deleted file mode 100644
index d523e81b2..000000000
--- a/tensor2tensor/layers/reversible_layers_test.py
+++ /dev/null
@@ -1,590 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 reversible layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import parameterized
-import numpy as np
-
-from tensor2tensor.layers import reversible_layers as reversible
-from tensor2tensor.utils import test_utils
-
-import tensorflow as tf
-import tensorflow_probability as tfp
-tf.compat.v1.enable_eager_execution()
-
-
-def _log_prob(self, x):
-  """Re-implementation of OneHotCategorical._log_prob for gradients wrt x."""
-  logits = self.logits_parameter()
-  x = self._maybe_assert_valid_sample(x, dtype=logits.dtype)
-  if (not x.shape.is_fully_defined() or
-      not logits.shape.is_fully_defined() or
-      x.shape != logits.shape):
-    logits = tf.ones_like(x, dtype=logits.dtype) * logits
-    x = tf.ones_like(logits, dtype=x.dtype) * x
-
-  logits_shape = tf.shape(tf.reduce_sum(logits, -1))
-  logits_2d = tf.reshape(logits, [-1, self.event_size])
-  x_2d = tf.reshape(x, [-1, self.event_size])
-  ret = -tf.nn.softmax_cross_entropy_with_logits_v2(
-      labels=x_2d, logits=logits_2d)
-  ret = tf.reshape(ret, logits_shape)
-  return ret
-
-tfp.distributions.OneHotCategorical._log_prob = _log_prob  # monkey patch
-
-
-class ReversibleLayersTest(parameterized.TestCase, tf.test.TestCase):
-
-  @parameterized.parameters(
-      (False,),
-      (True,),
-  )
-  @test_utils.run_in_graph_and_eager_modes()
-  def testDiscreteAutoregressiveFlowCall(self, loc_only):
-    batch_size = 3
-    vocab_size = 79
-    length = 5
-    if loc_only:
-      units = vocab_size
-      network = reversible.MADE(units, [])
-    else:
-      units = 2 * vocab_size
-      mask = tf.reshape([0] * vocab_size + [-1e10] + [0] * (vocab_size - 1),
-                        [1, 1, 2 * vocab_size])
-      network_ = reversible.MADE(units, [])
-      network = lambda inputs: mask + network_(inputs)
-    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
-    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
-    layer = reversible.DiscreteAutoregressiveFlow(network, 1.)
-    outputs = layer(inputs)
-    self.evaluate(tf.global_variables_initializer())
-    outputs_val = self.evaluate(outputs)
-    self.assertEqual(outputs_val.shape, (batch_size, length, vocab_size))
-    self.assertAllGreaterEqual(outputs_val, 0)
-    self.assertAllLessEqual(outputs_val, vocab_size - 1)
-
-  @parameterized.parameters(
-      (False,),
-      (True,),
-  )
-  @test_utils.run_in_graph_and_eager_modes()
-  def testDiscreteAutoregressiveFlowSample(self, loc_only):
-    batch_size = 5
-    length = 2
-    vocab_size = 2
-    if loc_only:
-      units = vocab_size
-      network = reversible.MADE(units, [])
-    else:
-      units = 2 * vocab_size
-      mask = tf.reshape([0] * vocab_size + [-1e10] + [0] * (vocab_size - 1),
-                        [1, 1, 2 * vocab_size])
-      network_ = reversible.MADE(units, [])
-      network = lambda inputs: mask + network_(inputs)
-    layer = reversible.DiscreteAutoregressiveFlow(network, 1.)
-    logits = tf.tile(tf.random_normal([length, vocab_size])[tf.newaxis],
-                     [batch_size, 1, 1])
-    base = tfp.edward2.OneHotCategorical(logits=logits, dtype=tf.float32)
-    outputs = layer(base)
-    _ = outputs.value  # need to do this to instantiate tf.variables
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(outputs)
-    self.assertEqual(res.shape, (batch_size, length, vocab_size))
-    self.assertAllGreaterEqual(res, 0)
-    self.assertAllLessEqual(res, vocab_size - 1)
-
-  @parameterized.parameters(
-      (False,),
-      (True,),
-  )
-  @test_utils.run_in_graph_and_eager_modes()
-  def testDiscreteAutoregressiveFlowInverse(self, loc_only):
-    batch_size = 2
-    vocab_size = 79
-    length = 5
-    if loc_only:
-      units = vocab_size
-      network = reversible.MADE(units, [])
-    else:
-      units = 2 * vocab_size
-      mask = tf.reshape([0] * vocab_size + [-1e10] + [0] * (vocab_size - 1),
-                        [1, 1, 2 * vocab_size])
-      network_ = reversible.MADE(units, [])
-      network = lambda inputs: mask + network_(inputs)
-    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
-    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
-    layer = reversible.DiscreteAutoregressiveFlow(network, 1.)
-    rev_fwd_inputs = layer.reverse(layer(inputs))
-    fwd_rev_inputs = layer(layer.reverse(inputs))
-    self.evaluate(tf.global_variables_initializer())
-    inputs_val, rev_fwd_inputs_val, fwd_rev_inputs_val = self.evaluate(
-        [inputs, rev_fwd_inputs, fwd_rev_inputs])
-    self.assertAllClose(inputs_val, rev_fwd_inputs_val)
-    self.assertAllClose(inputs_val, fwd_rev_inputs_val)
-
-  @parameterized.parameters(
-      (False,),
-      (True,),
-  )
-  @test_utils.run_in_graph_and_eager_modes()
-  def testDiscreteAutoregressiveFlowRandomVariable(self, loc_only):
-    batch_size = 2
-    length = 4
-    vocab_size = 5
-    if loc_only:
-      units = vocab_size
-      network = reversible.MADE(units, [])
-    else:
-      units = 2 * vocab_size
-      mask = tf.reshape([0] * vocab_size + [-1e10] + [0] * (vocab_size - 1),
-                        [1, 1, 2 * vocab_size])
-      network_ = reversible.MADE(units, [])
-      network = lambda inputs: mask + network_(inputs)
-    base = tfp.edward2.OneHotCategorical(logits=tf.random_normal([batch_size,
-                                                                  length,
-                                                                  vocab_size]),
-                                         dtype=tf.float32)
-    flow = reversible.DiscreteAutoregressiveFlow(network, 1.)
-    flow_rv = flow(base)
-    self.assertEqual(flow_rv.dtype, tf.float32)
-
-    self.evaluate(tf.global_variables_initializer())
-    res = self.evaluate(flow_rv)
-    self.assertEqual(res.shape, (batch_size, length, vocab_size))
-    self.assertAllGreaterEqual(res, 0)
-    self.assertAllLessEqual(res, vocab_size - 1)
-
-    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
-    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
-    outputs = flow(inputs)
-    rev_outputs = flow.reverse(outputs)
-    inputs_val, rev_outputs_val = self.evaluate([inputs, rev_outputs])
-    self.assertAllClose(inputs_val, rev_outputs_val)
-
-    inputs_log_prob = base.distribution.log_prob(inputs)
-    outputs_log_prob = flow_rv.distribution.log_prob(outputs)
-    res1, res2 = self.evaluate([inputs_log_prob, outputs_log_prob])
-    self.assertEqual(res1.shape, (batch_size, length))
-    self.assertAllClose(res1, res2)
-
-  @parameterized.parameters(
-      (False,),
-      (True,),
-  )
-  @test_utils.run_in_graph_mode_only()
-  def testDiscreteAutoregressiveFlowReverseGradients(self, loc_only):
-    batch_size = 2
-    length = 4
-    vocab_size = 2
-    if loc_only:
-      units = vocab_size
-      network = reversible.MADE(units, [16, 16])
-    else:
-      units = 2 * vocab_size
-      mask = tf.reshape([0] * vocab_size + [-1e10] + [0] * (vocab_size - 1),
-                        [1, 1, 2 * vocab_size])
-      network_ = reversible.MADE(units, [16, 16])
-      network = lambda inputs: mask + network_(inputs)
-    base = tfp.edward2.OneHotCategorical(
-        logits=tf.random_normal([batch_size, length, vocab_size]))
-    flow = reversible.DiscreteAutoregressiveFlow(network, 1.)
-    flow_rv = flow(base)
-    features = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
-    features = tf.one_hot(features, depth=vocab_size, dtype=tf.float32)
-    loss = -tf.reduce_sum(flow_rv.distribution.log_prob(features))
-    grads = tf.gradients(loss, tf.trainable_variables())
-    self.evaluate(tf.global_variables_initializer())
-    _ = self.evaluate(grads)
-    for grad in grads:
-      self.assertIsNotNone(grad)
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testOneHotAddExactHard(self):
-    inputs = tf.constant([[0., 1., 0.],
-                          [0., 0., 1.]])
-    shift = tf.constant([[0., 1., 0.],
-                         [1., 0., 0.]])
-
-    outputs = reversible.one_hot_add(inputs, shift)
-    outputs_val = self.evaluate(outputs)
-    self.assertAllEqual(outputs_val, np.array([[0., 0., 1.],
-                                               [0., 0., 1.]]))
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testOneHotMinusExactHard(self):
-    inputs = tf.constant([[0., 1., 0.],
-                          [0., 0., 1.]])
-    shift = tf.constant([[0., 1., 0.],
-                         [1., 0., 0.]])
-
-    outputs = reversible.one_hot_minus(inputs, shift)
-    outputs_val = self.evaluate(outputs)
-    self.assertAllEqual(outputs_val, np.array([[1., 0., 0.],
-                                               [0., 0., 1.]]))
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testOneHotMultiplyExactHard(self):
-    inputs = tf.constant([[0., 1., 0.],
-                          [0., 0., 1.]])
-    scale = tf.constant([[0., 1., 0.],
-                         [0., 0., 1.]])
-
-    outputs = reversible.one_hot_multiply(inputs, scale)
-    outputs_val = self.evaluate(outputs)
-    self.assertAllEqual(outputs_val, np.array([[0., 1., 0.],
-                                               [0., 1., 0.]]))
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testOneHotAddExactSoft(self):
-    inputs = tf.constant([[0., 1., 0.],
-                          [0., 0., 1.]])
-    shift = tf.constant([[0.1, 0.6, 0.3],
-                         [0.2, 0.4, 0.4]])
-
-    outputs = reversible.one_hot_add(inputs, shift)
-
-    shift_zero = inputs
-    shift_one = np.array([[0., 0., 1.],
-                          [1., 0., 0.]])
-    shift_two = np.array([[1., 0., 0.],
-                          [0., 1., 0.]])
-    expected_outputs = (shift[..., 0][..., tf.newaxis] * shift_zero +
-                        shift[..., 1][..., tf.newaxis] * shift_one +
-                        shift[..., 2][..., tf.newaxis] * shift_two)
-
-    actual_outputs_val, expected_outputs_val = self.evaluate([
-        outputs, expected_outputs])
-    self.assertAllClose(actual_outputs_val, expected_outputs_val)
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testOneHotMinusExactSoft(self):
-    inputs = tf.constant([[0., 1., 0.],
-                          [0., 0., 1.]])
-    shift = tf.constant([[0.1, 0.6, 0.3],
-                         [0.2, 0.4, 0.4]])
-
-    outputs = reversible.one_hot_minus(inputs, shift)
-
-    shift_zero = inputs
-    shift_one = np.array([[1., 0., 0.],
-                          [0., 1., 0.]])
-    shift_two = np.array([[0., 0., 1.],
-                          [1., 0., 0.]])
-    expected_outputs = (shift[..., 0][..., tf.newaxis] * shift_zero +
-                        shift[..., 1][..., tf.newaxis] * shift_one +
-                        shift[..., 2][..., tf.newaxis] * shift_two)
-
-    actual_outputs_val, expected_outputs_val = self.evaluate([
-        outputs, expected_outputs])
-    self.assertAllEqual(actual_outputs_val, expected_outputs_val)
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testOneHotMultiplyExactSoft(self):
-    inputs = tf.constant([[0., 1., 0.],
-                          [0., 0., 1.]])
-    scale = tf.constant([[0.1, 0.6, 0.3],
-                         [0.2, 0.4, 0.4]])
-
-    outputs = reversible.one_hot_multiply(inputs, scale)
-
-    scale_zero = np.array([[0., 0., 0.],
-                           [0., 0., 0.]])
-    scale_one = inputs
-    scale_two = np.array([[0., 0., 1.],
-                          [0., 1., 0.]])
-    expected_outputs = (scale[..., 0][..., tf.newaxis] * scale_zero +
-                        scale[..., 1][..., tf.newaxis] * scale_one +
-                        scale[..., 2][..., tf.newaxis] * scale_two)
-
-    actual_outputs_val, expected_outputs_val = self.evaluate([
-        outputs, expected_outputs])
-    self.assertAllEqual(actual_outputs_val, expected_outputs_val)
-
-  @parameterized.parameters(
-      (reversible.one_hot_add,),
-      (reversible.one_hot_minus,),
-  )
-  @test_utils.run_in_graph_and_eager_modes()
-  def testOneHotAddShapeHard(self, one_hot_add_fn):
-    batch_size = 2
-    length = 4
-    vocab_size = 5
-    inputs = tf.random_uniform(
-        [batch_size, length], minval=0, maxval=vocab_size, dtype=tf.int32)
-    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
-    shift = tf.random_uniform(
-        [batch_size, length], minval=0, maxval=vocab_size, dtype=tf.int32)
-    shift = tf.one_hot(shift, depth=vocab_size)
-
-    outputs = one_hot_add_fn(inputs, shift)
-    outputs_val = self.evaluate(outputs)
-    self.assertEqual(outputs_val.shape, (batch_size, length, vocab_size))
-
-  @parameterized.parameters(
-      (reversible.one_hot_add,),
-      (reversible.one_hot_minus,),
-  )
-  @test_utils.run_in_graph_and_eager_modes()
-  def testOneHotAddShapeSoft(self, one_hot_add_fn):
-    batch_size = 2
-    length = 4
-    vocab_size = 5
-    inputs = tf.random_uniform([batch_size, length, vocab_size])
-    shift = tf.random_uniform([batch_size, length, vocab_size])
-
-    outputs = one_hot_add_fn(inputs, shift)
-    outputs_val = self.evaluate(outputs)
-    self.assertEqual(outputs_val.shape, (batch_size, length, vocab_size))
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testMultiplicativeInverse(self):
-    batch_size = 3
-    vocab_size = 79
-    length = 5
-    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
-    one_hot_inputs = tf.one_hot(inputs, depth=vocab_size)
-
-    one_hot_inv = reversible.multiplicative_inverse(one_hot_inputs, vocab_size)
-    inv_inputs = tf.argmax(one_hot_inv, axis=-1)
-    inputs_inv_inputs = tf.floormod(inputs * inv_inputs, vocab_size)
-    inputs_inv_inputs_val = self.evaluate(inputs_inv_inputs)
-    self.assertAllEqual(inputs_inv_inputs_val, np.ones((batch_size, length)))
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testDiscreteBipartiteFlowCall(self):
-    batch_size = 3
-    vocab_size = 79
-    length = 5
-    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
-    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
-    layer = reversible.DiscreteBipartiteFlow(
-        tf.identity,
-        mask=tf.random_uniform([length], minval=0, maxval=2, dtype=tf.int32),
-        temperature=1.)
-    outputs = layer(inputs)
-    self.evaluate(tf.global_variables_initializer())
-    outputs_val = self.evaluate(outputs)
-    self.assertEqual(outputs_val.shape, (batch_size, length, vocab_size))
-    self.assertAllGreaterEqual(outputs_val, 0)
-    self.assertAllLessEqual(outputs_val, vocab_size - 1)
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testDiscreteBipartiteFlowInverse(self):
-    batch_size = 2
-    vocab_size = 79
-    length = 5
-    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
-    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
-    layer = reversible.DiscreteBipartiteFlow(
-        tf.identity,
-        mask=tf.random_uniform([length], minval=0, maxval=2, dtype=tf.int32),
-        temperature=1.)
-    rev_fwd_inputs = layer.reverse(layer(inputs))
-    fwd_rev_inputs = layer(layer.reverse(inputs))
-    self.evaluate(tf.global_variables_initializer())
-    inputs_val, rev_fwd_inputs_val, fwd_rev_inputs_val = self.evaluate(
-        [inputs, rev_fwd_inputs, fwd_rev_inputs])
-    self.assertAllClose(inputs_val, rev_fwd_inputs_val)
-    self.assertAllClose(inputs_val, fwd_rev_inputs_val)
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testSinkhornAutoregressiveFlowCall(self):
-    batch_size = 3
-    vocab_size = 79
-    length = 5
-    units = vocab_size ** 2
-    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
-    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
-    layer = reversible.SinkhornAutoregressiveFlow(
-        reversible.MADE(units, []), 1.)
-    outputs = layer(inputs)
-    self.evaluate(tf.global_variables_initializer())
-    outputs_val = self.evaluate(outputs)
-    self.assertEqual(outputs_val.shape, (batch_size, length, vocab_size))
-    self.assertAllGreaterEqual(outputs_val, 0)
-    self.assertAllLessEqual(outputs_val, vocab_size - 1)
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testDiscreteSinkhornFlowInverse(self):
-    batch_size = 2
-    vocab_size = 79
-    length = 5
-    units = vocab_size ** 2
-    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
-    inputs = tf.one_hot(inputs, depth=vocab_size, dtype=tf.float32)
-    layer = reversible.SinkhornAutoregressiveFlow(
-        reversible.MADE(units, []), 1.)
-    rev_fwd_inputs = layer.reverse(layer(inputs))
-    fwd_rev_inputs = layer(layer.reverse(inputs))
-    self.evaluate(tf.global_variables_initializer())
-    inputs_val, rev_fwd_inputs_val, fwd_rev_inputs_val = self.evaluate(
-        [inputs, rev_fwd_inputs, fwd_rev_inputs])
-    self.assertAllEqual(inputs_val, rev_fwd_inputs_val)
-    self.assertAllEqual(inputs_val, fwd_rev_inputs_val)
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testApproximatelyStochastic(self):
-    rng = np.random.RandomState(0)
-    tf.set_random_seed(1)
-    for dims in [2, 5, 10]:
-      for batch_size in [1, 2, 10]:
-        log_alpha = rng.randn(batch_size, dims, dims)
-        result = reversible.sinkhorn(log_alpha)
-        result_val = self.evaluate(result)
-        self.assertAllClose(np.sum(result_val, 1),
-                            np.tile([1.0], (batch_size, dims)),
-                            atol=1e-3)
-        self.assertAllClose(np.sum(result_val, 2),
-                            np.tile([1.0], (batch_size, dims)),
-                            atol=1e-3)
-
-  @test_utils.run_in_graph_mode_only()
-  def test_soft_to_hard_permutation(self):
-    """The solution of the matching for the identity matrix is range(N)."""
-    dims = 10
-    identity = np.eye(dims)
-    result_matching = reversible.soft_to_hard_permutation(identity)
-    result_matching_val = self.evaluate(result_matching)
-    self.assertAllEqual(result_matching_val[0], np.eye(dims))
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testActNorm(self):
-    np.random.seed(83243)
-    batch_size = 25
-    length = 15
-    channels = 4
-    inputs = 3. + 0.8 * np.random.randn(batch_size, length, channels)
-    inputs = tf.cast(inputs, tf.float32)
-    layer = reversible.ActNorm()
-    outputs = layer(inputs)
-    mean, variance = tf.nn.moments(outputs, axes=[0, 1])
-    self.evaluate(tf.global_variables_initializer())
-    mean_val, variance_val = self.evaluate([mean, variance])
-    self.assertAllClose(mean_val, np.zeros(channels), atol=1e-3)
-    self.assertAllClose(variance_val, np.ones(channels), atol=1e-3)
-
-    inputs = 3. + 0.8 * np.random.randn(batch_size, length, channels)
-    inputs = tf.cast(inputs, tf.float32)
-    outputs = layer(inputs)
-    mean, variance = tf.nn.moments(outputs, axes=[0, 1])
-    self.evaluate(tf.global_variables_initializer())
-    mean_val, variance_val = self.evaluate([mean, variance])
-    self.assertAllClose(mean_val, np.zeros(channels), atol=0.25)
-    self.assertAllClose(variance_val, np.ones(channels), atol=0.25)
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testMADELeftToRight(self):
-    np.random.seed(83243)
-    batch_size = 2
-    length = 3
-    channels = 1
-    units = 5
-    network = reversible.MADE(units, [4], activation=tf.nn.relu)
-    inputs = tf.zeros([batch_size, length, channels])
-    outputs = network(inputs)
-
-    num_weights = sum([np.prod(weight.shape) for weight in network.weights])
-    # Disable lint error for open-source. pylint: disable=g-generic-assert
-    self.assertEqual(len(network.weights), 4)
-    # pylint: enable=g-generic-assert
-    self.assertEqual(num_weights, (3*1*4 + 4) + (4*3*5 + 3*5))
-
-    self.evaluate(tf.global_variables_initializer())
-    outputs_val = self.evaluate(outputs)
-    self.assertAllEqual(outputs_val[:, 0, :], np.zeros((batch_size, units)))
-    self.assertEqual(outputs_val.shape, (batch_size, length, units))
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testMADERightToLeft(self):
-    np.random.seed(1328)
-    batch_size = 2
-    length = 3
-    channels = 5
-    units = 1
-    network = reversible.MADE(units, [4, 3],
-                              input_order='right-to-left',
-                              activation=tf.nn.relu,
-                              use_bias=False)
-    inputs = tf.zeros([batch_size, length, channels])
-    outputs = network(inputs)
-
-    num_weights = sum([np.prod(weight.shape) for weight in network.weights])
-    # Disable lint error for open-source. pylint: disable=g-generic-assert
-    self.assertEqual(len(network.weights), 3)
-    # pylint: enable=g-generic-assert
-    self.assertEqual(num_weights, 3*5*4 + 4*3 + 3*3*1)
-
-    self.evaluate(tf.global_variables_initializer())
-    outputs_val = self.evaluate(outputs)
-    self.assertAllEqual(outputs_val[:, -1, :], np.zeros((batch_size, units)))
-    self.assertEqual(outputs_val.shape, (batch_size, length, units))
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testMADENoHidden(self):
-    np.random.seed(532)
-    batch_size = 2
-    length = 3
-    channels = 5
-    units = 4
-    network = reversible.MADE(units, [], input_order='left-to-right')
-    inputs = tf.zeros([batch_size, length, channels])
-    outputs = network(inputs)
-
-    num_weights = sum([np.prod(weight.shape) for weight in network.weights])
-    # Disable lint error for open-source. pylint: disable=g-generic-assert
-    self.assertEqual(len(network.weights), 2)
-    # pylint: enable=g-generic-assert
-    self.assertEqual(num_weights, 3*5*3*4 + 3*4)
-
-    self.evaluate(tf.global_variables_initializer())
-    outputs_val = self.evaluate(outputs)
-    self.assertAllEqual(outputs_val[:, 0, :], np.zeros((batch_size, units)))
-    self.assertEqual(outputs_val.shape, (batch_size, length, units))
-
-  @test_utils.run_in_graph_and_eager_modes()
-  def testTransformedRandomVariable(self):
-    class Exp(tf.keras.layers.Layer):
-      """Exponential activation function for reversible networks."""
-
-      def __call__(self, inputs, *args, **kwargs):
-        if not isinstance(inputs, tfp.edward2.RandomVariable):
-          return super(Exp, self).__call__(inputs, *args, **kwargs)
-        return reversible.TransformedRandomVariable(inputs, self)
-
-      def call(self, inputs):
-        return tf.exp(inputs)
-
-      def reverse(self, inputs):
-        return tf.log(inputs)
-
-      def log_det_jacobian(self, inputs):
-        return -tf.log(inputs)
-
-    x = tfp.edward2.Normal(0., 1.)
-    y = Exp()(x)
-    y_sample = self.evaluate(y.distribution.sample())
-    y_log_prob = self.evaluate(y.distribution.log_prob(y_sample))
-    self.assertGreater(y_sample, 0.)
-    self.assertTrue(np.isfinite(y_log_prob))
-
-
-if __name__ == '__main__':
-  tf.test.main()

From ba101df57edba2d10e943ca96bc49111fcb2ede6 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 23 Aug 2019 11:48:15 -0700
Subject: [PATCH 2312/2720] Save optimizer slots in PPO

PiperOrigin-RevId: 265102491
---
 tensor2tensor/trax/rl/ppo.py         | 38 ++++++++++++++++++----------
 tensor2tensor/trax/rl/ppo_trainer.py | 20 +++++++--------
 2 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index 530dda7d5..4156394a9 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -784,37 +784,47 @@ def compute_stats(reward_dict):
   }, state
 
 
-def maybe_restore_params(output_dir, policy_and_value_net_params, state):
-  """Maybe restore the params from the checkpoint dir.
+def maybe_restore_opt_state(output_dir, policy_and_value_opt_state,
+                            policy_and_value_state):
+  """Maybe restore the optimization state from the checkpoint dir.
+
+  Optimization state includes parameters and optimizer slots.
 
   Args:
     output_dir: Directory where saved model checkpoints are stored.
-    policy_and_value_net_params: Default params, returned if model is'nt found.
-    state: policy state.
+    policy_and_value_opt_state: Default optimization state, returned if model
+      isn't found.
+    policy_and_value_state: state of the policy and value network.
 
   Returns:
-    tuple (restore (bool), params, state, iter (int), opt_step (int)) where iter
-    is the epoch from which we restored the params, 0 is restore = False, and
-    opt_step is the total optimization step (sum of all optimization steps made
-    up to the current epoch).
+    tuple (restored (bool), opt_state, state, epoch (int),
+    opt_step (int)) where epoch is the epoch from which we restored the
+    optimization state, 0 is restored = False, and opt_step is the total
+    optimization step (sum of all optimization steps made up to the current
+    epoch).
   """
+  restored = False
+  epoch = 0
+  total_opt_step = 0
   model_files = gfile.glob(os.path.join(output_dir, "model-??????.pkl"))
   for model_file in reversed(sorted(model_files)):
     logging.info("Trying to restore model from %s", model_file)
     try:
       with gfile.GFile(model_file, "rb") as f:
-        loaded_policy_and_value_net_params, loaded_state, total_opt_step = (
+        policy_and_value_opt_state, policy_and_value_state, total_opt_step = (
             pickle.load(f))
-        policy_and_value_net_params = loaded_policy_and_value_net_params
-        state = loaded_state
       model_file_basename = os.path.basename(model_file)  # model-??????.pkl
-      i = int(filter(str.isdigit, model_file_basename))
-      return True, policy_and_value_net_params, state, i, total_opt_step
+      restored = True
+      epoch = int(filter(str.isdigit, model_file_basename))
+      break
     except EOFError as e:
       logging.error("Unable to load model from: %s with %s", model_file, e)
       # Try an older version.
       continue
-  return False, policy_and_value_net_params, state, 0, 0
+  return (
+      restored, policy_and_value_opt_state, policy_and_value_state, epoch,
+      total_opt_step,
+  )
 
 
 def write_eval_reward_summaries(reward_stats_by_mode, summary_writer, epoch):
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index d84cddb6e..2a9992e2e 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -153,22 +153,22 @@ def __init__(
     )
     self._policy_and_value_net_apply = jit(policy_and_value_net_apply)
 
-    # Maybe restore the policy params. If there is nothing to restore, then
-    # iteration = 0 and policy_and_value_net_params are returned as is.
-    (restored, policy_and_value_net_params, self._model_state, self._epoch,
-     self._total_opt_step) = ppo.maybe_restore_params(
-         output_dir, policy_and_value_net_params, self._model_state)
+    # Initialize the optimizer.
+    (policy_and_value_opt_state, self._policy_and_value_opt_update,
+     self._policy_and_value_get_params) = ppo.optimizer_fn(
+         policy_and_value_optimizer, policy_and_value_net_params)
+
+    # Maybe restore the optimization state. If there is nothing to restore, then
+    # iteration = 0 and policy_and_value_opt_state is returned as is.
+    (restored, self._policy_and_value_opt_state, self._model_state, self._epoch,
+     self._total_opt_step) = ppo.maybe_restore_opt_state(
+         output_dir, policy_and_value_opt_state, self._model_state)
 
     if restored:
       logging.info("Restored parameters from iteration [%d]", self._epoch)
       # We should start from the next iteration.
       self._epoch += 1
 
-    # Initialize the optimizer.
-    (self._policy_and_value_opt_state, self._policy_and_value_opt_update,
-     self._policy_and_value_get_params) = ppo.optimizer_fn(
-         policy_and_value_optimizer, policy_and_value_net_params)
-
     # Create summary writers and history.
     self._train_sw = jaxboard.SummaryWriter(
         os.path.join(self._output_dir, "train"))

From df665a3757d5cb90d112c198d8138b553c75fb14 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 23 Aug 2019 14:36:00 -0700
Subject: [PATCH 2313/2720] Pass in the correct output_dir to trainer_class.

Sometimes the output_dir needs to be formatted according to the env_name, the
work id, especially in sweeps.

This possibly got messup somewhere along the way.

PiperOrigin-RevId: 265136977
---
 tensor2tensor/trax/rl_trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/rl_trainer.py b/tensor2tensor/trax/rl_trainer.py
index 39833c5e0..a7a5514a9 100644
--- a/tensor2tensor/trax/rl_trainer.py
+++ b/tensor2tensor/trax/rl_trainer.py
@@ -201,7 +201,7 @@ def run_training_loop():
     logging.info("Starting the training loop.")
 
     trainer = trainer_class(
-        output_dir=FLAGS.output_dir,
+        output_dir=output_dir,
         train_env=train_env,
         eval_env=eval_env,
     )

From a21ad0077190b371f177cd05ec5d1475e1780d01 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Sat, 24 Aug 2019 15:41:34 -0700
Subject: [PATCH 2314/2720] Dump trajectories from RL algorithms

The code is included in BaseTrainer, which allows dumping trajectories from any RL algorithm. It is reused for collecting trajectories in SimPLe.

Also fixed training step calculation in SimPLe, which was wrong before.

PiperOrigin-RevId: 265264061
---
 tensor2tensor/trax/rl/base_trainer.py        |  65 ++++++++-
 tensor2tensor/trax/rl/base_trainer_test.py   | 144 +++++++++++++++++++
 tensor2tensor/trax/rl/simple_trainer.py      |  75 +++++++---
 tensor2tensor/trax/rl/simple_trainer_test.py |   1 +
 4 files changed, 263 insertions(+), 22 deletions(-)
 create mode 100644 tensor2tensor/trax/rl/base_trainer_test.py

diff --git a/tensor2tensor/trax/rl/base_trainer.py b/tensor2tensor/trax/rl/base_trainer.py
index dbdb652e2..3921ffd57 100644
--- a/tensor2tensor/trax/rl/base_trainer.py
+++ b/tensor2tensor/trax/rl/base_trainer.py
@@ -19,19 +19,40 @@
 from __future__ import division
 from __future__ import print_function
 
+import os
+
 from absl import logging
+import cloudpickle as pickle
 from tensorflow.io import gfile
 
 
 class BaseTrainer(object):
   """Base class for RL trainers."""
 
-  def __init__(self, train_env, eval_env, output_dir):
-    # Train and eval envs are settable.
+  def __init__(
+      self, train_env, eval_env, output_dir,
+      trajectory_dump_dir=None, trajectory_dump_min_count_per_shard=16,
+  ):
+    """Base class constructor.
+
+    Args:
+      train_env: EnvProblem to use for training. Settable.
+      eval_env: EnvProblem to use for evaluation. Settable.
+      output_dir: Directory to save checkpoints and metrics to.
+      trajectory_dump_dir: Directory to dump trajectories to. Trajectories
+        are saved in shards of name <epoch>.pkl under this directory. Settable.
+      trajectory_dump_min_count_per_shard: Minimum number of trajectories to
+        collect before dumping in a new shard. Sharding is for efficient
+        shuffling for model training in SimPLe.
+    """
     self.train_env = train_env
     self.eval_env = eval_env
     self._output_dir = output_dir
     gfile.makedirs(self._output_dir)
+    self.trajectory_dump_dir = trajectory_dump_dir
+    self._trajectory_dump_min_count_per_shard = (
+        trajectory_dump_min_count_per_shard)
+    self._trajectory_buffer = []
 
   @property
   def epoch(self):
@@ -49,10 +70,50 @@ def save(self):
   def flush_summaries(self):
     raise NotImplementedError
 
+  def dump_trajectories(self, force=False):
+    """Dumps trajectories in a new shard.
+
+    Should be called at most once per epoch.
+
+    Args:
+      force: (bool) Whether to complete unfinished trajectories and create
+        a new shard even if we have not reached the minimum size.
+    """
+    if self.trajectory_dump_dir is None:
+      return
+    gfile.makedirs(self.trajectory_dump_dir)
+
+    trajectories = self.train_env.trajectories
+    if force:
+      trajectories.complete_all_trajectories()
+
+    # complete_all_trajectories() also adds trajectories that were just reset.
+    # We don't want them since they have just the initial observation and no
+    # actions, so we filter them out.
+    def has_any_action(trajectory):
+      return (
+          trajectory.time_steps and trajectory.time_steps[0].action is not None)
+    self._trajectory_buffer.extend(
+        filter(has_any_action, trajectories.completed_trajectories))
+
+    trajectories.clear_completed_trajectories()
+    ready = (
+        len(self._trajectory_buffer) >=
+        self._trajectory_dump_min_count_per_shard
+    )
+    if ready or force:
+      shard_path = os.path.join(
+          self.trajectory_dump_dir, "{}.pkl".format(self.epoch))
+      with gfile.GFile(shard_path, "wb") as f:
+        pickle.dump(self._trajectory_buffer, f)
+      self._trajectory_buffer = []
+
   def training_loop(self, n_epochs):
     logging.info("Starting the RL training loop.")
     for _ in range(self.epoch, n_epochs):
       self.train_epoch()
+      self.dump_trajectories()
     self.save()
+    self.dump_trajectories(force=True)
     self.evaluate()
     self.flush_summaries()
diff --git a/tensor2tensor/trax/rl/base_trainer_test.py b/tensor2tensor/trax/rl/base_trainer_test.py
new file mode 100644
index 000000000..13307c77c
--- /dev/null
+++ b/tensor2tensor/trax/rl/base_trainer_test.py
@@ -0,0 +1,144 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.rl.base_trainer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+import cloudpickle as pickle
+import numpy as np
+
+from tensor2tensor.envs import gym_env_problem
+from tensor2tensor.trax.rl import base_trainer
+from tensorflow import test
+
+
+class FakeTrainer(base_trainer.BaseTrainer):
+  """Fake Trainer.
+
+  Adds one complete and one incomplete trajectory every epoch.
+  """
+
+  def __init__(self, *args, **kwargs):
+    super(FakeTrainer, self).__init__(*args, **kwargs)
+    self._epoch = 0
+    self._should_reset = True
+
+  @property
+  def epoch(self):
+    return self._epoch
+
+  def train_epoch(self):
+    trajectories = self.train_env.trajectories
+    if self._should_reset:
+      trajectories.reset(indices=np.arange(2), observations=np.zeros(2))
+    self._should_reset = False
+    trajectories.step(
+        observations=np.zeros(2),
+        raw_rewards=np.zeros(2),
+        processed_rewards=np.zeros(2),
+        dones=np.array([False, True]),
+        actions=np.zeros(2),
+    )
+    # Reset the trajectories that are done, as
+    # env_problem_utils.play_env_problem_with_policy does.
+    trajectories.reset(indices=np.array([1]), observations=np.zeros(1))
+    self._epoch += 1
+
+  def evaluate(self):
+    pass
+
+  def save(self):
+    pass
+
+  def flush_summaries(self):
+    pass
+
+
+class BaseTrainerTest(test.TestCase):
+
+  def _make_trainer(self, min_count_per_shard):
+    train_env = gym_env_problem.GymEnvProblem(
+        base_env_name="Acrobot-v1", batch_size=2)
+    eval_env = gym_env_problem.GymEnvProblem(
+        base_env_name="Acrobot-v1", batch_size=1)
+    temp_dir = self.get_temp_dir()
+    return FakeTrainer(
+        train_env, eval_env,
+        output_dir=temp_dir,
+        trajectory_dump_dir=temp_dir,
+        trajectory_dump_min_count_per_shard=min_count_per_shard,
+    )
+
+  def _assert_no_shard_exists(self, trajectory_dir):
+    self.assertFalse(os.listdir(trajectory_dir))
+
+  def _assert_single_shard_exists_and_has_trajectories(
+      self, trajectory_dir, expected_trajectory_lengths):
+    shard_filenames = os.listdir(trajectory_dir)
+    self.assertEqual(len(shard_filenames), 1)
+    shard_path = os.path.join(trajectory_dir, shard_filenames[0])
+    with open(shard_path, "rb") as f:
+      trajectories = pickle.load(f)
+    actual_trajectory_lengths = [
+        len(trajectory.time_steps) for trajectory in trajectories]
+    self.assertEqual(
+        list(sorted(actual_trajectory_lengths)),
+        list(sorted(expected_trajectory_lengths)),
+    )
+
+  def test_dumps_full_shard(self):
+    trainer = self._make_trainer(min_count_per_shard=2)
+    trajectory_dir = self.get_temp_dir()
+
+    # Add one complete trajectory to the buffer. Should not dump yet.
+    trainer.train_epoch()
+    trainer.dump_trajectories()
+    self._assert_no_shard_exists(trajectory_dir)
+
+    # Add the second complete trajectory. Now we should dump.
+    trainer.train_epoch()
+    trainer.dump_trajectories()
+    self._assert_single_shard_exists_and_has_trajectories(
+        trajectory_dir, [2, 2])
+
+  def test_dumps_incomplete_trajectories_when_force_is_true(self):
+    trainer = self._make_trainer(min_count_per_shard=2)
+    trajectory_dir = self.get_temp_dir()
+
+    # Add one complete and one incomplete trajectory to the buffer. Should dump.
+    trainer.train_epoch()
+    trainer.dump_trajectories(force=True)
+    self._assert_single_shard_exists_and_has_trajectories(
+        trajectory_dir, [2, 2])
+
+  def test_dumps_incomplete_shard_when_force_is_true(self):
+    trainer = self._make_trainer(min_count_per_shard=4)
+    trajectory_dir = self.get_temp_dir()
+
+    # Add one complete and one incomplete trajectory to the buffer. Should dump,
+    # even though we don't have a full shard yet.
+    trainer.train_epoch()
+    trainer.dump_trajectories(force=True)
+    self._assert_single_shard_exists_and_has_trajectories(
+        trajectory_dir, [2, 2])
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index 07b447188..ada78297f 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -24,11 +24,13 @@
 import random
 
 from absl import logging
+import cloudpickle as pickle
 from jax import numpy as np
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.rl import base_trainer
 from tensor2tensor.trax.rl import simulated_env_problem
+from tensorflow.io import gfile
 
 
 class SimPLe(base_trainer.BaseTrainer):
@@ -41,14 +43,18 @@ def __init__(
       output_dir,
       policy_trainer_class,
       n_real_epochs=10,
-      data_eval_frac=0.05,
+      data_eval_frac=0.125,
       model_train_batch_size=64,
+      n_model_train_steps=1000,
       simulated_env_problem_class=(
           simulated_env_problem.SerializedSequenceSimulatedEnvProblem),
       simulated_batch_size=16,
       n_simulated_epochs=1000,
+      trajectory_dump_dir=None,
+      **kwargs
   ):
-    super(SimPLe, self).__init__(train_env, eval_env, output_dir)
+    super(SimPLe, self).__init__(
+        train_env, eval_env, output_dir, **kwargs)
     self._policy_dir = os.path.join(output_dir, "policy")
     self._policy_trainer = policy_trainer_class(
         train_env=train_env,
@@ -57,9 +63,8 @@ def __init__(
     )
     self._n_real_epochs = n_real_epochs
     self._model_train_batch_size = model_train_batch_size
+    self._n_model_train_steps = n_model_train_steps
     self._data_eval_frac = data_eval_frac
-    self._train_trajectories = []
-    self._eval_trajectories = []
     self._model_dir = os.path.join(output_dir, "model")
     self._sim_env = simulated_env_problem_class(
         batch_size=None,
@@ -72,17 +77,26 @@ def __init__(
     )
     self._simulated_batch_size = simulated_batch_size
     self._n_simulated_epochs = n_simulated_epochs
-    self._epoch = 0
+
+    # If trajectory_dump_dir is not provided explicitly, save the trajectories
+    # in output_dir.
+    if trajectory_dump_dir is None:
+      trajectory_dump_dir = os.path.join(output_dir, "trajectories")
+    self._trajectory_dump_root_dir = trajectory_dump_dir
+
+    self._simple_epoch = 0
+    self._policy_epoch = 0
+    self._model_train_step = 0
 
   @property
   def epoch(self):
-    return self._epoch
+    return self._simple_epoch
 
   def train_epoch(self):
     self.collect_trajectories()
     self.train_model()
     self.train_policy()
-    self._epoch += 1
+    self._simple_epoch += 1
 
   def evaluate(self):
     self._policy_trainer.evaluate()
@@ -96,16 +110,29 @@ def flush_summaries(self):
     pass
 
   def collect_trajectories(self):
-    logging.info("Epoch %d: collecting data", self._epoch)
+    logging.info("Epoch %d: collecting data", self._simple_epoch)
 
     self._policy_trainer.train_env = self.train_env
-    self._policy_trainer.training_loop(self._n_real_epochs)
-    self.train_env.trajectories.complete_all_trajectories()
-    trajectories = self.train_env.trajectories.completed_trajectories
-    pivot = int(len(trajectories) * (1 - self._data_eval_frac))
-    self._train_trajectories.extend(trajectories[:pivot])
-    self._eval_trajectories.extend(trajectories[pivot:])
-    # TODO(pkozakowski): Save trajectories to disk. Support restoring.
+    self._policy_trainer.trajectory_dump_dir = os.path.join(
+        self._trajectory_dump_root_dir, str(self.epoch))
+    self._policy_epoch += self._n_real_epochs
+    self._policy_trainer.training_loop(self._policy_epoch)
+
+  def _load_trajectories(self, trajectory_dir):
+    train_trajectories = []
+    eval_trajectories = []
+    # Search the entire directory subtree for trajectories.
+    for (subdir, _, filenames) in gfile.walk(trajectory_dir):
+      for filename in filenames:
+        shard_path = os.path.join(subdir, filename)
+        with gfile.GFile(shard_path, "rb") as f:
+          trajectories = pickle.load(f)
+          pivot = int(len(trajectories) * (1 - self._data_eval_frac))
+          train_trajectories.extend(trajectories[:pivot])
+          eval_trajectories.extend(trajectories[pivot:])
+    assert train_trajectories, "Haven't found any training data."
+    assert eval_trajectories, "Haven't found any evaluation data."
+    return (train_trajectories, eval_trajectories)
 
   def _data_stream(self, trajectories, batch_size):
     def make_batch(examples):
@@ -139,12 +166,16 @@ def make_batch(examples):
         yield make_batch(example_list)
 
   def train_model(self):
-    logging.info("Epoch %d: training model", self._epoch)
+    logging.info("Epoch %d: training model", self._simple_epoch)
 
+    # Load data from all epochs.
+    # TODO(pkozakowski): Handle the case when the data won't fit in the memory.
+    (train_trajectories, eval_trajectories) = self._load_trajectories(
+        self._trajectory_dump_root_dir)
     train_stream = lambda: self._data_stream(  # pylint: disable=g-long-lambda
-        self._train_trajectories, self._model_train_batch_size)
+        train_trajectories, self._model_train_batch_size)
     eval_stream = lambda: self._data_stream(  # pylint: disable=g-long-lambda
-        self._eval_trajectories, self._model_train_batch_size)
+        eval_trajectories, self._model_train_batch_size)
     # Ignore n_devices for now.
     inputs = lambda _: trax_inputs.Inputs(  # pylint: disable=g-long-lambda
         train_stream=train_stream,
@@ -153,19 +184,23 @@ def train_model(self):
         input_shape=self._sim_env.model_input_shape,
         input_dtype=self._sim_env.model_input_dtype,
     )
+
+    self._model_train_step += self._n_model_train_steps
     trax.train(
         model=self._sim_env.model,
         inputs=inputs,
+        train_steps=self._model_train_step,
         output_dir=self._model_dir,
         has_weights=True,
     )
 
   def train_policy(self):
-    logging.info("Epoch %d: training policy", self._epoch)
+    logging.info("Epoch %d: training policy", self._simple_epoch)
 
     self._sim_env.initialize(
         batch_size=self._simulated_batch_size,
         history_stream=itertools.repeat(None),
     )
     self._policy_trainer.train_env = self._sim_env
-    self._policy_trainer.training_loop(self._n_simulated_epochs)
+    self._policy_epoch += self._n_simulated_epochs
+    self._policy_trainer.training_loop(self._policy_epoch)
diff --git a/tensor2tensor/trax/rl/simple_trainer_test.py b/tensor2tensor/trax/rl/simple_trainer_test.py
index 052f27000..927418855 100644
--- a/tensor2tensor/trax/rl/simple_trainer_test.py
+++ b/tensor2tensor/trax/rl/simple_trainer_test.py
@@ -70,6 +70,7 @@ def test_training_loop_acrobot(self):
             n_optimizer_steps=1,
         ),
         n_real_epochs=1,
+        data_eval_frac=0.5,
         model_train_batch_size=2,
         simulated_env_problem_class=functools.partial(
             simulated_env_problem.SerializedSequenceSimulatedEnvProblem,

From 172f8257b35d754bf360069d4f7369a35997e057 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sun, 25 Aug 2019 16:56:47 -0700
Subject: [PATCH 2315/2720] Fixes a bug which didn't swap input and target
 vocabulary size properly when the problem is reversed.

PiperOrigin-RevId: 265361799
---
 tensor2tensor/data_generators/problem.py      |  2 --
 tensor2tensor/data_generators/problem_test.py | 21 +++++++++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index e1cb5d943..2133d9356 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -985,8 +985,6 @@ def _reverse_problem_hparams(p_hparams):
     if "target" in feature_name and reversed_feature_name in p.vocab_size:
       reversed_vocab_size[feature_name] = p.vocab_size[reversed_feature_name]
       reversed_vocab_size[reversed_feature_name] = p.vocab_size[feature_name]
-    else:
-      reversed_vocab_size[feature_name] = p.vocab_size[feature_name]
 
   p.vocab_size = reversed_vocab_size
 
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 7bf0137a9..63ee900ff 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -189,6 +189,27 @@ def testServingInputFnUseTpu(self):
       self.assertEqual(output_shape[0], batch_size)
       self.assertEqual(output_shape[1], max_length)
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testInputAndTargetVocabSizesAreReversed(self):
+
+    class WasReversedTestProblem(problem_module.Problem):
+
+      def __init__(self, input_vocab_size, target_vocab_size, was_reversed):
+        super(WasReversedTestProblem, self).__init__(was_reversed, False)
+        self.input_vocab_size = input_vocab_size
+        self.target_vocab_size = target_vocab_size
+
+      def hparams(self, defaults, model_hparams):
+        hp = defaults
+        hp.vocab_size = {"targets": self.target_vocab_size,
+                         "inputs": self.input_vocab_size}
+
+    problem = WasReversedTestProblem(input_vocab_size=1,
+                                     target_vocab_size=3,
+                                     was_reversed=True)
+    p_hparams = problem.get_hparams()
+    self.assertEqual(p_hparams.vocab_size["inputs"], 3)
+    self.assertEqual(p_hparams.vocab_size["targets"], 1)
 
 if __name__ == "__main__":
   tf.test.main()

From 960b61d1cd58c843bbec46e3965436dfec6d60e3 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 26 Aug 2019 11:50:57 -0700
Subject: [PATCH 2316/2720] Delete tensor2tensor/layers/bayes_test.py from
 oss_tests.sh -- this was deleted recently.

PiperOrigin-RevId: 265506476
---
 oss_scripts/oss_tests.sh | 1 -
 1 file changed, 1 deletion(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 889236f08..0e3710f6f 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -98,7 +98,6 @@ set_status
 # These tests enable eager, so are tested separately.
 pytest --disable-warnings \
   tensor2tensor/data_generators/problem_test.py \
-  tensor2tensor/layers/bayes_test.py \
   tensor2tensor/layers/common_attention_test.py \
   tensor2tensor/layers/common_layers_test.py \
   tensor2tensor/layers/common_video_test.py \

From 23f03a2ee5fda24234e8cc0573e3f727435a9428 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 26 Aug 2019 12:57:47 -0700
Subject: [PATCH 2317/2720] Delete gaussian_process_test.py from oss_tests,
 this was deleted recently.

PiperOrigin-RevId: 265521065
---
 oss_scripts/oss_tests.sh | 1 -
 1 file changed, 1 deletion(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 0e3710f6f..553acb3ac 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -102,7 +102,6 @@ pytest --disable-warnings \
   tensor2tensor/layers/common_layers_test.py \
   tensor2tensor/layers/common_video_test.py \
   tensor2tensor/layers/discretization_test.py \
-  tensor2tensor/layers/gaussian_process_test.py \
   tensor2tensor/layers/latent_layers_test.py \
   tensor2tensor/layers/modalities_test.py \
   tensor2tensor/layers/ngram_test.py \

From 7fa4b1aa1a45f54b9a44ca40d59eb02a541eaaaa Mon Sep 17 00:00:00 2001
From: Tanguy Urvoy <tanguy.urvoy@orange.com>
Date: Mon, 26 Aug 2019 22:05:45 +0200
Subject: [PATCH 2318/2720] Update generator_utils.py (#1674)

Hi,
`isinstance(v[0], six.integer_types)` is False for `np.int64` type causing algorithmic_sort_problem data generation to fail.
---
 tensor2tensor/data_generators/generator_utils.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 43d4ff14d..4773db565 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -38,6 +38,7 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
+import numpy as np
 
 UNSHUFFLED_SUFFIX = "-unshuffled"
 
@@ -48,7 +49,7 @@ def to_example(dictionary):
   for (k, v) in six.iteritems(dictionary):
     if not v:
       raise ValueError("Empty generated field: %s" % str((k, v)))
-    if isinstance(v[0], six.integer_types):
+    if isinstance(v[0], six.integer_types) or isinstance(v[0], np.int64):
       features[k] = tf.train.Feature(int64_list=tf.train.Int64List(value=v))
     elif isinstance(v[0], float):
       features[k] = tf.train.Feature(float_list=tf.train.FloatList(value=v))

From 6f4f72f5ae50c4de75551f15b797f0432ce07d88 Mon Sep 17 00:00:00 2001
From: Yacine BENAFFANE <yacine.benaffane@gmail.com>
Date: Mon, 26 Aug 2019 16:08:05 -0400
Subject: [PATCH 2319/2720] Transformer tutorial (#1675)

* Add Colaboratory notebook

* Add a link to the notebook

* Add an Apache license
---
 README.md                                     |    2 +
 .../notebooks/Transformer_translate.ipynb     | 1102 +++++++++++++++++
 2 files changed, 1104 insertions(+)
 create mode 100644 tensor2tensor/notebooks/Transformer_translate.ipynb

diff --git a/README.md b/README.md
index d51eff0a9..f0ee8bb61 100644
--- a/README.md
+++ b/README.md
@@ -227,6 +227,8 @@ which is close to state-of-the art. If training on a single GPU, try the
 or larger data-sets (e.g., for English-French), try the big model
 with `--hparams_set=transformer_big`.
 
+See this [example](https://github.com/Styleoshin/tensor2tensor/blob/Transformer_tutorial/tensor2tensor/notebooks/Transformer_translate.ipynb) to know how the translation works.
+
 ## Basics
 
 ### Walkthrough
diff --git a/tensor2tensor/notebooks/Transformer_translate.ipynb b/tensor2tensor/notebooks/Transformer_translate.ipynb
new file mode 100644
index 000000000..07c350351
--- /dev/null
+++ b/tensor2tensor/notebooks/Transformer_translate.ipynb
@@ -0,0 +1,1102 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "name": "Transformer_translate.ipynb",
+      "version": "0.3.2",
+      "provenance": [],
+      "collapsed_sections": [],
+      "toc_visible": true,
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/Styleoshin/tensor2tensor/blob/Transformer_tutorial/tensor2tensor/notebooks/Transformer_translate.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e7PMze9tKHX9",
+        "colab_type": "text"
+      },
+      "source": [
+        "# Welcome to the [Tensor2Tensor](https://github.com/tensorflow/tensor2tensor) Colab\n",
+        "\n",
+        "Tensor2Tensor, or T2T for short, is a library of deep learning models and datasets designed to make deep learning more accessible and [accelerate ML research](https://research.googleblog.com/2017/06/accelerating-deep-learning-research.html). In this notebook we will see how to use this library for a translation task by exploring the necessary steps. We will see how to define a problem, generate the data, train the model and test the quality of it, and we will translate our sequences and we visualize the attention. We will also see how to download a pre-trained model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "KC8jNpnyKJdm",
+        "colab_type": "code",
+        "cellView": "form",
+        "colab": {}
+      },
+      "source": [
+        "#@title\n",
+        "# Copyright 2018 Google LLC.\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",
+        "# https://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."
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "AYUy570fKRcw",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "# Install deps\n",
+        "!pip install -q -U tensor2tensor"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "hEhFfyVNbB_D",
+        "colab_type": "text"
+      },
+      "source": [
+        "#1. Initialization\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "i23pCAVwegx3",
+        "colab_type": "text"
+      },
+      "source": [
+        "##1.1. Make some directories"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "oUf4e18_8E31",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "import tensorflow as tf\n",
+        "import os\n",
+        "\n",
+        "DATA_DIR = os.path.expanduser(\"/t2t/data\") # This folder contain the data\n",
+        "TMP_DIR = os.path.expanduser(\"/t2t/tmp\")\n",
+        "TRAIN_DIR = os.path.expanduser(\"/t2t/train\") # This folder contain the model\n",
+        "EXPORT_DIR = os.path.expanduser(\"/t2t/export\") # This folder contain the exported model for production\n",
+        "TRANSLATIONS_DIR = os.path.expanduser(\"/t2t/translation\") # This folder contain  all translated sequence\n",
+        "EVENT_DIR = os.path.expanduser(\"/t2t/event\") # Test the BLEU score\n",
+        "USR_DIR = os.path.expanduser(\"/t2t/user\") # This folder contains our data that we want to add\n",
+        " \n",
+        "tf.gfile.MakeDirs(DATA_DIR)\n",
+        "tf.gfile.MakeDirs(TMP_DIR)\n",
+        "tf.gfile.MakeDirs(TRAIN_DIR)\n",
+        "tf.gfile.MakeDirs(EXPORT_DIR)\n",
+        "tf.gfile.MakeDirs(TRANSLATIONS_DIR)\n",
+        "tf.gfile.MakeDirs(EVENT_DIR)\n",
+        "tf.gfile.MakeDirs(USR_DIR)"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HIuzsMzgbLv9",
+        "colab_type": "text"
+      },
+      "source": [
+        "## 1.2. Init parameters\n",
+        "\n",
+        "\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "ZQaURmfKBGus",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "PROBLEM = \"translate_enfr_wmt32k\" # We chose a problem translation English to French with 32.768 vocabulary\n",
+        "MODEL = \"transformer\" # Our model\n",
+        "HPARAMS = \"transformer_big\" # Hyperparameters for the model by default \n",
+        "                            # If you have a one gpu, use transformer_big_single_gpu"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "EikK-hW5m-ax",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "#Show all problems and models \n",
+        "\n",
+        "from tensor2tensor.utils import registry\n",
+        "from tensor2tensor import problems\n",
+        "\n",
+        "problems.available() #Show all problems\n",
+        "registry.list_models() #Show all registered models\n",
+        "\n",
+        "#or\n",
+        "\n",
+        "#Command line\n",
+        "!t2t-trainer --registry_help #Show all problems\n",
+        "!t2t-trainer --problems_help #Show all models"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "78kBAIMQbeO6",
+        "colab_type": "text"
+      },
+      "source": [
+        "# 2. Data generation \n",
+        "\n",
+        "Generate the data (download the dataset and generate the data).\n",
+        "\n",
+        "---\n",
+        "\n",
+        " You can choose between command line or code."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "CrDy3V7ibpQH",
+        "colab_type": "text"
+      },
+      "source": [
+        "## 2.1. Generate with terminal\n",
+        "For more information: https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_datagen.py"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "0Dfr8nFXmg1o",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!t2t-datagen \\\n",
+        "  --data_dir=$DATA_DIR \\\n",
+        "  --tmp_dir=$TMP_DIR \\\n",
+        "  --problem=$PROBLEM \\\n",
+        "  --t2t_usr_dir=$USR_DIR"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tMvCiiBtbuzh",
+        "colab_type": "text"
+      },
+      "source": [
+        "## 2.2. Generate with code"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "Of5bHYVJmbwH",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "t2t_problem = problems.problem(PROBLEM)\n",
+        "t2t_problem.generate_data(DATA_DIR, TMP_DIR) "
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "UkSwoqBzb47T",
+        "colab_type": "text"
+      },
+      "source": [
+        "# 3. Train the model\n",
+        "\n",
+        "\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1JVF2PJn7ByQ",
+        "colab_type": "text"
+      },
+      "source": [
+        "##3.1. Init parameters\n",
+        "\n",
+        "You can choose between command line or code.\n",
+        "\n",
+        "---\n",
+        "\n",
+        " batch_size :  a great value of preference.\n",
+        "\n",
+        "---\n",
+        "train_steps : research paper mentioned 300k steps with 8 gpu on big transformer. So if you have 1 gpu, you will need to train the model x8 more. (https://arxiv.org/abs/1706.03762 for more information).\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "yw6HgVWA7AQF",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "train_steps = 300000 # Total number of train steps for all Epochs\n",
+        "eval_steps = 100 # Number of steps to perform for each evaluation\n",
+        "batch_size = 4096\n",
+        "save_checkpoints_steps = 1000\n",
+        "ALPHA = 0.1\n",
+        "schedule = \"continuous_train_and_eval\""
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ze_YvVnIfD8z",
+        "colab_type": "text"
+      },
+      "source": [
+        "You can choose schedule :\n",
+        " \n",
+        "\n",
+        "*  train. Bad quality\n",
+        "*  continuous_train_and_eval (default)\n",
+        "*   train_and_eval\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-zAub7Ggb8tj",
+        "colab_type": "text"
+      },
+      "source": [
+        "##3.2. Train with terminal\n",
+        "https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_trainer.py\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "kSYAi4BsnpSD",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!t2t-trainer \\\n",
+        "  --data_dir=$DATA_DIR \\\n",
+        "  --problem=$PROBLEM \\\n",
+        "  --model=$MODEL \\\n",
+        "  --hparams_set=$HPARAMS \\\n",
+        "  --hparams=\"batch_size=$batch_size\" \\\n",
+        "  --schedule=$schedule\\\n",
+        "  --output_dir=$TRAIN_DIR \\\n",
+        "  --train_steps=$train_steps \\\n",
+        "  --worker-gpu=1 \\ \n",
+        "  --eval_steps=$eval_steps "
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bNfNBWtNVMwO",
+        "colab_type": "text"
+      },
+      "source": [
+        "  --worker-gpu = 1, for train on 1 gpu (facultative).\n",
+        "\n",
+        "---\n",
+        "\n",
+        "For distributed training see: https://github.com/tensorflow/tensor2tensor/blob/master/docs/distributed_training.md\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "nnSoC1AUcLG6",
+        "colab_type": "text"
+      },
+      "source": [
+        "##3.3. Train with code\n",
+        "create_hparams : https://github.com/tensorflow/tensor2tensor/blob/28adf2690c551ef0f570d41bef2019d9c502ec7e/tensor2tensor/utils/hparams_lib.py#L42\n",
+        "\n",
+        "---\n",
+        "Change hyper parameters :\n",
+        "https://github.com/tensorflow/tensor2tensor/blob/28adf2690c551ef0f570d41bef2019d9c502ec7e/tensor2tensor/models/transformer.py#L1627\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "RJ91vQ2hyIPx",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "from tensor2tensor.utils.trainer_lib import create_run_config, create_experiment\n",
+        "from tensor2tensor.utils.trainer_lib import create_hparams\n",
+        "from tensor2tensor.utils import registry\n",
+        "from tensor2tensor import models\n",
+        "from tensor2tensor import problems\n",
+        "\n",
+        "# Init Hparams object from T2T Problem\n",
+        "hparams = create_hparams(HPARAMS)\n",
+        "\n",
+        "# Make Changes to Hparams\n",
+        "hparams.batch_size = batch_size\n",
+        "hparams.learning_rate = ALPHA\n",
+        "#hparams.max_length = 256\n",
+        "\n",
+        "# Can see all Hparams with code below\n",
+        "#print(json.loads(hparams.to_json())"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KZX1cwK3TEXs",
+        "colab_type": "text"
+      },
+      "source": [
+        "create_run_config : https://github.com/tensorflow/tensor2tensor/blob/28adf2690c551ef0f570d41bef2019d9c502ec7e/tensor2tensor/utils/trainer_lib.py#L105\n",
+        "\n",
+        "---\n",
+        "\n",
+        "\n",
+        "create_experiment : https://github.com/tensorflow/tensor2tensor/blob/28adf2690c551ef0f570d41bef2019d9c502ec7e/tensor2tensor/utils/trainer_lib.py#L611"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "yByKcs7XvAXL",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "RUN_CONFIG = create_run_config(\n",
+        "      model_dir=TRAIN_DIR,\n",
+        "      model_name=MODEL,\n",
+        "      save_checkpoints_steps= save_checkpoints_steps\n",
+        ")\n",
+        "\n",
+        "tensorflow_exp_fn = create_experiment(\n",
+        "        run_config=RUN_CONFIG,\n",
+        "        hparams=hparams,\n",
+        "        model_name=MODEL,\n",
+        "        problem_name=PROBLEM,\n",
+        "        data_dir=DATA_DIR, \n",
+        "        train_steps=train_steps, \n",
+        "        eval_steps=eval_steps, \n",
+        "        #use_xla=True # For acceleration\n",
+        "    ) \n",
+        "\n",
+        "tensorflow_exp_fn.train_and_evaluate()"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "03xuR70jce_2",
+        "colab_type": "text"
+      },
+      "source": [
+        "#4. See the BLEU score"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "MiwyVWPhhGrk",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "#INIT FILE FOR TRANSLATE\n",
+        "\n",
+        "SOURCE_TEST_TRANSLATE_DIR = TMP_DIR+\"/dev/newstest2014-fren-src.en.sgm\"\n",
+        "REFERENCE_TEST_TRANSLATE_DIR = TMP_DIR+\"/dev/newstest2014-fren-ref.en.sgm\"\n",
+        "BEAM_SIZE=1"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "agnSg_89cr63",
+        "colab_type": "text"
+      },
+      "source": [
+        "##4.1. Translate all\n",
+        "https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_translate_all.py"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "Jrt5fwqsg3pl",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!t2t-translate-all \\\n",
+        "  --source=$SOURCE_TEST_TRANSLATE_DIR \\\n",
+        "  --model_dir=$TRAIN_DIR \\\n",
+        "  --translations_dir=$TRANSLATIONS_DIR \\\n",
+        "  --data_dir=$DATA_DIR \\\n",
+        "  --problem=$PROBLEM \\\n",
+        "  --hparams_set=$HPARAMS \\\n",
+        "  --output_dir=$TRAIN_DIR \\\n",
+        "  --t2t_usr_dir=$USR_DIR \\\n",
+        "  --beam_size=$BEAM_SIZE \\\n",
+        "  --model=$MODEL"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "O-pKKU2Acv8Q",
+        "colab_type": "text"
+      },
+      "source": [
+        "##4.2. Test the BLEU score\n",
+        "The BLEU score for all translations: https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_bleu.py#L68\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "EULP9TdPc58d",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!t2t-bleu \\\n",
+        "   --translations_dir=$TRANSLATIONS_DIR \\\n",
+        "   --model_dir=$TRAIN_DIR \\\n",
+        "   --data_dir=$DATA_DIR \\\n",
+        "   --problem=$PROBLEM \\\n",
+        "   --hparams_set=$HPARAMS \\\n",
+        "   --source=$SOURCE_TEST_TRANSLATE_DIR \\\n",
+        "   --reference=$REFERENCE_TEST_TRANSLATE_DIR \\\n",
+        "   --event_dir=$EVENT_DIR"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "13j50bpAc-bM",
+        "colab_type": "text"
+      },
+      "source": [
+        "#5. Prediction of sentence\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8WHPnqxhdQl6",
+        "colab_type": "text"
+      },
+      "source": [
+        "##5.1. Predict with terminal\n",
+        "https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_decoder.py"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "3SD-XhImnwpo",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!echo \"the business of the house\" > \"inputs.en\"\n",
+        "!echo -e \"les affaires de la maison\" > \"reference.fr\" # You can add other references\n",
+        "\n",
+        "!t2t-decoder \\\n",
+        "  --data_dir=$DATA_DIR \\\n",
+        "  --problem=$PROBLEM \\\n",
+        "  --model=$MODEL \\\n",
+        "  --hparams_set=$HPARAMS \\\n",
+        "  --output_dir=$TRAIN_DIR \\\n",
+        "  --decode_hparams=\"beam_size=1,alpha=$ALPHA\" \\\n",
+        "  --decode_from_file=\"inputs.en\" \\\n",
+        "  --decode_to_file=\"outputs.fr\"\n",
+        "\n",
+        "# See the translations\n",
+        "!cat outputs.fr"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "sGOC25N4dWdM",
+        "colab_type": "text"
+      },
+      "source": [
+        "##5.2. Predict with code"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "S6u4QmhPIbDx",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "import tensorflow as tf\n",
+        "\n",
+        "#After training the model, re-run the environment but run this code in first, then predict.\n",
+        "\n",
+        "tfe = tf.contrib.eager\n",
+        "tfe.enable_eager_execution()\n",
+        "Modes = tf.estimator.ModeKeys"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "PaCkILfjz9x3",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "#Config\n",
+        "\n",
+        "from tensor2tensor import models\n",
+        "from tensor2tensor import problems\n",
+        "from tensor2tensor.layers import common_layers\n",
+        "from tensor2tensor.utils import trainer_lib\n",
+        "from tensor2tensor.utils import t2t_model\n",
+        "from tensor2tensor.utils import registry\n",
+        "from tensor2tensor.utils import metrics\n",
+        "import numpy as np\n",
+        "\n",
+        "enfr_problem = problems.problem(PROBLEM)\n",
+        "\n",
+        "# Copy the vocab file locally so we can encode inputs and decode model outputs\n",
+        "vocab_name = \"vocab.translate_enfr_wmt32k.32768.subwords\"\n",
+        "vocab_file = os.path.join(DATA_DIR, vocab_name)\n",
+        "\n",
+        "# Get the encoders from the problem\n",
+        "encoders = enfr_problem.feature_encoders(DATA_DIR)\n",
+        "\n",
+        "ckpt_path = tf.train.latest_checkpoint(os.path.join(TRAIN_DIR))\n",
+        "print(ckpt_path)\n",
+        "\n",
+        "def translate(inputs):\n",
+        "  encoded_inputs = encode(inputs)\n",
+        "  with tfe.restore_variables_on_create(ckpt_path):\n",
+        "    model_output = translate_model.infer(encoded_inputs)[\"outputs\"]\n",
+        "  return decode(model_output)\n",
+        "\n",
+        "def encode(input_str, output_str=None):\n",
+        "  \"\"\"Input str to features dict, ready for inference\"\"\"\n",
+        "  inputs = encoders[\"inputs\"].encode(input_str) + [1]  # add EOS id\n",
+        "  batch_inputs = tf.reshape(inputs, [1, -1, 1])  # Make it 3D.\n",
+        "  return {\"inputs\": batch_inputs}\n",
+        "\n",
+        "def decode(integers):\n",
+        "  \"\"\"List of ints to str\"\"\"\n",
+        "  integers = list(np.squeeze(integers))\n",
+        "  if 1 in integers:\n",
+        "    integers = integers[:integers.index(1)]\n",
+        "  return encoders[\"inputs\"].decode(np.squeeze(integers))"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "5zE8yHLUA2He",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "#Predict \n",
+        "\n",
+        "hparams = trainer_lib.create_hparams(HPARAMS, data_dir=DATA_DIR, problem_name=PROBLEM)\n",
+        "translate_model = registry.model(MODEL)(hparams, Modes.PREDICT)\n",
+        "\n",
+        "inputs = \"the aniamal didn't cross the river because it was too tired\"\n",
+        "ref = \"l'animal n'a pas traversé la rue parcequ'il etait trop fatigué\" ## this just a reference for evaluate the quality of the traduction\n",
+        "outputs = translate(inputs)\n",
+        "\n",
+        "print(\"Inputs: %s\" % inputs)\n",
+        "print(\"Outputs: %s\" % outputs)\n",
+        "\n",
+        "file_input = open(\"outputs.fr\",\"w+\")\n",
+        "file_input.write(outputs)\n",
+        "file_input.close()\n",
+        "\n",
+        "file_output = open(\"reference.fr\",\"w+\")\n",
+        "file_output.write(ref)\n",
+        "file_output.close()"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "y6jbQ6FoRsmG",
+        "colab_type": "text"
+      },
+      "source": [
+        "##5.3. Evaluate the BLEU Score\n",
+        "BLEU score for a sequence translation: https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_bleu.py#L24"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "il2oevmXRrbf",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!t2t-bleu \\\n",
+        "    --translation=outputs.fr \\\n",
+        "    --reference=reference.fr"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "FXegHzD1I67e",
+        "colab_type": "text"
+      },
+      "source": [
+        "#6. Attention visualization\n",
+        "We need to have a predicted sentence with code."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ISHauPT8I-3S",
+        "colab_type": "text"
+      },
+      "source": [
+        "##6.1. Attention utils\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "2RHCTrc9I55K",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "from tensor2tensor.visualization import attention\n",
+        "from tensor2tensor.data_generators import text_encoder\n",
+        "\n",
+        "SIZE = 35\n",
+        "\n",
+        "def encode_eval(input_str, output_str):\n",
+        "  inputs = tf.reshape(encoders[\"inputs\"].encode(input_str) + [1], [1, -1, 1, 1])  # Make it 3D.\n",
+        "  outputs = tf.reshape(encoders[\"inputs\"].encode(output_str) + [1], [1, -1, 1, 1])  # Make it 3D.\n",
+        "  return {\"inputs\": inputs, \"targets\": outputs}\n",
+        "\n",
+        "def get_att_mats():\n",
+        "  enc_atts = []\n",
+        "  dec_atts = []\n",
+        "  encdec_atts = []\n",
+        "\n",
+        "  for i in range(hparams.num_hidden_layers):\n",
+        "    enc_att = translate_model.attention_weights[\n",
+        "      \"transformer/body/encoder/layer_%i/self_attention/multihead_attention/dot_product_attention\" % i][0]\n",
+        "    dec_att = translate_model.attention_weights[\n",
+        "      \"transformer/body/decoder/layer_%i/self_attention/multihead_attention/dot_product_attention\" % i][0]\n",
+        "    encdec_att = translate_model.attention_weights[\n",
+        "      \"transformer/body/decoder/layer_%i/encdec_attention/multihead_attention/dot_product_attention\" % i][0]\n",
+        "    enc_atts.append(resize(enc_att))\n",
+        "    dec_atts.append(resize(dec_att))\n",
+        "    encdec_atts.append(resize(encdec_att))\n",
+        "  return enc_atts, dec_atts, encdec_atts\n",
+        "\n",
+        "def resize(np_mat):\n",
+        "  # Sum across heads\n",
+        "  np_mat = np_mat[:, :SIZE, :SIZE]\n",
+        "  row_sums = np.sum(np_mat, axis=0)\n",
+        "  # Normalize\n",
+        "  layer_mat = np_mat / row_sums[np.newaxis, :]\n",
+        "  lsh = layer_mat.shape\n",
+        "  # Add extra dim for viz code to work.\n",
+        "  layer_mat = np.reshape(layer_mat, (1, lsh[0], lsh[1], lsh[2]))\n",
+        "  return layer_mat\n",
+        "\n",
+        "def to_tokens(ids):\n",
+        "  ids = np.squeeze(ids)\n",
+        "  subtokenizer = hparams.problem_hparams.vocabulary['targets']\n",
+        "  tokens = []\n",
+        "  for _id in ids:\n",
+        "    if _id == 0:\n",
+        "      tokens.append('<PAD>')\n",
+        "    elif _id == 1:\n",
+        "      tokens.append('<EOS>')\n",
+        "    elif _id == -1:\n",
+        "      tokens.append('<NULL>')\n",
+        "    else:\n",
+        "        tokens.append(subtokenizer._subtoken_id_to_subtoken_string(_id))\n",
+        "  return tokens\n",
+        "\n",
+        "def call_html():\n",
+        "  import IPython\n",
+        "  display(IPython.core.display.HTML('''\n",
+        "        <script src=\"/static/components/requirejs/require.js\"></script>\n",
+        "        <script>\n",
+        "          requirejs.config({\n",
+        "            paths: {\n",
+        "              base: '/static/base',\n",
+        "              \"d3\": \"https://cdnjs.cloudflare.com/ajax/libs/d3/3.5.8/d3.min\",\n",
+        "              jquery: '//ajax.googleapis.com/ajax/libs/jquery/2.0.0/jquery.min',\n",
+        "            },\n",
+        "          });\n",
+        "        </script>\n",
+        "        '''))"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9PGwUbJuJHJS",
+        "colab_type": "text"
+      },
+      "source": [
+        "##6.2 Display Attention"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "ijTOlrt8JI4t",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "import numpy as np\n",
+        "\n",
+        "# Convert inputs and outputs to subwords\n",
+        "\n",
+        "inp_text = to_tokens(encoders[\"inputs\"].encode(inputs))\n",
+        "out_text = to_tokens(encoders[\"inputs\"].encode(outputs))\n",
+        "\n",
+        "hparams = trainer_lib.create_hparams(HPARAMS, data_dir=DATA_DIR, problem_name=PROBLEM)\n",
+        "\n",
+        "# Run eval to collect attention weights\n",
+        "example = encode_eval(inputs, outputs)\n",
+        "with tfe.restore_variables_on_create(tf.train.latest_checkpoint(ckpt_path)):\n",
+        "  translate_model.set_mode(Modes.EVAL)\n",
+        "  translate_model(example)\n",
+        "# Get normalized attention weights for each layer\n",
+        "enc_atts, dec_atts, encdec_atts = get_att_mats()\n",
+        "\n",
+        "call_html()\n",
+        "attention.show(inp_text, out_text, enc_atts, dec_atts, encdec_atts)"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "r8yAQUDZdm1p",
+        "colab_type": "text"
+      },
+      "source": [
+        "#7. Export the model\n",
+        "For more information: https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/serving"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "c2yulC7J8_I9",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "#export Model\n",
+        "!t2t-exporter \\\n",
+        "  --data_dir=$DATA_DIR \\\n",
+        "  --output_dir=$TRAIN_DIR \\\n",
+        "  --problem=$PROBLEM \\\n",
+        "  --model=$MODEL \\\n",
+        "  --hparams_set=$HPARAMS \\\n",
+        "  --decode_hparams=\"beam_size=1,alpha=$ALPHA\" \\\n",
+        "  --export_dir=$EXPORT_DIR"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2ltjEr3JX5-e",
+        "colab_type": "text"
+      },
+      "source": [
+        "#8.Load pretrained model from Google Storage\n",
+        "We use the pretrained model En-De translation."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "QgY3Fw261bZC",
+        "colab_type": "text"
+      },
+      "source": [
+        "##8.1. See existing content storaged"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "7P7aJClG0t8c",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "print(\"checkpoint: \")\n",
+        "!gsutil ls \"gs://tensor2tensor-checkpoints\"\n",
+        "\n",
+        "print(\"data: \")\n",
+        "!gsutil ls \"gs://tensor2tensor-data\""
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "wP8jrR5bbu7e",
+        "colab_type": "text"
+      },
+      "source": [
+        "##8.2. Init model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "AnYU7lrazkMm",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "PROBLEM_PRETRAINED = \"translate_ende_wmt32k\"\n",
+        "MODEL_PRETRAINED = \"transformer\" \n",
+        "HPARAMS_PRETRAINED = \"transformer_base\""
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DTgPvq4q1VAr",
+        "colab_type": "text"
+      },
+      "source": [
+        "##8.3. Load content from google storage"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "FrxOAVcyinll",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "import tensorflow as tf\n",
+        "import os\n",
+        "\n",
+        "\n",
+        "DATA_DIR_PRETRAINED = os.path.expanduser(\"/t2t/data_pretrained\")\n",
+        "CHECKPOINT_DIR_PRETRAINED = os.path.expanduser(\"/t2t/checkpoints_pretrained\")\n",
+        "\n",
+        "tf.gfile.MakeDirs(DATA_DIR_PRETRAINED)\n",
+        "tf.gfile.MakeDirs(CHECKPOINT_DIR_PRETRAINED)\n",
+        "\n",
+        "\n",
+        "gs_data_dir = \"gs://tensor2tensor-data/\"\n",
+        "vocab_name = \"vocab.translate_ende_wmt32k.32768.subwords\"\n",
+        "vocab_file = os.path.join(gs_data_dir, vocab_name)\n",
+        "\n",
+        "gs_ckpt_dir = \"gs://tensor2tensor-checkpoints/\"\n",
+        "ckpt_name = \"transformer_ende_test\"\n",
+        "gs_ckpt = os.path.join(gs_ckpt_dir, ckpt_name)\n",
+        "\n",
+        "TRAIN_DIR_PRETRAINED = os.path.join(CHECKPOINT_DIR_PRETRAINED, ckpt_name)\n",
+        "\n",
+        "!gsutil cp {vocab_file} {DATA_DIR_PRETRAINED}\n",
+        "!gsutil -q cp -R {gs_ckpt} {CHECKPOINT_DIR_PRETRAINED}\n",
+        "\n",
+        "CHECKPOINT_NAME_PRETRAINED = tf.train.latest_checkpoint(TRAIN_DIR_PRETRAINED) # for translate with code\n"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "LP6cro9Xbygf",
+        "colab_type": "text"
+      },
+      "source": [
+        "##8.4. Translate"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "CBoNpy5HbzoF",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!echo \"the business of the house\" > \"inputs.en\"\n",
+        "!echo -e \"das Geschäft des Hauses\" > \"reference.de\"\n",
+        "\n",
+        "!t2t-decoder \\\n",
+        "  --data_dir=$DATA_DIR_PRETRAINED \\\n",
+        "  --problem=$PROBLEM_PRETRAINED  \\\n",
+        "  --model=$MODEL_PRETRAINED  \\\n",
+        "  --hparams_set=$HPARAMS_PRETRAINED \\\n",
+        "  --output_dir=$TRAIN_DIR_PRETRAINED  \\\n",
+        "  --decode_hparams=\"beam_size=1\" \\\n",
+        "  --decode_from_file=\"inputs.en\" \\\n",
+        "  --decode_to_file=\"outputs.de\"\n",
+        "\n",
+        "# See the translations\n",
+        "!cat outputs.de\n",
+        "\n",
+        "!t2t-bleu \\\n",
+        "    --translation=outputs.de \\\n",
+        "    --reference=reference.de"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bKI4WF0DgoFd",
+        "colab_type": "text"
+      },
+      "source": [
+        "#9.  Add your dataset/problem\n",
+        "To add a new dataset/problem, subclass Problem and register it with @registry.register_problem. See TranslateEnfrWmt8k for an example: \n",
+        "https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/translate_enfr.py\n",
+        "\n",
+        "---\n",
+        "Adding your own components: https://github.com/tensorflow/tensor2tensor#adding-your-own-components\n",
+        "\n",
+        "---\n",
+        "\n",
+        "See this example: https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/test_data/example_usr_dir"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "mB1SIrJNqy1N",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "from tensor2tensor.utils import registry\n",
+        "\n",
+        "@registry.register_problem\n",
+        "class MyTranslateEnFr(translate_enfr.TranslateEnfrWmt8k):\n",
+        "\n",
+        "  def generator(self, data_dir, tmp_dir, train):\n",
+        "   #your code"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    }
+  ]
+}
\ No newline at end of file

From 54a9ef75bb9500e0ee3a4e4a48b246dec2c7fb8e Mon Sep 17 00:00:00 2001
From: Tanguy Urvoy <tanguy.urvoy@orange.com>
Date: Mon, 26 Aug 2019 13:20:37 -0700
Subject: [PATCH 2320/2720] Merge of PR #1674

PiperOrigin-RevId: 265525791
---
 README.md                                     |    2 -
 .../data_generators/generator_utils.py        |    5 +-
 .../notebooks/Transformer_translate.ipynb     | 1102 -----------------
 3 files changed, 3 insertions(+), 1106 deletions(-)
 delete mode 100644 tensor2tensor/notebooks/Transformer_translate.ipynb

diff --git a/README.md b/README.md
index f0ee8bb61..d51eff0a9 100644
--- a/README.md
+++ b/README.md
@@ -227,8 +227,6 @@ which is close to state-of-the art. If training on a single GPU, try the
 or larger data-sets (e.g., for English-French), try the big model
 with `--hparams_set=transformer_big`.
 
-See this [example](https://github.com/Styleoshin/tensor2tensor/blob/Transformer_tutorial/tensor2tensor/notebooks/Transformer_translate.ipynb) to know how the translation works.
-
 ## Basics
 
 ### Walkthrough
diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 4773db565..92e330743 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -28,6 +28,7 @@
 import stat
 import tarfile
 import tempfile
+import numpy as np
 import requests
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
@@ -38,7 +39,6 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
-import numpy as np
 
 UNSHUFFLED_SUFFIX = "-unshuffled"
 
@@ -49,7 +49,8 @@ def to_example(dictionary):
   for (k, v) in six.iteritems(dictionary):
     if not v:
       raise ValueError("Empty generated field: %s" % str((k, v)))
-    if isinstance(v[0], six.integer_types) or isinstance(v[0], np.int64):
+    if (isinstance(v[0], six.integer_types) or
+        np.issubdtype(type(v[0]), np.integer)):
       features[k] = tf.train.Feature(int64_list=tf.train.Int64List(value=v))
     elif isinstance(v[0], float):
       features[k] = tf.train.Feature(float_list=tf.train.FloatList(value=v))
diff --git a/tensor2tensor/notebooks/Transformer_translate.ipynb b/tensor2tensor/notebooks/Transformer_translate.ipynb
deleted file mode 100644
index 07c350351..000000000
--- a/tensor2tensor/notebooks/Transformer_translate.ipynb
+++ /dev/null
@@ -1,1102 +0,0 @@
-{
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "name": "Transformer_translate.ipynb",
-      "version": "0.3.2",
-      "provenance": [],
-      "collapsed_sections": [],
-      "toc_visible": true,
-      "include_colab_link": true
-    },
-    "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
-    },
-    "accelerator": "GPU"
-  },
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "view-in-github",
-        "colab_type": "text"
-      },
-      "source": [
-        "<a href=\"https://colab.research.google.com/github/Styleoshin/tensor2tensor/blob/Transformer_tutorial/tensor2tensor/notebooks/Transformer_translate.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "e7PMze9tKHX9",
-        "colab_type": "text"
-      },
-      "source": [
-        "# Welcome to the [Tensor2Tensor](https://github.com/tensorflow/tensor2tensor) Colab\n",
-        "\n",
-        "Tensor2Tensor, or T2T for short, is a library of deep learning models and datasets designed to make deep learning more accessible and [accelerate ML research](https://research.googleblog.com/2017/06/accelerating-deep-learning-research.html). In this notebook we will see how to use this library for a translation task by exploring the necessary steps. We will see how to define a problem, generate the data, train the model and test the quality of it, and we will translate our sequences and we visualize the attention. We will also see how to download a pre-trained model."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "KC8jNpnyKJdm",
-        "colab_type": "code",
-        "cellView": "form",
-        "colab": {}
-      },
-      "source": [
-        "#@title\n",
-        "# Copyright 2018 Google LLC.\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",
-        "# https://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."
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "AYUy570fKRcw",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "# Install deps\n",
-        "!pip install -q -U tensor2tensor"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "hEhFfyVNbB_D",
-        "colab_type": "text"
-      },
-      "source": [
-        "#1. Initialization\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "i23pCAVwegx3",
-        "colab_type": "text"
-      },
-      "source": [
-        "##1.1. Make some directories"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "oUf4e18_8E31",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "import tensorflow as tf\n",
-        "import os\n",
-        "\n",
-        "DATA_DIR = os.path.expanduser(\"/t2t/data\") # This folder contain the data\n",
-        "TMP_DIR = os.path.expanduser(\"/t2t/tmp\")\n",
-        "TRAIN_DIR = os.path.expanduser(\"/t2t/train\") # This folder contain the model\n",
-        "EXPORT_DIR = os.path.expanduser(\"/t2t/export\") # This folder contain the exported model for production\n",
-        "TRANSLATIONS_DIR = os.path.expanduser(\"/t2t/translation\") # This folder contain  all translated sequence\n",
-        "EVENT_DIR = os.path.expanduser(\"/t2t/event\") # Test the BLEU score\n",
-        "USR_DIR = os.path.expanduser(\"/t2t/user\") # This folder contains our data that we want to add\n",
-        " \n",
-        "tf.gfile.MakeDirs(DATA_DIR)\n",
-        "tf.gfile.MakeDirs(TMP_DIR)\n",
-        "tf.gfile.MakeDirs(TRAIN_DIR)\n",
-        "tf.gfile.MakeDirs(EXPORT_DIR)\n",
-        "tf.gfile.MakeDirs(TRANSLATIONS_DIR)\n",
-        "tf.gfile.MakeDirs(EVENT_DIR)\n",
-        "tf.gfile.MakeDirs(USR_DIR)"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "HIuzsMzgbLv9",
-        "colab_type": "text"
-      },
-      "source": [
-        "## 1.2. Init parameters\n",
-        "\n",
-        "\n",
-        "\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "ZQaURmfKBGus",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "PROBLEM = \"translate_enfr_wmt32k\" # We chose a problem translation English to French with 32.768 vocabulary\n",
-        "MODEL = \"transformer\" # Our model\n",
-        "HPARAMS = \"transformer_big\" # Hyperparameters for the model by default \n",
-        "                            # If you have a one gpu, use transformer_big_single_gpu"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "EikK-hW5m-ax",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "#Show all problems and models \n",
-        "\n",
-        "from tensor2tensor.utils import registry\n",
-        "from tensor2tensor import problems\n",
-        "\n",
-        "problems.available() #Show all problems\n",
-        "registry.list_models() #Show all registered models\n",
-        "\n",
-        "#or\n",
-        "\n",
-        "#Command line\n",
-        "!t2t-trainer --registry_help #Show all problems\n",
-        "!t2t-trainer --problems_help #Show all models"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "78kBAIMQbeO6",
-        "colab_type": "text"
-      },
-      "source": [
-        "# 2. Data generation \n",
-        "\n",
-        "Generate the data (download the dataset and generate the data).\n",
-        "\n",
-        "---\n",
-        "\n",
-        " You can choose between command line or code."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "CrDy3V7ibpQH",
-        "colab_type": "text"
-      },
-      "source": [
-        "## 2.1. Generate with terminal\n",
-        "For more information: https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_datagen.py"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "0Dfr8nFXmg1o",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "!t2t-datagen \\\n",
-        "  --data_dir=$DATA_DIR \\\n",
-        "  --tmp_dir=$TMP_DIR \\\n",
-        "  --problem=$PROBLEM \\\n",
-        "  --t2t_usr_dir=$USR_DIR"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "tMvCiiBtbuzh",
-        "colab_type": "text"
-      },
-      "source": [
-        "## 2.2. Generate with code"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "Of5bHYVJmbwH",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "t2t_problem = problems.problem(PROBLEM)\n",
-        "t2t_problem.generate_data(DATA_DIR, TMP_DIR) "
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "UkSwoqBzb47T",
-        "colab_type": "text"
-      },
-      "source": [
-        "# 3. Train the model\n",
-        "\n",
-        "\n",
-        "\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "1JVF2PJn7ByQ",
-        "colab_type": "text"
-      },
-      "source": [
-        "##3.1. Init parameters\n",
-        "\n",
-        "You can choose between command line or code.\n",
-        "\n",
-        "---\n",
-        "\n",
-        " batch_size :  a great value of preference.\n",
-        "\n",
-        "---\n",
-        "train_steps : research paper mentioned 300k steps with 8 gpu on big transformer. So if you have 1 gpu, you will need to train the model x8 more. (https://arxiv.org/abs/1706.03762 for more information).\n",
-        "\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "yw6HgVWA7AQF",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "train_steps = 300000 # Total number of train steps for all Epochs\n",
-        "eval_steps = 100 # Number of steps to perform for each evaluation\n",
-        "batch_size = 4096\n",
-        "save_checkpoints_steps = 1000\n",
-        "ALPHA = 0.1\n",
-        "schedule = \"continuous_train_and_eval\""
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ze_YvVnIfD8z",
-        "colab_type": "text"
-      },
-      "source": [
-        "You can choose schedule :\n",
-        " \n",
-        "\n",
-        "*  train. Bad quality\n",
-        "*  continuous_train_and_eval (default)\n",
-        "*   train_and_eval\n",
-        "\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "-zAub7Ggb8tj",
-        "colab_type": "text"
-      },
-      "source": [
-        "##3.2. Train with terminal\n",
-        "https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_trainer.py\n",
-        "\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "kSYAi4BsnpSD",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "!t2t-trainer \\\n",
-        "  --data_dir=$DATA_DIR \\\n",
-        "  --problem=$PROBLEM \\\n",
-        "  --model=$MODEL \\\n",
-        "  --hparams_set=$HPARAMS \\\n",
-        "  --hparams=\"batch_size=$batch_size\" \\\n",
-        "  --schedule=$schedule\\\n",
-        "  --output_dir=$TRAIN_DIR \\\n",
-        "  --train_steps=$train_steps \\\n",
-        "  --worker-gpu=1 \\ \n",
-        "  --eval_steps=$eval_steps "
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "bNfNBWtNVMwO",
-        "colab_type": "text"
-      },
-      "source": [
-        "  --worker-gpu = 1, for train on 1 gpu (facultative).\n",
-        "\n",
-        "---\n",
-        "\n",
-        "For distributed training see: https://github.com/tensorflow/tensor2tensor/blob/master/docs/distributed_training.md\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "nnSoC1AUcLG6",
-        "colab_type": "text"
-      },
-      "source": [
-        "##3.3. Train with code\n",
-        "create_hparams : https://github.com/tensorflow/tensor2tensor/blob/28adf2690c551ef0f570d41bef2019d9c502ec7e/tensor2tensor/utils/hparams_lib.py#L42\n",
-        "\n",
-        "---\n",
-        "Change hyper parameters :\n",
-        "https://github.com/tensorflow/tensor2tensor/blob/28adf2690c551ef0f570d41bef2019d9c502ec7e/tensor2tensor/models/transformer.py#L1627\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "RJ91vQ2hyIPx",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "from tensor2tensor.utils.trainer_lib import create_run_config, create_experiment\n",
-        "from tensor2tensor.utils.trainer_lib import create_hparams\n",
-        "from tensor2tensor.utils import registry\n",
-        "from tensor2tensor import models\n",
-        "from tensor2tensor import problems\n",
-        "\n",
-        "# Init Hparams object from T2T Problem\n",
-        "hparams = create_hparams(HPARAMS)\n",
-        "\n",
-        "# Make Changes to Hparams\n",
-        "hparams.batch_size = batch_size\n",
-        "hparams.learning_rate = ALPHA\n",
-        "#hparams.max_length = 256\n",
-        "\n",
-        "# Can see all Hparams with code below\n",
-        "#print(json.loads(hparams.to_json())"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "KZX1cwK3TEXs",
-        "colab_type": "text"
-      },
-      "source": [
-        "create_run_config : https://github.com/tensorflow/tensor2tensor/blob/28adf2690c551ef0f570d41bef2019d9c502ec7e/tensor2tensor/utils/trainer_lib.py#L105\n",
-        "\n",
-        "---\n",
-        "\n",
-        "\n",
-        "create_experiment : https://github.com/tensorflow/tensor2tensor/blob/28adf2690c551ef0f570d41bef2019d9c502ec7e/tensor2tensor/utils/trainer_lib.py#L611"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "yByKcs7XvAXL",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "RUN_CONFIG = create_run_config(\n",
-        "      model_dir=TRAIN_DIR,\n",
-        "      model_name=MODEL,\n",
-        "      save_checkpoints_steps= save_checkpoints_steps\n",
-        ")\n",
-        "\n",
-        "tensorflow_exp_fn = create_experiment(\n",
-        "        run_config=RUN_CONFIG,\n",
-        "        hparams=hparams,\n",
-        "        model_name=MODEL,\n",
-        "        problem_name=PROBLEM,\n",
-        "        data_dir=DATA_DIR, \n",
-        "        train_steps=train_steps, \n",
-        "        eval_steps=eval_steps, \n",
-        "        #use_xla=True # For acceleration\n",
-        "    ) \n",
-        "\n",
-        "tensorflow_exp_fn.train_and_evaluate()"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "03xuR70jce_2",
-        "colab_type": "text"
-      },
-      "source": [
-        "#4. See the BLEU score"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "MiwyVWPhhGrk",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "#INIT FILE FOR TRANSLATE\n",
-        "\n",
-        "SOURCE_TEST_TRANSLATE_DIR = TMP_DIR+\"/dev/newstest2014-fren-src.en.sgm\"\n",
-        "REFERENCE_TEST_TRANSLATE_DIR = TMP_DIR+\"/dev/newstest2014-fren-ref.en.sgm\"\n",
-        "BEAM_SIZE=1"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "agnSg_89cr63",
-        "colab_type": "text"
-      },
-      "source": [
-        "##4.1. Translate all\n",
-        "https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_translate_all.py"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "Jrt5fwqsg3pl",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "!t2t-translate-all \\\n",
-        "  --source=$SOURCE_TEST_TRANSLATE_DIR \\\n",
-        "  --model_dir=$TRAIN_DIR \\\n",
-        "  --translations_dir=$TRANSLATIONS_DIR \\\n",
-        "  --data_dir=$DATA_DIR \\\n",
-        "  --problem=$PROBLEM \\\n",
-        "  --hparams_set=$HPARAMS \\\n",
-        "  --output_dir=$TRAIN_DIR \\\n",
-        "  --t2t_usr_dir=$USR_DIR \\\n",
-        "  --beam_size=$BEAM_SIZE \\\n",
-        "  --model=$MODEL"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "O-pKKU2Acv8Q",
-        "colab_type": "text"
-      },
-      "source": [
-        "##4.2. Test the BLEU score\n",
-        "The BLEU score for all translations: https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_bleu.py#L68\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "EULP9TdPc58d",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "!t2t-bleu \\\n",
-        "   --translations_dir=$TRANSLATIONS_DIR \\\n",
-        "   --model_dir=$TRAIN_DIR \\\n",
-        "   --data_dir=$DATA_DIR \\\n",
-        "   --problem=$PROBLEM \\\n",
-        "   --hparams_set=$HPARAMS \\\n",
-        "   --source=$SOURCE_TEST_TRANSLATE_DIR \\\n",
-        "   --reference=$REFERENCE_TEST_TRANSLATE_DIR \\\n",
-        "   --event_dir=$EVENT_DIR"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "13j50bpAc-bM",
-        "colab_type": "text"
-      },
-      "source": [
-        "#5. Prediction of sentence\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "8WHPnqxhdQl6",
-        "colab_type": "text"
-      },
-      "source": [
-        "##5.1. Predict with terminal\n",
-        "https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_decoder.py"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "3SD-XhImnwpo",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "!echo \"the business of the house\" > \"inputs.en\"\n",
-        "!echo -e \"les affaires de la maison\" > \"reference.fr\" # You can add other references\n",
-        "\n",
-        "!t2t-decoder \\\n",
-        "  --data_dir=$DATA_DIR \\\n",
-        "  --problem=$PROBLEM \\\n",
-        "  --model=$MODEL \\\n",
-        "  --hparams_set=$HPARAMS \\\n",
-        "  --output_dir=$TRAIN_DIR \\\n",
-        "  --decode_hparams=\"beam_size=1,alpha=$ALPHA\" \\\n",
-        "  --decode_from_file=\"inputs.en\" \\\n",
-        "  --decode_to_file=\"outputs.fr\"\n",
-        "\n",
-        "# See the translations\n",
-        "!cat outputs.fr"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sGOC25N4dWdM",
-        "colab_type": "text"
-      },
-      "source": [
-        "##5.2. Predict with code"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "S6u4QmhPIbDx",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "import tensorflow as tf\n",
-        "\n",
-        "#After training the model, re-run the environment but run this code in first, then predict.\n",
-        "\n",
-        "tfe = tf.contrib.eager\n",
-        "tfe.enable_eager_execution()\n",
-        "Modes = tf.estimator.ModeKeys"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "PaCkILfjz9x3",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "#Config\n",
-        "\n",
-        "from tensor2tensor import models\n",
-        "from tensor2tensor import problems\n",
-        "from tensor2tensor.layers import common_layers\n",
-        "from tensor2tensor.utils import trainer_lib\n",
-        "from tensor2tensor.utils import t2t_model\n",
-        "from tensor2tensor.utils import registry\n",
-        "from tensor2tensor.utils import metrics\n",
-        "import numpy as np\n",
-        "\n",
-        "enfr_problem = problems.problem(PROBLEM)\n",
-        "\n",
-        "# Copy the vocab file locally so we can encode inputs and decode model outputs\n",
-        "vocab_name = \"vocab.translate_enfr_wmt32k.32768.subwords\"\n",
-        "vocab_file = os.path.join(DATA_DIR, vocab_name)\n",
-        "\n",
-        "# Get the encoders from the problem\n",
-        "encoders = enfr_problem.feature_encoders(DATA_DIR)\n",
-        "\n",
-        "ckpt_path = tf.train.latest_checkpoint(os.path.join(TRAIN_DIR))\n",
-        "print(ckpt_path)\n",
-        "\n",
-        "def translate(inputs):\n",
-        "  encoded_inputs = encode(inputs)\n",
-        "  with tfe.restore_variables_on_create(ckpt_path):\n",
-        "    model_output = translate_model.infer(encoded_inputs)[\"outputs\"]\n",
-        "  return decode(model_output)\n",
-        "\n",
-        "def encode(input_str, output_str=None):\n",
-        "  \"\"\"Input str to features dict, ready for inference\"\"\"\n",
-        "  inputs = encoders[\"inputs\"].encode(input_str) + [1]  # add EOS id\n",
-        "  batch_inputs = tf.reshape(inputs, [1, -1, 1])  # Make it 3D.\n",
-        "  return {\"inputs\": batch_inputs}\n",
-        "\n",
-        "def decode(integers):\n",
-        "  \"\"\"List of ints to str\"\"\"\n",
-        "  integers = list(np.squeeze(integers))\n",
-        "  if 1 in integers:\n",
-        "    integers = integers[:integers.index(1)]\n",
-        "  return encoders[\"inputs\"].decode(np.squeeze(integers))"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "5zE8yHLUA2He",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "#Predict \n",
-        "\n",
-        "hparams = trainer_lib.create_hparams(HPARAMS, data_dir=DATA_DIR, problem_name=PROBLEM)\n",
-        "translate_model = registry.model(MODEL)(hparams, Modes.PREDICT)\n",
-        "\n",
-        "inputs = \"the aniamal didn't cross the river because it was too tired\"\n",
-        "ref = \"l'animal n'a pas traversé la rue parcequ'il etait trop fatigué\" ## this just a reference for evaluate the quality of the traduction\n",
-        "outputs = translate(inputs)\n",
-        "\n",
-        "print(\"Inputs: %s\" % inputs)\n",
-        "print(\"Outputs: %s\" % outputs)\n",
-        "\n",
-        "file_input = open(\"outputs.fr\",\"w+\")\n",
-        "file_input.write(outputs)\n",
-        "file_input.close()\n",
-        "\n",
-        "file_output = open(\"reference.fr\",\"w+\")\n",
-        "file_output.write(ref)\n",
-        "file_output.close()"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "y6jbQ6FoRsmG",
-        "colab_type": "text"
-      },
-      "source": [
-        "##5.3. Evaluate the BLEU Score\n",
-        "BLEU score for a sequence translation: https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_bleu.py#L24"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "il2oevmXRrbf",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "!t2t-bleu \\\n",
-        "    --translation=outputs.fr \\\n",
-        "    --reference=reference.fr"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "FXegHzD1I67e",
-        "colab_type": "text"
-      },
-      "source": [
-        "#6. Attention visualization\n",
-        "We need to have a predicted sentence with code."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ISHauPT8I-3S",
-        "colab_type": "text"
-      },
-      "source": [
-        "##6.1. Attention utils\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "2RHCTrc9I55K",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "from tensor2tensor.visualization import attention\n",
-        "from tensor2tensor.data_generators import text_encoder\n",
-        "\n",
-        "SIZE = 35\n",
-        "\n",
-        "def encode_eval(input_str, output_str):\n",
-        "  inputs = tf.reshape(encoders[\"inputs\"].encode(input_str) + [1], [1, -1, 1, 1])  # Make it 3D.\n",
-        "  outputs = tf.reshape(encoders[\"inputs\"].encode(output_str) + [1], [1, -1, 1, 1])  # Make it 3D.\n",
-        "  return {\"inputs\": inputs, \"targets\": outputs}\n",
-        "\n",
-        "def get_att_mats():\n",
-        "  enc_atts = []\n",
-        "  dec_atts = []\n",
-        "  encdec_atts = []\n",
-        "\n",
-        "  for i in range(hparams.num_hidden_layers):\n",
-        "    enc_att = translate_model.attention_weights[\n",
-        "      \"transformer/body/encoder/layer_%i/self_attention/multihead_attention/dot_product_attention\" % i][0]\n",
-        "    dec_att = translate_model.attention_weights[\n",
-        "      \"transformer/body/decoder/layer_%i/self_attention/multihead_attention/dot_product_attention\" % i][0]\n",
-        "    encdec_att = translate_model.attention_weights[\n",
-        "      \"transformer/body/decoder/layer_%i/encdec_attention/multihead_attention/dot_product_attention\" % i][0]\n",
-        "    enc_atts.append(resize(enc_att))\n",
-        "    dec_atts.append(resize(dec_att))\n",
-        "    encdec_atts.append(resize(encdec_att))\n",
-        "  return enc_atts, dec_atts, encdec_atts\n",
-        "\n",
-        "def resize(np_mat):\n",
-        "  # Sum across heads\n",
-        "  np_mat = np_mat[:, :SIZE, :SIZE]\n",
-        "  row_sums = np.sum(np_mat, axis=0)\n",
-        "  # Normalize\n",
-        "  layer_mat = np_mat / row_sums[np.newaxis, :]\n",
-        "  lsh = layer_mat.shape\n",
-        "  # Add extra dim for viz code to work.\n",
-        "  layer_mat = np.reshape(layer_mat, (1, lsh[0], lsh[1], lsh[2]))\n",
-        "  return layer_mat\n",
-        "\n",
-        "def to_tokens(ids):\n",
-        "  ids = np.squeeze(ids)\n",
-        "  subtokenizer = hparams.problem_hparams.vocabulary['targets']\n",
-        "  tokens = []\n",
-        "  for _id in ids:\n",
-        "    if _id == 0:\n",
-        "      tokens.append('<PAD>')\n",
-        "    elif _id == 1:\n",
-        "      tokens.append('<EOS>')\n",
-        "    elif _id == -1:\n",
-        "      tokens.append('<NULL>')\n",
-        "    else:\n",
-        "        tokens.append(subtokenizer._subtoken_id_to_subtoken_string(_id))\n",
-        "  return tokens\n",
-        "\n",
-        "def call_html():\n",
-        "  import IPython\n",
-        "  display(IPython.core.display.HTML('''\n",
-        "        <script src=\"/static/components/requirejs/require.js\"></script>\n",
-        "        <script>\n",
-        "          requirejs.config({\n",
-        "            paths: {\n",
-        "              base: '/static/base',\n",
-        "              \"d3\": \"https://cdnjs.cloudflare.com/ajax/libs/d3/3.5.8/d3.min\",\n",
-        "              jquery: '//ajax.googleapis.com/ajax/libs/jquery/2.0.0/jquery.min',\n",
-        "            },\n",
-        "          });\n",
-        "        </script>\n",
-        "        '''))"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "9PGwUbJuJHJS",
-        "colab_type": "text"
-      },
-      "source": [
-        "##6.2 Display Attention"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "ijTOlrt8JI4t",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "import numpy as np\n",
-        "\n",
-        "# Convert inputs and outputs to subwords\n",
-        "\n",
-        "inp_text = to_tokens(encoders[\"inputs\"].encode(inputs))\n",
-        "out_text = to_tokens(encoders[\"inputs\"].encode(outputs))\n",
-        "\n",
-        "hparams = trainer_lib.create_hparams(HPARAMS, data_dir=DATA_DIR, problem_name=PROBLEM)\n",
-        "\n",
-        "# Run eval to collect attention weights\n",
-        "example = encode_eval(inputs, outputs)\n",
-        "with tfe.restore_variables_on_create(tf.train.latest_checkpoint(ckpt_path)):\n",
-        "  translate_model.set_mode(Modes.EVAL)\n",
-        "  translate_model(example)\n",
-        "# Get normalized attention weights for each layer\n",
-        "enc_atts, dec_atts, encdec_atts = get_att_mats()\n",
-        "\n",
-        "call_html()\n",
-        "attention.show(inp_text, out_text, enc_atts, dec_atts, encdec_atts)"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "r8yAQUDZdm1p",
-        "colab_type": "text"
-      },
-      "source": [
-        "#7. Export the model\n",
-        "For more information: https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/serving"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "c2yulC7J8_I9",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "#export Model\n",
-        "!t2t-exporter \\\n",
-        "  --data_dir=$DATA_DIR \\\n",
-        "  --output_dir=$TRAIN_DIR \\\n",
-        "  --problem=$PROBLEM \\\n",
-        "  --model=$MODEL \\\n",
-        "  --hparams_set=$HPARAMS \\\n",
-        "  --decode_hparams=\"beam_size=1,alpha=$ALPHA\" \\\n",
-        "  --export_dir=$EXPORT_DIR"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "2ltjEr3JX5-e",
-        "colab_type": "text"
-      },
-      "source": [
-        "#8.Load pretrained model from Google Storage\n",
-        "We use the pretrained model En-De translation."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "QgY3Fw261bZC",
-        "colab_type": "text"
-      },
-      "source": [
-        "##8.1. See existing content storaged"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "7P7aJClG0t8c",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "print(\"checkpoint: \")\n",
-        "!gsutil ls \"gs://tensor2tensor-checkpoints\"\n",
-        "\n",
-        "print(\"data: \")\n",
-        "!gsutil ls \"gs://tensor2tensor-data\""
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "wP8jrR5bbu7e",
-        "colab_type": "text"
-      },
-      "source": [
-        "##8.2. Init model"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "AnYU7lrazkMm",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "PROBLEM_PRETRAINED = \"translate_ende_wmt32k\"\n",
-        "MODEL_PRETRAINED = \"transformer\" \n",
-        "HPARAMS_PRETRAINED = \"transformer_base\""
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "DTgPvq4q1VAr",
-        "colab_type": "text"
-      },
-      "source": [
-        "##8.3. Load content from google storage"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "FrxOAVcyinll",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "import tensorflow as tf\n",
-        "import os\n",
-        "\n",
-        "\n",
-        "DATA_DIR_PRETRAINED = os.path.expanduser(\"/t2t/data_pretrained\")\n",
-        "CHECKPOINT_DIR_PRETRAINED = os.path.expanduser(\"/t2t/checkpoints_pretrained\")\n",
-        "\n",
-        "tf.gfile.MakeDirs(DATA_DIR_PRETRAINED)\n",
-        "tf.gfile.MakeDirs(CHECKPOINT_DIR_PRETRAINED)\n",
-        "\n",
-        "\n",
-        "gs_data_dir = \"gs://tensor2tensor-data/\"\n",
-        "vocab_name = \"vocab.translate_ende_wmt32k.32768.subwords\"\n",
-        "vocab_file = os.path.join(gs_data_dir, vocab_name)\n",
-        "\n",
-        "gs_ckpt_dir = \"gs://tensor2tensor-checkpoints/\"\n",
-        "ckpt_name = \"transformer_ende_test\"\n",
-        "gs_ckpt = os.path.join(gs_ckpt_dir, ckpt_name)\n",
-        "\n",
-        "TRAIN_DIR_PRETRAINED = os.path.join(CHECKPOINT_DIR_PRETRAINED, ckpt_name)\n",
-        "\n",
-        "!gsutil cp {vocab_file} {DATA_DIR_PRETRAINED}\n",
-        "!gsutil -q cp -R {gs_ckpt} {CHECKPOINT_DIR_PRETRAINED}\n",
-        "\n",
-        "CHECKPOINT_NAME_PRETRAINED = tf.train.latest_checkpoint(TRAIN_DIR_PRETRAINED) # for translate with code\n"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "LP6cro9Xbygf",
-        "colab_type": "text"
-      },
-      "source": [
-        "##8.4. Translate"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "CBoNpy5HbzoF",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "!echo \"the business of the house\" > \"inputs.en\"\n",
-        "!echo -e \"das Geschäft des Hauses\" > \"reference.de\"\n",
-        "\n",
-        "!t2t-decoder \\\n",
-        "  --data_dir=$DATA_DIR_PRETRAINED \\\n",
-        "  --problem=$PROBLEM_PRETRAINED  \\\n",
-        "  --model=$MODEL_PRETRAINED  \\\n",
-        "  --hparams_set=$HPARAMS_PRETRAINED \\\n",
-        "  --output_dir=$TRAIN_DIR_PRETRAINED  \\\n",
-        "  --decode_hparams=\"beam_size=1\" \\\n",
-        "  --decode_from_file=\"inputs.en\" \\\n",
-        "  --decode_to_file=\"outputs.de\"\n",
-        "\n",
-        "# See the translations\n",
-        "!cat outputs.de\n",
-        "\n",
-        "!t2t-bleu \\\n",
-        "    --translation=outputs.de \\\n",
-        "    --reference=reference.de"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "bKI4WF0DgoFd",
-        "colab_type": "text"
-      },
-      "source": [
-        "#9.  Add your dataset/problem\n",
-        "To add a new dataset/problem, subclass Problem and register it with @registry.register_problem. See TranslateEnfrWmt8k for an example: \n",
-        "https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/translate_enfr.py\n",
-        "\n",
-        "---\n",
-        "Adding your own components: https://github.com/tensorflow/tensor2tensor#adding-your-own-components\n",
-        "\n",
-        "---\n",
-        "\n",
-        "See this example: https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/test_data/example_usr_dir"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "mB1SIrJNqy1N",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "from tensor2tensor.utils import registry\n",
-        "\n",
-        "@registry.register_problem\n",
-        "class MyTranslateEnFr(translate_enfr.TranslateEnfrWmt8k):\n",
-        "\n",
-        "  def generator(self, data_dir, tmp_dir, train):\n",
-        "   #your code"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    }
-  ]
-}
\ No newline at end of file

From 15a148d42bf63c58ae86cd6423feede39638a34e Mon Sep 17 00:00:00 2001
From: Yacine BENAFFANE <yacine.benaffane@gmail.com>
Date: Mon, 26 Aug 2019 13:22:08 -0700
Subject: [PATCH 2321/2720] Merge of PR #1675

PiperOrigin-RevId: 265526105
---
 README.md                                     |    2 +
 docs/walkthrough.md                           |    2 +
 .../notebooks/Transformer_translate.ipynb     | 1102 +++++++++++++++++
 3 files changed, 1106 insertions(+)
 create mode 100644 tensor2tensor/notebooks/Transformer_translate.ipynb

diff --git a/README.md b/README.md
index d51eff0a9..66ece0c34 100644
--- a/README.md
+++ b/README.md
@@ -227,6 +227,8 @@ which is close to state-of-the art. If training on a single GPU, try the
 or larger data-sets (e.g., for English-French), try the big model
 with `--hparams_set=transformer_big`.
 
+See this [example](https://github.com/tensorflow/tensor2tensor/blob/Transformer_tutorial/tensor2tensor/notebooks/Transformer_translate.ipynb) to know how the translation works.
+
 ## Basics
 
 ### Walkthrough
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index d51eff0a9..66ece0c34 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -227,6 +227,8 @@ which is close to state-of-the art. If training on a single GPU, try the
 or larger data-sets (e.g., for English-French), try the big model
 with `--hparams_set=transformer_big`.
 
+See this [example](https://github.com/tensorflow/tensor2tensor/blob/Transformer_tutorial/tensor2tensor/notebooks/Transformer_translate.ipynb) to know how the translation works.
+
 ## Basics
 
 ### Walkthrough
diff --git a/tensor2tensor/notebooks/Transformer_translate.ipynb b/tensor2tensor/notebooks/Transformer_translate.ipynb
new file mode 100644
index 000000000..07c350351
--- /dev/null
+++ b/tensor2tensor/notebooks/Transformer_translate.ipynb
@@ -0,0 +1,1102 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "name": "Transformer_translate.ipynb",
+      "version": "0.3.2",
+      "provenance": [],
+      "collapsed_sections": [],
+      "toc_visible": true,
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/Styleoshin/tensor2tensor/blob/Transformer_tutorial/tensor2tensor/notebooks/Transformer_translate.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e7PMze9tKHX9",
+        "colab_type": "text"
+      },
+      "source": [
+        "# Welcome to the [Tensor2Tensor](https://github.com/tensorflow/tensor2tensor) Colab\n",
+        "\n",
+        "Tensor2Tensor, or T2T for short, is a library of deep learning models and datasets designed to make deep learning more accessible and [accelerate ML research](https://research.googleblog.com/2017/06/accelerating-deep-learning-research.html). In this notebook we will see how to use this library for a translation task by exploring the necessary steps. We will see how to define a problem, generate the data, train the model and test the quality of it, and we will translate our sequences and we visualize the attention. We will also see how to download a pre-trained model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "KC8jNpnyKJdm",
+        "colab_type": "code",
+        "cellView": "form",
+        "colab": {}
+      },
+      "source": [
+        "#@title\n",
+        "# Copyright 2018 Google LLC.\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",
+        "# https://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."
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "AYUy570fKRcw",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "# Install deps\n",
+        "!pip install -q -U tensor2tensor"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "hEhFfyVNbB_D",
+        "colab_type": "text"
+      },
+      "source": [
+        "#1. Initialization\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "i23pCAVwegx3",
+        "colab_type": "text"
+      },
+      "source": [
+        "##1.1. Make some directories"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "oUf4e18_8E31",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "import tensorflow as tf\n",
+        "import os\n",
+        "\n",
+        "DATA_DIR = os.path.expanduser(\"/t2t/data\") # This folder contain the data\n",
+        "TMP_DIR = os.path.expanduser(\"/t2t/tmp\")\n",
+        "TRAIN_DIR = os.path.expanduser(\"/t2t/train\") # This folder contain the model\n",
+        "EXPORT_DIR = os.path.expanduser(\"/t2t/export\") # This folder contain the exported model for production\n",
+        "TRANSLATIONS_DIR = os.path.expanduser(\"/t2t/translation\") # This folder contain  all translated sequence\n",
+        "EVENT_DIR = os.path.expanduser(\"/t2t/event\") # Test the BLEU score\n",
+        "USR_DIR = os.path.expanduser(\"/t2t/user\") # This folder contains our data that we want to add\n",
+        " \n",
+        "tf.gfile.MakeDirs(DATA_DIR)\n",
+        "tf.gfile.MakeDirs(TMP_DIR)\n",
+        "tf.gfile.MakeDirs(TRAIN_DIR)\n",
+        "tf.gfile.MakeDirs(EXPORT_DIR)\n",
+        "tf.gfile.MakeDirs(TRANSLATIONS_DIR)\n",
+        "tf.gfile.MakeDirs(EVENT_DIR)\n",
+        "tf.gfile.MakeDirs(USR_DIR)"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HIuzsMzgbLv9",
+        "colab_type": "text"
+      },
+      "source": [
+        "## 1.2. Init parameters\n",
+        "\n",
+        "\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "ZQaURmfKBGus",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "PROBLEM = \"translate_enfr_wmt32k\" # We chose a problem translation English to French with 32.768 vocabulary\n",
+        "MODEL = \"transformer\" # Our model\n",
+        "HPARAMS = \"transformer_big\" # Hyperparameters for the model by default \n",
+        "                            # If you have a one gpu, use transformer_big_single_gpu"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "EikK-hW5m-ax",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "#Show all problems and models \n",
+        "\n",
+        "from tensor2tensor.utils import registry\n",
+        "from tensor2tensor import problems\n",
+        "\n",
+        "problems.available() #Show all problems\n",
+        "registry.list_models() #Show all registered models\n",
+        "\n",
+        "#or\n",
+        "\n",
+        "#Command line\n",
+        "!t2t-trainer --registry_help #Show all problems\n",
+        "!t2t-trainer --problems_help #Show all models"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "78kBAIMQbeO6",
+        "colab_type": "text"
+      },
+      "source": [
+        "# 2. Data generation \n",
+        "\n",
+        "Generate the data (download the dataset and generate the data).\n",
+        "\n",
+        "---\n",
+        "\n",
+        " You can choose between command line or code."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "CrDy3V7ibpQH",
+        "colab_type": "text"
+      },
+      "source": [
+        "## 2.1. Generate with terminal\n",
+        "For more information: https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_datagen.py"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "0Dfr8nFXmg1o",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!t2t-datagen \\\n",
+        "  --data_dir=$DATA_DIR \\\n",
+        "  --tmp_dir=$TMP_DIR \\\n",
+        "  --problem=$PROBLEM \\\n",
+        "  --t2t_usr_dir=$USR_DIR"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tMvCiiBtbuzh",
+        "colab_type": "text"
+      },
+      "source": [
+        "## 2.2. Generate with code"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "Of5bHYVJmbwH",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "t2t_problem = problems.problem(PROBLEM)\n",
+        "t2t_problem.generate_data(DATA_DIR, TMP_DIR) "
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "UkSwoqBzb47T",
+        "colab_type": "text"
+      },
+      "source": [
+        "# 3. Train the model\n",
+        "\n",
+        "\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1JVF2PJn7ByQ",
+        "colab_type": "text"
+      },
+      "source": [
+        "##3.1. Init parameters\n",
+        "\n",
+        "You can choose between command line or code.\n",
+        "\n",
+        "---\n",
+        "\n",
+        " batch_size :  a great value of preference.\n",
+        "\n",
+        "---\n",
+        "train_steps : research paper mentioned 300k steps with 8 gpu on big transformer. So if you have 1 gpu, you will need to train the model x8 more. (https://arxiv.org/abs/1706.03762 for more information).\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "yw6HgVWA7AQF",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "train_steps = 300000 # Total number of train steps for all Epochs\n",
+        "eval_steps = 100 # Number of steps to perform for each evaluation\n",
+        "batch_size = 4096\n",
+        "save_checkpoints_steps = 1000\n",
+        "ALPHA = 0.1\n",
+        "schedule = \"continuous_train_and_eval\""
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ze_YvVnIfD8z",
+        "colab_type": "text"
+      },
+      "source": [
+        "You can choose schedule :\n",
+        " \n",
+        "\n",
+        "*  train. Bad quality\n",
+        "*  continuous_train_and_eval (default)\n",
+        "*   train_and_eval\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-zAub7Ggb8tj",
+        "colab_type": "text"
+      },
+      "source": [
+        "##3.2. Train with terminal\n",
+        "https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_trainer.py\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "kSYAi4BsnpSD",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!t2t-trainer \\\n",
+        "  --data_dir=$DATA_DIR \\\n",
+        "  --problem=$PROBLEM \\\n",
+        "  --model=$MODEL \\\n",
+        "  --hparams_set=$HPARAMS \\\n",
+        "  --hparams=\"batch_size=$batch_size\" \\\n",
+        "  --schedule=$schedule\\\n",
+        "  --output_dir=$TRAIN_DIR \\\n",
+        "  --train_steps=$train_steps \\\n",
+        "  --worker-gpu=1 \\ \n",
+        "  --eval_steps=$eval_steps "
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bNfNBWtNVMwO",
+        "colab_type": "text"
+      },
+      "source": [
+        "  --worker-gpu = 1, for train on 1 gpu (facultative).\n",
+        "\n",
+        "---\n",
+        "\n",
+        "For distributed training see: https://github.com/tensorflow/tensor2tensor/blob/master/docs/distributed_training.md\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "nnSoC1AUcLG6",
+        "colab_type": "text"
+      },
+      "source": [
+        "##3.3. Train with code\n",
+        "create_hparams : https://github.com/tensorflow/tensor2tensor/blob/28adf2690c551ef0f570d41bef2019d9c502ec7e/tensor2tensor/utils/hparams_lib.py#L42\n",
+        "\n",
+        "---\n",
+        "Change hyper parameters :\n",
+        "https://github.com/tensorflow/tensor2tensor/blob/28adf2690c551ef0f570d41bef2019d9c502ec7e/tensor2tensor/models/transformer.py#L1627\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "RJ91vQ2hyIPx",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "from tensor2tensor.utils.trainer_lib import create_run_config, create_experiment\n",
+        "from tensor2tensor.utils.trainer_lib import create_hparams\n",
+        "from tensor2tensor.utils import registry\n",
+        "from tensor2tensor import models\n",
+        "from tensor2tensor import problems\n",
+        "\n",
+        "# Init Hparams object from T2T Problem\n",
+        "hparams = create_hparams(HPARAMS)\n",
+        "\n",
+        "# Make Changes to Hparams\n",
+        "hparams.batch_size = batch_size\n",
+        "hparams.learning_rate = ALPHA\n",
+        "#hparams.max_length = 256\n",
+        "\n",
+        "# Can see all Hparams with code below\n",
+        "#print(json.loads(hparams.to_json())"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KZX1cwK3TEXs",
+        "colab_type": "text"
+      },
+      "source": [
+        "create_run_config : https://github.com/tensorflow/tensor2tensor/blob/28adf2690c551ef0f570d41bef2019d9c502ec7e/tensor2tensor/utils/trainer_lib.py#L105\n",
+        "\n",
+        "---\n",
+        "\n",
+        "\n",
+        "create_experiment : https://github.com/tensorflow/tensor2tensor/blob/28adf2690c551ef0f570d41bef2019d9c502ec7e/tensor2tensor/utils/trainer_lib.py#L611"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "yByKcs7XvAXL",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "RUN_CONFIG = create_run_config(\n",
+        "      model_dir=TRAIN_DIR,\n",
+        "      model_name=MODEL,\n",
+        "      save_checkpoints_steps= save_checkpoints_steps\n",
+        ")\n",
+        "\n",
+        "tensorflow_exp_fn = create_experiment(\n",
+        "        run_config=RUN_CONFIG,\n",
+        "        hparams=hparams,\n",
+        "        model_name=MODEL,\n",
+        "        problem_name=PROBLEM,\n",
+        "        data_dir=DATA_DIR, \n",
+        "        train_steps=train_steps, \n",
+        "        eval_steps=eval_steps, \n",
+        "        #use_xla=True # For acceleration\n",
+        "    ) \n",
+        "\n",
+        "tensorflow_exp_fn.train_and_evaluate()"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "03xuR70jce_2",
+        "colab_type": "text"
+      },
+      "source": [
+        "#4. See the BLEU score"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "MiwyVWPhhGrk",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "#INIT FILE FOR TRANSLATE\n",
+        "\n",
+        "SOURCE_TEST_TRANSLATE_DIR = TMP_DIR+\"/dev/newstest2014-fren-src.en.sgm\"\n",
+        "REFERENCE_TEST_TRANSLATE_DIR = TMP_DIR+\"/dev/newstest2014-fren-ref.en.sgm\"\n",
+        "BEAM_SIZE=1"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "agnSg_89cr63",
+        "colab_type": "text"
+      },
+      "source": [
+        "##4.1. Translate all\n",
+        "https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_translate_all.py"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "Jrt5fwqsg3pl",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!t2t-translate-all \\\n",
+        "  --source=$SOURCE_TEST_TRANSLATE_DIR \\\n",
+        "  --model_dir=$TRAIN_DIR \\\n",
+        "  --translations_dir=$TRANSLATIONS_DIR \\\n",
+        "  --data_dir=$DATA_DIR \\\n",
+        "  --problem=$PROBLEM \\\n",
+        "  --hparams_set=$HPARAMS \\\n",
+        "  --output_dir=$TRAIN_DIR \\\n",
+        "  --t2t_usr_dir=$USR_DIR \\\n",
+        "  --beam_size=$BEAM_SIZE \\\n",
+        "  --model=$MODEL"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "O-pKKU2Acv8Q",
+        "colab_type": "text"
+      },
+      "source": [
+        "##4.2. Test the BLEU score\n",
+        "The BLEU score for all translations: https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_bleu.py#L68\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "EULP9TdPc58d",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!t2t-bleu \\\n",
+        "   --translations_dir=$TRANSLATIONS_DIR \\\n",
+        "   --model_dir=$TRAIN_DIR \\\n",
+        "   --data_dir=$DATA_DIR \\\n",
+        "   --problem=$PROBLEM \\\n",
+        "   --hparams_set=$HPARAMS \\\n",
+        "   --source=$SOURCE_TEST_TRANSLATE_DIR \\\n",
+        "   --reference=$REFERENCE_TEST_TRANSLATE_DIR \\\n",
+        "   --event_dir=$EVENT_DIR"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "13j50bpAc-bM",
+        "colab_type": "text"
+      },
+      "source": [
+        "#5. Prediction of sentence\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8WHPnqxhdQl6",
+        "colab_type": "text"
+      },
+      "source": [
+        "##5.1. Predict with terminal\n",
+        "https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_decoder.py"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "3SD-XhImnwpo",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!echo \"the business of the house\" > \"inputs.en\"\n",
+        "!echo -e \"les affaires de la maison\" > \"reference.fr\" # You can add other references\n",
+        "\n",
+        "!t2t-decoder \\\n",
+        "  --data_dir=$DATA_DIR \\\n",
+        "  --problem=$PROBLEM \\\n",
+        "  --model=$MODEL \\\n",
+        "  --hparams_set=$HPARAMS \\\n",
+        "  --output_dir=$TRAIN_DIR \\\n",
+        "  --decode_hparams=\"beam_size=1,alpha=$ALPHA\" \\\n",
+        "  --decode_from_file=\"inputs.en\" \\\n",
+        "  --decode_to_file=\"outputs.fr\"\n",
+        "\n",
+        "# See the translations\n",
+        "!cat outputs.fr"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "sGOC25N4dWdM",
+        "colab_type": "text"
+      },
+      "source": [
+        "##5.2. Predict with code"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "S6u4QmhPIbDx",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "import tensorflow as tf\n",
+        "\n",
+        "#After training the model, re-run the environment but run this code in first, then predict.\n",
+        "\n",
+        "tfe = tf.contrib.eager\n",
+        "tfe.enable_eager_execution()\n",
+        "Modes = tf.estimator.ModeKeys"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "PaCkILfjz9x3",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "#Config\n",
+        "\n",
+        "from tensor2tensor import models\n",
+        "from tensor2tensor import problems\n",
+        "from tensor2tensor.layers import common_layers\n",
+        "from tensor2tensor.utils import trainer_lib\n",
+        "from tensor2tensor.utils import t2t_model\n",
+        "from tensor2tensor.utils import registry\n",
+        "from tensor2tensor.utils import metrics\n",
+        "import numpy as np\n",
+        "\n",
+        "enfr_problem = problems.problem(PROBLEM)\n",
+        "\n",
+        "# Copy the vocab file locally so we can encode inputs and decode model outputs\n",
+        "vocab_name = \"vocab.translate_enfr_wmt32k.32768.subwords\"\n",
+        "vocab_file = os.path.join(DATA_DIR, vocab_name)\n",
+        "\n",
+        "# Get the encoders from the problem\n",
+        "encoders = enfr_problem.feature_encoders(DATA_DIR)\n",
+        "\n",
+        "ckpt_path = tf.train.latest_checkpoint(os.path.join(TRAIN_DIR))\n",
+        "print(ckpt_path)\n",
+        "\n",
+        "def translate(inputs):\n",
+        "  encoded_inputs = encode(inputs)\n",
+        "  with tfe.restore_variables_on_create(ckpt_path):\n",
+        "    model_output = translate_model.infer(encoded_inputs)[\"outputs\"]\n",
+        "  return decode(model_output)\n",
+        "\n",
+        "def encode(input_str, output_str=None):\n",
+        "  \"\"\"Input str to features dict, ready for inference\"\"\"\n",
+        "  inputs = encoders[\"inputs\"].encode(input_str) + [1]  # add EOS id\n",
+        "  batch_inputs = tf.reshape(inputs, [1, -1, 1])  # Make it 3D.\n",
+        "  return {\"inputs\": batch_inputs}\n",
+        "\n",
+        "def decode(integers):\n",
+        "  \"\"\"List of ints to str\"\"\"\n",
+        "  integers = list(np.squeeze(integers))\n",
+        "  if 1 in integers:\n",
+        "    integers = integers[:integers.index(1)]\n",
+        "  return encoders[\"inputs\"].decode(np.squeeze(integers))"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "5zE8yHLUA2He",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "#Predict \n",
+        "\n",
+        "hparams = trainer_lib.create_hparams(HPARAMS, data_dir=DATA_DIR, problem_name=PROBLEM)\n",
+        "translate_model = registry.model(MODEL)(hparams, Modes.PREDICT)\n",
+        "\n",
+        "inputs = \"the aniamal didn't cross the river because it was too tired\"\n",
+        "ref = \"l'animal n'a pas traversé la rue parcequ'il etait trop fatigué\" ## this just a reference for evaluate the quality of the traduction\n",
+        "outputs = translate(inputs)\n",
+        "\n",
+        "print(\"Inputs: %s\" % inputs)\n",
+        "print(\"Outputs: %s\" % outputs)\n",
+        "\n",
+        "file_input = open(\"outputs.fr\",\"w+\")\n",
+        "file_input.write(outputs)\n",
+        "file_input.close()\n",
+        "\n",
+        "file_output = open(\"reference.fr\",\"w+\")\n",
+        "file_output.write(ref)\n",
+        "file_output.close()"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "y6jbQ6FoRsmG",
+        "colab_type": "text"
+      },
+      "source": [
+        "##5.3. Evaluate the BLEU Score\n",
+        "BLEU score for a sequence translation: https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_bleu.py#L24"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "il2oevmXRrbf",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!t2t-bleu \\\n",
+        "    --translation=outputs.fr \\\n",
+        "    --reference=reference.fr"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "FXegHzD1I67e",
+        "colab_type": "text"
+      },
+      "source": [
+        "#6. Attention visualization\n",
+        "We need to have a predicted sentence with code."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ISHauPT8I-3S",
+        "colab_type": "text"
+      },
+      "source": [
+        "##6.1. Attention utils\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "2RHCTrc9I55K",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "from tensor2tensor.visualization import attention\n",
+        "from tensor2tensor.data_generators import text_encoder\n",
+        "\n",
+        "SIZE = 35\n",
+        "\n",
+        "def encode_eval(input_str, output_str):\n",
+        "  inputs = tf.reshape(encoders[\"inputs\"].encode(input_str) + [1], [1, -1, 1, 1])  # Make it 3D.\n",
+        "  outputs = tf.reshape(encoders[\"inputs\"].encode(output_str) + [1], [1, -1, 1, 1])  # Make it 3D.\n",
+        "  return {\"inputs\": inputs, \"targets\": outputs}\n",
+        "\n",
+        "def get_att_mats():\n",
+        "  enc_atts = []\n",
+        "  dec_atts = []\n",
+        "  encdec_atts = []\n",
+        "\n",
+        "  for i in range(hparams.num_hidden_layers):\n",
+        "    enc_att = translate_model.attention_weights[\n",
+        "      \"transformer/body/encoder/layer_%i/self_attention/multihead_attention/dot_product_attention\" % i][0]\n",
+        "    dec_att = translate_model.attention_weights[\n",
+        "      \"transformer/body/decoder/layer_%i/self_attention/multihead_attention/dot_product_attention\" % i][0]\n",
+        "    encdec_att = translate_model.attention_weights[\n",
+        "      \"transformer/body/decoder/layer_%i/encdec_attention/multihead_attention/dot_product_attention\" % i][0]\n",
+        "    enc_atts.append(resize(enc_att))\n",
+        "    dec_atts.append(resize(dec_att))\n",
+        "    encdec_atts.append(resize(encdec_att))\n",
+        "  return enc_atts, dec_atts, encdec_atts\n",
+        "\n",
+        "def resize(np_mat):\n",
+        "  # Sum across heads\n",
+        "  np_mat = np_mat[:, :SIZE, :SIZE]\n",
+        "  row_sums = np.sum(np_mat, axis=0)\n",
+        "  # Normalize\n",
+        "  layer_mat = np_mat / row_sums[np.newaxis, :]\n",
+        "  lsh = layer_mat.shape\n",
+        "  # Add extra dim for viz code to work.\n",
+        "  layer_mat = np.reshape(layer_mat, (1, lsh[0], lsh[1], lsh[2]))\n",
+        "  return layer_mat\n",
+        "\n",
+        "def to_tokens(ids):\n",
+        "  ids = np.squeeze(ids)\n",
+        "  subtokenizer = hparams.problem_hparams.vocabulary['targets']\n",
+        "  tokens = []\n",
+        "  for _id in ids:\n",
+        "    if _id == 0:\n",
+        "      tokens.append('<PAD>')\n",
+        "    elif _id == 1:\n",
+        "      tokens.append('<EOS>')\n",
+        "    elif _id == -1:\n",
+        "      tokens.append('<NULL>')\n",
+        "    else:\n",
+        "        tokens.append(subtokenizer._subtoken_id_to_subtoken_string(_id))\n",
+        "  return tokens\n",
+        "\n",
+        "def call_html():\n",
+        "  import IPython\n",
+        "  display(IPython.core.display.HTML('''\n",
+        "        <script src=\"/static/components/requirejs/require.js\"></script>\n",
+        "        <script>\n",
+        "          requirejs.config({\n",
+        "            paths: {\n",
+        "              base: '/static/base',\n",
+        "              \"d3\": \"https://cdnjs.cloudflare.com/ajax/libs/d3/3.5.8/d3.min\",\n",
+        "              jquery: '//ajax.googleapis.com/ajax/libs/jquery/2.0.0/jquery.min',\n",
+        "            },\n",
+        "          });\n",
+        "        </script>\n",
+        "        '''))"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9PGwUbJuJHJS",
+        "colab_type": "text"
+      },
+      "source": [
+        "##6.2 Display Attention"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "ijTOlrt8JI4t",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "import numpy as np\n",
+        "\n",
+        "# Convert inputs and outputs to subwords\n",
+        "\n",
+        "inp_text = to_tokens(encoders[\"inputs\"].encode(inputs))\n",
+        "out_text = to_tokens(encoders[\"inputs\"].encode(outputs))\n",
+        "\n",
+        "hparams = trainer_lib.create_hparams(HPARAMS, data_dir=DATA_DIR, problem_name=PROBLEM)\n",
+        "\n",
+        "# Run eval to collect attention weights\n",
+        "example = encode_eval(inputs, outputs)\n",
+        "with tfe.restore_variables_on_create(tf.train.latest_checkpoint(ckpt_path)):\n",
+        "  translate_model.set_mode(Modes.EVAL)\n",
+        "  translate_model(example)\n",
+        "# Get normalized attention weights for each layer\n",
+        "enc_atts, dec_atts, encdec_atts = get_att_mats()\n",
+        "\n",
+        "call_html()\n",
+        "attention.show(inp_text, out_text, enc_atts, dec_atts, encdec_atts)"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "r8yAQUDZdm1p",
+        "colab_type": "text"
+      },
+      "source": [
+        "#7. Export the model\n",
+        "For more information: https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/serving"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "c2yulC7J8_I9",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "#export Model\n",
+        "!t2t-exporter \\\n",
+        "  --data_dir=$DATA_DIR \\\n",
+        "  --output_dir=$TRAIN_DIR \\\n",
+        "  --problem=$PROBLEM \\\n",
+        "  --model=$MODEL \\\n",
+        "  --hparams_set=$HPARAMS \\\n",
+        "  --decode_hparams=\"beam_size=1,alpha=$ALPHA\" \\\n",
+        "  --export_dir=$EXPORT_DIR"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2ltjEr3JX5-e",
+        "colab_type": "text"
+      },
+      "source": [
+        "#8.Load pretrained model from Google Storage\n",
+        "We use the pretrained model En-De translation."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "QgY3Fw261bZC",
+        "colab_type": "text"
+      },
+      "source": [
+        "##8.1. See existing content storaged"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "7P7aJClG0t8c",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "print(\"checkpoint: \")\n",
+        "!gsutil ls \"gs://tensor2tensor-checkpoints\"\n",
+        "\n",
+        "print(\"data: \")\n",
+        "!gsutil ls \"gs://tensor2tensor-data\""
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "wP8jrR5bbu7e",
+        "colab_type": "text"
+      },
+      "source": [
+        "##8.2. Init model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "AnYU7lrazkMm",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "PROBLEM_PRETRAINED = \"translate_ende_wmt32k\"\n",
+        "MODEL_PRETRAINED = \"transformer\" \n",
+        "HPARAMS_PRETRAINED = \"transformer_base\""
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DTgPvq4q1VAr",
+        "colab_type": "text"
+      },
+      "source": [
+        "##8.3. Load content from google storage"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "FrxOAVcyinll",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "import tensorflow as tf\n",
+        "import os\n",
+        "\n",
+        "\n",
+        "DATA_DIR_PRETRAINED = os.path.expanduser(\"/t2t/data_pretrained\")\n",
+        "CHECKPOINT_DIR_PRETRAINED = os.path.expanduser(\"/t2t/checkpoints_pretrained\")\n",
+        "\n",
+        "tf.gfile.MakeDirs(DATA_DIR_PRETRAINED)\n",
+        "tf.gfile.MakeDirs(CHECKPOINT_DIR_PRETRAINED)\n",
+        "\n",
+        "\n",
+        "gs_data_dir = \"gs://tensor2tensor-data/\"\n",
+        "vocab_name = \"vocab.translate_ende_wmt32k.32768.subwords\"\n",
+        "vocab_file = os.path.join(gs_data_dir, vocab_name)\n",
+        "\n",
+        "gs_ckpt_dir = \"gs://tensor2tensor-checkpoints/\"\n",
+        "ckpt_name = \"transformer_ende_test\"\n",
+        "gs_ckpt = os.path.join(gs_ckpt_dir, ckpt_name)\n",
+        "\n",
+        "TRAIN_DIR_PRETRAINED = os.path.join(CHECKPOINT_DIR_PRETRAINED, ckpt_name)\n",
+        "\n",
+        "!gsutil cp {vocab_file} {DATA_DIR_PRETRAINED}\n",
+        "!gsutil -q cp -R {gs_ckpt} {CHECKPOINT_DIR_PRETRAINED}\n",
+        "\n",
+        "CHECKPOINT_NAME_PRETRAINED = tf.train.latest_checkpoint(TRAIN_DIR_PRETRAINED) # for translate with code\n"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "LP6cro9Xbygf",
+        "colab_type": "text"
+      },
+      "source": [
+        "##8.4. Translate"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "CBoNpy5HbzoF",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!echo \"the business of the house\" > \"inputs.en\"\n",
+        "!echo -e \"das Geschäft des Hauses\" > \"reference.de\"\n",
+        "\n",
+        "!t2t-decoder \\\n",
+        "  --data_dir=$DATA_DIR_PRETRAINED \\\n",
+        "  --problem=$PROBLEM_PRETRAINED  \\\n",
+        "  --model=$MODEL_PRETRAINED  \\\n",
+        "  --hparams_set=$HPARAMS_PRETRAINED \\\n",
+        "  --output_dir=$TRAIN_DIR_PRETRAINED  \\\n",
+        "  --decode_hparams=\"beam_size=1\" \\\n",
+        "  --decode_from_file=\"inputs.en\" \\\n",
+        "  --decode_to_file=\"outputs.de\"\n",
+        "\n",
+        "# See the translations\n",
+        "!cat outputs.de\n",
+        "\n",
+        "!t2t-bleu \\\n",
+        "    --translation=outputs.de \\\n",
+        "    --reference=reference.de"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bKI4WF0DgoFd",
+        "colab_type": "text"
+      },
+      "source": [
+        "#9.  Add your dataset/problem\n",
+        "To add a new dataset/problem, subclass Problem and register it with @registry.register_problem. See TranslateEnfrWmt8k for an example: \n",
+        "https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/translate_enfr.py\n",
+        "\n",
+        "---\n",
+        "Adding your own components: https://github.com/tensorflow/tensor2tensor#adding-your-own-components\n",
+        "\n",
+        "---\n",
+        "\n",
+        "See this example: https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/test_data/example_usr_dir"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "mB1SIrJNqy1N",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "from tensor2tensor.utils import registry\n",
+        "\n",
+        "@registry.register_problem\n",
+        "class MyTranslateEnFr(translate_enfr.TranslateEnfrWmt8k):\n",
+        "\n",
+        "  def generator(self, data_dir, tmp_dir, train):\n",
+        "   #your code"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    }
+  ]
+}
\ No newline at end of file

From d02dd84633b8ac1366e28468cd0dfe70e3b36821 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 26 Aug 2019 13:39:27 -0700
Subject: [PATCH 2322/2720] Fix notebook link.

PiperOrigin-RevId: 265529943
---
 README.md           | 2 +-
 docs/walkthrough.md | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 66ece0c34..3ad60b8f4 100644
--- a/README.md
+++ b/README.md
@@ -227,7 +227,7 @@ which is close to state-of-the art. If training on a single GPU, try the
 or larger data-sets (e.g., for English-French), try the big model
 with `--hparams_set=transformer_big`.
 
-See this [example](https://github.com/tensorflow/tensor2tensor/blob/Transformer_tutorial/tensor2tensor/notebooks/Transformer_translate.ipynb) to know how the translation works.
+See this [example](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/Transformer_translate.ipynb) to know how the translation works.
 
 ## Basics
 
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 66ece0c34..3ad60b8f4 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -227,7 +227,7 @@ which is close to state-of-the art. If training on a single GPU, try the
 or larger data-sets (e.g., for English-French), try the big model
 with `--hparams_set=transformer_big`.
 
-See this [example](https://github.com/tensorflow/tensor2tensor/blob/Transformer_tutorial/tensor2tensor/notebooks/Transformer_translate.ipynb) to know how the translation works.
+See this [example](https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/Transformer_translate.ipynb) to know how the translation works.
 
 ## Basics
 

From ab10c536ab8189264aa7b81869bec349cdedd154 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Mon, 26 Aug 2019 14:57:37 -0700
Subject: [PATCH 2323/2720] Fix restarts in PPO

Fixing a bug left after a previous change. Added a test as well.

PiperOrigin-RevId: 265547360
---
 tensor2tensor/trax/rl/ppo_trainer.py      |  2 +-
 tensor2tensor/trax/rl/ppo_trainer_test.py | 46 ++++++++++++++++++-----
 2 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index 2a9992e2e..b36ecf1e8 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -489,7 +489,7 @@ def save(self):
     params_file = os.path.join(self._output_dir, "model-%06d.pkl" % self._epoch)
     with gfile.GFile(params_file, "wb") as f:
       pickle.dump(
-          (self._policy_and_value_net_params, self._model_state,
+          (self._policy_and_value_opt_state, self._model_state,
            self._total_opt_step), f)
     # Remove the old model files.
     for path in old_model_files:
diff --git a/tensor2tensor/trax/rl/ppo_trainer_test.py b/tensor2tensor/trax/rl/ppo_trainer_test.py
index a8b48ec47..e3ca84fc7 100644
--- a/tensor2tensor/trax/rl/ppo_trainer_test.py
+++ b/tensor2tensor/trax/rl/ppo_trainer_test.py
@@ -69,12 +69,10 @@ def tmp_dir(self):
     yield tmp
     gfile.rmtree(tmp)
 
-  def _run_training_loop(self, train_env, eval_env, output_dir, model=None):
+  def _make_trainer(self, train_env, eval_env, output_dir, model=None):
     if model is None:
       model = lambda: [layers.Dense(1)]
-    n_epochs = 2
-    # Run the training loop.
-    trainer = ppo_trainer.PPO(
+    return ppo_trainer.PPO(
         train_env=train_env,
         eval_env=eval_env,
         policy_and_value_model=model,
@@ -83,19 +81,19 @@ def _run_training_loop(self, train_env, eval_env, output_dir, model=None):
         random_seed=0,
         boundary=2,
     )
-    trainer.training_loop(n_epochs=n_epochs)
 
   def test_training_loop_cartpole(self):
     with self.tmp_dir() as output_dir:
-      self._run_training_loop(
+      trainer = self._make_trainer(
           train_env=self.get_wrapped_env("CartPole-v0", 2),
           eval_env=self.get_wrapped_env("CartPole-v0", 2),
           output_dir=output_dir,
       )
+      trainer.training_loop(n_epochs=2)
 
   def test_training_loop_cartpole_transformer(self):
     with self.tmp_dir() as output_dir:
-      self._run_training_loop(
+      trainer = self._make_trainer(
           train_env=self.get_wrapped_env("CartPole-v0", 2),
           eval_env=self.get_wrapped_env("CartPole-v0", 2),
           output_dir=output_dir,
@@ -109,6 +107,7 @@ def test_training_loop_cartpole_transformer(self):
               mode="train",
           ),
       )
+      trainer.training_loop(n_epochs=2)
 
   def test_training_loop_onlinetune(self):
     with self.tmp_dir() as output_dir:
@@ -128,11 +127,12 @@ def test_training_loop_onlinetune(self):
       gin.bind_parameter("OnlineTuneEnv.eval_steps", 2)
       gin.bind_parameter(
           "OnlineTuneEnv.output_dir", os.path.join(output_dir, "envs"))
-      self._run_training_loop(
+      trainer = self._make_trainer(
           train_env=self.get_wrapped_env("OnlineTuneEnv-v0", 2),
           eval_env=self.get_wrapped_env("OnlineTuneEnv-v0", 2),
           output_dir=output_dir,
       )
+      trainer.training_loop(n_epochs=2)
 
   def test_training_loop_simulated(self):
     n_actions = 5
@@ -204,11 +204,39 @@ def loss(params, batch, model_predict, state, rng, **kwargs):
           output_dir=output_dir,
       )
 
-      self._run_training_loop(
+      trainer = self._make_trainer(
           train_env=env_fn(),
           eval_env=env_fn(),
           output_dir=output_dir,
       )
+      trainer.training_loop(n_epochs=2)
+
+  def test_restarts(self):
+    with self.tmp_dir() as output_dir:
+      train_env = self.get_wrapped_env("CartPole-v0", 2)
+      eval_env = self.get_wrapped_env("CartPole-v0", 2)
+
+      # Train for 1 epoch and save.
+      trainer = self._make_trainer(
+          train_env=train_env,
+          eval_env=eval_env,
+          output_dir=output_dir,
+      )
+      self.assertEqual(trainer.epoch, 0)
+      trainer.training_loop(n_epochs=1)
+
+      # Restore from the saved state.
+      trainer = self._make_trainer(
+          train_env=train_env,
+          eval_env=eval_env,
+          output_dir=output_dir,
+      )
+      # This is 2 instead of 1 because epoch calculation is a little weird right
+      # now.
+      # TODO(pkozakowski): Fix.
+      self.assertEqual(trainer.epoch, 2)
+      # Check that we can continue training from the restored checkpoint.
+      trainer.training_loop(n_epochs=3)
 
 
 if __name__ == "__main__":

From 8f96713b9c64858ea2e8ae6fcf3284bdda32d958 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Mon, 26 Aug 2019 15:13:13 -0700
Subject: [PATCH 2324/2720] Initialize BatchNorm running variance to 1 to
 prevent numerical issues

PiperOrigin-RevId: 265551074
---
 tensor2tensor/trax/layers/normalization.py      | 2 +-
 tensor2tensor/trax/layers/normalization_test.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/layers/normalization.py b/tensor2tensor/trax/layers/normalization.py
index c5f0cca9a..05cef6276 100644
--- a/tensor2tensor/trax/layers/normalization.py
+++ b/tensor2tensor/trax/layers/normalization.py
@@ -51,7 +51,7 @@ def get_stats_axis(i, d):
         return d
     stats_shape = tuple(get_stats_axis(i, d) for i, d in enumerate(input_shape))
     running_mean = np.zeros(stats_shape, dtype=np.float32)
-    running_var = np.zeros(stats_shape, dtype=np.float32)
+    running_var = np.ones(stats_shape, dtype=np.float32)
     num_batches = np.zeros((), dtype=np.int32)
     return (beta, gamma), (running_mean, running_var, num_batches)
 
diff --git a/tensor2tensor/trax/layers/normalization_test.py b/tensor2tensor/trax/layers/normalization_test.py
index d5a8067ee..d5ab55b76 100644
--- a/tensor2tensor/trax/layers/normalization_test.py
+++ b/tensor2tensor/trax/layers/normalization_test.py
@@ -47,7 +47,7 @@ def test_batch_norm(self):
     layer = normalization.BatchNorm(axis=(0, 1, 2))
     params, state = layer.initialize(input_shape, input_dtype, rng)
     onp.testing.assert_allclose(state[0], 0)
-    onp.testing.assert_allclose(state[1], 0)
+    onp.testing.assert_allclose(state[1], 1)
     self.assertEqual(state[2], 0)
     out, state = layer(inp1, params, state)
     onp.testing.assert_allclose(state[0], m1)

From 03897eb71a22d110d24be3bc14b1217844dd9409 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 26 Aug 2019 16:36:00 -0700
Subject: [PATCH 2325/2720] Use continuous eval streams in Trax. Use
 momentum=0.999 by default in batch norm.

PiperOrigin-RevId: 265567912
---
 tensor2tensor/trax/layers/normalization.py    | 13 +++-----
 .../trax/layers/normalization_test.py         | 26 +++------------
 tensor2tensor/trax/trax.py                    | 32 +++++++++++++------
 3 files changed, 31 insertions(+), 40 deletions(-)

diff --git a/tensor2tensor/trax/layers/normalization.py b/tensor2tensor/trax/layers/normalization.py
index 05cef6276..b58f362f0 100644
--- a/tensor2tensor/trax/layers/normalization.py
+++ b/tensor2tensor/trax/layers/normalization.py
@@ -27,7 +27,7 @@ class BatchNorm(base.Layer):
   """Batch normalization."""
 
   def __init__(self, axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True,
-               momentum=None, mode='train'):
+               momentum=0.999, mode='train'):
     super(BatchNorm, self).__init__()
     self._axis = axis
     self._epsilon = epsilon
@@ -66,15 +66,10 @@ def call(self, x, params, state, **unused_kwargs):
       m1 = np.mean(x**2, self._axis, keepdims=True)
       var = m1 - mean**2
       num_batches = num_batches + 1
-      if self._momentum is None:
-        # A simple average over all batches seen so far
-        exponential_average_factor = 1.0 / num_batches
-      else:
-        exponential_average_factor = self._momentum
       def average(factor, new, old):
-        return (factor * new + (1 - factor) * old).astype(old.dtype)
-      running_mean = average(exponential_average_factor, mean, running_mean)
-      running_var = average(exponential_average_factor, var, running_var)
+        return (factor * old + (1 - factor) * new).astype(old.dtype)
+      running_mean = average(self._momentum, mean, running_mean)
+      running_var = average(self._momentum, var, running_var)
       state = (running_mean, running_var, num_batches)
     else:
       mean = running_mean
diff --git a/tensor2tensor/trax/layers/normalization_test.py b/tensor2tensor/trax/layers/normalization_test.py
index d5ab55b76..00595b197 100644
--- a/tensor2tensor/trax/layers/normalization_test.py
+++ b/tensor2tensor/trax/layers/normalization_test.py
@@ -42,37 +42,19 @@ def test_batch_norm(self):
     rng = backend.random.get_prng(0)
     inp1 = np.reshape(np.arange(np.prod(input_shape), dtype=input_dtype),
                       input_shape)
-    m1 = 11.5
-    v1 = 47.9167
+    m1 = 11.5  # Mean of this random input.
+    v1 = 47.9167  # Variance of this random input.
     layer = normalization.BatchNorm(axis=(0, 1, 2))
     params, state = layer.initialize(input_shape, input_dtype, rng)
     onp.testing.assert_allclose(state[0], 0)
     onp.testing.assert_allclose(state[1], 1)
     self.assertEqual(state[2], 0)
     out, state = layer(inp1, params, state)
-    onp.testing.assert_allclose(state[0], m1)
-    onp.testing.assert_allclose(state[1], v1, rtol=1e-6)
+    onp.testing.assert_allclose(state[0], m1 * 0.001)
+    onp.testing.assert_allclose(state[1], 0.999 + v1 * 0.001, rtol=1e-6)
     self.assertEqual(state[2], 1)
     onp.testing.assert_allclose(out, (inp1 - m1) / np.sqrt(v1 + eps),
                                 rtol=1e-6)
-    inp2 = inp1 * 2 + 3
-    m2 = m1 * 2 + 3
-    v2 = v1 * 4
-    m12 = (m1 + m2) / 2
-    v12 = (v1 + v2) / 2
-    out, state = layer(inp2, params, state)
-    onp.testing.assert_allclose(state[0], m12)
-    onp.testing.assert_allclose(state[1], v12, rtol=1e-6)
-    self.assertEqual(state[2], 2)
-    onp.testing.assert_allclose(out, (inp2 - m2) / np.sqrt(v2 + eps),
-                                rtol=1e-6)
-    layer = normalization.BatchNorm(axis=(0, 1, 2), mode="eval")
-    inp3 = inp1 * 5 + 7
-    out, state_unchanged = layer(inp3, params, state)
-    for i in range(3):
-      onp.testing.assert_allclose(state_unchanged[i], state[i])
-    onp.testing.assert_allclose(out, (inp3 - m12) / np.sqrt(v12 + eps),
-                                rtol=1e-6)
 
   def test_layer_norm_shape(self):
     input_shape = (29, 5, 7, 20)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 8b970a1a7..a87980b94 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -239,15 +239,15 @@ def _print_n_params(opt_state, n_devices, step):
 }
 
 
-def evaluate_train_and_eval(step, inputs, predict_fn, eval_steps, state, rng,
-                            has_weights,
+def evaluate_train_and_eval(step, eval_stream, train_eval_stream,
+                            predict_fn, eval_steps, state, rng, has_weights,
                             train_sw=None, eval_sw=None, history=None):
   """Evalaute on train and eval data, and log metrics."""
   step_log(step, "Evaluation")
   metrics_list = []
-  for input_stream in [inputs.train_eval_stream, inputs.eval_stream]:
+  for input_stream in [eval_stream, train_eval_stream]:
     metrics, state = evaluate(  # pylint: disable=g-complex-comprehension
-        itertools.islice(input_stream(), eval_steps),
+        itertools.islice(input_stream, eval_steps),
         predict_fn,
         _METRICS,
         state,
@@ -292,13 +292,14 @@ def evaluate(inputs_stream, predict_fn, metric_fns, state, rng, has_weights):
   return {m: v / count for (m, v) in six.iteritems(metrics)}, state
 
 
-def evaluate_loss_train_and_eval(step, inputs, compute_loss_fn, eval_steps,
+def evaluate_loss_train_and_eval(step, eval_stream, train_eval_stream,
+                                 compute_loss_fn, eval_steps,
                                  state, rngs, has_weights,
                                  train_sw=None, eval_sw=None, history=None):
   """More efficient evaluation that logs only the loss on train & eval data."""
   step_log(step, "Evaluation")
   train_eval_metrics = []
-  for input_stream in [inputs.train_eval_stream, inputs.eval_stream]:
+  for input_stream in [train_eval_stream, eval_stream]:
     total = 0.0
     count = 0.0
     for inp in itertools.islice(input_stream(), eval_steps):
@@ -499,6 +500,13 @@ def reshape_by_device(x, n_devices):
       x, lambda x: _reshape_by_device_single(x, n_devices))
 
 
+def _repeat_stream(stream):
+  """Repeat a stream indefinitely."""
+  while True:
+    for example in stream():
+      yield example
+
+
 @gin.configurable(whitelist=[])
 class Trainer(object):
   """Trax trainer.
@@ -601,8 +609,12 @@ def reset(self, output_dir):
     self._train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
     self._eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
 
-    # Reset the training stream.
+    # Reset the train and eval streams.
     self._train_stream = self._inputs.train_stream()
+    # TODO(lukaszkaiser): add an option to evaluate exactly on the full eval
+    #   set by adding a padding and stopping the stream when too large.
+    self._eval_stream = _repeat_stream(self._inputs.eval_stream)
+    self._train_eval_stream = _repeat_stream(self._inputs.train_eval_stream)
 
     # Restore the training state.
     state = restore_state(output_dir)
@@ -729,7 +741,8 @@ def evaluate(self, eval_steps):
     _, rng = jax_random.split(self._rngs[0])
     _, _, self._model_state = evaluate_train_and_eval(
         step=self._step,
-        inputs=self._inputs,
+        eval_stream=self._eval_stream,
+        train_eval_stream=self._train_eval_stream,
         predict_fn=functools.partial(self._jit_model_predict_eval,
                                      params=self._opt_state[0]),
         eval_steps=eval_steps,
@@ -789,7 +802,8 @@ def evaluate(self, eval_steps):
     # Evaluate only the loss function (a more efficient, jitted, implementation)
     self._model_state = evaluate_loss_train_and_eval(
         step=self._step,
-        inputs=self._inputs,
+        eval_stream=self._eval_stream,
+        train_eval_stream=self._train_eval_stream,
         compute_loss_fn=functools.partial(self._jit_compute_loss,
                                           self._opt_state),
         eval_steps=eval_steps,

From 4cedd2e0487f8618e52a97176789255b5a018153 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 26 Aug 2019 16:47:46 -0700
Subject: [PATCH 2326/2720] Correct a typo from previous commit.

PiperOrigin-RevId: 265570170
---
 tensor2tensor/trax/trax.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index a87980b94..3b9c9b956 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -302,7 +302,7 @@ def evaluate_loss_train_and_eval(step, eval_stream, train_eval_stream,
   for input_stream in [train_eval_stream, eval_stream]:
     total = 0.0
     count = 0.0
-    for inp in itertools.islice(input_stream(), eval_steps):
+    for inp in itertools.islice(input_stream, eval_steps):
       loss_values, state, rngs = compute_loss_fn(inp, state, rngs, has_weights)
       total += float(numpy.mean(loss_values))
       count += 1.0

From 604d5a71db53a1f2abecdd084a3708ab5e9d0201 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 26 Aug 2019 22:47:48 -0700
Subject: [PATCH 2327/2720] Fix typo in docstring.

PiperOrigin-RevId: 265614544
---
 tensor2tensor/trax/optimizers/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/optimizers/base.py b/tensor2tensor/trax/optimizers/base.py
index af592d0f5..992331e8d 100644
--- a/tensor2tensor/trax/optimizers/base.py
+++ b/tensor2tensor/trax/optimizers/base.py
@@ -167,7 +167,7 @@ def update(self, step, grads, params, slots, opt_params):
 
 
 class Momentum(Optimizer):
-  """Nestrov momentum optimizer."""
+  """Nesterov momentum optimizer."""
 
   def __init__(self, learning_rate, mass=0.9):  # pylint: disable=useless-super-delegation
     super(Momentum, self).__init__(learning_rate, mass)

From ac59eec33b90b4685f8b6c634aa0d0f80529dd93 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Mon, 26 Aug 2019 23:04:13 -0700
Subject: [PATCH 2328/2720] Decouple lengths of input and target sequences.

The sequence lengths are now specified to utils.run.sequence_length as a dictionary containing the keys "inputs" and "targets", for example, {"inputs": 1024, "targets": 256"}.  For datasets with no inputs (i.e. language modeling), the dictionary need not contain an "inputs" key.

Also, remove from defaults.gin the %sequence_length macro, which used to control both the runtime sequence length and the sizes of the positional embedding tables.  They now have to be specified separately - for example:

# runtime length
utils.run.sequence_length = {"inputs": 512, "targets": 512} # runtime length
# size of the positional embedding tables
Unitransformer.max_length = 512

PiperOrigin-RevId: 265616475
---
 .../data_generators/generator_utils.py        |  2 +-
 .../data_generators/ops/pack_sequences_ops.cc | 38 ++++++++++++-------
 .../ops/pack_sequences_ops_test.py            |  3 +-
 3 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 92e330743..4c9394c61 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -758,7 +758,7 @@ def map_fn_custom(x):
     """Map-function."""
     (k1_packed, k1_segmengation, k1_position,
      k2_packed, k2_segmentation, k2_position) = (
-         pack_sequences_ops.pack_sequences2(x[k1], x[k2], length))
+         pack_sequences_ops.pack_sequences2(x[k1], x[k2], length, length))
     packed = {
         k1: k1_packed,
         k1 + "_segmentation": k1_segmengation,
diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
index 13b5906a1..9fcb79f43 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
@@ -20,7 +20,8 @@ using ::tensorflow::shape_inference::InferenceContext;
 REGISTER_OP("PackSequences2")
     .Input("inputs: int64")
     .Input("targets: int64")
-    .Input("max_length: int32")
+    .Input("inputs_max_length: int32")
+    .Input("targets_max_length: int32")
     .Output("inputs_packed: int64")
     .Output("inputs_segmentation: int32")
     .Output("inputs_position: int32")
@@ -44,12 +45,13 @@ class PackSequences2Op : public OpKernel {
   void Compute(OpKernelContext* ctx) override {
     auto inputs = ctx->input(0).matrix<int64>();
     auto targets = ctx->input(1).matrix<int64>();
-    int max_length = ctx->input(2).scalar<int32>()();
+    int inputs_max_length = ctx->input(2).scalar<int32>()();
+    int targets_max_length = ctx->input(3).scalar<int32>()();
     int n = inputs.dimension(0);
     std::vector<int> inputs_lengths(n);
     std::vector<int> targets_lengths(n);
     int padded_inputs_length = min(static_cast<int>(inputs.dimension(1)),
-                                   max_length);
+                                   inputs_max_length);
     for (int i = 0; i < n; i++) {
       for (int j = 0; j < padded_inputs_length; j++) {
           if (inputs(i, j) != 0)
@@ -57,7 +59,7 @@ class PackSequences2Op : public OpKernel {
       }
     }
     int padded_targets_length = min(static_cast<int>(targets.dimension(1)),
-                                    max_length);
+                                    targets_max_length);
     for (int i = 0; i < n; i++) {
       for (int j = 0; j < padded_targets_length; j++) {
           if (targets(i, j) != 0)
@@ -80,9 +82,9 @@ class PackSequences2Op : public OpKernel {
           break;
         } else if (
             (combined_inputs_length[combined_id] + inputs_length
-             <= max_length) &&
+             <= inputs_max_length) &&
             (combined_targets_length[combined_id] + targets_length
-             <= max_length)) {
+             <= targets_max_length)) {
           combined_inputs_length[combined_id] += inputs_length;
           combined_targets_length[combined_id] += targets_length;
           combined_sequence_ids[combined_id].push_back(seq_id);
@@ -91,40 +93,48 @@ class PackSequences2Op : public OpKernel {
       }
     }
 
-    auto output_shape = TensorShape(
-        {static_cast<int64>(num_combined), static_cast<int64>(max_length)});
+    auto output_shape_inputs = TensorShape(
+        {static_cast<int64>(num_combined),
+         static_cast<int64>(inputs_max_length)});
+    auto output_shape_targets = TensorShape(
+        {static_cast<int64>(num_combined),
+         static_cast<int64>(targets_max_length)});
 
     Tensor* inputs_packed;
-    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, output_shape, &inputs_packed));
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(
+        0, output_shape_inputs, &inputs_packed));
     auto inputs_packed_m = inputs_packed->matrix<int64>();
     inputs_packed_m.setZero();
 
     Tensor* inputs_segmentation;
     OP_REQUIRES_OK(
-        ctx, ctx->allocate_output(1, output_shape, &inputs_segmentation));
+        ctx, ctx->allocate_output(
+            1, output_shape_inputs, &inputs_segmentation));
     auto inputs_segmentation_m = inputs_segmentation->matrix<int32>();
     inputs_segmentation_m.setZero();
 
     Tensor* inputs_position;
     OP_REQUIRES_OK(
-        ctx, ctx->allocate_output(2, output_shape, &inputs_position));
+        ctx, ctx->allocate_output(2, output_shape_inputs, &inputs_position));
     auto inputs_position_m = inputs_position->matrix<int32>();
     inputs_position_m.setZero();
 
     Tensor* targets_packed;
-    OP_REQUIRES_OK(ctx, ctx->allocate_output(3, output_shape, &targets_packed));
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(
+        3, output_shape_targets, &targets_packed));
     auto targets_packed_m = targets_packed->matrix<int64>();
     targets_packed_m.setZero();
 
     Tensor* targets_segmentation;
     OP_REQUIRES_OK(
-        ctx, ctx->allocate_output(4, output_shape, &targets_segmentation));
+        ctx, ctx->allocate_output(
+            4, output_shape_targets, &targets_segmentation));
     auto targets_segmentation_m = targets_segmentation->matrix<int32>();
     targets_segmentation_m.setZero();
 
     Tensor* targets_position;
     OP_REQUIRES_OK(
-        ctx, ctx->allocate_output(5, output_shape, &targets_position));
+        ctx, ctx->allocate_output(5, output_shape_targets, &targets_position));
     auto targets_position_m = targets_position->matrix<int32>();
     targets_position_m.setZero();
 
diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
index 77b42e7da..3531d7631 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
@@ -39,7 +39,8 @@ def test_pack_sequences(self):
     max_length = 5
     (inputs_packed, inputs_segmentation, inputs_position,
      targets_packed, targets_segmentation, targets_position) = (
-         pack_sequences_ops.pack_sequences2(inputs, targets, max_length))
+         pack_sequences_ops.pack_sequences2(
+             inputs, targets, max_length, max_length))
     self.assertAllEqual(
         inputs_packed, [
             [1, 2, 3, 4, 5],

From 453b8721f969b97f29df649ff3eedbaec9b052cc Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 27 Aug 2019 10:05:47 -0700
Subject: [PATCH 2329/2720] Make Reformer run again

PiperOrigin-RevId: 265707949
---
 tensor2tensor/trax/layers/base.py             | 13 ++--
 tensor2tensor/trax/layers/base_test.py        |  4 +-
 tensor2tensor/trax/layers/reversible.py       | 35 ++++++-----
 .../models/research/transformer_revnet.py     | 62 ++++++++++++-------
 tensor2tensor/trax/trax.py                    | 11 +++-
 5 files changed, 75 insertions(+), 50 deletions(-)

diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 842ec3025..8bd41544b 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -257,12 +257,12 @@ def __call__(self, x, params=(), state=(), **kwargs):
       #   remove the constraints on state below when this feature is added to
       #   JAX.
 
-      assert state is (), (  # pylint: disable=literal-comparison
+      assert not jax.tree_util.tree_leaves(state), (
           'Custom gradients require trivial start state. Got %s' % str(state))
 
       def check_end_state(output_state):
         output, state = output_state
-        assert state is (), (  # pylint: disable=literal-comparison
+        assert not jax.tree_util.tree_leaves(state), (
             'Custom gradients require trivial end state. Got %s' % str(state))
         return output
 
@@ -271,20 +271,21 @@ def check_end_state(output_state):
       # Note that we capture the kwargs and don't calculate gradients wrt. them.
       @jax.custom_transforms
       def do_call(y, params):
-        return check_end_state(self.call(y, params=params, state=(), **kwargs))
+        return check_end_state(self.call(y, params=params, state=state,
+                                         **kwargs))
 
       # This is the custom gradient (vector-jacobian product in JAX) function.
       # For the exact specification of this custom transformation see this link:
       # https://jax.readthedocs.io/en/latest/jax.html#jax.defjvp_all
       def do_call_vjp(y, params):
-        output = check_end_state(self.call(y, params=params, state=(),
+        output = check_end_state(self.call(y, params=params, state=state,
                                            **kwargs))
         def vjpfun(grad):
-          return self.custom_grad(y, output, grad, params, **kwargs)
+          return self.custom_grad(y, output, grad, params, state, **kwargs)
         return output, vjpfun
 
       jax.defvjp_all(do_call, do_call_vjp)
-      return do_call(x, params), ()
+      return do_call(x, params), state
 
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback()
diff --git a/tensor2tensor/trax/layers/base_test.py b/tensor2tensor/trax/layers/base_test.py
index 103040c7e..26140248c 100644
--- a/tensor2tensor/trax/layers/base_test.py
+++ b/tensor2tensor/trax/layers/base_test.py
@@ -50,7 +50,7 @@ def new_parameters(self, input_shapes, input_dtype, rng):
       def has_custom_grad(self):
         return True
 
-      def custom_grad(self, inputs, output, ct, params, **kwargs):
+      def custom_grad(self, inputs, output, ct, params, state, **kwargs):
         return (backend.numpy.zeros_like(ct), ())
 
     layer = IdWithZeroGrad()
@@ -80,7 +80,7 @@ def new_parameters(self, input_shapes, input_dtype, rng):
       def has_custom_grad(self):
         return True
 
-      def custom_grad(self, inputs, output, ct, params, **kwargs):
+      def custom_grad(self, inputs, output, ct, params, state, **kwargs):
         return (inputs, ())
 
     layer = IdWithIdGrad()
diff --git a/tensor2tensor/trax/layers/reversible.py b/tensor2tensor/trax/layers/reversible.py
index f10e45a50..98038b8c0 100644
--- a/tensor2tensor/trax/layers/reversible.py
+++ b/tensor2tensor/trax/layers/reversible.py
@@ -28,11 +28,11 @@
 class ReversibleLayer(base.Layer):
   """Reversible Layer."""
 
-  def reverse(self, output, params=(), **kwargs):
+  def reverse(self, output, params=(), state=(), **kwargs):
     """Reverse this layer: compute input given output."""
     raise NotImplementedError
 
-  def reverse_and_grad(self, output, grad, params=(), **kwargs):
+  def reverse_and_grad(self, output, grad, params=(), state=(), **kwargs):
     """Backward pass: computes the inverse of a layer and propagates gradients.
 
     While you may choose to only implement reverse, some layers implement this
@@ -44,6 +44,7 @@ def reverse_and_grad(self, output, grad, params=(), **kwargs):
       grad: gradient signal (cotangent) computed based on subsequent layers.
         The structure and shape must match the output.
       params: layer parameters
+      state: start state
       **kwargs: kwargs for the layer
 
     Returns:
@@ -53,9 +54,10 @@ def reverse_and_grad(self, output, grad, params=(), **kwargs):
     """
     # Note: jax.vjp does not allow us to use **kwargs in the signature here.
     def _do_call(x, params):
-      return super(ReversibleLayer, self).call(x, params=params, **kwargs)[0]
+      return super(ReversibleLayer, self).call(
+          x, params=params, state=state, **kwargs)[0]
 
-    reconstructed_x = self.reverse(output, params, **kwargs)
+    reconstructed_x = self.reverse(output, params, state, **kwargs)
     _, vjpfun = jax.vjp(_do_call, reconstructed_x, params)
     x_params_grad = vjpfun(grad)
     return reconstructed_x, x_params_grad
@@ -64,18 +66,19 @@ def _do_call(x, params):
   def has_custom_grad(self):
     return True
 
-  def custom_grad(self, inputs, output, ct, params, **kwargs):
+  def custom_grad(self, inputs, output, ct, params, state, **kwargs):
     del inputs
-    _, inputs_params_ct = self.reverse_and_grad(output, ct, params, **kwargs)
+    _, inputs_params_ct = self.reverse_and_grad(output, ct, params, state,
+                                                **kwargs)
     return inputs_params_ct
 
 
 class ReversibleSwap(ReversibleLayer, cb.Swap):
   """Swap the first two element on the stack."""
 
-  def reverse(self, output, params=(), **kwargs):
-    # Swap is its own inverse
-    return self.call(output, params, **kwargs)
+  def reverse(self, output, params=(), state=(), **kwargs):
+    # Swap is its own inverse, except that reverse doesn't return the state.
+    return self.call(output, params, state, **kwargs)[0]
 
 
 class ReversibleSerial(ReversibleLayer, cb.Serial):
@@ -91,19 +94,20 @@ def __init__(self, *layers):
             'Sub-layer {} of ReversibleSerial is not reversible: {}'.format(
                 i, layer))
 
-  def reverse(self, output, params=(), **kwargs):
+  def reverse(self, output, params=(), state=(), **kwargs):
     rng = kwargs.pop('rng', None)
     rngs = (None,) * self._n_layers
     if rng is not None:
       rngs = backend.random.split(rng, self._n_layers)
 
     layer_val = output
-    for layer, p, rng in reversed(zip(self.sublayers(), params, rngs)):
-      layer_val = layer.reverse(layer_val, p, rng=rng, **kwargs)
+    for layer, p, s, rng in reversed(zip(self.sublayers(),
+                                         params, state, rngs)):
+      layer_val = layer.reverse(layer_val, p, s, rng=rng, **kwargs)
 
     return layer_val
 
-  def reverse_and_grad(self, output, ct, params=(), **kwargs):
+  def reverse_and_grad(self, output, ct, params=(), state=(), **kwargs):
     rng = kwargs.pop('rng', None)
     rngs = (None,) * self._n_layers
     if rng is not None:
@@ -112,9 +116,10 @@ def reverse_and_grad(self, output, ct, params=(), **kwargs):
     layer_val = output
     layer_ct = ct
     params_ct = []
-    for layer, p, rng in reversed(zip(self.sublayers(), params, rngs)):
+    for layer, p, s, rng in reversed(zip(self.sublayers(),
+                                         params, state, rngs)):
       layer_val, layer_ct = layer.reverse_and_grad(
-          layer_val, layer_ct, p, rng=rng, **kwargs)
+          layer_val, layer_ct, p, s, rng=rng, **kwargs)
       layer_ct, p_ct = layer_ct
       params_ct.insert(0, p_ct)
 
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index ff5880b0f..adaaf480b 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -70,6 +70,8 @@ def call(self, inputs, params=(), state=(), **kwargs):
     results = [self._layer(x, params=params, state=state, rng=r, **kwargs)
                for x, r in zip(inputs, rngs)]
     result_outputs, result_states = zip(*results)
+    # TODO(kitaev): think about how to merge state across copies in the map.
+    result_states = result_states[0]
     return tuple(result_outputs), tuple(result_states)
 
   def new_parameters(self, input_shape, input_dtype, rng):
@@ -181,7 +183,7 @@ def call(self, inputs, params=(), state=(), **kwargs):
     res = [backend.numpy.concatenate(ys, -1) for ys in zip(x1_split, x2_split)]
     return tuple(res), state
 
-  def reverse(self, output, params=(), **kwargs):
+  def reverse(self, output, params=(), state=(), **kwargs):
     del params, kwargs
 
     x1_split = []
@@ -196,7 +198,7 @@ def reverse(self, output, params=(), **kwargs):
 
     return (x1, x2)
 
-  def reverse_and_grad(self, output, ct, params=(), **kwargs):
+  def reverse_and_grad(self, output, ct, params=(), state=(), **kwargs):
     del params, kwargs
     return self.reverse(output), (self.reverse(ct), ())
 
@@ -241,7 +243,7 @@ def __init__(self, residual_layers):
     self.subtract_top = tl.Parallel(tl.SubtractTop(), [])
     self.reverse_layers = [self.compute_residual, self.subtract_top]
 
-  def reverse(self, output, params=(), **kwargs):
+  def reverse(self, output, params=(), state=(), **kwargs):
     reconstructed_x = output
     rng = kwargs.pop('rng', None)
     rngs = (None,) * self._n_layers
@@ -249,29 +251,32 @@ def reverse(self, output, params=(), **kwargs):
       rngs = backend.random.split(rng, self._n_layers)
     # Note that self.sublayers() aligns exactly with self.reverse_layers in
     # terms of parameter and rng usage, so no re-ordering is required.
-    for layer, p, rng in zip(self.reverse_layers, params, rngs):
-      reconstructed_x = layer(reconstructed_x, p, rng=rng, **kwargs)
+    for layer, p, s, rng in zip(self.reverse_layers, params, state, rngs):
+      reconstructed_x, _ = layer(reconstructed_x, p, s, rng=rng, **kwargs)
     return reconstructed_x
 
-  def reverse_and_grad(self, output, ct, params=(), **kwargs):
+  def reverse_and_grad(self, output, ct, params=(), state=(), **kwargs):
     rng = kwargs.pop('rng', None)
     rngs = (None,) * self._n_layers
     if rng is not None:
       rngs = backend.random.split(rng, self._n_layers)
 
     def call_compute_residual(x, params):
-      return self.compute_residual(x, params, rng=rngs[0], **kwargs)
+      res, _ = self.compute_residual(x, params, state[0], rng=rngs[0], **kwargs)
+      return res
 
     assert len(ct) == 2
     ct = ((ct[0], ct[0], ct[1]))
 
     stack_with_residual, vjpfun = jax.vjp(
         call_compute_residual, output, params[0])
-    reconstructed_x = self.subtract_top(
-        stack_with_residual, params[-1], rng=rngs[-1], **kwargs)
+    reconstructed_x, _ = self.subtract_top(
+        stack_with_residual, params[-1], state[-1], rng=rngs[-1], **kwargs)
 
     x_ct, residual_params_ct = vjpfun(ct)
-    return reconstructed_x, (x_ct, (residual_params_ct, ()))
+    assert not jax.tree_util.tree_leaves(params[-1])
+    add_top_params_ct = params[-1]
+    return reconstructed_x, (x_ct, [residual_params_ct, add_top_params_ct])
 
 
 class ComputeAttentionHeads(tl.Layer):
@@ -388,7 +393,8 @@ def call(self, inputs, params=(), state=(), rng=None, **kwargs):
   def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
     # Simultaneous forward pass and backprop through the attention mechanism.
     def do_call(x):
-      return self.call(x, params, **kwargs)
+      res, _ = self.call(x, params, **kwargs)
+      return res
     output, vjpfun = jax.vjp(do_call, inputs)
     return output, vjpfun(ct)[0]
 
@@ -697,7 +703,7 @@ def __init__(self, pre_attention, attention, post_attention):
         self.subtract_top,
     ]
 
-  def reverse(self, output, params=(), **kwargs):
+  def reverse(self, output, params=(), state=(), **kwargs):
     rng = kwargs.pop('rng', None)
     rngs = (None,) * self._n_layers
     if rng is not None:
@@ -706,11 +712,12 @@ def reverse(self, output, params=(), **kwargs):
     reconstructed_x = output
     # Note that self.sublayers() aligns exactly with self.reverse_layers in
     # terms of parameter and rng usage, so no re-ordering is required.
-    for layer, p, rng in zip(self.reverse_layers, params, rngs):
-      reconstructed_x = layer.reverse(reconstructed_x, p, rng=rng, **kwargs)
+    for layer, p, s, rng in zip(self.reverse_layers, params, state, rngs):
+      reconstructed_x, _ = layer.reverse(reconstructed_x, p, s, rng=rng,
+                                         **kwargs)
     return reconstructed_x
 
-  def reverse_and_grad(self, output, ct, params=(), **kwargs):
+  def reverse_and_grad(self, output, ct, params=(), state=(), **kwargs):
     rng = kwargs.pop('rng', None)
     rngs = (None,) * self._n_layers
     if rng is not None:
@@ -719,7 +726,8 @@ def reverse_and_grad(self, output, ct, params=(), **kwargs):
     # Forward pass through self.pre_attention, while preparing for
     # later backprop.
     def call_pre_attention(x, params):
-      return self.pre_attention(x, params, rng=rngs[0], **kwargs)
+      res, _ = self.pre_attention(x, params, state[0], rng=rngs[0], **kwargs)
+      return res
     stack, pre_attention_vjpfun = jax.vjp(call_pre_attention, output, params[0])
 
     # Backprop through adding the residual
@@ -728,7 +736,9 @@ def call_pre_attention(x, params):
 
     # Backprop through self.post_attention with respect to the inputs only
     def call_post_attention(x):
-      return self.post_attention(x, params[2], rng=rngs[2], **kwargs)
+      res, _ = self.post_attention(x, params[2], state[2], rng=rngs[2],
+                                   **kwargs)
+      return res
     # Note: these are *not* the actual inputs to self.post_attention.
     # If self.post_attention is not linear, we will get incorrect gradients.
     dummy_inputs = (stack[-3], stack[-2], stack[-1])
@@ -738,7 +748,7 @@ def call_post_attention(x):
     # Simultaneous forward pass and backprop through the attention mechanism
     stack, ct = self.attention.forward_and_vjp(
         stack, ct, rng=rngs[1], **kwargs)
-    attention_params_ct = ()
+    attention_params_ct = params[1]  # Note: this assumes that params are empty.
 
     # Backprop through self.pre_attention
     x_ct, pre_attention_params_ct = pre_attention_vjpfun(ct)
@@ -746,20 +756,24 @@ def call_post_attention(x):
     # Forward pass for self.post_attention, and backprop with respect to the
     # parameters only
     def call_post_attention2(params):
-      return self.post_attention(stack, params, rng=rngs[2], **kwargs)
+      res, _ = self.post_attention(stack, params, state[2], rng=rngs[2],
+                                   **kwargs)
+      return res
     stack, post_attention_vjpfun = jax.vjp(call_post_attention2, params[2])
     (post_attention_params_ct,) = post_attention_vjpfun(saved_ct)
 
     # Forward pass through subtracting the residual
-    reconstructed_x = self.subtract_top(
-        stack, params[-1], rng=rngs[-1], **kwargs)
+    reconstructed_x, _ = self.subtract_top(
+        stack, params[-1], state[-1], rng=rngs[-1], **kwargs)
 
-    params_ct = (
+    assert not jax.tree_util.tree_leaves(params[-1])
+    add_top_params_ct = params[-1]
+    params_ct = [
         pre_attention_params_ct,
         attention_params_ct,
         post_attention_params_ct,
-        (),
-        )
+        add_top_params_ct,
+    ]
 
     return reconstructed_x, (x_ct, params_ct)
 
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 3b9c9b956..2b48454be 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -297,13 +297,15 @@ def evaluate_loss_train_and_eval(step, eval_stream, train_eval_stream,
                                  state, rngs, has_weights,
                                  train_sw=None, eval_sw=None, history=None):
   """More efficient evaluation that logs only the loss on train & eval data."""
+  assert not has_weights, (
+      "MemoryEfficientTrainer doesn't support has_weights")
   step_log(step, "Evaluation")
   train_eval_metrics = []
   for input_stream in [train_eval_stream, eval_stream]:
     total = 0.0
     count = 0.0
     for inp in itertools.islice(input_stream, eval_steps):
-      loss_values, state, rngs = compute_loss_fn(inp, state, rngs, has_weights)
+      loss_values, state, rngs = compute_loss_fn(inp, state, rngs)
       total += float(numpy.mean(loss_values))
       count += 1.0
     metrics = {"loss": total / count}
@@ -453,7 +455,8 @@ def _jit_compute_loss_fn(predict_fn, loss_fn, n_devices, jit=True):
   if n_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
     def single_compute_loss(opt_state, batch, state, rng):
       rng, subrng = jax_random.split(rng[0])
-      return loss_fn(opt_state[0], batch, predict_fn, state, rng), [subrng]
+      loss_val, state = loss_fn(opt_state[0], batch, predict_fn, state, rng)
+      return loss_val, state, [subrng]
     if jit:
       return backend.jit(single_compute_loss)
     else:
@@ -797,10 +800,12 @@ def __init__(self, *args, **kwargs):
     # we only implement computing the loss, and not any other metrics.
     self._jit_compute_loss = _jit_compute_loss_fn(
         self._model_predict_eval, self._loss_fn, self._n_devices)
+    assert not self._has_weights, (
+        "MemoryEfficientTrainer doesn't support has_weights")
 
   def evaluate(self, eval_steps):
     # Evaluate only the loss function (a more efficient, jitted, implementation)
-    self._model_state = evaluate_loss_train_and_eval(
+    _, _, self._model_state = evaluate_loss_train_and_eval(
         step=self._step,
         eval_stream=self._eval_stream,
         train_eval_stream=self._train_eval_stream,

From ffe4cf4051a8fc517daf7fde0c1b74dfc7a8eb19 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 27 Aug 2019 10:31:22 -0700
Subject: [PATCH 2330/2720] Point out the difference between positional
 embedding from the published paper.

PiperOrigin-RevId: 265713853
---
 tensor2tensor/layers/common_attention.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 33116776b..5a1d6e43a 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -447,6 +447,8 @@ def get_timing_signal_1d(length,
   inv_timescales = min_timescale * tf.exp(
       tf.to_float(tf.range(num_timescales)) * -log_timescale_increment)
   scaled_time = tf.expand_dims(position, 1) * tf.expand_dims(inv_timescales, 0)
+  # Please note that this slightly differs from the published paper.
+  # See a discussion here: https://github.com/tensorflow/tensor2tensor/pull/177
   signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
   signal = tf.pad(signal, [[0, 0], [0, tf.mod(channels, 2)]])
   signal = tf.reshape(signal, [1, length, channels])

From 5ca390bb8fc6304662f39f65a7ef4ededa0c3a12 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 27 Aug 2019 15:39:13 -0700
Subject: [PATCH 2331/2720] Unswap train and eval streams in evaluation

PiperOrigin-RevId: 265782808
---
 tensor2tensor/trax/trax.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 2b48454be..9da70f266 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -245,7 +245,7 @@ def evaluate_train_and_eval(step, eval_stream, train_eval_stream,
   """Evalaute on train and eval data, and log metrics."""
   step_log(step, "Evaluation")
   metrics_list = []
-  for input_stream in [eval_stream, train_eval_stream]:
+  for input_stream in [train_eval_stream, eval_stream]:
     metrics, state = evaluate(  # pylint: disable=g-complex-comprehension
         itertools.islice(input_stream, eval_steps),
         predict_fn,
@@ -254,6 +254,7 @@ def evaluate_train_and_eval(step, eval_stream, train_eval_stream,
         rng,
         has_weights)
     metrics_list.append(metrics)
+  # Unpack in the same order we've iterated over streams in the loop above.
   train_metrics, eval_metrics = metrics_list  # pylint: disable=unbalanced-tuple-unpacking
   if train_sw:
     log_metrics(train_metrics, train_sw, "train", step, history=history)
@@ -310,6 +311,7 @@ def evaluate_loss_train_and_eval(step, eval_stream, train_eval_stream,
       count += 1.0
     metrics = {"loss": total / count}
     train_eval_metrics.append(metrics)
+  # Unpack in the same order we've iterated over streams in the loop above.
   train_metrics, eval_metrics = train_eval_metrics  # pylint: disable=unbalanced-tuple-unpacking
   if train_sw:
     log_metrics(train_metrics, train_sw, "train", step, history=history)

From 470a9d0266ea6c9bb8cedee68acea8deecf4d99c Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 27 Aug 2019 15:54:04 -0700
Subject: [PATCH 2332/2720] Reset the new environment after changing it in PPO.

In SimPLe, we change the PPO training environment from real to simulated and back. When we do that, we need to reset the new environment. This is not done automatically, because we can have trajectories spanning over multiple PPO epochs.

PiperOrigin-RevId: 265785727
---
 tensor2tensor/trax/rl/ppo_trainer.py    | 24 +++++++++++++++++++++++-
 tensor2tensor/trax/rl/simple_trainer.py |  2 +-
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index b36ecf1e8..9f8e0c95f 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -111,6 +111,10 @@ def __init__(
       eval_temperatures: Sequence of temperatures to try for categorical
         sampling during evaluation.
     """
+    # Set in base class constructor.
+    self._train_env = None
+    self._should_reset = None
+
     super(PPO, self).__init__(train_env, eval_env, output_dir)
 
     self._n_optimizer_steps = n_optimizer_steps
@@ -177,11 +181,29 @@ def __init__(
     self._eval_sw = jaxboard.SummaryWriter(
         os.path.join(self._output_dir, "eval"))
 
-    self._should_reset = True
     self._n_trajectories_done = 0
 
     self._last_saved_at = 0
 
+  @property
+  def train_env(self):
+    return self._train_env
+
+  @train_env.setter
+  def train_env(self, new_train_env):
+    if self._train_env is not None:
+      def assert_same_space(space1, space2):
+        assert space1.shape == space2.shape
+        assert space1.dtype == space2.dtype
+      assert_same_space(
+          new_train_env.observation_space, self._train_env.observation_space)
+      assert_same_space(
+          new_train_env.action_space, self._train_env.action_space)
+      # We don't check the reward range, because PPO will work either way.
+
+    self._train_env = new_train_env
+    self._should_reset = True
+
   @property
   def epoch(self):
     return self._epoch
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index ada78297f..4d9f7a34d 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -25,7 +25,7 @@
 
 from absl import logging
 import cloudpickle as pickle
-from jax import numpy as np
+import numpy as np
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.rl import base_trainer

From 117e3e9969045280cb8a1938ed829961be337cb3 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 27 Aug 2019 16:29:06 -0700
Subject: [PATCH 2333/2720] Adds support for ImageModality to Parallel
 Scheduled Sampling

PiperOrigin-RevId: 265792695
---
 tensor2tensor/utils/t2t_model.py | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index e7b361ca3..e05d277dd 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1821,11 +1821,15 @@ def maybe_scheduled_sampling(self, features, logits, losses):
 
     # Only do scheduled sampling on language tasks.
     modality = problem_hparams.modality["targets"]
-    if modality not in [modalities.ModalityType.SYMBOL,
-                        modalities.ModalityType.SYMBOL_WEIGHTS_ALL]:
+    if modality not in [
+        modalities.ModalityType.SYMBOL,
+        modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
+        modalities.ModalityType.IMAGE
+    ]:
       assert hparams.scheduled_sampling_prob == 0, (
-          "Scheduled sampling only applies to ModalityType.{SYMBOL, "
-          "SYMBOL_WEIGHTS_ALL}. Set hparams.scheduled_sampling_prob == 0.0.")
+          "Scheduled sampling only applies to ModalityType.(SYMBOL, "
+          "SYMBOL_WEIGHTS_ALL, IMAGE). Found {modality}. Set "
+          "hparams.scheduled_sampling_prob == 0.0.").format(modality=modality)
       return (logits, losses)
 
     # Only do scheduled sampling when training.
@@ -1875,11 +1879,12 @@ def is_later_timestep(x, pass_idx):
       """Constructs mask based on timestep."""
       assert x.shape.ndims == 4, x.shape
       x_shape = tf.shape(x)
-      batch_size = x_shape[0]
       num_timesteps = x_shape[1]
       timesteps = tf.range(num_timesteps)
       timesteps = tf.reshape(timesteps, [1, num_timesteps, 1, 1])
-      timesteps = tf.tile(timesteps, [batch_size, 1, 1, 1])
+      # The following is a bit untrue. For images, "num_timesteps" actually
+      # represents image height, not time. We ignore that fact here.
+      timesteps = tf.broadcast_to(timesteps, x_shape)
       return tf.greater_equal(timesteps, pass_idx)
 
     # TODO(duckworthd): Move to scheduled_sampling.py.

From 76803debef448a0f2c56b40b3427cef2e9c8b5be Mon Sep 17 00:00:00 2001
From: joao guilherme <joaoguilhermearujo@gmail.com>
Date: Tue, 27 Aug 2019 21:03:54 -0300
Subject: [PATCH 2334/2720] WIP: Added new initializers (#1666)

* Added new initializers

* Fixed glorot uniform

* Fixed order of init arguments

* Added uniform random and documentation

* Added requested changes to initializers.py

* Added requested changes to initializers_test.py

* Added choice of input axis and output axis

* Added choice of input axis and output axis

* fixed errors and implemented requested changes

* fixed typo

* tests passing

* fixed get fans

* Fixed numpy weird behavior

* Fixed typo
---
 tensor2tensor/trax/layers/initializers.py     | 107 +++++++++++++++---
 .../trax/layers/initializers_test.py          |  48 ++++++++
 2 files changed, 142 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/trax/layers/initializers.py b/tensor2tensor/trax/layers/initializers.py
index 9505d13f9..05459ad75 100644
--- a/tensor2tensor/trax/layers/initializers.py
+++ b/tensor2tensor/trax/layers/initializers.py
@@ -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.
-
 """Trax initializers."""
 
 from __future__ import absolute_import
@@ -23,28 +22,110 @@
 from tensor2tensor.trax import backend
 
 
+def _get_fans(shape, out_dim=-1, in_dim=-2):
+  #temporary fix until numpy.delete supports negative indices
+  if out_dim < 0:
+    out_dim += len(shape)
+  if in_dim < 0:
+    in_dim += len(shape)
+
+  receptive_field = backend.numpy.prod(onp.delete(shape, [in_dim, out_dim]))
+  if len(shape) >= 2:
+    fan_in, fan_out = shape[in_dim], shape[out_dim]
+  elif len(shape) == 1:
+    fan_in, fan_out = shape[0]
+  else:
+    fan_in, fan_out = 1.
+    fan_in *= receptive_field
+    fan_out *= receptive_field
+  return fan_in, fan_out
+
+
 def RandomNormalInitializer(stddev=1e-2):
   """An initializer function for random normal coefficients."""
+
   def Init(shape, rng):
     return (stddev * backend.random.normal(rng, shape)).astype('float32')
+
   return Init
 
 
-def GlorotNormalInitializer(out_dim=0, in_dim=1, scale=onp.sqrt(2)):
-  """An initializer function for random Glorot-scaled coefficients."""
+def RandomUniformInitializer(lim=1.0):
+  """An initializer function for random uniform coefficients."""
+
   def Init(shape, rng):
-    fan_in, fan_out = shape[in_dim], shape[out_dim]
-    size = onp.prod(onp.delete(shape, [in_dim, out_dim]))
-    std = scale / backend.numpy.sqrt((fan_in + fan_out) / 2. * size)
-    return (std * backend.random.normal(rng, shape)).astype('float32')
+    return (backend.random.uniform(rng, shape, backend.numpy.float32, -lim,
+                                   lim))
+
   return Init
 
 
-def GlorotUniformInitializer(out_dim=0, in_dim=1):
-  """An initializer function for random uniform Glorot-scaled coefficients."""
+def VarianceScalingInitializer(out_dim, in_dim, scale, mode, distribution):
+  """Initializer capable of adapting its scale to the shape of weights tensors."""
+  if scale <= 0.:
+    raise ValueError('scale must be positive float, {} given'.format(scale))
+  if mode not in {'fan_in', 'fan_out', 'fan_avg'}:
+    raise ValueError(
+        'Invalid mode argument:, {}, must be either fan_in, fan_out or fan_avg'
+        .format(mode))
+
   def Init(shape, rng):
-    fan_in, fan_out = shape[in_dim], shape[out_dim]
-    std = backend.numpy.sqrt(2.0 / (fan_in + fan_out))
-    a = backend.numpy.sqrt(3.0) * std
-    return backend.random.uniform(rng, shape, minval=-a, maxval=a)
+    fan_in, fan_out = _get_fans(shape, out_dim, in_dim)
+    gain = scale
+    if mode == 'fan_in':
+      gain /= fan_in
+    elif mode == 'fan_out':
+      gain /= fan_out
+    elif mode == 'fan_avg':
+      gain /= (fan_in + fan_out) / 2
+    if distribution == 'truncated_normal':
+      # constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.)
+      stddev = backend.numpy.sqrt(gain) / .87962566103423978
+      return (backend.random.truncated_normal(rng, -2, 2, shape) *
+              stddev).astype('float32')
+    elif distribution == 'normal':
+      return (backend.random.normal(rng, shape) *
+              backend.numpy.sqrt(gain)).astype('float32')
+    elif distribution == 'uniform':
+      lim = backend.numpy.sqrt(3. * gain)
+      return (backend.random.uniform(rng, shape, backend.numpy.float32, -lim,
+                                     lim))
+    else:
+      raise ValueError('invalid distribution for variance scaling Initializer')
+
   return Init
+
+
+def GlorotNormalInitializer(out_dim=-1, in_dim=-2, scale=1.):
+  """An initializer function for random Glorot-scaled coefficients."""
+  return VarianceScalingInitializer(out_dim, in_dim, scale, 'fan_avg', 'normal')
+
+
+def GlorotUniformInitializer(out_dim=-1, in_dim=-2, scale=1.):
+  """An initializer function for random uniform Glorot-scaled coefficients."""
+  return VarianceScalingInitializer(out_dim, in_dim, scale, 'fan_avg',
+                                    'uniform')
+
+
+def LeCunNormalInitializer(out_dim=-1, in_dim=-2, scale=1.):
+  """An initializer function for random LeCun-scaled coefficients."""
+  return VarianceScalingInitializer(out_dim, in_dim, scale, 'fan_in', 'normal')
+
+
+def LeCunUniformInitializer(out_dim=-1, in_dim=-2, scale=1.):
+  """An initializer function for random uniform LeCun-scaled coefficients."""
+  return VarianceScalingInitializer(out_dim, in_dim, scale, 'fan_in', 'uniform')
+
+
+def KaimingNormalInitializer(out_dim=-1, in_dim=-2, param=0.):
+  """An initializer function for random Kaiming-scaled coefficients."""
+  return VarianceScalingInitializer(out_dim, in_dim,
+                                    2.0 / backend.numpy.sqrt(1 + param**2),
+                                    'fan_in', 'normal')
+
+
+def KaimingUniformInitializer(out_dim=-1, in_dim=-2, param=0.):
+  """An initializer function for random uniform Kaiming-scaled coefficients."""
+  return VarianceScalingInitializer(out_dim, in_dim,
+                                    2.0 / backend.numpy.sqrt(1 + param**2),
+                                    'fan_in', 'uniform')
diff --git a/tensor2tensor/trax/layers/initializers_test.py b/tensor2tensor/trax/layers/initializers_test.py
index bb5ad1b0c..d28c992b6 100644
--- a/tensor2tensor/trax/layers/initializers_test.py
+++ b/tensor2tensor/trax/layers/initializers_test.py
@@ -32,5 +32,53 @@ def test_random_normal(self):
     self.assertEqual(tuple(init_value.shape), input_shape)
 
 
+  def test_random_uniform(self):
+    initializer = initializers.RandomUniformInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+  def test_glorot_normal(self):
+    initializer = initializers.GlorotNormalInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+
+  def test_glorot_uniform(self):
+    initializer = initializers.GlorotUniformInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+
+  def test_lecun_normal(self):
+    initializer = initializers.LeCunNormalInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+
+  def test_lecun_uniform(self):
+    initializer = initializers.LeCunUniformInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+
+  def test_kaiming_normal(self):
+    initializer = initializers.KaimingNormalInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+
+  def test_kaiming_uniform(self):
+    initializer = initializers.KaimingUniformInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+
 if __name__ == "__main__":
   absltest.main()

From 488bdaa173f17eec7e9313e2cc14e36a2105fdc0 Mon Sep 17 00:00:00 2001
From: joao guilherme <joaoguilhermearujo@gmail.com>
Date: Tue, 27 Aug 2019 17:04:20 -0700
Subject: [PATCH 2335/2720] Merge of PR #1666

PiperOrigin-RevId: 265800010
---
 tensor2tensor/trax/layers/initializers.py      | 17 +++++++++++------
 tensor2tensor/trax/layers/initializers_test.py |  6 ------
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/trax/layers/initializers.py b/tensor2tensor/trax/layers/initializers.py
index 05459ad75..97d334a56 100644
--- a/tensor2tensor/trax/layers/initializers.py
+++ b/tensor2tensor/trax/layers/initializers.py
@@ -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.
+
 """Trax initializers."""
 
 from __future__ import absolute_import
@@ -22,8 +23,9 @@
 from tensor2tensor.trax import backend
 
 
-def _get_fans(shape, out_dim=-1, in_dim=-2):
-  #temporary fix until numpy.delete supports negative indices
+def _GetFans(shape, out_dim=-1, in_dim=-2):
+  """Get the fan-in and fan-out sizes for the given shape and dims."""
+  # Temporary fix until numpy.delete supports negative indices.
   if out_dim < 0:
     out_dim += len(shape)
   if in_dim < 0:
@@ -33,9 +35,11 @@ def _get_fans(shape, out_dim=-1, in_dim=-2):
   if len(shape) >= 2:
     fan_in, fan_out = shape[in_dim], shape[out_dim]
   elif len(shape) == 1:
-    fan_in, fan_out = shape[0]
+    fan_in = shape[0]
+    fan_out = shape[0]
   else:
-    fan_in, fan_out = 1.
+    fan_in = 1.
+    fan_out = 1.
     fan_in *= receptive_field
     fan_out *= receptive_field
   return fan_in, fan_out
@@ -61,7 +65,7 @@ def Init(shape, rng):
 
 
 def VarianceScalingInitializer(out_dim, in_dim, scale, mode, distribution):
-  """Initializer capable of adapting its scale to the shape of weights tensors."""
+  """Initializer capable of adapting its scale to the shape of weights."""
   if scale <= 0.:
     raise ValueError('scale must be positive float, {} given'.format(scale))
   if mode not in {'fan_in', 'fan_out', 'fan_avg'}:
@@ -70,7 +74,8 @@ def VarianceScalingInitializer(out_dim, in_dim, scale, mode, distribution):
         .format(mode))
 
   def Init(shape, rng):
-    fan_in, fan_out = _get_fans(shape, out_dim, in_dim)
+    """The initializer function."""
+    fan_in, fan_out = _GetFans(shape, out_dim, in_dim)
     gain = scale
     if mode == 'fan_in':
       gain /= fan_in
diff --git a/tensor2tensor/trax/layers/initializers_test.py b/tensor2tensor/trax/layers/initializers_test.py
index d28c992b6..3a5586e96 100644
--- a/tensor2tensor/trax/layers/initializers_test.py
+++ b/tensor2tensor/trax/layers/initializers_test.py
@@ -31,7 +31,6 @@ def test_random_normal(self):
     init_value = initializer(input_shape, random.get_prng(0))
     self.assertEqual(tuple(init_value.shape), input_shape)
 
-
   def test_random_uniform(self):
     initializer = initializers.RandomUniformInitializer()
     input_shape = (29, 5, 7, 20)
@@ -44,35 +43,30 @@ def test_glorot_normal(self):
     init_value = initializer(input_shape, random.get_prng(0))
     self.assertEqual(tuple(init_value.shape), input_shape)
 
-
   def test_glorot_uniform(self):
     initializer = initializers.GlorotUniformInitializer()
     input_shape = (29, 5, 7, 20)
     init_value = initializer(input_shape, random.get_prng(0))
     self.assertEqual(tuple(init_value.shape), input_shape)
 
-
   def test_lecun_normal(self):
     initializer = initializers.LeCunNormalInitializer()
     input_shape = (29, 5, 7, 20)
     init_value = initializer(input_shape, random.get_prng(0))
     self.assertEqual(tuple(init_value.shape), input_shape)
 
-
   def test_lecun_uniform(self):
     initializer = initializers.LeCunUniformInitializer()
     input_shape = (29, 5, 7, 20)
     init_value = initializer(input_shape, random.get_prng(0))
     self.assertEqual(tuple(init_value.shape), input_shape)
 
-
   def test_kaiming_normal(self):
     initializer = initializers.KaimingNormalInitializer()
     input_shape = (29, 5, 7, 20)
     init_value = initializer(input_shape, random.get_prng(0))
     self.assertEqual(tuple(init_value.shape), input_shape)
 
-
   def test_kaiming_uniform(self):
     initializer = initializers.KaimingUniformInitializer()
     input_shape = (29, 5, 7, 20)

From 8edb68ca31a02fe96a591e2ca4bd38a0e447277b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 27 Aug 2019 19:34:28 -0700
Subject: [PATCH 2336/2720] Move memory-efficient attention to attention.py

PiperOrigin-RevId: 265819997
---
 .../transformer_revnet_imagenet64_8gb.gin     |  14 +-
 tensor2tensor/trax/layers/__init__.py         |   5 +
 tensor2tensor/trax/layers/attention.py        | 269 ++++++++++++-
 tensor2tensor/trax/models/__init__.py         |   7 -
 .../models/research/transformer_revnet.py     | 369 +-----------------
 tensor2tensor/trax/models/transformer.py      |   6 +-
 6 files changed, 289 insertions(+), 381 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
index 26ba18308..1ba24f292 100644
--- a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
@@ -32,14 +32,14 @@ train.optimizer = @trax.optimizers.Adafactor
 train.train_steps = 500000
 train.trainer_class = @MemoryEfficientTrainer
 
-# Parameters for DotProductAttention:
+# Parameters for DotProductCausalAttention:
 # ==============================================================================
-DotProductAttention.dropout = 0.0
+DotProductCausalAttention.dropout = 0.0
 
-# Parameters for MemoryEfficientDotProductAttention:
+# Parameters for MemoryEfficientCausalAttention:
 # ==============================================================================
-MemoryEfficientDotProductAttention.dropout = 0.0
-MemoryEfficientDotProductAttention.loop_stride = 512
+MemoryEfficientCausalAttention.dropout = 0.0
+MemoryEfficientCausalAttention.loop_stride = 512
 
 # Parameters for DummyHashedAttention:
 # ==============================================================================
@@ -60,5 +60,5 @@ TransformerRevnetLM.n_layers = 6
 TransformerRevnetLM.vocab_size = 256
 TransformerRevnetLM.n_chunks = 16
 TransformerRevnetLM.n_attention_chunks = 1
-TransformerRevnetLM.attention_type = \
-    @trax.models.MemoryEfficientDotProductAttention
+TransformerRevnetLM.attention_type = @trax.layers.MemoryEfficientCausalAttention
+
diff --git a/tensor2tensor/trax/layers/__init__.py b/tensor2tensor/trax/layers/__init__.py
index f6a956b02..03427d6bb 100644
--- a/tensor2tensor/trax/layers/__init__.py
+++ b/tensor2tensor/trax/layers/__init__.py
@@ -50,3 +50,8 @@ def layer_configure(*args, **kwargs):
 LogSoftmax = layer_configure(LogSoftmax)
 Softmax = layer_configure(Softmax)
 Softplus = layer_configure(Softplus)
+
+DotProductCausalAttention = layer_configure(
+    DotProductCausalAttention, blacklist=["mode"])
+MemoryEfficientCausalAttention = layer_configure(
+    MemoryEfficientCausalAttention, blacklist=["mode"])
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index f1ecd26b7..e50b9a99f 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -26,6 +26,7 @@
 from tensor2tensor.trax.layers import base
 from tensor2tensor.trax.layers import combinators as cb
 from tensor2tensor.trax.layers import core
+from tensor2tensor.trax.layers import initializers as init
 
 
 @base.layer()
@@ -194,9 +195,12 @@ def Attention(d_feature, n_heads=1, dropout=0.0, mode='train'):
   ]
 
 
-def CausalAttention(d_feature, n_heads=1, dropout=0.0, mode='train'):
+def BasicCausalAttention(d_feature, n_heads=1, dropout=0.0, mode='train'):
   """Transformer-style multi-headed causal attention.
 
+  This implementation is less configurable than the CausalAttention layer
+  defined below, but it shares code with the non-causal attention.
+
   # TODO(jonni,lukaszkaiser): standardize and improve layer comments.
   Accepts inputs of the form x and constructs (q, k, v) and causal mask from x.
 
@@ -225,3 +229,266 @@ def ShiftRight(x, **unused_kwargs):
   padded = np.pad(x, pad_widths, mode='constant',
                   constant_values=x.dtype.type(0))
   return padded[:, :-1]
+
+
+class ComputeAttentionHeads(base.Layer):
+  """Computes queries/keys/values via linear projection.
+
+  The output shape is (n_batch * n_heads, seqlen, d_head); the batch and head
+  dimensions are fused to allow for more efficient memory layouts.
+  """
+
+  def __init__(self, n_heads=1, d_head=64,
+               kernel_initializer=init.GlorotUniformInitializer()):
+    super(ComputeAttentionHeads, self).__init__()
+    self._n_heads = n_heads
+    self._d_head = d_head
+    self._kernel_initializer = kernel_initializer
+    # The lack of a bias term here is consistent with the tensor2tensor
+    # implementation, and shouldn't have an effect on modeling quality.
+    # Note that AttentionQKV above is different in that it uses a bias term.
+
+  def call(self, x, params, state, **kwargs):
+    del kwargs
+    seqlen = x.shape[1]
+    res = np.dot(x, params)
+
+    # n_batch, seqlen, n_heads*d_head -> n_batch, seqlen, n_heads, d_head
+    res = np.reshape(res, (x.shape[0], seqlen, self._n_heads, self._d_head))
+    # n_batch, seqlen, n_heads, d_head -> n_batch, n_heads, seqlen, d_head
+    res = np.transpose(res, (0, 2, 1, 3))
+    # n_batch, n_heads, seqlen, d_head -> n_batch*n_heads, seqlen, d_head
+    res = np.reshape(res, (-1, seqlen, self._d_head))
+
+    return res, state
+
+  def new_parameters(self, input_shape, input_dtype, rng):
+    del input_dtype
+    w = self._kernel_initializer(
+        (input_shape[-1], self._n_heads * self._d_head), rng)
+    return w, ()
+
+
+class ComputeAttentionOutput(base.Layer):
+  """Joins outputs from different heads via linear projection."""
+
+  def __init__(self, n_heads=1, d_model=1024,
+               kernel_initializer=init.GlorotUniformInitializer()):
+    super(ComputeAttentionOutput, self).__init__()
+    self._n_heads = n_heads
+    self._d_model = d_model
+    self._kernel_initializer = kernel_initializer
+    # The lack of a bias term here is consistent with the tensor2tensor
+    # implementation, and shouldn't have an effect on modeling quality.
+    # Note that AttentionQKV above is different in that it uses a bias term.
+
+  def call(self, x, params, state, **kwargs):
+    del kwargs
+    seqlen = x.shape[1]
+    d_head = x.shape[2]
+
+    x = np.reshape(x, (-1, self._n_heads, seqlen, d_head))
+    x = np.transpose(x, (0, 2, 1, 3))  # -> n_batch, seqlen, n_heads, d_head
+    x = np.reshape(x, (-1, seqlen, self._n_heads * d_head))
+
+    return np.dot(x, params), state
+
+  def new_parameters(self, input_shape, input_dtype, rng):
+    del input_dtype
+    w = self._kernel_initializer(
+        (input_shape[-1] * self._n_heads, self._d_model), rng)
+    return w, ()
+
+
+class BaseCausalAttention(base.Layer):
+  """Base class for variants of causal self-attention.
+
+  This class sets up an API that includes forward_and_vjp, which is required to
+  implement MemoryEfficientCausalAttention.
+  """
+
+  def call(self, inputs, params=(), state=(), rng=None, **kwargs):
+    raise NotImplementedError()
+
+  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
+    raise NotImplementedError()
+
+  def new_parameters(self, input_shapes, input_dtype, rng):
+    return (), ()
+
+  def n_inputs(self):
+    return 3
+
+  def n_outputs(self):
+    return 1
+
+
+class DotProductCausalAttention(BaseCausalAttention):
+  """A standard (non-memory-efficient) dot product attention implementation."""
+
+  def __init__(self, dropout, mode):
+    super(DotProductCausalAttention, self).__init__()
+    self._dropout = dropout
+    self._mode = mode
+
+  def call(self, inputs, params=(), state=(), rng=None, **kwargs):
+    del params
+    q, k, v = inputs
+    mask_size = q.shape[-2]
+    mask = np.tril(np.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
+    res = DotProductAttention(
+        q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=rng)
+    return res, state
+
+  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
+    assert backend.get_name() == 'jax', (
+        'JAX backend is required to use forward_and_vjp.')
+    # Simultaneous forward pass and backprop through the attention mechanism.
+    def do_call(x):  # pylint: disable=invalid-name
+      res, _ = self.call(x, params, **kwargs)
+      return res
+    output, vjpfun = jax.vjp(do_call, inputs)
+    return output, vjpfun(ct)[0]
+
+
+class MemoryEfficientCausalAttention(BaseCausalAttention):
+  """Memory-efficient dot product attention."""
+
+  def __init__(self, loop_stride, dropout, mode):
+    assert backend.get_name() == 'jax', (
+        'JAX backend is required to use MemoryEfficientCausalAttention.')
+    super(MemoryEfficientCausalAttention, self).__init__()
+    self._loop_stride = loop_stride
+    if dropout >= 1.0:
+      raise ValueError('Dropout rates must be lower than 1.')
+    if mode == 'train':
+      self.dropout = dropout
+    else:
+      self.dropout = None
+
+  def call(self, inputs, params=(), state=(), **kwargs):
+    output, _ = self.forward_and_vjp(inputs, None, params=params, **kwargs)
+    return output, state
+
+  def forward_and_vjp(self, inputs, ct, params=(), rng=None, **kwargs):
+    # This is the core of the memory-efficient attention implementation, where
+    # we use the jax.lax.while_loop primitive to compute attention for a small
+    # set of query positions at a time. Note how in the backwards pass, we
+    # compute both the forward direction (to recover the previous layer's
+    # activations) and the backward direction simultaneously. This allows us to
+    # only use a single loop, where the inner portion of the loop does a slice
+    # of the forward+backward joint computation. Unfortunately we have had to
+    # introduce a large number of wrapper classes (including
+    # ReversibleAttentionHalfResidual and ApplyAttentionWrapper) for the sole
+    # purpose of connecting this implementation of forward_and_vjp with the core
+    # backprop implementation.
+
+    query, key, value = inputs
+    depth = np.shape(query)[-1]
+    do_backprop = ct is not None
+
+    def make_mask(N, M, k):  # pylint: disable=invalid-name
+      x = np.arange(N, dtype=np.int32)
+      y = np.arange(M, dtype=np.int32)
+      mask = jax.lax.lt(
+          (jax.lax.broadcast_in_dim(
+              x, shape=(N, M), broadcast_dimensions=(0,)) + k),
+          jax.lax.broadcast(y, [N]))
+      mask = jax.lax.convert_element_type(mask, np.float32)
+      return mask
+
+    def forward_slice(query_slice, q_loop_idx, key, value):  # pylint: disable=invalid-name
+      """Forward pass for a subset of the query vectors."""
+      dots = np.matmul(
+          query_slice, np.swapaxes(key, -1, -2)) / np.sqrt(depth)
+
+      # Causal masking
+      mask = make_mask(dots.shape[-2], dots.shape[-1], q_loop_idx)
+      dots = dots - 1e9 * mask
+
+      # Softmax.
+      dots = np.exp(dots - dots.max(axis=-1, keepdims=True))
+      dots = dots / dots.sum(axis=-1, keepdims=True)
+
+      if self.dropout is not None and self.dropout > 0.0:
+        # Dropout is broadcast across the batch+head dimension
+        dropout_shape = (1, dots.shape[-2], dots.shape[-1])
+        slice_rng = jax.random.fold_in(rng, q_loop_idx)
+        keep_prob = jax.lax.tie_in(dots, 1.0 - self.dropout)
+        keep = backend.random.bernoulli(slice_rng, keep_prob, dropout_shape)
+        multiplier = keep.astype(dots.dtype) / jax.lax.tie_in(keep, keep_prob)
+        dots = dots * multiplier
+
+      out_slice = np.matmul(dots, value)
+      return out_slice
+
+    def forward_and_vjp_slice(query_slice, q_loop_idx, key, value, ct_slice):  # pylint: disable=invalid-name
+      # Capture q_loop_idx to avoid calculated gradients wrt. it.
+      def forward_slice_with_q_loop_idx(query_slice, key, value):  # pylint: disable=invalid-name
+        return forward_slice(query_slice, q_loop_idx, key, value)
+
+      output_slice, vjpfun = jax.vjp(
+          forward_slice_with_q_loop_idx, query_slice, key, value)
+      return output_slice, vjpfun(ct_slice)
+
+    q_loop_idx = np.zeros((), dtype=np.int32)
+    q_loop_max = query.shape[-2]
+    q_loop_stride = self._loop_stride
+    assert q_loop_max % q_loop_stride == 0, (
+        'Stride must evenly divide the number of query elements.')
+
+    out_accum = np.zeros_like(query)
+    if do_backprop:
+      query_ct_accum = np.zeros_like(query)
+      key_ct_accum = np.zeros_like(key)
+      value_ct_accum = np.zeros_like(value)
+      init_vals = (
+          q_loop_idx, out_accum,
+          query_ct_accum, key_ct_accum, value_ct_accum)
+    else:
+      init_vals = (q_loop_idx, out_accum)
+
+    def cond_fun(vals):  # pylint: disable=invalid-name
+      q_loop_idx = vals[0]
+      return jax.lax.lt(q_loop_idx, q_loop_max)
+
+    def body_fun(vals):  # pylint: disable=invalid-name
+      """Compute a slice of the attention mechanism."""
+      if do_backprop:
+        (q_loop_idx, out_accum,
+         query_ct_accum, key_ct_accum, value_ct_accum) = vals
+      else:
+        q_loop_idx, out_accum = vals
+
+      query_slice = jax.lax.dynamic_slice_in_dim(
+          query, q_loop_idx, q_loop_stride, axis=-2)
+
+      if do_backprop:
+        ct_slice = jax.lax.dynamic_slice_in_dim(
+            ct, q_loop_idx, q_loop_stride, axis=-2)
+        out_slice, partial_ct = forward_and_vjp_slice(
+            query_slice, q_loop_idx, key, value, ct_slice)
+        query_ct_accum = jax.lax.dynamic_update_slice_in_dim(
+            query_ct_accum, partial_ct[0], q_loop_idx, axis=-2)
+        key_ct_accum = key_ct_accum + partial_ct[1]
+        value_ct_accum = value_ct_accum + partial_ct[2]
+      else:
+        out_slice = forward_slice(query_slice, q_loop_idx, key, value)
+
+      out_accum = jax.lax.dynamic_update_slice_in_dim(
+          out_accum, out_slice, q_loop_idx, axis=-2)
+      q_loop_idx = q_loop_idx + q_loop_stride
+
+      if do_backprop:
+        return (q_loop_idx, out_accum,
+                query_ct_accum, key_ct_accum, value_ct_accum)
+      else:
+        return (q_loop_idx, out_accum)
+
+    final_vals = jax.lax.while_loop(cond_fun, body_fun, init_vals)
+
+    if not do_backprop:
+      return final_vals[1], None
+    else:
+      return final_vals[1], final_vals[2:]
+
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index d1ee95066..b1a1ac895 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -49,10 +49,3 @@ def model_configure(*args, **kwargs):
 TransformerLM = model_configure(transformer.TransformerLM)
 TransformerRevnetLM = model_configure(transformer_revnet.TransformerRevnetLM)
 WideResnet = model_configure(resnet.WideResnet)
-
-DotProductAttention = model_configure(
-    transformer_revnet.DotProductAttention, blacklist=["mode"])
-MemoryEfficientDotProductAttention = model_configure(
-    transformer_revnet.MemoryEfficientDotProductAttention, blacklist=["mode"])
-DummyHashedAttention = model_configure(
-    transformer_revnet.DummyHashedAttention, blacklist=["mode"])
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index adaaf480b..62d334e2f 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -19,11 +19,9 @@
 from __future__ import print_function
 
 import jax
-import numpy as onp
 
 from tensor2tensor.trax import backend
 from tensor2tensor.trax import layers as tl
-from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.layers.combinators import _pop_rng_and_split
 
 
@@ -279,73 +277,6 @@ def call_compute_residual(x, params):
     return reconstructed_x, (x_ct, [residual_params_ct, add_top_params_ct])
 
 
-class ComputeAttentionHeads(tl.Layer):
-  """Computes queries/keys/values via linear projection.
-
-  The output shape is (n_batch * n_heads, seqlen, d_head); the batch and head
-  dimensions are fused to allow for more efficient memory layouts.
-  """
-
-  def __init__(self, n_heads=1, d_head=64,
-               kernel_initializer=tl.initializers.GlorotUniformInitializer()):
-    super(ComputeAttentionHeads, self).__init__()
-    self._n_heads = n_heads
-    self._d_head = d_head
-    self._kernel_initializer = kernel_initializer
-    # The lack of a bias term here is consistent with the tensor2tensor
-    # implementation, and shouldn't have an effect on modeling quality.
-
-  def call(self, x, params, state, **kwargs):
-    del kwargs
-    seqlen = x.shape[1]
-    res = np.dot(x, params)
-
-    # n_batch, seqlen, n_heads*d_head -> n_batch, seqlen, n_heads, d_head
-    res = np.reshape(res, (x.shape[0], seqlen, self._n_heads, self._d_head))
-    # n_batch, seqlen, n_heads, d_head -> n_batch, n_heads, seqlen, d_head
-    res = np.transpose(res, (0, 2, 1, 3))
-    # n_batch, n_heads, seqlen, d_head -> n_batch*n_heads, seqlen, d_head
-    res = np.reshape(res, (-1, seqlen, self._d_head))
-
-    return res, state
-
-  def new_parameters(self, input_shape, input_dtype, rng):
-    del input_dtype
-    w = self._kernel_initializer(
-        (input_shape[-1], self._n_heads * self._d_head), rng)
-    return w, ()
-
-
-class ComputeAttentionOutput(tl.Layer):
-  """Joins outputs from different heads via linear projection."""
-
-  def __init__(self, n_heads=1, d_model=1024,
-               kernel_initializer=tl.initializers.GlorotUniformInitializer()):
-    super(ComputeAttentionOutput, self).__init__()
-    self._n_heads = n_heads
-    self._d_model = d_model
-    self._kernel_initializer = kernel_initializer
-    # The lack of a bias term here is consistent with the tensor2tensor
-    # implementation, and shouldn't have an effect on modeling quality.
-
-  def call(self, x, params, state, **kwargs):
-    del kwargs
-    seqlen = x.shape[1]
-    d_head = x.shape[2]
-
-    x = np.reshape(x, (-1, self._n_heads, seqlen, d_head))
-    x = np.transpose(x, (0, 2, 1, 3))  # -> n_batch, seqlen, n_heads, d_head
-    x = np.reshape(x, (-1, seqlen, self._n_heads * d_head))
-
-    return np.dot(x, params), state
-
-  def new_parameters(self, input_shape, input_dtype, rng):
-    del input_dtype
-    w = self._kernel_initializer(
-        (input_shape[-1] * self._n_heads, self._d_model), rng)
-    return w, ()
-
-
 class ApplyAttentionWrapper(tl.Parallel):
   """Same as tl.Parallel(attention, [], []), but implements forward_and_vjp.
 
@@ -369,296 +300,6 @@ def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
     return (out,) + passthrough, qkv_ct + passthrough_ct
 
 
-class DotProductAttention(tl.Layer):
-  """A standard (non-memory-efficient) dot product attention implementation.
-
-  This class sets up the API that is required to implement
-  MemoryEfficientDotProductAttention.
-  """
-
-  def __init__(self, dropout, mode):
-    super(DotProductAttention, self).__init__()
-    self._dropout = dropout
-    self._mode = mode
-
-  def call(self, inputs, params=(), state=(), rng=None, **kwargs):
-    del params
-    q, k, v = inputs
-    mask_size = q.shape[-2]
-    mask = np.tril(np.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
-    res = tl.DotProductAttention(
-        q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=rng)
-    return res, state
-
-  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
-    # Simultaneous forward pass and backprop through the attention mechanism.
-    def do_call(x):
-      res, _ = self.call(x, params, **kwargs)
-      return res
-    output, vjpfun = jax.vjp(do_call, inputs)
-    return output, vjpfun(ct)[0]
-
-  def new_parameters(self, input_shapes, input_dtype, rng):
-    return (), ()
-
-  def n_inputs(self):
-    return 3
-
-  def n_outputs(self):
-    return 1
-
-
-class MemoryEfficientDotProductAttention(DotProductAttention):
-  """Memory-efficient dot product attention."""
-
-  def __init__(self, loop_stride, dropout, mode):
-    super(MemoryEfficientDotProductAttention, self).__init__(dropout, mode)
-    self._loop_stride = loop_stride
-    if dropout >= 1.0:
-      raise ValueError('Dropout rates must be lower than 1.')
-    if mode == 'train':
-      self.dropout = dropout
-    else:
-      self.dropout = None
-
-  def call(self, inputs, params=(), state=(), **kwargs):
-    output, _ = self.forward_and_vjp(inputs, None, params=params, **kwargs)
-    return output, state
-
-  def forward_and_vjp(self, inputs, ct, params=(), rng=None, **kwargs):
-    # This is the core of the memory-efficient attention implementation, where
-    # we use the jax.lax.while_loop primitive to compute attention for a small
-    # set of query positions at a time. Note how in the backwards pass, we
-    # compute both the forward direction (to recover the previous layer's
-    # activations) and the backward direction simultaneously. This allows us to
-    # only use a single loop, where the inner portion of the loop does a slice
-    # of the forward+backward joint computation. Unfortunately we have had to
-    # introduce a large number of wrapper classes (including
-    # ReversibleAttentionHalfResidual and ApplyAttentionWrapper) for the sole
-    # purpose of connecting this implementation of forward_and_vjp with the core
-    # backprop implementation.
-
-    query, key, value = inputs
-    depth = np.shape(query)[-1]
-    do_backprop = ct is not None
-
-    def make_mask(N, M, k):
-      x = np.arange(N, dtype=np.int32)
-      y = np.arange(M, dtype=np.int32)
-      mask = jax.lax.lt(
-          (jax.lax.broadcast_in_dim(
-              x, shape=(N, M), broadcast_dimensions=(0,)) + k),
-          jax.lax.broadcast(y, [N]))
-      mask = jax.lax.convert_element_type(mask, np.float32)
-      return mask
-
-    def forward_slice(query_slice, q_loop_idx, key, value):
-      """Forward pass for a subset of the query vectors."""
-      dots = np.matmul(
-          query_slice, np.swapaxes(key, -1, -2)) / np.sqrt(depth)
-
-      # Causal masking
-      mask = make_mask(dots.shape[-2], dots.shape[-1], q_loop_idx)
-      dots = dots - 1e9 * mask
-
-      # Softmax.
-      dots = np.exp(dots - dots.max(axis=-1, keepdims=True))
-      dots = dots / dots.sum(axis=-1, keepdims=True)
-
-      if self.dropout is not None and self.dropout > 0.0:
-        # Dropout is broadcast across the batch+head dimension
-        dropout_shape = (1, dots.shape[-2], dots.shape[-1])
-        slice_rng = jax.random.fold_in(rng, q_loop_idx)
-        keep_prob = jax.lax.tie_in(dots, 1.0 - self.dropout)
-        keep = backend.random.bernoulli(slice_rng, keep_prob, dropout_shape)
-        multiplier = keep.astype(dots.dtype) / jax.lax.tie_in(keep, keep_prob)
-        dots = dots * multiplier
-
-      out_slice = np.matmul(dots, value)
-      return out_slice
-
-    def forward_and_vjp_slice(query_slice, q_loop_idx, key, value, ct_slice):
-      # Capture q_loop_idx to avoid calculated gradients wrt. it.
-      def forward_slice_with_q_loop_idx(query_slice, key, value):
-        return forward_slice(query_slice, q_loop_idx, key, value)
-
-      output_slice, vjpfun = jax.vjp(
-          forward_slice_with_q_loop_idx, query_slice, key, value)
-      return output_slice, vjpfun(ct_slice)
-
-    q_loop_idx = np.zeros((), dtype=np.int32)
-    q_loop_max = query.shape[-2]
-    q_loop_stride = self._loop_stride
-    assert q_loop_max % q_loop_stride == 0, (
-        'Stride must evenly divide the number of query elements.')
-
-    out_accum = np.zeros_like(query)
-    if do_backprop:
-      query_ct_accum = np.zeros_like(query)
-      key_ct_accum = np.zeros_like(key)
-      value_ct_accum = np.zeros_like(value)
-      init_vals = (
-          q_loop_idx, out_accum,
-          query_ct_accum, key_ct_accum, value_ct_accum)
-    else:
-      init_vals = (q_loop_idx, out_accum)
-
-    def cond_fun(vals):
-      q_loop_idx = vals[0]
-      return jax.lax.lt(q_loop_idx, q_loop_max)
-
-    def body_fun(vals):
-      """Compute a slice of the attention mechanism."""
-      if do_backprop:
-        (q_loop_idx, out_accum,
-         query_ct_accum, key_ct_accum, value_ct_accum) = vals
-      else:
-        q_loop_idx, out_accum = vals
-
-      query_slice = jax.lax.dynamic_slice_in_dim(
-          query, q_loop_idx, q_loop_stride, axis=-2)
-
-      if do_backprop:
-        ct_slice = jax.lax.dynamic_slice_in_dim(
-            ct, q_loop_idx, q_loop_stride, axis=-2)
-        out_slice, partial_ct = forward_and_vjp_slice(
-            query_slice, q_loop_idx, key, value, ct_slice)
-        query_ct_accum = jax.lax.dynamic_update_slice_in_dim(
-            query_ct_accum, partial_ct[0], q_loop_idx, axis=-2)
-        key_ct_accum = key_ct_accum + partial_ct[1]
-        value_ct_accum = value_ct_accum + partial_ct[2]
-      else:
-        out_slice = forward_slice(query_slice, q_loop_idx, key, value)
-
-      out_accum = jax.lax.dynamic_update_slice_in_dim(
-          out_accum, out_slice, q_loop_idx, axis=-2)
-      q_loop_idx = q_loop_idx + q_loop_stride
-
-      if do_backprop:
-        return (q_loop_idx, out_accum,
-                query_ct_accum, key_ct_accum, value_ct_accum)
-      else:
-        return (q_loop_idx, out_accum)
-
-    final_vals = jax.lax.while_loop(cond_fun, body_fun, init_vals)
-
-    if not do_backprop:
-      return final_vals[1], None
-    else:
-      return final_vals[1], final_vals[2:]
-
-
-class DummyHashedAttention(DotProductAttention):
-  """A stand-in for hash-based attention, but without a real hash function."""
-
-  def __init__(self, dropout, mode, n_bins=64):
-    super(DummyHashedAttention, self).__init__(dropout, mode)
-    self.n_bins = n_bins
-
-  def call(self, inputs, params=(), state=(), **kwargs):
-    output, _ = self.forward_and_vjp(inputs, None, params=params, **kwargs)
-    return output, state
-
-  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
-    del params, kwargs
-    q, k, v = inputs
-    # q/k/v are n_batch*n_heads, seqlen, d_head
-
-    assert k.shape[-2] % self.n_bins == 0
-    bin_size = int(k.shape[-2] // self.n_bins)
-
-    # q_bins/kv_bins are n_batch*n_heads, seqlen
-    # They specify which hash bucket the query/key/value vectors fall in. For
-    # now, instead of hashing we just put consecutive items in the same bucket.
-    q_bins = np.arange(q.shape[-2], dtype=np.int32) // bin_size
-    q_bins = jax.lax.tie_in(q, q_bins)
-    q_bins = q_bins[None, :]
-    q_bins = np.broadcast_to(q_bins, q.shape[:-1])
-    q_bins = -q_bins
-    kv_bins = q_bins * 2
-
-    # q_t/kv_t are n_batch*n_heads, seqlen
-    q_t = jax.lax.tie_in(q, np.arange(q.shape[-2]))
-    q_t = np.reshape(q_t, (1, q_t.shape[0]))
-    q_t = np.broadcast_to(q_t, q.shape[:-1])
-    kv_t = q_t
-
-    def chunk_scalars(x):
-      return np.reshape(x, (x.shape[0], self.n_bins, -1))
-
-    def chunk_vectors(x):
-      return np.reshape(
-          x, (x.shape[0], self.n_bins, -1, x.shape[-1]))
-
-    def unchunk_vectors(x):
-      return np.reshape(x, (x.shape[0], -1, x.shape[-1]))
-
-   # Sort everything by bin number (variables starting with "s" are sorted)
-    _, sq_t = jax.lax.sort_key_val(q_bins, q_t, dimension=-1)
-
-    sq = np.take_along_axis(q, sq_t[:, :, None], axis=-2)
-    if ct is not None:
-      so_ct = np.take_along_axis(ct, sq_t[:, :, None], axis=-2)
-
-    _, skv_t = jax.lax.sort_key_val(kv_bins, kv_t, dimension=-1)
-    sk = np.take_along_axis(k, skv_t[:, :, None], axis=-2)
-    sv = np.take_along_axis(v, skv_t[:, :, None], axis=-2)
-
-    @jax.jit
-    def binned_attn(sq, sk, sv):
-      """Performs attention on sorted queries/keys/values."""
-      # Split off a "bin" axis so that attention only occurs whithin chunks.
-      bq_t = chunk_scalars(sq_t)
-      bkv_t = chunk_scalars(skv_t)
-      bq = chunk_vectors(sq)
-      bk = chunk_vectors(sk)
-      bv = chunk_vectors(sv)
-
-      dots = np.matmul(bq, np.swapaxes(bk, -1, -2)) / np.sqrt(bq.shape[-1])
-
-      # Causal masking
-      mask = jax.lax.convert_element_type(
-          jax.lax.lt(bq_t[:, :, :, None], bkv_t[:, :, None, :]),
-          np.float32)
-      dots = dots - 1e9 * mask
-
-      # Softmax.
-      dots = np.exp(dots - dots.max(axis=-1, keepdims=True))
-      dots = dots / dots.sum(axis=-1, keepdims=True)
-      bo = np.matmul(dots, bv)
-
-      so = unchunk_vectors(bo)
-      return so
-
-    @jax.jit
-    def binned_attn_vjp(sq, sk, sv, so_ct):
-      so, vjpfun = jax.vjp(binned_attn, sq, sk, sv)
-      sqkv_ct = vjpfun(so_ct)
-      return so, sqkv_ct
-
-    if ct is None:
-      so = binned_attn(sq, sk, sv)
-      _, undo_q_sort = jax.lax.sort_key_val(sq_t, q_t, dimension=-1)
-      out = np.take_along_axis(so, undo_q_sort[:, :, None], axis=-2)
-      return out, None
-    else:
-      # Jax can construct a backward pass automatically, but it's about 2x
-      # slower than writing our own. The main reason is that the backward pass
-      # of gather is in general a scatter operation, but we know we're dealing
-      # with permutations so we use gather for the backward pass too.
-      so, (sq_ct, sk_ct, sv_ct) = binned_attn_vjp(sq, sk, sv, so_ct)
-
-      _, undo_q_sort = jax.lax.sort_key_val(sq_t, q_t, dimension=-1)
-      out = np.take_along_axis(so, undo_q_sort[:, :, None], axis=-2)
-      q_ct = np.take_along_axis(sq_ct, undo_q_sort[:, :, None], axis=-2)
-
-      _, undo_kv_sort = jax.lax.sort_key_val(skv_t, kv_t, dimension=-1)
-      k_ct = np.take_along_axis(sk_ct, undo_kv_sort[:, :, None], axis=-2)
-      v_ct = np.take_along_axis(sv_ct, undo_kv_sort[:, :, None], axis=-2)
-
-      return out, (q_ct, k_ct, v_ct)
-
-
 class ReversibleAttentionHalfResidual(tl.ReversibleLayer, tl.Serial):
   """Half of a RevNet-style residual that performs attention.
 
@@ -803,9 +444,9 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
       tl.LayerNorm(),
       tl.Dup(), tl.Dup(),
       tl.Parallel(
-          [ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key)],
-          [ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key)],
-          [ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value)],
+          [tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key)],
+          [tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key)],
+          [tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value)],
       ),
   ]
 
@@ -814,7 +455,7 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
   # ReversibleAttentionHalfResidual requires that post_attention be linear in
   # its input (so the backward pass can be computed without knowing the input)
   post_attention = [
-      ComputeAttentionOutput(n_heads=n_heads, d_model=d_model),
+      tl.ComputeAttentionOutput(n_heads=n_heads, d_model=d_model),
       Unchunk(n_sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
       BroadcastedDropout(rate=dropout, mode=mode),  # pylint: disable=no-value-for-parameter
   ]
@@ -841,7 +482,7 @@ def TransformerRevnetLM(vocab_size,
                         max_len=2048,
                         n_chunks=32,
                         n_attention_chunks=8,
-                        attention_type=DotProductAttention,
+                        attention_type=tl.DotProductCausalAttention,
                         mode='train'):
   """Reversible transformer language model (only uses a decoder, no encoder).
 
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 3d2604dca..af64e32f9 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -128,7 +128,8 @@ def DecoderBlock(d_model, d_ff, n_heads, dropout, mode):
   """
   self_attention = [
       tl.LayerNorm(),  # vec
-      tl.CausalAttention(d_model, n_heads=n_heads, dropout=dropout, mode=mode),
+      tl.BasicCausalAttention(
+          d_model, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Dropout(rate=dropout, mode=mode),  # vec
   ]
   feed_forward = [
@@ -237,7 +238,8 @@ def EncoderDecoder(d_model, d_ff, n_heads, dropout, mode):
   """
   decoder_self_attention = [                    #        vecs_d   pmask vecs_e
       tl.LayerNorm(),                           #        vecs_d   ..... ......
-      tl.CausalAttention(d_model, n_heads=n_heads, dropout=dropout, mode=mode),
+      tl.BasicCausalAttention(
+          d_model, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Dropout(rate=dropout, mode=mode),      # vecs_d          ..... ......
   ]
   decoder_to_encoder_attention = [        # vecs_d        masks         vecs_e

From 57bb1a5852e180a45eda781348eb830084d7f9f3 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 28 Aug 2019 08:16:07 -0700
Subject: [PATCH 2337/2720] Minor redivision of labor between Layer class and
 layer decorator.

PiperOrigin-RevId: 265912586
---
 tensor2tensor/trax/layers/base.py | 32 ++++++++++++++-----------------
 1 file changed, 14 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 8bd41544b..96e607308 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -65,8 +65,9 @@ class Layer(object):
   outputs are spliced back into the stack.
   """
 
-  def __init__(self, **kwargs):
-    self._init_kwargs = kwargs  # can be used in creating a generic decorator
+  def __init__(self, n_inputs=1, n_outputs=1):
+    self._n_inputs = n_inputs
+    self._n_outputs = n_outputs
     self._params = ()  # cached parameters
     self._caller = _find_frame(inspect.stack())  # for custom error messages
     self._init_finished = False
@@ -131,11 +132,11 @@ def new_parameters(self, input_shapes, input_dtype, rng):
 
   def n_inputs(self):
     """Specifies how many data tensors this layer expects as input."""
-    return 1  # Default is one input; subclasses can override.
+    return self._n_inputs
 
   def n_outputs(self):
     """Specifies how many data tensors this layer promises as output."""
-    return 1  # Default is one output: subclasses can override.
+    return self._n_outputs
 
   def sublayers(self):
     """Returns the sublayers contained in / managed by this layer."""
@@ -423,25 +424,21 @@ def _validate_call_input(x, n_inputs):
           ' ({})'.format(len(x), n_inputs))
 
 
-def layer(new_parameters=None, n_inputs=1, n_outputs=1):
+def layer(n_inputs=1, n_outputs=1, new_parameters=None):
   """Decorates a function to make it the call method of a new Layer class."""
   # TODO(jonni): Consider renaming new_parameters to new_parameters_fn.
 
   def _build_layer_class(raw_call_fn):
     """Returns a Layer class built around the given call function."""
 
-    def _n_inputs(self):
-      del self
-      return n_inputs
-
-    def _n_outputs(self):
-      del self
-      return n_outputs
+    def _init(self, **kwargs):
+      self._kwargs = kwargs  # pylint: disable=protected-access
+      Layer.__init__(self, n_inputs=n_inputs, n_outputs=n_outputs)
 
     def _new_parameters(self, input_shapes, input_dtype, rng):
       if new_parameters is None:
         return (), ()
-      kwargs = self._init_kwargs  # pylint: disable=protected-access
+      kwargs = self._kwargs  # pylint: disable=protected-access
       return new_parameters(input_shapes, input_dtype, rng, **kwargs), ()
 
     def _is_empty(raw_output):
@@ -451,7 +448,7 @@ def _is_empty(raw_output):
     def _call_with_context(self, x, params=(), state=(), **kwargs):
       """Calls raw_call_fn with extra keyword args from Layer.__init__."""
       merged_kwargs = kwargs.copy()
-      merged_kwargs.update(self._init_kwargs)  # pylint: disable=protected-access
+      merged_kwargs.update(self._kwargs)  # pylint: disable=protected-access
 
       _validate_call_input(x, n_inputs)
       raw_output = raw_call_fn(x, params=params, **merged_kwargs)
@@ -462,10 +459,9 @@ def _call_with_context(self, x, params=(), state=(), **kwargs):
     _call_with_context.__doc__ = raw_call_fn.__doc__
     _new_parameters.__doc__ = new_parameters.__doc__  # None.__doc__ is None
     cls = type(raw_call_fn.__name__, (Layer,),
-               {'call': _call_with_context,
-                'new_parameters': _new_parameters,
-                'n_inputs': _n_inputs,
-                'n_outputs': _n_outputs})
+               {'__init__': _init,
+                'call': _call_with_context,
+                'new_parameters': _new_parameters})
     return cls
 
   return _build_layer_class

From 19152b03a99d37f32723ada483c6f3a00b29e530 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 28 Aug 2019 16:19:36 -0700
Subject: [PATCH 2338/2720] Tweak test to use different PRNG keys for different
 models.

PiperOrigin-RevId: 266018110
---
 tensor2tensor/trax/layers/core_test.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/trax/layers/core_test.py b/tensor2tensor/trax/layers/core_test.py
index 83356a5af..e2b625807 100644
--- a/tensor2tensor/trax/layers/core_test.py
+++ b/tensor2tensor/trax/layers/core_test.py
@@ -80,9 +80,10 @@ def test_dense_param_sharing(self):
     model1 = combinators.Serial(core.Dense(32), core.Dense(32))
     layer = core.Dense(32)
     model2 = combinators.Serial(layer, layer)
-    rng = backend.random.get_prng(0)
-    params1, _ = model1.initialize((1, 32), onp.float32, rng)
-    params2, _ = model2.initialize((1, 32), onp.float32, rng)
+
+    rng1, rng2 = backend.random.split(backend.random.get_prng(0), 2)
+    params1, _ = model1.initialize((1, 32), onp.float32, rng1)
+    params2, _ = model2.initialize((1, 32), onp.float32, rng2)
     # The first parameters have 2 kernels of size (32, 32).
     self.assertEqual((32, 32), params1[0][0].shape)
     self.assertEqual((32, 32), params1[1][0].shape)

From ed1b9458c773cd0b538aa16336c5d4de569d55f3 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 28 Aug 2019 19:34:25 -0700
Subject: [PATCH 2339/2720] Small corrections to revnet-transformer after
 recent changes.

PiperOrigin-RevId: 266047440
---
 .../trax/configs/transformer_revnet_imagenet64_8gb.gin    | 8 ++++----
 tensor2tensor/trax/layers/attention.py                    | 4 +---
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
index 1ba24f292..25d609052 100644
--- a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
@@ -6,7 +6,7 @@ import tensor2tensor.trax.trax
 # Parameters for batch_fun:
 # ==============================================================================
 batch_fun.batch_size_per_device = 8
-batch_fun.eval_batch_size = 8
+batch_fun.eval_batch_size = 128
 batch_fun.max_eval_length = 12288  # 64 * 64 * 3
 
 # Parameters for inputs:
@@ -24,7 +24,7 @@ MultifactorSchedule.warmup_steps = 8000
 
 # Parameters for train:
 # ==============================================================================
-train.eval_frequency = 100
+train.eval_frequency = 200
 train.eval_steps = 8
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerRevnetLM
@@ -43,8 +43,8 @@ MemoryEfficientCausalAttention.loop_stride = 512
 
 # Parameters for DummyHashedAttention:
 # ==============================================================================
-DummyHashedAttention.dropout = 0.0
-DummyHashedAttention.n_bins = 64
+# DummyHashedAttention.dropout = 0.0
+# DummyHashedAttention.n_bins = 64
 
 # Parameters for TransformerRevnetLM:
 # ==============================================================================
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index e50b9a99f..f1a038f4a 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -407,8 +407,7 @@ def forward_slice(query_slice, q_loop_idx, key, value):  # pylint: disable=inval
       dots = dots - 1e9 * mask
 
       # Softmax.
-      dots = np.exp(dots - dots.max(axis=-1, keepdims=True))
-      dots = dots / dots.sum(axis=-1, keepdims=True)
+      dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))
 
       if self.dropout is not None and self.dropout > 0.0:
         # Dropout is broadcast across the batch+head dimension
@@ -491,4 +490,3 @@ def body_fun(vals):  # pylint: disable=invalid-name
       return final_vals[1], None
     else:
       return final_vals[1], final_vals[2:]
-

From 4e3d675da70e80fe6069be5cb513afcf12a446b6 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 28 Aug 2019 20:59:27 -0700
Subject: [PATCH 2340/2720] Standardize attention between TransformerLM and
 TransformerRevnetLM

PiperOrigin-RevId: 266056388
---
 .../trax/configs/transformer_big_lm1b_8gb.gin |   5 +
 .../trax/configs/transformer_imagenet64.gin   |  51 ++++++
 .../trax/configs/transformer_imdb_8gb.gin     |   5 +
 .../trax/configs/transformer_lm1b_8gb.gin     |   5 +
 .../configs/transformer_lm1b_8gb_testing.gin  |   5 +
 tensor2tensor/trax/layers/attention.py        | 158 +++++++++++++-----
 tensor2tensor/trax/layers/base.py             |   4 +-
 .../models/research/transformer_revnet.py     |  45 +++--
 tensor2tensor/trax/models/transformer.py      |  36 +++-
 9 files changed, 243 insertions(+), 71 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/transformer_imagenet64.gin

diff --git a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
index 63b693f7b..46e28ad13 100644
--- a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
@@ -40,8 +40,13 @@ train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.SM3
 train.train_steps = 500000
 
+# Parameters for DotProductCausalAttention:
+# ==============================================================================
+DotProductCausalAttention.dropout = 0.1
+
 # Parameters for TransformerLM:
 # ==============================================================================
+TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
 TransformerLM.d_model = 1024
 TransformerLM.d_ff = 8192
 TransformerLM.dropout = 0.1
diff --git a/tensor2tensor/trax/configs/transformer_imagenet64.gin b/tensor2tensor/trax/configs/transformer_imagenet64.gin
new file mode 100644
index 000000000..0f56cf59a
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_imagenet64.gin
@@ -0,0 +1,51 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 2
+batch_fun.eval_batch_size = 128
+batch_fun.max_eval_length = 12288  # 64 * 64 * 3
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
+inputs.input_name = 'targets'
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 2.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 200
+train.eval_steps = 8
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerLM
+train.optimizer = @trax.optimizers.Adafactor
+train.train_steps = 500000
+train.trainer_class = @MemoryEfficientTrainer
+
+# Parameters for MemoryEfficientCausalAttention:
+# ==============================================================================
+MemoryEfficientCausalAttention.dropout = 0.0
+MemoryEfficientCausalAttention.loop_stride = 512
+
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.attention_type = @trax.layers.MemoryEfficientCausalAttention
+TransformerLM.d_model = 1024
+TransformerLM.d_ff = 2048
+TransformerLM.dropout = 0.0
+TransformerLM.max_len = 12288  # 64 * 64 * 3
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 4
+TransformerLM.n_layers = 3
+TransformerLM.vocab_size = 256
+
diff --git a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
index 58391c63d..ba44965bb 100644
--- a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
@@ -36,8 +36,13 @@ train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerEncoder
 train.train_steps = 1000
 
+# Parameters for DotProductCausalAttention:
+# ==============================================================================
+DotProductCausalAttention.dropout = 0.1
+
 # Parameters for TransformerLM:
 # ==============================================================================
+TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
 TransformerEncoder.d_model = 512
 TransformerEncoder.d_ff = 2048
 TransformerEncoder.dropout = 0.1
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index 5549858d7..33c32a465 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -40,8 +40,13 @@ train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.SM3
 train.train_steps = 500000
 
+# Parameters for DotProductCausalAttention:
+# ==============================================================================
+DotProductCausalAttention.dropout = 0.1
+
 # Parameters for TransformerLM:
 # ==============================================================================
+TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
 TransformerLM.d_model = 512
 TransformerLM.d_ff = 2048
 TransformerLM.dropout = 0.1
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
index 56f779b52..2f8f91ec6 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
@@ -40,8 +40,13 @@ train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.Adam
 train.train_steps = 100000
 
+# Parameters for DotProductCausalAttention:
+# ==============================================================================
+DotProductCausalAttention.dropout = 0.1
+
 # Parameters for TransformerLM:
 # ==============================================================================
+TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
 TransformerLM.d_model = 512
 TransformerLM.d_ff = 2048
 TransformerLM.dropout = 0.1
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index f1a038f4a..b1648b0bc 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -29,6 +29,16 @@
 from tensor2tensor.trax.layers import initializers as init
 
 
+@base.layer()
+def ShiftRight(x, **unused_kwargs):
+  """Layer to shift the tensor to the right by padding on axis 1."""
+  pad_widths = [(0, 0)] * len(x.shape)
+  pad_widths[1] = (1, 0)  # Padding on axis=1
+  padded = np.pad(x, pad_widths, mode='constant',
+                  constant_values=x.dtype.type(0))
+  return padded[:, :-1]
+
+
 @base.layer()
 def CausalMask(x, params, axis=-1, **kwargs):
   del params, kwargs
@@ -221,16 +231,6 @@ def BasicCausalAttention(d_feature, n_heads=1, dropout=0.0, mode='train'):
   ]
 
 
-@base.layer()
-def ShiftRight(x, **unused_kwargs):
-  """Layer to shift the tensor to the right by padding on axis 1."""
-  pad_widths = [(0, 0)] * len(x.shape)
-  pad_widths[1] = (1, 0)  # Padding on axis=1
-  padded = np.pad(x, pad_widths, mode='constant',
-                  constant_values=x.dtype.type(0))
-  return padded[:, :-1]
-
-
 class ComputeAttentionHeads(base.Layer):
   """Computes queries/keys/values via linear projection.
 
@@ -301,32 +301,46 @@ def new_parameters(self, input_shape, input_dtype, rng):
 
 
 class BaseCausalAttention(base.Layer):
-  """Base class for variants of causal self-attention.
+  """Base class for variants of causal self-attention."""
 
-  This class sets up an API that includes forward_and_vjp, which is required to
-  implement MemoryEfficientCausalAttention.
-  """
+  def __init__(self):
+    super(BaseCausalAttention, self).__init__(n_inputs=3, n_outputs=1)
 
   def call(self, inputs, params=(), state=(), rng=None, **kwargs):
+    """Forward pass for the attention layer."""
     raise NotImplementedError()
 
-  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
+  def call_and_grad(self, inputs, grad, **kwargs):
+    """Performs both forward and backward pass for the attention layer.
+
+    This is used in reversible models: for the backward pass of a reversible
+    model, we need to compute both the forward direction (to recover the
+    previous layer's activations) and the backward direction simultaneously.
+    Some computation can be shared between the forward and backward directions,
+    which makes it more efficient to implement them jointly.
+
+    This method assumes that the layer is stateless and has no parameters.
+
+    Args:
+      inputs: A tuple (q, k, v), where each element has shape
+          n_batch*n_heads, seqlen, d_head
+      grad: gradient signal for the layer output.
+      **kwargs: kwargs for the layer
+
+    Returns:
+      A nested-tuple structure (output, (q_grad, k_grad, v_grad)) that contains
+      the output of the forward pass and the gradient signal for each input.
+    """
     raise NotImplementedError()
 
   def new_parameters(self, input_shapes, input_dtype, rng):
     return (), ()
 
-  def n_inputs(self):
-    return 3
-
-  def n_outputs(self):
-    return 1
-
 
 class DotProductCausalAttention(BaseCausalAttention):
   """A standard (non-memory-efficient) dot product attention implementation."""
 
-  def __init__(self, dropout, mode):
+  def __init__(self, dropout=0.0, mode='train'):
     super(DotProductCausalAttention, self).__init__()
     self._dropout = dropout
     self._mode = mode
@@ -335,24 +349,41 @@ def call(self, inputs, params=(), state=(), rng=None, **kwargs):
     del params
     q, k, v = inputs
     mask_size = q.shape[-2]
-    mask = np.tril(np.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
+    # Not all backends define np.tril. However, using onp.tril is inefficient in
+    # that it creates a large global constant. TODO(kitaev): try to find an
+    # alternative that works across all backends.
+    if backend.get_name() == 'jax':
+      mask = np.tril(np.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
+    else:
+      mask = onp.tril(onp.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
     res = DotProductAttention(
         q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=rng)
     return res, state
 
-  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
+  def call_and_grad(self, inputs, ct, **kwargs):
     assert backend.get_name() == 'jax', (
-        'JAX backend is required to use forward_and_vjp.')
+        'JAX backend is required to use call_and_grad.')
     # Simultaneous forward pass and backprop through the attention mechanism.
     def do_call(x):  # pylint: disable=invalid-name
-      res, _ = self.call(x, params, **kwargs)
+      res, _ = self.call(x, **kwargs)
       return res
     output, vjpfun = jax.vjp(do_call, inputs)
     return output, vjpfun(ct)[0]
 
 
 class MemoryEfficientCausalAttention(BaseCausalAttention):
-  """Memory-efficient dot product attention."""
+  """Memory-efficient dot product attention.
+
+  This layer performs causal attention on long sequences without running out
+  of memory. Instead of computing dot products for all query-key pairs at once,
+  it uses a loop to compute attention for a small set of query positions at a
+  time. The "loop_stride" parameter controls how many query positions are
+  considered at each iteration of the loop.
+
+  Note that this class does not slice along the batch/head dimension. Looping
+  over batch elements and heads instead of query positions is also a viable
+  option. We haven't implemented it, but it may perform well, too.
+  """
 
   def __init__(self, loop_stride, dropout, mode):
     assert backend.get_name() == 'jax', (
@@ -367,27 +398,37 @@ def __init__(self, loop_stride, dropout, mode):
       self.dropout = None
 
   def call(self, inputs, params=(), state=(), **kwargs):
-    output, _ = self.forward_and_vjp(inputs, None, params=params, **kwargs)
+    del params
+    output, _ = self.call_and_grad(inputs, None, **kwargs)
     return output, state
 
-  def forward_and_vjp(self, inputs, ct, params=(), rng=None, **kwargs):
-    # This is the core of the memory-efficient attention implementation, where
-    # we use the jax.lax.while_loop primitive to compute attention for a small
-    # set of query positions at a time. Note how in the backwards pass, we
-    # compute both the forward direction (to recover the previous layer's
-    # activations) and the backward direction simultaneously. This allows us to
-    # only use a single loop, where the inner portion of the loop does a slice
-    # of the forward+backward joint computation. Unfortunately we have had to
-    # introduce a large number of wrapper classes (including
-    # ReversibleAttentionHalfResidual and ApplyAttentionWrapper) for the sole
-    # purpose of connecting this implementation of forward_and_vjp with the core
-    # backprop implementation.
+  def has_custom_grad(self):
+    return True
+
+  def custom_grad(self, inputs, output, ct, params=(), state=(), **kwargs):
+    del output, params, state
+    _, inputs_ct = self.call_and_grad(inputs, ct, **kwargs)
+    return inputs_ct, ()
 
+  def call_and_grad(self, inputs, ct, rng=None, **kwargs):
+    del kwargs
     query, key, value = inputs
     depth = np.shape(query)[-1]
     do_backprop = ct is not None
+    # jax uses the term cotangent (ct) to refer to gradient signals, and
+    # vector-Jacobian product (vjp) for back-propagation through a layer.
 
     def make_mask(N, M, k):  # pylint: disable=invalid-name
+      """Constructs a slice of the causal attention mask.
+
+      Args:
+        N: number of query positions
+        M: number of key positions
+        k: position of the initial query element
+
+      Returns:
+        N x M mask, where 1.0 indicates that attention is not allowed.
+      """
       x = np.arange(N, dtype=np.int32)
       y = np.arange(M, dtype=np.int32)
       mask = jax.lax.lt(
@@ -490,3 +531,40 @@ def body_fun(vals):  # pylint: disable=invalid-name
       return final_vals[1], None
     else:
       return final_vals[1], final_vals[2:]
+
+
+def CausalAttention(d_feature, n_heads=1,
+                    d_attention_key=None, d_attention_value=None,
+                    attention_type=DotProductCausalAttention, mode='train'):
+  """Transformer-style multi-headed causal attention.
+
+  Args:
+    d_feature: int:  dimensionality of feature embedding
+    n_heads: int: number of attention heads
+    d_attention_key: int: depth of key vector for each attention head
+        (default is d_feature // n_heads)
+    d_attention_value: int: depth of value vector for each attention head
+        (default is d_feature // n_heads)
+    attention_type: subclass of tl.BaseCausalAttention: attention class to use
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    Multi-headed self-attention result.
+  """
+  if d_attention_key is None:
+    assert d_feature % n_heads == 0
+    d_attention_key = d_feature // n_heads
+  if d_attention_value is None:
+    assert d_feature % n_heads == 0
+    d_attention_value = d_feature // n_heads
+
+  return [
+      cb.Dup(), cb.Dup(),
+      cb.Parallel(
+          ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
+          ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
+          ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
+      ),
+      attention_type(mode=mode),
+      ComputeAttentionOutput(n_heads=n_heads, d_model=d_feature),
+  ]
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 96e607308..a6eb653fa 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -161,8 +161,8 @@ def custom_grad(self, inputs, output, grad, params, state, **kwargs):
 
     Returns:
       The custom gradient signal for the input. Note that we need to return
-      a gradient for each argument of call, so it will usually be a triple
-      of signals: the gradient for inputs, parameters, and kwargs.
+      a gradient for each argument of call, so it will usually be a tuple
+      of signals: the gradient for inputs and parameters.
     """
     raise NotImplementedError
 
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 62d334e2f..d3c180416 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -278,25 +278,21 @@ def call_compute_residual(x, params):
 
 
 class ApplyAttentionWrapper(tl.Parallel):
-  """Same as tl.Parallel(attention, [], []), but implements forward_and_vjp.
-
-  See MemoryEfficientDotProductAttention for why this is needed.
-  """
+  """Same as tl.Parallel(attention, [], []), but implements call_and_grad."""
 
   def __init__(self, attention):
-    assert hasattr(attention, 'forward_and_vjp')
+    assert hasattr(attention, 'call_and_grad')
     super(ApplyAttentionWrapper, self).__init__(attention, [], [])
     self.attention = attention
 
-  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
+  def call_and_grad(self, inputs, ct, **kwargs):
     # Simultaneous forward pass and backprop through the attention mechanism.
     qkv = inputs[:3]
     passthrough = inputs[3:]
     out_ct = ct[0]
     passthrough_ct = ct[1:]
 
-    out, qkv_ct = self.attention.forward_and_vjp(
-        qkv, out_ct, params=(), **kwargs)
+    out, qkv_ct = self.attention.call_and_grad(qkv, out_ct, **kwargs)
     return (out,) + passthrough, qkv_ct + passthrough_ct
 
 
@@ -306,15 +302,16 @@ class ReversibleAttentionHalfResidual(tl.ReversibleLayer, tl.Serial):
   If inputs are (x1, x2), then outputs are (x1 + z, x2) where:
   z = post_attention(attention(pre_attention(x1)))
 
+  Other than an efficiency optimization, this layer is equivalent to
+  ReversibleHalfResidual([pre_attention, attention, post_attention]).
+
   The post_attention layers must be linear in their input (typically they will
-  consists of reshaping and dense linear layers). This allows back-propagating
-  the gradient signal from the output of ReversibleAttentionHalfResidual to the
-  output of the "attention" portion based only on the network parameters.
-
-  The forward pass is equivalent to using
-  ReversibleHalfResidual([pre_attention, attention, post_attention]), but the
-  backward pass uses attention.forward_and_vjp. See
-  MemoryEfficientDotProductAttention for why forward_and_vjp is helpful.
+  consists of reshaping and dense linear layers), which allows the following
+  optimization. We can back-propagate the gradient signal from the output of
+  ReversibleAttentionHalfResidual to the output of the "attention" portion based
+  only on the network parameters. Then, attention.call_and_grad can be used to
+  recover the output of the "attention" portion while simultaneously performing
+  the backward pass, which allows shared computation between the two directions.
   """
 
   def __init__(self, pre_attention, attention, post_attention):
@@ -324,7 +321,7 @@ def __init__(self, pre_attention, attention, post_attention):
         tl.Swap(),
         tl.Parallel(pre_attention, [], []),
     ])
-    assert hasattr(attention, 'forward_and_vjp')
+    assert hasattr(attention, 'call_and_grad')
     self.attention = ApplyAttentionWrapper(attention)
     self.post_attention = tl.Parallel(post_attention, [], [])
 
@@ -387,9 +384,9 @@ def call_post_attention(x):
     (ct,) = post_attention_vjpfun(ct)
 
     # Simultaneous forward pass and backprop through the attention mechanism
-    stack, ct = self.attention.forward_and_vjp(
-        stack, ct, rng=rngs[1], **kwargs)
-    attention_params_ct = params[1]  # Note: this assumes that params are empty.
+    stack, ct = self.attention.call_and_grad(stack, ct, rng=rngs[1], **kwargs)
+    assert not jax.tree_util.tree_leaves(params[1])
+    attention_params_ct = params[1]  # This is valid when params is empty.
 
     # Backprop through self.pre_attention
     x_ct, pre_attention_params_ct = pre_attention_vjpfun(ct)
@@ -431,7 +428,7 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
     d_attention_value: int: depth of value vector for each attention head
     n_heads: int: number of attention heads
     n_attention_chunks: int: number of chunks for attention
-    attention_type: class: attention class to use, such as DotProductAttention.
+    attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
 
@@ -444,9 +441,9 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
       tl.LayerNorm(),
       tl.Dup(), tl.Dup(),
       tl.Parallel(
-          [tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key)],
-          [tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key)],
-          [tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value)],
+          tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
+          tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
+          tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
       ),
   ]
 
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index af64e32f9..1dc2e0325 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -111,7 +111,8 @@ def TransformerEncoder(vocab_size,
   ])
 
 
-def DecoderBlock(d_model, d_ff, n_heads, dropout, mode):
+def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
+                 attention_type, dropout, mode):
   """Returns a layer sequence that implements a Transformer decoder block.
 
   The input to the layer sequence is an activation tensor.
@@ -120,6 +121,9 @@ def DecoderBlock(d_model, d_ff, n_heads, dropout, mode):
     d_model: int:  depth of embedding
     d_ff: int: depth of feed-forward layer
     n_heads: int: number of attention heads
+    d_attention_key: int: depth of key vector for each attention head
+    d_attention_value: int: depth of value vector for each attention head
+    attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
 
@@ -128,8 +132,10 @@ def DecoderBlock(d_model, d_ff, n_heads, dropout, mode):
   """
   self_attention = [
       tl.LayerNorm(),  # vec
-      tl.BasicCausalAttention(
-          d_model, n_heads=n_heads, dropout=dropout, mode=mode),
+      tl.CausalAttention(
+          d_model, n_heads=n_heads, d_attention_key=d_attention_key,
+          d_attention_value=d_attention_value, attention_type=attention_type,
+          mode=mode),
       tl.Dropout(rate=dropout, mode=mode),  # vec
   ]
   feed_forward = [
@@ -145,6 +151,9 @@ def TransformerDecoder(d_model=512,
                        d_ff=2048,
                        n_layers=6,
                        n_heads=8,
+                       d_attention_key=None,
+                       d_attention_value=None,
+                       attention_type=tl.DotProductCausalAttention,
                        dropout=0.1,
                        max_len=2048,
                        mode='train'):
@@ -159,6 +168,11 @@ def TransformerDecoder(d_model=512,
     d_ff: int: depth of feed-forward layer
     n_layers: int: number of encoder/decoder layers
     n_heads: int: number of attention heads
+    d_attention_key: int: depth of key vector for each attention head
+        (default is d_model // n_heads)
+    d_attention_value: int: depth of value vector for each attention head
+        (default is d_model // n_heads)
+    attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
     max_len: int: maximum symbol length for positional encoding
     mode: str: 'train' or 'eval'
@@ -170,7 +184,9 @@ def TransformerDecoder(d_model=512,
   return tl.Model(                  # vecs
       tl.PositionalEncoding(max_len=max_len),
       tl.Dense(d_model),            # vecs
-      [DecoderBlock(d_model, d_ff, n_heads, dropout, mode)
+      [DecoderBlock(  # pylint: disable=g-complex-comprehension
+          d_model, d_ff, n_heads, d_attention_key, d_attention_value,
+          attention_type, dropout, mode)
        for _ in range(n_layers)],   # vecs
       tl.LayerNorm(),               # vecs
   )
@@ -181,6 +197,9 @@ def TransformerLM(vocab_size,
                   d_ff=2048,
                   n_layers=6,
                   n_heads=8,
+                  d_attention_key=None,
+                  d_attention_value=None,
+                  attention_type=tl.DotProductCausalAttention,
                   dropout=0.1,
                   max_len=2048,
                   mode='train'):
@@ -195,6 +214,11 @@ def TransformerLM(vocab_size,
     d_ff: int: depth of feed-forward layer
     n_layers: int: number of encoder/decoder layers
     n_heads: int: number of attention heads
+    d_attention_key: int: depth of key vector for each attention head
+        (default is d_model // n_heads)
+    d_attention_value: int: depth of value vector for each attention head
+        (default is d_model // n_heads)
+    attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
     max_len: int: maximum symbol length for positional encoding
     mode: str: 'train' or 'eval'
@@ -211,7 +235,9 @@ def TransformerLM(vocab_size,
   return tl.Model(                  # tokens
       tl.ShiftRight(),              # toks
       embedder,                     # vecs
-      [DecoderBlock(d_model, d_ff, n_heads, dropout, mode)
+      [DecoderBlock(  # pylint: disable=g-complex-comprehension
+          d_model, d_ff, n_heads, d_attention_key, d_attention_value,
+          attention_type, dropout, mode)
        for _ in range(n_layers)],   # vecs
       tl.LayerNorm(),               # vecs
       tl.Dense(vocab_size),         # vecs

From b7ff8c4c227d45dfd6f0ae73c60349d532061bc0 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 28 Aug 2019 21:25:23 -0700
Subject: [PATCH 2341/2720] Standardize attention between TransformerLM and
 TransformerRevnetLM

PiperOrigin-RevId: 266059301
---
 .../trax/configs/transformer_big_lm1b_8gb.gin |   5 -
 .../trax/configs/transformer_imagenet64.gin   |  51 ------
 .../trax/configs/transformer_imdb_8gb.gin     |   5 -
 .../trax/configs/transformer_lm1b_8gb.gin     |   5 -
 .../configs/transformer_lm1b_8gb_testing.gin  |   5 -
 tensor2tensor/trax/layers/attention.py        | 158 +++++-------------
 tensor2tensor/trax/layers/base.py             |   4 +-
 .../models/research/transformer_revnet.py     |  45 ++---
 tensor2tensor/trax/models/transformer.py      |  36 +---
 9 files changed, 71 insertions(+), 243 deletions(-)
 delete mode 100644 tensor2tensor/trax/configs/transformer_imagenet64.gin

diff --git a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
index 46e28ad13..63b693f7b 100644
--- a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
@@ -40,13 +40,8 @@ train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.SM3
 train.train_steps = 500000
 
-# Parameters for DotProductCausalAttention:
-# ==============================================================================
-DotProductCausalAttention.dropout = 0.1
-
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
 TransformerLM.d_model = 1024
 TransformerLM.d_ff = 8192
 TransformerLM.dropout = 0.1
diff --git a/tensor2tensor/trax/configs/transformer_imagenet64.gin b/tensor2tensor/trax/configs/transformer_imagenet64.gin
deleted file mode 100644
index 0f56cf59a..000000000
--- a/tensor2tensor/trax/configs/transformer_imagenet64.gin
+++ /dev/null
@@ -1,51 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fun:
-# ==============================================================================
-batch_fun.batch_size_per_device = 2
-batch_fun.eval_batch_size = 128
-batch_fun.max_eval_length = 12288  # 64 * 64 * 3
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
-inputs.input_name = 'targets'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 2.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 200
-train.eval_steps = 8
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerLM
-train.optimizer = @trax.optimizers.Adafactor
-train.train_steps = 500000
-train.trainer_class = @MemoryEfficientTrainer
-
-# Parameters for MemoryEfficientCausalAttention:
-# ==============================================================================
-MemoryEfficientCausalAttention.dropout = 0.0
-MemoryEfficientCausalAttention.loop_stride = 512
-
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = @trax.layers.MemoryEfficientCausalAttention
-TransformerLM.d_model = 1024
-TransformerLM.d_ff = 2048
-TransformerLM.dropout = 0.0
-TransformerLM.max_len = 12288  # 64 * 64 * 3
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 4
-TransformerLM.n_layers = 3
-TransformerLM.vocab_size = 256
-
diff --git a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
index ba44965bb..58391c63d 100644
--- a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
@@ -36,13 +36,8 @@ train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerEncoder
 train.train_steps = 1000
 
-# Parameters for DotProductCausalAttention:
-# ==============================================================================
-DotProductCausalAttention.dropout = 0.1
-
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
 TransformerEncoder.d_model = 512
 TransformerEncoder.d_ff = 2048
 TransformerEncoder.dropout = 0.1
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index 33c32a465..5549858d7 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -40,13 +40,8 @@ train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.SM3
 train.train_steps = 500000
 
-# Parameters for DotProductCausalAttention:
-# ==============================================================================
-DotProductCausalAttention.dropout = 0.1
-
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
 TransformerLM.d_model = 512
 TransformerLM.d_ff = 2048
 TransformerLM.dropout = 0.1
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
index 2f8f91ec6..56f779b52 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
@@ -40,13 +40,8 @@ train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.Adam
 train.train_steps = 100000
 
-# Parameters for DotProductCausalAttention:
-# ==============================================================================
-DotProductCausalAttention.dropout = 0.1
-
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
 TransformerLM.d_model = 512
 TransformerLM.d_ff = 2048
 TransformerLM.dropout = 0.1
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index b1648b0bc..f1a038f4a 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -29,16 +29,6 @@
 from tensor2tensor.trax.layers import initializers as init
 
 
-@base.layer()
-def ShiftRight(x, **unused_kwargs):
-  """Layer to shift the tensor to the right by padding on axis 1."""
-  pad_widths = [(0, 0)] * len(x.shape)
-  pad_widths[1] = (1, 0)  # Padding on axis=1
-  padded = np.pad(x, pad_widths, mode='constant',
-                  constant_values=x.dtype.type(0))
-  return padded[:, :-1]
-
-
 @base.layer()
 def CausalMask(x, params, axis=-1, **kwargs):
   del params, kwargs
@@ -231,6 +221,16 @@ def BasicCausalAttention(d_feature, n_heads=1, dropout=0.0, mode='train'):
   ]
 
 
+@base.layer()
+def ShiftRight(x, **unused_kwargs):
+  """Layer to shift the tensor to the right by padding on axis 1."""
+  pad_widths = [(0, 0)] * len(x.shape)
+  pad_widths[1] = (1, 0)  # Padding on axis=1
+  padded = np.pad(x, pad_widths, mode='constant',
+                  constant_values=x.dtype.type(0))
+  return padded[:, :-1]
+
+
 class ComputeAttentionHeads(base.Layer):
   """Computes queries/keys/values via linear projection.
 
@@ -301,46 +301,32 @@ def new_parameters(self, input_shape, input_dtype, rng):
 
 
 class BaseCausalAttention(base.Layer):
-  """Base class for variants of causal self-attention."""
+  """Base class for variants of causal self-attention.
 
-  def __init__(self):
-    super(BaseCausalAttention, self).__init__(n_inputs=3, n_outputs=1)
+  This class sets up an API that includes forward_and_vjp, which is required to
+  implement MemoryEfficientCausalAttention.
+  """
 
   def call(self, inputs, params=(), state=(), rng=None, **kwargs):
-    """Forward pass for the attention layer."""
     raise NotImplementedError()
 
-  def call_and_grad(self, inputs, grad, **kwargs):
-    """Performs both forward and backward pass for the attention layer.
-
-    This is used in reversible models: for the backward pass of a reversible
-    model, we need to compute both the forward direction (to recover the
-    previous layer's activations) and the backward direction simultaneously.
-    Some computation can be shared between the forward and backward directions,
-    which makes it more efficient to implement them jointly.
-
-    This method assumes that the layer is stateless and has no parameters.
-
-    Args:
-      inputs: A tuple (q, k, v), where each element has shape
-          n_batch*n_heads, seqlen, d_head
-      grad: gradient signal for the layer output.
-      **kwargs: kwargs for the layer
-
-    Returns:
-      A nested-tuple structure (output, (q_grad, k_grad, v_grad)) that contains
-      the output of the forward pass and the gradient signal for each input.
-    """
+  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
     raise NotImplementedError()
 
   def new_parameters(self, input_shapes, input_dtype, rng):
     return (), ()
 
+  def n_inputs(self):
+    return 3
+
+  def n_outputs(self):
+    return 1
+
 
 class DotProductCausalAttention(BaseCausalAttention):
   """A standard (non-memory-efficient) dot product attention implementation."""
 
-  def __init__(self, dropout=0.0, mode='train'):
+  def __init__(self, dropout, mode):
     super(DotProductCausalAttention, self).__init__()
     self._dropout = dropout
     self._mode = mode
@@ -349,41 +335,24 @@ def call(self, inputs, params=(), state=(), rng=None, **kwargs):
     del params
     q, k, v = inputs
     mask_size = q.shape[-2]
-    # Not all backends define np.tril. However, using onp.tril is inefficient in
-    # that it creates a large global constant. TODO(kitaev): try to find an
-    # alternative that works across all backends.
-    if backend.get_name() == 'jax':
-      mask = np.tril(np.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
-    else:
-      mask = onp.tril(onp.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
+    mask = np.tril(np.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
     res = DotProductAttention(
         q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=rng)
     return res, state
 
-  def call_and_grad(self, inputs, ct, **kwargs):
+  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
     assert backend.get_name() == 'jax', (
-        'JAX backend is required to use call_and_grad.')
+        'JAX backend is required to use forward_and_vjp.')
     # Simultaneous forward pass and backprop through the attention mechanism.
     def do_call(x):  # pylint: disable=invalid-name
-      res, _ = self.call(x, **kwargs)
+      res, _ = self.call(x, params, **kwargs)
       return res
     output, vjpfun = jax.vjp(do_call, inputs)
     return output, vjpfun(ct)[0]
 
 
 class MemoryEfficientCausalAttention(BaseCausalAttention):
-  """Memory-efficient dot product attention.
-
-  This layer performs causal attention on long sequences without running out
-  of memory. Instead of computing dot products for all query-key pairs at once,
-  it uses a loop to compute attention for a small set of query positions at a
-  time. The "loop_stride" parameter controls how many query positions are
-  considered at each iteration of the loop.
-
-  Note that this class does not slice along the batch/head dimension. Looping
-  over batch elements and heads instead of query positions is also a viable
-  option. We haven't implemented it, but it may perform well, too.
-  """
+  """Memory-efficient dot product attention."""
 
   def __init__(self, loop_stride, dropout, mode):
     assert backend.get_name() == 'jax', (
@@ -398,37 +367,27 @@ def __init__(self, loop_stride, dropout, mode):
       self.dropout = None
 
   def call(self, inputs, params=(), state=(), **kwargs):
-    del params
-    output, _ = self.call_and_grad(inputs, None, **kwargs)
+    output, _ = self.forward_and_vjp(inputs, None, params=params, **kwargs)
     return output, state
 
-  def has_custom_grad(self):
-    return True
-
-  def custom_grad(self, inputs, output, ct, params=(), state=(), **kwargs):
-    del output, params, state
-    _, inputs_ct = self.call_and_grad(inputs, ct, **kwargs)
-    return inputs_ct, ()
+  def forward_and_vjp(self, inputs, ct, params=(), rng=None, **kwargs):
+    # This is the core of the memory-efficient attention implementation, where
+    # we use the jax.lax.while_loop primitive to compute attention for a small
+    # set of query positions at a time. Note how in the backwards pass, we
+    # compute both the forward direction (to recover the previous layer's
+    # activations) and the backward direction simultaneously. This allows us to
+    # only use a single loop, where the inner portion of the loop does a slice
+    # of the forward+backward joint computation. Unfortunately we have had to
+    # introduce a large number of wrapper classes (including
+    # ReversibleAttentionHalfResidual and ApplyAttentionWrapper) for the sole
+    # purpose of connecting this implementation of forward_and_vjp with the core
+    # backprop implementation.
 
-  def call_and_grad(self, inputs, ct, rng=None, **kwargs):
-    del kwargs
     query, key, value = inputs
     depth = np.shape(query)[-1]
     do_backprop = ct is not None
-    # jax uses the term cotangent (ct) to refer to gradient signals, and
-    # vector-Jacobian product (vjp) for back-propagation through a layer.
 
     def make_mask(N, M, k):  # pylint: disable=invalid-name
-      """Constructs a slice of the causal attention mask.
-
-      Args:
-        N: number of query positions
-        M: number of key positions
-        k: position of the initial query element
-
-      Returns:
-        N x M mask, where 1.0 indicates that attention is not allowed.
-      """
       x = np.arange(N, dtype=np.int32)
       y = np.arange(M, dtype=np.int32)
       mask = jax.lax.lt(
@@ -531,40 +490,3 @@ def body_fun(vals):  # pylint: disable=invalid-name
       return final_vals[1], None
     else:
       return final_vals[1], final_vals[2:]
-
-
-def CausalAttention(d_feature, n_heads=1,
-                    d_attention_key=None, d_attention_value=None,
-                    attention_type=DotProductCausalAttention, mode='train'):
-  """Transformer-style multi-headed causal attention.
-
-  Args:
-    d_feature: int:  dimensionality of feature embedding
-    n_heads: int: number of attention heads
-    d_attention_key: int: depth of key vector for each attention head
-        (default is d_feature // n_heads)
-    d_attention_value: int: depth of value vector for each attention head
-        (default is d_feature // n_heads)
-    attention_type: subclass of tl.BaseCausalAttention: attention class to use
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    Multi-headed self-attention result.
-  """
-  if d_attention_key is None:
-    assert d_feature % n_heads == 0
-    d_attention_key = d_feature // n_heads
-  if d_attention_value is None:
-    assert d_feature % n_heads == 0
-    d_attention_value = d_feature // n_heads
-
-  return [
-      cb.Dup(), cb.Dup(),
-      cb.Parallel(
-          ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
-          ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
-          ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
-      ),
-      attention_type(mode=mode),
-      ComputeAttentionOutput(n_heads=n_heads, d_model=d_feature),
-  ]
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index a6eb653fa..96e607308 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -161,8 +161,8 @@ def custom_grad(self, inputs, output, grad, params, state, **kwargs):
 
     Returns:
       The custom gradient signal for the input. Note that we need to return
-      a gradient for each argument of call, so it will usually be a tuple
-      of signals: the gradient for inputs and parameters.
+      a gradient for each argument of call, so it will usually be a triple
+      of signals: the gradient for inputs, parameters, and kwargs.
     """
     raise NotImplementedError
 
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index d3c180416..62d334e2f 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -278,21 +278,25 @@ def call_compute_residual(x, params):
 
 
 class ApplyAttentionWrapper(tl.Parallel):
-  """Same as tl.Parallel(attention, [], []), but implements call_and_grad."""
+  """Same as tl.Parallel(attention, [], []), but implements forward_and_vjp.
+
+  See MemoryEfficientDotProductAttention for why this is needed.
+  """
 
   def __init__(self, attention):
-    assert hasattr(attention, 'call_and_grad')
+    assert hasattr(attention, 'forward_and_vjp')
     super(ApplyAttentionWrapper, self).__init__(attention, [], [])
     self.attention = attention
 
-  def call_and_grad(self, inputs, ct, **kwargs):
+  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
     # Simultaneous forward pass and backprop through the attention mechanism.
     qkv = inputs[:3]
     passthrough = inputs[3:]
     out_ct = ct[0]
     passthrough_ct = ct[1:]
 
-    out, qkv_ct = self.attention.call_and_grad(qkv, out_ct, **kwargs)
+    out, qkv_ct = self.attention.forward_and_vjp(
+        qkv, out_ct, params=(), **kwargs)
     return (out,) + passthrough, qkv_ct + passthrough_ct
 
 
@@ -302,16 +306,15 @@ class ReversibleAttentionHalfResidual(tl.ReversibleLayer, tl.Serial):
   If inputs are (x1, x2), then outputs are (x1 + z, x2) where:
   z = post_attention(attention(pre_attention(x1)))
 
-  Other than an efficiency optimization, this layer is equivalent to
-  ReversibleHalfResidual([pre_attention, attention, post_attention]).
-
   The post_attention layers must be linear in their input (typically they will
-  consists of reshaping and dense linear layers), which allows the following
-  optimization. We can back-propagate the gradient signal from the output of
-  ReversibleAttentionHalfResidual to the output of the "attention" portion based
-  only on the network parameters. Then, attention.call_and_grad can be used to
-  recover the output of the "attention" portion while simultaneously performing
-  the backward pass, which allows shared computation between the two directions.
+  consists of reshaping and dense linear layers). This allows back-propagating
+  the gradient signal from the output of ReversibleAttentionHalfResidual to the
+  output of the "attention" portion based only on the network parameters.
+
+  The forward pass is equivalent to using
+  ReversibleHalfResidual([pre_attention, attention, post_attention]), but the
+  backward pass uses attention.forward_and_vjp. See
+  MemoryEfficientDotProductAttention for why forward_and_vjp is helpful.
   """
 
   def __init__(self, pre_attention, attention, post_attention):
@@ -321,7 +324,7 @@ def __init__(self, pre_attention, attention, post_attention):
         tl.Swap(),
         tl.Parallel(pre_attention, [], []),
     ])
-    assert hasattr(attention, 'call_and_grad')
+    assert hasattr(attention, 'forward_and_vjp')
     self.attention = ApplyAttentionWrapper(attention)
     self.post_attention = tl.Parallel(post_attention, [], [])
 
@@ -384,9 +387,9 @@ def call_post_attention(x):
     (ct,) = post_attention_vjpfun(ct)
 
     # Simultaneous forward pass and backprop through the attention mechanism
-    stack, ct = self.attention.call_and_grad(stack, ct, rng=rngs[1], **kwargs)
-    assert not jax.tree_util.tree_leaves(params[1])
-    attention_params_ct = params[1]  # This is valid when params is empty.
+    stack, ct = self.attention.forward_and_vjp(
+        stack, ct, rng=rngs[1], **kwargs)
+    attention_params_ct = params[1]  # Note: this assumes that params are empty.
 
     # Backprop through self.pre_attention
     x_ct, pre_attention_params_ct = pre_attention_vjpfun(ct)
@@ -428,7 +431,7 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
     d_attention_value: int: depth of value vector for each attention head
     n_heads: int: number of attention heads
     n_attention_chunks: int: number of chunks for attention
-    attention_type: subclass of tl.BaseCausalAttention: attention class to use
+    attention_type: class: attention class to use, such as DotProductAttention.
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
 
@@ -441,9 +444,9 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
       tl.LayerNorm(),
       tl.Dup(), tl.Dup(),
       tl.Parallel(
-          tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
-          tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
-          tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
+          [tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key)],
+          [tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key)],
+          [tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value)],
       ),
   ]
 
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 1dc2e0325..af64e32f9 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -111,8 +111,7 @@ def TransformerEncoder(vocab_size,
   ])
 
 
-def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-                 attention_type, dropout, mode):
+def DecoderBlock(d_model, d_ff, n_heads, dropout, mode):
   """Returns a layer sequence that implements a Transformer decoder block.
 
   The input to the layer sequence is an activation tensor.
@@ -121,9 +120,6 @@ def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
     d_model: int:  depth of embedding
     d_ff: int: depth of feed-forward layer
     n_heads: int: number of attention heads
-    d_attention_key: int: depth of key vector for each attention head
-    d_attention_value: int: depth of value vector for each attention head
-    attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
 
@@ -132,10 +128,8 @@ def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
   """
   self_attention = [
       tl.LayerNorm(),  # vec
-      tl.CausalAttention(
-          d_model, n_heads=n_heads, d_attention_key=d_attention_key,
-          d_attention_value=d_attention_value, attention_type=attention_type,
-          mode=mode),
+      tl.BasicCausalAttention(
+          d_model, n_heads=n_heads, dropout=dropout, mode=mode),
       tl.Dropout(rate=dropout, mode=mode),  # vec
   ]
   feed_forward = [
@@ -151,9 +145,6 @@ def TransformerDecoder(d_model=512,
                        d_ff=2048,
                        n_layers=6,
                        n_heads=8,
-                       d_attention_key=None,
-                       d_attention_value=None,
-                       attention_type=tl.DotProductCausalAttention,
                        dropout=0.1,
                        max_len=2048,
                        mode='train'):
@@ -168,11 +159,6 @@ def TransformerDecoder(d_model=512,
     d_ff: int: depth of feed-forward layer
     n_layers: int: number of encoder/decoder layers
     n_heads: int: number of attention heads
-    d_attention_key: int: depth of key vector for each attention head
-        (default is d_model // n_heads)
-    d_attention_value: int: depth of value vector for each attention head
-        (default is d_model // n_heads)
-    attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
     max_len: int: maximum symbol length for positional encoding
     mode: str: 'train' or 'eval'
@@ -184,9 +170,7 @@ def TransformerDecoder(d_model=512,
   return tl.Model(                  # vecs
       tl.PositionalEncoding(max_len=max_len),
       tl.Dense(d_model),            # vecs
-      [DecoderBlock(  # pylint: disable=g-complex-comprehension
-          d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-          attention_type, dropout, mode)
+      [DecoderBlock(d_model, d_ff, n_heads, dropout, mode)
        for _ in range(n_layers)],   # vecs
       tl.LayerNorm(),               # vecs
   )
@@ -197,9 +181,6 @@ def TransformerLM(vocab_size,
                   d_ff=2048,
                   n_layers=6,
                   n_heads=8,
-                  d_attention_key=None,
-                  d_attention_value=None,
-                  attention_type=tl.DotProductCausalAttention,
                   dropout=0.1,
                   max_len=2048,
                   mode='train'):
@@ -214,11 +195,6 @@ def TransformerLM(vocab_size,
     d_ff: int: depth of feed-forward layer
     n_layers: int: number of encoder/decoder layers
     n_heads: int: number of attention heads
-    d_attention_key: int: depth of key vector for each attention head
-        (default is d_model // n_heads)
-    d_attention_value: int: depth of value vector for each attention head
-        (default is d_model // n_heads)
-    attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
     max_len: int: maximum symbol length for positional encoding
     mode: str: 'train' or 'eval'
@@ -235,9 +211,7 @@ def TransformerLM(vocab_size,
   return tl.Model(                  # tokens
       tl.ShiftRight(),              # toks
       embedder,                     # vecs
-      [DecoderBlock(  # pylint: disable=g-complex-comprehension
-          d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-          attention_type, dropout, mode)
+      [DecoderBlock(d_model, d_ff, n_heads, dropout, mode)
        for _ in range(n_layers)],   # vecs
       tl.LayerNorm(),               # vecs
       tl.Dense(vocab_size),         # vecs

From 7f2acbf98f7e44b3646bd3f4f6a33823a9e675c9 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 29 Aug 2019 11:21:09 -0700
Subject: [PATCH 2342/2720] Ensure modality is swapped in reversed problems.

PiperOrigin-RevId: 266185066
---
 tensor2tensor/data_generators/problem.py      | 15 ++++++++++-----
 tensor2tensor/data_generators/problem_test.py | 19 +++++++++++++++++++
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 2133d9356..68745257c 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -969,12 +969,17 @@ def _reverse_problem_hparams(p_hparams):
   # In the future, remove need for this behavior.
   reversed_modality = {}
   for feature_name in p.modality:
-    reversed_feature_name = feature_name.replace("target", "input")
-    if "target" in feature_name and reversed_feature_name in p.modality:
-      reversed_modality[feature_name] = p.modality[reversed_feature_name]
-      reversed_modality[reversed_feature_name] = p.modality[feature_name]
-    else:
+    # Copy feature as-is.
+    if "target" not in feature_name and "input" not in feature_name:
       reversed_modality[feature_name] = p.modality[feature_name]
+    else:
+      # Change "target" to "input" and vice-versa for this feature.
+      if "target" in feature_name:
+        reversed_feature_name = feature_name.replace("target", "input")
+      else:
+        assert "input" in feature_name, feature_name
+        reversed_feature_name = feature_name.replace("input", "target")
+      reversed_modality[reversed_feature_name] = p.modality[feature_name]
 
   p.modality = reversed_modality
 
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 63ee900ff..51aeb4f34 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -211,5 +211,24 @@ def hparams(self, defaults, model_hparams):
     self.assertEqual(p_hparams.vocab_size["inputs"], 3)
     self.assertEqual(p_hparams.vocab_size["targets"], 1)
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testInputAndTargetModalitiesAreReversed(self):
+
+    class WasReversedTestProblem(problem_module.Problem):
+
+      def __init__(self, was_reversed):
+        super(WasReversedTestProblem, self).__init__(was_reversed, False)
+
+      def hparams(self, defaults, model_hparams):
+        hp = defaults
+        hp.modality["inputs"] = "inputs_modality"
+        hp.modality["targets"] = "targets_modality"
+
+    problem = WasReversedTestProblem(was_reversed=True)
+    p_hparams = problem.get_hparams()
+    self.assertEqual(p_hparams.modality["inputs"], "targets_modality")
+    self.assertEqual(p_hparams.modality["targets"], "inputs_modality")
+
+
 if __name__ == "__main__":
   tf.test.main()

From 2db3e775cdd2c14812810de7e6b2f0a79b7f8a7e Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 29 Aug 2019 13:16:27 -0700
Subject: [PATCH 2343/2720] Add comments to the SimPLe loop

PiperOrigin-RevId: 266212175
---
 tensor2tensor/trax/rl/simple_trainer.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index 4d9f7a34d..66e7131bc 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -93,9 +93,13 @@ def epoch(self):
     return self._simple_epoch
 
   def train_epoch(self):
+    # Collect trajectories by running the policy in the real environment.
     self.collect_trajectories()
+    # Train the model of the environment on the collected trajectories.
     self.train_model()
+    # Train the policy inside the simulated environment generated by the model.
     self.train_policy()
+
     self._simple_epoch += 1
 
   def evaluate(self):

From 512da60974e491247debe94e6080b08991977316 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 29 Aug 2019 14:55:19 -0700
Subject: [PATCH 2344/2720] Add algorithm name to SimPLe and PPO logs

PiperOrigin-RevId: 266235923
---
 tensor2tensor/trax/rl/ppo_trainer.py    | 12 +++++++-----
 tensor2tensor/trax/rl/simple_trainer.py |  6 +++---
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index 9f8e0c95f..8889e992d 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -221,7 +221,7 @@ def train_epoch(self):
     policy_eval_time = ppo.get_time(policy_eval_start_time)
 
     trajectory_collection_start_time = time.time()
-    logging.vlog(1, "Epoch [% 6d] collecting trajectories.", self._epoch)
+    logging.vlog(1, "PPO epoch [% 6d]: collecting trajectories.", self._epoch)
     self._rng, key = jax_random.split(self._rng)
     trajs, n_done, timing_info, self._model_state = ppo.collect_trajectories(
         self.train_env,
@@ -452,7 +452,7 @@ def train_epoch(self):
     epoch_time = ppo.get_time(epoch_start_time)
 
     logging.info(
-        "Epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
+        "PPO epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
         " Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)]", self._epoch,
         min_reward, max_reward, avg_reward, loss_combined, loss_value, loss_ppo,
         entropy_bonus)
@@ -479,7 +479,9 @@ def train_epoch(self):
         for k, v in sorted(timing_dict.items())
     ]
     logging.info(
-        "Epoch [% 6d], Timings: \n%s", self._epoch, "\n".join(timing_info_list))
+        "PPO epoch [% 6d], Timings: \n%s", self._epoch,
+        "\n".join(timing_info_list)
+    )
 
     self._epoch += 1
 
@@ -489,7 +491,7 @@ def train_epoch(self):
 
   def evaluate(self):
     """Evaluate the agent."""
-    logging.vlog(1, "Epoch [% 6d] evaluating policy.", self._epoch)
+    logging.vlog(1, "PPO epoch [% 6d]: evaluating policy.", self._epoch)
     self._rng, key = jax_random.split(self._rng, num=2)
     reward_stats, self._model_state = ppo.evaluate_policy(
         self.eval_env,
@@ -505,7 +507,7 @@ def evaluate(self):
 
   def save(self):
     """Save the agent parameters."""
-    logging.vlog(1, "Epoch [% 6d] saving model.", self._epoch)
+    logging.vlog(1, "PPO epoch [% 6d]: saving model.", self._epoch)
     old_model_files = gfile.glob(
         os.path.join(self._output_dir, "model-??????.pkl"))
     params_file = os.path.join(self._output_dir, "model-%06d.pkl" % self._epoch)
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index 66e7131bc..f20060bea 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -114,7 +114,7 @@ def flush_summaries(self):
     pass
 
   def collect_trajectories(self):
-    logging.info("Epoch %d: collecting data", self._simple_epoch)
+    logging.info("SimPLe epoch [% 6d]: collecting data.", self._simple_epoch)
 
     self._policy_trainer.train_env = self.train_env
     self._policy_trainer.trajectory_dump_dir = os.path.join(
@@ -170,7 +170,7 @@ def make_batch(examples):
         yield make_batch(example_list)
 
   def train_model(self):
-    logging.info("Epoch %d: training model", self._simple_epoch)
+    logging.info("SimPLe epoch [% 6d]: training model.", self._simple_epoch)
 
     # Load data from all epochs.
     # TODO(pkozakowski): Handle the case when the data won't fit in the memory.
@@ -199,7 +199,7 @@ def train_model(self):
     )
 
   def train_policy(self):
-    logging.info("Epoch %d: training policy", self._simple_epoch)
+    logging.info("SimPLe epoch [% 6d]: training policy.", self._simple_epoch)
 
     self._sim_env.initialize(
         batch_size=self._simulated_batch_size,

From a67313d2e071e2a7f3efa98dcffd0649ad3af9ca Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 29 Aug 2019 15:39:39 -0700
Subject: [PATCH 2345/2720] Expose core aspects of layers consistently as
 (read-only) properties.

PiperOrigin-RevId: 266245575
---
 tensor2tensor/trax/layers/attention.py        |  9 +--
 tensor2tensor/trax/layers/base.py             | 19 ++++--
 tensor2tensor/trax/layers/combinators.py      | 67 +++++++------------
 .../models/research/transformer_revnet.py     | 30 +--------
 4 files changed, 41 insertions(+), 84 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index f1a038f4a..19a88e30f 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -307,6 +307,9 @@ class BaseCausalAttention(base.Layer):
   implement MemoryEfficientCausalAttention.
   """
 
+  def __init__(self):
+    super(BaseCausalAttention, self).__init__(n_inputs=3)
+
   def call(self, inputs, params=(), state=(), rng=None, **kwargs):
     raise NotImplementedError()
 
@@ -316,12 +319,6 @@ def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
   def new_parameters(self, input_shapes, input_dtype, rng):
     return (), ()
 
-  def n_inputs(self):
-    return 3
-
-  def n_outputs(self):
-    return 1
-
 
 class DotProductCausalAttention(BaseCausalAttention):
   """A standard (non-memory-efficient) dot product attention implementation."""
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 96e607308..ec680b9e9 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -35,11 +35,12 @@ class Layer(object):
 
   A layer is a function from zero or more inputs to zero or more outputs,
   possibly with trainable parameters. A layer is either atomic or composed
-  of sublayers. All layers provide accessor methods for these aspects:
+  of sublayers. These aspects of a layer are set via a layer's constructor,
+  and can be inspected via read-only properties:
 
-    - n_inputs()
-    - n_outputs()
-    - sublayers()
+    - n_inputs
+    - n_outputs
+    - sublayers
 
   The inputs to a layer are activation tensors, packaged according to how many
   there are:
@@ -68,14 +69,15 @@ class Layer(object):
   def __init__(self, n_inputs=1, n_outputs=1):
     self._n_inputs = n_inputs
     self._n_outputs = n_outputs
+    self._sublayers = ()  # Default is no sublayers.
     self._params = ()  # cached parameters
     self._caller = _find_frame(inspect.stack())  # for custom error messages
     self._init_finished = False
 
   def __repr__(self):
     class_str = self.__class__.__name__
-    fields_str = 'in={},out={}'.format(self.n_inputs(), self.n_outputs())
-    objs = self.sublayers()
+    fields_str = 'in={},out={}'.format(self.n_inputs, self.n_outputs)
+    objs = self.sublayers
     if objs:
       objs_str = ', '.join(str(x) for x in objs)
       return '{}[{},layers=[{}]]'.format(class_str, fields_str, objs_str)
@@ -130,17 +132,20 @@ def new_parameters(self, input_shapes, input_dtype, rng):
     """
     raise NotImplementedError
 
+  @property
   def n_inputs(self):
     """Specifies how many data tensors this layer expects as input."""
     return self._n_inputs
 
+  @property
   def n_outputs(self):
     """Specifies how many data tensors this layer promises as output."""
     return self._n_outputs
 
+  @property
   def sublayers(self):
     """Returns the sublayers contained in / managed by this layer."""
-    return ()  # Default is no sublayers; subclasses can override.
+    return self._sublayers
 
   @property
   def has_custom_grad(self):
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 800f48a7e..d754c016e 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -106,10 +106,10 @@ class Serial(base.Layer):
   sublayer k, following sublayer j, gets called with the data stack in the
   state left after layer j has applied. The Serial combinator then:
 
-    - takes N_in items off the top of the stack (N_in = k.n_inputs()) and calls
+    - takes N_in items off the top of the stack (N_in = k.n_inputs) and calls
       layer k, passing those items as arguments; and
 
-    - takes layer k's N_out return values (N_out = k.n_outputs()) and pushes
+    - takes layer k's N_out return values (N_out = k.n_outputs) and pushes
       them onto the data stack.
   """
 
@@ -144,34 +144,25 @@ def _n_inputs_n_outputs(self, layers):
     running_max = 0
     running_total = 0
     for layer in layers:
-      running_total += layer.n_inputs()
+      running_total += layer.n_inputs
       running_max = max(running_max, running_total)
-      running_total -= layer.n_outputs()
+      running_total -= layer.n_outputs
     return running_max, (running_max - running_total)
 
-  def n_inputs(self):
-    return self._n_inputs
-
-  def n_outputs(self):
-    return self._n_outputs
-
-  def sublayers(self):
-    return self._sublayers
-
   def _validate_call_inputs(self, xs):
     if not isinstance(xs, tuple) and self._n_inputs != 1:
       raise TypeError(
           'Serial.call input must be a tuple; instead got {}'.format(xs))
     len_xs = 1 if isinstance(xs, np.ndarray) else len(xs)
-    if len_xs < self.n_inputs():
+    if len_xs < self.n_inputs:
       raise ValueError(
           'number of inputs ({}) to Serial.call less than n_inputs'
-          ' ({})'.format(len(xs), self.n_inputs()))
+          ' ({})'.format(len(xs), self.n_inputs))
 
   def call(self, xs, params=(), state=(), **kwargs):
     self._validate_call_inputs(xs)
     rngs = _pop_rng_and_split(kwargs, self._n_layers)
-    if not self._sublayers:  # No-op: leave args unchanged.
+    if not self.sublayers:  # No-op: leave args unchanged.
       return (xs, state)
 
     stack = xs
@@ -183,11 +174,11 @@ def call(self, xs, params=(), state=(), **kwargs):
     if n_layers != 1 and len(state) != n_layers:
       raise ValueError('length of state ({}) not equal to number of layers '
                        '({})'.format(len(state), n_layers))
-    for layer, p, s, rng in zip(self._sublayers, params, state, rngs):
+    for layer, p, s, rng in zip(self.sublayers, params, state, rngs):
       is_stack_just_one_item = (_count_items(stack) == 1)
 
       # Give layer its args from the stack; treat 1-arg layer specially.
-      n_in = layer.n_inputs()
+      n_in = layer.n_inputs
       if n_in == 1 and is_stack_just_one_item:
         inputs = stack
       elif n_in == 1:
@@ -199,7 +190,7 @@ def call(self, xs, params=(), state=(), **kwargs):
 
       # Push outputs onto remaining stack (if any).
       if n_in < _count_items(stack):
-        if layer.n_outputs() == 1:
+        if layer.n_outputs == 1:
           outputs = (outputs,)
         stack = outputs + stack[n_in:]
       else:
@@ -216,12 +207,12 @@ def MakeShapeType(shape, dtype):
     params = []
     states = []
     pseudo_xs = MakeShapeType(input_shape, input_dtype)
-    for layer in self._sublayers:
+    for layer in self.sublayers:
       rng, layer_rng = backend.random.split(rng)
 
       # Give layer its args from pseudo_xs; treat 1-arg layer specially.
       is_stack_just_one_item = (_count_items(pseudo_xs) == 1)
-      n_in = layer.n_inputs()
+      n_in = layer.n_inputs
       if n_in == 1 and is_stack_just_one_item:
         inputs = pseudo_xs
       elif n_in == 1:
@@ -238,7 +229,7 @@ def MakeShapeType(shape, dtype):
 
       # Push outputs onto remaining pseudo_xs (if any).
       if n_in < _count_items(pseudo_xs):
-        if layer.n_outputs() == 1:
+        if layer.n_outputs == 1:
           outputs = (outputs,)
         pseudo_xs = outputs + pseudo_xs[n_in:]
       else:
@@ -333,13 +324,10 @@ class Concatenate(base.Layer):
   """Concatenates n tensors into a single tensor."""
 
   def __init__(self, n_items=2, axis=-1):
-    super(Concatenate, self).__init__()
+    super(Concatenate, self).__init__(n_inputs=n_items)
     self._n_items = n_items
     self._axis = axis
 
-  def n_inputs(self):
-    return self._n_items
-
   def new_parameters(self, input_shape, input_dtype, rng):
     return (), ()
 
@@ -390,8 +378,8 @@ def __init__(self, *layers):
     layers = self._validate(layers)
     self._n_layers = len(layers)
     self._sublayers = layers
-    self._n_inputs = sum(x.n_inputs() for x in layers)
-    self._n_outputs = sum(x.n_outputs() for x in layers)
+    self._n_inputs = sum(x.n_inputs for x in layers)
+    self._n_outputs = sum(x.n_outputs for x in layers)
 
   def _validate(self, layers):
     if not layers or len(layers) < 2:
@@ -409,21 +397,12 @@ def _validate(self, layers):
           raise ValueError(
               'Found nonlayer object ({}) in layers list: [{}].'.format(
                   obj, layers))
-      if layers[i].n_inputs() == 0:
+      if layers[i].n_inputs == 0:
         raise ValueError(
             'Sublayer with n_inputs = 0 not allowed in Parallel:'
             ' {}'.format(layers[i]))
     return layers
 
-  def n_inputs(self):
-    return self._n_inputs
-
-  def n_outputs(self):
-    return self._n_outputs
-
-  def sublayers(self):
-    return self._sublayers
-
   def _allot_to_sublayers(self, inputs):
     """Divides Parallel's inputs for use by the sublayers.
 
@@ -437,8 +416,8 @@ def _allot_to_sublayers(self, inputs):
     """
     start, end = 0, 0
     sub_inputs = []
-    for layer in self._sublayers:
-      n_in = layer.n_inputs()
+    for layer in self.sublayers:
+      n_in = layer.n_inputs
       end = start + n_in
       if n_in == 1:
         sub_inputs.append(inputs[start])
@@ -448,7 +427,7 @@ def _allot_to_sublayers(self, inputs):
     return tuple(sub_inputs)
 
   def call(self, inputs, params=(), state=(), **kwargs):
-    n_layers, layers = self._n_layers, self._sublayers
+    n_layers, layers = self._n_layers, self.sublayers
     sublayer_inputs = self._allot_to_sublayers(inputs)
     rngs = _pop_rng_and_split(kwargs, n_layers)
     assert len(sublayer_inputs) == n_layers
@@ -460,12 +439,12 @@ def call(self, inputs, params=(), state=(), **kwargs):
     for layer, x, p, s, r in zip(layers, sublayer_inputs, params, state, rngs):
       # Note that zip silently truncates its result if lengths don't match.
       sub_outputs, s = layer(x, params=p, state=s, rng=r, **kwargs)
-      if layer.n_outputs() == 1:
+      if layer.n_outputs == 1:
         outputs.append(sub_outputs)
       else:
         outputs.extend(sub_outputs)
       new_state.append(s)
-    output = outputs[0] if self.n_outputs() == 1 else tuple(outputs)
+    output = outputs[0] if self.n_outputs == 1 else tuple(outputs)
     return output, new_state
 
   def new_parameters(self, input_shapes, input_dtypes, rng):
@@ -473,7 +452,7 @@ def new_parameters(self, input_shapes, input_dtypes, rng):
     sublayer_dtypes = self._allot_to_sublayers(input_dtypes)
     rngs = backend.random.split(rng, self._n_layers)
     inits = [layer.initialize(shape, dtype, rng) for layer, shape, dtype, rng
-             in zip(self._sublayers, sublayer_shapes, sublayer_dtypes, rngs)]
+             in zip(self.sublayers, sublayer_shapes, sublayer_dtypes, rngs)]
     if not inits:
       return (), ()
     else:
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 62d334e2f..9079da7c5 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -43,7 +43,7 @@ def __init__(self, layer, n_sections=1, check_shapes=True):
     Returns:
       A new layer representing mapping layer to all elements of the input.
     """
-    super(Map, self).__init__()
+    super(Map, self).__init__(n_inputs=n_sections, n_outputs=n_sections)
     if layer is None or isinstance(layer, (list, tuple)):
       layer = tl.Serial(layer)
     self._layer = layer
@@ -55,14 +55,6 @@ def __init__(self, layer, n_sections=1, check_shapes=True):
     self._check_shapes = check_shapes
     self._n_sections = n_sections
 
-  def n_inputs(self):
-    """Specifies how many data tensors this layer expects as input."""
-    return self._n_sections
-
-  def n_outputs(self):
-    """Specifies how many data tensors this layer promises as output."""
-    return self._n_sections
-
   def call(self, inputs, params=(), state=(), **kwargs):
     rngs = _pop_rng_and_split(kwargs, len(inputs))
     results = [self._layer(x, params=params, state=state, rng=r, **kwargs)
@@ -119,7 +111,7 @@ class Split(tl.Layer):
   """Splits the input into sections along an axis."""
 
   def __init__(self, n_sections=2, axis=-1):
-    super(Split, self).__init__()
+    super(Split, self).__init__(n_outputs=n_sections)
     self._n_sections = n_sections
     self._axis = axis
 
@@ -131,14 +123,6 @@ def call(self, inputs, params=(), state=(), **kwargs):
   def new_parameters(self, input_shapes, input_dtype, rng):
     return (), ()
 
-  def n_inputs(self):
-    """Specifies how many data tensors this layer expects as input."""
-    return 1
-
-  def n_outputs(self):
-    """Specifies how many data tensors this layer promises as output."""
-    return self._n_sections
-
 
 class SplitForOutput(tl.ReversibleLayer):
   """Splits activations into sections (for use right before the output layer).
@@ -156,18 +140,10 @@ class SplitForOutput(tl.ReversibleLayer):
   """
 
   def __init__(self, n_sections=2, axis=-2):
-    super(SplitForOutput, self).__init__()
+    super(SplitForOutput, self).__init__(n_inputs=2, n_outputs=n_sections)
     self._n_sections = n_sections
     self._axis = axis
 
-  def n_inputs(self):
-    """Specifies how many data tensors this layer expects as input."""
-    return 2
-
-  def n_outputs(self):
-    """Specifies how many data tensors this layer promises as output."""
-    return self._n_sections
-
   def new_parameters(self, input_shape, input_dtype, rng):
     return (), ()
 

From 55c7bf2c01299a9a9836be31596b50172274c26d Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 29 Aug 2019 15:49:04 -0700
Subject: [PATCH 2346/2720] Add an upper bound for the learning rate in
 OnlineTuneEnv

PiperOrigin-RevId: 266247524
---
 tensor2tensor/trax/rl/envs/online_tune_env.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index bfe16218a..759cbae6d 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -61,6 +61,7 @@ def __init__(self,
                eval_steps=10,
                env_steps=100,
                start_lr=0.001,
+               max_lr=10.0,
                # Don't save checkpoints by default, as they tend to use a lot of
                # space.
                should_save_checkpoints=False):
@@ -83,6 +84,7 @@ def __init__(self,
     self._eval_steps = eval_steps
     self._env_steps = env_steps
     self._start_lr = start_lr
+    self._max_lr = max_lr
 
     self._output_dir = output_dir
     gfile.makedirs(self._output_dir)
@@ -163,7 +165,8 @@ def step(self, action):
         metric since the last step. done is set after reaching self.env_steps
         environment steps. info is an empty dict.
     """
-    self._current_lr *= self._action_multipliers[action]
+    self._current_lr = min(
+        self._current_lr * self._action_multipliers[action], self._max_lr)
     last_metric_value = self._current_metric_value(self._reward_metric)
     self._trainer.train_epoch(self._train_steps, self._eval_steps)
     self._step += 1

From 3552c66aaf5a9331265375aee7a907948fb450d2 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 29 Aug 2019 16:22:42 -0700
Subject: [PATCH 2347/2720] Standardize attention between TransformerLM and
 TransformerRevnetLM

PiperOrigin-RevId: 266254395
---
 .../trax/configs/transformer_big_lm1b_8gb.gin |   5 +
 .../trax/configs/transformer_imagenet64.gin   |  51 ++++++
 .../trax/configs/transformer_imdb_8gb.gin     |   5 +
 .../trax/configs/transformer_lm1b_8gb.gin     |   5 +
 .../configs/transformer_lm1b_8gb_testing.gin  |   5 +
 tensor2tensor/trax/layers/attention.py        | 151 ++++++++++++++----
 tensor2tensor/trax/layers/base.py             |   4 +-
 .../models/research/transformer_revnet.py     |  45 +++---
 tensor2tensor/trax/models/transformer.py      |  36 ++++-
 9 files changed, 241 insertions(+), 66 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/transformer_imagenet64.gin

diff --git a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
index 63b693f7b..46e28ad13 100644
--- a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
@@ -40,8 +40,13 @@ train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.SM3
 train.train_steps = 500000
 
+# Parameters for DotProductCausalAttention:
+# ==============================================================================
+DotProductCausalAttention.dropout = 0.1
+
 # Parameters for TransformerLM:
 # ==============================================================================
+TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
 TransformerLM.d_model = 1024
 TransformerLM.d_ff = 8192
 TransformerLM.dropout = 0.1
diff --git a/tensor2tensor/trax/configs/transformer_imagenet64.gin b/tensor2tensor/trax/configs/transformer_imagenet64.gin
new file mode 100644
index 000000000..0f56cf59a
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_imagenet64.gin
@@ -0,0 +1,51 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 2
+batch_fun.eval_batch_size = 128
+batch_fun.max_eval_length = 12288  # 64 * 64 * 3
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
+inputs.input_name = 'targets'
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 2.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 200
+train.eval_steps = 8
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerLM
+train.optimizer = @trax.optimizers.Adafactor
+train.train_steps = 500000
+train.trainer_class = @MemoryEfficientTrainer
+
+# Parameters for MemoryEfficientCausalAttention:
+# ==============================================================================
+MemoryEfficientCausalAttention.dropout = 0.0
+MemoryEfficientCausalAttention.loop_stride = 512
+
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.attention_type = @trax.layers.MemoryEfficientCausalAttention
+TransformerLM.d_model = 1024
+TransformerLM.d_ff = 2048
+TransformerLM.dropout = 0.0
+TransformerLM.max_len = 12288  # 64 * 64 * 3
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 4
+TransformerLM.n_layers = 3
+TransformerLM.vocab_size = 256
+
diff --git a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
index 58391c63d..ba44965bb 100644
--- a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
@@ -36,8 +36,13 @@ train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerEncoder
 train.train_steps = 1000
 
+# Parameters for DotProductCausalAttention:
+# ==============================================================================
+DotProductCausalAttention.dropout = 0.1
+
 # Parameters for TransformerLM:
 # ==============================================================================
+TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
 TransformerEncoder.d_model = 512
 TransformerEncoder.d_ff = 2048
 TransformerEncoder.dropout = 0.1
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index 5549858d7..33c32a465 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -40,8 +40,13 @@ train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.SM3
 train.train_steps = 500000
 
+# Parameters for DotProductCausalAttention:
+# ==============================================================================
+DotProductCausalAttention.dropout = 0.1
+
 # Parameters for TransformerLM:
 # ==============================================================================
+TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
 TransformerLM.d_model = 512
 TransformerLM.d_ff = 2048
 TransformerLM.dropout = 0.1
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
index 56f779b52..2f8f91ec6 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
@@ -40,8 +40,13 @@ train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.Adam
 train.train_steps = 100000
 
+# Parameters for DotProductCausalAttention:
+# ==============================================================================
+DotProductCausalAttention.dropout = 0.1
+
 # Parameters for TransformerLM:
 # ==============================================================================
+TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
 TransformerLM.d_model = 512
 TransformerLM.d_ff = 2048
 TransformerLM.dropout = 0.1
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 19a88e30f..8f5b4b559 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -29,6 +29,16 @@
 from tensor2tensor.trax.layers import initializers as init
 
 
+@base.layer()
+def ShiftRight(x, **unused_kwargs):
+  """Layer to shift the tensor to the right by padding on axis 1."""
+  pad_widths = [(0, 0)] * len(x.shape)
+  pad_widths[1] = (1, 0)  # Padding on axis=1
+  padded = np.pad(x, pad_widths, mode='constant',
+                  constant_values=x.dtype.type(0))
+  return padded[:, :-1]
+
+
 @base.layer()
 def CausalMask(x, params, axis=-1, **kwargs):
   del params, kwargs
@@ -221,16 +231,6 @@ def BasicCausalAttention(d_feature, n_heads=1, dropout=0.0, mode='train'):
   ]
 
 
-@base.layer()
-def ShiftRight(x, **unused_kwargs):
-  """Layer to shift the tensor to the right by padding on axis 1."""
-  pad_widths = [(0, 0)] * len(x.shape)
-  pad_widths[1] = (1, 0)  # Padding on axis=1
-  padded = np.pad(x, pad_widths, mode='constant',
-                  constant_values=x.dtype.type(0))
-  return padded[:, :-1]
-
-
 class ComputeAttentionHeads(base.Layer):
   """Computes queries/keys/values via linear projection.
 
@@ -301,19 +301,36 @@ def new_parameters(self, input_shape, input_dtype, rng):
 
 
 class BaseCausalAttention(base.Layer):
-  """Base class for variants of causal self-attention.
-
-  This class sets up an API that includes forward_and_vjp, which is required to
-  implement MemoryEfficientCausalAttention.
-  """
+  """Base class for variants of causal self-attention."""
 
   def __init__(self):
     super(BaseCausalAttention, self).__init__(n_inputs=3)
 
   def call(self, inputs, params=(), state=(), rng=None, **kwargs):
+    """Forward pass for the attention layer."""
     raise NotImplementedError()
 
-  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
+  def call_and_grad(self, inputs, grad, **kwargs):
+    """Performs both forward and backward pass for the attention layer.
+
+    This is used in reversible models: for the backward pass of a reversible
+    model, we need to compute both the forward direction (to recover the
+    previous layer's activations) and the backward direction simultaneously.
+    Some computation can be shared between the forward and backward directions,
+    which makes it more efficient to implement them jointly.
+
+    This method assumes that the layer is stateless and has no parameters.
+
+    Args:
+      inputs: A tuple (q, k, v), where each element has shape
+          n_batch*n_heads, seqlen, d_head
+      grad: gradient signal for the layer output.
+      **kwargs: kwargs for the layer
+
+    Returns:
+      A nested-tuple structure (output, (q_grad, k_grad, v_grad)) that contains
+      the output of the forward pass and the gradient signal for each input.
+    """
     raise NotImplementedError()
 
   def new_parameters(self, input_shapes, input_dtype, rng):
@@ -323,7 +340,7 @@ def new_parameters(self, input_shapes, input_dtype, rng):
 class DotProductCausalAttention(BaseCausalAttention):
   """A standard (non-memory-efficient) dot product attention implementation."""
 
-  def __init__(self, dropout, mode):
+  def __init__(self, dropout=0.0, mode='train'):
     super(DotProductCausalAttention, self).__init__()
     self._dropout = dropout
     self._mode = mode
@@ -332,24 +349,41 @@ def call(self, inputs, params=(), state=(), rng=None, **kwargs):
     del params
     q, k, v = inputs
     mask_size = q.shape[-2]
-    mask = np.tril(np.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
+    # Not all backends define np.tril. However, using onp.tril is inefficient in
+    # that it creates a large global constant. TODO(kitaev): try to find an
+    # alternative that works across all backends.
+    if backend.get_name() == 'jax':
+      mask = np.tril(np.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
+    else:
+      mask = onp.tril(onp.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
     res = DotProductAttention(
         q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=rng)
     return res, state
 
-  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
+  def call_and_grad(self, inputs, ct, **kwargs):
     assert backend.get_name() == 'jax', (
-        'JAX backend is required to use forward_and_vjp.')
+        'JAX backend is required to use call_and_grad.')
     # Simultaneous forward pass and backprop through the attention mechanism.
     def do_call(x):  # pylint: disable=invalid-name
-      res, _ = self.call(x, params, **kwargs)
+      res, _ = self.call(x, **kwargs)
       return res
     output, vjpfun = jax.vjp(do_call, inputs)
     return output, vjpfun(ct)[0]
 
 
 class MemoryEfficientCausalAttention(BaseCausalAttention):
-  """Memory-efficient dot product attention."""
+  """Memory-efficient dot product attention.
+
+  This layer performs causal attention on long sequences without running out
+  of memory. Instead of computing dot products for all query-key pairs at once,
+  it uses a loop to compute attention for a small set of query positions at a
+  time. The "loop_stride" parameter controls how many query positions are
+  considered at each iteration of the loop.
+
+  Note that this class does not slice along the batch/head dimension. Looping
+  over batch elements and heads instead of query positions is also a viable
+  option. We haven't implemented it, but it may perform well, too.
+  """
 
   def __init__(self, loop_stride, dropout, mode):
     assert backend.get_name() == 'jax', (
@@ -364,27 +398,37 @@ def __init__(self, loop_stride, dropout, mode):
       self.dropout = None
 
   def call(self, inputs, params=(), state=(), **kwargs):
-    output, _ = self.forward_and_vjp(inputs, None, params=params, **kwargs)
+    del params
+    output, _ = self.call_and_grad(inputs, None, **kwargs)
     return output, state
 
-  def forward_and_vjp(self, inputs, ct, params=(), rng=None, **kwargs):
-    # This is the core of the memory-efficient attention implementation, where
-    # we use the jax.lax.while_loop primitive to compute attention for a small
-    # set of query positions at a time. Note how in the backwards pass, we
-    # compute both the forward direction (to recover the previous layer's
-    # activations) and the backward direction simultaneously. This allows us to
-    # only use a single loop, where the inner portion of the loop does a slice
-    # of the forward+backward joint computation. Unfortunately we have had to
-    # introduce a large number of wrapper classes (including
-    # ReversibleAttentionHalfResidual and ApplyAttentionWrapper) for the sole
-    # purpose of connecting this implementation of forward_and_vjp with the core
-    # backprop implementation.
+  def has_custom_grad(self):
+    return True
+
+  def custom_grad(self, inputs, output, ct, params=(), state=(), **kwargs):
+    del output, params, state
+    _, inputs_ct = self.call_and_grad(inputs, ct, **kwargs)
+    return inputs_ct, ()
 
+  def call_and_grad(self, inputs, ct, rng=None, **kwargs):
+    del kwargs
     query, key, value = inputs
     depth = np.shape(query)[-1]
     do_backprop = ct is not None
+    # jax uses the term cotangent (ct) to refer to gradient signals, and
+    # vector-Jacobian product (vjp) for back-propagation through a layer.
 
     def make_mask(N, M, k):  # pylint: disable=invalid-name
+      """Constructs a slice of the causal attention mask.
+
+      Args:
+        N: number of query positions
+        M: number of key positions
+        k: position of the initial query element
+
+      Returns:
+        N x M mask, where 1.0 indicates that attention is not allowed.
+      """
       x = np.arange(N, dtype=np.int32)
       y = np.arange(M, dtype=np.int32)
       mask = jax.lax.lt(
@@ -487,3 +531,40 @@ def body_fun(vals):  # pylint: disable=invalid-name
       return final_vals[1], None
     else:
       return final_vals[1], final_vals[2:]
+
+
+def CausalAttention(d_feature, n_heads=1,
+                    d_attention_key=None, d_attention_value=None,
+                    attention_type=DotProductCausalAttention, mode='train'):
+  """Transformer-style multi-headed causal attention.
+
+  Args:
+    d_feature: int:  dimensionality of feature embedding
+    n_heads: int: number of attention heads
+    d_attention_key: int: depth of key vector for each attention head
+        (default is d_feature // n_heads)
+    d_attention_value: int: depth of value vector for each attention head
+        (default is d_feature // n_heads)
+    attention_type: subclass of tl.BaseCausalAttention: attention class to use
+    mode: str: 'train' or 'eval'
+
+  Returns:
+    Multi-headed self-attention result.
+  """
+  if d_attention_key is None:
+    assert d_feature % n_heads == 0
+    d_attention_key = d_feature // n_heads
+  if d_attention_value is None:
+    assert d_feature % n_heads == 0
+    d_attention_value = d_feature // n_heads
+
+  return [
+      cb.Dup(), cb.Dup(),
+      cb.Parallel(
+          ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
+          ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
+          ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
+      ),
+      attention_type(mode=mode),
+      ComputeAttentionOutput(n_heads=n_heads, d_model=d_feature),
+  ]
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index ec680b9e9..93b9a6127 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -166,8 +166,8 @@ def custom_grad(self, inputs, output, grad, params, state, **kwargs):
 
     Returns:
       The custom gradient signal for the input. Note that we need to return
-      a gradient for each argument of call, so it will usually be a triple
-      of signals: the gradient for inputs, parameters, and kwargs.
+      a gradient for each argument of call, so it will usually be a tuple
+      of signals: the gradient for inputs and parameters.
     """
     raise NotImplementedError
 
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 9079da7c5..b98a3101b 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -254,25 +254,21 @@ def call_compute_residual(x, params):
 
 
 class ApplyAttentionWrapper(tl.Parallel):
-  """Same as tl.Parallel(attention, [], []), but implements forward_and_vjp.
-
-  See MemoryEfficientDotProductAttention for why this is needed.
-  """
+  """Same as tl.Parallel(attention, [], []), but implements call_and_grad."""
 
   def __init__(self, attention):
-    assert hasattr(attention, 'forward_and_vjp')
+    assert hasattr(attention, 'call_and_grad')
     super(ApplyAttentionWrapper, self).__init__(attention, [], [])
     self.attention = attention
 
-  def forward_and_vjp(self, inputs, ct, params=(), **kwargs):
+  def call_and_grad(self, inputs, ct, **kwargs):
     # Simultaneous forward pass and backprop through the attention mechanism.
     qkv = inputs[:3]
     passthrough = inputs[3:]
     out_ct = ct[0]
     passthrough_ct = ct[1:]
 
-    out, qkv_ct = self.attention.forward_and_vjp(
-        qkv, out_ct, params=(), **kwargs)
+    out, qkv_ct = self.attention.call_and_grad(qkv, out_ct, **kwargs)
     return (out,) + passthrough, qkv_ct + passthrough_ct
 
 
@@ -282,15 +278,16 @@ class ReversibleAttentionHalfResidual(tl.ReversibleLayer, tl.Serial):
   If inputs are (x1, x2), then outputs are (x1 + z, x2) where:
   z = post_attention(attention(pre_attention(x1)))
 
+  Other than an efficiency optimization, this layer is equivalent to
+  ReversibleHalfResidual([pre_attention, attention, post_attention]).
+
   The post_attention layers must be linear in their input (typically they will
-  consists of reshaping and dense linear layers). This allows back-propagating
-  the gradient signal from the output of ReversibleAttentionHalfResidual to the
-  output of the "attention" portion based only on the network parameters.
-
-  The forward pass is equivalent to using
-  ReversibleHalfResidual([pre_attention, attention, post_attention]), but the
-  backward pass uses attention.forward_and_vjp. See
-  MemoryEfficientDotProductAttention for why forward_and_vjp is helpful.
+  consists of reshaping and dense linear layers), which allows the following
+  optimization. We can back-propagate the gradient signal from the output of
+  ReversibleAttentionHalfResidual to the output of the "attention" portion based
+  only on the network parameters. Then, attention.call_and_grad can be used to
+  recover the output of the "attention" portion while simultaneously performing
+  the backward pass, which allows shared computation between the two directions.
   """
 
   def __init__(self, pre_attention, attention, post_attention):
@@ -300,7 +297,7 @@ def __init__(self, pre_attention, attention, post_attention):
         tl.Swap(),
         tl.Parallel(pre_attention, [], []),
     ])
-    assert hasattr(attention, 'forward_and_vjp')
+    assert hasattr(attention, 'call_and_grad')
     self.attention = ApplyAttentionWrapper(attention)
     self.post_attention = tl.Parallel(post_attention, [], [])
 
@@ -363,9 +360,9 @@ def call_post_attention(x):
     (ct,) = post_attention_vjpfun(ct)
 
     # Simultaneous forward pass and backprop through the attention mechanism
-    stack, ct = self.attention.forward_and_vjp(
-        stack, ct, rng=rngs[1], **kwargs)
-    attention_params_ct = params[1]  # Note: this assumes that params are empty.
+    stack, ct = self.attention.call_and_grad(stack, ct, rng=rngs[1], **kwargs)
+    assert not jax.tree_util.tree_leaves(params[1])
+    attention_params_ct = params[1]  # This is valid when params is empty.
 
     # Backprop through self.pre_attention
     x_ct, pre_attention_params_ct = pre_attention_vjpfun(ct)
@@ -407,7 +404,7 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
     d_attention_value: int: depth of value vector for each attention head
     n_heads: int: number of attention heads
     n_attention_chunks: int: number of chunks for attention
-    attention_type: class: attention class to use, such as DotProductAttention.
+    attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
 
@@ -420,9 +417,9 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
       tl.LayerNorm(),
       tl.Dup(), tl.Dup(),
       tl.Parallel(
-          [tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key)],
-          [tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key)],
-          [tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value)],
+          tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
+          tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
+          tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
       ),
   ]
 
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index af64e32f9..1dc2e0325 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -111,7 +111,8 @@ def TransformerEncoder(vocab_size,
   ])
 
 
-def DecoderBlock(d_model, d_ff, n_heads, dropout, mode):
+def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
+                 attention_type, dropout, mode):
   """Returns a layer sequence that implements a Transformer decoder block.
 
   The input to the layer sequence is an activation tensor.
@@ -120,6 +121,9 @@ def DecoderBlock(d_model, d_ff, n_heads, dropout, mode):
     d_model: int:  depth of embedding
     d_ff: int: depth of feed-forward layer
     n_heads: int: number of attention heads
+    d_attention_key: int: depth of key vector for each attention head
+    d_attention_value: int: depth of value vector for each attention head
+    attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
     mode: str: 'train' or 'eval'
 
@@ -128,8 +132,10 @@ def DecoderBlock(d_model, d_ff, n_heads, dropout, mode):
   """
   self_attention = [
       tl.LayerNorm(),  # vec
-      tl.BasicCausalAttention(
-          d_model, n_heads=n_heads, dropout=dropout, mode=mode),
+      tl.CausalAttention(
+          d_model, n_heads=n_heads, d_attention_key=d_attention_key,
+          d_attention_value=d_attention_value, attention_type=attention_type,
+          mode=mode),
       tl.Dropout(rate=dropout, mode=mode),  # vec
   ]
   feed_forward = [
@@ -145,6 +151,9 @@ def TransformerDecoder(d_model=512,
                        d_ff=2048,
                        n_layers=6,
                        n_heads=8,
+                       d_attention_key=None,
+                       d_attention_value=None,
+                       attention_type=tl.DotProductCausalAttention,
                        dropout=0.1,
                        max_len=2048,
                        mode='train'):
@@ -159,6 +168,11 @@ def TransformerDecoder(d_model=512,
     d_ff: int: depth of feed-forward layer
     n_layers: int: number of encoder/decoder layers
     n_heads: int: number of attention heads
+    d_attention_key: int: depth of key vector for each attention head
+        (default is d_model // n_heads)
+    d_attention_value: int: depth of value vector for each attention head
+        (default is d_model // n_heads)
+    attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
     max_len: int: maximum symbol length for positional encoding
     mode: str: 'train' or 'eval'
@@ -170,7 +184,9 @@ def TransformerDecoder(d_model=512,
   return tl.Model(                  # vecs
       tl.PositionalEncoding(max_len=max_len),
       tl.Dense(d_model),            # vecs
-      [DecoderBlock(d_model, d_ff, n_heads, dropout, mode)
+      [DecoderBlock(  # pylint: disable=g-complex-comprehension
+          d_model, d_ff, n_heads, d_attention_key, d_attention_value,
+          attention_type, dropout, mode)
        for _ in range(n_layers)],   # vecs
       tl.LayerNorm(),               # vecs
   )
@@ -181,6 +197,9 @@ def TransformerLM(vocab_size,
                   d_ff=2048,
                   n_layers=6,
                   n_heads=8,
+                  d_attention_key=None,
+                  d_attention_value=None,
+                  attention_type=tl.DotProductCausalAttention,
                   dropout=0.1,
                   max_len=2048,
                   mode='train'):
@@ -195,6 +214,11 @@ def TransformerLM(vocab_size,
     d_ff: int: depth of feed-forward layer
     n_layers: int: number of encoder/decoder layers
     n_heads: int: number of attention heads
+    d_attention_key: int: depth of key vector for each attention head
+        (default is d_model // n_heads)
+    d_attention_value: int: depth of value vector for each attention head
+        (default is d_model // n_heads)
+    attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
     max_len: int: maximum symbol length for positional encoding
     mode: str: 'train' or 'eval'
@@ -211,7 +235,9 @@ def TransformerLM(vocab_size,
   return tl.Model(                  # tokens
       tl.ShiftRight(),              # toks
       embedder,                     # vecs
-      [DecoderBlock(d_model, d_ff, n_heads, dropout, mode)
+      [DecoderBlock(  # pylint: disable=g-complex-comprehension
+          d_model, d_ff, n_heads, d_attention_key, d_attention_value,
+          attention_type, dropout, mode)
        for _ in range(n_layers)],   # vecs
       tl.LayerNorm(),               # vecs
       tl.Dense(vocab_size),         # vecs

From f86e40a33fbf6dbaac9500765fcb03697a9f194d Mon Sep 17 00:00:00 2001
From: Vladimir Vlasov <mr.voov@gmail.com>
Date: Fri, 30 Aug 2019 01:42:41 +0200
Subject: [PATCH 2348/2720] save attention weights for relative attention v2
 (#1682)

---
 tensor2tensor/layers/common_attention.py | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 5a1d6e43a..ce8a7178c 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1865,6 +1865,7 @@ def dot_product_self_attention_relative_v2(q,
                                            max_relative_position=None,
                                            dropout_rate=0.0,
                                            image_shapes=None,
+                                           save_weights_to=None,
                                            name=None,
                                            make_image_summary=True,
                                            dropout_broadcast_dims=None,
@@ -1886,6 +1887,9 @@ def dot_product_self_attention_relative_v2(q,
       to look back - changing this invalidates checkpoints
     dropout_rate: a floating point number.
     image_shapes: optional tuple of integer scalars.
+    save_weights_to: an optional dictionary to capture attention weights
+      for visualization; the weights tensor will be appended there under
+      a string key created from the variable scope (including name).
     name: an optional string.
     make_image_summary: Whether to make an attention image summary.
     dropout_broadcast_dims:  an optional list of integers less than 4
@@ -1908,7 +1912,7 @@ def dot_product_self_attention_relative_v2(q,
   with tf.variable_scope(
       name,
       default_name="dot_product_self_attention_relative_v2",
-      values=[q, k, v]):
+      values=[q, k, v]) as scope:
 
     # This calculation only works for self attention.
     # q, k and v must therefore have the same shape.
@@ -1932,6 +1936,9 @@ def dot_product_self_attention_relative_v2(q,
       logits += bias
 
     weights = tf.nn.softmax(logits, name="attention_weights")
+    if save_weights_to is not None:
+      save_weights_to[scope.name] = weights
+      save_weights_to[scope.name + "/logits"] = logits
     # Dropping out the attention links for each of the heads.
     weights = common_layers.dropout_with_broadcast_dims(
         weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
@@ -2036,7 +2043,7 @@ def get_relative_embeddings_left_right(max_relative_position, length, depth,
 
 def dot_product_unmasked_self_attention_relative_v2(
     q, k, v, bias, max_relative_position=None, dropout_rate=0.0,
-    image_shapes=None, name=None, make_image_summary=True,
+    image_shapes=None, save_weights_to=None, name=None, make_image_summary=True,
     dropout_broadcast_dims=None, heads_share_relative_embedding=False,
     add_relative_to_values=False):
   """Calculate relative position-aware dot-product self-attention.
@@ -2053,6 +2060,9 @@ def dot_product_unmasked_self_attention_relative_v2(
       Changing this invalidates checkpoints.
     dropout_rate: a floating point number.
     image_shapes: optional tuple of integer scalars.
+    save_weights_to: an optional dictionary to capture attention weights
+      for visualization; the weights tensor will be appended there under
+      a string key created from the variable scope (including name).
     name: an optional string.
     make_image_summary: Whether to make an attention image summary.
     dropout_broadcast_dims:  an optional list of integers less than 4
@@ -2076,7 +2086,7 @@ def dot_product_unmasked_self_attention_relative_v2(
   with tf.variable_scope(
       name,
       default_name="dot_product_unmasked_self_attention_relative_v2",
-      values=[q, k, v]):
+      values=[q, k, v]) as scope:
 
     # This calculation only works for self attention.
     # q, k and v must therefore have the same shape.
@@ -2104,6 +2114,9 @@ def dot_product_unmasked_self_attention_relative_v2(
     if bias is not None:
       logits += bias
     weights = tf.nn.softmax(logits, name="attention_weights")
+    if save_weights_to is not None:
+      save_weights_to[scope.name] = weights
+      save_weights_to[scope.name + "/logits"] = logits
     # dropping out the attention links for each of the heads
     weights = common_layers.dropout_with_broadcast_dims(
         weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
@@ -4664,6 +4677,7 @@ def multihead_attention(query_antecedent,
           max_relative_position,
           dropout_rate,
           image_shapes,
+          save_weights_to=save_weights_to,
           make_image_summary=make_image_summary,
           dropout_broadcast_dims=dropout_broadcast_dims,
           heads_share_relative_embedding=heads_share_relative_embedding,
@@ -4677,6 +4691,7 @@ def multihead_attention(query_antecedent,
           max_relative_position,
           dropout_rate,
           image_shapes,
+          save_weights_to=save_weights_to,
           make_image_summary=make_image_summary,
           dropout_broadcast_dims=dropout_broadcast_dims,
           heads_share_relative_embedding=heads_share_relative_embedding,

From af684bb7df810076bee8d64f762fbeb12bc52751 Mon Sep 17 00:00:00 2001
From: Sepehr Sameni <Sepehr.Sameni@gmail.com>
Date: Fri, 30 Aug 2019 04:16:56 +0430
Subject: [PATCH 2349/2720] correct return shape of rel_pos2abs_pos() (#1686)

---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index ce8a7178c..391987458 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2744,7 +2744,7 @@ def _relative_position_to_absolute_position_unmasked(x):
                           length, 2 * length - 1]
 
   Returns:
-    A Tensor of shape [batch (or batch*num_blocks), heads, length, length-1]
+    A Tensor of shape [batch (or batch*num_blocks), heads, length, length]
   """
   x_shape = common_layers.shape_list(x)
   batch = x_shape[0]

From ee8b8c81905267d87759308a3bf08cd1f1f1f9b0 Mon Sep 17 00:00:00 2001
From: Vladimir Vlasov <mr.voov@gmail.com>
Date: Thu, 29 Aug 2019 16:43:06 -0700
Subject: [PATCH 2350/2720] Merge of PR #1682

PiperOrigin-RevId: 266259504
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 391987458..ce8a7178c 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2744,7 +2744,7 @@ def _relative_position_to_absolute_position_unmasked(x):
                           length, 2 * length - 1]
 
   Returns:
-    A Tensor of shape [batch (or batch*num_blocks), heads, length, length]
+    A Tensor of shape [batch (or batch*num_blocks), heads, length, length-1]
   """
   x_shape = common_layers.shape_list(x)
   batch = x_shape[0]

From 32929305e1a4ec926eff24123758b794df35492b Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 29 Aug 2019 16:45:49 -0700
Subject: [PATCH 2351/2720] Fiddle around with the test gin configuration.

PiperOrigin-RevId: 266260056
---
 tensor2tensor/trax/rl/configs/atari_regression_test.gin | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/rl/configs/atari_regression_test.gin b/tensor2tensor/trax/rl/configs/atari_regression_test.gin
index 98a2c896d..54a96d953 100644
--- a/tensor2tensor/trax/rl/configs/atari_regression_test.gin
+++ b/tensor2tensor/trax/rl/configs/atari_regression_test.gin
@@ -23,7 +23,7 @@ PPO.policy_and_value_model = @trax.models.AtariCnn
 # Parameters for train_rl:
 # ==============================================================================
 train_rl.env_name = "PongNoFrameskip-v4"
-train_rl.n_epochs = 10000
+train_rl.n_epochs = 4000
 train_rl.clip_rewards = True
 train_rl.max_timestep = 10000
 train_rl.rendered_env = True

From c771c7f8e9ab5858d49048557f40ec18e85b391b Mon Sep 17 00:00:00 2001
From: Sepehr Sameni <Sepehr.Sameni@gmail.com>
Date: Thu, 29 Aug 2019 17:01:19 -0700
Subject: [PATCH 2352/2720] Merge of PR #1686

PiperOrigin-RevId: 266263195
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index ce8a7178c..391987458 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2744,7 +2744,7 @@ def _relative_position_to_absolute_position_unmasked(x):
                           length, 2 * length - 1]
 
   Returns:
-    A Tensor of shape [batch (or batch*num_blocks), heads, length, length-1]
+    A Tensor of shape [batch (or batch*num_blocks), heads, length, length]
   """
   x_shape = common_layers.shape_list(x)
   batch = x_shape[0]

From 63112b49296acaea91f3721cc23194ab34199d4b Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 29 Aug 2019 17:12:55 -0700
Subject: [PATCH 2353/2720] Report more metrics in PPO

PiperOrigin-RevId: 266265617
---
 tensor2tensor/trax/rl/ppo.py         | 51 ++++++++++----
 tensor2tensor/trax/rl/ppo_test.py    |  8 +--
 tensor2tensor/trax/rl/ppo_trainer.py | 99 +++++++++++++++-------------
 3 files changed, 97 insertions(+), 61 deletions(-)

diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index 4156394a9..959a27d8f 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -389,7 +389,9 @@ def value_loss_given_predictions(value_prediction,
       well. This is from the OpenAI baselines implementation.
 
   Returns:
-    The average L2 value loss, averaged over instances where reward_mask is 1.
+    Pair (value_loss, summaries), where value_loss is the average L2 value loss,
+      averaged over instances where reward_mask is 1. Summaries is a dict of
+      summaries collected during value loss computation.
   """
 
   B, T = rewards.shape  # pylint: disable=invalid-name
@@ -412,7 +414,13 @@ def value_loss_given_predictions(value_prediction,
     loss = np.maximum(v_clipped_loss, loss)
 
   # Take an average on only the points where mask != 0.
-  return np.sum(loss) / np.sum(reward_mask)
+  value_loss = np.sum(loss) / np.sum(reward_mask)
+
+  summaries = {
+      "value_loss": value_loss,
+  }
+
+  return (value_loss, summaries)
 
 
 def deltas(predicted_values, rewards, mask, gamma=0.99):
@@ -560,7 +568,9 @@ def ppo_loss_given_predictions(log_probab_actions_new,
       td_deltas, reward_mask, lambda_=lambda_, gamma=gamma)
 
   # Normalize the advantages.
-  advantages = (advantages - np.mean(advantages)) / np.std(advantages)
+  advantage_mean = np.mean(advantages)
+  advantage_std = np.std(advantages)
+  advantages = (advantages - advantage_mean) / (advantage_std + 1e-8)
 
   # (B, T)
   ratios = compute_probab_ratios(log_probab_actions_new, log_probab_actions_old,
@@ -576,7 +586,15 @@ def ppo_loss_given_predictions(log_probab_actions_new,
   average_objective = np.sum(objective) / np.sum(reward_mask)
 
   # Loss is negative objective.
-  return -average_objective
+  ppo_loss = -average_objective
+
+  summaries = {
+      "ppo_loss": ppo_loss,
+      "advantage_mean": advantage_mean,
+      "advantage_std": advantage_std,
+  }
+
+  return (ppo_loss, summaries)
 
 
 @jit
@@ -593,14 +611,14 @@ def combined_loss_given_predictions(log_probab_actions_new,
                                     c1=1.0,
                                     c2=0.01):
   """Computes the combined (clipped loss + value loss) given predictions."""
-  loss_value = value_loss_given_predictions(
+  (value_loss, value_summaries) = value_loss_given_predictions(
       value_prediction_new,
       padded_rewards,
       reward_mask,
       gamma=gamma,
       value_prediction_old=value_prediction_old,
       epsilon=epsilon)
-  loss_ppo = ppo_loss_given_predictions(
+  (ppo_loss, ppo_summaries) = ppo_loss_given_predictions(
       log_probab_actions_new,
       log_probab_actions_old,
       value_prediction_old,
@@ -611,8 +629,16 @@ def combined_loss_given_predictions(log_probab_actions_new,
       lambda_=lambda_,
       epsilon=epsilon)
   entropy_bonus = masked_entropy(log_probab_actions_new, reward_mask)
-  return (loss_ppo + (c1 * loss_value) - (c2 * entropy_bonus), loss_ppo,
-          loss_value, entropy_bonus)
+  combined_loss_ = ppo_loss + (c1 * value_loss) - (c2 * entropy_bonus)
+
+  summaries = {
+      "combined_loss": combined_loss_,
+      "entropy_bonus": entropy_bonus,
+  }
+  for loss_summaries in (value_summaries, ppo_summaries):
+    summaries.update(loss_summaries)
+
+  return (combined_loss_, (ppo_loss, value_loss, entropy_bonus), summaries)
 
 
 @functools.partial(jit, static_argnums=(3,))
@@ -636,8 +662,7 @@ def combined_loss(new_params,
       policy_and_value_net_apply(padded_observations, new_params, state,
                                  rng=rng))
 
-  # (combined_loss, ppo_loss, value_loss, entropy_bonus)
-  return combined_loss_given_predictions(
+  (loss, component_losses, summaries) = combined_loss_given_predictions(
       log_probab_actions_new,
       log_probab_actions_old,
       value_predictions_new,
@@ -649,7 +674,9 @@ def combined_loss(new_params,
       lambda_=lambda_,
       epsilon=epsilon,
       c1=c1,
-      c2=c2), state
+      c2=c2,
+  )
+  return (loss, component_losses, summaries, state)
 
 
 @functools.partial(jit, static_argnums=(2, 3, 4))
@@ -676,7 +703,7 @@ def policy_and_value_opt_step(i,
   # Combined loss function given the new params.
   def policy_and_value_loss(params, state):
     """Returns the combined loss given just parameters."""
-    (loss, _, _, _), state = combined_loss(
+    (loss, _, _, state) = combined_loss(
         params,
         log_probab_actions_old,
         value_predictions_old,
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
index c7ddc314e..b0aa4e86b 100644
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -194,7 +194,7 @@ def value_net_apply(observations, params, rng=None):
     value_prediction = value_net_apply(random_observations, [])
 
     with jax.disable_jit():
-      value_loss = ppo.value_loss_given_predictions(
+      (value_loss, _) = ppo.value_loss_given_predictions(
           value_prediction,
           rewards,
           rewards_mask,
@@ -419,10 +419,10 @@ def test_combined_loss(self):
     c1 = 1.0
     c2 = 0.01
 
-    value_loss_1 = ppo.value_loss_given_predictions(
+    (value_loss_1, _) = ppo.value_loss_given_predictions(
         value_predictions_new, rewards, mask, gamma=gamma,
         value_prediction_old=value_predictions_old, epsilon=epsilon)
-    ppo_loss_1 = ppo.ppo_loss_given_predictions(
+    (ppo_loss_1, _) = ppo.ppo_loss_given_predictions(
         new_log_probabs,
         old_log_probabs,
         value_predictions_old,
@@ -433,7 +433,7 @@ def test_combined_loss(self):
         lambda_=lambda_,
         epsilon=epsilon)
 
-    (combined_loss, ppo_loss_2, value_loss_2, entropy_bonus), state = (
+    (combined_loss, (ppo_loss_2, value_loss_2, entropy_bonus), _, state) = (
         ppo.combined_loss(new_params,
                           old_log_probabs,
                           value_predictions_old,
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index 8889e992d..d86468177 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -319,35 +319,36 @@ def train_epoch(self):
     self._rng, key1 = jax_random.split(self._rng, num=2)
     logging.vlog(2, "Starting to compute P&V loss.")
     loss_compute_start_time = time.time()
-    (cur_combined_loss, cur_ppo_loss, cur_value_loss,
-     entropy_bonus), self._model_state = (
-         ppo.combined_loss(
-             self._policy_and_value_net_params,
-             log_probabs_traj,
-             value_predictions_traj,
-             self._policy_and_value_net_apply,
-             padded_observations,
-             padded_actions,
-             padded_rewards,
-             reward_mask,
-             gamma=self._gamma,
-             lambda_=self._lambda_,
-             c1=self._c1,
-             c2=self._c2,
-             state=self._model_state,
-             rng=key1))
+    (cur_combined_loss, component_losses, summaries, self._model_state) = (
+        ppo.combined_loss(
+            self._policy_and_value_net_params,
+            log_probabs_traj,
+            value_predictions_traj,
+            self._policy_and_value_net_apply,
+            padded_observations,
+            padded_actions,
+            padded_rewards,
+            reward_mask,
+            gamma=self._gamma,
+            lambda_=self._lambda_,
+            c1=self._c1,
+            c2=self._c2,
+            state=self._model_state,
+            rng=key1))
     loss_compute_time = ppo.get_time(loss_compute_start_time)
+    (cur_ppo_loss, cur_value_loss, cur_entropy_bonus) = component_losses
     logging.vlog(
         1,
         "Calculating P&V loss [%10.2f(%10.2f, %10.2f, %10.2f)] took %0.2f msec.",
-        cur_combined_loss, cur_value_loss, cur_ppo_loss, entropy_bonus,
+        cur_combined_loss, cur_ppo_loss, cur_value_loss, cur_entropy_bonus,
         ppo.get_time(loss_compute_start_time))
 
     self._rng, key1 = jax_random.split(self._rng, num=2)
     logging.vlog(1, "Policy and Value Optimization")
     optimization_start_time = time.time()
     keys = jax_random.split(key1, num=self._n_optimizer_steps)
-    for (j, key) in enumerate(keys):
+    opt_step = 0
+    for key in keys:
       k1, k2, k3 = jax_random.split(key, num=3)
       t = time.time()
       # Update the optimizer state.
@@ -378,6 +379,7 @@ def train_epoch(self):
               lambda_=self._lambda_,
               state=self._model_state,
               rng=k1))
+      opt_step += 1
       self._total_opt_step += 1
 
       # Compute the approx KL for early stopping.
@@ -392,38 +394,38 @@ def train_epoch(self):
       early_stopping = approx_kl > 1.5 * self._target_kl
       if early_stopping:
         logging.vlog(
-            1, "Early stopping policy and value optimization at iter: %d, "
-            "with approx_kl: %0.2f", j, approx_kl)
+            1, "Early stopping policy and value optimization after %d steps, "
+            "with approx_kl: %0.2f", opt_step, approx_kl)
         # We don't return right-away, we want the below to execute on the last
         # iteration.
 
       t2 = time.time()
-      if (((j + 1) % self._print_every_optimizer_steps == 0) or
-          (j == self._n_optimizer_steps - 1) or early_stopping):
+      if (opt_step % self._print_every_optimizer_steps == 0 or
+          opt_step == self._n_optimizer_steps or early_stopping):
         # Compute and log the loss.
-        (loss_combined, loss_ppo, loss_value,
-         entropy_bonus), self._model_state = (
-             ppo.combined_loss(
-                 self._policy_and_value_net_params,
-                 log_probabs_traj,
-                 value_predictions_traj,
-                 self._policy_and_value_net_apply,
-                 padded_observations,
-                 padded_actions,
-                 padded_rewards,
-                 reward_mask,
-                 gamma=self._gamma,
-                 lambda_=self._lambda_,
-                 c1=self._c1,
-                 c2=self._c2,
-                 state=self._model_state,
-                 rng=k3))
+        (combined_loss, component_losses, _, self._model_state) = (
+            ppo.combined_loss(
+                self._policy_and_value_net_params,
+                log_probabs_traj,
+                value_predictions_traj,
+                self._policy_and_value_net_apply,
+                padded_observations,
+                padded_actions,
+                padded_rewards,
+                reward_mask,
+                gamma=self._gamma,
+                lambda_=self._lambda_,
+                c1=self._c1,
+                c2=self._c2,
+                state=self._model_state,
+                rng=k3))
         logging.vlog(1, "One Policy and Value grad desc took: %0.2f msec",
                      ppo.get_time(t, t2))
+        (ppo_loss, value_loss, entropy_bonus) = component_losses
         logging.vlog(
             1, "Combined Loss(value, ppo, entropy_bonus) [%10.2f] ->"
-            " [%10.2f(%10.2f,%10.2f,%10.2f)]", cur_combined_loss, loss_combined,
-            loss_value, loss_ppo, entropy_bonus)
+            " [%10.2f(%10.2f,%10.2f,%10.2f)]", cur_combined_loss, combined_loss,
+            ppo_loss, value_loss, entropy_bonus)
 
       if early_stopping:
         break
@@ -432,7 +434,14 @@ def train_epoch(self):
 
     logging.vlog(
         1, "Total Combined Loss reduction [%0.2f]%%",
-        (100 * (cur_combined_loss - loss_combined) / np.abs(cur_combined_loss)))
+        (100 * (cur_combined_loss - combined_loss) / np.abs(cur_combined_loss)))
+
+    summaries.update({
+        "n_optimizer_steps": opt_step,
+        "approx_kl": approx_kl,
+    })
+    for (name, value) in summaries.items():
+      self._train_sw.scalar("train/{}".format(name), value, step=self._epoch)
 
     # Save parameters every time we see the end of at least a fraction of batch
     # number of trajectories that are done (not completed -- completed includes
@@ -453,8 +462,8 @@ def train_epoch(self):
 
     logging.info(
         "PPO epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
-        " Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)]", self._epoch,
-        min_reward, max_reward, avg_reward, loss_combined, loss_value, loss_ppo,
+        " Loss(ppo, value, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)]", self._epoch,
+        min_reward, max_reward, avg_reward, combined_loss, ppo_loss, value_loss,
         entropy_bonus)
 
     timing_dict = {

From 36df324c7bf1b55f06fec2b37eb7b378634aa8a3 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 29 Aug 2019 17:15:05 -0700
Subject: [PATCH 2354/2720] Fix a NaN bug in SimPLe, caused by incomplete
 batches in the training input stream

PiperOrigin-RevId: 266265965
---
 tensor2tensor/trax/rl/simple_trainer.py | 40 +++++++++++++------------
 1 file changed, 21 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index f20060bea..479057afe 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -139,6 +139,20 @@ def _load_trajectories(self, trajectory_dir):
     return (train_trajectories, eval_trajectories)
 
   def _data_stream(self, trajectories, batch_size):
+    def generate_examples():
+      """Creates an infinite stream of shuffled examples."""
+      examples = [
+          example  # pylint: disable=g-complex-comprehension
+          for trajectory_examples in map(
+              self._sim_env.trajectory_to_training_examples, trajectories)
+          for example in trajectory_examples
+      ]
+      assert examples
+      while True:
+        random.shuffle(examples)
+        for example in examples:
+          yield example
+
     def make_batch(examples):
       """Stack a structure of np arrays nested in lists/tuples."""
       assert examples
@@ -148,26 +162,14 @@ def make_batch(examples):
             for i in range(len(examples[0]))
         )
       else:
-        batch = np.stack(examples, axis=0)
-        pad_width = (
-            [(0, batch_size - len(examples))] +
-            [(0, 0)] * (len(batch.shape) - 1)
-        )
-        # Pad with zeros. This doesn't change anything, because we have weights
-        # in the examples.
-        return np.pad(batch, pad_width, mode="constant")
-
-    examples = [
-        example  # pylint: disable=g-complex-comprehension
-        for trajectory_examples in map(
-            self._sim_env.trajectory_to_training_examples, trajectories)
-        for example in trajectory_examples
-    ]
+        return np.stack(examples, axis=0)
+
+    # Take consecutive batches from an infinite stream. This way there are no
+    # incomplete batches. We might get duplicate examples in the same batch, but
+    # that should be very rare.
+    example_stream = generate_examples()
     while True:
-      random.shuffle(examples)
-      for from_index in range(0, len(examples), batch_size):
-        example_list = examples[from_index:(from_index + batch_size)]
-        yield make_batch(example_list)
+      yield make_batch(list(itertools.islice(example_stream, batch_size)))
 
   def train_model(self):
     logging.info("SimPLe epoch [% 6d]: training model.", self._simple_epoch)

From ade39365355afc24129bb99e2d2c8238a06e5011 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 29 Aug 2019 21:16:49 -0700
Subject: [PATCH 2355/2720] Bound spaces in BoxSpaceSerializer

PiperOrigin-RevId: 266298501
---
 tensor2tensor/trax/rl/space_serializer.py     | 25 +++++++++--
 .../trax/rl/space_serializer_test.py          | 44 ++++++++++++-------
 2 files changed, 50 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/trax/rl/space_serializer.py b/tensor2tensor/trax/rl/space_serializer.py
index 220fb7d5b..0f6209148 100644
--- a/tensor2tensor/trax/rl/space_serializer.py
+++ b/tensor2tensor/trax/rl/space_serializer.py
@@ -19,6 +19,9 @@
 from __future__ import division
 from __future__ import print_function
 
+import copy
+
+from absl import logging
 import gin
 import gym
 import numpy as np
@@ -86,7 +89,7 @@ def create(space, vocab_size):
   }[type(space)](space, vocab_size)
 
 
-@gin.configurable(whitelist=["precision"])
+@gin.configurable(blacklist=["space", "vocab_size"])
 class BoxSpaceSerializer(SpaceSerializer):
   """Serializer for gym.spaces.Box.
 
@@ -96,11 +99,25 @@ class BoxSpaceSerializer(SpaceSerializer):
 
   space_type = gym.spaces.Box
 
-  def __init__(self, space, vocab_size, precision):
-    super(BoxSpaceSerializer, self).__init__(space, vocab_size)
-    assert space.is_bounded(), "Only bounded spaces are supported."
+  def __init__(self, space, vocab_size, precision=2, max_range=(-100.0, 100.0)):
     self._precision = precision
 
+    # Some gym envs (e.g. CartPole) have unreasonably high bounds for
+    # observations. We clip so we can represent them.
+    bounded_space = copy.copy(space)
+    (min_low, max_high) = max_range
+    bounded_space.low = np.maximum(space.low, min_low)
+    bounded_space.high = np.minimum(space.high, max_high)
+    if (not np.allclose(bounded_space.low, space.low) or
+        not np.allclose(bounded_space.high, space.high)):
+      logging.warning(
+          "Space limits %s, %s out of bounds %s. Clipping to %s, %s.",
+          str(space.low), str(space.high), str(max_range),
+          str(bounded_space.low), str(bounded_space.high)
+      )
+
+    super(BoxSpaceSerializer, self).__init__(bounded_space, vocab_size)
+
   def serialize(self, data):
     array = data
     batch_size = array.shape[0]
diff --git a/tensor2tensor/trax/rl/space_serializer_test.py b/tensor2tensor/trax/rl/space_serializer_test.py
index 288747c24..5f98b31f3 100644
--- a/tensor2tensor/trax/rl/space_serializer_test.py
+++ b/tensor2tensor/trax/rl/space_serializer_test.py
@@ -29,37 +29,51 @@
 
 class BoxSpaceSerializerTest(test.TestCase):
 
-  def setUp(self):
-    super(BoxSpaceSerializerTest, self).setUp()
+  def _make_space_and_serializer(self, low=-10, high=10, shape=(2,)):
     # Enough precision to represent float32s accurately.
     gin.bind_parameter("BoxSpaceSerializer.precision", 4)
-    self._space = gym.spaces.Box(low=-10, high=10, shape=(2,))
-    self._serializer = space_serializer.create(
-        self._space,
+    space = gym.spaces.Box(low=low, high=high, shape=shape)
+    serializer = space_serializer.create(
+        space,
         # Weird vocab_size to test that it doesn't only work with powers of 2.
         vocab_size=257)
+    return (space, serializer)
 
-  def _sample_batch(self):
-    return np.reshape(self._space.sample(), (1,) + self._space.shape)
+  def _sample_batch(self, space):
+    return np.reshape(space.sample(), (1,) + space.shape)
 
   def test_representation_length(self):
-    input_array = self._sample_batch()
-    representation = self._serializer.serialize(input_array)
+    (space, serializer) = self._make_space_and_serializer()
+    input_array = self._sample_batch(space)
+    representation = serializer.serialize(input_array)
     self.assertEqual(
-        representation.shape, (1, self._serializer.representation_length))
+        representation.shape, (1, serializer.representation_length))
 
   def test_commutes(self):
-    input_array = self._sample_batch()
-    representation = self._serializer.serialize(input_array)
-    output_array = self._serializer.deserialize(representation)
+    (space, serializer) = self._make_space_and_serializer()
+    input_array = self._sample_batch(space)
+    representation = serializer.serialize(input_array)
+    output_array = serializer.deserialize(representation)
     np.testing.assert_array_almost_equal(input_array, output_array)
 
   def test_representation_changes(self):
-    array1 = self._sample_batch()
+    (space, serializer) = self._make_space_and_serializer()
+    array1 = self._sample_batch(space)
     array2 = -array1
-    (repr1, repr2) = tuple(map(self._serializer.serialize, (array1, array2)))
+    (repr1, repr2) = tuple(map(serializer.serialize, (array1, array2)))
     self.assertFalse(np.array_equal(repr1, repr2))
 
+  def test_bounds_space(self):
+    gin.bind_parameter("BoxSpaceSerializer.max_range", (-10.0, 10.0))
+    (_, serializer) = self._make_space_and_serializer(
+        # Too wide range to represent, need to clip.
+        low=-1e18, high=1e18,
+        shape=(1,))
+    input_array = np.array([[1.2345]])
+    representation = serializer.serialize(input_array)
+    output_array = serializer.deserialize(representation)
+    np.testing.assert_array_almost_equal(input_array, output_array)
+
 
 class DiscreteSpaceSerializerTest(test.TestCase):
 

From 324ba280167a4e64210d80aff2e5650183764afb Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 29 Aug 2019 21:22:51 -0700
Subject: [PATCH 2356/2720] Add a dense layer before positional encoding in
 TransformerDecoder

This way the dimensionality of the positional encoding does not depend on the input size, and also the model can learn how to combine the input vectors with the positional encoding. Also added dropout to make the model closer to TransformerLM.

PiperOrigin-RevId: 266299209
---
 tensor2tensor/trax/models/transformer.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 1dc2e0325..edc60ade2 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -182,8 +182,9 @@ def TransformerDecoder(d_model=512,
     a continuous tensor.
   """
   return tl.Model(                  # vecs
-      tl.PositionalEncoding(max_len=max_len),
       tl.Dense(d_model),            # vecs
+      tl.Dropout(rate=dropout, mode=mode),
+      tl.PositionalEncoding(max_len=max_len),
       [DecoderBlock(  # pylint: disable=g-complex-comprehension
           d_model, d_ff, n_heads, d_attention_key, d_attention_value,
           attention_type, dropout, mode)

From d973bc8517b83bf2f135e5f19e8ca0badaf210e9 Mon Sep 17 00:00:00 2001
From: Vinh Nguyen <vinh.nguyenx@gmail.com>
Date: Sat, 31 Aug 2019 02:00:24 +1000
Subject: [PATCH 2357/2720] Refine automatic mixed precision support via hyper
 param (#1681)

* moving gpu_auto_mixed_precision parameter to hparams

* fix param naming

* move gpu automixed precision training to trainer flag

* remove unused os lib

* adding automatic mixed precision support as a hparam

* revert t2t_trainer changes

* adding gpu_automatic_mixed_precision flag to trainer
---
 tensor2tensor/bin/t2t_trainer.py       | 6 +++++-
 tensor2tensor/layers/common_hparams.py | 4 +++-
 tensor2tensor/utils/optimize.py        | 7 ++-----
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index d146f84a7..74173c530 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -141,7 +141,9 @@
 flags.DEFINE_integer("log_step_count_steps", 100,
                      "Number of local steps after which progress is printed "
                      "out")
-
+flags.DEFINE_bool("gpu_automatic_mixed_precision", False,
+                  "Whether to employ GPU automatic mixed precision training "
+                  "(via graph rewrite and dynamic loss scaling).")
 
 
 def set_hparams_from_args(args):
@@ -385,6 +387,8 @@ def main(argv):
     set_hparams_from_args(argv[1:])
   if FLAGS.schedule != "run_std_server":
     hparams = create_hparams()
+  if FLAGS.gpu_automatic_mixed_precision:
+    setattr(hparams, "gpu_automatic_mixed_precision", True)
 
   if FLAGS.schedule == "train" or FLAGS.schedule == "train_eval_and_decode":
     mlperf_log.transformer_print(key=mlperf_log.RUN_START, hparams=hparams)
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 2af2e56f3..16401437f 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -347,7 +347,9 @@ def basic_params1():
       num_area_layers=0,
       max_area_width=1,
       max_area_height=1,
-      memory_height=1
+      memory_height=1,
+      # Whether to use GPU automatic mixed precision (via graph rewrite)
+      gpu_automatic_mixed_precision=False
   )
 
 
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 1df3fbbdd..cf1150e4b 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -18,7 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import os
 import numpy as np
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import adafactor as adafactor_lib
@@ -45,8 +44,7 @@ def optimize(loss,
              learning_rate,
              hparams,
              use_tpu=False,
-             variables=None,
-             gpu_auto_mixed_precision=False):
+             variables=None):
   """Minimize loss."""
   loss = weight_decay_and_noise(loss, hparams, learning_rate)
   loss = tf.identity(loss, name="total_loss")
@@ -71,8 +69,7 @@ def optimize(loss,
   opt = ConditionalOptimizer(hparams.optimizer, learning_rate, hparams, use_tpu)
   if use_tpu:
     opt = tf.contrib.tpu.CrossShardOptimizer(opt)
-  if gpu_auto_mixed_precision or os.environ.get(
-      "TF_ENABLE_AUTO_MIXED_PRECISION", "0") == "1":
+  if getattr(hparams, "gpu_automatic_mixed_precision", False):
     if use_tpu:
       raise RuntimeError("GPU auto mixed precision cannot be used with TPU")
     elif _mixed_precision_is_enabled(hparams):

From 1a2542ed09421aab5529d1f0d5a3a1333a07174c Mon Sep 17 00:00:00 2001
From: Vinh Nguyen <vinh.nguyenx@gmail.com>
Date: Fri, 30 Aug 2019 09:00:47 -0700
Subject: [PATCH 2358/2720] Merge of PR #1681

PiperOrigin-RevId: 266390503
---
 tensor2tensor/bin/t2t_trainer.py       | 1 +
 tensor2tensor/layers/common_hparams.py | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 74173c530..0d7c7529a 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -146,6 +146,7 @@
                   "(via graph rewrite and dynamic loss scaling).")
 
 
+
 def set_hparams_from_args(args):
   """Set hparams overrides from unparsed args list."""
   if not args:
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 16401437f..eb5d687a0 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -349,7 +349,7 @@ def basic_params1():
       max_area_height=1,
       memory_height=1,
       # Whether to use GPU automatic mixed precision (via graph rewrite)
-      gpu_automatic_mixed_precision=False
+      gpu_automatic_mixed_precision=False,
   )
 
 

From ab6cf3b75c7e68cf65b9f13ffe0a2f300a9f618f Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 30 Aug 2019 10:42:34 -0700
Subject: [PATCH 2359/2720] Rename flag tf_eager to enable_eager_execution to
 reduce cognitive load, this is also used elsewhere in the code.

PiperOrigin-RevId: 266411414
---
 tensor2tensor/trax/trainer.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index 1eb71b950..c274b6b9f 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -47,7 +47,8 @@
                           "Configuration parameters (gin string).")
 flags.DEFINE_integer("log_level", logging.INFO, "Log level.")
 flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
-flags.DEFINE_bool("tf_eager", True, "Whether we're running TF in eager mode.")
+flags.DEFINE_bool("enable_eager_execution", True,
+                  "Whether we're running TF in eager mode.")
 flags.DEFINE_bool("tf_xla", True, "Whether to turn on XLA for TF.")
 flags.DEFINE_bool("tf_opt_pin_to_host", False, "Whether to turn on TF "
                   "pin-to-host optimization.")
@@ -94,7 +95,7 @@ def _setup_gin():
 def main(_):
   logging.set_verbosity(FLAGS.log_level)
 
-  if FLAGS.tf_eager:
+  if FLAGS.enable_eager_execution:
     tf.enable_eager_execution()
 
   if FLAGS.tf_xla:

From 595249ae667c689ffc60d27156ba36ec1670f5e2 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 30 Aug 2019 14:18:32 -0700
Subject: [PATCH 2360/2720] Small corrections: self.sublayers() ->
 self.sublayers since it's a property now.

PiperOrigin-RevId: 266454473
---
 tensor2tensor/trax/layers/reversible.py                  | 6 +++---
 tensor2tensor/trax/models/research/transformer_revnet.py | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/trax/layers/reversible.py b/tensor2tensor/trax/layers/reversible.py
index 98038b8c0..1c70d2270 100644
--- a/tensor2tensor/trax/layers/reversible.py
+++ b/tensor2tensor/trax/layers/reversible.py
@@ -88,7 +88,7 @@ def __init__(self, *layers):
     super(ReversibleSerial, self).__init__(*layers)
 
     # Note that sublayers has already been flattened to remove nested lists.
-    for i, layer in enumerate(self.sublayers()):
+    for i, layer in enumerate(self.sublayers):
       if not isinstance(layer, ReversibleLayer):
         raise ValueError(
             'Sub-layer {} of ReversibleSerial is not reversible: {}'.format(
@@ -101,7 +101,7 @@ def reverse(self, output, params=(), state=(), **kwargs):
       rngs = backend.random.split(rng, self._n_layers)
 
     layer_val = output
-    for layer, p, s, rng in reversed(zip(self.sublayers(),
+    for layer, p, s, rng in reversed(zip(self.sublayers,
                                          params, state, rngs)):
       layer_val = layer.reverse(layer_val, p, s, rng=rng, **kwargs)
 
@@ -116,7 +116,7 @@ def reverse_and_grad(self, output, ct, params=(), state=(), **kwargs):
     layer_val = output
     layer_ct = ct
     params_ct = []
-    for layer, p, s, rng in reversed(zip(self.sublayers(),
+    for layer, p, s, rng in reversed(zip(self.sublayers,
                                          params, state, rngs)):
       layer_val, layer_ct = layer.reverse_and_grad(
           layer_val, layer_ct, p, s, rng=rng, **kwargs)
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index b98a3101b..c64052873 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -223,7 +223,7 @@ def reverse(self, output, params=(), state=(), **kwargs):
     rngs = (None,) * self._n_layers
     if rng is not None:
       rngs = backend.random.split(rng, self._n_layers)
-    # Note that self.sublayers() aligns exactly with self.reverse_layers in
+    # Note that self.sublayers aligns exactly with self.reverse_layers in
     # terms of parameter and rng usage, so no re-ordering is required.
     for layer, p, s, rng in zip(self.reverse_layers, params, state, rngs):
       reconstructed_x, _ = layer(reconstructed_x, p, s, rng=rng, **kwargs)
@@ -324,7 +324,7 @@ def reverse(self, output, params=(), state=(), **kwargs):
       rngs = backend.random.split(rng, self._n_layers)
 
     reconstructed_x = output
-    # Note that self.sublayers() aligns exactly with self.reverse_layers in
+    # Note that self.sublayers aligns exactly with self.reverse_layers in
     # terms of parameter and rng usage, so no re-ordering is required.
     for layer, p, s, rng in zip(self.reverse_layers, params, state, rngs):
       reconstructed_x, _ = layer.reverse(reconstructed_x, p, s, rng=rng,

From 90db045a32f31ce690b4689a99e255063f2172e3 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 30 Aug 2019 14:42:01 -0700
Subject: [PATCH 2361/2720] Add tests for the SimPLe training example
 serialization

PiperOrigin-RevId: 266459008
---
 .../trax/rl/simulated_env_problem.py          |   1 +
 .../trax/rl/simulated_env_problem_test.py     | 123 ++++++++++++++++--
 2 files changed, 111 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index 07703ee9c..a31b8df76 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -446,6 +446,7 @@ def concat_and_pad(arrays):
           mode="constant",
       )
     (reprs, weights) = map(concat_and_pad, (reprs, weights))
+    reprs = reprs.astype(self.model_input_dtype)
     return [(reprs, reprs, weights)]  # (inputs, targets, weights)
 
   @property
diff --git a/tensor2tensor/trax/rl/simulated_env_problem_test.py b/tensor2tensor/trax/rl/simulated_env_problem_test.py
index f824bf741..38f1a37f4 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem_test.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem_test.py
@@ -26,6 +26,7 @@
 import mock
 import numpy as np
 
+from tensor2tensor.envs import trajectory
 from tensor2tensor.trax import backend
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.rl import simulated_env_problem
@@ -121,6 +122,37 @@ def test_takes_new_history(self):
 
 class SerializedSequenceSimulatedEnvProblemTest(test.TestCase):
 
+  def _make_env(
+      self, observation_space, action_space, vocab_size,
+      predict_fn=None, reward_fn=None, done_fn=None,
+      batch_size=None, max_trajectory_length=None,
+  ):
+    mock_model_fn = mock.MagicMock()
+    if predict_fn is not None:
+      mock_model_fn.return_value = predict_fn
+    return simulated_env_problem.SerializedSequenceSimulatedEnvProblem(
+        model=mock_model_fn,
+        reward_fn=reward_fn,
+        done_fn=done_fn,
+        vocab_size=vocab_size,
+        max_trajectory_length=3,
+        batch_size=batch_size,
+        observation_space=observation_space,
+        action_space=action_space,
+        reward_range=(-1, 1),
+        discrete_rewards=False,
+        history_stream=itertools.repeat(None),
+        output_dir=None,
+    )
+
+  def _make_trajectory(self, observations, actions):
+    assert len(observations) == len(actions) + 1
+    t = trajectory.Trajectory()
+    for (obs, act) in zip(observations, actions):
+      t.add_time_step(observation=obs, action=act, done=False)
+    t.add_time_step(observation=observations[-1], done=True)
+    return t
+
   @mock.patch.object(trax, "restore_state", autospec=True)
   def test_communicates_with_model(self, mock_restore_state):
     gin.bind_parameter("BoxSpaceSerializer.precision", 1)
@@ -138,33 +170,28 @@ def make_prediction(symbol):
       # (4 obs symbols + 1 action symbol) * 3 timesteps = 15.
       return np.array([[log_probs] * 15]), ()
 
-    mock_model_fn = mock.MagicMock()
-    mock_model = mock_model_fn.return_value
-    mock_model.side_effect = map(make_prediction, symbols)
+    mock_predict_fn = mock.MagicMock()
+    mock_predict_fn.side_effect = map(make_prediction, symbols)
 
     with backend.use_backend("numpy"):
       # (model_params, opt_state)
       mock_restore_state.return_value.params = (None, None)
-      env = simulated_env_problem.SerializedSequenceSimulatedEnvProblem(
-          model=mock_model_fn,
+      env = self._make_env(
+          predict_fn=mock_predict_fn,
           reward_fn=(lambda _1, _2: np.array([0.5])),
           done_fn=(lambda _1, _2: np.array([False])),
           vocab_size=vocab_size,
-          max_trajectory_length=3,
           batch_size=1,
+          max_trajectory_length=3,
           observation_space=gym.spaces.Box(low=0, high=5, shape=(4,)),
           action_space=gym.spaces.Discrete(2),
-          reward_range=(-1, 1),
-          discrete_rewards=False,
-          history_stream=itertools.repeat(None),
-          output_dir=None,
       )
       obs1 = env.reset()
-      ((inputs,), _) = mock_model.call_args
+      ((inputs,), _) = mock_predict_fn.call_args
 
       act1 = 0
       (obs2, reward, done, _) = env.step(np.array([act1]))
-      ((inputs,), _) = mock_model.call_args
+      ((inputs,), _) = mock_predict_fn.call_args
       self.assertEqual(inputs[0, 4], act1)
       np.testing.assert_array_equal(inputs[0, :4], symbols[:4])
       np.testing.assert_array_equal(obs1, obs2)
@@ -173,13 +200,83 @@ def make_prediction(symbol):
 
       act2 = 1
       (obs3, reward, done, _) = env.step(np.array([act2]))
-      ((inputs,), _) = mock_model.call_args
+      ((inputs,), _) = mock_predict_fn.call_args
       self.assertEqual(inputs[0, 9], act2)
       np.testing.assert_array_equal(inputs[0, 5:9], symbols[4:8])
       self.assertFalse(np.array_equal(obs2, obs3))
       np.testing.assert_array_equal(reward, [0.5])
       np.testing.assert_array_equal(done, [True])
 
+  def test_makes_training_example(self):
+    env = self._make_env(
+        vocab_size=2,
+        observation_space=gym.spaces.Discrete(2),
+        action_space=gym.spaces.Discrete(2),
+        max_trajectory_length=3,
+    )
+    t = self._make_trajectory(observations=[0, 1, 0], actions=[1, 0])
+    examples = env.trajectory_to_training_examples(t)
+
+    # There should be 1 example with the whole trajectory.
+    self.assertEqual(len(examples), 1)
+    [(inputs, targets, weights)] = examples
+    # inputs == targets for autoregressive sequence prediction.
+    np.testing.assert_array_equal(inputs, targets)
+    # Assert array shapes and datatypes.
+    self.assertEqual(inputs.shape, env.model_input_shape)
+    self.assertEqual(inputs.dtype, env.model_input_dtype)
+    self.assertEqual(weights.shape, env.model_input_shape)
+    # Actions should be masked out.
+    self.assertEqual(np.min(weights), 0)
+    # At least part of the observation should have full weight.
+    self.assertEqual(np.max(weights), 1)
+
+  def test_makes_training_examples_from_trajectories_of_different_lengths(self):
+    env = self._make_env(
+        vocab_size=2,
+        observation_space=gym.spaces.Discrete(2),
+        action_space=gym.spaces.Discrete(2),
+        max_trajectory_length=3,
+    )
+    t1 = self._make_trajectory(observations=[0, 1], actions=[1])
+    [(x1, _, w1)] = env.trajectory_to_training_examples(t1)
+    t2 = self._make_trajectory(observations=[0, 1, 0], actions=[1, 0])
+    [(x2, _, w2)] = env.trajectory_to_training_examples(t2)
+
+    # Examples should be padded to the same shape.
+    self.assertEqual(x1.shape, x2.shape)
+    self.assertEqual(w1.shape, w2.shape)
+    # Cumulative weight should increase with trajectory length.
+    self.assertGreater(np.sum(w2), np.sum(w1))
+
+  def test_masked_representation_changes_with_observation(self):
+    env = self._make_env(
+        vocab_size=2,
+        observation_space=gym.spaces.Discrete(2),
+        action_space=gym.spaces.Discrete(2),
+        max_trajectory_length=3,
+    )
+    t1 = self._make_trajectory(observations=[0, 1], actions=[1])
+    [(x1, _, w1)] = env.trajectory_to_training_examples(t1)
+    t2 = self._make_trajectory(observations=[0, 0], actions=[1])
+    [(x2, _, w2)] = env.trajectory_to_training_examples(t2)
+
+    self.assertFalse(np.array_equal(x1 * w1, x2 * w2))
+
+  def test_masked_representation_doesnt_change_with_action(self):
+    env = self._make_env(
+        vocab_size=2,
+        observation_space=gym.spaces.Discrete(2),
+        action_space=gym.spaces.Discrete(2),
+        max_trajectory_length=3,
+    )
+    t1 = self._make_trajectory(observations=[0, 1], actions=[1])
+    [(x1, _, w1)] = env.trajectory_to_training_examples(t1)
+    t2 = self._make_trajectory(observations=[0, 1], actions=[0])
+    [(x2, _, w2)] = env.trajectory_to_training_examples(t2)
+
+    np.testing.assert_array_equal(x1 * w1, x2 * w2)
+
 
 if __name__ == "__main__":
   test.main()

From e47dc6583ff16f6e9d6b0e234b3e2748f5900020 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 30 Aug 2019 15:27:22 -0700
Subject: [PATCH 2362/2720] s/assertTrue(isinstance(o,
 CLS))/assertIsInstance(a, CLS)/g

PiperOrigin-RevId: 266467537
---
 .../data_generators/image_utils_test.py       |  2 +-
 tensor2tensor/envs/gym_env_problem_test.py    |  2 +-
 tensor2tensor/rl/gym_utils_test.py            | 10 +++---
 tensor2tensor/utils/hparam_test.py            | 36 +++++++++----------
 tensor2tensor/utils/registry_test.py          |  2 +-
 5 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/tensor2tensor/data_generators/image_utils_test.py b/tensor2tensor/data_generators/image_utils_test.py
index dfdf843a0..0dd020c24 100644
--- a/tensor2tensor/data_generators/image_utils_test.py
+++ b/tensor2tensor/data_generators/image_utils_test.py
@@ -144,7 +144,7 @@ def testConvertPredictionsToImageSummaries(self):
           decode_hooks)
       self.assertEqual(len(summaries), summaries_length)
       if summaries:
-        self.assertTrue(isinstance(summaries[0], tf.Summary.Value))
+        self.assertIsInstance(summaries[0], tf.Summary.Value)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/envs/gym_env_problem_test.py b/tensor2tensor/envs/gym_env_problem_test.py
index fe88d9b04..d75e81645 100644
--- a/tensor2tensor/envs/gym_env_problem_test.py
+++ b/tensor2tensor/envs/gym_env_problem_test.py
@@ -50,7 +50,7 @@ def test_setup(self):
 
     # Expectations on the observation space.
     observation_space = ep.observation_space
-    self.assertTrue(isinstance(observation_space, Box))
+    self.assertIsInstance(observation_space, Box)
     self.assertEqual(observation_space.shape, (4,))
     self.assertEqual(observation_space.dtype, np.float32)
 
diff --git a/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index 007010e10..5d3d46f16 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -67,19 +67,19 @@ class GymUtilsTest(tf.test.TestCase):
   # Just make an environment and expect to get one.
   def test_making_simple_env(self):
     env = gym_utils.make_gym_env("CartPole-v0")
-    self.assertTrue(isinstance(env, gym.Env))
+    self.assertIsInstance(env, gym.Env)
 
   # Make a time-wrapped environment and expect to get one.
   def test_making_timewrapped_env(self):
     env = gym_utils.make_gym_env("CartPole-v0", rl_env_max_episode_steps=1000)
-    self.assertTrue(isinstance(env, gym.Env))
-    self.assertTrue(isinstance(env, gym.wrappers.TimeLimit))
+    self.assertIsInstance(env, gym.Env)
+    self.assertIsInstance(env, gym.wrappers.TimeLimit)
     self.assertEqual(1000, env._max_episode_steps)
 
   # Make an instance of the environment without a TimeLimit
   def test_unlimited_env(self):
     env = gym_utils.make_gym_env("CartPole-v0", rl_env_max_episode_steps=None)
-    self.assertTrue(isinstance(env, gym.Env))
+    self.assertIsInstance(env, gym.Env)
     self.assertNotIsInstance(env, gym.wrappers.TimeLimit)
 
   def test_rendered_env(self):
@@ -99,7 +99,7 @@ def test_gym_registration(self):
     self.assertEqual("T2TEnv-SimpleEnv-v0", reg_id)
 
     # Most basic check.
-    self.assertTrue(isinstance(env, gym.Env))
+    self.assertIsInstance(env, gym.Env)
 
     # Just make sure we got the same environment.
     self.assertTrue(
diff --git a/tensor2tensor/utils/hparam_test.py b/tensor2tensor/utils/hparam_test.py
index 85790ded2..bea9ee36f 100644
--- a/tensor2tensor/utils/hparam_test.py
+++ b/tensor2tensor/utils/hparam_test.py
@@ -197,7 +197,7 @@ def testParseValuesWithIndexAssigment1(self):
     """Assignment to an index position."""
     parse_dict = hparam.parse_values('arr[1]=10', {'arr': int})
     self.assertEqual(len(parse_dict), 1)
-    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertIsInstance(parse_dict['arr'], dict)
     self.assertDictEqual(parse_dict['arr'], {1: 10})
 
   def testParseValuesWithIndexAssigment1_IgnoreUnknown(self):
@@ -205,14 +205,14 @@ def testParseValuesWithIndexAssigment1_IgnoreUnknown(self):
     parse_dict = hparam.parse_values(
         'arr[1]=10,b=5', {'arr': int}, ignore_unknown=True)
     self.assertEqual(len(parse_dict), 1)
-    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertIsInstance(parse_dict['arr'], dict)
     self.assertDictEqual(parse_dict['arr'], {1: 10})
 
   def testParseValuesWithIndexAssigment2(self):
     """Assignment to multiple index positions."""
     parse_dict = hparam.parse_values('arr[0]=10,arr[5]=20', {'arr': int})
     self.assertEqual(len(parse_dict), 1)
-    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertIsInstance(parse_dict['arr'], dict)
     self.assertDictEqual(parse_dict['arr'], {0: 10, 5: 20})
 
   def testParseValuesWithIndexAssigment2_IgnoreUnknown(self):
@@ -220,7 +220,7 @@ def testParseValuesWithIndexAssigment2_IgnoreUnknown(self):
     parse_dict = hparam.parse_values(
         'arr[0]=10,arr[5]=20,foo=bar', {'arr': int}, ignore_unknown=True)
     self.assertEqual(len(parse_dict), 1)
-    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertIsInstance(parse_dict['arr'], dict)
     self.assertDictEqual(parse_dict['arr'], {0: 10, 5: 20})
 
   def testParseValuesWithIndexAssigment3(self):
@@ -229,9 +229,9 @@ def testParseValuesWithIndexAssigment3(self):
                                      {'arr': int,
                                       'L': int})
     self.assertEqual(len(parse_dict), 2)
-    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertIsInstance(parse_dict['arr'], dict)
     self.assertDictEqual(parse_dict['arr'], {0: 10, 1: 20})
-    self.assertTrue(isinstance(parse_dict['L'], dict))
+    self.assertIsInstance(parse_dict['L'], dict)
     self.assertDictEqual(parse_dict['L'], {5: 100, 10: 200})
 
   def testParseValuesWithIndexAssigment3_IgnoreUnknown(self):
@@ -240,9 +240,9 @@ def testParseValuesWithIndexAssigment3_IgnoreUnknown(self):
         'arr[0]=10,C=5,arr[1]=20,B[0]=kkk,L[5]=100,L[10]=200',
         {'arr': int, 'L': int}, ignore_unknown=True)
     self.assertEqual(len(parse_dict), 2)
-    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertIsInstance(parse_dict['arr'], dict)
     self.assertDictEqual(parse_dict['arr'], {0: 10, 1: 20})
-    self.assertTrue(isinstance(parse_dict['L'], dict))
+    self.assertIsInstance(parse_dict['L'], dict)
     self.assertDictEqual(parse_dict['L'], {5: 100, 10: 200})
 
   def testParseValuesWithIndexAssigment4(self):
@@ -252,7 +252,7 @@ def testParseValuesWithIndexAssigment4(self):
                                       'y': int,
                                       'arr': int})
     self.assertEqual(len(parse_dict), 3)
-    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertIsInstance(parse_dict['arr'], dict)
     self.assertDictEqual(parse_dict['arr'], {1: 20})
     self.assertEqual(parse_dict['x'], 10)
     self.assertEqual(parse_dict['y'], 30)
@@ -263,7 +263,7 @@ def testParseValuesWithIndexAssigment4_IgnoreUnknown(self):
         'x=10,foo[0]=bar,arr[1]=20,zzz=78,y=30',
         {'x': int, 'y': int, 'arr': int}, ignore_unknown=True)
     self.assertEqual(len(parse_dict), 3)
-    self.assertTrue(isinstance(parse_dict['arr'], dict))
+    self.assertIsInstance(parse_dict['arr'], dict)
     self.assertDictEqual(parse_dict['arr'], {1: 20})
     self.assertEqual(parse_dict['x'], 10)
     self.assertEqual(parse_dict['y'], 30)
@@ -277,13 +277,13 @@ def testParseValuesWithIndexAssigment5(self):
         'd': float
     })
     self.assertEqual(set(parse_dict.keys()), {'a', 'b', 'c', 'd'})
-    self.assertTrue(isinstance(parse_dict['a'], dict))
+    self.assertIsInstance(parse_dict['a'], dict)
     self.assertDictEqual(parse_dict['a'], {0: 5})
-    self.assertTrue(isinstance(parse_dict['b'], dict))
+    self.assertIsInstance(parse_dict['b'], dict)
     self.assertDictEqual(parse_dict['b'], {1: True})
-    self.assertTrue(isinstance(parse_dict['c'], dict))
+    self.assertIsInstance(parse_dict['c'], dict)
     self.assertDictEqual(parse_dict['c'], {2: 'abc'})
-    self.assertTrue(isinstance(parse_dict['d'], dict))
+    self.assertIsInstance(parse_dict['d'], dict)
     self.assertDictEqual(parse_dict['d'], {3: 3.14})
 
   def testParseValuesWithIndexAssigment5_IgnoreUnknown(self):
@@ -293,13 +293,13 @@ def testParseValuesWithIndexAssigment5_IgnoreUnknown(self):
         {'a': int, 'b': bool, 'c': str, 'd': float},
         ignore_unknown=True)
     self.assertEqual(set(parse_dict.keys()), {'a', 'b', 'c', 'd'})
-    self.assertTrue(isinstance(parse_dict['a'], dict))
+    self.assertIsInstance(parse_dict['a'], dict)
     self.assertDictEqual(parse_dict['a'], {0: 5})
-    self.assertTrue(isinstance(parse_dict['b'], dict))
+    self.assertIsInstance(parse_dict['b'], dict)
     self.assertDictEqual(parse_dict['b'], {1: True})
-    self.assertTrue(isinstance(parse_dict['c'], dict))
+    self.assertIsInstance(parse_dict['c'], dict)
     self.assertDictEqual(parse_dict['c'], {2: 'abc'})
-    self.assertTrue(isinstance(parse_dict['d'], dict))
+    self.assertIsInstance(parse_dict['d'], dict)
     self.assertDictEqual(parse_dict['d'], {3: 3.14})
 
   def testParseValuesWithBadIndexAssigment1(self):
diff --git a/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index e50d1ce6a..adf4bfccd 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -117,7 +117,7 @@ def initialize(self, batch_size):
     self.assertEqual(batch_size, ep.batch_size)
 
     # assert on the type.
-    self.assertTrue(isinstance(ep, EnvProb))
+    self.assertIsInstance(ep, EnvProb)
 
 
 class ModelRegistryTest(tf.test.TestCase):

From 309511cae54aee1b72d88ba8c9513c50cb5e398e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 30 Aug 2019 15:41:48 -0700
Subject: [PATCH 2363/2720] Rename another usage of the tf_eager flag

PiperOrigin-RevId: 266470013
---
 tensor2tensor/trax/trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index c274b6b9f..7801f4711 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -111,7 +111,7 @@ def main(_):
 
   _setup_gin()
 
-  if FLAGS.tf_eager and backend.get_name() in ("numpy", "jax"):
+  if FLAGS.enable_eager_execution and backend.get_name() in ("numpy", "jax"):
     # Numpy backend doesn't benefit from having the input pipeline run on GPU,
     # and jax backend has GPU memory contention if TF uses the GPU. Gin must be
     # set up first before determining the backend.

From 39959804676e31dc14043acd1c58940abea76c4b Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 30 Aug 2019 15:50:38 -0700
Subject: [PATCH 2364/2720] Change env_problem and gym_utils to use
 absl.logging instead of tf's logging.

PiperOrigin-RevId: 266471505
---
 tensor2tensor/envs/env_problem.py | 9 +++++----
 tensor2tensor/rl/gym_utils.py     | 6 +++---
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 0fc27d9ce..0f10b1b78 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -23,6 +23,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from absl import logging
 from gym.core import Env
 import numpy as np
 import six
@@ -277,10 +278,10 @@ def num_rewards(self):
     # Pre-conditions: reward range is finite.
     #               : processed rewards are discrete.
     if not self.is_reward_range_finite:
-      tf.logging.warn("Infinite reward range, `num_rewards returning None`")
+      logging.warn("Infinite reward range, `num_rewards returning None`")
       return None
     if not self.is_processed_rewards_discrete:
-      tf.logging.warn(
+      logging.warn(
           "Processed rewards are not discrete, `num_rewards` returning None")
       return None
 
@@ -360,7 +361,7 @@ def reset(self, indices=None):
 
     # If this is empty (not None) then don't do anything, no env was done.
     if indices.size == 0:
-      tf.logging.warning(
+      logging.warning(
           "`reset` called with empty indices array, this is a no-op.")
       return None
 
@@ -617,7 +618,7 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
     num_completed_trajectories = self.trajectories.num_completed_trajectories
     num_shards = len(files_list)
     if num_completed_trajectories < num_shards:
-      tf.logging.warning(
+      logging.warning(
           "Number of completed trajectories [%d] is less than "
           "the number of shards [%d], some shards maybe empty.",
           num_completed_trajectories, num_shards)
diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index d1c7bada3..f0fc2e219 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -19,11 +19,11 @@
 from __future__ import division
 from __future__ import print_function
 
+from absl import logging
 import gym
 import gym.wrappers
 import numpy as np
 from PIL import Image
-import tensorflow as tf
 
 
 class StickyActionEnv(gym.Wrapper):
@@ -232,7 +232,7 @@ def register_gym_env(class_entry_point, version="v0", kwargs=None):
   env_name = "T2TEnv-{}-{}".format(class_name, version)
   gym.envs.register(id=env_name, entry_point=class_entry_point, kwargs=kwargs)
 
-  tf.logging.info("Entry Point [%s] registered with id [%s]", class_entry_point,
-                  env_name)
+  logging.info(
+      "Entry Point [%s] registered with id [%s]", class_entry_point, env_name)
 
   return env_name, gym.make(env_name)

From 439ad6855fe4375277d8653416b6b1602c315f84 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Sat, 31 Aug 2019 18:11:13 -0700
Subject: [PATCH 2365/2720] Allow starting the SimPLe loop from training the
 model on precollected data.

In later epochs, mix new data with precollected data. Also added some tests and moved the data pipeline to a separate module so it can be tested independently.

PiperOrigin-RevId: 266594480
---
 tensor2tensor/trax/rl/simple.py              |  92 ++++++++++++++
 tensor2tensor/trax/rl/simple_test.py         | 107 ++++++++++++++++
 tensor2tensor/trax/rl/simple_trainer.py      | 121 +++++++++----------
 tensor2tensor/trax/rl/simple_trainer_test.py |   2 +-
 4 files changed, 257 insertions(+), 65 deletions(-)
 create mode 100644 tensor2tensor/trax/rl/simple.py
 create mode 100644 tensor2tensor/trax/rl/simple_test.py

diff --git a/tensor2tensor/trax/rl/simple.py b/tensor2tensor/trax/rl/simple.py
new file mode 100644
index 000000000..152d346cd
--- /dev/null
+++ b/tensor2tensor/trax/rl/simple.py
@@ -0,0 +1,92 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 helper functions."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import itertools
+import os
+import random
+
+import cloudpickle as pickle
+import numpy as np
+from tensorflow.io import gfile
+
+
+def load_trajectories(trajectory_dir, eval_frac):
+  """Loads trajectories from a possibly nested directory of pickles."""
+  train_trajectories = []
+  eval_trajectories = []
+  # Search the entire directory subtree for trajectories.
+  for (subdir, _, filenames) in gfile.walk(trajectory_dir):
+    for filename in filenames:
+      shard_path = os.path.join(subdir, filename)
+      with gfile.GFile(shard_path, "rb") as f:
+        trajectories = pickle.load(f)
+        pivot = int(len(trajectories) * (1 - eval_frac))
+        train_trajectories.extend(trajectories[:pivot])
+        eval_trajectories.extend(trajectories[pivot:])
+  assert train_trajectories, "Haven't found any training data."
+  assert eval_trajectories, "Haven't found any evaluation data."
+  return (train_trajectories, eval_trajectories)
+
+
+def generate_examples(trajectories, trajectory_to_training_examples_fn):
+  """Generates an infinite stream of shuffled examples out of trajectories."""
+  examples = [
+      example  # pylint: disable=g-complex-comprehension
+      for trajectory_examples in map(
+          trajectory_to_training_examples_fn, trajectories)
+      for example in trajectory_examples
+  ]
+  assert examples
+  while True:
+    random.shuffle(examples)
+    for example in examples:
+      yield example
+
+
+def mix_streams(stream1, stream2, mix_prob):
+  """Mixes two streams together with a fixed probability."""
+  while True:
+    # In the corner cases (mix_prob = 0 or 1) mixing the other stream never
+    # happens, because random() samples from the semi-open interval [0, 1).
+    if random.random() < mix_prob:
+      yield next(stream1)
+    else:
+      yield next(stream2)
+
+
+def batch_stream(stream, batch_size):
+  """Batches a stream of training examples."""
+  def make_batch(examples):
+    """Stacks a structure of numpy arrays nested in lists/tuples."""
+    assert examples
+    if isinstance(examples[0], (list, tuple)):
+      return type(examples[0])(
+          make_batch([example[i] for example in examples])
+          for i in range(len(examples[0]))
+      )
+    else:
+      return np.stack(examples, axis=0)
+
+  # Take consecutive batches from an infinite stream. This way there are no
+  # incomplete batches. We might get duplicate examples in the same batch, but
+  # that should be very rare.
+  while True:
+    yield make_batch(list(itertools.islice(stream, batch_size)))
diff --git a/tensor2tensor/trax/rl/simple_test.py b/tensor2tensor/trax/rl/simple_test.py
new file mode 100644
index 000000000..b44b6853d
--- /dev/null
+++ b/tensor2tensor/trax/rl/simple_test.py
@@ -0,0 +1,107 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.rl.simple."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import itertools
+import os
+
+import cloudpickle as pickle
+import numpy as np
+from tensor2tensor.envs import trajectory
+from tensor2tensor.trax.rl import simple
+from tensorflow import test
+from tensorflow.io import gfile
+
+
+class SimpleTest(test.TestCase):
+
+  def _make_singleton_trajectory(self, observation):
+    t = trajectory.Trajectory()
+    t.add_time_step(observation=observation)
+    return t
+
+  def _dump_trajectory_pickle(self, observations, path):
+    trajectories = list(map(self._make_singleton_trajectory, observations))
+    with gfile.GFile(path, "wb") as f:
+      pickle.dump(trajectories, f)
+
+  def test_loads_trajectories(self):
+    temp_dir = self.get_temp_dir()
+    # Dump two trajectory pickles with given observations.
+    self._dump_trajectory_pickle(
+        observations=[0, 1, 2, 3], path=os.path.join(temp_dir, "0.pkl"))
+    self._dump_trajectory_pickle(
+        observations=[4, 5, 6, 7], path=os.path.join(temp_dir, "1.pkl"))
+    (train_trajs, eval_trajs) = simple.load_trajectories(
+        temp_dir, eval_frac=0.25)
+    extract_obs = lambda t: t.last_time_step.observation
+    # The order of pickles is undefined, so we compare sets.
+    actual_train_obs = set(map(extract_obs, train_trajs))
+    actual_eval_obs = set(map(extract_obs, eval_trajs))
+
+    # First 3 trajectories from each pickle go to train, the last one to eval.
+    expected_train_obs = {0, 1, 2, 4, 5, 6}
+    expected_eval_obs = {3, 7}
+    self.assertEqual(actual_train_obs, expected_train_obs)
+    self.assertEqual(actual_eval_obs, expected_eval_obs)
+
+  def test_generates_examples(self):
+    observations = [0, 1, 2, 3]
+    trajectories = map(self._make_singleton_trajectory, observations)
+    trajectory_to_training_examples = lambda t: [t.last_time_step.observation]
+    stream = simple.generate_examples(
+        trajectories, trajectory_to_training_examples)
+
+    # The examples are shuffled, so we compare sets.
+    self.assertEqual(
+        set(itertools.islice(stream, len(observations))), set(observations))
+    # The stream is infinite, so we should be able to take a next element.
+    self.assertIn(next(stream), observations)
+
+  def test_mixes_streams_with_prob_one(self):
+    # Mix infinite streams of 0s and 1s.
+    stream = simple.mix_streams(
+        itertools.repeat(0), itertools.repeat(1), mix_prob=1.0)
+    # Mixed stream should have only 0s.
+    self.assertEqual(set(itertools.islice(stream, 100)), {0})
+
+  def test_mixes_streams_with_prob_zero(self):
+    stream = simple.mix_streams(
+        itertools.repeat(0), itertools.repeat(1), mix_prob=0.0)
+    # Mixed stream should have only 1s.
+    self.assertEqual(set(itertools.islice(stream, 100)), {1})
+
+  def test_mixes_streams_with_prob_half(self):
+    stream = simple.mix_streams(
+        itertools.repeat(0), itertools.repeat(1), mix_prob=0.5)
+    # Mixed stream should have both 0s and 1s.
+    self.assertEqual(set(itertools.islice(stream, 100)), {0, 1})
+
+  def test_batches_stream(self):
+    stream = iter([(0, 1), (2, 3), (4, 5), (6, 7)])
+    batched_stream = simple.batch_stream(stream, batch_size=2)
+    np.testing.assert_equal(
+        next(batched_stream), (np.array([0, 2]), np.array([1, 3])))
+    np.testing.assert_equal(
+        next(batched_stream), (np.array([4, 6]), np.array([5, 7])))
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index 479057afe..450a87f1a 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -19,18 +19,16 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
 import itertools
 import os
-import random
 
 from absl import logging
-import cloudpickle as pickle
-import numpy as np
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.rl import base_trainer
+from tensor2tensor.trax.rl import simple
 from tensor2tensor.trax.rl import simulated_env_problem
-from tensorflow.io import gfile
 
 
 class SimPLe(base_trainer.BaseTrainer):
@@ -51,6 +49,8 @@ def __init__(
       simulated_batch_size=16,
       n_simulated_epochs=1000,
       trajectory_dump_dir=None,
+      initial_trajectory_dir=None,
+      initial_trajectory_mix_prob=0.5,
       **kwargs
   ):
     super(SimPLe, self).__init__(
@@ -84,6 +84,9 @@ def __init__(
       trajectory_dump_dir = os.path.join(output_dir, "trajectories")
     self._trajectory_dump_root_dir = trajectory_dump_dir
 
+    self._initial_trajectory_dir = initial_trajectory_dir
+    self._initial_trajectory_mix_prob = initial_trajectory_mix_prob
+
     self._simple_epoch = 0
     self._policy_epoch = 0
     self._model_train_step = 0
@@ -122,71 +125,15 @@ def collect_trajectories(self):
     self._policy_epoch += self._n_real_epochs
     self._policy_trainer.training_loop(self._policy_epoch)
 
-  def _load_trajectories(self, trajectory_dir):
-    train_trajectories = []
-    eval_trajectories = []
-    # Search the entire directory subtree for trajectories.
-    for (subdir, _, filenames) in gfile.walk(trajectory_dir):
-      for filename in filenames:
-        shard_path = os.path.join(subdir, filename)
-        with gfile.GFile(shard_path, "rb") as f:
-          trajectories = pickle.load(f)
-          pivot = int(len(trajectories) * (1 - self._data_eval_frac))
-          train_trajectories.extend(trajectories[:pivot])
-          eval_trajectories.extend(trajectories[pivot:])
-    assert train_trajectories, "Haven't found any training data."
-    assert eval_trajectories, "Haven't found any evaluation data."
-    return (train_trajectories, eval_trajectories)
-
-  def _data_stream(self, trajectories, batch_size):
-    def generate_examples():
-      """Creates an infinite stream of shuffled examples."""
-      examples = [
-          example  # pylint: disable=g-complex-comprehension
-          for trajectory_examples in map(
-              self._sim_env.trajectory_to_training_examples, trajectories)
-          for example in trajectory_examples
-      ]
-      assert examples
-      while True:
-        random.shuffle(examples)
-        for example in examples:
-          yield example
-
-    def make_batch(examples):
-      """Stack a structure of np arrays nested in lists/tuples."""
-      assert examples
-      if isinstance(examples[0], (list, tuple)):
-        return type(examples[0])(
-            make_batch([example[i] for example in examples])
-            for i in range(len(examples[0]))
-        )
-      else:
-        return np.stack(examples, axis=0)
-
-    # Take consecutive batches from an infinite stream. This way there are no
-    # incomplete batches. We might get duplicate examples in the same batch, but
-    # that should be very rare.
-    example_stream = generate_examples()
-    while True:
-      yield make_batch(list(itertools.islice(example_stream, batch_size)))
-
   def train_model(self):
     logging.info("SimPLe epoch [% 6d]: training model.", self._simple_epoch)
 
-    # Load data from all epochs.
-    # TODO(pkozakowski): Handle the case when the data won't fit in the memory.
-    (train_trajectories, eval_trajectories) = self._load_trajectories(
-        self._trajectory_dump_root_dir)
-    train_stream = lambda: self._data_stream(  # pylint: disable=g-long-lambda
-        train_trajectories, self._model_train_batch_size)
-    eval_stream = lambda: self._data_stream(  # pylint: disable=g-long-lambda
-        eval_trajectories, self._model_train_batch_size)
+    (train_stream, eval_stream) = self._make_input_streams()
     # Ignore n_devices for now.
     inputs = lambda _: trax_inputs.Inputs(  # pylint: disable=g-long-lambda
-        train_stream=train_stream,
-        train_eval_stream=train_stream,
-        eval_stream=eval_stream,
+        train_stream=(lambda: train_stream),
+        train_eval_stream=(lambda: train_stream),
+        eval_stream=(lambda: eval_stream),
         input_shape=self._sim_env.model_input_shape,
         input_dtype=self._sim_env.model_input_dtype,
     )
@@ -208,5 +155,51 @@ def train_policy(self):
         history_stream=itertools.repeat(None),
     )
     self._policy_trainer.train_env = self._sim_env
+    # Don't dump trajectories from the simulated environment.
+    self._policy_trainer.trajectory_dump_dir = None
     self._policy_epoch += self._n_simulated_epochs
     self._policy_trainer.training_loop(self._policy_epoch)
+
+  def _make_input_streams(self):
+    def make_example_streams(trajectory_dir):
+      (train_trajs, eval_trajs) = simple.load_trajectories(
+          trajectory_dir, eval_frac=self._data_eval_frac)
+      generate_examples = functools.partial(
+          simple.generate_examples,
+          trajectory_to_training_examples_fn=(
+              self._sim_env.trajectory_to_training_examples),
+      )
+      return tuple(map(generate_examples, (train_trajs, eval_trajs)))
+
+    # We mix two data sources: trajectories collected in this SimPLe training
+    # loop ("own" data) and trajectories collected before, outside of this
+    # training loop ("initial" data).
+    mix_prob = self._initial_trajectory_mix_prob
+
+    if self._initial_trajectory_dir is None:
+      (init_train_stream, init_eval_stream) = (None, None)
+      mix_prob = 0.0  # Take just our own collected data.
+    else:
+      # Load the initial, precollected data.
+      (init_train_stream, init_eval_stream) = make_example_streams(
+          self._initial_trajectory_dir)
+
+    if self._simple_epoch == 0 and self._initial_trajectory_dir is not None:
+      # We start the loop with training the model, so we don't have our own
+      # collected data yet.
+      (own_train_stream, own_eval_stream) = (None, None)
+      mix_prob = 1.0  # Take just the initial data.
+    else:
+      # Load trajectories collected in all epochs so far.
+      (own_train_stream, own_eval_stream) = make_example_streams(
+          self._trajectory_dump_root_dir)
+
+    def mix_and_batch(streams):
+      (init_stream, own_stream) = streams
+      mixed_stream = simple.mix_streams(init_stream, own_stream, mix_prob)
+      return simple.batch_stream(mixed_stream, self._model_train_batch_size)
+
+    return tuple(map(mix_and_batch, (
+        (init_train_stream, own_train_stream),
+        (init_eval_stream, own_eval_stream),
+    )))
diff --git a/tensor2tensor/trax/rl/simple_trainer_test.py b/tensor2tensor/trax/rl/simple_trainer_test.py
index 927418855..6a9878e90 100644
--- a/tensor2tensor/trax/rl/simple_trainer_test.py
+++ b/tensor2tensor/trax/rl/simple_trainer_test.py
@@ -53,7 +53,6 @@ def _make_wrapped_env(self, name, max_episode_steps=2):
 
   def test_training_loop_acrobot(self):
     gin.bind_parameter("BoxSpaceSerializer.precision", 2)
-    gin.bind_parameter("trax.train.train_steps", 1)
     gin.bind_parameter("trax.train.eval_steps", 1)
     trainer = trainers.SimPLe(
         train_env=self._make_wrapped_env("Acrobot-v1"),
@@ -72,6 +71,7 @@ def test_training_loop_acrobot(self):
         n_real_epochs=1,
         data_eval_frac=0.5,
         model_train_batch_size=2,
+        n_model_train_steps=1,
         simulated_env_problem_class=functools.partial(
             simulated_env_problem.SerializedSequenceSimulatedEnvProblem,
             model=functools.partial(

From 2036ffe309b86bda367b1e687fafb114534500f9 Mon Sep 17 00:00:00 2001
From: Zi Yang <ziy@google.com>
Date: Sun, 1 Sep 2019 04:11:36 -0700
Subject: [PATCH 2366/2720] Added batch_config (with
 batch_op_allowed_batch_sizes) in model export V2.

PiperOrigin-RevId: 266637506
---
 tensor2tensor/utils/trainer_lib.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 42c8daff8..5a108f53c 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -349,13 +349,18 @@ def guarantee_const_scope():
       def tpu_model_fn(features, labels, mode, params):
         """Wrapper model_fn with tpu.rewrite / TPUPartitionedCall."""
         if mode == tf.estimator.ModeKeys.PREDICT and params["use_tpu"]:
+          batch_config = tpu_estimator.BatchConfig(
+              num_batch_threads=2,
+              max_batch_size=predict_batch_size,
+              batch_timeout_micros=60 * 1000,
+              allowed_batch_sizes=[predict_batch_size])
           return tpu_estimator.model_fn_inference_on_tpu(
               maybe_use_guarantee_const_getter_model_fn,
               features=features,
               labels=labels,
               config=None,
               params=params,
-              batch_config=None)
+              batch_config=batch_config)
         else:
           return model_fn(features, labels, mode, params)
 

From 8a95e96f31d1beccb8efbb2290f6a271600eb3f3 Mon Sep 17 00:00:00 2001
From: Utku Evci <evcu@google.com>
Date: Mon, 2 Sep 2019 13:35:05 -0700
Subject: [PATCH 2367/2720] switching network backbone to tf.keras.

PiperOrigin-RevId: 266812514
---
 tensor2tensor/rl/batch_dqn_agent_test.py | 38 ++++++++++++------------
 1 file changed, 19 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/rl/batch_dqn_agent_test.py b/tensor2tensor/rl/batch_dqn_agent_test.py
index 365d7b054..87152f49d 100644
--- a/tensor2tensor/rl/batch_dqn_agent_test.py
+++ b/tensor2tensor/rl/batch_dqn_agent_test.py
@@ -23,15 +23,13 @@
 
 from absl import flags
 from dopamine.agents.dqn import dqn_agent
+from dopamine.discrete_domains import atari_lib
 import numpy as np
 
 from tensor2tensor.rl import dopamine_connector
 
 import tensorflow as tf
 
-
-slim = tf.contrib.slim
-
 FLAGS = flags.FLAGS
 
 
@@ -40,6 +38,7 @@ class BatchDQNAgentTest(tf.test.TestCase):
   # dqn_agent_test.py
 
   def setUp(self):
+    super(BatchDQNAgentTest, self).setUp()
     self._test_subdir = os.path.join('/tmp/dopamine_tests', 'ckpts')
     shutil.rmtree(self._test_subdir, ignore_errors=True)
     os.makedirs(self._test_subdir)
@@ -60,26 +59,27 @@ def setUp(self):
   def _create_test_agent(self, sess):
     stack_size = self.stack_size
 
-    class MockDQNAgent(dopamine_connector.BatchDQNAgent):
+    class MockDQNNetwork(tf.keras.Model):
+      """The Keras network used in tests."""
 
-      def _network_template(self, state):
-        # This dummy network allows us to deterministically anticipate that
-        # action 0 will be selected by an argmax.
-        inputs = tf.constant(
-            np.zeros((state.shape[0], stack_size)), dtype=tf.float32)
+      def __init__(self, num_actions, **kwargs):
         # This weights_initializer gives action 0 a higher weight, ensuring
         # that it gets picked by the argmax.
+        super(MockDQNNetwork, self).__init__(**kwargs)
         weights_initializer = np.tile(
-            np.arange(self.num_actions, 0, -1), (stack_size, 1))
-        q = slim.fully_connected(
-            inputs,
-            self.num_actions,
-            weights_initializer=tf.constant_initializer(weights_initializer),
-            biases_initializer=tf.ones_initializer(),
-            activation_fn=None)
-        return self._get_network_type()(q)
-
-    agent = MockDQNAgent(
+            np.arange(num_actions, 0, -1), (stack_size, 1))
+        self.layer = tf.keras.layers.Dense(
+            num_actions,
+            kernel_initializer=tf.constant_initializer(weights_initializer),
+            bias_initializer=tf.ones_initializer())
+
+      def call(self, state):
+        inputs = tf.constant(
+            np.zeros((state.shape[0], stack_size)), dtype=tf.float32)
+        return atari_lib.DQNNetworkType(self.layer((inputs)))
+
+    agent = dopamine_connector.BatchDQNAgent(
+        network=MockDQNNetwork,
         replay_capacity=100,
         buffer_batch_size=8,
         generates_trainable_dones=True,

From a41ffb598eddf65ec31b41aa8afa56bc59694173 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 3 Sep 2019 11:21:01 -0700
Subject: [PATCH 2368/2720] Introduce hashed attention

PiperOrigin-RevId: 266974122
---
 .../trax/configs/transformer_imagenet64.gin   |   7 +-
 tensor2tensor/trax/layers/__init__.py         |   2 +
 tensor2tensor/trax/layers/attention.py        | 175 ++++++++++++++++++
 tensor2tensor/trax/layers/attention_test.py   |   9 +
 4 files changed, 192 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/configs/transformer_imagenet64.gin b/tensor2tensor/trax/configs/transformer_imagenet64.gin
index 0f56cf59a..3f299776f 100644
--- a/tensor2tensor/trax/configs/transformer_imagenet64.gin
+++ b/tensor2tensor/trax/configs/transformer_imagenet64.gin
@@ -36,10 +36,15 @@ train.trainer_class = @MemoryEfficientTrainer
 MemoryEfficientCausalAttention.dropout = 0.0
 MemoryEfficientCausalAttention.loop_stride = 512
 
+# Parameters for MergedHashedCausalAttention:
+# ==============================================================================
+MergedHashedCausalAttention.dropout = 0.0
+MergedHashedCausalAttention.n_bins = 16
+MergedHashedCausalAttention.bin_by_time = True
 
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.attention_type = @trax.layers.MemoryEfficientCausalAttention
+TransformerLM.attention_type = @trax.layers.MergedHashedCausalAttention
 TransformerLM.d_model = 1024
 TransformerLM.d_ff = 2048
 TransformerLM.dropout = 0.0
diff --git a/tensor2tensor/trax/layers/__init__.py b/tensor2tensor/trax/layers/__init__.py
index 03427d6bb..56258fede 100644
--- a/tensor2tensor/trax/layers/__init__.py
+++ b/tensor2tensor/trax/layers/__init__.py
@@ -55,3 +55,5 @@ def layer_configure(*args, **kwargs):
     DotProductCausalAttention, blacklist=["mode"])
 MemoryEfficientCausalAttention = layer_configure(
     MemoryEfficientCausalAttention, blacklist=["mode"])
+MergedHashedCausalAttention = layer_configure(
+    MergedHashedCausalAttention, blacklist=["mode"])
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 8f5b4b559..d723b6e93 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -533,6 +533,181 @@ def body_fun(vals):  # pylint: disable=invalid-name
       return final_vals[1], final_vals[2:]
 
 
+class MergedHashedCausalAttention(BaseCausalAttention):
+  """Hash-based causal attention."""
+
+  def __init__(self, dropout, mode, n_bins=64, bin_by_time=False):
+    del dropout, mode
+    super(MergedHashedCausalAttention, self).__init__()
+    self.n_bins = n_bins
+    self.bin_by_time = bin_by_time
+
+  def call(self, inputs, params=(), state=(), **kwargs):
+    del params
+    output, _ = self.call_and_grad(inputs, None, **kwargs)
+    return output, state
+
+  def has_custom_grad(self):
+    return True
+
+  def custom_grad(self, inputs, output, ct, params=(), state=(), **kwargs):
+    del output, params, state
+    _, inputs_ct = self.call_and_grad(inputs, ct, **kwargs)
+    return inputs_ct, ()
+
+  def bin_vectors_by_time(self, vecs):
+    seqlen = vecs.shape[-2]
+    assert seqlen % self.n_bins == 0
+    bin_size = int(seqlen // self.n_bins)
+
+    bins = np.arange(seqlen, dtype=np.int32) // bin_size
+    bins = jax.lax.tie_in(vecs, bins)
+    bins = bins[None, :]
+    bins = np.broadcast_to(bins, vecs.shape[:-1])
+    return bins
+
+  def make_unit_length(self, x, epsilon=1e-6):
+    variance = np.mean(x**2, axis=-1, keepdims=True)
+    norm_inputs = x / np.sqrt(variance + epsilon)
+    return norm_inputs
+
+  def hash_vectors(self, vecs, rng):
+    if self.bin_by_time:
+      # Instead of hashing, put chunks of consecutive items in the same bin.
+      # This exists as a sanity check for the other parts of this class.
+      return self.bin_vectors_by_time(vecs)
+
+    # See https://arxiv.org/pdf/1509.02897.pdf
+    assert self.n_bins % 2 == 0
+    random_rotation = jax.random.normal(
+        rng, (vecs.shape[-1], self.n_bins//2)).astype('float32')
+
+    # TODO(kitaev): making the vectors unit-length here is probably redundant.
+    vecs = self.make_unit_length(vecs)
+    rotated_vecs = np.matmul(vecs, random_rotation)
+    rotated_vecs = self.make_unit_length(rotated_vecs)
+    rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)
+    bins = np.argmax(rotated_vecs, axis=-1)
+    return bins
+
+  def call_and_grad(self, inputs, ct, rng=None, **kwargs):
+    del kwargs
+    # We use the same vector as both a query and a key. For now we haven't
+    # adjusted any of the surrounding code, so we still get a separate "key"
+    # input that we ignore.
+    qk, ignored_k, v = inputs
+    seqlen = qk.shape[-2]
+    # qk/v are n_batch*n_heads, seqlen, d_head
+
+    # bins are n_batch*n_heads, seqlen
+    # They specify which hash bucket the query/key/value vectors fall in.
+    bins = self.hash_vectors(qk, rng=rng)
+
+    # joint_t is n_batch*n_heads, seqlen
+    joint_t = jax.lax.tie_in(qk, np.arange(seqlen))
+    joint_t = np.reshape(joint_t, (1, seqlen))
+    joint_t = np.broadcast_to(joint_t, qk.shape[:-1])
+
+    assert int((self.n_bins + 1) * seqlen) < 2 ** 31, (
+        'Potential 32-bit integer overflow; please double-check the code.')
+    joint_bins_and_t = seqlen * bins + joint_t
+
+    def chunk_scalars(x):  # pylint: disable=invalid-name
+      return np.reshape(x, (x.shape[0], self.n_bins, -1))
+
+    def chunk_vectors(x):  # pylint: disable=invalid-name
+      return np.reshape(
+          x, (x.shape[0], self.n_bins, -1, x.shape[-1]))
+
+    def unchunk_vectors(x):  # pylint: disable=invalid-name
+      return np.reshape(x, (x.shape[0], -1, x.shape[-1]))
+
+    # Sort everything by bin number, with a secondary sort by time
+    # (variables starting with "s" are sorted)
+    _, sjoint_t = jax.lax.sort_key_val(
+        joint_bins_and_t, joint_t, dimension=-1)
+
+    sqk = np.take_along_axis(qk, sjoint_t[:, :, None], axis=-2)
+    sv = np.take_along_axis(v, sjoint_t[:, :, None], axis=-2)
+
+    if ct is not None:
+      so_ct = np.take_along_axis(ct, sjoint_t[:, :, None], axis=-2)
+
+    @jax.jit
+    def binned_attn(sqk, sv):  # pylint: disable=invalid-name
+      """Performs attention on sorted queries/keys/values."""
+      # Split off a "bin" axis so that attention only occurs whithin chunks.
+      bq_t = bkv_t = chunk_scalars(sjoint_t)
+      bqk = chunk_vectors(sqk)
+      bv = chunk_vectors(sv)
+
+      # Hashing operates on unit-length vectors. Unnormalized query vectors are
+      # fine because they effectively provide a learnable temperature for the
+      # attention softmax, but normalizing keys is needed so that similarity for
+      # the purposes of attention correctly corresponds to hash locality.
+      bq = bqk
+      bk = self.make_unit_length(bqk)
+
+      # Allow each chunk to attend within itself, and also one chunk back. Chunk
+      # boundaries might occur in the middle of a sequence of items from the
+      # same bin, so this increases the chances of attending to relevant items.
+      # TODO(kitaev): benchmark whether XLA pad operation is noticeably faster.
+      bk_extra = np.concatenate([bk[:, -1:, :, :], bk[:, :-1, :, :]], axis=1)
+      bk = np.concatenate([bk, bk_extra], axis=2)
+      bv_extra = np.concatenate([bv[:, -1:, :, :], bv[:, :-1, :, :]], axis=1)
+      bv = np.concatenate([bv, bv_extra], axis=2)
+      bkv_t_extra = np.concatenate([bkv_t[:, -1:, :], bkv_t[:, :-1, :]], axis=1)
+      bkv_t = np.concatenate([bkv_t, bkv_t_extra], axis=2)
+
+      # Dot-product attention.
+      dots = np.matmul(bq, np.swapaxes(bk, -1, -2)) / np.sqrt(bq.shape[-1])
+
+      # Causal masking
+      mask = jax.lax.convert_element_type(
+          jax.lax.lt(bq_t[:, :, :, None], bkv_t[:, :, None, :]),
+          np.float32)
+      dots = dots - 1e9 * mask
+
+      # Mask out attention to self except when no other targets are available.
+      self_mask = jax.lax.broadcasted_eye(dots.dtype, dots.shape, (2, 3))
+      self_mask = jax.lax.tie_in(dots, self_mask)
+      dots = dots - 32 * self_mask
+
+      # Softmax.
+      dots = np.exp(dots - dots.max(axis=-1, keepdims=True))
+      dots = dots / dots.sum(axis=-1, keepdims=True)
+      bo = np.matmul(dots, bv)
+
+      so = unchunk_vectors(bo)
+      return so
+
+    @jax.jit
+    def binned_attn_vjp(sqk, sv, so_ct):  # pylint: disable=invalid-name
+      so, vjpfun = jax.vjp(binned_attn, sqk, sv)
+      sqkv_ct = vjpfun(so_ct)
+      return so, sqkv_ct
+
+    if ct is None:
+      so = binned_attn(sqk, sv)
+      _, undo_sort = jax.lax.sort_key_val(sjoint_t, joint_t, dimension=-1)
+      out = np.take_along_axis(so, undo_sort[:, :, None], axis=-2)
+      return out, None
+    else:
+      # Jax can construct a backward pass automatically, but it's about 2x
+      # slower than writing our own. The main reason is that the backward pass
+      # of gather is in general a scatter operation, but we know we're dealing
+      # with permutations so we use gather for the backward pass too.
+      so, (sqk_ct, sv_ct) = binned_attn_vjp(sqk, sv, so_ct)
+
+      _, undo_sort = jax.lax.sort_key_val(sjoint_t, joint_t, dimension=-1)
+      out = np.take_along_axis(so, undo_sort[:, :, None], axis=-2)
+
+      qk_ct = np.take_along_axis(sqk_ct, undo_sort[:, :, None], axis=-2)
+      v_ct = np.take_along_axis(sv_ct, undo_sort[:, :, None], axis=-2)
+
+      return out, (qk_ct, np.zeros_like(ignored_k), v_ct)
+
+
 def CausalAttention(d_feature, n_heads=1,
                     d_attention_key=None, d_attention_value=None,
                     attention_type=DotProductCausalAttention, mode='train'):
diff --git a/tensor2tensor/trax/layers/attention_test.py b/tensor2tensor/trax/layers/attention_test.py
index cb42732ed..6d3f0576a 100644
--- a/tensor2tensor/trax/layers/attention_test.py
+++ b/tensor2tensor/trax/layers/attention_test.py
@@ -21,6 +21,7 @@
 
 import numpy as onp
 from tensor2tensor.trax.layers import attention
+from tensor2tensor.trax.layers import base
 from tensorflow import test
 
 
@@ -62,6 +63,14 @@ def test_shift_right_float(self):
                                     [6., 6.5, 7.]]]),
                         output_np)
 
+  def test_merged_hashed_causal_attention(self):
+    qkv_shape = (3, 32, 8)
+    input_shape = (qkv_shape, qkv_shape, qkv_shape)
+    layer = attention.MemoryEfficientCausalAttention(
+        loop_stride=16, dropout=0.1, mode='train')
+    final_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual((3, 32, 8), final_shape)
+
 
 if __name__ == '__main__':
   test.main()

From 948720432abfe6360db39295f4e6c4a7a95a6e04 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 3 Sep 2019 17:01:15 -0700
Subject: [PATCH 2369/2720] Disable policy evaluation when training on the
 simulated environment

PiperOrigin-RevId: 267046832
---
 tensor2tensor/trax/rl/base_trainer.py   | 5 +++--
 tensor2tensor/trax/rl/simple_trainer.py | 2 +-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/trax/rl/base_trainer.py b/tensor2tensor/trax/rl/base_trainer.py
index 3921ffd57..3b63e73a9 100644
--- a/tensor2tensor/trax/rl/base_trainer.py
+++ b/tensor2tensor/trax/rl/base_trainer.py
@@ -108,12 +108,13 @@ def has_any_action(trajectory):
         pickle.dump(self._trajectory_buffer, f)
       self._trajectory_buffer = []
 
-  def training_loop(self, n_epochs):
+  def training_loop(self, n_epochs, evaluate=True):
     logging.info("Starting the RL training loop.")
     for _ in range(self.epoch, n_epochs):
       self.train_epoch()
       self.dump_trajectories()
     self.save()
     self.dump_trajectories(force=True)
-    self.evaluate()
+    if evaluate:
+      self.evaluate()
     self.flush_summaries()
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index 450a87f1a..f47a3ce5d 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -158,7 +158,7 @@ def train_policy(self):
     # Don't dump trajectories from the simulated environment.
     self._policy_trainer.trajectory_dump_dir = None
     self._policy_epoch += self._n_simulated_epochs
-    self._policy_trainer.training_loop(self._policy_epoch)
+    self._policy_trainer.training_loop(self._policy_epoch, evaluate=False)
 
   def _make_input_streams(self):
     def make_example_streams(trajectory_dir):

From 5dca26abe9e9e827b48bc1cdfb36418f2b86be69 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 3 Sep 2019 17:01:27 -0700
Subject: [PATCH 2370/2720] Copy last observations when stepping the
 environment

PiperOrigin-RevId: 267046874
---
 tensor2tensor/trax/rl/simulated_env_problem.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index a31b8df76..734bf7b12 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -415,7 +415,9 @@ def _step_model(self, predict_fn, actions, rng):
     observation = self._predict_obs(predict_fn, rng)
     reward = self._reward_fn(self._last_observations, observation)
     done = self._done_fn(self._last_observations, observation)
-    self._last_observations = observation
+    # Copy the last observations, so that we don't overwrite data stored in a
+    # trajectory when resetting the environment (see _reset_model).
+    self._last_observations = np.copy(observation)
     done = np.logical_or(done, self._steps == self._max_trajectory_length - 1)
     return (observation, reward, done)
 

From f81fae74a2721229287473545093f5a9adca5fbb Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 3 Sep 2019 17:11:52 -0700
Subject: [PATCH 2371/2720] Remove redundancy in Serial constructor.

PiperOrigin-RevId: 267048999
---
 tensor2tensor/trax/layers/combinators.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index d754c016e..08832ea16 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -120,17 +120,14 @@ def __init__(self, *layers):
     self._sublayers = layers
     self._n_layers = len(layers)
 
-    if not layers:
-      self._n_inputs = 1
-      self._n_outputs = 1
-    else:
+    if layers:
       self._n_inputs, self._n_outputs = self._n_inputs_n_outputs(layers)
 
   def _ensure_flat(self, layers):
     """Ensures that layers is a single flat list of Layer instances."""
     del self
     if len(layers) == 1 and layers[0] is None:
-      layers = []
+      layers = ()
     else:
       layers = _deep_flatten(layers)
     for obj in layers:

From 7cb690820dc578a259a1476d51d03b94a965a28d Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 3 Sep 2019 17:12:02 -0700
Subject: [PATCH 2372/2720] Pass trajectory_dump_dir in rl_trainer and extend
 the trajectory shard when dumping again in the same epoch

PiperOrigin-RevId: 267049030
---
 tensor2tensor/trax/rl/base_trainer.py | 6 ++++++
 tensor2tensor/trax/rl/ppo_trainer.py  | 4 +++-
 tensor2tensor/trax/rl_trainer.py      | 8 +++++++-
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/rl/base_trainer.py b/tensor2tensor/trax/rl/base_trainer.py
index 3b63e73a9..8f671954c 100644
--- a/tensor2tensor/trax/rl/base_trainer.py
+++ b/tensor2tensor/trax/rl/base_trainer.py
@@ -104,6 +104,12 @@ def has_any_action(trajectory):
     if ready or force:
       shard_path = os.path.join(
           self.trajectory_dump_dir, "{}.pkl".format(self.epoch))
+      if gfile.exists(shard_path):
+        # Since we do an extra dump at the end of the training loop, we
+        # sometimes dump 2 times in the same epoch. When this happens, merge the
+        # two sets of trajectories.
+        with gfile.GFile(shard_path, "rb") as f:
+          self._trajectory_buffer = pickle.load(f) + self._trajectory_buffer
       with gfile.GFile(shard_path, "wb") as f:
         pickle.dump(self._trajectory_buffer, f)
       self._trajectory_buffer = []
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index d86468177..fc34c611a 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -76,6 +76,7 @@ def __init__(
       n_evals=1,
       len_history_for_policy=4,
       eval_temperatures=(1.0, 0.5),
+      **kwargs
   ):
     """Creates the PPO trainer.
 
@@ -110,12 +111,13 @@ def __init__(
       len_history_for_policy: How much of history to give to the policy.
       eval_temperatures: Sequence of temperatures to try for categorical
         sampling during evaluation.
+      **kwargs: Additional keyword arguments passed to the base class.
     """
     # Set in base class constructor.
     self._train_env = None
     self._should_reset = None
 
-    super(PPO, self).__init__(train_env, eval_env, output_dir)
+    super(PPO, self).__init__(train_env, eval_env, output_dir, **kwargs)
 
     self._n_optimizer_steps = n_optimizer_steps
     self._print_every_optimizer_steps = print_every_optimizer_steps
diff --git a/tensor2tensor/trax/rl_trainer.py b/tensor2tensor/trax/rl_trainer.py
index a7a5514a9..6cf101091 100644
--- a/tensor2tensor/trax/rl_trainer.py
+++ b/tensor2tensor/trax/rl_trainer.py
@@ -71,6 +71,8 @@
 flags.DEFINE_integer("eval_batch_size", 4, "Batch size for evaluation.")
 flags.DEFINE_boolean("parallelize_envs", False,
                      "If true, sets parallelism to number of cpu cores.")
+flags.DEFINE_string(
+    "trajectory_dump_dir", "", "Directory to dump trajectories to.")
 
 
 # TODO(afrozm): Find a better way to do these configurations.
@@ -118,7 +120,7 @@ def make_env(name, batch_size, max_timestep, clip_rewards, rendered_env,
 
 # Not just "train" to avoid a conflict with trax.train in GIN files.
 @gin.configurable(blacklist=[
-    "output_dir", "train_batch_size", "eval_batch_size"])
+    "output_dir", "train_batch_size", "eval_batch_size", "trajectory_dump_dir"])
 def train_rl(
     output_dir,
     train_batch_size,
@@ -130,6 +132,7 @@ def train_rl(
     resize_dims=(105, 80),
     trainer_class=rl_trainers.PPO,
     n_epochs=10000,
+    trajectory_dump_dir=None,
 ):
   """Train the RL agent.
 
@@ -147,6 +150,7 @@ def train_rl(
       observations to.
     trainer_class: RLTrainer class to use.
     n_epochs: Number epochs to run the training for.
+    trajectory_dump_dir: Directory to dump trajectories to.
   """
 
   if FLAGS.jax_debug_nans:
@@ -204,6 +208,7 @@ def run_training_loop():
         output_dir=output_dir,
         train_env=train_env,
         eval_env=eval_env,
+        trajectory_dump_dir=trajectory_dump_dir,
     )
     trainer.training_loop(n_epochs=n_epochs)
 
@@ -225,6 +230,7 @@ def main(argv):
       output_dir=FLAGS.output_dir,
       train_batch_size=FLAGS.train_batch_size,
       eval_batch_size=FLAGS.eval_batch_size,
+      trajectory_dump_dir=(FLAGS.trajectory_dump_dir or None),
   )
 
 

From 3573b185ba259343de44b1a8b041471c4a6511c4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 3 Sep 2019 22:29:34 -0700
Subject: [PATCH 2373/2720] Remove unused variable.

PiperOrigin-RevId: 267084230
---
 tensor2tensor/trax/layers/normalization.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/layers/normalization.py b/tensor2tensor/trax/layers/normalization.py
index b58f362f0..6ff9648ae 100644
--- a/tensor2tensor/trax/layers/normalization.py
+++ b/tensor2tensor/trax/layers/normalization.py
@@ -104,9 +104,9 @@ def average(factor, new, old):
 
 
 # Layer normalization.
-def _layer_norm_params(input_shape, input_dtype, rng, epsilon=1e-6):
+def _layer_norm_params(input_shape, input_dtype, rng):
   """Helper: create layer norm parameters."""
-  del input_dtype, rng, epsilon
+  del input_dtype, rng
   features = input_shape[-1]
   scale = np.ones(features)
   bias = np.zeros(features)

From 2e59196d8f5a69e930b9b1af6e5abaeb947ca7f2 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 4 Sep 2019 14:00:56 -0700
Subject: [PATCH 2374/2720] Implement digit weighting in SimPLe

PiperOrigin-RevId: 267226067
---
 tensor2tensor/trax/rl/simulated_env_problem.py | 14 +++++++++-----
 tensor2tensor/trax/rl/space_serializer.py      | 16 +++++++++++++++-
 tensor2tensor/trax/rl/space_serializer_test.py |  8 ++++++++
 3 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index 734bf7b12..54fae3695 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -313,7 +313,7 @@ class SerializedSequenceSimulatedEnvProblem(SimulatedEnvProblem):
 
   def __init__(self, model, reward_fn, done_fn, vocab_size,
                max_trajectory_length, observation_space, action_space,
-               *args, **kwargs):
+               significance_decay=1.0, **kwargs):
     """Initializes the env.
 
     Args:
@@ -328,13 +328,15 @@ def __init__(self, model, reward_fn, done_fn, vocab_size,
         the model.
       observation_space: (gym.Space) Observation space.
       action_space: (gym.Space) Action space.
-      *args: (tuple) Positional arguments passed to the base class.
+      significance_decay: (float) Decay for training weights of progressively
+        less significant symbols in the representation.
       **kwargs: (dict) Keyword arguments passed to the base class.
     """
     self._reward_fn = reward_fn
     self._done_fn = done_fn
     self._vocab_size = vocab_size
     self._max_trajectory_length = max_trajectory_length
+    self._significance_decay = significance_decay
     self._history = None
     self._steps = None
     self._observation_space = None
@@ -353,7 +355,6 @@ def __init__(self, model, reward_fn, done_fn, vocab_size,
     # TransformerLM).
     model = functools.partial(model, vocab_size=vocab_size)
     super(SerializedSequenceSimulatedEnvProblem, self).__init__(
-        *args,
         model=model,
         observation_space=observation_space,
         action_space=action_space,
@@ -429,8 +430,11 @@ def trajectory_to_training_examples(self, trajectory):
       obs_repr = self._obs_serializer.serialize(
           np.array([time_step.observation]))[0]
       reprs.append(obs_repr)
-      # TODO(pkozakowski): Digit weighting.
-      weights.append(np.ones_like(obs_repr))
+      # significance_map is an array of the same size as the representation,
+      # indicating the significance of each symbol. See
+      # SpaceSerializer.significance_map.
+      weights.append(
+          self._significance_decay ** self._obs_serializer.significance_map)
       if time_step.action is not None:
         action_repr = self._action_serializer.serialize(
             np.array([time_step.action]))[0]
diff --git a/tensor2tensor/trax/rl/space_serializer.py b/tensor2tensor/trax/rl/space_serializer.py
index 0f6209148..ce35d5994 100644
--- a/tensor2tensor/trax/rl/space_serializer.py
+++ b/tensor2tensor/trax/rl/space_serializer.py
@@ -35,10 +35,15 @@ class SpaceSerializer(object):
       to. Should be defined in subclasses.
     representation_length: (int) Number of symbols in the representation of
       every element of the space.
+    significance_map: (np.ndarray) Integer array of the same size as the
+      discrete representation, where elements describe the significance of
+      symbols, e.g. in fixed-precision encoding. 0 is the most significant
+      symbol, 1 the second most significant etc.
   """
 
   space_type = None
   representation_length = None
+  significance_map = None
 
   def __init__(self, space, vocab_size):
     """Creates a SpaceSerializer.
@@ -54,7 +59,7 @@ def __init__(self, space, vocab_size):
     self._vocab_size = vocab_size
 
   def serialize(self, data):
-    """Serializes a batch of space elements into a discrete sequences.
+    """Serializes a batch of space elements into discrete sequences.
 
     Should be defined in subclasses.
 
@@ -146,6 +151,11 @@ def deserialize(self, representation):
   def representation_length(self):
     return self._precision * self._space.low.size
 
+  @property
+  def significance_map(self):
+    return np.reshape(np.broadcast_to(
+        np.arange(self._precision), self._space.shape + (self._precision,)), -1)
+
 
 class DiscreteSpaceSerializer(SpaceSerializer):
   """Serializer for gym.spaces.Discrete.
@@ -166,3 +176,7 @@ def serialize(self, data):
 
   def deserialize(self, representation):
     return np.reshape(representation, -1)
+
+  @property
+  def significance_map(self):
+    return np.zeros(1, dtype=np.int32)
diff --git a/tensor2tensor/trax/rl/space_serializer_test.py b/tensor2tensor/trax/rl/space_serializer_test.py
index 5f98b31f3..148bdfce7 100644
--- a/tensor2tensor/trax/rl/space_serializer_test.py
+++ b/tensor2tensor/trax/rl/space_serializer_test.py
@@ -74,6 +74,11 @@ def test_bounds_space(self):
     output_array = serializer.deserialize(representation)
     np.testing.assert_array_almost_equal(input_array, output_array)
 
+  def test_significance_map(self):
+    (_, serializer) = self._make_space_and_serializer(shape=(2,))
+    np.testing.assert_array_equal(
+        serializer.significance_map, [0, 1, 2, 3, 0, 1, 2, 3])
+
 
 class DiscreteSpaceSerializerTest(test.TestCase):
 
@@ -103,6 +108,9 @@ def test_representation_changes(self):
     (repr1, repr2) = tuple(map(self._serializer.serialize, (array1, array2)))
     self.assertFalse(np.array_equal(repr1, repr2))
 
+  def test_significance_map(self):
+    np.testing.assert_array_equal(self._serializer.significance_map, [0])
+
 
 if __name__ == "__main__":
   test.main()

From 39232f4e4330f409f4d3abd6cf0a134db2f2a4db Mon Sep 17 00:00:00 2001
From: evalphobia <evalphobia@gmail.com>
Date: Thu, 5 Sep 2019 06:11:29 +0900
Subject: [PATCH 2375/2720] Fix serving response from Cloud ML Engine (#1688)

---
 tensor2tensor/serving/serving_utils.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index cfc2f3b5b..ba535ca88 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -23,6 +23,7 @@
 import functools
 from googleapiclient import discovery
 import grpc
+import numpy as np
 
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
 from tensor2tensor.data_generators import text_encoder
@@ -140,8 +141,12 @@ def _make_cloud_mlengine_request(examples):
             }
         } for ex in examples]
     }
-    prediction = api.projects().predict(body=input_data, name=parent).execute()
-    return prediction["predictions"]
+    response = api.projects().predict(body=input_data, name=parent).execute()
+    predictions = response["predictions"]
+    for prediction in predictions:
+      prediction["outputs"] = np.array([prediction["outputs"]])
+      prediction["scores"] = np.array(prediction["scores"])
+    return predictions
 
   return _make_cloud_mlengine_request
 

From 0daff201d0a37f34a020528dffd110131fa024ac Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 4 Sep 2019 14:41:10 -0700
Subject: [PATCH 2376/2720] Clip observations in OnlineTuneEnv

This is because the loss sometimes jumps to very high values, far outside the usual range. Clipping lets us be more robust to those jumps.

PiperOrigin-RevId: 267234718
---
 tensor2tensor/trax/rl/envs/online_tune_env.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index 759cbae6d..29e708570 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -62,6 +62,7 @@ def __init__(self,
                env_steps=100,
                start_lr=0.001,
                max_lr=10.0,
+               observation_range=(0.0, 5.0),
                # Don't save checkpoints by default, as they tend to use a lot of
                # space.
                should_save_checkpoints=False):
@@ -94,8 +95,10 @@ def __init__(self,
     # observation_metrics plus optionally the learning rate.
     observation_dim = (
         len(self._observation_metrics) + int(self._include_lr_in_observation))
+    self._observation_range = observation_range
+    (low, high) = self._observation_range
     self.observation_space = gym.spaces.Box(
-        low=float("-inf"), high=float("+inf"), shape=(observation_dim,))
+        low=low, high=high, shape=(observation_dim,))
 
   @property
   def _next_trajectory_dir(self):
@@ -127,7 +130,7 @@ def _current_metric_value(self, metric):
     metric_sequence = self._trainer.state.history.get(*metric)
     assert metric_sequence
     (_, metric_value) = metric_sequence[-1]
-    return metric_value
+    return np.clip(metric_value, *self._observation_range)
 
   @property
   def _current_observation(self):

From 5b9cc180bf60892ff1cabd29c820a9cb3ab06e19 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 4 Sep 2019 15:49:42 -0700
Subject: [PATCH 2377/2720] Internal

PiperOrigin-RevId: 267248921
---
 tensor2tensor/data_generators/text_problems.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 4c2d2cbf4..4067a65aa 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -645,6 +645,14 @@ def text2text_txt_iterator_with_label(source_txt_path, target_txt_path):
     yield {"inputs": inputs, "targets": targets, "extra_label": [extra_label]}
 
 
+def text2text_txt_iterator_with_index(source_txt_path, target_txt_path):
+  """Yield dicts for Text2TextProblem.generate_samples from lines of files."""
+  for (idx, (inputs, targets)) in enumerate(zip(
+      txt_line_iterator(source_txt_path),
+      txt_line_iterator(target_txt_path))):
+    yield {"inputs": inputs, "targets": targets, "idx": [idx]}
+
+
 def text2text_distill_iterator(source_txt_path, target_txt_path,
                                distill_txt_path):
   """Yield dicts for Text2TextProblem.generate_samples from lines of files."""

From d36a768e9e425d5272e1e7afa8a43496dd5e416c Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 4 Sep 2019 16:12:25 -0700
Subject: [PATCH 2378/2720] Enable reporting training reward in evaluation
 metrics

This is so that we don't have to run a separate evaluation in environments where episodes are very expensive (e.g. OnlineTuneEnv). Computing reward during training does the same thing, modulo truncation and temperature.

PiperOrigin-RevId: 267253713
---
 tensor2tensor/trax/rl/base_trainer.py   |  4 ++--
 tensor2tensor/trax/rl/ppo_trainer.py    | 22 +++++++++++++++++-----
 tensor2tensor/trax/rl/simple_trainer.py |  8 ++++----
 3 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/trax/rl/base_trainer.py b/tensor2tensor/trax/rl/base_trainer.py
index 8f671954c..ecd9be22d 100644
--- a/tensor2tensor/trax/rl/base_trainer.py
+++ b/tensor2tensor/trax/rl/base_trainer.py
@@ -58,7 +58,7 @@ def __init__(
   def epoch(self):
     raise NotImplementedError
 
-  def train_epoch(self):
+  def train_epoch(self, evaluate=True):
     raise NotImplementedError
 
   def evaluate(self):
@@ -117,7 +117,7 @@ def has_any_action(trajectory):
   def training_loop(self, n_epochs, evaluate=True):
     logging.info("Starting the RL training loop.")
     for _ in range(self.epoch, n_epochs):
-      self.train_epoch()
+      self.train_epoch(evaluate=evaluate)
       self.dump_trajectories()
     self.save()
     self.dump_trajectories(force=True)
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index fc34c611a..ede56f5d7 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -76,6 +76,7 @@ def __init__(
       n_evals=1,
       len_history_for_policy=4,
       eval_temperatures=(1.0, 0.5),
+      separate_eval=True,
       **kwargs
   ):
     """Creates the PPO trainer.
@@ -111,6 +112,9 @@ def __init__(
       len_history_for_policy: How much of history to give to the policy.
       eval_temperatures: Sequence of temperatures to try for categorical
         sampling during evaluation.
+      separate_eval: Whether to run separate evaluation using a set of
+        temperatures. If False, the training reward is reported as evaluation
+        reward with temperature 1.0.
       **kwargs: Additional keyword arguments passed to the base class.
     """
     # Set in base class constructor.
@@ -134,6 +138,7 @@ def __init__(
     self._n_evals = n_evals
     self._len_history_for_policy = len_history_for_policy
     self._eval_temperatures = eval_temperatures
+    self._separate_eval = separate_eval
 
     assert isinstance(self.train_env.action_space, gym.spaces.Discrete)
     n_actions = self.train_env.action_space.n
@@ -210,13 +215,13 @@ def assert_same_space(space1, space2):
   def epoch(self):
     return self._epoch
 
-  def train_epoch(self):
+  def train_epoch(self, evaluate=True):
     """Train one PPO epoch."""
     epoch_start_time = time.time()
 
     # Evaluate the policy.
     policy_eval_start_time = time.time()
-    if (self._epoch + 1) % self._eval_every_n == 0:
+    if evaluate and (self._epoch + 1) % self._eval_every_n == 0:
       self._rng, key = jax_random.split(self._rng, num=2)
       self.evaluate()
 
@@ -242,12 +247,17 @@ def train_epoch(self):
     logging.vlog(1, "Collecting trajectories took %0.2f msec.",
                  trajectory_collection_time)
 
-    avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
-    max_reward = max(np.sum(traj[2]) for traj in trajs)
-    min_reward = min(np.sum(traj[2]) for traj in trajs)
+    rewards = np.array([np.sum(traj[2]) for traj in trajs])
+    avg_reward = np.mean(rewards)
+    std_reward = np.std(rewards)
+    max_reward = np.max(rewards)
+    min_reward = np.min(rewards)
 
     self._train_sw.scalar(
         "train/reward_mean_truncated", avg_reward, step=self._epoch)
+    if evaluate and not self._separate_eval:
+      metrics = {"raw": {1.0: {"mean": avg_reward, "std": std_reward}}}
+      ppo.write_eval_reward_summaries(metrics, self._eval_sw, self._epoch)
 
     logging.vlog(1, "Rewards avg=[%0.2f], max=[%0.2f], min=[%0.2f], all=%s",
                  avg_reward, max_reward, min_reward,
@@ -502,6 +512,8 @@ def train_epoch(self):
 
   def evaluate(self):
     """Evaluate the agent."""
+    if not self._separate_eval:
+      return
     logging.vlog(1, "PPO epoch [% 6d]: evaluating policy.", self._epoch)
     self._rng, key = jax_random.split(self._rng, num=2)
     reward_stats, self._model_state = ppo.evaluate_policy(
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index f47a3ce5d..c29bed624 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -95,9 +95,9 @@ def __init__(
   def epoch(self):
     return self._simple_epoch
 
-  def train_epoch(self):
+  def train_epoch(self, evaluate=True):
     # Collect trajectories by running the policy in the real environment.
-    self.collect_trajectories()
+    self.collect_trajectories(evaluate=evaluate)
     # Train the model of the environment on the collected trajectories.
     self.train_model()
     # Train the policy inside the simulated environment generated by the model.
@@ -116,14 +116,14 @@ def flush_summaries(self):
     # TODO(pkozakowski): Report some metrics, timing?
     pass
 
-  def collect_trajectories(self):
+  def collect_trajectories(self, evaluate):
     logging.info("SimPLe epoch [% 6d]: collecting data.", self._simple_epoch)
 
     self._policy_trainer.train_env = self.train_env
     self._policy_trainer.trajectory_dump_dir = os.path.join(
         self._trajectory_dump_root_dir, str(self.epoch))
     self._policy_epoch += self._n_real_epochs
-    self._policy_trainer.training_loop(self._policy_epoch)
+    self._policy_trainer.training_loop(self._policy_epoch, evaluate=evaluate)
 
   def train_model(self):
     logging.info("SimPLe epoch [% 6d]: training model.", self._simple_epoch)

From 0ee133bf3aa6dc7b02cf7fdc8c8a32f1e3fa07b0 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 4 Sep 2019 16:59:29 -0700
Subject: [PATCH 2379/2720] Update gin configs for PPO + OnlineTune

Also changed names of configs for distributed training, so the name reflects what's inside - now we have env_online_tune_{model}_{dataset}.gin for environment configs and {algo}_online_tune.gin for RL configs. This way we can use the same RL config with multiple different env configs.

PiperOrigin-RevId: 267262029
---
 ...> env_online_tune_wide_resnet_cifar10.gin} |  0
 ...resnet_cifar10.gin => ppo_online_tune.gin} | 27 +++++++++++++------
 ...> ppo_online_tune_wide_resnet_cifar10.gin} | 26 ++++++++++++------
 tensor2tensor/trax/rl_trainer.py              |  1 +
 4 files changed, 38 insertions(+), 16 deletions(-)
 rename tensor2tensor/trax/rl/configs/{lr_learning_env_params_online_tune_wide_resnet_cifar10.gin => env_online_tune_wide_resnet_cifar10.gin} (100%)
 rename tensor2tensor/trax/rl/configs/{lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin => ppo_online_tune.gin} (50%)
 rename tensor2tensor/trax/rl/configs/{online_tune_wide_resnet_cifar10.gin => ppo_online_tune_wide_resnet_cifar10.gin} (78%)

diff --git a/tensor2tensor/trax/rl/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
similarity index 100%
rename from tensor2tensor/trax/rl/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin
rename to tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
diff --git a/tensor2tensor/trax/rl/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
similarity index 50%
rename from tensor2tensor/trax/rl/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
rename to tensor2tensor/trax/rl/configs/ppo_online_tune.gin
index 86b3692e6..665e07052 100644
--- a/tensor2tensor/trax/rl/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
@@ -1,13 +1,25 @@
 import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.rl.trainers
 
-# Parameters for FrameStackMLP:
+# Parameters for Adam:
 # ==============================================================================
-FrameStackMLP.n_frames = 4
+Adam.learning_rate = 1e-3
+Adam.b1 = 0.9
+Adam.b2 = 0.999
+Adam.weight_decay_rate = 0.0
+
+# Parameters for TransformerDecoder:
+# ==============================================================================
+TransformerDecoder.d_model = 64
+TransformerDecoder.d_ff = 128
+TransformerDecoder.dropout = 0.0
+TransformerDecoder.n_heads = 2
+TransformerDecoder.n_layers = 1
 
 # Parameters for PPO:
 # ==============================================================================
-PPO.n_optimizer_steps = 30
+PPO.n_optimizer_steps = 10
 PPO.target_kl = 0.1
 PPO.boundary = 128
 PPO.max_timestep = 128
@@ -17,12 +29,11 @@ PPO.gamma = 0.99
 PPO.lambda_ = 0.95
 PPO.c1 = 1.0
 PPO.c2 = 0.01
-PPO.eval_every_n = 10
 PPO.done_frac_for_policy_save = 0
-PPO.n_evals = 1
-PPO.len_history_for_policy = 4
-PPO.eval_temperatures = (1.0, 0.5)
-PPO.policy_and_value_model = @trax.models.FrameStackMLP
+PPO.len_history_for_policy = None
+PPO.separate_eval = False
+PPO.policy_and_value_model = @trax.models.TransformerDecoder
+PPO.policy_and_value_optimizer = @trax.optimizers.Adam
 
 # Parameters for train_rl:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rl/configs/online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin
similarity index 78%
rename from tensor2tensor/trax/rl/configs/online_tune_wide_resnet_cifar10.gin
rename to tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin
index f89769344..fdb27d4b9 100644
--- a/tensor2tensor/trax/rl/configs/online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin
@@ -24,9 +24,20 @@ Momentum.mass = 0.9
 # ==============================================================================
 shuffle_and_batch_data.preprocess_fun = @trax.inputs.cifar10_no_augmentation_preprocess
 
-# Parameters for FrameStackMLP:
+# Parameters for Adam:
 # ==============================================================================
-FrameStackMLP.n_frames = 4
+Adam.learning_rate = 1e-3
+Adam.b1 = 0.9
+Adam.b2 = 0.999
+Adam.weight_decay_rate = 0.0
+
+# Parameters for TransformerDecoder:
+# ==============================================================================
+TransformerDecoder.d_model = 64
+TransformerDecoder.d_ff = 128
+TransformerDecoder.dropout = 0.0
+TransformerDecoder.n_heads = 2
+TransformerDecoder.n_layers = 1
 
 # Parameters for WideResnet:
 # ==============================================================================
@@ -52,7 +63,7 @@ OnlineTuneEnv.observation_metrics = [
 
 # Parameters for PPO:
 # ==============================================================================
-PPO.n_optimizer_steps = 30
+PPO.n_optimizer_steps = 10
 PPO.target_kl = 0.1
 PPO.boundary = 128
 PPO.max_timestep = 128
@@ -62,12 +73,11 @@ PPO.gamma = 0.99
 PPO.lambda_ = 0.95
 PPO.c1 = 1.0
 PPO.c2 = 0.01
-PPO.eval_every_n = 10
 PPO.done_frac_for_policy_save = 0
-PPO.n_evals = 1
-PPO.len_history_for_policy = 4
-PPO.eval_temperatures = (1.0, 0.5)
-PPO.policy_and_value_model = @trax.models.FrameStackMLP
+PPO.len_history_for_policy = None
+PPO.separate_eval = False
+PPO.policy_and_value_model = @trax.models.TransformerDecoder
+PPO.policy_and_value_optimizer = @trax.optimizers.Adam
 
 # Parameters for train_rl:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rl_trainer.py b/tensor2tensor/trax/rl_trainer.py
index 6cf101091..21651607c 100644
--- a/tensor2tensor/trax/rl_trainer.py
+++ b/tensor2tensor/trax/rl_trainer.py
@@ -48,6 +48,7 @@
 from tensor2tensor.envs import rendered_env_problem
 from tensor2tensor.rl import gym_utils
 from tensor2tensor.rl.google import atari_utils  # GOOGLE-INTERNAL:
+from tensor2tensor.trax import rl  # pylint: disable=unused-import
 from tensor2tensor.trax.rl import envs as rl_envs  # pylint: disable=unused-import
 from tensor2tensor.trax.rl import trainers as rl_trainers
 

From a6a9990bff89a06f92f27631308687cc887edb88 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 4 Sep 2019 18:11:41 -0700
Subject: [PATCH 2380/2720] Adjust batch sizes of wide resnet to use just under
 8GB of memory.

PiperOrigin-RevId: 267273510
---
 tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin    | 6 +++---
 .../rl/configs/env_online_tune_wide_resnet_cifar10.gin    | 8 ++++----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
index 74eabe53f..18e65f4a7 100644
--- a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
+++ b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
@@ -6,10 +6,10 @@ import tensor2tensor.trax.trax
 
 # Parameters for batch_fun:
 # ==============================================================================
-batch_fun.batch_size = 32
+batch_fun.batch_size_per_device = 256
 batch_fun.bucket_length = 32
 batch_fun.buckets = None
-batch_fun.eval_batch_size = 32
+batch_fun.eval_batch_size = 512
 
 # Parameters for inputs:
 # ==============================================================================
@@ -39,7 +39,7 @@ WideResnet.n_output_classes = 10
 # Parameters for train:
 # ==============================================================================
 train.eval_frequency = 2000
-train.eval_steps = 20
+train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.WideResnet
 train.optimizer = @trax.optimizers.Momentum
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
index 42260dd64..c460570ed 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
@@ -6,10 +6,10 @@ import tensor2tensor.trax.rl.envs
 
 # Parameters for batch_fun:
 # ==============================================================================
-batch_fun.batch_size = 32
+batch_fun.batch_size_per_device = 256
 batch_fun.bucket_length = 32
 batch_fun.buckets = None
-batch_fun.eval_batch_size = 32
+batch_fun.eval_batch_size = 512
 
 # Parameters for inputs:
 # ==============================================================================
@@ -36,8 +36,8 @@ OnlineTuneEnv.inputs = @trax.inputs.inputs
 OnlineTuneEnv.model = @trax.models.WideResnet
 OnlineTuneEnv.optimizer = @trax.optimizers.Momentum
 OnlineTuneEnv.start_lr = 0.01
-OnlineTuneEnv.train_steps = 500
-OnlineTuneEnv.eval_steps = 50
+OnlineTuneEnv.train_steps = 100
+OnlineTuneEnv.eval_steps = 10
 OnlineTuneEnv.env_steps = 100
 OnlineTuneEnv.observation_metrics = [
     ("train", "metrics/accuracy"),

From 8e6928b97bc97c27cee5e1bf395eec6fda6d0954 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 4 Sep 2019 18:13:53 -0700
Subject: [PATCH 2381/2720] Temporarily turn off an assertion in
 complete_all_trajectories()

PiperOrigin-RevId: 267273760
---
 tensor2tensor/envs/trajectory.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index 1e6787588..73169e84b 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -280,8 +280,12 @@ def complete_all_trajectories(self):
     """Essentially same as reset, but we don't have observations."""
     for index in range(self.batch_size):
       trajectory = self._trajectories[index]
-      assert trajectory.is_active
-      self._complete_trajectory(trajectory, index)
+      # TODO(pkozakowski): This assertion breaks something in SimPLe trajectory
+      # collection code - we're probably doing something wrong there. Commenting
+      # out the assertion as a temporary measure.
+      # assert trajectory.is_active
+      if trajectory.is_active:
+        self._complete_trajectory(trajectory, index)
 
   def step(self, observations, raw_rewards, processed_rewards, dones, actions,
            infos=None):

From 6277da74f1f0fe5ed79451a79a66623374c92352 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 4 Sep 2019 18:22:10 -0700
Subject: [PATCH 2382/2720] Implement model evaluation in SimPLe

PiperOrigin-RevId: 267274723
---
 tensor2tensor/trax/rl/simple.py         | 107 ++++++++++++++++++++++++
 tensor2tensor/trax/rl/simple_trainer.py |  67 ++++++++++++---
 2 files changed, 162 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/trax/rl/simple.py b/tensor2tensor/trax/rl/simple.py
index 152d346cd..711131213 100644
--- a/tensor2tensor/trax/rl/simple.py
+++ b/tensor2tensor/trax/rl/simple.py
@@ -23,8 +23,11 @@
 import os
 import random
 
+from absl import logging
 import cloudpickle as pickle
 import numpy as np
+from tensor2tensor.envs import env_problem_utils
+from tensor2tensor.envs import trajectory
 from tensorflow.io import gfile
 
 
@@ -90,3 +93,107 @@ def make_batch(examples):
   # that should be very rare.
   while True:
     yield make_batch(list(itertools.islice(stream, batch_size)))
+
+
+# TODO(pkozakowski): This is mostly a simplified version of
+# env_problem_utils.play_env_problem_with_policy, generalized to work with
+# policies not being neural networks. Unify if possible.
+def play_env_problem(env, policy):
+  """Plays an EnvProblem using a given policy function."""
+  trajectories = [trajectory.Trajectory() for _ in range(env.batch_size)]
+  observations = env.reset()
+  for (traj, observation) in zip(trajectories, observations):
+    traj.add_time_step(observation=observation)
+
+  done_so_far = np.array([False] * env.batch_size)
+  while not np.all(done_so_far):
+    padded_observations, _ = env.trajectories.observations_np(
+        len_history_for_policy=None)
+    actions = policy(padded_observations)
+    (observations, rewards, dones, _) = env.step(actions)
+    for (traj, observation, action, reward, done) in zip(
+        trajectories, observations, actions, rewards, dones
+    ):
+      if not traj.done:
+        traj.change_last_time_step(action=action)
+        traj.add_time_step(
+            observation=observation, raw_reward=reward, done=done)
+      env.reset(indices=env_problem_utils.done_indices(dones))
+    done_so_far = np.logical_or(done_so_far, dones)
+  return trajectories
+
+
+def calculate_observation_error(real_trajectories, sim_trajectories):
+  """Calculates MSE of observations in two trajectories."""
+  def padded_obs(traj, length_difference):
+    return np.pad(
+        traj.observations_np,
+        pad_width=((0, max(length_difference, 0)), (0, 0)),
+        mode="edge",
+    )
+
+  def calculate_for_single_pair(real_trajectory, sim_trajectory):
+    diff = sim_trajectory.num_time_steps - real_trajectory.num_time_steps
+    padded_real_obs = padded_obs(real_trajectory, diff)
+    padded_sim_obs = padded_obs(sim_trajectory, -diff)
+    x = np.sum((padded_real_obs - padded_sim_obs) ** 2, axis=0)
+    return x
+
+  return np.mean([
+      calculate_for_single_pair(real_traj, sim_traj)
+      for (real_traj, sim_traj) in zip(real_trajectories, sim_trajectories)
+  ], axis=0)
+
+
+def plot_observation_error(real_trajectories, sim_trajectories, mpl_plt):
+  """Plots observations from two trajectories on the same graph."""
+  assert len(real_trajectories) == len(sim_trajectories)
+  assert real_trajectories
+  obs_dim = real_trajectories[0].last_time_step.observation.shape[0]
+  (w, h) = mpl_plt.rcParams["figure.figsize"]
+  ncols = len(real_trajectories)
+  nrows = obs_dim
+  (_, axes) = mpl_plt.subplots(
+      nrows, ncols, figsize=(w * ncols, h * nrows))
+  for (traj_index, (real_traj, sim_traj)) in enumerate(
+      zip(real_trajectories, sim_trajectories)
+  ):
+    for dim_index in range(obs_dim):
+      for (traj, label) in ((real_traj, "real"), (sim_traj, "simulated")):
+        obs = traj.observations_np
+        ax = axes[dim_index, traj_index]
+        ax.plot(np.arange(obs.shape[0]), obs[:, dim_index], label=label)
+        ax.legend()
+
+
+def evaluate_model(sim_env, real_trajectories, mpl_plt, n_to_plot=3):
+  """Reports the observation error metric and the corresponding plot."""
+  if len(sim_env.observation_space.shape) != 1:
+    logging.warning(
+        "Could not evaluate the model - only environments with vector "
+        "observation spaces are supported."
+    )
+    return
+
+  assert len(real_trajectories) == sim_env.batch_size
+
+  step = [0]
+  def policy(observations):
+    del observations
+    def get_action(traj):
+      if step[0] < traj.num_time_steps:
+        return traj.time_steps[step[0]].action or 0
+      else:
+        return 0
+    actions = np.array([get_action(traj) for traj in real_trajectories])
+    step[0] += 1
+    return actions
+
+  sim_trajectories = play_env_problem(sim_env, policy)
+  obs_errors = calculate_observation_error(real_trajectories, sim_trajectories)
+  plot_observation_error(
+      real_trajectories[:n_to_plot], sim_trajectories[:n_to_plot], mpl_plt)
+  return {
+      "observation_error/{}".format(i): obs_error
+      for (i, obs_error) in enumerate(obs_errors)
+  }
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index c29bed624..539307c0c 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -22,9 +22,12 @@
 import functools
 import itertools
 import os
+import random
 
 from absl import logging
+from matplotlib import pyplot as plt
 from tensor2tensor.trax import inputs as trax_inputs
+from tensor2tensor.trax import jaxboard
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.rl import base_trainer
 from tensor2tensor.trax.rl import simple
@@ -87,6 +90,8 @@ def __init__(
     self._initial_trajectory_dir = initial_trajectory_dir
     self._initial_trajectory_mix_prob = initial_trajectory_mix_prob
 
+    self._summary_writer = jaxboard.SummaryWriter(self._output_dir)
+
     self._simple_epoch = 0
     self._policy_epoch = 0
     self._model_train_step = 0
@@ -97,9 +102,12 @@ def epoch(self):
 
   def train_epoch(self, evaluate=True):
     # Collect trajectories by running the policy in the real environment.
-    self.collect_trajectories(evaluate=evaluate)
+    if self._simple_epoch > 0 or not self._has_initial_data:
+      self.collect_trajectories(evaluate=evaluate)
     # Train the model of the environment on the collected trajectories.
     self.train_model()
+    if evaluate:
+      self.evaluate_model()
     # Train the policy inside the simulated environment generated by the model.
     self.train_policy()
 
@@ -113,8 +121,7 @@ def save(self):
     pass
 
   def flush_summaries(self):
-    # TODO(pkozakowski): Report some metrics, timing?
-    pass
+    self._summary_writer.flush()
 
   def collect_trajectories(self, evaluate):
     logging.info("SimPLe epoch [% 6d]: collecting data.", self._simple_epoch)
@@ -160,6 +167,14 @@ def train_policy(self):
     self._policy_epoch += self._n_simulated_epochs
     self._policy_trainer.training_loop(self._policy_epoch, evaluate=False)
 
+  @property
+  def _has_own_data(self):
+    return self._simple_epoch > 0 or self._initial_trajectory_dir is None
+
+  @property
+  def _has_initial_data(self):
+    return self._initial_trajectory_dir is not None
+
   def _make_input_streams(self):
     def make_example_streams(trajectory_dir):
       (train_trajs, eval_trajs) = simple.load_trajectories(
@@ -176,23 +191,23 @@ def make_example_streams(trajectory_dir):
     # training loop ("initial" data).
     mix_prob = self._initial_trajectory_mix_prob
 
-    if self._initial_trajectory_dir is None:
-      (init_train_stream, init_eval_stream) = (None, None)
-      mix_prob = 0.0  # Take just our own collected data.
-    else:
+    if self._has_initial_data:
       # Load the initial, precollected data.
       (init_train_stream, init_eval_stream) = make_example_streams(
           self._initial_trajectory_dir)
+    else:
+      (init_train_stream, init_eval_stream) = (None, None)
+      mix_prob = 0.0  # Take just our own collected data.
 
-    if self._simple_epoch == 0 and self._initial_trajectory_dir is not None:
+    if self._has_own_data:
+      # Load trajectories collected in all epochs so far.
+      (own_train_stream, own_eval_stream) = make_example_streams(
+          self._trajectory_dump_root_dir)
+    else:
       # We start the loop with training the model, so we don't have our own
       # collected data yet.
       (own_train_stream, own_eval_stream) = (None, None)
       mix_prob = 1.0  # Take just the initial data.
-    else:
-      # Load trajectories collected in all epochs so far.
-      (own_train_stream, own_eval_stream) = make_example_streams(
-          self._trajectory_dump_root_dir)
 
     def mix_and_batch(streams):
       (init_stream, own_stream) = streams
@@ -203,3 +218,31 @@ def mix_and_batch(streams):
         (init_train_stream, own_train_stream),
         (init_eval_stream, own_eval_stream),
     )))
+
+  def evaluate_model(self):
+    logging.info("SimPLe epoch [% 6d]: evaluating model.", self._simple_epoch)
+
+    self._sim_env.initialize(
+        batch_size=self._simulated_batch_size,
+        history_stream=itertools.repeat(None),
+    )
+
+    if self._has_own_data:
+      trajectory_dir = self._trajectory_dump_root_dir
+    else:
+      trajectory_dir = self._initial_trajectory_dir
+
+    (_, eval_trajectories) = simple.load_trajectories(
+        trajectory_dir, eval_frac=self._data_eval_frac)
+    chosen_trajectories = [
+        random.choice(eval_trajectories)
+        for _ in range(self._sim_env.batch_size)
+    ]
+    summaries = simple.evaluate_model(self._sim_env, chosen_trajectories, plt)
+    if summaries is not None:
+      for (name, value) in summaries.items():
+        self._summary_writer.scalar(
+            "simple/{}".format(name), value, step=self._simple_epoch)
+      self._summary_writer.plot(
+          "simple/model_eval_plot", plt, step=self._simple_epoch)
+      self.flush_summaries()

From 7c45b49e0df458b97e5f88b6b220943075e2dcb9 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 4 Sep 2019 19:02:02 -0700
Subject: [PATCH 2383/2720] Add a config for SimPLe + OnlineTuneEnv

PiperOrigin-RevId: 267279270
---
 tensor2tensor/trax/rl/__init__.py             |  2 +
 .../env_online_tune_wide_resnet_cifar10.gin   |  6 --
 .../ppo_online_tune_wide_resnet_cifar10.gin   |  6 --
 .../trax/rl/configs/simple_online_tune.gin    | 74 +++++++++++++++++++
 tensor2tensor/trax/rl/envs/online_tune_env.py |  7 +-
 tensor2tensor/trax/rl/simple_trainer.py       |  7 +-
 .../trax/rl/simulated_env_problem.py          | 20 +++++
 7 files changed, 108 insertions(+), 14 deletions(-)
 create mode 100644 tensor2tensor/trax/rl/configs/simple_online_tune.gin

diff --git a/tensor2tensor/trax/rl/__init__.py b/tensor2tensor/trax/rl/__init__.py
index b86f23d3b..529ab0e55 100644
--- a/tensor2tensor/trax/rl/__init__.py
+++ b/tensor2tensor/trax/rl/__init__.py
@@ -48,3 +48,5 @@ def configure_simulated_env_problem(*args, **kwargs):
 cartpole_reward_fn = configure_rl(simulated_env_problem.cartpole_reward_fn)
 acrobot_done_fn = configure_rl(simulated_env_problem.acrobot_done_fn)
 acrobot_reward_fn = configure_rl(simulated_env_problem.acrobot_reward_fn)
+onlinetune_done_fn = configure_rl(simulated_env_problem.onlinetune_done_fn)
+onlinetune_reward_fn = configure_rl(simulated_env_problem.onlinetune_reward_fn)
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
index c460570ed..330eb255d 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
@@ -39,9 +39,3 @@ OnlineTuneEnv.start_lr = 0.01
 OnlineTuneEnv.train_steps = 100
 OnlineTuneEnv.eval_steps = 10
 OnlineTuneEnv.env_steps = 100
-OnlineTuneEnv.observation_metrics = [
-    ("train", "metrics/accuracy"),
-    ("train", "metrics/loss"),
-    ("eval", "metrics/accuracy"),
-    ("eval", "metrics/loss"),
-]
diff --git a/tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin
index fdb27d4b9..b0adfdb9c 100644
--- a/tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin
@@ -54,12 +54,6 @@ OnlineTuneEnv.start_lr = 0.01
 OnlineTuneEnv.train_steps = 500
 OnlineTuneEnv.eval_steps = 50
 OnlineTuneEnv.env_steps = 100
-OnlineTuneEnv.observation_metrics = [
-    ("train", "metrics/accuracy"),
-    ("train", "metrics/loss"),
-    ("eval", "metrics/accuracy"),
-    ("eval", "metrics/loss"),
-]
 
 # Parameters for PPO:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rl/configs/simple_online_tune.gin b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
new file mode 100644
index 000000000..0fb13480a
--- /dev/null
+++ b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
@@ -0,0 +1,74 @@
+import tensor2tensor.trax.models
+import tensor2tensor.trax.rl
+import tensor2tensor.trax.rl.space_serializer
+import tensor2tensor.trax.rl.trainers
+
+# Parameters for BoxSpaceSerializer:
+# ==============================================================================
+BoxSpaceSerializer.precision = 2
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 0.1
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 1000
+
+# Parameters for TransformerDecoder:
+# ==============================================================================
+TransformerDecoder.d_model = 64
+TransformerDecoder.d_ff = 128
+TransformerDecoder.dropout = 0.0
+TransformerDecoder.n_heads = 2
+TransformerDecoder.n_layers = 1
+
+# Parameters for PPO:
+# ==============================================================================
+PPO.n_optimizer_steps = 10
+PPO.target_kl = 0.1
+PPO.boundary = 128
+PPO.max_timestep = 128
+PPO.max_timestep_eval = 128
+PPO.random_seed = None
+PPO.gamma = 0.99
+PPO.lambda_ = 0.95
+PPO.c1 = 1.0
+PPO.c2 = 0.01
+PPO.done_frac_for_policy_save = 0
+PPO.len_history_for_policy = None
+PPO.separate_eval = False
+PPO.policy_and_value_model = @trax.models.TransformerDecoder
+
+# Parameters for SerializedSequenceSimulatedEnvProblem:
+# ==============================================================================
+SerializedSequenceSimulatedEnvProblem.model = @trax.models.TransformerLM
+SerializedSequenceSimulatedEnvProblem.reward_fn = @trax.rl.onlinetune_reward_fn
+SerializedSequenceSimulatedEnvProblem.done_fn = @trax.rl.onlinetune_done_fn
+SerializedSequenceSimulatedEnvProblem.vocab_size = 128
+SerializedSequenceSimulatedEnvProblem.max_trajectory_length = 101
+SerializedSequenceSimulatedEnvProblem.significance_decay = 1.0
+
+# Parameters for SimPLe:
+# ==============================================================================
+SimPLe.policy_trainer_class = @trax.rl.trainers.PPO
+SimPLe.n_real_epochs = 2
+SimPLe.n_model_train_steps = 100000
+SimPLe.model_train_batch_size = 128
+SimPLe.simulated_env_problem_class = @trax.rl.SerializedSequenceSimulatedEnvProblem
+SimPLe.simulated_batch_size = 128
+SimPLe.n_simulated_epochs = 50
+SimPLe.initial_trajectory_mix_prob = 0.5
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.d_model = 256
+TransformerLM.d_ff = 512
+TransformerLM.n_layers = 2
+TransformerLM.n_heads = 2
+TransformerLM.dropout = 0.1
+TransformerLM.max_len = 1024
+
+# Parameters for train_rl:
+# ==============================================================================
+train_rl.env_name = "ClientEnv-v0"
+train_rl.n_epochs = 100
+train_rl.trainer_class = @trax.rl.trainers.SimPLe
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index 29e708570..70d6afbb0 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -54,7 +54,12 @@ def __init__(self,
                optimizer=trax_opt.SM3,
                inputs=trax_inputs.inputs,
                action_multipliers=None,
-               observation_metrics=(("eval", "metrics/accuracy"),),
+               observation_metrics=(
+                   ("train", "metrics/accuracy"),
+                   ("train", "metrics/loss"),
+                   ("eval", "metrics/accuracy"),
+                   ("eval", "metrics/loss"),
+               ),
                include_lr_in_observation=False,
                reward_metric=("eval", "metrics/accuracy"),
                train_steps=100,
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index 539307c0c..32eb92944 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -149,7 +149,12 @@ def train_model(self):
     trax.train(
         model=self._sim_env.model,
         inputs=inputs,
-        train_steps=self._model_train_step,
+        # TODO(pkozakowski): Currently trax.train trains the model for
+        # train_steps more steps, whereas it should train up to train_steps
+        # total steps in order for the restarts to work properly. Change the
+        # argument once this behavior is changed.
+        # train_steps=self._model_train_step,
+        train_steps=self._n_model_train_steps,
         output_dir=self._model_dir,
         has_weights=True,
     )
diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index 54fae3695..818fc42f6 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -488,3 +488,23 @@ def acrobot_done_fn(previous_observation, current_observation):
 def acrobot_reward_fn(previous_observation, current_observation):
   done = acrobot_done_fn(previous_observation, current_observation)
   return -1.0 + done  # -1 reward for every timestep until the end.
+
+
+def onlinetune_done_fn(previous_observation, current_observation):
+  del previous_observation
+  del current_observation
+  # Never return "done" from the environment, rely on max_trajectory_length
+  # instead.
+  return False
+
+
+def onlinetune_reward_fn(
+    previous_observation,
+    current_observation,
+    # 2 is the evaluation accuracy metric in the default settings of
+    # OnlineTuneEnv.
+    dim_index=2,
+):
+  prev = previous_observation[:, dim_index]
+  cur = current_observation[:, dim_index]
+  return cur - prev

From e5326c19222f30e1df5f37759911ed7f99b556f0 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 4 Sep 2019 19:11:25 -0700
Subject: [PATCH 2384/2720] Remove unused artifacts from learning_rate module
 and increase wide resnet lr after last batch size increase.

PiperOrigin-RevId: 267280644
---
 .../trax/configs/wide_resnet_cifar10_8gb.gin        |  4 ++--
 tensor2tensor/trax/learning_rate.py                 | 13 -------------
 2 files changed, 2 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
index 18e65f4a7..473b0a928 100644
--- a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
+++ b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
@@ -18,7 +18,7 @@ inputs.dataset_name = 'cifar10'
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-EvalAdjustingSchedule.constant = 1.0
+EvalAdjustingSchedule.constant = 4.0
 MultifactorSchedule.factors = 'constant * linear_warmup'
 MultifactorSchedule.warmup_steps = 400
 
@@ -38,7 +38,7 @@ WideResnet.n_output_classes = 10
 
 # Parameters for train:
 # ==============================================================================
-train.eval_frequency = 2000
+train.eval_frequency = 200
 train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.WideResnet
diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index 4363eddf8..cb5ae605f 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -30,14 +30,6 @@
 from tensor2tensor.trax.backend import numpy as np
 
 
-# A dictionary to memoize results of the MultifactorSchedule below.
-# We memoize because returning exactly the same function objects allows
-# later (in the training loop) to optimize re-compiling the function (for
-# running on an accelerator) only when it changes. Note that this does not
-# affect correctness, it is done purely for performance reasons.
-_memoized_multifactor_schedules = {}
-
-
 @gin.configurable(blacklist=["history"])
 def MultifactorSchedule(history=None,
                         factors="constant * linear_warmup",
@@ -66,10 +58,6 @@ def MultifactorSchedule(history=None,
   """
   del history
 
-  cache_args = (factors, constant, warmup_steps)
-  if cache_args in _memoized_multifactor_schedules:
-    return _memoized_multifactor_schedules[cache_args]
-
   factors = [n.strip() for n in factors.split("*")]
 
   def learning_rate(step):  # pylint: disable=invalid-name
@@ -88,7 +76,6 @@ def learning_rate(step):  # pylint: disable=invalid-name
         raise ValueError("Unknown factor %s." % name)
     return ret
 
-  _memoized_multifactor_schedules[cache_args] = learning_rate
   return learning_rate
 
 

From 44e30ad91bf864b4cbe622d8d780c678aea585db Mon Sep 17 00:00:00 2001
From: Katherine Lee <katherinelee@google.com>
Date: Thu, 5 Sep 2019 12:01:51 -0700
Subject: [PATCH 2385/2720] Update avg_checkpoint.py to use compat.v1.

PiperOrigin-RevId: 267422314
---
 tensor2tensor/utils/avg_checkpoints.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index 10e58bd31..453ee7cc0 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -114,4 +114,4 @@ def main(_):
 
 
 if __name__ == "__main__":
-  tf.app.run()
+  tf.compat.v1.app.run()

From 448ce0f56d7f67f7b610c1db72e2ea9529c44511 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 5 Sep 2019 12:11:42 -0700
Subject: [PATCH 2386/2720] [t2t] Remove the workaround for "ValueError:
 prediction values with "inputs: DatasetToSingleElement:0" must be from the
 default graph.".

If not run on TPU this workaround is unnecessary but sometimes breaks running code with an error: "Cannot use 'while/Pad' as input to 'Identity' because 'while/Pad' is in a while loop." (encountered while running e.g. universal_transformer).

PiperOrigin-RevId: 267424524
---
 tensor2tensor/utils/t2t_model.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index e05d277dd..224a1f236 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1701,12 +1701,6 @@ def estimator_spec_predict(self, features, use_tpu=False):
       outputs = infer_out
       scores = None
 
-    # Workaround for "ValueError: prediction values must be from the default
-    # graph" during TPU model exporting.
-    # TODO(b/130501786): remove tf.identity once default graph mismatch is fixed
-    for name, feature in features.items():
-      features[name] = tf.identity(feature)
-
     inputs = features.get("inputs")
     if inputs is None:
       inputs = features["targets"]

From f6d8e0a2c56c294061a55fca1748633d1ac8a9b3 Mon Sep 17 00:00:00 2001
From: joao guilherme <joaoguilhermearujo@gmail.com>
Date: Thu, 5 Sep 2019 20:47:16 -0300
Subject: [PATCH 2387/2720] Activations (#1684)

* Added new initializers

* Fixed glorot uniform

* Fixed order of init arguments

* Added uniform random and documentation

* Added requested changes to initializers.py

* Added requested changes to initializers_test.py

* Added choice of input axis and output axis

* Added choice of input axis and output axis

* fixed errors and implemented requested changes

* fixed typo

* tests passing

* fixed get fans

* Fixed numpy weird behavior

* Fixed typo

* Use scipy.special.expit for sigmoid

* Use clip instead of maximum, see colab on pull request

* Added rectifiers

* Added Gaussian rectifier

* Added Gaussian rectifier
---
 tensor2tensor/trax/backend.py                 |  2 +
 tensor2tensor/trax/layers/core.py             | 47 ++++++++++++++----
 tensor2tensor/trax/layers/initializers.py     |  1 -
 .../trax/layers/initializers_test.py          | 48 +++++++++++++++++++
 4 files changed, 87 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index 866850f9b..04d87c707 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -158,6 +158,8 @@ def jax_randint(key, shape, minval, maxval, dtype=onp.int32):
     "name": "jax",
     "np": jnp,
     "logsumexp": jax_special.logsumexp,
+    "expit": jax_special.expit,
+    "erf": jax_special.erf,
     "conv": jax_conv,
     "avg_pool": jax_avg_pool,
     "max_pool": jax_max_pool,
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 3e39a0f85..0c2827040 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -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.
-
 """Trax layers library."""
 
 from __future__ import absolute_import
@@ -30,12 +29,38 @@
 
 @base.layer()
 def Relu(x, **unused_kwargs):
-  return np.maximum(x, np.zeros_like(x))
+  return np.clip(x, a_min=0.)
+
+
+@base.layer()
+def ParametricRelu(x, a=1., **unused_kwargs):
+  return np.clip(a * x, a_min=0.)
+
+
+@base.layer()
+def LeakyRelu(x, a=0.01, **unused_kwargs):
+  return np.where(x >= 0, x, a * x)
+
+
+@base.layer()
+def Elu(x, a=1., **unused_kwargs):
+  return np.where(x > 0, x, a * np.expm1(x))
+
+
+@base.layer()
+def Selu(x,
+         alpha=1.6732632423543772848170429916717,
+         lmbda=1.0507009873554804934193349852946):
+  return lmbda * np.where(x > 0, x, alpha * np.expm1(x))
+
+@base.layer()
+def Gelu(x, **unused_kwargs):
+  return x * backend.erf(x)
 
 
 @base.layer()
 def Sigmoid(x, **unused_kwargs):
-  return 1. / (1. + np.exp(-x))
+  return backend.expit(x)
 
 
 @base.layer()
@@ -87,7 +112,8 @@ def ToFloat(x, **unused_kwargs):
 class Dense(base.Layer):
   """Layer constructor function for a dense (fully-connected) layer."""
 
-  def __init__(self, n_units,
+  def __init__(self,
+               n_units,
                kernel_initializer=init.GlorotUniformInitializer(),
                bias_initializer=init.RandomNormalInitializer(1e-6)):
     super(Dense, self).__init__()
@@ -111,7 +137,9 @@ def new_parameters(self, input_shape, input_dtype, rng):
 class Embedding(base.Layer):
   """Layer constructor function for an embedding layer."""
 
-  def __init__(self, d_feature, vocab_size,
+  def __init__(self,
+               d_feature,
+               vocab_size,
                kernel_initializer=init.GlorotUniformInitializer()):
     super(Embedding, self).__init__()
     self._d_feature = d_feature  # feature dimensionality
@@ -124,8 +152,8 @@ def call(self, x, params, state, **kwargs):
 
   def new_parameters(self, input_shape, input_dtype, rng):
     del input_dtype
-    return self._kernel_initializer(
-        (self._vocab_size, self._d_feature), rng), ()
+    return self._kernel_initializer((self._vocab_size, self._d_feature),
+                                    rng), ()
 
 
 # Flatten.
@@ -133,9 +161,8 @@ def new_parameters(self, input_shape, input_dtype, rng):
 def Flatten(x, params, n_axes_to_keep=1, **kwargs):
   del params, kwargs
   if n_axes_to_keep >= len(x.shape):
-    raise ValueError(
-        "n_axes_to_keep[%d] should be less than input's rank[%d]" %
-        (n_axes_to_keep, len(x.shape)))
+    raise ValueError("n_axes_to_keep[%d] should be less than input's rank[%d]" %
+                     (n_axes_to_keep, len(x.shape)))
   return np.reshape(x, (x.shape[:n_axes_to_keep] + (-1,)))
 
 
diff --git a/tensor2tensor/trax/layers/initializers.py b/tensor2tensor/trax/layers/initializers.py
index 97d334a56..b63a719ae 100644
--- a/tensor2tensor/trax/layers/initializers.py
+++ b/tensor2tensor/trax/layers/initializers.py
@@ -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.
-
 """Trax initializers."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/trax/layers/initializers_test.py b/tensor2tensor/trax/layers/initializers_test.py
index 3a5586e96..a4f3669a9 100644
--- a/tensor2tensor/trax/layers/initializers_test.py
+++ b/tensor2tensor/trax/layers/initializers_test.py
@@ -74,5 +74,53 @@ def test_kaiming_uniform(self):
     self.assertEqual(tuple(init_value.shape), input_shape)
 
 
+  def test_random_uniform(self):
+    initializer = initializers.RandomUniformInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+  def test_glorot_normal(self):
+    initializer = initializers.GlorotNormalInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+
+  def test_glorot_uniform(self):
+    initializer = initializers.GlorotUniformInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+
+  def test_lecun_normal(self):
+    initializer = initializers.LeCunNormalInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+
+  def test_lecun_uniform(self):
+    initializer = initializers.LeCunUniformInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+
+  def test_kaiming_normal(self):
+    initializer = initializers.KaimingNormalInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+
+  def test_kaiming_uniform(self):
+    initializer = initializers.KaimingUniformInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
+
+
 if __name__ == "__main__":
   absltest.main()

From 60d5a1df373d4eaa23419c6b0ec340111c59ef3b Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 5 Sep 2019 16:47:34 -0700
Subject: [PATCH 2388/2720] Merge of PR #1684

PiperOrigin-RevId: 267485270
---
 tensor2tensor/trax/backend.py                 |  9 ++++
 tensor2tensor/trax/layers/core.py             |  6 ++-
 tensor2tensor/trax/layers/initializers.py     |  1 +
 .../trax/layers/initializers_test.py          | 48 -------------------
 4 files changed, 14 insertions(+), 50 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index 04d87c707..226d9b553 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -184,6 +184,7 @@ def jax_randint(key, shape, minval, maxval, dtype=onp.int32):
     "jit": (lambda f: f),
     "random_get_prng": lambda seed: None,
     "random_split": lambda prng, num=2: (None,) * num,
+    "expit": (lambda x: 1. / (1. + onp.exp(-x))),
 }
 
 
@@ -195,6 +196,14 @@ def logsumexp(*args, **kwargs):
   return backend()["logsumexp"](*args, **kwargs)
 
 
+def expit(*args, **kwargs):
+  return backend()["expit"](*args, **kwargs)
+
+
+def erf(*args, **kwargs):
+  return backend()["erf"](*args, **kwargs)
+
+
 def conv(*args, **kwargs):
   return backend()["conv"](*args, **kwargs)
 
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 0c2827040..117f30388 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -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.
+
 """Trax layers library."""
 
 from __future__ import absolute_import
@@ -29,12 +30,12 @@
 
 @base.layer()
 def Relu(x, **unused_kwargs):
-  return np.clip(x, a_min=0.)
+  return np.maximum(x, np.zeros_like(x))
 
 
 @base.layer()
 def ParametricRelu(x, a=1., **unused_kwargs):
-  return np.clip(a * x, a_min=0.)
+  return np.maximum(a * x, np.zeros_like(x))
 
 
 @base.layer()
@@ -53,6 +54,7 @@ def Selu(x,
          lmbda=1.0507009873554804934193349852946):
   return lmbda * np.where(x > 0, x, alpha * np.expm1(x))
 
+
 @base.layer()
 def Gelu(x, **unused_kwargs):
   return x * backend.erf(x)
diff --git a/tensor2tensor/trax/layers/initializers.py b/tensor2tensor/trax/layers/initializers.py
index b63a719ae..97d334a56 100644
--- a/tensor2tensor/trax/layers/initializers.py
+++ b/tensor2tensor/trax/layers/initializers.py
@@ -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.
+
 """Trax initializers."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/trax/layers/initializers_test.py b/tensor2tensor/trax/layers/initializers_test.py
index a4f3669a9..6a0bef4c4 100644
--- a/tensor2tensor/trax/layers/initializers_test.py
+++ b/tensor2tensor/trax/layers/initializers_test.py
@@ -31,49 +31,12 @@ def test_random_normal(self):
     init_value = initializer(input_shape, random.get_prng(0))
     self.assertEqual(tuple(init_value.shape), input_shape)
 
-  def test_random_uniform(self):
-    initializer = initializers.RandomUniformInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-  def test_glorot_normal(self):
-    initializer = initializers.GlorotNormalInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-  def test_glorot_uniform(self):
-    initializer = initializers.GlorotUniformInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-  def test_lecun_normal(self):
-    initializer = initializers.LeCunNormalInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
   def test_lecun_uniform(self):
     initializer = initializers.LeCunUniformInitializer()
     input_shape = (29, 5, 7, 20)
     init_value = initializer(input_shape, random.get_prng(0))
     self.assertEqual(tuple(init_value.shape), input_shape)
 
-  def test_kaiming_normal(self):
-    initializer = initializers.KaimingNormalInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-  def test_kaiming_uniform(self):
-    initializer = initializers.KaimingUniformInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-
   def test_random_uniform(self):
     initializer = initializers.RandomUniformInitializer()
     input_shape = (29, 5, 7, 20)
@@ -86,35 +49,24 @@ def test_glorot_normal(self):
     init_value = initializer(input_shape, random.get_prng(0))
     self.assertEqual(tuple(init_value.shape), input_shape)
 
-
   def test_glorot_uniform(self):
     initializer = initializers.GlorotUniformInitializer()
     input_shape = (29, 5, 7, 20)
     init_value = initializer(input_shape, random.get_prng(0))
     self.assertEqual(tuple(init_value.shape), input_shape)
 
-
   def test_lecun_normal(self):
     initializer = initializers.LeCunNormalInitializer()
     input_shape = (29, 5, 7, 20)
     init_value = initializer(input_shape, random.get_prng(0))
     self.assertEqual(tuple(init_value.shape), input_shape)
 
-
-  def test_lecun_uniform(self):
-    initializer = initializers.LeCunUniformInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-
   def test_kaiming_normal(self):
     initializer = initializers.KaimingNormalInitializer()
     input_shape = (29, 5, 7, 20)
     init_value = initializer(input_shape, random.get_prng(0))
     self.assertEqual(tuple(init_value.shape), input_shape)
 
-
   def test_kaiming_uniform(self):
     initializer = initializers.KaimingUniformInitializer()
     input_shape = (29, 5, 7, 20)

From c6b175c1578a5af5e90a68649de032abbc2d3146 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 6 Sep 2019 09:59:36 -0700
Subject: [PATCH 2389/2720] Add default augmentation for CIFAR images and
 adjust configs.

PiperOrigin-RevId: 267618835
---
 .../trax/configs/wide_resnet_cifar10_8gb.gin  |  4 +--
 tensor2tensor/trax/inputs.py                  | 33 +++++++++++++++++++
 .../env_online_tune_wide_resnet_cifar10.gin   |  4 +--
 3 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
index 473b0a928..8c5aab04d 100644
--- a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
+++ b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
@@ -28,12 +28,12 @@ Momentum.mass = 0.9
 
 # Parameters for preprocess_fun:
 # ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.cifar10_no_augmentation_preprocess
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.cifar10_augmentation_preprocess
 
 # Parameters for WideResnet:
 # ==============================================================================
 WideResnet.widen_factor = 10
-WideResnet.n_blocks = 3
+WideResnet.n_blocks = 4
 WideResnet.n_output_classes = 10
 
 # Parameters for train:
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 8a60b17a5..30e76da7f 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -373,6 +373,39 @@ def cast_image(features, targets):
   return dataset
 
 
+@gin.configurable(blacklist=["dataset", "training"])
+def cifar10_augmentation_preprocess(dataset, training):
+  """Preprocessing for cifar10 with augmentation (see below)."""
+
+  def augment_image(image):
+    """Image augmentation suitable for CIFAR-10/100.
+
+    As described in https://arxiv.org/pdf/1608.06993v3.pdf (page 5).
+
+    Args:
+      image: a Tensor.
+    Returns:
+      Tensor of the same shape as image.
+    """
+    image = tf.image.resize_image_with_crop_or_pad(image, 40, 40)
+    image = tf.random_crop(image, [32, 32, 3])
+    image = tf.image.random_flip_left_right(image)
+    return image
+
+  def augment(features, targets):
+    features["image"] = augment_image(features["image"])
+    return features, targets
+
+  def cast_image(features, targets):
+    features["image"] = tf.cast(features["image"], tf.float32) / 255.0
+    return features, targets
+
+  if training:
+    dataset = dataset.map(augment)
+  dataset = dataset.map(cast_image)
+  return dataset
+
+
 def no_preprocess(dataset, training):
   del training
   return dataset
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
index 330eb255d..e67ab1b0d 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
@@ -22,12 +22,12 @@ Momentum.mass = 0.9
 
 # Parameters for shuffle_and_batch_data:
 # ==============================================================================
-shuffle_and_batch_data.preprocess_fun = @trax.inputs.cifar10_no_augmentation_preprocess
+shuffle_and_batch_data.preprocess_fun = @trax.inputs.cifar10_augmentation_preprocess
 
 # Parameters for WideResnet:
 # ==============================================================================
 WideResnet.widen_factor = 10
-WideResnet.n_blocks = 3
+WideResnet.n_blocks = 4
 WideResnet.n_output_classes = 10
 
 # Parameters for OnlineTuneEnv:

From afbb8c95d5846758a6d293e4aa880af8682e2306 Mon Sep 17 00:00:00 2001
From: Allen Lavoie <allenl@google.com>
Date: Fri, 6 Sep 2019 11:28:27 -0700
Subject: [PATCH 2390/2720] Handle sparse resource variable gradients in
 MultiStepAdamOptimizer

As for reference variables it's just making them dense. But at least it won't throw an exception.

PiperOrigin-RevId: 267638658
---
 tensor2tensor/utils/multistep_optimizer.py      | 13 +++++++++++++
 tensor2tensor/utils/multistep_optimizer_test.py | 10 ++++++++++
 2 files changed, 23 insertions(+)

diff --git a/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
index 2e997e2bc..ee80e659a 100644
--- a/tensor2tensor/utils/multistep_optimizer.py
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -100,6 +100,19 @@ def _apply_sparse(self, grad, var):
     return self._apply_cond(
         super(MultistepAdamOptimizer, self)._apply_dense, dense_grad, var)
 
+  def _resource_apply_sparse_duplicate_indices(self, grad, var, indices):
+    tf.logging.warning("MultistepAdamOptimizer does not support sparse updates")
+    # Note that conversion to a dense Tensor handles duplicate `indices`
+    # correctly (summing them). A real sparse implementation will probably want
+    # to override _resource_apply_sparse instead so it gets them de-duplicated
+    # automatically.
+    dense_grad = tf.convert_to_tensor(
+        tf.IndexedSlices(values=grad, indices=indices,
+                         dense_shape=tf.shape(var)))
+    return self._apply_cond(
+        super(MultistepAdamOptimizer, self)._resource_apply_dense,
+        dense_grad, var)
+
   def _finish(self, update_ops, name_scope):
     """Updates beta_power variables every n batches and incrs counter."""
     iter_ = self._get_iter_variable()
diff --git a/tensor2tensor/utils/multistep_optimizer_test.py b/tensor2tensor/utils/multistep_optimizer_test.py
index 05b9fb059..c11fb21b6 100644
--- a/tensor2tensor/utils/multistep_optimizer_test.py
+++ b/tensor2tensor/utils/multistep_optimizer_test.py
@@ -93,6 +93,16 @@ def testMultistep(self):
                 self.evaluate(singlestep_var1),
                 self.evaluate(multistep_var1))
 
+  def testResourceVariables(self):
+    v1 = tf.Variable([1., 2.], use_resource=True)
+    v2 = tf.Variable([3., 4.], use_resource=True)
+    with tf.GradientTape() as tape:
+      tape.watch([v1, v2])
+      loss = tf.reduce_sum(tf.gather(params=v1, indices=[0]) + v2)
+    v1_grad, v2_grad = tape.gradient(loss, [v1, v2])
+    multistep_opt = multistep_optimizer.MultistepAdamOptimizer(0.1)
+    multistep_opt.apply_gradients(((v1_grad, v1), (v2_grad, v2)))
+
 
 if __name__ == '__main__':
   tf.test.main()

From 0f52d84904ffb81297d553af9a58b51e33fed23c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 6 Sep 2019 12:44:03 -0700
Subject: [PATCH 2391/2720] Qualify uses of std::min and std::max.

These functions are currently used through the using declarations in third_party/stl/gcc3/algorithm, which harms portability.

PiperOrigin-RevId: 267652802
---
 .../data_generators/ops/pack_sequences_ops.cc         | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
index 9fcb79f43..9a5ed57eb 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
@@ -50,16 +50,16 @@ class PackSequences2Op : public OpKernel {
     int n = inputs.dimension(0);
     std::vector<int> inputs_lengths(n);
     std::vector<int> targets_lengths(n);
-    int padded_inputs_length = min(static_cast<int>(inputs.dimension(1)),
-                                   inputs_max_length);
+    int padded_inputs_length =
+        std::min(static_cast<int>(inputs.dimension(1)), inputs_max_length);
     for (int i = 0; i < n; i++) {
       for (int j = 0; j < padded_inputs_length; j++) {
           if (inputs(i, j) != 0)
             inputs_lengths[i]++;
       }
     }
-    int padded_targets_length = min(static_cast<int>(targets.dimension(1)),
-                                    targets_max_length);
+    int padded_targets_length =
+        std::min(static_cast<int>(targets.dimension(1)), targets_max_length);
     for (int i = 0; i < n; i++) {
       for (int j = 0; j < padded_targets_length; j++) {
           if (targets(i, j) != 0)
@@ -73,7 +73,8 @@ class PackSequences2Op : public OpKernel {
     for (int seq_id = 0; seq_id < n; seq_id++) {
       int inputs_length = inputs_lengths[seq_id];
       int targets_length = targets_lengths[seq_id];
-      for (int combined_id = max(0, num_combined - 10); true; combined_id++) {
+      for (int combined_id = std::max(0, num_combined - 10); true;
+           combined_id++) {
         if (combined_id == num_combined) {
           combined_inputs_length.push_back(inputs_length);
           combined_targets_length.push_back(targets_length);

From 943623b0bdd102bc98ebec472a451231a8181418 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 6 Sep 2019 15:17:32 -0700
Subject: [PATCH 2392/2720] Add the option of using holdout parts of training
 set for validation.

PiperOrigin-RevId: 267683594
---
 tensor2tensor/trax/inputs.py                  | 118 ++++++++++--------
 .../env_online_tune_wide_resnet_cifar10.gin   |   5 +
 2 files changed, 69 insertions(+), 54 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 30e76da7f..8c903ba22 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -48,9 +48,9 @@
 
 
 Inputs = collections.namedtuple(
-    "_Inputs",
-    ["train_stream", "train_eval_stream", "eval_stream",
-     "input_shape", "input_dtype"]
+    '_Inputs',
+    ['train_stream', 'train_eval_stream', 'eval_stream',
+     'input_shape', 'input_dtype']
 )
 
 # How many examples from the stream to skip at random during training.
@@ -59,7 +59,7 @@
 _MAX_SKIP_EXAMPLES = 1e5
 
 
-@gin.configurable(blacklist=["n_devices"])
+@gin.configurable(blacklist=['n_devices'])
 def inputs(n_devices, dataset_name, data_dir=None, input_name=None,
            n_chunks=0, append_targets=False):
   """Make Inputs for built-in datasets.
@@ -77,7 +77,7 @@ def inputs(n_devices, dataset_name, data_dir=None, input_name=None,
   Returns:
     trax.inputs.Inputs
   """
-  assert data_dir, "Must provide a data directory"
+  assert data_dir, 'Must provide a data directory'
   data_dir = os.path.expanduser(data_dir)
 
   (train_batches, train_eval_batches, eval_batches,
@@ -112,7 +112,7 @@ def numpy_stream(dataset):
                 input_shape=input_shape, input_dtype=input_dtype)
 
 
-@gin.configurable(blacklist=["n_devices"])
+@gin.configurable(blacklist=['n_devices'])
 def random_inputs(
     n_devices,
     input_shape=gin.REQUIRED, input_dtype=np.int32, input_range=(0, 255),
@@ -133,11 +133,11 @@ def random_inputs(
   """
   if input_shape[0] % n_devices != 0:
     tf.logging.fatal(
-        "n_devices[%d] should divide the first dimension of input_shape[%s]",
+        'n_devices[%d] should divide the first dimension of input_shape[%s]',
         n_devices, input_shape)
   if output_shape[0] % n_devices != 0:
     tf.logging.fatal(
-        "n_devices[%d] should divide the first dimension of output_shape[%s]",
+        'n_devices[%d] should divide the first dimension of output_shape[%s]',
         n_devices, output_shape)
 
   def random_minibatches():
@@ -166,7 +166,7 @@ def dataset_to_stream(dataset, input_name, n_chunks=0, append_targets=False):
   for example in backend.dataset_as_numpy(dataset):
     inp, out = example[0][input_name], example[1]
     # All input-pipeline processing should be on CPU.
-    with tf.device("cpu:0"):
+    with tf.device('cpu:0'):
       # Some accelerators don't handle uint8 well, cast to int.
       if isinstance(inp, np.uint8):
         inp = inp.astype(np.int32)
@@ -182,9 +182,10 @@ def dataset_to_stream(dataset, input_name, n_chunks=0, append_targets=False):
     yield inp, out
 
 
-@gin.configurable(whitelist=["train_shuffle_files", "test_shuffle_files"])
-def train_and_eval_dataset(dataset_name, data_dir, train_shuffle_files=True,
-                           test_shuffle_files=False):
+@gin.configurable(whitelist=['train_shuffle_files', 'eval_shuffle_files',
+                             'eval_holdout_size'])
+def train_and_eval_dataset(dataset_name, data_dir, eval_holdout_size=0,
+                           train_shuffle_files=True, eval_shuffle_files=False):
   """Return train and evaluation datasets, feature info and supervised keys.
 
   Args:
@@ -192,9 +193,11 @@ def train_and_eval_dataset(dataset_name, data_dir, train_shuffle_files=True,
       then we'll search T2T Problem registry for it, otherwise we assume it
       is a dataset from TFDS and load it from there.
     data_dir: directory where the data is located.
+    eval_holdout_size: float from 0 to <1; if >0 use this much of training data
+      for evaluation (instead of looking for a pre-specified VALIDATION split).
     train_shuffle_files: Boolean determining whether or not to shuffle the train
       files at startup. Set to False if you want data determinism.
-    test_shuffle_files: Boolean determining whether or not to shuffle the test
+    eval_shuffle_files: Boolean determining whether or not to shuffle the test
       files at startup. Set to False if you want data determinism.
 
   Returns:
@@ -206,24 +209,31 @@ def train_and_eval_dataset(dataset_name, data_dir, train_shuffle_files=True,
      * supervised_keys: information what's the input and what's the target,
          ie., a pair of lists with input and target feature names.
   """
-  if dataset_name.startswith("t2t_"):
+  if dataset_name.startswith('t2t_'):
     return _train_and_eval_dataset_v1(dataset_name[4:], data_dir)
   dataset_builder = tfds.builder(dataset_name, data_dir=data_dir)
   info = dataset_builder.info
   splits = dataset_builder.info.splits
   if tfds.Split.TRAIN not in splits:
-    raise ValueError("To train we require a train split in the dataset.")
-  if tfds.Split.VALIDATION not in splits and "test" not in splits:
-    raise ValueError("We require a validation or test split in the dataset.")
-  eval_split = tfds.Split.VALIDATION
-  if tfds.Split.VALIDATION not in splits:
-    eval_split = tfds.Split.TEST
+    raise ValueError('To train we require a train split in the dataset.')
+  train_split = tfds.Split.TRAIN
+  if eval_holdout_size > 0:
+    holdout_percentage = int(eval_holdout_size * 100.0)
+    train_percentage = 100 - holdout_percentage
+    train_split = tfds.Split.TRAIN.subsplit(tfds.percent[:train_percentage])
+    eval_split = tfds.Split.TRAIN.subsplit(tfds.percent[train_percentage:])
+  else:
+    if tfds.Split.VALIDATION not in splits and 'test' not in splits:
+      raise ValueError('We require a validation or test split in the dataset.')
+    eval_split = tfds.Split.VALIDATION
+    if tfds.Split.VALIDATION not in splits:
+      eval_split = tfds.Split.TEST
   train = tfds.load(
-      name=dataset_name, split=tfds.Split.TRAIN, data_dir=data_dir,
-      as_dataset_kwargs={"shuffle_files": train_shuffle_files})
+      name=dataset_name, split=train_split, data_dir=data_dir,
+      shuffle_files=train_shuffle_files)
   valid = tfds.load(
       name=dataset_name, split=eval_split, data_dir=data_dir,
-      as_dataset_kwargs={"shuffle_files": test_shuffle_files})
+      shuffle_files=eval_shuffle_files)
   keys = None
   if info.supervised_keys:
     keys = ([info.supervised_keys[0]], [info.supervised_keys[1]])
@@ -233,12 +243,12 @@ def train_and_eval_dataset(dataset_name, data_dir, train_shuffle_files=True,
 def _make_info(shape_list, n_classes, dtype):
   """Create an info-like tuple for feature given some shapes and vocab size."""
   feature_info = collections.namedtuple(
-      "FeatureInfo", ["shape", "n_classes", "dtype"])
+      'FeatureInfo', ['shape', 'n_classes', 'dtype'])
   cur_shape = list(shape_list[0])
   # We need to merge the provided shapes, put None where they disagree.
   for shape in shape_list:
     if len(shape) != len(cur_shape):
-      raise ValueError("Shapes need to have the same number of dimensions.")
+      raise ValueError('Shapes need to have the same number of dimensions.')
     for i in range(len(shape)):
       if cur_shape[i] is not None:
         if shape[i] != cur_shape[i]:
@@ -248,7 +258,7 @@ def _make_info(shape_list, n_classes, dtype):
 
 def _select_features(example, feature_list=None):
   """Select a subset of features from the example dict."""
-  feature_list = feature_list or ["inputs", "targets"]
+  feature_list = feature_list or ['inputs', 'targets']
   return {f: example[f] for f in feature_list if f in example}
 
 
@@ -261,7 +271,7 @@ def _eager_dataset_iterator(dataset):
 
 def _train_and_eval_dataset_v1(problem_name, data_dir):
   """Return train and evaluation datasets, feature info and supervised keys."""
-  with tf.device("cpu:0"):
+  with tf.device('cpu:0'):
     problem = t2t_problems.problem(problem_name)
     train_dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, data_dir)
     train_dataset = train_dataset.map(_select_features)
@@ -280,25 +290,25 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
       example2 = sess.run(example_tensor)
       example3 = sess.run(example_tensor)
       examples = [example1, example2, example3]
-  # We use "inputs" as input except for purely auto-regressive tasks like
-  # language models where "targets" are used as input_key.
-  input_key = "inputs" if "inputs" in examples[0] else "targets"
-  supervised_keys = ([input_key], ["targets"])
+  # We use 'inputs' as input except for purely auto-regressive tasks like
+  # language models where 'targets' are used as input_key.
+  input_key = 'inputs' if 'inputs' in examples[0] else 'targets'
+  supervised_keys = ([input_key], ['targets'])
   for example in examples:
     input_shapes.append(list(example[input_key].shape))
-    target_shapes.append(list(example["targets"].shape))
+    target_shapes.append(list(example['targets'].shape))
   input_vocab_size = hparams.vocab_size[input_key]
-  target_vocab_size = hparams.vocab_size["targets"]
+  target_vocab_size = hparams.vocab_size['targets']
   input_dtype = examples[0][input_key].dtype
-  target_dtype = examples[0]["targets"].dtype
+  target_dtype = examples[0]['targets'].dtype
   input_info = _make_info(input_shapes, input_vocab_size, input_dtype)
   target_info = _make_info(target_shapes, target_vocab_size, target_dtype)
-  info = {input_key: input_info, "targets": target_info}
+  info = {input_key: input_info, 'targets': target_info}
   return train_dataset, eval_dataset, info, supervised_keys
 
 
-@gin.configurable(blacklist=["dataset", "training", "shapes",
-                             "target_names", "n_devices"])
+@gin.configurable(blacklist=['dataset', 'training', 'shapes',
+                             'target_names', 'n_devices'])
 def batch_fun(dataset, training, shapes, target_names, n_devices,
               batch_size_per_device=32, batch_size=None, eval_batch_size=32,
               bucket_length=32, buckets=None,
@@ -319,8 +329,8 @@ def batch_fun(dataset, training, shapes, target_names, n_devices,
     for dim in target_shape:
       if dim is None:
         variable_target_shapes = True
-    tf.logging.info("Heuristically setting bucketing to %s based on shapes "
-                    "of target tensors." % variable_target_shapes)
+    tf.logging.info('Heuristically setting bucketing to %s based on shapes '
+                    'of target tensors.' % variable_target_shapes)
     if variable_target_shapes:
       bucket_boundaries = [bucket_length // 4, bucket_length // 2,
                            bucket_length, bucket_length * 2,
@@ -343,12 +353,12 @@ def batch_fun(dataset, training, shapes, target_names, n_devices,
       buckets = (bucket_boundaries, bucket_batch_sizes)
 
   if buckets:
-    tf.logging.info("Bucketing with buckets %s." % str(buckets))
+    tf.logging.info('Bucketing with buckets %s.' % str(buckets))
     def example_length(example_inputs, target):
       """The length function used by bucket_by_sequence_length to bucket."""
       other_length = 0
       if buckets_include_inputs_in_length:
-        other_length = tf.shape(example_inputs["inputs"])[0]
+        other_length = tf.shape(example_inputs['inputs'])[0]
       return tf.maximum(tf.shape(target)[0], other_length)
     boundaries, batch_sizes = buckets
     dataset = dataset.apply(tf.data.experimental.bucket_by_sequence_length(
@@ -361,19 +371,19 @@ def example_length(example_inputs, target):
   return dataset
 
 
-@gin.configurable(blacklist=["dataset", "training"])
+@gin.configurable(blacklist=['dataset', 'training'])
 def cifar10_no_augmentation_preprocess(dataset, training):
   del training
 
   def cast_image(features, targets):
-    features["image"] = tf.cast(features["image"], tf.float32) / 255.0
+    features['image'] = tf.cast(features['image'], tf.float32) / 255.0
     return features, targets
 
   dataset = dataset.map(cast_image)
   return dataset
 
 
-@gin.configurable(blacklist=["dataset", "training"])
+@gin.configurable(blacklist=['dataset', 'training'])
 def cifar10_augmentation_preprocess(dataset, training):
   """Preprocessing for cifar10 with augmentation (see below)."""
 
@@ -393,11 +403,11 @@ def augment_image(image):
     return image
 
   def augment(features, targets):
-    features["image"] = augment_image(features["image"])
+    features['image'] = augment_image(features['image'])
     return features, targets
 
   def cast_image(features, targets):
-    features["image"] = tf.cast(features["image"], tf.float32) / 255.0
+    features['image'] = tf.cast(features['image'], tf.float32) / 255.0
     return features, targets
 
   if training:
@@ -411,25 +421,25 @@ def no_preprocess(dataset, training):
   return dataset
 
 
-@gin.configurable(blacklist=["dataset", "training"])
+@gin.configurable(blacklist=['dataset', 'training'])
 def concat_preprocess(dataset, training, pad_symbol=0):
   """Pre-processing function that concatenates input and target for LM."""
   del training
 
   def concat(features, targets):
-    inp = features["inputs"]
+    inp = features['inputs']
     pad = tf.expand_dims(tf.zeros_like(inp[0]) + pad_symbol, axis=0)
     concat = tf.concat([pad, inp, pad, targets], axis=0)
     # Note: we're updating existing features dictionary here, so make sure
     # it is not re-used in some other ways outside of this function.
-    features["inputs"] = concat
+    features['inputs'] = concat
     return features, concat
 
   dataset = dataset.map(concat)
   return dataset
 
 
-@gin.configurable(blacklist=["dataset", "training"])
+@gin.configurable(blacklist=['dataset', 'training'])
 def lm1b_preprocess(dataset, training,
                     max_target_length=-1, max_eval_target_length=-1):
   """Preprocessing for LM1B: filter out targets exceeding maximum length."""
@@ -450,16 +460,16 @@ def eval_target_right_length(_, target):
 
 
 # TODO(lukaszkaiser): find a single more abstract way of text pre-processing.
-@gin.configurable(blacklist=["dataset", "training"])
+@gin.configurable(blacklist=['dataset', 'training'])
 def wmt_preprocess(dataset, training, max_length=-1, max_eval_length=-1):
   """Preprocessing for LM1B: filter out targets exceeding maximum length."""
 
   def train_right_length(example, target):
-    l = tf.maximum(tf.shape(example["inputs"])[0], tf.shape(target)[0])
+    l = tf.maximum(tf.shape(example['inputs'])[0], tf.shape(target)[0])
     return tf.less(l, max_length + 1)
 
   def eval_right_length(example, target):
-    l = tf.maximum(tf.shape(example["inputs"])[0], tf.shape(target)[0])
+    l = tf.maximum(tf.shape(example['inputs'])[0], tf.shape(target)[0])
     return tf.less(l, max_eval_length + 1)
 
   if max_length > 0 and training:
@@ -471,7 +481,7 @@ def eval_right_length(example, target):
   return dataset
 
 
-@gin.configurable(whitelist=["preprocess_fun", "shuffle_buffer_size"])
+@gin.configurable(whitelist=['preprocess_fun', 'shuffle_buffer_size'])
 def shuffle_and_batch_data(dataset,
                            target_names,
                            features_info,
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
index e67ab1b0d..3016320bd 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
@@ -16,6 +16,11 @@ batch_fun.eval_batch_size = 512
 inputs.data_dir = None
 inputs.dataset_name = 'cifar10'
 
+# Parameters for train_and_eval_dataset:
+# ==============================================================================
+train_and_eval_dataset.eval_holdout_size = 0.05
+train_and_eval_dataset.eval_shuffle_files = True
+
 # Parameters for Momentum:
 # ==============================================================================
 Momentum.mass = 0.9

From 32fe4b38cefb1997f6561e3bb94a220b8e81a498 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 6 Sep 2019 15:35:44 -0700
Subject: [PATCH 2393/2720] Repeat the evaluation dataset, so there are no
 incomplete batches

PiperOrigin-RevId: 267687032
---
 tensor2tensor/trax/inputs.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 8c903ba22..c2b12ac7b 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -499,8 +499,12 @@ def append_targets(example):
       targets[name] = example[name]
     return (example, targets)
   dataset = dataset.map(append_targets)
+  # TODO(pkozakowski): Repeat both the training and evaluation set, so we don't
+  # have incomplete batches during evaluation. This will be a problem when we
+  # add an option to evaluate on the whole dataset, then we'll need to think of
+  # a different solution.
+  dataset = dataset.repeat()
   if training:
-    dataset = dataset.repeat()
     # Skip a random fraction at the beginning of the stream.  The skip is
     # essential for synchronous highly-parallel training to avoid multiple
     # replicas reading the same data in lock-step.

From 38c7bb01f33bd85fccc584e10fa17543bdf25757 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 6 Sep 2019 16:30:35 -0700
Subject: [PATCH 2394/2720] Add orthogonal initializer to Trax.

PiperOrigin-RevId: 267696959
---
 tensor2tensor/trax/layers/initializers.py     | 37 +++++++++++++++++++
 .../trax/layers/initializers_test.py          |  5 +++
 2 files changed, 42 insertions(+)

diff --git a/tensor2tensor/trax/layers/initializers.py b/tensor2tensor/trax/layers/initializers.py
index 97d334a56..b5c8c0d56 100644
--- a/tensor2tensor/trax/layers/initializers.py
+++ b/tensor2tensor/trax/layers/initializers.py
@@ -134,3 +134,40 @@ def KaimingUniformInitializer(out_dim=-1, in_dim=-2, param=0.):
   return VarianceScalingInitializer(out_dim, in_dim,
                                     2.0 / backend.numpy.sqrt(1 + param**2),
                                     'fan_in', 'uniform')
+
+
+def OrthogonalInitializer(stddev=1.0):
+  """Orthogonal Initializer."""
+  def Init(shape, rng):
+    """The orthogonal initializer function."""
+    # Have at least 2 elements in shape.
+    cur_shape = list(shape)
+    while len(cur_shape) < 2:
+      cur_shape = [1] + cur_shape
+
+    # Flatten the input shape with the last dimension remaining.
+    n_rows = 1
+    for dim in cur_shape[:-1]:
+      n_rows *= dim
+    n_cols = cur_shape[-1]
+    flat_shape = (n_cols, n_rows) if n_rows < n_cols else (n_rows, n_cols)
+
+    # Generate a random matrix
+    a = backend.random.normal(rng, flat_shape, dtype=backend.numpy.float32)
+
+    # Compute the qr factorization
+    q, r = backend.numpy.linalg.qr(a)
+
+    # Make Q uniform
+    d = backend.numpy.diag(r)
+    q *= backend.numpy.sign(d)
+
+    # Transpose and reshape back q if needed.
+    if n_rows < n_cols:
+      q = backend.numpy.transpose(q)
+    q = backend.numpy.reshape(q, shape)
+
+    # Return scaled as requested.
+    return stddev * q
+
+  return Init
diff --git a/tensor2tensor/trax/layers/initializers_test.py b/tensor2tensor/trax/layers/initializers_test.py
index 6a0bef4c4..469fb29fa 100644
--- a/tensor2tensor/trax/layers/initializers_test.py
+++ b/tensor2tensor/trax/layers/initializers_test.py
@@ -73,6 +73,11 @@ def test_kaiming_uniform(self):
     init_value = initializer(input_shape, random.get_prng(0))
     self.assertEqual(tuple(init_value.shape), input_shape)
 
+  def test_orthogonal(self):
+    initializer = initializers.OrthogonalInitializer()
+    input_shape = (29, 5, 7, 20)
+    init_value = initializer(input_shape, random.get_prng(0))
+    self.assertEqual(tuple(init_value.shape), input_shape)
 
 if __name__ == "__main__":
   absltest.main()

From ddd02b62a50b9a570d97250883d986b20e3cd8f9 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 6 Sep 2019 22:37:09 -0700
Subject: [PATCH 2395/2720] Add a basic MNIST-MLP config and use Adafactor by
 default in Trax.

PiperOrigin-RevId: 267733310
---
 tensor2tensor/trax/configs/mlp_mnist.gin | 35 ++++++++++++++++++++++++
 tensor2tensor/trax/learning_rate.py      |  2 +-
 tensor2tensor/trax/optimizers/base.py    |  3 +-
 tensor2tensor/trax/trax.py               |  2 +-
 4 files changed, 39 insertions(+), 3 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/mlp_mnist.gin

diff --git a/tensor2tensor/trax/configs/mlp_mnist.gin b/tensor2tensor/trax/configs/mlp_mnist.gin
new file mode 100644
index 000000000..a8bee6abd
--- /dev/null
+++ b/tensor2tensor/trax/configs/mlp_mnist.gin
@@ -0,0 +1,35 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.learning_rate
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 256
+batch_fun.eval_batch_size = 256
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 'mnist'
+
+# Parameters for MLP:
+# ==============================================================================
+MLP.d_hidden = 512
+MLP.n_hidden_layers = 2
+MLP.n_output_classes = 10
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 0.1
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 400
+
+# Parameters for train:
+# ==============================================================================
+train.optimizer = @trax.optimizers.Adafactor
+train.eval_frequency = 200
+train.eval_steps = 10
+train.model = @trax.models.MLP
+train.train_steps = 2000
diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index cb5ae605f..176e363f7 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -32,7 +32,7 @@
 
 @gin.configurable(blacklist=["history"])
 def MultifactorSchedule(history=None,
-                        factors="constant * linear_warmup",
+                        factors="constant * linear_warmup * rsqrt_decay",
                         constant=0.1,
                         warmup_steps=400,
                         decay_factor=0.5,
diff --git a/tensor2tensor/trax/optimizers/base.py b/tensor2tensor/trax/optimizers/base.py
index 992331e8d..21a2cdcae 100644
--- a/tensor2tensor/trax/optimizers/base.py
+++ b/tensor2tensor/trax/optimizers/base.py
@@ -347,7 +347,8 @@ def update(self, step, grads, params, slots, opt_params):
       updates.append(new_m)
 
     new_params = (1 - weight_decay_rate) * params - subtrahend
-    return new_params, updates
+    # TODO(lukaszkaiser): why is the astype needed here? Check and correct.
+    return new_params.astype(params.dtype), updates
 
 
 class SM3(Optimizer):
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 9da70f266..8930fdcbd 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -833,7 +833,7 @@ def train(output_dir,
           model=gin.REQUIRED,
           loss_fn=loss,
           inputs=trax_inputs.inputs,
-          optimizer=trax_opt.SM3,
+          optimizer=trax_opt.Adafactor,
           lr_schedule=lr.MultifactorSchedule,
           trainer_class=Trainer,
           train_steps=1000,

From 39f18f1cee6b6d460d031d21b1964e0ba2221650 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 6 Sep 2019 23:48:00 -0700
Subject: [PATCH 2396/2720] s/tf.logging/absl.logging/gym_env_problem.py

PiperOrigin-RevId: 267737846
---
 tensor2tensor/envs/gym_env_problem.py | 32 ++++++++++++++-------------
 1 file changed, 17 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/envs/gym_env_problem.py b/tensor2tensor/envs/gym_env_problem.py
index 3c8437091..eb4f3b619 100644
--- a/tensor2tensor/envs/gym_env_problem.py
+++ b/tensor2tensor/envs/gym_env_problem.py
@@ -26,12 +26,11 @@
 import copy
 import multiprocessing.pool
 import time
-
+from absl import logging
 import gym
 import numpy as np
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import trajectory
-import tensorflow as tf
 
 
 class GymEnvProblem(env_problem.EnvProblem):
@@ -76,7 +75,10 @@ class GymEnvProblem(env_problem.EnvProblem):
   the following properties: observation_space, action_space, reward_range.
   """
 
-  def __init__(self, base_env_name=None, env_wrapper_fn=None, reward_range=None,
+  def __init__(self,
+               base_env_name=None,
+               env_wrapper_fn=None,
+               reward_range=None,
                **kwargs):
     """Initializes this class by creating the envs and managing trajectories.
 
@@ -133,8 +135,8 @@ def _verify_same_spaces(self):
       raise ValueError("Environments not initialized.")
 
     if not isinstance(self._envs, list):
-      tf.logging.warning("Not checking observation and action space "
-                         "compatibility across envs, since there is just one.")
+      logging.warning("Not checking observation and action space "
+                      "compatibility across envs, since there is just one.")
       return
 
     # NOTE: We compare string representations of observation_space and
@@ -146,20 +148,20 @@ def _verify_same_spaces(self):
         for env in self._envs):
       err_str = ("All environments should have the same observation space, but "
                  "don't.")
-      tf.logging.error(err_str)
+      logging.error(err_str)
       # Log all observation spaces.
       for i, env in enumerate(self._envs):
-        tf.logging.error("Env[%d] has observation space [%s]", i,
-                         env.observation_space)
+        logging.error("Env[%d] has observation space [%s]", i,
+                      env.observation_space)
       raise ValueError(err_str)
 
     if not all(
         str(env.action_space) == str(self.action_space) for env in self._envs):
       err_str = "All environments should have the same action space, but don't."
-      tf.logging.error(err_str)
+      logging.error(err_str)
       # Log all action spaces.
       for i, env in enumerate(self._envs):
-        tf.logging.error("Env[%d] has action space [%s]", i, env.action_space)
+        logging.error("Env[%d] has action space [%s]", i, env.action_space)
       raise ValueError(err_str)
 
   def initialize_environments(self,
@@ -235,14 +237,14 @@ def reward_range(self):
 
   def seed(self, seed=None):
     if not self._envs:
-      tf.logging.info("`seed` called on non-existent envs, doing nothing.")
+      logging.info("`seed` called on non-existent envs, doing nothing.")
       return None
 
     if not isinstance(self._envs, list):
-      tf.logging.warning("`seed` called on non-list envs, doing nothing.")
+      logging.warning("`seed` called on non-list envs, doing nothing.")
       return None
 
-    tf.logging.warning(
+    logging.warning(
         "Called `seed` on EnvProblem, calling seed on the underlying envs.")
     for env in self._envs:
       env.seed(seed)
@@ -251,11 +253,11 @@ def seed(self, seed=None):
 
   def close(self):
     if not self._envs:
-      tf.logging.info("`close` called on non-existent envs, doing nothing.")
+      logging.info("`close` called on non-existent envs, doing nothing.")
       return
 
     if not isinstance(self._envs, list):
-      tf.logging.warning("`close` called on non-list envs, doing nothing.")
+      logging.warning("`close` called on non-list envs, doing nothing.")
       return
 
     # Call close on all the envs one by one.

From 648184c46abb86239d8fdb317d7921d4bae8800c Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Sat, 7 Sep 2019 15:31:57 -0700
Subject: [PATCH 2397/2720] Clean up model evaluation in SimPLe and add tests

PiperOrigin-RevId: 267798484
---
 tensor2tensor/trax/rl/simple.py      |  82 +++++++----
 tensor2tensor/trax/rl/simple_test.py | 197 +++++++++++++++++++++++++++
 2 files changed, 253 insertions(+), 26 deletions(-)

diff --git a/tensor2tensor/trax/rl/simple.py b/tensor2tensor/trax/rl/simple.py
index 711131213..d5f025bf7 100644
--- a/tensor2tensor/trax/rl/simple.py
+++ b/tensor2tensor/trax/rl/simple.py
@@ -97,8 +97,10 @@ def make_batch(examples):
 
 # TODO(pkozakowski): This is mostly a simplified version of
 # env_problem_utils.play_env_problem_with_policy, generalized to work with
-# policies not being neural networks. Unify if possible.
-def play_env_problem(env, policy):
+# policies not being neural networks. Another difference is that it always
+# collects exactly one trajectory from each environment in the batch. Unify if
+# possible.
+def play_env_problem(env, policy_fn):
   """Plays an EnvProblem using a given policy function."""
   trajectories = [trajectory.Trajectory() for _ in range(env.batch_size)]
   observations = env.reset()
@@ -109,7 +111,7 @@ def play_env_problem(env, policy):
   while not np.all(done_so_far):
     padded_observations, _ = env.trajectories.observations_np(
         len_history_for_policy=None)
-    actions = policy(padded_observations)
+    actions = policy_fn(padded_observations)
     (observations, rewards, dones, _) = env.step(actions)
     for (traj, observation, action, reward, done) in zip(
         trajectories, observations, actions, rewards, dones
@@ -125,19 +127,22 @@ def play_env_problem(env, policy):
 
 def calculate_observation_error(real_trajectories, sim_trajectories):
   """Calculates MSE of observations in two trajectories."""
-  def padded_obs(traj, length_difference):
-    return np.pad(
-        traj.observations_np,
-        pad_width=((0, max(length_difference, 0)), (0, 0)),
-        mode="edge",
-    )
+  def pad_or_truncate(observations, desired_length):
+    (current_length, _) = observations.shape
+    if current_length < desired_length:
+      return np.pad(
+          observations,
+          pad_width=((0, desired_length - current_length), (0, 0)),
+          mode="edge",
+      )
+    else:
+      return observations[:desired_length, :]
 
   def calculate_for_single_pair(real_trajectory, sim_trajectory):
-    diff = sim_trajectory.num_time_steps - real_trajectory.num_time_steps
-    padded_real_obs = padded_obs(real_trajectory, diff)
-    padded_sim_obs = padded_obs(sim_trajectory, -diff)
-    x = np.sum((padded_real_obs - padded_sim_obs) ** 2, axis=0)
-    return x
+    real_obs = real_trajectory.observations_np
+    sim_obs = pad_or_truncate(
+        sim_trajectory.observations_np, real_trajectory.num_time_steps)
+    return np.sum((real_obs - sim_obs) ** 2, axis=0)
 
   return np.mean([
       calculate_for_single_pair(real_traj, sim_traj)
@@ -162,10 +167,40 @@ def plot_observation_error(real_trajectories, sim_trajectories, mpl_plt):
       for (traj, label) in ((real_traj, "real"), (sim_traj, "simulated")):
         obs = traj.observations_np
         ax = axes[dim_index, traj_index]
+        ax.set_title("trajectory {}, observation dimension {}".format(
+            traj_index, dim_index))
         ax.plot(np.arange(obs.shape[0]), obs[:, dim_index], label=label)
         ax.legend()
 
 
+class ReplayPolicy(object):
+  """Policy function repeating actions from a given batch of trajectories."""
+
+  def __init__(self, trajectories, out_of_bounds_action):
+    """Creates ReplayPolicy.
+
+    Args:
+      trajectories: Batch of trajectories to repeat actions from.
+      out_of_bounds_action: Action to play after the replayed trajectory ends.
+    """
+    self._trajectories = trajectories
+    self._out_of_bounds_action = out_of_bounds_action
+    self._step = 0
+
+  def __call__(self, observations):
+    del observations
+
+    def get_action(traj):
+      if self._step < traj.num_time_steps:
+        action = traj.time_steps[self._step].action
+      else:
+        action = None
+      return action or self._out_of_bounds_action
+    actions = np.array(list(map(get_action, self._trajectories)))
+    self._step += 1
+    return actions
+
+
 def evaluate_model(sim_env, real_trajectories, mpl_plt, n_to_plot=3):
   """Reports the observation error metric and the corresponding plot."""
   if len(sim_env.observation_space.shape) != 1:
@@ -177,19 +212,14 @@ def evaluate_model(sim_env, real_trajectories, mpl_plt, n_to_plot=3):
 
   assert len(real_trajectories) == sim_env.batch_size
 
-  step = [0]
-  def policy(observations):
-    del observations
-    def get_action(traj):
-      if step[0] < traj.num_time_steps:
-        return traj.time_steps[step[0]].action or 0
-      else:
-        return 0
-    actions = np.array([get_action(traj) for traj in real_trajectories])
-    step[0] += 1
-    return actions
+  policy_fn = ReplayPolicy(
+      real_trajectories,
+      # Does not matter which action we play after the real trajetory ends, we
+      # cut the simulated one to match the real one anyway.
+      out_of_bounds_action=sim_env.action_space.sample(),
+  )
 
-  sim_trajectories = play_env_problem(sim_env, policy)
+  sim_trajectories = play_env_problem(sim_env, policy_fn)
   obs_errors = calculate_observation_error(real_trajectories, sim_trajectories)
   plot_observation_error(
       real_trajectories[:n_to_plot], sim_trajectories[:n_to_plot], mpl_plt)
diff --git a/tensor2tensor/trax/rl/simple_test.py b/tensor2tensor/trax/rl/simple_test.py
index b44b6853d..8d5efea2d 100644
--- a/tensor2tensor/trax/rl/simple_test.py
+++ b/tensor2tensor/trax/rl/simple_test.py
@@ -23,9 +23,18 @@
 import os
 
 import cloudpickle as pickle
+import gin
+import gym
+from matplotlib import pyplot as plt
+import mock
 import numpy as np
+
 from tensor2tensor.envs import trajectory
+from tensor2tensor.trax import backend
+from tensor2tensor.trax import trax
 from tensor2tensor.trax.rl import simple
+from tensor2tensor.trax.rl import simulated_env_problem
+from tensor2tensor.trax.rl import space_serializer  # pylint: disable=unused-import
 from tensorflow import test
 from tensorflow.io import gfile
 
@@ -102,6 +111,194 @@ def test_batches_stream(self):
     np.testing.assert_equal(
         next(batched_stream), (np.array([4, 6]), np.array([5, 7])))
 
+  def test_plays_env_problem(self):
+    # Shape: (time, trajectory).
+    observations = np.array([[0, 1], [2, 3], [4, 5]])
+    rewards = np.array([[0, 1], [1, 0]])
+    actions = np.array([[1, 2], [2, 0]])
+    # We end the second environment 2 times, but we shouldn't collect the second
+    # trajectory.
+    dones = np.array([[False, True], [True, True]])
+    infos = [{}, {}]
+
+    mock_env = mock.MagicMock()
+    mock_env.batch_size = 2
+    # (observations, lengths)
+    mock_env.trajectories.observations_np.return_value = (None, None)
+    mock_env.reset.return_value = observations[0]
+    mock_env.step.side_effect = zip(observations[1:], rewards, dones, infos)
+
+    mock_policy_fn = mock.MagicMock()
+    mock_policy_fn.side_effect = actions
+
+    trajectories = simple.play_env_problem(mock_env, mock_policy_fn)
+    self.assertEqual(len(trajectories), 2)
+    expected_lengths = [3, 2]
+    for (i, (traj, expected_length)) in enumerate(
+        zip(trajectories, expected_lengths)):
+      self.assertEqual(traj.num_time_steps, expected_length)
+      np.testing.assert_array_equal(
+          traj.observations_np, observations[:expected_length, i])
+      np.testing.assert_array_equal(
+          traj.raw_rewards_np, rewards[:(expected_length - 1), i])
+      np.testing.assert_array_equal(
+          traj.actions_np, actions[:(expected_length - 1), i])
+
+  def _make_trajectory(self, observations=None, actions=None):
+    t = trajectory.Trajectory()
+    if observations is None:
+      observations = itertools.repeat(None)
+    if actions is None:
+      actions = itertools.repeat(None)
+    for (observation, action) in zip(observations, actions):
+      t.add_time_step(observation=observation, action=action)
+    return t
+
+  def test_replay_policy(self):
+    trajectories = [
+        self._make_trajectory(actions=actions)
+        for actions in map(np.array, [[1, 2], [3]])
+    ]
+    policy_fn = simple.ReplayPolicy(trajectories, out_of_bounds_action=0)
+    np.testing.assert_array_equal(policy_fn(None), [1, 3])
+    np.testing.assert_array_equal(policy_fn(None), [2, 0])
+
+  def test_observation_error_zero_for_same_trajectories(self):
+    observations = np.array([[0], [2], [1]])
+    (traj1, traj2) = map(self._make_trajectory, (observations, observations))
+    error = simple.calculate_observation_error([traj1], [traj2])
+    np.testing.assert_array_almost_equal(error, [0])
+
+  def test_observation_error_positive_for_different_trajectories(self):
+    observations1 = np.array([[1], [2], [3]])
+    observations2 = np.array([[0], [2], [3]])
+    (traj1, traj2) = map(self._make_trajectory, (observations1, observations2))
+    error = simple.calculate_observation_error([traj1], [traj2])
+    np.testing.assert_array_less([0], error)
+
+  def test_observation_error_dims_correspond_to_observation_dims(self):
+    observations1 = np.array([[0, 1, 0], [0, 2, 0], [0, 3, 0]])
+    observations2 = np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]])
+    (traj1, traj2) = map(self._make_trajectory, (observations1, observations2))
+    error = simple.calculate_observation_error([traj1], [traj2])
+    self.assertEqual(error.shape, (3,))
+    np.testing.assert_array_almost_equal(error[0], 0)
+    self.assertFalse(np.allclose(error[1], 0))
+    np.testing.assert_array_almost_equal(error[2], 0)
+
+  def test_observation_error_increases_with_distance(self):
+    observations_zero = np.array([[0], [0], [0]])
+    observations_positive = np.array([[3], [2], [1]])
+    (traj_zero, traj_positive, traj_negative) = map(
+        self._make_trajectory,
+        (observations_zero, observations_positive, -observations_positive),
+    )
+    error_small = simple.calculate_observation_error(
+        [traj_zero], [traj_positive])
+    error_big = simple.calculate_observation_error(
+        [traj_positive], [traj_negative])
+    np.testing.assert_array_less(error_small, error_big)
+
+  def test_observation_error_increases_with_real_trajectory_length(self):
+    observations_real_short = np.array([[1], [2]])
+    observations_real_long = np.array([[1], [2], [3]])
+    observations_sim = np.array([[0], [1]])
+    (traj_real_short, traj_real_long, traj_sim) = map(
+        self._make_trajectory,
+        (observations_real_short, observations_real_long, observations_sim),
+    )
+    error_small = simple.calculate_observation_error(
+        real_trajectories=[traj_real_short], sim_trajectories=[traj_sim])
+    error_big = simple.calculate_observation_error(
+        real_trajectories=[traj_real_long], sim_trajectories=[traj_sim])
+    np.testing.assert_array_less(error_small, error_big)
+
+  def test_observation_error_same_when_sim_trajectory_longer(self):
+    observations_real = np.array([[0], [1]])
+    observations_sim_short = np.array([[1], [2]])
+    observations_sim_long = np.array([[1], [2], [3]])
+    (traj_real, traj_sim_short, traj_sim_long) = map(
+        self._make_trajectory,
+        (observations_real, observations_sim_short, observations_sim_long),
+    )
+    error1 = simple.calculate_observation_error(
+        real_trajectories=[traj_real], sim_trajectories=[traj_sim_short])
+    error2 = simple.calculate_observation_error(
+        real_trajectories=[traj_real], sim_trajectories=[traj_sim_long])
+    np.testing.assert_array_almost_equal(error1, error2)
+
+  def test_observation_error_reduces_over_trajectories(self):
+    observations1 = np.array([[1], [2], [3]])
+    observations2 = np.array([[0], [2], [3]])
+    (traj1, traj2) = map(self._make_trajectory, (observations1, observations2))
+    error = simple.calculate_observation_error([traj1, traj1], [traj2, traj2])
+    self.assertEqual(error.shape, (1,))
+
+  @staticmethod
+  @mock.patch.object(trax, "restore_state", autospec=True)
+  def _make_env(
+      mock_restore_state, observation_space, action_space,
+      max_trajectory_length, batch_size,
+  ):
+    # (model_params, opt_state)
+    mock_restore_state.return_value.params = (None, None)
+
+    gin.bind_parameter("BoxSpaceSerializer.precision", 1)
+
+    seq_length = max_trajectory_length * int(
+        np.prod(observation_space.shape) + np.prod(action_space.shape))
+    predict_output = (np.array([[[0.0]] * seq_length]), ())
+    mock_model_fn = mock.MagicMock()
+    mock_model_fn.return_value.side_effect = itertools.repeat(predict_output)
+
+    return simulated_env_problem.SerializedSequenceSimulatedEnvProblem(
+        model=mock_model_fn,
+        reward_fn=(lambda _1, _2: np.zeros(batch_size)),
+        done_fn=(lambda _1, _2: np.full((batch_size,), False)),
+        vocab_size=1,
+        max_trajectory_length=max_trajectory_length,
+        batch_size=batch_size,
+        observation_space=observation_space,
+        action_space=action_space,
+        reward_range=(-1, 1),
+        discrete_rewards=False,
+        history_stream=itertools.repeat(None),
+        output_dir=None,
+    )
+
+  def test_evaluates_model_with_vector_observation_space(self):
+    with backend.use_backend("numpy"):
+      env = self._make_env(  # pylint: disable=no-value-for-parameter
+          observation_space=gym.spaces.Box(shape=(2,), low=0, high=1),
+          action_space=gym.spaces.Discrete(n=1),
+          max_trajectory_length=2,
+          batch_size=3,
+      )
+      trajectories = [
+          self._make_trajectory(observations, actions)  # pylint: disable=g-complex-comprehension
+          for (observations, actions) in [
+              (np.array([[0, 1]]), np.array([0])),
+              (np.array([[1, 2], [3, 4]]), np.array([0, 0])),
+              (np.array([[1, 2], [3, 4], [5, 6]]), np.array([0, 0, 0])),
+          ]
+      ]
+      metrics = simple.evaluate_model(env, trajectories, plt)
+      self.assertIsNotNone(metrics)
+      self.assertEqual(len(metrics), 2)
+
+  def test_fails_to_evaluate_model_with_matrix_observation_space(self):
+    with backend.use_backend("numpy"):
+      env = self._make_env(  # pylint: disable=no-value-for-parameter
+          observation_space=gym.spaces.Box(shape=(2, 2), low=0, high=1),
+          action_space=gym.spaces.Discrete(n=1),
+          max_trajectory_length=2,
+          batch_size=1,
+      )
+      trajectories = [
+          self._make_trajectory(np.array([[0, 1], [2, 3]]), np.array([0]))]
+      metrics = simple.evaluate_model(env, trajectories, plt)
+      self.assertIsNone(metrics)
+
 
 if __name__ == "__main__":
   test.main()

From 236b8921e30ae85b2a34afcd4892d9a87d01e224 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Sat, 7 Sep 2019 23:58:25 -0700
Subject: [PATCH 2398/2720] Make trax.train() train only up to train_steps in
 case of a restart

PiperOrigin-RevId: 267829806
---
 tensor2tensor/trax/rl/simple_trainer.py |  7 +---
 tensor2tensor/trax/trax.py              | 44 +++++++++++++++----------
 tensor2tensor/trax/trax_test.py         | 20 ++++++++++-
 3 files changed, 46 insertions(+), 25 deletions(-)

diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index 32eb92944..539307c0c 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -149,12 +149,7 @@ def train_model(self):
     trax.train(
         model=self._sim_env.model,
         inputs=inputs,
-        # TODO(pkozakowski): Currently trax.train trains the model for
-        # train_steps more steps, whereas it should train up to train_steps
-        # total steps in order for the restarts to work properly. Change the
-        # argument once this behavior is changed.
-        # train_steps=self._model_train_step,
-        train_steps=self._n_model_train_steps,
+        train_steps=self._model_train_step,
         output_dir=self._model_dir,
         has_weights=True,
     )
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 8930fdcbd..5f4ca6715 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -358,28 +358,36 @@ def get_random_number_generator_and_set_seed(seed=None):
 # * Allow disabling eval
 
 
-def epochs(steps=None, epoch_steps=1):
-  """Iterator over epochs until steps is reached. 1-indexed.
+def epochs(total_steps, steps_to_skip, epoch_steps):
+  """Generates the number of steps in each epoch before reaching total_steps.
 
   Args:
-    steps: int, total number of steps. Infinite if None.
-    epoch_steps: int, number of steps per epoch. Can also be an iterable<int> to
-      enable variable length epochs.
+    total_steps: int, total number of steps.
+    steps_to_skip: int, number of steps to skip because of a restart.
+    epoch_steps: iterable of int, numbers of steps in each epoch.
 
   Yields:
-    (epoch: int, epoch id, epoch_steps: int, number of steps in this epoch)
+    epoch_steps: int, number of steps in this epoch
   """
-  try:
-    iter(epoch_steps)
-  except TypeError:
-    epoch_steps = itertools.repeat(epoch_steps)
-
-  step = 0
-  for epoch, epoch_steps in enumerate(epoch_steps):
-    epoch_steps = min(epoch_steps, steps - step)
-    yield (epoch + 1, epoch_steps)
-    step += epoch_steps
-    if steps and step >= steps:
+  steps_to_go = total_steps - steps_to_skip
+  epoch_steps = iter(epoch_steps)
+
+  # Remove the desired number of steps from the stream.
+  for steps_this_epoch in epoch_steps:
+    if steps_this_epoch > steps_to_skip:
+      # Put back the number of steps left in the unfinished epoch.
+      epoch_steps = itertools.chain(
+          [steps_this_epoch - steps_to_skip], epoch_steps)
+    if steps_this_epoch >= steps_to_skip:
+      break
+    steps_to_skip -= steps_this_epoch
+
+  # Yield the remaining steps per epoch up to total_steps.
+  for steps_this_epoch in epoch_steps:
+    steps_this_epoch = min(steps_this_epoch, steps_to_go)
+    yield steps_this_epoch
+    steps_to_go -= steps_this_epoch
+    if steps_to_go == 0:
       break
 
 
@@ -885,7 +893,7 @@ def train(output_dir,
   step_log(trainer.step,
            "Starting training using %d devices" % trainer.n_devices)
 
-  for _, epoch_steps in epochs(train_steps, epoch_steps):
+  for epoch_steps in epochs(train_steps, trainer.step, epoch_steps):
     trainer.train_epoch(epoch_steps, eval_steps)
 
     # Update learning rate with new history
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index 60f83ab90..34fce4440 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -185,7 +185,7 @@ def test_train_restart(self, backend_name):
       state = trax.train(output_dir,
                          model=model_fn,
                          inputs=inputs,
-                         train_steps=train_steps,
+                         train_steps=(2 * train_steps),
                          eval_steps=eval_steps)
 
       # Assert total train steps
@@ -242,6 +242,24 @@ def test_reset_twice(self, backend_name):
       trainer.evaluate(1)
 
 
+class EpochsTest(test.TestCase):
+
+  def test_cuts_epoch_when_total_steps_reached(self):
+    epoch_steps = trax.epochs(
+        total_steps=5, steps_to_skip=0, epoch_steps=[1, 2, 3])
+    self.assertEqual(list(epoch_steps), [1, 2, 2])
+
+  def test_skips_full_epoch(self):
+    epoch_steps = trax.epochs(
+        total_steps=4, steps_to_skip=2, epoch_steps=[2, 2])
+    self.assertEqual(list(epoch_steps), [2])
+
+  def test_skips_part_of_epoch(self):
+    epoch_steps = trax.epochs(
+        total_steps=4, steps_to_skip=1, epoch_steps=[2, 2])
+    self.assertEqual(list(epoch_steps), [1, 2])
+
+
 MASKED_MEAN_TEST_BACKENDS = ["numpy"]
 
 

From 2f68a91f2abe3dffbbf84d950335ab9f9d7427d5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 9 Sep 2019 09:47:45 -0700
Subject: [PATCH 2399/2720] Use logsumexp for softmax in
 MergedHashedCausalAttention

PiperOrigin-RevId: 268016575
---
 tensor2tensor/trax/layers/attention.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index d723b6e93..32a6910d4 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -674,8 +674,7 @@ def binned_attn(sqk, sv):  # pylint: disable=invalid-name
       dots = dots - 32 * self_mask
 
       # Softmax.
-      dots = np.exp(dots - dots.max(axis=-1, keepdims=True))
-      dots = dots / dots.sum(axis=-1, keepdims=True)
+      dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))
       bo = np.matmul(dots, bv)
 
       so = unchunk_vectors(bo)

From 51fc796aad0259bc9622950fc2f0750e25c9af59 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Mon, 9 Sep 2019 10:01:13 -0700
Subject: [PATCH 2400/2720] Make PPO saving work predictably with
 BaseTrainer.training_loop()

It previously bumped the epoch counter when saving after an epoch finished, breaking restarts in SimPLe. Also added a parameter save_every_n to PPO, so we can save checkpoints more often than evaluate.

PiperOrigin-RevId: 268019096
---
 .../trax/rl/configs/ppo_online_tune.gin       |  1 +
 .../trax/rl/configs/simple_online_tune.gin    |  1 +
 tensor2tensor/trax/rl/ppo.py                  | 13 ++---
 tensor2tensor/trax/rl/ppo_trainer.py          | 50 ++++++++++---------
 tensor2tensor/trax/rl/ppo_trainer_test.py     | 10 ++--
 5 files changed, 39 insertions(+), 36 deletions(-)

diff --git a/tensor2tensor/trax/rl/configs/ppo_online_tune.gin b/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
index 665e07052..edc1dc6a2 100644
--- a/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
+++ b/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
@@ -32,6 +32,7 @@ PPO.c2 = 0.01
 PPO.done_frac_for_policy_save = 0
 PPO.len_history_for_policy = None
 PPO.separate_eval = False
+PPO.save_every_n = 1
 PPO.policy_and_value_model = @trax.models.TransformerDecoder
 PPO.policy_and_value_optimizer = @trax.optimizers.Adam
 
diff --git a/tensor2tensor/trax/rl/configs/simple_online_tune.gin b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
index 0fb13480a..70d2e63e2 100644
--- a/tensor2tensor/trax/rl/configs/simple_online_tune.gin
+++ b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
@@ -36,6 +36,7 @@ PPO.c2 = 0.01
 PPO.done_frac_for_policy_save = 0
 PPO.len_history_for_policy = None
 PPO.separate_eval = False
+PPO.save_every_n = 1
 PPO.policy_and_value_model = @trax.models.TransformerDecoder
 
 # Parameters for SerializedSequenceSimulatedEnvProblem:
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index 959a27d8f..d42d9f4a4 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -824,13 +824,11 @@ def maybe_restore_opt_state(output_dir, policy_and_value_opt_state,
     policy_and_value_state: state of the policy and value network.
 
   Returns:
-    tuple (restored (bool), opt_state, state, epoch (int),
-    opt_step (int)) where epoch is the epoch from which we restored the
-    optimization state, 0 is restored = False, and opt_step is the total
-    optimization step (sum of all optimization steps made up to the current
-    epoch).
+    tuple (opt_state, state, epoch (int), opt_step (int)) where epoch is the
+    epoch from which we restored the optimization state, 0 if no checkpoint was
+    found, and opt_step is the total optimization step (sum of all optimization
+    steps made up to the current epoch).
   """
-  restored = False
   epoch = 0
   total_opt_step = 0
   model_files = gfile.glob(os.path.join(output_dir, "model-??????.pkl"))
@@ -841,7 +839,6 @@ def maybe_restore_opt_state(output_dir, policy_and_value_opt_state,
         policy_and_value_opt_state, policy_and_value_state, total_opt_step = (
             pickle.load(f))
       model_file_basename = os.path.basename(model_file)  # model-??????.pkl
-      restored = True
       epoch = int(filter(str.isdigit, model_file_basename))
       break
     except EOFError as e:
@@ -849,7 +846,7 @@ def maybe_restore_opt_state(output_dir, policy_and_value_opt_state,
       # Try an older version.
       continue
   return (
-      restored, policy_and_value_opt_state, policy_and_value_state, epoch,
+      policy_and_value_opt_state, policy_and_value_state, epoch,
       total_opt_step,
   )
 
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index ede56f5d7..2693c3b40 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -72,6 +72,7 @@ def __init__(
       c1=1.0,
       c2=0.01,
       eval_every_n=1000,
+      save_every_n=1000,
       done_frac_for_policy_save=0.5,
       n_evals=1,
       len_history_for_policy=4,
@@ -106,6 +107,7 @@ def __init__(
       c1: Value loss coefficient.
       c2: Entropy loss coefficient.
       eval_every_n: How frequently to eval the policy.
+      save_every_n: How frequently to save the policy.
       done_frac_for_policy_save: Fraction of the trajectories that should be
         done to checkpoint the policy.
       n_evals: Number of times to evaluate.
@@ -134,6 +136,7 @@ def __init__(
     self._c1 = c1
     self._c2 = c2
     self._eval_every_n = eval_every_n
+    self._save_every_n = save_every_n
     self._done_frac_for_policy_save = done_frac_for_policy_save
     self._n_evals = n_evals
     self._len_history_for_policy = len_history_for_policy
@@ -170,15 +173,13 @@ def __init__(
          policy_and_value_optimizer, policy_and_value_net_params)
 
     # Maybe restore the optimization state. If there is nothing to restore, then
-    # iteration = 0 and policy_and_value_opt_state is returned as is.
-    (restored, self._policy_and_value_opt_state, self._model_state, self._epoch,
+    # epoch = 0 and policy_and_value_opt_state is returned as is.
+    (self._policy_and_value_opt_state, self._model_state, self._epoch,
      self._total_opt_step) = ppo.maybe_restore_opt_state(
          output_dir, policy_and_value_opt_state, self._model_state)
 
-    if restored:
-      logging.info("Restored parameters from iteration [%d]", self._epoch)
-      # We should start from the next iteration.
-      self._epoch += 1
+    if self._epoch > 0:
+      logging.info("Restored parameters from epoch [%d]", self._epoch)
 
     # Create summary writers and history.
     self._train_sw = jaxboard.SummaryWriter(
@@ -455,29 +456,33 @@ def train_epoch(self, evaluate=True):
     for (name, value) in summaries.items():
       self._train_sw.scalar("train/{}".format(name), value, step=self._epoch)
 
+    logging.info(
+        "PPO epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
+        " Loss(ppo, value, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)]", self._epoch,
+        min_reward, max_reward, avg_reward, combined_loss, ppo_loss, value_loss,
+        entropy_bonus)
+
+    # Bump the epoch counter before saving a checkpoint, so that a call to
+    # save() after the training loop is a no-op if a checkpoint was saved last
+    # epoch - otherwise it would bump the epoch counter on the checkpoint.
+    last_epoch = self._epoch
+    self._epoch += 1
+
     # Save parameters every time we see the end of at least a fraction of batch
     # number of trajectories that are done (not completed -- completed includes
     # truncated and done).
     # Also don't save too frequently, enforce a minimum gap.
-    # Or if this is the last iteration.
     policy_save_start_time = time.time()
     self._n_trajectories_done += n_done
     # TODO(afrozm): Refactor to trax.save_state.
-    if ((self._n_trajectories_done >=
-         self._done_frac_for_policy_save * self.train_env.batch_size) and
-        (self._epoch - self._last_saved_at > self._eval_every_n) and
-        (((self._epoch + 1) % self._eval_every_n == 0))):
+    if (self._n_trajectories_done >=
+        self._done_frac_for_policy_save * self.train_env.batch_size and
+        self._epoch % self._save_every_n == 0):
       self.save()
     policy_save_time = ppo.get_time(policy_save_start_time)
 
     epoch_time = ppo.get_time(epoch_start_time)
 
-    logging.info(
-        "PPO epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
-        " Loss(ppo, value, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)]", self._epoch,
-        min_reward, max_reward, avg_reward, combined_loss, ppo_loss, value_loss,
-        entropy_bonus)
-
     timing_dict = {
         "epoch": epoch_time,
         "policy_eval": policy_eval_time,
@@ -492,7 +497,7 @@ def train_epoch(self, evaluate=True):
     timing_dict.update(timing_info)
 
     for k, v in timing_dict.items():
-      self._timing_sw.scalar("timing/%s" % k, v, step=self._epoch)
+      self._timing_sw.scalar("timing/%s" % k, v, step=last_epoch)
 
     max_key_len = max(len(k) for k in timing_dict)
     timing_info_list = [
@@ -500,14 +505,12 @@ def train_epoch(self, evaluate=True):
         for k, v in sorted(timing_dict.items())
     ]
     logging.info(
-        "PPO epoch [% 6d], Timings: \n%s", self._epoch,
+        "PPO epoch [% 6d], Timings: \n%s", last_epoch,
         "\n".join(timing_info_list)
     )
 
-    self._epoch += 1
-
     # Flush summary writers once in a while.
-    if (self._epoch + 1) % 1000 == 0:
+    if self._epoch % 1000 == 0:
       self.flush_summaries()
 
   def evaluate(self):
@@ -540,7 +543,8 @@ def save(self):
            self._total_opt_step), f)
     # Remove the old model files.
     for path in old_model_files:
-      gfile.remove(path)
+      if path != params_file:
+        gfile.remove(path)
     # Reset this number.
     self._n_trajectories_done = 0
     self._last_saved_at = self._epoch
diff --git a/tensor2tensor/trax/rl/ppo_trainer_test.py b/tensor2tensor/trax/rl/ppo_trainer_test.py
index e3ca84fc7..a65853fa8 100644
--- a/tensor2tensor/trax/rl/ppo_trainer_test.py
+++ b/tensor2tensor/trax/rl/ppo_trainer_test.py
@@ -80,6 +80,7 @@ def _make_trainer(self, train_env, eval_env, output_dir, model=None):
         output_dir=output_dir,
         random_seed=0,
         boundary=2,
+        save_every_n=1,
     )
 
   def test_training_loop_cartpole(self):
@@ -224,6 +225,7 @@ def test_restarts(self):
       )
       self.assertEqual(trainer.epoch, 0)
       trainer.training_loop(n_epochs=1)
+      self.assertEqual(trainer.epoch, 1)
 
       # Restore from the saved state.
       trainer = self._make_trainer(
@@ -231,12 +233,10 @@ def test_restarts(self):
           eval_env=eval_env,
           output_dir=output_dir,
       )
-      # This is 2 instead of 1 because epoch calculation is a little weird right
-      # now.
-      # TODO(pkozakowski): Fix.
-      self.assertEqual(trainer.epoch, 2)
+      self.assertEqual(trainer.epoch, 1)
       # Check that we can continue training from the restored checkpoint.
-      trainer.training_loop(n_epochs=3)
+      trainer.training_loop(n_epochs=2)
+      self.assertEqual(trainer.epoch, 2)
 
 
 if __name__ == "__main__":

From 01fbe2b8076fae667b08c03750600d3687f6be41 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 9 Sep 2019 12:56:10 -0700
Subject: [PATCH 2401/2720] Attention with multiple rounds of hashing

PiperOrigin-RevId: 268056963
---
 .../trax/configs/transformer_imagenet64.gin   |  17 +-
 tensor2tensor/trax/layers/__init__.py         |   2 +
 tensor2tensor/trax/layers/attention.py        | 193 ++++++++++++++++++
 3 files changed, 208 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_imagenet64.gin b/tensor2tensor/trax/configs/transformer_imagenet64.gin
index 3f299776f..096de4739 100644
--- a/tensor2tensor/trax/configs/transformer_imagenet64.gin
+++ b/tensor2tensor/trax/configs/transformer_imagenet64.gin
@@ -6,7 +6,7 @@ import tensor2tensor.trax.trax
 # Parameters for batch_fun:
 # ==============================================================================
 batch_fun.batch_size_per_device = 2
-batch_fun.eval_batch_size = 128
+batch_fun.eval_batch_size = 16
 batch_fun.max_eval_length = 12288  # 64 * 64 * 3
 
 # Parameters for inputs:
@@ -23,8 +23,8 @@ MultifactorSchedule.warmup_steps = 8000
 
 # Parameters for train:
 # ==============================================================================
-train.eval_frequency = 200
-train.eval_steps = 8
+train.eval_frequency = 500
+train.eval_steps = 64
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.Adafactor
@@ -42,9 +42,18 @@ MergedHashedCausalAttention.dropout = 0.0
 MergedHashedCausalAttention.n_bins = 16
 MergedHashedCausalAttention.bin_by_time = True
 
+# Parameters for MergedMultiHashedCausalAttention:
+# ==============================================================================
+MergedMultiHashedCausalAttention.dropout = 0.0
+MergedMultiHashedCausalAttention.n_bins = 64
+MergedMultiHashedCausalAttention.n_hashes = 4
+MergedMultiHashedCausalAttention.bin_by_time = False
+
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.attention_type = @trax.layers.MergedHashedCausalAttention
+TransformerLM.attention_type = @trax.layers.MergedMultiHashedCausalAttention
+TransformerLM.d_attention_key = 64
+TransformerLM.d_attention_value = 256
 TransformerLM.d_model = 1024
 TransformerLM.d_ff = 2048
 TransformerLM.dropout = 0.0
diff --git a/tensor2tensor/trax/layers/__init__.py b/tensor2tensor/trax/layers/__init__.py
index 56258fede..427b18d5a 100644
--- a/tensor2tensor/trax/layers/__init__.py
+++ b/tensor2tensor/trax/layers/__init__.py
@@ -57,3 +57,5 @@ def layer_configure(*args, **kwargs):
     MemoryEfficientCausalAttention, blacklist=["mode"])
 MergedHashedCausalAttention = layer_configure(
     MergedHashedCausalAttention, blacklist=["mode"])
+MergedMultiHashedCausalAttention = layer_configure(
+    MergedMultiHashedCausalAttention, blacklist=["mode"])
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 32a6910d4..3df498207 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -29,6 +29,10 @@
 from tensor2tensor.trax.layers import initializers as init
 
 
+# Layers are always CamelCase, but functions in general are snake_case
+# pylint: disable=invalid-name
+
+
 @base.layer()
 def ShiftRight(x, **unused_kwargs):
   """Layer to shift the tensor to the right by padding on axis 1."""
@@ -707,6 +711,195 @@ def binned_attn_vjp(sqk, sv, so_ct):  # pylint: disable=invalid-name
       return out, (qk_ct, np.zeros_like(ignored_k), v_ct)
 
 
+class MergedMultiHashedCausalAttention(BaseCausalAttention):
+  """Hash-based causal attention, with multiple hashes."""
+  # TODO(kitaev): Adapt this layer for use in a reversible network. At the
+  # moment that isn't supported because there's no call_and_grad implementation.
+
+  def __init__(self, dropout, mode, n_bins=64, n_hashes=1, bin_by_time=False):
+    del dropout, mode
+    super(MergedMultiHashedCausalAttention, self).__init__()
+    self.n_bins = n_bins
+    self.n_hashes = n_hashes
+    self.bin_by_time = bin_by_time
+
+  def bin_vectors_by_time(self, vecs):
+    seqlen = vecs.shape[-2]
+    assert seqlen % self.n_bins == 0
+    bin_size = int(seqlen // self.n_bins)
+
+    bins = np.arange(seqlen, dtype=np.int32) // bin_size
+    bins = jax.lax.tie_in(vecs, bins)
+    bins = bins[None, :]
+    bins = np.broadcast_to(bins, vecs.shape[:-1])
+    return bins
+
+  def make_unit_length(self, x, epsilon=1e-6):
+    variance = np.mean(x**2, axis=-1, keepdims=True)
+    norm_inputs = x / np.sqrt(variance + epsilon)
+    return norm_inputs
+
+  def hash_vectors(self, vecs, rng):
+    if self.bin_by_time:
+      # Instead of hashing, put chunks of consecutive items in the same bin.
+      # This exists as a sanity check for the other parts of this class.
+      return self.bin_vectors_by_time(vecs)
+
+    # See https://arxiv.org/pdf/1509.02897.pdf
+    assert self.n_bins % 2 == 0
+    random_rotation = jax.random.normal(
+        rng, (vecs.shape[-1], self.n_bins//2)).astype('float32')
+
+    # TODO(kitaev): making the vectors unit-length here is probably redundant.
+    vecs = self.make_unit_length(vecs)
+    rotated_vecs = np.matmul(vecs, random_rotation)
+    rotated_vecs = self.make_unit_length(rotated_vecs)
+    rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)
+    bins = np.argmax(rotated_vecs, axis=-1)
+    return bins
+
+  def call(self, inputs, params=(), state=(), rng=None, **kwargs):
+    del params, kwargs
+    # We use the same vector as both a query and a key. For now we haven't
+    # adjusted any of the surrounding code, so we still get a separate "key"
+    # input that we ignore.
+    qk, _, v = inputs
+    seqlen = qk.shape[-2]
+
+    # qk/v are n_hashes*n_batch*n_heads, seqlen, d_head
+    # TODO(kitaev): is it faster to fuse this tiling into gather/scatter ops?
+    qk = np.tile(qk, (self.n_hashes, 1, 1))
+    v = np.tile(v, (self.n_hashes, 1, 1))
+
+    # bins are n_hashes*n_batch*n_heads, seqlen
+    # They specify which hash bucket the query/key/value vectors fall in.
+    bins = self.hash_vectors(qk, rng=rng)
+
+    # joint_t is n_hashes*n_batch*n_heads, seqlen
+    joint_t = jax.lax.tie_in(qk, np.arange(seqlen))
+    joint_t = np.reshape(joint_t, (1, seqlen))
+    joint_t = np.broadcast_to(joint_t, qk.shape[:-1])
+
+    assert int((self.n_bins + 1) * seqlen) < 2 ** 31, (
+        'Potential 32-bit integer overflow; please double-check the code.')
+    joint_bins_and_t = seqlen * bins + joint_t
+
+    def chunk_scalars(x):  # pylint: disable=invalid-name
+      return np.reshape(x, (x.shape[0], self.n_bins, -1))
+
+    def chunk_vectors(x):  # pylint: disable=invalid-name
+      return np.reshape(
+          x, (x.shape[0], self.n_bins, -1, x.shape[-1]))
+
+    def unchunk_vectors(x):  # pylint: disable=invalid-name
+      return np.reshape(x, (x.shape[0], -1, x.shape[-1]))
+
+    # Sort everything by bin number, with a secondary sort by time
+    # (variables starting with "s" are sorted)
+    _, sjoint_t = jax.lax.sort_key_val(
+        joint_bins_and_t, joint_t, dimension=-1)
+    _, undo_sort = jax.lax.sort_key_val(sjoint_t, joint_t, dimension=-1)
+    # TODO(kitaev): why does jax flag integer indices as differentiable?
+    # If we don't call stop_gradient here, custom gradients below won't work
+    # because the primitive functions close over "differentiable" variables.
+    sjoint_t = jax.lax.stop_gradient(sjoint_t)
+    undo_sort = jax.lax.stop_gradient(undo_sort)
+
+    # The backward pass of gather is in general a scatter operation, but we know
+    # we're dealing with permutations so we use gather for the backward pass
+    # too. This custom gradient should be about 2x faster than having jax infer
+    # one that uses scatter ops instead.
+    def permute_impl(vecs):
+      assert len(vecs.shape) == 3
+      return np.take_along_axis(vecs, sjoint_t[:, :, None], axis=-2)
+
+    def unpermute_impl(vecs):
+      assert len(vecs.shape) == 3
+      return np.take_along_axis(vecs, undo_sort[:, :, None], axis=-2)
+
+    @jax.custom_transforms
+    def permute(vecs):
+      return permute_impl(vecs)
+
+    def permute_vjp(vecs):
+      out_vecs = permute_impl(vecs)
+      def vjpfun(grad):
+        return (unpermute_impl(grad),)
+      return out_vecs, vjpfun
+
+    @jax.custom_transforms
+    def unpermute(vecs):
+      return unpermute_impl(vecs)
+
+    def unpermute_vjp(vecs):
+      out_vecs = unpermute_impl(vecs)
+      def vjpfun(grad):
+        return (permute_impl(grad),)
+      return out_vecs, vjpfun
+
+    jax.defvjp_all(permute, permute_vjp)
+    jax.defvjp_all(unpermute, unpermute_vjp)
+
+    sqk = permute(qk)
+    sv = permute(v)
+
+    # Split off a "bin" axis so that attention only occurs within chunks.
+    bq_t = bkv_t = chunk_scalars(sjoint_t)
+    bqk = chunk_vectors(sqk)
+    bv = chunk_vectors(sv)
+
+    # Hashing operates on unit-length vectors. Unnormalized query vectors are
+    # fine because they effectively provide a learnable temperature for the
+    # attention softmax, but normalizing keys is needed so that similarity for
+    # the purposes of attention correctly corresponds to hash locality.
+    bq = bqk
+    bk = self.make_unit_length(bqk)
+
+    # Allow each chunk to attend within itself, and also one chunk back. Chunk
+    # boundaries might occur in the middle of a sequence of items from the
+    # same bin, so this increases the chances of attending to relevant items.
+    # TODO(kitaev): benchmark whether XLA pad operation is noticeably faster.
+    bk_extra = np.concatenate([bk[:, -1:, :, :], bk[:, :-1, :, :]], axis=1)
+    bk = np.concatenate([bk, bk_extra], axis=2)
+    bv_extra = np.concatenate([bv[:, -1:, :, :], bv[:, :-1, :, :]], axis=1)
+    bv = np.concatenate([bv, bv_extra], axis=2)
+    bkv_t_extra = np.concatenate([bkv_t[:, -1:, :], bkv_t[:, :-1, :]], axis=1)
+    bkv_t = np.concatenate([bkv_t, bkv_t_extra], axis=2)
+
+    # Dot-product attention.
+    dots = np.matmul(bq, np.swapaxes(bk, -1, -2)) / np.sqrt(bq.shape[-1])
+
+    # Causal masking
+    mask = jax.lax.convert_element_type(
+        jax.lax.lt(bq_t[:, :, :, None], bkv_t[:, :, None, :]),
+        np.float32)
+    dots = dots - 1e9 * mask
+
+    # Mask out attention to self except when no other targets are available.
+    self_mask = jax.lax.broadcasted_eye(dots.dtype, dots.shape, (2, 3))
+    self_mask = jax.lax.tie_in(dots, self_mask)
+    dots = dots - 32 * self_mask
+
+    # Softmax.
+    dots_logsumexp = backend.logsumexp(dots, axis=-1, keepdims=True)
+    dots = np.exp(dots - dots_logsumexp)
+
+    bo = np.matmul(dots, bv)
+    so = unchunk_vectors(bo)
+    slogits = unchunk_vectors(dots_logsumexp)
+
+    o = unpermute(so)
+    logits = unpermute(slogits)
+
+    o = np.reshape(o, (self.n_hashes, -1, seqlen, o.shape[-1]))
+    logits = np.reshape(logits, (self.n_hashes, -1, seqlen, 1))
+    probs = np.exp(logits - backend.logsumexp(logits, axis=0, keepdims=True))
+    out = np.sum(o * probs, axis=0)
+    assert out.shape == inputs[2].shape
+
+    return out, state
+
+
 def CausalAttention(d_feature, n_heads=1,
                     d_attention_key=None, d_attention_value=None,
                     attention_type=DotProductCausalAttention, mode='train'):

From 58454ef3ad1b026380d8eb5eb3ccce6efacb9dd8 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 9 Sep 2019 14:32:35 -0700
Subject: [PATCH 2402/2720] Pass targets as inputs each time to model train
 step (so we can later make loss a layer).

PiperOrigin-RevId: 268078775
---
 ...former_wmt_ende_16gb_adafactor_testing.gin |  1 -
 .../transformer_wmt_ende_8gb_adafactor.gin    |  1 -
 .../configs/transformer_wmt_ende_8gb_adam.gin |  1 -
 .../configs/transformer_wmt_ende_8gb_sm3.gin  |  1 -
 tensor2tensor/trax/inputs.py                  | 46 ++++++++-------
 tensor2tensor/trax/rl/ppo_trainer_test.py     |  2 +
 tensor2tensor/trax/rl/simple_trainer.py       |  3 +
 tensor2tensor/trax/trax.py                    | 59 +++++++++++++++----
 tensor2tensor/trax/trax_test.py               |  4 +-
 9 files changed, 83 insertions(+), 35 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin
index 9577bb91b..5ee37f868 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin
@@ -14,7 +14,6 @@ batch_fun.buckets_include_inputs_in_length=True
 # ==============================================================================
 inputs.data_dir = None
 inputs.dataset_name = 't2t_translate_ende_wmt32k'
-inputs.append_targets = True
 
 # Parameters for mask:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
index e2b9ce3b6..d99183860 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
@@ -14,7 +14,6 @@ batch_fun.buckets_include_inputs_in_length=True
 # ==============================================================================
 inputs.data_dir = None
 inputs.dataset_name = 't2t_translate_ende_wmt32k'
-inputs.append_targets = True
 
 # Parameters for mask:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
index e167181e9..f9bda05e2 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
@@ -14,7 +14,6 @@ batch_fun.buckets_include_inputs_in_length=True
 # ==============================================================================
 inputs.data_dir = None
 inputs.dataset_name = 't2t_translate_ende_wmt32k'
-inputs.append_targets = True
 
 # Parameters for mask:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin
index a5ae925e9..8546162c5 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin
@@ -14,7 +14,6 @@ batch_fun.buckets_include_inputs_in_length=True
 # ==============================================================================
 inputs.data_dir = None
 inputs.dataset_name = 't2t_translate_ende_wmt32k'
-inputs.append_targets = True
 
 # Parameters for mask:
 # ==============================================================================
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index c2b12ac7b..634a8a07b 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -45,12 +45,14 @@
 # * input_shape: the shape of inputs
 #     the [...] above, without batch size
 # * input_dtype: the data type of inputs
-
+# * target_shape: the shape of targets
+#     the [...] above, without batch size
+# * target_dtype: the data type of targets
 
 Inputs = collections.namedtuple(
     '_Inputs',
     ['train_stream', 'train_eval_stream', 'eval_stream',
-     'input_shape', 'input_dtype']
+     'input_shape', 'input_dtype', 'target_shape', 'target_dtype']
 )
 
 # How many examples from the stream to skip at random during training.
@@ -61,7 +63,7 @@
 
 @gin.configurable(blacklist=['n_devices'])
 def inputs(n_devices, dataset_name, data_dir=None, input_name=None,
-           n_chunks=0, append_targets=False):
+           n_chunks=0):
   """Make Inputs for built-in datasets.
 
   Args:
@@ -71,8 +73,6 @@ def inputs(n_devices, dataset_name, data_dir=None, input_name=None,
     data_dir: data directory.
     input_name: optional, name of the inputs from the dictionary.
     n_chunks: optional, into how many pieces should we chunk (large inputs).
-    append_targets: optional, instead of inputs return a pair (inputs, targets)
-      which is useful for autoregressive models.
 
   Returns:
     trax.inputs.Inputs
@@ -81,35 +81,37 @@ def inputs(n_devices, dataset_name, data_dir=None, input_name=None,
   data_dir = os.path.expanduser(data_dir)
 
   (train_batches, train_eval_batches, eval_batches,
-   input_name, input_shape, input_dtype) = _train_and_eval_batches(
+   input_name, input_shape, input_dtype,
+   target_shape, target_dtype) = _train_and_eval_batches(
        dataset_name, data_dir, input_name, n_devices)
 
   if isinstance(input_dtype, tf.DType):
     input_dtype = input_dtype.as_numpy_dtype
+  if isinstance(target_dtype, tf.DType):
+    target_dtype = target_dtype.as_numpy_dtype
 
   if input_dtype == np.uint8:  # TPUs don't like uint8s, we cast to ints.
     input_dtype = np.int32
+  if target_dtype == np.uint8:
+    target_dtype = np.int32
 
   def numpy_stream(dataset):
-    return dataset_to_stream(
-        dataset, input_name,
-        n_chunks=n_chunks, append_targets=append_targets)
+    return dataset_to_stream(dataset, input_name, n_chunks=n_chunks)
 
   if n_chunks > 0:
     length = input_shape[0]
     input_shape = tuple(
         [tuple([length // n_chunks] + list(input_shape)[1:])] * n_chunks)
     input_dtype = tuple([input_dtype] * n_chunks)
-  if append_targets:
-    # TODO(lukaszkaiser): remove the assumption that input and target
-    # shapes are the same, which is used below for now.
-    input_shape = (input_shape, input_shape)
-    input_dtype = (input_dtype, input_dtype)
+    target_shape = tuple(
+        [tuple([length // n_chunks] + list(target_shape)[1:])] * n_chunks)
+    target_dtype = tuple([target_dtype] * n_chunks)
 
   return Inputs(train_stream=lambda: numpy_stream(train_batches),
                 train_eval_stream=lambda: numpy_stream(train_eval_batches),
                 eval_stream=lambda: numpy_stream(eval_batches),
-                input_shape=input_shape, input_dtype=input_dtype)
+                input_shape=input_shape, input_dtype=input_dtype,
+                target_shape=target_shape, target_dtype=target_dtype)
 
 
 @gin.configurable(blacklist=['n_devices'])
@@ -154,14 +156,17 @@ def random_minibatches():
       yield inp, out
 
   input_shape_without_batch = list(input_shape)[1:]
+  output_shape_without_batch = list(output_shape)[1:]
   return Inputs(train_stream=random_minibatches,
                 train_eval_stream=random_minibatches,
                 eval_stream=random_minibatches,
                 input_shape=input_shape_without_batch,
-                input_dtype=input_dtype)
+                input_dtype=input_dtype,
+                target_shape=output_shape_without_batch,
+                target_dtype=output_dtype)
 
 
-def dataset_to_stream(dataset, input_name, n_chunks=0, append_targets=False):
+def dataset_to_stream(dataset, input_name, n_chunks=0):
   """Takes a tf.Dataset and creates a numpy stream of ready batches."""
   for example in backend.dataset_as_numpy(dataset):
     inp, out = example[0][input_name], example[1]
@@ -177,8 +182,6 @@ def dataset_to_stream(dataset, input_name, n_chunks=0, append_targets=False):
       if n_chunks > 0:
         inp = tuple(np.split(inp, n_chunks, axis=1))
         out = tuple(np.split(out, n_chunks, axis=1))
-    if append_targets:
-      inp = (inp, out)
     yield inp, out
 
 
@@ -534,5 +537,8 @@ def _train_and_eval_batches(dataset, data_dir, input_name, n_devices):
   input_name = input_name or input_names[0]
   input_shape = features_info[input_name].shape
   input_dtype = features_info[input_name].dtype
+  target_shape = features_info[target_names[0]].shape
+  target_dtype = features_info[target_names[0]].dtype
   return (train_batches, train_eval_batches, eval_batches,
-          input_name, list(input_shape), input_dtype)
+          input_name, list(input_shape), input_dtype,
+          list(target_shape), target_dtype)
diff --git a/tensor2tensor/trax/rl/ppo_trainer_test.py b/tensor2tensor/trax/rl/ppo_trainer_test.py
index a65853fa8..4a13f781a 100644
--- a/tensor2tensor/trax/rl/ppo_trainer_test.py
+++ b/tensor2tensor/trax/rl/ppo_trainer_test.py
@@ -170,6 +170,8 @@ def inputs(n_devices):
           eval_stream=lambda: stream,
           input_shape=(history_shape[1:], action_shape[1:]),
           input_dtype=(np.float32, np.int32),
+          target_shape=(obs_shape[1:], reward_shape[1:]),
+          target_dtype=(np.float32, np.float32),
       )
 
     def loss(params, batch, model_predict, state, rng, **kwargs):
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index 539307c0c..c03645eef 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -143,6 +143,9 @@ def train_model(self):
         eval_stream=(lambda: eval_stream),
         input_shape=self._sim_env.model_input_shape,
         input_dtype=self._sim_env.model_input_dtype,
+        # TODO(lukaszkaiser): correct those, they may differ from inputs.
+        target_shape=self._sim_env.model_input_shape,
+        target_dtype=self._sim_env.model_input_dtype
     )
 
     self._model_train_step += self._n_model_train_steps
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 5f4ca6715..5668eeea9 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -124,10 +124,31 @@ def neg_log_perplexity(batch, model_predictions, has_weights):
   return masked_mean(xent, targets, weights)
 
 
+def _stack_inputs_targets_and_get_predictions(inputs_and_targets):
+  """Helper to stack inputs and targets and retrieve predictions from output."""
+  # Inputs and targets can be lists - we build a flat one to input to the model.
+  model_inp = []
+  for x in inputs_and_targets:
+    if not isinstance(x, (list, tuple)):
+      model_inp.append(x)
+    else:
+      model_inp.extend(x)
+  # We retrieve as many predictions from model output as many there were inputs.
+  inp = inputs_and_targets[0]
+  inp_len = len(inp) if isinstance(inp, (list, tuple)) else 1
+  get_pred = lambda x: x[0] if inp_len == 1 else x[:inp_len]
+  return tuple(model_inp), get_pred
+
+
 def loss(params, batch, model_predict, state, rng, has_weights):
   """Calculate loss."""
   inputs, targets, weights = unpack_batch(batch, has_weights)
-  predictions, state = model_predict(inputs, params, state, rng=rng)
+  model_input, get_preds = _stack_inputs_targets_and_get_predictions(
+      [inputs, targets])
+  # Call model, predictions will be the returned stack, usually consisting of
+  # the prediction tensor and the targets.
+  predictions, state = model_predict(model_input, params, state, rng=rng)
+  predictions = get_preds(predictions)
   predictions, targets, weights = _make_list(predictions, targets, weights)
   xent = []
   for (pred, target) in zip(predictions, targets):
@@ -287,9 +308,12 @@ def evaluate(inputs_stream, predict_fn, metric_fns, state, rng, has_weights):
   for inp in inputs_stream:
     count += 1
     rng, subrng = jax_random.split(rng)
-    preds, state = predict_fn(inp[0], state=state, rng=subrng)
+    model_inp, get_preds = _stack_inputs_targets_and_get_predictions(inp)
+    # Call model, preds will be the returned stack, usually (pred, targets).
+    preds, state = predict_fn(model_inp, state=state, rng=subrng)
+    pred = get_preds(preds)
     for m, f in six.iteritems(metric_fns):
-      metrics[m] += f(inp, preds, has_weights=has_weights)
+      metrics[m] += f(inp, pred, has_weights=has_weights)
   return {m: v / count for (m, v) in six.iteritems(metrics)}, state
 
 
@@ -563,24 +587,39 @@ def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs,
     if isinstance(first_shape, (list, tuple)):
       model_input_shape = tuple(
           tuple([None] + list(shape)) for shape in inputs.input_shape)
+      model_target_shape = tuple(
+          tuple([None] + list(shape)) for shape in inputs.target_shape)
     else:  # Otherwise just add [None] to the input shape.
       model_input_shape = tuple([None] + list(inputs.input_shape))
-    # Change all None to 1 in input shape.
+      model_target_shape = tuple([None] + list(inputs.target_shape))
+    # Change all None to 1 in input and target shape.
     model_input_shape = layers.nested_map(
         model_input_shape, lambda x: x if x else 1)
-    def initialize(input_shape, input_dtype, init_rng):
+    model_target_shape = layers.nested_map(
+        model_target_shape, lambda x: x if x else 1)
+    def initialize(input_shape, input_dtype, target_shape, target_dtype, rng):
+      """Helper to initialize the model."""
+      # Combine inputs and targets on the stack.
+      if not isinstance(input_dtype, (list, tuple)):
+        input_dtype = [input_dtype]
+        input_shape = [input_shape]
+      if not isinstance(target_dtype, (list, tuple)):
+        target_dtype = [target_dtype]
+        target_shape = [target_shape]
+      full_type = list(input_dtype) + list(target_dtype)
+      full_shape = list(input_shape) + list(target_shape)
       # We need to create a new model instance and not reuse `model_train` here,
       # because `m.initialize` puts cached parameter values in `m` and hence the
       # next call of `m.initialize` will give wrong results.
-      params, state = model(mode="train").initialize(input_shape, input_dtype,
-                                                     init_rng)
+      params, state = model(mode="train").initialize(full_shape, full_type, rng)
       (slots, opt_params) = opt.tree_init(params)
       return (OptState(params, slots, opt_params), state)
     if _is_jit_init():
       # JIT parameter initialization to avoid memory fragmentation
-      initialize = backend.jit(initialize, static_argnums=(0, 1))
+      initialize = backend.jit(initialize, static_argnums=(0, 1, 2, 3))
     self._initialize = lambda: initialize(  # pylint: disable=g-long-lambda
-        model_input_shape, self._inputs.input_dtype, init_rng)
+        model_input_shape, self._inputs.input_dtype,
+        model_target_shape, self._inputs.target_dtype, init_rng)
 
     # jit model_predict and update so they're fast
     self._jit_model_predict_eval = _jit_predict_fn(
@@ -779,7 +818,7 @@ def save_computation_graphs(self, save_backward_graph):
       next_train_batch = reshape_by_device(next_train_batch, self._n_devices)
     params = self._opt_state[0]
     forward_computation = jax.xla_computation(self._model_predict_eval)(
-        next_train_batch[0], params=params, state=self._model_state,
+        next_train_batch, params=params, state=self._model_state,
         rng=self._rngs[0])
     with gfile.GFile(os.path.join(output_dir, "forward.txt"), "w") as f:
       f.write(forward_computation.GetHloText())
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index 34fce4440..a6ada4ca8 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -67,7 +67,9 @@ def input_stream():
       train_eval_stream=input_stream,
       eval_stream=input_stream,
       input_shape=input_shape,
-      input_dtype=np.float32)
+      input_dtype=np.float32,
+      target_shape=(),
+      target_dtype=np.int32)
 
 
 BACKENDS = ["jax"]

From 270b9ed82b279911e05fe4c2be8b9c893996c4f9 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 9 Sep 2019 14:44:26 -0700
Subject: [PATCH 2403/2720] [async-rl CL 1/n] : Three functions to deal with
 ppo's policy model files.  - `get_policy_model_files` : returns all the
 policy model files, sorted on    decreasing epoch.  -
 `get_epoch_from_policy_model_file` : returns the epoch number.  -
 `get_policy_model_file_from_epoch` : given an epoch, returns the policy file 
   corresponding to it.

PiperOrigin-RevId: 268081382
---
 tensor2tensor/trax/rl/ppo.py         | 22 +++++++++++++---
 tensor2tensor/trax/rl/ppo_test.py    | 39 ++++++++++++++++++++++++++++
 tensor2tensor/trax/rl/ppo_trainer.py |  3 +--
 3 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index d42d9f4a4..d49398ab6 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -52,6 +52,7 @@
 import collections
 import functools
 import os
+import re
 import time
 
 from absl import logging
@@ -811,6 +812,21 @@ def compute_stats(reward_dict):
   }, state
 
 
+def get_policy_model_files(output_dir):
+  return list(
+      reversed(
+          sorted(gfile.glob(os.path.join(output_dir, "model-??????.pkl")))))
+
+
+def get_epoch_from_policy_model_file(policy_model_file):
+  base_name = os.path.basename(policy_model_file)
+  return int(re.match(r"model-(\d+).pkl", base_name).groups()[0])
+
+
+def get_policy_model_file_from_epoch(output_dir, epoch):
+  return os.path.join(output_dir, "model-%06d.pkl" % epoch)
+
+
 def maybe_restore_opt_state(output_dir, policy_and_value_opt_state,
                             policy_and_value_state):
   """Maybe restore the optimization state from the checkpoint dir.
@@ -831,15 +847,13 @@ def maybe_restore_opt_state(output_dir, policy_and_value_opt_state,
   """
   epoch = 0
   total_opt_step = 0
-  model_files = gfile.glob(os.path.join(output_dir, "model-??????.pkl"))
-  for model_file in reversed(sorted(model_files)):
+  for model_file in get_policy_model_files(output_dir):
     logging.info("Trying to restore model from %s", model_file)
     try:
       with gfile.GFile(model_file, "rb") as f:
         policy_and_value_opt_state, policy_and_value_state, total_opt_step = (
             pickle.load(f))
-      model_file_basename = os.path.basename(model_file)  # model-??????.pkl
-      epoch = int(filter(str.isdigit, model_file_basename))
+      epoch = get_epoch_from_policy_model_file(model_file)
       break
     except EOFError as e:
       logging.error("Unable to load model from: %s with %s", model_file, e)
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
index b0aa4e86b..2e53384cf 100644
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.rl import ppo
 from tensorflow import test
+from tensorflow.io import gfile
 
 
 class PpoTest(test.TestCase):
@@ -34,6 +35,44 @@ def setUp(self):
     super(PpoTest, self).setUp()
     self.rng_key = trax.get_random_number_generator_and_set_seed(0)
 
+  def test_get_policy_model_files(self):
+    output_dir = self.get_temp_dir()
+
+    def write_policy_model_file(epoch):
+      with gfile.GFile(
+          ppo.get_policy_model_file_from_epoch(output_dir, epoch), "w") as f:
+        f.write("some data")
+
+    epochs = [200, 100, 300]
+
+    # 300, 200, 100
+    expected_policy_model_files = [
+        output_dir + "/model-000300.pkl",
+        output_dir + "/model-000200.pkl",
+        output_dir + "/model-000100.pkl",
+    ]
+
+    for epoch in epochs:
+      write_policy_model_file(epoch)
+
+    policy_model_files = ppo.get_policy_model_files(output_dir)
+
+    self.assertEqual(expected_policy_model_files, policy_model_files)
+
+    gfile.rmtree(output_dir)
+
+  def test_get_epoch_from_policy_model_file(self):
+    self.assertEqual(0,
+                     ppo.get_epoch_from_policy_model_file("model-000000.pkl"))
+    self.assertEqual(123456,
+                     ppo.get_epoch_from_policy_model_file("model-123456.pkl"))
+
+  def test_get_policy_model_file_from_epoch(self):
+    self.assertEqual("/tmp/model-000000.pkl",
+                     ppo.get_policy_model_file_from_epoch("/tmp", 0))
+    self.assertEqual("/tmp/model-123456.pkl",
+                     ppo.get_policy_model_file_from_epoch("/tmp", 123456))
+
   def test_policy_and_value_net(self):
     observation_shape = (3, 4, 5)
     batch_observation_shape = (1, 1) + observation_shape
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index 2693c3b40..bfa8bdec1 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -534,8 +534,7 @@ def evaluate(self):
   def save(self):
     """Save the agent parameters."""
     logging.vlog(1, "PPO epoch [% 6d]: saving model.", self._epoch)
-    old_model_files = gfile.glob(
-        os.path.join(self._output_dir, "model-??????.pkl"))
+    old_model_files = ppo.get_policy_model_files(self._output_dir)
     params_file = os.path.join(self._output_dir, "model-%06d.pkl" % self._epoch)
     with gfile.GFile(params_file, "wb") as f:
       pickle.dump(

From 31544451f5f74e3e277e8d4d41d35873206d4a53 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 9 Sep 2019 15:07:34 -0700
Subject: [PATCH 2404/2720] [async-rl CL 2/n] : Unify make_env in rl_trainer
 and env_problem_utils.

PiperOrigin-RevId: 268086511
---
 tensor2tensor/envs/env_problem_utils.py       |   5 +-
 tensor2tensor/envs/env_service_server.py      |   3 +-
 .../trax/rl/envs/env_service_server.py        |   3 +-
 tensor2tensor/trax/rl_trainer.py              | 100 ++++++------------
 4 files changed, 36 insertions(+), 75 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index bfe232f7b..32d81b436 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -200,8 +200,7 @@ def gumbel_sample(log_probs):
 def make_env(batch_size=1,
              env_problem_name="",
              resize=True,
-             resized_height=105,
-             resized_width=80,
+             resize_dims=(105, 80),
              max_timestep="None",
              clip_rewards=True,
              parallelism=1,
@@ -232,7 +231,7 @@ def make_env(batch_size=1,
           "rl_env_max_episode_steps": max_timestep,
           "maxskip_env": True,
           "rendered_env": True,
-          "rendered_env_resize_to": (resized_height, resized_width),
+          "rendered_env_resize_to": resize_dims,
           "sticky_actions": False,
           "output_dtype": np.int32 if use_tpu else None,
       })
diff --git a/tensor2tensor/envs/env_service_server.py b/tensor2tensor/envs/env_service_server.py
index 24b268ce9..83f61ebf5 100644
--- a/tensor2tensor/envs/env_service_server.py
+++ b/tensor2tensor/envs/env_service_server.py
@@ -60,8 +60,7 @@ def main(argv):
       batch_size=1,
       env_problem_name=FLAGS.env_problem_name,
       resize=FLAGS.resize,
-      resized_height=FLAGS.resized_height,
-      resized_width=FLAGS.resized_width,
+      resize_dims=(FLAGS.resized_height, FLAGS.resized_width),
       max_timestep=FLAGS.max_timestep,
       clip_rewards=FLAGS.clip_rewards)
 
diff --git a/tensor2tensor/trax/rl/envs/env_service_server.py b/tensor2tensor/trax/rl/envs/env_service_server.py
index 15a5bbd4b..b81cfd878 100644
--- a/tensor2tensor/trax/rl/envs/env_service_server.py
+++ b/tensor2tensor/trax/rl/envs/env_service_server.py
@@ -100,8 +100,7 @@ def main(argv):
       batch_size=1,
       env_problem_name=FLAGS.env_problem_name,
       resize=FLAGS.resize,
-      resized_height=FLAGS.resized_height,
-      resized_width=FLAGS.resized_width,
+      resize_dims=(FLAGS.resized_height, FLAGS.resized_width),
       max_timestep=FLAGS.max_timestep,
       clip_rewards=FLAGS.clip_rewards,
       **env_kwargs)
diff --git a/tensor2tensor/trax/rl_trainer.py b/tensor2tensor/trax/rl_trainer.py
index 21651607c..04de7d56a 100644
--- a/tensor2tensor/trax/rl_trainer.py
+++ b/tensor2tensor/trax/rl_trainer.py
@@ -32,7 +32,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import functools
 import multiprocessing
 import os
 
@@ -42,11 +41,8 @@
 import gin
 import jax
 from jax.config import config
-import numpy as onp
 from tensor2tensor import envs  # pylint: disable=unused-import
-from tensor2tensor.envs import gym_env_problem
-from tensor2tensor.envs import rendered_env_problem
-from tensor2tensor.rl import gym_utils
+from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.rl.google import atari_utils  # GOOGLE-INTERNAL:
 from tensor2tensor.trax import rl  # pylint: disable=unused-import
 from tensor2tensor.trax.rl import envs as rl_envs  # pylint: disable=unused-import
@@ -72,56 +68,18 @@
 flags.DEFINE_integer("eval_batch_size", 4, "Batch size for evaluation.")
 flags.DEFINE_boolean("parallelize_envs", False,
                      "If true, sets parallelism to number of cpu cores.")
-flags.DEFINE_string(
-    "trajectory_dump_dir", "", "Directory to dump trajectories to.")
-
+flags.DEFINE_string("trajectory_dump_dir", "",
+                    "Directory to dump trajectories to.")
 
 # TODO(afrozm): Find a better way to do these configurations.
 flags.DEFINE_string("train_server_bns", "", "Train Server's BNS.")
 flags.DEFINE_string("eval_server_bns", "", "Eval Server's BNS.")
 
 
-def make_env(name, batch_size, max_timestep, clip_rewards, rendered_env,
-             resize_dims, **env_kwargs):
-  """Creates the env."""
-
-  if clip_rewards:
-    env_kwargs.update({"reward_range": (-1, 1), "discrete_rewards": True})
-  else:
-    env_kwargs.update({"discrete_rewards": False})
-
-  # TODO(afrozm): Should we leave out some cores?
-  parallelism = multiprocessing.cpu_count() if FLAGS.parallelize_envs else 1
-
-  # No resizing needed, so let's be on the normal EnvProblem.
-  if not rendered_env:
-    return gym_env_problem.GymEnvProblem(
-        base_env_name=name,
-        batch_size=batch_size,
-        parallelism=parallelism,
-        **env_kwargs)
-
-  wrapper_fn = functools.partial(
-      gym_utils.gym_env_wrapper, **{
-          "rl_env_max_episode_steps": max_timestep,
-          "maxskip_env": True,
-          "rendered_env": True,
-          "rendered_env_resize_to": resize_dims,
-          "sticky_actions": False,
-          "output_dtype": onp.int32 if FLAGS.use_tpu else None,
-      })
-
-  return rendered_env_problem.RenderedEnvProblem(
-      base_env_name=name,
-      batch_size=batch_size,
-      parallelism=parallelism,
-      env_wrapper_fn=wrapper_fn,
-      **env_kwargs)
-
-
 # Not just "train" to avoid a conflict with trax.train in GIN files.
 @gin.configurable(blacklist=[
-    "output_dir", "train_batch_size", "eval_batch_size", "trajectory_dump_dir"])
+    "output_dir", "train_batch_size", "eval_batch_size", "trajectory_dump_dir"
+])
 def train_rl(
     output_dir,
     train_batch_size,
@@ -145,8 +103,8 @@ def train_rl(
     max_timestep: Int or None, the maximum number of timesteps in a trajectory.
       The environment is wrapped in a TimeLimit wrapper.
     clip_rewards: Whether to clip and discretize the rewards.
-    rendered_env: Whether the environment has visual input. If so,
-      a RenderedEnvProblem will be used.
+    rendered_env: Whether the environment has visual input. If so, a
+      RenderedEnvProblem will be used.
     resize_dims: Pair (height, width), dimensions to resize the visual
       observations to.
     trainer_class: RLTrainer class to use.
@@ -168,12 +126,8 @@ def train_rl(
   eval_env_kwargs = {}
   if "OnlineTuneEnv" in env_name:
     # TODO(pkozakowski): Separate env output dirs by train/eval and epoch.
-    train_env_kwargs = {
-        "output_dir": os.path.join(output_dir, "envs/train")
-    }
-    eval_env_kwargs = {
-        "output_dir": os.path.join(output_dir, "envs/eval")
-    }
+    train_env_kwargs = {"output_dir": os.path.join(output_dir, "envs/train")}
+    eval_env_kwargs = {"output_dir": os.path.join(output_dir, "envs/eval")}
 
   if "ClientEnv" in env_name:
     train_env_kwargs["per_env_kwargs"] = [{
@@ -184,21 +138,31 @@ def train_rl(
         "remote_env_address": os.path.join(FLAGS.eval_server_bns, str(replica))
     } for replica in range(eval_batch_size)]
 
-  common_env_kwargs = {
-      "name": env_name,
-      "max_timestep": max_timestep,
-      "clip_rewards": clip_rewards,
-      "rendered_env": rendered_env,
-      "resize_dims": resize_dims,
-  }
-  train_env_kwargs.update(common_env_kwargs)
-  eval_env_kwargs.update(common_env_kwargs)
-
-  # Make an env here.
-  train_env = make_env(batch_size=train_batch_size, **train_env_kwargs)
+  # TODO(afrozm): Should we leave out some cores?
+  parallelism = multiprocessing.cpu_count() if FLAGS.parallelize_envs else 1
+
+  train_env = env_problem_utils.make_env(
+      batch_size=train_batch_size,
+      env_problem_name=env_name,
+      resize=rendered_env,
+      resize_dims=resize_dims,
+      max_timestep=max_timestep,
+      clip_rewards=clip_rewards,
+      parallelism=parallelism,
+      use_tpu=FLAGS.use_tpu,
+      **train_env_kwargs)
   assert train_env
 
-  eval_env = make_env(batch_size=eval_batch_size, **eval_env_kwargs)
+  eval_env = env_problem_utils.make_env(
+      batch_size=eval_batch_size,
+      env_problem_name=env_name,
+      resize=rendered_env,
+      resize_dims=resize_dims,
+      max_timestep=max_timestep,
+      clip_rewards=clip_rewards,
+      parallelism=parallelism,
+      use_tpu=FLAGS.use_tpu,
+      **eval_env_kwargs)
   assert eval_env
 
   def run_training_loop():

From f077f56eccf1a63a28cc03b919182cb9abb26930 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 9 Sep 2019 15:32:27 -0700
Subject: [PATCH 2405/2720] Rewrite parts of _jit_..._fn for further clarity.
 (Pure refactor.)

PiperOrigin-RevId: 268091490
---
 tensor2tensor/trax/trax.py | 31 ++++++++++++-------------------
 1 file changed, 12 insertions(+), 19 deletions(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 5668eeea9..e91b42f2e 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -417,13 +417,10 @@ def epochs(total_steps, steps_to_skip, epoch_steps):
 
 @gin.configurable
 def _jit_predict_fn(model_predict, n_devices, jit=True):
-  """Use jit on model_predict if required."""
+  """Returns a JIT-compiled predict function (unless jit=False)."""
 
   if n_devices == 1:
-    if jit:
-      return backend.jit(model_predict)
-    else:
-      return model_predict
+    return backend.jit(model_predict) if jit else model_predict
 
   # Multi-devices, pmap and run.
   @functools.partial(backend.pmap, axis_name="batch")
@@ -449,28 +446,26 @@ def combine(x):
 
 @gin.configurable
 def _jit_update_fn(predict_fn, loss_fn, optimizer, n_devices, jit=True):
-  """Get jit-ed update function for loss, optimizer, learning rate function."""
+  """Returns a (JIT-compiled) function that computes updates for one step."""
   if n_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
     def single_update(i, opt_state, batch, state, rng):
       params, slots, opt_params = opt_state
       rng, subrng = jax_random.split(rng[0])
-      grads, state = backend.grad(loss_fn, has_aux=True)(params, batch,
-                                                         predict_fn, state, rng)
+      grad_fn = backend.grad(loss_fn, has_aux=True)
+      grads, state = grad_fn(params, batch, predict_fn, state, rng)
       return optimizer.tree_update(
           i, grads, params, slots, opt_params), state, [subrng]
-    if jit:
-      return backend.jit(single_update)
-    else:
-      return single_update
+    return backend.jit(single_update) if jit else single_update
 
+  # Else, for n_devices > 1:
   @functools.partial(backend.pmap, axis_name="batch")
   def mapped_update(i, opt_state, batch, state, rng):
     """This is a multi-device version of the update function above."""
     # We assume all tensors have the first dimension = n_devices.
     params, slots, opt_params = opt_state
     rng, subrng = jax_random.split(rng)
-    grads, state = backend.grad(loss_fn, has_aux=True)(params, batch,
-                                                       predict_fn, state, rng)
+    grad_fn = backend.grad(loss_fn, has_aux=True)
+    grads, state = grad_fn(params, batch, predict_fn, state, rng)
     grads = jax.tree_util.tree_map(
         lambda g: lax.psum(g, "batch"), grads)
     return optimizer.tree_update(
@@ -485,17 +480,15 @@ def update(i, opt_state, batch, state, rng):
 
 @gin.configurable
 def _jit_compute_loss_fn(predict_fn, loss_fn, n_devices, jit=True):
-  """Get jit-ed function that computes the loss."""
+  """Returns a (JIT-compiled) function that computes the loss for one step."""
   if n_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
     def single_compute_loss(opt_state, batch, state, rng):
       rng, subrng = jax_random.split(rng[0])
       loss_val, state = loss_fn(opt_state[0], batch, predict_fn, state, rng)
       return loss_val, state, [subrng]
-    if jit:
-      return backend.jit(single_compute_loss)
-    else:
-      return single_compute_loss
+    return backend.jit(single_compute_loss) if jit else single_compute_loss
 
+  # Else, for n_devices > 1:
   @functools.partial(backend.pmap, axis_name="batch")
   def mapped_compute_loss(opt_state, batch, state, rng):
     """This is a multi-device version of the update function above."""

From dfcf88cb9b2ac695b1ca6be46b4ec29190d093b7 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 9 Sep 2019 17:36:49 -0700
Subject: [PATCH 2406/2720] Add a setting to share keys and values in causal
 attention.

PiperOrigin-RevId: 268116208
---
 tensor2tensor/trax/layers/attention.py        | 52 +++++++++++++++----
 .../models/research/transformer_revnet.py     | 40 +++++++++-----
 tensor2tensor/trax/models/transformer.py      | 13 +++--
 3 files changed, 79 insertions(+), 26 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 3df498207..b8e3a990e 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -235,6 +235,25 @@ def BasicCausalAttention(d_feature, n_heads=1, dropout=0.0, mode='train'):
   ]
 
 
+class ShiftRightLearned(base.Layer):
+  """Layer constructor function for shifting right by a learned vector."""
+
+  def __init__(self, initializer=init.RandomNormalInitializer(0.01)):
+    super(ShiftRightLearned, self).__init__()
+    self._initializer = initializer
+
+  def call(self, x, params, state, **kwargs):
+    del kwargs
+    c = backend.numpy.reshape(params, [1, 1, -1])
+    c += backend.numpy.zeros((x.shape[0], 1, x.shape[2]), dtype=x.dtype)
+    return backend.numpy.concatenate([c, x], axis=1)[:, :-1, :], state
+
+  def new_parameters(self, input_shape, input_dtype, rng):
+    del input_dtype
+    b = self._initializer((input_shape[-1],), rng)
+    return b, ()
+
+
 class ComputeAttentionHeads(base.Layer):
   """Computes queries/keys/values via linear projection.
 
@@ -902,7 +921,8 @@ def vjpfun(grad):
 
 def CausalAttention(d_feature, n_heads=1,
                     d_attention_key=None, d_attention_value=None,
-                    attention_type=DotProductCausalAttention, mode='train'):
+                    attention_type=DotProductCausalAttention,
+                    share_kv=False, mode='train'):
   """Transformer-style multi-headed causal attention.
 
   Args:
@@ -912,7 +932,8 @@ def CausalAttention(d_feature, n_heads=1,
         (default is d_feature // n_heads)
     d_attention_value: int: depth of value vector for each attention head
         (default is d_feature // n_heads)
-    attention_type: subclass of tl.BaseCausalAttention: attention class to use
+    attention_type: subclass of BaseCausalAttention: attention class to use
+    share_kv: bool, whether to share keys and values
     mode: str: 'train' or 'eval'
 
   Returns:
@@ -925,13 +946,26 @@ def CausalAttention(d_feature, n_heads=1,
     assert d_feature % n_heads == 0
     d_attention_value = d_feature // n_heads
 
-  return [
-      cb.Dup(), cb.Dup(),
-      cb.Parallel(
-          ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
-          ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
-          ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
-      ),
+  if share_kv:
+    pre_attention = [
+        cb.Dup(),
+        cb.Parallel(
+            ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
+            ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
+        ),
+        cb.Dup(),
+    ]
+  else:
+    pre_attention = [
+        cb.Dup(), cb.Dup(),
+        cb.Parallel(
+            ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
+            ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
+            ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
+        ),
+    ]
+
+  return pre_attention + [
       attention_type(mode=mode),
       ComputeAttentionOutput(n_heads=n_heads, d_model=d_feature),
   ]
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index c64052873..71cd89bbc 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -394,7 +394,7 @@ def call_post_attention2(params):
 
 def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
                  n_heads, n_attention_chunks, attention_type,
-                 dropout, mode):
+                 dropout, share_kv, mode):
   """Reversible transformer decoder layer.
 
   Args:
@@ -406,22 +406,34 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
     n_attention_chunks: int: number of chunks for attention
     attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
+    share_kv: string, whether to share keys and values
     mode: str: 'train' or 'eval'
 
   Returns:
     the layer.
   """
-
-  pre_attention = [
-      Chunk(n_sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
-      tl.LayerNorm(),
-      tl.Dup(), tl.Dup(),
-      tl.Parallel(
-          tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
-          tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
-          tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
-      ),
-  ]
+  if share_kv:
+    pre_attention = [
+        Chunk(n_sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
+        tl.LayerNorm(),
+        tl.Dup(),
+        tl.Parallel(
+            tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
+            tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
+        ),
+        tl.Dup(),
+    ]
+  else:
+    pre_attention = [
+        Chunk(n_sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
+        tl.LayerNorm(),
+        tl.Dup(), tl.Dup(),
+        tl.Parallel(
+            tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
+            tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
+            tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
+        ),
+    ]
 
   attention = attention_type(mode=mode)
 
@@ -456,6 +468,7 @@ def TransformerRevnetLM(vocab_size,
                         n_chunks=32,
                         n_attention_chunks=8,
                         attention_type=tl.DotProductCausalAttention,
+                        share_kv=False,
                         mode='train'):
   """Reversible transformer language model (only uses a decoder, no encoder).
 
@@ -472,6 +485,7 @@ def TransformerRevnetLM(vocab_size,
     n_chunks: int: number of chunks (must match input pipeline)
     n_attention_chunks: int: number of chunks for attention
     attention_type: class: attention class to use, such as DotProductAttention.
+    share_kv: bool, whether to share keys and values.
     mode: str: 'train' or 'eval'
 
   Returns:
@@ -492,7 +506,7 @@ def TransformerRevnetLM(vocab_size,
           DecoderBlock(d_model, d_ff,
                        d_attention_key, d_attention_value, n_heads,
                        n_attention_chunks, attention_type,
-                       dropout, mode)
+                       dropout, share_kv, mode)
           for _ in range(n_layers)
       ] + [
           SplitForOutput(n_sections=n_chunks, axis=-2),  # pylint: disable=no-value-for-parameter
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index edc60ade2..f180c2a57 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -112,7 +112,7 @@ def TransformerEncoder(vocab_size,
 
 
 def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-                 attention_type, dropout, mode):
+                 attention_type, dropout, share_kv, mode):
   """Returns a layer sequence that implements a Transformer decoder block.
 
   The input to the layer sequence is an activation tensor.
@@ -125,6 +125,7 @@ def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
     d_attention_value: int: depth of value vector for each attention head
     attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
+    share_kv: bool, whether to share keys and values
     mode: str: 'train' or 'eval'
 
   Returns:
@@ -135,7 +136,7 @@ def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
       tl.CausalAttention(
           d_model, n_heads=n_heads, d_attention_key=d_attention_key,
           d_attention_value=d_attention_value, attention_type=attention_type,
-          mode=mode),
+          share_kv=share_kv, mode=mode),
       tl.Dropout(rate=dropout, mode=mode),  # vec
   ]
   feed_forward = [
@@ -155,6 +156,7 @@ def TransformerDecoder(d_model=512,
                        d_attention_value=None,
                        attention_type=tl.DotProductCausalAttention,
                        dropout=0.1,
+                       share_kv=False,
                        max_len=2048,
                        mode='train'):
   """Returns a Transformer decoder model.
@@ -174,6 +176,7 @@ def TransformerDecoder(d_model=512,
         (default is d_model // n_heads)
     attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
+    share_kv: bool, whether to share keys and values in decoder attention
     max_len: int: maximum symbol length for positional encoding
     mode: str: 'train' or 'eval'
 
@@ -187,7 +190,7 @@ def TransformerDecoder(d_model=512,
       tl.PositionalEncoding(max_len=max_len),
       [DecoderBlock(  # pylint: disable=g-complex-comprehension
           d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-          attention_type, dropout, mode)
+          attention_type, dropout, share_kv, mode)
        for _ in range(n_layers)],   # vecs
       tl.LayerNorm(),               # vecs
   )
@@ -202,6 +205,7 @@ def TransformerLM(vocab_size,
                   d_attention_value=None,
                   attention_type=tl.DotProductCausalAttention,
                   dropout=0.1,
+                  share_kv=False,
                   max_len=2048,
                   mode='train'):
   """Returns a Transformer language model.
@@ -221,6 +225,7 @@ def TransformerLM(vocab_size,
         (default is d_model // n_heads)
     attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
+    share_kv: bool, whether to share keys and values in decoder attention
     max_len: int: maximum symbol length for positional encoding
     mode: str: 'train' or 'eval'
 
@@ -238,7 +243,7 @@ def TransformerLM(vocab_size,
       embedder,                     # vecs
       [DecoderBlock(  # pylint: disable=g-complex-comprehension
           d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-          attention_type, dropout, mode)
+          attention_type, dropout, share_kv, mode)
        for _ in range(n_layers)],   # vecs
       tl.LayerNorm(),               # vecs
       tl.Dense(vocab_size),         # vecs

From 6561075293da462447c73c98a0e045d17c6d32b1 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 9 Sep 2019 17:44:07 -0700
Subject: [PATCH 2407/2720] [async-rl CL 3/n] : Add method to `BatchTrajectory`
 to initialize from a directory.

PiperOrigin-RevId: 268117379
---
 tensor2tensor/envs/trajectory.py      | 90 +++++++++++++++++++++++++--
 tensor2tensor/envs/trajectory_test.py | 10 +++
 2 files changed, 94 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index 73169e84b..d7622e4bf 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -23,8 +23,15 @@
 from __future__ import division
 from __future__ import print_function
 
+import os
+import re
+import cloudpickle as pickle
 import numpy as np
 from tensor2tensor.envs import time_step
+from tensorflow.io import gfile
+
+TRAJECTORY_FILE_REGEXP = r"trajectory_epoch_(\d{6})_env_id_(\d{6})_r_(.*).pkl"
+TRAJECTORY_FILE_GLOB = r"trajectory_epoch_*_env_id_*_r_*.pkl"
 
 
 class Trajectory(object):
@@ -156,15 +163,20 @@ def as_numpy(self):
 class BatchTrajectory(object):
   """Basically a batch of active trajectories and a list of completed ones."""
 
-  def __init__(self, batch_size=1):
+  def __init__(self,
+               batch_size=1,
+               trajectories=None,
+               completed_trajectories=None):
     self.batch_size = batch_size
 
     # Stores trajectories that are currently active, i.e. aren't done or reset.
-    self._trajectories = [Trajectory() for _ in range(self.batch_size)]
+    self._trajectories = trajectories or [
+        Trajectory() for _ in range(self.batch_size)
+    ]
 
     # Stores trajectories that are completed.
     # NOTE: We don't track the index this came from, as it's not needed, right?
-    self._completed_trajectories = []
+    self._completed_trajectories = completed_trajectories or []
 
   def reset_batch_trajectories(self):
     self.__init__(batch_size=self.batch_size)
@@ -287,7 +299,12 @@ def complete_all_trajectories(self):
       if trajectory.is_active:
         self._complete_trajectory(trajectory, index)
 
-  def step(self, observations, raw_rewards, processed_rewards, dones, actions,
+  def step(self,
+           observations,
+           raw_rewards,
+           processed_rewards,
+           dones,
+           actions,
            infos=None):
     """Record the information obtained from taking a step in all envs.
 
@@ -350,8 +367,7 @@ def extract_info_at_index(infos, index):
 
       # To this trajectory's last time-step, set actions.
       trajectory.change_last_time_step(
-          action=actions[index],
-          info=extract_info_at_index(infos, index))
+          action=actions[index], info=extract_info_at_index(infos, index))
 
       # Create a new time-step to add observation, done & rewards (no actions).
       trajectory.add_time_step(
@@ -434,3 +450,65 @@ def padding_config(obs):
         np.pad(obs, padding_config(obs), "constant")
         for obs in list_observations_np_ts
     ]), trajectory_lengths
+
+  @staticmethod
+  def parse_trajectory_file_name(trajectory_file_name):
+    """Parse out the trajectory file's groups and return to caller."""
+    base_trajectory_file_name = os.path.basename(trajectory_file_name)
+    compiled_regexp = re.compile(TRAJECTORY_FILE_REGEXP)
+    r = compiled_regexp.match(base_trajectory_file_name)
+    if not r:
+      return None
+    g = r.groups()
+    if len(g) is not compiled_regexp.groups:
+      return None
+    return (int(g[0]), int(g[1])) + g[2:]
+
+  # TODO(afrozm): Test this function.
+  @staticmethod
+  def load_from_directory(trajectory_dir, epoch=None, n_trajectories=None):
+    """Load trajectories from specified dir and epoch."""
+
+    trajectory_file_glob = TRAJECTORY_FILE_GLOB
+
+    # If there is a desired epoch, modify the glob to get that instead.
+    if epoch:
+      trajectory_file_glob = trajectory_file_glob.replace(
+          "epoch_*", "epoch_%06d" % epoch)
+
+    trajectory_files = gfile.glob(
+        os.path.join(trajectory_dir, trajectory_file_glob))
+
+    if not trajectory_files:
+      return None
+
+    # We read and load all the files, revisit if this becomes a problem.
+    trajectories_buffer = []
+    completed_trajectories_buffer = []
+    for trajectory_file in trajectory_files:
+      with gfile.GFile(trajectory_file, "rb") as f:
+        list_trajectories = pickle.load(f)
+        assert isinstance(list_trajectories, list)
+        if not list_trajectories:
+          continue
+        assert isinstance(list_trajectories[0], Trajectory)
+        for trajectory in list_trajectories:
+          if trajectory.done:
+            completed_trajectories_buffer.append(trajectory)
+          else:
+            trajectories_buffer.append(trajectory)
+
+    if not trajectories_buffer and not completed_trajectories_buffer:
+      return None
+
+    # Randomly sample `n_trajectories` if needed.
+    n_trajectories = None if not n_trajectories else int(n_trajectories)
+    if n_trajectories and n_trajectories > 0:
+      trajectories_buffer = list(
+          np.random.choice(trajectories_buffer, int(trajectories_buffer)))
+
+    # Construct and return a new BatchTrajectory object.
+    return BatchTrajectory(
+        batch_size=len(trajectories_buffer),
+        trajectories=trajectories_buffer,
+        completed_trajectories=completed_trajectories_buffer)
diff --git a/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index 8658aee21..129bab8b3 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -467,5 +467,15 @@ def test_observations_np(self):
       for ts in range(lengths[b], len(padded_obs_np[b])):
         self.assertAllEqual(zero_obs, padded_obs_np[b][ts])
 
+  def test_parse_trajectory_file_name(self):
+    self.assertEqual(
+        (12, 13, "abc"),
+        trajectory.BatchTrajectory.parse_trajectory_file_name(
+            "/tmp/trajectory_epoch_000012_env_id_000013_r_abc.pkl"))
+
+    self.assertIsNone(
+        trajectory.BatchTrajectory.parse_trajectory_file_name(
+            "/tmp/trajectory_epoch_000012_env_id_000013.pkl"))
+
 if __name__ == "__main__":
   tf.test.main()

From d828c87916109f2dc6f21dbbf0eab7a77aacfc3f Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Mon, 9 Sep 2019 18:43:44 -0700
Subject: [PATCH 2408/2720] Pass the current learning rate to LR schedules

This will be needed for policy-controlled schedules, as actions are relative changes in the learning rate.

PiperOrigin-RevId: 268125833
---
 tensor2tensor/trax/trax.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index e91b42f2e..8d946a90e 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -758,11 +758,7 @@ def train_epoch(self, epoch_steps, eval_steps):
 
       # LR log
       if self._step == 1 or self._step % 10 == 0:
-        # TODO(lukaszkaiser): it makes no sense to use an accelerator (e.g. TPU)
-        # in op-by-op mode just to compute the learning rate. However, there
-        # should be a cleaner approach that forceably swapping out the backend.
-        with backend.use_backend("numpy"):
-          self._train_sw.scalar("training/learning rate", self.learning_rate)
+        self._train_sw.scalar("training/learning_rate", self.learning_rate)
 
     # Timer
     epoch_time = time.time() - start_time
@@ -783,6 +779,7 @@ def train_epoch(self, epoch_steps, eval_steps):
     self._eval_sw.flush()
 
   def evaluate(self, eval_steps):
+    """Evaluate the model and log metrics."""
     _, rng = jax_random.split(self._rngs[0])
     _, _, self._model_state = evaluate_train_and_eval(
         step=self._step,
@@ -798,6 +795,10 @@ def evaluate(self, eval_steps):
         history=self._history,
         has_weights=self._has_weights)
 
+    # Save the learning rate in the history
+    self._history.append("train", "training/learning_rate", self._step,
+                         self.learning_rate)
+
   def update_learning_rate(self):
     self._lr_fn = self._lr_schedule(self._history)
 

From 840248c5b58b2b91de42239bacddc2647944c4e8 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 10 Sep 2019 07:48:37 -0700
Subject: [PATCH 2409/2720] fix documentation regarding padding/dropping long
 and short sequences.

PiperOrigin-RevId: 268222948
---
 tensor2tensor/models/transformer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 37140c9d8..75fd1826f 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -2495,7 +2495,7 @@ def update_hparams_for_tpu(hparams):
   # to a longer length, e.g. the "_packed" problems.
   #
   # For problems with variable sequence lengths, this parameter controls the
-  # maximum sequence length.  Shorter sequences are dropped and longer ones
+  # maximum sequence length. Longer sequences are dropped and shorter ones
   # are padded.
   #
   # For problems with fixed sequence lengths - e.g. the "_packed" problems,

From 3eba3dff29087c0531274ec1bca8e8453ac3a13d Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 10 Sep 2019 10:57:46 -0700
Subject: [PATCH 2410/2720] Move network initialization out of
 policy_and_value_net.

It doesn't make sense to call initialize() in a case when we always want to restore parameters from a checkpoint (e.g. when using the policy as a learning rate schedule).

PiperOrigin-RevId: 268261947
---
 tensor2tensor/trax/rl/ppo.py         | 17 +++++----------
 tensor2tensor/trax/rl/ppo_test.py    | 32 +++++++++++++++-------------
 tensor2tensor/trax/rl/ppo_trainer.py | 18 +++++++---------
 3 files changed, 30 insertions(+), 37 deletions(-)

diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index d49398ab6..c87000a3e 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -63,18 +63,14 @@
 from jax import numpy as np
 from jax import random as jax_random
 import numpy as onp
+
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.trax import layers as tl
 from tensorflow.io import gfile
 
 
-def policy_and_value_net(rng_key,
-                         batch_observations_shape,
-                         observations_dtype,
-                         n_actions,
-                         bottom_layers_fn=(),
-                         two_towers=True):
+def policy_and_value_net(n_actions, bottom_layers_fn, two_towers):
   """A policy and value net function."""
 
   # Layers.
@@ -100,10 +96,7 @@ def policy_and_value_net(rng_key,
             [tl.Dense(1)],
         )
     ]
-  net = tl.Model(layers)
-  params, state = net.initialize(batch_observations_shape, observations_dtype,
-                                 rng_key)
-  return params, state, net
+  return tl.Model(layers)
 
 
 def optimizer_fn(optimizer, net_params):
@@ -827,8 +820,8 @@ def get_policy_model_file_from_epoch(output_dir, epoch):
   return os.path.join(output_dir, "model-%06d.pkl" % epoch)
 
 
-def maybe_restore_opt_state(output_dir, policy_and_value_opt_state,
-                            policy_and_value_state):
+def maybe_restore_opt_state(output_dir, policy_and_value_opt_state=None,
+                            policy_and_value_state=None):
   """Maybe restore the optimization state from the checkpoint dir.
 
   Optimization state includes parameters and optimizer slots.
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
index 2e53384cf..c250a401a 100644
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -77,14 +77,16 @@ def test_policy_and_value_net(self):
     observation_shape = (3, 4, 5)
     batch_observation_shape = (1, 1) + observation_shape
     n_actions = 2
-    pnv_params, pnv_state, pnv_apply = ppo.policy_and_value_net(
-        self.rng_key, batch_observation_shape, np.float32, n_actions,
-        lambda: [layers.Flatten(n_axes_to_keep=2)])
+    pnv_model = ppo.policy_and_value_net(
+        n_actions, lambda: [layers.Flatten(n_axes_to_keep=2)], two_towers=True)
+    pnv_params, pnv_state = pnv_model.initialize(
+        batch_observation_shape, np.float32, self.rng_key)
+
     batch = 2
     time_steps = 10
     batch_of_observations = np.random.uniform(
         size=(batch, time_steps) + observation_shape)
-    pnv_output, _ = pnv_apply(batch_of_observations, pnv_params, pnv_state)
+    pnv_output, _ = pnv_model(batch_of_observations, pnv_params, pnv_state)
 
     # Output is a list, first is probab of actions and the next is value output.
     self.assertEqual(2, len(pnv_output))
@@ -431,13 +433,13 @@ def test_combined_loss(self):
     B, T, A, OBS = 2, 10, 2, (28, 28, 3)  # pylint: disable=invalid-name
     batch_observation_shape = (1, 1) + OBS
 
-    old_params, _, _ = ppo.policy_and_value_net(
-        key1, batch_observation_shape, np.float32, A,
-        lambda: [layers.Flatten(n_axes_to_keep=2)])
+    net = ppo.policy_and_value_net(
+        A, lambda: [layers.Flatten(n_axes_to_keep=2)], two_towers=True)
 
-    new_params, state, net_apply = ppo.policy_and_value_net(
-        key2, batch_observation_shape, np.float32, A,
-        lambda: [layers.Flatten(n_axes_to_keep=2)])
+    old_params, _ = net.initialize(
+        batch_observation_shape, np.float32, key1)
+    new_params, state = net.initialize(
+        batch_observation_shape, np.float32, key2)
 
     # Generate a batch of observations.
 
@@ -447,10 +449,10 @@ def test_combined_loss(self):
     mask = np.ones_like(rewards)
 
     # Just test that this computes at all.
-    (new_log_probabs, value_predictions_new), _ = net_apply(observations,
-                                                            new_params, state)
-    (old_log_probabs, value_predictions_old), _ = net_apply(observations,
-                                                            old_params, state)
+    (new_log_probabs, value_predictions_new), _ = net(observations, new_params,
+                                                      state)
+    (old_log_probabs, value_predictions_old), _ = net(observations, old_params,
+                                                      state)
 
     gamma = 0.99
     lambda_ = 0.95
@@ -476,7 +478,7 @@ def test_combined_loss(self):
         ppo.combined_loss(new_params,
                           old_log_probabs,
                           value_predictions_old,
-                          net_apply,
+                          net,
                           observations,
                           actions,
                           rewards,
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index bfa8bdec1..b2dd7826e 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -155,17 +155,15 @@ def __init__(
     self._rng, key1 = jax_random.split(self._rng, num=2)
 
     # Initialize the policy and value network.
-    policy_and_value_net_params, self._model_state, policy_and_value_net_apply = (
-        ppo.policy_and_value_net(
-            rng_key=key1,
-            batch_observations_shape=batch_observations_shape,
-            observations_dtype=observations_dtype,
-            n_actions=n_actions,
-            bottom_layers_fn=policy_and_value_model,
-            two_towers=policy_and_value_two_towers,
-        )
+    policy_and_value_net = ppo.policy_and_value_net(
+        n_actions=n_actions,
+        bottom_layers_fn=policy_and_value_model,
+        two_towers=policy_and_value_two_towers,
     )
-    self._policy_and_value_net_apply = jit(policy_and_value_net_apply)
+    self._policy_and_value_net_apply = jit(policy_and_value_net)
+    policy_and_value_net_params, self._model_state = (
+        policy_and_value_net.initialize(
+            batch_observations_shape, observations_dtype, key1))
 
     # Initialize the optimizer.
     (policy_and_value_opt_state, self._policy_and_value_opt_update,

From c3fbb61047a217e8e19691639f79980ca533af78 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 10 Sep 2019 13:47:51 -0700
Subject: [PATCH 2411/2720] [async-rl CL 4/n] : Add `collect_trajectories`
 method to ppo_trainer and use that in train and eval.  - Additionally a
 BaseTrainer._async_mode and BaseTrainer._async_mode_trajectory_subdir were
 also added.  - `collect_trajectories_async` which collects trajectories in
 async mode is unimplemented for now.

PiperOrigin-RevId: 268300991
---
 tensor2tensor/trax/rl/base_trainer.py   |   8 +
 tensor2tensor/trax/rl/ppo.py            |  90 ++++-------
 tensor2tensor/trax/rl/ppo_trainer.py    | 196 +++++++++++++++---------
 tensor2tensor/trax/rl/simple_trainer.py |  62 ++++----
 4 files changed, 190 insertions(+), 166 deletions(-)

diff --git a/tensor2tensor/trax/rl/base_trainer.py b/tensor2tensor/trax/rl/base_trainer.py
index ecd9be22d..6f71cbe97 100644
--- a/tensor2tensor/trax/rl/base_trainer.py
+++ b/tensor2tensor/trax/rl/base_trainer.py
@@ -32,6 +32,7 @@ class BaseTrainer(object):
   def __init__(
       self, train_env, eval_env, output_dir,
       trajectory_dump_dir=None, trajectory_dump_min_count_per_shard=16,
+      async_mode=False, async_mode_trajectory_subdir="trajectories",
   ):
     """Base class constructor.
 
@@ -44,6 +45,10 @@ def __init__(
       trajectory_dump_min_count_per_shard: Minimum number of trajectories to
         collect before dumping in a new shard. Sharding is for efficient
         shuffling for model training in SimPLe.
+      async_mode: (bool) If True, this means we are in async mode and we read
+        trajectories from a location rather than interact with the environment.
+      async_mode_trajectory_subdir: (string) The subdir of output_dir to search
+        for trajectories in async mode.
     """
     self.train_env = train_env
     self.eval_env = eval_env
@@ -54,6 +59,9 @@ def __init__(
         trajectory_dump_min_count_per_shard)
     self._trajectory_buffer = []
 
+    self._async_mode = async_mode
+    self._async_mode_trajectory_subdir = async_mode_trajectory_subdir
+
   @property
   def epoch(self):
     raise NotImplementedError
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index c87000a3e..e9753326b 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -61,7 +61,6 @@
 from jax import jit
 from jax import lax
 from jax import numpy as np
-from jax import random as jax_random
 import numpy as onp
 
 from tensor2tensor.envs import env_problem
@@ -83,7 +82,9 @@ def policy_and_value_net(n_actions, bottom_layers_fn, two_towers):
     layers = [
         tl.Dup(),
         tl.Parallel(
-            [bottom_layers_fn(), tl.Dense(n_actions), tl.LogSoftmax()],
+            [bottom_layers_fn(),
+             tl.Dense(n_actions),
+             tl.LogSoftmax()],
             [bottom_layers_fn(), tl.Dense(1)],
         )
     ]
@@ -142,6 +143,7 @@ def collect_trajectories(env,
                          len_history_for_policy=32,
                          boundary=32,
                          state=None,
+                         temperature=1.0,
                          rng=None):
   """Collect trajectories with the given policy net and behaviour.
 
@@ -158,6 +160,7 @@ def collect_trajectories(env,
       applying the policy on. If None, use the full history.
     boundary: int, pad the sequences to the multiples of this number.
     state: state for `policy_fn`.
+    temperature: (float) temperature to sample action from policy_fn.
     rng: jax rng, splittable.
 
   Returns:
@@ -180,12 +183,13 @@ def collect_trajectories(env,
       len_history_for_policy=len_history_for_policy,
       boundary=boundary,
       state=state,
+      temperature=temperature,
       rng=rng)
   # Skip returning raw_rewards here, since they aren't used.
 
   # t is the return value of Trajectory.as_numpy, so:
   # (observation, action, processed_reward, raw_reward, infos)
-  return [(t[0], t[1], t[2], t[4]) for t in trajs], n_done, timing_info, state
+  return trajs, n_done, timing_info, state
 
 
 # This function can probably be simplified, ask how?
@@ -653,8 +657,8 @@ def combined_loss(new_params,
                   rng=None):
   """Computes the combined (clipped loss + value loss) given observations."""
   (log_probab_actions_new, value_predictions_new), state = (
-      policy_and_value_net_apply(padded_observations, new_params, state,
-                                 rng=rng))
+      policy_and_value_net_apply(
+          padded_observations, new_params, state, rng=rng))
 
   (loss, component_losses, summaries) = combined_loss_given_predictions(
       log_probab_actions_new,
@@ -766,45 +770,6 @@ def masked_entropy(log_probs, mask):
   return -(np.sum(lp * p) / np.sum(mask))
 
 
-def evaluate_policy(eval_env,
-                    get_predictions,
-                    temperatures,
-                    max_timestep=20000,
-                    n_evals=1,
-                    len_history_for_policy=32,
-                    state=None,
-                    rng=None):
-  """Evaluate the policy."""
-
-  processed_reward_sums = collections.defaultdict(list)
-  raw_reward_sums = collections.defaultdict(list)
-  for eval_rng in jax_random.split(rng, num=n_evals):
-    for temperature in temperatures:
-      trajs, _, _, state = env_problem_utils.play_env_problem_with_policy(
-          eval_env,
-          get_predictions,
-          num_trajectories=eval_env.batch_size,
-          max_timestep=max_timestep,
-          reset=True,
-          temperature=temperature,
-          state=state,
-          rng=eval_rng,
-          len_history_for_policy=len_history_for_policy)
-      processed_reward_sums[temperature].extend(sum(traj[2]) for traj in trajs)
-      raw_reward_sums[temperature].extend(sum(traj[3]) for traj in trajs)
-
-  # Return the mean and standard deviation for each temperature.
-  def compute_stats(reward_dict):
-    return {
-        temperature: {"mean": onp.mean(rewards), "std": onp.std(rewards)}
-        for (temperature, rewards) in reward_dict.items()
-    }
-  return {
-      "processed": compute_stats(processed_reward_sums),
-      "raw": compute_stats(raw_reward_sums),
-  }, state
-
-
 def get_policy_model_files(output_dir):
   return list(
       reversed(
@@ -820,7 +785,8 @@ def get_policy_model_file_from_epoch(output_dir, epoch):
   return os.path.join(output_dir, "model-%06d.pkl" % epoch)
 
 
-def maybe_restore_opt_state(output_dir, policy_and_value_opt_state=None,
+def maybe_restore_opt_state(output_dir,
+                            policy_and_value_opt_state=None,
                             policy_and_value_state=None):
   """Maybe restore the optimization state from the checkpoint dir.
 
@@ -853,7 +819,9 @@ def maybe_restore_opt_state(output_dir, policy_and_value_opt_state=None,
       # Try an older version.
       continue
   return (
-      policy_and_value_opt_state, policy_and_value_state, epoch,
+      policy_and_value_opt_state,
+      policy_and_value_state,
+      epoch,
       total_opt_step,
   )
 
@@ -862,17 +830,13 @@ def write_eval_reward_summaries(reward_stats_by_mode, summary_writer, epoch):
   """Writes evaluation reward statistics to summary and logs them.
 
   Args:
-    reward_stats_by_mode: Nested dict of structure:
-      {
+    reward_stats_by_mode: Nested dict of structure: {
           "raw": {
               <temperature 1>: {
                   "mean": <reward mean>,
-                  "std": <reward std>,
-              },
-              <temperature 2>: ...
-          },
-          "processed": ...
-      }
+                  "std": <reward std>, },
+              <temperature 2>: ... },
+          "processed": ... }
     summary_writer: jaxboard.SummaryWriter.
     epoch: Current epoch number.
   """
@@ -881,11 +845,13 @@ def write_eval_reward_summaries(reward_stats_by_mode, summary_writer, epoch):
       for (stat_name, stat) in reward_stats.items():
         summary_writer.scalar(
             "eval/{reward_mode}_reward_{stat_name}/"
-            "temperature_{temperature}".format(reward_mode=reward_mode,
-                                               stat_name=stat_name,
-                                               temperature=temperature),
-            stat, step=epoch)
-      logging.info("Epoch [% 6d] Policy Evaluation (%s reward) "
-                   "[temperature %.2f] = %10.2f (+/- %.2f)",
-                   epoch, reward_mode, temperature,
-                   reward_stats["mean"], reward_stats["std"])
+            "temperature_{temperature}".format(
+                reward_mode=reward_mode,
+                stat_name=stat_name,
+                temperature=temperature),
+            stat,
+            step=epoch)
+      logging.info(
+          "Epoch [% 6d] Policy Evaluation (%s reward) "
+          "[temperature %.2f] = %10.2f (+/- %.2f)", epoch, reward_mode,
+          temperature, reward_stats["mean"], reward_stats["std"])
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index b2dd7826e..19c1eaf06 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import functools
 import os
 import time
@@ -29,6 +30,7 @@
 from jax import jit
 from jax import numpy as np
 from jax import random as jax_random
+import numpy as onp
 from tensor2tensor.trax import jaxboard
 from tensor2tensor.trax import models as trax_models
 from tensor2tensor.trax import optimizers as trax_opt
@@ -37,7 +39,6 @@
 from tensor2tensor.trax.rl import ppo
 from tensorflow.io import gfile
 
-
 DEBUG_LOGGING = False
 GAMMA = 0.99
 LAMBDA = 0.95
@@ -51,35 +52,33 @@
 class PPO(base_trainer.BaseTrainer):
   """PPO trainer."""
 
-  def __init__(
-      self,
-      train_env,
-      eval_env,
-      output_dir,
-      policy_and_value_model=trax_models.FrameStackMLP,
-      policy_and_value_optimizer=functools.partial(
-          trax_opt.Adam, learning_rate=1e-3),
-      policy_and_value_two_towers=False,
-      n_optimizer_steps=N_OPTIMIZER_STEPS,
-      print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
-      target_kl=0.01,
-      boundary=20,
-      max_timestep=None,
-      max_timestep_eval=20000,
-      random_seed=None,
-      gamma=GAMMA,
-      lambda_=LAMBDA,
-      c1=1.0,
-      c2=0.01,
-      eval_every_n=1000,
-      save_every_n=1000,
-      done_frac_for_policy_save=0.5,
-      n_evals=1,
-      len_history_for_policy=4,
-      eval_temperatures=(1.0, 0.5),
-      separate_eval=True,
-      **kwargs
-  ):
+  def __init__(self,
+               train_env,
+               eval_env,
+               output_dir,
+               policy_and_value_model=trax_models.FrameStackMLP,
+               policy_and_value_optimizer=functools.partial(
+                   trax_opt.Adam, learning_rate=1e-3),
+               policy_and_value_two_towers=False,
+               n_optimizer_steps=N_OPTIMIZER_STEPS,
+               print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
+               target_kl=0.01,
+               boundary=20,
+               max_timestep=None,
+               max_timestep_eval=20000,
+               random_seed=None,
+               gamma=GAMMA,
+               lambda_=LAMBDA,
+               c1=1.0,
+               c2=0.01,
+               eval_every_n=1000,
+               save_every_n=1000,
+               done_frac_for_policy_save=0.5,
+               n_evals=1,
+               len_history_for_policy=4,
+               eval_temperatures=(1.0, 0.5),
+               separate_eval=True,
+               **kwargs):
     """Creates the PPO trainer.
 
     Args:
@@ -97,8 +96,8 @@ def __init__(
       target_kl: Policy iteration early stopping. Set to infinity to disable
         early stopping.
       boundary: We pad trajectories at integer multiples of this number.
-      max_timestep: If set to an integer, maximum number of time-steps in
-        a trajectory. Used in the collect procedure.
+      max_timestep: If set to an integer, maximum number of time-steps in a
+        trajectory. Used in the collect procedure.
       max_timestep_eval: If set to an integer, maximum number of time-steps in
         an evaluation trajectory. Used in the collect procedure.
       random_seed: Random seed.
@@ -162,8 +161,8 @@ def __init__(
     )
     self._policy_and_value_net_apply = jit(policy_and_value_net)
     policy_and_value_net_params, self._model_state = (
-        policy_and_value_net.initialize(
-            batch_observations_shape, observations_dtype, key1))
+        policy_and_value_net.initialize(batch_observations_shape,
+                                        observations_dtype, key1))
 
     # Initialize the optimizer.
     (policy_and_value_opt_state, self._policy_and_value_opt_update,
@@ -198,13 +197,15 @@ def train_env(self):
   @train_env.setter
   def train_env(self, new_train_env):
     if self._train_env is not None:
+
       def assert_same_space(space1, space2):
         assert space1.shape == space2.shape
         assert space1.dtype == space2.dtype
-      assert_same_space(
-          new_train_env.observation_space, self._train_env.observation_space)
-      assert_same_space(
-          new_train_env.action_space, self._train_env.action_space)
+
+      assert_same_space(new_train_env.observation_space,
+                        self._train_env.observation_space)
+      assert_same_space(new_train_env.action_space,
+                        self._train_env.action_space)
       # We don't check the reward range, because PPO will work either way.
 
     self._train_env = new_train_env
@@ -214,6 +215,48 @@ def assert_same_space(space1, space2):
   def epoch(self):
     return self._epoch
 
+  def collect_trajectories_async(self, env, train=True, n_trajectories=1):
+    """Collects trajectories in an async manner."""
+
+    assert self._async_mode
+    del env
+    del train
+    del n_trajectories
+    raise NotImplementedError
+
+  def collect_trajectories(self, train=True, temperature=1.0):
+    self._rng, key = jax_random.split(self._rng)
+
+    env = self.train_env
+    max_timestep = self._max_timestep
+    should_reset = self._should_reset
+    if not train:  # eval
+      env = self.eval_env
+      max_timestep = self._max_timestep_eval
+      should_reset = True
+
+    n_trajectories = env.batch_size
+
+    # If async, read the required trajectories for the epoch.
+    if self._async_mode:
+      return self.collect_trajectories_async(
+          env, train=train, n_trajectories=n_trajectories)
+
+    trajs, n_done, timing_info, self._model_state = ppo.collect_trajectories(
+        env,
+        policy_fn=self._get_predictions,
+        n_trajectories=n_trajectories,
+        max_timestep=max_timestep,
+        state=self._model_state,
+        rng=key,
+        len_history_for_policy=self._len_history_for_policy,
+        boundary=self._boundary,
+        reset=should_reset,
+        temperature=temperature,
+    )
+
+    return trajs, n_done, timing_info, self._model_state
+
   def train_epoch(self, evaluate=True):
     """Train one PPO epoch."""
     epoch_start_time = time.time()
@@ -229,17 +272,9 @@ def train_epoch(self, evaluate=True):
     trajectory_collection_start_time = time.time()
     logging.vlog(1, "PPO epoch [% 6d]: collecting trajectories.", self._epoch)
     self._rng, key = jax_random.split(self._rng)
-    trajs, n_done, timing_info, self._model_state = ppo.collect_trajectories(
-        self.train_env,
-        policy_fn=self._get_predictions,
-        n_trajectories=self.train_env.batch_size,
-        max_timestep=self._max_timestep,
-        state=self._model_state,
-        rng=key,
-        len_history_for_policy=self._len_history_for_policy,
-        boundary=self._boundary,
-        reset=self._should_reset,
-    )
+    trajs, n_done, timing_info, self._model_state = self.collect_trajectories(
+        train=True, temperature=1.0)
+    trajs = [(t[0], t[1], t[2], t[4]) for t in trajs]
     self._should_reset = False
     trajectory_collection_time = ppo.get_time(trajectory_collection_start_time)
 
@@ -270,8 +305,8 @@ def train_epoch(self, evaluate=True):
     logging.vlog(2, "Trajectory Lengths: %s", [len(traj[0]) for traj in trajs])
 
     padding_start_time = time.time()
-    (_, reward_mask, padded_observations, padded_actions,
-     padded_rewards, padded_infos) = ppo.pad_trajectories(
+    (_, reward_mask, padded_observations, padded_actions, padded_rewards,
+     padded_infos) = ppo.pad_trajectories(
          trajs, boundary=self._boundary)
     padding_time = ppo.get_time(padding_start_time)
 
@@ -287,8 +322,8 @@ def train_epoch(self, evaluate=True):
     assert (B, T) == padded_rewards.shape
     assert (B, T) == reward_mask.shape
     assert (B, T + 1) == padded_observations.shape[:2]
-    assert ((B, T + 1) + self.train_env.observation_space.shape ==
-            padded_observations.shape)
+    assert ((B, T + 1) +
+            self.train_env.observation_space.shape == padded_observations.shape)
 
     log_prob_recompute_start_time = time.time()
     assert ("log_prob_actions" in padded_infos and
@@ -395,9 +430,11 @@ def train_epoch(self, evaluate=True):
 
       # Compute the approx KL for early stopping.
       (log_probab_actions_new, _), self._model_state = (
-          self._policy_and_value_net_apply(padded_observations,
-                                           self._policy_and_value_net_params,
-                                           self._model_state, rng=k2))
+          self._policy_and_value_net_apply(
+              padded_observations,
+              self._policy_and_value_net_params,
+              self._model_state,
+              rng=k2))
 
       approx_kl = ppo.approximate_kl(log_probab_actions_new, log_probabs_traj,
                                      reward_mask)
@@ -502,10 +539,8 @@ def train_epoch(self, evaluate=True):
         "%s : % 10.2f" % (k.rjust(max_key_len + 1), v)
         for k, v in sorted(timing_dict.items())
     ]
-    logging.info(
-        "PPO epoch [% 6d], Timings: \n%s", last_epoch,
-        "\n".join(timing_info_list)
-    )
+    logging.info("PPO epoch [% 6d], Timings: \n%s", last_epoch,
+                 "\n".join(timing_info_list))
 
     # Flush summary writers once in a while.
     if self._epoch % 1000 == 0:
@@ -516,16 +551,32 @@ def evaluate(self):
     if not self._separate_eval:
       return
     logging.vlog(1, "PPO epoch [% 6d]: evaluating policy.", self._epoch)
-    self._rng, key = jax_random.split(self._rng, num=2)
-    reward_stats, self._model_state = ppo.evaluate_policy(
-        self.eval_env,
-        self._get_predictions,
-        temperatures=self._eval_temperatures,
-        max_timestep=self._max_timestep_eval,
-        n_evals=self._n_evals,
-        len_history_for_policy=self._len_history_for_policy,
-        state=self._model_state,
-        rng=key)
+
+    processed_reward_sums = collections.defaultdict(list)
+    raw_reward_sums = collections.defaultdict(list)
+    for _ in range(self._n_evals):
+      for temperature in self._eval_temperatures:
+        trajs, _, _, self._model_state = self.collect_trajectories(
+            train=False, temperature=temperature)
+
+        processed_reward_sums[temperature].extend(
+            sum(traj[2]) for traj in trajs)
+        raw_reward_sums[temperature].extend(sum(traj[3]) for traj in trajs)
+
+    # Return the mean and standard deviation for each temperature.
+    def compute_stats(reward_dict):
+      return {
+          temperature: {  # pylint: disable=g-complex-comprehension
+              "mean": onp.mean(rewards),
+              "std": onp.std(rewards)
+          } for (temperature, rewards) in reward_dict.items()
+      }
+
+    reward_stats = {
+        "processed": compute_stats(processed_reward_sums),
+        "raw": compute_stats(raw_reward_sums),
+    }
+
     ppo.write_eval_reward_summaries(
         reward_stats, self._eval_sw, epoch=self._epoch)
 
@@ -535,9 +586,8 @@ def save(self):
     old_model_files = ppo.get_policy_model_files(self._output_dir)
     params_file = os.path.join(self._output_dir, "model-%06d.pkl" % self._epoch)
     with gfile.GFile(params_file, "wb") as f:
-      pickle.dump(
-          (self._policy_and_value_opt_state, self._model_state,
-           self._total_opt_step), f)
+      pickle.dump((self._policy_and_value_opt_state, self._model_state,
+                   self._total_opt_step), f)
     # Remove the old model files.
     for path in old_model_files:
       if path != params_file:
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index c03645eef..c0513f265 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -37,32 +37,31 @@
 class SimPLe(base_trainer.BaseTrainer):
   """SimPLe trainer."""
 
-  def __init__(
-      self,
-      train_env,
-      eval_env,
-      output_dir,
-      policy_trainer_class,
-      n_real_epochs=10,
-      data_eval_frac=0.125,
-      model_train_batch_size=64,
-      n_model_train_steps=1000,
-      simulated_env_problem_class=(
-          simulated_env_problem.SerializedSequenceSimulatedEnvProblem),
-      simulated_batch_size=16,
-      n_simulated_epochs=1000,
-      trajectory_dump_dir=None,
-      initial_trajectory_dir=None,
-      initial_trajectory_mix_prob=0.5,
-      **kwargs
-  ):
-    super(SimPLe, self).__init__(
-        train_env, eval_env, output_dir, **kwargs)
+  def __init__(self,
+               train_env,
+               eval_env,
+               output_dir,
+               policy_trainer_class,
+               n_real_epochs=10,
+               data_eval_frac=0.125,
+               model_train_batch_size=64,
+               n_model_train_steps=1000,
+               simulated_env_problem_class=(
+                   simulated_env_problem.SerializedSequenceSimulatedEnvProblem),
+               simulated_batch_size=16,
+               n_simulated_epochs=1000,
+               trajectory_dump_dir=None,
+               initial_trajectory_dir=None,
+               initial_trajectory_mix_prob=0.5,
+               **kwargs):
+    super(SimPLe, self).__init__(train_env, eval_env, output_dir, **kwargs)
     self._policy_dir = os.path.join(output_dir, "policy")
     self._policy_trainer = policy_trainer_class(
         train_env=train_env,
         eval_env=eval_env,
         output_dir=self._policy_dir,
+        async_mode=self._async_mode,
+        async_mode_trajectory_subdir=self._async_mode_trajectory_subdir,
     )
     self._n_real_epochs = n_real_epochs
     self._model_train_batch_size = model_train_batch_size
@@ -145,8 +144,7 @@ def train_model(self):
         input_dtype=self._sim_env.model_input_dtype,
         # TODO(lukaszkaiser): correct those, they may differ from inputs.
         target_shape=self._sim_env.model_input_shape,
-        target_dtype=self._sim_env.model_input_dtype
-    )
+        target_dtype=self._sim_env.model_input_dtype)
 
     self._model_train_step += self._n_model_train_steps
     trax.train(
@@ -179,6 +177,7 @@ def _has_initial_data(self):
     return self._initial_trajectory_dir is not None
 
   def _make_input_streams(self):
+
     def make_example_streams(trajectory_dir):
       (train_trajs, eval_trajs) = simple.load_trajectories(
           trajectory_dir, eval_frac=self._data_eval_frac)
@@ -196,16 +195,16 @@ def make_example_streams(trajectory_dir):
 
     if self._has_initial_data:
       # Load the initial, precollected data.
-      (init_train_stream, init_eval_stream) = make_example_streams(
-          self._initial_trajectory_dir)
+      (init_train_stream,
+       init_eval_stream) = make_example_streams(self._initial_trajectory_dir)
     else:
       (init_train_stream, init_eval_stream) = (None, None)
       mix_prob = 0.0  # Take just our own collected data.
 
     if self._has_own_data:
       # Load trajectories collected in all epochs so far.
-      (own_train_stream, own_eval_stream) = make_example_streams(
-          self._trajectory_dump_root_dir)
+      (own_train_stream,
+       own_eval_stream) = make_example_streams(self._trajectory_dump_root_dir)
     else:
       # We start the loop with training the model, so we don't have our own
       # collected data yet.
@@ -217,10 +216,11 @@ def mix_and_batch(streams):
       mixed_stream = simple.mix_streams(init_stream, own_stream, mix_prob)
       return simple.batch_stream(mixed_stream, self._model_train_batch_size)
 
-    return tuple(map(mix_and_batch, (
-        (init_train_stream, own_train_stream),
-        (init_eval_stream, own_eval_stream),
-    )))
+    return tuple(
+        map(mix_and_batch, (
+            (init_train_stream, own_train_stream),
+            (init_eval_stream, own_eval_stream),
+        )))
 
   def evaluate_model(self):
     logging.info("SimPLe epoch [% 6d]: evaluating model.", self._simple_epoch)

From d9a6d7b6327fdaf69f9f5e230977236b45a127ab Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 10 Sep 2019 14:52:06 -0700
Subject: [PATCH 2412/2720] Add enwik8_l65k problem

PiperOrigin-RevId: 268315466
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/enwik8.py       | 128 ++++++++++++++++++
 2 files changed, 129 insertions(+)
 create mode 100644 tensor2tensor/data_generators/enwik8.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 19ccffc78..a865a0ceb 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -39,6 +39,7 @@
     "tensor2tensor.data_generators.cola",
     "tensor2tensor.data_generators.common_voice",
     "tensor2tensor.data_generators.desc2code",
+    "tensor2tensor.data_generators.enwik8",
     "tensor2tensor.data_generators.fsns",
     "tensor2tensor.data_generators.function_docstring",
     "tensor2tensor.data_generators.gene_expression",
diff --git a/tensor2tensor/data_generators/enwik8.py b/tensor2tensor/data_generators/enwik8.py
new file mode 100644
index 000000000..8fa41a272
--- /dev/null
+++ b/tensor2tensor/data_generators/enwik8.py
@@ -0,0 +1,128 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Data generators for enwik8 data-set."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import zipfile
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import registry
+
+import tensorflow as tf
+
+
+def _maybe_download_corpus(tmp_dir):
+  """Download and unpack the corpus.
+
+  Args:
+    tmp_dir: directory containing dataset.
+
+  Returns:
+    path to entire corpus as a text file.
+  """
+  corpus_url = "http://mattmahoney.net/dc/enwik8.zip"
+  corpus_filename = os.path.basename(corpus_url)
+  compressed_filepath = generator_utils.maybe_download(
+      tmp_dir, corpus_filename, corpus_url)
+
+  zip_ref = zipfile.ZipFile(compressed_filepath, "r")
+  zip_ref.extractall(tmp_dir)
+  zip_ref.close()
+
+  return os.path.join(tmp_dir, "enwik8")
+
+
+@registry.register_problem
+class Enwik8L65k(text_problems.Text2SelfProblem):
+  """Enwiki8, with examples up to 65,536 characters long."""
+
+  DUPE_FACTOR = 4
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.CHARACTER
+
+  def global_task_id(self):
+    return problem.TaskID.EN_CHR
+
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each."""
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 16,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }, {
+        "split": problem.DatasetSplit.TEST,
+        "shards": 1,
+    }]
+
+  def max_length(self, model_hparams):
+    return self.sequence_length
+
+  @property
+  def sequence_length(self):
+    """Length of each example (number of characters)."""
+    return 65536
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    filepath = _maybe_download_corpus(tmp_dir)
+    with tf.io.gfile.GFile(filepath) as f:
+      data = f.read()
+
+    tf.logging.info("Length of enwik8 = %d", len(data))
+
+    num_test_chars = 5000000
+
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      part = data[: -2 * num_test_chars]
+    elif dataset_split == problem.DatasetSplit.EVAL:
+      part = data[-2 * num_test_chars: -num_test_chars]
+    elif dataset_split == problem.DatasetSplit.TEST:
+      part = data[-num_test_chars:]
+    else:
+      raise ValueError("Undefined dataset_split")
+
+    tf.logging.info("Length of split '%s' = %d", dataset_split, len(part))
+
+    # TODO(kitaev): Better handling of evaluation data, to ensure that there is
+    # always context available.
+    if dataset_split == problem.DatasetSplit.TRAIN:
+      offset = self.sequence_length // self.DUPE_FACTOR
+      for start in range(0, len(part), offset):
+        yield {"targets": part[start:start+self.sequence_length]}
+    else:
+      for start in range(0, len(part), self.sequence_length):
+        yield {"targets": part[start:start+self.sequence_length]}
+
+  def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
+    generator = self.generate_samples(data_dir, tmp_dir, dataset_split)
+    vocab = self.get_or_create_vocab(data_dir, tmp_dir)
+    for sample in generator:
+      sample["targets"] = vocab.encode(sample["targets"])
+      yield sample

From 211605d86723703ea8510f18640808061c7089c5 Mon Sep 17 00:00:00 2001
From: Daniel De Freitas Adiwardana <adiwardana@google.com>
Date: Tue, 10 Sep 2019 16:03:31 -0700
Subject: [PATCH 2413/2720] Adding hparams.num_trainable_top_decoder_layers to
 evolved_transformer for allowing fine tuning of only N top decoder layers.

PiperOrigin-RevId: 268330777
---
 tensor2tensor/models/evolved_transformer.py   |  41 +-
 .../models/evolved_transformer_test.py        | 532 +++++++++++++++++-
 2 files changed, 544 insertions(+), 29 deletions(-)

diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 2da51fbd2..3a2a497ef 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -26,6 +26,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
 from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
 
 import tensorflow as tf
 
@@ -72,6 +73,30 @@ def __init__(self, *args, **kwargs):
     self._decoder_function = evolved_transformer_decoder
     self._init_cache_fn = _init_evolved_transformer_cache
 
+    # -1 means train all weights.
+    if self.hparams.get("num_trainable_top_decoder_layers", -1) < 0:
+      t2t_model.log_info(
+          "num_trainable_top_decoder_layers is negative so training all weights."
+      )
+    elif self.hparams.shared_embedding_and_softmax_weights:
+      t2t_model.log_info(
+          "Setting hparams.shared_embedding_and_softmax_weights to False, "
+          "because hparam.num_trainable_top_decoder_layers is being used.")
+
+      # When hparam.num_trainable_top_decoder_layers is set to N >= 0 we will
+      # freeze (not train) every variable except the N top decoder layers and
+      # the (pre-)softmax matrix. For any N >= 0 we will freeze the encoder and
+      # input/target embeddings. This also means we will not share the
+      # (pre-)softmax matrix with input/target embeddings otherwise they will be
+      # trained as well.
+      self.hparams.shared_embedding_and_softmax_weights = False
+
+      # If hparams.shared_embedding_and_softmax_weights was previously True,
+      # then input and target embeddings were being shared.
+      # To make sure it they embeddings continue to be shared, we need to set
+      # hparams.shared_embedding to True.
+      self.hparams.shared_embedding = True
+
 
 def evolved_transformer_encoder(encoder_input,
                                 encoder_self_attention_bias,
@@ -289,6 +314,12 @@ def evolved_transformer_decoder(decoder_input,
   """
   del losses
 
+  num_trainable_top_decoder_layers = hparams.get(
+      "num_trainable_top_decoder_layers", -1)  # -1 means train all weights.
+
+  if num_trainable_top_decoder_layers >= 0:
+    encoder_output = tf.stop_gradient(encoder_output)
+
   attention_dropout_broadcast_dims = (
       common_layers.comma_separated_string_to_integer_list(
           getattr(hparams, "attention_dropout_broadcast_dims", "")))
@@ -296,7 +327,10 @@ def evolved_transformer_decoder(decoder_input,
   with tf.variable_scope(name):
     hidden_state = decoder_input
 
-    for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
+    num_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
+    for layer in range(num_layers):
+      if num_trainable_top_decoder_layers == num_layers - layer:
+        hidden_state = tf.stop_gradient(hidden_state)
       layer_name = "layer_%d" % layer
       layer_cache = cache[layer_name] if cache is not None else None
       with tf.variable_scope(layer_name):
@@ -586,7 +620,10 @@ def evolved_transformer_decoder(decoder_input,
           hidden_state = common_layers.layer_postprocess(
               residual_state, hidden_state, hparams)
 
-    return common_layers.layer_preprocess(hidden_state, hparams)
+    decoder_output = common_layers.layer_preprocess(hidden_state, hparams)
+    if num_trainable_top_decoder_layers == 0:
+      decoder_output = tf.stop_gradient(decoder_output)
+    return decoder_output
 
 
 def _add_attend_to_encoder_cache(cache, attention_name, hparams, num_layers,
diff --git a/tensor2tensor/models/evolved_transformer_test.py b/tensor2tensor/models/evolved_transformer_test.py
index 0761a06b4..f989390f1 100644
--- a/tensor2tensor/models/evolved_transformer_test.py
+++ b/tensor2tensor/models/evolved_transformer_test.py
@@ -19,7 +19,6 @@
 from __future__ import division
 from __future__ import print_function
 import numpy as np
-
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.models import evolved_transformer
 from tensor2tensor.models import transformer
@@ -33,12 +32,44 @@
 DECODE_LENGTH = 3
 
 
-def get_model(hparams, has_input=True):
+def print_vars(all_vars=None):
+  """Print info about a list of variables."""
+  if not all_vars:
+    all_vars = tf.trainable_variables()
+  tf.logging.info("Format: <name>, <shape>, <(soft) device placement>")
+  for var in all_vars:
+    tf.logging.info("  %s, %s, %s" %
+                    (var.name, str(var.get_shape()), var.op.device))
+
+
+def get_var(name):
+  """Get trainable variable by name."""
+  variables = [var for var in tf.trainable_variables() if var.name == name]
+  if len(variables) == 1:
+    return variables[0]
+  raise ValueError("`name` must match exactly one variable. '%s' matched %d" %
+                   (name, len(variables)))
+
+
+def get_vars(names):
+  """Get trainable variables by name."""
+  return [get_var(name) for name in names]
+
+
+def assert_with_message(assert_method, a, b, message):
+  try:
+    assert_method(a, b)
+  except AssertionError as e:
+    tf.logging.error(message)
+    raise e
+
+
+def get_model(hparams, has_input=True, num_decoder_layers=1):
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.hidden_size = 4
   hparams.num_heads = 1
   hparams.num_encoder_layers = 1
-  hparams.num_decoder_layers = 1
+  hparams.num_decoder_layers = num_decoder_layers
 
   p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE, VOCAB_SIZE,
                                                    hparams)
@@ -56,8 +87,9 @@ def get_model(hparams, has_input=True):
   if has_input:
     features["inputs"] = tf.constant(inputs, dtype=tf.int32, name="inputs")
 
-  return (evolved_transformer.EvolvedTransformer(
-      hparams, tf.estimator.ModeKeys.TRAIN, p_hparams), features)
+  return (evolved_transformer.EvolvedTransformer(hparams,
+                                                 tf.estimator.ModeKeys.TRAIN,
+                                                 p_hparams), features)
 
 
 class EvolvedTransformerTest(tf.test.TestCase):
@@ -91,8 +123,8 @@ def testSlowVsFast(self):
     model.set_mode(tf.estimator.ModeKeys.PREDICT)
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
-      greedy_result = model._slow_greedy_infer(
-          features, decode_length)["outputs"]
+      greedy_result = model._slow_greedy_infer(features,
+                                               decode_length)["outputs"]
       greedy_result = tf.squeeze(greedy_result, axis=[2, 3])
 
       fast_result = model._greedy_infer(features, decode_length)["outputs"]
@@ -105,8 +137,7 @@ def testSlowVsFast(self):
     self.assertAllClose(greedy_res, fast_res)
 
   def testSlowVsFastNoInput(self):
-    model, features = get_model(
-        transformer.transformer_tiny(), has_input=False)
+    model, features = get_model(transformer.transformer_tiny(), has_input=False)
 
     decode_length = DECODE_LENGTH
 
@@ -126,8 +157,7 @@ def testSlowVsFastNoInput(self):
     model.set_mode(tf.estimator.ModeKeys.PREDICT)
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
-      slow_result = model._slow_greedy_infer(
-          features, decode_length)["outputs"]
+      slow_result = model._slow_greedy_infer(features, decode_length)["outputs"]
       slow_result = tf.squeeze(slow_result, axis=[2, 3])
 
       fast_result = model._greedy_infer(features, decode_length)["outputs"]
@@ -161,17 +191,11 @@ def testBeamVsFast(self):
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
       beam_result = model._beam_decode_slow(
-          features,
-          decode_length,
-          beam_size=4,
-          top_beams=1,
+          features, decode_length, beam_size=4, top_beams=1,
           alpha=1.0)["outputs"]
 
       fast_result = model._beam_decode(
-          features,
-          decode_length,
-          beam_size=4,
-          top_beams=1,
+          features, decode_length, beam_size=4, top_beams=1,
           alpha=1.0)["outputs"]
 
     with self.test_session():
@@ -212,12 +236,12 @@ def testGreedySlowTPUVsNonTPU(self):
     model, features = self._create_greedy_infer_model()
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
-      slow_result_non_tpu = model._slow_greedy_infer(
-          features, decode_length)["outputs"]
+      slow_result_non_tpu = model._slow_greedy_infer(features,
+                                                     decode_length)["outputs"]
       slow_result_non_tpu = tf.squeeze(slow_result_non_tpu, axis=[2, 3])
 
-      slow_result_tpu = model._slow_greedy_infer_tpu(
-          features, decode_length)["outputs"]
+      slow_result_tpu = model._slow_greedy_infer_tpu(features,
+                                                     decode_length)["outputs"]
       slow_result_tpu = tf.squeeze(slow_result_tpu, axis=[2, 3])
 
     with self.test_session():
@@ -254,8 +278,8 @@ def testGreedyTPUSlowVsFast(self):
     model, features = self._create_greedy_infer_model()
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
-      slow_result = model._slow_greedy_infer_tpu(
-          features, decode_length)["outputs"]
+      slow_result = model._slow_greedy_infer_tpu(features,
+                                                 decode_length)["outputs"]
       slow_result = tf.squeeze(slow_result, axis=[2, 3])
 
       fast_result = model._greedy_infer(
@@ -265,10 +289,464 @@ def testGreedyTPUSlowVsFast(self):
       slow_res = slow_result.eval()
       fast_res = fast_result.eval()
 
-    self.assertEqual(fast_res.shape,
-                     (BATCH_SIZE, INPUT_LENGTH + decode_length))
+    self.assertEqual(fast_res.shape, (BATCH_SIZE, INPUT_LENGTH + decode_length))
     self.assertAllClose(fast_res, slow_res)
 
+  def testFrozenWeightsUnchangedByTraining(self):
+    # Arrange.
+    hparams = transformer.transformer_tiny()
+    hparams.add_hparam("num_trainable_top_decoder_layers", 1)
+    model, features = get_model(hparams, num_decoder_layers=3)
+    out_logits, _ = model(features)
+    out_logits = tf.squeeze(out_logits, axis=[2, 3])
+    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
+        logits=tf.reshape(out_logits, [-1, VOCAB_SIZE]),
+        labels=tf.reshape(features["targets"], [-1]))
+    loss = tf.reduce_mean(loss)
+    apply_grad = tf.train.AdamOptimizer(0.001).minimize(loss)
+    frozen_names = [
+        "evolved_transformer/symbol_modality_10_4/shared/weights_0:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_1:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_2:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_3:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_4:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_5:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_6:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_7:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_8:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_9:0",
+        "evolved_transformer/body/target_space_embedding/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/gated_linear_unit/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/encoder/layer_0/gated_linear_unit/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/encoder/layer_0/gated_linear_unit/dense/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/gated_linear_unit/dense/bias:0",
+        "evolved_transformer/body/encoder/layer_0/gated_linear_unit/dense_1/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/gated_linear_unit/dense_1/bias:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/dense/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/dense/bias:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/standard_conv_3x1/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/standard_conv_3x1/bias:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/separable_conv_9x1/depthwise_kernel:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/separable_conv_9x1/pointwise_kernel:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/separable_conv_9x1/bias:0",
+        "evolved_transformer/body/encoder/layer_0/self_attention/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/encoder/layer_0/self_attention/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/encoder/layer_0/self_attention/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/self_attention/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/self_attention/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/self_attention/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/dense_layers/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/encoder/layer_0/dense_layers/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/encoder/layer_0/dense_layers/dense/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/dense_layers/dense/bias:0",
+        "evolved_transformer/body/encoder/layer_0/dense_layers/dense_1/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/dense_layers/dense_1/bias:0",
+        "evolved_transformer/body/encoder/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/encoder/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/16_head_self_attention/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_0/16_head_self_attention/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/16_head_self_attention/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/16_head_self_attention/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/16_head_self_attention/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/16_head_self_attention/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/first_attend_to_encoder/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/first_attend_to_encoder/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/first_attend_to_encoder/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/first_attend_to_encoder/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv11x1/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv11x1/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv11x1/bias:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv_7x1_1/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv_7x1_1/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv_7x1_1/bias:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv_7x1_2/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv_7x1_2/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv_7x1_2/bias:0",
+        "evolved_transformer/body/decoder/layer_0/self_attention/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_0/self_attention/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/self_attention/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/self_attention/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/self_attention/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/self_attention/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/second_attend_to_encoder/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_0/second_attend_to_encoder/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/second_attend_to_encoder/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/second_attend_to_encoder/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/second_attend_to_encoder/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/second_attend_to_encoder/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/dense/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/dense/bias:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/layer_prepostprocess_1/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/layer_prepostprocess_1/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/dense_1/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/dense_1/bias:0",
+        "evolved_transformer/body/decoder/layer_1/16_head_self_attention/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_1/16_head_self_attention/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_1/16_head_self_attention/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/16_head_self_attention/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/16_head_self_attention/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/16_head_self_attention/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/first_attend_to_encoder/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/first_attend_to_encoder/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/first_attend_to_encoder/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/first_attend_to_encoder/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv11x1/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv11x1/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv11x1/bias:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv_7x1_1/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv_7x1_1/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv_7x1_1/bias:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv_7x1_2/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv_7x1_2/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv_7x1_2/bias:0",
+        "evolved_transformer/body/decoder/layer_1/self_attention/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_1/self_attention/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_1/self_attention/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/self_attention/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/self_attention/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/self_attention/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/second_attend_to_encoder/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_1/second_attend_to_encoder/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_1/second_attend_to_encoder/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/second_attend_to_encoder/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/second_attend_to_encoder/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/second_attend_to_encoder/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/dense/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/dense/bias:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/layer_prepostprocess_1/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/layer_prepostprocess_1/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/dense_1/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/dense_1/bias:0",
+    ]
+    train_names = [
+        "evolved_transformer/body/decoder/layer_2/16_head_self_attention/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_2/16_head_self_attention/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_2/16_head_self_attention/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/16_head_self_attention/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/16_head_self_attention/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/16_head_self_attention/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/first_attend_to_encoder/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/first_attend_to_encoder/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/first_attend_to_encoder/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/first_attend_to_encoder/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv11x1/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv11x1/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv11x1/bias:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv_7x1_1/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv_7x1_1/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv_7x1_1/bias:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv_7x1_2/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv_7x1_2/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv_7x1_2/bias:0",
+        "evolved_transformer/body/decoder/layer_2/self_attention/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_2/self_attention/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_2/self_attention/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/self_attention/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/self_attention/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/self_attention/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/second_attend_to_encoder/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_2/second_attend_to_encoder/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_2/second_attend_to_encoder/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/second_attend_to_encoder/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/second_attend_to_encoder/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/second_attend_to_encoder/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/dense/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/dense/bias:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/layer_prepostprocess_1/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/layer_prepostprocess_1/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/dense_1/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/dense_1/bias:0",
+        "evolved_transformer/body/decoder/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/symbol_modality_10_4/softmax/weights_1:0",
+        "evolved_transformer/symbol_modality_10_4/softmax/weights_2:0",
+        "evolved_transformer/symbol_modality_10_4/softmax/weights_3:0",
+        "evolved_transformer/symbol_modality_10_4/softmax/weights_4:0",
+        "evolved_transformer/symbol_modality_10_4/softmax/weights_5:0",
+        "evolved_transformer/symbol_modality_10_4/softmax/weights_6:0",
+        "evolved_transformer/symbol_modality_10_4/softmax/weights_7:0",
+        "evolved_transformer/symbol_modality_10_4/softmax/weights_8:0",
+        "evolved_transformer/symbol_modality_10_4/softmax/weights_9:0",
+    ]
+    frozen_vars = get_vars(frozen_names)
+    train_vars = get_vars(train_names)
+    print_vars()
+
+    # Act.
+    with self.test_session() as session:
+      tf.global_variables_initializer().run()
+      frozen_values_before = session.run(frozen_vars)
+      train_values_before = session.run(train_vars)
+      for _ in range(10):  # Arbitrary number of training steps.
+        apply_grad.run()
+      frozen_values_after = session.run(frozen_vars)
+      train_values_after = session.run(train_vars)
+
+    # Assert.
+    self.assertTrue(
+        model._original_hparams.shared_embedding_and_softmax_weights)
+    self.assertFalse(model.hparams.shared_embedding_and_softmax_weights)
+    self.assertTrue(model.hparams.shared_embedding)
+    for name, before, after in zip(frozen_names, frozen_values_before,
+                                   frozen_values_after):
+      assert_with_message(
+          self.assertAllClose, before, after,
+          "%s should be frozen, but changed after training." % name)
+    for name, before, after in zip(train_names, train_values_before,
+                                   train_values_after):
+      assert_with_message(
+          self.assertNotAllClose, before, after,
+          "%s should be trainable, but did not change after training." % name)
+
+  def testAllWeightsTrainableByDefault(self):
+    # Arrange.
+    model, features = get_model(
+        transformer.transformer_tiny(), num_decoder_layers=3)
+    out_logits, _ = model(features)
+    out_logits = tf.squeeze(out_logits, axis=[2, 3])
+    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
+        logits=tf.reshape(out_logits, [-1, VOCAB_SIZE]),
+        labels=tf.reshape(features["targets"], [-1]))
+    loss = tf.reduce_mean(loss)
+    apply_grad = tf.train.AdamOptimizer(0.001).minimize(loss)
+    var_names = [
+        "evolved_transformer/symbol_modality_10_4/shared/weights_0:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_1:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_2:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_3:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_4:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_5:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_6:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_7:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_8:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_9:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_10:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_11:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_12:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_13:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_14:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_15:0",
+        "evolved_transformer/body/target_space_embedding/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/gated_linear_unit/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/encoder/layer_0/gated_linear_unit/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/encoder/layer_0/gated_linear_unit/dense/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/gated_linear_unit/dense/bias:0",
+        "evolved_transformer/body/encoder/layer_0/gated_linear_unit/dense_1/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/gated_linear_unit/dense_1/bias:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/dense/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/dense/bias:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/standard_conv_3x1/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/standard_conv_3x1/bias:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/separable_conv_9x1/depthwise_kernel:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/separable_conv_9x1/pointwise_kernel:0",
+        "evolved_transformer/body/encoder/layer_0/conv_branches/separable_conv_9x1/bias:0",
+        "evolved_transformer/body/encoder/layer_0/self_attention/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/encoder/layer_0/self_attention/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/encoder/layer_0/self_attention/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/self_attention/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/self_attention/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/self_attention/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/dense_layers/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/encoder/layer_0/dense_layers/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/encoder/layer_0/dense_layers/dense/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/dense_layers/dense/bias:0",
+        "evolved_transformer/body/encoder/layer_0/dense_layers/dense_1/kernel:0",
+        "evolved_transformer/body/encoder/layer_0/dense_layers/dense_1/bias:0",
+        "evolved_transformer/body/encoder/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/encoder/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/16_head_self_attention/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_0/16_head_self_attention/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/16_head_self_attention/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/16_head_self_attention/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/16_head_self_attention/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/16_head_self_attention/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/first_attend_to_encoder/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/first_attend_to_encoder/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/first_attend_to_encoder/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/first_attend_to_encoder/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv11x1/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv11x1/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv11x1/bias:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv_7x1_1/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv_7x1_1/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv_7x1_1/bias:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv_7x1_2/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv_7x1_2/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_0/conv_branches/separable_conv_7x1_2/bias:0",
+        "evolved_transformer/body/decoder/layer_0/self_attention/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_0/self_attention/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/self_attention/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/self_attention/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/self_attention/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/self_attention/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/second_attend_to_encoder/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_0/second_attend_to_encoder/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/second_attend_to_encoder/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/second_attend_to_encoder/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/second_attend_to_encoder/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/second_attend_to_encoder/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/dense/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/dense/bias:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/layer_prepostprocess_1/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/layer_prepostprocess_1/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/dense_1/kernel:0",
+        "evolved_transformer/body/decoder/layer_0/dense_layers/dense_1/bias:0",
+        "evolved_transformer/body/decoder/layer_1/16_head_self_attention/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_1/16_head_self_attention/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_1/16_head_self_attention/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/16_head_self_attention/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/16_head_self_attention/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/16_head_self_attention/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/first_attend_to_encoder/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/first_attend_to_encoder/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/first_attend_to_encoder/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/first_attend_to_encoder/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv11x1/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv11x1/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv11x1/bias:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv_7x1_1/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv_7x1_1/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv_7x1_1/bias:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv_7x1_2/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv_7x1_2/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_1/conv_branches/separable_conv_7x1_2/bias:0",
+        "evolved_transformer/body/decoder/layer_1/self_attention/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_1/self_attention/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_1/self_attention/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/self_attention/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/self_attention/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/self_attention/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/second_attend_to_encoder/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_1/second_attend_to_encoder/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_1/second_attend_to_encoder/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/second_attend_to_encoder/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/second_attend_to_encoder/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/second_attend_to_encoder/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/dense/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/dense/bias:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/layer_prepostprocess_1/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/layer_prepostprocess_1/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/dense_1/kernel:0",
+        "evolved_transformer/body/decoder/layer_1/dense_layers/dense_1/bias:0",
+        "evolved_transformer/body/decoder/layer_2/16_head_self_attention/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_2/16_head_self_attention/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_2/16_head_self_attention/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/16_head_self_attention/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/16_head_self_attention/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/16_head_self_attention/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/first_attend_to_encoder/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/first_attend_to_encoder/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/first_attend_to_encoder/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/first_attend_to_encoder/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv11x1/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv11x1/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv11x1/bias:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv_7x1_1/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv_7x1_1/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv_7x1_1/bias:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/layer_prepostprocess_1/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv_7x1_2/depthwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv_7x1_2/pointwise_kernel:0",
+        "evolved_transformer/body/decoder/layer_2/conv_branches/separable_conv_7x1_2/bias:0",
+        "evolved_transformer/body/decoder/layer_2/self_attention/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_2/self_attention/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_2/self_attention/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/self_attention/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/self_attention/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/self_attention/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/second_attend_to_encoder/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_2/second_attend_to_encoder/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_2/second_attend_to_encoder/multihead_attention/q/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/second_attend_to_encoder/multihead_attention/k/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/second_attend_to_encoder/multihead_attention/v/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/second_attend_to_encoder/multihead_attention/output_transform/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/dense/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/dense/bias:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/layer_prepostprocess_1/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/layer_prepostprocess_1/layer_norm/layer_norm_bias:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/dense_1/kernel:0",
+        "evolved_transformer/body/decoder/layer_2/dense_layers/dense_1/bias:0",
+        "evolved_transformer/body/decoder/layer_prepostprocess/layer_norm/layer_norm_scale:0",
+        "evolved_transformer/body/decoder/layer_prepostprocess/layer_norm/layer_norm_bias:0",
+    ]
+    variables = get_vars(var_names)
+    print_vars()
+
+    # Act.
+    with self.test_session() as session:
+      tf.global_variables_initializer().run()
+      values_before = session.run(variables)
+      for _ in range(10):  # Arbitrary number of training steps.
+        apply_grad.run()
+      values_after = session.run(variables)
+
+    # Assert.
+    self.assertTrue(
+        model._original_hparams.shared_embedding_and_softmax_weights)
+    self.assertTrue(model.hparams.shared_embedding_and_softmax_weights)
+    self.assertFalse(model.hparams.shared_embedding)
+    self.assertSameElements(var_names,
+                            [var.name for var in tf.trainable_variables()])
+    empty_vars = {
+        "evolved_transformer/symbol_modality_10_4/shared/weights_10:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_11:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_12:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_13:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_14:0",
+        "evolved_transformer/symbol_modality_10_4/shared/weights_15:0"
+    }
+    for name, before, after in zip(var_names, values_before, values_after):
+      if name in empty_vars:
+        self.assertEqual(before.size, after.size)
+        self.assertEqual(before.size, 0)
+      else:
+        assert_with_message(
+            self.assertNotAllClose, before, after,
+            "%s should be trainable, but did not change after training." % name)
+
 
 if __name__ == "__main__":
   tf.test.main()

From 9791182d9997c660374b8df6868185a45a92cfcc Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 10 Sep 2019 17:03:45 -0700
Subject: [PATCH 2414/2720] Randomize rotations and add multi-hash
 call_and_grad

PiperOrigin-RevId: 268343039
---
 tensor2tensor/trax/layers/attention.py | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index b8e3a990e..37c2cb5b3 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -601,9 +601,11 @@ def hash_vectors(self, vecs, rng):
       return self.bin_vectors_by_time(vecs)
 
     # See https://arxiv.org/pdf/1509.02897.pdf
+    # It's not clear whether sampling a different random rotation for each head
+    # and batch element matters here, but see MergedMultiHashedCausalAttention.
     assert self.n_bins % 2 == 0
     random_rotation = jax.random.normal(
-        rng, (vecs.shape[-1], self.n_bins//2)).astype('float32')
+        rng, (vecs.shape[0], vecs.shape[-1], self.n_bins//2)).astype('float32')
 
     # TODO(kitaev): making the vectors unit-length here is probably redundant.
     vecs = self.make_unit_length(vecs)
@@ -732,8 +734,6 @@ def binned_attn_vjp(sqk, sv, so_ct):  # pylint: disable=invalid-name
 
 class MergedMultiHashedCausalAttention(BaseCausalAttention):
   """Hash-based causal attention, with multiple hashes."""
-  # TODO(kitaev): Adapt this layer for use in a reversible network. At the
-  # moment that isn't supported because there's no call_and_grad implementation.
 
   def __init__(self, dropout, mode, n_bins=64, n_hashes=1, bin_by_time=False):
     del dropout, mode
@@ -765,9 +765,13 @@ def hash_vectors(self, vecs, rng):
       return self.bin_vectors_by_time(vecs)
 
     # See https://arxiv.org/pdf/1509.02897.pdf
+    # We sample a different random rotation for each batch element, head, and
+    # (crucially) each round of hashing. All of these are part of dimension 0
+    # of vecs. Applying multiple hashes to the same input is important because
+    # it increases the probability of being in the same bin as relevant items.
     assert self.n_bins % 2 == 0
     random_rotation = jax.random.normal(
-        rng, (vecs.shape[-1], self.n_bins//2)).astype('float32')
+        rng, (vecs.shape[0], vecs.shape[-1], self.n_bins//2)).astype('float32')
 
     # TODO(kitaev): making the vectors unit-length here is probably redundant.
     vecs = self.make_unit_length(vecs)
@@ -918,6 +922,16 @@ def vjpfun(grad):
 
     return out, state
 
+  def call_and_grad(self, inputs, ct, rng=None, **kwargs):
+    # TODO(kitaev): is there a manual implementation of call_and_grad that's
+    # faster than having jax infer one? Or are the permute/unpermute custom
+    # gradients defined in call() sufficient for reasonable speed?
+    def _do_call(x):
+      return self.call(x, params=(), state=(), rng=rng, **kwargs)[0]
+
+    output, vjpfun = jax.vjp(_do_call, inputs)
+    return output, vjpfun(ct)[0]
+
 
 def CausalAttention(d_feature, n_heads=1,
                     d_attention_key=None, d_attention_value=None,

From 9b5f169417d9f746ea84136fa35fac52e606ca90 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 10 Sep 2019 17:20:58 -0700
Subject: [PATCH 2415/2720] [async-rl CL 5/n] : Update
 `BatchTrajectory.load_from_directory` and add tests.  - `load_from_directory`
 can now wait for the required number of trajectories.

PiperOrigin-RevId: 268346090
---
 tensor2tensor/envs/trajectory.py      | 133 +++++++++++++++++++-------
 tensor2tensor/envs/trajectory_test.py |  97 +++++++++++++++++--
 2 files changed, 187 insertions(+), 43 deletions(-)

diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index d7622e4bf..27c8b406b 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -22,16 +22,24 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
 import os
+import pickle
 import re
-import cloudpickle as pickle
+import sys
+import time
+from absl import logging
+import cloudpickle
 import numpy as np
 from tensor2tensor.envs import time_step
 from tensorflow.io import gfile
 
-TRAJECTORY_FILE_REGEXP = r"trajectory_epoch_(\d{6})_env_id_(\d{6})_r_(.*).pkl"
-TRAJECTORY_FILE_GLOB = r"trajectory_epoch_*_env_id_*_r_*.pkl"
+TRAJECTORY_FILE_FORMAT = r"trajectory_epoch_{epoch}_env_id_{env_id}_temperature_{temperature}_r_{r}.pkl"
+
+
+def _get_pickle_module():
+  if sys.version_info[0] < 3:
+    return cloudpickle
+  return pickle
 
 
 class Trajectory(object):
@@ -455,60 +463,111 @@ def padding_config(obs):
   def parse_trajectory_file_name(trajectory_file_name):
     """Parse out the trajectory file's groups and return to caller."""
     base_trajectory_file_name = os.path.basename(trajectory_file_name)
-    compiled_regexp = re.compile(TRAJECTORY_FILE_REGEXP)
+    trajectory_file_regexp = TRAJECTORY_FILE_FORMAT.format(
+        epoch="(.*)",
+        env_id="(.*)",
+        temperature="(.*)",
+        r="(.*)",
+    )
+    compiled_regexp = re.compile(trajectory_file_regexp)
     r = compiled_regexp.match(base_trajectory_file_name)
     if not r:
       return None
     g = r.groups()
     if len(g) is not compiled_regexp.groups:
       return None
-    return (int(g[0]), int(g[1])) + g[2:]
+    # epoch, env_id, temp, random string
+    try:
+      epoch = int(g[0])
+      env_id = int(g[1])
+      temperature = float(g[2])
+      random_string = g[3]
+    except ValueError:
+      logging.error("Trajectory file name isn't parseable: %s",
+                    base_trajectory_file_name)
+      return None
+    return epoch, env_id, temperature, random_string
 
-  # TODO(afrozm): Test this function.
   @staticmethod
-  def load_from_directory(trajectory_dir, epoch=None, n_trajectories=None):
-    """Load trajectories from specified dir and epoch."""
+  def load_from_directory(trajectory_dir,
+                          epoch=None,
+                          temperature=None,
+                          n_trajectories=None,
+                          up_sample=False,
+                          wait_time_secs=10,
+                          max_tries=10):
+    """Load trajectories from specified dir and epoch.
+
+    Args:
+      trajectory_dir: (string) directory to find trajectories.
+      epoch: (int) epoch for which to load trajectories, if None we don't filter
+        on an epoch.
+      temperature: (float) this is used to filter the trajectory files, if None
+        we don't filter on temperature.
+      n_trajectories: (int) This is the batch size of the returned
+        BatchTrajectory object if one is returned. If set to None, then the
+        number of trajectories becomes the batch size. If set to some number,
+        then we wait for those many trajectory files to be available.
+      up_sample: (bool) If there are fewer than required (n_trajectories) number
+        of incomplete trajectories, then we upsample to make up the numbers.
+      wait_time_secs: (float) Waiting time, with exponential backoff to wait for
+        min_trajectories.
+      max_tries: (int) The number of tries to get min_trajectories trajectories.
 
-    trajectory_file_glob = TRAJECTORY_FILE_GLOB
+    Returns:
+      A BatchTrajectory object with all the constraints satisfied or None.
+    """
 
-    # If there is a desired epoch, modify the glob to get that instead.
-    if epoch:
-      trajectory_file_glob = trajectory_file_glob.replace(
-          "epoch_*", "epoch_%06d" % epoch)
+    # Modify the format to get a glob with desired epoch and temperature.
+    trajectory_file_glob = TRAJECTORY_FILE_FORMAT.format(
+        epoch=epoch if epoch is not None else "*",
+        env_id="*",
+        temperature=temperature if temperature is not None else "*",
+        r="*",
+    )
 
     trajectory_files = gfile.glob(
         os.path.join(trajectory_dir, trajectory_file_glob))
 
-    if not trajectory_files:
-      return None
+    if n_trajectories:
+      # We need to get `n_trajectories` number of `trajectory_files`.
+      # This works out to a maximum ~3hr waiting period.
+      while max_tries > 0 and len(trajectory_files) < n_trajectories:
+        logging.info(
+            "Sleeping for %s seconds while waiting for %s trajectories, found "
+            "%s right now.", wait_time_secs, n_trajectories,
+            len(trajectory_files))
+        time.sleep(wait_time_secs)
+        max_tries -= 1
+        wait_time_secs *= 2  # exponential backoff.
+        trajectory_files = gfile.glob(
+            os.path.join(trajectory_dir, trajectory_file_glob))
+
+      # We can't get the required number of files and we can't up-sample either.
+      if (len(trajectory_files) < n_trajectories) and not up_sample:
+        return None
+
+      # Sample up or down as the case maybe.
+      trajectory_files = list(
+          np.random.choice(trajectory_files, n_trajectories))
 
     # We read and load all the files, revisit if this becomes a problem.
     trajectories_buffer = []
-    completed_trajectories_buffer = []
     for trajectory_file in trajectory_files:
       with gfile.GFile(trajectory_file, "rb") as f:
-        list_trajectories = pickle.load(f)
-        assert isinstance(list_trajectories, list)
-        if not list_trajectories:
-          continue
-        assert isinstance(list_trajectories[0], Trajectory)
-        for trajectory in list_trajectories:
-          if trajectory.done:
-            completed_trajectories_buffer.append(trajectory)
-          else:
-            trajectories_buffer.append(trajectory)
-
-    if not trajectories_buffer and not completed_trajectories_buffer:
+        trajectory = _get_pickle_module().load(f)
+        assert isinstance(trajectory, Trajectory)
+        trajectories_buffer.append(trajectory)
+
+    if not trajectories_buffer:
       return None
 
-    # Randomly sample `n_trajectories` if needed.
-    n_trajectories = None if not n_trajectories else int(n_trajectories)
-    if n_trajectories and n_trajectories > 0:
-      trajectories_buffer = list(
-          np.random.choice(trajectories_buffer, int(trajectories_buffer)))
+    # If n_trajectories wasn't set, then set to the number of trajectories we're
+    # returning.
+    n_trajectories = n_trajectories or len(trajectories_buffer)
 
     # Construct and return a new BatchTrajectory object.
     return BatchTrajectory(
-        batch_size=len(trajectories_buffer),
-        trajectories=trajectories_buffer,
-        completed_trajectories=completed_trajectories_buffer)
+        batch_size=n_trajectories,
+        trajectories=[Trajectory() for _ in range(n_trajectories)],
+        completed_trajectories=trajectories_buffer)
diff --git a/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index 129bab8b3..40ceab7f9 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -19,9 +19,12 @@
 from __future__ import division
 from __future__ import print_function
 
+import os
 import numpy as np
+from tensor2tensor.envs import time_step
 from tensor2tensor.envs import trajectory
 import tensorflow as tf
+from tensorflow.io import gfile
 
 
 class TrajectoryTest(tf.test.TestCase):
@@ -138,15 +141,18 @@ def test_as_numpy(self):
     observations = np.random.uniform(size=(ts,) + shape)
     actions = np.random.choice(range(num_actions), size=(ts - 1,))
     rewards = np.random.choice([-1, 0, 1], size=(ts - 1,))
-    squares = np.arange(ts - 1) ** 2
-    cubes = np.arange(ts - 1) ** 3
+    squares = np.arange(ts - 1)**2
+    cubes = np.arange(ts - 1)**3
 
     def get_info(i):
       return {"sq": squares[i], "cu": cubes[i]}
 
     # First time-step has no reward.
-    t.add_time_step(observation=observations[0], done=False, action=actions[0],
-                    info=get_info(0))
+    t.add_time_step(
+        observation=observations[0],
+        done=False,
+        action=actions[0],
+        info=get_info(0))
     for i in range(1, ts - 1):
       t.add_time_step(
           observation=observations[i],
@@ -469,13 +475,92 @@ def test_observations_np(self):
 
   def test_parse_trajectory_file_name(self):
     self.assertEqual(
-        (12, 13, "abc"),
+        (12, 13, 1.0, "abc"),
         trajectory.BatchTrajectory.parse_trajectory_file_name(
-            "/tmp/trajectory_epoch_000012_env_id_000013_r_abc.pkl"))
+            "/tmp/trajectory_epoch_000012_env_id_000013_temperature_1.0_r_abc.pkl"
+        ))
 
     self.assertIsNone(
         trajectory.BatchTrajectory.parse_trajectory_file_name(
             "/tmp/trajectory_epoch_000012_env_id_000013.pkl"))
 
+  def test_load_from_directory(self):
+    output_dir = self.get_temp_dir()
+
+    epochs = [0, 1, 2]
+    env_ids = [0, 1, 2]
+    temperatures = [0.5, 1.0]
+    random_strings = ["a", "b"]
+
+    # Write some trajectories.
+    # There are 3x3x2x2 (36) trajectories, and of them 3x2x2 (12) are done.
+    for epoch in epochs:
+      for env_id in env_ids:
+        for temperature in temperatures:
+          for random_string in random_strings:
+            traj = trajectory.Trajectory(time_steps=[
+                time_step.TimeStep(
+                    observation=epoch,
+                    done=(epoch == 0),
+                    raw_reward=1.0,
+                    processed_reward=1.0,
+                    action=env_id,
+                    info={})
+            ])
+
+            trajectory_file_name = trajectory.TRAJECTORY_FILE_FORMAT.format(
+                epoch=epoch,
+                env_id=env_id,
+                temperature=temperature,
+                r=random_string)
+
+            with gfile.GFile(
+                os.path.join(output_dir, trajectory_file_name), "w") as f:
+              trajectory._get_pickle_module().dump(traj, f)
+
+    # Load everything and check.
+    bt = trajectory.BatchTrajectory.load_from_directory(output_dir)
+
+    self.assertIsInstance(bt, trajectory.BatchTrajectory)
+    self.assertEqual(36, bt.num_completed_trajectories)
+    self.assertEqual(36, bt.batch_size)
+
+    bt = trajectory.BatchTrajectory.load_from_directory(output_dir, epoch=0)
+    self.assertEqual(12, bt.num_completed_trajectories)
+    self.assertEqual(12, bt.batch_size)
+
+    # Get 100 trajectories, but there aren't any.
+    bt = trajectory.BatchTrajectory.load_from_directory(
+        output_dir, epoch=0, n_trajectories=100, max_tries=0)
+    self.assertIsNone(bt)
+
+    bt = trajectory.BatchTrajectory.load_from_directory(
+        output_dir, epoch=0, temperature=0.5)
+    self.assertEqual(6, bt.num_completed_trajectories)
+    self.assertEqual(6, bt.batch_size)
+
+    bt = trajectory.BatchTrajectory.load_from_directory(output_dir, epoch=1)
+    self.assertEqual(12, bt.num_completed_trajectories)
+    self.assertEqual(12, bt.batch_size)
+
+    # Constraints cannot be satisfied.
+    bt = trajectory.BatchTrajectory.load_from_directory(
+        output_dir, epoch=1, n_trajectories=100, up_sample=False, max_tries=0)
+    self.assertIsNone(bt)
+
+    # Constraints can be satisfied.
+    bt = trajectory.BatchTrajectory.load_from_directory(
+        output_dir, epoch=1, n_trajectories=100, up_sample=True, max_tries=0)
+    self.assertEqual(100, bt.num_completed_trajectories)
+    self.assertEqual(100, bt.batch_size)
+
+    bt = trajectory.BatchTrajectory.load_from_directory(
+        output_dir, epoch=1, n_trajectories=10)
+    self.assertEqual(10, bt.num_completed_trajectories)
+    self.assertEqual(10, bt.batch_size)
+
+    gfile.rmtree(output_dir)
+
+
 if __name__ == "__main__":
   tf.test.main()

From b02c5038d986944ad8edc50b6da80763f08fb7c3 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 10 Sep 2019 17:34:57 -0700
Subject: [PATCH 2416/2720] Move functions for making observations out of the
 history outside of OnlineTuneEnv and write tests

PiperOrigin-RevId: 268348239
---
 tensor2tensor/trax/rl/envs/online_tune.py     | 45 ++++++++++
 tensor2tensor/trax/rl/envs/online_tune_env.py | 40 +++++----
 .../trax/rl/envs/online_tune_test.py          | 88 +++++++++++++++++++
 3 files changed, 156 insertions(+), 17 deletions(-)
 create mode 100644 tensor2tensor/trax/rl/envs/online_tune.py
 create mode 100644 tensor2tensor/trax/rl/envs/online_tune_test.py

diff --git a/tensor2tensor/trax/rl/envs/online_tune.py b/tensor2tensor/trax/rl/envs/online_tune.py
new file mode 100644
index 000000000..8f72b6a3c
--- /dev/null
+++ b/tensor2tensor/trax/rl/envs/online_tune.py
@@ -0,0 +1,45 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Utility functions for OnlineTuneEnv."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+
+def historical_metric_values(history, metric, observation_range):
+  """Converts a metric stream from a trax History object into a numpy array."""
+  metric_sequence = history.get(*metric)
+  metric_values = np.array([
+      metric_value for (_, metric_value) in metric_sequence
+  ])
+  return np.clip(metric_values, *observation_range)
+
+
+def history_to_observations(history, metrics, observation_range, include_lr):
+  """Converts a trax History object into a sequence of observations."""
+  observation_dimensions = [
+      historical_metric_values(history, metric, observation_range)
+      for metric in metrics
+  ]
+  if include_lr:
+    # Logartihm of the learning rate.
+    observation_dimensions.append(np.log(historical_metric_values(
+        history, ("train", "training/learning_rate"), observation_range
+    )))
+  return np.stack(observation_dimensions, axis=1)
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index 70d6afbb0..cc181bfc2 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -19,15 +19,14 @@
 from __future__ import division
 from __future__ import print_function
 
-import math
 import os
 
 import gym
-import numpy as np
 
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
+from tensor2tensor.trax.rl.envs import online_tune
 from tensorflow.io import gfile
 
 
@@ -51,7 +50,7 @@ def __init__(self,
                output_dir,
                trainer_class=trax.Trainer,
                loss_fn=trax.loss,
-               optimizer=trax_opt.SM3,
+               optimizer=trax_opt.Adafactor,
                inputs=trax_inputs.inputs,
                action_multipliers=None,
                observation_metrics=(
@@ -131,20 +130,27 @@ def int_or_none(s):
 
     return os.path.join(self._output_dir, str(next_trajectory_id))
 
-  def _current_metric_value(self, metric):
-    metric_sequence = self._trainer.state.history.get(*metric)
-    assert metric_sequence
-    (_, metric_value) = metric_sequence[-1]
-    return np.clip(metric_value, *self._observation_range)
+  @property
+  def _current_reward_metric(self):
+    metric_values = online_tune.historical_metric_values(
+        self._trainer.state.history,
+        self._reward_metric,
+        self._observation_range,
+    )
+    assert metric_values.shape[0] > 0, (
+        "No values in history for metric {}.".format(self._reward_metric))
+    return metric_values[-1]
 
   @property
   def _current_observation(self):
-    observation = list(
-        map(self._current_metric_value, self._observation_metrics))
-    if self._include_lr_in_observation:
-      # Logartihm of the learning rate.
-      observation.append(math.log(self._current_lr))
-    return np.array(observation)
+    observations = online_tune.history_to_observations(
+        self._trainer.state.history,
+        self._observation_metrics,
+        self._observation_range,
+        self._include_lr_in_observation,
+    )
+    assert observations.shape[0] > 0, "No values in history for any metric."
+    return observations[-1, :]
 
   @property
   def trainer(self):
@@ -175,11 +181,11 @@ def step(self, action):
     """
     self._current_lr = min(
         self._current_lr * self._action_multipliers[action], self._max_lr)
-    last_metric_value = self._current_metric_value(self._reward_metric)
+    last_reward_metric = self._current_reward_metric
     self._trainer.train_epoch(self._train_steps, self._eval_steps)
     self._step += 1
-    current_metric_value = self._current_metric_value(self._reward_metric)
+    current_reward_metric = self._current_reward_metric
     observation = self._current_observation
-    reward = current_metric_value - last_metric_value
+    reward = current_reward_metric - last_reward_metric
     done = self._step == self._env_steps
     return (observation, reward, done, {})
diff --git a/tensor2tensor/trax/rl/envs/online_tune_test.py b/tensor2tensor/trax/rl/envs/online_tune_test.py
new file mode 100644
index 000000000..0f02215d8
--- /dev/null
+++ b/tensor2tensor/trax/rl/envs/online_tune_test.py
@@ -0,0 +1,88 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.rl.online_tune."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensor2tensor.trax import history as trax_history
+from tensor2tensor.trax.rl.envs import online_tune
+from tensorflow import test
+
+
+class OnlineTuneTest(test.TestCase):
+
+  def _append_metrics(self, h, metric, values):
+    for (i, value) in enumerate(values):
+      h.append(*metric, step=i, value=value)
+
+  def test_retrieves_historical_metric_values(self):
+    history = trax_history.History()
+    self._append_metrics(history, ("train", "accuracy"), [0.1, 0.73])
+    metric_values = online_tune.historical_metric_values(
+        history, metric=("train", "accuracy"), observation_range=(0, 5))
+    np.testing.assert_array_equal(metric_values, [0.1, 0.73])
+
+  def test_clips_historical_metric_values(self):
+    history = trax_history.History()
+    self._append_metrics(history, ("train", "loss"), [-10, 10])
+    metric_values = online_tune.historical_metric_values(
+        history, metric=("train", "loss"), observation_range=(-1, 1))
+    np.testing.assert_array_equal(metric_values, [-1, 1])
+
+  def test_converts_history_to_observations_without_learning_rate(self):
+    history = trax_history.History()
+    self._append_metrics(history, ("train", "loss"), [3.0, 1.07])
+    self._append_metrics(history, ("eval", "accuracy"), [0.12, 0.68])
+    observations = online_tune.history_to_observations(
+        history,
+        metrics=(("eval", "accuracy"), ("train", "loss")),
+        observation_range=(0, 5),
+        include_lr=False,
+    )
+    np.testing.assert_array_equal(observations, [[0.12, 3.0], [0.68, 1.07]])
+
+  def test_converts_history_to_observations_with_learning_rate(self):
+    history = trax_history.History()
+    self._append_metrics(
+        history, ("train", "training/learning_rate"), [1e-3, 1e-4])
+    observations = online_tune.history_to_observations(
+        history,
+        metrics=(),
+        observation_range=(0, 5),
+        include_lr=True,
+    )
+    self.assertEqual(observations.shape, (2, 1))
+    ((log_lr_1,), (log_lr_2,)) = observations
+    self.assertGreater(log_lr_1, log_lr_2)
+
+  def test_clips_observations(self):
+    history = trax_history.History()
+    self._append_metrics(history, ("eval", "loss"), [-10, 10])
+    observations = online_tune.history_to_observations(
+        history,
+        metrics=(("eval", "loss"),),
+        observation_range=(-2, 2),
+        include_lr=False,
+    )
+    np.testing.assert_array_equal(observations, [[-2], [2]])
+
+
+if __name__ == "__main__":
+  test.main()

From ac49451685995dad311e87f52de8dc12fe819782 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 10 Sep 2019 17:46:26 -0700
Subject: [PATCH 2417/2720] Decrease the initial wait_time_secs to 0.1 and bump
 up max_tries to 17 to maintain total ~3 hour wait time.

PiperOrigin-RevId: 268350140
---
 tensor2tensor/envs/trajectory.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index 27c8b406b..e3352433d 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -494,8 +494,8 @@ def load_from_directory(trajectory_dir,
                           temperature=None,
                           n_trajectories=None,
                           up_sample=False,
-                          wait_time_secs=10,
-                          max_tries=10):
+                          wait_time_secs=0.1,
+                          max_tries=17):
     """Load trajectories from specified dir and epoch.
 
     Args:

From fa19052fe6327205ee0fbe487a416d0de68d198b Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 10 Sep 2019 18:23:46 -0700
Subject: [PATCH 2418/2720] [async-rl CL 6/n] : Implement ppo_trainer's
 collect_trajectories_async.  - Its mainly a thin wrapper over
 `BatchTrajectory.load_from_directory`  - I believe this should work from the
 perspective of the actual RL trainer.  - Now we need the server that is going
 to collect the data.

PiperOrigin-RevId: 268356027
---
 tensor2tensor/envs/env_problem_utils.py | 12 +++++--
 tensor2tensor/trax/rl/base_trainer.py   |  6 +---
 tensor2tensor/trax/rl/ppo_trainer.py    | 48 +++++++++++++++++++++----
 tensor2tensor/trax/rl/simple_trainer.py |  1 -
 4 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 32d81b436..9f15483de 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -52,6 +52,13 @@ def play_env_problem_randomly(env_problem, num_steps):
     env_problem.reset(indices=done_indices(dones))
 
 
+def get_completed_trajectories_from_env(env, n_trajectories):
+  completed_trajectories = []
+  for trajectory in env.trajectories.completed_trajectories[:n_trajectories]:
+    completed_trajectories.append(trajectory.as_numpy)
+  return completed_trajectories
+
+
 def play_env_problem_with_policy(env,
                                  policy_fun,
                                  num_trajectories=1,
@@ -183,9 +190,8 @@ def gumbel_sample(log_probs):
 
   # We have the trajectories we need, return a list of triples:
   # (observations, actions, rewards)
-  completed_trajectories = []
-  for trajectory in env.trajectories.completed_trajectories[:num_trajectories]:
-    completed_trajectories.append(trajectory.as_numpy)
+  completed_trajectories = get_completed_trajectories_from_env(
+      env, num_trajectories)
 
   timing_info = {
       "trajectory_collection/policy_application": policy_application_total_time,
diff --git a/tensor2tensor/trax/rl/base_trainer.py b/tensor2tensor/trax/rl/base_trainer.py
index 6f71cbe97..2d1d55b99 100644
--- a/tensor2tensor/trax/rl/base_trainer.py
+++ b/tensor2tensor/trax/rl/base_trainer.py
@@ -32,7 +32,7 @@ class BaseTrainer(object):
   def __init__(
       self, train_env, eval_env, output_dir,
       trajectory_dump_dir=None, trajectory_dump_min_count_per_shard=16,
-      async_mode=False, async_mode_trajectory_subdir="trajectories",
+      async_mode=False,
   ):
     """Base class constructor.
 
@@ -47,8 +47,6 @@ def __init__(
         shuffling for model training in SimPLe.
       async_mode: (bool) If True, this means we are in async mode and we read
         trajectories from a location rather than interact with the environment.
-      async_mode_trajectory_subdir: (string) The subdir of output_dir to search
-        for trajectories in async mode.
     """
     self.train_env = train_env
     self.eval_env = eval_env
@@ -58,9 +56,7 @@ def __init__(
     self._trajectory_dump_min_count_per_shard = (
         trajectory_dump_min_count_per_shard)
     self._trajectory_buffer = []
-
     self._async_mode = async_mode
-    self._async_mode_trajectory_subdir = async_mode_trajectory_subdir
 
   @property
   def epoch(self):
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index 19c1eaf06..bc9acc1b2 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -31,6 +31,8 @@
 from jax import numpy as np
 from jax import random as jax_random
 import numpy as onp
+from tensor2tensor.envs import env_problem_utils
+from tensor2tensor.envs import trajectory
 from tensor2tensor.trax import jaxboard
 from tensor2tensor.trax import models as trax_models
 from tensor2tensor.trax import optimizers as trax_opt
@@ -215,14 +217,45 @@ def assert_same_space(space1, space2):
   def epoch(self):
     return self._epoch
 
-  def collect_trajectories_async(self, env, train=True, n_trajectories=1):
+  def collect_trajectories_async(self,
+                                 env,
+                                 train=True,
+                                 n_trajectories=1,
+                                 temperature=1.0):
     """Collects trajectories in an async manner."""
 
     assert self._async_mode
-    del env
-    del train
-    del n_trajectories
-    raise NotImplementedError
+
+    trajectory_dir = os.path.join(self._output_dir, "trajectories",
+                                  "train" if train else "eval")
+    epoch = self.epoch
+
+    logging.info(
+        "Loading [%s] trajectories from dir [%s] for epoch [%s] and temperature"
+        " [%s]", n_trajectories, trajectory_dir, epoch, temperature)
+
+    bt = trajectory.BatchTrajectory.load_from_directory(
+        trajectory_dir,
+        epoch=epoch,
+        temperature=temperature,
+        n_trajectories=n_trajectories)
+
+    if bt is None:
+      logging.error(
+          "Couldn't load [%s] trajectories from dir [%s] for epoch [%s] and "
+          "temperature [%s]", n_trajectories, trajectory_dir, epoch,
+          temperature)
+      assert bt
+
+    # Doing this is important, since we want to modify `env` so that it looks
+    # like `env` was actually played and the trajectories came from it.
+    env.trajectories = bt
+
+    trajs = env_problem_utils.get_completed_trajectories_from_env(
+        env, n_trajectories)
+    n_done = len(trajs)
+    timing_info = {}
+    return trajs, n_done, timing_info, self._model_state
 
   def collect_trajectories(self, train=True, temperature=1.0):
     self._rng, key = jax_random.split(self._rng)
@@ -240,7 +273,10 @@ def collect_trajectories(self, train=True, temperature=1.0):
     # If async, read the required trajectories for the epoch.
     if self._async_mode:
       return self.collect_trajectories_async(
-          env, train=train, n_trajectories=n_trajectories)
+          env,
+          train=train,
+          n_trajectories=n_trajectories,
+          temperature=temperature)
 
     trajs, n_done, timing_info, self._model_state = ppo.collect_trajectories(
         env,
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index c0513f265..edd371b26 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -61,7 +61,6 @@ def __init__(self,
         eval_env=eval_env,
         output_dir=self._policy_dir,
         async_mode=self._async_mode,
-        async_mode_trajectory_subdir=self._async_mode_trajectory_subdir,
     )
     self._n_real_epochs = n_real_epochs
     self._model_train_batch_size = model_train_batch_size

From cf13ad30aee8ed2478cac3af74c4e544dd2646fe Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 10 Sep 2019 18:42:53 -0700
Subject: [PATCH 2419/2720] Add learning rate update to OnlineTune helper
 functions

PiperOrigin-RevId: 268358416
---
 tensor2tensor/trax/rl/envs/online_tune.py     | 16 +++++++++++--
 tensor2tensor/trax/rl/envs/online_tune_env.py |  8 +++++--
 .../trax/rl/envs/online_tune_env_test.py      |  7 ++++++
 .../trax/rl/envs/online_tune_test.py          | 23 +++++++++++++++++++
 4 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/trax/rl/envs/online_tune.py b/tensor2tensor/trax/rl/envs/online_tune.py
index 8f72b6a3c..a4631edba 100644
--- a/tensor2tensor/trax/rl/envs/online_tune.py
+++ b/tensor2tensor/trax/rl/envs/online_tune.py
@@ -22,7 +22,11 @@
 import numpy as np
 
 
-def historical_metric_values(history, metric, observation_range):
+LEARNING_RATE_METRIC = ("train", "training/learning_rate")
+
+
+def historical_metric_values(
+    history, metric, observation_range=(-np.inf, np.inf)):
   """Converts a metric stream from a trax History object into a numpy array."""
   metric_sequence = history.get(*metric)
   metric_values = np.array([
@@ -40,6 +44,14 @@ def history_to_observations(history, metrics, observation_range, include_lr):
   if include_lr:
     # Logartihm of the learning rate.
     observation_dimensions.append(np.log(historical_metric_values(
-        history, ("train", "training/learning_rate"), observation_range
+        history, LEARNING_RATE_METRIC, observation_range
     )))
   return np.stack(observation_dimensions, axis=1)
+
+
+def new_learning_rate(action, history, action_multipliers, max_lr):
+  """Calculates a new learning rate based on an action."""
+  learning_rates = historical_metric_values(history, LEARNING_RATE_METRIC)
+  assert learning_rates.shape[0] > 0, "No last learning rate found in history."
+  current_lr = learning_rates[-1]
+  return min(current_lr * action_multipliers[action], max_lr)
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index cc181bfc2..eb00d0e49 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -179,8 +179,12 @@ def step(self, action):
         metric since the last step. done is set after reaching self.env_steps
         environment steps. info is an empty dict.
     """
-    self._current_lr = min(
-        self._current_lr * self._action_multipliers[action], self._max_lr)
+    self._current_lr = online_tune.new_learning_rate(
+        action,
+        self._trainer.state.history,
+        self._action_multipliers,
+        self._max_lr,
+    )
     last_reward_metric = self._current_reward_metric
     self._trainer.train_epoch(self._train_steps, self._eval_steps)
     self._step += 1
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env_test.py b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
index df9e1ab69..2fb7f84e1 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env_test.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import models
 from tensor2tensor.trax import trax
+from tensor2tensor.trax.rl.envs import online_tune
 from tensor2tensor.trax.rl.envs import online_tune_env
 from tensorflow import test
 from tensorflow.io import gfile
@@ -59,6 +60,12 @@ def evaluate(self, eval_steps):
         metric=METRIC,
         step=self.step,
         value=self.metrics_to_report.pop(0))
+    (lr_mode, lr_metric) = online_tune.LEARNING_RATE_METRIC
+    self.state.history.append(
+        mode=lr_mode,
+        metric=lr_metric,
+        step=self.step,
+        value=self.learning_rate)
 
 
 class OnlineTuneTest(test.TestCase):
diff --git a/tensor2tensor/trax/rl/envs/online_tune_test.py b/tensor2tensor/trax/rl/envs/online_tune_test.py
index 0f02215d8..69b09a836 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_test.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_test.py
@@ -83,6 +83,29 @@ def test_clips_observations(self):
     )
     np.testing.assert_array_equal(observations, [[-2], [2]])
 
+  def test_calculates_new_learning_rate(self):
+    history = trax_history.History()
+    self._append_metrics(
+        history, online_tune.LEARNING_RATE_METRIC, [1e-2, 1e-3])
+    new_lr = online_tune.new_learning_rate(
+        action=2,
+        history=history,
+        action_multipliers=(0.5, 1.0, 2.0),
+        max_lr=1.0,
+    )
+    np.testing.assert_almost_equal(new_lr, 2e-3)
+
+  def test_clips_new_learning_rate(self):
+    history = trax_history.History()
+    self._append_metrics(history, online_tune.LEARNING_RATE_METRIC, [1e-3])
+    new_lr = online_tune.new_learning_rate(
+        action=0,
+        history=history,
+        action_multipliers=(4.0, 1.0, 0.25),
+        max_lr=3e-3,
+    )
+    np.testing.assert_almost_equal(new_lr, 3e-3)
+
 
 if __name__ == "__main__":
   test.main()

From 7d5700082d303623b445566ab910cef235bb0a69 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 10 Sep 2019 21:01:09 -0700
Subject: [PATCH 2420/2720] Move policy saving to PPO helpers.

This will be needed to test PolicySchedule.

PiperOrigin-RevId: 268374694
---
 tensor2tensor/trax/rl/ppo.py         | 17 +++++++++++++++++
 tensor2tensor/trax/rl/ppo_test.py    | 11 +++++++++++
 tensor2tensor/trax/rl/ppo_trainer.py | 18 +++++++-----------
 3 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index e9753326b..641517b2e 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -826,6 +826,23 @@ def maybe_restore_opt_state(output_dir,
   )
 
 
+def save_opt_state(output_dir,
+                   policy_and_value_opt_state,
+                   policy_and_value_state,
+                   epoch,
+                   total_opt_step):
+  """Saves the policy and value network optimization state etc."""
+  old_model_files = get_policy_model_files(output_dir)
+  params_file = os.path.join(output_dir, "model-%06d.pkl" % epoch)
+  with gfile.GFile(params_file, "wb") as f:
+    pickle.dump(
+        (policy_and_value_opt_state, policy_and_value_state, total_opt_step), f)
+  # Remove the old model files.
+  for path in old_model_files:
+    if path != params_file:
+      gfile.remove(path)
+
+
 def write_eval_reward_summaries(reward_stats_by_mode, summary_writer, epoch):
   """Writes evaluation reward statistics to summary and logs them.
 
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
index c250a401a..3be53bb4f 100644
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -537,5 +537,16 @@ def plp(p):
                     -np.sum(filtered_log_probs) / 5.0,
                     1e-6)
 
+  def test_saves_and_restores_opt_state(self):
+    opt_state = 123
+    state = 456
+    epoch = 7
+    opt_step = 89
+    output_dir = self.get_temp_dir()
+    ppo.save_opt_state(output_dir, opt_state, state, epoch, opt_step)
+    restored_data = ppo.maybe_restore_opt_state(output_dir)
+    self.assertEqual(restored_data, (opt_state, state, epoch, opt_step))
+
+
 if __name__ == "__main__":
   test.main()
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index bc9acc1b2..3a4e62f5b 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -25,7 +25,6 @@
 import time
 
 from absl import logging
-import cloudpickle as pickle
 import gym
 from jax import jit
 from jax import numpy as np
@@ -39,7 +38,6 @@
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.rl import base_trainer
 from tensor2tensor.trax.rl import ppo
-from tensorflow.io import gfile
 
 DEBUG_LOGGING = False
 GAMMA = 0.99
@@ -619,15 +617,13 @@ def compute_stats(reward_dict):
   def save(self):
     """Save the agent parameters."""
     logging.vlog(1, "PPO epoch [% 6d]: saving model.", self._epoch)
-    old_model_files = ppo.get_policy_model_files(self._output_dir)
-    params_file = os.path.join(self._output_dir, "model-%06d.pkl" % self._epoch)
-    with gfile.GFile(params_file, "wb") as f:
-      pickle.dump((self._policy_and_value_opt_state, self._model_state,
-                   self._total_opt_step), f)
-    # Remove the old model files.
-    for path in old_model_files:
-      if path != params_file:
-        gfile.remove(path)
+    ppo.save_opt_state(
+        self._output_dir,
+        self._policy_and_value_opt_state,
+        self._model_state,
+        self._epoch,
+        self._total_opt_step,
+    )
     # Reset this number.
     self._n_trajectories_done = 0
     self._last_saved_at = self._epoch

From 9db59c5cba0c847b0e4e76d174f19147e044a0c1 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 10 Sep 2019 22:06:17 -0700
Subject: [PATCH 2421/2720] Automatically download and prepare TFDS and T2T
 datasets in Trax when data_dir is unspecified.

PiperOrigin-RevId: 268382292
---
 tensor2tensor/trax/inputs.py | 36 ++++++++++++++++++++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 634a8a07b..325121824 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -61,6 +61,39 @@
 _MAX_SKIP_EXAMPLES = 1e5
 
 
+def download_and_prepare(dataset_name, data_dir):
+  """Downloads and prepares T2T or TFDS dataset.
+
+  Args:
+    dataset_name: tfds dataset or t2t problem name prefixed by "t2t_".
+    data_dir: location of existing dataset or None.
+
+  Returns:
+    data_dir: path string of downloaded data.
+  """
+  if not data_dir:
+    data_dir = os.path.expanduser('~/tensorflow_datasets/')
+    dl_dir = os.path.join(data_dir, 'download')
+    tf.logging.info(
+        ('No dataset directory provided. '
+         'Downloading and generating dataset for %s inside data directory %s '
+         'For large datasets it is better to prepare datasets manually!')
+        % (dataset_name, data_dir))
+    if dataset_name.startswith('t2t_'):
+      # Download and run dataset generator for T2T problem.
+      data_dir = os.path.join(data_dir, dataset_name)
+      tf.gfile.MakeDirs(data_dir)
+      tf.gfile.MakeDirs(dl_dir)
+      t2t_problems.problem(dataset_name[4:]).generate_data(data_dir, dl_dir)
+    else:
+      # Download and prepare TFDS dataset.
+      tfds_builder = tfds.builder(dataset_name)
+      tfds_builder.download_and_prepare(download_dir=dl_dir)
+  else:
+    data_dir = os.path.expanduser(data_dir)
+  return data_dir
+
+
 @gin.configurable(blacklist=['n_devices'])
 def inputs(n_devices, dataset_name, data_dir=None, input_name=None,
            n_chunks=0):
@@ -77,8 +110,7 @@ def inputs(n_devices, dataset_name, data_dir=None, input_name=None,
   Returns:
     trax.inputs.Inputs
   """
-  assert data_dir, 'Must provide a data directory'
-  data_dir = os.path.expanduser(data_dir)
+  data_dir = download_and_prepare(dataset_name, data_dir)
 
   (train_batches, train_eval_batches, eval_batches,
    input_name, input_shape, input_dtype,

From 14eb9ca43b6a378b5bc4cbe510edc4f5f2067abd Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 10 Sep 2019 23:45:54 -0700
Subject: [PATCH 2422/2720] Small readability fix in download_and_prepare
 routine.

PiperOrigin-RevId: 268394334
---
 tensor2tensor/trax/inputs.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index 325121824..c18150001 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -84,7 +84,8 @@ def download_and_prepare(dataset_name, data_dir):
       data_dir = os.path.join(data_dir, dataset_name)
       tf.gfile.MakeDirs(data_dir)
       tf.gfile.MakeDirs(dl_dir)
-      t2t_problems.problem(dataset_name[4:]).generate_data(data_dir, dl_dir)
+      t2t_problems.problem(
+          dataset_name[len('t2t_'):]).generate_data(data_dir, dl_dir)
     else:
       # Download and prepare TFDS dataset.
       tfds_builder = tfds.builder(dataset_name)

From 8276312fed74963dd607b7d623849f4779afcbc5 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 11 Sep 2019 12:17:11 -0700
Subject: [PATCH 2423/2720] Put gumbel_sample() and a function for getting the
 pickle module to a new module, trax.utils.

Now trax.rl works on Python 3.

PiperOrigin-RevId: 268514255
---
 tensor2tensor/trax/rl/base_trainer.py         |  7 +--
 tensor2tensor/trax/rl/ppo.py                  |  8 ++--
 tensor2tensor/trax/rl/simple.py               |  5 ++-
 tensor2tensor/trax/rl/simple_test.py          |  5 ++-
 .../trax/rl/simulated_env_problem.py          |  8 +---
 tensor2tensor/trax/trax.py                    | 13 ++----
 tensor2tensor/trax/utils.py                   | 43 +++++++++++++++++++
 7 files changed, 64 insertions(+), 25 deletions(-)
 create mode 100644 tensor2tensor/trax/utils.py

diff --git a/tensor2tensor/trax/rl/base_trainer.py b/tensor2tensor/trax/rl/base_trainer.py
index 2d1d55b99..622f106a7 100644
--- a/tensor2tensor/trax/rl/base_trainer.py
+++ b/tensor2tensor/trax/rl/base_trainer.py
@@ -22,7 +22,7 @@
 import os
 
 from absl import logging
-import cloudpickle as pickle
+from tensor2tensor.trax import utils
 from tensorflow.io import gfile
 
 
@@ -83,6 +83,7 @@ def dump_trajectories(self, force=False):
       force: (bool) Whether to complete unfinished trajectories and create
         a new shard even if we have not reached the minimum size.
     """
+    pkl_module = utils.get_pickle_module()
     if self.trajectory_dump_dir is None:
       return
     gfile.makedirs(self.trajectory_dump_dir)
@@ -113,9 +114,9 @@ def has_any_action(trajectory):
         # sometimes dump 2 times in the same epoch. When this happens, merge the
         # two sets of trajectories.
         with gfile.GFile(shard_path, "rb") as f:
-          self._trajectory_buffer = pickle.load(f) + self._trajectory_buffer
+          self._trajectory_buffer = pkl_module.load(f) + self._trajectory_buffer
       with gfile.GFile(shard_path, "wb") as f:
-        pickle.dump(self._trajectory_buffer, f)
+        pkl_module.dump(self._trajectory_buffer, f)
       self._trajectory_buffer = []
 
   def training_loop(self, n_epochs, evaluate=True):
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index 641517b2e..469625aa0 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -56,7 +56,6 @@
 import time
 
 from absl import logging
-import cloudpickle as pickle
 from jax import grad
 from jax import jit
 from jax import lax
@@ -66,6 +65,7 @@
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.trax import layers as tl
+from tensor2tensor.trax import utils
 from tensorflow.io import gfile
 
 
@@ -804,6 +804,7 @@ def maybe_restore_opt_state(output_dir,
     found, and opt_step is the total optimization step (sum of all optimization
     steps made up to the current epoch).
   """
+  pkl_module = utils.get_pickle_module()
   epoch = 0
   total_opt_step = 0
   for model_file in get_policy_model_files(output_dir):
@@ -811,7 +812,7 @@ def maybe_restore_opt_state(output_dir,
     try:
       with gfile.GFile(model_file, "rb") as f:
         policy_and_value_opt_state, policy_and_value_state, total_opt_step = (
-            pickle.load(f))
+            pkl_module.load(f))
       epoch = get_epoch_from_policy_model_file(model_file)
       break
     except EOFError as e:
@@ -832,10 +833,11 @@ def save_opt_state(output_dir,
                    epoch,
                    total_opt_step):
   """Saves the policy and value network optimization state etc."""
+  pkl_module = utils.get_pickle_module()
   old_model_files = get_policy_model_files(output_dir)
   params_file = os.path.join(output_dir, "model-%06d.pkl" % epoch)
   with gfile.GFile(params_file, "wb") as f:
-    pickle.dump(
+    pkl_module.dump(
         (policy_and_value_opt_state, policy_and_value_state, total_opt_step), f)
   # Remove the old model files.
   for path in old_model_files:
diff --git a/tensor2tensor/trax/rl/simple.py b/tensor2tensor/trax/rl/simple.py
index d5f025bf7..57feaa2b9 100644
--- a/tensor2tensor/trax/rl/simple.py
+++ b/tensor2tensor/trax/rl/simple.py
@@ -24,15 +24,16 @@
 import random
 
 from absl import logging
-import cloudpickle as pickle
 import numpy as np
 from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.envs import trajectory
+from tensor2tensor.trax import utils
 from tensorflow.io import gfile
 
 
 def load_trajectories(trajectory_dir, eval_frac):
   """Loads trajectories from a possibly nested directory of pickles."""
+  pkl_module = utils.get_pickle_module()
   train_trajectories = []
   eval_trajectories = []
   # Search the entire directory subtree for trajectories.
@@ -40,7 +41,7 @@ def load_trajectories(trajectory_dir, eval_frac):
     for filename in filenames:
       shard_path = os.path.join(subdir, filename)
       with gfile.GFile(shard_path, "rb") as f:
-        trajectories = pickle.load(f)
+        trajectories = pkl_module.load(f)
         pivot = int(len(trajectories) * (1 - eval_frac))
         train_trajectories.extend(trajectories[:pivot])
         eval_trajectories.extend(trajectories[pivot:])
diff --git a/tensor2tensor/trax/rl/simple_test.py b/tensor2tensor/trax/rl/simple_test.py
index 8d5efea2d..957fcf6cf 100644
--- a/tensor2tensor/trax/rl/simple_test.py
+++ b/tensor2tensor/trax/rl/simple_test.py
@@ -22,7 +22,6 @@
 import itertools
 import os
 
-import cloudpickle as pickle
 import gin
 import gym
 from matplotlib import pyplot as plt
@@ -32,6 +31,7 @@
 from tensor2tensor.envs import trajectory
 from tensor2tensor.trax import backend
 from tensor2tensor.trax import trax
+from tensor2tensor.trax import utils
 from tensor2tensor.trax.rl import simple
 from tensor2tensor.trax.rl import simulated_env_problem
 from tensor2tensor.trax.rl import space_serializer  # pylint: disable=unused-import
@@ -47,9 +47,10 @@ def _make_singleton_trajectory(self, observation):
     return t
 
   def _dump_trajectory_pickle(self, observations, path):
+    pkl_module = utils.get_pickle_module()
     trajectories = list(map(self._make_singleton_trajectory, observations))
     with gfile.GFile(path, "wb") as f:
-      pickle.dump(trajectories, f)
+      pkl_module.dump(trajectories, f)
 
   def test_loads_trajectories(self):
     temp_dir = self.get_temp_dir()
diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index 818fc42f6..454f9c865 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -27,6 +27,7 @@
 from tensor2tensor.envs import env_problem
 from tensor2tensor.trax import backend
 from tensor2tensor.trax import trax
+from tensor2tensor.trax import utils
 from tensor2tensor.trax.backend import random as jax_random
 from tensor2tensor.trax.rl import space_serializer
 
@@ -384,18 +385,13 @@ def _action_repr_indices(self):
     return index_range_2d(begin_indices, self._action_repr_length)
 
   def _predict_obs(self, predict_fn, rng):
-    def gumbel_sample(log_probs):
-      u = np.random.uniform(low=1e-6, high=1.0 - 1e-6, size=log_probs.shape)
-      g = -np.log(-np.log(u))
-      return np.argmax(log_probs + g, axis=-1)
-
     for (i, subrng) in enumerate(jax_random.split(rng, self._obs_repr_length)):
       symbol_index = self._steps * self._step_repr_length + i
       log_probs, self._model_state = predict_fn(self._history,
                                                 state=self._model_state,
                                                 rng=subrng)
       log_probs = log_probs[:, symbol_index, :]
-      self._history[:, symbol_index] = gumbel_sample(log_probs)
+      self._history[:, symbol_index] = utils.gumbel_sample(log_probs)
 
     obs_repr = self._history[self._obs_repr_indices]
     return self._obs_serializer.deserialize(obs_repr)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 8d946a90e..1cf009140 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -23,13 +23,11 @@
 import functools
 import itertools
 import os
-import pickle
 import random
 import sys
 import time
 
 from absl import logging
-import cloudpickle
 
 import gin
 
@@ -45,6 +43,7 @@
 from tensor2tensor.trax import layers
 from tensor2tensor.trax import learning_rate as lr
 from tensor2tensor.trax import optimizers as trax_opt
+from tensor2tensor.trax import utils
 from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.backend import random as jax_random
 
@@ -189,8 +188,9 @@ def restore_state(output_dir):
     return State(step=None, opt_state=None, history=trax_history.History(),
                  model_state=None)
 
+  pkl_module = utils.get_pickle_module()
   with gfile.GFile(params_file, "rb") as f:
-    (opt_state, step, history, model_state) = pickle.load(f)
+    (opt_state, step, history, model_state) = pkl_module.load(f)
   log("Model loaded from %s at step %d" % (params_file, step))
   logging.debug("From loaded model : history = %s", history)
   return State(step=step, opt_state=OptState(*opt_state), history=history,
@@ -209,12 +209,7 @@ def _save_gin(output_dir, sw=None):
 
 def save_state(state, output_dir, keep=False):
   """Save State and optionally gin config."""
-  # TODO(gilmer, lukaszkaiser): figure out how to use cloudpickle in python3.
-  # Currently the code throws an error when run in python3.
-  if sys.version_info[0] < 3:
-    pkl_module = cloudpickle
-  else:
-    pkl_module = pickle
+  pkl_module = utils.get_pickle_module()
   params_file = os.path.join(output_dir, "model.pkl")
   with gfile.GFile(params_file, "wb") as f:
     pkl_module.dump((tuple(state.opt_state), state.step, state.history,
diff --git a/tensor2tensor/trax/utils.py b/tensor2tensor/trax/utils.py
new file mode 100644
index 000000000..902e73ba4
--- /dev/null
+++ b/tensor2tensor/trax/utils.py
@@ -0,0 +1,43 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Utility functions."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import pickle
+import sys
+
+import cloudpickle
+import numpy as np
+
+
+def get_pickle_module():
+  """Returns the appropriate pickle module based on Python version."""
+  # TODO(gilmer, lukaszkaiser): figure out how to use cloudpickle in python3.
+  # Currently the code throws an error when run in python3.
+  if sys.version_info[0] < 3:
+    return cloudpickle
+  else:
+    return pickle
+
+
+def gumbel_sample(log_probs):
+  """Gumbel sampling from a categorical distribution."""
+  u = np.random.uniform(low=1e-6, high=1.0 - 1e-6, size=log_probs.shape)
+  g = -np.log(-np.log(u))
+  return np.argmax(log_probs + g, axis=-1)

From 35639ad87cdc6d8aaab778b895af7bdb422a8a04 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 11 Sep 2019 15:04:46 -0700
Subject: [PATCH 2424/2720] Implement a policy-controlled LR schedule.

PiperOrigin-RevId: 268551358
---
 tensor2tensor/trax/learning_rate.py           |  99 ++++++++++++++++
 tensor2tensor/trax/learning_rate_test.py      |  87 ++++++++++++++
 tensor2tensor/trax/rl/envs/online_tune_env.py |   2 +-
 .../trax/rl/envs/online_tune_env_test.py      |   2 +-
 tensor2tensor/trax/rl/online_tune.py          |  57 +++++++++
 tensor2tensor/trax/rl/online_tune_test.py     | 111 ++++++++++++++++++
 6 files changed, 356 insertions(+), 2 deletions(-)
 create mode 100644 tensor2tensor/trax/learning_rate_test.py
 create mode 100644 tensor2tensor/trax/rl/online_tune.py
 create mode 100644 tensor2tensor/trax/rl/online_tune_test.py

diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index 176e363f7..2201837ba 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -26,8 +26,18 @@
 from __future__ import division
 from __future__ import print_function
 
+import random
+import time
+
+from absl import logging
 import gin
+
+from tensor2tensor.trax import models as trax_models
+from tensor2tensor.trax import utils
 from tensor2tensor.trax.backend import numpy as np
+from tensor2tensor.trax.backend import random as jax_random
+from tensor2tensor.trax.rl import online_tune
+from tensor2tensor.trax.rl import ppo
 
 
 @gin.configurable(blacklist=["history"])
@@ -128,3 +138,92 @@ def EvalAdjustingSchedule(history,
       steps_without_improvement = 0
 
   return MultifactorSchedule(history, constant=adjusted)
+
+
+@gin.configurable(blacklist=["history"])
+def PolicySchedule(
+    history,
+    observation_metrics=(
+        ("train", "metrics/accuracy"),
+        ("train", "metrics/loss"),
+        ("eval", "metrics/accuracy"),
+        ("eval", "metrics/loss"),
+    ),
+    include_lr_in_observation=False,
+    observation_range=(0.0, 5.0),
+    start_lr=0.001,
+    max_lr=10.0,
+    action_multipliers=(1.0 / 1.5, 1.0 / 1.25, 1.0, 1.25, 1.5),
+    policy_and_value_model=trax_models.FrameStackMLP,
+    policy_and_value_two_towers=False,
+    policy_dir=gin.REQUIRED,
+):
+  """Learning rate schedule controlled by a learned policy.
+
+  Args:
+    history: the history of training and evaluation (History object).
+    observation_metrics: list of pairs (mode, metric), as in the History object.
+    include_lr_in_observation: bool, whether to include the learning rate in
+      observations.
+    observation_range: tuple (low, high), range to clip the observation to.
+    start_lr: starting learning rate.
+    max_lr: maximum value to clip the learning rate to.
+    action_multipliers: sequence of LR multipliers that policy actions
+      correspond to.
+    policy_and_value_model: Trax model to use as the policy.
+    policy_and_value_two_towers: bool, whether the action distribution and value
+      prediction is computed by separate model towers.
+    policy_dir: directory with the policy checkpoint.
+
+  Returns:
+    a function learning_rate(step): float -> float, the step-dependent lr.
+  """
+
+  # Turn the history into observations for the policy. If we don't have any,
+  # return the initial learning rate.
+  start_time = time.time()
+  observations = online_tune.history_to_observations(
+      history, observation_metrics, observation_range, include_lr_in_observation
+  )
+  logging.vlog(
+      1, "Building observations took %0.2f sec.", time.time() - start_time)
+  if observations.shape[0] == 0:
+    return lambda _: start_lr
+
+  # Build the policy network and load its parameters.
+  start_time = time.time()
+  net = ppo.policy_and_value_net(
+      n_actions=len(action_multipliers),
+      bottom_layers_fn=policy_and_value_model,
+      two_towers=policy_and_value_two_towers,
+  )
+  logging.vlog(
+      1, "Building the policy network took %0.2f sec.", time.time() - start_time
+  )
+  start_time = time.time()
+  # (opt_state, state, epoch, opt_step)
+  (opt_state, state, _, _) = ppo.maybe_restore_opt_state(policy_dir)
+  assert opt_state is not None, "Policy checkpoint not found."
+  (params, _) = opt_state
+  logging.vlog(
+      1, "Restoring the policy parameters took %0.2f sec.",
+      time.time() - start_time
+  )
+
+  # Run the policy and sample an action.
+  seed = random.randint(0, 2**31 - 1)
+  rng = jax_random.get_prng(seed=seed)
+  start_time = time.time()
+  # ((log_probs, value_preds), state). We have no way to pass state to the next
+  # step, but that should be fine.
+  ((log_probs, _), _) = net(np.array([observations]), params, state, rng=rng)
+  logging.vlog(
+      1, "Running the policy took %0.2f sec.", time.time() - start_time
+  )
+  # Sample from the action distribution for the last timestep.
+  action = utils.gumbel_sample(log_probs[0, -1, :])
+
+  # Get a new learning rate.
+  new_lr = online_tune.new_learning_rate(
+      action, history, action_multipliers, max_lr)
+  return lambda _: new_lr
diff --git a/tensor2tensor/trax/learning_rate_test.py b/tensor2tensor/trax/learning_rate_test.py
new file mode 100644
index 000000000..6ba477016
--- /dev/null
+++ b/tensor2tensor/trax/learning_rate_test.py
@@ -0,0 +1,87 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.learning_rate."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as onp
+
+from tensor2tensor.trax import history as trax_history
+from tensor2tensor.trax import learning_rate
+from tensor2tensor.trax.backend import numpy as np
+from tensor2tensor.trax.backend import random as jax_random
+from tensor2tensor.trax.models import atari_cnn
+from tensor2tensor.trax.rl import online_tune
+from tensor2tensor.trax.rl import ppo
+from tensorflow import test
+
+
+class PolicyScheduleTest(test.TestCase):
+
+  def _make_schedule(
+      self,
+      history,
+      start_lr=1e-3,
+      observation_metrics=(("eval", "metrics/accuracy"),),
+      action_multipliers=(1.0,),
+  ):
+    policy_and_value_model = atari_cnn.FrameStackMLP
+    net = ppo.policy_and_value_net(
+        n_actions=len(action_multipliers),
+        bottom_layers_fn=policy_and_value_model,
+        two_towers=False,
+    )
+    rng = jax_random.get_prng(seed=0)
+    obs_dim = len(observation_metrics)
+    (params, state) = net.initialize((1, 1, obs_dim), np.float32, rng)
+    policy_dir = self.get_temp_dir()
+    # Optimizer slots should not be used for anything.
+    slots = None
+    opt_state = (params, slots)
+    ppo.save_opt_state(policy_dir, opt_state, state, epoch=0, total_opt_step=0)
+    return learning_rate.PolicySchedule(
+        history,
+        observation_metrics=observation_metrics,
+        include_lr_in_observation=False,
+        action_multipliers=action_multipliers,
+        start_lr=start_lr,
+        policy_and_value_model=policy_and_value_model,
+        policy_and_value_two_towers=False,
+        policy_dir=policy_dir,
+    )
+
+  def test_returns_start_lr_when_there_are_no_metrics(self):
+    history = trax_history.History()
+    schedule = self._make_schedule(history, start_lr=1e-3)
+    self.assertEqual(schedule(0), 1e-3)
+
+  def test_changes_lr_when_there_are_some_metrics(self):
+    history = trax_history.History()
+    history.append("eval", "metrics/accuracy", step=0, value=0.8)
+    history.append(*online_tune.LEARNING_RATE_METRIC, step=0, value=1e-4)
+    schedule = self._make_schedule(
+        history,
+        observation_metrics=(("eval", "metrics/accuracy"),),
+        action_multipliers=(0.5, 2.0),
+    )
+    self.assertTrue(
+        onp.allclose(schedule(123), 5e-5) or onp.allclose(schedule(123), 2e-4))
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index eb00d0e49..3cf7203d4 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -26,7 +26,7 @@
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
-from tensor2tensor.trax.rl.envs import online_tune
+from tensor2tensor.trax.rl import online_tune
 from tensorflow.io import gfile
 
 
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env_test.py b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
index 2fb7f84e1..64e981d82 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env_test.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
@@ -26,7 +26,7 @@
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import models
 from tensor2tensor.trax import trax
-from tensor2tensor.trax.rl.envs import online_tune
+from tensor2tensor.trax.rl import online_tune
 from tensor2tensor.trax.rl.envs import online_tune_env
 from tensorflow import test
 from tensorflow.io import gfile
diff --git a/tensor2tensor/trax/rl/online_tune.py b/tensor2tensor/trax/rl/online_tune.py
new file mode 100644
index 000000000..a4631edba
--- /dev/null
+++ b/tensor2tensor/trax/rl/online_tune.py
@@ -0,0 +1,57 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Utility functions for OnlineTuneEnv."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+
+LEARNING_RATE_METRIC = ("train", "training/learning_rate")
+
+
+def historical_metric_values(
+    history, metric, observation_range=(-np.inf, np.inf)):
+  """Converts a metric stream from a trax History object into a numpy array."""
+  metric_sequence = history.get(*metric)
+  metric_values = np.array([
+      metric_value for (_, metric_value) in metric_sequence
+  ])
+  return np.clip(metric_values, *observation_range)
+
+
+def history_to_observations(history, metrics, observation_range, include_lr):
+  """Converts a trax History object into a sequence of observations."""
+  observation_dimensions = [
+      historical_metric_values(history, metric, observation_range)
+      for metric in metrics
+  ]
+  if include_lr:
+    # Logartihm of the learning rate.
+    observation_dimensions.append(np.log(historical_metric_values(
+        history, LEARNING_RATE_METRIC, observation_range
+    )))
+  return np.stack(observation_dimensions, axis=1)
+
+
+def new_learning_rate(action, history, action_multipliers, max_lr):
+  """Calculates a new learning rate based on an action."""
+  learning_rates = historical_metric_values(history, LEARNING_RATE_METRIC)
+  assert learning_rates.shape[0] > 0, "No last learning rate found in history."
+  current_lr = learning_rates[-1]
+  return min(current_lr * action_multipliers[action], max_lr)
diff --git a/tensor2tensor/trax/rl/online_tune_test.py b/tensor2tensor/trax/rl/online_tune_test.py
new file mode 100644
index 000000000..93486383c
--- /dev/null
+++ b/tensor2tensor/trax/rl/online_tune_test.py
@@ -0,0 +1,111 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.online_tune."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensor2tensor.trax import history as trax_history
+from tensor2tensor.trax.rl import online_tune
+from tensorflow import test
+
+
+class OnlineTuneTest(test.TestCase):
+
+  def _append_metrics(self, h, metric, values):
+    for (i, value) in enumerate(values):
+      h.append(*metric, step=i, value=value)
+
+  def test_retrieves_historical_metric_values(self):
+    history = trax_history.History()
+    self._append_metrics(history, ("train", "accuracy"), [0.1, 0.73])
+    metric_values = online_tune.historical_metric_values(
+        history, metric=("train", "accuracy"), observation_range=(0, 5))
+    np.testing.assert_array_equal(metric_values, [0.1, 0.73])
+
+  def test_clips_historical_metric_values(self):
+    history = trax_history.History()
+    self._append_metrics(history, ("train", "loss"), [-10, 10])
+    metric_values = online_tune.historical_metric_values(
+        history, metric=("train", "loss"), observation_range=(-1, 1))
+    np.testing.assert_array_equal(metric_values, [-1, 1])
+
+  def test_converts_history_to_observations_without_learning_rate(self):
+    history = trax_history.History()
+    self._append_metrics(history, ("train", "loss"), [3.0, 1.07])
+    self._append_metrics(history, ("eval", "accuracy"), [0.12, 0.68])
+    observations = online_tune.history_to_observations(
+        history,
+        metrics=(("eval", "accuracy"), ("train", "loss")),
+        observation_range=(0, 5),
+        include_lr=False,
+    )
+    np.testing.assert_array_equal(observations, [[0.12, 3.0], [0.68, 1.07]])
+
+  def test_converts_history_to_observations_with_learning_rate(self):
+    history = trax_history.History()
+    self._append_metrics(
+        history, ("train", "training/learning_rate"), [1e-3, 1e-4])
+    observations = online_tune.history_to_observations(
+        history,
+        metrics=(),
+        observation_range=(0, 5),
+        include_lr=True,
+    )
+    self.assertEqual(observations.shape, (2, 1))
+    ((log_lr_1,), (log_lr_2,)) = observations
+    self.assertGreater(log_lr_1, log_lr_2)
+
+  def test_clips_observations(self):
+    history = trax_history.History()
+    self._append_metrics(history, ("eval", "loss"), [-10, 10])
+    observations = online_tune.history_to_observations(
+        history,
+        metrics=(("eval", "loss"),),
+        observation_range=(-2, 2),
+        include_lr=False,
+    )
+    np.testing.assert_array_equal(observations, [[-2], [2]])
+
+  def test_calculates_new_learning_rate(self):
+    history = trax_history.History()
+    self._append_metrics(
+        history, online_tune.LEARNING_RATE_METRIC, [1e-2, 1e-3])
+    new_lr = online_tune.new_learning_rate(
+        action=2,
+        history=history,
+        action_multipliers=(0.5, 1.0, 2.0),
+        max_lr=1.0,
+    )
+    np.testing.assert_almost_equal(new_lr, 2e-3)
+
+  def test_clips_new_learning_rate(self):
+    history = trax_history.History()
+    self._append_metrics(history, online_tune.LEARNING_RATE_METRIC, [1e-3])
+    new_lr = online_tune.new_learning_rate(
+        action=0,
+        history=history,
+        action_multipliers=(4.0, 1.0, 0.25),
+        max_lr=3e-3,
+    )
+    np.testing.assert_almost_equal(new_lr, 3e-3)
+
+
+if __name__ == "__main__":
+  test.main()

From c06cf3f40108b7be9928e0ab20e54623f057feaf Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 11 Sep 2019 16:24:50 -0700
Subject: [PATCH 2425/2720] Make Reformer config smaller and allow to use a
 single rng for all steps.

PiperOrigin-RevId: 268567430
---
 .../transformer_revnet_imagenet64_8gb.gin     | 29 ++++++++++++-------
 tensor2tensor/trax/layers/attention.py        | 29 ++++++++++++++++---
 2 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
index 25d609052..5d5ece224 100644
--- a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
@@ -5,8 +5,8 @@ import tensor2tensor.trax.trax
 
 # Parameters for batch_fun:
 # ==============================================================================
-batch_fun.batch_size_per_device = 8
-batch_fun.eval_batch_size = 128
+batch_fun.batch_size_per_device = 2
+batch_fun.eval_batch_size = 16
 batch_fun.max_eval_length = 12288  # 64 * 64 * 3
 
 # Parameters for inputs:
@@ -41,24 +41,33 @@ DotProductCausalAttention.dropout = 0.0
 MemoryEfficientCausalAttention.dropout = 0.0
 MemoryEfficientCausalAttention.loop_stride = 512
 
-# Parameters for DummyHashedAttention:
+# Parameters for MergedHashedCausalAttention:
 # ==============================================================================
-# DummyHashedAttention.dropout = 0.0
-# DummyHashedAttention.n_bins = 64
+MergedHashedCausalAttention.dropout = 0.0
+MergedHashedCausalAttention.n_bins = 32
+MergedHashedCausalAttention.bin_by_time = True
+MergedHashedCausalAttention.one_rng = False
+
+# Parameters for MergedMultiHashedCausalAttention:
+# ==============================================================================
+MergedMultiHashedCausalAttention.dropout = 0.0
+MergedMultiHashedCausalAttention.n_bins = 64
+MergedMultiHashedCausalAttention.n_hashes = 4
+MergedMultiHashedCausalAttention.bin_by_time = False
+MergedHashedCausalAttention.one_rng = True
 
 # Parameters for TransformerRevnetLM:
 # ==============================================================================
 TransformerRevnetLM.d_model = 1024
 TransformerRevnetLM.d_ff = 2048
-TransformerRevnetLM.d_attention_key = 32
-TransformerRevnetLM.d_attention_value = 32
+TransformerRevnetLM.d_attention_key = 64
+TransformerRevnetLM.d_attention_value = 64
 TransformerRevnetLM.dropout = 0.0
 TransformerRevnetLM.max_len = 12288  # 64 * 64 * 3
 TransformerRevnetLM.mode = 'train'
 TransformerRevnetLM.n_heads = 4
-TransformerRevnetLM.n_layers = 6
+TransformerRevnetLM.n_layers = 4
 TransformerRevnetLM.vocab_size = 256
 TransformerRevnetLM.n_chunks = 16
 TransformerRevnetLM.n_attention_chunks = 1
-TransformerRevnetLM.attention_type = @trax.layers.MemoryEfficientCausalAttention
-
+TransformerRevnetLM.attention_type = @trax.layers.MergedMultiHashedCausalAttention
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 37c2cb5b3..104edba32 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import random
 import jax
 import numpy as onp
 
@@ -559,11 +560,17 @@ def body_fun(vals):  # pylint: disable=invalid-name
 class MergedHashedCausalAttention(BaseCausalAttention):
   """Hash-based causal attention."""
 
-  def __init__(self, dropout, mode, n_bins=64, bin_by_time=False):
+  def __init__(self, dropout, mode, n_bins=64,
+               bin_by_time=False, one_rng=False):
     del dropout, mode
     super(MergedHashedCausalAttention, self).__init__()
     self.n_bins = n_bins
     self.bin_by_time = bin_by_time
+    seed = random.randint(0, 2**31 - 1)
+    self._one_rng = one_rng
+    self._prng = None
+    if one_rng:
+      self._prng = backend.random.get_prng(seed)
 
   def call(self, inputs, params=(), state=(), **kwargs):
     del params
@@ -604,8 +611,12 @@ def hash_vectors(self, vecs, rng):
     # It's not clear whether sampling a different random rotation for each head
     # and batch element matters here, but see MergedMultiHashedCausalAttention.
     assert self.n_bins % 2 == 0
+    rot_rng = rng
+    if self._one_rng:
+      rot_rng = jax.lax.tie_in(vecs, self._prng)
     random_rotation = jax.random.normal(
-        rng, (vecs.shape[0], vecs.shape[-1], self.n_bins//2)).astype('float32')
+        rot_rng,
+        (vecs.shape[0], vecs.shape[-1], self.n_bins//2)).astype('float32')
 
     # TODO(kitaev): making the vectors unit-length here is probably redundant.
     vecs = self.make_unit_length(vecs)
@@ -735,12 +746,18 @@ def binned_attn_vjp(sqk, sv, so_ct):  # pylint: disable=invalid-name
 class MergedMultiHashedCausalAttention(BaseCausalAttention):
   """Hash-based causal attention, with multiple hashes."""
 
-  def __init__(self, dropout, mode, n_bins=64, n_hashes=1, bin_by_time=False):
+  def __init__(self, dropout, mode, n_bins=64, n_hashes=1,
+               bin_by_time=False, one_rng=False):
     del dropout, mode
     super(MergedMultiHashedCausalAttention, self).__init__()
     self.n_bins = n_bins
     self.n_hashes = n_hashes
     self.bin_by_time = bin_by_time
+    seed = random.randint(0, 2**31 - 1)
+    self._one_rng = one_rng
+    self._prng = None
+    if one_rng:
+      self._prng = backend.random.get_prng(seed)
 
   def bin_vectors_by_time(self, vecs):
     seqlen = vecs.shape[-2]
@@ -770,8 +787,12 @@ def hash_vectors(self, vecs, rng):
     # of vecs. Applying multiple hashes to the same input is important because
     # it increases the probability of being in the same bin as relevant items.
     assert self.n_bins % 2 == 0
+    rot_rng = rng
+    if self._one_rng:
+      rot_rng = jax.lax.tie_in(vecs, self._prng)
     random_rotation = jax.random.normal(
-        rng, (vecs.shape[0], vecs.shape[-1], self.n_bins//2)).astype('float32')
+        rot_rng,
+        (vecs.shape[0], vecs.shape[-1], self.n_bins//2)).astype('float32')
 
     # TODO(kitaev): making the vectors unit-length here is probably redundant.
     vecs = self.make_unit_length(vecs)

From 77f18662431e019f066731911a8d8141d42ccea5 Mon Sep 17 00:00:00 2001
From: Dmitrii Murygin <dmitrij.murygin7@gmail.com>
Date: Thu, 12 Sep 2019 21:33:03 +0300
Subject: [PATCH 2426/2720] Update test file (#1698)

---
 tensor2tensor/data_generators/wikisum/utils_test.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/wikisum/utils_test.py b/tensor2tensor/data_generators/wikisum/utils_test.py
index 36397bae6..e713d1938 100644
--- a/tensor2tensor/data_generators/wikisum/utils_test.py
+++ b/tensor2tensor/data_generators/wikisum/utils_test.py
@@ -24,23 +24,24 @@
 
 import tensorflow as tf
 
-pkg_dir, _ = os.path.split(__file__)
+pkg_dir = os.path.abspath(__file__)
+pkg_dir, _ = os.path.split(pkg_dir)
 _TESTDATA = os.path.join(pkg_dir, "test_data")
 
 
 def _get_testdata(filename):
-  with tf.gfile.Open(os.path.join(_TESTDATA, filename)) as f:
+  with tf.io.gfile.GFile(filename) as f:
     return f.read()
 
 
 class UtilsTest(tf.test.TestCase):
 
   def test_filter_paragraph(self):
-    for bad in tf.gfile.Glob(os.path.join(_TESTDATA, "para_bad*.txt")):
+    for bad in tf.io.gfile.glob(os.path.join(_TESTDATA, "para_bad*.txt")):
       for p in _get_testdata(bad).split("\n"):
         self.assertTrue(utils.filter_paragraph(p),
                         msg="Didn't filter %s" % p)
-    for good in tf.gfile.Glob(os.path.join(_TESTDATA, "para_good*.txt")):
+    for good in tf.io.gfile.glob(os.path.join(_TESTDATA, "para_good*.txt")):
       for p in _get_testdata(good).split("\n"):
         p = _get_testdata(good)
       self.assertFalse(utils.filter_paragraph(p), msg="Filtered %s" % p)

From bade297a0860ab208ac87496613522278cbf6f5e Mon Sep 17 00:00:00 2001
From: Dmitrii Murygin <dmitrij.murygin7@gmail.com>
Date: Thu, 12 Sep 2019 21:33:34 +0300
Subject: [PATCH 2427/2720] Fixes three utils test files (#1699)

---
 tensor2tensor/utils/checkpoint_compatibility_test.py | 3 ++-
 tensor2tensor/utils/hparams_lib_test.py              | 3 ++-
 tensor2tensor/utils/trainer_lib_test.py              | 3 ++-
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/checkpoint_compatibility_test.py b/tensor2tensor/utils/checkpoint_compatibility_test.py
index 6657adb9c..759105b99 100644
--- a/tensor2tensor/utils/checkpoint_compatibility_test.py
+++ b/tensor2tensor/utils/checkpoint_compatibility_test.py
@@ -41,7 +41,8 @@
 
 
 def get_data_dir():
-  pkg, _ = os.path.split(__file__)
+  pkg = os.path.abspath(__file__)
+  pkg, _ = os.path.split(pkg)
   pkg, _ = os.path.split(pkg)
   return os.path.join(pkg, "test_data")
 
diff --git a/tensor2tensor/utils/hparams_lib_test.py b/tensor2tensor/utils/hparams_lib_test.py
index 651743806..589caee1b 100644
--- a/tensor2tensor/utils/hparams_lib_test.py
+++ b/tensor2tensor/utils/hparams_lib_test.py
@@ -30,7 +30,8 @@ class HparamsLibTest(tf.test.TestCase):
 
   def testCreateHparamsFromJson(self):
     # Get json_path
-    pkg, _ = os.path.split(__file__)
+    pkg = os.path.abspath(__file__)
+    pkg, _ = os.path.split(pkg)
     pkg, _ = os.path.split(pkg)
     json_path = os.path.join(
         pkg, "test_data", "transformer_test_ckpt", "hparams.json")
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index d45550fcf..0c7ffaeb4 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -147,7 +147,8 @@ def body(args, mb=model.body):
 
   def testCreateHparams(self):
     # Get json_path
-    pkg, _ = os.path.split(__file__)
+    pkg = os.path.abspath(__file__)
+    pkg, _ = os.path.split(pkg)
     pkg, _ = os.path.split(pkg)
     json_path = os.path.join(
         pkg, "test_data", "transformer_test_ckpt", "hparams.json")

From ebb54cc412d8652389058be1f5385580077f2233 Mon Sep 17 00:00:00 2001
From: Dmitrii Murygin <dmitrij.murygin7@gmail.com>
Date: Thu, 12 Sep 2019 11:33:20 -0700
Subject: [PATCH 2428/2720] Merge of PR #1698

PiperOrigin-RevId: 268729413
---
 tensor2tensor/utils/checkpoint_compatibility_test.py | 3 +--
 tensor2tensor/utils/hparams_lib_test.py              | 3 +--
 tensor2tensor/utils/trainer_lib_test.py              | 3 +--
 3 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/utils/checkpoint_compatibility_test.py b/tensor2tensor/utils/checkpoint_compatibility_test.py
index 759105b99..6657adb9c 100644
--- a/tensor2tensor/utils/checkpoint_compatibility_test.py
+++ b/tensor2tensor/utils/checkpoint_compatibility_test.py
@@ -41,8 +41,7 @@
 
 
 def get_data_dir():
-  pkg = os.path.abspath(__file__)
-  pkg, _ = os.path.split(pkg)
+  pkg, _ = os.path.split(__file__)
   pkg, _ = os.path.split(pkg)
   return os.path.join(pkg, "test_data")
 
diff --git a/tensor2tensor/utils/hparams_lib_test.py b/tensor2tensor/utils/hparams_lib_test.py
index 589caee1b..651743806 100644
--- a/tensor2tensor/utils/hparams_lib_test.py
+++ b/tensor2tensor/utils/hparams_lib_test.py
@@ -30,8 +30,7 @@ class HparamsLibTest(tf.test.TestCase):
 
   def testCreateHparamsFromJson(self):
     # Get json_path
-    pkg = os.path.abspath(__file__)
-    pkg, _ = os.path.split(pkg)
+    pkg, _ = os.path.split(__file__)
     pkg, _ = os.path.split(pkg)
     json_path = os.path.join(
         pkg, "test_data", "transformer_test_ckpt", "hparams.json")
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index 0c7ffaeb4..d45550fcf 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -147,8 +147,7 @@ def body(args, mb=model.body):
 
   def testCreateHparams(self):
     # Get json_path
-    pkg = os.path.abspath(__file__)
-    pkg, _ = os.path.split(pkg)
+    pkg, _ = os.path.split(__file__)
     pkg, _ = os.path.split(pkg)
     json_path = os.path.join(
         pkg, "test_data", "transformer_test_ckpt", "hparams.json")

From 63c9e981d019e95c84a3954347a278e91c5c7d22 Mon Sep 17 00:00:00 2001
From: Dmitrii Murygin <dmitrij.murygin7@gmail.com>
Date: Thu, 12 Sep 2019 11:33:53 -0700
Subject: [PATCH 2429/2720] Merge of PR #1699

PiperOrigin-RevId: 268729536
---
 tensor2tensor/utils/checkpoint_compatibility_test.py | 3 ++-
 tensor2tensor/utils/hparams_lib_test.py              | 3 ++-
 tensor2tensor/utils/trainer_lib_test.py              | 3 ++-
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/checkpoint_compatibility_test.py b/tensor2tensor/utils/checkpoint_compatibility_test.py
index 6657adb9c..759105b99 100644
--- a/tensor2tensor/utils/checkpoint_compatibility_test.py
+++ b/tensor2tensor/utils/checkpoint_compatibility_test.py
@@ -41,7 +41,8 @@
 
 
 def get_data_dir():
-  pkg, _ = os.path.split(__file__)
+  pkg = os.path.abspath(__file__)
+  pkg, _ = os.path.split(pkg)
   pkg, _ = os.path.split(pkg)
   return os.path.join(pkg, "test_data")
 
diff --git a/tensor2tensor/utils/hparams_lib_test.py b/tensor2tensor/utils/hparams_lib_test.py
index 651743806..589caee1b 100644
--- a/tensor2tensor/utils/hparams_lib_test.py
+++ b/tensor2tensor/utils/hparams_lib_test.py
@@ -30,7 +30,8 @@ class HparamsLibTest(tf.test.TestCase):
 
   def testCreateHparamsFromJson(self):
     # Get json_path
-    pkg, _ = os.path.split(__file__)
+    pkg = os.path.abspath(__file__)
+    pkg, _ = os.path.split(pkg)
     pkg, _ = os.path.split(pkg)
     json_path = os.path.join(
         pkg, "test_data", "transformer_test_ckpt", "hparams.json")
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index d45550fcf..0c7ffaeb4 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -147,7 +147,8 @@ def body(args, mb=model.body):
 
   def testCreateHparams(self):
     # Get json_path
-    pkg, _ = os.path.split(__file__)
+    pkg = os.path.abspath(__file__)
+    pkg, _ = os.path.split(pkg)
     pkg, _ = os.path.split(pkg)
     json_path = os.path.join(
         pkg, "test_data", "transformer_test_ckpt", "hparams.json")

From 85b5816786ec9c6f90df9db76b241fdfb2b4e72e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 12 Sep 2019 11:35:31 -0700
Subject: [PATCH 2430/2720] Update imagenet64 config and add enwik8 config

PiperOrigin-RevId: 268729887
---
 .../trax/configs/transformer_enwik8.gin       | 68 +++++++++++++++++++
 .../trax/configs/transformer_imagenet64.gin   |  7 +-
 2 files changed, 73 insertions(+), 2 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/transformer_enwik8.gin

diff --git a/tensor2tensor/trax/configs/transformer_enwik8.gin b/tensor2tensor/trax/configs/transformer_enwik8.gin
new file mode 100644
index 000000000..ed38015e1
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_enwik8.gin
@@ -0,0 +1,68 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 1
+batch_fun.eval_batch_size = 8
+batch_fun.max_eval_length = 65536
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_enwik8_l65k'
+inputs.input_name = 'targets'
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 2.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 100
+train.eval_steps = 8
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerLM
+train.optimizer = @trax.optimizers.Adafactor
+train.train_steps = 200000
+train.trainer_class = @MemoryEfficientTrainer
+
+# Parameters for MemoryEfficientCausalAttention:
+# ==============================================================================
+MemoryEfficientCausalAttention.dropout = 0.0
+MemoryEfficientCausalAttention.loop_stride = 256
+
+# Parameters for MergedHashedCausalAttention:
+# ==============================================================================
+MergedHashedCausalAttention.dropout = 0.0
+MergedHashedCausalAttention.n_bins = 512
+MergedHashedCausalAttention.bin_by_time = True
+MergedMultiHashedCausalAttention.one_rng = False
+
+# Parameters for MergedMultiHashedCausalAttention:
+# ==============================================================================
+MergedMultiHashedCausalAttention.dropout = 0.0
+MergedMultiHashedCausalAttention.n_bins = 512
+MergedMultiHashedCausalAttention.n_hashes = 2
+MergedMultiHashedCausalAttention.bin_by_time = False
+MergedMultiHashedCausalAttention.one_rng = False
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.attention_type = @trax.layers.MergedMultiHashedCausalAttention
+TransformerLM.d_attention_key = 64
+TransformerLM.d_attention_value = 64
+TransformerLM.d_model = 512
+TransformerLM.d_ff = 2048
+TransformerLM.dropout = 0.0
+TransformerLM.max_len = 65536
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 4
+TransformerLM.n_layers = 3
+TransformerLM.share_kv = True
+TransformerLM.vocab_size = 258  # Includes pad token and unused EOS token
+
diff --git a/tensor2tensor/trax/configs/transformer_imagenet64.gin b/tensor2tensor/trax/configs/transformer_imagenet64.gin
index 096de4739..ca64be8fb 100644
--- a/tensor2tensor/trax/configs/transformer_imagenet64.gin
+++ b/tensor2tensor/trax/configs/transformer_imagenet64.gin
@@ -41,19 +41,21 @@ MemoryEfficientCausalAttention.loop_stride = 512
 MergedHashedCausalAttention.dropout = 0.0
 MergedHashedCausalAttention.n_bins = 16
 MergedHashedCausalAttention.bin_by_time = True
+MergedMultiHashedCausalAttention.one_rng = False
 
 # Parameters for MergedMultiHashedCausalAttention:
 # ==============================================================================
 MergedMultiHashedCausalAttention.dropout = 0.0
 MergedMultiHashedCausalAttention.n_bins = 64
-MergedMultiHashedCausalAttention.n_hashes = 4
+MergedMultiHashedCausalAttention.n_hashes = 2
 MergedMultiHashedCausalAttention.bin_by_time = False
+MergedMultiHashedCausalAttention.one_rng = False
 
 # Parameters for TransformerLM:
 # ==============================================================================
 TransformerLM.attention_type = @trax.layers.MergedMultiHashedCausalAttention
 TransformerLM.d_attention_key = 64
-TransformerLM.d_attention_value = 256
+TransformerLM.d_attention_value = 64
 TransformerLM.d_model = 1024
 TransformerLM.d_ff = 2048
 TransformerLM.dropout = 0.0
@@ -61,5 +63,6 @@ TransformerLM.max_len = 12288  # 64 * 64 * 3
 TransformerLM.mode = 'train'
 TransformerLM.n_heads = 4
 TransformerLM.n_layers = 3
+TransformerLM.share_kv = True
 TransformerLM.vocab_size = 256
 

From f1a2cfb41e51a886cd9e0f244a89b7e79ee554f4 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 12 Sep 2019 17:24:04 -0700
Subject: [PATCH 2431/2720] Set BatchNorm momentum in WideResNet to 0.9

PiperOrigin-RevId: 268802491
---
 tensor2tensor/trax/models/resnet.py | 29 ++++++++++++++++++-----------
 1 file changed, 18 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
index 3551e10a5..a9c644cc9 100644
--- a/tensor2tensor/trax/models/resnet.py
+++ b/tensor2tensor/trax/models/resnet.py
@@ -111,37 +111,41 @@ def Resnet50(d_hidden=64, n_output_classes=1001, mode='train'):
   )
 
 
-def WideResnetBlock(channels, strides=(1, 1), mode='train'):
+def WideResnetBlock(channels, strides=(1, 1), bn_momentum=0.9, mode='train'):
   """WideResnet convolutional block."""
   return [
-      tl.BatchNorm(mode=mode),
+      tl.BatchNorm(momentum=bn_momentum, mode=mode),
       tl.Relu(),
       tl.Conv(channels, (3, 3), strides, padding='SAME'),
-      tl.BatchNorm(mode=mode),
+      tl.BatchNorm(momentum=bn_momentum, mode=mode),
       tl.Relu(),
       tl.Conv(channels, (3, 3), padding='SAME'),
   ]
 
 
-def WideResnetGroup(n, channels, strides=(1, 1), mode='train'):
+def WideResnetGroup(n, channels, strides=(1, 1), bn_momentum=0.9, mode='train'):
   shortcut = [
       tl.Conv(channels, (3, 3), strides, padding='SAME'),
   ]
   return [
-      tl.Residual(WideResnetBlock(channels, strides, mode=mode),
+      tl.Residual(WideResnetBlock(channels, strides, bn_momentum=bn_momentum,
+                                  mode=mode),
                   shortcut=shortcut),
-      tl.Residual([WideResnetBlock(channels, (1, 1), mode=mode)
+      tl.Residual([WideResnetBlock(channels, (1, 1), bn_momentum=bn_momentum,
+                                   mode=mode)
                    for _ in range(n - 1)]),
   ]
 
 
-def WideResnet(n_blocks=3, widen_factor=1, n_output_classes=10, mode='train'):
+def WideResnet(n_blocks=3, widen_factor=1, n_output_classes=10, bn_momentum=0.9,
+               mode='train'):
   """WideResnet from https://arxiv.org/pdf/1605.07146.pdf.
 
   Args:
     n_blocks: int, number of blocks in a group. total layers = 6n + 4.
     widen_factor: int, widening factor of each group. k=1 is vanilla resnet.
     n_output_classes: int, number of distinct output classes.
+    bn_momentum: float, momentum in BatchNorm.
     mode: Whether we are training or evaluating or doing inference.
 
   Returns:
@@ -150,10 +154,13 @@ def WideResnet(n_blocks=3, widen_factor=1, n_output_classes=10, mode='train'):
   return tl.Model(
       tl.ToFloat(),
       tl.Conv(16, (3, 3), padding='SAME'),
-      WideResnetGroup(n_blocks, 16 * widen_factor, mode=mode),
-      WideResnetGroup(n_blocks, 32 * widen_factor, (2, 2), mode=mode),
-      WideResnetGroup(n_blocks, 64 * widen_factor, (2, 2), mode=mode),
-      tl.BatchNorm(mode=mode),
+      WideResnetGroup(n_blocks, 16 * widen_factor, bn_momentum=bn_momentum,
+                      mode=mode),
+      WideResnetGroup(n_blocks, 32 * widen_factor, (2, 2),
+                      bn_momentum=bn_momentum, mode=mode),
+      WideResnetGroup(n_blocks, 64 * widen_factor, (2, 2),
+                      bn_momentum=bn_momentum, mode=mode),
+      tl.BatchNorm(momentum=bn_momentum, mode=mode),
       tl.Relu(),
       tl.AvgPool(pool_size=(8, 8)),
       tl.Flatten(),

From f23d147ef8910c94caf982c387ae2ecc122b2cdd Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Thu, 12 Sep 2019 21:27:59 -0700
Subject: [PATCH 2432/2720] increase window size of partially-built sequences
 in c++ packing op - results in more efficient packing.

PiperOrigin-RevId: 268831197
---
 tensor2tensor/data_generators/ops/pack_sequences_ops.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
index 9a5ed57eb..76ed97b4f 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
@@ -73,7 +73,7 @@ class PackSequences2Op : public OpKernel {
     for (int seq_id = 0; seq_id < n; seq_id++) {
       int inputs_length = inputs_lengths[seq_id];
       int targets_length = targets_lengths[seq_id];
-      for (int combined_id = std::max(0, num_combined - 10); true;
+      for (int combined_id = std::max(0, num_combined - 1000); true;
            combined_id++) {
         if (combined_id == num_combined) {
           combined_inputs_length.push_back(inputs_length);

From 8cb395e8036397aae1fa7383972cccea9f59e141 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 13 Sep 2019 08:34:01 -0700
Subject: [PATCH 2433/2720] Make PPO work with MultiDiscrete action spaces.

This is so that we can control more than one hyperparameter at a time, e.g. learning rate and weight decay. Different controls in the same timestep are drawn from independent categorical distributions and are assumed to have equal numbers of categories.

PiperOrigin-RevId: 268910551
---
 tensor2tensor/envs/env_problem_utils.py      |  20 ++--
 tensor2tensor/envs/env_problem_utils_test.py |   4 +-
 tensor2tensor/trax/learning_rate.py          |   3 +-
 tensor2tensor/trax/learning_rate_test.py     |   1 +
 tensor2tensor/trax/rl/envs/fake_env.py       |   9 +-
 tensor2tensor/trax/rl/ppo.py                 | 100 +++++++++++--------
 tensor2tensor/trax/rl/ppo_test.py            |  82 ++++++++-------
 tensor2tensor/trax/rl/ppo_trainer.py         |  31 ++++--
 tensor2tensor/trax/rl/ppo_trainer_test.py    |  14 +++
 9 files changed, 167 insertions(+), 97 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 9f15483de..dd3795765 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -21,6 +21,8 @@
 
 import functools
 import time
+
+import gym
 import numpy as np
 
 from tensor2tensor.envs import gym_env_problem
@@ -75,7 +77,7 @@ def play_env_problem_with_policy(env,
   Args:
     env: environment object, should be a subclass of env_problem.EnvProblem.
     policy_fun: callable, taking in observations((B, T) + OBS) and returning
-      back log-probabilities (B, T, A).
+      back log-probabilities (B, T, C, A).
     num_trajectories: int, number of trajectories to collect.
     max_timestep: int or None, if not None or a negative number, we cut any
       trajectory that exceeds this time put it in the completed bin, and *dont*
@@ -99,7 +101,7 @@ def gumbel_sample(log_probs):
     """Gumbel sampling."""
     u = np.random.uniform(low=1e-6, high=1.0 - 1e-6, size=log_probs.shape)
     g = -np.log(-np.log(u))
-    return np.argmax((log_probs / temperature) + g, axis=1)
+    return np.argmax((log_probs / temperature) + g, axis=-1)
 
   # We need to reset all environments, if we're coming here the first time.
   if reset or max_timestep is None or max_timestep <= 0:
@@ -132,21 +134,25 @@ def gumbel_sample(log_probs):
     policy_application_total_time += (time.time() - t1)
 
     assert (B, T) == log_prob_actions.shape[:2]
-    A = log_prob_actions.shape[2]  # pylint: disable=invalid-name
+    C = log_prob_actions.shape[2]  # pylint: disable=invalid-name
+    A = log_prob_actions.shape[3]  # pylint: disable=invalid-name
 
     # We need the log_probs of those actions that correspond to the last actual
     # time-step.
     index = lengths - 1  # Since we want to index using lengths.
-    log_probs = log_prob_actions[np.arange(B)[:, None], index[:, None],
-                                 np.arange(A)]
+    log_probs = log_prob_actions[np.arange(B)[:, None, None],
+                                 index[:, None, None],
+                                 np.arange(C)[:, None], np.arange(A)]
     value_preds = value_predictions[np.arange(B)[:, None], index[:, None],
                                     np.arange(1)]
-    assert (B, A) == log_probs.shape, \
-        "B=%d, A=%d, log_probs.shape=%s" % (B, A, log_probs.shape)
+    assert (B, C, A) == log_probs.shape, \
+        "B=%d, C=%d, A=%d, log_probs.shape=%s" % (B, C, A, log_probs.shape)
     assert (B, 1) == value_preds.shape, \
         "B=%d, value_preds.shape=%s" % (B, value_preds.shape)
 
     actions = gumbel_sample(log_probs)
+    if isinstance(env.action_space, gym.spaces.Discrete):
+      actions = np.squeeze(actions, axis=1)
 
     # Step through the env.
     t1 = time.time()
diff --git a/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index 0bda965ad..6304a3d16 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -62,10 +62,10 @@ def policy_fun(observations, state=None, rng=None):
       self.assertFalse(
           np.all(observations[:, :len_history_for_policy, ...] == 0))
       a = env.action_space.n
-      p = np.random.uniform(size=(b, t, a))
+      p = np.random.uniform(size=(b, t, 1, a))
       p = np.exp(p)
       p = p / np.sum(p, axis=-1, keepdims=True)
-      return np.log(p), np.log(p), state, rng
+      return np.log(p), np.mean(p, axis=-1), state, rng
 
     max_timestep = 15
     num_trajectories = 2
diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index 2201837ba..a10e27863 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -193,6 +193,7 @@ def PolicySchedule(
   # Build the policy network and load its parameters.
   start_time = time.time()
   net = ppo.policy_and_value_net(
+      n_controls=1,
       n_actions=len(action_multipliers),
       bottom_layers_fn=policy_and_value_model,
       two_towers=policy_and_value_two_towers,
@@ -225,5 +226,5 @@ def PolicySchedule(
 
   # Get a new learning rate.
   new_lr = online_tune.new_learning_rate(
-      action, history, action_multipliers, max_lr)
+      action.item(), history, action_multipliers, max_lr)
   return lambda _: new_lr
diff --git a/tensor2tensor/trax/learning_rate_test.py b/tensor2tensor/trax/learning_rate_test.py
index 6ba477016..73063aab4 100644
--- a/tensor2tensor/trax/learning_rate_test.py
+++ b/tensor2tensor/trax/learning_rate_test.py
@@ -43,6 +43,7 @@ def _make_schedule(
     policy_and_value_model = atari_cnn.FrameStackMLP
     net = ppo.policy_and_value_net(
         n_actions=len(action_multipliers),
+        n_controls=1,
         bottom_layers_fn=policy_and_value_model,
         two_towers=False,
     )
diff --git a/tensor2tensor/trax/rl/envs/fake_env.py b/tensor2tensor/trax/rl/envs/fake_env.py
index cff0fac7b..2d18c3f82 100644
--- a/tensor2tensor/trax/rl/envs/fake_env.py
+++ b/tensor2tensor/trax/rl/envs/fake_env.py
@@ -28,19 +28,23 @@
 import numpy as np
 
 
-class FakeEnv(object):
+class FakeEnv(gym.Env):
   """A fake env which is either done with a specific action or a time-step."""
 
   def __init__(self,
                input_shape=(4,),
                n_actions=2,
+               n_controls=1,
                done_time_step=None,
                done_action=None):
     self._input_shape = input_shape
     self._done_time_step = done_time_step
     self._done_action = done_action
     self._t = 0
-    self.action_space = gym.spaces.Discrete(n_actions)
+    if n_controls == 1:
+      self.action_space = gym.spaces.Discrete(n_actions)
+    else:
+      self.action_space = gym.spaces.MultiDiscrete([n_actions] * n_controls)
     self.observation_space = gym.spaces.Box(
         low=-1.0, high=1.0, shape=input_shape)
 
@@ -52,6 +56,7 @@ def reset(self):
     return self._get_random_observation()
 
   def step(self, action):
+    assert self.action_space.contains(action)
     done = False
     if self._done_action is not None:
       done = action == self._done_action
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index 469625aa0..175a445d6 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -22,13 +22,14 @@
              time-step dimension.
 OBS, tuple - shape of a singular observation from the environment.
              Ex: For CartPole-v0 this is (4,) and Pong-v0 it's (210, 160, 3)
+C, scalar  - Number of controls, i.e. independent groups of actions.
 A, scalar  - Number of actions, assuming a discrete space.
 
 Policy and Value function signatures:
 
-Policy            Function :: [B, T] + OBS ->  [B, T, A]
+Policy            Function :: [B, T] + OBS ->  [B, T, C, A]
 Value             Function :: [B, T] + OBS ->  [B, T, 1]
-Policy and Value  Function :: [B, T] + OBS -> ([B, T, A], [B, T, 1])
+Policy and Value  Function :: [B, T] + OBS -> ([B, T, C, A], [B, T, 1])
 
 i.e. the policy net should take a batch of *trajectories* and at each time-step
 in each batch deliver a probability distribution over actions.
@@ -69,7 +70,7 @@
 from tensorflow.io import gfile
 
 
-def policy_and_value_net(n_actions, bottom_layers_fn, two_towers):
+def policy_and_value_net(n_actions, n_controls, bottom_layers_fn, two_towers):
   """A policy and value net function."""
 
   # Layers.
@@ -78,12 +79,20 @@ def policy_and_value_net(n_actions, bottom_layers_fn, two_towers):
   # other computes the value function.
   # NOTE: The LogSoftmax instead of the Softmax because of numerical stability.
 
+  @tl.layer()
+  def SplitActions(x, **unused_kwargs):  # pylint: disable=invalid-name
+    """Splits logits for actions in different controls."""
+    return np.reshape(x, x.shape[:-1] + (n_controls, n_actions))
+
+  n_logits = n_controls * n_actions
+
   if two_towers:
     layers = [
         tl.Dup(),
         tl.Parallel(
             [bottom_layers_fn(),
-             tl.Dense(n_actions),
+             tl.Dense(n_logits),
+             SplitActions(),  # pylint: disable=no-value-for-parameter
              tl.LogSoftmax()],
             [bottom_layers_fn(), tl.Dense(1)],
         )
@@ -93,7 +102,9 @@ def policy_and_value_net(n_actions, bottom_layers_fn, two_towers):
         bottom_layers_fn(),
         tl.Dup(),
         tl.Parallel(
-            [tl.Dense(n_actions), tl.LogSoftmax()],
+            [tl.Dense(n_logits),
+             SplitActions(),  # pylint: disable=no-value-for-parameter
+             tl.LogSoftmax()],
             [tl.Dense(1)],
         )
     ]
@@ -168,7 +179,7 @@ def collect_trajectories(env,
     trajectory: list of (observation, action, reward) tuples, where each element
     `i` is a tuple of numpy arrays with shapes as follows:
     observation[i] = (B, T_i + 1)
-    action[i] = (B, T_i)
+    action[i] = (B, T_i, C)
     reward[i] = (B, T_i)
   """
 
@@ -472,54 +483,56 @@ def gae_advantages(td_deltas, mask, lambda_=0.95, gamma=0.99):
   return rewards_to_go(td_deltas, mask, lambda_ * gamma)
 
 
-def chosen_probabs(probab_observations, actions):
+def chosen_probabs(probab_actions, actions):
   """Picks out the probabilities of the actions along batch and time-steps.
 
   Args:
-    probab_observations: ndarray of shape `[B, T+1, A]`, where
-      probab_observations[b, t, i] contains the log-probability of action = i at
+    probab_actions: ndarray of shape `[B, T+1, C, A]`, where
+      probab_actions[b, t, i] contains the log-probability of action = i at
       the t^th time-step in the b^th trajectory.
-    actions: ndarray of shape `[B, T]`, with each entry in [0, A) denoting which
-      action was chosen in the b^th trajectory's t^th time-step.
+    actions: ndarray of shape `[B, T, C]`, with each entry in [0, A) denoting
+      which action was chosen in the b^th trajectory's t^th time-step.
 
   Returns:
-    `[B, T]` ndarray with the log-probabilities of the chosen actions.
+    `[B, T, C]` ndarray with the log-probabilities of the chosen actions.
   """
-  B, T = actions.shape  # pylint: disable=invalid-name
-  assert (B, T + 1) == probab_observations.shape[:2]
-  return probab_observations[np.arange(B)[:, None], np.arange(T), actions]
+  B, T, C = actions.shape  # pylint: disable=invalid-name
+  assert (B, T + 1, C) == probab_actions.shape[:3]
+  return probab_actions[
+      np.arange(B)[:, None, None], np.arange(T)[:, None], np.arange(C), actions]
 
 
 def compute_probab_ratios(p_new, p_old, actions, reward_mask):
   """Computes the probability ratios for each time-step in a trajectory.
 
   Args:
-    p_new: ndarray of shape [B, T+1, A] of the log-probabilities that the policy
-      network assigns to all the actions at each time-step in each batch using
-      the old parameters.
-    p_old: ndarray of shape [B, T+1, A], same as above, but using old policy
+    p_new: ndarray of shape [B, T+1, C, A] of the log-probabilities that the
+      policy network assigns to all the actions at each time-step in each batch
+      using the old parameters.
+    p_old: ndarray of shape [B, T+1, C, A], same as above, but using old policy
       network parameters.
-    actions: ndarray of shape [B, T] where each element is from [0, A).
+    actions: ndarray of shape [B, T, C] where each element is from [0, A).
     reward_mask: ndarray of shape [B, T] masking over probabilities.
 
   Returns:
-    probab_ratios: ndarray of shape [B, T], where
-    probab_ratios_{b,t} = p_new_{b,t,action_{b,t}} / p_old_{b,t,action_{b,t}}
+    probab_ratios: ndarray of shape [B, T, C], where
+    probab_ratios_{b,t, c} = p_new_{b,t,c,action_{b,t,c}} /
+                             p_old_{b,t,c,action_{b,t,c}}
   """
 
-  B, T = actions.shape  # pylint: disable=invalid-name
-  assert (B, T + 1) == p_old.shape[:2]
-  assert (B, T + 1) == p_new.shape[:2]
+  B, T, C = actions.shape  # pylint: disable=invalid-name
+  assert (B, T + 1, C) == p_old.shape[:3]
+  assert (B, T + 1, C) == p_new.shape[:3]
 
   logp_old = chosen_probabs(p_old, actions)
   logp_new = chosen_probabs(p_new, actions)
 
-  assert (B, T) == logp_old.shape
-  assert (B, T) == logp_new.shape
+  assert (B, T, C) == logp_old.shape
+  assert (B, T, C) == logp_new.shape
 
   # Since these are log-probabilities, we just subtract them.
-  probab_ratios = np.exp(logp_new - logp_old) * reward_mask
-  assert (B, T) == probab_ratios.shape
+  probab_ratios = np.exp(logp_new - logp_old) * reward_mask[:, :, None]
+  assert (B, T, C) == probab_ratios.shape
   return probab_ratios
 
 
@@ -528,10 +541,11 @@ def clipped_probab_ratios(probab_ratios, epsilon=0.2):
 
 
 def clipped_objective(probab_ratios, advantages, reward_mask, epsilon=0.2):
+  advantages = advantages[:, :, None]
   return np.minimum(
       probab_ratios * advantages,
       clipped_probab_ratios(probab_ratios, epsilon=epsilon) *
-      advantages) * reward_mask
+      advantages) * reward_mask[:, :, None]
 
 
 @jit
@@ -546,13 +560,15 @@ def ppo_loss_given_predictions(log_probab_actions_new,
                                epsilon=0.2):
   """PPO objective, with an eventual minus sign, given predictions."""
   B, T = padded_rewards.shape  # pylint: disable=invalid-name
-  assert (B, T) == padded_actions.shape
+  _, _, C, A = log_probab_actions_old.shape  # pylint: disable=invalid-name
+
+  assert (B, T) == padded_rewards.shape
+  assert (B, T, C) == padded_actions.shape
   assert (B, T) == reward_mask.shape
 
-  _, _, A = log_probab_actions_old.shape  # pylint: disable=invalid-name
   assert (B, T + 1, 1) == value_predictions_old.shape
-  assert (B, T + 1, A) == log_probab_actions_old.shape
-  assert (B, T + 1, A) == log_probab_actions_new.shape
+  assert (B, T + 1, C, A) == log_probab_actions_old.shape
+  assert (B, T + 1, C, A) == log_probab_actions_new.shape
 
   # (B, T)
   td_deltas = deltas(
@@ -573,12 +589,12 @@ def ppo_loss_given_predictions(log_probab_actions_new,
   # (B, T)
   ratios = compute_probab_ratios(log_probab_actions_new, log_probab_actions_old,
                                  padded_actions, reward_mask)
-  assert (B, T) == ratios.shape
+  assert (B, T, C) == ratios.shape
 
   # (B, T)
   objective = clipped_objective(
       ratios, advantages, reward_mask, epsilon=epsilon)
-  assert (B, T) == objective.shape
+  assert (B, T, C) == objective.shape
 
   # ()
   average_objective = np.sum(objective) / np.sum(reward_mask)
@@ -735,8 +751,8 @@ def approximate_kl(log_prob_new, log_prob_old, mask):
   """Computes the approximate KL divergence between the old and new log-probs.
 
   Args:
-    log_prob_new: (B, T+1, A) log probs new
-    log_prob_old: (B, T+1, A) log probs old
+    log_prob_new: (B, T+1, C, A) log probs new
+    log_prob_old: (B, T+1, C, A) log probs old
     mask: (B, T)
 
   Returns:
@@ -746,7 +762,7 @@ def approximate_kl(log_prob_new, log_prob_old, mask):
   # Cut the last time-step out.
   diff = diff[:, :-1]
   # Mask out the irrelevant part.
-  diff *= mask[:, :, np.newaxis]  # make mask (B, T, 1)
+  diff *= mask[:, :, np.newaxis, np.newaxis]  # make mask (B, T, 1, 1)
   # Average on non-masked part.
   return np.sum(diff) / np.sum(mask)
 
@@ -755,7 +771,7 @@ def masked_entropy(log_probs, mask):
   """Computes the entropy for the given log-probs.
 
   Args:
-    log_probs: (B, T+1, A) log probs
+    log_probs: (B, T+1, C, A) log probs
     mask: (B, T) mask.
 
   Returns:
@@ -764,8 +780,8 @@ def masked_entropy(log_probs, mask):
   # Cut the last time-step out.
   lp = log_probs[:, :-1]
   # Mask out the irrelevant part.
-  lp *= mask[:, :, np.newaxis]  # make mask (B, T, 1)
-  p = np.exp(lp) * mask[:, :, np.newaxis]  # (B, T, 1)
+  lp *= mask[:, :, np.newaxis, np.newaxis]  # make mask (B, T, 1, 1)
+  p = np.exp(lp) * mask[:, :, np.newaxis, np.newaxis]  # (B, T, 1, 1)
   # Average on non-masked part and take negative.
   return -(np.sum(lp * p) / np.sum(mask))
 
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
index 3be53bb4f..aa25e51b9 100644
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -77,8 +77,13 @@ def test_policy_and_value_net(self):
     observation_shape = (3, 4, 5)
     batch_observation_shape = (1, 1) + observation_shape
     n_actions = 2
+    n_controls = 3
     pnv_model = ppo.policy_and_value_net(
-        n_actions, lambda: [layers.Flatten(n_axes_to_keep=2)], two_towers=True)
+        n_controls=n_controls,
+        n_actions=n_actions,
+        bottom_layers_fn=lambda: [layers.Flatten(n_axes_to_keep=2)],
+        two_towers=True,
+    )
     pnv_params, pnv_state = pnv_model.initialize(
         batch_observation_shape, np.float32, self.rng_key)
 
@@ -90,7 +95,8 @@ def test_policy_and_value_net(self):
 
     # Output is a list, first is probab of actions and the next is value output.
     self.assertEqual(2, len(pnv_output))
-    self.assertEqual((batch, time_steps, n_actions), pnv_output[0].shape)
+    self.assertEqual(
+        (batch, time_steps, n_controls, n_actions), pnv_output[0].shape)
     self.assertEqual((batch, time_steps, 1), pnv_output[1].shape)
 
   def test_pad_trajectories(self):
@@ -308,18 +314,19 @@ def test_gae_advantages(self):
     self.assertAllEqual(expected_gae_advantages, gae_advantages)
 
   def test_chosen_probabs(self):
-    # Shape (2, 2+1, 3)
+    # Shape (2, 2+1, 1, 3)
     probab_observations = np.array(
-        [[[0.1, 0.2, 0.7], [0.4, 0.1, 0.5], [0.2, 0.4, 0.4]],
-         [[0.3, 0.1, 0.6], [0.1, 0.1, 0.8], [0.2, 0.4, 0.4]]]
+        [[[[0.1, 0.2, 0.7]], [[0.4, 0.1, 0.5]], [[0.2, 0.4, 0.4]]],
+         [[[0.3, 0.1, 0.6]], [[0.1, 0.1, 0.8]], [[0.2, 0.4, 0.4]]]]
     )
 
-    # Shape (2, 2)
-    actions = np.array([[1, 2], [0, 1]])
+    # Shape (2, 2, 1)
+    actions = np.array([[[1], [2]], [[0], [1]]])
 
     chosen_probabs = ppo.chosen_probabs(probab_observations, actions)
 
-    self.assertAllEqual(np.array([[0.2, 0.5], [0.3, 0.1]]), chosen_probabs)
+    self.assertAllEqual(
+        np.array([[[0.2], [0.5]], [[0.3], [0.1]]]), chosen_probabs)
 
   def test_compute_probab_ratios(self):
     p_old = np.array([[
@@ -334,7 +341,7 @@ def test_compute_probab_ratios(self):
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
         [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-    ]])
+    ]])[:, :, None]
 
     p_new = np.array([[
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
@@ -348,9 +355,9 @@ def test_compute_probab_ratios(self):
         [np.log(0.3), np.log(0.1), np.log(0.3), np.log(0.3)],
         [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
         [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
-    ]])
+    ]])[:, :, None]
 
-    actions = np.array([[1, 2, 0, 1], [0, 3, 3, 0]])
+    actions = np.array([[1, 2, 0, 1], [0, 3, 3, 0]])[:, :, None]
 
     mask = np.array([[1, 1, 0, 0], [1, 1, 1, 0]])
 
@@ -360,7 +367,7 @@ def test_compute_probab_ratios(self):
         np.array([
             [0.1 / 0.2, 0.1 / 0.4, 0.0, 0.0],
             [0.1 / 0.3, 0.6 / 0.4, 0.3 / 0.1, 0.0],
-        ]), probab_ratios)
+        ])[:, :, None], probab_ratios)
 
   def test_clipped_probab_ratios(self):
     probab_ratios = np.array([
@@ -378,8 +385,8 @@ def test_clipped_probab_ratios(self):
 
   def test_clipped_objective(self):
     probab_ratios = np.array([
-        [1.5, 2.0, 0.5, 0.7],
-        [2.5, 2.0, 0.1, 1.0],
+        [[1.5], [2.0], [0.5], [0.7]],
+        [[2.5], [2.0], [0.1], [1.0]],
     ])
 
     advantages = np.array([
@@ -392,35 +399,36 @@ def test_clipped_objective(self):
     epsilon = 0.1
 
     clipped_probab_ratios = np.array([
-        [1.1, 1.1, 0.9, 0.9],
-        [1.1, 1.1, 0.9, 1.0],
+        [[1.1], [1.1], [0.9], [0.9]],
+        [[1.1], [1.1], [0.9], [1.0]],
     ])
 
     unused_advantages_x_probab_ratios = np.array([
-        [0.15, -0.2, 0.25, 0.49],
-        [5.00, -4.0, 0.20, 2.00]
+        [[0.15], [-0.2], [0.25], [0.49]],
+        [[5.00], [-4.0], [0.20], [2.00]]
     ])
 
     unused_advantages_x_clipped_probab_ratios = np.array([
-        [0.11, -0.11, 0.45, 0.63],
-        [2.20, -2.20, 1.80, 2.00]
+        [[0.11], [-0.11], [0.45], [0.63]],
+        [[2.20], [-2.20], [1.80], [2.00]]
     ])
 
     unused_minimums = np.array([
-        [0.11, -0.2, 0.25, 0.49],
-        [2.20, -4.0, 0.20, 2.00]
+        [[0.11], [-0.2], [0.25], [0.49]],
+        [[2.20], [-4.0], [0.20], [2.00]]
     ])
 
     # minimums * mask
     objective = np.array([
-        [0.11, -0.2, 0.0, 0.],
-        [2.20, -4.0, 0.2, 0.]
+        [[0.11], [-0.2], [0.0], [0.]],
+        [[2.20], [-4.0], [0.2], [0.]]
     ])
 
     # Assert that we computed things correctly in this test.
     self.assertAllClose(
-        np.minimum(probab_ratios * advantages,
-                   clipped_probab_ratios * advantages) * mask,
+        np.minimum(probab_ratios * advantages[:, :, None],
+                   clipped_probab_ratios * advantages[:, :, None]) *
+        mask[:, :, None],
         objective)
 
     self.assertAllClose(
@@ -434,7 +442,11 @@ def test_combined_loss(self):
     batch_observation_shape = (1, 1) + OBS
 
     net = ppo.policy_and_value_net(
-        A, lambda: [layers.Flatten(n_axes_to_keep=2)], two_towers=True)
+        n_controls=1,
+        n_actions=A,
+        bottom_layers_fn=lambda: [layers.Flatten(n_axes_to_keep=2)],
+        two_towers=True,
+    )
 
     old_params, _ = net.initialize(
         batch_observation_shape, np.float32, key1)
@@ -444,7 +456,7 @@ def test_combined_loss(self):
     # Generate a batch of observations.
 
     observations = np.random.uniform(size=(B, T + 1) + OBS)
-    actions = np.random.randint(0, A, size=(B, T))
+    actions = np.random.randint(0, A, size=(B, T, 1))
     rewards = np.random.uniform(0, 1, size=(B, T))
     mask = np.ones_like(rewards)
 
@@ -500,7 +512,7 @@ def test_combined_loss(self):
                     1e-6)
 
   def test_masked_entropy(self):
-    # (2, 4+1, 4)
+    # (2, 4+1, 1, 4)
     log_probs = np.array([[
         [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
         [np.log(0.4), np.log(0.1), np.log(0.4), np.log(0.1)],
@@ -513,7 +525,7 @@ def test_masked_entropy(self):
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
         [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-    ]])
+    ]])[:, :, None, :]
 
     # (2, 4)
     mask = np.array([
@@ -526,11 +538,11 @@ def plp(p):
 
     # Removing the last time-step and the masked stuff, gets us this.
     filtered_log_probs = np.array([[
-        [plp(0.1), plp(0.2), plp(0.6), plp(0.1)],
-        [plp(0.4), plp(0.1), plp(0.4), plp(0.1)],
-        [plp(0.3), plp(0.1), plp(0.5), plp(0.1)],
-        [plp(0.1), plp(0.1), plp(0.4), plp(0.4)],
-        [plp(0.3), plp(0.1), plp(0.5), plp(0.1)],
+        [[plp(0.1), plp(0.2), plp(0.6), plp(0.1)]],
+        [[plp(0.4), plp(0.1), plp(0.4), plp(0.1)]],
+        [[plp(0.3), plp(0.1), plp(0.5), plp(0.1)]],
+        [[plp(0.1), plp(0.1), plp(0.4), plp(0.4)]],
+        [[plp(0.3), plp(0.1), plp(0.5), plp(0.1)]],
     ]])
 
     self.assertNear(ppo.masked_entropy(log_probs, mask),
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index 3a4e62f5b..e31d9f9f8 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -142,8 +142,18 @@ def __init__(self,
     self._eval_temperatures = eval_temperatures
     self._separate_eval = separate_eval
 
-    assert isinstance(self.train_env.action_space, gym.spaces.Discrete)
-    n_actions = self.train_env.action_space.n
+    action_space = self.train_env.action_space
+    assert isinstance(
+        action_space, (gym.spaces.Discrete, gym.spaces.MultiDiscrete))
+    if isinstance(action_space, gym.spaces.Discrete):
+      n_actions = action_space.n
+      n_controls = 1
+    else:
+      (n_controls,) = action_space.nvec.shape
+      assert n_controls > 0
+      assert onp.min(action_space.nvec) == onp.max(action_space.nvec), (
+          "Every control must have the same number of actions.")
+      n_actions = action_space.nvec[0]
 
     # Batch Observations Shape = [1, 1] + OBS, because we will eventually call
     # policy and value networks on shape [B, T] +_OBS
@@ -156,6 +166,7 @@ def __init__(self,
     # Initialize the policy and value network.
     policy_and_value_net = ppo.policy_and_value_net(
         n_actions=n_actions,
+        n_controls=n_controls,
         bottom_layers_fn=policy_and_value_model,
         two_towers=policy_and_value_two_towers,
     )
@@ -351,8 +362,12 @@ def train_epoch(self, evaluate=True):
     logging.vlog(1, "Padded Actions' shape [%s]", str(padded_actions.shape))
     logging.vlog(1, "Padded Rewards' shape [%s]", str(padded_rewards.shape))
 
+    if padded_actions.ndim == 2:
+      # Add control axis.
+      padded_actions = np.expand_dims(padded_actions, axis=-1)
+
     # Some assertions.
-    B, T = padded_actions.shape  # pylint: disable=invalid-name
+    B, T, C = padded_actions.shape  # pylint: disable=invalid-name
     assert (B, T) == padded_rewards.shape
     assert (B, T) == reward_mask.shape
     assert (B, T + 1) == padded_observations.shape[:2]
@@ -367,12 +382,12 @@ def train_epoch(self, evaluate=True):
     actual_log_probabs_traj = padded_infos["log_prob_actions"]
     actual_value_predictions_traj = padded_infos["value_predictions"]
 
-    assert (B, T) == actual_log_probabs_traj.shape[:2]
-    A = actual_log_probabs_traj.shape[2]  # pylint: disable=invalid-name
+    assert (B, T, C) == actual_log_probabs_traj.shape[:3]
+    A = actual_log_probabs_traj.shape[3]  # pylint: disable=invalid-name
     assert (B, T, 1) == actual_value_predictions_traj.shape
 
-    # TODO(afrozm): log-probabs doesn't need to be (B, T+1, A) it can do with
-    # (B, T, A), so make that change throughout.
+    # TODO(afrozm): log-probabs doesn't need to be (B, T+1, C, A) it can do with
+    # (B, T, C, A), so make that change throughout.
 
     # NOTE: We don't have the log-probabs and value-predictions for the last
     # observation, so we re-calculate for everything, but use the original ones
@@ -382,7 +397,7 @@ def train_epoch(self, evaluate=True):
     log_probabs_traj, value_predictions_traj, self._model_state, _ = (
         self._get_predictions(padded_observations, self._model_state, rng=key))
 
-    assert (B, T + 1, A) == log_probabs_traj.shape
+    assert (B, T + 1, C, A) == log_probabs_traj.shape
     assert (B, T + 1, 1) == value_predictions_traj.shape
 
     # Concatenate the last time-step's log-probabs and value predictions to the
diff --git a/tensor2tensor/trax/rl/ppo_trainer_test.py b/tensor2tensor/trax/rl/ppo_trainer_test.py
index 4a13f781a..da11555fd 100644
--- a/tensor2tensor/trax/rl/ppo_trainer_test.py
+++ b/tensor2tensor/trax/rl/ppo_trainer_test.py
@@ -240,6 +240,20 @@ def test_restarts(self):
       trainer.training_loop(n_epochs=2)
       self.assertEqual(trainer.epoch, 2)
 
+  def test_training_loop_multi_control(self):
+    gym.register(
+        "FakeEnv-v0",
+        entry_point="tensor2tensor.trax.rl.envs.fake_env:FakeEnv",
+        kwargs={"n_actions": 3, "n_controls": 2},
+    )
+    with self.tmp_dir() as output_dir:
+      trainer = self._make_trainer(
+          train_env=self.get_wrapped_env("FakeEnv-v0", 2),
+          eval_env=self.get_wrapped_env("FakeEnv-v0", 2),
+          output_dir=output_dir,
+      )
+      trainer.training_loop(n_epochs=2)
+
 
 if __name__ == "__main__":
   test.main()

From 9f87f153f3ab59d7d273d5b1ddabbd161638be20 Mon Sep 17 00:00:00 2001
From: Richard Csaky <ricsinaruto@hotmail.com>
Date: Fri, 13 Sep 2019 19:49:16 +0300
Subject: [PATCH 2434/2720] 4 new dialog problems (#1642)

* mostly working

* fixed opensubtitles

* fixed almost everything

* fixed everything

* fixing conflict

* re-added problems

* fixed docstrings
---
 tensor2tensor/data_generators/all_problems.py |   6 +-
 .../data_generators/dialog_abstract.py        | 389 ++++++++++++++++++
 .../data_generators/dialog_cornell.py         | 161 ++++++++
 .../data_generators/dialog_dailydialog.py     | 127 ++++++
 .../data_generators/dialog_opensubtitles.py   | 241 +++++++++++
 .../data_generators/dialog_personachat.py     | 195 +++++++++
 6 files changed, 1117 insertions(+), 2 deletions(-)
 create mode 100644 tensor2tensor/data_generators/dialog_abstract.py
 create mode 100644 tensor2tensor/data_generators/dialog_cornell.py
 create mode 100644 tensor2tensor/data_generators/dialog_dailydialog.py
 create mode 100644 tensor2tensor/data_generators/dialog_opensubtitles.py
 create mode 100644 tensor2tensor/data_generators/dialog_personachat.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index a865a0ceb..d4dd3f7cb 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -19,7 +19,6 @@
 from __future__ import print_function
 
 import importlib
-import six
 from six.moves import range  # pylint: disable=redefined-builtin
 
 MODULES = [
@@ -39,6 +38,10 @@
     "tensor2tensor.data_generators.cola",
     "tensor2tensor.data_generators.common_voice",
     "tensor2tensor.data_generators.desc2code",
+    "tensor2tensor.data_generators.dialog_cornell",
+    "tensor2tensor.data_generators.dialog_dailydialog",
+    "tensor2tensor.data_generators.dialog_opensubtitles",
+    "tensor2tensor.data_generators.dialog_personachat",
     "tensor2tensor.data_generators.enwik8",
     "tensor2tensor.data_generators.fsns",
     "tensor2tensor.data_generators.function_docstring",
@@ -104,7 +107,6 @@
 ALL_MODULES = list(MODULES)
 
 
-
 def _is_import_err_msg(err_str, module):
   parts = module.split(".")
   suffixes = [".".join(parts[i:]) for i in range(len(parts))]
diff --git a/tensor2tensor/data_generators/dialog_abstract.py b/tensor2tensor/data_generators/dialog_abstract.py
new file mode 100644
index 000000000..3d8fcd233
--- /dev/null
+++ b/tensor2tensor/data_generators/dialog_abstract.py
@@ -0,0 +1,389 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import re
+import tarfile
+import zipfile
+
+import requests
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators.text_problems import VocabType
+from tensor2tensor.layers import modalities
+from tensor2tensor.utils import metrics
+import tensorflow as tf
+
+# End-of-sentence marker.
+EOS = text_encoder.EOS_ID
+
+
+# An abstract base class for word based chatbot problems.
+class DialogAbstract(text_problems.Text2TextProblem):
+
+  @property
+  def vocab_type(self):
+    return text_problems.VocabType.TOKEN
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  @property
+  def vocab_file(self):
+    return self.vocab_filename
+
+  @property
+  def vocab_filename(self):
+    return 'vocab.chatbot.' + str(self.targeted_vocab_size)
+
+  @property
+  def oov_token(self):
+    return '<unk>'
+
+  @property
+  def use_subword_tokenizer(self):
+    return False
+
+  @property
+  def input_space_id(self):
+    return problem.SpaceID.EN_TOK
+
+  @property
+  def target_space_id(self):
+    return problem.SpaceID.EN_TOK
+
+  @property
+  def targeted_vocab_size(self):
+    return 2**14
+
+  @property
+  def targeted_dataset_size(self):
+    # Number of utterance pairs in the full dataset.
+    # If it's 0, then the full size of the dataset is used.
+    return 0
+
+  @property
+  def dataset_split(self):
+    return {'train': 80, 'val': 10, 'test': 10}
+
+  @property
+  def dataset_splits(self):
+    return [{
+        'split': problem.DatasetSplit.TRAIN,
+        'shards': 1,
+    }, {
+        'split': problem.DatasetSplit.EVAL,
+        'shards': 1,
+    }, {
+        'split': problem.DatasetSplit.TEST,
+        'shards': 1,
+    }]
+
+  @property
+  def data_dir(self):
+    return ''
+
+  @property
+  def raw_data_dir(self):
+    return ''
+
+  @property
+  def raw_data(self):
+    return ''
+
+  @property
+  def zipped_data(self):
+    return ''
+
+  @property
+  def url(self):
+    return ''
+
+  @data_dir.setter
+  def data_dir(self, value):
+    self._data_dir = value
+
+  @raw_data_dir.setter
+  def raw_data_dir(self, value):
+    self._raw_data_dir = value
+
+  @raw_data.setter
+  def raw_data(self, value):
+    self._raw_data = value
+
+  @zipped_data.setter
+  def zipped_data(self, value):
+    self._zipped_data = value
+
+  @url.setter
+  def url(self, value):
+    self._url = value
+
+  # Main function where the preprocessing of the data starts.
+  def preprocess_data(self, train_mode):
+    return NotImplementedError
+
+  # This should also be overriden if the data_pipeline_status is used.
+  def create_data(self, train_mode):
+    pass
+
+  def data_pipeline_status(self, train_mode):
+    """Check at which part of the pipeline are we at.
+
+    This function first checks recursively at which point in the
+    data processing point are we (what files can be found on the disk),
+    and then proceeds from there.
+
+    Args:
+      train_mode: string, whether we are in train, dev or test mode
+    """
+
+    # Build the source and target paths.
+    sourcepath = os.path.join(self._data_dir, train_mode + 'Source.txt')
+    targetpath = os.path.join(self._data_dir, train_mode + 'Target.txt')
+
+    # If raw data dir doesn't exist, create it.
+    if not os.path.exists(self._raw_data_dir):
+      os.makedirs(self._raw_data_dir)
+
+    # Check whether sourcePath.txt exists.
+    if (os.path.isfile(sourcepath) and os.path.isfile(targetpath) and
+        os.path.isfile(os.path.join(self._data_dir, self.vocab_file))):
+      print('problem_log: Source, target and vocab files exist in ' +
+            self._data_dir + ', proceeding with data generation. ' +
+            'If you want to rebuild these files, delete them first.')
+      return
+
+    # Check whether the raw data is extracted to the raw_data_dir folder.
+    elif os.path.exists(self._raw_data):
+      print('problem_log: No source, target or vocab files found in ' +
+            self._data_dir + '.')
+      print('problem_log: Extracted raw data is in ' + self._raw_data_dir +
+            '. Proceeding with creating source, target and vocab files.')
+      self.create_data(train_mode)
+
+    # Check whether the data is downloaded in the raw_data_dir_folder.
+    elif os.path.exists(self._zipped_data):
+      print('problem_log: No source, target or vocab files found in ' +
+            self._data_dir + '.')
+      print('problem_log: No extracted raw data found in ' +
+            self._raw_data_dir + '.')
+      print('problem_log: Unextracted raw data is in ' + self._raw_data_dir +
+            '. Extracting and creating source, target and vocab files.')
+      self.extract_data(train_mode)
+
+    else:
+      print('problem_log: No source, target or vocab files found in ' +
+            self._data_dir + '.')
+      print('problem_log: No raw data found in ' + self._raw_data_dir +
+            '. Proceeding with downloading the data, extracting it, ' +
+            'and creating source, target and vocab files.')
+      self.download_data(train_mode)
+
+  def download_data(self, train_mode):
+    """Download data from official sources.
+
+    Args:
+      train_mode: string, whether we are in train, dev or test mode
+    """
+
+    # Open the url and download the data with progress bars.
+    data_stream = requests.get(self._url, stream=True)
+    with open(self._zipped_data, 'wb') as f:
+      for chunk in data_stream.iter_content(1024):
+        if chunk:
+          f.write(chunk)
+          f.flush()
+
+    # Next step is extracting the data.
+    print('problem_log: Extracting data to ' + self._zipped_data + '.')
+    self.extract_data(train_mode)
+
+  def extract_data(self, train_mode):
+    """Extract data and go to the next step.
+
+    Args:
+      train_mode:  string, whether we are in train, dev or test mode
+    """
+
+    if self._zipped_data[-2:] == 'gz':
+      zip_file = tarfile.open(self._zipped_data, 'r:gz')
+    elif self._zipped_data[-3:] == 'zip':
+      zip_file = zipfile.ZipFile(self._zipped_data, 'r')
+    else:
+      print('problem_log: ' + self._zipped_data +
+            ' is not a .zip or .gz file, so I can\'t extract it.')
+
+    zip_file.extractall(self._raw_data_dir)
+    zip_file.close()
+
+    # Next step is creating the source, target and vocab files.
+    print('problem_log: Creating ' +
+          train_mode + ' files in ' + self._data_dir)
+    self.create_data(train_mode)
+
+  # hparams for the problem.
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    p.stop_at_eos = int(True)
+
+    p.modality = {'targets': modalities.ModalityType.SYMBOL}
+    if self.has_inputs:
+      p.modality['inputs'] = modalities.ModalityType.SYMBOL
+      p.vocab_size = {'inputs': self._encoders['inputs'].vocab_size}
+    p.vocab_size['targets'] = self._encoders['inputs'].vocab_size
+
+    if self.vocab_type == VocabType.CHARACTER:
+      p.loss_multiplier = 2.0
+
+    if self.packed_length:
+      if self.has_inputs:
+        p.modality['inputs_segmentation'] = modalities.ModalityType.IDENTITY
+        p.modality['inputs_position'] = modalities.ModalityType.IDENTITY
+        p.vocab_size['inputs_segmentation'] = None
+        p.vocab_size['inputs_position'] = None
+      p.modality['targets_segmentation'] = modalities.ModalityType.IDENTITY
+      p.modality['targets_position'] = modalities.ModalityType.IDENTITY
+      p.vocab_size['targets_segmentation'] = None
+      p.vocab_size['targets_position'] = None
+
+  # What evaluation metrics to use with this problem.
+  def eval_metrics(self):
+    return [metrics.Metrics.ACC, metrics.Metrics.ACC_TOP5,
+            metrics.Metrics.ACC_PER_SEQ,
+            metrics.Metrics.NEG_LOG_PERPLEXITY,
+            metrics.Metrics.APPROX_BLEU]
+
+  # Override this, to start with preprocessing.
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
+    self.data_dir = data_dir
+    # Determine whether we are in training or validation mode.
+    self.mode = {problem.DatasetSplit.TRAIN: 'train',
+                 problem.DatasetSplit.EVAL: 'dev',
+                 problem.DatasetSplit.TEST: 'test'}
+    filepath_fns = {problem.DatasetSplit.TRAIN: self.training_filepaths,
+                    problem.DatasetSplit.EVAL: self.dev_filepaths,
+                    problem.DatasetSplit.TEST: self.test_filepaths}
+
+    split_paths = [(split['split'], filepath_fns[split['split']](
+        data_dir, split['shards'], shuffled=self.already_shuffled))
+                   for split in self.dataset_splits]
+    all_paths = []
+    for _, paths in split_paths:
+      all_paths.extend(paths)
+
+    if self.is_generate_per_split:
+      for split, paths in split_paths:
+        # Create the source and target txt files from the raw data.
+        self.preprocess_data(self.mode[split])
+        generator_utils.generate_files(
+            self.generate_encoded_samples(data_dir, tmp_dir, split), paths)
+    else:
+      self.preprocess_data(self.mode[problem.DatasetSplit.TRAIN])
+      generator_utils.generate_files(
+          self.generate_encoded_samples(
+              data_dir, tmp_dir, problem.DatasetSplit.TRAIN), all_paths)
+
+    generator_utils.shuffle_dataset(all_paths, extra_fn=self._pack_fn())
+
+  def generate_samples(self, data_dir, tmp_dir, data_split):
+    """This function generates train and validation pairs in t2t-datagen style.
+
+    The function assumes that if you have data at one level of the pipeline,
+    you don't want to re-generate it, so for example if the 4 txt files exist,
+    the function continues by generating the t2t-datagen format files.
+    So if you want to re-download or re-generate data,
+    you have to delete it first from the appropriate directories.
+
+    Args:
+      data_dir: string, Directory where the data will be generated. The raw
+                        data has to be downloaded one directory level higher.
+      data_split: stromg, which data split to generate samples for
+
+    Yields:
+      dict
+    """
+
+    self.data_dir = data_dir
+    print('problem_log: ' +
+          self.mode[data_split] + ' data generation activated.')
+
+    s_path = os.path.join(data_dir, self.mode[data_split] + 'Source.txt')
+    t_path = os.path.join(data_dir, self.mode[data_split] + 'Target.txt')
+
+    # Open the files and yield source-target lines.
+    with tf.gfile.GFile(s_path, mode='r') as source_file:
+      with tf.gfile.GFile(t_path, mode='r') as target_file:
+        source, target = source_file.readline(), target_file.readline()
+        while source and target:
+          yield {'inputs': source.strip(), 'targets': target.strip()}
+          source, target = source_file.readline(), target_file.readline()
+
+  def save_vocab(self, vocab):
+    """Save the vocabulary to a file.
+
+    Args:
+      vocab: dict
+    """
+    voc_file = open(os.path.join(self._data_dir, self.vocab_file), 'w')
+
+    # Put the reserved tokens in.
+    voc_file.write('<pad>\n')
+    voc_file.write('<EOS>\n')
+    for word, _ in vocab.most_common(self.targeted_vocab_size - 3):
+      voc_file.write(word + '\n')
+    voc_file.write('<unk>')
+
+    voc_file.close()
+
+  # Open the 6 files to write the processed data into.
+  def open_6_files(self):
+    trainsource = open(os.path.join(self._data_dir, 'trainSource.txt'), 'w')
+    traintarget = open(os.path.join(self._data_dir, 'trainTarget.txt'), 'w')
+    devsource = open(os.path.join(self._data_dir, 'devSource.txt'), 'w')
+    devtarget = open(os.path.join(self._data_dir, 'devTarget.txt'), 'w')
+    testsource = open(os.path.join(self._data_dir, 'testSource.txt'), 'w')
+    testtarget = open(os.path.join(self._data_dir, 'testTarget.txt'), 'w')
+
+    return trainsource, traintarget, devsource, \
+        devtarget, testsource, testtarget
+
+  # Close the 6 files to write the processed data into.
+  def close_n_files(self, files):
+    for f in files:
+      f.close()
+
+  def clean_line(self, line):
+    """Clean a line with some regex rules.
+
+    Args:
+      line: string, line to be processed and returned
+
+    Returns:
+      string
+    """
+
+    # 2 functions for more complex replacing.
+    def replace(matchobj):
+      return re.sub("'", " '", str(matchobj.group(0)))
+
+    def replace_null(matchobj):
+      return re.sub("'", '', str(matchobj.group(0)))
+
+    # Keep some special tokens.
+    line = re.sub("[^a-z .?!'0-9]", '', line)
+    line = re.sub('[.]', ' . ', line)
+    line = re.sub('[?]', ' ? ', line)
+    line = re.sub('[!]', ' ! ', line)
+
+    # Take care of apostrophes.
+    line = re.sub("[ ]'[ ]", ' ', line)
+    line = re.sub(" '[a-z]", replace_null, line)
+    line = re.sub("n't", " n't", line)
+    line = re.sub("[^ n]'[^ t]", replace, line)
+
+    return line
diff --git a/tensor2tensor/data_generators/dialog_cornell.py b/tensor2tensor/data_generators/dialog_cornell.py
new file mode 100644
index 000000000..867b835f3
--- /dev/null
+++ b/tensor2tensor/data_generators/dialog_cornell.py
@@ -0,0 +1,161 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from collections import Counter
+import os
+import re
+
+from tensor2tensor.data_generators import dialog_abstract
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.utils import registry
+
+
+# End-of-sentence marker.
+EOS = text_encoder.EOS_ID
+
+
+@registry.register_problem
+class DialogCornell32k(dialog_abstract.DialogAbstract):
+  """Implements the chatbot problem with Cornell Movie Dialog dataset.
+
+  https://www.cs.cornell.edu/~cristian/Cornell_Movie-Dialogs_Corpus.html
+  """
+
+  @property
+  def targeted_vocab_size(self):
+    return 2**15
+
+  def preprocess_data(self, train_mode):
+    """Main function where the preprocessing of the data starts.
+
+    Args:
+      train_mode: string, whether we are in train, dev or test mode
+    """
+
+    # Set the raw data directory and data.
+    self.raw_data_dir = os.path.join('/'.join(self._data_dir.split('/')[:-1]),
+                                     'raw_data')
+    self.raw_data = os.path.join(self._raw_data_dir,
+                                 'cornell movie-dialogs corpus')
+    self.zipped_data = os.path.join(self._raw_data_dir,
+                                    'cornell_movie_dialogs_corpus.zip')
+
+    # Create the download url.
+    self.url = ('http://www.cs.cornell.edu/~cristian/data/' +
+                'cornell_movie_dialogs_corpus.zip')
+
+    # Check at which part of the pipeline are we at.
+    self.data_pipeline_status(train_mode)
+
+  def create_data(self, train_mode):
+    """Create the source, target and vocab files.
+
+    Args:
+      train_mode: string, whether we are in train, dev or test mode
+    """
+
+    # Open the 6 files.
+    trainsource, traintarget, devsource, devtarget, testsource, testtarget = \
+        self.open_6_files()
+
+    # Open the raw data.
+    movie_lines = open(
+        os.path.join(self._raw_data, 'movie_lines.txt'), errors='ignore')
+    dialog_list = self.extract_dialog_ids()
+
+    vocabulary = Counter()
+    line_dict = {}
+    number_of_lines = 0
+    # Iterate through file.
+    for line in movie_lines:
+      if number_of_lines % 10000 == 0:
+        print('problem_log: Parsed ' + str(number_of_lines) + ' lines.')
+
+      line = line.split(' +++$+++ ')
+      dialog_id = line[0]
+      line = line[4].lower()
+
+      # Do some cleaning.
+      line = self.clean_line(line)
+      line_dict[dialog_id] = line
+
+      number_of_lines += 1
+      # Check if we reached the desired dataset size.
+      if (self.targeted_dataset_size != 0 and
+          self.targeted_dataset_size < number_of_lines):
+        break
+
+    counter = 0
+    dataset_split_counter = 0
+    # Save the actual dialogs.
+    for dialog in dialog_list:
+      if counter % 10000 == 0:
+        print('problem_log: Saved ' +
+              str(counter) + '/' + str(len(dialog_list)) + ' dialogs.')
+
+      dataset_split_counter += 1
+      i = 0
+      # Save one utterance.
+      for utterance in dialog:
+        if (utterance != dialog[-1] and
+            dialog[i + 1] != 'L211194' and
+            dialog[i + 1] != 'L1045'):
+          source_line = line_dict[utterance] + '\n'
+          target_line = line_dict[dialog[i + 1]] + '\n'
+
+          # Save to the files according to dataset split.
+          if dataset_split_counter <= self.dataset_split['train']:
+            # Build vocabulary.
+            words = source_line.split()
+            for word in words:
+              vocabulary[word] = vocabulary.get(word, 0) + 1
+
+            trainsource.write(source_line)
+            traintarget.write(target_line)
+
+          elif dataset_split_counter <= (self.dataset_split['train'] +
+                                         self.dataset_split['val']):
+            devsource.write(source_line)
+            devtarget.write(target_line)
+          else:
+            testsource.write(source_line)
+            testtarget.write(target_line)
+        i += 1
+
+      # Reset the split counter if we reached 100%.
+      if dataset_split_counter == 100:
+        dataset_split_counter = 0
+      counter += 1
+
+    # Close the files.
+    self.close_n_files([trainsource,
+                        traintarget,
+                        devsource,
+                        devtarget,
+                        testsource,
+                        testtarget])
+    movie_lines.close()
+
+    # Save the vocabulary.
+    self.save_vocab(vocabulary)
+
+  # Extract the dialog ids from the dialog file.
+  def extract_dialog_ids(self):
+    dialogs = open(os.path.join(self._raw_data, 'movie_conversations.txt'),
+                   errors='ignore')
+
+    dialog_list = []
+    # Each line contains a dialog.
+    for line in dialogs:
+      line = line.split(' +++$+++ ')
+      line = line[3].split(',')
+
+      i = 0
+      for item in line:
+        line[i] = re.sub('[^A-Z0-9]', '', item)
+        i += 1
+      dialog_list.append(line)
+
+    dialogs.close()
+    return dialog_list
diff --git a/tensor2tensor/data_generators/dialog_dailydialog.py b/tensor2tensor/data_generators/dialog_dailydialog.py
new file mode 100644
index 000000000..98d04b05f
--- /dev/null
+++ b/tensor2tensor/data_generators/dialog_dailydialog.py
@@ -0,0 +1,127 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from collections import Counter
+import os
+
+from tensor2tensor.data_generators import dialog_abstract
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.utils import registry
+
+
+# End-of-sentence marker.
+EOS = text_encoder.EOS_ID
+
+
+@registry.register_problem
+class DialogDailydialog16k(dialog_abstract.DialogAbstract):
+  """A class implementing a simple chatbot problem for the DailyDialog dataset.
+
+  https://arxiv.org/abs/1710.03957
+  This version doesn't use any auxiliary information.
+  """
+
+  def preprocess_data(self, train_mode):
+    """Main function where the preprocessing of the data starts.
+
+    Args:
+      train_mode: string, whether we are in train, dev or test mode
+    """
+
+    # Set the raw data directory and data.
+    self.raw_data_dir = os.path.join('/'.join(self._data_dir.split('/')[:-1]),
+                                     'raw_data')
+    self.raw_data = os.path.join(self._raw_data_dir, 'ijcnlp_dailydialog')
+    self.zipped_data = os.path.join(self._raw_data_dir,
+                                    'ijcnlp_dailydialog.zip')
+
+    # Create the download url.
+    self.url = 'http://yanran.li/files/ijcnlp_dailydialog.zip'
+
+    # Check at which part of the pipeline are we at.
+    self.data_pipeline_status(train_mode)
+
+  def create_data(self, train_mode):
+    """Create the source, target and vocab files.
+
+    Args:
+      train_mode: string, whether we are in train, dev or test mode
+    """
+
+    # Open the 6 files.
+    trainsource, traintarget, devsource, devtarget, testsource, testtarget = \
+        self.open_6_files()
+
+    # Open the raw data.
+    dialogs = open(
+        os.path.join(self._raw_data, 'dialogues_text.txt'), errors='ignore')
+
+    vocabulary = Counter()
+    number_of_dialogs = 0
+    line_counter = 0
+    dataset_split_counter = 0
+    # Iterate through the file.
+    for dialog in dialogs:
+      dataset_split_counter += 1
+      if number_of_dialogs % 1000 == 0:
+        print('problem_log: Parsed ' + str(number_of_dialogs) + ' dialogs.')
+
+      # Utterances are separated by the __eou__ token.
+      utterances = dialog.split('__eou__')[:-1]
+
+      # Check which file we should write to.
+      if dataset_split_counter <= self.dataset_split['train']:
+        source_file = trainsource
+        target_file = traintarget
+      elif dataset_split_counter <= (self.dataset_split['train'] +
+                                     self.dataset_split['val']):
+        source_file = devsource
+        target_file = devtarget
+      else:
+        source_file = testsource
+        target_file = testtarget
+
+      # Clean the utterances.
+      i = 0
+      for utterance in utterances:
+        line_counter += 1
+        utterance = self.clean_line(utterance.lower())
+        i += 1
+
+        # Build vocabulary.
+        if dataset_split_counter <= self.dataset_split['train']:
+          words = utterance.split()
+          for word in words:
+            if word in vocabulary:
+              vocabulary[word] += 1
+            else:
+              vocabulary[word] = 1
+
+        # Write to files.
+        if i != len(utterances):
+          source_file.write(utterance + '\n')
+        if i != 1:
+          target_file.write(utterance + '\n')
+
+      number_of_dialogs += 1
+      # Reset the split counter if we reached 100%.
+      if dataset_split_counter == 100:
+        dataset_split_counter = 0
+
+      # Check if we reached the desired dataset size.
+      if (self.targeted_dataset_size != 0 and
+          self.targeted_dataset_size < line_counter):
+        break
+
+    # Close the files.
+    self.close_n_files([trainsource,
+                        traintarget,
+                        devsource,
+                        devtarget,
+                        testsource,
+                        testtarget])
+    dialogs.close()
+
+    # Save the vocabulary.
+    self.save_vocab(vocabulary)
diff --git a/tensor2tensor/data_generators/dialog_opensubtitles.py b/tensor2tensor/data_generators/dialog_opensubtitles.py
new file mode 100644
index 000000000..4807553b1
--- /dev/null
+++ b/tensor2tensor/data_generators/dialog_opensubtitles.py
@@ -0,0 +1,241 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from collections import Counter
+import os
+import re
+import zipfile
+
+from tensor2tensor.data_generators import dialog_abstract
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.utils import registry
+
+
+# End-of-sentence marker.
+EOS = text_encoder.EOS_ID
+
+
+@registry.register_problem
+class DialogOpensubtitles64k2009(dialog_abstract.DialogAbstract):
+  """A class implementing the chatbot problem for the OpenSubtitles dataset.
+
+  http://opus.nlpl.eu/OpenSubtitles-v2018.php
+  """
+
+  @property
+  def targeted_vocab_size(self):
+    return 2**16
+
+  @property
+  def dataset_version(self):
+    # Year of the opensubtitles dataset creation.
+    return 2009
+
+  def extract_data(self, train_mode):
+    """Extract data and go to the next step.
+
+    Args:
+      train_mode: string, whether we are in train, dev or test mode
+    """
+
+    if self._zipped_data[-3:] == 'zip' or self._zipped_data[-2:] == 'gz':
+      zip_file = zipfile.ZipFile(self._zipped_data, 'r')
+    else:
+      print('problem_log: ' + self._zipped_data +
+            ' is not a .zip or .gz file, so I can\'t extract it.')
+
+    zip_file.extractall(self._raw_data_dir)
+    zip_file.close()
+
+    # Next step is creating the source, target and vocab files.
+    print('problem_log: Creating ' +
+          train_mode + ' files in ' + self._data_dir)
+    self.create_data(train_mode)
+
+  def preprocess_data(self, train_mode):
+    """Main function where the preprocessing of the data starts.
+
+    Args:
+      train_mode: string, whether we are in train, dev or test mode
+    """
+
+    year = '' if self.dataset_version == 2009 else str(self.dataset_version)
+    # Set the raw data directory and data.
+    self.raw_data_dir = os.path.join('/'.join(self._data_dir.split('/')[:-1]),
+                                     'raw_data_' + str(self.dataset_version))
+    self.raw_data = os.path.join(self._raw_data_dir, 'OpenSubtitles' + year)
+    self.zipped_data = os.path.join(self._raw_data_dir, 'en.tar.gz')
+
+    # Create the download url.
+    self.url = ('http://opus.nlpl.eu/download.php?f=OpenSubtitles' +
+                str(year) + '/en.tar.gz')
+
+    # Check at which part of the pipeline are we at.
+    self.data_pipeline_status(train_mode)
+
+  def create_data(self, train_mode):
+    """Create the source, target and vocab files.
+
+    Args:
+      train_mode: string, whether we are in train, dev or test mode
+    """
+
+    # open the 6 files
+    trainsource, traintarget, devsource, devtarget, testsource, testtarget = \
+        self.open_6_files()
+
+    conv_id = 0
+    number_of_lines = 0
+    dataset_split_counter = 0
+    vocabulary = Counter()
+    # Dind all the files.
+    for root, _, files in os.walk(self._raw_data_dir):
+      for f in files:
+        if conv_id % 100 == 0:
+          print('problem_log: Parsed ' + str(conv_id) + ' files.')
+
+        source_lines = ''
+        target_lines = ''
+        conv_id += 1
+        dataset_split_counter += 1
+
+        # Open one .xml file and parse it.
+        with open(os.path.join(root, f), 'r', errors='ignore') as txt_file:
+          words = ''
+          line_id = 1
+
+          # Parse one line.
+          for line in txt_file:
+            line = str(line)
+
+            # Check if it's a new sentence.
+            if line.find('<s id="') != -1:
+              if words:
+                # Do some cleaning.
+                words = self.clean_line(words)
+
+                # Build the vocabulary.
+                if dataset_split_counter <= self.dataset_split['train']:
+                  word_list = words.split()
+                  for word in word_list:
+                    vocabulary[word] = vocabulary.get(word, 0) + 1
+
+                # Add the previous line.
+                source_lines += words + '\n'
+                if line_id != 1:
+                  target_lines += words + '\n'
+                line_id += 1
+              words = ''
+
+            else:
+              index = line.find('<w id="')
+              if index >= 0:
+                line = line[index:]
+                word = line[line.find('>') + 1:line.find('</w')]
+                words = words + ' ' + word.replace('\t', ' ')
+
+          # Delete the final source sentence, since it doesn't have a target.
+          source_lines = '\n'.join(source_lines.split('\n')[:-2]) + '\n'
+
+        # Save the dialog according to the dataset split.
+        if dataset_split_counter <= self.dataset_split['train']:
+          trainsource.write(source_lines)
+          traintarget.write(target_lines)
+        elif dataset_split_counter <= (self.dataset_split['train'] +
+                                       self.dataset_split['val']):
+          devsource.write(source_lines)
+          devtarget.write(target_lines)
+        else:
+          testsource.write(source_lines)
+          testtarget.write(target_lines)
+
+        # Reset the split counter if we reached 100%.
+        if dataset_split_counter == 100:
+          dataset_split_counter = 0
+
+        # Check if we reached the desired dataset size.
+        number_of_lines += line_id
+        if (self.targeted_dataset_size != 0 and
+            self.targeted_dataset_size < number_of_lines):
+          break
+      else:
+        continue
+      break
+
+    # Close the files.
+    self.close_n_files([trainsource,
+                        traintarget,
+                        devsource,
+                        devtarget,
+                        testsource,
+                        testtarget])
+    # Save the vocabulary.
+    self.save_vocab(vocabulary)
+
+  def clean_line(self, line):
+    """Clean a line with some regex rules.
+
+    Args:
+      line: string, line to be processed and returned
+
+    Returns:
+      string
+    """
+
+    line = line.lower()
+    line = re.sub("[^a-z .!?'\t\\\\]", '', line)
+    line = re.sub("\\\\['] ", " '", line)
+    line = re.sub('[\\\\]', ' ', line)
+    line = re.sub('[.]', ' . ', line)
+    line = re.sub('[?]', ' ? ', line)
+    line = re.sub('[!]', ' ! ', line)
+    line = re.sub("[ ]'[ ]", ' ', line)
+    line = re.sub("n't", " n't", line)
+
+    return line
+
+
+@registry.register_problem
+class DialogOpensubtitles64k2011(DialogOpensubtitles64k2009):
+
+  @property
+  def dataset_version(self):
+    # Year of the opensubtitles dataset creation.
+    return 2011
+
+
+@registry.register_problem
+class DialogOpensubtitles64k2012(DialogOpensubtitles64k2009):
+
+  @property
+  def dataset_version(self):
+    # Year of the opensubtitles dataset creation.
+    return 2012
+
+
+@registry.register_problem
+class DialogOpensubtitles64k2013(DialogOpensubtitles64k2009):
+
+  @property
+  def dataset_version(self):
+    # Year of the opensubtitles dataset creation.
+    return 2013
+
+
+@registry.register_problem
+class DialogOpensubtitles64k2016(DialogOpensubtitles64k2009):
+
+  @property
+  def dataset_version(self):
+    # Year of the opensubtitles dataset creation.
+    return 2016
+
+
+@registry.register_problem
+class DialogOpensubtitles64k2018(DialogOpensubtitles64k2009):
+
+  @property
+  def dataset_version(self):
+    # Year of the opensubtitles dataset creation.
+    return 2018
diff --git a/tensor2tensor/data_generators/dialog_personachat.py b/tensor2tensor/data_generators/dialog_personachat.py
new file mode 100644
index 000000000..01a6bf746
--- /dev/null
+++ b/tensor2tensor/data_generators/dialog_personachat.py
@@ -0,0 +1,195 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from collections import Counter
+import os
+import tarfile
+import zipfile
+
+from tensor2tensor.data_generators import dialog_abstract
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+FLAGS = tf.flags.FLAGS
+
+# End-of-sentence marker.
+EOS = text_encoder.EOS_ID
+
+
+@registry.register_problem
+class DialogPersonachat16k(dialog_abstract.DialogAbstract):
+  """Implements a simple chatbot for the original Persona-chat dataset.
+
+  The personas are not used in this class, only the raw dialogs.
+  https://github.com/facebookresearch/ParlAI/tree/master/projects/personachat
+  """
+
+  def preprocess_data(self, train_mode):
+    """Main function where the preprocessing of the data starts.
+
+    Args:
+      train_mode: string, whether we are in train, dev or test mode
+    """
+
+    # Set the raw data directory and data.
+    self.raw_data_dir = os.path.join('/'.join(self._data_dir.split('/')[:-1]),
+                                     'raw_data')
+    self.raw_data = os.path.join(self._raw_data_dir, 'ConvAI2')
+    self.zipped_data = os.path.join(self._raw_data_dir, 'convai2.tar.gz')
+
+    # Create the download url.
+    self.url = 'http://parl.ai/downloads/convai2/convai2_fix_723.tgz'
+
+    # Check at which part of the pipeline are we at.
+    self.data_pipeline_status(train_mode)
+
+  def extract_data(self, train_mode):
+    """Extract data and go to the next step.
+
+    Args:
+      train_mode: string, whether we are in train, dev or test mode
+    """
+
+    if self._zipped_data[-2:] == 'gz':
+      zip_file = tarfile.open(self._zipped_data, 'r:gz')
+    elif self._zipped_data[-3:] == 'zip':
+      zip_file = zipfile.ZipFile(self._zipped_data, 'r')
+    else:
+      print('problem_log: ' + self._zipped_data +
+            ' is not a .zip or .gz file, so I can\'t extract it.')
+
+    zip_file.extractall(self._raw_data)
+    zip_file.close()
+
+    # Next step is creating the source, target and vocab files.
+    print('problem_log: Creating ' +
+          train_mode + ' files in ' + self._data_dir + '.')
+    self.create_data(train_mode)
+
+  def create_data(self, train_mode):
+    """Create the source, target and vocab files.
+
+    Args:
+      train_mode: string, whether we are in train, dev or test mode
+    """
+
+    # Open the 6 files.
+    trainsource, traintarget, devsource, devtarget, testsource, testtarget = \
+        self.open_6_files()
+
+    # Open the raw data.
+    train_dialogs = open(
+        os.path.join(self._raw_data, 'train_none_original_no_cands.txt'),
+        errors='ignore')
+    valid_dialogs = open(
+        os.path.join(self._raw_data, 'valid_none_original_no_cands.txt'),
+        errors='ignore')
+    filenames = [train_dialogs, valid_dialogs]
+
+    # Copy the data to a new file.
+    with open(os.path.join(self._raw_data,
+                           'full_none_original_no_cands.txt'), 'w') as outfile:
+      for fname in filenames:
+        with fname as infile:
+          outfile.write(infile.read())
+    train_dialogs.close()
+    valid_dialogs.close()
+
+    # Open the big file.
+    dialogs = open(
+        os.path.join(self._raw_data, 'full_none_original_no_cands.txt'),
+        errors='ignore')
+
+    number_of_lines = 0
+    current_dialog = ''
+    dialog_list = []
+    dialog_silenced = False
+    # Iterate through the file and build list of dialogs separated by __eou__.
+    for line in dialogs:
+      if number_of_lines % 10000 == 0:
+        print('problem_log: Parsed ' + str(number_of_lines) + ' lines.')
+
+      dialog_id = line.split()[0]
+      # Check if this is a refurbished line.
+      if ('__SILENCE__' not in line and
+          ((dialog_silenced and dialog_id == '1') or not dialog_silenced)):
+        dialog_silenced = False
+        number_of_lines += 1
+
+        # Get the utterances.
+        source = ' '.join(line.split('\t')[0].split()[1:])
+        target = line.split('\t')[1].strip('\n')
+        source = self.clean_line(source.lower())
+        target = self.clean_line(target.lower())
+
+        # Whether this is a new dialog.
+        if dialog_id == '1' and current_dialog:
+          dialog_list.append(current_dialog)
+          current_dialog = source + '__eou__' + target + '__eou__'
+        else:
+          current_dialog += source + '__eou__' + target + '__eou__'
+      else:
+        dialog_silenced = True
+
+      if (self.targeted_dataset_size != 0 and
+          self.targeted_dataset_size < number_of_lines):
+        break
+    dialogs.close()
+
+    vocabulary = Counter()
+    number_of_dialogs = 0
+    dataset_split_counter = 0
+    # Build the dataset.
+    for dialog in dialog_list:
+      if number_of_dialogs % 1000 == 0:
+        print('problem_log: Parsed ' + str(number_of_dialogs) + ' dialogs.')
+
+      # Check which file we should write to.
+      if dataset_split_counter <= self.dataset_split['train']:
+        source_file = trainsource
+        target_file = traintarget
+      elif dataset_split_counter <= (self.dataset_split['train'] +
+                                     self.dataset_split['val']):
+        source_file = devsource
+        target_file = devtarget
+      else:
+        source_file = testsource
+        target_file = testtarget
+
+      utterances = dialog.split('__eou__')[:-1]
+      i = 0
+      # Loop through the dialog.
+      for utterance in utterances:
+        i += 1
+        # Build vocabulary.
+        if dataset_split_counter <= self.dataset_split['train']:
+          words = utterance.split()
+          for word in words:
+            if word in vocabulary:
+              vocabulary[word] += 1
+            else:
+              vocabulary[word] = 1
+
+        # Write to files.
+        if i != len(utterances):
+          source_file.write(utterance + '\n')
+        if i != 1:
+          target_file.write(utterance + '\n')
+
+      dataset_split_counter += 1
+      number_of_dialogs += 1
+      # Reset the split counter if we reached 100%.
+      if dataset_split_counter == 100:
+        dataset_split_counter = 0
+
+    # Close the files.
+    self.close_n_files([trainsource,
+                        traintarget,
+                        devsource,
+                        devtarget,
+                        testsource,
+                        testtarget])
+    # Save the vocabulary.
+    self.save_vocab(vocabulary)

From 19cae7303343a92ebeb37df3398c2c312486302e Mon Sep 17 00:00:00 2001
From: Richard Csaky <ricsinaruto@hotmail.com>
Date: Fri, 13 Sep 2019 09:49:36 -0700
Subject: [PATCH 2435/2720] Merge of PR #1642

PiperOrigin-RevId: 268923712
---
 tensor2tensor/data_generators/all_problems.py |  2 ++
 .../data_generators/dialog_abstract.py        | 21 ++++++++++++++++++-
 .../data_generators/dialog_cornell.py         | 21 +++++++++++++++++--
 .../data_generators/dialog_dailydialog.py     | 21 +++++++++++++++++--
 .../data_generators/dialog_opensubtitles.py   | 21 +++++++++++++++++--
 .../data_generators/dialog_personachat.py     | 21 +++++++++++++++++--
 6 files changed, 98 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index d4dd3f7cb..3efdc9cd3 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -19,6 +19,7 @@
 from __future__ import print_function
 
 import importlib
+import six
 from six.moves import range  # pylint: disable=redefined-builtin
 
 MODULES = [
@@ -107,6 +108,7 @@
 ALL_MODULES = list(MODULES)
 
 
+
 def _is_import_err_msg(err_str, module):
   parts = module.split(".")
   suffixes = [".".join(parts[i:]) for i in range(len(parts))]
diff --git a/tensor2tensor/data_generators/dialog_abstract.py b/tensor2tensor/data_generators/dialog_abstract.py
index 3d8fcd233..16e791e0a 100644
--- a/tensor2tensor/data_generators/dialog_abstract.py
+++ b/tensor2tensor/data_generators/dialog_abstract.py
@@ -1,3 +1,20 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Abstract class for dialog problems."""
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
@@ -23,6 +40,7 @@
 
 # An abstract base class for word based chatbot problems.
 class DialogAbstract(text_problems.Text2TextProblem):
+  """Abstract class for dialog problems."""
 
   @property
   def vocab_type(self):
@@ -302,7 +320,8 @@ def generate_samples(self, data_dir, tmp_dir, data_split):
     Args:
       data_dir: string, Directory where the data will be generated. The raw
                         data has to be downloaded one directory level higher.
-      data_split: stromg, which data split to generate samples for
+      tmp_dir: string, temp directory.
+      data_split: string, which data split to generate samples for
 
     Yields:
       dict
diff --git a/tensor2tensor/data_generators/dialog_cornell.py b/tensor2tensor/data_generators/dialog_cornell.py
index 867b835f3..12363dd36 100644
--- a/tensor2tensor/data_generators/dialog_cornell.py
+++ b/tensor2tensor/data_generators/dialog_cornell.py
@@ -1,8 +1,25 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Cornell Movie Dialog Dataset."""
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from collections import Counter
+import collections
 import os
 import re
 
@@ -64,7 +81,7 @@ def create_data(self, train_mode):
         os.path.join(self._raw_data, 'movie_lines.txt'), errors='ignore')
     dialog_list = self.extract_dialog_ids()
 
-    vocabulary = Counter()
+    vocabulary = collections.Counter()
     line_dict = {}
     number_of_lines = 0
     # Iterate through file.
diff --git a/tensor2tensor/data_generators/dialog_dailydialog.py b/tensor2tensor/data_generators/dialog_dailydialog.py
index 98d04b05f..ea16ff3bb 100644
--- a/tensor2tensor/data_generators/dialog_dailydialog.py
+++ b/tensor2tensor/data_generators/dialog_dailydialog.py
@@ -1,8 +1,25 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""DailyDialog dataset."""
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from collections import Counter
+import collections
 import os
 
 from tensor2tensor.data_generators import dialog_abstract
@@ -57,7 +74,7 @@ def create_data(self, train_mode):
     dialogs = open(
         os.path.join(self._raw_data, 'dialogues_text.txt'), errors='ignore')
 
-    vocabulary = Counter()
+    vocabulary = collections.Counter()
     number_of_dialogs = 0
     line_counter = 0
     dataset_split_counter = 0
diff --git a/tensor2tensor/data_generators/dialog_opensubtitles.py b/tensor2tensor/data_generators/dialog_opensubtitles.py
index 4807553b1..172bf0bb3 100644
--- a/tensor2tensor/data_generators/dialog_opensubtitles.py
+++ b/tensor2tensor/data_generators/dialog_opensubtitles.py
@@ -1,8 +1,25 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""OpenSubtitles dataset."""
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from collections import Counter
+import collections
 import os
 import re
 import zipfile
@@ -88,7 +105,7 @@ def create_data(self, train_mode):
     conv_id = 0
     number_of_lines = 0
     dataset_split_counter = 0
-    vocabulary = Counter()
+    vocabulary = collections.Counter()
     # Dind all the files.
     for root, _, files in os.walk(self._raw_data_dir):
       for f in files:
diff --git a/tensor2tensor/data_generators/dialog_personachat.py b/tensor2tensor/data_generators/dialog_personachat.py
index 01a6bf746..0d848772c 100644
--- a/tensor2tensor/data_generators/dialog_personachat.py
+++ b/tensor2tensor/data_generators/dialog_personachat.py
@@ -1,8 +1,25 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Persona-chat dataset."""
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from collections import Counter
+import collections
 import os
 import tarfile
 import zipfile
@@ -138,7 +155,7 @@ def create_data(self, train_mode):
         break
     dialogs.close()
 
-    vocabulary = Counter()
+    vocabulary = collections.Counter()
     number_of_dialogs = 0
     dataset_split_counter = 0
     # Build the dataset.

From 4806e77ee9d3012831e253b72fd2766b6d9ea706 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 13 Sep 2019 10:41:17 -0700
Subject: [PATCH 2436/2720] Correct the Nesterov momentum optimizer and add
 weight decay.

PiperOrigin-RevId: 268934546
---
 .../trax/configs/wide_resnet_cifar10_8gb.gin          |  1 +
 tensor2tensor/trax/optimizers/base.py                 | 11 ++++++-----
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
index 8c5aab04d..e02837f48 100644
--- a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
+++ b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
@@ -25,6 +25,7 @@ MultifactorSchedule.warmup_steps = 400
 # Parameters for Momentum:
 # ==============================================================================
 Momentum.mass = 0.9
+Momentum.weight_decay_rate = 5e-4
 
 # Parameters for preprocess_fun:
 # ==============================================================================
diff --git a/tensor2tensor/trax/optimizers/base.py b/tensor2tensor/trax/optimizers/base.py
index 21a2cdcae..fd68cd64f 100644
--- a/tensor2tensor/trax/optimizers/base.py
+++ b/tensor2tensor/trax/optimizers/base.py
@@ -169,17 +169,18 @@ def update(self, step, grads, params, slots, opt_params):
 class Momentum(Optimizer):
   """Nesterov momentum optimizer."""
 
-  def __init__(self, learning_rate, mass=0.9):  # pylint: disable=useless-super-delegation
-    super(Momentum, self).__init__(learning_rate, mass)
+  def __init__(self, learning_rate, mass=0.9, weight_decay_rate=1e-5):  # pylint: disable=useless-super-delegation
+    super(Momentum, self).__init__(learning_rate, mass, weight_decay_rate)
 
   def init(self, params):
     return np.zeros_like(params)
 
   def update(self, step, grads, params, velocity, opt_params):
     del step
-    (learning_rate, mass) = opt_params
-    new_velocity = mass * velocity - (1. - mass) * grads
-    new_params = params + (learning_rate * new_velocity).astype(params.dtype)
+    (learning_rate, mass, weight_decay_rate) = opt_params
+    new_velocity = mass * velocity + grads
+    new_params = (1 - weight_decay_rate) * params - (
+        learning_rate * (mass * new_velocity + grads)).astype(params.dtype)
     return (new_params, new_velocity)
 
 

From 5343e728c3441acc20082fd60ddd02d9b6cbdf42 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 13 Sep 2019 16:05:10 -0700
Subject: [PATCH 2437/2720] Make optimizer parameters a dict.

PiperOrigin-RevId: 269000247
---
 tensor2tensor/trax/backend.py         |  2 +
 tensor2tensor/trax/optimizers/base.py | 76 ++++++++++++++++++++-------
 tensor2tensor/trax/trax.py            |  4 +-
 3 files changed, 61 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
index 226d9b553..69997ec24 100644
--- a/tensor2tensor/trax/backend.py
+++ b/tensor2tensor/trax/backend.py
@@ -100,6 +100,8 @@ def nested_map(x, f):
     return [nested_map(y, f) for y in x]
   if isinstance(x, tuple):
     return tuple([nested_map(y, f) for y in x])
+  if isinstance(x, dict):
+    return {k: nested_map(v, f) for (k, v) in x.items()}
   return f(x)
 
 
diff --git a/tensor2tensor/trax/optimizers/base.py b/tensor2tensor/trax/optimizers/base.py
index fd68cd64f..13631a538 100644
--- a/tensor2tensor/trax/optimizers/base.py
+++ b/tensor2tensor/trax/optimizers/base.py
@@ -65,7 +65,7 @@ def tree_unflatten(flat, tree):
 class Optimizer(object):
   """Optimizer object, base class. Maps per-parameter functions to trees."""
 
-  def __init__(self, learning_rate, *init_opt_params):
+  def __init__(self, learning_rate, **init_opt_params):
     """Initialize the optimizer.
 
     Takes the initial optimizer parameters as positional arguments. They are fed
@@ -77,10 +77,12 @@ def __init__(self, learning_rate, *init_opt_params):
 
     Args:
       learning_rate: The initial learning rate.
-      *init_opt_params: Initial values of any additional optimizer parameters.
+      **init_opt_params: Initial values of any additional optimizer parameters.
     """
-    self._init_opt_params = tuple(
-        map(np.array, (learning_rate,) + init_opt_params))
+    init_opt_params["learning_rate"] = learning_rate
+    self._init_opt_params = {
+        name: np.array(value) for (name, value) in init_opt_params.items()
+    }
 
   def init(self, params):
     """Create optimizer slots for the given parameters."""
@@ -162,7 +164,7 @@ def init(self, params):
   def update(self, step, grads, params, slots, opt_params):
     del step
     del slots
-    (learning_rate,) = opt_params
+    learning_rate = opt_params["learning_rate"]
     return params - (learning_rate * grads).astype(params.dtype), None
 
 
@@ -170,14 +172,20 @@ class Momentum(Optimizer):
   """Nesterov momentum optimizer."""
 
   def __init__(self, learning_rate, mass=0.9, weight_decay_rate=1e-5):  # pylint: disable=useless-super-delegation
-    super(Momentum, self).__init__(learning_rate, mass, weight_decay_rate)
+    super(Momentum, self).__init__(
+        learning_rate=learning_rate,
+        mass=mass,
+        weight_decay_rate=weight_decay_rate,
+    )
 
   def init(self, params):
     return np.zeros_like(params)
 
   def update(self, step, grads, params, velocity, opt_params):
     del step
-    (learning_rate, mass, weight_decay_rate) = opt_params
+    learning_rate = opt_params["learning_rate"]
+    mass = opt_params["mass"]
+    weight_decay_rate = opt_params["weight_decay_rate"]
     new_velocity = mass * velocity + grads
     new_params = (1 - weight_decay_rate) * params - (
         learning_rate * (mass * new_velocity + grads)).astype(params.dtype)
@@ -188,14 +196,20 @@ class RMSProp(Optimizer):
   """RMSProp optimizer."""
 
   def __init__(self, learning_rate, gamma=0.9, eps=1e-8):  # pylint: disable=useless-super-delegation
-    super(RMSProp, self).__init__(learning_rate, gamma, eps)
+    super(RMSProp, self).__init__(
+        learning_rate=learning_rate,
+        gamma=gamma,
+        eps=eps,
+    )
 
   def init(self, params):
     return np.ones_like(params)
 
   def update(self, step, grads, params, avg_sq_grad, opt_params):
     del step
-    (learning_rate, gamma, eps) = opt_params
+    learning_rate = opt_params["learning_rate"]
+    gamma = opt_params["gamma"]
+    eps = opt_params["eps"]
     avg_sq_grad = avg_sq_grad * gamma + grads**2 * (1. - gamma)
     params = params - (learning_rate * grads /
                        (np.sqrt(avg_sq_grad) + eps)).astype(params.dtype)
@@ -219,7 +233,13 @@ def __init__(self, learning_rate, weight_decay_rate=1e-5,  # pylint: disable=use
       eps: optional, a positive scalar value for epsilon, a small constant for
         numerical stability (default 1e-8).
     """
-    super(Adam, self).__init__(learning_rate, weight_decay_rate, b1, b2, eps)
+    super(Adam, self).__init__(
+        learning_rate=learning_rate,
+        weight_decay_rate=weight_decay_rate,
+        b1=b1,
+        b2=b2,
+        eps=eps,
+    )
 
   def init(self, params):
     m = np.zeros_like(params)
@@ -228,7 +248,11 @@ def init(self, params):
 
   def update(self, step, grads, params, slots, opt_params):
     m, v = slots
-    learning_rate, weight_decay_rate, b1, b2, eps = opt_params
+    learning_rate = opt_params["learning_rate"]
+    weight_decay_rate = opt_params["weight_decay_rate"]
+    b1 = opt_params["b1"]
+    b2 = opt_params["b2"]
+    eps = opt_params["eps"]
     m = (1 - b1) * grads + b1 * m  # First  moment estimate.
     v = (1 - b2) * (grads ** 2) + b2 * v  # Second moment estimate.
     mhat = m / (1 - b1 ** (step + 1))  # Bias correction.
@@ -281,8 +305,14 @@ def __init__(self,
     self._do_momentum = do_momentum
     # Dynamically configurable parameters will be passed to the update function.
     super(Adafactor, self).__init__(
-        learning_rate, beta1, decay_rate, clipping_threshold,
-        weight_decay_rate, epsilon1, epsilon2)
+        learning_rate=learning_rate,
+        beta1=beta1,
+        decay_rate=decay_rate,
+        clipping_threshold=clipping_threshold,
+        weight_decay_rate=weight_decay_rate,
+        epsilon1=epsilon1,
+        epsilon2=epsilon2,
+    )
 
   @staticmethod
   def _decay_rate_pow(i, exponent=0.8):
@@ -307,8 +337,13 @@ def init(self, params):
 
   def update(self, step, grads, params, slots, opt_params):
     updates = []
-    (learning_rate, beta1, decay_rate, clipping_threshold,
-     weight_decay_rate, epsilon1, epsilon2) = opt_params
+    learning_rate = opt_params["learning_rate"]
+    beta1 = opt_params["beta1"]
+    decay_rate = opt_params["decay_rate"]
+    clipping_threshold = opt_params["clipping_threshold"]
+    weight_decay_rate = opt_params["weight_decay_rate"]
+    epsilon1 = opt_params["epsilon1"]
+    epsilon2 = opt_params["epsilon2"]
     decay_rate = self._decay_rate_pow(step, exponent=decay_rate)
     update_scale = learning_rate
     if self._multiply_by_parameter_scale:
@@ -365,14 +400,18 @@ def __init__(self, learning_rate, momentum=0.9):  # pylint: disable=useless-supe
       learning_rate: a postitive scalar value for the initial learning rate.
       momentum: optional, a positive scalar value for momentum
     """
-    super(SM3, self).__init__(learning_rate, momentum)
+    super(SM3, self).__init__(
+        learning_rate=learning_rate,
+        momentum=momentum,
+    )
 
   def init(self, params):
     vs = [np.zeros(sz, dtype=params.dtype) for sz in params.shape]
     return (np.zeros_like(params), vs)
 
   def _update_diagonal(self, grads, params, m, v, opt_params):
-    (learning_rate, momentum) = opt_params
+    learning_rate = opt_params["learning_rate"]
+    momentum = opt_params["momentum"]
     v[0] += grads * grads
     preconditioner = np.where(v[0] > 0, 1.0 / np.sqrt(v[0]),
                               np.zeros_like(v[0]))
@@ -395,7 +434,8 @@ def _minimum(self, tensor_list):
 
   def _update_sketched(self, grads, params, m, v, opt_params):
     """Update for higher-rank parameters."""
-    (learning_rate, momentum) = opt_params
+    learning_rate = opt_params["learning_rate"]
+    momentum = opt_params["momentum"]
     shape = params.shape
     rank = len(shape)
     reshaped_accumulators = [np.reshape(v[i], self._expanded_shape(shape, i))
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 1cf009140..51a8041ac 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -722,9 +722,7 @@ def _train_step(self, next_train_batch):
     # Calculate the current learning rate.
     learning_rate = self._maybe_replicate(np.array(self.learning_rate))
     opt_state = self._opt_state
-    opt_params = opt_state.opt_params
-    opt_params = (learning_rate,) + opt_params[1:]
-    opt_state = opt_state._replace(opt_params=opt_params)
+    opt_state.opt_params["learning_rate"] = learning_rate
 
     # Run the update.
     (params, slots), self._model_state, self._rngs = self._jit_update_fn(

From 5fc3dc1527d0c524526c29363f655dc2be4f9646 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 13 Sep 2019 17:30:59 -0700
Subject: [PATCH 2438/2720] Return is missing in one of the reward models.

PiperOrigin-RevId: 269014010
---
 tensor2tensor/models/video/sv2p.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 551dd6171..ea8f26e15 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -215,6 +215,7 @@ def reward_prediction_big(
       x = tfcl.layer_norm(x)
       x = tfl.conv2d(x, conv_size[3], [3, 3], strides=(2, 2),
                      activation=tf.nn.relu, name="reward_conv3")
+    return x
 
   def get_extra_loss(self,
                      latent_means=None, latent_stds=None,

From ac433ea406f2fa2b17329994f0b2e67b02819c05 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 13 Sep 2019 18:23:09 -0700
Subject: [PATCH 2439/2720] Internal change

PiperOrigin-RevId: 269020489
---
 tensor2tensor/envs/env_problem.py             |   5 +
 tensor2tensor/envs/env_problem_utils.py       |  30 ++-
 tensor2tensor/envs/trajectory.py              |  23 ++-
 tensor2tensor/envs/trajectory_test.py         |   2 +-
 .../trax/rl/configs/ppo_online_tune.gin       |   5 -
 .../rl/envs/async_trajectory_collector.py     | 192 +++++++++++++++++
 .../rl/envs/async_trajectory_collector_lib.py | 195 ++++++++++++++++++
 .../async_trajectory_collector_lib_test.py    |  64 ++++++
 tensor2tensor/trax/rl/ppo.py                  |  20 +-
 tensor2tensor/trax/rl/ppo_trainer.py          |  74 ++++---
 tensor2tensor/trax/rl_trainer.py              |   5 +-
 11 files changed, 565 insertions(+), 50 deletions(-)
 create mode 100644 tensor2tensor/trax/rl/envs/async_trajectory_collector.py
 create mode 100644 tensor2tensor/trax/rl/envs/async_trajectory_collector_lib.py
 create mode 100644 tensor2tensor/trax/rl/envs/async_trajectory_collector_lib_test.py

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 0f10b1b78..43afb1a75 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -151,6 +151,11 @@ def batch_size(self):
   def trajectories(self):
     return self._trajectories
 
+  @trajectories.setter
+  def trajectories(self, trajectories_):
+    assert self.trajectories.batch_size == trajectories_.batch_size
+    self._trajectories = trajectories_
+
   def initialize(self, batch_size=1, **kwargs):
     self.initialize_environments(batch_size=batch_size, **kwargs)
 
diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index dd3795765..6fbfdc6a3 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -54,7 +54,16 @@ def play_env_problem_randomly(env_problem, num_steps):
     env_problem.reset(indices=done_indices(dones))
 
 
-def get_completed_trajectories_from_env(env, n_trajectories):
+def get_completed_trajectories_from_env(env,
+                                        n_trajectories,
+                                        raw_trajectory=False):
+  """Returns completed `n_trajectories` from `env`."""
+
+  # Just the raw trajectories.
+  if raw_trajectory:
+    return env.trajectories.completed_trajectories[:n_trajectories]
+
+  # The numpy version of the above.
   completed_trajectories = []
   for trajectory in env.trajectories.completed_trajectories[:n_trajectories]:
     completed_trajectories.append(trajectory.as_numpy)
@@ -71,7 +80,9 @@ def play_env_problem_with_policy(env,
                                  temperature=1.0,
                                  boundary=32,
                                  len_history_for_policy=32,
-                                 num_to_keep=1):
+                                 num_to_keep=1,
+                                 abort_fn=None,
+                                 raw_trajectory=False):
   """Plays the given env with the policy function to collect trajectories.
 
   Args:
@@ -91,6 +102,12 @@ def play_env_problem_with_policy(env,
     len_history_for_policy: int or None, the maximum history to keep for
       applying the policy on. If None, use the whole history.
     num_to_keep: int, while truncating trajectory how many time-steps to keep.
+    abort_fn: callable, If not None, then at every step call and abort the
+      trajectory collection if it returns True, if so reset the env and return
+      None.
+    raw_trajectory: bool, if True a list of trajectory.Trajectory objects is
+      returned, otherwise a list of numpy representations of
+      `trajectory.Trajectory` is returned.
 
   Returns:
     A tuple, (trajectories, number of completed trajectories). Where
@@ -116,6 +133,13 @@ def gumbel_sample(log_probs):
   env_actions_total_time = 0
   bare_env_run_time = 0
   while env.trajectories.num_completed_trajectories < num_trajectories:
+    # Check if we should abort and return nothing.
+    if abort_fn and abort_fn():
+      # We should also reset the environment, since it will have some
+      # trajectories (complete and incomplete) that we want to discard.
+      env.reset()
+      return None, 0, {}, state
+
     # Get all the observations for all the active trajectories.
     # Shape is (B, T) + OBS
     # Bucket on whatever length is needed.
@@ -197,7 +221,7 @@ def gumbel_sample(log_probs):
   # We have the trajectories we need, return a list of triples:
   # (observations, actions, rewards)
   completed_trajectories = get_completed_trajectories_from_env(
-      env, num_trajectories)
+      env, num_trajectories, raw_trajectory=raw_trajectory)
 
   timing_info = {
       "trajectory_collection/policy_application": policy_application_total_time,
diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index e3352433d..cf4705a1c 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -36,7 +36,7 @@
 TRAJECTORY_FILE_FORMAT = r"trajectory_epoch_{epoch}_env_id_{env_id}_temperature_{temperature}_r_{r}.pkl"
 
 
-def _get_pickle_module():
+def get_pickle_module():
   if sys.version_info[0] < 3:
     return cloudpickle
   return pickle
@@ -494,8 +494,9 @@ def load_from_directory(trajectory_dir,
                           temperature=None,
                           n_trajectories=None,
                           up_sample=False,
-                          wait_time_secs=0.1,
-                          max_tries=17):
+                          sleep_time_secs=0.1,
+                          max_tries=100,
+                          wait_forever=False):
     """Load trajectories from specified dir and epoch.
 
     Args:
@@ -510,9 +511,10 @@ def load_from_directory(trajectory_dir,
         then we wait for those many trajectory files to be available.
       up_sample: (bool) If there are fewer than required (n_trajectories) number
         of incomplete trajectories, then we upsample to make up the numbers.
-      wait_time_secs: (float) Waiting time, with exponential backoff to wait for
-        min_trajectories.
+      sleep_time_secs: (float) Sleep time, to wait for min_trajectories. We
+        exponentially back-off this up till a maximum of 10 seconds.
       max_tries: (int) The number of tries to get min_trajectories trajectories.
+      wait_forever: (bool) If true, overrides max_tries and waits forever.
 
     Returns:
       A BatchTrajectory object with all the constraints satisfied or None.
@@ -532,14 +534,15 @@ def load_from_directory(trajectory_dir,
     if n_trajectories:
       # We need to get `n_trajectories` number of `trajectory_files`.
       # This works out to a maximum ~3hr waiting period.
-      while max_tries > 0 and len(trajectory_files) < n_trajectories:
+      while ((max_tries > 0 or wait_forever) and
+             len(trajectory_files) < n_trajectories):
         logging.info(
             "Sleeping for %s seconds while waiting for %s trajectories, found "
-            "%s right now.", wait_time_secs, n_trajectories,
+            "%s right now.", sleep_time_secs, n_trajectories,
             len(trajectory_files))
-        time.sleep(wait_time_secs)
+        time.sleep(sleep_time_secs)
         max_tries -= 1
-        wait_time_secs *= 2  # exponential backoff.
+        sleep_time_secs = min(10.0, sleep_time_secs * 2)
         trajectory_files = gfile.glob(
             os.path.join(trajectory_dir, trajectory_file_glob))
 
@@ -555,7 +558,7 @@ def load_from_directory(trajectory_dir,
     trajectories_buffer = []
     for trajectory_file in trajectory_files:
       with gfile.GFile(trajectory_file, "rb") as f:
-        trajectory = _get_pickle_module().load(f)
+        trajectory = get_pickle_module().load(f)
         assert isinstance(trajectory, Trajectory)
         trajectories_buffer.append(trajectory)
 
diff --git a/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index 40ceab7f9..6bc3dc442 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -516,7 +516,7 @@ def test_load_from_directory(self):
 
             with gfile.GFile(
                 os.path.join(output_dir, trajectory_file_name), "w") as f:
-              trajectory._get_pickle_module().dump(traj, f)
+              trajectory.get_pickle_module().dump(traj, f)
 
     # Load everything and check.
     bt = trajectory.BatchTrajectory.load_from_directory(output_dir)
diff --git a/tensor2tensor/trax/rl/configs/ppo_online_tune.gin b/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
index edc1dc6a2..23a656769 100644
--- a/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
+++ b/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
@@ -35,8 +35,3 @@ PPO.separate_eval = False
 PPO.save_every_n = 1
 PPO.policy_and_value_model = @trax.models.TransformerDecoder
 PPO.policy_and_value_optimizer = @trax.optimizers.Adam
-
-# Parameters for train_rl:
-# ==============================================================================
-train_rl.env_name = "ClientEnv-v0"
-train_rl.n_epochs = 1000
diff --git a/tensor2tensor/trax/rl/envs/async_trajectory_collector.py b/tensor2tensor/trax/rl/envs/async_trajectory_collector.py
new file mode 100644
index 000000000..f94e7162a
--- /dev/null
+++ b/tensor2tensor/trax/rl/envs/async_trajectory_collector.py
@@ -0,0 +1,192 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""A trajectory collector that polls on policy files and keeps collecting trajectories."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import multiprocessing
+import os
+from absl import app
+from absl import flags
+from absl import logging
+import gin
+import jax
+from jax.config import config
+from tensor2tensor.envs import env_problem_utils
+from tensor2tensor.rl.google import atari_utils  # GOOGLE-INTERNAL:
+from tensor2tensor.trax.rl.envs import async_trajectory_collector_lib as async_lib
+import tensorflow as tf
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_multi_string("config_file", None,
+                          "Configuration file with parameters (.gin).")
+flags.DEFINE_multi_string("config", None,
+                          "Configuration parameters (gin string).")
+flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
+flags.DEFINE_bool("xm", False, "Copy atari roms?")
+
+flags.DEFINE_bool(
+    "try_abort", True,
+    "Should we try to abort a trajectory collection if a newer "
+    "policy is available.")
+
+flags.DEFINE_string("output_dir", "", "Output dir.")
+
+flags.DEFINE_boolean(
+    "jax_debug_nans", False,
+    "Setting to true will help to debug nans and disable jit.")
+flags.DEFINE_boolean("disable_jit", False, "Setting to true will disable jit.")
+
+flags.DEFINE_boolean("parallelize_envs", False,
+                     "If true, sets parallelism to number of cpu cores.")
+
+flags.DEFINE_integer("replica", 0, "Basically to append to trajectory name.")
+flags.DEFINE_bool("enable_eager_execution", False, "")
+
+flags.DEFINE_integer(
+    "max_trajectories_to_collect", -1,
+    "-1 for infinite, otherwise whatever number was specified.")
+
+
+# TODO(afrozm): This code snippet is strewn across many places, unify it.
+def initialize_gin():
+  gin_configs = FLAGS.config or []
+  gin.parse_config_files_and_bindings(FLAGS.config_file, gin_configs)
+
+
+def get_output_dir():
+  """Return output_dir."""
+  output_dir = FLAGS.output_dir
+  return output_dir
+
+
+def update_jax_config():
+  """Update JAX config based on flags."""
+
+  if FLAGS.jax_debug_nans:
+    config.update("jax_debug_nans", True)
+
+  if FLAGS.use_tpu:
+    config.update("jax_platform_name", "tpu")
+  else:
+    config.update("jax_platform_name", "gpu")
+
+
+@gin.configurable(blacklist=[
+    "output_dir",
+])
+def create_envs_and_collect_trajectories(
+    output_dir,
+    env_name="OnlineTuneEnv-v0",
+    max_timestep=None,
+    clip_rewards=False,
+    rendered_env=False,
+    resize_dims=(105, 80),
+):
+  """Creates the envs and continuously collects trajectories."""
+
+
+  train_batch_size = 1
+  eval_batch_size = 1
+
+  # TODO(pkozakowski): Find a better way to determine this.
+  train_env_kwargs = {}
+  eval_env_kwargs = {}
+  if "OnlineTuneEnv" in env_name:
+    # TODO(pkozakowski): Separate env output dirs by train/eval and epoch.
+    train_env_kwargs = {"output_dir": os.path.join(output_dir, "envs/train")}
+    eval_env_kwargs = {"output_dir": os.path.join(output_dir, "envs/eval")}
+
+  if "ClientEnv" in env_name:
+    train_env_kwargs["per_env_kwargs"] = [{
+        "remote_env_address": os.path.join(FLAGS.train_server_bns, str(replica))
+    } for replica in range(train_batch_size)]
+
+    eval_env_kwargs["per_env_kwargs"] = [{
+        "remote_env_address": os.path.join(FLAGS.eval_server_bns, str(replica))
+    } for replica in range(eval_batch_size)]
+
+  parallelism = multiprocessing.cpu_count() if FLAGS.parallelize_envs else 1
+  train_parallelism = min(train_batch_size, parallelism)
+  eval_parallelism = min(eval_batch_size, parallelism)
+
+  train_env = env_problem_utils.make_env(
+      batch_size=train_batch_size,
+      env_problem_name=env_name,
+      resize=rendered_env,
+      resize_dims=resize_dims,
+      max_timestep=max_timestep,
+      clip_rewards=clip_rewards,
+      parallelism=train_parallelism,
+      use_tpu=FLAGS.use_tpu,
+      **train_env_kwargs)
+  assert train_env
+
+  eval_env = env_problem_utils.make_env(
+      batch_size=eval_batch_size,
+      env_problem_name=env_name,
+      resize=rendered_env,
+      resize_dims=resize_dims,
+      max_timestep=max_timestep,
+      clip_rewards=clip_rewards,
+      parallelism=eval_parallelism,
+      use_tpu=FLAGS.use_tpu,
+      **eval_env_kwargs)
+  assert eval_env
+
+  def run_collect_loop():
+    async_lib.continuously_collect_trajectories(
+        output_dir,
+        train_env,
+        eval_env,
+        trajectory_dump_dir=None,
+        env_id=FLAGS.replica,
+        try_abort=FLAGS.try_abort,
+        max_trajectories_to_collect=(None
+                                     if FLAGS.max_trajectories_to_collect < 0
+                                     else FLAGS.max_trajectories_to_collect))
+
+  if FLAGS.jax_debug_nans or FLAGS.disable_jit:
+    with jax.disable_jit():
+      run_collect_loop()
+  else:
+    run_collect_loop()
+
+
+def main(argv):
+  del argv
+
+  if FLAGS.enable_eager_execution:
+    tf.enable_eager_execution()
+
+  logging.info("Initializing Gin.")
+  initialize_gin()
+
+  logging.info("Update JAX config.")
+  update_jax_config()
+
+  logging.info("Getting output_dir")
+  output_dir = get_output_dir()
+  logging.info("Got output_dir = %s", output_dir)
+
+  logging.info("Starting Trajectory collection.")
+  create_envs_and_collect_trajectories(output_dir)
+
+
+if __name__ == "__main__":
+  app.run(main)
diff --git a/tensor2tensor/trax/rl/envs/async_trajectory_collector_lib.py b/tensor2tensor/trax/rl/envs/async_trajectory_collector_lib.py
new file mode 100644
index 000000000..447a53819
--- /dev/null
+++ b/tensor2tensor/trax/rl/envs/async_trajectory_collector_lib.py
@@ -0,0 +1,195 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Uitlity functions for the async trajectory collector."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import random
+import time
+
+from absl import logging
+from tensor2tensor.envs import trajectory
+from tensor2tensor.trax.rl import ppo
+from tensor2tensor.trax.rl import trainers as rl_trainers
+from tensorflow.io import gfile
+
+LARGE_MAX_TRIES_FOR_POLICY_FILE = 100
+
+
+# TODO(afrozm): Is there a better way to poll for a file on CNS?
+def get_newer_policy_model_file(output_dir,
+                                min_epoch=-1,
+                                sleep_time_secs=0.1,
+                                max_sleep_time_secs=1.0,
+                                max_tries=1,
+                                wait_forever=False,):
+  """Gets a policy model file subject to availability and wait time."""
+
+  while max_tries or wait_forever:
+    max_tries -= 1
+    policy_files = ppo.get_policy_model_files(output_dir)
+
+    def do_wait(t):
+      time.sleep(t)
+      t *= 2
+      return min(t, max_sleep_time_secs)
+
+    # No policy files at all.
+    if not policy_files:
+      logging.info("There are no policy files in [%s], waiting for %s secs.",
+                   output_dir, sleep_time_secs)
+      sleep_time_secs = do_wait(sleep_time_secs)
+      continue
+
+    # Check if we have a newer epoch.
+    policy_file = policy_files[0]
+    epoch = ppo.get_epoch_from_policy_model_file(policy_file)
+
+    # We don't - wait.
+    if epoch <= min_epoch:
+      logging.info("epoch [%s] <= min_epoch [%s], waiting for %s secs.", epoch,
+                   min_epoch, sleep_time_secs)
+      sleep_time_secs = do_wait(sleep_time_secs)
+      continue
+
+    # We do have a new file, return it.
+    policy_file = policy_files[0]
+    epoch = ppo.get_epoch_from_policy_model_file(policy_file)
+    logging.info("Found epoch [%s] and policy file [%s]", epoch, policy_file)
+    return policy_file, epoch
+
+  # Exhausted our waiting limit.
+  return None
+
+
+def dump_trajectory(output_dir, epoch, env_id, temperature, random_string,
+                    trajs):
+  """Write the trajectory to disk."""
+
+  assert 1 == len(trajs)
+  traj = trajs[0]
+
+  trajectory_file_name = trajectory.TRAJECTORY_FILE_FORMAT.format(
+      epoch=epoch, env_id=env_id, temperature=temperature, r=random_string)
+
+  with gfile.GFile(os.path.join(output_dir, trajectory_file_name), "w") as f:
+    trajectory.get_pickle_module().dump(traj, f)
+
+
+def continuously_collect_trajectories(output_dir,
+                                      train_env,
+                                      eval_env,
+                                      trajectory_dump_dir=None,
+                                      env_id=None,
+                                      max_trajectories_to_collect=None,
+                                      try_abort=True):
+  """Instantiates a PPO trainer and collects trajectories."""
+
+  # Make the PPO trainer.
+  ppo_trainer = rl_trainers.PPO(
+      output_dir=output_dir,
+      train_env=train_env,
+      eval_env=eval_env,
+      trajectory_dump_dir=trajectory_dump_dir,
+  )
+
+  # TODO(afrozm): Update base_trainer interface to support SimPLe as well.
+  assert isinstance(ppo_trainer, rl_trainers.PPO)
+
+  assert env_id is not None
+
+  # Get an initial policy and wait a forever to get it if needed.
+  policy_and_epoch = get_newer_policy_model_file(output_dir, wait_forever=True)
+  assert policy_and_epoch
+  policy_file, epoch = policy_and_epoch
+  logging.info("Read initial policy for epoch [%s] -> [%s]", epoch, policy_file)
+
+  # Returns immediately if there is a newer epoch available.
+  def is_newer_policy_file_available(epoch_, sleep_time_secs_=0.1):
+    return get_newer_policy_model_file(
+        output_dir, min_epoch=epoch_, sleep_time_secs=sleep_time_secs_)
+
+  assert 1 == train_env.batch_size
+  assert 1 == eval_env.batch_size
+
+  temperature = 1.0
+
+  trajectories_collected = 0
+
+  train_env_trajectory_dump_dir = os.path.join(output_dir, "trajectories/train")
+  eval_env_trajectory_dump_dir = os.path.join(output_dir, "trajectories/eval")
+
+  gfile.makedirs(train_env_trajectory_dump_dir)
+  gfile.makedirs(eval_env_trajectory_dump_dir)
+
+  while max_trajectories_to_collect is None or trajectories_collected < int(
+      max_trajectories_to_collect):
+    logging.info("Collecting a trajectory, trajectories_collected = %s",
+                 trajectories_collected)
+
+    # Abort function -- if something newever is available, then abort the
+    # current computation and reload.
+
+    # Useful if env.step is long.
+    def long_abort_fn():
+      # We want this to be as quick as possible.
+      return is_newer_policy_file_available(epoch, 0) is not None
+
+    abort_fn = long_abort_fn if try_abort else None
+
+    # Collect a training trajectory.
+    trajs, n_done, unused_timing_info, unused_model_state = (
+        ppo_trainer.collect_trajectories(train=True,
+                                         temperature=temperature,
+                                         abort_fn=abort_fn,
+                                         raw_trajectory=True))
+
+    if trajs and n_done > 0:
+      assert 1 == n_done
+      trajectories_collected += n_done
+
+      # Write the trajectory down.
+      logging.info(
+          "Dumping the collected trajectory, trajectories_collected = %s",
+          trajectories_collected)
+      dump_trajectory(train_env_trajectory_dump_dir, epoch, env_id, temperature,
+                      str(random.randint(0, 2**31 - 1)), trajs)
+    else:
+      logging.info("Computation was aborted, a new policy is available.")
+
+    # This maybe useless, since `abort_fn` will take care of it. We might want
+    # to have this here if abort_fn is False always.
+    # Do we have a newer policy?
+    policy_file_and_epoch = is_newer_policy_file_available(epoch)
+    if policy_file_and_epoch is None:
+      # Continue churning out these policies.
+      logging.info("We don't have a newer policy, continuing with the old one.")
+      continue
+
+    # We have a newer policy, read it and update the parameters.
+    policy_file, epoch = policy_file_and_epoch
+    logging.info(
+        "We have a newer policy epoch [%s], file [%s], updating parameters.",
+        epoch, policy_file)
+    ppo_trainer.update_optimization_state(
+        output_dir, policy_and_value_opt_state=None)
+    logging.info("Parameters of PPOTrainer updated.")
+
+    # Check that the epochs match.
+    assert epoch == ppo_trainer.epoch
diff --git a/tensor2tensor/trax/rl/envs/async_trajectory_collector_lib_test.py b/tensor2tensor/trax/rl/envs/async_trajectory_collector_lib_test.py
new file mode 100644
index 000000000..39ba95614
--- /dev/null
+++ b/tensor2tensor/trax/rl/envs/async_trajectory_collector_lib_test.py
@@ -0,0 +1,64 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from tensor2tensor.trax.rl import ppo
+from tensor2tensor.trax.rl.envs import async_trajectory_collector_lib as async_lib
+from tensorflow import test
+from tensorflow.io import gfile
+
+
+class AsyncTrajectoryCollectorLibTest(test.TestCase):
+
+  def test_get_newer_policy_model_file(self):
+    output_dir = self.get_temp_dir()
+
+    def write_policy_model_file(epoch):
+      fname = ppo.get_policy_model_file_from_epoch(output_dir, epoch)
+      with gfile.GFile(fname, "w") as f:
+        f.write("some data")
+      return fname
+
+    # No file exists currently.
+    self.assertIsNone(async_lib.get_newer_policy_model_file(output_dir))
+
+    # Write a policy model file.
+    epoch = 0
+    policy_model_filename = write_policy_model_file(epoch)
+
+    # See that we get it.
+    actual_policy_file, actual_epoch = (
+        async_lib.get_newer_policy_model_file(output_dir, min_epoch=-1))
+
+    self.assertEqual(actual_policy_file, policy_model_filename)
+    self.assertEqual(actual_epoch, epoch)
+
+    # If we now ask for a larger epoch, we don't get it.
+    self.assertIsNone(
+        async_lib.get_newer_policy_model_file(output_dir, min_epoch=0))
+
+    # Write a newer epoch and expect to get that with appropriate min_epoch.
+    epoch = 1
+    policy_model_filename = write_policy_model_file(epoch)
+    actual_policy_file, actual_epoch = (
+        async_lib.get_newer_policy_model_file(output_dir, min_epoch=0))
+    self.assertEqual(actual_policy_file, policy_model_filename)
+    self.assertEqual(actual_epoch, epoch)
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index 175a445d6..cd0fad361 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -155,7 +155,9 @@ def collect_trajectories(env,
                          boundary=32,
                          state=None,
                          temperature=1.0,
-                         rng=None):
+                         rng=None,
+                         abort_fn=None,
+                         raw_trajectory=False,):
   """Collect trajectories with the given policy net and behaviour.
 
   Args:
@@ -173,6 +175,12 @@ def collect_trajectories(env,
     state: state for `policy_fn`.
     temperature: (float) temperature to sample action from policy_fn.
     rng: jax rng, splittable.
+    abort_fn: callable, If not None, then at every env step call and abort the
+      trajectory collection if it returns True, if so reset the env and return
+      None.
+    raw_trajectory: bool, if True a list of trajectory.Trajectory objects is
+      returned, otherwise a list of numpy representations of
+      `trajectory.Trajectory` is returned.
 
   Returns:
     A tuple (trajectory, number of trajectories that are done)
@@ -195,7 +203,10 @@ def collect_trajectories(env,
       boundary=boundary,
       state=state,
       temperature=temperature,
-      rng=rng)
+      rng=rng,
+      abort_fn=abort_fn,
+      raw_trajectory=raw_trajectory,
+  )
   # Skip returning raw_rewards here, since they aren't used.
 
   # t is the return value of Trajectory.as_numpy, so:
@@ -855,8 +866,9 @@ def save_opt_state(output_dir,
   with gfile.GFile(params_file, "wb") as f:
     pkl_module.dump(
         (policy_and_value_opt_state, policy_and_value_state, total_opt_step), f)
-  # Remove the old model files.
-  for path in old_model_files:
+  # Remove the old model files, leave the latest one (it might be in the
+  # process of getting read async) -- this will get cleaned up later.
+  for path in old_model_files[1:]:
     if path != params_file:
       gfile.remove(path)
 
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index e31d9f9f8..820b81120 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -180,14 +180,12 @@ def __init__(self,
      self._policy_and_value_get_params) = ppo.optimizer_fn(
          policy_and_value_optimizer, policy_and_value_net_params)
 
-    # Maybe restore the optimization state. If there is nothing to restore, then
-    # epoch = 0 and policy_and_value_opt_state is returned as is.
-    (self._policy_and_value_opt_state, self._model_state, self._epoch,
-     self._total_opt_step) = ppo.maybe_restore_opt_state(
-         output_dir, policy_and_value_opt_state, self._model_state)
-
-    if self._epoch > 0:
-      logging.info("Restored parameters from epoch [%d]", self._epoch)
+    # Restore the optimizer state.
+    self._policy_and_value_opt_state = policy_and_value_opt_state
+    self._epoch = 0
+    self._total_opt_step = 0
+    self.update_optimization_state(
+        output_dir, policy_and_value_opt_state=policy_and_value_opt_state)
 
     # Create summary writers and history.
     self._train_sw = jaxboard.SummaryWriter(
@@ -200,6 +198,21 @@ def __init__(self,
     self._n_trajectories_done = 0
 
     self._last_saved_at = 0
+    if self._async_mode:
+      logging.info("Saving model on startup to have a model policy file.")
+      self.save()
+
+  # Maybe restore the optimization state. If there is nothing to restore, then
+  # epoch = 0 and policy_and_value_opt_state is returned as is.
+  def update_optimization_state(self,
+                                output_dir,
+                                policy_and_value_opt_state=None):
+    (self._policy_and_value_opt_state, self._model_state, self._epoch,
+     self._total_opt_step) = ppo.maybe_restore_opt_state(
+         output_dir, policy_and_value_opt_state, self._model_state)
+
+    if self._epoch > 0:
+      logging.info("Restored parameters from epoch [%d]", self._epoch)
 
   @property
   def train_env(self):
@@ -235,6 +248,7 @@ def collect_trajectories_async(self,
 
     assert self._async_mode
 
+    # trajectories/train and trajectories/eval are the two subdirectories.
     trajectory_dir = os.path.join(self._output_dir, "trajectories",
                                   "train" if train else "eval")
     epoch = self.epoch
@@ -266,7 +280,11 @@ def collect_trajectories_async(self,
     timing_info = {}
     return trajs, n_done, timing_info, self._model_state
 
-  def collect_trajectories(self, train=True, temperature=1.0):
+  def collect_trajectories(self,
+                           train=True,
+                           temperature=1.0,
+                           abort_fn=None,
+                           raw_trajectory=False):
     self._rng, key = jax_random.split(self._rng)
 
     env = self.train_env
@@ -281,24 +299,29 @@ def collect_trajectories(self, train=True, temperature=1.0):
 
     # If async, read the required trajectories for the epoch.
     if self._async_mode:
-      return self.collect_trajectories_async(
+      trajs, n_done, timing_info, self._model_state = self.collect_trajectories_async(
           env,
           train=train,
           n_trajectories=n_trajectories,
           temperature=temperature)
-
-    trajs, n_done, timing_info, self._model_state = ppo.collect_trajectories(
-        env,
-        policy_fn=self._get_predictions,
-        n_trajectories=n_trajectories,
-        max_timestep=max_timestep,
-        state=self._model_state,
-        rng=key,
-        len_history_for_policy=self._len_history_for_policy,
-        boundary=self._boundary,
-        reset=should_reset,
-        temperature=temperature,
-    )
+    else:
+      trajs, n_done, timing_info, self._model_state = ppo.collect_trajectories(
+          env,
+          policy_fn=self._get_predictions,
+          n_trajectories=n_trajectories,
+          max_timestep=max_timestep,
+          state=self._model_state,
+          rng=key,
+          len_history_for_policy=self._len_history_for_policy,
+          boundary=self._boundary,
+          reset=should_reset,
+          temperature=temperature,
+          abort_fn=abort_fn,
+          raw_trajectory=raw_trajectory,
+      )
+
+    if train:
+      self._n_trajectories_done += n_done
 
     return trajs, n_done, timing_info, self._model_state
 
@@ -317,7 +340,7 @@ def train_epoch(self, evaluate=True):
     trajectory_collection_start_time = time.time()
     logging.vlog(1, "PPO epoch [% 6d]: collecting trajectories.", self._epoch)
     self._rng, key = jax_random.split(self._rng)
-    trajs, n_done, timing_info, self._model_state = self.collect_trajectories(
+    trajs, _, timing_info, self._model_state = self.collect_trajectories(
         train=True, temperature=1.0)
     trajs = [(t[0], t[1], t[2], t[4]) for t in trajs]
     self._should_reset = False
@@ -557,11 +580,10 @@ def train_epoch(self, evaluate=True):
     # truncated and done).
     # Also don't save too frequently, enforce a minimum gap.
     policy_save_start_time = time.time()
-    self._n_trajectories_done += n_done
     # TODO(afrozm): Refactor to trax.save_state.
     if (self._n_trajectories_done >=
         self._done_frac_for_policy_save * self.train_env.batch_size and
-        self._epoch % self._save_every_n == 0):
+        self._epoch % self._save_every_n == 0) or self._async_mode:
       self.save()
     policy_save_time = ppo.get_time(policy_save_start_time)
 
diff --git a/tensor2tensor/trax/rl_trainer.py b/tensor2tensor/trax/rl_trainer.py
index 04de7d56a..f38888b83 100644
--- a/tensor2tensor/trax/rl_trainer.py
+++ b/tensor2tensor/trax/rl_trainer.py
@@ -75,6 +75,8 @@
 flags.DEFINE_string("train_server_bns", "", "Train Server's BNS.")
 flags.DEFINE_string("eval_server_bns", "", "Eval Server's BNS.")
 
+flags.DEFINE_bool("async_mode", False, "Async mode.")
+
 
 # Not just "train" to avoid a conflict with trax.train in GIN files.
 @gin.configurable(blacklist=[
@@ -84,7 +86,7 @@ def train_rl(
     output_dir,
     train_batch_size,
     eval_batch_size,
-    env_name="Acrobot-v1",
+    env_name="ClientEnv-v0",
     max_timestep=None,
     clip_rewards=False,
     rendered_env=False,
@@ -174,6 +176,7 @@ def run_training_loop():
         train_env=train_env,
         eval_env=eval_env,
         trajectory_dump_dir=trajectory_dump_dir,
+        async_mode=FLAGS.async_mode,
     )
     trainer.training_loop(n_epochs=n_epochs)
 

From 5e50f36e1a4a0f4efb7effeec541e516bee00852 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Sat, 14 Sep 2019 00:13:11 -0700
Subject: [PATCH 2440/2720] Extend OnlineTuneEnv to control arbitrary optimizer
 parameters.

As a side-effect, other optimizer parameters are also stored in History and in TensorBoard. This will be convenient for monitoring. Learning rate schedules can now also control other optimizer parameters - they return a dict with overrides. Haven't updated PolicySchedule to reflect that yet, this will come next.

PiperOrigin-RevId: 269050937
---
 tensor2tensor/trax/learning_rate.py           |  29 +++--
 tensor2tensor/trax/learning_rate_test.py      |   8 +-
 .../env_online_tune_wide_resnet_cifar10.gin   |   7 +-
 .../trax/rl/configs/ppo_online_tune.gin       |   2 +-
 tensor2tensor/trax/rl/envs/online_tune_env.py |  56 +++++---
 .../trax/rl/envs/online_tune_env_test.py      |  46 +++++--
 tensor2tensor/trax/rl/online_tune.py          |  94 +++++++++++---
 tensor2tensor/trax/rl/online_tune_test.py     | 121 ++++++++++++++----
 tensor2tensor/trax/trax.py                    |  28 ++--
 9 files changed, 284 insertions(+), 107 deletions(-)

diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index a10e27863..279a2bece 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -16,10 +16,10 @@
 """trax learning rate schedules.
 
 The learning rate schedules here all have the signature:
-  lr: history -> (step -> lr)
+  lr: history -> (step -> {"learning_rate": lr})
 
 That is, they are functions that take a trax.history.History and return a
-function that takes a step and returns a learning rate.
+function that takes a step and returns a dict with entry "learning_rate".
 """
 
 from __future__ import absolute_import
@@ -64,7 +64,8 @@ def MultifactorSchedule(history=None,
     steps_per_decay: How often to decay the learning rate.
 
   Returns:
-    a function learning_rate(step): float -> float, the step-dependent lr.
+    a function learning_rate(step): float -> {"learning_rate": float}, the
+    step-dependent lr.
   """
   del history
 
@@ -84,7 +85,7 @@ def learning_rate(step):  # pylint: disable=invalid-name
         ret *= (decay_factor ** (step//steps_per_decay))
       else:
         raise ValueError("Unknown factor %s." % name)
-    return ret
+    return {"learning_rate": ret}
 
   return learning_rate
 
@@ -115,7 +116,8 @@ def EvalAdjustingSchedule(history,
     metric: which evaluation metric to use for adjustments.
 
   Returns:
-    a function learning_rate(step): float -> float, the step-dependent lr.
+    a function learning_rate(step): float -> {"learning_rate": float}, the
+    step-dependent lr.
   """
   metrics = history.get(history_mode, metric)
   adjusted = constant
@@ -176,14 +178,20 @@ def PolicySchedule(
     policy_dir: directory with the policy checkpoint.
 
   Returns:
-    a function learning_rate(step): float -> float, the step-dependent lr.
+    a function learning_rate(step): float -> {"learning_rate": float}, the
+    step-dependent lr.
   """
 
   # Turn the history into observations for the policy. If we don't have any,
   # return the initial learning rate.
   start_time = time.time()
+  lr_config = ("learning_rate", start_lr, (1e-9, max_lr), False)
+  if include_lr_in_observation:
+    control_configs = (lr_config,)
+  else:
+    control_configs = None
   observations = online_tune.history_to_observations(
-      history, observation_metrics, observation_range, include_lr_in_observation
+      history, observation_metrics, observation_range, control_configs
   )
   logging.vlog(
       1, "Building observations took %0.2f sec.", time.time() - start_time)
@@ -225,6 +233,7 @@ def PolicySchedule(
   action = utils.gumbel_sample(log_probs[0, -1, :])
 
   # Get a new learning rate.
-  new_lr = online_tune.new_learning_rate(
-      action.item(), history, action_multipliers, max_lr)
-  return lambda _: new_lr
+  new_lr = online_tune.update_control(
+      lr_config, action.item(), history, action_multipliers
+  )
+  return lambda _: {"learning_rate": new_lr}
diff --git a/tensor2tensor/trax/learning_rate_test.py b/tensor2tensor/trax/learning_rate_test.py
index 73063aab4..fef08dd9d 100644
--- a/tensor2tensor/trax/learning_rate_test.py
+++ b/tensor2tensor/trax/learning_rate_test.py
@@ -74,14 +74,18 @@ def test_returns_start_lr_when_there_are_no_metrics(self):
   def test_changes_lr_when_there_are_some_metrics(self):
     history = trax_history.History()
     history.append("eval", "metrics/accuracy", step=0, value=0.8)
-    history.append(*online_tune.LEARNING_RATE_METRIC, step=0, value=1e-4)
+    history.append(
+        *online_tune.control_metric("learning_rate"), step=0, value=1e-4
+    )
     schedule = self._make_schedule(
         history,
         observation_metrics=(("eval", "metrics/accuracy"),),
         action_multipliers=(0.5, 2.0),
     )
+    new_lr = schedule(123)["learning_rate"]
     self.assertTrue(
-        onp.allclose(schedule(123), 5e-5) or onp.allclose(schedule(123), 2e-4))
+        onp.allclose(new_lr, 5e-5) or onp.allclose(new_lr, 2e-4)
+    )
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
index 3016320bd..7eb617244 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
@@ -40,7 +40,12 @@ WideResnet.n_output_classes = 10
 OnlineTuneEnv.inputs = @trax.inputs.inputs
 OnlineTuneEnv.model = @trax.models.WideResnet
 OnlineTuneEnv.optimizer = @trax.optimizers.Momentum
-OnlineTuneEnv.start_lr = 0.01
+OnlineTuneEnv.control_configs = (
+    ("learning_rate", 0.1, (1e-9, 10.0), False),
+    ("weight_decay_rate", 1e-5, (1e-9, 0.1), False),
+)
+OnlineTuneEnv.include_controls_in_observation = False
+OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
 OnlineTuneEnv.train_steps = 100
 OnlineTuneEnv.eval_steps = 10
 OnlineTuneEnv.env_steps = 100
diff --git a/tensor2tensor/trax/rl/configs/ppo_online_tune.gin b/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
index 23a656769..eade5f618 100644
--- a/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
+++ b/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
@@ -25,7 +25,7 @@ PPO.boundary = 128
 PPO.max_timestep = 128
 PPO.max_timestep_eval = 128
 PPO.random_seed = None
-PPO.gamma = 0.99
+PPO.gamma = 1.0
 PPO.lambda_ = 0.95
 PPO.c1 = 1.0
 PPO.c2 = 0.01
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index 3cf7203d4..1a25057e2 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -59,13 +59,17 @@ def __init__(self,
                    ("eval", "metrics/accuracy"),
                    ("eval", "metrics/loss"),
                ),
-               include_lr_in_observation=False,
+               include_controls_in_observation=False,
                reward_metric=("eval", "metrics/accuracy"),
                train_steps=100,
                eval_steps=10,
                env_steps=100,
-               start_lr=0.001,
-               max_lr=10.0,
+               # This is a tuple instead of a dict because the controls are
+               # ordered in the action space.
+               control_configs=(
+                   # (name, start, (low, high), flip)
+                   ("learning_rate", 1e-3, (1e-9, 10.0), False),
+               ),
                observation_range=(0.0, 5.0),
                # Don't save checkpoints by default, as they tend to use a lot of
                # space.
@@ -77,28 +81,31 @@ def __init__(self,
         model=model,
         loss_fn=loss_fn,
         optimizer=optimizer,
-        lr_schedule=(lambda history: lambda step: self._current_lr),
+        lr_schedule=(lambda history: lambda step: self._current_controls),
         inputs=inputs,
         should_save=should_save_checkpoints,
     )
     self._action_multipliers = action_multipliers
     self._observation_metrics = observation_metrics
-    self._include_lr_in_observation = include_lr_in_observation
+    self._include_controls_in_observation = include_controls_in_observation
     self._reward_metric = reward_metric
     self._train_steps = train_steps
     self._eval_steps = eval_steps
     self._env_steps = env_steps
-    self._start_lr = start_lr
-    self._max_lr = max_lr
+    self._control_configs = control_configs
 
     self._output_dir = output_dir
     gfile.makedirs(self._output_dir)
-    # Action is an index in self._action_multipliers.
-    self.action_space = gym.spaces.Discrete(len(self._action_multipliers))
+    # Actions are indices in self._action_multipliers.
+    self.action_space = gym.spaces.MultiDiscrete(
+        [len(self._action_multipliers)] * len(self._control_configs)
+    )
     # Observation is a vector with the values of the metrics specified in
-    # observation_metrics plus optionally the learning rate.
+    # observation_metrics plus optionally the current controls.
     observation_dim = (
-        len(self._observation_metrics) + int(self._include_lr_in_observation))
+        len(self._observation_metrics) +
+        int(self._include_controls_in_observation) * len(self._control_configs)
+    )
     self._observation_range = observation_range
     (low, high) = self._observation_range
     self.observation_space = gym.spaces.Box(
@@ -135,7 +142,6 @@ def _current_reward_metric(self):
     metric_values = online_tune.historical_metric_values(
         self._trainer.state.history,
         self._reward_metric,
-        self._observation_range,
     )
     assert metric_values.shape[0] > 0, (
         "No values in history for metric {}.".format(self._reward_metric))
@@ -147,7 +153,8 @@ def _current_observation(self):
         self._trainer.state.history,
         self._observation_metrics,
         self._observation_range,
-        self._include_lr_in_observation,
+        self._control_configs if self._include_controls_in_observation
+        else None,
     )
     assert observations.shape[0] > 0, "No values in history for any metric."
     return observations[-1, :]
@@ -159,7 +166,10 @@ def trainer(self):
     return self._trainer
 
   def reset(self):
-    self._current_lr = self._start_lr
+    self._current_controls = {
+        name: start_value
+        for (name, start_value, _, _) in self._control_configs
+    }
     self._step = 0
     self._trainer.reset(output_dir=self._next_trajectory_dir)
     self._trainer.evaluate(self._eval_steps)
@@ -179,12 +189,18 @@ def step(self, action):
         metric since the last step. done is set after reaching self.env_steps
         environment steps. info is an empty dict.
     """
-    self._current_lr = online_tune.new_learning_rate(
-        action,
-        self._trainer.state.history,
-        self._action_multipliers,
-        self._max_lr,
-    )
+    self._current_controls = {
+        # name: value
+        control_config[0]: online_tune.update_control(  # pylint: disable=g-complex-comprehension
+            control_config,
+            control_action,
+            self._trainer.state.history,
+            self._action_multipliers,
+        )
+        for (control_action, control_config) in zip(
+            action, self._control_configs
+        )
+    }
     last_reward_metric = self._current_reward_metric
     self._trainer.train_epoch(self._train_steps, self._eval_steps)
     self._step += 1
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env_test.py b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
index 64e981d82..ad6467676 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env_test.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
@@ -25,6 +25,7 @@
 
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import models
+from tensor2tensor.trax import optimizers
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.rl import online_tune
 from tensor2tensor.trax.rl.envs import online_tune_env
@@ -39,7 +40,7 @@ class MockTrainer(trax.Trainer):
 
   def __init__(self, metrics_to_report, *args, **kwargs):
     super(MockTrainer, self).__init__(*args, **kwargs)
-    self.learning_rates = []
+    self.controls = []
     self.init_metrics_to_report = metrics_to_report
     self.metrics_to_report = None
 
@@ -50,7 +51,7 @@ def reset(self, output_dir):
 
   def train_epoch(self, epoch_steps, eval_steps):
     del epoch_steps
-    self.learning_rates.append(self.learning_rate)
+    self.controls.append(self.optimizer_params)
     self.evaluate(eval_steps)
 
   def evaluate(self, eval_steps):
@@ -60,18 +61,21 @@ def evaluate(self, eval_steps):
         metric=METRIC,
         step=self.step,
         value=self.metrics_to_report.pop(0))
-    (lr_mode, lr_metric) = online_tune.LEARNING_RATE_METRIC
-    self.state.history.append(
-        mode=lr_mode,
-        metric=lr_metric,
-        step=self.step,
-        value=self.learning_rate)
+    for (name, value) in self.optimizer_params.items():
+      (mode, metric) = online_tune.control_metric(name)
+      self.state.history.append(
+          mode=mode,
+          metric=metric,
+          step=self.step,
+          value=value)
 
 
 class OnlineTuneTest(test.TestCase):
 
   @staticmethod
-  def _create_env(output_dir, metrics_to_report=(0.0,), action_multipliers=()):
+  def _create_env(
+      output_dir, metrics_to_report=(0.0,), action_multipliers=(1,)
+  ):
     return online_tune_env.OnlineTuneEnv(
         trainer_class=functools.partial(MockTrainer, metrics_to_report),
         model=functools.partial(
@@ -82,6 +86,13 @@ def _create_env(output_dir, metrics_to_report=(0.0,), action_multipliers=()):
             input_dtype=np.float32,
             output_shape=(1, 1),
             output_dtype=np.float32),
+        optimizer=optimizers.Momentum,
+        control_configs=(
+            ("learning_rate", 1e-3, (1e-9, 10.0), False),
+            ("weight_decay_rate", 1e-5, (1e-9, 0.1), False),
+        ),
+        observation_range=(-1, 1),
+        include_controls_in_observation=False,
         output_dir=output_dir,
         action_multipliers=action_multipliers,
         observation_metrics=[(HISTORY_MODE, METRIC)],
@@ -93,12 +104,16 @@ def _create_env(output_dir, metrics_to_report=(0.0,), action_multipliers=()):
   def test_communicates_with_trainer(self):
     action_multipliers = [0.8, 1.0, 1.25]
     metrics_to_report = [0.1, 0.5, 0.8, 0.9]
-    actions_to_take = [0, 1, 2]
+    actions_to_take = [[0, 1], [1, 2], [2, 0]]
     expected_observations = np.expand_dims(metrics_to_report, axis=1)
     # Metric difference in consecutive timesteps.
     expected_rewards = [0.4, 0.3, 0.1]
     expected_dones = [False, False, True]
-    expected_learning_rates = [0.0008, 0.0008, 0.001]
+    expected_controls = [
+        {"learning_rate": 0.0008, "weight_decay_rate": 1e-5},
+        {"learning_rate": 0.0008, "weight_decay_rate": 1.25e-5},
+        {"learning_rate": 0.001, "weight_decay_rate": 1e-5},
+    ]
 
     env = self._create_env(
         output_dir=self.get_temp_dir(),
@@ -116,8 +131,13 @@ def test_communicates_with_trainer(self):
     np.testing.assert_allclose(actual_observations, expected_observations)
     np.testing.assert_allclose(actual_rewards, expected_rewards)
     self.assertEqual(actual_dones, expected_dones)
-    np.testing.assert_allclose(env.trainer.learning_rates,
-                               expected_learning_rates)
+    def get_control(name, controls):
+      return [control[name] for control in controls]
+    for name in ("learning_rate", "weight_decay_rate"):
+      np.testing.assert_allclose(
+          get_control(name, env.trainer.controls),
+          get_control(name, expected_controls),
+      )
 
   def test_creates_new_trajectory_dirs(self):
     output_dir = self.get_temp_dir()
diff --git a/tensor2tensor/trax/rl/online_tune.py b/tensor2tensor/trax/rl/online_tune.py
index a4631edba..a470f71dd 100644
--- a/tensor2tensor/trax/rl/online_tune.py
+++ b/tensor2tensor/trax/rl/online_tune.py
@@ -22,36 +22,90 @@
 import numpy as np
 
 
-LEARNING_RATE_METRIC = ("train", "training/learning_rate")
-
-
-def historical_metric_values(
-    history, metric, observation_range=(-np.inf, np.inf)):
+def historical_metric_values(history, metric):
   """Converts a metric stream from a trax History object into a numpy array."""
   metric_sequence = history.get(*metric)
-  metric_values = np.array([
+  return np.array([
       metric_value for (_, metric_value) in metric_sequence
   ])
-  return np.clip(metric_values, *observation_range)
 
 
-def history_to_observations(history, metrics, observation_range, include_lr):
+def metric_to_observation(metric_values, metric_range):
+  """Clips and scales the metric to the [-1, 1] interval."""
+  (low, high) = metric_range
+  clipped_values = np.clip(metric_values, low, high)
+  return (clipped_values - low) / (high - low) * 2 - 1
+
+
+def control_to_observation(control_values, control_config):
+  """Flips, logarithms, clips and scales the control to the [-1, 1] interval."""
+  (_, _, (low, high), flip) = control_config
+  def transform(x):
+    return np.log(maybe_flip(x, flip))
+  (log_control_values, log_low, log_high) = map(
+      transform, (control_values, low, high)
+  )
+  if flip:
+    (log_low, log_high) = (log_high, log_low)
+  return metric_to_observation(log_control_values, (log_low, log_high))
+
+
+def control_metric(name):
+  """Returns the (mode, metric) pair in History for the given control."""
+  return ("train", "training/{}".format(name))
+
+
+def maybe_flip(value, flip):
+  """Flips a control (or not).
+
+  Meant to translate controls that naturally take values close to 1
+  (e.g. momentum) to a space where multiplication makes sense (i.e. close to 0).
+
+  Args:
+    value: float or numpy array, value of the control.
+    flip: bool, whether to flip or not.
+
+  Returns:
+    Either value or 1 - value based on flip.
+  """
+  if flip:
+    value = 1 - value
+  return value
+
+
+def history_to_observations(
+    history, metrics, observation_range, control_configs=None):
   """Converts a trax History object into a sequence of observations."""
   observation_dimensions = [
-      historical_metric_values(history, metric, observation_range)
+      metric_to_observation(  # pylint: disable=g-complex-comprehension
+          historical_metric_values(history, metric), observation_range
+      )
       for metric in metrics
   ]
-  if include_lr:
-    # Logartihm of the learning rate.
-    observation_dimensions.append(np.log(historical_metric_values(
-        history, LEARNING_RATE_METRIC, observation_range
-    )))
+  if control_configs is not None:
+    for control_config in control_configs:
+      (control_name, _, _, _) = control_config
+      observation_dimensions.append(control_to_observation(
+          historical_metric_values(history, control_metric(control_name)),
+          control_config,
+      ))
   return np.stack(observation_dimensions, axis=1)
 
 
-def new_learning_rate(action, history, action_multipliers, max_lr):
-  """Calculates a new learning rate based on an action."""
-  learning_rates = historical_metric_values(history, LEARNING_RATE_METRIC)
-  assert learning_rates.shape[0] > 0, "No last learning rate found in history."
-  current_lr = learning_rates[-1]
-  return min(current_lr * action_multipliers[action], max_lr)
+def update_control(control_config, action, history, action_multipliers):
+  """Calculates a new value of a control based on an action."""
+  (name, _, (low, high), flip) = control_config
+  metric = control_metric(name)
+  control_values = historical_metric_values(history, metric)
+  assert control_values.shape[0] > 0, (
+      "No last control {} found in history.".format(name))
+  current_control = control_values[-1]
+  (current_control, low, high) = maybe_flip(
+      np.array([current_control, low, high]), flip
+  )
+  if flip:
+    (low, high) = (high, low)
+  new_control = np.clip(
+      current_control * action_multipliers[action], low, high
+  )
+  return maybe_flip(new_control, flip)
diff --git a/tensor2tensor/trax/rl/online_tune_test.py b/tensor2tensor/trax/rl/online_tune_test.py
index 93486383c..4337b24e6 100644
--- a/tensor2tensor/trax/rl/online_tune_test.py
+++ b/tensor2tensor/trax/rl/online_tune_test.py
@@ -36,29 +36,63 @@ def test_retrieves_historical_metric_values(self):
     history = trax_history.History()
     self._append_metrics(history, ("train", "accuracy"), [0.1, 0.73])
     metric_values = online_tune.historical_metric_values(
-        history, metric=("train", "accuracy"), observation_range=(0, 5))
+        history, metric=("train", "accuracy")
+    )
     np.testing.assert_array_equal(metric_values, [0.1, 0.73])
 
-  def test_clips_historical_metric_values(self):
-    history = trax_history.History()
-    self._append_metrics(history, ("train", "loss"), [-10, 10])
-    metric_values = online_tune.historical_metric_values(
-        history, metric=("train", "loss"), observation_range=(-1, 1))
-    np.testing.assert_array_equal(metric_values, [-1, 1])
-
-  def test_converts_history_to_observations_without_learning_rate(self):
+  def test_metric_to_observation_rescales(self):
+    metric = np.random.uniform(low=-10, high=10, size=(100,))
+    observation = online_tune.metric_to_observation(metric, (-10, 10))
+    self.assertLess(-1, np.min(observation))
+    self.assertLess(np.min(observation), -0.9)
+    self.assertLess(0.9, np.max(observation))
+    self.assertLess(np.max(observation), 1.0)
+
+  def test_metric_to_observation_clips(self):
+    metric = np.random.uniform(low=-10, high=10, size=(100,))
+    observation = online_tune.metric_to_observation(metric, (-2, 2))
+    self.assertEqual(np.min(observation), -1)
+    self.assertEqual(np.max(observation), 1)
+
+  def test_converts_control_to_log_scale_without_flipping(self):
+    config = ("weight_decay", None, (1e-5, 0.1), False)
+    controls = np.array([0.01, 0.02, 0.04])
+    obs = online_tune.control_to_observation(controls, config)
+    np.testing.assert_almost_equal(obs[1] - obs[0], obs[2] - obs[1])
+
+  def test_converts_control_to_log_scale_with_flipping(self):
+    config = ("momentum", None, (0.5, 0.99), True)
+    controls = np.array([0.98, 0.96, 0.92])
+    obs = online_tune.control_to_observation(controls, config)
+    np.testing.assert_almost_equal(obs[1] - obs[0], obs[2] - obs[1])
+
+  def test_clips_control_without_flipping(self):
+    config = ("weight_decay", None, (1e-5, 0.1), False)
+    controls = np.array([0.0, 0.2])
+    obs = online_tune.control_to_observation(controls, config)
+    np.testing.assert_equal(obs, [-1, 1])
+
+  def test_clips_control_with_flipping(self):
+    config = ("momentum", None, (0.5, 0.99), True)
+    controls = np.array([0.4, 1.0])
+    obs = online_tune.control_to_observation(controls, config)
+    np.testing.assert_equal(obs, [1, -1])
+
+  def test_converts_history_to_observations_without_controls(self):
     history = trax_history.History()
-    self._append_metrics(history, ("train", "loss"), [3.0, 1.07])
+    self._append_metrics(history, ("train", "loss"), [1.0, 0.07])
     self._append_metrics(history, ("eval", "accuracy"), [0.12, 0.68])
     observations = online_tune.history_to_observations(
         history,
         metrics=(("eval", "accuracy"), ("train", "loss")),
-        observation_range=(0, 5),
-        include_lr=False,
+        observation_range=(-1, 1),
+        control_configs=None,
+    )
+    np.testing.assert_array_almost_equal(
+        observations, [[0.12, 1.0], [0.68, 0.07]]
     )
-    np.testing.assert_array_equal(observations, [[0.12, 3.0], [0.68, 1.07]])
 
-  def test_converts_history_to_observations_with_learning_rate(self):
+  def test_converts_history_to_observations_with_controls(self):
     history = trax_history.History()
     self._append_metrics(
         history, ("train", "training/learning_rate"), [1e-3, 1e-4])
@@ -66,7 +100,9 @@ def test_converts_history_to_observations_with_learning_rate(self):
         history,
         metrics=(),
         observation_range=(0, 5),
-        include_lr=True,
+        control_configs=(
+            ("learning_rate", None, (1e-9, 10.0), False),
+        ),
     )
     self.assertEqual(observations.shape, (2, 1))
     ((log_lr_1,), (log_lr_2,)) = observations
@@ -79,32 +115,61 @@ def test_clips_observations(self):
         history,
         metrics=(("eval", "loss"),),
         observation_range=(-2, 2),
-        include_lr=False,
+        control_configs=None,
+    )
+    np.testing.assert_array_equal(observations, [[-1], [1]])
+
+  def test_updates_control_without_flipping(self):
+    config = ("learning_rate", None, (1e-9, 10.0), False)
+    history = trax_history.History()
+    self._append_metrics(
+        history, online_tune.control_metric("learning_rate"), [1e-2, 1e-3])
+    new_control = online_tune.update_control(
+        control_config=config,
+        action=2,
+        history=history,
+        action_multipliers=(0.5, 1.0, 2.0),
     )
-    np.testing.assert_array_equal(observations, [[-2], [2]])
+    np.testing.assert_almost_equal(new_control, 2e-3)
 
-  def test_calculates_new_learning_rate(self):
+  def test_updates_control_with_flipping(self):
+    config = ("momentum", None, (0.5, 0.99), True)
     history = trax_history.History()
     self._append_metrics(
-        history, online_tune.LEARNING_RATE_METRIC, [1e-2, 1e-3])
-    new_lr = online_tune.new_learning_rate(
+        history, online_tune.control_metric("momentum"), [0.96, 0.98])
+    new_control = online_tune.update_control(
+        control_config=config,
+        action=0,
+        history=history,
+        action_multipliers=(0.5, 1.0, 2.0),
+    )
+    np.testing.assert_almost_equal(new_control, 0.99)
+
+  def test_clips_updated_control_without_flipping(self):
+    config = ("learning_rate", None, (1e-9, 10.0), False)
+    history = trax_history.History()
+    self._append_metrics(
+        history, online_tune.control_metric("learning_rate"), [7.0])
+    new_control = online_tune.update_control(
+        control_config=config,
         action=2,
         history=history,
         action_multipliers=(0.5, 1.0, 2.0),
-        max_lr=1.0,
     )
-    np.testing.assert_almost_equal(new_lr, 2e-3)
+    np.testing.assert_almost_equal(new_control, 10.0)
 
-  def test_clips_new_learning_rate(self):
+  def test_clips_updated_control_with_flipping(self):
+    config = ("momentum", None, (0.5, 0.99), True)
     history = trax_history.History()
-    self._append_metrics(history, online_tune.LEARNING_RATE_METRIC, [1e-3])
-    new_lr = online_tune.new_learning_rate(
+    self._append_metrics(
+        history, online_tune.control_metric("momentum"), [0.985])
+    new_control = online_tune.update_control(
+        control_config=config,
         action=0,
         history=history,
-        action_multipliers=(4.0, 1.0, 0.25),
-        max_lr=3e-3,
+        action_multipliers=(0.5, 1.0, 2.0),
     )
-    np.testing.assert_almost_equal(new_lr, 3e-3)
+    np.testing.assert_almost_equal(new_control, 0.99)
 
 
 if __name__ == "__main__":
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 51a8041ac..c533702f3 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -675,7 +675,7 @@ def reset(self, output_dir):
     if not state.opt_state:
       self._maybe_save_state(keep=False)
 
-    self.update_learning_rate()
+    self.update_optimizer_params()
 
   @property
   def step(self):
@@ -692,7 +692,7 @@ def state(self):
         model_state=self._model_state)
 
   @property
-  def learning_rate(self):
+  def optimizer_params(self):
     # TODO(lukaszkaiser): it makes no sense to use an accelerator (e.g. TPU)
     # in op-by-op mode just to compute the learning rate. However, there
     # should be a cleaner approach that forceably swapping out the backend.
@@ -720,9 +720,11 @@ def print_n_params(self):
   def _train_step(self, next_train_batch):
     """Run one training step and update self._opt_state."""
     # Calculate the current learning rate.
-    learning_rate = self._maybe_replicate(np.array(self.learning_rate))
+    opt_param_updates = layers.nested_map(
+        self.optimizer_params, lambda x: self._maybe_replicate(np.array(x))
+    )
     opt_state = self._opt_state
-    opt_state.opt_params["learning_rate"] = learning_rate
+    opt_state.opt_params.update(opt_param_updates)
 
     # Run the update.
     (params, slots), self._model_state, self._rngs = self._jit_update_fn(
@@ -749,9 +751,10 @@ def train_epoch(self, epoch_steps, eval_steps):
       if self._step in self._save_steps:
         self._maybe_save_state(keep=True)
 
-      # LR log
+      # Log optimizer params (learning rate etc.)
       if self._step == 1 or self._step % 10 == 0:
-        self._train_sw.scalar("training/learning_rate", self.learning_rate)
+        for (name, value) in self.optimizer_params.items():
+          self._train_sw.scalar("training/{}".format(name), value)
 
     # Timer
     epoch_time = time.time() - start_time
@@ -788,11 +791,12 @@ def evaluate(self, eval_steps):
         history=self._history,
         has_weights=self._has_weights)
 
-    # Save the learning rate in the history
-    self._history.append("train", "training/learning_rate", self._step,
-                         self.learning_rate)
+    # Save the optimizer params in the history
+    for (name, value) in self.optimizer_params.items():
+      self._history.append("train", "training/{}".format(name), self._step,
+                           value)
 
-  def update_learning_rate(self):
+  def update_optimizer_params(self):
     self._lr_fn = self._lr_schedule(self._history)
 
   def save_computation_graphs(self, save_backward_graph):
@@ -922,8 +926,8 @@ def train(output_dir,
   for epoch_steps in epochs(train_steps, trainer.step, epoch_steps):
     trainer.train_epoch(epoch_steps, eval_steps)
 
-    # Update learning rate with new history
-    trainer.update_learning_rate()
+    # Update optimizer parameters with new history
+    trainer.update_optimizer_params()
 
     # Bookkeeping we do at the first step
     if trainer.step == 1:

From 4804b8a5bd78f8eb86ed9cffb4009e57844e7be2 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 14 Sep 2019 09:49:30 -0700
Subject: [PATCH 2441/2720] [async-rl CL 8/n] : Wait to collect trajectories
 for rl_trainer.

This comes in handy when the rl_trainer itself restarts due to rescheduling
and its wait time is exhausted.

A `wait_forever` is appropriate here, since otherwise no progress can be made.

PiperOrigin-RevId: 269091806
---
 tensor2tensor/trax/rl/ppo_trainer.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index 820b81120..b3b09bbda 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -261,6 +261,7 @@ def collect_trajectories_async(self,
         trajectory_dir,
         epoch=epoch,
         temperature=temperature,
+        wait_forever=True,
         n_trajectories=n_trajectories)
 
     if bt is None:

From 4f537a85fb35233f7a41dd0ca281836a9ad52b9d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 14 Sep 2019 14:05:21 -0700
Subject: [PATCH 2442/2720] Actions can have more than 1 dimension now.

PiperOrigin-RevId: 269108474
---
 tensor2tensor/trax/rl/ppo.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index cd0fad361..ddb7ae417 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -294,11 +294,12 @@ def pad_trajectories(trajectories, boundary=20):
     padded_observations.append(padded_obs)
 
     # Now pad actions and rewards.
-    assert a.ndim == 1 and r.ndim == 1
-    padding_config = ((0, num_to_pad, 0),)
-
+    padding_config = tuple([(0, num_to_pad, 0)] + [(0, 0, 0)] * (a.ndim - 1))
     padded_action = lax.pad(a, action_padding_value, padding_config)
     padded_actions.append(padded_action)
+
+    assert r.ndim == 1
+    padding_config = ((0, num_to_pad, 0),)
     padded_reward = lax.pad(r, reward_padding_value, padding_config)
     padded_rewards.append(padded_reward)
 

From 38066343c9c76bad3d37549eaf2fef1f28619152 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Sat, 14 Sep 2019 22:40:35 -0700
Subject: [PATCH 2443/2720] Initialize Trainer in OnlineTuneEnv lazily to
 prevent long startup in the async setup.

PiperOrigin-RevId: 269142699
---
 tensor2tensor/trax/rl/envs/online_tune_env.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index 1a25057e2..bec69faea 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
 import os
 
 import gym
@@ -77,7 +78,11 @@ def __init__(self,
     if action_multipliers is None:
       action_multipliers = self.DEFAULT_ACTION_MULTIPLIERS
     self._model = model
-    self._trainer = trainer_class(
+    # Initialize Trainer in OnlineTuneEnv lazily to prevent long startup in the
+    # async setup, where we just use the environments as containers for
+    # trajectories.
+    self._trainer_fn = functools.partial(
+        trainer_class,
         model=model,
         loss_fn=loss_fn,
         optimizer=optimizer,
@@ -85,6 +90,7 @@ def __init__(self,
         inputs=inputs,
         should_save=should_save_checkpoints,
     )
+    self._trainer = None
     self._action_multipliers = action_multipliers
     self._observation_metrics = observation_metrics
     self._include_controls_in_observation = include_controls_in_observation
@@ -166,6 +172,8 @@ def trainer(self):
     return self._trainer
 
   def reset(self):
+    if self._trainer is None:
+      self._trainer = self._trainer_fn()
     self._current_controls = {
         name: start_value
         for (name, start_value, _, _) in self._control_configs

From 5a307b4f2ba93fa88e4e574dd065b893beb963f7 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Sat, 14 Sep 2019 23:43:28 -0700
Subject: [PATCH 2444/2720] Implement a serializer for MultiDiscrete spaces

PiperOrigin-RevId: 269146599
---
 tensor2tensor/trax/rl/space_serializer.py     | 32 +++++++++++++++++++
 .../trax/rl/space_serializer_test.py          | 32 +++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/tensor2tensor/trax/rl/space_serializer.py b/tensor2tensor/trax/rl/space_serializer.py
index ce35d5994..683aaf65d 100644
--- a/tensor2tensor/trax/rl/space_serializer.py
+++ b/tensor2tensor/trax/rl/space_serializer.py
@@ -91,6 +91,7 @@ def create(space, vocab_size):
   return {
       gym.spaces.Box: BoxSpaceSerializer,
       gym.spaces.Discrete: DiscreteSpaceSerializer,
+      gym.spaces.MultiDiscrete: MultiDiscreteSpaceSerializer,
   }[type(space)](space, vocab_size)
 
 
@@ -180,3 +181,34 @@ def deserialize(self, representation):
   @property
   def significance_map(self):
     return np.zeros(1, dtype=np.int32)
+
+
+class MultiDiscreteSpaceSerializer(SpaceSerializer):
+  """Serializer for gym.spaces.MultiDiscrete.
+
+  Assumes that the number of categories in each dimension fits in the number of
+  symbols.
+  """
+
+  space_type = gym.spaces.MultiDiscrete
+
+  def __init__(self, space, vocab_size):
+    super(MultiDiscreteSpaceSerializer, self).__init__(space, vocab_size)
+    assert np.max(space.nvec) <= vocab_size, (
+        "MultiDiscrete maximum number of categories should fit in the number "
+        "of symbols."
+    )
+
+  def serialize(self, data):
+    return data
+
+  def deserialize(self, representation):
+    return representation
+
+  @property
+  def representation_length(self):
+    return len(self._space.nvec)
+
+  @property
+  def significance_map(self):
+    return np.zeros(self.representation_length, dtype=np.int32)
diff --git a/tensor2tensor/trax/rl/space_serializer_test.py b/tensor2tensor/trax/rl/space_serializer_test.py
index 148bdfce7..eed62b3d4 100644
--- a/tensor2tensor/trax/rl/space_serializer_test.py
+++ b/tensor2tensor/trax/rl/space_serializer_test.py
@@ -112,5 +112,37 @@ def test_significance_map(self):
     np.testing.assert_array_equal(self._serializer.significance_map, [0])
 
 
+class MultiDiscreteSpaceSerializerTest(test.TestCase):
+
+  def setUp(self):
+    super(MultiDiscreteSpaceSerializerTest, self).setUp()
+    self._space = gym.spaces.MultiDiscrete(nvec=[2, 2])
+    self._serializer = space_serializer.create(self._space, vocab_size=2)
+
+  def _sample_batch(self):
+    return np.reshape(self._space.sample(), (1,) + self._space.shape)
+
+  def test_representation_length(self):
+    input_array = self._sample_batch()
+    representation = self._serializer.serialize(input_array)
+    self.assertEqual(
+        representation.shape, (1, self._serializer.representation_length))
+
+  def test_commutes(self):
+    input_array = self._sample_batch()
+    representation = self._serializer.serialize(input_array)
+    output_array = self._serializer.deserialize(representation)
+    np.testing.assert_array_almost_equal(input_array, output_array)
+
+  def test_representation_changes(self):
+    array1 = self._sample_batch()
+    array2 = 1 - array1
+    (repr1, repr2) = tuple(map(self._serializer.serialize, (array1, array2)))
+    self.assertFalse(np.array_equal(repr1, repr2))
+
+  def test_significance_map(self):
+    np.testing.assert_array_equal(self._serializer.significance_map, [0, 0])
+
+
 if __name__ == "__main__":
   test.main()

From 53294b0d55393cd61064b45fe21c4c6ce90351b2 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Sun, 15 Sep 2019 19:36:28 -0700
Subject: [PATCH 2445/2720] Make SimPLe robust to corrupted trajectory shards.

PiperOrigin-RevId: 269234105
---
 tensor2tensor/trax/rl/simple.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/rl/simple.py b/tensor2tensor/trax/rl/simple.py
index 57feaa2b9..35b6d501f 100644
--- a/tensor2tensor/trax/rl/simple.py
+++ b/tensor2tensor/trax/rl/simple.py
@@ -40,11 +40,17 @@ def load_trajectories(trajectory_dir, eval_frac):
   for (subdir, _, filenames) in gfile.walk(trajectory_dir):
     for filename in filenames:
       shard_path = os.path.join(subdir, filename)
-      with gfile.GFile(shard_path, "rb") as f:
-        trajectories = pkl_module.load(f)
+      try:
+        with gfile.GFile(shard_path, "rb") as f:
+          trajectories = pkl_module.load(f)
         pivot = int(len(trajectories) * (1 - eval_frac))
         train_trajectories.extend(trajectories[:pivot])
         eval_trajectories.extend(trajectories[pivot:])
+      except EOFError:
+        logging.warning(
+            "Could not load trajectories from a corrupted shard %s.",
+            shard_path,
+        )
   assert train_trajectories, "Haven't found any training data."
   assert eval_trajectories, "Haven't found any evaluation data."
   return (train_trajectories, eval_trajectories)

From c999853dc43975d25d13b33fcd4a7365c775cb3d Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Sun, 15 Sep 2019 19:50:10 -0700
Subject: [PATCH 2446/2720] Log timing information in SimPLe.

PiperOrigin-RevId: 269235280
---
 tensor2tensor/trax/rl/simple_trainer.py | 27 +++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index edd371b26..394eac484 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -23,6 +23,7 @@
 import itertools
 import os
 import random
+import time
 
 from absl import logging
 from matplotlib import pyplot as plt
@@ -123,6 +124,7 @@ def flush_summaries(self):
 
   def collect_trajectories(self, evaluate):
     logging.info("SimPLe epoch [% 6d]: collecting data.", self._simple_epoch)
+    start_time = time.time()
 
     self._policy_trainer.train_env = self.train_env
     self._policy_trainer.trajectory_dump_dir = os.path.join(
@@ -130,8 +132,12 @@ def collect_trajectories(self, evaluate):
     self._policy_epoch += self._n_real_epochs
     self._policy_trainer.training_loop(self._policy_epoch, evaluate=evaluate)
 
+    logging.vlog(
+        1, "Collecting trajectories took %0.2f sec.", time.time() - start_time)
+
   def train_model(self):
     logging.info("SimPLe epoch [% 6d]: training model.", self._simple_epoch)
+    start_time = time.time()
 
     (train_stream, eval_stream) = self._make_input_streams()
     # Ignore n_devices for now.
@@ -154,8 +160,12 @@ def train_model(self):
         has_weights=True,
     )
 
+    logging.vlog(
+        1, "Training model took %0.2f sec.", time.time() - start_time)
+
   def train_policy(self):
     logging.info("SimPLe epoch [% 6d]: training policy.", self._simple_epoch)
+    start_time = time.time()
 
     self._sim_env.initialize(
         batch_size=self._simulated_batch_size,
@@ -167,6 +177,9 @@ def train_policy(self):
     self._policy_epoch += self._n_simulated_epochs
     self._policy_trainer.training_loop(self._policy_epoch, evaluate=False)
 
+    logging.vlog(
+        1, "Training policy took %0.2f sec.", time.time() - start_time)
+
   @property
   def _has_own_data(self):
     return self._simple_epoch > 0 or self._initial_trajectory_dir is None
@@ -193,17 +206,27 @@ def make_example_streams(trajectory_dir):
     mix_prob = self._initial_trajectory_mix_prob
 
     if self._has_initial_data:
+      start_time = time.time()
       # Load the initial, precollected data.
       (init_train_stream,
        init_eval_stream) = make_example_streams(self._initial_trajectory_dir)
+      logging.vlog(
+          1, "Loading initial trajectories took %0.2f sec.",
+          time.time() - start_time
+      )
     else:
       (init_train_stream, init_eval_stream) = (None, None)
       mix_prob = 0.0  # Take just our own collected data.
 
     if self._has_own_data:
+      start_time = time.time()
       # Load trajectories collected in all epochs so far.
       (own_train_stream,
        own_eval_stream) = make_example_streams(self._trajectory_dump_root_dir)
+      logging.vlog(
+          1, "Loading own trajectories took %0.2f sec.",
+          time.time() - start_time
+      )
     else:
       # We start the loop with training the model, so we don't have our own
       # collected data yet.
@@ -223,6 +246,7 @@ def mix_and_batch(streams):
 
   def evaluate_model(self):
     logging.info("SimPLe epoch [% 6d]: evaluating model.", self._simple_epoch)
+    start_time = time.time()
 
     self._sim_env.initialize(
         batch_size=self._simulated_batch_size,
@@ -248,3 +272,6 @@ def evaluate_model(self):
       self._summary_writer.plot(
           "simple/model_eval_plot", plt, step=self._simple_epoch)
       self.flush_summaries()
+
+    logging.vlog(
+        1, "Evaluating model took %0.2f sec.", time.time() - start_time)

From e68147c47fb5ebea5bc05113669481201b18d3a9 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Sun, 15 Sep 2019 20:11:07 -0700
Subject: [PATCH 2447/2720] Train longer on the initial set of trajectories,
 allow starting from an initial model checkpoint and update the gin config for
 SimPLe.

PiperOrigin-RevId: 269237329
---
 .../trax/rl/configs/simple_online_tune.gin    | 43 +++++++++++++------
 tensor2tensor/trax/rl/simple_trainer.py       | 39 ++++++++++++++---
 tensor2tensor/trax/rl/simple_trainer_test.py  |  3 +-
 3 files changed, 64 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/trax/rl/configs/simple_online_tune.gin b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
index 70d2e63e2..4df69ceb5 100644
--- a/tensor2tensor/trax/rl/configs/simple_online_tune.gin
+++ b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
@@ -1,4 +1,6 @@
 import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
 import tensor2tensor.trax.rl
 import tensor2tensor.trax.rl.space_serializer
 import tensor2tensor.trax.rl.trainers
@@ -9,9 +11,16 @@ BoxSpaceSerializer.precision = 2
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-MultifactorSchedule.constant = 0.1
+MultifactorSchedule.constant = 3.0
 MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 1000
+MultifactorSchedule.warmup_steps = 10000
+
+# Parameters for Adam:
+# ==============================================================================
+Adam.learning_rate = 1e-3
+Adam.b1 = 0.9
+Adam.b2 = 0.999
+Adam.weight_decay_rate = 0.0
 
 # Parameters for TransformerDecoder:
 # ==============================================================================
@@ -29,7 +38,7 @@ PPO.boundary = 128
 PPO.max_timestep = 128
 PPO.max_timestep_eval = 128
 PPO.random_seed = None
-PPO.gamma = 0.99
+PPO.gamma = 1.0
 PPO.lambda_ = 0.95
 PPO.c1 = 1.0
 PPO.c2 = 0.01
@@ -38,6 +47,8 @@ PPO.len_history_for_policy = None
 PPO.separate_eval = False
 PPO.save_every_n = 1
 PPO.policy_and_value_model = @trax.models.TransformerDecoder
+PPO.policy_and_value_optimizer = @trax.optimizers.Adam
+PPO.trajectory_dump_min_count_per_shard = 8
 
 # Parameters for SerializedSequenceSimulatedEnvProblem:
 # ==============================================================================
@@ -46,28 +57,34 @@ SerializedSequenceSimulatedEnvProblem.reward_fn = @trax.rl.onlinetune_reward_fn
 SerializedSequenceSimulatedEnvProblem.done_fn = @trax.rl.onlinetune_done_fn
 SerializedSequenceSimulatedEnvProblem.vocab_size = 128
 SerializedSequenceSimulatedEnvProblem.max_trajectory_length = 101
-SerializedSequenceSimulatedEnvProblem.significance_decay = 1.0
+SerializedSequenceSimulatedEnvProblem.significance_decay = 0.8
 
 # Parameters for SimPLe:
 # ==============================================================================
 SimPLe.policy_trainer_class = @trax.rl.trainers.PPO
-SimPLe.n_real_epochs = 2
-SimPLe.n_model_train_steps = 100000
-SimPLe.model_train_batch_size = 128
+SimPLe.n_real_epochs = 1
+SimPLe.n_model_initial_train_steps = 50000
+SimPLe.n_model_train_steps_per_epoch = 20000
+SimPLe.model_train_batch_size = 64
 SimPLe.simulated_env_problem_class = @trax.rl.SerializedSequenceSimulatedEnvProblem
 SimPLe.simulated_batch_size = 128
-SimPLe.n_simulated_epochs = 50
+SimPLe.n_simulated_epochs = 30
 SimPLe.initial_trajectory_mix_prob = 0.5
 
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.d_model = 256
-TransformerLM.d_ff = 512
-TransformerLM.n_layers = 2
-TransformerLM.n_heads = 2
-TransformerLM.dropout = 0.1
+TransformerLM.d_model = 512
+TransformerLM.d_ff = 1024
+TransformerLM.n_layers = 3
+TransformerLM.n_heads = 4
+TransformerLM.dropout = 0.3
 TransformerLM.max_len = 1024
 
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 1000
+train.optimizer = @trax.optimizers.Adafactor
+
 # Parameters for train_rl:
 # ==============================================================================
 train_rl.env_name = "ClientEnv-v0"
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index 394eac484..a68f32736 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -33,6 +33,7 @@
 from tensor2tensor.trax.rl import base_trainer
 from tensor2tensor.trax.rl import simple
 from tensor2tensor.trax.rl import simulated_env_problem
+from tensorflow.io import gfile
 
 
 class SimPLe(base_trainer.BaseTrainer):
@@ -46,7 +47,8 @@ def __init__(self,
                n_real_epochs=10,
                data_eval_frac=0.125,
                model_train_batch_size=64,
-               n_model_train_steps=1000,
+               n_model_initial_train_steps=1000,
+               n_model_train_steps_per_epoch=1000,
                simulated_env_problem_class=(
                    simulated_env_problem.SerializedSequenceSimulatedEnvProblem),
                simulated_batch_size=16,
@@ -54,6 +56,7 @@ def __init__(self,
                trajectory_dump_dir=None,
                initial_trajectory_dir=None,
                initial_trajectory_mix_prob=0.5,
+               initial_model=None,
                **kwargs):
     super(SimPLe, self).__init__(train_env, eval_env, output_dir, **kwargs)
     self._policy_dir = os.path.join(output_dir, "policy")
@@ -65,9 +68,20 @@ def __init__(self,
     )
     self._n_real_epochs = n_real_epochs
     self._model_train_batch_size = model_train_batch_size
-    self._n_model_train_steps = n_model_train_steps
+    self._n_model_initial_train_steps = n_model_initial_train_steps
+    self._n_model_train_steps_per_epoch = n_model_train_steps_per_epoch
     self._data_eval_frac = data_eval_frac
     self._model_dir = os.path.join(output_dir, "model")
+
+    gfile.makedirs(self._model_dir)
+    if initial_model is not None:
+      gfile.copy(
+          initial_model,
+          os.path.join(self._model_dir, "model.pkl"),
+          overwrite=True,
+      )
+    self._initial_model = initial_model
+
     self._sim_env = simulated_env_problem_class(
         batch_size=None,
         observation_space=train_env.observation_space,
@@ -104,9 +118,10 @@ def train_epoch(self, evaluate=True):
     if self._simple_epoch > 0 or not self._has_initial_data:
       self.collect_trajectories(evaluate=evaluate)
     # Train the model of the environment on the collected trajectories.
-    self.train_model()
-    if evaluate:
-      self.evaluate_model()
+    if self._simple_epoch > 0 or not self._initial_model:
+      skipped = self.train_model()
+      if evaluate and not skipped:
+        self.evaluate_model()
     # Train the policy inside the simulated environment generated by the model.
     self.train_policy()
 
@@ -136,6 +151,11 @@ def collect_trajectories(self, evaluate):
         1, "Collecting trajectories took %0.2f sec.", time.time() - start_time)
 
   def train_model(self):
+    """Train the model.
+
+    Returns:
+      whether the training was skipped due to a restart.
+    """
     logging.info("SimPLe epoch [% 6d]: training model.", self._simple_epoch)
     start_time = time.time()
 
@@ -151,8 +171,12 @@ def train_model(self):
         target_shape=self._sim_env.model_input_shape,
         target_dtype=self._sim_env.model_input_dtype)
 
-    self._model_train_step += self._n_model_train_steps
-    trax.train(
+    if self._simple_epoch == 0:
+      train_steps = self._n_model_initial_train_steps
+    else:
+      train_steps = self._n_model_train_steps_per_epoch
+    self._model_train_step += train_steps
+    state = trax.train(
         model=self._sim_env.model,
         inputs=inputs,
         train_steps=self._model_train_step,
@@ -162,6 +186,7 @@ def train_model(self):
 
     logging.vlog(
         1, "Training model took %0.2f sec.", time.time() - start_time)
+    return state.step > self._model_train_step
 
   def train_policy(self):
     logging.info("SimPLe epoch [% 6d]: training policy.", self._simple_epoch)
diff --git a/tensor2tensor/trax/rl/simple_trainer_test.py b/tensor2tensor/trax/rl/simple_trainer_test.py
index 6a9878e90..317db27b8 100644
--- a/tensor2tensor/trax/rl/simple_trainer_test.py
+++ b/tensor2tensor/trax/rl/simple_trainer_test.py
@@ -71,7 +71,8 @@ def test_training_loop_acrobot(self):
         n_real_epochs=1,
         data_eval_frac=0.5,
         model_train_batch_size=2,
-        n_model_train_steps=1,
+        n_model_initial_train_steps=1,
+        n_model_train_steps_per_epoch=1,
         simulated_env_problem_class=functools.partial(
             simulated_env_problem.SerializedSequenceSimulatedEnvProblem,
             model=functools.partial(

From 7ff1a59c4203ffbffc7b0c27e3f8393d09ec8c16 Mon Sep 17 00:00:00 2001
From: Zi Yang <ziy@google.com>
Date: Mon, 16 Sep 2019 01:05:04 -0700
Subject: [PATCH 2448/2720] Reenabled the workaround for "ValueError:
 prediction values with "inputs: DatasetToSingleElement:0" must be from the
 default graph" with use_tpu check.

PiperOrigin-RevId: 269274531
---
 tensor2tensor/utils/t2t_model.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 224a1f236..7acba41e3 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1701,6 +1701,13 @@ def estimator_spec_predict(self, features, use_tpu=False):
       outputs = infer_out
       scores = None
 
+    # Workaround for "ValueError: prediction values must be from the default
+    # graph" during TPU model exporting.
+    # TODO(b/130501786): remove tf.identity once default graph mismatch is fixed
+    if use_tpu:
+      for name, feature in features.items():
+        features[name] = tf.identity(feature)
+
     inputs = features.get("inputs")
     if inputs is None:
       inputs = features["targets"]

From ce56e181d9c7d992116bdc174891939f17f21fda Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 16 Sep 2019 14:49:23 -0700
Subject: [PATCH 2449/2720] More options to experiment with multi-hash
 attention and inputs code for synthetic problems.

PiperOrigin-RevId: 269426144
---
 .../trax/configs/transformer_copy.gin         | 70 +++++++++++++++++++
 .../trax/configs/transformer_imagenet64.gin   |  9 ++-
 tensor2tensor/trax/inputs.py                  | 51 ++++++++++++++
 tensor2tensor/trax/layers/attention.py        | 37 ++++++++--
 4 files changed, 157 insertions(+), 10 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/transformer_copy.gin

diff --git a/tensor2tensor/trax/configs/transformer_copy.gin b/tensor2tensor/trax/configs/transformer_copy.gin
new file mode 100644
index 000000000..e1f58afa4
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_copy.gin
@@ -0,0 +1,70 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 32
+batch_fun.eval_batch_size = 32
+batch_fun.max_eval_length = 1024
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 'sequence_copy'
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 1.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for sequence_copy_inputs:
+# ==============================================================================
+sequence_copy_inputs.vocab_size = 128
+sequence_copy_inputs.batch_size = 32
+sequence_copy_inputs.train_lengths = [1024]
+sequence_copy_inputs.eval_lengths = [1024]
+sequence_copy_inputs.reverse = False
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 500
+train.eval_steps = 64
+train.inputs = @trax.inputs.sequence_copy_inputs
+train.model = @trax.models.TransformerLM
+train.optimizer = @trax.optimizers.Adafactor
+train.train_steps = 50000
+train.has_weights = True
+
+# Parameters for MemoryEfficientCausalAttention:
+# ==============================================================================
+MemoryEfficientCausalAttention.dropout = 0.0
+MemoryEfficientCausalAttention.loop_stride = 512
+
+# Parameters for MergedMultiHashedCausalAttention:
+# ==============================================================================
+MergedMultiHashedCausalAttention.dropout = 0.0
+MergedMultiHashedCausalAttention.n_bins = 32
+MergedMultiHashedCausalAttention.n_hashes = 2
+MergedMultiHashedCausalAttention.n_buckets_per_bin = 1
+MergedMultiHashedCausalAttention.bin_by_time = False
+MergedMultiHashedCausalAttention.one_rng = False
+MergedMultiHashedCausalAttention.drop_for_hash_rate = 0.0
+MergedMultiHashedCausalAttention.hard_k = 16
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.attention_type = @trax.layers.MergedMultiHashedCausalAttention
+TransformerLM.d_attention_key = 64
+TransformerLM.d_attention_value = 64
+TransformerLM.d_model = 256
+TransformerLM.d_ff = 256
+TransformerLM.dropout = 0.0
+TransformerLM.max_len = 1024
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 4
+TransformerLM.n_layers = 1
+TransformerLM.share_kv = True
+TransformerLM.vocab_size = 128
diff --git a/tensor2tensor/trax/configs/transformer_imagenet64.gin b/tensor2tensor/trax/configs/transformer_imagenet64.gin
index ca64be8fb..bfd8246ee 100644
--- a/tensor2tensor/trax/configs/transformer_imagenet64.gin
+++ b/tensor2tensor/trax/configs/transformer_imagenet64.gin
@@ -5,8 +5,8 @@ import tensor2tensor.trax.trax
 
 # Parameters for batch_fun:
 # ==============================================================================
-batch_fun.batch_size_per_device = 2
-batch_fun.eval_batch_size = 16
+batch_fun.batch_size_per_device = 1
+batch_fun.eval_batch_size = 8
 batch_fun.max_eval_length = 12288  # 64 * 64 * 3
 
 # Parameters for inputs:
@@ -47,9 +47,12 @@ MergedMultiHashedCausalAttention.one_rng = False
 # ==============================================================================
 MergedMultiHashedCausalAttention.dropout = 0.0
 MergedMultiHashedCausalAttention.n_bins = 64
-MergedMultiHashedCausalAttention.n_hashes = 2
+MergedMultiHashedCausalAttention.n_hashes = 4
+MergedMultiHashedCausalAttention.n_buckets_per_bin = 2
 MergedMultiHashedCausalAttention.bin_by_time = False
 MergedMultiHashedCausalAttention.one_rng = False
+MergedMultiHashedCausalAttention.drop_for_hash_rate = 0.1
+MergedMultiHashedCausalAttention.hard_k = 32
 
 # Parameters for TransformerLM:
 # ==============================================================================
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index c18150001..e962c965e 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -199,6 +199,57 @@ def random_minibatches():
                 target_dtype=output_dtype)
 
 
+@gin.configurable(blacklist=['n_devices'])
+def sequence_copy_inputs(
+    n_devices, vocab_size=gin.REQUIRED, batch_size=gin.REQUIRED,
+    train_lengths=gin.REQUIRED, eval_lengths=gin.REQUIRED, reverse=False):
+  """Inputs for the sequence copy problem: 0w0w for w in [1..vocab_size-1]*.
+
+  Args:
+    n_devices: how many devices to build the inputs for.
+    vocab_size: how many symbols to use.
+    batch_size: how large are the batches.
+    train_lengths: lengths of w for training.
+    eval_lengths: lengths of w for eval.
+    reverse: bool (optional, false by default): reverse the second sequence.
+
+  Returns:
+    trax.inputs.Inputs
+  """
+  assert batch_size % n_devices == 0
+  def random_minibatches(length_list):
+    """Generate a stream of random mini-batches."""
+    while True:
+      length = random.choice(length_list)
+      assert length % 2 == 0
+      w_length = (length // 2) - 1
+      w = onp.random.randint(low=1, high=vocab_size-1,
+                             size=(batch_size, w_length))
+      zero = onp.zeros([batch_size, 1], onp.int32)
+      loss_weights = onp.concatenate([onp.zeros((batch_size, w_length+2)),
+                                      onp.ones((batch_size, w_length))], axis=1)
+      if reverse:
+        x = onp.concatenate([zero, w, zero, np.flip(w, axis=1)], axis=1)
+      else:
+        x = onp.concatenate([zero, w, zero, w], axis=1)
+      yield (x, x, loss_weights)  # Here inputs and targets are the same.
+
+  # If there's only one length, make the shape known.
+  example_length = None
+  if (len(train_lengths) == 1 and len(eval_lengths) == 1 and
+      train_lengths[0] == eval_lengths[0]):
+    example_length = train_lengths[0]
+
+  return Inputs(
+      train_stream=lambda: random_minibatches(train_lengths),
+      train_eval_stream=lambda: random_minibatches(train_lengths),
+      eval_stream=lambda: random_minibatches(eval_lengths),
+      input_shape=(example_length,),
+      input_dtype=onp.int32,
+      target_shape=(example_length,),
+      target_dtype=onp.int32)
+
+
 def dataset_to_stream(dataset, input_name, n_chunks=0):
   """Takes a tf.Dataset and creates a numpy stream of ready batches."""
   for example in backend.dataset_as_numpy(dataset):
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 104edba32..f237d7a2b 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -747,14 +747,19 @@ class MergedMultiHashedCausalAttention(BaseCausalAttention):
   """Hash-based causal attention, with multiple hashes."""
 
   def __init__(self, dropout, mode, n_bins=64, n_hashes=1,
-               bin_by_time=False, one_rng=False):
-    del dropout, mode
+               n_buckets_per_bin=1, bin_by_time=False, one_rng=False,
+               drop_for_hash_rate=0.0, hard_k=0):
+    del dropout
+    self._mode = mode
     super(MergedMultiHashedCausalAttention, self).__init__()
     self.n_bins = n_bins
     self.n_hashes = n_hashes
+    self.n_buckets_per_bin = n_buckets_per_bin
     self.bin_by_time = bin_by_time
     seed = random.randint(0, 2**31 - 1)
     self._one_rng = one_rng
+    self._drop_for_hash_rate = drop_for_hash_rate
+    self._hard_k = hard_k
     self._prng = None
     if one_rng:
       self._prng = backend.random.get_prng(seed)
@@ -775,6 +780,13 @@ def make_unit_length(self, x, epsilon=1e-6):
     norm_inputs = x / np.sqrt(variance + epsilon)
     return norm_inputs
 
+  def drop_for_hash(self, x, rng):
+    rate = self._drop_for_hash_rate
+    if self._mode == 'train' and rate > 0.0:
+      keep = backend.random.bernoulli(rng, 1.0 - rate, x.shape)
+      return np.where(keep, x / (1.0 - rate), np.zeros_like(x))
+    return x
+
   def hash_vectors(self, vecs, rng):
     if self.bin_by_time:
       # Instead of hashing, put chunks of consecutive items in the same bin.
@@ -786,18 +798,20 @@ def hash_vectors(self, vecs, rng):
     # (crucially) each round of hashing. All of these are part of dimension 0
     # of vecs. Applying multiple hashes to the same input is important because
     # it increases the probability of being in the same bin as relevant items.
-    assert self.n_bins % 2 == 0
+    n_buckets = self.n_buckets_per_bin * self.n_bins
+    assert n_buckets % 2 == 0
     rot_rng = rng
     if self._one_rng:
       rot_rng = jax.lax.tie_in(vecs, self._prng)
     random_rotation = jax.random.normal(
         rot_rng,
-        (vecs.shape[0], vecs.shape[-1], self.n_bins//2)).astype('float32')
+        (vecs.shape[0], vecs.shape[-1], n_buckets // 2)).astype('float32')
 
     # TODO(kitaev): making the vectors unit-length here is probably redundant.
-    vecs = self.make_unit_length(vecs)
+    # vecs = self.make_unit_length(vecs)
+    rng, subrng = backend.random.split(rng)
+    vecs = self.drop_for_hash(vecs, subrng)
     rotated_vecs = np.matmul(vecs, random_rotation)
-    rotated_vecs = self.make_unit_length(rotated_vecs)
     rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)
     bins = np.argmax(rotated_vecs, axis=-1)
     return bins
@@ -824,7 +838,7 @@ def call(self, inputs, params=(), state=(), rng=None, **kwargs):
     joint_t = np.reshape(joint_t, (1, seqlen))
     joint_t = np.broadcast_to(joint_t, qk.shape[:-1])
 
-    assert int((self.n_bins + 1) * seqlen) < 2 ** 31, (
+    assert int((self.n_buckets_per_bin * self.n_bins + 1) * seqlen) < 2 ** 31, (
         'Potential 32-bit integer overflow; please double-check the code.')
     joint_bins_and_t = seqlen * bins + joint_t
 
@@ -928,6 +942,15 @@ def vjpfun(grad):
     dots_logsumexp = backend.logsumexp(dots, axis=-1, keepdims=True)
     dots = np.exp(dots - dots_logsumexp)
 
+    if self._hard_k > 0:
+      top_k = np.sort(dots)[..., -self._hard_k]  # Get the top-kth weight.
+      top_k = jax.lax.stop_gradient(top_k)
+      dots -= top_k[..., np.newaxis]  # Subtract (be 0 for lower ones).
+      dots = np.maximum(dots, 0)
+      dots_sum = np.sum(dots, axis=-1, keepdims=True)  # Sum to re-normalize.
+      dots_logsumexp += np.log(dots_sum)  # Add it to the weight.
+      dots /= dots_sum  # Re-normalize.
+
     bo = np.matmul(dots, bv)
     so = unchunk_vectors(bo)
     slogits = unchunk_vectors(dots_logsumexp)

From 043fdf585109e8fdfe8e1e78e2feaaa5ec166ee1 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 16 Sep 2019 22:22:22 -0700
Subject: [PATCH 2450/2720] Mask out duplicate attention between hash rounds

PiperOrigin-RevId: 269493726
---
 tensor2tensor/trax/layers/attention.py | 49 ++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index f237d7a2b..f57a2b94b 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -748,6 +748,7 @@ class MergedMultiHashedCausalAttention(BaseCausalAttention):
 
   def __init__(self, dropout, mode, n_bins=64, n_hashes=1,
                n_buckets_per_bin=1, bin_by_time=False, one_rng=False,
+               allow_duplicate_attention=False,
                drop_for_hash_rate=0.0, hard_k=0):
     del dropout
     self._mode = mode
@@ -763,6 +764,7 @@ def __init__(self, dropout, mode, n_bins=64, n_hashes=1,
     self._prng = None
     if one_rng:
       self._prng = backend.random.get_prng(seed)
+    self._allow_duplicate_attention = allow_duplicate_attention
 
   def bin_vectors_by_time(self, vecs):
     seqlen = vecs.shape[-2]
@@ -938,6 +940,53 @@ def vjpfun(grad):
     self_mask = jax.lax.tie_in(dots, self_mask)
     dots = dots - 32 * self_mask
 
+    # Mask out later rounds attending to the same items as previous rounds
+    if self.n_hashes > 1 and not self._allow_duplicate_attention:
+      chunks = undo_sort // bq_t.shape[2]  # n_hashes*n_batch*n_heads, seqlen
+      chunks = np.reshape(chunks, (self.n_hashes, -1, seqlen))
+      chunks = np.moveaxis(chunks, 0, -1)
+      chunks = np.tile(chunks, (self.n_hashes, 1, 1))
+      # chunks is now n_hashes*n_batch*n_heads, seqlen, n_hashes
+      schunks = np.take_along_axis(chunks, sjoint_t[:, :, None], axis=-2)
+      bchunks = chunk_vectors(schunks)
+
+      # Queries/keys have shape (n_hashes*n_batch*n_heads, n_bins, binlen). For
+      # each query/key vector, the chunks numbers it's mapped to across each of
+      # the rounds of hashing are stored in bchunks, which has shape
+      # (n_hashes*n_batch*n_heads, n_bins, binlen, n_hashes). Query-key pairs
+      # that fall in the same or consecutive chunks in one hashing round will be
+      # masked out for subsequent rounds.
+      round_counter = jax.lax.tie_in(bchunks, np.arange(self.n_hashes))
+      cur_round = np.tile(
+          np.reshape(round_counter, (-1, 1)),
+          (1, bchunks.shape[0] // self.n_hashes))
+      # cur_round (shape n_hashes*n_batch*n_heads, 1, 1, 1) specifies which
+      # round of hashing a query-key pair belongs to. This shape broadcasts with
+      # (shape n_hashes*n_batch*n_heads, n_bins, binlen, n_hashes). The first
+      # n_batch*n_heads elements along dimension 0 are hash round 0, the next
+      # are round 1, etc.
+      cur_round = np.reshape(cur_round, (-1, 1, 1, 1))
+      # past_round (shape 1, 1, 1, n_hashes) contains round numbers, where the
+      # last dimension has one entry per hashing round.
+      past_round = np.reshape(round_counter, (1, 1, 1, -1))
+      # Set query chunk numbers for future rounds to an out-of-bounds (negative)
+      # value, so they don't match with any keys.
+      bq_chunks = np.where(past_round < cur_round, bchunks, -bchunks)
+
+      bkv_chunks_extra = np.concatenate(
+          [bchunks[:, -1:, :, :], bchunks[:, :-1, :, :]], axis=1)
+      bkv_chunks = np.concatenate([bchunks, bkv_chunks_extra], axis=2)
+
+      dup_mask = np.any(
+          jax.lax.eq(bq_chunks[:, :, :, None, :], bkv_chunks[:, :, None, :, :]),
+          axis=-1)
+      dup_mask = dup_mask | np.any(
+          jax.lax.eq(
+              bq_chunks[:, :, :, None, :], bkv_chunks[:, :, None, :, :] + 1),
+          axis=-1)
+      dup_mask = jax.lax.convert_element_type(dup_mask, np.float32)
+      dots = dots - 30 * dup_mask
+
     # Softmax.
     dots_logsumexp = backend.logsumexp(dots, axis=-1, keepdims=True)
     dots = np.exp(dots - dots_logsumexp)

From a64762863d16427096719c62dd9407268e40ce70 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 16 Sep 2019 22:33:31 -0700
Subject: [PATCH 2451/2720] Add share_qk and hard_k flags to
 MemoryEfficientCausalAttention

PiperOrigin-RevId: 269495331
---
 tensor2tensor/trax/layers/attention.py | 45 +++++++++++++++++++++++++-
 1 file changed, 44 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index f57a2b94b..e25dc4198 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -409,7 +409,7 @@ class MemoryEfficientCausalAttention(BaseCausalAttention):
   option. We haven't implemented it, but it may perform well, too.
   """
 
-  def __init__(self, loop_stride, dropout, mode):
+  def __init__(self, loop_stride, dropout, mode, share_qk=False, hard_k=0):
     assert backend.get_name() == 'jax', (
         'JAX backend is required to use MemoryEfficientCausalAttention.')
     super(MemoryEfficientCausalAttention, self).__init__()
@@ -420,6 +420,8 @@ def __init__(self, loop_stride, dropout, mode):
       self.dropout = dropout
     else:
       self.dropout = None
+    self._share_qk = share_qk
+    self._hard_k = hard_k
 
   def call(self, inputs, params=(), state=(), **kwargs):
     del params
@@ -434,6 +436,11 @@ def custom_grad(self, inputs, output, ct, params=(), state=(), **kwargs):
     _, inputs_ct = self.call_and_grad(inputs, ct, **kwargs)
     return inputs_ct, ()
 
+  def make_unit_length(self, x, epsilon=1e-6):
+    variance = np.mean(x**2, axis=-1, keepdims=True)
+    norm_inputs = x / np.sqrt(variance + epsilon)
+    return norm_inputs
+
   def call_and_grad(self, inputs, ct, rng=None, **kwargs):
     del kwargs
     query, key, value = inputs
@@ -462,8 +469,31 @@ def make_mask(N, M, k):  # pylint: disable=invalid-name
       mask = jax.lax.convert_element_type(mask, np.float32)
       return mask
 
+    def make_self_mask(N, M, k):  # pylint: disable=invalid-name
+      """Masks out elements attending to self.
+
+      Args:
+        N: number of query positions
+        M: number of key positions
+        k: position of the initial query element
+
+      Returns:
+        N x M mask, where 1.0 indicates that attention is not allowed.
+      """
+      x = np.arange(N, dtype=np.int32)
+      y = np.arange(M, dtype=np.int32)
+      mask = jax.lax.eq(
+          (jax.lax.broadcast_in_dim(
+              x, shape=(N, M), broadcast_dimensions=(0,)) + k),
+          jax.lax.broadcast(y, [N]))
+      mask = jax.lax.convert_element_type(mask, np.float32)
+      return mask
+
     def forward_slice(query_slice, q_loop_idx, key, value):  # pylint: disable=invalid-name
       """Forward pass for a subset of the query vectors."""
+      if self._share_qk:
+        key = self.make_unit_length(key)
+
       dots = np.matmul(
           query_slice, np.swapaxes(key, -1, -2)) / np.sqrt(depth)
 
@@ -471,6 +501,11 @@ def forward_slice(query_slice, q_loop_idx, key, value):  # pylint: disable=inval
       mask = make_mask(dots.shape[-2], dots.shape[-1], q_loop_idx)
       dots = dots - 1e9 * mask
 
+      # Mask out attention to self except when no other targets are available.
+      if self._share_qk:
+        self_mask = make_self_mask(dots.shape[-2], dots.shape[-1], q_loop_idx)
+        dots = dots - 32 * self_mask
+
       # Softmax.
       dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))
 
@@ -483,6 +518,14 @@ def forward_slice(query_slice, q_loop_idx, key, value):  # pylint: disable=inval
         multiplier = keep.astype(dots.dtype) / jax.lax.tie_in(keep, keep_prob)
         dots = dots * multiplier
 
+      if self._hard_k > 0:
+        top_k = np.sort(dots)[..., -self._hard_k]  # Get the top-kth weight.
+        top_k = jax.lax.stop_gradient(top_k)
+        dots -= top_k[..., np.newaxis]  # Subtract (be 0 for lower ones).
+        dots = np.maximum(dots, 0)
+        dots_sum = np.sum(dots, axis=-1, keepdims=True)  # Re-normalize.
+        dots /= dots_sum  # Re-normalize.
+
       out_slice = np.matmul(dots, value)
       return out_slice
 

From 20713dae19c65bd228f9fb285ad43cb4d8cfdbea Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 17 Sep 2019 15:31:16 -0700
Subject: [PATCH 2452/2720] Move sequence serialization functions to a new
 module.

This will be useful for extending policy networks to work on serialized sequences.

PiperOrigin-RevId: 269669250
---
 tensor2tensor/trax/rl/serialization_utils.py  | 184 ++++++++++++++++++
 .../trax/rl/serialization_utils_test.py       | 169 ++++++++++++++++
 .../trax/rl/simulated_env_problem.py          |  55 +++---
 3 files changed, 377 insertions(+), 31 deletions(-)
 create mode 100644 tensor2tensor/trax/rl/serialization_utils.py
 create mode 100644 tensor2tensor/trax/rl/serialization_utils_test.py

diff --git a/tensor2tensor/trax/rl/serialization_utils.py b/tensor2tensor/trax/rl/serialization_utils.py
new file mode 100644
index 000000000..e175f4460
--- /dev/null
+++ b/tensor2tensor/trax/rl/serialization_utils.py
@@ -0,0 +1,184 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Utilities for serializing trajectories into discrete sequences."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+
+def serialize_observations_and_actions(
+    observations,
+    actions,
+    mask,
+    observation_serializer,
+    action_serializer,
+    representation_length,
+):
+  """Serializes observations and actions into a discrete sequence.
+
+  Args:
+    observations: Array (B, T + 1, ...), of observations, where B is the batch
+      size and T is the number of timesteps excluding the last observation.
+    actions: Array (B, T, ...) of actions.
+    mask: Binary array (B, T) indicating where each sequence ends (1s while
+      it continues).
+    observation_serializer: SpaceSerializer for observations.
+    action_serializer: SpaceSerializer for actions.
+    representation_length: Number of symbols in the serialized sequence. The
+      sequence is padded up to this number.
+  Returns:
+    Pair (representation, mask), where representation is the serialized sequence
+    of shape (B, R) where R = representation_length, and mask is a binary array
+    of shape (B, R) indicating where each sequence ends.
+  """
+  (batch_size, n_timesteps) = actions.shape[:2]
+  assert observations.shape[:2] == (batch_size, n_timesteps + 1)
+  assert mask.shape == (batch_size, n_timesteps)
+
+  reprs = []
+  for t in range(n_timesteps):
+    reprs.append(observation_serializer.serialize(observations[:, t, ...]))
+    reprs.append(action_serializer.serialize(actions[:, t, ...]))
+  reprs.append(observation_serializer.serialize(observations[:, -1, ...]))
+  reprs = np.concatenate(reprs, axis=1)
+  assert reprs.shape[1] <= representation_length
+  reprs = np.pad(
+      reprs,
+      pad_width=((0, 0), (0, representation_length - reprs.shape[1])),
+      mode="constant",
+  )
+
+  obs_repr_length = observation_serializer.representation_length
+  act_repr_length = action_serializer.representation_length
+  step_repr_length = obs_repr_length + act_repr_length
+  seq_lengths = np.sum(mask, axis=1).astype(np.int32)
+  repr_lengths = seq_lengths * step_repr_length + obs_repr_length
+  repr_mask = np.zeros((batch_size, representation_length), dtype=np.int32)
+  for (i, repr_length) in enumerate(repr_lengths):
+    repr_mask[i, :repr_length] = 1
+
+  return (reprs, repr_mask)
+
+
+def observation_mask(
+    observation_serializer, action_serializer, representation_length
+):
+  """Calculates an observation mask for a serialized sequence.
+
+  Args:
+    observation_serializer: SpaceSerializer for observations.
+    action_serializer: SpaceSerializer for actions.
+    representation_length: Number of symbols in the serialized sequence. The
+      mask is padded up to this number.
+
+  Returns:
+    Binary mask indicating which symbols in the representation correspond to
+    observations.
+  """
+  mask = np.zeros(representation_length)
+  obs_repr_length = observation_serializer.representation_length
+  step_repr_length = obs_repr_length + action_serializer.representation_length
+  for step_start_index in range(0, representation_length, step_repr_length):
+    mask[step_start_index:(step_start_index + obs_repr_length)] = 1
+  return mask
+
+
+def action_mask(
+    observation_serializer, action_serializer, representation_length
+):
+  """Calculates an action mask for a serialized sequence.
+
+  Args:
+    observation_serializer: SpaceSerializer for observations.
+    action_serializer: SpaceSerializer for actions.
+    representation_length: Number of symbols in the serialized sequence. The
+      mask is padded up to this number.
+
+  Returns:
+    Binary mask indicating which symbols in the representation correspond to
+    actions.
+  """
+  return 1 - observation_mask(
+      observation_serializer, action_serializer, representation_length
+  )
+
+
+def significance_map(
+    observation_serializer, action_serializer, representation_length
+):
+  """Calculates a significance map for the entire serialized sequence.
+
+  See SpaceSerializer.significance_map.
+
+  Args:
+    observation_serializer: SpaceSerializer for observations.
+    action_serializer: SpaceSerializer for actions.
+    representation_length: Number of symbols in the serialized sequence. The
+      significance map is padded up to this number.
+
+  Returns:
+    Significance map for the entire serialized sequence.
+  """
+  sig_map = np.zeros(representation_length, dtype=np.int32)
+  obs_repr_length = observation_serializer.representation_length
+  act_repr_length = action_serializer.representation_length
+  step_repr_length = obs_repr_length + act_repr_length
+  for step_start_index in range(0, representation_length, step_repr_length):
+    act_start_index = step_start_index + obs_repr_length
+    step_end_index = step_start_index + step_repr_length
+    limit = representation_length - step_start_index
+    sig_map[step_start_index:act_start_index] = (
+        observation_serializer.significance_map[:limit]
+    )
+    limit = representation_length - act_start_index
+    sig_map[act_start_index:step_end_index] = (
+        action_serializer.significance_map[:limit]
+    )
+  return sig_map
+
+
+def rewards_to_actions_map(
+    observation_serializer,
+    action_serializer,
+    n_timesteps,
+    representation_length,
+):
+  """Calculates a mapping between the rewards and the serialized sequence.
+
+  Used to broadcast advantages over the log-probabilities of corresponding
+  actions.
+
+  Args:
+    observation_serializer: SpaceSerializer for observations.
+    action_serializer: SpaceSerializer for actions.
+    n_timesteps: Number of timesteps (length of the reward sequence).
+    representation_length: Number of symbols in the serialized sequence.
+
+  Returns:
+    Array (T, R) translating from the reward sequence to actions in the
+    representation.
+  """
+  r2a_map = np.zeros((n_timesteps, representation_length))
+  obs_repr_length = observation_serializer.representation_length
+  act_repr_length = action_serializer.representation_length
+  step_repr_length = obs_repr_length + act_repr_length
+  for t in range(n_timesteps):
+    act_start_index = t * step_repr_length + obs_repr_length
+    r2a_map[t, act_start_index:(act_start_index + act_repr_length)] = 1
+  return r2a_map
diff --git a/tensor2tensor/trax/rl/serialization_utils_test.py b/tensor2tensor/trax/rl/serialization_utils_test.py
new file mode 100644
index 000000000..46d8985e0
--- /dev/null
+++ b/tensor2tensor/trax/rl/serialization_utils_test.py
@@ -0,0 +1,169 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.trax.rl.serialization_utils."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gin
+import gym
+import numpy as np
+
+from tensor2tensor.trax.rl import serialization_utils
+from tensor2tensor.trax.rl import space_serializer
+from tensorflow import test
+
+
+class SerializationTest(test.TestCase):
+
+  def setUp(self):
+    super(SerializationTest, self).setUp()
+    self._serializer = space_serializer.create(
+        gym.spaces.Discrete(2), vocab_size=2
+    )
+    self._repr_length = 100
+    self._serialization_utils_kwargs = {
+        "observation_serializer": self._serializer,
+        "action_serializer": self._serializer,
+        "representation_length": self._repr_length,
+    }
+
+  def test_serializes_observations_and_actions(self):
+    (reprs, mask) = serialization_utils.serialize_observations_and_actions(
+        observations=np.array([[0, 1]]),
+        actions=np.array([[0]]),
+        mask=np.array([[1]]),
+        **self._serialization_utils_kwargs
+    )
+    self.assertEqual(reprs.shape, (1, self._repr_length))
+    self.assertEqual(mask.shape, (1, self._repr_length))
+    self.assertGreater(np.sum(mask), 0)
+    self.assertEqual(np.max(mask), 1)
+
+  def test_masks_length(self):
+    (reprs, mask) = serialization_utils.serialize_observations_and_actions(
+        observations=np.array([[0, 1, 0], [0, 1, 0], [0, 1, 1]]),
+        actions=np.array([[0, 0], [0, 1], [0, 0]]),
+        mask=np.array([[1, 0], [1, 1], [1, 1]]),
+        **self._serialization_utils_kwargs
+    )
+    # Trajectories 1 and 2 are longer than 0.
+    self.assertGreater(np.sum(mask[1]), np.sum(mask[0]))
+    self.assertGreater(np.sum(mask[2]), np.sum(mask[0]))
+    # Trajectory 0 is a common prefix of 1 and 2. 1 and 2 are different.
+    np.testing.assert_array_equal(reprs[0] * mask[0], reprs[1] * mask[0])
+    np.testing.assert_array_equal(reprs[0] * mask[0], reprs[2] * mask[0])
+    self.assertFalse(np.array_equal(reprs[1] * mask[1], reprs[2] * mask[2]))
+    # Trajectories should be padded with 0s.
+    np.testing.assert_array_equal(
+        reprs * (1 - mask), np.zeros((3, self._repr_length))
+    )
+
+  def test_observation_and_action_masks_are_valid_and_complementary(self):
+    obs_mask = serialization_utils.observation_mask(
+        **self._serialization_utils_kwargs
+    )
+    self.assertEqual(obs_mask.shape, (self._repr_length,))
+    self.assertEqual(np.min(obs_mask), 0)
+    self.assertEqual(np.max(obs_mask), 1)
+
+    act_mask = serialization_utils.action_mask(
+        **self._serialization_utils_kwargs
+    )
+    self.assertEqual(act_mask.shape, (self._repr_length,))
+    self.assertEqual(np.min(act_mask), 0)
+    self.assertEqual(np.max(act_mask), 1)
+
+    np.testing.assert_array_equal(
+        obs_mask + act_mask, np.ones(self._repr_length)
+    )
+
+  def test_masks_observations(self):
+    (reprs, _) = serialization_utils.serialize_observations_and_actions(
+        # Observations are different, actions are the same.
+        observations=np.array([[0, 1], [1, 1]]),
+        actions=np.array([[0], [0]]),
+        mask=np.array([[1], [1]]),
+        **self._serialization_utils_kwargs
+    )
+    obs_mask = serialization_utils.observation_mask(
+        **self._serialization_utils_kwargs
+    )
+    act_mask = serialization_utils.action_mask(
+        **self._serialization_utils_kwargs
+    )
+
+    self.assertFalse(np.array_equal(reprs[0] * obs_mask, reprs[1] * obs_mask))
+    np.testing.assert_array_equal(reprs[0] * act_mask, reprs[1] * act_mask)
+
+  def test_masks_actions(self):
+    (reprs, _) = serialization_utils.serialize_observations_and_actions(
+        # Observations are the same, actions are different.
+        observations=np.array([[0, 1], [0, 1]]),
+        actions=np.array([[0], [1]]),
+        mask=np.array([[1], [1]]),
+        **self._serialization_utils_kwargs
+    )
+    obs_mask = serialization_utils.observation_mask(
+        **self._serialization_utils_kwargs
+    )
+    act_mask = serialization_utils.action_mask(
+        **self._serialization_utils_kwargs
+    )
+
+    np.testing.assert_array_equal(reprs[0] * obs_mask, reprs[1] * obs_mask)
+    self.assertFalse(np.array_equal(reprs[0] * act_mask, reprs[1] * act_mask))
+
+  def test_significance_map(self):
+    gin.bind_parameter("BoxSpaceSerializer.precision", 3)
+    significance_map = serialization_utils.significance_map(
+        observation_serializer=space_serializer.create(
+            gym.spaces.Box(low=0, high=1, shape=(2,)), vocab_size=2
+        ),
+        action_serializer=space_serializer.create(
+            gym.spaces.MultiDiscrete(nvec=[2, 2]), vocab_size=2
+        ),
+        representation_length=20,
+    )
+    np.testing.assert_array_equal(
+        significance_map,
+        # obs1, act1, obs2, act2, obs3 cut after 4th symbol.
+        [0, 1, 2, 0, 1, 2, 0, 0, 0, 1, 2, 0, 1, 2, 0, 0, 0, 1, 2, 0],
+    )
+
+  def test_rewards_to_actions_map(self):
+    rewards = np.array([1, 2, 3])
+    r2a_map = serialization_utils.rewards_to_actions_map(
+        observation_serializer=space_serializer.create(
+            gym.spaces.MultiDiscrete(nvec=[2, 2, 2]), vocab_size=2
+        ),
+        action_serializer=space_serializer.create(
+            gym.spaces.MultiDiscrete(nvec=[2, 2]), vocab_size=2
+        ),
+        n_timesteps=len(rewards),
+        representation_length=16,
+    )
+    broadcast_rewards = np.dot(rewards, r2a_map)
+    np.testing.assert_array_equal(
+        broadcast_rewards,
+        # obs1, act1, obs2, act2, obs3 cut after 1st symbol.
+        [0, 0, 0, 1, 1, 0, 0, 0, 2, 2, 0, 0, 0, 3, 3, 0],
+    )
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index 454f9c865..2dbc752b0 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -29,6 +29,7 @@
 from tensor2tensor.trax import trax
 from tensor2tensor.trax import utils
 from tensor2tensor.trax.backend import random as jax_random
+from tensor2tensor.trax.rl import serialization_utils
 from tensor2tensor.trax.rl import space_serializer
 
 
@@ -419,37 +420,29 @@ def _step_model(self, predict_fn, actions, rng):
     return (observation, reward, done)
 
   def trajectory_to_training_examples(self, trajectory):
-    reprs = []
-    weights = []
-    for time_step in trajectory.time_steps:
-      # Serializers work on batches.
-      obs_repr = self._obs_serializer.serialize(
-          np.array([time_step.observation]))[0]
-      reprs.append(obs_repr)
-      # significance_map is an array of the same size as the representation,
-      # indicating the significance of each symbol. See
-      # SpaceSerializer.significance_map.
-      weights.append(
-          self._significance_decay ** self._obs_serializer.significance_map)
-      if time_step.action is not None:
-        action_repr = self._action_serializer.serialize(
-            np.array([time_step.action]))[0]
-        reprs.append(action_repr)
-        weights.append(np.zeros_like(action_repr))
-
-    def concat_and_pad(arrays):
-      (desired_length,) = self.model_input_shape
-      flat_array = np.concatenate(arrays, axis=0)
-      (actual_length,) = flat_array.shape
-      assert actual_length <= desired_length
-      return np.pad(
-          flat_array,
-          pad_width=((0, desired_length - actual_length),),
-          mode="constant",
-      )
-    (reprs, weights) = map(concat_and_pad, (reprs, weights))
-    reprs = reprs.astype(self.model_input_dtype)
-    return [(reprs, reprs, weights)]  # (inputs, targets, weights)
+    (repr_length,) = self.model_input_shape
+    seq_mask = np.ones((1, trajectory.num_time_steps - 1))
+    (reprs, repr_mask) = serialization_utils.serialize_observations_and_actions(
+        # Serialization works on batches, so we add a singleton batch dimension.
+        trajectory.observations_np[None, ...],
+        trajectory.actions_np[None, ...],
+        seq_mask,
+        self._obs_serializer,
+        self._action_serializer,
+        repr_length,
+    )
+    reprs = reprs[0, ...].astype(self.model_input_dtype)
+    sig_weights = (
+        self._significance_decay ** serialization_utils.significance_map(
+            self._obs_serializer, self._action_serializer, repr_length
+        )[None, ...]
+    )
+    obs_mask = serialization_utils.observation_mask(
+        self._obs_serializer, self._action_serializer, repr_length
+    )
+    weights = (sig_weights * obs_mask * repr_mask)[0, ...]
+    # (inputs, targets, weights)
+    return [(reprs, reprs, weights)]
 
   @property
   def model_input_shape(self):

From f6db92a905ccbf7047cae3f187a2b6a09ff36e1a Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 17 Sep 2019 16:01:57 -0700
Subject: [PATCH 2453/2720] Add rl libs to async_trajectory_collector.

PiperOrigin-RevId: 269675330
---
 tensor2tensor/trax/rl/configs/simple_online_tune.gin     | 5 -----
 tensor2tensor/trax/rl/envs/async_trajectory_collector.py | 2 ++
 tensor2tensor/trax/rl/ppo.py                             | 9 ++++++---
 tensor2tensor/trax/rl/simple.py                          | 6 +++---
 4 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/trax/rl/configs/simple_online_tune.gin b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
index 4df69ceb5..8378311d7 100644
--- a/tensor2tensor/trax/rl/configs/simple_online_tune.gin
+++ b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
@@ -85,8 +85,3 @@ TransformerLM.max_len = 1024
 train.eval_frequency = 1000
 train.optimizer = @trax.optimizers.Adafactor
 
-# Parameters for train_rl:
-# ==============================================================================
-train_rl.env_name = "ClientEnv-v0"
-train_rl.n_epochs = 100
-train_rl.trainer_class = @trax.rl.trainers.SimPLe
diff --git a/tensor2tensor/trax/rl/envs/async_trajectory_collector.py b/tensor2tensor/trax/rl/envs/async_trajectory_collector.py
index f94e7162a..7eb79131c 100644
--- a/tensor2tensor/trax/rl/envs/async_trajectory_collector.py
+++ b/tensor2tensor/trax/rl/envs/async_trajectory_collector.py
@@ -28,6 +28,8 @@
 from jax.config import config
 from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.rl.google import atari_utils  # GOOGLE-INTERNAL:
+from tensor2tensor.trax import rl  # pylint: disable=unused-import
+from tensor2tensor.trax.rl import envs as rl_envs  # pylint: disable=unused-import
 from tensor2tensor.trax.rl.envs import async_trajectory_collector_lib as async_lib
 import tensorflow as tf
 
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index ddb7ae417..3e9e49aa4 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -855,6 +855,9 @@ def maybe_restore_opt_state(output_dir,
   )
 
 
+LAST_N_POLICY_MODELS_TO_KEEP = 5
+
+
 def save_opt_state(output_dir,
                    policy_and_value_opt_state,
                    policy_and_value_state,
@@ -867,9 +870,9 @@ def save_opt_state(output_dir,
   with gfile.GFile(params_file, "wb") as f:
     pkl_module.dump(
         (policy_and_value_opt_state, policy_and_value_state, total_opt_step), f)
-  # Remove the old model files, leave the latest one (it might be in the
-  # process of getting read async) -- this will get cleaned up later.
-  for path in old_model_files[1:]:
+  # Keep the last k model files lying around (note k > 1 because the latest
+  # model file might be in the process of getting read async).
+  for path in old_model_files[LAST_N_POLICY_MODELS_TO_KEEP:]:
     if path != params_file:
       gfile.remove(path)
 
diff --git a/tensor2tensor/trax/rl/simple.py b/tensor2tensor/trax/rl/simple.py
index 35b6d501f..f09f616a2 100644
--- a/tensor2tensor/trax/rl/simple.py
+++ b/tensor2tensor/trax/rl/simple.py
@@ -51,9 +51,9 @@ def load_trajectories(trajectory_dir, eval_frac):
             "Could not load trajectories from a corrupted shard %s.",
             shard_path,
         )
-  assert train_trajectories, "Haven't found any training data."
-  assert eval_trajectories, "Haven't found any evaluation data."
-  return (train_trajectories, eval_trajectories)
+  assert train_trajectories, "Can't find training data in %s" % trajectory_dir
+  assert eval_trajectories, "Can't find evaluation data in %s" % trajectory_dir
+  return train_trajectories, eval_trajectories
 
 
 def generate_examples(trajectories, trajectory_to_training_examples_fn):

From 77086dae3388593887f08199e97e2df6a242706b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 18 Sep 2019 08:26:23 -0700
Subject: [PATCH 2454/2720] Internal

PiperOrigin-RevId: 269810727
---
 tensor2tensor/utils/expert_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index c5b7f1875..f177f2ac9 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -65,7 +65,7 @@ def decorated(*args, **kwargs):
 
 
 def add_var_scope(scope=None):
-  return add_scope(scope, scope_fn=tf.variable_scope)
+  return add_scope(scope, scope_fn=tf.compat.v1.variable_scope)
 
 
 def add_name_scope(scope=None):

From 5fc58e361337adace0584a5e099b65f54f7ce2a1 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 18 Sep 2019 09:01:55 -0700
Subject: [PATCH 2455/2720] Since actions can now be np.array, "if action" is
 ambiguous and crashes.

PiperOrigin-RevId: 269817463
---
 tensor2tensor/trax/rl/simple.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/trax/rl/simple.py b/tensor2tensor/trax/rl/simple.py
index f09f616a2..8c3868160 100644
--- a/tensor2tensor/trax/rl/simple.py
+++ b/tensor2tensor/trax/rl/simple.py
@@ -198,11 +198,11 @@ def __call__(self, observations):
     del observations
 
     def get_action(traj):
+      action = None
       if self._step < traj.num_time_steps:
         action = traj.time_steps[self._step].action
-      else:
-        action = None
-      return action or self._out_of_bounds_action
+        # PS: action can still be None, if this is the last time-step in traj.
+      return action if action is not None else self._out_of_bounds_action
     actions = np.array(list(map(get_action, self._trajectories)))
     self._step += 1
     return actions

From 34cfea2406c3c5ce71212330a94c4a33d88d0c68 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 18 Sep 2019 12:33:28 -0700
Subject: [PATCH 2456/2720] Untie the reward and action sequence in PPO, so
 that controls are flattened into time.

This is an incremental change towards enabling the policy to run on a serialized representation. It introuduces rewards_to_actions, a linear map between the reward sequence and the action sequence, used to scatter advantages over action log-probabilities.
Breaking changes:
 - max_timesteps is now mandatory
 - commented out collecting log probabilities and value predictions during stepping the environment, so PPO recomputes this data during training, hence might not work with anymore with non-deterministic policies (but we don't use them anyway)

PiperOrigin-RevId: 269865846
---
 tensor2tensor/envs/env_problem_utils.py      |  37 +--
 tensor2tensor/envs/env_problem_utils_test.py |   6 +-
 tensor2tensor/trax/rl/ppo.py                 | 240 ++++++++++---------
 tensor2tensor/trax/rl/ppo_test.py            |  80 +++----
 tensor2tensor/trax/rl/ppo_trainer.py         | 121 +++++++---
 tensor2tensor/trax/rl/ppo_trainer_test.py    |   2 +-
 6 files changed, 285 insertions(+), 201 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 6fbfdc6a3..99d0e62eb 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -72,6 +72,7 @@ def get_completed_trajectories_from_env(env,
 
 def play_env_problem_with_policy(env,
                                  policy_fun,
+                                 action_index_fn,
                                  num_trajectories=1,
                                  max_timestep=None,
                                  reset=True,
@@ -87,8 +88,10 @@ def play_env_problem_with_policy(env,
 
   Args:
     env: environment object, should be a subclass of env_problem.EnvProblem.
-    policy_fun: callable, taking in observations((B, T) + OBS) and returning
-      back log-probabilities (B, T, C, A).
+    policy_fun: callable, taking in observations((B, RT) + OBS) and returning
+      back log-probabilities (B, AT, A).
+    action_index_fn: function converting timestep indices into indices in the
+      log-probability array.
     num_trajectories: int, number of trajectories to collect.
     max_timestep: int or None, if not None or a negative number, we cut any
       trajectory that exceeds this time put it in the completed bin, and *dont*
@@ -141,13 +144,13 @@ def gumbel_sample(log_probs):
       return None, 0, {}, state
 
     # Get all the observations for all the active trajectories.
-    # Shape is (B, T) + OBS
+    # Shape is (B, RT) + OBS
     # Bucket on whatever length is needed.
     padded_observations, lengths = env.trajectories.observations_np(
         boundary=boundary,
         len_history_for_policy=len_history_for_policy)
 
-    B, T = padded_observations.shape[:2]  # pylint: disable=invalid-name
+    B = padded_observations.shape[0]  # pylint: disable=invalid-name
 
     assert B == env.batch_size
     assert (B,) == lengths.shape
@@ -157,22 +160,22 @@ def gumbel_sample(log_probs):
         padded_observations, state=state, rng=rng)
     policy_application_total_time += (time.time() - t1)
 
-    assert (B, T) == log_prob_actions.shape[:2]
-    C = log_prob_actions.shape[2]  # pylint: disable=invalid-name
-    A = log_prob_actions.shape[3]  # pylint: disable=invalid-name
+    assert B == log_prob_actions.shape[0]
+    (_, A) = log_prob_actions.shape[1:]  # pylint: disable=invalid-name
 
     # We need the log_probs of those actions that correspond to the last actual
     # time-step.
     index = lengths - 1  # Since we want to index using lengths.
-    log_probs = log_prob_actions[np.arange(B)[:, None, None],
-                                 index[:, None, None],
-                                 np.arange(C)[:, None], np.arange(A)]
-    value_preds = value_predictions[np.arange(B)[:, None], index[:, None],
-                                    np.arange(1)]
-    assert (B, C, A) == log_probs.shape, \
-        "B=%d, C=%d, A=%d, log_probs.shape=%s" % (B, C, A, log_probs.shape)
-    assert (B, 1) == value_preds.shape, \
-        "B=%d, value_preds.shape=%s" % (B, value_preds.shape)
+    pred_index = action_index_fn(index)
+    log_probs = log_prob_actions[
+        np.arange(B)[:, None, None],
+        pred_index[:, :, None],
+        np.arange(A),
+    ]
+    value_preds = value_predictions[np.arange(B)[:, None], pred_index]
+
+    assert B == log_probs.shape[0]
+    assert A == log_probs.shape[2]
 
     actions = gumbel_sample(log_probs)
     if isinstance(env.action_space, gym.spaces.Discrete):
@@ -184,7 +187,7 @@ def gumbel_sample(log_probs):
         actions,
         infos={
             "log_prob_actions": log_probs,
-            "value_predictions": value_preds
+            "value_predictions": value_preds,
         })
     env_actions_total_time += (time.time() - t1)
     bare_env_run_time += sum(
diff --git a/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index 6304a3d16..cb15941c8 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -62,16 +62,20 @@ def policy_fun(observations, state=None, rng=None):
       self.assertFalse(
           np.all(observations[:, :len_history_for_policy, ...] == 0))
       a = env.action_space.n
-      p = np.random.uniform(size=(b, t, 1, a))
+      p = np.random.uniform(size=(b, t, a))
       p = np.exp(p)
       p = p / np.sum(p, axis=-1, keepdims=True)
       return np.log(p), np.mean(p, axis=-1), state, rng
 
+    def action_index_fn(index):
+      return index[:, None]
+
     max_timestep = 15
     num_trajectories = 2
     trajectories, _, _, _ = env_problem_utils.play_env_problem_with_policy(
         env,
         policy_fun,
+        action_index_fn=action_index_fn,
         num_trajectories=num_trajectories,
         max_timestep=max_timestep,
         len_history_for_policy=len_history_for_policy)
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index 3e9e49aa4..db5b2e1ef 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -17,19 +17,20 @@
 
 Notation:
 
-B, scalar  - batch size
-T, scalar  - number of time-steps in a trajectory, or the value of the padded
-             time-step dimension.
-OBS, tuple - shape of a singular observation from the environment.
+B, scalar   - batch size
+RT, scalar  - (reward time) number of time-steps in a trajectory, or the size
+              of the padded reward sequence.
+AT, scalar  - (action time) number of controls in a trajectory, or the size
+              of the policy network output.
+OBS, tuple  - shape of a singular observation from the environment.
              Ex: For CartPole-v0 this is (4,) and Pong-v0 it's (210, 160, 3)
-C, scalar  - Number of controls, i.e. independent groups of actions.
-A, scalar  - Number of actions, assuming a discrete space.
+A, scalar   - Number of actions, assuming a discrete space.
 
 Policy and Value function signatures:
 
-Policy            Function :: [B, T] + OBS ->  [B, T, C, A]
-Value             Function :: [B, T] + OBS ->  [B, T, 1]
-Policy and Value  Function :: [B, T] + OBS -> ([B, T, C, A], [B, T, 1])
+Policy            Function :: [B, RT + 1] + OBS ->  [B, AT, A]
+Value             Function :: [B, RT + 1] + OBS ->  [B, AT]
+Policy and Value  Function :: [B, RT + 1] + OBS -> ([B, AT, A], [B, AT])
 
 i.e. the policy net should take a batch of *trajectories* and at each time-step
 in each batch deliver a probability distribution over actions.
@@ -80,9 +81,9 @@ def policy_and_value_net(n_actions, n_controls, bottom_layers_fn, two_towers):
   # NOTE: The LogSoftmax instead of the Softmax because of numerical stability.
 
   @tl.layer()
-  def SplitActions(x, **unused_kwargs):  # pylint: disable=invalid-name
-    """Splits logits for actions in different controls."""
-    return np.reshape(x, x.shape[:-1] + (n_controls, n_actions))
+  def FlattenControlsIntoTime(x, **unused_kwargs):  # pylint: disable=invalid-name
+    """Splits logits for actions in different controls and flattens controls."""
+    return np.reshape(x, (x.shape[0], -1, n_actions))
 
   n_logits = n_controls * n_actions
 
@@ -92,9 +93,9 @@ def SplitActions(x, **unused_kwargs):  # pylint: disable=invalid-name
         tl.Parallel(
             [bottom_layers_fn(),
              tl.Dense(n_logits),
-             SplitActions(),  # pylint: disable=no-value-for-parameter
+             FlattenControlsIntoTime(),  # pylint: disable=no-value-for-parameter
              tl.LogSoftmax()],
-            [bottom_layers_fn(), tl.Dense(1)],
+            [bottom_layers_fn(), tl.Dense(n_controls), tl.Flatten()],
         )
     ]
   else:
@@ -103,9 +104,9 @@ def SplitActions(x, **unused_kwargs):  # pylint: disable=invalid-name
         tl.Dup(),
         tl.Parallel(
             [tl.Dense(n_logits),
-             SplitActions(),  # pylint: disable=no-value-for-parameter
+             FlattenControlsIntoTime(),  # pylint: disable=no-value-for-parameter
              tl.LogSoftmax()],
-            [tl.Dense(1)],
+            [tl.Dense(n_controls), tl.Flatten()],
         )
     ]
   return tl.Model(layers)
@@ -148,6 +149,7 @@ def get_params(opt_state):
 # Any other option?
 def collect_trajectories(env,
                          policy_fn,
+                         action_index_fn,
                          n_trajectories=1,
                          max_timestep=None,
                          reset=True,
@@ -162,7 +164,9 @@ def collect_trajectories(env,
 
   Args:
     env: A gym env interface, for now this is not-batched.
-    policy_fn: observations(B,T+1) -> log-probabs(B,T+1, A) callable.
+    policy_fn: observations(B,RT+1) -> log-probabs(B, AT, A) callable.
+    action_index_fn: function converting timestep indices into indices in the
+      log-probability array.
     n_trajectories: int, number of trajectories.
     max_timestep: int or None, the index of the maximum time-step at which we
       return the trajectory, None for ending a trajectory only when env returns
@@ -187,7 +191,7 @@ def collect_trajectories(env,
     trajectory: list of (observation, action, reward) tuples, where each element
     `i` is a tuple of numpy arrays with shapes as follows:
     observation[i] = (B, T_i + 1)
-    action[i] = (B, T_i, C)
+    action[i] = (B, T_i)
     reward[i] = (B, T_i)
   """
 
@@ -196,6 +200,7 @@ def collect_trajectories(env,
   trajs, n_done, timing_info, state = env_problem_utils.play_env_problem_with_policy(
       env,
       policy_fn,
+      action_index_fn,
       num_trajectories=n_trajectories,
       max_timestep=max_timestep,
       reset=reset,
@@ -247,9 +252,9 @@ def pad_trajectories(trajectories, boundary=20):
 
   Returns:
     tuple: (padding lengths, reward_mask, padded_observations, padded_actions,
-        padded_rewards) where padded_observations is shaped (B, T+1) + OBS and
-        padded_actions, padded_rewards & reward_mask are shaped (B, T).
-        Where T is max(t) rounded up to an integer multiple of boundary.
+        padded_rewards) where padded_observations is shaped (B, RT+1) + OBS and
+        padded_actions, padded_rewards & reward_mask are shaped (B, RT).
+        Where RT is max(t) rounded up to an integer multiple of boundary.
         padded_length is how much padding we've added and
         reward_mask is 1s for actual rewards and 0s for the padding.
   """
@@ -332,16 +337,16 @@ def rewards_to_go(rewards, mask, gamma=0.99):
   r2g_t = \sum_{l=0}^{\infty} (\gamma^{l} * reward_{t+l})
 
   Args:
-    rewards: np.ndarray of shape (B, T) of rewards.
-    mask: np.ndarray of shape (B, T) of mask for the rewards.
+    rewards: np.ndarray of shape (B, RT) of rewards.
+    mask: np.ndarray of shape (B, RT) of mask for the rewards.
     gamma: float, discount factor.
 
   Returns:
-    rewards to go, np.ndarray of shape (B, T).
+    rewards to go, np.ndarray of shape (B, RT).
   """
-  B, T = rewards.shape  # pylint: disable=invalid-name,unused-variable
+  B, RT = rewards.shape  # pylint: disable=invalid-name,unused-variable
 
-  masked_rewards = rewards * mask  # (B, T)
+  masked_rewards = rewards * mask  # (B, RT)
 
   # The lax.scan version of this is slow, but we still show it here for
   # completeness.
@@ -379,14 +384,14 @@ def rewards_to_go(rewards, mask, gamma=0.99):
   r2gs = [masked_rewards[:, -1]]
 
   # Go from T-2 down to 0.
-  for t in reversed(range(T - 1)):
+  for t in reversed(range(RT - 1)):
     r2gs.append(masked_rewards[:, t] + (gamma * r2gs[-1]))
 
-  # The list should have length T.
-  assert T == len(r2gs)
+  # The list should have length RT.
+  assert RT == len(r2gs)
 
-  # First we stack them in the correct way to make it (B, T), but these are
-  # still from newest (T-1) to oldest (0), so then we flip it on time axis.
+  # First we stack them in the correct way to make it (B, RT), but these are
+  # still from newest (RT-1) to oldest (0), so then we flip it on time axis.
   return np.flip(np.stack(r2gs, axis=1), axis=1)
 
 
@@ -400,12 +405,12 @@ def value_loss_given_predictions(value_prediction,
   """Computes the value loss given the prediction of the value function.
 
   Args:
-    value_prediction: np.ndarray of shape (B, T+1, 1)
-    rewards: np.ndarray of shape (B, T) of rewards.
-    reward_mask: np.ndarray of shape (B, T), the mask over rewards.
+    value_prediction: np.ndarray of shape (B, RT+1, 1)
+    rewards: np.ndarray of shape (B, RT) of rewards.
+    reward_mask: np.ndarray of shape (B, RT), the mask over rewards.
     gamma: float, discount factor.
     epsilon: float, clip-fraction, used if value_value_prediction_old isn't None
-    value_prediction_old: np.ndarray of shape (B, T+1, 1) of value predictions
+    value_prediction_old: np.ndarray of shape (B, RT+1, 1) of value predictions
       using the old parameters. If provided, we incorporate this in the loss as
       well. This is from the OpenAI baselines implementation.
 
@@ -415,19 +420,17 @@ def value_loss_given_predictions(value_prediction,
       summaries collected during value loss computation.
   """
 
-  B, T = rewards.shape  # pylint: disable=invalid-name
-  assert (B, T) == reward_mask.shape
-  assert (B, T + 1, 1) == value_prediction.shape
+  B, RT = rewards.shape  # pylint: disable=invalid-name
+  assert (B, RT) == reward_mask.shape
+  assert (B, RT + 1) == value_prediction.shape
 
-  value_prediction = np.squeeze(value_prediction, axis=2)  # (B, T+1)
-  value_prediction = value_prediction[:, :-1] * reward_mask  # (B, T)
-  r2g = rewards_to_go(rewards, reward_mask, gamma=gamma)  # (B, T)
+  value_prediction = value_prediction[:, :-1] * reward_mask  # (B, RT)
+  r2g = rewards_to_go(rewards, reward_mask, gamma=gamma)  # (B, RT)
   loss = (value_prediction - r2g)**2
 
   # From the baselines implementation.
   if value_prediction_old is not None:
-    value_prediction_old = np.squeeze(value_prediction_old, axis=2)  # (B, T+1)
-    value_prediction_old = value_prediction_old[:, :-1] * reward_mask  # (B, T)
+    value_prediction_old = value_prediction_old[:, :-1] * reward_mask  # (B, RT)
 
     v_clipped = value_prediction_old + np.clip(
         value_prediction - value_prediction_old, -epsilon, epsilon)
@@ -452,19 +455,19 @@ def deltas(predicted_values, rewards, mask, gamma=0.99):
   delta_{b,t} = r_{b,t} + \gamma * v_{b,t+1} - v_{b,t}.
 
   Args:
-    predicted_values: ndarray of shape (B, T+1). NOTE: Expects axis 2 was
-      squeezed. These represent V(s_bt) for b < B and t < T+1
-    rewards: ndarray of shape (B, T) of rewards.
-    mask: ndarray of shape (B, T) of mask for rewards.
+    predicted_values: ndarray of shape (B, RT+1). NOTE: Expects axis 2 was
+      squeezed. These represent V(s_bt) for b < B and t < RT+1
+    rewards: ndarray of shape (B, RT) of rewards.
+    mask: ndarray of shape (B, RT) of mask for rewards.
     gamma: float, discount factor.
 
   Returns:
-    ndarray of shape (B, T) of one-step TD-residuals.
+    ndarray of shape (B, RT) of one-step TD-residuals.
   """
 
-  # Predicted values at time t, cutting off the last to have shape (B, T).
+  # Predicted values at time t, cutting off the last to have shape (B, RT).
   predicted_values_bt = predicted_values[:, :-1]
-  # Predicted values at time t+1, by cutting off the first to have shape (B, T)
+  # Predicted values at time t+1, by cutting off the first to have shape (B, RT)
   predicted_values_btplus1 = predicted_values[:, 1:]
   # Return the deltas as defined above.
   return (rewards +
@@ -481,7 +484,7 @@ def gae_advantages(td_deltas, mask, lambda_=0.95, gamma=0.99):
   Internally we just call rewards_to_go, since it is the same computation.
 
   Args:
-    td_deltas: np.ndarray of shape (B, T) of one step TD-residuals.
+    td_deltas: np.ndarray of shape (B, RT) of one step TD-residuals.
     mask: np.ndarray of shape (B, T) of mask for the residuals. It maybe the
       case that the `td_deltas` are already masked correctly since they are
       produced by `deltas(...)`
@@ -499,52 +502,51 @@ def chosen_probabs(probab_actions, actions):
   """Picks out the probabilities of the actions along batch and time-steps.
 
   Args:
-    probab_actions: ndarray of shape `[B, T+1, C, A]`, where
+    probab_actions: ndarray of shape `[B, AT, A]`, where
       probab_actions[b, t, i] contains the log-probability of action = i at
       the t^th time-step in the b^th trajectory.
-    actions: ndarray of shape `[B, T, C]`, with each entry in [0, A) denoting
+    actions: ndarray of shape `[B, AT]`, with each entry in [0, A) denoting
       which action was chosen in the b^th trajectory's t^th time-step.
 
   Returns:
-    `[B, T, C]` ndarray with the log-probabilities of the chosen actions.
+    `[B, AT, A]` ndarray with the log-probabilities of the chosen actions.
   """
-  B, T, C = actions.shape  # pylint: disable=invalid-name
-  assert (B, T + 1, C) == probab_actions.shape[:3]
-  return probab_actions[
-      np.arange(B)[:, None, None], np.arange(T)[:, None], np.arange(C), actions]
+  B, AT = actions.shape  # pylint: disable=invalid-name
+  assert (B, AT) == probab_actions.shape[:2]
+  return probab_actions[np.arange(B)[:, None], np.arange(AT), actions]
 
 
-def compute_probab_ratios(p_new, p_old, actions, reward_mask):
+def compute_probab_ratios(p_new, p_old, actions, action_mask):
   """Computes the probability ratios for each time-step in a trajectory.
 
   Args:
-    p_new: ndarray of shape [B, T+1, C, A] of the log-probabilities that the
+    p_new: ndarray of shape [B, AT, A] of the log-probabilities that the
       policy network assigns to all the actions at each time-step in each batch
       using the old parameters.
-    p_old: ndarray of shape [B, T+1, C, A], same as above, but using old policy
+    p_old: ndarray of shape [B, AT, A], same as above, but using old policy
       network parameters.
-    actions: ndarray of shape [B, T, C] where each element is from [0, A).
-    reward_mask: ndarray of shape [B, T] masking over probabilities.
+    actions: ndarray of shape [B, AT] where each element is from [0, A).
+    action_mask: ndarray of shape [B, T] masking over probabilities.
 
   Returns:
-    probab_ratios: ndarray of shape [B, T, C], where
-    probab_ratios_{b,t, c} = p_new_{b,t,c,action_{b,t,c}} /
-                             p_old_{b,t,c,action_{b,t,c}}
+    probab_ratios: ndarray of shape [B, AT], where
+    probab_ratios_{b,t,} = p_new_{b,t,action_{b,t}} /
+                           p_old_{b,t,action_{b,t}}
   """
 
-  B, T, C = actions.shape  # pylint: disable=invalid-name
-  assert (B, T + 1, C) == p_old.shape[:3]
-  assert (B, T + 1, C) == p_new.shape[:3]
+  B, AT = actions.shape  # pylint: disable=invalid-name
+  assert (B, AT) == p_old.shape[:2]
+  assert (B, AT) == p_new.shape[:2]
 
   logp_old = chosen_probabs(p_old, actions)
   logp_new = chosen_probabs(p_new, actions)
 
-  assert (B, T, C) == logp_old.shape
-  assert (B, T, C) == logp_new.shape
+  assert (B, AT) == logp_old.shape
+  assert (B, AT) == logp_new.shape
 
   # Since these are log-probabilities, we just subtract them.
-  probab_ratios = np.exp(logp_new - logp_old) * reward_mask[:, :, None]
-  assert (B, T, C) == probab_ratios.shape
+  probab_ratios = np.exp(logp_new - logp_old) * action_mask
+  assert (B, AT) == probab_ratios.shape
   return probab_ratios
 
 
@@ -552,12 +554,12 @@ def clipped_probab_ratios(probab_ratios, epsilon=0.2):
   return np.clip(probab_ratios, 1 - epsilon, 1 + epsilon)
 
 
-def clipped_objective(probab_ratios, advantages, reward_mask, epsilon=0.2):
-  advantages = advantages[:, :, None]
+def clipped_objective(probab_ratios, advantages, action_mask, epsilon=0.2):
+  advantages = advantages
   return np.minimum(
       probab_ratios * advantages,
       clipped_probab_ratios(probab_ratios, epsilon=epsilon) *
-      advantages) * reward_mask[:, :, None]
+      advantages) * action_mask
 
 
 @jit
@@ -565,31 +567,34 @@ def ppo_loss_given_predictions(log_probab_actions_new,
                                log_probab_actions_old,
                                value_predictions_old,
                                padded_actions,
+                               rewards_to_actions,
                                padded_rewards,
                                reward_mask,
                                gamma=0.99,
                                lambda_=0.95,
                                epsilon=0.2):
   """PPO objective, with an eventual minus sign, given predictions."""
-  B, T = padded_rewards.shape  # pylint: disable=invalid-name
-  _, _, C, A = log_probab_actions_old.shape  # pylint: disable=invalid-name
+  B, RT = padded_rewards.shape  # pylint: disable=invalid-name
+  _, AT, A = log_probab_actions_old.shape  # pylint: disable=invalid-name
 
-  assert (B, T) == padded_rewards.shape
-  assert (B, T, C) == padded_actions.shape
-  assert (B, T) == reward_mask.shape
+  assert (B, RT) == padded_rewards.shape
+  assert (B, AT) == padded_actions.shape
+  assert (B, RT) == reward_mask.shape
 
-  assert (B, T + 1, 1) == value_predictions_old.shape
-  assert (B, T + 1, C, A) == log_probab_actions_old.shape
-  assert (B, T + 1, C, A) == log_probab_actions_new.shape
+  assert (B, RT + 1) == value_predictions_old.shape
+  assert (B, AT, A) == log_probab_actions_old.shape
+  assert (B, AT, A) == log_probab_actions_new.shape
 
-  # (B, T)
+  assert (RT + 1, AT) == rewards_to_actions.shape
+
+  # (B, RT)
   td_deltas = deltas(
-      np.squeeze(value_predictions_old, axis=2),  # (B, T+1)
+      value_predictions_old,  # (B, RT+1)
       padded_rewards,
       reward_mask,
       gamma=gamma)
 
-  # (B, T)
+  # (B, RT)
   advantages = gae_advantages(
       td_deltas, reward_mask, lambda_=lambda_, gamma=gamma)
 
@@ -598,18 +603,26 @@ def ppo_loss_given_predictions(log_probab_actions_new,
   advantage_std = np.std(advantages)
   advantages = (advantages - advantage_mean) / (advantage_std + 1e-8)
 
-  # (B, T)
+  # Scatter advantages over padded_actions.
+  # rewards_to_actions is RT + 1 -> AT, so we pad the advantages and the reward
+  # mask by 1.
+  advantages = np.dot(np.pad(advantages, ((0, 0), (0, 1))), rewards_to_actions)
+  action_mask = np.dot(
+      np.pad(reward_mask, ((0, 0), (0, 1))), rewards_to_actions
+  )
+
+  # (B, AT)
   ratios = compute_probab_ratios(log_probab_actions_new, log_probab_actions_old,
-                                 padded_actions, reward_mask)
-  assert (B, T, C) == ratios.shape
+                                 padded_actions, action_mask)
+  assert (B, AT) == ratios.shape
 
-  # (B, T)
+  # (B, AT)
   objective = clipped_objective(
-      ratios, advantages, reward_mask, epsilon=epsilon)
-  assert (B, T, C) == objective.shape
+      ratios, advantages, action_mask, epsilon=epsilon)
+  assert (B, AT) == objective.shape
 
   # ()
-  average_objective = np.sum(objective) / np.sum(reward_mask)
+  average_objective = np.sum(objective) / np.sum(action_mask)
 
   # Loss is negative objective.
   ppo_loss = -average_objective
@@ -629,6 +642,7 @@ def combined_loss_given_predictions(log_probab_actions_new,
                                     value_prediction_new,
                                     value_prediction_old,
                                     padded_actions,
+                                    rewards_to_actions,
                                     padded_rewards,
                                     reward_mask,
                                     gamma=0.99,
@@ -637,6 +651,14 @@ def combined_loss_given_predictions(log_probab_actions_new,
                                     c1=1.0,
                                     c2=0.01):
   """Computes the combined (clipped loss + value loss) given predictions."""
+  # Sum values over symbols in an action's representation, because it's a simple
+  # way of going from AT to RT+1 and does not decrease the expressive power.
+  value_prediction_old = np.dot(
+      value_prediction_old, rewards_to_actions.transpose()
+  )
+  value_prediction_new = np.dot(
+      value_prediction_new, rewards_to_actions.transpose()
+  )
   (value_loss, value_summaries) = value_loss_given_predictions(
       value_prediction_new,
       padded_rewards,
@@ -649,12 +671,16 @@ def combined_loss_given_predictions(log_probab_actions_new,
       log_probab_actions_old,
       value_prediction_old,
       padded_actions,
+      rewards_to_actions,
       padded_rewards,
       reward_mask,
       gamma=gamma,
       lambda_=lambda_,
       epsilon=epsilon)
-  entropy_bonus = masked_entropy(log_probab_actions_new, reward_mask)
+  # Pad the reward mask to be compatible with rewards_to_actions.
+  padded_reward_mask = np.pad(reward_mask, ((0, 0), (0, 1)))
+  action_mask = np.dot(padded_reward_mask, rewards_to_actions)
+  entropy_bonus = masked_entropy(log_probab_actions_new, action_mask)
   combined_loss_ = ppo_loss + (c1 * value_loss) - (c2 * entropy_bonus)
 
   summaries = {
@@ -674,6 +700,7 @@ def combined_loss(new_params,
                   policy_and_value_net_apply,
                   padded_observations,
                   padded_actions,
+                  rewards_to_actions,
                   padded_rewards,
                   reward_mask,
                   gamma=0.99,
@@ -694,6 +721,7 @@ def combined_loss(new_params,
       value_predictions_new,
       value_predictions_old,
       padded_actions,
+      rewards_to_actions,
       padded_rewards,
       reward_mask,
       gamma=gamma,
@@ -715,6 +743,7 @@ def policy_and_value_opt_step(i,
                               value_predictions_old,
                               padded_observations,
                               padded_actions,
+                              rewards_to_actions,
                               padded_rewards,
                               reward_mask,
                               c1=1.0,
@@ -736,6 +765,7 @@ def policy_and_value_loss(params, state):
         policy_and_value_net_apply,
         padded_observations,
         padded_actions,
+        rewards_to_actions,
         padded_rewards,
         reward_mask,
         c1=c1,
@@ -763,18 +793,16 @@ def approximate_kl(log_prob_new, log_prob_old, mask):
   """Computes the approximate KL divergence between the old and new log-probs.
 
   Args:
-    log_prob_new: (B, T+1, C, A) log probs new
-    log_prob_old: (B, T+1, C, A) log probs old
-    mask: (B, T)
+    log_prob_new: (B, AT, A) log probs new
+    log_prob_old: (B, AT, A) log probs old
+    mask: (B, AT)
 
   Returns:
     Approximate KL.
   """
   diff = log_prob_old - log_prob_new
-  # Cut the last time-step out.
-  diff = diff[:, :-1]
   # Mask out the irrelevant part.
-  diff *= mask[:, :, np.newaxis, np.newaxis]  # make mask (B, T, 1, 1)
+  diff *= mask[:, :, np.newaxis]  # make mask (B, RT, 1)
   # Average on non-masked part.
   return np.sum(diff) / np.sum(mask)
 
@@ -783,17 +811,15 @@ def masked_entropy(log_probs, mask):
   """Computes the entropy for the given log-probs.
 
   Args:
-    log_probs: (B, T+1, C, A) log probs
-    mask: (B, T) mask.
+    log_probs: (B, AT, A) log probs
+    mask: (B, AT) mask.
 
   Returns:
     Entropy.
   """
-  # Cut the last time-step out.
-  lp = log_probs[:, :-1]
   # Mask out the irrelevant part.
-  lp *= mask[:, :, np.newaxis, np.newaxis]  # make mask (B, T, 1, 1)
-  p = np.exp(lp) * mask[:, :, np.newaxis, np.newaxis]  # (B, T, 1, 1)
+  lp = log_probs * mask[:, :, np.newaxis]  # make mask (B, RT, 1)
+  p = np.exp(lp) * mask[:, :, np.newaxis]  # (B, RT, 1)
   # Average on non-masked part and take negative.
   return -(np.sum(lp * p) / np.sum(mask))
 
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
index aa25e51b9..5b6d6a67d 100644
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -96,8 +96,8 @@ def test_policy_and_value_net(self):
     # Output is a list, first is probab of actions and the next is value output.
     self.assertEqual(2, len(pnv_output))
     self.assertEqual(
-        (batch, time_steps, n_controls, n_actions), pnv_output[0].shape)
-    self.assertEqual((batch, time_steps, 1), pnv_output[1].shape)
+        (batch, time_steps * n_controls, n_actions), pnv_output[0].shape)
+    self.assertEqual((batch, time_steps * n_controls), pnv_output[1].shape)
 
   def test_pad_trajectories(self):
     observation_shape = (2, 3, 4)
@@ -235,7 +235,7 @@ def value_net_apply(observations, params, rng=None):
       B, T_p_1, OBS = (observations.shape[0], observations.shape[1],
                        observations.shape[2:])
       del OBS
-      return np.ones((B, T_p_1, 1))
+      return np.ones((B, T_p_1))
       # pylint: enable=invalid-name
 
     value_prediction = value_net_apply(random_observations, [])
@@ -314,19 +314,19 @@ def test_gae_advantages(self):
     self.assertAllEqual(expected_gae_advantages, gae_advantages)
 
   def test_chosen_probabs(self):
-    # Shape (2, 2+1, 1, 3)
+    # Shape (2, 2, 3)
     probab_observations = np.array(
-        [[[[0.1, 0.2, 0.7]], [[0.4, 0.1, 0.5]], [[0.2, 0.4, 0.4]]],
-         [[[0.3, 0.1, 0.6]], [[0.1, 0.1, 0.8]], [[0.2, 0.4, 0.4]]]]
+        [[[0.1, 0.2, 0.7], [0.4, 0.1, 0.5]],
+         [[0.3, 0.1, 0.6], [0.1, 0.1, 0.8]]]
     )
 
     # Shape (2, 2, 1)
-    actions = np.array([[[1], [2]], [[0], [1]]])
+    actions = np.array([[1, 2], [0, 1]])
 
     chosen_probabs = ppo.chosen_probabs(probab_observations, actions)
 
     self.assertAllEqual(
-        np.array([[[0.2], [0.5]], [[0.3], [0.1]]]), chosen_probabs)
+        np.array([[0.2, 0.5], [0.3, 0.1]]), chosen_probabs)
 
   def test_compute_probab_ratios(self):
     p_old = np.array([[
@@ -334,30 +334,26 @@ def test_compute_probab_ratios(self):
         [np.log(0.4), np.log(0.1), np.log(0.4), np.log(0.1)],
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
         [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
     ], [
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
         [np.log(0.1), np.log(0.1), np.log(0.4), np.log(0.4)],
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
         [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-    ]])[:, :, None]
+    ]])
 
     p_new = np.array([[
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
         [np.log(0.4), np.log(0.1), np.log(0.1), np.log(0.3)],
         [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
     ], [
         [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
         [np.log(0.1), np.log(0.1), np.log(0.2), np.log(0.6)],
         [np.log(0.3), np.log(0.1), np.log(0.3), np.log(0.3)],
         [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
-        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
-    ]])[:, :, None]
+    ]])
 
-    actions = np.array([[1, 2, 0, 1], [0, 3, 3, 0]])[:, :, None]
+    actions = np.array([[1, 2, 0, 1], [0, 3, 3, 0]])
 
     mask = np.array([[1, 1, 0, 0], [1, 1, 1, 0]])
 
@@ -367,7 +363,7 @@ def test_compute_probab_ratios(self):
         np.array([
             [0.1 / 0.2, 0.1 / 0.4, 0.0, 0.0],
             [0.1 / 0.3, 0.6 / 0.4, 0.3 / 0.1, 0.0],
-        ])[:, :, None], probab_ratios)
+        ]), probab_ratios)
 
   def test_clipped_probab_ratios(self):
     probab_ratios = np.array([
@@ -385,8 +381,8 @@ def test_clipped_probab_ratios(self):
 
   def test_clipped_objective(self):
     probab_ratios = np.array([
-        [[1.5], [2.0], [0.5], [0.7]],
-        [[2.5], [2.0], [0.1], [1.0]],
+        [1.5, 2.0, 0.5, 0.7],
+        [2.5, 2.0, 0.1, 1.0],
     ])
 
     advantages = np.array([
@@ -399,36 +395,35 @@ def test_clipped_objective(self):
     epsilon = 0.1
 
     clipped_probab_ratios = np.array([
-        [[1.1], [1.1], [0.9], [0.9]],
-        [[1.1], [1.1], [0.9], [1.0]],
+        [1.1, 1.1, 0.9, 0.9],
+        [1.1, 1.1, 0.9, 1.0],
     ])
 
     unused_advantages_x_probab_ratios = np.array([
-        [[0.15], [-0.2], [0.25], [0.49]],
-        [[5.00], [-4.0], [0.20], [2.00]]
+        [0.15, -0.2, 0.25, 0.49],
+        [5.00, -4.0, 0.20, 2.00]
     ])
 
     unused_advantages_x_clipped_probab_ratios = np.array([
-        [[0.11], [-0.11], [0.45], [0.63]],
-        [[2.20], [-2.20], [1.80], [2.00]]
+        [0.11, -0.11, 0.45, 0.63],
+        [2.20, -2.20, .80, 2.00]
     ])
 
     unused_minimums = np.array([
-        [[0.11], [-0.2], [0.25], [0.49]],
-        [[2.20], [-4.0], [0.20], [2.00]]
+        [0.11, -0.2, 0.25, 0.49],
+        [2.20, -4.0, 0.20, 2.00]
     ])
 
     # minimums * mask
     objective = np.array([
-        [[0.11], [-0.2], [0.0], [0.]],
-        [[2.20], [-4.0], [0.2], [0.]]
+        [0.11, -0.2, 0.0, 0.],
+        [2.20, -4.0, 0.2, 0.]
     ])
 
     # Assert that we computed things correctly in this test.
     self.assertAllClose(
-        np.minimum(probab_ratios * advantages[:, :, None],
-                   clipped_probab_ratios * advantages[:, :, None]) *
-        mask[:, :, None],
+        np.minimum(probab_ratios * advantages,
+                   clipped_probab_ratios * advantages) * mask,
         objective)
 
     self.assertAllClose(
@@ -456,7 +451,7 @@ def test_combined_loss(self):
     # Generate a batch of observations.
 
     observations = np.random.uniform(size=(B, T + 1) + OBS)
-    actions = np.random.randint(0, A, size=(B, T, 1))
+    actions = np.random.randint(0, A, size=(B, T + 1))
     rewards = np.random.uniform(0, 1, size=(B, T))
     mask = np.ones_like(rewards)
 
@@ -472,6 +467,7 @@ def test_combined_loss(self):
     c1 = 1.0
     c2 = 0.01
 
+    rewards_to_actions = np.eye(value_predictions_old.shape[1])
     (value_loss_1, _) = ppo.value_loss_given_predictions(
         value_predictions_new, rewards, mask, gamma=gamma,
         value_prediction_old=value_predictions_old, epsilon=epsilon)
@@ -480,6 +476,7 @@ def test_combined_loss(self):
         old_log_probabs,
         value_predictions_old,
         actions,
+        rewards_to_actions,
         rewards,
         mask,
         gamma=gamma,
@@ -493,6 +490,7 @@ def test_combined_loss(self):
                           net,
                           observations,
                           actions,
+                          rewards_to_actions,
                           rewards,
                           mask,
                           gamma=gamma,
@@ -512,7 +510,7 @@ def test_combined_loss(self):
                     1e-6)
 
   def test_masked_entropy(self):
-    # (2, 4+1, 1, 4)
+    # (2, 4+1, 4)
     log_probs = np.array([[
         [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
         [np.log(0.4), np.log(0.1), np.log(0.4), np.log(0.1)],
@@ -525,12 +523,12 @@ def test_masked_entropy(self):
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
         [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
         [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-    ]])[:, :, None, :]
+    ]])
 
     # (2, 4)
     mask = np.array([
-        [1, 1, 0, 0],
-        [1, 1, 1, 0]
+        [1, 1, 0, 0, 0],
+        [1, 1, 1, 0, 0]
     ])
 
     def plp(p):
@@ -538,11 +536,11 @@ def plp(p):
 
     # Removing the last time-step and the masked stuff, gets us this.
     filtered_log_probs = np.array([[
-        [[plp(0.1), plp(0.2), plp(0.6), plp(0.1)]],
-        [[plp(0.4), plp(0.1), plp(0.4), plp(0.1)]],
-        [[plp(0.3), plp(0.1), plp(0.5), plp(0.1)]],
-        [[plp(0.1), plp(0.1), plp(0.4), plp(0.4)]],
-        [[plp(0.3), plp(0.1), plp(0.5), plp(0.1)]],
+        [plp(0.1), plp(0.2), plp(0.6), plp(0.1)],
+        [plp(0.4), plp(0.1), plp(0.4), plp(0.1)],
+        [plp(0.3), plp(0.1), plp(0.5), plp(0.1)],
+        [plp(0.1), plp(0.1), plp(0.4), plp(0.4)],
+        [plp(0.3), plp(0.1), plp(0.5), plp(0.1)],
     ]])
 
     self.assertNear(ppo.masked_entropy(log_probs, mask),
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index b3b09bbda..5e51a1bb6 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -64,7 +64,7 @@ def __init__(self,
                print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
                target_kl=0.01,
                boundary=20,
-               max_timestep=None,
+               max_timestep=100,
                max_timestep_eval=20000,
                random_seed=None,
                gamma=GAMMA,
@@ -202,6 +202,22 @@ def __init__(self,
       logging.info("Saving model on startup to have a model policy file.")
       self.save()
 
+    if isinstance(self.train_env.action_space, gym.spaces.Discrete):
+      n_controls = 1
+    else:
+      n_controls = len(self.train_env.action_space.nvec)
+
+    # Linear map from the reward sequence to the action sequence, used for
+    # scattering advantages over action log-probs and some other things.
+    # It has one more timestep at the and, so it's compatible with the value
+    # predictions.
+    rewards_to_actions = np.eye(max_timestep + 1)[:, None, :]
+    rewards_to_actions = np.broadcast_to(
+        rewards_to_actions, (max_timestep + 1, n_controls, max_timestep + 1)
+    )
+    rewards_to_actions = np.reshape(rewards_to_actions, (max_timestep + 1, -1))
+    self._rewards_to_actions = rewards_to_actions
+
   # Maybe restore the optimization state. If there is nothing to restore, then
   # epoch = 0 and policy_and_value_opt_state is returned as is.
   def update_optimization_state(self,
@@ -309,6 +325,7 @@ def collect_trajectories(self,
       trajs, n_done, timing_info, self._model_state = ppo.collect_trajectories(
           env,
           policy_fn=self._get_predictions,
+          action_index_fn=self._action_index_fn,
           n_trajectories=n_trajectories,
           max_timestep=max_timestep,
           state=self._model_state,
@@ -373,42 +390,47 @@ def train_epoch(self, evaluate=True):
                  min(len(traj[0]) for traj in trajs))
     logging.vlog(2, "Trajectory Lengths: %s", [len(traj[0]) for traj in trajs])
 
-    padding_start_time = time.time()
-    (_, reward_mask, padded_observations, padded_actions, padded_rewards,
-     padded_infos) = ppo.pad_trajectories(
-         trajs, boundary=self._boundary)
-    padding_time = ppo.get_time(padding_start_time)
+    preprocessing_start_time = time.time()
+    (padded_observations, padded_actions, padded_rewards, reward_mask,
+     padded_infos) = self._preprocess_trajectories(trajs)
+    preprocessing_time = ppo.get_time(preprocessing_start_time)
 
-    logging.vlog(1, "Padding trajectories took %0.2f msec.",
-                 ppo.get_time(padding_start_time))
+    logging.vlog(1, "Preprocessing trajectories took %0.2f msec.",
+                 ppo.get_time(preprocessing_start_time))
     logging.vlog(1, "Padded Observations' shape [%s]",
                  str(padded_observations.shape))
     logging.vlog(1, "Padded Actions' shape [%s]", str(padded_actions.shape))
     logging.vlog(1, "Padded Rewards' shape [%s]", str(padded_rewards.shape))
 
-    if padded_actions.ndim == 2:
-      # Add control axis.
-      padded_actions = np.expand_dims(padded_actions, axis=-1)
-
     # Some assertions.
-    B, T, C = padded_actions.shape  # pylint: disable=invalid-name
-    assert (B, T) == padded_rewards.shape
-    assert (B, T) == reward_mask.shape
-    assert (B, T + 1) == padded_observations.shape[:2]
-    assert ((B, T + 1) +
-            self.train_env.observation_space.shape == padded_observations.shape)
+    B, RT = padded_rewards.shape  # pylint: disable=invalid-name
+    B, AT = padded_actions.shape  # pylint: disable=invalid-name
+    assert (B, RT) == reward_mask.shape
+    assert B == padded_observations.shape[0]
+    assert (
+        self.train_env.observation_space.shape == padded_observations.shape[2:]
+    )
 
     log_prob_recompute_start_time = time.time()
-    assert ("log_prob_actions" in padded_infos and
-            "value_predictions" in padded_infos)
+    # TODO(pkozakowski): The following commented out code collects the network
+    # predictions made while stepping the environment and uses them in PPO
+    # training, so that we can use non-deterministic networks (e.g. with
+    # dropout). This does not work well with serialization, so instead we
+    # recompute all network predictions. Let's figure out a solution that will
+    # work with both serialized sequences and non-deterministic networks.
+
+    # assert ("log_prob_actions" in padded_infos and
+    #         "value_predictions" in padded_infos)
     # These are the actual log-probabs and value predictions seen while picking
     # the actions.
-    actual_log_probabs_traj = padded_infos["log_prob_actions"]
-    actual_value_predictions_traj = padded_infos["value_predictions"]
+    # actual_log_probabs_traj = padded_infos["log_prob_actions"]
+    # actual_value_predictions_traj = padded_infos["value_predictions"]
 
-    assert (B, T, C) == actual_log_probabs_traj.shape[:3]
-    A = actual_log_probabs_traj.shape[3]  # pylint: disable=invalid-name
-    assert (B, T, 1) == actual_value_predictions_traj.shape
+    # assert (B, T, C) == actual_log_probabs_traj.shape[:3]
+    # A = actual_log_probabs_traj.shape[3]  # pylint: disable=invalid-name
+    # assert (B, T, 1) == actual_value_predictions_traj.shape
+
+    del padded_infos
 
     # TODO(afrozm): log-probabs doesn't need to be (B, T+1, C, A) it can do with
     # (B, T, C, A), so make that change throughout.
@@ -421,16 +443,18 @@ def train_epoch(self, evaluate=True):
     log_probabs_traj, value_predictions_traj, self._model_state, _ = (
         self._get_predictions(padded_observations, self._model_state, rng=key))
 
-    assert (B, T + 1, C, A) == log_probabs_traj.shape
-    assert (B, T + 1, 1) == value_predictions_traj.shape
+    assert (B, AT) == log_probabs_traj.shape[:2]
+    assert (B, AT) == value_predictions_traj.shape
+
+    # TODO(pkozakowski): Commented out for the same reason as before.
 
     # Concatenate the last time-step's log-probabs and value predictions to the
     # actual log-probabs and value predictions and use those going forward.
-    log_probabs_traj = np.concatenate(
-        (actual_log_probabs_traj, log_probabs_traj[:, -1:, :]), axis=1)
-    value_predictions_traj = np.concatenate(
-        (actual_value_predictions_traj, value_predictions_traj[:, -1:, :]),
-        axis=1)
+    # log_probabs_traj = np.concatenate(
+    #     (actual_log_probabs_traj, log_probabs_traj[:, -1:, :]), axis=1)
+    # value_predictions_traj = np.concatenate(
+    #     (actual_value_predictions_traj, value_predictions_traj[:, -1:, :]),
+    #     axis=1)
 
     log_prob_recompute_time = ppo.get_time(log_prob_recompute_start_time)
 
@@ -446,6 +470,7 @@ def train_epoch(self, evaluate=True):
             self._policy_and_value_net_apply,
             padded_observations,
             padded_actions,
+            self._rewards_to_actions,
             padded_rewards,
             reward_mask,
             gamma=self._gamma,
@@ -490,6 +515,7 @@ def train_epoch(self, evaluate=True):
               value_predictions_traj,
               padded_observations,
               padded_actions,
+              self._rewards_to_actions,
               padded_rewards,
               reward_mask,
               c1=self._c1,
@@ -509,8 +535,11 @@ def train_epoch(self, evaluate=True):
               self._model_state,
               rng=k2))
 
+      action_mask = np.dot(
+          np.pad(reward_mask, ((0, 0), (0, 1))), self._rewards_to_actions
+      )
       approx_kl = ppo.approximate_kl(log_probab_actions_new, log_probabs_traj,
-                                     reward_mask)
+                                     action_mask)
 
       early_stopping = approx_kl > 1.5 * self._target_kl
       if early_stopping:
@@ -532,6 +561,7 @@ def train_epoch(self, evaluate=True):
                 self._policy_and_value_net_apply,
                 padded_observations,
                 padded_actions,
+                self._rewards_to_actions,
                 padded_rewards,
                 reward_mask,
                 gamma=self._gamma,
@@ -594,7 +624,7 @@ def train_epoch(self, evaluate=True):
         "epoch": epoch_time,
         "policy_eval": policy_eval_time,
         "trajectory_collection": trajectory_collection_time,
-        "padding": padding_time,
+        "preprocessing": preprocessing_time,
         "log_prob_recompute": log_prob_recompute_time,
         "loss_compute": loss_compute_time,
         "optimization": optimization_time,
@@ -675,6 +705,16 @@ def flush_summaries(self):
   def _policy_and_value_net_params(self):
     return self._policy_and_value_get_params(self._policy_and_value_opt_state)
 
+  # Prepares the trajectories for policy training.
+  def _preprocess_trajectories(self, trajectories):
+    (_, reward_mask, observations, actions, rewards,
+     infos) = ppo.pad_trajectories(trajectories, boundary=self._max_timestep)
+    # Add one timestep at the end, so it's compatible with
+    # self._rewards_to_actions.
+    actions = np.pad(actions, ((0, 0), (0, 1)) + ((0, 0),) * (actions.ndim - 2))
+    actions = np.reshape(actions, (actions.shape[0], -1))
+    return (observations, actions, rewards, reward_mask, infos)
+
   # A function to get the policy and value predictions.
   def _get_predictions(self, observations, state, rng=None):
     """Returns log-probs, value predictions and key back."""
@@ -684,3 +724,16 @@ def _get_predictions(self, observations, state, rng=None):
         observations, self._policy_and_value_net_params, state, rng=key1)
 
     return log_probs, value_preds, state, key
+
+  def _action_index_fn(self, index):
+    # Project the one-hot position in the reward sequence onto the action
+    # sequence to figure out which actions correspond to that position.
+    one_hot_index = np.eye(self._rewards_to_actions.shape[0])[index]
+    action_mask = np.dot(one_hot_index, self._rewards_to_actions)
+    # Compute the number of symbols in an action. It's just the number of 1s in
+    # the mask.
+    action_length = int(np.sum(action_mask[0]))
+    # Argmax stops on the first occurrence, so we use it to find the first 1 in
+    # the mask.
+    action_start_index = np.argmax(action_mask, axis=1)
+    return action_start_index[:, None] + np.arange(action_length)[None, :]
diff --git a/tensor2tensor/trax/rl/ppo_trainer_test.py b/tensor2tensor/trax/rl/ppo_trainer_test.py
index da11555fd..39b99894c 100644
--- a/tensor2tensor/trax/rl/ppo_trainer_test.py
+++ b/tensor2tensor/trax/rl/ppo_trainer_test.py
@@ -104,7 +104,7 @@ def test_training_loop_cartpole_transformer(self):
               d_ff=1,
               n_layers=1,
               n_heads=1,
-              max_len=64,
+              max_len=128,
               mode="train",
           ),
       )

From 0df000a2aedfd20e930b818e979e9222a555417e Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 18 Sep 2019 12:35:17 -0700
Subject: [PATCH 2457/2720] Set the async mode to False while using
 SimulatedEnv in SimPLe, restore to original value when using the real env.

Otherwise it expects trajectories from SimulatedEnv but reads them from the
real envs (since they are still churning out trajectories).

PiperOrigin-RevId: 269866203
---
 tensor2tensor/trax/rl/base_trainer.py   | 10 ++++++++++
 tensor2tensor/trax/rl/simple_trainer.py | 15 ++++++++++++---
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/trax/rl/base_trainer.py b/tensor2tensor/trax/rl/base_trainer.py
index 622f106a7..39e79f8fd 100644
--- a/tensor2tensor/trax/rl/base_trainer.py
+++ b/tensor2tensor/trax/rl/base_trainer.py
@@ -58,6 +58,16 @@ def __init__(
     self._trajectory_buffer = []
     self._async_mode = async_mode
 
+  @property
+  def async_mode(self):
+    return self._async_mode
+
+  @async_mode.setter
+  def async_mode(self, async_mode):
+    logging.vlog(1, "Changing async mode from %s to: %s",
+                 self._async_mode, async_mode)
+    self._async_mode = async_mode
+
   @property
   def epoch(self):
     raise NotImplementedError
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index a68f32736..c615fc8bb 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -114,15 +114,19 @@ def epoch(self):
     return self._simple_epoch
 
   def train_epoch(self, evaluate=True):
-    # Collect trajectories by running the policy in the real environment.
     if self._simple_epoch > 0 or not self._has_initial_data:
+      logging.info(
+          "Collect trajectories by running the policy in the real environment.")
       self.collect_trajectories(evaluate=evaluate)
-    # Train the model of the environment on the collected trajectories.
     if self._simple_epoch > 0 or not self._initial_model:
+      logging.info(
+          "Train the model of the environment on the collected trajectories.")
       skipped = self.train_model()
       if evaluate and not skipped:
+        logging.info("Evaluate the trained model.")
         self.evaluate_model()
-    # Train the policy inside the simulated environment generated by the model.
+    logging.info("Train the policy inside the simulated environment generated "
+                 "by the model.")
     self.train_policy()
 
     self._simple_epoch += 1
@@ -196,11 +200,16 @@ def train_policy(self):
         batch_size=self._simulated_batch_size,
         history_stream=itertools.repeat(None),
     )
+    # We never want async mode in the simulated env.
+    original_async_mode = self._policy_trainer.async_mode
+    self._policy_trainer.async_mode = False
     self._policy_trainer.train_env = self._sim_env
     # Don't dump trajectories from the simulated environment.
     self._policy_trainer.trajectory_dump_dir = None
     self._policy_epoch += self._n_simulated_epochs
     self._policy_trainer.training_loop(self._policy_epoch, evaluate=False)
+    # Revert back to the original async mode in the policy trainer.
+    self._policy_trainer.async_mode = original_async_mode
 
     logging.vlog(
         1, "Training policy took %0.2f sec.", time.time() - start_time)

From a2d1848aa0c008332f882411f4e458f2bd41bffa Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 18 Sep 2019 15:04:39 -0700
Subject: [PATCH 2458/2720] Update observation_space in OnlineTuneEnv.

PiperOrigin-RevId: 269899347
---
 tensor2tensor/trax/rl/envs/online_tune_env.py      | 10 +++++-----
 tensor2tensor/trax/rl/envs/online_tune_env_test.py |  2 +-
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index bec69faea..74104d7fe 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -71,7 +71,7 @@ def __init__(self,
                    # (name, start, (low, high), flip)
                    ("learning_rate", 1e-3, (1e-9, 10.0), False),
                ),
-               observation_range=(0.0, 5.0),
+               metric_range=(0.0, 5.0),
                # Don't save checkpoints by default, as they tend to use a lot of
                # space.
                should_save_checkpoints=False):
@@ -99,6 +99,7 @@ def __init__(self,
     self._eval_steps = eval_steps
     self._env_steps = env_steps
     self._control_configs = control_configs
+    self._metric_range = metric_range
 
     self._output_dir = output_dir
     gfile.makedirs(self._output_dir)
@@ -112,10 +113,9 @@ def __init__(self,
         len(self._observation_metrics) +
         int(self._include_controls_in_observation) * len(self._control_configs)
     )
-    self._observation_range = observation_range
-    (low, high) = self._observation_range
     self.observation_space = gym.spaces.Box(
-        low=low, high=high, shape=(observation_dim,))
+        # Observations are clipped to metric_range and rescaled to [-1, 1].
+        low=-1, high=1, shape=(observation_dim,))
 
   @property
   def _next_trajectory_dir(self):
@@ -158,7 +158,7 @@ def _current_observation(self):
     observations = online_tune.history_to_observations(
         self._trainer.state.history,
         self._observation_metrics,
-        self._observation_range,
+        self._metric_range,
         self._control_configs if self._include_controls_in_observation
         else None,
     )
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env_test.py b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
index ad6467676..b2b95c2d9 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env_test.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
@@ -91,7 +91,7 @@ def _create_env(
             ("learning_rate", 1e-3, (1e-9, 10.0), False),
             ("weight_decay_rate", 1e-5, (1e-9, 0.1), False),
         ),
-        observation_range=(-1, 1),
+        metric_range=(-1, 1),
         include_controls_in_observation=False,
         output_dir=output_dir,
         action_multipliers=action_multipliers,

From d600c8bb196193596fdb38c2b6e5393c4e240564 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 18 Sep 2019 15:59:31 -0700
Subject: [PATCH 2459/2720] Enable the policy to run on a serialized
 representation.

PiperOrigin-RevId: 269910480
---
 tensor2tensor/envs/env_problem_utils.py      |  22 +--
 tensor2tensor/envs/env_problem_utils_test.py |  11 +-
 tensor2tensor/trax/learning_rate.py          |   5 +
 tensor2tensor/trax/learning_rate_test.py     |   1 +
 tensor2tensor/trax/models/transformer.py     |  23 ++-
 tensor2tensor/trax/rl/ppo.py                 |  23 ++-
 tensor2tensor/trax/rl/ppo_test.py            |   2 +
 tensor2tensor/trax/rl/ppo_trainer.py         | 150 +++++++++++++++----
 tensor2tensor/trax/rl/ppo_trainer_test.py    |  25 +++-
 tensor2tensor/trax/rl/serialization_utils.py |   2 +-
 tensor2tensor/trax/rl/space_serializer.py    |   4 +-
 11 files changed, 191 insertions(+), 77 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 99d0e62eb..2520cc74e 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -72,7 +72,6 @@ def get_completed_trajectories_from_env(env,
 
 def play_env_problem_with_policy(env,
                                  policy_fun,
-                                 action_index_fn,
                                  num_trajectories=1,
                                  max_timestep=None,
                                  reset=True,
@@ -90,8 +89,6 @@ def play_env_problem_with_policy(env,
     env: environment object, should be a subclass of env_problem.EnvProblem.
     policy_fun: callable, taking in observations((B, RT) + OBS) and returning
       back log-probabilities (B, AT, A).
-    action_index_fn: function converting timestep indices into indices in the
-      log-probability array.
     num_trajectories: int, number of trajectories to collect.
     max_timestep: int or None, if not None or a negative number, we cut any
       trajectory that exceeds this time put it in the completed bin, and *dont*
@@ -156,26 +153,11 @@ def gumbel_sample(log_probs):
     assert (B,) == lengths.shape
 
     t1 = time.time()
-    log_prob_actions, value_predictions, state, rng = policy_fun(
-        padded_observations, state=state, rng=rng)
+    log_probs, value_preds, state, rng = policy_fun(
+        padded_observations, lengths, state=state, rng=rng)
     policy_application_total_time += (time.time() - t1)
 
-    assert B == log_prob_actions.shape[0]
-    (_, A) = log_prob_actions.shape[1:]  # pylint: disable=invalid-name
-
-    # We need the log_probs of those actions that correspond to the last actual
-    # time-step.
-    index = lengths - 1  # Since we want to index using lengths.
-    pred_index = action_index_fn(index)
-    log_probs = log_prob_actions[
-        np.arange(B)[:, None, None],
-        pred_index[:, :, None],
-        np.arange(A),
-    ]
-    value_preds = value_predictions[np.arange(B)[:, None], pred_index]
-
     assert B == log_probs.shape[0]
-    assert A == log_probs.shape[2]
 
     actions = gumbel_sample(log_probs)
     if isinstance(env.action_space, gym.spaces.Discrete):
diff --git a/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index cb15941c8..572294bca 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -53,8 +53,9 @@ def test_play_env_problem_with_policy(self):
     # Let's make sure that at-most 4 observations come to the policy function.
     len_history_for_policy = 4
 
-    def policy_fun(observations, state=None, rng=None):
-      b, t = observations.shape[:2]
+    def policy_fun(observations, lengths, state=None, rng=None):
+      del lengths
+      b = observations.shape[0]
       # Assert that observations from time-step len_history_for_policy onwards
       # are zeros.
       self.assertTrue(
@@ -62,20 +63,16 @@ def policy_fun(observations, state=None, rng=None):
       self.assertFalse(
           np.all(observations[:, :len_history_for_policy, ...] == 0))
       a = env.action_space.n
-      p = np.random.uniform(size=(b, t, a))
+      p = np.random.uniform(size=(b, 1, a))
       p = np.exp(p)
       p = p / np.sum(p, axis=-1, keepdims=True)
       return np.log(p), np.mean(p, axis=-1), state, rng
 
-    def action_index_fn(index):
-      return index[:, None]
-
     max_timestep = 15
     num_trajectories = 2
     trajectories, _, _, _ = env_problem_utils.play_env_problem_with_policy(
         env,
         policy_fun,
-        action_index_fn=action_index_fn,
         num_trajectories=num_trajectories,
         max_timestep=max_timestep,
         len_history_for_policy=len_history_for_policy)
diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index 279a2bece..0ff4ed801 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -158,6 +158,7 @@ def PolicySchedule(
     action_multipliers=(1.0 / 1.5, 1.0 / 1.25, 1.0, 1.25, 1.5),
     policy_and_value_model=trax_models.FrameStackMLP,
     policy_and_value_two_towers=False,
+    policy_and_value_vocab_size=None,
     policy_dir=gin.REQUIRED,
 ):
   """Learning rate schedule controlled by a learned policy.
@@ -175,6 +176,9 @@ def PolicySchedule(
     policy_and_value_model: Trax model to use as the policy.
     policy_and_value_two_towers: bool, whether the action distribution and value
       prediction is computed by separate model towers.
+    policy_and_value_vocab_size: Vocabulary size of a policy and value network
+      operating on serialized representation. If None, use raw continuous
+      representation.
     policy_dir: directory with the policy checkpoint.
 
   Returns:
@@ -203,6 +207,7 @@ def PolicySchedule(
   net = ppo.policy_and_value_net(
       n_controls=1,
       n_actions=len(action_multipliers),
+      vocab_size=policy_and_value_vocab_size,
       bottom_layers_fn=policy_and_value_model,
       two_towers=policy_and_value_two_towers,
   )
diff --git a/tensor2tensor/trax/learning_rate_test.py b/tensor2tensor/trax/learning_rate_test.py
index fef08dd9d..d2fabea05 100644
--- a/tensor2tensor/trax/learning_rate_test.py
+++ b/tensor2tensor/trax/learning_rate_test.py
@@ -44,6 +44,7 @@ def _make_schedule(
     net = ppo.policy_and_value_net(
         n_actions=len(action_multipliers),
         n_controls=1,
+        vocab_size=None,
         bottom_layers_fn=policy_and_value_model,
         two_towers=False,
     )
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index f180c2a57..aaf729fa0 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -18,6 +18,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
+
 from tensor2tensor.trax import layers as tl
 
 
@@ -148,7 +150,8 @@ def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
   ]
 
 
-def TransformerDecoder(d_model=512,
+def TransformerDecoder(vocab_size=None,
+                       d_model=512,
                        d_ff=2048,
                        n_layers=6,
                        n_heads=8,
@@ -161,11 +164,13 @@ def TransformerDecoder(d_model=512,
                        mode='train'):
   """Returns a Transformer decoder model.
 
-  The input to the model is a continuous tensor. Does not shift the input to the
-  right, i.e. the output for timestep t is based on inputs up to timestep t
-  inclusively.
+  The input to the model is either continuous or discrete - controlled by
+  vocab_size. Does not shift the input to the right, i.e. the output for
+  timestep t is based on inputs up to timestep t inclusively.
 
   Args:
+    vocab_size: int or None: vocab size if running on discrete input, None
+        otherwise.
     d_model: int:  depth of embedding
     d_ff: int: depth of feed-forward layer
     n_layers: int: number of encoder/decoder layers
@@ -181,11 +186,15 @@ def TransformerDecoder(d_model=512,
     mode: str: 'train' or 'eval'
 
   Returns:
-    A Transformer decoder as a layer that maps from a continuous tensor to
-    a continuous tensor.
+    A Transformer decoder as a layer that maps from a continuous or discrete
+    tensor to a continuous tensor.
   """
+  if vocab_size is None:
+    input_layer = tl.Dense
+  else:
+    input_layer = functools.partial(tl.Embedding, vocab_size=vocab_size)
   return tl.Model(                  # vecs
-      tl.Dense(d_model),            # vecs
+      input_layer(d_model),         # vecs
       tl.Dropout(rate=dropout, mode=mode),
       tl.PositionalEncoding(max_len=max_len),
       [DecoderBlock(  # pylint: disable=g-complex-comprehension
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index db5b2e1ef..2518ece0e 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -71,7 +71,9 @@
 from tensorflow.io import gfile
 
 
-def policy_and_value_net(n_actions, n_controls, bottom_layers_fn, two_towers):
+def policy_and_value_net(
+    n_actions, n_controls, vocab_size, bottom_layers_fn, two_towers
+):
   """A policy and value net function."""
 
   # Layers.
@@ -85,13 +87,22 @@ def FlattenControlsIntoTime(x, **unused_kwargs):  # pylint: disable=invalid-name
     """Splits logits for actions in different controls and flattens controls."""
     return np.reshape(x, (x.shape[0], -1, n_actions))
 
-  n_logits = n_controls * n_actions
+  if vocab_size is None:
+    # In continuous policies every element of the output sequence corresponds to
+    # an observation.
+    n_logits = n_controls * n_actions
+    kwargs = {}
+  else:
+    # In discrete policies every element of the output sequence corresponds to
+    # a symbol in the discrete representation, and each control takes 1 symbol.
+    n_logits = n_actions
+    kwargs = {"vocab_size": vocab_size}
 
   if two_towers:
     layers = [
         tl.Dup(),
         tl.Parallel(
-            [bottom_layers_fn(),
+            [bottom_layers_fn(**kwargs),
              tl.Dense(n_logits),
              FlattenControlsIntoTime(),  # pylint: disable=no-value-for-parameter
              tl.LogSoftmax()],
@@ -100,7 +111,7 @@ def FlattenControlsIntoTime(x, **unused_kwargs):  # pylint: disable=invalid-name
     ]
   else:
     layers = [
-        bottom_layers_fn(),
+        bottom_layers_fn(**kwargs),
         tl.Dup(),
         tl.Parallel(
             [tl.Dense(n_logits),
@@ -149,7 +160,6 @@ def get_params(opt_state):
 # Any other option?
 def collect_trajectories(env,
                          policy_fn,
-                         action_index_fn,
                          n_trajectories=1,
                          max_timestep=None,
                          reset=True,
@@ -165,8 +175,6 @@ def collect_trajectories(env,
   Args:
     env: A gym env interface, for now this is not-batched.
     policy_fn: observations(B,RT+1) -> log-probabs(B, AT, A) callable.
-    action_index_fn: function converting timestep indices into indices in the
-      log-probability array.
     n_trajectories: int, number of trajectories.
     max_timestep: int or None, the index of the maximum time-step at which we
       return the trajectory, None for ending a trajectory only when env returns
@@ -200,7 +208,6 @@ def collect_trajectories(env,
   trajs, n_done, timing_info, state = env_problem_utils.play_env_problem_with_policy(
       env,
       policy_fn,
-      action_index_fn,
       num_trajectories=n_trajectories,
       max_timestep=max_timestep,
       reset=reset,
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
index 5b6d6a67d..d71772647 100644
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -81,6 +81,7 @@ def test_policy_and_value_net(self):
     pnv_model = ppo.policy_and_value_net(
         n_controls=n_controls,
         n_actions=n_actions,
+        vocab_size=None,
         bottom_layers_fn=lambda: [layers.Flatten(n_axes_to_keep=2)],
         two_towers=True,
     )
@@ -439,6 +440,7 @@ def test_combined_loss(self):
     net = ppo.policy_and_value_net(
         n_controls=1,
         n_actions=A,
+        vocab_size=None,
         bottom_layers_fn=lambda: [layers.Flatten(n_axes_to_keep=2)],
         two_towers=True,
     )
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index 5e51a1bb6..c809d1b5e 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -38,6 +38,8 @@
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.rl import base_trainer
 from tensor2tensor.trax.rl import ppo
+from tensor2tensor.trax.rl import serialization_utils
+from tensor2tensor.trax.rl import space_serializer
 
 DEBUG_LOGGING = False
 GAMMA = 0.99
@@ -60,6 +62,7 @@ def __init__(self,
                policy_and_value_optimizer=functools.partial(
                    trax_opt.Adam, learning_rate=1e-3),
                policy_and_value_two_towers=False,
+               policy_and_value_vocab_size=None,
                n_optimizer_steps=N_OPTIMIZER_STEPS,
                print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
                target_kl=0.01,
@@ -90,6 +93,9 @@ def __init__(self,
       policy_and_value_optimizer: Function defining the optimizer.
       policy_and_value_two_towers: Whether to use two separate models as the
         policy and value networks. If False, share their parameters.
+      policy_and_value_vocab_size: Vocabulary size of a policy and value network
+        operating on serialized representation. If None, use raw continuous
+        representation.
       n_optimizer_steps: Number of optimizer steps.
       print_every_optimizer_steps: How often to log during the policy
         optimization process.
@@ -154,26 +160,31 @@ def __init__(self,
       assert onp.min(action_space.nvec) == onp.max(action_space.nvec), (
           "Every control must have the same number of actions.")
       n_actions = action_space.nvec[0]
-
-    # Batch Observations Shape = [1, 1] + OBS, because we will eventually call
-    # policy and value networks on shape [B, T] +_OBS
-    batch_observations_shape = (1, 1) + self.train_env.observation_space.shape
-    observations_dtype = self.train_env.observation_space.dtype
+    self._n_actions = n_actions
+    self._n_controls = n_controls
 
     self._rng = trax.get_random_number_generator_and_set_seed(random_seed)
     self._rng, key1 = jax_random.split(self._rng, num=2)
 
+    vocab_size = policy_and_value_vocab_size
+    self._serialized_sequence_policy = vocab_size is not None
+    if self._serialized_sequence_policy:
+      self._serialization_kwargs = self._init_serialization(vocab_size)
+    else:
+      self._serialization_kwargs = {}
+
     # Initialize the policy and value network.
     policy_and_value_net = ppo.policy_and_value_net(
         n_actions=n_actions,
         n_controls=n_controls,
+        vocab_size=vocab_size,
         bottom_layers_fn=policy_and_value_model,
         two_towers=policy_and_value_two_towers,
     )
     self._policy_and_value_net_apply = jit(policy_and_value_net)
+    (batch_obs_shape, obs_dtype) = self._batch_obs_shape_and_dtype
     policy_and_value_net_params, self._model_state = (
-        policy_and_value_net.initialize(batch_observations_shape,
-                                        observations_dtype, key1))
+        policy_and_value_net.initialize(batch_obs_shape, obs_dtype, key1))
 
     # Initialize the optimizer.
     (policy_and_value_opt_state, self._policy_and_value_opt_update,
@@ -202,21 +213,53 @@ def __init__(self,
       logging.info("Saving model on startup to have a model policy file.")
       self.save()
 
-    if isinstance(self.train_env.action_space, gym.spaces.Discrete):
-      n_controls = 1
-    else:
-      n_controls = len(self.train_env.action_space.nvec)
+    self._rewards_to_actions = self._init_rewards_to_actions()
+
+  def _init_serialization(self, vocab_size):
+    obs_serializer = space_serializer.create(
+        self.train_env.observation_space, vocab_size=vocab_size
+    )
+    act_serializer = space_serializer.create(
+        self.train_env.action_space, vocab_size=vocab_size
+    )
+    repr_length = (
+        obs_serializer.representation_length +
+        act_serializer.representation_length
+    ) * (self._max_timestep + 1)
+    return {
+        "observation_serializer": obs_serializer,
+        "action_serializer": act_serializer,
+        "representation_length": repr_length,
+    }
 
+  def _init_rewards_to_actions(self):
     # Linear map from the reward sequence to the action sequence, used for
     # scattering advantages over action log-probs and some other things.
-    # It has one more timestep at the and, so it's compatible with the value
+    # It has one more timestep at the end, so it's compatible with the value
     # predictions.
-    rewards_to_actions = np.eye(max_timestep + 1)[:, None, :]
-    rewards_to_actions = np.broadcast_to(
-        rewards_to_actions, (max_timestep + 1, n_controls, max_timestep + 1)
-    )
-    rewards_to_actions = np.reshape(rewards_to_actions, (max_timestep + 1, -1))
-    self._rewards_to_actions = rewards_to_actions
+    if not self._serialized_sequence_policy:
+      rewards_to_actions = np.eye(self._max_timestep + 1)[:, None, :]
+      rewards_to_actions = np.broadcast_to(
+          rewards_to_actions,
+          (self._max_timestep + 1, self._n_controls, self._max_timestep + 1),
+      )
+      return np.reshape(rewards_to_actions, (self._max_timestep + 1, -1))
+    else:
+      return serialization_utils.rewards_to_actions_map(
+          n_timesteps=(self._max_timestep + 1), **self._serialization_kwargs
+      )
+
+  @property
+  def _batch_obs_shape_and_dtype(self):
+    if not self._serialized_sequence_policy:
+      # Batch Observations Shape = [1, 1] + OBS, because we will eventually call
+      # policy and value networks on shape [B, T] +_OBS
+      shape = (1, 1) + self.train_env.observation_space.shape
+      dtype = self.train_env.observation_space.dtype
+    else:
+      shape = (1, 1)
+      dtype = np.int32
+    return (shape, dtype)
 
   # Maybe restore the optimization state. If there is nothing to restore, then
   # epoch = 0 and policy_and_value_opt_state is returned as is.
@@ -324,8 +367,7 @@ def collect_trajectories(self,
     else:
       trajs, n_done, timing_info, self._model_state = ppo.collect_trajectories(
           env,
-          policy_fn=self._get_predictions,
-          action_index_fn=self._action_index_fn,
+          policy_fn=self._policy_fun,
           n_trajectories=n_trajectories,
           max_timestep=max_timestep,
           state=self._model_state,
@@ -407,9 +449,6 @@ def train_epoch(self, evaluate=True):
     B, AT = padded_actions.shape  # pylint: disable=invalid-name
     assert (B, RT) == reward_mask.shape
     assert B == padded_observations.shape[0]
-    assert (
-        self.train_env.observation_space.shape == padded_observations.shape[2:]
-    )
 
     log_prob_recompute_start_time = time.time()
     # TODO(pkozakowski): The following commented out code collects the network
@@ -707,14 +746,37 @@ def _policy_and_value_net_params(self):
 
   # Prepares the trajectories for policy training.
   def _preprocess_trajectories(self, trajectories):
-    (_, reward_mask, observations, actions, rewards,
-     infos) = ppo.pad_trajectories(trajectories, boundary=self._max_timestep)
-    # Add one timestep at the end, so it's compatible with
-    # self._rewards_to_actions.
-    actions = np.pad(actions, ((0, 0), (0, 1)) + ((0, 0),) * (actions.ndim - 2))
-    actions = np.reshape(actions, (actions.shape[0], -1))
+    (_, reward_mask, observations, actions, rewards, infos) = (
+        ppo.pad_trajectories(trajectories, boundary=self._max_timestep)
+    )
+    assert self.train_env.observation_space.shape == observations.shape[2:]
+    if not self._serialized_sequence_policy:
+      # Add one timestep at the end, so it's compatible with
+      # self._rewards_to_actions.
+      pad_width = ((0, 0), (0, 1)) + ((0, 0),) * (actions.ndim - 2)
+      actions = np.pad(actions, pad_width)
+      actions = np.reshape(actions, (actions.shape[0], -1))
+    else:
+      (observations, actions) = self._serialize_trajectories(
+          observations, actions, reward_mask
+      )
     return (observations, actions, rewards, reward_mask, infos)
 
+  def _serialize_trajectories(self, observations, actions, reward_mask):
+    (reprs, _) = serialization_utils.serialize_observations_and_actions(
+        observations=observations,
+        actions=actions,
+        mask=reward_mask,
+        **self._serialization_kwargs
+    )
+    # Mask out actions in the representation - otherwise we sample an action
+    # based on itself.
+    observations = reprs * serialization_utils.observation_mask(
+        **self._serialization_kwargs
+    )
+    actions = reprs
+    return (observations, actions)
+
   # A function to get the policy and value predictions.
   def _get_predictions(self, observations, state, rng=None):
     """Returns log-probs, value predictions and key back."""
@@ -725,10 +787,36 @@ def _get_predictions(self, observations, state, rng=None):
 
     return log_probs, value_preds, state, key
 
-  def _action_index_fn(self, index):
+  def _policy_fun(self, observations, lengths, state, rng):
+    (batch_size, n_timesteps) = observations.shape[:2]
+    if self._serialized_sequence_policy:
+      actions = np.zeros(
+          (batch_size, n_timesteps - 1) + self.train_env.action_space.shape,
+          dtype=self.train_env.action_space.dtype,
+      )
+      reward_mask = np.ones((batch_size, n_timesteps - 1), dtype=np.int32)
+      (observations, _) = self._serialize_trajectories(
+          observations, actions, reward_mask
+      )
+    (log_probs, value_preds, state, rng) = self._get_predictions(
+        observations, state=state, rng=rng
+    )
+    # We need the log_probs of those actions that correspond to the last actual
+    # time-step.
+    index = lengths - 1  # Since we want to index using lengths.
+    pred_index = self._calc_action_index(index)
+    log_probs = log_probs[
+        np.arange(batch_size)[:, None, None],
+        pred_index[:, :, None],
+        np.arange(self._n_actions),
+    ]
+    value_preds = value_preds[np.arange(batch_size)[:, None], pred_index]
+    return (log_probs, value_preds, state, rng)
+
+  def _calc_action_index(self, reward_index):
     # Project the one-hot position in the reward sequence onto the action
     # sequence to figure out which actions correspond to that position.
-    one_hot_index = np.eye(self._rewards_to_actions.shape[0])[index]
+    one_hot_index = np.eye(self._rewards_to_actions.shape[0])[reward_index]
     action_mask = np.dot(one_hot_index, self._rewards_to_actions)
     # Compute the number of symbols in an action. It's just the number of 1s in
     # the mask.
diff --git a/tensor2tensor/trax/rl/ppo_trainer_test.py b/tensor2tensor/trax/rl/ppo_trainer_test.py
index 39b99894c..295e79c5c 100644
--- a/tensor2tensor/trax/rl/ppo_trainer_test.py
+++ b/tensor2tensor/trax/rl/ppo_trainer_test.py
@@ -69,7 +69,9 @@ def tmp_dir(self):
     yield tmp
     gfile.rmtree(tmp)
 
-  def _make_trainer(self, train_env, eval_env, output_dir, model=None):
+  def _make_trainer(
+      self, train_env, eval_env, output_dir, model=None, **kwargs
+  ):
     if model is None:
       model = lambda: [layers.Dense(1)]
     return ppo_trainer.PPO(
@@ -81,6 +83,7 @@ def _make_trainer(self, train_env, eval_env, output_dir, model=None):
         random_seed=0,
         boundary=2,
         save_every_n=1,
+        **kwargs
     )
 
   def test_training_loop_cartpole(self):
@@ -254,6 +257,26 @@ def test_training_loop_multi_control(self):
       )
       trainer.training_loop(n_epochs=2)
 
+  def test_training_loop_cartpole_serialized(self):
+    gin.bind_parameter("BoxSpaceSerializer.precision", 1)
+    with self.tmp_dir() as output_dir:
+      trainer = self._make_trainer(
+          train_env=self.get_wrapped_env("CartPole-v0", 2),
+          eval_env=self.get_wrapped_env("CartPole-v0", 2),
+          output_dir=output_dir,
+          model=functools.partial(
+              models.TransformerDecoder,
+              d_model=1,
+              d_ff=1,
+              n_layers=1,
+              n_heads=1,
+              max_len=1024,
+              mode="train",
+          ),
+          policy_and_value_vocab_size=4,
+      )
+      trainer.training_loop(n_epochs=2)
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensor2tensor/trax/rl/serialization_utils.py b/tensor2tensor/trax/rl/serialization_utils.py
index e175f4460..7bd5a2df2 100644
--- a/tensor2tensor/trax/rl/serialization_utils.py
+++ b/tensor2tensor/trax/rl/serialization_utils.py
@@ -91,7 +91,7 @@ def observation_mask(
     Binary mask indicating which symbols in the representation correspond to
     observations.
   """
-  mask = np.zeros(representation_length)
+  mask = np.zeros(representation_length, dtype=np.int32)
   obs_repr_length = observation_serializer.representation_length
   step_repr_length = obs_repr_length + action_serializer.representation_length
   for step_start_index in range(0, representation_length, step_repr_length):
diff --git a/tensor2tensor/trax/rl/space_serializer.py b/tensor2tensor/trax/rl/space_serializer.py
index 683aaf65d..3c4809c23 100644
--- a/tensor2tensor/trax/rl/space_serializer.py
+++ b/tensor2tensor/trax/rl/space_serializer.py
@@ -173,7 +173,7 @@ def __init__(self, space, vocab_size):
         "Discrete space size should fit in the number of symbols.")
 
   def serialize(self, data):
-    return np.reshape(data, (-1, 1))
+    return np.reshape(data, (-1, 1)).astype(np.int32)
 
   def deserialize(self, representation):
     return np.reshape(representation, -1)
@@ -200,7 +200,7 @@ def __init__(self, space, vocab_size):
     )
 
   def serialize(self, data):
-    return data
+    return data.astype(np.int32)
 
   def deserialize(self, representation):
     return representation

From e2e9a5ff9975ab5699c9c174012ba6290f3ead8a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 18 Sep 2019 18:53:53 -0700
Subject: [PATCH 2460/2720] Make dropout ranges part of model state and update
 state in Trainer step.

PiperOrigin-RevId: 269938775
---
 tensor2tensor/trax/layers/core.py        | 45 ++++++++++++++--------
 tensor2tensor/trax/models/transformer.py | 49 +++++++++++++-----------
 tensor2tensor/trax/trax.py               | 23 +++++++++++
 3 files changed, 79 insertions(+), 38 deletions(-)

diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 117f30388..fc4d8f2c1 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -168,22 +168,37 @@ def Flatten(x, params, n_axes_to_keep=1, **kwargs):
   return np.reshape(x, (x.shape[:n_axes_to_keep] + (-1,)))
 
 
-@base.layer()
-def Dropout(x, params, rate=0.0, mode='train', rng=None, **kwargs):
-  """Layer construction function for a dropout layer with given rate."""
-  del params, kwargs
-  if rng is None:
-    msg = ('Dropout layer requires apply_fn to be called with a rng keyword '
-           'argument. That is, instead of `Dropout(params, inputs)`, call '
-           'it like `Dropout(params, inputs, rng=key)`.')
-    raise ValueError(msg)
-  if rate >= 1.0:
-    raise ValueError('Dropout rate (%f) must be lower than 1.' % rate)
-  if mode == 'train' and rate > 0.0:
+class Dropout(base.Layer):
+  """Dropout."""
+
+  def __init__(self, rate=0.0, name='dropout', mode='train'):
+    super(Dropout, self).__init__()
+    self._initial_rate = rate
+    # TODO(lukaszkaiser): remove the name property by the end of September'19.
+    # It's only needed for a specific purpose in the short term, will go.
+    self._name = 'dropout_' + name
+    self._mode = mode
+
+  def new_parameters(self, input_shape, input_dtype, rng):
+    """Initialize dropout parameters and state."""
+    del input_shape, input_dtype, rng
+    return (), {self._name: np.array(self._initial_rate)}
+
+  def call(self, x, params, state, rng=None, **unused_kwargs):
+    """Execute dropout."""
+    del params
+    rate = self._initial_rate
+    if isinstance(state, dict) and self._name in state:
+      rate = state[self._name]
+    if rng is None:
+      msg = ('Dropout layer requires apply_fn to be called with a rng keyword '
+             'argument. That is, instead of `Dropout(params, inputs)`, call '
+             'it like `Dropout(params, inputs, rng=key)`.')
+      raise ValueError(msg)
+    if self._mode != 'train':
+      return x, state
     keep = backend.random.bernoulli(rng, 1.0 - rate, x.shape)
-    return np.where(keep, x / (1.0 - rate), np.zeros_like(x))
-  else:
-    return x
+    return np.where(keep, x / (1.0 - rate), np.zeros_like(x)), state
 
 
 @base.layer()
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index aaf729fa0..f40d558bc 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -23,19 +23,19 @@
 from tensor2tensor.trax import layers as tl
 
 
-def FeedForward(d_model, d_ff, dropout, mode):
+def FeedForward(d_model, d_ff, dropout, layer_idx, mode):
   """Feed-forward block with layer normalization at start."""
   return [
       tl.LayerNorm(),
       tl.Dense(d_ff),
       tl.Relu(),
-      tl.Dropout(rate=dropout, mode=mode),
+      tl.Dropout(rate=dropout, name='ff_middle_%d' % layer_idx, mode=mode),
       tl.Dense(d_model),
-      tl.Dropout(rate=dropout, mode=mode),
+      tl.Dropout(rate=dropout, name='ff_final_%d' % layer_idx, mode=mode),
   ]
 
 
-def EncoderBlock(d_model, d_ff, n_heads, dropout, mode):
+def EncoderBlock(d_model, d_ff, n_heads, dropout, layer_idx, mode):
   """Returns a layer sequence that implements a Transformer encoder block.
 
   The input to the layer sequence is a pair, (activations, mask), where the
@@ -47,6 +47,7 @@ def EncoderBlock(d_model, d_ff, n_heads, dropout, mode):
     d_ff: int: depth of feed-forward layer
     n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
+    layer_idx: which layer are we at (for bookkeeping)
     mode: str: 'train' or 'eval'
 
   Returns:
@@ -56,10 +57,10 @@ def EncoderBlock(d_model, d_ff, n_heads, dropout, mode):
   attention = [
       tl.LayerNorm(),
       tl.Attention(d_model, n_heads=n_heads, dropout=dropout, mode=mode),
-      tl.Dropout(rate=dropout, mode=mode),
+      tl.Dropout(rate=dropout, name='enc_attn_dropout', mode=mode),
   ]
   feed_forward = [
-      FeedForward(d_model, d_ff, dropout, mode=mode),
+      FeedForward(d_model, d_ff, dropout, layer_idx=layer_idx, mode=mode),
   ]
   return [
       tl.Residual(attention),
@@ -97,14 +98,14 @@ def TransformerEncoder(vocab_size,
   """
   embedder = [
       tl.Embedding(d_model, vocab_size),
-      tl.Dropout(rate=dropout, mode=mode),
+      tl.Dropout(rate=dropout, name='emb_dropout', mode=mode),
       tl.PositionalEncoding(max_len=max_len),
   ]
   return tl.Model([                             #      tokens
       tl.Dup(),                                 # toks toks
       tl.Parallel(embedder, tl.PaddingMask()),  # vecs mask
-      [EncoderBlock(d_model, d_ff, n_heads, dropout, mode)
-       for _ in range(n_layers)],               # vecs mask
+      [EncoderBlock(d_model, d_ff, n_heads, dropout, i, mode)
+       for i in range(n_layers)],               # vecs mask
       tl.Parallel([], tl.Drop()),               # ____  0
       tl.LayerNorm(),                           # vecs
       tl.Mean(axis=1),  # Average on length.    # vecs
@@ -114,7 +115,7 @@ def TransformerEncoder(vocab_size,
 
 
 def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-                 attention_type, dropout, share_kv, mode):
+                 attention_type, dropout, share_kv, layer_idx, mode):
   """Returns a layer sequence that implements a Transformer decoder block.
 
   The input to the layer sequence is an activation tensor.
@@ -128,6 +129,7 @@ def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
     attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
     share_kv: bool, whether to share keys and values
+    layer_idx: which layer are we at (for bookkeeping)
     mode: str: 'train' or 'eval'
 
   Returns:
@@ -139,10 +141,10 @@ def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
           d_model, n_heads=n_heads, d_attention_key=d_attention_key,
           d_attention_value=d_attention_value, attention_type=attention_type,
           share_kv=share_kv, mode=mode),
-      tl.Dropout(rate=dropout, mode=mode),  # vec
+      tl.Dropout(rate=dropout, name='attention_%d' % layer_idx, mode=mode),
   ]
   feed_forward = [
-      FeedForward(d_model, d_ff, dropout, mode=mode),
+      FeedForward(d_model, d_ff, dropout, layer_idx=layer_idx, mode=mode),
   ]
   return [
       tl.Residual(self_attention),
@@ -199,8 +201,8 @@ def TransformerDecoder(vocab_size=None,
       tl.PositionalEncoding(max_len=max_len),
       [DecoderBlock(  # pylint: disable=g-complex-comprehension
           d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-          attention_type, dropout, share_kv, mode)
-       for _ in range(n_layers)],   # vecs
+          attention_type, dropout, share_kv, i, mode)
+       for i in range(n_layers)],   # vecs
       tl.LayerNorm(),               # vecs
   )
 
@@ -244,7 +246,7 @@ def TransformerLM(vocab_size,
   """
   embedder = [
       tl.Embedding(d_model, vocab_size),
-      tl.Dropout(rate=dropout, mode=mode),
+      tl.Dropout(rate=dropout, name='embedding', mode=mode),
       tl.PositionalEncoding(max_len=max_len),
   ]
   return tl.Model(                  # tokens
@@ -252,15 +254,15 @@ def TransformerLM(vocab_size,
       embedder,                     # vecs
       [DecoderBlock(  # pylint: disable=g-complex-comprehension
           d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-          attention_type, dropout, share_kv, mode)
-       for _ in range(n_layers)],   # vecs
+          attention_type, dropout, share_kv, i, mode)
+       for i in range(n_layers)],   # vecs
       tl.LayerNorm(),               # vecs
       tl.Dense(vocab_size),         # vecs
       tl.LogSoftmax(),              # vecs
   )
 
 
-def EncoderDecoder(d_model, d_ff, n_heads, dropout, mode):
+def EncoderDecoder(d_model, d_ff, n_heads, dropout, layer_idx, mode):
   """Transformer encoder-decoder layer.
 
   The input is a triple (decoder_input, mask, encoder) where the mask is
@@ -272,6 +274,7 @@ def EncoderDecoder(d_model, d_ff, n_heads, dropout, mode):
     d_ff: int: depth of feed-forward layer
     n_heads: int: number of attention heads
     dropout: float: dropout rate (how much to drop out)
+    layer_idx: which layer are we at (for bookkeeping)
     mode: str: 'train' or 'eval'
 
   Returns:
@@ -293,7 +296,7 @@ def EncoderDecoder(d_model, d_ff, n_heads, dropout, mode):
       tl.Dropout(rate=dropout, mode=mode),  # vecs_d mask vecs_e
   ]
   feed_forward = [
-      FeedForward(d_model, d_ff, dropout, mode=mode),
+      FeedForward(d_model, d_ff, dropout, layer_idx=layer_idx, mode=mode),
   ]
   return [                                        # vecs_d masks vecs_e
       tl.Residual(decoder_self_attention),        # vecs_d masks vecs_e
@@ -348,12 +351,12 @@ def Transformer(input_vocab_size,
     ]
 
   encoder_stack = (  # masks vectors --> masks vectors
-      [EncoderBlock(d_model, d_ff, n_heads, dropout, mode)
-       for _ in range(n_layers)])
+      [EncoderBlock(d_model, d_ff, n_heads, dropout, i, mode)
+       for i in range(n_layers)])
 
   encoder_decoder_stack = (  # vecs_d masks vecs_e --> vecs_d masks vecs_e
-      [EncoderDecoder(d_model, d_ff, n_heads, dropout, mode)
-       for _ in range(n_layers)])
+      [EncoderDecoder(d_model, d_ff, n_heads, dropout, i, mode)
+       for i in range(n_layers)])
 
   # Input: encoder_side_tokens, decoder_side_tokens
   return tl.Model(  # tokens_e tokens_d
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index c533702f3..aa35544bd 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -717,6 +717,28 @@ def save_gin(self):
   def print_n_params(self):
     _print_n_params(self._opt_state, self._n_devices, self._step)
 
+  def _map_to_state_dicts(self, f):
+    """Map the function f to all dicts in model state."""
+    def nested_map(x, f):
+      if isinstance(x, list):
+        return [nested_map(y, f) for y in x]
+      if isinstance(x, tuple):
+        return tuple([nested_map(y, f) for y in x])
+      if isinstance(x, dict) and len(x) == 1:
+        return f(x)
+      return x
+    return nested_map(self._model_state, f)
+
+  def _state_dicts_update(self, state_dict):
+    assert len(state_dict.keys()) == 1
+    key = list(state_dict.keys())[0]
+    value = np.array(state_dict[key])
+    return {key: np.array(self.update_model_state(key, value))}
+
+  def update_model_state(self, key, value):
+    del key
+    return value
+
   def _train_step(self, next_train_batch):
     """Run one training step and update self._opt_state."""
     # Calculate the current learning rate.
@@ -729,6 +751,7 @@ def _train_step(self, next_train_batch):
     # Run the update.
     (params, slots), self._model_state, self._rngs = self._jit_update_fn(
         self._step, opt_state, next_train_batch, self._model_state, self._rngs)
+    self._model_state = self._map_to_state_dicts(self._state_dicts_update)
     self._opt_state = opt_state._replace(params=params, slots=slots)
     self._step += 1
 

From 67e31d16fed00d6b318410acede3a116b5db534e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 18 Sep 2019 19:29:23 -0700
Subject: [PATCH 2461/2720] Add a multi-device put function for
 memory-efficient parameter replication in Trax.

PiperOrigin-RevId: 269942713
---
 tensor2tensor/trax/trax.py | 43 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 42 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index aa35544bd..02e513bb1 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -525,6 +525,44 @@ def reshape_by_device(x, n_devices):
       x, lambda x: _reshape_by_device_single(x, n_devices))
 
 
+def multi_device_put(x, devices=None, reuse=True):
+  """Memory efficient multi-device replication in JAX.
+
+  Args:
+    x: jax DeviceArray or numpy ndarray to be replicated.
+    devices: a jax.devices() list or subset thereof of devices to
+      replicate onto.  Should match the list passed to any pmaps
+      ingesting the replicated array.
+    reuse: bool. If x is a DeviceArray whether to reuse its backing
+      device_buffer in the resulting ShardedDeviceArray.
+
+  Returns:
+    A ShardedDeviceArray with dtype = x.dtype and shape =
+    (n_devices,) + x.shape that's backed by replica
+    device_buffers on each device.
+  """
+  # Convert _FilledConstants that don't have device_buffer, etc.
+  if type(x) != jax.xla.DeviceArray:  # pylint: disable=unidiomatic-typecheck
+    x = np.array(x)
+  if not devices:
+    devices = jax.devices()
+  n_devices = len(devices)
+  x_aval = jax.xla.abstractify(x)
+  broadcast_x_aval = jax.abstract_arrays.ShapedArray(
+      (n_devices,) + x_aval.shape,
+      x_aval.dtype)
+  if reuse:
+    other_device_ordinals = [dv.id for dv in jax.devices()
+                             if dv != x.device_buffer.device()]
+    broadcast_buffers = ([x.device_buffer,] +
+                         [jax.xla.xc.Buffer.from_pyval(x, device=i)
+                          for i in other_device_ordinals])
+  else:
+    broadcast_buffers = [jax.xla.xc.Buffer.from_pyval(x, device=i)
+                         for i in range(n_devices)]
+  return jax.pxla.ShardedDeviceArray(broadcast_x_aval, broadcast_buffers)
+
+
 def _repeat_stream(stream):
   """Repeat a stream indefinitely."""
   while True:
@@ -701,7 +739,10 @@ def optimizer_params(self):
 
   def _maybe_replicate(self, x):
     if self._n_devices > 1:
-      return np.broadcast_to(x, (self._n_devices,) + x.shape)
+      if backend.get_name() == "jax":
+        return multi_device_put(x)
+      else:
+        return np.broadcast_to(x, (self._n_devices,) + x.shape)
     else:
       return x
 

From 54e448e58787df22f8b6673067e2dfe52d7290b2 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 18 Sep 2019 20:50:24 -0700
Subject: [PATCH 2462/2720] Add a config for reversible transformer with merged
 hashed attention on enwik8 for Trax.

PiperOrigin-RevId: 269951277
---
 .../configs/transformer_revnet_enwik8.gin     | 56 +++++++++++++++++++
 1 file changed, 56 insertions(+)
 create mode 100644 tensor2tensor/trax/configs/transformer_revnet_enwik8.gin

diff --git a/tensor2tensor/trax/configs/transformer_revnet_enwik8.gin b/tensor2tensor/trax/configs/transformer_revnet_enwik8.gin
new file mode 100644
index 000000000..6bb673fb6
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_revnet_enwik8.gin
@@ -0,0 +1,56 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 1
+batch_fun.eval_batch_size = 8
+batch_fun.max_eval_length = 65536
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_enwik8_l65k'
+inputs.input_name = 'targets'
+inputs.n_chunks = 16
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 2.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 200
+train.eval_steps = 8
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerRevnetLM
+train.optimizer = @trax.optimizers.Adafactor
+train.train_steps = 20000
+train.trainer_class = @MemoryEfficientTrainer
+
+# Parameters for MergedHashedCausalAttention:
+# ==============================================================================
+MergedHashedCausalAttention.dropout = 0.0
+MergedHashedCausalAttention.n_bins = 512
+MergedHashedCausalAttention.bin_by_time = True
+MergedHashedCausalAttention.one_rng = False
+
+# Parameters for TransformerRevnetLM:
+# ==============================================================================
+TransformerRevnetLM.d_model = 512
+TransformerRevnetLM.d_ff = 2048
+TransformerRevnetLM.d_attention_key = 64
+TransformerRevnetLM.d_attention_value = 64
+TransformerRevnetLM.dropout = 0.1
+TransformerRevnetLM.max_len = 65536
+TransformerRevnetLM.mode = 'train'
+TransformerRevnetLM.n_heads = 4
+TransformerRevnetLM.n_layers = 12
+TransformerRevnetLM.vocab_size = 258
+TransformerRevnetLM.n_chunks = 16
+TransformerRevnetLM.n_attention_chunks = 1
+TransformerRevnetLM.attention_type = @trax.layers.MergedHashedCausalAttention

From 190532f278cf3a74c0b7430613114930e7ae3e46 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 18 Sep 2019 23:44:14 -0700
Subject: [PATCH 2463/2720] Add an option for initializing the policy network
 using parameters of a pretrained world model.

In SimPLe, this initialization will occur only before training on simulated environment, so we don't need to pass this parameter to the async workers.

PiperOrigin-RevId: 269971904
---
 tensor2tensor/trax/rl/ppo.py              | 14 ++++++++
 tensor2tensor/trax/rl/ppo_test.py         | 44 +++++++++++++++++++++++
 tensor2tensor/trax/rl/ppo_trainer.py      |  7 ++++
 tensor2tensor/trax/rl/ppo_trainer_test.py | 11 +++---
 tensor2tensor/trax/rl/simple_trainer.py   |  6 +++-
 5 files changed, 76 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index 2518ece0e..48b676028 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -910,6 +910,20 @@ def save_opt_state(output_dir,
       gfile.remove(path)
 
 
+def init_policy_from_world_model_checkpoint(policy_params, model_output_dir):
+  """Initializes policy parameters from world model parameters."""
+  pkl_module = utils.get_pickle_module()
+  params_file = os.path.join(model_output_dir, "model.pkl")
+  # Don't use trax.restore_state to avoid a circular import.
+  with gfile.GFile(params_file, "rb") as f:
+    model_params = pkl_module.load(f)[0][0]
+  # TODO(pkozakowski): The following, brittle line of code is hardcoded for
+  # transplanting parameters from TransformerLM to TransformerDecoder-based
+  # policy network of the same configuration. Figure out a more general method.
+  policy_params[0] = model_params[1:-2]
+  return policy_params
+
+
 def write_eval_reward_summaries(reward_stats_by_mode, summary_writer, epoch):
   """Writes evaluation reward statistics to summary and logs them.
 
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
index d71772647..850bc6329 100644
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -19,10 +19,13 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
+
 import jax
 from jax import random as jax_random
 import numpy as np
 from tensor2tensor.trax import layers
+from tensor2tensor.trax import models
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.rl import ppo
 from tensorflow import test
@@ -559,6 +562,47 @@ def test_saves_and_restores_opt_state(self):
     restored_data = ppo.maybe_restore_opt_state(output_dir)
     self.assertEqual(restored_data, (opt_state, state, epoch, opt_step))
 
+  def test_inits_policy_by_world_model_checkpoint(self):
+    transformer_kwargs = {
+        "d_model": 1,
+        "d_ff": 1,
+        "n_layers": 1,
+        "n_heads": 1,
+        "max_len": 128,
+        "mode": "train",
+    }
+    rng = jax_random.PRNGKey(123)
+    init_kwargs = {
+        "input_shapes": (1, 1),
+        "input_dtype": np.int32,
+        "rng": rng,
+    }
+    model = models.TransformerLM(vocab_size=4, **transformer_kwargs)
+    (model_params, _) = model.initialize(**init_kwargs)
+    policy = ppo.policy_and_value_net(
+        n_actions=3,
+        n_controls=2,
+        vocab_size=4,
+        bottom_layers_fn=functools.partial(
+            models.TransformerDecoder, **transformer_kwargs
+        ),
+        two_towers=False,
+    )
+    (policy_params, policy_state) = policy.initialize(**init_kwargs)
+    output_dir = self.get_temp_dir()
+    # Initialize state by restoring from a nonexistent checkpoint.
+    trax_state = trax.restore_state(output_dir)
+    trax_state = trax_state._replace(opt_state=(model_params, None))
+    # Save world model parameters.
+    trax.save_state(trax_state, output_dir)
+    # Initialize policy parameters from world model parameters.
+    new_policy_params = ppo.init_policy_from_world_model_checkpoint(
+        policy_params, output_dir
+    )
+    # Try to run the policy with new parameters.
+    observations = np.zeros((1, 100), dtype=np.int32)
+    policy(observations, new_policy_params, state=policy_state, rng=rng)
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index c809d1b5e..e4d5c54d4 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -81,6 +81,7 @@ def __init__(self,
                len_history_for_policy=4,
                eval_temperatures=(1.0, 0.5),
                separate_eval=True,
+               init_policy_from_world_model_output_dir=None,
                **kwargs):
     """Creates the PPO trainer.
 
@@ -122,6 +123,8 @@ def __init__(self,
       separate_eval: Whether to run separate evaluation using a set of
         temperatures. If False, the training reward is reported as evaluation
         reward with temperature 1.0.
+      init_policy_from_world_model_output_dir: Model output dir for initializing
+        the policy. If None, initialize randomly.
       **kwargs: Additional keyword arguments passed to the base class.
     """
     # Set in base class constructor.
@@ -185,6 +188,10 @@ def __init__(self,
     (batch_obs_shape, obs_dtype) = self._batch_obs_shape_and_dtype
     policy_and_value_net_params, self._model_state = (
         policy_and_value_net.initialize(batch_obs_shape, obs_dtype, key1))
+    if init_policy_from_world_model_output_dir is not None:
+      policy_and_value_net_params = ppo.init_policy_from_world_model_checkpoint(
+          policy_and_value_net_params, init_policy_from_world_model_output_dir
+      )
 
     # Initialize the optimizer.
     (policy_and_value_opt_state, self._policy_and_value_opt_update,
diff --git a/tensor2tensor/trax/rl/ppo_trainer_test.py b/tensor2tensor/trax/rl/ppo_trainer_test.py
index 295e79c5c..c93804a8d 100644
--- a/tensor2tensor/trax/rl/ppo_trainer_test.py
+++ b/tensor2tensor/trax/rl/ppo_trainer_test.py
@@ -81,6 +81,7 @@ def _make_trainer(
         n_optimizer_steps=1,
         output_dir=output_dir,
         random_seed=0,
+        max_timestep=3,
         boundary=2,
         save_every_n=1,
         **kwargs
@@ -127,16 +128,16 @@ def test_training_loop_onlinetune(self):
           output_shape=(1, 1),
           output_dtype=np.float32,
       ))
-      gin.bind_parameter("OnlineTuneEnv.train_steps", 2)
-      gin.bind_parameter("OnlineTuneEnv.eval_steps", 2)
+      gin.bind_parameter("OnlineTuneEnv.train_steps", 1)
+      gin.bind_parameter("OnlineTuneEnv.eval_steps", 1)
       gin.bind_parameter(
           "OnlineTuneEnv.output_dir", os.path.join(output_dir, "envs"))
       trainer = self._make_trainer(
-          train_env=self.get_wrapped_env("OnlineTuneEnv-v0", 2),
-          eval_env=self.get_wrapped_env("OnlineTuneEnv-v0", 2),
+          train_env=self.get_wrapped_env("OnlineTuneEnv-v0", 1),
+          eval_env=self.get_wrapped_env("OnlineTuneEnv-v0", 1),
           output_dir=output_dir,
       )
-      trainer.training_loop(n_epochs=2)
+      trainer.training_loop(n_epochs=1)
 
   def test_training_loop_simulated(self):
     n_actions = 5
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index c615fc8bb..bad077cc5 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -57,21 +57,25 @@ def __init__(self,
                initial_trajectory_dir=None,
                initial_trajectory_mix_prob=0.5,
                initial_model=None,
+               init_policy_from_world_model=False,
                **kwargs):
     super(SimPLe, self).__init__(train_env, eval_env, output_dir, **kwargs)
     self._policy_dir = os.path.join(output_dir, "policy")
+    self._model_dir = os.path.join(output_dir, "model")
     self._policy_trainer = policy_trainer_class(
         train_env=train_env,
         eval_env=eval_env,
         output_dir=self._policy_dir,
         async_mode=self._async_mode,
+        init_policy_from_world_model_output_dir=(
+            self._model_dir if init_policy_from_world_model else None
+        ),
     )
     self._n_real_epochs = n_real_epochs
     self._model_train_batch_size = model_train_batch_size
     self._n_model_initial_train_steps = n_model_initial_train_steps
     self._n_model_train_steps_per_epoch = n_model_train_steps_per_epoch
     self._data_eval_frac = data_eval_frac
-    self._model_dir = os.path.join(output_dir, "model")
 
     gfile.makedirs(self._model_dir)
     if initial_model is not None:

From bbadf12fe46a2f074c7c1f1f6f818b977c83086c Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 19 Sep 2019 01:09:03 -0700
Subject: [PATCH 2464/2720] Add a config for training translation as a langauge
 model with Trax.

PiperOrigin-RevId: 269982969
---
 .../configs/transformer_lm_wmt_ende_8gb.gin   | 64 +++++++++++++++++++
 tensor2tensor/trax/inputs.py                  | 24 ++++++-
 2 files changed, 86 insertions(+), 2 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/transformer_lm_wmt_ende_8gb.gin

diff --git a/tensor2tensor/trax/configs/transformer_lm_wmt_ende_8gb.gin b/tensor2tensor/trax/configs/transformer_lm_wmt_ende_8gb.gin
new file mode 100644
index 000000000..7fa1f2504
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_lm_wmt_ende_8gb.gin
@@ -0,0 +1,64 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 64
+batch_fun.eval_batch_size = 64
+batch_fun.bucket_length=64
+batch_fun.max_eval_length = 1024
+batch_fun.buckets_include_inputs_in_length=True
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_translate_ende_wmt32k'
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 1.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 10000
+
+# Parameters for Adafactor:
+# ==============================================================================
+Adafactor.beta1 = 0.0
+Adafactor.decay_rate = 0.8
+Adafactor.clipping_threshold = 1.0
+Adafactor.epsilon1 = 1e-30
+Adafactor.epsilon2 = 0.001
+Adafactor.factored = True
+Adafactor.multiply_by_parameter_scale = True
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_concat_preprocess
+wmt_concat_preprocess.max_length = 255
+wmt_concat_preprocess.max_eval_length = 511
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 1000
+train.eval_steps = 10
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerLM
+train.train_steps = 500000
+train.optimizer = @trax.optimizers.Adafactor
+train.has_weights = True
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.d_model = 512
+TransformerLM.d_ff = 2048
+TransformerLM.dropout = 0.1
+TransformerLM.max_len = 2048
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 8
+TransformerLM.n_layers = 6
+TransformerLM.vocab_size = 33300
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index e962c965e..da4d2c5c5 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -253,7 +253,9 @@ def random_minibatches(length_list):
 def dataset_to_stream(dataset, input_name, n_chunks=0):
   """Takes a tf.Dataset and creates a numpy stream of ready batches."""
   for example in backend.dataset_as_numpy(dataset):
-    inp, out = example[0][input_name], example[1]
+    features = example[0]
+    inp, out = features[input_name], example[1]
+    mask = features['mask'] if 'mask' in features else None
     # All input-pipeline processing should be on CPU.
     with tf.device('cpu:0'):
       # Some accelerators don't handle uint8 well, cast to int.
@@ -266,7 +268,7 @@ def dataset_to_stream(dataset, input_name, n_chunks=0):
       if n_chunks > 0:
         inp = tuple(np.split(inp, n_chunks, axis=1))
         out = tuple(np.split(out, n_chunks, axis=1))
-    yield inp, out
+    yield (inp, out) if mask is None else (inp, out, mask)
 
 
 @gin.configurable(whitelist=['train_shuffle_files', 'eval_shuffle_files',
@@ -568,6 +570,24 @@ def eval_right_length(example, target):
   return dataset
 
 
+@gin.configurable(blacklist=['dataset', 'training'])
+def wmt_concat_preprocess(dataset, training, max_length=-1, max_eval_length=-1):
+  """Preprocessing for WMT: filter exceeding maximum length and concatenate."""
+  dataset = wmt_preprocess(dataset, training, max_length, max_eval_length)
+
+  def concat_and_add_mask(features, targets):
+    inp = features['inputs']
+    pad = tf.expand_dims(tf.zeros_like(inp[0]), axis=0)
+    concat = tf.concat([inp, pad, targets], axis=0)
+    mask = tf.concat([tf.zeros_like(inp), pad, tf.ones_like(targets)], axis=0)
+    features['inputs'] = concat
+    features['mask'] = mask
+    return features, concat
+
+  dataset = dataset.map(concat_and_add_mask)
+  return dataset
+
+
 @gin.configurable(whitelist=['preprocess_fun', 'shuffle_buffer_size'])
 def shuffle_and_batch_data(dataset,
                            target_names,

From 574539fdff0cc0c240d738ea5dd754097553a02c Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 19 Sep 2019 14:07:42 -0700
Subject: [PATCH 2465/2720] Control model state from OnlineTuneEnv.

PiperOrigin-RevId: 270121144
---
 tensor2tensor/trax/rl/envs/online_tune_env.py |  2 +
 .../trax/rl/envs/online_tune_env_test.py      |  4 +-
 tensor2tensor/trax/trax.py                    | 48 ++++++++++++++-----
 3 files changed, 40 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index 74104d7fe..76065c535 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -71,6 +71,7 @@ def __init__(self,
                    # (name, start, (low, high), flip)
                    ("learning_rate", 1e-3, (1e-9, 10.0), False),
                ),
+               nontrainable_param_map=None,
                metric_range=(0.0, 5.0),
                # Don't save checkpoints by default, as they tend to use a lot of
                # space.
@@ -89,6 +90,7 @@ def __init__(self,
         lr_schedule=(lambda history: lambda step: self._current_controls),
         inputs=inputs,
         should_save=should_save_checkpoints,
+        nontrainable_param_map=nontrainable_param_map,
     )
     self._trainer = None
     self._action_multipliers = action_multipliers
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env_test.py b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
index b2b95c2d9..4f83fb9a3 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env_test.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
@@ -51,7 +51,7 @@ def reset(self, output_dir):
 
   def train_epoch(self, epoch_steps, eval_steps):
     del epoch_steps
-    self.controls.append(self.optimizer_params)
+    self.controls.append(self.nontrainable_params)
     self.evaluate(eval_steps)
 
   def evaluate(self, eval_steps):
@@ -61,7 +61,7 @@ def evaluate(self, eval_steps):
         metric=METRIC,
         step=self.step,
         value=self.metrics_to_report.pop(0))
-    for (name, value) in self.optimizer_params.items():
+    for (name, value) in self.nontrainable_params.items():
       (mode, metric) = online_tune.control_metric(name)
       self.state.history.append(
           mode=mode,
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 02e513bb1..88eaabd52 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -580,7 +580,8 @@ class Trainer(object):
 
   def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs,
                output_dir=None, random_seed=None, n_devices=None,
-               save_steps=None, should_save=True, has_weights=False):
+               save_steps=None, should_save=True, has_weights=False,
+               nontrainable_param_map=None):
     if save_steps is None:
       save_steps = []
     self._save_steps = save_steps
@@ -655,8 +656,15 @@ def initialize(input_shape, input_dtype, target_shape, target_dtype, rng):
     self._model_train = model_train
     self._model_predict_eval = model_predict_eval
     self._loss_fn = loss_fn
+    # TODO(pkozakowski): "Learning rate schedules" are currently able to control
+    # control all optimizer parameters and model state, so let's rename them
+    # accordingly.
     self._lr_schedule = lr_schedule
 
+    if nontrainable_param_map is None:
+      nontrainable_param_map = {}
+    self._nontrainable_param_map = nontrainable_param_map
+
     # Those fields will be set in reset().
     self._output_dir = None
     self._train_sw = None
@@ -713,7 +721,7 @@ def reset(self, output_dir):
     if not state.opt_state:
       self._maybe_save_state(keep=False)
 
-    self.update_optimizer_params()
+    self.update_nontrainable_params()
 
   @property
   def step(self):
@@ -730,7 +738,7 @@ def state(self):
         model_state=self._model_state)
 
   @property
-  def optimizer_params(self):
+  def nontrainable_params(self):
     # TODO(lukaszkaiser): it makes no sense to use an accelerator (e.g. TPU)
     # in op-by-op mode just to compute the learning rate. However, there
     # should be a cleaner approach that forceably swapping out the backend.
@@ -777,14 +785,30 @@ def _state_dicts_update(self, state_dict):
     return {key: np.array(self.update_model_state(key, value))}
 
   def update_model_state(self, key, value):
-    del key
+    """Updates model state based on nontrainable_params."""
+    # Translate model state keys to nontrainable param names.
+    if key in self._nontrainable_param_map:
+      param_name = self._nontrainable_param_map[key]
+    else:
+      # If a key is not in mapping, it stays the same.
+      param_name = key
+    if param_name in self.nontrainable_params:
+      if self._step == 0:
+        log("Mapping model state key {} to nontrainable param {}.".format(
+            key, param_name
+        ))
+        return self._maybe_replicate(
+            np.array(self.nontrainable_params[param_name])
+        )
     return value
 
   def _train_step(self, next_train_batch):
     """Run one training step and update self._opt_state."""
-    # Calculate the current learning rate.
+    # Calculate the current optimizer parameters.
+    # TODO(pkozakowski): Optimizer parameters get polluted with model state,
+    # which doesn't break anything but is weird. Filter it out.
     opt_param_updates = layers.nested_map(
-        self.optimizer_params, lambda x: self._maybe_replicate(np.array(x))
+        self.nontrainable_params, lambda x: self._maybe_replicate(np.array(x))
     )
     opt_state = self._opt_state
     opt_state.opt_params.update(opt_param_updates)
@@ -815,9 +839,9 @@ def train_epoch(self, epoch_steps, eval_steps):
       if self._step in self._save_steps:
         self._maybe_save_state(keep=True)
 
-      # Log optimizer params (learning rate etc.)
+      # Log nontrainable params (learning rate, dropout etc.)
       if self._step == 1 or self._step % 10 == 0:
-        for (name, value) in self.optimizer_params.items():
+        for (name, value) in self.nontrainable_params.items():
           self._train_sw.scalar("training/{}".format(name), value)
 
     # Timer
@@ -856,11 +880,11 @@ def evaluate(self, eval_steps):
         has_weights=self._has_weights)
 
     # Save the optimizer params in the history
-    for (name, value) in self.optimizer_params.items():
+    for (name, value) in self.nontrainable_params.items():
       self._history.append("train", "training/{}".format(name), self._step,
                            value)
 
-  def update_optimizer_params(self):
+  def update_nontrainable_params(self):
     self._lr_fn = self._lr_schedule(self._history)
 
   def save_computation_graphs(self, save_backward_graph):
@@ -990,8 +1014,8 @@ def train(output_dir,
   for epoch_steps in epochs(train_steps, trainer.step, epoch_steps):
     trainer.train_epoch(epoch_steps, eval_steps)
 
-    # Update optimizer parameters with new history
-    trainer.update_optimizer_params()
+    # Update nontrainable parameters with new history
+    trainer.update_nontrainable_params()
 
     # Bookkeeping we do at the first step
     if trainer.step == 1:

From 3881bd5a840d1f21d2ec0a8f6bd5bfb13c4cdcff Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 19 Sep 2019 15:50:01 -0700
Subject: [PATCH 2466/2720] Redirect OnlineTuneEnv's output_dir in async mode.

PiperOrigin-RevId: 270143202
---
 tensor2tensor/trax/rl/envs/async_trajectory_collector.py | 9 ++++++---
 tensor2tensor/trax/rl_trainer.py                         | 9 ++++++---
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/trax/rl/envs/async_trajectory_collector.py b/tensor2tensor/trax/rl/envs/async_trajectory_collector.py
index 7eb79131c..339f3d34b 100644
--- a/tensor2tensor/trax/rl/envs/async_trajectory_collector.py
+++ b/tensor2tensor/trax/rl/envs/async_trajectory_collector.py
@@ -48,6 +48,7 @@
     "policy is available.")
 
 flags.DEFINE_string("output_dir", "", "Output dir.")
+flags.DEFINE_string("envs_output_dir", "", "Output dir for the envs.")
 
 flags.DEFINE_boolean(
     "jax_debug_nans", False,
@@ -110,9 +111,11 @@ def create_envs_and_collect_trajectories(
   train_env_kwargs = {}
   eval_env_kwargs = {}
   if "OnlineTuneEnv" in env_name:
-    # TODO(pkozakowski): Separate env output dirs by train/eval and epoch.
-    train_env_kwargs = {"output_dir": os.path.join(output_dir, "envs/train")}
-    eval_env_kwargs = {"output_dir": os.path.join(output_dir, "envs/eval")}
+    envs_output_dir = FLAGS.envs_output_dir or os.path.join(output_dir, "envs")
+    train_env_output_dir = os.path.join(envs_output_dir, "train")
+    eval_env_output_dir = os.path.join(envs_output_dir, "eval")
+    train_env_kwargs = {"output_dir": train_env_output_dir}
+    eval_env_kwargs = {"output_dir": eval_env_output_dir}
 
   if "ClientEnv" in env_name:
     train_env_kwargs["per_env_kwargs"] = [{
diff --git a/tensor2tensor/trax/rl_trainer.py b/tensor2tensor/trax/rl_trainer.py
index f38888b83..65c4b7ed1 100644
--- a/tensor2tensor/trax/rl_trainer.py
+++ b/tensor2tensor/trax/rl_trainer.py
@@ -57,6 +57,7 @@
 flags.DEFINE_boolean("disable_jit", False, "Setting to true will disable jit.")
 
 flags.DEFINE_string("output_dir", "", "Output dir.")
+flags.DEFINE_string("envs_output_dir", "", "Output dir for the envs.")
 flags.DEFINE_multi_string("config_file", None,
                           "Configuration file with parameters (.gin).")
 flags.DEFINE_multi_string("config", None,
@@ -127,9 +128,11 @@ def train_rl(
   train_env_kwargs = {}
   eval_env_kwargs = {}
   if "OnlineTuneEnv" in env_name:
-    # TODO(pkozakowski): Separate env output dirs by train/eval and epoch.
-    train_env_kwargs = {"output_dir": os.path.join(output_dir, "envs/train")}
-    eval_env_kwargs = {"output_dir": os.path.join(output_dir, "envs/eval")}
+    envs_output_dir = FLAGS.envs_output_dir or os.path.join(output_dir, "envs")
+    train_env_output_dir = os.path.join(envs_output_dir, "train")
+    eval_env_output_dir = os.path.join(envs_output_dir, "eval")
+    train_env_kwargs = {"output_dir": train_env_output_dir}
+    eval_env_kwargs = {"output_dir": eval_env_output_dir}
 
   if "ClientEnv" in env_name:
     train_env_kwargs["per_env_kwargs"] = [{

From 6680649e959c18e2d2e6e3bf01c1ba6efe97eabd Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 19 Sep 2019 16:59:00 -0700
Subject: [PATCH 2467/2720] Add gin configs for TransformerLM tuning on various
 tasks.

PiperOrigin-RevId: 270156843
---
 ...nline_tune_transformer_imagenet64_16gb.gin | 105 ++++++++++++++++++
 .../env_online_tune_transformer_lm1b_16gb.gin | 102 +++++++++++++++++
 ...line_tune_transformer_lm_wmt_ende_16gb.gin |  99 +++++++++++++++++
 .../trax/rl/configs/ppo_online_tune.gin       |   2 +-
 tensor2tensor/trax/rl/envs/online_tune_env.py |   4 +-
 5 files changed, 310 insertions(+), 2 deletions(-)
 create mode 100644 tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
 create mode 100644 tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
 create mode 100644 tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin

diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
new file mode 100644
index 000000000..8c6ee7f64
--- /dev/null
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
@@ -0,0 +1,105 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.rl
+import tensor2tensor.trax.rl.envs
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 1
+batch_fun.eval_batch_size = 16
+batch_fun.max_eval_length = 12288  # 64 * 64 * 3
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
+inputs.input_name = 'targets'
+
+# Parameters for train_and_eval_dataset:
+# ==============================================================================
+train_and_eval_dataset.eval_holdout_size = 0.05
+train_and_eval_dataset.eval_shuffle_files = True
+
+# Parameters for MemoryEfficientCausalAttention:
+# ==============================================================================
+MemoryEfficientCausalAttention.dropout = 0.0
+MemoryEfficientCausalAttention.loop_stride = 512
+
+# Parameters for MergedHashedCausalAttention:
+# ==============================================================================
+MergedHashedCausalAttention.dropout = 0.0
+MergedHashedCausalAttention.n_bins = 16
+MergedHashedCausalAttention.bin_by_time = True
+MergedMultiHashedCausalAttention.one_rng = False
+
+# Parameters for MergedMultiHashedCausalAttention:
+# ==============================================================================
+MergedMultiHashedCausalAttention.dropout = 0.0
+MergedMultiHashedCausalAttention.n_bins = 64
+MergedMultiHashedCausalAttention.n_hashes = 2
+MergedMultiHashedCausalAttention.n_buckets_per_bin = 2
+MergedMultiHashedCausalAttention.bin_by_time = False
+MergedMultiHashedCausalAttention.one_rng = False
+MergedMultiHashedCausalAttention.drop_for_hash_rate = 0.1
+MergedMultiHashedCausalAttention.hard_k = 32
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.attention_type = @trax.layers.MergedMultiHashedCausalAttention
+TransformerLM.d_attention_key = 64
+TransformerLM.d_attention_value = 64
+TransformerLM.d_model = 1024
+TransformerLM.d_ff = 2048
+TransformerLM.dropout = 0.0
+TransformerLM.max_len = 12288  # 64 * 64 * 3
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 4
+TransformerLM.n_layers = 3
+TransformerLM.share_kv = True
+TransformerLM.vocab_size = 256
+
+# Parameters for OnlineTuneEnv:
+# ==============================================================================
+OnlineTuneEnv.inputs = @trax.inputs.inputs
+OnlineTuneEnv.model = @trax.models.TransformerLM
+OnlineTuneEnv.optimizer = @trax.optimizers.Adafactor
+OnlineTuneEnv.control_configs = (
+    ("learning_rate", 1e-3, (1e-9, 1e-2), False),
+    ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
+
+    ("dropout_embedding", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_initial", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_initial", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_initial", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_middle", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_middle", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_middle", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_final", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_final", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_final", 0.1, (0.0, 0.9), True),
+)
+OnlineTuneEnv.nontrainable_param_map = {
+    # "dropout_{layer_type}_{block_index}": "dropout_{layer_type}_{block_group}"
+
+    "dropout_attention_0": "dropout_attention_initial",
+    "dropout_ff_middle_0": "dropout_ff_middle_initial",
+    "dropout_ff_final_0": "dropout_ff_final_initial",
+
+    "dropout_attention_1": "dropout_attention_middle",
+    "dropout_ff_middle_1": "dropout_ff_middle_middle",
+    "dropout_ff_final_1": "dropout_ff_final_middle",
+
+    "dropout_attention_2": "dropout_attention_final",
+    "dropout_ff_middle_2": "dropout_ff_middle_final",
+    "dropout_ff_final_2": "dropout_ff_final_final",
+}
+OnlineTuneEnv.include_controls_in_observation = True
+OnlineTuneEnv.train_steps = 150
+OnlineTuneEnv.eval_steps = 2
+OnlineTuneEnv.env_steps = 100
+OnlineTuneEnv.metric_range = (0.0, 10.0)
+OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
new file mode 100644
index 000000000..7bc28c4c1
--- /dev/null
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
@@ -0,0 +1,102 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.rl
+import tensor2tensor.trax.rl.envs
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 256
+batch_fun.eval_batch_size = 256
+batch_fun.max_eval_length = 2048
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_languagemodel_lm1b32k'
+inputs.input_name = 'targets'
+
+# Parameters for train_and_eval_dataset:
+# ==============================================================================
+train_and_eval_dataset.eval_holdout_size = 0.05
+train_and_eval_dataset.eval_shuffle_files = True
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
+lm1b_preprocess.max_target_length = 512
+lm1b_preprocess.max_eval_target_length = 2048
+
+# Parameters for DotProductCausalAttention:
+# ==============================================================================
+DotProductCausalAttention.dropout = 0.1
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
+TransformerLM.d_model = 512
+TransformerLM.d_ff = 2048
+TransformerLM.dropout = 0.1
+TransformerLM.max_len = 2048
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 8
+TransformerLM.n_layers = 6
+TransformerLM.vocab_size = 32000
+
+# Parameters for OnlineTuneEnv:
+# ==============================================================================
+OnlineTuneEnv.inputs = @trax.inputs.inputs
+OnlineTuneEnv.model = @trax.models.TransformerLM
+OnlineTuneEnv.optimizer = @trax.optimizers.Adafactor
+OnlineTuneEnv.control_configs = (
+    ("learning_rate", 1e-3, (1e-9, 1e-2), False),
+    ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
+
+    ("dropout_embedding", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_initial", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_initial", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_initial", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_middle", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_middle", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_middle", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_final", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_final", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_final", 0.1, (0.0, 0.9), True),
+)
+OnlineTuneEnv.nontrainable_param_map = {
+    # "dropout_{layer_type}_{block_index}": "dropout_{layer_type}_{block_group}"
+
+    "dropout_attention_0": "dropout_attention_initial",
+    "dropout_ff_middle_0": "dropout_ff_middle_initial",
+    "dropout_ff_final_0": "dropout_ff_final_initial",
+
+    "dropout_attention_1": "dropout_attention_middle",
+    "dropout_ff_middle_1": "dropout_ff_middle_middle",
+    "dropout_ff_final_1": "dropout_ff_final_middle",
+    "dropout_attention_2": "dropout_attention_middle",
+    "dropout_ff_middle_2": "dropout_ff_middle_middle",
+    "dropout_ff_final_2": "dropout_ff_final_middle",
+    "dropout_attention_3": "dropout_attention_middle",
+    "dropout_ff_middle_3": "dropout_ff_middle_middle",
+    "dropout_ff_final_3": "dropout_ff_final_middle",
+    "dropout_attention_4": "dropout_attention_middle",
+    "dropout_ff_middle_4": "dropout_ff_middle_middle",
+    "dropout_ff_final_4": "dropout_ff_final_middle",
+
+    "dropout_attention_5": "dropout_attention_final",
+    "dropout_ff_middle_5": "dropout_ff_middle_final",
+    "dropout_ff_final_5": "dropout_ff_final_final",
+}
+OnlineTuneEnv.include_controls_in_observation = True
+OnlineTuneEnv.train_steps = 300
+OnlineTuneEnv.eval_steps = 1
+OnlineTuneEnv.env_steps = 100
+OnlineTuneEnv.metric_range = (0.0, 10.0)
+OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
new file mode 100644
index 000000000..4a91735f4
--- /dev/null
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
@@ -0,0 +1,99 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.rl
+import tensor2tensor.trax.rl.envs
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 128
+batch_fun.eval_batch_size = 128
+batch_fun.bucket_length = 64
+batch_fun.max_eval_length = 512
+batch_fun.buckets_include_inputs_in_length = True
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_translate_ende_wmt32k'
+
+# Parameters for train_and_eval_dataset:
+# ==============================================================================
+train_and_eval_dataset.eval_holdout_size = 0.05
+train_and_eval_dataset.eval_shuffle_files = True
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_concat_preprocess
+wmt_concat_preprocess.max_length = 255
+wmt_concat_preprocess.max_eval_length = 255
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.d_model = 512
+TransformerLM.d_ff = 2048
+TransformerLM.dropout = 0.1
+TransformerLM.max_len = 2048
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 8
+TransformerLM.n_layers = 6
+TransformerLM.vocab_size = 33300
+
+# Parameters for OnlineTuneEnv:
+# ==============================================================================
+OnlineTuneEnv.inputs = @trax.inputs.inputs
+OnlineTuneEnv.model = @trax.models.TransformerLM
+OnlineTuneEnv.optimizer = @trax.optimizers.Adafactor
+OnlineTuneEnv.control_configs = (
+    ("learning_rate", 1e-3, (1e-9, 1e-2), False),
+    ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
+
+    ("dropout_embedding", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_initial", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_initial", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_initial", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_middle", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_middle", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_middle", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_final", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_final", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_final", 0.1, (0.0, 0.9), True),
+)
+OnlineTuneEnv.nontrainable_param_map = {
+    # "dropout_{layer_type}_{block_index}": "dropout_{layer_type}_{block_group}"
+
+    "dropout_attention_0": "dropout_attention_initial",
+    "dropout_ff_middle_0": "dropout_ff_middle_initial",
+    "dropout_ff_final_0": "dropout_ff_final_initial",
+
+    "dropout_attention_1": "dropout_attention_middle",
+    "dropout_ff_middle_1": "dropout_ff_middle_middle",
+    "dropout_ff_final_1": "dropout_ff_final_middle",
+    "dropout_attention_2": "dropout_attention_middle",
+    "dropout_ff_middle_2": "dropout_ff_middle_middle",
+    "dropout_ff_final_2": "dropout_ff_final_middle",
+    "dropout_attention_3": "dropout_attention_middle",
+    "dropout_ff_middle_3": "dropout_ff_middle_middle",
+    "dropout_ff_final_3": "dropout_ff_final_middle",
+    "dropout_attention_4": "dropout_attention_middle",
+    "dropout_ff_middle_4": "dropout_ff_middle_middle",
+    "dropout_ff_final_4": "dropout_ff_final_middle",
+
+    "dropout_attention_5": "dropout_attention_final",
+    "dropout_ff_middle_5": "dropout_ff_middle_final",
+    "dropout_ff_final_5": "dropout_ff_final_final",
+}
+OnlineTuneEnv.include_controls_in_observation = True
+OnlineTuneEnv.train_steps = 500
+OnlineTuneEnv.eval_steps = 1
+OnlineTuneEnv.env_steps = 100
+OnlineTuneEnv.metric_range = (0.0, 10.0)
+OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
+OnlineTuneEnv.has_weights = True
diff --git a/tensor2tensor/trax/rl/configs/ppo_online_tune.gin b/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
index eade5f618..5b57f4fd2 100644
--- a/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
+++ b/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
@@ -28,7 +28,7 @@ PPO.random_seed = None
 PPO.gamma = 1.0
 PPO.lambda_ = 0.95
 PPO.c1 = 1.0
-PPO.c2 = 0.01
+PPO.c2 = 0.03
 PPO.done_frac_for_policy_save = 0
 PPO.len_history_for_policy = None
 PPO.separate_eval = False
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index 76065c535..b13cf1c32 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -75,7 +75,8 @@ def __init__(self,
                metric_range=(0.0, 5.0),
                # Don't save checkpoints by default, as they tend to use a lot of
                # space.
-               should_save_checkpoints=False):
+               should_save_checkpoints=False,
+               has_weights=False):
     if action_multipliers is None:
       action_multipliers = self.DEFAULT_ACTION_MULTIPLIERS
     self._model = model
@@ -91,6 +92,7 @@ def __init__(self,
         inputs=inputs,
         should_save=should_save_checkpoints,
         nontrainable_param_map=nontrainable_param_map,
+        has_weights=has_weights,
     )
     self._trainer = None
     self._action_multipliers = action_multipliers

From fe03eb5e94409b4efe84a0126ef5d70862ea0301 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 19 Sep 2019 17:05:47 -0700
Subject: [PATCH 2468/2720] Optimize hashed attention

PiperOrigin-RevId: 270158366
---
 .../trax/configs/transformer_imagenet64.gin   |  10 +-
 tensor2tensor/trax/layers/__init__.py         |   2 +
 tensor2tensor/trax/layers/attention.py        | 307 ++++++++++++++++++
 3 files changed, 318 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/configs/transformer_imagenet64.gin b/tensor2tensor/trax/configs/transformer_imagenet64.gin
index bfd8246ee..341e7ac54 100644
--- a/tensor2tensor/trax/configs/transformer_imagenet64.gin
+++ b/tensor2tensor/trax/configs/transformer_imagenet64.gin
@@ -54,9 +54,17 @@ MergedMultiHashedCausalAttention.one_rng = False
 MergedMultiHashedCausalAttention.drop_for_hash_rate = 0.1
 MergedMultiHashedCausalAttention.hard_k = 32
 
+# Parameters for MergedMultiHashedCausalAttentionV2:
+# ==============================================================================
+MergedMultiHashedCausalAttentionV2.dropout = 0.0
+MergedMultiHashedCausalAttentionV2.n_bins = 96
+MergedMultiHashedCausalAttentionV2.n_buckets = 256
+MergedMultiHashedCausalAttentionV2.n_hashes = 4
+MergedMultiHashedCausalAttentionV2.one_rng = True
+
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.attention_type = @trax.layers.MergedMultiHashedCausalAttention
+TransformerLM.attention_type = @trax.layers.MergedMultiHashedCausalAttentionV2
 TransformerLM.d_attention_key = 64
 TransformerLM.d_attention_value = 64
 TransformerLM.d_model = 1024
diff --git a/tensor2tensor/trax/layers/__init__.py b/tensor2tensor/trax/layers/__init__.py
index 427b18d5a..6e7bbc124 100644
--- a/tensor2tensor/trax/layers/__init__.py
+++ b/tensor2tensor/trax/layers/__init__.py
@@ -59,3 +59,5 @@ def layer_configure(*args, **kwargs):
     MergedHashedCausalAttention, blacklist=["mode"])
 MergedMultiHashedCausalAttention = layer_configure(
     MergedMultiHashedCausalAttention, blacklist=["mode"])
+MergedMultiHashedCausalAttentionV2 = layer_configure(
+    MergedMultiHashedCausalAttentionV2, blacklist=["mode"])
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index e25dc4198..d59504a76 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -1069,6 +1069,313 @@ def _do_call(x):
     return output, vjpfun(ct)[0]
 
 
+class MergedMultiHashedCausalAttentionV2(BaseCausalAttention):
+  """Hash-based causal attention, with multiple hashes (faster version)."""
+
+  def __init__(self, dropout, mode, n_bins=64, n_hashes=1, n_buckets=64,
+               one_rng=False, allow_duplicate_attention=False,
+               attend_across_buckets=False):
+    del dropout, mode
+    super(MergedMultiHashedCausalAttentionV2, self).__init__()
+    assert n_buckets >= n_bins, 'This setting is not recommended: too few bins.'
+    self.n_bins = n_bins
+    self.n_hashes = n_hashes
+    self.n_buckets = n_buckets
+
+    self._one_rng = one_rng
+    self._prng = None
+    if one_rng:
+      seed = random.randint(0, 2**31 - 1)
+      self._prng = backend.random.get_prng(seed)
+
+    self._allow_duplicate_attention = allow_duplicate_attention
+    self._attend_across_buckets = attend_across_buckets
+
+  def call(self, inputs, params=(), state=(), rng=None, **kwargs):
+    del params, kwargs
+    output, _ = self.batch_call_and_or_grad(inputs[0], inputs[2], rng=rng)
+    return output, state
+
+  def call_and_grad(self, inputs, ct, rng=None, **kwargs):
+    del kwargs
+    output, (qk_ct, v_ct) = self.batch_call_and_or_grad(
+        inputs[0], inputs[2], ct=ct, rng=rng)
+    return output, (qk_ct, np.zeros_like(inputs[1]), v_ct)
+
+  def has_custom_grad(self):
+    return True
+
+  def custom_grad(self, inputs, output, ct, params=(), state=(), rng=None,
+                  **kwargs):
+    del output, params, state
+    _, (qk_ct, v_ct) = self.batch_call_and_or_grad(
+        inputs[0], inputs[2], return_output=False, ct=ct, rng=rng)
+    inputs_ct = (qk_ct, np.zeros_like(inputs[1]), v_ct)
+    return inputs_ct, ()
+
+  def batch_call_and_or_grad(self, qk, v, ct=None, return_output=True,
+                             rng=None):
+    assert return_output or ct is not None, 'No work to perform!'
+
+    # The approach here is to perform attention for one batch element and head
+    # at a time. Note that there is absolutely no interaction across examples or
+    # heads: this layer has no parameters, and hashing patterns are also
+    # different across examples/heads. As a result, batching doesn't give any
+    # performance gains except in the case of accelerator under-utilization. We
+    # assume that hash-based attention will be applied primarily to long
+    # sequences, where unbatched attention for a single head has sufficient
+    # computation to fill up the accelerator.
+
+    batch_loop_idx = np.zeros((), dtype=np.int32)
+    batch_loop_max = qk.shape[0]
+
+    init_vals = (batch_loop_idx,)
+    if return_output:
+      out_accum = np.zeros_like(qk)
+      init_vals = init_vals + (out_accum,)
+    if ct is not None:
+      qk_ct_accum = np.zeros_like(qk)
+      v_ct_accum = np.zeros_like(v)
+      init_vals = init_vals + (qk_ct_accum, v_ct_accum)
+
+    def cond_fun(vals):
+      batch_loop_idx = vals[0]
+      return jax.lax.lt(batch_loop_idx, batch_loop_max)
+
+    def body_fun(vals):
+      """Performs attention for a single batch element and head."""
+      batch_loop_idx = vals[0]
+      if self._prng is None:
+        hash_rng = jax.random.fold_in(rng, batch_loop_idx)
+      else:
+        # TODO(kitaev): Maybe use the same RNG across examples (but not heads)?
+        hash_rng = jax.random.fold_in(self._prng, batch_loop_idx)
+      qk_slice = jax.lax.dynamic_index_in_dim(
+          qk, batch_loop_idx, axis=0, keepdims=False)
+      v_slice = jax.lax.dynamic_index_in_dim(
+          v, batch_loop_idx, axis=0, keepdims=False)
+
+      if ct is None:
+        out_slice = self.single_call(qk_slice, v_slice, hash_rng=hash_rng)
+      else:
+        def _do_single_call(qk_slice, v_slice):
+          return self.single_call(qk_slice, v_slice, hash_rng=hash_rng)
+        ct_slice = jax.lax.dynamic_index_in_dim(
+            ct, batch_loop_idx, axis=0, keepdims=False)
+        out_slice, vjpfun = jax.vjp(_do_single_call, qk_slice, v_slice)
+        qk_ct_slice, v_ct_slice = vjpfun(ct_slice)
+
+      new_vals = (batch_loop_idx + 1,)
+      if return_output:
+        out_accum = vals[1]
+        out_accum = jax.lax.dynamic_update_index_in_dim(
+            out_accum, out_slice, batch_loop_idx, axis=0)
+        new_vals = new_vals + (out_accum,)
+      if ct is not None:
+        qk_ct_accum, v_ct_accum = vals[-2:]
+        qk_ct_accum = jax.lax.dynamic_update_index_in_dim(
+            qk_ct_accum, qk_ct_slice, batch_loop_idx, axis=0)
+        v_ct_accum = jax.lax.dynamic_update_index_in_dim(
+            v_ct_accum, v_ct_slice, batch_loop_idx, axis=0)
+        new_vals = new_vals + (qk_ct_accum, v_ct_accum)
+
+      return new_vals
+
+    final_vals = jax.lax.while_loop(cond_fun, body_fun, init_vals)
+
+    if return_output:
+      out = final_vals[1]
+    else:
+      out = None
+
+    if ct is not None:
+      input_ct = final_vals[-2:]
+    else:
+      input_ct = None
+
+    return out, input_ct
+
+  def make_unit_length(self, x, epsilon=1e-6):
+    variance = np.mean(x**2, axis=-1, keepdims=True)
+    norm_inputs = x / np.sqrt(variance + epsilon)
+    return norm_inputs
+
+  def hash_vectors(self, vecs, rng):
+    # See https://arxiv.org/pdf/1509.02897.pdf
+    # We sample a different random rotation for each round of hashing to
+    # decrease the probability of hash misses.
+    assert self.n_buckets % 2 == 0
+    random_rotations = jax.random.normal(
+        jax.lax.tie_in(vecs, rng),
+        (vecs.shape[-1], self.n_hashes, self.n_buckets // 2)).astype('float32')
+
+    rotated_vecs = np.einsum('tf,fhb->htb', vecs, random_rotations)
+    rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)
+    buckets = np.argmax(rotated_vecs, axis=-1)
+    # buckets is now (self.n_hashes, seqlen). Next we add offsets so that bucket
+    # numbers from different hashing rounds don't overlap.
+    offsets = jax.lax.tie_in(buckets, np.arange(self.n_hashes))
+    offsets = np.reshape(offsets * self.n_buckets, (-1, 1))
+    buckets = np.reshape(buckets + offsets, (-1,))
+    return buckets
+
+  def single_call(self, qk, v, hash_rng=None):
+    # We use the same vector as both a query and a key.
+    seqlen = qk.shape[-2]
+
+    buckets = self.hash_vectors(qk, rng=hash_rng)
+    assert int(buckets.shape[0]) == self.n_hashes * seqlen
+
+    ticker = jax.lax.tie_in(qk, np.arange(self.n_hashes * seqlen))
+    buckets_and_t = seqlen * buckets + (ticker % seqlen)
+    buckets_and_t = jax.lax.stop_gradient(buckets_and_t)
+
+    # Hash-based sort ("s" at the start of variable names means "sorted")
+    sbuckets_and_t, sticker = jax.lax.sort_key_val(
+        buckets_and_t, ticker, dimension=-1)
+    _, undo_sort = jax.lax.sort_key_val(sticker, ticker, dimension=-1)
+    sbuckets_and_t = jax.lax.stop_gradient(sbuckets_and_t)
+    sticker = jax.lax.stop_gradient(sticker)
+    undo_sort = jax.lax.stop_gradient(undo_sort)
+
+    st = (sticker % seqlen)
+    sqk = np.take(qk, st, axis=0)
+    sv = np.take(v, st, axis=0)
+
+    # Split off a "bin" axis so that attention only occurs within chunks.
+    bq_t = bkv_t = np.reshape(st, (self.n_hashes * self.n_bins, -1))
+    bqk = np.reshape(sqk, (self.n_hashes * self.n_bins, -1, sqk.shape[-1]))
+    bv = np.reshape(sv, (self.n_hashes * self.n_bins, -1, sv.shape[-1]))
+    bq_buckets = bkv_buckets = np.reshape(
+        sbuckets_and_t // seqlen, (self.n_hashes * self.n_bins, -1))
+
+    # Hashing operates on unit-length vectors. Unnormalized query vectors are
+    # fine because they effectively provide a learnable temperature for the
+    # attention softmax, but normalizing keys is needed so that similarity for
+    # the purposes of attention correctly corresponds to hash locality.
+    bq = bqk
+    bk = self.make_unit_length(bqk)
+
+    # Allow each chunk to attend within itself, and also one chunk back. Chunk
+    # boundaries might occur in the middle of a sequence of items from the
+    # same bucket, so this increases the chances of attending to relevant items.
+    # TODO(kitaev): benchmark whether XLA pad operation is noticeably faster.
+    def look_one_back(x):
+      if len(x.shape) == 2:
+        x_extra = np.concatenate([x[-1:, :], x[:-1, :]], axis=0)
+      else:
+        x_extra = np.concatenate([x[-1:, :, :], x[:-1, :, :]], axis=0)
+      return np.concatenate([x, x_extra], axis=1)
+
+    bk = look_one_back(bk)
+    bv = look_one_back(bv)
+    bkv_t = look_one_back(bkv_t)
+    bkv_buckets = look_one_back(bkv_buckets)
+
+    # Dot-product attention.
+    dots = np.matmul(bq, np.swapaxes(bk, -1, -2)) / np.sqrt(bq.shape[-1])
+
+    # Causal masking
+    mask = jax.lax.convert_element_type(
+        jax.lax.lt(bq_t[:, :, None], bkv_t[:, None, :]),
+        np.float32)
+    dots = dots - 1e9 * mask
+
+    # Mask out attention to self except when no other targets are available.
+    self_mask = jax.lax.convert_element_type(
+        jax.lax.eq(bq_t[:, :, None], bkv_t[:, None, :]),
+        np.float32)
+    dots = dots - 1e6 * self_mask
+
+    # Mask out attention to other hash buckets.
+    if not self._attend_across_buckets:
+      bucket_mask = jax.lax.convert_element_type(
+          jax.lax.ne(bq_buckets[:, :, None], bkv_buckets[:, None, :]),
+          np.float32)
+      dots = dots - 1e5 * bucket_mask
+
+    # Don't double-count query-key pairs across multiple rounds of hashing.
+    # The approach here is to count how many times a query-key pair is repeated,
+    # and to lower its log-prob correspondingly at each repetition.
+    if not self._allow_duplicate_attention:
+      locs1 = undo_sort // bq_t.shape[-1]
+      locs2 = (locs1 + 1) % (self.n_hashes * self.n_bins)
+      if not self._attend_across_buckets:
+        locs1 = buckets * (self.n_hashes * self.n_bins) + locs1
+        locs2 = buckets * (self.n_hashes * self.n_bins) + locs2
+      locs = np.moveaxis(np.concatenate([
+          np.reshape(locs1, (self.n_hashes, seqlen)),
+          np.reshape(locs2, (self.n_hashes, seqlen)),
+      ], 0), 0, -1)  # produces shape (seqlen, 2 * self.n_hashes)
+      slocs = np.take(locs, st, axis=0)
+      b_locs = np.reshape(
+          slocs, (self.n_hashes * self.n_bins, -1, 2 * self.n_hashes))
+      # Queries always use the primary location (based on loc1).
+      bq_locs = np.broadcast_to(
+          b_locs[:, :, None, :self.n_hashes],
+          b_locs.shape[:2] + (2, self.n_hashes))
+      bq_locs = np.reshape(bq_locs, b_locs.shape)
+      bkv_locs = look_one_back(b_locs)
+
+      dup_counts = np.sum(
+          jax.lax.convert_element_type(
+              jax.lax.eq(bq_locs[:, :, None, :], bkv_locs[:, None, :, :]),
+              np.float32),
+          axis=-1)
+      assert dup_counts.shape == dots.shape
+      dots = dots - jax.lax.stop_gradient(np.log(dup_counts + 1e-9))
+
+    # Softmax.
+    dots_logsumexp = backend.logsumexp(dots, axis=-1, keepdims=True)
+    dots = np.exp(dots - dots_logsumexp)
+
+    bo = np.matmul(dots, bv)
+    so = np.reshape(bo, (-1, bo.shape[-1]))
+    slogits = np.reshape(dots_logsumexp, (-1,))
+
+    def unsort_for_output_impl(so, slogits):
+      o = np.take(so, undo_sort, axis=0)
+      # Sorting is considerably faster than gather, but first we need to get the
+      # XLA compiler to abandon the idea of fusing this sort with the input sort
+      # (which introduces a computation cycle and leads to a crash).
+      # TODO(kitaev): remove "sticker_" variable if XLA is fixed.
+      sticker_ = sticker + jax.lax.convert_element_type(
+          slogits[0] > 0, sticker.dtype)
+      _, logits = jax.lax.sort_key_val(sticker_, slogits, dimension=-1)
+      return o, logits
+
+    def unsort_for_output_vjp(so, slogits):
+      """Custom gradient for unsort_for_output."""
+      so = jax.lax.stop_gradient(so)
+      slogits = jax.lax.stop_gradient(slogits)
+      o, logits = unsort_for_output_impl(so, slogits)
+      def vjpfun(o_logits_grads):
+        so_grad = np.take(o_logits_grads[0], sticker, axis=0)
+        # TODO(kitaev): this exists to match the forward pass, but I'm not sure
+        # if it's actually required.
+        buckets_and_t_ = buckets_and_t + jax.lax.convert_element_type(
+            o_logits_grads[1][0] > 0, buckets_and_t.dtype)
+        _, slogits_grad = jax.lax.sort_key_val(
+            buckets_and_t_, o_logits_grads[1], dimension=-1)
+        return (so_grad, slogits_grad)
+      return (o, logits), vjpfun
+
+    unsort_for_output = jax.custom_transforms(unsort_for_output_impl)
+    jax.defvjp_all(unsort_for_output, unsort_for_output_vjp)
+    o, logits = unsort_for_output_impl(so, slogits)
+
+    if self.n_hashes == 1:
+      out = o
+    else:
+      o = np.reshape(o, (self.n_hashes, seqlen, o.shape[-1]))
+      logits = np.reshape(logits, (self.n_hashes, seqlen, 1))
+      probs = np.exp(logits - backend.logsumexp(logits, axis=0, keepdims=True))
+      out = np.sum(o * probs, axis=0)
+
+    assert out.shape == v.shape
+    return out
+
+
 def CausalAttention(d_feature, n_heads=1,
                     d_attention_key=None, d_attention_value=None,
                     attention_type=DotProductCausalAttention,

From adb9f391e1be6fe56f68706d0e622b5fabec5430 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 19 Sep 2019 17:44:06 -0700
Subject: [PATCH 2469/2720] Add hard_k hparam to multi-round hashed attention

PiperOrigin-RevId: 270164357
---
 tensor2tensor/trax/layers/attention.py | 53 +++++++++++++++++++++++---
 1 file changed, 47 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index d59504a76..9d60f5956 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -1074,7 +1074,7 @@ class MergedMultiHashedCausalAttentionV2(BaseCausalAttention):
 
   def __init__(self, dropout, mode, n_bins=64, n_hashes=1, n_buckets=64,
                one_rng=False, allow_duplicate_attention=False,
-               attend_across_buckets=False):
+               attend_across_buckets=False, hard_k=0):
     del dropout, mode
     super(MergedMultiHashedCausalAttentionV2, self).__init__()
     assert n_buckets >= n_bins, 'This setting is not recommended: too few bins.'
@@ -1090,6 +1090,7 @@ def __init__(self, dropout, mode, n_bins=64, n_hashes=1, n_buckets=64,
 
     self._allow_duplicate_attention = allow_duplicate_attention
     self._attend_across_buckets = attend_across_buckets
+    self._hard_k = hard_k
 
   def call(self, inputs, params=(), state=(), rng=None, **kwargs):
     del params, kwargs
@@ -1295,8 +1296,11 @@ def look_one_back(x):
       dots = dots - 1e5 * bucket_mask
 
     # Don't double-count query-key pairs across multiple rounds of hashing.
-    # The approach here is to count how many times a query-key pair is repeated,
-    # and to lower its log-prob correspondingly at each repetition.
+    # There are two possible strategies here. (1) The default is to count how
+    # many times a query-key pair is repeated, and to lower its log-prob
+    # correspondingly at each repetition. (2) When hard_k is set, the code
+    # instead masks all but the first occurence of each query-key pair.
+    # TODO(kitaev): is one strategy faster or more numerically stable?
     if not self._allow_duplicate_attention:
       locs1 = undo_sort // bq_t.shape[-1]
       locs2 = (locs1 + 1) % (self.n_hashes * self.n_bins)
@@ -1310,9 +1314,19 @@ def look_one_back(x):
       slocs = np.take(locs, st, axis=0)
       b_locs = np.reshape(
           slocs, (self.n_hashes * self.n_bins, -1, 2 * self.n_hashes))
-      # Queries always use the primary location (based on loc1).
+      # Queries always use the primary location (based on locs1).
+      b_locs1 = b_locs[:, :, None, :self.n_hashes]
+      if self._hard_k > 0:
+        range_n_hashes = jax.lax.tie_in(b_locs, np.arange(self.n_hashes))
+        nouse_locs = (range_n_hashes[:, None] > range_n_hashes[None, :])
+        nouse_locs = 2 * nouse_locs - 1  # 1 = use, -1 = don't use
+        nouse_locs = np.reshape(
+            np.broadcast_to(nouse_locs[:, None, :],
+                            (self.n_hashes, self.n_bins, self.n_hashes)),
+            (self.n_hashes * self.n_bins, 1, 1, self.n_hashes))
+        b_locs1 = b_locs1 * nouse_locs
       bq_locs = np.broadcast_to(
-          b_locs[:, :, None, :self.n_hashes],
+          b_locs1,
           b_locs.shape[:2] + (2, self.n_hashes))
       bq_locs = np.reshape(bq_locs, b_locs.shape)
       bkv_locs = look_one_back(b_locs)
@@ -1323,7 +1337,34 @@ def look_one_back(x):
               np.float32),
           axis=-1)
       assert dup_counts.shape == dots.shape
-      dots = dots - jax.lax.stop_gradient(np.log(dup_counts + 1e-9))
+      if self._hard_k > 0:
+        dots = dots - 1e5 * jax.lax.stop_gradient(dup_counts)
+      else:
+        dots = dots - jax.lax.stop_gradient(np.log(dup_counts + 1e-9))
+
+    # Each query only attends to the top k most relevant keys.
+    if self._hard_k > 0:
+      b_top_dots = np.sort(dots)[..., -self._hard_k:]  # Get the top k dots.
+      b_top_dots = jax.lax.stop_gradient(b_top_dots)
+      s_top_dots = np.reshape(b_top_dots, (-1, self._hard_k))
+      top_dots = np.take(s_top_dots, undo_sort, axis=0)
+
+      merged_top_dots = np.moveaxis(
+          np.reshape(top_dots, (self.n_hashes, seqlen, self._hard_k)), 0, -1)
+      merged_top_dots = np.reshape(merged_top_dots, (seqlen, -1))
+
+      dots_thresh = np.sort(merged_top_dots)[:, -self._hard_k]
+      # It's possible to compute the partition function at this point, but right
+      # now this codepath isn't set up for backprop, and there might also be
+      # issues computing it this way if two dot-products are exactly equal.
+
+      sdots_thresh = dots_thresh[st]
+      bdots_thresh = np.reshape(sdots_thresh, (self.n_hashes * self.n_bins, -1))
+      bdots_thresh = jax.lax.stop_gradient(bdots_thresh)
+
+      top_k_mask = jax.lax.convert_element_type(
+          dots < bdots_thresh[..., None], np.float32)
+      dots = dots - 1e5 * jax.lax.stop_gradient(top_k_mask)
 
     # Softmax.
     dots_logsumexp = backend.logsumexp(dots, axis=-1, keepdims=True)

From 369854260d7970276fac6fd27293009a84c571b8 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 19 Sep 2019 21:33:32 -0700
Subject: [PATCH 2470/2720] Hashed attention that uses top-k buckets

PiperOrigin-RevId: 270192777
---
 tensor2tensor/trax/layers/attention.py | 41 +++++++++++++++++++-------
 1 file changed, 31 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 9d60f5956..cb4514722 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -1074,10 +1074,13 @@ class MergedMultiHashedCausalAttentionV2(BaseCausalAttention):
 
   def __init__(self, dropout, mode, n_bins=64, n_hashes=1, n_buckets=64,
                one_rng=False, allow_duplicate_attention=False,
-               attend_across_buckets=False, hard_k=0):
+               attend_across_buckets=False, hard_k=0, rehash_each_round=True):
     del dropout, mode
     super(MergedMultiHashedCausalAttentionV2, self).__init__()
     assert n_buckets >= n_bins, 'This setting is not recommended: too few bins.'
+    assert rehash_each_round or allow_duplicate_attention, (
+        'The setting {allow_duplicate_attention=False, rehash_each_round=False}'
+        ' is not implemented.')
     self.n_bins = n_bins
     self.n_hashes = n_hashes
     self.n_buckets = n_buckets
@@ -1091,6 +1094,7 @@ def __init__(self, dropout, mode, n_bins=64, n_hashes=1, n_buckets=64,
     self._allow_duplicate_attention = allow_duplicate_attention
     self._attend_across_buckets = attend_across_buckets
     self._hard_k = hard_k
+    self._rehash_each_round = rehash_each_round
 
   def call(self, inputs, params=(), state=(), rng=None, **kwargs):
     del params, kwargs
@@ -1206,18 +1210,35 @@ def hash_vectors(self, vecs, rng):
     # We sample a different random rotation for each round of hashing to
     # decrease the probability of hash misses.
     assert self.n_buckets % 2 == 0
-    random_rotations = jax.random.normal(
-        jax.lax.tie_in(vecs, rng),
-        (vecs.shape[-1], self.n_hashes, self.n_buckets // 2)).astype('float32')
+    random_rotations_shape = (
+        vecs.shape[-1],
+        self.n_hashes if self._rehash_each_round else 1,
+        self.n_buckets // 2)
 
+    random_rotations = jax.random.normal(
+        jax.lax.tie_in(vecs, rng), random_rotations_shape).astype('float32')
     rotated_vecs = np.einsum('tf,fhb->htb', vecs, random_rotations)
     rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)
-    buckets = np.argmax(rotated_vecs, axis=-1)
-    # buckets is now (self.n_hashes, seqlen). Next we add offsets so that bucket
-    # numbers from different hashing rounds don't overlap.
-    offsets = jax.lax.tie_in(buckets, np.arange(self.n_hashes))
-    offsets = np.reshape(offsets * self.n_buckets, (-1, 1))
-    buckets = np.reshape(buckets + offsets, (-1,))
+
+    if self._rehash_each_round:
+      buckets = np.argmax(rotated_vecs, axis=-1)
+      # buckets is now (self.n_hashes, seqlen). Next we add offsets so that
+      # bucket numbers from different hashing rounds don't overlap.
+      offsets = jax.lax.tie_in(buckets, np.arange(self.n_hashes))
+      offsets = np.reshape(offsets * self.n_buckets, (-1, 1))
+      buckets = np.reshape(buckets + offsets, (-1,))
+    else:
+      # In this configuration, we map each item to the top self.n_hashes buckets
+      rotated_vecs = np.squeeze(rotated_vecs, 0)
+      bucket_range = jax.lax.tie_in(vecs, np.arange(rotated_vecs.shape[-1]))
+      bucket_range = np.reshape(bucket_range, (1, -1))
+      bucket_range = np.broadcast_to(bucket_range, rotated_vecs.shape)
+
+      _, buckets = jax.lax.sort_key_val(
+          rotated_vecs, bucket_range, dimension=-1)
+      buckets = buckets[:, -self.n_hashes:]
+      buckets = np.reshape(np.moveaxis(buckets, 0, -1), (-1,))
+
     return buckets
 
   def single_call(self, qk, v, hash_rng=None):

From 0bc60f10e6b1c6fcc4f8772921fa3adbd7d5a979 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 20 Sep 2019 12:56:31 -0700
Subject: [PATCH 2471/2720] Add support for hash dropout in V2 hashed
 attention.

PiperOrigin-RevId: 270331326
---
 tensor2tensor/trax/layers/attention.py | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index cb4514722..535b32028 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -1074,8 +1074,10 @@ class MergedMultiHashedCausalAttentionV2(BaseCausalAttention):
 
   def __init__(self, dropout, mode, n_bins=64, n_hashes=1, n_buckets=64,
                one_rng=False, allow_duplicate_attention=False,
-               attend_across_buckets=False, hard_k=0, rehash_each_round=True):
-    del dropout, mode
+               attend_across_buckets=False, hard_k=0,
+               rehash_each_round=True, drop_for_hash_rate=0.0):
+    del dropout
+    self._mode = mode
     super(MergedMultiHashedCausalAttentionV2, self).__init__()
     assert n_buckets >= n_bins, 'This setting is not recommended: too few bins.'
     assert rehash_each_round or allow_duplicate_attention, (
@@ -1084,7 +1086,7 @@ def __init__(self, dropout, mode, n_bins=64, n_hashes=1, n_buckets=64,
     self.n_bins = n_bins
     self.n_hashes = n_hashes
     self.n_buckets = n_buckets
-
+    self._drop_for_hash_rate = drop_for_hash_rate
     self._one_rng = one_rng
     self._prng = None
     if one_rng:
@@ -1205,6 +1207,13 @@ def make_unit_length(self, x, epsilon=1e-6):
     norm_inputs = x / np.sqrt(variance + epsilon)
     return norm_inputs
 
+  def drop_for_hash(self, x, rng):
+    rate = self._drop_for_hash_rate
+    if self._mode == 'train' and rate > 0.0:
+      keep = backend.random.bernoulli(rng, 1.0 - rate, x.shape)
+      return np.where(keep, x / (1.0 - rate), np.zeros_like(x))
+    return x
+
   def hash_vectors(self, vecs, rng):
     # See https://arxiv.org/pdf/1509.02897.pdf
     # We sample a different random rotation for each round of hashing to
@@ -1215,9 +1224,14 @@ def hash_vectors(self, vecs, rng):
         self.n_hashes if self._rehash_each_round else 1,
         self.n_buckets // 2)
 
+    rng = jax.lax.tie_in(vecs, rng)
+    rng, subrng = backend.random.split(rng)
     random_rotations = jax.random.normal(
-        jax.lax.tie_in(vecs, rng), random_rotations_shape).astype('float32')
-    rotated_vecs = np.einsum('tf,fhb->htb', vecs, random_rotations)
+        rng, random_rotations_shape).astype('float32')
+    # TODO(lukaszkaiser): the dropout mask will be used for all rounds of
+    # hashing, so it's shared between them. Check if that's what we want.
+    dropped_vecs = self.drop_for_hash(vecs, subrng)
+    rotated_vecs = np.einsum('tf,fhb->htb', dropped_vecs, random_rotations)
     rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)
 
     if self._rehash_each_round:

From fb98ddec7c67768bf474dbaa7ddf7bf6611e6735 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 20 Sep 2019 14:28:51 -0700
Subject: [PATCH 2472/2720] New base config for reformer experiments

PiperOrigin-RevId: 270350727
---
 .../trax/configs/reformer_imagenet64.gin      | 127 ++++++++++++++++++
 .../trax/configs/transformer_enwik8.gin       |   3 +-
 .../trax/configs/transformer_imagenet64.gin   |   3 +-
 .../configs/transformer_revnet_enwik8.gin     |   1 +
 tensor2tensor/trax/layers/attention.py        |   6 +-
 .../models/research/transformer_revnet.py     |  12 +-
 tensor2tensor/trax/models/transformer.py      |  18 +--
 7 files changed, 150 insertions(+), 20 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/reformer_imagenet64.gin

diff --git a/tensor2tensor/trax/configs/reformer_imagenet64.gin b/tensor2tensor/trax/configs/reformer_imagenet64.gin
new file mode 100644
index 000000000..2653cf01c
--- /dev/null
+++ b/tensor2tensor/trax/configs/reformer_imagenet64.gin
@@ -0,0 +1,127 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters that will vary between experiments:
+# ==============================================================================
+train.model = @trax.models.TransformerLM
+
+TransformerLM.n_layers = 3
+TransformerRevnetLM.n_layers = 3
+
+# Attention types:
+# MemoryEfficientCausalAttention: full attention
+# MergedHashedCausalAttention: timebin
+# MergedMultiHashedCausalAttentionV2: hashbin
+TransformerLM.attention_type = @MergedMultiHashedCausalAttentionV2
+TransformerRevnetLM.attention_type = @MergedMultiHashedCausalAttentionV2
+
+# These three should all be equal, and MergedMultiHashedCausalAttentionV2
+# requires that share_qk be set to True.
+MemoryEfficientCausalAttention.share_qk = True
+TransformerLM.share_qk = True
+TransformerRevnetLM.share_qk = True
+
+# These four should all be equal
+TransformerLM.d_attention_key = 64
+TransformerLM.d_attention_value = 64
+TransformerRevnetLM.d_attention_key = 64
+TransformerRevnetLM.d_attention_value = 64
+
+# Use MergedHashedCausalAttention for timebin
+MergedHashedCausalAttention.n_bins = 64
+
+# Use MergedMultiHashedCausalAttentionV2 for hashbin
+MergedMultiHashedCausalAttentionV2.n_bins = 96
+MergedMultiHashedCausalAttentionV2.n_buckets = 192  # Always 2 * n_bins
+MergedMultiHashedCausalAttentionV2.n_hashes = 2
+MergedMultiHashedCausalAttentionV2.drop_for_hash_rate = 0.0
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 1
+batch_fun.eval_batch_size = 8
+batch_fun.max_eval_length = 12288  # 64 * 64 * 3
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
+inputs.input_name = 'targets'
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 2.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 500
+train.eval_steps = 64
+train.inputs = @trax.inputs.inputs
+# train.model: see top
+train.optimizer = @trax.optimizers.Adafactor
+train.train_steps = 500000
+train.trainer_class = @MemoryEfficientTrainer
+train.save_steps = \
+    [1000, 5000, 10000, 20000, 40000, 60000, 80000,
+     100000, 200000, 300000, 400000, 500000]
+
+# Parameters for MemoryEfficientCausalAttention:
+# ==============================================================================
+MemoryEfficientCausalAttention.dropout = 0.0
+MemoryEfficientCausalAttention.loop_stride = 512
+# MemoryEfficientCausalAttention.share_qk: see top
+
+# Parameters for MergedHashedCausalAttention:
+# ==============================================================================
+MergedHashedCausalAttention.dropout = 0.0
+# MergedHashedCausalAttention.n_bins: see top
+MergedHashedCausalAttention.bin_by_time = True
+MergedHashedCausalAttention.one_rng = False
+
+# Parameters for MergedMultiHashedCausalAttentionV2:
+# ==============================================================================
+MergedMultiHashedCausalAttentionV2.allow_duplicate_attention = False
+MergedMultiHashedCausalAttentionV2.attend_across_buckets = False
+MergedMultiHashedCausalAttentionV2.rehash_each_round = True
+# MergedMultiHashedCausalAttentionV2.n_bins: see top
+# MergedMultiHashedCausalAttentionV2.n_buckets: see top
+# MergedMultiHashedCausalAttentionV2.n_hashes: see top
+MergedMultiHashedCausalAttentionV2.one_rng = False
+MergedMultiHashedCausalAttentionV2.hard_k = 0
+MergedMultiHashedCausalAttentionV2.dropout = 0.0
+# MergedMultiHashedCausalAttentionV2.drop_for_hash_rate: see top
+
+# Parameters for TransformerLM:
+# ==============================================================================
+# TransformerLM.attention_type: see top
+# TransformerLM.d_attention_key: see top
+# TransformerLM.d_attention_value: see top
+TransformerLM.d_model = 1024
+TransformerLM.d_ff = 4096
+TransformerLM.dropout = 0.0
+TransformerLM.max_len = 12288  # 64 * 64 * 3
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 8
+# TransformerLM.n_layers: see top
+# TransformerLM.share_qk: see top
+TransformerLM.vocab_size = 256
+
+# Parameters for TransformerRevnetLM:
+# ==============================================================================
+# TransformerRevnetLM.attention_type: see top
+# TransformerRevnetLM.d_attention_key: see top
+# TransformerRevnetLM.d_attention_value: see top
+TransformerRevnetLM.d_model = 1024
+TransformerRevnetLM.d_ff = 4096
+TransformerRevnetLM.dropout = 0.0
+TransformerRevnetLM.max_len = 12288  # 64 * 64 * 3
+TransformerRevnetLM.mode = 'train'
+TransformerRevnetLM.n_heads = 8
+# TransformerRevnetLM.n_layers: see top
+TransformerRevnetLM.vocab_size = 256
+TransformerRevnetLM.n_chunks = 16
+TransformerRevnetLM.n_attention_chunks = 1
diff --git a/tensor2tensor/trax/configs/transformer_enwik8.gin b/tensor2tensor/trax/configs/transformer_enwik8.gin
index ed38015e1..ce1727082 100644
--- a/tensor2tensor/trax/configs/transformer_enwik8.gin
+++ b/tensor2tensor/trax/configs/transformer_enwik8.gin
@@ -35,6 +35,7 @@ train.trainer_class = @MemoryEfficientTrainer
 # ==============================================================================
 MemoryEfficientCausalAttention.dropout = 0.0
 MemoryEfficientCausalAttention.loop_stride = 256
+MemoryEfficientCausalAttention.share_qk = True  # matches TransformerLM.share_qk
 
 # Parameters for MergedHashedCausalAttention:
 # ==============================================================================
@@ -63,6 +64,6 @@ TransformerLM.max_len = 65536
 TransformerLM.mode = 'train'
 TransformerLM.n_heads = 4
 TransformerLM.n_layers = 3
-TransformerLM.share_kv = True
+TransformerLM.share_qk = True
 TransformerLM.vocab_size = 258  # Includes pad token and unused EOS token
 
diff --git a/tensor2tensor/trax/configs/transformer_imagenet64.gin b/tensor2tensor/trax/configs/transformer_imagenet64.gin
index 341e7ac54..63fd2f380 100644
--- a/tensor2tensor/trax/configs/transformer_imagenet64.gin
+++ b/tensor2tensor/trax/configs/transformer_imagenet64.gin
@@ -35,6 +35,7 @@ train.trainer_class = @MemoryEfficientTrainer
 # ==============================================================================
 MemoryEfficientCausalAttention.dropout = 0.0
 MemoryEfficientCausalAttention.loop_stride = 512
+MemoryEfficientCausalAttention.share_qk = True  # matches TransformerLM.share_qk
 
 # Parameters for MergedHashedCausalAttention:
 # ==============================================================================
@@ -74,6 +75,6 @@ TransformerLM.max_len = 12288  # 64 * 64 * 3
 TransformerLM.mode = 'train'
 TransformerLM.n_heads = 4
 TransformerLM.n_layers = 3
-TransformerLM.share_kv = True
+TransformerLM.share_qk = True
 TransformerLM.vocab_size = 256
 
diff --git a/tensor2tensor/trax/configs/transformer_revnet_enwik8.gin b/tensor2tensor/trax/configs/transformer_revnet_enwik8.gin
index 6bb673fb6..7ab04d905 100644
--- a/tensor2tensor/trax/configs/transformer_revnet_enwik8.gin
+++ b/tensor2tensor/trax/configs/transformer_revnet_enwik8.gin
@@ -53,4 +53,5 @@ TransformerRevnetLM.n_layers = 12
 TransformerRevnetLM.vocab_size = 258
 TransformerRevnetLM.n_chunks = 16
 TransformerRevnetLM.n_attention_chunks = 1
+TransformerRevnetLM.share_qk = True
 TransformerRevnetLM.attention_type = @trax.layers.MergedHashedCausalAttention
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 535b32028..84379fb3c 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -1455,7 +1455,7 @@ def vjpfun(o_logits_grads):
 def CausalAttention(d_feature, n_heads=1,
                     d_attention_key=None, d_attention_value=None,
                     attention_type=DotProductCausalAttention,
-                    share_kv=False, mode='train'):
+                    share_qk=False, mode='train'):
   """Transformer-style multi-headed causal attention.
 
   Args:
@@ -1466,7 +1466,7 @@ def CausalAttention(d_feature, n_heads=1,
     d_attention_value: int: depth of value vector for each attention head
         (default is d_feature // n_heads)
     attention_type: subclass of BaseCausalAttention: attention class to use
-    share_kv: bool, whether to share keys and values
+    share_qk: bool, whether to share queries and keys
     mode: str: 'train' or 'eval'
 
   Returns:
@@ -1479,7 +1479,7 @@ def CausalAttention(d_feature, n_heads=1,
     assert d_feature % n_heads == 0
     d_attention_value = d_feature // n_heads
 
-  if share_kv:
+  if share_qk:
     pre_attention = [
         cb.Dup(),
         cb.Parallel(
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 71cd89bbc..867341101 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -394,7 +394,7 @@ def call_post_attention2(params):
 
 def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
                  n_heads, n_attention_chunks, attention_type,
-                 dropout, share_kv, mode):
+                 dropout, share_qk, mode):
   """Reversible transformer decoder layer.
 
   Args:
@@ -406,13 +406,13 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
     n_attention_chunks: int: number of chunks for attention
     attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
-    share_kv: string, whether to share keys and values
+    share_qk: string, whether to share queries and keys
     mode: str: 'train' or 'eval'
 
   Returns:
     the layer.
   """
-  if share_kv:
+  if share_qk:
     pre_attention = [
         Chunk(n_sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
         tl.LayerNorm(),
@@ -468,7 +468,7 @@ def TransformerRevnetLM(vocab_size,
                         n_chunks=32,
                         n_attention_chunks=8,
                         attention_type=tl.DotProductCausalAttention,
-                        share_kv=False,
+                        share_qk=False,
                         mode='train'):
   """Reversible transformer language model (only uses a decoder, no encoder).
 
@@ -485,7 +485,7 @@ def TransformerRevnetLM(vocab_size,
     n_chunks: int: number of chunks (must match input pipeline)
     n_attention_chunks: int: number of chunks for attention
     attention_type: class: attention class to use, such as DotProductAttention.
-    share_kv: bool, whether to share keys and values.
+    share_qk: bool, whether to share queries and keys.
     mode: str: 'train' or 'eval'
 
   Returns:
@@ -506,7 +506,7 @@ def TransformerRevnetLM(vocab_size,
           DecoderBlock(d_model, d_ff,
                        d_attention_key, d_attention_value, n_heads,
                        n_attention_chunks, attention_type,
-                       dropout, share_kv, mode)
+                       dropout, share_qk, mode)
           for _ in range(n_layers)
       ] + [
           SplitForOutput(n_sections=n_chunks, axis=-2),  # pylint: disable=no-value-for-parameter
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index f40d558bc..17a94e7b8 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -115,7 +115,7 @@ def TransformerEncoder(vocab_size,
 
 
 def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-                 attention_type, dropout, share_kv, layer_idx, mode):
+                 attention_type, dropout, share_qk, layer_idx, mode):
   """Returns a layer sequence that implements a Transformer decoder block.
 
   The input to the layer sequence is an activation tensor.
@@ -128,7 +128,7 @@ def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
     d_attention_value: int: depth of value vector for each attention head
     attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
-    share_kv: bool, whether to share keys and values
+    share_qk: bool, whether to share queries and keys
     layer_idx: which layer are we at (for bookkeeping)
     mode: str: 'train' or 'eval'
 
@@ -140,7 +140,7 @@ def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
       tl.CausalAttention(
           d_model, n_heads=n_heads, d_attention_key=d_attention_key,
           d_attention_value=d_attention_value, attention_type=attention_type,
-          share_kv=share_kv, mode=mode),
+          share_qk=share_qk, mode=mode),
       tl.Dropout(rate=dropout, name='attention_%d' % layer_idx, mode=mode),
   ]
   feed_forward = [
@@ -161,7 +161,7 @@ def TransformerDecoder(vocab_size=None,
                        d_attention_value=None,
                        attention_type=tl.DotProductCausalAttention,
                        dropout=0.1,
-                       share_kv=False,
+                       share_qk=False,
                        max_len=2048,
                        mode='train'):
   """Returns a Transformer decoder model.
@@ -183,7 +183,7 @@ def TransformerDecoder(vocab_size=None,
         (default is d_model // n_heads)
     attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
-    share_kv: bool, whether to share keys and values in decoder attention
+    share_qk: bool, whether to share queries and keys in decoder attention
     max_len: int: maximum symbol length for positional encoding
     mode: str: 'train' or 'eval'
 
@@ -201,7 +201,7 @@ def TransformerDecoder(vocab_size=None,
       tl.PositionalEncoding(max_len=max_len),
       [DecoderBlock(  # pylint: disable=g-complex-comprehension
           d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-          attention_type, dropout, share_kv, i, mode)
+          attention_type, dropout, share_qk, i, mode)
        for i in range(n_layers)],   # vecs
       tl.LayerNorm(),               # vecs
   )
@@ -216,7 +216,7 @@ def TransformerLM(vocab_size,
                   d_attention_value=None,
                   attention_type=tl.DotProductCausalAttention,
                   dropout=0.1,
-                  share_kv=False,
+                  share_qk=False,
                   max_len=2048,
                   mode='train'):
   """Returns a Transformer language model.
@@ -236,7 +236,7 @@ def TransformerLM(vocab_size,
         (default is d_model // n_heads)
     attention_type: subclass of tl.BaseCausalAttention: attention class to use
     dropout: float: dropout rate (how much to drop out)
-    share_kv: bool, whether to share keys and values in decoder attention
+    share_qk: bool, whether to share queries and keys in decoder attention
     max_len: int: maximum symbol length for positional encoding
     mode: str: 'train' or 'eval'
 
@@ -254,7 +254,7 @@ def TransformerLM(vocab_size,
       embedder,                     # vecs
       [DecoderBlock(  # pylint: disable=g-complex-comprehension
           d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-          attention_type, dropout, share_kv, i, mode)
+          attention_type, dropout, share_qk, i, mode)
        for i in range(n_layers)],   # vecs
       tl.LayerNorm(),               # vecs
       tl.Dense(vocab_size),         # vecs

From 2a5617f136c5575f59c2c7927c114580007832b1 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 20 Sep 2019 14:54:52 -0700
Subject: [PATCH 2473/2720] internal fixes.

PiperOrigin-RevId: 270356107
---
 tensor2tensor/trax/trainer.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index 7801f4711..682b70cd3 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -93,6 +93,7 @@ def _setup_gin():
 
 
 def main(_):
+
   logging.set_verbosity(FLAGS.log_level)
 
   if FLAGS.enable_eager_execution:

From 7484887f48c08d712eda8eb0800f59f5a242aea3 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 20 Sep 2019 16:23:29 -0700
Subject: [PATCH 2474/2720] Add base Reformer config for enwik8

PiperOrigin-RevId: 270373209
---
 .../trax/configs/reformer_enwik8.gin          | 131 ++++++++++++++++++
 .../trax/configs/reformer_imagenet64.gin      |   1 +
 2 files changed, 132 insertions(+)
 create mode 100644 tensor2tensor/trax/configs/reformer_enwik8.gin

diff --git a/tensor2tensor/trax/configs/reformer_enwik8.gin b/tensor2tensor/trax/configs/reformer_enwik8.gin
new file mode 100644
index 000000000..1eb3e715c
--- /dev/null
+++ b/tensor2tensor/trax/configs/reformer_enwik8.gin
@@ -0,0 +1,131 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters that will vary between experiments:
+# ==============================================================================
+train.model = @trax.models.TransformerLM
+# inputs.n_chunks = 16  # Uncomment this line iff using TransformerRevnetLM
+
+TransformerLM.n_layers = 3
+TransformerRevnetLM.n_layers = 3
+
+TransformerLM.dropout = 0.1
+TransformerRevnetLM.dropout = 0.1
+
+# Attention types:
+# MemoryEfficientCausalAttention: full attention
+# MergedHashedCausalAttention: timebin
+# MergedMultiHashedCausalAttentionV2: hashbin
+TransformerLM.attention_type = @MergedHashedCausalAttention
+TransformerRevnetLM.attention_type = @MergedHashedCausalAttention
+
+# These three should all be equal, and MergedMultiHashedCausalAttentionV2
+# requires that share_qk be set to True.
+MemoryEfficientCausalAttention.share_qk = True
+TransformerLM.share_qk = True
+TransformerRevnetLM.share_qk = True
+
+# These four should all be equal
+TransformerLM.d_attention_key = 64
+TransformerLM.d_attention_value = 64
+TransformerRevnetLM.d_attention_key = 64
+TransformerRevnetLM.d_attention_value = 64
+
+# Use MergedHashedCausalAttention for timebin
+MergedHashedCausalAttention.n_bins = 512
+
+# Use MergedMultiHashedCausalAttentionV2 for hashbin
+MergedMultiHashedCausalAttentionV2.n_bins = 512
+MergedMultiHashedCausalAttentionV2.n_buckets = 1024  # Always 2 * n_bins
+MergedMultiHashedCausalAttentionV2.n_hashes = 2
+MergedMultiHashedCausalAttentionV2.drop_for_hash_rate = 0.0
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 1
+batch_fun.eval_batch_size = 8
+batch_fun.max_eval_length = 65536
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_enwik8_l65k'
+inputs.input_name = 'targets'
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 2.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 500
+train.eval_steps = 8
+train.inputs = @trax.inputs.inputs
+# train.model: see top
+train.optimizer = @trax.optimizers.Adafactor
+train.train_steps = 60000
+train.trainer_class = @MemoryEfficientTrainer
+train.save_steps = \
+    [1000, 5000, 10000, 15000, 20000, 25000, 30000, 35000, 40000, 45000, 50000,
+     55000, 60000]
+
+# Parameters for MemoryEfficientCausalAttention:
+# ==============================================================================
+MemoryEfficientCausalAttention.dropout = 0.0
+MemoryEfficientCausalAttention.loop_stride = 256
+# MemoryEfficientCausalAttention.share_qk: see top
+
+# Parameters for MergedHashedCausalAttention:
+# ==============================================================================
+MergedHashedCausalAttention.dropout = 0.0
+# MergedHashedCausalAttention.n_bins: see top
+MergedHashedCausalAttention.bin_by_time = True
+MergedHashedCausalAttention.one_rng = False
+
+# Parameters for MergedMultiHashedCausalAttentionV2:
+# ==============================================================================
+MergedMultiHashedCausalAttentionV2.allow_duplicate_attention = False
+MergedMultiHashedCausalAttentionV2.attend_across_buckets = False
+MergedMultiHashedCausalAttentionV2.rehash_each_round = True
+# MergedMultiHashedCausalAttentionV2.n_bins: see top
+# MergedMultiHashedCausalAttentionV2.n_buckets: see top
+# MergedMultiHashedCausalAttentionV2.n_hashes: see top
+MergedMultiHashedCausalAttentionV2.one_rng = False
+MergedMultiHashedCausalAttentionV2.hard_k = 0
+MergedMultiHashedCausalAttentionV2.dropout = 0.0
+# MergedMultiHashedCausalAttentionV2.drop_for_hash_rate: see top
+
+# Parameters for TransformerLM:
+# ==============================================================================
+# TransformerLM.attention_type: see top
+# TransformerLM.d_attention_key: see top
+# TransformerLM.d_attention_value: see top
+TransformerLM.d_model = 1024
+TransformerLM.d_ff = 4096
+# TransformerLM.dropout: see top
+TransformerLM.max_len = 65536
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 8
+# TransformerLM.n_layers: see top
+# TransformerLM.share_qk: see top
+TransformerLM.vocab_size = 258  # Includes pad token and unused EOS token
+
+# Parameters for TransformerRevnetLM:
+# ==============================================================================
+# TransformerRevnetLM.attention_type: see top
+# TransformerRevnetLM.d_attention_key: see top
+# TransformerRevnetLM.d_attention_value: see top
+TransformerRevnetLM.d_model = 1024
+TransformerRevnetLM.d_ff = 4096
+# TransformerRevnetLM.dropout: see top
+TransformerRevnetLM.max_len = 65536
+TransformerRevnetLM.mode = 'train'
+TransformerRevnetLM.n_heads = 8
+# TransformerRevnetLM.n_layers: see top
+TransformerRevnetLM.vocab_size = 258  # Includes pad token and unused EOS token
+TransformerRevnetLM.n_chunks = 16
+TransformerRevnetLM.n_attention_chunks = 1
diff --git a/tensor2tensor/trax/configs/reformer_imagenet64.gin b/tensor2tensor/trax/configs/reformer_imagenet64.gin
index 2653cf01c..6f0cff10d 100644
--- a/tensor2tensor/trax/configs/reformer_imagenet64.gin
+++ b/tensor2tensor/trax/configs/reformer_imagenet64.gin
@@ -6,6 +6,7 @@ import tensor2tensor.trax.trax
 # Parameters that will vary between experiments:
 # ==============================================================================
 train.model = @trax.models.TransformerLM
+# inputs.n_chunks = 16  # Uncomment this line iff using TransformerRevnetLM
 
 TransformerLM.n_layers = 3
 TransformerRevnetLM.n_layers = 3

From 44da27423239fee8fd39a02833a777fd239a4bf4 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 20 Sep 2019 17:48:58 -0700
Subject: [PATCH 2475/2720] Configs for reformer experiments and one more
 tie-in.

PiperOrigin-RevId: 270385869
---
 .../trax/configs/reformer_base_sweep.yaml     |   2 +
 .../trax/configs/reformer_enwik8.gin          |  20 +--
 .../trax/configs/reformer_enwik8_rev.gin      | 133 ++++++++++++++++++
 .../trax/configs/reformer_hash_sweep.yaml     |   1 +
 .../trax/configs/reformer_imagenet64.gin      |  20 +--
 .../trax/configs/reformer_imagenet64_rev.gin  | 130 +++++++++++++++++
 tensor2tensor/trax/layers/attention.py        |   8 +-
 7 files changed, 292 insertions(+), 22 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/reformer_base_sweep.yaml
 create mode 100644 tensor2tensor/trax/configs/reformer_enwik8_rev.gin
 create mode 100644 tensor2tensor/trax/configs/reformer_hash_sweep.yaml
 create mode 100644 tensor2tensor/trax/configs/reformer_imagenet64_rev.gin

diff --git a/tensor2tensor/trax/configs/reformer_base_sweep.yaml b/tensor2tensor/trax/configs/reformer_base_sweep.yaml
new file mode 100644
index 000000000..490518f1e
--- /dev/null
+++ b/tensor2tensor/trax/configs/reformer_base_sweep.yaml
@@ -0,0 +1,2 @@
+share_qk: [True, False]
+attn_kv: [64, 128]
diff --git a/tensor2tensor/trax/configs/reformer_enwik8.gin b/tensor2tensor/trax/configs/reformer_enwik8.gin
index 1eb3e715c..b711b2895 100644
--- a/tensor2tensor/trax/configs/reformer_enwik8.gin
+++ b/tensor2tensor/trax/configs/reformer_enwik8.gin
@@ -18,20 +18,22 @@ TransformerRevnetLM.dropout = 0.1
 # MemoryEfficientCausalAttention: full attention
 # MergedHashedCausalAttention: timebin
 # MergedMultiHashedCausalAttentionV2: hashbin
-TransformerLM.attention_type = @MergedHashedCausalAttention
-TransformerRevnetLM.attention_type = @MergedHashedCausalAttention
+TransformerLM.attention_type = @MemoryEfficientCausalAttention
+TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
 
 # These three should all be equal, and MergedMultiHashedCausalAttentionV2
 # requires that share_qk be set to True.
-MemoryEfficientCausalAttention.share_qk = True
-TransformerLM.share_qk = True
-TransformerRevnetLM.share_qk = True
+share_qk = True
+MemoryEfficientCausalAttention.share_qk = %share_qk
+TransformerLM.share_qk = %share_qk
+TransformerRevnetLM.share_qk = %share_qk
 
 # These four should all be equal
-TransformerLM.d_attention_key = 64
-TransformerLM.d_attention_value = 64
-TransformerRevnetLM.d_attention_key = 64
-TransformerRevnetLM.d_attention_value = 64
+attn_kv = 64
+TransformerLM.d_attention_key = %attn_kv
+TransformerLM.d_attention_value = %attn_kv
+TransformerRevnetLM.d_attention_key = %attn_kv
+TransformerRevnetLM.d_attention_value = %attn_kv
 
 # Use MergedHashedCausalAttention for timebin
 MergedHashedCausalAttention.n_bins = 512
diff --git a/tensor2tensor/trax/configs/reformer_enwik8_rev.gin b/tensor2tensor/trax/configs/reformer_enwik8_rev.gin
new file mode 100644
index 000000000..c9f73fe7a
--- /dev/null
+++ b/tensor2tensor/trax/configs/reformer_enwik8_rev.gin
@@ -0,0 +1,133 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters that will vary between experiments:
+# ==============================================================================
+train.model = @trax.models.TransformerRevnetLM
+inputs.n_chunks = 16  # Uncomment this line iff using TransformerRevnetLM
+
+TransformerLM.n_layers = 3
+TransformerRevnetLM.n_layers = 3
+
+TransformerLM.dropout = 0.1
+TransformerRevnetLM.dropout = 0.1
+
+# Attention types:
+# MemoryEfficientCausalAttention: full attention
+# MergedHashedCausalAttention: timebin
+# MergedMultiHashedCausalAttentionV2: hashbin
+TransformerLM.attention_type = @MemoryEfficientCausalAttention
+TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
+
+# These three should all be equal, and MergedMultiHashedCausalAttentionV2
+# requires that share_qk be set to True.
+share_qk = True
+MemoryEfficientCausalAttention.share_qk = %share_qk
+TransformerLM.share_qk = %share_qk
+TransformerRevnetLM.share_qk = %share_qk
+
+# These four should all be equal
+attn_kv = 64
+TransformerLM.d_attention_key = %attn_kv
+TransformerLM.d_attention_value = %attn_kv
+TransformerRevnetLM.d_attention_key = %attn_kv
+TransformerRevnetLM.d_attention_value = %attn_kv
+
+# Use MergedHashedCausalAttention for timebin
+MergedHashedCausalAttention.n_bins = 512
+
+# Use MergedMultiHashedCausalAttentionV2 for hashbin
+MergedMultiHashedCausalAttentionV2.n_bins = 512
+MergedMultiHashedCausalAttentionV2.n_buckets = 1024  # Always 2 * n_bins
+MergedMultiHashedCausalAttentionV2.n_hashes = 2
+MergedMultiHashedCausalAttentionV2.drop_for_hash_rate = 0.0
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 1
+batch_fun.eval_batch_size = 8
+batch_fun.max_eval_length = 65536
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_enwik8_l65k'
+inputs.input_name = 'targets'
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 2.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 500
+train.eval_steps = 8
+train.inputs = @trax.inputs.inputs
+# train.model: see top
+train.optimizer = @trax.optimizers.Adafactor
+train.train_steps = 60000
+train.trainer_class = @MemoryEfficientTrainer
+train.save_steps = \
+    [1000, 5000, 10000, 15000, 20000, 25000, 30000, 35000, 40000, 45000, 50000,
+     55000, 60000]
+
+# Parameters for MemoryEfficientCausalAttention:
+# ==============================================================================
+MemoryEfficientCausalAttention.dropout = 0.0
+MemoryEfficientCausalAttention.loop_stride = 256
+# MemoryEfficientCausalAttention.share_qk: see top
+
+# Parameters for MergedHashedCausalAttention:
+# ==============================================================================
+MergedHashedCausalAttention.dropout = 0.0
+# MergedHashedCausalAttention.n_bins: see top
+MergedHashedCausalAttention.bin_by_time = True
+MergedHashedCausalAttention.one_rng = False
+
+# Parameters for MergedMultiHashedCausalAttentionV2:
+# ==============================================================================
+MergedMultiHashedCausalAttentionV2.allow_duplicate_attention = False
+MergedMultiHashedCausalAttentionV2.attend_across_buckets = False
+MergedMultiHashedCausalAttentionV2.rehash_each_round = True
+# MergedMultiHashedCausalAttentionV2.n_bins: see top
+# MergedMultiHashedCausalAttentionV2.n_buckets: see top
+# MergedMultiHashedCausalAttentionV2.n_hashes: see top
+MergedMultiHashedCausalAttentionV2.one_rng = False
+MergedMultiHashedCausalAttentionV2.hard_k = 0
+MergedMultiHashedCausalAttentionV2.dropout = 0.0
+# MergedMultiHashedCausalAttentionV2.drop_for_hash_rate: see top
+
+# Parameters for TransformerLM:
+# ==============================================================================
+# TransformerLM.attention_type: see top
+# TransformerLM.d_attention_key: see top
+# TransformerLM.d_attention_value: see top
+TransformerLM.d_model = 1024
+TransformerLM.d_ff = 4096
+# TransformerLM.dropout: see top
+TransformerLM.max_len = 65536
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 8
+# TransformerLM.n_layers: see top
+# TransformerLM.share_qk: see top
+TransformerLM.vocab_size = 258  # Includes pad token and unused EOS token
+
+# Parameters for TransformerRevnetLM:
+# ==============================================================================
+# TransformerRevnetLM.attention_type: see top
+# TransformerRevnetLM.d_attention_key: see top
+# TransformerRevnetLM.d_attention_value: see top
+TransformerRevnetLM.d_model = 1024
+TransformerRevnetLM.d_ff = 4096
+# TransformerRevnetLM.dropout: see top
+TransformerRevnetLM.max_len = 65536
+TransformerRevnetLM.mode = 'train'
+TransformerRevnetLM.n_heads = 8
+# TransformerRevnetLM.n_layers: see top
+TransformerRevnetLM.vocab_size = 258  # Includes pad token and unused EOS token
+TransformerRevnetLM.n_chunks = 16
+TransformerRevnetLM.n_attention_chunks = 1
diff --git a/tensor2tensor/trax/configs/reformer_hash_sweep.yaml b/tensor2tensor/trax/configs/reformer_hash_sweep.yaml
new file mode 100644
index 000000000..405d33f78
--- /dev/null
+++ b/tensor2tensor/trax/configs/reformer_hash_sweep.yaml
@@ -0,0 +1 @@
+attn_kv: [64, 128]
diff --git a/tensor2tensor/trax/configs/reformer_imagenet64.gin b/tensor2tensor/trax/configs/reformer_imagenet64.gin
index 6f0cff10d..7a3b74074 100644
--- a/tensor2tensor/trax/configs/reformer_imagenet64.gin
+++ b/tensor2tensor/trax/configs/reformer_imagenet64.gin
@@ -15,20 +15,22 @@ TransformerRevnetLM.n_layers = 3
 # MemoryEfficientCausalAttention: full attention
 # MergedHashedCausalAttention: timebin
 # MergedMultiHashedCausalAttentionV2: hashbin
-TransformerLM.attention_type = @MergedMultiHashedCausalAttentionV2
-TransformerRevnetLM.attention_type = @MergedMultiHashedCausalAttentionV2
+TransformerLM.attention_type = @MemoryEfficientCausalAttention
+TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
 
 # These three should all be equal, and MergedMultiHashedCausalAttentionV2
 # requires that share_qk be set to True.
-MemoryEfficientCausalAttention.share_qk = True
-TransformerLM.share_qk = True
-TransformerRevnetLM.share_qk = True
+share_qk = True
+MemoryEfficientCausalAttention.share_qk = %share_qk
+TransformerLM.share_qk = %share_qk
+TransformerRevnetLM.share_qk = %share_qk
 
 # These four should all be equal
-TransformerLM.d_attention_key = 64
-TransformerLM.d_attention_value = 64
-TransformerRevnetLM.d_attention_key = 64
-TransformerRevnetLM.d_attention_value = 64
+attn_kv = 64
+TransformerLM.d_attention_key = %attn_kv
+TransformerLM.d_attention_value = %attn_kv
+TransformerRevnetLM.d_attention_key = %attn_kv
+TransformerRevnetLM.d_attention_value = %attn_kv
 
 # Use MergedHashedCausalAttention for timebin
 MergedHashedCausalAttention.n_bins = 64
diff --git a/tensor2tensor/trax/configs/reformer_imagenet64_rev.gin b/tensor2tensor/trax/configs/reformer_imagenet64_rev.gin
new file mode 100644
index 000000000..0a131044b
--- /dev/null
+++ b/tensor2tensor/trax/configs/reformer_imagenet64_rev.gin
@@ -0,0 +1,130 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters that will vary between experiments:
+# ==============================================================================
+train.model = @trax.models.TransformerRevnetLM
+inputs.n_chunks = 16  # Uncomment this line iff using TransformerRevnetLM
+
+TransformerLM.n_layers = 3
+TransformerRevnetLM.n_layers = 3
+
+# Attention types:
+# MemoryEfficientCausalAttention: full attention
+# MergedHashedCausalAttention: timebin
+# MergedMultiHashedCausalAttentionV2: hashbin
+TransformerLM.attention_type = @MemoryEfficientCausalAttention
+TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
+
+# These three should all be equal, and MergedMultiHashedCausalAttentionV2
+# requires that share_qk be set to True.
+share_qk = True
+MemoryEfficientCausalAttention.share_qk = %share_qk
+TransformerLM.share_qk = %share_qk
+TransformerRevnetLM.share_qk = %share_qk
+
+# These four should all be equal
+attn_kv = 64
+TransformerLM.d_attention_key = %attn_kv
+TransformerLM.d_attention_value = %attn_kv
+TransformerRevnetLM.d_attention_key = %attn_kv
+TransformerRevnetLM.d_attention_value = %attn_kv
+
+# Use MergedHashedCausalAttention for timebin
+MergedHashedCausalAttention.n_bins = 64
+
+# Use MergedMultiHashedCausalAttentionV2 for hashbin
+MergedMultiHashedCausalAttentionV2.n_bins = 96
+MergedMultiHashedCausalAttentionV2.n_buckets = 192  # Always 2 * n_bins
+MergedMultiHashedCausalAttentionV2.n_hashes = 2
+MergedMultiHashedCausalAttentionV2.drop_for_hash_rate = 0.0
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 1
+batch_fun.eval_batch_size = 8
+batch_fun.max_eval_length = 12288  # 64 * 64 * 3
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
+inputs.input_name = 'targets'
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 2.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 500
+train.eval_steps = 64
+train.inputs = @trax.inputs.inputs
+# train.model: see top
+train.optimizer = @trax.optimizers.Adafactor
+train.train_steps = 500000
+train.trainer_class = @MemoryEfficientTrainer
+train.save_steps = \
+    [1000, 5000, 10000, 20000, 40000, 60000, 80000,
+     100000, 200000, 300000, 400000, 500000]
+
+# Parameters for MemoryEfficientCausalAttention:
+# ==============================================================================
+MemoryEfficientCausalAttention.dropout = 0.0
+MemoryEfficientCausalAttention.loop_stride = 512
+# MemoryEfficientCausalAttention.share_qk: see top
+
+# Parameters for MergedHashedCausalAttention:
+# ==============================================================================
+MergedHashedCausalAttention.dropout = 0.0
+# MergedHashedCausalAttention.n_bins: see top
+MergedHashedCausalAttention.bin_by_time = True
+MergedHashedCausalAttention.one_rng = False
+
+# Parameters for MergedMultiHashedCausalAttentionV2:
+# ==============================================================================
+MergedMultiHashedCausalAttentionV2.allow_duplicate_attention = False
+MergedMultiHashedCausalAttentionV2.attend_across_buckets = False
+MergedMultiHashedCausalAttentionV2.rehash_each_round = True
+# MergedMultiHashedCausalAttentionV2.n_bins: see top
+# MergedMultiHashedCausalAttentionV2.n_buckets: see top
+# MergedMultiHashedCausalAttentionV2.n_hashes: see top
+MergedMultiHashedCausalAttentionV2.one_rng = False
+MergedMultiHashedCausalAttentionV2.hard_k = 0
+MergedMultiHashedCausalAttentionV2.dropout = 0.0
+# MergedMultiHashedCausalAttentionV2.drop_for_hash_rate: see top
+
+# Parameters for TransformerLM:
+# ==============================================================================
+# TransformerLM.attention_type: see top
+# TransformerLM.d_attention_key: see top
+# TransformerLM.d_attention_value: see top
+TransformerLM.d_model = 1024
+TransformerLM.d_ff = 4096
+TransformerLM.dropout = 0.0
+TransformerLM.max_len = 12288  # 64 * 64 * 3
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 8
+# TransformerLM.n_layers: see top
+# TransformerLM.share_qk: see top
+TransformerLM.vocab_size = 256
+
+# Parameters for TransformerRevnetLM:
+# ==============================================================================
+# TransformerRevnetLM.attention_type: see top
+# TransformerRevnetLM.d_attention_key: see top
+# TransformerRevnetLM.d_attention_value: see top
+TransformerRevnetLM.d_model = 1024
+TransformerRevnetLM.d_ff = 4096
+TransformerRevnetLM.dropout = 0.0
+TransformerRevnetLM.max_len = 12288  # 64 * 64 * 3
+TransformerRevnetLM.mode = 'train'
+TransformerRevnetLM.n_heads = 8
+# TransformerRevnetLM.n_layers: see top
+TransformerRevnetLM.vocab_size = 256
+TransformerRevnetLM.n_chunks = 16
+TransformerRevnetLM.n_attention_chunks = 1
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 84379fb3c..906499cc3 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -460,8 +460,8 @@ def make_mask(N, M, k):  # pylint: disable=invalid-name
       Returns:
         N x M mask, where 1.0 indicates that attention is not allowed.
       """
-      x = np.arange(N, dtype=np.int32)
-      y = np.arange(M, dtype=np.int32)
+      x = jax.lax.tie_in(k, np.arange(N, dtype=np.int32))
+      y = jax.lax.tie_in(k, np.arange(M, dtype=np.int32))
       mask = jax.lax.lt(
           (jax.lax.broadcast_in_dim(
               x, shape=(N, M), broadcast_dimensions=(0,)) + k),
@@ -480,8 +480,8 @@ def make_self_mask(N, M, k):  # pylint: disable=invalid-name
       Returns:
         N x M mask, where 1.0 indicates that attention is not allowed.
       """
-      x = np.arange(N, dtype=np.int32)
-      y = np.arange(M, dtype=np.int32)
+      x = jax.lax.tie_in(k, np.arange(N, dtype=np.int32))
+      y = jax.lax.tie_in(k, np.arange(M, dtype=np.int32))
       mask = jax.lax.eq(
           (jax.lax.broadcast_in_dim(
               x, shape=(N, M), broadcast_dimensions=(0,)) + k),

From 4343cfeb38a4d25352d6202b96a4e1721d7d2973 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sat, 21 Sep 2019 16:15:05 -0700
Subject: [PATCH 2476/2720] Stash hash buckets in reversible transformer

PiperOrigin-RevId: 270479770
---
 tensor2tensor/trax/layers/attention.py | 36 ++++++++++++++++++++++----
 tensor2tensor/trax/layers/base.py      | 18 +++++++++++--
 2 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 906499cc3..7434cb82a 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -1123,6 +1123,14 @@ def custom_grad(self, inputs, output, ct, params=(), state=(), rng=None,
   def batch_call_and_or_grad(self, qk, v, ct=None, return_output=True,
                              rng=None):
     assert return_output or ct is not None, 'No work to perform!'
+    # pylint: disable=protected-access
+    stash_buckets = (return_output and ct is None
+                     and base.Layer._STASH_IN is not None)
+    if return_output and ct is not None and base.Layer._STASH_OUT is not None:
+      buckets = base.Layer._STASH_OUT.pop(self)
+    else:
+      buckets = None
+    # pylint: enable=protected-access
 
     # The approach here is to perform attention for one batch element and head
     # at a time. Note that there is absolutely no interaction across examples or
@@ -1140,6 +1148,10 @@ def batch_call_and_or_grad(self, qk, v, ct=None, return_output=True,
     if return_output:
       out_accum = np.zeros_like(qk)
       init_vals = init_vals + (out_accum,)
+    if stash_buckets:
+      buckets_accum = np.zeros(
+          [qk.shape[0], self.n_hashes * qk.shape[1]], dtype=np.int32)
+      init_vals = init_vals + (buckets_accum,)
     if ct is not None:
       qk_ct_accum = np.zeros_like(qk)
       v_ct_accum = np.zeros_like(v)
@@ -1162,11 +1174,19 @@ def body_fun(vals):
       v_slice = jax.lax.dynamic_index_in_dim(
           v, batch_loop_idx, axis=0, keepdims=False)
 
+      if buckets is None:
+        buckets_slice = self.hash_vectors(qk_slice, rng=hash_rng)
+      else:
+        buckets_slice = jax.lax.dynamic_index_in_dim(
+            buckets, batch_loop_idx, axis=0, keepdims=False)
+
       if ct is None:
-        out_slice = self.single_call(qk_slice, v_slice, hash_rng=hash_rng)
+        out_slice = self.single_call(
+            qk_slice, v_slice, buckets_slice, hash_rng=hash_rng)
       else:
         def _do_single_call(qk_slice, v_slice):
-          return self.single_call(qk_slice, v_slice, hash_rng=hash_rng)
+          return self.single_call(
+              qk_slice, v_slice, buckets_slice, hash_rng=hash_rng)
         ct_slice = jax.lax.dynamic_index_in_dim(
             ct, batch_loop_idx, axis=0, keepdims=False)
         out_slice, vjpfun = jax.vjp(_do_single_call, qk_slice, v_slice)
@@ -1178,6 +1198,11 @@ def _do_single_call(qk_slice, v_slice):
         out_accum = jax.lax.dynamic_update_index_in_dim(
             out_accum, out_slice, batch_loop_idx, axis=0)
         new_vals = new_vals + (out_accum,)
+      if stash_buckets:
+        buckets_accum = vals[2]
+        buckets_accum = jax.lax.dynamic_update_index_in_dim(
+            buckets_accum, buckets_slice, batch_loop_idx, axis=0)
+        new_vals = new_vals + (buckets_accum,)
       if ct is not None:
         qk_ct_accum, v_ct_accum = vals[-2:]
         qk_ct_accum = jax.lax.dynamic_update_index_in_dim(
@@ -1195,6 +1220,9 @@ def _do_single_call(qk_slice, v_slice):
     else:
       out = None
 
+    if stash_buckets:
+      base.Layer._STASH_IN[self] = final_vals[2]  # pylint: disable=protected-access
+
     if ct is not None:
       input_ct = final_vals[-2:]
     else:
@@ -1255,11 +1283,9 @@ def hash_vectors(self, vecs, rng):
 
     return buckets
 
-  def single_call(self, qk, v, hash_rng=None):
+  def single_call(self, qk, v, buckets, hash_rng=None):
     # We use the same vector as both a query and a key.
     seqlen = qk.shape[-2]
-
-    buckets = self.hash_vectors(qk, rng=hash_rng)
     assert int(buckets.shape[0]) == self.n_hashes * seqlen
 
     ticker = jax.lax.tie_in(qk, np.arange(self.n_hashes * seqlen))
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 93b9a6127..6908cd6ad 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -239,6 +239,10 @@ def initialize(self, input_shapes, input_dtype, rng):
       name, trace = self.__class__.__name__, _short_traceback(skip=3)
       raise LayerError(name, 'initialize', self._caller, input_shapes, trace)
 
+  # XXX(kitaev):
+  _STASH_IN = None
+  _STASH_OUT = None
+
   def __call__(self, x, params=(), state=(), **kwargs):
     try:
       # If params are nothing, we may be reusing this layer.
@@ -251,7 +255,7 @@ def __call__(self, x, params=(), state=(), **kwargs):
         # In this case, we're called for the first time: cache parameters.
         self._params = params
 
-      if not self.has_custom_grad:
+      if not self.has_custom_grad or Layer._STASH_IN is not None:
         return self.call(x, params=params, state=state, **kwargs)
 
       # Custom gradients part.
@@ -284,10 +288,20 @@ def do_call(y, params):
       # For the exact specification of this custom transformation see this link:
       # https://jax.readthedocs.io/en/latest/jax.html#jax.defjvp_all
       def do_call_vjp(y, params):
+        """Custom gradient (vjp) function."""
+        stash = None
+        if Layer._STASH_IN is None:
+          Layer._STASH_IN = stash = {}
         output = check_end_state(self.call(y, params=params, state=state,
                                            **kwargs))
+        if stash is not None:
+          Layer._STASH_IN = None
         def vjpfun(grad):
-          return self.custom_grad(y, output, grad, params, state, **kwargs)
+          assert Layer._STASH_OUT is None
+          Layer._STASH_OUT = stash
+          res = self.custom_grad(y, output, grad, params, state, **kwargs)
+          Layer._STASH_OUT = None
+          return res
         return output, vjpfun
 
       jax.defvjp_all(do_call, do_call_vjp)

From 7a7c067e73791e6353b2a29abf4793f8b55e025a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Sat, 21 Sep 2019 18:13:35 -0700
Subject: [PATCH 2477/2720] Use hashed attention by default in configs, 128 kv
 size, add sweeps for nhashes and a large sweep.

PiperOrigin-RevId: 270487144
---
 .../trax/configs/reformer_enwik8.gin          |  8 ++++---
 .../trax/configs/reformer_enwik8_rev.gin      |  8 ++++---
 .../trax/configs/reformer_hash_sweep.yaml     |  2 +-
 .../trax/configs/reformer_imagenet64.gin      |  8 ++++---
 .../trax/configs/reformer_imagenet64_rev.gin  |  8 ++++---
 .../trax/configs/reformer_large_sweep.yaml    |  3 +++
 .../trax/configs/transformer_copy.gin         | 24 ++++++++++---------
 7 files changed, 37 insertions(+), 24 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/reformer_large_sweep.yaml

diff --git a/tensor2tensor/trax/configs/reformer_enwik8.gin b/tensor2tensor/trax/configs/reformer_enwik8.gin
index b711b2895..63b032c17 100644
--- a/tensor2tensor/trax/configs/reformer_enwik8.gin
+++ b/tensor2tensor/trax/configs/reformer_enwik8.gin
@@ -18,8 +18,10 @@ TransformerRevnetLM.dropout = 0.1
 # MemoryEfficientCausalAttention: full attention
 # MergedHashedCausalAttention: timebin
 # MergedMultiHashedCausalAttentionV2: hashbin
-TransformerLM.attention_type = @MemoryEfficientCausalAttention
-TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
+# TransformerLM.attention_type = @MemoryEfficientCausalAttention
+# TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
+TransformerLM.attention_type = @MergedMultiHashedCausalAttentionV2
+TransformerRevnetLM.attention_type = @MergedMultiHashedCausalAttentionV2
 
 # These three should all be equal, and MergedMultiHashedCausalAttentionV2
 # requires that share_qk be set to True.
@@ -29,7 +31,7 @@ TransformerLM.share_qk = %share_qk
 TransformerRevnetLM.share_qk = %share_qk
 
 # These four should all be equal
-attn_kv = 64
+attn_kv = 128
 TransformerLM.d_attention_key = %attn_kv
 TransformerLM.d_attention_value = %attn_kv
 TransformerRevnetLM.d_attention_key = %attn_kv
diff --git a/tensor2tensor/trax/configs/reformer_enwik8_rev.gin b/tensor2tensor/trax/configs/reformer_enwik8_rev.gin
index c9f73fe7a..f762b8835 100644
--- a/tensor2tensor/trax/configs/reformer_enwik8_rev.gin
+++ b/tensor2tensor/trax/configs/reformer_enwik8_rev.gin
@@ -18,8 +18,10 @@ TransformerRevnetLM.dropout = 0.1
 # MemoryEfficientCausalAttention: full attention
 # MergedHashedCausalAttention: timebin
 # MergedMultiHashedCausalAttentionV2: hashbin
-TransformerLM.attention_type = @MemoryEfficientCausalAttention
-TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
+# TransformerLM.attention_type = @MemoryEfficientCausalAttention
+# TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
+TransformerLM.attention_type = @MergedMultiHashedCausalAttentionV2
+TransformerRevnetLM.attention_type = @MergedMultiHashedCausalAttentionV2
 
 # These three should all be equal, and MergedMultiHashedCausalAttentionV2
 # requires that share_qk be set to True.
@@ -29,7 +31,7 @@ TransformerLM.share_qk = %share_qk
 TransformerRevnetLM.share_qk = %share_qk
 
 # These four should all be equal
-attn_kv = 64
+attn_kv = 128
 TransformerLM.d_attention_key = %attn_kv
 TransformerLM.d_attention_value = %attn_kv
 TransformerRevnetLM.d_attention_key = %attn_kv
diff --git a/tensor2tensor/trax/configs/reformer_hash_sweep.yaml b/tensor2tensor/trax/configs/reformer_hash_sweep.yaml
index 405d33f78..94216b406 100644
--- a/tensor2tensor/trax/configs/reformer_hash_sweep.yaml
+++ b/tensor2tensor/trax/configs/reformer_hash_sweep.yaml
@@ -1 +1 @@
-attn_kv: [64, 128]
+MergedMultiHashedCausalAttentionV2.n_hashes: [2, 4, 8, 16]
diff --git a/tensor2tensor/trax/configs/reformer_imagenet64.gin b/tensor2tensor/trax/configs/reformer_imagenet64.gin
index 7a3b74074..ac78c3d56 100644
--- a/tensor2tensor/trax/configs/reformer_imagenet64.gin
+++ b/tensor2tensor/trax/configs/reformer_imagenet64.gin
@@ -15,8 +15,10 @@ TransformerRevnetLM.n_layers = 3
 # MemoryEfficientCausalAttention: full attention
 # MergedHashedCausalAttention: timebin
 # MergedMultiHashedCausalAttentionV2: hashbin
-TransformerLM.attention_type = @MemoryEfficientCausalAttention
-TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
+# TransformerLM.attention_type = @MemoryEfficientCausalAttention
+# TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
+TransformerLM.attention_type = @MergedMultiHashedCausalAttentionV2
+TransformerRevnetLM.attention_type = @MergedMultiHashedCausalAttentionV2
 
 # These three should all be equal, and MergedMultiHashedCausalAttentionV2
 # requires that share_qk be set to True.
@@ -26,7 +28,7 @@ TransformerLM.share_qk = %share_qk
 TransformerRevnetLM.share_qk = %share_qk
 
 # These four should all be equal
-attn_kv = 64
+attn_kv = 128
 TransformerLM.d_attention_key = %attn_kv
 TransformerLM.d_attention_value = %attn_kv
 TransformerRevnetLM.d_attention_key = %attn_kv
diff --git a/tensor2tensor/trax/configs/reformer_imagenet64_rev.gin b/tensor2tensor/trax/configs/reformer_imagenet64_rev.gin
index 0a131044b..6d15419cb 100644
--- a/tensor2tensor/trax/configs/reformer_imagenet64_rev.gin
+++ b/tensor2tensor/trax/configs/reformer_imagenet64_rev.gin
@@ -15,8 +15,10 @@ TransformerRevnetLM.n_layers = 3
 # MemoryEfficientCausalAttention: full attention
 # MergedHashedCausalAttention: timebin
 # MergedMultiHashedCausalAttentionV2: hashbin
-TransformerLM.attention_type = @MemoryEfficientCausalAttention
-TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
+# TransformerLM.attention_type = @MemoryEfficientCausalAttention
+# TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
+TransformerLM.attention_type = @MergedMultiHashedCausalAttentionV2
+TransformerRevnetLM.attention_type = @MergedMultiHashedCausalAttentionV2
 
 # These three should all be equal, and MergedMultiHashedCausalAttentionV2
 # requires that share_qk be set to True.
@@ -26,7 +28,7 @@ TransformerLM.share_qk = %share_qk
 TransformerRevnetLM.share_qk = %share_qk
 
 # These four should all be equal
-attn_kv = 64
+attn_kv = 128
 TransformerLM.d_attention_key = %attn_kv
 TransformerLM.d_attention_value = %attn_kv
 TransformerRevnetLM.d_attention_key = %attn_kv
diff --git a/tensor2tensor/trax/configs/reformer_large_sweep.yaml b/tensor2tensor/trax/configs/reformer_large_sweep.yaml
new file mode 100644
index 000000000..9cbf67088
--- /dev/null
+++ b/tensor2tensor/trax/configs/reformer_large_sweep.yaml
@@ -0,0 +1,3 @@
+MergedMultiHashedCausalAttentionV2.n_hashes: [2, 4]
+TransformerRevnetLM.n_layers: [12, 16, 20, 24]
+MultifactorSchedule.constant: [0.3, 1.0]
diff --git a/tensor2tensor/trax/configs/transformer_copy.gin b/tensor2tensor/trax/configs/transformer_copy.gin
index e1f58afa4..15c67e650 100644
--- a/tensor2tensor/trax/configs/transformer_copy.gin
+++ b/tensor2tensor/trax/configs/transformer_copy.gin
@@ -43,20 +43,22 @@ train.has_weights = True
 MemoryEfficientCausalAttention.dropout = 0.0
 MemoryEfficientCausalAttention.loop_stride = 512
 
-# Parameters for MergedMultiHashedCausalAttention:
+# Parameters for MergedMultiHashedCausalAttentionV2:
 # ==============================================================================
-MergedMultiHashedCausalAttention.dropout = 0.0
-MergedMultiHashedCausalAttention.n_bins = 32
-MergedMultiHashedCausalAttention.n_hashes = 2
-MergedMultiHashedCausalAttention.n_buckets_per_bin = 1
-MergedMultiHashedCausalAttention.bin_by_time = False
-MergedMultiHashedCausalAttention.one_rng = False
-MergedMultiHashedCausalAttention.drop_for_hash_rate = 0.0
-MergedMultiHashedCausalAttention.hard_k = 16
+MergedMultiHashedCausalAttentionV2.allow_duplicate_attention = False
+MergedMultiHashedCausalAttentionV2.attend_across_buckets = False
+MergedMultiHashedCausalAttentionV2.rehash_each_round = True
+# MergedMultiHashedCausalAttentionV2.n_bins: 32
+# MergedMultiHashedCausalAttentionV2.n_buckets: 64
+# MergedMultiHashedCausalAttentionV2.n_hashes: 4
+MergedMultiHashedCausalAttentionV2.one_rng = False
+MergedMultiHashedCausalAttentionV2.hard_k = 0
+MergedMultiHashedCausalAttentionV2.dropout = 0.0
+# MergedMultiHashedCausalAttentionV2.drop_for_hash_rate: 0.0
 
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.attention_type = @trax.layers.MergedMultiHashedCausalAttention
+TransformerLM.attention_type = @trax.layers.MergedMultiHashedCausalAttentionV2
 TransformerLM.d_attention_key = 64
 TransformerLM.d_attention_value = 64
 TransformerLM.d_model = 256
@@ -66,5 +68,5 @@ TransformerLM.max_len = 1024
 TransformerLM.mode = 'train'
 TransformerLM.n_heads = 4
 TransformerLM.n_layers = 1
-TransformerLM.share_kv = True
+TransformerLM.share_qk = True
 TransformerLM.vocab_size = 128

From 033381958d6a7bf937f8babdf7bfa8da7ae613cb Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Sat, 21 Sep 2019 19:14:23 -0700
Subject: [PATCH 2478/2720] Implement fast inference for TransformerLM.

Also updated the simulated environment to use it.

PiperOrigin-RevId: 270491153
---
 tensor2tensor/trax/layers/attention.py        | 111 +++++++++++++-----
 tensor2tensor/trax/models/transformer.py      |   6 +-
 tensor2tensor/trax/models/transformer_test.py |  33 ++++++
 tensor2tensor/trax/rl/simple_test.py          |   5 +-
 .../trax/rl/simulated_env_problem.py          |  30 +++--
 .../trax/rl/simulated_env_problem_test.py     |  10 +-
 6 files changed, 145 insertions(+), 50 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 7434cb82a..d2d2cc42c 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -35,8 +35,12 @@
 
 
 @base.layer()
-def ShiftRight(x, **unused_kwargs):
+def ShiftRight(x, mode='train', **unused_kwargs):
   """Layer to shift the tensor to the right by padding on axis 1."""
+  if mode == 'predict':
+    # Do nothing in predict mode, as then the sequence length is 1.
+    return x
+
   pad_widths = [(0, 0)] * len(x.shape)
   pad_widths[1] = (1, 0)  # Padding on axis=1
   padded = np.pad(x, pad_widths, mode='constant',
@@ -67,27 +71,38 @@ def EncoderDecoderMask(x, **unused_kwargs):
   return padding_mask + np.zeros((1, 1, decoder_input.shape[1], 1))
 
 
-# Positional encoding.
-def _positional_encoding_new_params(  # pylint: disable=invalid-name
-    input_shape, input_dtype, rng, max_len=2048):
-  """Helper: create positional encoding parameters."""
-  del input_dtype, rng
-  d_feature = input_shape[-1]
-  pe = onp.zeros((max_len, d_feature), dtype=onp.float32)
-  position = onp.arange(0, max_len)[:, onp.newaxis]
-  div_term = onp.exp(
-      onp.arange(0, d_feature, 2) * -(onp.log(10000.0) / d_feature))
-  pe[:, 0::2] = onp.sin(position * div_term)
-  pe[:, 1::2] = onp.cos(position * div_term)
-  pe = pe[onp.newaxis, :, :]  # [1, max_len, d_feature]
-  return np.array(pe)  # These are trainable parameters, initialized as above.
-
-
-@base.layer(new_parameters=_positional_encoding_new_params)
-def PositionalEncoding(x, params, **unused_kwargs):
+class PositionalEncoding(base.Layer):
   """Implements bare positional encoding."""
-  symbol_size = np.shape(x)[1]
-  return x + params[:, :symbol_size, :]
+
+  def __init__(self, max_len=2048, mode='train'):
+    super(PositionalEncoding, self).__init__()
+    self._max_len = max_len
+    self._mode = mode
+
+  def call(self, inputs, params, state, **kwargs):
+    if self._mode in ('train', 'eval'):
+      x = inputs
+      symbol_size = np.shape(x)[1]
+      return (x + params[:, :symbol_size, :], state)
+    else:
+      assert self._mode == 'predict'
+      # Fast inference: return consectutive elements of the encoding sequence,
+      # storing the index in state.
+      return (inputs + np.expand_dims(params[:, state, :], 1), state + 1)
+
+  def new_parameters(self, input_shape, input_dtype, rng):
+    del input_dtype, rng
+    d_feature = input_shape[-1]
+    pe = onp.zeros((self._max_len, d_feature), dtype=onp.float32)
+    position = onp.arange(0, self._max_len)[:, onp.newaxis]
+    div_term = onp.exp(
+        onp.arange(0, d_feature, 2) * -(onp.log(10000.0) / d_feature))
+    pe[:, 0::2] = onp.sin(position * div_term)
+    pe[:, 1::2] = onp.cos(position * div_term)
+    pe = pe[onp.newaxis, :, :]  # [1, self._max_len, d_feature]
+    pe = np.array(pe)  # These are trainable parameters, initialized as above.
+    state = 0 if self._mode == 'predict' else ()
+    return (pe, state)
 
 
 def DotProductAttention(query, key, value, mask, dropout, mode, rng):
@@ -372,14 +387,34 @@ def __init__(self, dropout=0.0, mode='train'):
   def call(self, inputs, params=(), state=(), rng=None, **kwargs):
     del params
     q, k, v = inputs
-    mask_size = q.shape[-2]
-    # Not all backends define np.tril. However, using onp.tril is inefficient in
-    # that it creates a large global constant. TODO(kitaev): try to find an
-    # alternative that works across all backends.
-    if backend.get_name() == 'jax':
-      mask = np.tril(np.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
+    if self._mode in ('train', 'eval'):
+      mask_size = q.shape[-2]
+      # Not all backends define np.tril. However, using onp.tril is inefficient
+      # in that it creates a large global constant. TODO(kitaev): try to find an
+      # alternative that works across all backends.
+      if backend.get_name() == 'jax':
+        mask = np.tril(
+            np.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
+      else:
+        mask = onp.tril(
+            onp.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
     else:
-      mask = onp.tril(onp.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
+      assert self._mode == 'predict'
+      assert backend.get_name() == 'jax', (
+          'JAX backend is required to use the predict mode.')
+      for x in (q, k, v):
+        assert x.shape[1] == 1, (
+            'In predict mode the input sequence must be of length 1.')
+      # Fast inference: run with only 1 query in each step, storing the sequence
+      # of keys and values calculated so far in state.
+      (new_k, new_v) = (k, v)
+      (k, v, mask, index) = state
+      k = jax.ops.index_update(k, jax.ops.index[:, index, :], new_k[:, 0, :])
+      v = jax.ops.index_update(v, jax.ops.index[:, index, :], new_v[:, 0, :])
+      new_mask = jax.ops.index_update(mask, jax.ops.index[:, :, index], 1)
+      state = (k, v, new_mask, index + 1)
+      mask = new_mask
+
     res = DotProductAttention(
         q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=rng)
     return res, state
@@ -394,6 +429,26 @@ def do_call(x):  # pylint: disable=invalid-name
     output, vjpfun = jax.vjp(do_call, inputs)
     return output, vjpfun(ct)[0]
 
+  def new_parameters(self, input_shapes, input_dtype, rng):
+    if self._mode in ('train', 'eval'):
+      return (), ()
+
+    assert self._mode == 'predict'
+    # Buffer length is hardcoded for now. TODO(pkozakowski): Pass it from the
+    # model.
+    max_len = 2048
+    ((batch_size, _, _), _, _) = input_shapes
+    def initial_state(shape, dtype):
+      (_, _, depth) = shape
+      return np.zeros((batch_size, max_len, depth), dtype=dtype)
+    (_, k, v) = tuple(
+        initial_state(shape, dtype)
+        for (shape, dtype) in zip(input_shapes, input_dtype)
+    )
+    mask = np.zeros((batch_size, 1, max_len))
+    index = 0
+    return (), (k, v, mask, index)
+
 
 class MemoryEfficientCausalAttention(BaseCausalAttention):
   """Memory-efficient dot product attention.
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 17a94e7b8..7388e8786 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -238,7 +238,7 @@ def TransformerLM(vocab_size,
     dropout: float: dropout rate (how much to drop out)
     share_qk: bool, whether to share queries and keys in decoder attention
     max_len: int: maximum symbol length for positional encoding
-    mode: str: 'train' or 'eval'
+    mode: str: 'train', 'eval' or 'predict', predict mode is for fast inference
 
   Returns:
     A Transformer language model as a layer that maps from a tensor of tokens
@@ -247,10 +247,10 @@ def TransformerLM(vocab_size,
   embedder = [
       tl.Embedding(d_model, vocab_size),
       tl.Dropout(rate=dropout, name='embedding', mode=mode),
-      tl.PositionalEncoding(max_len=max_len),
+      tl.PositionalEncoding(max_len=max_len, mode=mode),
   ]
   return tl.Model(                  # tokens
-      tl.ShiftRight(),              # toks
+      tl.ShiftRight(mode=mode),     # toks
       embedder,                     # vecs
       [DecoderBlock(  # pylint: disable=g-complex-comprehension
           d_model, d_ff, n_heads, d_attention_key, d_attention_value,
diff --git a/tensor2tensor/trax/models/transformer_test.py b/tensor2tensor/trax/models/transformer_test.py
index 8b3a89d9c..5d578d4bc 100644
--- a/tensor2tensor/trax/models/transformer_test.py
+++ b/tensor2tensor/trax/models/transformer_test.py
@@ -19,10 +19,14 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
+
 from absl.testing import absltest
 from absl.testing import parameterized
+import numpy as onp
 from tensor2tensor.trax import backend
 from tensor2tensor.trax import layers as tl
+from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.models import transformer
 
 
@@ -53,6 +57,35 @@ def _test_transformer_forward_shape(self, input_vocab_size,
                              else input_vocab_size]))
     self.assertEqual(expected_shape, final_shape)
 
+  def test_transformer_lm_fast_inference(self):
+    with backend.use_backend('jax'):
+      vocab_size = 16
+      model_fn = functools.partial(
+          transformer.TransformerLM,
+          vocab_size=vocab_size, d_model=4, d_ff=8, n_layers=2, n_heads=2)
+      model_slow = model_fn(mode='eval')
+      model_fast = model_fn(mode='predict')
+      rng = backend.random.get_prng(0)
+      batch_size = 2
+      (params, state_slow) = model_slow.initialize(
+          input_shapes=(batch_size, 1), input_dtype=np.int32, rng=rng)
+      (_, state_fast) = model_fast.initialize(
+          input_shapes=(batch_size, 1), input_dtype=np.int32, rng=rng)
+
+      max_length = 5
+      buf = onp.zeros((batch_size, max_length), dtype=np.int32)
+      next_sym = onp.zeros((batch_size, 1), dtype=onp.int32)
+
+      for index in range(max_length):
+        (logits_slow, state_slow) = model_slow(
+            buf, params=params, state=state_slow, rng=rng)
+        (logits_fast, state_fast) = model_fast(
+            next_sym, params=params, state=state_fast, rng=rng)
+        onp.testing.assert_array_almost_equal(
+            logits_slow[:, index, :], logits_fast[:, 0, :])
+        next_sym = onp.random.randint(vocab_size, size=(batch_size, 1))
+        buf[:, index] = next_sym[:, 0]
+
   @parameterized.named_parameters(
       ('same_vocab', 16, None),
       ('same_size', 16, 16),
diff --git a/tensor2tensor/trax/rl/simple_test.py b/tensor2tensor/trax/rl/simple_test.py
index 957fcf6cf..41d65e5f2 100644
--- a/tensor2tensor/trax/rl/simple_test.py
+++ b/tensor2tensor/trax/rl/simple_test.py
@@ -246,11 +246,10 @@ def _make_env(
 
     gin.bind_parameter("BoxSpaceSerializer.precision", 1)
 
-    seq_length = max_trajectory_length * int(
-        np.prod(observation_space.shape) + np.prod(action_space.shape))
-    predict_output = (np.array([[[0.0]] * seq_length]), ())
+    predict_output = (np.array([[[0.0]]] * batch_size), ())
     mock_model_fn = mock.MagicMock()
     mock_model_fn.return_value.side_effect = itertools.repeat(predict_output)
+    mock_model_fn.return_value.initialize.return_value = ((), ())
 
     return simulated_env_problem.SerializedSequenceSimulatedEnvProblem(
         model=mock_model_fn,
diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index 2dbc752b0..77c62c735 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -57,10 +57,10 @@ def __init__(self, model, batch_size, observation_space, action_space,
         model. The format is implementation-specific.
       output_dir: (str) Output dir.
     """
-    # TODO(pkozakowski): At some point we will have a "predict" mode which we
-    # should use here. When this happens, change the mode.
     self._model = model
-    self._model_predict = backend.jit(self._model(mode="eval"))
+    model_predict = self._model(mode="predict")
+    self._model_predict = backend.jit(model_predict)
+    self._model_initialize = model_predict.initialize
     self._observation_space = observation_space
     self._action_space = action_space
     self._reward_range = reward_range
@@ -372,8 +372,13 @@ def initialize_environments(self, batch_size=1, **kwargs):
     self._steps = np.zeros(batch_size, dtype=np.int32)
     self._last_observations = np.full(
         (batch_size,) + self._observation_space.shape, np.nan)
+    self._last_symbols = np.zeros((batch_size, 1), dtype=np.int32)
     super(SerializedSequenceSimulatedEnvProblem, self).initialize_environments(
         batch_size=batch_size, **kwargs)
+    (subrng, self._rng) = jax_random.split(self._rng)
+    (_, self._init_model_state) = self._model_initialize(
+        input_shapes=(batch_size, 1), input_dtype=np.int32, rng=subrng
+    )
 
   @property
   def _obs_repr_indices(self):
@@ -388,11 +393,12 @@ def _action_repr_indices(self):
   def _predict_obs(self, predict_fn, rng):
     for (i, subrng) in enumerate(jax_random.split(rng, self._obs_repr_length)):
       symbol_index = self._steps * self._step_repr_length + i
-      log_probs, self._model_state = predict_fn(self._history,
-                                                state=self._model_state,
-                                                rng=subrng)
-      log_probs = log_probs[:, symbol_index, :]
-      self._history[:, symbol_index] = utils.gumbel_sample(log_probs)
+      log_probs, self._model_state = predict_fn(
+          self._last_symbols, state=self._model_state, rng=subrng,
+      )
+      log_probs = log_probs
+      self._last_symbols = utils.gumbel_sample(log_probs)
+      self._history[:, symbol_index] = self._last_symbols[:, 0]
 
     obs_repr = self._history[self._obs_repr_indices]
     return self._obs_serializer.deserialize(obs_repr)
@@ -401,6 +407,14 @@ def _reset_model(self, predict_fn, indices, history, rng):
     # TODO(pkozakowski): Random starts.
     del history
 
+    indices = np.array(indices)
+    assert indices.shape[0] in (0, self._history.shape[0]), (
+        # TODO(pkozakowski): Lift this requirement.
+        "Only resetting all envs at once is supported."
+    )
+
+    self._model_state = self._init_model_state
+    self._last_symbols[indices] = 0
     self._steps[indices] = 0
     observation = self._predict_obs(predict_fn, rng)[indices]
     self._last_observations[indices] = observation
diff --git a/tensor2tensor/trax/rl/simulated_env_problem_test.py b/tensor2tensor/trax/rl/simulated_env_problem_test.py
index 38f1a37f4..ac80bb5e2 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem_test.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem_test.py
@@ -130,6 +130,7 @@ def _make_env(
     mock_model_fn = mock.MagicMock()
     if predict_fn is not None:
       mock_model_fn.return_value = predict_fn
+      mock_model_fn.return_value.initialize.return_value = ((), ())
     return simulated_env_problem.SerializedSequenceSimulatedEnvProblem(
         model=mock_model_fn,
         reward_fn=reward_fn,
@@ -168,7 +169,7 @@ def make_prediction(symbol):
       one_hot = np.eye(vocab_size)[symbol]
       log_probs = (1 - one_hot) * -100.0  # Virtually deterministic.
       # (4 obs symbols + 1 action symbol) * 3 timesteps = 15.
-      return np.array([[log_probs] * 15]), ()
+      return np.array([[log_probs]]), ()
 
     mock_predict_fn = mock.MagicMock()
     mock_predict_fn.side_effect = map(make_prediction, symbols)
@@ -187,22 +188,15 @@ def make_prediction(symbol):
           action_space=gym.spaces.Discrete(2),
       )
       obs1 = env.reset()
-      ((inputs,), _) = mock_predict_fn.call_args
 
       act1 = 0
       (obs2, reward, done, _) = env.step(np.array([act1]))
-      ((inputs,), _) = mock_predict_fn.call_args
-      self.assertEqual(inputs[0, 4], act1)
-      np.testing.assert_array_equal(inputs[0, :4], symbols[:4])
       np.testing.assert_array_equal(obs1, obs2)
       np.testing.assert_array_equal(reward, [0.5])
       np.testing.assert_array_equal(done, [False])
 
       act2 = 1
       (obs3, reward, done, _) = env.step(np.array([act2]))
-      ((inputs,), _) = mock_predict_fn.call_args
-      self.assertEqual(inputs[0, 9], act2)
-      np.testing.assert_array_equal(inputs[0, 5:9], symbols[4:8])
       self.assertFalse(np.array_equal(obs2, obs3))
       np.testing.assert_array_equal(reward, [0.5])
       np.testing.assert_array_equal(done, [True])

From 1632111c8dd86a0a65d9c4a231b97214a6662dba Mon Sep 17 00:00:00 2001
From: Shawn Simister <simister@google.com>
Date: Sun, 22 Sep 2019 16:23:07 -0700
Subject: [PATCH 2479/2720] Refactoring to use named tuples in the controller
 to make it easier to read. Also refactoring Neural Stack cell so that it can
 be extended to a double ended queue later on.

PiperOrigin-RevId: 270571957
---
 tensor2tensor/models/research/neural_stack.py | 280 +++++++++++-------
 .../models/research/neural_stack_test.py      | 170 +++++------
 2 files changed, 249 insertions(+), 201 deletions(-)

diff --git a/tensor2tensor/models/research/neural_stack.py b/tensor2tensor/models/research/neural_stack.py
index b347e6032..5c138e026 100644
--- a/tensor2tensor/models/research/neural_stack.py
+++ b/tensor2tensor/models/research/neural_stack.py
@@ -27,6 +27,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
+
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import registry
@@ -34,6 +36,17 @@
 
 import tensorflow as tf
 
+# This is the interface between the RNN controller and the neural stack.
+NeuralStackControllerInterface = collections.namedtuple(
+    "NeuralStackControllerInterface",
+    "push_strengths, pop_strengths, write_values, outputs, state")
+
+# This is recurrent state of the neural stack RNN cell.
+NeuralStackState = collections.namedtuple(
+    "NeuralStackState",
+    "controller_state, read_values, memory_values, read_strengths, " +
+    "write_strengths")
+
 
 class NeuralStackCell(tf.nn.rnn_cell.RNNCell):
   """An RNN cell base class that can implement a stack or queue.
@@ -72,17 +85,17 @@ def state_size(self):
     return (tf.TensorShape([self._num_units]),
             tf.TensorShape([self._num_read_heads, self._embedding_size]),
             tf.TensorShape([self._memory_size, self._embedding_size]),
-            tf.TensorShape([self._num_read_heads, self._memory_size, 1]),
+            tf.TensorShape([1, self._memory_size, 1]),
             tf.TensorShape([self._num_write_heads, self._memory_size, 1]))
 
   @property
   def output_size(self):
-    return tf.TensorShape([self._num_read_heads, self._embedding_size])
+    return tf.TensorShape([1, self._embedding_size])
 
   def initialize_write_strengths(self, batch_size):
     """Initialize write strengths to write to the first memory address.
 
-    This is exposed as it's own function so that it can be overridden to provide
+    This is exposed as its own function so that it can be overridden to provide
     alternate write adressing schemes.
 
     Args:
@@ -104,26 +117,47 @@ def zero_state(self, batch_size, dtype):
       dtype: The default datatype to initialize to.
 
     Returns:
-      (controller_state.shape,
-       read_values.shape,
-       memory_values.shape,
-       read_strengths.shape,
-       write_strengths.shape)
+      A new NeuralStackState tuple.
     """
-    state = list(super(NeuralStackCell, self).zero_state(batch_size, dtype))
-    state[4] = self.initialize_write_strengths(batch_size)
-    return tuple(state)
-
-  def build_read_mask(self):
+    parent_state = NeuralStackState(*super(NeuralStackCell, self).zero_state(
+        batch_size, dtype))
+    return NeuralStackState(
+        controller_state=parent_state.controller_state,
+        read_values=parent_state.read_values,
+        memory_values=parent_state.memory_values,
+        read_strengths=parent_state.read_strengths,
+        write_strengths=self.initialize_write_strengths(batch_size))
+
+  def get_read_mask(self, read_head_index):
     """Creates a mask which allows us to attenuate subsequent read strengths.
 
-    This is exposed as it's own function so that it can be overridden to provide
+    This is exposed as its own function so that it can be overridden to provide
     alternate read adressing schemes.
 
+    Args:
+      read_head_index: Identifies which read head we're getting the mask for.
+
     Returns:
-      A tf.float32 tensor of shape [1, memory_size, memory_size]
+      A tf.float32 tensor of shape [1, 1, memory_size, memory_size]
     """
-    return common_layers.mask_pos_gt(self._memory_size, self._memory_size)
+    return tf.expand_dims(
+        common_layers.mask_pos_lt(self._memory_size, self._memory_size), axis=0)
+
+  def get_write_head_offset(self, write_head_index):
+    """Lookup the offset to shift the write head at each step.
+
+    By default, we move each write head forward by 1.
+
+    This is exposed as its own function so that it can be overridden to provide
+    alternate write adressing schemes.
+
+    Args:
+      write_head_index: Identifies which write head we're getting the index for.
+
+    Returns:
+      An integer offset to move the write head at each step.
+    """
+    return 1
 
   def add_scalar_projection(self, name, size):
     """A helper function for mapping scalar controller outputs.
@@ -193,145 +227,165 @@ def build_controller(self):
           "output", self._num_read_heads)
 
   def build(self, _):
-    """Build the controller, read mask and write shift convolutional filter.
-
-    The write shift convolutional filter is a simple 3x3 convolution which is
-    used to advance the read heads to the next memory address at each step. This
-    filter can be changed to move the read heads in other ways.
+    """Build the controller.
     """
-    self.read_mask = self.build_read_mask()
-    self.write_shift_convolution = tf.reshape(tf.one_hot([[3]], depth=9),
-                                              shape=[3, 3, 1, 1])
     self.build_controller()
-
     self.built = True
 
-  def call_controller(self, inputs, state, batch_size):
+  def get_controller_shape(self, batch_size):
+    """Define the output shapes of the neural stack controller.
+
+    Making this a separate functions so that it can be used in unit tests.
+
+    Args:
+      batch_size: The size of the current batch of data.
+
+    Returns:
+      A tuple of shapes for each output returned from the controller.
+    """
+    return (
+        # push_strengths,
+        [batch_size, self._num_write_heads, 1, 1],
+        # pop_strengths
+        [batch_size, self._num_write_heads, 1, 1],
+        # write_values
+        [batch_size, self._num_write_heads, self._embedding_size],
+        # outputs
+        [batch_size, 1, self._embedding_size],
+        # state
+        [batch_size, self._num_units])
+
+  def call_controller(self, inputs, read_values, prev_state, batch_size):
     """Make a call to the neural stack controller.
 
     See Section 3.1 of Grefenstette et al., 2015.
 
     Args:
-      inputs: The combined inputs to the controller consisting of the current
-         input value concatenated with the read values from the previous
-         timestep with shape [batch_size, (num_write_heads + num_read_heads)
-         * embedding_size].
-      state: The hidden state from the previous time step.
+      inputs: The inputs to the neural stack cell should be a tf.float32 tensor
+        with shape [batch_size, num_write_heads, embedding_size]
+      read_values: The values of the read heads at the previous timestep.
+      prev_state: The hidden state from the previous time step.
       batch_size: The size of the current batch of input values.
 
     Returns:
-      A tuple of outputs and the new hidden state value:
-      (push_strengths, pop_strengths, write_values, outputs, state)
+      A tuple of outputs and the new NeuralStackControllerInterface.
     """
     with tf.name_scope("controller"):
+      # Concatenate the current input values with the read value from the
+      # previous timestep before feeding them into the controller.
+      controller_inputs = tf.concat([
+          tf.contrib.layers.flatten(inputs),
+          tf.contrib.layers.flatten(read_values),
+      ], axis=1)
+
       rnn_input = tf.tanh(tf.nn.bias_add(tf.matmul(
-          inputs, self._input_proj), self._input_bias))
+          controller_inputs, self._input_proj), self._input_bias))
 
-      (rnn_output, state) = self.rnn(rnn_input, state)
+      (rnn_output, state) = self.rnn(rnn_input, prev_state)
 
-      push_strengths = tf.reshape(
-          tf.sigmoid(tf.nn.bias_add(tf.matmul(
-              rnn_output, self._push_proj), self._push_bias)),
-          shape=[batch_size, self._num_write_heads, 1, 1])
+      push_strengths = tf.sigmoid(tf.nn.bias_add(tf.matmul(
+          rnn_output, self._push_proj), self._push_bias))
 
-      pop_strengths = tf.reshape(
-          tf.sigmoid(tf.nn.bias_add(tf.matmul(
-              rnn_output, self._pop_proj), self._pop_bias)),
-          shape=[batch_size, self._num_write_heads, 1, 1])
+      pop_strengths = tf.sigmoid(tf.nn.bias_add(tf.matmul(
+          rnn_output, self._pop_proj), self._pop_bias))
 
-      write_values = tf.reshape(
-          tf.tanh(tf.nn.bias_add(tf.matmul(
-              rnn_output, self._value_proj), self._value_bias)),
-          shape=[batch_size, self._num_read_heads, self._embedding_size])
+      write_values = tf.tanh(tf.nn.bias_add(tf.matmul(
+          rnn_output, self._value_proj), self._value_bias))
 
-      outputs = tf.reshape(
-          tf.tanh(tf.nn.bias_add(tf.matmul(
-              rnn_output, self._output_proj), self._output_bias)),
-          shape=[batch_size, self._num_read_heads, self._embedding_size])
+      outputs = tf.tanh(tf.nn.bias_add(tf.matmul(
+          rnn_output, self._output_proj), self._output_bias))
 
-    return push_strengths, pop_strengths, write_values, outputs, state
+      # Reshape all the outputs according to the shapes specified by
+      # get_controller_shape()
+      projected_outputs = [push_strengths,
+                           pop_strengths,
+                           write_values,
+                           outputs,
+                           state]
+      next_state = [
+          tf.reshape(output, shape=output_shape) for output, output_shape
+          in zip(projected_outputs, self.get_controller_shape(batch_size))]
+      return NeuralStackControllerInterface(*next_state)
 
-  def call(self, inputs, state):
+  def call(self, inputs, prev_state):
     """Evaluates one timestep of the current neural stack cell.
 
     See section 3.4 of Grefenstette et al., 2015.
 
     Args:
       inputs: The inputs to the neural stack cell should be a tf.float32 tensor
-        with shape [batch_size, max_timesteps, 1, embedding_size]
-      state: The tuple of state values from the previous timestep.
+        with shape [batch_size, embedding_size]
+      prev_state: The NeuralStackState from the previous timestep.
 
     Returns:
-      The output value of the stack as well as the new tuple of state values.
-      (outputs, (controller_state, read_values, memory_values, read_strengths,
-                 write_strengths))
+      A tuple of the output of the stack as well as the new NeuralStackState.
     """
-    (controller_state,
-     read_values,
-     memory_values,
-     read_strengths,
-     write_strengths) = state
-
     batch_size = tf.shape(inputs)[0]
 
-    # Concatenate the current input value with the read value from  the previous
-    # timestep before feeding them into the controller.
-    controller_inputs = tf.concat([
-        tf.reshape(
-            read_values,
-            shape=[batch_size, self._num_read_heads * self._embedding_size]),
-        tf.reshape(
-            inputs,
-            shape=[batch_size, self._num_write_heads * self._embedding_size])
-    ], axis=1)
-
     # Call the controller and get controller interface values.
-    with tf.control_dependencies([read_strengths]):
-      (push_strengths, pop_strengths,
-       write_values, outputs, controller_state) = self.call_controller(
-           controller_inputs, controller_state, batch_size)
+    with tf.control_dependencies([prev_state.read_strengths]):
+      controller_output = self.call_controller(
+          inputs, prev_state.read_values, prev_state.controller_state,
+          batch_size)
 
     # Always write input values to memory regardless of push strength.
     # See Equation-1 in Grefenstette et al., 2015.
-    memory_values += tf.reduce_sum(
-        tf.expand_dims(write_values, axis=1) * write_strengths, axis=1)
+    new_memory_values = prev_state.memory_values + tf.reduce_sum(
+        tf.expand_dims(controller_output.write_values, axis=1) *
+        prev_state.write_strengths,
+        axis=1)
 
     # Attenuate the read strengths of existing memory values depending on the
     # current pop strength.
     # See Equation-2 in Grefenstette et al., 2015.
-    read_strengths = tf.nn.relu(
-        read_strengths - tf.nn.relu(pop_strengths - tf.reduce_sum(
-            tf.reshape(read_strengths,
-                       shape=[batch_size, 1, 1, self._memory_size]) *
-            self.read_mask, axis=3, keepdims=True)))
-
-    # Set read strength for the current timestep based on the push strength.
-    read_strengths = read_strengths + push_strengths * write_strengths
+    new_read_strengths = prev_state.read_strengths
+    for h in range(self._num_read_heads - 1, -1, -1):
+      new_read_strengths = tf.nn.relu(
+          new_read_strengths -
+          tf.nn.relu(controller_output.pop_strengths - tf.expand_dims(
+              tf.reduce_sum(new_read_strengths * self.get_read_mask(h), axis=2),
+              axis=3)))
+
+    # Combine all write heads and their associated push values into a single set
+    # of read weights.
+    new_read_strengths += tf.reduce_sum(
+        controller_output.push_strengths * prev_state.write_strengths,
+        axis=1, keep_dims=True)
 
     # Calculate the "top" value of the stack by looking at read strengths.
     # See Equation-3 in Grefenstette et al., 2015.
-    read_values = tf.reduce_sum(
+    new_read_values = tf.reduce_sum(
         tf.minimum(
-            read_strengths,
-            tf.nn.relu(1 - tf.reshape(
-                tf.reduce_sum(read_strengths * self.read_mask,
-                              axis=2,
-                              keepdims=True),
-                shape=[
-                    batch_size, self._num_read_heads, self._memory_size, 1
-                ]))) * tf.expand_dims(memory_values, axis=1),
+            new_read_strengths,
+            tf.nn.relu(1 - tf.expand_dims(
+                tf.reduce_sum(
+                    new_read_strengths * tf.concat([
+                        self.get_read_mask(h)
+                        for h in range(self._num_read_heads)
+                    ], axis=1),
+                    axis=2),
+                axis=3))
+        ) * tf.expand_dims(new_memory_values, axis=1),
         axis=2)
 
-    # Shift the write strengths forward by one memory address for the next step.
-    write_strengths = tf.nn.conv2d(
-        write_strengths, self.write_shift_convolution, [1, 1, 1, 1],
-        padding="SAME")
+    # Temporarily split write strengths apart so they can be shifted in
+    # different directions.
+    write_strengths_by_head = tf.split(prev_state.write_strengths,
+                                       self._num_write_heads,
+                                       axis=1)
+    # Shift the write strengths for each write head in the direction indicated
+    # by get_write_head_offset().
+    new_write_strengths = tf.concat([
+        tf.roll(write_strength, shift=self.get_write_head_offset(h), axis=2)
+        for h, write_strength in enumerate(write_strengths_by_head)
+    ], axis=1)
 
-    return (outputs, (controller_state,
-                      read_values,
-                      memory_values,
-                      read_strengths,
-                      write_strengths))
+    return (controller_output.outputs, NeuralStackState(
+        controller_state=controller_output.state,
+        read_values=new_read_values,
+        memory_values=new_memory_values,
+        read_strengths=new_read_strengths,
+        write_strengths=new_write_strengths))
 
 
 class NeuralQueueCell(NeuralStackCell):
@@ -340,13 +394,17 @@ class NeuralQueueCell(NeuralStackCell):
   See section 3.2 of Grefenstette et al., 2015.
   """
 
-  def build_read_mask(self):
+  def get_read_mask(self, read_head_index):
     """Uses mask_pos_lt() instead of mask_pos_gt() to reverse read values.
 
+    Args:
+      read_head_index: Identifies which read head we're getting the mask for.
+
     Returns:
-      A tf.float32 tensor of shape [1, memory_size, memory_size].
+      A tf.float32 tensor of shape [1, 1, memory_size, memory_size].
     """
-    return common_layers.mask_pos_lt(self._memory_size, self._memory_size)
+    return tf.expand_dims(
+        common_layers.mask_pos_gt(self._memory_size, self._memory_size), axis=0)
 
 
 @registry.register_model
@@ -357,7 +415,7 @@ class NeuralStackModel(t2t_model.T2TModel):
   def cell(self, hidden_size):
     """Build an RNN cell.
 
-    This is exposed as it's own function so that it can be overridden to provide
+    This is exposed as its own function so that it can be overridden to provide
     different types of RNN cells.
 
     Args:
diff --git a/tensor2tensor/models/research/neural_stack_test.py b/tensor2tensor/models/research/neural_stack_test.py
index 79a31662b..e32001580 100644
--- a/tensor2tensor/models/research/neural_stack_test.py
+++ b/tensor2tensor/models/research/neural_stack_test.py
@@ -40,123 +40,89 @@ def build_fake_controller(cell):
       trainable=False)
 
 
-def call_fake_controller(push_values, pop_values, read_values, output_values):
+def call_fake_controller(push_values, pop_values, write_values, output_values):
   """Mock a RNN controller from a set of expected outputs.
 
   Args:
     push_values: Expected controller push values.
     pop_values: Expected controller pop values.
-    read_values: Expected controller read values.
+    write_values: Expected controller write values.
     output_values: Expected controller output values.
 
   Returns:
     A callable which behaves like the call method of an NeuralStackCell.
   """
-  def call(cell, inputs, state, batch_size):
+  def call(cell, inputs, prev_read_values, controller_state, batch_size):
     del inputs
+    del prev_read_values
     del batch_size
-    next_step = tf.assign_add(cell.current_step, tf.constant(1))
-    return (
-        tf.slice(tf.constant(push_values), [next_step, 0], [1, -1]),
-        tf.slice(tf.constant(pop_values), [next_step, 0], [1, -1]),
-        tf.slice(tf.constant(read_values), [next_step, 0, 0], [1, -1, -1]),
-        tf.slice(tf.constant(output_values), [next_step, 0, 0], [1, -1, -1]),
-        state
+    next_step = tf.constant(0)
+    if hasattr(cell, "current_step"):
+      next_step = tf.assign_add(cell.current_step, tf.constant(1))
+    return neural_stack.NeuralStackControllerInterface(
+        push_strengths=tf.slice(tf.constant(push_values),
+                                [next_step, 0, 0, 0],
+                                [1, -1, -1, -1]),
+        pop_strengths=tf.slice(tf.constant(pop_values),
+                               [next_step, 0, 0, 0],
+                               [1, -1, -1, -1]),
+        write_values=tf.slice(tf.constant(write_values),
+                              [next_step, 0, 0],
+                              [1, -1, -1]),
+        outputs=tf.slice(tf.constant(output_values),
+                         [next_step, 0, 0],
+                         [1, -1, -1]),
+        state=controller_state
     )
   return call
 
 
+def assert_controller_shapes(test, controller_outputs, controller_shapes):
+  for name, output, shape in zip(controller_outputs._fields, controller_outputs,
+                                 controller_shapes):
+    test.assertEqual(shape, output.shape, "%s shapes don't match" % name)
+
+
+def assert_cell_shapes(test, controller_outputs, zero_state):
+  for name, output, state in zip(controller_outputs._fields, controller_outputs,
+                                 zero_state):
+    test.assertEqual(state.shape, output.shape, "%s shapes don't match" % name)
+
+
 class NeuralStackCellTest(tf.test.TestCase):
 
-  def test_controller_shapes(self):
+  def test_cell_shapes(self):
     """Check that all the NeuralStackCell tensor shapes are correct.
     """
-
     batch_size = 5
     embedding_size = 3
     memory_size = 6
     num_units = 8
 
     stack = neural_stack.NeuralStackCell(num_units, memory_size, embedding_size)
-
     stack.build(None)
 
-    self.assertEqual([1, embedding_size], stack.output_size)
-    self.assertEqual([1, memory_size, memory_size], stack.read_mask.shape)
-    self.assertEqual([3, 3, 1, 1], stack.write_shift_convolution.shape)
+    self.assertEqual([1, 1, memory_size, memory_size],
+                     stack.get_read_mask(0).shape)
 
     stack_input = tf.zeros([batch_size, 1, embedding_size], dtype=tf.float32)
-
     zero_state = stack.zero_state(batch_size, tf.float32)
-
-    (controller_state,
-     previous_values,
-     memory_values,
-     read_strengths,
-     write_strengths) = zero_state
-
-    self.assertEqual([batch_size, num_units], controller_state.shape)
-    self.assertEqual([batch_size, 1, embedding_size], previous_values.shape)
-    self.assertEqual([batch_size, memory_size, embedding_size],
-                     memory_values.shape)
-    self.assertEqual([batch_size, 1, memory_size, 1], read_strengths.shape)
-    self.assertEqual([batch_size, 1, memory_size, 1], write_strengths.shape)
-
-    rnn_input = tf.concat([
-        tf.reshape(
-            previous_values,
-            shape=[batch_size, embedding_size]),
-        tf.reshape(
-            stack_input,
-            shape=[batch_size, embedding_size])
-    ], axis=1)
-    self.assertEqual([batch_size, 2 * embedding_size], rnn_input.shape)
-
-    (push_strengths,
-     pop_strengths,
-     new_values,
-     outputs,
-     controller_next_state) = stack.call_controller(rnn_input,
-                                                    controller_state,
-                                                    batch_size)
-
-    self.assertEqual([batch_size, 1, 1, 1], push_strengths.shape)
-    self.assertEqual([batch_size, 1, 1, 1], pop_strengths.shape)
-    self.assertEqual([batch_size, 1, embedding_size], new_values.shape)
-    self.assertEqual([batch_size, 1, embedding_size], outputs.shape)
-    self.assertEqual([batch_size, num_units], controller_next_state.shape)
-
-    (outputs, (controller_next_state,
-               read_values,
-               next_memory_values,
-               next_read_strengths,
-               next_write_strengths)) = stack.call(stack_input, zero_state)
-
-    self.assertEqual([batch_size, 1, embedding_size], outputs.shape)
-    self.assertEqual([batch_size, num_units], controller_next_state.shape)
-    self.assertEqual([batch_size, 1, embedding_size], read_values.shape)
-    self.assertEqual([batch_size, memory_size, embedding_size],
-                     next_memory_values.shape)
-    self.assertEqual([batch_size, 1, memory_size, 1], next_read_strengths.shape)
-    self.assertEqual([batch_size, 1, memory_size, 1],
-                     next_write_strengths.shape)
+    (outputs, (stack_next_state)) = stack.call(stack_input, zero_state)
 
     # Make sure that stack output shapes match stack input shapes
-    self.assertEqual(controller_next_state.shape, controller_state.shape)
-    self.assertEqual(read_values.shape, previous_values.shape)
-    self.assertEqual(next_memory_values.shape, memory_values.shape)
-    self.assertEqual(next_read_strengths.shape, read_strengths.shape)
-    self.assertEqual(next_write_strengths.shape, write_strengths.shape)
+    self.assertEqual(outputs.shape, stack_input.shape)
+
+    assert_cell_shapes(self, stack_next_state, zero_state)
 
   @mock.patch.object(neural_stack.NeuralStackCell, "build_controller",
                      build_fake_controller)
   @mock.patch.object(neural_stack.NeuralStackCell, "call_controller",
                      call_fake_controller(
-                         push_values=[[1.0], [1.0], [0.0]],
-                         pop_values=[[0.0], [0.0], [1.0]],
-                         read_values=[[[1.0, 0.0, 0.0]],
-                                      [[0.0, 1.0, 0.0]],
-                                      [[0.0, 0.0, 1.0]]],
+                         push_values=[[[[1.0]]], [[[1.0]]], [[[0.0]]]],
+                         pop_values=[[[[0.0]]], [[[0.0]]], [[[1.0]]]],
+                         write_values=[[[1.0, 0.0, 0.0]],
+                                       [[0.0, 1.0, 0.0]],
+                                       [[0.0, 0.0, 1.0]]],
                          output_values=[[[0.0, 0.0, 0.0]],
                                         [[0.0, 0.0, 0.0]],
                                         [[0.0, 0.0, 0.0]]]))
@@ -179,8 +145,20 @@ def test_push_pop(self):
         [[[0.0], [0.0], [0.], [1.0], [0.0], [0.0]]]])
     expected_top = np.array([[[1.0, 0.0, 0.0]]])
 
-    stack = neural_stack.NeuralStackCell(8, 6, 3)
+    batch_size = 1
+    embedding_size = 3
+    memory_size = 6
+    num_units = 8
+
+    stack = neural_stack.NeuralStackCell(num_units, memory_size, embedding_size)
     stack_input = tf.constant(input_values, dtype=tf.float32)
+
+    stack_zero_state = tf.zeros([batch_size, num_units])
+    controller_outputs = stack.call_controller(None, None, stack_zero_state,
+                                               batch_size)
+    assert_controller_shapes(self, controller_outputs,
+                             stack.get_controller_shape(batch_size))
+
     (outputs, state) = tf.nn.dynamic_rnn(cell=stack,
                                          inputs=stack_input,
                                          time_major=False,
@@ -191,10 +169,10 @@ def test_push_pop(self):
       _, state_vals = sess.run([outputs, state])
       (_, stack_top, values, read_strengths, write_strengths) = state_vals
 
-      self.assertAllClose(expected_top, stack_top)
       self.assertAllClose(expected_values, values)
-      self.assertAllClose(expected_read_strengths, read_strengths)
       self.assertAllClose(expected_write_strengths, write_strengths)
+      self.assertAllClose(expected_read_strengths, read_strengths)
+      self.assertAllClose(expected_top, stack_top)
 
 
 class NeuralQueueCellTest(tf.test.TestCase):
@@ -203,11 +181,11 @@ class NeuralQueueCellTest(tf.test.TestCase):
                      build_fake_controller)
   @mock.patch.object(neural_stack.NeuralQueueCell, "call_controller",
                      call_fake_controller(
-                         push_values=[[1.0], [1.0], [0.0]],
-                         pop_values=[[0.0], [0.0], [1.0]],
-                         read_values=[[[1.0, 0.0, 0.0]],
-                                      [[0.0, 1.0, 0.0]],
-                                      [[0.0, 0.0, 1.0]]],
+                         push_values=[[[[1.0]]], [[[1.0]]], [[[0.0]]]],
+                         pop_values=[[[[0.0]]], [[[0.0]]], [[[1.0]]]],
+                         write_values=[[[1.0, 0.0, 0.0]],
+                                       [[0.0, 1.0, 0.0]],
+                                       [[0.0, 0.0, 1.0]]],
                          output_values=[[[0.0, 0.0, 0.0]],
                                         [[0.0, 0.0, 0.0]],
                                         [[0.0, 0.0, 0.0]]]))
@@ -229,8 +207,20 @@ def test_enqueue_dequeue(self):
         [[[0.0], [0.0], [0.0], [1.0], [0.0], [0.0]]]])
     expected_front = np.array([[[0.0, 1.0, 0.0]]])
 
-    queue = neural_stack.NeuralQueueCell(8, 6, 3)
+    batch_size = 1
+    num_units = 8
+    embedding_size = 3
+    memory_size = 6
+
+    queue = neural_stack.NeuralQueueCell(num_units, memory_size, embedding_size)
     rnn_input = tf.constant(input_values, dtype=tf.float32)
+
+    queue_zero_state = tf.zeros([batch_size, num_units])
+    controller_outputs = queue.call_controller(None, None, queue_zero_state,
+                                               batch_size)
+    assert_controller_shapes(self, controller_outputs,
+                             queue.get_controller_shape(batch_size))
+
     (outputs, state) = tf.nn.dynamic_rnn(cell=queue,
                                          inputs=rnn_input,
                                          time_major=False,
@@ -241,10 +231,10 @@ def test_enqueue_dequeue(self):
       _, state_vals = sess.run([outputs, state])
       (_, queue_front, values, read_strengths, write_strengths) = state_vals
 
-      self.assertAllClose(expected_front, queue_front)
       self.assertAllClose(expected_values, values)
-      self.assertAllClose(expected_read_strengths, read_strengths)
       self.assertAllClose(expected_write_strengths, write_strengths)
+      self.assertAllClose(expected_read_strengths, read_strengths)
+      self.assertAllClose(expected_front, queue_front)
 
 
 class NeuralStackModelTest(tf.test.TestCase):

From 4fce0b79c5500a01c81925959a1915b696be3dce Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 24 Sep 2019 11:20:16 -0700
Subject: [PATCH 2480/2720] Update PolicySchedule to control arbitrary
 nontrainable parameters.

PiperOrigin-RevId: 270947079
---
 tensor2tensor/trax/learning_rate.py      | 57 ++++++++++++++----------
 tensor2tensor/trax/learning_rate_test.py |  6 +--
 2 files changed, 36 insertions(+), 27 deletions(-)

diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index 0ff4ed801..7b047a420 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -151,56 +151,59 @@ def PolicySchedule(
         ("eval", "metrics/accuracy"),
         ("eval", "metrics/loss"),
     ),
-    include_lr_in_observation=False,
-    observation_range=(0.0, 5.0),
-    start_lr=0.001,
-    max_lr=10.0,
+    include_controls_in_observation=False,
+    control_configs=(
+        # (name, start, (low, high), flip)
+        ("learning_rate", 1e-3, (1e-9, 10.0), False),
+    ),
+    metric_range=(0.0, 5.0),
     action_multipliers=(1.0 / 1.5, 1.0 / 1.25, 1.0, 1.25, 1.5),
     policy_and_value_model=trax_models.FrameStackMLP,
     policy_and_value_two_towers=False,
     policy_and_value_vocab_size=None,
     policy_dir=gin.REQUIRED,
+    temperature=1.0,
 ):
   """Learning rate schedule controlled by a learned policy.
 
   Args:
     history: the history of training and evaluation (History object).
     observation_metrics: list of pairs (mode, metric), as in the History object.
-    include_lr_in_observation: bool, whether to include the learning rate in
+    include_controls_in_observation: bool, whether to include the controls in
       observations.
-    observation_range: tuple (low, high), range to clip the observation to.
-    start_lr: starting learning rate.
-    max_lr: maximum value to clip the learning rate to.
+    control_configs: control configs, see trax.rl.envs.OnlineTuneEnv.
+    metric_range: tuple (low, high), range to clip the metrics to.
     action_multipliers: sequence of LR multipliers that policy actions
       correspond to.
     policy_and_value_model: Trax model to use as the policy.
     policy_and_value_two_towers: bool, whether the action distribution and value
       prediction is computed by separate model towers.
-    policy_and_value_vocab_size: Vocabulary size of a policy and value network
+    policy_and_value_vocab_size: vocabulary size of a policy and value network
       operating on serialized representation. If None, use raw continuous
       representation.
     policy_dir: directory with the policy checkpoint.
+    temperature: temperature for sampling from the policy.
 
   Returns:
-    a function learning_rate(step): float -> {"learning_rate": float}, the
-    step-dependent lr.
+    a function nontrainable_params(step): float -> {"name": float}, the
+    step-dependent schedule for nontrainable parameters.
   """
 
   # Turn the history into observations for the policy. If we don't have any,
   # return the initial learning rate.
   start_time = time.time()
-  lr_config = ("learning_rate", start_lr, (1e-9, max_lr), False)
-  if include_lr_in_observation:
-    control_configs = (lr_config,)
-  else:
-    control_configs = None
   observations = online_tune.history_to_observations(
-      history, observation_metrics, observation_range, control_configs
+      history, observation_metrics, metric_range,
+      control_configs if include_controls_in_observation else None
   )
   logging.vlog(
       1, "Building observations took %0.2f sec.", time.time() - start_time)
   if observations.shape[0] == 0:
-    return lambda _: start_lr
+    controls = {
+        name: start_value
+        for (name, start_value, _, _) in control_configs
+    }
+    return lambda _: controls
 
   # Build the policy network and load its parameters.
   start_time = time.time()
@@ -235,10 +238,16 @@ def PolicySchedule(
       1, "Running the policy took %0.2f sec.", time.time() - start_time
   )
   # Sample from the action distribution for the last timestep.
-  action = utils.gumbel_sample(log_probs[0, -1, :])
-
-  # Get a new learning rate.
-  new_lr = online_tune.update_control(
-      lr_config, action.item(), history, action_multipliers
+  action = utils.gumbel_sample(
+      log_probs[0, -len(control_configs):, :] / temperature
   )
-  return lambda _: {"learning_rate": new_lr}
+
+  # Get new controls.
+  controls = {
+      # name: value
+      control_config[0]: online_tune.update_control(  # pylint: disable=g-complex-comprehension
+          control_config, control_action, history, action_multipliers
+      )
+      for (control_action, control_config) in zip(action, control_configs)
+  }
+  return lambda _: controls
diff --git a/tensor2tensor/trax/learning_rate_test.py b/tensor2tensor/trax/learning_rate_test.py
index d2fabea05..a480d0e65 100644
--- a/tensor2tensor/trax/learning_rate_test.py
+++ b/tensor2tensor/trax/learning_rate_test.py
@@ -59,9 +59,9 @@ def _make_schedule(
     return learning_rate.PolicySchedule(
         history,
         observation_metrics=observation_metrics,
-        include_lr_in_observation=False,
+        include_controls_in_observation=False,
         action_multipliers=action_multipliers,
-        start_lr=start_lr,
+        control_configs=(("learning_rate", start_lr, (1e-9, 1.0), False),),
         policy_and_value_model=policy_and_value_model,
         policy_and_value_two_towers=False,
         policy_dir=policy_dir,
@@ -70,7 +70,7 @@ def _make_schedule(
   def test_returns_start_lr_when_there_are_no_metrics(self):
     history = trax_history.History()
     schedule = self._make_schedule(history, start_lr=1e-3)
-    self.assertEqual(schedule(0), 1e-3)
+    self.assertEqual(schedule(0)["learning_rate"], 1e-3)
 
   def test_changes_lr_when_there_are_some_metrics(self):
     history = trax_history.History()

From 9dff225df2296d748cdf5856e5c9bba1e64f090a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 25 Sep 2019 00:57:05 -0700
Subject: [PATCH 2481/2720] Add mixture of Gaussian PDF computations (both
 diagonal and full).

PiperOrigin-RevId: 271071908
---
 tensor2tensor/trax/layers/core.py      | 39 ++++++++++++++++++++++++++
 tensor2tensor/trax/layers/core_test.py | 19 +++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index fc4d8f2c1..7f1128c56 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -227,3 +227,42 @@ def one_hot(x, size, dtype=np.float32):  # pylint: disable=invalid-name
 def Mean(x, params, axis=-1, keepdims=False, **kwargs):
   del params, kwargs
   return np.mean(x, axis=axis, keepdims=keepdims)
+
+
+def log_gaussian_pdf(x, mu, sigma):  # pylint: disable=invalid-name
+  """Compute log N(x | mu, sigma)."""
+  a = mu.shape[-1] * np.log(2 * np.pi)
+  _, b = np.linalg.slogdet(sigma)
+  y = np.linalg.solve(sigma, x - mu)
+  y = np.expand_dims(y, axis=-1)
+  xm = np.expand_dims(x - mu, axis=-2)
+  c = np.matmul(xm, y)
+  c = np.squeeze(np.squeeze(c, axis=-1), axis=-1)
+  return -0.5 * (a + b + c)
+
+
+def log_gaussian_diag_pdf(x, mu, diag_sigma):  # pylint: disable=invalid-name
+  """Compute log N(x | mu, eye(diag_sigma))."""
+  a = mu.shape[-1] * np.log(2 * np.pi)
+  b = np.sum(np.log(diag_sigma), axis=-1)
+  y = x - mu / diag_sigma
+  y = np.expand_dims(y, axis=-1)
+  xm = np.expand_dims(x - mu, axis=-2)
+  c = np.matmul(xm, y)
+  c = np.squeeze(np.squeeze(c, axis=-1), axis=-1)
+  return -0.5 * (a + b + c)
+
+
+def multigaussian_loss(preds, targets, ngauss=1):  # pylint: disable=invalid-name
+  """Compute mixture of gaussians loss."""
+  ndims = targets.shape[-1]
+  logits = preds[:, :ngauss]
+  mus = preds[:, ngauss:ngauss*(ndims + 1)]
+  sigmas = preds[:, ngauss(ndims + 1):]
+  sigmas = sigmas * sigmas + 1e-6  # Make positive.
+  loglogits = logits - backend.logsumexp(logits, axis=-1, keepdims=True)
+  mus = np.reshape(mus, [-1, ngauss, ndims])
+  sigmas = np.reshape(sigmas, [-1, ngauss, ndims])
+  targets = np.reshape(targets, [-1, 1, ndims])
+  glogprobs = log_gaussian_diag_pdf(targets, mus, sigmas)
+  return backend.logsumexp(loglogits + glogprobs, axis=-1)
diff --git a/tensor2tensor/trax/layers/core_test.py b/tensor2tensor/trax/layers/core_test.py
index e2b625807..1d0d4e734 100644
--- a/tensor2tensor/trax/layers/core_test.py
+++ b/tensor2tensor/trax/layers/core_test.py
@@ -101,6 +101,25 @@ def test_dropout(self):
         core.Dropout(rate=0.1, mode="eval"), input_shape)
     self.assertEqual(final_shape, output_shape)
 
+  def test_log_gaussian_pdf(self):
+    x = onp.zeros((2, 5), dtype=onp.float32)
+    mu = x
+    dsigma = onp.eye(5)[None, :, :]
+    sigma = onp.concatenate([dsigma, 2*dsigma], axis=0)
+    prob = core.log_gaussian_pdf(x, mu, sigma)
+    self.assertEqual(prob.shape, (2,))
+    self.assertEqual(int(prob[0]), -4)
+    self.assertEqual(int(prob[1]), -6)
+
+  def test_log_gaussian_diag_pdf(self):
+    x = onp.zeros((2, 5), dtype=onp.float32)
+    mu = x
+    sigma = onp.ones((5,))[None, :]
+    sigma = onp.concatenate([sigma, 2*sigma], axis=0)
+    prob = core.log_gaussian_diag_pdf(x, mu, sigma)
+    self.assertEqual(prob.shape, (2,))
+    self.assertEqual(int(prob[0]), -4)
+    self.assertEqual(int(prob[1]), -6)
 
 if __name__ == "__main__":
   absltest.main()

From b1c43708005bad463902b1210a99079d04186a21 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 25 Sep 2019 11:12:39 -0700
Subject: [PATCH 2482/2720] Clarify Layer.__call__ by extracting a method for
 custom gradients.

PiperOrigin-RevId: 271163869
---
 tensor2tensor/trax/layers/base.py | 99 ++++++++++++++++---------------
 1 file changed, 51 insertions(+), 48 deletions(-)

diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 6908cd6ad..b2b553d6e 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -257,59 +257,62 @@ def __call__(self, x, params=(), state=(), **kwargs):
 
       if not self.has_custom_grad or Layer._STASH_IN is not None:
         return self.call(x, params=params, state=state, **kwargs)
+      else:
+        return self._do_custom_gradients(x, params, state, **kwargs)
 
-      # Custom gradients part.
-      assert backend.get_name() == 'jax', (
-          'Custom gradients are only supported in JAX for now.')
+    except Exception:
+      name, trace = self.__class__.__name__, _short_traceback()
+      raise LayerError(name, 'call', self._caller, shapes(x), trace)
 
-      # TODO(wangpeng): JAX doesn't support custom grads for functions with
-      #   auxiliary output yet (https://github.com/google/jax/issues/844). Will
-      #   remove the constraints on state below when this feature is added to
-      #   JAX.
+  def _do_custom_gradients(self, x, params, state, **kwargs):
+    """Calls this layer for a forward pass, but with custom gradients."""
+    assert backend.get_name() == 'jax', (
+        'Custom gradients are only supported in JAX for now.')
 
-      assert not jax.tree_util.tree_leaves(state), (
-          'Custom gradients require trivial start state. Got %s' % str(state))
-
-      def check_end_state(output_state):
-        output, state = output_state
-        assert not jax.tree_util.tree_leaves(state), (
-            'Custom gradients require trivial end state. Got %s' % str(state))
-        return output
-
-      # See this link for how custom transformations are defined in JAX:
-      # https://jax.readthedocs.io/en/latest/jax.html#jax.custom_transforms
-      # Note that we capture the kwargs and don't calculate gradients wrt. them.
-      @jax.custom_transforms
-      def do_call(y, params):
-        return check_end_state(self.call(y, params=params, state=state,
-                                         **kwargs))
+    # TODO(wangpeng): JAX doesn't support custom grads for functions with
+    #   auxiliary output yet (https://github.com/google/jax/issues/844). Will
+    #   remove the constraints on state below when this feature is added to
+    #   JAX.
 
-      # This is the custom gradient (vector-jacobian product in JAX) function.
-      # For the exact specification of this custom transformation see this link:
-      # https://jax.readthedocs.io/en/latest/jax.html#jax.defjvp_all
-      def do_call_vjp(y, params):
-        """Custom gradient (vjp) function."""
-        stash = None
-        if Layer._STASH_IN is None:
-          Layer._STASH_IN = stash = {}
-        output = check_end_state(self.call(y, params=params, state=state,
-                                           **kwargs))
-        if stash is not None:
-          Layer._STASH_IN = None
-        def vjpfun(grad):
-          assert Layer._STASH_OUT is None
-          Layer._STASH_OUT = stash
-          res = self.custom_grad(y, output, grad, params, state, **kwargs)
-          Layer._STASH_OUT = None
-          return res
-        return output, vjpfun
-
-      jax.defvjp_all(do_call, do_call_vjp)
-      return do_call(x, params), state
+    assert not jax.tree_util.tree_leaves(state), (
+        'Custom gradients require trivial start state. Got %s' % str(state))
 
-    except Exception:
-      name, trace = self.__class__.__name__, _short_traceback()
-      raise LayerError(name, 'call', self._caller, shapes(x), trace)
+    def check_end_state(output_state):
+      output, state = output_state
+      assert not jax.tree_util.tree_leaves(state), (
+          'Custom gradients require trivial end state. Got %s' % str(state))
+      return output
+
+    # See this link for how custom transformations are defined in JAX:
+    # https://jax.readthedocs.io/en/latest/jax.html#jax.custom_transforms
+    # Note that we capture the kwargs and don't calculate gradients wrt. them.
+    @jax.custom_transforms
+    def do_call(y, params):
+      return check_end_state(self.call(y, params=params, state=state,
+                                       **kwargs))
+
+    # This is the custom gradient (vector-jacobian product in JAX) function.
+    # For the exact specification of this custom transformation see this link:
+    # https://jax.readthedocs.io/en/latest/jax.html#jax.defjvp_all
+    def do_call_vjp(y, params):
+      """Custom gradient (vjp) function."""
+      stash = None
+      if Layer._STASH_IN is None:
+        Layer._STASH_IN = stash = {}
+      output = check_end_state(self.call(y, params=params, state=state,
+                                         **kwargs))
+      if stash is not None:
+        Layer._STASH_IN = None
+      def vjpfun(grad):
+        assert Layer._STASH_OUT is None
+        Layer._STASH_OUT = stash
+        res = self.custom_grad(y, output, grad, params, state, **kwargs)
+        Layer._STASH_OUT = None
+        return res
+      return output, vjpfun
+
+    jax.defvjp_all(do_call, do_call_vjp)
+    return do_call(x, params), state
 
 
 class LayerError(Exception):

From 049b9d8fe681989ad69383ee04fb32b321b4f564 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Wed, 25 Sep 2019 11:46:57 -0700
Subject: [PATCH 2483/2720] Update Transformer copy config to make it easier to
 play with settings.

PiperOrigin-RevId: 271171680
---
 .../trax/configs/transformer_copy.gin         | 32 +++++++++++--------
 1 file changed, 18 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_copy.gin b/tensor2tensor/trax/configs/transformer_copy.gin
index 15c67e650..a7982405f 100644
--- a/tensor2tensor/trax/configs/transformer_copy.gin
+++ b/tensor2tensor/trax/configs/transformer_copy.gin
@@ -3,11 +3,15 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
+n_symbols = 128
+length = 1024
+batch = 32
+
 # Parameters for batch_fun:
 # ==============================================================================
-batch_fun.batch_size_per_device = 32
-batch_fun.eval_batch_size = 32
-batch_fun.max_eval_length = 1024
+batch_fun.batch_size_per_device = %batch
+batch_fun.eval_batch_size = %batch
+batch_fun.max_eval_length = %length
 
 # Parameters for inputs:
 # ==============================================================================
@@ -22,10 +26,10 @@ MultifactorSchedule.warmup_steps = 8000
 
 # Parameters for sequence_copy_inputs:
 # ==============================================================================
-sequence_copy_inputs.vocab_size = 128
-sequence_copy_inputs.batch_size = 32
-sequence_copy_inputs.train_lengths = [1024]
-sequence_copy_inputs.eval_lengths = [1024]
+sequence_copy_inputs.vocab_size = %n_symbols
+sequence_copy_inputs.batch_size = %batch
+sequence_copy_inputs.train_lengths = [%length]
+sequence_copy_inputs.eval_lengths = [%length]
 sequence_copy_inputs.reverse = False
 
 # Parameters for train:
@@ -35,7 +39,7 @@ train.eval_steps = 64
 train.inputs = @trax.inputs.sequence_copy_inputs
 train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.Adafactor
-train.train_steps = 50000
+train.train_steps = 150000
 train.has_weights = True
 
 # Parameters for MemoryEfficientCausalAttention:
@@ -48,13 +52,13 @@ MemoryEfficientCausalAttention.loop_stride = 512
 MergedMultiHashedCausalAttentionV2.allow_duplicate_attention = False
 MergedMultiHashedCausalAttentionV2.attend_across_buckets = False
 MergedMultiHashedCausalAttentionV2.rehash_each_round = True
-# MergedMultiHashedCausalAttentionV2.n_bins: 32
-# MergedMultiHashedCausalAttentionV2.n_buckets: 64
-# MergedMultiHashedCausalAttentionV2.n_hashes: 4
+MergedMultiHashedCausalAttentionV2.n_bins = 32
+MergedMultiHashedCausalAttentionV2.n_buckets = 64
+MergedMultiHashedCausalAttentionV2.n_hashes = 4
 MergedMultiHashedCausalAttentionV2.one_rng = False
 MergedMultiHashedCausalAttentionV2.hard_k = 0
 MergedMultiHashedCausalAttentionV2.dropout = 0.0
-# MergedMultiHashedCausalAttentionV2.drop_for_hash_rate: 0.0
+MergedMultiHashedCausalAttentionV2.drop_for_hash_rate = 0.0
 
 # Parameters for TransformerLM:
 # ==============================================================================
@@ -64,9 +68,9 @@ TransformerLM.d_attention_value = 64
 TransformerLM.d_model = 256
 TransformerLM.d_ff = 256
 TransformerLM.dropout = 0.0
-TransformerLM.max_len = 1024
+TransformerLM.max_len = %length
 TransformerLM.mode = 'train'
 TransformerLM.n_heads = 4
 TransformerLM.n_layers = 1
 TransformerLM.share_qk = True
-TransformerLM.vocab_size = 128
+TransformerLM.vocab_size = %n_symbols

From 05f222d27a4885550450d9ba26987f78af5f9ecd Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 26 Sep 2019 11:13:57 -0700
Subject: [PATCH 2484/2720] Adjust learning rate of ResNet to correct for the
 change in Momentum optimizer.

PiperOrigin-RevId: 271390165
---
 tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin         | 2 +-
 tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
index 94e9f6b37..86415eb1a 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
@@ -18,7 +18,7 @@ inputs.dataset_name = 't2t_image_imagenet224'
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-EvalAdjustingSchedule.constant = 1.0
+EvalAdjustingSchedule.constant = 0.2
 MultifactorSchedule.factors = 'constant * linear_warmup'
 MultifactorSchedule.warmup_steps = 400
 
diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
index 6ecd54a04..d4085acea 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
@@ -18,7 +18,7 @@ inputs.dataset_name = 't2t_image_imagenet224'
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-EvalAdjustingSchedule.constant = 1.0
+EvalAdjustingSchedule.constant = 0.2
 MultifactorSchedule.factors = 'constant * linear_warmup'
 MultifactorSchedule.warmup_steps = 400
 

From c290b168333b16b3a6f1d9e9e1e1acc7354d9ffb Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 27 Sep 2019 09:17:14 -0700
Subject: [PATCH 2485/2720] Rename trainer-internal function to clarify its
 intent.

PiperOrigin-RevId: 271584428
---
 tensor2tensor/trax/trax.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 88eaabd52..e16e27969 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -624,8 +624,9 @@ def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs,
         model_input_shape, lambda x: x if x else 1)
     model_target_shape = layers.nested_map(
         model_target_shape, lambda x: x if x else 1)
-    def initialize(input_shape, input_dtype, target_shape, target_dtype, rng):
-      """Helper to initialize the model."""
+    def new_opt_state_and_model_state(input_shape, input_dtype, target_shape,
+                                      target_dtype, rng):
+      """Returns optimizer and model states suitable for training a model."""
       # Combine inputs and targets on the stack.
       if not isinstance(input_dtype, (list, tuple)):
         input_dtype = [input_dtype]
@@ -643,10 +644,12 @@ def initialize(input_shape, input_dtype, target_shape, target_dtype, rng):
       return (OptState(params, slots, opt_params), state)
     if _is_jit_init():
       # JIT parameter initialization to avoid memory fragmentation
-      initialize = backend.jit(initialize, static_argnums=(0, 1, 2, 3))
-    self._initialize = lambda: initialize(  # pylint: disable=g-long-lambda
-        model_input_shape, self._inputs.input_dtype,
-        model_target_shape, self._inputs.target_dtype, init_rng)
+      new_opt_state_and_model_state = backend.jit(new_opt_state_and_model_state,
+                                                  static_argnums=(0, 1, 2, 3))
+    self._new_opt_state_and_model_state = (
+        lambda: new_opt_state_and_model_state(  # pylint: disable=g-long-lambda
+            model_input_shape, self._inputs.input_dtype,
+            model_target_shape, self._inputs.target_dtype, init_rng))
 
     # jit model_predict and update so they're fast
     self._jit_model_predict_eval = _jit_predict_fn(
@@ -712,7 +715,7 @@ def reset(self, output_dir):
       opt_state = state.opt_state
       model_state = state.model_state
     else:
-      opt_state, model_state = self._initialize()
+      opt_state, model_state = self._new_opt_state_and_model_state()
       model_state = layers.nested_map(
           model_state, self._maybe_replicate)
     self._opt_state = OptState(*layers.nested_map(

From 5ba0a4e838c211226686204bdedd0fb6c0067b22 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 27 Sep 2019 13:57:49 -0700
Subject: [PATCH 2486/2720] Run PPO optimization on minibatches, so we can use
 heavier policy networks.

PiperOrigin-RevId: 271641725
---
 tensor2tensor/trax/rl/ppo.py              | 14 ++++++++++++
 tensor2tensor/trax/rl/ppo_test.py         | 26 +++++++++++++++++++++++
 tensor2tensor/trax/rl/ppo_trainer.py      | 26 +++++++++++++++--------
 tensor2tensor/trax/rl/ppo_trainer_test.py | 16 ++++++++++++--
 4 files changed, 71 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index 48b676028..e25836a3b 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -53,6 +53,7 @@
 
 import collections
 import functools
+import itertools
 import os
 import re
 import time
@@ -953,3 +954,16 @@ def write_eval_reward_summaries(reward_stats_by_mode, summary_writer, epoch):
           "Epoch [% 6d] Policy Evaluation (%s reward) "
           "[temperature %.2f] = %10.2f (+/- %.2f)", epoch, reward_mode,
           temperature, reward_stats["mean"], reward_stats["std"])
+
+
+def shuffled_index_batches(dataset_size, batch_size):
+  """Generates batches of shuffled indices over a dataset."""
+  def shuffled_indices():
+    while True:
+      perm = onp.random.permutation(dataset_size)
+      for x in perm:
+        yield x
+
+  indices = shuffled_indices()
+  while True:
+    yield onp.array(list(itertools.islice(indices, int(batch_size))))
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
index 850bc6329..05af01074 100644
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import functools
+import itertools
 
 import jax
 from jax import random as jax_random
@@ -603,6 +604,31 @@ def test_inits_policy_by_world_model_checkpoint(self):
     observations = np.zeros((1, 100), dtype=np.int32)
     policy(observations, new_policy_params, state=policy_state, rng=rng)
 
+  def test_shuffled_index_batches_generates_valid_batch(self):
+    dataset_size = 16
+    batch_size = 4
+    stream = ppo.shuffled_index_batches(dataset_size, batch_size)
+    batch = next(stream)
+    self.assertEqual(batch.shape, (batch_size,))
+    # Assert that all indices are different.
+    self.assertEqual(len(set(batch)), batch_size)
+
+  def test_shuffled_index_batches_generates_all_indices(self):
+    dataset_size = 16
+    batch_size = 4
+    stream = ppo.shuffled_index_batches(dataset_size, batch_size)
+    indices = np.reshape(
+        list(itertools.islice(stream, dataset_size // batch_size)), -1
+    )
+    self.assertEqual(set(indices), set(range(dataset_size)))
+
+  def test_shuffled_index_batches_gives_different_permutations(self):
+    dataset_size = 256
+    batch_size = 8
+    stream1 = ppo.shuffled_index_batches(dataset_size, batch_size)
+    stream2 = ppo.shuffled_index_batches(dataset_size, batch_size)
+    self.assertFalse(np.array_equal(next(stream1), next(stream2)))
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index e4d5c54d4..84e7c642c 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -64,6 +64,7 @@ def __init__(self,
                policy_and_value_two_towers=False,
                policy_and_value_vocab_size=None,
                n_optimizer_steps=N_OPTIMIZER_STEPS,
+               optimizer_batch_size=64,
                print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
                target_kl=0.01,
                boundary=20,
@@ -98,6 +99,7 @@ def __init__(self,
         operating on serialized representation. If None, use raw continuous
         representation.
       n_optimizer_steps: Number of optimizer steps.
+      optimizer_batch_size: Batch size of an optimizer step.
       print_every_optimizer_steps: How often to log during the policy
         optimization process.
       target_kl: Policy iteration early stopping. Set to infinity to disable
@@ -134,6 +136,7 @@ def __init__(self,
     super(PPO, self).__init__(train_env, eval_env, output_dir, **kwargs)
 
     self._n_optimizer_steps = n_optimizer_steps
+    self._optimizer_batch_size = optimizer_batch_size
     self._print_every_optimizer_steps = print_every_optimizer_steps
     self._target_kl = target_kl
     self._boundary = boundary
@@ -538,10 +541,14 @@ def train_epoch(self, evaluate=True):
     optimization_start_time = time.time()
     keys = jax_random.split(key1, num=self._n_optimizer_steps)
     opt_step = 0
-    for key in keys:
+    opt_batch_size = min(self._optimizer_batch_size, B)
+    index_batches = ppo.shuffled_index_batches(
+        dataset_size=B, batch_size=opt_batch_size
+    )
+    for (index_batch, key) in zip(index_batches, keys):
       k1, k2, k3 = jax_random.split(key, num=3)
       t = time.time()
-      # Update the optimizer state.
+      # Update the optimizer state on the sampled minibatch.
       self._policy_and_value_opt_state, self._model_state = (
           ppo.policy_and_value_opt_step(
               # We pass the optimizer slots between PPO epochs, so we need to
@@ -557,13 +564,13 @@ def train_epoch(self, evaluate=True):
               self._policy_and_value_opt_update,
               self._policy_and_value_get_params,
               self._policy_and_value_net_apply,
-              log_probabs_traj,
-              value_predictions_traj,
-              padded_observations,
-              padded_actions,
+              log_probabs_traj[index_batch],
+              value_predictions_traj[index_batch],
+              padded_observations[index_batch],
+              padded_actions[index_batch],
               self._rewards_to_actions,
-              padded_rewards,
-              reward_mask,
+              padded_rewards[index_batch],
+              reward_mask[index_batch],
               c1=self._c1,
               c2=self._c2,
               gamma=self._gamma,
@@ -573,7 +580,8 @@ def train_epoch(self, evaluate=True):
       opt_step += 1
       self._total_opt_step += 1
 
-      # Compute the approx KL for early stopping.
+      # Compute the approx KL for early stopping. Use the whole dataset - as we
+      # only do inference, it should fit in the memory.
       (log_probab_actions_new, _), self._model_state = (
           self._policy_and_value_net_apply(
               padded_observations,
diff --git a/tensor2tensor/trax/rl/ppo_trainer_test.py b/tensor2tensor/trax/rl/ppo_trainer_test.py
index c93804a8d..e67df34b1 100644
--- a/tensor2tensor/trax/rl/ppo_trainer_test.py
+++ b/tensor2tensor/trax/rl/ppo_trainer_test.py
@@ -46,7 +46,9 @@
 
 class PpoTrainerTest(test.TestCase):
 
-  def get_wrapped_env(self, name="CartPole-v0", max_episode_steps=2):
+  def get_wrapped_env(
+      self, name="CartPole-v0", max_episode_steps=2, batch_size=1
+  ):
     wrapper_fn = functools.partial(
         gym_utils.gym_env_wrapper,
         **{
@@ -59,7 +61,7 @@ def get_wrapped_env(self, name="CartPole-v0", max_episode_steps=2):
         })
 
     return gym_env_problem.GymEnvProblem(base_env_name=name,
-                                         batch_size=1,
+                                         batch_size=batch_size,
                                          env_wrapper_fn=wrapper_fn,
                                          discrete_rewards=False)
 
@@ -278,6 +280,16 @@ def test_training_loop_cartpole_serialized(self):
       )
       trainer.training_loop(n_epochs=2)
 
+  def test_training_loop_cartpole_minibatch(self):
+    with self.tmp_dir() as output_dir:
+      trainer = self._make_trainer(
+          train_env=self.get_wrapped_env("CartPole-v0", 2, batch_size=4),
+          eval_env=self.get_wrapped_env("CartPole-v0", 2),
+          output_dir=output_dir,
+          optimizer_batch_size=2,
+      )
+      trainer.training_loop(n_epochs=2)
+
 
 if __name__ == "__main__":
   test.main()

From e29caddf0b6e5afe2494b038a4fa8a8aae825591 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 27 Sep 2019 14:20:27 -0700
Subject: [PATCH 2487/2720] Initialize the policy trainer in SimPLe lazily.

PiperOrigin-RevId: 271646475
---
 tensor2tensor/trax/rl/simple_trainer.py | 33 ++++++++++++++++---------
 1 file changed, 22 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index bad077cc5..d1cde1644 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -62,7 +62,11 @@ def __init__(self,
     super(SimPLe, self).__init__(train_env, eval_env, output_dir, **kwargs)
     self._policy_dir = os.path.join(output_dir, "policy")
     self._model_dir = os.path.join(output_dir, "model")
-    self._policy_trainer = policy_trainer_class(
+    # Initialize the policy trainer lazily, so in case of initializing the
+    # policy from world model checkpoint, the trainer will try to load the
+    # checkpoint _after_ it's been created in train_model().
+    self._policy_trainer_fn = functools.partial(
+        policy_trainer_class,
         train_env=train_env,
         eval_env=eval_env,
         output_dir=self._policy_dir,
@@ -71,6 +75,7 @@ def __init__(self,
             self._model_dir if init_policy_from_world_model else None
         ),
     )
+    self._policy_trainer = None
     self._n_real_epochs = n_real_epochs
     self._model_train_batch_size = model_train_batch_size
     self._n_model_initial_train_steps = n_model_initial_train_steps
@@ -113,6 +118,12 @@ def __init__(self,
     self._policy_epoch = 0
     self._model_train_step = 0
 
+  @property
+  def policy_trainer(self):
+    if self._policy_trainer is None:
+      self._policy_trainer = self._policy_trainer_fn()
+    return self._policy_trainer
+
   @property
   def epoch(self):
     return self._simple_epoch
@@ -136,7 +147,7 @@ def train_epoch(self, evaluate=True):
     self._simple_epoch += 1
 
   def evaluate(self):
-    self._policy_trainer.evaluate()
+    self.policy_trainer.evaluate()
 
   def save(self):
     # Nothing to do, as we save stuff continuously.
@@ -149,11 +160,11 @@ def collect_trajectories(self, evaluate):
     logging.info("SimPLe epoch [% 6d]: collecting data.", self._simple_epoch)
     start_time = time.time()
 
-    self._policy_trainer.train_env = self.train_env
-    self._policy_trainer.trajectory_dump_dir = os.path.join(
+    self.policy_trainer.train_env = self.train_env
+    self.policy_trainer.trajectory_dump_dir = os.path.join(
         self._trajectory_dump_root_dir, str(self.epoch))
     self._policy_epoch += self._n_real_epochs
-    self._policy_trainer.training_loop(self._policy_epoch, evaluate=evaluate)
+    self.policy_trainer.training_loop(self._policy_epoch, evaluate=evaluate)
 
     logging.vlog(
         1, "Collecting trajectories took %0.2f sec.", time.time() - start_time)
@@ -205,15 +216,15 @@ def train_policy(self):
         history_stream=itertools.repeat(None),
     )
     # We never want async mode in the simulated env.
-    original_async_mode = self._policy_trainer.async_mode
-    self._policy_trainer.async_mode = False
-    self._policy_trainer.train_env = self._sim_env
+    original_async_mode = self.policy_trainer.async_mode
+    self.policy_trainer.async_mode = False
+    self.policy_trainer.train_env = self._sim_env
     # Don't dump trajectories from the simulated environment.
-    self._policy_trainer.trajectory_dump_dir = None
+    self.policy_trainer.trajectory_dump_dir = None
     self._policy_epoch += self._n_simulated_epochs
-    self._policy_trainer.training_loop(self._policy_epoch, evaluate=False)
+    self.policy_trainer.training_loop(self._policy_epoch, evaluate=False)
     # Revert back to the original async mode in the policy trainer.
-    self._policy_trainer.async_mode = original_async_mode
+    self.policy_trainer.async_mode = original_async_mode
 
     logging.vlog(
         1, "Training policy took %0.2f sec.", time.time() - start_time)

From a80c332c4efb28720cdf4dc47b2ff3e140921d14 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 27 Sep 2019 14:58:14 -0700
Subject: [PATCH 2488/2720] Cache the initial trajectories in memory to save
 time on loading them.

PiperOrigin-RevId: 271654265
---
 tensor2tensor/trax/rl/simple_trainer.py | 43 ++++++++++++++++---------
 1 file changed, 28 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index d1cde1644..01b512345 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -90,6 +90,7 @@ def __init__(self,
           overwrite=True,
       )
     self._initial_model = initial_model
+    self._initial_trajectories = None
 
     self._sim_env = simulated_env_problem_class(
         batch_size=None,
@@ -237,11 +238,27 @@ def _has_own_data(self):
   def _has_initial_data(self):
     return self._initial_trajectory_dir is not None
 
-  def _make_input_streams(self):
+  def _load_trajectories(self, initial):
+    # Cache the initial trajectories in memory, as loading them can take a lot
+    # of time and they don't change.
+    if initial:
+      if self._initial_trajectories is not None:
+        return self._initial_trajectories
+      trajectory_dir = self._initial_trajectory_dir
+    else:
+      trajectory_dir = self._trajectory_dump_root_dir
+
+    trajectories = simple.load_trajectories(
+        trajectory_dir, self._data_eval_frac
+    )
+
+    if initial:
+      self._initial_trajectories = trajectories
+    return trajectories
 
-    def make_example_streams(trajectory_dir):
-      (train_trajs, eval_trajs) = simple.load_trajectories(
-          trajectory_dir, eval_frac=self._data_eval_frac)
+  def _make_input_streams(self):
+    def make_example_streams(initial):
+      (train_trajs, eval_trajs) = self._load_trajectories(initial)
       generate_examples = functools.partial(
           simple.generate_examples,
           trajectory_to_training_examples_fn=(
@@ -257,8 +274,7 @@ def make_example_streams(trajectory_dir):
     if self._has_initial_data:
       start_time = time.time()
       # Load the initial, precollected data.
-      (init_train_stream,
-       init_eval_stream) = make_example_streams(self._initial_trajectory_dir)
+      (init_train_stream, init_eval_stream) = make_example_streams(initial=True)
       logging.vlog(
           1, "Loading initial trajectories took %0.2f sec.",
           time.time() - start_time
@@ -270,8 +286,7 @@ def make_example_streams(trajectory_dir):
     if self._has_own_data:
       start_time = time.time()
       # Load trajectories collected in all epochs so far.
-      (own_train_stream,
-       own_eval_stream) = make_example_streams(self._trajectory_dump_root_dir)
+      (own_train_stream, own_eval_stream) = make_example_streams(initial=False)
       logging.vlog(
           1, "Loading own trajectories took %0.2f sec.",
           time.time() - start_time
@@ -302,13 +317,11 @@ def evaluate_model(self):
         history_stream=itertools.repeat(None),
     )
 
-    if self._has_own_data:
-      trajectory_dir = self._trajectory_dump_root_dir
-    else:
-      trajectory_dir = self._initial_trajectory_dir
-
-    (_, eval_trajectories) = simple.load_trajectories(
-        trajectory_dir, eval_frac=self._data_eval_frac)
+    (_, eval_trajectories) = self._load_trajectories(
+        # If we have any trajectories collected in this run, evaluate on them.
+        # Otherwise, use the initial dataset.
+        initial=(not self._has_own_data)
+    )
     chosen_trajectories = [
         random.choice(eval_trajectories)
         for _ in range(self._sim_env.batch_size)

From df87f90f61812e75356100bf9eb85c54f3035cc2 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 27 Sep 2019 14:58:26 -0700
Subject: [PATCH 2489/2720] Zero out all observations if we get any NaNs.

Now we are NaN-proof!

PiperOrigin-RevId: 271654300
---
 tensor2tensor/trax/rl/online_tune.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/trax/rl/online_tune.py b/tensor2tensor/trax/rl/online_tune.py
index a470f71dd..1252370d2 100644
--- a/tensor2tensor/trax/rl/online_tune.py
+++ b/tensor2tensor/trax/rl/online_tune.py
@@ -25,9 +25,14 @@
 def historical_metric_values(history, metric):
   """Converts a metric stream from a trax History object into a numpy array."""
   metric_sequence = history.get(*metric)
-  return np.array([
+  metric_values = np.array([
       metric_value for (_, metric_value) in metric_sequence
   ])
+  if np.any(np.isnan(metric_values)):
+    # Zero out all observations if any element is NaN. This way the agent
+    # doesn't get any rewards, so it learns to avoid those regions.
+    metric_values[:] = 0.0
+  return metric_values
 
 
 def metric_to_observation(metric_values, metric_range):

From 40bdf03e5b1025eb78cca6d19908cf7c08f35f91 Mon Sep 17 00:00:00 2001
From: Thang Luong <thangluong@google.com>
Date: Fri, 27 Sep 2019 15:02:42 -0700
Subject: [PATCH 2490/2720] Add a unit test for attention_bias_local.

PiperOrigin-RevId: 271655306
---
 tensor2tensor/layers/common_attention_test.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index f747059b8..250e0c864 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -33,6 +33,21 @@
 
 class CommonAttentionTest(parameterized.TestCase, tf.test.TestCase):
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testAttentionBiasLocal(self):
+    length = 5
+    bias = common_attention.attention_bias_local(length, 0, 0)
+    # For length = 5
+    # [[[[-0.e+00 -1.e+09 -1.e+09 -1.e+09 -1.e+09]
+    #    [-1.e+09 -0.e+00 -1.e+09 -1.e+09 -1.e+09]
+    #    [-1.e+09 -1.e+09 -0.e+00 -1.e+09 -1.e+09]
+    #    [-1.e+09 -1.e+09 -1.e+09 -0.e+00 -1.e+09]
+    #    [-1.e+09 -1.e+09 -1.e+09 -1.e+09 -0.e+00]]]]
+    res = self.evaluate(bias)
+    expected_res = -1e9 * np.ones((length, length)) - -1e9 * np.identity(length)
+    expected_res = np.reshape(expected_res, (1, 1, length, length))
+    self.assertAllClose(res, expected_res)
+
   @test_utils.run_in_graph_and_eager_modes()
   def testAddPositionalEmbedding(self):
     x = np.random.rand(5, 3, 12)

From b9c512922db0103d495c7c935546f8d3861a9fce Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 27 Sep 2019 15:21:48 -0700
Subject: [PATCH 2491/2720] Update the config for SimPLe and solve name
 conflicts using gin scopes.

PiperOrigin-RevId: 271658940
---
 .../trax/rl/configs/simple_online_tune.gin    | 50 +++++++++++--------
 tensor2tensor/trax/rl/simple_trainer.py       | 16 +++---
 2 files changed, 38 insertions(+), 28 deletions(-)

diff --git a/tensor2tensor/trax/rl/configs/simple_online_tune.gin b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
index 8378311d7..abcdd0717 100644
--- a/tensor2tensor/trax/rl/configs/simple_online_tune.gin
+++ b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
@@ -11,9 +11,9 @@ BoxSpaceSerializer.precision = 2
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-MultifactorSchedule.constant = 3.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 10000
+world_model/MultifactorSchedule.constant = 3.0
+world_model/MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+world_model/MultifactorSchedule.warmup_steps = 10000
 
 # Parameters for Adam:
 # ==============================================================================
@@ -26,22 +26,23 @@ Adam.weight_decay_rate = 0.0
 # ==============================================================================
 TransformerDecoder.d_model = 64
 TransformerDecoder.d_ff = 128
-TransformerDecoder.dropout = 0.0
+TransformerDecoder.n_layers = 2
 TransformerDecoder.n_heads = 2
-TransformerDecoder.n_layers = 1
+TransformerDecoder.dropout = 0.0
 
 # Parameters for PPO:
 # ==============================================================================
 PPO.n_optimizer_steps = 10
+PPO.optimizer_batch_size = 128
 PPO.target_kl = 0.1
-PPO.boundary = 128
-PPO.max_timestep = 128
-PPO.max_timestep_eval = 128
+PPO.boundary = 100
+PPO.max_timestep = 100
+PPO.max_timestep_eval = 100
 PPO.random_seed = None
 PPO.gamma = 1.0
 PPO.lambda_ = 0.95
 PPO.c1 = 1.0
-PPO.c2 = 0.01
+PPO.c2 = 0.1
 PPO.done_frac_for_policy_save = 0
 PPO.len_history_for_policy = None
 PPO.separate_eval = False
@@ -49,10 +50,16 @@ PPO.save_every_n = 1
 PPO.policy_and_value_model = @trax.models.TransformerDecoder
 PPO.policy_and_value_optimizer = @trax.optimizers.Adam
 PPO.trajectory_dump_min_count_per_shard = 8
+PPO.print_every_optimizer_steps = 1
+
+## Parameters for MemoryEfficientCausalAttention:
+## ==============================================================================
+world_model/MemoryEfficientCausalAttention.dropout = 0.3
+world_model/MemoryEfficientCausalAttention.loop_stride = 2
 
 # Parameters for SerializedSequenceSimulatedEnvProblem:
 # ==============================================================================
-SerializedSequenceSimulatedEnvProblem.model = @trax.models.TransformerLM
+SerializedSequenceSimulatedEnvProblem.model = @world_model/trax.models.TransformerLM
 SerializedSequenceSimulatedEnvProblem.reward_fn = @trax.rl.onlinetune_reward_fn
 SerializedSequenceSimulatedEnvProblem.done_fn = @trax.rl.onlinetune_done_fn
 SerializedSequenceSimulatedEnvProblem.vocab_size = 128
@@ -68,20 +75,21 @@ SimPLe.n_model_train_steps_per_epoch = 20000
 SimPLe.model_train_batch_size = 64
 SimPLe.simulated_env_problem_class = @trax.rl.SerializedSequenceSimulatedEnvProblem
 SimPLe.simulated_batch_size = 128
-SimPLe.n_simulated_epochs = 30
-SimPLe.initial_trajectory_mix_prob = 0.5
+SimPLe.n_simulated_epochs = 50
+SimPLe.initial_trajectory_mix_prob = 0.9
+SimPLe.init_policy_from_world_model = False
 
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.d_model = 512
-TransformerLM.d_ff = 1024
-TransformerLM.n_layers = 3
-TransformerLM.n_heads = 4
-TransformerLM.dropout = 0.3
-TransformerLM.max_len = 1024
+world_model/TransformerLM.attention_type = @world_model/trax.layers.MemoryEfficientCausalAttention
+world_model/TransformerLM.d_model = 256
+world_model/TransformerLM.d_ff = 512
+world_model/TransformerLM.n_layers = 3
+world_model/TransformerLM.n_heads = 4
+world_model/TransformerLM.dropout = 0.3
+world_model/TransformerLM.max_len = 2048
 
 # Parameters for train:
 # ==============================================================================
-train.eval_frequency = 1000
-train.optimizer = @trax.optimizers.Adafactor
-
+world_model/train.eval_frequency = 1000
+world_model/train.optimizer = @trax.optimizers.Adafactor
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
index 01b512345..c5edc996a 100644
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ b/tensor2tensor/trax/rl/simple_trainer.py
@@ -26,6 +26,7 @@
 import time
 
 from absl import logging
+import gin
 from matplotlib import pyplot as plt
 from tensor2tensor.trax import inputs as trax_inputs
 from tensor2tensor.trax import jaxboard
@@ -196,13 +197,14 @@ def train_model(self):
     else:
       train_steps = self._n_model_train_steps_per_epoch
     self._model_train_step += train_steps
-    state = trax.train(
-        model=self._sim_env.model,
-        inputs=inputs,
-        train_steps=self._model_train_step,
-        output_dir=self._model_dir,
-        has_weights=True,
-    )
+    with gin.config_scope("world_model"):
+      state = trax.train(
+          model=self._sim_env.model,
+          inputs=inputs,
+          train_steps=self._model_train_step,
+          output_dir=self._model_dir,
+          has_weights=True,
+      )
 
     logging.vlog(
         1, "Training model took %0.2f sec.", time.time() - start_time)

From 87e0f6207cacf74cd734e9ed197b7687746f93d3 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 27 Sep 2019 17:18:35 -0700
Subject: [PATCH 2492/2720] Pass the correct number of controls to the policy
 network in PolicySchedule.

PiperOrigin-RevId: 271678993
---
 tensor2tensor/trax/learning_rate.py      | 11 ++++++-
 tensor2tensor/trax/learning_rate_test.py | 37 ++++++++++++++++++++----
 tensor2tensor/trax/rl/ppo.py             | 14 +++++----
 3 files changed, 50 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index 7b047a420..f45f7188a 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -22,6 +22,9 @@
 function that takes a step and returns a dict with entry "learning_rate".
 """
 
+# TODO(pkozakowski): Revisit the decision to control nontrainable parameters
+# using LR schedules, or at least rename the module.
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
@@ -205,10 +208,13 @@ def PolicySchedule(
     }
     return lambda _: controls
 
+  assert policy_and_value_vocab_size is None, (
+      "Serialized policies are not supported yet."
+  )
   # Build the policy network and load its parameters.
   start_time = time.time()
   net = ppo.policy_and_value_net(
-      n_controls=1,
+      n_controls=len(control_configs),
       n_actions=len(action_multipliers),
       vocab_size=policy_and_value_vocab_size,
       bottom_layers_fn=policy_and_value_model,
@@ -238,6 +244,9 @@ def PolicySchedule(
       1, "Running the policy took %0.2f sec.", time.time() - start_time
   )
   # Sample from the action distribution for the last timestep.
+  assert log_probs.shape == (
+      1, len(control_configs) * observations.shape[0], len(action_multipliers)
+  )
   action = utils.gumbel_sample(
       log_probs[0, -len(control_configs):, :] / temperature
   )
diff --git a/tensor2tensor/trax/learning_rate_test.py b/tensor2tensor/trax/learning_rate_test.py
index a480d0e65..4bc0aafd8 100644
--- a/tensor2tensor/trax/learning_rate_test.py
+++ b/tensor2tensor/trax/learning_rate_test.py
@@ -36,14 +36,14 @@ class PolicyScheduleTest(test.TestCase):
   def _make_schedule(
       self,
       history,
-      start_lr=1e-3,
+      control_configs,
       observation_metrics=(("eval", "metrics/accuracy"),),
       action_multipliers=(1.0,),
   ):
     policy_and_value_model = atari_cnn.FrameStackMLP
     net = ppo.policy_and_value_net(
         n_actions=len(action_multipliers),
-        n_controls=1,
+        n_controls=len(control_configs),
         vocab_size=None,
         bottom_layers_fn=policy_and_value_model,
         two_towers=False,
@@ -61,7 +61,7 @@ def _make_schedule(
         observation_metrics=observation_metrics,
         include_controls_in_observation=False,
         action_multipliers=action_multipliers,
-        control_configs=(("learning_rate", start_lr, (1e-9, 1.0), False),),
+        control_configs=control_configs,
         policy_and_value_model=policy_and_value_model,
         policy_and_value_two_towers=False,
         policy_dir=policy_dir,
@@ -69,8 +69,12 @@ def _make_schedule(
 
   def test_returns_start_lr_when_there_are_no_metrics(self):
     history = trax_history.History()
-    schedule = self._make_schedule(history, start_lr=1e-3)
-    self.assertEqual(schedule(0)["learning_rate"], 1e-3)
+    start_lr = 1e-3
+    schedule = self._make_schedule(
+        history,
+        control_configs=(("learning_rate", start_lr, (1e-9, 1.0), False),),
+    )
+    self.assertEqual(schedule(0)["learning_rate"], start_lr)
 
   def test_changes_lr_when_there_are_some_metrics(self):
     history = trax_history.History()
@@ -80,6 +84,7 @@ def test_changes_lr_when_there_are_some_metrics(self):
     )
     schedule = self._make_schedule(
         history,
+        control_configs=(("learning_rate", 1e-3, (1e-9, 1.0), False),),
         observation_metrics=(("eval", "metrics/accuracy"),),
         action_multipliers=(0.5, 2.0),
     )
@@ -88,6 +93,28 @@ def test_changes_lr_when_there_are_some_metrics(self):
         onp.allclose(new_lr, 5e-5) or onp.allclose(new_lr, 2e-4)
     )
 
+  def test_works_with_multiple_controls(self):
+    history = trax_history.History()
+    history.append("eval", "metrics/accuracy", step=0, value=0.8)
+    history.append(
+        *online_tune.control_metric("learning_rate"), step=0, value=1e-4
+    )
+    history.append(
+        *online_tune.control_metric("weight_decay_rate"), step=0, value=1e-5
+    )
+    schedule = self._make_schedule(
+        history,
+        observation_metrics=(("eval", "metrics/accuracy"),),
+        control_configs=(
+            ("learning_rate", 1e-3, (1e-9, 1.0), False),
+            ("weight_decay_rate", 1e-5, (1e-9, 1.0), False),
+        ),
+        action_multipliers=(1.0,),
+    )
+    new_controls = schedule(123)
+    self.assertIn("learning_rate", new_controls)
+    self.assertIn("weight_decay_rate", new_controls)
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index e25836a3b..7b591ed9d 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -91,12 +91,12 @@ def FlattenControlsIntoTime(x, **unused_kwargs):  # pylint: disable=invalid-name
   if vocab_size is None:
     # In continuous policies every element of the output sequence corresponds to
     # an observation.
-    n_logits = n_controls * n_actions
+    n_preds_per_input = n_controls
     kwargs = {}
   else:
     # In discrete policies every element of the output sequence corresponds to
     # a symbol in the discrete representation, and each control takes 1 symbol.
-    n_logits = n_actions
+    n_preds_per_input = 1
     kwargs = {"vocab_size": vocab_size}
 
   if two_towers:
@@ -104,10 +104,12 @@ def FlattenControlsIntoTime(x, **unused_kwargs):  # pylint: disable=invalid-name
         tl.Dup(),
         tl.Parallel(
             [bottom_layers_fn(**kwargs),
-             tl.Dense(n_logits),
+             tl.Dense(n_preds_per_input * n_actions),
              FlattenControlsIntoTime(),  # pylint: disable=no-value-for-parameter
              tl.LogSoftmax()],
-            [bottom_layers_fn(), tl.Dense(n_controls), tl.Flatten()],
+            [bottom_layers_fn(**kwargs),
+             tl.Dense(n_preds_per_input),
+             tl.Flatten()],
         )
     ]
   else:
@@ -115,10 +117,10 @@ def FlattenControlsIntoTime(x, **unused_kwargs):  # pylint: disable=invalid-name
         bottom_layers_fn(**kwargs),
         tl.Dup(),
         tl.Parallel(
-            [tl.Dense(n_logits),
+            [tl.Dense(n_preds_per_input * n_actions),
              FlattenControlsIntoTime(),  # pylint: disable=no-value-for-parameter
              tl.LogSoftmax()],
-            [tl.Dense(n_controls), tl.Flatten()],
+            [tl.Dense(n_preds_per_input), tl.Flatten()],
         )
     ]
   return tl.Model(layers)

From 03bac1a833af654b4d3f08c78e071ee2f15e5e3f Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 27 Sep 2019 17:23:05 -0700
Subject: [PATCH 2493/2720] Allow overriding model keyword arguments for
 inference in SimulatedEnvProblem.

PiperOrigin-RevId: 271679587
---
 tensor2tensor/trax/rl/configs/simple_online_tune.gin |  3 +++
 tensor2tensor/trax/rl/simulated_env_problem.py       | 11 +++++++++--
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/rl/configs/simple_online_tune.gin b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
index abcdd0717..5a9878af6 100644
--- a/tensor2tensor/trax/rl/configs/simple_online_tune.gin
+++ b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
@@ -65,6 +65,9 @@ SerializedSequenceSimulatedEnvProblem.done_fn = @trax.rl.onlinetune_done_fn
 SerializedSequenceSimulatedEnvProblem.vocab_size = 128
 SerializedSequenceSimulatedEnvProblem.max_trajectory_length = 101
 SerializedSequenceSimulatedEnvProblem.significance_decay = 0.8
+SerializedSequenceSimulatedEnvProblem.model_predict_kwargs = {
+    "attention_type": @trax.layers.DotProductCausalAttention,
+}
 
 # Parameters for SimPLe:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index 77c62c735..a0c139a69 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -43,7 +43,8 @@ class SimulatedEnvProblem(env_problem.EnvProblem):
   """
 
   def __init__(self, model, batch_size, observation_space, action_space,
-               reward_range, discrete_rewards, history_stream, output_dir):
+               reward_range, discrete_rewards, history_stream, output_dir,
+               model_predict_kwargs=None):
     """Initializes the env.
 
     Args:
@@ -56,9 +57,15 @@ def __init__(self, model, batch_size, observation_space, action_space,
       history_stream: Iterator yielding batches of initial input data for the
         model. The format is implementation-specific.
       output_dir: (str) Output dir.
+      model_predict_kwargs: (dict) Additional model keyword arguments for
+        inference. Useful when different config is needed for training and
+        inference, e.g. train with memory efficient attention and predict with
+        the regular one.
     """
     self._model = model
-    model_predict = self._model(mode="predict")
+    if model_predict_kwargs is None:
+      model_predict_kwargs = {}
+    model_predict = self._model(mode="predict", **model_predict_kwargs)
     self._model_predict = backend.jit(model_predict)
     self._model_initialize = model_predict.initialize
     self._observation_space = observation_space

From 2e56e6609faa1909f08f644e8b1bc0682427a6c9 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Fri, 27 Sep 2019 18:02:31 -0700
Subject: [PATCH 2494/2720] Undo rescaling the observations to the [-1, 1]
 interval. Instead rescale the controls to the observation range.

PiperOrigin-RevId: 271684386
---
 tensor2tensor/trax/learning_rate.py           |  6 +--
 ...nline_tune_transformer_imagenet64_16gb.gin |  2 +-
 .../env_online_tune_transformer_lm1b_16gb.gin |  2 +-
 ...line_tune_transformer_lm_wmt_ende_16gb.gin |  2 +-
 tensor2tensor/trax/rl/envs/online_tune_env.py | 13 ++++---
 .../trax/rl/envs/online_tune_env_test.py      |  1 -
 tensor2tensor/trax/rl/online_tune.py          | 25 ++++++------
 tensor2tensor/trax/rl/online_tune_test.py     | 38 +++++++++----------
 8 files changed, 45 insertions(+), 44 deletions(-)

diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index f45f7188a..a0ba582c7 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -159,7 +159,7 @@ def PolicySchedule(
         # (name, start, (low, high), flip)
         ("learning_rate", 1e-3, (1e-9, 10.0), False),
     ),
-    metric_range=(0.0, 5.0),
+    observation_range=(0.0, 10.0),
     action_multipliers=(1.0 / 1.5, 1.0 / 1.25, 1.0, 1.25, 1.5),
     policy_and_value_model=trax_models.FrameStackMLP,
     policy_and_value_two_towers=False,
@@ -175,7 +175,7 @@ def PolicySchedule(
     include_controls_in_observation: bool, whether to include the controls in
       observations.
     control_configs: control configs, see trax.rl.envs.OnlineTuneEnv.
-    metric_range: tuple (low, high), range to clip the metrics to.
+    observation_range: tuple (low, high), range to clip the metrics to.
     action_multipliers: sequence of LR multipliers that policy actions
       correspond to.
     policy_and_value_model: Trax model to use as the policy.
@@ -196,7 +196,7 @@ def PolicySchedule(
   # return the initial learning rate.
   start_time = time.time()
   observations = online_tune.history_to_observations(
-      history, observation_metrics, metric_range,
+      history, observation_metrics, observation_range,
       control_configs if include_controls_in_observation else None
   )
   logging.vlog(
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
index 8c6ee7f64..3a33ddb93 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
@@ -101,5 +101,5 @@ OnlineTuneEnv.include_controls_in_observation = True
 OnlineTuneEnv.train_steps = 150
 OnlineTuneEnv.eval_steps = 2
 OnlineTuneEnv.env_steps = 100
-OnlineTuneEnv.metric_range = (0.0, 10.0)
+OnlineTuneEnv.observation_range = (0.0, 10.0)
 OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
index 7bc28c4c1..318f2b458 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
@@ -98,5 +98,5 @@ OnlineTuneEnv.include_controls_in_observation = True
 OnlineTuneEnv.train_steps = 300
 OnlineTuneEnv.eval_steps = 1
 OnlineTuneEnv.env_steps = 100
-OnlineTuneEnv.metric_range = (0.0, 10.0)
+OnlineTuneEnv.observation_range = (0.0, 10.0)
 OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
index 4a91735f4..b32d106c9 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
@@ -94,6 +94,6 @@ OnlineTuneEnv.include_controls_in_observation = True
 OnlineTuneEnv.train_steps = 500
 OnlineTuneEnv.eval_steps = 1
 OnlineTuneEnv.env_steps = 100
-OnlineTuneEnv.metric_range = (0.0, 10.0)
+OnlineTuneEnv.observation_range = (0.0, 10.0)
 OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
 OnlineTuneEnv.has_weights = True
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index b13cf1c32..dba9e67a9 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -72,7 +72,7 @@ def __init__(self,
                    ("learning_rate", 1e-3, (1e-9, 10.0), False),
                ),
                nontrainable_param_map=None,
-               metric_range=(0.0, 5.0),
+               observation_range=(0.0, 10.0),
                # Don't save checkpoints by default, as they tend to use a lot of
                # space.
                should_save_checkpoints=False,
@@ -103,7 +103,7 @@ def __init__(self,
     self._eval_steps = eval_steps
     self._env_steps = env_steps
     self._control_configs = control_configs
-    self._metric_range = metric_range
+    self._observation_range = observation_range
 
     self._output_dir = output_dir
     gfile.makedirs(self._output_dir)
@@ -117,9 +117,12 @@ def __init__(self,
         len(self._observation_metrics) +
         int(self._include_controls_in_observation) * len(self._control_configs)
     )
+
+    (obs_low, obs_high) = observation_range
     self.observation_space = gym.spaces.Box(
-        # Observations are clipped to metric_range and rescaled to [-1, 1].
-        low=-1, high=1, shape=(observation_dim,))
+        # Observations are clipped to this range.
+        low=obs_low, high=obs_high, shape=(observation_dim,),
+    )
 
   @property
   def _next_trajectory_dir(self):
@@ -162,7 +165,7 @@ def _current_observation(self):
     observations = online_tune.history_to_observations(
         self._trainer.state.history,
         self._observation_metrics,
-        self._metric_range,
+        self._observation_range,
         self._control_configs if self._include_controls_in_observation
         else None,
     )
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env_test.py b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
index 4f83fb9a3..aafddfe1a 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env_test.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
@@ -91,7 +91,6 @@ def _create_env(
             ("learning_rate", 1e-3, (1e-9, 10.0), False),
             ("weight_decay_rate", 1e-5, (1e-9, 0.1), False),
         ),
-        metric_range=(-1, 1),
         include_controls_in_observation=False,
         output_dir=output_dir,
         action_multipliers=action_multipliers,
diff --git a/tensor2tensor/trax/rl/online_tune.py b/tensor2tensor/trax/rl/online_tune.py
index 1252370d2..905fe59b6 100644
--- a/tensor2tensor/trax/rl/online_tune.py
+++ b/tensor2tensor/trax/rl/online_tune.py
@@ -35,15 +35,8 @@ def historical_metric_values(history, metric):
   return metric_values
 
 
-def metric_to_observation(metric_values, metric_range):
-  """Clips and scales the metric to the [-1, 1] interval."""
-  (low, high) = metric_range
-  clipped_values = np.clip(metric_values, low, high)
-  return (clipped_values - low) / (high - low) * 2 - 1
-
-
-def control_to_observation(control_values, control_config):
-  """Flips, logarithms, clips and scales the control to the [-1, 1] interval."""
+def control_to_observation(control_values, control_config, observation_range):
+  """Flips, logarithms, clips and scales the control to observation_range."""
   (_, _, (low, high), flip) = control_config
   def transform(x):
     return np.log(maybe_flip(x, flip))
@@ -52,7 +45,13 @@ def transform(x):
   )
   if flip:
     (log_low, log_high) = (log_high, log_low)
-  return metric_to_observation(log_control_values, (log_low, log_high))
+  log_control_values = np.clip(log_control_values, log_low, log_high)
+  # Rescale the log control values to the observation range.
+  (obs_low, obs_high) = observation_range
+  return (
+      (log_control_values - log_low) / (log_high - log_low) *
+      (obs_high - obs_low) + obs_low
+  )
 
 
 def control_metric(name):
@@ -81,10 +80,9 @@ def maybe_flip(value, flip):
 def history_to_observations(
     history, metrics, observation_range, control_configs=None):
   """Converts a trax History object into a sequence of observations."""
+  (obs_low, obs_high) = observation_range
   observation_dimensions = [
-      metric_to_observation(  # pylint: disable=g-complex-comprehension
-          historical_metric_values(history, metric), observation_range
-      )
+      np.clip(historical_metric_values(history, metric), obs_low, obs_high)
       for metric in metrics
   ]
   if control_configs is not None:
@@ -93,6 +91,7 @@ def history_to_observations(
       observation_dimensions.append(control_to_observation(
           historical_metric_values(history, control_metric(control_name)),
           control_config,
+          observation_range,
       ))
   return np.stack(observation_dimensions, axis=1)
 
diff --git a/tensor2tensor/trax/rl/online_tune_test.py b/tensor2tensor/trax/rl/online_tune_test.py
index 4337b24e6..54d5691d9 100644
--- a/tensor2tensor/trax/rl/online_tune_test.py
+++ b/tensor2tensor/trax/rl/online_tune_test.py
@@ -40,44 +40,44 @@ def test_retrieves_historical_metric_values(self):
     )
     np.testing.assert_array_equal(metric_values, [0.1, 0.73])
 
-  def test_metric_to_observation_rescales(self):
-    metric = np.random.uniform(low=-10, high=10, size=(100,))
-    observation = online_tune.metric_to_observation(metric, (-10, 10))
-    self.assertLess(-1, np.min(observation))
-    self.assertLess(np.min(observation), -0.9)
-    self.assertLess(0.9, np.max(observation))
-    self.assertLess(np.max(observation), 1.0)
-
-  def test_metric_to_observation_clips(self):
-    metric = np.random.uniform(low=-10, high=10, size=(100,))
-    observation = online_tune.metric_to_observation(metric, (-2, 2))
-    self.assertEqual(np.min(observation), -1)
-    self.assertEqual(np.max(observation), 1)
-
   def test_converts_control_to_log_scale_without_flipping(self):
     config = ("weight_decay", None, (1e-5, 0.1), False)
     controls = np.array([0.01, 0.02, 0.04])
-    obs = online_tune.control_to_observation(controls, config)
+    obs_range = (-1, 1)
+    obs = online_tune.control_to_observation(controls, config, obs_range)
     np.testing.assert_almost_equal(obs[1] - obs[0], obs[2] - obs[1])
 
   def test_converts_control_to_log_scale_with_flipping(self):
     config = ("momentum", None, (0.5, 0.99), True)
     controls = np.array([0.98, 0.96, 0.92])
-    obs = online_tune.control_to_observation(controls, config)
+    obs_range = (-1, 1)
+    obs = online_tune.control_to_observation(controls, config, obs_range)
     np.testing.assert_almost_equal(obs[1] - obs[0], obs[2] - obs[1])
 
   def test_clips_control_without_flipping(self):
     config = ("weight_decay", None, (1e-5, 0.1), False)
     controls = np.array([0.0, 0.2])
-    obs = online_tune.control_to_observation(controls, config)
+    obs_range = (-1, 1)
+    obs = online_tune.control_to_observation(controls, config, obs_range)
     np.testing.assert_equal(obs, [-1, 1])
 
   def test_clips_control_with_flipping(self):
     config = ("momentum", None, (0.5, 0.99), True)
     controls = np.array([0.4, 1.0])
-    obs = online_tune.control_to_observation(controls, config)
+    obs_range = (-1, 1)
+    obs = online_tune.control_to_observation(controls, config, obs_range)
     np.testing.assert_equal(obs, [1, -1])
 
+  def test_rescales_control(self):
+    config = ("weight_decay", None, (1e-5, 0.1), False)
+    controls = np.array([4e-4, 3e-3, 2e-2])
+    (obs_low, obs_high) = (103, 104)
+    obs = online_tune.control_to_observation(
+        controls, config, observation_range=(obs_low, obs_high),
+    )
+    np.testing.assert_array_less(obs, [obs_high] * 3)
+    np.testing.assert_array_less([obs_low] * 3, obs)
+
   def test_converts_history_to_observations_without_controls(self):
     history = trax_history.History()
     self._append_metrics(history, ("train", "loss"), [1.0, 0.07])
@@ -117,7 +117,7 @@ def test_clips_observations(self):
         observation_range=(-2, 2),
         control_configs=None,
     )
-    np.testing.assert_array_equal(observations, [[-1], [1]])
+    np.testing.assert_array_equal(observations, [[-2], [2]])
 
   def test_updates_control_without_flipping(self):
     config = ("learning_rate", None, (1e-9, 10.0), False)

From e2254bc24f4ab49201bddf3270b11bdab6c9e57e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 27 Sep 2019 18:10:31 -0700
Subject: [PATCH 2495/2720] Change names of key attributes of base Layer class.

PiperOrigin-RevId: 271685286
---
 tensor2tensor/trax/layers/README.md           |   8 +-
 tensor2tensor/trax/layers/attention.py        |  86 +++++++--------
 tensor2tensor/trax/layers/base.py             | 104 +++++++++---------
 tensor2tensor/trax/layers/base_test.py        |  16 +--
 tensor2tensor/trax/layers/combinators.py      |  27 ++---
 tensor2tensor/trax/layers/convolution.py      |   8 +-
 tensor2tensor/trax/layers/core.py             |  12 +-
 tensor2tensor/trax/layers/core_test.py        |   4 +-
 tensor2tensor/trax/layers/demo.ipynb          |   2 +-
 tensor2tensor/trax/layers/normalization.py    |   6 +-
 .../trax/layers/normalization_test.py         |   2 +-
 tensor2tensor/trax/layers/reversible.py       |  12 +-
 tensor2tensor/trax/learning_rate_test.py      |   2 +-
 tensor2tensor/trax/models/atari_cnn_test.py   |   4 +-
 .../models/research/transformer_revnet.py     |  34 +++---
 tensor2tensor/trax/models/transformer_test.py |   4 +-
 tensor2tensor/trax/rl/ppo_test.py             |  10 +-
 tensor2tensor/trax/rl/ppo_trainer.py          |   2 +-
 tensor2tensor/trax/rl/simple_test.py          |   2 +-
 .../trax/rl/simulated_env_problem.py          |   2 +-
 .../trax/rl/simulated_env_problem_test.py     |   2 +-
 tensor2tensor/trax/trax.py                    |   3 +-
 22 files changed, 179 insertions(+), 173 deletions(-)

diff --git a/tensor2tensor/trax/layers/README.md b/tensor2tensor/trax/layers/README.md
index fbdaedd50..5efe6e86b 100644
--- a/tensor2tensor/trax/layers/README.md
+++ b/tensor2tensor/trax/layers/README.md
@@ -7,13 +7,13 @@
 All layers inherit form the Layer class and need to implement 3 functions:
 
 ```python
-def call(self, params, inputs, **kwargs):
+def forward(self, params, inputs, **kwargs):
 """Call this layer using the given parameters on the given inputs."""
 
 def output_shape(self, input_shape):
 """The shape of the output given the shape of the input."""
 
-def new_parameters(self, input_shape, rng):
+def new_params_and_state(self, input_shape, rng):
 """Create new parameters given the shape of the input."""
 ```
 
@@ -24,7 +24,7 @@ and call functions to be used as follows.
 layer = MyLayer()
 x = np.zeros(10)
 rng = random.get_prng(0)
-params = layer.initialize(x.shape, x.dtype, rng)
+params = layer.initialize_once(x.shape, x.dtype, rng)
 output = layer(x, params, rng=rng)
 ```
 
@@ -48,7 +48,7 @@ standard_mlp = layers.Serial(layers.Dense(10), layers.Dense(10))
 layer = Dense(10)
 shared_parameters_mlp = layers.Serial(layer, layer)
 ```
-For this reason, if you call `layer.initialize(...)` for the second time
+For this reason, if you call `layer.initialize_once(...)` for the second time
 on an already initialized layer, it will return `()`.
 
 ## Core layers
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index d2d2cc42c..afdeb8082 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -79,7 +79,7 @@ def __init__(self, max_len=2048, mode='train'):
     self._max_len = max_len
     self._mode = mode
 
-  def call(self, inputs, params, state, **kwargs):
+  def forward(self, inputs, params, state, **kwargs):
     if self._mode in ('train', 'eval'):
       x = inputs
       symbol_size = np.shape(x)[1]
@@ -90,7 +90,7 @@ def call(self, inputs, params, state, **kwargs):
       # storing the index in state.
       return (inputs + np.expand_dims(params[:, state, :], 1), state + 1)
 
-  def new_parameters(self, input_shape, input_dtype, rng):
+  def new_params_and_state(self, input_shape, input_dtype, rng):
     del input_dtype, rng
     d_feature = input_shape[-1]
     pe = onp.zeros((self._max_len, d_feature), dtype=onp.float32)
@@ -258,13 +258,13 @@ def __init__(self, initializer=init.RandomNormalInitializer(0.01)):
     super(ShiftRightLearned, self).__init__()
     self._initializer = initializer
 
-  def call(self, x, params, state, **kwargs):
+  def forward(self, x, params, state, **kwargs):
     del kwargs
     c = backend.numpy.reshape(params, [1, 1, -1])
     c += backend.numpy.zeros((x.shape[0], 1, x.shape[2]), dtype=x.dtype)
     return backend.numpy.concatenate([c, x], axis=1)[:, :-1, :], state
 
-  def new_parameters(self, input_shape, input_dtype, rng):
+  def new_params_and_state(self, input_shape, input_dtype, rng):
     del input_dtype
     b = self._initializer((input_shape[-1],), rng)
     return b, ()
@@ -287,7 +287,7 @@ def __init__(self, n_heads=1, d_head=64,
     # implementation, and shouldn't have an effect on modeling quality.
     # Note that AttentionQKV above is different in that it uses a bias term.
 
-  def call(self, x, params, state, **kwargs):
+  def forward(self, x, params, state, **kwargs):
     del kwargs
     seqlen = x.shape[1]
     res = np.dot(x, params)
@@ -301,7 +301,7 @@ def call(self, x, params, state, **kwargs):
 
     return res, state
 
-  def new_parameters(self, input_shape, input_dtype, rng):
+  def new_params_and_state(self, input_shape, input_dtype, rng):
     del input_dtype
     w = self._kernel_initializer(
         (input_shape[-1], self._n_heads * self._d_head), rng)
@@ -321,7 +321,7 @@ def __init__(self, n_heads=1, d_model=1024,
     # implementation, and shouldn't have an effect on modeling quality.
     # Note that AttentionQKV above is different in that it uses a bias term.
 
-  def call(self, x, params, state, **kwargs):
+  def forward(self, x, params, state, **kwargs):
     del kwargs
     seqlen = x.shape[1]
     d_head = x.shape[2]
@@ -332,7 +332,7 @@ def call(self, x, params, state, **kwargs):
 
     return np.dot(x, params), state
 
-  def new_parameters(self, input_shape, input_dtype, rng):
+  def new_params_and_state(self, input_shape, input_dtype, rng):
     del input_dtype
     w = self._kernel_initializer(
         (input_shape[-1] * self._n_heads, self._d_model), rng)
@@ -345,11 +345,11 @@ class BaseCausalAttention(base.Layer):
   def __init__(self):
     super(BaseCausalAttention, self).__init__(n_inputs=3)
 
-  def call(self, inputs, params=(), state=(), rng=None, **kwargs):
+  def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
     """Forward pass for the attention layer."""
     raise NotImplementedError()
 
-  def call_and_grad(self, inputs, grad, **kwargs):
+  def forward_and_backward(self, inputs, grad, **kwargs):
     """Performs both forward and backward pass for the attention layer.
 
     This is used in reversible models: for the backward pass of a reversible
@@ -372,7 +372,7 @@ def call_and_grad(self, inputs, grad, **kwargs):
     """
     raise NotImplementedError()
 
-  def new_parameters(self, input_shapes, input_dtype, rng):
+  def new_params_and_state(self, input_shapes, input_dtype, rng):
     return (), ()
 
 
@@ -384,7 +384,7 @@ def __init__(self, dropout=0.0, mode='train'):
     self._dropout = dropout
     self._mode = mode
 
-  def call(self, inputs, params=(), state=(), rng=None, **kwargs):
+  def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
     del params
     q, k, v = inputs
     if self._mode in ('train', 'eval'):
@@ -419,17 +419,17 @@ def call(self, inputs, params=(), state=(), rng=None, **kwargs):
         q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=rng)
     return res, state
 
-  def call_and_grad(self, inputs, ct, **kwargs):
+  def forward_and_backward(self, inputs, ct, **kwargs):
     assert backend.get_name() == 'jax', (
-        'JAX backend is required to use call_and_grad.')
+        'JAX backend is required to use forward_and_backward.')
     # Simultaneous forward pass and backprop through the attention mechanism.
-    def do_call(x):  # pylint: disable=invalid-name
-      res, _ = self.call(x, **kwargs)
+    def _do_forward(x):  # pylint: disable=invalid-name
+      res, _ = self.forward(x, **kwargs)
       return res
-    output, vjpfun = jax.vjp(do_call, inputs)
+    output, vjpfun = jax.vjp(_do_forward, inputs)
     return output, vjpfun(ct)[0]
 
-  def new_parameters(self, input_shapes, input_dtype, rng):
+  def new_params_and_state(self, input_shapes, input_dtype, rng):
     if self._mode in ('train', 'eval'):
       return (), ()
 
@@ -478,17 +478,17 @@ def __init__(self, loop_stride, dropout, mode, share_qk=False, hard_k=0):
     self._share_qk = share_qk
     self._hard_k = hard_k
 
-  def call(self, inputs, params=(), state=(), **kwargs):
+  def forward(self, inputs, params=(), state=(), **kwargs):
     del params
-    output, _ = self.call_and_grad(inputs, None, **kwargs)
+    output, _ = self.forward_and_backward(inputs, None, **kwargs)
     return output, state
 
-  def has_custom_grad(self):
+  def has_backward(self):
     return True
 
-  def custom_grad(self, inputs, output, ct, params=(), state=(), **kwargs):
+  def backward(self, inputs, output, ct, params=(), state=(), **kwargs):
     del output, params, state
-    _, inputs_ct = self.call_and_grad(inputs, ct, **kwargs)
+    _, inputs_ct = self.forward_and_backward(inputs, ct, **kwargs)
     return inputs_ct, ()
 
   def make_unit_length(self, x, epsilon=1e-6):
@@ -496,7 +496,7 @@ def make_unit_length(self, x, epsilon=1e-6):
     norm_inputs = x / np.sqrt(variance + epsilon)
     return norm_inputs
 
-  def call_and_grad(self, inputs, ct, rng=None, **kwargs):
+  def forward_and_backward(self, inputs, ct, rng=None, **kwargs):
     del kwargs
     query, key, value = inputs
     depth = np.shape(query)[-1]
@@ -670,17 +670,17 @@ def __init__(self, dropout, mode, n_bins=64,
     if one_rng:
       self._prng = backend.random.get_prng(seed)
 
-  def call(self, inputs, params=(), state=(), **kwargs):
+  def forward(self, inputs, params=(), state=(), **kwargs):
     del params
-    output, _ = self.call_and_grad(inputs, None, **kwargs)
+    output, _ = self.forward_and_backward(inputs, None, **kwargs)
     return output, state
 
-  def has_custom_grad(self):
+  def has_backward(self):
     return True
 
-  def custom_grad(self, inputs, output, ct, params=(), state=(), **kwargs):
+  def backward(self, inputs, output, ct, params=(), state=(), **kwargs):
     del output, params, state
-    _, inputs_ct = self.call_and_grad(inputs, ct, **kwargs)
+    _, inputs_ct = self.forward_and_backward(inputs, ct, **kwargs)
     return inputs_ct, ()
 
   def bin_vectors_by_time(self, vecs):
@@ -724,7 +724,7 @@ def hash_vectors(self, vecs, rng):
     bins = np.argmax(rotated_vecs, axis=-1)
     return bins
 
-  def call_and_grad(self, inputs, ct, rng=None, **kwargs):
+  def forward_and_backward(self, inputs, ct, rng=None, **kwargs):
     del kwargs
     # We use the same vector as both a query and a key. For now we haven't
     # adjusted any of the surrounding code, so we still get a separate "key"
@@ -916,7 +916,7 @@ def hash_vectors(self, vecs, rng):
     bins = np.argmax(rotated_vecs, axis=-1)
     return bins
 
-  def call(self, inputs, params=(), state=(), rng=None, **kwargs):
+  def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
     del params, kwargs
     # We use the same vector as both a query and a key. For now we haven't
     # adjusted any of the surrounding code, so we still get a separate "key"
@@ -1113,14 +1113,14 @@ def vjpfun(grad):
 
     return out, state
 
-  def call_and_grad(self, inputs, ct, rng=None, **kwargs):
-    # TODO(kitaev): is there a manual implementation of call_and_grad that's
-    # faster than having jax infer one? Or are the permute/unpermute custom
-    # gradients defined in call() sufficient for reasonable speed?
-    def _do_call(x):
-      return self.call(x, params=(), state=(), rng=rng, **kwargs)[0]
+  def forward_and_backward(self, inputs, ct, rng=None, **kwargs):
+    # TODO(kitaev): is there a manual implementation of forward_and_backward
+    # that's faster than having jax infer one? Or are the permute/unpermute
+    # custom gradients defined in forward() sufficient for reasonable speed?
+    def _do_forward(x):
+      return self.forward(x, params=(), state=(), rng=rng, **kwargs)[0]
 
-    output, vjpfun = jax.vjp(_do_call, inputs)
+    output, vjpfun = jax.vjp(_do_forward, inputs)
     return output, vjpfun(ct)[0]
 
 
@@ -1153,22 +1153,22 @@ def __init__(self, dropout, mode, n_bins=64, n_hashes=1, n_buckets=64,
     self._hard_k = hard_k
     self._rehash_each_round = rehash_each_round
 
-  def call(self, inputs, params=(), state=(), rng=None, **kwargs):
+  def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
     del params, kwargs
     output, _ = self.batch_call_and_or_grad(inputs[0], inputs[2], rng=rng)
     return output, state
 
-  def call_and_grad(self, inputs, ct, rng=None, **kwargs):
+  def forward_and_backward(self, inputs, ct, rng=None, **kwargs):
     del kwargs
     output, (qk_ct, v_ct) = self.batch_call_and_or_grad(
         inputs[0], inputs[2], ct=ct, rng=rng)
     return output, (qk_ct, np.zeros_like(inputs[1]), v_ct)
 
-  def has_custom_grad(self):
+  def has_backward(self):
     return True
 
-  def custom_grad(self, inputs, output, ct, params=(), state=(), rng=None,
-                  **kwargs):
+  def backward(self, inputs, output, ct, params=(), state=(), rng=None,
+               **kwargs):
     del output, params, state
     _, (qk_ct, v_ct) = self.batch_call_and_or_grad(
         inputs[0], inputs[2], return_output=False, ct=ct, rng=rng)
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index b2b553d6e..5117903a9 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -84,7 +84,7 @@ def __repr__(self):
     else:
       return '{}[{}]'.format(class_str, fields_str)
 
-  def call(self, inputs, params=(), state=(), **kwargs):
+  def forward(self, inputs, params=(), state=(), **kwargs):
     """Applies this layer to given activation tensors, using trainable params.
 
     Args:
@@ -110,8 +110,7 @@ def call(self, inputs, params=(), state=(), **kwargs):
     """
     raise NotImplementedError
 
-  # TODO(wangpeng): Should be called `new_parameters_and_state`.
-  def new_parameters(self, input_shapes, input_dtype, rng):
+  def new_params_and_state(self, input_shapes, input_dtype, rng):
     """Creates layer-specific parameters based on data shape, dtype and rng.
 
     Args:
@@ -148,11 +147,11 @@ def sublayers(self):
     return self._sublayers
 
   @property
-  def has_custom_grad(self):
+  def has_backward(self):
     """Whether to use custom gradients (in which case, see below)."""
     return False
 
-  def custom_grad(self, inputs, output, grad, params, state, **kwargs):
+  def backward(self, inputs, output, grad, params, state, **kwargs):
     """Custom backward pass to propagate gradients in a custom way.
 
     Args:
@@ -166,20 +165,20 @@ def custom_grad(self, inputs, output, grad, params, state, **kwargs):
 
     Returns:
       The custom gradient signal for the input. Note that we need to return
-      a gradient for each argument of call, so it will usually be a tuple
+      a gradient for each argument of forward, so it will usually be a tuple
       of signals: the gradient for inputs and parameters.
     """
     raise NotImplementedError
 
   # End of subclassing interface, all functions below are internal.
 
-  def pseudo_call(self, pseudo_inputs, params, state):
+  def pseudo_forward(self, pseudo_inputs, params, state):
     """Computes shapes and types this layer would produce for the given inputs.
 
     Args:
       pseudo_inputs: A ShapeType instance (input data minus the actual values)
           or a tuple of ShapeType instances, following the same conventions as
-          Layer.call's input arg.
+          Layer.forward's input arg.
       params: Parameters for this layer.
       state: start state.
 
@@ -194,7 +193,7 @@ def pseudo_call(self, pseudo_inputs, params, state):
       # stored in global memory.
       rng = ShapeType(shape=(2,), dtype=onp.uint32)
       def call_on_input(x, params, state, rng):
-        return self.call(x, params=params, state=state, rng=rng)
+        return self.forward(x, params=params, state=state, rng=rng)
       params_shapes = nested_map(
           params, lambda x: ShapeType(shape=x.shape, dtype=x.dtype))
       s = backend.eval_on_shapes(call_on_input)(pseudo_inputs,
@@ -202,13 +201,14 @@ def call_on_input(x, params, state, rng):
       return s
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback(skip=3)
-      raise LayerError(name, 'pseudo_call', self._caller, pseudo_inputs, trace)
+      raise LayerError(name, 'pseudo_forward', self._caller, pseudo_inputs,
+                       trace)
 
-  def initialize(self, input_shapes, input_dtype, rng):
+  def initialize_once(self, input_shapes, input_dtype, rng):
     """Initialize the layer given an input shape, dtype and rng.
 
-    Returns new_parameters(input_shapes, rng) on the first call and () on any
-    subsequent call, as the layer is already initialized. This is used for
+    Returns new_params_and_state(input_shapes, rng) on the first call and () on
+    any subsequent call, as the layer is already initialized. This is used for
     networks that share parameters, so the layer only produces them once.
 
     Args:
@@ -223,12 +223,12 @@ def initialize(self, input_shapes, input_dtype, rng):
     """
     try:
       # Initialize params once; store them for use when this layer is called.
-      # Needs to call new_parameters regardless of _init_finished because state
-      # also needs to be initialized. After jitting, graph pruning should be
-      # able to remove unnecessary computation.
+      # Needs to call new_params_and_state regardless of _init_finished because
+      # state also needs to be initialized. After jitting, graph pruning should
+      # be able to remove unnecessary computation.
       # TODO(lukaszkaiser): Revisit this decision and see whether layers sharing
       #   params should also share states.
-      params, state = self.new_parameters(input_shapes, input_dtype, rng)
+      params, state = self.new_params_and_state(input_shapes, input_dtype, rng)
       if not self._init_finished:
         self._init_finished = True
         self._params = params
@@ -237,7 +237,8 @@ def initialize(self, input_shapes, input_dtype, rng):
       return (params, state)
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback(skip=3)
-      raise LayerError(name, 'initialize', self._caller, input_shapes, trace)
+      raise LayerError(name, 'initialize_once', self._caller, input_shapes,
+                       trace)
 
   # XXX(kitaev):
   _STASH_IN = None
@@ -255,14 +256,14 @@ def __call__(self, x, params=(), state=(), **kwargs):
         # In this case, we're called for the first time: cache parameters.
         self._params = params
 
-      if not self.has_custom_grad or Layer._STASH_IN is not None:
-        return self.call(x, params=params, state=state, **kwargs)
+      if not self.has_backward or Layer._STASH_IN is not None:
+        return self.forward(x, params=params, state=state, **kwargs)
       else:
         return self._do_custom_gradients(x, params, state, **kwargs)
 
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback()
-      raise LayerError(name, 'call', self._caller, shapes(x), trace)
+      raise LayerError(name, 'forward', self._caller, shapes(x), trace)
 
   def _do_custom_gradients(self, x, params, state, **kwargs):
     """Calls this layer for a forward pass, but with custom gradients."""
@@ -287,32 +288,32 @@ def check_end_state(output_state):
     # https://jax.readthedocs.io/en/latest/jax.html#jax.custom_transforms
     # Note that we capture the kwargs and don't calculate gradients wrt. them.
     @jax.custom_transforms
-    def do_call(y, params):
-      return check_end_state(self.call(y, params=params, state=state,
-                                       **kwargs))
+    def _do_forward(y, params):
+      return check_end_state(self.forward(y, params=params, state=state,
+                                          **kwargs))
 
     # This is the custom gradient (vector-jacobian product in JAX) function.
     # For the exact specification of this custom transformation see this link:
     # https://jax.readthedocs.io/en/latest/jax.html#jax.defjvp_all
-    def do_call_vjp(y, params):
+    def do_forward_vjp(y, params):
       """Custom gradient (vjp) function."""
       stash = None
       if Layer._STASH_IN is None:
         Layer._STASH_IN = stash = {}
-      output = check_end_state(self.call(y, params=params, state=state,
-                                         **kwargs))
+      output = check_end_state(self.forward(y, params=params, state=state,
+                                            **kwargs))
       if stash is not None:
         Layer._STASH_IN = None
       def vjpfun(grad):
         assert Layer._STASH_OUT is None
         Layer._STASH_OUT = stash
-        res = self.custom_grad(y, output, grad, params, state, **kwargs)
+        res = self.backward(y, output, grad, params, state, **kwargs)
         Layer._STASH_OUT = None
         return res
       return output, vjpfun
 
-    jax.defvjp_all(do_call, do_call_vjp)
-    return do_call(x, params), state
+    jax.defvjp_all(_do_forward, do_forward_vjp)
+    return _do_forward(x, params), state
 
 
 class LayerError(Exception):
@@ -325,7 +326,7 @@ class LayerError(Exception):
   def __init__(self, layer_name, function_name, caller,
                input_shapes, traceback_string):
     self._layer_name = layer_name
-    self._function_name = function_name  # Is it call or initialize?
+    self._function_name = function_name  # Is it forward or initialize_once?
     self._caller = caller  # Python inspect object with init caller info.
     self._traceback = traceback_string
     self._input_shapes = input_shapes
@@ -435,7 +436,7 @@ def _short_traceback(skip=3):
   return '\n'.join(res)
 
 
-def _validate_call_input(x, n_inputs):
+def _validate_forward_input(x, n_inputs):
   if n_inputs != 1:
     if not isinstance(x, tuple):
       raise TypeError(
@@ -446,44 +447,45 @@ def _validate_call_input(x, n_inputs):
           ' ({})'.format(len(x), n_inputs))
 
 
-def layer(n_inputs=1, n_outputs=1, new_parameters=None):
+def layer(n_inputs=1, n_outputs=1, new_params_and_state_fn=None):
   """Decorates a function to make it the call method of a new Layer class."""
-  # TODO(jonni): Consider renaming new_parameters to new_parameters_fn.
 
-  def _build_layer_class(raw_call_fn):
-    """Returns a Layer class built around the given call function."""
+  def _build_layer_class(raw_fn):
+    """Returns a Layer class built around the given forward function."""
 
     def _init(self, **kwargs):
       self._kwargs = kwargs  # pylint: disable=protected-access
       Layer.__init__(self, n_inputs=n_inputs, n_outputs=n_outputs)
 
-    def _new_parameters(self, input_shapes, input_dtype, rng):
-      if new_parameters is None:
+    def _new_params_and_state(self, input_shapes, input_dtype, rng):
+      if new_params_and_state_fn is None:
         return (), ()
       kwargs = self._kwargs  # pylint: disable=protected-access
-      return new_parameters(input_shapes, input_dtype, rng, **kwargs), ()
+      return (
+          new_params_and_state_fn(input_shapes, input_dtype, rng, **kwargs), ())
 
     def _is_empty(raw_output):
       return raw_output is None or (isinstance(raw_output, (list, tuple))
                                     and len(raw_output) == 0)  # pylint: disable=g-explicit-length-test
 
-    def _call_with_context(self, x, params=(), state=(), **kwargs):
-      """Calls raw_call_fn with extra keyword args from Layer.__init__."""
+    def _forward(self, x, params=(), state=(), **kwargs):
+      """Calls raw_fn with extra keyword args from Layer.__init__."""
       merged_kwargs = kwargs.copy()
       merged_kwargs.update(self._kwargs)  # pylint: disable=protected-access
 
-      _validate_call_input(x, n_inputs)
-      raw_output = raw_call_fn(x, params=params, **merged_kwargs)
+      _validate_forward_input(x, n_inputs)
+      raw_output = raw_fn(x, params=params, **merged_kwargs)
       output = () if _is_empty(raw_output) else raw_output
       return (output, state)
 
     # Set docstrings and create the class.
-    _call_with_context.__doc__ = raw_call_fn.__doc__
-    _new_parameters.__doc__ = new_parameters.__doc__  # None.__doc__ is None
-    cls = type(raw_call_fn.__name__, (Layer,),
+    _forward.__doc__ = raw_fn.__doc__
+    _new_params_and_state.__doc__ = new_params_and_state_fn.__doc__
+    # Note: None.__doc__ is None
+    cls = type(raw_fn.__name__, (Layer,),
                {'__init__': _init,
-                'call': _call_with_context,
-                'new_parameters': _new_parameters})
+                'forward': _forward,
+                'new_params_and_state': _new_params_and_state})
     return cls
 
   return _build_layer_class
@@ -523,7 +525,7 @@ def _is_tuple_of_shapes(shape):
 
 
 def check_shape_agreement(layer_fn, input_shapes, integer_inputs=False):
-  """Checks if the layer's call output agrees its pseudo_call predictions.
+  """Checks if the layer's forward output agrees its pseudo_forward predictions.
 
   This function helps test layer mechanics and inter-layer connections that
   aren't dependent on specific data values.
@@ -548,8 +550,8 @@ def check_shape_agreement(layer_fn, input_shapes, integer_inputs=False):
     input_dtype = tuple(input_dtype for _ in input_shapes)
   else:
     pseudo_data = ShapeType(input_shapes, input_dtype)
-  params, state = layer_fn.initialize(input_shapes, input_dtype, rng1)
-  pseudo_output, _ = layer_fn.pseudo_call(pseudo_data, params, state)
+  params, state = layer_fn.initialize_once(input_shapes, input_dtype, rng1)
+  pseudo_output, _ = layer_fn.pseudo_forward(pseudo_data, params, state)
   if isinstance(pseudo_output, tuple):
     output_shape = tuple(x.shape for x in pseudo_output)
   else:
diff --git a/tensor2tensor/trax/layers/base_test.py b/tensor2tensor/trax/layers/base_test.py
index 26140248c..aa49869fc 100644
--- a/tensor2tensor/trax/layers/base_test.py
+++ b/tensor2tensor/trax/layers/base_test.py
@@ -38,19 +38,19 @@ def test_custom_zero_grad(self):
 
     class IdWithZeroGrad(base.Layer):
 
-      def call(self, x, params, **kwargs):
+      def forward(self, x, params, **kwargs):
         del params, kwargs
         return x, ()
 
-      def new_parameters(self, input_shapes, input_dtype, rng):
+      def new_params_and_state(self, input_shapes, input_dtype, rng):
         del input_shapes, input_dtype, rng
         return (), ()
 
       @property
-      def has_custom_grad(self):
+      def has_backward(self):
         return True
 
-      def custom_grad(self, inputs, output, ct, params, state, **kwargs):
+      def backward(self, inputs, output, ct, params, state, **kwargs):
         return (backend.numpy.zeros_like(ct), ())
 
     layer = IdWithZeroGrad()
@@ -68,19 +68,19 @@ def test_custom_id_grad(self):
 
     class IdWithIdGrad(base.Layer):
 
-      def call(self, x, params, **kwargs):
+      def forward(self, x, params, **kwargs):
         del params, kwargs
         return x, ()
 
-      def new_parameters(self, input_shapes, input_dtype, rng):
+      def new_params_and_state(self, input_shapes, input_dtype, rng):
         del input_shapes, input_dtype, rng
         return (), ()
 
       @property
-      def has_custom_grad(self):
+      def has_backward(self):
         return True
 
-      def custom_grad(self, inputs, output, ct, params, state, **kwargs):
+      def backward(self, inputs, output, ct, params, state, **kwargs):
         return (inputs, ())
 
     layer = IdWithIdGrad()
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 08832ea16..c66680a47 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -146,18 +146,18 @@ def _n_inputs_n_outputs(self, layers):
       running_total -= layer.n_outputs
     return running_max, (running_max - running_total)
 
-  def _validate_call_inputs(self, xs):
+  def _validate_forward_inputs(self, xs):
     if not isinstance(xs, tuple) and self._n_inputs != 1:
       raise TypeError(
-          'Serial.call input must be a tuple; instead got {}'.format(xs))
+          'Serial.forward input must be a tuple; instead got {}'.format(xs))
     len_xs = 1 if isinstance(xs, np.ndarray) else len(xs)
     if len_xs < self.n_inputs:
       raise ValueError(
-          'number of inputs ({}) to Serial.call less than n_inputs'
+          'number of inputs ({}) to Serial.forward less than n_inputs'
           ' ({})'.format(len(xs), self.n_inputs))
 
-  def call(self, xs, params=(), state=(), **kwargs):
-    self._validate_call_inputs(xs)
+  def forward(self, xs, params=(), state=(), **kwargs):
+    self._validate_forward_inputs(xs)
     rngs = _pop_rng_and_split(kwargs, self._n_layers)
     if not self.sublayers:  # No-op: leave args unchanged.
       return (xs, state)
@@ -195,7 +195,7 @@ def call(self, xs, params=(), state=(), **kwargs):
 
     return stack, new_state
 
-  def new_parameters(self, input_shape, input_dtype, rng):
+  def new_params_and_state(self, input_shape, input_dtype, rng):
     def MakeShapeType(shape, dtype):
       if isinstance(dtype, (list, tuple)):
         return tuple(MakeShapeType(s, t) for s, t in zip(shape, dtype))
@@ -219,10 +219,10 @@ def MakeShapeType(shape, dtype):
 
       in_shape = base.nested_map(inputs, lambda x: x.shape)
       in_dtype = base.nested_map(inputs, lambda x: x.dtype)
-      param, state = layer.initialize(in_shape, in_dtype, layer_rng)
+      param, state = layer.initialize_once(in_shape, in_dtype, layer_rng)
       pparam = layer._params   # pylint: disable=protected-access
 
-      outputs, _ = layer.pseudo_call(inputs, pparam, state)
+      outputs, _ = layer.pseudo_forward(inputs, pparam, state)
 
       # Push outputs onto remaining pseudo_xs (if any).
       if n_in < _count_items(pseudo_xs):
@@ -325,10 +325,10 @@ def __init__(self, n_items=2, axis=-1):
     self._n_items = n_items
     self._axis = axis
 
-  def new_parameters(self, input_shape, input_dtype, rng):
+  def new_params_and_state(self, input_shape, input_dtype, rng):
     return (), ()
 
-  def call(self, xs, params=(), state=(), **kwargs):
+  def forward(self, xs, params=(), state=(), **kwargs):
     del params, kwargs
     return backend.numpy.concatenate(xs, self._axis), state
 
@@ -423,7 +423,7 @@ def _allot_to_sublayers(self, inputs):
       start = end
     return tuple(sub_inputs)
 
-  def call(self, inputs, params=(), state=(), **kwargs):
+  def forward(self, inputs, params=(), state=(), **kwargs):
     n_layers, layers = self._n_layers, self.sublayers
     sublayer_inputs = self._allot_to_sublayers(inputs)
     rngs = _pop_rng_and_split(kwargs, n_layers)
@@ -444,11 +444,12 @@ def call(self, inputs, params=(), state=(), **kwargs):
     output = outputs[0] if self.n_outputs == 1 else tuple(outputs)
     return output, new_state
 
-  def new_parameters(self, input_shapes, input_dtypes, rng):
+  def new_params_and_state(self, input_shapes, input_dtypes, rng):
     sublayer_shapes = self._allot_to_sublayers(input_shapes)
     sublayer_dtypes = self._allot_to_sublayers(input_dtypes)
     rngs = backend.random.split(rng, self._n_layers)
-    inits = [layer.initialize(shape, dtype, rng) for layer, shape, dtype, rng
+    inits = [layer.initialize_once(shape, dtype, rng)
+             for layer, shape, dtype, rng
              in zip(self.sublayers, sublayer_shapes, sublayer_dtypes, rngs)]
     if not inits:
       return (), ()
diff --git a/tensor2tensor/trax/layers/convolution.py b/tensor2tensor/trax/layers/convolution.py
index aa9fd4efc..d54a8d6f3 100644
--- a/tensor2tensor/trax/layers/convolution.py
+++ b/tensor2tensor/trax/layers/convolution.py
@@ -56,7 +56,7 @@ def _check_nhwc(self):
     msg = 'Convolutions on more than 4 dimensions only supported in NHWC.'
     assert self._lhs_spec == self._out_spec == 'NHWC', msg
 
-  def call(self, x, params=(), state=(), **kwargs):
+  def forward(self, x, params=(), state=(), **kwargs):
     del kwargs
     w, b = params
     x_shape = list(x.shape)
@@ -78,7 +78,7 @@ def _kernel_shape(self, input_shape):
             input_shape[self._lhs_spec.index('C')] if c == 'I' else
             next(kernel_size_iter) for c in self._rhs_spec]
 
-  def new_parameters(self, input_shape, input_dtype, rng):
+  def new_params_and_state(self, input_shape, input_dtype, rng):
     del input_dtype
     if len(input_shape) > 4:
       self._check_nhwc()
@@ -112,7 +112,7 @@ def __init__(self,
         kernel_initializer=kernel_initializer,
         bias_initializer=bias_initializer)
 
-  def call(self, x, params=(), **kwargs):
+  def forward(self, x, params=(), **kwargs):
     assert self._padding == 'VALID'
     # Left pad with 0s. Applying an unmasked valid convolution on top of this
     # yields a causal convolution.
@@ -122,5 +122,5 @@ def call(self, x, params=(), **kwargs):
     pad = effective_kernel_size - 1
     x_leftpad = np.pad(x, pad_width=[[0, 0], [pad, 0], [0, 0]], mode='constant')
 
-    res = super(CausalConv, self).call(x_leftpad, params)
+    res = super(CausalConv, self).forward(x_leftpad, params)
     return res
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 7f1128c56..f815590bf 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -123,12 +123,12 @@ def __init__(self,
     self._kernel_initializer = kernel_initializer
     self._bias_initializer = bias_initializer
 
-  def call(self, x, params, state, **kwargs):
+  def forward(self, x, params, state, **kwargs):
     del kwargs
     w, b = params
     return np.dot(x, w) + b, state
 
-  def new_parameters(self, input_shape, input_dtype, rng):
+  def new_params_and_state(self, input_shape, input_dtype, rng):
     del input_dtype
     rng1, rng2 = backend.random.split(rng, 2)
     w = self._kernel_initializer((input_shape[-1], self._n_units), rng1)
@@ -148,11 +148,11 @@ def __init__(self,
     self._vocab_size = vocab_size
     self._kernel_initializer = kernel_initializer
 
-  def call(self, x, params, state, **kwargs):
+  def forward(self, x, params, state, **kwargs):
     del kwargs
     return np.take(params, x, axis=0), state
 
-  def new_parameters(self, input_shape, input_dtype, rng):
+  def new_params_and_state(self, input_shape, input_dtype, rng):
     del input_dtype
     return self._kernel_initializer((self._vocab_size, self._d_feature),
                                     rng), ()
@@ -179,12 +179,12 @@ def __init__(self, rate=0.0, name='dropout', mode='train'):
     self._name = 'dropout_' + name
     self._mode = mode
 
-  def new_parameters(self, input_shape, input_dtype, rng):
+  def new_params_and_state(self, input_shape, input_dtype, rng):
     """Initialize dropout parameters and state."""
     del input_shape, input_dtype, rng
     return (), {self._name: np.array(self._initial_rate)}
 
-  def call(self, x, params, state, rng=None, **unused_kwargs):
+  def forward(self, x, params, state, rng=None, **unused_kwargs):
     """Execute dropout."""
     del params
     rate = self._initial_rate
diff --git a/tensor2tensor/trax/layers/core_test.py b/tensor2tensor/trax/layers/core_test.py
index 1d0d4e734..2e7b93a87 100644
--- a/tensor2tensor/trax/layers/core_test.py
+++ b/tensor2tensor/trax/layers/core_test.py
@@ -82,8 +82,8 @@ def test_dense_param_sharing(self):
     model2 = combinators.Serial(layer, layer)
 
     rng1, rng2 = backend.random.split(backend.random.get_prng(0), 2)
-    params1, _ = model1.initialize((1, 32), onp.float32, rng1)
-    params2, _ = model2.initialize((1, 32), onp.float32, rng2)
+    params1, _ = model1.initialize_once((1, 32), onp.float32, rng1)
+    params2, _ = model2.initialize_once((1, 32), onp.float32, rng2)
     # The first parameters have 2 kernels of size (32, 32).
     self.assertEqual((32, 32), params1[0][0].shape)
     self.assertEqual((32, 32), params1[1][0].shape)
diff --git a/tensor2tensor/trax/layers/demo.ipynb b/tensor2tensor/trax/layers/demo.ipynb
index 1e3d0303f..147e1a71d 100644
--- a/tensor2tensor/trax/layers/demo.ipynb
+++ b/tensor2tensor/trax/layers/demo.ipynb
@@ -156,7 +156,7 @@
         "x  = onp.arange(-7, 8).reshape(3, -1)\n",
         "rng = backend.random.get_prng(0)\n",
         "layer = tl.Relu()\n",
-        "params = layer.initialize(x.shape, x.dtype, rng)\n",
+        "params = layer.initialize_once(x.shape, x.dtype, rng)\n",
         "output = layer(x, params, rng=rng)\n",
         "print(x)\n",
         "print(output)"
diff --git a/tensor2tensor/trax/layers/normalization.py b/tensor2tensor/trax/layers/normalization.py
index 6ff9648ae..5c11a28d6 100644
--- a/tensor2tensor/trax/layers/normalization.py
+++ b/tensor2tensor/trax/layers/normalization.py
@@ -36,7 +36,7 @@ def __init__(self, axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True,
     self._momentum = momentum
     self._mode = mode
 
-  def new_parameters(self, input_shape, input_dtype, rng):
+  def new_params_and_state(self, input_shape, input_dtype, rng):
     """Helper to initialize batch norm params."""
     del input_dtype, rng
     axis = self._axis
@@ -55,7 +55,7 @@ def get_stats_axis(i, d):
     num_batches = np.zeros((), dtype=np.int32)
     return (beta, gamma), (running_mean, running_var, num_batches)
 
-  def call(self, x, params, state, **unused_kwargs):
+  def forward(self, x, params, state, **unused_kwargs):
     """Layer construction function for a batch normalization layer."""
 
     running_mean, running_var, num_batches = state
@@ -113,7 +113,7 @@ def _layer_norm_params(input_shape, input_dtype, rng):
   return (scale, bias)
 
 
-@base.layer(new_parameters=_layer_norm_params)
+@base.layer(new_params_and_state_fn=_layer_norm_params)
 def LayerNorm(x, params, epsilon=1e-6, **unused_kwargs):  # pylint: disable=invalid-name
   (scale, bias) = params
   mean = np.mean(x, axis=-1, keepdims=True)
diff --git a/tensor2tensor/trax/layers/normalization_test.py b/tensor2tensor/trax/layers/normalization_test.py
index 00595b197..0a96a9c0c 100644
--- a/tensor2tensor/trax/layers/normalization_test.py
+++ b/tensor2tensor/trax/layers/normalization_test.py
@@ -45,7 +45,7 @@ def test_batch_norm(self):
     m1 = 11.5  # Mean of this random input.
     v1 = 47.9167  # Variance of this random input.
     layer = normalization.BatchNorm(axis=(0, 1, 2))
-    params, state = layer.initialize(input_shape, input_dtype, rng)
+    params, state = layer.initialize_once(input_shape, input_dtype, rng)
     onp.testing.assert_allclose(state[0], 0)
     onp.testing.assert_allclose(state[1], 1)
     self.assertEqual(state[2], 0)
diff --git a/tensor2tensor/trax/layers/reversible.py b/tensor2tensor/trax/layers/reversible.py
index 1c70d2270..ff56072c9 100644
--- a/tensor2tensor/trax/layers/reversible.py
+++ b/tensor2tensor/trax/layers/reversible.py
@@ -53,20 +53,20 @@ def reverse_and_grad(self, output, grad, params=(), state=(), **kwargs):
       gradient signal for the parameters.
     """
     # Note: jax.vjp does not allow us to use **kwargs in the signature here.
-    def _do_call(x, params):
-      return super(ReversibleLayer, self).call(
+    def _do_forward(x, params):
+      return super(ReversibleLayer, self).forward(
           x, params=params, state=state, **kwargs)[0]
 
     reconstructed_x = self.reverse(output, params, state, **kwargs)
-    _, vjpfun = jax.vjp(_do_call, reconstructed_x, params)
+    _, vjpfun = jax.vjp(_do_forward, reconstructed_x, params)
     x_params_grad = vjpfun(grad)
     return reconstructed_x, x_params_grad
 
   @property
-  def has_custom_grad(self):
+  def has_backward(self):
     return True
 
-  def custom_grad(self, inputs, output, ct, params, state, **kwargs):
+  def backward(self, inputs, output, ct, params, state, **kwargs):
     del inputs
     _, inputs_params_ct = self.reverse_and_grad(output, ct, params, state,
                                                 **kwargs)
@@ -78,7 +78,7 @@ class ReversibleSwap(ReversibleLayer, cb.Swap):
 
   def reverse(self, output, params=(), state=(), **kwargs):
     # Swap is its own inverse, except that reverse doesn't return the state.
-    return self.call(output, params, state, **kwargs)[0]
+    return self.forward(output, params, state, **kwargs)[0]
 
 
 class ReversibleSerial(ReversibleLayer, cb.Serial):
diff --git a/tensor2tensor/trax/learning_rate_test.py b/tensor2tensor/trax/learning_rate_test.py
index 4bc0aafd8..904310026 100644
--- a/tensor2tensor/trax/learning_rate_test.py
+++ b/tensor2tensor/trax/learning_rate_test.py
@@ -50,7 +50,7 @@ def _make_schedule(
     )
     rng = jax_random.get_prng(seed=0)
     obs_dim = len(observation_metrics)
-    (params, state) = net.initialize((1, 1, obs_dim), np.float32, rng)
+    (params, state) = net.initialize_once((1, 1, obs_dim), np.float32, rng)
     policy_dir = self.get_temp_dir()
     # Optimizer slots should not be used for anything.
     slots = None
diff --git a/tensor2tensor/trax/models/atari_cnn_test.py b/tensor2tensor/trax/models/atari_cnn_test.py
index 007ed6557..66b8d59a4 100644
--- a/tensor2tensor/trax/models/atari_cnn_test.py
+++ b/tensor2tensor/trax/models/atari_cnn_test.py
@@ -37,7 +37,7 @@ def test_computes(self):
         hidden_sizes=hidden_size, output_size=output_size)
     B, T, OBS = 2, 2, (28, 28, 3)  # pylint: disable=invalid-name
     rng_key, key = jax_random.split(rng_key)
-    params, state = model.initialize((1, 1) + OBS, onp.float32, key)
+    params, state = model.initialize_once((1, 1) + OBS, onp.float32, key)
     x = onp.arange(B * (T + 1) * functools.reduce(op.mul, OBS)).reshape(
         B, T + 1, *OBS)
     rng_key, key = jax_random.split(rng_key)
@@ -55,7 +55,7 @@ def test_computes(self):
         hidden_sizes=hidden_size, output_size=output_size)
     B, T, OBS = 2, 2, 3  # pylint: disable=invalid-name
     rng_key, key = jax_random.split(rng_key)
-    params, state = model.initialize((1, 1, OBS), onp.float32, key)
+    params, state = model.initialize_once((1, 1, OBS), onp.float32, key)
     x = onp.arange(B * (T + 1) * OBS).reshape(
         B, T + 1, OBS)
     rng_key, key = jax_random.split(rng_key)
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 867341101..daa398edb 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -55,7 +55,7 @@ def __init__(self, layer, n_sections=1, check_shapes=True):
     self._check_shapes = check_shapes
     self._n_sections = n_sections
 
-  def call(self, inputs, params=(), state=(), **kwargs):
+  def forward(self, inputs, params=(), state=(), **kwargs):
     rngs = _pop_rng_and_split(kwargs, len(inputs))
     results = [self._layer(x, params=params, state=state, rng=r, **kwargs)
                for x, r in zip(inputs, rngs)]
@@ -64,14 +64,14 @@ def call(self, inputs, params=(), state=(), **kwargs):
     result_states = result_states[0]
     return tuple(result_outputs), tuple(result_states)
 
-  def new_parameters(self, input_shape, input_dtype, rng):
+  def new_params_and_state(self, input_shape, input_dtype, rng):
     first_shape = input_shape[0]
     if self._check_shapes:
       for shape in input_shape:
         if shape != first_shape:
           raise ValueError('Map layer can only be applied to list of elements '
                            'with the same shapes. Shapes: %s' % str(shape))
-    return self._layer.initialize(first_shape, input_dtype[0], rng)
+    return self._layer.initialize_once(first_shape, input_dtype[0], rng)
 
 
 @tl.layer()
@@ -115,12 +115,12 @@ def __init__(self, n_sections=2, axis=-1):
     self._n_sections = n_sections
     self._axis = axis
 
-  def call(self, inputs, params=(), state=(), **kwargs):
+  def forward(self, inputs, params=(), state=(), **kwargs):
     del params, kwargs
     res = tuple(backend.numpy.split(inputs, self._n_sections, self._axis))
     return res, state
 
-  def new_parameters(self, input_shapes, input_dtype, rng):
+  def new_params_and_state(self, input_shapes, input_dtype, rng):
     return (), ()
 
 
@@ -144,10 +144,10 @@ def __init__(self, n_sections=2, axis=-2):
     self._n_sections = n_sections
     self._axis = axis
 
-  def new_parameters(self, input_shape, input_dtype, rng):
+  def new_params_and_state(self, input_shape, input_dtype, rng):
     return (), ()
 
-  def call(self, inputs, params=(), state=(), **kwargs):
+  def forward(self, inputs, params=(), state=(), **kwargs):
     del params, kwargs
     x1, x2 = inputs
 
@@ -254,21 +254,21 @@ def call_compute_residual(x, params):
 
 
 class ApplyAttentionWrapper(tl.Parallel):
-  """Same as tl.Parallel(attention, [], []), but implements call_and_grad."""
+  """Like tl.Parallel(attention, [], []) but implements forward_and_backward."""
 
   def __init__(self, attention):
-    assert hasattr(attention, 'call_and_grad')
+    assert hasattr(attention, 'forward_and_backward')
     super(ApplyAttentionWrapper, self).__init__(attention, [], [])
     self.attention = attention
 
-  def call_and_grad(self, inputs, ct, **kwargs):
+  def forward_and_backward(self, inputs, ct, **kwargs):
     # Simultaneous forward pass and backprop through the attention mechanism.
     qkv = inputs[:3]
     passthrough = inputs[3:]
     out_ct = ct[0]
     passthrough_ct = ct[1:]
 
-    out, qkv_ct = self.attention.call_and_grad(qkv, out_ct, **kwargs)
+    out, qkv_ct = self.attention.forward_and_backward(qkv, out_ct, **kwargs)
     return (out,) + passthrough, qkv_ct + passthrough_ct
 
 
@@ -285,9 +285,10 @@ class ReversibleAttentionHalfResidual(tl.ReversibleLayer, tl.Serial):
   consists of reshaping and dense linear layers), which allows the following
   optimization. We can back-propagate the gradient signal from the output of
   ReversibleAttentionHalfResidual to the output of the "attention" portion based
-  only on the network parameters. Then, attention.call_and_grad can be used to
-  recover the output of the "attention" portion while simultaneously performing
-  the backward pass, which allows shared computation between the two directions.
+  only on the network parameters. Then, attention.forward_and_backward can be
+  used to recover the output of the "attention" portion while simultaneously
+  performing the backward pass, which allows shared computation between the two
+  directions.
   """
 
   def __init__(self, pre_attention, attention, post_attention):
@@ -297,7 +298,7 @@ def __init__(self, pre_attention, attention, post_attention):
         tl.Swap(),
         tl.Parallel(pre_attention, [], []),
     ])
-    assert hasattr(attention, 'call_and_grad')
+    assert hasattr(attention, 'forward_and_backward')
     self.attention = ApplyAttentionWrapper(attention)
     self.post_attention = tl.Parallel(post_attention, [], [])
 
@@ -360,7 +361,8 @@ def call_post_attention(x):
     (ct,) = post_attention_vjpfun(ct)
 
     # Simultaneous forward pass and backprop through the attention mechanism
-    stack, ct = self.attention.call_and_grad(stack, ct, rng=rngs[1], **kwargs)
+    stack, ct = self.attention.forward_and_backward(stack, ct, rng=rngs[1],
+                                                    **kwargs)
     assert not jax.tree_util.tree_leaves(params[1])
     attention_params_ct = params[1]  # This is valid when params is empty.
 
diff --git a/tensor2tensor/trax/models/transformer_test.py b/tensor2tensor/trax/models/transformer_test.py
index 5d578d4bc..760a400cb 100644
--- a/tensor2tensor/trax/models/transformer_test.py
+++ b/tensor2tensor/trax/models/transformer_test.py
@@ -67,9 +67,9 @@ def test_transformer_lm_fast_inference(self):
       model_fast = model_fn(mode='predict')
       rng = backend.random.get_prng(0)
       batch_size = 2
-      (params, state_slow) = model_slow.initialize(
+      (params, state_slow) = model_slow.initialize_once(
           input_shapes=(batch_size, 1), input_dtype=np.int32, rng=rng)
-      (_, state_fast) = model_fast.initialize(
+      (_, state_fast) = model_fast.initialize_once(
           input_shapes=(batch_size, 1), input_dtype=np.int32, rng=rng)
 
       max_length = 5
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
index 05af01074..5a348a623 100644
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -89,7 +89,7 @@ def test_policy_and_value_net(self):
         bottom_layers_fn=lambda: [layers.Flatten(n_axes_to_keep=2)],
         two_towers=True,
     )
-    pnv_params, pnv_state = pnv_model.initialize(
+    pnv_params, pnv_state = pnv_model.initialize_once(
         batch_observation_shape, np.float32, self.rng_key)
 
     batch = 2
@@ -449,9 +449,9 @@ def test_combined_loss(self):
         two_towers=True,
     )
 
-    old_params, _ = net.initialize(
+    old_params, _ = net.initialize_once(
         batch_observation_shape, np.float32, key1)
-    new_params, state = net.initialize(
+    new_params, state = net.initialize_once(
         batch_observation_shape, np.float32, key2)
 
     # Generate a batch of observations.
@@ -579,7 +579,7 @@ def test_inits_policy_by_world_model_checkpoint(self):
         "rng": rng,
     }
     model = models.TransformerLM(vocab_size=4, **transformer_kwargs)
-    (model_params, _) = model.initialize(**init_kwargs)
+    (model_params, _) = model.initialize_once(**init_kwargs)
     policy = ppo.policy_and_value_net(
         n_actions=3,
         n_controls=2,
@@ -589,7 +589,7 @@ def test_inits_policy_by_world_model_checkpoint(self):
         ),
         two_towers=False,
     )
-    (policy_params, policy_state) = policy.initialize(**init_kwargs)
+    (policy_params, policy_state) = policy.initialize_once(**init_kwargs)
     output_dir = self.get_temp_dir()
     # Initialize state by restoring from a nonexistent checkpoint.
     trax_state = trax.restore_state(output_dir)
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index 84e7c642c..b58febe5a 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -190,7 +190,7 @@ def __init__(self,
     self._policy_and_value_net_apply = jit(policy_and_value_net)
     (batch_obs_shape, obs_dtype) = self._batch_obs_shape_and_dtype
     policy_and_value_net_params, self._model_state = (
-        policy_and_value_net.initialize(batch_obs_shape, obs_dtype, key1))
+        policy_and_value_net.initialize_once(batch_obs_shape, obs_dtype, key1))
     if init_policy_from_world_model_output_dir is not None:
       policy_and_value_net_params = ppo.init_policy_from_world_model_checkpoint(
           policy_and_value_net_params, init_policy_from_world_model_output_dir
diff --git a/tensor2tensor/trax/rl/simple_test.py b/tensor2tensor/trax/rl/simple_test.py
index 41d65e5f2..6383fef5a 100644
--- a/tensor2tensor/trax/rl/simple_test.py
+++ b/tensor2tensor/trax/rl/simple_test.py
@@ -249,7 +249,7 @@ def _make_env(
     predict_output = (np.array([[[0.0]]] * batch_size), ())
     mock_model_fn = mock.MagicMock()
     mock_model_fn.return_value.side_effect = itertools.repeat(predict_output)
-    mock_model_fn.return_value.initialize.return_value = ((), ())
+    mock_model_fn.return_value.initialize_once.return_value = ((), ())
 
     return simulated_env_problem.SerializedSequenceSimulatedEnvProblem(
         model=mock_model_fn,
diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index a0c139a69..f989bcb14 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -67,7 +67,7 @@ def __init__(self, model, batch_size, observation_space, action_space,
       model_predict_kwargs = {}
     model_predict = self._model(mode="predict", **model_predict_kwargs)
     self._model_predict = backend.jit(model_predict)
-    self._model_initialize = model_predict.initialize
+    self._model_initialize = model_predict.initialize_once
     self._observation_space = observation_space
     self._action_space = action_space
     self._reward_range = reward_range
diff --git a/tensor2tensor/trax/rl/simulated_env_problem_test.py b/tensor2tensor/trax/rl/simulated_env_problem_test.py
index ac80bb5e2..d56e3545e 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem_test.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem_test.py
@@ -130,7 +130,7 @@ def _make_env(
     mock_model_fn = mock.MagicMock()
     if predict_fn is not None:
       mock_model_fn.return_value = predict_fn
-      mock_model_fn.return_value.initialize.return_value = ((), ())
+      mock_model_fn.return_value.initialize_once.return_value = ((), ())
     return simulated_env_problem.SerializedSequenceSimulatedEnvProblem(
         model=mock_model_fn,
         reward_fn=reward_fn,
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index e16e27969..84763c21b 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -639,7 +639,8 @@ def new_opt_state_and_model_state(input_shape, input_dtype, target_shape,
       # We need to create a new model instance and not reuse `model_train` here,
       # because `m.initialize` puts cached parameter values in `m` and hence the
       # next call of `m.initialize` will give wrong results.
-      params, state = model(mode="train").initialize(full_shape, full_type, rng)
+      params, state = model(mode="train").initialize_once(full_shape, full_type,
+                                                          rng)
       (slots, opt_params) = opt.tree_init(params)
       return (OptState(params, slots, opt_params), state)
     if _is_jit_init():

From beb485c871aa5e7c8d18f2e984a8b57cac79c3d1 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sat, 28 Sep 2019 07:16:04 -0700
Subject: [PATCH 2496/2720] Give Layer.__call__ option to take params & state
 from object itself.

PiperOrigin-RevId: 271745879
---
 tensor2tensor/trax/layers/base.py             | 170 +++++++++++-------
 tensor2tensor/trax/layers/base_test.py        |   4 +-
 tensor2tensor/trax/layers/combinators.py      |  34 +++-
 tensor2tensor/trax/layers/demo.ipynb          |   4 +-
 .../trax/layers/normalization_test.py         |   2 +-
 tensor2tensor/trax/learning_rate.py           |   3 +-
 tensor2tensor/trax/models/atari_cnn_test.py   |  10 +-
 tensor2tensor/trax/rl/ppo.py                  |   2 +-
 tensor2tensor/trax/rl/ppo_test.py             |  14 +-
 tensor2tensor/trax/rl/ppo_trainer.py          |   7 +-
 tensor2tensor/trax/trax.py                    |   5 +-
 tensor2tensor/trax/trax_test.py               |   4 +-
 12 files changed, 168 insertions(+), 91 deletions(-)

diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 5117903a9..908e7bb39 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -33,37 +33,37 @@
 class Layer(object):
   """Base class for composable layers in a deep learning network.
 
-  A layer is a function from zero or more inputs to zero or more outputs,
-  possibly with trainable parameters. A layer is either atomic or composed
-  of sublayers. These aspects of a layer are set via a layer's constructor,
-  and can be inspected via read-only properties:
+  A layer is a part of a trainable network that can compute a function from
+  zero or more inputs to zero or more outputs. It may make use of trainable
+  parameters as well as non-parameter state for its computation. A layer is
+  either atomic or composed of sublayers. All layers provide accessors for
+  these aspects:
 
-    - n_inputs
-    - n_outputs
-    - sublayers
+    - n_inputs: int
+    - n_outputs: int
+    - params: tuple (empty if the layer has no parameters)
+    - state: tuple (empty if the layer has no non-parameter state)
+    - sublayers: tuple (empty if the layer has no sublayers)
 
-  The inputs to a layer are activation tensors, packaged according to how many
-  there are:
+  The inputs to a layer are tensors, packaged according to how many there are:
 
     - n_inputs = 0: an empty tuple ()
-    _ n_inputs = 1: the activation tensor (NOT wrapped in a tuple)
-    _ n_inputs > 1: a tuple of activation tensors
+    - n_inputs = 1: one tensor (NOT wrapped in a tuple)
+    - n_inputs > 1: a tuple of tensors
 
-  (The special treatment for the single-input case is intended as a
-  simplification for layer writers; this design choice may be revisited in the
-  future.)
+  (The special treatment for the single-input case is meant to simplify the
+  work of layer writers; this design choice may be revisited in the future.)
 
-  The outputs from a layer are also activations tensors, packaged the same as
-  layer inputs:
+  The outputs from a layer are also tensors, packaged the same as layer inputs:
 
     - n_outputs = 0: an empty tuple ()
-    _ n_outputs = 1: the activation tensor (NOT wrapped in a tuple)
-    _ n_outputs > 1: a tuple of activation tensors
+    - n_outputs = 1: the tensor (NOT wrapped in a tuple)
+    - n_outputs > 1: a tuple of tensors
 
-  The runtime maintains a data stack with which layer calls are composed. One
-  can therefore view each layer as a function from stack state to stack state,
-  where the function's inputs are a slice from the stack, and the function's
-  outputs are spliced back into the stack.
+  The Trax runtime maintains a data stack with which layer calls are composed.
+  One can therefore view each layer as a function from stack state to stack
+  state, where the function's inputs are a slice from the stack, and the
+  function's outputs are spliced back into the stack.
   """
 
   def __init__(self, n_inputs=1, n_outputs=1):
@@ -71,6 +71,7 @@ def __init__(self, n_inputs=1, n_outputs=1):
     self._n_outputs = n_outputs
     self._sublayers = ()  # Default is no sublayers.
     self._params = ()  # cached parameters
+    self._state = ()
     self._caller = _find_frame(inspect.stack())  # for custom error messages
     self._init_finished = False
 
@@ -80,72 +81,89 @@ def __repr__(self):
     objs = self.sublayers
     if objs:
       objs_str = ', '.join(str(x) for x in objs)
-      return '{}[{},layers=[{}]]'.format(class_str, fields_str, objs_str)
+      return '{}{{{},sublayers=[{}]}}'.format(class_str, fields_str, objs_str)
     else:
-      return '{}[{}]'.format(class_str, fields_str)
+      return '{}{{{}}}'.format(class_str, fields_str)
 
   def forward(self, inputs, params=(), state=(), **kwargs):
-    """Applies this layer to given activation tensors, using trainable params.
+    """Uses this layer as part of a forward pass through the model.
+
+    Authors of new Layer subclasses should override this method to define the
+    forward computation that their layer performs.
 
     Args:
-      inputs: Data tensors, matching the number (n_inputs) expected by this
+      inputs: Input tensors, matching the number (n_inputs) expected by this
           layer. Specifically:
             - n_inputs = 0: an empty tuple ()
-            - n_inputs = 1: a data tensor (NOT wrapped in a tuple)
-            - n_inputs > 1: a tuple of data tensors, with n_inputs items
+            - n_inputs = 1: a tensor (NOT wrapped in a tuple)
+            - n_inputs > 1: a tuple of tensors, with n_inputs items
       params: A tuple of trainable parameters, with one element for this layer
-          and one for each of this layer's sublayers. If a layer (or sublayer)
-          has no trainable parameters, the corresponding params element is an
-          empty tuple.
-      state: start state.
-      **kwargs: Layer-specific keyword args.
+          if this layer has no sublayers, or one for each sublayer if this
+          layer has sublayers. If a layer (or sublayer) has no trainable
+          parameters, the corresponding params element is an empty tuple.
+      state: Layer-specific non-parameter state that can update between batches.
+      **kwargs: Often empty; main current use is to carry a PRNG key for random
+          number generation, using the keyword 'rng'.
 
     Returns:
-      Data tensors, matching the number (n_outputs) promised by this layer.
+      Tensors, matching the number (n_outputs) promised by this layer.
       Specifically:
         - n_outputs = 0: an empty tuple
-        - n_outputs = 1: a data tensor (NOT wrapped in a tuple)
-        - n_outputs > 1: a tuple of data tensors, with n_outputs items
-      A tuple of activation tensors, one for each output.
+        - n_outputs = 1: one tensor (NOT wrapped in a tuple)
+        - n_outputs > 1: a tuple of tensors, with n_outputs items
     """
     raise NotImplementedError
 
   def new_params_and_state(self, input_shapes, input_dtype, rng):
-    """Creates layer-specific parameters based on data shape, dtype and rng.
+    """Returns a (params, state) pair suitable for initializing this layer.
 
-    Args:
-      input_shapes: A tuple, depending on the number of inputs (n_inputs)
-          expected by this layer:
-            - n_inputs = 0: an empty tuple ()
-            - n_inputs = 1: a tuple representing the shape of the input
-            - n_inputs > 1: a tuple of shape tuples, one for each input
-          For example:
-            - 0 inputs: ()
-            - 1 input: (210, 160, 3) [NOTE: no tuple wrapping the shape]
-            - 2 inputs: ((210, 160, 3), (105, 80, 3))
-      input_dtype: numpy dtype of the input.
-      rng: A random number generator.
+    Authors of new Layer subclasses should override this method if their layer
+    uses trainable parameters or has non-parameter state that gets updated
+    between batches. The default implementation works for layers that have
+    no parameters or state.
 
-    Returns:
-      The newly created parameters for this layer.
+    Args:
+      input_shapes: A tuple representing a shape (if this layer takes one input)
+          or a tuple of shapes (if this layer takes more than one input).
+          For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
+      input_dtype: Numpy dtype(s) for each of the inputs.
+      rng: A PRNG key for random number generation.
     """
     raise NotImplementedError
 
   @property
   def n_inputs(self):
-    """Specifies how many data tensors this layer expects as input."""
+    """Returns how many tensors this layer expects as input."""
     return self._n_inputs
 
   @property
   def n_outputs(self):
-    """Specifies how many data tensors this layer promises as output."""
+    """Returns how many tensors this layer promises as output."""
     return self._n_outputs
 
   @property
   def sublayers(self):
-    """Returns the sublayers contained in / managed by this layer."""
+    """Returns a tuple containing this layer's sublayers; may be empty."""
     return self._sublayers
 
+  @property
+  def params(self):
+    """Returns a tuple containing this layer's parameters; may be empty."""
+    return self._params
+
+  @params.setter
+  def params(self, params):
+    self._params = params
+
+  @property
+  def state(self):
+    """Returns a tuple containing this layer's state; may be empty."""
+    return self._state
+
+  @state.setter
+  def state(self, state):
+    self._state = state
+
   @property
   def has_backward(self):
     """Whether to use custom gradients (in which case, see below)."""
@@ -232,6 +250,7 @@ def initialize_once(self, input_shapes, input_dtype, rng):
       if not self._init_finished:
         self._init_finished = True
         self._params = params
+        self._state = state
       else:
         params = ()
       return (params, state)
@@ -244,7 +263,33 @@ def initialize_once(self, input_shapes, input_dtype, rng):
   _STASH_IN = None
   _STASH_OUT = None
 
-  def __call__(self, x, params=(), state=(), **kwargs):
+  def __call__(self, x, **kwargs):
+    """Makes Layer instances callable; for use in tests or interactive settings.
+
+    This convenience method helps library users play with, test, or otherwise
+    probe the behavior of layers outside of a full training environment. It
+    presents the layer as callable function from inputs to outputs, with the
+    option of manually specifying parameters and non-parameter state per
+    individual call. For convenience, parameters and non-parameter state are
+    cached per layer instance, starting from default values of () and (), and
+    acquiring non-empty values either by initialization or from values
+    explicitly provided via the params and state keyword arguments.
+
+    Args:
+      x: 0 or more input tensors, formatted the same as the inputs to
+          Layer.forward.
+      **kwargs: Additional keyword arguments if needed/desired for this layer.
+          Three possible keyword arguments are especially relevant:
+            - params=... will override any cached params values
+            - state=... will override any cached state values
+            - rng=... will supply a PRNG key for use by the layer
+
+    Returns:
+      An (outputs, state) tuple. The outputs part of the tuple is formatted the
+          same as the outputs from Layer.forward.
+    """
+    params = kwargs.pop('params', self.params)
+    state = kwargs.pop('state', self.state)
     try:
       # If params are nothing, we may be reusing this layer.
       # Use the cached parameters to calculate the value.
@@ -524,15 +569,14 @@ def _is_tuple_of_shapes(shape):
   return isinstance(shape, tuple) and isinstance(shape[0], tuple)
 
 
-def check_shape_agreement(layer_fn, input_shapes, integer_inputs=False):
-  """Checks if the layer's forward output agrees its pseudo_forward predictions.
+def check_shape_agreement(layer_obj, input_shapes, integer_inputs=False):
+  """Checks if the layer's call output agrees its pseudo_forward predictions.
 
   This function helps test layer mechanics and inter-layer connections that
   aren't dependent on specific data values.
 
   Args:
-    layer_fn: A Layer instance, viewed as a function from input shapes to
-        output shapes.
+    layer_obj: A Layer instance.
     input_shapes: A tuple representing a shape (if the layer takes one input)
         or a tuple of shapes (if this layer takes more than one input).
         For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
@@ -550,15 +594,15 @@ def check_shape_agreement(layer_fn, input_shapes, integer_inputs=False):
     input_dtype = tuple(input_dtype for _ in input_shapes)
   else:
     pseudo_data = ShapeType(input_shapes, input_dtype)
-  params, state = layer_fn.initialize_once(input_shapes, input_dtype, rng1)
-  pseudo_output, _ = layer_fn.pseudo_forward(pseudo_data, params, state)
+  params, state = layer_obj.initialize_once(input_shapes, input_dtype, rng1)
+  pseudo_output, _ = layer_obj.pseudo_forward(pseudo_data, params, state)
   if isinstance(pseudo_output, tuple):
     output_shape = tuple(x.shape for x in pseudo_output)
   else:
     output_shape = pseudo_output.shape
 
   random_input = _random_values(input_shapes, rng2, integer_inputs)
-  real_output, _ = layer_fn(random_input, params, state=state, rng=rng3)
+  real_output, _ = layer_obj(random_input, params=params, state=state, rng=rng3)
   result_shape = shapes(real_output)
 
   msg = 'output shape %s != real result shape %s' % (output_shape, result_shape)
diff --git a/tensor2tensor/trax/layers/base_test.py b/tensor2tensor/trax/layers/base_test.py
index aa49869fc..a8bb04d11 100644
--- a/tensor2tensor/trax/layers/base_test.py
+++ b/tensor2tensor/trax/layers/base_test.py
@@ -59,7 +59,7 @@ def backward(self, inputs, output, ct, params, state, **kwargs):
     input_shape = (9, 17)
     random_input = backend.random.uniform(rng, input_shape, minval=-1.0,
                                           maxval=1.0)
-    f = lambda x: backend.numpy.mean(layer(x, params, rng=rng)[0])
+    f = lambda x: backend.numpy.mean(layer(x, params=params, rng=rng)[0])
     grad = backend.grad(f)(random_input)
     self.assertEqual(grad.shape, input_shape)  # Gradient for each input.
     self.assertEqual(sum(sum(grad * grad)), 0.0)  # Each one is 0.
@@ -89,7 +89,7 @@ def backward(self, inputs, output, ct, params, state, **kwargs):
     input_shape = (9, 17)
     random_input = backend.random.uniform(rng, input_shape, minval=-1.0,
                                           maxval=1.0)
-    f = lambda x: backend.numpy.mean(layer(x, params, rng=rng)[0])
+    f = lambda x: backend.numpy.mean(layer(x, params=params, rng=rng)[0])
     grad = backend.grad(f)(random_input)
     self.assertEqual(grad.shape, input_shape)  # Gradient for each input.
     self.assertEqual(sum(sum(grad)), sum(sum(random_input)))  # Same as input.
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index c66680a47..c17455f64 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -156,6 +156,22 @@ def _validate_forward_inputs(self, xs):
           'number of inputs ({}) to Serial.forward less than n_inputs'
           ' ({})'.format(len(xs), self.n_inputs))
 
+  @base.Layer.params.setter
+  def params(self, params):
+    """Recursively sets params on this layer and all sublayers."""
+    self._params = params
+    assert len(params) == self._n_layers
+    for layer, sublayer_params in zip(self.sublayers, params):
+      layer.params = sublayer_params
+
+  @base.Layer.state.setter
+  def state(self, state):
+    """Recursively sets non-param state on this layer and all sublayers."""
+    self._state = state
+    assert len(state) == self._n_layers
+    for layer, sublayer_state in zip(self.sublayers, state):
+      layer.state = sublayer_state
+
   def forward(self, xs, params=(), state=(), **kwargs):
     self._validate_forward_inputs(xs)
     rngs = _pop_rng_and_split(kwargs, self._n_layers)
@@ -182,7 +198,7 @@ def forward(self, xs, params=(), state=(), **kwargs):
         inputs = stack[0]
       else:
         inputs = stack[:n_in]
-      outputs, s = layer(inputs, p, state=s, rng=rng, **kwargs)
+      outputs, s = layer(inputs, params=p, state=s, rng=rng, **kwargs)
       new_state.append(s)
 
       # Push outputs onto remaining stack (if any).
@@ -423,6 +439,22 @@ def _allot_to_sublayers(self, inputs):
       start = end
     return tuple(sub_inputs)
 
+  @base.Layer.params.setter
+  def params(self, params):
+    """Recursively sets params on this layer and all sublayers."""
+    self._params = params
+    assert len(params) == self._n_layers
+    for layer, sublayer_params in zip(self.sublayers, params):
+      layer.params = sublayer_params
+
+  @base.Layer.state.setter
+  def state(self, state):
+    """Recursively sets non-param state on this layer and all sublayers."""
+    self._state = state
+    assert len(state) == self._n_layers
+    for layer, sublayer_state in zip(self.sublayers, state):
+      layer.state = sublayer_state
+
   def forward(self, inputs, params=(), state=(), **kwargs):
     n_layers, layers = self._n_layers, self.sublayers
     sublayer_inputs = self._allot_to_sublayers(inputs)
diff --git a/tensor2tensor/trax/layers/demo.ipynb b/tensor2tensor/trax/layers/demo.ipynb
index 147e1a71d..ecb186560 100644
--- a/tensor2tensor/trax/layers/demo.ipynb
+++ b/tensor2tensor/trax/layers/demo.ipynb
@@ -156,8 +156,8 @@
         "x  = onp.arange(-7, 8).reshape(3, -1)\n",
         "rng = backend.random.get_prng(0)\n",
         "layer = tl.Relu()\n",
-        "params = layer.initialize_once(x.shape, x.dtype, rng)\n",
-        "output = layer(x, params, rng=rng)\n",
+        "params, _ = layer.initialize_once(x.shape, x.dtype, rng)\n",
+        "output = layer(x)\n",
         "print(x)\n",
         "print(output)"
       ]
diff --git a/tensor2tensor/trax/layers/normalization_test.py b/tensor2tensor/trax/layers/normalization_test.py
index 0a96a9c0c..e430d9f0d 100644
--- a/tensor2tensor/trax/layers/normalization_test.py
+++ b/tensor2tensor/trax/layers/normalization_test.py
@@ -49,7 +49,7 @@ def test_batch_norm(self):
     onp.testing.assert_allclose(state[0], 0)
     onp.testing.assert_allclose(state[1], 1)
     self.assertEqual(state[2], 0)
-    out, state = layer(inp1, params, state)
+    out, state = layer(inp1, params=params, state=state)
     onp.testing.assert_allclose(state[0], m1 * 0.001)
     onp.testing.assert_allclose(state[1], 0.999 + v1 * 0.001, rtol=1e-6)
     self.assertEqual(state[2], 1)
diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index a0ba582c7..0d0028357 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -239,7 +239,8 @@ def PolicySchedule(
   start_time = time.time()
   # ((log_probs, value_preds), state). We have no way to pass state to the next
   # step, but that should be fine.
-  ((log_probs, _), _) = net(np.array([observations]), params, state, rng=rng)
+  ((log_probs, _), _) = (
+      net(np.array([observations]), params=params, state=state, rng=rng))
   logging.vlog(
       1, "Running the policy took %0.2f sec.", time.time() - start_time
   )
diff --git a/tensor2tensor/trax/models/atari_cnn_test.py b/tensor2tensor/trax/models/atari_cnn_test.py
index 66b8d59a4..88e3eb52d 100644
--- a/tensor2tensor/trax/models/atari_cnn_test.py
+++ b/tensor2tensor/trax/models/atari_cnn_test.py
@@ -37,11 +37,10 @@ def test_computes(self):
         hidden_sizes=hidden_size, output_size=output_size)
     B, T, OBS = 2, 2, (28, 28, 3)  # pylint: disable=invalid-name
     rng_key, key = jax_random.split(rng_key)
-    params, state = model.initialize_once((1, 1) + OBS, onp.float32, key)
+    _, _ = model.initialize_once((1, 1) + OBS, onp.float32, key)
     x = onp.arange(B * (T + 1) * functools.reduce(op.mul, OBS)).reshape(
         B, T + 1, *OBS)
-    rng_key, key = jax_random.split(rng_key)
-    y, _ = model(x, params, state=state, rng=key)
+    y, _ = model(x)
     self.assertEqual((B, T + 1, output_size), y.shape)
 
 
@@ -55,11 +54,10 @@ def test_computes(self):
         hidden_sizes=hidden_size, output_size=output_size)
     B, T, OBS = 2, 2, 3  # pylint: disable=invalid-name
     rng_key, key = jax_random.split(rng_key)
-    params, state = model.initialize_once((1, 1, OBS), onp.float32, key)
+    _, _ = model.initialize_once((1, 1, OBS), onp.float32, key)
     x = onp.arange(B * (T + 1) * OBS).reshape(
         B, T + 1, OBS)
-    rng_key, key = jax_random.split(rng_key)
-    y, _ = model(x, params, state=state, rng=key)
+    y, _ = model(x)
     self.assertEqual((B, T + 1, output_size), y.shape)
 
 
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index 7b591ed9d..5e631a2fe 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -723,7 +723,7 @@ def combined_loss(new_params,
   """Computes the combined (clipped loss + value loss) given observations."""
   (log_probab_actions_new, value_predictions_new), state = (
       policy_and_value_net_apply(
-          padded_observations, new_params, state, rng=rng))
+          padded_observations, params=new_params, state=state, rng=rng))
 
   (loss, component_losses, summaries) = combined_loss_given_predictions(
       log_probab_actions_new,
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
index 5a348a623..b9f9e83d2 100644
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -89,14 +89,14 @@ def test_policy_and_value_net(self):
         bottom_layers_fn=lambda: [layers.Flatten(n_axes_to_keep=2)],
         two_towers=True,
     )
-    pnv_params, pnv_state = pnv_model.initialize_once(
+    _, _ = pnv_model.initialize_once(
         batch_observation_shape, np.float32, self.rng_key)
 
     batch = 2
     time_steps = 10
     batch_of_observations = np.random.uniform(
         size=(batch, time_steps) + observation_shape)
-    pnv_output, _ = pnv_model(batch_of_observations, pnv_params, pnv_state)
+    pnv_output, _ = pnv_model(batch_of_observations)
 
     # Output is a list, first is probab of actions and the next is value output.
     self.assertEqual(2, len(pnv_output))
@@ -462,10 +462,10 @@ def test_combined_loss(self):
     mask = np.ones_like(rewards)
 
     # Just test that this computes at all.
-    (new_log_probabs, value_predictions_new), _ = net(observations, new_params,
-                                                      state)
-    (old_log_probabs, value_predictions_old), _ = net(observations, old_params,
-                                                      state)
+    (new_log_probabs, value_predictions_new), _ = (
+        net(observations, param=new_params, state=state))
+    (old_log_probabs, value_predictions_old), _ = (
+        net(observations, params=old_params, state=state))
 
     gamma = 0.99
     lambda_ = 0.95
@@ -602,7 +602,7 @@ def test_inits_policy_by_world_model_checkpoint(self):
     )
     # Try to run the policy with new parameters.
     observations = np.zeros((1, 100), dtype=np.int32)
-    policy(observations, new_policy_params, state=policy_state, rng=rng)
+    policy(observations, params=new_policy_params, state=policy_state, rng=rng)
 
   def test_shuffled_index_batches_generates_valid_batch(self):
     dataset_size = 16
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index b58febe5a..49a5f539d 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -585,8 +585,8 @@ def train_epoch(self, evaluate=True):
       (log_probab_actions_new, _), self._model_state = (
           self._policy_and_value_net_apply(
               padded_observations,
-              self._policy_and_value_net_params,
-              self._model_state,
+              params=self._policy_and_value_net_params,
+              state=self._model_state,
               rng=k2))
 
       action_mask = np.dot(
@@ -798,7 +798,8 @@ def _get_predictions(self, observations, state, rng=None):
     key, key1 = jax_random.split(rng, num=2)
 
     (log_probs, value_preds), state = self._policy_and_value_net_apply(
-        observations, self._policy_and_value_net_params, state, rng=key1)
+        observations, params=self._policy_and_value_net_params, state=state,
+        rng=key1)
 
     return log_probs, value_preds, state, key
 
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 84763c21b..c6d1e4b79 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -146,7 +146,8 @@ def loss(params, batch, model_predict, state, rng, has_weights):
       [inputs, targets])
   # Call model, predictions will be the returned stack, usually consisting of
   # the prediction tensor and the targets.
-  predictions, state = model_predict(model_input, params, state, rng=rng)
+  predictions, state = model_predict(model_input, params=params, state=state,
+                                     rng=rng)
   predictions = get_preds(predictions)
   predictions, targets, weights = _make_list(predictions, targets, weights)
   xent = []
@@ -420,7 +421,7 @@ def _jit_predict_fn(model_predict, n_devices, jit=True):
   # Multi-devices, pmap and run.
   @functools.partial(backend.pmap, axis_name="batch")
   def mapped_predict(x, params, state, rng):
-    return model_predict(x, params, state, rng=rng)
+    return model_predict(x, params=params, state=state, rng=rng)
 
   def predict(x, params=(), state=(), rng=None):
     """Predict function jited and parallelized as requested."""
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index a6ada4ca8..13a153208 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -124,7 +124,7 @@ def model_fn(mode="train"):
       # Predict with final params
       inputs = inputs(1).train_stream()
       model = layers.Serial(model_fn())
-      model(next(inputs)[0], state.opt_state.params)
+      model(next(inputs)[0], params=state.opt_state.params)
 
   @parameterized.parameters(BACKENDS)
   def test_train_eval_predict_sm3(self, backend_name):
@@ -160,7 +160,7 @@ def test_train_eval_predict_sm3(self, backend_name):
       # Predict with final params
       inputs = inputs(1).train_stream()
       model = layers.Serial(model_fn())
-      model(next(inputs)[0], state.opt_state.params)
+      model(next(inputs)[0], params=state.opt_state.params)
 
   @parameterized.parameters(BACKENDS)
   def test_train_restart(self, backend_name):

From 5f359eee7ede1ff1b0623a1215a41c2893223faa Mon Sep 17 00:00:00 2001
From: Pierre Ruyssen <pierrot@google.com>
Date: Sat, 28 Sep 2019 12:28:55 -0700
Subject: [PATCH 2497/2720] Internal change

PiperOrigin-RevId: 271766307
---
 tensor2tensor/data_generators/moving_mnist.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/moving_mnist.py b/tensor2tensor/data_generators/moving_mnist.py
index f2140089f..507207623 100644
--- a/tensor2tensor/data_generators/moving_mnist.py
+++ b/tensor2tensor/data_generators/moving_mnist.py
@@ -122,7 +122,8 @@ def map_fn(self, image, label):
     return sequence.image_sequence
 
   def get_train_iterator(self):
-    mnist_ds = tfds.load("mnist", split=tfds.Split.TRAIN, as_supervised=True)
+    mnist_ds = tfds.load("mnist:3.*.*", split=tfds.Split.TRAIN,
+                         as_supervised=True)
     mnist_ds = mnist_ds.repeat()
     moving_mnist_ds = mnist_ds.map(self.map_fn).batch(2)
     moving_mnist_ds = moving_mnist_ds.map(lambda x: tf.reduce_max(x, axis=0))

From 8c41386a7a9eaa079f9c50e0a4524bf01fa408c2 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sat, 28 Sep 2019 18:12:31 -0700
Subject: [PATCH 2498/2720] Change Layer.__call__ to return results rather than
 (results, state).

PiperOrigin-RevId: 271788568
---
 tensor2tensor/trax/layers/README.md           | 12 +++---
 tensor2tensor/trax/layers/attention.py        | 19 +++++-----
 tensor2tensor/trax/layers/attention_test.py   |  4 +-
 tensor2tensor/trax/layers/base.py             | 20 +++++-----
 tensor2tensor/trax/layers/base_test.py        | 12 ++----
 tensor2tensor/trax/layers/combinators.py      | 14 +++----
 tensor2tensor/trax/layers/convolution.py      |  2 +-
 tensor2tensor/trax/layers/core.py             | 38 +++++++++----------
 tensor2tensor/trax/layers/core_test.py        |  2 +-
 .../trax/layers/normalization_test.py         |  6 ++-
 tensor2tensor/trax/learning_rate.py           |  2 +-
 tensor2tensor/trax/models/atari_cnn_test.py   |  4 +-
 .../models/research/transformer_revnet.py     | 34 +++++++++--------
 tensor2tensor/trax/models/transformer_test.py | 12 ++----
 tensor2tensor/trax/rl/ppo.py                  |  2 +-
 tensor2tensor/trax/rl/ppo_test.py             |  8 ++--
 tensor2tensor/trax/rl/ppo_trainer.py          |  4 +-
 tensor2tensor/trax/rl/simple_test.py          |  2 +-
 .../trax/rl/simulated_env_problem.py          |  4 +-
 .../trax/rl/simulated_env_problem_test.py     |  4 +-
 tensor2tensor/trax/trax.py                    | 13 +++----
 tensor2tensor/trax/trax_test.py               |  7 ++--
 22 files changed, 108 insertions(+), 117 deletions(-)

diff --git a/tensor2tensor/trax/layers/README.md b/tensor2tensor/trax/layers/README.md
index 5efe6e86b..a9e4deabe 100644
--- a/tensor2tensor/trax/layers/README.md
+++ b/tensor2tensor/trax/layers/README.md
@@ -17,21 +17,21 @@ def new_params_and_state(self, input_shape, rng):
 """Create new parameters given the shape of the input."""
 ```
 
-The base layer class wraps these functions and provides initialization
+The base Layer class wraps these functions and provides initialization
 and call functions to be used as follows.
 
 ```python
 layer = MyLayer()
 x = np.zeros(10)
 rng = random.get_prng(0)
-params = layer.initialize_once(x.shape, x.dtype, rng)
-output = layer(x, params, rng=rng)
+layer.initialize_once(x.shape, x.dtype, rng)
+output = layer(x)
 ```
 
 ## Decorator
 
-To create simple layers, especially ones without parameters and where
-the output shape is the same as the input shape, use the layer decorator.
+To create simple layers, especially ones without parameters, use the layer
+decorator.
 
 ```python
 @base.layer()
@@ -49,7 +49,7 @@ layer = Dense(10)
 shared_parameters_mlp = layers.Serial(layer, layer)
 ```
 For this reason, if you call `layer.initialize_once(...)` for the second time
-on an already initialized layer, it will return `()`.
+on an already initialized layer, it will not re-initialize the layer.
 
 ## Core layers
 
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index afdeb8082..7a0e31359 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -79,7 +79,7 @@ def __init__(self, max_len=2048, mode='train'):
     self._max_len = max_len
     self._mode = mode
 
-  def forward(self, inputs, params, state, **kwargs):
+  def forward(self, inputs, params=(), state=(), **kwargs):
     if self._mode in ('train', 'eval'):
       x = inputs
       symbol_size = np.shape(x)[1]
@@ -100,9 +100,9 @@ def new_params_and_state(self, input_shape, input_dtype, rng):
     pe[:, 0::2] = onp.sin(position * div_term)
     pe[:, 1::2] = onp.cos(position * div_term)
     pe = pe[onp.newaxis, :, :]  # [1, self._max_len, d_feature]
-    pe = np.array(pe)  # These are trainable parameters, initialized as above.
+    params = np.array(pe)  # These are trainable parameters, initialized above.
     state = 0 if self._mode == 'predict' else ()
-    return (pe, state)
+    return params, state
 
 
 def DotProductAttention(query, key, value, mask, dropout, mode, rng):
@@ -258,7 +258,7 @@ def __init__(self, initializer=init.RandomNormalInitializer(0.01)):
     super(ShiftRightLearned, self).__init__()
     self._initializer = initializer
 
-  def forward(self, x, params, state, **kwargs):
+  def forward(self, x, params=(), state=(), **kwargs):
     del kwargs
     c = backend.numpy.reshape(params, [1, 1, -1])
     c += backend.numpy.zeros((x.shape[0], 1, x.shape[2]), dtype=x.dtype)
@@ -287,7 +287,7 @@ def __init__(self, n_heads=1, d_head=64,
     # implementation, and shouldn't have an effect on modeling quality.
     # Note that AttentionQKV above is different in that it uses a bias term.
 
-  def forward(self, x, params, state, **kwargs):
+  def forward(self, x, params=(), state=(), **kwargs):
     del kwargs
     seqlen = x.shape[1]
     res = np.dot(x, params)
@@ -321,7 +321,7 @@ def __init__(self, n_heads=1, d_model=1024,
     # implementation, and shouldn't have an effect on modeling quality.
     # Note that AttentionQKV above is different in that it uses a bias term.
 
-  def forward(self, x, params, state, **kwargs):
+  def forward(self, x, params=(), state=(), **kwargs):
     del kwargs
     seqlen = x.shape[1]
     d_head = x.shape[2]
@@ -372,9 +372,6 @@ def forward_and_backward(self, inputs, grad, **kwargs):
     """
     raise NotImplementedError()
 
-  def new_params_and_state(self, input_shapes, input_dtype, rng):
-    return (), ()
-
 
 class DotProductCausalAttention(BaseCausalAttention):
   """A standard (non-memory-efficient) dot product attention implementation."""
@@ -447,7 +444,9 @@ def initial_state(shape, dtype):
     )
     mask = np.zeros((batch_size, 1, max_len))
     index = 0
-    return (), (k, v, mask, index)
+    params = ()
+    state = (k, v, mask, index)
+    return params, state
 
 
 class MemoryEfficientCausalAttention(BaseCausalAttention):
diff --git a/tensor2tensor/trax/layers/attention_test.py b/tensor2tensor/trax/layers/attention_test.py
index 6d3f0576a..e3452e49f 100644
--- a/tensor2tensor/trax/layers/attention_test.py
+++ b/tensor2tensor/trax/layers/attention_test.py
@@ -31,7 +31,7 @@ def test_shift_right(self):
     # Test shifts right on axis=1
     layer = attention.ShiftRight()
     input_np = onp.arange(2*3*3).reshape(2, 3, 3)
-    output_np, _ = layer(input_np)
+    output_np = layer(input_np)
     self.assertEqual(input_np.shape, output_np.shape)
     self.assertAllEqual(onp.array([[[0, 0, 0],
                                     [0, 1, 2],
@@ -50,7 +50,7 @@ def test_shift_right_float(self):
     input_np /= 2.0
     self.assertEqual(input_np.dtype, onp.float32)
 
-    output_np, _ = layer(input_np)
+    output_np = layer(input_np)
     self.assertEqual(input_np.shape, output_np.shape)
     self.assertEqual(output_np.dtype, onp.float32)
 
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 908e7bb39..0136f1afd 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -114,7 +114,7 @@ def forward(self, inputs, params=(), state=(), **kwargs):
     """
     raise NotImplementedError
 
-  def new_params_and_state(self, input_shapes, input_dtype, rng):
+  def new_params_and_state(self, input_shape, input_dtype, rng):
     """Returns a (params, state) pair suitable for initializing this layer.
 
     Authors of new Layer subclasses should override this method if their layer
@@ -123,13 +123,14 @@ def new_params_and_state(self, input_shapes, input_dtype, rng):
     no parameters or state.
 
     Args:
-      input_shapes: A tuple representing a shape (if this layer takes one input)
+      input_shape: A tuple representing a shape (if this layer takes one input)
           or a tuple of shapes (if this layer takes more than one input).
           For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
       input_dtype: Numpy dtype(s) for each of the inputs.
       rng: A PRNG key for random number generation.
     """
-    raise NotImplementedError
+    del input_shape, input_dtype, rng
+    return (), ()
 
   @property
   def n_inputs(self):
@@ -285,8 +286,8 @@ def __call__(self, x, **kwargs):
             - rng=... will supply a PRNG key for use by the layer
 
     Returns:
-      An (outputs, state) tuple. The outputs part of the tuple is formatted the
-          same as the outputs from Layer.forward.
+      0 or more output tensors, formatted the same as the outputs from
+          Layer.forward.
     """
     params = kwargs.pop('params', self.params)
     state = kwargs.pop('state', self.state)
@@ -302,10 +303,11 @@ def __call__(self, x, **kwargs):
         self._params = params
 
       if not self.has_backward or Layer._STASH_IN is not None:
-        return self.forward(x, params=params, state=state, **kwargs)
+        outputs, s = self.forward(x, params=params, state=state, **kwargs)
       else:
-        return self._do_custom_gradients(x, params, state, **kwargs)
-
+        outputs, s = self._do_custom_gradients(x, params, state, **kwargs)
+      self._state = s
+      return outputs
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback()
       raise LayerError(name, 'forward', self._caller, shapes(x), trace)
@@ -602,7 +604,7 @@ def check_shape_agreement(layer_obj, input_shapes, integer_inputs=False):
     output_shape = pseudo_output.shape
 
   random_input = _random_values(input_shapes, rng2, integer_inputs)
-  real_output, _ = layer_obj(random_input, params=params, state=state, rng=rng3)
+  real_output = layer_obj(random_input, params=params, state=state, rng=rng3)
   result_shape = shapes(real_output)
 
   msg = 'output shape %s != real result shape %s' % (output_shape, result_shape)
diff --git a/tensor2tensor/trax/layers/base_test.py b/tensor2tensor/trax/layers/base_test.py
index a8bb04d11..a2ab373d4 100644
--- a/tensor2tensor/trax/layers/base_test.py
+++ b/tensor2tensor/trax/layers/base_test.py
@@ -38,14 +38,10 @@ def test_custom_zero_grad(self):
 
     class IdWithZeroGrad(base.Layer):
 
-      def forward(self, x, params, **kwargs):
-        del params, kwargs
+      def forward(self, x, params=(), state=(), **kwargs):
+        del kwargs
         return x, ()
 
-      def new_params_and_state(self, input_shapes, input_dtype, rng):
-        del input_shapes, input_dtype, rng
-        return (), ()
-
       @property
       def has_backward(self):
         return True
@@ -68,8 +64,8 @@ def test_custom_id_grad(self):
 
     class IdWithIdGrad(base.Layer):
 
-      def forward(self, x, params, **kwargs):
-        del params, kwargs
+      def forward(self, x, params=(), state=(), **kwargs):
+        del kwargs
         return x, ()
 
       def new_params_and_state(self, input_shapes, input_dtype, rng):
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index c17455f64..8d4b8f950 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -111,6 +111,9 @@ class Serial(base.Layer):
 
     - takes layer k's N_out return values (N_out = k.n_outputs) and pushes
       them onto the data stack.
+
+  A Serial instance with no sublayers acts as a special-case (but useful)
+  1-input 1-output no-op.
   """
 
   def __init__(self, *layers):
@@ -198,8 +201,8 @@ def forward(self, xs, params=(), state=(), **kwargs):
         inputs = stack[0]
       else:
         inputs = stack[:n_in]
-      outputs, s = layer(inputs, params=p, state=s, rng=rng, **kwargs)
-      new_state.append(s)
+      outputs = layer(inputs, params=p, state=s, rng=rng, **kwargs)
+      new_state.append(layer.state)
 
       # Push outputs onto remaining stack (if any).
       if n_in < _count_items(stack):
@@ -341,9 +344,6 @@ def __init__(self, n_items=2, axis=-1):
     self._n_items = n_items
     self._axis = axis
 
-  def new_params_and_state(self, input_shape, input_dtype, rng):
-    return (), ()
-
   def forward(self, xs, params=(), state=(), **kwargs):
     del params, kwargs
     return backend.numpy.concatenate(xs, self._axis), state
@@ -467,12 +467,12 @@ def forward(self, inputs, params=(), state=(), **kwargs):
     new_state = []
     for layer, x, p, s, r in zip(layers, sublayer_inputs, params, state, rngs):
       # Note that zip silently truncates its result if lengths don't match.
-      sub_outputs, s = layer(x, params=p, state=s, rng=r, **kwargs)
+      sub_outputs = layer(x, params=p, state=s, rng=r, **kwargs)
       if layer.n_outputs == 1:
         outputs.append(sub_outputs)
       else:
         outputs.extend(sub_outputs)
-      new_state.append(s)
+      new_state.append(layer.state)
     output = outputs[0] if self.n_outputs == 1 else tuple(outputs)
     return output, new_state
 
diff --git a/tensor2tensor/trax/layers/convolution.py b/tensor2tensor/trax/layers/convolution.py
index d54a8d6f3..6a7239feb 100644
--- a/tensor2tensor/trax/layers/convolution.py
+++ b/tensor2tensor/trax/layers/convolution.py
@@ -112,7 +112,7 @@ def __init__(self,
         kernel_initializer=kernel_initializer,
         bias_initializer=bias_initializer)
 
-  def forward(self, x, params=(), **kwargs):
+  def forward(self, x, params=(), state=(), **kwargs):
     assert self._padding == 'VALID'
     # Left pad with 0s. Applying an unmasked valid convolution on top of this
     # yields a causal convolution.
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index f815590bf..a36c74a84 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -88,16 +88,14 @@ def Exp(x, **unused_kwargs):
 
 
 @base.layer()
-def LogSoftmax(x, params, axis=-1, **kwargs):
+def LogSoftmax(x, axis=-1, **unused_kwargs):
   """Apply log softmax to x: log-normalize along the given axis."""
-  del params, kwargs
   return x - backend.logsumexp(x, axis, keepdims=True)
 
 
 @base.layer()
-def Softmax(x, params, axis=-1, **kwargs):
+def Softmax(x, axis=-1, **unused_kwargs):
   """Apply softmax to x: exponentiate and normalize along the given axis."""
-  del params, kwargs
   return np.exp(x - backend.logsumexp(x, axis, keepdims=True))
 
 
@@ -112,7 +110,7 @@ def ToFloat(x, **unused_kwargs):
 
 
 class Dense(base.Layer):
-  """Layer constructor function for a dense (fully-connected) layer."""
+  """A dense (a.k.a. fully-connected, affine) layer."""
 
   def __init__(self,
                n_units,
@@ -123,7 +121,7 @@ def __init__(self,
     self._kernel_initializer = kernel_initializer
     self._bias_initializer = bias_initializer
 
-  def forward(self, x, params, state, **kwargs):
+  def forward(self, x, params=(), state=(), **kwargs):
     del kwargs
     w, b = params
     return np.dot(x, w) + b, state
@@ -148,20 +146,20 @@ def __init__(self,
     self._vocab_size = vocab_size
     self._kernel_initializer = kernel_initializer
 
-  def forward(self, x, params, state, **kwargs):
+  def forward(self, x, params=(), state=(), **kwargs):
     del kwargs
     return np.take(params, x, axis=0), state
 
   def new_params_and_state(self, input_shape, input_dtype, rng):
-    del input_dtype
-    return self._kernel_initializer((self._vocab_size, self._d_feature),
-                                    rng), ()
+    del input_shape, input_dtype
+    out_dim = (self._vocab_size, self._d_feature)
+    params = self._kernel_initializer(out_dim, rng)
+    return params, ()
 
 
 # Flatten.
 @base.layer()
-def Flatten(x, params, n_axes_to_keep=1, **kwargs):
-  del params, kwargs
+def Flatten(x, n_axes_to_keep=1, **unused_kwargs):
   if n_axes_to_keep >= len(x.shape):
     raise ValueError("n_axes_to_keep[%d] should be less than input's rank[%d]" %
                      (n_axes_to_keep, len(x.shape)))
@@ -180,11 +178,12 @@ def __init__(self, rate=0.0, name='dropout', mode='train'):
     self._mode = mode
 
   def new_params_and_state(self, input_shape, input_dtype, rng):
-    """Initialize dropout parameters and state."""
     del input_shape, input_dtype, rng
-    return (), {self._name: np.array(self._initial_rate)}
+    params = ()
+    state = {self._name: np.array(self._initial_rate)}
+    return params, state
 
-  def forward(self, x, params, state, rng=None, **unused_kwargs):
+  def forward(self, x, params=(), state=(), rng=None, **kwargs):
     """Execute dropout."""
     del params
     rate = self._initial_rate
@@ -202,14 +201,12 @@ def forward(self, x, params, state, rng=None, **unused_kwargs):
 
 
 @base.layer()
-def Div(x, params, divisor=1.0, **kwargs):
-  del params, kwargs
+def Div(x, divisor=1.0, **unused_kwargs):
   return x / divisor
 
 
 @base.layer()
-def AddConstant(x, params, constant=0.0, **unused_kwargs):
-  del params
+def AddConstant(x, constant=0.0, **unused_kwargs):
   return x + constant
 
 
@@ -224,8 +221,7 @@ def one_hot(x, size, dtype=np.float32):  # pylint: disable=invalid-name
 
 # Mean.
 @base.layer()
-def Mean(x, params, axis=-1, keepdims=False, **kwargs):
-  del params, kwargs
+def Mean(x, axis=-1, keepdims=False, **unused_kwargs):
   return np.mean(x, axis=axis, keepdims=keepdims)
 
 
diff --git a/tensor2tensor/trax/layers/core_test.py b/tensor2tensor/trax/layers/core_test.py
index 2e7b93a87..b4df4dfba 100644
--- a/tensor2tensor/trax/layers/core_test.py
+++ b/tensor2tensor/trax/layers/core_test.py
@@ -61,7 +61,7 @@ def test_flatten_n(self):
   def test_div(self):
     layer = core.Div(divisor=2.0)
     input_np = onp.array([[1, 2, 3], [4, 5, 6]], dtype=onp.float32)
-    output_np, _ = layer(input_np)
+    output_np = layer(input_np)
     # absltest doesn't have ndarray equalities.
     expected_output_np = input_np / 2.0
     self.assertAlmostEqual(
diff --git a/tensor2tensor/trax/layers/normalization_test.py b/tensor2tensor/trax/layers/normalization_test.py
index e430d9f0d..80edbd663 100644
--- a/tensor2tensor/trax/layers/normalization_test.py
+++ b/tensor2tensor/trax/layers/normalization_test.py
@@ -45,11 +45,13 @@ def test_batch_norm(self):
     m1 = 11.5  # Mean of this random input.
     v1 = 47.9167  # Variance of this random input.
     layer = normalization.BatchNorm(axis=(0, 1, 2))
-    params, state = layer.initialize_once(input_shape, input_dtype, rng)
+    _, _ = layer.initialize_once(input_shape, input_dtype, rng)
+    state = layer.state
     onp.testing.assert_allclose(state[0], 0)
     onp.testing.assert_allclose(state[1], 1)
     self.assertEqual(state[2], 0)
-    out, state = layer(inp1, params=params, state=state)
+    out = layer(inp1)
+    state = layer.state
     onp.testing.assert_allclose(state[0], m1 * 0.001)
     onp.testing.assert_allclose(state[1], 0.999 + v1 * 0.001, rtol=1e-6)
     self.assertEqual(state[2], 1)
diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index 0d0028357..02704831c 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -239,7 +239,7 @@ def PolicySchedule(
   start_time = time.time()
   # ((log_probs, value_preds), state). We have no way to pass state to the next
   # step, but that should be fine.
-  ((log_probs, _), _) = (
+  (log_probs, _) = (
       net(np.array([observations]), params=params, state=state, rng=rng))
   logging.vlog(
       1, "Running the policy took %0.2f sec.", time.time() - start_time
diff --git a/tensor2tensor/trax/models/atari_cnn_test.py b/tensor2tensor/trax/models/atari_cnn_test.py
index 88e3eb52d..5ee4d4cea 100644
--- a/tensor2tensor/trax/models/atari_cnn_test.py
+++ b/tensor2tensor/trax/models/atari_cnn_test.py
@@ -40,7 +40,7 @@ def test_computes(self):
     _, _ = model.initialize_once((1, 1) + OBS, onp.float32, key)
     x = onp.arange(B * (T + 1) * functools.reduce(op.mul, OBS)).reshape(
         B, T + 1, *OBS)
-    y, _ = model(x)
+    y = model(x)
     self.assertEqual((B, T + 1, output_size), y.shape)
 
 
@@ -57,7 +57,7 @@ def test_computes(self):
     _, _ = model.initialize_once((1, 1, OBS), onp.float32, key)
     x = onp.arange(B * (T + 1) * OBS).reshape(
         B, T + 1, OBS)
-    y, _ = model(x)
+    y = model(x)
     self.assertEqual((B, T + 1, output_size), y.shape)
 
 
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index daa398edb..3cfb8488e 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -59,10 +59,8 @@ def forward(self, inputs, params=(), state=(), **kwargs):
     rngs = _pop_rng_and_split(kwargs, len(inputs))
     results = [self._layer(x, params=params, state=state, rng=r, **kwargs)
                for x, r in zip(inputs, rngs)]
-    result_outputs, result_states = zip(*results)
     # TODO(kitaev): think about how to merge state across copies in the map.
-    result_states = result_states[0]
-    return tuple(result_outputs), tuple(result_states)
+    return tuple(results), self._layer.state
 
   def new_params_and_state(self, input_shape, input_dtype, rng):
     first_shape = input_shape[0]
@@ -226,7 +224,8 @@ def reverse(self, output, params=(), state=(), **kwargs):
     # Note that self.sublayers aligns exactly with self.reverse_layers in
     # terms of parameter and rng usage, so no re-ordering is required.
     for layer, p, s, rng in zip(self.reverse_layers, params, state, rngs):
-      reconstructed_x, _ = layer(reconstructed_x, p, s, rng=rng, **kwargs)
+      reconstructed_x = layer(reconstructed_x, params=p, state=s, rng=rng,
+                              **kwargs)
     return reconstructed_x
 
   def reverse_and_grad(self, output, ct, params=(), state=(), **kwargs):
@@ -236,7 +235,8 @@ def reverse_and_grad(self, output, ct, params=(), state=(), **kwargs):
       rngs = backend.random.split(rng, self._n_layers)
 
     def call_compute_residual(x, params):
-      res, _ = self.compute_residual(x, params, state[0], rng=rngs[0], **kwargs)
+      res = self.compute_residual(x, params=params, state=state[0], rng=rngs[0],
+                                  **kwargs)
       return res
 
     assert len(ct) == 2
@@ -244,8 +244,9 @@ def call_compute_residual(x, params):
 
     stack_with_residual, vjpfun = jax.vjp(
         call_compute_residual, output, params[0])
-    reconstructed_x, _ = self.subtract_top(
-        stack_with_residual, params[-1], state[-1], rng=rngs[-1], **kwargs)
+    reconstructed_x = self.subtract_top(
+        stack_with_residual, params=params[-1], state=state[-1], rng=rngs[-1],
+        **kwargs)
 
     x_ct, residual_params_ct = vjpfun(ct)
     assert not jax.tree_util.tree_leaves(params[-1])
@@ -328,8 +329,8 @@ def reverse(self, output, params=(), state=(), **kwargs):
     # Note that self.sublayers aligns exactly with self.reverse_layers in
     # terms of parameter and rng usage, so no re-ordering is required.
     for layer, p, s, rng in zip(self.reverse_layers, params, state, rngs):
-      reconstructed_x, _ = layer.reverse(reconstructed_x, p, s, rng=rng,
-                                         **kwargs)
+      reconstructed_x = layer.reverse(reconstructed_x, params=p, state=s,
+                                      rng=rng, **kwargs)
     return reconstructed_x
 
   def reverse_and_grad(self, output, ct, params=(), state=(), **kwargs):
@@ -341,7 +342,8 @@ def reverse_and_grad(self, output, ct, params=(), state=(), **kwargs):
     # Forward pass through self.pre_attention, while preparing for
     # later backprop.
     def call_pre_attention(x, params):
-      res, _ = self.pre_attention(x, params, state[0], rng=rngs[0], **kwargs)
+      res = self.pre_attention(x, params=params, state=state[0], rng=rngs[0],
+                               **kwargs)
       return res
     stack, pre_attention_vjpfun = jax.vjp(call_pre_attention, output, params[0])
 
@@ -351,8 +353,8 @@ def call_pre_attention(x, params):
 
     # Backprop through self.post_attention with respect to the inputs only
     def call_post_attention(x):
-      res, _ = self.post_attention(x, params[2], state[2], rng=rngs[2],
-                                   **kwargs)
+      res = self.post_attention(x, params=params[2], state=state[2],
+                                rng=rngs[2], **kwargs)
       return res
     # Note: these are *not* the actual inputs to self.post_attention.
     # If self.post_attention is not linear, we will get incorrect gradients.
@@ -372,15 +374,15 @@ def call_post_attention(x):
     # Forward pass for self.post_attention, and backprop with respect to the
     # parameters only
     def call_post_attention2(params):
-      res, _ = self.post_attention(stack, params, state[2], rng=rngs[2],
-                                   **kwargs)
+      res = self.post_attention(stack, params=params, state=state[2],
+                                rng=rngs[2], **kwargs)
       return res
     stack, post_attention_vjpfun = jax.vjp(call_post_attention2, params[2])
     (post_attention_params_ct,) = post_attention_vjpfun(saved_ct)
 
     # Forward pass through subtracting the residual
-    reconstructed_x, _ = self.subtract_top(
-        stack, params[-1], state[-1], rng=rngs[-1], **kwargs)
+    reconstructed_x = self.subtract_top(
+        stack, params=params[-1], state=state[-1], rng=rngs[-1], **kwargs)
 
     assert not jax.tree_util.tree_leaves(params[-1])
     add_top_params_ct = params[-1]
diff --git a/tensor2tensor/trax/models/transformer_test.py b/tensor2tensor/trax/models/transformer_test.py
index 760a400cb..5b87a2026 100644
--- a/tensor2tensor/trax/models/transformer_test.py
+++ b/tensor2tensor/trax/models/transformer_test.py
@@ -67,20 +67,16 @@ def test_transformer_lm_fast_inference(self):
       model_fast = model_fn(mode='predict')
       rng = backend.random.get_prng(0)
       batch_size = 2
-      (params, state_slow) = model_slow.initialize_once(
-          input_shapes=(batch_size, 1), input_dtype=np.int32, rng=rng)
-      (_, state_fast) = model_fast.initialize_once(
-          input_shapes=(batch_size, 1), input_dtype=np.int32, rng=rng)
+      _, _ = model_slow.initialize_once((batch_size, 1), np.int32, rng)
+      _, _ = model_fast.initialize_once((batch_size, 1), np.int32, rng)
 
       max_length = 5
       buf = onp.zeros((batch_size, max_length), dtype=np.int32)
       next_sym = onp.zeros((batch_size, 1), dtype=onp.int32)
 
       for index in range(max_length):
-        (logits_slow, state_slow) = model_slow(
-            buf, params=params, state=state_slow, rng=rng)
-        (logits_fast, state_fast) = model_fast(
-            next_sym, params=params, state=state_fast, rng=rng)
+        logits_slow = model_slow(buf, rng=rng)
+        logits_fast = model_fast(next_sym, rng=rng)
         onp.testing.assert_array_almost_equal(
             logits_slow[:, index, :], logits_fast[:, 0, :])
         next_sym = onp.random.randint(vocab_size, size=(batch_size, 1))
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index 5e631a2fe..99d4741bc 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -721,7 +721,7 @@ def combined_loss(new_params,
                   state=None,
                   rng=None):
   """Computes the combined (clipped loss + value loss) given observations."""
-  (log_probab_actions_new, value_predictions_new), state = (
+  (log_probab_actions_new, value_predictions_new) = (
       policy_and_value_net_apply(
           padded_observations, params=new_params, state=state, rng=rng))
 
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
index b9f9e83d2..e0046b103 100644
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -96,7 +96,7 @@ def test_policy_and_value_net(self):
     time_steps = 10
     batch_of_observations = np.random.uniform(
         size=(batch, time_steps) + observation_shape)
-    pnv_output, _ = pnv_model(batch_of_observations)
+    pnv_output = pnv_model(batch_of_observations)
 
     # Output is a list, first is probab of actions and the next is value output.
     self.assertEqual(2, len(pnv_output))
@@ -462,9 +462,9 @@ def test_combined_loss(self):
     mask = np.ones_like(rewards)
 
     # Just test that this computes at all.
-    (new_log_probabs, value_predictions_new), _ = (
-        net(observations, param=new_params, state=state))
-    (old_log_probabs, value_predictions_old), _ = (
+    (new_log_probabs, value_predictions_new) = (
+        net(observations, params=new_params, state=state))
+    (old_log_probabs, value_predictions_old) = (
         net(observations, params=old_params, state=state))
 
     gamma = 0.99
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
index 49a5f539d..af2d1e6fc 100644
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ b/tensor2tensor/trax/rl/ppo_trainer.py
@@ -582,7 +582,7 @@ def train_epoch(self, evaluate=True):
 
       # Compute the approx KL for early stopping. Use the whole dataset - as we
       # only do inference, it should fit in the memory.
-      (log_probab_actions_new, _), self._model_state = (
+      (log_probab_actions_new, _) = (
           self._policy_and_value_net_apply(
               padded_observations,
               params=self._policy_and_value_net_params,
@@ -797,7 +797,7 @@ def _get_predictions(self, observations, state, rng=None):
     """Returns log-probs, value predictions and key back."""
     key, key1 = jax_random.split(rng, num=2)
 
-    (log_probs, value_preds), state = self._policy_and_value_net_apply(
+    (log_probs, value_preds) = self._policy_and_value_net_apply(
         observations, params=self._policy_and_value_net_params, state=state,
         rng=key1)
 
diff --git a/tensor2tensor/trax/rl/simple_test.py b/tensor2tensor/trax/rl/simple_test.py
index 6383fef5a..799f4654c 100644
--- a/tensor2tensor/trax/rl/simple_test.py
+++ b/tensor2tensor/trax/rl/simple_test.py
@@ -246,7 +246,7 @@ def _make_env(
 
     gin.bind_parameter("BoxSpaceSerializer.precision", 1)
 
-    predict_output = (np.array([[[0.0]]] * batch_size), ())
+    predict_output = (np.array([[[0.0]]] * batch_size))
     mock_model_fn = mock.MagicMock()
     mock_model_fn.return_value.side_effect = itertools.repeat(predict_output)
     mock_model_fn.return_value.initialize_once.return_value = ((), ())
diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index f989bcb14..5536158a3 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -267,7 +267,7 @@ def _reset_model(self, predict_fn, indices, history, rng):
     return history[:, -1, ...]
 
   def _step_model(self, predict_fn, actions, rng):
-    (observation, reward), self._model_state = predict_fn(
+    (observation, reward) = predict_fn(
         (self._history, actions), state=self._model_state, rng=rng)
 
     # Roll the history one timestep back and append the new observation.
@@ -400,7 +400,7 @@ def _action_repr_indices(self):
   def _predict_obs(self, predict_fn, rng):
     for (i, subrng) in enumerate(jax_random.split(rng, self._obs_repr_length)):
       symbol_index = self._steps * self._step_repr_length + i
-      log_probs, self._model_state = predict_fn(
+      log_probs = predict_fn(
           self._last_symbols, state=self._model_state, rng=subrng,
       )
       log_probs = log_probs
diff --git a/tensor2tensor/trax/rl/simulated_env_problem_test.py b/tensor2tensor/trax/rl/simulated_env_problem_test.py
index d56e3545e..214c2aae5 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem_test.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem_test.py
@@ -64,7 +64,7 @@ def mock_transition(inputs, *args, **kwargs):
       (observations, actions) = inputs
       new_observations = observations[:, -1] + actions
       rewards = np.array([[int(new_observations % 2 == 0)]])
-      return (new_observations, rewards), ()
+      return (new_observations, rewards)
 
     mock_model_fn = mock.MagicMock()
     mock_model_fn.return_value.side_effect = mock_transition
@@ -169,7 +169,7 @@ def make_prediction(symbol):
       one_hot = np.eye(vocab_size)[symbol]
       log_probs = (1 - one_hot) * -100.0  # Virtually deterministic.
       # (4 obs symbols + 1 action symbol) * 3 timesteps = 15.
-      return np.array([[log_probs]]), ()
+      return np.array([[log_probs]])
 
     mock_predict_fn = mock.MagicMock()
     mock_predict_fn.side_effect = map(make_prediction, symbols)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index c6d1e4b79..7b43e1192 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -146,9 +146,8 @@ def loss(params, batch, model_predict, state, rng, has_weights):
       [inputs, targets])
   # Call model, predictions will be the returned stack, usually consisting of
   # the prediction tensor and the targets.
-  predictions, state = model_predict(model_input, params=params, state=state,
-                                     rng=rng)
-  predictions = get_preds(predictions)
+  outputs = model_predict(model_input, params=params, state=state, rng=rng)
+  predictions = get_preds(outputs)
   predictions, targets, weights = _make_list(predictions, targets, weights)
   xent = []
   for (pred, target) in zip(predictions, targets):
@@ -306,8 +305,8 @@ def evaluate(inputs_stream, predict_fn, metric_fns, state, rng, has_weights):
     rng, subrng = jax_random.split(rng)
     model_inp, get_preds = _stack_inputs_targets_and_get_predictions(inp)
     # Call model, preds will be the returned stack, usually (pred, targets).
-    preds, state = predict_fn(model_inp, state=state, rng=subrng)
-    pred = get_preds(preds)
+    outputs = predict_fn(model_inp, state=state, rng=subrng)
+    pred = get_preds(outputs)
     for m, f in six.iteritems(metric_fns):
       metrics[m] += f(inp, pred, has_weights=has_weights)
   return {m: v / count for (m, v) in six.iteritems(metrics)}, state
@@ -425,7 +424,7 @@ def mapped_predict(x, params, state, rng):
 
   def predict(x, params=(), state=(), rng=None):
     """Predict function jited and parallelized as requested."""
-    pred, state = mapped_predict(
+    pred = mapped_predict(
         reshape_by_device(x, n_devices),
         params,
         state,
@@ -435,7 +434,7 @@ def predict(x, params=(), state=(), rng=None):
     def combine(x):
       batch_size = x.shape[0] * x.shape[1]
       return np.reshape(x, [batch_size] + list(x.shape[2:]))
-    return layers.nested_map(pred, combine), state
+    return layers.nested_map(pred, combine)
 
   return predict
 
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index 13a153208..e892f89a7 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -95,13 +95,12 @@ def test_train_eval_predict(self, backend_name):
       train_steps = 2
       eval_steps = 2
       # Adds Dropout and BatchNorm to test state handling.
-      mlp = functools.partial(models.MLP,
-                              d_hidden=16,
-                              n_output_classes=n_classes)
       def model_fn(mode="train"):
         return layers.Model(layers.Dropout(mode=mode, rate=0.1),
                             layers.BatchNorm(mode=mode),
-                            mlp(mode=mode))
+                            models.MLP(d_hidden=16,
+                                       n_output_classes=n_classes,
+                                       mode=mode))
 
       inputs = lambda _: test_inputs(n_classes)
 

From 694b00a4b7d23d2c3bda991526334b07124d3193 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Sun, 29 Sep 2019 07:44:12 -0700
Subject: [PATCH 2499/2720] Add test for Transformer-Revnet.

PiperOrigin-RevId: 271844734
---
 .../research/transformer_revnet_test.py       | 44 +++++++++++++++++++
 1 file changed, 44 insertions(+)
 create mode 100644 tensor2tensor/trax/models/research/transformer_revnet_test.py

diff --git a/tensor2tensor/trax/models/research/transformer_revnet_test.py b/tensor2tensor/trax/models/research/transformer_revnet_test.py
new file mode 100644
index 000000000..5642d54bd
--- /dev/null
+++ b/tensor2tensor/trax/models/research/transformer_revnet_test.py
@@ -0,0 +1,44 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 Transformer-Revnet models."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from tensor2tensor.trax import layers as tl
+from tensor2tensor.trax.models.research import transformer_revnet
+
+
+class TransformerRevnetTest(parameterized.TestCase):
+
+  def test_transformer_lm_forward_shape(self):
+    """Run the TransformerRevnet LM forward and check output shape."""
+    vocab_size = 16
+    input_shape = ((1, 8), (1, 8))
+    model = transformer_revnet.TransformerRevnetLM(
+        vocab_size, d_model=32, d_ff=64,
+        d_attention_key=16, d_attention_value=16, n_layers=1, n_heads=2,
+        max_len=16, n_chunks=2, n_attention_chunks=1)
+    final_shape = tl.check_shape_agreement(
+        model, tuple(input_shape), integer_inputs=True)
+    self.assertEqual(((1, 8, 16), (1, 8, 16)), final_shape)
+
+
+if __name__ == '__main__':
+  absltest.main()

From 2f35350c02afc805f9d04aac0c36302df6f8fb3b Mon Sep 17 00:00:00 2001
From: RJ Skerry-Ryan <rjryan@google.com>
Date: Sun, 29 Sep 2019 14:38:14 -0700
Subject: [PATCH 2500/2720] Rewrite references to tf.contrib.signal to
 tf.signal.

PiperOrigin-RevId: 271871580
---
 tensor2tensor/layers/common_audio.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/common_audio.py b/tensor2tensor/layers/common_audio.py
index 8c388728e..545e831b7 100644
--- a/tensor2tensor/layers/common_audio.py
+++ b/tensor2tensor/layers/common_audio.py
@@ -56,7 +56,7 @@ def compute_mel_filterbank_features(
     waveforms,
     sample_rate=16000, dither=1.0 / np.iinfo(np.int16).max, preemphasis=0.97,
     frame_length=25, frame_step=10, fft_length=None,
-    window_fn=functools.partial(tf.contrib.signal.hann_window, periodic=True),
+    window_fn=functools.partial(tf.signal.hann_window, periodic=True),
     lower_edge_hertz=80.0, upper_edge_hertz=7600.0, num_mel_bins=80,
     log_noise_floor=1e-3, apply_mask=True):
   """Implement mel-filterbank extraction using tf ops.
@@ -101,7 +101,7 @@ def compute_mel_filterbank_features(
   if fft_length is None:
     fft_length = int(2**(np.ceil(np.log2(frame_length))))
 
-  stfts = tf.contrib.signal.stft(
+  stfts = tf.signal.stft(
       waveforms,
       frame_length=frame_length,
       frame_step=frame_step,
@@ -121,7 +121,7 @@ def compute_mel_filterbank_features(
   # Warp the linear-scale, magnitude spectrograms into the mel-scale.
   num_spectrogram_bins = magnitude_spectrograms.shape[-1].value
   linear_to_mel_weight_matrix = (
-      tf.contrib.signal.linear_to_mel_weight_matrix(
+      tf.signal.linear_to_mel_weight_matrix(
           num_mel_bins, num_spectrogram_bins, sample_rate, lower_edge_hertz,
           upper_edge_hertz))
   mel_spectrograms = tf.tensordot(

From 019b6be6cf43ab0326c495284606724cd7527309 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sun, 29 Sep 2019 23:42:55 -0700
Subject: [PATCH 2501/2720] Internal code clean-ups, around initialization and
 parameters.

PiperOrigin-RevId: 271916998
---
 tensor2tensor/trax/layers/README.md           | 14 ++--
 tensor2tensor/trax/layers/base.py             | 74 +++++++++++++------
 tensor2tensor/trax/layers/base_test.py        | 12 +--
 tensor2tensor/trax/layers/combinators.py      | 10 ++-
 tensor2tensor/trax/layers/core.py             |  2 +-
 tensor2tensor/trax/layers/demo.ipynb          |  2 +-
 tensor2tensor/trax/layers/normalization.py    | 11 ++-
 tensor2tensor/trax/layers/reversible.py       |  2 +-
 .../models/research/transformer_revnet.py     | 18 +++--
 tensor2tensor/trax/rl/ppo_trainer_test.py     |  2 +-
 10 files changed, 92 insertions(+), 55 deletions(-)

diff --git a/tensor2tensor/trax/layers/README.md b/tensor2tensor/trax/layers/README.md
index a9e4deabe..4782fa2dc 100644
--- a/tensor2tensor/trax/layers/README.md
+++ b/tensor2tensor/trax/layers/README.md
@@ -4,17 +4,15 @@
 
 ## Base layer structure
 
-All layers inherit form the Layer class and need to implement 3 functions:
+All layers inherit from the Layer class and generally need to implement 2
+methods:
 
 ```python
-def forward(self, params, inputs, **kwargs):
-"""Call this layer using the given parameters on the given inputs."""
+def forward(self, inputs, params=(), state=(), **kwargs):
+  """Computes the layer's output as part of a forward pass through the model."""
 
-def output_shape(self, input_shape):
-"""The shape of the output given the shape of the input."""
-
-def new_params_and_state(self, input_shape, rng):
-"""Create new parameters given the shape of the input."""
+def new_params_and_state(self, input_shape, input_dtype, rng):
+  """Returns a (params, state) pair suitable for initializing this layer."""
 ```
 
 The base Layer class wraps these functions and provides initialization
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 0136f1afd..09ea3b16a 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -67,6 +67,12 @@ class Layer(object):
   """
 
   def __init__(self, n_inputs=1, n_outputs=1):
+    """Creates a partially initialized, unconnected layer instance.
+
+    Args:
+      n_inputs: Number of inputs expected by this layer.
+      n_outputs: Number of outputs promised by this layer.
+    """
     self._n_inputs = n_inputs
     self._n_outputs = n_outputs
     self._sublayers = ()  # Default is no sublayers.
@@ -86,7 +92,7 @@ def __repr__(self):
       return '{}{{{}}}'.format(class_str, fields_str)
 
   def forward(self, inputs, params=(), state=(), **kwargs):
-    """Uses this layer as part of a forward pass through the model.
+    """Computes this layer's output as part of a forward pass through the model.
 
     Authors of new Layer subclasses should override this method to define the
     forward computation that their layer performs.
@@ -167,14 +173,18 @@ def state(self, state):
 
   @property
   def has_backward(self):
-    """Whether to use custom gradients (in which case, see below)."""
+    """Returns True if this layer provides its own (custom) backward pass code.
+
+    A layer subclass that provides custom backward pass code (for custom
+    gradients) must override this method to return True.
+    """
     return False
 
   def backward(self, inputs, output, grad, params, state, **kwargs):
     """Custom backward pass to propagate gradients in a custom way.
 
     Args:
-      inputs: Input activations; can be a (possibly nested) tuple.
+      inputs: Input tensors; can be a (possibly nested) tuple.
       output: The result of running this layer on inputs.
       grad: gradient signal (called cotangent in jax) computed based on
         subsequent layers. The structure and shape must match output.
@@ -202,9 +212,11 @@ def pseudo_forward(self, pseudo_inputs, params, state):
       state: start state.
 
     Returns:
-      A ShapeType instance representing the shape and type of the output (if
-      this layer has one output) or a tuple of ShapeType instances (if this
-      layer has more than one output).
+      A tuple of (output, state).
+
+      The output part of the tuple is a ShapeType instance representing the
+      shape and type of the output (if this layer has one output) or a tuple
+      of ShapeType instances (if this layer has more than one output).
     """
     try:
       # Beware: using an actual RNG (as opposed to this ShapeType stub) would
@@ -224,21 +236,22 @@ def call_on_input(x, params, state, rng):
                        trace)
 
   def initialize_once(self, input_shapes, input_dtype, rng):
-    """Initialize the layer given an input shape, dtype and rng.
+    """Initializes this layer and its sublayers recursively.
 
-    Returns new_params_and_state(input_shapes, rng) on the first call and () on
-    any subsequent call, as the layer is already initialized. This is used for
-    networks that share parameters, so the layer only produces them once.
+    This method is designed to initialize each layer instance once, even if the
+    same layer instance occurs in multiple places in the network. This enables
+    weight sharing to be implemented as layer sharing.
 
     Args:
       input_shapes: A tuple representing a shape (if this layer takes one input)
           or a tuple of shapes (if this layer takes more than one input).
           For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
-      input_dtype: numpy dtype of the input.
-      rng: A random number generator.
+      input_dtype: Numpy dtype(s) for each of the inputs.
+      rng: A PRNG key for random number generation.
 
     Returns:
-      Newly created parameters on the first call and () on all subsequent calls.
+      A (params, state) tuple, in which params contains newly created parameters
+          on the first call and () on all subsequent calls.
     """
     try:
       # Initialize params once; store them for use when this layer is called.
@@ -291,6 +304,25 @@ def __call__(self, x, **kwargs):
     """
     params = kwargs.pop('params', self.params)
     state = kwargs.pop('state', self.state)
+    outputs, _ = self.apply_forward(x, params=params, state=state, **kwargs)
+    return outputs
+
+  def apply_forward(self, x, params=(), state=(), **kwargs):
+    """Applies this layer as part of a forward pass; an internal system method.
+
+    This method is reserved for handling plumbing and other internal affairs
+    as needed by the overall library. Trax library users should use or override
+    the `forward` method instead.
+
+    Args:
+      x: See Layer.forward inputs.
+      params: See Layer.forward.
+      state: See Layer.forward.
+      **kwargs: See Layer.forward.
+
+    Returns:
+      See Layer.forward.
+    """
     try:
       # If params are nothing, we may be reusing this layer.
       # Use the cached parameters to calculate the value.
@@ -307,10 +339,11 @@ def __call__(self, x, **kwargs):
       else:
         outputs, s = self._do_custom_gradients(x, params, state, **kwargs)
       self._state = s
-      return outputs
+      return outputs, s
+
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback()
-      raise LayerError(name, 'forward', self._caller, shapes(x), trace)
+      raise LayerError(name, 'apply_forward', self._caller, shapes(x), trace)
 
   def _do_custom_gradients(self, x, params, state, **kwargs):
     """Calls this layer for a forward pass, but with custom gradients."""
@@ -373,7 +406,7 @@ class LayerError(Exception):
   def __init__(self, layer_name, function_name, caller,
                input_shapes, traceback_string):
     self._layer_name = layer_name
-    self._function_name = function_name  # Is it forward or initialize_once?
+    self._function_name = function_name
     self._caller = caller  # Python inspect object with init caller info.
     self._traceback = traceback_string
     self._input_shapes = input_shapes
@@ -495,10 +528,10 @@ def _validate_forward_input(x, n_inputs):
 
 
 def layer(n_inputs=1, n_outputs=1, new_params_and_state_fn=None):
-  """Decorates a function to make it the call method of a new Layer class."""
+  """Returns a decorator that converts a function into a Layer class builder."""
 
   def _build_layer_class(raw_fn):
-    """Returns a Layer class built around the given forward function."""
+    """Returns a Layer class whose callable instances execute the function."""
 
     def _init(self, **kwargs):
       self._kwargs = kwargs  # pylint: disable=protected-access
@@ -508,15 +541,14 @@ def _new_params_and_state(self, input_shapes, input_dtype, rng):
       if new_params_and_state_fn is None:
         return (), ()
       kwargs = self._kwargs  # pylint: disable=protected-access
-      return (
-          new_params_and_state_fn(input_shapes, input_dtype, rng, **kwargs), ())
+      return new_params_and_state_fn(input_shapes, input_dtype, rng, **kwargs)
 
     def _is_empty(raw_output):
       return raw_output is None or (isinstance(raw_output, (list, tuple))
                                     and len(raw_output) == 0)  # pylint: disable=g-explicit-length-test
 
     def _forward(self, x, params=(), state=(), **kwargs):
-      """Calls raw_fn with extra keyword args from Layer.__init__."""
+      """Uses this layer as part of a forward pass through the model."""
       merged_kwargs = kwargs.copy()
       merged_kwargs.update(self._kwargs)  # pylint: disable=protected-access
 
diff --git a/tensor2tensor/trax/layers/base_test.py b/tensor2tensor/trax/layers/base_test.py
index a2ab373d4..66aefee9d 100644
--- a/tensor2tensor/trax/layers/base_test.py
+++ b/tensor2tensor/trax/layers/base_test.py
@@ -51,11 +51,11 @@ def backward(self, inputs, output, ct, params, state, **kwargs):
 
     layer = IdWithZeroGrad()
     rng = backend.random.get_prng(0)
-    params = ()
     input_shape = (9, 17)
     random_input = backend.random.uniform(rng, input_shape, minval=-1.0,
                                           maxval=1.0)
-    f = lambda x: backend.numpy.mean(layer(x, params=params, rng=rng)[0])
+    layer.initialize_once(input_shape, random_input.dtype, rng)
+    f = lambda x: backend.numpy.mean(layer(x))
     grad = backend.grad(f)(random_input)
     self.assertEqual(grad.shape, input_shape)  # Gradient for each input.
     self.assertEqual(sum(sum(grad * grad)), 0.0)  # Each one is 0.
@@ -68,10 +68,6 @@ def forward(self, x, params=(), state=(), **kwargs):
         del kwargs
         return x, ()
 
-      def new_params_and_state(self, input_shapes, input_dtype, rng):
-        del input_shapes, input_dtype, rng
-        return (), ()
-
       @property
       def has_backward(self):
         return True
@@ -81,11 +77,11 @@ def backward(self, inputs, output, ct, params, state, **kwargs):
 
     layer = IdWithIdGrad()
     rng = backend.random.get_prng(0)
-    params = ()
     input_shape = (9, 17)
     random_input = backend.random.uniform(rng, input_shape, minval=-1.0,
                                           maxval=1.0)
-    f = lambda x: backend.numpy.mean(layer(x, params=params, rng=rng)[0])
+    layer.initialize_once(input_shape, random_input.dtype, rng)
+    f = lambda x: backend.numpy.mean(layer(x))
     grad = backend.grad(f)(random_input)
     self.assertEqual(grad.shape, input_shape)  # Gradient for each input.
     self.assertEqual(sum(sum(grad)), sum(sum(random_input)))  # Same as input.
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 8d4b8f950..f33b061f3 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -201,8 +201,9 @@ def forward(self, xs, params=(), state=(), **kwargs):
         inputs = stack[0]
       else:
         inputs = stack[:n_in]
-      outputs = layer(inputs, params=p, state=s, rng=rng, **kwargs)
-      new_state.append(layer.state)
+      outputs, s = layer.apply_forward(inputs, params=p, state=s, rng=rng,
+                                       **kwargs)
+      new_state.append(s)
 
       # Push outputs onto remaining stack (if any).
       if n_in < _count_items(stack):
@@ -467,12 +468,13 @@ def forward(self, inputs, params=(), state=(), **kwargs):
     new_state = []
     for layer, x, p, s, r in zip(layers, sublayer_inputs, params, state, rngs):
       # Note that zip silently truncates its result if lengths don't match.
-      sub_outputs = layer(x, params=p, state=s, rng=r, **kwargs)
+      sub_outputs, sub_state = layer.apply_forward(x, params=p, state=s, rng=r,
+                                                   **kwargs)
       if layer.n_outputs == 1:
         outputs.append(sub_outputs)
       else:
         outputs.extend(sub_outputs)
-      new_state.append(layer.state)
+      new_state.append(sub_state)
     output = outputs[0] if self.n_outputs == 1 else tuple(outputs)
     return output, new_state
 
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index a36c74a84..1cb8fc0b4 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -185,7 +185,7 @@ def new_params_and_state(self, input_shape, input_dtype, rng):
 
   def forward(self, x, params=(), state=(), rng=None, **kwargs):
     """Execute dropout."""
-    del params
+    del kwargs
     rate = self._initial_rate
     if isinstance(state, dict) and self._name in state:
       rate = state[self._name]
diff --git a/tensor2tensor/trax/layers/demo.ipynb b/tensor2tensor/trax/layers/demo.ipynb
index ecb186560..8317b683b 100644
--- a/tensor2tensor/trax/layers/demo.ipynb
+++ b/tensor2tensor/trax/layers/demo.ipynb
@@ -156,7 +156,7 @@
         "x  = onp.arange(-7, 8).reshape(3, -1)\n",
         "rng = backend.random.get_prng(0)\n",
         "layer = tl.Relu()\n",
-        "params, _ = layer.initialize_once(x.shape, x.dtype, rng)\n",
+        "layer.initialize_once(x.shape, x.dtype, rng)\n",
         "output = layer(x)\n",
         "print(x)\n",
         "print(output)"
diff --git a/tensor2tensor/trax/layers/normalization.py b/tensor2tensor/trax/layers/normalization.py
index 5c11a28d6..fd7d4908c 100644
--- a/tensor2tensor/trax/layers/normalization.py
+++ b/tensor2tensor/trax/layers/normalization.py
@@ -53,7 +53,9 @@ def get_stats_axis(i, d):
     running_mean = np.zeros(stats_shape, dtype=np.float32)
     running_var = np.ones(stats_shape, dtype=np.float32)
     num_batches = np.zeros((), dtype=np.int32)
-    return (beta, gamma), (running_mean, running_var, num_batches)
+    params = (beta, gamma)
+    state = (running_mean, running_var, num_batches)
+    return params, state
 
   def forward(self, x, params, state, **unused_kwargs):
     """Layer construction function for a batch normalization layer."""
@@ -104,16 +106,17 @@ def average(factor, new, old):
 
 
 # Layer normalization.
-def _layer_norm_params(input_shape, input_dtype, rng):
+def _layer_norm_params_and_state(input_shape, input_dtype, rng):
   """Helper: create layer norm parameters."""
   del input_dtype, rng
   features = input_shape[-1]
   scale = np.ones(features)
   bias = np.zeros(features)
-  return (scale, bias)
+  params = (scale, bias)
+  return params, ()
 
 
-@base.layer(new_params_and_state_fn=_layer_norm_params)
+@base.layer(new_params_and_state_fn=_layer_norm_params_and_state)
 def LayerNorm(x, params, epsilon=1e-6, **unused_kwargs):  # pylint: disable=invalid-name
   (scale, bias) = params
   mean = np.mean(x, axis=-1, keepdims=True)
diff --git a/tensor2tensor/trax/layers/reversible.py b/tensor2tensor/trax/layers/reversible.py
index ff56072c9..5b52aeaee 100644
--- a/tensor2tensor/trax/layers/reversible.py
+++ b/tensor2tensor/trax/layers/reversible.py
@@ -78,7 +78,7 @@ class ReversibleSwap(ReversibleLayer, cb.Swap):
 
   def reverse(self, output, params=(), state=(), **kwargs):
     # Swap is its own inverse, except that reverse doesn't return the state.
-    return self.forward(output, params, state, **kwargs)[0]
+    return self.forward(output, params=params, state=state, **kwargs)[0]
 
 
 class ReversibleSerial(ReversibleLayer, cb.Serial):
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 3cfb8488e..430635125 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -71,6 +71,18 @@ def new_params_and_state(self, input_shape, input_dtype, rng):
                            'with the same shapes. Shapes: %s' % str(shape))
     return self._layer.initialize_once(first_shape, input_dtype[0], rng)
 
+  @tl.Layer.params.setter
+  def params(self, params):
+    self._params = params
+    assert len(params) == 1
+    self._layer.params = params[0]
+
+  @tl.Layer.state.setter
+  def state(self, state):
+    self._state = state
+    assert len(state) == 1
+    self._layer.state = state[0]
+
 
 @tl.layer()
 def BroadcastedDropout(x, params, rate=0.0, mode='train', broadcast_dims=(-2,),
@@ -118,9 +130,6 @@ def forward(self, inputs, params=(), state=(), **kwargs):
     res = tuple(backend.numpy.split(inputs, self._n_sections, self._axis))
     return res, state
 
-  def new_params_and_state(self, input_shapes, input_dtype, rng):
-    return (), ()
-
 
 class SplitForOutput(tl.ReversibleLayer):
   """Splits activations into sections (for use right before the output layer).
@@ -142,9 +151,6 @@ def __init__(self, n_sections=2, axis=-2):
     self._n_sections = n_sections
     self._axis = axis
 
-  def new_params_and_state(self, input_shape, input_dtype, rng):
-    return (), ()
-
   def forward(self, inputs, params=(), state=(), **kwargs):
     del params, kwargs
     x1, x2 = inputs
diff --git a/tensor2tensor/trax/rl/ppo_trainer_test.py b/tensor2tensor/trax/rl/ppo_trainer_test.py
index e67df34b1..38b7c0bfc 100644
--- a/tensor2tensor/trax/rl/ppo_trainer_test.py
+++ b/tensor2tensor/trax/rl/ppo_trainer_test.py
@@ -75,7 +75,7 @@ def _make_trainer(
       self, train_env, eval_env, output_dir, model=None, **kwargs
   ):
     if model is None:
-      model = lambda: [layers.Dense(1)]
+      model = lambda: layers.Serial(layers.Dense(1))
     return ppo_trainer.PPO(
         train_env=train_env,
         eval_env=eval_env,

From b8c4998141770c71d50c4ee3e956b0e92dc431eb Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 30 Sep 2019 11:36:13 -0700
Subject: [PATCH 2502/2720] Rename batch_fun to batch_fn, to align with other
 ..._fn names.

PiperOrigin-RevId: 272028547
---
 tensor2tensor/trax/configs/mlp_mnist.gin       |  6 +++---
 .../position_lookup_transformer_copy.gin       |  6 +++---
 tensor2tensor/trax/configs/reformer_enwik8.gin |  8 ++++----
 .../trax/configs/reformer_enwik8_rev.gin       |  8 ++++----
 .../trax/configs/reformer_imagenet64.gin       |  8 ++++----
 .../trax/configs/reformer_imagenet64_rev.gin   |  8 ++++----
 .../trax/configs/resnet50_imagenet_8gb.gin     | 10 +++++-----
 .../configs/resnet50_imagenet_8gb_testing.gin  | 10 +++++-----
 .../trax/configs/transformer_big_lm1b_8gb.gin  |  8 ++++----
 .../trax/configs/transformer_copy.gin          |  8 ++++----
 .../trax/configs/transformer_enwik8.gin        |  8 ++++----
 .../trax/configs/transformer_imagenet64.gin    |  8 ++++----
 .../trax/configs/transformer_imdb_8gb.gin      |  8 ++++----
 .../trax/configs/transformer_lm1b_8gb.gin      |  8 ++++----
 .../configs/transformer_lm1b_8gb_testing.gin   |  8 ++++----
 .../configs/transformer_lm_wmt_ende_8gb.gin    | 12 ++++++------
 .../trax/configs/transformer_revnet_enwik8.gin |  8 ++++----
 .../transformer_revnet_imagenet64_8gb.gin      |  8 ++++----
 ...sformer_wmt_ende_16gb_adafactor_testing.gin | 10 +++++-----
 .../transformer_wmt_ende_8gb_adafactor.gin     | 10 +++++-----
 .../configs/transformer_wmt_ende_8gb_adam.gin  | 10 +++++-----
 .../configs/transformer_wmt_ende_8gb_sm3.gin   | 10 +++++-----
 .../trax/configs/wide_resnet_cifar10_8gb.gin   | 10 +++++-----
 tensor2tensor/trax/inputs.py                   | 18 +++++++++---------
 tensor2tensor/trax/inputs_test.py              |  8 ++++----
 ...online_tune_transformer_imagenet64_16gb.gin |  8 ++++----
 .../env_online_tune_transformer_lm1b_16gb.gin  |  8 ++++----
 ...nline_tune_transformer_lm_wmt_ende_16gb.gin | 12 ++++++------
 .../env_online_tune_wide_resnet_cifar10.gin    | 10 +++++-----
 .../ppo_online_tune_wide_resnet_cifar10.gin    | 10 +++++-----
 30 files changed, 136 insertions(+), 136 deletions(-)

diff --git a/tensor2tensor/trax/configs/mlp_mnist.gin b/tensor2tensor/trax/configs/mlp_mnist.gin
index a8bee6abd..75cbf7ba9 100644
--- a/tensor2tensor/trax/configs/mlp_mnist.gin
+++ b/tensor2tensor/trax/configs/mlp_mnist.gin
@@ -4,10 +4,10 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 256
-batch_fun.eval_batch_size = 256
+batch_fn.batch_size_per_device = 256
+batch_fn.eval_batch_size = 256
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin b/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
index f6aef3874..af0bedd66 100644
--- a/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
+++ b/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
@@ -3,10 +3,10 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 32
-batch_fun.eval_batch_size = 32
+batch_fn.batch_size_per_device = 32
+batch_fn.eval_batch_size = 32
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/reformer_enwik8.gin b/tensor2tensor/trax/configs/reformer_enwik8.gin
index 63b032c17..aa514ab62 100644
--- a/tensor2tensor/trax/configs/reformer_enwik8.gin
+++ b/tensor2tensor/trax/configs/reformer_enwik8.gin
@@ -46,11 +46,11 @@ MergedMultiHashedCausalAttentionV2.n_buckets = 1024  # Always 2 * n_bins
 MergedMultiHashedCausalAttentionV2.n_hashes = 2
 MergedMultiHashedCausalAttentionV2.drop_for_hash_rate = 0.0
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 1
-batch_fun.eval_batch_size = 8
-batch_fun.max_eval_length = 65536
+batch_fn.batch_size_per_device = 1
+batch_fn.eval_batch_size = 8
+batch_fn.max_eval_length = 65536
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/reformer_enwik8_rev.gin b/tensor2tensor/trax/configs/reformer_enwik8_rev.gin
index f762b8835..e82007c0c 100644
--- a/tensor2tensor/trax/configs/reformer_enwik8_rev.gin
+++ b/tensor2tensor/trax/configs/reformer_enwik8_rev.gin
@@ -46,11 +46,11 @@ MergedMultiHashedCausalAttentionV2.n_buckets = 1024  # Always 2 * n_bins
 MergedMultiHashedCausalAttentionV2.n_hashes = 2
 MergedMultiHashedCausalAttentionV2.drop_for_hash_rate = 0.0
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 1
-batch_fun.eval_batch_size = 8
-batch_fun.max_eval_length = 65536
+batch_fn.batch_size_per_device = 1
+batch_fn.eval_batch_size = 8
+batch_fn.max_eval_length = 65536
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/reformer_imagenet64.gin b/tensor2tensor/trax/configs/reformer_imagenet64.gin
index ac78c3d56..e1e9bf6be 100644
--- a/tensor2tensor/trax/configs/reformer_imagenet64.gin
+++ b/tensor2tensor/trax/configs/reformer_imagenet64.gin
@@ -43,11 +43,11 @@ MergedMultiHashedCausalAttentionV2.n_buckets = 192  # Always 2 * n_bins
 MergedMultiHashedCausalAttentionV2.n_hashes = 2
 MergedMultiHashedCausalAttentionV2.drop_for_hash_rate = 0.0
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 1
-batch_fun.eval_batch_size = 8
-batch_fun.max_eval_length = 12288  # 64 * 64 * 3
+batch_fn.batch_size_per_device = 1
+batch_fn.eval_batch_size = 8
+batch_fn.max_eval_length = 12288  # 64 * 64 * 3
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/reformer_imagenet64_rev.gin b/tensor2tensor/trax/configs/reformer_imagenet64_rev.gin
index 6d15419cb..673df196d 100644
--- a/tensor2tensor/trax/configs/reformer_imagenet64_rev.gin
+++ b/tensor2tensor/trax/configs/reformer_imagenet64_rev.gin
@@ -43,11 +43,11 @@ MergedMultiHashedCausalAttentionV2.n_buckets = 192  # Always 2 * n_bins
 MergedMultiHashedCausalAttentionV2.n_hashes = 2
 MergedMultiHashedCausalAttentionV2.drop_for_hash_rate = 0.0
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 1
-batch_fun.eval_batch_size = 8
-batch_fun.max_eval_length = 12288  # 64 * 64 * 3
+batch_fn.batch_size_per_device = 1
+batch_fn.eval_batch_size = 8
+batch_fn.max_eval_length = 12288  # 64 * 64 * 3
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
index 86415eb1a..7eb13a225 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
@@ -4,12 +4,12 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 32
-batch_fun.bucket_length = 32
-batch_fun.buckets = None
-batch_fun.eval_batch_size = 32
+batch_fn.batch_size_per_device = 32
+batch_fn.bucket_length = 32
+batch_fn.buckets = None
+batch_fn.eval_batch_size = 32
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
index d4085acea..6a49f89fd 100644
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
+++ b/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
@@ -4,12 +4,12 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 32
-batch_fun.bucket_length = 32
-batch_fun.buckets = None
-batch_fun.eval_batch_size = 32
+batch_fn.batch_size_per_device = 32
+batch_fn.bucket_length = 32
+batch_fn.buckets = None
+batch_fn.eval_batch_size = 32
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
index 46e28ad13..c3d110531 100644
--- a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
@@ -3,11 +3,11 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 32
-batch_fun.eval_batch_size = 64
-batch_fun.max_eval_length = 512
+batch_fn.batch_size_per_device = 32
+batch_fn.eval_batch_size = 64
+batch_fn.max_eval_length = 512
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_copy.gin b/tensor2tensor/trax/configs/transformer_copy.gin
index a7982405f..cc7d0f70d 100644
--- a/tensor2tensor/trax/configs/transformer_copy.gin
+++ b/tensor2tensor/trax/configs/transformer_copy.gin
@@ -7,11 +7,11 @@ n_symbols = 128
 length = 1024
 batch = 32
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = %batch
-batch_fun.eval_batch_size = %batch
-batch_fun.max_eval_length = %length
+batch_fn.batch_size_per_device = %batch
+batch_fn.eval_batch_size = %batch
+batch_fn.max_eval_length = %length
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_enwik8.gin b/tensor2tensor/trax/configs/transformer_enwik8.gin
index ce1727082..364e41ae8 100644
--- a/tensor2tensor/trax/configs/transformer_enwik8.gin
+++ b/tensor2tensor/trax/configs/transformer_enwik8.gin
@@ -3,11 +3,11 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 1
-batch_fun.eval_batch_size = 8
-batch_fun.max_eval_length = 65536
+batch_fn.batch_size_per_device = 1
+batch_fn.eval_batch_size = 8
+batch_fn.max_eval_length = 65536
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_imagenet64.gin b/tensor2tensor/trax/configs/transformer_imagenet64.gin
index 63fd2f380..5c9735bf3 100644
--- a/tensor2tensor/trax/configs/transformer_imagenet64.gin
+++ b/tensor2tensor/trax/configs/transformer_imagenet64.gin
@@ -3,11 +3,11 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 1
-batch_fun.eval_batch_size = 8
-batch_fun.max_eval_length = 12288  # 64 * 64 * 3
+batch_fn.batch_size_per_device = 1
+batch_fn.eval_batch_size = 8
+batch_fn.max_eval_length = 12288  # 64 * 64 * 3
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
index ba44965bb..b4cd87f81 100644
--- a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
@@ -3,11 +3,11 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 128
-batch_fun.eval_batch_size = 128
-batch_fun.max_eval_length = 2048
+batch_fn.batch_size_per_device = 128
+batch_fn.eval_batch_size = 128
+batch_fn.max_eval_length = 2048
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index 33c32a465..d44cf8d12 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -3,11 +3,11 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 128
-batch_fun.eval_batch_size = 128
-batch_fun.max_eval_length = 2048
+batch_fn.batch_size_per_device = 128
+batch_fn.eval_batch_size = 128
+batch_fn.max_eval_length = 2048
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
index 2f8f91ec6..e0bd85751 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
@@ -3,11 +3,11 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 128
-batch_fun.eval_batch_size = 128
-batch_fun.max_eval_length = 2048
+batch_fn.batch_size_per_device = 128
+batch_fn.eval_batch_size = 128
+batch_fn.max_eval_length = 2048
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_lm_wmt_ende_8gb.gin b/tensor2tensor/trax/configs/transformer_lm_wmt_ende_8gb.gin
index 7fa1f2504..ed43f7919 100644
--- a/tensor2tensor/trax/configs/transformer_lm_wmt_ende_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm_wmt_ende_8gb.gin
@@ -3,13 +3,13 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 64
-batch_fun.eval_batch_size = 64
-batch_fun.bucket_length=64
-batch_fun.max_eval_length = 1024
-batch_fun.buckets_include_inputs_in_length=True
+batch_fn.batch_size_per_device = 64
+batch_fn.eval_batch_size = 64
+batch_fn.bucket_length=64
+batch_fn.max_eval_length = 1024
+batch_fn.buckets_include_inputs_in_length=True
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_revnet_enwik8.gin b/tensor2tensor/trax/configs/transformer_revnet_enwik8.gin
index 7ab04d905..4e863d73d 100644
--- a/tensor2tensor/trax/configs/transformer_revnet_enwik8.gin
+++ b/tensor2tensor/trax/configs/transformer_revnet_enwik8.gin
@@ -3,11 +3,11 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 1
-batch_fun.eval_batch_size = 8
-batch_fun.max_eval_length = 65536
+batch_fn.batch_size_per_device = 1
+batch_fn.eval_batch_size = 8
+batch_fn.max_eval_length = 65536
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
index 5d5ece224..c5fa8e880 100644
--- a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
@@ -3,11 +3,11 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 2
-batch_fun.eval_batch_size = 16
-batch_fun.max_eval_length = 12288  # 64 * 64 * 3
+batch_fn.batch_size_per_device = 2
+batch_fn.eval_batch_size = 16
+batch_fn.max_eval_length = 12288  # 64 * 64 * 3
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin
index 5ee37f868..198e25e64 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin
@@ -3,12 +3,12 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 128
-batch_fun.eval_batch_size = 64
-batch_fun.max_eval_length = 1024
-batch_fun.buckets_include_inputs_in_length=True
+batch_fn.batch_size_per_device = 128
+batch_fn.eval_batch_size = 64
+batch_fn.max_eval_length = 1024
+batch_fn.buckets_include_inputs_in_length=True
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
index d99183860..06e7d9072 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
@@ -3,12 +3,12 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 64
-batch_fun.eval_batch_size = 64
-batch_fun.max_eval_length = 1024
-batch_fun.buckets_include_inputs_in_length=True
+batch_fn.batch_size_per_device = 64
+batch_fn.eval_batch_size = 64
+batch_fn.max_eval_length = 1024
+batch_fn.buckets_include_inputs_in_length=True
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
index f9bda05e2..e4a010b33 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
@@ -3,12 +3,12 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 64
-batch_fun.eval_batch_size = 64
-batch_fun.max_eval_length = 1024
-batch_fun.buckets_include_inputs_in_length=True
+batch_fn.batch_size_per_device = 64
+batch_fn.eval_batch_size = 64
+batch_fn.max_eval_length = 1024
+batch_fn.buckets_include_inputs_in_length=True
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin
index 8546162c5..310896205 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin
@@ -3,12 +3,12 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 64
-batch_fun.eval_batch_size = 64
-batch_fun.max_eval_length = 1024
-batch_fun.buckets_include_inputs_in_length=True
+batch_fn.batch_size_per_device = 64
+batch_fn.eval_batch_size = 64
+batch_fn.max_eval_length = 1024
+batch_fn.buckets_include_inputs_in_length=True
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
index e02837f48..97f29e528 100644
--- a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
+++ b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
@@ -4,12 +4,12 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 256
-batch_fun.bucket_length = 32
-batch_fun.buckets = None
-batch_fun.eval_batch_size = 512
+batch_fn.batch_size_per_device = 256
+batch_fn.bucket_length = 32
+batch_fn.buckets = None
+batch_fn.eval_batch_size = 512
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
index da4d2c5c5..6875c8445 100644
--- a/tensor2tensor/trax/inputs.py
+++ b/tensor2tensor/trax/inputs.py
@@ -35,13 +35,13 @@
 # Inputs is the trax tuple defining the input streams and shapes.
 # * train_stream: training data that will be used for training
 #     may include all the augmentation or selection the training wants
-#     the shape of examples is [batch_fun.batch_size, ...]
+#     the shape of examples is [batch_fn.batch_size, ...]
 # * train_eval_stream: training data used for evaluation
 #     examples from training data but usually without augmentation
-#     the shape of examples is [batch_fun.eval_batch_size, ...]
+#     the shape of examples is [batch_fn.eval_batch_size, ...]
 # * eval_stream: evaluation data stream
 #     examples from evaluation data, usually without augmentation
-#     the shape of examples is [batch_fun.eval_batch_size, ...]
+#     the shape of examples is [batch_fn.eval_batch_size, ...]
 # * input_shape: the shape of inputs
 #     the [...] above, without batch size
 # * input_dtype: the data type of inputs
@@ -398,11 +398,11 @@ def _train_and_eval_dataset_v1(problem_name, data_dir):
 
 @gin.configurable(blacklist=['dataset', 'training', 'shapes',
                              'target_names', 'n_devices'])
-def batch_fun(dataset, training, shapes, target_names, n_devices,
-              batch_size_per_device=32, batch_size=None, eval_batch_size=32,
-              bucket_length=32, buckets=None,
-              buckets_include_inputs_in_length=False,
-              batch_shuffle_size=128, max_eval_length=None):
+def batch_fn(dataset, training, shapes, target_names, n_devices,
+             batch_size_per_device=32, batch_size=None, eval_batch_size=32,
+             bucket_length=32, buckets=None,
+             buckets_include_inputs_in_length=False,
+             batch_shuffle_size=128, max_eval_length=None):
   """Batching function."""
   del target_names
   # Batch size is batch_size_per_device * n_devices unless given directly.
@@ -620,7 +620,7 @@ def append_targets(example):
   shapes = {k: features_info[k].shape for k in features_info}
   shapes = (shapes, shapes[target_names[0]])
   dataset = dataset.shuffle(shuffle_buffer_size)
-  dataset = batch_fun(dataset, training, shapes, target_names, n_devices)
+  dataset = batch_fn(dataset, training, shapes, target_names, n_devices)
   return dataset.prefetch(2)
 
 
diff --git a/tensor2tensor/trax/inputs_test.py b/tensor2tensor/trax/inputs_test.py
index 80c76b01d..9c145f9f3 100644
--- a/tensor2tensor/trax/inputs_test.py
+++ b/tensor2tensor/trax/inputs_test.py
@@ -44,10 +44,10 @@ class InputsTest(tf.test.TestCase):
   def setUp(self):
     gin.clear_config()
 
-  def test_batch_fun(self):
+  def test_batch_fn(self):
     dataset = test_dataset_ints([32])
     dataset = dataset.repeat(10)
-    batches = inputs.batch_fun(
+    batches = inputs.batch_fn(
         dataset, True, ([None], [None]), [], 1, batch_size=10)
     count = 0
     for example in tfds.as_numpy(batches):
@@ -55,10 +55,10 @@ def test_batch_fun(self):
       self.assertEqual(example[0].shape[0], 10)  # Batch size = 10.
     self.assertEqual(count, 1)  # Just one batch here.
 
-  def test_batch_fun_n_devices(self):
+  def test_batch_fn_n_devices(self):
     dataset = test_dataset_ints([32])
     dataset = dataset.repeat(9)
-    batches = inputs.batch_fun(
+    batches = inputs.batch_fn(
         dataset, True, ([None], [None]), [], 9, batch_size=10)
     count = 0
     for example in tfds.as_numpy(batches):
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
index 3a33ddb93..b63462595 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
@@ -4,11 +4,11 @@ import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.rl
 import tensor2tensor.trax.rl.envs
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 1
-batch_fun.eval_batch_size = 16
-batch_fun.max_eval_length = 12288  # 64 * 64 * 3
+batch_fn.batch_size_per_device = 1
+batch_fn.eval_batch_size = 16
+batch_fn.max_eval_length = 12288  # 64 * 64 * 3
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
index 318f2b458..2de8d8960 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
@@ -4,11 +4,11 @@ import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.rl
 import tensor2tensor.trax.rl.envs
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 256
-batch_fun.eval_batch_size = 256
-batch_fun.max_eval_length = 2048
+batch_fn.batch_size_per_device = 256
+batch_fn.eval_batch_size = 256
+batch_fn.max_eval_length = 2048
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
index b32d106c9..129bab5ae 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
@@ -4,13 +4,13 @@ import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.rl
 import tensor2tensor.trax.rl.envs
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 128
-batch_fun.eval_batch_size = 128
-batch_fun.bucket_length = 64
-batch_fun.max_eval_length = 512
-batch_fun.buckets_include_inputs_in_length = True
+batch_fn.batch_size_per_device = 128
+batch_fn.eval_batch_size = 128
+batch_fn.bucket_length = 64
+batch_fn.max_eval_length = 512
+batch_fn.buckets_include_inputs_in_length = True
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
index 7eb617244..aebc9955c 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
@@ -4,12 +4,12 @@ import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.rl
 import tensor2tensor.trax.rl.envs
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 256
-batch_fun.bucket_length = 32
-batch_fun.buckets = None
-batch_fun.eval_batch_size = 512
+batch_fn.batch_size_per_device = 256
+batch_fn.bucket_length = 32
+batch_fn.buckets = None
+batch_fn.eval_batch_size = 512
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin
index b0adfdb9c..040da7703 100644
--- a/tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin
@@ -4,12 +4,12 @@ import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.rl.envs
 import tensor2tensor.trax.rl.trainers
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size = 32
-batch_fun.bucket_length = 32
-batch_fun.buckets = None
-batch_fun.eval_batch_size = 32
+batch_fn.batch_size = 32
+batch_fn.bucket_length = 32
+batch_fn.buckets = None
+batch_fn.eval_batch_size = 32
 
 # Parameters for inputs:
 # ==============================================================================

From bc8f167136068ebc95f7e9e179c3ab9b9682f9e6 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Mon, 30 Sep 2019 11:55:21 -0700
Subject: [PATCH 2503/2720] Update/add OnlineTune config files.

PiperOrigin-RevId: 272032842
---
 .../trax/configs/transformer_lm1b_16gb.gin    | 130 +++++++++++++++++
 .../configs/transformer_lm_wmt_ende_16gb.gin  | 137 ++++++++++++++++++
 .../trax/configs/transformer_ptb_16gb.gin     | 131 +++++++++++++++++
 .../trax/configs/wide_resnet_cifar10_8gb.gin  |  39 ++++-
 ...nline_tune_transformer_imagenet64_16gb.gin |   2 +-
 .../env_online_tune_transformer_lm1b_16gb.gin |   8 +-
 ...line_tune_transformer_lm_wmt_ende_16gb.gin |  10 +-
 .../env_online_tune_transformer_ptb_16gb.gin  |  97 +++++++++++++
 ...v_online_tune_wide_resnet_cifar10_8gb.gin} |   1 +
 .../trax/rl/configs/ppo_online_tune.gin       |   2 +-
 tensor2tensor/trax/rl/envs/online_tune_env.py |   3 +-
 tensor2tensor/trax/trax.py                    |   8 +-
 12 files changed, 551 insertions(+), 17 deletions(-)
 create mode 100644 tensor2tensor/trax/configs/transformer_lm1b_16gb.gin
 create mode 100644 tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin
 create mode 100644 tensor2tensor/trax/configs/transformer_ptb_16gb.gin
 create mode 100644 tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin
 rename tensor2tensor/trax/rl/configs/{env_online_tune_wide_resnet_cifar10.gin => env_online_tune_wide_resnet_cifar10_8gb.gin} (98%)

diff --git a/tensor2tensor/trax/configs/transformer_lm1b_16gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_16gb.gin
new file mode 100644
index 000000000..ea2766d0f
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_lm1b_16gb.gin
@@ -0,0 +1,130 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 256
+batch_fun.eval_batch_size = 256
+batch_fun.max_eval_length = 2048
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_languagemodel_lm1b32k'
+inputs.input_name = 'targets'
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 1.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
+lm1b_preprocess.max_target_length = 512
+lm1b_preprocess.max_eval_target_length = 2048
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 500
+train.eval_steps = 1
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerLM
+train.optimizer = @trax.optimizers.Adafactor
+train.train_steps = 50000
+
+# Parameters for DotProductCausalAttention:
+# ==============================================================================
+DotProductCausalAttention.dropout = 0.1
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
+TransformerLM.d_model = 512
+TransformerLM.d_ff = 2048
+TransformerLM.dropout = 0.1
+TransformerLM.max_len = 2048
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 8
+TransformerLM.n_layers = 6
+TransformerLM.vocab_size = 32000
+
+# ==============================================================================
+# Parameters for the RL hyperparameter tuner; turn on with
+# train.lr_schedule=@learning_rate.PolicySchedule and set
+# PolicySchedule.policy_dir.
+# ==============================================================================
+
+# Parameters for PolicySchedule:
+# ==============================================================================
+PolicySchedule.observation_metrics = (
+    ("train", "metrics/accuracy"),
+    ("train", "metrics/loss"),
+    ("eval", "metrics/accuracy"),
+    ("eval", "metrics/loss"),
+)
+PolicySchedule.include_controls_in_observation = False
+PolicySchedule.control_configs = (
+    ("learning_rate", 1e-3, (1e-9, 1e-2), False),
+    ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
+
+    ("dropout_embedding", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_initial", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_initial", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_initial", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_middle", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_middle", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_middle", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_final", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_final", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_final", 0.1, (0.0, 0.9), True),
+)
+PolicySchedule.observation_range = (0.0, 10.0)
+PolicySchedule.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
+PolicySchedule.policy_and_value_model = @trax.models.TransformerDecoder
+PolicySchedule.policy_and_value_two_towers = False
+
+# Parameters for train:
+# ==============================================================================
+train.nontrainable_param_map = {
+    # "dropout_{layer_type}_{block_index}": "dropout_{layer_type}_{block_group}"
+
+    "dropout_attention_0": "dropout_attention_initial",
+    "dropout_ff_middle_0": "dropout_ff_middle_initial",
+    "dropout_ff_final_0": "dropout_ff_final_initial",
+
+    "dropout_attention_1": "dropout_attention_middle",
+    "dropout_ff_middle_1": "dropout_ff_middle_middle",
+    "dropout_ff_final_1": "dropout_ff_final_middle",
+    "dropout_attention_2": "dropout_attention_middle",
+    "dropout_ff_middle_2": "dropout_ff_middle_middle",
+    "dropout_ff_final_2": "dropout_ff_final_middle",
+    "dropout_attention_3": "dropout_attention_middle",
+    "dropout_ff_middle_3": "dropout_ff_middle_middle",
+    "dropout_ff_final_3": "dropout_ff_final_middle",
+    "dropout_attention_4": "dropout_attention_middle",
+    "dropout_ff_middle_4": "dropout_ff_middle_middle",
+    "dropout_ff_final_4": "dropout_ff_final_middle",
+
+    "dropout_attention_5": "dropout_attention_final",
+    "dropout_ff_middle_5": "dropout_ff_middle_final",
+    "dropout_ff_final_5": "dropout_ff_final_final",
+}
+
+# Parameters for TransformerDecoder:
+# ==============================================================================
+TransformerDecoder.d_model = 64
+TransformerDecoder.d_ff = 128
+TransformerDecoder.dropout = 0.0
+TransformerDecoder.n_heads = 2
+TransformerDecoder.n_layers = 2
diff --git a/tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin b/tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin
new file mode 100644
index 000000000..b61558110
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin
@@ -0,0 +1,137 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 64
+batch_fun.eval_batch_size = 64
+batch_fun.bucket_length=64
+batch_fun.max_eval_length = 1024
+batch_fun.buckets_include_inputs_in_length=True
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_translate_ende_wmt32k'
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 1.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 10000
+
+# Parameters for Adafactor:
+# ==============================================================================
+Adafactor.beta1 = 0.0
+Adafactor.decay_rate = 0.8
+Adafactor.clipping_threshold = 1.0
+Adafactor.epsilon1 = 1e-30
+Adafactor.epsilon2 = 0.001
+Adafactor.factored = True
+Adafactor.multiply_by_parameter_scale = True
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_concat_preprocess
+wmt_concat_preprocess.max_length = 255
+wmt_concat_preprocess.max_eval_length = 511
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 500
+train.eval_steps = 1
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerLM
+train.train_steps = 50000
+train.optimizer = @trax.optimizers.Adafactor
+train.has_weights = True
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.d_model = 512
+TransformerLM.d_ff = 2048
+TransformerLM.dropout = 0.1
+TransformerLM.max_len = 2048
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 8
+TransformerLM.n_layers = 6
+TransformerLM.vocab_size = 33300
+
+# ==============================================================================
+# Parameters for the RL hyperparameter tuner; turn on with
+# train.lr_schedule=@learning_rate.PolicySchedule and set
+# PolicySchedule.policy_dir.
+# ==============================================================================
+
+# Parameters for PolicySchedule:
+# ==============================================================================
+PolicySchedule.observation_metrics = (
+    ("train", "metrics/accuracy"),
+    ("train", "metrics/loss"),
+    ("eval", "metrics/accuracy"),
+    ("eval", "metrics/loss"),
+)
+PolicySchedule.include_controls_in_observation = False
+PolicySchedule.control_configs = (
+    ("learning_rate", 1e-3, (1e-9, 1e-2), False),
+    ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
+
+    ("dropout_embedding", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_initial", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_initial", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_initial", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_middle", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_middle", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_middle", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_final", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_final", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_final", 0.1, (0.0, 0.9), True),
+)
+PolicySchedule.observation_range = (0.0, 10.0)
+PolicySchedule.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
+PolicySchedule.policy_and_value_model = @trax.models.TransformerDecoder
+PolicySchedule.policy_and_value_two_towers = False
+
+# Parameters for train:
+# ==============================================================================
+train.nontrainable_param_map = {
+    # "dropout_{layer_type}_{block_index}": "dropout_{layer_type}_{block_group}"
+
+    "dropout_attention_0": "dropout_attention_initial",
+    "dropout_ff_middle_0": "dropout_ff_middle_initial",
+    "dropout_ff_final_0": "dropout_ff_final_initial",
+
+    "dropout_attention_1": "dropout_attention_middle",
+    "dropout_ff_middle_1": "dropout_ff_middle_middle",
+    "dropout_ff_final_1": "dropout_ff_final_middle",
+    "dropout_attention_2": "dropout_attention_middle",
+    "dropout_ff_middle_2": "dropout_ff_middle_middle",
+    "dropout_ff_final_2": "dropout_ff_final_middle",
+    "dropout_attention_3": "dropout_attention_middle",
+    "dropout_ff_middle_3": "dropout_ff_middle_middle",
+    "dropout_ff_final_3": "dropout_ff_final_middle",
+    "dropout_attention_4": "dropout_attention_middle",
+    "dropout_ff_middle_4": "dropout_ff_middle_middle",
+    "dropout_ff_final_4": "dropout_ff_final_middle",
+
+    "dropout_attention_5": "dropout_attention_final",
+    "dropout_ff_middle_5": "dropout_ff_middle_final",
+    "dropout_ff_final_5": "dropout_ff_final_final",
+}
+
+# Parameters for TransformerDecoder:
+# ==============================================================================
+TransformerDecoder.d_model = 64
+TransformerDecoder.d_ff = 128
+TransformerDecoder.dropout = 0.0
+TransformerDecoder.n_heads = 2
+TransformerDecoder.n_layers = 2
diff --git a/tensor2tensor/trax/configs/transformer_ptb_16gb.gin b/tensor2tensor/trax/configs/transformer_ptb_16gb.gin
new file mode 100644
index 000000000..ed6c337e7
--- /dev/null
+++ b/tensor2tensor/trax/configs/transformer_ptb_16gb.gin
@@ -0,0 +1,131 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.learning_rate
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 64
+batch_fun.eval_batch_size = 512
+batch_fun.max_eval_length = 2048
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_languagemodel_ptb10k'
+inputs.input_name = 'targets'
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 1.0
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
+lm1b_preprocess.max_target_length = 512
+lm1b_preprocess.max_eval_target_length = 2048
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 200
+train.eval_steps = 2
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerLM
+train.optimizer = @trax.optimizers.Adafactor
+train.train_steps = 20000
+
+# Parameters for DotProductCausalAttention:
+# ==============================================================================
+DotProductCausalAttention.dropout = 0.5
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
+TransformerLM.d_model = 512
+TransformerLM.d_ff = 2048
+TransformerLM.dropout = 0.5
+TransformerLM.max_len = 2048
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 8
+TransformerLM.n_layers = 6
+TransformerLM.vocab_size = 10240
+
+# ==============================================================================
+# Parameters for the RL hyperparameter tuner; turn on with
+# train.lr_schedule=@learning_rate.PolicySchedule and set
+# PolicySchedule.policy_dir.
+# ==============================================================================
+
+# Parameters for PolicySchedule:
+# ==============================================================================
+PolicySchedule.observation_metrics = (
+    ("train", "metrics/accuracy"),
+    ("train", "metrics/loss"),
+    ("eval", "metrics/accuracy"),
+    ("eval", "metrics/loss"),
+)
+PolicySchedule.include_controls_in_observation = False
+PolicySchedule.control_configs = (
+    ("learning_rate", 1e-3, (1e-9, 1e-2), False),
+    ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
+
+    ("dropout_embedding", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_initial", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_initial", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_initial", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_middle", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_middle", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_middle", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_final", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_final", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_final", 0.1, (0.0, 0.9), True),
+)
+PolicySchedule.observation_range = (0.0, 10.0)
+PolicySchedule.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
+PolicySchedule.policy_and_value_model = @trax.models.TransformerDecoder
+PolicySchedule.policy_and_value_two_towers = False
+
+# Parameters for train:
+# ==============================================================================
+train.nontrainable_param_map = {
+    # "dropout_{layer_type}_{block_index}": "dropout_{layer_type}_{block_group}"
+
+    "dropout_attention_0": "dropout_attention_initial",
+    "dropout_ff_middle_0": "dropout_ff_middle_initial",
+    "dropout_ff_final_0": "dropout_ff_final_initial",
+
+    "dropout_attention_1": "dropout_attention_middle",
+    "dropout_ff_middle_1": "dropout_ff_middle_middle",
+    "dropout_ff_final_1": "dropout_ff_final_middle",
+    "dropout_attention_2": "dropout_attention_middle",
+    "dropout_ff_middle_2": "dropout_ff_middle_middle",
+    "dropout_ff_final_2": "dropout_ff_final_middle",
+    "dropout_attention_3": "dropout_attention_middle",
+    "dropout_ff_middle_3": "dropout_ff_middle_middle",
+    "dropout_ff_final_3": "dropout_ff_final_middle",
+    "dropout_attention_4": "dropout_attention_middle",
+    "dropout_ff_middle_4": "dropout_ff_middle_middle",
+    "dropout_ff_final_4": "dropout_ff_final_middle",
+
+    "dropout_attention_5": "dropout_attention_final",
+    "dropout_ff_middle_5": "dropout_ff_middle_final",
+    "dropout_ff_final_5": "dropout_ff_final_final",
+}
+
+# Parameters for TransformerDecoder:
+# ==============================================================================
+TransformerDecoder.d_model = 64
+TransformerDecoder.d_ff = 128
+TransformerDecoder.dropout = 0.0
+TransformerDecoder.n_heads = 2
+TransformerDecoder.n_layers = 2
diff --git a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
index 97f29e528..7edefa0e4 100644
--- a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
+++ b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
@@ -18,7 +18,7 @@ inputs.dataset_name = 'cifar10'
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-EvalAdjustingSchedule.constant = 4.0
+EvalAdjustingSchedule.constant = 0.5
 MultifactorSchedule.factors = 'constant * linear_warmup'
 MultifactorSchedule.warmup_steps = 400
 
@@ -39,10 +39,43 @@ WideResnet.n_output_classes = 10
 
 # Parameters for train:
 # ==============================================================================
-train.eval_frequency = 200
+train.eval_frequency = 100
 train.eval_steps = 10
 train.inputs = @trax.inputs.inputs
 train.model = @trax.models.WideResnet
 train.optimizer = @trax.optimizers.Momentum
-train.train_steps = 1000000
+train.train_steps = 10000
 train.lr_schedule = @learning_rate.EvalAdjustingSchedule
+
+# ==============================================================================
+# Parameters for the RL hyperparameter tuner; turn on with
+# train.lr_schedule=@learning_rate.PolicySchedule and set
+# PolicySchedule.policy_dir.
+# ==============================================================================
+
+# Parameters for PolicySchedule:
+# ==============================================================================
+PolicySchedule.observation_metrics = (
+    ("train", "metrics/accuracy"),
+    ("train", "metrics/loss"),
+    ("eval", "metrics/accuracy"),
+    ("eval", "metrics/loss"),
+)
+PolicySchedule.include_controls_in_observation = False
+PolicySchedule.control_configs = (
+    ("learning_rate", 0.1, (1e-9, 10.0), False),
+    ("weight_decay_rate", 1e-5, (1e-9, 0.1), False),
+    ("mass", 0.9, (0.0, 0.99), True),
+)
+PolicySchedule.observation_range = (0.0, 10.0)
+PolicySchedule.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
+PolicySchedule.policy_and_value_model = @trax.models.TransformerDecoder
+PolicySchedule.policy_and_value_two_towers = False
+
+# Parameters for TransformerDecoder:
+# ==============================================================================
+TransformerDecoder.d_model = 64
+TransformerDecoder.d_ff = 128
+TransformerDecoder.dropout = 0.0
+TransformerDecoder.n_heads = 2
+TransformerDecoder.n_layers = 2
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
index b63462595..3aeb46b66 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
@@ -97,7 +97,7 @@ OnlineTuneEnv.nontrainable_param_map = {
     "dropout_ff_middle_2": "dropout_ff_middle_final",
     "dropout_ff_final_2": "dropout_ff_final_final",
 }
-OnlineTuneEnv.include_controls_in_observation = True
+OnlineTuneEnv.include_controls_in_observation = False
 OnlineTuneEnv.train_steps = 150
 OnlineTuneEnv.eval_steps = 2
 OnlineTuneEnv.env_steps = 100
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
index 2de8d8960..732260655 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
@@ -52,6 +52,9 @@ TransformerLM.vocab_size = 32000
 OnlineTuneEnv.inputs = @trax.inputs.inputs
 OnlineTuneEnv.model = @trax.models.TransformerLM
 OnlineTuneEnv.optimizer = @trax.optimizers.Adafactor
+OnlineTuneEnv.train_steps = 500
+OnlineTuneEnv.eval_steps = 1
+OnlineTuneEnv.env_steps = 100
 OnlineTuneEnv.control_configs = (
     ("learning_rate", 1e-3, (1e-9, 1e-2), False),
     ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
@@ -94,9 +97,6 @@ OnlineTuneEnv.nontrainable_param_map = {
     "dropout_ff_middle_5": "dropout_ff_middle_final",
     "dropout_ff_final_5": "dropout_ff_final_final",
 }
-OnlineTuneEnv.include_controls_in_observation = True
-OnlineTuneEnv.train_steps = 300
-OnlineTuneEnv.eval_steps = 1
-OnlineTuneEnv.env_steps = 100
+OnlineTuneEnv.include_controls_in_observation = False
 OnlineTuneEnv.observation_range = (0.0, 10.0)
 OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
index 129bab5ae..8f1eaf9d1 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
@@ -48,6 +48,10 @@ TransformerLM.vocab_size = 33300
 OnlineTuneEnv.inputs = @trax.inputs.inputs
 OnlineTuneEnv.model = @trax.models.TransformerLM
 OnlineTuneEnv.optimizer = @trax.optimizers.Adafactor
+OnlineTuneEnv.train_steps = 500
+OnlineTuneEnv.eval_steps = 1
+OnlineTuneEnv.env_steps = 100
+OnlineTuneEnv.has_weights = True
 OnlineTuneEnv.control_configs = (
     ("learning_rate", 1e-3, (1e-9, 1e-2), False),
     ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
@@ -90,10 +94,6 @@ OnlineTuneEnv.nontrainable_param_map = {
     "dropout_ff_middle_5": "dropout_ff_middle_final",
     "dropout_ff_final_5": "dropout_ff_final_final",
 }
-OnlineTuneEnv.include_controls_in_observation = True
-OnlineTuneEnv.train_steps = 500
-OnlineTuneEnv.eval_steps = 1
-OnlineTuneEnv.env_steps = 100
+OnlineTuneEnv.include_controls_in_observation = False
 OnlineTuneEnv.observation_range = (0.0, 10.0)
 OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
-OnlineTuneEnv.has_weights = True
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin
new file mode 100644
index 000000000..6d05acab7
--- /dev/null
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin
@@ -0,0 +1,97 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.rl
+import tensor2tensor.trax.rl.envs
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 64
+batch_fun.eval_batch_size = 512
+batch_fun.max_eval_length = 2048
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_languagemodel_ptb10k'
+inputs.input_name = 'targets'
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
+lm1b_preprocess.max_target_length = 512
+lm1b_preprocess.max_eval_target_length = 2048
+
+# Parameters for DotProductCausalAttention:
+# ==============================================================================
+DotProductCausalAttention.dropout = 0.1
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
+TransformerLM.d_model = 512
+TransformerLM.d_ff = 2048
+TransformerLM.dropout = 0.1
+TransformerLM.max_len = 2048
+TransformerLM.mode = 'train'
+TransformerLM.n_heads = 8
+TransformerLM.n_layers = 6
+TransformerLM.vocab_size = 10240
+
+# Parameters for OnlineTuneEnv:
+# ==============================================================================
+OnlineTuneEnv.inputs = @trax.inputs.inputs
+OnlineTuneEnv.model = @trax.models.TransformerLM
+OnlineTuneEnv.optimizer = @trax.optimizers.Adafactor
+OnlineTuneEnv.train_steps = 200
+OnlineTuneEnv.eval_steps = 2
+OnlineTuneEnv.env_steps = 100
+OnlineTuneEnv.control_configs = (
+    ("learning_rate", 1e-3, (1e-9, 1e-2), False),
+    ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
+
+    ("dropout_embedding", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_initial", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_initial", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_initial", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_middle", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_middle", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_middle", 0.1, (0.0, 0.9), True),
+
+    ("dropout_attention_final", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_middle_final", 0.1, (0.0, 0.9), True),
+    ("dropout_ff_final_final", 0.1, (0.0, 0.9), True),
+)
+OnlineTuneEnv.nontrainable_param_map = {
+    # "dropout_{layer_type}_{block_index}": "dropout_{layer_type}_{block_group}"
+
+    "dropout_attention_0": "dropout_attention_initial",
+    "dropout_ff_middle_0": "dropout_ff_middle_initial",
+    "dropout_ff_final_0": "dropout_ff_final_initial",
+
+    "dropout_attention_1": "dropout_attention_middle",
+    "dropout_ff_middle_1": "dropout_ff_middle_middle",
+    "dropout_ff_final_1": "dropout_ff_final_middle",
+    "dropout_attention_2": "dropout_attention_middle",
+    "dropout_ff_middle_2": "dropout_ff_middle_middle",
+    "dropout_ff_final_2": "dropout_ff_final_middle",
+    "dropout_attention_3": "dropout_attention_middle",
+    "dropout_ff_middle_3": "dropout_ff_middle_middle",
+    "dropout_ff_final_3": "dropout_ff_final_middle",
+    "dropout_attention_4": "dropout_attention_middle",
+    "dropout_ff_middle_4": "dropout_ff_middle_middle",
+    "dropout_ff_final_4": "dropout_ff_final_middle",
+
+    "dropout_attention_5": "dropout_attention_final",
+    "dropout_ff_middle_5": "dropout_ff_middle_final",
+    "dropout_ff_final_5": "dropout_ff_final_final",
+}
+OnlineTuneEnv.include_controls_in_observation = False
+OnlineTuneEnv.observation_range = (0.0, 10.0)
+OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10_8gb.gin
similarity index 98%
rename from tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
rename to tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10_8gb.gin
index aebc9955c..9b8a8e6d5 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10_8gb.gin
@@ -43,6 +43,7 @@ OnlineTuneEnv.optimizer = @trax.optimizers.Momentum
 OnlineTuneEnv.control_configs = (
     ("learning_rate", 0.1, (1e-9, 10.0), False),
     ("weight_decay_rate", 1e-5, (1e-9, 0.1), False),
+    ("mass", 0.9, (0.0, 0.99), True),
 )
 OnlineTuneEnv.include_controls_in_observation = False
 OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
diff --git a/tensor2tensor/trax/rl/configs/ppo_online_tune.gin b/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
index 5b57f4fd2..bfaa6e4fe 100644
--- a/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
+++ b/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
@@ -28,7 +28,7 @@ PPO.random_seed = None
 PPO.gamma = 1.0
 PPO.lambda_ = 0.95
 PPO.c1 = 1.0
-PPO.c2 = 0.03
+PPO.c2 = 0.1
 PPO.done_frac_for_policy_save = 0
 PPO.len_history_for_policy = None
 PPO.separate_eval = False
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index dba9e67a9..a66d59a57 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -25,6 +25,7 @@
 import gym
 
 from tensor2tensor.trax import inputs as trax_inputs
+from tensor2tensor.trax import models as trax_models
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
 from tensor2tensor.trax.rl import online_tune
@@ -47,8 +48,8 @@ class OnlineTuneEnv(gym.Env):
   DEFAULT_ACTION_MULTIPLIERS = [1.0 / 1.5, 1.0 / 1.25, 1.0, 1.25, 1.5]
 
   def __init__(self,
-               model,
                output_dir,
+               model=trax_models.TransformerLM,
                trainer_class=trax.Trainer,
                loss_fn=trax.loss,
                optimizer=trax_opt.Adafactor,
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 7b43e1192..2775d458f 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -974,7 +974,8 @@ def train(output_dir,
           random_seed=None,
           save_graphs=True,
           save_backward_graph=False,
-          has_weights=False):
+          has_weights=False,
+          nontrainable_param_map=None):
   """Train the model on the inputs.
 
   Args:
@@ -999,13 +1000,16 @@ def train(output_dir,
     save_graphs: bool, if True, save computation graph to file.
     save_backward_graph: bool, if True, save backward graph to file too.
     has_weights: bool, whether weights are included in the inputs.
+    nontrainable_param_map: dict, mapping from model nontrainable parameter
+      names to control names in PolicySchedule.
   Returns:
     trax.State
   """
   trainer = trainer_class(model, loss_fn, optimizer, lr_schedule, inputs,
                           output_dir,
                           random_seed=random_seed, n_devices=n_devices,
-                          save_steps=save_steps, has_weights=has_weights)
+                          save_steps=save_steps, has_weights=has_weights,
+                          nontrainable_param_map=nontrainable_param_map)
 
   epoch_steps = [train_steps]  # Only training if eval_frequency is 0 or None
   if eval_frequency and eval_steps > 0:

From 949e98501cfca97816ace3950255d0c2204e570c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 30 Sep 2019 12:04:59 -0700
Subject: [PATCH 2504/2720] Unify configs for reformer experiments

PiperOrigin-RevId: 272035162
---
 .../trax/configs/reformer_enwik8.gin          |  67 ++++-----
 .../trax/configs/reformer_enwik8_rev.gin      | 135 ------------------
 .../trax/configs/reformer_imagenet64.gin      |  59 +++-----
 .../trax/configs/reformer_imagenet64_rev.gin  | 132 -----------------
 .../trax/configs/transformer_enwik8.gin       |  69 ---------
 .../trax/configs/transformer_imagenet64.gin   |  80 -----------
 .../configs/transformer_revnet_enwik8.gin     |  57 --------
 .../transformer_revnet_imagenet64_8gb.gin     |  73 ----------
 tensor2tensor/trax/layers/combinators.py      |  14 ++
 .../models/research/transformer_revnet.py     |  14 --
 tensor2tensor/trax/models/transformer.py      |  13 +-
 11 files changed, 71 insertions(+), 642 deletions(-)
 delete mode 100644 tensor2tensor/trax/configs/reformer_enwik8_rev.gin
 delete mode 100644 tensor2tensor/trax/configs/reformer_imagenet64_rev.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_enwik8.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_imagenet64.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_revnet_enwik8.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin

diff --git a/tensor2tensor/trax/configs/reformer_enwik8.gin b/tensor2tensor/trax/configs/reformer_enwik8.gin
index aa514ab62..d90119520 100644
--- a/tensor2tensor/trax/configs/reformer_enwik8.gin
+++ b/tensor2tensor/trax/configs/reformer_enwik8.gin
@@ -5,42 +5,20 @@ import tensor2tensor.trax.trax
 
 # Parameters that will vary between experiments:
 # ==============================================================================
-train.model = @trax.models.TransformerLM
-# inputs.n_chunks = 16  # Uncomment this line iff using TransformerRevnetLM
+train.model = @trax.models.TransformerRevnetLM
+attn_type = @MergedMultiHashedCausalAttentionV2
+share_qk = True  # Required when using hashed attention
+attn_kv = 64
+n_layers = 3
+dropout = 0.1
 
-TransformerLM.n_layers = 3
-TransformerRevnetLM.n_layers = 3
-
-TransformerLM.dropout = 0.1
-TransformerRevnetLM.dropout = 0.1
-
-# Attention types:
 # MemoryEfficientCausalAttention: full attention
-# MergedHashedCausalAttention: timebin
-# MergedMultiHashedCausalAttentionV2: hashbin
-# TransformerLM.attention_type = @MemoryEfficientCausalAttention
-# TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
-TransformerLM.attention_type = @MergedMultiHashedCausalAttentionV2
-TransformerRevnetLM.attention_type = @MergedMultiHashedCausalAttentionV2
-
-# These three should all be equal, and MergedMultiHashedCausalAttentionV2
-# requires that share_qk be set to True.
-share_qk = True
-MemoryEfficientCausalAttention.share_qk = %share_qk
-TransformerLM.share_qk = %share_qk
-TransformerRevnetLM.share_qk = %share_qk
-
-# These four should all be equal
-attn_kv = 128
-TransformerLM.d_attention_key = %attn_kv
-TransformerLM.d_attention_value = %attn_kv
-TransformerRevnetLM.d_attention_key = %attn_kv
-TransformerRevnetLM.d_attention_value = %attn_kv
+# (no hparams to vary between experiments)
 
-# Use MergedHashedCausalAttention for timebin
+# MergedHashedCausalAttention: use for timebin
 MergedHashedCausalAttention.n_bins = 512
 
-# Use MergedMultiHashedCausalAttentionV2 for hashbin
+# MergedMultiHashedCausalAttentionV2: locality-sensitive hashing (LSH) attention
 MergedMultiHashedCausalAttentionV2.n_bins = 512
 MergedMultiHashedCausalAttentionV2.n_buckets = 1024  # Always 2 * n_bins
 MergedMultiHashedCausalAttentionV2.n_hashes = 2
@@ -57,6 +35,7 @@ batch_fn.max_eval_length = 65536
 inputs.data_dir = None
 inputs.dataset_name = 't2t_enwik8_l65k'
 inputs.input_name = 'targets'
+inputs.n_chunks = 16
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
@@ -81,7 +60,7 @@ train.save_steps = \
 # ==============================================================================
 MemoryEfficientCausalAttention.dropout = 0.0
 MemoryEfficientCausalAttention.loop_stride = 256
-# MemoryEfficientCausalAttention.share_qk: see top
+MemoryEfficientCausalAttention.share_qk = %share_qk
 
 # Parameters for MergedHashedCausalAttention:
 # ==============================================================================
@@ -105,31 +84,33 @@ MergedMultiHashedCausalAttentionV2.dropout = 0.0
 
 # Parameters for TransformerLM:
 # ==============================================================================
-# TransformerLM.attention_type: see top
-# TransformerLM.d_attention_key: see top
-# TransformerLM.d_attention_value: see top
+TransformerLM.attention_type = %attn_type
+TransformerLM.d_attention_key = %attn_kv
+TransformerLM.d_attention_value = %attn_kv
 TransformerLM.d_model = 1024
 TransformerLM.d_ff = 4096
-# TransformerLM.dropout: see top
+TransformerLM.dropout = %dropout
 TransformerLM.max_len = 65536
 TransformerLM.mode = 'train'
 TransformerLM.n_heads = 8
-# TransformerLM.n_layers: see top
-# TransformerLM.share_qk: see top
+TransformerLM.n_layers = %n_layers
+TransformerLM.n_chunks = 16
+TransformerLM.share_qk = %share_qk
 TransformerLM.vocab_size = 258  # Includes pad token and unused EOS token
 
 # Parameters for TransformerRevnetLM:
 # ==============================================================================
-# TransformerRevnetLM.attention_type: see top
-# TransformerRevnetLM.d_attention_key: see top
-# TransformerRevnetLM.d_attention_value: see top
+TransformerRevnetLM.attention_type = %attn_type
+TransformerRevnetLM.d_attention_key = %attn_kv
+TransformerRevnetLM.d_attention_value = %attn_kv
 TransformerRevnetLM.d_model = 1024
 TransformerRevnetLM.d_ff = 4096
-# TransformerRevnetLM.dropout: see top
+TransformerRevnetLM.dropout = %dropout
 TransformerRevnetLM.max_len = 65536
 TransformerRevnetLM.mode = 'train'
 TransformerRevnetLM.n_heads = 8
-# TransformerRevnetLM.n_layers: see top
+TransformerRevnetLM.n_layers = %n_layers
 TransformerRevnetLM.vocab_size = 258  # Includes pad token and unused EOS token
 TransformerRevnetLM.n_chunks = 16
 TransformerRevnetLM.n_attention_chunks = 1
+TransformerRevnetLM.share_qk = %share_qk
diff --git a/tensor2tensor/trax/configs/reformer_enwik8_rev.gin b/tensor2tensor/trax/configs/reformer_enwik8_rev.gin
deleted file mode 100644
index e82007c0c..000000000
--- a/tensor2tensor/trax/configs/reformer_enwik8_rev.gin
+++ /dev/null
@@ -1,135 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters that will vary between experiments:
-# ==============================================================================
-train.model = @trax.models.TransformerRevnetLM
-inputs.n_chunks = 16  # Uncomment this line iff using TransformerRevnetLM
-
-TransformerLM.n_layers = 3
-TransformerRevnetLM.n_layers = 3
-
-TransformerLM.dropout = 0.1
-TransformerRevnetLM.dropout = 0.1
-
-# Attention types:
-# MemoryEfficientCausalAttention: full attention
-# MergedHashedCausalAttention: timebin
-# MergedMultiHashedCausalAttentionV2: hashbin
-# TransformerLM.attention_type = @MemoryEfficientCausalAttention
-# TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
-TransformerLM.attention_type = @MergedMultiHashedCausalAttentionV2
-TransformerRevnetLM.attention_type = @MergedMultiHashedCausalAttentionV2
-
-# These three should all be equal, and MergedMultiHashedCausalAttentionV2
-# requires that share_qk be set to True.
-share_qk = True
-MemoryEfficientCausalAttention.share_qk = %share_qk
-TransformerLM.share_qk = %share_qk
-TransformerRevnetLM.share_qk = %share_qk
-
-# These four should all be equal
-attn_kv = 128
-TransformerLM.d_attention_key = %attn_kv
-TransformerLM.d_attention_value = %attn_kv
-TransformerRevnetLM.d_attention_key = %attn_kv
-TransformerRevnetLM.d_attention_value = %attn_kv
-
-# Use MergedHashedCausalAttention for timebin
-MergedHashedCausalAttention.n_bins = 512
-
-# Use MergedMultiHashedCausalAttentionV2 for hashbin
-MergedMultiHashedCausalAttentionV2.n_bins = 512
-MergedMultiHashedCausalAttentionV2.n_buckets = 1024  # Always 2 * n_bins
-MergedMultiHashedCausalAttentionV2.n_hashes = 2
-MergedMultiHashedCausalAttentionV2.drop_for_hash_rate = 0.0
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 1
-batch_fn.eval_batch_size = 8
-batch_fn.max_eval_length = 65536
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_enwik8_l65k'
-inputs.input_name = 'targets'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 2.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 500
-train.eval_steps = 8
-train.inputs = @trax.inputs.inputs
-# train.model: see top
-train.optimizer = @trax.optimizers.Adafactor
-train.train_steps = 60000
-train.trainer_class = @MemoryEfficientTrainer
-train.save_steps = \
-    [1000, 5000, 10000, 15000, 20000, 25000, 30000, 35000, 40000, 45000, 50000,
-     55000, 60000]
-
-# Parameters for MemoryEfficientCausalAttention:
-# ==============================================================================
-MemoryEfficientCausalAttention.dropout = 0.0
-MemoryEfficientCausalAttention.loop_stride = 256
-# MemoryEfficientCausalAttention.share_qk: see top
-
-# Parameters for MergedHashedCausalAttention:
-# ==============================================================================
-MergedHashedCausalAttention.dropout = 0.0
-# MergedHashedCausalAttention.n_bins: see top
-MergedHashedCausalAttention.bin_by_time = True
-MergedHashedCausalAttention.one_rng = False
-
-# Parameters for MergedMultiHashedCausalAttentionV2:
-# ==============================================================================
-MergedMultiHashedCausalAttentionV2.allow_duplicate_attention = False
-MergedMultiHashedCausalAttentionV2.attend_across_buckets = False
-MergedMultiHashedCausalAttentionV2.rehash_each_round = True
-# MergedMultiHashedCausalAttentionV2.n_bins: see top
-# MergedMultiHashedCausalAttentionV2.n_buckets: see top
-# MergedMultiHashedCausalAttentionV2.n_hashes: see top
-MergedMultiHashedCausalAttentionV2.one_rng = False
-MergedMultiHashedCausalAttentionV2.hard_k = 0
-MergedMultiHashedCausalAttentionV2.dropout = 0.0
-# MergedMultiHashedCausalAttentionV2.drop_for_hash_rate: see top
-
-# Parameters for TransformerLM:
-# ==============================================================================
-# TransformerLM.attention_type: see top
-# TransformerLM.d_attention_key: see top
-# TransformerLM.d_attention_value: see top
-TransformerLM.d_model = 1024
-TransformerLM.d_ff = 4096
-# TransformerLM.dropout: see top
-TransformerLM.max_len = 65536
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 8
-# TransformerLM.n_layers: see top
-# TransformerLM.share_qk: see top
-TransformerLM.vocab_size = 258  # Includes pad token and unused EOS token
-
-# Parameters for TransformerRevnetLM:
-# ==============================================================================
-# TransformerRevnetLM.attention_type: see top
-# TransformerRevnetLM.d_attention_key: see top
-# TransformerRevnetLM.d_attention_value: see top
-TransformerRevnetLM.d_model = 1024
-TransformerRevnetLM.d_ff = 4096
-# TransformerRevnetLM.dropout: see top
-TransformerRevnetLM.max_len = 65536
-TransformerRevnetLM.mode = 'train'
-TransformerRevnetLM.n_heads = 8
-# TransformerRevnetLM.n_layers: see top
-TransformerRevnetLM.vocab_size = 258  # Includes pad token and unused EOS token
-TransformerRevnetLM.n_chunks = 16
-TransformerRevnetLM.n_attention_chunks = 1
diff --git a/tensor2tensor/trax/configs/reformer_imagenet64.gin b/tensor2tensor/trax/configs/reformer_imagenet64.gin
index e1e9bf6be..006f57660 100644
--- a/tensor2tensor/trax/configs/reformer_imagenet64.gin
+++ b/tensor2tensor/trax/configs/reformer_imagenet64.gin
@@ -5,39 +5,19 @@ import tensor2tensor.trax.trax
 
 # Parameters that will vary between experiments:
 # ==============================================================================
-train.model = @trax.models.TransformerLM
-# inputs.n_chunks = 16  # Uncomment this line iff using TransformerRevnetLM
+train.model = @trax.models.TransformerRevnetLM
+attn_type = @MergedMultiHashedCausalAttentionV2
+share_qk = True  # Required when using hashed attention
+attn_kv = 64
+n_layers = 3
 
-TransformerLM.n_layers = 3
-TransformerRevnetLM.n_layers = 3
-
-# Attention types:
 # MemoryEfficientCausalAttention: full attention
-# MergedHashedCausalAttention: timebin
-# MergedMultiHashedCausalAttentionV2: hashbin
-# TransformerLM.attention_type = @MemoryEfficientCausalAttention
-# TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
-TransformerLM.attention_type = @MergedMultiHashedCausalAttentionV2
-TransformerRevnetLM.attention_type = @MergedMultiHashedCausalAttentionV2
-
-# These three should all be equal, and MergedMultiHashedCausalAttentionV2
-# requires that share_qk be set to True.
-share_qk = True
-MemoryEfficientCausalAttention.share_qk = %share_qk
-TransformerLM.share_qk = %share_qk
-TransformerRevnetLM.share_qk = %share_qk
+# (no hparams to vary between experiments)
 
-# These four should all be equal
-attn_kv = 128
-TransformerLM.d_attention_key = %attn_kv
-TransformerLM.d_attention_value = %attn_kv
-TransformerRevnetLM.d_attention_key = %attn_kv
-TransformerRevnetLM.d_attention_value = %attn_kv
-
-# Use MergedHashedCausalAttention for timebin
+# MergedHashedCausalAttention: use for timebin
 MergedHashedCausalAttention.n_bins = 64
 
-# Use MergedMultiHashedCausalAttentionV2 for hashbin
+# MergedMultiHashedCausalAttentionV2: locality-sensitive hashing (LSH) attention
 MergedMultiHashedCausalAttentionV2.n_bins = 96
 MergedMultiHashedCausalAttentionV2.n_buckets = 192  # Always 2 * n_bins
 MergedMultiHashedCausalAttentionV2.n_hashes = 2
@@ -54,6 +34,7 @@ batch_fn.max_eval_length = 12288  # 64 * 64 * 3
 inputs.data_dir = None
 inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
 inputs.input_name = 'targets'
+inputs.n_chunks = 16
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
@@ -78,7 +59,7 @@ train.save_steps = \
 # ==============================================================================
 MemoryEfficientCausalAttention.dropout = 0.0
 MemoryEfficientCausalAttention.loop_stride = 512
-# MemoryEfficientCausalAttention.share_qk: see top
+MemoryEfficientCausalAttention.share_qk = %share_qk
 
 # Parameters for MergedHashedCausalAttention:
 # ==============================================================================
@@ -102,31 +83,33 @@ MergedMultiHashedCausalAttentionV2.dropout = 0.0
 
 # Parameters for TransformerLM:
 # ==============================================================================
-# TransformerLM.attention_type: see top
-# TransformerLM.d_attention_key: see top
-# TransformerLM.d_attention_value: see top
+TransformerLM.attention_type = %attn_type
+TransformerLM.d_attention_key = %attn_kv
+TransformerLM.d_attention_value = %attn_kv
 TransformerLM.d_model = 1024
 TransformerLM.d_ff = 4096
 TransformerLM.dropout = 0.0
 TransformerLM.max_len = 12288  # 64 * 64 * 3
 TransformerLM.mode = 'train'
 TransformerLM.n_heads = 8
-# TransformerLM.n_layers: see top
-# TransformerLM.share_qk: see top
+TransformerLM.n_layers = %n_layers
+TransformerLM.n_chunks = 16
+TransformerLM.share_qk = %share_qk
 TransformerLM.vocab_size = 256
 
 # Parameters for TransformerRevnetLM:
 # ==============================================================================
-# TransformerRevnetLM.attention_type: see top
-# TransformerRevnetLM.d_attention_key: see top
-# TransformerRevnetLM.d_attention_value: see top
+TransformerRevnetLM.attention_type = %attn_type
+TransformerRevnetLM.d_attention_key = %attn_kv
+TransformerRevnetLM.d_attention_value = %attn_kv
 TransformerRevnetLM.d_model = 1024
 TransformerRevnetLM.d_ff = 4096
 TransformerRevnetLM.dropout = 0.0
 TransformerRevnetLM.max_len = 12288  # 64 * 64 * 3
 TransformerRevnetLM.mode = 'train'
 TransformerRevnetLM.n_heads = 8
-# TransformerRevnetLM.n_layers: see top
+TransformerRevnetLM.n_layers = %n_layers
 TransformerRevnetLM.vocab_size = 256
 TransformerRevnetLM.n_chunks = 16
 TransformerRevnetLM.n_attention_chunks = 1
+TransformerRevnetLM.share_qk = %share_qk
diff --git a/tensor2tensor/trax/configs/reformer_imagenet64_rev.gin b/tensor2tensor/trax/configs/reformer_imagenet64_rev.gin
deleted file mode 100644
index 673df196d..000000000
--- a/tensor2tensor/trax/configs/reformer_imagenet64_rev.gin
+++ /dev/null
@@ -1,132 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters that will vary between experiments:
-# ==============================================================================
-train.model = @trax.models.TransformerRevnetLM
-inputs.n_chunks = 16  # Uncomment this line iff using TransformerRevnetLM
-
-TransformerLM.n_layers = 3
-TransformerRevnetLM.n_layers = 3
-
-# Attention types:
-# MemoryEfficientCausalAttention: full attention
-# MergedHashedCausalAttention: timebin
-# MergedMultiHashedCausalAttentionV2: hashbin
-# TransformerLM.attention_type = @MemoryEfficientCausalAttention
-# TransformerRevnetLM.attention_type = @MemoryEfficientCausalAttention
-TransformerLM.attention_type = @MergedMultiHashedCausalAttentionV2
-TransformerRevnetLM.attention_type = @MergedMultiHashedCausalAttentionV2
-
-# These three should all be equal, and MergedMultiHashedCausalAttentionV2
-# requires that share_qk be set to True.
-share_qk = True
-MemoryEfficientCausalAttention.share_qk = %share_qk
-TransformerLM.share_qk = %share_qk
-TransformerRevnetLM.share_qk = %share_qk
-
-# These four should all be equal
-attn_kv = 128
-TransformerLM.d_attention_key = %attn_kv
-TransformerLM.d_attention_value = %attn_kv
-TransformerRevnetLM.d_attention_key = %attn_kv
-TransformerRevnetLM.d_attention_value = %attn_kv
-
-# Use MergedHashedCausalAttention for timebin
-MergedHashedCausalAttention.n_bins = 64
-
-# Use MergedMultiHashedCausalAttentionV2 for hashbin
-MergedMultiHashedCausalAttentionV2.n_bins = 96
-MergedMultiHashedCausalAttentionV2.n_buckets = 192  # Always 2 * n_bins
-MergedMultiHashedCausalAttentionV2.n_hashes = 2
-MergedMultiHashedCausalAttentionV2.drop_for_hash_rate = 0.0
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 1
-batch_fn.eval_batch_size = 8
-batch_fn.max_eval_length = 12288  # 64 * 64 * 3
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
-inputs.input_name = 'targets'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 2.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 500
-train.eval_steps = 64
-train.inputs = @trax.inputs.inputs
-# train.model: see top
-train.optimizer = @trax.optimizers.Adafactor
-train.train_steps = 500000
-train.trainer_class = @MemoryEfficientTrainer
-train.save_steps = \
-    [1000, 5000, 10000, 20000, 40000, 60000, 80000,
-     100000, 200000, 300000, 400000, 500000]
-
-# Parameters for MemoryEfficientCausalAttention:
-# ==============================================================================
-MemoryEfficientCausalAttention.dropout = 0.0
-MemoryEfficientCausalAttention.loop_stride = 512
-# MemoryEfficientCausalAttention.share_qk: see top
-
-# Parameters for MergedHashedCausalAttention:
-# ==============================================================================
-MergedHashedCausalAttention.dropout = 0.0
-# MergedHashedCausalAttention.n_bins: see top
-MergedHashedCausalAttention.bin_by_time = True
-MergedHashedCausalAttention.one_rng = False
-
-# Parameters for MergedMultiHashedCausalAttentionV2:
-# ==============================================================================
-MergedMultiHashedCausalAttentionV2.allow_duplicate_attention = False
-MergedMultiHashedCausalAttentionV2.attend_across_buckets = False
-MergedMultiHashedCausalAttentionV2.rehash_each_round = True
-# MergedMultiHashedCausalAttentionV2.n_bins: see top
-# MergedMultiHashedCausalAttentionV2.n_buckets: see top
-# MergedMultiHashedCausalAttentionV2.n_hashes: see top
-MergedMultiHashedCausalAttentionV2.one_rng = False
-MergedMultiHashedCausalAttentionV2.hard_k = 0
-MergedMultiHashedCausalAttentionV2.dropout = 0.0
-# MergedMultiHashedCausalAttentionV2.drop_for_hash_rate: see top
-
-# Parameters for TransformerLM:
-# ==============================================================================
-# TransformerLM.attention_type: see top
-# TransformerLM.d_attention_key: see top
-# TransformerLM.d_attention_value: see top
-TransformerLM.d_model = 1024
-TransformerLM.d_ff = 4096
-TransformerLM.dropout = 0.0
-TransformerLM.max_len = 12288  # 64 * 64 * 3
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 8
-# TransformerLM.n_layers: see top
-# TransformerLM.share_qk: see top
-TransformerLM.vocab_size = 256
-
-# Parameters for TransformerRevnetLM:
-# ==============================================================================
-# TransformerRevnetLM.attention_type: see top
-# TransformerRevnetLM.d_attention_key: see top
-# TransformerRevnetLM.d_attention_value: see top
-TransformerRevnetLM.d_model = 1024
-TransformerRevnetLM.d_ff = 4096
-TransformerRevnetLM.dropout = 0.0
-TransformerRevnetLM.max_len = 12288  # 64 * 64 * 3
-TransformerRevnetLM.mode = 'train'
-TransformerRevnetLM.n_heads = 8
-# TransformerRevnetLM.n_layers: see top
-TransformerRevnetLM.vocab_size = 256
-TransformerRevnetLM.n_chunks = 16
-TransformerRevnetLM.n_attention_chunks = 1
diff --git a/tensor2tensor/trax/configs/transformer_enwik8.gin b/tensor2tensor/trax/configs/transformer_enwik8.gin
deleted file mode 100644
index 364e41ae8..000000000
--- a/tensor2tensor/trax/configs/transformer_enwik8.gin
+++ /dev/null
@@ -1,69 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 1
-batch_fn.eval_batch_size = 8
-batch_fn.max_eval_length = 65536
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_enwik8_l65k'
-inputs.input_name = 'targets'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 2.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 100
-train.eval_steps = 8
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerLM
-train.optimizer = @trax.optimizers.Adafactor
-train.train_steps = 200000
-train.trainer_class = @MemoryEfficientTrainer
-
-# Parameters for MemoryEfficientCausalAttention:
-# ==============================================================================
-MemoryEfficientCausalAttention.dropout = 0.0
-MemoryEfficientCausalAttention.loop_stride = 256
-MemoryEfficientCausalAttention.share_qk = True  # matches TransformerLM.share_qk
-
-# Parameters for MergedHashedCausalAttention:
-# ==============================================================================
-MergedHashedCausalAttention.dropout = 0.0
-MergedHashedCausalAttention.n_bins = 512
-MergedHashedCausalAttention.bin_by_time = True
-MergedMultiHashedCausalAttention.one_rng = False
-
-# Parameters for MergedMultiHashedCausalAttention:
-# ==============================================================================
-MergedMultiHashedCausalAttention.dropout = 0.0
-MergedMultiHashedCausalAttention.n_bins = 512
-MergedMultiHashedCausalAttention.n_hashes = 2
-MergedMultiHashedCausalAttention.bin_by_time = False
-MergedMultiHashedCausalAttention.one_rng = False
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = @trax.layers.MergedMultiHashedCausalAttention
-TransformerLM.d_attention_key = 64
-TransformerLM.d_attention_value = 64
-TransformerLM.d_model = 512
-TransformerLM.d_ff = 2048
-TransformerLM.dropout = 0.0
-TransformerLM.max_len = 65536
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 4
-TransformerLM.n_layers = 3
-TransformerLM.share_qk = True
-TransformerLM.vocab_size = 258  # Includes pad token and unused EOS token
-
diff --git a/tensor2tensor/trax/configs/transformer_imagenet64.gin b/tensor2tensor/trax/configs/transformer_imagenet64.gin
deleted file mode 100644
index 5c9735bf3..000000000
--- a/tensor2tensor/trax/configs/transformer_imagenet64.gin
+++ /dev/null
@@ -1,80 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 1
-batch_fn.eval_batch_size = 8
-batch_fn.max_eval_length = 12288  # 64 * 64 * 3
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
-inputs.input_name = 'targets'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 2.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 500
-train.eval_steps = 64
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerLM
-train.optimizer = @trax.optimizers.Adafactor
-train.train_steps = 500000
-train.trainer_class = @MemoryEfficientTrainer
-
-# Parameters for MemoryEfficientCausalAttention:
-# ==============================================================================
-MemoryEfficientCausalAttention.dropout = 0.0
-MemoryEfficientCausalAttention.loop_stride = 512
-MemoryEfficientCausalAttention.share_qk = True  # matches TransformerLM.share_qk
-
-# Parameters for MergedHashedCausalAttention:
-# ==============================================================================
-MergedHashedCausalAttention.dropout = 0.0
-MergedHashedCausalAttention.n_bins = 16
-MergedHashedCausalAttention.bin_by_time = True
-MergedMultiHashedCausalAttention.one_rng = False
-
-# Parameters for MergedMultiHashedCausalAttention:
-# ==============================================================================
-MergedMultiHashedCausalAttention.dropout = 0.0
-MergedMultiHashedCausalAttention.n_bins = 64
-MergedMultiHashedCausalAttention.n_hashes = 4
-MergedMultiHashedCausalAttention.n_buckets_per_bin = 2
-MergedMultiHashedCausalAttention.bin_by_time = False
-MergedMultiHashedCausalAttention.one_rng = False
-MergedMultiHashedCausalAttention.drop_for_hash_rate = 0.1
-MergedMultiHashedCausalAttention.hard_k = 32
-
-# Parameters for MergedMultiHashedCausalAttentionV2:
-# ==============================================================================
-MergedMultiHashedCausalAttentionV2.dropout = 0.0
-MergedMultiHashedCausalAttentionV2.n_bins = 96
-MergedMultiHashedCausalAttentionV2.n_buckets = 256
-MergedMultiHashedCausalAttentionV2.n_hashes = 4
-MergedMultiHashedCausalAttentionV2.one_rng = True
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = @trax.layers.MergedMultiHashedCausalAttentionV2
-TransformerLM.d_attention_key = 64
-TransformerLM.d_attention_value = 64
-TransformerLM.d_model = 1024
-TransformerLM.d_ff = 2048
-TransformerLM.dropout = 0.0
-TransformerLM.max_len = 12288  # 64 * 64 * 3
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 4
-TransformerLM.n_layers = 3
-TransformerLM.share_qk = True
-TransformerLM.vocab_size = 256
-
diff --git a/tensor2tensor/trax/configs/transformer_revnet_enwik8.gin b/tensor2tensor/trax/configs/transformer_revnet_enwik8.gin
deleted file mode 100644
index 4e863d73d..000000000
--- a/tensor2tensor/trax/configs/transformer_revnet_enwik8.gin
+++ /dev/null
@@ -1,57 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 1
-batch_fn.eval_batch_size = 8
-batch_fn.max_eval_length = 65536
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_enwik8_l65k'
-inputs.input_name = 'targets'
-inputs.n_chunks = 16
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 2.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 200
-train.eval_steps = 8
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerRevnetLM
-train.optimizer = @trax.optimizers.Adafactor
-train.train_steps = 20000
-train.trainer_class = @MemoryEfficientTrainer
-
-# Parameters for MergedHashedCausalAttention:
-# ==============================================================================
-MergedHashedCausalAttention.dropout = 0.0
-MergedHashedCausalAttention.n_bins = 512
-MergedHashedCausalAttention.bin_by_time = True
-MergedHashedCausalAttention.one_rng = False
-
-# Parameters for TransformerRevnetLM:
-# ==============================================================================
-TransformerRevnetLM.d_model = 512
-TransformerRevnetLM.d_ff = 2048
-TransformerRevnetLM.d_attention_key = 64
-TransformerRevnetLM.d_attention_value = 64
-TransformerRevnetLM.dropout = 0.1
-TransformerRevnetLM.max_len = 65536
-TransformerRevnetLM.mode = 'train'
-TransformerRevnetLM.n_heads = 4
-TransformerRevnetLM.n_layers = 12
-TransformerRevnetLM.vocab_size = 258
-TransformerRevnetLM.n_chunks = 16
-TransformerRevnetLM.n_attention_chunks = 1
-TransformerRevnetLM.share_qk = True
-TransformerRevnetLM.attention_type = @trax.layers.MergedHashedCausalAttention
diff --git a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin b/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
deleted file mode 100644
index c5fa8e880..000000000
--- a/tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin
+++ /dev/null
@@ -1,73 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 2
-batch_fn.eval_batch_size = 16
-batch_fn.max_eval_length = 12288  # 64 * 64 * 3
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
-inputs.input_name = 'targets'
-inputs.n_chunks = 16
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 2.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 200
-train.eval_steps = 8
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerRevnetLM
-train.optimizer = @trax.optimizers.Adafactor
-train.train_steps = 500000
-train.trainer_class = @MemoryEfficientTrainer
-
-# Parameters for DotProductCausalAttention:
-# ==============================================================================
-DotProductCausalAttention.dropout = 0.0
-
-# Parameters for MemoryEfficientCausalAttention:
-# ==============================================================================
-MemoryEfficientCausalAttention.dropout = 0.0
-MemoryEfficientCausalAttention.loop_stride = 512
-
-# Parameters for MergedHashedCausalAttention:
-# ==============================================================================
-MergedHashedCausalAttention.dropout = 0.0
-MergedHashedCausalAttention.n_bins = 32
-MergedHashedCausalAttention.bin_by_time = True
-MergedHashedCausalAttention.one_rng = False
-
-# Parameters for MergedMultiHashedCausalAttention:
-# ==============================================================================
-MergedMultiHashedCausalAttention.dropout = 0.0
-MergedMultiHashedCausalAttention.n_bins = 64
-MergedMultiHashedCausalAttention.n_hashes = 4
-MergedMultiHashedCausalAttention.bin_by_time = False
-MergedHashedCausalAttention.one_rng = True
-
-# Parameters for TransformerRevnetLM:
-# ==============================================================================
-TransformerRevnetLM.d_model = 1024
-TransformerRevnetLM.d_ff = 2048
-TransformerRevnetLM.d_attention_key = 64
-TransformerRevnetLM.d_attention_value = 64
-TransformerRevnetLM.dropout = 0.0
-TransformerRevnetLM.max_len = 12288  # 64 * 64 * 3
-TransformerRevnetLM.mode = 'train'
-TransformerRevnetLM.n_heads = 4
-TransformerRevnetLM.n_layers = 4
-TransformerRevnetLM.vocab_size = 256
-TransformerRevnetLM.n_chunks = 16
-TransformerRevnetLM.n_attention_chunks = 1
-TransformerRevnetLM.attention_type = @trax.layers.MergedMultiHashedCausalAttention
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index f33b061f3..7a8d2d794 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -350,6 +350,20 @@ def forward(self, xs, params=(), state=(), **kwargs):
     return backend.numpy.concatenate(xs, self._axis), state
 
 
+class Split(base.Layer):
+  """Splits the input into sections along an axis."""
+
+  def __init__(self, n_sections=2, axis=-1):
+    super(Split, self).__init__(n_outputs=n_sections)
+    self._n_sections = n_sections
+    self._axis = axis
+
+  def forward(self, inputs, params=(), state=(), **kwargs):
+    del params, kwargs
+    res = tuple(backend.numpy.split(inputs, self._n_sections, self._axis))
+    return res, state
+
+
 class Parallel(base.Layer):
   """Combinator that applies a list of layers in parallel to its inputs.
 
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/transformer_revnet.py
index 430635125..abe968762 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/transformer_revnet.py
@@ -117,20 +117,6 @@ def FeedForward(d_model, d_ff, dropout, mode):
   ]
 
 
-class Split(tl.Layer):
-  """Splits the input into sections along an axis."""
-
-  def __init__(self, n_sections=2, axis=-1):
-    super(Split, self).__init__(n_outputs=n_sections)
-    self._n_sections = n_sections
-    self._axis = axis
-
-  def forward(self, inputs, params=(), state=(), **kwargs):
-    del params, kwargs
-    res = tuple(backend.numpy.split(inputs, self._n_sections, self._axis))
-    return res, state
-
-
 class SplitForOutput(tl.ReversibleLayer):
   """Splits activations into sections (for use right before the output layer).
 
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 7388e8786..300a381bb 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -218,6 +218,7 @@ def TransformerLM(vocab_size,
                   dropout=0.1,
                   share_qk=False,
                   max_len=2048,
+                  n_chunks=0,
                   mode='train'):
   """Returns a Transformer language model.
 
@@ -238,18 +239,27 @@ def TransformerLM(vocab_size,
     dropout: float: dropout rate (how much to drop out)
     share_qk: bool, whether to share queries and keys in decoder attention
     max_len: int: maximum symbol length for positional encoding
+    n_chunks: int: number of chunks (must match input pipeline)
     mode: str: 'train', 'eval' or 'predict', predict mode is for fast inference
 
   Returns:
     A Transformer language model as a layer that maps from a tensor of tokens
     to activations over a vocab set.
   """
+  if n_chunks == 0:
+    concatenate_chunks = split_chunks = []
+  else:
+    concatenate_chunks = tl.Concatenate(n_items=n_chunks)
+    split_chunks = tl.Split(n_sections=n_chunks, axis=-2)
+
   embedder = [
       tl.Embedding(d_model, vocab_size),
       tl.Dropout(rate=dropout, name='embedding', mode=mode),
       tl.PositionalEncoding(max_len=max_len, mode=mode),
   ]
-  return tl.Model(                  # tokens
+
+  return tl.Model(                  # tokens (or chunked tuple of tokens)
+      concatenate_chunks,           # tokens
       tl.ShiftRight(mode=mode),     # toks
       embedder,                     # vecs
       [DecoderBlock(  # pylint: disable=g-complex-comprehension
@@ -259,6 +269,7 @@ def TransformerLM(vocab_size,
       tl.LayerNorm(),               # vecs
       tl.Dense(vocab_size),         # vecs
       tl.LogSoftmax(),              # vecs
+      split_chunks,                 # vecs (or chunked tuple of vecs)
   )
 
 

From 08a851cbf3860c4413bf435f43247cc84fee9fe6 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 30 Sep 2019 12:35:30 -0700
Subject: [PATCH 2505/2720] Rename classes used in reformer

PiperOrigin-RevId: 272041361
---
 .../trax/configs/reformer_enwik8.gin          |  81 ++--
 .../trax/configs/reformer_imagenet64.gin      |  81 ++--
 .../trax/configs/transformer_copy.gin         |  24 +-
 tensor2tensor/trax/layers/__init__.py         |  10 +-
 tensor2tensor/trax/layers/attention.py        | 389 ++----------------
 tensor2tensor/trax/models/__init__.py         |   4 +-
 .../{transformer_revnet.py => reformer.py}    |  28 +-
 ...former_revnet_test.py => reformer_test.py} |  10 +-
 8 files changed, 158 insertions(+), 469 deletions(-)
 rename tensor2tensor/trax/models/research/{transformer_revnet.py => reformer.py} (97%)
 rename tensor2tensor/trax/models/research/{transformer_revnet_test.py => reformer_test.py} (81%)

diff --git a/tensor2tensor/trax/configs/reformer_enwik8.gin b/tensor2tensor/trax/configs/reformer_enwik8.gin
index d90119520..17b3e9883 100644
--- a/tensor2tensor/trax/configs/reformer_enwik8.gin
+++ b/tensor2tensor/trax/configs/reformer_enwik8.gin
@@ -5,9 +5,9 @@ import tensor2tensor.trax.trax
 
 # Parameters that will vary between experiments:
 # ==============================================================================
-train.model = @trax.models.TransformerRevnetLM
-attn_type = @MergedMultiHashedCausalAttentionV2
-share_qk = True  # Required when using hashed attention
+train.model = @trax.models.ReformerLM
+attn_type = @TimeBinCausalAttention
+share_qk = True  # Required when using LSHCausalAttention
 attn_kv = 64
 n_layers = 3
 dropout = 0.1
@@ -15,14 +15,14 @@ dropout = 0.1
 # MemoryEfficientCausalAttention: full attention
 # (no hparams to vary between experiments)
 
-# MergedHashedCausalAttention: use for timebin
-MergedHashedCausalAttention.n_bins = 512
+# TimeBinCausalAttention: attend to nearby items
+TimeBinCausalAttention.n_bins = 512
 
-# MergedMultiHashedCausalAttentionV2: locality-sensitive hashing (LSH) attention
-MergedMultiHashedCausalAttentionV2.n_bins = 512
-MergedMultiHashedCausalAttentionV2.n_buckets = 1024  # Always 2 * n_bins
-MergedMultiHashedCausalAttentionV2.n_hashes = 2
-MergedMultiHashedCausalAttentionV2.drop_for_hash_rate = 0.0
+# LSHCausalAttention: locality-sensitive hashing (LSH) attention
+LSHCausalAttention.n_bins = 512
+LSHCausalAttention.n_buckets = 1024  # Always 2 * n_bins
+LSHCausalAttention.n_hashes = 2
+LSHCausalAttention.drop_for_hash_rate = 0.0
 
 # Parameters for batch_fn:
 # ==============================================================================
@@ -62,25 +62,24 @@ MemoryEfficientCausalAttention.dropout = 0.0
 MemoryEfficientCausalAttention.loop_stride = 256
 MemoryEfficientCausalAttention.share_qk = %share_qk
 
-# Parameters for MergedHashedCausalAttention:
+# Parameters for TimeBinCausalAttention:
 # ==============================================================================
-MergedHashedCausalAttention.dropout = 0.0
-# MergedHashedCausalAttention.n_bins: see top
-MergedHashedCausalAttention.bin_by_time = True
-MergedHashedCausalAttention.one_rng = False
+TimeBinCausalAttention.dropout = 0.0
+# TimeBinCausalAttention.n_bins: see top
+TimeBinCausalAttention.share_qk = %share_qk
 
-# Parameters for MergedMultiHashedCausalAttentionV2:
+# Parameters for LSHCausalAttention:
 # ==============================================================================
-MergedMultiHashedCausalAttentionV2.allow_duplicate_attention = False
-MergedMultiHashedCausalAttentionV2.attend_across_buckets = False
-MergedMultiHashedCausalAttentionV2.rehash_each_round = True
-# MergedMultiHashedCausalAttentionV2.n_bins: see top
-# MergedMultiHashedCausalAttentionV2.n_buckets: see top
-# MergedMultiHashedCausalAttentionV2.n_hashes: see top
-MergedMultiHashedCausalAttentionV2.one_rng = False
-MergedMultiHashedCausalAttentionV2.hard_k = 0
-MergedMultiHashedCausalAttentionV2.dropout = 0.0
-# MergedMultiHashedCausalAttentionV2.drop_for_hash_rate: see top
+LSHCausalAttention.allow_duplicate_attention = False
+LSHCausalAttention.attend_across_buckets = False
+LSHCausalAttention.rehash_each_round = True
+# LSHCausalAttention.n_bins: see top
+# LSHCausalAttention.n_buckets: see top
+# LSHCausalAttention.n_hashes: see top
+LSHCausalAttention.one_rng = False
+LSHCausalAttention.hard_k = 0
+LSHCausalAttention.dropout = 0.0
+# LSHCausalAttention.drop_for_hash_rate: see top
 
 # Parameters for TransformerLM:
 # ==============================================================================
@@ -98,19 +97,19 @@ TransformerLM.n_chunks = 16
 TransformerLM.share_qk = %share_qk
 TransformerLM.vocab_size = 258  # Includes pad token and unused EOS token
 
-# Parameters for TransformerRevnetLM:
+# Parameters for ReformerLM:
 # ==============================================================================
-TransformerRevnetLM.attention_type = %attn_type
-TransformerRevnetLM.d_attention_key = %attn_kv
-TransformerRevnetLM.d_attention_value = %attn_kv
-TransformerRevnetLM.d_model = 1024
-TransformerRevnetLM.d_ff = 4096
-TransformerRevnetLM.dropout = %dropout
-TransformerRevnetLM.max_len = 65536
-TransformerRevnetLM.mode = 'train'
-TransformerRevnetLM.n_heads = 8
-TransformerRevnetLM.n_layers = %n_layers
-TransformerRevnetLM.vocab_size = 258  # Includes pad token and unused EOS token
-TransformerRevnetLM.n_chunks = 16
-TransformerRevnetLM.n_attention_chunks = 1
-TransformerRevnetLM.share_qk = %share_qk
+ReformerLM.attention_type = %attn_type
+ReformerLM.d_attention_key = %attn_kv
+ReformerLM.d_attention_value = %attn_kv
+ReformerLM.d_model = 1024
+ReformerLM.d_ff = 4096
+ReformerLM.dropout = %dropout
+ReformerLM.max_len = 65536
+ReformerLM.mode = 'train'
+ReformerLM.n_heads = 8
+ReformerLM.n_layers = %n_layers
+ReformerLM.vocab_size = 258  # Includes pad token and unused EOS token
+ReformerLM.n_chunks = 16
+ReformerLM.n_attention_chunks = 1
+ReformerLM.share_qk = %share_qk
diff --git a/tensor2tensor/trax/configs/reformer_imagenet64.gin b/tensor2tensor/trax/configs/reformer_imagenet64.gin
index 006f57660..9d9ae11c2 100644
--- a/tensor2tensor/trax/configs/reformer_imagenet64.gin
+++ b/tensor2tensor/trax/configs/reformer_imagenet64.gin
@@ -5,23 +5,23 @@ import tensor2tensor.trax.trax
 
 # Parameters that will vary between experiments:
 # ==============================================================================
-train.model = @trax.models.TransformerRevnetLM
-attn_type = @MergedMultiHashedCausalAttentionV2
-share_qk = True  # Required when using hashed attention
+train.model = @trax.models.ReformerLM
+attn_type = @TimeBinCausalAttention
+share_qk = True  # Required when using LSHCausalAttention
 attn_kv = 64
 n_layers = 3
 
 # MemoryEfficientCausalAttention: full attention
 # (no hparams to vary between experiments)
 
-# MergedHashedCausalAttention: use for timebin
-MergedHashedCausalAttention.n_bins = 64
+# TimeBinCausalAttention: attend to nearby items
+TimeBinCausalAttention.n_bins = 64
 
-# MergedMultiHashedCausalAttentionV2: locality-sensitive hashing (LSH) attention
-MergedMultiHashedCausalAttentionV2.n_bins = 96
-MergedMultiHashedCausalAttentionV2.n_buckets = 192  # Always 2 * n_bins
-MergedMultiHashedCausalAttentionV2.n_hashes = 2
-MergedMultiHashedCausalAttentionV2.drop_for_hash_rate = 0.0
+# LSHCausalAttention: locality-sensitive hashing (LSH) attention
+LSHCausalAttention.n_bins = 96
+LSHCausalAttention.n_buckets = 192  # Always 2 * n_bins
+LSHCausalAttention.n_hashes = 2
+LSHCausalAttention.drop_for_hash_rate = 0.0
 
 # Parameters for batch_fn:
 # ==============================================================================
@@ -61,25 +61,24 @@ MemoryEfficientCausalAttention.dropout = 0.0
 MemoryEfficientCausalAttention.loop_stride = 512
 MemoryEfficientCausalAttention.share_qk = %share_qk
 
-# Parameters for MergedHashedCausalAttention:
+# Parameters for TimeBinCausalAttention:
 # ==============================================================================
-MergedHashedCausalAttention.dropout = 0.0
-# MergedHashedCausalAttention.n_bins: see top
-MergedHashedCausalAttention.bin_by_time = True
-MergedHashedCausalAttention.one_rng = False
+TimeBinCausalAttention.dropout = 0.0
+# TimeBinCausalAttention.n_bins: see top
+TimeBinCausalAttention.share_qk = %share_qk
 
-# Parameters for MergedMultiHashedCausalAttentionV2:
+# Parameters for LSHCausalAttention:
 # ==============================================================================
-MergedMultiHashedCausalAttentionV2.allow_duplicate_attention = False
-MergedMultiHashedCausalAttentionV2.attend_across_buckets = False
-MergedMultiHashedCausalAttentionV2.rehash_each_round = True
-# MergedMultiHashedCausalAttentionV2.n_bins: see top
-# MergedMultiHashedCausalAttentionV2.n_buckets: see top
-# MergedMultiHashedCausalAttentionV2.n_hashes: see top
-MergedMultiHashedCausalAttentionV2.one_rng = False
-MergedMultiHashedCausalAttentionV2.hard_k = 0
-MergedMultiHashedCausalAttentionV2.dropout = 0.0
-# MergedMultiHashedCausalAttentionV2.drop_for_hash_rate: see top
+LSHCausalAttention.allow_duplicate_attention = False
+LSHCausalAttention.attend_across_buckets = False
+LSHCausalAttention.rehash_each_round = True
+# LSHCausalAttention.n_bins: see top
+# LSHCausalAttention.n_buckets: see top
+# LSHCausalAttention.n_hashes: see top
+LSHCausalAttention.one_rng = False
+LSHCausalAttention.hard_k = 0
+LSHCausalAttention.dropout = 0.0
+# LSHCausalAttention.drop_for_hash_rate: see top
 
 # Parameters for TransformerLM:
 # ==============================================================================
@@ -97,19 +96,19 @@ TransformerLM.n_chunks = 16
 TransformerLM.share_qk = %share_qk
 TransformerLM.vocab_size = 256
 
-# Parameters for TransformerRevnetLM:
+# Parameters for ReformerLM:
 # ==============================================================================
-TransformerRevnetLM.attention_type = %attn_type
-TransformerRevnetLM.d_attention_key = %attn_kv
-TransformerRevnetLM.d_attention_value = %attn_kv
-TransformerRevnetLM.d_model = 1024
-TransformerRevnetLM.d_ff = 4096
-TransformerRevnetLM.dropout = 0.0
-TransformerRevnetLM.max_len = 12288  # 64 * 64 * 3
-TransformerRevnetLM.mode = 'train'
-TransformerRevnetLM.n_heads = 8
-TransformerRevnetLM.n_layers = %n_layers
-TransformerRevnetLM.vocab_size = 256
-TransformerRevnetLM.n_chunks = 16
-TransformerRevnetLM.n_attention_chunks = 1
-TransformerRevnetLM.share_qk = %share_qk
+ReformerLM.attention_type = %attn_type
+ReformerLM.d_attention_key = %attn_kv
+ReformerLM.d_attention_value = %attn_kv
+ReformerLM.d_model = 1024
+ReformerLM.d_ff = 4096
+ReformerLM.dropout = 0.0
+ReformerLM.max_len = 12288  # 64 * 64 * 3
+ReformerLM.mode = 'train'
+ReformerLM.n_heads = 8
+ReformerLM.n_layers = %n_layers
+ReformerLM.vocab_size = 256
+ReformerLM.n_chunks = 16
+ReformerLM.n_attention_chunks = 1
+ReformerLM.share_qk = %share_qk
diff --git a/tensor2tensor/trax/configs/transformer_copy.gin b/tensor2tensor/trax/configs/transformer_copy.gin
index cc7d0f70d..2cc2d29fc 100644
--- a/tensor2tensor/trax/configs/transformer_copy.gin
+++ b/tensor2tensor/trax/configs/transformer_copy.gin
@@ -47,22 +47,22 @@ train.has_weights = True
 MemoryEfficientCausalAttention.dropout = 0.0
 MemoryEfficientCausalAttention.loop_stride = 512
 
-# Parameters for MergedMultiHashedCausalAttentionV2:
+# Parameters for LSHCausalAttention:
 # ==============================================================================
-MergedMultiHashedCausalAttentionV2.allow_duplicate_attention = False
-MergedMultiHashedCausalAttentionV2.attend_across_buckets = False
-MergedMultiHashedCausalAttentionV2.rehash_each_round = True
-MergedMultiHashedCausalAttentionV2.n_bins = 32
-MergedMultiHashedCausalAttentionV2.n_buckets = 64
-MergedMultiHashedCausalAttentionV2.n_hashes = 4
-MergedMultiHashedCausalAttentionV2.one_rng = False
-MergedMultiHashedCausalAttentionV2.hard_k = 0
-MergedMultiHashedCausalAttentionV2.dropout = 0.0
-MergedMultiHashedCausalAttentionV2.drop_for_hash_rate = 0.0
+LSHCausalAttention.allow_duplicate_attention = False
+LSHCausalAttention.attend_across_buckets = False
+LSHCausalAttention.rehash_each_round = True
+LSHCausalAttention.n_bins = 32
+LSHCausalAttention.n_buckets = 64
+LSHCausalAttention.n_hashes = 4
+LSHCausalAttention.one_rng = False
+LSHCausalAttention.hard_k = 0
+LSHCausalAttention.dropout = 0.0
+LSHCausalAttention.drop_for_hash_rate = 0.0
 
 # Parameters for TransformerLM:
 # ==============================================================================
-TransformerLM.attention_type = @trax.layers.MergedMultiHashedCausalAttentionV2
+TransformerLM.attention_type = @trax.layers.LSHCausalAttention
 TransformerLM.d_attention_key = 64
 TransformerLM.d_attention_value = 64
 TransformerLM.d_model = 256
diff --git a/tensor2tensor/trax/layers/__init__.py b/tensor2tensor/trax/layers/__init__.py
index 6e7bbc124..7032b5f79 100644
--- a/tensor2tensor/trax/layers/__init__.py
+++ b/tensor2tensor/trax/layers/__init__.py
@@ -55,9 +55,7 @@ def layer_configure(*args, **kwargs):
     DotProductCausalAttention, blacklist=["mode"])
 MemoryEfficientCausalAttention = layer_configure(
     MemoryEfficientCausalAttention, blacklist=["mode"])
-MergedHashedCausalAttention = layer_configure(
-    MergedHashedCausalAttention, blacklist=["mode"])
-MergedMultiHashedCausalAttention = layer_configure(
-    MergedMultiHashedCausalAttention, blacklist=["mode"])
-MergedMultiHashedCausalAttentionV2 = layer_configure(
-    MergedMultiHashedCausalAttentionV2, blacklist=["mode"])
+TimeBinCausalAttention = layer_configure(
+    TimeBinCausalAttention, blacklist=["mode"])
+LSHCausalAttention = layer_configure(
+    LSHCausalAttention, blacklist=["mode"])
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 7a0e31359..0c6156bf5 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -654,20 +654,15 @@ def body_fun(vals):  # pylint: disable=invalid-name
       return final_vals[1], final_vals[2:]
 
 
-class MergedHashedCausalAttention(BaseCausalAttention):
-  """Hash-based causal attention."""
+class TimeBinCausalAttention(BaseCausalAttention):
+  """Causal attention where only nearby chunks of items attend to each other."""
+  # TODO(kitaev): rewrite this class to use reshapes only, rather than sorting.
 
-  def __init__(self, dropout, mode, n_bins=64,
-               bin_by_time=False, one_rng=False):
+  def __init__(self, dropout, mode, n_bins=64, share_qk=False):
     del dropout, mode
-    super(MergedHashedCausalAttention, self).__init__()
+    super(TimeBinCausalAttention, self).__init__()
     self.n_bins = n_bins
-    self.bin_by_time = bin_by_time
-    seed = random.randint(0, 2**31 - 1)
-    self._one_rng = one_rng
-    self._prng = None
-    if one_rng:
-      self._prng = backend.random.get_prng(seed)
+    self._share_qk = share_qk
 
   def forward(self, inputs, params=(), state=(), **kwargs):
     del params
@@ -698,48 +693,23 @@ def make_unit_length(self, x, epsilon=1e-6):
     norm_inputs = x / np.sqrt(variance + epsilon)
     return norm_inputs
 
-  def hash_vectors(self, vecs, rng):
-    if self.bin_by_time:
-      # Instead of hashing, put chunks of consecutive items in the same bin.
-      # This exists as a sanity check for the other parts of this class.
-      return self.bin_vectors_by_time(vecs)
-
-    # See https://arxiv.org/pdf/1509.02897.pdf
-    # It's not clear whether sampling a different random rotation for each head
-    # and batch element matters here, but see MergedMultiHashedCausalAttention.
-    assert self.n_bins % 2 == 0
-    rot_rng = rng
-    if self._one_rng:
-      rot_rng = jax.lax.tie_in(vecs, self._prng)
-    random_rotation = jax.random.normal(
-        rot_rng,
-        (vecs.shape[0], vecs.shape[-1], self.n_bins//2)).astype('float32')
-
-    # TODO(kitaev): making the vectors unit-length here is probably redundant.
-    vecs = self.make_unit_length(vecs)
-    rotated_vecs = np.matmul(vecs, random_rotation)
-    rotated_vecs = self.make_unit_length(rotated_vecs)
-    rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)
-    bins = np.argmax(rotated_vecs, axis=-1)
-    return bins
-
-  def forward_and_backward(self, inputs, ct, rng=None, **kwargs):
+  def forward_and_backward(self, inputs, ct, **kwargs):
     del kwargs
     # We use the same vector as both a query and a key. For now we haven't
     # adjusted any of the surrounding code, so we still get a separate "key"
     # input that we ignore.
-    qk, ignored_k, v = inputs
-    seqlen = qk.shape[-2]
-    # qk/v are n_batch*n_heads, seqlen, d_head
+    q, k, v = inputs
+    seqlen = q.shape[-2]
+    # q/k/v are n_batch*n_heads, seqlen, d_head
 
     # bins are n_batch*n_heads, seqlen
     # They specify which hash bucket the query/key/value vectors fall in.
-    bins = self.hash_vectors(qk, rng=rng)
+    bins = self.bin_vectors_by_time(q)
 
     # joint_t is n_batch*n_heads, seqlen
-    joint_t = jax.lax.tie_in(qk, np.arange(seqlen))
+    joint_t = jax.lax.tie_in(q, np.arange(seqlen))
     joint_t = np.reshape(joint_t, (1, seqlen))
-    joint_t = np.broadcast_to(joint_t, qk.shape[:-1])
+    joint_t = np.broadcast_to(joint_t, q.shape[:-1])
 
     assert int((self.n_bins + 1) * seqlen) < 2 ** 31, (
         'Potential 32-bit integer overflow; please double-check the code.')
@@ -760,27 +730,27 @@ def unchunk_vectors(x):  # pylint: disable=invalid-name
     _, sjoint_t = jax.lax.sort_key_val(
         joint_bins_and_t, joint_t, dimension=-1)
 
-    sqk = np.take_along_axis(qk, sjoint_t[:, :, None], axis=-2)
+    sq = np.take_along_axis(q, sjoint_t[:, :, None], axis=-2)
+    if self._share_qk:
+      sk = sq
+    else:
+      sk = np.take_along_axis(k, sjoint_t[:, :, None], axis=-2)
     sv = np.take_along_axis(v, sjoint_t[:, :, None], axis=-2)
 
     if ct is not None:
       so_ct = np.take_along_axis(ct, sjoint_t[:, :, None], axis=-2)
 
     @jax.jit
-    def binned_attn(sqk, sv):  # pylint: disable=invalid-name
+    def binned_attn(sq, sk, sv):  # pylint: disable=invalid-name
       """Performs attention on sorted queries/keys/values."""
       # Split off a "bin" axis so that attention only occurs whithin chunks.
       bq_t = bkv_t = chunk_scalars(sjoint_t)
-      bqk = chunk_vectors(sqk)
+      bq = chunk_vectors(sq)
+      bk = chunk_vectors(sk)
+      if self._share_qk:
+        bk = self.make_unit_length(bk)
       bv = chunk_vectors(sv)
 
-      # Hashing operates on unit-length vectors. Unnormalized query vectors are
-      # fine because they effectively provide a learnable temperature for the
-      # attention softmax, but normalizing keys is needed so that similarity for
-      # the purposes of attention correctly corresponds to hash locality.
-      bq = bqk
-      bk = self.make_unit_length(bqk)
-
       # Allow each chunk to attend within itself, and also one chunk back. Chunk
       # boundaries might occur in the middle of a sequence of items from the
       # same bin, so this increases the chances of attending to relevant items.
@@ -802,9 +772,10 @@ def binned_attn(sqk, sv):  # pylint: disable=invalid-name
       dots = dots - 1e9 * mask
 
       # Mask out attention to self except when no other targets are available.
-      self_mask = jax.lax.broadcasted_eye(dots.dtype, dots.shape, (2, 3))
-      self_mask = jax.lax.tie_in(dots, self_mask)
-      dots = dots - 32 * self_mask
+      if self._share_qk:
+        self_mask = jax.lax.broadcasted_eye(dots.dtype, dots.shape, (2, 3))
+        self_mask = jax.lax.tie_in(dots, self_mask)
+        dots = dots - 32 * self_mask
 
       # Softmax.
       dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))
@@ -814,13 +785,13 @@ def binned_attn(sqk, sv):  # pylint: disable=invalid-name
       return so
 
     @jax.jit
-    def binned_attn_vjp(sqk, sv, so_ct):  # pylint: disable=invalid-name
-      so, vjpfun = jax.vjp(binned_attn, sqk, sv)
+    def binned_attn_vjp(sq, sk, sv, so_ct):  # pylint: disable=invalid-name
+      so, vjpfun = jax.vjp(binned_attn, sq, sk, sv)
       sqkv_ct = vjpfun(so_ct)
       return so, sqkv_ct
 
     if ct is None:
-      so = binned_attn(sqk, sv)
+      so = binned_attn(sq, sk, sv)
       _, undo_sort = jax.lax.sort_key_val(sjoint_t, joint_t, dimension=-1)
       out = np.take_along_axis(so, undo_sort[:, :, None], axis=-2)
       return out, None
@@ -829,302 +800,24 @@ def binned_attn_vjp(sqk, sv, so_ct):  # pylint: disable=invalid-name
       # slower than writing our own. The main reason is that the backward pass
       # of gather is in general a scatter operation, but we know we're dealing
       # with permutations so we use gather for the backward pass too.
-      so, (sqk_ct, sv_ct) = binned_attn_vjp(sqk, sv, so_ct)
+      so, (sq_ct, sk_ct, sv_ct) = binned_attn_vjp(sq, sk, sv, so_ct)
 
       _, undo_sort = jax.lax.sort_key_val(sjoint_t, joint_t, dimension=-1)
       out = np.take_along_axis(so, undo_sort[:, :, None], axis=-2)
 
-      qk_ct = np.take_along_axis(sqk_ct, undo_sort[:, :, None], axis=-2)
+      if self._share_qk:
+        q_ct = np.take_along_axis(sq_ct + sk_ct, undo_sort[:, :, None], axis=-2)
+        k_ct = np.zeros_like(k)
+      else:
+        q_ct = np.take_along_axis(sq_ct, undo_sort[:, :, None], axis=-2)
+        k_ct = np.take_along_axis(sk_ct, undo_sort[:, :, None], axis=-2)
       v_ct = np.take_along_axis(sv_ct, undo_sort[:, :, None], axis=-2)
 
-      return out, (qk_ct, np.zeros_like(ignored_k), v_ct)
-
-
-class MergedMultiHashedCausalAttention(BaseCausalAttention):
-  """Hash-based causal attention, with multiple hashes."""
-
-  def __init__(self, dropout, mode, n_bins=64, n_hashes=1,
-               n_buckets_per_bin=1, bin_by_time=False, one_rng=False,
-               allow_duplicate_attention=False,
-               drop_for_hash_rate=0.0, hard_k=0):
-    del dropout
-    self._mode = mode
-    super(MergedMultiHashedCausalAttention, self).__init__()
-    self.n_bins = n_bins
-    self.n_hashes = n_hashes
-    self.n_buckets_per_bin = n_buckets_per_bin
-    self.bin_by_time = bin_by_time
-    seed = random.randint(0, 2**31 - 1)
-    self._one_rng = one_rng
-    self._drop_for_hash_rate = drop_for_hash_rate
-    self._hard_k = hard_k
-    self._prng = None
-    if one_rng:
-      self._prng = backend.random.get_prng(seed)
-    self._allow_duplicate_attention = allow_duplicate_attention
-
-  def bin_vectors_by_time(self, vecs):
-    seqlen = vecs.shape[-2]
-    assert seqlen % self.n_bins == 0
-    bin_size = int(seqlen // self.n_bins)
-
-    bins = np.arange(seqlen, dtype=np.int32) // bin_size
-    bins = jax.lax.tie_in(vecs, bins)
-    bins = bins[None, :]
-    bins = np.broadcast_to(bins, vecs.shape[:-1])
-    return bins
-
-  def make_unit_length(self, x, epsilon=1e-6):
-    variance = np.mean(x**2, axis=-1, keepdims=True)
-    norm_inputs = x / np.sqrt(variance + epsilon)
-    return norm_inputs
-
-  def drop_for_hash(self, x, rng):
-    rate = self._drop_for_hash_rate
-    if self._mode == 'train' and rate > 0.0:
-      keep = backend.random.bernoulli(rng, 1.0 - rate, x.shape)
-      return np.where(keep, x / (1.0 - rate), np.zeros_like(x))
-    return x
-
-  def hash_vectors(self, vecs, rng):
-    if self.bin_by_time:
-      # Instead of hashing, put chunks of consecutive items in the same bin.
-      # This exists as a sanity check for the other parts of this class.
-      return self.bin_vectors_by_time(vecs)
-
-    # See https://arxiv.org/pdf/1509.02897.pdf
-    # We sample a different random rotation for each batch element, head, and
-    # (crucially) each round of hashing. All of these are part of dimension 0
-    # of vecs. Applying multiple hashes to the same input is important because
-    # it increases the probability of being in the same bin as relevant items.
-    n_buckets = self.n_buckets_per_bin * self.n_bins
-    assert n_buckets % 2 == 0
-    rot_rng = rng
-    if self._one_rng:
-      rot_rng = jax.lax.tie_in(vecs, self._prng)
-    random_rotation = jax.random.normal(
-        rot_rng,
-        (vecs.shape[0], vecs.shape[-1], n_buckets // 2)).astype('float32')
-
-    # TODO(kitaev): making the vectors unit-length here is probably redundant.
-    # vecs = self.make_unit_length(vecs)
-    rng, subrng = backend.random.split(rng)
-    vecs = self.drop_for_hash(vecs, subrng)
-    rotated_vecs = np.matmul(vecs, random_rotation)
-    rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)
-    bins = np.argmax(rotated_vecs, axis=-1)
-    return bins
-
-  def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
-    del params, kwargs
-    # We use the same vector as both a query and a key. For now we haven't
-    # adjusted any of the surrounding code, so we still get a separate "key"
-    # input that we ignore.
-    qk, _, v = inputs
-    seqlen = qk.shape[-2]
-
-    # qk/v are n_hashes*n_batch*n_heads, seqlen, d_head
-    # TODO(kitaev): is it faster to fuse this tiling into gather/scatter ops?
-    qk = np.tile(qk, (self.n_hashes, 1, 1))
-    v = np.tile(v, (self.n_hashes, 1, 1))
-
-    # bins are n_hashes*n_batch*n_heads, seqlen
-    # They specify which hash bucket the query/key/value vectors fall in.
-    bins = self.hash_vectors(qk, rng=rng)
-
-    # joint_t is n_hashes*n_batch*n_heads, seqlen
-    joint_t = jax.lax.tie_in(qk, np.arange(seqlen))
-    joint_t = np.reshape(joint_t, (1, seqlen))
-    joint_t = np.broadcast_to(joint_t, qk.shape[:-1])
-
-    assert int((self.n_buckets_per_bin * self.n_bins + 1) * seqlen) < 2 ** 31, (
-        'Potential 32-bit integer overflow; please double-check the code.')
-    joint_bins_and_t = seqlen * bins + joint_t
-
-    def chunk_scalars(x):  # pylint: disable=invalid-name
-      return np.reshape(x, (x.shape[0], self.n_bins, -1))
-
-    def chunk_vectors(x):  # pylint: disable=invalid-name
-      return np.reshape(
-          x, (x.shape[0], self.n_bins, -1, x.shape[-1]))
-
-    def unchunk_vectors(x):  # pylint: disable=invalid-name
-      return np.reshape(x, (x.shape[0], -1, x.shape[-1]))
-
-    # Sort everything by bin number, with a secondary sort by time
-    # (variables starting with "s" are sorted)
-    _, sjoint_t = jax.lax.sort_key_val(
-        joint_bins_and_t, joint_t, dimension=-1)
-    _, undo_sort = jax.lax.sort_key_val(sjoint_t, joint_t, dimension=-1)
-    # TODO(kitaev): why does jax flag integer indices as differentiable?
-    # If we don't call stop_gradient here, custom gradients below won't work
-    # because the primitive functions close over "differentiable" variables.
-    sjoint_t = jax.lax.stop_gradient(sjoint_t)
-    undo_sort = jax.lax.stop_gradient(undo_sort)
-
-    # The backward pass of gather is in general a scatter operation, but we know
-    # we're dealing with permutations so we use gather for the backward pass
-    # too. This custom gradient should be about 2x faster than having jax infer
-    # one that uses scatter ops instead.
-    def permute_impl(vecs):
-      assert len(vecs.shape) == 3
-      return np.take_along_axis(vecs, sjoint_t[:, :, None], axis=-2)
-
-    def unpermute_impl(vecs):
-      assert len(vecs.shape) == 3
-      return np.take_along_axis(vecs, undo_sort[:, :, None], axis=-2)
-
-    @jax.custom_transforms
-    def permute(vecs):
-      return permute_impl(vecs)
-
-    def permute_vjp(vecs):
-      out_vecs = permute_impl(vecs)
-      def vjpfun(grad):
-        return (unpermute_impl(grad),)
-      return out_vecs, vjpfun
-
-    @jax.custom_transforms
-    def unpermute(vecs):
-      return unpermute_impl(vecs)
-
-    def unpermute_vjp(vecs):
-      out_vecs = unpermute_impl(vecs)
-      def vjpfun(grad):
-        return (permute_impl(grad),)
-      return out_vecs, vjpfun
-
-    jax.defvjp_all(permute, permute_vjp)
-    jax.defvjp_all(unpermute, unpermute_vjp)
-
-    sqk = permute(qk)
-    sv = permute(v)
-
-    # Split off a "bin" axis so that attention only occurs within chunks.
-    bq_t = bkv_t = chunk_scalars(sjoint_t)
-    bqk = chunk_vectors(sqk)
-    bv = chunk_vectors(sv)
-
-    # Hashing operates on unit-length vectors. Unnormalized query vectors are
-    # fine because they effectively provide a learnable temperature for the
-    # attention softmax, but normalizing keys is needed so that similarity for
-    # the purposes of attention correctly corresponds to hash locality.
-    bq = bqk
-    bk = self.make_unit_length(bqk)
-
-    # Allow each chunk to attend within itself, and also one chunk back. Chunk
-    # boundaries might occur in the middle of a sequence of items from the
-    # same bin, so this increases the chances of attending to relevant items.
-    # TODO(kitaev): benchmark whether XLA pad operation is noticeably faster.
-    bk_extra = np.concatenate([bk[:, -1:, :, :], bk[:, :-1, :, :]], axis=1)
-    bk = np.concatenate([bk, bk_extra], axis=2)
-    bv_extra = np.concatenate([bv[:, -1:, :, :], bv[:, :-1, :, :]], axis=1)
-    bv = np.concatenate([bv, bv_extra], axis=2)
-    bkv_t_extra = np.concatenate([bkv_t[:, -1:, :], bkv_t[:, :-1, :]], axis=1)
-    bkv_t = np.concatenate([bkv_t, bkv_t_extra], axis=2)
-
-    # Dot-product attention.
-    dots = np.matmul(bq, np.swapaxes(bk, -1, -2)) / np.sqrt(bq.shape[-1])
-
-    # Causal masking
-    mask = jax.lax.convert_element_type(
-        jax.lax.lt(bq_t[:, :, :, None], bkv_t[:, :, None, :]),
-        np.float32)
-    dots = dots - 1e9 * mask
-
-    # Mask out attention to self except when no other targets are available.
-    self_mask = jax.lax.broadcasted_eye(dots.dtype, dots.shape, (2, 3))
-    self_mask = jax.lax.tie_in(dots, self_mask)
-    dots = dots - 32 * self_mask
-
-    # Mask out later rounds attending to the same items as previous rounds
-    if self.n_hashes > 1 and not self._allow_duplicate_attention:
-      chunks = undo_sort // bq_t.shape[2]  # n_hashes*n_batch*n_heads, seqlen
-      chunks = np.reshape(chunks, (self.n_hashes, -1, seqlen))
-      chunks = np.moveaxis(chunks, 0, -1)
-      chunks = np.tile(chunks, (self.n_hashes, 1, 1))
-      # chunks is now n_hashes*n_batch*n_heads, seqlen, n_hashes
-      schunks = np.take_along_axis(chunks, sjoint_t[:, :, None], axis=-2)
-      bchunks = chunk_vectors(schunks)
-
-      # Queries/keys have shape (n_hashes*n_batch*n_heads, n_bins, binlen). For
-      # each query/key vector, the chunks numbers it's mapped to across each of
-      # the rounds of hashing are stored in bchunks, which has shape
-      # (n_hashes*n_batch*n_heads, n_bins, binlen, n_hashes). Query-key pairs
-      # that fall in the same or consecutive chunks in one hashing round will be
-      # masked out for subsequent rounds.
-      round_counter = jax.lax.tie_in(bchunks, np.arange(self.n_hashes))
-      cur_round = np.tile(
-          np.reshape(round_counter, (-1, 1)),
-          (1, bchunks.shape[0] // self.n_hashes))
-      # cur_round (shape n_hashes*n_batch*n_heads, 1, 1, 1) specifies which
-      # round of hashing a query-key pair belongs to. This shape broadcasts with
-      # (shape n_hashes*n_batch*n_heads, n_bins, binlen, n_hashes). The first
-      # n_batch*n_heads elements along dimension 0 are hash round 0, the next
-      # are round 1, etc.
-      cur_round = np.reshape(cur_round, (-1, 1, 1, 1))
-      # past_round (shape 1, 1, 1, n_hashes) contains round numbers, where the
-      # last dimension has one entry per hashing round.
-      past_round = np.reshape(round_counter, (1, 1, 1, -1))
-      # Set query chunk numbers for future rounds to an out-of-bounds (negative)
-      # value, so they don't match with any keys.
-      bq_chunks = np.where(past_round < cur_round, bchunks, -bchunks)
-
-      bkv_chunks_extra = np.concatenate(
-          [bchunks[:, -1:, :, :], bchunks[:, :-1, :, :]], axis=1)
-      bkv_chunks = np.concatenate([bchunks, bkv_chunks_extra], axis=2)
-
-      dup_mask = np.any(
-          jax.lax.eq(bq_chunks[:, :, :, None, :], bkv_chunks[:, :, None, :, :]),
-          axis=-1)
-      dup_mask = dup_mask | np.any(
-          jax.lax.eq(
-              bq_chunks[:, :, :, None, :], bkv_chunks[:, :, None, :, :] + 1),
-          axis=-1)
-      dup_mask = jax.lax.convert_element_type(dup_mask, np.float32)
-      dots = dots - 30 * dup_mask
-
-    # Softmax.
-    dots_logsumexp = backend.logsumexp(dots, axis=-1, keepdims=True)
-    dots = np.exp(dots - dots_logsumexp)
-
-    if self._hard_k > 0:
-      top_k = np.sort(dots)[..., -self._hard_k]  # Get the top-kth weight.
-      top_k = jax.lax.stop_gradient(top_k)
-      dots -= top_k[..., np.newaxis]  # Subtract (be 0 for lower ones).
-      dots = np.maximum(dots, 0)
-      dots_sum = np.sum(dots, axis=-1, keepdims=True)  # Sum to re-normalize.
-      dots_logsumexp += np.log(dots_sum)  # Add it to the weight.
-      dots /= dots_sum  # Re-normalize.
-
-    bo = np.matmul(dots, bv)
-    so = unchunk_vectors(bo)
-    slogits = unchunk_vectors(dots_logsumexp)
-
-    o = unpermute(so)
-    logits = unpermute(slogits)
-
-    o = np.reshape(o, (self.n_hashes, -1, seqlen, o.shape[-1]))
-    logits = np.reshape(logits, (self.n_hashes, -1, seqlen, 1))
-    probs = np.exp(logits - backend.logsumexp(logits, axis=0, keepdims=True))
-    out = np.sum(o * probs, axis=0)
-    assert out.shape == inputs[2].shape
-
-    return out, state
-
-  def forward_and_backward(self, inputs, ct, rng=None, **kwargs):
-    # TODO(kitaev): is there a manual implementation of forward_and_backward
-    # that's faster than having jax infer one? Or are the permute/unpermute
-    # custom gradients defined in forward() sufficient for reasonable speed?
-    def _do_forward(x):
-      return self.forward(x, params=(), state=(), rng=rng, **kwargs)[0]
-
-    output, vjpfun = jax.vjp(_do_forward, inputs)
-    return output, vjpfun(ct)[0]
+      return out, (q_ct, k_ct, v_ct)
 
 
-class MergedMultiHashedCausalAttentionV2(BaseCausalAttention):
-  """Hash-based causal attention, with multiple hashes (faster version)."""
+class LSHCausalAttention(BaseCausalAttention):
+  """Causal attention based on locality-sensitive hashing."""
 
   def __init__(self, dropout, mode, n_bins=64, n_hashes=1, n_buckets=64,
                one_rng=False, allow_duplicate_attention=False,
@@ -1132,7 +825,7 @@ def __init__(self, dropout, mode, n_bins=64, n_hashes=1, n_buckets=64,
                rehash_each_round=True, drop_for_hash_rate=0.0):
     del dropout
     self._mode = mode
-    super(MergedMultiHashedCausalAttentionV2, self).__init__()
+    super(LSHCausalAttention, self).__init__()
     assert n_buckets >= n_bins, 'This setting is not recommended: too few bins.'
     assert rehash_each_round or allow_duplicate_attention, (
         'The setting {allow_duplicate_attention=False, rehash_each_round=False}'
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
index b1a1ac895..ebf81e273 100644
--- a/tensor2tensor/trax/models/__init__.py
+++ b/tensor2tensor/trax/models/__init__.py
@@ -26,7 +26,7 @@
 from tensor2tensor.trax.models import resnet
 from tensor2tensor.trax.models import transformer
 from tensor2tensor.trax.models.research import position_lookup_transformer
-from tensor2tensor.trax.models.research import transformer_revnet
+from tensor2tensor.trax.models.research import reformer
 
 
 # Ginify
@@ -42,10 +42,10 @@ def model_configure(*args, **kwargs):
 NeuralGPU = model_configure(neural_gpu.NeuralGPU)
 PositionLookupTransformerLM = model_configure(
     position_lookup_transformer.PositionLookupTransformerLM)
+ReformerLM = model_configure(reformer.ReformerLM)
 Resnet50 = model_configure(resnet.Resnet50)
 Transformer = model_configure(transformer.Transformer)
 TransformerDecoder = model_configure(transformer.TransformerDecoder)
 TransformerEncoder = model_configure(transformer.TransformerEncoder)
 TransformerLM = model_configure(transformer.TransformerLM)
-TransformerRevnetLM = model_configure(transformer_revnet.TransformerRevnetLM)
 WideResnet = model_configure(resnet.WideResnet)
diff --git a/tensor2tensor/trax/models/research/transformer_revnet.py b/tensor2tensor/trax/models/research/reformer.py
similarity index 97%
rename from tensor2tensor/trax/models/research/transformer_revnet.py
rename to tensor2tensor/trax/models/research/reformer.py
index abe968762..56b8a5aa5 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet.py
+++ b/tensor2tensor/trax/models/research/reformer.py
@@ -452,20 +452,20 @@ def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
   ]
 
 
-def TransformerRevnetLM(vocab_size,
-                        d_model=512,
-                        d_ff=2048,
-                        d_attention_key=64,
-                        d_attention_value=64,
-                        n_layers=6,
-                        n_heads=8,
-                        dropout=0.1,
-                        max_len=2048,
-                        n_chunks=32,
-                        n_attention_chunks=8,
-                        attention_type=tl.DotProductCausalAttention,
-                        share_qk=False,
-                        mode='train'):
+def ReformerLM(vocab_size,
+               d_model=512,
+               d_ff=2048,
+               d_attention_key=64,
+               d_attention_value=64,
+               n_layers=6,
+               n_heads=8,
+               dropout=0.1,
+               max_len=2048,
+               n_chunks=32,
+               n_attention_chunks=8,
+               attention_type=tl.DotProductCausalAttention,
+               share_qk=False,
+               mode='train'):
   """Reversible transformer language model (only uses a decoder, no encoder).
 
   Args:
diff --git a/tensor2tensor/trax/models/research/transformer_revnet_test.py b/tensor2tensor/trax/models/research/reformer_test.py
similarity index 81%
rename from tensor2tensor/trax/models/research/transformer_revnet_test.py
rename to tensor2tensor/trax/models/research/reformer_test.py
index 5642d54bd..799939748 100644
--- a/tensor2tensor/trax/models/research/transformer_revnet_test.py
+++ b/tensor2tensor/trax/models/research/reformer_test.py
@@ -22,16 +22,16 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 from tensor2tensor.trax import layers as tl
-from tensor2tensor.trax.models.research import transformer_revnet
+from tensor2tensor.trax.models.research import reformer
 
 
-class TransformerRevnetTest(parameterized.TestCase):
+class ReformerTest(parameterized.TestCase):
 
-  def test_transformer_lm_forward_shape(self):
-    """Run the TransformerRevnet LM forward and check output shape."""
+  def test_reformer_lm_forward_shape(self):
+    """Run the ReformerLM forward and check output shape."""
     vocab_size = 16
     input_shape = ((1, 8), (1, 8))
-    model = transformer_revnet.TransformerRevnetLM(
+    model = reformer.ReformerLM(
         vocab_size, d_model=32, d_ff=64,
         d_attention_key=16, d_attention_value=16, n_layers=1, n_heads=2,
         max_len=16, n_chunks=2, n_attention_chunks=1)

From 58f69690fbc6c7dea2fee9fb203a52bc6fb50bf1 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Mon, 30 Sep 2019 13:49:19 -0700
Subject: [PATCH 2506/2720] Fix a bug arising from not feeding the action
 symbols to the model during decoding.

PiperOrigin-RevId: 272057330
---
 .../trax/rl/simulated_env_problem.py          | 89 +++++++------------
 .../trax/rl/simulated_env_problem_test.py     | 42 ++++++---
 2 files changed, 61 insertions(+), 70 deletions(-)

diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index 5536158a3..b994b91ec 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -75,7 +75,7 @@ def __init__(self, model, batch_size, observation_space, action_space,
 
     self._predict_fn = None
     self._rng = None
-    self._model_state = None
+    self._model_state_override = None
     self._history_stream = None
 
     # Call the super's ctor. It will use some of the member fields, so we call
@@ -102,13 +102,19 @@ def initialize_environments(self,
     """
     del parallelism
 
-    model_state = trax.restore_state(self._output_dir)
-    model_params = model_state.opt_state.params
-    self._model_state = model_state.model_state
-    self._predict_fn = functools.partial(
-        self._model_predict,
-        params=model_params,
-    )
+    trax_state = trax.restore_state(self._output_dir)
+    model_params = trax_state.opt_state.params
+
+    # For initializing model state and resetting it.
+    self._model_state_override = trax_state.model_state
+
+    def predict_fn(*args, **kwargs):
+      kwargs["params"] = model_params
+      if self._model_state_override is not None:
+        kwargs["state"] = self._model_state_override
+      return self._model_predict(*args, **kwargs)
+
+    self._predict_fn = predict_fn
     self._history_stream = history_stream
 
     self._steps = np.zeros(batch_size, dtype=np.int32)
@@ -267,8 +273,7 @@ def _reset_model(self, predict_fn, indices, history, rng):
     return history[:, -1, ...]
 
   def _step_model(self, predict_fn, actions, rng):
-    (observation, reward) = predict_fn(
-        (self._history, actions), state=self._model_state, rng=rng)
+    (observation, reward) = predict_fn((self._history, actions), rng=rng)
 
     # Roll the history one timestep back and append the new observation.
     self._history = np.roll(self._history, shift=-1, axis=1)
@@ -285,24 +290,6 @@ def _step_model(self, predict_fn, actions, rng):
     return (observation, reward, done)
 
 
-def index_range_2d(begin_indices, length):
-  # Take all indices along the first dimension. Add another axis that'll
-  # broadcast along the second one.
-  first_dim = np.arange(len(begin_indices))[:, None]
-  # Take a range of indices along the second dimension. Offset it by
-  # begin_indices.
-  # TODO(pkozakowski): This materializes all indices of elements along the
-  # second dimension. Do it more efficiently if needed.
-  second_dim = np.arange(length)[None, :] + begin_indices[:, None]
-  return (first_dim, second_dim)
-
-
-def index_slice(indices):
-  first_dim = np.arange(len(indices))[:, None]
-  second_dim = indices[:, None]
-  return (first_dim, second_dim)
-
-
 class SerializedSequenceSimulatedEnvProblem(SimulatedEnvProblem):
   """SimulatedEnvProblem running a model operating on sequences of symbols.
 
@@ -346,7 +333,6 @@ def __init__(self, model, reward_fn, done_fn, vocab_size,
     self._vocab_size = vocab_size
     self._max_trajectory_length = max_trajectory_length
     self._significance_decay = significance_decay
-    self._history = None
     self._steps = None
     self._observation_space = None
     self._action_space = None
@@ -357,8 +343,8 @@ def __init__(self, model, reward_fn, done_fn, vocab_size,
     self._action_serializer = space_serializer.create(
         action_space, self._vocab_size)
     self._obs_repr_length = self._obs_serializer.representation_length
-    self._action_repr_length = self._action_serializer.representation_length
-    self._step_repr_length = self._obs_repr_length + self._action_repr_length
+    self._act_repr_length = self._action_serializer.representation_length
+    self._step_repr_length = self._obs_repr_length + self._act_repr_length
 
     # We assume that the model takes vocab_size as an argument (e.g.
     # TransformerLM).
@@ -372,10 +358,6 @@ def __init__(self, model, reward_fn, done_fn, vocab_size,
 
   def initialize_environments(self, batch_size=1, **kwargs):
     """Initializes the environments."""
-    self._history = np.zeros((
-        batch_size,
-        self._max_trajectory_length * self._step_repr_length
-    ), dtype=np.int32)
     self._steps = np.zeros(batch_size, dtype=np.int32)
     self._last_observations = np.full(
         (batch_size,) + self._observation_space.shape, np.nan)
@@ -387,40 +369,34 @@ def initialize_environments(self, batch_size=1, **kwargs):
         input_shapes=(batch_size, 1), input_dtype=np.int32, rng=subrng
     )
 
-  @property
-  def _obs_repr_indices(self):
-    begin_indices = self._step_repr_length * self._steps
-    return index_range_2d(begin_indices, self._obs_repr_length)
-
-  @property
-  def _action_repr_indices(self):
-    begin_indices = self._step_repr_length * self._steps + self._obs_repr_length
-    return index_range_2d(begin_indices, self._action_repr_length)
-
   def _predict_obs(self, predict_fn, rng):
+    obs_repr = np.zeros(
+        (self._steps.shape[0], self._obs_repr_length), dtype=np.int32,
+    )
     for (i, subrng) in enumerate(jax_random.split(rng, self._obs_repr_length)):
-      symbol_index = self._steps * self._step_repr_length + i
-      log_probs = predict_fn(
-          self._last_symbols, state=self._model_state, rng=subrng,
-      )
-      log_probs = log_probs
+      log_probs = predict_fn(self._last_symbols, rng=subrng)
       self._last_symbols = utils.gumbel_sample(log_probs)
-      self._history[:, symbol_index] = self._last_symbols[:, 0]
-
-    obs_repr = self._history[self._obs_repr_indices]
+      obs_repr[:, i] = self._last_symbols[:, 0]
     return self._obs_serializer.deserialize(obs_repr)
 
+  def _consume_act(self, actions, predict_fn, rng):
+    act_repr = self._action_serializer.serialize(actions)
+    for (i, subrng) in enumerate(jax_random.split(rng, self._act_repr_length)):
+      # Run the network to update the inference buffers, but ignore the result.
+      predict_fn(self._last_symbols, rng=subrng)
+      self._last_symbols = act_repr[:, i:(i + 1)]
+
   def _reset_model(self, predict_fn, indices, history, rng):
     # TODO(pkozakowski): Random starts.
     del history
 
     indices = np.array(indices)
-    assert indices.shape[0] in (0, self._history.shape[0]), (
+    assert indices.shape[0] in (0, self._steps.shape[0]), (
         # TODO(pkozakowski): Lift this requirement.
         "Only resetting all envs at once is supported."
     )
 
-    self._model_state = self._init_model_state
+    self._model_state_override = self._init_model_state
     self._last_symbols[indices] = 0
     self._steps[indices] = 0
     observation = self._predict_obs(predict_fn, rng)[indices]
@@ -428,8 +404,7 @@ def _reset_model(self, predict_fn, indices, history, rng):
     return observation
 
   def _step_model(self, predict_fn, actions, rng):
-    action_repr = self._action_serializer.serialize(actions)
-    self._history[self._action_repr_indices] = action_repr
+    self._consume_act(actions, predict_fn, rng)
     self._steps += 1
     observation = self._predict_obs(predict_fn, rng)
     reward = self._reward_fn(self._last_observations, observation)
diff --git a/tensor2tensor/trax/rl/simulated_env_problem_test.py b/tensor2tensor/trax/rl/simulated_env_problem_test.py
index 214c2aae5..20e6e2c42 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem_test.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem_test.py
@@ -136,7 +136,7 @@ def _make_env(
         reward_fn=reward_fn,
         done_fn=done_fn,
         vocab_size=vocab_size,
-        max_trajectory_length=3,
+        max_trajectory_length=max_trajectory_length,
         batch_size=batch_size,
         observation_space=observation_space,
         action_space=action_space,
@@ -159,11 +159,12 @@ def test_communicates_with_model(self, mock_restore_state):
     gin.bind_parameter("BoxSpaceSerializer.precision", 1)
     vocab_size = 16
     # Mock model predicting a fixed sequence of symbols. It is made such that
-    # the first two observations are equal and the last one is different.
+    # the first two observations are different and the last one is equal to the
+    # first.
     symbols = [
-        1, 1, 2, 2,  # obs1
-        1, 1, 2, 2,  # obs2
-        1, 2, 2, 1,  # obs3
+        1, 1, 2, 2, 0, 0,  # obs1 act1
+        1, 2, 2, 1, 0, 0,  # obs2 act2
+        1, 1, 2, 2,        # obs3
     ]
     def make_prediction(symbol):
       one_hot = np.eye(vocab_size)[symbol]
@@ -185,19 +186,34 @@ def make_prediction(symbol):
           batch_size=1,
           max_trajectory_length=3,
           observation_space=gym.spaces.Box(low=0, high=5, shape=(4,)),
-          action_space=gym.spaces.Discrete(2),
+          action_space=gym.spaces.MultiDiscrete(nvec=[2, 2]),
       )
-      obs1 = env.reset()
 
-      act1 = 0
-      (obs2, reward, done, _) = env.step(np.array([act1]))
-      np.testing.assert_array_equal(obs1, obs2)
+      def assert_input_suffix(expected_symbols):
+        actual_symbols = np.array([
+            symbol.item() for ((symbol,), _) in mock_predict_fn.call_args_list[
+                -len(expected_symbols):
+            ]
+        ])
+        np.testing.assert_array_equal(actual_symbols, expected_symbols)
+
+      actions = [[0, 1], [1, 0]]
+
+      obs1 = env.reset()
+      assert_input_suffix(symbols[:3])
+
+      (obs2, reward, done, _) = env.step(np.array([actions[0]]))
+      # Symbols going into the decoder when predicting the next observation are:
+      # the last symbol of the previous observation, all action symbols, all
+      # symbols but the last one of the next observation.
+      assert_input_suffix([symbols[3]] + actions[0] + symbols[6:9])
+      self.assertFalse(np.array_equal(obs1, obs2))
       np.testing.assert_array_equal(reward, [0.5])
       np.testing.assert_array_equal(done, [False])
 
-      act2 = 1
-      (obs3, reward, done, _) = env.step(np.array([act2]))
-      self.assertFalse(np.array_equal(obs2, obs3))
+      (obs3, reward, done, _) = env.step(np.array([actions[1]]))
+      assert_input_suffix([symbols[9]] + actions[1] + symbols[12:15])
+      np.testing.assert_array_equal(obs1, obs3)
       np.testing.assert_array_equal(reward, [0.5])
       np.testing.assert_array_equal(done, [True])
 

From 3867955682d5b1779df85eccd03d61d36fe0eae7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 30 Sep 2019 13:56:16 -0700
Subject: [PATCH 2507/2720] Adjust relative masking strengths

PiperOrigin-RevId: 272058835
---
 tensor2tensor/trax/layers/attention.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 0c6156bf5..6006478d1 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -558,7 +558,7 @@ def forward_slice(query_slice, q_loop_idx, key, value):  # pylint: disable=inval
       # Mask out attention to self except when no other targets are available.
       if self._share_qk:
         self_mask = make_self_mask(dots.shape[-2], dots.shape[-1], q_loop_idx)
-        dots = dots - 32 * self_mask
+        dots = dots - 1e5 * self_mask
 
       # Softmax.
       dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))
@@ -775,7 +775,7 @@ def binned_attn(sq, sk, sv):  # pylint: disable=invalid-name
       if self._share_qk:
         self_mask = jax.lax.broadcasted_eye(dots.dtype, dots.shape, (2, 3))
         self_mask = jax.lax.tie_in(dots, self_mask)
-        dots = dots - 32 * self_mask
+        dots = dots - 1e5 * self_mask
 
       # Softmax.
       dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))
@@ -1094,14 +1094,14 @@ def look_one_back(x):
     self_mask = jax.lax.convert_element_type(
         jax.lax.eq(bq_t[:, :, None], bkv_t[:, None, :]),
         np.float32)
-    dots = dots - 1e6 * self_mask
+    dots = dots - 1e5 * self_mask
 
     # Mask out attention to other hash buckets.
     if not self._attend_across_buckets:
       bucket_mask = jax.lax.convert_element_type(
           jax.lax.ne(bq_buckets[:, :, None], bkv_buckets[:, None, :]),
           np.float32)
-      dots = dots - 1e5 * bucket_mask
+      dots = dots - 1e7 * bucket_mask
 
     # Don't double-count query-key pairs across multiple rounds of hashing.
     # There are two possible strategies here. (1) The default is to count how
@@ -1146,7 +1146,7 @@ def look_one_back(x):
           axis=-1)
       assert dup_counts.shape == dots.shape
       if self._hard_k > 0:
-        dots = dots - 1e5 * jax.lax.stop_gradient(dup_counts)
+        dots = dots - 1e7 * jax.lax.stop_gradient(dup_counts)
       else:
         dots = dots - jax.lax.stop_gradient(np.log(dup_counts + 1e-9))
 
@@ -1172,7 +1172,7 @@ def look_one_back(x):
 
       top_k_mask = jax.lax.convert_element_type(
           dots < bdots_thresh[..., None], np.float32)
-      dots = dots - 1e5 * jax.lax.stop_gradient(top_k_mask)
+      dots = dots - 1e7 * jax.lax.stop_gradient(top_k_mask)
 
     # Softmax.
     dots_logsumexp = backend.logsumexp(dots, axis=-1, keepdims=True)

From 331d4fb0298ebd7adb241ce6a7502cb51cbc5400 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 30 Sep 2019 15:23:32 -0700
Subject: [PATCH 2508/2720] Rollback of the sync-RL pipeline, we'll use async
 going forward.

PiperOrigin-RevId: 272078120
---
 oss_scripts/generate_py_proto.sh              |  78 --
 tensor2tensor/envs/__init__.py                |  39 +-
 tensor2tensor/envs/client_env.py              | 139 ---
 tensor2tensor/envs/client_env_test.py         |  99 --
 tensor2tensor/envs/env_service.proto          | 110 ---
 tensor2tensor/envs/env_service_client.py      |  63 --
 .../envs/env_service_generated_pb2.py         | 924 ------------------
 .../envs/env_service_generated_pb2_grpc.py    | 131 ---
 .../envs/env_service_serialization.py         |  89 --
 .../envs/env_service_serialization_test.py    |  95 --
 tensor2tensor/envs/env_service_server.py      |  72 --
 tensor2tensor/envs/env_service_servicer.py    | 111 ---
 .../envs/env_service_servicer_test.py         |  95 --
 tensor2tensor/envs/server_utils.py            |  48 -
 .../trax/rl/envs/env_service_server.py        | 113 ---
 15 files changed, 3 insertions(+), 2203 deletions(-)
 delete mode 100755 oss_scripts/generate_py_proto.sh
 delete mode 100644 tensor2tensor/envs/client_env.py
 delete mode 100644 tensor2tensor/envs/client_env_test.py
 delete mode 100644 tensor2tensor/envs/env_service.proto
 delete mode 100644 tensor2tensor/envs/env_service_client.py
 delete mode 100644 tensor2tensor/envs/env_service_generated_pb2.py
 delete mode 100644 tensor2tensor/envs/env_service_generated_pb2_grpc.py
 delete mode 100644 tensor2tensor/envs/env_service_serialization.py
 delete mode 100644 tensor2tensor/envs/env_service_serialization_test.py
 delete mode 100644 tensor2tensor/envs/env_service_server.py
 delete mode 100644 tensor2tensor/envs/env_service_servicer.py
 delete mode 100644 tensor2tensor/envs/env_service_servicer_test.py
 delete mode 100644 tensor2tensor/envs/server_utils.py
 delete mode 100644 tensor2tensor/trax/rl/envs/env_service_server.py

diff --git a/oss_scripts/generate_py_proto.sh b/oss_scripts/generate_py_proto.sh
deleted file mode 100755
index b72c0b75e..000000000
--- a/oss_scripts/generate_py_proto.sh
+++ /dev/null
@@ -1,78 +0,0 @@
-#!/bin/bash
-
-# This script use the protoc compiler to generate the python code of the
-# all of our proto files.
-
-
-# Function to prepend a pylint directive to skip the generated python file.
-function pylint_skip_file() {
-  local file_name=$1
-  printf "%s\n%s" "# pylint: skip-file" "$(cat ${file_name})" > ${file_name}
-}
-
-
-# Setup tmp directories
-TMP_DIR=$(mktemp -d)
-TMP_TF_DIR=${TMP_DIR}/tensorflow
-TMP_T2T_DIR="$PWD"
-
-echo "Temporary directory created: "
-echo ${TMP_DIR}
-
-
-TMP_T2T_PROTO_DIR="${TMP_T2T_DIR}/tensor2tensor/envs"
-ENV_SERVICE_PROTO="${TMP_T2T_PROTO_DIR}/env_service.proto"
-if [ ! -f ${ENV_SERVICE_PROTO} ]; then
-    echo "${ENV_SERVICE_PROTO} not found."
-    echo "Please run this script from the appropriate root directory."
-fi
-
-# Clone tensorflow repository.
-git clone https://github.com/tensorflow/tensorflow.git ${TMP_TF_DIR}
-
-# Install gRPC tools.
-pip install grpcio-tools
-
-# Invoke the grpc protoc compiler on env_service.proto
-python -m grpc_tools.protoc \
-  --proto_path=${TMP_TF_DIR}/ \
-  --proto_path=${TMP_T2T_DIR}/ \
-  --python_out=${TMP_T2T_DIR}/ \
-  --grpc_python_out=${TMP_T2T_DIR}/ \
-  ${ENV_SERVICE_PROTO}
-
-# Add pylint ignore and name the file as generated.
-GENERATED_ENV_SERVICE_PY="${TMP_T2T_PROTO_DIR}/env_service_generated_pb2.py"
-GENERATED_ENV_SERVICE_GRPC_PY="${TMP_T2T_PROTO_DIR}/env_service_generated_pb2_grpc.py"
-mv ${TMP_T2T_PROTO_DIR}/env_service_pb2.py ${GENERATED_ENV_SERVICE_PY}
-mv ${TMP_T2T_PROTO_DIR}/env_service_pb2_grpc.py ${GENERATED_ENV_SERVICE_GRPC_PY}
-pylint_skip_file "${GENERATED_ENV_SERVICE_PY}"
-pylint_skip_file "${GENERATED_ENV_SERVICE_GRPC_PY}"
-
-
-LICENSING_TEXT=$(cat <<-END
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-END
-)
-
-function add_licensing_text() {
-  local file_name=$1
-  printf "%s\n%s" "${LICENSING_TEXT}" "$(cat ${file_name})" > ${file_name}
-}
-
-add_licensing_text "${GENERATED_ENV_SERVICE_PY}"
-add_licensing_text "${GENERATED_ENV_SERVICE_GRPC_PY}"
-
diff --git a/tensor2tensor/envs/__init__.py b/tensor2tensor/envs/__init__.py
index 3c335c2f2..8c036a176 100644
--- a/tensor2tensor/envs/__init__.py
+++ b/tensor2tensor/envs/__init__.py
@@ -15,43 +15,10 @@
 
 """Environments defined in T2T. Imports here force registration."""
 
-# Proto imports.
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
-
-# pylint: disable=g-import-not-at-top,g-statement-before-imports
-def _get_env_service():
-  from tensor2tensor.envs import env_service_generated_pb2 as env_service_pb2_
-  return env_service_pb2_
-
-
-def _get_env_service_grpc():
-  from tensor2tensor.envs import env_service_generated_pb2_grpc as env_service_pb2_grpc_
-  return env_service_pb2_grpc_
-# pylint: enable=g-import-not-at-top
-
-
-env_service_pb2 = _get_env_service()  # pylint: disable=invalid-name
-env_service_pb2_grpc = _get_env_service_grpc()  # pylint: disable=invalid-name
-del _get_env_service, _get_env_service_grpc
-# pylint: enable=g-statement-before-imports
-
-from gym.envs.registration import register
-
-from tensor2tensor.envs import client_env
 from tensor2tensor.envs import gym_env_problem
 from tensor2tensor.envs import tic_tac_toe_env
 from tensor2tensor.envs import tic_tac_toe_env_problem
-
-
-def register_env(env_class):
-  register(
-      id="{}-v0".format(env_class.__name__),
-      entry_point="tensor2tensor.envs:{}".format(env_class.__name__),
-  )
-  return env_class
-
-
-# TODO(afrozm): Register TicTacToeEnv the same way.
-# register_env(tic_tac_toe_env.TicTacToeEnv)
-ClientEnv = register_env(client_env.ClientEnv)  # pylint: disable=invalid-name
-
diff --git a/tensor2tensor/envs/client_env.py b/tensor2tensor/envs/client_env.py
deleted file mode 100644
index 46e57e136..000000000
--- a/tensor2tensor/envs/client_env.py
+++ /dev/null
@@ -1,139 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Client Env that connects to a distributed env."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl import logging
-import grpc
-import gym
-import numpy as np
-from tensor2tensor.envs import env_service_pb2
-from tensor2tensor.envs import env_service_pb2_grpc
-from tensor2tensor.envs import env_service_serialization as serialization
-
-
-class ClientEnv(gym.Env):
-  """Creates a connection to a remote env, and calls RPC methods on it."""
-
-  @staticmethod
-  def create_channel(remote_env_address):
-    return grpc.insecure_channel(remote_env_address)  # pylint: disable=unreachable
-
-  @staticmethod
-  def run_step(stub, discrete_action):
-    action_proto = env_service_pb2.Action(discrete_action=discrete_action)
-    step_request = env_service_pb2.StepRequest()
-    step_request.action.CopyFrom(action_proto)
-    return stub.Step(step_request)
-
-  @staticmethod
-  def run_reset(stub):
-    return stub.Reset(env_service_pb2.ResetRequest())
-
-  @staticmethod
-  def run_close(stub, channel):
-    close_response = stub.Close(env_service_pb2.CloseRequest())
-    channel.close()
-    return close_response
-
-  @staticmethod
-  def run_render(stub, mode="rgb_array"):
-    return stub.Render(env_service_pb2.RenderRequest(mode=mode))
-
-  @staticmethod
-  def run_get_env_info(stub):
-    env_info_response = stub.GetEnvInfo(env_service_pb2.EnvInfoRequest())
-    gym_observation_space = serialization.proto_to_gym_space(
-        env_info_response.observation_space)
-    gym_action_space = serialization.proto_to_gym_space(
-        env_info_response.action_space)
-    reward_range = (env_info_response.reward_range.low,
-                    env_info_response.reward_range.high)
-    return (gym_action_space, gym_observation_space, reward_range,
-            env_info_response.batch_size)
-
-  def __init__(self, remote_env_address=None, stub=None):
-    self._channel = None
-    self._stub = None
-    self._remote_env_address = None
-
-    if stub is not None:
-      self._stub = stub
-    else:
-      assert remote_env_address is not None
-      logging.vlog(1, "Making a ClientEnv with remote address: [%s]",
-                   remote_env_address)
-      self._remote_env_address = remote_env_address
-      self.initialize_stub()
-
-    assert self._stub is not None
-
-    # We now have to do an RPC to determine spaces and reward range.
-    #
-    # NOTE: If all these are same across replicas, then we technically only need
-    # to do this once on the 'master' replica (say 0), but `GymEnvProblem`
-    # checks that they are all the same.
-    (self.action_space, self.observation_space, self.reward_range,
-     self._server_env_batch_size) = (
-         ClientEnv.run_get_env_info(self._stub))
-
-  def initialize_stub(self):
-    self._channel = ClientEnv.create_channel(self._remote_env_address)
-    # TODO(afrozm): Why is this done?
-    grpc.channel_ready_future(self._channel).result()
-    self._stub = env_service_pb2_grpc.EnvServiceStub(self._channel)
-
-  def _maybe_squeeze_array(self, np_array):
-    # Usually this client is talking to a server env that is running a single
-    # element batch, if so, this client should strip out the batch dimension
-    # before reporting the observation upstream (since this is a plain gym env,
-    # not an EnvProblem), the upstream EnvProblem will then batch across
-    # multiple ClientEnvs.
-    if isinstance(
-        np_array, np.ndarray
-    ) and self._server_env_batch_size == 1 and np_array.shape[0] == 1:
-      np_array = np.squeeze(np_array, axis=0)
-    return np_array
-
-  def reset(self):
-    # Run the RPC.
-    reset_response_proto = ClientEnv.run_reset(self._stub)
-    # Convert the TensorProto to numpy.
-    obs_np = serialization.tensor_proto_to_numpy_array(
-        reset_response_proto.observation.observation)
-    return self._maybe_squeeze_array(obs_np)
-
-  def close(self):
-    ClientEnv.run_close(self._stub, self._channel)
-
-  def render(self, mode="rgb_array"):
-    render_response = ClientEnv.run_render(self._stub, mode=mode)
-    if not render_response:
-      return
-    # Parse out the numpy array.
-    return serialization.tensor_proto_to_numpy_array(
-        render_response.observation.observation)
-
-  def step(self, action):
-    step_response = ClientEnv.run_step(self._stub, action)
-    observation = self._maybe_squeeze_array(
-        serialization.tensor_proto_to_numpy_array(
-            step_response.observation.observation))
-    info = {k: v for k, v in step_response.info.info_map.items()}
-    return observation, step_response.reward, step_response.done, info
diff --git a/tensor2tensor/envs/client_env_test.py b/tensor2tensor/envs/client_env_test.py
deleted file mode 100644
index aa06b9661..000000000
--- a/tensor2tensor/envs/client_env_test.py
+++ /dev/null
@@ -1,99 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.envs.client_env."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gym
-import mock
-import numpy as np
-from tensor2tensor.envs import client_env
-from tensor2tensor.envs import env_service_pb2
-from tensor2tensor.envs import env_service_serialization
-from tensorflow import test
-
-
-class ClientEnvTest(test.TestCase):
-
-  def configure_env_info_on_mock(self, mock_obj):
-    env_info_response = env_service_pb2.EnvInfoResponse()
-    env_info_response.observation_space.box.CopyFrom(
-        env_service_serialization.gym_space_to_proto(
-            gym.spaces.Box(low=0, high=255, shape=(28, 28, 3))).box)
-    env_info_response.action_space.discrete.num_actions = 6
-    env_info_response.reward_range.low = -1
-    env_info_response.reward_range.high = 1
-    env_info_response.batch_size = 1
-    mock_obj.GetEnvInfo.return_value = env_info_response
-
-  def test_get_env_info(self):
-    mock_stub = mock.Mock()
-    self.configure_env_info_on_mock(mock_stub)
-
-    env = client_env.ClientEnv(stub=mock_stub)
-
-    self.assertIsInstance(env.action_space, gym.spaces.Discrete)
-    self.assertIsInstance(env.observation_space, gym.spaces.Box)
-
-    self.assertEqual(6, env.action_space.n)
-    self.assertEqual((28, 28, 3), env.observation_space.shape)
-    self.assertEqual((-1, 1), env.reward_range)
-
-  def test_reset(self):
-    mock_stub = mock.Mock()
-    self.configure_env_info_on_mock(mock_stub)
-    obs_np = np.random.uniform(size=(1, 28, 28, 3))
-    reset_response = env_service_pb2.ResetResponse()
-    reset_response.observation.CopyFrom(
-        env_service_serialization.numpy_array_to_observation(obs_np))
-    mock_stub.Reset.return_value = reset_response
-
-    env = client_env.ClientEnv(stub=mock_stub)
-
-    self.assertAllEqual(np.squeeze(obs_np, axis=0), env.reset())
-
-  def test_step(self):
-    mock_stub = mock.Mock()
-    self.configure_env_info_on_mock(mock_stub)
-    obs_np = np.random.uniform(size=(1, 28, 28, 3))
-    reward = 0.5
-    done = True
-    step_response = env_service_pb2.StepResponse(reward=reward, done=done)
-    step_response.observation.CopyFrom(
-        env_service_serialization.numpy_array_to_observation(obs_np))
-    step_response.info.info_map["k1"] = 1
-    step_response.info.info_map["k2"] = 2
-    mock_stub.Step.return_value = step_response
-
-    action = 4
-    step_request = env_service_pb2.StepRequest(
-        action=env_service_pb2.Action(discrete_action=action))
-
-    env = client_env.ClientEnv(stub=mock_stub)
-    step_retval = env.step(action)
-
-    mock_stub.Step.assert_called_with(step_request)
-    self.assertAllEqual(np.squeeze(obs_np, axis=0), step_retval[0])
-    self.assertEqual(reward, step_retval[1])
-    self.assertEqual(done, step_retval[2])
-    self.assertEqual(1, step_retval[3]["k1"])
-    self.assertEqual(2, step_retval[3]["k2"])
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/envs/env_service.proto b/tensor2tensor/envs/env_service.proto
deleted file mode 100644
index 46556b2e0..000000000
--- a/tensor2tensor/envs/env_service.proto
+++ /dev/null
@@ -1,110 +0,0 @@
-syntax = "proto3";
-option cc_enable_arenas = true;
-
-package tensor2tensor.trax.rlax.envs;
-
-import "tensorflow/core/framework/tensor.proto";
-import "tensorflow/core/framework/tensor_shape.proto";
-import "tensorflow/core/framework/types.proto";
-
-// We use tensorflow.TensorProto to represent numpy arrays.
-
-message Action {
-  oneof payload {
-    int64 discrete_action = 1;
-    tensorflow.TensorProto continuous_action = 2;
-  }
-}
-
-message Observation {
-  tensorflow.TensorProto observation = 1;
-}
-
-message Info {
-  map<string, double> info_map = 1;
-}
-
-message StepRequest {
-  Action action = 1;
-}
-
-message StepResponse {
-  Observation observation = 1;
-  double reward = 2;
-  bool done = 3;
-  Info info = 4;
-}
-
-message ResetRequest {}
-
-message ResetResponse {
-  Observation observation = 1;
-}
-
-message CloseRequest {}
-message CloseResponse {}
-
-message RenderRequest {
-  string mode = 1;
-}
-
-message RenderResponse {
-  Observation observation = 1;
-}
-
-message EnvInfoRequest {}
-
-message SpaceBox {
-  tensorflow.DataType dtype = 1;
-  tensorflow.TensorShapeProto shape = 2;
-  tensorflow.TensorProto low = 3;
-  tensorflow.TensorProto high = 4;
-}
-
-message SpaceDiscrete {
-  int32 num_actions = 1;
-}
-
-message GymSpace {
-  // TODO(afrozm): Move spaces can be added as needed, or composite spaces.
-  oneof gym_space {
-    bool unimplemented_space = 1;
-    SpaceBox box = 2;
-    SpaceDiscrete discrete = 3;
-  }
-}
-
-message RewardRange {
-  double low = 1;
-  double high = 2;
-}
-
-message EnvInfoResponse {
-  GymSpace observation_space = 1;
-  GymSpace action_space = 2;
-  RewardRange reward_range = 3;
-  int64 batch_size = 4;
-}
-
-service EnvService {
-  // Reset
-  rpc Reset(ResetRequest) returns (ResetResponse) {
-  }
-
-  // Step
-  rpc Step(StepRequest) returns (StepResponse) {
-  }
-
-  // Close
-  rpc Close(CloseRequest) returns (CloseResponse) {
-  }
-
-  // Render
-  rpc Render(RenderRequest) returns (RenderResponse) {
-  }
-
-  // Observation and Action Space.
-  rpc GetEnvInfo(EnvInfoRequest) returns (EnvInfoResponse) {
-  }
-}
-
diff --git a/tensor2tensor/envs/env_service_client.py b/tensor2tensor/envs/env_service_client.py
deleted file mode 100644
index b49e86db2..000000000
--- a/tensor2tensor/envs/env_service_client.py
+++ /dev/null
@@ -1,63 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 client binary that talks to remote envs, for debugging."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-import pdb
-from absl import app
-from absl import flags
-import numpy as np  # pylint: disable=unused-import
-from tensor2tensor import envs  # pylint: disable=unused-import
-from tensor2tensor.envs import client_env
-from tensor2tensor.envs import env_problem_utils
-
-FLAGS = flags.FLAGS
-
-flags.DEFINE_string("server_bns", "", "Server's BNS.")
-flags.DEFINE_integer("replicas", 0, "Number of replicas in the server.")
-
-
-def main(argv):
-  del argv
-
-  if FLAGS.replicas == 0:
-    env = client_env.ClientEnv(FLAGS.server_bns)
-    pdb.set_trace()
-    env.close()
-    return
-
-  # Replicated server.
-  per_env_kwargs = [{
-      "remote_env_address": os.path.join(FLAGS.server_bns, str(replica))
-  } for replica in range(FLAGS.replicas)]
-  env = env_problem_utils.make_env(
-      batch_size=FLAGS.replicas,
-      env_problem_name="ClientEnv-v0",
-      resize=False,
-      parallelism=FLAGS.replicas,
-      per_env_kwargs=per_env_kwargs)
-
-  pdb.set_trace()
-
-  env.close()
-
-
-if __name__ == "__main__":
-  app.run(main)
diff --git a/tensor2tensor/envs/env_service_generated_pb2.py b/tensor2tensor/envs/env_service_generated_pb2.py
deleted file mode 100644
index 4efd96d89..000000000
--- a/tensor2tensor/envs/env_service_generated_pb2.py
+++ /dev/null
@@ -1,924 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-
-# pylint: skip-file
-# -*- coding: utf-8 -*-
-# Generated by the protocol buffer compiler.  DO NOT EDIT!
-# source: tensor2tensor/envs/env_service.proto
-
-import sys
-_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
-from google.protobuf import descriptor as _descriptor
-from google.protobuf import message as _message
-from google.protobuf import reflection as _reflection
-from google.protobuf import symbol_database as _symbol_database
-# @@protoc_insertion_point(imports)
-
-_sym_db = _symbol_database.Default()
-
-
-from tensorflow.core.framework import tensor_pb2 as tensorflow_dot_core_dot_framework_dot_tensor__pb2
-from tensorflow.core.framework import tensor_shape_pb2 as tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2
-from tensorflow.core.framework import types_pb2 as tensorflow_dot_core_dot_framework_dot_types__pb2
-
-
-DESCRIPTOR = _descriptor.FileDescriptor(
-  name='tensor2tensor/envs/env_service.proto',
-  package='tensor2tensor.trax.rlax.envs',
-  syntax='proto3',
-  serialized_options=_b('\370\001\001'),
-  serialized_pb=_b('\n$tensor2tensor/envs/env_service.proto\x12\x1ctensor2tensor.trax.rlax.envs\x1a&tensorflow/core/framework/tensor.proto\x1a,tensorflow/core/framework/tensor_shape.proto\x1a%tensorflow/core/framework/types.proto\"d\n\x06\x41\x63tion\x12\x19\n\x0f\x64iscrete_action\x18\x01 \x01(\x03H\x00\x12\x34\n\x11\x63ontinuous_action\x18\x02 \x01(\x0b\x32\x17.tensorflow.TensorProtoH\x00\x42\t\n\x07payload\";\n\x0bObservation\x12,\n\x0bobservation\x18\x01 \x01(\x0b\x32\x17.tensorflow.TensorProto\"y\n\x04Info\x12\x41\n\x08info_map\x18\x01 \x03(\x0b\x32/.tensor2tensor.trax.rlax.envs.Info.InfoMapEntry\x1a.\n\x0cInfoMapEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x01:\x02\x38\x01\"C\n\x0bStepRequest\x12\x34\n\x06\x61\x63tion\x18\x01 \x01(\x0b\x32$.tensor2tensor.trax.rlax.envs.Action\"\x9e\x01\n\x0cStepResponse\x12>\n\x0bobservation\x18\x01 \x01(\x0b\x32).tensor2tensor.trax.rlax.envs.Observation\x12\x0e\n\x06reward\x18\x02 \x01(\x01\x12\x0c\n\x04\x64one\x18\x03 \x01(\x08\x12\x30\n\x04info\x18\x04 \x01(\x0b\x32\".tensor2tensor.trax.rlax.envs.Info\"\x0e\n\x0cResetRequest\"O\n\rResetResponse\x12>\n\x0bobservation\x18\x01 \x01(\x0b\x32).tensor2tensor.trax.rlax.envs.Observation\"\x0e\n\x0c\x43loseRequest\"\x0f\n\rCloseResponse\"\x1d\n\rRenderRequest\x12\x0c\n\x04mode\x18\x01 \x01(\t\"P\n\x0eRenderResponse\x12>\n\x0bobservation\x18\x01 \x01(\x0b\x32).tensor2tensor.trax.rlax.envs.Observation\"\x10\n\x0e\x45nvInfoRequest\"\xa9\x01\n\x08SpaceBox\x12#\n\x05\x64type\x18\x01 \x01(\x0e\x32\x14.tensorflow.DataType\x12+\n\x05shape\x18\x02 \x01(\x0b\x32\x1c.tensorflow.TensorShapeProto\x12$\n\x03low\x18\x03 \x01(\x0b\x32\x17.tensorflow.TensorProto\x12%\n\x04high\x18\x04 \x01(\x0b\x32\x17.tensorflow.TensorProto\"$\n\rSpaceDiscrete\x12\x13\n\x0bnum_actions\x18\x01 \x01(\x05\"\xae\x01\n\x08GymSpace\x12\x1d\n\x13unimplemented_space\x18\x01 \x01(\x08H\x00\x12\x35\n\x03\x62ox\x18\x02 \x01(\x0b\x32&.tensor2tensor.trax.rlax.envs.SpaceBoxH\x00\x12?\n\x08\x64iscrete\x18\x03 \x01(\x0b\x32+.tensor2tensor.trax.rlax.envs.SpaceDiscreteH\x00\x42\x0b\n\tgym_space\"(\n\x0bRewardRange\x12\x0b\n\x03low\x18\x01 \x01(\x01\x12\x0c\n\x04high\x18\x02 \x01(\x01\"\xe7\x01\n\x0f\x45nvInfoResponse\x12\x41\n\x11observation_space\x18\x01 \x01(\x0b\x32&.tensor2tensor.trax.rlax.envs.GymSpace\x12<\n\x0c\x61\x63tion_space\x18\x02 \x01(\x0b\x32&.tensor2tensor.trax.rlax.envs.GymSpace\x12?\n\x0creward_range\x18\x03 \x01(\x0b\x32).tensor2tensor.trax.rlax.envs.RewardRange\x12\x12\n\nbatch_size\x18\x04 \x01(\x03\x32\x89\x04\n\nEnvService\x12\x62\n\x05Reset\x12*.tensor2tensor.trax.rlax.envs.ResetRequest\x1a+.tensor2tensor.trax.rlax.envs.ResetResponse\"\x00\x12_\n\x04Step\x12).tensor2tensor.trax.rlax.envs.StepRequest\x1a*.tensor2tensor.trax.rlax.envs.StepResponse\"\x00\x12\x62\n\x05\x43lose\x12*.tensor2tensor.trax.rlax.envs.CloseRequest\x1a+.tensor2tensor.trax.rlax.envs.CloseResponse\"\x00\x12\x65\n\x06Render\x12+.tensor2tensor.trax.rlax.envs.RenderRequest\x1a,.tensor2tensor.trax.rlax.envs.RenderResponse\"\x00\x12k\n\nGetEnvInfo\x12,.tensor2tensor.trax.rlax.envs.EnvInfoRequest\x1a-.tensor2tensor.trax.rlax.envs.EnvInfoResponse\"\x00\x42\x03\xf8\x01\x01\x62\x06proto3')
-  ,
-  dependencies=[tensorflow_dot_core_dot_framework_dot_tensor__pb2.DESCRIPTOR,tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2.DESCRIPTOR,tensorflow_dot_core_dot_framework_dot_types__pb2.DESCRIPTOR,])
-
-
-
-
-_ACTION = _descriptor.Descriptor(
-  name='Action',
-  full_name='tensor2tensor.trax.rlax.envs.Action',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-    _descriptor.FieldDescriptor(
-      name='discrete_action', full_name='tensor2tensor.trax.rlax.envs.Action.discrete_action', index=0,
-      number=1, type=3, cpp_type=2, label=1,
-      has_default_value=False, default_value=0,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-    _descriptor.FieldDescriptor(
-      name='continuous_action', full_name='tensor2tensor.trax.rlax.envs.Action.continuous_action', index=1,
-      number=2, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-    _descriptor.OneofDescriptor(
-      name='payload', full_name='tensor2tensor.trax.rlax.envs.Action.payload',
-      index=0, containing_type=None, fields=[]),
-  ],
-  serialized_start=195,
-  serialized_end=295,
-)
-
-
-_OBSERVATION = _descriptor.Descriptor(
-  name='Observation',
-  full_name='tensor2tensor.trax.rlax.envs.Observation',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-    _descriptor.FieldDescriptor(
-      name='observation', full_name='tensor2tensor.trax.rlax.envs.Observation.observation', index=0,
-      number=1, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=297,
-  serialized_end=356,
-)
-
-
-_INFO_INFOMAPENTRY = _descriptor.Descriptor(
-  name='InfoMapEntry',
-  full_name='tensor2tensor.trax.rlax.envs.Info.InfoMapEntry',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-    _descriptor.FieldDescriptor(
-      name='key', full_name='tensor2tensor.trax.rlax.envs.Info.InfoMapEntry.key', index=0,
-      number=1, type=9, cpp_type=9, label=1,
-      has_default_value=False, default_value=_b("").decode('utf-8'),
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-    _descriptor.FieldDescriptor(
-      name='value', full_name='tensor2tensor.trax.rlax.envs.Info.InfoMapEntry.value', index=1,
-      number=2, type=1, cpp_type=5, label=1,
-      has_default_value=False, default_value=float(0),
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=_b('8\001'),
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=433,
-  serialized_end=479,
-)
-
-_INFO = _descriptor.Descriptor(
-  name='Info',
-  full_name='tensor2tensor.trax.rlax.envs.Info',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-    _descriptor.FieldDescriptor(
-      name='info_map', full_name='tensor2tensor.trax.rlax.envs.Info.info_map', index=0,
-      number=1, type=11, cpp_type=10, label=3,
-      has_default_value=False, default_value=[],
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-  ],
-  extensions=[
-  ],
-  nested_types=[_INFO_INFOMAPENTRY, ],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=358,
-  serialized_end=479,
-)
-
-
-_STEPREQUEST = _descriptor.Descriptor(
-  name='StepRequest',
-  full_name='tensor2tensor.trax.rlax.envs.StepRequest',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-    _descriptor.FieldDescriptor(
-      name='action', full_name='tensor2tensor.trax.rlax.envs.StepRequest.action', index=0,
-      number=1, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=481,
-  serialized_end=548,
-)
-
-
-_STEPRESPONSE = _descriptor.Descriptor(
-  name='StepResponse',
-  full_name='tensor2tensor.trax.rlax.envs.StepResponse',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-    _descriptor.FieldDescriptor(
-      name='observation', full_name='tensor2tensor.trax.rlax.envs.StepResponse.observation', index=0,
-      number=1, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-    _descriptor.FieldDescriptor(
-      name='reward', full_name='tensor2tensor.trax.rlax.envs.StepResponse.reward', index=1,
-      number=2, type=1, cpp_type=5, label=1,
-      has_default_value=False, default_value=float(0),
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-    _descriptor.FieldDescriptor(
-      name='done', full_name='tensor2tensor.trax.rlax.envs.StepResponse.done', index=2,
-      number=3, type=8, cpp_type=7, label=1,
-      has_default_value=False, default_value=False,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-    _descriptor.FieldDescriptor(
-      name='info', full_name='tensor2tensor.trax.rlax.envs.StepResponse.info', index=3,
-      number=4, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=551,
-  serialized_end=709,
-)
-
-
-_RESETREQUEST = _descriptor.Descriptor(
-  name='ResetRequest',
-  full_name='tensor2tensor.trax.rlax.envs.ResetRequest',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=711,
-  serialized_end=725,
-)
-
-
-_RESETRESPONSE = _descriptor.Descriptor(
-  name='ResetResponse',
-  full_name='tensor2tensor.trax.rlax.envs.ResetResponse',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-    _descriptor.FieldDescriptor(
-      name='observation', full_name='tensor2tensor.trax.rlax.envs.ResetResponse.observation', index=0,
-      number=1, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=727,
-  serialized_end=806,
-)
-
-
-_CLOSEREQUEST = _descriptor.Descriptor(
-  name='CloseRequest',
-  full_name='tensor2tensor.trax.rlax.envs.CloseRequest',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=808,
-  serialized_end=822,
-)
-
-
-_CLOSERESPONSE = _descriptor.Descriptor(
-  name='CloseResponse',
-  full_name='tensor2tensor.trax.rlax.envs.CloseResponse',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=824,
-  serialized_end=839,
-)
-
-
-_RENDERREQUEST = _descriptor.Descriptor(
-  name='RenderRequest',
-  full_name='tensor2tensor.trax.rlax.envs.RenderRequest',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-    _descriptor.FieldDescriptor(
-      name='mode', full_name='tensor2tensor.trax.rlax.envs.RenderRequest.mode', index=0,
-      number=1, type=9, cpp_type=9, label=1,
-      has_default_value=False, default_value=_b("").decode('utf-8'),
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=841,
-  serialized_end=870,
-)
-
-
-_RENDERRESPONSE = _descriptor.Descriptor(
-  name='RenderResponse',
-  full_name='tensor2tensor.trax.rlax.envs.RenderResponse',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-    _descriptor.FieldDescriptor(
-      name='observation', full_name='tensor2tensor.trax.rlax.envs.RenderResponse.observation', index=0,
-      number=1, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=872,
-  serialized_end=952,
-)
-
-
-_ENVINFOREQUEST = _descriptor.Descriptor(
-  name='EnvInfoRequest',
-  full_name='tensor2tensor.trax.rlax.envs.EnvInfoRequest',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=954,
-  serialized_end=970,
-)
-
-
-_SPACEBOX = _descriptor.Descriptor(
-  name='SpaceBox',
-  full_name='tensor2tensor.trax.rlax.envs.SpaceBox',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-    _descriptor.FieldDescriptor(
-      name='dtype', full_name='tensor2tensor.trax.rlax.envs.SpaceBox.dtype', index=0,
-      number=1, type=14, cpp_type=8, label=1,
-      has_default_value=False, default_value=0,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-    _descriptor.FieldDescriptor(
-      name='shape', full_name='tensor2tensor.trax.rlax.envs.SpaceBox.shape', index=1,
-      number=2, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-    _descriptor.FieldDescriptor(
-      name='low', full_name='tensor2tensor.trax.rlax.envs.SpaceBox.low', index=2,
-      number=3, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-    _descriptor.FieldDescriptor(
-      name='high', full_name='tensor2tensor.trax.rlax.envs.SpaceBox.high', index=3,
-      number=4, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=973,
-  serialized_end=1142,
-)
-
-
-_SPACEDISCRETE = _descriptor.Descriptor(
-  name='SpaceDiscrete',
-  full_name='tensor2tensor.trax.rlax.envs.SpaceDiscrete',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-    _descriptor.FieldDescriptor(
-      name='num_actions', full_name='tensor2tensor.trax.rlax.envs.SpaceDiscrete.num_actions', index=0,
-      number=1, type=5, cpp_type=1, label=1,
-      has_default_value=False, default_value=0,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=1144,
-  serialized_end=1180,
-)
-
-
-_GYMSPACE = _descriptor.Descriptor(
-  name='GymSpace',
-  full_name='tensor2tensor.trax.rlax.envs.GymSpace',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-    _descriptor.FieldDescriptor(
-      name='unimplemented_space', full_name='tensor2tensor.trax.rlax.envs.GymSpace.unimplemented_space', index=0,
-      number=1, type=8, cpp_type=7, label=1,
-      has_default_value=False, default_value=False,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-    _descriptor.FieldDescriptor(
-      name='box', full_name='tensor2tensor.trax.rlax.envs.GymSpace.box', index=1,
-      number=2, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-    _descriptor.FieldDescriptor(
-      name='discrete', full_name='tensor2tensor.trax.rlax.envs.GymSpace.discrete', index=2,
-      number=3, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-    _descriptor.OneofDescriptor(
-      name='gym_space', full_name='tensor2tensor.trax.rlax.envs.GymSpace.gym_space',
-      index=0, containing_type=None, fields=[]),
-  ],
-  serialized_start=1183,
-  serialized_end=1357,
-)
-
-
-_REWARDRANGE = _descriptor.Descriptor(
-  name='RewardRange',
-  full_name='tensor2tensor.trax.rlax.envs.RewardRange',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-    _descriptor.FieldDescriptor(
-      name='low', full_name='tensor2tensor.trax.rlax.envs.RewardRange.low', index=0,
-      number=1, type=1, cpp_type=5, label=1,
-      has_default_value=False, default_value=float(0),
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-    _descriptor.FieldDescriptor(
-      name='high', full_name='tensor2tensor.trax.rlax.envs.RewardRange.high', index=1,
-      number=2, type=1, cpp_type=5, label=1,
-      has_default_value=False, default_value=float(0),
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=1359,
-  serialized_end=1399,
-)
-
-
-_ENVINFORESPONSE = _descriptor.Descriptor(
-  name='EnvInfoResponse',
-  full_name='tensor2tensor.trax.rlax.envs.EnvInfoResponse',
-  filename=None,
-  file=DESCRIPTOR,
-  containing_type=None,
-  fields=[
-    _descriptor.FieldDescriptor(
-      name='observation_space', full_name='tensor2tensor.trax.rlax.envs.EnvInfoResponse.observation_space', index=0,
-      number=1, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-    _descriptor.FieldDescriptor(
-      name='action_space', full_name='tensor2tensor.trax.rlax.envs.EnvInfoResponse.action_space', index=1,
-      number=2, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-    _descriptor.FieldDescriptor(
-      name='reward_range', full_name='tensor2tensor.trax.rlax.envs.EnvInfoResponse.reward_range', index=2,
-      number=3, type=11, cpp_type=10, label=1,
-      has_default_value=False, default_value=None,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-    _descriptor.FieldDescriptor(
-      name='batch_size', full_name='tensor2tensor.trax.rlax.envs.EnvInfoResponse.batch_size', index=3,
-      number=4, type=3, cpp_type=2, label=1,
-      has_default_value=False, default_value=0,
-      message_type=None, enum_type=None, containing_type=None,
-      is_extension=False, extension_scope=None,
-      serialized_options=None, file=DESCRIPTOR),
-  ],
-  extensions=[
-  ],
-  nested_types=[],
-  enum_types=[
-  ],
-  serialized_options=None,
-  is_extendable=False,
-  syntax='proto3',
-  extension_ranges=[],
-  oneofs=[
-  ],
-  serialized_start=1402,
-  serialized_end=1633,
-)
-
-_ACTION.fields_by_name['continuous_action'].message_type = tensorflow_dot_core_dot_framework_dot_tensor__pb2._TENSORPROTO
-_ACTION.oneofs_by_name['payload'].fields.append(
-  _ACTION.fields_by_name['discrete_action'])
-_ACTION.fields_by_name['discrete_action'].containing_oneof = _ACTION.oneofs_by_name['payload']
-_ACTION.oneofs_by_name['payload'].fields.append(
-  _ACTION.fields_by_name['continuous_action'])
-_ACTION.fields_by_name['continuous_action'].containing_oneof = _ACTION.oneofs_by_name['payload']
-_OBSERVATION.fields_by_name['observation'].message_type = tensorflow_dot_core_dot_framework_dot_tensor__pb2._TENSORPROTO
-_INFO_INFOMAPENTRY.containing_type = _INFO
-_INFO.fields_by_name['info_map'].message_type = _INFO_INFOMAPENTRY
-_STEPREQUEST.fields_by_name['action'].message_type = _ACTION
-_STEPRESPONSE.fields_by_name['observation'].message_type = _OBSERVATION
-_STEPRESPONSE.fields_by_name['info'].message_type = _INFO
-_RESETRESPONSE.fields_by_name['observation'].message_type = _OBSERVATION
-_RENDERRESPONSE.fields_by_name['observation'].message_type = _OBSERVATION
-_SPACEBOX.fields_by_name['dtype'].enum_type = tensorflow_dot_core_dot_framework_dot_types__pb2._DATATYPE
-_SPACEBOX.fields_by_name['shape'].message_type = tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2._TENSORSHAPEPROTO
-_SPACEBOX.fields_by_name['low'].message_type = tensorflow_dot_core_dot_framework_dot_tensor__pb2._TENSORPROTO
-_SPACEBOX.fields_by_name['high'].message_type = tensorflow_dot_core_dot_framework_dot_tensor__pb2._TENSORPROTO
-_GYMSPACE.fields_by_name['box'].message_type = _SPACEBOX
-_GYMSPACE.fields_by_name['discrete'].message_type = _SPACEDISCRETE
-_GYMSPACE.oneofs_by_name['gym_space'].fields.append(
-  _GYMSPACE.fields_by_name['unimplemented_space'])
-_GYMSPACE.fields_by_name['unimplemented_space'].containing_oneof = _GYMSPACE.oneofs_by_name['gym_space']
-_GYMSPACE.oneofs_by_name['gym_space'].fields.append(
-  _GYMSPACE.fields_by_name['box'])
-_GYMSPACE.fields_by_name['box'].containing_oneof = _GYMSPACE.oneofs_by_name['gym_space']
-_GYMSPACE.oneofs_by_name['gym_space'].fields.append(
-  _GYMSPACE.fields_by_name['discrete'])
-_GYMSPACE.fields_by_name['discrete'].containing_oneof = _GYMSPACE.oneofs_by_name['gym_space']
-_ENVINFORESPONSE.fields_by_name['observation_space'].message_type = _GYMSPACE
-_ENVINFORESPONSE.fields_by_name['action_space'].message_type = _GYMSPACE
-_ENVINFORESPONSE.fields_by_name['reward_range'].message_type = _REWARDRANGE
-DESCRIPTOR.message_types_by_name['Action'] = _ACTION
-DESCRIPTOR.message_types_by_name['Observation'] = _OBSERVATION
-DESCRIPTOR.message_types_by_name['Info'] = _INFO
-DESCRIPTOR.message_types_by_name['StepRequest'] = _STEPREQUEST
-DESCRIPTOR.message_types_by_name['StepResponse'] = _STEPRESPONSE
-DESCRIPTOR.message_types_by_name['ResetRequest'] = _RESETREQUEST
-DESCRIPTOR.message_types_by_name['ResetResponse'] = _RESETRESPONSE
-DESCRIPTOR.message_types_by_name['CloseRequest'] = _CLOSEREQUEST
-DESCRIPTOR.message_types_by_name['CloseResponse'] = _CLOSERESPONSE
-DESCRIPTOR.message_types_by_name['RenderRequest'] = _RENDERREQUEST
-DESCRIPTOR.message_types_by_name['RenderResponse'] = _RENDERRESPONSE
-DESCRIPTOR.message_types_by_name['EnvInfoRequest'] = _ENVINFOREQUEST
-DESCRIPTOR.message_types_by_name['SpaceBox'] = _SPACEBOX
-DESCRIPTOR.message_types_by_name['SpaceDiscrete'] = _SPACEDISCRETE
-DESCRIPTOR.message_types_by_name['GymSpace'] = _GYMSPACE
-DESCRIPTOR.message_types_by_name['RewardRange'] = _REWARDRANGE
-DESCRIPTOR.message_types_by_name['EnvInfoResponse'] = _ENVINFORESPONSE
-_sym_db.RegisterFileDescriptor(DESCRIPTOR)
-
-Action = _reflection.GeneratedProtocolMessageType('Action', (_message.Message,), {
-  'DESCRIPTOR' : _ACTION,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.Action)
-  })
-_sym_db.RegisterMessage(Action)
-
-Observation = _reflection.GeneratedProtocolMessageType('Observation', (_message.Message,), {
-  'DESCRIPTOR' : _OBSERVATION,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.Observation)
-  })
-_sym_db.RegisterMessage(Observation)
-
-Info = _reflection.GeneratedProtocolMessageType('Info', (_message.Message,), {
-
-  'InfoMapEntry' : _reflection.GeneratedProtocolMessageType('InfoMapEntry', (_message.Message,), {
-    'DESCRIPTOR' : _INFO_INFOMAPENTRY,
-    '__module__' : 'tensor2tensor.envs.env_service_pb2'
-    # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.Info.InfoMapEntry)
-    })
-  ,
-  'DESCRIPTOR' : _INFO,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.Info)
-  })
-_sym_db.RegisterMessage(Info)
-_sym_db.RegisterMessage(Info.InfoMapEntry)
-
-StepRequest = _reflection.GeneratedProtocolMessageType('StepRequest', (_message.Message,), {
-  'DESCRIPTOR' : _STEPREQUEST,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.StepRequest)
-  })
-_sym_db.RegisterMessage(StepRequest)
-
-StepResponse = _reflection.GeneratedProtocolMessageType('StepResponse', (_message.Message,), {
-  'DESCRIPTOR' : _STEPRESPONSE,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.StepResponse)
-  })
-_sym_db.RegisterMessage(StepResponse)
-
-ResetRequest = _reflection.GeneratedProtocolMessageType('ResetRequest', (_message.Message,), {
-  'DESCRIPTOR' : _RESETREQUEST,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.ResetRequest)
-  })
-_sym_db.RegisterMessage(ResetRequest)
-
-ResetResponse = _reflection.GeneratedProtocolMessageType('ResetResponse', (_message.Message,), {
-  'DESCRIPTOR' : _RESETRESPONSE,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.ResetResponse)
-  })
-_sym_db.RegisterMessage(ResetResponse)
-
-CloseRequest = _reflection.GeneratedProtocolMessageType('CloseRequest', (_message.Message,), {
-  'DESCRIPTOR' : _CLOSEREQUEST,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.CloseRequest)
-  })
-_sym_db.RegisterMessage(CloseRequest)
-
-CloseResponse = _reflection.GeneratedProtocolMessageType('CloseResponse', (_message.Message,), {
-  'DESCRIPTOR' : _CLOSERESPONSE,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.CloseResponse)
-  })
-_sym_db.RegisterMessage(CloseResponse)
-
-RenderRequest = _reflection.GeneratedProtocolMessageType('RenderRequest', (_message.Message,), {
-  'DESCRIPTOR' : _RENDERREQUEST,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.RenderRequest)
-  })
-_sym_db.RegisterMessage(RenderRequest)
-
-RenderResponse = _reflection.GeneratedProtocolMessageType('RenderResponse', (_message.Message,), {
-  'DESCRIPTOR' : _RENDERRESPONSE,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.RenderResponse)
-  })
-_sym_db.RegisterMessage(RenderResponse)
-
-EnvInfoRequest = _reflection.GeneratedProtocolMessageType('EnvInfoRequest', (_message.Message,), {
-  'DESCRIPTOR' : _ENVINFOREQUEST,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.EnvInfoRequest)
-  })
-_sym_db.RegisterMessage(EnvInfoRequest)
-
-SpaceBox = _reflection.GeneratedProtocolMessageType('SpaceBox', (_message.Message,), {
-  'DESCRIPTOR' : _SPACEBOX,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.SpaceBox)
-  })
-_sym_db.RegisterMessage(SpaceBox)
-
-SpaceDiscrete = _reflection.GeneratedProtocolMessageType('SpaceDiscrete', (_message.Message,), {
-  'DESCRIPTOR' : _SPACEDISCRETE,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.SpaceDiscrete)
-  })
-_sym_db.RegisterMessage(SpaceDiscrete)
-
-GymSpace = _reflection.GeneratedProtocolMessageType('GymSpace', (_message.Message,), {
-  'DESCRIPTOR' : _GYMSPACE,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.GymSpace)
-  })
-_sym_db.RegisterMessage(GymSpace)
-
-RewardRange = _reflection.GeneratedProtocolMessageType('RewardRange', (_message.Message,), {
-  'DESCRIPTOR' : _REWARDRANGE,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.RewardRange)
-  })
-_sym_db.RegisterMessage(RewardRange)
-
-EnvInfoResponse = _reflection.GeneratedProtocolMessageType('EnvInfoResponse', (_message.Message,), {
-  'DESCRIPTOR' : _ENVINFORESPONSE,
-  '__module__' : 'tensor2tensor.envs.env_service_pb2'
-  # @@protoc_insertion_point(class_scope:tensor2tensor.trax.rlax.envs.EnvInfoResponse)
-  })
-_sym_db.RegisterMessage(EnvInfoResponse)
-
-
-DESCRIPTOR._options = None
-_INFO_INFOMAPENTRY._options = None
-
-_ENVSERVICE = _descriptor.ServiceDescriptor(
-  name='EnvService',
-  full_name='tensor2tensor.trax.rlax.envs.EnvService',
-  file=DESCRIPTOR,
-  index=0,
-  serialized_options=None,
-  serialized_start=1636,
-  serialized_end=2157,
-  methods=[
-  _descriptor.MethodDescriptor(
-    name='Reset',
-    full_name='tensor2tensor.trax.rlax.envs.EnvService.Reset',
-    index=0,
-    containing_service=None,
-    input_type=_RESETREQUEST,
-    output_type=_RESETRESPONSE,
-    serialized_options=None,
-  ),
-  _descriptor.MethodDescriptor(
-    name='Step',
-    full_name='tensor2tensor.trax.rlax.envs.EnvService.Step',
-    index=1,
-    containing_service=None,
-    input_type=_STEPREQUEST,
-    output_type=_STEPRESPONSE,
-    serialized_options=None,
-  ),
-  _descriptor.MethodDescriptor(
-    name='Close',
-    full_name='tensor2tensor.trax.rlax.envs.EnvService.Close',
-    index=2,
-    containing_service=None,
-    input_type=_CLOSEREQUEST,
-    output_type=_CLOSERESPONSE,
-    serialized_options=None,
-  ),
-  _descriptor.MethodDescriptor(
-    name='Render',
-    full_name='tensor2tensor.trax.rlax.envs.EnvService.Render',
-    index=3,
-    containing_service=None,
-    input_type=_RENDERREQUEST,
-    output_type=_RENDERRESPONSE,
-    serialized_options=None,
-  ),
-  _descriptor.MethodDescriptor(
-    name='GetEnvInfo',
-    full_name='tensor2tensor.trax.rlax.envs.EnvService.GetEnvInfo',
-    index=4,
-    containing_service=None,
-    input_type=_ENVINFOREQUEST,
-    output_type=_ENVINFORESPONSE,
-    serialized_options=None,
-  ),
-])
-_sym_db.RegisterServiceDescriptor(_ENVSERVICE)
-
-DESCRIPTOR.services_by_name['EnvService'] = _ENVSERVICE
-
-# @@protoc_insertion_point(module_scope)
\ No newline at end of file
diff --git a/tensor2tensor/envs/env_service_generated_pb2_grpc.py b/tensor2tensor/envs/env_service_generated_pb2_grpc.py
deleted file mode 100644
index 378939f00..000000000
--- a/tensor2tensor/envs/env_service_generated_pb2_grpc.py
+++ /dev/null
@@ -1,131 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-
-# pylint: skip-file
-# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
-import grpc
-
-from tensor2tensor.envs import env_service_pb2 as tensor2tensor_dot_envs_dot_env__service__pb2
-
-
-class EnvServiceStub(object):
-  # missing associated documentation comment in .proto file
-  pass
-
-  def __init__(self, channel):
-    """Constructor.
-
-    Args:
-      channel: A grpc.Channel.
-    """
-    self.Reset = channel.unary_unary(
-        '/tensor2tensor.trax.rlax.envs.EnvService/Reset',
-        request_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.ResetRequest.SerializeToString,
-        response_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.ResetResponse.FromString,
-        )
-    self.Step = channel.unary_unary(
-        '/tensor2tensor.trax.rlax.envs.EnvService/Step',
-        request_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.StepRequest.SerializeToString,
-        response_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.StepResponse.FromString,
-        )
-    self.Close = channel.unary_unary(
-        '/tensor2tensor.trax.rlax.envs.EnvService/Close',
-        request_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.CloseRequest.SerializeToString,
-        response_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.CloseResponse.FromString,
-        )
-    self.Render = channel.unary_unary(
-        '/tensor2tensor.trax.rlax.envs.EnvService/Render',
-        request_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.RenderRequest.SerializeToString,
-        response_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.RenderResponse.FromString,
-        )
-    self.GetEnvInfo = channel.unary_unary(
-        '/tensor2tensor.trax.rlax.envs.EnvService/GetEnvInfo',
-        request_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.EnvInfoRequest.SerializeToString,
-        response_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.EnvInfoResponse.FromString,
-        )
-
-
-class EnvServiceServicer(object):
-  # missing associated documentation comment in .proto file
-  pass
-
-  def Reset(self, request, context):
-    """Reset
-    """
-    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-    context.set_details('Method not implemented!')
-    raise NotImplementedError('Method not implemented!')
-
-  def Step(self, request, context):
-    """Step
-    """
-    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-    context.set_details('Method not implemented!')
-    raise NotImplementedError('Method not implemented!')
-
-  def Close(self, request, context):
-    """Close
-    """
-    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-    context.set_details('Method not implemented!')
-    raise NotImplementedError('Method not implemented!')
-
-  def Render(self, request, context):
-    """Render
-    """
-    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-    context.set_details('Method not implemented!')
-    raise NotImplementedError('Method not implemented!')
-
-  def GetEnvInfo(self, request, context):
-    """Observation and Action Space.
-    """
-    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-    context.set_details('Method not implemented!')
-    raise NotImplementedError('Method not implemented!')
-
-
-def add_EnvServiceServicer_to_server(servicer, server):
-  rpc_method_handlers = {
-      'Reset': grpc.unary_unary_rpc_method_handler(
-          servicer.Reset,
-          request_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.ResetRequest.FromString,
-          response_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.ResetResponse.SerializeToString,
-      ),
-      'Step': grpc.unary_unary_rpc_method_handler(
-          servicer.Step,
-          request_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.StepRequest.FromString,
-          response_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.StepResponse.SerializeToString,
-      ),
-      'Close': grpc.unary_unary_rpc_method_handler(
-          servicer.Close,
-          request_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.CloseRequest.FromString,
-          response_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.CloseResponse.SerializeToString,
-      ),
-      'Render': grpc.unary_unary_rpc_method_handler(
-          servicer.Render,
-          request_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.RenderRequest.FromString,
-          response_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.RenderResponse.SerializeToString,
-      ),
-      'GetEnvInfo': grpc.unary_unary_rpc_method_handler(
-          servicer.GetEnvInfo,
-          request_deserializer=tensor2tensor_dot_envs_dot_env__service__pb2.EnvInfoRequest.FromString,
-          response_serializer=tensor2tensor_dot_envs_dot_env__service__pb2.EnvInfoResponse.SerializeToString,
-      ),
-  }
-  generic_handler = grpc.method_handlers_generic_handler(
-      'tensor2tensor.trax.rlax.envs.EnvService', rpc_method_handlers)
-  server.add_generic_rpc_handlers((generic_handler,))
\ No newline at end of file
diff --git a/tensor2tensor/envs/env_service_serialization.py b/tensor2tensor/envs/env_service_serialization.py
deleted file mode 100644
index 51c3f8226..000000000
--- a/tensor2tensor/envs/env_service_serialization.py
+++ /dev/null
@@ -1,89 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Utilities for serializing numpy arrays, gym spaces and envs."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gym
-import numpy as np
-from tensor2tensor.envs import env_service_pb2
-from tensorflow.python.framework import tensor_util  # pylint: disable=g-direct-tensorflow-import
-
-
-def numpy_array_to_observation(array):
-  obs = env_service_pb2.Observation()
-  obs.observation.CopyFrom(tensor_util.make_tensor_proto(array))
-  return obs
-
-
-def tensor_proto_to_numpy_array(tensor_proto):
-  return tensor_util.MakeNdarray(tensor_proto)
-
-
-def step_request_from_discrete_action(action):
-  action_proto = env_service_pb2.Action(discrete_action=action)
-  step_request = env_service_pb2.StepRequest()
-  step_request.action.CopyFrom(action_proto)
-  return step_request
-
-
-def gym_space_to_proto(gym_space):
-  """Converts a gym space to `env_service_pb2.GymSpace`."""
-
-  if isinstance(gym_space, gym.spaces.Discrete):
-    return env_service_pb2.GymSpace(
-        discrete=env_service_pb2.SpaceDiscrete(num_actions=gym_space.n))
-  elif isinstance(gym_space, gym.spaces.Box):
-    space_proto = env_service_pb2.GymSpace()
-    box_proto = space_proto.box
-
-    # Set low & high first, we can set shape and type from it later.
-    box_proto.low.CopyFrom(tensor_util.make_tensor_proto(gym_space.low))
-    box_proto.high.CopyFrom(tensor_util.make_tensor_proto(gym_space.high))
-
-    # dtype and shape.
-    box_proto.dtype = box_proto.low.dtype
-    box_proto.shape.CopyFrom(box_proto.low.tensor_shape)
-
-    return space_proto
-
-  # A space that we haven't implemented.
-  return env_service_pb2.GymSpace(unimplemented_space=True)
-
-
-def proto_to_gym_space(gym_space_proto):
-  """Converts a `env_service_pb2.GymSpace` to a `gym.spaces`."""
-
-  if gym_space_proto.unimplemented_space:
-    return None
-
-  if gym_space_proto.HasField("discrete"):
-    return gym.spaces.Discrete(gym_space_proto.discrete.num_actions)
-
-  assert gym_space_proto.HasField("box")
-
-  low_np = tensor_proto_to_numpy_array(gym_space_proto.box.low)
-  high_np = tensor_proto_to_numpy_array(gym_space_proto.box.high)
-
-  return gym.spaces.Box(low=low_np, high=high_np, dtype=low_np.dtype)
-
-
-def reward_range_to_proto(reward_range=None):
-  if reward_range is None:
-    reward_range = (-np.inf, np.inf)
-  return env_service_pb2.RewardRange(low=reward_range[0], high=reward_range[1])
diff --git a/tensor2tensor/envs/env_service_serialization_test.py b/tensor2tensor/envs/env_service_serialization_test.py
deleted file mode 100644
index 92b3d6feb..000000000
--- a/tensor2tensor/envs/env_service_serialization_test.py
+++ /dev/null
@@ -1,95 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.envs.env_service_serialization."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gym
-import numpy as np
-
-from tensor2tensor.envs import env_service_serialization as utils
-
-from tensorflow import test
-from tensorflow.core.framework import types_pb2  # pylint: disable=g-direct-tensorflow-import
-
-
-class UtilsTest(test.TestCase):
-
-  def test_conversion(self):
-    np_a = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
-    obs = utils.numpy_array_to_observation(np_a)
-
-    tp_a = obs.observation
-    np_tp_a = utils.tensor_proto_to_numpy_array(tp_a)
-
-    np.testing.assert_array_equal(np_a, np_tp_a)
-
-  def test_step_request_from_discrete_action(self):
-    discrete_action = 6
-    step_request = utils.step_request_from_discrete_action(discrete_action)
-    action_request = step_request.action
-    self.assertTrue(action_request.HasField("discrete_action"))
-    self.assertEqual("discrete_action", action_request.WhichOneof("payload"))
-    self.assertEqual(discrete_action, action_request.discrete_action)
-
-  def test_gym_space_to_proto_discrete(self):
-    num_actions = 77
-    space = gym.spaces.Discrete(num_actions)
-    space_proto = utils.gym_space_to_proto(space)
-
-    self.assertFalse(space_proto.HasField("box"))
-    self.assertTrue(space_proto.HasField("discrete"))
-    self.assertEqual(num_actions, space_proto.discrete.num_actions)
-
-  def test_gym_space_to_proto_box(self):
-    space = gym.spaces.Box(low=0, high=255, shape=(28, 29, 3), dtype=np.uint8)
-    space_proto = utils.gym_space_to_proto(space)
-
-    self.assertTrue(space_proto.HasField("box"))
-    self.assertEqual(types_pb2.DT_UINT8, space_proto.box.dtype)
-
-    self.assertEqual(28, space_proto.box.shape.dim[0].size)
-    self.assertEqual(29, space_proto.box.shape.dim[1].size)
-    self.assertEqual(3, space_proto.box.shape.dim[2].size)
-
-  def test_proto_to_gym_space_discrete(self):
-    num_actions = 77
-    space = gym.spaces.Discrete(num_actions)
-    space_proto = utils.gym_space_to_proto(space)
-    space_gym = utils.proto_to_gym_space(space_proto)
-    space_gym.n = num_actions
-
-  def test_proto_to_gym_space_box(self):
-    space = gym.spaces.Box(low=-1.0, high=1.0, shape=(28, 29), dtype=np.float32)
-    space_proto = utils.gym_space_to_proto(space)
-    space_gym = utils.proto_to_gym_space(space_proto)
-    self.assertEqual(np.float32, space_gym.dtype)
-    self.assertAllEqual(space.shape, space_gym.shape)
-
-  def test_reward_range_to_proto(self):
-    reward_proto = utils.reward_range_to_proto((-12, +13))
-    self.assertEqual(-12, reward_proto.low)
-    self.assertEqual(+13, reward_proto.high)
-
-    reward_proto = utils.reward_range_to_proto()
-    self.assertEqual(-np.inf, reward_proto.low)
-    self.assertEqual(np.inf, reward_proto.high)
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/envs/env_service_server.py b/tensor2tensor/envs/env_service_server.py
deleted file mode 100644
index 83f61ebf5..000000000
--- a/tensor2tensor/envs/env_service_server.py
+++ /dev/null
@@ -1,72 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Server that acts as a remote env."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-from absl import app
-from absl import flags
-from absl import logging
-from tensor2tensor.envs import env_problem_utils
-from tensor2tensor.envs import server_utils
-
-FLAGS = flags.FLAGS
-
-flags.DEFINE_bool("xm", False, "Copy atari roms?")
-flags.DEFINE_integer("env_service_port", 7777, "Port on which to run.")
-flags.DEFINE_string("env_problem_name", None, "Name of the EnvProblem to make.")
-flags.DEFINE_string("max_timestep",
-                    None,
-                    "If set to an integer, maximum number of time-steps in a "
-                    "trajectory. The bare env is TimeLimit wrapped.")
-flags.DEFINE_boolean("resize", False, "If true, resize the game frame")
-flags.DEFINE_integer("resized_height", 105, "Resized height of the game frame.")
-flags.DEFINE_integer("resized_width", 80, "Resized width of the game frame.")
-flags.DEFINE_string("output_dir", "", "Output dir.")
-flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
-flags.DEFINE_integer("replica", 0, "Basically to append to output_dir")
-flags.DEFINE_bool("clip_rewards",
-                  True,
-                  "Whether to clip and discretize the rewards.")
-
-# Since we're only dealing with 1 GPU machines here.
-_MAX_CONCURRENCY = 1
-_ADDRESS_FORMAT = "[::]:{}"
-
-
-def main(argv):
-  del argv
-  output_dir = FLAGS.output_dir
-
-  output_dir = os.path.join(output_dir, str(FLAGS.replica))
-
-  env = env_problem_utils.make_env(
-      batch_size=1,
-      env_problem_name=FLAGS.env_problem_name,
-      resize=FLAGS.resize,
-      resize_dims=(FLAGS.resized_height, FLAGS.resized_width),
-      max_timestep=FLAGS.max_timestep,
-      clip_rewards=FLAGS.clip_rewards)
-
-  logging.info("Replica[%s] is ready to serve requests.", FLAGS.replica)
-  server_utils.serve(output_dir, env, FLAGS.env_service_port)
-
-
-if __name__ == "__main__":
-  app.run(main)
diff --git a/tensor2tensor/envs/env_service_servicer.py b/tensor2tensor/envs/env_service_servicer.py
deleted file mode 100644
index c7bc72c7d..000000000
--- a/tensor2tensor/envs/env_service_servicer.py
+++ /dev/null
@@ -1,111 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Implementation of the EnvService RPC."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl import logging
-import grpc
-import numpy as np
-from tensor2tensor.envs import env_service_pb2
-from tensor2tensor.envs import env_service_pb2_grpc
-from tensor2tensor.envs import env_service_serialization as serialization
-
-
-class EnvServiceServicer(env_service_pb2_grpc.EnvServiceServicer):
-  """Implementation of the EnvService service defined in env_service.proto."""
-
-  def __init__(self, env):
-    self._env = env
-
-  def Reset(self, request, context):
-    """Reset."""
-    logging.vlog(1, "EnvServiceServicer is being reset.")
-
-    obs = self._env.reset()
-    reset_response = env_service_pb2.ResetResponse()
-    # Anything more efficient?
-    reset_response.observation.CopyFrom(
-        serialization.numpy_array_to_observation(obs))
-
-    return reset_response
-
-  def Step(self, step_request, context):
-    """Step."""
-    logging.vlog(1, "EnvServiceServicer is being stepped.")
-
-    step_response = env_service_pb2.StepResponse()
-    action = step_request.action
-
-    if "discrete_action" != action.WhichOneof("payload"):
-      context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-      context.set_details("Method not implemented for non-discrete actions!")
-      return step_response
-
-    obs, rewards, dones, infos = self._env.step(
-        np.array([action.discrete_action]))
-
-    step_response.observation.CopyFrom(
-        serialization.numpy_array_to_observation(obs))
-    step_response.reward = rewards
-    step_response.done = dones
-
-    # TODO(afrozm): Take care of this later. `info` is an np array of dicts.
-    if len(infos) > 1:
-      logging.error("Skipping adding the info for other elements in batch.")
-
-    for k, v in infos[0].items():
-      step_response.info.info_map[k] = v
-
-    return step_response
-
-  def Close(self, request, context):
-    """Close."""
-
-    self._env.close()
-    return env_service_pb2.CloseResponse()
-
-  def Render(self, request, context):
-    """Render."""
-
-    mode = request.mode or "rgb_array"
-    rendered_value = self._env.render(mode=mode)
-    response = env_service_pb2.RenderResponse()
-    if (rendered_value is not None) and isinstance(rendered_value, np.ndarray):
-      response.observation = serialization.numpy_array_to_observation(
-          rendered_value)
-
-    return response
-
-  def GetEnvInfo(self, request, context):
-    # Request is empty.
-    del request
-    del context
-
-    response = env_service_pb2.EnvInfoResponse()
-
-    response.observation_space.CopyFrom(
-        serialization.gym_space_to_proto(self._env.observation_space))
-    response.action_space.CopyFrom(
-        serialization.gym_space_to_proto(self._env.action_space))
-    response.reward_range.CopyFrom(
-        serialization.reward_range_to_proto(self._env.reward_range))
-    # Usually these envs aren't batched envs, in that case batch size = 1.
-    response.batch_size = getattr(self._env, "batch_size", 1)
-
-    return response
diff --git a/tensor2tensor/envs/env_service_servicer_test.py b/tensor2tensor/envs/env_service_servicer_test.py
deleted file mode 100644
index 11285088a..000000000
--- a/tensor2tensor/envs/env_service_servicer_test.py
+++ /dev/null
@@ -1,95 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.envs.env_service_servicer."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gym
-import mock
-import numpy as np
-from tensor2tensor.envs import env_service_pb2
-from tensor2tensor.envs import env_service_serialization
-from tensor2tensor.envs import env_service_servicer
-from tensorflow import test
-
-
-class EnvServiceServicerTest(test.TestCase):
-
-  def test_get_env_info(self):
-    env = gym.make("CartPole-v0")
-    env_ss = env_service_servicer.EnvServiceServicer(env)
-    env_info = env_ss.GetEnvInfo(None, None)
-
-    self.assertIsInstance(env_info, env_service_pb2.EnvInfoResponse)
-
-    self.assertTrue(env_info.observation_space.HasField("box"))
-    self.assertTrue(env_info.action_space.HasField("discrete"))
-
-    self.assertEqual(1, len(env_info.observation_space.box.shape.dim))
-    self.assertEqual(4, env_info.observation_space.box.shape.dim[0].size)
-    self.assertEqual(2, env_info.action_space.discrete.num_actions)
-
-    self.assertEqual(-np.inf, env_info.reward_range.low)
-    self.assertEqual(np.inf, env_info.reward_range.high)
-
-    self.assertEqual(1, env_info.batch_size)
-
-  def test_reset(self):
-    # Set expectation on a mock.
-    reset_obs = np.array([0.1, 0.2, 0.3, 0.4])
-    env = mock.Mock()
-    env.reset.return_value = reset_obs
-
-    # Call reset.
-    env_ss = env_service_servicer.EnvServiceServicer(env)
-    reset_response = env_ss.Reset(None, None)
-
-    # Assert the set expectation.
-    self.assertIsInstance(reset_response, env_service_pb2.ResetResponse)
-    self.assertAllEqual(
-        reset_obs,
-        env_service_serialization.tensor_proto_to_numpy_array(
-            reset_response.observation.observation))
-
-  def test_step(self):
-    action = 3
-    step_obs = np.array([1.1, 1.2, 1.3, 1.4])
-    reward = 1.2
-    done = False
-    info = {"k1": 1, "k2": 2}
-
-    env = mock.Mock()
-    env.step.return_value = (step_obs, reward, done, [info])
-
-    env_ss = env_service_servicer.EnvServiceServicer(env)
-    step_request = env_service_pb2.StepRequest(
-        action=env_service_pb2.Action(discrete_action=action))
-    step_response = env_ss.Step(step_request, None)
-
-    self.assertAllEqual(
-        step_obs,
-        env_service_serialization.tensor_proto_to_numpy_array(
-            step_response.observation.observation))
-    self.assertEqual(reward, step_response.reward)
-    self.assertEqual(done, step_response.done)
-    self.assertEqual(info["k1"], step_response.info.info_map["k1"])
-    self.assertEqual(info["k2"], step_response.info.info_map["k2"])
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/envs/server_utils.py b/tensor2tensor/envs/server_utils.py
deleted file mode 100644
index 825cb52a3..000000000
--- a/tensor2tensor/envs/server_utils.py
+++ /dev/null
@@ -1,48 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Utilities for env_service_server.py."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import time
-from absl import logging
-from concurrent import futures
-import grpc
-
-from tensor2tensor.envs import env_service_pb2_grpc
-from tensor2tensor.envs import env_service_servicer
-
-# Since we're only dealing with 1 GPU machines here.
-_MAX_CONCURRENCY = 1
-_ADDRESS_FORMAT = "[::]:{}"
-
-
-def add_port(server, port):
-  return server.add_insecure_port(_ADDRESS_FORMAT.format(port))  # pylint: disable=unreachable
-
-
-def serve(output_dir, env, port):
-  del output_dir  # may use later.
-  server = grpc.server(futures.ThreadPoolExecutor(max_workers=_MAX_CONCURRENCY))
-  servicer = env_service_servicer.EnvServiceServicer(env)
-  env_service_pb2_grpc.add_EnvServiceServicer_to_server(servicer, server)
-  serving_port = add_port(server, port)
-  server.start()
-  logging.info("Starting server on port %s", serving_port)
-  while True:
-    time.sleep(60 * 60 * 24)  # sleep for a day only to sleep again.
diff --git a/tensor2tensor/trax/rl/envs/env_service_server.py b/tensor2tensor/trax/rl/envs/env_service_server.py
deleted file mode 100644
index b81cfd878..000000000
--- a/tensor2tensor/trax/rl/envs/env_service_server.py
+++ /dev/null
@@ -1,113 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Server that acts as a remote env.
-
-NOTE: This is a fork from T2T's `env_service_server.py` since we need to
-link in some TRAX specific envs and gin configuration. This also enables
-eager execution.
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-
-from absl import app
-from absl import flags
-from absl import logging
-import gin
-from tensor2tensor.envs import env_problem_utils
-from tensor2tensor.envs import server_utils
-from tensor2tensor.rl.google import atari_utils
-from tensor2tensor.trax.rl import envs  # pylint: disable=unused-import
-import tensorflow as tf
-
-
-
-FLAGS = flags.FLAGS
-
-flags.DEFINE_bool("xm", False, "Copy atari roms?")
-flags.DEFINE_integer("env_service_port", 7777, "Port on which to run.")
-flags.DEFINE_string("env_problem_name", None, "Name of the EnvProblem to make.")
-flags.DEFINE_string("max_timestep",
-                    None,
-                    "If set to an integer, maximum number of time-steps in a "
-                    "trajectory. The bare env is TimeLimit wrapped.")
-flags.DEFINE_boolean("resize", False, "If true, resize the game frame")
-flags.DEFINE_integer("resized_height", 105, "Resized height of the game frame.")
-flags.DEFINE_integer("resized_width", 80, "Resized width of the game frame.")
-flags.DEFINE_string("output_dir", "", "Output dir.")
-flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
-flags.DEFINE_integer("replica", 0, "Basically to append to output_dir")
-flags.DEFINE_bool("clip_rewards",
-                  True,
-                  "Whether to clip and discretize the rewards.")
-
-# Gin related flags.
-flags.DEFINE_multi_string("gin_config_file",
-                          None,
-                          "Configuration file with parameters (.gin).")
-flags.DEFINE_multi_string("gin_config_string",
-                          [],
-                          "Configuration parameters (gin string).")
-
-
-# TODO(afrozm): Check this.
-flags.DEFINE_bool("enable_eager_execution", False, "")
-
-# Since we're only dealing with 1 GPU machines here.
-_MAX_CONCURRENCY = 1
-_ADDRESS_FORMAT = "[::]:{}"
-
-
-def initialize_gin():
-  gin_bindings = FLAGS.gin_config_string
-  if not (FLAGS.gin_config_file or gin_bindings):
-    return
-  gin.parse_config_files_and_bindings(FLAGS.gin_config_file, gin_bindings)
-
-
-def main(argv):
-  del argv
-
-  if FLAGS.enable_eager_execution:
-    tf.enable_eager_execution()
-
-  output_dir = FLAGS.output_dir
-
-  # Initialize Gin.
-  initialize_gin()
-
-  output_dir = os.path.join(output_dir, str(FLAGS.replica))
-
-  env_kwargs = {"output_dir": output_dir}
-
-  env = env_problem_utils.make_env(
-      batch_size=1,
-      env_problem_name=FLAGS.env_problem_name,
-      resize=FLAGS.resize,
-      resize_dims=(FLAGS.resized_height, FLAGS.resized_width),
-      max_timestep=FLAGS.max_timestep,
-      clip_rewards=FLAGS.clip_rewards,
-      **env_kwargs)
-
-  logging.info("Replica[%s] is ready to serve requests.", FLAGS.replica)
-  server_utils.serve(output_dir, env, FLAGS.env_service_port)
-
-
-if __name__ == "__main__":
-  app.run(main)

From f3f018fe6d7d2c62a115040c3f19d6016678bf6f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 30 Sep 2019 17:07:53 -0700
Subject: [PATCH 2509/2720] Don't sort inside TimeBinCausalAttention

PiperOrigin-RevId: 272099375
---
 tensor2tensor/trax/layers/attention.py | 195 ++++++++-----------------
 1 file changed, 62 insertions(+), 133 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 6006478d1..c1f556360 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -656,164 +656,93 @@ def body_fun(vals):  # pylint: disable=invalid-name
 
 class TimeBinCausalAttention(BaseCausalAttention):
   """Causal attention where only nearby chunks of items attend to each other."""
-  # TODO(kitaev): rewrite this class to use reshapes only, rather than sorting.
 
   def __init__(self, dropout, mode, n_bins=64, share_qk=False):
-    del dropout, mode
     super(TimeBinCausalAttention, self).__init__()
     self.n_bins = n_bins
     self._share_qk = share_qk
+    if dropout >= 1.0:
+      raise ValueError('Dropout rates must be lower than 1.')
+    if mode == 'train':
+      self.dropout = dropout
+    else:
+      self.dropout = None
 
-  def forward(self, inputs, params=(), state=(), **kwargs):
-    del params
-    output, _ = self.forward_and_backward(inputs, None, **kwargs)
-    return output, state
-
-  def has_backward(self):
-    return True
-
-  def backward(self, inputs, output, ct, params=(), state=(), **kwargs):
-    del output, params, state
-    _, inputs_ct = self.forward_and_backward(inputs, ct, **kwargs)
-    return inputs_ct, ()
-
-  def bin_vectors_by_time(self, vecs):
-    seqlen = vecs.shape[-2]
-    assert seqlen % self.n_bins == 0
-    bin_size = int(seqlen // self.n_bins)
-
-    bins = np.arange(seqlen, dtype=np.int32) // bin_size
-    bins = jax.lax.tie_in(vecs, bins)
-    bins = bins[None, :]
-    bins = np.broadcast_to(bins, vecs.shape[:-1])
-    return bins
+  def forward_and_backward(self, inputs, ct, **kwargs):
+    assert backend.get_name() == 'jax', (
+        'JAX backend is required to use forward_and_backward.')
+    # Simultaneous forward pass and backprop through the attention mechanism.
+    def _do_forward(x):  # pylint: disable=invalid-name
+      res, _ = self.forward(x, **kwargs)
+      return res
+    output, vjpfun = jax.vjp(_do_forward, inputs)
+    return output, vjpfun(ct)[0]
 
   def make_unit_length(self, x, epsilon=1e-6):
     variance = np.mean(x**2, axis=-1, keepdims=True)
     norm_inputs = x / np.sqrt(variance + epsilon)
     return norm_inputs
 
-  def forward_and_backward(self, inputs, ct, **kwargs):
-    del kwargs
-    # We use the same vector as both a query and a key. For now we haven't
-    # adjusted any of the surrounding code, so we still get a separate "key"
-    # input that we ignore.
+  def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
+    del params, kwargs
     q, k, v = inputs
     seqlen = q.shape[-2]
     # q/k/v are n_batch*n_heads, seqlen, d_head
+    t = jax.lax.tie_in(q, np.arange(seqlen))
 
-    # bins are n_batch*n_heads, seqlen
-    # They specify which hash bucket the query/key/value vectors fall in.
-    bins = self.bin_vectors_by_time(q)
-
-    # joint_t is n_batch*n_heads, seqlen
-    joint_t = jax.lax.tie_in(q, np.arange(seqlen))
-    joint_t = np.reshape(joint_t, (1, seqlen))
-    joint_t = np.broadcast_to(joint_t, q.shape[:-1])
-
-    assert int((self.n_bins + 1) * seqlen) < 2 ** 31, (
-        'Potential 32-bit integer overflow; please double-check the code.')
-    joint_bins_and_t = seqlen * bins + joint_t
-
-    def chunk_scalars(x):  # pylint: disable=invalid-name
-      return np.reshape(x, (x.shape[0], self.n_bins, -1))
-
-    def chunk_vectors(x):  # pylint: disable=invalid-name
-      return np.reshape(
-          x, (x.shape[0], self.n_bins, -1, x.shape[-1]))
-
-    def unchunk_vectors(x):  # pylint: disable=invalid-name
-      return np.reshape(x, (x.shape[0], -1, x.shape[-1]))
-
-    # Sort everything by bin number, with a secondary sort by time
-    # (variables starting with "s" are sorted)
-    _, sjoint_t = jax.lax.sort_key_val(
-        joint_bins_and_t, joint_t, dimension=-1)
-
-    sq = np.take_along_axis(q, sjoint_t[:, :, None], axis=-2)
+    # Split off a "bin" axis for chunks of consecutive items.
+    bq_t = np.reshape(t, (self.n_bins, -1))
+    bq = np.reshape(q, (q.shape[0], self.n_bins, -1, q.shape[-1]))
     if self._share_qk:
-      sk = sq
+      bk = self.make_unit_length(bq)
     else:
-      sk = np.take_along_axis(k, sjoint_t[:, :, None], axis=-2)
-    sv = np.take_along_axis(v, sjoint_t[:, :, None], axis=-2)
+      bk = np.reshape(k, (k.shape[0], self.n_bins, -1, k.shape[-1]))
+    bv = np.reshape(v, (v.shape[0], self.n_bins, -1, v.shape[-1]))
 
-    if ct is not None:
-      so_ct = np.take_along_axis(ct, sjoint_t[:, :, None], axis=-2)
-
-    @jax.jit
-    def binned_attn(sq, sk, sv):  # pylint: disable=invalid-name
-      """Performs attention on sorted queries/keys/values."""
-      # Split off a "bin" axis so that attention only occurs whithin chunks.
-      bq_t = bkv_t = chunk_scalars(sjoint_t)
-      bq = chunk_vectors(sq)
-      bk = chunk_vectors(sk)
-      if self._share_qk:
-        bk = self.make_unit_length(bk)
-      bv = chunk_vectors(sv)
-
-      # Allow each chunk to attend within itself, and also one chunk back. Chunk
-      # boundaries might occur in the middle of a sequence of items from the
-      # same bin, so this increases the chances of attending to relevant items.
-      # TODO(kitaev): benchmark whether XLA pad operation is noticeably faster.
-      bk_extra = np.concatenate([bk[:, -1:, :, :], bk[:, :-1, :, :]], axis=1)
-      bk = np.concatenate([bk, bk_extra], axis=2)
-      bv_extra = np.concatenate([bv[:, -1:, :, :], bv[:, :-1, :, :]], axis=1)
-      bv = np.concatenate([bv, bv_extra], axis=2)
-      bkv_t_extra = np.concatenate([bkv_t[:, -1:, :], bkv_t[:, :-1, :]], axis=1)
-      bkv_t = np.concatenate([bkv_t, bkv_t_extra], axis=2)
-
-      # Dot-product attention.
-      dots = np.matmul(bq, np.swapaxes(bk, -1, -2)) / np.sqrt(bq.shape[-1])
+    # Allow each chunk to attend within itself, and also one chunk back.
+    def look_one_back(x):
+      if len(x.shape) == 2:
+        x_extra = np.concatenate([x[-1:, :], x[:-1, :]], axis=0)
+        return np.concatenate([x, x_extra], axis=1)
+      else:
+        assert len(x.shape) == 4
+        x_extra = np.concatenate([x[:, -1:, :, :], x[:, :-1, :, :]], axis=1)
+        return np.concatenate([x, x_extra], axis=2)
 
-      # Causal masking
-      mask = jax.lax.convert_element_type(
-          jax.lax.lt(bq_t[:, :, :, None], bkv_t[:, :, None, :]),
-          np.float32)
-      dots = dots - 1e9 * mask
+    bkv_t = look_one_back(bq_t)
+    bk = look_one_back(bk)
+    bv = look_one_back(bv)
 
-      # Mask out attention to self except when no other targets are available.
-      if self._share_qk:
-        self_mask = jax.lax.broadcasted_eye(dots.dtype, dots.shape, (2, 3))
-        self_mask = jax.lax.tie_in(dots, self_mask)
-        dots = dots - 1e5 * self_mask
+    # Dot-product attention.
+    dots = np.matmul(bq, np.swapaxes(bk, -1, -2)) / np.sqrt(bq.shape[-1])
 
-      # Softmax.
-      dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))
-      bo = np.matmul(dots, bv)
-
-      so = unchunk_vectors(bo)
-      return so
-
-    @jax.jit
-    def binned_attn_vjp(sq, sk, sv, so_ct):  # pylint: disable=invalid-name
-      so, vjpfun = jax.vjp(binned_attn, sq, sk, sv)
-      sqkv_ct = vjpfun(so_ct)
-      return so, sqkv_ct
-
-    if ct is None:
-      so = binned_attn(sq, sk, sv)
-      _, undo_sort = jax.lax.sort_key_val(sjoint_t, joint_t, dimension=-1)
-      out = np.take_along_axis(so, undo_sort[:, :, None], axis=-2)
-      return out, None
-    else:
-      # Jax can construct a backward pass automatically, but it's about 2x
-      # slower than writing our own. The main reason is that the backward pass
-      # of gather is in general a scatter operation, but we know we're dealing
-      # with permutations so we use gather for the backward pass too.
-      so, (sq_ct, sk_ct, sv_ct) = binned_attn_vjp(sq, sk, sv, so_ct)
+    # Causal masking
+    mask = jax.lax.convert_element_type(
+        jax.lax.lt(bq_t[None, :, :, None], bkv_t[None, :, None, :]),
+        np.float32)
+    dots = dots - 1e9 * mask
 
-      _, undo_sort = jax.lax.sort_key_val(sjoint_t, joint_t, dimension=-1)
-      out = np.take_along_axis(so, undo_sort[:, :, None], axis=-2)
+    # Mask out attention to self except when no other targets are available.
+    if self._share_qk:
+      self_mask = jax.lax.broadcasted_eye(dots.dtype, dots.shape, (2, 3))
+      self_mask = jax.lax.tie_in(dots, self_mask)
+      dots = dots - 1e5 * self_mask
+
+    if self.dropout is not None and self.dropout > 0.0:
+      # Dropout is broadcast across the batch+head dimension
+      dropout_shape = (1, dots.shape[-3], dots.shape[-2], dots.shape[-1])
+      keep_prob = jax.lax.tie_in(dots, 1.0 - self.dropout)
+      keep = backend.random.bernoulli(rng, keep_prob, dropout_shape)
+      multiplier = keep.astype(dots.dtype) / jax.lax.tie_in(keep, keep_prob)
+      dots = dots * multiplier
 
-      if self._share_qk:
-        q_ct = np.take_along_axis(sq_ct + sk_ct, undo_sort[:, :, None], axis=-2)
-        k_ct = np.zeros_like(k)
-      else:
-        q_ct = np.take_along_axis(sq_ct, undo_sort[:, :, None], axis=-2)
-        k_ct = np.take_along_axis(sk_ct, undo_sort[:, :, None], axis=-2)
-      v_ct = np.take_along_axis(sv_ct, undo_sort[:, :, None], axis=-2)
+    # Softmax.
+    dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))
+    bo = np.matmul(dots, bv)
 
-      return out, (q_ct, k_ct, v_ct)
+    output = np.reshape(bo, (bo.shape[0], -1, bo.shape[-1]))
+    assert output.shape == v.shape
+    return output, state
 
 
 class LSHCausalAttention(BaseCausalAttention):

From 6222d0c45894220db96b35edb7afc6f4193553fd Mon Sep 17 00:00:00 2001
From: Peng Wang <wangpeng@google.com>
Date: Mon, 30 Sep 2019 17:12:46 -0700
Subject: [PATCH 2510/2720] Added a centralized switch in tf-numpy to
 disable/enable float64.

PiperOrigin-RevId: 272100203
---
 tensor2tensor/trax/learning_rate.py | 1 +
 tensor2tensor/trax/trainer.py       | 3 +++
 tensor2tensor/trax/trax.py          | 2 +-
 tensor2tensor/trax/trax_test.py     | 7 +++++--
 4 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
index 02704831c..aa53758c9 100644
--- a/tensor2tensor/trax/learning_rate.py
+++ b/tensor2tensor/trax/learning_rate.py
@@ -88,6 +88,7 @@ def learning_rate(step):  # pylint: disable=invalid-name
         ret *= (decay_factor ** (step//steps_per_decay))
       else:
         raise ValueError("Unknown factor %s." % name)
+    ret = np.asarray(ret, dtype=np.float32)
     return {"learning_rate": ret}
 
   return learning_rate
diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
index 682b70cd3..4e634b19d 100644
--- a/tensor2tensor/trax/trainer.py
+++ b/tensor2tensor/trax/trainer.py
@@ -92,6 +92,8 @@ def _setup_gin():
   gin.parse_config_files_and_bindings(FLAGS.config_file, configs)
 
 
+
+
 def main(_):
 
   logging.set_verbosity(FLAGS.log_level)
@@ -110,6 +112,7 @@ def main(_):
       {"layout_optimizer": FLAGS.tf_opt_layout}
   )
 
+
   _setup_gin()
 
   if FLAGS.enable_eager_execution and backend.get_name() in ("numpy", "jax"):
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 2775d458f..69f9c4097 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -93,7 +93,7 @@ def masked_mean(inputs, targets, weights, mask_id=None):
   if mask_id is not None:
     weights = [w * (1.0 - np.equal(t, mask_id).astype(np.float32))
                for t, w in zip(targets, weights)]
-  weight_sums = [np.float32(t.size) if np.isscalar(w) else np.sum(w)
+  weight_sums = [int(t.size) if np.isscalar(w) else np.sum(w)
                  for w, t in zip(weights, targets)]
   return sum([np.sum(x * w) / (length * s)
               for x, w, s in zip(inputs, weights, weight_sums)])
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index e892f89a7..a10266e46 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -43,10 +43,10 @@
 from tensorflow.io import gfile
 
 
-def test_inputs(n_classes, with_weights=False):
+
+def test_inputs(n_classes, with_weights=False, input_shape=(6, 6, 3)):
   """Make trax.inputs.Inputs."""
   batch_size = 2 * xla_bridge.device_count()
-  input_shape = (6, 6, 3)
 
   def input_stream():
     key = backend.random.get_prng(0)
@@ -72,6 +72,8 @@ def input_stream():
       target_dtype=np.int32)
 
 
+
+
 BACKENDS = ["jax"]
 
 
@@ -243,6 +245,7 @@ def test_reset_twice(self, backend_name):
       trainer.evaluate(1)
 
 
+
 class EpochsTest(test.TestCase):
 
   def test_cuts_epoch_when_total_steps_reached(self):

From bf33311314005528482ea50b098d1aca8da85d84 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Mon, 30 Sep 2019 19:42:04 -0700
Subject: [PATCH 2511/2720] Rename batch_fun to batch_fn in OnlineTune
 baselines.

PiperOrigin-RevId: 272120070
---
 tensor2tensor/trax/configs/transformer_lm1b_16gb.gin |  8 ++++----
 .../trax/configs/transformer_lm_wmt_ende_16gb.gin    | 12 ++++++------
 tensor2tensor/trax/configs/transformer_ptb_16gb.gin  |  8 ++++----
 .../configs/env_online_tune_transformer_ptb_16gb.gin |  8 ++++----
 4 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_lm1b_16gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_16gb.gin
index ea2766d0f..8f4bd0a36 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_16gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_16gb.gin
@@ -3,11 +3,11 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 256
-batch_fun.eval_batch_size = 256
-batch_fun.max_eval_length = 2048
+batch_fn.batch_size_per_device = 256
+batch_fn.eval_batch_size = 256
+batch_fn.max_eval_length = 2048
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin b/tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin
index b61558110..9882f2e9a 100644
--- a/tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin
@@ -3,13 +3,13 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 64
-batch_fun.eval_batch_size = 64
-batch_fun.bucket_length=64
-batch_fun.max_eval_length = 1024
-batch_fun.buckets_include_inputs_in_length=True
+batch_fn.batch_size_per_device = 64
+batch_fn.eval_batch_size = 64
+batch_fn.bucket_length=64
+batch_fn.max_eval_length = 1024
+batch_fn.buckets_include_inputs_in_length=True
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_ptb_16gb.gin b/tensor2tensor/trax/configs/transformer_ptb_16gb.gin
index ed6c337e7..69561693e 100644
--- a/tensor2tensor/trax/configs/transformer_ptb_16gb.gin
+++ b/tensor2tensor/trax/configs/transformer_ptb_16gb.gin
@@ -4,11 +4,11 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 64
-batch_fun.eval_batch_size = 512
-batch_fun.max_eval_length = 2048
+batch_fn.batch_size_per_device = 64
+batch_fn.eval_batch_size = 512
+batch_fn.max_eval_length = 2048
 
 # Parameters for inputs:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin
index 6d05acab7..f7de3538d 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin
@@ -4,11 +4,11 @@ import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.rl
 import tensor2tensor.trax.rl.envs
 
-# Parameters for batch_fun:
+# Parameters for batch_fn:
 # ==============================================================================
-batch_fun.batch_size_per_device = 64
-batch_fun.eval_batch_size = 512
-batch_fun.max_eval_length = 2048
+batch_fn.batch_size_per_device = 64
+batch_fn.eval_batch_size = 512
+batch_fn.max_eval_length = 2048
 
 # Parameters for inputs:
 # ==============================================================================

From 6f7a221534e7550c8f57020f1c32c65c115fe0e3 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 1 Oct 2019 10:28:51 -0700
Subject: [PATCH 2512/2720] Revise docstring for base Layer class.

PiperOrigin-RevId: 272240953
---
 tensor2tensor/trax/layers/base.py | 43 ++++++++++++++++++++-----------
 1 file changed, 28 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index 09ea3b16a..bfc552dbb 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -33,17 +33,29 @@
 class Layer(object):
   """Base class for composable layers in a deep learning network.
 
-  A layer is a part of a trainable network that can compute a function from
-  zero or more inputs to zero or more outputs. It may make use of trainable
-  parameters as well as non-parameter state for its computation. A layer is
-  either atomic or composed of sublayers. All layers provide accessors for
-  these aspects:
-
-    - n_inputs: int
-    - n_outputs: int
-    - params: tuple (empty if the layer has no parameters)
-    - state: tuple (empty if the layer has no non-parameter state)
-    - sublayers: tuple (empty if the layer has no sublayers)
+  Layers are the basic building blocks for deep learning models. A Trax layer
+  computes a function from zero or more inputs to zero or more outputs,
+  optionally using trainable parameters (common) and non-parameter state (not
+  common). Authors of new layer subclasses typically override at most two
+  methods of the base `Layer` class:
+
+    forward(inputs, params=(), state=(), **kwargs):
+      Computes this layer's output as part of a forward pass through the model.
+
+    new_params_and_state(self, input_shape, input_dtype, rng):
+      Returns a (params, state) pair suitable for initializing this layer.
+
+  A small subset of layer types are combinators -- they organize the computation
+  of their sublayers, e.g., applying their sublayers in series or in parallel.
+
+  All layers have the following properties, with default values implemented
+  in the base `Layer` class:
+
+    - n_inputs: int (default 1)
+    - n_outputs: int (default 1)
+    - params: tuple (default empty -- the layer has no parameters)
+    - state: tuple (default empty -- the layer has no non-parameter state)
+    - sublayers: tuple (default empty -- the layer has no sublayers)
 
   The inputs to a layer are tensors, packaged according to how many there are:
 
@@ -51,7 +63,7 @@ class Layer(object):
     - n_inputs = 1: one tensor (NOT wrapped in a tuple)
     - n_inputs > 1: a tuple of tensors
 
-  (The special treatment for the single-input case is meant to simplify the
+  (The special treatment of the single-input case is meant to simplify the
   work of layer writers; this design choice may be revisited in the future.)
 
   The outputs from a layer are also tensors, packaged the same as layer inputs:
@@ -61,9 +73,10 @@ class Layer(object):
     - n_outputs > 1: a tuple of tensors
 
   The Trax runtime maintains a data stack with which layer calls are composed.
-  One can therefore view each layer as a function from stack state to stack
-  state, where the function's inputs are a slice from the stack, and the
-  function's outputs are spliced back into the stack.
+  For more complex data network architectures, possibly involving multiple data
+  flows, one can view each layer as a function from stack state to stack state,
+  where the function's inputs are a slice from the stack, and the function's
+  outputs are spliced back into the stack.
   """
 
   def __init__(self, n_inputs=1, n_outputs=1):

From 67bdc3ebd14339aadc37f5c59f64f98b0d563b57 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 1 Oct 2019 10:47:33 -0700
Subject: [PATCH 2513/2720] Back to passing around state in the world model.
 Inference works now.

PiperOrigin-RevId: 272245726
---
 .../trax/rl/simulated_env_problem.py          | 24 ++++++++++---------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index b994b91ec..54108eaef 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -66,8 +66,12 @@ def __init__(self, model, batch_size, observation_space, action_space,
     if model_predict_kwargs is None:
       model_predict_kwargs = {}
     model_predict = self._model(mode="predict", **model_predict_kwargs)
-    self._model_predict = backend.jit(model_predict)
+    def predict_with_state(*args, **kwargs):
+      output = model_predict(*args, **kwargs)
+      return (output, model_predict.state)
+    self._model_predict = backend.jit(predict_with_state)
     self._model_initialize = model_predict.initialize_once
+
     self._observation_space = observation_space
     self._action_space = action_space
     self._reward_range = reward_range
@@ -75,7 +79,7 @@ def __init__(self, model, batch_size, observation_space, action_space,
 
     self._predict_fn = None
     self._rng = None
-    self._model_state_override = None
+    self._model_state = None
     self._history_stream = None
 
     # Call the super's ctor. It will use some of the member fields, so we call
@@ -104,15 +108,13 @@ def initialize_environments(self,
 
     trax_state = trax.restore_state(self._output_dir)
     model_params = trax_state.opt_state.params
+    self._model_state = trax_state.model_state
 
-    # For initializing model state and resetting it.
-    self._model_state_override = trax_state.model_state
-
-    def predict_fn(*args, **kwargs):
-      kwargs["params"] = model_params
-      if self._model_state_override is not None:
-        kwargs["state"] = self._model_state_override
-      return self._model_predict(*args, **kwargs)
+    def predict_fn(inputs, rng):
+      (output, self._model_state) = self._model_predict(
+          inputs, params=model_params, state=self._model_state, rng=rng
+      )
+      return output
 
     self._predict_fn = predict_fn
     self._history_stream = history_stream
@@ -396,7 +398,7 @@ def _reset_model(self, predict_fn, indices, history, rng):
         "Only resetting all envs at once is supported."
     )
 
-    self._model_state_override = self._init_model_state
+    self._model_state = self._init_model_state
     self._last_symbols[indices] = 0
     self._steps[indices] = 0
     observation = self._predict_obs(predict_fn, rng)[indices]

From d7ac88f3ea5773f447c397b7fa6de5c6ebe37aca Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 1 Oct 2019 11:18:04 -0700
Subject: [PATCH 2514/2720] Internal

PiperOrigin-RevId: 272253582
---
 tensor2tensor/models/evolved_transformer.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 3a2a497ef..130d49206 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -71,7 +71,7 @@ def __init__(self, *args, **kwargs):
     super(EvolvedTransformer, self).__init__(*args, **kwargs)
     self._encoder_function = evolved_transformer_encoder
     self._decoder_function = evolved_transformer_decoder
-    self._init_cache_fn = _init_evolved_transformer_cache
+    self._init_cache_fn = init_evolved_transformer_cache
 
     # -1 means train all weights.
     if self.hparams.get("num_trainable_top_decoder_layers", -1) < 0:
@@ -654,10 +654,10 @@ def _add_attend_to_encoder_cache(cache, attention_name, hparams, num_layers,
   return cache
 
 
-def _init_evolved_transformer_cache(cache, hparams, batch_size,
-                                    attention_init_length, encoder_output,
-                                    encoder_decoder_attention_bias,
-                                    scope_prefix):
+def init_evolved_transformer_cache(cache, hparams, batch_size,
+                                   attention_init_length, encoder_output,
+                                   encoder_decoder_attention_bias,
+                                   scope_prefix):
   """Create the initial cache for Evolved Transformer fast decoding."""
   key_channels = hparams.attention_key_channels or hparams.hidden_size
   value_channels = hparams.attention_value_channels or hparams.hidden_size

From 6e4cd2362b4de49608d689b4af5b4cdec35d2539 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 1 Oct 2019 13:19:17 -0700
Subject: [PATCH 2515/2720] Make losses and metrics instances of layers in
 Trax.

PiperOrigin-RevId: 272279474
---
 .../position_lookup_transformer_copy.gin      |   5 +-
 .../trax/configs/transformer_big_lm1b_8gb.gin |   5 +-
 .../trax/configs/transformer_lm1b_8gb.gin     |   5 +-
 .../configs/transformer_lm1b_8gb_testing.gin  |   5 +-
 ...former_wmt_ende_16gb_adafactor_testing.gin |   5 +-
 .../transformer_wmt_ende_8gb_adafactor.gin    |   5 +-
 .../configs/transformer_wmt_ende_8gb_adam.gin |   5 +-
 .../configs/transformer_wmt_ende_8gb_sm3.gin  |   5 +-
 tensor2tensor/trax/layers/__init__.py         |   1 +
 tensor2tensor/trax/layers/base.py             |  14 +-
 tensor2tensor/trax/layers/combinators.py      |  24 ++
 tensor2tensor/trax/layers/core.py             |   6 +
 tensor2tensor/trax/layers/metrics.py          | 124 +++++++
 tensor2tensor/trax/layers/metrics_test.py     |  90 +++++
 tensor2tensor/trax/models/transformer.py      |   3 +-
 tensor2tensor/trax/models/transformer_test.py |   2 +-
 tensor2tensor/trax/rl/envs/online_tune_env.py |   3 +-
 tensor2tensor/trax/rl/ppo_trainer_test.py     |  25 +-
 .../trax/rl/simulated_env_problem.py          |   8 +-
 tensor2tensor/trax/trax.py                    | 345 +++++++-----------
 tensor2tensor/trax/trax_test.py               |  45 +--
 21 files changed, 423 insertions(+), 307 deletions(-)
 create mode 100644 tensor2tensor/trax/layers/metrics.py
 create mode 100644 tensor2tensor/trax/layers/metrics_test.py

diff --git a/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin b/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
index af0bedd66..d4fde63c0 100644
--- a/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
+++ b/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
@@ -12,10 +12,6 @@ batch_fn.eval_batch_size = 32
 # ==============================================================================
 inputs.dataset_name = 't2t_algorithmic_identity_vocab95_train20_eval30'
 
-# Parameters for masked_mean:
-# ==============================================================================
-masked_mean.mask_id = 0
-
 # Parameters for preprocess_fun:
 # ==============================================================================
 shuffle_and_batch_data.preprocess_fun=@trax.inputs.concat_preprocess
@@ -43,3 +39,4 @@ train.eval_steps = 10
 train.model = @trax.models.PositionLookupTransformerLM
 train.optimizer = @trax.optimizers.Adam
 train.train_steps = 100000
+train.mask_id = 0
\ No newline at end of file
diff --git a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
index c3d110531..1125eb032 100644
--- a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
@@ -15,10 +15,6 @@ inputs.data_dir = None
 inputs.dataset_name = 't2t_languagemodel_lm1b32k'
 inputs.input_name = 'targets'
 
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
 # Parameters for MultifactorSchedule:
 # ==============================================================================
 MultifactorSchedule.constant = 0.1
@@ -39,6 +35,7 @@ train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.SM3
 train.train_steps = 500000
+train.mask_id = 0
 
 # Parameters for DotProductCausalAttention:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
index d44cf8d12..0cfd5c434 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
@@ -15,10 +15,6 @@ inputs.data_dir = None
 inputs.dataset_name = 't2t_languagemodel_lm1b32k'
 inputs.input_name = 'targets'
 
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
 # Parameters for MultifactorSchedule:
 # ==============================================================================
 MultifactorSchedule.constant = 0.3
@@ -39,6 +35,7 @@ train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.SM3
 train.train_steps = 500000
+train.mask_id = 0
 
 # Parameters for DotProductCausalAttention:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
index e0bd85751..a832553af 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
@@ -15,10 +15,6 @@ inputs.data_dir = None
 inputs.dataset_name = 't2t_languagemodel_lm1b32k'
 inputs.input_name = 'targets'
 
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
 # Parameters for MultifactorSchedule:
 # ==============================================================================
 MultifactorSchedule.constant = 0.1
@@ -39,6 +35,7 @@ train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.Adam
 train.train_steps = 100000
+train.mask_id = 0
 
 # Parameters for DotProductCausalAttention:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin
index 198e25e64..40451973d 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin
@@ -15,10 +15,6 @@ batch_fn.buckets_include_inputs_in_length=True
 inputs.data_dir = None
 inputs.dataset_name = 't2t_translate_ende_wmt32k'
 
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
 # Parameters for MultifactorSchedule:
 # ==============================================================================
 MultifactorSchedule.constant = 1.0
@@ -49,6 +45,7 @@ train.inputs = @trax.inputs.inputs
 train.model = @trax.models.Transformer
 train.train_steps = 100000
 train.optimizer = @trax.optimizers.Adafactor
+train.mask_id = 0
 
 # Parameters for Transformer:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
index 06e7d9072..db95a6236 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
@@ -15,10 +15,6 @@ batch_fn.buckets_include_inputs_in_length=True
 inputs.data_dir = None
 inputs.dataset_name = 't2t_translate_ende_wmt32k'
 
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
 # Parameters for MultifactorSchedule:
 # ==============================================================================
 MultifactorSchedule.constant = 1.0
@@ -49,6 +45,7 @@ train.inputs = @trax.inputs.inputs
 train.model = @trax.models.Transformer
 train.train_steps = 500000
 train.optimizer = @trax.optimizers.Adafactor
+train.mask_id = 0
 
 # Parameters for Transformer:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
index e4a010b33..2638003ac 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
@@ -15,10 +15,6 @@ batch_fn.buckets_include_inputs_in_length=True
 inputs.data_dir = None
 inputs.dataset_name = 't2t_translate_ende_wmt32k'
 
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
 # Parameters for MultifactorSchedule:
 # ==============================================================================
 # 0.044 ~= 512^-0.5 = d_model^-0.5
@@ -46,6 +42,7 @@ train.inputs = @trax.inputs.inputs
 train.model = @trax.models.Transformer
 train.train_steps = 500000
 train.optimizer = @trax.optimizers.Adam
+train.mask_id = 0
 
 # Parameters for Transformer:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin
index 310896205..04a4356b4 100644
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin
+++ b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin
@@ -15,10 +15,6 @@ batch_fn.buckets_include_inputs_in_length=True
 inputs.data_dir = None
 inputs.dataset_name = 't2t_translate_ende_wmt32k'
 
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
 # Parameters for MultifactorSchedule:
 # ==============================================================================
 MultifactorSchedule.constant = 0.1
@@ -43,6 +39,7 @@ train.inputs = @trax.inputs.inputs
 train.model = @trax.models.Transformer
 train.train_steps = 500000
 train.optimizer = @trax.optimizers.SM3
+train.mask_id = 0
 
 # Parameters for Transformer:
 # ==============================================================================
diff --git a/tensor2tensor/trax/layers/__init__.py b/tensor2tensor/trax/layers/__init__.py
index 7032b5f79..66643e9d9 100644
--- a/tensor2tensor/trax/layers/__init__.py
+++ b/tensor2tensor/trax/layers/__init__.py
@@ -28,6 +28,7 @@
 from tensor2tensor.trax.layers.convolution import *
 from tensor2tensor.trax.layers.core import *
 from tensor2tensor.trax.layers.initializers import *
+from tensor2tensor.trax.layers.metrics import *
 from tensor2tensor.trax.layers.normalization import *
 from tensor2tensor.trax.layers.pooling import *
 from tensor2tensor.trax.layers.reversible import *
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
index bfc552dbb..658966cfb 100644
--- a/tensor2tensor/trax/layers/base.py
+++ b/tensor2tensor/trax/layers/base.py
@@ -246,7 +246,7 @@ def call_on_input(x, params, state, rng):
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback(skip=3)
       raise LayerError(name, 'pseudo_forward', self._caller, pseudo_inputs,
-                       trace)
+                       None, trace)
 
   def initialize_once(self, input_shapes, input_dtype, rng):
     """Initializes this layer and its sublayers recursively.
@@ -284,7 +284,7 @@ def initialize_once(self, input_shapes, input_dtype, rng):
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback(skip=3)
       raise LayerError(name, 'initialize_once', self._caller, input_shapes,
-                       trace)
+                       input_dtype, trace)
 
   # XXX(kitaev):
   _STASH_IN = None
@@ -356,7 +356,8 @@ def apply_forward(self, x, params=(), state=(), **kwargs):
 
     except Exception:
       name, trace = self.__class__.__name__, _short_traceback()
-      raise LayerError(name, 'apply_forward', self._caller, shapes(x), trace)
+      raise LayerError(name, 'apply_forward', self._caller,
+                       shapes(x), None, trace)
 
   def _do_custom_gradients(self, x, params, state, **kwargs):
     """Calls this layer for a forward pass, but with custom gradients."""
@@ -417,22 +418,27 @@ class LayerError(Exception):
   """
 
   def __init__(self, layer_name, function_name, caller,
-               input_shapes, traceback_string):
+               input_shapes, input_types, traceback_string):
     self._layer_name = layer_name
     self._function_name = function_name
     self._caller = caller  # Python inspect object with init caller info.
     self._traceback = traceback_string
     self._input_shapes = input_shapes
+    self._input_types = input_types
     super(LayerError, self).__init__(self.message)
 
   @property
   def message(self):
+    """Create error message."""
     prefix = 'Exception passing through layer '
     prefix += '%s (in %s):\n' % (self._layer_name, self._function_name)
     short_path = '[...]/' + '/'.join(self._caller.filename.split('/')[-3:])
     caller = '  layer created in file %s, line %d\n' % (short_path,
                                                         self._caller.lineno)
     shapes_str = '  layer input shapes: %s\n\n' % str(self._input_shapes)
+    if self._input_types is not None:
+      types_str = '  layer input types: %s\n' % str(self._input_types)
+      shapes_str = types_str + shapes_str
     return prefix + caller + shapes_str + self._traceback
 
 
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
index 7a8d2d794..311ed2278 100644
--- a/tensor2tensor/trax/layers/combinators.py
+++ b/tensor2tensor/trax/layers/combinators.py
@@ -269,6 +269,30 @@ def Swap(xs, **unused_kwargs):
   return (xs[1], xs[0])
 
 
+def Dup2():
+  """Copy first 2 elements of the stack: (a, b, ...) -> (a, b, a, b, ...)."""
+  return Serial([
+      # Stack is (a, b, ...)
+      Parallel(Dup(), Dup()),  # pylint: disable=no-value-for-parameter
+      # Stack is (a, a, b, b, ...)
+      Parallel([], Swap()),  # pylint: disable=no-value-for-parameter
+      # Stack is (a, b, a, b, ...)
+  ])
+
+
+def Dup3():
+  """Copy 3 elements of the stack: (a, b, c, ...) -> (a, b, c, a, b, c, ...)."""
+  return Serial([
+      # Stack is (a, b, c, ...)
+      Parallel(Dup(), Dup(), Dup()),  # pylint: disable=no-value-for-parameter
+      # Stack is (a, a, b, b, c, c, ...)
+      Parallel([], Swap(), Swap()),  # pylint: disable=no-value-for-parameter
+      # Stack is (a, b, a, c, b, c, ...)
+      Parallel([], [], Swap()),  # pylint: disable=no-value-for-parameter
+      # Stack is (a, b, c, a, b, c, ...)
+  ])
+
+
 @base.layer(n_outputs=0)
 def Drop(x, **unused_kwargs):
   """Drops one element."""
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index 1cb8fc0b4..d6a338a9c 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -210,6 +210,12 @@ def AddConstant(x, constant=0.0, **unused_kwargs):
   return x + constant
 
 
+@base.layer()
+def MulConstant(x, params, constant=1.0, **unused_kwargs):
+  del params
+  return x * constant
+
+
 def one_hot(x, size, dtype=np.float32):  # pylint: disable=invalid-name
   """Make a n+1 dim one-hot array from n dim int-categorical array."""
   arange_size = np.arange(size)
diff --git a/tensor2tensor/trax/layers/metrics.py b/tensor2tensor/trax/layers/metrics.py
new file mode 100644
index 000000000..e79c83cb2
--- /dev/null
+++ b/tensor2tensor/trax/layers/metrics.py
@@ -0,0 +1,124 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Trax metrics layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.trax.backend import numpy as np
+from tensor2tensor.trax.layers import base
+from tensor2tensor.trax.layers import combinators as cb
+from tensor2tensor.trax.layers import core
+
+
+@base.layer(n_inputs=2, n_outputs=1)
+def CrossEntropy(x, axis=-1, **kw):
+  del kw
+  prediction, target = x
+  return np.sum(prediction * core.one_hot(target, prediction.shape[-1]),
+                axis=axis)
+
+
+@base.layer(n_inputs=2, n_outputs=1)
+def L2(x, axis=-1, **kw):
+  del kw
+  prediction, target = x
+  return np.sum((prediction - target)**2, axis=axis)
+
+
+@base.layer(n_inputs=2, n_outputs=1)
+def Accuracy(x, axis=-1, **kw):
+  del kw
+  prediction, target = x
+  predicted_class = np.argmax(prediction, axis=axis)
+  return np.equal(predicted_class, target)
+
+
+@base.layer()
+def WeightMask(target, mask_id=0, **kw):
+  del kw
+  if mask_id is None:
+    return np.ones_like(target)
+  return 1.0 - np.equal(target, mask_id).astype(np.float32)
+
+
+@base.layer(n_inputs=2, n_outputs=1)
+def WeightedMean(x, **kw):
+  del kw
+  metric, weights = x
+  weights_sum = np.sum(weights)
+  return np.sum(metric * weights) / weights_sum
+
+
+def MaskedScalar(metric_layer, mask_id=None, has_weights=False):
+  """Metric as scalar compatible with Trax masking."""
+  # Stack of (inputs, targets) --> (metric, weight-mask).
+  metric_and_mask = [
+      cb.Parallel(
+          [],
+          cb.Dup()  # Duplicate targets
+      ),
+      cb.Parallel(
+          metric_layer,  # Metric: (inputs, targets) --> metric
+          WeightMask(mask_id=mask_id)  # pylint: disable=no-value-for-parameter
+      )
+  ]
+  if not has_weights:
+    # Take (metric, weight-mask) and return the weighted mean.
+    return cb.Serial([metric_and_mask, WeightedMean()])  # pylint: disable=no-value-for-parameter
+  return cb.Serial([
+      metric_and_mask,
+      cb.Parallel(
+          [],
+          cb.Multiply()  # Multiply given weights by mask_id weights
+      ),
+      WeightedMean()  # pylint: disable=no-value-for-parameter
+  ])
+
+
+def CrossEntropyScalar(mask_id=None, has_weights=False):
+  """Cross-entropy as scalar compatible with Trax masking."""
+  return MaskedScalar(CrossEntropy(), mask_id=mask_id, has_weights=has_weights)  # pylint: disable=no-value-for-parameter
+
+
+NegLogPerplexityScalar = CrossEntropyScalar
+
+
+def CrossEntropyLossScalar(mask_id=None, has_weights=False):
+  """Cross-entropy loss as scalar compatible with Trax masking."""
+  return cb.Serial(
+      CrossEntropyScalar(mask_id=mask_id, has_weights=has_weights),
+      core.MulConstant(constant=-1.0)
+  )
+
+
+def L2Scalar(mask_id=None, has_weights=False):
+  """L2 as scalar compatible with Trax masking."""
+  return MaskedScalar(L2(), mask_id=mask_id, has_weights=has_weights)  # pylint: disable=no-value-for-parameter
+
+
+def L2LossScalar(mask_id=None, has_weights=False):
+  """L2 loss as scalar compatible with Trax masking."""
+  return cb.Serial(
+      L2Scalar(mask_id=mask_id, has_weights=has_weights),
+      core.MulConstant(constant=-1.0)
+  )
+
+
+def AccuracyScalar(mask_id=None, has_weights=False):
+  """Accuracy as scalar compatible with Trax masking."""
+  return MaskedScalar(Accuracy(), mask_id=mask_id, has_weights=has_weights)  # pylint: disable=no-value-for-parameter
diff --git a/tensor2tensor/trax/layers/metrics_test.py b/tensor2tensor/trax/layers/metrics_test.py
new file mode 100644
index 000000000..8eded2671
--- /dev/null
+++ b/tensor2tensor/trax/layers/metrics_test.py
@@ -0,0 +1,90 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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 metrics layers."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+import numpy as onp
+from tensor2tensor.trax import backend
+from tensor2tensor.trax.layers import base
+from tensor2tensor.trax.layers import metrics
+
+
+class MetricsLayerTest(absltest.TestCase):
+
+  def test_cross_entropy(self):
+    input_shape = ((29, 4, 4, 20), (29, 4, 4))
+    result_shape = base.check_shape_agreement(
+        metrics.CrossEntropy(), input_shape)
+    self.assertEqual(result_shape, (29, 4, 4))
+
+  def test_accuracy(self):
+    input_shape = ((29, 4, 4, 20), (29, 4, 4))
+    result_shape = base.check_shape_agreement(
+        metrics.Accuracy(), input_shape)
+    self.assertEqual(result_shape, (29, 4, 4))
+
+  def test_weight_mask(self):
+    input_shape = (29, 4, 4, 20)
+    result_shape = base.check_shape_agreement(
+        metrics.WeightMask(), input_shape)
+    self.assertEqual(result_shape, input_shape)
+
+  def test_weighted_mean_shape(self):
+    input_shape = ((29, 4, 4, 20), (29, 4, 4, 20))
+    result_shape = base.check_shape_agreement(
+        metrics.WeightedMean(), input_shape)
+    self.assertEqual(result_shape, ())
+
+  def test_weighted_mean_semantics(self):
+    inputs = onp.array([1, 2, 3], dtype=onp.float32)
+    weights1 = onp.array([1, 1, 1], dtype=onp.float32)
+    layer = metrics.WeightedMean()
+    rng = backend.random.get_prng(0)
+    layer.initialize_once((inputs.shape, weights1.shape),
+                          (inputs.dtype, weights1.dtype), rng)
+    mean1 = layer((inputs, weights1))
+    onp.testing.assert_allclose(mean1, 2.0)
+    weights2 = onp.array([0, 0, 1], dtype=onp.float32)
+    mean2 = layer((inputs, weights2))
+    onp.testing.assert_allclose(mean2, 3.0)
+    weights3 = onp.array([1, 0, 0], dtype=onp.float32)
+    mean3 = layer((inputs, weights3))
+    onp.testing.assert_allclose(mean3, 1.0)
+
+  def test_cross_entropy_scalar(self):
+    input_shape = ((29, 4, 4, 20), (29, 4, 4))
+    result_shape = base.check_shape_agreement(
+        metrics.CrossEntropyScalar(), input_shape)
+    self.assertEqual(result_shape, ())
+
+  def test_cross_entropy_loss_scalar(self):
+    input_shape = ((29, 4, 4, 20), (29, 4, 4))
+    result_shape = base.check_shape_agreement(
+        metrics.CrossEntropyLossScalar(), input_shape)
+    self.assertEqual(result_shape, ())
+
+  def test_accuracy_scalar(self):
+    input_shape = ((29, 4, 4, 20), (29, 4, 4))
+    result_shape = base.check_shape_agreement(
+        metrics.AccuracyScalar(), input_shape)
+    self.assertEqual(result_shape, ())
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
index 300a381bb..640d54bf7 100644
--- a/tensor2tensor/trax/models/transformer.py
+++ b/tensor2tensor/trax/models/transformer.py
@@ -371,7 +371,8 @@ def Transformer(input_vocab_size,
 
   # Input: encoder_side_tokens, decoder_side_tokens
   return tl.Model(  # tokens_e tokens_d
-      tl.Swap(),    # toks_d toks_e
+      tl.Parallel([], tl.Dup()),    # toks_e toks_d toks_d (for loss)
+      tl.Swap(),    # toks_d toks_e ....
 
       # Encode.
       tl.Parallel(                                       # toks_d        toks_e
diff --git a/tensor2tensor/trax/models/transformer_test.py b/tensor2tensor/trax/models/transformer_test.py
index 5b87a2026..70266a8a3 100644
--- a/tensor2tensor/trax/models/transformer_test.py
+++ b/tensor2tensor/trax/models/transformer_test.py
@@ -55,7 +55,7 @@ def _test_transformer_forward_shape(self, input_vocab_size,
     expected_shape = (tuple(single_input_shape +
                             [output_vocab_size if output_vocab_size is not None
                              else input_vocab_size]))
-    self.assertEqual(expected_shape, final_shape)
+    self.assertEqual(expected_shape, final_shape[0])
 
   def test_transformer_lm_fast_inference(self):
     with backend.use_backend('jax'):
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index a66d59a57..804872e3a 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -25,6 +25,7 @@
 import gym
 
 from tensor2tensor.trax import inputs as trax_inputs
+from tensor2tensor.trax import layers
 from tensor2tensor.trax import models as trax_models
 from tensor2tensor.trax import optimizers as trax_opt
 from tensor2tensor.trax import trax
@@ -51,7 +52,7 @@ def __init__(self,
                output_dir,
                model=trax_models.TransformerLM,
                trainer_class=trax.Trainer,
-               loss_fn=trax.loss,
+               loss_fn=layers.CrossEntropyLossScalar,
                optimizer=trax_opt.Adafactor,
                inputs=trax_inputs.inputs,
                action_multipliers=None,
diff --git a/tensor2tensor/trax/rl/ppo_trainer_test.py b/tensor2tensor/trax/rl/ppo_trainer_test.py
index 38b7c0bfc..134f5207d 100644
--- a/tensor2tensor/trax/rl/ppo_trainer_test.py
+++ b/tensor2tensor/trax/rl/ppo_trainer_test.py
@@ -166,10 +166,10 @@ def model(mode):
 
     def inputs(n_devices):
       del n_devices
-      stream = itertools.repeat((
-          (np.zeros(history_shape), np.zeros(action_shape, dtype=np.int32)),
-          (np.zeros(obs_shape), np.zeros(reward_shape)),
-      ))
+      stream = itertools.repeat(
+          (np.zeros(history_shape), np.zeros(action_shape, dtype=np.int32),
+           np.zeros(obs_shape), np.zeros(reward_shape))
+      )
       return trax_inputs.Inputs(
           train_stream=lambda: stream,
           train_eval_stream=lambda: stream,
@@ -180,9 +180,20 @@ def inputs(n_devices):
           target_dtype=(np.float32, np.float32),
       )
 
-    def loss(params, batch, model_predict, state, rng, **kwargs):
-      del params, batch, model_predict, rng, kwargs
-      return 0.0, state
+    def loss(mask_id=None, has_weights=False):
+      """Cross-entropy loss as scalar compatible with Trax masking."""
+      return layers.Serial(
+          # Swap from (pred-obs, pred-reward, target-obs, target-reward)
+          # to (pred-obs, target-obs, pred-reward, target-reward).
+          layers.Parallel([], layers.Swap()),
+          # Cross-entropy loss for obs, L2 loss on reward.
+          layers.Parallel(layers.CrossEntropyLossScalar(mask_id, has_weights),
+                          layers.L2LossScalar(mask_id, has_weights)),
+          # Add both losses.
+          layers.Add(),
+          # Zero out in this test.
+          layers.MulConstant(constant=0.0)
+      )
 
     with self.tmp_dir() as output_dir:
       # Run fake training just to save the parameters.
diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
index 54108eaef..42d1c73b7 100644
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ b/tensor2tensor/trax/rl/simulated_env_problem.py
@@ -107,8 +107,11 @@ def initialize_environments(self,
     del parallelism
 
     trax_state = trax.restore_state(self._output_dir)
-    model_params = trax_state.opt_state.params
-    self._model_state = trax_state.model_state
+    # TODO(lukaszkaiser): both model state and parameters by default include
+    # the loss layer. Currently, we access the pure-model parameters by just
+    # indexing, [0] here. But we should make it more explicit in a better API.
+    model_params = trax_state.opt_state.params[0]
+    self._model_state = trax_state.model_state[0]
 
     def predict_fn(inputs, rng):
       (output, self._model_state) = self._model_predict(
@@ -118,7 +121,6 @@ def predict_fn(inputs, rng):
 
     self._predict_fn = predict_fn
     self._history_stream = history_stream
-
     self._steps = np.zeros(batch_size, dtype=np.int32)
 
   @property
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 69f9c4097..995095070 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -51,78 +51,6 @@
 from tensorflow.io import gfile
 
 
-def unpack_batch(batch, has_weights):
-  """Unpacks a training batch into inputs, targets and weights."""
-  if has_weights:
-    assert len(batch) == 3  # (inputs, targets, weights)
-    return batch
-  else:
-    inputs, targets = batch
-    if isinstance(targets, (list, tuple)):
-      # If weights are not provided, use scalar 1s and rely on broadcasting.
-      weights = [1.0] * len(targets)
-    else:
-      weights = 1.0
-    return inputs, targets, weights
-
-
-def _make_list(predictions, targets, weights):
-  """Make predictions, targets and weights lists, check they match on length."""
-  #  Our models sometimes return predictions in lists, make it a list always.
-  # TODO(lukaszkaiser): make abstractions for nested structures and refactor.
-  if not isinstance(predictions, (list, tuple)):
-    if isinstance(targets, (list, tuple)):
-      raise ValueError("Targets are a list or tuple but predictions are not.")
-    if isinstance(weights, (list, tuple)):
-      raise ValueError("Weights are a list or tuple but predictions are not.")
-    predictions, targets, weights = [predictions], [targets], [weights]
-  if len(predictions) != len(targets):
-    raise ValueError("Predictions and targets have different lengths.")
-  if len(predictions) != len(weights):
-    raise ValueError("Predictions and weights have different lengths.")
-  return list(predictions), list(targets), list(weights)
-
-
-@gin.configurable(blacklist=["inputs", "targets", "weights"])
-def masked_mean(inputs, targets, weights, mask_id=None):
-  """Weighted mean of the inputs, excluding where targets == mask_id."""
-  inputs = [x.astype(np.float32) for x in inputs]
-  # We assume all elements in the list contribute equally.
-  # TODO(lukaszkaiser): remove this assumption (e.g., when masks differ).
-  length = len(inputs)
-  if mask_id is not None:
-    weights = [w * (1.0 - np.equal(t, mask_id).astype(np.float32))
-               for t, w in zip(targets, weights)]
-  weight_sums = [int(t.size) if np.isscalar(w) else np.sum(w)
-                 for w, t in zip(weights, targets)]
-  return sum([np.sum(x * w) / (length * s)
-              for x, w, s in zip(inputs, weights, weight_sums)])
-
-
-def accuracy(batch, model_predictions, has_weights):
-  """Calculate accuracy."""
-  _, targets, weights = unpack_batch(batch, has_weights)
-  model_predictions, targets, weights = _make_list(
-      model_predictions, targets, weights)
-  correct = []
-  for (prediction, target) in zip(model_predictions, targets):
-    predicted_class = np.argmax(prediction, axis=-1)
-    correct.append(np.equal(predicted_class, target))
-  return masked_mean(correct, targets, weights)
-
-
-def neg_log_perplexity(batch, model_predictions, has_weights):
-  """Calculate negative log perplexity."""
-  _, targets, weights = unpack_batch(batch, has_weights)
-  model_predictions, targets, weights = _make_list(
-      model_predictions, targets, weights)
-  xent = []
-  for (prediction, target) in zip(model_predictions, targets):
-    hot_target = layers.one_hot(target, prediction.shape[-1])
-    xent.append(np.sum(prediction * hot_target, axis=-1))
-  return masked_mean(xent, targets, weights)
-
-
 def _stack_inputs_targets_and_get_predictions(inputs_and_targets):
   """Helper to stack inputs and targets and retrieve predictions from output."""
   # Inputs and targets can be lists - we build a flat one to input to the model.
@@ -139,22 +67,6 @@ def _stack_inputs_targets_and_get_predictions(inputs_and_targets):
   return tuple(model_inp), get_pred
 
 
-def loss(params, batch, model_predict, state, rng, has_weights):
-  """Calculate loss."""
-  inputs, targets, weights = unpack_batch(batch, has_weights)
-  model_input, get_preds = _stack_inputs_targets_and_get_predictions(
-      [inputs, targets])
-  # Call model, predictions will be the returned stack, usually consisting of
-  # the prediction tensor and the targets.
-  outputs = model_predict(model_input, params=params, state=state, rng=rng)
-  predictions = get_preds(outputs)
-  predictions, targets, weights = _make_list(predictions, targets, weights)
-  xent = []
-  for (pred, target) in zip(predictions, targets):
-    xent.append(np.sum(pred * layers.one_hot(target, pred.shape[-1]), axis=-1))
-  return - masked_mean(xent, targets, weights), state
-
-
 def log(s, stdout=True):
   logging.info(s)
   if stdout:
@@ -249,49 +161,23 @@ def _print_n_params(opt_state, n_devices, step):
 
 # Metrics to calculate and report.
 _METRICS = {
-    "accuracy": accuracy,
-    "neg_log_perplexity": neg_log_perplexity,
-    "loss": lambda *args, **kwargs: - neg_log_perplexity(*args, **kwargs),
+    "accuracy": layers.AccuracyScalar,
+    "neg_log_perplexity": layers.NegLogPerplexityScalar,
+    "loss": layers.CrossEntropyLossScalar,
 }
 
 
-def evaluate_train_and_eval(step, eval_stream, train_eval_stream,
-                            predict_fn, eval_steps, state, rng, has_weights,
-                            train_sw=None, eval_sw=None, history=None):
-  """Evalaute on train and eval data, and log metrics."""
-  step_log(step, "Evaluation")
-  metrics_list = []
-  for input_stream in [train_eval_stream, eval_stream]:
-    metrics, state = evaluate(  # pylint: disable=g-complex-comprehension
-        itertools.islice(input_stream, eval_steps),
-        predict_fn,
-        _METRICS,
-        state,
-        rng,
-        has_weights)
-    metrics_list.append(metrics)
-  # Unpack in the same order we've iterated over streams in the loop above.
-  train_metrics, eval_metrics = metrics_list  # pylint: disable=unbalanced-tuple-unpacking
-  if train_sw:
-    log_metrics(train_metrics, train_sw, "train", step, history=history)
-  if eval_sw:
-    log_metrics(eval_metrics, eval_sw, "eval", step, history=history)
-  step_log(step, "Finished evaluation")
-  return train_metrics, eval_metrics, state
-
-
-def evaluate(inputs_stream, predict_fn, metric_fns, state, rng, has_weights):
+def evaluation_round(inputs_stream, metric_names, eval_fn, params, state, rng):
   """Evaluate.
 
   Args:
     inputs_stream: iterable of inputs to evaluate on.
-    predict_fn: function from inputs to predictions. params should already be
-      partially applied.
-    metric_fns: dict from metric name to metric function, which takes inputs
-      and predictions and returns a scalar metric value.
-    state: start state for `predict_fn`.
+    metric_names: list of strings, the order in which eval_fn returns metrics.
+    eval_fn: metric function, which takes inputs and predictions (and
+      params, state, rng) and returns a tuple of scalar metric values.
+    params: params for each f in eval_fns.
+    state: state for each f in eval_fns.
     rng: random number generator.
-    has_weights: bool, whether weights are included in the inputs.
 
   Returns:
     metrics: dict from metric name to metric value averaged over the number of
@@ -303,43 +189,16 @@ def evaluate(inputs_stream, predict_fn, metric_fns, state, rng, has_weights):
   for inp in inputs_stream:
     count += 1
     rng, subrng = jax_random.split(rng)
-    model_inp, get_preds = _stack_inputs_targets_and_get_predictions(inp)
-    # Call model, preds will be the returned stack, usually (pred, targets).
-    outputs = predict_fn(model_inp, state=state, rng=subrng)
-    pred = get_preds(outputs)
-    for m, f in six.iteritems(metric_fns):
-      metrics[m] += f(inp, pred, has_weights=has_weights)
+    metric_values = eval_fn(inp, params=params, state=state, rng=subrng)
+    try:
+      metric_values = list(metric_values)
+    except TypeError:
+      metric_values = [float(metric_values)]
+    for m, v in zip(metric_names, metric_values):
+      metrics[m] += v
   return {m: v / count for (m, v) in six.iteritems(metrics)}, state
 
 
-def evaluate_loss_train_and_eval(step, eval_stream, train_eval_stream,
-                                 compute_loss_fn, eval_steps,
-                                 state, rngs, has_weights,
-                                 train_sw=None, eval_sw=None, history=None):
-  """More efficient evaluation that logs only the loss on train & eval data."""
-  assert not has_weights, (
-      "MemoryEfficientTrainer doesn't support has_weights")
-  step_log(step, "Evaluation")
-  train_eval_metrics = []
-  for input_stream in [train_eval_stream, eval_stream]:
-    total = 0.0
-    count = 0.0
-    for inp in itertools.islice(input_stream, eval_steps):
-      loss_values, state, rngs = compute_loss_fn(inp, state, rngs)
-      total += float(numpy.mean(loss_values))
-      count += 1.0
-    metrics = {"loss": total / count}
-    train_eval_metrics.append(metrics)
-  # Unpack in the same order we've iterated over streams in the loop above.
-  train_metrics, eval_metrics = train_eval_metrics  # pylint: disable=unbalanced-tuple-unpacking
-  if train_sw:
-    log_metrics(train_metrics, train_sw, "train", step, history=history)
-  if eval_sw:
-    log_metrics(eval_metrics, eval_sw, "eval", step, history=history)
-  step_log(step, "Finished evaluation")
-  return train_metrics, eval_metrics, state
-
-
 def log_metrics(metrics, summ_writer, log_prefix, step, history=None):
   """Log metrics to summary writer and history."""
   rjust_len = max([0] + [len(name) for name in metrics])
@@ -365,18 +224,6 @@ def get_random_number_generator_and_set_seed(seed=None):
   return jax_random.get_prng(seed)
 
 
-# TODO(trax):
-# * Make configurable:
-#   * loss
-#   * metrics
-# * Training loop callbacks/hooks/...
-# * Save/restore: pickle unsafe. Use np.array.savez + MessagePack?
-# * Move metrics to metrics.py
-# * Setup namedtuples for interfaces (e.g. lr fun constructors can take a
-#   LearningRateInit, metric funs, etc.).
-# * Allow disabling eval
-
-
 def epochs(total_steps, steps_to_skip, epoch_steps):
   """Generates the number of steps in each epoch before reaching total_steps.
 
@@ -411,8 +258,9 @@ def epochs(total_steps, steps_to_skip, epoch_steps):
 
 
 @gin.configurable
-def _jit_predict_fn(model_predict, n_devices, jit=True):
+def _jit_predict_fn(model_predict, metric_fn, n_devices, jit=True):
   """Returns a JIT-compiled predict function (unless jit=False)."""
+  model_predict = layers.Serial([model_predict, metric_fn])
 
   if n_devices == 1:
     return backend.jit(model_predict) if jit else model_predict
@@ -432,8 +280,12 @@ def predict(x, params=(), state=(), rng=None):
     # Need to reduce the [device, per-device-batch, ...] tensors back to
     # a [batch, ...] tensor. The tensors may be nested.
     def combine(x):
-      batch_size = x.shape[0] * x.shape[1]
-      return np.reshape(x, [batch_size] + list(x.shape[2:]))
+      if len(x.shape) > 1:
+        batch_size = x.shape[0] * x.shape[1]
+        return np.reshape(x, [batch_size] + list(x.shape[2:]))
+      # TODO(lukaszkaiser): is returning averages for scalars the right choice?
+      # If it is only scalar, return the average.
+      return np.mean(x, axis=0)
     return layers.nested_map(pred, combine)
 
   return predict
@@ -442,12 +294,17 @@ def combine(x):
 @gin.configurable
 def _jit_update_fn(predict_fn, loss_fn, optimizer, n_devices, jit=True):
   """Returns a (JIT-compiled) function that computes updates for one step."""
+  model_and_loss = layers.Serial([predict_fn, loss_fn])
+  # Gradients are always wrt. the first argument, so putting params first.
+  def model_and_loss_call(params, batch, state, rng):
+    res = model_and_loss(batch, params=params, state=state, rng=rng)
+    return res, model_and_loss.state
   if n_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
     def single_update(i, opt_state, batch, state, rng):
       params, slots, opt_params = opt_state
       rng, subrng = jax_random.split(rng[0])
-      grad_fn = backend.grad(loss_fn, has_aux=True)
-      grads, state = grad_fn(params, batch, predict_fn, state, rng)
+      grad_fn = backend.grad(model_and_loss_call, has_aux=True)
+      grads, state = grad_fn(params, batch, state, rng)
       return optimizer.tree_update(
           i, grads, params, slots, opt_params), state, [subrng]
     return backend.jit(single_update) if jit else single_update
@@ -459,8 +316,8 @@ def mapped_update(i, opt_state, batch, state, rng):
     # We assume all tensors have the first dimension = n_devices.
     params, slots, opt_params = opt_state
     rng, subrng = jax_random.split(rng)
-    grad_fn = backend.grad(loss_fn, has_aux=True)
-    grads, state = grad_fn(params, batch, predict_fn, state, rng)
+    grad_fn = backend.grad(model_and_loss_call, has_aux=True)
+    grads, state = grad_fn(params, batch, state, rng)
     grads = jax.tree_util.tree_map(
         lambda g: lax.psum(g, "batch"), grads)
     return optimizer.tree_update(
@@ -581,13 +438,14 @@ class Trainer(object):
   def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs,
                output_dir=None, random_seed=None, n_devices=None,
                save_steps=None, should_save=True, has_weights=False,
-               nontrainable_param_map=None):
+               nontrainable_param_map=None, mask_id=None):
     if save_steps is None:
       save_steps = []
     self._save_steps = save_steps
     self._should_save = should_save
     self._has_weights = has_weights
-    loss_fn = functools.partial(loss_fn, has_weights=self._has_weights)
+    self._mask_id = mask_id
+    loss_fn = loss_fn(has_weights=has_weights, mask_id=mask_id)
     device_count = jax.lib.xla_bridge.device_count()
     n_devices = n_devices or device_count
     # TODO(lukaszkaiser): remove this restriction when possible.
@@ -636,11 +494,14 @@ def new_opt_state_and_model_state(input_shape, input_dtype, target_shape,
         target_shape = [target_shape]
       full_type = list(input_dtype) + list(target_dtype)
       full_shape = list(input_shape) + list(target_shape)
+      if self._has_weights:
+        full_shape += list(target_shape)
+        full_type += [np.float32 for _ in target_dtype]
       # We need to create a new model instance and not reuse `model_train` here,
       # because `m.initialize` puts cached parameter values in `m` and hence the
       # next call of `m.initialize` will give wrong results.
-      params, state = model(mode="train").initialize_once(full_shape, full_type,
-                                                          rng)
+      m = layers.Serial([model(mode="train"), loss_fn])
+      params, state = m.initialize_once(full_shape, full_type, rng)
       (slots, opt_params) = opt.tree_init(params)
       return (OptState(params, slots, opt_params), state)
     if _is_jit_init():
@@ -653,8 +514,40 @@ def new_opt_state_and_model_state(input_shape, input_dtype, target_shape,
             model_target_shape, self._inputs.target_dtype, init_rng))
 
     # jit model_predict and update so they're fast
-    self._jit_model_predict_eval = _jit_predict_fn(
-        model_predict_eval, n_devices)
+    # TODO(lukaszkaiser): the code below creates a layer computing
+    # multiple metrics from a single model output; re-factor for clarity.
+    dup_layer = layers.Dup3() if self._has_weights else layers.Dup2()
+    def lower(layer):
+      """Apply layer below the current inputs, targets, and possibly weights."""
+      if self._has_weights:
+        # Apply layer below inputs, targets, and loss weights.
+        return layers.Parallel([], [], [], layer)
+      else:
+        # Apply layer below inputs and targets.
+        return layers.Parallel([], [], layer)
+    metrics_layer = []
+    self._metrics = list(sorted(_METRICS.keys()))
+    for i, m in enumerate(reversed(self._metrics)):
+      metric = _METRICS[m](has_weights=self._has_weights, mask_id=self._mask_id)
+      if i != len(self._metrics) - 1:
+        metrics_layer.append(dup_layer)
+        metrics_layer.append(lower(metric))
+      else:
+        metrics_layer.append(metric)
+    # TODO(lukaszkaiser): clean this up once layer API stabilizes.
+    # For now, we need to initialize metric layers somehow, so here we go.
+    # We assume that they do not have any parameters, so this is a dummy.
+    dummy_shape = ((1, 2), (1,), (1,)) if self._has_weights else ((1, 2), (1,))
+    dummy_type = [np.float32] * (3 if self._has_weights else 2)
+    metrics_layer = layers.Serial(metrics_layer)
+    metrics_params, metrics_state = metrics_layer.initialize_once(
+        dummy_shape, tuple(dummy_type), init_rng)
+    self._metrics_params = layers.nested_map(
+        metrics_params, self._maybe_replicate)
+    self._metrics_state = layers.nested_map(
+        metrics_state, self._maybe_replicate)
+    self._jit_eval = _jit_predict_fn(
+        model_predict_eval, metrics_layer, n_devices)
     self._jit_update_fn = _jit_update_fn(model_train, loss_fn, opt, n_devices)
 
     self._model_train = model_train
@@ -869,19 +762,25 @@ def train_epoch(self, epoch_steps, eval_steps):
   def evaluate(self, eval_steps):
     """Evaluate the model and log metrics."""
     _, rng = jax_random.split(self._rngs[0])
-    _, _, self._model_state = evaluate_train_and_eval(
-        step=self._step,
-        eval_stream=self._eval_stream,
-        train_eval_stream=self._train_eval_stream,
-        predict_fn=functools.partial(self._jit_model_predict_eval,
-                                     params=self._opt_state[0]),
-        eval_steps=eval_steps,
-        state=self._model_state,
-        rng=rng,
-        train_sw=self._train_sw,
-        eval_sw=self._eval_sw,
-        history=self._history,
-        has_weights=self._has_weights)
+    # TODO(lukaszkaiser): both model state and parameters by default include
+    # the loss layer. Currently, we access the pure-model parameters by just
+    # indexing, [0] here. But we should make it more explicit in a better API.
+    params = (self._opt_state[0][0], self._metrics_params)
+    state = (self._model_state[0], self._metrics_state)
+    step_log(self._step, "Evaluation")
+    train_eval_slice = itertools.islice(self._train_eval_stream, eval_steps)
+    train_metrics, _ = evaluation_round(
+        train_eval_slice, self._metrics, self._jit_eval, params, state, rng)
+    if self._train_sw:
+      log_metrics(train_metrics, self._train_sw, "train",
+                  self._step, history=self._history)
+    eval_slice = itertools.islice(self._eval_stream, eval_steps)
+    eval_metrics, _ = evaluation_round(
+        eval_slice, self._metrics, self._jit_eval, params, state, rng)
+    if self._eval_sw:
+      log_metrics(eval_metrics, self._eval_sw, "eval",
+                  self._step, history=self._history)
+    step_log(self._step, "Finished evaluation")
 
     # Save the optimizer params in the history
     for (name, value) in self.nontrainable_params.items():
@@ -899,9 +798,9 @@ def save_computation_graphs(self, save_backward_graph):
     output_dir = self._output_dir
     if self._n_devices > 1:
       next_train_batch = reshape_by_device(next_train_batch, self._n_devices)
-    params = self._opt_state[0]
+    params = self._opt_state[0][0]
     forward_computation = jax.xla_computation(self._model_predict_eval)(
-        next_train_batch, params=params, state=self._model_state,
+        next_train_batch, params=params, state=self._model_state[0],
         rng=self._rngs[0])
     with gfile.GFile(os.path.join(output_dir, "forward.txt"), "w") as f:
       f.write(forward_computation.GetHloText())
@@ -936,20 +835,33 @@ def __init__(self, *args, **kwargs):
         "MemoryEfficientTrainer doesn't support has_weights")
 
   def evaluate(self, eval_steps):
-    # Evaluate only the loss function (a more efficient, jitted, implementation)
-    _, _, self._model_state = evaluate_loss_train_and_eval(
-        step=self._step,
-        eval_stream=self._eval_stream,
-        train_eval_stream=self._train_eval_stream,
-        compute_loss_fn=functools.partial(self._jit_compute_loss,
-                                          self._opt_state),
-        eval_steps=eval_steps,
-        state=self._model_state,
-        rngs=self._rngs,
-        has_weights=self._has_weights,
-        train_sw=self._train_sw,
-        eval_sw=self._eval_sw,
-        history=self._history)
+    """Evaluate only the loss function (efficient, jitted, implementation)."""
+    assert not self._has_weights, (
+        "MemoryEfficientTrainer doesn't support has_weights")
+    step = self._step
+    rngs = self._rngs
+    state = self._model_state
+    history = self._history
+    compute_loss_fn = functools.partial(self._jit_compute_loss,
+                                        self._opt_state)
+    step_log(step, "Evaluation")
+    train_eval_metrics = []
+    for input_stream in [self._train_eval_stream, self._eval_stream]:
+      total = 0.0
+      count = 0.0
+      for inp in itertools.islice(input_stream, eval_steps):
+        loss_values, state, rngs = compute_loss_fn(inp, state, rngs)
+        total += float(numpy.mean(loss_values))
+        count += 1.0
+      metrics = {"loss": total / count}
+      train_eval_metrics.append(metrics)
+    # Unpack in the same order we've iterated over streams in the loop above.
+    train_metrics, eval_metrics = train_eval_metrics  # pylint: disable=unbalanced-tuple-unpacking
+    if self._train_sw:
+      log_metrics(train_metrics, self._train_sw, "train", step, history=history)
+    if self._eval_sw:
+      log_metrics(eval_metrics, self._eval_sw, "eval", step, history=history)
+    step_log(step, "Finished evaluation")
 
   def save_computation_graphs(self, save_backward_graph):
     # TODO(kitaev): implement saving graphs while making sure that no op-by-op
@@ -961,7 +873,7 @@ def save_computation_graphs(self, save_backward_graph):
 @gin.configurable(blacklist=["output_dir"])
 def train(output_dir,
           model=gin.REQUIRED,
-          loss_fn=loss,
+          loss_fn=layers.CrossEntropyLossScalar,
           inputs=trax_inputs.inputs,
           optimizer=trax_opt.Adafactor,
           lr_schedule=lr.MultifactorSchedule,
@@ -975,7 +887,8 @@ def train(output_dir,
           save_graphs=True,
           save_backward_graph=False,
           has_weights=False,
-          nontrainable_param_map=None):
+          nontrainable_param_map=None,
+          mask_id=None):
   """Train the model on the inputs.
 
   Args:
@@ -1002,14 +915,18 @@ def train(output_dir,
     has_weights: bool, whether weights are included in the inputs.
     nontrainable_param_map: dict, mapping from model nontrainable parameter
       names to control names in PolicySchedule.
+    mask_id: id to mask out (None by default).
+
   Returns:
     trax.State
   """
+  # TODO(lukaszkaiser): remove has_weights and mask_id later (configure loss).
   trainer = trainer_class(model, loss_fn, optimizer, lr_schedule, inputs,
                           output_dir,
                           random_seed=random_seed, n_devices=n_devices,
                           save_steps=save_steps, has_weights=has_weights,
-                          nontrainable_param_map=nontrainable_param_map)
+                          nontrainable_param_map=nontrainable_param_map,
+                          mask_id=mask_id)
 
   epoch_steps = [train_steps]  # Only training if eval_frequency is 0 or None
   if eval_frequency and eval_steps > 0:
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
index a10266e46..8c9041aab 100644
--- a/tensor2tensor/trax/trax_test.py
+++ b/tensor2tensor/trax/trax_test.py
@@ -27,7 +27,6 @@
 from jax import test_util  # pylint: disable=unused-import
 from jax.config import config
 from jax.lib import xla_bridge
-import numpy as onp
 
 from tensor2tensor.trax import backend
 from tensor2tensor.trax import inputs as inputs_lib
@@ -233,7 +232,7 @@ def test_reset_twice(self, backend_name):
 
       trainer = trax.Trainer(
           model=model_fn,
-          loss_fn=trax.loss,
+          loss_fn=layers.CrossEntropyLossScalar,
           optimizer=trax_opt.SM3,
           lr_schedule=lr.MultifactorSchedule,
           inputs=inputs,
@@ -264,48 +263,6 @@ def test_skips_part_of_epoch(self):
     self.assertEqual(list(epoch_steps), [1, 2])
 
 
-MASKED_MEAN_TEST_BACKENDS = ["numpy"]
-
-
-class MaskedMeanTest(test.TestCase, parameterized.TestCase):
-
-  @parameterized.parameters(MASKED_MEAN_TEST_BACKENDS)
-  def test_computes_basic_mean(self, backend_name):
-    with backend.use_backend(backend_name):
-      inputs = [np.array([1, 2, 3])]
-      targets = [np.zeros(3)]
-      weights = [1]
-      mean = trax.masked_mean(inputs, targets, weights)
-      onp.testing.assert_allclose(mean, 2)
-
-  @parameterized.parameters(MASKED_MEAN_TEST_BACKENDS)
-  def test_computes_mean_with_weights(self, backend_name):
-    with backend.use_backend(backend_name):
-      inputs = [np.array([1, 2, 3])]
-      targets = [np.zeros(3)]
-      weights = [np.array([3, 1, 0])]
-      mean = trax.masked_mean(inputs, targets, weights)
-      onp.testing.assert_allclose(mean, 1.25)
-
-  @parameterized.parameters(MASKED_MEAN_TEST_BACKENDS)
-  def test_computes_mean_with_mask(self, backend_name):
-    with backend.use_backend(backend_name):
-      inputs = [np.array([1, 2, 3])]
-      targets = [np.array([1, 0, 0])]
-      weights = [1]
-      mean = trax.masked_mean(inputs, targets, weights, mask_id=1)
-      onp.testing.assert_allclose(mean, 2.5)
-
-  @parameterized.parameters(MASKED_MEAN_TEST_BACKENDS)
-  def test_computes_mean_with_weights_and_mask(self, backend_name):
-    with backend.use_backend(backend_name):
-      inputs = [np.array([1, 2, 4])]
-      targets = [np.array([1, 0, 0])]
-      weights = [np.array([10, 4, 1])]
-      mean = trax.masked_mean(inputs, targets, weights, mask_id=1)
-      onp.testing.assert_allclose(mean, 2.4)
-
-
 if __name__ == "__main__":
   config.config_with_absl()
   test.main()

From 7c2012c932d09f78f0804ffd220ce0edfc7a9088 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 1 Oct 2019 15:03:54 -0700
Subject: [PATCH 2516/2720] Delete grpcio from setup.py, since we deleted the
 gRPC code.

PiperOrigin-RevId: 272303379
---
 setup.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/setup.py b/setup.py
index cfc678d98..55cd65e54 100644
--- a/setup.py
+++ b/setup.py
@@ -40,7 +40,6 @@
         'gevent',
         'gin-config',
         'google-api-python-client',
-        'grpcio',
         'gunicorn',
         'gym',
         'h5py',

From 71e10f397694477b8230b482661f5f8833d26a12 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 1 Oct 2019 16:56:45 -0700
Subject: [PATCH 2517/2720] Fix a corner case in BoxSpaceSerializer, causing a
 bug in SimPLe training data generation.

PiperOrigin-RevId: 272326328
---
 tensor2tensor/trax/rl/space_serializer.py     |  4 ++-
 .../trax/rl/space_serializer_test.py          | 28 ++++++++++++++-----
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/tensor2tensor/trax/rl/space_serializer.py b/tensor2tensor/trax/rl/space_serializer.py
index 3c4809c23..b6c166ce6 100644
--- a/tensor2tensor/trax/rl/space_serializer.py
+++ b/tensor2tensor/trax/rl/space_serializer.py
@@ -131,7 +131,9 @@ def serialize(self, data):
     digits = []
     for digit_index in range(-1, -self._precision - 1, -1):
       threshold = self._vocab_size ** digit_index
-      digit = np.array(array / threshold).astype(np.int32) % self._vocab_size
+      digit = np.array(array / threshold).astype(np.int32)
+      # For the corner case of x == high.
+      digit[digit == self._vocab_size] -= 1
       digits.append(digit)
       array -= digit * threshold
     digits = np.stack(digits, axis=-1)
diff --git a/tensor2tensor/trax/rl/space_serializer_test.py b/tensor2tensor/trax/rl/space_serializer_test.py
index eed62b3d4..b80bc85f6 100644
--- a/tensor2tensor/trax/rl/space_serializer_test.py
+++ b/tensor2tensor/trax/rl/space_serializer_test.py
@@ -29,14 +29,16 @@
 
 class BoxSpaceSerializerTest(test.TestCase):
 
-  def _make_space_and_serializer(self, low=-10, high=10, shape=(2,)):
-    # Enough precision to represent float32s accurately.
-    gin.bind_parameter("BoxSpaceSerializer.precision", 4)
+  def _make_space_and_serializer(
+      self, low=-10, high=10, shape=(2,),
+      # Weird vocab_size to test that it doesn't only work with powers of 2.
+      vocab_size=257,
+      # Enough precision to represent float32s accurately.
+      precision=4,
+  ):
+    gin.bind_parameter("BoxSpaceSerializer.precision", precision)
     space = gym.spaces.Box(low=low, high=high, shape=shape)
-    serializer = space_serializer.create(
-        space,
-        # Weird vocab_size to test that it doesn't only work with powers of 2.
-        vocab_size=257)
+    serializer = space_serializer.create(space, vocab_size=vocab_size)
     return (space, serializer)
 
   def _sample_batch(self, space):
@@ -79,6 +81,18 @@ def test_significance_map(self):
     np.testing.assert_array_equal(
         serializer.significance_map, [0, 1, 2, 3, 0, 1, 2, 3])
 
+  def test_serializes_boundaries(self):
+    vocab_size = 256
+    precision = 4
+    (_, serializer) = self._make_space_and_serializer(
+        low=-1, high=1, shape=(1,), vocab_size=vocab_size, precision=precision,
+    )
+    input_array = np.array([[-1, 1]])
+    representation = serializer.serialize(input_array)
+    np.testing.assert_array_equal(
+        representation, [[0] * precision + [vocab_size - 1] * precision]
+    )
+
 
 class DiscreteSpaceSerializerTest(test.TestCase):
 

From 76872f613d3ab9d79c768850fe370bfe88014e89 Mon Sep 17 00:00:00 2001
From: Shawn Simister <simister@google.com>
Date: Tue, 1 Oct 2019 17:24:26 -0700
Subject: [PATCH 2518/2720] Extend Neural Stack to support Deque by reading and
 writing in both directions.

PiperOrigin-RevId: 272331503
---
 tensor2tensor/models/research/neural_stack.py | 149 ++++++++++++++--
 .../models/research/neural_stack_test.py      | 159 +++++++++++++++++-
 2 files changed, 288 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/models/research/neural_stack.py b/tensor2tensor/models/research/neural_stack.py
index 5c138e026..8909e9748 100644
--- a/tensor2tensor/models/research/neural_stack.py
+++ b/tensor2tensor/models/research/neural_stack.py
@@ -140,8 +140,12 @@ def get_read_mask(self, read_head_index):
     Returns:
       A tf.float32 tensor of shape [1, 1, memory_size, memory_size]
     """
-    return tf.expand_dims(
-        common_layers.mask_pos_lt(self._memory_size, self._memory_size), axis=0)
+    if read_head_index == 0:
+      return tf.expand_dims(
+          common_layers.mask_pos_lt(self._memory_size, self._memory_size),
+          axis=0)
+    else:
+      raise ValueError("Read head index must be 0 for stack.")
 
   def get_write_head_offset(self, write_head_index):
     """Lookup the offset to shift the write head at each step.
@@ -157,7 +161,10 @@ def get_write_head_offset(self, write_head_index):
     Returns:
       An integer offset to move the write head at each step.
     """
-    return 1
+    if write_head_index == 0:
+      return 1
+    else:
+      raise ValueError("Write head index must be 0 for stack.")
 
   def add_scalar_projection(self, name, size):
     """A helper function for mapping scalar controller outputs.
@@ -209,8 +216,7 @@ def build_controller(self):
       self.rnn = tf.contrib.rnn.BasicRNNCell(self._num_units)
       self._input_proj = self.add_variable(
           "input_projection_weights",
-          shape=[(self._embedding_size * self._num_read_heads) +
-                 (self._embedding_size * self._num_write_heads),
+          shape=[self._embedding_size * (self._num_read_heads + 1),
                  self._num_units],
           dtype=self.dtype)
       self._input_bias = self.add_variable(
@@ -224,7 +230,7 @@ def build_controller(self):
       self._value_proj, self._value_bias = self.add_vector_projection(
           "value", self._num_write_heads)
       self._output_proj, self._output_bias = self.add_vector_projection(
-          "output", self._num_read_heads)
+          "output", 1)
 
   def build(self, _):
     """Build the controller.
@@ -255,14 +261,14 @@ def get_controller_shape(self, batch_size):
         # state
         [batch_size, self._num_units])
 
-  def call_controller(self, inputs, read_values, prev_state, batch_size):
+  def call_controller(self, input_value, read_values, prev_state, batch_size):
     """Make a call to the neural stack controller.
 
     See Section 3.1 of Grefenstette et al., 2015.
 
     Args:
-      inputs: The inputs to the neural stack cell should be a tf.float32 tensor
-        with shape [batch_size, num_write_heads, embedding_size]
+      input_value: The input to the neural stack cell should be a tf.float32
+        tensor with shape [batch_size, 1, embedding_size]
       read_values: The values of the read heads at the previous timestep.
       prev_state: The hidden state from the previous time step.
       batch_size: The size of the current batch of input values.
@@ -271,10 +277,10 @@ def call_controller(self, inputs, read_values, prev_state, batch_size):
       A tuple of outputs and the new NeuralStackControllerInterface.
     """
     with tf.name_scope("controller"):
-      # Concatenate the current input values with the read value from the
+      # Concatenate the current input value with the read values from the
       # previous timestep before feeding them into the controller.
       controller_inputs = tf.concat([
-          tf.contrib.layers.flatten(inputs),
+          tf.contrib.layers.flatten(input_value),
           tf.contrib.layers.flatten(read_values),
       ], axis=1)
 
@@ -331,7 +337,7 @@ def call(self, inputs, prev_state):
     # Always write input values to memory regardless of push strength.
     # See Equation-1 in Grefenstette et al., 2015.
     new_memory_values = prev_state.memory_values + tf.reduce_sum(
-        tf.expand_dims(controller_output.write_values, axis=1) *
+        tf.expand_dims(controller_output.write_values, axis=2) *
         prev_state.write_strengths,
         axis=1)
 
@@ -340,9 +346,11 @@ def call(self, inputs, prev_state):
     # See Equation-2 in Grefenstette et al., 2015.
     new_read_strengths = prev_state.read_strengths
     for h in range(self._num_read_heads - 1, -1, -1):
-      new_read_strengths = tf.nn.relu(
-          new_read_strengths -
-          tf.nn.relu(controller_output.pop_strengths - tf.expand_dims(
+      new_read_strengths = tf.nn.relu(new_read_strengths - tf.nn.relu(
+          tf.slice(controller_output.pop_strengths,
+                   [0, h, 0, 0],
+                   [-1, 1, -1, -1]) -
+          tf.expand_dims(
               tf.reduce_sum(new_read_strengths * self.get_read_mask(h), axis=2),
               axis=3)))
 
@@ -403,8 +411,78 @@ def get_read_mask(self, read_head_index):
     Returns:
       A tf.float32 tensor of shape [1, 1, memory_size, memory_size].
     """
+    if read_head_index == 0:
+      return tf.expand_dims(
+          common_layers.mask_pos_gt(self._memory_size, self._memory_size),
+          axis=0)
+    else:
+      raise ValueError("Read head index must be 0 for queue.")
+
+
+class NeuralDequeCell(NeuralStackCell):
+  """An subclass of the NeuralStackCell which reads/writes in both directions.
+
+  See section 3.3 of Grefenstette et al., 2015.
+  """
+
+  def __init__(self, num_units, memory_size, embedding_size, reuse=None):
+    # Override constructor to set 2 read/write heads.
+    super(NeuralDequeCell, self).__init__(num_units,
+                                          memory_size,
+                                          embedding_size,
+                                          num_read_heads=2,
+                                          num_write_heads=2,
+                                          reuse=reuse)
+
+  def get_read_mask(self, read_head_index):
+    if read_head_index == 0:
+      # Use the same read mask as the queue for the bottom of the deque.
+      return tf.expand_dims(
+          common_layers.mask_pos_gt(self._memory_size, self._memory_size),
+          axis=0)
+    elif read_head_index == 1:
+      # Use the same read mask as the stack for the top of the deque.
+      return tf.expand_dims(
+          common_layers.mask_pos_lt(self._memory_size, self._memory_size),
+          axis=0)
+    else:
+      raise ValueError("Read head index must be either 0 or 1 for deque.")
+
+  def get_write_head_offset(self, write_head_index):
+    if write_head_index == 0:
+      # Move the bottom write position back at each timestep.
+      return -1
+    elif write_head_index == 1:
+      # Move the top write position forward at each timestep.
+      return 1
+    else:
+      raise ValueError("Write head index must be 0 or 1 for deque.")
+
+  def initialize_write_strengths(self, batch_size):
+    """Initialize write strengths which write in both directions.
+
+    Unlike in Grefenstette et al., It's writing out from the center of the
+    memory so that it doesn't need to shift the entire memory forward at each
+    step.
+
+    Args:
+      batch_size: The size of the current batch.
+
+    Returns:
+      A tf.float32 tensor of shape [num_write_heads, memory_size, 1].
+    """
+    memory_center = self._memory_size // 2
     return tf.expand_dims(
-        common_layers.mask_pos_gt(self._memory_size, self._memory_size), axis=0)
+        tf.concat([
+            # The write strength for the deque bottom.
+            # Should be shifted back at each timestep.
+            tf.one_hot([[memory_center - 1]] * batch_size,
+                       depth=self._memory_size, dtype=tf.float32),
+            # The write strength for the deque top.
+            # Should be shifted forward at each timestep.
+            tf.one_hot([[memory_center]] * batch_size,
+                       depth=self._memory_size, dtype=tf.float32)
+        ], axis=1), axis=3)
 
 
 @registry.register_model
@@ -502,6 +580,25 @@ def cell(self, hidden_size):
                            self._hparams.embedding_size)
 
 
+@registry.register_model
+class NeuralDequeModel(NeuralStackModel):
+  """Subclass of NeuralStackModel which implements a double-ended queue.
+  """
+
+  def cell(self, hidden_size):
+    """Build a NeuralDequeCell instead of a NeuralStackCell.
+
+    Args:
+      hidden_size: The hidden size of the cell.
+
+    Returns:
+      A new NeuralDequeCell with the given hidden size.
+    """
+    return NeuralDequeCell(hidden_size,
+                           self._hparams.memory_size,
+                           self._hparams.embedding_size)
+
+
 @registry.register_hparams
 def lstm_transduction():
   """HParams for LSTM base on transduction tasks."""
@@ -540,3 +637,23 @@ def neural_stack():
   hparams.add_hparam("embedding_size", 64)
   hparams.hidden_size = hparams.embedding_size
   return hparams
+
+
+@registry.register_hparams
+def neural_deque():
+  """HParams for neural deques."""
+  hparams = common_hparams.basic_params1()
+  hparams.daisy_chain_variables = False
+  hparams.batch_size = 10
+  hparams.clip_grad_norm = 1.0
+  hparams.initializer = "uniform_unit_scaling"
+  hparams.initializer_gain = 1.0
+  hparams.optimizer = "RMSProp"
+  hparams.learning_rate = 0.0001
+  hparams.weight_decay = 0.0
+
+  hparams.add_hparam("controller_layer_sizes", [256, 512])
+  hparams.add_hparam("memory_size", 256)
+  hparams.add_hparam("embedding_size", 64)
+  hparams.hidden_size = hparams.embedding_size
+  return hparams
diff --git a/tensor2tensor/models/research/neural_stack_test.py b/tensor2tensor/models/research/neural_stack_test.py
index e32001580..a43ba6e23 100644
--- a/tensor2tensor/models/research/neural_stack_test.py
+++ b/tensor2tensor/models/research/neural_stack_test.py
@@ -83,10 +83,10 @@ def assert_controller_shapes(test, controller_outputs, controller_shapes):
     test.assertEqual(shape, output.shape, "%s shapes don't match" % name)
 
 
-def assert_cell_shapes(test, controller_outputs, zero_state):
-  for name, output, state in zip(controller_outputs._fields, controller_outputs,
-                                 zero_state):
-    test.assertEqual(state.shape, output.shape, "%s shapes don't match" % name)
+def assert_cell_shapes(test, output_state, zero_state):
+  for name, output, zero in zip(output_state._fields, output_state,
+                                zero_state):
+    test.assertEqual(zero.shape, output.shape, "%s shapes don't match" % name)
 
 
 class NeuralStackCellTest(tf.test.TestCase):
@@ -128,6 +128,11 @@ def test_cell_shapes(self):
                                         [[0.0, 0.0, 0.0]]]))
   def test_push_pop(self):
     """Test pushing a popping from a NeuralStackCell.
+
+    The sequence of operations is:
+      push([1.0, 0.0, 0.0])
+      push([0.0, 1.0, 0.0])
+      pop()
     """
     input_values = np.array([[[[1.0, 0.0, 0.0]],
                               [[0.0, 1.0, 0.0]],
@@ -191,6 +196,11 @@ class NeuralQueueCellTest(tf.test.TestCase):
                                         [[0.0, 0.0, 0.0]]]))
   def test_enqueue_dequeue(self):
     """Test enqueueing a dequeueing from a NeuralQueueCell.
+
+    The sequence of operations is:
+      enqueue([1.0, 0.0, 0.0])
+      enqueue([0.0, 1.0, 0.0])
+      dequeue()
     """
     input_values = np.array([[[[1.0, 0.0, 0.0]],
                               [[0.0, 1.0, 0.0]],
@@ -237,6 +247,147 @@ def test_enqueue_dequeue(self):
       self.assertAllClose(expected_front, queue_front)
 
 
+class NeuralDequeCellTest(tf.test.TestCase):
+
+  def test_cell_shapes(self):
+    """Check that all the NeuralStackCell tensor shapes are correct.
+    """
+    batch_size = 5
+    embedding_size = 4
+    memory_size = 12
+    num_units = 8
+
+    deque = neural_stack.NeuralDequeCell(num_units, memory_size, embedding_size)
+    deque.build(None)
+
+    self.assertEqual([1, 1, memory_size, memory_size],
+                     deque.get_read_mask(0).shape)
+    self.assertEqual([1, 1, memory_size, memory_size],
+                     deque.get_read_mask(1).shape)
+
+    deque_input = tf.zeros([batch_size, 1, embedding_size], dtype=tf.float32)
+    zero_state = deque.zero_state(batch_size, tf.float32)
+    (outputs, (deque_next_state)) = deque.call(deque_input, zero_state)
+
+    # Make sure that deque output shapes match deque input shapes
+    self.assertEqual(outputs.shape, deque_input.shape)
+
+    assert_cell_shapes(self, deque_next_state, zero_state)
+
+  @mock.patch.object(neural_stack.NeuralDequeCell, "build_controller",
+                     build_fake_controller)
+  @mock.patch.object(neural_stack.NeuralDequeCell, "call_controller",
+                     call_fake_controller(
+                         push_values=[[[[1.0]], [[0.0]]],
+                                      [[[1.0]], [[0.0]]],
+                                      [[[1.0]], [[0.0]]],
+                                      [[[0.0]], [[1.0]]],
+                                      [[[0.0]], [[0.0]]],
+                                      [[[0.0]], [[0.0]]]],
+                         pop_values=[[[[0.0]], [[0.0]]],
+                                     [[[0.0]], [[0.0]]],
+                                     [[[0.0]], [[0.0]]],
+                                     [[[0.0]], [[0.0]]],
+                                     [[[0.0]], [[1.0]]],
+                                     [[[0.0]], [[1.0]]]],
+                         write_values=[[[1.0, 0.0, 0.0, 0.0],
+                                        [0.0, 0.0, 0.0, 0.0]],
+                                       [[0.0, 1.0, 0.0, 0.0],
+                                        [0.0, 0.0, 0.0, 0.0]],
+                                       [[0.0, 0.0, 1.0, 0.0],
+                                        [0.0, 0.0, 0.0, 0.0]],
+                                       [[0.0, 0.0, 0.0, 0.0],
+                                        [0.0, 0.0, 0.0, 1.0]],
+                                       [[0.0, 0.0, 0.0, 0.0],
+                                        [0.0, 0.0, 0.0, 0.0]],
+                                       [[0.0, 0.0, 0.0, 0.0],
+                                        [0.0, 0.0, 0.0, 0.0]]],
+                         output_values=[[[0.0, 0.0, 0.0, 0.0]],
+                                        [[0.0, 0.0, 0.0, 0.0]],
+                                        [[0.0, 0.0, 0.0, 0.0]],
+                                        [[0.0, 0.0, 0.0, 0.0]],
+                                        [[0.0, 0.0, 0.0, 0.0]],
+                                        [[0.0, 0.0, 0.0, 0.0]]]))
+  def test_enqueue_dequeue(self):
+    """Test enqueueing a dequeueing from a NeuralDequeCell.
+
+    The sequence of operations is:
+      enqueue_bottom([1.0, 0.0, 0.0, 0.0])
+      enqueue_bottom([0.0, 1.0, 0.0, 0.0])
+      enqueue_bottom([0.0, 0.0, 1.0, 0.0])
+      enqueue_top([0.0, 0.0, 0.0, 1.0])
+      dequeue_top()
+      dequeue_top()
+    """
+    input_values = np.array([[[[1.0, 0.0, 0.0, 0.0]],
+                              [[0.0, 1.0, 0.0, 0.0]],
+                              [[0.0, 0.0, 1.0, 0.0]],
+                              [[0.0, 0.0, 0.0, 1.0]],
+                              [[0.0, 0.0, 0.0, 0.0]],
+                              [[0.0, 0.0, 0.0, 0.0]]]])
+
+    expected_values = np.array([[[0.0, 0.0, 0.0, 0.0],
+                                 [0.0, 0.0, 0.0, 0.0],
+                                 [0.0, 0.0, 0.0, 0.0],
+                                 [0.0, 0.0, 1.0, 0.0],
+                                 [0.0, 1.0, 0.0, 0.0],
+                                 [1.0, 0.0, 0.0, 0.0],
+                                 [0.0, 0.0, 0.0, 0.0],
+                                 [0.0, 0.0, 0.0, 0.0],
+                                 [0.0, 0.0, 0.0, 0.0],
+                                 [0.0, 0.0, 0.0, 1.0],
+                                 [0.0, 0.0, 0.0, 0.0],
+                                 [0.0, 0.0, 0.0, 0.0]]])
+
+    expected_read_strengths = np.array([[[[0.0], [0.0], [0.0], [1.0], [1.0],
+                                          [0.0], [0.0], [0.0], [0.0], [0.0],
+                                          [0.0], [0.0]]]])
+
+    expected_write_strengths = np.array([[[[0.0], [0.0], [0.0], [0.0], [0.0],
+                                           [0.0], [0.0], [0.0], [0.0], [0.0],
+                                           [0.0], [1.0]],
+                                          [[1.0], [0.0], [0.0], [0.0], [0.0],
+                                           [0.0], [0.0], [0.0], [0.0], [0.0],
+                                           [0.0], [0.0]]]])
+
+    expected_read_values = np.array([[[0.0, 0.0, 1.0, 0.0],
+                                      [0.0, 1.0, 0.0, 0.0]]])
+
+    batch_size = input_values.shape[0]
+    memory_size = input_values.shape[1] * 2
+    embedding_size = input_values.shape[3]
+    num_units = 8
+
+    deque = neural_stack.NeuralDequeCell(num_units, memory_size, embedding_size)
+    rnn_input = tf.constant(input_values, dtype=tf.float32)
+
+    deque_zero_state = tf.zeros([batch_size, num_units])
+    controller_outputs = deque.call_controller(None, None,
+                                               deque_zero_state,
+                                               batch_size)
+    assert_controller_shapes(self, controller_outputs,
+                             deque.get_controller_shape(batch_size))
+
+    (outputs, state) = tf.nn.dynamic_rnn(cell=deque,
+                                         inputs=rnn_input,
+                                         time_major=False,
+                                         dtype=tf.float32)
+
+    with self.test_session() as sess:
+      sess.run(tf.global_variables_initializer())
+      _, state_vals = sess.run([outputs, state])
+      (_, read_values,
+       memory_values,
+       read_strengths,
+       write_strengths) = state_vals
+
+      print(read_values)
+      self.assertAllClose(expected_values, memory_values)
+      self.assertAllClose(expected_write_strengths, write_strengths)
+      self.assertAllClose(expected_read_strengths, read_strengths)
+      self.assertAllClose(expected_read_values, read_values)
+
+
 class NeuralStackModelTest(tf.test.TestCase):
 
   def test_model_shapes(self):

From 2df12d5d2054178b18f6faabf7ebecd9a0177284 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 1 Oct 2019 18:19:47 -0700
Subject: [PATCH 2519/2720] Add a bin_len parametrization to
 TimeBinCausalAttention and add padding to support variable-length sequences.

Both parametrizations are allowed, so no existing configs need changing.

PiperOrigin-RevId: 272339876
---
 tensor2tensor/trax/layers/attention.py      | 38 +++++++++++++++++----
 tensor2tensor/trax/layers/attention_test.py | 16 +++++++++
 2 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index c1f556360..2ec8cf86a 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -18,7 +18,9 @@
 from __future__ import division
 from __future__ import print_function
 
+import math
 import random
+
 import jax
 import numpy as onp
 
@@ -657,8 +659,12 @@ def body_fun(vals):  # pylint: disable=invalid-name
 class TimeBinCausalAttention(BaseCausalAttention):
   """Causal attention where only nearby chunks of items attend to each other."""
 
-  def __init__(self, dropout, mode, n_bins=64, share_qk=False):
+  def __init__(self, dropout, mode, bin_length=None, n_bins=None,
+               share_qk=False):
     super(TimeBinCausalAttention, self).__init__()
+    if (bin_length is None) == (n_bins is None):
+      raise ValueError('Exactly one of {bin_length, n_bins} must be set.')
+    self.bin_length = bin_length
     self.n_bins = n_bins
     self._share_qk = share_qk
     if dropout >= 1.0:
@@ -683,21 +689,41 @@ def make_unit_length(self, x, epsilon=1e-6):
     norm_inputs = x / np.sqrt(variance + epsilon)
     return norm_inputs
 
+  def _pad_inputs(self, inputs):
+    seq_len = inputs[0].shape[-2]
+    n_bins = self.n_bins
+    bin_length = self.bin_length
+    if n_bins is None:
+      n_bins = int(math.ceil(seq_len / bin_length))
+    else:
+      bin_length = int(math.ceil(seq_len / n_bins))
+    pad_len = n_bins * bin_length - seq_len
+
+    def pad_input(x):
+      pad_widths = [(0, 0)] * len(x.shape)
+      pad_widths[-2] = (0, pad_len)  # Padding on axis=-2
+      return np.pad(x, pad_widths, mode='constant',
+                    constant_values=x.dtype.type(0))
+
+    padded_inputs = tuple(map(pad_input, inputs))
+    return (padded_inputs, seq_len, n_bins)
+
   def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
     del params, kwargs
+    (inputs, original_len, n_bins) = self._pad_inputs(inputs)
     q, k, v = inputs
     seqlen = q.shape[-2]
     # q/k/v are n_batch*n_heads, seqlen, d_head
     t = jax.lax.tie_in(q, np.arange(seqlen))
 
     # Split off a "bin" axis for chunks of consecutive items.
-    bq_t = np.reshape(t, (self.n_bins, -1))
-    bq = np.reshape(q, (q.shape[0], self.n_bins, -1, q.shape[-1]))
+    bq_t = np.reshape(t, (n_bins, -1))
+    bq = np.reshape(q, (q.shape[0], n_bins, -1, q.shape[-1]))
     if self._share_qk:
       bk = self.make_unit_length(bq)
     else:
-      bk = np.reshape(k, (k.shape[0], self.n_bins, -1, k.shape[-1]))
-    bv = np.reshape(v, (v.shape[0], self.n_bins, -1, v.shape[-1]))
+      bk = np.reshape(k, (k.shape[0], n_bins, -1, k.shape[-1]))
+    bv = np.reshape(v, (v.shape[0], n_bins, -1, v.shape[-1]))
 
     # Allow each chunk to attend within itself, and also one chunk back.
     def look_one_back(x):
@@ -742,7 +768,7 @@ def look_one_back(x):
 
     output = np.reshape(bo, (bo.shape[0], -1, bo.shape[-1]))
     assert output.shape == v.shape
-    return output, state
+    return output[..., :original_len, :], state
 
 
 class LSHCausalAttention(BaseCausalAttention):
diff --git a/tensor2tensor/trax/layers/attention_test.py b/tensor2tensor/trax/layers/attention_test.py
index e3452e49f..c1b44421a 100644
--- a/tensor2tensor/trax/layers/attention_test.py
+++ b/tensor2tensor/trax/layers/attention_test.py
@@ -71,6 +71,22 @@ def test_merged_hashed_causal_attention(self):
     final_shape = base.check_shape_agreement(layer, input_shape)
     self.assertEqual((3, 32, 8), final_shape)
 
+  def test_time_bin_causal_attention_bin_length(self):
+    qkv_shape = (3, 57, 8)
+    input_shape = (qkv_shape, qkv_shape, qkv_shape)
+    layer = attention.TimeBinCausalAttention(
+        bin_length=16, dropout=0.1, mode='train')
+    final_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual((3, 57, 8), final_shape)
+
+  def test_time_bin_causal_attention_n_bins(self):
+    qkv_shape = (3, 57, 8)
+    input_shape = (qkv_shape, qkv_shape, qkv_shape)
+    layer = attention.TimeBinCausalAttention(
+        n_bins=4, dropout=0.1, mode='train')
+    final_shape = base.check_shape_agreement(layer, input_shape)
+    self.assertEqual((3, 57, 8), final_shape)
+
 
 if __name__ == '__main__':
   test.main()

From dbff3ac5218f41febc1146af8baf8f2d6716264a Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 2 Oct 2019 11:31:51 -0700
Subject: [PATCH 2520/2720] Internal

PiperOrigin-RevId: 272482483
---
 tensor2tensor/models/evolved_transformer.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 130d49206..56a72b832 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -96,6 +96,7 @@ def __init__(self, *args, **kwargs):
       # To make sure it they embeddings continue to be shared, we need to set
       # hparams.shared_embedding to True.
       self.hparams.shared_embedding = True
+    self._init_cache_fn = init_evolved_transformer_cache
 
 
 def evolved_transformer_encoder(encoder_input,

From c9ba7ec4476949939453e7cead91a6a4896297e9 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 2 Oct 2019 11:51:46 -0700
Subject: [PATCH 2521/2720] Update mask_id in the remaining configs.

PiperOrigin-RevId: 272486903
---
 tensor2tensor/trax/configs/transformer_lm1b_16gb.gin         | 5 +----
 tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin  | 5 +----
 tensor2tensor/trax/configs/transformer_ptb_16gb.gin          | 5 +----
 .../rl/configs/env_online_tune_transformer_lm1b_16gb.gin     | 5 +----
 .../configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin | 5 +----
 .../trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin | 5 +----
 tensor2tensor/trax/rl/envs/online_tune_env.py                | 4 +++-
 7 files changed, 9 insertions(+), 25 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_lm1b_16gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_16gb.gin
index 8f4bd0a36..d0725cdbb 100644
--- a/tensor2tensor/trax/configs/transformer_lm1b_16gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm1b_16gb.gin
@@ -15,10 +15,6 @@ inputs.data_dir = None
 inputs.dataset_name = 't2t_languagemodel_lm1b32k'
 inputs.input_name = 'targets'
 
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
 # Parameters for MultifactorSchedule:
 # ==============================================================================
 MultifactorSchedule.constant = 1.0
@@ -39,6 +35,7 @@ train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.Adafactor
 train.train_steps = 50000
+train.mask_id = 0
 
 # Parameters for DotProductCausalAttention:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin b/tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin
index 9882f2e9a..d4ea4390e 100644
--- a/tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin
+++ b/tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin
@@ -16,10 +16,6 @@ batch_fn.buckets_include_inputs_in_length=True
 inputs.data_dir = None
 inputs.dataset_name = 't2t_translate_ende_wmt32k'
 
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
 # Parameters for MultifactorSchedule:
 # ==============================================================================
 MultifactorSchedule.constant = 1.0
@@ -51,6 +47,7 @@ train.model = @trax.models.TransformerLM
 train.train_steps = 50000
 train.optimizer = @trax.optimizers.Adafactor
 train.has_weights = True
+train.mask_id = 0
 
 # Parameters for TransformerLM:
 # ==============================================================================
diff --git a/tensor2tensor/trax/configs/transformer_ptb_16gb.gin b/tensor2tensor/trax/configs/transformer_ptb_16gb.gin
index 69561693e..9d1b796ff 100644
--- a/tensor2tensor/trax/configs/transformer_ptb_16gb.gin
+++ b/tensor2tensor/trax/configs/transformer_ptb_16gb.gin
@@ -16,10 +16,6 @@ inputs.data_dir = None
 inputs.dataset_name = 't2t_languagemodel_ptb10k'
 inputs.input_name = 'targets'
 
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
 # Parameters for MultifactorSchedule:
 # ==============================================================================
 MultifactorSchedule.constant = 1.0
@@ -40,6 +36,7 @@ train.inputs = @trax.inputs.inputs
 train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.Adafactor
 train.train_steps = 20000
+train.mask_id = 0
 
 # Parameters for DotProductCausalAttention:
 # ==============================================================================
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
index 732260655..c765da790 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
@@ -21,10 +21,6 @@ inputs.input_name = 'targets'
 train_and_eval_dataset.eval_holdout_size = 0.05
 train_and_eval_dataset.eval_shuffle_files = True
 
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
 # Parameters for preprocess_fun:
 # ==============================================================================
 shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
@@ -100,3 +96,4 @@ OnlineTuneEnv.nontrainable_param_map = {
 OnlineTuneEnv.include_controls_in_observation = False
 OnlineTuneEnv.observation_range = (0.0, 10.0)
 OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
+OnlineTuneEnv.mask_id = 0
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
index 8f1eaf9d1..f4cef2897 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
@@ -22,10 +22,6 @@ inputs.dataset_name = 't2t_translate_ende_wmt32k'
 train_and_eval_dataset.eval_holdout_size = 0.05
 train_and_eval_dataset.eval_shuffle_files = True
 
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
 # Parameters for preprocess_fun:
 # ==============================================================================
 shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_concat_preprocess
@@ -97,3 +93,4 @@ OnlineTuneEnv.nontrainable_param_map = {
 OnlineTuneEnv.include_controls_in_observation = False
 OnlineTuneEnv.observation_range = (0.0, 10.0)
 OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
+OnlineTuneEnv.mask_id = 0
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin
index f7de3538d..34e42e9cb 100644
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin
+++ b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin
@@ -16,10 +16,6 @@ inputs.data_dir = None
 inputs.dataset_name = 't2t_languagemodel_ptb10k'
 inputs.input_name = 'targets'
 
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
 # Parameters for preprocess_fun:
 # ==============================================================================
 shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
@@ -95,3 +91,4 @@ OnlineTuneEnv.nontrainable_param_map = {
 OnlineTuneEnv.include_controls_in_observation = False
 OnlineTuneEnv.observation_range = (0.0, 10.0)
 OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
+OnlineTuneEnv.mask_id = 0
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index 804872e3a..2423c50d5 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -78,7 +78,8 @@ def __init__(self,
                # Don't save checkpoints by default, as they tend to use a lot of
                # space.
                should_save_checkpoints=False,
-               has_weights=False):
+               has_weights=False,
+               mask_id=None):
     if action_multipliers is None:
       action_multipliers = self.DEFAULT_ACTION_MULTIPLIERS
     self._model = model
@@ -95,6 +96,7 @@ def __init__(self,
         should_save=should_save_checkpoints,
         nontrainable_param_map=nontrainable_param_map,
         has_weights=has_weights,
+        mask_id=mask_id,
     )
     self._trainer = None
     self._action_multipliers = action_multipliers

From 1843c72d1d5faf4c085bb198b5dde0908f4081d0 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Wed, 2 Oct 2019 17:18:10 -0700
Subject: [PATCH 2522/2720] Implement fast inference for
 TimeBinCausalAttention.

Also added a test for consistency between {DotProduct,TimeBin}CausalAttention and added a bunch of comments.

PiperOrigin-RevId: 272554499
---
 tensor2tensor/trax/layers/attention.py        | 139 ++++++++++++++----
 tensor2tensor/trax/layers/attention_test.py   |  14 ++
 tensor2tensor/trax/models/transformer_test.py |  37 +++--
 3 files changed, 145 insertions(+), 45 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 2ec8cf86a..a276a8b7c 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -375,6 +375,39 @@ def forward_and_backward(self, inputs, grad, **kwargs):
     raise NotImplementedError()
 
 
+def _fast_inference_init_state(input_shapes, input_dtypes, buffer_length):
+  """Initializes state of a causal attention layer for fast inference."""
+  ((batch_size, _, _), _, _) = input_shapes
+  def init_buffer(shape, dtype):
+    (_, _, depth) = shape
+    return np.zeros((batch_size, buffer_length, depth), dtype=dtype)
+  (_, k, v) = tuple(
+      init_buffer(shape, dtype)
+      for (shape, dtype) in zip(input_shapes, input_dtypes)
+  )
+  mask = np.zeros((batch_size, 1, buffer_length))
+  index = 0
+  state = (k, v, mask, index)
+  return state
+
+
+def _fast_inference_update_state(inputs, state):
+  """Updates state of a causal attention layer for fast inference."""
+  assert backend.get_name() == 'jax', (
+      'JAX backend is required to use the predict mode.')
+  for x in inputs:
+    assert x.shape[1] == 1, (
+        'In predict mode the input sequence must be of length 1.')
+  # Fast inference: run with only 1 query in each step, storing the sequence
+  # of keys and values calculated so far in state.
+  (_, new_k, new_v) = inputs
+  (ks, vs, mask, index) = state
+  ks = jax.ops.index_update(ks, jax.ops.index[:, index, :], new_k[:, 0, :])
+  vs = jax.ops.index_update(vs, jax.ops.index[:, index, :], new_v[:, 0, :])
+  mask = jax.ops.index_update(mask, jax.ops.index[:, :, index], 1)
+  return (ks, vs, mask, index + 1)
+
+
 class DotProductCausalAttention(BaseCausalAttention):
   """A standard (non-memory-efficient) dot product attention implementation."""
 
@@ -399,20 +432,8 @@ def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
             onp.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
     else:
       assert self._mode == 'predict'
-      assert backend.get_name() == 'jax', (
-          'JAX backend is required to use the predict mode.')
-      for x in (q, k, v):
-        assert x.shape[1] == 1, (
-            'In predict mode the input sequence must be of length 1.')
-      # Fast inference: run with only 1 query in each step, storing the sequence
-      # of keys and values calculated so far in state.
-      (new_k, new_v) = (k, v)
-      (k, v, mask, index) = state
-      k = jax.ops.index_update(k, jax.ops.index[:, index, :], new_k[:, 0, :])
-      v = jax.ops.index_update(v, jax.ops.index[:, index, :], new_v[:, 0, :])
-      new_mask = jax.ops.index_update(mask, jax.ops.index[:, :, index], 1)
-      state = (k, v, new_mask, index + 1)
-      mask = new_mask
+      state = _fast_inference_update_state(inputs, state)
+      (k, v, mask, _) = state
 
     res = DotProductAttention(
         q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=rng)
@@ -433,21 +454,11 @@ def new_params_and_state(self, input_shapes, input_dtype, rng):
       return (), ()
 
     assert self._mode == 'predict'
+    params = ()
     # Buffer length is hardcoded for now. TODO(pkozakowski): Pass it from the
     # model.
     max_len = 2048
-    ((batch_size, _, _), _, _) = input_shapes
-    def initial_state(shape, dtype):
-      (_, _, depth) = shape
-      return np.zeros((batch_size, max_len, depth), dtype=dtype)
-    (_, k, v) = tuple(
-        initial_state(shape, dtype)
-        for (shape, dtype) in zip(input_shapes, input_dtype)
-    )
-    mask = np.zeros((batch_size, 1, max_len))
-    index = 0
-    params = ()
-    state = (k, v, mask, index)
+    state = _fast_inference_init_state(input_shapes, input_dtype, max_len)
     return params, state
 
 
@@ -659,7 +670,7 @@ def body_fun(vals):  # pylint: disable=invalid-name
 class TimeBinCausalAttention(BaseCausalAttention):
   """Causal attention where only nearby chunks of items attend to each other."""
 
-  def __init__(self, dropout, mode, bin_length=None, n_bins=None,
+  def __init__(self, mode, dropout=0.0, bin_length=None, n_bins=None,
                share_qk=False):
     super(TimeBinCausalAttention, self).__init__()
     if (bin_length is None) == (n_bins is None):
@@ -672,7 +683,8 @@ def __init__(self, dropout, mode, bin_length=None, n_bins=None,
     if mode == 'train':
       self.dropout = dropout
     else:
-      self.dropout = None
+      self.dropout = 0.0
+    self._mode = mode
 
   def forward_and_backward(self, inputs, ct, **kwargs):
     assert backend.get_name() == 'jax', (
@@ -710,10 +722,19 @@ def pad_input(x):
 
   def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
     del params, kwargs
+    if self._mode in ('train', 'eval'):
+      output = self._forward_train_eval(inputs, rng)
+      return (output, state)
+    else:
+      assert self._mode == 'predict'
+      return self._forward_predict(inputs, state, rng)
+
+  def _forward_train_eval(self, inputs, rng):
     (inputs, original_len, n_bins) = self._pad_inputs(inputs)
     q, k, v = inputs
     seqlen = q.shape[-2]
     # q/k/v are n_batch*n_heads, seqlen, d_head
+    # Time indices for causal masking.
     t = jax.lax.tie_in(q, np.arange(seqlen))
 
     # Split off a "bin" axis for chunks of consecutive items.
@@ -727,6 +748,7 @@ def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
 
     # Allow each chunk to attend within itself, and also one chunk back.
     def look_one_back(x):
+      # Output: pairs [ bin_i bin_{i-1} ] concatenated on the time axis.
       if len(x.shape) == 2:
         x_extra = np.concatenate([x[-1:, :], x[:-1, :]], axis=0)
         return np.concatenate([x, x_extra], axis=1)
@@ -742,7 +764,7 @@ def look_one_back(x):
     # Dot-product attention.
     dots = np.matmul(bq, np.swapaxes(bk, -1, -2)) / np.sqrt(bq.shape[-1])
 
-    # Causal masking
+    # Causal masking based on the time indices.
     mask = jax.lax.convert_element_type(
         jax.lax.lt(bq_t[None, :, :, None], bkv_t[None, :, None, :]),
         np.float32)
@@ -754,7 +776,7 @@ def look_one_back(x):
       self_mask = jax.lax.tie_in(dots, self_mask)
       dots = dots - 1e5 * self_mask
 
-    if self.dropout is not None and self.dropout > 0.0:
+    if self.dropout > 0.0:
       # Dropout is broadcast across the batch+head dimension
       dropout_shape = (1, dots.shape[-3], dots.shape[-2], dots.shape[-1])
       keep_prob = jax.lax.tie_in(dots, 1.0 - self.dropout)
@@ -768,7 +790,62 @@ def look_one_back(x):
 
     output = np.reshape(bo, (bo.shape[0], -1, bo.shape[-1]))
     assert output.shape == v.shape
-    return output[..., :original_len, :], state
+    return output[..., :original_len, :]
+
+  def _forward_predict(self, inputs, state, rng):
+    state = _fast_inference_update_state(inputs, state)
+
+    (q, _, _) = inputs
+    (ks, vs, mask, index) = state
+    output = DotProductAttention(
+        q, ks, vs, mask, dropout=self.dropout, mode=self._mode, rng=rng
+    )
+
+    def roll_state(state):
+      """Rolls the buffers backward to make space for new data."""
+      (ks, vs, mask, index) = state
+      # Move the second bin into the first one's place in both buffers.
+      def roll_buffer(buf):
+        return jax.ops.index_update(
+            buf,
+            jax.ops.index[:, :self.bin_length, :],
+            buf[:, self.bin_length:, :],
+        )
+      (ks, vs) = map(roll_buffer, (ks, vs))
+      # Zero out the second bin in the mask.
+      mask = jax.ops.index_update(
+          mask, jax.ops.index[:, :, self.bin_length:], 0
+      )
+      # Update the index to match the rolled buffers.
+      index -= self.bin_length
+      return (ks, vs, mask, index)
+
+    # Once we get to the end of the buffer, move the second bin back to make
+    # space for new data: [ bin_i bin_{i+1} | ] -> [ bin_{i+1} | bin_{i+1} ],
+    # where | is where index points at in the buffer.
+    state = jax.lax.cond(
+        pred=(index == 2 * self.bin_length),
+        true_operand=state,
+        true_fun=roll_state,
+        false_operand=state,
+        false_fun=(lambda x: x),
+    )
+    return (output, state)
+
+  def new_params_and_state(self, input_shapes, input_dtype, rng):
+    if self._mode in ('train', 'eval'):
+      return (), ()
+
+    assert self._mode == 'predict'
+    assert self.bin_length is not None, (
+        'For fast inference, TimeBinCausalAttention must be parameterized by '
+        'bin_length.'
+    )
+    params = ()
+    state = _fast_inference_init_state(
+        input_shapes, input_dtype, 2 * self.bin_length
+    )
+    return params, state
 
 
 class LSHCausalAttention(BaseCausalAttention):
diff --git a/tensor2tensor/trax/layers/attention_test.py b/tensor2tensor/trax/layers/attention_test.py
index c1b44421a..f2ca0f73e 100644
--- a/tensor2tensor/trax/layers/attention_test.py
+++ b/tensor2tensor/trax/layers/attention_test.py
@@ -87,6 +87,20 @@ def test_time_bin_causal_attention_n_bins(self):
     final_shape = base.check_shape_agreement(layer, input_shape)
     self.assertEqual((3, 57, 8), final_shape)
 
+  def test_time_bin_and_dot_product_causal_attention_are_consistent(self):
+    dot_product_layer = attention.DotProductCausalAttention(
+        dropout=0.0, mode='train')
+    time_bin_layer = attention.TimeBinCausalAttention(
+        bin_length=4, dropout=0.0, mode='train')
+
+    # Exactly 2 bins.
+    input_shape = (3, 8, 8)
+    inputs = [onp.random.uniform(size=input_shape) for _ in range(3)]
+
+    dot_product_output = dot_product_layer(inputs)
+    time_bin_output = time_bin_layer(inputs)
+    onp.testing.assert_array_almost_equal(dot_product_output, time_bin_output)
+
 
 if __name__ == '__main__':
   test.main()
diff --git a/tensor2tensor/trax/models/transformer_test.py b/tensor2tensor/trax/models/transformer_test.py
index 70266a8a3..2abd9adcd 100644
--- a/tensor2tensor/trax/models/transformer_test.py
+++ b/tensor2tensor/trax/models/transformer_test.py
@@ -57,24 +57,35 @@ def _test_transformer_forward_shape(self, input_vocab_size,
                              else input_vocab_size]))
     self.assertEqual(expected_shape, final_shape[0])
 
-  def test_transformer_lm_fast_inference(self):
+  @parameterized.named_parameters(
+      ('same_vocab', 16, None),
+      ('same_size', 16, 16),
+      ('different_size', 16, 50))
+  def test_transformer_forward_shape(self, input_vocab_size, output_vocab_size):
+    """Run the Transformer forward and check output shape."""
+    self._test_transformer_forward_shape(input_vocab_size, output_vocab_size)
+
+
+  def _test_fast_inference(self, attention_type, length):
     with backend.use_backend('jax'):
       vocab_size = 16
       model_fn = functools.partial(
           transformer.TransformerLM,
-          vocab_size=vocab_size, d_model=4, d_ff=8, n_layers=2, n_heads=2)
+          vocab_size=vocab_size, d_model=4, d_ff=8, n_layers=2, n_heads=2,
+          attention_type=attention_type,
+      )
       model_slow = model_fn(mode='eval')
       model_fast = model_fn(mode='predict')
       rng = backend.random.get_prng(0)
       batch_size = 2
-      _, _ = model_slow.initialize_once((batch_size, 1), np.int32, rng)
-      _, _ = model_fast.initialize_once((batch_size, 1), np.int32, rng)
+      # Given the same rng, both models initialize with the same parameters.
+      model_slow.initialize_once((batch_size, 1), np.int32, rng)
+      model_fast.initialize_once((batch_size, 1), np.int32, rng)
 
-      max_length = 5
-      buf = onp.zeros((batch_size, max_length), dtype=np.int32)
+      buf = onp.zeros((batch_size, length), dtype=np.int32)
       next_sym = onp.zeros((batch_size, 1), dtype=onp.int32)
 
-      for index in range(max_length):
+      for index in range(length):
         logits_slow = model_slow(buf, rng=rng)
         logits_fast = model_fast(next_sym, rng=rng)
         onp.testing.assert_array_almost_equal(
@@ -82,14 +93,12 @@ def test_transformer_lm_fast_inference(self):
         next_sym = onp.random.randint(vocab_size, size=(batch_size, 1))
         buf[:, index] = next_sym[:, 0]
 
-  @parameterized.named_parameters(
-      ('same_vocab', 16, None),
-      ('same_size', 16, 16),
-      ('different_size', 16, 50))
-  def test_transformer_forward_shape(self, input_vocab_size, output_vocab_size):
-    """Run the Transformer forward and check output shape."""
-    self._test_transformer_forward_shape(input_vocab_size, output_vocab_size)
+  def test_dot_product_causal_attention_fast_inference(self):
+    self._test_fast_inference(tl.DotProductCausalAttention, length=5)
 
+  def test_time_bin_causal_attention_fast_inference(self):
+    attention = functools.partial(tl.TimeBinCausalAttention, bin_length=2)
+    self._test_fast_inference(attention, length=7)
 
 if __name__ == '__main__':
   absltest.main()

From 8be915afc7570f08d19c47aea57541250a32028b Mon Sep 17 00:00:00 2001
From: pb <9627490+przemb@users.noreply.github.com>
Date: Thu, 3 Oct 2019 03:07:45 +0200
Subject: [PATCH 2523/2720] Deep_discriminator - missed relu (#1720)

---
 tensor2tensor/layers/common_layers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index a59d508fd..6d0af1e6a 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3724,7 +3724,7 @@ def double_discriminator(x, filters1=128, filters2=None,
       tf.reshape(net, [batch_size, -1])
     net = tf.nn.relu(net)
     net = layers().Conv2D(
-        filters2, kernel_size, strides=strides, padding="SAME", name="conv2")(x)
+        filters2, kernel_size, strides=strides, padding="SAME", name="conv2")(net)
     if pure_mean:
       net2 = tf.reduce_mean(net, [1, 2])
     else:

From 5b3409165d52b31241652608564d2fc2783321e9 Mon Sep 17 00:00:00 2001
From: pb <9627490+przemb@users.noreply.github.com>
Date: Wed, 2 Oct 2019 18:28:47 -0700
Subject: [PATCH 2524/2720] Merge of PR #1720

PiperOrigin-RevId: 272565130
---
 tensor2tensor/layers/common_layers.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 6d0af1e6a..e772870b1 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -3724,7 +3724,8 @@ def double_discriminator(x, filters1=128, filters2=None,
       tf.reshape(net, [batch_size, -1])
     net = tf.nn.relu(net)
     net = layers().Conv2D(
-        filters2, kernel_size, strides=strides, padding="SAME", name="conv2")(net)
+        filters2, kernel_size, strides=strides, padding="SAME",
+        name="conv2")(net)
     if pure_mean:
       net2 = tf.reduce_mean(net, [1, 2])
     else:

From 1fa0b46132e15c866127b6723dfa84ae87743e9f Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 2 Oct 2019 18:54:10 -0700
Subject: [PATCH 2525/2720] Bump up tf version in setup.py, missed this the
 last time.

PiperOrigin-RevId: 272568653
---
 setup.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index 55cd65e54..f5274d52c 100644
--- a/setup.py
+++ b/setup.py
@@ -60,8 +60,8 @@
         'tqdm',
     ],
     extras_require={
-        'tensorflow': ['tensorflow>=1.13.1'],
-        'tensorflow_gpu': ['tensorflow-gpu>=1.13.1'],
+        'tensorflow': ['tensorflow>=1.14.0'],
+        'tensorflow_gpu': ['tensorflow-gpu>=1.14.0'],
         'tensorflow-hub': ['tensorflow-hub>=0.1.1'],
         'tests': [
             'absl-py',

From a2fff1c0da9a4e7ae08aa2f33a58ef873975ddce Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 2 Oct 2019 19:06:56 -0700
Subject: [PATCH 2526/2720] Repurpose demo notebook as a practical introduction
 to Trax layers.

PiperOrigin-RevId: 272570383
---
 tensor2tensor/trax/layers/demo.ipynb  | 189 ------
 tensor2tensor/trax/layers/intro.ipynb | 860 ++++++++++++++++++++++++++
 2 files changed, 860 insertions(+), 189 deletions(-)
 delete mode 100644 tensor2tensor/trax/layers/demo.ipynb
 create mode 100644 tensor2tensor/trax/layers/intro.ipynb

diff --git a/tensor2tensor/trax/layers/demo.ipynb b/tensor2tensor/trax/layers/demo.ipynb
deleted file mode 100644
index 8317b683b..000000000
--- a/tensor2tensor/trax/layers/demo.ipynb
+++ /dev/null
@@ -1,189 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "7yuytuIllsv1"
-      },
-      "source": [
-        "# Trax Layers\n",
-        "\n",
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 1,
-      "metadata": {
-        "colab": {
-          "height": 119
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 88481,
-          "status": "ok",
-          "timestamp": 1563927238895,
-          "user": {
-            "displayName": "Lukasz Kaiser",
-            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mC8pChl87HbK_eOtVhtNPwUVx8btvfyYzH9UHn3=s64",
-            "userId": "13267693649565518272"
-          },
-          "user_tz": 420
-        },
-        "id": "oILRLCWN_16u",
-        "outputId": "3f750014-c633-4162-ad07-f3c56c273304"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "T2T: skipped importing 1 data_generators modules. OK if no other errors. Depend on _heavy or problem-specific py_binary targets if trying to use a module that was skipped.\n"
-          ]
-        },
-        {
-          "name": "stderr",
-          "output_type": "stream",
-          "text": [
-            "matplotlib.use() must be called *before* pylab, matplotlib.pyplot,\n",
-            "or matplotlib.backends is imported for the first time.\n",
-            "\n"
-          ]
-        }
-      ],
-      "source": [
-        "#@title Licence and python imports.\n",
-        "# Copyright 2018 Google LLC.\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",
-        "# https://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",
-        "import datetime\n",
-        "import numpy as onp\n",
-        "\n",
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 5,
-      "metadata": {
-        "cellView": "both",
-        "colab": {
-          "height": 51
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 2992,
-          "status": "ok",
-          "timestamp": 1563927313403,
-          "user": {
-            "displayName": "Lukasz Kaiser",
-            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mC8pChl87HbK_eOtVhtNPwUVx8btvfyYzH9UHn3=s64",
-            "userId": "13267693649565518272"
-          },
-          "user_tz": 420
-        },
-        "id": "vlGjGoGMTt-D",
-        "outputId": "7a7b5a1e-c01e-4a5e-eeb1-a88d9500aad4"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "/bin/sh: pip: command not found\n",
-            "/bin/sh: pip: command not found\n"
-          ]
-        }
-      ],
-      "source": [
-        "#@title Install TensorFlow, Tensor2Tensor and Trax.\n",
-        "\n",
-        "! pip install -q tensorflow\n",
-        "! pip install -q -U tensor2tensor\n",
-        "\n",
-        "from tensor2tensor.trax import trax\n",
-        "from tensor2tensor.trax import backend\n",
-        "from tensor2tensor.trax import layers as tl"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 6,
-      "metadata": {
-        "colab": {
-          "height": 68
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 434,
-          "status": "ok",
-          "timestamp": 1563927323194,
-          "user": {
-            "displayName": "Lukasz Kaiser",
-            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mC8pChl87HbK_eOtVhtNPwUVx8btvfyYzH9UHn3=s64",
-            "userId": "13267693649565518272"
-          },
-          "user_tz": 420
-        },
-        "id": "V09viOSEQvQe",
-        "outputId": "822ec44b-1e2e-4b6e-9fc7-9fc29f5d3783"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "[[-7 -6 -5 -4 -3]\n",
-            " [-2 -1  0  1  2]\n",
-            " [ 3  4  5  6  7]]\n",
-            "[[0 0 0 0 0]\n",
-            " [0 0 0 1 2]\n",
-            " [3 4 5 6 7]]\n"
-          ]
-        }
-      ],
-      "source": [
-        "x  = onp.arange(-7, 8).reshape(3, -1)\n",
-        "rng = backend.random.get_prng(0)\n",
-        "layer = tl.Relu()\n",
-        "layer.initialize_once(x.shape, x.dtype, rng)\n",
-        "output = layer(x)\n",
-        "print(x)\n",
-        "print(output)"
-      ]
-    }
-  ],
-  "metadata": {
-    "colab": {
-      "collapsed_sections": [],
-      "last_runtime": {
-        "build_target": "//learning/deepmind/dm_python:dm_notebook",
-        "kind": "private"
-      },
-      "name": "Trax Layers",
-      "provenance": [
-        {
-          "file_id": "1EH76AWQ_pvT4i8ZXfkv-SCV4MrmllEl5",
-          "timestamp": 1563927451951
-        }
-      ],
-      "version": "0.3.2"
-    },
-    "kernelspec": {
-      "display_name": "Python 2",
-      "name": "python2"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}
diff --git a/tensor2tensor/trax/layers/intro.ipynb b/tensor2tensor/trax/layers/intro.ipynb
new file mode 100644
index 000000000..680b18b0b
--- /dev/null
+++ b/tensor2tensor/trax/layers/intro.ipynb
@@ -0,0 +1,860 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "7yuytuIllsv1"
+      },
+      "source": [
+        "# A Conceptual, Practical Introduction to Trax Layers\n",
+        "\n",
+        "This notebook introduces the core concepts and programming components of the Trax library through a series of code samples and explanations. The topics covered in following sections are:\n",
+        "  - **layers**: the basic building blocks and how to combine them into networks\n",
+        "  - **data flows**, data stack: how the Trax runtime moves data through the layers\n",
+        "  - **models**: how to train, evaluate, and run predictions with Trax models\n",
+        "  - **new layer classes**: how to define and test your own Layer classes\n",
+        "\n",
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "BIl27504La0G"
+      },
+      "source": [
+        "## General Setup\n",
+        "Execute the following few cells (once) before running any of the code samples in this notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "oILRLCWN_16u"
+      },
+      "outputs": [],
+      "source": [
+        "#@title\n",
+        "# Copyright 2018 Google LLC.\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",
+        "# https://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",
+        "import numpy as onp\n",
+        "\n",
+        "\n",
+        "\n",
+        "# Import Trax\n",
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {
+        "cellView": "both",
+        "colab": {
+          "height": 51
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 267,
+          "status": "ok",
+          "timestamp": 1569980014037,
+          "user": {
+            "displayName": "Jonni Kanerva",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
+            "userId": "12646438700820120918"
+          },
+          "user_tz": 420
+        },
+        "id": "vlGjGoGMTt-D",
+        "outputId": "97840e48-52ee-4ae1-87f1-83c1b0251055"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "/bin/sh: pip: command not found\n",
+            "/bin/sh: pip: command not found\n"
+          ]
+        }
+      ],
+      "source": [
+        "#@title Run for installation.\n",
+        "\n",
+        "! pip install -q -U tensor2tensor\n",
+        "! pip install -q tensorflow\n",
+        "\n",
+        "from tensor2tensor.trax import backend\n",
+        "from tensor2tensor.trax import layers as tl"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "bYWNWL9MJHv9"
+      },
+      "outputs": [],
+      "source": [
+        "onp.set_printoptions(precision=3)  # Less visual noise in the numerical outputs."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "-LQ89rFFsEdk"
+      },
+      "source": [
+        "# Layers\n",
+        "\n",
+        "The Layer class represents Trax's concept of a layer, as summarized in the start of the class's docstring:\n",
+        "```\n",
+        "class Layer(object):\n",
+        "  \"\"\"Base class for composable layers in a deep learning network.\n",
+        "\n",
+        "  Layers are the basic building blocks for deep learning models. A Trax layer\n",
+        "  computes a function from zero or more inputs to zero or more outputs,\n",
+        "  optionally using trainable parameters (common) and non-parameter state (not\n",
+        "  common). Authors of new layer subclasses typically override at most two\n",
+        "  methods of the base `Layer` class:\n",
+        "\n",
+        "    forward(inputs, params=(), state=(), **kwargs):\n",
+        "      Computes this layer's output as part of a forward pass through the model.\n",
+        "\n",
+        "    new_params_and_state(self, input_shape, input_dtype, rng):\n",
+        "      Returns a (params, state) pair suitable for initializing this layer.\n",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "LyLVtdxorDPO"
+      },
+      "source": [
+        "## A layer computes a function."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "ntZ4_eNQldzL"
+      },
+      "source": [
+        "A layer computes a function from zero or more inputs to zero or more outputs. The inputs and outputs are NumPy arrays or JAX objects wrapping NumPy arrays.\n",
+        "\n",
+        "The simplest layers, those with no parameters or state, can be used without initialization. You can think of them (and test them) like simple mathematical functions. For ease of testing and interactive exploration, layer\n",
+        "objects implement the `__call__ ` method, so you can call them directly on input data:\n",
+        "\n",
+        "```\n",
+        "y = layer(x)\n",
+        "```\n",
+        "\n",
+        "\n",
+        "### Example 1. tl.Relu [1 input, 1 output]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 4,
+      "metadata": {
+        "colab": {
+          "height": 204
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 1224,
+          "status": "ok",
+          "timestamp": 1569980015601,
+          "user": {
+            "displayName": "Jonni Kanerva",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
+            "userId": "12646438700820120918"
+          },
+          "user_tz": 420
+        },
+        "id": "V09viOSEQvQe",
+        "outputId": "b5f206b9-29ad-42b1-cd64-2a52cac4911e"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "type(x): \u003ctype 'numpy.ndarray'\u003e\n",
+            "\n",
+            "x:\n",
+            "[[-7. -6. -5. -4. -3.]\n",
+            " [-2. -1.  0.  1.  2.]\n",
+            " [ 3.  4.  5.  6.  7.]]\n",
+            "\n",
+            "relu(x):\n",
+            "[[0. 0. 0. 0. 0.]\n",
+            " [0. 0. 0. 1. 2.]\n",
+            " [3. 4. 5. 6. 7.]]\n"
+          ]
+        }
+      ],
+      "source": [
+        "x = onp.arange(-7, 8).reshape(3, -1).astype(onp.float32)\n",
+        "\n",
+        "# Create a layer object (a Relu instance) and apply the layer to data x.\n",
+        "relu = tl.Relu()\n",
+        "y = relu(x)\n",
+        "\n",
+        "# Show the inputs and outputs.\n",
+        "template = ('type(x): {}\\n\\n'\n",
+        "            'x:\\n{}\\n\\n'\n",
+        "            'relu(x):\\n{}')\n",
+        "print(template.format(type(x), x, y))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "7sYxIT8crFVE"
+      },
+      "source": [
+        "### Example 2. tl.Concatenate [2 inputs, 1 output]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 5,
+      "metadata": {
+        "colab": {
+          "height": 391
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 1251,
+          "status": "ok",
+          "timestamp": 1569980016872,
+          "user": {
+            "displayName": "Jonni Kanerva",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
+            "userId": "12646438700820120918"
+          },
+          "user_tz": 420
+        },
+        "id": "LMPPNWXLoOZI",
+        "outputId": "310d6c78-7dfc-48c5-a5e9-0f0b09abb2bc"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "x1:\n",
+            "[[-7. -6. -5. -4. -3.]\n",
+            " [-2. -1.  0.  1.  2.]\n",
+            " [ 3.  4.  5.  6.  7.]]\n",
+            "\n",
+            "x2:\n",
+            "[[-70. -60. -50. -40. -30.]\n",
+            " [-20. -10.   0.  10.  20.]\n",
+            " [ 30.  40.  50.  60.  70.]]\n",
+            "\n",
+            "concatenate_0([x1, x2]):\n",
+            "[[ -7.  -6.  -5.  -4.  -3.]\n",
+            " [ -2.  -1.   0.   1.   2.]\n",
+            " [  3.   4.   5.   6.   7.]\n",
+            " [-70. -60. -50. -40. -30.]\n",
+            " [-20. -10.   0.  10.  20.]\n",
+            " [ 30.  40.  50.  60.  70.]]\n",
+            "\n",
+            "concatenate_1([x1, x2]):\n",
+            "[[ -7.  -6.  -5.  -4.  -3. -70. -60. -50. -40. -30.]\n",
+            " [ -2.  -1.   0.   1.   2. -20. -10.   0.  10.  20.]\n",
+            " [  3.   4.   5.   6.   7.  30.  40.  50.  60.  70.]]\n"
+          ]
+        }
+      ],
+      "source": [
+        "x1 = onp.arange(-7, 8).reshape(3, -1).astype(onp.float32)\n",
+        "x2 = x1 * 10\n",
+        "\n",
+        "concatenate_0 = tl.Concatenate(axis=0)\n",
+        "y = concatenate_0([x1, x2])\n",
+        "\n",
+        "template = ('x1:\\n{}\\n\\n'\n",
+        "            'x2:\\n{}\\n\\n'\n",
+        "            'concatenate_0([x1, x2]):\\n{}')\n",
+        "print(template.format(x1, x2, y))\n",
+        "\n",
+        "concatenate_1 = tl.Concatenate(axis=1)\n",
+        "y = concatenate_1([x1, x2])\n",
+        "\n",
+        "template = ('\\nconcatenate_1([x1, x2]):\\n{}')\n",
+        "print(template.format(y))\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "1oZv3R8bRMvF"
+      },
+      "source": [
+        "## Layers are trainable."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "3d64M7wLryji"
+      },
+      "source": [
+        "Most layer types are trainable: they include parameters that modify the computation of outputs from inputs, and they use back-progagated gradients to update those parameters.\n",
+        "\n",
+        "Before use, trainable layers must have their parameters initialized, typically using a PRNG (pseudo-random number generator) key for random number generation. Trax's model trainers take care of this behind the scenes, but if you are using a layer in insolation, you have to do the initialization yourself. For this, use the `initialize_once` method:\n",
+        "\n",
+        "```\n",
+        "  def initialize_once(self, input_shapes, input_dtype, rng):\n",
+        "    \"\"\"Initializes this layer and its sublayers recursively.\n",
+        "\n",
+        "    This method is designed to initialize each layer instance once, even if the\n",
+        "    same layer instance occurs in multiple places in the network. This enables\n",
+        "    weight sharing to be implemented as layer sharing.\n",
+        "\n",
+        "    ...\n",
+        "```\n",
+        "\n",
+        "**Example 2.** tl.LayerNorm [1 input, 1 output, has parameters]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 6,
+      "metadata": {
+        "colab": {
+          "height": 221
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 2200,
+          "status": "ok",
+          "timestamp": 1569980019093,
+          "user": {
+            "displayName": "Jonni Kanerva",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
+            "userId": "12646438700820120918"
+          },
+          "user_tz": 420
+        },
+        "id": "Ie7iyX91qAx2",
+        "outputId": "b89ecf84-456f-4507-a6d3-a3e36f629ecb"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "x:\n",
+            "[[-7. -6. -5. -4. -3.]\n",
+            " [-2. -1.  0.  1.  2.]\n",
+            " [ 3.  4.  5.  6.  7.]]\n",
+            "\n",
+            "layer_norm(x):\n",
+            "[[-1.414 -0.707  0.     0.707  1.414]\n",
+            " [-1.414 -0.707  0.     0.707  1.414]\n",
+            " [-1.414 -0.707  0.     0.707  1.414]]\n",
+            "\n",
+            "layer_norm.params:\n",
+            "(_FilledConstant([1., 1., 1., 1., 1.], dtype=float32), _FilledConstant([0., 0., 0., 0., 0.], dtype=float32))\n"
+          ]
+        }
+      ],
+      "source": [
+        "prng_key = backend.random.get_prng(0)\n",
+        "input_dtype = onp.float32\n",
+        "input_shape = (3, 5)\n",
+        "x = onp.arange(-7, 8).reshape(*input_shape).astype(input_dtype)\n",
+        "\n",
+        "layer_norm = tl.LayerNorm()\n",
+        "layer_norm.initialize_once(input_shape, input_dtype, prng_key)\n",
+        "y = layer_norm(x)\n",
+        "\n",
+        "template = ('x:\\n{}\\n\\n'\n",
+        "            'layer_norm(x):\\n{}\\n\\n'\n",
+        "            'layer_norm.params:\\n{}')\n",
+        "print(template.format(x, y, layer_norm.params))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "ZWZUXEJAofH-"
+      },
+      "source": [
+        "## Layers combine into layers."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "d47gVdGV1vWw"
+      },
+      "source": [
+        "The Trax library builders encourage users, where possible, to build new layers as combinations of existing layers. The library provides a small set of _combinator_ layers for this: layer objects that make a list of layers behave as a single layer, computing outputs from inputs and updating parameters from gradients.\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "vC1ymG2j0iyp"
+      },
+      "source": [
+        "## Combine with Serial(...)\n",
+        "\n",
+        "The most common way to combine layers is serially, using the `Serial` class:\n",
+        "```\n",
+        "class Serial(base.Layer):\n",
+        "  \"\"\"Combinator that applies layers serially (by function composition).\n",
+        "\n",
+        "  A Serial combinator uses stack semantics to manage data for its sublayers.\n",
+        "  Each sublayer sees only the inputs it needs and returns only the outputs it\n",
+        "  has generated. The sublayers interact via the data stack. For instance, a\n",
+        "  sublayer k, following sublayer j, gets called with the data stack in the\n",
+        "  state left after layer j has applied. The Serial combinator then:\n",
+        "\n",
+        "    - takes N_in items off the top of the stack (N_in = k.n_inputs) and calls\n",
+        "      layer k, passing those items as arguments; and\n",
+        "\n",
+        "    - takes layer k's N_out return values (N_out = k.n_outputs) and pushes\n",
+        "      them onto the data stack.\n",
+        "\n",
+        "  ...\n",
+        "```\n",
+        "As described above, the output of one layer is the input of the next, which amounts to function composition:\n",
+        "\n",
+        "```\n",
+        "#  h(.) = g(f(.))\n",
+        "layer_h = Serial(\n",
+        "    layer_f,\n",
+        "    layer_g,\n",
+        ")\n",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 7,
+      "metadata": {
+        "colab": {
+          "height": 170
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 2492,
+          "status": "ok",
+          "timestamp": 1569980021601,
+          "user": {
+            "displayName": "Jonni Kanerva",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
+            "userId": "12646438700820120918"
+          },
+          "user_tz": 420
+        },
+        "id": "dW5fpusjvjmh",
+        "outputId": "ff8d5d4e-00af-487e-d287-fc77dab366c8"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "x:\n",
+            "[[-7. -6. -5. -4. -3.]\n",
+            " [-2. -1.  0.  1.  2.]\n",
+            " [ 3.  4.  5.  6.  7.]]\n",
+            "\n",
+            "layer_norm(relu(x)):\n",
+            "[[ 0.     0.     0.     0.     0.   ]\n",
+            " [-0.75  -0.75  -0.75   0.5    1.75 ]\n",
+            " [-1.414 -0.707  0.     0.707  1.414]]\n"
+          ]
+        }
+      ],
+      "source": [
+        "prng_key = backend.random.get_prng(0)\n",
+        "input_dtype = onp.float32\n",
+        "input_shape = (3, 5)\n",
+        "x = onp.arange(-7, 8).reshape(*input_shape).astype(input_dtype)\n",
+        "\n",
+        "layer_block = tl.Serial(\n",
+        "    tl.Relu(),\n",
+        "    tl.LayerNorm(),\n",
+        ")\n",
+        "layer_block.initialize_once(input_shape, input_dtype, prng_key)\n",
+        "y = layer_block(x)\n",
+        "\n",
+        "template = ('x:\\n{}\\n\\n'\n",
+        "            'layer_norm(relu(x)):\\n{}')\n",
+        "print(template.format(x, y))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "PqVNdoONcTp0"
+      },
+      "source": [
+        "## Combine with Parallel(...)\n",
+        "\n",
+        "\u003c... and need to introduce Dup(...) ...\u003e"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 8,
+      "metadata": {
+        "colab": {
+          "height": 255
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 2050,
+          "status": "ok",
+          "timestamp": 1569980023681,
+          "user": {
+            "displayName": "Jonni Kanerva",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
+            "userId": "12646438700820120918"
+          },
+          "user_tz": 420
+        },
+        "id": "teOhSJ8A8Es3",
+        "outputId": "6759189c-2ef7-4530-c9ee-88b0571208c5"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "x:\n",
+            "[[-7. -6. -5. -4. -3.]\n",
+            " [-2. -1.  0.  1.  2.]\n",
+            " [ 3.  4.  5.  6.  7.]]\n",
+            "\n",
+            "relu(x):\n",
+            "[[0. 0. 0. 0. 0.]\n",
+            " [0. 0. 0. 1. 2.]\n",
+            " [3. 4. 5. 6. 7.]]\n",
+            "\n",
+            "layer_norm(x):\n",
+            "[[-1.414 -0.707  0.     0.707  1.414]\n",
+            " [-1.414 -0.707  0.     0.707  1.414]\n",
+            " [-1.414 -0.707  0.     0.707  1.414]]\n"
+          ]
+        }
+      ],
+      "source": [
+        "prng_key = backend.random.get_prng(0)\n",
+        "input_dtype = onp.float32\n",
+        "input_shape = (3, 5)\n",
+        "x = onp.arange(-7, 8).reshape(*input_shape).astype(input_dtype)\n",
+        "\n",
+        "layer_block = tl.Serial(\n",
+        "    tl.Dup(),\n",
+        "    tl.Parallel(tl.Relu(), tl.LayerNorm()),\n",
+        ")\n",
+        "layer_block.initialize_once(input_shape, input_dtype, prng_key)\n",
+        "y = layer_block(x)\n",
+        "(y0, y1) = y  # two outputs from layer\n",
+        "\n",
+        "template = ('x:\\n{}\\n\\n'\n",
+        "            'relu(x):\\n{}\\n\\n'\n",
+        "            'layer_norm(x):\\n{}')\n",
+        "print(template.format(x, y0, y1))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 9,
+      "metadata": {
+        "colab": {
+          "height": 170
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 2912,
+          "status": "ok",
+          "timestamp": 1569980026618,
+          "user": {
+            "displayName": "Jonni Kanerva",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
+            "userId": "12646438700820120918"
+          },
+          "user_tz": 420
+        },
+        "id": "2EIewVHC9QQZ",
+        "outputId": "0194d386-4e06-4bd3-d2a1-24f668b268ac"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "x:\n",
+            "[[-7. -6. -5. -4. -3.]\n",
+            " [-2. -1.  0.  1.  2.]\n",
+            " [ 3.  4.  5.  6.  7.]]\n",
+            "\n",
+            "concatenate(relu(x), layer_norm(x)):\n",
+            "[[ 0.     0.     0.     0.     0.    -1.414 -0.707  0.     0.707  1.414]\n",
+            " [ 0.     0.     0.     1.     2.    -1.414 -0.707  0.     0.707  1.414]\n",
+            " [ 3.     4.     5.     6.     7.    -1.414 -0.707  0.     0.707  1.414]]\n"
+          ]
+        }
+      ],
+      "source": [
+        "prng_key = backend.random.get_prng(0)\n",
+        "input_dtype = onp.float32\n",
+        "input_shape = (3, 5)\n",
+        "x = onp.arange(-7, 8).reshape(*input_shape).astype(input_dtype)\n",
+        "\n",
+        "layer_block = tl.Serial(\n",
+        "    tl.Dup(),\n",
+        "    tl.Parallel(tl.Relu(), tl.LayerNorm()),\n",
+        "    tl.Concatenate(),\n",
+        ")\n",
+        "layer_block.initialize_once(input_shape, input_dtype, prng_key)\n",
+        "y = layer_block(x)\n",
+        "\n",
+        "template = ('x:\\n{}\\n\\n'\n",
+        "            'concatenate(relu(x), layer_norm(x)):\\n{}')\n",
+        "print(template.format(x, y))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 10,
+      "metadata": {
+        "colab": {
+          "height": 442
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 6372,
+          "status": "ok",
+          "timestamp": 1569980033028,
+          "user": {
+            "displayName": "Jonni Kanerva",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
+            "userId": "12646438700820120918"
+          },
+          "user_tz": 420
+        },
+        "id": "1byhBplSby2a",
+        "outputId": "6d1cc0da-a0a4-4b45-a654-da800634edb9"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "x:\n",
+            "[[-7. -6. -5. -4. -3.]\n",
+            " [-2. -1.  0.  1.  2.]\n",
+            " [ 3.  4.  5.  6.  7.]]\n",
+            "\n",
+            "dense(dense(x)):\n",
+            "[[ 5.83  -6.634  3.101 -3.777  6.508]\n",
+            " [ 0.025 -1.503  0.041 -0.626  0.387]\n",
+            " [-5.78   3.628 -3.018  2.526 -5.734]]\n",
+            "\n",
+            "params:\n",
+            "[(DeviceArray([[-0.355,  0.284, -0.737,  0.309, -0.613],\n",
+            "             [ 0.574, -0.127,  0.149,  0.763, -0.566],\n",
+            "             [-0.083, -0.454,  0.19 ,  0.018, -0.399],\n",
+            "             [-0.706,  0.454, -0.636,  0.022, -0.084],\n",
+            "             [ 0.422, -0.509, -0.415,  0.638, -0.469]], dtype=float32), DeviceArray([-2.072e-07, -1.046e-06,  2.496e-06,  2.705e-07, -1.035e-08],\n",
+            "            dtype=float32)), (DeviceArray([[ 0.459, -0.581,  0.693, -0.163,  0.77 ],\n",
+            "             [ 0.016,  0.735,  0.619,  0.585,  0.337],\n",
+            "             [ 0.617,  0.345, -0.162, -0.039, -0.073],\n",
+            "             [-0.122,  0.103,  0.477,  0.499,  0.199],\n",
+            "             [-0.01 , -0.713,  0.639,  0.056,  0.678]], dtype=float32), DeviceArray([-1.936e-07, -1.772e-06, -6.521e-07, -1.623e-06, -4.441e-07],\n",
+            "            dtype=float32))]\n",
+            "\n",
+            "state:\n",
+            "[(), ()]\n"
+          ]
+        }
+      ],
+      "source": [
+        "prng_key = backend.random.get_prng(0)\n",
+        "input_dtype = onp.float32\n",
+        "input_shape = (3, 5)\n",
+        "x = onp.arange(-7, 8).reshape(*input_shape).astype(input_dtype)\n",
+        "\n",
+        "layer_block = tl.Serial(\n",
+        "    tl.Dense(5),\n",
+        "    tl.Dense(5),\n",
+        ")\n",
+        "layer_block.initialize_once(input_shape, input_dtype, prng_key)\n",
+        "y = layer_block(x)\n",
+        "\n",
+        "template = ('x:\\n{}\\n\\n'\n",
+        "            'dense(dense(x)):\\n{}\\n\\n'\n",
+        "            'params:\\n{}\\n\\n'\n",
+        "            'state:\\n{}')\n",
+        "print(template.format(x, y, layer_block.params, layer_block.state))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 11,
+      "metadata": {
+        "colab": {
+          "height": 357
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 1047,
+          "status": "ok",
+          "timestamp": 1569980034098,
+          "user": {
+            "displayName": "Jonni Kanerva",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
+            "userId": "12646438700820120918"
+          },
+          "user_tz": 420
+        },
+        "id": "aRKnfN4adLqi",
+        "outputId": "568f4974-b9bf-4e5f-b9c4-dc6413068462"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "x:\n",
+            "[[-7. -6. -5. -4. -3.]\n",
+            " [-2. -1.  0.  1.  2.]\n",
+            " [ 3.  4.  5.  6.  7.]]\n",
+            "\n",
+            "dense(dense(x)):\n",
+            "[[ 10.543 -12.855   1.589   8.106  -8.503]\n",
+            " [ -0.278  -0.161  -0.644  -0.195   0.103]\n",
+            " [-11.099  12.533  -2.877  -8.497   8.709]]\n",
+            "\n",
+            "params:\n",
+            "[(DeviceArray([[-0.355,  0.284, -0.737,  0.309, -0.613],\n",
+            "             [ 0.574, -0.127,  0.149,  0.763, -0.566],\n",
+            "             [-0.083, -0.454,  0.19 ,  0.018, -0.399],\n",
+            "             [-0.706,  0.454, -0.636,  0.022, -0.084],\n",
+            "             [ 0.422, -0.509, -0.415,  0.638, -0.469]], dtype=float32), DeviceArray([-2.072e-07, -1.046e-06,  2.496e-06,  2.705e-07, -1.035e-08],\n",
+            "            dtype=float32)), ()]\n",
+            "\n",
+            "state:\n",
+            "[(), ()]\n"
+          ]
+        }
+      ],
+      "source": [
+        "prng_key = backend.random.get_prng(0)\n",
+        "input_dtype = onp.float32\n",
+        "input_shape = (3, 5)\n",
+        "x = onp.arange(-7, 8).reshape(*input_shape).astype(input_dtype)\n",
+        "\n",
+        "dense = tl.Dense(5)\n",
+        "layer_block = tl.Serial(\n",
+        "    dense,\n",
+        "    dense,\n",
+        ")\n",
+        "layer_block.initialize_once(input_shape, input_dtype, prng_key)\n",
+        "y = layer_block(x)\n",
+        "\n",
+        "template = ('x:\\n{}\\n\\n'\n",
+        "            'dense(dense(x)):\\n{}\\n\\n'\n",
+        "            'params:\\n{}\\n\\n'\n",
+        "            'state:\\n{}')\n",
+        "print(template.format(x, y, layer_block.params, layer_block.state))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "rwgiP0tK1H6p"
+      },
+      "source": [
+        "# Data Flows, Data Stack"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "llAH3cdE1UeU"
+      },
+      "source": [
+        "# Training and Using Models"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "65ite-671cTT"
+      },
+      "source": [
+        "# Defining Your Own Layer Classes"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [],
+      "last_runtime": {
+        "build_target": "//learning/deepmind/dm_python:dm_notebook",
+        "kind": "private"
+      },
+      "name": "A Conceptual, Practical Introduction to Trax Layers",
+      "provenance": [
+        {
+          "file_id": "1sF8QbqJ19ZU6oy5z4GUTt4lgUCjqO6kt",
+          "timestamp": 1569980697572
+        },
+        {
+          "file_id": "1EH76AWQ_pvT4i8ZXfkv-SCV4MrmllEl5",
+          "timestamp": 1563927451951
+        }
+      ]
+    },
+    "kernelspec": {
+      "display_name": "Python 2",
+      "name": "python2"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

From 1e5e7b4a602b0f629c6e127422db692556b9c3cd Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 2 Oct 2019 19:07:00 -0700
Subject: [PATCH 2527/2720] Move the tf.autograph.to_graph annotation just
 before using the _scan_step_fn.

Intended to fix this -
https://travis-ci.org/tensorflow/tensor2tensor/jobs/587720856

I can't figure out why this became a problem suddenly.

PiperOrigin-RevId: 272570396
---
 tensor2tensor/data_generators/generator_utils.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 4c9394c61..30e7bb5eb 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -969,7 +969,7 @@ def _scanning_pack(self, dataset):
 
     initial_state = self._scan_initial_state()
     step_fn = functools.partial(
-        _scan_step_fn, packed_length=self._packed_length,
+        tf.autograph.to_graph(_scan_step_fn), packed_length=self._packed_length,
         queue_size=self._queue_size, spacing=self._spacing,
         num_sequences=self._num_sequences, token_dtype=self._token_dtype)
 
@@ -1020,7 +1020,6 @@ def _finalize(self, _, contents):
             "segment": segment, "position": position}
 
 
-@tf.autograph.to_graph
 def _scan_step_fn(state, example, packed_length, queue_size, spacing,
                   num_sequences, token_dtype):  # pylint: disable=g-doc-args
   """Transform function used by tf.data.experimental.scan to process an example.

From d7d65f11faa40a44e9518c706c5accf8f83fb633 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 2 Oct 2019 19:33:41 -0700
Subject: [PATCH 2528/2720] Pin tfp to 0.7.0 -- their latest version 0.8.0 pins
 cloudpickle to 1.1.1, which gym doesn't like (it needs ~1.2)

Also upgrade pip/setuptools in travis, this gives better errors so I was able
to pin this down.

PiperOrigin-RevId: 272572911
---
 oss_scripts/oss_pip_install.sh | 4 ++++
 setup.py                       | 4 ++--
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/oss_scripts/oss_pip_install.sh b/oss_scripts/oss_pip_install.sh
index 5246a8363..86c939660 100755
--- a/oss_scripts/oss_pip_install.sh
+++ b/oss_scripts/oss_pip_install.sh
@@ -5,6 +5,10 @@ set -e  # fail and exit on any command erroring
 
 : "${TF_VERSION:?}"
 
+# Make sure we have the latest pip and setuptools installed.
+pip install -q -U pip
+pip install -q -U setuptools
+
 # Make sure we have the latest version of numpy - avoid problems we were
 # seeing with Python 3
 pip install -q -U numpy
diff --git a/setup.py b/setup.py
index f5274d52c..e2415edbd 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.14.0',
+    version='1.14.1',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',
@@ -56,7 +56,7 @@
         'sympy',
         'tensorflow-datasets',
         'tensorflow-gan',
-        'tensorflow-probability',
+        'tensorflow-probability==0.7.0',
         'tqdm',
     ],
     extras_require={

From 5913ad62a95d0d136da5b730953d99ce1b86d5ce Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 2 Oct 2019 19:33:51 -0700
Subject: [PATCH 2529/2720] Disable DatasetPacking test (in Travis) till we can
 figure out why autograph errors out.

https://travis-ci.org/afrozenator/tensor2tensor/jobs/592856608

PiperOrigin-RevId: 272572935
---
 oss_scripts/oss_tests.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 553acb3ac..82bdf784a 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -49,6 +49,7 @@ pytest --disable-warnings \
   --ignore=tensor2tensor/data_generators/ops/pack_sequences_ops_test.py \
   --ignore=tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py \
   --ignore=tensor2tensor/data_generators/problem_test.py \
+  --deselect=tensor2tensor/data_generators/generator_utils_test.py::GeneratorUtilsTest.testDatasetPacking \
   tensor2tensor/data_generators
 set_status
 

From 8e238927a4cc58ca5cdfd7a02fd698a960ebb4aa Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 2 Oct 2019 22:27:49 -0700
Subject: [PATCH 2530/2720] Fix deselect syntax :/

PiperOrigin-RevId: 272590689
---
 oss_scripts/oss_tests.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 82bdf784a..2c956e138 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -49,7 +49,7 @@ pytest --disable-warnings \
   --ignore=tensor2tensor/data_generators/ops/pack_sequences_ops_test.py \
   --ignore=tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py \
   --ignore=tensor2tensor/data_generators/problem_test.py \
-  --deselect=tensor2tensor/data_generators/generator_utils_test.py::GeneratorUtilsTest.testDatasetPacking \
+  --deselect=tensor2tensor/data_generators/generator_utils_test.py::GeneratorUtilsTest::testDatasetPacking \
   tensor2tensor/data_generators
 set_status
 

From 9f29518aa672a1d26f47ad94f4eb1dc987d9682b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 3 Oct 2019 10:07:39 -0700
Subject: [PATCH 2531/2720] Remove unused params arg from MulConstant.

PiperOrigin-RevId: 272685588
---
 tensor2tensor/trax/layers/core.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
index d6a338a9c..dd62ad85e 100644
--- a/tensor2tensor/trax/layers/core.py
+++ b/tensor2tensor/trax/layers/core.py
@@ -211,8 +211,7 @@ def AddConstant(x, constant=0.0, **unused_kwargs):
 
 
 @base.layer()
-def MulConstant(x, params, constant=1.0, **unused_kwargs):
-  del params
+def MulConstant(x, constant=1.0, **unused_kwargs):
   return x * constant
 
 

From 176148c9e2a0ff520dd2f108a9182075874a442b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 3 Oct 2019 11:42:15 -0700
Subject: [PATCH 2532/2720] Fix attention rng mismatch between forward and
 reverse direction

PiperOrigin-RevId: 272707157
---
 tensor2tensor/trax/layers/attention.py        |  3 +-
 tensor2tensor/trax/layers/reversible.py       |  8 +--
 .../trax/models/research/reformer.py          |  8 ++-
 .../trax/models/research/reformer_test.py     | 60 +++++++++++++++++++
 4 files changed, 72 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index a276a8b7c..32f2b6640 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -344,7 +344,8 @@ def new_params_and_state(self, input_shape, input_dtype, rng):
 class BaseCausalAttention(base.Layer):
   """Base class for variants of causal self-attention."""
 
-  def __init__(self):
+  def __init__(self, mode='train'):
+    del mode
     super(BaseCausalAttention, self).__init__(n_inputs=3)
 
   def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
diff --git a/tensor2tensor/trax/layers/reversible.py b/tensor2tensor/trax/layers/reversible.py
index 5b52aeaee..29f845244 100644
--- a/tensor2tensor/trax/layers/reversible.py
+++ b/tensor2tensor/trax/layers/reversible.py
@@ -101,8 +101,8 @@ def reverse(self, output, params=(), state=(), **kwargs):
       rngs = backend.random.split(rng, self._n_layers)
 
     layer_val = output
-    for layer, p, s, rng in reversed(zip(self.sublayers,
-                                         params, state, rngs)):
+    for layer, p, s, rng in reversed(list(zip(self.sublayers,
+                                              params, state, rngs))):
       layer_val = layer.reverse(layer_val, p, s, rng=rng, **kwargs)
 
     return layer_val
@@ -116,8 +116,8 @@ def reverse_and_grad(self, output, ct, params=(), state=(), **kwargs):
     layer_val = output
     layer_ct = ct
     params_ct = []
-    for layer, p, s, rng in reversed(zip(self.sublayers,
-                                         params, state, rngs)):
+    for layer, p, s, rng in reversed(list(zip(self.sublayers,
+                                              params, state, rngs))):
       layer_val, layer_ct = layer.reverse_and_grad(
           layer_val, layer_ct, p, s, rng=rng, **kwargs)
       layer_ct, p_ct = layer_ct
diff --git a/tensor2tensor/trax/models/research/reformer.py b/tensor2tensor/trax/models/research/reformer.py
index 56b8a5aa5..913343e67 100644
--- a/tensor2tensor/trax/models/research/reformer.py
+++ b/tensor2tensor/trax/models/research/reformer.py
@@ -254,14 +254,18 @@ def __init__(self, attention):
     super(ApplyAttentionWrapper, self).__init__(attention, [], [])
     self.attention = attention
 
-  def forward_and_backward(self, inputs, ct, **kwargs):
+  def forward_and_backward(self, inputs, ct, rng=None, **kwargs):
     # Simultaneous forward pass and backprop through the attention mechanism.
     qkv = inputs[:3]
     passthrough = inputs[3:]
     out_ct = ct[0]
     passthrough_ct = ct[1:]
+    if rng is not None:
+      # Adjust RNG to match the forward pass.
+      rng = backend.random.split(rng, self._n_layers)[0]
 
-    out, qkv_ct = self.attention.forward_and_backward(qkv, out_ct, **kwargs)
+    out, qkv_ct = self.attention.forward_and_backward(
+        qkv, out_ct, rng=rng, **kwargs)
     return (out,) + passthrough, qkv_ct + passthrough_ct
 
 
diff --git a/tensor2tensor/trax/models/research/reformer_test.py b/tensor2tensor/trax/models/research/reformer_test.py
index 799939748..0a8bcdad8 100644
--- a/tensor2tensor/trax/models/research/reformer_test.py
+++ b/tensor2tensor/trax/models/research/reformer_test.py
@@ -21,10 +21,43 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
+import jax
+import numpy as onp
+
+from tensor2tensor.trax import backend
 from tensor2tensor.trax import layers as tl
+from tensor2tensor.trax.backend import numpy as np
 from tensor2tensor.trax.models.research import reformer
 
 
+class PoisonOnRNGMismatchAttention(tl.BaseCausalAttention):
+  """Fills gradients with NaNs if reverse rng does not match forward rng."""
+
+  # pylint: disable=protected-access
+  def forward_and_backward(self, inputs, ct, rng=None, **kwargs):
+    assert backend.get_name() == 'jax', (
+        'JAX backend is required to use forward_and_backward.')
+
+    if ct is not None and tl.Layer._STASH_OUT is not None:
+      recovered_rng = tl.Layer._STASH_OUT.pop(self)
+      is_same = (rng[0] == recovered_rng[0]) & (rng[1] == recovered_rng[1])
+      is_same = is_same.astype(np.float32)
+      # Divides by zero if rngs are not the same, which results in NaNs.
+      inputs = (inputs[0] / is_same, inputs[1] / is_same, inputs[2] / is_same)
+
+    def _do_forward(x):  # pylint: disable=invalid-name
+      res, _ = self.forward(x, rng=rng, **kwargs)
+      return res
+    output, vjpfun = jax.vjp(_do_forward, inputs)
+    return output, vjpfun(ct)[0]
+
+  def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
+    if tl.Layer._STASH_IN is not None:
+      tl.Layer._STASH_IN[self] = rng
+    return inputs[2], state
+  # pylint: enable=protected-access
+
+
 class ReformerTest(parameterized.TestCase):
 
   def test_reformer_lm_forward_shape(self):
@@ -39,6 +72,33 @@ def test_reformer_lm_forward_shape(self):
         model, tuple(input_shape), integer_inputs=True)
     self.assertEqual(((1, 8, 16), (1, 8, 16)), final_shape)
 
+  def test_reformer_rng_consistency(self):
+    with backend.use_backend('jax'):
+      vocab_size = 16
+      batch_size = 1
+      input_shape = ((batch_size, 8), (batch_size, 8))
+      model = reformer.ReformerLM(
+          vocab_size, d_model=32, d_ff=64,
+          d_attention_key=16, d_attention_value=16, n_layers=1, n_heads=2,
+          max_len=16, n_chunks=2, n_attention_chunks=1, mode='train',
+          attention_type=PoisonOnRNGMismatchAttention)
+
+      rng = backend.random.get_prng(0)
+      params, state = model.initialize_once(
+          input_shape, (np.int32, np.int32), rng)
+
+      def dummy_loss_fn(params):
+        inputs = (np.zeros(input_shape[0], dtype=np.int32),) * 2
+        output = model(inputs, params=params, state=state, rng=rng)
+        dummy_loss = backend.numpy.sum(output[0])
+        return dummy_loss
+
+      grad_fn = backend.grad(dummy_loss_fn)
+      grads = grad_fn(params)
+      # PoisonOnRNGMismatchAttention uses NaNs to signal an rng mismatch.
+      for grad in jax.tree_util.tree_leaves(grads):
+        assert onp.all(onp.isfinite(grad))
+
 
 if __name__ == '__main__':
   absltest.main()

From 232225f06efe19e59ecb322f507b048eef499bde Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 3 Oct 2019 12:36:55 -0700
Subject: [PATCH 2533/2720] Execute `jupyter nbconvert` on python3 instead of
 python2.

python2 fails with:
`RuntimeError: Kernel died before replying to kernel_info`

python3 seems to pass.

PiperOrigin-RevId: 272718077
---
 oss_scripts/oss_tests.sh | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index 2c956e138..ede0fa1f8 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -166,9 +166,18 @@ then
     # Can't add disable warning here since it parses flags.
     pytest tensor2tensor/rl/trainer_model_based_test.py
     set_status
-    jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/hello_t2t.ipynb
+fi
+
+if [[ "$TRAVIS_PYTHON_VERSION" == "3.6" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]
+then
+    jupyter nbconvert --ExecutePreprocessor.kernel_name=python3 \
+      --ExecutePreprocessor.timeout=600 --to notebook --execute \
+      tensor2tensor/notebooks/hello_t2t.ipynb;
     set_status
-    jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute tensor2tensor/notebooks/t2t_problem.ipynb;
+
+    jupyter nbconvert --ExecutePreprocessor.kernel_name=python3 \
+      --ExecutePreprocessor.timeout=600 --to notebook --execute \
+      tensor2tensor/notebooks/t2t_problem.ipynb;
     set_status
 fi
 

From e73eb746f27bc92d85ca7e46c5fa15e41a6ede06 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 3 Oct 2019 17:14:39 -0700
Subject: [PATCH 2534/2720] Don't write summaries by default in OnlineTuneEnv.
 This takes up a lot of space and we can generate those metrics in the
 evaluator afterwards.

PiperOrigin-RevId: 272775151
---
 tensor2tensor/trax/rl/envs/online_tune_env.py |  5 ++-
 tensor2tensor/trax/trax.py                    | 38 +++++++++----------
 2 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
index 2423c50d5..df3513cc8 100644
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ b/tensor2tensor/trax/rl/envs/online_tune_env.py
@@ -78,6 +78,8 @@ def __init__(self,
                # Don't save checkpoints by default, as they tend to use a lot of
                # space.
                should_save_checkpoints=False,
+               # Same here.
+               should_write_summaries=False,
                has_weights=False,
                mask_id=None):
     if action_multipliers is None:
@@ -93,7 +95,8 @@ def __init__(self,
         optimizer=optimizer,
         lr_schedule=(lambda history: lambda step: self._current_controls),
         inputs=inputs,
-        should_save=should_save_checkpoints,
+        should_save_checkpoints=should_save_checkpoints,
+        should_write_summaries=should_write_summaries,
         nontrainable_param_map=nontrainable_param_map,
         has_weights=has_weights,
         mask_id=mask_id,
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
index 995095070..8ba8b0d3a 100644
--- a/tensor2tensor/trax/trax.py
+++ b/tensor2tensor/trax/trax.py
@@ -437,12 +437,14 @@ class Trainer(object):
 
   def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs,
                output_dir=None, random_seed=None, n_devices=None,
-               save_steps=None, should_save=True, has_weights=False,
+               save_steps=None, should_save_checkpoints=True,
+               should_write_summaries=True, has_weights=False,
                nontrainable_param_map=None, mask_id=None):
     if save_steps is None:
       save_steps = []
     self._save_steps = save_steps
-    self._should_save = should_save
+    self._should_save_checkpoints = should_save_checkpoints
+    self._should_write_summaries = should_write_summaries
     self._has_weights = has_weights
     self._mask_id = mask_id
     loss_fn = loss_fn(has_weights=has_weights, mask_id=mask_id)
@@ -589,8 +591,9 @@ def reset(self, output_dir):
     self._output_dir = output_dir
     gfile.makedirs(output_dir)
     # Create summary writers and history.
-    self._train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
-    self._eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
+    if self._should_write_summaries:
+      self._train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
+      self._eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
 
     # Reset the train and eval streams.
     self._train_stream = self._inputs.train_stream()
@@ -652,7 +655,7 @@ def _maybe_replicate(self, x):
       return x
 
   def _maybe_save_state(self, keep):
-    if self._should_save:
+    if self._should_save_checkpoints:
       _save_replicated(self._opt_state, self._step, self._history,
                        self._model_state, self._n_devices, self._output_dir,
                        keep)
@@ -737,7 +740,7 @@ def train_epoch(self, epoch_steps, eval_steps):
         self._maybe_save_state(keep=True)
 
       # Log nontrainable params (learning rate, dropout etc.)
-      if self._step == 1 or self._step % 10 == 0:
+      if (self._step == 1 or self._step % 10 == 0) and self._train_sw:
         for (name, value) in self.nontrainable_params.items():
           self._train_sw.scalar("training/{}".format(name), value)
 
@@ -745,7 +748,7 @@ def train_epoch(self, epoch_steps, eval_steps):
     epoch_time = time.time() - start_time
     step_log(self._step, "Ran %d train steps in %0.2f secs" %
              (epoch_steps, epoch_time))
-    if epoch_steps > 1:
+    if epoch_steps > 1 and self._train_sw:
       self._train_sw.scalar("training/steps per second",
                             epoch_steps / epoch_time, step=self._step)
 
@@ -756,8 +759,9 @@ def train_epoch(self, epoch_steps, eval_steps):
     self._maybe_save_state(keep=False)
 
     # Flush summary writers
-    self._train_sw.flush()
-    self._eval_sw.flush()
+    if self._train_sw:
+      self._train_sw.flush()
+      self._eval_sw.flush()
 
   def evaluate(self, eval_steps):
     """Evaluate the model and log metrics."""
@@ -771,15 +775,13 @@ def evaluate(self, eval_steps):
     train_eval_slice = itertools.islice(self._train_eval_stream, eval_steps)
     train_metrics, _ = evaluation_round(
         train_eval_slice, self._metrics, self._jit_eval, params, state, rng)
-    if self._train_sw:
-      log_metrics(train_metrics, self._train_sw, "train",
-                  self._step, history=self._history)
+    log_metrics(train_metrics, self._train_sw, "train",
+                self._step, history=self._history)
     eval_slice = itertools.islice(self._eval_stream, eval_steps)
     eval_metrics, _ = evaluation_round(
         eval_slice, self._metrics, self._jit_eval, params, state, rng)
-    if self._eval_sw:
-      log_metrics(eval_metrics, self._eval_sw, "eval",
-                  self._step, history=self._history)
+    log_metrics(eval_metrics, self._eval_sw, "eval",
+                self._step, history=self._history)
     step_log(self._step, "Finished evaluation")
 
     # Save the optimizer params in the history
@@ -857,10 +859,8 @@ def evaluate(self, eval_steps):
       train_eval_metrics.append(metrics)
     # Unpack in the same order we've iterated over streams in the loop above.
     train_metrics, eval_metrics = train_eval_metrics  # pylint: disable=unbalanced-tuple-unpacking
-    if self._train_sw:
-      log_metrics(train_metrics, self._train_sw, "train", step, history=history)
-    if self._eval_sw:
-      log_metrics(eval_metrics, self._eval_sw, "eval", step, history=history)
+    log_metrics(train_metrics, self._train_sw, "train", step, history=history)
+    log_metrics(eval_metrics, self._eval_sw, "eval", step, history=history)
     step_log(step, "Finished evaluation")
 
   def save_computation_graphs(self, save_backward_graph):

From ec5c72e108817dcab850ac49ff787185423bc65b Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Thu, 3 Oct 2019 17:16:18 -0700
Subject: [PATCH 2535/2720] Allow to factorize hash creation to lower hashing
 complexity.

PiperOrigin-RevId: 272775435
---
 .../trax/configs/transformer_copy.gin         | 19 +++----
 tensor2tensor/trax/layers/attention.py        | 53 +++++++++++++++++--
 2 files changed, 59 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/trax/configs/transformer_copy.gin b/tensor2tensor/trax/configs/transformer_copy.gin
index 2cc2d29fc..86bfbead8 100644
--- a/tensor2tensor/trax/configs/transformer_copy.gin
+++ b/tensor2tensor/trax/configs/transformer_copy.gin
@@ -4,12 +4,12 @@ import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
 n_symbols = 128
-length = 1024
-batch = 32
+length = 4096
+batch = 16
 
 # Parameters for batch_fn:
 # ==============================================================================
-batch_fn.batch_size_per_device = %batch
+batch_fn.batch_size = %batch
 batch_fn.eval_batch_size = %batch
 batch_fn.max_eval_length = %length
 
@@ -39,7 +39,7 @@ train.eval_steps = 64
 train.inputs = @trax.inputs.sequence_copy_inputs
 train.model = @trax.models.TransformerLM
 train.optimizer = @trax.optimizers.Adafactor
-train.train_steps = 150000
+train.train_steps = 50000
 train.has_weights = True
 
 # Parameters for MemoryEfficientCausalAttention:
@@ -50,15 +50,16 @@ MemoryEfficientCausalAttention.loop_stride = 512
 # Parameters for LSHCausalAttention:
 # ==============================================================================
 LSHCausalAttention.allow_duplicate_attention = False
-LSHCausalAttention.attend_across_buckets = False
+LSHCausalAttention.attend_across_buckets = True
 LSHCausalAttention.rehash_each_round = True
-LSHCausalAttention.n_bins = 32
-LSHCausalAttention.n_buckets = 64
-LSHCausalAttention.n_hashes = 4
+LSHCausalAttention.n_bins = 64
+LSHCausalAttention.n_buckets = 128
+LSHCausalAttention.n_hashes = 8
 LSHCausalAttention.one_rng = False
 LSHCausalAttention.hard_k = 0
 LSHCausalAttention.dropout = 0.0
-LSHCausalAttention.drop_for_hash_rate = 0.0
+LSHCausalAttention.drop_for_hash_rate = 0.1
+LSHCausalAttention.factorize_hash = True
 
 # Parameters for TransformerLM:
 # ==============================================================================
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
index 32f2b6640..0a3b7a49c 100644
--- a/tensor2tensor/trax/layers/attention.py
+++ b/tensor2tensor/trax/layers/attention.py
@@ -854,7 +854,7 @@ class LSHCausalAttention(BaseCausalAttention):
 
   def __init__(self, dropout, mode, n_bins=64, n_hashes=1, n_buckets=64,
                one_rng=False, allow_duplicate_attention=False,
-               attend_across_buckets=False, hard_k=0,
+               attend_across_buckets=False, hard_k=0, factorize_hash=False,
                rehash_each_round=True, drop_for_hash_rate=0.0):
     del dropout
     self._mode = mode
@@ -868,6 +868,7 @@ def __init__(self, dropout, mode, n_bins=64, n_hashes=1, n_buckets=64,
     self.n_buckets = n_buckets
     self._drop_for_hash_rate = drop_for_hash_rate
     self._one_rng = one_rng
+    self._factorize_hash = factorize_hash
     self._prng = None
     if one_rng:
       seed = random.randint(0, 2**31 - 1)
@@ -1027,10 +1028,38 @@ def hash_vectors(self, vecs, rng):
     # We sample a different random rotation for each round of hashing to
     # decrease the probability of hash misses.
     assert self.n_buckets % 2 == 0
+
+    # If we factorize the hash, find a factor dividing n_buckets nicely.
+    rot_size, factor_list = self.n_buckets, [self.n_buckets]
+    if self._factorize_hash:
+      # If we are given a list of factors, verify it and use later.
+      if isinstance(self._factorize_hash, list):
+        rot_size, product = 0, 1
+        factor_list = self._factorize_hash
+        for factor in factor_list:
+          assert factor % 2 == 0
+          product *= factor
+          rot_size += factor
+        assert product == self.n_buckets
+      else:  # Find one factor if just set to True.
+        # We want to represent self.n_buckets = factor * rest so that
+        # (1) both factor and rest are even, and (2) factor + rest is minimal.
+        # To compute this we start from factor = sqrt(n_buckets) and go down
+        # with it until we find one that satisfies the constraints above.
+        factor = int(math.sqrt(self.n_buckets))
+        while factor > 0 and not (
+            self.n_buckets % factor == 0 and
+            factor % 2 == 0 and
+            (self.n_buckets // factor) % 2 == 0):
+          factor -= 1
+        if factor > 2:  # Factor of 2 does not warrant the effort.
+          rot_size = factor + (self.n_buckets // factor)
+          factor_list = [factor, self.n_buckets // factor]
+
     random_rotations_shape = (
         vecs.shape[-1],
         self.n_hashes if self._rehash_each_round else 1,
-        self.n_buckets // 2)
+        rot_size // 2)
 
     rng = jax.lax.tie_in(vecs, rng)
     rng, subrng = backend.random.split(rng)
@@ -1040,16 +1069,32 @@ def hash_vectors(self, vecs, rng):
     # hashing, so it's shared between them. Check if that's what we want.
     dropped_vecs = self.drop_for_hash(vecs, subrng)
     rotated_vecs = np.einsum('tf,fhb->htb', dropped_vecs, random_rotations)
-    rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)
 
     if self._rehash_each_round:
-      buckets = np.argmax(rotated_vecs, axis=-1)
+      if self._factorize_hash and len(factor_list) > 1:
+        # We factorized self.n_buckets as the product of factor_list.
+        # Get the buckets for them and combine.
+        buckets, cur_sum, cur_product = None, 0, 1
+        for factor in factor_list:
+          rv = rotated_vecs[..., cur_sum:cur_sum + (factor // 2)]
+          cur_sum += factor // 2
+          rv = np.concatenate([rv, -rv], axis=-1)
+          if buckets is None:
+            buckets = np.argmax(rv, axis=-1)
+          else:
+            buckets += cur_product * np.argmax(rv, axis=-1)
+          cur_product *= factor
+      else:
+        rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)
+        buckets = np.argmax(rotated_vecs, axis=-1)
       # buckets is now (self.n_hashes, seqlen). Next we add offsets so that
       # bucket numbers from different hashing rounds don't overlap.
       offsets = jax.lax.tie_in(buckets, np.arange(self.n_hashes))
       offsets = np.reshape(offsets * self.n_buckets, (-1, 1))
       buckets = np.reshape(buckets + offsets, (-1,))
     else:
+      assert not self._factorize_hash
+      rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)
       # In this configuration, we map each item to the top self.n_hashes buckets
       rotated_vecs = np.squeeze(rotated_vecs, 0)
       bucket_range = jax.lax.tie_in(vecs, np.arange(rotated_vecs.shape[-1]))

From e258b4d8bd40ed634ffece644f4f05c37a7afc03 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 3 Oct 2019 17:17:48 -0700
Subject: [PATCH 2536/2720] Update the world model -> policy parameter
 initialization code after the change making loss functions layers.

PiperOrigin-RevId: 272775707
---
 tensor2tensor/trax/rl/ppo.py      |  2 +-
 tensor2tensor/trax/rl/ppo_test.py | 25 +++++++++++++++++--------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
index 99d4741bc..15338d1fa 100644
--- a/tensor2tensor/trax/rl/ppo.py
+++ b/tensor2tensor/trax/rl/ppo.py
@@ -923,7 +923,7 @@ def init_policy_from_world_model_checkpoint(policy_params, model_output_dir):
   # TODO(pkozakowski): The following, brittle line of code is hardcoded for
   # transplanting parameters from TransformerLM to TransformerDecoder-based
   # policy network of the same configuration. Figure out a more general method.
-  policy_params[0] = model_params[1:-2]
+  policy_params[0] = model_params[0][1:-2]
   return policy_params
 
 
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
index e0046b103..1e48dc87d 100644
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ b/tensor2tensor/trax/rl/ppo_test.py
@@ -25,6 +25,7 @@
 import jax
 from jax import random as jax_random
 import numpy as np
+from tensor2tensor.trax import inputs
 from tensor2tensor.trax import layers
 from tensor2tensor.trax import models
 from tensor2tensor.trax import trax
@@ -578,8 +579,21 @@ def test_inits_policy_by_world_model_checkpoint(self):
         "input_dtype": np.int32,
         "rng": rng,
     }
-    model = models.TransformerLM(vocab_size=4, **transformer_kwargs)
-    (model_params, _) = model.initialize_once(**init_kwargs)
+    model_fn = functools.partial(
+        models.TransformerLM, vocab_size=4, **transformer_kwargs
+    )
+    output_dir = self.get_temp_dir()
+    # Initialize a world model checkpoint by running the trainer.
+    trax.train(
+        output_dir,
+        model=model_fn,
+        inputs=functools.partial(
+            inputs.random_inputs, input_shape=(1, 1), output_shape=(1, 1)
+        ),
+        train_steps=1,
+        eval_steps=1,
+    )
+
     policy = ppo.policy_and_value_net(
         n_actions=3,
         n_controls=2,
@@ -590,12 +604,7 @@ def test_inits_policy_by_world_model_checkpoint(self):
         two_towers=False,
     )
     (policy_params, policy_state) = policy.initialize_once(**init_kwargs)
-    output_dir = self.get_temp_dir()
-    # Initialize state by restoring from a nonexistent checkpoint.
-    trax_state = trax.restore_state(output_dir)
-    trax_state = trax_state._replace(opt_state=(model_params, None))
-    # Save world model parameters.
-    trax.save_state(trax_state, output_dir)
+
     # Initialize policy parameters from world model parameters.
     new_policy_params = ppo.init_policy_from_world_model_checkpoint(
         policy_params, output_dir

From 4ea0ef7f956055c55ab2692177573312c89570e1 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Thu, 3 Oct 2019 19:16:38 -0700
Subject: [PATCH 2537/2720] Switch {MemoryEfficient -> TimeBin}CausalAttention
 and add a config for SimPLe with a serialized policy.

PiperOrigin-RevId: 272791159
---
 .../trax/rl/configs/simple_online_tune.gin    |  17 ++-
 .../configs/simple_online_tune_serialized.gin | 100 ++++++++++++++++++
 2 files changed, 107 insertions(+), 10 deletions(-)
 create mode 100644 tensor2tensor/trax/rl/configs/simple_online_tune_serialized.gin

diff --git a/tensor2tensor/trax/rl/configs/simple_online_tune.gin b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
index 5a9878af6..ac67788f0 100644
--- a/tensor2tensor/trax/rl/configs/simple_online_tune.gin
+++ b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
@@ -11,7 +11,7 @@ BoxSpaceSerializer.precision = 2
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
-world_model/MultifactorSchedule.constant = 3.0
+world_model/MultifactorSchedule.constant = 1.0
 world_model/MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
 world_model/MultifactorSchedule.warmup_steps = 10000
 
@@ -52,10 +52,10 @@ PPO.policy_and_value_optimizer = @trax.optimizers.Adam
 PPO.trajectory_dump_min_count_per_shard = 8
 PPO.print_every_optimizer_steps = 1
 
-## Parameters for MemoryEfficientCausalAttention:
+## Parameters for TimeBinCausalAttention:
 ## ==============================================================================
-world_model/MemoryEfficientCausalAttention.dropout = 0.3
-world_model/MemoryEfficientCausalAttention.loop_stride = 2
+world_model/TimeBinCausalAttention.dropout = 0.1
+world_model/TimeBinCausalAttention.bin_length = 512
 
 # Parameters for SerializedSequenceSimulatedEnvProblem:
 # ==============================================================================
@@ -65,16 +65,13 @@ SerializedSequenceSimulatedEnvProblem.done_fn = @trax.rl.onlinetune_done_fn
 SerializedSequenceSimulatedEnvProblem.vocab_size = 128
 SerializedSequenceSimulatedEnvProblem.max_trajectory_length = 101
 SerializedSequenceSimulatedEnvProblem.significance_decay = 0.8
-SerializedSequenceSimulatedEnvProblem.model_predict_kwargs = {
-    "attention_type": @trax.layers.DotProductCausalAttention,
-}
 
 # Parameters for SimPLe:
 # ==============================================================================
 SimPLe.policy_trainer_class = @trax.rl.trainers.PPO
 SimPLe.n_real_epochs = 1
 SimPLe.n_model_initial_train_steps = 50000
-SimPLe.n_model_train_steps_per_epoch = 20000
+SimPLe.n_model_train_steps_per_epoch = 10000
 SimPLe.model_train_batch_size = 64
 SimPLe.simulated_env_problem_class = @trax.rl.SerializedSequenceSimulatedEnvProblem
 SimPLe.simulated_batch_size = 128
@@ -84,12 +81,12 @@ SimPLe.init_policy_from_world_model = False
 
 # Parameters for TransformerLM:
 # ==============================================================================
-world_model/TransformerLM.attention_type = @world_model/trax.layers.MemoryEfficientCausalAttention
+world_model/TransformerLM.attention_type = @world_model/trax.layers.TimeBinCausalAttention
 world_model/TransformerLM.d_model = 256
 world_model/TransformerLM.d_ff = 512
 world_model/TransformerLM.n_layers = 3
 world_model/TransformerLM.n_heads = 4
-world_model/TransformerLM.dropout = 0.3
+world_model/TransformerLM.dropout = 0.1
 world_model/TransformerLM.max_len = 2048
 
 # Parameters for train:
diff --git a/tensor2tensor/trax/rl/configs/simple_online_tune_serialized.gin b/tensor2tensor/trax/rl/configs/simple_online_tune_serialized.gin
new file mode 100644
index 000000000..181efb64c
--- /dev/null
+++ b/tensor2tensor/trax/rl/configs/simple_online_tune_serialized.gin
@@ -0,0 +1,100 @@
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+import tensor2tensor.trax.rl
+import tensor2tensor.trax.rl.space_serializer
+import tensor2tensor.trax.rl.trainers
+
+# Parameters for BoxSpaceSerializer:
+# ==============================================================================
+BoxSpaceSerializer.precision = 2
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+world_model/MultifactorSchedule.constant = 1.0
+world_model/MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+world_model/MultifactorSchedule.warmup_steps = 10000
+
+# Parameters for Adam:
+# ==============================================================================
+Adam.learning_rate = 1e-3
+Adam.b1 = 0.9
+Adam.b2 = 0.999
+Adam.weight_decay_rate = 0.0
+
+# Parameters for TransformerDecoder:
+# ==============================================================================
+TransformerDecoder.attention_type = @policy/trax.layers.TimeBinCausalAttention
+TransformerDecoder.d_model = 256
+TransformerDecoder.d_ff = 512
+TransformerDecoder.n_layers = 3
+TransformerDecoder.n_heads = 4
+TransformerDecoder.dropout = 0.0
+
+# Parameters for PPO:
+# ==============================================================================
+PPO.n_optimizer_steps = 20
+PPO.optimizer_batch_size = 64
+PPO.target_kl = 0.1
+PPO.boundary = 100
+PPO.max_timestep = 100
+PPO.max_timestep_eval = 100
+PPO.random_seed = None
+PPO.gamma = 1.0
+PPO.lambda_ = 0.95
+PPO.c1 = 1.0
+PPO.c2 = 0.1
+PPO.done_frac_for_policy_save = 0
+PPO.len_history_for_policy = None
+PPO.separate_eval = False
+PPO.save_every_n = 1
+PPO.policy_and_value_model = @trax.models.TransformerDecoder
+PPO.policy_and_value_optimizer = @trax.optimizers.Adam
+PPO.policy_and_value_vocab_size = 128
+PPO.trajectory_dump_min_count_per_shard = 8
+PPO.print_every_optimizer_steps = 1
+
+## Parameters for TimeBinCausalAttention:
+## ==============================================================================
+world_model/TimeBinCausalAttention.dropout = 0.1
+world_model/TimeBinCausalAttention.bin_length = 512
+
+policy/TimeBinCausalAttention.dropout = 0.0
+policy/TimeBinCausalAttention.bin_length = 512
+
+# Parameters for SerializedSequenceSimulatedEnvProblem:
+# ==============================================================================
+SerializedSequenceSimulatedEnvProblem.model = @world_model/trax.models.TransformerLM
+SerializedSequenceSimulatedEnvProblem.reward_fn = @trax.rl.onlinetune_reward_fn
+SerializedSequenceSimulatedEnvProblem.done_fn = @trax.rl.onlinetune_done_fn
+SerializedSequenceSimulatedEnvProblem.vocab_size = 128
+SerializedSequenceSimulatedEnvProblem.max_trajectory_length = 101
+SerializedSequenceSimulatedEnvProblem.significance_decay = 0.8
+
+# Parameters for SimPLe:
+# ==============================================================================
+SimPLe.policy_trainer_class = @trax.rl.trainers.PPO
+SimPLe.n_real_epochs = 1
+SimPLe.n_model_initial_train_steps = 50000
+SimPLe.n_model_train_steps_per_epoch = 10000
+SimPLe.model_train_batch_size = 64
+SimPLe.simulated_env_problem_class = @trax.rl.SerializedSequenceSimulatedEnvProblem
+SimPLe.simulated_batch_size = 128
+SimPLe.n_simulated_epochs = 50
+SimPLe.initial_trajectory_mix_prob = 0.9
+SimPLe.init_policy_from_world_model = True
+
+# Parameters for TransformerLM:
+# ==============================================================================
+world_model/TransformerLM.attention_type = @world_model/trax.layers.TimeBinCausalAttention
+world_model/TransformerLM.d_model = 256
+world_model/TransformerLM.d_ff = 512
+world_model/TransformerLM.n_layers = 3
+world_model/TransformerLM.n_heads = 4
+world_model/TransformerLM.dropout = 0.1
+world_model/TransformerLM.max_len = 2048
+
+# Parameters for train:
+# ==============================================================================
+world_model/train.eval_frequency = 1000
+world_model/train.optimizer = @trax.optimizers.Adafactor

From e9d6b10cd1e74f0966900065d1006f1af814499e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 4 Oct 2019 09:00:34 -0700
Subject: [PATCH 2538/2720] General progress on the first section (Layers) of
 the intro notebook.

PiperOrigin-RevId: 272886830
---
 tensor2tensor/trax/layers/intro.ipynb | 484 ++++++++++++--------------
 1 file changed, 229 insertions(+), 255 deletions(-)

diff --git a/tensor2tensor/trax/layers/intro.ipynb b/tensor2tensor/trax/layers/intro.ipynb
index 680b18b0b..0658786f1 100644
--- a/tensor2tensor/trax/layers/intro.ipynb
+++ b/tensor2tensor/trax/layers/intro.ipynb
@@ -11,7 +11,7 @@
         "\n",
         "This notebook introduces the core concepts and programming components of the Trax library through a series of code samples and explanations. The topics covered in following sections are:\n",
         "  - **layers**: the basic building blocks and how to combine them into networks\n",
-        "  - **data flows**, data stack: how the Trax runtime moves data through the layers\n",
+        "  - **data flows, data stack**: how the Trax runtime moves data through the layers\n",
         "  - **models**: how to train, evaluate, and run predictions with Trax models\n",
         "  - **new layer classes**: how to define and test your own Layer classes\n",
         "\n",
@@ -70,18 +70,18 @@
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 267,
+          "elapsed": 383,
           "status": "ok",
-          "timestamp": 1569980014037,
+          "timestamp": 1570168980195,
           "user": {
-            "displayName": "Jonni Kanerva",
-            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
-            "userId": "12646438700820120918"
+            "displayName": "",
+            "photoUrl": "",
+            "userId": ""
           },
           "user_tz": 420
         },
         "id": "vlGjGoGMTt-D",
-        "outputId": "97840e48-52ee-4ae1-87f1-83c1b0251055"
+        "outputId": "6d2ecf3d-3eb8-48a7-ad12-ebefe83afaf1"
       },
       "outputs": [
         {
@@ -113,7 +113,17 @@
       },
       "outputs": [],
       "source": [
-        "onp.set_printoptions(precision=3)  # Less visual noise in the numerical outputs."
+        "onp.set_printoptions(precision=3)  # Less visual noise in the numerical outputs.\n",
+        "\n",
+        "def show_layer_properties(layer_obj, layer_name):\n",
+        "  template = ('{}.n_inputs:  {}\\n'\n",
+        "              '{}.n_outputs: {}\\n'\n",
+        "              '{}.sublayers: {}\\n'\n",
+        "              '{}.params:    {}\\n')\n",
+        "  print(template.format(layer_name, layer_obj.n_inputs,\n",
+        "                        layer_name, layer_obj.n_outputs,\n",
+        "                        layer_name, layer_obj.sublayers,\n",
+        "                        layer_name, layer_obj.params))  "
       ]
     },
     {
@@ -163,13 +173,16 @@
       "source": [
         "A layer computes a function from zero or more inputs to zero or more outputs. The inputs and outputs are NumPy arrays or JAX objects wrapping NumPy arrays.\n",
         "\n",
-        "The simplest layers, those with no parameters or state, can be used without initialization. You can think of them (and test them) like simple mathematical functions. For ease of testing and interactive exploration, layer\n",
+        "The simplest layers, those with no parameters, state or sublayers, can be used without initialization. You can think of them (and test them) like simple mathematical functions. For ease of testing and interactive exploration, layer\n",
         "objects implement the `__call__ ` method, so you can call them directly on input data:\n",
-        "\n",
         "```\n",
         "y = layer(x)\n",
         "```\n",
         "\n",
+        "Layers are also objects, so you can inspect their properties. For example:\n",
+        "```\n",
+        "print('Number of inputs required by this layer: {}'.format(layer.n_inputs))\n",
+        "```\n",
         "\n",
         "### Example 1. tl.Relu [1 input, 1 output]"
       ]
@@ -179,30 +192,28 @@
       "execution_count": 4,
       "metadata": {
         "colab": {
-          "height": 204
+          "height": 221
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 1224,
+          "elapsed": 1543,
           "status": "ok",
-          "timestamp": 1569980015601,
+          "timestamp": 1570168982080,
           "user": {
-            "displayName": "Jonni Kanerva",
-            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
-            "userId": "12646438700820120918"
+            "displayName": "",
+            "photoUrl": "",
+            "userId": ""
           },
           "user_tz": 420
         },
         "id": "V09viOSEQvQe",
-        "outputId": "b5f206b9-29ad-42b1-cd64-2a52cac4911e"
+        "outputId": "b7c1c085-3b54-4673-f284-99d6440f8a52"
       },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "type(x): \u003ctype 'numpy.ndarray'\u003e\n",
-            "\n",
             "x:\n",
             "[[-7. -6. -5. -4. -3.]\n",
             " [-2. -1.  0.  1.  2.]\n",
@@ -211,7 +222,10 @@
             "relu(x):\n",
             "[[0. 0. 0. 0. 0.]\n",
             " [0. 0. 0. 1. 2.]\n",
-            " [3. 4. 5. 6. 7.]]\n"
+            " [3. 4. 5. 6. 7.]]\n",
+            "\n",
+            "number of inputs expected by this layer: 1\n",
+            "number of outputs promised by this layer: 1\n"
           ]
         }
       ],
@@ -222,11 +236,12 @@
         "relu = tl.Relu()\n",
         "y = relu(x)\n",
         "\n",
-        "# Show the inputs and outputs.\n",
-        "template = ('type(x): {}\\n\\n'\n",
-        "            'x:\\n{}\\n\\n'\n",
-        "            'relu(x):\\n{}')\n",
-        "print(template.format(type(x), x, y))"
+        "# Show input, output, and two layer properties.\n",
+        "template = ('x:\\n{}\\n\\n'\n",
+        "            'relu(x):\\n{}\\n\\n'\n",
+        "            'number of inputs expected by this layer: {}\\n'\n",
+        "            'number of outputs promised by this layer: {}')\n",
+        "print(template.format(x, y, relu.n_inputs, relu.n_outputs))"
       ]
     },
     {
@@ -244,22 +259,22 @@
       "execution_count": 5,
       "metadata": {
         "colab": {
-          "height": 391
+          "height": 442
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 1251,
+          "elapsed": 1558,
           "status": "ok",
-          "timestamp": 1569980016872,
+          "timestamp": 1570168983657,
           "user": {
-            "displayName": "Jonni Kanerva",
-            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
-            "userId": "12646438700820120918"
+            "displayName": "",
+            "photoUrl": "",
+            "userId": ""
           },
           "user_tz": 420
         },
         "id": "LMPPNWXLoOZI",
-        "outputId": "310d6c78-7dfc-48c5-a5e9-0f0b09abb2bc"
+        "outputId": "24398ccb-9cda-4bdd-c0f0-4904c02a215e"
       },
       "outputs": [
         {
@@ -276,7 +291,7 @@
             " [-20. -10.   0.  10.  20.]\n",
             " [ 30.  40.  50.  60.  70.]]\n",
             "\n",
-            "concatenate_0([x1, x2]):\n",
+            "concat0([x1, x2]):\n",
             "[[ -7.  -6.  -5.  -4.  -3.]\n",
             " [ -2.  -1.   0.   1.   2.]\n",
             " [  3.   4.   5.   6.   7.]\n",
@@ -284,30 +299,35 @@
             " [-20. -10.   0.  10.  20.]\n",
             " [ 30.  40.  50.  60.  70.]]\n",
             "\n",
-            "concatenate_1([x1, x2]):\n",
+            "concat1([x1, x2]):\n",
             "[[ -7.  -6.  -5.  -4.  -3. -70. -60. -50. -40. -30.]\n",
             " [ -2.  -1.   0.   1.   2. -20. -10.   0.  10.  20.]\n",
-            " [  3.   4.   5.   6.   7.  30.  40.  50.  60.  70.]]\n"
+            " [  3.   4.   5.   6.   7.  30.  40.  50.  60.  70.]]\n",
+            "\n",
+            "concat0: Concatenate{in=2,out=1}\n",
+            "concat1: Concatenate{in=2,out=1}\n"
           ]
         }
       ],
       "source": [
         "x1 = onp.arange(-7, 8).reshape(3, -1).astype(onp.float32)\n",
-        "x2 = x1 * 10\n",
+        "x2 = 10 * x1\n",
         "\n",
-        "concatenate_0 = tl.Concatenate(axis=0)\n",
-        "y = concatenate_0([x1, x2])\n",
+        "concat0 = tl.Concatenate(axis=0)\n",
+        "concat1 = tl.Concatenate(axis=1)\n",
+        "\n",
+        "y0 = concat0([x1, x2])\n",
+        "y1 = concat1([x1, x2])\n",
         "\n",
         "template = ('x1:\\n{}\\n\\n'\n",
         "            'x2:\\n{}\\n\\n'\n",
-        "            'concatenate_0([x1, x2]):\\n{}')\n",
-        "print(template.format(x1, x2, y))\n",
-        "\n",
-        "concatenate_1 = tl.Concatenate(axis=1)\n",
-        "y = concatenate_1([x1, x2])\n",
+        "            'concat0([x1, x2]):\\n{}\\n\\n'\n",
+        "            'concat1([x1, x2]):\\n{}\\n')\n",
+        "print(template.format(x1, x2, y0, y1))\n",
         "\n",
-        "template = ('\\nconcatenate_1([x1, x2]):\\n{}')\n",
-        "print(template.format(y))\n"
+        "# Print abbreviated object representations (useful for debugging).\n",
+        "print('concat0: {}'.format(concat0))\n",
+        "print('concat1: {}'.format(concat1))"
       ]
     },
     {
@@ -342,7 +362,7 @@
         "    ...\n",
         "```\n",
         "\n",
-        "**Example 2.** tl.LayerNorm [1 input, 1 output, has parameters]"
+        "### Example 3. tl.LayerNorm [1 input, 1 output, has parameters]"
       ]
     },
     {
@@ -354,18 +374,18 @@
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 2200,
+          "elapsed": 2555,
           "status": "ok",
-          "timestamp": 1569980019093,
+          "timestamp": 1570168986228,
           "user": {
-            "displayName": "Jonni Kanerva",
-            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
-            "userId": "12646438700820120918"
+            "displayName": "",
+            "photoUrl": "",
+            "userId": ""
           },
           "user_tz": 420
         },
         "id": "Ie7iyX91qAx2",
-        "outputId": "b89ecf84-456f-4507-a6d3-a3e36f629ecb"
+        "outputId": "3fe02659-481b-4912-c7eb-85eb01cfadd6"
       },
       "outputs": [
         {
@@ -388,19 +408,17 @@
         }
       ],
       "source": [
-        "prng_key = backend.random.get_prng(0)\n",
-        "input_dtype = onp.float32\n",
-        "input_shape = (3, 5)\n",
-        "x = onp.arange(-7, 8).reshape(*input_shape).astype(input_dtype)\n",
+        "prng_key = backend.random.get_prng(0)  # Used below for layer initialization.\n",
+        "x = onp.arange(-7, 8).reshape(3, -1).astype(onp.float32)\n",
         "\n",
         "layer_norm = tl.LayerNorm()\n",
-        "layer_norm.initialize_once(input_shape, input_dtype, prng_key)\n",
+        "layer_norm.initialize_once(x.shape, x.dtype, prng_key)\n",
         "y = layer_norm(x)\n",
         "\n",
         "template = ('x:\\n{}\\n\\n'\n",
-        "            'layer_norm(x):\\n{}\\n\\n'\n",
-        "            'layer_norm.params:\\n{}')\n",
-        "print(template.format(x, y, layer_norm.params))"
+        "            'layer_norm(x):\\n{}\\n')\n",
+        "print(template.format(x, y))\n",
+        "print('layer_norm.params:\\n{}'.format(layer_norm.params))"
       ]
     },
     {
@@ -420,7 +438,7 @@
         "id": "d47gVdGV1vWw"
       },
       "source": [
-        "The Trax library builders encourage users, where possible, to build new layers as combinations of existing layers. The library provides a small set of _combinator_ layers for this: layer objects that make a list of layers behave as a single layer, computing outputs from inputs and updating parameters from gradients.\n",
+        "The Trax library authors encourage users, where possible, to build new layers as combinations of existing layers. The library provides a small set of _combinator_ layers for this: layer objects that make a list of layers behave as a single layer (a unit able to compute outputs from inputs, update parameters from gradients, and combine with yet more layers).\n",
         "\n"
       ]
     },
@@ -460,7 +478,9 @@
         "    layer_f,\n",
         "    layer_g,\n",
         ")\n",
-        "```"
+        "```\n",
+        "\n",
+        "### Example 4. y = layer_norm(relu(x)) [1 input, 1 output, has parameters]"
       ]
     },
     {
@@ -472,18 +492,18 @@
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 2492,
+          "elapsed": 1664,
           "status": "ok",
-          "timestamp": 1569980021601,
+          "timestamp": 1570168987915,
           "user": {
-            "displayName": "Jonni Kanerva",
-            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
-            "userId": "12646438700820120918"
+            "displayName": "",
+            "photoUrl": "",
+            "userId": ""
           },
           "user_tz": 420
         },
         "id": "dW5fpusjvjmh",
-        "outputId": "ff8d5d4e-00af-487e-d287-fc77dab366c8"
+        "outputId": "207f6a59-b767-414f-a836-ec342157ef51"
       },
       "outputs": [
         {
@@ -495,7 +515,7 @@
             " [-2. -1.  0.  1.  2.]\n",
             " [ 3.  4.  5.  6.  7.]]\n",
             "\n",
-            "layer_norm(relu(x)):\n",
+            "layer_block(x):\n",
             "[[ 0.     0.     0.     0.     0.   ]\n",
             " [-0.75  -0.75  -0.75   0.5    1.75 ]\n",
             " [-1.414 -0.707  0.     0.707  1.414]]\n"
@@ -504,32 +524,30 @@
       ],
       "source": [
         "prng_key = backend.random.get_prng(0)\n",
-        "input_dtype = onp.float32\n",
-        "input_shape = (3, 5)\n",
-        "x = onp.arange(-7, 8).reshape(*input_shape).astype(input_dtype)\n",
+        "x = onp.arange(-7, 8).reshape(3, -1).astype(onp.float32)\n",
         "\n",
         "layer_block = tl.Serial(\n",
         "    tl.Relu(),\n",
         "    tl.LayerNorm(),\n",
         ")\n",
-        "layer_block.initialize_once(input_shape, input_dtype, prng_key)\n",
+        "layer_block.initialize_once(x.shape, x.dtype, prng_key)\n",
         "y = layer_block(x)\n",
         "\n",
         "template = ('x:\\n{}\\n\\n'\n",
-        "            'layer_norm(relu(x)):\\n{}')\n",
-        "print(template.format(x, y))"
+        "            'layer_block(x):\\n{}')\n",
+        "print(template.format(x, y,))"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
         "colab_type": "text",
-        "id": "PqVNdoONcTp0"
+        "id": "bRtmN6ckQO1q"
       },
       "source": [
-        "## Combine with Parallel(...)\n",
+        "And we can inspect the block as a whole, as if it were just another layer:\n",
         "\n",
-        "\u003c... and need to introduce Dup(...) ...\u003e"
+        "### Example 5. Inspecting a Serial layer."
       ]
     },
     {
@@ -537,267 +555,223 @@
       "execution_count": 8,
       "metadata": {
         "colab": {
-          "height": 255
+          "height": 102
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 2050,
+          "elapsed": 289,
           "status": "ok",
-          "timestamp": 1569980023681,
+          "timestamp": 1570168988225,
           "user": {
-            "displayName": "Jonni Kanerva",
-            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
-            "userId": "12646438700820120918"
+            "displayName": "",
+            "photoUrl": "",
+            "userId": ""
           },
           "user_tz": 420
         },
-        "id": "teOhSJ8A8Es3",
-        "outputId": "6759189c-2ef7-4530-c9ee-88b0571208c5"
+        "id": "D6BpYddZQ1eu",
+        "outputId": "03a99733-cd84-4639-fb1f-8dfacebf5b07"
       },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "x:\n",
-            "[[-7. -6. -5. -4. -3.]\n",
-            " [-2. -1.  0.  1.  2.]\n",
-            " [ 3.  4.  5.  6.  7.]]\n",
-            "\n",
-            "relu(x):\n",
-            "[[0. 0. 0. 0. 0.]\n",
-            " [0. 0. 0. 1. 2.]\n",
-            " [3. 4. 5. 6. 7.]]\n",
+            "layer_block:\n",
+            "Serial{in=1,out=1,sublayers=[Relu{in=1,out=1}, LayerNorm{in=1,out=1}]}\n",
             "\n",
-            "layer_norm(x):\n",
-            "[[-1.414 -0.707  0.     0.707  1.414]\n",
-            " [-1.414 -0.707  0.     0.707  1.414]\n",
-            " [-1.414 -0.707  0.     0.707  1.414]]\n"
+            "layer_block.params:\n",
+            "[(), (_FilledConstant([1., 1., 1., 1., 1.], dtype=float32), _FilledConstant([0., 0., 0., 0., 0.], dtype=float32))]\n"
           ]
         }
       ],
       "source": [
-        "prng_key = backend.random.get_prng(0)\n",
-        "input_dtype = onp.float32\n",
-        "input_shape = (3, 5)\n",
-        "x = onp.arange(-7, 8).reshape(*input_shape).astype(input_dtype)\n",
+        "print('layer_block:\\n{}\\n'.format(layer_block))\n",
         "\n",
-        "layer_block = tl.Serial(\n",
-        "    tl.Dup(),\n",
-        "    tl.Parallel(tl.Relu(), tl.LayerNorm()),\n",
-        ")\n",
-        "layer_block.initialize_once(input_shape, input_dtype, prng_key)\n",
-        "y = layer_block(x)\n",
-        "(y0, y1) = y  # two outputs from layer\n",
-        "\n",
-        "template = ('x:\\n{}\\n\\n'\n",
-        "            'relu(x):\\n{}\\n\\n'\n",
-        "            'layer_norm(x):\\n{}')\n",
-        "print(template.format(x, y0, y1))"
+        "print('layer_block.params:\\n{}'.format(layer_block.params))"
       ]
     },
     {
-      "cell_type": "code",
-      "execution_count": 9,
+      "cell_type": "markdown",
       "metadata": {
-        "colab": {
-          "height": 170
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 2912,
-          "status": "ok",
-          "timestamp": 1569980026618,
-          "user": {
-            "displayName": "Jonni Kanerva",
-            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
-            "userId": "12646438700820120918"
-          },
-          "user_tz": 420
-        },
-        "id": "2EIewVHC9QQZ",
-        "outputId": "0194d386-4e06-4bd3-d2a1-24f668b268ac"
+        "colab_type": "text",
+        "id": "PqVNdoONcTp0"
       },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "x:\n",
-            "[[-7. -6. -5. -4. -3.]\n",
-            " [-2. -1.  0.  1.  2.]\n",
-            " [ 3.  4.  5.  6.  7.]]\n",
-            "\n",
-            "concatenate(relu(x), layer_norm(x)):\n",
-            "[[ 0.     0.     0.     0.     0.    -1.414 -0.707  0.     0.707  1.414]\n",
-            " [ 0.     0.     0.     1.     2.    -1.414 -0.707  0.     0.707  1.414]\n",
-            " [ 3.     4.     5.     6.     7.    -1.414 -0.707  0.     0.707  1.414]]\n"
-          ]
-        }
-      ],
       "source": [
-        "prng_key = backend.random.get_prng(0)\n",
-        "input_dtype = onp.float32\n",
-        "input_shape = (3, 5)\n",
-        "x = onp.arange(-7, 8).reshape(*input_shape).astype(input_dtype)\n",
+        "## Combine with Parallel(...)\n",
         "\n",
-        "layer_block = tl.Serial(\n",
-        "    tl.Dup(),\n",
-        "    tl.Parallel(tl.Relu(), tl.LayerNorm()),\n",
-        "    tl.Concatenate(),\n",
-        ")\n",
-        "layer_block.initialize_once(input_shape, input_dtype, prng_key)\n",
-        "y = layer_block(x)\n",
+        "The `Parallel` combinator arranges layers into separate computational channels, each with its own inputs/outputs and gradient flows:\n",
+        "```\n",
+        "class Parallel(base.Layer):\n",
+        "  \"\"\"Combinator that applies a list of layers in parallel to its inputs.\n",
         "\n",
-        "template = ('x:\\n{}\\n\\n'\n",
-        "            'concatenate(relu(x), layer_norm(x)):\\n{}')\n",
-        "print(template.format(x, y))"
+        "  Layers in the list apply to successive spans of inputs, where the spans are\n",
+        "  determined how many inputs each layer takes. The resulting output is the\n",
+        "  (flattened) concatenation of the resepective layer outputs.\n",
+        "\n",
+        "  For example, suppose one has three layers:\n",
+        "\n",
+        "    - F: 1 input, 1 output\n",
+        "    - G: 3 inputs, 1 output\n",
+        "    - H: 2 inputs, 2 outputs (h1, h2)\n",
+        "\n",
+        "  Then Parallel(F, G, H) will take 6 inputs and give 4 outputs:\n",
+        "\n",
+        "    - inputs: a, b, c, d, e, f\n",
+        "    - outputs: F(a), G(b, c, d), h1, h2\n",
+        "```\n",
+        "\n",
+        "Separate (parallel) computation channels make sense when each channel can do its work (computing outputs from inputs) independent of the inputs and outputs of the others.\n",
+        "\n",
+        "As a simplistic example, consider writing a converter from three-digit octal (base 8) numerals to their corresponding values. For instance, to do conversions such as\n",
+        "```\n",
+        "123 (octal) = 1 * 8^2 + 2 * 8^1 + 3 * 8^0 =  83 (decimal)\n",
+        "345 (octal) = 3 * 8^2 + 4 * 8^1 + 6 * 8^0 = 229 (decimal)\n",
+        "567 (octal) = 5 * 8^2 + 6 * 8^1 + 7 * 8^0 = 375 (decimal)\n",
+        "```\n",
+        "the digits can first be converted independently, according to their place value (multiply by 64, multiply by 8, or multiply by 1). The following code runs the 64's-place digits ([1, 3, 5]) through one layer, the 8's-place digits ([2, 4, 6]) through a different layer, and the 1's-place digits ([3, 5, 7]) through yet a different layer. These three layers are combined in a Parallel layer:\n",
+        "\n",
+        "### Example 6. Processing octal digits in parallel."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 10,
+      "execution_count": 9,
       "metadata": {
         "colab": {
-          "height": 442
+          "height": 204
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 6372,
+          "elapsed": 2224,
           "status": "ok",
-          "timestamp": 1569980033028,
+          "timestamp": 1570168990465,
           "user": {
-            "displayName": "Jonni Kanerva",
-            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
-            "userId": "12646438700820120918"
+            "displayName": "",
+            "photoUrl": "",
+            "userId": ""
           },
           "user_tz": 420
         },
-        "id": "1byhBplSby2a",
-        "outputId": "6d1cc0da-a0a4-4b45-a654-da800634edb9"
+        "id": "uQMqq3h_b2jQ",
+        "outputId": "f3a43cae-e271-493a-f74a-31e1ff971bc1"
       },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "x:\n",
-            "[[-7. -6. -5. -4. -3.]\n",
-            " [-2. -1.  0.  1.  2.]\n",
-            " [ 3.  4.  5.  6.  7.]]\n",
-            "\n",
-            "dense(dense(x)):\n",
-            "[[ 5.83  -6.634  3.101 -3.777  6.508]\n",
-            " [ 0.025 -1.503  0.041 -0.626  0.387]\n",
-            " [-5.78   3.628 -3.018  2.526 -5.734]]\n",
+            "inputs:\n",
+            "(array([1, 3, 5]), array([2, 4, 6]), array([3, 5, 7]))\n",
             "\n",
-            "params:\n",
-            "[(DeviceArray([[-0.355,  0.284, -0.737,  0.309, -0.613],\n",
-            "             [ 0.574, -0.127,  0.149,  0.763, -0.566],\n",
-            "             [-0.083, -0.454,  0.19 ,  0.018, -0.399],\n",
-            "             [-0.706,  0.454, -0.636,  0.022, -0.084],\n",
-            "             [ 0.422, -0.509, -0.415,  0.638, -0.469]], dtype=float32), DeviceArray([-2.072e-07, -1.046e-06,  2.496e-06,  2.705e-07, -1.035e-08],\n",
-            "            dtype=float32)), (DeviceArray([[ 0.459, -0.581,  0.693, -0.163,  0.77 ],\n",
-            "             [ 0.016,  0.735,  0.619,  0.585,  0.337],\n",
-            "             [ 0.617,  0.345, -0.162, -0.039, -0.073],\n",
-            "             [-0.122,  0.103,  0.477,  0.499,  0.199],\n",
-            "             [-0.01 , -0.713,  0.639,  0.056,  0.678]], dtype=float32), DeviceArray([-1.936e-07, -1.772e-06, -6.521e-07, -1.623e-06, -4.441e-07],\n",
-            "            dtype=float32))]\n",
+            "octal_place_values(inputs):\n",
+            "(array([ 64., 192., 320.]), array([16., 32., 48.]), array([3., 5., 7.]))\n",
             "\n",
-            "state:\n",
-            "[(), ()]\n"
+            "octal_place_values.n_inputs:  3\n",
+            "octal_place_values.n_outputs: 3\n",
+            "octal_place_values.sublayers: [MulConstant{in=1,out=1}, MulConstant{in=1,out=1}, MulConstant{in=1,out=1}]\n",
+            "octal_place_values.params:    ((), (), ())\n",
+            "\n"
           ]
         }
       ],
       "source": [
         "prng_key = backend.random.get_prng(0)\n",
-        "input_dtype = onp.float32\n",
-        "input_shape = (3, 5)\n",
-        "x = onp.arange(-7, 8).reshape(*input_shape).astype(input_dtype)\n",
-        "\n",
-        "layer_block = tl.Serial(\n",
-        "    tl.Dense(5),\n",
-        "    tl.Dense(5),\n",
-        ")\n",
-        "layer_block.initialize_once(input_shape, input_dtype, prng_key)\n",
-        "y = layer_block(x)\n",
+        "place_64_digits = onp.array([1, 3, 5])\n",
+        "place_8_digits = onp.array([2, 4, 6])\n",
+        "place_1_digits = onp.array([3, 5, 7])\n",
+        "inputs = (place_64_digits, place_8_digits, place_1_digits)\n",
+        "input_shapes = [[3]] * 3\n",
+        "input_dtypes = [onp.int32] * 3\n",
+        "\n",
+        "# Create three simple layers, each for converting a different digit in base 8.\n",
+        "sixty_fours = tl.MulConstant(constant=64.0)  # 8^2: 100 in base 8\n",
+        "eights = tl.MulConstant(constant=8.0)  # 8^1: 10 in base 8\n",
+        "ones = tl.MulConstant(constant=1.0)  # 8^0: 1 in base 8\n",
+        "\n",
+        "# Create a combined layer to convert digits to values (using big-endian base 8),\n",
+        "# initialize it, and apply it.\n",
+        "octal_place_values = tl.Parallel(sixty_fours, eights, ones)\n",
+        "octal_place_values.initialize_once(input_shapes, input_dtypes, prng_key)\n",
+        "outputs = octal_place_values(inputs)\n",
+        "\n",
+        "# Show inputs, outputs, and properties.\n",
+        "template = ('inputs:\\n{}\\n\\n'\n",
+        "            'octal_place_values(inputs):\\n{}\\n')\n",
+        "print(template.format(inputs, outputs))\n",
+        "show_layer_properties(octal_place_values, 'octal_place_values')"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "q_xcWide3e5f"
+      },
+      "source": [
+        "To complete the example, the three outputs (values for the different digits) are combined by successive pairwise additions:\n",
         "\n",
-        "template = ('x:\\n{}\\n\\n'\n",
-        "            'dense(dense(x)):\\n{}\\n\\n'\n",
-        "            'params:\\n{}\\n\\n'\n",
-        "            'state:\\n{}')\n",
-        "print(template.format(x, y, layer_block.params, layer_block.state))"
+        "### Example 6'. Combining outputs from upstream parallel layers."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 11,
+      "execution_count": 10,
       "metadata": {
         "colab": {
-          "height": 357
+          "height": 275
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 1047,
+          "elapsed": 2139,
           "status": "ok",
-          "timestamp": 1569980034098,
+          "timestamp": 1570168992621,
           "user": {
-            "displayName": "Jonni Kanerva",
-            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAWyI3hqIhQiLmO-Yzb2o-qkz0xCUECEG0fptCO=s64",
-            "userId": "12646438700820120918"
+            "displayName": "",
+            "photoUrl": "",
+            "userId": ""
           },
           "user_tz": 420
         },
-        "id": "aRKnfN4adLqi",
-        "outputId": "568f4974-b9bf-4e5f-b9c4-dc6413068462"
+        "id": "ZDCkrvUp3u0-",
+        "outputId": "696f21aa-5dad-4284-bfdd-ae637e2ce53f"
       },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "x:\n",
-            "[[-7. -6. -5. -4. -3.]\n",
-            " [-2. -1.  0.  1.  2.]\n",
-            " [ 3.  4.  5.  6.  7.]]\n",
+            "inputs:\n",
+            "(array([1, 3, 5]), array([2, 4, 6]), array([3, 5, 7]))\n",
             "\n",
-            "dense(dense(x)):\n",
-            "[[ 10.543 -12.855   1.589   8.106  -8.503]\n",
-            " [ -0.278  -0.161  -0.644  -0.195   0.103]\n",
-            " [-11.099  12.533  -2.877  -8.497   8.709]]\n",
+            "octal_place_values(inputs):\n",
+            "(array([ 64., 192., 320.]), array([16., 32., 48.]), array([3., 5., 7.]))\n",
             "\n",
-            "params:\n",
-            "[(DeviceArray([[-0.355,  0.284, -0.737,  0.309, -0.613],\n",
-            "             [ 0.574, -0.127,  0.149,  0.763, -0.566],\n",
-            "             [-0.083, -0.454,  0.19 ,  0.018, -0.399],\n",
-            "             [-0.706,  0.454, -0.636,  0.022, -0.084],\n",
-            "             [ 0.422, -0.509, -0.415,  0.638, -0.469]], dtype=float32), DeviceArray([-2.072e-07, -1.046e-06,  2.496e-06,  2.705e-07, -1.035e-08],\n",
-            "            dtype=float32)), ()]\n",
+            "evaluate_octal(inputs):\n",
+            "[ 83. 229. 375.]\n",
             "\n",
-            "state:\n",
-            "[(), ()]\n"
+            "evaluate_octal.n_inputs:  3\n",
+            "evaluate_octal.n_outputs: 1\n",
+            "evaluate_octal.sublayers: [Parallel{in=3,out=3,sublayers=[MulConstant{in=1,out=1}, MulConstant{in=1,out=1}, MulConstant{in=1,out=1}]}, Add{in=2,out=1}, Add{in=2,out=1}]\n",
+            "evaluate_octal.params:    [(), (), ()]\n",
+            "\n"
           ]
         }
       ],
       "source": [
-        "prng_key = backend.random.get_prng(0)\n",
-        "input_dtype = onp.float32\n",
-        "input_shape = (3, 5)\n",
-        "x = onp.arange(-7, 8).reshape(*input_shape).astype(input_dtype)\n",
-        "\n",
-        "dense = tl.Dense(5)\n",
-        "layer_block = tl.Serial(\n",
-        "    dense,\n",
-        "    dense,\n",
+        "evaluate_octal = tl.Serial(\n",
+        "    octal_place_values,\n",
+        "    tl.Add(),  # Adds the 64's-place values and the 8's-place values.\n",
+        "    tl.Add(),  # Adds the 1's-place values to the sums from the previous Add.\n",
         ")\n",
-        "layer_block.initialize_once(input_shape, input_dtype, prng_key)\n",
-        "y = layer_block(x)\n",
-        "\n",
-        "template = ('x:\\n{}\\n\\n'\n",
-        "            'dense(dense(x)):\\n{}\\n\\n'\n",
-        "            'params:\\n{}\\n\\n'\n",
-        "            'state:\\n{}')\n",
-        "print(template.format(x, y, layer_block.params, layer_block.state))"
+        "evaluate_octal.initialize_once(input_shapes, input_dtypes, prng_key)\n",
+        "y = evaluate_octal(inputs)\n",
+        "\n",
+        "template = ('inputs:\\n{}\\n\\n'\n",
+        "            'octal_place_values(inputs):\\n{}\\n\\n'\n",
+        "            'evaluate_octal(inputs):\\n{}\\n')\n",
+        "print(template.format(inputs, outputs, y))\n",
+        "show_layer_properties(evaluate_octal, 'evaluate_octal')"
       ]
     },
     {

From a8024e8f225f38824d17dc10ec8eaa4fd2c50eeb Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 4 Oct 2019 12:44:42 -0700
Subject: [PATCH 2539/2720] Update Position Lookup Transformer config to use
 recent Trax additions.

PiperOrigin-RevId: 272937269
---
 .../position_lookup_transformer_copy.gin      | 32 +++++++++++++------
 1 file changed, 22 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin b/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
index d4fde63c0..036310433 100644
--- a/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
+++ b/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
@@ -3,18 +3,18 @@ import tensor2tensor.trax.models
 import tensor2tensor.trax.optimizers
 import tensor2tensor.trax.trax
 
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 32
-batch_fn.eval_batch_size = 32
-
 # Parameters for inputs:
 # ==============================================================================
-inputs.dataset_name = 't2t_algorithmic_identity_vocab95_train20_eval30'
+inputs.data_dir = None
+inputs.dataset_name = 'position_lookup_copy'
 
-# Parameters for preprocess_fun:
+# Parameters for sequence_copy_inputs:
 # ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.concat_preprocess
+sequence_copy_inputs.vocab_size = 128
+sequence_copy_inputs.batch_size = 16
+sequence_copy_inputs.train_lengths = [20, 30, 40]
+sequence_copy_inputs.eval_lengths = [60]
+sequence_copy_inputs.reverse = False
 
 # Parameters for MultifactorSchedule:
 # ==============================================================================
@@ -27,16 +27,28 @@ MultifactorSchedule.warmup_steps = 16000
 PositionLookupTransformerLM.d_model = 256
 PositionLookupTransformerLM.d_ff = 512
 PositionLookupTransformerLM.dropout = 0.01
-PositionLookupTransformerLM.max_len = 100
+PositionLookupTransformerLM.max_len = 62
 PositionLookupTransformerLM.n_heads = 4
 PositionLookupTransformerLM.n_layers = 3
 PositionLookupTransformerLM.vocab_size = 128
 
+# Parameters for TransformerLM: (same as above, for easy comparisons)
+# ==============================================================================
+TransformerLM.d_model = 256
+TransformerLM.d_ff = 512
+TransformerLM.dropout = 0.01
+TransformerLM.max_len = 62
+TransformerLM.n_heads = 4
+TransformerLM.n_layers = 3
+TransformerLM.vocab_size = 128
+
 # Parameters for train:
 # ==============================================================================
+train.inputs = @trax.inputs.sequence_copy_inputs
 train.eval_frequency = 1000
 train.eval_steps = 10
 train.model = @trax.models.PositionLookupTransformerLM
 train.optimizer = @trax.optimizers.Adam
 train.train_steps = 100000
-train.mask_id = 0
\ No newline at end of file
+train.mask_id = 0
+train.has_weights = True

From 6c7c601b8c4429dcc81ab3ec828daddea5ff2b67 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 5 Oct 2019 12:36:19 -0700
Subject: [PATCH 2540/2720] Move Trax out of T2T into its own repo.

PiperOrigin-RevId: 273078371
---
 tensor2tensor/trax/README.md                  |   65 -
 tensor2tensor/trax/__init__.py                |   15 -
 tensor2tensor/trax/backend.py                 |  308 ----
 tensor2tensor/trax/backend_test.py            |   75 -
 tensor2tensor/trax/configs/mlp_mnist.gin      |   35 -
 .../position_lookup_transformer_copy.gin      |   54 -
 .../trax/configs/reformer_base_sweep.yaml     |    2 -
 .../trax/configs/reformer_enwik8.gin          |  115 --
 .../trax/configs/reformer_hash_sweep.yaml     |    1 -
 .../trax/configs/reformer_imagenet64.gin      |  114 --
 .../trax/configs/reformer_large_sweep.yaml    |    3 -
 .../trax/configs/resnet50_imagenet_8gb.gin    |   44 -
 .../configs/resnet50_imagenet_8gb_testing.gin |   44 -
 .../trax/configs/transformer_big_lm1b_8gb.gin |   54 -
 .../trax/configs/transformer_copy.gin         |   77 -
 .../trax/configs/transformer_imdb_8gb.gin     |   54 -
 .../trax/configs/transformer_lm1b_16gb.gin    |  127 --
 .../trax/configs/transformer_lm1b_8gb.gin     |   54 -
 .../configs/transformer_lm1b_8gb_testing.gin  |   54 -
 .../configs/transformer_lm_wmt_ende_16gb.gin  |  134 --
 .../configs/transformer_lm_wmt_ende_8gb.gin   |   64 -
 .../trax/configs/transformer_ptb_16gb.gin     |  128 --
 ...former_wmt_ende_16gb_adafactor_testing.gin |   59 -
 .../transformer_wmt_ende_8gb_adafactor.gin    |   59 -
 .../configs/transformer_wmt_ende_8gb_adam.gin |   56 -
 .../configs/transformer_wmt_ende_8gb_sm3.gin  |   53 -
 .../trax/configs/wide_resnet_cifar10_8gb.gin  |   81 -
 tensor2tensor/trax/history.py                 |   78 -
 tensor2tensor/trax/inputs.py                  |  648 --------
 tensor2tensor/trax/inputs_test.py             |   72 -
 tensor2tensor/trax/jaxboard.py                |  350 -----
 tensor2tensor/trax/layers/README.md           |   60 -
 tensor2tensor/trax/layers/__init__.py         |   62 -
 tensor2tensor/trax/layers/attention.py        | 1355 -----------------
 tensor2tensor/trax/layers/attention_test.py   |  106 --
 tensor2tensor/trax/layers/base.py             |  664 --------
 tensor2tensor/trax/layers/base_test.py        |   90 --
 tensor2tensor/trax/layers/combinators.py      |  539 -------
 tensor2tensor/trax/layers/combinators_test.py |  116 --
 tensor2tensor/trax/layers/convolution.py      |  126 --
 tensor2tensor/trax/layers/convolution_test.py |   53 -
 tensor2tensor/trax/layers/core.py             |  269 ----
 tensor2tensor/trax/layers/core_test.py        |  125 --
 tensor2tensor/trax/layers/initializers.py     |  173 ---
 .../trax/layers/initializers_test.py          |   83 -
 tensor2tensor/trax/layers/intro.ipynb         |  834 ----------
 tensor2tensor/trax/layers/metrics.py          |  124 --
 tensor2tensor/trax/layers/metrics_test.py     |   90 --
 tensor2tensor/trax/layers/normalization.py    |  125 --
 .../trax/layers/normalization_test.py         |   69 -
 tensor2tensor/trax/layers/pooling.py          |   44 -
 tensor2tensor/trax/layers/pooling_test.py     |   36 -
 tensor2tensor/trax/layers/reversible.py       |  126 --
 tensor2tensor/trax/layers/reversible_test.py  |   37 -
 tensor2tensor/trax/layers/rnn.py              |  128 --
 tensor2tensor/trax/layers/rnn_test.py         |   45 -
 tensor2tensor/trax/learning_rate.py           |  264 ----
 tensor2tensor/trax/learning_rate_test.py      |  120 --
 tensor2tensor/trax/models/__init__.py         |   51 -
 tensor2tensor/trax/models/atari_cnn.py        |   79 -
 tensor2tensor/trax/models/atari_cnn_test.py   |   65 -
 tensor2tensor/trax/models/mlp.py              |   38 -
 tensor2tensor/trax/models/mlp_test.py         |   39 -
 tensor2tensor/trax/models/neural_gpu.py       |   82 -
 tensor2tensor/trax/models/neural_gpu_test.py  |   39 -
 .../trax/models/research/__init__.py          |   15 -
 .../research/position_lookup_transformer.py   |  341 -----
 .../trax/models/research/reformer.py          |  522 -------
 .../trax/models/research/reformer_test.py     |  104 --
 tensor2tensor/trax/models/resnet.py           |  169 --
 tensor2tensor/trax/models/resnet_test.py      |   45 -
 tensor2tensor/trax/models/transformer.py      |  395 -----
 tensor2tensor/trax/models/transformer_test.py |  104 --
 .../trax/notebooks/trax_demo_iclr2019.ipynb   |  854 -----------
 tensor2tensor/trax/optimizers/__init__.py     |   37 -
 tensor2tensor/trax/optimizers/base.py         |  465 ------
 tensor2tensor/trax/rl/__init__.py             |   52 -
 tensor2tensor/trax/rl/base_trainer.py         |  141 --
 tensor2tensor/trax/rl/base_trainer_test.py    |  144 --
 tensor2tensor/trax/rl/configs/acrobot.gin     |   30 -
 .../trax/rl/configs/acrobot_transformer.gin   |   34 -
 tensor2tensor/trax/rl/configs/atari.gin       |   30 -
 .../trax/rl/configs/atari_regression_test.gin |   30 -
 ...nline_tune_transformer_imagenet64_16gb.gin |  105 --
 .../env_online_tune_transformer_lm1b_16gb.gin |   99 --
 ...line_tune_transformer_lm_wmt_ende_16gb.gin |   96 --
 .../env_online_tune_transformer_ptb_16gb.gin  |   94 --
 ...nv_online_tune_wide_resnet_cifar10_8gb.gin |   52 -
 .../trax/rl/configs/ppo_online_tune.gin       |   37 -
 .../ppo_online_tune_wide_resnet_cifar10.gin   |   79 -
 .../trax/rl/configs/simple_online_tune.gin    |   95 --
 .../configs/simple_online_tune_serialized.gin |  100 --
 tensor2tensor/trax/rl/envs/__init__.py        |   34 -
 .../rl/envs/async_trajectory_collector.py     |  197 ---
 .../rl/envs/async_trajectory_collector_lib.py |  195 ---
 .../async_trajectory_collector_lib_test.py    |   64 -
 tensor2tensor/trax/rl/envs/fake_env.py        |   68 -
 tensor2tensor/trax/rl/envs/fake_env_test.py   |   65 -
 tensor2tensor/trax/rl/envs/online_tune.py     |   57 -
 tensor2tensor/trax/rl/envs/online_tune_env.py |  233 ---
 .../trax/rl/envs/online_tune_env_test.py      |  152 --
 .../trax/rl/envs/online_tune_test.py          |  111 --
 tensor2tensor/trax/rl/online_tune.py          |  115 --
 tensor2tensor/trax/rl/online_tune_test.py     |  176 ---
 tensor2tensor/trax/rl/ppo.py                  |  971 ------------
 tensor2tensor/trax/rl/ppo_test.py             |  643 --------
 tensor2tensor/trax/rl/ppo_trainer.py          |  843 ----------
 tensor2tensor/trax/rl/ppo_trainer_test.py     |  306 ----
 tensor2tensor/trax/rl/serialization_utils.py  |  184 ---
 .../trax/rl/serialization_utils_test.py       |  169 --
 tensor2tensor/trax/rl/simple.py               |  236 ---
 tensor2tensor/trax/rl/simple_test.py          |  304 ----
 tensor2tensor/trax/rl/simple_trainer.py       |  341 -----
 tensor2tensor/trax/rl/simple_trainer_test.py  |   96 --
 .../trax/rl/simulated_env_problem.py          |  499 ------
 .../trax/rl/simulated_env_problem_test.py     |  292 ----
 tensor2tensor/trax/rl/space_serializer.py     |  216 ---
 .../trax/rl/space_serializer_test.py          |  162 --
 tensor2tensor/trax/rl/trainers.py             |   37 -
 tensor2tensor/trax/rl_trainer.py              |  209 ---
 tensor2tensor/trax/trainer.py                 |  137 --
 tensor2tensor/trax/trax.py                    |  958 ------------
 tensor2tensor/trax/trax_test.py               |  268 ----
 tensor2tensor/trax/utils.py                   |   43 -
 124 files changed, 21983 deletions(-)
 delete mode 100644 tensor2tensor/trax/README.md
 delete mode 100644 tensor2tensor/trax/__init__.py
 delete mode 100644 tensor2tensor/trax/backend.py
 delete mode 100644 tensor2tensor/trax/backend_test.py
 delete mode 100644 tensor2tensor/trax/configs/mlp_mnist.gin
 delete mode 100644 tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
 delete mode 100644 tensor2tensor/trax/configs/reformer_base_sweep.yaml
 delete mode 100644 tensor2tensor/trax/configs/reformer_enwik8.gin
 delete mode 100644 tensor2tensor/trax/configs/reformer_hash_sweep.yaml
 delete mode 100644 tensor2tensor/trax/configs/reformer_imagenet64.gin
 delete mode 100644 tensor2tensor/trax/configs/reformer_large_sweep.yaml
 delete mode 100644 tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
 delete mode 100644 tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_copy.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_imdb_8gb.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_lm1b_16gb.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_lm_wmt_ende_8gb.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_ptb_16gb.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
 delete mode 100644 tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin
 delete mode 100644 tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
 delete mode 100644 tensor2tensor/trax/history.py
 delete mode 100644 tensor2tensor/trax/inputs.py
 delete mode 100644 tensor2tensor/trax/inputs_test.py
 delete mode 100644 tensor2tensor/trax/jaxboard.py
 delete mode 100644 tensor2tensor/trax/layers/README.md
 delete mode 100644 tensor2tensor/trax/layers/__init__.py
 delete mode 100644 tensor2tensor/trax/layers/attention.py
 delete mode 100644 tensor2tensor/trax/layers/attention_test.py
 delete mode 100644 tensor2tensor/trax/layers/base.py
 delete mode 100644 tensor2tensor/trax/layers/base_test.py
 delete mode 100644 tensor2tensor/trax/layers/combinators.py
 delete mode 100644 tensor2tensor/trax/layers/combinators_test.py
 delete mode 100644 tensor2tensor/trax/layers/convolution.py
 delete mode 100644 tensor2tensor/trax/layers/convolution_test.py
 delete mode 100644 tensor2tensor/trax/layers/core.py
 delete mode 100644 tensor2tensor/trax/layers/core_test.py
 delete mode 100644 tensor2tensor/trax/layers/initializers.py
 delete mode 100644 tensor2tensor/trax/layers/initializers_test.py
 delete mode 100644 tensor2tensor/trax/layers/intro.ipynb
 delete mode 100644 tensor2tensor/trax/layers/metrics.py
 delete mode 100644 tensor2tensor/trax/layers/metrics_test.py
 delete mode 100644 tensor2tensor/trax/layers/normalization.py
 delete mode 100644 tensor2tensor/trax/layers/normalization_test.py
 delete mode 100644 tensor2tensor/trax/layers/pooling.py
 delete mode 100644 tensor2tensor/trax/layers/pooling_test.py
 delete mode 100644 tensor2tensor/trax/layers/reversible.py
 delete mode 100644 tensor2tensor/trax/layers/reversible_test.py
 delete mode 100644 tensor2tensor/trax/layers/rnn.py
 delete mode 100644 tensor2tensor/trax/layers/rnn_test.py
 delete mode 100644 tensor2tensor/trax/learning_rate.py
 delete mode 100644 tensor2tensor/trax/learning_rate_test.py
 delete mode 100644 tensor2tensor/trax/models/__init__.py
 delete mode 100644 tensor2tensor/trax/models/atari_cnn.py
 delete mode 100644 tensor2tensor/trax/models/atari_cnn_test.py
 delete mode 100644 tensor2tensor/trax/models/mlp.py
 delete mode 100644 tensor2tensor/trax/models/mlp_test.py
 delete mode 100644 tensor2tensor/trax/models/neural_gpu.py
 delete mode 100644 tensor2tensor/trax/models/neural_gpu_test.py
 delete mode 100644 tensor2tensor/trax/models/research/__init__.py
 delete mode 100644 tensor2tensor/trax/models/research/position_lookup_transformer.py
 delete mode 100644 tensor2tensor/trax/models/research/reformer.py
 delete mode 100644 tensor2tensor/trax/models/research/reformer_test.py
 delete mode 100644 tensor2tensor/trax/models/resnet.py
 delete mode 100644 tensor2tensor/trax/models/resnet_test.py
 delete mode 100644 tensor2tensor/trax/models/transformer.py
 delete mode 100644 tensor2tensor/trax/models/transformer_test.py
 delete mode 100644 tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb
 delete mode 100644 tensor2tensor/trax/optimizers/__init__.py
 delete mode 100644 tensor2tensor/trax/optimizers/base.py
 delete mode 100644 tensor2tensor/trax/rl/__init__.py
 delete mode 100644 tensor2tensor/trax/rl/base_trainer.py
 delete mode 100644 tensor2tensor/trax/rl/base_trainer_test.py
 delete mode 100644 tensor2tensor/trax/rl/configs/acrobot.gin
 delete mode 100644 tensor2tensor/trax/rl/configs/acrobot_transformer.gin
 delete mode 100644 tensor2tensor/trax/rl/configs/atari.gin
 delete mode 100644 tensor2tensor/trax/rl/configs/atari_regression_test.gin
 delete mode 100644 tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
 delete mode 100644 tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
 delete mode 100644 tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
 delete mode 100644 tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin
 delete mode 100644 tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10_8gb.gin
 delete mode 100644 tensor2tensor/trax/rl/configs/ppo_online_tune.gin
 delete mode 100644 tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin
 delete mode 100644 tensor2tensor/trax/rl/configs/simple_online_tune.gin
 delete mode 100644 tensor2tensor/trax/rl/configs/simple_online_tune_serialized.gin
 delete mode 100644 tensor2tensor/trax/rl/envs/__init__.py
 delete mode 100644 tensor2tensor/trax/rl/envs/async_trajectory_collector.py
 delete mode 100644 tensor2tensor/trax/rl/envs/async_trajectory_collector_lib.py
 delete mode 100644 tensor2tensor/trax/rl/envs/async_trajectory_collector_lib_test.py
 delete mode 100644 tensor2tensor/trax/rl/envs/fake_env.py
 delete mode 100644 tensor2tensor/trax/rl/envs/fake_env_test.py
 delete mode 100644 tensor2tensor/trax/rl/envs/online_tune.py
 delete mode 100644 tensor2tensor/trax/rl/envs/online_tune_env.py
 delete mode 100644 tensor2tensor/trax/rl/envs/online_tune_env_test.py
 delete mode 100644 tensor2tensor/trax/rl/envs/online_tune_test.py
 delete mode 100644 tensor2tensor/trax/rl/online_tune.py
 delete mode 100644 tensor2tensor/trax/rl/online_tune_test.py
 delete mode 100644 tensor2tensor/trax/rl/ppo.py
 delete mode 100644 tensor2tensor/trax/rl/ppo_test.py
 delete mode 100644 tensor2tensor/trax/rl/ppo_trainer.py
 delete mode 100644 tensor2tensor/trax/rl/ppo_trainer_test.py
 delete mode 100644 tensor2tensor/trax/rl/serialization_utils.py
 delete mode 100644 tensor2tensor/trax/rl/serialization_utils_test.py
 delete mode 100644 tensor2tensor/trax/rl/simple.py
 delete mode 100644 tensor2tensor/trax/rl/simple_test.py
 delete mode 100644 tensor2tensor/trax/rl/simple_trainer.py
 delete mode 100644 tensor2tensor/trax/rl/simple_trainer_test.py
 delete mode 100644 tensor2tensor/trax/rl/simulated_env_problem.py
 delete mode 100644 tensor2tensor/trax/rl/simulated_env_problem_test.py
 delete mode 100644 tensor2tensor/trax/rl/space_serializer.py
 delete mode 100644 tensor2tensor/trax/rl/space_serializer_test.py
 delete mode 100644 tensor2tensor/trax/rl/trainers.py
 delete mode 100644 tensor2tensor/trax/rl_trainer.py
 delete mode 100644 tensor2tensor/trax/trainer.py
 delete mode 100644 tensor2tensor/trax/trax.py
 delete mode 100644 tensor2tensor/trax/trax_test.py
 delete mode 100644 tensor2tensor/trax/utils.py

diff --git a/tensor2tensor/trax/README.md b/tensor2tensor/trax/README.md
deleted file mode 100644
index 459b9b916..000000000
--- a/tensor2tensor/trax/README.md
+++ /dev/null
@@ -1,65 +0,0 @@
-## `trax`: Train Neural Nets with JAX
-
-![train tracks](https://images.pexels.com/photos/461772/pexels-photo-461772.jpeg?dl&fit=crop&crop=entropy&w=640&h=426)
-
-### `trax`: T2T Radically Simpler with JAX
-
-*Why?* Because T2T has gotten too complex. We are simplifying the main code too,
-but we wanted to try a more radical step. So you can write code as in pure
-NumPy and debug directly. So you can easily pinpoint each line where things
-happen and understand each function. But we also want it to run fast on
-accelerators, and that's possible with [JAX](https://github.com/google/jax).
-
-*Status:* preview; things work: models train, checkpoints are saved, TensorBoard
-has summaries, you can decode. But we are changing a lot every day for now.
-Please let us know what we should add, delete, keep, change. We plan to move
-the best parts into core JAX.
-
-*Entrypoints:*
-
-* Script: `trainer.py`
-* Main library entrypoint: `trax.train`
-
-### Examples
-
-#### Example Colab
-
-See our example constructing language models from scratch in a GPU-backed colab notebook at
-[Trax Demo](https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb)
-
-#### MLP on MNIST
-
-
-```
-python -m tensor2tensor.trax.trainer \
-  --dataset=mnist \
-  --model=MLP \
-  --config="train.train_steps=1000"
-```
-
-#### Resnet50 on Imagenet
-
-
-```
-python -m tensor2tensor.trax.trainer \
-  --config_file=$PWD/trax/configs/resnet50_imagenet_8gb.gin
-```
-
-#### TransformerDecoder on LM1B
-
-
-```
-python -m tensor2tensor.trax.trainer \
-  --config_file=$PWD/trax/configs/transformer_lm1b_8gb.gin
-```
-
-### How `trax` differs from T2T
-
-* Configuration is done with [`gin`](https://github.com/google/gin-config).
-  `trainer.py` takes `--config_file` as well as `--config` for file overrides.
-* Models are defined with [`stax`](https://github.com/google/jax/blob/master/jax/experimental/stax.py) in
-  `models/`. They are made gin-configurable in `models/__init__.py`.
-* Datasets are simple iterators over batches. Datasets from
-  [`tensorflow/datasets`](https://github.com/tensorflow/datasets)
-  and [`tensor2tensor`](https://github.com/tensorflow/tensor2tensor)
-  are built-in and can be addressed by name.
diff --git a/tensor2tensor/trax/__init__.py b/tensor2tensor/trax/__init__.py
deleted file mode 100644
index 4872e5d5d..000000000
--- a/tensor2tensor/trax/__init__.py
+++ /dev/null
@@ -1,15 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
diff --git a/tensor2tensor/trax/backend.py b/tensor2tensor/trax/backend.py
deleted file mode 100644
index 69997ec24..000000000
--- a/tensor2tensor/trax/backend.py
+++ /dev/null
@@ -1,308 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Trax backend: all the primitive functions needed."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import contextlib
-import gin
-
-import jax
-from jax import lax
-from jax import random as jax_random
-import jax.numpy as jnp
-import jax.scipy.special as jax_special
-import numpy as onp
-import tensorflow_datasets as tfds
-
-
-
-def jax_conv(inp, fltr, window_strides, padding, dimension_numbers,
-             filter_dilation=None):
-  """A wrapper around `lax.conv_general_dilated`.
-
-  It requires `dimension_numbers` and disallows `inp_dilation`.
-
-  Args:
-    inp: an (N+2)-D array. The input of the convolution.
-    fltr: an (N+2)-D array. The filter (i.e. kernel) of the convolution.
-    window_strides: the strides for moving the convolution window.
-    padding: a string, either "VALID" or "SAME". The padding algorithm.
-    dimension_numbers: a tuple of three strings encoding the data format of
-      input, filter and output. "I" means input; "O" means output; "C" means
-      channel; other characters such as "W", "H" and "D" means spatial
-      dimensions.
-    filter_dilation: the dilation rates for the filter. Dilating the filter
-      means adding "holes" to the filter.
-
-  Returns:
-    An (N+2)-D array. The convolution result.
-  """
-  return lax.conv_general_dilated(inp, fltr, window_strides, padding,
-                                  lhs_dilation=None,
-                                  rhs_dilation=filter_dilation,
-                                  dimension_numbers=dimension_numbers)
-
-
-def _pooling_general(inputs, reducer, init_val, rescaler=None,
-                     pool_size=(2, 2), strides=None, padding="VALID"):
-  """Helper: general pooling computation used in pooling layers later."""
-  spatial_strides = strides or (1,) * len(pool_size)
-  rescale = rescaler(pool_size, spatial_strides, padding) if rescaler else None
-  dims = (1,) + pool_size + (1,)  # NHWC
-  strides = (1,) + spatial_strides + (1,)
-  out = lax.reduce_window(inputs, init_val, reducer, dims, strides, padding)
-  return rescale(out, inputs) if rescale else out
-
-
-def jax_max_pool(x, pool_size, strides, padding):
-  return _pooling_general(x, lax.max, -jnp.inf, pool_size=pool_size,
-                          strides=strides, padding=padding)
-
-
-def jax_sum_pool(x, pool_size, strides, padding):
-  return _pooling_general(x, lax.add, 0., pool_size=pool_size,
-                          strides=strides, padding=padding)
-
-
-def _normalize_by_window_size(dims, spatial_strides, padding):  # pylint: disable=invalid-name
-  def rescale(outputs, inputs):
-    one = jnp.ones(inputs.shape[1:-1], dtype=inputs.dtype)
-    window_sizes = lax.reduce_window(
-        one, 0., lax.add, dims, spatial_strides, padding)
-    return outputs / window_sizes[..., jnp.newaxis]
-  return rescale
-
-
-def jax_avg_pool(x, pool_size, strides, padding):
-  return _pooling_general(x, lax.add, 0., _normalize_by_window_size,
-                          pool_size, strides=strides, padding=padding)
-
-
-def nested_map(x, f):
-  """Map the function f to the nested structure x (dicts, tuples, lists)."""
-  if isinstance(x, list):
-    return [nested_map(y, f) for y in x]
-  if isinstance(x, tuple):
-    return tuple([nested_map(y, f) for y in x])
-  if isinstance(x, dict):
-    return {k: nested_map(v, f) for (k, v) in x.items()}
-  return f(x)
-
-
-class ShapeType(object):
-  """Store shape and type."""
-
-  def __init__(self, shape, dtype):
-    self.shape = shape
-    self.dtype = dtype
-
-  def __repr__(self):
-    return "[shape:" + str(self.shape) + ", dtype:" + str(self.dtype) + "]"
-
-
-def jax_eval_on_shapes(f):
-  """Returns a function that evaluates `f` given input shapes and dtypes.
-
-  It transforms function `f` to a function that performs the same computation as
-  `f` but only on shapes and dtypes (a.k.a. shape inference).
-
-  Args:
-    f: the function to be transformed.
-
-  Returns:
-    A function whose input arguments can be either the same as `f`'s or only
-    their shapes/dtypes represented by `ShapeType`, and whose return values are
-    `ShapeType`s with the same nested structure as `f`'s return values.
-  """
-  def shape_fun(*args, **kwargs):
-    jax_shapes = jax.eval_shape(f, *args, **kwargs)
-    return nested_map(jax_shapes, lambda x: ShapeType(x.shape, x.dtype))
-  return shape_fun
-
-
-# The default value of dtype is different from jax_random.randint
-def jax_randint(key, shape, minval, maxval, dtype=onp.int32):
-  """Sample uniform random values in [minval, maxval) with given shape/dtype.
-
-  Args:
-    key: a PRNGKey used as the random key.
-    shape: a tuple of nonnegative integers representing the shape.
-    minval: int or array of ints broadcast-compatible with ``shape``, a minimum
-      (inclusive) value for the range.
-    maxval: int or array of ints broadcast-compatible with  ``shape``, a maximum
-      (exclusive) value for the range.
-    dtype: optional, an int dtype for the returned values (default int32).
-
-  Returns:
-    A random array with the specified shape and dtype.
-  """
-  return jax_random.randint(key, shape, minval=minval, maxval=maxval,
-                            dtype=dtype)
-
-
-_JAX_BACKEND = {
-    "name": "jax",
-    "np": jnp,
-    "logsumexp": jax_special.logsumexp,
-    "expit": jax_special.expit,
-    "erf": jax_special.erf,
-    "conv": jax_conv,
-    "avg_pool": jax_avg_pool,
-    "max_pool": jax_max_pool,
-    "sum_pool": jax_sum_pool,
-    "jit": jax.jit,
-    "grad": jax.grad,
-    "pmap": jax.pmap,
-    "eval_on_shapes": jax_eval_on_shapes,
-    "random_uniform": jax_random.uniform,
-    "random_randint": jax_randint,
-    "random_normal": jax_random.normal,
-    "random_bernoulli": jax_random.bernoulli,
-    "random_get_prng": jax.jit(jax_random.PRNGKey),
-    "random_split": jax_random.split,
-    "dataset_as_numpy": tfds.as_numpy,
-}
-
-
-_NUMPY_BACKEND = {
-    "name": "numpy",
-    "np": onp,
-    "jit": (lambda f: f),
-    "random_get_prng": lambda seed: None,
-    "random_split": lambda prng, num=2: (None,) * num,
-    "expit": (lambda x: 1. / (1. + onp.exp(-x))),
-}
-
-
-def get_name():
-  return backend()["name"]
-
-
-def logsumexp(*args, **kwargs):
-  return backend()["logsumexp"](*args, **kwargs)
-
-
-def expit(*args, **kwargs):
-  return backend()["expit"](*args, **kwargs)
-
-
-def erf(*args, **kwargs):
-  return backend()["erf"](*args, **kwargs)
-
-
-def conv(*args, **kwargs):
-  return backend()["conv"](*args, **kwargs)
-
-
-def avg_pool(*args, **kwargs):
-  return backend()["avg_pool"](*args, **kwargs)
-
-
-def max_pool(*args, **kwargs):
-  return backend()["max_pool"](*args, **kwargs)
-
-
-def sum_pool(*args, **kwargs):
-  return backend()["sum_pool"](*args, **kwargs)
-
-
-def jit(*args, **kwargs):
-  return backend()["jit"](*args, **kwargs)
-
-
-def grad(*args, **kwargs):
-  return backend()["grad"](*args, **kwargs)
-
-
-def pmap(*args, **kwargs):
-  return backend()["pmap"](*args, **kwargs)
-
-
-def eval_on_shapes(*args, **kwargs):
-  return backend()["eval_on_shapes"](*args, **kwargs)
-
-
-def dataset_as_numpy(*args, **kwargs):
-  return backend()["dataset_as_numpy"](*args, **kwargs)
-
-
-# For numpy and random modules, we need to call "backend()" lazily, only when
-# the function is called -- so that it can be set by gin configs.
-# (Otherwise, backend() is called on import before gin-config is parsed.)
-# To do that, we make objects to encapsulated these modules.
-
-
-class RandomBackend(object):
-  """Backend providing random functions."""
-
-  def get_prng(self, seed):
-    return backend()["random_get_prng"](seed)
-
-  def split(self, prng, num=2):
-    return backend()["random_split"](prng, num)
-
-  def uniform(self, *args, **kwargs):
-    return backend()["random_uniform"](*args, **kwargs)
-
-  def randint(self, *args, **kwargs):
-    return backend()["random_randint"](*args, **kwargs)
-
-  def normal(self, *args, **kwargs):
-    return backend()["random_normal"](*args, **kwargs)
-
-  def bernoulli(self, *args, **kwargs):
-    return backend()["random_bernoulli"](*args, **kwargs)
-
-
-random = RandomBackend()
-
-
-# A class that just forwards attribute accesses to backend's numpy object.
-class NumpyBackend(object):
-
-  def __getattr__(self, attr):
-    return getattr(backend()["np"], attr)
-
-
-numpy = NumpyBackend()
-
-
-
-
-override_backend_name = None
-
-
-@gin.configurable()
-def backend(name="jax"):
-  name = name if not override_backend_name else override_backend_name
-  if name == "numpy":
-    return _NUMPY_BACKEND
-  return _JAX_BACKEND
-
-
-@contextlib.contextmanager
-def use_backend(name):
-  global override_backend_name
-  prev_name = override_backend_name
-  override_backend_name = name
-  # Run the decorated function in try-finally in case it throws, e.g. for tests.
-  try:
-    yield
-  finally:
-    override_backend_name = prev_name
diff --git a/tensor2tensor/trax/backend_test.py b/tensor2tensor/trax/backend_test.py
deleted file mode 100644
index a0939e540..000000000
--- a/tensor2tensor/trax/backend_test.py
+++ /dev/null
@@ -1,75 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.backend."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gin
-import jax.numpy as jnp
-import numpy as onp
-from tensor2tensor.trax import backend as backend_lib
-from tensorflow import test
-
-
-class BackendTest(test.TestCase):
-
-  def setUp(self):
-    gin.clear_config()
-
-  def override_gin(self, bindings):
-    gin.parse_config_files_and_bindings(None, bindings)
-
-  def test_backend_imports_correctly(self):
-    backend = backend_lib.backend()
-    self.assertEqual(jnp, backend["np"])
-    self.assertNotEqual(onp, backend["np"])
-
-    self.override_gin("backend.name = 'numpy'")
-
-    backend = backend_lib.backend()
-    self.assertNotEqual(jnp, backend["np"])
-    self.assertEqual(onp, backend["np"])
-
-  def test_numpy_backend_delegation(self):
-    # Assert that we are getting JAX's numpy backend.
-    backend = backend_lib.backend()
-    numpy = backend_lib.numpy
-    self.assertEqual(jnp, backend["np"])
-
-    # Assert that `numpy` calls the appropriate gin configured functions and
-    # properties.
-    self.assertTrue(numpy.isinf(numpy.inf))
-    self.assertEqual(jnp.isinf, numpy.isinf)
-    self.assertEqual(jnp.inf, numpy.inf)
-
-    # Assert that we will now get the pure numpy backend.
-
-    self.override_gin("backend.name = 'numpy'")
-
-    backend = backend_lib.backend()
-    numpy = backend_lib.numpy
-    self.assertEqual(onp, backend["np"])
-
-    # Assert that `numpy` calls the appropriate gin configured functions and
-    # properties.
-    self.assertTrue(numpy.isinf(numpy.inf))
-    self.assertEqual(onp.isinf, numpy.isinf)
-    self.assertEqual(onp.inf, numpy.inf)
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/configs/mlp_mnist.gin b/tensor2tensor/trax/configs/mlp_mnist.gin
deleted file mode 100644
index 75cbf7ba9..000000000
--- a/tensor2tensor/trax/configs/mlp_mnist.gin
+++ /dev/null
@@ -1,35 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.learning_rate
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 256
-batch_fn.eval_batch_size = 256
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 'mnist'
-
-# Parameters for MLP:
-# ==============================================================================
-MLP.d_hidden = 512
-MLP.n_hidden_layers = 2
-MLP.n_output_classes = 10
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 0.1
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 400
-
-# Parameters for train:
-# ==============================================================================
-train.optimizer = @trax.optimizers.Adafactor
-train.eval_frequency = 200
-train.eval_steps = 10
-train.model = @trax.models.MLP
-train.train_steps = 2000
diff --git a/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin b/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
deleted file mode 100644
index 036310433..000000000
--- a/tensor2tensor/trax/configs/position_lookup_transformer_copy.gin
+++ /dev/null
@@ -1,54 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 'position_lookup_copy'
-
-# Parameters for sequence_copy_inputs:
-# ==============================================================================
-sequence_copy_inputs.vocab_size = 128
-sequence_copy_inputs.batch_size = 16
-sequence_copy_inputs.train_lengths = [20, 30, 40]
-sequence_copy_inputs.eval_lengths = [60]
-sequence_copy_inputs.reverse = False
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 0.05
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 16000
-
-# Parameters for PositionLookupTransformerLM:
-# ==============================================================================
-PositionLookupTransformerLM.d_model = 256
-PositionLookupTransformerLM.d_ff = 512
-PositionLookupTransformerLM.dropout = 0.01
-PositionLookupTransformerLM.max_len = 62
-PositionLookupTransformerLM.n_heads = 4
-PositionLookupTransformerLM.n_layers = 3
-PositionLookupTransformerLM.vocab_size = 128
-
-# Parameters for TransformerLM: (same as above, for easy comparisons)
-# ==============================================================================
-TransformerLM.d_model = 256
-TransformerLM.d_ff = 512
-TransformerLM.dropout = 0.01
-TransformerLM.max_len = 62
-TransformerLM.n_heads = 4
-TransformerLM.n_layers = 3
-TransformerLM.vocab_size = 128
-
-# Parameters for train:
-# ==============================================================================
-train.inputs = @trax.inputs.sequence_copy_inputs
-train.eval_frequency = 1000
-train.eval_steps = 10
-train.model = @trax.models.PositionLookupTransformerLM
-train.optimizer = @trax.optimizers.Adam
-train.train_steps = 100000
-train.mask_id = 0
-train.has_weights = True
diff --git a/tensor2tensor/trax/configs/reformer_base_sweep.yaml b/tensor2tensor/trax/configs/reformer_base_sweep.yaml
deleted file mode 100644
index 490518f1e..000000000
--- a/tensor2tensor/trax/configs/reformer_base_sweep.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-share_qk: [True, False]
-attn_kv: [64, 128]
diff --git a/tensor2tensor/trax/configs/reformer_enwik8.gin b/tensor2tensor/trax/configs/reformer_enwik8.gin
deleted file mode 100644
index 17b3e9883..000000000
--- a/tensor2tensor/trax/configs/reformer_enwik8.gin
+++ /dev/null
@@ -1,115 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters that will vary between experiments:
-# ==============================================================================
-train.model = @trax.models.ReformerLM
-attn_type = @TimeBinCausalAttention
-share_qk = True  # Required when using LSHCausalAttention
-attn_kv = 64
-n_layers = 3
-dropout = 0.1
-
-# MemoryEfficientCausalAttention: full attention
-# (no hparams to vary between experiments)
-
-# TimeBinCausalAttention: attend to nearby items
-TimeBinCausalAttention.n_bins = 512
-
-# LSHCausalAttention: locality-sensitive hashing (LSH) attention
-LSHCausalAttention.n_bins = 512
-LSHCausalAttention.n_buckets = 1024  # Always 2 * n_bins
-LSHCausalAttention.n_hashes = 2
-LSHCausalAttention.drop_for_hash_rate = 0.0
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 1
-batch_fn.eval_batch_size = 8
-batch_fn.max_eval_length = 65536
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_enwik8_l65k'
-inputs.input_name = 'targets'
-inputs.n_chunks = 16
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 2.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 500
-train.eval_steps = 8
-train.inputs = @trax.inputs.inputs
-# train.model: see top
-train.optimizer = @trax.optimizers.Adafactor
-train.train_steps = 60000
-train.trainer_class = @MemoryEfficientTrainer
-train.save_steps = \
-    [1000, 5000, 10000, 15000, 20000, 25000, 30000, 35000, 40000, 45000, 50000,
-     55000, 60000]
-
-# Parameters for MemoryEfficientCausalAttention:
-# ==============================================================================
-MemoryEfficientCausalAttention.dropout = 0.0
-MemoryEfficientCausalAttention.loop_stride = 256
-MemoryEfficientCausalAttention.share_qk = %share_qk
-
-# Parameters for TimeBinCausalAttention:
-# ==============================================================================
-TimeBinCausalAttention.dropout = 0.0
-# TimeBinCausalAttention.n_bins: see top
-TimeBinCausalAttention.share_qk = %share_qk
-
-# Parameters for LSHCausalAttention:
-# ==============================================================================
-LSHCausalAttention.allow_duplicate_attention = False
-LSHCausalAttention.attend_across_buckets = False
-LSHCausalAttention.rehash_each_round = True
-# LSHCausalAttention.n_bins: see top
-# LSHCausalAttention.n_buckets: see top
-# LSHCausalAttention.n_hashes: see top
-LSHCausalAttention.one_rng = False
-LSHCausalAttention.hard_k = 0
-LSHCausalAttention.dropout = 0.0
-# LSHCausalAttention.drop_for_hash_rate: see top
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = %attn_type
-TransformerLM.d_attention_key = %attn_kv
-TransformerLM.d_attention_value = %attn_kv
-TransformerLM.d_model = 1024
-TransformerLM.d_ff = 4096
-TransformerLM.dropout = %dropout
-TransformerLM.max_len = 65536
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 8
-TransformerLM.n_layers = %n_layers
-TransformerLM.n_chunks = 16
-TransformerLM.share_qk = %share_qk
-TransformerLM.vocab_size = 258  # Includes pad token and unused EOS token
-
-# Parameters for ReformerLM:
-# ==============================================================================
-ReformerLM.attention_type = %attn_type
-ReformerLM.d_attention_key = %attn_kv
-ReformerLM.d_attention_value = %attn_kv
-ReformerLM.d_model = 1024
-ReformerLM.d_ff = 4096
-ReformerLM.dropout = %dropout
-ReformerLM.max_len = 65536
-ReformerLM.mode = 'train'
-ReformerLM.n_heads = 8
-ReformerLM.n_layers = %n_layers
-ReformerLM.vocab_size = 258  # Includes pad token and unused EOS token
-ReformerLM.n_chunks = 16
-ReformerLM.n_attention_chunks = 1
-ReformerLM.share_qk = %share_qk
diff --git a/tensor2tensor/trax/configs/reformer_hash_sweep.yaml b/tensor2tensor/trax/configs/reformer_hash_sweep.yaml
deleted file mode 100644
index 94216b406..000000000
--- a/tensor2tensor/trax/configs/reformer_hash_sweep.yaml
+++ /dev/null
@@ -1 +0,0 @@
-MergedMultiHashedCausalAttentionV2.n_hashes: [2, 4, 8, 16]
diff --git a/tensor2tensor/trax/configs/reformer_imagenet64.gin b/tensor2tensor/trax/configs/reformer_imagenet64.gin
deleted file mode 100644
index 9d9ae11c2..000000000
--- a/tensor2tensor/trax/configs/reformer_imagenet64.gin
+++ /dev/null
@@ -1,114 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters that will vary between experiments:
-# ==============================================================================
-train.model = @trax.models.ReformerLM
-attn_type = @TimeBinCausalAttention
-share_qk = True  # Required when using LSHCausalAttention
-attn_kv = 64
-n_layers = 3
-
-# MemoryEfficientCausalAttention: full attention
-# (no hparams to vary between experiments)
-
-# TimeBinCausalAttention: attend to nearby items
-TimeBinCausalAttention.n_bins = 64
-
-# LSHCausalAttention: locality-sensitive hashing (LSH) attention
-LSHCausalAttention.n_bins = 96
-LSHCausalAttention.n_buckets = 192  # Always 2 * n_bins
-LSHCausalAttention.n_hashes = 2
-LSHCausalAttention.drop_for_hash_rate = 0.0
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 1
-batch_fn.eval_batch_size = 8
-batch_fn.max_eval_length = 12288  # 64 * 64 * 3
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
-inputs.input_name = 'targets'
-inputs.n_chunks = 16
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 2.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 500
-train.eval_steps = 64
-train.inputs = @trax.inputs.inputs
-# train.model: see top
-train.optimizer = @trax.optimizers.Adafactor
-train.train_steps = 500000
-train.trainer_class = @MemoryEfficientTrainer
-train.save_steps = \
-    [1000, 5000, 10000, 20000, 40000, 60000, 80000,
-     100000, 200000, 300000, 400000, 500000]
-
-# Parameters for MemoryEfficientCausalAttention:
-# ==============================================================================
-MemoryEfficientCausalAttention.dropout = 0.0
-MemoryEfficientCausalAttention.loop_stride = 512
-MemoryEfficientCausalAttention.share_qk = %share_qk
-
-# Parameters for TimeBinCausalAttention:
-# ==============================================================================
-TimeBinCausalAttention.dropout = 0.0
-# TimeBinCausalAttention.n_bins: see top
-TimeBinCausalAttention.share_qk = %share_qk
-
-# Parameters for LSHCausalAttention:
-# ==============================================================================
-LSHCausalAttention.allow_duplicate_attention = False
-LSHCausalAttention.attend_across_buckets = False
-LSHCausalAttention.rehash_each_round = True
-# LSHCausalAttention.n_bins: see top
-# LSHCausalAttention.n_buckets: see top
-# LSHCausalAttention.n_hashes: see top
-LSHCausalAttention.one_rng = False
-LSHCausalAttention.hard_k = 0
-LSHCausalAttention.dropout = 0.0
-# LSHCausalAttention.drop_for_hash_rate: see top
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = %attn_type
-TransformerLM.d_attention_key = %attn_kv
-TransformerLM.d_attention_value = %attn_kv
-TransformerLM.d_model = 1024
-TransformerLM.d_ff = 4096
-TransformerLM.dropout = 0.0
-TransformerLM.max_len = 12288  # 64 * 64 * 3
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 8
-TransformerLM.n_layers = %n_layers
-TransformerLM.n_chunks = 16
-TransformerLM.share_qk = %share_qk
-TransformerLM.vocab_size = 256
-
-# Parameters for ReformerLM:
-# ==============================================================================
-ReformerLM.attention_type = %attn_type
-ReformerLM.d_attention_key = %attn_kv
-ReformerLM.d_attention_value = %attn_kv
-ReformerLM.d_model = 1024
-ReformerLM.d_ff = 4096
-ReformerLM.dropout = 0.0
-ReformerLM.max_len = 12288  # 64 * 64 * 3
-ReformerLM.mode = 'train'
-ReformerLM.n_heads = 8
-ReformerLM.n_layers = %n_layers
-ReformerLM.vocab_size = 256
-ReformerLM.n_chunks = 16
-ReformerLM.n_attention_chunks = 1
-ReformerLM.share_qk = %share_qk
diff --git a/tensor2tensor/trax/configs/reformer_large_sweep.yaml b/tensor2tensor/trax/configs/reformer_large_sweep.yaml
deleted file mode 100644
index 9cbf67088..000000000
--- a/tensor2tensor/trax/configs/reformer_large_sweep.yaml
+++ /dev/null
@@ -1,3 +0,0 @@
-MergedMultiHashedCausalAttentionV2.n_hashes: [2, 4]
-TransformerRevnetLM.n_layers: [12, 16, 20, 24]
-MultifactorSchedule.constant: [0.3, 1.0]
diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
deleted file mode 100644
index 7eb13a225..000000000
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb.gin
+++ /dev/null
@@ -1,44 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.learning_rate
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 32
-batch_fn.bucket_length = 32
-batch_fn.buckets = None
-batch_fn.eval_batch_size = 32
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_image_imagenet224'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-EvalAdjustingSchedule.constant = 0.2
-MultifactorSchedule.factors = 'constant * linear_warmup'
-MultifactorSchedule.warmup_steps = 400
-
-# Parameters for Momentum:
-# ==============================================================================
-Momentum.mass = 0.9
-
-
-# Parameters for Resnet50:
-# ==============================================================================
-Resnet50.d_hidden = 64
-Resnet50.n_output_classes = 1001
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 2000
-train.eval_steps = 20
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.Resnet50
-train.optimizer = @trax.optimizers.Momentum
-train.train_steps = 1000000
-train.lr_schedule = @learning_rate.EvalAdjustingSchedule
-
diff --git a/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin b/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
deleted file mode 100644
index 6a49f89fd..000000000
--- a/tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin
+++ /dev/null
@@ -1,44 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.learning_rate
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 32
-batch_fn.bucket_length = 32
-batch_fn.buckets = None
-batch_fn.eval_batch_size = 32
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_image_imagenet224'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-EvalAdjustingSchedule.constant = 0.2
-MultifactorSchedule.factors = 'constant * linear_warmup'
-MultifactorSchedule.warmup_steps = 400
-
-# Parameters for Momentum:
-# ==============================================================================
-Momentum.mass = 0.9
-
-
-# Parameters for Resnet50:
-# ==============================================================================
-Resnet50.d_hidden = 64
-Resnet50.n_output_classes = 1001
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 2000
-train.eval_steps = 20
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.Resnet50
-train.optimizer = @trax.optimizers.Momentum
-train.train_steps = 100000
-train.lr_schedule = @learning_rate.EvalAdjustingSchedule
-
diff --git a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
deleted file mode 100644
index 1125eb032..000000000
--- a/tensor2tensor/trax/configs/transformer_big_lm1b_8gb.gin
+++ /dev/null
@@ -1,54 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 32
-batch_fn.eval_batch_size = 64
-batch_fn.max_eval_length = 512
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_languagemodel_lm1b32k'
-inputs.input_name = 'targets'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 0.1
-MultifactorSchedule.factors = 'constant * linear_warmup'
-MultifactorSchedule.warmup_steps = 16000
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
-lm1b_preprocess.max_target_length = 512
-lm1b_preprocess.max_eval_target_length = 512
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 1000
-train.eval_steps = 10
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerLM
-train.optimizer = @trax.optimizers.SM3
-train.train_steps = 500000
-train.mask_id = 0
-
-# Parameters for DotProductCausalAttention:
-# ==============================================================================
-DotProductCausalAttention.dropout = 0.1
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
-TransformerLM.d_model = 1024
-TransformerLM.d_ff = 8192
-TransformerLM.dropout = 0.1
-TransformerLM.max_len = 2048
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 8
-TransformerLM.n_layers = 8
-TransformerLM.vocab_size = 32000
diff --git a/tensor2tensor/trax/configs/transformer_copy.gin b/tensor2tensor/trax/configs/transformer_copy.gin
deleted file mode 100644
index 86bfbead8..000000000
--- a/tensor2tensor/trax/configs/transformer_copy.gin
+++ /dev/null
@@ -1,77 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-n_symbols = 128
-length = 4096
-batch = 16
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size = %batch
-batch_fn.eval_batch_size = %batch
-batch_fn.max_eval_length = %length
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 'sequence_copy'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 1.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for sequence_copy_inputs:
-# ==============================================================================
-sequence_copy_inputs.vocab_size = %n_symbols
-sequence_copy_inputs.batch_size = %batch
-sequence_copy_inputs.train_lengths = [%length]
-sequence_copy_inputs.eval_lengths = [%length]
-sequence_copy_inputs.reverse = False
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 500
-train.eval_steps = 64
-train.inputs = @trax.inputs.sequence_copy_inputs
-train.model = @trax.models.TransformerLM
-train.optimizer = @trax.optimizers.Adafactor
-train.train_steps = 50000
-train.has_weights = True
-
-# Parameters for MemoryEfficientCausalAttention:
-# ==============================================================================
-MemoryEfficientCausalAttention.dropout = 0.0
-MemoryEfficientCausalAttention.loop_stride = 512
-
-# Parameters for LSHCausalAttention:
-# ==============================================================================
-LSHCausalAttention.allow_duplicate_attention = False
-LSHCausalAttention.attend_across_buckets = True
-LSHCausalAttention.rehash_each_round = True
-LSHCausalAttention.n_bins = 64
-LSHCausalAttention.n_buckets = 128
-LSHCausalAttention.n_hashes = 8
-LSHCausalAttention.one_rng = False
-LSHCausalAttention.hard_k = 0
-LSHCausalAttention.dropout = 0.0
-LSHCausalAttention.drop_for_hash_rate = 0.1
-LSHCausalAttention.factorize_hash = True
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = @trax.layers.LSHCausalAttention
-TransformerLM.d_attention_key = 64
-TransformerLM.d_attention_value = 64
-TransformerLM.d_model = 256
-TransformerLM.d_ff = 256
-TransformerLM.dropout = 0.0
-TransformerLM.max_len = %length
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 4
-TransformerLM.n_layers = 1
-TransformerLM.share_qk = True
-TransformerLM.vocab_size = %n_symbols
diff --git a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin b/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
deleted file mode 100644
index b4cd87f81..000000000
--- a/tensor2tensor/trax/configs/transformer_imdb_8gb.gin
+++ /dev/null
@@ -1,54 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 128
-batch_fn.eval_batch_size = 128
-batch_fn.max_eval_length = 2048
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_sentiment_imdb'
-inputs.input_name = 'targets'
-
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 0.1
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
-lm1b_preprocess.max_target_length = 512
-lm1b_preprocess.max_eval_target_length = 2048
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 100
-train.eval_steps = 10
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerEncoder
-train.train_steps = 1000
-
-# Parameters for DotProductCausalAttention:
-# ==============================================================================
-DotProductCausalAttention.dropout = 0.1
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
-TransformerEncoder.d_model = 512
-TransformerEncoder.d_ff = 2048
-TransformerEncoder.dropout = 0.1
-TransformerEncoder.max_len = 2048
-TransformerEncoder.mode = 'train'
-TransformerEncoder.n_classes = 10
-TransformerEncoder.n_heads = 8
-TransformerEncoder.n_layers = 6
-TransformerEncoder.vocab_size = 32000
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_16gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_16gb.gin
deleted file mode 100644
index d0725cdbb..000000000
--- a/tensor2tensor/trax/configs/transformer_lm1b_16gb.gin
+++ /dev/null
@@ -1,127 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 256
-batch_fn.eval_batch_size = 256
-batch_fn.max_eval_length = 2048
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_languagemodel_lm1b32k'
-inputs.input_name = 'targets'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 1.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
-lm1b_preprocess.max_target_length = 512
-lm1b_preprocess.max_eval_target_length = 2048
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 500
-train.eval_steps = 1
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerLM
-train.optimizer = @trax.optimizers.Adafactor
-train.train_steps = 50000
-train.mask_id = 0
-
-# Parameters for DotProductCausalAttention:
-# ==============================================================================
-DotProductCausalAttention.dropout = 0.1
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
-TransformerLM.d_model = 512
-TransformerLM.d_ff = 2048
-TransformerLM.dropout = 0.1
-TransformerLM.max_len = 2048
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 8
-TransformerLM.n_layers = 6
-TransformerLM.vocab_size = 32000
-
-# ==============================================================================
-# Parameters for the RL hyperparameter tuner; turn on with
-# train.lr_schedule=@learning_rate.PolicySchedule and set
-# PolicySchedule.policy_dir.
-# ==============================================================================
-
-# Parameters for PolicySchedule:
-# ==============================================================================
-PolicySchedule.observation_metrics = (
-    ("train", "metrics/accuracy"),
-    ("train", "metrics/loss"),
-    ("eval", "metrics/accuracy"),
-    ("eval", "metrics/loss"),
-)
-PolicySchedule.include_controls_in_observation = False
-PolicySchedule.control_configs = (
-    ("learning_rate", 1e-3, (1e-9, 1e-2), False),
-    ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
-
-    ("dropout_embedding", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_initial", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_initial", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_initial", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_middle", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_middle", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_middle", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_final", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_final", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_final", 0.1, (0.0, 0.9), True),
-)
-PolicySchedule.observation_range = (0.0, 10.0)
-PolicySchedule.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
-PolicySchedule.policy_and_value_model = @trax.models.TransformerDecoder
-PolicySchedule.policy_and_value_two_towers = False
-
-# Parameters for train:
-# ==============================================================================
-train.nontrainable_param_map = {
-    # "dropout_{layer_type}_{block_index}": "dropout_{layer_type}_{block_group}"
-
-    "dropout_attention_0": "dropout_attention_initial",
-    "dropout_ff_middle_0": "dropout_ff_middle_initial",
-    "dropout_ff_final_0": "dropout_ff_final_initial",
-
-    "dropout_attention_1": "dropout_attention_middle",
-    "dropout_ff_middle_1": "dropout_ff_middle_middle",
-    "dropout_ff_final_1": "dropout_ff_final_middle",
-    "dropout_attention_2": "dropout_attention_middle",
-    "dropout_ff_middle_2": "dropout_ff_middle_middle",
-    "dropout_ff_final_2": "dropout_ff_final_middle",
-    "dropout_attention_3": "dropout_attention_middle",
-    "dropout_ff_middle_3": "dropout_ff_middle_middle",
-    "dropout_ff_final_3": "dropout_ff_final_middle",
-    "dropout_attention_4": "dropout_attention_middle",
-    "dropout_ff_middle_4": "dropout_ff_middle_middle",
-    "dropout_ff_final_4": "dropout_ff_final_middle",
-
-    "dropout_attention_5": "dropout_attention_final",
-    "dropout_ff_middle_5": "dropout_ff_middle_final",
-    "dropout_ff_final_5": "dropout_ff_final_final",
-}
-
-# Parameters for TransformerDecoder:
-# ==============================================================================
-TransformerDecoder.d_model = 64
-TransformerDecoder.d_ff = 128
-TransformerDecoder.dropout = 0.0
-TransformerDecoder.n_heads = 2
-TransformerDecoder.n_layers = 2
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
deleted file mode 100644
index 0cfd5c434..000000000
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb.gin
+++ /dev/null
@@ -1,54 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 128
-batch_fn.eval_batch_size = 128
-batch_fn.max_eval_length = 2048
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_languagemodel_lm1b32k'
-inputs.input_name = 'targets'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 0.3
-MultifactorSchedule.factors = 'constant * linear_warmup'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
-lm1b_preprocess.max_target_length = 512
-lm1b_preprocess.max_eval_target_length = 2048
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 1000
-train.eval_steps = 10
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerLM
-train.optimizer = @trax.optimizers.SM3
-train.train_steps = 500000
-train.mask_id = 0
-
-# Parameters for DotProductCausalAttention:
-# ==============================================================================
-DotProductCausalAttention.dropout = 0.1
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
-TransformerLM.d_model = 512
-TransformerLM.d_ff = 2048
-TransformerLM.dropout = 0.1
-TransformerLM.max_len = 2048
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 8
-TransformerLM.n_layers = 6
-TransformerLM.vocab_size = 32000
diff --git a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin b/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
deleted file mode 100644
index a832553af..000000000
--- a/tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin
+++ /dev/null
@@ -1,54 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 128
-batch_fn.eval_batch_size = 128
-batch_fn.max_eval_length = 2048
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_languagemodel_lm1b32k'
-inputs.input_name = 'targets'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 0.1
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
-lm1b_preprocess.max_target_length = 512
-lm1b_preprocess.max_eval_target_length = 2048
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 1000
-train.eval_steps = 10
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerLM
-train.optimizer = @trax.optimizers.Adam
-train.train_steps = 100000
-train.mask_id = 0
-
-# Parameters for DotProductCausalAttention:
-# ==============================================================================
-DotProductCausalAttention.dropout = 0.1
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
-TransformerLM.d_model = 512
-TransformerLM.d_ff = 2048
-TransformerLM.dropout = 0.1
-TransformerLM.max_len = 2048
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 8
-TransformerLM.n_layers = 6
-TransformerLM.vocab_size = 32000
diff --git a/tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin b/tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin
deleted file mode 100644
index d4ea4390e..000000000
--- a/tensor2tensor/trax/configs/transformer_lm_wmt_ende_16gb.gin
+++ /dev/null
@@ -1,134 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 64
-batch_fn.eval_batch_size = 64
-batch_fn.bucket_length=64
-batch_fn.max_eval_length = 1024
-batch_fn.buckets_include_inputs_in_length=True
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_translate_ende_wmt32k'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 1.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 10000
-
-# Parameters for Adafactor:
-# ==============================================================================
-Adafactor.beta1 = 0.0
-Adafactor.decay_rate = 0.8
-Adafactor.clipping_threshold = 1.0
-Adafactor.epsilon1 = 1e-30
-Adafactor.epsilon2 = 0.001
-Adafactor.factored = True
-Adafactor.multiply_by_parameter_scale = True
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_concat_preprocess
-wmt_concat_preprocess.max_length = 255
-wmt_concat_preprocess.max_eval_length = 511
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 500
-train.eval_steps = 1
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerLM
-train.train_steps = 50000
-train.optimizer = @trax.optimizers.Adafactor
-train.has_weights = True
-train.mask_id = 0
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.d_model = 512
-TransformerLM.d_ff = 2048
-TransformerLM.dropout = 0.1
-TransformerLM.max_len = 2048
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 8
-TransformerLM.n_layers = 6
-TransformerLM.vocab_size = 33300
-
-# ==============================================================================
-# Parameters for the RL hyperparameter tuner; turn on with
-# train.lr_schedule=@learning_rate.PolicySchedule and set
-# PolicySchedule.policy_dir.
-# ==============================================================================
-
-# Parameters for PolicySchedule:
-# ==============================================================================
-PolicySchedule.observation_metrics = (
-    ("train", "metrics/accuracy"),
-    ("train", "metrics/loss"),
-    ("eval", "metrics/accuracy"),
-    ("eval", "metrics/loss"),
-)
-PolicySchedule.include_controls_in_observation = False
-PolicySchedule.control_configs = (
-    ("learning_rate", 1e-3, (1e-9, 1e-2), False),
-    ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
-
-    ("dropout_embedding", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_initial", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_initial", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_initial", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_middle", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_middle", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_middle", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_final", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_final", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_final", 0.1, (0.0, 0.9), True),
-)
-PolicySchedule.observation_range = (0.0, 10.0)
-PolicySchedule.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
-PolicySchedule.policy_and_value_model = @trax.models.TransformerDecoder
-PolicySchedule.policy_and_value_two_towers = False
-
-# Parameters for train:
-# ==============================================================================
-train.nontrainable_param_map = {
-    # "dropout_{layer_type}_{block_index}": "dropout_{layer_type}_{block_group}"
-
-    "dropout_attention_0": "dropout_attention_initial",
-    "dropout_ff_middle_0": "dropout_ff_middle_initial",
-    "dropout_ff_final_0": "dropout_ff_final_initial",
-
-    "dropout_attention_1": "dropout_attention_middle",
-    "dropout_ff_middle_1": "dropout_ff_middle_middle",
-    "dropout_ff_final_1": "dropout_ff_final_middle",
-    "dropout_attention_2": "dropout_attention_middle",
-    "dropout_ff_middle_2": "dropout_ff_middle_middle",
-    "dropout_ff_final_2": "dropout_ff_final_middle",
-    "dropout_attention_3": "dropout_attention_middle",
-    "dropout_ff_middle_3": "dropout_ff_middle_middle",
-    "dropout_ff_final_3": "dropout_ff_final_middle",
-    "dropout_attention_4": "dropout_attention_middle",
-    "dropout_ff_middle_4": "dropout_ff_middle_middle",
-    "dropout_ff_final_4": "dropout_ff_final_middle",
-
-    "dropout_attention_5": "dropout_attention_final",
-    "dropout_ff_middle_5": "dropout_ff_middle_final",
-    "dropout_ff_final_5": "dropout_ff_final_final",
-}
-
-# Parameters for TransformerDecoder:
-# ==============================================================================
-TransformerDecoder.d_model = 64
-TransformerDecoder.d_ff = 128
-TransformerDecoder.dropout = 0.0
-TransformerDecoder.n_heads = 2
-TransformerDecoder.n_layers = 2
diff --git a/tensor2tensor/trax/configs/transformer_lm_wmt_ende_8gb.gin b/tensor2tensor/trax/configs/transformer_lm_wmt_ende_8gb.gin
deleted file mode 100644
index ed43f7919..000000000
--- a/tensor2tensor/trax/configs/transformer_lm_wmt_ende_8gb.gin
+++ /dev/null
@@ -1,64 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 64
-batch_fn.eval_batch_size = 64
-batch_fn.bucket_length=64
-batch_fn.max_eval_length = 1024
-batch_fn.buckets_include_inputs_in_length=True
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_translate_ende_wmt32k'
-
-# Parameters for mask:
-# ==============================================================================
-masked_mean.mask_id = 0
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 1.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 10000
-
-# Parameters for Adafactor:
-# ==============================================================================
-Adafactor.beta1 = 0.0
-Adafactor.decay_rate = 0.8
-Adafactor.clipping_threshold = 1.0
-Adafactor.epsilon1 = 1e-30
-Adafactor.epsilon2 = 0.001
-Adafactor.factored = True
-Adafactor.multiply_by_parameter_scale = True
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_concat_preprocess
-wmt_concat_preprocess.max_length = 255
-wmt_concat_preprocess.max_eval_length = 511
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 1000
-train.eval_steps = 10
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerLM
-train.train_steps = 500000
-train.optimizer = @trax.optimizers.Adafactor
-train.has_weights = True
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.d_model = 512
-TransformerLM.d_ff = 2048
-TransformerLM.dropout = 0.1
-TransformerLM.max_len = 2048
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 8
-TransformerLM.n_layers = 6
-TransformerLM.vocab_size = 33300
diff --git a/tensor2tensor/trax/configs/transformer_ptb_16gb.gin b/tensor2tensor/trax/configs/transformer_ptb_16gb.gin
deleted file mode 100644
index 9d1b796ff..000000000
--- a/tensor2tensor/trax/configs/transformer_ptb_16gb.gin
+++ /dev/null
@@ -1,128 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.learning_rate
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 64
-batch_fn.eval_batch_size = 512
-batch_fn.max_eval_length = 2048
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_languagemodel_ptb10k'
-inputs.input_name = 'targets'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 1.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
-lm1b_preprocess.max_target_length = 512
-lm1b_preprocess.max_eval_target_length = 2048
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 200
-train.eval_steps = 2
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.TransformerLM
-train.optimizer = @trax.optimizers.Adafactor
-train.train_steps = 20000
-train.mask_id = 0
-
-# Parameters for DotProductCausalAttention:
-# ==============================================================================
-DotProductCausalAttention.dropout = 0.5
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
-TransformerLM.d_model = 512
-TransformerLM.d_ff = 2048
-TransformerLM.dropout = 0.5
-TransformerLM.max_len = 2048
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 8
-TransformerLM.n_layers = 6
-TransformerLM.vocab_size = 10240
-
-# ==============================================================================
-# Parameters for the RL hyperparameter tuner; turn on with
-# train.lr_schedule=@learning_rate.PolicySchedule and set
-# PolicySchedule.policy_dir.
-# ==============================================================================
-
-# Parameters for PolicySchedule:
-# ==============================================================================
-PolicySchedule.observation_metrics = (
-    ("train", "metrics/accuracy"),
-    ("train", "metrics/loss"),
-    ("eval", "metrics/accuracy"),
-    ("eval", "metrics/loss"),
-)
-PolicySchedule.include_controls_in_observation = False
-PolicySchedule.control_configs = (
-    ("learning_rate", 1e-3, (1e-9, 1e-2), False),
-    ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
-
-    ("dropout_embedding", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_initial", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_initial", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_initial", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_middle", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_middle", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_middle", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_final", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_final", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_final", 0.1, (0.0, 0.9), True),
-)
-PolicySchedule.observation_range = (0.0, 10.0)
-PolicySchedule.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
-PolicySchedule.policy_and_value_model = @trax.models.TransformerDecoder
-PolicySchedule.policy_and_value_two_towers = False
-
-# Parameters for train:
-# ==============================================================================
-train.nontrainable_param_map = {
-    # "dropout_{layer_type}_{block_index}": "dropout_{layer_type}_{block_group}"
-
-    "dropout_attention_0": "dropout_attention_initial",
-    "dropout_ff_middle_0": "dropout_ff_middle_initial",
-    "dropout_ff_final_0": "dropout_ff_final_initial",
-
-    "dropout_attention_1": "dropout_attention_middle",
-    "dropout_ff_middle_1": "dropout_ff_middle_middle",
-    "dropout_ff_final_1": "dropout_ff_final_middle",
-    "dropout_attention_2": "dropout_attention_middle",
-    "dropout_ff_middle_2": "dropout_ff_middle_middle",
-    "dropout_ff_final_2": "dropout_ff_final_middle",
-    "dropout_attention_3": "dropout_attention_middle",
-    "dropout_ff_middle_3": "dropout_ff_middle_middle",
-    "dropout_ff_final_3": "dropout_ff_final_middle",
-    "dropout_attention_4": "dropout_attention_middle",
-    "dropout_ff_middle_4": "dropout_ff_middle_middle",
-    "dropout_ff_final_4": "dropout_ff_final_middle",
-
-    "dropout_attention_5": "dropout_attention_final",
-    "dropout_ff_middle_5": "dropout_ff_middle_final",
-    "dropout_ff_final_5": "dropout_ff_final_final",
-}
-
-# Parameters for TransformerDecoder:
-# ==============================================================================
-TransformerDecoder.d_model = 64
-TransformerDecoder.d_ff = 128
-TransformerDecoder.dropout = 0.0
-TransformerDecoder.n_heads = 2
-TransformerDecoder.n_layers = 2
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin
deleted file mode 100644
index 40451973d..000000000
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_16gb_adafactor_testing.gin
+++ /dev/null
@@ -1,59 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 128
-batch_fn.eval_batch_size = 64
-batch_fn.max_eval_length = 1024
-batch_fn.buckets_include_inputs_in_length=True
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_translate_ende_wmt32k'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 1.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 20000
-
-# Parameters for Adafactor:
-# ==============================================================================
-Adafactor.beta1 = 0.0
-Adafactor.decay_rate = 0.8
-Adafactor.clipping_threshold = 1.0
-Adafactor.epsilon1 = 1e-30
-Adafactor.epsilon2 = 0.001
-Adafactor.factored = True
-Adafactor.multiply_by_parameter_scale = True
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_preprocess
-wmt_preprocess.max_length = 512
-wmt_preprocess.max_eval_length = 1024
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 1000
-train.eval_steps = 10
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.Transformer
-train.train_steps = 100000
-train.optimizer = @trax.optimizers.Adafactor
-train.mask_id = 0
-
-# Parameters for Transformer:
-# ==============================================================================
-Transformer.d_model = 512
-Transformer.d_ff = 2048
-Transformer.dropout = 0.1
-Transformer.max_len = 2048
-Transformer.mode = 'train'
-Transformer.n_heads = 8
-Transformer.n_layers = 6
-Transformer.input_vocab_size = 33300
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
deleted file mode 100644
index db95a6236..000000000
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adafactor.gin
+++ /dev/null
@@ -1,59 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 64
-batch_fn.eval_batch_size = 64
-batch_fn.max_eval_length = 1024
-batch_fn.buckets_include_inputs_in_length=True
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_translate_ende_wmt32k'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 1.0
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 10000
-
-# Parameters for Adafactor:
-# ==============================================================================
-Adafactor.beta1 = 0.0
-Adafactor.decay_rate = 0.8
-Adafactor.clipping_threshold = 1.0
-Adafactor.epsilon1 = 1e-30
-Adafactor.epsilon2 = 0.001
-Adafactor.factored = True
-Adafactor.multiply_by_parameter_scale = True
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_preprocess
-wmt_preprocess.max_length = 512
-wmt_preprocess.max_eval_length = 1024
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 1000
-train.eval_steps = 10
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.Transformer
-train.train_steps = 500000
-train.optimizer = @trax.optimizers.Adafactor
-train.mask_id = 0
-
-# Parameters for Transformer:
-# ==============================================================================
-Transformer.d_model = 512
-Transformer.d_ff = 2048
-Transformer.dropout = 0.1
-Transformer.max_len = 2048
-Transformer.mode = 'train'
-Transformer.n_heads = 8
-Transformer.n_layers = 6
-Transformer.input_vocab_size = 33300
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
deleted file mode 100644
index 2638003ac..000000000
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_adam.gin
+++ /dev/null
@@ -1,56 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 64
-batch_fn.eval_batch_size = 64
-batch_fn.max_eval_length = 1024
-batch_fn.buckets_include_inputs_in_length=True
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_translate_ende_wmt32k'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-# 0.044 ~= 512^-0.5 = d_model^-0.5
-MultifactorSchedule.constant = 0.044
-MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for Adam:
-# ==============================================================================
-Adam.b1 = 0.9
-Adam.b2 = 0.98
-Adam.eps = 1e-9
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_preprocess
-wmt_preprocess.max_length = 512
-wmt_preprocess.max_eval_length = 1024
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 1000
-train.eval_steps = 10
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.Transformer
-train.train_steps = 500000
-train.optimizer = @trax.optimizers.Adam
-train.mask_id = 0
-
-# Parameters for Transformer:
-# ==============================================================================
-Transformer.d_model= 512
-Transformer.d_ff = 2048
-Transformer.dropout = 0.1
-Transformer.max_len = 2048
-Transformer.mode = 'train'
-Transformer.n_heads = 8
-Transformer.n_layers = 6
-Transformer.input_vocab_size = 33300
diff --git a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin b/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin
deleted file mode 100644
index 04a4356b4..000000000
--- a/tensor2tensor/trax/configs/transformer_wmt_ende_8gb_sm3.gin
+++ /dev/null
@@ -1,53 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 64
-batch_fn.eval_batch_size = 64
-batch_fn.max_eval_length = 1024
-batch_fn.buckets_include_inputs_in_length=True
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_translate_ende_wmt32k'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-MultifactorSchedule.constant = 0.1
-MultifactorSchedule.factors = 'constant * linear_warmup'
-MultifactorSchedule.warmup_steps = 8000
-
-# Parameters for SM3:
-# ==============================================================================
-SM3.momentum = 0.9
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_preprocess
-wmt_preprocess.max_length = 512
-wmt_preprocess.max_eval_length = 1024
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 1000
-train.eval_steps = 10
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.Transformer
-train.train_steps = 500000
-train.optimizer = @trax.optimizers.SM3
-train.mask_id = 0
-
-# Parameters for Transformer:
-# ==============================================================================
-Transformer.d_model= 512
-Transformer.d_ff = 2048
-Transformer.dropout = 0.1
-Transformer.max_len = 2048
-Transformer.mode = 'train'
-Transformer.n_heads = 8
-Transformer.n_layers = 6
-Transformer.input_vocab_size = 33300
diff --git a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin b/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
deleted file mode 100644
index 7edefa0e4..000000000
--- a/tensor2tensor/trax/configs/wide_resnet_cifar10_8gb.gin
+++ /dev/null
@@ -1,81 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.learning_rate
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 256
-batch_fn.bucket_length = 32
-batch_fn.buckets = None
-batch_fn.eval_batch_size = 512
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 'cifar10'
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-EvalAdjustingSchedule.constant = 0.5
-MultifactorSchedule.factors = 'constant * linear_warmup'
-MultifactorSchedule.warmup_steps = 400
-
-# Parameters for Momentum:
-# ==============================================================================
-Momentum.mass = 0.9
-Momentum.weight_decay_rate = 5e-4
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.cifar10_augmentation_preprocess
-
-# Parameters for WideResnet:
-# ==============================================================================
-WideResnet.widen_factor = 10
-WideResnet.n_blocks = 4
-WideResnet.n_output_classes = 10
-
-# Parameters for train:
-# ==============================================================================
-train.eval_frequency = 100
-train.eval_steps = 10
-train.inputs = @trax.inputs.inputs
-train.model = @trax.models.WideResnet
-train.optimizer = @trax.optimizers.Momentum
-train.train_steps = 10000
-train.lr_schedule = @learning_rate.EvalAdjustingSchedule
-
-# ==============================================================================
-# Parameters for the RL hyperparameter tuner; turn on with
-# train.lr_schedule=@learning_rate.PolicySchedule and set
-# PolicySchedule.policy_dir.
-# ==============================================================================
-
-# Parameters for PolicySchedule:
-# ==============================================================================
-PolicySchedule.observation_metrics = (
-    ("train", "metrics/accuracy"),
-    ("train", "metrics/loss"),
-    ("eval", "metrics/accuracy"),
-    ("eval", "metrics/loss"),
-)
-PolicySchedule.include_controls_in_observation = False
-PolicySchedule.control_configs = (
-    ("learning_rate", 0.1, (1e-9, 10.0), False),
-    ("weight_decay_rate", 1e-5, (1e-9, 0.1), False),
-    ("mass", 0.9, (0.0, 0.99), True),
-)
-PolicySchedule.observation_range = (0.0, 10.0)
-PolicySchedule.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
-PolicySchedule.policy_and_value_model = @trax.models.TransformerDecoder
-PolicySchedule.policy_and_value_two_towers = False
-
-# Parameters for TransformerDecoder:
-# ==============================================================================
-TransformerDecoder.d_model = 64
-TransformerDecoder.d_ff = 128
-TransformerDecoder.dropout = 0.0
-TransformerDecoder.n_heads = 2
-TransformerDecoder.n_layers = 2
diff --git a/tensor2tensor/trax/history.py b/tensor2tensor/trax/history.py
deleted file mode 100644
index 4ffc9af3f..000000000
--- a/tensor2tensor/trax/history.py
+++ /dev/null
@@ -1,78 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""trax history."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import collections
-
-from absl import logging
-
-
-class History(object):
-  """History of metrics.
-
-  History contains the metrics recorded during training and evaluation.
-  Save data with history.append and get a sequence of data by calling
-  history.get.
-
-  For example:
-  history.append("train", "metrics/accuracy", 1, 0.04)
-  history.append("train", "metrics/accuracy", 1000, 0.31)
-  history.get("train", "metrics/accuracy")
-  # returns [(1, 0.04), (1000, 0.31)]
-  """
-
-  def __init__(self):
-    # Structure is
-    # values = {
-    #   "mode1": {
-    #     "metric1": [val1, val2],
-    #     ...
-    #   },
-    #   "mode2": ...
-    # }
-    self._values = {}
-
-  def append(self, mode, metric, step, value):
-    """Append (step, value) pair to history for the given mode and metric."""
-    if mode not in self._values:
-      self._values[mode] = collections.defaultdict(list)
-    self._values[mode][metric].append((step, value))
-
-  def get(self, mode, metric):
-    """Get the history for the given metric and mode."""
-    if mode not in self._values:
-      logging.info("Metric %s not found for mode %s", metric, mode)
-      return []
-    return list(self._values[mode][metric])
-
-  @property
-  def modes(self):
-    """Current tracked modes."""
-    return sorted(list(self._values.keys()))
-
-  def metrics_for_mode(self, mode):
-    """Metrics available for a given mode."""
-    if mode not in self._values:
-      logging.info("Mode %s not found", mode)
-      return []
-    return sorted(list(self._values[mode].keys()))
-
-  def __str__(self):
-    return str(self._values)
diff --git a/tensor2tensor/trax/inputs.py b/tensor2tensor/trax/inputs.py
deleted file mode 100644
index 6875c8445..000000000
--- a/tensor2tensor/trax/inputs.py
+++ /dev/null
@@ -1,648 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""trax input pipeline."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import collections
-import os
-import random
-
-import gin
-import numpy as onp
-
-from tensor2tensor import problems_colab as t2t_problems
-from tensor2tensor.trax import backend
-from tensor2tensor.trax.backend import numpy as np
-import tensorflow as tf
-import tensorflow_datasets as tfds
-
-# Inputs is the trax tuple defining the input streams and shapes.
-# * train_stream: training data that will be used for training
-#     may include all the augmentation or selection the training wants
-#     the shape of examples is [batch_fn.batch_size, ...]
-# * train_eval_stream: training data used for evaluation
-#     examples from training data but usually without augmentation
-#     the shape of examples is [batch_fn.eval_batch_size, ...]
-# * eval_stream: evaluation data stream
-#     examples from evaluation data, usually without augmentation
-#     the shape of examples is [batch_fn.eval_batch_size, ...]
-# * input_shape: the shape of inputs
-#     the [...] above, without batch size
-# * input_dtype: the data type of inputs
-# * target_shape: the shape of targets
-#     the [...] above, without batch size
-# * target_dtype: the data type of targets
-
-Inputs = collections.namedtuple(
-    '_Inputs',
-    ['train_stream', 'train_eval_stream', 'eval_stream',
-     'input_shape', 'input_dtype', 'target_shape', 'target_dtype']
-)
-
-# How many examples from the stream to skip at random during training.
-# For now, we skip at most 100K examples for efficiency.
-# TODO(lukaszkaiser): can we improve efficiency, should that be changed?
-_MAX_SKIP_EXAMPLES = 1e5
-
-
-def download_and_prepare(dataset_name, data_dir):
-  """Downloads and prepares T2T or TFDS dataset.
-
-  Args:
-    dataset_name: tfds dataset or t2t problem name prefixed by "t2t_".
-    data_dir: location of existing dataset or None.
-
-  Returns:
-    data_dir: path string of downloaded data.
-  """
-  if not data_dir:
-    data_dir = os.path.expanduser('~/tensorflow_datasets/')
-    dl_dir = os.path.join(data_dir, 'download')
-    tf.logging.info(
-        ('No dataset directory provided. '
-         'Downloading and generating dataset for %s inside data directory %s '
-         'For large datasets it is better to prepare datasets manually!')
-        % (dataset_name, data_dir))
-    if dataset_name.startswith('t2t_'):
-      # Download and run dataset generator for T2T problem.
-      data_dir = os.path.join(data_dir, dataset_name)
-      tf.gfile.MakeDirs(data_dir)
-      tf.gfile.MakeDirs(dl_dir)
-      t2t_problems.problem(
-          dataset_name[len('t2t_'):]).generate_data(data_dir, dl_dir)
-    else:
-      # Download and prepare TFDS dataset.
-      tfds_builder = tfds.builder(dataset_name)
-      tfds_builder.download_and_prepare(download_dir=dl_dir)
-  else:
-    data_dir = os.path.expanduser(data_dir)
-  return data_dir
-
-
-@gin.configurable(blacklist=['n_devices'])
-def inputs(n_devices, dataset_name, data_dir=None, input_name=None,
-           n_chunks=0):
-  """Make Inputs for built-in datasets.
-
-  Args:
-    n_devices: how many devices to build the inputs for.
-    dataset_name: a TFDS or T2T dataset name. If it's a T2T dataset name, prefix
-      with "t2t_".
-    data_dir: data directory.
-    input_name: optional, name of the inputs from the dictionary.
-    n_chunks: optional, into how many pieces should we chunk (large inputs).
-
-  Returns:
-    trax.inputs.Inputs
-  """
-  data_dir = download_and_prepare(dataset_name, data_dir)
-
-  (train_batches, train_eval_batches, eval_batches,
-   input_name, input_shape, input_dtype,
-   target_shape, target_dtype) = _train_and_eval_batches(
-       dataset_name, data_dir, input_name, n_devices)
-
-  if isinstance(input_dtype, tf.DType):
-    input_dtype = input_dtype.as_numpy_dtype
-  if isinstance(target_dtype, tf.DType):
-    target_dtype = target_dtype.as_numpy_dtype
-
-  if input_dtype == np.uint8:  # TPUs don't like uint8s, we cast to ints.
-    input_dtype = np.int32
-  if target_dtype == np.uint8:
-    target_dtype = np.int32
-
-  def numpy_stream(dataset):
-    return dataset_to_stream(dataset, input_name, n_chunks=n_chunks)
-
-  if n_chunks > 0:
-    length = input_shape[0]
-    input_shape = tuple(
-        [tuple([length // n_chunks] + list(input_shape)[1:])] * n_chunks)
-    input_dtype = tuple([input_dtype] * n_chunks)
-    target_shape = tuple(
-        [tuple([length // n_chunks] + list(target_shape)[1:])] * n_chunks)
-    target_dtype = tuple([target_dtype] * n_chunks)
-
-  return Inputs(train_stream=lambda: numpy_stream(train_batches),
-                train_eval_stream=lambda: numpy_stream(train_eval_batches),
-                eval_stream=lambda: numpy_stream(eval_batches),
-                input_shape=input_shape, input_dtype=input_dtype,
-                target_shape=target_shape, target_dtype=target_dtype)
-
-
-@gin.configurable(blacklist=['n_devices'])
-def random_inputs(
-    n_devices,
-    input_shape=gin.REQUIRED, input_dtype=np.int32, input_range=(0, 255),
-    output_shape=gin.REQUIRED, output_dtype=np.int32, output_range=(0, 9)):
-  """Make random Inputs for debugging.
-
-  Args:
-    n_devices: how many devices to build the inputs for.
-    input_shape: the shape of inputs (including batch dimension).
-    input_dtype: the type of the inputs (int32 by default).
-    input_range: the range of inputs (defaults to (0, 255)).
-    output_shape: the shape of outputs (including batch dimension).
-    output_dtype: the type of the outputs (int32 by default).
-    output_range: the range of outputs (defaults to (0, 9)).
-
-  Returns:
-    trax.inputs.Inputs
-  """
-  if input_shape[0] % n_devices != 0:
-    tf.logging.fatal(
-        'n_devices[%d] should divide the first dimension of input_shape[%s]',
-        n_devices, input_shape)
-  if output_shape[0] % n_devices != 0:
-    tf.logging.fatal(
-        'n_devices[%d] should divide the first dimension of output_shape[%s]',
-        n_devices, output_shape)
-
-  def random_minibatches():
-    """Generate a stream of random mini-batches."""
-    if input_dtype in [np.float16, np.float32, np.float64]:
-      rand = onp.random.uniform
-    else:
-      rand = onp.random.random_integers
-    while True:
-      inp = rand(input_range[0], input_range[1], input_shape)
-      inp = inp.astype(input_dtype)
-      out = rand(output_range[0], output_range[1], output_shape)
-      out = out.astype(output_dtype)
-      yield inp, out
-
-  input_shape_without_batch = list(input_shape)[1:]
-  output_shape_without_batch = list(output_shape)[1:]
-  return Inputs(train_stream=random_minibatches,
-                train_eval_stream=random_minibatches,
-                eval_stream=random_minibatches,
-                input_shape=input_shape_without_batch,
-                input_dtype=input_dtype,
-                target_shape=output_shape_without_batch,
-                target_dtype=output_dtype)
-
-
-@gin.configurable(blacklist=['n_devices'])
-def sequence_copy_inputs(
-    n_devices, vocab_size=gin.REQUIRED, batch_size=gin.REQUIRED,
-    train_lengths=gin.REQUIRED, eval_lengths=gin.REQUIRED, reverse=False):
-  """Inputs for the sequence copy problem: 0w0w for w in [1..vocab_size-1]*.
-
-  Args:
-    n_devices: how many devices to build the inputs for.
-    vocab_size: how many symbols to use.
-    batch_size: how large are the batches.
-    train_lengths: lengths of w for training.
-    eval_lengths: lengths of w for eval.
-    reverse: bool (optional, false by default): reverse the second sequence.
-
-  Returns:
-    trax.inputs.Inputs
-  """
-  assert batch_size % n_devices == 0
-  def random_minibatches(length_list):
-    """Generate a stream of random mini-batches."""
-    while True:
-      length = random.choice(length_list)
-      assert length % 2 == 0
-      w_length = (length // 2) - 1
-      w = onp.random.randint(low=1, high=vocab_size-1,
-                             size=(batch_size, w_length))
-      zero = onp.zeros([batch_size, 1], onp.int32)
-      loss_weights = onp.concatenate([onp.zeros((batch_size, w_length+2)),
-                                      onp.ones((batch_size, w_length))], axis=1)
-      if reverse:
-        x = onp.concatenate([zero, w, zero, np.flip(w, axis=1)], axis=1)
-      else:
-        x = onp.concatenate([zero, w, zero, w], axis=1)
-      yield (x, x, loss_weights)  # Here inputs and targets are the same.
-
-  # If there's only one length, make the shape known.
-  example_length = None
-  if (len(train_lengths) == 1 and len(eval_lengths) == 1 and
-      train_lengths[0] == eval_lengths[0]):
-    example_length = train_lengths[0]
-
-  return Inputs(
-      train_stream=lambda: random_minibatches(train_lengths),
-      train_eval_stream=lambda: random_minibatches(train_lengths),
-      eval_stream=lambda: random_minibatches(eval_lengths),
-      input_shape=(example_length,),
-      input_dtype=onp.int32,
-      target_shape=(example_length,),
-      target_dtype=onp.int32)
-
-
-def dataset_to_stream(dataset, input_name, n_chunks=0):
-  """Takes a tf.Dataset and creates a numpy stream of ready batches."""
-  for example in backend.dataset_as_numpy(dataset):
-    features = example[0]
-    inp, out = features[input_name], example[1]
-    mask = features['mask'] if 'mask' in features else None
-    # All input-pipeline processing should be on CPU.
-    with tf.device('cpu:0'):
-      # Some accelerators don't handle uint8 well, cast to int.
-      if isinstance(inp, np.uint8):
-        inp = inp.astype(np.int32)
-      if isinstance(out, np.uint8):
-        out = out.astype(np.int32)
-      if len(out.shape) > 1 and out.shape[-1] == 1:
-        out = np.squeeze(out, axis=-1)
-      if n_chunks > 0:
-        inp = tuple(np.split(inp, n_chunks, axis=1))
-        out = tuple(np.split(out, n_chunks, axis=1))
-    yield (inp, out) if mask is None else (inp, out, mask)
-
-
-@gin.configurable(whitelist=['train_shuffle_files', 'eval_shuffle_files',
-                             'eval_holdout_size'])
-def train_and_eval_dataset(dataset_name, data_dir, eval_holdout_size=0,
-                           train_shuffle_files=True, eval_shuffle_files=False):
-  """Return train and evaluation datasets, feature info and supervised keys.
-
-  Args:
-    dataset_name: a string, the name of the dataset; if it starts with "t2t_"
-      then we'll search T2T Problem registry for it, otherwise we assume it
-      is a dataset from TFDS and load it from there.
-    data_dir: directory where the data is located.
-    eval_holdout_size: float from 0 to <1; if >0 use this much of training data
-      for evaluation (instead of looking for a pre-specified VALIDATION split).
-    train_shuffle_files: Boolean determining whether or not to shuffle the train
-      files at startup. Set to False if you want data determinism.
-    eval_shuffle_files: Boolean determining whether or not to shuffle the test
-      files at startup. Set to False if you want data determinism.
-
-  Returns:
-    a 4-tuple consisting of:
-     * the train tf.Dataset
-     * the eval tf.Dataset
-     * information about features: a python dictionary with feature names
-         as keys and an object as value that provides .shape and .n_classes.
-     * supervised_keys: information what's the input and what's the target,
-         ie., a pair of lists with input and target feature names.
-  """
-  if dataset_name.startswith('t2t_'):
-    return _train_and_eval_dataset_v1(dataset_name[4:], data_dir)
-  dataset_builder = tfds.builder(dataset_name, data_dir=data_dir)
-  info = dataset_builder.info
-  splits = dataset_builder.info.splits
-  if tfds.Split.TRAIN not in splits:
-    raise ValueError('To train we require a train split in the dataset.')
-  train_split = tfds.Split.TRAIN
-  if eval_holdout_size > 0:
-    holdout_percentage = int(eval_holdout_size * 100.0)
-    train_percentage = 100 - holdout_percentage
-    train_split = tfds.Split.TRAIN.subsplit(tfds.percent[:train_percentage])
-    eval_split = tfds.Split.TRAIN.subsplit(tfds.percent[train_percentage:])
-  else:
-    if tfds.Split.VALIDATION not in splits and 'test' not in splits:
-      raise ValueError('We require a validation or test split in the dataset.')
-    eval_split = tfds.Split.VALIDATION
-    if tfds.Split.VALIDATION not in splits:
-      eval_split = tfds.Split.TEST
-  train = tfds.load(
-      name=dataset_name, split=train_split, data_dir=data_dir,
-      shuffle_files=train_shuffle_files)
-  valid = tfds.load(
-      name=dataset_name, split=eval_split, data_dir=data_dir,
-      shuffle_files=eval_shuffle_files)
-  keys = None
-  if info.supervised_keys:
-    keys = ([info.supervised_keys[0]], [info.supervised_keys[1]])
-  return train, valid, info.features, keys
-
-
-def _make_info(shape_list, n_classes, dtype):
-  """Create an info-like tuple for feature given some shapes and vocab size."""
-  feature_info = collections.namedtuple(
-      'FeatureInfo', ['shape', 'n_classes', 'dtype'])
-  cur_shape = list(shape_list[0])
-  # We need to merge the provided shapes, put None where they disagree.
-  for shape in shape_list:
-    if len(shape) != len(cur_shape):
-      raise ValueError('Shapes need to have the same number of dimensions.')
-    for i in range(len(shape)):
-      if cur_shape[i] is not None:
-        if shape[i] != cur_shape[i]:
-          cur_shape[i] = None
-  return feature_info(cur_shape, n_classes, dtype)
-
-
-def _select_features(example, feature_list=None):
-  """Select a subset of features from the example dict."""
-  feature_list = feature_list or ['inputs', 'targets']
-  return {f: example[f] for f in feature_list if f in example}
-
-
-def _eager_dataset_iterator(dataset):
-  for item in dataset:
-    flat = tf.nest.flatten(item)
-    flat = [el.numpy() for el in flat]
-    yield tf.nest.pack_sequence_as(item, flat)
-
-
-def _train_and_eval_dataset_v1(problem_name, data_dir):
-  """Return train and evaluation datasets, feature info and supervised keys."""
-  with tf.device('cpu:0'):
-    problem = t2t_problems.problem(problem_name)
-    train_dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, data_dir)
-    train_dataset = train_dataset.map(_select_features)
-    eval_dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, data_dir)
-    eval_dataset = eval_dataset.map(_select_features)
-    hparams = problem.get_hparams()
-    # We take a few training examples to guess the shapes.
-    input_shapes, target_shapes, examples = [], [], []
-    if tf.executing_eagerly():
-      for example in _eager_dataset_iterator(train_dataset.take(3)):
-        examples.append(example)
-    else:
-      example_tensor = train_dataset.make_one_shot_iterator().get_next()
-      sess = tf.Session()
-      example1 = sess.run(example_tensor)
-      example2 = sess.run(example_tensor)
-      example3 = sess.run(example_tensor)
-      examples = [example1, example2, example3]
-  # We use 'inputs' as input except for purely auto-regressive tasks like
-  # language models where 'targets' are used as input_key.
-  input_key = 'inputs' if 'inputs' in examples[0] else 'targets'
-  supervised_keys = ([input_key], ['targets'])
-  for example in examples:
-    input_shapes.append(list(example[input_key].shape))
-    target_shapes.append(list(example['targets'].shape))
-  input_vocab_size = hparams.vocab_size[input_key]
-  target_vocab_size = hparams.vocab_size['targets']
-  input_dtype = examples[0][input_key].dtype
-  target_dtype = examples[0]['targets'].dtype
-  input_info = _make_info(input_shapes, input_vocab_size, input_dtype)
-  target_info = _make_info(target_shapes, target_vocab_size, target_dtype)
-  info = {input_key: input_info, 'targets': target_info}
-  return train_dataset, eval_dataset, info, supervised_keys
-
-
-@gin.configurable(blacklist=['dataset', 'training', 'shapes',
-                             'target_names', 'n_devices'])
-def batch_fn(dataset, training, shapes, target_names, n_devices,
-             batch_size_per_device=32, batch_size=None, eval_batch_size=32,
-             bucket_length=32, buckets=None,
-             buckets_include_inputs_in_length=False,
-             batch_shuffle_size=128, max_eval_length=None):
-  """Batching function."""
-  del target_names
-  # Batch size is batch_size_per_device * n_devices unless given directly.
-  batch_size = batch_size or batch_size_per_device * n_devices
-  # If bucketing is not specified, check if target shapes are variable.
-  cur_batch_size = batch_size if training else eval_batch_size
-  # Make cur_batch_size divisible by n_devices.
-  cur_batch_size = max(cur_batch_size // n_devices, 1) * n_devices
-  # Create heuristic buckets is none are specified.
-  if buckets is None:
-    variable_target_shapes = False
-    target_shape = shapes[1]
-    for dim in target_shape:
-      if dim is None:
-        variable_target_shapes = True
-    tf.logging.info('Heuristically setting bucketing to %s based on shapes '
-                    'of target tensors.' % variable_target_shapes)
-    if variable_target_shapes:
-      bucket_boundaries = [bucket_length // 4, bucket_length // 2,
-                           bucket_length, bucket_length * 2,
-                           bucket_length * 4, bucket_length * 8,
-                           bucket_length * 16]
-      if not training:
-        max_eval_length = max_eval_length or bucket_length * 32
-        bucket_boundaries[-1] = max_eval_length
-      # We will pad to boundaries which pads to bucket_boundary - 1: add 1 here.
-      bucket_boundaries = [b + 1 for b in bucket_boundaries]
-      bucket_batch_sizes = [cur_batch_size * 4, cur_batch_size * 2,
-                            cur_batch_size, cur_batch_size // 2,
-                            cur_batch_size // 4, cur_batch_size // 8,
-                            cur_batch_size // 16, 1]
-      if not training:
-        bucket_batch_sizes[-2] = cur_batch_size // max_eval_length
-      # Make batch sizes divisible by n_devices.
-      bucket_batch_sizes = [max(b // n_devices, 1) * n_devices
-                            for b in bucket_batch_sizes]
-      buckets = (bucket_boundaries, bucket_batch_sizes)
-
-  if buckets:
-    tf.logging.info('Bucketing with buckets %s.' % str(buckets))
-    def example_length(example_inputs, target):
-      """The length function used by bucket_by_sequence_length to bucket."""
-      other_length = 0
-      if buckets_include_inputs_in_length:
-        other_length = tf.shape(example_inputs['inputs'])[0]
-      return tf.maximum(tf.shape(target)[0], other_length)
-    boundaries, batch_sizes = buckets
-    dataset = dataset.apply(tf.data.experimental.bucket_by_sequence_length(
-        example_length, boundaries, batch_sizes,
-        pad_to_bucket_boundary=True))
-  else:
-    dataset = dataset.padded_batch(cur_batch_size, shapes)
-  if training:
-    return dataset.shuffle(batch_shuffle_size)
-  return dataset
-
-
-@gin.configurable(blacklist=['dataset', 'training'])
-def cifar10_no_augmentation_preprocess(dataset, training):
-  del training
-
-  def cast_image(features, targets):
-    features['image'] = tf.cast(features['image'], tf.float32) / 255.0
-    return features, targets
-
-  dataset = dataset.map(cast_image)
-  return dataset
-
-
-@gin.configurable(blacklist=['dataset', 'training'])
-def cifar10_augmentation_preprocess(dataset, training):
-  """Preprocessing for cifar10 with augmentation (see below)."""
-
-  def augment_image(image):
-    """Image augmentation suitable for CIFAR-10/100.
-
-    As described in https://arxiv.org/pdf/1608.06993v3.pdf (page 5).
-
-    Args:
-      image: a Tensor.
-    Returns:
-      Tensor of the same shape as image.
-    """
-    image = tf.image.resize_image_with_crop_or_pad(image, 40, 40)
-    image = tf.random_crop(image, [32, 32, 3])
-    image = tf.image.random_flip_left_right(image)
-    return image
-
-  def augment(features, targets):
-    features['image'] = augment_image(features['image'])
-    return features, targets
-
-  def cast_image(features, targets):
-    features['image'] = tf.cast(features['image'], tf.float32) / 255.0
-    return features, targets
-
-  if training:
-    dataset = dataset.map(augment)
-  dataset = dataset.map(cast_image)
-  return dataset
-
-
-def no_preprocess(dataset, training):
-  del training
-  return dataset
-
-
-@gin.configurable(blacklist=['dataset', 'training'])
-def concat_preprocess(dataset, training, pad_symbol=0):
-  """Pre-processing function that concatenates input and target for LM."""
-  del training
-
-  def concat(features, targets):
-    inp = features['inputs']
-    pad = tf.expand_dims(tf.zeros_like(inp[0]) + pad_symbol, axis=0)
-    concat = tf.concat([pad, inp, pad, targets], axis=0)
-    # Note: we're updating existing features dictionary here, so make sure
-    # it is not re-used in some other ways outside of this function.
-    features['inputs'] = concat
-    return features, concat
-
-  dataset = dataset.map(concat)
-  return dataset
-
-
-@gin.configurable(blacklist=['dataset', 'training'])
-def lm1b_preprocess(dataset, training,
-                    max_target_length=-1, max_eval_target_length=-1):
-  """Preprocessing for LM1B: filter out targets exceeding maximum length."""
-
-  def target_right_length(_, target):
-    return tf.less(tf.shape(target)[0], max_target_length + 1)
-
-  def eval_target_right_length(_, target):
-    return tf.less(tf.shape(target)[0], max_eval_target_length + 1)
-
-  if max_target_length > 0 and training:
-    dataset = dataset.filter(target_right_length)
-
-  if max_eval_target_length > 0 and not training:
-    dataset = dataset.filter(eval_target_right_length)
-
-  return dataset
-
-
-# TODO(lukaszkaiser): find a single more abstract way of text pre-processing.
-@gin.configurable(blacklist=['dataset', 'training'])
-def wmt_preprocess(dataset, training, max_length=-1, max_eval_length=-1):
-  """Preprocessing for LM1B: filter out targets exceeding maximum length."""
-
-  def train_right_length(example, target):
-    l = tf.maximum(tf.shape(example['inputs'])[0], tf.shape(target)[0])
-    return tf.less(l, max_length + 1)
-
-  def eval_right_length(example, target):
-    l = tf.maximum(tf.shape(example['inputs'])[0], tf.shape(target)[0])
-    return tf.less(l, max_eval_length + 1)
-
-  if max_length > 0 and training:
-    dataset = dataset.filter(train_right_length)
-
-  if max_eval_length > 0 and not training:
-    dataset = dataset.filter(eval_right_length)
-
-  return dataset
-
-
-@gin.configurable(blacklist=['dataset', 'training'])
-def wmt_concat_preprocess(dataset, training, max_length=-1, max_eval_length=-1):
-  """Preprocessing for WMT: filter exceeding maximum length and concatenate."""
-  dataset = wmt_preprocess(dataset, training, max_length, max_eval_length)
-
-  def concat_and_add_mask(features, targets):
-    inp = features['inputs']
-    pad = tf.expand_dims(tf.zeros_like(inp[0]), axis=0)
-    concat = tf.concat([inp, pad, targets], axis=0)
-    mask = tf.concat([tf.zeros_like(inp), pad, tf.ones_like(targets)], axis=0)
-    features['inputs'] = concat
-    features['mask'] = mask
-    return features, concat
-
-  dataset = dataset.map(concat_and_add_mask)
-  return dataset
-
-
-@gin.configurable(whitelist=['preprocess_fun', 'shuffle_buffer_size'])
-def shuffle_and_batch_data(dataset,
-                           target_names,
-                           features_info,
-                           training,
-                           n_devices,
-                           shuffle_buffer_size=1024,
-                           preprocess_fun=no_preprocess):
-  """Shuffle and batch the given dataset."""
-  def append_targets(example):
-    """Append targets to the example dictionary. Needed for Keras."""
-    if len(target_names) == 1:
-      return (example, example[target_names[0]])
-    targets = {}
-    for name in target_names:
-      targets[name] = example[name]
-    return (example, targets)
-  dataset = dataset.map(append_targets)
-  # TODO(pkozakowski): Repeat both the training and evaluation set, so we don't
-  # have incomplete batches during evaluation. This will be a problem when we
-  # add an option to evaluate on the whole dataset, then we'll need to think of
-  # a different solution.
-  dataset = dataset.repeat()
-  if training:
-    # Skip a random fraction at the beginning of the stream.  The skip is
-    # essential for synchronous highly-parallel training to avoid multiple
-    # replicas reading the same data in lock-step.
-    dataset = dataset.skip(random.randint(0, _MAX_SKIP_EXAMPLES))
-  dataset = preprocess_fun(dataset, training)
-  shapes = {k: features_info[k].shape for k in features_info}
-  shapes = (shapes, shapes[target_names[0]])
-  dataset = dataset.shuffle(shuffle_buffer_size)
-  dataset = batch_fn(dataset, training, shapes, target_names, n_devices)
-  return dataset.prefetch(2)
-
-
-def _train_and_eval_batches(dataset, data_dir, input_name, n_devices):
-  """Return train and eval batches with input name and shape."""
-  (train_data, eval_data, features_info, keys) = train_and_eval_dataset(
-      dataset, data_dir)
-  input_names, target_names = keys[0], keys[1]
-  train_batches = shuffle_and_batch_data(
-      train_data, target_names, features_info, training=True,
-      n_devices=n_devices)
-  train_eval_batches = shuffle_and_batch_data(  # Data for eval-on-train.
-      train_data, target_names, features_info, training=False,
-      n_devices=n_devices)
-  eval_batches = shuffle_and_batch_data(
-      eval_data, target_names, features_info, training=False,
-      n_devices=n_devices)
-  input_name = input_name or input_names[0]
-  input_shape = features_info[input_name].shape
-  input_dtype = features_info[input_name].dtype
-  target_shape = features_info[target_names[0]].shape
-  target_dtype = features_info[target_names[0]].dtype
-  return (train_batches, train_eval_batches, eval_batches,
-          input_name, list(input_shape), input_dtype,
-          list(target_shape), target_dtype)
diff --git a/tensor2tensor/trax/inputs_test.py b/tensor2tensor/trax/inputs_test.py
deleted file mode 100644
index 9c145f9f3..000000000
--- a/tensor2tensor/trax/inputs_test.py
+++ /dev/null
@@ -1,72 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.inputs."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gin
-import numpy as np
-from tensor2tensor.trax import inputs
-import tensorflow as tf
-import tensorflow_datasets as tfds
-
-
-def test_dataset_ints(lengths):
-  """Create a test dataset of int64 tensors of shape [length]."""
-  def generator():
-    """Sample generator of sequences of shape [length] of type int64."""
-    for length in lengths:
-      x = np.zeros([length], dtype=np.int64)
-      yield (x, x)  # Inputs and targets are the same here.
-  types = (tf.int64, tf.int64)
-  shapes = (tf.TensorShape([None]), tf.TensorShape([None]))
-  return tf.data.Dataset.from_generator(
-      generator, output_types=types, output_shapes=shapes)
-
-
-class InputsTest(tf.test.TestCase):
-
-  def setUp(self):
-    gin.clear_config()
-
-  def test_batch_fn(self):
-    dataset = test_dataset_ints([32])
-    dataset = dataset.repeat(10)
-    batches = inputs.batch_fn(
-        dataset, True, ([None], [None]), [], 1, batch_size=10)
-    count = 0
-    for example in tfds.as_numpy(batches):
-      count += 1
-      self.assertEqual(example[0].shape[0], 10)  # Batch size = 10.
-    self.assertEqual(count, 1)  # Just one batch here.
-
-  def test_batch_fn_n_devices(self):
-    dataset = test_dataset_ints([32])
-    dataset = dataset.repeat(9)
-    batches = inputs.batch_fn(
-        dataset, True, ([None], [None]), [], 9, batch_size=10)
-    count = 0
-    for example in tfds.as_numpy(batches):
-      count += 1
-      # Batch size adjusted to be divisible by n_devices.
-      self.assertEqual(example[0].shape[0], 9)
-    self.assertEqual(count, 1)  # Just one batch here.
-
-
-if __name__ == "__main__":
-  tf.test.main()
diff --git a/tensor2tensor/trax/jaxboard.py b/tensor2tensor/trax/jaxboard.py
deleted file mode 100644
index 5ea33cee6..000000000
--- a/tensor2tensor/trax/jaxboard.py
+++ /dev/null
@@ -1,350 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Write Summaries from JAX for use with Tensorboard.
-
-See jaxboard_demo.py for example usage.
-"""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import io
-import struct
-import time
-import warnings
-import wave
-import matplotlib as mpl
-# Necessary to prevent attempted Tk import:
-with warnings.catch_warnings():
-  warnings.simplefilter('ignore')
-  mpl.use('Agg')
-# pylint: disable=g-import-not-at-top
-import matplotlib.pyplot as plt
-import numpy as onp
-import tensorflow as tf
-from tensorflow import HistogramProto
-from tensorflow import Summary
-from tensorflow import SummaryMetadata
-from tensorflow.io import gfile
-
-# pylint: disable=g-direct-tensorflow-import
-from tensorflow.core.util import event_pb2
-from tensorflow.python.summary.writer.event_file_writer import EventFileWriter
-# pylint: enable=g-direct-tensorflow-import
-
-
-def _pack_images(images, rows, cols):
-  """Helper utility to make a tiled field of images from numpy arrays.
-
-  Args:
-    images: Image tensor in shape [N, W, H, C].
-    rows: Number of images per row in tiled image.
-    cols: Number of images per column in tiled image.
-
-  Returns:
-    A tiled image of shape [W * rows, H * cols, C].
-    Truncates incomplete rows.
-  """
-  shape = onp.shape(images)
-  width, height, depth = shape[-3:]
-  images = onp.reshape(images, (-1, width, height, depth))
-  batch = onp.shape(images)[0]
-  rows = onp.minimum(rows, batch)
-  cols = onp.minimum(batch // rows, cols)
-  images = images[:rows * cols]
-  images = onp.reshape(images, (rows, cols, width, height, depth))
-  images = onp.transpose(images, [0, 2, 1, 3, 4])
-  images = onp.reshape(images, [rows * width, cols * height, depth])
-  return images
-
-
-class SummaryWriter(object):
-  """Saves data in event and summary protos for tensorboard."""
-
-  def __init__(self, log_dir):
-    """Create a new SummaryWriter.
-
-    Args:
-      log_dir: path to record tfevents files in.
-    """
-    # If needed, create log_dir directory as well as missing parent directories.
-    if not gfile.isdir(log_dir):
-      gfile.makedirs(log_dir)
-
-    self._event_writer = EventFileWriter(log_dir, 10, 120, None)
-    self._step = 0
-    self._closed = False
-
-  def add_summary(self, summary, step):
-    event = event_pb2.Event(summary=summary)
-    event.wall_time = time.time()
-    if step is not None:
-      event.step = int(step)
-    self._event_writer.add_event(event)
-
-  def close(self):
-    """Close SummaryWriter. Final!"""
-    if not self._closed:
-      self._event_writer.close()
-      self._closed = True
-      del self._event_writer
-
-  def __del__(self):  # safe?
-    self.close()
-
-  def flush(self):
-    self._event_writer.flush()
-
-  def scalar(self, tag, value, step=None):
-    """Saves scalar value.
-
-    Args:
-      tag: str: label for this data
-      value: int/float: number to log
-      step: int: training step
-    """
-    value = float(onp.array(value))
-    if step is None:
-      step = self._step
-    else:
-      self._step = step
-    summary = Summary(value=[Summary.Value(tag=tag, simple_value=value)])
-    self.add_summary(summary, step)
-
-  def image(self, tag, image, step=None):
-    """Saves RGB image summary from onp.ndarray [H,W], [H,W,1], or [H,W,3].
-
-    Args:
-      tag: str: label for this data
-      image: ndarray: [H,W], [H,W,1], [H,W,3] save image in greyscale or colors/
-      step: int: training step
-    """
-    image = onp.array(image)
-    if step is None:
-      step = self._step
-    else:
-      self._step = step
-    if len(onp.shape(image)) == 2:
-      image = image[:, :, onp.newaxis]
-    if onp.shape(image)[-1] == 1:
-      image = onp.repeat(image, 3, axis=-1)
-    image_strio = io.BytesIO()
-    plt.imsave(image_strio, image, format='png')
-    image_summary = Summary.Image(
-        encoded_image_string=image_strio.getvalue(),
-        colorspace=3,
-        height=image.shape[0],
-        width=image.shape[1])
-    summary = Summary(value=[Summary.Value(tag=tag, image=image_summary)])
-    self.add_summary(summary, step)
-
-  def images(self, tag, images, step=None, rows=None, cols=None):
-    """Saves (rows, cols) tiled images from onp.ndarray.
-
-    If either rows or cols aren't given, they are determined automatically
-    from the size of the image batch, if neither are given a long column
-    of images is produced. This truncates the image batch rather than padding
-    if it doesn't fill the final row.
-
-    Args:
-      tag: str: label for this data
-      images: ndarray: [N,H,W,1] or [N,H,W,3] to tile in 2d
-      step: int: training step
-      rows: int: number of rows in tile
-      cols: int: number of columns in tile
-    """
-    images = onp.array(images)
-    if step is None:
-      step = self._step
-    else:
-      self._step = step
-    n_images = onp.shape(images)[0]
-    if rows is None and cols is None:
-      rows = 1
-      cols = n_images
-    elif rows is None:
-      rows = n_images // cols
-    elif cols is None:
-      cols = n_images // rows
-    tiled_images = _pack_images(images, rows, cols)
-    self.image(tag, tiled_images, step=step)
-
-  def plot(self, tag, mpl_plt, step=None, close_plot=True):
-    """Saves matplotlib plot output to summary image.
-
-    Args:
-      tag: str: label for this data
-      mpl_plt: matplotlib stateful pyplot object with prepared plotting state
-      step: int: training step
-      close_plot: bool: automatically closes plot
-    """
-    if step is None:
-      step = self._step
-    else:
-      self._step = step
-    fig = mpl_plt.get_current_fig_manager()
-    img_w, img_h = fig.canvas.get_width_height()
-    image_buf = io.BytesIO()
-    mpl_plt.savefig(image_buf, format='png')
-    image_summary = Summary.Image(
-        encoded_image_string=image_buf.getvalue(),
-        colorspace=4,  # RGBA
-        height=img_h,
-        width=img_w)
-    summary = Summary(value=[Summary.Value(tag=tag, image=image_summary)])
-    self.add_summary(summary, step)
-    if close_plot:
-      mpl_plt.close()
-
-  def audio(self, tag, audiodata, step=None, sample_rate=44100):
-    """Saves audio.
-
-    NB: single channel only right now.
-
-    Args:
-      tag: str: label for this data
-      audiodata: ndarray [Nsamples,]: data between (-1.0,1.0) to save as wave
-      step: int: training step
-      sample_rate: sample rate of passed in audio buffer
-    """
-    audiodata = onp.array(audiodata)
-    if step is None:
-      step = self._step
-    else:
-      self._step = step
-    audiodata = onp.clip(onp.squeeze(audiodata), -1, 1)
-    if audiodata.ndim != 1:
-      raise ValueError('Audio data must be 1D.')
-    sample_list = (32767.0 * audiodata).astype(int).tolist()
-    wio = io.BytesIO()
-    wav_buf = wave.open(wio, 'wb')
-    wav_buf.setnchannels(1)
-    wav_buf.setsampwidth(2)
-    wav_buf.setframerate(sample_rate)
-    enc = b''.join([struct.pack('<h', v) for v in sample_list])
-    wav_buf.writeframes(enc)
-    wav_buf.close()
-    encoded_audio_bytes = wio.getvalue()
-    wio.close()
-    audio = Summary.Audio(
-        sample_rate=sample_rate,
-        num_channels=1,
-        length_frames=len(sample_list),
-        encoded_audio_string=encoded_audio_bytes,
-        content_type='audio/wav')
-    summary = Summary(value=[Summary.Value(tag=tag, audio=audio)])
-    self.add_summary(summary, step)
-
-  def histogram(self, tag, values, bins, step=None):
-    """Saves histogram of values.
-
-    Args:
-      tag: str: label for this data
-      values: ndarray: will be flattened by this routine
-      bins: number of bins in histogram, or array of bins for onp.histogram
-      step: int: training step
-    """
-    if step is None:
-      step = self._step
-    else:
-      self._step = step
-    values = onp.array(values)
-    bins = onp.array(bins)
-    values = onp.reshape(values, -1)
-    counts, limits = onp.histogram(values, bins=bins)
-    # boundary logic
-    cum_counts = onp.cumsum(onp.greater(counts, 0, dtype=onp.int32))
-    start, end = onp.searchsorted(
-        cum_counts, [0, cum_counts[-1] - 1], side='right')
-    start, end = int(start), int(end) + 1
-    counts = (
-        counts[start -
-               1:end] if start > 0 else onp.concatenate([[0], counts[:end]]))
-    limits = limits[start:end + 1]
-    sum_sq = values.dot(values)
-    histo = HistogramProto(
-        min=values.min(),
-        max=values.max(),
-        num=len(values),
-        sum=values.sum(),
-        sum_squares=sum_sq,
-        bucket_limit=limits.tolist(),
-        bucket=counts.tolist())
-    summary = Summary(value=[Summary.Value(tag=tag, histo=histo)])
-    self.add_summary(summary, step)
-
-  def text(self, tag, textdata, step=None):
-    """Saves a text summary.
-
-    Args:
-      tag: str: label for this data
-      textdata: string, or 1D/2D list/numpy array of strings
-      step: int: training step
-    Note: markdown formatting is rendered by tensorboard.
-    """
-    if step is None:
-      step = self._step
-    else:
-      self._step = step
-    smd = SummaryMetadata(
-        plugin_data=SummaryMetadata.PluginData(plugin_name='text'))
-    if isinstance(textdata, (str, bytes)):
-      tensor = tf.make_tensor_proto(
-          values=[textdata.encode(encoding='utf_8')], shape=(1,))
-    else:
-      textdata = onp.array(textdata)  # convert lists, jax arrays, etc.
-      datashape = onp.shape(textdata)
-      if len(datashape) == 1:
-        tensor = tf.make_tensor_proto(
-            values=[td.encode(encoding='utf_8') for td in textdata],
-            shape=(datashape[0],))
-      elif len(datashape) == 2:
-        tensor = tf.make_tensor_proto(
-            values=[
-                td.encode(encoding='utf_8') for td in onp.reshape(textdata, -1)
-            ],
-            shape=(datashape[0], datashape[1]))
-    summary = Summary(
-        value=[Summary.Value(tag=tag, metadata=smd, tensor=tensor)])
-    self.add_summary(summary, step)
-
-
-# Copied from gin/tf/utils.py:GinConfigSaverHook
-def markdownify_operative_config_str(string):
-  """Convert an operative config string to markdown format."""
-
-  # TODO(b/37527917): Total hack below. Implement more principled formatting.
-  def process(line):
-    """Convert a single line to markdown format."""
-    if not line.startswith('#'):
-      return '    ' + line
-
-    line = line[2:]
-    if line.startswith('===='):
-      return ''
-    if line.startswith('None'):
-      return '    # None.'
-    if line.endswith(':'):
-      return '#### ' + line
-    return line
-
-  output_lines = []
-  for line in string.splitlines():
-    procd_line = process(line)
-    if procd_line is not None:
-      output_lines.append(procd_line)
-
-  return '\n'.join(output_lines)
diff --git a/tensor2tensor/trax/layers/README.md b/tensor2tensor/trax/layers/README.md
deleted file mode 100644
index 4782fa2dc..000000000
--- a/tensor2tensor/trax/layers/README.md
+++ /dev/null
@@ -1,60 +0,0 @@
-# Trax Layers
-
-
-
-## Base layer structure
-
-All layers inherit from the Layer class and generally need to implement 2
-methods:
-
-```python
-def forward(self, inputs, params=(), state=(), **kwargs):
-  """Computes the layer's output as part of a forward pass through the model."""
-
-def new_params_and_state(self, input_shape, input_dtype, rng):
-  """Returns a (params, state) pair suitable for initializing this layer."""
-```
-
-The base Layer class wraps these functions and provides initialization
-and call functions to be used as follows.
-
-```python
-layer = MyLayer()
-x = np.zeros(10)
-rng = random.get_prng(0)
-layer.initialize_once(x.shape, x.dtype, rng)
-output = layer(x)
-```
-
-## Decorator
-
-To create simple layers, especially ones without parameters, use the layer
-decorator.
-
-```python
-@base.layer()
-def Relu(x, **unused_kwargs):
-  return np.maximum(x, 0.)
-```
-
-## Parameter sharing
-
-Parameters are shared when the same layer object is used.
-
-```python
-standard_mlp = layers.Serial(layers.Dense(10), layers.Dense(10))
-layer = Dense(10)
-shared_parameters_mlp = layers.Serial(layer, layer)
-```
-For this reason, if you call `layer.initialize_once(...)` for the second time
-on an already initialized layer, it will not re-initialize the layer.
-
-## Core layers
-
-* Dense
-* Conv
-
-## Layer composition
-
-* Serial
-* Parallel
diff --git a/tensor2tensor/trax/layers/__init__.py b/tensor2tensor/trax/layers/__init__.py
deleted file mode 100644
index 66643e9d9..000000000
--- a/tensor2tensor/trax/layers/__init__.py
+++ /dev/null
@@ -1,62 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Layers defined in trax."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gin
-# We create a flat layers.* namespace for uniform calling conventions as we
-# upstream changes.
-# pylint: disable=wildcard-import
-from tensor2tensor.trax.layers.attention import *
-from tensor2tensor.trax.layers.base import *
-from tensor2tensor.trax.layers.combinators import *
-from tensor2tensor.trax.layers.convolution import *
-from tensor2tensor.trax.layers.core import *
-from tensor2tensor.trax.layers.initializers import *
-from tensor2tensor.trax.layers.metrics import *
-from tensor2tensor.trax.layers.normalization import *
-from tensor2tensor.trax.layers.pooling import *
-from tensor2tensor.trax.layers.reversible import *
-from tensor2tensor.trax.layers.rnn import *
-
-
-# Ginify
-def layer_configure(*args, **kwargs):
-  kwargs["module"] = "trax.layers"
-  return gin.external_configurable(*args, **kwargs)
-
-# pylint: disable=used-before-assignment
-# pylint: disable=invalid-name
-Relu = layer_configure(Relu)
-Sigmoid = layer_configure(Sigmoid)
-Tanh = layer_configure(Tanh)
-HardSigmoid = layer_configure(HardSigmoid)
-HardTanh = layer_configure(HardTanh)
-Exp = layer_configure(Exp)
-LogSoftmax = layer_configure(LogSoftmax)
-Softmax = layer_configure(Softmax)
-Softplus = layer_configure(Softplus)
-
-DotProductCausalAttention = layer_configure(
-    DotProductCausalAttention, blacklist=["mode"])
-MemoryEfficientCausalAttention = layer_configure(
-    MemoryEfficientCausalAttention, blacklist=["mode"])
-TimeBinCausalAttention = layer_configure(
-    TimeBinCausalAttention, blacklist=["mode"])
-LSHCausalAttention = layer_configure(
-    LSHCausalAttention, blacklist=["mode"])
diff --git a/tensor2tensor/trax/layers/attention.py b/tensor2tensor/trax/layers/attention.py
deleted file mode 100644
index 0a3b7a49c..000000000
--- a/tensor2tensor/trax/layers/attention.py
+++ /dev/null
@@ -1,1355 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Attention Layers."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import math
-import random
-
-import jax
-import numpy as onp
-
-from tensor2tensor.trax import backend
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import combinators as cb
-from tensor2tensor.trax.layers import core
-from tensor2tensor.trax.layers import initializers as init
-
-
-# Layers are always CamelCase, but functions in general are snake_case
-# pylint: disable=invalid-name
-
-
-@base.layer()
-def ShiftRight(x, mode='train', **unused_kwargs):
-  """Layer to shift the tensor to the right by padding on axis 1."""
-  if mode == 'predict':
-    # Do nothing in predict mode, as then the sequence length is 1.
-    return x
-
-  pad_widths = [(0, 0)] * len(x.shape)
-  pad_widths[1] = (1, 0)  # Padding on axis=1
-  padded = np.pad(x, pad_widths, mode='constant',
-                  constant_values=x.dtype.type(0))
-  return padded[:, :-1]
-
-
-@base.layer()
-def CausalMask(x, params, axis=-1, **kwargs):
-  del params, kwargs
-  size = x.shape[axis]
-  return onp.tril(onp.ones((1, size, size), dtype=onp.bool_), k=0)
-
-
-@base.layer()
-def PaddingMask(x, params, pad=0, **kwargs):
-  del params, kwargs
-  return np.reshape(x != pad, (x.shape[0], 1, 1, x.shape[-1]))
-
-
-@base.layer(n_inputs=2)
-def EncoderDecoderMask(x, **unused_kwargs):
-  """Makes encoder-decoder mask from decoder input and a padding mask."""
-  decoder_input, padding_mask = x
-  padding_mask = np.reshape(
-      padding_mask, (padding_mask.shape[0], 1, 1, padding_mask.shape[-1]))
-  # Final mask shape is [batch, 1 for heads, decoder-len, encoder-len].
-  return padding_mask + np.zeros((1, 1, decoder_input.shape[1], 1))
-
-
-class PositionalEncoding(base.Layer):
-  """Implements bare positional encoding."""
-
-  def __init__(self, max_len=2048, mode='train'):
-    super(PositionalEncoding, self).__init__()
-    self._max_len = max_len
-    self._mode = mode
-
-  def forward(self, inputs, params=(), state=(), **kwargs):
-    if self._mode in ('train', 'eval'):
-      x = inputs
-      symbol_size = np.shape(x)[1]
-      return (x + params[:, :symbol_size, :], state)
-    else:
-      assert self._mode == 'predict'
-      # Fast inference: return consectutive elements of the encoding sequence,
-      # storing the index in state.
-      return (inputs + np.expand_dims(params[:, state, :], 1), state + 1)
-
-  def new_params_and_state(self, input_shape, input_dtype, rng):
-    del input_dtype, rng
-    d_feature = input_shape[-1]
-    pe = onp.zeros((self._max_len, d_feature), dtype=onp.float32)
-    position = onp.arange(0, self._max_len)[:, onp.newaxis]
-    div_term = onp.exp(
-        onp.arange(0, d_feature, 2) * -(onp.log(10000.0) / d_feature))
-    pe[:, 0::2] = onp.sin(position * div_term)
-    pe[:, 1::2] = onp.cos(position * div_term)
-    pe = pe[onp.newaxis, :, :]  # [1, self._max_len, d_feature]
-    params = np.array(pe)  # These are trainable parameters, initialized above.
-    state = 0 if self._mode == 'predict' else ()
-    return params, state
-
-
-def DotProductAttention(query, key, value, mask, dropout, mode, rng):
-  """Core dot product self-attention.
-
-  Args:
-    query: array of representations
-    key: array of representations
-    value: array of representations
-    mask: attention-mask, gates attention
-    dropout: float: dropout rate
-    mode: 'eval' or 'train': whether to use dropout
-    rng: JAX PRNGKey: subkey for disposable use
-
-  Returns:
-    Self attention for q, k, v arrays.
-  """
-  depth = np.shape(query)[-1]
-  dots = np.matmul(query, np.swapaxes(key, -1, -2)) / np.sqrt(depth)
-  if mask is not None:
-    # TODO(kitaev): workaround for https://github.com/google/jax/issues/850
-    # We must ensure that both mask and the -1e9 constant have a data dependency
-    # on the input. Broadcasted copies of these use a lot of memory, so they
-    # should be computed at runtime (rather than being global constants).
-    if backend.get_name() == 'jax':
-      mask = jax.lax.tie_in(dots, mask)
-    dots = np.where(mask, dots, np.full_like(dots, -1e9))
-  # Softmax.
-  dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))
-  if dropout >= 1.0:
-    raise ValueError('Dropout rates must be lower than 1.')
-  if dropout is not None and dropout > 0.0 and mode == 'train':
-    keep = backend.random.bernoulli(rng, 1.0 - dropout, dots.shape)
-    dots = np.where(keep, dots / (1.0 - dropout), np.zeros_like(dots))
-  out = np.matmul(dots, value)
-  return out
-
-
-@base.layer(n_inputs=4, n_outputs=2)
-def PureAttention(x, params, n_heads=1, dropout=0.0, mode='train', **kwargs):
-  """Pure transformer-style multi-headed attention.
-
-  Args:
-    x: inputs (q, k, v, mask)
-    params: parameters (none)
-    n_heads: int: number of attention heads
-    dropout: float: dropout rate
-    mode: str: 'train' or 'eval'
-    **kwargs: other arguments including the rng
-
-  Returns:
-    Pure Multi-headed attention result, and the mask.
-  """
-  del params
-  rng = kwargs.get('rng', None)
-  q, k, v, mask = x
-  d_feature = q.shape[-1]
-  assert d_feature % n_heads == 0
-  d_head = d_feature // n_heads
-  nbatch = np.shape(q)[0]
-  # nbatch, seqlen, d_feature --> nbatch, n_heads, seqlen, d_head
-  def SplitHeads(x):
-    return np.transpose(
-        np.reshape(x, (nbatch, -1, n_heads, d_head)), (0, 2, 1, 3))
-  # nbatch, n_heads, seqlen, d_head --> nbatch, seqlen, d_feature
-  def JoinHeads(x):  # pylint: disable=invalid-name
-    return np.reshape(
-        np.transpose(x, (0, 2, 1, 3)), (nbatch, -1, n_heads * d_head))
-  # Split heads, dot-product attention, rejoin heads.
-  res = JoinHeads(
-      DotProductAttention(
-          SplitHeads(q), SplitHeads(k), SplitHeads(v), mask,
-          dropout=dropout, mode=mode, rng=rng))
-  return res, mask  # Keep the mask.
-
-
-def AttentionQKV(d_feature, n_heads=1, dropout=0.0, mode='train'):
-  """Transformer-style multi-headed attention.
-
-  Accepts inputs of the form q, k, v, mask.
-
-  Args:
-    d_feature: int:  dimensionality of feature embedding
-    n_heads: int: number of attention heads
-    dropout: float: dropout rate
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    Multi-headed self-attention result and the mask.
-  """
-  return [
-      cb.Parallel(
-          core.Dense(d_feature),
-          core.Dense(d_feature),
-          core.Dense(d_feature),
-      ),
-      PureAttention(  # pylint: disable=no-value-for-parameter
-          n_heads=n_heads, dropout=dropout, mode=mode),
-      core.Dense(d_feature),
-  ]
-
-
-def Attention(d_feature, n_heads=1, dropout=0.0, mode='train'):
-  """Transformer-style multi-headed attention.
-
-  Accepts inputs of the form (x, mask) and constructs (q, k, v) from x.
-
-  Args:
-    d_feature: int:  dimensionality of feature embedding
-    n_heads: int: number of attention heads
-    dropout: float: dropout rate
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    Multi-headed self-attention result and the mask.
-  """
-  return [
-      cb.Dup(), cb.Dup(),
-      AttentionQKV(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
-  ]
-
-
-def BasicCausalAttention(d_feature, n_heads=1, dropout=0.0, mode='train'):
-  """Transformer-style multi-headed causal attention.
-
-  This implementation is less configurable than the CausalAttention layer
-  defined below, but it shares code with the non-causal attention.
-
-  # TODO(jonni,lukaszkaiser): standardize and improve layer comments.
-  Accepts inputs of the form x and constructs (q, k, v) and causal mask from x.
-
-  Args:
-    d_feature: int:  dimensionality of feature embedding
-    n_heads: int: number of attention heads
-    dropout: float: dropout rate
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    Multi-headed self-attention result.
-  """
-  return [
-      cb.Dup(),
-      cb.Parallel([], CausalMask(axis=-2)),  # pylint: disable=no-value-for-parameter
-      Attention(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
-      cb.Parallel([], cb.Drop()),  # x
-  ]
-
-
-class ShiftRightLearned(base.Layer):
-  """Layer constructor function for shifting right by a learned vector."""
-
-  def __init__(self, initializer=init.RandomNormalInitializer(0.01)):
-    super(ShiftRightLearned, self).__init__()
-    self._initializer = initializer
-
-  def forward(self, x, params=(), state=(), **kwargs):
-    del kwargs
-    c = backend.numpy.reshape(params, [1, 1, -1])
-    c += backend.numpy.zeros((x.shape[0], 1, x.shape[2]), dtype=x.dtype)
-    return backend.numpy.concatenate([c, x], axis=1)[:, :-1, :], state
-
-  def new_params_and_state(self, input_shape, input_dtype, rng):
-    del input_dtype
-    b = self._initializer((input_shape[-1],), rng)
-    return b, ()
-
-
-class ComputeAttentionHeads(base.Layer):
-  """Computes queries/keys/values via linear projection.
-
-  The output shape is (n_batch * n_heads, seqlen, d_head); the batch and head
-  dimensions are fused to allow for more efficient memory layouts.
-  """
-
-  def __init__(self, n_heads=1, d_head=64,
-               kernel_initializer=init.GlorotUniformInitializer()):
-    super(ComputeAttentionHeads, self).__init__()
-    self._n_heads = n_heads
-    self._d_head = d_head
-    self._kernel_initializer = kernel_initializer
-    # The lack of a bias term here is consistent with the tensor2tensor
-    # implementation, and shouldn't have an effect on modeling quality.
-    # Note that AttentionQKV above is different in that it uses a bias term.
-
-  def forward(self, x, params=(), state=(), **kwargs):
-    del kwargs
-    seqlen = x.shape[1]
-    res = np.dot(x, params)
-
-    # n_batch, seqlen, n_heads*d_head -> n_batch, seqlen, n_heads, d_head
-    res = np.reshape(res, (x.shape[0], seqlen, self._n_heads, self._d_head))
-    # n_batch, seqlen, n_heads, d_head -> n_batch, n_heads, seqlen, d_head
-    res = np.transpose(res, (0, 2, 1, 3))
-    # n_batch, n_heads, seqlen, d_head -> n_batch*n_heads, seqlen, d_head
-    res = np.reshape(res, (-1, seqlen, self._d_head))
-
-    return res, state
-
-  def new_params_and_state(self, input_shape, input_dtype, rng):
-    del input_dtype
-    w = self._kernel_initializer(
-        (input_shape[-1], self._n_heads * self._d_head), rng)
-    return w, ()
-
-
-class ComputeAttentionOutput(base.Layer):
-  """Joins outputs from different heads via linear projection."""
-
-  def __init__(self, n_heads=1, d_model=1024,
-               kernel_initializer=init.GlorotUniformInitializer()):
-    super(ComputeAttentionOutput, self).__init__()
-    self._n_heads = n_heads
-    self._d_model = d_model
-    self._kernel_initializer = kernel_initializer
-    # The lack of a bias term here is consistent with the tensor2tensor
-    # implementation, and shouldn't have an effect on modeling quality.
-    # Note that AttentionQKV above is different in that it uses a bias term.
-
-  def forward(self, x, params=(), state=(), **kwargs):
-    del kwargs
-    seqlen = x.shape[1]
-    d_head = x.shape[2]
-
-    x = np.reshape(x, (-1, self._n_heads, seqlen, d_head))
-    x = np.transpose(x, (0, 2, 1, 3))  # -> n_batch, seqlen, n_heads, d_head
-    x = np.reshape(x, (-1, seqlen, self._n_heads * d_head))
-
-    return np.dot(x, params), state
-
-  def new_params_and_state(self, input_shape, input_dtype, rng):
-    del input_dtype
-    w = self._kernel_initializer(
-        (input_shape[-1] * self._n_heads, self._d_model), rng)
-    return w, ()
-
-
-class BaseCausalAttention(base.Layer):
-  """Base class for variants of causal self-attention."""
-
-  def __init__(self, mode='train'):
-    del mode
-    super(BaseCausalAttention, self).__init__(n_inputs=3)
-
-  def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
-    """Forward pass for the attention layer."""
-    raise NotImplementedError()
-
-  def forward_and_backward(self, inputs, grad, **kwargs):
-    """Performs both forward and backward pass for the attention layer.
-
-    This is used in reversible models: for the backward pass of a reversible
-    model, we need to compute both the forward direction (to recover the
-    previous layer's activations) and the backward direction simultaneously.
-    Some computation can be shared between the forward and backward directions,
-    which makes it more efficient to implement them jointly.
-
-    This method assumes that the layer is stateless and has no parameters.
-
-    Args:
-      inputs: A tuple (q, k, v), where each element has shape
-          n_batch*n_heads, seqlen, d_head
-      grad: gradient signal for the layer output.
-      **kwargs: kwargs for the layer
-
-    Returns:
-      A nested-tuple structure (output, (q_grad, k_grad, v_grad)) that contains
-      the output of the forward pass and the gradient signal for each input.
-    """
-    raise NotImplementedError()
-
-
-def _fast_inference_init_state(input_shapes, input_dtypes, buffer_length):
-  """Initializes state of a causal attention layer for fast inference."""
-  ((batch_size, _, _), _, _) = input_shapes
-  def init_buffer(shape, dtype):
-    (_, _, depth) = shape
-    return np.zeros((batch_size, buffer_length, depth), dtype=dtype)
-  (_, k, v) = tuple(
-      init_buffer(shape, dtype)
-      for (shape, dtype) in zip(input_shapes, input_dtypes)
-  )
-  mask = np.zeros((batch_size, 1, buffer_length))
-  index = 0
-  state = (k, v, mask, index)
-  return state
-
-
-def _fast_inference_update_state(inputs, state):
-  """Updates state of a causal attention layer for fast inference."""
-  assert backend.get_name() == 'jax', (
-      'JAX backend is required to use the predict mode.')
-  for x in inputs:
-    assert x.shape[1] == 1, (
-        'In predict mode the input sequence must be of length 1.')
-  # Fast inference: run with only 1 query in each step, storing the sequence
-  # of keys and values calculated so far in state.
-  (_, new_k, new_v) = inputs
-  (ks, vs, mask, index) = state
-  ks = jax.ops.index_update(ks, jax.ops.index[:, index, :], new_k[:, 0, :])
-  vs = jax.ops.index_update(vs, jax.ops.index[:, index, :], new_v[:, 0, :])
-  mask = jax.ops.index_update(mask, jax.ops.index[:, :, index], 1)
-  return (ks, vs, mask, index + 1)
-
-
-class DotProductCausalAttention(BaseCausalAttention):
-  """A standard (non-memory-efficient) dot product attention implementation."""
-
-  def __init__(self, dropout=0.0, mode='train'):
-    super(DotProductCausalAttention, self).__init__()
-    self._dropout = dropout
-    self._mode = mode
-
-  def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
-    del params
-    q, k, v = inputs
-    if self._mode in ('train', 'eval'):
-      mask_size = q.shape[-2]
-      # Not all backends define np.tril. However, using onp.tril is inefficient
-      # in that it creates a large global constant. TODO(kitaev): try to find an
-      # alternative that works across all backends.
-      if backend.get_name() == 'jax':
-        mask = np.tril(
-            np.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
-      else:
-        mask = onp.tril(
-            onp.ones((1, mask_size, mask_size), dtype=onp.bool_), k=0)
-    else:
-      assert self._mode == 'predict'
-      state = _fast_inference_update_state(inputs, state)
-      (k, v, mask, _) = state
-
-    res = DotProductAttention(
-        q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=rng)
-    return res, state
-
-  def forward_and_backward(self, inputs, ct, **kwargs):
-    assert backend.get_name() == 'jax', (
-        'JAX backend is required to use forward_and_backward.')
-    # Simultaneous forward pass and backprop through the attention mechanism.
-    def _do_forward(x):  # pylint: disable=invalid-name
-      res, _ = self.forward(x, **kwargs)
-      return res
-    output, vjpfun = jax.vjp(_do_forward, inputs)
-    return output, vjpfun(ct)[0]
-
-  def new_params_and_state(self, input_shapes, input_dtype, rng):
-    if self._mode in ('train', 'eval'):
-      return (), ()
-
-    assert self._mode == 'predict'
-    params = ()
-    # Buffer length is hardcoded for now. TODO(pkozakowski): Pass it from the
-    # model.
-    max_len = 2048
-    state = _fast_inference_init_state(input_shapes, input_dtype, max_len)
-    return params, state
-
-
-class MemoryEfficientCausalAttention(BaseCausalAttention):
-  """Memory-efficient dot product attention.
-
-  This layer performs causal attention on long sequences without running out
-  of memory. Instead of computing dot products for all query-key pairs at once,
-  it uses a loop to compute attention for a small set of query positions at a
-  time. The "loop_stride" parameter controls how many query positions are
-  considered at each iteration of the loop.
-
-  Note that this class does not slice along the batch/head dimension. Looping
-  over batch elements and heads instead of query positions is also a viable
-  option. We haven't implemented it, but it may perform well, too.
-  """
-
-  def __init__(self, loop_stride, dropout, mode, share_qk=False, hard_k=0):
-    assert backend.get_name() == 'jax', (
-        'JAX backend is required to use MemoryEfficientCausalAttention.')
-    super(MemoryEfficientCausalAttention, self).__init__()
-    self._loop_stride = loop_stride
-    if dropout >= 1.0:
-      raise ValueError('Dropout rates must be lower than 1.')
-    if mode == 'train':
-      self.dropout = dropout
-    else:
-      self.dropout = None
-    self._share_qk = share_qk
-    self._hard_k = hard_k
-
-  def forward(self, inputs, params=(), state=(), **kwargs):
-    del params
-    output, _ = self.forward_and_backward(inputs, None, **kwargs)
-    return output, state
-
-  def has_backward(self):
-    return True
-
-  def backward(self, inputs, output, ct, params=(), state=(), **kwargs):
-    del output, params, state
-    _, inputs_ct = self.forward_and_backward(inputs, ct, **kwargs)
-    return inputs_ct, ()
-
-  def make_unit_length(self, x, epsilon=1e-6):
-    variance = np.mean(x**2, axis=-1, keepdims=True)
-    norm_inputs = x / np.sqrt(variance + epsilon)
-    return norm_inputs
-
-  def forward_and_backward(self, inputs, ct, rng=None, **kwargs):
-    del kwargs
-    query, key, value = inputs
-    depth = np.shape(query)[-1]
-    do_backprop = ct is not None
-    # jax uses the term cotangent (ct) to refer to gradient signals, and
-    # vector-Jacobian product (vjp) for back-propagation through a layer.
-
-    def make_mask(N, M, k):  # pylint: disable=invalid-name
-      """Constructs a slice of the causal attention mask.
-
-      Args:
-        N: number of query positions
-        M: number of key positions
-        k: position of the initial query element
-
-      Returns:
-        N x M mask, where 1.0 indicates that attention is not allowed.
-      """
-      x = jax.lax.tie_in(k, np.arange(N, dtype=np.int32))
-      y = jax.lax.tie_in(k, np.arange(M, dtype=np.int32))
-      mask = jax.lax.lt(
-          (jax.lax.broadcast_in_dim(
-              x, shape=(N, M), broadcast_dimensions=(0,)) + k),
-          jax.lax.broadcast(y, [N]))
-      mask = jax.lax.convert_element_type(mask, np.float32)
-      return mask
-
-    def make_self_mask(N, M, k):  # pylint: disable=invalid-name
-      """Masks out elements attending to self.
-
-      Args:
-        N: number of query positions
-        M: number of key positions
-        k: position of the initial query element
-
-      Returns:
-        N x M mask, where 1.0 indicates that attention is not allowed.
-      """
-      x = jax.lax.tie_in(k, np.arange(N, dtype=np.int32))
-      y = jax.lax.tie_in(k, np.arange(M, dtype=np.int32))
-      mask = jax.lax.eq(
-          (jax.lax.broadcast_in_dim(
-              x, shape=(N, M), broadcast_dimensions=(0,)) + k),
-          jax.lax.broadcast(y, [N]))
-      mask = jax.lax.convert_element_type(mask, np.float32)
-      return mask
-
-    def forward_slice(query_slice, q_loop_idx, key, value):  # pylint: disable=invalid-name
-      """Forward pass for a subset of the query vectors."""
-      if self._share_qk:
-        key = self.make_unit_length(key)
-
-      dots = np.matmul(
-          query_slice, np.swapaxes(key, -1, -2)) / np.sqrt(depth)
-
-      # Causal masking
-      mask = make_mask(dots.shape[-2], dots.shape[-1], q_loop_idx)
-      dots = dots - 1e9 * mask
-
-      # Mask out attention to self except when no other targets are available.
-      if self._share_qk:
-        self_mask = make_self_mask(dots.shape[-2], dots.shape[-1], q_loop_idx)
-        dots = dots - 1e5 * self_mask
-
-      # Softmax.
-      dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))
-
-      if self.dropout is not None and self.dropout > 0.0:
-        # Dropout is broadcast across the batch+head dimension
-        dropout_shape = (1, dots.shape[-2], dots.shape[-1])
-        slice_rng = jax.random.fold_in(rng, q_loop_idx)
-        keep_prob = jax.lax.tie_in(dots, 1.0 - self.dropout)
-        keep = backend.random.bernoulli(slice_rng, keep_prob, dropout_shape)
-        multiplier = keep.astype(dots.dtype) / jax.lax.tie_in(keep, keep_prob)
-        dots = dots * multiplier
-
-      if self._hard_k > 0:
-        top_k = np.sort(dots)[..., -self._hard_k]  # Get the top-kth weight.
-        top_k = jax.lax.stop_gradient(top_k)
-        dots -= top_k[..., np.newaxis]  # Subtract (be 0 for lower ones).
-        dots = np.maximum(dots, 0)
-        dots_sum = np.sum(dots, axis=-1, keepdims=True)  # Re-normalize.
-        dots /= dots_sum  # Re-normalize.
-
-      out_slice = np.matmul(dots, value)
-      return out_slice
-
-    def forward_and_vjp_slice(query_slice, q_loop_idx, key, value, ct_slice):  # pylint: disable=invalid-name
-      # Capture q_loop_idx to avoid calculated gradients wrt. it.
-      def forward_slice_with_q_loop_idx(query_slice, key, value):  # pylint: disable=invalid-name
-        return forward_slice(query_slice, q_loop_idx, key, value)
-
-      output_slice, vjpfun = jax.vjp(
-          forward_slice_with_q_loop_idx, query_slice, key, value)
-      return output_slice, vjpfun(ct_slice)
-
-    q_loop_idx = np.zeros((), dtype=np.int32)
-    q_loop_max = query.shape[-2]
-    q_loop_stride = self._loop_stride
-    assert q_loop_max % q_loop_stride == 0, (
-        'Stride must evenly divide the number of query elements.')
-
-    out_accum = np.zeros_like(query)
-    if do_backprop:
-      query_ct_accum = np.zeros_like(query)
-      key_ct_accum = np.zeros_like(key)
-      value_ct_accum = np.zeros_like(value)
-      init_vals = (
-          q_loop_idx, out_accum,
-          query_ct_accum, key_ct_accum, value_ct_accum)
-    else:
-      init_vals = (q_loop_idx, out_accum)
-
-    def cond_fun(vals):  # pylint: disable=invalid-name
-      q_loop_idx = vals[0]
-      return jax.lax.lt(q_loop_idx, q_loop_max)
-
-    def body_fun(vals):  # pylint: disable=invalid-name
-      """Compute a slice of the attention mechanism."""
-      if do_backprop:
-        (q_loop_idx, out_accum,
-         query_ct_accum, key_ct_accum, value_ct_accum) = vals
-      else:
-        q_loop_idx, out_accum = vals
-
-      query_slice = jax.lax.dynamic_slice_in_dim(
-          query, q_loop_idx, q_loop_stride, axis=-2)
-
-      if do_backprop:
-        ct_slice = jax.lax.dynamic_slice_in_dim(
-            ct, q_loop_idx, q_loop_stride, axis=-2)
-        out_slice, partial_ct = forward_and_vjp_slice(
-            query_slice, q_loop_idx, key, value, ct_slice)
-        query_ct_accum = jax.lax.dynamic_update_slice_in_dim(
-            query_ct_accum, partial_ct[0], q_loop_idx, axis=-2)
-        key_ct_accum = key_ct_accum + partial_ct[1]
-        value_ct_accum = value_ct_accum + partial_ct[2]
-      else:
-        out_slice = forward_slice(query_slice, q_loop_idx, key, value)
-
-      out_accum = jax.lax.dynamic_update_slice_in_dim(
-          out_accum, out_slice, q_loop_idx, axis=-2)
-      q_loop_idx = q_loop_idx + q_loop_stride
-
-      if do_backprop:
-        return (q_loop_idx, out_accum,
-                query_ct_accum, key_ct_accum, value_ct_accum)
-      else:
-        return (q_loop_idx, out_accum)
-
-    final_vals = jax.lax.while_loop(cond_fun, body_fun, init_vals)
-
-    if not do_backprop:
-      return final_vals[1], None
-    else:
-      return final_vals[1], final_vals[2:]
-
-
-class TimeBinCausalAttention(BaseCausalAttention):
-  """Causal attention where only nearby chunks of items attend to each other."""
-
-  def __init__(self, mode, dropout=0.0, bin_length=None, n_bins=None,
-               share_qk=False):
-    super(TimeBinCausalAttention, self).__init__()
-    if (bin_length is None) == (n_bins is None):
-      raise ValueError('Exactly one of {bin_length, n_bins} must be set.')
-    self.bin_length = bin_length
-    self.n_bins = n_bins
-    self._share_qk = share_qk
-    if dropout >= 1.0:
-      raise ValueError('Dropout rates must be lower than 1.')
-    if mode == 'train':
-      self.dropout = dropout
-    else:
-      self.dropout = 0.0
-    self._mode = mode
-
-  def forward_and_backward(self, inputs, ct, **kwargs):
-    assert backend.get_name() == 'jax', (
-        'JAX backend is required to use forward_and_backward.')
-    # Simultaneous forward pass and backprop through the attention mechanism.
-    def _do_forward(x):  # pylint: disable=invalid-name
-      res, _ = self.forward(x, **kwargs)
-      return res
-    output, vjpfun = jax.vjp(_do_forward, inputs)
-    return output, vjpfun(ct)[0]
-
-  def make_unit_length(self, x, epsilon=1e-6):
-    variance = np.mean(x**2, axis=-1, keepdims=True)
-    norm_inputs = x / np.sqrt(variance + epsilon)
-    return norm_inputs
-
-  def _pad_inputs(self, inputs):
-    seq_len = inputs[0].shape[-2]
-    n_bins = self.n_bins
-    bin_length = self.bin_length
-    if n_bins is None:
-      n_bins = int(math.ceil(seq_len / bin_length))
-    else:
-      bin_length = int(math.ceil(seq_len / n_bins))
-    pad_len = n_bins * bin_length - seq_len
-
-    def pad_input(x):
-      pad_widths = [(0, 0)] * len(x.shape)
-      pad_widths[-2] = (0, pad_len)  # Padding on axis=-2
-      return np.pad(x, pad_widths, mode='constant',
-                    constant_values=x.dtype.type(0))
-
-    padded_inputs = tuple(map(pad_input, inputs))
-    return (padded_inputs, seq_len, n_bins)
-
-  def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
-    del params, kwargs
-    if self._mode in ('train', 'eval'):
-      output = self._forward_train_eval(inputs, rng)
-      return (output, state)
-    else:
-      assert self._mode == 'predict'
-      return self._forward_predict(inputs, state, rng)
-
-  def _forward_train_eval(self, inputs, rng):
-    (inputs, original_len, n_bins) = self._pad_inputs(inputs)
-    q, k, v = inputs
-    seqlen = q.shape[-2]
-    # q/k/v are n_batch*n_heads, seqlen, d_head
-    # Time indices for causal masking.
-    t = jax.lax.tie_in(q, np.arange(seqlen))
-
-    # Split off a "bin" axis for chunks of consecutive items.
-    bq_t = np.reshape(t, (n_bins, -1))
-    bq = np.reshape(q, (q.shape[0], n_bins, -1, q.shape[-1]))
-    if self._share_qk:
-      bk = self.make_unit_length(bq)
-    else:
-      bk = np.reshape(k, (k.shape[0], n_bins, -1, k.shape[-1]))
-    bv = np.reshape(v, (v.shape[0], n_bins, -1, v.shape[-1]))
-
-    # Allow each chunk to attend within itself, and also one chunk back.
-    def look_one_back(x):
-      # Output: pairs [ bin_i bin_{i-1} ] concatenated on the time axis.
-      if len(x.shape) == 2:
-        x_extra = np.concatenate([x[-1:, :], x[:-1, :]], axis=0)
-        return np.concatenate([x, x_extra], axis=1)
-      else:
-        assert len(x.shape) == 4
-        x_extra = np.concatenate([x[:, -1:, :, :], x[:, :-1, :, :]], axis=1)
-        return np.concatenate([x, x_extra], axis=2)
-
-    bkv_t = look_one_back(bq_t)
-    bk = look_one_back(bk)
-    bv = look_one_back(bv)
-
-    # Dot-product attention.
-    dots = np.matmul(bq, np.swapaxes(bk, -1, -2)) / np.sqrt(bq.shape[-1])
-
-    # Causal masking based on the time indices.
-    mask = jax.lax.convert_element_type(
-        jax.lax.lt(bq_t[None, :, :, None], bkv_t[None, :, None, :]),
-        np.float32)
-    dots = dots - 1e9 * mask
-
-    # Mask out attention to self except when no other targets are available.
-    if self._share_qk:
-      self_mask = jax.lax.broadcasted_eye(dots.dtype, dots.shape, (2, 3))
-      self_mask = jax.lax.tie_in(dots, self_mask)
-      dots = dots - 1e5 * self_mask
-
-    if self.dropout > 0.0:
-      # Dropout is broadcast across the batch+head dimension
-      dropout_shape = (1, dots.shape[-3], dots.shape[-2], dots.shape[-1])
-      keep_prob = jax.lax.tie_in(dots, 1.0 - self.dropout)
-      keep = backend.random.bernoulli(rng, keep_prob, dropout_shape)
-      multiplier = keep.astype(dots.dtype) / jax.lax.tie_in(keep, keep_prob)
-      dots = dots * multiplier
-
-    # Softmax.
-    dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))
-    bo = np.matmul(dots, bv)
-
-    output = np.reshape(bo, (bo.shape[0], -1, bo.shape[-1]))
-    assert output.shape == v.shape
-    return output[..., :original_len, :]
-
-  def _forward_predict(self, inputs, state, rng):
-    state = _fast_inference_update_state(inputs, state)
-
-    (q, _, _) = inputs
-    (ks, vs, mask, index) = state
-    output = DotProductAttention(
-        q, ks, vs, mask, dropout=self.dropout, mode=self._mode, rng=rng
-    )
-
-    def roll_state(state):
-      """Rolls the buffers backward to make space for new data."""
-      (ks, vs, mask, index) = state
-      # Move the second bin into the first one's place in both buffers.
-      def roll_buffer(buf):
-        return jax.ops.index_update(
-            buf,
-            jax.ops.index[:, :self.bin_length, :],
-            buf[:, self.bin_length:, :],
-        )
-      (ks, vs) = map(roll_buffer, (ks, vs))
-      # Zero out the second bin in the mask.
-      mask = jax.ops.index_update(
-          mask, jax.ops.index[:, :, self.bin_length:], 0
-      )
-      # Update the index to match the rolled buffers.
-      index -= self.bin_length
-      return (ks, vs, mask, index)
-
-    # Once we get to the end of the buffer, move the second bin back to make
-    # space for new data: [ bin_i bin_{i+1} | ] -> [ bin_{i+1} | bin_{i+1} ],
-    # where | is where index points at in the buffer.
-    state = jax.lax.cond(
-        pred=(index == 2 * self.bin_length),
-        true_operand=state,
-        true_fun=roll_state,
-        false_operand=state,
-        false_fun=(lambda x: x),
-    )
-    return (output, state)
-
-  def new_params_and_state(self, input_shapes, input_dtype, rng):
-    if self._mode in ('train', 'eval'):
-      return (), ()
-
-    assert self._mode == 'predict'
-    assert self.bin_length is not None, (
-        'For fast inference, TimeBinCausalAttention must be parameterized by '
-        'bin_length.'
-    )
-    params = ()
-    state = _fast_inference_init_state(
-        input_shapes, input_dtype, 2 * self.bin_length
-    )
-    return params, state
-
-
-class LSHCausalAttention(BaseCausalAttention):
-  """Causal attention based on locality-sensitive hashing."""
-
-  def __init__(self, dropout, mode, n_bins=64, n_hashes=1, n_buckets=64,
-               one_rng=False, allow_duplicate_attention=False,
-               attend_across_buckets=False, hard_k=0, factorize_hash=False,
-               rehash_each_round=True, drop_for_hash_rate=0.0):
-    del dropout
-    self._mode = mode
-    super(LSHCausalAttention, self).__init__()
-    assert n_buckets >= n_bins, 'This setting is not recommended: too few bins.'
-    assert rehash_each_round or allow_duplicate_attention, (
-        'The setting {allow_duplicate_attention=False, rehash_each_round=False}'
-        ' is not implemented.')
-    self.n_bins = n_bins
-    self.n_hashes = n_hashes
-    self.n_buckets = n_buckets
-    self._drop_for_hash_rate = drop_for_hash_rate
-    self._one_rng = one_rng
-    self._factorize_hash = factorize_hash
-    self._prng = None
-    if one_rng:
-      seed = random.randint(0, 2**31 - 1)
-      self._prng = backend.random.get_prng(seed)
-
-    self._allow_duplicate_attention = allow_duplicate_attention
-    self._attend_across_buckets = attend_across_buckets
-    self._hard_k = hard_k
-    self._rehash_each_round = rehash_each_round
-
-  def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
-    del params, kwargs
-    output, _ = self.batch_call_and_or_grad(inputs[0], inputs[2], rng=rng)
-    return output, state
-
-  def forward_and_backward(self, inputs, ct, rng=None, **kwargs):
-    del kwargs
-    output, (qk_ct, v_ct) = self.batch_call_and_or_grad(
-        inputs[0], inputs[2], ct=ct, rng=rng)
-    return output, (qk_ct, np.zeros_like(inputs[1]), v_ct)
-
-  def has_backward(self):
-    return True
-
-  def backward(self, inputs, output, ct, params=(), state=(), rng=None,
-               **kwargs):
-    del output, params, state
-    _, (qk_ct, v_ct) = self.batch_call_and_or_grad(
-        inputs[0], inputs[2], return_output=False, ct=ct, rng=rng)
-    inputs_ct = (qk_ct, np.zeros_like(inputs[1]), v_ct)
-    return inputs_ct, ()
-
-  def batch_call_and_or_grad(self, qk, v, ct=None, return_output=True,
-                             rng=None):
-    assert return_output or ct is not None, 'No work to perform!'
-    # pylint: disable=protected-access
-    stash_buckets = (return_output and ct is None
-                     and base.Layer._STASH_IN is not None)
-    if return_output and ct is not None and base.Layer._STASH_OUT is not None:
-      buckets = base.Layer._STASH_OUT.pop(self)
-    else:
-      buckets = None
-    # pylint: enable=protected-access
-
-    # The approach here is to perform attention for one batch element and head
-    # at a time. Note that there is absolutely no interaction across examples or
-    # heads: this layer has no parameters, and hashing patterns are also
-    # different across examples/heads. As a result, batching doesn't give any
-    # performance gains except in the case of accelerator under-utilization. We
-    # assume that hash-based attention will be applied primarily to long
-    # sequences, where unbatched attention for a single head has sufficient
-    # computation to fill up the accelerator.
-
-    batch_loop_idx = np.zeros((), dtype=np.int32)
-    batch_loop_max = qk.shape[0]
-
-    init_vals = (batch_loop_idx,)
-    if return_output:
-      out_accum = np.zeros_like(qk)
-      init_vals = init_vals + (out_accum,)
-    if stash_buckets:
-      buckets_accum = np.zeros(
-          [qk.shape[0], self.n_hashes * qk.shape[1]], dtype=np.int32)
-      init_vals = init_vals + (buckets_accum,)
-    if ct is not None:
-      qk_ct_accum = np.zeros_like(qk)
-      v_ct_accum = np.zeros_like(v)
-      init_vals = init_vals + (qk_ct_accum, v_ct_accum)
-
-    def cond_fun(vals):
-      batch_loop_idx = vals[0]
-      return jax.lax.lt(batch_loop_idx, batch_loop_max)
-
-    def body_fun(vals):
-      """Performs attention for a single batch element and head."""
-      batch_loop_idx = vals[0]
-      if self._prng is None:
-        hash_rng = jax.random.fold_in(rng, batch_loop_idx)
-      else:
-        # TODO(kitaev): Maybe use the same RNG across examples (but not heads)?
-        hash_rng = jax.random.fold_in(self._prng, batch_loop_idx)
-      qk_slice = jax.lax.dynamic_index_in_dim(
-          qk, batch_loop_idx, axis=0, keepdims=False)
-      v_slice = jax.lax.dynamic_index_in_dim(
-          v, batch_loop_idx, axis=0, keepdims=False)
-
-      if buckets is None:
-        buckets_slice = self.hash_vectors(qk_slice, rng=hash_rng)
-      else:
-        buckets_slice = jax.lax.dynamic_index_in_dim(
-            buckets, batch_loop_idx, axis=0, keepdims=False)
-
-      if ct is None:
-        out_slice = self.single_call(
-            qk_slice, v_slice, buckets_slice, hash_rng=hash_rng)
-      else:
-        def _do_single_call(qk_slice, v_slice):
-          return self.single_call(
-              qk_slice, v_slice, buckets_slice, hash_rng=hash_rng)
-        ct_slice = jax.lax.dynamic_index_in_dim(
-            ct, batch_loop_idx, axis=0, keepdims=False)
-        out_slice, vjpfun = jax.vjp(_do_single_call, qk_slice, v_slice)
-        qk_ct_slice, v_ct_slice = vjpfun(ct_slice)
-
-      new_vals = (batch_loop_idx + 1,)
-      if return_output:
-        out_accum = vals[1]
-        out_accum = jax.lax.dynamic_update_index_in_dim(
-            out_accum, out_slice, batch_loop_idx, axis=0)
-        new_vals = new_vals + (out_accum,)
-      if stash_buckets:
-        buckets_accum = vals[2]
-        buckets_accum = jax.lax.dynamic_update_index_in_dim(
-            buckets_accum, buckets_slice, batch_loop_idx, axis=0)
-        new_vals = new_vals + (buckets_accum,)
-      if ct is not None:
-        qk_ct_accum, v_ct_accum = vals[-2:]
-        qk_ct_accum = jax.lax.dynamic_update_index_in_dim(
-            qk_ct_accum, qk_ct_slice, batch_loop_idx, axis=0)
-        v_ct_accum = jax.lax.dynamic_update_index_in_dim(
-            v_ct_accum, v_ct_slice, batch_loop_idx, axis=0)
-        new_vals = new_vals + (qk_ct_accum, v_ct_accum)
-
-      return new_vals
-
-    final_vals = jax.lax.while_loop(cond_fun, body_fun, init_vals)
-
-    if return_output:
-      out = final_vals[1]
-    else:
-      out = None
-
-    if stash_buckets:
-      base.Layer._STASH_IN[self] = final_vals[2]  # pylint: disable=protected-access
-
-    if ct is not None:
-      input_ct = final_vals[-2:]
-    else:
-      input_ct = None
-
-    return out, input_ct
-
-  def make_unit_length(self, x, epsilon=1e-6):
-    variance = np.mean(x**2, axis=-1, keepdims=True)
-    norm_inputs = x / np.sqrt(variance + epsilon)
-    return norm_inputs
-
-  def drop_for_hash(self, x, rng):
-    rate = self._drop_for_hash_rate
-    if self._mode == 'train' and rate > 0.0:
-      keep = backend.random.bernoulli(rng, 1.0 - rate, x.shape)
-      return np.where(keep, x / (1.0 - rate), np.zeros_like(x))
-    return x
-
-  def hash_vectors(self, vecs, rng):
-    # See https://arxiv.org/pdf/1509.02897.pdf
-    # We sample a different random rotation for each round of hashing to
-    # decrease the probability of hash misses.
-    assert self.n_buckets % 2 == 0
-
-    # If we factorize the hash, find a factor dividing n_buckets nicely.
-    rot_size, factor_list = self.n_buckets, [self.n_buckets]
-    if self._factorize_hash:
-      # If we are given a list of factors, verify it and use later.
-      if isinstance(self._factorize_hash, list):
-        rot_size, product = 0, 1
-        factor_list = self._factorize_hash
-        for factor in factor_list:
-          assert factor % 2 == 0
-          product *= factor
-          rot_size += factor
-        assert product == self.n_buckets
-      else:  # Find one factor if just set to True.
-        # We want to represent self.n_buckets = factor * rest so that
-        # (1) both factor and rest are even, and (2) factor + rest is minimal.
-        # To compute this we start from factor = sqrt(n_buckets) and go down
-        # with it until we find one that satisfies the constraints above.
-        factor = int(math.sqrt(self.n_buckets))
-        while factor > 0 and not (
-            self.n_buckets % factor == 0 and
-            factor % 2 == 0 and
-            (self.n_buckets // factor) % 2 == 0):
-          factor -= 1
-        if factor > 2:  # Factor of 2 does not warrant the effort.
-          rot_size = factor + (self.n_buckets // factor)
-          factor_list = [factor, self.n_buckets // factor]
-
-    random_rotations_shape = (
-        vecs.shape[-1],
-        self.n_hashes if self._rehash_each_round else 1,
-        rot_size // 2)
-
-    rng = jax.lax.tie_in(vecs, rng)
-    rng, subrng = backend.random.split(rng)
-    random_rotations = jax.random.normal(
-        rng, random_rotations_shape).astype('float32')
-    # TODO(lukaszkaiser): the dropout mask will be used for all rounds of
-    # hashing, so it's shared between them. Check if that's what we want.
-    dropped_vecs = self.drop_for_hash(vecs, subrng)
-    rotated_vecs = np.einsum('tf,fhb->htb', dropped_vecs, random_rotations)
-
-    if self._rehash_each_round:
-      if self._factorize_hash and len(factor_list) > 1:
-        # We factorized self.n_buckets as the product of factor_list.
-        # Get the buckets for them and combine.
-        buckets, cur_sum, cur_product = None, 0, 1
-        for factor in factor_list:
-          rv = rotated_vecs[..., cur_sum:cur_sum + (factor // 2)]
-          cur_sum += factor // 2
-          rv = np.concatenate([rv, -rv], axis=-1)
-          if buckets is None:
-            buckets = np.argmax(rv, axis=-1)
-          else:
-            buckets += cur_product * np.argmax(rv, axis=-1)
-          cur_product *= factor
-      else:
-        rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)
-        buckets = np.argmax(rotated_vecs, axis=-1)
-      # buckets is now (self.n_hashes, seqlen). Next we add offsets so that
-      # bucket numbers from different hashing rounds don't overlap.
-      offsets = jax.lax.tie_in(buckets, np.arange(self.n_hashes))
-      offsets = np.reshape(offsets * self.n_buckets, (-1, 1))
-      buckets = np.reshape(buckets + offsets, (-1,))
-    else:
-      assert not self._factorize_hash
-      rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)
-      # In this configuration, we map each item to the top self.n_hashes buckets
-      rotated_vecs = np.squeeze(rotated_vecs, 0)
-      bucket_range = jax.lax.tie_in(vecs, np.arange(rotated_vecs.shape[-1]))
-      bucket_range = np.reshape(bucket_range, (1, -1))
-      bucket_range = np.broadcast_to(bucket_range, rotated_vecs.shape)
-
-      _, buckets = jax.lax.sort_key_val(
-          rotated_vecs, bucket_range, dimension=-1)
-      buckets = buckets[:, -self.n_hashes:]
-      buckets = np.reshape(np.moveaxis(buckets, 0, -1), (-1,))
-
-    return buckets
-
-  def single_call(self, qk, v, buckets, hash_rng=None):
-    # We use the same vector as both a query and a key.
-    seqlen = qk.shape[-2]
-    assert int(buckets.shape[0]) == self.n_hashes * seqlen
-
-    ticker = jax.lax.tie_in(qk, np.arange(self.n_hashes * seqlen))
-    buckets_and_t = seqlen * buckets + (ticker % seqlen)
-    buckets_and_t = jax.lax.stop_gradient(buckets_and_t)
-
-    # Hash-based sort ("s" at the start of variable names means "sorted")
-    sbuckets_and_t, sticker = jax.lax.sort_key_val(
-        buckets_and_t, ticker, dimension=-1)
-    _, undo_sort = jax.lax.sort_key_val(sticker, ticker, dimension=-1)
-    sbuckets_and_t = jax.lax.stop_gradient(sbuckets_and_t)
-    sticker = jax.lax.stop_gradient(sticker)
-    undo_sort = jax.lax.stop_gradient(undo_sort)
-
-    st = (sticker % seqlen)
-    sqk = np.take(qk, st, axis=0)
-    sv = np.take(v, st, axis=0)
-
-    # Split off a "bin" axis so that attention only occurs within chunks.
-    bq_t = bkv_t = np.reshape(st, (self.n_hashes * self.n_bins, -1))
-    bqk = np.reshape(sqk, (self.n_hashes * self.n_bins, -1, sqk.shape[-1]))
-    bv = np.reshape(sv, (self.n_hashes * self.n_bins, -1, sv.shape[-1]))
-    bq_buckets = bkv_buckets = np.reshape(
-        sbuckets_and_t // seqlen, (self.n_hashes * self.n_bins, -1))
-
-    # Hashing operates on unit-length vectors. Unnormalized query vectors are
-    # fine because they effectively provide a learnable temperature for the
-    # attention softmax, but normalizing keys is needed so that similarity for
-    # the purposes of attention correctly corresponds to hash locality.
-    bq = bqk
-    bk = self.make_unit_length(bqk)
-
-    # Allow each chunk to attend within itself, and also one chunk back. Chunk
-    # boundaries might occur in the middle of a sequence of items from the
-    # same bucket, so this increases the chances of attending to relevant items.
-    # TODO(kitaev): benchmark whether XLA pad operation is noticeably faster.
-    def look_one_back(x):
-      if len(x.shape) == 2:
-        x_extra = np.concatenate([x[-1:, :], x[:-1, :]], axis=0)
-      else:
-        x_extra = np.concatenate([x[-1:, :, :], x[:-1, :, :]], axis=0)
-      return np.concatenate([x, x_extra], axis=1)
-
-    bk = look_one_back(bk)
-    bv = look_one_back(bv)
-    bkv_t = look_one_back(bkv_t)
-    bkv_buckets = look_one_back(bkv_buckets)
-
-    # Dot-product attention.
-    dots = np.matmul(bq, np.swapaxes(bk, -1, -2)) / np.sqrt(bq.shape[-1])
-
-    # Causal masking
-    mask = jax.lax.convert_element_type(
-        jax.lax.lt(bq_t[:, :, None], bkv_t[:, None, :]),
-        np.float32)
-    dots = dots - 1e9 * mask
-
-    # Mask out attention to self except when no other targets are available.
-    self_mask = jax.lax.convert_element_type(
-        jax.lax.eq(bq_t[:, :, None], bkv_t[:, None, :]),
-        np.float32)
-    dots = dots - 1e5 * self_mask
-
-    # Mask out attention to other hash buckets.
-    if not self._attend_across_buckets:
-      bucket_mask = jax.lax.convert_element_type(
-          jax.lax.ne(bq_buckets[:, :, None], bkv_buckets[:, None, :]),
-          np.float32)
-      dots = dots - 1e7 * bucket_mask
-
-    # Don't double-count query-key pairs across multiple rounds of hashing.
-    # There are two possible strategies here. (1) The default is to count how
-    # many times a query-key pair is repeated, and to lower its log-prob
-    # correspondingly at each repetition. (2) When hard_k is set, the code
-    # instead masks all but the first occurence of each query-key pair.
-    # TODO(kitaev): is one strategy faster or more numerically stable?
-    if not self._allow_duplicate_attention:
-      locs1 = undo_sort // bq_t.shape[-1]
-      locs2 = (locs1 + 1) % (self.n_hashes * self.n_bins)
-      if not self._attend_across_buckets:
-        locs1 = buckets * (self.n_hashes * self.n_bins) + locs1
-        locs2 = buckets * (self.n_hashes * self.n_bins) + locs2
-      locs = np.moveaxis(np.concatenate([
-          np.reshape(locs1, (self.n_hashes, seqlen)),
-          np.reshape(locs2, (self.n_hashes, seqlen)),
-      ], 0), 0, -1)  # produces shape (seqlen, 2 * self.n_hashes)
-      slocs = np.take(locs, st, axis=0)
-      b_locs = np.reshape(
-          slocs, (self.n_hashes * self.n_bins, -1, 2 * self.n_hashes))
-      # Queries always use the primary location (based on locs1).
-      b_locs1 = b_locs[:, :, None, :self.n_hashes]
-      if self._hard_k > 0:
-        range_n_hashes = jax.lax.tie_in(b_locs, np.arange(self.n_hashes))
-        nouse_locs = (range_n_hashes[:, None] > range_n_hashes[None, :])
-        nouse_locs = 2 * nouse_locs - 1  # 1 = use, -1 = don't use
-        nouse_locs = np.reshape(
-            np.broadcast_to(nouse_locs[:, None, :],
-                            (self.n_hashes, self.n_bins, self.n_hashes)),
-            (self.n_hashes * self.n_bins, 1, 1, self.n_hashes))
-        b_locs1 = b_locs1 * nouse_locs
-      bq_locs = np.broadcast_to(
-          b_locs1,
-          b_locs.shape[:2] + (2, self.n_hashes))
-      bq_locs = np.reshape(bq_locs, b_locs.shape)
-      bkv_locs = look_one_back(b_locs)
-
-      dup_counts = np.sum(
-          jax.lax.convert_element_type(
-              jax.lax.eq(bq_locs[:, :, None, :], bkv_locs[:, None, :, :]),
-              np.float32),
-          axis=-1)
-      assert dup_counts.shape == dots.shape
-      if self._hard_k > 0:
-        dots = dots - 1e7 * jax.lax.stop_gradient(dup_counts)
-      else:
-        dots = dots - jax.lax.stop_gradient(np.log(dup_counts + 1e-9))
-
-    # Each query only attends to the top k most relevant keys.
-    if self._hard_k > 0:
-      b_top_dots = np.sort(dots)[..., -self._hard_k:]  # Get the top k dots.
-      b_top_dots = jax.lax.stop_gradient(b_top_dots)
-      s_top_dots = np.reshape(b_top_dots, (-1, self._hard_k))
-      top_dots = np.take(s_top_dots, undo_sort, axis=0)
-
-      merged_top_dots = np.moveaxis(
-          np.reshape(top_dots, (self.n_hashes, seqlen, self._hard_k)), 0, -1)
-      merged_top_dots = np.reshape(merged_top_dots, (seqlen, -1))
-
-      dots_thresh = np.sort(merged_top_dots)[:, -self._hard_k]
-      # It's possible to compute the partition function at this point, but right
-      # now this codepath isn't set up for backprop, and there might also be
-      # issues computing it this way if two dot-products are exactly equal.
-
-      sdots_thresh = dots_thresh[st]
-      bdots_thresh = np.reshape(sdots_thresh, (self.n_hashes * self.n_bins, -1))
-      bdots_thresh = jax.lax.stop_gradient(bdots_thresh)
-
-      top_k_mask = jax.lax.convert_element_type(
-          dots < bdots_thresh[..., None], np.float32)
-      dots = dots - 1e7 * jax.lax.stop_gradient(top_k_mask)
-
-    # Softmax.
-    dots_logsumexp = backend.logsumexp(dots, axis=-1, keepdims=True)
-    dots = np.exp(dots - dots_logsumexp)
-
-    bo = np.matmul(dots, bv)
-    so = np.reshape(bo, (-1, bo.shape[-1]))
-    slogits = np.reshape(dots_logsumexp, (-1,))
-
-    def unsort_for_output_impl(so, slogits):
-      o = np.take(so, undo_sort, axis=0)
-      # Sorting is considerably faster than gather, but first we need to get the
-      # XLA compiler to abandon the idea of fusing this sort with the input sort
-      # (which introduces a computation cycle and leads to a crash).
-      # TODO(kitaev): remove "sticker_" variable if XLA is fixed.
-      sticker_ = sticker + jax.lax.convert_element_type(
-          slogits[0] > 0, sticker.dtype)
-      _, logits = jax.lax.sort_key_val(sticker_, slogits, dimension=-1)
-      return o, logits
-
-    def unsort_for_output_vjp(so, slogits):
-      """Custom gradient for unsort_for_output."""
-      so = jax.lax.stop_gradient(so)
-      slogits = jax.lax.stop_gradient(slogits)
-      o, logits = unsort_for_output_impl(so, slogits)
-      def vjpfun(o_logits_grads):
-        so_grad = np.take(o_logits_grads[0], sticker, axis=0)
-        # TODO(kitaev): this exists to match the forward pass, but I'm not sure
-        # if it's actually required.
-        buckets_and_t_ = buckets_and_t + jax.lax.convert_element_type(
-            o_logits_grads[1][0] > 0, buckets_and_t.dtype)
-        _, slogits_grad = jax.lax.sort_key_val(
-            buckets_and_t_, o_logits_grads[1], dimension=-1)
-        return (so_grad, slogits_grad)
-      return (o, logits), vjpfun
-
-    unsort_for_output = jax.custom_transforms(unsort_for_output_impl)
-    jax.defvjp_all(unsort_for_output, unsort_for_output_vjp)
-    o, logits = unsort_for_output_impl(so, slogits)
-
-    if self.n_hashes == 1:
-      out = o
-    else:
-      o = np.reshape(o, (self.n_hashes, seqlen, o.shape[-1]))
-      logits = np.reshape(logits, (self.n_hashes, seqlen, 1))
-      probs = np.exp(logits - backend.logsumexp(logits, axis=0, keepdims=True))
-      out = np.sum(o * probs, axis=0)
-
-    assert out.shape == v.shape
-    return out
-
-
-def CausalAttention(d_feature, n_heads=1,
-                    d_attention_key=None, d_attention_value=None,
-                    attention_type=DotProductCausalAttention,
-                    share_qk=False, mode='train'):
-  """Transformer-style multi-headed causal attention.
-
-  Args:
-    d_feature: int:  dimensionality of feature embedding
-    n_heads: int: number of attention heads
-    d_attention_key: int: depth of key vector for each attention head
-        (default is d_feature // n_heads)
-    d_attention_value: int: depth of value vector for each attention head
-        (default is d_feature // n_heads)
-    attention_type: subclass of BaseCausalAttention: attention class to use
-    share_qk: bool, whether to share queries and keys
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    Multi-headed self-attention result.
-  """
-  if d_attention_key is None:
-    assert d_feature % n_heads == 0
-    d_attention_key = d_feature // n_heads
-  if d_attention_value is None:
-    assert d_feature % n_heads == 0
-    d_attention_value = d_feature // n_heads
-
-  if share_qk:
-    pre_attention = [
-        cb.Dup(),
-        cb.Parallel(
-            ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
-            ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
-        ),
-        cb.Dup(),
-    ]
-  else:
-    pre_attention = [
-        cb.Dup(), cb.Dup(),
-        cb.Parallel(
-            ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
-            ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
-            ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
-        ),
-    ]
-
-  return pre_attention + [
-      attention_type(mode=mode),
-      ComputeAttentionOutput(n_heads=n_heads, d_model=d_feature),
-  ]
diff --git a/tensor2tensor/trax/layers/attention_test.py b/tensor2tensor/trax/layers/attention_test.py
deleted file mode 100644
index f2ca0f73e..000000000
--- a/tensor2tensor/trax/layers/attention_test.py
+++ /dev/null
@@ -1,106 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.layers.attention."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as onp
-from tensor2tensor.trax.layers import attention
-from tensor2tensor.trax.layers import base
-from tensorflow import test
-
-
-class AttentionTest(test.TestCase):
-
-  def test_shift_right(self):
-    # Test shifts right on axis=1
-    layer = attention.ShiftRight()
-    input_np = onp.arange(2*3*3).reshape(2, 3, 3)
-    output_np = layer(input_np)
-    self.assertEqual(input_np.shape, output_np.shape)
-    self.assertAllEqual(onp.array([[[0, 0, 0],
-                                    [0, 1, 2],
-                                    [3, 4, 5]],
-
-                                   [[0, 0, 0],
-                                    [9, 10, 11],
-                                    [12, 13, 14]]]),
-                        output_np)
-
-  def test_shift_right_float(self):
-    layer = attention.ShiftRight()
-    input_np = onp.arange(2*3*3).reshape(2, 3, 3).astype(onp.float32)
-    # Test on a float array.
-    input_np = input_np.astype(onp.float32)
-    input_np /= 2.0
-    self.assertEqual(input_np.dtype, onp.float32)
-
-    output_np = layer(input_np)
-    self.assertEqual(input_np.shape, output_np.shape)
-    self.assertEqual(output_np.dtype, onp.float32)
-
-    self.assertAllEqual(onp.array([[[0., 0., 0.],
-                                    [0., 0.5, 1.],
-                                    [1.5, 2., 2.5]],
-
-                                   [[0., 0., 0.],
-                                    [4.5, 5., 5.5],
-                                    [6., 6.5, 7.]]]),
-                        output_np)
-
-  def test_merged_hashed_causal_attention(self):
-    qkv_shape = (3, 32, 8)
-    input_shape = (qkv_shape, qkv_shape, qkv_shape)
-    layer = attention.MemoryEfficientCausalAttention(
-        loop_stride=16, dropout=0.1, mode='train')
-    final_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual((3, 32, 8), final_shape)
-
-  def test_time_bin_causal_attention_bin_length(self):
-    qkv_shape = (3, 57, 8)
-    input_shape = (qkv_shape, qkv_shape, qkv_shape)
-    layer = attention.TimeBinCausalAttention(
-        bin_length=16, dropout=0.1, mode='train')
-    final_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual((3, 57, 8), final_shape)
-
-  def test_time_bin_causal_attention_n_bins(self):
-    qkv_shape = (3, 57, 8)
-    input_shape = (qkv_shape, qkv_shape, qkv_shape)
-    layer = attention.TimeBinCausalAttention(
-        n_bins=4, dropout=0.1, mode='train')
-    final_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual((3, 57, 8), final_shape)
-
-  def test_time_bin_and_dot_product_causal_attention_are_consistent(self):
-    dot_product_layer = attention.DotProductCausalAttention(
-        dropout=0.0, mode='train')
-    time_bin_layer = attention.TimeBinCausalAttention(
-        bin_length=4, dropout=0.0, mode='train')
-
-    # Exactly 2 bins.
-    input_shape = (3, 8, 8)
-    inputs = [onp.random.uniform(size=input_shape) for _ in range(3)]
-
-    dot_product_output = dot_product_layer(inputs)
-    time_bin_output = time_bin_layer(inputs)
-    onp.testing.assert_array_almost_equal(dot_product_output, time_bin_output)
-
-
-if __name__ == '__main__':
-  test.main()
diff --git a/tensor2tensor/trax/layers/base.py b/tensor2tensor/trax/layers/base.py
deleted file mode 100644
index 658966cfb..000000000
--- a/tensor2tensor/trax/layers/base.py
+++ /dev/null
@@ -1,664 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 layer class."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import inspect
-import traceback
-
-import jax
-
-import numpy as onp
-from tensor2tensor.trax import backend
-from tensor2tensor.trax.backend import nested_map
-from tensor2tensor.trax.backend import ShapeType
-
-
-class Layer(object):
-  """Base class for composable layers in a deep learning network.
-
-  Layers are the basic building blocks for deep learning models. A Trax layer
-  computes a function from zero or more inputs to zero or more outputs,
-  optionally using trainable parameters (common) and non-parameter state (not
-  common). Authors of new layer subclasses typically override at most two
-  methods of the base `Layer` class:
-
-    forward(inputs, params=(), state=(), **kwargs):
-      Computes this layer's output as part of a forward pass through the model.
-
-    new_params_and_state(self, input_shape, input_dtype, rng):
-      Returns a (params, state) pair suitable for initializing this layer.
-
-  A small subset of layer types are combinators -- they organize the computation
-  of their sublayers, e.g., applying their sublayers in series or in parallel.
-
-  All layers have the following properties, with default values implemented
-  in the base `Layer` class:
-
-    - n_inputs: int (default 1)
-    - n_outputs: int (default 1)
-    - params: tuple (default empty -- the layer has no parameters)
-    - state: tuple (default empty -- the layer has no non-parameter state)
-    - sublayers: tuple (default empty -- the layer has no sublayers)
-
-  The inputs to a layer are tensors, packaged according to how many there are:
-
-    - n_inputs = 0: an empty tuple ()
-    - n_inputs = 1: one tensor (NOT wrapped in a tuple)
-    - n_inputs > 1: a tuple of tensors
-
-  (The special treatment of the single-input case is meant to simplify the
-  work of layer writers; this design choice may be revisited in the future.)
-
-  The outputs from a layer are also tensors, packaged the same as layer inputs:
-
-    - n_outputs = 0: an empty tuple ()
-    - n_outputs = 1: the tensor (NOT wrapped in a tuple)
-    - n_outputs > 1: a tuple of tensors
-
-  The Trax runtime maintains a data stack with which layer calls are composed.
-  For more complex data network architectures, possibly involving multiple data
-  flows, one can view each layer as a function from stack state to stack state,
-  where the function's inputs are a slice from the stack, and the function's
-  outputs are spliced back into the stack.
-  """
-
-  def __init__(self, n_inputs=1, n_outputs=1):
-    """Creates a partially initialized, unconnected layer instance.
-
-    Args:
-      n_inputs: Number of inputs expected by this layer.
-      n_outputs: Number of outputs promised by this layer.
-    """
-    self._n_inputs = n_inputs
-    self._n_outputs = n_outputs
-    self._sublayers = ()  # Default is no sublayers.
-    self._params = ()  # cached parameters
-    self._state = ()
-    self._caller = _find_frame(inspect.stack())  # for custom error messages
-    self._init_finished = False
-
-  def __repr__(self):
-    class_str = self.__class__.__name__
-    fields_str = 'in={},out={}'.format(self.n_inputs, self.n_outputs)
-    objs = self.sublayers
-    if objs:
-      objs_str = ', '.join(str(x) for x in objs)
-      return '{}{{{},sublayers=[{}]}}'.format(class_str, fields_str, objs_str)
-    else:
-      return '{}{{{}}}'.format(class_str, fields_str)
-
-  def forward(self, inputs, params=(), state=(), **kwargs):
-    """Computes this layer's output as part of a forward pass through the model.
-
-    Authors of new Layer subclasses should override this method to define the
-    forward computation that their layer performs.
-
-    Args:
-      inputs: Input tensors, matching the number (n_inputs) expected by this
-          layer. Specifically:
-            - n_inputs = 0: an empty tuple ()
-            - n_inputs = 1: a tensor (NOT wrapped in a tuple)
-            - n_inputs > 1: a tuple of tensors, with n_inputs items
-      params: A tuple of trainable parameters, with one element for this layer
-          if this layer has no sublayers, or one for each sublayer if this
-          layer has sublayers. If a layer (or sublayer) has no trainable
-          parameters, the corresponding params element is an empty tuple.
-      state: Layer-specific non-parameter state that can update between batches.
-      **kwargs: Often empty; main current use is to carry a PRNG key for random
-          number generation, using the keyword 'rng'.
-
-    Returns:
-      Tensors, matching the number (n_outputs) promised by this layer.
-      Specifically:
-        - n_outputs = 0: an empty tuple
-        - n_outputs = 1: one tensor (NOT wrapped in a tuple)
-        - n_outputs > 1: a tuple of tensors, with n_outputs items
-    """
-    raise NotImplementedError
-
-  def new_params_and_state(self, input_shape, input_dtype, rng):
-    """Returns a (params, state) pair suitable for initializing this layer.
-
-    Authors of new Layer subclasses should override this method if their layer
-    uses trainable parameters or has non-parameter state that gets updated
-    between batches. The default implementation works for layers that have
-    no parameters or state.
-
-    Args:
-      input_shape: A tuple representing a shape (if this layer takes one input)
-          or a tuple of shapes (if this layer takes more than one input).
-          For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
-      input_dtype: Numpy dtype(s) for each of the inputs.
-      rng: A PRNG key for random number generation.
-    """
-    del input_shape, input_dtype, rng
-    return (), ()
-
-  @property
-  def n_inputs(self):
-    """Returns how many tensors this layer expects as input."""
-    return self._n_inputs
-
-  @property
-  def n_outputs(self):
-    """Returns how many tensors this layer promises as output."""
-    return self._n_outputs
-
-  @property
-  def sublayers(self):
-    """Returns a tuple containing this layer's sublayers; may be empty."""
-    return self._sublayers
-
-  @property
-  def params(self):
-    """Returns a tuple containing this layer's parameters; may be empty."""
-    return self._params
-
-  @params.setter
-  def params(self, params):
-    self._params = params
-
-  @property
-  def state(self):
-    """Returns a tuple containing this layer's state; may be empty."""
-    return self._state
-
-  @state.setter
-  def state(self, state):
-    self._state = state
-
-  @property
-  def has_backward(self):
-    """Returns True if this layer provides its own (custom) backward pass code.
-
-    A layer subclass that provides custom backward pass code (for custom
-    gradients) must override this method to return True.
-    """
-    return False
-
-  def backward(self, inputs, output, grad, params, state, **kwargs):
-    """Custom backward pass to propagate gradients in a custom way.
-
-    Args:
-      inputs: Input tensors; can be a (possibly nested) tuple.
-      output: The result of running this layer on inputs.
-      grad: gradient signal (called cotangent in jax) computed based on
-        subsequent layers. The structure and shape must match output.
-      params: layer parameters
-      state: start state.
-      **kwargs: kwargs for the layer
-
-    Returns:
-      The custom gradient signal for the input. Note that we need to return
-      a gradient for each argument of forward, so it will usually be a tuple
-      of signals: the gradient for inputs and parameters.
-    """
-    raise NotImplementedError
-
-  # End of subclassing interface, all functions below are internal.
-
-  def pseudo_forward(self, pseudo_inputs, params, state):
-    """Computes shapes and types this layer would produce for the given inputs.
-
-    Args:
-      pseudo_inputs: A ShapeType instance (input data minus the actual values)
-          or a tuple of ShapeType instances, following the same conventions as
-          Layer.forward's input arg.
-      params: Parameters for this layer.
-      state: start state.
-
-    Returns:
-      A tuple of (output, state).
-
-      The output part of the tuple is a ShapeType instance representing the
-      shape and type of the output (if this layer has one output) or a tuple
-      of ShapeType instances (if this layer has more than one output).
-    """
-    try:
-      # Beware: using an actual RNG (as opposed to this ShapeType stub) would
-      # cause a large number of dropout masks to be computed and permanently
-      # stored in global memory.
-      rng = ShapeType(shape=(2,), dtype=onp.uint32)
-      def call_on_input(x, params, state, rng):
-        return self.forward(x, params=params, state=state, rng=rng)
-      params_shapes = nested_map(
-          params, lambda x: ShapeType(shape=x.shape, dtype=x.dtype))
-      s = backend.eval_on_shapes(call_on_input)(pseudo_inputs,
-                                                params_shapes, state, rng)
-      return s
-    except Exception:
-      name, trace = self.__class__.__name__, _short_traceback(skip=3)
-      raise LayerError(name, 'pseudo_forward', self._caller, pseudo_inputs,
-                       None, trace)
-
-  def initialize_once(self, input_shapes, input_dtype, rng):
-    """Initializes this layer and its sublayers recursively.
-
-    This method is designed to initialize each layer instance once, even if the
-    same layer instance occurs in multiple places in the network. This enables
-    weight sharing to be implemented as layer sharing.
-
-    Args:
-      input_shapes: A tuple representing a shape (if this layer takes one input)
-          or a tuple of shapes (if this layer takes more than one input).
-          For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
-      input_dtype: Numpy dtype(s) for each of the inputs.
-      rng: A PRNG key for random number generation.
-
-    Returns:
-      A (params, state) tuple, in which params contains newly created parameters
-          on the first call and () on all subsequent calls.
-    """
-    try:
-      # Initialize params once; store them for use when this layer is called.
-      # Needs to call new_params_and_state regardless of _init_finished because
-      # state also needs to be initialized. After jitting, graph pruning should
-      # be able to remove unnecessary computation.
-      # TODO(lukaszkaiser): Revisit this decision and see whether layers sharing
-      #   params should also share states.
-      params, state = self.new_params_and_state(input_shapes, input_dtype, rng)
-      if not self._init_finished:
-        self._init_finished = True
-        self._params = params
-        self._state = state
-      else:
-        params = ()
-      return (params, state)
-    except Exception:
-      name, trace = self.__class__.__name__, _short_traceback(skip=3)
-      raise LayerError(name, 'initialize_once', self._caller, input_shapes,
-                       input_dtype, trace)
-
-  # XXX(kitaev):
-  _STASH_IN = None
-  _STASH_OUT = None
-
-  def __call__(self, x, **kwargs):
-    """Makes Layer instances callable; for use in tests or interactive settings.
-
-    This convenience method helps library users play with, test, or otherwise
-    probe the behavior of layers outside of a full training environment. It
-    presents the layer as callable function from inputs to outputs, with the
-    option of manually specifying parameters and non-parameter state per
-    individual call. For convenience, parameters and non-parameter state are
-    cached per layer instance, starting from default values of () and (), and
-    acquiring non-empty values either by initialization or from values
-    explicitly provided via the params and state keyword arguments.
-
-    Args:
-      x: 0 or more input tensors, formatted the same as the inputs to
-          Layer.forward.
-      **kwargs: Additional keyword arguments if needed/desired for this layer.
-          Three possible keyword arguments are especially relevant:
-            - params=... will override any cached params values
-            - state=... will override any cached state values
-            - rng=... will supply a PRNG key for use by the layer
-
-    Returns:
-      0 or more output tensors, formatted the same as the outputs from
-          Layer.forward.
-    """
-    params = kwargs.pop('params', self.params)
-    state = kwargs.pop('state', self.state)
-    outputs, _ = self.apply_forward(x, params=params, state=state, **kwargs)
-    return outputs
-
-  def apply_forward(self, x, params=(), state=(), **kwargs):
-    """Applies this layer as part of a forward pass; an internal system method.
-
-    This method is reserved for handling plumbing and other internal affairs
-    as needed by the overall library. Trax library users should use or override
-    the `forward` method instead.
-
-    Args:
-      x: See Layer.forward inputs.
-      params: See Layer.forward.
-      state: See Layer.forward.
-      **kwargs: See Layer.forward.
-
-    Returns:
-      See Layer.forward.
-    """
-    try:
-      # If params are nothing, we may be reusing this layer.
-      # Use the cached parameters to calculate the value.
-      # Note: to make sure jit tracers can decide this branch in python we
-      #   use "params is ()" instead of, e.g., "not params" or "params == ()".
-      if params is ():  # pylint: disable=literal-comparison
-        params = self._params
-      else:
-        # In this case, we're called for the first time: cache parameters.
-        self._params = params
-
-      if not self.has_backward or Layer._STASH_IN is not None:
-        outputs, s = self.forward(x, params=params, state=state, **kwargs)
-      else:
-        outputs, s = self._do_custom_gradients(x, params, state, **kwargs)
-      self._state = s
-      return outputs, s
-
-    except Exception:
-      name, trace = self.__class__.__name__, _short_traceback()
-      raise LayerError(name, 'apply_forward', self._caller,
-                       shapes(x), None, trace)
-
-  def _do_custom_gradients(self, x, params, state, **kwargs):
-    """Calls this layer for a forward pass, but with custom gradients."""
-    assert backend.get_name() == 'jax', (
-        'Custom gradients are only supported in JAX for now.')
-
-    # TODO(wangpeng): JAX doesn't support custom grads for functions with
-    #   auxiliary output yet (https://github.com/google/jax/issues/844). Will
-    #   remove the constraints on state below when this feature is added to
-    #   JAX.
-
-    assert not jax.tree_util.tree_leaves(state), (
-        'Custom gradients require trivial start state. Got %s' % str(state))
-
-    def check_end_state(output_state):
-      output, state = output_state
-      assert not jax.tree_util.tree_leaves(state), (
-          'Custom gradients require trivial end state. Got %s' % str(state))
-      return output
-
-    # See this link for how custom transformations are defined in JAX:
-    # https://jax.readthedocs.io/en/latest/jax.html#jax.custom_transforms
-    # Note that we capture the kwargs and don't calculate gradients wrt. them.
-    @jax.custom_transforms
-    def _do_forward(y, params):
-      return check_end_state(self.forward(y, params=params, state=state,
-                                          **kwargs))
-
-    # This is the custom gradient (vector-jacobian product in JAX) function.
-    # For the exact specification of this custom transformation see this link:
-    # https://jax.readthedocs.io/en/latest/jax.html#jax.defjvp_all
-    def do_forward_vjp(y, params):
-      """Custom gradient (vjp) function."""
-      stash = None
-      if Layer._STASH_IN is None:
-        Layer._STASH_IN = stash = {}
-      output = check_end_state(self.forward(y, params=params, state=state,
-                                            **kwargs))
-      if stash is not None:
-        Layer._STASH_IN = None
-      def vjpfun(grad):
-        assert Layer._STASH_OUT is None
-        Layer._STASH_OUT = stash
-        res = self.backward(y, output, grad, params, state, **kwargs)
-        Layer._STASH_OUT = None
-        return res
-      return output, vjpfun
-
-    jax.defvjp_all(_do_forward, do_forward_vjp)
-    return _do_forward(x, params), state
-
-
-class LayerError(Exception):
-  """Exception raised in the layer stack.
-
-  Attributes:
-    message: the message corresponding to this exception.
-  """
-
-  def __init__(self, layer_name, function_name, caller,
-               input_shapes, input_types, traceback_string):
-    self._layer_name = layer_name
-    self._function_name = function_name
-    self._caller = caller  # Python inspect object with init caller info.
-    self._traceback = traceback_string
-    self._input_shapes = input_shapes
-    self._input_types = input_types
-    super(LayerError, self).__init__(self.message)
-
-  @property
-  def message(self):
-    """Create error message."""
-    prefix = 'Exception passing through layer '
-    prefix += '%s (in %s):\n' % (self._layer_name, self._function_name)
-    short_path = '[...]/' + '/'.join(self._caller.filename.split('/')[-3:])
-    caller = '  layer created in file %s, line %d\n' % (short_path,
-                                                        self._caller.lineno)
-    shapes_str = '  layer input shapes: %s\n\n' % str(self._input_shapes)
-    if self._input_types is not None:
-      types_str = '  layer input types: %s\n' % str(self._input_types)
-      shapes_str = types_str + shapes_str
-    return prefix + caller + shapes_str + self._traceback
-
-
-def _apply_to_first_n(f, x, n):
-  """Helper: apply f to first n elements on the stack x if n > 0."""
-  if n < 1:
-    return f(x)
-  argument, rest = x[:n], x[n:]
-  if n == 1:
-    argument = argument[0]
-  result = f(argument)
-  if not rest:
-    return result
-  if n == 1:
-    result = [result]
-  result = list(result) + list(rest)
-  if isinstance(x, tuple):
-    result = tuple(result)
-  return result
-
-
-def nested_reduce(x, f):
-  """Fold the function f to the nested structure x (dicts, tuples, lists)."""
-  if isinstance(x, list):
-    return f([nested_reduce(y, f) for y in x])
-  if isinstance(x, tuple):
-    return f([nested_reduce(y, f) for y in x])
-  return x
-
-
-def shapes(x):
-  """Get a structure of shapes for a structure of nested arrays."""
-  def shape(x):
-    try:
-      return tuple([int(i) for i in x.shape])
-    except Exception:  # pylint: disable=broad-except
-      return []
-  return nested_map(x, shape)
-
-
-def sizes(x):
-  """Get a structure of sizes for a structure of nested arrays."""
-  def size(x):
-    try:
-      return x.size
-    except Exception:  # pylint: disable=broad-except
-      return 0
-  return nested_map(x, size)
-
-
-def _find_frame(stack, start=0):
-  """Find the frame with the caller on the stack."""
-  # We want to find the first place where the layer was called
-  # that is *not* an __init__ function of an inheriting layer.
-  frame = inspect.getframeinfo(stack[start][0])
-  # If we are in an init, move on.
-  if frame.function == '__init__':
-    return _find_frame(stack, start + 1)
-  return frame
-
-
-def _shorten_file_path(line):
-  """Shorten file path in error lines for more readable tracebacks."""
-  start = line.lower().find('file')
-  if start < 0:
-    return line
-  first_quote = line.find('"', start)
-  if first_quote < 0:
-    return line
-  second_quote = line.find('"', first_quote + 1)
-  if second_quote < 0:
-    return line
-  path = line[first_quote + 1:second_quote]
-  new_path = '/'.join(path.split('/')[-3:])
-  return line[:first_quote] + '[...]/' + new_path + line[second_quote + 1:]
-
-
-def _short_traceback(skip=3):
-  """Cleaned-up form of traceback."""
-  counter, res = 0, []
-  # Skipping 3 lines by default: the top (useless) and self-call.
-  lines = traceback.format_exc().splitlines()[skip:]
-  for l in lines:
-    res.append(_shorten_file_path(l))
-    if counter % 2 == 1:
-      res.append('')
-    counter += 1
-    # If we see a LayerError, the traceback has already been processed.
-    if l.startswith('LayerError'):
-      # Skip 4 back except last as these are internal base-layer calls.
-      res = res[:-4] + [res[-1]]
-      res += lines[counter:]
-      break
-  return '\n'.join(res)
-
-
-def _validate_forward_input(x, n_inputs):
-  if n_inputs != 1:
-    if not isinstance(x, tuple):
-      raise TypeError(
-          'expected input to be a tuple; instead received {}'.format(type(x)))
-    if len(x) != n_inputs:
-      raise ValueError(
-          'input tuple length ({}) does not equal required number of inputs'
-          ' ({})'.format(len(x), n_inputs))
-
-
-def layer(n_inputs=1, n_outputs=1, new_params_and_state_fn=None):
-  """Returns a decorator that converts a function into a Layer class builder."""
-
-  def _build_layer_class(raw_fn):
-    """Returns a Layer class whose callable instances execute the function."""
-
-    def _init(self, **kwargs):
-      self._kwargs = kwargs  # pylint: disable=protected-access
-      Layer.__init__(self, n_inputs=n_inputs, n_outputs=n_outputs)
-
-    def _new_params_and_state(self, input_shapes, input_dtype, rng):
-      if new_params_and_state_fn is None:
-        return (), ()
-      kwargs = self._kwargs  # pylint: disable=protected-access
-      return new_params_and_state_fn(input_shapes, input_dtype, rng, **kwargs)
-
-    def _is_empty(raw_output):
-      return raw_output is None or (isinstance(raw_output, (list, tuple))
-                                    and len(raw_output) == 0)  # pylint: disable=g-explicit-length-test
-
-    def _forward(self, x, params=(), state=(), **kwargs):
-      """Uses this layer as part of a forward pass through the model."""
-      merged_kwargs = kwargs.copy()
-      merged_kwargs.update(self._kwargs)  # pylint: disable=protected-access
-
-      _validate_forward_input(x, n_inputs)
-      raw_output = raw_fn(x, params=params, **merged_kwargs)
-      output = () if _is_empty(raw_output) else raw_output
-      return (output, state)
-
-    # Set docstrings and create the class.
-    _forward.__doc__ = raw_fn.__doc__
-    _new_params_and_state.__doc__ = new_params_and_state_fn.__doc__
-    # Note: None.__doc__ is None
-    cls = type(raw_fn.__name__, (Layer,),
-               {'__init__': _init,
-                'forward': _forward,
-                'new_params_and_state': _new_params_and_state})
-    return cls
-
-  return _build_layer_class
-
-
-def _random_values(input_shapes, rng, integer_inputs=False):
-  """Creates random floats or ints of the given shape.
-
-  Args:
-    input_shapes: A tuple representing a shape (if the layer takes one input)
-        or a tuple of shapes (if this layer takes more than one input).
-        For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
-    rng: A random number generator.
-    integer_inputs: If True, use numpy int32 to produce the random data, else
-        use float32.
-
-  Returns:
-    Random values with the shape and type specified.
-  """
-  if isinstance(input_shapes[0], int):
-    # Non-nested shape, create a random tuple.
-    if not integer_inputs:
-      return backend.random.uniform(rng, input_shapes, minval=-1.0, maxval=1.0)
-    return backend.random.bernoulli(rng, 0.5, input_shapes).astype(onp.int32)
-  elif isinstance(input_shapes, tuple):  # Nested shape: tuple.
-    return tuple(_random_values(x, rng, integer_inputs) for x in input_shapes)
-  else:
-    raise TypeError(type(input_shapes))
-
-
-def _is_tuple_of_shapes(shape):
-  # TODO(jonni): Find better way to distinguish a shape from a tuple of shapes.
-  if not isinstance(shape, tuple):
-    raise TypeError('shape must be a tuple or tuple of tuples, instead got:'
-                    ' {}'.format(shape))
-  return isinstance(shape, tuple) and isinstance(shape[0], tuple)
-
-
-def check_shape_agreement(layer_obj, input_shapes, integer_inputs=False):
-  """Checks if the layer's call output agrees its pseudo_forward predictions.
-
-  This function helps test layer mechanics and inter-layer connections that
-  aren't dependent on specific data values.
-
-  Args:
-    layer_obj: A Layer instance.
-    input_shapes: A tuple representing a shape (if the layer takes one input)
-        or a tuple of shapes (if this layer takes more than one input).
-        For example: (210, 160, 3) or ((210, 160, 3), (105, 80, 3)).
-    integer_inputs: If True, use numpy int32 as the type for the pseudo-data,
-        else use float32.
-
-  Returns:
-    A tuple representing either a single shape (if the layer has one output) or
-    a tuple of shape tuples (if the layer has more than one output).
-  """
-  rng1, rng2, rng3 = backend.random.split(backend.random.get_prng(0), 3)
-  input_dtype = onp.int32 if integer_inputs else onp.float32
-  if _is_tuple_of_shapes(input_shapes):
-    pseudo_data = tuple(ShapeType(x, input_dtype) for x in input_shapes)
-    input_dtype = tuple(input_dtype for _ in input_shapes)
-  else:
-    pseudo_data = ShapeType(input_shapes, input_dtype)
-  params, state = layer_obj.initialize_once(input_shapes, input_dtype, rng1)
-  pseudo_output, _ = layer_obj.pseudo_forward(pseudo_data, params, state)
-  if isinstance(pseudo_output, tuple):
-    output_shape = tuple(x.shape for x in pseudo_output)
-  else:
-    output_shape = pseudo_output.shape
-
-  random_input = _random_values(input_shapes, rng2, integer_inputs)
-  real_output = layer_obj(random_input, params=params, state=state, rng=rng3)
-  result_shape = shapes(real_output)
-
-  msg = 'output shape %s != real result shape %s' % (output_shape, result_shape)
-  assert output_shape == result_shape, msg
-  # TODO(jonni): Remove this assert? It makes test logs harder to read.
-  return output_shape
diff --git a/tensor2tensor/trax/layers/base_test.py b/tensor2tensor/trax/layers/base_test.py
deleted file mode 100644
index 66aefee9d..000000000
--- a/tensor2tensor/trax/layers/base_test.py
+++ /dev/null
@@ -1,90 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 base layer."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import absltest
-from tensor2tensor.trax import backend
-from tensor2tensor.trax.layers import base
-
-
-class BaseLayerTest(absltest.TestCase):
-
-  def test_layer_decorator_and_shape_agreement(self):
-    @base.layer()
-    def add_one(x, **unused_kwargs):
-      return x + 1
-
-    output_shape = base.check_shape_agreement(
-        add_one(), (12, 17))  # pylint: disable=no-value-for-parameter
-    self.assertEqual(output_shape, (12, 17))
-
-  def test_custom_zero_grad(self):
-
-    class IdWithZeroGrad(base.Layer):
-
-      def forward(self, x, params=(), state=(), **kwargs):
-        del kwargs
-        return x, ()
-
-      @property
-      def has_backward(self):
-        return True
-
-      def backward(self, inputs, output, ct, params, state, **kwargs):
-        return (backend.numpy.zeros_like(ct), ())
-
-    layer = IdWithZeroGrad()
-    rng = backend.random.get_prng(0)
-    input_shape = (9, 17)
-    random_input = backend.random.uniform(rng, input_shape, minval=-1.0,
-                                          maxval=1.0)
-    layer.initialize_once(input_shape, random_input.dtype, rng)
-    f = lambda x: backend.numpy.mean(layer(x))
-    grad = backend.grad(f)(random_input)
-    self.assertEqual(grad.shape, input_shape)  # Gradient for each input.
-    self.assertEqual(sum(sum(grad * grad)), 0.0)  # Each one is 0.
-
-  def test_custom_id_grad(self):
-
-    class IdWithIdGrad(base.Layer):
-
-      def forward(self, x, params=(), state=(), **kwargs):
-        del kwargs
-        return x, ()
-
-      @property
-      def has_backward(self):
-        return True
-
-      def backward(self, inputs, output, ct, params, state, **kwargs):
-        return (inputs, ())
-
-    layer = IdWithIdGrad()
-    rng = backend.random.get_prng(0)
-    input_shape = (9, 17)
-    random_input = backend.random.uniform(rng, input_shape, minval=-1.0,
-                                          maxval=1.0)
-    layer.initialize_once(input_shape, random_input.dtype, rng)
-    f = lambda x: backend.numpy.mean(layer(x))
-    grad = backend.grad(f)(random_input)
-    self.assertEqual(grad.shape, input_shape)  # Gradient for each input.
-    self.assertEqual(sum(sum(grad)), sum(sum(random_input)))  # Same as input.
-
-if __name__ == "__main__":
-  absltest.main()
diff --git a/tensor2tensor/trax/layers/combinators.py b/tensor2tensor/trax/layers/combinators.py
deleted file mode 100644
index 311ed2278..000000000
--- a/tensor2tensor/trax/layers/combinators.py
+++ /dev/null
@@ -1,539 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Combinators for composing layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.trax import backend
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.layers import base
-
-
-def Model(*layers):
-  """Ensures that a layer or list of layers can be treated as a model.
-
-  Currently, any subclass of base.Layer can be treated as a model.
-
-  Args:
-    *layers: One or more layer objects. In fuller detail, the list may contain
-        nested sublists, and the top-level list can also be a tuple.
-
-  Returns:
-    A single object that treated as a model, e.g., trained or evaluated.
-  """
-  return Serial(*layers)
-
-
-def _deep_flatten(items):  # pylint: disable=invalid-name
-  """Returns a list of objects, flattening sublists/subtuples along the way.
-
-  Example: _deep_flatten([1, (2, 3, (4, 5), [6, 7]), [[[8]]]]) would return
-  the list [1, 2, 3, 4, 5, 6, 7, 8].
-
-  Args:
-    items: An iterable. If elements of this iterable are lists or tuples, they
-        will be (recursively) flattened until non-list non-tuple objects are
-        reached.
-
-  Returns:
-    A list of non-list, non-tuple objects.
-  """
-  def _flat_gen(xs):  # pylint: disable=invalid-name
-    for x in xs:
-      if isinstance(x, (list, tuple)):
-        for y in _flat_gen(x):
-          yield y
-      else:
-        yield x
-  return list(_flat_gen(items))
-
-
-def _ensure_sublayers(layers):  # pylint: disable=invalid-name
-  """Ensures that elements in a layer list are layers.
-
-  Args:
-    layers: A tuple or list whose elements can each be a layer, tuple, or list,
-        and so on recursively.
-
-  Returns:
-    An analogous collection of layers in which embedded layer lists are
-    wrapped in Serial layer instances.
-  """
-  if not layers:  # None or an empty list can signal a no-op.
-    return Serial(None)  # no-op, but still handles shapes and initialization
-  elif isinstance(layers, (list, tuple)):
-    sublayers_not_lists = []
-    for layer in layers:
-      sublayers_not_lists.append(
-          Serial(layer) if isinstance(layer, (list, tuple)) else layer)
-    return sublayers_not_lists
-  else:
-    raise TypeError(type(layers))
-
-
-def _pop_rng_and_split(args_dict, n_copies):  # pylint: disable=invalid-name
-  rng = args_dict.pop('rng', None)
-  if rng is None:
-    return (None,) * n_copies
-  return backend.random.split(rng, n_copies)
-
-
-def _count_items(xs):  # pylint: disable=invalid-name
-  return len(xs) if isinstance(xs, (list, tuple)) else 1
-
-
-class Serial(base.Layer):
-  """Combinator that applies layers serially (by function composition).
-
-  A Serial combinator uses stack semantics to manage data for its sublayers.
-  Each sublayer sees only the inputs it needs and returns only the outputs it
-  has generated. The sublayers interact via the data stack. For instance, a
-  sublayer k, following sublayer j, gets called with the data stack in the
-  state left after layer j has applied. The Serial combinator then:
-
-    - takes N_in items off the top of the stack (N_in = k.n_inputs) and calls
-      layer k, passing those items as arguments; and
-
-    - takes layer k's N_out return values (N_out = k.n_outputs) and pushes
-      them onto the data stack.
-
-  A Serial instance with no sublayers acts as a special-case (but useful)
-  1-input 1-output no-op.
-  """
-
-  def __init__(self, *layers):
-    super(Serial, self).__init__()
-
-    layers = self._ensure_flat(layers)
-    self._sublayers = layers
-    self._n_layers = len(layers)
-
-    if layers:
-      self._n_inputs, self._n_outputs = self._n_inputs_n_outputs(layers)
-
-  def _ensure_flat(self, layers):
-    """Ensures that layers is a single flat list of Layer instances."""
-    del self
-    if len(layers) == 1 and layers[0] is None:
-      layers = ()
-    else:
-      layers = _deep_flatten(layers)
-    for obj in layers:
-      if not isinstance(obj, base.Layer):
-        raise ValueError(
-            'Found nonlayer object ({}) in layers: {}.'.format(obj, layers))
-    return layers
-
-  def _n_inputs_n_outputs(self, layers):
-    del self
-    running_max = 0
-    running_total = 0
-    for layer in layers:
-      running_total += layer.n_inputs
-      running_max = max(running_max, running_total)
-      running_total -= layer.n_outputs
-    return running_max, (running_max - running_total)
-
-  def _validate_forward_inputs(self, xs):
-    if not isinstance(xs, tuple) and self._n_inputs != 1:
-      raise TypeError(
-          'Serial.forward input must be a tuple; instead got {}'.format(xs))
-    len_xs = 1 if isinstance(xs, np.ndarray) else len(xs)
-    if len_xs < self.n_inputs:
-      raise ValueError(
-          'number of inputs ({}) to Serial.forward less than n_inputs'
-          ' ({})'.format(len(xs), self.n_inputs))
-
-  @base.Layer.params.setter
-  def params(self, params):
-    """Recursively sets params on this layer and all sublayers."""
-    self._params = params
-    assert len(params) == self._n_layers
-    for layer, sublayer_params in zip(self.sublayers, params):
-      layer.params = sublayer_params
-
-  @base.Layer.state.setter
-  def state(self, state):
-    """Recursively sets non-param state on this layer and all sublayers."""
-    self._state = state
-    assert len(state) == self._n_layers
-    for layer, sublayer_state in zip(self.sublayers, state):
-      layer.state = sublayer_state
-
-  def forward(self, xs, params=(), state=(), **kwargs):
-    self._validate_forward_inputs(xs)
-    rngs = _pop_rng_and_split(kwargs, self._n_layers)
-    if not self.sublayers:  # No-op: leave args unchanged.
-      return (xs, state)
-
-    stack = xs
-    new_state = []
-    n_layers = self._n_layers
-    if n_layers != 1 and len(params) != n_layers:
-      raise ValueError('number of params ({}) not equal to number of layers '
-                       '({})'.format(len(params), n_layers))
-    if n_layers != 1 and len(state) != n_layers:
-      raise ValueError('length of state ({}) not equal to number of layers '
-                       '({})'.format(len(state), n_layers))
-    for layer, p, s, rng in zip(self.sublayers, params, state, rngs):
-      is_stack_just_one_item = (_count_items(stack) == 1)
-
-      # Give layer its args from the stack; treat 1-arg layer specially.
-      n_in = layer.n_inputs
-      if n_in == 1 and is_stack_just_one_item:
-        inputs = stack
-      elif n_in == 1:
-        inputs = stack[0]
-      else:
-        inputs = stack[:n_in]
-      outputs, s = layer.apply_forward(inputs, params=p, state=s, rng=rng,
-                                       **kwargs)
-      new_state.append(s)
-
-      # Push outputs onto remaining stack (if any).
-      if n_in < _count_items(stack):
-        if layer.n_outputs == 1:
-          outputs = (outputs,)
-        stack = outputs + stack[n_in:]
-      else:
-        stack = outputs  # NOTE: can be single value or tuple.
-
-    return stack, new_state
-
-  def new_params_and_state(self, input_shape, input_dtype, rng):
-    def MakeShapeType(shape, dtype):
-      if isinstance(dtype, (list, tuple)):
-        return tuple(MakeShapeType(s, t) for s, t in zip(shape, dtype))
-      return base.ShapeType(shape=shape, dtype=dtype)
-
-    params = []
-    states = []
-    pseudo_xs = MakeShapeType(input_shape, input_dtype)
-    for layer in self.sublayers:
-      rng, layer_rng = backend.random.split(rng)
-
-      # Give layer its args from pseudo_xs; treat 1-arg layer specially.
-      is_stack_just_one_item = (_count_items(pseudo_xs) == 1)
-      n_in = layer.n_inputs
-      if n_in == 1 and is_stack_just_one_item:
-        inputs = pseudo_xs
-      elif n_in == 1:
-        inputs = pseudo_xs[0]
-      else:
-        inputs = pseudo_xs[:n_in]
-
-      in_shape = base.nested_map(inputs, lambda x: x.shape)
-      in_dtype = base.nested_map(inputs, lambda x: x.dtype)
-      param, state = layer.initialize_once(in_shape, in_dtype, layer_rng)
-      pparam = layer._params   # pylint: disable=protected-access
-
-      outputs, _ = layer.pseudo_forward(inputs, pparam, state)
-
-      # Push outputs onto remaining pseudo_xs (if any).
-      if n_in < _count_items(pseudo_xs):
-        if layer.n_outputs == 1:
-          outputs = (outputs,)
-        pseudo_xs = outputs + pseudo_xs[n_in:]
-      else:
-        pseudo_xs = outputs  # NOTE: can be single value or tuple.
-
-      params.append(param)
-      states.append(state)
-    return params, states
-
-
-@base.layer(n_outputs=2)
-def Dup(x, **unused_kwargs):
-  """Duplicates (copies) an element."""
-  return (x, x)
-
-
-@base.layer(n_inputs=2, n_outputs=2)
-def Swap(xs, **unused_kwargs):
-  """Swaps two elements."""
-  return (xs[1], xs[0])
-
-
-def Dup2():
-  """Copy first 2 elements of the stack: (a, b, ...) -> (a, b, a, b, ...)."""
-  return Serial([
-      # Stack is (a, b, ...)
-      Parallel(Dup(), Dup()),  # pylint: disable=no-value-for-parameter
-      # Stack is (a, a, b, b, ...)
-      Parallel([], Swap()),  # pylint: disable=no-value-for-parameter
-      # Stack is (a, b, a, b, ...)
-  ])
-
-
-def Dup3():
-  """Copy 3 elements of the stack: (a, b, c, ...) -> (a, b, c, a, b, c, ...)."""
-  return Serial([
-      # Stack is (a, b, c, ...)
-      Parallel(Dup(), Dup(), Dup()),  # pylint: disable=no-value-for-parameter
-      # Stack is (a, a, b, b, c, c, ...)
-      Parallel([], Swap(), Swap()),  # pylint: disable=no-value-for-parameter
-      # Stack is (a, b, a, c, b, c, ...)
-      Parallel([], [], Swap()),  # pylint: disable=no-value-for-parameter
-      # Stack is (a, b, c, a, b, c, ...)
-  ])
-
-
-@base.layer(n_outputs=0)
-def Drop(x, **unused_kwargs):
-  """Drops one element."""
-  del x  # Just for the compiler.
-  return ()
-
-
-@base.layer(n_inputs=0)
-def FlattenList(xs, **unused_kwargs):
-  """Flatten lists."""
-  # TODO(jonni): Consider renaming layer to DeepFlatten.
-  return tuple(_deep_flatten(xs))
-
-
-def _nested_op(inputs, op):  # pylint: disable=invalid-name
-  """Helper: apply op over a list of arrays or nested arrays."""
-  # If input is a dictionary, apply to the values (ignore keys).
-  if isinstance(inputs, dict):
-    return _nested_op(list(inputs.values()), op)
-  # First the simple non-nested case.
-  if not isinstance(inputs[0], (list, tuple)):
-    return op(inputs)
-  # In the nested case, sum on each axis separately.
-  result_list = []
-  for i in range(len(inputs[0])):
-    result_list.append(_nested_op([x[i] for x in inputs], op=op))
-  if isinstance(inputs[0], list):
-    return result_list
-  return tuple(result_list)
-
-
-@base.layer(n_inputs=2)
-def Add(xs, **unused_kwargs):
-  """Adds two tensors."""
-  return xs[0] + xs[1]
-
-
-@base.layer(n_inputs=2)
-def SubtractTop(xs, **unused_kwargs):
-  """Subtracts the first tensor from the second."""
-  return xs[1] - xs[0]
-
-
-@base.layer(n_inputs=2)
-def Multiply(xs, **unused_kwargs):
-  """Multiplies two tensors."""
-  return xs[0] * xs[1]
-
-
-@base.layer(n_inputs=3)
-def Gate(xs, **unused_kwargs):
-  """Implements a gating function on a (memory, gate, candidate) tuple.
-
-  Final update is memory * gate + (1-gate) * candidate
-
-  This gating equation may also be referred to as Highway Network.
-  Highway Networks: https://arxiv.org/abs/1505.00387
-
-  Args:
-    xs: A tuple of memory, gate, candidate
-
-  Returns:
-    The result of applying gating.
-  """
-  state, gate, candidate = xs
-  return gate * state + (1.0 - gate) * candidate
-
-
-class Concatenate(base.Layer):
-  """Concatenates n tensors into a single tensor."""
-
-  def __init__(self, n_items=2, axis=-1):
-    super(Concatenate, self).__init__(n_inputs=n_items)
-    self._n_items = n_items
-    self._axis = axis
-
-  def forward(self, xs, params=(), state=(), **kwargs):
-    del params, kwargs
-    return backend.numpy.concatenate(xs, self._axis), state
-
-
-class Split(base.Layer):
-  """Splits the input into sections along an axis."""
-
-  def __init__(self, n_sections=2, axis=-1):
-    super(Split, self).__init__(n_outputs=n_sections)
-    self._n_sections = n_sections
-    self._axis = axis
-
-  def forward(self, inputs, params=(), state=(), **kwargs):
-    del params, kwargs
-    res = tuple(backend.numpy.split(inputs, self._n_sections, self._axis))
-    return res, state
-
-
-class Parallel(base.Layer):
-  """Combinator that applies a list of layers in parallel to its inputs.
-
-  Layers in the list apply to successive spans of inputs, where the spans are
-  determined how many inputs each layer takes. The resulting output is the
-  (flattened) concatenation of the resepective layer outputs.
-
-  For example, suppose one has three layers:
-
-    - F: 1 input, 1 output
-    - G: 3 inputs, 1 output
-    - H: 2 inputs, 2 outputs (h1, h2)
-
-  Then Parallel(F, G, H) will take 6 inputs and give 4 outputs:
-
-    - inputs: a, b, c, d, e, f
-    - outputs: F(a), G(b, c, d), h1, h2
-
-  As an important special case, a None argument to Parallel acts as if it takes
-  one argument, which it leaves unchanged. (It acts as a one-arg no-op.) For
-  example:
-
-    Parallel(None, F)
-
-  creates a layer that passes its first input unchanged and applies F to the
-  following input(s).
-  """
-
-  def __init__(self, *layers):
-    """The constructor.
-
-    Args:
-      *layers: A list of layers.
-
-    Returns:
-      A new layer in which each of the given layers applies to its corresponding
-      span of elements in the dataflow stack.
-    """
-    super(Parallel, self).__init__()
-    layers = self._validate(layers)
-    self._n_layers = len(layers)
-    self._sublayers = layers
-    self._n_inputs = sum(x.n_inputs for x in layers)
-    self._n_outputs = sum(x.n_outputs for x in layers)
-
-  def _validate(self, layers):
-    if not layers or len(layers) < 2:
-      raise ValueError(
-          'layers ({}) must be a list with at least two elements'.format(
-              layers))
-    layers = list(layers)  # Ensure we can modify layers.
-    for i, obj in enumerate(layers):
-      if obj is None or obj == []:  # pylint: disable=g-explicit-bool-comparison
-        layers[i] = Serial(None)
-      elif isinstance(obj, (list, tuple)):
-        layers[i] = Serial(obj)
-      else:
-        if not isinstance(obj, base.Layer):
-          raise ValueError(
-              'Found nonlayer object ({}) in layers list: [{}].'.format(
-                  obj, layers))
-      if layers[i].n_inputs == 0:
-        raise ValueError(
-            'Sublayer with n_inputs = 0 not allowed in Parallel:'
-            ' {}'.format(layers[i]))
-    return layers
-
-  def _allot_to_sublayers(self, inputs):
-    """Divides Parallel's inputs for use by the sublayers.
-
-    Args:
-      inputs: Tuple of elements.
-
-    Returns:
-      A tuple that partitions this layer's inputs among its sublayers.
-      Sublayers that take one argument get that argument directly. All other
-      sublayers get a tuple of items.
-    """
-    start, end = 0, 0
-    sub_inputs = []
-    for layer in self.sublayers:
-      n_in = layer.n_inputs
-      end = start + n_in
-      if n_in == 1:
-        sub_inputs.append(inputs[start])
-      else:
-        sub_inputs.append(inputs[start:end])
-      start = end
-    return tuple(sub_inputs)
-
-  @base.Layer.params.setter
-  def params(self, params):
-    """Recursively sets params on this layer and all sublayers."""
-    self._params = params
-    assert len(params) == self._n_layers
-    for layer, sublayer_params in zip(self.sublayers, params):
-      layer.params = sublayer_params
-
-  @base.Layer.state.setter
-  def state(self, state):
-    """Recursively sets non-param state on this layer and all sublayers."""
-    self._state = state
-    assert len(state) == self._n_layers
-    for layer, sublayer_state in zip(self.sublayers, state):
-      layer.state = sublayer_state
-
-  def forward(self, inputs, params=(), state=(), **kwargs):
-    n_layers, layers = self._n_layers, self.sublayers
-    sublayer_inputs = self._allot_to_sublayers(inputs)
-    rngs = _pop_rng_and_split(kwargs, n_layers)
-    assert len(sublayer_inputs) == n_layers
-    assert len(params) == n_layers
-    assert len(state) == n_layers
-    assert len(rngs) == n_layers
-    outputs = []
-    new_state = []
-    for layer, x, p, s, r in zip(layers, sublayer_inputs, params, state, rngs):
-      # Note that zip silently truncates its result if lengths don't match.
-      sub_outputs, sub_state = layer.apply_forward(x, params=p, state=s, rng=r,
-                                                   **kwargs)
-      if layer.n_outputs == 1:
-        outputs.append(sub_outputs)
-      else:
-        outputs.extend(sub_outputs)
-      new_state.append(sub_state)
-    output = outputs[0] if self.n_outputs == 1 else tuple(outputs)
-    return output, new_state
-
-  def new_params_and_state(self, input_shapes, input_dtypes, rng):
-    sublayer_shapes = self._allot_to_sublayers(input_shapes)
-    sublayer_dtypes = self._allot_to_sublayers(input_dtypes)
-    rngs = backend.random.split(rng, self._n_layers)
-    inits = [layer.initialize_once(shape, dtype, rng)
-             for layer, shape, dtype, rng
-             in zip(self.sublayers, sublayer_shapes, sublayer_dtypes, rngs)]
-    if not inits:
-      return (), ()
-    else:
-      return tuple(zip(*inits))
-
-
-def Residual(*layers, **kwargs):
-  """Constructs a residual version of layers, summing input to layers output."""
-  shortcut = kwargs.get('shortcut')  # default None signals no-op
-  return [
-      Dup(),  # pylint: disable=no-value-for-parameter
-      Parallel(shortcut, layers),
-      Add(),  # pylint: disable=no-value-for-parameter
-  ]
diff --git a/tensor2tensor/trax/layers/combinators_test.py b/tensor2tensor/trax/layers/combinators_test.py
deleted file mode 100644
index cf4d7907b..000000000
--- a/tensor2tensor/trax/layers/combinators_test.py
+++ /dev/null
@@ -1,116 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 combinator layers."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import absltest
-from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import combinators as cb
-from tensor2tensor.trax.layers import core
-
-
-class CombinatorLayerTest(absltest.TestCase):
-
-  def test_drop(self):
-    layer = cb.Drop()
-    input_shape = (3, 2)
-    expected_shape = ()
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
-  def test_dup(self):
-    layer = cb.Dup()
-    input_shape = (3, 2)
-    expected_shape = ((3, 2), (3, 2))
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
-  def test_swap(self):
-    layer = cb.Swap()
-    input_shape = ((3, 2), (4, 7))
-    expected_shape = ((4, 7), (3, 2))
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
-  def test_serial_no_op(self):
-    layer = cb.Serial(None)
-    input_shape = ((3, 2), (4, 7))
-    expected_shape = ((3, 2), (4, 7))
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
-  def test_serial_no_op_list(self):
-    layer = cb.Serial([])
-    input_shape = ((3, 2), (4, 7))
-    expected_shape = ((3, 2), (4, 7))
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
-  def test_serial_one_in_one_out(self):
-    layer = cb.Serial(core.Div(divisor=2.0))
-    input_shape = (3, 2)
-    expected_shape = (3, 2)
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
-  def test_serial_div_div(self):
-    layer = cb.Serial(core.Div(divisor=2.0), core.Div(divisor=5.0))
-    input_shape = (3, 2)
-    expected_shape = (3, 2)
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
-  def test_serial_dup_dup(self):
-    layer = cb.Serial(cb.Dup(), cb.Dup())
-    input_shape = (3, 2)
-    expected_shape = ((3, 2), (3, 2), (3, 2))
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
-  def test_parallel_dup_dup(self):
-    layer = cb.Parallel(cb.Dup(), cb.Dup())
-    input_shape = ((3, 2), (4, 7))
-    expected_shape = ((3, 2), (3, 2), (4, 7), (4, 7))
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
-  def test_parallel_div_div(self):
-    layer = cb.Parallel(core.Div(divisor=0.5), core.Div(divisor=3.0))
-    input_shape = ((3, 2), (4, 7))
-    expected_shape = ((3, 2), (4, 7))
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
-  def test_parallel_no_ops(self):
-    layer = cb.Parallel([], None)
-    input_shape = ((3, 2), (4, 7))
-    expected_shape = ((3, 2), (4, 7))
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
-  def test_branch_op_not_defined(self):
-    with self.assertRaises(AttributeError):
-      cb.Branch([], [])
-
-  def test_select_op_not_defined(self):
-    input_shape = ((3, 2), (4, 7))
-    with self.assertRaises(AttributeError):
-      cb.Select(1, input_shape)
-
-if __name__ == '__main__':
-  absltest.main()
diff --git a/tensor2tensor/trax/layers/convolution.py b/tensor2tensor/trax/layers/convolution.py
deleted file mode 100644
index 6a7239feb..000000000
--- a/tensor2tensor/trax/layers/convolution.py
+++ /dev/null
@@ -1,126 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Trax convolution layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import itertools
-import operator
-
-import six
-
-from tensor2tensor.trax import backend
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import initializers as init
-
-
-class Conv(base.Layer):
-  """Layer constructor function for a general convolution layer."""
-
-  def __init__(self, filters, kernel_size, strides=None, padding='VALID',
-               dimension_numbers=('NHWC', 'HWIO', 'NHWC'),
-               kernel_initializer=None,
-               bias_initializer=init.RandomNormalInitializer(1e-6)):
-    super(Conv, self).__init__()
-    self._filters = filters
-    self._kernel_size = kernel_size
-    self._padding = padding
-    self._dimension_numbers = dimension_numbers
-    self._lhs_spec, self._rhs_spec, self._out_spec = dimension_numbers
-    self._one = (1,) * len(kernel_size)
-    self._strides = strides or self._one
-    self._bias_initializer = bias_initializer
-    rhs_spec = self._rhs_spec
-    self._kernel_initializer = kernel_initializer
-    if kernel_initializer is None:
-      self._kernel_initializer = init.GlorotNormalInitializer(
-          rhs_spec.index('O'), rhs_spec.index('I'))
-
-  def _check_nhwc(self):
-    msg = 'Convolutions on more than 4 dimensions only supported in NHWC.'
-    assert self._lhs_spec == self._out_spec == 'NHWC', msg
-
-  def forward(self, x, params=(), state=(), **kwargs):
-    del kwargs
-    w, b = params
-    x_shape = list(x.shape)
-    if len(x_shape) > 4:
-      self._check_nhwc()
-      new_batch_dim = six.moves.reduce(operator.mul, x_shape[:-3])
-      x = np.reshape(x, [new_batch_dim] + x_shape[-3:])
-    res = backend.conv(
-        x, w, self._strides, self._padding, self._dimension_numbers,
-        self._one) + b
-    if len(x_shape) > 4:
-      res = np.reshape(res, x_shape[:-3] + list(res.shape[-3:]))
-    return res, state
-
-  def _kernel_shape(self, input_shape):
-    """Helper to calculate the kernel shape."""
-    kernel_size_iter = iter(self._kernel_size)
-    return [self._filters if c == 'O' else
-            input_shape[self._lhs_spec.index('C')] if c == 'I' else
-            next(kernel_size_iter) for c in self._rhs_spec]
-
-  def new_params_and_state(self, input_shape, input_dtype, rng):
-    del input_dtype
-    if len(input_shape) > 4:
-      self._check_nhwc()
-      new_batch_dim = six.moves.reduce(operator.mul, input_shape[:-3])
-      input_shape = [new_batch_dim] + list(input_shape[-3:])
-    kernel_shape = self._kernel_shape(input_shape)
-    bias_shape = [self._filters if c == 'C' else 1 for c in self._out_spec]
-    bias_shape = tuple(itertools.dropwhile(lambda x: x == 1, bias_shape))
-    w = self._kernel_initializer(kernel_shape, rng)
-    b = self._bias_initializer(bias_shape, rng)
-    return (w, b), ()
-
-
-class CausalConv(Conv):
-  """Causal (masked) convolution for [batch x time x depth] sequences.
-
-  Maintains causality along time axis. Used in language modeling tasks.
-  """
-
-  def __init__(self,
-               filters,
-               kernel_width=3,
-               kernel_initializer=None,
-               bias_initializer=init.RandomNormalInitializer(1e-6)):
-    super(CausalConv, self).__init__(
-        filters=filters,
-        kernel_size=(kernel_width,),
-        strides=None,
-        padding='VALID',
-        dimension_numbers=('NWC', 'WIO', 'NWC'),
-        kernel_initializer=kernel_initializer,
-        bias_initializer=bias_initializer)
-
-  def forward(self, x, params=(), state=(), **kwargs):
-    assert self._padding == 'VALID'
-    # Left pad with 0s. Applying an unmasked valid convolution on top of this
-    # yields a causal convolution.
-    # TODO(ddohan): Support strided and dilated convolutions.
-    rate = 1
-    effective_kernel_size = int((self._kernel_size[0] - 1) * rate + 1)
-    pad = effective_kernel_size - 1
-    x_leftpad = np.pad(x, pad_width=[[0, 0], [pad, 0], [0, 0]], mode='constant')
-
-    res = super(CausalConv, self).forward(x_leftpad, params)
-    return res
diff --git a/tensor2tensor/trax/layers/convolution_test.py b/tensor2tensor/trax/layers/convolution_test.py
deleted file mode 100644
index b17f4721f..000000000
--- a/tensor2tensor/trax/layers/convolution_test.py
+++ /dev/null
@@ -1,53 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 convolution layers."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import absltest
-from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import convolution
-
-
-class ConvolutionLayerTest(absltest.TestCase):
-
-  def test_conv(self):
-    input_shape = (29, 5, 5, 20)
-    result_shape = base.check_shape_agreement(
-        convolution.Conv(30, (3, 3)), input_shape)
-    self.assertEqual(result_shape, (29, 3, 3, 30))
-
-  def test_conv_rebatch(self):
-    input_shape = (3, 29, 5, 5, 20)
-    result_shape = base.check_shape_agreement(
-        convolution.Conv(30, (3, 3)), input_shape)
-    self.assertEqual(result_shape, (3, 29, 3, 3, 30))
-
-
-class CausalConvolutionTest(absltest.TestCase):
-
-  def test_causal_conv(self):
-    input_shape = (29, 5, 20)
-    conv = convolution.CausalConv(filters=30, kernel_width=3)
-    result_shape = base.check_shape_agreement(conv, input_shape)
-    self.assertEqual(result_shape, (29, 5, 30))
-
-    # TODO(ddohan): How to test for causality? Gradient check between positions?
-
-
-if __name__ == "__main__":
-  absltest.main()
diff --git a/tensor2tensor/trax/layers/core.py b/tensor2tensor/trax/layers/core.py
deleted file mode 100644
index dd62ad85e..000000000
--- a/tensor2tensor/trax/layers/core.py
+++ /dev/null
@@ -1,269 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Trax layers library."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import jax
-import numpy as onp
-
-from tensor2tensor.trax import backend
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import initializers as init
-
-
-@base.layer()
-def Relu(x, **unused_kwargs):
-  return np.maximum(x, np.zeros_like(x))
-
-
-@base.layer()
-def ParametricRelu(x, a=1., **unused_kwargs):
-  return np.maximum(a * x, np.zeros_like(x))
-
-
-@base.layer()
-def LeakyRelu(x, a=0.01, **unused_kwargs):
-  return np.where(x >= 0, x, a * x)
-
-
-@base.layer()
-def Elu(x, a=1., **unused_kwargs):
-  return np.where(x > 0, x, a * np.expm1(x))
-
-
-@base.layer()
-def Selu(x,
-         alpha=1.6732632423543772848170429916717,
-         lmbda=1.0507009873554804934193349852946):
-  return lmbda * np.where(x > 0, x, alpha * np.expm1(x))
-
-
-@base.layer()
-def Gelu(x, **unused_kwargs):
-  return x * backend.erf(x)
-
-
-@base.layer()
-def Sigmoid(x, **unused_kwargs):
-  return backend.expit(x)
-
-
-@base.layer()
-def Tanh(x, **unused_kwargs):
-  return np.tanh(x)
-
-
-@base.layer()
-def HardSigmoid(x, **unused_kwargs):
-  """Linear approximation to sigmoid."""
-  return np.maximum(0, np.minimum(1, (1 + x)))
-
-
-@base.layer()
-def HardTanh(x, **unused_kwargs):
-  """Linear approximation to tanh."""
-  return np.maximum(-1, np.minimum(1, x))
-
-
-@base.layer()
-def Exp(x, **unused_kwargs):
-  return np.exp(x)
-
-
-@base.layer()
-def LogSoftmax(x, axis=-1, **unused_kwargs):
-  """Apply log softmax to x: log-normalize along the given axis."""
-  return x - backend.logsumexp(x, axis, keepdims=True)
-
-
-@base.layer()
-def Softmax(x, axis=-1, **unused_kwargs):
-  """Apply softmax to x: exponentiate and normalize along the given axis."""
-  return np.exp(x - backend.logsumexp(x, axis, keepdims=True))
-
-
-@base.layer()
-def Softplus(x, **unused_kwargs):
-  return np.logaddexp(x, 0.)
-
-
-@base.layer()
-def ToFloat(x, **unused_kwargs):
-  return x.astype(onp.float32)
-
-
-class Dense(base.Layer):
-  """A dense (a.k.a. fully-connected, affine) layer."""
-
-  def __init__(self,
-               n_units,
-               kernel_initializer=init.GlorotUniformInitializer(),
-               bias_initializer=init.RandomNormalInitializer(1e-6)):
-    super(Dense, self).__init__()
-    self._n_units = n_units
-    self._kernel_initializer = kernel_initializer
-    self._bias_initializer = bias_initializer
-
-  def forward(self, x, params=(), state=(), **kwargs):
-    del kwargs
-    w, b = params
-    return np.dot(x, w) + b, state
-
-  def new_params_and_state(self, input_shape, input_dtype, rng):
-    del input_dtype
-    rng1, rng2 = backend.random.split(rng, 2)
-    w = self._kernel_initializer((input_shape[-1], self._n_units), rng1)
-    b = self._bias_initializer((self._n_units,), rng2)
-    return (w, b), ()
-
-
-class Embedding(base.Layer):
-  """Layer constructor function for an embedding layer."""
-
-  def __init__(self,
-               d_feature,
-               vocab_size,
-               kernel_initializer=init.GlorotUniformInitializer()):
-    super(Embedding, self).__init__()
-    self._d_feature = d_feature  # feature dimensionality
-    self._vocab_size = vocab_size
-    self._kernel_initializer = kernel_initializer
-
-  def forward(self, x, params=(), state=(), **kwargs):
-    del kwargs
-    return np.take(params, x, axis=0), state
-
-  def new_params_and_state(self, input_shape, input_dtype, rng):
-    del input_shape, input_dtype
-    out_dim = (self._vocab_size, self._d_feature)
-    params = self._kernel_initializer(out_dim, rng)
-    return params, ()
-
-
-# Flatten.
-@base.layer()
-def Flatten(x, n_axes_to_keep=1, **unused_kwargs):
-  if n_axes_to_keep >= len(x.shape):
-    raise ValueError("n_axes_to_keep[%d] should be less than input's rank[%d]" %
-                     (n_axes_to_keep, len(x.shape)))
-  return np.reshape(x, (x.shape[:n_axes_to_keep] + (-1,)))
-
-
-class Dropout(base.Layer):
-  """Dropout."""
-
-  def __init__(self, rate=0.0, name='dropout', mode='train'):
-    super(Dropout, self).__init__()
-    self._initial_rate = rate
-    # TODO(lukaszkaiser): remove the name property by the end of September'19.
-    # It's only needed for a specific purpose in the short term, will go.
-    self._name = 'dropout_' + name
-    self._mode = mode
-
-  def new_params_and_state(self, input_shape, input_dtype, rng):
-    del input_shape, input_dtype, rng
-    params = ()
-    state = {self._name: np.array(self._initial_rate)}
-    return params, state
-
-  def forward(self, x, params=(), state=(), rng=None, **kwargs):
-    """Execute dropout."""
-    del kwargs
-    rate = self._initial_rate
-    if isinstance(state, dict) and self._name in state:
-      rate = state[self._name]
-    if rng is None:
-      msg = ('Dropout layer requires apply_fn to be called with a rng keyword '
-             'argument. That is, instead of `Dropout(params, inputs)`, call '
-             'it like `Dropout(params, inputs, rng=key)`.')
-      raise ValueError(msg)
-    if self._mode != 'train':
-      return x, state
-    keep = backend.random.bernoulli(rng, 1.0 - rate, x.shape)
-    return np.where(keep, x / (1.0 - rate), np.zeros_like(x)), state
-
-
-@base.layer()
-def Div(x, divisor=1.0, **unused_kwargs):
-  return x / divisor
-
-
-@base.layer()
-def AddConstant(x, constant=0.0, **unused_kwargs):
-  return x + constant
-
-
-@base.layer()
-def MulConstant(x, constant=1.0, **unused_kwargs):
-  return x * constant
-
-
-def one_hot(x, size, dtype=np.float32):  # pylint: disable=invalid-name
-  """Make a n+1 dim one-hot array from n dim int-categorical array."""
-  arange_size = np.arange(size)
-  if backend.get_name() == 'jax':
-    # Work around a jax broadcasting issue.
-    arange_size = jax.lax.tie_in(x, arange_size)
-  return np.array(x[..., np.newaxis] == arange_size, dtype)
-
-
-# Mean.
-@base.layer()
-def Mean(x, axis=-1, keepdims=False, **unused_kwargs):
-  return np.mean(x, axis=axis, keepdims=keepdims)
-
-
-def log_gaussian_pdf(x, mu, sigma):  # pylint: disable=invalid-name
-  """Compute log N(x | mu, sigma)."""
-  a = mu.shape[-1] * np.log(2 * np.pi)
-  _, b = np.linalg.slogdet(sigma)
-  y = np.linalg.solve(sigma, x - mu)
-  y = np.expand_dims(y, axis=-1)
-  xm = np.expand_dims(x - mu, axis=-2)
-  c = np.matmul(xm, y)
-  c = np.squeeze(np.squeeze(c, axis=-1), axis=-1)
-  return -0.5 * (a + b + c)
-
-
-def log_gaussian_diag_pdf(x, mu, diag_sigma):  # pylint: disable=invalid-name
-  """Compute log N(x | mu, eye(diag_sigma))."""
-  a = mu.shape[-1] * np.log(2 * np.pi)
-  b = np.sum(np.log(diag_sigma), axis=-1)
-  y = x - mu / diag_sigma
-  y = np.expand_dims(y, axis=-1)
-  xm = np.expand_dims(x - mu, axis=-2)
-  c = np.matmul(xm, y)
-  c = np.squeeze(np.squeeze(c, axis=-1), axis=-1)
-  return -0.5 * (a + b + c)
-
-
-def multigaussian_loss(preds, targets, ngauss=1):  # pylint: disable=invalid-name
-  """Compute mixture of gaussians loss."""
-  ndims = targets.shape[-1]
-  logits = preds[:, :ngauss]
-  mus = preds[:, ngauss:ngauss*(ndims + 1)]
-  sigmas = preds[:, ngauss(ndims + 1):]
-  sigmas = sigmas * sigmas + 1e-6  # Make positive.
-  loglogits = logits - backend.logsumexp(logits, axis=-1, keepdims=True)
-  mus = np.reshape(mus, [-1, ngauss, ndims])
-  sigmas = np.reshape(sigmas, [-1, ngauss, ndims])
-  targets = np.reshape(targets, [-1, 1, ndims])
-  glogprobs = log_gaussian_diag_pdf(targets, mus, sigmas)
-  return backend.logsumexp(loglogits + glogprobs, axis=-1)
diff --git a/tensor2tensor/trax/layers/core_test.py b/tensor2tensor/trax/layers/core_test.py
deleted file mode 100644
index b4df4dfba..000000000
--- a/tensor2tensor/trax/layers/core_test.py
+++ /dev/null
@@ -1,125 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 core layers."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import absltest
-import numpy as onp
-from tensor2tensor.trax import backend
-from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import combinators
-from tensor2tensor.trax.layers import core
-
-
-class CoreLayerTest(absltest.TestCase):
-
-  def test_flatten_n(self):
-    input_shape = (29, 87, 10, 20, 30)
-
-    layer = core.Flatten()
-    expected_shape = (29, 87 * 10 * 20 * 30)
-    actual_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(actual_shape, expected_shape)
-
-    layer = core.Flatten(n_axes_to_keep=2)
-    expected_shape = (29, 87, 10 * 20 * 30)
-    actual_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(actual_shape, expected_shape)
-
-    layer = core.Flatten(n_axes_to_keep=3)
-    expected_shape = (29, 87, 10, 20 * 30)
-    actual_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(actual_shape, expected_shape)
-
-    layer = core.Flatten(n_axes_to_keep=4)
-    expected_shape = (29, 87, 10, 20, 30)
-    actual_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(actual_shape, expected_shape)
-
-    # Not enough dimensions.
-    with self.assertRaises(base.LayerError):
-      base.check_shape_agreement(core.Flatten(n_axes_to_keep=5), input_shape)
-
-    with self.assertRaises(base.LayerError):
-      base.check_shape_agreement(core.Flatten(n_axes_to_keep=6), input_shape)
-
-  def test_div(self):
-    layer = core.Div(divisor=2.0)
-    input_np = onp.array([[1, 2, 3], [4, 5, 6]], dtype=onp.float32)
-    output_np = layer(input_np)
-    # absltest doesn't have ndarray equalities.
-    expected_output_np = input_np / 2.0
-    self.assertAlmostEqual(
-        0.0,
-        onp.sum((output_np - expected_output_np) ** 2),
-        delta=1e-6)
-
-  def test_div_shapes(self):
-    layer = core.Div(divisor=2.0)
-    input_shape = (3, 2)
-    expected_shape = (3, 2)
-    output_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, expected_shape)
-
-  def test_dense_param_sharing(self):
-    model1 = combinators.Serial(core.Dense(32), core.Dense(32))
-    layer = core.Dense(32)
-    model2 = combinators.Serial(layer, layer)
-
-    rng1, rng2 = backend.random.split(backend.random.get_prng(0), 2)
-    params1, _ = model1.initialize_once((1, 32), onp.float32, rng1)
-    params2, _ = model2.initialize_once((1, 32), onp.float32, rng2)
-    # The first parameters have 2 kernels of size (32, 32).
-    self.assertEqual((32, 32), params1[0][0].shape)
-    self.assertEqual((32, 32), params1[1][0].shape)
-    # The second parameters have 1 kernel of size (32, 32) and an empty dict.
-    self.assertEqual((32, 32), params2[0][0].shape)
-    self.assertEqual((), params2[1])
-
-  def test_dropout(self):
-    input_shape = (8, 7, 9)
-    output_shape = (8, 7, 9)
-    final_shape = base.check_shape_agreement(
-        core.Dropout(rate=0.1, mode="train"), input_shape)
-    self.assertEqual(final_shape, output_shape)
-    final_shape = base.check_shape_agreement(
-        core.Dropout(rate=0.1, mode="eval"), input_shape)
-    self.assertEqual(final_shape, output_shape)
-
-  def test_log_gaussian_pdf(self):
-    x = onp.zeros((2, 5), dtype=onp.float32)
-    mu = x
-    dsigma = onp.eye(5)[None, :, :]
-    sigma = onp.concatenate([dsigma, 2*dsigma], axis=0)
-    prob = core.log_gaussian_pdf(x, mu, sigma)
-    self.assertEqual(prob.shape, (2,))
-    self.assertEqual(int(prob[0]), -4)
-    self.assertEqual(int(prob[1]), -6)
-
-  def test_log_gaussian_diag_pdf(self):
-    x = onp.zeros((2, 5), dtype=onp.float32)
-    mu = x
-    sigma = onp.ones((5,))[None, :]
-    sigma = onp.concatenate([sigma, 2*sigma], axis=0)
-    prob = core.log_gaussian_diag_pdf(x, mu, sigma)
-    self.assertEqual(prob.shape, (2,))
-    self.assertEqual(int(prob[0]), -4)
-    self.assertEqual(int(prob[1]), -6)
-
-if __name__ == "__main__":
-  absltest.main()
diff --git a/tensor2tensor/trax/layers/initializers.py b/tensor2tensor/trax/layers/initializers.py
deleted file mode 100644
index b5c8c0d56..000000000
--- a/tensor2tensor/trax/layers/initializers.py
+++ /dev/null
@@ -1,173 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Trax initializers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as onp
-from tensor2tensor.trax import backend
-
-
-def _GetFans(shape, out_dim=-1, in_dim=-2):
-  """Get the fan-in and fan-out sizes for the given shape and dims."""
-  # Temporary fix until numpy.delete supports negative indices.
-  if out_dim < 0:
-    out_dim += len(shape)
-  if in_dim < 0:
-    in_dim += len(shape)
-
-  receptive_field = backend.numpy.prod(onp.delete(shape, [in_dim, out_dim]))
-  if len(shape) >= 2:
-    fan_in, fan_out = shape[in_dim], shape[out_dim]
-  elif len(shape) == 1:
-    fan_in = shape[0]
-    fan_out = shape[0]
-  else:
-    fan_in = 1.
-    fan_out = 1.
-    fan_in *= receptive_field
-    fan_out *= receptive_field
-  return fan_in, fan_out
-
-
-def RandomNormalInitializer(stddev=1e-2):
-  """An initializer function for random normal coefficients."""
-
-  def Init(shape, rng):
-    return (stddev * backend.random.normal(rng, shape)).astype('float32')
-
-  return Init
-
-
-def RandomUniformInitializer(lim=1.0):
-  """An initializer function for random uniform coefficients."""
-
-  def Init(shape, rng):
-    return (backend.random.uniform(rng, shape, backend.numpy.float32, -lim,
-                                   lim))
-
-  return Init
-
-
-def VarianceScalingInitializer(out_dim, in_dim, scale, mode, distribution):
-  """Initializer capable of adapting its scale to the shape of weights."""
-  if scale <= 0.:
-    raise ValueError('scale must be positive float, {} given'.format(scale))
-  if mode not in {'fan_in', 'fan_out', 'fan_avg'}:
-    raise ValueError(
-        'Invalid mode argument:, {}, must be either fan_in, fan_out or fan_avg'
-        .format(mode))
-
-  def Init(shape, rng):
-    """The initializer function."""
-    fan_in, fan_out = _GetFans(shape, out_dim, in_dim)
-    gain = scale
-    if mode == 'fan_in':
-      gain /= fan_in
-    elif mode == 'fan_out':
-      gain /= fan_out
-    elif mode == 'fan_avg':
-      gain /= (fan_in + fan_out) / 2
-    if distribution == 'truncated_normal':
-      # constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.)
-      stddev = backend.numpy.sqrt(gain) / .87962566103423978
-      return (backend.random.truncated_normal(rng, -2, 2, shape) *
-              stddev).astype('float32')
-    elif distribution == 'normal':
-      return (backend.random.normal(rng, shape) *
-              backend.numpy.sqrt(gain)).astype('float32')
-    elif distribution == 'uniform':
-      lim = backend.numpy.sqrt(3. * gain)
-      return (backend.random.uniform(rng, shape, backend.numpy.float32, -lim,
-                                     lim))
-    else:
-      raise ValueError('invalid distribution for variance scaling Initializer')
-
-  return Init
-
-
-def GlorotNormalInitializer(out_dim=-1, in_dim=-2, scale=1.):
-  """An initializer function for random Glorot-scaled coefficients."""
-  return VarianceScalingInitializer(out_dim, in_dim, scale, 'fan_avg', 'normal')
-
-
-def GlorotUniformInitializer(out_dim=-1, in_dim=-2, scale=1.):
-  """An initializer function for random uniform Glorot-scaled coefficients."""
-  return VarianceScalingInitializer(out_dim, in_dim, scale, 'fan_avg',
-                                    'uniform')
-
-
-def LeCunNormalInitializer(out_dim=-1, in_dim=-2, scale=1.):
-  """An initializer function for random LeCun-scaled coefficients."""
-  return VarianceScalingInitializer(out_dim, in_dim, scale, 'fan_in', 'normal')
-
-
-def LeCunUniformInitializer(out_dim=-1, in_dim=-2, scale=1.):
-  """An initializer function for random uniform LeCun-scaled coefficients."""
-  return VarianceScalingInitializer(out_dim, in_dim, scale, 'fan_in', 'uniform')
-
-
-def KaimingNormalInitializer(out_dim=-1, in_dim=-2, param=0.):
-  """An initializer function for random Kaiming-scaled coefficients."""
-  return VarianceScalingInitializer(out_dim, in_dim,
-                                    2.0 / backend.numpy.sqrt(1 + param**2),
-                                    'fan_in', 'normal')
-
-
-def KaimingUniformInitializer(out_dim=-1, in_dim=-2, param=0.):
-  """An initializer function for random uniform Kaiming-scaled coefficients."""
-  return VarianceScalingInitializer(out_dim, in_dim,
-                                    2.0 / backend.numpy.sqrt(1 + param**2),
-                                    'fan_in', 'uniform')
-
-
-def OrthogonalInitializer(stddev=1.0):
-  """Orthogonal Initializer."""
-  def Init(shape, rng):
-    """The orthogonal initializer function."""
-    # Have at least 2 elements in shape.
-    cur_shape = list(shape)
-    while len(cur_shape) < 2:
-      cur_shape = [1] + cur_shape
-
-    # Flatten the input shape with the last dimension remaining.
-    n_rows = 1
-    for dim in cur_shape[:-1]:
-      n_rows *= dim
-    n_cols = cur_shape[-1]
-    flat_shape = (n_cols, n_rows) if n_rows < n_cols else (n_rows, n_cols)
-
-    # Generate a random matrix
-    a = backend.random.normal(rng, flat_shape, dtype=backend.numpy.float32)
-
-    # Compute the qr factorization
-    q, r = backend.numpy.linalg.qr(a)
-
-    # Make Q uniform
-    d = backend.numpy.diag(r)
-    q *= backend.numpy.sign(d)
-
-    # Transpose and reshape back q if needed.
-    if n_rows < n_cols:
-      q = backend.numpy.transpose(q)
-    q = backend.numpy.reshape(q, shape)
-
-    # Return scaled as requested.
-    return stddev * q
-
-  return Init
diff --git a/tensor2tensor/trax/layers/initializers_test.py b/tensor2tensor/trax/layers/initializers_test.py
deleted file mode 100644
index 469fb29fa..000000000
--- a/tensor2tensor/trax/layers/initializers_test.py
+++ /dev/null
@@ -1,83 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 initializers."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import absltest
-from tensor2tensor.trax.backend import random
-from tensor2tensor.trax.layers import initializers
-
-
-class InitializersTest(absltest.TestCase):
-
-  def test_random_normal(self):
-    initializer = initializers.RandomNormalInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-  def test_lecun_uniform(self):
-    initializer = initializers.LeCunUniformInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-  def test_random_uniform(self):
-    initializer = initializers.RandomUniformInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-  def test_glorot_normal(self):
-    initializer = initializers.GlorotNormalInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-  def test_glorot_uniform(self):
-    initializer = initializers.GlorotUniformInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-  def test_lecun_normal(self):
-    initializer = initializers.LeCunNormalInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-  def test_kaiming_normal(self):
-    initializer = initializers.KaimingNormalInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-  def test_kaiming_uniform(self):
-    initializer = initializers.KaimingUniformInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-  def test_orthogonal(self):
-    initializer = initializers.OrthogonalInitializer()
-    input_shape = (29, 5, 7, 20)
-    init_value = initializer(input_shape, random.get_prng(0))
-    self.assertEqual(tuple(init_value.shape), input_shape)
-
-if __name__ == "__main__":
-  absltest.main()
diff --git a/tensor2tensor/trax/layers/intro.ipynb b/tensor2tensor/trax/layers/intro.ipynb
deleted file mode 100644
index 0658786f1..000000000
--- a/tensor2tensor/trax/layers/intro.ipynb
+++ /dev/null
@@ -1,834 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "7yuytuIllsv1"
-      },
-      "source": [
-        "# A Conceptual, Practical Introduction to Trax Layers\n",
-        "\n",
-        "This notebook introduces the core concepts and programming components of the Trax library through a series of code samples and explanations. The topics covered in following sections are:\n",
-        "  - **layers**: the basic building blocks and how to combine them into networks\n",
-        "  - **data flows, data stack**: how the Trax runtime moves data through the layers\n",
-        "  - **models**: how to train, evaluate, and run predictions with Trax models\n",
-        "  - **new layer classes**: how to define and test your own Layer classes\n",
-        "\n",
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "BIl27504La0G"
-      },
-      "source": [
-        "## General Setup\n",
-        "Execute the following few cells (once) before running any of the code samples in this notebook."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 0,
-      "metadata": {
-        "colab": {},
-        "colab_type": "code",
-        "id": "oILRLCWN_16u"
-      },
-      "outputs": [],
-      "source": [
-        "#@title\n",
-        "# Copyright 2018 Google LLC.\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",
-        "# https://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",
-        "import numpy as onp\n",
-        "\n",
-        "\n",
-        "\n",
-        "# Import Trax\n",
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 2,
-      "metadata": {
-        "cellView": "both",
-        "colab": {
-          "height": 51
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 383,
-          "status": "ok",
-          "timestamp": 1570168980195,
-          "user": {
-            "displayName": "",
-            "photoUrl": "",
-            "userId": ""
-          },
-          "user_tz": 420
-        },
-        "id": "vlGjGoGMTt-D",
-        "outputId": "6d2ecf3d-3eb8-48a7-ad12-ebefe83afaf1"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "/bin/sh: pip: command not found\n",
-            "/bin/sh: pip: command not found\n"
-          ]
-        }
-      ],
-      "source": [
-        "#@title Run for installation.\n",
-        "\n",
-        "! pip install -q -U tensor2tensor\n",
-        "! pip install -q tensorflow\n",
-        "\n",
-        "from tensor2tensor.trax import backend\n",
-        "from tensor2tensor.trax import layers as tl"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 0,
-      "metadata": {
-        "colab": {},
-        "colab_type": "code",
-        "id": "bYWNWL9MJHv9"
-      },
-      "outputs": [],
-      "source": [
-        "onp.set_printoptions(precision=3)  # Less visual noise in the numerical outputs.\n",
-        "\n",
-        "def show_layer_properties(layer_obj, layer_name):\n",
-        "  template = ('{}.n_inputs:  {}\\n'\n",
-        "              '{}.n_outputs: {}\\n'\n",
-        "              '{}.sublayers: {}\\n'\n",
-        "              '{}.params:    {}\\n')\n",
-        "  print(template.format(layer_name, layer_obj.n_inputs,\n",
-        "                        layer_name, layer_obj.n_outputs,\n",
-        "                        layer_name, layer_obj.sublayers,\n",
-        "                        layer_name, layer_obj.params))  "
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "-LQ89rFFsEdk"
-      },
-      "source": [
-        "# Layers\n",
-        "\n",
-        "The Layer class represents Trax's concept of a layer, as summarized in the start of the class's docstring:\n",
-        "```\n",
-        "class Layer(object):\n",
-        "  \"\"\"Base class for composable layers in a deep learning network.\n",
-        "\n",
-        "  Layers are the basic building blocks for deep learning models. A Trax layer\n",
-        "  computes a function from zero or more inputs to zero or more outputs,\n",
-        "  optionally using trainable parameters (common) and non-parameter state (not\n",
-        "  common). Authors of new layer subclasses typically override at most two\n",
-        "  methods of the base `Layer` class:\n",
-        "\n",
-        "    forward(inputs, params=(), state=(), **kwargs):\n",
-        "      Computes this layer's output as part of a forward pass through the model.\n",
-        "\n",
-        "    new_params_and_state(self, input_shape, input_dtype, rng):\n",
-        "      Returns a (params, state) pair suitable for initializing this layer.\n",
-        "```"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "LyLVtdxorDPO"
-      },
-      "source": [
-        "## A layer computes a function."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "ntZ4_eNQldzL"
-      },
-      "source": [
-        "A layer computes a function from zero or more inputs to zero or more outputs. The inputs and outputs are NumPy arrays or JAX objects wrapping NumPy arrays.\n",
-        "\n",
-        "The simplest layers, those with no parameters, state or sublayers, can be used without initialization. You can think of them (and test them) like simple mathematical functions. For ease of testing and interactive exploration, layer\n",
-        "objects implement the `__call__ ` method, so you can call them directly on input data:\n",
-        "```\n",
-        "y = layer(x)\n",
-        "```\n",
-        "\n",
-        "Layers are also objects, so you can inspect their properties. For example:\n",
-        "```\n",
-        "print('Number of inputs required by this layer: {}'.format(layer.n_inputs))\n",
-        "```\n",
-        "\n",
-        "### Example 1. tl.Relu [1 input, 1 output]"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 4,
-      "metadata": {
-        "colab": {
-          "height": 221
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 1543,
-          "status": "ok",
-          "timestamp": 1570168982080,
-          "user": {
-            "displayName": "",
-            "photoUrl": "",
-            "userId": ""
-          },
-          "user_tz": 420
-        },
-        "id": "V09viOSEQvQe",
-        "outputId": "b7c1c085-3b54-4673-f284-99d6440f8a52"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "x:\n",
-            "[[-7. -6. -5. -4. -3.]\n",
-            " [-2. -1.  0.  1.  2.]\n",
-            " [ 3.  4.  5.  6.  7.]]\n",
-            "\n",
-            "relu(x):\n",
-            "[[0. 0. 0. 0. 0.]\n",
-            " [0. 0. 0. 1. 2.]\n",
-            " [3. 4. 5. 6. 7.]]\n",
-            "\n",
-            "number of inputs expected by this layer: 1\n",
-            "number of outputs promised by this layer: 1\n"
-          ]
-        }
-      ],
-      "source": [
-        "x = onp.arange(-7, 8).reshape(3, -1).astype(onp.float32)\n",
-        "\n",
-        "# Create a layer object (a Relu instance) and apply the layer to data x.\n",
-        "relu = tl.Relu()\n",
-        "y = relu(x)\n",
-        "\n",
-        "# Show input, output, and two layer properties.\n",
-        "template = ('x:\\n{}\\n\\n'\n",
-        "            'relu(x):\\n{}\\n\\n'\n",
-        "            'number of inputs expected by this layer: {}\\n'\n",
-        "            'number of outputs promised by this layer: {}')\n",
-        "print(template.format(x, y, relu.n_inputs, relu.n_outputs))"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "7sYxIT8crFVE"
-      },
-      "source": [
-        "### Example 2. tl.Concatenate [2 inputs, 1 output]"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 5,
-      "metadata": {
-        "colab": {
-          "height": 442
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 1558,
-          "status": "ok",
-          "timestamp": 1570168983657,
-          "user": {
-            "displayName": "",
-            "photoUrl": "",
-            "userId": ""
-          },
-          "user_tz": 420
-        },
-        "id": "LMPPNWXLoOZI",
-        "outputId": "24398ccb-9cda-4bdd-c0f0-4904c02a215e"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "x1:\n",
-            "[[-7. -6. -5. -4. -3.]\n",
-            " [-2. -1.  0.  1.  2.]\n",
-            " [ 3.  4.  5.  6.  7.]]\n",
-            "\n",
-            "x2:\n",
-            "[[-70. -60. -50. -40. -30.]\n",
-            " [-20. -10.   0.  10.  20.]\n",
-            " [ 30.  40.  50.  60.  70.]]\n",
-            "\n",
-            "concat0([x1, x2]):\n",
-            "[[ -7.  -6.  -5.  -4.  -3.]\n",
-            " [ -2.  -1.   0.   1.   2.]\n",
-            " [  3.   4.   5.   6.   7.]\n",
-            " [-70. -60. -50. -40. -30.]\n",
-            " [-20. -10.   0.  10.  20.]\n",
-            " [ 30.  40.  50.  60.  70.]]\n",
-            "\n",
-            "concat1([x1, x2]):\n",
-            "[[ -7.  -6.  -5.  -4.  -3. -70. -60. -50. -40. -30.]\n",
-            " [ -2.  -1.   0.   1.   2. -20. -10.   0.  10.  20.]\n",
-            " [  3.   4.   5.   6.   7.  30.  40.  50.  60.  70.]]\n",
-            "\n",
-            "concat0: Concatenate{in=2,out=1}\n",
-            "concat1: Concatenate{in=2,out=1}\n"
-          ]
-        }
-      ],
-      "source": [
-        "x1 = onp.arange(-7, 8).reshape(3, -1).astype(onp.float32)\n",
-        "x2 = 10 * x1\n",
-        "\n",
-        "concat0 = tl.Concatenate(axis=0)\n",
-        "concat1 = tl.Concatenate(axis=1)\n",
-        "\n",
-        "y0 = concat0([x1, x2])\n",
-        "y1 = concat1([x1, x2])\n",
-        "\n",
-        "template = ('x1:\\n{}\\n\\n'\n",
-        "            'x2:\\n{}\\n\\n'\n",
-        "            'concat0([x1, x2]):\\n{}\\n\\n'\n",
-        "            'concat1([x1, x2]):\\n{}\\n')\n",
-        "print(template.format(x1, x2, y0, y1))\n",
-        "\n",
-        "# Print abbreviated object representations (useful for debugging).\n",
-        "print('concat0: {}'.format(concat0))\n",
-        "print('concat1: {}'.format(concat1))"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "1oZv3R8bRMvF"
-      },
-      "source": [
-        "## Layers are trainable."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "3d64M7wLryji"
-      },
-      "source": [
-        "Most layer types are trainable: they include parameters that modify the computation of outputs from inputs, and they use back-progagated gradients to update those parameters.\n",
-        "\n",
-        "Before use, trainable layers must have their parameters initialized, typically using a PRNG (pseudo-random number generator) key for random number generation. Trax's model trainers take care of this behind the scenes, but if you are using a layer in insolation, you have to do the initialization yourself. For this, use the `initialize_once` method:\n",
-        "\n",
-        "```\n",
-        "  def initialize_once(self, input_shapes, input_dtype, rng):\n",
-        "    \"\"\"Initializes this layer and its sublayers recursively.\n",
-        "\n",
-        "    This method is designed to initialize each layer instance once, even if the\n",
-        "    same layer instance occurs in multiple places in the network. This enables\n",
-        "    weight sharing to be implemented as layer sharing.\n",
-        "\n",
-        "    ...\n",
-        "```\n",
-        "\n",
-        "### Example 3. tl.LayerNorm [1 input, 1 output, has parameters]"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 6,
-      "metadata": {
-        "colab": {
-          "height": 221
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 2555,
-          "status": "ok",
-          "timestamp": 1570168986228,
-          "user": {
-            "displayName": "",
-            "photoUrl": "",
-            "userId": ""
-          },
-          "user_tz": 420
-        },
-        "id": "Ie7iyX91qAx2",
-        "outputId": "3fe02659-481b-4912-c7eb-85eb01cfadd6"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "x:\n",
-            "[[-7. -6. -5. -4. -3.]\n",
-            " [-2. -1.  0.  1.  2.]\n",
-            " [ 3.  4.  5.  6.  7.]]\n",
-            "\n",
-            "layer_norm(x):\n",
-            "[[-1.414 -0.707  0.     0.707  1.414]\n",
-            " [-1.414 -0.707  0.     0.707  1.414]\n",
-            " [-1.414 -0.707  0.     0.707  1.414]]\n",
-            "\n",
-            "layer_norm.params:\n",
-            "(_FilledConstant([1., 1., 1., 1., 1.], dtype=float32), _FilledConstant([0., 0., 0., 0., 0.], dtype=float32))\n"
-          ]
-        }
-      ],
-      "source": [
-        "prng_key = backend.random.get_prng(0)  # Used below for layer initialization.\n",
-        "x = onp.arange(-7, 8).reshape(3, -1).astype(onp.float32)\n",
-        "\n",
-        "layer_norm = tl.LayerNorm()\n",
-        "layer_norm.initialize_once(x.shape, x.dtype, prng_key)\n",
-        "y = layer_norm(x)\n",
-        "\n",
-        "template = ('x:\\n{}\\n\\n'\n",
-        "            'layer_norm(x):\\n{}\\n')\n",
-        "print(template.format(x, y))\n",
-        "print('layer_norm.params:\\n{}'.format(layer_norm.params))"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "ZWZUXEJAofH-"
-      },
-      "source": [
-        "## Layers combine into layers."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "d47gVdGV1vWw"
-      },
-      "source": [
-        "The Trax library authors encourage users, where possible, to build new layers as combinations of existing layers. The library provides a small set of _combinator_ layers for this: layer objects that make a list of layers behave as a single layer (a unit able to compute outputs from inputs, update parameters from gradients, and combine with yet more layers).\n",
-        "\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "vC1ymG2j0iyp"
-      },
-      "source": [
-        "## Combine with Serial(...)\n",
-        "\n",
-        "The most common way to combine layers is serially, using the `Serial` class:\n",
-        "```\n",
-        "class Serial(base.Layer):\n",
-        "  \"\"\"Combinator that applies layers serially (by function composition).\n",
-        "\n",
-        "  A Serial combinator uses stack semantics to manage data for its sublayers.\n",
-        "  Each sublayer sees only the inputs it needs and returns only the outputs it\n",
-        "  has generated. The sublayers interact via the data stack. For instance, a\n",
-        "  sublayer k, following sublayer j, gets called with the data stack in the\n",
-        "  state left after layer j has applied. The Serial combinator then:\n",
-        "\n",
-        "    - takes N_in items off the top of the stack (N_in = k.n_inputs) and calls\n",
-        "      layer k, passing those items as arguments; and\n",
-        "\n",
-        "    - takes layer k's N_out return values (N_out = k.n_outputs) and pushes\n",
-        "      them onto the data stack.\n",
-        "\n",
-        "  ...\n",
-        "```\n",
-        "As described above, the output of one layer is the input of the next, which amounts to function composition:\n",
-        "\n",
-        "```\n",
-        "#  h(.) = g(f(.))\n",
-        "layer_h = Serial(\n",
-        "    layer_f,\n",
-        "    layer_g,\n",
-        ")\n",
-        "```\n",
-        "\n",
-        "### Example 4. y = layer_norm(relu(x)) [1 input, 1 output, has parameters]"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 7,
-      "metadata": {
-        "colab": {
-          "height": 170
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 1664,
-          "status": "ok",
-          "timestamp": 1570168987915,
-          "user": {
-            "displayName": "",
-            "photoUrl": "",
-            "userId": ""
-          },
-          "user_tz": 420
-        },
-        "id": "dW5fpusjvjmh",
-        "outputId": "207f6a59-b767-414f-a836-ec342157ef51"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "x:\n",
-            "[[-7. -6. -5. -4. -3.]\n",
-            " [-2. -1.  0.  1.  2.]\n",
-            " [ 3.  4.  5.  6.  7.]]\n",
-            "\n",
-            "layer_block(x):\n",
-            "[[ 0.     0.     0.     0.     0.   ]\n",
-            " [-0.75  -0.75  -0.75   0.5    1.75 ]\n",
-            " [-1.414 -0.707  0.     0.707  1.414]]\n"
-          ]
-        }
-      ],
-      "source": [
-        "prng_key = backend.random.get_prng(0)\n",
-        "x = onp.arange(-7, 8).reshape(3, -1).astype(onp.float32)\n",
-        "\n",
-        "layer_block = tl.Serial(\n",
-        "    tl.Relu(),\n",
-        "    tl.LayerNorm(),\n",
-        ")\n",
-        "layer_block.initialize_once(x.shape, x.dtype, prng_key)\n",
-        "y = layer_block(x)\n",
-        "\n",
-        "template = ('x:\\n{}\\n\\n'\n",
-        "            'layer_block(x):\\n{}')\n",
-        "print(template.format(x, y,))"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "bRtmN6ckQO1q"
-      },
-      "source": [
-        "And we can inspect the block as a whole, as if it were just another layer:\n",
-        "\n",
-        "### Example 5. Inspecting a Serial layer."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 8,
-      "metadata": {
-        "colab": {
-          "height": 102
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 289,
-          "status": "ok",
-          "timestamp": 1570168988225,
-          "user": {
-            "displayName": "",
-            "photoUrl": "",
-            "userId": ""
-          },
-          "user_tz": 420
-        },
-        "id": "D6BpYddZQ1eu",
-        "outputId": "03a99733-cd84-4639-fb1f-8dfacebf5b07"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "layer_block:\n",
-            "Serial{in=1,out=1,sublayers=[Relu{in=1,out=1}, LayerNorm{in=1,out=1}]}\n",
-            "\n",
-            "layer_block.params:\n",
-            "[(), (_FilledConstant([1., 1., 1., 1., 1.], dtype=float32), _FilledConstant([0., 0., 0., 0., 0.], dtype=float32))]\n"
-          ]
-        }
-      ],
-      "source": [
-        "print('layer_block:\\n{}\\n'.format(layer_block))\n",
-        "\n",
-        "print('layer_block.params:\\n{}'.format(layer_block.params))"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "PqVNdoONcTp0"
-      },
-      "source": [
-        "## Combine with Parallel(...)\n",
-        "\n",
-        "The `Parallel` combinator arranges layers into separate computational channels, each with its own inputs/outputs and gradient flows:\n",
-        "```\n",
-        "class Parallel(base.Layer):\n",
-        "  \"\"\"Combinator that applies a list of layers in parallel to its inputs.\n",
-        "\n",
-        "  Layers in the list apply to successive spans of inputs, where the spans are\n",
-        "  determined how many inputs each layer takes. The resulting output is the\n",
-        "  (flattened) concatenation of the resepective layer outputs.\n",
-        "\n",
-        "  For example, suppose one has three layers:\n",
-        "\n",
-        "    - F: 1 input, 1 output\n",
-        "    - G: 3 inputs, 1 output\n",
-        "    - H: 2 inputs, 2 outputs (h1, h2)\n",
-        "\n",
-        "  Then Parallel(F, G, H) will take 6 inputs and give 4 outputs:\n",
-        "\n",
-        "    - inputs: a, b, c, d, e, f\n",
-        "    - outputs: F(a), G(b, c, d), h1, h2\n",
-        "```\n",
-        "\n",
-        "Separate (parallel) computation channels make sense when each channel can do its work (computing outputs from inputs) independent of the inputs and outputs of the others.\n",
-        "\n",
-        "As a simplistic example, consider writing a converter from three-digit octal (base 8) numerals to their corresponding values. For instance, to do conversions such as\n",
-        "```\n",
-        "123 (octal) = 1 * 8^2 + 2 * 8^1 + 3 * 8^0 =  83 (decimal)\n",
-        "345 (octal) = 3 * 8^2 + 4 * 8^1 + 6 * 8^0 = 229 (decimal)\n",
-        "567 (octal) = 5 * 8^2 + 6 * 8^1 + 7 * 8^0 = 375 (decimal)\n",
-        "```\n",
-        "the digits can first be converted independently, according to their place value (multiply by 64, multiply by 8, or multiply by 1). The following code runs the 64's-place digits ([1, 3, 5]) through one layer, the 8's-place digits ([2, 4, 6]) through a different layer, and the 1's-place digits ([3, 5, 7]) through yet a different layer. These three layers are combined in a Parallel layer:\n",
-        "\n",
-        "### Example 6. Processing octal digits in parallel."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 9,
-      "metadata": {
-        "colab": {
-          "height": 204
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 2224,
-          "status": "ok",
-          "timestamp": 1570168990465,
-          "user": {
-            "displayName": "",
-            "photoUrl": "",
-            "userId": ""
-          },
-          "user_tz": 420
-        },
-        "id": "uQMqq3h_b2jQ",
-        "outputId": "f3a43cae-e271-493a-f74a-31e1ff971bc1"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "inputs:\n",
-            "(array([1, 3, 5]), array([2, 4, 6]), array([3, 5, 7]))\n",
-            "\n",
-            "octal_place_values(inputs):\n",
-            "(array([ 64., 192., 320.]), array([16., 32., 48.]), array([3., 5., 7.]))\n",
-            "\n",
-            "octal_place_values.n_inputs:  3\n",
-            "octal_place_values.n_outputs: 3\n",
-            "octal_place_values.sublayers: [MulConstant{in=1,out=1}, MulConstant{in=1,out=1}, MulConstant{in=1,out=1}]\n",
-            "octal_place_values.params:    ((), (), ())\n",
-            "\n"
-          ]
-        }
-      ],
-      "source": [
-        "prng_key = backend.random.get_prng(0)\n",
-        "place_64_digits = onp.array([1, 3, 5])\n",
-        "place_8_digits = onp.array([2, 4, 6])\n",
-        "place_1_digits = onp.array([3, 5, 7])\n",
-        "inputs = (place_64_digits, place_8_digits, place_1_digits)\n",
-        "input_shapes = [[3]] * 3\n",
-        "input_dtypes = [onp.int32] * 3\n",
-        "\n",
-        "# Create three simple layers, each for converting a different digit in base 8.\n",
-        "sixty_fours = tl.MulConstant(constant=64.0)  # 8^2: 100 in base 8\n",
-        "eights = tl.MulConstant(constant=8.0)  # 8^1: 10 in base 8\n",
-        "ones = tl.MulConstant(constant=1.0)  # 8^0: 1 in base 8\n",
-        "\n",
-        "# Create a combined layer to convert digits to values (using big-endian base 8),\n",
-        "# initialize it, and apply it.\n",
-        "octal_place_values = tl.Parallel(sixty_fours, eights, ones)\n",
-        "octal_place_values.initialize_once(input_shapes, input_dtypes, prng_key)\n",
-        "outputs = octal_place_values(inputs)\n",
-        "\n",
-        "# Show inputs, outputs, and properties.\n",
-        "template = ('inputs:\\n{}\\n\\n'\n",
-        "            'octal_place_values(inputs):\\n{}\\n')\n",
-        "print(template.format(inputs, outputs))\n",
-        "show_layer_properties(octal_place_values, 'octal_place_values')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "q_xcWide3e5f"
-      },
-      "source": [
-        "To complete the example, the three outputs (values for the different digits) are combined by successive pairwise additions:\n",
-        "\n",
-        "### Example 6'. Combining outputs from upstream parallel layers."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 10,
-      "metadata": {
-        "colab": {
-          "height": 275
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 2139,
-          "status": "ok",
-          "timestamp": 1570168992621,
-          "user": {
-            "displayName": "",
-            "photoUrl": "",
-            "userId": ""
-          },
-          "user_tz": 420
-        },
-        "id": "ZDCkrvUp3u0-",
-        "outputId": "696f21aa-5dad-4284-bfdd-ae637e2ce53f"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "inputs:\n",
-            "(array([1, 3, 5]), array([2, 4, 6]), array([3, 5, 7]))\n",
-            "\n",
-            "octal_place_values(inputs):\n",
-            "(array([ 64., 192., 320.]), array([16., 32., 48.]), array([3., 5., 7.]))\n",
-            "\n",
-            "evaluate_octal(inputs):\n",
-            "[ 83. 229. 375.]\n",
-            "\n",
-            "evaluate_octal.n_inputs:  3\n",
-            "evaluate_octal.n_outputs: 1\n",
-            "evaluate_octal.sublayers: [Parallel{in=3,out=3,sublayers=[MulConstant{in=1,out=1}, MulConstant{in=1,out=1}, MulConstant{in=1,out=1}]}, Add{in=2,out=1}, Add{in=2,out=1}]\n",
-            "evaluate_octal.params:    [(), (), ()]\n",
-            "\n"
-          ]
-        }
-      ],
-      "source": [
-        "evaluate_octal = tl.Serial(\n",
-        "    octal_place_values,\n",
-        "    tl.Add(),  # Adds the 64's-place values and the 8's-place values.\n",
-        "    tl.Add(),  # Adds the 1's-place values to the sums from the previous Add.\n",
-        ")\n",
-        "evaluate_octal.initialize_once(input_shapes, input_dtypes, prng_key)\n",
-        "y = evaluate_octal(inputs)\n",
-        "\n",
-        "template = ('inputs:\\n{}\\n\\n'\n",
-        "            'octal_place_values(inputs):\\n{}\\n\\n'\n",
-        "            'evaluate_octal(inputs):\\n{}\\n')\n",
-        "print(template.format(inputs, outputs, y))\n",
-        "show_layer_properties(evaluate_octal, 'evaluate_octal')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "rwgiP0tK1H6p"
-      },
-      "source": [
-        "# Data Flows, Data Stack"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "llAH3cdE1UeU"
-      },
-      "source": [
-        "# Training and Using Models"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "65ite-671cTT"
-      },
-      "source": [
-        "# Defining Your Own Layer Classes"
-      ]
-    }
-  ],
-  "metadata": {
-    "colab": {
-      "collapsed_sections": [],
-      "last_runtime": {
-        "build_target": "//learning/deepmind/dm_python:dm_notebook",
-        "kind": "private"
-      },
-      "name": "A Conceptual, Practical Introduction to Trax Layers",
-      "provenance": [
-        {
-          "file_id": "1sF8QbqJ19ZU6oy5z4GUTt4lgUCjqO6kt",
-          "timestamp": 1569980697572
-        },
-        {
-          "file_id": "1EH76AWQ_pvT4i8ZXfkv-SCV4MrmllEl5",
-          "timestamp": 1563927451951
-        }
-      ]
-    },
-    "kernelspec": {
-      "display_name": "Python 2",
-      "name": "python2"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}
diff --git a/tensor2tensor/trax/layers/metrics.py b/tensor2tensor/trax/layers/metrics.py
deleted file mode 100644
index e79c83cb2..000000000
--- a/tensor2tensor/trax/layers/metrics.py
+++ /dev/null
@@ -1,124 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Trax metrics layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import combinators as cb
-from tensor2tensor.trax.layers import core
-
-
-@base.layer(n_inputs=2, n_outputs=1)
-def CrossEntropy(x, axis=-1, **kw):
-  del kw
-  prediction, target = x
-  return np.sum(prediction * core.one_hot(target, prediction.shape[-1]),
-                axis=axis)
-
-
-@base.layer(n_inputs=2, n_outputs=1)
-def L2(x, axis=-1, **kw):
-  del kw
-  prediction, target = x
-  return np.sum((prediction - target)**2, axis=axis)
-
-
-@base.layer(n_inputs=2, n_outputs=1)
-def Accuracy(x, axis=-1, **kw):
-  del kw
-  prediction, target = x
-  predicted_class = np.argmax(prediction, axis=axis)
-  return np.equal(predicted_class, target)
-
-
-@base.layer()
-def WeightMask(target, mask_id=0, **kw):
-  del kw
-  if mask_id is None:
-    return np.ones_like(target)
-  return 1.0 - np.equal(target, mask_id).astype(np.float32)
-
-
-@base.layer(n_inputs=2, n_outputs=1)
-def WeightedMean(x, **kw):
-  del kw
-  metric, weights = x
-  weights_sum = np.sum(weights)
-  return np.sum(metric * weights) / weights_sum
-
-
-def MaskedScalar(metric_layer, mask_id=None, has_weights=False):
-  """Metric as scalar compatible with Trax masking."""
-  # Stack of (inputs, targets) --> (metric, weight-mask).
-  metric_and_mask = [
-      cb.Parallel(
-          [],
-          cb.Dup()  # Duplicate targets
-      ),
-      cb.Parallel(
-          metric_layer,  # Metric: (inputs, targets) --> metric
-          WeightMask(mask_id=mask_id)  # pylint: disable=no-value-for-parameter
-      )
-  ]
-  if not has_weights:
-    # Take (metric, weight-mask) and return the weighted mean.
-    return cb.Serial([metric_and_mask, WeightedMean()])  # pylint: disable=no-value-for-parameter
-  return cb.Serial([
-      metric_and_mask,
-      cb.Parallel(
-          [],
-          cb.Multiply()  # Multiply given weights by mask_id weights
-      ),
-      WeightedMean()  # pylint: disable=no-value-for-parameter
-  ])
-
-
-def CrossEntropyScalar(mask_id=None, has_weights=False):
-  """Cross-entropy as scalar compatible with Trax masking."""
-  return MaskedScalar(CrossEntropy(), mask_id=mask_id, has_weights=has_weights)  # pylint: disable=no-value-for-parameter
-
-
-NegLogPerplexityScalar = CrossEntropyScalar
-
-
-def CrossEntropyLossScalar(mask_id=None, has_weights=False):
-  """Cross-entropy loss as scalar compatible with Trax masking."""
-  return cb.Serial(
-      CrossEntropyScalar(mask_id=mask_id, has_weights=has_weights),
-      core.MulConstant(constant=-1.0)
-  )
-
-
-def L2Scalar(mask_id=None, has_weights=False):
-  """L2 as scalar compatible with Trax masking."""
-  return MaskedScalar(L2(), mask_id=mask_id, has_weights=has_weights)  # pylint: disable=no-value-for-parameter
-
-
-def L2LossScalar(mask_id=None, has_weights=False):
-  """L2 loss as scalar compatible with Trax masking."""
-  return cb.Serial(
-      L2Scalar(mask_id=mask_id, has_weights=has_weights),
-      core.MulConstant(constant=-1.0)
-  )
-
-
-def AccuracyScalar(mask_id=None, has_weights=False):
-  """Accuracy as scalar compatible with Trax masking."""
-  return MaskedScalar(Accuracy(), mask_id=mask_id, has_weights=has_weights)  # pylint: disable=no-value-for-parameter
diff --git a/tensor2tensor/trax/layers/metrics_test.py b/tensor2tensor/trax/layers/metrics_test.py
deleted file mode 100644
index 8eded2671..000000000
--- a/tensor2tensor/trax/layers/metrics_test.py
+++ /dev/null
@@ -1,90 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 metrics layers."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import absltest
-import numpy as onp
-from tensor2tensor.trax import backend
-from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import metrics
-
-
-class MetricsLayerTest(absltest.TestCase):
-
-  def test_cross_entropy(self):
-    input_shape = ((29, 4, 4, 20), (29, 4, 4))
-    result_shape = base.check_shape_agreement(
-        metrics.CrossEntropy(), input_shape)
-    self.assertEqual(result_shape, (29, 4, 4))
-
-  def test_accuracy(self):
-    input_shape = ((29, 4, 4, 20), (29, 4, 4))
-    result_shape = base.check_shape_agreement(
-        metrics.Accuracy(), input_shape)
-    self.assertEqual(result_shape, (29, 4, 4))
-
-  def test_weight_mask(self):
-    input_shape = (29, 4, 4, 20)
-    result_shape = base.check_shape_agreement(
-        metrics.WeightMask(), input_shape)
-    self.assertEqual(result_shape, input_shape)
-
-  def test_weighted_mean_shape(self):
-    input_shape = ((29, 4, 4, 20), (29, 4, 4, 20))
-    result_shape = base.check_shape_agreement(
-        metrics.WeightedMean(), input_shape)
-    self.assertEqual(result_shape, ())
-
-  def test_weighted_mean_semantics(self):
-    inputs = onp.array([1, 2, 3], dtype=onp.float32)
-    weights1 = onp.array([1, 1, 1], dtype=onp.float32)
-    layer = metrics.WeightedMean()
-    rng = backend.random.get_prng(0)
-    layer.initialize_once((inputs.shape, weights1.shape),
-                          (inputs.dtype, weights1.dtype), rng)
-    mean1 = layer((inputs, weights1))
-    onp.testing.assert_allclose(mean1, 2.0)
-    weights2 = onp.array([0, 0, 1], dtype=onp.float32)
-    mean2 = layer((inputs, weights2))
-    onp.testing.assert_allclose(mean2, 3.0)
-    weights3 = onp.array([1, 0, 0], dtype=onp.float32)
-    mean3 = layer((inputs, weights3))
-    onp.testing.assert_allclose(mean3, 1.0)
-
-  def test_cross_entropy_scalar(self):
-    input_shape = ((29, 4, 4, 20), (29, 4, 4))
-    result_shape = base.check_shape_agreement(
-        metrics.CrossEntropyScalar(), input_shape)
-    self.assertEqual(result_shape, ())
-
-  def test_cross_entropy_loss_scalar(self):
-    input_shape = ((29, 4, 4, 20), (29, 4, 4))
-    result_shape = base.check_shape_agreement(
-        metrics.CrossEntropyLossScalar(), input_shape)
-    self.assertEqual(result_shape, ())
-
-  def test_accuracy_scalar(self):
-    input_shape = ((29, 4, 4, 20), (29, 4, 4))
-    result_shape = base.check_shape_agreement(
-        metrics.AccuracyScalar(), input_shape)
-    self.assertEqual(result_shape, ())
-
-
-if __name__ == "__main__":
-  absltest.main()
diff --git a/tensor2tensor/trax/layers/normalization.py b/tensor2tensor/trax/layers/normalization.py
deleted file mode 100644
index fd7d4908c..000000000
--- a/tensor2tensor/trax/layers/normalization.py
+++ /dev/null
@@ -1,125 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Trax normalization layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.layers import base
-
-
-class BatchNorm(base.Layer):
-  """Batch normalization."""
-
-  def __init__(self, axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True,
-               momentum=0.999, mode='train'):
-    super(BatchNorm, self).__init__()
-    self._axis = axis
-    self._epsilon = epsilon
-    self._center = center
-    self._scale = scale
-    self._momentum = momentum
-    self._mode = mode
-
-  def new_params_and_state(self, input_shape, input_dtype, rng):
-    """Helper to initialize batch norm params."""
-    del input_dtype, rng
-    axis = self._axis
-    axis = (axis,) if np.isscalar(axis) else axis
-    shape = tuple(d for i, d in enumerate(input_shape) if i not in axis)
-    beta = np.zeros(shape, dtype='float32') if self._center else ()
-    gamma = np.ones(shape, dtype='float32') if self._scale else ()
-    def get_stats_axis(i, d):
-      if i in axis:
-        return 1
-      else:
-        return d
-    stats_shape = tuple(get_stats_axis(i, d) for i, d in enumerate(input_shape))
-    running_mean = np.zeros(stats_shape, dtype=np.float32)
-    running_var = np.ones(stats_shape, dtype=np.float32)
-    num_batches = np.zeros((), dtype=np.int32)
-    params = (beta, gamma)
-    state = (running_mean, running_var, num_batches)
-    return params, state
-
-  def forward(self, x, params, state, **unused_kwargs):
-    """Layer construction function for a batch normalization layer."""
-
-    running_mean, running_var, num_batches = state
-
-    if self._mode == 'train':
-      mean = np.mean(x, self._axis, keepdims=True)
-      # Fast but less numerically-stable variance calculation than np.var.
-      m1 = np.mean(x**2, self._axis, keepdims=True)
-      var = m1 - mean**2
-      num_batches = num_batches + 1
-      def average(factor, new, old):
-        return (factor * old + (1 - factor) * new).astype(old.dtype)
-      running_mean = average(self._momentum, mean, running_mean)
-      running_var = average(self._momentum, var, running_var)
-      state = (running_mean, running_var, num_batches)
-    else:
-      mean = running_mean
-      var = running_var
-
-    z = (x - mean.astype(x.dtype)) / np.sqrt(var +
-                                             self._epsilon).astype(x.dtype)
-
-    # Expand the parameters to have the right axes.
-    beta, gamma = params
-    # TODO(phawkins): np.expand_dims should accept an axis tuple.
-    # (https://github.com/numpy/numpy/issues/12290)
-    ed = tuple(None if i in self._axis else slice(None)
-               for i in range(np.ndim(x)))
-    beta = beta[ed]
-    gamma = gamma[ed]
-
-    # Return the z rescaled by the parameters if requested.
-    if self._center and self._scale:
-      output = gamma * z + beta
-    elif self._center:
-      output = z + beta
-    elif self._scale:
-      output = gamma * z
-    else:
-      output = z
-    assert output.dtype == x.dtype, ('The dtype of the output (%s) of batch '
-                                     'norm is not the same as the input (%s). '
-                                     'Batch norm should not change the dtype' %
-                                     (output.dtype, x.dtype))
-    return output, state
-
-
-# Layer normalization.
-def _layer_norm_params_and_state(input_shape, input_dtype, rng):
-  """Helper: create layer norm parameters."""
-  del input_dtype, rng
-  features = input_shape[-1]
-  scale = np.ones(features)
-  bias = np.zeros(features)
-  params = (scale, bias)
-  return params, ()
-
-
-@base.layer(new_params_and_state_fn=_layer_norm_params_and_state)
-def LayerNorm(x, params, epsilon=1e-6, **unused_kwargs):  # pylint: disable=invalid-name
-  (scale, bias) = params
-  mean = np.mean(x, axis=-1, keepdims=True)
-  variance = np.mean((x - mean)**2, axis=-1, keepdims=True)
-  norm_inputs = (x - mean) / np.sqrt(variance + epsilon)
-  return norm_inputs * scale + bias
diff --git a/tensor2tensor/trax/layers/normalization_test.py b/tensor2tensor/trax/layers/normalization_test.py
deleted file mode 100644
index 80edbd663..000000000
--- a/tensor2tensor/trax/layers/normalization_test.py
+++ /dev/null
@@ -1,69 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 normalization layers."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import absltest
-import numpy as onp
-
-from tensor2tensor.trax import backend
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import normalization
-
-
-class NormalizationLayerTest(absltest.TestCase):
-
-  def test_batch_norm_shape(self):
-    input_shape = (29, 5, 7, 20)
-    result_shape = base.check_shape_agreement(
-        normalization.BatchNorm(), input_shape)
-    self.assertEqual(result_shape, input_shape)
-
-  def test_batch_norm(self):
-    input_shape = (2, 3, 4)
-    input_dtype = np.float32
-    eps = 1e-5
-    rng = backend.random.get_prng(0)
-    inp1 = np.reshape(np.arange(np.prod(input_shape), dtype=input_dtype),
-                      input_shape)
-    m1 = 11.5  # Mean of this random input.
-    v1 = 47.9167  # Variance of this random input.
-    layer = normalization.BatchNorm(axis=(0, 1, 2))
-    _, _ = layer.initialize_once(input_shape, input_dtype, rng)
-    state = layer.state
-    onp.testing.assert_allclose(state[0], 0)
-    onp.testing.assert_allclose(state[1], 1)
-    self.assertEqual(state[2], 0)
-    out = layer(inp1)
-    state = layer.state
-    onp.testing.assert_allclose(state[0], m1 * 0.001)
-    onp.testing.assert_allclose(state[1], 0.999 + v1 * 0.001, rtol=1e-6)
-    self.assertEqual(state[2], 1)
-    onp.testing.assert_allclose(out, (inp1 - m1) / np.sqrt(v1 + eps),
-                                rtol=1e-6)
-
-  def test_layer_norm_shape(self):
-    input_shape = (29, 5, 7, 20)
-    result_shape = base.check_shape_agreement(
-        normalization.LayerNorm(), input_shape)
-    self.assertEqual(result_shape, input_shape)
-
-
-if __name__ == "__main__":
-  absltest.main()
diff --git a/tensor2tensor/trax/layers/pooling.py b/tensor2tensor/trax/layers/pooling.py
deleted file mode 100644
index 3782491cf..000000000
--- a/tensor2tensor/trax/layers/pooling.py
+++ /dev/null
@@ -1,44 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Trax pooling layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.trax import backend
-from tensor2tensor.trax.layers import base
-
-
-@base.layer()
-def MaxPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
-  del params, kw
-  return backend.max_pool(x, pool_size=pool_size, strides=strides,
-                          padding=padding)
-
-
-@base.layer()
-def SumPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
-  del params, kw
-  return backend.sum_pool(x, pool_size=pool_size, strides=strides,
-                          padding=padding)
-
-
-@base.layer()
-def AvgPool(x, params, pool_size=(2, 2), strides=None, padding='VALID', **kw):
-  del params, kw
-  return backend.avg_pool(x, pool_size=pool_size, strides=strides,
-                          padding=padding)
diff --git a/tensor2tensor/trax/layers/pooling_test.py b/tensor2tensor/trax/layers/pooling_test.py
deleted file mode 100644
index 8924a8b5e..000000000
--- a/tensor2tensor/trax/layers/pooling_test.py
+++ /dev/null
@@ -1,36 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 conv layers."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import absltest
-from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import pooling
-
-
-class PoolingLayerTest(absltest.TestCase):
-
-  def test_avg_pool(self):
-    input_shape = (29, 4, 4, 20)
-    result_shape = base.check_shape_agreement(
-        pooling.AvgPool(pool_size=(2, 2), strides=(2, 2)), input_shape)
-    self.assertEqual(result_shape, (29, 2, 2, 20))
-
-
-if __name__ == "__main__":
-  absltest.main()
diff --git a/tensor2tensor/trax/layers/reversible.py b/tensor2tensor/trax/layers/reversible.py
deleted file mode 100644
index 29f845244..000000000
--- a/tensor2tensor/trax/layers/reversible.py
+++ /dev/null
@@ -1,126 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Implementations of reversible layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import jax
-from tensor2tensor.trax import backend
-from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import combinators as cb
-
-
-class ReversibleLayer(base.Layer):
-  """Reversible Layer."""
-
-  def reverse(self, output, params=(), state=(), **kwargs):
-    """Reverse this layer: compute input given output."""
-    raise NotImplementedError
-
-  def reverse_and_grad(self, output, grad, params=(), state=(), **kwargs):
-    """Backward pass: computes the inverse of a layer and propagates gradients.
-
-    While you may choose to only implement reverse, some layers implement this
-    function directly as computation may be shared between reversing and
-    computing gradients.
-
-    Args:
-      output: Output activations; can be a (possibly nested) tuple.
-      grad: gradient signal (cotangent) computed based on subsequent layers.
-        The structure and shape must match the output.
-      params: layer parameters
-      state: start state
-      **kwargs: kwargs for the layer
-
-    Returns:
-      A tuple (x, (x_grad, params_grad)), where x is the reconstructed input,
-      x_grad is the gradient signal for the input, and params_grad is the
-      gradient signal for the parameters.
-    """
-    # Note: jax.vjp does not allow us to use **kwargs in the signature here.
-    def _do_forward(x, params):
-      return super(ReversibleLayer, self).forward(
-          x, params=params, state=state, **kwargs)[0]
-
-    reconstructed_x = self.reverse(output, params, state, **kwargs)
-    _, vjpfun = jax.vjp(_do_forward, reconstructed_x, params)
-    x_params_grad = vjpfun(grad)
-    return reconstructed_x, x_params_grad
-
-  @property
-  def has_backward(self):
-    return True
-
-  def backward(self, inputs, output, ct, params, state, **kwargs):
-    del inputs
-    _, inputs_params_ct = self.reverse_and_grad(output, ct, params, state,
-                                                **kwargs)
-    return inputs_params_ct
-
-
-class ReversibleSwap(ReversibleLayer, cb.Swap):
-  """Swap the first two element on the stack."""
-
-  def reverse(self, output, params=(), state=(), **kwargs):
-    # Swap is its own inverse, except that reverse doesn't return the state.
-    return self.forward(output, params=params, state=state, **kwargs)[0]
-
-
-class ReversibleSerial(ReversibleLayer, cb.Serial):
-  """A reversible version of tl.Serial (requires reversible sub-layers)."""
-
-  def __init__(self, *layers):
-    super(ReversibleSerial, self).__init__(*layers)
-
-    # Note that sublayers has already been flattened to remove nested lists.
-    for i, layer in enumerate(self.sublayers):
-      if not isinstance(layer, ReversibleLayer):
-        raise ValueError(
-            'Sub-layer {} of ReversibleSerial is not reversible: {}'.format(
-                i, layer))
-
-  def reverse(self, output, params=(), state=(), **kwargs):
-    rng = kwargs.pop('rng', None)
-    rngs = (None,) * self._n_layers
-    if rng is not None:
-      rngs = backend.random.split(rng, self._n_layers)
-
-    layer_val = output
-    for layer, p, s, rng in reversed(list(zip(self.sublayers,
-                                              params, state, rngs))):
-      layer_val = layer.reverse(layer_val, p, s, rng=rng, **kwargs)
-
-    return layer_val
-
-  def reverse_and_grad(self, output, ct, params=(), state=(), **kwargs):
-    rng = kwargs.pop('rng', None)
-    rngs = (None,) * self._n_layers
-    if rng is not None:
-      rngs = backend.random.split(rng, self._n_layers)
-
-    layer_val = output
-    layer_ct = ct
-    params_ct = []
-    for layer, p, s, rng in reversed(list(zip(self.sublayers,
-                                              params, state, rngs))):
-      layer_val, layer_ct = layer.reverse_and_grad(
-          layer_val, layer_ct, p, s, rng=rng, **kwargs)
-      layer_ct, p_ct = layer_ct
-      params_ct.insert(0, p_ct)
-
-    return layer_val, (layer_ct, params_ct)
diff --git a/tensor2tensor/trax/layers/reversible_test.py b/tensor2tensor/trax/layers/reversible_test.py
deleted file mode 100644
index 6f63963b4..000000000
--- a/tensor2tensor/trax/layers/reversible_test.py
+++ /dev/null
@@ -1,37 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 reversible layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import absltest
-from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import reversible
-
-
-class ReversibleLayerTest(absltest.TestCase):
-
-  def test_reversible_swap(self):
-    layer = reversible.ReversibleSwap()
-    input_shape = ((2, 3), (3, 3))
-    final_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(final_shape, input_shape[::-1])
-
-
-if __name__ == '__main__':
-  absltest.main()
diff --git a/tensor2tensor/trax/layers/rnn.py b/tensor2tensor/trax/layers/rnn.py
deleted file mode 100644
index f42ebfe35..000000000
--- a/tensor2tensor/trax/layers/rnn.py
+++ /dev/null
@@ -1,128 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Implementations of common recurrent neural network cells (RNNs)."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.trax.layers import combinators as cb
-from tensor2tensor.trax.layers import convolution
-from tensor2tensor.trax.layers import core
-
-
-def GRUCell(n_units):
-  """Builds a traditional GRU cell with dense internal transformations.
-
-  Gated Recurrent Unit paper: https://arxiv.org/abs/1412.3555
-
-
-  Args:
-    n_units: Number of hidden units.
-
-  Returns:
-    A Stax model representing a traditional GRU RNN cell.
-  """
-  return GeneralGRUCell(
-      candidate_transform=lambda: core.Dense(n_units),
-      memory_transform_fn=None,
-      gate_nonlinearity=core.Sigmoid,
-      candidate_nonlinearity=core.Tanh)
-
-
-def ConvGRUCell(n_units, kernel_size=(3, 3)):
-  """Builds a convolutional GRU.
-
-  Paper: https://arxiv.org/abs/1511.06432.
-
-  Args:
-    n_units: Number of hidden units
-    kernel_size: Kernel size for convolution
-
-  Returns:
-    A Stax model representing a GRU cell with convolution transforms.
-  """
-
-  def BuildConv():
-    return convolution.Conv(
-        filters=n_units, kernel_size=kernel_size, padding='SAME')
-
-  return GeneralGRUCell(
-      candidate_transform=BuildConv,
-      memory_transform_fn=None,
-      gate_nonlinearity=core.Sigmoid,
-      candidate_nonlinearity=core.Tanh)
-
-
-def GeneralGRUCell(candidate_transform,
-                   memory_transform_fn=None,
-                   gate_nonlinearity=core.Sigmoid,
-                   candidate_nonlinearity=core.Tanh,
-                   dropout_rate_c=0.1,
-                   sigmoid_bias=0.5):
-  r"""Parametrized Gated Recurrent Unit (GRU) cell construction.
-
-  GRU update equations:
-  $$ Update gate: u_t = \sigmoid(U' * s_{t-1} + B') $$
-  $$ Reset gate: r_t = \sigmoid(U'' * s_{t-1} + B'') $$
-  $$ Candidate memory: c_t = \tanh(U * (r_t \odot s_{t-1}) + B) $$
-  $$ New State: s_t = u_t \odot s_{t-1} + (1 - u_t) \odot c_t $$
-
-  See combinators.Gate for details on the gating function.
-
-
-  Args:
-    candidate_transform: Transform to apply inside the Candidate branch. Applied
-      before nonlinearities.
-    memory_transform_fn: Optional transformation on the memory before gating.
-    gate_nonlinearity: Function to use as gate activation. Allows trying
-      alternatives to Sigmoid, such as HardSigmoid.
-    candidate_nonlinearity: Nonlinearity to apply after candidate branch. Allows
-      trying alternatives to traditional Tanh, such as HardTanh
-    dropout_rate_c: Amount of dropout on the transform (c) gate. Dropout works
-      best in a GRU when applied exclusively to this branch.
-    sigmoid_bias: Constant to add before sigmoid gates. Generally want to start
-      off with a positive bias.
-
-  Returns:
-    A model representing a GRU cell with specified transforms.
-  """
-  gate_block = [  # u_t
-      candidate_transform(),
-      core.AddConstant(constant=sigmoid_bias),
-      gate_nonlinearity(),
-  ]
-  reset_block = [  # r_t
-      candidate_transform(),
-      core.AddConstant(constant=sigmoid_bias),  # Want bias to start positive.
-      gate_nonlinearity(),
-  ]
-  candidate_block = [
-      cb.Dup(),
-      reset_block,
-      cb.Multiply(),  # Gate S{t-1} with sigmoid(candidate_transform(S{t-1}))
-      candidate_transform(),  # Final projection + tanh to get Ct
-      candidate_nonlinearity(),  # Candidate gate
-
-      # Only apply dropout on the C gate. Paper reports 0.1 as a good default.
-      core.Dropout(rate=dropout_rate_c)
-  ]
-  memory_transform = memory_transform_fn() if memory_transform_fn else []
-  return cb.Model(
-      cb.Dup(), cb.Dup(),
-      cb.Parallel(memory_transform, gate_block, candidate_block),
-      cb.Gate(),
-  )
diff --git a/tensor2tensor/trax/layers/rnn_test.py b/tensor2tensor/trax/layers/rnn_test.py
deleted file mode 100644
index c88ab4b39..000000000
--- a/tensor2tensor/trax/layers/rnn_test.py
+++ /dev/null
@@ -1,45 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 rnn layers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import absltest
-from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.layers import rnn
-
-
-class RnnLayerTest(absltest.TestCase):
-
-  def _test_cell_runs(self, layer, input_shape, output_shape):
-    final_shape = base.check_shape_agreement(layer, input_shape)
-    self.assertEqual(output_shape, final_shape)
-
-  def test_conv_gru_cell(self):
-    self._test_cell_runs(
-        rnn.ConvGRUCell(9, kernel_size=(3, 3)),
-        input_shape=(8, 1, 7, 9),
-        output_shape=(8, 1, 7, 9))
-
-  def test_gru_cell(self):
-    self._test_cell_runs(
-        rnn.GRUCell(9), input_shape=(8, 7, 9), output_shape=(8, 7, 9))
-
-
-if __name__ == '__main__':
-  absltest.main()
diff --git a/tensor2tensor/trax/learning_rate.py b/tensor2tensor/trax/learning_rate.py
deleted file mode 100644
index aa53758c9..000000000
--- a/tensor2tensor/trax/learning_rate.py
+++ /dev/null
@@ -1,264 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""trax learning rate schedules.
-
-The learning rate schedules here all have the signature:
-  lr: history -> (step -> {"learning_rate": lr})
-
-That is, they are functions that take a trax.history.History and return a
-function that takes a step and returns a dict with entry "learning_rate".
-"""
-
-# TODO(pkozakowski): Revisit the decision to control nontrainable parameters
-# using LR schedules, or at least rename the module.
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import random
-import time
-
-from absl import logging
-import gin
-
-from tensor2tensor.trax import models as trax_models
-from tensor2tensor.trax import utils
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.backend import random as jax_random
-from tensor2tensor.trax.rl import online_tune
-from tensor2tensor.trax.rl import ppo
-
-
-@gin.configurable(blacklist=["history"])
-def MultifactorSchedule(history=None,
-                        factors="constant * linear_warmup * rsqrt_decay",
-                        constant=0.1,
-                        warmup_steps=400,
-                        decay_factor=0.5,
-                        steps_per_decay=20000):
-  """Factor-based learning rate schedule.
-
-  Interprets factors in the factors string which can consist of:
-  * constant: interpreted as the constant value,
-  * linear_warmup: interpreted as linear warmup until warmup_steps,
-  * rsqrt_decay: divide by square root of max(step, warmup_steps)
-  * decay_every: Every k steps decay the learning rate by decay_factor.
-
-  Args:
-    history: the history of training and evaluation (History object).
-    factors: a string with factors separated by "*" that defines the schedule.
-    constant: float, the starting constant for the learning rate schedule.
-    warmup_steps: how many steps to warm up for in the warmup schedule.
-    decay_factor: The amount to decay the learning rate by.
-    steps_per_decay: How often to decay the learning rate.
-
-  Returns:
-    a function learning_rate(step): float -> {"learning_rate": float}, the
-    step-dependent lr.
-  """
-  del history
-
-  factors = [n.strip() for n in factors.split("*")]
-
-  def learning_rate(step):  # pylint: disable=invalid-name
-    """Step to learning rate function."""
-    ret = 1.0
-    for name in factors:
-      if name == "constant":
-        ret *= constant
-      elif name == "linear_warmup":
-        ret *= np.minimum(1.0, step / warmup_steps)
-      elif name == "rsqrt_decay":
-        ret /= np.sqrt(np.maximum(step, warmup_steps))
-      elif name == "decay_every":
-        ret *= (decay_factor ** (step//steps_per_decay))
-      else:
-        raise ValueError("Unknown factor %s." % name)
-    ret = np.asarray(ret, dtype=np.float32)
-    return {"learning_rate": ret}
-
-  return learning_rate
-
-
-@gin.configurable(blacklist=["history"])
-def EvalAdjustingSchedule(history,
-                          constant=0.1,
-                          steps_to_decrease=20,
-                          improvement_margin=0.001,
-                          decrease_rate=1.5,
-                          history_mode="eval",
-                          metric="metrics/accuracy"):
-  """Learning rate that decreases when eval metric stalls.
-
-  If the chosen metric does not improve by improvement_margin for as many as
-  steps_to_decrease steps, then the constant gets decreased by decrease rate.
-  Finally, the MultifactorSchedule gets called with the adjusted constant.
-
-  Args:
-    history: trax.history.History, the history of training and evaluation.
-    constant: float, the starting constant for the learning rate schedule.
-    steps_to_decrease: int, after how many steps without improvement
-      should we decrease the constant.
-    improvement_margin: how much we need to improve to consider the metric
-      improved.
-    decrease_rate: by what fraction to decrease (i.e. lr /= decrease_rate).
-    history_mode: str, which mode of the history to use.
-    metric: which evaluation metric to use for adjustments.
-
-  Returns:
-    a function learning_rate(step): float -> {"learning_rate": float}, the
-    step-dependent lr.
-  """
-  metrics = history.get(history_mode, metric)
-  adjusted = constant
-  if len(metrics) < 2:
-    return MultifactorSchedule(history, constant=adjusted)
-
-  steps_without_improvement = 0
-  cur = metrics.pop()[1]  # The most-recent value of the metric.
-  while len(metrics) > 1:
-    # The one-before value of metrics as .pop() removes one element each time.
-    prev = metrics.pop()[1]
-    if cur < prev * (1 + improvement_margin):
-      steps_without_improvement += 1
-    else:
-      cur = prev
-      steps_without_improvement = 0
-    if steps_without_improvement >= steps_to_decrease:
-      adjusted /= decrease_rate
-      cur = prev
-      steps_without_improvement = 0
-
-  return MultifactorSchedule(history, constant=adjusted)
-
-
-@gin.configurable(blacklist=["history"])
-def PolicySchedule(
-    history,
-    observation_metrics=(
-        ("train", "metrics/accuracy"),
-        ("train", "metrics/loss"),
-        ("eval", "metrics/accuracy"),
-        ("eval", "metrics/loss"),
-    ),
-    include_controls_in_observation=False,
-    control_configs=(
-        # (name, start, (low, high), flip)
-        ("learning_rate", 1e-3, (1e-9, 10.0), False),
-    ),
-    observation_range=(0.0, 10.0),
-    action_multipliers=(1.0 / 1.5, 1.0 / 1.25, 1.0, 1.25, 1.5),
-    policy_and_value_model=trax_models.FrameStackMLP,
-    policy_and_value_two_towers=False,
-    policy_and_value_vocab_size=None,
-    policy_dir=gin.REQUIRED,
-    temperature=1.0,
-):
-  """Learning rate schedule controlled by a learned policy.
-
-  Args:
-    history: the history of training and evaluation (History object).
-    observation_metrics: list of pairs (mode, metric), as in the History object.
-    include_controls_in_observation: bool, whether to include the controls in
-      observations.
-    control_configs: control configs, see trax.rl.envs.OnlineTuneEnv.
-    observation_range: tuple (low, high), range to clip the metrics to.
-    action_multipliers: sequence of LR multipliers that policy actions
-      correspond to.
-    policy_and_value_model: Trax model to use as the policy.
-    policy_and_value_two_towers: bool, whether the action distribution and value
-      prediction is computed by separate model towers.
-    policy_and_value_vocab_size: vocabulary size of a policy and value network
-      operating on serialized representation. If None, use raw continuous
-      representation.
-    policy_dir: directory with the policy checkpoint.
-    temperature: temperature for sampling from the policy.
-
-  Returns:
-    a function nontrainable_params(step): float -> {"name": float}, the
-    step-dependent schedule for nontrainable parameters.
-  """
-
-  # Turn the history into observations for the policy. If we don't have any,
-  # return the initial learning rate.
-  start_time = time.time()
-  observations = online_tune.history_to_observations(
-      history, observation_metrics, observation_range,
-      control_configs if include_controls_in_observation else None
-  )
-  logging.vlog(
-      1, "Building observations took %0.2f sec.", time.time() - start_time)
-  if observations.shape[0] == 0:
-    controls = {
-        name: start_value
-        for (name, start_value, _, _) in control_configs
-    }
-    return lambda _: controls
-
-  assert policy_and_value_vocab_size is None, (
-      "Serialized policies are not supported yet."
-  )
-  # Build the policy network and load its parameters.
-  start_time = time.time()
-  net = ppo.policy_and_value_net(
-      n_controls=len(control_configs),
-      n_actions=len(action_multipliers),
-      vocab_size=policy_and_value_vocab_size,
-      bottom_layers_fn=policy_and_value_model,
-      two_towers=policy_and_value_two_towers,
-  )
-  logging.vlog(
-      1, "Building the policy network took %0.2f sec.", time.time() - start_time
-  )
-  start_time = time.time()
-  # (opt_state, state, epoch, opt_step)
-  (opt_state, state, _, _) = ppo.maybe_restore_opt_state(policy_dir)
-  assert opt_state is not None, "Policy checkpoint not found."
-  (params, _) = opt_state
-  logging.vlog(
-      1, "Restoring the policy parameters took %0.2f sec.",
-      time.time() - start_time
-  )
-
-  # Run the policy and sample an action.
-  seed = random.randint(0, 2**31 - 1)
-  rng = jax_random.get_prng(seed=seed)
-  start_time = time.time()
-  # ((log_probs, value_preds), state). We have no way to pass state to the next
-  # step, but that should be fine.
-  (log_probs, _) = (
-      net(np.array([observations]), params=params, state=state, rng=rng))
-  logging.vlog(
-      1, "Running the policy took %0.2f sec.", time.time() - start_time
-  )
-  # Sample from the action distribution for the last timestep.
-  assert log_probs.shape == (
-      1, len(control_configs) * observations.shape[0], len(action_multipliers)
-  )
-  action = utils.gumbel_sample(
-      log_probs[0, -len(control_configs):, :] / temperature
-  )
-
-  # Get new controls.
-  controls = {
-      # name: value
-      control_config[0]: online_tune.update_control(  # pylint: disable=g-complex-comprehension
-          control_config, control_action, history, action_multipliers
-      )
-      for (control_action, control_config) in zip(action, control_configs)
-  }
-  return lambda _: controls
diff --git a/tensor2tensor/trax/learning_rate_test.py b/tensor2tensor/trax/learning_rate_test.py
deleted file mode 100644
index 904310026..000000000
--- a/tensor2tensor/trax/learning_rate_test.py
+++ /dev/null
@@ -1,120 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.learning_rate."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as onp
-
-from tensor2tensor.trax import history as trax_history
-from tensor2tensor.trax import learning_rate
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.backend import random as jax_random
-from tensor2tensor.trax.models import atari_cnn
-from tensor2tensor.trax.rl import online_tune
-from tensor2tensor.trax.rl import ppo
-from tensorflow import test
-
-
-class PolicyScheduleTest(test.TestCase):
-
-  def _make_schedule(
-      self,
-      history,
-      control_configs,
-      observation_metrics=(("eval", "metrics/accuracy"),),
-      action_multipliers=(1.0,),
-  ):
-    policy_and_value_model = atari_cnn.FrameStackMLP
-    net = ppo.policy_and_value_net(
-        n_actions=len(action_multipliers),
-        n_controls=len(control_configs),
-        vocab_size=None,
-        bottom_layers_fn=policy_and_value_model,
-        two_towers=False,
-    )
-    rng = jax_random.get_prng(seed=0)
-    obs_dim = len(observation_metrics)
-    (params, state) = net.initialize_once((1, 1, obs_dim), np.float32, rng)
-    policy_dir = self.get_temp_dir()
-    # Optimizer slots should not be used for anything.
-    slots = None
-    opt_state = (params, slots)
-    ppo.save_opt_state(policy_dir, opt_state, state, epoch=0, total_opt_step=0)
-    return learning_rate.PolicySchedule(
-        history,
-        observation_metrics=observation_metrics,
-        include_controls_in_observation=False,
-        action_multipliers=action_multipliers,
-        control_configs=control_configs,
-        policy_and_value_model=policy_and_value_model,
-        policy_and_value_two_towers=False,
-        policy_dir=policy_dir,
-    )
-
-  def test_returns_start_lr_when_there_are_no_metrics(self):
-    history = trax_history.History()
-    start_lr = 1e-3
-    schedule = self._make_schedule(
-        history,
-        control_configs=(("learning_rate", start_lr, (1e-9, 1.0), False),),
-    )
-    self.assertEqual(schedule(0)["learning_rate"], start_lr)
-
-  def test_changes_lr_when_there_are_some_metrics(self):
-    history = trax_history.History()
-    history.append("eval", "metrics/accuracy", step=0, value=0.8)
-    history.append(
-        *online_tune.control_metric("learning_rate"), step=0, value=1e-4
-    )
-    schedule = self._make_schedule(
-        history,
-        control_configs=(("learning_rate", 1e-3, (1e-9, 1.0), False),),
-        observation_metrics=(("eval", "metrics/accuracy"),),
-        action_multipliers=(0.5, 2.0),
-    )
-    new_lr = schedule(123)["learning_rate"]
-    self.assertTrue(
-        onp.allclose(new_lr, 5e-5) or onp.allclose(new_lr, 2e-4)
-    )
-
-  def test_works_with_multiple_controls(self):
-    history = trax_history.History()
-    history.append("eval", "metrics/accuracy", step=0, value=0.8)
-    history.append(
-        *online_tune.control_metric("learning_rate"), step=0, value=1e-4
-    )
-    history.append(
-        *online_tune.control_metric("weight_decay_rate"), step=0, value=1e-5
-    )
-    schedule = self._make_schedule(
-        history,
-        observation_metrics=(("eval", "metrics/accuracy"),),
-        control_configs=(
-            ("learning_rate", 1e-3, (1e-9, 1.0), False),
-            ("weight_decay_rate", 1e-5, (1e-9, 1.0), False),
-        ),
-        action_multipliers=(1.0,),
-    )
-    new_controls = schedule(123)
-    self.assertIn("learning_rate", new_controls)
-    self.assertIn("weight_decay_rate", new_controls)
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/models/__init__.py b/tensor2tensor/trax/models/__init__.py
deleted file mode 100644
index ebf81e273..000000000
--- a/tensor2tensor/trax/models/__init__.py
+++ /dev/null
@@ -1,51 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Models defined in trax."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gin
-
-from tensor2tensor.trax.models import atari_cnn
-from tensor2tensor.trax.models import mlp
-from tensor2tensor.trax.models import neural_gpu
-from tensor2tensor.trax.models import resnet
-from tensor2tensor.trax.models import transformer
-from tensor2tensor.trax.models.research import position_lookup_transformer
-from tensor2tensor.trax.models.research import reformer
-
-
-# Ginify
-def model_configure(*args, **kwargs):
-  kwargs["module"] = "trax.models"
-  return gin.external_configurable(*args, **kwargs)
-
-
-# pylint: disable=invalid-name
-AtariCnn = model_configure(atari_cnn.AtariCnn)
-FrameStackMLP = model_configure(atari_cnn.FrameStackMLP)
-MLP = model_configure(mlp.MLP)
-NeuralGPU = model_configure(neural_gpu.NeuralGPU)
-PositionLookupTransformerLM = model_configure(
-    position_lookup_transformer.PositionLookupTransformerLM)
-ReformerLM = model_configure(reformer.ReformerLM)
-Resnet50 = model_configure(resnet.Resnet50)
-Transformer = model_configure(transformer.Transformer)
-TransformerDecoder = model_configure(transformer.TransformerDecoder)
-TransformerEncoder = model_configure(transformer.TransformerEncoder)
-TransformerLM = model_configure(transformer.TransformerLM)
-WideResnet = model_configure(resnet.WideResnet)
diff --git a/tensor2tensor/trax/models/atari_cnn.py b/tensor2tensor/trax/models/atari_cnn.py
deleted file mode 100644
index d50a295d9..000000000
--- a/tensor2tensor/trax/models/atari_cnn.py
+++ /dev/null
@@ -1,79 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 net for playing Atari games using PPO."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.trax import layers as tl
-
-
-def FrameStack(n_frames):
-  """Stacks a fixed number of frames along the dimension 1."""
-  # Input shape: (B, T, ..., C).
-  # Output shape: (B, T, ..., C * n_frames).
-  assert n_frames >= 1
-  if n_frames == 1:
-    return ()
-  return (
-      # Make n_frames copies of the input sequence.
-      [tl.Dup()] * (n_frames - 1),
-      # Shift copies to the right by [0, .., n_frames - 1] frames.
-      tl.Parallel(*map(_shift_right, range(n_frames))),
-      # Concatenate along the channel dimension.
-      tl.Concatenate(n_items=n_frames, axis=-1),
-  )
-
-
-def AtariCnn(n_frames=4, hidden_sizes=(32, 32), output_size=128, mode='train'):
-  """An Atari CNN."""
-  del mode
-
-  # TODO(jonni): Include link to paper?
-  # Input shape: (B, T, H, W, C)
-  # Output shape: (B, T, output_size)
-  return tl.Model(
-      tl.ToFloat(),
-      tl.Div(divisor=255.0),
-
-      # Set up n_frames successive game frames, concatenated on the last axis.
-      FrameStack(n_frames=n_frames),  # (B, T, H, W, 4C)
-
-      tl.Conv(hidden_sizes[0], (5, 5), (2, 2), 'SAME'),
-      tl.Relu(),
-      tl.Conv(hidden_sizes[1], (5, 5), (2, 2), 'SAME'),
-      tl.Relu(),
-      tl.Flatten(n_axes_to_keep=2),  # B, T and rest.
-      tl.Dense(output_size),
-      tl.Relu(),
-  )
-
-
-def FrameStackMLP(n_frames=4, hidden_sizes=(64,), output_size=64,
-                  mode='train'):
-  """MLP operating on a fixed number of last frames."""
-  del mode
-
-  return tl.Model(
-      FrameStack(n_frames=n_frames),
-      [[tl.Dense(d_hidden), tl.Relu()] for d_hidden in hidden_sizes],
-      tl.Dense(output_size),
-  )
-
-
-def _shift_right(n):  # pylint: disable=invalid-name
-  return [tl.ShiftRight()] * n
diff --git a/tensor2tensor/trax/models/atari_cnn_test.py b/tensor2tensor/trax/models/atari_cnn_test.py
deleted file mode 100644
index 5ee4d4cea..000000000
--- a/tensor2tensor/trax/models/atari_cnn_test.py
+++ /dev/null
@@ -1,65 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.models.atari_cnn."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import functools
-import operator as op
-import numpy as onp
-from tensor2tensor.trax.backend import random as jax_random
-from tensor2tensor.trax.models import atari_cnn
-from tensorflow import test
-
-
-class AtariCnnTest(test.TestCase):
-
-  def test_computes(self):
-    rng_key = jax_random.get_prng(0)
-    hidden_size = (4, 4)
-    output_size = 6
-    model = atari_cnn.AtariCnn(
-        hidden_sizes=hidden_size, output_size=output_size)
-    B, T, OBS = 2, 2, (28, 28, 3)  # pylint: disable=invalid-name
-    rng_key, key = jax_random.split(rng_key)
-    _, _ = model.initialize_once((1, 1) + OBS, onp.float32, key)
-    x = onp.arange(B * (T + 1) * functools.reduce(op.mul, OBS)).reshape(
-        B, T + 1, *OBS)
-    y = model(x)
-    self.assertEqual((B, T + 1, output_size), y.shape)
-
-
-class FrameStackMLPTest(test.TestCase):
-
-  def test_computes(self):
-    rng_key = jax_random.get_prng(0)
-    hidden_size = (4, 4)
-    output_size = 6
-    model = atari_cnn.FrameStackMLP(
-        hidden_sizes=hidden_size, output_size=output_size)
-    B, T, OBS = 2, 2, 3  # pylint: disable=invalid-name
-    rng_key, key = jax_random.split(rng_key)
-    _, _ = model.initialize_once((1, 1, OBS), onp.float32, key)
-    x = onp.arange(B * (T + 1) * OBS).reshape(
-        B, T + 1, OBS)
-    y = model(x)
-    self.assertEqual((B, T + 1, output_size), y.shape)
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/models/mlp.py b/tensor2tensor/trax/models/mlp.py
deleted file mode 100644
index f6f2a9575..000000000
--- a/tensor2tensor/trax/models/mlp.py
+++ /dev/null
@@ -1,38 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""MLP."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.trax import layers as tl
-
-
-def MLP(n_hidden_layers=2,
-        d_hidden=512,
-        activation_fn=tl.Relu,
-        n_output_classes=10,
-        mode="train"):
-  """A multi-layer feedforward (perceptron) network."""
-  del mode
-
-  return tl.Model(
-      tl.Flatten(),
-      [[tl.Dense(d_hidden), activation_fn()] for _ in range(n_hidden_layers)],
-      tl.Dense(n_output_classes),
-      tl.LogSoftmax(),
-  )
diff --git a/tensor2tensor/trax/models/mlp_test.py b/tensor2tensor/trax/models/mlp_test.py
deleted file mode 100644
index 2ba7568d0..000000000
--- a/tensor2tensor/trax/models/mlp_test.py
+++ /dev/null
@@ -1,39 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 MLP."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import absltest
-from tensor2tensor.trax import backend
-from tensor2tensor.trax import layers as tl
-from tensor2tensor.trax.models import mlp
-
-
-class MLPTest(absltest.TestCase):
-
-  def test_mlp_forward_shape(self):
-    """Run the MLP model forward and check output shape."""
-    input_shape = (3, 28, 28, 1)
-    model = mlp.MLP(d_hidden=32, n_output_classes=10)
-    final_shape = tl.check_shape_agreement(model, input_shape)
-    self.assertEqual((3, 10), final_shape)
-
-
-if __name__ == '__main__':
-  absltest.main()
diff --git a/tensor2tensor/trax/models/neural_gpu.py b/tensor2tensor/trax/models/neural_gpu.py
deleted file mode 100644
index b89c1f9b6..000000000
--- a/tensor2tensor/trax/models/neural_gpu.py
+++ /dev/null
@@ -1,82 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Implementation of the improved Neural GPU (NGPU)."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.trax import layers as tl
-from tensor2tensor.trax.backend import numpy as np
-
-
-# TODO(ddohan): Combinator to add saturation costs to loss
-def SaturationCost(x, limit=0.9):
-  return np.minimum(0, np.abs(x) - limit)
-
-
-@tl.layer()
-def DiagonalGate(x, params, **kwargs):
-  """Split channels in 3 parts. Shifts 1st and 3rd sections to left/right."""
-  del params
-  del kwargs
-  # x : [batch, 1, length, depth]
-  x = np.pad(
-      x, [(0, 0), (0, 0), (1, 1), (0, 0)], mode='constant', constant_values=0.0)
-  depth = x.shape[-1] // 3
-  assert 3 * depth == x.shape[-1], ('Depth must be divisible by 3', depth,
-                                    x.shape)
-  xs = [
-      x[:, :, :-2, :depth], x[:, :, 1:-1, depth:2 * depth],
-      x[:, :, 2:, 2 * depth:3 * depth]
-  ]
-  return np.concatenate(xs, axis=3)
-
-
-def ConvDiagonalGRU(units, kernel_size=(3, 3)):
-  """Build convolutional GRU with diagonal gating as in ImprovedNGPU."""
-
-  def BuildConv():
-    return tl.Conv(filters=units, kernel_size=kernel_size, padding='SAME')
-
-  return tl.GeneralGRUCell(
-      candidate_transform=BuildConv,
-      memory_transform_fn=DiagonalGate,
-      gate_nonlinearity=tl.HardSigmoid,
-      candidate_nonlinearity=tl.HardTanh)
-
-
-def NeuralGPU(d_feature=96, steps=16, vocab_size=2, mode='train'):
-  """Implementation of Neural GPU: https://arxiv.org/abs/1702.08727.
-
-  Args:
-    d_feature: Number of memory channels (dimensionality of feature embedding).
-    steps: Number of times depthwise recurrence steps.
-    vocab_size: Vocabulary size.
-    mode: Whether we are training or evaluating or doing inference.
-
-  Returns:
-    A NeuralGPU Stax model.
-  """
-  del mode
-
-  core = ConvDiagonalGRU(units=d_feature)
-  return tl.Model(
-      tl.Embedding(d_feature=d_feature, vocab_size=vocab_size),
-      [core] * steps,
-      tl.Dense(vocab_size),
-      tl.LogSoftmax(),
-  )
diff --git a/tensor2tensor/trax/models/neural_gpu_test.py b/tensor2tensor/trax/models/neural_gpu_test.py
deleted file mode 100644
index 8583376d3..000000000
--- a/tensor2tensor/trax/models/neural_gpu_test.py
+++ /dev/null
@@ -1,39 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 trax.models.neural_gpu."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import absltest
-from tensor2tensor.trax.layers import base
-from tensor2tensor.trax.models import neural_gpu
-
-
-class NeuralGPUTest(absltest.TestCase):
-
-  def test_ngpu(self):
-    vocab_size = 2
-    input_shape = [3, 5, 7]
-    model = neural_gpu.NeuralGPU(d_feature=30, steps=4, vocab_size=vocab_size)
-    final_shape = base.check_shape_agreement(
-        model, tuple(input_shape), integer_inputs=True)
-    self.assertEqual(tuple(input_shape + [vocab_size]), final_shape)
-
-
-if __name__ == '__main__':
-  absltest.main()
diff --git a/tensor2tensor/trax/models/research/__init__.py b/tensor2tensor/trax/models/research/__init__.py
deleted file mode 100644
index 4872e5d5d..000000000
--- a/tensor2tensor/trax/models/research/__init__.py
+++ /dev/null
@@ -1,15 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
diff --git a/tensor2tensor/trax/models/research/position_lookup_transformer.py b/tensor2tensor/trax/models/research/position_lookup_transformer.py
deleted file mode 100644
index 969812745..000000000
--- a/tensor2tensor/trax/models/research/position_lookup_transformer.py
+++ /dev/null
@@ -1,341 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Deep Lookups for Transformer Positions."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as onp
-
-from tensor2tensor.trax import layers as tl
-from tensor2tensor.trax.backend import numpy as np
-
-
-# pylint: disable=g-complex-comprehension
-# pylint: disable=no-value-for-parameter
-
-POS_VECTOR_SIZE = 32
-_ABSOLUTE_MAX_LEN = 10000
-_POSITIONS = onp.random.uniform(size=[_ABSOLUTE_MAX_LEN, POS_VECTOR_SIZE])
-
-
-def Dup2():
-  """Copy first 2 elements of the stack: (a, b, ...) -> (a, b, a, b, ...)."""
-  return [                              # Stack is (a, b, ...)
-      tl.Parallel(tl.Dup(), tl.Dup()),  # Stack is (a, a, b, b, ...)
-      tl.Parallel([], tl.Swap())        # Stack is (a, b, a, b, ...)
-  ]
-
-
-@tl.layer()
-def NewPositionalEncoding(x, positions=None, **kwargs):
-  """Implements new positional encoding."""
-  del kwargs
-  x_length = np.shape(x)[1]
-  pos = np.array(positions)[np.newaxis, :x_length, :]
-  pos += np.zeros((np.shape(x)[0], 1, 1))  # Broadcast on batch.
-  res = np.concatenate([x, pos], axis=2)
-  return res
-
-
-@tl.layer(n_inputs=1, n_outputs=2)
-def CutAtPosition(x, **unused_kwargs):
-  """Splits x into a pair (x[:position], position)."""
-  return tuple([x[:, :, :-POS_VECTOR_SIZE], x[:, :, -POS_VECTOR_SIZE:]])
-
-
-@tl.layer()
-def MixHeadsPos(x, h=8, **unused_kwargs):
-  """Mix x = (x0, p) into x0_h1, p, x0_h2, p, ...."""
-  head_size = (x.shape[2] - POS_VECTOR_SIZE) // h
-  p = x[:, :, -POS_VECTOR_SIZE:]
-  res, idx = [], 0
-  for _ in range(h):
-    res.append(x[:, :, idx:idx+head_size])
-    res.append(p)
-    idx += head_size
-  return np.concatenate(res, axis=-1)
-
-
-@tl.layer()
-def CombineHeadsPos(x, h=8, **unused_kwargs):
-  """Mix x = (x0, p0, ..., xH, pH) into x0, ...., xH, p_combined.
-
-  The positions are added as vectors.
-
-  Args:
-    x: input vector, concatenated (x0, p0, ..., xH, pH).
-    h: number of heads.
-
-  Returns:
-    the vector with combined positions.
-  """
-  head_size = int((x.shape[2] / h) - POS_VECTOR_SIZE)
-  res, positions, idx = [], [], 0
-  for _ in range(h):
-    res.append(x[:, :, idx:idx+head_size])
-    idx += head_size
-    positions.append(x[:, :, idx:idx+POS_VECTOR_SIZE])
-    idx += POS_VECTOR_SIZE
-  combined_position = sum(positions)
-  res.append(combined_position)
-  return np.concatenate(res, axis=-1)
-
-
-@tl.layer()
-def CopyHeadsPos(x, h=8, **unused_kwargs):
-  """Mix x = (x, p) into x_h1, p_h1, x_h2, p_h2, ...."""
-  head_size = (x.shape[2] - h*POS_VECTOR_SIZE) // h
-  p = x[:, :, -h*POS_VECTOR_SIZE:]
-  res, idx = [], 0
-  for i in range(h):
-    res.append(x[:, :, idx:idx+head_size])
-    res.append(p[:, :, i*POS_VECTOR_SIZE:(i+1)*POS_VECTOR_SIZE])
-    idx += head_size
-  return np.concatenate(res, axis=-1)
-
-
-def DeepFlatten(xs):
-  for x in xs:
-    if isinstance(x, (list, tuple)):
-      for y in DeepFlatten(x):
-        yield y
-    else:
-      yield x
-
-
-def PreservePosition(layer):
-  """Execute layer without position but preserve it in parallel."""
-  return tl.Serial(
-      CutAtPosition(),
-      layer,
-      tl.Concatenate(n_items=2)
-  )
-
-
-def ApplyAndQueryPositions(layer, pos):
-  """Execute layer without position and pos-layers on positions.
-
-  This takes an embedding including position x = (emb, p), and
-  outputs layer(emb).pos1(x, p).....layer(emb).posn(x, p)
-  where pos=[pos1...posn].
-
-  Args:
-    layer: layer to be executed without position information.
-    pos: list of layers to be applied to positions.
-
-  Returns:
-    the result of this application.
-  """
-  n_heads = len(pos)
-  return tl.Serial(
-      tl.Dup(),                    # (x, x)
-      CutAtPosition(),          # (x_content, x_position, x)
-      tl.Parallel([], tl.Swap()),  # (x_content, x, x_position)
-      [tl.Parallel([], Dup2()) for _ in range(n_heads - 1)],
-      # Now the stack is x_content, (x, x_position) * n_heads.
-      tl.Parallel(*([layer] + pos)),
-      tl.Concatenate(n_items=n_heads + 1)
-  )
-
-
-@tl.layer()
-def QueryPositionKV(x, keys=None, values=None, binary=False, **unused_kwargs):
-  """Query a table with a position vector."""
-  if keys is None:
-    return x
-  k = np.array(keys)
-  v = np.array(values)
-  q = x
-  if binary:
-    q = np.concatenate([x, x], axis=-1)
-  return tl.DotProductAttention(q, k, v, None, None, None, None)
-
-
-def LearnedQP(keys=None, values=None, binary=False):
-  """Get (query, pos), make learned weight of qeury and return with pos."""
-  return tl.Parallel(
-      tl.Dense(1),
-      QueryPositionKV(keys=keys, values=values, binary=binary),
-  )
-
-
-@tl.layer(n_inputs=10, n_outputs=1)
-def Softmax5Branches(x_list, n_branches=2, **unused_kwargs):
-  """Softmax xs.
-
-  The input xs is a list of embeddings and weights of the form
-  w_1 e_1 .... w_n e_n (followed by optional rest that is preserved).
-
-  Args:
-    x_list: the input weights and embeddings.
-    n_branches: what part of the list to use.
-
-  Returns:
-    softmax(w) * e for the joint weights w and embeddings e.
-  """
-  assert n_branches == 5
-  softmax_activations = [x_list[2*i] for i in range(n_branches)]
-  max_sa = softmax_activations[0]
-  for x in softmax_activations:
-    max_sa = np.maximum(max_sa, x)
-  softmax_activations = [x - max_sa for x in softmax_activations]
-  softmax_activations = [np.exp(x) for x in softmax_activations]
-  sum_sa = sum(softmax_activations)
-  softmax_activations = [x / sum_sa for x in softmax_activations]
-  res = sum([x_list[2*i+1] * softmax_activations[i] for i in range(n_branches)])
-  return res
-
-
-def SumLearnedPick(positions):
-  """Get a pair (vec, pos) and pick new pos."""
-  succ_keys = positions[:-1, :]
-  succ_values = positions[1:, :]
-  subtract_1_keys = positions[1:, :]
-  subtract_1_values = positions[:-1, :]
-  l = int(positions.shape[0]) // 2
-  add_keys = np.array([np.concatenate([positions[i, :], positions[j, :]])
-                       for i in range(l) for j in range(l)])
-  add_values = np.array([positions[i + j, :]
-                         for i in range(l) for j in range(l)])
-  # TODO(lukaszkaiser): try this below: "for j in range(i) for i in range(2*l)"
-  sub_keys = np.array([np.concatenate([positions[i, :], positions[j, :]])
-                       for j in range(l) for i in range(l)])
-  sub_values = np.array([positions[max(i - j, 0), :]
-                         for j in range(l) for i in range(l)])
-  return tl.Serial(
-      Dup2(), Dup2(), Dup2(), Dup2(),
-      tl.Parallel(
-          LearnedQP(),
-          LearnedQP(keys=succ_keys, values=succ_values),
-          LearnedQP(keys=subtract_1_keys, values=subtract_1_values),
-          LearnedQP(keys=add_keys, values=add_values, binary=True),
-          LearnedQP(keys=sub_keys, values=sub_values, binary=True),
-      ),
-      Softmax5Branches(n_branches=5)
-  )
-
-
-def AttentionPosition(positions, d_model, n_heads=8, dropout=0.0,
-                      mode='train'):
-  """Transformer-style multi-headed attention."""
-  return tl.Serial(
-      tl.Dup(),
-      tl.Dup(),
-      tl.Parallel(
-          ApplyAndQueryPositions(tl.Dense(d_model),
-                                 pos=[SumLearnedPick(positions)
-                                      for _ in range(n_heads)]),
-          PreservePosition(tl.Dense(d_model)),
-          PreservePosition(tl.Dense(d_model)),
-      ),
-      tl.Parallel(
-          CopyHeadsPos(h=n_heads),
-          MixHeadsPos(h=n_heads),
-          MixHeadsPos(h=n_heads),
-      ),
-      tl.PureAttention(d_model=d_model, n_heads=n_heads, dropout=dropout,
-                       mode=mode),
-      tl.Parallel([], tl.Drop()),  # Drop the mask.
-      CombineHeadsPos(h=n_heads),
-      PreservePosition(tl.Dense(d_model)),
-  )
-
-
-def ResidualFeedForward(d_model,
-                        d_ff,
-                        dropout,
-                        mode):
-  """Residual feed-forward layer with normalization at start."""
-  stack = tl.Serial(
-      tl.LayerNorm(),
-      tl.Dense(d_ff),
-      tl.Relu(),
-      tl.Dropout(rate=dropout, mode=mode),
-      tl.Dense(d_model),
-      tl.Dropout(rate=dropout, mode=mode)
-  )
-  return tl.Residual(PreservePosition(stack))
-
-
-def DecoderLayer(positions,
-                 d_model,
-                 d_ff,
-                 n_heads,
-                 dropout,
-                 mode):
-  """Transformer decoder layer.
-
-  Args:
-    positions: random vectors for positions
-    d_model: int:  depth of embedding
-    d_ff: int: depth of feed-forward layer
-    n_heads: int: number of attention heads
-    dropout: float: dropout rate (how much to drop out)
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    the layer.
-  """
-  return [
-      tl.Residual(  # Self-attention block.
-          PreservePosition(tl.LayerNorm()),
-          tl.Dup(),
-          tl.Parallel([],  # activation for (q, k, v)
-                      tl.CausalMask(axis=-2)),  # attention mask
-          AttentionPosition(positions, d_model, n_heads=n_heads,
-                            dropout=dropout, mode=mode),
-          PreservePosition(tl.Dropout(rate=dropout, mode=mode))
-      ),
-      ResidualFeedForward(d_model, d_ff, dropout, mode=mode)
-  ]
-
-
-def PositionLookupTransformerLM(vocab_size=128,
-                                d_model=256,
-                                d_ff=512,
-                                n_layers=3,
-                                n_heads=4,
-                                dropout=0.1,
-                                max_len=100,
-                                mode='train'):
-  """Transformer language model (only uses the decoder part of Transformer).
-
-  Args:
-    vocab_size: int: vocab size
-    d_model: int:  depth of embedding
-    d_ff: int: depth of feed-forward layer
-    n_layers: int: number of layers
-    n_heads: int: number of attention heads
-    dropout: float: dropout rate (how much to drop out)
-    max_len: maximal length
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    the layer.
-  """
-  positions = _POSITIONS[:max_len, :]
-  return tl.Serial(
-      tl.ShiftRight(),
-      tl.Embedding(d_model, vocab_size),
-      tl.Dropout(rate=dropout, mode=mode),
-      NewPositionalEncoding(positions=positions),
-      [DecoderLayer(positions, d_model, d_ff, n_heads, dropout, mode)
-       for _ in range(n_layers)],
-      PreservePosition(tl.LayerNorm()),
-      tl.Dense(vocab_size),
-      tl.LogSoftmax()
-  )
diff --git a/tensor2tensor/trax/models/research/reformer.py b/tensor2tensor/trax/models/research/reformer.py
deleted file mode 100644
index 913343e67..000000000
--- a/tensor2tensor/trax/models/research/reformer.py
+++ /dev/null
@@ -1,522 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Transformer Models."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import jax
-
-from tensor2tensor.trax import backend
-from tensor2tensor.trax import layers as tl
-from tensor2tensor.trax.layers.combinators import _pop_rng_and_split
-
-
-# Layers are always CamelCase, but functions in general are snake_case
-# pylint: disable=invalid-name
-
-
-class Map(tl.Layer):
-  """Combinator for applying a layer to a list or tuple."""
-
-  def __init__(self, layer, n_sections=1, check_shapes=True):
-    """Initialize the combinator.
-
-    Args:
-      layer: a layer to apply to each element.
-      n_sections: how many sections to map to (default: 1).
-      check_shapes: whether to check that shapes are identical (default: true).
-
-    Returns:
-      A new layer representing mapping layer to all elements of the input.
-    """
-    super(Map, self).__init__(n_inputs=n_sections, n_outputs=n_sections)
-    if layer is None or isinstance(layer, (list, tuple)):
-      layer = tl.Serial(layer)
-    self._layer = layer
-    # Generally a Map should be applied to lists where all elements have
-    # the same shape -- because self._layer will only be initialized once
-    # and it could have different parameters for different shapes. But there
-    # are valid cases -- e.g., when self._layer has no parameters -- where we
-    # can apply Map to different shapes -- set check_shapes=False in such cases.
-    self._check_shapes = check_shapes
-    self._n_sections = n_sections
-
-  def forward(self, inputs, params=(), state=(), **kwargs):
-    rngs = _pop_rng_and_split(kwargs, len(inputs))
-    results = [self._layer(x, params=params, state=state, rng=r, **kwargs)
-               for x, r in zip(inputs, rngs)]
-    # TODO(kitaev): think about how to merge state across copies in the map.
-    return tuple(results), self._layer.state
-
-  def new_params_and_state(self, input_shape, input_dtype, rng):
-    first_shape = input_shape[0]
-    if self._check_shapes:
-      for shape in input_shape:
-        if shape != first_shape:
-          raise ValueError('Map layer can only be applied to list of elements '
-                           'with the same shapes. Shapes: %s' % str(shape))
-    return self._layer.initialize_once(first_shape, input_dtype[0], rng)
-
-  @tl.Layer.params.setter
-  def params(self, params):
-    self._params = params
-    assert len(params) == 1
-    self._layer.params = params[0]
-
-  @tl.Layer.state.setter
-  def state(self, state):
-    self._state = state
-    assert len(state) == 1
-    self._layer.state = state[0]
-
-
-@tl.layer()
-def BroadcastedDropout(x, params, rate=0.0, mode='train', broadcast_dims=(-2,),
-                       rng=None, **kwargs):
-  """Dropout, with broadcasting to save memory."""
-  del params, kwargs
-  if rng is None:
-    raise ValueError('BroadcastedDropout requires rng kwarg.')
-  if rate >= 1.0:
-    raise ValueError('Dropout rate (%f) must be lower than 1.' % rate)
-  if mode == 'train' and rate > 0.0:
-    noise_shape = list(x.shape)
-    for dim in broadcast_dims:
-      noise_shape[dim] = 1
-    keep_prob = jax.lax.tie_in(rng, 1.0 - rate)
-    keep = backend.random.bernoulli(rng, keep_prob, tuple(noise_shape))
-    multiplier = keep.astype(x.dtype) / jax.lax.tie_in(keep, keep_prob)
-    return x * multiplier
-  else:
-    return x
-
-
-def FeedForward(d_model, d_ff, dropout, mode):
-  """Feed-forward block with layer normalization at start."""
-  return [
-      tl.LayerNorm(),
-      tl.Dense(d_ff),
-      BroadcastedDropout(rate=dropout, mode=mode),  # pylint: disable=no-value-for-parameter
-      tl.Relu(),
-      tl.Dense(d_model),
-      BroadcastedDropout(rate=dropout, mode=mode),  # pylint: disable=no-value-for-parameter
-  ]
-
-
-class SplitForOutput(tl.ReversibleLayer):
-  """Splits activations into sections (for use right before the output layer).
-
-  After the reversible portion of the network, there is a final output portion
-  that's non-reversible (which at minimum includes normalization, output
-  projection, and log-softmax). The output portion needs to operate on chunks
-  of the sequence to avoid running out of memory for large vocabulary sizes.
-
-  This layer concatenates the two subparts of the activations along the feature
-  dimension, and then splits into chunks along the time dimension. We implement
-  it is a subclass of tl.ReversibleLayer because we want to ensure that multiple
-  copies of the activations don't exist simultaneously except in the middle of a
-  memory copy operation.
-  """
-
-  def __init__(self, n_sections=2, axis=-2):
-    super(SplitForOutput, self).__init__(n_inputs=2, n_outputs=n_sections)
-    self._n_sections = n_sections
-    self._axis = axis
-
-  def forward(self, inputs, params=(), state=(), **kwargs):
-    del params, kwargs
-    x1, x2 = inputs
-
-    x1_split = backend.numpy.split(x1, self._n_sections, self._axis)
-    x2_split = backend.numpy.split(x2, self._n_sections, self._axis)
-
-    res = [backend.numpy.concatenate(ys, -1) for ys in zip(x1_split, x2_split)]
-    return tuple(res), state
-
-  def reverse(self, output, params=(), state=(), **kwargs):
-    del params, kwargs
-
-    x1_split = []
-    x2_split = []
-    for y in output:
-      y1, y2 = backend.numpy.split(y, 2, -1)
-      x1_split.append(y1)
-      x2_split.append(y2)
-
-    x1 = backend.numpy.concatenate(x1_split, self._axis)
-    x2 = backend.numpy.concatenate(x2_split, self._axis)
-
-    return (x1, x2)
-
-  def reverse_and_grad(self, output, ct, params=(), state=(), **kwargs):
-    del params, kwargs
-    return self.reverse(output), (self.reverse(ct), ())
-
-
-@tl.layer()
-def Chunk(x, params, n_sections=2, **kwargs):
-  del params, kwargs
-  assert x.shape[1] % n_sections == 0
-  return backend.numpy.reshape(x, (
-      x.shape[0] * n_sections,
-      x.shape[1] // n_sections,
-      ) + x.shape[2:])
-
-
-@tl.layer()
-def Unchunk(x, params, n_sections=2, **kwargs):
-  del params, kwargs
-  assert x.shape[0] % n_sections == 0
-  return backend.numpy.reshape(x, (
-      x.shape[0] // n_sections,
-      x.shape[1] * n_sections,
-      ) + x.shape[2:])
-
-
-class ReversibleHalfResidual(tl.ReversibleLayer, tl.Serial):
-  """Half of a RevNet-style residual (only updates part of the hidden state)."""
-
-  def __init__(self, residual_layers):
-    self.compute_residual = tl.Serial([
-        # (x1_or_y1, x2) -> (x2, x1_or_y1, x2)
-        tl.Parallel([], tl.Dup()),
-        tl.Swap(),
-        tl.Parallel(residual_layers, [], []),
-    ])
-
-    layers = [
-        self.compute_residual,
-        tl.Parallel(tl.Add(), [])
-    ]
-    super(ReversibleHalfResidual, self).__init__(layers)
-
-    self.subtract_top = tl.Parallel(tl.SubtractTop(), [])
-    self.reverse_layers = [self.compute_residual, self.subtract_top]
-
-  def reverse(self, output, params=(), state=(), **kwargs):
-    reconstructed_x = output
-    rng = kwargs.pop('rng', None)
-    rngs = (None,) * self._n_layers
-    if rng is not None:
-      rngs = backend.random.split(rng, self._n_layers)
-    # Note that self.sublayers aligns exactly with self.reverse_layers in
-    # terms of parameter and rng usage, so no re-ordering is required.
-    for layer, p, s, rng in zip(self.reverse_layers, params, state, rngs):
-      reconstructed_x = layer(reconstructed_x, params=p, state=s, rng=rng,
-                              **kwargs)
-    return reconstructed_x
-
-  def reverse_and_grad(self, output, ct, params=(), state=(), **kwargs):
-    rng = kwargs.pop('rng', None)
-    rngs = (None,) * self._n_layers
-    if rng is not None:
-      rngs = backend.random.split(rng, self._n_layers)
-
-    def call_compute_residual(x, params):
-      res = self.compute_residual(x, params=params, state=state[0], rng=rngs[0],
-                                  **kwargs)
-      return res
-
-    assert len(ct) == 2
-    ct = ((ct[0], ct[0], ct[1]))
-
-    stack_with_residual, vjpfun = jax.vjp(
-        call_compute_residual, output, params[0])
-    reconstructed_x = self.subtract_top(
-        stack_with_residual, params=params[-1], state=state[-1], rng=rngs[-1],
-        **kwargs)
-
-    x_ct, residual_params_ct = vjpfun(ct)
-    assert not jax.tree_util.tree_leaves(params[-1])
-    add_top_params_ct = params[-1]
-    return reconstructed_x, (x_ct, [residual_params_ct, add_top_params_ct])
-
-
-class ApplyAttentionWrapper(tl.Parallel):
-  """Like tl.Parallel(attention, [], []) but implements forward_and_backward."""
-
-  def __init__(self, attention):
-    assert hasattr(attention, 'forward_and_backward')
-    super(ApplyAttentionWrapper, self).__init__(attention, [], [])
-    self.attention = attention
-
-  def forward_and_backward(self, inputs, ct, rng=None, **kwargs):
-    # Simultaneous forward pass and backprop through the attention mechanism.
-    qkv = inputs[:3]
-    passthrough = inputs[3:]
-    out_ct = ct[0]
-    passthrough_ct = ct[1:]
-    if rng is not None:
-      # Adjust RNG to match the forward pass.
-      rng = backend.random.split(rng, self._n_layers)[0]
-
-    out, qkv_ct = self.attention.forward_and_backward(
-        qkv, out_ct, rng=rng, **kwargs)
-    return (out,) + passthrough, qkv_ct + passthrough_ct
-
-
-class ReversibleAttentionHalfResidual(tl.ReversibleLayer, tl.Serial):
-  """Half of a RevNet-style residual that performs attention.
-
-  If inputs are (x1, x2), then outputs are (x1 + z, x2) where:
-  z = post_attention(attention(pre_attention(x1)))
-
-  Other than an efficiency optimization, this layer is equivalent to
-  ReversibleHalfResidual([pre_attention, attention, post_attention]).
-
-  The post_attention layers must be linear in their input (typically they will
-  consists of reshaping and dense linear layers), which allows the following
-  optimization. We can back-propagate the gradient signal from the output of
-  ReversibleAttentionHalfResidual to the output of the "attention" portion based
-  only on the network parameters. Then, attention.forward_and_backward can be
-  used to recover the output of the "attention" portion while simultaneously
-  performing the backward pass, which allows shared computation between the two
-  directions.
-  """
-
-  def __init__(self, pre_attention, attention, post_attention):
-    self.pre_attention = tl.Serial([
-        # (x1_or_y1, x2) -> (x2, x1_or_y1, x2)
-        tl.Parallel([], tl.Dup()),
-        tl.Swap(),
-        tl.Parallel(pre_attention, [], []),
-    ])
-    assert hasattr(attention, 'forward_and_backward')
-    self.attention = ApplyAttentionWrapper(attention)
-    self.post_attention = tl.Parallel(post_attention, [], [])
-
-    layers = [
-        self.pre_attention,
-        self.attention,
-        self.post_attention,
-        tl.Parallel(tl.Add(), []),
-    ]
-    super(ReversibleAttentionHalfResidual, self).__init__(layers)
-
-    self.subtract_top = tl.Parallel(tl.SubtractTop(), [])
-    self.reverse_layers = [
-        self.pre_attention,
-        self.attention,
-        self.post_attention,
-        self.subtract_top,
-    ]
-
-  def reverse(self, output, params=(), state=(), **kwargs):
-    rng = kwargs.pop('rng', None)
-    rngs = (None,) * self._n_layers
-    if rng is not None:
-      rngs = backend.random.split(rng, self._n_layers)
-
-    reconstructed_x = output
-    # Note that self.sublayers aligns exactly with self.reverse_layers in
-    # terms of parameter and rng usage, so no re-ordering is required.
-    for layer, p, s, rng in zip(self.reverse_layers, params, state, rngs):
-      reconstructed_x = layer.reverse(reconstructed_x, params=p, state=s,
-                                      rng=rng, **kwargs)
-    return reconstructed_x
-
-  def reverse_and_grad(self, output, ct, params=(), state=(), **kwargs):
-    rng = kwargs.pop('rng', None)
-    rngs = (None,) * self._n_layers
-    if rng is not None:
-      rngs = backend.random.split(rng, self._n_layers)
-
-    # Forward pass through self.pre_attention, while preparing for
-    # later backprop.
-    def call_pre_attention(x, params):
-      res = self.pre_attention(x, params=params, state=state[0], rng=rngs[0],
-                               **kwargs)
-      return res
-    stack, pre_attention_vjpfun = jax.vjp(call_pre_attention, output, params[0])
-
-    # Backprop through adding the residual
-    assert len(ct) == 2
-    ct = saved_ct = (ct[0], ct[0], ct[1])
-
-    # Backprop through self.post_attention with respect to the inputs only
-    def call_post_attention(x):
-      res = self.post_attention(x, params=params[2], state=state[2],
-                                rng=rngs[2], **kwargs)
-      return res
-    # Note: these are *not* the actual inputs to self.post_attention.
-    # If self.post_attention is not linear, we will get incorrect gradients.
-    dummy_inputs = (stack[-3], stack[-2], stack[-1])
-    _, post_attention_vjpfun = jax.vjp(call_post_attention, dummy_inputs)
-    (ct,) = post_attention_vjpfun(ct)
-
-    # Simultaneous forward pass and backprop through the attention mechanism
-    stack, ct = self.attention.forward_and_backward(stack, ct, rng=rngs[1],
-                                                    **kwargs)
-    assert not jax.tree_util.tree_leaves(params[1])
-    attention_params_ct = params[1]  # This is valid when params is empty.
-
-    # Backprop through self.pre_attention
-    x_ct, pre_attention_params_ct = pre_attention_vjpfun(ct)
-
-    # Forward pass for self.post_attention, and backprop with respect to the
-    # parameters only
-    def call_post_attention2(params):
-      res = self.post_attention(stack, params=params, state=state[2],
-                                rng=rngs[2], **kwargs)
-      return res
-    stack, post_attention_vjpfun = jax.vjp(call_post_attention2, params[2])
-    (post_attention_params_ct,) = post_attention_vjpfun(saved_ct)
-
-    # Forward pass through subtracting the residual
-    reconstructed_x = self.subtract_top(
-        stack, params=params[-1], state=state[-1], rng=rngs[-1], **kwargs)
-
-    assert not jax.tree_util.tree_leaves(params[-1])
-    add_top_params_ct = params[-1]
-    params_ct = [
-        pre_attention_params_ct,
-        attention_params_ct,
-        post_attention_params_ct,
-        add_top_params_ct,
-    ]
-
-    return reconstructed_x, (x_ct, params_ct)
-
-
-def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
-                 n_heads, n_attention_chunks, attention_type,
-                 dropout, share_qk, mode):
-  """Reversible transformer decoder layer.
-
-  Args:
-    d_model: int:  depth of embedding
-    d_ff: int: depth of feed-forward layer
-    d_attention_key: int: depth of key vector for each attention head
-    d_attention_value: int: depth of value vector for each attention head
-    n_heads: int: number of attention heads
-    n_attention_chunks: int: number of chunks for attention
-    attention_type: subclass of tl.BaseCausalAttention: attention class to use
-    dropout: float: dropout rate (how much to drop out)
-    share_qk: string, whether to share queries and keys
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    the layer.
-  """
-  if share_qk:
-    pre_attention = [
-        Chunk(n_sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
-        tl.LayerNorm(),
-        tl.Dup(),
-        tl.Parallel(
-            tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
-            tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
-        ),
-        tl.Dup(),
-    ]
-  else:
-    pre_attention = [
-        Chunk(n_sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
-        tl.LayerNorm(),
-        tl.Dup(), tl.Dup(),
-        tl.Parallel(
-            tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
-            tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key),
-            tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value),
-        ),
-    ]
-
-  attention = attention_type(mode=mode)
-
-  # ReversibleAttentionHalfResidual requires that post_attention be linear in
-  # its input (so the backward pass can be computed without knowing the input)
-  post_attention = [
-      tl.ComputeAttentionOutput(n_heads=n_heads, d_model=d_model),
-      Unchunk(n_sections=n_attention_chunks),  # pylint: disable=no-value-for-parameter
-      BroadcastedDropout(rate=dropout, mode=mode),  # pylint: disable=no-value-for-parameter
-  ]
-
-  feed_forward = [
-      FeedForward(d_model, d_ff, dropout, mode=mode),
-  ]
-  return [
-      ReversibleAttentionHalfResidual(pre_attention, attention, post_attention),
-      tl.ReversibleSwap(),
-      ReversibleHalfResidual(feed_forward),
-      tl.ReversibleSwap(),
-  ]
-
-
-def ReformerLM(vocab_size,
-               d_model=512,
-               d_ff=2048,
-               d_attention_key=64,
-               d_attention_value=64,
-               n_layers=6,
-               n_heads=8,
-               dropout=0.1,
-               max_len=2048,
-               n_chunks=32,
-               n_attention_chunks=8,
-               attention_type=tl.DotProductCausalAttention,
-               share_qk=False,
-               mode='train'):
-  """Reversible transformer language model (only uses a decoder, no encoder).
-
-  Args:
-    vocab_size: int: vocab size
-    d_model: int:  depth of *each half* of the two-part features
-    d_ff: int: depth of feed-forward layer
-    d_attention_key: int: depth of key vector for each attention head
-    d_attention_value: int: depth of value vector for each attention head
-    n_layers: int: number of decoder layers
-    n_heads: int: number of attention heads
-    dropout: float: dropout rate (how much to drop out)
-    max_len: int: maximum symbol length for positional encoding
-    n_chunks: int: number of chunks (must match input pipeline)
-    n_attention_chunks: int: number of chunks for attention
-    attention_type: class: attention class to use, such as DotProductAttention.
-    share_qk: bool, whether to share queries and keys.
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    the layer.
-  """
-  positional_embedder = [
-      tl.Embedding(d_model, vocab_size),
-      BroadcastedDropout(rate=dropout, mode=mode),  # pylint: disable=no-value-for-parameter
-      tl.PositionalEncoding(max_len=max_len),
-  ]
-  return tl.Model(
-      tl.Concatenate(n_items=n_chunks),
-      tl.ShiftRight(),
-      positional_embedder,
-      tl.Dup(),
-      tl.ReversibleSerial([
-          # pylint: disable=g-complex-comprehension
-          DecoderBlock(d_model, d_ff,
-                       d_attention_key, d_attention_value, n_heads,
-                       n_attention_chunks, attention_type,
-                       dropout, share_qk, mode)
-          for _ in range(n_layers)
-      ] + [
-          SplitForOutput(n_sections=n_chunks, axis=-2),  # pylint: disable=no-value-for-parameter
-      ]),
-      Map([
-          # TODO(kitaev): Test whether dropout should go before or after the
-          # LayerNorm, and whether dropout broadcasting is needed here.
-          tl.LayerNorm(),
-          BroadcastedDropout(rate=dropout, mode=mode),  # pylint: disable=no-value-for-parameter
-          tl.Dense(vocab_size),
-          tl.LogSoftmax(),
-      ], n_sections=n_chunks),
-  )
diff --git a/tensor2tensor/trax/models/research/reformer_test.py b/tensor2tensor/trax/models/research/reformer_test.py
deleted file mode 100644
index 0a8bcdad8..000000000
--- a/tensor2tensor/trax/models/research/reformer_test.py
+++ /dev/null
@@ -1,104 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 Transformer-Revnet models."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import absltest
-from absl.testing import parameterized
-import jax
-import numpy as onp
-
-from tensor2tensor.trax import backend
-from tensor2tensor.trax import layers as tl
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.models.research import reformer
-
-
-class PoisonOnRNGMismatchAttention(tl.BaseCausalAttention):
-  """Fills gradients with NaNs if reverse rng does not match forward rng."""
-
-  # pylint: disable=protected-access
-  def forward_and_backward(self, inputs, ct, rng=None, **kwargs):
-    assert backend.get_name() == 'jax', (
-        'JAX backend is required to use forward_and_backward.')
-
-    if ct is not None and tl.Layer._STASH_OUT is not None:
-      recovered_rng = tl.Layer._STASH_OUT.pop(self)
-      is_same = (rng[0] == recovered_rng[0]) & (rng[1] == recovered_rng[1])
-      is_same = is_same.astype(np.float32)
-      # Divides by zero if rngs are not the same, which results in NaNs.
-      inputs = (inputs[0] / is_same, inputs[1] / is_same, inputs[2] / is_same)
-
-    def _do_forward(x):  # pylint: disable=invalid-name
-      res, _ = self.forward(x, rng=rng, **kwargs)
-      return res
-    output, vjpfun = jax.vjp(_do_forward, inputs)
-    return output, vjpfun(ct)[0]
-
-  def forward(self, inputs, params=(), state=(), rng=None, **kwargs):
-    if tl.Layer._STASH_IN is not None:
-      tl.Layer._STASH_IN[self] = rng
-    return inputs[2], state
-  # pylint: enable=protected-access
-
-
-class ReformerTest(parameterized.TestCase):
-
-  def test_reformer_lm_forward_shape(self):
-    """Run the ReformerLM forward and check output shape."""
-    vocab_size = 16
-    input_shape = ((1, 8), (1, 8))
-    model = reformer.ReformerLM(
-        vocab_size, d_model=32, d_ff=64,
-        d_attention_key=16, d_attention_value=16, n_layers=1, n_heads=2,
-        max_len=16, n_chunks=2, n_attention_chunks=1)
-    final_shape = tl.check_shape_agreement(
-        model, tuple(input_shape), integer_inputs=True)
-    self.assertEqual(((1, 8, 16), (1, 8, 16)), final_shape)
-
-  def test_reformer_rng_consistency(self):
-    with backend.use_backend('jax'):
-      vocab_size = 16
-      batch_size = 1
-      input_shape = ((batch_size, 8), (batch_size, 8))
-      model = reformer.ReformerLM(
-          vocab_size, d_model=32, d_ff=64,
-          d_attention_key=16, d_attention_value=16, n_layers=1, n_heads=2,
-          max_len=16, n_chunks=2, n_attention_chunks=1, mode='train',
-          attention_type=PoisonOnRNGMismatchAttention)
-
-      rng = backend.random.get_prng(0)
-      params, state = model.initialize_once(
-          input_shape, (np.int32, np.int32), rng)
-
-      def dummy_loss_fn(params):
-        inputs = (np.zeros(input_shape[0], dtype=np.int32),) * 2
-        output = model(inputs, params=params, state=state, rng=rng)
-        dummy_loss = backend.numpy.sum(output[0])
-        return dummy_loss
-
-      grad_fn = backend.grad(dummy_loss_fn)
-      grads = grad_fn(params)
-      # PoisonOnRNGMismatchAttention uses NaNs to signal an rng mismatch.
-      for grad in jax.tree_util.tree_leaves(grads):
-        assert onp.all(onp.isfinite(grad))
-
-
-if __name__ == '__main__':
-  absltest.main()
diff --git a/tensor2tensor/trax/models/resnet.py b/tensor2tensor/trax/models/resnet.py
deleted file mode 100644
index a9c644cc9..000000000
--- a/tensor2tensor/trax/models/resnet.py
+++ /dev/null
@@ -1,169 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""ResNet."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.trax import layers as tl
-
-
-def ConvBlock(kernel_size, filters, strides, mode='train'):
-  """ResNet convolutional striding block."""
-  # TODO(jonni): Use good defaults so Resnet50 code is cleaner / less redundant.
-  ks = kernel_size
-  filters1, filters2, filters3 = filters
-  main = [
-      tl.Conv(filters1, (1, 1), strides),
-      tl.BatchNorm(mode=mode),
-      tl.Relu(),
-      tl.Conv(filters2, (ks, ks), padding='SAME'),
-      tl.BatchNorm(mode=mode),
-      tl.Relu(),
-      tl.Conv(filters3, (1, 1)),
-      tl.BatchNorm(mode=mode),
-  ]
-  shortcut = [
-      tl.Conv(filters3, (1, 1), strides),
-      tl.BatchNorm(mode=mode),
-  ]
-  return [
-      tl.Residual(main, shortcut=shortcut),
-      tl.Relu(),
-  ]
-
-
-def IdentityBlock(kernel_size, filters, mode='train'):
-  """ResNet identical size block."""
-  # TODO(jonni): Use good defaults so Resnet50 code is cleaner / less redundant.
-  ks = kernel_size
-  filters1, filters2, filters3 = filters
-  main = [
-      tl.Conv(filters1, (1, 1)),
-      tl.BatchNorm(mode=mode),
-      tl.Relu(),
-      tl.Conv(filters2, (ks, ks), padding='SAME'),
-      tl.BatchNorm(mode=mode),
-      tl.Relu(),
-      tl.Conv(filters3, (1, 1)),
-      tl.BatchNorm(mode=mode),
-  ]
-  return [
-      tl.Residual(main),
-      tl.Relu(),
-  ]
-
-
-def Resnet50(d_hidden=64, n_output_classes=1001, mode='train'):
-  """ResNet.
-
-  Args:
-    d_hidden: Dimensionality of the first hidden layer (multiplied later).
-    n_output_classes: Number of distinct output classes.
-    mode: Whether we are training or evaluating or doing inference.
-
-  Returns:
-    The list of layers comprising a ResNet model with the given parameters.
-  """
-  return tl.Model(
-      tl.ToFloat(),
-      tl.Conv(d_hidden, (7, 7), (2, 2), 'SAME'),
-      tl.BatchNorm(mode=mode),
-      tl.Relu(),
-      tl.MaxPool(pool_size=(3, 3), strides=(2, 2)),
-      ConvBlock(3, [d_hidden, d_hidden, 4 * d_hidden], (1, 1), mode=mode),
-      IdentityBlock(3, [d_hidden, d_hidden, 4 * d_hidden], mode=mode),
-      IdentityBlock(3, [d_hidden, d_hidden, 4 * d_hidden], mode=mode),
-      ConvBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden], (2, 2),
-                mode=mode),
-      IdentityBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden], mode=mode),
-      IdentityBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden], mode=mode),
-      IdentityBlock(3, [2 * d_hidden, 2 * d_hidden, 8 * d_hidden], mode=mode),
-      ConvBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden], (2, 2),
-                mode=mode),
-      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden], mode=mode),
-      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden], mode=mode),
-      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden], mode=mode),
-      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden], mode=mode),
-      IdentityBlock(3, [4 * d_hidden, 4 * d_hidden, 16 * d_hidden], mode=mode),
-      ConvBlock(3, [8 * d_hidden, 8 * d_hidden, 32 * d_hidden], (2, 2),
-                mode=mode),
-      IdentityBlock(3, [8 * d_hidden, 8 * d_hidden, 32 * d_hidden], mode=mode),
-      IdentityBlock(3, [8 * d_hidden, 8 * d_hidden, 32 * d_hidden], mode=mode),
-      tl.AvgPool(pool_size=(7, 7)),
-      tl.Flatten(),
-      tl.Dense(n_output_classes),
-      tl.LogSoftmax(),
-  )
-
-
-def WideResnetBlock(channels, strides=(1, 1), bn_momentum=0.9, mode='train'):
-  """WideResnet convolutional block."""
-  return [
-      tl.BatchNorm(momentum=bn_momentum, mode=mode),
-      tl.Relu(),
-      tl.Conv(channels, (3, 3), strides, padding='SAME'),
-      tl.BatchNorm(momentum=bn_momentum, mode=mode),
-      tl.Relu(),
-      tl.Conv(channels, (3, 3), padding='SAME'),
-  ]
-
-
-def WideResnetGroup(n, channels, strides=(1, 1), bn_momentum=0.9, mode='train'):
-  shortcut = [
-      tl.Conv(channels, (3, 3), strides, padding='SAME'),
-  ]
-  return [
-      tl.Residual(WideResnetBlock(channels, strides, bn_momentum=bn_momentum,
-                                  mode=mode),
-                  shortcut=shortcut),
-      tl.Residual([WideResnetBlock(channels, (1, 1), bn_momentum=bn_momentum,
-                                   mode=mode)
-                   for _ in range(n - 1)]),
-  ]
-
-
-def WideResnet(n_blocks=3, widen_factor=1, n_output_classes=10, bn_momentum=0.9,
-               mode='train'):
-  """WideResnet from https://arxiv.org/pdf/1605.07146.pdf.
-
-  Args:
-    n_blocks: int, number of blocks in a group. total layers = 6n + 4.
-    widen_factor: int, widening factor of each group. k=1 is vanilla resnet.
-    n_output_classes: int, number of distinct output classes.
-    bn_momentum: float, momentum in BatchNorm.
-    mode: Whether we are training or evaluating or doing inference.
-
-  Returns:
-    The list of layers comprising a WideResnet model with the given parameters.
-  """
-  return tl.Model(
-      tl.ToFloat(),
-      tl.Conv(16, (3, 3), padding='SAME'),
-      WideResnetGroup(n_blocks, 16 * widen_factor, bn_momentum=bn_momentum,
-                      mode=mode),
-      WideResnetGroup(n_blocks, 32 * widen_factor, (2, 2),
-                      bn_momentum=bn_momentum, mode=mode),
-      WideResnetGroup(n_blocks, 64 * widen_factor, (2, 2),
-                      bn_momentum=bn_momentum, mode=mode),
-      tl.BatchNorm(momentum=bn_momentum, mode=mode),
-      tl.Relu(),
-      tl.AvgPool(pool_size=(8, 8)),
-      tl.Flatten(),
-      tl.Dense(n_output_classes),
-      tl.LogSoftmax(),
-  )
diff --git a/tensor2tensor/trax/models/resnet_test.py b/tensor2tensor/trax/models/resnet_test.py
deleted file mode 100644
index 4751ac9a8..000000000
--- a/tensor2tensor/trax/models/resnet_test.py
+++ /dev/null
@@ -1,45 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 Resnet models."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import absltest
-from tensor2tensor.trax import backend
-from tensor2tensor.trax import layers as tl
-from tensor2tensor.trax.models import resnet
-
-
-class ResnetTest(absltest.TestCase):
-
-  def test_resnet(self):
-    input_shape = (3, 256, 256, 3)
-    model = resnet.Resnet50(d_hidden=8, n_output_classes=10)
-    final_shape = tl.check_shape_agreement(model, input_shape)
-    self.assertEqual((3, 10), final_shape)
-
-  def test_wide_resnet(self):
-    input_shape = (3, 32, 32, 3)
-    model = resnet.WideResnet(n_blocks=1, n_output_classes=10)
-    final_shape = tl.check_shape_agreement(model, input_shape)
-    self.assertEqual((3, 10), final_shape)
-
-
-
-if __name__ == '__main__':
-  absltest.main()
diff --git a/tensor2tensor/trax/models/transformer.py b/tensor2tensor/trax/models/transformer.py
deleted file mode 100644
index 640d54bf7..000000000
--- a/tensor2tensor/trax/models/transformer.py
+++ /dev/null
@@ -1,395 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Transformer Models."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import functools
-
-from tensor2tensor.trax import layers as tl
-
-
-def FeedForward(d_model, d_ff, dropout, layer_idx, mode):
-  """Feed-forward block with layer normalization at start."""
-  return [
-      tl.LayerNorm(),
-      tl.Dense(d_ff),
-      tl.Relu(),
-      tl.Dropout(rate=dropout, name='ff_middle_%d' % layer_idx, mode=mode),
-      tl.Dense(d_model),
-      tl.Dropout(rate=dropout, name='ff_final_%d' % layer_idx, mode=mode),
-  ]
-
-
-def EncoderBlock(d_model, d_ff, n_heads, dropout, layer_idx, mode):
-  """Returns a layer sequence that implements a Transformer encoder block.
-
-  The input to the layer sequence is a pair, (activations, mask), where the
-  mask was created from the original source tokens to prevent attending to the
-  padding part of the input.
-
-  Args:
-    d_model: int:  depth of embedding
-    d_ff: int: depth of feed-forward layer
-    n_heads: int: number of attention heads
-    dropout: float: dropout rate (how much to drop out)
-    layer_idx: which layer are we at (for bookkeeping)
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    A sequence of layers that maps an (activations, mask) pair to an
-    (activations, mask) pair.
-  """
-  attention = [
-      tl.LayerNorm(),
-      tl.Attention(d_model, n_heads=n_heads, dropout=dropout, mode=mode),
-      tl.Dropout(rate=dropout, name='enc_attn_dropout', mode=mode),
-  ]
-  feed_forward = [
-      FeedForward(d_model, d_ff, dropout, layer_idx=layer_idx, mode=mode),
-  ]
-  return [
-      tl.Residual(attention),
-      tl.Residual(feed_forward),
-  ]
-
-
-def TransformerEncoder(vocab_size,
-                       n_classes=10,
-                       d_model=512,
-                       d_ff=2048,
-                       n_layers=6,
-                       n_heads=8,
-                       dropout=0.1,
-                       max_len=2048,
-                       mode='train'):
-  """Returns a Transformer encoder model.
-
-  The input to the model is a tensor of tokens.
-
-  Args:
-    vocab_size: int: vocab size
-    n_classes: how many classes on output
-    d_model: int:  depth of embedding
-    d_ff: int: depth of feed-forward layer
-    n_layers: int: number of encoder/decoder layers
-    n_heads: int: number of attention heads
-    dropout: float: dropout rate (how much to drop out)
-    max_len: int: maximum symbol length for positional encoding
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    A Transformer model as a layer that maps from a tensor of tokens to
-    activations over a set of output classes.
-  """
-  embedder = [
-      tl.Embedding(d_model, vocab_size),
-      tl.Dropout(rate=dropout, name='emb_dropout', mode=mode),
-      tl.PositionalEncoding(max_len=max_len),
-  ]
-  return tl.Model([                             #      tokens
-      tl.Dup(),                                 # toks toks
-      tl.Parallel(embedder, tl.PaddingMask()),  # vecs mask
-      [EncoderBlock(d_model, d_ff, n_heads, dropout, i, mode)
-       for i in range(n_layers)],               # vecs mask
-      tl.Parallel([], tl.Drop()),               # ____  0
-      tl.LayerNorm(),                           # vecs
-      tl.Mean(axis=1),  # Average on length.    # vecs
-      tl.Dense(n_classes),                      # vecs
-      tl.LogSoftmax(),                          # vecs
-  ])
-
-
-def DecoderBlock(d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-                 attention_type, dropout, share_qk, layer_idx, mode):
-  """Returns a layer sequence that implements a Transformer decoder block.
-
-  The input to the layer sequence is an activation tensor.
-
-  Args:
-    d_model: int:  depth of embedding
-    d_ff: int: depth of feed-forward layer
-    n_heads: int: number of attention heads
-    d_attention_key: int: depth of key vector for each attention head
-    d_attention_value: int: depth of value vector for each attention head
-    attention_type: subclass of tl.BaseCausalAttention: attention class to use
-    dropout: float: dropout rate (how much to drop out)
-    share_qk: bool, whether to share queries and keys
-    layer_idx: which layer are we at (for bookkeeping)
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    A sequence of layers that maps an activation tensor to an activation tensor.
-  """
-  self_attention = [
-      tl.LayerNorm(),  # vec
-      tl.CausalAttention(
-          d_model, n_heads=n_heads, d_attention_key=d_attention_key,
-          d_attention_value=d_attention_value, attention_type=attention_type,
-          share_qk=share_qk, mode=mode),
-      tl.Dropout(rate=dropout, name='attention_%d' % layer_idx, mode=mode),
-  ]
-  feed_forward = [
-      FeedForward(d_model, d_ff, dropout, layer_idx=layer_idx, mode=mode),
-  ]
-  return [
-      tl.Residual(self_attention),
-      tl.Residual(feed_forward),
-  ]
-
-
-def TransformerDecoder(vocab_size=None,
-                       d_model=512,
-                       d_ff=2048,
-                       n_layers=6,
-                       n_heads=8,
-                       d_attention_key=None,
-                       d_attention_value=None,
-                       attention_type=tl.DotProductCausalAttention,
-                       dropout=0.1,
-                       share_qk=False,
-                       max_len=2048,
-                       mode='train'):
-  """Returns a Transformer decoder model.
-
-  The input to the model is either continuous or discrete - controlled by
-  vocab_size. Does not shift the input to the right, i.e. the output for
-  timestep t is based on inputs up to timestep t inclusively.
-
-  Args:
-    vocab_size: int or None: vocab size if running on discrete input, None
-        otherwise.
-    d_model: int:  depth of embedding
-    d_ff: int: depth of feed-forward layer
-    n_layers: int: number of encoder/decoder layers
-    n_heads: int: number of attention heads
-    d_attention_key: int: depth of key vector for each attention head
-        (default is d_model // n_heads)
-    d_attention_value: int: depth of value vector for each attention head
-        (default is d_model // n_heads)
-    attention_type: subclass of tl.BaseCausalAttention: attention class to use
-    dropout: float: dropout rate (how much to drop out)
-    share_qk: bool, whether to share queries and keys in decoder attention
-    max_len: int: maximum symbol length for positional encoding
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    A Transformer decoder as a layer that maps from a continuous or discrete
-    tensor to a continuous tensor.
-  """
-  if vocab_size is None:
-    input_layer = tl.Dense
-  else:
-    input_layer = functools.partial(tl.Embedding, vocab_size=vocab_size)
-  return tl.Model(                  # vecs
-      input_layer(d_model),         # vecs
-      tl.Dropout(rate=dropout, mode=mode),
-      tl.PositionalEncoding(max_len=max_len),
-      [DecoderBlock(  # pylint: disable=g-complex-comprehension
-          d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-          attention_type, dropout, share_qk, i, mode)
-       for i in range(n_layers)],   # vecs
-      tl.LayerNorm(),               # vecs
-  )
-
-
-def TransformerLM(vocab_size,
-                  d_model=512,
-                  d_ff=2048,
-                  n_layers=6,
-                  n_heads=8,
-                  d_attention_key=None,
-                  d_attention_value=None,
-                  attention_type=tl.DotProductCausalAttention,
-                  dropout=0.1,
-                  share_qk=False,
-                  max_len=2048,
-                  n_chunks=0,
-                  mode='train'):
-  """Returns a Transformer language model.
-
-  The input to the model is a tensor of tokens. (This model uses only the
-  decoder part of the overall Transformer.)
-
-  Args:
-    vocab_size: int: vocab size
-    d_model: int:  depth of embedding
-    d_ff: int: depth of feed-forward layer
-    n_layers: int: number of encoder/decoder layers
-    n_heads: int: number of attention heads
-    d_attention_key: int: depth of key vector for each attention head
-        (default is d_model // n_heads)
-    d_attention_value: int: depth of value vector for each attention head
-        (default is d_model // n_heads)
-    attention_type: subclass of tl.BaseCausalAttention: attention class to use
-    dropout: float: dropout rate (how much to drop out)
-    share_qk: bool, whether to share queries and keys in decoder attention
-    max_len: int: maximum symbol length for positional encoding
-    n_chunks: int: number of chunks (must match input pipeline)
-    mode: str: 'train', 'eval' or 'predict', predict mode is for fast inference
-
-  Returns:
-    A Transformer language model as a layer that maps from a tensor of tokens
-    to activations over a vocab set.
-  """
-  if n_chunks == 0:
-    concatenate_chunks = split_chunks = []
-  else:
-    concatenate_chunks = tl.Concatenate(n_items=n_chunks)
-    split_chunks = tl.Split(n_sections=n_chunks, axis=-2)
-
-  embedder = [
-      tl.Embedding(d_model, vocab_size),
-      tl.Dropout(rate=dropout, name='embedding', mode=mode),
-      tl.PositionalEncoding(max_len=max_len, mode=mode),
-  ]
-
-  return tl.Model(                  # tokens (or chunked tuple of tokens)
-      concatenate_chunks,           # tokens
-      tl.ShiftRight(mode=mode),     # toks
-      embedder,                     # vecs
-      [DecoderBlock(  # pylint: disable=g-complex-comprehension
-          d_model, d_ff, n_heads, d_attention_key, d_attention_value,
-          attention_type, dropout, share_qk, i, mode)
-       for i in range(n_layers)],   # vecs
-      tl.LayerNorm(),               # vecs
-      tl.Dense(vocab_size),         # vecs
-      tl.LogSoftmax(),              # vecs
-      split_chunks,                 # vecs (or chunked tuple of vecs)
-  )
-
-
-def EncoderDecoder(d_model, d_ff, n_heads, dropout, layer_idx, mode):
-  """Transformer encoder-decoder layer.
-
-  The input is a triple (decoder_input, mask, encoder) where the mask is
-  created from the original source to prevent attending to the padding part
-  of the encoder.
-
-  Args:
-    d_model: int:  depth of embedding
-    d_ff: int: depth of feed-forward layer
-    n_heads: int: number of attention heads
-    dropout: float: dropout rate (how much to drop out)
-    layer_idx: which layer are we at (for bookkeeping)
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    the layer, returning a triple (decoder_activations, mask, encoder).
-  """
-  decoder_self_attention = [                    #        vecs_d   pmask vecs_e
-      tl.LayerNorm(),                           #        vecs_d   ..... ......
-      tl.BasicCausalAttention(
-          d_model, n_heads=n_heads, dropout=dropout, mode=mode),
-      tl.Dropout(rate=dropout, mode=mode),      # vecs_d          ..... ......
-  ]
-  decoder_to_encoder_attention = [        # vecs_d        masks         vecs_e
-      tl.LayerNorm(),                     # vecs_d        masks         vecs_e
-      tl.Parallel([], [], tl.Dup()),      # ______        _____  vecs_e vecs_e
-      tl.Parallel([], tl.Swap()),         # ______        vecs_e masks  ......
-      tl.Parallel([], tl.Dup()),          # ______ vecs_e vecs_e .....  ......
-      tl.AttentionQKV(  # (q k v masks ... --> vecs_d masks ...)
-          d_model, n_heads=n_heads, dropout=dropout, mode=mode),
-      tl.Dropout(rate=dropout, mode=mode),  # vecs_d mask vecs_e
-  ]
-  feed_forward = [
-      FeedForward(d_model, d_ff, dropout, layer_idx=layer_idx, mode=mode),
-  ]
-  return [                                        # vecs_d masks vecs_e
-      tl.Residual(decoder_self_attention),        # vecs_d masks vecs_e
-      tl.Residual(decoder_to_encoder_attention),  # vecs_d masks vecs_e
-      tl.Residual(feed_forward),                  # vecs_d masks vecs_e
-  ]
-
-
-def Transformer(input_vocab_size,
-                output_vocab_size=None,
-                d_model=512,
-                d_ff=2048,
-                n_layers=6,
-                n_heads=8,
-                dropout=0.1,
-                max_len=2048,
-                mode='train'):
-  """Returns a Transformer model.
-
-  This model expects an input pair: target, source.
-
-  Args:
-    input_vocab_size: int: vocab size of the source.
-    output_vocab_size: int (optional): vocab size of the target. If None, the
-      source and target are assumed to have the same vocab.
-    d_model: int:  depth of embedding
-    d_ff: int: depth of feed-forward layer
-    n_layers: int: number of encoder/decoder layers
-    n_heads: int: number of attention heads
-    dropout: float: dropout rate (how much to drop out)
-    max_len: int: maximum symbol length for positional encoding
-    mode: str: 'train' or 'eval'
-
-  Returns:
-    A Transformer model as a layer that maps from a target, source pair to
-    activations over a vocab set.
-  """
-  in_embed = [                                    # tokens
-      tl.Embedding(d_model, input_vocab_size),  # vecs
-      tl.Dropout(rate=dropout, mode=mode),        # vecs
-      tl.PositionalEncoding(max_len=max_len),     # vecs
-  ]
-
-  if output_vocab_size is None:
-    output_vocab_size = input_vocab_size
-    out_embed = in_embed
-  else:
-    out_embed = [                                    # tokens
-        tl.Embedding(d_model, output_vocab_size),  # vecs
-        tl.Dropout(rate=dropout, mode=mode),         # vecs
-        tl.PositionalEncoding(max_len=max_len),      # vecs
-    ]
-
-  encoder_stack = (  # masks vectors --> masks vectors
-      [EncoderBlock(d_model, d_ff, n_heads, dropout, i, mode)
-       for i in range(n_layers)])
-
-  encoder_decoder_stack = (  # vecs_d masks vecs_e --> vecs_d masks vecs_e
-      [EncoderDecoder(d_model, d_ff, n_heads, dropout, i, mode)
-       for i in range(n_layers)])
-
-  # Input: encoder_side_tokens, decoder_side_tokens
-  return tl.Model(  # tokens_e tokens_d
-      tl.Parallel([], tl.Dup()),    # toks_e toks_d toks_d (for loss)
-      tl.Swap(),    # toks_d toks_e ....
-
-      # Encode.
-      tl.Parallel(                                       # toks_d        toks_e
-          [], [tl.Dup(),                                 # ______ toks_e toks_e
-               tl.Parallel(in_embed, tl.PaddingMask()),  # ______ vecs_e masks
-               encoder_stack,                            # ______ vecs_e masks
-               tl.LayerNorm(),                           # ______ vecs_e .....
-               tl.Swap()]),                              # ______ masks  vecs_e
-
-      # Decode.                                  #        toks_d masks vecs_e
-      tl.ShiftRight(),                           #        toks_d ..... ......
-      out_embed,                                 #        vecs_d ..... ......
-      tl.Dup(),                                  # vecs_d vecs_d ..... ......
-      tl.Parallel([], tl.EncoderDecoderMask()),  # ______    masks     ......
-      encoder_decoder_stack,                     # vecs_d    masks     vecs_e
-      tl.Parallel([], tl.Drop(), tl.Drop()),     # vecs_d
-      tl.LayerNorm(),                            # vecs_d
-      tl.Dense(output_vocab_size),               # vecs_d
-      tl.LogSoftmax(),                           # vecs_d
-  )
diff --git a/tensor2tensor/trax/models/transformer_test.py b/tensor2tensor/trax/models/transformer_test.py
deleted file mode 100644
index 2abd9adcd..000000000
--- a/tensor2tensor/trax/models/transformer_test.py
+++ /dev/null
@@ -1,104 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 Transformer models."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import functools
-
-from absl.testing import absltest
-from absl.testing import parameterized
-import numpy as onp
-from tensor2tensor.trax import backend
-from tensor2tensor.trax import layers as tl
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.models import transformer
-
-
-class TransformerTest(parameterized.TestCase):
-
-  def test_transformer_lm_forward_shape(self):
-    """Run the Transformer LM forward and check output shape."""
-    vocab_size = 16
-    input_shape = [3, 5]
-    model = transformer.TransformerLM(
-        vocab_size, d_model=32, d_ff=64, n_layers=2, n_heads=2)
-    final_shape = tl.check_shape_agreement(
-        model, tuple(input_shape), integer_inputs=True)
-    self.assertEqual(tuple(input_shape + [vocab_size]), final_shape)
-
-  def _test_transformer_forward_shape(self, input_vocab_size,
-                                      output_vocab_size):
-    """Run the Transformer forward and check output shape."""
-    single_input_shape = [3, 5]
-    input_shape = (tuple(single_input_shape), tuple(single_input_shape))
-    model = transformer.Transformer(
-        input_vocab_size, output_vocab_size,
-        d_model=32, d_ff=64, n_layers=2, n_heads=2)
-    final_shape = tl.check_shape_agreement(
-        model, input_shape, integer_inputs=True)
-    expected_shape = (tuple(single_input_shape +
-                            [output_vocab_size if output_vocab_size is not None
-                             else input_vocab_size]))
-    self.assertEqual(expected_shape, final_shape[0])
-
-  @parameterized.named_parameters(
-      ('same_vocab', 16, None),
-      ('same_size', 16, 16),
-      ('different_size', 16, 50))
-  def test_transformer_forward_shape(self, input_vocab_size, output_vocab_size):
-    """Run the Transformer forward and check output shape."""
-    self._test_transformer_forward_shape(input_vocab_size, output_vocab_size)
-
-
-  def _test_fast_inference(self, attention_type, length):
-    with backend.use_backend('jax'):
-      vocab_size = 16
-      model_fn = functools.partial(
-          transformer.TransformerLM,
-          vocab_size=vocab_size, d_model=4, d_ff=8, n_layers=2, n_heads=2,
-          attention_type=attention_type,
-      )
-      model_slow = model_fn(mode='eval')
-      model_fast = model_fn(mode='predict')
-      rng = backend.random.get_prng(0)
-      batch_size = 2
-      # Given the same rng, both models initialize with the same parameters.
-      model_slow.initialize_once((batch_size, 1), np.int32, rng)
-      model_fast.initialize_once((batch_size, 1), np.int32, rng)
-
-      buf = onp.zeros((batch_size, length), dtype=np.int32)
-      next_sym = onp.zeros((batch_size, 1), dtype=onp.int32)
-
-      for index in range(length):
-        logits_slow = model_slow(buf, rng=rng)
-        logits_fast = model_fast(next_sym, rng=rng)
-        onp.testing.assert_array_almost_equal(
-            logits_slow[:, index, :], logits_fast[:, 0, :])
-        next_sym = onp.random.randint(vocab_size, size=(batch_size, 1))
-        buf[:, index] = next_sym[:, 0]
-
-  def test_dot_product_causal_attention_fast_inference(self):
-    self._test_fast_inference(tl.DotProductCausalAttention, length=5)
-
-  def test_time_bin_causal_attention_fast_inference(self):
-    attention = functools.partial(tl.TimeBinCausalAttention, bin_length=2)
-    self._test_fast_inference(attention, length=7)
-
-if __name__ == '__main__':
-  absltest.main()
diff --git a/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb b/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb
deleted file mode 100644
index 86663d97f..000000000
--- a/tensor2tensor/trax/notebooks/trax_demo_iclr2019.ipynb
+++ /dev/null
@@ -1,854 +0,0 @@
-{
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "name": "Trax Demo",
-      "version": "0.3.2",
-      "provenance": [],
-      "collapsed_sections": []
-    },
-    "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
-    },
-    "accelerator": "GPU"
-  },
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ySEmBgmqMSIJ",
-        "colab_type": "text"
-      },
-      "source": [
-        "##### Copyright 2019 Google LLC.\n",
-        "\n",
-        "Licensed under the Apache License, Version 2.0 (the \"License\");\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",
-        "https://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."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "o4WGihMLneYq",
-        "colab_type": "text"
-      },
-      "source": [
-        "# Trax: Train Models in JAX\n",
-        "\n",
-        "[JAX](https://github.com/google/jax) allows you to write [numpy](https://www.numpy.org/) and run it fast on accelerators.\n",
-        "\n",
-        "This makes ML research more *fun* and *clear* so we made\n",
-        "* [Trax](https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/trax): a library of models in JAX.\n",
-        "\n",
-        "In this demo we show how to:\n",
-        "* Train a Trax model on a toy copy problem.\n",
-        "* Decode from a pre-trained [Transformer](https://arxiv.org/abs/1706.03762) language model.\n",
-        "* Define [Transformer](https://arxiv.org/abs/1706.03762) from scratch in Trax.\n",
-        "* Do research in Trax: play with hard attention to see how it impacts training and results.\n",
-        "\n",
-        "We would like your feedback!\n",
-        "* What are the parts you like or dislike in JAX and Trax?\n",
-        "* Will you start doing your research in Trax? If not, why? What would change your mind?\n",
-        "* What should we focus on? Speed, cleanliness, memory use?\n",
-        "* If you cannot tell us in person, please add your feedback on [this github issue](https://github.com/tensorflow/tensor2tensor/issues/1478).\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "8YQw0hySTVlK",
-        "colab_type": "text"
-      },
-      "source": [
-        "## Installs\n",
-        "\n",
-        "We install jax and trax and download a pretrained model and vocab file."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "vAWJVzYRnbDU",
-        "colab_type": "code",
-        "outputId": "6cdeff6f-3fc9-406f-feaf-fd1f8d9de775",
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 578
-        }
-      },
-      "source": [
-        "# Install JAX for GPU and Tensor2Tensor.\n",
-        "!pip install --upgrade -q https://storage.googleapis.com/jax-wheels/cuda100/jaxlib-0.1.14-cp36-none-linux_x86_64.whl\n",
-        "!pip install --upgrade -q jax==0.1.27\n",
-        "!pip install --upgrade -q tensor2tensor==1.13.4\n",
-        "# Grab language-model checkpoint and vocab file.\n",
-        "!rm -f model.pkl\n",
-        "!wget https://storage.googleapis.com/traxdemo/model.pkl\n",
-        "!wget https://storage.googleapis.com/traxdemo/vocab.lm1b.en.32768\n",
-        "# Show GPU type.\n",
-        "!nvidia-smi -L"
-      ],
-      "execution_count": 2,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "text": [
-            "\u001b[K     |████████████████████████████████| 44.6MB 1.2MB/s \n",
-            "\u001b[K     |████████████████████████████████| 174kB 3.5MB/s \n",
-            "\u001b[K     |████████████████████████████████| 61kB 24.4MB/s \n",
-            "\u001b[?25h  Building wheel for jax (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
-            "  Building wheel for opt-einsum (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
-            "\u001b[K     |████████████████████████████████| 1.4MB 3.4MB/s \n",
-            "\u001b[K     |████████████████████████████████| 686kB 45.8MB/s \n",
-            "\u001b[K     |████████████████████████████████| 143kB 40.2MB/s \n",
-            "\u001b[K     |████████████████████████████████| 296kB 32.6MB/s \n",
-            "\u001b[?25h  Building wheel for pypng (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
-            "--2019-05-14 22:57:21--  https://storage.googleapis.com/traxdemo/model.pkl\n",
-            "Resolving storage.googleapis.com (storage.googleapis.com)... 209.85.234.128, 2607:f8b0:4001:c12::80\n",
-            "Connecting to storage.googleapis.com (storage.googleapis.com)|209.85.234.128|:443... connected.\n",
-            "HTTP request sent, awaiting response... 200 OK\n",
-            "Length: 211170062 (201M) [application/octet-stream]\n",
-            "Saving to: ‘model.pkl’\n",
-            "\n",
-            "model.pkl           100%[===================>] 201.39M   101MB/s    in 2.0s    \n",
-            "\n",
-            "2019-05-14 22:57:23 (101 MB/s) - ‘model.pkl’ saved [211170062/211170062]\n",
-            "\n",
-            "--2019-05-14 22:57:23--  https://storage.googleapis.com/traxdemo/vocab.lm1b.en.32768\n",
-            "Resolving storage.googleapis.com (storage.googleapis.com)... 64.233.183.128, 2607:f8b0:4001:c07::80\n",
-            "Connecting to storage.googleapis.com (storage.googleapis.com)|64.233.183.128|:443... connected.\n",
-            "HTTP request sent, awaiting response... 200 OK\n",
-            "Length: 297760 (291K) [application/octet-stream]\n",
-            "Saving to: ‘vocab.lm1b.en.32768’\n",
-            "\n",
-            "vocab.lm1b.en.32768 100%[===================>] 290.78K  --.-KB/s    in 0.007s  \n",
-            "\n",
-            "2019-05-14 22:57:24 (40.8 MB/s) - ‘vocab.lm1b.en.32768’ saved [297760/297760]\n",
-            "\n",
-            "GPU 0: Tesla T4 (UUID: GPU-1959cc75-52ab-cf03-e5fa-36aee0d59bc5)\n"
-          ],
-          "name": "stdout"
-        }
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "vvFrqacVS6B6",
-        "colab_type": "text"
-      },
-      "source": [
-        "## Imports"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "dYq8J8uBn9ZC",
-        "colab_type": "code",
-        "outputId": "db8ca8de-164c-4355-8abb-a493e7f9f393",
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 136
-        }
-      },
-      "source": [
-        "from six.moves import cPickle\n",
-        "import os\n",
-        "import datetime\n",
-        "import random\n",
-        "\n",
-        "import numpy as onp\n",
-        "from matplotlib import pyplot as plt\n",
-        "\n",
-        "from jax.ops import index, index_update\n",
-        "\n",
-        "from tensor2tensor.trax import trax\n",
-        "from tensor2tensor.trax import layers as tl\n",
-        "from tensor2tensor.trax import inputs as trax_input\n",
-        "from tensor2tensor.trax import models as trax_models\n",
-        "from tensor2tensor.trax import optimizers as trax_optimizers\n",
-        "from tensor2tensor.trax import backend\n",
-        "from tensor2tensor.trax.backend import numpy as np\n",
-        "from tensor2tensor.trax.backend import random as trax_random"
-      ],
-      "execution_count": 3,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "text": [
-            "\n",
-            "WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
-            "For more information, please see:\n",
-            "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
-            "  * https://github.com/tensorflow/addons\n",
-            "If you depend on functionality not listed there, please file an issue.\n",
-            "\n"
-          ],
-          "name": "stdout"
-        }
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "zR6RVHx4lPzA",
-        "colab_type": "text"
-      },
-      "source": [
-        "# Toy Copy Problem\n",
-        "\n",
-        "Here we define batched random integer inputs for a trivial sequence-copy learning task."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "wGmWmpIslQYv",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "VOCAB_SIZE = 128\n",
-        "def toy_problem_inputs(num_devices, batch_size=64,\n",
-        "                       train_lengths=[10, 20], eval_lengths=[20]):\n",
-        "  \"\"\"Make Inputs for the toy problem of the language 0w0w for w in [1..127]*.\n",
-        "\n",
-        "  Args:\n",
-        "    num_devices: how many devices to build the inputs for (assert 1 for colab).\n",
-        "    batch_size: how large are the batches.\n",
-        "    train_lengths: lengths of w for training.\n",
-        "    eval_lengths: lengths of w for eval.\n",
-        "\n",
-        "  Returns:\n",
-        "    trax.inputs.Inputs\n",
-        "  \"\"\"\n",
-        "  assert num_devices == 1\n",
-        "  def random_minibatches(length_list):\n",
-        "    \"\"\"Generate a stream of random mini-batches.\"\"\"\n",
-        "    while True:\n",
-        "      length = random.choice(length_list)\n",
-        "      w = onp.random.randint(low=1, high=VOCAB_SIZE-1,\n",
-        "                            size=(batch_size, length // 2))\n",
-        "      zero = onp.zeros([batch_size, 1], onp.int32)\n",
-        "      x = onp.concatenate([zero, w, zero, w], axis=1)\n",
-        "      yield (x, x)  # In a language model input and output are the same.\n",
-        "\n",
-        "  return trax_input.Inputs(\n",
-        "      train_stream=lambda: random_minibatches(train_lengths),\n",
-        "      train_eval_stream=lambda: random_minibatches(train_lengths),\n",
-        "      eval_stream=lambda: random_minibatches(eval_lengths),\n",
-        "      input_shape=(None,))"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "eU0mpaf1lRky",
-        "colab_type": "code",
-        "outputId": "bf94086c-5d97-462b-b565-d4ba5f59b6c4",
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 51
-        }
-      },
-      "source": [
-        "inputs = toy_problem_inputs(1)\n",
-        "print(next(inputs.train_stream())[0][0])"
-      ],
-      "execution_count": 5,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "text": [
-            "[  0  68  91  99 107 115 113 111  17 102  48   0  68  91  99 107 115 113\n",
-            " 111  17 102  48]\n"
-          ],
-          "name": "stdout"
-        }
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "KvNaSWu5g2Vm",
-        "colab_type": "text"
-      },
-      "source": [
-        "## Baseline Transformer on Toy Problem"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "AGDmtrgcl73M",
-        "colab_type": "code",
-        "outputId": "4c0f12e9-10ec-4e67-9f15-d2cc7084c083",
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 748
-        }
-      },
-      "source": [
-        "timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M\")\n",
-        "output_dir = os.path.expanduser(\"~/trax_lm_%s\" % timestamp)\n",
-        "def model(mode):\n",
-        "  return trax_models.TransformerLM(\n",
-        "      VOCAB_SIZE, feature_depth=128,\n",
-        "      feedforward_depth=256, num_layers=3,\n",
-        "      num_heads=4, mode=mode)\n",
-        "_ = trax.train(model=model,\n",
-        "               inputs=toy_problem_inputs,\n",
-        "               output_dir=output_dir,\n",
-        "               train_steps=3000,\n",
-        "               eval_steps=10,\n",
-        "               eval_frequency=1000)"
-      ],
-      "execution_count": 5,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "text": [
-            "Step      0: Starting training using 1 devices\n",
-            "\n",
-            "Step      1: Ran 1 train steps in 36.77 secs\n",
-            "Step      1: Total trainable parameters size: 692736\n",
-            "Step      1: Evaluation\n",
-            "Step      1: train           accuracy |  0.00616714\n",
-            "Step      1: train neg_log_perplexity | -5.06836748\n",
-            "Step      1: train               loss |  5.06836748\n",
-            "Step      1: eval            accuracy |  0.00610795\n",
-            "Step      1: eval  neg_log_perplexity | -5.20451212\n",
-            "Step      1: eval                loss |  5.20451212\n",
-            "Step      1: Finished evaluation\n",
-            "\n",
-            "Step   1000: Ran 999 train steps in 89.13 secs\n",
-            "Step   1000: Evaluation\n",
-            "Step   1000: train           accuracy |  0.45719695\n",
-            "Step   1000: train neg_log_perplexity | -2.71764731\n",
-            "Step   1000: train               loss |  2.71764731\n",
-            "Step   1000: eval            accuracy |  0.41278410\n",
-            "Step   1000: eval  neg_log_perplexity | -2.94052887\n",
-            "Step   1000: eval                loss |  2.94052887\n",
-            "Step   1000: Finished evaluation\n",
-            "\n",
-            "Step   2000: Ran 1000 train steps in 15.61 secs\n",
-            "Step   2000: Evaluation\n",
-            "Step   2000: train           accuracy |  0.43169984\n",
-            "Step   2000: train neg_log_perplexity | -2.82782769\n",
-            "Step   2000: train               loss |  2.82782769\n",
-            "Step   2000: eval            accuracy |  0.41278410\n",
-            "Step   2000: eval  neg_log_perplexity | -2.92255998\n",
-            "Step   2000: eval                loss |  2.92255998\n",
-            "Step   2000: Finished evaluation\n",
-            "\n",
-            "Step   3000: Ran 1000 train steps in 15.64 secs\n",
-            "Step   3000: Evaluation\n",
-            "Step   3000: train           accuracy |  0.45053267\n",
-            "Step   3000: train neg_log_perplexity | -2.73254609\n",
-            "Step   3000: train               loss |  2.73254609\n",
-            "Step   3000: eval            accuracy |  0.41249999\n",
-            "Step   3000: eval  neg_log_perplexity | -2.92720962\n",
-            "Step   3000: eval                loss |  2.92720962\n",
-            "Step   3000: Finished evaluation\n",
-            "Step   3000: Training done\n"
-          ],
-          "name": "stdout"
-        }
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "eapBBkRUuho7",
-        "colab_type": "text"
-      },
-      "source": [
-        "# Decoding from a Pre-Trained Transformer Language Model"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "H6hVQ3v5iC00",
-        "colab_type": "code",
-        "outputId": "812949cc-4294-4a42-f55a-c40f65e151f8",
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 187
-        }
-      },
-      "source": [
-        "# load model checkpoint\n",
-        "with open(\"model.pkl\", \"rb\") as f:\n",
-        "   (params, step, history) = cPickle.load(f, encoding=\"latin1\")\n",
-        "\n",
-        "# lm1b subword vocab\n",
-        "def clean(x):\n",
-        "  return x[1:-2]\n",
-        "with open(\"vocab.lm1b.en.32768\", \"r\") as fp:\n",
-        "  vocab = list(map(clean, fp.readlines()))\n",
-        "vocab_map = {v:idx for idx,v in enumerate(vocab)}\n",
-        "\n",
-        "list(enumerate(vocab))[:10]"
-      ],
-      "execution_count": 6,
-      "outputs": [
-        {
-          "output_type": "execute_result",
-          "data": {
-            "text/plain": [
-              "[(0, '<pad>_'),\n",
-              " (1, '<EOS>_'),\n",
-              " (2, 'the_'),\n",
-              " (3, ' , _'),\n",
-              " (4, ' ._'),\n",
-              " (5, 'to_'),\n",
-              " (6, 'of_'),\n",
-              " (7, 'a_'),\n",
-              " (8, 'and_'),\n",
-              " (9, 'in_')]"
-            ]
-          },
-          "metadata": {
-            "tags": []
-          },
-          "execution_count": 6
-        }
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "W-7s9RXQNIru",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "tlm = trax_models.TransformerLM(\n",
-        "  dropout=0.1, \n",
-        "  feature_depth=512, \n",
-        "  feedforward_depth=2048, \n",
-        "  max_len=2048, \n",
-        "  mode='eval', \n",
-        "  num_heads=8, \n",
-        "  num_layers=6, \n",
-        "  vocab_size=32000)"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "iLdtplDpdTMr",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "def gumbel_sample(v, temperature=0.8):\n",
-        "  u = onp.random.uniform(low=1e-9, high=1.0, size=v.shape)\n",
-        "  g = -onp.log(-onp.log(u))\n",
-        "  return np.argmax(v + g * temperature)"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "IHSbtHzPjW6i",
-        "colab_type": "code",
-        "outputId": "7a8306b7-6c6b-41ba-c5aa-c1a76d9d8037",
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 102
-        }
-      },
-      "source": [
-        "prompt = \"Please_\"\n",
-        "num_samples = 5\n",
-        "max_length = 20\n",
-        "for _ in range(num_samples):\n",
-        "  enc = [vocab_map[w] for w in str.split(prompt)]\n",
-        "  pos = len(enc)\n",
-        "  rng = trax_random.get_prng(0)\n",
-        "  data = np.zeros((1, 50), dtype=np.int32)\n",
-        "  data = index_update(data, index[0, 0:pos], enc)\n",
-        "\n",
-        "  while pos < max_length:\n",
-        "    tmp = tlm(data, params=params, rng=rng)\n",
-        "    next_sym = gumbel_sample(tmp[0, pos])\n",
-        "    data = index_update(data, index[0, pos], next_sym)\n",
-        "    pos += 1\n",
-        "    if int(next_sym) == 1:\n",
-        "      break\n",
-        "\n",
-        "  print(\"\".join([vocab[idx] for idx in onp.array(data)[0, 0:pos]]))"
-      ],
-      "execution_count": 10,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "text": [
-            "Please_write_to_him_to_tell_him_about_the_Wallace_and_Gromit_films_. _and_to_give_him_this_\n",
-            "Please_do_not_turn_to_making_sure_your_children_are_already_in_school_or_that_you_have_school_ ._<EOS>_\n",
-            "Please_read_the_full_prospectus_to_see_if_the_proposed_transaction_may_be_accurate_ ._<EOS>_\n",
-            "Please_note_that_the_new_policy_has_been_strengthened_by_the_fact_that_Britney_Spears_ ' _mother_ , _Janet_Jackson_\n",
-            "Please_ , _please_aim_at_your_brother_ , _if_you_want_to_ ._<EOS>_\n"
-          ],
-          "name": "stdout"
-        }
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Ym8otS7HpUIO",
-        "colab_type": "text"
-      },
-      "source": [
-        "# Transformer from Scratch\n",
-        "\n",
-        "Here we re-implement multiheaded self-attention and a transformer language model from scratch using only a few simple linear primitives from trax.\n",
-        "\n",
-        "Note in particular the commented modifications in the core  __DotProductAttention__ function as an example of how easy it is to modify layers and models for research using Trax."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "uw-GIdm2p_4X",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "def DotProductAttention(query, key, value, mask, dropout, mode, rng, hard_k=4):\n",
-        "  \"\"\"Core dot product self-attention.\n",
-        "  Args:\n",
-        "    query: array of representations\n",
-        "    key: array of representations\n",
-        "    value: array of representations\n",
-        "    mask: attention-mask, gates attention\n",
-        "    dropout: float: dropout rate\n",
-        "    mode: 'eval' or 'train': whether to use dropout\n",
-        "    rng: JAX PRNGKey: subkey for disposable use\n",
-        "  Returns:\n",
-        "    Self attention for q, k, v arrays.\n",
-        "  \"\"\"\n",
-        "  depth = np.shape(query)[-1]\n",
-        "  dots = np.matmul(query, np.swapaxes(key, -1, -2)) / np.sqrt(depth)\n",
-        "  if mask is not None:\n",
-        "    dots = np.where(mask, dots, -1e9)\n",
-        "  # Softmax.\n",
-        "  dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))\n",
-        "  # ----------------------------------------------------------------------\n",
-        "  # As an example of a simple research modification, we modify the typical \n",
-        "  # dot-product attention mechanism with top-k \"hard attention\":\n",
-        "  # ----------------------------------------------------------------------\n",
-        "  if hard_k > 0:\n",
-        "    top_k = np.sort(dots)[..., -hard_k]  # Get the top-kth weight.\n",
-        "    dots -= top_k[..., np.newaxis]  # Subtract (be 0 for lower ones).\n",
-        "    dots = np.maximum(dots, 0)\n",
-        "    dots /= np.sum(dots, axis=-1, keepdims=True)  # Re-normalize.\n",
-        "  # ----------------------------------------------------------------------\n",
-        "  if dropout >= 1.0:\n",
-        "    raise ValueError('Dropout rates must be lower than 1.')\n",
-        "  if dropout is not None and dropout > 0.0 and mode == 'train':\n",
-        "    keep = backend.random.bernoulli(rng, 1.0 - dropout, dots.shape)\n",
-        "    dots = np.where(keep, dots / (1.0 - dropout), 0)\n",
-        "  out = np.matmul(dots, value)\n",
-        "  # Uncomment to see an example TRAX stack trace to this point:\n",
-        "  # ----------------------------------------------------------------------\n",
-        "  # raise ValueError(\"err\")\n",
-        "  # ----------------------------------------------------------------------\n",
-        "  return out\n",
-        "\n",
-        "\n",
-        "def _multihead_attention_output_shape(  # pylint: disable=invalid-name\n",
-        "    input_shapes, **unused_kwargs):\n",
-        "  \"\"\"Helper: calculate multihead attention output shape.\"\"\"\n",
-        "  q_shape = input_shapes[0][0]  # Inputs are ((q, k, v), mask).\n",
-        "  mask_shape = input_shapes[1]\n",
-        "  return q_shape, mask_shape\n",
-        "\n",
-        "\n",
-        "@tl.layer(output_shape=_multihead_attention_output_shape)\n",
-        "def PureMultiHeadedAttention(x, params, num_heads=8, dropout=0.0,\n",
-        "                             mode='train', **kwargs):\n",
-        "  \"\"\"Pure transformer-style multi-headed attention.\n",
-        "  Args:\n",
-        "    x: inputs ((q, k, v), mask)\n",
-        "    params: parameters (none)\n",
-        "    num_heads: int: number of attention heads\n",
-        "    dropout: float: dropout rate\n",
-        "    mode: str: 'train' or 'eval'\n",
-        "    **kwargs: other arguments including the rng\n",
-        "  Returns:\n",
-        "    Pure Multi-headed attention result, and the mask.\n",
-        "  \"\"\"\n",
-        "  del params\n",
-        "  rng = kwargs.get('rng', None)\n",
-        "  (q, k, v), mask = x\n",
-        "  feature_depth = q.shape[-1]\n",
-        "  assert feature_depth % num_heads == 0\n",
-        "  head_depth = feature_depth // num_heads\n",
-        "  nbatch = np.shape(q)[0]\n",
-        "  # nbatch, seqlen, feature_depth --> nbatch, num_heads, seqlen, head_depth\n",
-        "  def SplitHeads(x):\n",
-        "    return np.transpose(\n",
-        "        np.reshape(x, (nbatch, -1, num_heads, head_depth)), (0, 2, 1, 3))\n",
-        "  # nbatch, num_heads, seqlen, head_depth --> nbatch, seqlen, feature_depth\n",
-        "  def JoinHeads(x):  # pylint: disable=invalid-name\n",
-        "    return np.reshape(\n",
-        "        np.transpose(x, (0, 2, 1, 3)), (nbatch, -1, num_heads*head_depth))\n",
-        "  # Split heads, dot-product attention, rejoin heads.\n",
-        "  res = JoinHeads(\n",
-        "      DotProductAttention(\n",
-        "          SplitHeads(q), SplitHeads(k), SplitHeads(v), mask,\n",
-        "          dropout=dropout, mode=mode, rng=rng))\n",
-        "  return res, mask  # Keep the mask.\n",
-        "\n",
-        "\n",
-        "def MultiHeadedAttentionQKV(\n",
-        "    feature_depth, num_heads=8, dropout=0.0, mode='train'):\n",
-        "  \"\"\"Transformer-style multi-headed attention.\n",
-        "  Accepts inputs of the form (q, k, v), mask.\n",
-        "  Args:\n",
-        "    feature_depth: int:  depth of embedding\n",
-        "    num_heads: int: number of attention heads\n",
-        "    dropout: float: dropout rate\n",
-        "    mode: str: 'train' or 'eval'\n",
-        "  Returns:\n",
-        "    Multi-headed self-attention result and the mask.\n",
-        "  \"\"\"\n",
-        "  return tl.Serial(\n",
-        "      tl.Parallel(\n",
-        "          tl.Parallel(\n",
-        "              tl.Dense(feature_depth),\n",
-        "              tl.Dense(feature_depth),\n",
-        "              tl.Dense(feature_depth),\n",
-        "          ),\n",
-        "          tl.Copy()\n",
-        "      ),\n",
-        "      PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter\n",
-        "          feature_depth=feature_depth, num_heads=num_heads,\n",
-        "          dropout=dropout, mode=mode),\n",
-        "      tl.Parallel(tl.Dense(feature_depth), tl.Copy())\n",
-        "  )\n",
-        "\n",
-        "\n",
-        "def MultiHeadedAttention(\n",
-        "    feature_depth, num_heads=8, dropout=0.0, mode='train'):\n",
-        "  \"\"\"Transformer-style multi-headed attention.\n",
-        "  Accepts inputs of the form (x, mask) and constructs (q, k, v) from x.\n",
-        "  Args:\n",
-        "    feature_depth: int:  depth of embedding\n",
-        "    num_heads: int: number of attention heads\n",
-        "    dropout: float: dropout rate\n",
-        "    mode: str: 'train' or 'eval'\n",
-        "  Returns:\n",
-        "    Multi-headed self-attention layer.\n",
-        "  \"\"\"\n",
-        "  return tl.Serial(\n",
-        "      tl.Parallel(\n",
-        "          # q = k = v = first input\n",
-        "          tl.Branch(\n",
-        "              tl.Copy(), tl.Copy(), tl.Copy()),\n",
-        "          tl.Copy()  # pass the mask\n",
-        "      ),\n",
-        "      MultiHeadedAttentionQKV(  # pylint: disable=no-value-for-parameter\n",
-        "          feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),\n",
-        "  )"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "Ge42t7VZl-d2",
-        "colab_type": "code",
-        "colab": {}
-      },
-      "source": [
-        "def ResidualFeedForward(feature_depth,\n",
-        "                        feedforward_depth,\n",
-        "                        dropout,\n",
-        "                        mode):\n",
-        "  \"\"\"Residual feed-forward layer with normalization at start.\"\"\"\n",
-        "  return tl.Residual(\n",
-        "      tl.LayerNorm(),\n",
-        "      tl.Dense(feedforward_depth),\n",
-        "      tl.Relu(),\n",
-        "      tl.Dropout(rate=dropout, mode=mode),\n",
-        "      tl.Dense(feature_depth),\n",
-        "      tl.Dropout(rate=dropout, mode=mode)\n",
-        "  )\n",
-        "\n",
-        "\n",
-        "def DecoderLayer(feature_depth,\n",
-        "                 feedforward_depth,\n",
-        "                 num_heads,\n",
-        "                 dropout,\n",
-        "                 mode):\n",
-        "  \"\"\"Transformer decoder layer.\n",
-        "  Args:\n",
-        "    feature_depth: int:  depth of embedding\n",
-        "    feedforward_depth: int: depth of feed-forward layer\n",
-        "    num_heads: int: number of attention heads\n",
-        "    dropout: float: dropout rate (how much to drop out)\n",
-        "    mode: str: 'train' or 'eval'\n",
-        "  Returns:\n",
-        "    the layer.\n",
-        "  \"\"\"\n",
-        "  return tl.Serial(\n",
-        "      tl.Residual(  # Self-attention block.\n",
-        "          tl.LayerNorm(),\n",
-        "          tl.Branch(tl.Copy(), tl.CausalMask(axis=-2)),  # Create mask.\n",
-        "          # We replace the \"stock\" self-attention layer with the one defined\n",
-        "          # above:\n",
-        "          # tl.MultiHeadedAttention(feature_depth, num_heads=num_heads,\n",
-        "          #                         dropout=dropout, mode=mode),\n",
-        "          MultiHeadedAttention(feature_depth, num_heads=num_heads,\n",
-        "                                  dropout=dropout, mode=mode),\n",
-        "          tl.Select(0),  # Drop the mask.\n",
-        "          tl.Dropout(rate=dropout, mode=mode)\n",
-        "      ),\n",
-        "      ResidualFeedForward(feature_depth, feedforward_depth, dropout, mode=mode)\n",
-        "  )\n",
-        "\n",
-        "\n",
-        "def TransformerLM(vocab_size,\n",
-        "                  feature_depth=512,\n",
-        "                  feedforward_depth=2048,\n",
-        "                  num_layers=6,\n",
-        "                  num_heads=8,\n",
-        "                  dropout=0.1,\n",
-        "                  max_len=2048,\n",
-        "                  mode='train'):\n",
-        "  \"\"\"Transformer language model (only uses the decoder part of Transformer).\n",
-        "  Args:\n",
-        "    vocab_size: int: vocab size\n",
-        "    feature_depth: int:  depth of embedding\n",
-        "    feedforward_depth: int: depth of feed-forward layer\n",
-        "    num_layers: int: number of encoder/decoder layers\n",
-        "    num_heads: int: number of attention heads\n",
-        "    dropout: float: dropout rate (how much to drop out)\n",
-        "    max_len: int: maximum symbol length for positional encoding\n",
-        "    mode: str: 'train' or 'eval'\n",
-        "  Returns:\n",
-        "    the layer.\n",
-        "  \"\"\"\n",
-        "  return tl.Serial(\n",
-        "      tl.ShiftRight(),\n",
-        "      tl.Embedding(feature_depth, vocab_size),\n",
-        "      tl.Dropout(rate=dropout, mode=mode),\n",
-        "      tl.PositionalEncoding(max_len=max_len),\n",
-        "      tl.Serial(*[DecoderLayer(feature_depth, feedforward_depth, num_heads,\n",
-        "                               dropout, mode)\n",
-        "                  for _ in range(num_layers)]),\n",
-        "      tl.LayerNorm(),\n",
-        "      tl.Dense(vocab_size),\n",
-        "      tl.LogSoftmax()\n",
-        "  )"
-      ],
-      "execution_count": 0,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "WZxnwjAEqYDh",
-        "colab_type": "code",
-        "outputId": "f90e965d-2625-4e56-9038-65c087639051",
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 748
-        }
-      },
-      "source": [
-        "timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M\")\n",
-        "output_dir = os.path.expanduser(\"~/trax_lm_%s\" % timestamp)\n",
-        "def new_model(mode):\n",
-        "  return TransformerLM(\n",
-        "      VOCAB_SIZE, feature_depth=128,\n",
-        "      feedforward_depth=256, num_layers=3,\n",
-        "      num_heads=4, mode=mode)\n",
-        "_ = trax.train(model=new_model,\n",
-        "           inputs=toy_problem_inputs,\n",
-        "           output_dir=output_dir,\n",
-        "           train_steps=3000,\n",
-        "           eval_steps=10,\n",
-        "           eval_frequency=1000)"
-      ],
-      "execution_count": 22,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "text": [
-            "Step      0: Starting training using 1 devices\n",
-            "\n",
-            "Step      1: Ran 1 train steps in 42.29 secs\n",
-            "Step      1: Total trainable parameters size: 692736\n",
-            "Step      1: Evaluation\n",
-            "Step      1: train           accuracy |  0.00686553\n",
-            "Step      1: train neg_log_perplexity | -5.42891455\n",
-            "Step      1: train               loss |  5.42891455\n",
-            "Step      1: eval            accuracy |  0.00809659\n",
-            "Step      1: eval  neg_log_perplexity | -5.39403439\n",
-            "Step      1: eval                loss |  5.39403439\n",
-            "Step      1: Finished evaluation\n",
-            "\n",
-            "Step   1000: Ran 999 train steps in 109.64 secs\n",
-            "Step   1000: Evaluation\n",
-            "Step   1000: train           accuracy |  0.12875238\n",
-            "Step   1000: train neg_log_perplexity | -4.29979420\n",
-            "Step   1000: train               loss |  4.29979420\n",
-            "Step   1000: eval            accuracy |  0.09928977\n",
-            "Step   1000: eval  neg_log_perplexity | -4.45948172\n",
-            "Step   1000: eval                loss |  4.45948172\n",
-            "Step   1000: Finished evaluation\n",
-            "\n",
-            "Step   2000: Ran 1000 train steps in 16.89 secs\n",
-            "Step   2000: Evaluation\n",
-            "Step   2000: train           accuracy |  0.53104877\n",
-            "Step   2000: train neg_log_perplexity | -2.33383632\n",
-            "Step   2000: train               loss |  2.33383632\n",
-            "Step   2000: eval            accuracy |  0.54900569\n",
-            "Step   2000: eval  neg_log_perplexity | -2.24813342\n",
-            "Step   2000: eval                loss |  2.24813342\n",
-            "Step   2000: Finished evaluation\n",
-            "\n",
-            "Step   3000: Ran 1000 train steps in 16.91 secs\n",
-            "Step   3000: Evaluation\n",
-            "Step   3000: train           accuracy |  0.56715208\n",
-            "Step   3000: train neg_log_perplexity | -2.15219927\n",
-            "Step   3000: train               loss |  2.15219927\n",
-            "Step   3000: eval            accuracy |  0.54928976\n",
-            "Step   3000: eval  neg_log_perplexity | -2.25436211\n",
-            "Step   3000: eval                loss |  2.25436211\n",
-            "Step   3000: Finished evaluation\n",
-            "Step   3000: Training done\n"
-          ],
-          "name": "stdout"
-        }
-      ]
-    }
-  ]
-}
diff --git a/tensor2tensor/trax/optimizers/__init__.py b/tensor2tensor/trax/optimizers/__init__.py
deleted file mode 100644
index 5974e76b9..000000000
--- a/tensor2tensor/trax/optimizers/__init__.py
+++ /dev/null
@@ -1,37 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Optimizers defined in trax."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gin
-
-from tensor2tensor.trax.optimizers import base
-
-
-def opt_configure(*args, **kwargs):
-  kwargs["module"] = "trax.optimizers"
-  return gin.external_configurable(*args, **kwargs)
-
-# Optimizers (using upper-case names).
-# pylint: disable=invalid-name
-SGD = opt_configure(base.SGD)
-Momentum = opt_configure(base.Momentum)
-RMSProp = opt_configure(base.RMSProp)
-Adam = opt_configure(base.Adam)
-Adafactor = opt_configure(base.Adafactor)
-SM3 = opt_configure(base.SM3)
diff --git a/tensor2tensor/trax/optimizers/base.py b/tensor2tensor/trax/optimizers/base.py
deleted file mode 100644
index 13631a538..000000000
--- a/tensor2tensor/trax/optimizers/base.py
+++ /dev/null
@@ -1,465 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Trax base optimizer class."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.layers import base as layers
-
-
-def tree_flatten(tree):
-  """Flatten a tree into a list."""
-  if isinstance(tree, (list, tuple)):
-    # In python, sum of lists starting from [] is the concatenation.
-    return sum([tree_flatten(t) for t in tree], [])
-  if isinstance(tree, dict):
-    # Only use the values in case of a dictionary node.
-    return sum([tree_flatten(v) for v in tree.values()], [])
-  return [tree]
-
-
-def tree_unflatten(flat, tree):
-  """Unflatten a list into a tree given the tree shape as second argument.
-
-  Args:
-    flat: a flat list of elements to be assembled into a tree.
-    tree: a tree with the structure we want to have in the new tree.
-
-  Returns:
-    A pair (new_tree, rest_of_flat) where the new tree that has the structure
-    of tree but with leaves from flat, and the remaining elements of flat if
-    more were provided than the number of leaves of tree (useful for recursion).
-  """
-  if isinstance(tree, (list, tuple)):
-    new_tree, rest = [], flat
-    for t in tree:
-      new_t, rest = tree_unflatten(rest, t)
-      new_tree.append(new_t)
-    new_tree = tuple(new_tree) if isinstance(tree, tuple) else new_tree
-    return new_tree, rest
-  if isinstance(tree, dict):
-    new_tree, rest = {}, flat
-    for k in tree:
-      new_v, rest = tree_unflatten(rest, tree[k])
-      new_tree[k] = new_v
-    return new_tree, rest
-  return flat[0], flat[1:]
-
-
-class Optimizer(object):
-  """Optimizer object, base class. Maps per-parameter functions to trees."""
-
-  def __init__(self, learning_rate, **init_opt_params):
-    """Initialize the optimizer.
-
-    Takes the initial optimizer parameters as positional arguments. They are fed
-    back to the optimizer in tree_update, in the same order. They can be changed
-    between updates, e.g. for learning rate schedules.
-
-    The constructor should be overridden in derived classes to give names to the
-    optimizer parameters, so the gin configuration can set them.
-
-    Args:
-      learning_rate: The initial learning rate.
-      **init_opt_params: Initial values of any additional optimizer parameters.
-    """
-    init_opt_params["learning_rate"] = learning_rate
-    self._init_opt_params = {
-        name: np.array(value) for (name, value) in init_opt_params.items()
-    }
-
-  def init(self, params):
-    """Create optimizer slots for the given parameters."""
-    raise NotImplementedError
-
-  def update(self, step, grads, params, slots, opt_params):
-    """Update a single parameter array.
-
-    Args:
-      step: Current step.
-      grads: Gradients.
-      params: Parameters.
-      slots: Optimizer slots (e.g. gradient moments).
-      opt_params: Optimizer (hyper)parameters (e.g. learning rate, momentum).
-
-    Returns:
-      (new_params, new_slots)
-    """
-    raise NotImplementedError
-
-  # End subclass interface.
-
-  def tree_init(self, param_tree):
-    return (
-        [self.init(param) for param in tree_flatten(param_tree)],
-        self._init_opt_params,
-    )
-
-  def _update_and_check(self, step, grads, params, slots, opt_params):
-    """Update a single parameter array and check types."""
-    new_params, new_slots = self.update(
-        step, grads, params, slots, opt_params)
-    if isinstance(params, np.ndarray):
-      assert isinstance(new_params, np.ndarray), (
-          "The type of the new parameter values should be np.ndarray; got %s" %
-          type(new_params))
-      assert new_params.dtype == params.dtype, (
-          "The dtype of the new parameter values (%s) is not the same as the "
-          "old one (%s)" % (new_params.dtype, params.dtype))
-    return new_params, new_slots
-
-  def tree_update(self, step, grad_tree, param_tree, slots, opt_params):
-    grads_flat = tree_flatten(grad_tree)
-    params_flat = tree_flatten(param_tree)
-    updated_pairs = [
-        self._update_and_check(step, grad, param, slot, opt_params)
-        for (grad, param, slot) in zip(grads_flat, params_flat, slots)
-    ]
-    new_params_flat, new_slots = zip(*updated_pairs)
-    new_params, _ = tree_unflatten(new_params_flat, param_tree)
-    return new_params, new_slots
-
-
-# Utilities.
-
-
-def l2_norm(tree):
-  """Compute the l2 norm of a pytree of arrays. Useful for weight decay."""
-  leaves = tree_flatten(tree)
-  return np.sqrt(sum(np.vdot(x, x) for x in leaves))
-
-
-def clip_grads(grad_tree, max_norm):
-  """Clip gradients stored as a pytree of arrays to maximum norm `max_norm`."""
-  norm = l2_norm(grad_tree)
-  normalize = lambda g: np.where(norm < max_norm, g, g * (max_norm / norm))
-  return layers.nested_map(normalize, grad_tree)
-
-
-# Optimizers.
-
-
-class SGD(Optimizer):
-  """Plain SGD optimizer."""
-
-  def init(self, params):
-    return None
-
-  def update(self, step, grads, params, slots, opt_params):
-    del step
-    del slots
-    learning_rate = opt_params["learning_rate"]
-    return params - (learning_rate * grads).astype(params.dtype), None
-
-
-class Momentum(Optimizer):
-  """Nesterov momentum optimizer."""
-
-  def __init__(self, learning_rate, mass=0.9, weight_decay_rate=1e-5):  # pylint: disable=useless-super-delegation
-    super(Momentum, self).__init__(
-        learning_rate=learning_rate,
-        mass=mass,
-        weight_decay_rate=weight_decay_rate,
-    )
-
-  def init(self, params):
-    return np.zeros_like(params)
-
-  def update(self, step, grads, params, velocity, opt_params):
-    del step
-    learning_rate = opt_params["learning_rate"]
-    mass = opt_params["mass"]
-    weight_decay_rate = opt_params["weight_decay_rate"]
-    new_velocity = mass * velocity + grads
-    new_params = (1 - weight_decay_rate) * params - (
-        learning_rate * (mass * new_velocity + grads)).astype(params.dtype)
-    return (new_params, new_velocity)
-
-
-class RMSProp(Optimizer):
-  """RMSProp optimizer."""
-
-  def __init__(self, learning_rate, gamma=0.9, eps=1e-8):  # pylint: disable=useless-super-delegation
-    super(RMSProp, self).__init__(
-        learning_rate=learning_rate,
-        gamma=gamma,
-        eps=eps,
-    )
-
-  def init(self, params):
-    return np.ones_like(params)
-
-  def update(self, step, grads, params, avg_sq_grad, opt_params):
-    del step
-    learning_rate = opt_params["learning_rate"]
-    gamma = opt_params["gamma"]
-    eps = opt_params["eps"]
-    avg_sq_grad = avg_sq_grad * gamma + grads**2 * (1. - gamma)
-    params = params - (learning_rate * grads /
-                       (np.sqrt(avg_sq_grad) + eps)).astype(params.dtype)
-    return params, avg_sq_grad
-
-
-class Adam(Optimizer):
-  """Adam optimizer."""
-
-  def __init__(self, learning_rate, weight_decay_rate=1e-5,  # pylint: disable=useless-super-delegation
-               b1=0.9, b2=0.999, eps=1e-5):
-    """Create the Adam optimizer.
-
-    Args:
-      learning_rate: a postitive scalar value for the initial learning rate.
-      weight_decay_rate: rate at which to decay weights.
-      b1: optional, a positive scalar value for beta_1, the exponential decay
-        rate for the first moment estimates (default 0.9).
-      b2: optional, a positive scalar value for beta_2, the exponential decay
-         rate for the second moment estimates (default 0.999).
-      eps: optional, a positive scalar value for epsilon, a small constant for
-        numerical stability (default 1e-8).
-    """
-    super(Adam, self).__init__(
-        learning_rate=learning_rate,
-        weight_decay_rate=weight_decay_rate,
-        b1=b1,
-        b2=b2,
-        eps=eps,
-    )
-
-  def init(self, params):
-    m = np.zeros_like(params)
-    v = np.zeros_like(params)
-    return m, v
-
-  def update(self, step, grads, params, slots, opt_params):
-    m, v = slots
-    learning_rate = opt_params["learning_rate"]
-    weight_decay_rate = opt_params["weight_decay_rate"]
-    b1 = opt_params["b1"]
-    b2 = opt_params["b2"]
-    eps = opt_params["eps"]
-    m = (1 - b1) * grads + b1 * m  # First  moment estimate.
-    v = (1 - b2) * (grads ** 2) + b2 * v  # Second moment estimate.
-    mhat = m / (1 - b1 ** (step + 1))  # Bias correction.
-    vhat = v / (1 - b2 ** (step + 1))
-    params = (1 - weight_decay_rate) * params - (
-        learning_rate * mhat / (np.sqrt(vhat) + eps)).astype(params.dtype)
-    return params, (m, v)
-
-
-class Adafactor(Optimizer):
-  """Adafactor optimizer."""
-
-  def __init__(self,
-               learning_rate,
-               factored=True,
-               multiply_by_parameter_scale=True,
-               do_clipping=True,
-               do_momentum=False,
-               beta1=0.0,
-               decay_rate=0.8,
-               clipping_threshold=1.0,
-               weight_decay_rate=1e-5,
-               epsilon1=1e-30,
-               epsilon2=1e-3):
-    """Create the Adafactor optimizer.
-
-    Adafactor is described in https://arxiv.org/abs/1804.04235.
-
-    Args:
-      learning_rate: float: trax-provided learning rate.
-      factored: boolean: whether to use factored second-moment estimator for 2d
-        variables.
-      multiply_by_parameter_scale: boolean: if True, then scale provided
-        learning_rate by parameter norm. if False, provided learning_rate is
-        absolute step size.
-      do_clipping: whether to clip gradients; if True, set clipping_theshold.
-      do_momentum: whether to use momentum; if True, set beta1.
-      beta1: a float value between 0 and 1, enables momentum and uses extra
-        memory if nonzero!  Off by default.
-      decay_rate: float: controls second-moment exponential decay schedule.
-      clipping_threshold: an optional float >= 1, if None no update clipping.
-      weight_decay_rate: rate at which to decay weights.
-      epsilon1: Regularization constant for squared gradient.
-      epsilon2: Regularization constant for parameter scale.
-    """
-    # These 4 parameters are not configurable once the class is created.
-    self._factored = factored
-    self._multiply_by_parameter_scale = multiply_by_parameter_scale
-    self._do_clipping = do_clipping
-    self._do_momentum = do_momentum
-    # Dynamically configurable parameters will be passed to the update function.
-    super(Adafactor, self).__init__(
-        learning_rate=learning_rate,
-        beta1=beta1,
-        decay_rate=decay_rate,
-        clipping_threshold=clipping_threshold,
-        weight_decay_rate=weight_decay_rate,
-        epsilon1=epsilon1,
-        epsilon2=epsilon2,
-    )
-
-  @staticmethod
-  def _decay_rate_pow(i, exponent=0.8):
-    """Default Adafactor second-moment decay schedule."""
-    t = np.array(i, np.float32) + 1.0
-    return 1.0 - t**(-exponent)
-
-  def init(self, params):
-    shape = params.shape
-    slots = []
-    if self._factored and len(shape) >= 2:
-      v_row = np.zeros(shape[:-1], dtype=np.float32)
-      v_col = np.zeros(shape[:-2] + shape[-1:], dtype=np.float32)
-      slots.extend([v_row, v_col])
-    else:
-      v = np.zeros_like(params)
-      slots.append(v)
-    if self._do_momentum:
-      m = np.zeros_like(params)
-      slots.append(m)
-    return slots
-
-  def update(self, step, grads, params, slots, opt_params):
-    updates = []
-    learning_rate = opt_params["learning_rate"]
-    beta1 = opt_params["beta1"]
-    decay_rate = opt_params["decay_rate"]
-    clipping_threshold = opt_params["clipping_threshold"]
-    weight_decay_rate = opt_params["weight_decay_rate"]
-    epsilon1 = opt_params["epsilon1"]
-    epsilon2 = opt_params["epsilon2"]
-    decay_rate = self._decay_rate_pow(step, exponent=decay_rate)
-    update_scale = learning_rate
-    if self._multiply_by_parameter_scale:
-      update_scale *= np.maximum(
-          np.sqrt(np.mean(params * params)), epsilon2)
-    mixing_rate = 1.0 - decay_rate
-
-    grads_sqr = grads * grads + epsilon1
-    if self._factored and len(params.shape) >= 2:
-      v_row = slots.pop(0)
-      v_col = slots.pop(0)
-      new_v_row = decay_rate * v_row + mixing_rate * np.mean(grads_sqr, axis=-1)
-      new_v_col = decay_rate * v_col + mixing_rate * np.mean(grads_sqr, axis=-2)
-      updates.extend([new_v_row, new_v_col])
-      row_col_mean = np.mean(new_v_row, axis=-1, keepdims=True)
-      row_factor = (new_v_row / row_col_mean)**-0.5
-      col_factor = (new_v_col)**-0.5
-      y = (
-          grads * np.expand_dims(row_factor, axis=-1) *
-          np.expand_dims(col_factor, axis=-2))
-    else:
-      v = slots.pop(0)
-      new_v = decay_rate * v + mixing_rate * grads_sqr
-      updates.append(new_v)
-      y = grads * (new_v)**-0.5
-
-    if self._do_clipping:
-      clipping_denom = (
-          np.maximum(1.0, np.sqrt(np.mean(y * y)) / clipping_threshold))
-      y /= clipping_denom
-
-    subtrahend = update_scale * y
-    if self._do_momentum:
-      m = slots.pop(0)
-      new_m = beta1 * m + (1.0 - beta1) * subtrahend
-      subtrahend = new_m
-      updates.append(new_m)
-
-    new_params = (1 - weight_decay_rate) * params - subtrahend
-    # TODO(lukaszkaiser): why is the astype needed here? Check and correct.
-    return new_params.astype(params.dtype), updates
-
-
-class SM3(Optimizer):
-  """SM3 optimizer."""
-
-  def __init__(self, learning_rate, momentum=0.9):  # pylint: disable=useless-super-delegation
-    """Create the SM3 optimizer.
-
-    Memory-Efficient Adaptive Optimization for Large-Scale Learning.
-    https://arxiv.org/abs/1901.11150
-
-    Args:
-      learning_rate: a postitive scalar value for the initial learning rate.
-      momentum: optional, a positive scalar value for momentum
-    """
-    super(SM3, self).__init__(
-        learning_rate=learning_rate,
-        momentum=momentum,
-    )
-
-  def init(self, params):
-    vs = [np.zeros(sz, dtype=params.dtype) for sz in params.shape]
-    return (np.zeros_like(params), vs)
-
-  def _update_diagonal(self, grads, params, m, v, opt_params):
-    learning_rate = opt_params["learning_rate"]
-    momentum = opt_params["momentum"]
-    v[0] += grads * grads
-    preconditioner = np.where(v[0] > 0, 1.0 / np.sqrt(v[0]),
-                              np.zeros_like(v[0]))
-    preconditioned_grads = preconditioner * grads
-    m = (1 - momentum) * preconditioned_grads + momentum * m
-    params = params - (learning_rate * m).astype(params.dtype)
-    return params, (m, v)
-
-  def _expanded_shape(self, shape, axis):
-    # Replaces a `shape` of [M, N, K] with 1 in all dimensions except for i.
-    # For eg: i = 1 returns [1, N, 1].
-    rank = len(shape)
-    return [1] * axis + [shape[axis]] + [1] * (rank - axis - 1)
-
-  def _minimum(self, tensor_list):
-    minimum = tensor_list[0]
-    for i in range(1, len(tensor_list)):
-      minimum = np.minimum(minimum, tensor_list[i])
-    return minimum
-
-  def _update_sketched(self, grads, params, m, v, opt_params):
-    """Update for higher-rank parameters."""
-    learning_rate = opt_params["learning_rate"]
-    momentum = opt_params["momentum"]
-    shape = params.shape
-    rank = len(shape)
-    reshaped_accumulators = [np.reshape(v[i], self._expanded_shape(shape, i))
-                             for i in range(rank)]
-    current_accumulator = self._minimum(reshaped_accumulators)
-    current_accumulator += grads * grads
-    accumulator_inv_sqrt = np.where(current_accumulator > 0.0,
-                                    1.0 / np.sqrt(current_accumulator),
-                                    np.zeros_like(current_accumulator))
-    preconditioned_gradient = grads * accumulator_inv_sqrt
-    m = (1.0 - momentum) * preconditioned_gradient + momentum * m
-    params = params - (learning_rate * m).astype(params.dtype)
-    for i in range(len(v)):
-      axes = list(range(int(i))) + list(range(int(i) + 1, rank))
-      dim_accumulator = np.amax(current_accumulator, axis=axes)
-      v[i] = dim_accumulator
-    return params, (m, v)
-
-  def update(self, step, grads, params, slots, opt_params):
-    del step
-    m, v = slots
-    shape = params.shape
-    rank = len(shape)
-    if rank > 1:
-      return self._update_sketched(grads, params, m, v, opt_params)
-    else:
-      return self._update_diagonal(grads, params, m, v, opt_params)
diff --git a/tensor2tensor/trax/rl/__init__.py b/tensor2tensor/trax/rl/__init__.py
deleted file mode 100644
index 529ab0e55..000000000
--- a/tensor2tensor/trax/rl/__init__.py
+++ /dev/null
@@ -1,52 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Trax RL library."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gin
-
-from tensor2tensor.trax.rl import simulated_env_problem
-
-
-def configure_rl(*args, **kwargs):
-  kwargs["module"] = "trax.rl"
-  return gin.external_configurable(*args, **kwargs)
-
-
-def configure_simulated_env_problem(*args, **kwargs):
-  kwargs["blacklist"] = [
-      "batch_size", "observation_space", "action_space", "reward_range",
-      "discrete_rewards", "history_stream", "output_dir"]
-  return configure_rl(*args, **kwargs)
-
-
-# pylint: disable=invalid-name
-RawSimulatedEnvProblem = configure_simulated_env_problem(
-    simulated_env_problem.RawSimulatedEnvProblem)
-SerializedSequenceSimulatedEnvProblem = configure_simulated_env_problem(
-    simulated_env_problem.SerializedSequenceSimulatedEnvProblem)
-
-
-# pylint: disable=invalid-name
-cartpole_done_fn = configure_rl(simulated_env_problem.cartpole_done_fn)
-cartpole_reward_fn = configure_rl(simulated_env_problem.cartpole_reward_fn)
-acrobot_done_fn = configure_rl(simulated_env_problem.acrobot_done_fn)
-acrobot_reward_fn = configure_rl(simulated_env_problem.acrobot_reward_fn)
-onlinetune_done_fn = configure_rl(simulated_env_problem.onlinetune_done_fn)
-onlinetune_reward_fn = configure_rl(simulated_env_problem.onlinetune_reward_fn)
diff --git a/tensor2tensor/trax/rl/base_trainer.py b/tensor2tensor/trax/rl/base_trainer.py
deleted file mode 100644
index 39e79f8fd..000000000
--- a/tensor2tensor/trax/rl/base_trainer.py
+++ /dev/null
@@ -1,141 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 class for RL trainers."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-
-from absl import logging
-from tensor2tensor.trax import utils
-from tensorflow.io import gfile
-
-
-class BaseTrainer(object):
-  """Base class for RL trainers."""
-
-  def __init__(
-      self, train_env, eval_env, output_dir,
-      trajectory_dump_dir=None, trajectory_dump_min_count_per_shard=16,
-      async_mode=False,
-  ):
-    """Base class constructor.
-
-    Args:
-      train_env: EnvProblem to use for training. Settable.
-      eval_env: EnvProblem to use for evaluation. Settable.
-      output_dir: Directory to save checkpoints and metrics to.
-      trajectory_dump_dir: Directory to dump trajectories to. Trajectories
-        are saved in shards of name <epoch>.pkl under this directory. Settable.
-      trajectory_dump_min_count_per_shard: Minimum number of trajectories to
-        collect before dumping in a new shard. Sharding is for efficient
-        shuffling for model training in SimPLe.
-      async_mode: (bool) If True, this means we are in async mode and we read
-        trajectories from a location rather than interact with the environment.
-    """
-    self.train_env = train_env
-    self.eval_env = eval_env
-    self._output_dir = output_dir
-    gfile.makedirs(self._output_dir)
-    self.trajectory_dump_dir = trajectory_dump_dir
-    self._trajectory_dump_min_count_per_shard = (
-        trajectory_dump_min_count_per_shard)
-    self._trajectory_buffer = []
-    self._async_mode = async_mode
-
-  @property
-  def async_mode(self):
-    return self._async_mode
-
-  @async_mode.setter
-  def async_mode(self, async_mode):
-    logging.vlog(1, "Changing async mode from %s to: %s",
-                 self._async_mode, async_mode)
-    self._async_mode = async_mode
-
-  @property
-  def epoch(self):
-    raise NotImplementedError
-
-  def train_epoch(self, evaluate=True):
-    raise NotImplementedError
-
-  def evaluate(self):
-    raise NotImplementedError
-
-  def save(self):
-    raise NotImplementedError
-
-  def flush_summaries(self):
-    raise NotImplementedError
-
-  def dump_trajectories(self, force=False):
-    """Dumps trajectories in a new shard.
-
-    Should be called at most once per epoch.
-
-    Args:
-      force: (bool) Whether to complete unfinished trajectories and create
-        a new shard even if we have not reached the minimum size.
-    """
-    pkl_module = utils.get_pickle_module()
-    if self.trajectory_dump_dir is None:
-      return
-    gfile.makedirs(self.trajectory_dump_dir)
-
-    trajectories = self.train_env.trajectories
-    if force:
-      trajectories.complete_all_trajectories()
-
-    # complete_all_trajectories() also adds trajectories that were just reset.
-    # We don't want them since they have just the initial observation and no
-    # actions, so we filter them out.
-    def has_any_action(trajectory):
-      return (
-          trajectory.time_steps and trajectory.time_steps[0].action is not None)
-    self._trajectory_buffer.extend(
-        filter(has_any_action, trajectories.completed_trajectories))
-
-    trajectories.clear_completed_trajectories()
-    ready = (
-        len(self._trajectory_buffer) >=
-        self._trajectory_dump_min_count_per_shard
-    )
-    if ready or force:
-      shard_path = os.path.join(
-          self.trajectory_dump_dir, "{}.pkl".format(self.epoch))
-      if gfile.exists(shard_path):
-        # Since we do an extra dump at the end of the training loop, we
-        # sometimes dump 2 times in the same epoch. When this happens, merge the
-        # two sets of trajectories.
-        with gfile.GFile(shard_path, "rb") as f:
-          self._trajectory_buffer = pkl_module.load(f) + self._trajectory_buffer
-      with gfile.GFile(shard_path, "wb") as f:
-        pkl_module.dump(self._trajectory_buffer, f)
-      self._trajectory_buffer = []
-
-  def training_loop(self, n_epochs, evaluate=True):
-    logging.info("Starting the RL training loop.")
-    for _ in range(self.epoch, n_epochs):
-      self.train_epoch(evaluate=evaluate)
-      self.dump_trajectories()
-    self.save()
-    self.dump_trajectories(force=True)
-    if evaluate:
-      self.evaluate()
-    self.flush_summaries()
diff --git a/tensor2tensor/trax/rl/base_trainer_test.py b/tensor2tensor/trax/rl/base_trainer_test.py
deleted file mode 100644
index 13307c77c..000000000
--- a/tensor2tensor/trax/rl/base_trainer_test.py
+++ /dev/null
@@ -1,144 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.rl.base_trainer."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-
-import cloudpickle as pickle
-import numpy as np
-
-from tensor2tensor.envs import gym_env_problem
-from tensor2tensor.trax.rl import base_trainer
-from tensorflow import test
-
-
-class FakeTrainer(base_trainer.BaseTrainer):
-  """Fake Trainer.
-
-  Adds one complete and one incomplete trajectory every epoch.
-  """
-
-  def __init__(self, *args, **kwargs):
-    super(FakeTrainer, self).__init__(*args, **kwargs)
-    self._epoch = 0
-    self._should_reset = True
-
-  @property
-  def epoch(self):
-    return self._epoch
-
-  def train_epoch(self):
-    trajectories = self.train_env.trajectories
-    if self._should_reset:
-      trajectories.reset(indices=np.arange(2), observations=np.zeros(2))
-    self._should_reset = False
-    trajectories.step(
-        observations=np.zeros(2),
-        raw_rewards=np.zeros(2),
-        processed_rewards=np.zeros(2),
-        dones=np.array([False, True]),
-        actions=np.zeros(2),
-    )
-    # Reset the trajectories that are done, as
-    # env_problem_utils.play_env_problem_with_policy does.
-    trajectories.reset(indices=np.array([1]), observations=np.zeros(1))
-    self._epoch += 1
-
-  def evaluate(self):
-    pass
-
-  def save(self):
-    pass
-
-  def flush_summaries(self):
-    pass
-
-
-class BaseTrainerTest(test.TestCase):
-
-  def _make_trainer(self, min_count_per_shard):
-    train_env = gym_env_problem.GymEnvProblem(
-        base_env_name="Acrobot-v1", batch_size=2)
-    eval_env = gym_env_problem.GymEnvProblem(
-        base_env_name="Acrobot-v1", batch_size=1)
-    temp_dir = self.get_temp_dir()
-    return FakeTrainer(
-        train_env, eval_env,
-        output_dir=temp_dir,
-        trajectory_dump_dir=temp_dir,
-        trajectory_dump_min_count_per_shard=min_count_per_shard,
-    )
-
-  def _assert_no_shard_exists(self, trajectory_dir):
-    self.assertFalse(os.listdir(trajectory_dir))
-
-  def _assert_single_shard_exists_and_has_trajectories(
-      self, trajectory_dir, expected_trajectory_lengths):
-    shard_filenames = os.listdir(trajectory_dir)
-    self.assertEqual(len(shard_filenames), 1)
-    shard_path = os.path.join(trajectory_dir, shard_filenames[0])
-    with open(shard_path, "rb") as f:
-      trajectories = pickle.load(f)
-    actual_trajectory_lengths = [
-        len(trajectory.time_steps) for trajectory in trajectories]
-    self.assertEqual(
-        list(sorted(actual_trajectory_lengths)),
-        list(sorted(expected_trajectory_lengths)),
-    )
-
-  def test_dumps_full_shard(self):
-    trainer = self._make_trainer(min_count_per_shard=2)
-    trajectory_dir = self.get_temp_dir()
-
-    # Add one complete trajectory to the buffer. Should not dump yet.
-    trainer.train_epoch()
-    trainer.dump_trajectories()
-    self._assert_no_shard_exists(trajectory_dir)
-
-    # Add the second complete trajectory. Now we should dump.
-    trainer.train_epoch()
-    trainer.dump_trajectories()
-    self._assert_single_shard_exists_and_has_trajectories(
-        trajectory_dir, [2, 2])
-
-  def test_dumps_incomplete_trajectories_when_force_is_true(self):
-    trainer = self._make_trainer(min_count_per_shard=2)
-    trajectory_dir = self.get_temp_dir()
-
-    # Add one complete and one incomplete trajectory to the buffer. Should dump.
-    trainer.train_epoch()
-    trainer.dump_trajectories(force=True)
-    self._assert_single_shard_exists_and_has_trajectories(
-        trajectory_dir, [2, 2])
-
-  def test_dumps_incomplete_shard_when_force_is_true(self):
-    trainer = self._make_trainer(min_count_per_shard=4)
-    trajectory_dir = self.get_temp_dir()
-
-    # Add one complete and one incomplete trajectory to the buffer. Should dump,
-    # even though we don't have a full shard yet.
-    trainer.train_epoch()
-    trainer.dump_trajectories(force=True)
-    self._assert_single_shard_exists_and_has_trajectories(
-        trajectory_dir, [2, 2])
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/rl/configs/acrobot.gin b/tensor2tensor/trax/rl/configs/acrobot.gin
deleted file mode 100644
index e122270ae..000000000
--- a/tensor2tensor/trax/rl/configs/acrobot.gin
+++ /dev/null
@@ -1,30 +0,0 @@
-import tensor2tensor.trax.models
-import tensor2tensor.trax.rl.trainers
-
-# Parameters for FrameStackMLP:
-# ==============================================================================
-FrameStackMLP.n_frames = 1
-
-# Parameters for PPO:
-# ==============================================================================
-PPO.n_optimizer_steps = 32
-PPO.target_kl = 1000  # Virtually infinite.
-PPO.boundary = 512
-PPO.max_timestep = 512
-PPO.max_timestep_eval = 20000
-PPO.random_seed = None
-PPO.gamma = 0.99
-PPO.lambda_ = 0.95
-PPO.c1 = 1.0
-PPO.c2 = 0.0
-PPO.eval_every_n = 500
-PPO.done_frac_for_policy_save = 0.9
-PPO.n_evals = 16
-PPO.len_history_for_policy = 1
-PPO.eval_temperatures = (1.0, 0.5)
-PPO.policy_and_value_model = @trax.models.FrameStackMLP
-
-# Parameters for train_rl:
-# ==============================================================================
-train_rl.env_name = "Acrobot-v1"
-train_rl.n_epochs = 40000
diff --git a/tensor2tensor/trax/rl/configs/acrobot_transformer.gin b/tensor2tensor/trax/rl/configs/acrobot_transformer.gin
deleted file mode 100644
index 16343023d..000000000
--- a/tensor2tensor/trax/rl/configs/acrobot_transformer.gin
+++ /dev/null
@@ -1,34 +0,0 @@
-import tensor2tensor.trax.models
-import tensor2tensor.trax.rl.trainers
-
-# Parameters for TransformerDecoder:
-# ==============================================================================
-TransformerDecoder.d_model = 64
-TransformerDecoder.d_ff = 128
-TransformerDecoder.dropout = 0.1
-TransformerDecoder.n_heads = 2
-TransformerDecoder.n_layers = 1
-
-# Parameters for PPO:
-# ==============================================================================
-PPO.n_optimizer_steps = 32
-PPO.target_kl = 1000  # Virtually infinite.
-PPO.boundary = 512
-PPO.max_timestep = 512
-PPO.max_timestep_eval = 20000
-PPO.random_seed = None
-PPO.gamma = 0.99
-PPO.lambda_ = 0.95
-PPO.c1 = 1.0
-PPO.c2 = 0.0
-PPO.eval_every_n = 500
-PPO.done_frac_for_policy_save = 0.9
-PPO.n_evals = 16
-PPO.len_history_for_policy = None
-PPO.eval_temperatures = (1.0, 0.5)
-PPO.policy_and_value_model = @trax.models.TransformerDecoder
-
-# Parameters for train_rl:
-# ==============================================================================
-train_rl.env_name = "Acrobot-v1"
-train_rl.n_epochs = 40000
diff --git a/tensor2tensor/trax/rl/configs/atari.gin b/tensor2tensor/trax/rl/configs/atari.gin
deleted file mode 100644
index e4c0ab2c1..000000000
--- a/tensor2tensor/trax/rl/configs/atari.gin
+++ /dev/null
@@ -1,30 +0,0 @@
-import tensor2tensor.trax.models
-import tensor2tensor.trax.rl.trainers
-
-# Parameters for PPO:
-# ==============================================================================
-PPO.n_optimizer_steps = 4
-PPO.target_kl = 0.01
-PPO.boundary = 20
-PPO.max_timestep = 128
-PPO.max_timestep_eval = 20000
-PPO.random_seed = None
-PPO.gamma = 0.99
-PPO.lambda_ = 0.95
-PPO.c1 = 1.0
-PPO.c2 = 0.01
-PPO.eval_every_n = 500
-PPO.done_frac_for_policy_save = 0.9
-PPO.n_evals = 16
-PPO.len_history_for_policy = 4
-PPO.eval_temperatures = (1.0, 0.5)
-PPO.policy_and_value_model = @trax.models.AtariCnn
-
-# Parameters for train_rl:
-# ==============================================================================
-train_rl.env_name = "PongNoFrameskip-v4"
-train_rl.n_epochs = 40000
-train_rl.clip_rewards = True
-train_rl.max_timestep = 10000
-train_rl.rendered_env = True
-train_rl.resize_dims = (105, 80)
diff --git a/tensor2tensor/trax/rl/configs/atari_regression_test.gin b/tensor2tensor/trax/rl/configs/atari_regression_test.gin
deleted file mode 100644
index 54a96d953..000000000
--- a/tensor2tensor/trax/rl/configs/atari_regression_test.gin
+++ /dev/null
@@ -1,30 +0,0 @@
-import tensor2tensor.trax.models
-import tensor2tensor.trax.rl.trainers
-
-# Parameters for PPO:
-# ==============================================================================
-PPO.n_optimizer_steps = 30
-PPO.target_kl = 0.01
-PPO.boundary = 20
-PPO.max_timestep = 128
-PPO.max_timestep_eval = 20000
-PPO.random_seed = None
-PPO.gamma = 0.99
-PPO.lambda_ = 0.95
-PPO.c1 = 1.0
-PPO.c2 = 0.01
-PPO.eval_every_n = 500
-PPO.done_frac_for_policy_save = 0.9
-PPO.n_evals = 16
-PPO.len_history_for_policy = 4
-PPO.eval_temperatures = (1.0, 0.5)
-PPO.policy_and_value_model = @trax.models.AtariCnn
-
-# Parameters for train_rl:
-# ==============================================================================
-train_rl.env_name = "PongNoFrameskip-v4"
-train_rl.n_epochs = 4000
-train_rl.clip_rewards = True
-train_rl.max_timestep = 10000
-train_rl.rendered_env = True
-train_rl.resize_dims = (105, 80)
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
deleted file mode 100644
index 3aeb46b66..000000000
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_imagenet64_16gb.gin
+++ /dev/null
@@ -1,105 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.rl
-import tensor2tensor.trax.rl.envs
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 1
-batch_fn.eval_batch_size = 16
-batch_fn.max_eval_length = 12288  # 64 * 64 * 3
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_image_imagenet64_gen_flat_rev'
-inputs.input_name = 'targets'
-
-# Parameters for train_and_eval_dataset:
-# ==============================================================================
-train_and_eval_dataset.eval_holdout_size = 0.05
-train_and_eval_dataset.eval_shuffle_files = True
-
-# Parameters for MemoryEfficientCausalAttention:
-# ==============================================================================
-MemoryEfficientCausalAttention.dropout = 0.0
-MemoryEfficientCausalAttention.loop_stride = 512
-
-# Parameters for MergedHashedCausalAttention:
-# ==============================================================================
-MergedHashedCausalAttention.dropout = 0.0
-MergedHashedCausalAttention.n_bins = 16
-MergedHashedCausalAttention.bin_by_time = True
-MergedMultiHashedCausalAttention.one_rng = False
-
-# Parameters for MergedMultiHashedCausalAttention:
-# ==============================================================================
-MergedMultiHashedCausalAttention.dropout = 0.0
-MergedMultiHashedCausalAttention.n_bins = 64
-MergedMultiHashedCausalAttention.n_hashes = 2
-MergedMultiHashedCausalAttention.n_buckets_per_bin = 2
-MergedMultiHashedCausalAttention.bin_by_time = False
-MergedMultiHashedCausalAttention.one_rng = False
-MergedMultiHashedCausalAttention.drop_for_hash_rate = 0.1
-MergedMultiHashedCausalAttention.hard_k = 32
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = @trax.layers.MergedMultiHashedCausalAttention
-TransformerLM.d_attention_key = 64
-TransformerLM.d_attention_value = 64
-TransformerLM.d_model = 1024
-TransformerLM.d_ff = 2048
-TransformerLM.dropout = 0.0
-TransformerLM.max_len = 12288  # 64 * 64 * 3
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 4
-TransformerLM.n_layers = 3
-TransformerLM.share_kv = True
-TransformerLM.vocab_size = 256
-
-# Parameters for OnlineTuneEnv:
-# ==============================================================================
-OnlineTuneEnv.inputs = @trax.inputs.inputs
-OnlineTuneEnv.model = @trax.models.TransformerLM
-OnlineTuneEnv.optimizer = @trax.optimizers.Adafactor
-OnlineTuneEnv.control_configs = (
-    ("learning_rate", 1e-3, (1e-9, 1e-2), False),
-    ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
-
-    ("dropout_embedding", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_initial", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_initial", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_initial", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_middle", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_middle", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_middle", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_final", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_final", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_final", 0.1, (0.0, 0.9), True),
-)
-OnlineTuneEnv.nontrainable_param_map = {
-    # "dropout_{layer_type}_{block_index}": "dropout_{layer_type}_{block_group}"
-
-    "dropout_attention_0": "dropout_attention_initial",
-    "dropout_ff_middle_0": "dropout_ff_middle_initial",
-    "dropout_ff_final_0": "dropout_ff_final_initial",
-
-    "dropout_attention_1": "dropout_attention_middle",
-    "dropout_ff_middle_1": "dropout_ff_middle_middle",
-    "dropout_ff_final_1": "dropout_ff_final_middle",
-
-    "dropout_attention_2": "dropout_attention_final",
-    "dropout_ff_middle_2": "dropout_ff_middle_final",
-    "dropout_ff_final_2": "dropout_ff_final_final",
-}
-OnlineTuneEnv.include_controls_in_observation = False
-OnlineTuneEnv.train_steps = 150
-OnlineTuneEnv.eval_steps = 2
-OnlineTuneEnv.env_steps = 100
-OnlineTuneEnv.observation_range = (0.0, 10.0)
-OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
deleted file mode 100644
index c765da790..000000000
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm1b_16gb.gin
+++ /dev/null
@@ -1,99 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.rl
-import tensor2tensor.trax.rl.envs
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 256
-batch_fn.eval_batch_size = 256
-batch_fn.max_eval_length = 2048
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_languagemodel_lm1b32k'
-inputs.input_name = 'targets'
-
-# Parameters for train_and_eval_dataset:
-# ==============================================================================
-train_and_eval_dataset.eval_holdout_size = 0.05
-train_and_eval_dataset.eval_shuffle_files = True
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
-lm1b_preprocess.max_target_length = 512
-lm1b_preprocess.max_eval_target_length = 2048
-
-# Parameters for DotProductCausalAttention:
-# ==============================================================================
-DotProductCausalAttention.dropout = 0.1
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
-TransformerLM.d_model = 512
-TransformerLM.d_ff = 2048
-TransformerLM.dropout = 0.1
-TransformerLM.max_len = 2048
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 8
-TransformerLM.n_layers = 6
-TransformerLM.vocab_size = 32000
-
-# Parameters for OnlineTuneEnv:
-# ==============================================================================
-OnlineTuneEnv.inputs = @trax.inputs.inputs
-OnlineTuneEnv.model = @trax.models.TransformerLM
-OnlineTuneEnv.optimizer = @trax.optimizers.Adafactor
-OnlineTuneEnv.train_steps = 500
-OnlineTuneEnv.eval_steps = 1
-OnlineTuneEnv.env_steps = 100
-OnlineTuneEnv.control_configs = (
-    ("learning_rate", 1e-3, (1e-9, 1e-2), False),
-    ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
-
-    ("dropout_embedding", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_initial", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_initial", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_initial", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_middle", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_middle", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_middle", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_final", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_final", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_final", 0.1, (0.0, 0.9), True),
-)
-OnlineTuneEnv.nontrainable_param_map = {
-    # "dropout_{layer_type}_{block_index}": "dropout_{layer_type}_{block_group}"
-
-    "dropout_attention_0": "dropout_attention_initial",
-    "dropout_ff_middle_0": "dropout_ff_middle_initial",
-    "dropout_ff_final_0": "dropout_ff_final_initial",
-
-    "dropout_attention_1": "dropout_attention_middle",
-    "dropout_ff_middle_1": "dropout_ff_middle_middle",
-    "dropout_ff_final_1": "dropout_ff_final_middle",
-    "dropout_attention_2": "dropout_attention_middle",
-    "dropout_ff_middle_2": "dropout_ff_middle_middle",
-    "dropout_ff_final_2": "dropout_ff_final_middle",
-    "dropout_attention_3": "dropout_attention_middle",
-    "dropout_ff_middle_3": "dropout_ff_middle_middle",
-    "dropout_ff_final_3": "dropout_ff_final_middle",
-    "dropout_attention_4": "dropout_attention_middle",
-    "dropout_ff_middle_4": "dropout_ff_middle_middle",
-    "dropout_ff_final_4": "dropout_ff_final_middle",
-
-    "dropout_attention_5": "dropout_attention_final",
-    "dropout_ff_middle_5": "dropout_ff_middle_final",
-    "dropout_ff_final_5": "dropout_ff_final_final",
-}
-OnlineTuneEnv.include_controls_in_observation = False
-OnlineTuneEnv.observation_range = (0.0, 10.0)
-OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
-OnlineTuneEnv.mask_id = 0
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
deleted file mode 100644
index f4cef2897..000000000
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_lm_wmt_ende_16gb.gin
+++ /dev/null
@@ -1,96 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.rl
-import tensor2tensor.trax.rl.envs
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 128
-batch_fn.eval_batch_size = 128
-batch_fn.bucket_length = 64
-batch_fn.max_eval_length = 512
-batch_fn.buckets_include_inputs_in_length = True
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_translate_ende_wmt32k'
-
-# Parameters for train_and_eval_dataset:
-# ==============================================================================
-train_and_eval_dataset.eval_holdout_size = 0.05
-train_and_eval_dataset.eval_shuffle_files = True
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.wmt_concat_preprocess
-wmt_concat_preprocess.max_length = 255
-wmt_concat_preprocess.max_eval_length = 255
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.d_model = 512
-TransformerLM.d_ff = 2048
-TransformerLM.dropout = 0.1
-TransformerLM.max_len = 2048
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 8
-TransformerLM.n_layers = 6
-TransformerLM.vocab_size = 33300
-
-# Parameters for OnlineTuneEnv:
-# ==============================================================================
-OnlineTuneEnv.inputs = @trax.inputs.inputs
-OnlineTuneEnv.model = @trax.models.TransformerLM
-OnlineTuneEnv.optimizer = @trax.optimizers.Adafactor
-OnlineTuneEnv.train_steps = 500
-OnlineTuneEnv.eval_steps = 1
-OnlineTuneEnv.env_steps = 100
-OnlineTuneEnv.has_weights = True
-OnlineTuneEnv.control_configs = (
-    ("learning_rate", 1e-3, (1e-9, 1e-2), False),
-    ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
-
-    ("dropout_embedding", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_initial", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_initial", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_initial", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_middle", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_middle", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_middle", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_final", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_final", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_final", 0.1, (0.0, 0.9), True),
-)
-OnlineTuneEnv.nontrainable_param_map = {
-    # "dropout_{layer_type}_{block_index}": "dropout_{layer_type}_{block_group}"
-
-    "dropout_attention_0": "dropout_attention_initial",
-    "dropout_ff_middle_0": "dropout_ff_middle_initial",
-    "dropout_ff_final_0": "dropout_ff_final_initial",
-
-    "dropout_attention_1": "dropout_attention_middle",
-    "dropout_ff_middle_1": "dropout_ff_middle_middle",
-    "dropout_ff_final_1": "dropout_ff_final_middle",
-    "dropout_attention_2": "dropout_attention_middle",
-    "dropout_ff_middle_2": "dropout_ff_middle_middle",
-    "dropout_ff_final_2": "dropout_ff_final_middle",
-    "dropout_attention_3": "dropout_attention_middle",
-    "dropout_ff_middle_3": "dropout_ff_middle_middle",
-    "dropout_ff_final_3": "dropout_ff_final_middle",
-    "dropout_attention_4": "dropout_attention_middle",
-    "dropout_ff_middle_4": "dropout_ff_middle_middle",
-    "dropout_ff_final_4": "dropout_ff_final_middle",
-
-    "dropout_attention_5": "dropout_attention_final",
-    "dropout_ff_middle_5": "dropout_ff_middle_final",
-    "dropout_ff_final_5": "dropout_ff_final_final",
-}
-OnlineTuneEnv.include_controls_in_observation = False
-OnlineTuneEnv.observation_range = (0.0, 10.0)
-OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
-OnlineTuneEnv.mask_id = 0
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin
deleted file mode 100644
index 34e42e9cb..000000000
--- a/tensor2tensor/trax/rl/configs/env_online_tune_transformer_ptb_16gb.gin
+++ /dev/null
@@ -1,94 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.rl
-import tensor2tensor.trax.rl.envs
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 64
-batch_fn.eval_batch_size = 512
-batch_fn.max_eval_length = 2048
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 't2t_languagemodel_ptb10k'
-inputs.input_name = 'targets'
-
-# Parameters for preprocess_fun:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
-lm1b_preprocess.max_target_length = 512
-lm1b_preprocess.max_eval_target_length = 2048
-
-# Parameters for DotProductCausalAttention:
-# ==============================================================================
-DotProductCausalAttention.dropout = 0.1
-
-# Parameters for TransformerLM:
-# ==============================================================================
-TransformerLM.attention_type = @trax.layers.DotProductCausalAttention
-TransformerLM.d_model = 512
-TransformerLM.d_ff = 2048
-TransformerLM.dropout = 0.1
-TransformerLM.max_len = 2048
-TransformerLM.mode = 'train'
-TransformerLM.n_heads = 8
-TransformerLM.n_layers = 6
-TransformerLM.vocab_size = 10240
-
-# Parameters for OnlineTuneEnv:
-# ==============================================================================
-OnlineTuneEnv.inputs = @trax.inputs.inputs
-OnlineTuneEnv.model = @trax.models.TransformerLM
-OnlineTuneEnv.optimizer = @trax.optimizers.Adafactor
-OnlineTuneEnv.train_steps = 200
-OnlineTuneEnv.eval_steps = 2
-OnlineTuneEnv.env_steps = 100
-OnlineTuneEnv.control_configs = (
-    ("learning_rate", 1e-3, (1e-9, 1e-2), False),
-    ("weight_decay_rate", 1e-5, (1e-9, 1e-3), False),
-
-    ("dropout_embedding", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_initial", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_initial", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_initial", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_middle", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_middle", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_middle", 0.1, (0.0, 0.9), True),
-
-    ("dropout_attention_final", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_middle_final", 0.1, (0.0, 0.9), True),
-    ("dropout_ff_final_final", 0.1, (0.0, 0.9), True),
-)
-OnlineTuneEnv.nontrainable_param_map = {
-    # "dropout_{layer_type}_{block_index}": "dropout_{layer_type}_{block_group}"
-
-    "dropout_attention_0": "dropout_attention_initial",
-    "dropout_ff_middle_0": "dropout_ff_middle_initial",
-    "dropout_ff_final_0": "dropout_ff_final_initial",
-
-    "dropout_attention_1": "dropout_attention_middle",
-    "dropout_ff_middle_1": "dropout_ff_middle_middle",
-    "dropout_ff_final_1": "dropout_ff_final_middle",
-    "dropout_attention_2": "dropout_attention_middle",
-    "dropout_ff_middle_2": "dropout_ff_middle_middle",
-    "dropout_ff_final_2": "dropout_ff_final_middle",
-    "dropout_attention_3": "dropout_attention_middle",
-    "dropout_ff_middle_3": "dropout_ff_middle_middle",
-    "dropout_ff_final_3": "dropout_ff_final_middle",
-    "dropout_attention_4": "dropout_attention_middle",
-    "dropout_ff_middle_4": "dropout_ff_middle_middle",
-    "dropout_ff_final_4": "dropout_ff_final_middle",
-
-    "dropout_attention_5": "dropout_attention_final",
-    "dropout_ff_middle_5": "dropout_ff_middle_final",
-    "dropout_ff_final_5": "dropout_ff_final_final",
-}
-OnlineTuneEnv.include_controls_in_observation = False
-OnlineTuneEnv.observation_range = (0.0, 10.0)
-OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
-OnlineTuneEnv.mask_id = 0
diff --git a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10_8gb.gin b/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10_8gb.gin
deleted file mode 100644
index 9b8a8e6d5..000000000
--- a/tensor2tensor/trax/rl/configs/env_online_tune_wide_resnet_cifar10_8gb.gin
+++ /dev/null
@@ -1,52 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.rl
-import tensor2tensor.trax.rl.envs
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size_per_device = 256
-batch_fn.bucket_length = 32
-batch_fn.buckets = None
-batch_fn.eval_batch_size = 512
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 'cifar10'
-
-# Parameters for train_and_eval_dataset:
-# ==============================================================================
-train_and_eval_dataset.eval_holdout_size = 0.05
-train_and_eval_dataset.eval_shuffle_files = True
-
-# Parameters for Momentum:
-# ==============================================================================
-Momentum.mass = 0.9
-
-# Parameters for shuffle_and_batch_data:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun = @trax.inputs.cifar10_augmentation_preprocess
-
-# Parameters for WideResnet:
-# ==============================================================================
-WideResnet.widen_factor = 10
-WideResnet.n_blocks = 4
-WideResnet.n_output_classes = 10
-
-# Parameters for OnlineTuneEnv:
-# ==============================================================================
-OnlineTuneEnv.inputs = @trax.inputs.inputs
-OnlineTuneEnv.model = @trax.models.WideResnet
-OnlineTuneEnv.optimizer = @trax.optimizers.Momentum
-OnlineTuneEnv.control_configs = (
-    ("learning_rate", 0.1, (1e-9, 10.0), False),
-    ("weight_decay_rate", 1e-5, (1e-9, 0.1), False),
-    ("mass", 0.9, (0.0, 0.99), True),
-)
-OnlineTuneEnv.include_controls_in_observation = False
-OnlineTuneEnv.action_multipliers = (0.5, 0.8, 0.95, 1.0, 1.05, 1.25, 2.0)
-OnlineTuneEnv.train_steps = 100
-OnlineTuneEnv.eval_steps = 10
-OnlineTuneEnv.env_steps = 100
diff --git a/tensor2tensor/trax/rl/configs/ppo_online_tune.gin b/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
deleted file mode 100644
index bfaa6e4fe..000000000
--- a/tensor2tensor/trax/rl/configs/ppo_online_tune.gin
+++ /dev/null
@@ -1,37 +0,0 @@
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.rl.trainers
-
-# Parameters for Adam:
-# ==============================================================================
-Adam.learning_rate = 1e-3
-Adam.b1 = 0.9
-Adam.b2 = 0.999
-Adam.weight_decay_rate = 0.0
-
-# Parameters for TransformerDecoder:
-# ==============================================================================
-TransformerDecoder.d_model = 64
-TransformerDecoder.d_ff = 128
-TransformerDecoder.dropout = 0.0
-TransformerDecoder.n_heads = 2
-TransformerDecoder.n_layers = 1
-
-# Parameters for PPO:
-# ==============================================================================
-PPO.n_optimizer_steps = 10
-PPO.target_kl = 0.1
-PPO.boundary = 128
-PPO.max_timestep = 128
-PPO.max_timestep_eval = 128
-PPO.random_seed = None
-PPO.gamma = 1.0
-PPO.lambda_ = 0.95
-PPO.c1 = 1.0
-PPO.c2 = 0.1
-PPO.done_frac_for_policy_save = 0
-PPO.len_history_for_policy = None
-PPO.separate_eval = False
-PPO.save_every_n = 1
-PPO.policy_and_value_model = @trax.models.TransformerDecoder
-PPO.policy_and_value_optimizer = @trax.optimizers.Adam
diff --git a/tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin b/tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin
deleted file mode 100644
index 040da7703..000000000
--- a/tensor2tensor/trax/rl/configs/ppo_online_tune_wide_resnet_cifar10.gin
+++ /dev/null
@@ -1,79 +0,0 @@
-import tensor2tensor.trax.inputs
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.rl.envs
-import tensor2tensor.trax.rl.trainers
-
-# Parameters for batch_fn:
-# ==============================================================================
-batch_fn.batch_size = 32
-batch_fn.bucket_length = 32
-batch_fn.buckets = None
-batch_fn.eval_batch_size = 32
-
-# Parameters for inputs:
-# ==============================================================================
-inputs.data_dir = None
-inputs.dataset_name = 'cifar10'
-
-# Parameters for Momentum:
-# ==============================================================================
-Momentum.mass = 0.9
-
-# Parameters for shuffle_and_batch_data:
-# ==============================================================================
-shuffle_and_batch_data.preprocess_fun = @trax.inputs.cifar10_no_augmentation_preprocess
-
-# Parameters for Adam:
-# ==============================================================================
-Adam.learning_rate = 1e-3
-Adam.b1 = 0.9
-Adam.b2 = 0.999
-Adam.weight_decay_rate = 0.0
-
-# Parameters for TransformerDecoder:
-# ==============================================================================
-TransformerDecoder.d_model = 64
-TransformerDecoder.d_ff = 128
-TransformerDecoder.dropout = 0.0
-TransformerDecoder.n_heads = 2
-TransformerDecoder.n_layers = 1
-
-# Parameters for WideResnet:
-# ==============================================================================
-WideResnet.widen_factor = 10
-WideResnet.n_blocks = 3
-WideResnet.n_output_classes = 10
-
-# Parameters for OnlineTuneEnv:
-# ==============================================================================
-OnlineTuneEnv.inputs = @trax.inputs.inputs
-OnlineTuneEnv.model = @trax.models.WideResnet
-OnlineTuneEnv.optimizer = @trax.optimizers.Momentum
-OnlineTuneEnv.start_lr = 0.01
-OnlineTuneEnv.train_steps = 500
-OnlineTuneEnv.eval_steps = 50
-OnlineTuneEnv.env_steps = 100
-
-# Parameters for PPO:
-# ==============================================================================
-PPO.n_optimizer_steps = 10
-PPO.target_kl = 0.1
-PPO.boundary = 128
-PPO.max_timestep = 128
-PPO.max_timestep_eval = 128
-PPO.random_seed = None
-PPO.gamma = 0.99
-PPO.lambda_ = 0.95
-PPO.c1 = 1.0
-PPO.c2 = 0.01
-PPO.done_frac_for_policy_save = 0
-PPO.len_history_for_policy = None
-PPO.separate_eval = False
-PPO.policy_and_value_model = @trax.models.TransformerDecoder
-PPO.policy_and_value_optimizer = @trax.optimizers.Adam
-
-# Parameters for train_rl:
-# ==============================================================================
-train_rl.env_name = "OnlineTuneEnv-v0"
-train_rl.n_epochs = 1000
diff --git a/tensor2tensor/trax/rl/configs/simple_online_tune.gin b/tensor2tensor/trax/rl/configs/simple_online_tune.gin
deleted file mode 100644
index ac67788f0..000000000
--- a/tensor2tensor/trax/rl/configs/simple_online_tune.gin
+++ /dev/null
@@ -1,95 +0,0 @@
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-import tensor2tensor.trax.rl
-import tensor2tensor.trax.rl.space_serializer
-import tensor2tensor.trax.rl.trainers
-
-# Parameters for BoxSpaceSerializer:
-# ==============================================================================
-BoxSpaceSerializer.precision = 2
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-world_model/MultifactorSchedule.constant = 1.0
-world_model/MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-world_model/MultifactorSchedule.warmup_steps = 10000
-
-# Parameters for Adam:
-# ==============================================================================
-Adam.learning_rate = 1e-3
-Adam.b1 = 0.9
-Adam.b2 = 0.999
-Adam.weight_decay_rate = 0.0
-
-# Parameters for TransformerDecoder:
-# ==============================================================================
-TransformerDecoder.d_model = 64
-TransformerDecoder.d_ff = 128
-TransformerDecoder.n_layers = 2
-TransformerDecoder.n_heads = 2
-TransformerDecoder.dropout = 0.0
-
-# Parameters for PPO:
-# ==============================================================================
-PPO.n_optimizer_steps = 10
-PPO.optimizer_batch_size = 128
-PPO.target_kl = 0.1
-PPO.boundary = 100
-PPO.max_timestep = 100
-PPO.max_timestep_eval = 100
-PPO.random_seed = None
-PPO.gamma = 1.0
-PPO.lambda_ = 0.95
-PPO.c1 = 1.0
-PPO.c2 = 0.1
-PPO.done_frac_for_policy_save = 0
-PPO.len_history_for_policy = None
-PPO.separate_eval = False
-PPO.save_every_n = 1
-PPO.policy_and_value_model = @trax.models.TransformerDecoder
-PPO.policy_and_value_optimizer = @trax.optimizers.Adam
-PPO.trajectory_dump_min_count_per_shard = 8
-PPO.print_every_optimizer_steps = 1
-
-## Parameters for TimeBinCausalAttention:
-## ==============================================================================
-world_model/TimeBinCausalAttention.dropout = 0.1
-world_model/TimeBinCausalAttention.bin_length = 512
-
-# Parameters for SerializedSequenceSimulatedEnvProblem:
-# ==============================================================================
-SerializedSequenceSimulatedEnvProblem.model = @world_model/trax.models.TransformerLM
-SerializedSequenceSimulatedEnvProblem.reward_fn = @trax.rl.onlinetune_reward_fn
-SerializedSequenceSimulatedEnvProblem.done_fn = @trax.rl.onlinetune_done_fn
-SerializedSequenceSimulatedEnvProblem.vocab_size = 128
-SerializedSequenceSimulatedEnvProblem.max_trajectory_length = 101
-SerializedSequenceSimulatedEnvProblem.significance_decay = 0.8
-
-# Parameters for SimPLe:
-# ==============================================================================
-SimPLe.policy_trainer_class = @trax.rl.trainers.PPO
-SimPLe.n_real_epochs = 1
-SimPLe.n_model_initial_train_steps = 50000
-SimPLe.n_model_train_steps_per_epoch = 10000
-SimPLe.model_train_batch_size = 64
-SimPLe.simulated_env_problem_class = @trax.rl.SerializedSequenceSimulatedEnvProblem
-SimPLe.simulated_batch_size = 128
-SimPLe.n_simulated_epochs = 50
-SimPLe.initial_trajectory_mix_prob = 0.9
-SimPLe.init_policy_from_world_model = False
-
-# Parameters for TransformerLM:
-# ==============================================================================
-world_model/TransformerLM.attention_type = @world_model/trax.layers.TimeBinCausalAttention
-world_model/TransformerLM.d_model = 256
-world_model/TransformerLM.d_ff = 512
-world_model/TransformerLM.n_layers = 3
-world_model/TransformerLM.n_heads = 4
-world_model/TransformerLM.dropout = 0.1
-world_model/TransformerLM.max_len = 2048
-
-# Parameters for train:
-# ==============================================================================
-world_model/train.eval_frequency = 1000
-world_model/train.optimizer = @trax.optimizers.Adafactor
diff --git a/tensor2tensor/trax/rl/configs/simple_online_tune_serialized.gin b/tensor2tensor/trax/rl/configs/simple_online_tune_serialized.gin
deleted file mode 100644
index 181efb64c..000000000
--- a/tensor2tensor/trax/rl/configs/simple_online_tune_serialized.gin
+++ /dev/null
@@ -1,100 +0,0 @@
-import tensor2tensor.trax.models
-import tensor2tensor.trax.optimizers
-import tensor2tensor.trax.trax
-import tensor2tensor.trax.rl
-import tensor2tensor.trax.rl.space_serializer
-import tensor2tensor.trax.rl.trainers
-
-# Parameters for BoxSpaceSerializer:
-# ==============================================================================
-BoxSpaceSerializer.precision = 2
-
-# Parameters for MultifactorSchedule:
-# ==============================================================================
-world_model/MultifactorSchedule.constant = 1.0
-world_model/MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
-world_model/MultifactorSchedule.warmup_steps = 10000
-
-# Parameters for Adam:
-# ==============================================================================
-Adam.learning_rate = 1e-3
-Adam.b1 = 0.9
-Adam.b2 = 0.999
-Adam.weight_decay_rate = 0.0
-
-# Parameters for TransformerDecoder:
-# ==============================================================================
-TransformerDecoder.attention_type = @policy/trax.layers.TimeBinCausalAttention
-TransformerDecoder.d_model = 256
-TransformerDecoder.d_ff = 512
-TransformerDecoder.n_layers = 3
-TransformerDecoder.n_heads = 4
-TransformerDecoder.dropout = 0.0
-
-# Parameters for PPO:
-# ==============================================================================
-PPO.n_optimizer_steps = 20
-PPO.optimizer_batch_size = 64
-PPO.target_kl = 0.1
-PPO.boundary = 100
-PPO.max_timestep = 100
-PPO.max_timestep_eval = 100
-PPO.random_seed = None
-PPO.gamma = 1.0
-PPO.lambda_ = 0.95
-PPO.c1 = 1.0
-PPO.c2 = 0.1
-PPO.done_frac_for_policy_save = 0
-PPO.len_history_for_policy = None
-PPO.separate_eval = False
-PPO.save_every_n = 1
-PPO.policy_and_value_model = @trax.models.TransformerDecoder
-PPO.policy_and_value_optimizer = @trax.optimizers.Adam
-PPO.policy_and_value_vocab_size = 128
-PPO.trajectory_dump_min_count_per_shard = 8
-PPO.print_every_optimizer_steps = 1
-
-## Parameters for TimeBinCausalAttention:
-## ==============================================================================
-world_model/TimeBinCausalAttention.dropout = 0.1
-world_model/TimeBinCausalAttention.bin_length = 512
-
-policy/TimeBinCausalAttention.dropout = 0.0
-policy/TimeBinCausalAttention.bin_length = 512
-
-# Parameters for SerializedSequenceSimulatedEnvProblem:
-# ==============================================================================
-SerializedSequenceSimulatedEnvProblem.model = @world_model/trax.models.TransformerLM
-SerializedSequenceSimulatedEnvProblem.reward_fn = @trax.rl.onlinetune_reward_fn
-SerializedSequenceSimulatedEnvProblem.done_fn = @trax.rl.onlinetune_done_fn
-SerializedSequenceSimulatedEnvProblem.vocab_size = 128
-SerializedSequenceSimulatedEnvProblem.max_trajectory_length = 101
-SerializedSequenceSimulatedEnvProblem.significance_decay = 0.8
-
-# Parameters for SimPLe:
-# ==============================================================================
-SimPLe.policy_trainer_class = @trax.rl.trainers.PPO
-SimPLe.n_real_epochs = 1
-SimPLe.n_model_initial_train_steps = 50000
-SimPLe.n_model_train_steps_per_epoch = 10000
-SimPLe.model_train_batch_size = 64
-SimPLe.simulated_env_problem_class = @trax.rl.SerializedSequenceSimulatedEnvProblem
-SimPLe.simulated_batch_size = 128
-SimPLe.n_simulated_epochs = 50
-SimPLe.initial_trajectory_mix_prob = 0.9
-SimPLe.init_policy_from_world_model = True
-
-# Parameters for TransformerLM:
-# ==============================================================================
-world_model/TransformerLM.attention_type = @world_model/trax.layers.TimeBinCausalAttention
-world_model/TransformerLM.d_model = 256
-world_model/TransformerLM.d_ff = 512
-world_model/TransformerLM.n_layers = 3
-world_model/TransformerLM.n_heads = 4
-world_model/TransformerLM.dropout = 0.1
-world_model/TransformerLM.max_len = 2048
-
-# Parameters for train:
-# ==============================================================================
-world_model/train.eval_frequency = 1000
-world_model/train.optimizer = @trax.optimizers.Adafactor
diff --git a/tensor2tensor/trax/rl/envs/__init__.py b/tensor2tensor/trax/rl/envs/__init__.py
deleted file mode 100644
index 7fae2be45..000000000
--- a/tensor2tensor/trax/rl/envs/__init__.py
+++ /dev/null
@@ -1,34 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Environments defined in RL."""
-
-import gin
-from gym.envs.registration import register
-
-from tensor2tensor.trax.rl.envs import online_tune_env
-
-
-# Ginify and register in gym.
-def configure_and_register_env(env_class):
-  register(
-      id="{}-v0".format(env_class.__name__),
-      entry_point="tensor2tensor.trax.rl.envs:{}".format(env_class.__name__),
-  )
-  return gin.external_configurable(env_class, module="trax.rl.envs")
-
-
-# pylint: disable=invalid-name
-OnlineTuneEnv = configure_and_register_env(online_tune_env.OnlineTuneEnv)
diff --git a/tensor2tensor/trax/rl/envs/async_trajectory_collector.py b/tensor2tensor/trax/rl/envs/async_trajectory_collector.py
deleted file mode 100644
index 339f3d34b..000000000
--- a/tensor2tensor/trax/rl/envs/async_trajectory_collector.py
+++ /dev/null
@@ -1,197 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""A trajectory collector that polls on policy files and keeps collecting trajectories."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-import multiprocessing
-import os
-from absl import app
-from absl import flags
-from absl import logging
-import gin
-import jax
-from jax.config import config
-from tensor2tensor.envs import env_problem_utils
-from tensor2tensor.rl.google import atari_utils  # GOOGLE-INTERNAL:
-from tensor2tensor.trax import rl  # pylint: disable=unused-import
-from tensor2tensor.trax.rl import envs as rl_envs  # pylint: disable=unused-import
-from tensor2tensor.trax.rl.envs import async_trajectory_collector_lib as async_lib
-import tensorflow as tf
-
-FLAGS = flags.FLAGS
-
-flags.DEFINE_multi_string("config_file", None,
-                          "Configuration file with parameters (.gin).")
-flags.DEFINE_multi_string("config", None,
-                          "Configuration parameters (gin string).")
-flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
-flags.DEFINE_bool("xm", False, "Copy atari roms?")
-
-flags.DEFINE_bool(
-    "try_abort", True,
-    "Should we try to abort a trajectory collection if a newer "
-    "policy is available.")
-
-flags.DEFINE_string("output_dir", "", "Output dir.")
-flags.DEFINE_string("envs_output_dir", "", "Output dir for the envs.")
-
-flags.DEFINE_boolean(
-    "jax_debug_nans", False,
-    "Setting to true will help to debug nans and disable jit.")
-flags.DEFINE_boolean("disable_jit", False, "Setting to true will disable jit.")
-
-flags.DEFINE_boolean("parallelize_envs", False,
-                     "If true, sets parallelism to number of cpu cores.")
-
-flags.DEFINE_integer("replica", 0, "Basically to append to trajectory name.")
-flags.DEFINE_bool("enable_eager_execution", False, "")
-
-flags.DEFINE_integer(
-    "max_trajectories_to_collect", -1,
-    "-1 for infinite, otherwise whatever number was specified.")
-
-
-# TODO(afrozm): This code snippet is strewn across many places, unify it.
-def initialize_gin():
-  gin_configs = FLAGS.config or []
-  gin.parse_config_files_and_bindings(FLAGS.config_file, gin_configs)
-
-
-def get_output_dir():
-  """Return output_dir."""
-  output_dir = FLAGS.output_dir
-  return output_dir
-
-
-def update_jax_config():
-  """Update JAX config based on flags."""
-
-  if FLAGS.jax_debug_nans:
-    config.update("jax_debug_nans", True)
-
-  if FLAGS.use_tpu:
-    config.update("jax_platform_name", "tpu")
-  else:
-    config.update("jax_platform_name", "gpu")
-
-
-@gin.configurable(blacklist=[
-    "output_dir",
-])
-def create_envs_and_collect_trajectories(
-    output_dir,
-    env_name="OnlineTuneEnv-v0",
-    max_timestep=None,
-    clip_rewards=False,
-    rendered_env=False,
-    resize_dims=(105, 80),
-):
-  """Creates the envs and continuously collects trajectories."""
-
-
-  train_batch_size = 1
-  eval_batch_size = 1
-
-  # TODO(pkozakowski): Find a better way to determine this.
-  train_env_kwargs = {}
-  eval_env_kwargs = {}
-  if "OnlineTuneEnv" in env_name:
-    envs_output_dir = FLAGS.envs_output_dir or os.path.join(output_dir, "envs")
-    train_env_output_dir = os.path.join(envs_output_dir, "train")
-    eval_env_output_dir = os.path.join(envs_output_dir, "eval")
-    train_env_kwargs = {"output_dir": train_env_output_dir}
-    eval_env_kwargs = {"output_dir": eval_env_output_dir}
-
-  if "ClientEnv" in env_name:
-    train_env_kwargs["per_env_kwargs"] = [{
-        "remote_env_address": os.path.join(FLAGS.train_server_bns, str(replica))
-    } for replica in range(train_batch_size)]
-
-    eval_env_kwargs["per_env_kwargs"] = [{
-        "remote_env_address": os.path.join(FLAGS.eval_server_bns, str(replica))
-    } for replica in range(eval_batch_size)]
-
-  parallelism = multiprocessing.cpu_count() if FLAGS.parallelize_envs else 1
-  train_parallelism = min(train_batch_size, parallelism)
-  eval_parallelism = min(eval_batch_size, parallelism)
-
-  train_env = env_problem_utils.make_env(
-      batch_size=train_batch_size,
-      env_problem_name=env_name,
-      resize=rendered_env,
-      resize_dims=resize_dims,
-      max_timestep=max_timestep,
-      clip_rewards=clip_rewards,
-      parallelism=train_parallelism,
-      use_tpu=FLAGS.use_tpu,
-      **train_env_kwargs)
-  assert train_env
-
-  eval_env = env_problem_utils.make_env(
-      batch_size=eval_batch_size,
-      env_problem_name=env_name,
-      resize=rendered_env,
-      resize_dims=resize_dims,
-      max_timestep=max_timestep,
-      clip_rewards=clip_rewards,
-      parallelism=eval_parallelism,
-      use_tpu=FLAGS.use_tpu,
-      **eval_env_kwargs)
-  assert eval_env
-
-  def run_collect_loop():
-    async_lib.continuously_collect_trajectories(
-        output_dir,
-        train_env,
-        eval_env,
-        trajectory_dump_dir=None,
-        env_id=FLAGS.replica,
-        try_abort=FLAGS.try_abort,
-        max_trajectories_to_collect=(None
-                                     if FLAGS.max_trajectories_to_collect < 0
-                                     else FLAGS.max_trajectories_to_collect))
-
-  if FLAGS.jax_debug_nans or FLAGS.disable_jit:
-    with jax.disable_jit():
-      run_collect_loop()
-  else:
-    run_collect_loop()
-
-
-def main(argv):
-  del argv
-
-  if FLAGS.enable_eager_execution:
-    tf.enable_eager_execution()
-
-  logging.info("Initializing Gin.")
-  initialize_gin()
-
-  logging.info("Update JAX config.")
-  update_jax_config()
-
-  logging.info("Getting output_dir")
-  output_dir = get_output_dir()
-  logging.info("Got output_dir = %s", output_dir)
-
-  logging.info("Starting Trajectory collection.")
-  create_envs_and_collect_trajectories(output_dir)
-
-
-if __name__ == "__main__":
-  app.run(main)
diff --git a/tensor2tensor/trax/rl/envs/async_trajectory_collector_lib.py b/tensor2tensor/trax/rl/envs/async_trajectory_collector_lib.py
deleted file mode 100644
index 447a53819..000000000
--- a/tensor2tensor/trax/rl/envs/async_trajectory_collector_lib.py
+++ /dev/null
@@ -1,195 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Uitlity functions for the async trajectory collector."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-import random
-import time
-
-from absl import logging
-from tensor2tensor.envs import trajectory
-from tensor2tensor.trax.rl import ppo
-from tensor2tensor.trax.rl import trainers as rl_trainers
-from tensorflow.io import gfile
-
-LARGE_MAX_TRIES_FOR_POLICY_FILE = 100
-
-
-# TODO(afrozm): Is there a better way to poll for a file on CNS?
-def get_newer_policy_model_file(output_dir,
-                                min_epoch=-1,
-                                sleep_time_secs=0.1,
-                                max_sleep_time_secs=1.0,
-                                max_tries=1,
-                                wait_forever=False,):
-  """Gets a policy model file subject to availability and wait time."""
-
-  while max_tries or wait_forever:
-    max_tries -= 1
-    policy_files = ppo.get_policy_model_files(output_dir)
-
-    def do_wait(t):
-      time.sleep(t)
-      t *= 2
-      return min(t, max_sleep_time_secs)
-
-    # No policy files at all.
-    if not policy_files:
-      logging.info("There are no policy files in [%s], waiting for %s secs.",
-                   output_dir, sleep_time_secs)
-      sleep_time_secs = do_wait(sleep_time_secs)
-      continue
-
-    # Check if we have a newer epoch.
-    policy_file = policy_files[0]
-    epoch = ppo.get_epoch_from_policy_model_file(policy_file)
-
-    # We don't - wait.
-    if epoch <= min_epoch:
-      logging.info("epoch [%s] <= min_epoch [%s], waiting for %s secs.", epoch,
-                   min_epoch, sleep_time_secs)
-      sleep_time_secs = do_wait(sleep_time_secs)
-      continue
-
-    # We do have a new file, return it.
-    policy_file = policy_files[0]
-    epoch = ppo.get_epoch_from_policy_model_file(policy_file)
-    logging.info("Found epoch [%s] and policy file [%s]", epoch, policy_file)
-    return policy_file, epoch
-
-  # Exhausted our waiting limit.
-  return None
-
-
-def dump_trajectory(output_dir, epoch, env_id, temperature, random_string,
-                    trajs):
-  """Write the trajectory to disk."""
-
-  assert 1 == len(trajs)
-  traj = trajs[0]
-
-  trajectory_file_name = trajectory.TRAJECTORY_FILE_FORMAT.format(
-      epoch=epoch, env_id=env_id, temperature=temperature, r=random_string)
-
-  with gfile.GFile(os.path.join(output_dir, trajectory_file_name), "w") as f:
-    trajectory.get_pickle_module().dump(traj, f)
-
-
-def continuously_collect_trajectories(output_dir,
-                                      train_env,
-                                      eval_env,
-                                      trajectory_dump_dir=None,
-                                      env_id=None,
-                                      max_trajectories_to_collect=None,
-                                      try_abort=True):
-  """Instantiates a PPO trainer and collects trajectories."""
-
-  # Make the PPO trainer.
-  ppo_trainer = rl_trainers.PPO(
-      output_dir=output_dir,
-      train_env=train_env,
-      eval_env=eval_env,
-      trajectory_dump_dir=trajectory_dump_dir,
-  )
-
-  # TODO(afrozm): Update base_trainer interface to support SimPLe as well.
-  assert isinstance(ppo_trainer, rl_trainers.PPO)
-
-  assert env_id is not None
-
-  # Get an initial policy and wait a forever to get it if needed.
-  policy_and_epoch = get_newer_policy_model_file(output_dir, wait_forever=True)
-  assert policy_and_epoch
-  policy_file, epoch = policy_and_epoch
-  logging.info("Read initial policy for epoch [%s] -> [%s]", epoch, policy_file)
-
-  # Returns immediately if there is a newer epoch available.
-  def is_newer_policy_file_available(epoch_, sleep_time_secs_=0.1):
-    return get_newer_policy_model_file(
-        output_dir, min_epoch=epoch_, sleep_time_secs=sleep_time_secs_)
-
-  assert 1 == train_env.batch_size
-  assert 1 == eval_env.batch_size
-
-  temperature = 1.0
-
-  trajectories_collected = 0
-
-  train_env_trajectory_dump_dir = os.path.join(output_dir, "trajectories/train")
-  eval_env_trajectory_dump_dir = os.path.join(output_dir, "trajectories/eval")
-
-  gfile.makedirs(train_env_trajectory_dump_dir)
-  gfile.makedirs(eval_env_trajectory_dump_dir)
-
-  while max_trajectories_to_collect is None or trajectories_collected < int(
-      max_trajectories_to_collect):
-    logging.info("Collecting a trajectory, trajectories_collected = %s",
-                 trajectories_collected)
-
-    # Abort function -- if something newever is available, then abort the
-    # current computation and reload.
-
-    # Useful if env.step is long.
-    def long_abort_fn():
-      # We want this to be as quick as possible.
-      return is_newer_policy_file_available(epoch, 0) is not None
-
-    abort_fn = long_abort_fn if try_abort else None
-
-    # Collect a training trajectory.
-    trajs, n_done, unused_timing_info, unused_model_state = (
-        ppo_trainer.collect_trajectories(train=True,
-                                         temperature=temperature,
-                                         abort_fn=abort_fn,
-                                         raw_trajectory=True))
-
-    if trajs and n_done > 0:
-      assert 1 == n_done
-      trajectories_collected += n_done
-
-      # Write the trajectory down.
-      logging.info(
-          "Dumping the collected trajectory, trajectories_collected = %s",
-          trajectories_collected)
-      dump_trajectory(train_env_trajectory_dump_dir, epoch, env_id, temperature,
-                      str(random.randint(0, 2**31 - 1)), trajs)
-    else:
-      logging.info("Computation was aborted, a new policy is available.")
-
-    # This maybe useless, since `abort_fn` will take care of it. We might want
-    # to have this here if abort_fn is False always.
-    # Do we have a newer policy?
-    policy_file_and_epoch = is_newer_policy_file_available(epoch)
-    if policy_file_and_epoch is None:
-      # Continue churning out these policies.
-      logging.info("We don't have a newer policy, continuing with the old one.")
-      continue
-
-    # We have a newer policy, read it and update the parameters.
-    policy_file, epoch = policy_file_and_epoch
-    logging.info(
-        "We have a newer policy epoch [%s], file [%s], updating parameters.",
-        epoch, policy_file)
-    ppo_trainer.update_optimization_state(
-        output_dir, policy_and_value_opt_state=None)
-    logging.info("Parameters of PPOTrainer updated.")
-
-    # Check that the epochs match.
-    assert epoch == ppo_trainer.epoch
diff --git a/tensor2tensor/trax/rl/envs/async_trajectory_collector_lib_test.py b/tensor2tensor/trax/rl/envs/async_trajectory_collector_lib_test.py
deleted file mode 100644
index 39ba95614..000000000
--- a/tensor2tensor/trax/rl/envs/async_trajectory_collector_lib_test.py
+++ /dev/null
@@ -1,64 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-from tensor2tensor.trax.rl import ppo
-from tensor2tensor.trax.rl.envs import async_trajectory_collector_lib as async_lib
-from tensorflow import test
-from tensorflow.io import gfile
-
-
-class AsyncTrajectoryCollectorLibTest(test.TestCase):
-
-  def test_get_newer_policy_model_file(self):
-    output_dir = self.get_temp_dir()
-
-    def write_policy_model_file(epoch):
-      fname = ppo.get_policy_model_file_from_epoch(output_dir, epoch)
-      with gfile.GFile(fname, "w") as f:
-        f.write("some data")
-      return fname
-
-    # No file exists currently.
-    self.assertIsNone(async_lib.get_newer_policy_model_file(output_dir))
-
-    # Write a policy model file.
-    epoch = 0
-    policy_model_filename = write_policy_model_file(epoch)
-
-    # See that we get it.
-    actual_policy_file, actual_epoch = (
-        async_lib.get_newer_policy_model_file(output_dir, min_epoch=-1))
-
-    self.assertEqual(actual_policy_file, policy_model_filename)
-    self.assertEqual(actual_epoch, epoch)
-
-    # If we now ask for a larger epoch, we don't get it.
-    self.assertIsNone(
-        async_lib.get_newer_policy_model_file(output_dir, min_epoch=0))
-
-    # Write a newer epoch and expect to get that with appropriate min_epoch.
-    epoch = 1
-    policy_model_filename = write_policy_model_file(epoch)
-    actual_policy_file, actual_epoch = (
-        async_lib.get_newer_policy_model_file(output_dir, min_epoch=0))
-    self.assertEqual(actual_policy_file, policy_model_filename)
-    self.assertEqual(actual_epoch, epoch)
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/rl/envs/fake_env.py b/tensor2tensor/trax/rl/envs/fake_env.py
deleted file mode 100644
index 2d18c3f82..000000000
--- a/tensor2tensor/trax/rl/envs/fake_env.py
+++ /dev/null
@@ -1,68 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""A fake gym environment.
-
-Can specify either:
-1. A done action, i.e. the action on which the environment returns done.
-2. A done time-step, i.e. the time step at which the environment returns done.
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gym
-import numpy as np
-
-
-class FakeEnv(gym.Env):
-  """A fake env which is either done with a specific action or a time-step."""
-
-  def __init__(self,
-               input_shape=(4,),
-               n_actions=2,
-               n_controls=1,
-               done_time_step=None,
-               done_action=None):
-    self._input_shape = input_shape
-    self._done_time_step = done_time_step
-    self._done_action = done_action
-    self._t = 0
-    if n_controls == 1:
-      self.action_space = gym.spaces.Discrete(n_actions)
-    else:
-      self.action_space = gym.spaces.MultiDiscrete([n_actions] * n_controls)
-    self.observation_space = gym.spaces.Box(
-        low=-1.0, high=1.0, shape=input_shape)
-
-  def _get_random_observation(self):
-    return np.random.random(self._input_shape)
-
-  def reset(self):
-    self._t = 0
-    return self._get_random_observation()
-
-  def step(self, action):
-    assert self.action_space.contains(action)
-    done = False
-    if self._done_action is not None:
-      done = action == self._done_action
-    elif self._done_time_step is not None:
-      done = self._t == self._done_time_step
-
-    reward = -1.0 if not done else 1.0
-    self._t += 1
-    return self._get_random_observation(), reward, done, {}
diff --git a/tensor2tensor/trax/rl/envs/fake_env_test.py b/tensor2tensor/trax/rl/envs/fake_env_test.py
deleted file mode 100644
index cff4d481f..000000000
--- a/tensor2tensor/trax/rl/envs/fake_env_test.py
+++ /dev/null
@@ -1,65 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.rl.fake_env."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.trax.rl.envs import fake_env
-from tensorflow import test
-
-
-class FakeEnvTest(test.TestCase):
-
-  def test_done_action(self):
-    env = fake_env.FakeEnv(input_shape=(2, 3),
-                           n_actions=10,
-                           done_time_step=None,
-                           done_action=9)
-    env.reset()
-
-    # Actions 0 to 8
-    for action in range(9):
-      _, reward, done, _ = env.step(action)
-      self.assertFalse(done)
-      self.assertEqual(-1.0, reward)
-
-    _, reward, done, _ = env.step(9)
-    self.assertTrue(done)
-    self.assertEqual(1.0, reward)
-
-  def test_done_time_step(self):
-    env = fake_env.FakeEnv(input_shape=(2, 3),
-                           n_actions=10,
-                           done_time_step=10,
-                           done_action=None)
-    env.reset()
-
-    # Take 10 steps.
-    for _ in range(10):
-      _, reward, done, _ = env.step(0)
-      self.assertFalse(done)
-      self.assertEqual(-1.0, reward)
-
-    # Take final time-step, this is the time-step numbered 10 since time-steps
-    # are 0 indexed.
-    _, reward, done, _ = env.step(0)
-    self.assertTrue(done)
-    self.assertEqual(1.0, reward)
-
-if __name__ == '__main__':
-  test.main()
diff --git a/tensor2tensor/trax/rl/envs/online_tune.py b/tensor2tensor/trax/rl/envs/online_tune.py
deleted file mode 100644
index a4631edba..000000000
--- a/tensor2tensor/trax/rl/envs/online_tune.py
+++ /dev/null
@@ -1,57 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Utility functions for OnlineTuneEnv."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-
-
-LEARNING_RATE_METRIC = ("train", "training/learning_rate")
-
-
-def historical_metric_values(
-    history, metric, observation_range=(-np.inf, np.inf)):
-  """Converts a metric stream from a trax History object into a numpy array."""
-  metric_sequence = history.get(*metric)
-  metric_values = np.array([
-      metric_value for (_, metric_value) in metric_sequence
-  ])
-  return np.clip(metric_values, *observation_range)
-
-
-def history_to_observations(history, metrics, observation_range, include_lr):
-  """Converts a trax History object into a sequence of observations."""
-  observation_dimensions = [
-      historical_metric_values(history, metric, observation_range)
-      for metric in metrics
-  ]
-  if include_lr:
-    # Logartihm of the learning rate.
-    observation_dimensions.append(np.log(historical_metric_values(
-        history, LEARNING_RATE_METRIC, observation_range
-    )))
-  return np.stack(observation_dimensions, axis=1)
-
-
-def new_learning_rate(action, history, action_multipliers, max_lr):
-  """Calculates a new learning rate based on an action."""
-  learning_rates = historical_metric_values(history, LEARNING_RATE_METRIC)
-  assert learning_rates.shape[0] > 0, "No last learning rate found in history."
-  current_lr = learning_rates[-1]
-  return min(current_lr * action_multipliers[action], max_lr)
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env.py b/tensor2tensor/trax/rl/envs/online_tune_env.py
deleted file mode 100644
index df3513cc8..000000000
--- a/tensor2tensor/trax/rl/envs/online_tune_env.py
+++ /dev/null
@@ -1,233 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""An environment for tuning model hyperparameters during training."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import functools
-import os
-
-import gym
-
-from tensor2tensor.trax import inputs as trax_inputs
-from tensor2tensor.trax import layers
-from tensor2tensor.trax import models as trax_models
-from tensor2tensor.trax import optimizers as trax_opt
-from tensor2tensor.trax import trax
-from tensor2tensor.trax.rl import online_tune
-from tensorflow.io import gfile
-
-
-class OnlineTuneEnv(gym.Env):
-  """An environment for tuning model hyperparameters during training.
-
-  A rollout is one instance of training a specific model on a specific problem.
-  Observations are the values of some evaluation metric. Actions control
-  hyperparameter changes during training. Reward is the change of the evaluation
-  metric. One environment step corresponds to a fixed number of training steps.
-
-  For now we only support tuning the learning rate.
-  """
-
-  # Chosen so that the opposite actions cancel each other out, so random walk
-  # has a median of 1.
-  DEFAULT_ACTION_MULTIPLIERS = [1.0 / 1.5, 1.0 / 1.25, 1.0, 1.25, 1.5]
-
-  def __init__(self,
-               output_dir,
-               model=trax_models.TransformerLM,
-               trainer_class=trax.Trainer,
-               loss_fn=layers.CrossEntropyLossScalar,
-               optimizer=trax_opt.Adafactor,
-               inputs=trax_inputs.inputs,
-               action_multipliers=None,
-               observation_metrics=(
-                   ("train", "metrics/accuracy"),
-                   ("train", "metrics/loss"),
-                   ("eval", "metrics/accuracy"),
-                   ("eval", "metrics/loss"),
-               ),
-               include_controls_in_observation=False,
-               reward_metric=("eval", "metrics/accuracy"),
-               train_steps=100,
-               eval_steps=10,
-               env_steps=100,
-               # This is a tuple instead of a dict because the controls are
-               # ordered in the action space.
-               control_configs=(
-                   # (name, start, (low, high), flip)
-                   ("learning_rate", 1e-3, (1e-9, 10.0), False),
-               ),
-               nontrainable_param_map=None,
-               observation_range=(0.0, 10.0),
-               # Don't save checkpoints by default, as they tend to use a lot of
-               # space.
-               should_save_checkpoints=False,
-               # Same here.
-               should_write_summaries=False,
-               has_weights=False,
-               mask_id=None):
-    if action_multipliers is None:
-      action_multipliers = self.DEFAULT_ACTION_MULTIPLIERS
-    self._model = model
-    # Initialize Trainer in OnlineTuneEnv lazily to prevent long startup in the
-    # async setup, where we just use the environments as containers for
-    # trajectories.
-    self._trainer_fn = functools.partial(
-        trainer_class,
-        model=model,
-        loss_fn=loss_fn,
-        optimizer=optimizer,
-        lr_schedule=(lambda history: lambda step: self._current_controls),
-        inputs=inputs,
-        should_save_checkpoints=should_save_checkpoints,
-        should_write_summaries=should_write_summaries,
-        nontrainable_param_map=nontrainable_param_map,
-        has_weights=has_weights,
-        mask_id=mask_id,
-    )
-    self._trainer = None
-    self._action_multipliers = action_multipliers
-    self._observation_metrics = observation_metrics
-    self._include_controls_in_observation = include_controls_in_observation
-    self._reward_metric = reward_metric
-    self._train_steps = train_steps
-    self._eval_steps = eval_steps
-    self._env_steps = env_steps
-    self._control_configs = control_configs
-    self._observation_range = observation_range
-
-    self._output_dir = output_dir
-    gfile.makedirs(self._output_dir)
-    # Actions are indices in self._action_multipliers.
-    self.action_space = gym.spaces.MultiDiscrete(
-        [len(self._action_multipliers)] * len(self._control_configs)
-    )
-    # Observation is a vector with the values of the metrics specified in
-    # observation_metrics plus optionally the current controls.
-    observation_dim = (
-        len(self._observation_metrics) +
-        int(self._include_controls_in_observation) * len(self._control_configs)
-    )
-
-    (obs_low, obs_high) = observation_range
-    self.observation_space = gym.spaces.Box(
-        # Observations are clipped to this range.
-        low=obs_low, high=obs_high, shape=(observation_dim,),
-    )
-
-  @property
-  def _next_trajectory_dir(self):
-    """Assigns a new output dir for a trajectory under self._output_dir.
-
-    Directory names are consecutive integers starting from zero. New directory
-    index is assigned as the maximum of past indices plus one. Directories that
-    are not integers are ignored.
-
-    Returns:
-      A path of the new directory.
-    """
-    trajectory_dirs = gfile.listdir(self._output_dir)
-
-    def int_or_none(s):
-      try:
-        return int(s)
-      except TypeError:
-        return None
-
-    past_trajectory_ids = [
-        trajectory_id for trajectory_id in map(int_or_none, trajectory_dirs)
-        if trajectory_id is not None]
-    next_trajectory_id = max([-1] + past_trajectory_ids) + 1
-
-    return os.path.join(self._output_dir, str(next_trajectory_id))
-
-  @property
-  def _current_reward_metric(self):
-    metric_values = online_tune.historical_metric_values(
-        self._trainer.state.history,
-        self._reward_metric,
-    )
-    assert metric_values.shape[0] > 0, (
-        "No values in history for metric {}.".format(self._reward_metric))
-    return metric_values[-1]
-
-  @property
-  def _current_observation(self):
-    observations = online_tune.history_to_observations(
-        self._trainer.state.history,
-        self._observation_metrics,
-        self._observation_range,
-        self._control_configs if self._include_controls_in_observation
-        else None,
-    )
-    assert observations.shape[0] > 0, "No values in history for any metric."
-    return observations[-1, :]
-
-  @property
-  def trainer(self):
-    if self._trainer is None:
-      raise ValueError("The environment has to be reset first.")
-    return self._trainer
-
-  def reset(self):
-    if self._trainer is None:
-      self._trainer = self._trainer_fn()
-    self._current_controls = {
-        name: start_value
-        for (name, start_value, _, _) in self._control_configs
-    }
-    self._step = 0
-    self._trainer.reset(output_dir=self._next_trajectory_dir)
-    self._trainer.evaluate(self._eval_steps)
-    return self._current_observation
-
-  def step(self, action):
-    """Step the environment.
-
-    One environment step corresponds to self.train_steps training steps.
-
-    Args:
-      action: (int) Action to take. An index in self.action_multipliers.
-
-    Returns:
-      Tuple (observation, reward, done, info). observation is a singleton vector
-        with the current value of the metric. reward is the difference in the
-        metric since the last step. done is set after reaching self.env_steps
-        environment steps. info is an empty dict.
-    """
-    self._current_controls = {
-        # name: value
-        control_config[0]: online_tune.update_control(  # pylint: disable=g-complex-comprehension
-            control_config,
-            control_action,
-            self._trainer.state.history,
-            self._action_multipliers,
-        )
-        for (control_action, control_config) in zip(
-            action, self._control_configs
-        )
-    }
-    last_reward_metric = self._current_reward_metric
-    self._trainer.train_epoch(self._train_steps, self._eval_steps)
-    self._step += 1
-    current_reward_metric = self._current_reward_metric
-    observation = self._current_observation
-    reward = current_reward_metric - last_reward_metric
-    done = self._step == self._env_steps
-    return (observation, reward, done, {})
diff --git a/tensor2tensor/trax/rl/envs/online_tune_env_test.py b/tensor2tensor/trax/rl/envs/online_tune_env_test.py
deleted file mode 100644
index aafddfe1a..000000000
--- a/tensor2tensor/trax/rl/envs/online_tune_env_test.py
+++ /dev/null
@@ -1,152 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.rl.online_tune_env."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import functools
-
-import numpy as np
-
-from tensor2tensor.trax import inputs as trax_inputs
-from tensor2tensor.trax import models
-from tensor2tensor.trax import optimizers
-from tensor2tensor.trax import trax
-from tensor2tensor.trax.rl import online_tune
-from tensor2tensor.trax.rl.envs import online_tune_env
-from tensorflow import test
-from tensorflow.io import gfile
-
-HISTORY_MODE = "eval"
-METRIC = "metrics/accuracy"
-
-
-class MockTrainer(trax.Trainer):
-
-  def __init__(self, metrics_to_report, *args, **kwargs):
-    super(MockTrainer, self).__init__(*args, **kwargs)
-    self.controls = []
-    self.init_metrics_to_report = metrics_to_report
-    self.metrics_to_report = None
-
-  def reset(self, output_dir):
-    super(MockTrainer, self).reset(output_dir)
-    # Copy the sequence to a list so we can modify it later.
-    self.metrics_to_report = list(self.init_metrics_to_report)
-
-  def train_epoch(self, epoch_steps, eval_steps):
-    del epoch_steps
-    self.controls.append(self.nontrainable_params)
-    self.evaluate(eval_steps)
-
-  def evaluate(self, eval_steps):
-    del eval_steps
-    self.state.history.append(
-        mode=HISTORY_MODE,
-        metric=METRIC,
-        step=self.step,
-        value=self.metrics_to_report.pop(0))
-    for (name, value) in self.nontrainable_params.items():
-      (mode, metric) = online_tune.control_metric(name)
-      self.state.history.append(
-          mode=mode,
-          metric=metric,
-          step=self.step,
-          value=value)
-
-
-class OnlineTuneTest(test.TestCase):
-
-  @staticmethod
-  def _create_env(
-      output_dir, metrics_to_report=(0.0,), action_multipliers=(1,)
-  ):
-    return online_tune_env.OnlineTuneEnv(
-        trainer_class=functools.partial(MockTrainer, metrics_to_report),
-        model=functools.partial(
-            models.MLP, n_hidden_layers=0, n_output_classes=1),
-        inputs=functools.partial(
-            trax_inputs.random_inputs,
-            input_shape=(1, 1),
-            input_dtype=np.float32,
-            output_shape=(1, 1),
-            output_dtype=np.float32),
-        optimizer=optimizers.Momentum,
-        control_configs=(
-            ("learning_rate", 1e-3, (1e-9, 10.0), False),
-            ("weight_decay_rate", 1e-5, (1e-9, 0.1), False),
-        ),
-        include_controls_in_observation=False,
-        output_dir=output_dir,
-        action_multipliers=action_multipliers,
-        observation_metrics=[(HISTORY_MODE, METRIC)],
-        reward_metric=(HISTORY_MODE, METRIC),
-        train_steps=1,
-        eval_steps=1,
-        env_steps=(len(metrics_to_report) - 1))
-
-  def test_communicates_with_trainer(self):
-    action_multipliers = [0.8, 1.0, 1.25]
-    metrics_to_report = [0.1, 0.5, 0.8, 0.9]
-    actions_to_take = [[0, 1], [1, 2], [2, 0]]
-    expected_observations = np.expand_dims(metrics_to_report, axis=1)
-    # Metric difference in consecutive timesteps.
-    expected_rewards = [0.4, 0.3, 0.1]
-    expected_dones = [False, False, True]
-    expected_controls = [
-        {"learning_rate": 0.0008, "weight_decay_rate": 1e-5},
-        {"learning_rate": 0.0008, "weight_decay_rate": 1.25e-5},
-        {"learning_rate": 0.001, "weight_decay_rate": 1e-5},
-    ]
-
-    env = self._create_env(
-        output_dir=self.get_temp_dir(),
-        metrics_to_report=metrics_to_report,
-        action_multipliers=action_multipliers)
-    actual_observations = [env.reset()]
-    actual_rewards = []
-    actual_dones = []
-    for action in actions_to_take:
-      (observation, reward, done, _) = env.step(action)
-      actual_observations.append(observation)
-      actual_rewards.append(reward)
-      actual_dones.append(done)
-
-    np.testing.assert_allclose(actual_observations, expected_observations)
-    np.testing.assert_allclose(actual_rewards, expected_rewards)
-    self.assertEqual(actual_dones, expected_dones)
-    def get_control(name, controls):
-      return [control[name] for control in controls]
-    for name in ("learning_rate", "weight_decay_rate"):
-      np.testing.assert_allclose(
-          get_control(name, env.trainer.controls),
-          get_control(name, expected_controls),
-      )
-
-  def test_creates_new_trajectory_dirs(self):
-    output_dir = self.get_temp_dir()
-    env = self._create_env(output_dir=output_dir)
-    self.assertEqual(set(gfile.listdir(output_dir)), set())
-    env.reset()
-    self.assertEqual(set(gfile.listdir(output_dir)), {"0"})
-    env.reset()
-    self.assertEqual(set(gfile.listdir(output_dir)), {"0", "1"})
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/rl/envs/online_tune_test.py b/tensor2tensor/trax/rl/envs/online_tune_test.py
deleted file mode 100644
index 69b09a836..000000000
--- a/tensor2tensor/trax/rl/envs/online_tune_test.py
+++ /dev/null
@@ -1,111 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.rl.online_tune."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-
-from tensor2tensor.trax import history as trax_history
-from tensor2tensor.trax.rl.envs import online_tune
-from tensorflow import test
-
-
-class OnlineTuneTest(test.TestCase):
-
-  def _append_metrics(self, h, metric, values):
-    for (i, value) in enumerate(values):
-      h.append(*metric, step=i, value=value)
-
-  def test_retrieves_historical_metric_values(self):
-    history = trax_history.History()
-    self._append_metrics(history, ("train", "accuracy"), [0.1, 0.73])
-    metric_values = online_tune.historical_metric_values(
-        history, metric=("train", "accuracy"), observation_range=(0, 5))
-    np.testing.assert_array_equal(metric_values, [0.1, 0.73])
-
-  def test_clips_historical_metric_values(self):
-    history = trax_history.History()
-    self._append_metrics(history, ("train", "loss"), [-10, 10])
-    metric_values = online_tune.historical_metric_values(
-        history, metric=("train", "loss"), observation_range=(-1, 1))
-    np.testing.assert_array_equal(metric_values, [-1, 1])
-
-  def test_converts_history_to_observations_without_learning_rate(self):
-    history = trax_history.History()
-    self._append_metrics(history, ("train", "loss"), [3.0, 1.07])
-    self._append_metrics(history, ("eval", "accuracy"), [0.12, 0.68])
-    observations = online_tune.history_to_observations(
-        history,
-        metrics=(("eval", "accuracy"), ("train", "loss")),
-        observation_range=(0, 5),
-        include_lr=False,
-    )
-    np.testing.assert_array_equal(observations, [[0.12, 3.0], [0.68, 1.07]])
-
-  def test_converts_history_to_observations_with_learning_rate(self):
-    history = trax_history.History()
-    self._append_metrics(
-        history, ("train", "training/learning_rate"), [1e-3, 1e-4])
-    observations = online_tune.history_to_observations(
-        history,
-        metrics=(),
-        observation_range=(0, 5),
-        include_lr=True,
-    )
-    self.assertEqual(observations.shape, (2, 1))
-    ((log_lr_1,), (log_lr_2,)) = observations
-    self.assertGreater(log_lr_1, log_lr_2)
-
-  def test_clips_observations(self):
-    history = trax_history.History()
-    self._append_metrics(history, ("eval", "loss"), [-10, 10])
-    observations = online_tune.history_to_observations(
-        history,
-        metrics=(("eval", "loss"),),
-        observation_range=(-2, 2),
-        include_lr=False,
-    )
-    np.testing.assert_array_equal(observations, [[-2], [2]])
-
-  def test_calculates_new_learning_rate(self):
-    history = trax_history.History()
-    self._append_metrics(
-        history, online_tune.LEARNING_RATE_METRIC, [1e-2, 1e-3])
-    new_lr = online_tune.new_learning_rate(
-        action=2,
-        history=history,
-        action_multipliers=(0.5, 1.0, 2.0),
-        max_lr=1.0,
-    )
-    np.testing.assert_almost_equal(new_lr, 2e-3)
-
-  def test_clips_new_learning_rate(self):
-    history = trax_history.History()
-    self._append_metrics(history, online_tune.LEARNING_RATE_METRIC, [1e-3])
-    new_lr = online_tune.new_learning_rate(
-        action=0,
-        history=history,
-        action_multipliers=(4.0, 1.0, 0.25),
-        max_lr=3e-3,
-    )
-    np.testing.assert_almost_equal(new_lr, 3e-3)
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/rl/online_tune.py b/tensor2tensor/trax/rl/online_tune.py
deleted file mode 100644
index 905fe59b6..000000000
--- a/tensor2tensor/trax/rl/online_tune.py
+++ /dev/null
@@ -1,115 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Utility functions for OnlineTuneEnv."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-
-
-def historical_metric_values(history, metric):
-  """Converts a metric stream from a trax History object into a numpy array."""
-  metric_sequence = history.get(*metric)
-  metric_values = np.array([
-      metric_value for (_, metric_value) in metric_sequence
-  ])
-  if np.any(np.isnan(metric_values)):
-    # Zero out all observations if any element is NaN. This way the agent
-    # doesn't get any rewards, so it learns to avoid those regions.
-    metric_values[:] = 0.0
-  return metric_values
-
-
-def control_to_observation(control_values, control_config, observation_range):
-  """Flips, logarithms, clips and scales the control to observation_range."""
-  (_, _, (low, high), flip) = control_config
-  def transform(x):
-    return np.log(maybe_flip(x, flip))
-  (log_control_values, log_low, log_high) = map(
-      transform, (control_values, low, high)
-  )
-  if flip:
-    (log_low, log_high) = (log_high, log_low)
-  log_control_values = np.clip(log_control_values, log_low, log_high)
-  # Rescale the log control values to the observation range.
-  (obs_low, obs_high) = observation_range
-  return (
-      (log_control_values - log_low) / (log_high - log_low) *
-      (obs_high - obs_low) + obs_low
-  )
-
-
-def control_metric(name):
-  """Returns the (mode, metric) pair in History for the given control."""
-  return ("train", "training/{}".format(name))
-
-
-def maybe_flip(value, flip):
-  """Flips a control (or not).
-
-  Meant to translate controls that naturally take values close to 1
-  (e.g. momentum) to a space where multiplication makes sense (i.e. close to 0).
-
-  Args:
-    value: float or numpy array, value of the control.
-    flip: bool, whether to flip or not.
-
-  Returns:
-    Either value or 1 - value based on flip.
-  """
-  if flip:
-    value = 1 - value
-  return value
-
-
-def history_to_observations(
-    history, metrics, observation_range, control_configs=None):
-  """Converts a trax History object into a sequence of observations."""
-  (obs_low, obs_high) = observation_range
-  observation_dimensions = [
-      np.clip(historical_metric_values(history, metric), obs_low, obs_high)
-      for metric in metrics
-  ]
-  if control_configs is not None:
-    for control_config in control_configs:
-      (control_name, _, _, _) = control_config
-      observation_dimensions.append(control_to_observation(
-          historical_metric_values(history, control_metric(control_name)),
-          control_config,
-          observation_range,
-      ))
-  return np.stack(observation_dimensions, axis=1)
-
-
-def update_control(control_config, action, history, action_multipliers):
-  """Calculates a new value of a control based on an action."""
-  (name, _, (low, high), flip) = control_config
-  metric = control_metric(name)
-  control_values = historical_metric_values(history, metric)
-  assert control_values.shape[0] > 0, (
-      "No last control {} found in history.".format(name))
-  current_control = control_values[-1]
-  (current_control, low, high) = maybe_flip(
-      np.array([current_control, low, high]), flip
-  )
-  if flip:
-    (low, high) = (high, low)
-  new_control = np.clip(
-      current_control * action_multipliers[action], low, high
-  )
-  return maybe_flip(new_control, flip)
diff --git a/tensor2tensor/trax/rl/online_tune_test.py b/tensor2tensor/trax/rl/online_tune_test.py
deleted file mode 100644
index 54d5691d9..000000000
--- a/tensor2tensor/trax/rl/online_tune_test.py
+++ /dev/null
@@ -1,176 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.online_tune."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-
-from tensor2tensor.trax import history as trax_history
-from tensor2tensor.trax.rl import online_tune
-from tensorflow import test
-
-
-class OnlineTuneTest(test.TestCase):
-
-  def _append_metrics(self, h, metric, values):
-    for (i, value) in enumerate(values):
-      h.append(*metric, step=i, value=value)
-
-  def test_retrieves_historical_metric_values(self):
-    history = trax_history.History()
-    self._append_metrics(history, ("train", "accuracy"), [0.1, 0.73])
-    metric_values = online_tune.historical_metric_values(
-        history, metric=("train", "accuracy")
-    )
-    np.testing.assert_array_equal(metric_values, [0.1, 0.73])
-
-  def test_converts_control_to_log_scale_without_flipping(self):
-    config = ("weight_decay", None, (1e-5, 0.1), False)
-    controls = np.array([0.01, 0.02, 0.04])
-    obs_range = (-1, 1)
-    obs = online_tune.control_to_observation(controls, config, obs_range)
-    np.testing.assert_almost_equal(obs[1] - obs[0], obs[2] - obs[1])
-
-  def test_converts_control_to_log_scale_with_flipping(self):
-    config = ("momentum", None, (0.5, 0.99), True)
-    controls = np.array([0.98, 0.96, 0.92])
-    obs_range = (-1, 1)
-    obs = online_tune.control_to_observation(controls, config, obs_range)
-    np.testing.assert_almost_equal(obs[1] - obs[0], obs[2] - obs[1])
-
-  def test_clips_control_without_flipping(self):
-    config = ("weight_decay", None, (1e-5, 0.1), False)
-    controls = np.array([0.0, 0.2])
-    obs_range = (-1, 1)
-    obs = online_tune.control_to_observation(controls, config, obs_range)
-    np.testing.assert_equal(obs, [-1, 1])
-
-  def test_clips_control_with_flipping(self):
-    config = ("momentum", None, (0.5, 0.99), True)
-    controls = np.array([0.4, 1.0])
-    obs_range = (-1, 1)
-    obs = online_tune.control_to_observation(controls, config, obs_range)
-    np.testing.assert_equal(obs, [1, -1])
-
-  def test_rescales_control(self):
-    config = ("weight_decay", None, (1e-5, 0.1), False)
-    controls = np.array([4e-4, 3e-3, 2e-2])
-    (obs_low, obs_high) = (103, 104)
-    obs = online_tune.control_to_observation(
-        controls, config, observation_range=(obs_low, obs_high),
-    )
-    np.testing.assert_array_less(obs, [obs_high] * 3)
-    np.testing.assert_array_less([obs_low] * 3, obs)
-
-  def test_converts_history_to_observations_without_controls(self):
-    history = trax_history.History()
-    self._append_metrics(history, ("train", "loss"), [1.0, 0.07])
-    self._append_metrics(history, ("eval", "accuracy"), [0.12, 0.68])
-    observations = online_tune.history_to_observations(
-        history,
-        metrics=(("eval", "accuracy"), ("train", "loss")),
-        observation_range=(-1, 1),
-        control_configs=None,
-    )
-    np.testing.assert_array_almost_equal(
-        observations, [[0.12, 1.0], [0.68, 0.07]]
-    )
-
-  def test_converts_history_to_observations_with_controls(self):
-    history = trax_history.History()
-    self._append_metrics(
-        history, ("train", "training/learning_rate"), [1e-3, 1e-4])
-    observations = online_tune.history_to_observations(
-        history,
-        metrics=(),
-        observation_range=(0, 5),
-        control_configs=(
-            ("learning_rate", None, (1e-9, 10.0), False),
-        ),
-    )
-    self.assertEqual(observations.shape, (2, 1))
-    ((log_lr_1,), (log_lr_2,)) = observations
-    self.assertGreater(log_lr_1, log_lr_2)
-
-  def test_clips_observations(self):
-    history = trax_history.History()
-    self._append_metrics(history, ("eval", "loss"), [-10, 10])
-    observations = online_tune.history_to_observations(
-        history,
-        metrics=(("eval", "loss"),),
-        observation_range=(-2, 2),
-        control_configs=None,
-    )
-    np.testing.assert_array_equal(observations, [[-2], [2]])
-
-  def test_updates_control_without_flipping(self):
-    config = ("learning_rate", None, (1e-9, 10.0), False)
-    history = trax_history.History()
-    self._append_metrics(
-        history, online_tune.control_metric("learning_rate"), [1e-2, 1e-3])
-    new_control = online_tune.update_control(
-        control_config=config,
-        action=2,
-        history=history,
-        action_multipliers=(0.5, 1.0, 2.0),
-    )
-    np.testing.assert_almost_equal(new_control, 2e-3)
-
-  def test_updates_control_with_flipping(self):
-    config = ("momentum", None, (0.5, 0.99), True)
-    history = trax_history.History()
-    self._append_metrics(
-        history, online_tune.control_metric("momentum"), [0.96, 0.98])
-    new_control = online_tune.update_control(
-        control_config=config,
-        action=0,
-        history=history,
-        action_multipliers=(0.5, 1.0, 2.0),
-    )
-    np.testing.assert_almost_equal(new_control, 0.99)
-
-  def test_clips_updated_control_without_flipping(self):
-    config = ("learning_rate", None, (1e-9, 10.0), False)
-    history = trax_history.History()
-    self._append_metrics(
-        history, online_tune.control_metric("learning_rate"), [7.0])
-    new_control = online_tune.update_control(
-        control_config=config,
-        action=2,
-        history=history,
-        action_multipliers=(0.5, 1.0, 2.0),
-    )
-    np.testing.assert_almost_equal(new_control, 10.0)
-
-  def test_clips_updated_control_with_flipping(self):
-    config = ("momentum", None, (0.5, 0.99), True)
-    history = trax_history.History()
-    self._append_metrics(
-        history, online_tune.control_metric("momentum"), [0.985])
-    new_control = online_tune.update_control(
-        control_config=config,
-        action=0,
-        history=history,
-        action_multipliers=(0.5, 1.0, 2.0),
-    )
-    np.testing.assert_almost_equal(new_control, 0.99)
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/rl/ppo.py b/tensor2tensor/trax/rl/ppo.py
deleted file mode 100644
index 15338d1fa..000000000
--- a/tensor2tensor/trax/rl/ppo.py
+++ /dev/null
@@ -1,971 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""PPO in JAX.
-
-Notation:
-
-B, scalar   - batch size
-RT, scalar  - (reward time) number of time-steps in a trajectory, or the size
-              of the padded reward sequence.
-AT, scalar  - (action time) number of controls in a trajectory, or the size
-              of the policy network output.
-OBS, tuple  - shape of a singular observation from the environment.
-             Ex: For CartPole-v0 this is (4,) and Pong-v0 it's (210, 160, 3)
-A, scalar   - Number of actions, assuming a discrete space.
-
-Policy and Value function signatures:
-
-Policy            Function :: [B, RT + 1] + OBS ->  [B, AT, A]
-Value             Function :: [B, RT + 1] + OBS ->  [B, AT]
-Policy and Value  Function :: [B, RT + 1] + OBS -> ([B, AT, A], [B, AT])
-
-i.e. the policy net should take a batch of *trajectories* and at each time-step
-in each batch deliver a probability distribution over actions.
-
-NOTE: It doesn't return logits, rather the expectation is that it returns
-log-probabilities instead.
-
-NOTE: The policy and value functions need to take care to not take into account
-future time-steps while deciding the actions (or value) for the current
-time-step.
-
-Policy and Value Function produces a tuple of the expected output of a policy
-function and a value function.
-
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import collections
-import functools
-import itertools
-import os
-import re
-import time
-
-from absl import logging
-from jax import grad
-from jax import jit
-from jax import lax
-from jax import numpy as np
-import numpy as onp
-
-from tensor2tensor.envs import env_problem
-from tensor2tensor.envs import env_problem_utils
-from tensor2tensor.trax import layers as tl
-from tensor2tensor.trax import utils
-from tensorflow.io import gfile
-
-
-def policy_and_value_net(
-    n_actions, n_controls, vocab_size, bottom_layers_fn, two_towers
-):
-  """A policy and value net function."""
-
-  # Layers.
-
-  # Now, with the current logits, one head computes action probabilities and the
-  # other computes the value function.
-  # NOTE: The LogSoftmax instead of the Softmax because of numerical stability.
-
-  @tl.layer()
-  def FlattenControlsIntoTime(x, **unused_kwargs):  # pylint: disable=invalid-name
-    """Splits logits for actions in different controls and flattens controls."""
-    return np.reshape(x, (x.shape[0], -1, n_actions))
-
-  if vocab_size is None:
-    # In continuous policies every element of the output sequence corresponds to
-    # an observation.
-    n_preds_per_input = n_controls
-    kwargs = {}
-  else:
-    # In discrete policies every element of the output sequence corresponds to
-    # a symbol in the discrete representation, and each control takes 1 symbol.
-    n_preds_per_input = 1
-    kwargs = {"vocab_size": vocab_size}
-
-  if two_towers:
-    layers = [
-        tl.Dup(),
-        tl.Parallel(
-            [bottom_layers_fn(**kwargs),
-             tl.Dense(n_preds_per_input * n_actions),
-             FlattenControlsIntoTime(),  # pylint: disable=no-value-for-parameter
-             tl.LogSoftmax()],
-            [bottom_layers_fn(**kwargs),
-             tl.Dense(n_preds_per_input),
-             tl.Flatten()],
-        )
-    ]
-  else:
-    layers = [
-        bottom_layers_fn(**kwargs),
-        tl.Dup(),
-        tl.Parallel(
-            [tl.Dense(n_preds_per_input * n_actions),
-             FlattenControlsIntoTime(),  # pylint: disable=no-value-for-parameter
-             tl.LogSoftmax()],
-            [tl.Dense(n_preds_per_input), tl.Flatten()],
-        )
-    ]
-  return tl.Model(layers)
-
-
-def optimizer_fn(optimizer, net_params):
-  """Exposes a convenient interface for the optimizer.
-
-  Args:
-    optimizer: Optimizer class to use.
-    net_params: A nested structure of network parameters.
-
-  Returns:
-    A tuple (opt_state, opt_update, get_params), where:
-      opt_state: Pair (net_params, opt_slots) - initial optimization state.
-      opt_update: Function (step, grads, opt_state) -> opt_state doing one
-        optimization step.
-      get_params: Function opt_state -> net_params for extracting the network
-        parameters from the optimization state.
-  """
-  opt = optimizer()
-  (init_slots, init_nontrainable_slots) = opt.tree_init(net_params)
-  init_state = (net_params, init_slots)
-
-  def opt_update(step, grads, opt_state):
-    (params, slots) = opt_state
-    # Pass the initial nontrainable_slots as we don't tune them during training.
-    # (yet!)
-    return opt.tree_update(step, grads, params, slots, init_nontrainable_slots)
-
-  def get_params(opt_state):
-    (params, _) = opt_state
-    return params
-
-  return init_state, opt_update, get_params
-
-
-# Should this be collect 'n' trajectories, or
-# Run the env for 'n' steps and take completed trajectories, or
-# Any other option?
-def collect_trajectories(env,
-                         policy_fn,
-                         n_trajectories=1,
-                         max_timestep=None,
-                         reset=True,
-                         len_history_for_policy=32,
-                         boundary=32,
-                         state=None,
-                         temperature=1.0,
-                         rng=None,
-                         abort_fn=None,
-                         raw_trajectory=False,):
-  """Collect trajectories with the given policy net and behaviour.
-
-  Args:
-    env: A gym env interface, for now this is not-batched.
-    policy_fn: observations(B,RT+1) -> log-probabs(B, AT, A) callable.
-    n_trajectories: int, number of trajectories.
-    max_timestep: int or None, the index of the maximum time-step at which we
-      return the trajectory, None for ending a trajectory only when env returns
-      done.
-    reset: bool, true if we want to reset the envs. The envs are also reset if
-      max_max_timestep is None or < 0
-    len_history_for_policy: int or None, the maximum history to keep for
-      applying the policy on. If None, use the full history.
-    boundary: int, pad the sequences to the multiples of this number.
-    state: state for `policy_fn`.
-    temperature: (float) temperature to sample action from policy_fn.
-    rng: jax rng, splittable.
-    abort_fn: callable, If not None, then at every env step call and abort the
-      trajectory collection if it returns True, if so reset the env and return
-      None.
-    raw_trajectory: bool, if True a list of trajectory.Trajectory objects is
-      returned, otherwise a list of numpy representations of
-      `trajectory.Trajectory` is returned.
-
-  Returns:
-    A tuple (trajectory, number of trajectories that are done)
-    trajectory: list of (observation, action, reward) tuples, where each element
-    `i` is a tuple of numpy arrays with shapes as follows:
-    observation[i] = (B, T_i + 1)
-    action[i] = (B, T_i)
-    reward[i] = (B, T_i)
-  """
-
-  assert isinstance(env, env_problem.EnvProblem)
-  # This is an env_problem, run its collect function.
-  trajs, n_done, timing_info, state = env_problem_utils.play_env_problem_with_policy(
-      env,
-      policy_fn,
-      num_trajectories=n_trajectories,
-      max_timestep=max_timestep,
-      reset=reset,
-      len_history_for_policy=len_history_for_policy,
-      boundary=boundary,
-      state=state,
-      temperature=temperature,
-      rng=rng,
-      abort_fn=abort_fn,
-      raw_trajectory=raw_trajectory,
-  )
-  # Skip returning raw_rewards here, since they aren't used.
-
-  # t is the return value of Trajectory.as_numpy, so:
-  # (observation, action, processed_reward, raw_reward, infos)
-  return trajs, n_done, timing_info, state
-
-
-# This function can probably be simplified, ask how?
-# Can we do something much simpler than lax.pad, maybe np.pad?
-# Others?
-
-
-def get_padding_value(dtype):
-  """Returns the padding value given a dtype."""
-  padding_value = None
-  if dtype == np.uint8:
-    padding_value = np.uint8(0)
-  elif dtype == np.uint16:
-    padding_value = np.uint16(0)
-  elif dtype == np.float32 or dtype == np.float64:
-    padding_value = 0.0
-  else:
-    padding_value = 0
-  assert padding_value is not None
-  return padding_value
-
-
-# TODO(afrozm): Use np.pad instead and make jittable?
-def pad_trajectories(trajectories, boundary=20):
-  """Pad trajectories to a bucket length that is a multiple of boundary.
-
-  Args:
-    trajectories: list[(observation, actions, rewards)], where each observation
-      is shaped (t+1,) + OBS and actions & rewards are shaped (t,), with the
-      length of the list being B (batch size).
-    boundary: int, bucket length, the actions and rewards are padded to integer
-      multiples of boundary.
-
-  Returns:
-    tuple: (padding lengths, reward_mask, padded_observations, padded_actions,
-        padded_rewards) where padded_observations is shaped (B, RT+1) + OBS and
-        padded_actions, padded_rewards & reward_mask are shaped (B, RT).
-        Where RT is max(t) rounded up to an integer multiple of boundary.
-        padded_length is how much padding we've added and
-        reward_mask is 1s for actual rewards and 0s for the padding.
-  """
-
-  # Let's compute max(t) over all trajectories.
-  t_max = max(r.shape[0] for (_, _, r, _) in trajectories)
-
-  # t_max is rounded to the next multiple of `boundary`
-  boundary = int(boundary)
-  bucket_length = boundary * int(np.ceil(float(t_max) / boundary))
-
-  # So all obs will be padded to t_max + 1 and actions and rewards to t_max.
-  padded_observations = []
-  padded_actions = []
-  padded_rewards = []
-  padded_infos = collections.defaultdict(list)
-  padded_lengths = []
-  reward_masks = []
-
-  for (o, a, r, i) in trajectories:
-    # Determine the amount to pad, this holds true for obs, actions and rewards.
-    num_to_pad = bucket_length + 1 - o.shape[0]
-    padded_lengths.append(num_to_pad)
-    if num_to_pad == 0:
-      padded_observations.append(o)
-      padded_actions.append(a)
-      padded_rewards.append(r)
-      reward_masks.append(onp.ones_like(r, dtype=np.int32))
-      if i:
-        for k, v in i.items():
-          padded_infos[k].append(v)
-      continue
-
-    # First pad observations.
-    padding_config = tuple([(0, num_to_pad, 0)] + [(0, 0, 0)] * (o.ndim - 1))
-
-    padding_value = get_padding_value(o.dtype)
-    action_padding_value = get_padding_value(a.dtype)
-    reward_padding_value = get_padding_value(r.dtype)
-
-    padded_obs = lax.pad(o, padding_value, padding_config)
-    padded_observations.append(padded_obs)
-
-    # Now pad actions and rewards.
-    padding_config = tuple([(0, num_to_pad, 0)] + [(0, 0, 0)] * (a.ndim - 1))
-    padded_action = lax.pad(a, action_padding_value, padding_config)
-    padded_actions.append(padded_action)
-
-    assert r.ndim == 1
-    padding_config = ((0, num_to_pad, 0),)
-    padded_reward = lax.pad(r, reward_padding_value, padding_config)
-    padded_rewards.append(padded_reward)
-
-    # Also create the mask to use later.
-    reward_mask = onp.ones_like(r, dtype=np.int32)
-    reward_masks.append(lax.pad(reward_mask, 0, padding_config))
-
-    if i:
-      for k, v in i.items():
-        # Create a padding configuration for this value.
-        padding_config = [(0, num_to_pad, 0)] + [(0, 0, 0)] * (v.ndim - 1)
-        padded_infos[k].append(lax.pad(v, 0.0, tuple(padding_config)))
-
-  # Now stack these padded_infos if they exist.
-  stacked_padded_infos = None
-  if padded_infos:
-    stacked_padded_infos = {k: np.stack(v) for k, v in padded_infos.items()}
-
-  return padded_lengths, np.stack(reward_masks), np.stack(
-      padded_observations), np.stack(padded_actions), np.stack(
-          padded_rewards), stacked_padded_infos
-
-
-def rewards_to_go(rewards, mask, gamma=0.99):
-  r"""Computes rewards to go.
-
-  Reward to go is defined as follows, the discounted reward that we have to
-  yet collect, going forward from this point, i.e.:
-
-  r2g_t = \sum_{l=0}^{\infty} (\gamma^{l} * reward_{t+l})
-
-  Args:
-    rewards: np.ndarray of shape (B, RT) of rewards.
-    mask: np.ndarray of shape (B, RT) of mask for the rewards.
-    gamma: float, discount factor.
-
-  Returns:
-    rewards to go, np.ndarray of shape (B, RT).
-  """
-  B, RT = rewards.shape  # pylint: disable=invalid-name,unused-variable
-
-  masked_rewards = rewards * mask  # (B, RT)
-
-  # The lax.scan version of this is slow, but we still show it here for
-  # completeness.
-  #   rewards_rev = np.flip(masked_rewards, axis=1)  # (B, T) flipped on time.
-  #   rrt = np.transpose(rewards_rev)  # (T, B) transpose to scan over time.
-  #
-  #   def discounting_add(carry, reward):
-  #     x = reward + (gamma * carry)
-  #     return x, x
-  #
-  #   _, ys = lax.scan(discounting_add,
-  #                    np.zeros_like(rrt[0], dtype=np.float32),
-  #                    rrt.astype(np.float32))
-  #
-  #   # ys is (T, B) and T is in reverse order.
-  #   return np.flip(np.transpose(ys), axis=1)
-
-  # We use the following recurrence relation, derived from the equation above:
-  #
-  # r2g[t+1] = (r2g[t] - r[t]) / gamma
-  #
-  # This means we'll need to calculate r2g[0] first and then r2g[1] and so on ..
-  #
-  # **However** this leads to overflows for long sequences: r2g[t] - r[t] > 0
-  # and gamma < 1.0, so the division keeps increasing.
-  #
-  # So we just run the recurrence in reverse, i.e.
-  #
-  # r2g[t] = r[t] + (gamma*r2g[t+1])
-  #
-  # This is much better, but might have lost updates since the (small) rewards
-  # at earlier time-steps may get added to a (very?) large sum.
-
-  # Compute r2g_{T-1} at the start and then compute backwards in time.
-  r2gs = [masked_rewards[:, -1]]
-
-  # Go from T-2 down to 0.
-  for t in reversed(range(RT - 1)):
-    r2gs.append(masked_rewards[:, t] + (gamma * r2gs[-1]))
-
-  # The list should have length RT.
-  assert RT == len(r2gs)
-
-  # First we stack them in the correct way to make it (B, RT), but these are
-  # still from newest (RT-1) to oldest (0), so then we flip it on time axis.
-  return np.flip(np.stack(r2gs, axis=1), axis=1)
-
-
-@jit
-def value_loss_given_predictions(value_prediction,
-                                 rewards,
-                                 reward_mask,
-                                 gamma=0.99,
-                                 epsilon=0.2,
-                                 value_prediction_old=None):
-  """Computes the value loss given the prediction of the value function.
-
-  Args:
-    value_prediction: np.ndarray of shape (B, RT+1, 1)
-    rewards: np.ndarray of shape (B, RT) of rewards.
-    reward_mask: np.ndarray of shape (B, RT), the mask over rewards.
-    gamma: float, discount factor.
-    epsilon: float, clip-fraction, used if value_value_prediction_old isn't None
-    value_prediction_old: np.ndarray of shape (B, RT+1, 1) of value predictions
-      using the old parameters. If provided, we incorporate this in the loss as
-      well. This is from the OpenAI baselines implementation.
-
-  Returns:
-    Pair (value_loss, summaries), where value_loss is the average L2 value loss,
-      averaged over instances where reward_mask is 1. Summaries is a dict of
-      summaries collected during value loss computation.
-  """
-
-  B, RT = rewards.shape  # pylint: disable=invalid-name
-  assert (B, RT) == reward_mask.shape
-  assert (B, RT + 1) == value_prediction.shape
-
-  value_prediction = value_prediction[:, :-1] * reward_mask  # (B, RT)
-  r2g = rewards_to_go(rewards, reward_mask, gamma=gamma)  # (B, RT)
-  loss = (value_prediction - r2g)**2
-
-  # From the baselines implementation.
-  if value_prediction_old is not None:
-    value_prediction_old = value_prediction_old[:, :-1] * reward_mask  # (B, RT)
-
-    v_clipped = value_prediction_old + np.clip(
-        value_prediction - value_prediction_old, -epsilon, epsilon)
-    v_clipped_loss = (v_clipped - r2g)**2
-    loss = np.maximum(v_clipped_loss, loss)
-
-  # Take an average on only the points where mask != 0.
-  value_loss = np.sum(loss) / np.sum(reward_mask)
-
-  summaries = {
-      "value_loss": value_loss,
-  }
-
-  return (value_loss, summaries)
-
-
-def deltas(predicted_values, rewards, mask, gamma=0.99):
-  r"""Computes TD-residuals from V(s) and rewards.
-
-  Where a `delta`, i.e. a td-residual is defined as:
-
-  delta_{b,t} = r_{b,t} + \gamma * v_{b,t+1} - v_{b,t}.
-
-  Args:
-    predicted_values: ndarray of shape (B, RT+1). NOTE: Expects axis 2 was
-      squeezed. These represent V(s_bt) for b < B and t < RT+1
-    rewards: ndarray of shape (B, RT) of rewards.
-    mask: ndarray of shape (B, RT) of mask for rewards.
-    gamma: float, discount factor.
-
-  Returns:
-    ndarray of shape (B, RT) of one-step TD-residuals.
-  """
-
-  # Predicted values at time t, cutting off the last to have shape (B, RT).
-  predicted_values_bt = predicted_values[:, :-1]
-  # Predicted values at time t+1, by cutting off the first to have shape (B, RT)
-  predicted_values_btplus1 = predicted_values[:, 1:]
-  # Return the deltas as defined above.
-  return (rewards +
-          (gamma * predicted_values_btplus1) - predicted_values_bt) * mask
-
-
-def gae_advantages(td_deltas, mask, lambda_=0.95, gamma=0.99):
-  r"""Computes the GAE advantages given the one step TD-residuals.
-
-  The formula for a GAE advantage estimator is as follows:
-
-  A_{bt} = \sum_{l=0}^{\infty}(\gamma * \lambda)^{l}(\delta_{b,t+l}).
-
-  Internally we just call rewards_to_go, since it is the same computation.
-
-  Args:
-    td_deltas: np.ndarray of shape (B, RT) of one step TD-residuals.
-    mask: np.ndarray of shape (B, T) of mask for the residuals. It maybe the
-      case that the `td_deltas` are already masked correctly since they are
-      produced by `deltas(...)`
-    lambda_: float, lambda parameter for GAE estimators.
-    gamma: float, lambda parameter for GAE estimators.
-
-  Returns:
-    GAE advantage estimates.
-  """
-
-  return rewards_to_go(td_deltas, mask, lambda_ * gamma)
-
-
-def chosen_probabs(probab_actions, actions):
-  """Picks out the probabilities of the actions along batch and time-steps.
-
-  Args:
-    probab_actions: ndarray of shape `[B, AT, A]`, where
-      probab_actions[b, t, i] contains the log-probability of action = i at
-      the t^th time-step in the b^th trajectory.
-    actions: ndarray of shape `[B, AT]`, with each entry in [0, A) denoting
-      which action was chosen in the b^th trajectory's t^th time-step.
-
-  Returns:
-    `[B, AT, A]` ndarray with the log-probabilities of the chosen actions.
-  """
-  B, AT = actions.shape  # pylint: disable=invalid-name
-  assert (B, AT) == probab_actions.shape[:2]
-  return probab_actions[np.arange(B)[:, None], np.arange(AT), actions]
-
-
-def compute_probab_ratios(p_new, p_old, actions, action_mask):
-  """Computes the probability ratios for each time-step in a trajectory.
-
-  Args:
-    p_new: ndarray of shape [B, AT, A] of the log-probabilities that the
-      policy network assigns to all the actions at each time-step in each batch
-      using the old parameters.
-    p_old: ndarray of shape [B, AT, A], same as above, but using old policy
-      network parameters.
-    actions: ndarray of shape [B, AT] where each element is from [0, A).
-    action_mask: ndarray of shape [B, T] masking over probabilities.
-
-  Returns:
-    probab_ratios: ndarray of shape [B, AT], where
-    probab_ratios_{b,t,} = p_new_{b,t,action_{b,t}} /
-                           p_old_{b,t,action_{b,t}}
-  """
-
-  B, AT = actions.shape  # pylint: disable=invalid-name
-  assert (B, AT) == p_old.shape[:2]
-  assert (B, AT) == p_new.shape[:2]
-
-  logp_old = chosen_probabs(p_old, actions)
-  logp_new = chosen_probabs(p_new, actions)
-
-  assert (B, AT) == logp_old.shape
-  assert (B, AT) == logp_new.shape
-
-  # Since these are log-probabilities, we just subtract them.
-  probab_ratios = np.exp(logp_new - logp_old) * action_mask
-  assert (B, AT) == probab_ratios.shape
-  return probab_ratios
-
-
-def clipped_probab_ratios(probab_ratios, epsilon=0.2):
-  return np.clip(probab_ratios, 1 - epsilon, 1 + epsilon)
-
-
-def clipped_objective(probab_ratios, advantages, action_mask, epsilon=0.2):
-  advantages = advantages
-  return np.minimum(
-      probab_ratios * advantages,
-      clipped_probab_ratios(probab_ratios, epsilon=epsilon) *
-      advantages) * action_mask
-
-
-@jit
-def ppo_loss_given_predictions(log_probab_actions_new,
-                               log_probab_actions_old,
-                               value_predictions_old,
-                               padded_actions,
-                               rewards_to_actions,
-                               padded_rewards,
-                               reward_mask,
-                               gamma=0.99,
-                               lambda_=0.95,
-                               epsilon=0.2):
-  """PPO objective, with an eventual minus sign, given predictions."""
-  B, RT = padded_rewards.shape  # pylint: disable=invalid-name
-  _, AT, A = log_probab_actions_old.shape  # pylint: disable=invalid-name
-
-  assert (B, RT) == padded_rewards.shape
-  assert (B, AT) == padded_actions.shape
-  assert (B, RT) == reward_mask.shape
-
-  assert (B, RT + 1) == value_predictions_old.shape
-  assert (B, AT, A) == log_probab_actions_old.shape
-  assert (B, AT, A) == log_probab_actions_new.shape
-
-  assert (RT + 1, AT) == rewards_to_actions.shape
-
-  # (B, RT)
-  td_deltas = deltas(
-      value_predictions_old,  # (B, RT+1)
-      padded_rewards,
-      reward_mask,
-      gamma=gamma)
-
-  # (B, RT)
-  advantages = gae_advantages(
-      td_deltas, reward_mask, lambda_=lambda_, gamma=gamma)
-
-  # Normalize the advantages.
-  advantage_mean = np.mean(advantages)
-  advantage_std = np.std(advantages)
-  advantages = (advantages - advantage_mean) / (advantage_std + 1e-8)
-
-  # Scatter advantages over padded_actions.
-  # rewards_to_actions is RT + 1 -> AT, so we pad the advantages and the reward
-  # mask by 1.
-  advantages = np.dot(np.pad(advantages, ((0, 0), (0, 1))), rewards_to_actions)
-  action_mask = np.dot(
-      np.pad(reward_mask, ((0, 0), (0, 1))), rewards_to_actions
-  )
-
-  # (B, AT)
-  ratios = compute_probab_ratios(log_probab_actions_new, log_probab_actions_old,
-                                 padded_actions, action_mask)
-  assert (B, AT) == ratios.shape
-
-  # (B, AT)
-  objective = clipped_objective(
-      ratios, advantages, action_mask, epsilon=epsilon)
-  assert (B, AT) == objective.shape
-
-  # ()
-  average_objective = np.sum(objective) / np.sum(action_mask)
-
-  # Loss is negative objective.
-  ppo_loss = -average_objective
-
-  summaries = {
-      "ppo_loss": ppo_loss,
-      "advantage_mean": advantage_mean,
-      "advantage_std": advantage_std,
-  }
-
-  return (ppo_loss, summaries)
-
-
-@jit
-def combined_loss_given_predictions(log_probab_actions_new,
-                                    log_probab_actions_old,
-                                    value_prediction_new,
-                                    value_prediction_old,
-                                    padded_actions,
-                                    rewards_to_actions,
-                                    padded_rewards,
-                                    reward_mask,
-                                    gamma=0.99,
-                                    lambda_=0.95,
-                                    epsilon=0.2,
-                                    c1=1.0,
-                                    c2=0.01):
-  """Computes the combined (clipped loss + value loss) given predictions."""
-  # Sum values over symbols in an action's representation, because it's a simple
-  # way of going from AT to RT+1 and does not decrease the expressive power.
-  value_prediction_old = np.dot(
-      value_prediction_old, rewards_to_actions.transpose()
-  )
-  value_prediction_new = np.dot(
-      value_prediction_new, rewards_to_actions.transpose()
-  )
-  (value_loss, value_summaries) = value_loss_given_predictions(
-      value_prediction_new,
-      padded_rewards,
-      reward_mask,
-      gamma=gamma,
-      value_prediction_old=value_prediction_old,
-      epsilon=epsilon)
-  (ppo_loss, ppo_summaries) = ppo_loss_given_predictions(
-      log_probab_actions_new,
-      log_probab_actions_old,
-      value_prediction_old,
-      padded_actions,
-      rewards_to_actions,
-      padded_rewards,
-      reward_mask,
-      gamma=gamma,
-      lambda_=lambda_,
-      epsilon=epsilon)
-  # Pad the reward mask to be compatible with rewards_to_actions.
-  padded_reward_mask = np.pad(reward_mask, ((0, 0), (0, 1)))
-  action_mask = np.dot(padded_reward_mask, rewards_to_actions)
-  entropy_bonus = masked_entropy(log_probab_actions_new, action_mask)
-  combined_loss_ = ppo_loss + (c1 * value_loss) - (c2 * entropy_bonus)
-
-  summaries = {
-      "combined_loss": combined_loss_,
-      "entropy_bonus": entropy_bonus,
-  }
-  for loss_summaries in (value_summaries, ppo_summaries):
-    summaries.update(loss_summaries)
-
-  return (combined_loss_, (ppo_loss, value_loss, entropy_bonus), summaries)
-
-
-@functools.partial(jit, static_argnums=(3,))
-def combined_loss(new_params,
-                  log_probab_actions_old,
-                  value_predictions_old,
-                  policy_and_value_net_apply,
-                  padded_observations,
-                  padded_actions,
-                  rewards_to_actions,
-                  padded_rewards,
-                  reward_mask,
-                  gamma=0.99,
-                  lambda_=0.95,
-                  epsilon=0.2,
-                  c1=1.0,
-                  c2=0.01,
-                  state=None,
-                  rng=None):
-  """Computes the combined (clipped loss + value loss) given observations."""
-  (log_probab_actions_new, value_predictions_new) = (
-      policy_and_value_net_apply(
-          padded_observations, params=new_params, state=state, rng=rng))
-
-  (loss, component_losses, summaries) = combined_loss_given_predictions(
-      log_probab_actions_new,
-      log_probab_actions_old,
-      value_predictions_new,
-      value_predictions_old,
-      padded_actions,
-      rewards_to_actions,
-      padded_rewards,
-      reward_mask,
-      gamma=gamma,
-      lambda_=lambda_,
-      epsilon=epsilon,
-      c1=c1,
-      c2=c2,
-  )
-  return (loss, component_losses, summaries, state)
-
-
-@functools.partial(jit, static_argnums=(2, 3, 4))
-def policy_and_value_opt_step(i,
-                              opt_state,
-                              opt_update,
-                              get_params,
-                              policy_and_value_net_apply,
-                              log_probab_actions_old,
-                              value_predictions_old,
-                              padded_observations,
-                              padded_actions,
-                              rewards_to_actions,
-                              padded_rewards,
-                              reward_mask,
-                              c1=1.0,
-                              c2=0.01,
-                              gamma=0.99,
-                              lambda_=0.95,
-                              epsilon=0.1,
-                              state=None,
-                              rng=None):
-  """Policy and Value optimizer step."""
-
-  # Combined loss function given the new params.
-  def policy_and_value_loss(params, state):
-    """Returns the combined loss given just parameters."""
-    (loss, _, _, state) = combined_loss(
-        params,
-        log_probab_actions_old,
-        value_predictions_old,
-        policy_and_value_net_apply,
-        padded_observations,
-        padded_actions,
-        rewards_to_actions,
-        padded_rewards,
-        reward_mask,
-        c1=c1,
-        c2=c2,
-        gamma=gamma,
-        lambda_=lambda_,
-        epsilon=epsilon,
-        state=state,
-        rng=rng)
-    return loss, state
-
-  new_params = get_params(opt_state)
-  g, state = grad(policy_and_value_loss, has_aux=True)(new_params, state)
-  # TODO(afrozm): Maybe clip gradients?
-  return opt_update(i, g, opt_state), state
-
-
-def get_time(t1, t2=None):
-  if t2 is None:
-    t2 = time.time()
-  return round((t2 - t1) * 1000, 2)
-
-
-def approximate_kl(log_prob_new, log_prob_old, mask):
-  """Computes the approximate KL divergence between the old and new log-probs.
-
-  Args:
-    log_prob_new: (B, AT, A) log probs new
-    log_prob_old: (B, AT, A) log probs old
-    mask: (B, AT)
-
-  Returns:
-    Approximate KL.
-  """
-  diff = log_prob_old - log_prob_new
-  # Mask out the irrelevant part.
-  diff *= mask[:, :, np.newaxis]  # make mask (B, RT, 1)
-  # Average on non-masked part.
-  return np.sum(diff) / np.sum(mask)
-
-
-def masked_entropy(log_probs, mask):
-  """Computes the entropy for the given log-probs.
-
-  Args:
-    log_probs: (B, AT, A) log probs
-    mask: (B, AT) mask.
-
-  Returns:
-    Entropy.
-  """
-  # Mask out the irrelevant part.
-  lp = log_probs * mask[:, :, np.newaxis]  # make mask (B, RT, 1)
-  p = np.exp(lp) * mask[:, :, np.newaxis]  # (B, RT, 1)
-  # Average on non-masked part and take negative.
-  return -(np.sum(lp * p) / np.sum(mask))
-
-
-def get_policy_model_files(output_dir):
-  return list(
-      reversed(
-          sorted(gfile.glob(os.path.join(output_dir, "model-??????.pkl")))))
-
-
-def get_epoch_from_policy_model_file(policy_model_file):
-  base_name = os.path.basename(policy_model_file)
-  return int(re.match(r"model-(\d+).pkl", base_name).groups()[0])
-
-
-def get_policy_model_file_from_epoch(output_dir, epoch):
-  return os.path.join(output_dir, "model-%06d.pkl" % epoch)
-
-
-def maybe_restore_opt_state(output_dir,
-                            policy_and_value_opt_state=None,
-                            policy_and_value_state=None):
-  """Maybe restore the optimization state from the checkpoint dir.
-
-  Optimization state includes parameters and optimizer slots.
-
-  Args:
-    output_dir: Directory where saved model checkpoints are stored.
-    policy_and_value_opt_state: Default optimization state, returned if model
-      isn't found.
-    policy_and_value_state: state of the policy and value network.
-
-  Returns:
-    tuple (opt_state, state, epoch (int), opt_step (int)) where epoch is the
-    epoch from which we restored the optimization state, 0 if no checkpoint was
-    found, and opt_step is the total optimization step (sum of all optimization
-    steps made up to the current epoch).
-  """
-  pkl_module = utils.get_pickle_module()
-  epoch = 0
-  total_opt_step = 0
-  for model_file in get_policy_model_files(output_dir):
-    logging.info("Trying to restore model from %s", model_file)
-    try:
-      with gfile.GFile(model_file, "rb") as f:
-        policy_and_value_opt_state, policy_and_value_state, total_opt_step = (
-            pkl_module.load(f))
-      epoch = get_epoch_from_policy_model_file(model_file)
-      break
-    except EOFError as e:
-      logging.error("Unable to load model from: %s with %s", model_file, e)
-      # Try an older version.
-      continue
-  return (
-      policy_and_value_opt_state,
-      policy_and_value_state,
-      epoch,
-      total_opt_step,
-  )
-
-
-LAST_N_POLICY_MODELS_TO_KEEP = 5
-
-
-def save_opt_state(output_dir,
-                   policy_and_value_opt_state,
-                   policy_and_value_state,
-                   epoch,
-                   total_opt_step):
-  """Saves the policy and value network optimization state etc."""
-  pkl_module = utils.get_pickle_module()
-  old_model_files = get_policy_model_files(output_dir)
-  params_file = os.path.join(output_dir, "model-%06d.pkl" % epoch)
-  with gfile.GFile(params_file, "wb") as f:
-    pkl_module.dump(
-        (policy_and_value_opt_state, policy_and_value_state, total_opt_step), f)
-  # Keep the last k model files lying around (note k > 1 because the latest
-  # model file might be in the process of getting read async).
-  for path in old_model_files[LAST_N_POLICY_MODELS_TO_KEEP:]:
-    if path != params_file:
-      gfile.remove(path)
-
-
-def init_policy_from_world_model_checkpoint(policy_params, model_output_dir):
-  """Initializes policy parameters from world model parameters."""
-  pkl_module = utils.get_pickle_module()
-  params_file = os.path.join(model_output_dir, "model.pkl")
-  # Don't use trax.restore_state to avoid a circular import.
-  with gfile.GFile(params_file, "rb") as f:
-    model_params = pkl_module.load(f)[0][0]
-  # TODO(pkozakowski): The following, brittle line of code is hardcoded for
-  # transplanting parameters from TransformerLM to TransformerDecoder-based
-  # policy network of the same configuration. Figure out a more general method.
-  policy_params[0] = model_params[0][1:-2]
-  return policy_params
-
-
-def write_eval_reward_summaries(reward_stats_by_mode, summary_writer, epoch):
-  """Writes evaluation reward statistics to summary and logs them.
-
-  Args:
-    reward_stats_by_mode: Nested dict of structure: {
-          "raw": {
-              <temperature 1>: {
-                  "mean": <reward mean>,
-                  "std": <reward std>, },
-              <temperature 2>: ... },
-          "processed": ... }
-    summary_writer: jaxboard.SummaryWriter.
-    epoch: Current epoch number.
-  """
-  for (reward_mode, reward_stats_by_temp) in reward_stats_by_mode.items():
-    for (temperature, reward_stats) in reward_stats_by_temp.items():
-      for (stat_name, stat) in reward_stats.items():
-        summary_writer.scalar(
-            "eval/{reward_mode}_reward_{stat_name}/"
-            "temperature_{temperature}".format(
-                reward_mode=reward_mode,
-                stat_name=stat_name,
-                temperature=temperature),
-            stat,
-            step=epoch)
-      logging.info(
-          "Epoch [% 6d] Policy Evaluation (%s reward) "
-          "[temperature %.2f] = %10.2f (+/- %.2f)", epoch, reward_mode,
-          temperature, reward_stats["mean"], reward_stats["std"])
-
-
-def shuffled_index_batches(dataset_size, batch_size):
-  """Generates batches of shuffled indices over a dataset."""
-  def shuffled_indices():
-    while True:
-      perm = onp.random.permutation(dataset_size)
-      for x in perm:
-        yield x
-
-  indices = shuffled_indices()
-  while True:
-    yield onp.array(list(itertools.islice(indices, int(batch_size))))
diff --git a/tensor2tensor/trax/rl/ppo_test.py b/tensor2tensor/trax/rl/ppo_test.py
deleted file mode 100644
index 1e48dc87d..000000000
--- a/tensor2tensor/trax/rl/ppo_test.py
+++ /dev/null
@@ -1,643 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.rl.ppo."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import functools
-import itertools
-
-import jax
-from jax import random as jax_random
-import numpy as np
-from tensor2tensor.trax import inputs
-from tensor2tensor.trax import layers
-from tensor2tensor.trax import models
-from tensor2tensor.trax import trax
-from tensor2tensor.trax.rl import ppo
-from tensorflow import test
-from tensorflow.io import gfile
-
-
-class PpoTest(test.TestCase):
-
-  def setUp(self):
-    super(PpoTest, self).setUp()
-    self.rng_key = trax.get_random_number_generator_and_set_seed(0)
-
-  def test_get_policy_model_files(self):
-    output_dir = self.get_temp_dir()
-
-    def write_policy_model_file(epoch):
-      with gfile.GFile(
-          ppo.get_policy_model_file_from_epoch(output_dir, epoch), "w") as f:
-        f.write("some data")
-
-    epochs = [200, 100, 300]
-
-    # 300, 200, 100
-    expected_policy_model_files = [
-        output_dir + "/model-000300.pkl",
-        output_dir + "/model-000200.pkl",
-        output_dir + "/model-000100.pkl",
-    ]
-
-    for epoch in epochs:
-      write_policy_model_file(epoch)
-
-    policy_model_files = ppo.get_policy_model_files(output_dir)
-
-    self.assertEqual(expected_policy_model_files, policy_model_files)
-
-    gfile.rmtree(output_dir)
-
-  def test_get_epoch_from_policy_model_file(self):
-    self.assertEqual(0,
-                     ppo.get_epoch_from_policy_model_file("model-000000.pkl"))
-    self.assertEqual(123456,
-                     ppo.get_epoch_from_policy_model_file("model-123456.pkl"))
-
-  def test_get_policy_model_file_from_epoch(self):
-    self.assertEqual("/tmp/model-000000.pkl",
-                     ppo.get_policy_model_file_from_epoch("/tmp", 0))
-    self.assertEqual("/tmp/model-123456.pkl",
-                     ppo.get_policy_model_file_from_epoch("/tmp", 123456))
-
-  def test_policy_and_value_net(self):
-    observation_shape = (3, 4, 5)
-    batch_observation_shape = (1, 1) + observation_shape
-    n_actions = 2
-    n_controls = 3
-    pnv_model = ppo.policy_and_value_net(
-        n_controls=n_controls,
-        n_actions=n_actions,
-        vocab_size=None,
-        bottom_layers_fn=lambda: [layers.Flatten(n_axes_to_keep=2)],
-        two_towers=True,
-    )
-    _, _ = pnv_model.initialize_once(
-        batch_observation_shape, np.float32, self.rng_key)
-
-    batch = 2
-    time_steps = 10
-    batch_of_observations = np.random.uniform(
-        size=(batch, time_steps) + observation_shape)
-    pnv_output = pnv_model(batch_of_observations)
-
-    # Output is a list, first is probab of actions and the next is value output.
-    self.assertEqual(2, len(pnv_output))
-    self.assertEqual(
-        (batch, time_steps * n_controls, n_actions), pnv_output[0].shape)
-    self.assertEqual((batch, time_steps * n_controls), pnv_output[1].shape)
-
-  def test_pad_trajectories(self):
-    observation_shape = (2, 3, 4)
-    trajectories = []
-    n_trajectories = 7
-    n_actions = 10
-
-    # Time-steps are between [min_allowable_time_step, max_allowable_time_step]
-    max_allowable_time_step = 19
-    min_allowable_time_step = 5
-
-    # The actual max we see in the data.
-    max_time_step = -1
-
-    # Bucket length.
-    bucket_length = 15
-
-    # Make `n_trajectories` random trajectories.
-    for i in range(n_trajectories):
-      time_steps = np.random.randint(min_allowable_time_step,
-                                     max_allowable_time_step + 1)
-      if time_steps > max_time_step:
-        max_time_step = time_steps
-      observations = np.random.randint(
-          0, 255, size=(time_steps + 1,) + observation_shape).astype(np.uint8)
-      rewards = np.random.uniform(size=(time_steps,)).astype(np.float32)
-      actions = np.random.randint(
-          0, n_actions, size=(time_steps,)).astype(np.int32)
-      infos = {
-          "a": np.random.uniform(size=(time_steps,)).astype(np.float32),
-          "b": np.random.uniform(size=(time_steps,)).astype(np.float32)
-      }
-      trajectories.append((observations, rewards, actions, infos))
-
-    # Now pad these trajectories.
-    padded_trajectories = ppo.pad_trajectories(
-        trajectories, boundary=bucket_length)
-
-    # Expected padding.
-    i = 1
-    while i * bucket_length < max_time_step:
-      i += 1
-    expected_padding = i * bucket_length
-
-    # Get the padded objects.
-    (pad_lengths, reward_mask, padded_observations, padded_actions,
-     padded_rewards, padded_infos) = padded_trajectories
-
-    # Expectations on the padded shapes.
-    self.assertEqual(padded_observations.shape, (
-        n_trajectories,
-        expected_padding + 1,
-    ) + observation_shape)
-    self.assertEqual(padded_actions.shape, (n_trajectories, expected_padding))
-    self.assertEqual(padded_rewards.shape, (n_trajectories, expected_padding))
-    self.assertEqual(reward_mask.shape, (n_trajectories, expected_padding))
-
-    self.assertEqual(padded_infos["a"].shape,
-                     (n_trajectories, expected_padding))
-    self.assertEqual(padded_infos["b"].shape,
-                     (n_trajectories, expected_padding))
-
-    # Assert that the padding lengths and reward mask are consistent.
-    self.assertAllEqual(
-        np.full((n_trajectories,), expected_padding),
-        np.array(np.sum(reward_mask, axis=1)) + pad_lengths)
-
-  def test_rewards_to_go(self):
-    rewards = np.array([
-        [1, 2, 4, 8, 16, 32, 64, 128],
-        [1, 1, 1, 1, 1, 1, 1, 1],
-    ])
-
-    rewards_mask = np.array([
-        [1, 1, 1, 1, 1, 0, 0, 0],
-        [1, 1, 1, 1, 1, 1, 1, 0],
-    ])
-
-    gamma = 0.5
-
-    rewards_to_go = ppo.rewards_to_go(rewards, rewards_mask, gamma)
-
-    self.assertAllEqual(
-        np.array([
-            [5, 8, 12, 16, 16, 0, 0, 0],
-            [1.984375, 1.96875, 1.9375, 1.875, 1.75, 1.5, 1.0, 0],
-        ]), rewards_to_go)
-
-  def test_rewards_to_go_really_long_sequences(self):
-    T = 1200  # pylint: disable=invalid-name
-
-    rewards = np.random.uniform(1e-3, 1e-2, (1, T))
-
-    # Make a mask, clear out a fixed number `L` of 1s from the end.
-    L = 36  # pylint: disable=invalid-name
-    assert L < T
-    rewards_mask = np.ones_like(rewards)
-    rewards_mask[0, L:] = 0
-
-    gamma = 0.94
-
-    actual_r2g = ppo.rewards_to_go(rewards, rewards_mask, gamma).reshape(-1)
-
-    # Let's compute r2g the slow way.
-    masked_rewards = (rewards_mask * rewards).reshape(-1)
-    expected_r2g = np.zeros_like(masked_rewards)
-    for t in range(T):
-      for j in range(t, T):
-        expected_r2g[t] += (gamma**(j - t)) * masked_rewards[j]
-
-    self.assertAllClose(expected_r2g, actual_r2g)
-
-  def test_value_loss(self):
-    rewards = np.array([
-        [1, 2, 4, 8, 16, 32, 64, 128],
-        [1, 1, 1, 1, 1, 1, 1, 1],
-    ])
-
-    rewards_mask = np.array([
-        [1, 1, 1, 1, 1, 0, 0, 0],
-        [1, 1, 1, 1, 1, 1, 1, 0],
-    ])
-
-    gamma = 0.5
-
-    # Random observations and a value function that returns a constant value.
-    # NOTE: Observations have an extra time-step.
-    B, T = rewards.shape  # pylint: disable=invalid-name
-    observation_shape = (210, 160, 3)  # atari pong
-    random_observations = np.random.uniform(size=(B, T + 1) + observation_shape)
-
-    def value_net_apply(observations, params, rng=None):
-      del params, rng
-      # pylint: disable=invalid-name
-      B, T_p_1, OBS = (observations.shape[0], observations.shape[1],
-                       observations.shape[2:])
-      del OBS
-      return np.ones((B, T_p_1))
-      # pylint: enable=invalid-name
-
-    value_prediction = value_net_apply(random_observations, [])
-
-    with jax.disable_jit():
-      (value_loss, _) = ppo.value_loss_given_predictions(
-          value_prediction,
-          rewards,
-          rewards_mask,
-          gamma)
-
-    self.assertNear(53.3637084961, value_loss, 1e-6)
-
-  def test_deltas(self):
-    rewards = np.array([
-        [1, 2, 4, 8, 16, 32, 64, 128],
-        [1, 1, 1, 1, 1, 1, 1, 1],
-    ])
-
-    rewards_mask = np.array([
-        [1, 1, 1, 1, 1, 0, 0, 0],
-        [1, 1, 1, 1, 1, 1, 1, 0],
-    ])
-
-    B, T = rewards.shape  # pylint: disable=invalid-name
-
-    # Say, all predicted values are 1.
-    predicted_values = np.ones((B, T + 1))
-
-    gamma = 1.0
-
-    td_residuals = ppo.deltas(predicted_values, rewards, rewards_mask, gamma)
-
-    # With V(s) being the same for all s, td_residuals should be
-    # equal to the rewards + (\gamma - 1)*v(s), masked in the right places.
-    truncated_pv = predicted_values[:, :-1]
-    masked_rewards = rewards * rewards_mask
-    expected_residuals = (masked_rewards +
-                          (gamma - 1) * truncated_pv) * rewards_mask
-    self.assertAllEqual(expected_residuals, td_residuals)
-
-    gamma = 0.5
-    td_residuals = ppo.deltas(predicted_values, rewards, rewards_mask, gamma)
-    expected_residuals = (masked_rewards +
-                          (gamma - 1) * truncated_pv) * rewards_mask
-    self.assertAllEqual(expected_residuals, td_residuals)
-
-  def test_gae_advantages(self):
-    td_deltas = np.array([
-        [1, 2, 4, 8, 16, 32, 64, 128],
-        [1, 1, 1, 1, 1, 1, 1, 1],
-    ])
-
-    rewards_mask = np.array([
-        [1, 1, 1, 1, 1, 0, 0, 0],
-        [1, 1, 1, 1, 1, 1, 1, 0],
-    ])
-
-    gamma = 0.5
-    lambda_ = 1.0
-
-    expected_gae_advantages = np.array([
-        [5, 8, 12, 16, 16, 0, 0, 0],
-        [1.984375, 1.96875, 1.9375, 1.875, 1.75, 1.5, 1.0, 0],
-    ])
-
-    gae_advantages = ppo.gae_advantages(td_deltas * rewards_mask, rewards_mask,
-                                        lambda_, gamma)
-    self.assertAllEqual(expected_gae_advantages, gae_advantages)
-
-    gamma = 1.0
-    lambda_ = 0.5
-
-    gae_advantages = ppo.gae_advantages(td_deltas * rewards_mask, rewards_mask,
-                                        lambda_, gamma)
-    self.assertAllEqual(expected_gae_advantages, gae_advantages)
-
-  def test_chosen_probabs(self):
-    # Shape (2, 2, 3)
-    probab_observations = np.array(
-        [[[0.1, 0.2, 0.7], [0.4, 0.1, 0.5]],
-         [[0.3, 0.1, 0.6], [0.1, 0.1, 0.8]]]
-    )
-
-    # Shape (2, 2, 1)
-    actions = np.array([[1, 2], [0, 1]])
-
-    chosen_probabs = ppo.chosen_probabs(probab_observations, actions)
-
-    self.assertAllEqual(
-        np.array([[0.2, 0.5], [0.3, 0.1]]), chosen_probabs)
-
-  def test_compute_probab_ratios(self):
-    p_old = np.array([[
-        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
-        [np.log(0.4), np.log(0.1), np.log(0.4), np.log(0.1)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
-    ], [
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-        [np.log(0.1), np.log(0.1), np.log(0.4), np.log(0.4)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
-    ]])
-
-    p_new = np.array([[
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-        [np.log(0.4), np.log(0.1), np.log(0.1), np.log(0.3)],
-        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-    ], [
-        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
-        [np.log(0.1), np.log(0.1), np.log(0.2), np.log(0.6)],
-        [np.log(0.3), np.log(0.1), np.log(0.3), np.log(0.3)],
-        [np.log(0.1), np.log(0.2), np.log(0.1), np.log(0.6)],
-    ]])
-
-    actions = np.array([[1, 2, 0, 1], [0, 3, 3, 0]])
-
-    mask = np.array([[1, 1, 0, 0], [1, 1, 1, 0]])
-
-    probab_ratios = ppo.compute_probab_ratios(p_new, p_old, actions, mask)
-
-    self.assertAllClose(
-        np.array([
-            [0.1 / 0.2, 0.1 / 0.4, 0.0, 0.0],
-            [0.1 / 0.3, 0.6 / 0.4, 0.3 / 0.1, 0.0],
-        ]), probab_ratios)
-
-  def test_clipped_probab_ratios(self):
-    probab_ratios = np.array([
-        [1.5, 1.0, 0.5, 0.7],
-        [2.5, 2.0, 0.1, 1.0],
-    ])
-
-    clipped_probab_ratios = ppo.clipped_probab_ratios(probab_ratios, 0.1)
-
-    self.assertAllClose(
-        np.array([
-            [1.1, 1.0, 0.9, 0.9],
-            [1.1, 1.1, 0.9, 1.0],
-        ]), clipped_probab_ratios)
-
-  def test_clipped_objective(self):
-    probab_ratios = np.array([
-        [1.5, 2.0, 0.5, 0.7],
-        [2.5, 2.0, 0.1, 1.0],
-    ])
-
-    advantages = np.array([
-        [0.1, -0.1, 0.5, 0.7],
-        [2.0, -2.0, 2.0, 2.0],
-    ])
-
-    mask = np.array([[1, 1, 0, 0], [1, 1, 1, 0]])
-
-    epsilon = 0.1
-
-    clipped_probab_ratios = np.array([
-        [1.1, 1.1, 0.9, 0.9],
-        [1.1, 1.1, 0.9, 1.0],
-    ])
-
-    unused_advantages_x_probab_ratios = np.array([
-        [0.15, -0.2, 0.25, 0.49],
-        [5.00, -4.0, 0.20, 2.00]
-    ])
-
-    unused_advantages_x_clipped_probab_ratios = np.array([
-        [0.11, -0.11, 0.45, 0.63],
-        [2.20, -2.20, .80, 2.00]
-    ])
-
-    unused_minimums = np.array([
-        [0.11, -0.2, 0.25, 0.49],
-        [2.20, -4.0, 0.20, 2.00]
-    ])
-
-    # minimums * mask
-    objective = np.array([
-        [0.11, -0.2, 0.0, 0.],
-        [2.20, -4.0, 0.2, 0.]
-    ])
-
-    # Assert that we computed things correctly in this test.
-    self.assertAllClose(
-        np.minimum(probab_ratios * advantages,
-                   clipped_probab_ratios * advantages) * mask,
-        objective)
-
-    self.assertAllClose(
-        objective,
-        ppo.clipped_objective(probab_ratios, advantages, mask, epsilon))
-
-  def test_combined_loss(self):
-    self.rng_key, key1, key2 = jax_random.split(self.rng_key, num=3)
-
-    B, T, A, OBS = 2, 10, 2, (28, 28, 3)  # pylint: disable=invalid-name
-    batch_observation_shape = (1, 1) + OBS
-
-    net = ppo.policy_and_value_net(
-        n_controls=1,
-        n_actions=A,
-        vocab_size=None,
-        bottom_layers_fn=lambda: [layers.Flatten(n_axes_to_keep=2)],
-        two_towers=True,
-    )
-
-    old_params, _ = net.initialize_once(
-        batch_observation_shape, np.float32, key1)
-    new_params, state = net.initialize_once(
-        batch_observation_shape, np.float32, key2)
-
-    # Generate a batch of observations.
-
-    observations = np.random.uniform(size=(B, T + 1) + OBS)
-    actions = np.random.randint(0, A, size=(B, T + 1))
-    rewards = np.random.uniform(0, 1, size=(B, T))
-    mask = np.ones_like(rewards)
-
-    # Just test that this computes at all.
-    (new_log_probabs, value_predictions_new) = (
-        net(observations, params=new_params, state=state))
-    (old_log_probabs, value_predictions_old) = (
-        net(observations, params=old_params, state=state))
-
-    gamma = 0.99
-    lambda_ = 0.95
-    epsilon = 0.2
-    c1 = 1.0
-    c2 = 0.01
-
-    rewards_to_actions = np.eye(value_predictions_old.shape[1])
-    (value_loss_1, _) = ppo.value_loss_given_predictions(
-        value_predictions_new, rewards, mask, gamma=gamma,
-        value_prediction_old=value_predictions_old, epsilon=epsilon)
-    (ppo_loss_1, _) = ppo.ppo_loss_given_predictions(
-        new_log_probabs,
-        old_log_probabs,
-        value_predictions_old,
-        actions,
-        rewards_to_actions,
-        rewards,
-        mask,
-        gamma=gamma,
-        lambda_=lambda_,
-        epsilon=epsilon)
-
-    (combined_loss, (ppo_loss_2, value_loss_2, entropy_bonus), _, state) = (
-        ppo.combined_loss(new_params,
-                          old_log_probabs,
-                          value_predictions_old,
-                          net,
-                          observations,
-                          actions,
-                          rewards_to_actions,
-                          rewards,
-                          mask,
-                          gamma=gamma,
-                          lambda_=lambda_,
-                          epsilon=epsilon,
-                          c1=c1,
-                          c2=c2,
-                          state=state)
-    )
-
-    # Test that these compute at all and are self consistent.
-    self.assertGreater(entropy_bonus, 0.0)
-    self.assertNear(value_loss_1, value_loss_2, 1e-6)
-    self.assertNear(ppo_loss_1, ppo_loss_2, 1e-6)
-    self.assertNear(combined_loss,
-                    ppo_loss_2 + (c1 * value_loss_2) - (c2 * entropy_bonus),
-                    1e-6)
-
-  def test_masked_entropy(self):
-    # (2, 4+1, 4)
-    log_probs = np.array([[
-        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
-        [np.log(0.4), np.log(0.1), np.log(0.4), np.log(0.1)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-    ], [
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-        [np.log(0.1), np.log(0.1), np.log(0.4), np.log(0.4)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-        [np.log(0.1), np.log(0.2), np.log(0.6), np.log(0.1)],
-        [np.log(0.3), np.log(0.1), np.log(0.5), np.log(0.1)],
-    ]])
-
-    # (2, 4)
-    mask = np.array([
-        [1, 1, 0, 0, 0],
-        [1, 1, 1, 0, 0]
-    ])
-
-    def plp(p):
-      return p * np.log(p)
-
-    # Removing the last time-step and the masked stuff, gets us this.
-    filtered_log_probs = np.array([[
-        [plp(0.1), plp(0.2), plp(0.6), plp(0.1)],
-        [plp(0.4), plp(0.1), plp(0.4), plp(0.1)],
-        [plp(0.3), plp(0.1), plp(0.5), plp(0.1)],
-        [plp(0.1), plp(0.1), plp(0.4), plp(0.4)],
-        [plp(0.3), plp(0.1), plp(0.5), plp(0.1)],
-    ]])
-
-    self.assertNear(ppo.masked_entropy(log_probs, mask),
-                    -np.sum(filtered_log_probs) / 5.0,
-                    1e-6)
-
-  def test_saves_and_restores_opt_state(self):
-    opt_state = 123
-    state = 456
-    epoch = 7
-    opt_step = 89
-    output_dir = self.get_temp_dir()
-    ppo.save_opt_state(output_dir, opt_state, state, epoch, opt_step)
-    restored_data = ppo.maybe_restore_opt_state(output_dir)
-    self.assertEqual(restored_data, (opt_state, state, epoch, opt_step))
-
-  def test_inits_policy_by_world_model_checkpoint(self):
-    transformer_kwargs = {
-        "d_model": 1,
-        "d_ff": 1,
-        "n_layers": 1,
-        "n_heads": 1,
-        "max_len": 128,
-        "mode": "train",
-    }
-    rng = jax_random.PRNGKey(123)
-    init_kwargs = {
-        "input_shapes": (1, 1),
-        "input_dtype": np.int32,
-        "rng": rng,
-    }
-    model_fn = functools.partial(
-        models.TransformerLM, vocab_size=4, **transformer_kwargs
-    )
-    output_dir = self.get_temp_dir()
-    # Initialize a world model checkpoint by running the trainer.
-    trax.train(
-        output_dir,
-        model=model_fn,
-        inputs=functools.partial(
-            inputs.random_inputs, input_shape=(1, 1), output_shape=(1, 1)
-        ),
-        train_steps=1,
-        eval_steps=1,
-    )
-
-    policy = ppo.policy_and_value_net(
-        n_actions=3,
-        n_controls=2,
-        vocab_size=4,
-        bottom_layers_fn=functools.partial(
-            models.TransformerDecoder, **transformer_kwargs
-        ),
-        two_towers=False,
-    )
-    (policy_params, policy_state) = policy.initialize_once(**init_kwargs)
-
-    # Initialize policy parameters from world model parameters.
-    new_policy_params = ppo.init_policy_from_world_model_checkpoint(
-        policy_params, output_dir
-    )
-    # Try to run the policy with new parameters.
-    observations = np.zeros((1, 100), dtype=np.int32)
-    policy(observations, params=new_policy_params, state=policy_state, rng=rng)
-
-  def test_shuffled_index_batches_generates_valid_batch(self):
-    dataset_size = 16
-    batch_size = 4
-    stream = ppo.shuffled_index_batches(dataset_size, batch_size)
-    batch = next(stream)
-    self.assertEqual(batch.shape, (batch_size,))
-    # Assert that all indices are different.
-    self.assertEqual(len(set(batch)), batch_size)
-
-  def test_shuffled_index_batches_generates_all_indices(self):
-    dataset_size = 16
-    batch_size = 4
-    stream = ppo.shuffled_index_batches(dataset_size, batch_size)
-    indices = np.reshape(
-        list(itertools.islice(stream, dataset_size // batch_size)), -1
-    )
-    self.assertEqual(set(indices), set(range(dataset_size)))
-
-  def test_shuffled_index_batches_gives_different_permutations(self):
-    dataset_size = 256
-    batch_size = 8
-    stream1 = ppo.shuffled_index_batches(dataset_size, batch_size)
-    stream2 = ppo.shuffled_index_batches(dataset_size, batch_size)
-    self.assertFalse(np.array_equal(next(stream1), next(stream2)))
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/rl/ppo_trainer.py b/tensor2tensor/trax/rl/ppo_trainer.py
deleted file mode 100644
index af2d1e6fc..000000000
--- a/tensor2tensor/trax/rl/ppo_trainer.py
+++ /dev/null
@@ -1,843 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""PPO trainer."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import collections
-import functools
-import os
-import time
-
-from absl import logging
-import gym
-from jax import jit
-from jax import numpy as np
-from jax import random as jax_random
-import numpy as onp
-from tensor2tensor.envs import env_problem_utils
-from tensor2tensor.envs import trajectory
-from tensor2tensor.trax import jaxboard
-from tensor2tensor.trax import models as trax_models
-from tensor2tensor.trax import optimizers as trax_opt
-from tensor2tensor.trax import trax
-from tensor2tensor.trax.rl import base_trainer
-from tensor2tensor.trax.rl import ppo
-from tensor2tensor.trax.rl import serialization_utils
-from tensor2tensor.trax.rl import space_serializer
-
-DEBUG_LOGGING = False
-GAMMA = 0.99
-LAMBDA = 0.95
-EPSILON = 0.1
-EPOCHS = 50  # 100
-N_OPTIMIZER_STEPS = 100
-PRINT_EVERY_OPTIMIZER_STEP = 20
-BATCH_TRAJECTORIES = 32
-
-
-class PPO(base_trainer.BaseTrainer):
-  """PPO trainer."""
-
-  def __init__(self,
-               train_env,
-               eval_env,
-               output_dir,
-               policy_and_value_model=trax_models.FrameStackMLP,
-               policy_and_value_optimizer=functools.partial(
-                   trax_opt.Adam, learning_rate=1e-3),
-               policy_and_value_two_towers=False,
-               policy_and_value_vocab_size=None,
-               n_optimizer_steps=N_OPTIMIZER_STEPS,
-               optimizer_batch_size=64,
-               print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
-               target_kl=0.01,
-               boundary=20,
-               max_timestep=100,
-               max_timestep_eval=20000,
-               random_seed=None,
-               gamma=GAMMA,
-               lambda_=LAMBDA,
-               c1=1.0,
-               c2=0.01,
-               eval_every_n=1000,
-               save_every_n=1000,
-               done_frac_for_policy_save=0.5,
-               n_evals=1,
-               len_history_for_policy=4,
-               eval_temperatures=(1.0, 0.5),
-               separate_eval=True,
-               init_policy_from_world_model_output_dir=None,
-               **kwargs):
-    """Creates the PPO trainer.
-
-    Args:
-      train_env: gym.Env to use for training.
-      eval_env: gym.Env to use for evaluation.
-      output_dir: Output dir.
-      policy_and_value_model: Function defining the policy and value network,
-        without the policy and value heads.
-      policy_and_value_optimizer: Function defining the optimizer.
-      policy_and_value_two_towers: Whether to use two separate models as the
-        policy and value networks. If False, share their parameters.
-      policy_and_value_vocab_size: Vocabulary size of a policy and value network
-        operating on serialized representation. If None, use raw continuous
-        representation.
-      n_optimizer_steps: Number of optimizer steps.
-      optimizer_batch_size: Batch size of an optimizer step.
-      print_every_optimizer_steps: How often to log during the policy
-        optimization process.
-      target_kl: Policy iteration early stopping. Set to infinity to disable
-        early stopping.
-      boundary: We pad trajectories at integer multiples of this number.
-      max_timestep: If set to an integer, maximum number of time-steps in a
-        trajectory. Used in the collect procedure.
-      max_timestep_eval: If set to an integer, maximum number of time-steps in
-        an evaluation trajectory. Used in the collect procedure.
-      random_seed: Random seed.
-      gamma: Reward discount factor.
-      lambda_: N-step TD-error discount factor in GAE.
-      c1: Value loss coefficient.
-      c2: Entropy loss coefficient.
-      eval_every_n: How frequently to eval the policy.
-      save_every_n: How frequently to save the policy.
-      done_frac_for_policy_save: Fraction of the trajectories that should be
-        done to checkpoint the policy.
-      n_evals: Number of times to evaluate.
-      len_history_for_policy: How much of history to give to the policy.
-      eval_temperatures: Sequence of temperatures to try for categorical
-        sampling during evaluation.
-      separate_eval: Whether to run separate evaluation using a set of
-        temperatures. If False, the training reward is reported as evaluation
-        reward with temperature 1.0.
-      init_policy_from_world_model_output_dir: Model output dir for initializing
-        the policy. If None, initialize randomly.
-      **kwargs: Additional keyword arguments passed to the base class.
-    """
-    # Set in base class constructor.
-    self._train_env = None
-    self._should_reset = None
-
-    super(PPO, self).__init__(train_env, eval_env, output_dir, **kwargs)
-
-    self._n_optimizer_steps = n_optimizer_steps
-    self._optimizer_batch_size = optimizer_batch_size
-    self._print_every_optimizer_steps = print_every_optimizer_steps
-    self._target_kl = target_kl
-    self._boundary = boundary
-    self._max_timestep = max_timestep
-    self._max_timestep_eval = max_timestep_eval
-    self._gamma = gamma
-    self._lambda_ = lambda_
-    self._c1 = c1
-    self._c2 = c2
-    self._eval_every_n = eval_every_n
-    self._save_every_n = save_every_n
-    self._done_frac_for_policy_save = done_frac_for_policy_save
-    self._n_evals = n_evals
-    self._len_history_for_policy = len_history_for_policy
-    self._eval_temperatures = eval_temperatures
-    self._separate_eval = separate_eval
-
-    action_space = self.train_env.action_space
-    assert isinstance(
-        action_space, (gym.spaces.Discrete, gym.spaces.MultiDiscrete))
-    if isinstance(action_space, gym.spaces.Discrete):
-      n_actions = action_space.n
-      n_controls = 1
-    else:
-      (n_controls,) = action_space.nvec.shape
-      assert n_controls > 0
-      assert onp.min(action_space.nvec) == onp.max(action_space.nvec), (
-          "Every control must have the same number of actions.")
-      n_actions = action_space.nvec[0]
-    self._n_actions = n_actions
-    self._n_controls = n_controls
-
-    self._rng = trax.get_random_number_generator_and_set_seed(random_seed)
-    self._rng, key1 = jax_random.split(self._rng, num=2)
-
-    vocab_size = policy_and_value_vocab_size
-    self._serialized_sequence_policy = vocab_size is not None
-    if self._serialized_sequence_policy:
-      self._serialization_kwargs = self._init_serialization(vocab_size)
-    else:
-      self._serialization_kwargs = {}
-
-    # Initialize the policy and value network.
-    policy_and_value_net = ppo.policy_and_value_net(
-        n_actions=n_actions,
-        n_controls=n_controls,
-        vocab_size=vocab_size,
-        bottom_layers_fn=policy_and_value_model,
-        two_towers=policy_and_value_two_towers,
-    )
-    self._policy_and_value_net_apply = jit(policy_and_value_net)
-    (batch_obs_shape, obs_dtype) = self._batch_obs_shape_and_dtype
-    policy_and_value_net_params, self._model_state = (
-        policy_and_value_net.initialize_once(batch_obs_shape, obs_dtype, key1))
-    if init_policy_from_world_model_output_dir is not None:
-      policy_and_value_net_params = ppo.init_policy_from_world_model_checkpoint(
-          policy_and_value_net_params, init_policy_from_world_model_output_dir
-      )
-
-    # Initialize the optimizer.
-    (policy_and_value_opt_state, self._policy_and_value_opt_update,
-     self._policy_and_value_get_params) = ppo.optimizer_fn(
-         policy_and_value_optimizer, policy_and_value_net_params)
-
-    # Restore the optimizer state.
-    self._policy_and_value_opt_state = policy_and_value_opt_state
-    self._epoch = 0
-    self._total_opt_step = 0
-    self.update_optimization_state(
-        output_dir, policy_and_value_opt_state=policy_and_value_opt_state)
-
-    # Create summary writers and history.
-    self._train_sw = jaxboard.SummaryWriter(
-        os.path.join(self._output_dir, "train"))
-    self._timing_sw = jaxboard.SummaryWriter(
-        os.path.join(self._output_dir, "timing"))
-    self._eval_sw = jaxboard.SummaryWriter(
-        os.path.join(self._output_dir, "eval"))
-
-    self._n_trajectories_done = 0
-
-    self._last_saved_at = 0
-    if self._async_mode:
-      logging.info("Saving model on startup to have a model policy file.")
-      self.save()
-
-    self._rewards_to_actions = self._init_rewards_to_actions()
-
-  def _init_serialization(self, vocab_size):
-    obs_serializer = space_serializer.create(
-        self.train_env.observation_space, vocab_size=vocab_size
-    )
-    act_serializer = space_serializer.create(
-        self.train_env.action_space, vocab_size=vocab_size
-    )
-    repr_length = (
-        obs_serializer.representation_length +
-        act_serializer.representation_length
-    ) * (self._max_timestep + 1)
-    return {
-        "observation_serializer": obs_serializer,
-        "action_serializer": act_serializer,
-        "representation_length": repr_length,
-    }
-
-  def _init_rewards_to_actions(self):
-    # Linear map from the reward sequence to the action sequence, used for
-    # scattering advantages over action log-probs and some other things.
-    # It has one more timestep at the end, so it's compatible with the value
-    # predictions.
-    if not self._serialized_sequence_policy:
-      rewards_to_actions = np.eye(self._max_timestep + 1)[:, None, :]
-      rewards_to_actions = np.broadcast_to(
-          rewards_to_actions,
-          (self._max_timestep + 1, self._n_controls, self._max_timestep + 1),
-      )
-      return np.reshape(rewards_to_actions, (self._max_timestep + 1, -1))
-    else:
-      return serialization_utils.rewards_to_actions_map(
-          n_timesteps=(self._max_timestep + 1), **self._serialization_kwargs
-      )
-
-  @property
-  def _batch_obs_shape_and_dtype(self):
-    if not self._serialized_sequence_policy:
-      # Batch Observations Shape = [1, 1] + OBS, because we will eventually call
-      # policy and value networks on shape [B, T] +_OBS
-      shape = (1, 1) + self.train_env.observation_space.shape
-      dtype = self.train_env.observation_space.dtype
-    else:
-      shape = (1, 1)
-      dtype = np.int32
-    return (shape, dtype)
-
-  # Maybe restore the optimization state. If there is nothing to restore, then
-  # epoch = 0 and policy_and_value_opt_state is returned as is.
-  def update_optimization_state(self,
-                                output_dir,
-                                policy_and_value_opt_state=None):
-    (self._policy_and_value_opt_state, self._model_state, self._epoch,
-     self._total_opt_step) = ppo.maybe_restore_opt_state(
-         output_dir, policy_and_value_opt_state, self._model_state)
-
-    if self._epoch > 0:
-      logging.info("Restored parameters from epoch [%d]", self._epoch)
-
-  @property
-  def train_env(self):
-    return self._train_env
-
-  @train_env.setter
-  def train_env(self, new_train_env):
-    if self._train_env is not None:
-
-      def assert_same_space(space1, space2):
-        assert space1.shape == space2.shape
-        assert space1.dtype == space2.dtype
-
-      assert_same_space(new_train_env.observation_space,
-                        self._train_env.observation_space)
-      assert_same_space(new_train_env.action_space,
-                        self._train_env.action_space)
-      # We don't check the reward range, because PPO will work either way.
-
-    self._train_env = new_train_env
-    self._should_reset = True
-
-  @property
-  def epoch(self):
-    return self._epoch
-
-  def collect_trajectories_async(self,
-                                 env,
-                                 train=True,
-                                 n_trajectories=1,
-                                 temperature=1.0):
-    """Collects trajectories in an async manner."""
-
-    assert self._async_mode
-
-    # trajectories/train and trajectories/eval are the two subdirectories.
-    trajectory_dir = os.path.join(self._output_dir, "trajectories",
-                                  "train" if train else "eval")
-    epoch = self.epoch
-
-    logging.info(
-        "Loading [%s] trajectories from dir [%s] for epoch [%s] and temperature"
-        " [%s]", n_trajectories, trajectory_dir, epoch, temperature)
-
-    bt = trajectory.BatchTrajectory.load_from_directory(
-        trajectory_dir,
-        epoch=epoch,
-        temperature=temperature,
-        wait_forever=True,
-        n_trajectories=n_trajectories)
-
-    if bt is None:
-      logging.error(
-          "Couldn't load [%s] trajectories from dir [%s] for epoch [%s] and "
-          "temperature [%s]", n_trajectories, trajectory_dir, epoch,
-          temperature)
-      assert bt
-
-    # Doing this is important, since we want to modify `env` so that it looks
-    # like `env` was actually played and the trajectories came from it.
-    env.trajectories = bt
-
-    trajs = env_problem_utils.get_completed_trajectories_from_env(
-        env, n_trajectories)
-    n_done = len(trajs)
-    timing_info = {}
-    return trajs, n_done, timing_info, self._model_state
-
-  def collect_trajectories(self,
-                           train=True,
-                           temperature=1.0,
-                           abort_fn=None,
-                           raw_trajectory=False):
-    self._rng, key = jax_random.split(self._rng)
-
-    env = self.train_env
-    max_timestep = self._max_timestep
-    should_reset = self._should_reset
-    if not train:  # eval
-      env = self.eval_env
-      max_timestep = self._max_timestep_eval
-      should_reset = True
-
-    n_trajectories = env.batch_size
-
-    # If async, read the required trajectories for the epoch.
-    if self._async_mode:
-      trajs, n_done, timing_info, self._model_state = self.collect_trajectories_async(
-          env,
-          train=train,
-          n_trajectories=n_trajectories,
-          temperature=temperature)
-    else:
-      trajs, n_done, timing_info, self._model_state = ppo.collect_trajectories(
-          env,
-          policy_fn=self._policy_fun,
-          n_trajectories=n_trajectories,
-          max_timestep=max_timestep,
-          state=self._model_state,
-          rng=key,
-          len_history_for_policy=self._len_history_for_policy,
-          boundary=self._boundary,
-          reset=should_reset,
-          temperature=temperature,
-          abort_fn=abort_fn,
-          raw_trajectory=raw_trajectory,
-      )
-
-    if train:
-      self._n_trajectories_done += n_done
-
-    return trajs, n_done, timing_info, self._model_state
-
-  def train_epoch(self, evaluate=True):
-    """Train one PPO epoch."""
-    epoch_start_time = time.time()
-
-    # Evaluate the policy.
-    policy_eval_start_time = time.time()
-    if evaluate and (self._epoch + 1) % self._eval_every_n == 0:
-      self._rng, key = jax_random.split(self._rng, num=2)
-      self.evaluate()
-
-    policy_eval_time = ppo.get_time(policy_eval_start_time)
-
-    trajectory_collection_start_time = time.time()
-    logging.vlog(1, "PPO epoch [% 6d]: collecting trajectories.", self._epoch)
-    self._rng, key = jax_random.split(self._rng)
-    trajs, _, timing_info, self._model_state = self.collect_trajectories(
-        train=True, temperature=1.0)
-    trajs = [(t[0], t[1], t[2], t[4]) for t in trajs]
-    self._should_reset = False
-    trajectory_collection_time = ppo.get_time(trajectory_collection_start_time)
-
-    logging.vlog(1, "Collecting trajectories took %0.2f msec.",
-                 trajectory_collection_time)
-
-    rewards = np.array([np.sum(traj[2]) for traj in trajs])
-    avg_reward = np.mean(rewards)
-    std_reward = np.std(rewards)
-    max_reward = np.max(rewards)
-    min_reward = np.min(rewards)
-
-    self._train_sw.scalar(
-        "train/reward_mean_truncated", avg_reward, step=self._epoch)
-    if evaluate and not self._separate_eval:
-      metrics = {"raw": {1.0: {"mean": avg_reward, "std": std_reward}}}
-      ppo.write_eval_reward_summaries(metrics, self._eval_sw, self._epoch)
-
-    logging.vlog(1, "Rewards avg=[%0.2f], max=[%0.2f], min=[%0.2f], all=%s",
-                 avg_reward, max_reward, min_reward,
-                 [float(np.sum(traj[2])) for traj in trajs])
-
-    logging.vlog(1,
-                 "Trajectory Length average=[%0.2f], max=[%0.2f], min=[%0.2f]",
-                 float(sum(len(traj[0]) for traj in trajs)) / len(trajs),
-                 max(len(traj[0]) for traj in trajs),
-                 min(len(traj[0]) for traj in trajs))
-    logging.vlog(2, "Trajectory Lengths: %s", [len(traj[0]) for traj in trajs])
-
-    preprocessing_start_time = time.time()
-    (padded_observations, padded_actions, padded_rewards, reward_mask,
-     padded_infos) = self._preprocess_trajectories(trajs)
-    preprocessing_time = ppo.get_time(preprocessing_start_time)
-
-    logging.vlog(1, "Preprocessing trajectories took %0.2f msec.",
-                 ppo.get_time(preprocessing_start_time))
-    logging.vlog(1, "Padded Observations' shape [%s]",
-                 str(padded_observations.shape))
-    logging.vlog(1, "Padded Actions' shape [%s]", str(padded_actions.shape))
-    logging.vlog(1, "Padded Rewards' shape [%s]", str(padded_rewards.shape))
-
-    # Some assertions.
-    B, RT = padded_rewards.shape  # pylint: disable=invalid-name
-    B, AT = padded_actions.shape  # pylint: disable=invalid-name
-    assert (B, RT) == reward_mask.shape
-    assert B == padded_observations.shape[0]
-
-    log_prob_recompute_start_time = time.time()
-    # TODO(pkozakowski): The following commented out code collects the network
-    # predictions made while stepping the environment and uses them in PPO
-    # training, so that we can use non-deterministic networks (e.g. with
-    # dropout). This does not work well with serialization, so instead we
-    # recompute all network predictions. Let's figure out a solution that will
-    # work with both serialized sequences and non-deterministic networks.
-
-    # assert ("log_prob_actions" in padded_infos and
-    #         "value_predictions" in padded_infos)
-    # These are the actual log-probabs and value predictions seen while picking
-    # the actions.
-    # actual_log_probabs_traj = padded_infos["log_prob_actions"]
-    # actual_value_predictions_traj = padded_infos["value_predictions"]
-
-    # assert (B, T, C) == actual_log_probabs_traj.shape[:3]
-    # A = actual_log_probabs_traj.shape[3]  # pylint: disable=invalid-name
-    # assert (B, T, 1) == actual_value_predictions_traj.shape
-
-    del padded_infos
-
-    # TODO(afrozm): log-probabs doesn't need to be (B, T+1, C, A) it can do with
-    # (B, T, C, A), so make that change throughout.
-
-    # NOTE: We don't have the log-probabs and value-predictions for the last
-    # observation, so we re-calculate for everything, but use the original ones
-    # for all but the last time-step.
-    self._rng, key = jax_random.split(self._rng)
-
-    log_probabs_traj, value_predictions_traj, self._model_state, _ = (
-        self._get_predictions(padded_observations, self._model_state, rng=key))
-
-    assert (B, AT) == log_probabs_traj.shape[:2]
-    assert (B, AT) == value_predictions_traj.shape
-
-    # TODO(pkozakowski): Commented out for the same reason as before.
-
-    # Concatenate the last time-step's log-probabs and value predictions to the
-    # actual log-probabs and value predictions and use those going forward.
-    # log_probabs_traj = np.concatenate(
-    #     (actual_log_probabs_traj, log_probabs_traj[:, -1:, :]), axis=1)
-    # value_predictions_traj = np.concatenate(
-    #     (actual_value_predictions_traj, value_predictions_traj[:, -1:, :]),
-    #     axis=1)
-
-    log_prob_recompute_time = ppo.get_time(log_prob_recompute_start_time)
-
-    # Compute value and ppo losses.
-    self._rng, key1 = jax_random.split(self._rng, num=2)
-    logging.vlog(2, "Starting to compute P&V loss.")
-    loss_compute_start_time = time.time()
-    (cur_combined_loss, component_losses, summaries, self._model_state) = (
-        ppo.combined_loss(
-            self._policy_and_value_net_params,
-            log_probabs_traj,
-            value_predictions_traj,
-            self._policy_and_value_net_apply,
-            padded_observations,
-            padded_actions,
-            self._rewards_to_actions,
-            padded_rewards,
-            reward_mask,
-            gamma=self._gamma,
-            lambda_=self._lambda_,
-            c1=self._c1,
-            c2=self._c2,
-            state=self._model_state,
-            rng=key1))
-    loss_compute_time = ppo.get_time(loss_compute_start_time)
-    (cur_ppo_loss, cur_value_loss, cur_entropy_bonus) = component_losses
-    logging.vlog(
-        1,
-        "Calculating P&V loss [%10.2f(%10.2f, %10.2f, %10.2f)] took %0.2f msec.",
-        cur_combined_loss, cur_ppo_loss, cur_value_loss, cur_entropy_bonus,
-        ppo.get_time(loss_compute_start_time))
-
-    self._rng, key1 = jax_random.split(self._rng, num=2)
-    logging.vlog(1, "Policy and Value Optimization")
-    optimization_start_time = time.time()
-    keys = jax_random.split(key1, num=self._n_optimizer_steps)
-    opt_step = 0
-    opt_batch_size = min(self._optimizer_batch_size, B)
-    index_batches = ppo.shuffled_index_batches(
-        dataset_size=B, batch_size=opt_batch_size
-    )
-    for (index_batch, key) in zip(index_batches, keys):
-      k1, k2, k3 = jax_random.split(key, num=3)
-      t = time.time()
-      # Update the optimizer state on the sampled minibatch.
-      self._policy_and_value_opt_state, self._model_state = (
-          ppo.policy_and_value_opt_step(
-              # We pass the optimizer slots between PPO epochs, so we need to
-              # pass the optimization step as well, so for example the
-              # bias-correction in Adam is calculated properly. Alternatively we
-              # could reset the slots and the step in every PPO epoch, but then
-              # the moment estimates in adaptive optimizers would never have
-              # enough time to warm up. So it makes sense to reuse the slots,
-              # even though we're optimizing a different loss in every new
-              # epoch.
-              self._total_opt_step,
-              self._policy_and_value_opt_state,
-              self._policy_and_value_opt_update,
-              self._policy_and_value_get_params,
-              self._policy_and_value_net_apply,
-              log_probabs_traj[index_batch],
-              value_predictions_traj[index_batch],
-              padded_observations[index_batch],
-              padded_actions[index_batch],
-              self._rewards_to_actions,
-              padded_rewards[index_batch],
-              reward_mask[index_batch],
-              c1=self._c1,
-              c2=self._c2,
-              gamma=self._gamma,
-              lambda_=self._lambda_,
-              state=self._model_state,
-              rng=k1))
-      opt_step += 1
-      self._total_opt_step += 1
-
-      # Compute the approx KL for early stopping. Use the whole dataset - as we
-      # only do inference, it should fit in the memory.
-      (log_probab_actions_new, _) = (
-          self._policy_and_value_net_apply(
-              padded_observations,
-              params=self._policy_and_value_net_params,
-              state=self._model_state,
-              rng=k2))
-
-      action_mask = np.dot(
-          np.pad(reward_mask, ((0, 0), (0, 1))), self._rewards_to_actions
-      )
-      approx_kl = ppo.approximate_kl(log_probab_actions_new, log_probabs_traj,
-                                     action_mask)
-
-      early_stopping = approx_kl > 1.5 * self._target_kl
-      if early_stopping:
-        logging.vlog(
-            1, "Early stopping policy and value optimization after %d steps, "
-            "with approx_kl: %0.2f", opt_step, approx_kl)
-        # We don't return right-away, we want the below to execute on the last
-        # iteration.
-
-      t2 = time.time()
-      if (opt_step % self._print_every_optimizer_steps == 0 or
-          opt_step == self._n_optimizer_steps or early_stopping):
-        # Compute and log the loss.
-        (combined_loss, component_losses, _, self._model_state) = (
-            ppo.combined_loss(
-                self._policy_and_value_net_params,
-                log_probabs_traj,
-                value_predictions_traj,
-                self._policy_and_value_net_apply,
-                padded_observations,
-                padded_actions,
-                self._rewards_to_actions,
-                padded_rewards,
-                reward_mask,
-                gamma=self._gamma,
-                lambda_=self._lambda_,
-                c1=self._c1,
-                c2=self._c2,
-                state=self._model_state,
-                rng=k3))
-        logging.vlog(1, "One Policy and Value grad desc took: %0.2f msec",
-                     ppo.get_time(t, t2))
-        (ppo_loss, value_loss, entropy_bonus) = component_losses
-        logging.vlog(
-            1, "Combined Loss(value, ppo, entropy_bonus) [%10.2f] ->"
-            " [%10.2f(%10.2f,%10.2f,%10.2f)]", cur_combined_loss, combined_loss,
-            ppo_loss, value_loss, entropy_bonus)
-
-      if early_stopping:
-        break
-
-    optimization_time = ppo.get_time(optimization_start_time)
-
-    logging.vlog(
-        1, "Total Combined Loss reduction [%0.2f]%%",
-        (100 * (cur_combined_loss - combined_loss) / np.abs(cur_combined_loss)))
-
-    summaries.update({
-        "n_optimizer_steps": opt_step,
-        "approx_kl": approx_kl,
-    })
-    for (name, value) in summaries.items():
-      self._train_sw.scalar("train/{}".format(name), value, step=self._epoch)
-
-    logging.info(
-        "PPO epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
-        " Loss(ppo, value, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)]", self._epoch,
-        min_reward, max_reward, avg_reward, combined_loss, ppo_loss, value_loss,
-        entropy_bonus)
-
-    # Bump the epoch counter before saving a checkpoint, so that a call to
-    # save() after the training loop is a no-op if a checkpoint was saved last
-    # epoch - otherwise it would bump the epoch counter on the checkpoint.
-    last_epoch = self._epoch
-    self._epoch += 1
-
-    # Save parameters every time we see the end of at least a fraction of batch
-    # number of trajectories that are done (not completed -- completed includes
-    # truncated and done).
-    # Also don't save too frequently, enforce a minimum gap.
-    policy_save_start_time = time.time()
-    # TODO(afrozm): Refactor to trax.save_state.
-    if (self._n_trajectories_done >=
-        self._done_frac_for_policy_save * self.train_env.batch_size and
-        self._epoch % self._save_every_n == 0) or self._async_mode:
-      self.save()
-    policy_save_time = ppo.get_time(policy_save_start_time)
-
-    epoch_time = ppo.get_time(epoch_start_time)
-
-    timing_dict = {
-        "epoch": epoch_time,
-        "policy_eval": policy_eval_time,
-        "trajectory_collection": trajectory_collection_time,
-        "preprocessing": preprocessing_time,
-        "log_prob_recompute": log_prob_recompute_time,
-        "loss_compute": loss_compute_time,
-        "optimization": optimization_time,
-        "policy_save": policy_save_time,
-    }
-
-    timing_dict.update(timing_info)
-
-    for k, v in timing_dict.items():
-      self._timing_sw.scalar("timing/%s" % k, v, step=last_epoch)
-
-    max_key_len = max(len(k) for k in timing_dict)
-    timing_info_list = [
-        "%s : % 10.2f" % (k.rjust(max_key_len + 1), v)
-        for k, v in sorted(timing_dict.items())
-    ]
-    logging.info("PPO epoch [% 6d], Timings: \n%s", last_epoch,
-                 "\n".join(timing_info_list))
-
-    # Flush summary writers once in a while.
-    if self._epoch % 1000 == 0:
-      self.flush_summaries()
-
-  def evaluate(self):
-    """Evaluate the agent."""
-    if not self._separate_eval:
-      return
-    logging.vlog(1, "PPO epoch [% 6d]: evaluating policy.", self._epoch)
-
-    processed_reward_sums = collections.defaultdict(list)
-    raw_reward_sums = collections.defaultdict(list)
-    for _ in range(self._n_evals):
-      for temperature in self._eval_temperatures:
-        trajs, _, _, self._model_state = self.collect_trajectories(
-            train=False, temperature=temperature)
-
-        processed_reward_sums[temperature].extend(
-            sum(traj[2]) for traj in trajs)
-        raw_reward_sums[temperature].extend(sum(traj[3]) for traj in trajs)
-
-    # Return the mean and standard deviation for each temperature.
-    def compute_stats(reward_dict):
-      return {
-          temperature: {  # pylint: disable=g-complex-comprehension
-              "mean": onp.mean(rewards),
-              "std": onp.std(rewards)
-          } for (temperature, rewards) in reward_dict.items()
-      }
-
-    reward_stats = {
-        "processed": compute_stats(processed_reward_sums),
-        "raw": compute_stats(raw_reward_sums),
-    }
-
-    ppo.write_eval_reward_summaries(
-        reward_stats, self._eval_sw, epoch=self._epoch)
-
-  def save(self):
-    """Save the agent parameters."""
-    logging.vlog(1, "PPO epoch [% 6d]: saving model.", self._epoch)
-    ppo.save_opt_state(
-        self._output_dir,
-        self._policy_and_value_opt_state,
-        self._model_state,
-        self._epoch,
-        self._total_opt_step,
-    )
-    # Reset this number.
-    self._n_trajectories_done = 0
-    self._last_saved_at = self._epoch
-
-  def flush_summaries(self):
-    self._train_sw.flush()
-    self._timing_sw.flush()
-    self._eval_sw.flush()
-
-  @property
-  def _policy_and_value_net_params(self):
-    return self._policy_and_value_get_params(self._policy_and_value_opt_state)
-
-  # Prepares the trajectories for policy training.
-  def _preprocess_trajectories(self, trajectories):
-    (_, reward_mask, observations, actions, rewards, infos) = (
-        ppo.pad_trajectories(trajectories, boundary=self._max_timestep)
-    )
-    assert self.train_env.observation_space.shape == observations.shape[2:]
-    if not self._serialized_sequence_policy:
-      # Add one timestep at the end, so it's compatible with
-      # self._rewards_to_actions.
-      pad_width = ((0, 0), (0, 1)) + ((0, 0),) * (actions.ndim - 2)
-      actions = np.pad(actions, pad_width)
-      actions = np.reshape(actions, (actions.shape[0], -1))
-    else:
-      (observations, actions) = self._serialize_trajectories(
-          observations, actions, reward_mask
-      )
-    return (observations, actions, rewards, reward_mask, infos)
-
-  def _serialize_trajectories(self, observations, actions, reward_mask):
-    (reprs, _) = serialization_utils.serialize_observations_and_actions(
-        observations=observations,
-        actions=actions,
-        mask=reward_mask,
-        **self._serialization_kwargs
-    )
-    # Mask out actions in the representation - otherwise we sample an action
-    # based on itself.
-    observations = reprs * serialization_utils.observation_mask(
-        **self._serialization_kwargs
-    )
-    actions = reprs
-    return (observations, actions)
-
-  # A function to get the policy and value predictions.
-  def _get_predictions(self, observations, state, rng=None):
-    """Returns log-probs, value predictions and key back."""
-    key, key1 = jax_random.split(rng, num=2)
-
-    (log_probs, value_preds) = self._policy_and_value_net_apply(
-        observations, params=self._policy_and_value_net_params, state=state,
-        rng=key1)
-
-    return log_probs, value_preds, state, key
-
-  def _policy_fun(self, observations, lengths, state, rng):
-    (batch_size, n_timesteps) = observations.shape[:2]
-    if self._serialized_sequence_policy:
-      actions = np.zeros(
-          (batch_size, n_timesteps - 1) + self.train_env.action_space.shape,
-          dtype=self.train_env.action_space.dtype,
-      )
-      reward_mask = np.ones((batch_size, n_timesteps - 1), dtype=np.int32)
-      (observations, _) = self._serialize_trajectories(
-          observations, actions, reward_mask
-      )
-    (log_probs, value_preds, state, rng) = self._get_predictions(
-        observations, state=state, rng=rng
-    )
-    # We need the log_probs of those actions that correspond to the last actual
-    # time-step.
-    index = lengths - 1  # Since we want to index using lengths.
-    pred_index = self._calc_action_index(index)
-    log_probs = log_probs[
-        np.arange(batch_size)[:, None, None],
-        pred_index[:, :, None],
-        np.arange(self._n_actions),
-    ]
-    value_preds = value_preds[np.arange(batch_size)[:, None], pred_index]
-    return (log_probs, value_preds, state, rng)
-
-  def _calc_action_index(self, reward_index):
-    # Project the one-hot position in the reward sequence onto the action
-    # sequence to figure out which actions correspond to that position.
-    one_hot_index = np.eye(self._rewards_to_actions.shape[0])[reward_index]
-    action_mask = np.dot(one_hot_index, self._rewards_to_actions)
-    # Compute the number of symbols in an action. It's just the number of 1s in
-    # the mask.
-    action_length = int(np.sum(action_mask[0]))
-    # Argmax stops on the first occurrence, so we use it to find the first 1 in
-    # the mask.
-    action_start_index = np.argmax(action_mask, axis=1)
-    return action_start_index[:, None] + np.arange(action_length)[None, :]
diff --git a/tensor2tensor/trax/rl/ppo_trainer_test.py b/tensor2tensor/trax/rl/ppo_trainer_test.py
deleted file mode 100644
index 134f5207d..000000000
--- a/tensor2tensor/trax/rl/ppo_trainer_test.py
+++ /dev/null
@@ -1,306 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.rl.ppo's training_loop."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import contextlib
-import functools
-import itertools
-import os
-import tempfile
-
-import gin
-import gym
-import numpy as np
-
-from tensor2tensor.envs import gym_env_problem
-from tensor2tensor.rl import gym_utils
-from tensor2tensor.trax import inputs as trax_inputs
-from tensor2tensor.trax import layers
-from tensor2tensor.trax import learning_rate as lr
-from tensor2tensor.trax import models
-from tensor2tensor.trax import optimizers as trax_opt
-from tensor2tensor.trax import trax
-from tensor2tensor.trax.rl import envs  # pylint: disable=unused-import
-from tensor2tensor.trax.rl import ppo_trainer
-from tensor2tensor.trax.rl import simulated_env_problem
-from tensorflow import test
-from tensorflow.io import gfile
-
-
-class PpoTrainerTest(test.TestCase):
-
-  def get_wrapped_env(
-      self, name="CartPole-v0", max_episode_steps=2, batch_size=1
-  ):
-    wrapper_fn = functools.partial(
-        gym_utils.gym_env_wrapper,
-        **{
-            "rl_env_max_episode_steps": max_episode_steps,
-            "maxskip_env": False,
-            "rendered_env": False,
-            "rendered_env_resize_to": None,  # Do not resize frames
-            "sticky_actions": False,
-            "output_dtype": None,
-        })
-
-    return gym_env_problem.GymEnvProblem(base_env_name=name,
-                                         batch_size=batch_size,
-                                         env_wrapper_fn=wrapper_fn,
-                                         discrete_rewards=False)
-
-  @contextlib.contextmanager
-  def tmp_dir(self):
-    tmp = tempfile.mkdtemp(dir=self.get_temp_dir())
-    yield tmp
-    gfile.rmtree(tmp)
-
-  def _make_trainer(
-      self, train_env, eval_env, output_dir, model=None, **kwargs
-  ):
-    if model is None:
-      model = lambda: layers.Serial(layers.Dense(1))
-    return ppo_trainer.PPO(
-        train_env=train_env,
-        eval_env=eval_env,
-        policy_and_value_model=model,
-        n_optimizer_steps=1,
-        output_dir=output_dir,
-        random_seed=0,
-        max_timestep=3,
-        boundary=2,
-        save_every_n=1,
-        **kwargs
-    )
-
-  def test_training_loop_cartpole(self):
-    with self.tmp_dir() as output_dir:
-      trainer = self._make_trainer(
-          train_env=self.get_wrapped_env("CartPole-v0", 2),
-          eval_env=self.get_wrapped_env("CartPole-v0", 2),
-          output_dir=output_dir,
-      )
-      trainer.training_loop(n_epochs=2)
-
-  def test_training_loop_cartpole_transformer(self):
-    with self.tmp_dir() as output_dir:
-      trainer = self._make_trainer(
-          train_env=self.get_wrapped_env("CartPole-v0", 2),
-          eval_env=self.get_wrapped_env("CartPole-v0", 2),
-          output_dir=output_dir,
-          model=functools.partial(
-              models.TransformerDecoder,
-              d_model=1,
-              d_ff=1,
-              n_layers=1,
-              n_heads=1,
-              max_len=128,
-              mode="train",
-          ),
-      )
-      trainer.training_loop(n_epochs=2)
-
-  def test_training_loop_onlinetune(self):
-    with self.tmp_dir() as output_dir:
-      gin.bind_parameter("OnlineTuneEnv.model", functools.partial(
-          models.MLP,
-          n_hidden_layers=0,
-          n_output_classes=1,
-      ))
-      gin.bind_parameter("OnlineTuneEnv.inputs", functools.partial(
-          trax_inputs.random_inputs,
-          input_shape=(1, 1),
-          input_dtype=np.float32,
-          output_shape=(1, 1),
-          output_dtype=np.float32,
-      ))
-      gin.bind_parameter("OnlineTuneEnv.train_steps", 1)
-      gin.bind_parameter("OnlineTuneEnv.eval_steps", 1)
-      gin.bind_parameter(
-          "OnlineTuneEnv.output_dir", os.path.join(output_dir, "envs"))
-      trainer = self._make_trainer(
-          train_env=self.get_wrapped_env("OnlineTuneEnv-v0", 1),
-          eval_env=self.get_wrapped_env("OnlineTuneEnv-v0", 1),
-          output_dir=output_dir,
-      )
-      trainer.training_loop(n_epochs=1)
-
-  def test_training_loop_simulated(self):
-    n_actions = 5
-    history_shape = (3, 2, 3)
-    action_shape = (3,)
-    obs_shape = (3, 3)
-    reward_shape = (3, 1)
-
-    def model(mode):
-      del mode
-      return layers.Serial(
-          layers.Parallel(
-              layers.Flatten(),  # Observation stack.
-              layers.Embedding(d_feature=1, vocab_size=n_actions),  # Action.
-          ),
-          layers.Concatenate(),
-          layers.Dense(n_units=1),
-          layers.Dup(),
-          layers.Parallel(
-              layers.Dense(n_units=obs_shape[1]),  # New observation.
-              None,  # Reward.
-          )
-      )
-
-    def inputs(n_devices):
-      del n_devices
-      stream = itertools.repeat(
-          (np.zeros(history_shape), np.zeros(action_shape, dtype=np.int32),
-           np.zeros(obs_shape), np.zeros(reward_shape))
-      )
-      return trax_inputs.Inputs(
-          train_stream=lambda: stream,
-          train_eval_stream=lambda: stream,
-          eval_stream=lambda: stream,
-          input_shape=(history_shape[1:], action_shape[1:]),
-          input_dtype=(np.float32, np.int32),
-          target_shape=(obs_shape[1:], reward_shape[1:]),
-          target_dtype=(np.float32, np.float32),
-      )
-
-    def loss(mask_id=None, has_weights=False):
-      """Cross-entropy loss as scalar compatible with Trax masking."""
-      return layers.Serial(
-          # Swap from (pred-obs, pred-reward, target-obs, target-reward)
-          # to (pred-obs, target-obs, pred-reward, target-reward).
-          layers.Parallel([], layers.Swap()),
-          # Cross-entropy loss for obs, L2 loss on reward.
-          layers.Parallel(layers.CrossEntropyLossScalar(mask_id, has_weights),
-                          layers.L2LossScalar(mask_id, has_weights)),
-          # Add both losses.
-          layers.Add(),
-          # Zero out in this test.
-          layers.MulConstant(constant=0.0)
-      )
-
-    with self.tmp_dir() as output_dir:
-      # Run fake training just to save the parameters.
-      trainer = trax.Trainer(
-          model=model,
-          loss_fn=loss,
-          inputs=inputs,
-          optimizer=trax_opt.SM3,
-          lr_schedule=lr.MultifactorSchedule,
-          output_dir=output_dir,
-      )
-      trainer.train_epoch(epoch_steps=1, eval_steps=1)
-
-      # Repeat the history over and over again.
-      stream = itertools.repeat(np.zeros(history_shape))
-      env_fn = functools.partial(
-          simulated_env_problem.RawSimulatedEnvProblem,
-          model=model,
-          history_length=history_shape[1],
-          trajectory_length=3,
-          batch_size=history_shape[0],
-          observation_space=gym.spaces.Box(
-              low=-np.inf, high=np.inf, shape=(obs_shape[1],)),
-          action_space=gym.spaces.Discrete(n=n_actions),
-          reward_range=(-1, 1),
-          discrete_rewards=False,
-          history_stream=stream,
-          output_dir=output_dir,
-      )
-
-      trainer = self._make_trainer(
-          train_env=env_fn(),
-          eval_env=env_fn(),
-          output_dir=output_dir,
-      )
-      trainer.training_loop(n_epochs=2)
-
-  def test_restarts(self):
-    with self.tmp_dir() as output_dir:
-      train_env = self.get_wrapped_env("CartPole-v0", 2)
-      eval_env = self.get_wrapped_env("CartPole-v0", 2)
-
-      # Train for 1 epoch and save.
-      trainer = self._make_trainer(
-          train_env=train_env,
-          eval_env=eval_env,
-          output_dir=output_dir,
-      )
-      self.assertEqual(trainer.epoch, 0)
-      trainer.training_loop(n_epochs=1)
-      self.assertEqual(trainer.epoch, 1)
-
-      # Restore from the saved state.
-      trainer = self._make_trainer(
-          train_env=train_env,
-          eval_env=eval_env,
-          output_dir=output_dir,
-      )
-      self.assertEqual(trainer.epoch, 1)
-      # Check that we can continue training from the restored checkpoint.
-      trainer.training_loop(n_epochs=2)
-      self.assertEqual(trainer.epoch, 2)
-
-  def test_training_loop_multi_control(self):
-    gym.register(
-        "FakeEnv-v0",
-        entry_point="tensor2tensor.trax.rl.envs.fake_env:FakeEnv",
-        kwargs={"n_actions": 3, "n_controls": 2},
-    )
-    with self.tmp_dir() as output_dir:
-      trainer = self._make_trainer(
-          train_env=self.get_wrapped_env("FakeEnv-v0", 2),
-          eval_env=self.get_wrapped_env("FakeEnv-v0", 2),
-          output_dir=output_dir,
-      )
-      trainer.training_loop(n_epochs=2)
-
-  def test_training_loop_cartpole_serialized(self):
-    gin.bind_parameter("BoxSpaceSerializer.precision", 1)
-    with self.tmp_dir() as output_dir:
-      trainer = self._make_trainer(
-          train_env=self.get_wrapped_env("CartPole-v0", 2),
-          eval_env=self.get_wrapped_env("CartPole-v0", 2),
-          output_dir=output_dir,
-          model=functools.partial(
-              models.TransformerDecoder,
-              d_model=1,
-              d_ff=1,
-              n_layers=1,
-              n_heads=1,
-              max_len=1024,
-              mode="train",
-          ),
-          policy_and_value_vocab_size=4,
-      )
-      trainer.training_loop(n_epochs=2)
-
-  def test_training_loop_cartpole_minibatch(self):
-    with self.tmp_dir() as output_dir:
-      trainer = self._make_trainer(
-          train_env=self.get_wrapped_env("CartPole-v0", 2, batch_size=4),
-          eval_env=self.get_wrapped_env("CartPole-v0", 2),
-          output_dir=output_dir,
-          optimizer_batch_size=2,
-      )
-      trainer.training_loop(n_epochs=2)
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/rl/serialization_utils.py b/tensor2tensor/trax/rl/serialization_utils.py
deleted file mode 100644
index 7bd5a2df2..000000000
--- a/tensor2tensor/trax/rl/serialization_utils.py
+++ /dev/null
@@ -1,184 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Utilities for serializing trajectories into discrete sequences."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-
-
-def serialize_observations_and_actions(
-    observations,
-    actions,
-    mask,
-    observation_serializer,
-    action_serializer,
-    representation_length,
-):
-  """Serializes observations and actions into a discrete sequence.
-
-  Args:
-    observations: Array (B, T + 1, ...), of observations, where B is the batch
-      size and T is the number of timesteps excluding the last observation.
-    actions: Array (B, T, ...) of actions.
-    mask: Binary array (B, T) indicating where each sequence ends (1s while
-      it continues).
-    observation_serializer: SpaceSerializer for observations.
-    action_serializer: SpaceSerializer for actions.
-    representation_length: Number of symbols in the serialized sequence. The
-      sequence is padded up to this number.
-  Returns:
-    Pair (representation, mask), where representation is the serialized sequence
-    of shape (B, R) where R = representation_length, and mask is a binary array
-    of shape (B, R) indicating where each sequence ends.
-  """
-  (batch_size, n_timesteps) = actions.shape[:2]
-  assert observations.shape[:2] == (batch_size, n_timesteps + 1)
-  assert mask.shape == (batch_size, n_timesteps)
-
-  reprs = []
-  for t in range(n_timesteps):
-    reprs.append(observation_serializer.serialize(observations[:, t, ...]))
-    reprs.append(action_serializer.serialize(actions[:, t, ...]))
-  reprs.append(observation_serializer.serialize(observations[:, -1, ...]))
-  reprs = np.concatenate(reprs, axis=1)
-  assert reprs.shape[1] <= representation_length
-  reprs = np.pad(
-      reprs,
-      pad_width=((0, 0), (0, representation_length - reprs.shape[1])),
-      mode="constant",
-  )
-
-  obs_repr_length = observation_serializer.representation_length
-  act_repr_length = action_serializer.representation_length
-  step_repr_length = obs_repr_length + act_repr_length
-  seq_lengths = np.sum(mask, axis=1).astype(np.int32)
-  repr_lengths = seq_lengths * step_repr_length + obs_repr_length
-  repr_mask = np.zeros((batch_size, representation_length), dtype=np.int32)
-  for (i, repr_length) in enumerate(repr_lengths):
-    repr_mask[i, :repr_length] = 1
-
-  return (reprs, repr_mask)
-
-
-def observation_mask(
-    observation_serializer, action_serializer, representation_length
-):
-  """Calculates an observation mask for a serialized sequence.
-
-  Args:
-    observation_serializer: SpaceSerializer for observations.
-    action_serializer: SpaceSerializer for actions.
-    representation_length: Number of symbols in the serialized sequence. The
-      mask is padded up to this number.
-
-  Returns:
-    Binary mask indicating which symbols in the representation correspond to
-    observations.
-  """
-  mask = np.zeros(representation_length, dtype=np.int32)
-  obs_repr_length = observation_serializer.representation_length
-  step_repr_length = obs_repr_length + action_serializer.representation_length
-  for step_start_index in range(0, representation_length, step_repr_length):
-    mask[step_start_index:(step_start_index + obs_repr_length)] = 1
-  return mask
-
-
-def action_mask(
-    observation_serializer, action_serializer, representation_length
-):
-  """Calculates an action mask for a serialized sequence.
-
-  Args:
-    observation_serializer: SpaceSerializer for observations.
-    action_serializer: SpaceSerializer for actions.
-    representation_length: Number of symbols in the serialized sequence. The
-      mask is padded up to this number.
-
-  Returns:
-    Binary mask indicating which symbols in the representation correspond to
-    actions.
-  """
-  return 1 - observation_mask(
-      observation_serializer, action_serializer, representation_length
-  )
-
-
-def significance_map(
-    observation_serializer, action_serializer, representation_length
-):
-  """Calculates a significance map for the entire serialized sequence.
-
-  See SpaceSerializer.significance_map.
-
-  Args:
-    observation_serializer: SpaceSerializer for observations.
-    action_serializer: SpaceSerializer for actions.
-    representation_length: Number of symbols in the serialized sequence. The
-      significance map is padded up to this number.
-
-  Returns:
-    Significance map for the entire serialized sequence.
-  """
-  sig_map = np.zeros(representation_length, dtype=np.int32)
-  obs_repr_length = observation_serializer.representation_length
-  act_repr_length = action_serializer.representation_length
-  step_repr_length = obs_repr_length + act_repr_length
-  for step_start_index in range(0, representation_length, step_repr_length):
-    act_start_index = step_start_index + obs_repr_length
-    step_end_index = step_start_index + step_repr_length
-    limit = representation_length - step_start_index
-    sig_map[step_start_index:act_start_index] = (
-        observation_serializer.significance_map[:limit]
-    )
-    limit = representation_length - act_start_index
-    sig_map[act_start_index:step_end_index] = (
-        action_serializer.significance_map[:limit]
-    )
-  return sig_map
-
-
-def rewards_to_actions_map(
-    observation_serializer,
-    action_serializer,
-    n_timesteps,
-    representation_length,
-):
-  """Calculates a mapping between the rewards and the serialized sequence.
-
-  Used to broadcast advantages over the log-probabilities of corresponding
-  actions.
-
-  Args:
-    observation_serializer: SpaceSerializer for observations.
-    action_serializer: SpaceSerializer for actions.
-    n_timesteps: Number of timesteps (length of the reward sequence).
-    representation_length: Number of symbols in the serialized sequence.
-
-  Returns:
-    Array (T, R) translating from the reward sequence to actions in the
-    representation.
-  """
-  r2a_map = np.zeros((n_timesteps, representation_length))
-  obs_repr_length = observation_serializer.representation_length
-  act_repr_length = action_serializer.representation_length
-  step_repr_length = obs_repr_length + act_repr_length
-  for t in range(n_timesteps):
-    act_start_index = t * step_repr_length + obs_repr_length
-    r2a_map[t, act_start_index:(act_start_index + act_repr_length)] = 1
-  return r2a_map
diff --git a/tensor2tensor/trax/rl/serialization_utils_test.py b/tensor2tensor/trax/rl/serialization_utils_test.py
deleted file mode 100644
index 46d8985e0..000000000
--- a/tensor2tensor/trax/rl/serialization_utils_test.py
+++ /dev/null
@@ -1,169 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.rl.serialization_utils."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gin
-import gym
-import numpy as np
-
-from tensor2tensor.trax.rl import serialization_utils
-from tensor2tensor.trax.rl import space_serializer
-from tensorflow import test
-
-
-class SerializationTest(test.TestCase):
-
-  def setUp(self):
-    super(SerializationTest, self).setUp()
-    self._serializer = space_serializer.create(
-        gym.spaces.Discrete(2), vocab_size=2
-    )
-    self._repr_length = 100
-    self._serialization_utils_kwargs = {
-        "observation_serializer": self._serializer,
-        "action_serializer": self._serializer,
-        "representation_length": self._repr_length,
-    }
-
-  def test_serializes_observations_and_actions(self):
-    (reprs, mask) = serialization_utils.serialize_observations_and_actions(
-        observations=np.array([[0, 1]]),
-        actions=np.array([[0]]),
-        mask=np.array([[1]]),
-        **self._serialization_utils_kwargs
-    )
-    self.assertEqual(reprs.shape, (1, self._repr_length))
-    self.assertEqual(mask.shape, (1, self._repr_length))
-    self.assertGreater(np.sum(mask), 0)
-    self.assertEqual(np.max(mask), 1)
-
-  def test_masks_length(self):
-    (reprs, mask) = serialization_utils.serialize_observations_and_actions(
-        observations=np.array([[0, 1, 0], [0, 1, 0], [0, 1, 1]]),
-        actions=np.array([[0, 0], [0, 1], [0, 0]]),
-        mask=np.array([[1, 0], [1, 1], [1, 1]]),
-        **self._serialization_utils_kwargs
-    )
-    # Trajectories 1 and 2 are longer than 0.
-    self.assertGreater(np.sum(mask[1]), np.sum(mask[0]))
-    self.assertGreater(np.sum(mask[2]), np.sum(mask[0]))
-    # Trajectory 0 is a common prefix of 1 and 2. 1 and 2 are different.
-    np.testing.assert_array_equal(reprs[0] * mask[0], reprs[1] * mask[0])
-    np.testing.assert_array_equal(reprs[0] * mask[0], reprs[2] * mask[0])
-    self.assertFalse(np.array_equal(reprs[1] * mask[1], reprs[2] * mask[2]))
-    # Trajectories should be padded with 0s.
-    np.testing.assert_array_equal(
-        reprs * (1 - mask), np.zeros((3, self._repr_length))
-    )
-
-  def test_observation_and_action_masks_are_valid_and_complementary(self):
-    obs_mask = serialization_utils.observation_mask(
-        **self._serialization_utils_kwargs
-    )
-    self.assertEqual(obs_mask.shape, (self._repr_length,))
-    self.assertEqual(np.min(obs_mask), 0)
-    self.assertEqual(np.max(obs_mask), 1)
-
-    act_mask = serialization_utils.action_mask(
-        **self._serialization_utils_kwargs
-    )
-    self.assertEqual(act_mask.shape, (self._repr_length,))
-    self.assertEqual(np.min(act_mask), 0)
-    self.assertEqual(np.max(act_mask), 1)
-
-    np.testing.assert_array_equal(
-        obs_mask + act_mask, np.ones(self._repr_length)
-    )
-
-  def test_masks_observations(self):
-    (reprs, _) = serialization_utils.serialize_observations_and_actions(
-        # Observations are different, actions are the same.
-        observations=np.array([[0, 1], [1, 1]]),
-        actions=np.array([[0], [0]]),
-        mask=np.array([[1], [1]]),
-        **self._serialization_utils_kwargs
-    )
-    obs_mask = serialization_utils.observation_mask(
-        **self._serialization_utils_kwargs
-    )
-    act_mask = serialization_utils.action_mask(
-        **self._serialization_utils_kwargs
-    )
-
-    self.assertFalse(np.array_equal(reprs[0] * obs_mask, reprs[1] * obs_mask))
-    np.testing.assert_array_equal(reprs[0] * act_mask, reprs[1] * act_mask)
-
-  def test_masks_actions(self):
-    (reprs, _) = serialization_utils.serialize_observations_and_actions(
-        # Observations are the same, actions are different.
-        observations=np.array([[0, 1], [0, 1]]),
-        actions=np.array([[0], [1]]),
-        mask=np.array([[1], [1]]),
-        **self._serialization_utils_kwargs
-    )
-    obs_mask = serialization_utils.observation_mask(
-        **self._serialization_utils_kwargs
-    )
-    act_mask = serialization_utils.action_mask(
-        **self._serialization_utils_kwargs
-    )
-
-    np.testing.assert_array_equal(reprs[0] * obs_mask, reprs[1] * obs_mask)
-    self.assertFalse(np.array_equal(reprs[0] * act_mask, reprs[1] * act_mask))
-
-  def test_significance_map(self):
-    gin.bind_parameter("BoxSpaceSerializer.precision", 3)
-    significance_map = serialization_utils.significance_map(
-        observation_serializer=space_serializer.create(
-            gym.spaces.Box(low=0, high=1, shape=(2,)), vocab_size=2
-        ),
-        action_serializer=space_serializer.create(
-            gym.spaces.MultiDiscrete(nvec=[2, 2]), vocab_size=2
-        ),
-        representation_length=20,
-    )
-    np.testing.assert_array_equal(
-        significance_map,
-        # obs1, act1, obs2, act2, obs3 cut after 4th symbol.
-        [0, 1, 2, 0, 1, 2, 0, 0, 0, 1, 2, 0, 1, 2, 0, 0, 0, 1, 2, 0],
-    )
-
-  def test_rewards_to_actions_map(self):
-    rewards = np.array([1, 2, 3])
-    r2a_map = serialization_utils.rewards_to_actions_map(
-        observation_serializer=space_serializer.create(
-            gym.spaces.MultiDiscrete(nvec=[2, 2, 2]), vocab_size=2
-        ),
-        action_serializer=space_serializer.create(
-            gym.spaces.MultiDiscrete(nvec=[2, 2]), vocab_size=2
-        ),
-        n_timesteps=len(rewards),
-        representation_length=16,
-    )
-    broadcast_rewards = np.dot(rewards, r2a_map)
-    np.testing.assert_array_equal(
-        broadcast_rewards,
-        # obs1, act1, obs2, act2, obs3 cut after 1st symbol.
-        [0, 0, 0, 1, 1, 0, 0, 0, 2, 2, 0, 0, 0, 3, 3, 0],
-    )
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/rl/simple.py b/tensor2tensor/trax/rl/simple.py
deleted file mode 100644
index 8c3868160..000000000
--- a/tensor2tensor/trax/rl/simple.py
+++ /dev/null
@@ -1,236 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 helper functions."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import itertools
-import os
-import random
-
-from absl import logging
-import numpy as np
-from tensor2tensor.envs import env_problem_utils
-from tensor2tensor.envs import trajectory
-from tensor2tensor.trax import utils
-from tensorflow.io import gfile
-
-
-def load_trajectories(trajectory_dir, eval_frac):
-  """Loads trajectories from a possibly nested directory of pickles."""
-  pkl_module = utils.get_pickle_module()
-  train_trajectories = []
-  eval_trajectories = []
-  # Search the entire directory subtree for trajectories.
-  for (subdir, _, filenames) in gfile.walk(trajectory_dir):
-    for filename in filenames:
-      shard_path = os.path.join(subdir, filename)
-      try:
-        with gfile.GFile(shard_path, "rb") as f:
-          trajectories = pkl_module.load(f)
-        pivot = int(len(trajectories) * (1 - eval_frac))
-        train_trajectories.extend(trajectories[:pivot])
-        eval_trajectories.extend(trajectories[pivot:])
-      except EOFError:
-        logging.warning(
-            "Could not load trajectories from a corrupted shard %s.",
-            shard_path,
-        )
-  assert train_trajectories, "Can't find training data in %s" % trajectory_dir
-  assert eval_trajectories, "Can't find evaluation data in %s" % trajectory_dir
-  return train_trajectories, eval_trajectories
-
-
-def generate_examples(trajectories, trajectory_to_training_examples_fn):
-  """Generates an infinite stream of shuffled examples out of trajectories."""
-  examples = [
-      example  # pylint: disable=g-complex-comprehension
-      for trajectory_examples in map(
-          trajectory_to_training_examples_fn, trajectories)
-      for example in trajectory_examples
-  ]
-  assert examples
-  while True:
-    random.shuffle(examples)
-    for example in examples:
-      yield example
-
-
-def mix_streams(stream1, stream2, mix_prob):
-  """Mixes two streams together with a fixed probability."""
-  while True:
-    # In the corner cases (mix_prob = 0 or 1) mixing the other stream never
-    # happens, because random() samples from the semi-open interval [0, 1).
-    if random.random() < mix_prob:
-      yield next(stream1)
-    else:
-      yield next(stream2)
-
-
-def batch_stream(stream, batch_size):
-  """Batches a stream of training examples."""
-  def make_batch(examples):
-    """Stacks a structure of numpy arrays nested in lists/tuples."""
-    assert examples
-    if isinstance(examples[0], (list, tuple)):
-      return type(examples[0])(
-          make_batch([example[i] for example in examples])
-          for i in range(len(examples[0]))
-      )
-    else:
-      return np.stack(examples, axis=0)
-
-  # Take consecutive batches from an infinite stream. This way there are no
-  # incomplete batches. We might get duplicate examples in the same batch, but
-  # that should be very rare.
-  while True:
-    yield make_batch(list(itertools.islice(stream, batch_size)))
-
-
-# TODO(pkozakowski): This is mostly a simplified version of
-# env_problem_utils.play_env_problem_with_policy, generalized to work with
-# policies not being neural networks. Another difference is that it always
-# collects exactly one trajectory from each environment in the batch. Unify if
-# possible.
-def play_env_problem(env, policy_fn):
-  """Plays an EnvProblem using a given policy function."""
-  trajectories = [trajectory.Trajectory() for _ in range(env.batch_size)]
-  observations = env.reset()
-  for (traj, observation) in zip(trajectories, observations):
-    traj.add_time_step(observation=observation)
-
-  done_so_far = np.array([False] * env.batch_size)
-  while not np.all(done_so_far):
-    padded_observations, _ = env.trajectories.observations_np(
-        len_history_for_policy=None)
-    actions = policy_fn(padded_observations)
-    (observations, rewards, dones, _) = env.step(actions)
-    for (traj, observation, action, reward, done) in zip(
-        trajectories, observations, actions, rewards, dones
-    ):
-      if not traj.done:
-        traj.change_last_time_step(action=action)
-        traj.add_time_step(
-            observation=observation, raw_reward=reward, done=done)
-      env.reset(indices=env_problem_utils.done_indices(dones))
-    done_so_far = np.logical_or(done_so_far, dones)
-  return trajectories
-
-
-def calculate_observation_error(real_trajectories, sim_trajectories):
-  """Calculates MSE of observations in two trajectories."""
-  def pad_or_truncate(observations, desired_length):
-    (current_length, _) = observations.shape
-    if current_length < desired_length:
-      return np.pad(
-          observations,
-          pad_width=((0, desired_length - current_length), (0, 0)),
-          mode="edge",
-      )
-    else:
-      return observations[:desired_length, :]
-
-  def calculate_for_single_pair(real_trajectory, sim_trajectory):
-    real_obs = real_trajectory.observations_np
-    sim_obs = pad_or_truncate(
-        sim_trajectory.observations_np, real_trajectory.num_time_steps)
-    return np.sum((real_obs - sim_obs) ** 2, axis=0)
-
-  return np.mean([
-      calculate_for_single_pair(real_traj, sim_traj)
-      for (real_traj, sim_traj) in zip(real_trajectories, sim_trajectories)
-  ], axis=0)
-
-
-def plot_observation_error(real_trajectories, sim_trajectories, mpl_plt):
-  """Plots observations from two trajectories on the same graph."""
-  assert len(real_trajectories) == len(sim_trajectories)
-  assert real_trajectories
-  obs_dim = real_trajectories[0].last_time_step.observation.shape[0]
-  (w, h) = mpl_plt.rcParams["figure.figsize"]
-  ncols = len(real_trajectories)
-  nrows = obs_dim
-  (_, axes) = mpl_plt.subplots(
-      nrows, ncols, figsize=(w * ncols, h * nrows))
-  for (traj_index, (real_traj, sim_traj)) in enumerate(
-      zip(real_trajectories, sim_trajectories)
-  ):
-    for dim_index in range(obs_dim):
-      for (traj, label) in ((real_traj, "real"), (sim_traj, "simulated")):
-        obs = traj.observations_np
-        ax = axes[dim_index, traj_index]
-        ax.set_title("trajectory {}, observation dimension {}".format(
-            traj_index, dim_index))
-        ax.plot(np.arange(obs.shape[0]), obs[:, dim_index], label=label)
-        ax.legend()
-
-
-class ReplayPolicy(object):
-  """Policy function repeating actions from a given batch of trajectories."""
-
-  def __init__(self, trajectories, out_of_bounds_action):
-    """Creates ReplayPolicy.
-
-    Args:
-      trajectories: Batch of trajectories to repeat actions from.
-      out_of_bounds_action: Action to play after the replayed trajectory ends.
-    """
-    self._trajectories = trajectories
-    self._out_of_bounds_action = out_of_bounds_action
-    self._step = 0
-
-  def __call__(self, observations):
-    del observations
-
-    def get_action(traj):
-      action = None
-      if self._step < traj.num_time_steps:
-        action = traj.time_steps[self._step].action
-        # PS: action can still be None, if this is the last time-step in traj.
-      return action if action is not None else self._out_of_bounds_action
-    actions = np.array(list(map(get_action, self._trajectories)))
-    self._step += 1
-    return actions
-
-
-def evaluate_model(sim_env, real_trajectories, mpl_plt, n_to_plot=3):
-  """Reports the observation error metric and the corresponding plot."""
-  if len(sim_env.observation_space.shape) != 1:
-    logging.warning(
-        "Could not evaluate the model - only environments with vector "
-        "observation spaces are supported."
-    )
-    return
-
-  assert len(real_trajectories) == sim_env.batch_size
-
-  policy_fn = ReplayPolicy(
-      real_trajectories,
-      # Does not matter which action we play after the real trajetory ends, we
-      # cut the simulated one to match the real one anyway.
-      out_of_bounds_action=sim_env.action_space.sample(),
-  )
-
-  sim_trajectories = play_env_problem(sim_env, policy_fn)
-  obs_errors = calculate_observation_error(real_trajectories, sim_trajectories)
-  plot_observation_error(
-      real_trajectories[:n_to_plot], sim_trajectories[:n_to_plot], mpl_plt)
-  return {
-      "observation_error/{}".format(i): obs_error
-      for (i, obs_error) in enumerate(obs_errors)
-  }
diff --git a/tensor2tensor/trax/rl/simple_test.py b/tensor2tensor/trax/rl/simple_test.py
deleted file mode 100644
index 799f4654c..000000000
--- a/tensor2tensor/trax/rl/simple_test.py
+++ /dev/null
@@ -1,304 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.rl.simple."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import itertools
-import os
-
-import gin
-import gym
-from matplotlib import pyplot as plt
-import mock
-import numpy as np
-
-from tensor2tensor.envs import trajectory
-from tensor2tensor.trax import backend
-from tensor2tensor.trax import trax
-from tensor2tensor.trax import utils
-from tensor2tensor.trax.rl import simple
-from tensor2tensor.trax.rl import simulated_env_problem
-from tensor2tensor.trax.rl import space_serializer  # pylint: disable=unused-import
-from tensorflow import test
-from tensorflow.io import gfile
-
-
-class SimpleTest(test.TestCase):
-
-  def _make_singleton_trajectory(self, observation):
-    t = trajectory.Trajectory()
-    t.add_time_step(observation=observation)
-    return t
-
-  def _dump_trajectory_pickle(self, observations, path):
-    pkl_module = utils.get_pickle_module()
-    trajectories = list(map(self._make_singleton_trajectory, observations))
-    with gfile.GFile(path, "wb") as f:
-      pkl_module.dump(trajectories, f)
-
-  def test_loads_trajectories(self):
-    temp_dir = self.get_temp_dir()
-    # Dump two trajectory pickles with given observations.
-    self._dump_trajectory_pickle(
-        observations=[0, 1, 2, 3], path=os.path.join(temp_dir, "0.pkl"))
-    self._dump_trajectory_pickle(
-        observations=[4, 5, 6, 7], path=os.path.join(temp_dir, "1.pkl"))
-    (train_trajs, eval_trajs) = simple.load_trajectories(
-        temp_dir, eval_frac=0.25)
-    extract_obs = lambda t: t.last_time_step.observation
-    # The order of pickles is undefined, so we compare sets.
-    actual_train_obs = set(map(extract_obs, train_trajs))
-    actual_eval_obs = set(map(extract_obs, eval_trajs))
-
-    # First 3 trajectories from each pickle go to train, the last one to eval.
-    expected_train_obs = {0, 1, 2, 4, 5, 6}
-    expected_eval_obs = {3, 7}
-    self.assertEqual(actual_train_obs, expected_train_obs)
-    self.assertEqual(actual_eval_obs, expected_eval_obs)
-
-  def test_generates_examples(self):
-    observations = [0, 1, 2, 3]
-    trajectories = map(self._make_singleton_trajectory, observations)
-    trajectory_to_training_examples = lambda t: [t.last_time_step.observation]
-    stream = simple.generate_examples(
-        trajectories, trajectory_to_training_examples)
-
-    # The examples are shuffled, so we compare sets.
-    self.assertEqual(
-        set(itertools.islice(stream, len(observations))), set(observations))
-    # The stream is infinite, so we should be able to take a next element.
-    self.assertIn(next(stream), observations)
-
-  def test_mixes_streams_with_prob_one(self):
-    # Mix infinite streams of 0s and 1s.
-    stream = simple.mix_streams(
-        itertools.repeat(0), itertools.repeat(1), mix_prob=1.0)
-    # Mixed stream should have only 0s.
-    self.assertEqual(set(itertools.islice(stream, 100)), {0})
-
-  def test_mixes_streams_with_prob_zero(self):
-    stream = simple.mix_streams(
-        itertools.repeat(0), itertools.repeat(1), mix_prob=0.0)
-    # Mixed stream should have only 1s.
-    self.assertEqual(set(itertools.islice(stream, 100)), {1})
-
-  def test_mixes_streams_with_prob_half(self):
-    stream = simple.mix_streams(
-        itertools.repeat(0), itertools.repeat(1), mix_prob=0.5)
-    # Mixed stream should have both 0s and 1s.
-    self.assertEqual(set(itertools.islice(stream, 100)), {0, 1})
-
-  def test_batches_stream(self):
-    stream = iter([(0, 1), (2, 3), (4, 5), (6, 7)])
-    batched_stream = simple.batch_stream(stream, batch_size=2)
-    np.testing.assert_equal(
-        next(batched_stream), (np.array([0, 2]), np.array([1, 3])))
-    np.testing.assert_equal(
-        next(batched_stream), (np.array([4, 6]), np.array([5, 7])))
-
-  def test_plays_env_problem(self):
-    # Shape: (time, trajectory).
-    observations = np.array([[0, 1], [2, 3], [4, 5]])
-    rewards = np.array([[0, 1], [1, 0]])
-    actions = np.array([[1, 2], [2, 0]])
-    # We end the second environment 2 times, but we shouldn't collect the second
-    # trajectory.
-    dones = np.array([[False, True], [True, True]])
-    infos = [{}, {}]
-
-    mock_env = mock.MagicMock()
-    mock_env.batch_size = 2
-    # (observations, lengths)
-    mock_env.trajectories.observations_np.return_value = (None, None)
-    mock_env.reset.return_value = observations[0]
-    mock_env.step.side_effect = zip(observations[1:], rewards, dones, infos)
-
-    mock_policy_fn = mock.MagicMock()
-    mock_policy_fn.side_effect = actions
-
-    trajectories = simple.play_env_problem(mock_env, mock_policy_fn)
-    self.assertEqual(len(trajectories), 2)
-    expected_lengths = [3, 2]
-    for (i, (traj, expected_length)) in enumerate(
-        zip(trajectories, expected_lengths)):
-      self.assertEqual(traj.num_time_steps, expected_length)
-      np.testing.assert_array_equal(
-          traj.observations_np, observations[:expected_length, i])
-      np.testing.assert_array_equal(
-          traj.raw_rewards_np, rewards[:(expected_length - 1), i])
-      np.testing.assert_array_equal(
-          traj.actions_np, actions[:(expected_length - 1), i])
-
-  def _make_trajectory(self, observations=None, actions=None):
-    t = trajectory.Trajectory()
-    if observations is None:
-      observations = itertools.repeat(None)
-    if actions is None:
-      actions = itertools.repeat(None)
-    for (observation, action) in zip(observations, actions):
-      t.add_time_step(observation=observation, action=action)
-    return t
-
-  def test_replay_policy(self):
-    trajectories = [
-        self._make_trajectory(actions=actions)
-        for actions in map(np.array, [[1, 2], [3]])
-    ]
-    policy_fn = simple.ReplayPolicy(trajectories, out_of_bounds_action=0)
-    np.testing.assert_array_equal(policy_fn(None), [1, 3])
-    np.testing.assert_array_equal(policy_fn(None), [2, 0])
-
-  def test_observation_error_zero_for_same_trajectories(self):
-    observations = np.array([[0], [2], [1]])
-    (traj1, traj2) = map(self._make_trajectory, (observations, observations))
-    error = simple.calculate_observation_error([traj1], [traj2])
-    np.testing.assert_array_almost_equal(error, [0])
-
-  def test_observation_error_positive_for_different_trajectories(self):
-    observations1 = np.array([[1], [2], [3]])
-    observations2 = np.array([[0], [2], [3]])
-    (traj1, traj2) = map(self._make_trajectory, (observations1, observations2))
-    error = simple.calculate_observation_error([traj1], [traj2])
-    np.testing.assert_array_less([0], error)
-
-  def test_observation_error_dims_correspond_to_observation_dims(self):
-    observations1 = np.array([[0, 1, 0], [0, 2, 0], [0, 3, 0]])
-    observations2 = np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]])
-    (traj1, traj2) = map(self._make_trajectory, (observations1, observations2))
-    error = simple.calculate_observation_error([traj1], [traj2])
-    self.assertEqual(error.shape, (3,))
-    np.testing.assert_array_almost_equal(error[0], 0)
-    self.assertFalse(np.allclose(error[1], 0))
-    np.testing.assert_array_almost_equal(error[2], 0)
-
-  def test_observation_error_increases_with_distance(self):
-    observations_zero = np.array([[0], [0], [0]])
-    observations_positive = np.array([[3], [2], [1]])
-    (traj_zero, traj_positive, traj_negative) = map(
-        self._make_trajectory,
-        (observations_zero, observations_positive, -observations_positive),
-    )
-    error_small = simple.calculate_observation_error(
-        [traj_zero], [traj_positive])
-    error_big = simple.calculate_observation_error(
-        [traj_positive], [traj_negative])
-    np.testing.assert_array_less(error_small, error_big)
-
-  def test_observation_error_increases_with_real_trajectory_length(self):
-    observations_real_short = np.array([[1], [2]])
-    observations_real_long = np.array([[1], [2], [3]])
-    observations_sim = np.array([[0], [1]])
-    (traj_real_short, traj_real_long, traj_sim) = map(
-        self._make_trajectory,
-        (observations_real_short, observations_real_long, observations_sim),
-    )
-    error_small = simple.calculate_observation_error(
-        real_trajectories=[traj_real_short], sim_trajectories=[traj_sim])
-    error_big = simple.calculate_observation_error(
-        real_trajectories=[traj_real_long], sim_trajectories=[traj_sim])
-    np.testing.assert_array_less(error_small, error_big)
-
-  def test_observation_error_same_when_sim_trajectory_longer(self):
-    observations_real = np.array([[0], [1]])
-    observations_sim_short = np.array([[1], [2]])
-    observations_sim_long = np.array([[1], [2], [3]])
-    (traj_real, traj_sim_short, traj_sim_long) = map(
-        self._make_trajectory,
-        (observations_real, observations_sim_short, observations_sim_long),
-    )
-    error1 = simple.calculate_observation_error(
-        real_trajectories=[traj_real], sim_trajectories=[traj_sim_short])
-    error2 = simple.calculate_observation_error(
-        real_trajectories=[traj_real], sim_trajectories=[traj_sim_long])
-    np.testing.assert_array_almost_equal(error1, error2)
-
-  def test_observation_error_reduces_over_trajectories(self):
-    observations1 = np.array([[1], [2], [3]])
-    observations2 = np.array([[0], [2], [3]])
-    (traj1, traj2) = map(self._make_trajectory, (observations1, observations2))
-    error = simple.calculate_observation_error([traj1, traj1], [traj2, traj2])
-    self.assertEqual(error.shape, (1,))
-
-  @staticmethod
-  @mock.patch.object(trax, "restore_state", autospec=True)
-  def _make_env(
-      mock_restore_state, observation_space, action_space,
-      max_trajectory_length, batch_size,
-  ):
-    # (model_params, opt_state)
-    mock_restore_state.return_value.params = (None, None)
-
-    gin.bind_parameter("BoxSpaceSerializer.precision", 1)
-
-    predict_output = (np.array([[[0.0]]] * batch_size))
-    mock_model_fn = mock.MagicMock()
-    mock_model_fn.return_value.side_effect = itertools.repeat(predict_output)
-    mock_model_fn.return_value.initialize_once.return_value = ((), ())
-
-    return simulated_env_problem.SerializedSequenceSimulatedEnvProblem(
-        model=mock_model_fn,
-        reward_fn=(lambda _1, _2: np.zeros(batch_size)),
-        done_fn=(lambda _1, _2: np.full((batch_size,), False)),
-        vocab_size=1,
-        max_trajectory_length=max_trajectory_length,
-        batch_size=batch_size,
-        observation_space=observation_space,
-        action_space=action_space,
-        reward_range=(-1, 1),
-        discrete_rewards=False,
-        history_stream=itertools.repeat(None),
-        output_dir=None,
-    )
-
-  def test_evaluates_model_with_vector_observation_space(self):
-    with backend.use_backend("numpy"):
-      env = self._make_env(  # pylint: disable=no-value-for-parameter
-          observation_space=gym.spaces.Box(shape=(2,), low=0, high=1),
-          action_space=gym.spaces.Discrete(n=1),
-          max_trajectory_length=2,
-          batch_size=3,
-      )
-      trajectories = [
-          self._make_trajectory(observations, actions)  # pylint: disable=g-complex-comprehension
-          for (observations, actions) in [
-              (np.array([[0, 1]]), np.array([0])),
-              (np.array([[1, 2], [3, 4]]), np.array([0, 0])),
-              (np.array([[1, 2], [3, 4], [5, 6]]), np.array([0, 0, 0])),
-          ]
-      ]
-      metrics = simple.evaluate_model(env, trajectories, plt)
-      self.assertIsNotNone(metrics)
-      self.assertEqual(len(metrics), 2)
-
-  def test_fails_to_evaluate_model_with_matrix_observation_space(self):
-    with backend.use_backend("numpy"):
-      env = self._make_env(  # pylint: disable=no-value-for-parameter
-          observation_space=gym.spaces.Box(shape=(2, 2), low=0, high=1),
-          action_space=gym.spaces.Discrete(n=1),
-          max_trajectory_length=2,
-          batch_size=1,
-      )
-      trajectories = [
-          self._make_trajectory(np.array([[0, 1], [2, 3]]), np.array([0]))]
-      metrics = simple.evaluate_model(env, trajectories, plt)
-      self.assertIsNone(metrics)
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/rl/simple_trainer.py b/tensor2tensor/trax/rl/simple_trainer.py
deleted file mode 100644
index c5edc996a..000000000
--- a/tensor2tensor/trax/rl/simple_trainer.py
+++ /dev/null
@@ -1,341 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 trainer."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import functools
-import itertools
-import os
-import random
-import time
-
-from absl import logging
-import gin
-from matplotlib import pyplot as plt
-from tensor2tensor.trax import inputs as trax_inputs
-from tensor2tensor.trax import jaxboard
-from tensor2tensor.trax import trax
-from tensor2tensor.trax.rl import base_trainer
-from tensor2tensor.trax.rl import simple
-from tensor2tensor.trax.rl import simulated_env_problem
-from tensorflow.io import gfile
-
-
-class SimPLe(base_trainer.BaseTrainer):
-  """SimPLe trainer."""
-
-  def __init__(self,
-               train_env,
-               eval_env,
-               output_dir,
-               policy_trainer_class,
-               n_real_epochs=10,
-               data_eval_frac=0.125,
-               model_train_batch_size=64,
-               n_model_initial_train_steps=1000,
-               n_model_train_steps_per_epoch=1000,
-               simulated_env_problem_class=(
-                   simulated_env_problem.SerializedSequenceSimulatedEnvProblem),
-               simulated_batch_size=16,
-               n_simulated_epochs=1000,
-               trajectory_dump_dir=None,
-               initial_trajectory_dir=None,
-               initial_trajectory_mix_prob=0.5,
-               initial_model=None,
-               init_policy_from_world_model=False,
-               **kwargs):
-    super(SimPLe, self).__init__(train_env, eval_env, output_dir, **kwargs)
-    self._policy_dir = os.path.join(output_dir, "policy")
-    self._model_dir = os.path.join(output_dir, "model")
-    # Initialize the policy trainer lazily, so in case of initializing the
-    # policy from world model checkpoint, the trainer will try to load the
-    # checkpoint _after_ it's been created in train_model().
-    self._policy_trainer_fn = functools.partial(
-        policy_trainer_class,
-        train_env=train_env,
-        eval_env=eval_env,
-        output_dir=self._policy_dir,
-        async_mode=self._async_mode,
-        init_policy_from_world_model_output_dir=(
-            self._model_dir if init_policy_from_world_model else None
-        ),
-    )
-    self._policy_trainer = None
-    self._n_real_epochs = n_real_epochs
-    self._model_train_batch_size = model_train_batch_size
-    self._n_model_initial_train_steps = n_model_initial_train_steps
-    self._n_model_train_steps_per_epoch = n_model_train_steps_per_epoch
-    self._data_eval_frac = data_eval_frac
-
-    gfile.makedirs(self._model_dir)
-    if initial_model is not None:
-      gfile.copy(
-          initial_model,
-          os.path.join(self._model_dir, "model.pkl"),
-          overwrite=True,
-      )
-    self._initial_model = initial_model
-    self._initial_trajectories = None
-
-    self._sim_env = simulated_env_problem_class(
-        batch_size=None,
-        observation_space=train_env.observation_space,
-        action_space=train_env.action_space,
-        reward_range=train_env.reward_range,
-        discrete_rewards=train_env.discrete_rewards,
-        history_stream=None,  # TODO(pkozakowski): Support this.
-        output_dir=self._model_dir,
-    )
-    self._simulated_batch_size = simulated_batch_size
-    self._n_simulated_epochs = n_simulated_epochs
-
-    # If trajectory_dump_dir is not provided explicitly, save the trajectories
-    # in output_dir.
-    if trajectory_dump_dir is None:
-      trajectory_dump_dir = os.path.join(output_dir, "trajectories")
-    self._trajectory_dump_root_dir = trajectory_dump_dir
-
-    self._initial_trajectory_dir = initial_trajectory_dir
-    self._initial_trajectory_mix_prob = initial_trajectory_mix_prob
-
-    self._summary_writer = jaxboard.SummaryWriter(self._output_dir)
-
-    self._simple_epoch = 0
-    self._policy_epoch = 0
-    self._model_train_step = 0
-
-  @property
-  def policy_trainer(self):
-    if self._policy_trainer is None:
-      self._policy_trainer = self._policy_trainer_fn()
-    return self._policy_trainer
-
-  @property
-  def epoch(self):
-    return self._simple_epoch
-
-  def train_epoch(self, evaluate=True):
-    if self._simple_epoch > 0 or not self._has_initial_data:
-      logging.info(
-          "Collect trajectories by running the policy in the real environment.")
-      self.collect_trajectories(evaluate=evaluate)
-    if self._simple_epoch > 0 or not self._initial_model:
-      logging.info(
-          "Train the model of the environment on the collected trajectories.")
-      skipped = self.train_model()
-      if evaluate and not skipped:
-        logging.info("Evaluate the trained model.")
-        self.evaluate_model()
-    logging.info("Train the policy inside the simulated environment generated "
-                 "by the model.")
-    self.train_policy()
-
-    self._simple_epoch += 1
-
-  def evaluate(self):
-    self.policy_trainer.evaluate()
-
-  def save(self):
-    # Nothing to do, as we save stuff continuously.
-    pass
-
-  def flush_summaries(self):
-    self._summary_writer.flush()
-
-  def collect_trajectories(self, evaluate):
-    logging.info("SimPLe epoch [% 6d]: collecting data.", self._simple_epoch)
-    start_time = time.time()
-
-    self.policy_trainer.train_env = self.train_env
-    self.policy_trainer.trajectory_dump_dir = os.path.join(
-        self._trajectory_dump_root_dir, str(self.epoch))
-    self._policy_epoch += self._n_real_epochs
-    self.policy_trainer.training_loop(self._policy_epoch, evaluate=evaluate)
-
-    logging.vlog(
-        1, "Collecting trajectories took %0.2f sec.", time.time() - start_time)
-
-  def train_model(self):
-    """Train the model.
-
-    Returns:
-      whether the training was skipped due to a restart.
-    """
-    logging.info("SimPLe epoch [% 6d]: training model.", self._simple_epoch)
-    start_time = time.time()
-
-    (train_stream, eval_stream) = self._make_input_streams()
-    # Ignore n_devices for now.
-    inputs = lambda _: trax_inputs.Inputs(  # pylint: disable=g-long-lambda
-        train_stream=(lambda: train_stream),
-        train_eval_stream=(lambda: train_stream),
-        eval_stream=(lambda: eval_stream),
-        input_shape=self._sim_env.model_input_shape,
-        input_dtype=self._sim_env.model_input_dtype,
-        # TODO(lukaszkaiser): correct those, they may differ from inputs.
-        target_shape=self._sim_env.model_input_shape,
-        target_dtype=self._sim_env.model_input_dtype)
-
-    if self._simple_epoch == 0:
-      train_steps = self._n_model_initial_train_steps
-    else:
-      train_steps = self._n_model_train_steps_per_epoch
-    self._model_train_step += train_steps
-    with gin.config_scope("world_model"):
-      state = trax.train(
-          model=self._sim_env.model,
-          inputs=inputs,
-          train_steps=self._model_train_step,
-          output_dir=self._model_dir,
-          has_weights=True,
-      )
-
-    logging.vlog(
-        1, "Training model took %0.2f sec.", time.time() - start_time)
-    return state.step > self._model_train_step
-
-  def train_policy(self):
-    logging.info("SimPLe epoch [% 6d]: training policy.", self._simple_epoch)
-    start_time = time.time()
-
-    self._sim_env.initialize(
-        batch_size=self._simulated_batch_size,
-        history_stream=itertools.repeat(None),
-    )
-    # We never want async mode in the simulated env.
-    original_async_mode = self.policy_trainer.async_mode
-    self.policy_trainer.async_mode = False
-    self.policy_trainer.train_env = self._sim_env
-    # Don't dump trajectories from the simulated environment.
-    self.policy_trainer.trajectory_dump_dir = None
-    self._policy_epoch += self._n_simulated_epochs
-    self.policy_trainer.training_loop(self._policy_epoch, evaluate=False)
-    # Revert back to the original async mode in the policy trainer.
-    self.policy_trainer.async_mode = original_async_mode
-
-    logging.vlog(
-        1, "Training policy took %0.2f sec.", time.time() - start_time)
-
-  @property
-  def _has_own_data(self):
-    return self._simple_epoch > 0 or self._initial_trajectory_dir is None
-
-  @property
-  def _has_initial_data(self):
-    return self._initial_trajectory_dir is not None
-
-  def _load_trajectories(self, initial):
-    # Cache the initial trajectories in memory, as loading them can take a lot
-    # of time and they don't change.
-    if initial:
-      if self._initial_trajectories is not None:
-        return self._initial_trajectories
-      trajectory_dir = self._initial_trajectory_dir
-    else:
-      trajectory_dir = self._trajectory_dump_root_dir
-
-    trajectories = simple.load_trajectories(
-        trajectory_dir, self._data_eval_frac
-    )
-
-    if initial:
-      self._initial_trajectories = trajectories
-    return trajectories
-
-  def _make_input_streams(self):
-    def make_example_streams(initial):
-      (train_trajs, eval_trajs) = self._load_trajectories(initial)
-      generate_examples = functools.partial(
-          simple.generate_examples,
-          trajectory_to_training_examples_fn=(
-              self._sim_env.trajectory_to_training_examples),
-      )
-      return tuple(map(generate_examples, (train_trajs, eval_trajs)))
-
-    # We mix two data sources: trajectories collected in this SimPLe training
-    # loop ("own" data) and trajectories collected before, outside of this
-    # training loop ("initial" data).
-    mix_prob = self._initial_trajectory_mix_prob
-
-    if self._has_initial_data:
-      start_time = time.time()
-      # Load the initial, precollected data.
-      (init_train_stream, init_eval_stream) = make_example_streams(initial=True)
-      logging.vlog(
-          1, "Loading initial trajectories took %0.2f sec.",
-          time.time() - start_time
-      )
-    else:
-      (init_train_stream, init_eval_stream) = (None, None)
-      mix_prob = 0.0  # Take just our own collected data.
-
-    if self._has_own_data:
-      start_time = time.time()
-      # Load trajectories collected in all epochs so far.
-      (own_train_stream, own_eval_stream) = make_example_streams(initial=False)
-      logging.vlog(
-          1, "Loading own trajectories took %0.2f sec.",
-          time.time() - start_time
-      )
-    else:
-      # We start the loop with training the model, so we don't have our own
-      # collected data yet.
-      (own_train_stream, own_eval_stream) = (None, None)
-      mix_prob = 1.0  # Take just the initial data.
-
-    def mix_and_batch(streams):
-      (init_stream, own_stream) = streams
-      mixed_stream = simple.mix_streams(init_stream, own_stream, mix_prob)
-      return simple.batch_stream(mixed_stream, self._model_train_batch_size)
-
-    return tuple(
-        map(mix_and_batch, (
-            (init_train_stream, own_train_stream),
-            (init_eval_stream, own_eval_stream),
-        )))
-
-  def evaluate_model(self):
-    logging.info("SimPLe epoch [% 6d]: evaluating model.", self._simple_epoch)
-    start_time = time.time()
-
-    self._sim_env.initialize(
-        batch_size=self._simulated_batch_size,
-        history_stream=itertools.repeat(None),
-    )
-
-    (_, eval_trajectories) = self._load_trajectories(
-        # If we have any trajectories collected in this run, evaluate on them.
-        # Otherwise, use the initial dataset.
-        initial=(not self._has_own_data)
-    )
-    chosen_trajectories = [
-        random.choice(eval_trajectories)
-        for _ in range(self._sim_env.batch_size)
-    ]
-    summaries = simple.evaluate_model(self._sim_env, chosen_trajectories, plt)
-    if summaries is not None:
-      for (name, value) in summaries.items():
-        self._summary_writer.scalar(
-            "simple/{}".format(name), value, step=self._simple_epoch)
-      self._summary_writer.plot(
-          "simple/model_eval_plot", plt, step=self._simple_epoch)
-      self.flush_summaries()
-
-    logging.vlog(
-        1, "Evaluating model took %0.2f sec.", time.time() - start_time)
diff --git a/tensor2tensor/trax/rl/simple_trainer_test.py b/tensor2tensor/trax/rl/simple_trainer_test.py
deleted file mode 100644
index 317db27b8..000000000
--- a/tensor2tensor/trax/rl/simple_trainer_test.py
+++ /dev/null
@@ -1,96 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.rl.simple_trainer."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import functools
-
-import gin
-
-from tensor2tensor.envs import gym_env_problem
-from tensor2tensor.rl import gym_utils
-from tensor2tensor.trax import models
-from tensor2tensor.trax.rl import envs  # pylint: disable=unused-import
-from tensor2tensor.trax.rl import simulated_env_problem
-from tensor2tensor.trax.rl import trainers
-from tensorflow import test
-
-
-class SimpleTrainerTest(test.TestCase):
-
-  def _make_wrapped_env(self, name, max_episode_steps=2):
-    wrapper_fn = functools.partial(
-        gym_utils.gym_env_wrapper,
-        **{
-            "rl_env_max_episode_steps": max_episode_steps,
-            "maxskip_env": False,
-            "rendered_env": False,
-            "rendered_env_resize_to": None,  # Do not resize frames
-            "sticky_actions": False,
-            "output_dtype": None,
-        })
-
-    return gym_env_problem.GymEnvProblem(base_env_name=name,
-                                         batch_size=2,
-                                         env_wrapper_fn=wrapper_fn,
-                                         discrete_rewards=False)
-
-  def test_training_loop_acrobot(self):
-    gin.bind_parameter("BoxSpaceSerializer.precision", 2)
-    gin.bind_parameter("trax.train.eval_steps", 1)
-    trainer = trainers.SimPLe(
-        train_env=self._make_wrapped_env("Acrobot-v1"),
-        eval_env=self._make_wrapped_env("Acrobot-v1"),
-        output_dir=self.get_temp_dir(),
-        policy_trainer_class=functools.partial(
-            trainers.PPO,
-            policy_and_value_model=functools.partial(
-                models.FrameStackMLP,
-                n_frames=1,
-                hidden_sizes=(),
-                output_size=1,
-            ),
-            n_optimizer_steps=1,
-        ),
-        n_real_epochs=1,
-        data_eval_frac=0.5,
-        model_train_batch_size=2,
-        n_model_initial_train_steps=1,
-        n_model_train_steps_per_epoch=1,
-        simulated_env_problem_class=functools.partial(
-            simulated_env_problem.SerializedSequenceSimulatedEnvProblem,
-            model=functools.partial(
-                models.TransformerLM,
-                d_model=2,
-                n_layers=0,
-                max_len=64,
-            ),
-            reward_fn=simulated_env_problem.acrobot_reward_fn,
-            done_fn=simulated_env_problem.acrobot_done_fn,
-            vocab_size=4,
-            max_trajectory_length=4,
-        ),
-        simulated_batch_size=2,
-        n_simulated_epochs=1,
-    )
-    trainer.training_loop(n_epochs=1)
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/rl/simulated_env_problem.py b/tensor2tensor/trax/rl/simulated_env_problem.py
deleted file mode 100644
index 42d1c73b7..000000000
--- a/tensor2tensor/trax/rl/simulated_env_problem.py
+++ /dev/null
@@ -1,499 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""EnvProblem for environments simulated by a TRAX model."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import functools
-import random
-
-import numpy as np
-
-from tensor2tensor.envs import env_problem
-from tensor2tensor.trax import backend
-from tensor2tensor.trax import trax
-from tensor2tensor.trax import utils
-from tensor2tensor.trax.backend import random as jax_random
-from tensor2tensor.trax.rl import serialization_utils
-from tensor2tensor.trax.rl import space_serializer
-
-
-class SimulatedEnvProblem(env_problem.EnvProblem):
-  """EnvProblem base class for environments simulated by TRAX models.
-
-  The initial observations to start the model are taken from
-  initial_observation_stream. This iterator in incremented in every reset().
-
-  A checkpoint saved by the TRAX trainer should be available in output_dir.
-  """
-
-  def __init__(self, model, batch_size, observation_space, action_space,
-               reward_range, discrete_rewards, history_stream, output_dir,
-               model_predict_kwargs=None):
-    """Initializes the env.
-
-    Args:
-      model: TRAX model.
-      batch_size: (int) Number of simulated environments run in parallel.
-      observation_space: (gym.Space) Observation space.
-      action_space: (gym.Space) Action space.
-      reward_range: (tuple) Pair (min_reward, max_reward).
-      discrete_rewards: (bool) Whether to discretize the rewards.
-      history_stream: Iterator yielding batches of initial input data for the
-        model. The format is implementation-specific.
-      output_dir: (str) Output dir.
-      model_predict_kwargs: (dict) Additional model keyword arguments for
-        inference. Useful when different config is needed for training and
-        inference, e.g. train with memory efficient attention and predict with
-        the regular one.
-    """
-    self._model = model
-    if model_predict_kwargs is None:
-      model_predict_kwargs = {}
-    model_predict = self._model(mode="predict", **model_predict_kwargs)
-    def predict_with_state(*args, **kwargs):
-      output = model_predict(*args, **kwargs)
-      return (output, model_predict.state)
-    self._model_predict = backend.jit(predict_with_state)
-    self._model_initialize = model_predict.initialize_once
-
-    self._observation_space = observation_space
-    self._action_space = action_space
-    self._reward_range = reward_range
-    self._output_dir = output_dir
-
-    self._predict_fn = None
-    self._rng = None
-    self._model_state = None
-    self._history_stream = None
-
-    # Call the super's ctor. It will use some of the member fields, so we call
-    # it in the end.
-    super(SimulatedEnvProblem, self).__init__(
-        batch_size=batch_size,
-        discrete_rewards=discrete_rewards,
-        history_stream=history_stream,
-    )
-
-    self.seed()
-
-  def initialize_environments(self,
-                              history_stream,
-                              batch_size=1,
-                              parallelism=1):
-    """Initializes the environments.
-
-    Args:
-      history_stream: Iterator yielding batches of initial input data for the
-        model. The format is implementation-specific.
-      batch_size: (int) Number of environments in a batch.
-      parallelism: (int) Unused.
-    """
-    del parallelism
-
-    trax_state = trax.restore_state(self._output_dir)
-    # TODO(lukaszkaiser): both model state and parameters by default include
-    # the loss layer. Currently, we access the pure-model parameters by just
-    # indexing, [0] here. But we should make it more explicit in a better API.
-    model_params = trax_state.opt_state.params[0]
-    self._model_state = trax_state.model_state[0]
-
-    def predict_fn(inputs, rng):
-      (output, self._model_state) = self._model_predict(
-          inputs, params=model_params, state=self._model_state, rng=rng
-      )
-      return output
-
-    self._predict_fn = predict_fn
-    self._history_stream = history_stream
-    self._steps = np.zeros(batch_size, dtype=np.int32)
-
-  @property
-  def observation_space(self):
-    return self._observation_space
-
-  @property
-  def action_space(self):
-    return self._action_space
-
-  @property
-  def reward_range(self):
-    return self._reward_range
-
-  def seed(self, seed=None):
-    if seed is None:
-      seed = random.randint(0, 2**31 - 1)
-    self._rng = jax_random.get_prng(seed)
-    return super(SimulatedEnvProblem, self).seed(seed=seed)
-
-  def _reset_model(self, predict_fn, indices, history, rng):
-    """Resets the environments at the given indices.
-
-    Should be implemented in subclasses.
-
-    Args:
-      predict_fn: Function running prediction with the model.
-      indices: List of indices of underlying envs to call reset on.
-      history: Initial input data for the model.
-      rng: Jax RNG.
-
-    Returns:
-      np.ndarray of batched observations from the reset envs.
-    """
-    raise NotImplementedError
-
-  def _step_model(self, predict_fn, actions, rng):
-    """Takes a step in all environments.
-
-    Should be implemented in subclasses.
-
-    Args:
-      predict_fn: Function running prediction with the model.
-      actions: (np.ndarray) with first dimension equal to the batch size.
-      rng: Jax RNG.
-
-    Returns:
-      a tuple of batched raw observations, rewards and dones.
-    """
-    raise NotImplementedError
-
-  def trajectory_to_training_examples(self, trajectory):
-    raise NotImplementedError
-
-  @property
-  def model_input_shape(self):
-    raise NotImplementedError
-
-  @property
-  def model_input_dtype(self):
-    raise NotImplementedError
-
-  def _reset(self, indices):
-    """Resets environments at the given indices.
-
-    Args:
-      indices: list of indices of underlying envs to call reset on.
-
-    Returns:
-      np.ndarray of batched observations from the reset envs.
-    """
-    history = next(self._history_stream)
-    (subrng, self._rng) = jax_random.split(self._rng)
-    return self._reset_model(self._predict_fn, indices, history, subrng)
-
-  def _step(self, actions):
-    """Takes a step in all environments.
-
-    Args:
-      actions: (np.ndarray) with first dimension equal to the batch size.
-
-    Returns:
-      a tuple of batched raw observations, raw rewards, dones and infos.
-    """
-    # Predict the next observation.
-    (subrng, self._rng) = jax_random.split(self._rng)
-    (observation, reward, done) = self._step_model(
-        self._predict_fn, actions, subrng)
-    return (observation, reward, done, {})
-
-  @property
-  def model(self):
-    return self._model
-
-
-class RawSimulatedEnvProblem(SimulatedEnvProblem):
-  """SimulatedEnvProblem running a model operating on raw tensors.
-
-  Wraps an autoregressive TRAX model of signature
-  (observation_history, action) -> (observation, reward) in an EnvProblem.
-  The model is assumed to take a fixed number of last observations as input
-  and produce a single observation, which is fed back into the model in the
-  next environment step.
-
-  Shape requirements (without the batch dimension):
-    observation: Consistent with observation_space.
-    observation_history: (history_length,) + observation.shape.
-    action: Consistent with action_space.
-    reward: (1,). The singleton dimension is removed in step().
-  """
-
-  def __init__(self, history_length, trajectory_length, *args, **kwargs):
-    """Initializes the env.
-
-    Args:
-      history_length: (int) Number of last observations fed into the model.
-      trajectory_length: (int) Length of each trajectory unrolled from the
-        model.
-      *args: (tuple) Positional arguments passed to the base class.
-      **kwargs: (dict) Keyword arguments passed to the base class.
-    """
-    self._history_length = history_length
-    self._trajectory_length = trajectory_length
-    self._history = None
-    self._steps = None
-
-    super(RawSimulatedEnvProblem, self).__init__(*args, **kwargs)
-
-  def initialize_environments(self, batch_size=1, **kwargs):
-    """Initializes the environments."""
-    self._history = None
-    self._steps = np.zeros(batch_size)
-    return super(RawSimulatedEnvProblem, self).initialize_environments(
-        batch_size=batch_size, **kwargs)
-
-  def _reset_model(self, predict_fn, indices, history, rng):
-    del predict_fn
-    del rng
-    assert history.shape == ((self._batch_size, self._history_length) +
-                             self.observation_space.shape)
-
-    if self._history is None:
-      # At the first reset, all indices should be triggered.
-      assert set(indices) == set(range(self._batch_size))
-      self._history = np.array(history)
-    else:
-      history = history[indices, ...]
-      self._history[indices, ...] = history
-
-    # Reset the step counters.
-    self._steps[indices] = 0
-
-    # Return just the last timestep at the given indices.
-    return history[:, -1, ...]
-
-  def _step_model(self, predict_fn, actions, rng):
-    (observation, reward) = predict_fn((self._history, actions), rng=rng)
-
-    # Roll the history one timestep back and append the new observation.
-    self._history = np.roll(self._history, shift=-1, axis=1)
-    self._history[:, -1, ...] = observation
-
-    # Increment the step counters and determine which envs are done.
-    self._steps += 1
-    done = self._steps == self._trajectory_length
-
-    # Call copy() to get the data as numpy arrays.
-    observation = observation.copy()
-    # Reshape the rewards to get rid of the extra dimension.
-    reward = np.squeeze(reward.copy(), axis=1)
-    return (observation, reward, done)
-
-
-class SerializedSequenceSimulatedEnvProblem(SimulatedEnvProblem):
-  """SimulatedEnvProblem running a model operating on sequences of symbols.
-
-  Wraps an autoregressive TRAX model of signature past_symbols -> symbol_probs
-  in an EnvProblem. The model is assumed to take a sequence of symbols as input
-  and produce distributions over all symbols in the sequence. The next symbol
-  is sampled and fed back to the model in the next decoding step.
-
-  Shape requirements (without the batch dimension):
-    past_symbols: (max_trajectory_length * L,)
-    symbol_probs: (max_trajectory_length * L, vocab_size)
-  where L is the representation length of one environment step.
-
-  Observations, actions, rewards and done flags are (de)serialized from/to
-  sequences of symbols using an EnvSerializer passed to the constructor.
-  """
-
-  def __init__(self, model, reward_fn, done_fn, vocab_size,
-               max_trajectory_length, observation_space, action_space,
-               significance_decay=1.0, **kwargs):
-    """Initializes the env.
-
-    Args:
-      model: TRAX model to use for simulation. It's assumed to take keyword
-        arguments vocab_size and mode, where vocab_size is the number of symbols
-        in the vocabulary and mode is either "train" or "eval".
-
-      reward_fn: Function (previous_observation, current_observation) -> reward.
-      done_fn: Function (previous_observation, current_observation) -> done.
-      vocab_size: (int) Number of symbols in the vocabulary.
-      max_trajectory_length: (int) Maximum length of a trajectory unrolled from
-        the model.
-      observation_space: (gym.Space) Observation space.
-      action_space: (gym.Space) Action space.
-      significance_decay: (float) Decay for training weights of progressively
-        less significant symbols in the representation.
-      **kwargs: (dict) Keyword arguments passed to the base class.
-    """
-    self._reward_fn = reward_fn
-    self._done_fn = done_fn
-    self._vocab_size = vocab_size
-    self._max_trajectory_length = max_trajectory_length
-    self._significance_decay = significance_decay
-    self._steps = None
-    self._observation_space = None
-    self._action_space = None
-    self._last_observations = None
-
-    self._obs_serializer = space_serializer.create(
-        observation_space, self._vocab_size)
-    self._action_serializer = space_serializer.create(
-        action_space, self._vocab_size)
-    self._obs_repr_length = self._obs_serializer.representation_length
-    self._act_repr_length = self._action_serializer.representation_length
-    self._step_repr_length = self._obs_repr_length + self._act_repr_length
-
-    # We assume that the model takes vocab_size as an argument (e.g.
-    # TransformerLM).
-    model = functools.partial(model, vocab_size=vocab_size)
-    super(SerializedSequenceSimulatedEnvProblem, self).__init__(
-        model=model,
-        observation_space=observation_space,
-        action_space=action_space,
-        **kwargs
-    )
-
-  def initialize_environments(self, batch_size=1, **kwargs):
-    """Initializes the environments."""
-    self._steps = np.zeros(batch_size, dtype=np.int32)
-    self._last_observations = np.full(
-        (batch_size,) + self._observation_space.shape, np.nan)
-    self._last_symbols = np.zeros((batch_size, 1), dtype=np.int32)
-    super(SerializedSequenceSimulatedEnvProblem, self).initialize_environments(
-        batch_size=batch_size, **kwargs)
-    (subrng, self._rng) = jax_random.split(self._rng)
-    (_, self._init_model_state) = self._model_initialize(
-        input_shapes=(batch_size, 1), input_dtype=np.int32, rng=subrng
-    )
-
-  def _predict_obs(self, predict_fn, rng):
-    obs_repr = np.zeros(
-        (self._steps.shape[0], self._obs_repr_length), dtype=np.int32,
-    )
-    for (i, subrng) in enumerate(jax_random.split(rng, self._obs_repr_length)):
-      log_probs = predict_fn(self._last_symbols, rng=subrng)
-      self._last_symbols = utils.gumbel_sample(log_probs)
-      obs_repr[:, i] = self._last_symbols[:, 0]
-    return self._obs_serializer.deserialize(obs_repr)
-
-  def _consume_act(self, actions, predict_fn, rng):
-    act_repr = self._action_serializer.serialize(actions)
-    for (i, subrng) in enumerate(jax_random.split(rng, self._act_repr_length)):
-      # Run the network to update the inference buffers, but ignore the result.
-      predict_fn(self._last_symbols, rng=subrng)
-      self._last_symbols = act_repr[:, i:(i + 1)]
-
-  def _reset_model(self, predict_fn, indices, history, rng):
-    # TODO(pkozakowski): Random starts.
-    del history
-
-    indices = np.array(indices)
-    assert indices.shape[0] in (0, self._steps.shape[0]), (
-        # TODO(pkozakowski): Lift this requirement.
-        "Only resetting all envs at once is supported."
-    )
-
-    self._model_state = self._init_model_state
-    self._last_symbols[indices] = 0
-    self._steps[indices] = 0
-    observation = self._predict_obs(predict_fn, rng)[indices]
-    self._last_observations[indices] = observation
-    return observation
-
-  def _step_model(self, predict_fn, actions, rng):
-    self._consume_act(actions, predict_fn, rng)
-    self._steps += 1
-    observation = self._predict_obs(predict_fn, rng)
-    reward = self._reward_fn(self._last_observations, observation)
-    done = self._done_fn(self._last_observations, observation)
-    # Copy the last observations, so that we don't overwrite data stored in a
-    # trajectory when resetting the environment (see _reset_model).
-    self._last_observations = np.copy(observation)
-    done = np.logical_or(done, self._steps == self._max_trajectory_length - 1)
-    return (observation, reward, done)
-
-  def trajectory_to_training_examples(self, trajectory):
-    (repr_length,) = self.model_input_shape
-    seq_mask = np.ones((1, trajectory.num_time_steps - 1))
-    (reprs, repr_mask) = serialization_utils.serialize_observations_and_actions(
-        # Serialization works on batches, so we add a singleton batch dimension.
-        trajectory.observations_np[None, ...],
-        trajectory.actions_np[None, ...],
-        seq_mask,
-        self._obs_serializer,
-        self._action_serializer,
-        repr_length,
-    )
-    reprs = reprs[0, ...].astype(self.model_input_dtype)
-    sig_weights = (
-        self._significance_decay ** serialization_utils.significance_map(
-            self._obs_serializer, self._action_serializer, repr_length
-        )[None, ...]
-    )
-    obs_mask = serialization_utils.observation_mask(
-        self._obs_serializer, self._action_serializer, repr_length
-    )
-    weights = (sig_weights * obs_mask * repr_mask)[0, ...]
-    # (inputs, targets, weights)
-    return [(reprs, reprs, weights)]
-
-  @property
-  def model_input_shape(self):
-    return (self._max_trajectory_length * self._step_repr_length,)
-
-  @property
-  def model_input_dtype(self):
-    return np.int32
-
-
-def cartpole_done_fn(previous_observation, current_observation):
-  del previous_observation
-  x_threshold = 2.4
-  theta_threshold = 12 * 2 * np.pi / 360
-  x = current_observation[:, 0]
-  theta = current_observation[:, 2]
-  return np.logical_or(np.abs(x) > x_threshold, np.abs(theta) > theta_threshold)
-
-
-def cartpole_reward_fn(previous_observation, current_observation):
-  done = cartpole_done_fn(previous_observation, current_observation)
-  return 1.0 - done  # Unit reward for every timestep until the end.
-
-
-def acrobot_done_fn(previous_observation, current_observation):
-  del previous_observation
-  theta1 = current_observation[:, 0]
-  theta2 = current_observation[:, 1]
-  return -np.cos(theta1) - np.cos(theta2 + theta1) > 1.0
-
-
-def acrobot_reward_fn(previous_observation, current_observation):
-  done = acrobot_done_fn(previous_observation, current_observation)
-  return -1.0 + done  # -1 reward for every timestep until the end.
-
-
-def onlinetune_done_fn(previous_observation, current_observation):
-  del previous_observation
-  del current_observation
-  # Never return "done" from the environment, rely on max_trajectory_length
-  # instead.
-  return False
-
-
-def onlinetune_reward_fn(
-    previous_observation,
-    current_observation,
-    # 2 is the evaluation accuracy metric in the default settings of
-    # OnlineTuneEnv.
-    dim_index=2,
-):
-  prev = previous_observation[:, dim_index]
-  cur = current_observation[:, dim_index]
-  return cur - prev
diff --git a/tensor2tensor/trax/rl/simulated_env_problem_test.py b/tensor2tensor/trax/rl/simulated_env_problem_test.py
deleted file mode 100644
index 20e6e2c42..000000000
--- a/tensor2tensor/trax/rl/simulated_env_problem_test.py
+++ /dev/null
@@ -1,292 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.rl.simulated_env_problem."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import itertools
-
-import gin
-import gym
-import mock
-import numpy as np
-
-from tensor2tensor.envs import trajectory
-from tensor2tensor.trax import backend
-from tensor2tensor.trax import trax
-from tensor2tensor.trax.rl import simulated_env_problem
-from tensorflow import test
-
-
-class RawSimulatedEnvProblemTest(test.TestCase):
-
-  @staticmethod
-  @mock.patch.object(trax, "restore_state", autospec=True)
-  def _create_env(mock_restore_state, model, histories,
-                  trajectory_length):
-    # (model_params, opt_state)
-    mock_restore_state.return_value.params = (None, None)
-    space = gym.spaces.Discrete(100)
-    return simulated_env_problem.RawSimulatedEnvProblem(
-        model=model,
-        history_length=histories.shape[2],
-        trajectory_length=trajectory_length,
-        batch_size=1,
-        observation_space=space,
-        action_space=space,
-        reward_range=(-1, 1),
-        discrete_rewards=True,
-        history_stream=iter(histories),
-        output_dir=None,
-    )
-
-  def test_communicates_with_model(self):
-    # Mock model increasing the observation by action, reward is the parity of
-    # the new observation.
-    def mock_transition(inputs, *args, **kwargs):
-      del args
-      del kwargs
-      (observations, actions) = inputs
-      new_observations = observations[:, -1] + actions
-      rewards = np.array([[int(new_observations % 2 == 0)]])
-      return (new_observations, rewards)
-
-    mock_model_fn = mock.MagicMock()
-    mock_model_fn.return_value.side_effect = mock_transition
-    mock_model = mock_model_fn.return_value
-
-    actions_to_take = np.array([[1], [3]])
-    histories = np.array([[[0, 1, 2, 3]]])
-    expected_observations = np.array([[3], [4], [7]])
-    expected_rewards = np.array([[1], [0]])
-    expected_dones = np.array([[False], [True]])
-    expected_histories = np.array([[[0, 1, 2, 3]], [[1, 2, 3, 4]]])
-    expected_actions = actions_to_take
-
-    with backend.use_backend("numpy"):
-      env = self._create_env(  # pylint: disable=no-value-for-parameter
-          model=mock_model_fn,
-          histories=histories,
-          trajectory_length=len(actions_to_take),
-      )
-      actual_observations = [env.reset()]
-      actual_rewards = []
-      actual_dones = []
-      actual_histories = []
-      actual_actions = []
-      for action in actions_to_take:
-        (observation, reward, done, _) = env.step(action)
-        actual_observations.append(observation)
-        actual_rewards.append(reward)
-        actual_dones.append(done)
-        # Mock call is a tuple (args, kwargs). There is one positional argument,
-        # which is a tuple (history, action).
-        (((history, action),), _) = mock_model.call_args
-        actual_actions.append(action)
-        actual_histories.append(history)
-
-    np.testing.assert_array_equal(actual_observations, expected_observations)
-    np.testing.assert_array_equal(actual_rewards, expected_rewards)
-    np.testing.assert_array_equal(actual_dones, expected_dones)
-    np.testing.assert_array_equal(actual_histories, expected_histories)
-    np.testing.assert_array_equal(actual_actions, expected_actions)
-
-  def test_takes_new_history(self):
-    histories = np.array([[[0, 1, 2]], [[3, 4, 5]]])
-
-    with backend.use_backend("numpy"):
-      env = self._create_env(  # pylint: disable=no-value-for-parameter
-          model=mock.MagicMock(),
-          histories=histories,
-          trajectory_length=2,
-      )
-      env.reset()
-      observation = env.reset()
-      np.testing.assert_array_equal(observation, [5])
-
-
-class SerializedSequenceSimulatedEnvProblemTest(test.TestCase):
-
-  def _make_env(
-      self, observation_space, action_space, vocab_size,
-      predict_fn=None, reward_fn=None, done_fn=None,
-      batch_size=None, max_trajectory_length=None,
-  ):
-    mock_model_fn = mock.MagicMock()
-    if predict_fn is not None:
-      mock_model_fn.return_value = predict_fn
-      mock_model_fn.return_value.initialize_once.return_value = ((), ())
-    return simulated_env_problem.SerializedSequenceSimulatedEnvProblem(
-        model=mock_model_fn,
-        reward_fn=reward_fn,
-        done_fn=done_fn,
-        vocab_size=vocab_size,
-        max_trajectory_length=max_trajectory_length,
-        batch_size=batch_size,
-        observation_space=observation_space,
-        action_space=action_space,
-        reward_range=(-1, 1),
-        discrete_rewards=False,
-        history_stream=itertools.repeat(None),
-        output_dir=None,
-    )
-
-  def _make_trajectory(self, observations, actions):
-    assert len(observations) == len(actions) + 1
-    t = trajectory.Trajectory()
-    for (obs, act) in zip(observations, actions):
-      t.add_time_step(observation=obs, action=act, done=False)
-    t.add_time_step(observation=observations[-1], done=True)
-    return t
-
-  @mock.patch.object(trax, "restore_state", autospec=True)
-  def test_communicates_with_model(self, mock_restore_state):
-    gin.bind_parameter("BoxSpaceSerializer.precision", 1)
-    vocab_size = 16
-    # Mock model predicting a fixed sequence of symbols. It is made such that
-    # the first two observations are different and the last one is equal to the
-    # first.
-    symbols = [
-        1, 1, 2, 2, 0, 0,  # obs1 act1
-        1, 2, 2, 1, 0, 0,  # obs2 act2
-        1, 1, 2, 2,        # obs3
-    ]
-    def make_prediction(symbol):
-      one_hot = np.eye(vocab_size)[symbol]
-      log_probs = (1 - one_hot) * -100.0  # Virtually deterministic.
-      # (4 obs symbols + 1 action symbol) * 3 timesteps = 15.
-      return np.array([[log_probs]])
-
-    mock_predict_fn = mock.MagicMock()
-    mock_predict_fn.side_effect = map(make_prediction, symbols)
-
-    with backend.use_backend("numpy"):
-      # (model_params, opt_state)
-      mock_restore_state.return_value.params = (None, None)
-      env = self._make_env(
-          predict_fn=mock_predict_fn,
-          reward_fn=(lambda _1, _2: np.array([0.5])),
-          done_fn=(lambda _1, _2: np.array([False])),
-          vocab_size=vocab_size,
-          batch_size=1,
-          max_trajectory_length=3,
-          observation_space=gym.spaces.Box(low=0, high=5, shape=(4,)),
-          action_space=gym.spaces.MultiDiscrete(nvec=[2, 2]),
-      )
-
-      def assert_input_suffix(expected_symbols):
-        actual_symbols = np.array([
-            symbol.item() for ((symbol,), _) in mock_predict_fn.call_args_list[
-                -len(expected_symbols):
-            ]
-        ])
-        np.testing.assert_array_equal(actual_symbols, expected_symbols)
-
-      actions = [[0, 1], [1, 0]]
-
-      obs1 = env.reset()
-      assert_input_suffix(symbols[:3])
-
-      (obs2, reward, done, _) = env.step(np.array([actions[0]]))
-      # Symbols going into the decoder when predicting the next observation are:
-      # the last symbol of the previous observation, all action symbols, all
-      # symbols but the last one of the next observation.
-      assert_input_suffix([symbols[3]] + actions[0] + symbols[6:9])
-      self.assertFalse(np.array_equal(obs1, obs2))
-      np.testing.assert_array_equal(reward, [0.5])
-      np.testing.assert_array_equal(done, [False])
-
-      (obs3, reward, done, _) = env.step(np.array([actions[1]]))
-      assert_input_suffix([symbols[9]] + actions[1] + symbols[12:15])
-      np.testing.assert_array_equal(obs1, obs3)
-      np.testing.assert_array_equal(reward, [0.5])
-      np.testing.assert_array_equal(done, [True])
-
-  def test_makes_training_example(self):
-    env = self._make_env(
-        vocab_size=2,
-        observation_space=gym.spaces.Discrete(2),
-        action_space=gym.spaces.Discrete(2),
-        max_trajectory_length=3,
-    )
-    t = self._make_trajectory(observations=[0, 1, 0], actions=[1, 0])
-    examples = env.trajectory_to_training_examples(t)
-
-    # There should be 1 example with the whole trajectory.
-    self.assertEqual(len(examples), 1)
-    [(inputs, targets, weights)] = examples
-    # inputs == targets for autoregressive sequence prediction.
-    np.testing.assert_array_equal(inputs, targets)
-    # Assert array shapes and datatypes.
-    self.assertEqual(inputs.shape, env.model_input_shape)
-    self.assertEqual(inputs.dtype, env.model_input_dtype)
-    self.assertEqual(weights.shape, env.model_input_shape)
-    # Actions should be masked out.
-    self.assertEqual(np.min(weights), 0)
-    # At least part of the observation should have full weight.
-    self.assertEqual(np.max(weights), 1)
-
-  def test_makes_training_examples_from_trajectories_of_different_lengths(self):
-    env = self._make_env(
-        vocab_size=2,
-        observation_space=gym.spaces.Discrete(2),
-        action_space=gym.spaces.Discrete(2),
-        max_trajectory_length=3,
-    )
-    t1 = self._make_trajectory(observations=[0, 1], actions=[1])
-    [(x1, _, w1)] = env.trajectory_to_training_examples(t1)
-    t2 = self._make_trajectory(observations=[0, 1, 0], actions=[1, 0])
-    [(x2, _, w2)] = env.trajectory_to_training_examples(t2)
-
-    # Examples should be padded to the same shape.
-    self.assertEqual(x1.shape, x2.shape)
-    self.assertEqual(w1.shape, w2.shape)
-    # Cumulative weight should increase with trajectory length.
-    self.assertGreater(np.sum(w2), np.sum(w1))
-
-  def test_masked_representation_changes_with_observation(self):
-    env = self._make_env(
-        vocab_size=2,
-        observation_space=gym.spaces.Discrete(2),
-        action_space=gym.spaces.Discrete(2),
-        max_trajectory_length=3,
-    )
-    t1 = self._make_trajectory(observations=[0, 1], actions=[1])
-    [(x1, _, w1)] = env.trajectory_to_training_examples(t1)
-    t2 = self._make_trajectory(observations=[0, 0], actions=[1])
-    [(x2, _, w2)] = env.trajectory_to_training_examples(t2)
-
-    self.assertFalse(np.array_equal(x1 * w1, x2 * w2))
-
-  def test_masked_representation_doesnt_change_with_action(self):
-    env = self._make_env(
-        vocab_size=2,
-        observation_space=gym.spaces.Discrete(2),
-        action_space=gym.spaces.Discrete(2),
-        max_trajectory_length=3,
-    )
-    t1 = self._make_trajectory(observations=[0, 1], actions=[1])
-    [(x1, _, w1)] = env.trajectory_to_training_examples(t1)
-    t2 = self._make_trajectory(observations=[0, 1], actions=[0])
-    [(x2, _, w2)] = env.trajectory_to_training_examples(t2)
-
-    np.testing.assert_array_equal(x1 * w1, x2 * w2)
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/rl/space_serializer.py b/tensor2tensor/trax/rl/space_serializer.py
deleted file mode 100644
index b6c166ce6..000000000
--- a/tensor2tensor/trax/rl/space_serializer.py
+++ /dev/null
@@ -1,216 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Serialization of elements of Gym spaces into discrete sequences."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import copy
-
-from absl import logging
-import gin
-import gym
-import numpy as np
-
-
-class SpaceSerializer(object):
-  """Base class for Gym space serializers.
-
-  Attrs:
-    space_type: (type) Gym space class that this SpaceSerializer corresponds
-      to. Should be defined in subclasses.
-    representation_length: (int) Number of symbols in the representation of
-      every element of the space.
-    significance_map: (np.ndarray) Integer array of the same size as the
-      discrete representation, where elements describe the significance of
-      symbols, e.g. in fixed-precision encoding. 0 is the most significant
-      symbol, 1 the second most significant etc.
-  """
-
-  space_type = None
-  representation_length = None
-  significance_map = None
-
-  def __init__(self, space, vocab_size):
-    """Creates a SpaceSerializer.
-
-    Subclasses should retain the signature.
-
-    Args:
-      space: (gym.Space) Gym space of type self.space_type.
-      vocab_size: (int) Number of symbols in the vocabulary.
-    """
-    assert isinstance(space, self.space_type)
-    self._space = space
-    self._vocab_size = vocab_size
-
-  def serialize(self, data):
-    """Serializes a batch of space elements into discrete sequences.
-
-    Should be defined in subclasses.
-
-    Args:
-      data: A batch of batch_size elements of the Gym space to be serialized.
-
-    Returns:
-      int32 array of shape (batch_size, self.representation_length).
-    """
-    raise NotImplementedError
-
-  def deserialize(self, representation):
-    """Deserializes a batch of discrete sequences into space elements.
-
-    Should be defined in subclasses.
-
-    Args:
-      representation: int32 Numpy array of shape
-        (batch_size, self.representation_length) to be deserialized.
-
-    Returns:
-      A batch of batch_size deserialized elements of the Gym space.
-    """
-    raise NotImplementedError
-
-
-def create(space, vocab_size):
-  """Creates a SpaceSerializer for the given Gym space."""
-  return {
-      gym.spaces.Box: BoxSpaceSerializer,
-      gym.spaces.Discrete: DiscreteSpaceSerializer,
-      gym.spaces.MultiDiscrete: MultiDiscreteSpaceSerializer,
-  }[type(space)](space, vocab_size)
-
-
-@gin.configurable(blacklist=["space", "vocab_size"])
-class BoxSpaceSerializer(SpaceSerializer):
-  """Serializer for gym.spaces.Box.
-
-  Assumes that the space is bounded. Internally rescales it to the [0, 1]
-  interval and uses a fixed-precision encoding.
-  """
-
-  space_type = gym.spaces.Box
-
-  def __init__(self, space, vocab_size, precision=2, max_range=(-100.0, 100.0)):
-    self._precision = precision
-
-    # Some gym envs (e.g. CartPole) have unreasonably high bounds for
-    # observations. We clip so we can represent them.
-    bounded_space = copy.copy(space)
-    (min_low, max_high) = max_range
-    bounded_space.low = np.maximum(space.low, min_low)
-    bounded_space.high = np.minimum(space.high, max_high)
-    if (not np.allclose(bounded_space.low, space.low) or
-        not np.allclose(bounded_space.high, space.high)):
-      logging.warning(
-          "Space limits %s, %s out of bounds %s. Clipping to %s, %s.",
-          str(space.low), str(space.high), str(max_range),
-          str(bounded_space.low), str(bounded_space.high)
-      )
-
-    super(BoxSpaceSerializer, self).__init__(bounded_space, vocab_size)
-
-  def serialize(self, data):
-    array = data
-    batch_size = array.shape[0]
-    array = (array - self._space.low) / (self._space.high - self._space.low)
-    digits = []
-    for digit_index in range(-1, -self._precision - 1, -1):
-      threshold = self._vocab_size ** digit_index
-      digit = np.array(array / threshold).astype(np.int32)
-      # For the corner case of x == high.
-      digit[digit == self._vocab_size] -= 1
-      digits.append(digit)
-      array -= digit * threshold
-    digits = np.stack(digits, axis=-1)
-    return np.reshape(digits, (batch_size, -1))
-
-  def deserialize(self, representation):
-    digits = representation
-    batch_size = digits.shape[0]
-    digits = np.reshape(digits, (batch_size, -1, self._precision))
-    array = np.zeros(digits.shape[:-1])
-    for digit_index_in_seq in range(self._precision):
-      digit_index = -digit_index_in_seq - 1
-      array += self._vocab_size ** digit_index * digits[..., digit_index_in_seq]
-    array = np.reshape(array, (batch_size,) + self._space.shape)
-    return array * (self._space.high - self._space.low) + self._space.low
-
-  @property
-  def representation_length(self):
-    return self._precision * self._space.low.size
-
-  @property
-  def significance_map(self):
-    return np.reshape(np.broadcast_to(
-        np.arange(self._precision), self._space.shape + (self._precision,)), -1)
-
-
-class DiscreteSpaceSerializer(SpaceSerializer):
-  """Serializer for gym.spaces.Discrete.
-
-  Assumes that the size of the space fits in the number of symbols.
-  """
-
-  space_type = gym.spaces.Discrete
-  representation_length = 1
-
-  def __init__(self, space, vocab_size):
-    super(DiscreteSpaceSerializer, self).__init__(space, vocab_size)
-    assert space.n <= vocab_size, (
-        "Discrete space size should fit in the number of symbols.")
-
-  def serialize(self, data):
-    return np.reshape(data, (-1, 1)).astype(np.int32)
-
-  def deserialize(self, representation):
-    return np.reshape(representation, -1)
-
-  @property
-  def significance_map(self):
-    return np.zeros(1, dtype=np.int32)
-
-
-class MultiDiscreteSpaceSerializer(SpaceSerializer):
-  """Serializer for gym.spaces.MultiDiscrete.
-
-  Assumes that the number of categories in each dimension fits in the number of
-  symbols.
-  """
-
-  space_type = gym.spaces.MultiDiscrete
-
-  def __init__(self, space, vocab_size):
-    super(MultiDiscreteSpaceSerializer, self).__init__(space, vocab_size)
-    assert np.max(space.nvec) <= vocab_size, (
-        "MultiDiscrete maximum number of categories should fit in the number "
-        "of symbols."
-    )
-
-  def serialize(self, data):
-    return data.astype(np.int32)
-
-  def deserialize(self, representation):
-    return representation
-
-  @property
-  def representation_length(self):
-    return len(self._space.nvec)
-
-  @property
-  def significance_map(self):
-    return np.zeros(self.representation_length, dtype=np.int32)
diff --git a/tensor2tensor/trax/rl/space_serializer_test.py b/tensor2tensor/trax/rl/space_serializer_test.py
deleted file mode 100644
index b80bc85f6..000000000
--- a/tensor2tensor/trax/rl/space_serializer_test.py
+++ /dev/null
@@ -1,162 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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 tensor2tensor.trax.rl.space_serializer."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gin
-import gym
-import numpy as np
-
-from tensor2tensor.trax.rl import space_serializer
-from tensorflow import test
-
-
-class BoxSpaceSerializerTest(test.TestCase):
-
-  def _make_space_and_serializer(
-      self, low=-10, high=10, shape=(2,),
-      # Weird vocab_size to test that it doesn't only work with powers of 2.
-      vocab_size=257,
-      # Enough precision to represent float32s accurately.
-      precision=4,
-  ):
-    gin.bind_parameter("BoxSpaceSerializer.precision", precision)
-    space = gym.spaces.Box(low=low, high=high, shape=shape)
-    serializer = space_serializer.create(space, vocab_size=vocab_size)
-    return (space, serializer)
-
-  def _sample_batch(self, space):
-    return np.reshape(space.sample(), (1,) + space.shape)
-
-  def test_representation_length(self):
-    (space, serializer) = self._make_space_and_serializer()
-    input_array = self._sample_batch(space)
-    representation = serializer.serialize(input_array)
-    self.assertEqual(
-        representation.shape, (1, serializer.representation_length))
-
-  def test_commutes(self):
-    (space, serializer) = self._make_space_and_serializer()
-    input_array = self._sample_batch(space)
-    representation = serializer.serialize(input_array)
-    output_array = serializer.deserialize(representation)
-    np.testing.assert_array_almost_equal(input_array, output_array)
-
-  def test_representation_changes(self):
-    (space, serializer) = self._make_space_and_serializer()
-    array1 = self._sample_batch(space)
-    array2 = -array1
-    (repr1, repr2) = tuple(map(serializer.serialize, (array1, array2)))
-    self.assertFalse(np.array_equal(repr1, repr2))
-
-  def test_bounds_space(self):
-    gin.bind_parameter("BoxSpaceSerializer.max_range", (-10.0, 10.0))
-    (_, serializer) = self._make_space_and_serializer(
-        # Too wide range to represent, need to clip.
-        low=-1e18, high=1e18,
-        shape=(1,))
-    input_array = np.array([[1.2345]])
-    representation = serializer.serialize(input_array)
-    output_array = serializer.deserialize(representation)
-    np.testing.assert_array_almost_equal(input_array, output_array)
-
-  def test_significance_map(self):
-    (_, serializer) = self._make_space_and_serializer(shape=(2,))
-    np.testing.assert_array_equal(
-        serializer.significance_map, [0, 1, 2, 3, 0, 1, 2, 3])
-
-  def test_serializes_boundaries(self):
-    vocab_size = 256
-    precision = 4
-    (_, serializer) = self._make_space_and_serializer(
-        low=-1, high=1, shape=(1,), vocab_size=vocab_size, precision=precision,
-    )
-    input_array = np.array([[-1, 1]])
-    representation = serializer.serialize(input_array)
-    np.testing.assert_array_equal(
-        representation, [[0] * precision + [vocab_size - 1] * precision]
-    )
-
-
-class DiscreteSpaceSerializerTest(test.TestCase):
-
-  def setUp(self):
-    super(DiscreteSpaceSerializerTest, self).setUp()
-    self._space = gym.spaces.Discrete(n=2)
-    self._serializer = space_serializer.create(self._space, vocab_size=2)
-
-  def _sample_batch(self):
-    return np.reshape(self._space.sample(), (1,) + self._space.shape)
-
-  def test_representation_length(self):
-    input_array = self._sample_batch()
-    representation = self._serializer.serialize(input_array)
-    self.assertEqual(
-        representation.shape, (1, self._serializer.representation_length))
-
-  def test_commutes(self):
-    input_array = self._sample_batch()
-    representation = self._serializer.serialize(input_array)
-    output_array = self._serializer.deserialize(representation)
-    np.testing.assert_array_almost_equal(input_array, output_array)
-
-  def test_representation_changes(self):
-    array1 = self._sample_batch()
-    array2 = 1 - array1
-    (repr1, repr2) = tuple(map(self._serializer.serialize, (array1, array2)))
-    self.assertFalse(np.array_equal(repr1, repr2))
-
-  def test_significance_map(self):
-    np.testing.assert_array_equal(self._serializer.significance_map, [0])
-
-
-class MultiDiscreteSpaceSerializerTest(test.TestCase):
-
-  def setUp(self):
-    super(MultiDiscreteSpaceSerializerTest, self).setUp()
-    self._space = gym.spaces.MultiDiscrete(nvec=[2, 2])
-    self._serializer = space_serializer.create(self._space, vocab_size=2)
-
-  def _sample_batch(self):
-    return np.reshape(self._space.sample(), (1,) + self._space.shape)
-
-  def test_representation_length(self):
-    input_array = self._sample_batch()
-    representation = self._serializer.serialize(input_array)
-    self.assertEqual(
-        representation.shape, (1, self._serializer.representation_length))
-
-  def test_commutes(self):
-    input_array = self._sample_batch()
-    representation = self._serializer.serialize(input_array)
-    output_array = self._serializer.deserialize(representation)
-    np.testing.assert_array_almost_equal(input_array, output_array)
-
-  def test_representation_changes(self):
-    array1 = self._sample_batch()
-    array2 = 1 - array1
-    (repr1, repr2) = tuple(map(self._serializer.serialize, (array1, array2)))
-    self.assertFalse(np.array_equal(repr1, repr2))
-
-  def test_significance_map(self):
-    np.testing.assert_array_equal(self._serializer.significance_map, [0, 0])
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensor2tensor/trax/rl/trainers.py b/tensor2tensor/trax/rl/trainers.py
deleted file mode 100644
index 9368e3e61..000000000
--- a/tensor2tensor/trax/rl/trainers.py
+++ /dev/null
@@ -1,37 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Trainers defined in trax.rl."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import gin
-
-from tensor2tensor.trax.rl import ppo_trainer
-from tensor2tensor.trax.rl import simple_trainer
-
-
-# Ginify
-def trainer_configure(*args, **kwargs):
-  kwargs["module"] = "trax.rl.trainers"
-  kwargs["blacklist"] = ["train_env", "eval_env", "output_dir"]
-  return gin.external_configurable(*args, **kwargs)
-
-
-# pylint: disable=invalid-name
-PPO = trainer_configure(ppo_trainer.PPO)
-SimPLe = trainer_configure(simple_trainer.SimPLe)
diff --git a/tensor2tensor/trax/rl_trainer.py b/tensor2tensor/trax/rl_trainer.py
deleted file mode 100644
index 65c4b7ed1..000000000
--- a/tensor2tensor/trax/rl_trainer.py
+++ /dev/null
@@ -1,209 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-r"""Trainer for RL environments.
-
-For now we only support PPO as RL algorithm.
-
-Sample invocation:
-
-TRAIN_BATCH_SIZE=32
-python trax/rl_trainer.py \
-  --config_file=trax/rl/configs/acrobot.gin \
-  --train_batch_size=${TRAIN_BATCH_SIZE} \
-  --output_dir=${HOME}/ppo_acrobot \
-  --vmodule=*/tensor2tensor/*=1 \
-  --alsologtostderr
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import multiprocessing
-import os
-
-from absl import app
-from absl import flags
-from absl import logging
-import gin
-import jax
-from jax.config import config
-from tensor2tensor import envs  # pylint: disable=unused-import
-from tensor2tensor.envs import env_problem_utils
-from tensor2tensor.rl.google import atari_utils  # GOOGLE-INTERNAL:
-from tensor2tensor.trax import rl  # pylint: disable=unused-import
-from tensor2tensor.trax.rl import envs as rl_envs  # pylint: disable=unused-import
-from tensor2tensor.trax.rl import trainers as rl_trainers
-
-
-FLAGS = flags.FLAGS
-
-flags.DEFINE_boolean(
-    "jax_debug_nans", False,
-    "Setting to true will help to debug nans and disable jit.")
-flags.DEFINE_boolean("disable_jit", False, "Setting to true will disable jit.")
-
-flags.DEFINE_string("output_dir", "", "Output dir.")
-flags.DEFINE_string("envs_output_dir", "", "Output dir for the envs.")
-flags.DEFINE_multi_string("config_file", None,
-                          "Configuration file with parameters (.gin).")
-flags.DEFINE_multi_string("config", None,
-                          "Configuration parameters (gin string).")
-flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
-flags.DEFINE_bool("xm", False, "Copy atari roms?")
-flags.DEFINE_integer("train_batch_size", 32,
-                     "Number of parallel environments during training.")
-flags.DEFINE_integer("eval_batch_size", 4, "Batch size for evaluation.")
-flags.DEFINE_boolean("parallelize_envs", False,
-                     "If true, sets parallelism to number of cpu cores.")
-flags.DEFINE_string("trajectory_dump_dir", "",
-                    "Directory to dump trajectories to.")
-
-# TODO(afrozm): Find a better way to do these configurations.
-flags.DEFINE_string("train_server_bns", "", "Train Server's BNS.")
-flags.DEFINE_string("eval_server_bns", "", "Eval Server's BNS.")
-
-flags.DEFINE_bool("async_mode", False, "Async mode.")
-
-
-# Not just "train" to avoid a conflict with trax.train in GIN files.
-@gin.configurable(blacklist=[
-    "output_dir", "train_batch_size", "eval_batch_size", "trajectory_dump_dir"
-])
-def train_rl(
-    output_dir,
-    train_batch_size,
-    eval_batch_size,
-    env_name="ClientEnv-v0",
-    max_timestep=None,
-    clip_rewards=False,
-    rendered_env=False,
-    resize_dims=(105, 80),
-    trainer_class=rl_trainers.PPO,
-    n_epochs=10000,
-    trajectory_dump_dir=None,
-):
-  """Train the RL agent.
-
-  Args:
-    output_dir: Output directory.
-    train_batch_size: Number of parallel environments to use for training.
-    eval_batch_size: Number of parallel environments to use for evaluation.
-    env_name: Name of the environment.
-    max_timestep: Int or None, the maximum number of timesteps in a trajectory.
-      The environment is wrapped in a TimeLimit wrapper.
-    clip_rewards: Whether to clip and discretize the rewards.
-    rendered_env: Whether the environment has visual input. If so, a
-      RenderedEnvProblem will be used.
-    resize_dims: Pair (height, width), dimensions to resize the visual
-      observations to.
-    trainer_class: RLTrainer class to use.
-    n_epochs: Number epochs to run the training for.
-    trajectory_dump_dir: Directory to dump trajectories to.
-  """
-
-  if FLAGS.jax_debug_nans:
-    config.update("jax_debug_nans", True)
-
-  if FLAGS.use_tpu:
-    config.update("jax_platform_name", "tpu")
-  else:
-    config.update("jax_platform_name", "gpu")
-
-
-  # TODO(pkozakowski): Find a better way to determine this.
-  train_env_kwargs = {}
-  eval_env_kwargs = {}
-  if "OnlineTuneEnv" in env_name:
-    envs_output_dir = FLAGS.envs_output_dir or os.path.join(output_dir, "envs")
-    train_env_output_dir = os.path.join(envs_output_dir, "train")
-    eval_env_output_dir = os.path.join(envs_output_dir, "eval")
-    train_env_kwargs = {"output_dir": train_env_output_dir}
-    eval_env_kwargs = {"output_dir": eval_env_output_dir}
-
-  if "ClientEnv" in env_name:
-    train_env_kwargs["per_env_kwargs"] = [{
-        "remote_env_address": os.path.join(FLAGS.train_server_bns, str(replica))
-    } for replica in range(train_batch_size)]
-
-    eval_env_kwargs["per_env_kwargs"] = [{
-        "remote_env_address": os.path.join(FLAGS.eval_server_bns, str(replica))
-    } for replica in range(eval_batch_size)]
-
-  # TODO(afrozm): Should we leave out some cores?
-  parallelism = multiprocessing.cpu_count() if FLAGS.parallelize_envs else 1
-
-  train_env = env_problem_utils.make_env(
-      batch_size=train_batch_size,
-      env_problem_name=env_name,
-      resize=rendered_env,
-      resize_dims=resize_dims,
-      max_timestep=max_timestep,
-      clip_rewards=clip_rewards,
-      parallelism=parallelism,
-      use_tpu=FLAGS.use_tpu,
-      **train_env_kwargs)
-  assert train_env
-
-  eval_env = env_problem_utils.make_env(
-      batch_size=eval_batch_size,
-      env_problem_name=env_name,
-      resize=rendered_env,
-      resize_dims=resize_dims,
-      max_timestep=max_timestep,
-      clip_rewards=clip_rewards,
-      parallelism=parallelism,
-      use_tpu=FLAGS.use_tpu,
-      **eval_env_kwargs)
-  assert eval_env
-
-  def run_training_loop():
-    """Runs the training loop."""
-    logging.info("Starting the training loop.")
-
-    trainer = trainer_class(
-        output_dir=output_dir,
-        train_env=train_env,
-        eval_env=eval_env,
-        trajectory_dump_dir=trajectory_dump_dir,
-        async_mode=FLAGS.async_mode,
-    )
-    trainer.training_loop(n_epochs=n_epochs)
-
-  if FLAGS.jax_debug_nans or FLAGS.disable_jit:
-    with jax.disable_jit():
-      run_training_loop()
-  else:
-    run_training_loop()
-
-
-def main(argv):
-  del argv
-  logging.info("Starting RL training.")
-
-  gin_configs = FLAGS.config or []
-  gin.parse_config_files_and_bindings(FLAGS.config_file, gin_configs)
-
-  train_rl(
-      output_dir=FLAGS.output_dir,
-      train_batch_size=FLAGS.train_batch_size,
-      eval_batch_size=FLAGS.eval_batch_size,
-      trajectory_dump_dir=(FLAGS.trajectory_dump_dir or None),
-  )
-
-
-if __name__ == "__main__":
-  app.run(main)
diff --git a/tensor2tensor/trax/trainer.py b/tensor2tensor/trax/trainer.py
deleted file mode 100644
index 4e634b19d..000000000
--- a/tensor2tensor/trax/trainer.py
+++ /dev/null
@@ -1,137 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""trax trainer."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import datetime
-import os
-
-from absl import app
-from absl import flags
-from absl import logging
-
-import gin
-import jax
-from tensor2tensor.trax import backend
-from tensor2tensor.trax import trax
-
-import tensorflow as tf
-
-
-FLAGS = flags.FLAGS
-
-flags.DEFINE_string("dataset", None, "Which dataset to use.")
-flags.DEFINE_string("model", None, "Which model to train.")
-flags.DEFINE_string("data_dir", None, "Path to the directory with data.")
-flags.DEFINE_string("output_dir", None,
-                    "Path to the directory to save logs and checkpoints.")
-flags.DEFINE_multi_string("config_file", None,
-                          "Configuration file with parameters (.gin).")
-flags.DEFINE_multi_string("config", None,
-                          "Configuration parameters (gin string).")
-flags.DEFINE_integer("log_level", logging.INFO, "Log level.")
-flags.DEFINE_bool("use_tpu", False, "Whether we're running on TPU.")
-flags.DEFINE_bool("enable_eager_execution", True,
-                  "Whether we're running TF in eager mode.")
-flags.DEFINE_bool("tf_xla", True, "Whether to turn on XLA for TF.")
-flags.DEFINE_bool("tf_opt_pin_to_host", False, "Whether to turn on TF "
-                  "pin-to-host optimization.")
-flags.DEFINE_bool("tf_opt_layout", False, "Whether to turn on TF layout "
-                  "optimization.")
-
-
-def _default_output_dir():
-  """Default output directory."""
-  try:
-    dataset_name = gin.query_parameter("inputs.dataset_name")
-  except ValueError:
-    dataset_name = "random"
-  dir_name = "{model_name}_{dataset_name}_{timestamp}".format(
-      model_name=gin.query_parameter("train.model").configurable.name,
-      dataset_name=dataset_name,
-      timestamp=datetime.datetime.now().strftime("%Y%m%d_%H%M"),
-  )
-  dir_path = os.path.join("~", "trax", dir_name)
-  print()
-  trax.log("No --output_dir specified")
-  return dir_path
-
-
-def _setup_gin():
-  """Setup gin configuration."""
-  # Imports for configurables
-  # pylint: disable=g-import-not-at-top,unused-import,g-bad-import-order,reimported,unused-variable
-  from tensor2tensor.trax import models as _trax_models
-  from tensor2tensor.trax import optimizers as _trax_opt
-  # pylint: disable=g-import-not-at-top,unused-import,g-bad-import-order,reimported,unused-variable
-
-  configs = FLAGS.config or []
-  # Override with --dataset and --model
-  if FLAGS.dataset:
-    configs.append("inputs.dataset_name='%s'" % FLAGS.dataset)
-    if FLAGS.data_dir:
-      configs.append("inputs.data_dir='%s'" % FLAGS.data_dir)
-  if FLAGS.model:
-    configs.append("train.model=@trax.models.%s" % FLAGS.model)
-  gin.parse_config_files_and_bindings(FLAGS.config_file, configs)
-
-
-
-
-def main(_):
-
-  logging.set_verbosity(FLAGS.log_level)
-
-  if FLAGS.enable_eager_execution:
-    tf.enable_eager_execution()
-
-  if FLAGS.tf_xla:
-    tf.config.optimizer.set_jit(True)
-
-  tf.config.optimizer.set_experimental_options(
-      {"pin_to_host_optimization": FLAGS.tf_opt_pin_to_host}
-  )
-
-  tf.config.optimizer.set_experimental_options(
-      {"layout_optimizer": FLAGS.tf_opt_layout}
-  )
-
-
-  _setup_gin()
-
-  if FLAGS.enable_eager_execution and backend.get_name() in ("numpy", "jax"):
-    # Numpy backend doesn't benefit from having the input pipeline run on GPU,
-    # and jax backend has GPU memory contention if TF uses the GPU. Gin must be
-    # set up first before determining the backend.
-    tf.config.experimental.set_visible_devices([], "GPU")
-
-  # Setup output directory
-  output_dir = FLAGS.output_dir or _default_output_dir()
-  trax.log("Using --output_dir %s" % output_dir)
-  output_dir = os.path.expanduser(output_dir)
-
-  # If on TPU, let JAX know.
-  if FLAGS.use_tpu:
-    jax.config.update("jax_platform_name", "tpu")
-
-  trax.train(output_dir=output_dir)
-
-
-if __name__ == "__main__":
-  app.run(main)
diff --git a/tensor2tensor/trax/trax.py b/tensor2tensor/trax/trax.py
deleted file mode 100644
index 8ba8b0d3a..000000000
--- a/tensor2tensor/trax/trax.py
+++ /dev/null
@@ -1,958 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Trax main training functions."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import collections
-import functools
-import itertools
-import os
-import random
-import sys
-import time
-
-from absl import logging
-
-import gin
-
-import jax
-from jax import lax
-import numpy
-import six
-
-from tensor2tensor.trax import backend
-from tensor2tensor.trax import history as trax_history
-from tensor2tensor.trax import inputs as trax_inputs
-from tensor2tensor.trax import jaxboard
-from tensor2tensor.trax import layers
-from tensor2tensor.trax import learning_rate as lr
-from tensor2tensor.trax import optimizers as trax_opt
-from tensor2tensor.trax import utils
-from tensor2tensor.trax.backend import numpy as np
-from tensor2tensor.trax.backend import random as jax_random
-
-import tensorflow as tf
-from tensorflow.io import gfile
-
-
-def _stack_inputs_targets_and_get_predictions(inputs_and_targets):
-  """Helper to stack inputs and targets and retrieve predictions from output."""
-  # Inputs and targets can be lists - we build a flat one to input to the model.
-  model_inp = []
-  for x in inputs_and_targets:
-    if not isinstance(x, (list, tuple)):
-      model_inp.append(x)
-    else:
-      model_inp.extend(x)
-  # We retrieve as many predictions from model output as many there were inputs.
-  inp = inputs_and_targets[0]
-  inp_len = len(inp) if isinstance(inp, (list, tuple)) else 1
-  get_pred = lambda x: x[0] if inp_len == 1 else x[:inp_len]
-  return tuple(model_inp), get_pred
-
-
-def log(s, stdout=True):
-  logging.info(s)
-  if stdout:
-    print(s)
-    sys.stdout.flush()
-
-
-def step_log(step, s):
-  log("Step % 6d: %s" % (step, s))
-
-
-State = collections.namedtuple("_State", [
-    "step",       # Current training step number.
-    "opt_state",  # OptState.
-    "history",    # trax.history.History.
-    "model_state",
-])
-
-
-OptState = collections.namedtuple("_OptState", [
-    "params",      # Model parameters.
-    "slots",       # Per-parameter optimizer state, e.g. gradient moments.
-    "opt_params",  # Optimizer (hyper)parameters, e.g. learning rate, momentum.
-])
-
-
-def restore_state(output_dir):
-  """Restore State."""
-  params_file = os.path.join(output_dir, "model.pkl")
-  if not gfile.exists(params_file):
-    return State(step=None, opt_state=None, history=trax_history.History(),
-                 model_state=None)
-
-  pkl_module = utils.get_pickle_module()
-  with gfile.GFile(params_file, "rb") as f:
-    (opt_state, step, history, model_state) = pkl_module.load(f)
-  log("Model loaded from %s at step %d" % (params_file, step))
-  logging.debug("From loaded model : history = %s", history)
-  return State(step=step, opt_state=OptState(*opt_state), history=history,
-               model_state=model_state)
-
-
-def _save_gin(output_dir, sw=None):
-  config_path = os.path.join(output_dir, "config.gin")
-  config_str = gin.operative_config_str()
-  with gfile.GFile(config_path, "w") as f:
-    f.write(config_str)
-  if sw:
-    sw.text("gin_config",
-            jaxboard.markdownify_operative_config_str(config_str))
-
-
-def save_state(state, output_dir, keep=False):
-  """Save State and optionally gin config."""
-  pkl_module = utils.get_pickle_module()
-  params_file = os.path.join(output_dir, "model.pkl")
-  with gfile.GFile(params_file, "wb") as f:
-    pkl_module.dump((tuple(state.opt_state), state.step, state.history,
-                     state.model_state), f)
-  if keep:
-    params_file = os.path.join(output_dir, "model_{}.pkl".format(state.step))
-    with gfile.GFile(params_file, "wb") as f:
-      pkl_module.dump((tuple(state.opt_state), state.step, state.history,
-                       state.model_state), f)
-  log("Model saved to %s" % params_file, stdout=False)
-
-
-def _save_replicated(opt_state, step, history, model_state, n_devices,
-                     output_dir, keep):
-  """Save state but given a possibly replicated opt_state."""
-  if n_devices > 1:
-    first_replica = lambda x: x[0]
-    opt_state = OptState(*layers.nested_map(opt_state, first_replica))
-  # This line, while optional, allows JAX to transfer arrays from the device to
-  # the host in parallel, which is particularly important for cloud TPU.
-  if backend.get_name() == "jax":
-    opt_state = jax.device_get(opt_state)
-  save_state(State(opt_state=opt_state, step=step, history=history,
-                   model_state=model_state), output_dir, keep=keep)
-
-
-def _print_n_params(opt_state, n_devices, step):
-  """Print out the number of parameters."""
-  sizes = layers.sizes(opt_state.params)
-  if n_devices > 1:
-    unreplicate = lambda x: x[0]
-    single_params = layers.nested_map(opt_state.params, unreplicate)
-    sizes = layers.sizes(single_params)
-  total_size = layers.nested_reduce(sizes, sum)
-  step_log(step, "Total trainable parameters size: %d" % total_size)
-
-
-# Metrics to calculate and report.
-_METRICS = {
-    "accuracy": layers.AccuracyScalar,
-    "neg_log_perplexity": layers.NegLogPerplexityScalar,
-    "loss": layers.CrossEntropyLossScalar,
-}
-
-
-def evaluation_round(inputs_stream, metric_names, eval_fn, params, state, rng):
-  """Evaluate.
-
-  Args:
-    inputs_stream: iterable of inputs to evaluate on.
-    metric_names: list of strings, the order in which eval_fn returns metrics.
-    eval_fn: metric function, which takes inputs and predictions (and
-      params, state, rng) and returns a tuple of scalar metric values.
-    params: params for each f in eval_fns.
-    state: state for each f in eval_fns.
-    rng: random number generator.
-
-  Returns:
-    metrics: dict from metric name to metric value averaged over the number of
-      inputs.
-    state: end state for `predict_fn`.
-  """
-  metrics = collections.defaultdict(float)
-  count = 0
-  for inp in inputs_stream:
-    count += 1
-    rng, subrng = jax_random.split(rng)
-    metric_values = eval_fn(inp, params=params, state=state, rng=subrng)
-    try:
-      metric_values = list(metric_values)
-    except TypeError:
-      metric_values = [float(metric_values)]
-    for m, v in zip(metric_names, metric_values):
-      metrics[m] += v
-  return {m: v / count for (m, v) in six.iteritems(metrics)}, state
-
-
-def log_metrics(metrics, summ_writer, log_prefix, step, history=None):
-  """Log metrics to summary writer and history."""
-  rjust_len = max([0] + [len(name) for name in metrics])
-  for name, value in six.iteritems(metrics):
-    step_log(step, "%s %s | % .8f" % (
-        log_prefix.ljust(5), name.rjust(rjust_len), value))
-    full_name = "metrics/" + name
-    if history:
-      history.append(log_prefix, full_name, step, value)
-    if summ_writer:
-      summ_writer.scalar(full_name, value, step)
-
-
-def get_random_number_generator_and_set_seed(seed=None):
-  """Get a JAX random number generator and set random seed everywhere."""
-  random.seed(seed)
-  # While python random accepts None as seed and uses time/os seed then,
-  # some other functions expect integers so we create one here.
-  if seed is None:
-    seed = random.randint(0, 2**31 - 1)
-  tf.set_random_seed(seed)
-  numpy.random.seed(seed)
-  return jax_random.get_prng(seed)
-
-
-def epochs(total_steps, steps_to_skip, epoch_steps):
-  """Generates the number of steps in each epoch before reaching total_steps.
-
-  Args:
-    total_steps: int, total number of steps.
-    steps_to_skip: int, number of steps to skip because of a restart.
-    epoch_steps: iterable of int, numbers of steps in each epoch.
-
-  Yields:
-    epoch_steps: int, number of steps in this epoch
-  """
-  steps_to_go = total_steps - steps_to_skip
-  epoch_steps = iter(epoch_steps)
-
-  # Remove the desired number of steps from the stream.
-  for steps_this_epoch in epoch_steps:
-    if steps_this_epoch > steps_to_skip:
-      # Put back the number of steps left in the unfinished epoch.
-      epoch_steps = itertools.chain(
-          [steps_this_epoch - steps_to_skip], epoch_steps)
-    if steps_this_epoch >= steps_to_skip:
-      break
-    steps_to_skip -= steps_this_epoch
-
-  # Yield the remaining steps per epoch up to total_steps.
-  for steps_this_epoch in epoch_steps:
-    steps_this_epoch = min(steps_this_epoch, steps_to_go)
-    yield steps_this_epoch
-    steps_to_go -= steps_this_epoch
-    if steps_to_go == 0:
-      break
-
-
-@gin.configurable
-def _jit_predict_fn(model_predict, metric_fn, n_devices, jit=True):
-  """Returns a JIT-compiled predict function (unless jit=False)."""
-  model_predict = layers.Serial([model_predict, metric_fn])
-
-  if n_devices == 1:
-    return backend.jit(model_predict) if jit else model_predict
-
-  # Multi-devices, pmap and run.
-  @functools.partial(backend.pmap, axis_name="batch")
-  def mapped_predict(x, params, state, rng):
-    return model_predict(x, params=params, state=state, rng=rng)
-
-  def predict(x, params=(), state=(), rng=None):
-    """Predict function jited and parallelized as requested."""
-    pred = mapped_predict(
-        reshape_by_device(x, n_devices),
-        params,
-        state,
-        jax_random.split(rng, n_devices))
-    # Need to reduce the [device, per-device-batch, ...] tensors back to
-    # a [batch, ...] tensor. The tensors may be nested.
-    def combine(x):
-      if len(x.shape) > 1:
-        batch_size = x.shape[0] * x.shape[1]
-        return np.reshape(x, [batch_size] + list(x.shape[2:]))
-      # TODO(lukaszkaiser): is returning averages for scalars the right choice?
-      # If it is only scalar, return the average.
-      return np.mean(x, axis=0)
-    return layers.nested_map(pred, combine)
-
-  return predict
-
-
-@gin.configurable
-def _jit_update_fn(predict_fn, loss_fn, optimizer, n_devices, jit=True):
-  """Returns a (JIT-compiled) function that computes updates for one step."""
-  model_and_loss = layers.Serial([predict_fn, loss_fn])
-  # Gradients are always wrt. the first argument, so putting params first.
-  def model_and_loss_call(params, batch, state, rng):
-    res = model_and_loss(batch, params=params, state=state, rng=rng)
-    return res, model_and_loss.state
-  if n_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
-    def single_update(i, opt_state, batch, state, rng):
-      params, slots, opt_params = opt_state
-      rng, subrng = jax_random.split(rng[0])
-      grad_fn = backend.grad(model_and_loss_call, has_aux=True)
-      grads, state = grad_fn(params, batch, state, rng)
-      return optimizer.tree_update(
-          i, grads, params, slots, opt_params), state, [subrng]
-    return backend.jit(single_update) if jit else single_update
-
-  # Else, for n_devices > 1:
-  @functools.partial(backend.pmap, axis_name="batch")
-  def mapped_update(i, opt_state, batch, state, rng):
-    """This is a multi-device version of the update function above."""
-    # We assume all tensors have the first dimension = n_devices.
-    params, slots, opt_params = opt_state
-    rng, subrng = jax_random.split(rng)
-    grad_fn = backend.grad(model_and_loss_call, has_aux=True)
-    grads, state = grad_fn(params, batch, state, rng)
-    grads = jax.tree_util.tree_map(
-        lambda g: lax.psum(g, "batch"), grads)
-    return optimizer.tree_update(
-        i, grads, params, slots, opt_params), state, subrng
-
-  def update(i, opt_state, batch, state, rng):
-    return mapped_update(numpy.repeat(i, n_devices), opt_state, batch, state,
-                         rng)
-
-  return update
-
-
-@gin.configurable
-def _jit_compute_loss_fn(predict_fn, loss_fn, n_devices, jit=True):
-  """Returns a (JIT-compiled) function that computes the loss for one step."""
-  if n_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
-    def single_compute_loss(opt_state, batch, state, rng):
-      rng, subrng = jax_random.split(rng[0])
-      loss_val, state = loss_fn(opt_state[0], batch, predict_fn, state, rng)
-      return loss_val, state, [subrng]
-    return backend.jit(single_compute_loss) if jit else single_compute_loss
-
-  # Else, for n_devices > 1:
-  @functools.partial(backend.pmap, axis_name="batch")
-  def mapped_compute_loss(opt_state, batch, state, rng):
-    """This is a multi-device version of the update function above."""
-    # We assume all tensors have the first dimension = n_devices.
-    rng, subrng = jax_random.split(rng)
-    loss_val, state = loss_fn(opt_state[0], batch, predict_fn, state, rng)
-    return loss_val, state, subrng
-
-  def compute_loss(opt_state, batch, state, rng):
-    return mapped_compute_loss(
-        opt_state, reshape_by_device(batch, n_devices), state, rng)
-
-  return compute_loss
-
-
-@gin.configurable
-def _is_jit_init(value=True):
-  return value
-
-
-def _reshape_by_device_single(x, n_devices):
-  """Reshape x into a shape [n_devices, ...]."""
-  x_shape = list(x.shape)
-  batch_size = x_shape[0]
-  batch_size_per_device = batch_size // n_devices
-  # We require that n_devices divides batch_size evenly.
-  if batch_size_per_device * n_devices != batch_size:
-    logging.fatal(
-        "We require that n_devices[%d] divides batch_size[%d] evenly.",
-        n_devices, batch_size)
-  # New shape.
-  new_shape_prefix = [n_devices, batch_size_per_device]
-  return np.reshape(x, new_shape_prefix + x_shape[1:])
-
-
-def reshape_by_device(x, n_devices):
-  """Reshape possibly nested x into a shape [n_devices, ...]."""
-  return layers.nested_map(
-      x, lambda x: _reshape_by_device_single(x, n_devices))
-
-
-def multi_device_put(x, devices=None, reuse=True):
-  """Memory efficient multi-device replication in JAX.
-
-  Args:
-    x: jax DeviceArray or numpy ndarray to be replicated.
-    devices: a jax.devices() list or subset thereof of devices to
-      replicate onto.  Should match the list passed to any pmaps
-      ingesting the replicated array.
-    reuse: bool. If x is a DeviceArray whether to reuse its backing
-      device_buffer in the resulting ShardedDeviceArray.
-
-  Returns:
-    A ShardedDeviceArray with dtype = x.dtype and shape =
-    (n_devices,) + x.shape that's backed by replica
-    device_buffers on each device.
-  """
-  # Convert _FilledConstants that don't have device_buffer, etc.
-  if type(x) != jax.xla.DeviceArray:  # pylint: disable=unidiomatic-typecheck
-    x = np.array(x)
-  if not devices:
-    devices = jax.devices()
-  n_devices = len(devices)
-  x_aval = jax.xla.abstractify(x)
-  broadcast_x_aval = jax.abstract_arrays.ShapedArray(
-      (n_devices,) + x_aval.shape,
-      x_aval.dtype)
-  if reuse:
-    other_device_ordinals = [dv.id for dv in jax.devices()
-                             if dv != x.device_buffer.device()]
-    broadcast_buffers = ([x.device_buffer,] +
-                         [jax.xla.xc.Buffer.from_pyval(x, device=i)
-                          for i in other_device_ordinals])
-  else:
-    broadcast_buffers = [jax.xla.xc.Buffer.from_pyval(x, device=i)
-                         for i in range(n_devices)]
-  return jax.pxla.ShardedDeviceArray(broadcast_x_aval, broadcast_buffers)
-
-
-def _repeat_stream(stream):
-  """Repeat a stream indefinitely."""
-  while True:
-    for example in stream():
-      yield example
-
-
-@gin.configurable(whitelist=[])
-class Trainer(object):
-  """Trax trainer.
-
-  A trainer allows to make training steps, train for full epochs,
-  save the training state and access evaluation data.
-  """
-
-  def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs,
-               output_dir=None, random_seed=None, n_devices=None,
-               save_steps=None, should_save_checkpoints=True,
-               should_write_summaries=True, has_weights=False,
-               nontrainable_param_map=None, mask_id=None):
-    if save_steps is None:
-      save_steps = []
-    self._save_steps = save_steps
-    self._should_save_checkpoints = should_save_checkpoints
-    self._should_write_summaries = should_write_summaries
-    self._has_weights = has_weights
-    self._mask_id = mask_id
-    loss_fn = loss_fn(has_weights=has_weights, mask_id=mask_id)
-    device_count = jax.lib.xla_bridge.device_count()
-    n_devices = n_devices or device_count
-    # TODO(lukaszkaiser): remove this restriction when possible.
-    if n_devices != device_count:
-      raise ValueError("Jax cannot work yet with n_devices != all devices: "
-                       "%d != %d" % (n_devices, device_count))
-    self._n_devices = n_devices
-    rng = get_random_number_generator_and_set_seed(random_seed)
-    inputs = inputs(n_devices)
-    self._inputs = inputs
-
-    # Initialize the learning rate to a dummy value. It will be set in reset().
-    opt = optimizer(learning_rate=0.0)
-
-    # Setup the model.
-    model_train = model(mode="train")
-    model_predict_eval = model(mode="eval")
-
-    # Setup state.
-    rng, init_rng = jax_random.split(rng)
-    self._rngs = jax_random.split(rng, n_devices)
-    first_shape = inputs.input_shape[0]
-    # If the inputs are a tuple/list, add [None] (batch) to each element.
-    if isinstance(first_shape, (list, tuple)):
-      model_input_shape = tuple(
-          tuple([None] + list(shape)) for shape in inputs.input_shape)
-      model_target_shape = tuple(
-          tuple([None] + list(shape)) for shape in inputs.target_shape)
-    else:  # Otherwise just add [None] to the input shape.
-      model_input_shape = tuple([None] + list(inputs.input_shape))
-      model_target_shape = tuple([None] + list(inputs.target_shape))
-    # Change all None to 1 in input and target shape.
-    model_input_shape = layers.nested_map(
-        model_input_shape, lambda x: x if x else 1)
-    model_target_shape = layers.nested_map(
-        model_target_shape, lambda x: x if x else 1)
-    def new_opt_state_and_model_state(input_shape, input_dtype, target_shape,
-                                      target_dtype, rng):
-      """Returns optimizer and model states suitable for training a model."""
-      # Combine inputs and targets on the stack.
-      if not isinstance(input_dtype, (list, tuple)):
-        input_dtype = [input_dtype]
-        input_shape = [input_shape]
-      if not isinstance(target_dtype, (list, tuple)):
-        target_dtype = [target_dtype]
-        target_shape = [target_shape]
-      full_type = list(input_dtype) + list(target_dtype)
-      full_shape = list(input_shape) + list(target_shape)
-      if self._has_weights:
-        full_shape += list(target_shape)
-        full_type += [np.float32 for _ in target_dtype]
-      # We need to create a new model instance and not reuse `model_train` here,
-      # because `m.initialize` puts cached parameter values in `m` and hence the
-      # next call of `m.initialize` will give wrong results.
-      m = layers.Serial([model(mode="train"), loss_fn])
-      params, state = m.initialize_once(full_shape, full_type, rng)
-      (slots, opt_params) = opt.tree_init(params)
-      return (OptState(params, slots, opt_params), state)
-    if _is_jit_init():
-      # JIT parameter initialization to avoid memory fragmentation
-      new_opt_state_and_model_state = backend.jit(new_opt_state_and_model_state,
-                                                  static_argnums=(0, 1, 2, 3))
-    self._new_opt_state_and_model_state = (
-        lambda: new_opt_state_and_model_state(  # pylint: disable=g-long-lambda
-            model_input_shape, self._inputs.input_dtype,
-            model_target_shape, self._inputs.target_dtype, init_rng))
-
-    # jit model_predict and update so they're fast
-    # TODO(lukaszkaiser): the code below creates a layer computing
-    # multiple metrics from a single model output; re-factor for clarity.
-    dup_layer = layers.Dup3() if self._has_weights else layers.Dup2()
-    def lower(layer):
-      """Apply layer below the current inputs, targets, and possibly weights."""
-      if self._has_weights:
-        # Apply layer below inputs, targets, and loss weights.
-        return layers.Parallel([], [], [], layer)
-      else:
-        # Apply layer below inputs and targets.
-        return layers.Parallel([], [], layer)
-    metrics_layer = []
-    self._metrics = list(sorted(_METRICS.keys()))
-    for i, m in enumerate(reversed(self._metrics)):
-      metric = _METRICS[m](has_weights=self._has_weights, mask_id=self._mask_id)
-      if i != len(self._metrics) - 1:
-        metrics_layer.append(dup_layer)
-        metrics_layer.append(lower(metric))
-      else:
-        metrics_layer.append(metric)
-    # TODO(lukaszkaiser): clean this up once layer API stabilizes.
-    # For now, we need to initialize metric layers somehow, so here we go.
-    # We assume that they do not have any parameters, so this is a dummy.
-    dummy_shape = ((1, 2), (1,), (1,)) if self._has_weights else ((1, 2), (1,))
-    dummy_type = [np.float32] * (3 if self._has_weights else 2)
-    metrics_layer = layers.Serial(metrics_layer)
-    metrics_params, metrics_state = metrics_layer.initialize_once(
-        dummy_shape, tuple(dummy_type), init_rng)
-    self._metrics_params = layers.nested_map(
-        metrics_params, self._maybe_replicate)
-    self._metrics_state = layers.nested_map(
-        metrics_state, self._maybe_replicate)
-    self._jit_eval = _jit_predict_fn(
-        model_predict_eval, metrics_layer, n_devices)
-    self._jit_update_fn = _jit_update_fn(model_train, loss_fn, opt, n_devices)
-
-    self._model_train = model_train
-    self._model_predict_eval = model_predict_eval
-    self._loss_fn = loss_fn
-    # TODO(pkozakowski): "Learning rate schedules" are currently able to control
-    # control all optimizer parameters and model state, so let's rename them
-    # accordingly.
-    self._lr_schedule = lr_schedule
-
-    if nontrainable_param_map is None:
-      nontrainable_param_map = {}
-    self._nontrainable_param_map = nontrainable_param_map
-
-    # Those fields will be set in reset().
-    self._output_dir = None
-    self._train_sw = None
-    self._eval_sw = None
-    self._history = None
-    self._lr_fn = None
-    self._opt_state = None
-    self._step = None
-    self._model_state = None
-
-    if output_dir is not None:
-      self.reset(output_dir)
-
-  def reset(self, output_dir):
-    """Reset the model parameters.
-
-    Restores the parameters from the given output_dir if a checkpoint exists,
-    otherwise randomly initializes them.
-
-    Does not re-jit the model.
-
-    Args:
-      output_dir: Output directory.
-    """
-    self._output_dir = output_dir
-    gfile.makedirs(output_dir)
-    # Create summary writers and history.
-    if self._should_write_summaries:
-      self._train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
-      self._eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))
-
-    # Reset the train and eval streams.
-    self._train_stream = self._inputs.train_stream()
-    # TODO(lukaszkaiser): add an option to evaluate exactly on the full eval
-    #   set by adding a padding and stopping the stream when too large.
-    self._eval_stream = _repeat_stream(self._inputs.eval_stream)
-    self._train_eval_stream = _repeat_stream(self._inputs.train_eval_stream)
-
-    # Restore the training state.
-    state = restore_state(output_dir)
-    self._step = state.step or 0
-    history = state.history
-    self._lr_fn = self._lr_schedule(history)
-    self._history = history
-    if state.opt_state:
-      opt_state = state.opt_state
-      model_state = state.model_state
-    else:
-      opt_state, model_state = self._new_opt_state_and_model_state()
-      model_state = layers.nested_map(
-          model_state, self._maybe_replicate)
-    self._opt_state = OptState(*layers.nested_map(
-        opt_state, self._maybe_replicate))
-    self._model_state = model_state
-    if not state.opt_state:
-      self._maybe_save_state(keep=False)
-
-    self.update_nontrainable_params()
-
-  @property
-  def step(self):
-    return self._step
-
-  @property
-  def n_devices(self):
-    return self._n_devices
-
-  @property
-  def state(self):
-    return State(
-        opt_state=self._opt_state, step=self._step, history=self._history,
-        model_state=self._model_state)
-
-  @property
-  def nontrainable_params(self):
-    # TODO(lukaszkaiser): it makes no sense to use an accelerator (e.g. TPU)
-    # in op-by-op mode just to compute the learning rate. However, there
-    # should be a cleaner approach that forceably swapping out the backend.
-    with backend.use_backend("numpy"):
-      return self._lr_fn(self._step)
-
-  def _maybe_replicate(self, x):
-    if self._n_devices > 1:
-      if backend.get_name() == "jax":
-        return multi_device_put(x)
-      else:
-        return np.broadcast_to(x, (self._n_devices,) + x.shape)
-    else:
-      return x
-
-  def _maybe_save_state(self, keep):
-    if self._should_save_checkpoints:
-      _save_replicated(self._opt_state, self._step, self._history,
-                       self._model_state, self._n_devices, self._output_dir,
-                       keep)
-
-  def save_gin(self):
-    _save_gin(self._output_dir, self._train_sw)
-
-  def print_n_params(self):
-    _print_n_params(self._opt_state, self._n_devices, self._step)
-
-  def _map_to_state_dicts(self, f):
-    """Map the function f to all dicts in model state."""
-    def nested_map(x, f):
-      if isinstance(x, list):
-        return [nested_map(y, f) for y in x]
-      if isinstance(x, tuple):
-        return tuple([nested_map(y, f) for y in x])
-      if isinstance(x, dict) and len(x) == 1:
-        return f(x)
-      return x
-    return nested_map(self._model_state, f)
-
-  def _state_dicts_update(self, state_dict):
-    assert len(state_dict.keys()) == 1
-    key = list(state_dict.keys())[0]
-    value = np.array(state_dict[key])
-    return {key: np.array(self.update_model_state(key, value))}
-
-  def update_model_state(self, key, value):
-    """Updates model state based on nontrainable_params."""
-    # Translate model state keys to nontrainable param names.
-    if key in self._nontrainable_param_map:
-      param_name = self._nontrainable_param_map[key]
-    else:
-      # If a key is not in mapping, it stays the same.
-      param_name = key
-    if param_name in self.nontrainable_params:
-      if self._step == 0:
-        log("Mapping model state key {} to nontrainable param {}.".format(
-            key, param_name
-        ))
-        return self._maybe_replicate(
-            np.array(self.nontrainable_params[param_name])
-        )
-    return value
-
-  def _train_step(self, next_train_batch):
-    """Run one training step and update self._opt_state."""
-    # Calculate the current optimizer parameters.
-    # TODO(pkozakowski): Optimizer parameters get polluted with model state,
-    # which doesn't break anything but is weird. Filter it out.
-    opt_param_updates = layers.nested_map(
-        self.nontrainable_params, lambda x: self._maybe_replicate(np.array(x))
-    )
-    opt_state = self._opt_state
-    opt_state.opt_params.update(opt_param_updates)
-
-    # Run the update.
-    (params, slots), self._model_state, self._rngs = self._jit_update_fn(
-        self._step, opt_state, next_train_batch, self._model_state, self._rngs)
-    self._model_state = self._map_to_state_dicts(self._state_dicts_update)
-    self._opt_state = opt_state._replace(params=params, slots=slots)
-    self._step += 1
-
-  def train_epoch(self, epoch_steps, eval_steps):
-    """Train for one epoch."""
-    # Log separator
-    print()
-
-    # Timer
-    start_time = time.time()
-
-    for _ in range(epoch_steps):
-      # Train
-      next_train_batch = next(self._train_stream)
-      if self._n_devices > 1:  # TODO(lukaszkaiser): use everywhere if possible.
-        next_train_batch = reshape_by_device(next_train_batch, self._n_devices)
-
-      self._train_step(next_train_batch)
-
-      if self._step in self._save_steps:
-        self._maybe_save_state(keep=True)
-
-      # Log nontrainable params (learning rate, dropout etc.)
-      if (self._step == 1 or self._step % 10 == 0) and self._train_sw:
-        for (name, value) in self.nontrainable_params.items():
-          self._train_sw.scalar("training/{}".format(name), value)
-
-    # Timer
-    epoch_time = time.time() - start_time
-    step_log(self._step, "Ran %d train steps in %0.2f secs" %
-             (epoch_steps, epoch_time))
-    if epoch_steps > 1 and self._train_sw:
-      self._train_sw.scalar("training/steps per second",
-                            epoch_steps / epoch_time, step=self._step)
-
-    # Evaluate in parallel
-    self.evaluate(eval_steps)
-
-    # Save state
-    self._maybe_save_state(keep=False)
-
-    # Flush summary writers
-    if self._train_sw:
-      self._train_sw.flush()
-      self._eval_sw.flush()
-
-  def evaluate(self, eval_steps):
-    """Evaluate the model and log metrics."""
-    _, rng = jax_random.split(self._rngs[0])
-    # TODO(lukaszkaiser): both model state and parameters by default include
-    # the loss layer. Currently, we access the pure-model parameters by just
-    # indexing, [0] here. But we should make it more explicit in a better API.
-    params = (self._opt_state[0][0], self._metrics_params)
-    state = (self._model_state[0], self._metrics_state)
-    step_log(self._step, "Evaluation")
-    train_eval_slice = itertools.islice(self._train_eval_stream, eval_steps)
-    train_metrics, _ = evaluation_round(
-        train_eval_slice, self._metrics, self._jit_eval, params, state, rng)
-    log_metrics(train_metrics, self._train_sw, "train",
-                self._step, history=self._history)
-    eval_slice = itertools.islice(self._eval_stream, eval_steps)
-    eval_metrics, _ = evaluation_round(
-        eval_slice, self._metrics, self._jit_eval, params, state, rng)
-    log_metrics(eval_metrics, self._eval_sw, "eval",
-                self._step, history=self._history)
-    step_log(self._step, "Finished evaluation")
-
-    # Save the optimizer params in the history
-    for (name, value) in self.nontrainable_params.items():
-      self._history.append("train", "training/{}".format(name), self._step,
-                           value)
-
-  def update_nontrainable_params(self):
-    self._lr_fn = self._lr_schedule(self._history)
-
-  def save_computation_graphs(self, save_backward_graph):
-    """Dump computation graphs to files."""
-    if self._n_devices != 1:
-      return  # TODO(lukaszkaiser): make this work with more devices.
-    next_train_batch = next(self._train_stream)
-    output_dir = self._output_dir
-    if self._n_devices > 1:
-      next_train_batch = reshape_by_device(next_train_batch, self._n_devices)
-    params = self._opt_state[0][0]
-    forward_computation = jax.xla_computation(self._model_predict_eval)(
-        next_train_batch, params=params, state=self._model_state[0],
-        rng=self._rngs[0])
-    with gfile.GFile(os.path.join(output_dir, "forward.txt"), "w") as f:
-      f.write(forward_computation.GetHloText())
-    with gfile.GFile(os.path.join(output_dir, "forward.dot"), "w") as f:
-      f.write(forward_computation.GetHloDotGraph())
-    backward_computation = jax.xla_computation(self._jit_update_fn)(
-        self._step, self._opt_state, next_train_batch, self._model_state,
-        self._rngs)
-    with gfile.GFile(os.path.join(output_dir, "backward.txt"), "w") as f:
-      f.write(backward_computation.GetHloText())
-    if save_backward_graph:  # Backward graphs can be large so we guard it.
-      with gfile.GFile(os.path.join(output_dir, "backward.dot"), "w") as f:
-        f.write(backward_computation.GetHloDotGraph())
-
-
-@gin.configurable(whitelist=[])
-class MemoryEfficientTrainer(Trainer):
-  """Trax trainer that aims to minimize memory usage.
-  """
-  # TODO(kitaev): memory efficiency should be a feature of the main Trainer
-  # class, but there's a separate class for now because this trainer only
-  # supports evaluating the loss (and not any other metrics).
-
-  def __init__(self, *args, **kwargs):
-    super(MemoryEfficientTrainer, self).__init__(*args, **kwargs)
-    # Model predictions can use large amounts of memory. The memory-efficient
-    # approach is to compute metrics on each replica and then aggregate. For now
-    # we only implement computing the loss, and not any other metrics.
-    self._jit_compute_loss = _jit_compute_loss_fn(
-        self._model_predict_eval, self._loss_fn, self._n_devices)
-    assert not self._has_weights, (
-        "MemoryEfficientTrainer doesn't support has_weights")
-
-  def evaluate(self, eval_steps):
-    """Evaluate only the loss function (efficient, jitted, implementation)."""
-    assert not self._has_weights, (
-        "MemoryEfficientTrainer doesn't support has_weights")
-    step = self._step
-    rngs = self._rngs
-    state = self._model_state
-    history = self._history
-    compute_loss_fn = functools.partial(self._jit_compute_loss,
-                                        self._opt_state)
-    step_log(step, "Evaluation")
-    train_eval_metrics = []
-    for input_stream in [self._train_eval_stream, self._eval_stream]:
-      total = 0.0
-      count = 0.0
-      for inp in itertools.islice(input_stream, eval_steps):
-        loss_values, state, rngs = compute_loss_fn(inp, state, rngs)
-        total += float(numpy.mean(loss_values))
-        count += 1.0
-      metrics = {"loss": total / count}
-      train_eval_metrics.append(metrics)
-    # Unpack in the same order we've iterated over streams in the loop above.
-    train_metrics, eval_metrics = train_eval_metrics  # pylint: disable=unbalanced-tuple-unpacking
-    log_metrics(train_metrics, self._train_sw, "train", step, history=history)
-    log_metrics(eval_metrics, self._eval_sw, "eval", step, history=history)
-    step_log(step, "Finished evaluation")
-
-  def save_computation_graphs(self, save_backward_graph):
-    # TODO(kitaev): implement saving graphs while making sure that no op-by-op
-    # execution happens in the process.
-    del save_backward_graph
-    return
-
-
-@gin.configurable(blacklist=["output_dir"])
-def train(output_dir,
-          model=gin.REQUIRED,
-          loss_fn=layers.CrossEntropyLossScalar,
-          inputs=trax_inputs.inputs,
-          optimizer=trax_opt.Adafactor,
-          lr_schedule=lr.MultifactorSchedule,
-          trainer_class=Trainer,
-          train_steps=1000,
-          save_steps=None,
-          eval_steps=10,
-          eval_frequency=100,
-          n_devices=None,
-          random_seed=None,
-          save_graphs=True,
-          save_backward_graph=False,
-          has_weights=False,
-          nontrainable_param_map=None,
-          mask_id=None):
-  """Train the model on the inputs.
-
-  Args:
-    output_dir: Directory where to put the logs and checkpoints.
-    model: The model to train as a callable returning 2 callables, an init_fn
-      and apply_fn.
-    loss_fn: callable with signature: params, trax.inputs.Inputs, model, state,
-      rng -> loss.
-    inputs: callable returning trax.inputs.Inputs.
-    optimizer: The optimizer (see optimizers/base.py for signature).
-    lr_schedule: A learning rate schedule as a function that takes history and
-      returns a function from step to learning rate (a float).
-    trainer_class: The trainer class to use.
-    train_steps: int, total number of training steps.
-    save_steps: list of integers. Keep a model file at each of the supplied save
-      steps.
-    eval_steps: int, num of steps per evaluation. If None or 0, eval disabled.
-    eval_frequency: int, how often to run evaluation (every eval_frequency
-      steps). If None or 0, eval disabled.
-    n_devices: how many devices to use (if None, default, use all available)
-    random_seed: the random seed to use; time/os dependent if None (default).
-    save_graphs: bool, if True, save computation graph to file.
-    save_backward_graph: bool, if True, save backward graph to file too.
-    has_weights: bool, whether weights are included in the inputs.
-    nontrainable_param_map: dict, mapping from model nontrainable parameter
-      names to control names in PolicySchedule.
-    mask_id: id to mask out (None by default).
-
-  Returns:
-    trax.State
-  """
-  # TODO(lukaszkaiser): remove has_weights and mask_id later (configure loss).
-  trainer = trainer_class(model, loss_fn, optimizer, lr_schedule, inputs,
-                          output_dir,
-                          random_seed=random_seed, n_devices=n_devices,
-                          save_steps=save_steps, has_weights=has_weights,
-                          nontrainable_param_map=nontrainable_param_map,
-                          mask_id=mask_id)
-
-  epoch_steps = [train_steps]  # Only training if eval_frequency is 0 or None
-  if eval_frequency and eval_steps > 0:
-    epoch_steps = itertools.chain([1,  # first epoch only 1 step
-                                   eval_frequency - 1],
-                                  itertools.repeat(eval_frequency))
-  step_log(trainer.step,
-           "Starting training using %d devices" % trainer.n_devices)
-
-  for epoch_steps in epochs(train_steps, trainer.step, epoch_steps):
-    trainer.train_epoch(epoch_steps, eval_steps)
-
-    # Update nontrainable parameters with new history
-    trainer.update_nontrainable_params()
-
-    # Bookkeeping we do at the first step
-    if trainer.step == 1:
-      # Print number of parameters
-      trainer.print_n_params()
-
-      # Save computation graph (single-device only for now)
-      if (save_graphs and backend.get_name() == "jax"):
-        trainer.save_computation_graphs(save_backward_graph)
-
-      # Save Gin config
-      trainer.save_gin()
-
-  step_log(trainer.step, "Training done")
-  return trainer.state
diff --git a/tensor2tensor/trax/trax_test.py b/tensor2tensor/trax/trax_test.py
deleted file mode 100644
index 8c9041aab..000000000
--- a/tensor2tensor/trax/trax_test.py
+++ /dev/null
@@ -1,268 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""trax test."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import contextlib
-import functools
-import tempfile
-from absl.testing import parameterized
-
-from jax import test_util  # pylint: disable=unused-import
-from jax.config import config
-from jax.lib import xla_bridge
-
-from tensor2tensor.trax import backend
-from tensor2tensor.trax import inputs as inputs_lib
-from tensor2tensor.trax import layers
-from tensor2tensor.trax import learning_rate as lr
-from tensor2tensor.trax import models
-from tensor2tensor.trax import optimizers as trax_opt
-from tensor2tensor.trax import trax
-from tensor2tensor.trax.backend import numpy as np
-
-import tensorflow as tf
-from tensorflow import test
-from tensorflow.io import gfile
-
-
-
-def test_inputs(n_classes, with_weights=False, input_shape=(6, 6, 3)):
-  """Make trax.inputs.Inputs."""
-  batch_size = 2 * xla_bridge.device_count()
-
-  def input_stream():
-    key = backend.random.get_prng(0)
-    while True:
-      keys = backend.random.split(key, 4)
-      key = keys[0]
-      inputs = backend.random.uniform(keys[1], [batch_size] + list(input_shape))
-      targets = backend.random.randint(keys[2], [batch_size], dtype=np.int32,
-                                       minval=0, maxval=n_classes)
-      weights = backend.random.uniform(keys[3], [batch_size])
-      if with_weights:
-        yield inputs, targets, weights
-      else:
-        yield inputs, targets
-
-  return inputs_lib.Inputs(
-      train_stream=input_stream,
-      train_eval_stream=input_stream,
-      eval_stream=input_stream,
-      input_shape=input_shape,
-      input_dtype=np.float32,
-      target_shape=(),
-      target_dtype=np.int32)
-
-
-
-
-BACKENDS = ["jax"]
-
-
-class TraxTest(test.TestCase, parameterized.TestCase):
-
-  @contextlib.contextmanager
-  def tmp_dir(self):
-    tmp = tempfile.mkdtemp(dir=self.get_temp_dir())
-    yield tmp
-    gfile.rmtree(tmp)
-
-  # TODO(wangpeng): Remove `skipTest`'s when tf-numpy's `pmap` is in place
-
-  @parameterized.parameters(BACKENDS)
-  def test_train_eval_predict(self, backend_name):
-    if xla_bridge.device_count() > 1 and backend_name == "tf":
-      self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
-    with backend.use_backend(backend_name), self.tmp_dir() as output_dir:
-      # Prepare model and inputs
-      n_classes = 4
-      train_steps = 2
-      eval_steps = 2
-      # Adds Dropout and BatchNorm to test state handling.
-      def model_fn(mode="train"):
-        return layers.Model(layers.Dropout(mode=mode, rate=0.1),
-                            layers.BatchNorm(mode=mode),
-                            models.MLP(d_hidden=16,
-                                       n_output_classes=n_classes,
-                                       mode=mode))
-
-      inputs = lambda _: test_inputs(n_classes)
-
-      # Train and evaluate
-      state = trax.train(output_dir,
-                         model=model_fn,
-                         inputs=inputs,
-                         train_steps=train_steps,
-                         eval_steps=eval_steps)
-
-      # Assert total train steps
-      self.assertEqual(train_steps, state.step)
-
-      # Assert 2 evaluations ran
-      train_acc = state.history.get("train", "metrics/accuracy")
-      eval_acc = state.history.get("eval", "metrics/accuracy")
-      self.assertEqual(len(train_acc), len(eval_acc))
-      self.assertLen(eval_acc, 2)
-
-      # Predict with final params
-      inputs = inputs(1).train_stream()
-      model = layers.Serial(model_fn())
-      model(next(inputs)[0], params=state.opt_state.params)
-
-  @parameterized.parameters(BACKENDS)
-  def test_train_eval_predict_sm3(self, backend_name):
-    if xla_bridge.device_count() > 1 and backend_name == "tf":
-      self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
-    with backend.use_backend(backend_name), self.tmp_dir() as output_dir:
-      # Prepare model and inputs
-      n_classes = 4
-      train_steps = 2
-      eval_steps = 2
-      model_fn = functools.partial(models.MLP,
-                                   d_hidden=16,
-                                   n_output_classes=n_classes)
-      inputs = lambda _: test_inputs(n_classes)
-
-      # Train and evaluate
-      state = trax.train(output_dir,
-                         model=model_fn,
-                         inputs=inputs,
-                         train_steps=train_steps,
-                         eval_steps=eval_steps,
-                         optimizer=trax_opt.SM3)
-
-      # Assert total train steps
-      self.assertEqual(train_steps, state.step)
-
-      # Assert 2 evaluations ran
-      train_acc = state.history.get("train", "metrics/accuracy")
-      eval_acc = state.history.get("eval", "metrics/accuracy")
-      self.assertEqual(len(train_acc), len(eval_acc))
-      self.assertLen(eval_acc, 2)
-
-      # Predict with final params
-      inputs = inputs(1).train_stream()
-      model = layers.Serial(model_fn())
-      model(next(inputs)[0], params=state.opt_state.params)
-
-  @parameterized.parameters(BACKENDS)
-  def test_train_restart(self, backend_name):
-    if xla_bridge.device_count() > 1 and backend_name == "tf":
-      self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
-    with backend.use_backend(backend_name), self.tmp_dir() as output_dir:
-      # Prepare model and inputs
-      n_classes = 4
-      train_steps = 2
-      eval_steps = 2
-      model_fn = functools.partial(models.MLP,
-                                   d_hidden=16,
-                                   n_output_classes=n_classes)
-      inputs = lambda _: test_inputs(n_classes)
-
-      # Train and evaluate
-      trax.train(output_dir,
-                 model=model_fn,
-                 inputs=inputs,
-                 train_steps=train_steps,
-                 eval_steps=eval_steps)
-
-      # Restart training
-      state = trax.train(output_dir,
-                         model=model_fn,
-                         inputs=inputs,
-                         train_steps=(2 * train_steps),
-                         eval_steps=eval_steps)
-
-      # Assert total train steps
-      self.assertEqual(state.step, 2 * train_steps)
-
-  @parameterized.parameters(BACKENDS)
-  def test_train_with_weights(self, backend_name):
-    if xla_bridge.device_count() > 1 and backend_name == "tf":
-      self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
-    with backend.use_backend(backend_name), self.tmp_dir() as output_dir:
-      # Prepare model and inputs
-      n_classes = 4
-      train_steps = 2
-      eval_steps = 2
-      model_fn = functools.partial(models.MLP,
-                                   d_hidden=16,
-                                   n_output_classes=n_classes)
-      inputs = lambda _: test_inputs(n_classes, with_weights=True)
-
-      # Train and evaluate
-      state = trax.train(output_dir,
-                         model=model_fn,
-                         inputs=inputs,
-                         train_steps=train_steps,
-                         eval_steps=eval_steps,
-                         has_weights=True)
-
-      # Assert total train steps
-      self.assertEqual(state.step, train_steps)
-
-  @parameterized.parameters(BACKENDS)
-  def test_reset_twice(self, backend_name):
-    if xla_bridge.device_count() > 1 and backend_name == "tf":
-      self.skipTest("tf-numpy backend doesn't support multi-devices yet.")
-    with backend.use_backend(backend_name), self.tmp_dir() as output_dir1, \
-          self.tmp_dir() as output_dir2:
-      n_classes = 4
-      model_fn = functools.partial(models.MLP,
-                                   d_hidden=16,
-                                   n_output_classes=n_classes)
-      inputs = lambda _: test_inputs(n_classes)
-
-      trainer = trax.Trainer(
-          model=model_fn,
-          loss_fn=layers.CrossEntropyLossScalar,
-          optimizer=trax_opt.SM3,
-          lr_schedule=lr.MultifactorSchedule,
-          inputs=inputs,
-      )
-
-      trainer.reset(output_dir1)
-      trainer.evaluate(1)
-      trainer.reset(output_dir2)
-      trainer.evaluate(1)
-
-
-
-class EpochsTest(test.TestCase):
-
-  def test_cuts_epoch_when_total_steps_reached(self):
-    epoch_steps = trax.epochs(
-        total_steps=5, steps_to_skip=0, epoch_steps=[1, 2, 3])
-    self.assertEqual(list(epoch_steps), [1, 2, 2])
-
-  def test_skips_full_epoch(self):
-    epoch_steps = trax.epochs(
-        total_steps=4, steps_to_skip=2, epoch_steps=[2, 2])
-    self.assertEqual(list(epoch_steps), [2])
-
-  def test_skips_part_of_epoch(self):
-    epoch_steps = trax.epochs(
-        total_steps=4, steps_to_skip=1, epoch_steps=[2, 2])
-    self.assertEqual(list(epoch_steps), [1, 2])
-
-
-if __name__ == "__main__":
-  config.config_with_absl()
-  test.main()
diff --git a/tensor2tensor/trax/utils.py b/tensor2tensor/trax/utils.py
deleted file mode 100644
index 902e73ba4..000000000
--- a/tensor2tensor/trax/utils.py
+++ /dev/null
@@ -1,43 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Utility functions."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import pickle
-import sys
-
-import cloudpickle
-import numpy as np
-
-
-def get_pickle_module():
-  """Returns the appropriate pickle module based on Python version."""
-  # TODO(gilmer, lukaszkaiser): figure out how to use cloudpickle in python3.
-  # Currently the code throws an error when run in python3.
-  if sys.version_info[0] < 3:
-    return cloudpickle
-  else:
-    return pickle
-
-
-def gumbel_sample(log_probs):
-  """Gumbel sampling from a categorical distribution."""
-  u = np.random.uniform(low=1e-6, high=1.0 - 1e-6, size=log_probs.shape)
-  g = -np.log(-np.log(u))
-  return np.argmax(log_probs + g, axis=-1)

From 842826383a6daba4e23521c688cd91b55972bc5a Mon Sep 17 00:00:00 2001
From: Derek Murray <mrry@google.com>
Date: Mon, 7 Oct 2019 15:00:12 -0700
Subject: [PATCH 2541/2720] Internal change

PiperOrigin-RevId: 273383958
---
 tensor2tensor/utils/data_reader.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index 7f99eb585..c65878163 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -545,7 +545,7 @@ def collapse_nested_datasets(example):
           new_example[k] = v
         return tf.data.Dataset.from_tensor_slices(new_example)
 
-      dataset = dataset.apply(tf.data.experimental.unbatch())
+      dataset = dataset.unbatch()
       dataset = dataset.window(inferred_batch_size, inferred_batch_size,
                                chunk_stride)
       dataset = dataset.flat_map(collapse_nested_datasets)

From c395bfd58064546e69187424826eac03739f7928 Mon Sep 17 00:00:00 2001
From: Piotr Kozakowski <pkozakowski@google.com>
Date: Tue, 8 Oct 2019 15:23:37 -0700
Subject: [PATCH 2542/2720] Correct the dependencies for the MBRL notebook.

PiperOrigin-RevId: 273618262
---
 tensor2tensor/notebooks/hello_t2t-rl.ipynb | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/notebooks/hello_t2t-rl.ipynb b/tensor2tensor/notebooks/hello_t2t-rl.ipynb
index ed5e71331..47ea4bf84 100644
--- a/tensor2tensor/notebooks/hello_t2t-rl.ipynb
+++ b/tensor2tensor/notebooks/hello_t2t-rl.ipynb
@@ -85,8 +85,9 @@
     }
    ],
    "source": [
-    "!pip install -q -U tensor2tensor==1.13.1\n",
     "!pip install -q tensorflow==1.13.1\n",
+    "!pip install -q tensorflow_probability==0.6.0\n",
+    "!pip install -q tensor2tensor==1.13.1\n",
     "!pip install -q gym[atari]"
    ]
   },

From 42b35dd1625084fb7a734ea8e44ab26fcc0b0260 Mon Sep 17 00:00:00 2001
From: Mohammad Babaeizadeh <mbz@google.com>
Date: Fri, 11 Oct 2019 13:35:45 -0700
Subject: [PATCH 2543/2720] No round sampling for L2 models. This prevents the
 model to predict negative numbers.

PiperOrigin-RevId: 274238535
---
 tensor2tensor/models/video/base.py | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index bd31906c8..bdebf9632 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -381,6 +381,7 @@ def get_sampled_frame(self, pred_frame):
     # TODO(lukaszkaiser): the logic below heavily depend on the current
     # (a bit strange) video modalities - we should change that.
 
+    sampled_frame = pred_frame
     if self.is_per_pixel_softmax:
       frame_shape = common_layers.shape_list(pred_frame)
       target_shape = frame_shape[:-1] + [self.hparams.problem.num_channels]
@@ -389,13 +390,7 @@ def get_sampled_frame(self, pred_frame):
           sampled_frame, temperature=self.hparams.pixel_sampling_temperature)
       # TODO(lukaszkaiser): this should be consistent with modality.bottom()
       # sampled_frame = common_layers.standardize_images(sampled_frame)
-      sampled_frame = tf.to_float(sampled_frame)
-    else:
-      x = common_layers.convert_real_to_rgb(pred_frame)
-      x = x - tf.stop_gradient(x + tf.round(x))
-      x = common_layers.convert_rgb_to_real(x)
-      return x
-    return sampled_frame
+    return tf.to_float(sampled_frame)
 
   def __get_next_inputs(self, index, all_frames, all_actions, all_rewards):
     """Get inputs for next prediction iteration.

From b5d83812cbb33fa154450f38e8bed8e083d4b671 Mon Sep 17 00:00:00 2001
From: Gaurav Jain <gjn@google.com>
Date: Sat, 12 Oct 2019 01:24:25 -0700
Subject: [PATCH 2544/2720] Rename internal_convert_to_tensor for performance

Calling ops.internal_convert_to_tensor is more efficient than calling
ops.convert_to_tensor due to skipping the deprecated_argument_lookup and
also less python function calling overhead. We thus swap these functions
names so we can optimize most code paths.

PiperOrigin-RevId: 274321742
---
 tensor2tensor/layers/common_layers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index e772870b1..86b3a1b42 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2586,7 +2586,7 @@ def to_tensor(self):
 
 def _convert_factored_tensor_to_tensor(value, *args, **kwargs):
   # call ops.convert_to_tensor to handle optional arguments appropriately
-  return ops.internal_convert_to_tensor(value.to_tensor(), *args, **kwargs)
+  return ops.convert_to_tensor(value.to_tensor(), *args, **kwargs)
 
 
 tf.register_tensor_conversion_function(FactoredTensor,

From 3aca2ab360271a4684ffa7ac8767995e264cc558 Mon Sep 17 00:00:00 2001
From: David So <davidso@google.com>
Date: Mon, 14 Oct 2019 14:08:39 -0700
Subject: [PATCH 2545/2720] Add random seed for test consistency.

PiperOrigin-RevId: 274656425
---
 tensor2tensor/models/evolved_transformer_test.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensor2tensor/models/evolved_transformer_test.py b/tensor2tensor/models/evolved_transformer_test.py
index f989390f1..a7d782b25 100644
--- a/tensor2tensor/models/evolved_transformer_test.py
+++ b/tensor2tensor/models/evolved_transformer_test.py
@@ -103,6 +103,7 @@ def testEvolvedTransformer(self):
     self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, 1, 1, VOCAB_SIZE))
 
   def testSlowVsFast(self):
+    tf.set_random_seed(1234)
     model, features = get_model(transformer.transformer_tiny())
 
     decode_length = DECODE_LENGTH
@@ -253,6 +254,7 @@ def testGreedySlowTPUVsNonTPU(self):
     self.assertAllClose(slow_tpu_res, slow_non_tpu_res)
 
   def testGreedyFastTPUVsNonTPU(self):
+    tf.set_random_seed(1234)
     decode_length = DECODE_LENGTH
 
     model, features = self._create_greedy_infer_model()
@@ -273,6 +275,7 @@ def testGreedyFastTPUVsNonTPU(self):
     self.assertAllClose(fast_tpu_res, fast_non_tpu_res)
 
   def testGreedyTPUSlowVsFast(self):
+    tf.set_random_seed(1234)
     decode_length = DECODE_LENGTH
 
     model, features = self._create_greedy_infer_model()

From be35054ea6b4fd113ab2c137a1e4b41e952194ea Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 16 Oct 2019 12:26:18 -0700
Subject: [PATCH 2546/2720] Update t2t example notebooks to Python 3

PiperOrigin-RevId: 275088373
---
 tensor2tensor/notebooks/hello_t2t.ipynb   | 4 ++--
 tensor2tensor/notebooks/t2t_problem.ipynb | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/notebooks/hello_t2t.ipynb b/tensor2tensor/notebooks/hello_t2t.ipynb
index f2112adf5..acd9c3e15 100644
--- a/tensor2tensor/notebooks/hello_t2t.ipynb
+++ b/tensor2tensor/notebooks/hello_t2t.ipynb
@@ -16,8 +16,8 @@
       "collapsed_sections": []
     },
     "kernelspec": {
-      "display_name": "Python 2",
-      "name": "python2"
+      "display_name": "Python 3",
+      "name": "python3"
     }
   },
   "cells": [
diff --git a/tensor2tensor/notebooks/t2t_problem.ipynb b/tensor2tensor/notebooks/t2t_problem.ipynb
index 1eddf9e6b..592cbad39 100644
--- a/tensor2tensor/notebooks/t2t_problem.ipynb
+++ b/tensor2tensor/notebooks/t2t_problem.ipynb
@@ -561,8 +561,8 @@
       "views": {}
     },
     "kernelspec": {
-      "display_name": "Python 2",
-      "name": "python2"
+      "display_name": "Python 3",
+      "name": "python3"
     }
   },
   "nbformat": 4,

From e6da46ccc206c548b91f293b3bba4cff4b5d64a2 Mon Sep 17 00:00:00 2001
From: Sergei Lebedev <slebedev@google.com>
Date: Thu, 17 Oct 2019 14:43:29 -0700
Subject: [PATCH 2547/2720] Fixed cases where tf.TensorShape was constructed
 with float dimensions

This is a prerequisite for making TensorShape and Dimension more strict
about the types of their arguments.

PiperOrigin-RevId: 275342516
---
 .../models/neural_architecture_search/nas_layers.py       | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/models/neural_architecture_search/nas_layers.py b/tensor2tensor/models/neural_architecture_search/nas_layers.py
index ae8bb5ee7..f0220575c 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_layers.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers.py
@@ -297,8 +297,8 @@ def calculate_depthwise_channel_multiplier(input_depth, output_depth):
   # and output_depth % input_depth == 0. If this is the case then we
   # can satify the output_depth constraint, so the channel multiplier
   # will be set accordingly.
-  if ((output_depth >= input_depth) and (output_depth % input_depth == 0)):
-    return output_depth / input_depth
+  if output_depth >= input_depth and output_depth % input_depth == 0:
+    return output_depth // input_depth
   return 1
 
 
@@ -350,7 +350,7 @@ def _conv_function(self, input_tensor, output_depth, padding):
     channel_multiplier = calculate_depthwise_channel_multiplier(
         input_depth, output_depth)
 
-    num_input_variables = input_depth / self._num_repeat
+    num_input_variables = input_depth // self._num_repeat
     kernel_base = tf.get_variable(
         "kernel_base",
         [self._conv_width, 1, num_input_variables, channel_multiplier])
@@ -374,7 +374,7 @@ def _conv_function(self, input_tensor, output_depth, padding):
   def num_params(self, input_depth, output_depth, **unused_kwargs):
     channel_multiplier = calculate_depthwise_channel_multiplier(
         input_depth, output_depth)
-    return self._conv_width * (int(input_depth / self._num_repeat) + (
+    return self._conv_width * (input_depth // self._num_repeat + (
         input_depth % self._num_repeat)) * channel_multiplier
 
 

From 2330203cb267fa4efc4525ad5c7dffe0a7d2f2fc Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 18 Oct 2019 00:46:57 -0700
Subject: [PATCH 2548/2720] Fix deprecation warning.

PiperOrigin-RevId: 275421785
---
 tensor2tensor/utils/multistep_optimizer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
index ee80e659a..a9195b22d 100644
--- a/tensor2tensor/utils/multistep_optimizer.py
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -29,7 +29,7 @@
 import tensorflow as tf
 
 
-class MultistepAdamOptimizer(tf.train.AdamOptimizer):
+class MultistepAdamOptimizer(tf.compat.v1.train.AdamOptimizer):
   """Adam with SGD updates every n steps with accumulated gradients."""
 
   def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8,

From 6b0193cc65a7ad307271b218d1e5a9455f27ca2c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 18 Oct 2019 08:28:34 -0700
Subject: [PATCH 2549/2720] Fix deprecation warning.

PiperOrigin-RevId: 275476920
---
 tensor2tensor/utils/adafactor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index e578e6b19..bbb56f0e8 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -24,7 +24,7 @@
 import tensorflow as tf
 
 
-class AdafactorOptimizer(tf.train.Optimizer):
+class AdafactorOptimizer(tf.compat.v1.train.Optimizer):
   """Optimizer that implements the Adafactor algorithm.
 
   Adafactor is described in https://arxiv.org/abs/1804.04235.

From 1547c25571633f828ddd74accba76d07d8d043af Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 22 Oct 2019 08:46:58 -0700
Subject: [PATCH 2550/2720] Fix deprecation warning.

PiperOrigin-RevId: 276072314
---
 tensor2tensor/utils/yellowfin.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/yellowfin.py b/tensor2tensor/utils/yellowfin.py
index e21025f84..8896801c2 100644
--- a/tensor2tensor/utils/yellowfin.py
+++ b/tensor2tensor/utils/yellowfin.py
@@ -22,9 +22,9 @@
 
 
 # Values for gate_gradients.
-GATE_NONE = tf.train.Optimizer.GATE_NONE
-GATE_OP = tf.train.Optimizer.GATE_OP
-GATE_GRAPH = tf.train.Optimizer.GATE_GRAPH
+GATE_NONE = tf.compat.v1.train.Optimizer.GATE_NONE
+GATE_OP = tf.compat.v1.train.Optimizer.GATE_OP
+GATE_GRAPH = tf.compat.v1.train.Optimizer.GATE_GRAPH
 
 
 class YellowFinOptimizer(object):

From e4d9992831c9c2ea6728f45e13b094c517fc6b9f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 23 Oct 2019 14:56:16 -0700
Subject: [PATCH 2551/2720] Fix PY3 compatibility bug with map() input.

PiperOrigin-RevId: 276357497
---
 tensor2tensor/data_generators/generator_utils.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 30e7bb5eb..b66517a83 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -49,6 +49,10 @@ def to_example(dictionary):
   for (k, v) in six.iteritems(dictionary):
     if not v:
       raise ValueError("Empty generated field: %s" % str((k, v)))
+    # Subtly in PY2 vs PY3, map is not scriptable in py3. As a result,
+    # map objects will fail with TypeError, unless converted to a list.
+    if six.PY3 and isinstance(v, map):
+      v = list(v)
     if (isinstance(v[0], six.integer_types) or
         np.issubdtype(type(v[0]), np.integer)):
       features[k] = tf.train.Feature(int64_list=tf.train.Int64List(value=v))

From ae8dc45a25df3a39cf697e24461768aa051c97e6 Mon Sep 17 00:00:00 2001
From: Ian Simon <iansimon@google.com>
Date: Wed, 23 Oct 2019 16:44:59 -0700
Subject: [PATCH 2552/2720] fix shape assertion that causes relative attention
 to fail unnecessarily

PiperOrigin-RevId: 276378348
---
 tensor2tensor/layers/common_attention.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 391987458..4f6deb6aa 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1916,8 +1916,9 @@ def dot_product_self_attention_relative_v2(q,
 
     # This calculation only works for self attention.
     # q, k and v must therefore have the same shape.
+    # (Except v can have different depth.)
     q.get_shape().assert_is_compatible_with(k.get_shape())
-    q.get_shape().assert_is_compatible_with(v.get_shape())
+    q.get_shape()[:-1].assert_is_compatible_with(v.get_shape()[:-1])
 
     # Use separate embeddings suitable for keys and values.
     _, num_heads, length, depth_k = common_layers.shape_list(k)

From adb27c7e529f9eba7d01047767e7b61fe02b103d Mon Sep 17 00:00:00 2001
From: Kristy Choi <kechoi@google.com>
Date: Thu, 24 Oct 2019 10:47:26 -0700
Subject: [PATCH 2553/2720] Minor changes to tensor2tensor to handle cases
 where inputs and targets share the same symbol modality.

PiperOrigin-RevId: 276520696
---
 tensor2tensor/layers/transformer_layers.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index f61043fc8..236deb685 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -33,7 +33,8 @@ def layers():
 
 
 def transformer_prepare_encoder(inputs, target_space, hparams, features=None,
-                                type_ids=None, num_types=None):
+                                type_ids=None, num_types=None,
+                                reuse_target_embedding=tf.AUTO_REUSE):
   """Prepare one shard of the model for the encoder.
 
   Args:
@@ -45,6 +46,8 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None,
     type_ids: optional, an int64 Tensor of shape [batch, length] that allows
       for adding type embeddings, similar to positional embeddings.
     num_types: optional, an int that decides the number of types in type_ids.
+    reuse_target_embedding: option to reuse variable name in the case that
+      symbol modalities are reused between inputs/targets.
 
   Returns:
     encoder_input: a Tensor, bottom of encoder stack
@@ -98,7 +101,8 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None,
         32,
         ishape_static[-1],
         name="target_space_embedding",
-        dtype=hparams.get("activation_dtype", "float32"))
+        dtype=hparams.get("activation_dtype", "float32"),
+        reuse=reuse_target_embedding)
     emb_target_space = tf.reshape(emb_target_space, [1, 1, -1])
     encoder_input += emb_target_space
   if hparams.pos == "timing":

From f5d8c3e6957500b485f0b29f5b637e762c290a5a Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 25 Oct 2019 07:55:58 -0700
Subject: [PATCH 2554/2720] Fix deprecation warning.

PiperOrigin-RevId: 276687096
---
 tensor2tensor/layers/common_layers.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 86b3a1b42..478c0d7ce 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2737,7 +2737,7 @@ def _fn_with_custom_grad(fn, inputs, grad_fn, use_global_vars=False):
   Returns:
     fn(*inputs)
   """
-  vs = tf.get_variable_scope()
+  vs = tf.compat.v1.get_variable_scope()
   get_vars_fn = (
       vs.global_variables if use_global_vars else vs.trainable_variables)
   len_before_vars = len(get_vars_fn())
@@ -3064,7 +3064,7 @@ def grad_fn(inputs, variables, outputs, output_grads):
 
   @fn_with_custom_grad(grad_fn)
   def fn_with_recompute(*args):
-    cached_vs.append(tf.get_variable_scope())
+    cached_vs.append(tf.compat.v1.get_variable_scope())
     cached_arg_scope.append(tf.contrib.framework.current_arg_scope())
     return fn(*args)
 
@@ -3079,7 +3079,7 @@ def dense(x, units, **kwargs):
     # We need to find the layer parameters using scope name for the layer, so
     # check that the layer is named. Otherwise parameters for different layers
     # may get mixed up.
-    layer_name = tf.get_variable_scope().name
+    layer_name = tf.compat.v1.get_variable_scope().name
     if (not layer_name) or ("name" not in kwargs):
       raise ValueError(
           "Variable scope and layer name cannot be empty. Actual: "
@@ -3410,7 +3410,7 @@ def should_generate_summaries():
   if name_scope and "while/" in name_scope:
     # Summaries don't work well within tf.while_loop()
     return False
-  if tf.get_variable_scope().reuse:
+  if tf.compat.v1.get_variable_scope().reuse:
     # Avoid generating separate summaries for different data shards
     return False
   return True

From 687fbe09a1d6059f111dc12b317b096ba39e9784 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 25 Oct 2019 15:51:04 -0700
Subject: [PATCH 2555/2720] Fix deprecation warning.

PiperOrigin-RevId: 276776475
---
 tensor2tensor/utils/optimize.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index cf1150e4b..cdfd8e1b4 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -183,7 +183,7 @@ def _register_base_optimizer(name, opt):
   _register_base_optimizer(_name, _opt)
 
 
-class ConditionalOptimizer(tf.train.Optimizer):
+class ConditionalOptimizer(tf.compat.v1.train.Optimizer):
   """Conditional optimizer."""
 
   def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disable=super-init-not-called

From dc71240eb30f12f9da08acd170ee6f6834207224 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sun, 27 Oct 2019 18:48:28 -0700
Subject: [PATCH 2556/2720] Implementation of Neural Assistant model.

PiperOrigin-RevId: 276982866
---
 tensor2tensor/models/__init__.py         |   1 +
 tensor2tensor/models/neural_assistant.py | 452 +++++++++++++++++++++++
 2 files changed, 453 insertions(+)
 create mode 100644 tensor2tensor/models/neural_assistant.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 59ec38cb3..d07545379 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -34,6 +34,7 @@
 from tensor2tensor.models import mtf_resnet
 from tensor2tensor.models import mtf_transformer
 from tensor2tensor.models import mtf_transformer2
+from tensor2tensor.models import neural_assistant
 from tensor2tensor.models import neural_gpu
 from tensor2tensor.models import resnet
 from tensor2tensor.models import revnet
diff --git a/tensor2tensor/models/neural_assistant.py b/tensor2tensor/models/neural_assistant.py
new file mode 100644
index 000000000..454f467c6
--- /dev/null
+++ b/tensor2tensor/models/neural_assistant.py
@@ -0,0 +1,452 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Neural Assistant."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import six
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_layers
+from tensor2tensor.models import transformer
+from tensor2tensor.utils import registry
+import tensorflow as tf
+
+
+@registry.register_model
+class NeuralAssistant(transformer.Transformer):
+  """Attention net.  See file docstring."""
+
+  def __init__(self, *args, **kwargs):
+    super(NeuralAssistant, self).__init__(*args, **kwargs)
+    self.attention_weights = dict()  # For visualizing attention heads.
+
+    # Loss scheduling.
+    hparams = self._hparams
+    self.triple_num = hparams.train_triple_num
+
+  def model_fn(self, features):
+    with tf.variable_scope(tf.get_variable_scope(), use_resource=True) as vs:
+      self._add_variable_scope("model_fn", vs)
+      transformed_features = self.bottom(features)
+
+      if self.hparams.activation_dtype == "bfloat16":
+        for k, v in sorted(six.iteritems(transformed_features)):
+          if v.dtype == tf.float32:
+            transformed_features[k] = tf.cast(v, tf.bfloat16)
+
+      with tf.variable_scope("body") as body_vs:
+        self._add_variable_scope("body", body_vs)
+        body_out = self.body(transformed_features)
+      output, losses = self._normalize_body_output(body_out)
+
+      if "training" in losses:
+        tf.logging.info(
+            "Skipping T2TModel top and loss because training loss returned from body"
+        )
+        logits = output
+      else:
+        tf.logging.warn("The loss will be computed in model_fn now.")
+        logits = self.top(output, features)
+        losses["training"] = 0.0
+        cur_kb_loss = losses["kb_loss"]
+        cur_kb_loss_weight = self._hparams.kb_loss_weight
+        cur_lm_loss_weight = 1.0 - cur_kb_loss_weight
+        # Finalize loss
+        if (self._hparams.mode != tf.estimator.ModeKeys.PREDICT and
+            self._hparams.mode != "attack"):
+          lm_loss_num, lm_loss_denom = self.loss(logits, features)
+          total_loss = cur_kb_loss * cur_kb_loss_weight + (
+              lm_loss_num / lm_loss_denom) * cur_lm_loss_weight
+          tf.summary.scalar("kb_loss", cur_kb_loss)
+          tf.summary.scalar("lm_loss", (lm_loss_num / lm_loss_denom))
+          tf.summary.scalar("cur_kb_loss_weight",
+                            tf.reshape(cur_kb_loss_weight, []))
+          tf.logging.info("Loss computed " + str(total_loss))
+          losses = {"training": total_loss}
+
+      return logits, losses
+
+  def encode_knowledge_bottom(self, features):
+    tf.logging.info("Encoding knowledge " + str(self.triple_num))
+    hparams = self._hparams
+
+    # Make sure this is embeddings for triples
+    # <tf.float32>[batch_size, max_triple_num*max_triple_length, 1, emb_dim]
+    fact_embedding = features["encoded_triples"]
+    # [batch_size, max_triple_num*max_triple_length, emb_dim]
+    fact_embedding = tf.squeeze(fact_embedding, 2)
+
+    kb_shape = common_layers.shape_list(fact_embedding)
+    batch_size = kb_shape[0]
+    embed_dim = kb_shape[2]
+    max_triple_length = hparams.max_triple_length
+    fact_embedding = fact_embedding[:, :self.triple_num * max_triple_length, :]
+    # <tf.float32>[batch_size*max_triple_num, max_triple_length, emb_dim]
+    re_fact_embedding = tf.reshape(
+        fact_embedding, [batch_size * self.triple_num, -1, embed_dim],
+        name="reshape_fact_embedding")
+
+    # <tf.int64>[batch_size, max_triple_num]
+    input_fact_lengths = features["triple_lens"]
+    input_fact_lengths = input_fact_lengths[:, :self.triple_num]
+    # Stack the fact lengths.
+    # <tf.int64>[batch_size*max_triple_num]
+    re_fact_lengths = tf.reshape(
+        input_fact_lengths, [batch_size * self.triple_num, 1],
+        name="reshape_fact_lengths")
+
+    return re_fact_embedding, re_fact_lengths
+
+  def compute_knowledge_selection_and_loss(self, features, encoder_output,
+                                           fact_embedding, fact_lengths):
+    """Compute knowledge selection and loss.
+
+    Args:
+      features: features.
+      encoder_output: <tf.float32>[batch_size, input_length, hidden_dim]
+      fact_embedding: <tf.float32>[batch_size*max_triple_num, max_triple_length,
+        emb_dim]
+      fact_lengths: # <tf.int32>[batch_size*max_triple_num]
+
+    Returns:
+      knowledge_weights:
+      knowledge_loss:
+    """
+    hparams = self._hparams
+    encoder_output_shape = common_layers.shape_list(encoder_output)
+    encoder_hidden_dim = encoder_output_shape[-1]
+    inputs = features["inputs"]
+    # <tf.float32>[batch_size, input_length, emb_dim]
+    inputs = tf.squeeze(inputs, 2)
+    # <tf.float32>[batch_size, input_length]
+    context_padding = common_attention.embedding_to_padding(inputs)
+    # <tf.float32>[batch_size]
+    context_lens = tf.to_float(
+        common_attention.padding_to_length(context_padding))
+    # <tf.float32>[batch_size, 1]
+    context_lens = tf.expand_dims(context_lens, -1)
+    # Compute context vector summary.
+    # <tf.float32>[batch_size, hidden_dim]
+    context_vector_summary = compute_summary_embedding(encoder_output,
+                                                       context_lens, hparams)
+    knowledge_encoder_output = compute_average_embedding(
+        fact_embedding, fact_lengths)
+    # <tf.float32>[batch_size, triple_num, emb_dim]
+    knowledge_encoder_output = tf.reshape(
+        knowledge_encoder_output, [-1, self.triple_num, encoder_hidden_dim])
+    original_knowledge_encoder_output = knowledge_encoder_output
+    if hparams.similarity_fuction == "dot_product":
+      triple_logits = tf.squeeze(
+          tf.matmul(knowledge_encoder_output,
+                    tf.expand_dims(context_vector_summary, 2)), -1)
+    elif hparams.similarity_fuction == "bilinear":
+      # Tile the context vector summary.
+      # <tf.float32>[batch_size, max_triple_num*hidden_dim]
+      tiled_context_vector = tf.tile(context_vector_summary,
+                                     [1, self.triple_num])
+      # <tf.float32>[batch_size, max_triple_num, hidden_dim]
+      context_vector = tf.reshape(tiled_context_vector,
+                                  [-1, self.triple_num, encoder_hidden_dim])
+      # compute outer product
+      context_vector = tf.expand_dims(context_vector, -1)
+      knowledge_encoder_output = tf.expand_dims(knowledge_encoder_output, 2)
+      # <tf.float32>[batch_size, max_triple_num, hidden_dim, hidden_dim]
+      outer_product = tf.matmul(context_vector, knowledge_encoder_output)
+      outer_product = tf.reshape(
+          outer_product,
+          [-1, self.triple_num, encoder_hidden_dim * encoder_hidden_dim])
+      triple_logits = tf.squeeze(
+          tf.layers.dense(outer_product, 1, name="knolwedge_final_mlp"), -1)
+
+    avg_triple_loss = 0.0
+    triple_labels = features["triple_labels"]
+    triple_labels = triple_labels[:, :self.triple_num]
+    if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+      triple_losses = tf.nn.weighted_cross_entropy_with_logits(
+          labels=triple_labels,
+          logits=triple_logits,
+          pos_weight=hparams.pos_weight)
+      avg_triple_loss = tf.reduce_mean(triple_losses)
+      tf.summary.scalar("triple_loss", avg_triple_loss)
+
+    return triple_logits, avg_triple_loss, original_knowledge_encoder_output
+
+  def body(self, features):
+    """Transformer main model_fn.
+
+    Args:
+      features: Map of features to the model. Should contain the following:
+          "inputs": Transformer inputs [batch_size, input_length, hidden_dim]
+          "targets": Target decoder outputs. [batch_size, decoder_length,
+            hidden_dim]
+          "target_space_id": A scalar int from data_generators.problem.SpaceID.
+
+    Returns:
+      Final decoder representation. [batch_size, decoder_length, hidden_dim]
+    """
+    tf.logging.info("Using PgScratch BODY function.")
+    hparams = self._hparams
+
+    losses = {}
+    inputs = features["inputs"]
+    target_space = features["target_space_id"]
+    # encoder_output: <tf.float32>[batch_size, input_length, hidden_dim]
+    # encoder_decoder_attention_bias: <tf.float32>[batch_size, input_length]
+    encoder_output, encoder_decoder_attention_bias = self.encode(
+        inputs, target_space, hparams, features=features, losses=losses)
+
+    with tf.variable_scope("knowledge"):
+      with tf.name_scope("knowledge_encoding"):
+        # Encode knowledge.
+        # <tf.float32>[batch_size, max_triple_num, emb_dim]
+        fact_embedding, fact_lengths = self.encode_knowledge_bottom(features)
+        tf.logging.info("Encoded knowledge")
+
+      with tf.name_scope("knowledge_selection_and_loss"):
+        # Compute knowledge selection and loss.
+        triple_logits, avg_triple_selection_loss, knowledge_encoder_output = self.compute_knowledge_selection_and_loss(
+            features, encoder_output, fact_embedding, fact_lengths)
+        losses["kb_loss"] = avg_triple_selection_loss
+
+    if hparams.attend_kb:
+      tf.logging.info("ATTEND_KB is ACTIVE")
+      with tf.name_scope("knowledge_attention"):
+
+        knowledge_padding = tf.zeros_like(triple_logits, dtype=tf.float32)
+        knowledge_attention_bias = common_attention.attention_bias_ignore_padding(
+            knowledge_padding)
+        encoder_output = tf.concat([knowledge_encoder_output, encoder_output],
+                                   1)
+        encoder_decoder_attention_bias = tf.concat(
+            [knowledge_attention_bias, encoder_decoder_attention_bias], -1)
+
+    else:
+      tf.logging.info("ATTEND_KB is INACTIVE")
+
+    targets = features["targets"]
+    targets_shape = common_layers.shape_list(targets)
+    targets = common_layers.flatten4d3d(targets)
+
+    (decoder_input,
+     decoder_self_attention_bias) = transformer.transformer_prepare_decoder(
+         targets, hparams, features=features)
+
+    decode_kwargs = {}
+    decoder_output = self.decode(
+        decoder_input,
+        encoder_output,
+        encoder_decoder_attention_bias,
+        decoder_self_attention_bias,
+        hparams,
+        nonpadding=transformer.features_to_nonpadding(features, "targets"),
+        losses=losses,
+        **decode_kwargs)
+
+    expected_attentions = features.get("expected_attentions")
+    if expected_attentions is not None:
+      attention_loss = common_attention.encoder_decoder_attention_loss(
+          expected_attentions, self.attention_weights,
+          hparams.expected_attention_loss_type,
+          hparams.expected_attention_loss_multiplier)
+      return decoder_output, {"attention_loss": attention_loss}
+
+    ret = tf.reshape(decoder_output, targets_shape)
+    if losses:
+      return ret, losses
+    else:
+      return ret
+
+  def _normalize_body_output(self, body_out):
+    if len(body_out) == 2:
+      output, losses = body_out
+      if not isinstance(losses, dict):
+        losses = {"extra": tf.reduce_mean(losses)}
+    else:
+      output = body_out
+      losses = {"extra": 0.0}
+
+    return output, losses
+
+  def _beam_decode(self,
+                   features,
+                   decode_length,
+                   beam_size,
+                   top_beams,
+                   alpha,
+                   use_tpu=False):
+    """Beam search decoding.
+
+    Args:
+      features: an map of string to `Tensor`
+      decode_length: an integer.  How many additional timesteps to decode.
+      beam_size: number of beams.
+      top_beams: an integer. How many of the beams to return.
+      alpha: Float that controls the length penalty. larger the alpha, stronger
+        the preference for longer translations.
+      use_tpu: A bool, whether to do beam decode on TPU.
+
+    Returns:
+      A dict of decoding results {
+          "outputs": integer `Tensor` of decoded ids of shape
+              [batch_size, <= decode_length] if beam_size == 1 or
+              [batch_size, top_beams, <= decode_length]
+          "scores": decoding log probs from the beam search,
+              None if using greedy decoding (beam_size=1)
+      }
+    """
+    return super(transformer.Transformer,
+                 self)._beam_decode_slow(features, decode_length, beam_size,
+                                         top_beams, alpha, use_tpu)
+
+  def _greedy_infer(self, features, decode_length, use_tpu=False):
+    """Fast version of greedy decoding.
+
+    Args:
+      features: an map of string to `Tensor`
+      decode_length: an integer.  How many additional timesteps to decode.
+      use_tpu: A bool. Whether to build the inference graph for TPU.
+
+    Returns:
+      A dict of decoding results {
+          "outputs": integer `Tensor` of decoded ids of shape
+              [batch_size, <= decode_length] if beam_size == 1 or
+              [batch_size, top_beams, <= decode_length]
+          "scores": decoding log probs from the beam search,
+              None if using greedy decoding (beam_size=1)
+      }
+
+    Raises:
+      NotImplementedError: If there are multiple data shards.
+    """
+    return super(transformer.Transformer,
+                 self)._greedy_infer(features, decode_length)
+
+
+def compute_last_embedding(input_embeddings, input_lengths, hparams):
+  """Computes average of last K embedding.
+
+  Args:
+    input_embeddings: <tf.float32>[bs, max_seq_len, emb_dim]
+    input_lengths: <tf.int64>[bs, 1]
+    hparams: model hparams
+
+  Returns:
+    last_k_embedding: <tf.float32>[bs, emb_dim]
+  """
+  max_seq_len = tf.shape(input_embeddings)[1]
+  # <tf.float32>[bs, 1, max_seq_len]
+  mask = tf.sequence_mask(input_lengths, max_seq_len, dtype=tf.float32)
+  del_mask = tf.sequence_mask(
+      input_lengths - hparams.last_k, max_seq_len, dtype=tf.float32)
+  final_mask = mask - del_mask
+  # <tf.float32>[bs, 1, emb_dim]
+  sum_embedding = tf.matmul(final_mask, input_embeddings)
+  # <tf.float32>[bs, 1, emb_dim]
+  last_k_embedding = sum_embedding / tf.to_float(
+      tf.expand_dims(
+          tf.ones([tf.shape(input_embeddings)[0], 1]) * hparams.last_k, 2))
+  # <tf.float32>[bs, dim]
+  return tf.squeeze(last_k_embedding, 1)
+
+
+def compute_max_pool_embedding(input_embeddings, input_lengths):
+  """Computes max pool embedding.
+
+  Args:
+    input_embeddings: <tf.float32>[bs, max_seq_len, emb_dim]
+    input_lengths: <tf.int64>[bs, 1]
+
+  Returns:
+    max_pool_embedding: <tf.float32>[bs, emb_dim]
+  """
+  max_seq_len = tf.shape(input_embeddings)[1]
+  # <tf.float32>[bs, max_seq_len]
+  mask = 1.0 - tf.sequence_mask(input_lengths, max_seq_len, dtype=tf.float32)
+  mask = tf.squeeze(mask * (-1e-6), 1)
+  mask = tf.expand_dims(mask, 2)
+  # <tf.float32>[bs, emb_dim]
+  max_pool_embedding = tf.reduce_max(input_embeddings + mask, 1)
+  # <tf.float32>[bs, dim]
+  return max_pool_embedding
+
+
+def compute_average_embedding(input_embeddings, input_lengths):
+  """Computes bag-of-words embedding.
+
+  Args:
+    input_embeddings: <tf.float32>[bs, max_seq_len, emb_dim]
+    input_lengths: <tf.int64>[bs, 1]
+
+  Returns:
+    bow_embedding: <tf.float32>[bs, emb_dim]
+  """
+  max_seq_len = tf.shape(input_embeddings)[1]
+  # <tf.float32>[bs, 1, max_seq_len]
+  mask = tf.sequence_mask(input_lengths, max_seq_len, dtype=tf.float32)
+  # <tf.float32>[bs, 1, emb_dim]
+  sum_embedding = tf.matmul(mask, input_embeddings)
+  # <tf.float32>[bs, 1, emb_dim]
+  avg_embedding = sum_embedding / tf.to_float(tf.expand_dims(input_lengths, 2))
+  # <tf.float32>[bs, dim]
+  return tf.squeeze(avg_embedding, 1)
+
+
+def compute_summary_embedding(input_embeddings, input_lengths, hparams):
+  """Convert list of embedding to single embedding.
+
+  Args:
+    input_embeddings: <tf.float32>[bs, max_seq_len, emb_dim]
+    input_lengths: <tf.int64>[bs, 1]
+    hparams: model hparams
+
+  Returns:
+    embedding: <tf.float32>[bs, emb_dim]
+  """
+  if hparams.pool_technique == "average":
+    return compute_average_embedding(input_embeddings, input_lengths)
+  elif hparams.pool_technique == "max_pool":
+    return compute_max_pool_embedding(input_embeddings, input_lengths)
+  elif hparams.pool_technique == "last":
+    return compute_last_embedding(input_embeddings, input_lengths, hparams)
+
+
+@registry.register_hparams
+def neural_assistant_tiny():
+  """HParams for tiny neural_assistant model."""
+  hparams = transformer.transformer_tiny_tpu()
+  hparams.add_hparam("pos_weight", 1.0)  # weight for positive triples
+  hparams.add_hparam("similarity_fuction",
+                     "bilinear")  # dot_product or bilinear
+  hparams.add_hparam("pool_technique", "average")  # avg or max pool or last
+  hparams.add_hparam("last_k", 1)  # number of last indices for averaging
+  hparams.add_hparam("max_triple_length", 30)  # max length of every triple
+  hparams.add_hparam("train_triple_num",
+                     5000)  # max number of triples during training
+  hparams.add_hparam("attend_kb", True)  # if False, it's a transformer model
+  hparams.add_hparam("kb_loss_weight", 0.0)  # weight for distant supervision
+  hparams.add_hparam("test_triple_num",
+                     28483)  # max triples of KB
+  return hparams
+
+
+@registry.register_hparams
+def neural_assistant_tiny_ds():
+  """HParams for tiny neural_assistant model with distant supervision loss."""
+  hparams = neural_assistant_tiny()
+  hparams.kb_loss_weight = 0.2
+  return hparams

From ab918e0d9592394614aa2e10cfc8f23e8cb24dfc Mon Sep 17 00:00:00 2001
From: Blake Hechtman <blakehechtman@google.com>
Date: Tue, 29 Oct 2019 09:49:50 -0700
Subject: [PATCH 2557/2720] [TENSOR2TENSOR] Use an einsum instead of matmul.

PiperOrigin-RevId: 277303581
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 4f6deb6aa..f3c43443e 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1572,7 +1572,7 @@ def dot_product_attention(q,
   """
   with tf.variable_scope(
       name, default_name="dot_product_attention", values=[q, k, v]) as scope:
-    logits = tf.einsum("...kd,...qd->...qk", k, q)
+    logits = tf.matmul(q, k, transpose_b=True)  # [..., length_q, length_kv]
     if bias is not None:
       bias = common_layers.cast_like(bias, logits)
       logits += bias

From bcc430637dd542be944a5b91923854db9898a9df Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 31 Oct 2019 10:58:10 -0700
Subject: [PATCH 2558/2720] Fix deprecation warning.

PiperOrigin-RevId: 277757850
---
 tensor2tensor/layers/common_layers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 478c0d7ce..8eb6a8d19 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2996,7 +2996,7 @@ def ones_matrix_band_part(rows, cols, num_lower, num_upper, out_shape=None):
       band = band.reshape(out_shape)
     band = tf.constant(band, tf.float32)
   else:
-    band = tf.matrix_band_part(
+    band = tf.linalg.band_part(
         tf.ones([rows, cols]), tf.cast(num_lower, tf.int64),
         tf.cast(num_upper, tf.int64))
     if out_shape:

From 5185a4008fa5c4141d2372e800f82d7ffec9ead0 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 1 Nov 2019 06:58:14 -0700
Subject: [PATCH 2559/2720] Fix typo.

PiperOrigin-RevId: 277919398
---
 tensor2tensor/data_generators/problem.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 68745257c..15da7648d 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -395,7 +395,7 @@ def set_task_id(self, new_task_id):
   def preprocess(self, dataset, mode, hparams, interleave=True):
     """Runtime preprocessing on the whole dataset.
 
-    Return a tf.data.Datset -- the preprocessed version of the given one.
+    Return a tf.data.Dataset -- the preprocessed version of the given one.
     By default this function calls preprocess_example.
 
     Args:

From d63603265ab63e997525f143526bf6807d67c67f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 1 Nov 2019 14:23:29 -0700
Subject: [PATCH 2560/2720] Adding TPU and tiny hparams sets for
 evolved_transformer

PiperOrigin-RevId: 278000866
---
 tensor2tensor/models/evolved_transformer.py   | 25 +++++++++
 .../utils/partial_checkpoint_load_hook.py     | 55 +++++++++++++++++++
 2 files changed, 80 insertions(+)
 create mode 100644 tensor2tensor/utils/partial_checkpoint_load_hook.py

diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 56a72b832..5529fb03c 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -765,6 +765,15 @@ def add_evolved_transformer_hparams(hparams):
   return hparams
 
 
+@registry.register_hparams
+def evolved_transformer_tiny():
+  """Base parameters for Evolved Transformer model."""
+  hparams = add_evolved_transformer_hparams(transformer.transformer_tiny())
+  hparams.learning_rate_schedule = (
+      "constant*single_cycle_cos_decay")
+  return hparams
+
+
 @registry.register_hparams
 def evolved_transformer_base():
   """Base parameters for Evolved Transformer model."""
@@ -805,3 +814,19 @@ def evolved_transformer_big_tpu():
   hparams.learning_rate_schedule = (
       "constant*single_cycle_cos_decay")
   return hparams
+
+
+@registry.register_hparams
+def evolved_transformer_tpu_basic():
+  """Basic Seq2Seq TPU hyper-parameters."""
+  hparams = transformer.transformer_big_tpu()
+  hparams.add_hparam("print_vars", False)
+  hparams.batch_size = 8192
+  hparams.max_length = 256
+
+  # N < 0 means all weights in the model are trainable.
+  # N >= 0 means all weights are frozen except N top decoder layers +
+  # (pre-)softmax matrix (that projects from hidden size to vocab size).
+  hparams.add_hparam("num_trainable_top_decoder_layers", -1)
+
+  return hparams
diff --git a/tensor2tensor/utils/partial_checkpoint_load_hook.py b/tensor2tensor/utils/partial_checkpoint_load_hook.py
new file mode 100644
index 000000000..70b6c566b
--- /dev/null
+++ b/tensor2tensor/utils/partial_checkpoint_load_hook.py
@@ -0,0 +1,55 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Hook to partially load a checkpoint."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+
+class PartialCheckpointLoad(tf.train.SessionRunHook):
+  """Partially load train_variables from a checkpoint.
+
+  Hook used to load each variable saved in checkpoint into the graph. It
+  will ignore any additional variables present in the graph that are not
+  saved in the checkpoint. (Note: The loaded variables include ADAM/training
+  variables, if they exist in the checkpoint)
+  Can perform mapping if the base scopename for graph variables is different
+  from the checkpoint variables.
+  """
+
+  def __init__(self, hook_context, chk_scopename, graph_scopename):
+    """Initialize the hook with chkp directory and scopenames.
+
+    Args:
+      hook_context: HookContext object containing hparams.
+      chk_scopename: Base scopename of variables in the checkpoint being loaded
+      graph_scopename: Base scopename of variables in current graph
+    """
+    self.checkpoint_path = hook_context.hparams.partial_load_checkpoint
+    self.chk_scopename = chk_scopename
+    self.graph_scopename = graph_scopename
+
+  def begin(self):
+    # TODO(karishmamalkan): Add logging for when variables are loaded
+    variable_references = {var.name: var for var in tf.all_variables()}
+    variable_mappings = {}
+    vars_in_chk = tf.train.list_variables(self.checkpoint_path)
+    for (var, _) in vars_in_chk:
+      variable_mappings[var] = variable_references[
+          var.replace(self.chk_scopename, self.graph_scopename) + ":0"]
+    tf.train.init_from_checkpoint(self.checkpoint_path, variable_mappings)

From 06bc1eb47ba85d8b08b08157163fc75c93072800 Mon Sep 17 00:00:00 2001
From: Daniel De Freitas Adiwardana <adiwardana@google.com>
Date: Fri, 1 Nov 2019 14:32:27 -0700
Subject: [PATCH 2561/2720] Enabling t2t_decoder to read hparams from
 output_dir/hparams.json when available.

PiperOrigin-RevId: 278002517
---
 tensor2tensor/bin/t2t_decoder.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 3f9a56000..2013d8218 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -63,11 +63,15 @@
 
 
 def create_hparams():
+  hparams_path = None
+  if FLAGS.output_dir:
+    hparams_path = os.path.join(FLAGS.output_dir, "hparams.json")
   return trainer_lib.create_hparams(
       FLAGS.hparams_set,
       FLAGS.hparams,
       data_dir=os.path.expanduser(FLAGS.data_dir),
-      problem_name=FLAGS.problem)
+      problem_name=FLAGS.problem,
+      hparams_path=hparams_path)
 
 
 def create_decode_hparams():

From 2eebc98de294086a97440a44be2dadb05668e533 Mon Sep 17 00:00:00 2001
From: Jacob Burnim <jburnim@google.com>
Date: Fri, 1 Nov 2019 19:57:58 -0700
Subject: [PATCH 2562/2720] Update deprecated calls of tfp.distributions.* to
 tfp.math*.

PiperOrigin-RevId: 278049446
---
 tensor2tensor/layers/latent_layers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 2788f5e8f..77c51a0bb 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -733,7 +733,7 @@ def iaf_flow(one_hot_assignments,
     padded_assignments = tf.pad(
         one_hot_assignments, [[0, 0], [0, 0], [1, 0], [0, 0]])[:, :, :-1, :]
     scale_bijector = tfp.distributions.bijectors.Affine(
-        scale_tril=tfp.distributions.fill_triangular(scale_weights))
+        scale_tril=tfp.math.fill_triangular(scale_weights))
     scale = scale_bijector.forward(
         tf.transpose(padded_assignments, [0, 1, 3, 2]))
     # Transpose the bijector output since it performs a batch matmul.

From 9ed57fadc5df6a67bd2ddaf004e0e1dbd89c85ac Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 4 Nov 2019 09:17:25 -0800
Subject: [PATCH 2563/2720] Add metrics for TWO_CLASS_ACCURACY,
 TWO_CLASS_LOG_LIKELIHOOD, and UNPADDED_MSE.

PiperOrigin-RevId: 278394549
---
 tensor2tensor/utils/metrics.py      | 52 +++++++++++++++++++++++++++++
 tensor2tensor/utils/metrics_test.py | 33 ++++++++++++++++++
 2 files changed, 85 insertions(+)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 661566de1..40fff439f 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -45,6 +45,7 @@ class Metrics(object):
   APPROX_BLEU = "approx_bleu_score"
   APPROX_SARI = "approx_sari_score"
   RMSE = "rmse"
+  UNPADDED_MSE = "unpadded_mse"
   LOG_POISSON = "log_poisson"
   PEARSON = "pearson"
   R2 = "r_squared"
@@ -61,6 +62,8 @@ class Metrics(object):
   SIGMOID_RECALL_ONE_HOT = "sigmoid_recall_one_hot"
   SIGMOID_PRECISION_ONE_HOT = "sigmoid_precision_one_hot"
   SIGMOID_CROSS_ENTROPY_ONE_HOT = "sigmoid_cross_entropy_one_hot"
+  TWO_CLASS_ACCURACY = "two_class_accuracy"
+  TWO_CLASS_LOG_LIKELIHOOD = "two_class_log_likelihood"
   ROC_AUC = "roc_auc"
   IMAGE_SUMMARY = "image_summary"
   DMOL_PERPLEXITY = "disc_mol_neg_log_perplexity"
@@ -87,6 +90,15 @@ def padded_rmse(predictions, labels, weights_fn=common_layers.weights_all):
   return error_sqrt, tf.reduce_sum(weights)
 
 
+def unpadded_mse(predictions, labels, weights_fn=common_layers.weights_all):
+  predictions = tf.to_float(predictions)
+  labels = tf.to_float(labels)
+  weights = weights_fn(labels)
+  error = tf.pow(predictions - labels, 2)
+  mean_error = tf.reduce_mean(error * weights)
+  return mean_error, tf.reduce_sum(weights)
+
+
 def abs_error(predictions, labels, weights_fn=None):
   """Computes mean(abs(preds-target))."""
   del weights_fn  # Unused
@@ -163,6 +175,43 @@ def rounding_sequence_accuracy(predictions,
   return correct_seq, tf.constant(1.0)
 
 
+def two_class_accuracy(predictions, labels, weights_fn=None):
+  """Accuracy for two class classification with 0/1 labels."""
+  with tf.variable_scope("two_class_accuracy", values=[predictions, labels]):
+    del weights_fn
+    hard_predictions = tf.to_int32(tf.math.round(tf.squeeze(predictions)))
+    int_labels = tf.to_int32(labels)
+    _, accuracy = tf.metrics.accuracy(labels=int_labels,
+                                      predictions=hard_predictions)
+    return accuracy, tf.constant(1.0)
+
+
+def two_class_log_likelihood(predictions, labels, weights_fn=None):
+  """Log-likelihood for two class classification with 0/1 labels.
+
+  Args:
+    predictions: A float valued tensor of shape [`batch_size`].  Each
+      component should be between 0 and 1.
+    labels: An int valued tensor of shape [`batch_size`].  Each component
+      should either be 0 or 1.
+    weights_fn: unused.
+
+  Returns:
+    A pair, with the average log likelihood in the first component.
+  """
+  del weights_fn
+  float_labels = tf.cast(labels, dtype=tf.float64)
+  # likelihood should be just p for class 1, and 1 - p for class 0.
+  # signs is 1 for class 1, and -1 for class 0
+  signs = 2 * float_labels - tf.ones_like(float_labels)
+  # constant_term is 1 for class 0, and 0 for class 1.
+  constant_term = tf.ones_like(float_labels) - float_labels
+  likelihoods = constant_term + signs * tf.squeeze(predictions)
+  log_likelihoods = tf.log(likelihoods)
+  avg_log_likelihood = tf.reduce_mean(log_likelihoods)
+  return avg_log_likelihood, tf.constant(1.0)
+
+
 def padded_sequence_accuracy(predictions,
                              labels,
                              weights_fn=common_layers.weights_nonzero):
@@ -847,6 +896,7 @@ def pearson_correlation_coefficient(predictions, labels, weights_fn=None):
     Metrics.APPROX_BLEU: bleu_hook.bleu_score,
     Metrics.APPROX_SARI: sari_hook.sari_score,
     Metrics.RMSE: padded_rmse,
+    Metrics.UNPADDED_MSE: unpadded_mse,
     Metrics.LOG_POISSON: padded_log_poisson,
     Metrics.PEARSON: pearson_correlation_coefficient,
     Metrics.R2: padded_variance_explained,
@@ -861,6 +911,8 @@ def pearson_correlation_coefficient(predictions, labels, weights_fn=None):
     Metrics.SIGMOID_CROSS_ENTROPY_ONE_HOT: sigmoid_cross_entropy_one_hot,
     Metrics.SET_PRECISION: set_precision,
     Metrics.SET_RECALL: set_recall,
+    Metrics.TWO_CLASS_ACCURACY: two_class_accuracy,
+    Metrics.TWO_CLASS_LOG_LIKELIHOOD: two_class_log_likelihood,
     Metrics.ROC_AUC: roc_auc,
     Metrics.IMAGE_SUMMARY: image_summary,
     Metrics.DMOL_PERPLEXITY: dmol_neg_log_perplexity,
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 5e0037385..05c62b6c2 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -90,6 +90,27 @@ def testSequenceAccuracyMetric(self):
       actual = session.run(a)
     self.assertEqual(actual, expected)
 
+  def testTwoClassAccuracyMetric(self):
+    predictions = tf.constant([0.0, 0.2, 0.4, 0.6, 0.8, 1.0], dtype=tf.float32)
+    targets = tf.constant([0, 0, 1, 0, 1, 1], dtype=tf.int32)
+    expected = 2.0 / 3.0
+    with self.test_session() as session:
+      accuracy, _ = metrics.two_class_accuracy(predictions, targets)
+      session.run(tf.global_variables_initializer())
+      session.run(tf.local_variables_initializer())
+      actual = session.run(accuracy)
+    self.assertAlmostEqual(actual, expected)
+
+  def testTwoClassLogLikelihood(self):
+    predictions = np.array([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
+    targets = np.array([0, 0, 1, 0, 1, 1])
+    expected = (2.0 * np.log(0.8) + 2.0 * np.log(0.4)) / 6.0
+    with self.test_session() as session:
+      avg_log_likelihood, _ = metrics.two_class_log_likelihood(
+          predictions, targets)
+      actual = session.run(avg_log_likelihood)
+    self.assertAlmostEqual(actual, expected)
+
   def testRMSEMetric(self):
     predictions = np.full((10, 1), 1)  # All 1's
     targets = np.full((10, 1), 3)  # All 3's
@@ -102,6 +123,18 @@ def testRMSEMetric(self):
       actual = session.run(rmse)
     self.assertEqual(actual, expected)
 
+  def testUnpaddedRMSEMetric(self):
+    predictions = np.full((10, 1), 1)  # All 1's
+    targets = np.full((10, 1), 3)  # All 3's
+    expected = np.mean((predictions - targets)**2)  # MSE = 4.0
+    with self.test_session() as session:
+      mse, _ = metrics.unpadded_mse(
+          tf.constant(predictions, dtype=tf.int32),
+          tf.constant(targets, dtype=tf.int32))
+      session.run(tf.global_variables_initializer())
+      actual = session.run(mse)
+    self.assertEqual(actual, expected)
+
   def testSequenceEditDistanceMetric(self):
     predictions = np.array([[3, 4, 5, 1, 0, 0],
                             [2, 1, 3, 4, 0, 0],

From b679a887a762da7509fd7dcc9902a9da63d7e782 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 4 Nov 2019 12:19:42 -0800
Subject: [PATCH 2564/2720] Prevent float-type mismatches by converting
 predictions to float64.  (The labels had already been so converted.)

PiperOrigin-RevId: 278437095
---
 tensor2tensor/utils/metrics.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 40fff439f..fcdf1aab2 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -201,12 +201,13 @@ def two_class_log_likelihood(predictions, labels, weights_fn=None):
   """
   del weights_fn
   float_labels = tf.cast(labels, dtype=tf.float64)
+  float_predictions = tf.cast(tf.squeeze(predictions), dtype=tf.float64)
   # likelihood should be just p for class 1, and 1 - p for class 0.
   # signs is 1 for class 1, and -1 for class 0
   signs = 2 * float_labels - tf.ones_like(float_labels)
   # constant_term is 1 for class 0, and 0 for class 1.
   constant_term = tf.ones_like(float_labels) - float_labels
-  likelihoods = constant_term + signs * tf.squeeze(predictions)
+  likelihoods = constant_term + signs * float_predictions
   log_likelihoods = tf.log(likelihoods)
   avg_log_likelihood = tf.reduce_mean(log_likelihoods)
   return avg_log_likelihood, tf.constant(1.0)

From 8c0736a82e343f7c0cb75ebb2d2aeff57feecaa7 Mon Sep 17 00:00:00 2001
From: Marc van Zee <marcvanzee@google.com>
Date: Tue, 5 Nov 2019 09:04:51 -0800
Subject: [PATCH 2565/2720] Fix bug in universal transformer hyperparameter
 range.

PiperOrigin-RevId: 278638337
---
 tensor2tensor/models/research/universal_transformer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 185679cc0..28040584e 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -811,7 +811,7 @@ def universal_transformer_base_range(rhp):
   rhp.set_discrete("hidden_size", [1024, 2048, 4096])
   rhp.set_discrete("filter_size", [2048, 4096, 8192])
   rhp.set_discrete("num_heads", [8, 16, 32])
-  rhp.set_discrete("transformer_ffn_type", ["sepconv", "fc"])
+  rhp.set_categorical("transformer_ffn_type", ["sepconv", "fc"])
   rhp.set_float("learning_rate", 0.3, 3.0, scale=rhp.LOG_SCALE)
   rhp.set_float("weight_decay", 0.0, 2.0)
 
@@ -825,6 +825,6 @@ def adaptive_universal_transformer_base_range(rhp):
   rhp.set_discrete("hidden_size", [1024, 2048, 4096])
   rhp.set_discrete("filter_size", [2048, 4096, 8192])
   rhp.set_discrete("num_heads", [8, 16, 32])
-  rhp.set_discrete("transformer_ffn_type", ["sepconv", "fc"])
+  rhp.set_categorical("transformer_ffn_type", ["sepconv", "fc"])
   rhp.set_float("learning_rate", 0.3, 3.0, scale=rhp.LOG_SCALE)
   rhp.set_float("weight_decay", 0.0, 2.0)

From 4d96546af045b869e3a3f826a21b634e9631e43d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 11 Nov 2019 13:46:00 -0800
Subject: [PATCH 2566/2720] Enabled encoder-decoder attention for the
 multi-encoder case.

PiperOrigin-RevId: 279807999
---
 tensor2tensor/models/transformer.py | 71 +++++++++++++++--------------
 1 file changed, 37 insertions(+), 34 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 75fd1826f..eef24dba1 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1491,41 +1491,44 @@ def transformer_self_attention_layer(decoder_input,
               tf.estimator.ModeKeys.TRAIN) == tf.estimator.ModeKeys.TRAIN))
       x = common_layers.layer_postprocess(x, y, hparams)
     if encoder_output is not None:
+      if not isinstance(encoder_output, (list,)):
+        encoder_output = [encoder_output]
       with tf.variable_scope("encdec_attention"):
-        y = common_attention.multihead_attention(
-            common_layers.layer_preprocess(
-                x, hparams, layer_collection=layer_collection),
-            encoder_output,
-            encoder_decoder_attention_bias,
-            hparams.attention_key_channels or hparams.hidden_size,
-            hparams.attention_value_channels or hparams.hidden_size,
-            hparams.hidden_size,
-            hparams.num_heads,
-            hparams.attention_dropout,
-            max_relative_position=hparams.max_relative_position,
-            heads_share_relative_embedding=(
-                hparams.heads_share_relative_embedding),
-            add_relative_to_values=hparams.add_relative_to_values,
-            save_weights_to=save_weights_to,
-            cache=layer_cache,
-            make_image_summary=make_image_summary,
-            dropout_broadcast_dims=attention_dropout_broadcast_dims,
-            max_length=hparams.get("max_length"),
-            vars_3d=hparams.get("attention_variables_3d"),
-            activation_dtype=hparams.get("activation_dtype", "float32"),
-            weight_dtype=hparams.get("weight_dtype", "float32"),
-            layer_collection=layer_collection,
-            hard_attention_k=hparams.get("hard_attention_k", 0),
-            gumbel_noise_weight=hparams.get("gumbel_noise_weight", 0.0),
-            max_area_width=max_area_width,
-            max_area_height=max_area_height,
-            memory_height=memory_height,
-            area_key_mode=hparams.get("area_key_mode", "none"),
-            area_value_mode=hparams.get("area_value_mode", "none"),
-            training=(hparams.get(
-                "mode",
-                tf.estimator.ModeKeys.TRAIN) == tf.estimator.ModeKeys.TRAIN))
-        x = common_layers.layer_postprocess(x, y, hparams)
+        for enc_output in encoder_output:
+          y = common_attention.multihead_attention(
+              common_layers.layer_preprocess(
+                  x, hparams, layer_collection=layer_collection),
+              enc_output,
+              encoder_decoder_attention_bias,
+              hparams.attention_key_channels or hparams.hidden_size,
+              hparams.attention_value_channels or hparams.hidden_size,
+              hparams.hidden_size,
+              hparams.num_heads,
+              hparams.attention_dropout,
+              max_relative_position=hparams.max_relative_position,
+              heads_share_relative_embedding=(
+                  hparams.heads_share_relative_embedding),
+              add_relative_to_values=hparams.add_relative_to_values,
+              save_weights_to=save_weights_to,
+              cache=layer_cache,
+              make_image_summary=make_image_summary,
+              dropout_broadcast_dims=attention_dropout_broadcast_dims,
+              max_length=hparams.get("max_length"),
+              vars_3d=hparams.get("attention_variables_3d"),
+              activation_dtype=hparams.get("activation_dtype", "float32"),
+              weight_dtype=hparams.get("weight_dtype", "float32"),
+              layer_collection=layer_collection,
+              hard_attention_k=hparams.get("hard_attention_k", 0),
+              gumbel_noise_weight=hparams.get("gumbel_noise_weight", 0.0),
+              max_area_width=max_area_width,
+              max_area_height=max_area_height,
+              memory_height=memory_height,
+              area_key_mode=hparams.get("area_key_mode", "none"),
+              area_value_mode=hparams.get("area_value_mode", "none"),
+              training=(hparams.get(
+                  "mode",
+                  tf.estimator.ModeKeys.TRAIN) == tf.estimator.ModeKeys.TRAIN))
+          x = common_layers.layer_postprocess(x, y, hparams)
     return x, layer_cache
 
 

From f7d93b76d931e1a8a7cf9953a71db1d89bec3ab3 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Sun, 17 Nov 2019 18:41:35 -0800
Subject: [PATCH 2567/2720] More work on Trax imports: add optimizers and
 Trainer.

PiperOrigin-RevId: 280977657
---
 README.md           | 10 +++-------
 docs/walkthrough.md | 10 +++-------
 2 files changed, 6 insertions(+), 14 deletions(-)

diff --git a/README.md b/README.md
index 3ad60b8f4..4e704b130 100644
--- a/README.md
+++ b/README.md
@@ -16,14 +16,10 @@ welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CO
 of deep learning models and datasets designed to make deep learning more
 accessible and [accelerate ML
 research](https://research.googleblog.com/2017/06/accelerating-deep-learning-research.html).
-T2T is actively used and maintained by researchers and engineers within the
+T2T was developed by researchers and engineers in the
 [Google Brain team](https://research.google.com/teams/brain/) and a community
-of users. We're eager to collaborate with you too, so feel free to
-[open an issue on GitHub](https://github.com/tensorflow/tensor2tensor/issues)
-or send along a pull request (see [our contribution doc](CONTRIBUTING.md)).
-You can chat with us on
-[Gitter](https://gitter.im/tensor2tensor/Lobby) and join the
-[T2T Google Group](https://groups.google.com/forum/#!forum/tensor2tensor).
+of users. It is now in maintenance mode &mdash; we keep it running and welcome
+bug-fixes, but encourage users to use the successor library [Trax](https://github.com/google/trax).
 
 ### Quick Start
 
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 3ad60b8f4..4e704b130 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -16,14 +16,10 @@ welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CO
 of deep learning models and datasets designed to make deep learning more
 accessible and [accelerate ML
 research](https://research.googleblog.com/2017/06/accelerating-deep-learning-research.html).
-T2T is actively used and maintained by researchers and engineers within the
+T2T was developed by researchers and engineers in the
 [Google Brain team](https://research.google.com/teams/brain/) and a community
-of users. We're eager to collaborate with you too, so feel free to
-[open an issue on GitHub](https://github.com/tensorflow/tensor2tensor/issues)
-or send along a pull request (see [our contribution doc](CONTRIBUTING.md)).
-You can chat with us on
-[Gitter](https://gitter.im/tensor2tensor/Lobby) and join the
-[T2T Google Group](https://groups.google.com/forum/#!forum/tensor2tensor).
+of users. It is now in maintenance mode &mdash; we keep it running and welcome
+bug-fixes, but encourage users to use the successor library [Trax](https://github.com/google/trax).
 
 ### Quick Start
 

From db0007405dec43f3cfe9794fa555e78924aa53e9 Mon Sep 17 00:00:00 2001
From: Mohammad Taghi Saffar <msaffar@google.com>
Date: Mon, 18 Nov 2019 17:15:20 -0800
Subject: [PATCH 2568/2720] Fix features for decoding on TPUs.

PiperOrigin-RevId: 281188677
---
 tensor2tensor/utils/t2t_model.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 7acba41e3..fe33e6315 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1462,7 +1462,9 @@ def estimator_model_fn(cls,
       if use_tpu:
         inputs = features.get("inputs")
         if inputs is None:
-          inputs = features["targets"]
+          inputs = features.get("targets")
+        if inputs is None:
+          inputs = features["infer_targets"]
         shape = inputs.get_shape().as_list()
         if shape[0] is None:
           shape[0] = decode_hparams.batch_size or hparams.batch_size

From b3610a221f614c6c5cd5f11941d3c0775a48716f Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 20 Nov 2019 21:16:23 -0800
Subject: [PATCH 2569/2720]  - Bump up T2T version to 1.15.0 in lockstep with
 TF 1.15  - Add tf_slim as a requirement, for tf.contrib.slim replacement.  -
 Delete `attacks` and `pruning_strategies` they are aliases of `attack` and
 `pruning_strategy` respectively.

PiperOrigin-RevId: 281670284
---
 .travis.yml                                     |  4 ++--
 oss_scripts/oss_pip_install.sh                  |  2 +-
 setup.py                                        | 10 +++-------
 tensor2tensor/data_generators/allen_brain.py    |  3 ++-
 .../data_generators/bair_robot_pushing.py       |  7 ++++---
 tensor2tensor/data_generators/fsns.py           |  3 ++-
 tensor2tensor/data_generators/gym_env.py        |  3 ++-
 tensor2tensor/data_generators/image_utils.py    |  7 ++++---
 tensor2tensor/data_generators/moving_mnist.py   |  3 ++-
 tensor2tensor/data_generators/problem.py        | 11 ++++++-----
 tensor2tensor/data_generators/translate.py      |  5 +++--
 .../data_generators/video_generated.py          |  3 ++-
 tensor2tensor/data_generators/video_utils.py    |  7 ++++---
 tensor2tensor/data_generators/vqa.py            | 17 +++++++++--------
 tensor2tensor/envs/env_problem.py               |  3 ++-
 tensor2tensor/envs/rendered_env_problem.py      |  3 ++-
 tensor2tensor/utils/data_reader.py              |  3 ++-
 tensor2tensor/utils/registry.py                 |  7 -------
 18 files changed, 52 insertions(+), 49 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 45676b315..3f9d78e3b 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -14,11 +14,11 @@ env:
     - T2T_PROBLEM=algorithmic_reverse_binary40_test
     - T2T_DATA_DIR=/tmp/t2t-data
     - T2T_TRAIN_DIR=/tmp/t2t-train
-    - TF_LATEST="1.14.*"
+    - TF_LATEST="1.15.*"
     # This is necessary to have gsutil work with Python 2.7
     - BOTO_CONFIG=/dev/null
   matrix:
-    - TF_VERSION="1.14.*"
+    - TF_VERSION="1.15.*"
 install:
   - ./oss_scripts/oss_pip_install.sh
 script:
diff --git a/oss_scripts/oss_pip_install.sh b/oss_scripts/oss_pip_install.sh
index 86c939660..ad5dd5e41 100755
--- a/oss_scripts/oss_pip_install.sh
+++ b/oss_scripts/oss_pip_install.sh
@@ -20,6 +20,6 @@ t2t-trainer --registry_help 2>&1 >/dev/null
 t2t-datagen 2>&1 | grep translate_ende 2>&1 >/dev/null && echo passed
 
 # Then install the test dependencies
-pip install -q -e .[tests,allen,trax]
+pip install -q -e .[tests,allen]
 # Make sure to install the atari extras for gym
 pip install "gym[atari]"
diff --git a/setup.py b/setup.py
index e2415edbd..47a7ef51a 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.14.1',
+    version='1.15.0',
     description='Tensor2Tensor',
     author='Google Inc.',
     author_email='no-reply@google.com',
@@ -57,11 +57,11 @@
         'tensorflow-datasets',
         'tensorflow-gan',
         'tensorflow-probability==0.7.0',
+        'tf_slim',
         'tqdm',
     ],
     extras_require={
-        'tensorflow': ['tensorflow>=1.14.0'],
-        'tensorflow_gpu': ['tensorflow-gpu>=1.14.0'],
+        'tensorflow': ['tensorflow>=1.15.0'],
         'tensorflow-hub': ['tensorflow-hub>=0.1.1'],
         'tests': [
             'absl-py',
@@ -78,10 +78,6 @@
             # explicit pip install gym[atari] for the tests.
             # 'gym[atari]',
         ],
-        'trax': [
-            'jax',
-            'jaxlib',
-        ],
         'allen': ['Pillow==5.1.0', 'pandas==0.23.0'],
     },
     classifiers=[
diff --git a/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index 2f58de3e9..2a738885a 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -44,6 +44,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
+import tf_slim as slim
 
 _BASE_EXAMPLE_IMAGE_SIZE = 64
 
@@ -350,7 +351,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "targets":
-            tf.contrib.slim.tfexample_decoder.Image(
+            slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index eadd3275f..b4e094dfb 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -36,6 +36,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
+import tf_slim as slim
 
 DATA_URL = (
     "http://rail.eecs.berkeley.edu/datasets/bair_robot_pushing_dataset_v0.tar")
@@ -102,7 +103,7 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
+        "frame_number": slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
     }
     return data_fields, decoders
@@ -187,9 +188,9 @@ def extra_reading_spec(self):
         "action": tf.FixedLenFeature([4], tf.float32),
     }
     decoders = {
-        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
+        "frame_number": slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
-        "action": tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="action"),
+        "action": slim.tfexample_decoder.Tensor(tensor_key="action"),
     }
     return data_fields, decoders
 
diff --git a/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
index 8dff69547..4792263d7 100644
--- a/tensor2tensor/data_generators/fsns.py
+++ b/tensor2tensor/data_generators/fsns.py
@@ -28,6 +28,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
+import tf_slim as slim
 
 
 @registry.register_problem
@@ -76,5 +77,5 @@ def example_reading_spec(self):
         super(ImageFSNS, self).example_reading_spec())
     data_fields[label_key] = tf.VarLenFeature(tf.int64)
     data_items_to_decoders[
-        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 67aedef9b..c325772cf 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -36,6 +36,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
+import tf_slim as slim
 
 
 Frame = collections.namedtuple(
@@ -377,7 +378,7 @@ def extra_reading_spec(self):
         name: tf.FixedLenFeature([1], tf.int64) for name in field_names
     }
     decoders = {
-        name: tf.contrib.slim.tfexample_decoder.Tensor(tensor_key=name)
+        name: slim.tfexample_decoder.Tensor(tensor_key=name)
         for name in field_names
     }
     return (data_fields, decoders)
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 61766a89e..23c38c543 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -31,6 +31,7 @@
 from tensor2tensor.utils import metrics
 
 import tensorflow as tf
+import tf_slim as slim
 
 
 def matplotlib_pyplot():
@@ -172,7 +173,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "inputs":
-            tf.contrib.slim.tfexample_decoder.Image(
+            slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
@@ -239,7 +240,7 @@ def example_reading_spec(self):
     data_fields[label_key] = tf.FixedLenFeature((1,), tf.int64)
 
     data_items_to_decoders[
-        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
 
   def hparams(self, defaults, unused_model_hparams):
@@ -343,7 +344,7 @@ def example_reading_spec(self):
         super(Image2TextProblem, self).example_reading_spec())
     data_fields[label_key] = tf.VarLenFeature(tf.int64)
     data_items_to_decoders[
-        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
 
   def feature_encoders(self, data_dir):
diff --git a/tensor2tensor/data_generators/moving_mnist.py b/tensor2tensor/data_generators/moving_mnist.py
index 507207623..d8841c741 100644
--- a/tensor2tensor/data_generators/moving_mnist.py
+++ b/tensor2tensor/data_generators/moving_mnist.py
@@ -37,6 +37,7 @@
 import tensorflow as tf
 import tensorflow_datasets as tfds
 from tensorflow_datasets.video import moving_sequence
+import tf_slim as slim
 
 
 DATA_URL = (
@@ -94,7 +95,7 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
+        "frame_number": slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
     }
     return data_fields, decoders
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 15da7648d..c1ff4ef4f 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -32,6 +32,7 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
+import tf_slim as slim
 from tensorflow.contrib.tpu.python.tpu import tpu_config
 
 
@@ -198,7 +199,7 @@ class Problem(object):
         - Mutate defaults as needed
     * example_reading_spec
         - Specify the names and types of the features on disk.
-        - Specify tf.contrib.slim.tfexample_decoder
+        - Specify slim.tfexample_decoder
     * preprocess_example(example, mode, hparams)
         - Preprocess the example feature dict from feature name to Tensor or
           SparseTensor.
@@ -642,7 +643,7 @@ def dataset(self,
 
     data_filepattern = self.filepattern(data_dir, dataset_split, shard=shard)
     tf.logging.info("Reading data files from %s", data_filepattern)
-    data_files = sorted(tf.contrib.slim.parallel_reader.get_data_files(
+    data_files = sorted(slim.parallel_reader.get_data_files(
         data_filepattern))
 
     # Functions used in dataset transforms below. `filenames` can be either a
@@ -705,12 +706,12 @@ def decode_example(self, serialized_example):
     data_fields["batch_prediction_key"] = tf.FixedLenFeature([1], tf.int64, 0)
     if data_items_to_decoders is None:
       data_items_to_decoders = {
-          field: tf.contrib.slim.tfexample_decoder.Tensor(field)
+          field: slim.tfexample_decoder.Tensor(field)
           for field in data_fields
       }
 
-    decoder = tf.contrib.slim.tfexample_decoder.TFExampleDecoder(
-        data_fields, data_items_to_decoders)
+    decoder = slim.tfexample_decoder.TFExampleDecoder(data_fields,
+                                                      data_items_to_decoders)
 
     decode_items = list(sorted(data_items_to_decoders))
     decoded = decoder.decode(serialized_example, items=decode_items)
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index ff1014a41..b74c61ea5 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -32,6 +32,7 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
+import tf_slim as slim
 
 FLAGS = tf.flags.FLAGS
 
@@ -278,8 +279,8 @@ def example_reading_spec(self):
 
     # hack: ignoring true targets and putting dist_targets in targets
     data_items_to_decoders = {
-        "inputs": tf.contrib.slim.tfexample_decoder.Tensor("inputs"),
-        "targets": tf.contrib.slim.tfexample_decoder.Tensor("dist_targets"),
+        "inputs": slim.tfexample_decoder.Tensor("inputs"),
+        "targets": slim.tfexample_decoder.Tensor("dist_targets"),
     }
 
     return (data_fields, data_items_to_decoders)
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index aa0e5bcf3..e4f2e5661 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -28,6 +28,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
+import tf_slim as slim
 
 try:
   import matplotlib  # pylint: disable=g-import-not-at-top
@@ -85,7 +86,7 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
+        "frame_number": slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
     }
     return data_fields, decoders
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 4939efcc3..98683c0e6 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -35,6 +35,7 @@
 from tensor2tensor.utils import video_metrics
 
 import tensorflow as tf
+import tf_slim as slim
 
 flags = tf.flags
 FLAGS = flags.FLAGS
@@ -384,7 +385,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "frame":
-            tf.contrib.slim.tfexample_decoder.Image(
+            slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 shape=[self.frame_height, self.frame_width, self.num_channels],
@@ -676,7 +677,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "inputs":
-            tf.contrib.slim.tfexample_decoder.Image(
+            slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
@@ -766,7 +767,7 @@ def example_reading_spec(self):
         super(Video2ClassProblem, self).example_reading_spec())
     data_fields[label_key] = tf.FixedLenFeature((1,), tf.int64)
     data_items_to_decoders[
-        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
 
   def hparams(self, defaults, unused_model_hparams):
diff --git a/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
index 8a39f85f6..e0c75d388 100644
--- a/tensor2tensor/data_generators/vqa.py
+++ b/tensor2tensor/data_generators/vqa.py
@@ -39,6 +39,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
+import tf_slim as slim
 
 
 def _get_vqa_v2_annotations(directory,
@@ -217,10 +218,10 @@ def example_reading_spec(self):
         (), tf.int64, allow_missing=True)
 
     data_items_to_decoders[
-        "question"] = tf.contrib.slim.tfexample_decoder.Tensor(
+        "question"] = slim.tfexample_decoder.Tensor(
             "image/question")
     data_items_to_decoders[
-        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(
+        "targets"] = slim.tfexample_decoder.Tensor(
             "image/answer")
     return data_fields, data_items_to_decoders
 
@@ -338,23 +339,23 @@ def example_reading_spec(self):
         (), tf.int64, allow_missing=True)
 
     data_items_to_decoders[
-        "inputs"] = tf.contrib.slim.tfexample_decoder.Tensor(
+        "inputs"] = slim.tfexample_decoder.Tensor(
             "image/feature")
     data_items_to_decoders[
-        "question_id"] = tf.contrib.slim.tfexample_decoder.Tensor(
+        "question_id"] = slim.tfexample_decoder.Tensor(
             "image/question_id")
     data_items_to_decoders[
-        "image_id"] = tf.contrib.slim.tfexample_decoder.Tensor(
+        "image_id"] = slim.tfexample_decoder.Tensor(
             "image/image_id")
 
     data_items_to_decoders[
-        "spatial_feature"] = tf.contrib.slim.tfexample_decoder.Tensor(
+        "spatial_feature"] = slim.tfexample_decoder.Tensor(
             "image/spatial_feature")
     data_items_to_decoders[
-        "question"] = tf.contrib.slim.tfexample_decoder.Tensor(
+        "question"] = slim.tfexample_decoder.Tensor(
             "image/question")
     data_items_to_decoders[
-        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(
+        "targets"] = slim.tfexample_decoder.Tensor(
             "image/answer")
 
     return data_fields, data_items_to_decoders
diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 43afb1a75..7c02f43fd 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -33,6 +33,7 @@
 from tensor2tensor.envs import trajectory
 from tensor2tensor.layers import modalities
 import tensorflow as tf
+import tf_slim as slim
 
 # Names for data fields in stored tf.Examples.
 TIMESTEP_FIELD = "timestep"
@@ -476,7 +477,7 @@ def example_reading_spec(self):
     }
 
     data_items_to_decoders = {
-        field: tf.contrib.slim.tfexample_decoder.Tensor(field)
+        field: slim.tfexample_decoder.Tensor(field)
         for field in data_fields
     }
 
diff --git a/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
index 92a7a7e54..82eb2fa71 100644
--- a/tensor2tensor/envs/rendered_env_problem.py
+++ b/tensor2tensor/envs/rendered_env_problem.py
@@ -25,6 +25,7 @@
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import gym_env_problem
 import tensorflow as tf
+import tf_slim as slim
 
 _IMAGE_ENCODED_FIELD = "image/encoded"
 _IMAGE_FORMAT_FIELD = "image/format"
@@ -80,7 +81,7 @@ def example_reading_spec(self):
     # Add frame number spec and decoder.
     env_fields[_FRAME_NUMBER_FIELD] = tf.FixedLenFeature((1,), tf.int64)
     env_decoders[
-        _FRAME_NUMBER_FIELD] = tf.contrib.slim.tfexample_decoder.Tensor(
+        _FRAME_NUMBER_FIELD] = slim.tfexample_decoder.Tensor(
             _FRAME_NUMBER_FIELD)
 
     # Add video fields and decoders
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index c65878163..e80aa6dc4 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -28,6 +28,7 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
+import tf_slim as slim
 
 
 def cast_ints_to_int32(features):
@@ -379,7 +380,7 @@ def define_shapes(example):
     dataset = dataset.repeat()
 
   if is_training and skip_random_fraction_when_training:
-    data_files = tf.contrib.slim.parallel_reader.get_data_files(filepattern)
+    data_files = slim.parallel_reader.get_data_files(filepattern)
     #  In continuous_train_and_eval when switching between train and
     #  eval, this input_fn method gets called multiple times and it
     #  would give you the exact same samples from the last call
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 42f673517..7786b5ec7 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -546,13 +546,6 @@ def env_problem(env_problem_name, **kwargs):
 list_pruning_strategies = lambda: sorted(Registries.pruning_strategies)
 register_pruning_strategy = Registries.pruning_strategies.register
 
-# deprecated functions - plurals inconsistent with rest
-# deprecation decorators added 2019-01-25
-attacks = tf.contrib.framework.deprecated(None, "Use registry.attack")(attack)
-pruning_strategies = tf.contrib.framework.deprecated(
-    None, "Use registry.pruning_strategy")(
-        pruning_strategy)
-
 
 def display_list_by_prefix(names_list, starting_spaces=0):
   """Creates a help string for names_list grouped by prefix."""

From 67ddb40b3d804f415b6a33f476ac33dd415c33ac Mon Sep 17 00:00:00 2001
From: Sepehr Sameni <Sepehr.Sameni@gmail.com>
Date: Thu, 21 Nov 2019 17:57:51 +0000
Subject: [PATCH 2570/2720] use batch_size in _test_img2img_transformer (#1724)

---
 tensor2tensor/models/image_transformer_2d_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/image_transformer_2d_test.py b/tensor2tensor/models/image_transformer_2d_test.py
index 7f903fb15..7deddc870 100644
--- a/tensor2tensor/models/image_transformer_2d_test.py
+++ b/tensor2tensor/models/image_transformer_2d_test.py
@@ -35,8 +35,8 @@ def _test_img2img_transformer(self, net):
     hparams = image_transformer_2d.img2img_transformer2d_tiny()
     hparams.data_dir = ""
     p_hparams = registry.problem("image_celeba").get_hparams(hparams)
-    inputs = np.random.randint(256, size=(3, 4, 4, 3))
-    targets = np.random.randint(256, size=(3, 8, 8, 3))
+    inputs = np.random.randint(256, size=(batch_size, 4, 4, 3))
+    targets = np.random.randint(256, size=(batch_size, 8, 8, 3))
     with self.test_session() as session:
       features = {
           "inputs": tf.constant(inputs, dtype=tf.int32),

From c825d126643d1c8864d43d828df0e7a868eaa180 Mon Sep 17 00:00:00 2001
From: Seppo Enarvi <seppo.git@marjaniemi.com>
Date: Thu, 21 Nov 2019 18:58:13 +0100
Subject: [PATCH 2571/2720] Fix decoding in prepend mode (#1726)

* Create an integer problem_0_steps variable.

* Save inputs to the feature "partial_targets" when prepend_mode is not "none".

* Removed a second call to update_hparams_for_universal_transformer().

Fixes hyperparameter sets universal_transformer_big and universal_transformer_base_tpu.

* Fix a bug to make partial targets work for beam size > 1

The dimension of the multiplication of the partial targets was wrong:  (a, b, c, d) --> (a, b, c, d, a, b, c, d)
Correct multiplication needs to be: (a, b, c, d) --> (a, a, b, b, c, c, d, d)
This is because it is (batch_size * beam_size) instead of (beam_size * batch_size).

Basically, tf.tile needs to be replaced by tf.repeat which is introduced in tf 1.15. This is a workaround for tf 1.14.
---
 .../models/research/universal_transformer.py       |  2 --
 tensor2tensor/models/transformer.py                | 13 ++++++++-----
 tensor2tensor/utils/decoding.py                    | 14 ++++++++++++++
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 28040584e..a5a65504e 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -458,7 +458,6 @@ def universal_transformer_base():
 @registry.register_hparams
 def universal_transformer_base_tpu():
   hparams = universal_transformer_base()
-  hparams = update_hparams_for_universal_transformer(hparams)
   transformer.update_hparams_for_tpu(hparams)
   hparams.add_step_timing_signal = False
   return hparams
@@ -467,7 +466,6 @@ def universal_transformer_base_tpu():
 @registry.register_hparams
 def universal_transformer_big():
   hparams = universal_transformer_base()
-  hparams = update_hparams_for_universal_transformer(hparams)
   hparams.hidden_size = 2048
   hparams.filter_size = 8192
   return hparams
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index eef24dba1..51b60fc83 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -863,9 +863,15 @@ def symbols_to_logits_fn(ids, i, cache):
         vocab_size = tf.shape(ret)[1]
 
         def forced_logits():
+          # Workaround for: tf.one_hot(
+          #               tf.repeat(partial_targets[:, i], [beam_size]), vocab_size, 0.0,
+          #               -1e9)
+          # Can be replaced by the above in future versions (from tf 1.15).
           return tf.one_hot(
-              tf.tile(partial_targets[:, i], [beam_size]), vocab_size, 0.0,
-              -1e9)
+              tf.reshape(tf.tile(
+                  tf.reshape(partial_targets[:, i], [-1, 1]),
+                  [1, beam_size]), [-1]),
+              vocab_size, 0.0, -1e9)
 
         ret = tf.cond(
             tf.less(i, partial_targets_length), forced_logits, lambda: ret)
@@ -1168,9 +1174,6 @@ def fast_decode(encoder_output,
           "scores": decoding log probs from the beam search,
               None if using greedy decoding (beam_size=1)
       }
-
-    Raises:
-      NotImplementedError: If beam size > 1 with partial targets.
   """
   if encoder_output is not None:
     batch_size = common_layers.shape_list(encoder_output)[0]
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 656edfdfd..a63e83d5b 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -927,6 +927,13 @@ def _interactive_input_tensor_to_features_dict(feature_map, hparams):
   features["decode_length"] = (
       IMAGE_DECODE_LENGTH if input_is_image else inputs[1])
   features["inputs"] = x
+  # Save inputs to "partial_targets" when prepending inputs to targets. Also
+  # keep "inputs" as some models crash if they don't exist.
+  if getattr(hparams, "prepend_mode", "none") != "none":
+    shape = tf.shape(x)
+    partial_targets = tf.reshape(x, [shape[0], shape[1]])
+    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
+    features["partial_targets"] = partial_targets
   return features
 
 
@@ -957,6 +964,13 @@ def _decode_input_tensor_to_features_dict(feature_map, hparams):
   features["decode_length"] = (
       IMAGE_DECODE_LENGTH if input_is_image else tf.shape(x)[1] + 50)
   features["inputs"] = x
+  # Save inputs to "partial_targets" when prepending inputs to targets. Also
+  # keep "inputs" as some models crash if they don't exist.
+  if getattr(hparams, "prepend_mode", "none") != "none":
+    shape = tf.shape(x)
+    partial_targets = tf.reshape(x, [shape[0], shape[1]])
+    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
+    features["partial_targets"] = partial_targets
   return features
 
 

From d381f2bc30969cfaf93ba7987e25c17ce48cdac5 Mon Sep 17 00:00:00 2001
From: Prasasto Adi <prasastoadi@users.noreply.github.com>
Date: Fri, 22 Nov 2019 00:58:28 +0700
Subject: [PATCH 2572/2720] En-Id untokenized parallel corpora (#1733)

Change to untokenized parallel corpora due to t2t automatically tokenizes sentences.
---
 .../data_generators/translate_enid.py         | 26 +++++++++----------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/data_generators/translate_enid.py b/tensor2tensor/data_generators/translate_enid.py
index bf598cae2..1aa2d2a2b 100644
--- a/tensor2tensor/data_generators/translate_enid.py
+++ b/tensor2tensor/data_generators/translate_enid.py
@@ -39,24 +39,24 @@
 # http://www.panl10n.net/english/outputs/Indonesia/BPPT/0902/BPPTIndToEngCorpusHalfM.zip # pylint: disable=line-too-long
 _ENID_TRAIN_DATASETS = [
     [
-        _REPO + "IWSLT17.train.en-id.tok.tgz",
-        ("IWSLT17.train.en-id.tok.en", "IWSLT17.train.en-id.tok.id")
+        _REPO + "IWSLT17.train.en-id.tgz",
+        ("IWSLT17.train.en-id.en", "IWSLT17.train.en-id.id")
     ],
     [
-        _REPO + "PANL-BPPT-ECO-EN-ID-150Kw.tok.tgz",
-        ("PANL-BPPT-ECO-EN-150Kw.tok.txt", "PANL-BPPT-ECO-ID-150Kw.tok.txt")
+        _REPO + "PANL-BPPT-ECO-EN-ID-150Kw.tgz",
+        ("PANL-BPPT-ECO-EN-150Kw.txt", "PANL-BPPT-ECO-ID-150Kw.txt")
     ],
     [
-        _REPO + "PANL-BPPT-INT-EN-ID-150Kw.tok.tgz",
-        ("PANL-BPPT-INT-EN-150Kw.tok.txt", "PANL-BPPT-INT-ID-150Kw.tok.txt")
+        _REPO + "PANL-BPPT-INT-EN-ID-150Kw.tgz",
+        ("PANL-BPPT-INT-EN-150Kw.txt", "PANL-BPPT-INT-ID-150Kw.txt")
     ],
     [
-        _REPO + "PANL-BPPT-SCI-EN-ID-100Kw.tok.tgz",
-        ("PANL-BPPT-SCI-EN-100Kw.tok.txt", "PANL-BPPT-SCI-ID-100Kw.tok.txt")
+        _REPO + "PANL-BPPT-SCI-EN-ID-100Kw.tgz",
+        ("PANL-BPPT-SCI-EN-100Kw.txt", "PANL-BPPT-SCI-ID-100Kw.txt")
     ],
     [
-        _REPO + "PANL-BPPT-SPO-EN-ID-100Kw.tok.tgz",
-        ("PANL-BPPT-SPO-EN-100Kw.tok.txt", "PANL-BPPT-SPO-ID-100Kw.tok.txt")
+        _REPO + "PANL-BPPT-SPO-EN-ID-100Kw.tgz",
+        ("PANL-BPPT-SPO-EN-100Kw.txt", "PANL-BPPT-SPO-ID-100Kw.txt")
     ],
 ]
 
@@ -65,9 +65,9 @@
 # https://wit3.fbk.eu/mt.php?release=2017-01-more
 _ENID_TEST_DATASETS = [
     [
-        _REPO + "IWSLT17.TED.tst2017plus.en-id.tok.tgz",
-        ("IWSLT17.TED.tst2017plus.en-id.tok.en",
-         "IWSLT17.TED.tst2017plus.en-id.tok.id")
+        _REPO + "IWSLT17.TED.tst2017plus.en-id.tgz",
+        ("IWSLT17.TED.tst2017plus.en-id.en",
+         "IWSLT17.TED.tst2017plus.en-id.id")
     ]
 ]
 

From b84d5e8f88c7465e6cbf6a86ade049b9011f98fd Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 21 Nov 2019 10:21:28 -0800
Subject: [PATCH 2573/2720] Revert "Fix decoding in prepend mode (#1726)"
 (#1749)

This reverts commit c825d126643d1c8864d43d828df0e7a868eaa180.
---
 .../models/research/universal_transformer.py       |  2 ++
 tensor2tensor/models/transformer.py                | 13 +++++--------
 tensor2tensor/utils/decoding.py                    | 14 --------------
 3 files changed, 7 insertions(+), 22 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index a5a65504e..28040584e 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -458,6 +458,7 @@ def universal_transformer_base():
 @registry.register_hparams
 def universal_transformer_base_tpu():
   hparams = universal_transformer_base()
+  hparams = update_hparams_for_universal_transformer(hparams)
   transformer.update_hparams_for_tpu(hparams)
   hparams.add_step_timing_signal = False
   return hparams
@@ -466,6 +467,7 @@ def universal_transformer_base_tpu():
 @registry.register_hparams
 def universal_transformer_big():
   hparams = universal_transformer_base()
+  hparams = update_hparams_for_universal_transformer(hparams)
   hparams.hidden_size = 2048
   hparams.filter_size = 8192
   return hparams
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 51b60fc83..eef24dba1 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -863,15 +863,9 @@ def symbols_to_logits_fn(ids, i, cache):
         vocab_size = tf.shape(ret)[1]
 
         def forced_logits():
-          # Workaround for: tf.one_hot(
-          #               tf.repeat(partial_targets[:, i], [beam_size]), vocab_size, 0.0,
-          #               -1e9)
-          # Can be replaced by the above in future versions (from tf 1.15).
           return tf.one_hot(
-              tf.reshape(tf.tile(
-                  tf.reshape(partial_targets[:, i], [-1, 1]),
-                  [1, beam_size]), [-1]),
-              vocab_size, 0.0, -1e9)
+              tf.tile(partial_targets[:, i], [beam_size]), vocab_size, 0.0,
+              -1e9)
 
         ret = tf.cond(
             tf.less(i, partial_targets_length), forced_logits, lambda: ret)
@@ -1174,6 +1168,9 @@ def fast_decode(encoder_output,
           "scores": decoding log probs from the beam search,
               None if using greedy decoding (beam_size=1)
       }
+
+    Raises:
+      NotImplementedError: If beam size > 1 with partial targets.
   """
   if encoder_output is not None:
     batch_size = common_layers.shape_list(encoder_output)[0]
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index a63e83d5b..656edfdfd 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -927,13 +927,6 @@ def _interactive_input_tensor_to_features_dict(feature_map, hparams):
   features["decode_length"] = (
       IMAGE_DECODE_LENGTH if input_is_image else inputs[1])
   features["inputs"] = x
-  # Save inputs to "partial_targets" when prepending inputs to targets. Also
-  # keep "inputs" as some models crash if they don't exist.
-  if getattr(hparams, "prepend_mode", "none") != "none":
-    shape = tf.shape(x)
-    partial_targets = tf.reshape(x, [shape[0], shape[1]])
-    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
-    features["partial_targets"] = partial_targets
   return features
 
 
@@ -964,13 +957,6 @@ def _decode_input_tensor_to_features_dict(feature_map, hparams):
   features["decode_length"] = (
       IMAGE_DECODE_LENGTH if input_is_image else tf.shape(x)[1] + 50)
   features["inputs"] = x
-  # Save inputs to "partial_targets" when prepending inputs to targets. Also
-  # keep "inputs" as some models crash if they don't exist.
-  if getattr(hparams, "prepend_mode", "none") != "none":
-    shape = tf.shape(x)
-    partial_targets = tf.reshape(x, [shape[0], shape[1]])
-    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
-    features["partial_targets"] = partial_targets
   return features
 
 

From e86b7e355633322e4712e9b6ded5e14172fe9009 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 21 Nov 2019 09:56:07 -0800
Subject: [PATCH 2574/2720] Remove pylint as a t2t dependency, this is
 preventing pip install of our test dependencies.

Also we don't use it anymore.

PiperOrigin-RevId: 281771877
---
 setup.py                                      |  1 -
 .../data_generators/translate_enid.py         | 26 +++++++++----------
 .../models/image_transformer_2d_test.py       |  4 +--
 3 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/setup.py b/setup.py
index 47a7ef51a..de1a41645 100644
--- a/setup.py
+++ b/setup.py
@@ -70,7 +70,6 @@
             'attrs>=17.4.0',
             'pytest>=3.8.0',
             'mock',
-            'pylint',
             'jupyter',
             'matplotlib',
             # Need atari extras for Travis tests, but because gym is already in
diff --git a/tensor2tensor/data_generators/translate_enid.py b/tensor2tensor/data_generators/translate_enid.py
index 1aa2d2a2b..bf598cae2 100644
--- a/tensor2tensor/data_generators/translate_enid.py
+++ b/tensor2tensor/data_generators/translate_enid.py
@@ -39,24 +39,24 @@
 # http://www.panl10n.net/english/outputs/Indonesia/BPPT/0902/BPPTIndToEngCorpusHalfM.zip # pylint: disable=line-too-long
 _ENID_TRAIN_DATASETS = [
     [
-        _REPO + "IWSLT17.train.en-id.tgz",
-        ("IWSLT17.train.en-id.en", "IWSLT17.train.en-id.id")
+        _REPO + "IWSLT17.train.en-id.tok.tgz",
+        ("IWSLT17.train.en-id.tok.en", "IWSLT17.train.en-id.tok.id")
     ],
     [
-        _REPO + "PANL-BPPT-ECO-EN-ID-150Kw.tgz",
-        ("PANL-BPPT-ECO-EN-150Kw.txt", "PANL-BPPT-ECO-ID-150Kw.txt")
+        _REPO + "PANL-BPPT-ECO-EN-ID-150Kw.tok.tgz",
+        ("PANL-BPPT-ECO-EN-150Kw.tok.txt", "PANL-BPPT-ECO-ID-150Kw.tok.txt")
     ],
     [
-        _REPO + "PANL-BPPT-INT-EN-ID-150Kw.tgz",
-        ("PANL-BPPT-INT-EN-150Kw.txt", "PANL-BPPT-INT-ID-150Kw.txt")
+        _REPO + "PANL-BPPT-INT-EN-ID-150Kw.tok.tgz",
+        ("PANL-BPPT-INT-EN-150Kw.tok.txt", "PANL-BPPT-INT-ID-150Kw.tok.txt")
     ],
     [
-        _REPO + "PANL-BPPT-SCI-EN-ID-100Kw.tgz",
-        ("PANL-BPPT-SCI-EN-100Kw.txt", "PANL-BPPT-SCI-ID-100Kw.txt")
+        _REPO + "PANL-BPPT-SCI-EN-ID-100Kw.tok.tgz",
+        ("PANL-BPPT-SCI-EN-100Kw.tok.txt", "PANL-BPPT-SCI-ID-100Kw.tok.txt")
     ],
     [
-        _REPO + "PANL-BPPT-SPO-EN-ID-100Kw.tgz",
-        ("PANL-BPPT-SPO-EN-100Kw.txt", "PANL-BPPT-SPO-ID-100Kw.txt")
+        _REPO + "PANL-BPPT-SPO-EN-ID-100Kw.tok.tgz",
+        ("PANL-BPPT-SPO-EN-100Kw.tok.txt", "PANL-BPPT-SPO-ID-100Kw.tok.txt")
     ],
 ]
 
@@ -65,9 +65,9 @@
 # https://wit3.fbk.eu/mt.php?release=2017-01-more
 _ENID_TEST_DATASETS = [
     [
-        _REPO + "IWSLT17.TED.tst2017plus.en-id.tgz",
-        ("IWSLT17.TED.tst2017plus.en-id.en",
-         "IWSLT17.TED.tst2017plus.en-id.id")
+        _REPO + "IWSLT17.TED.tst2017plus.en-id.tok.tgz",
+        ("IWSLT17.TED.tst2017plus.en-id.tok.en",
+         "IWSLT17.TED.tst2017plus.en-id.tok.id")
     ]
 ]
 
diff --git a/tensor2tensor/models/image_transformer_2d_test.py b/tensor2tensor/models/image_transformer_2d_test.py
index 7deddc870..7f903fb15 100644
--- a/tensor2tensor/models/image_transformer_2d_test.py
+++ b/tensor2tensor/models/image_transformer_2d_test.py
@@ -35,8 +35,8 @@ def _test_img2img_transformer(self, net):
     hparams = image_transformer_2d.img2img_transformer2d_tiny()
     hparams.data_dir = ""
     p_hparams = registry.problem("image_celeba").get_hparams(hparams)
-    inputs = np.random.randint(256, size=(batch_size, 4, 4, 3))
-    targets = np.random.randint(256, size=(batch_size, 8, 8, 3))
+    inputs = np.random.randint(256, size=(3, 4, 4, 3))
+    targets = np.random.randint(256, size=(3, 8, 8, 3))
     with self.test_session() as session:
       features = {
           "inputs": tf.constant(inputs, dtype=tf.int32),

From 37ad151bdae3e4b8bdfc44559332f9cbffb2cb83 Mon Sep 17 00:00:00 2001
From: Prasasto Adi <prasastoadi@users.noreply.github.com>
Date: Thu, 21 Nov 2019 09:58:59 -0800
Subject: [PATCH 2575/2720] Merge of PR #1733

PiperOrigin-RevId: 281772422
---
 .../data_generators/translate_enid.py         | 26 +++++++++----------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/data_generators/translate_enid.py b/tensor2tensor/data_generators/translate_enid.py
index bf598cae2..1aa2d2a2b 100644
--- a/tensor2tensor/data_generators/translate_enid.py
+++ b/tensor2tensor/data_generators/translate_enid.py
@@ -39,24 +39,24 @@
 # http://www.panl10n.net/english/outputs/Indonesia/BPPT/0902/BPPTIndToEngCorpusHalfM.zip # pylint: disable=line-too-long
 _ENID_TRAIN_DATASETS = [
     [
-        _REPO + "IWSLT17.train.en-id.tok.tgz",
-        ("IWSLT17.train.en-id.tok.en", "IWSLT17.train.en-id.tok.id")
+        _REPO + "IWSLT17.train.en-id.tgz",
+        ("IWSLT17.train.en-id.en", "IWSLT17.train.en-id.id")
     ],
     [
-        _REPO + "PANL-BPPT-ECO-EN-ID-150Kw.tok.tgz",
-        ("PANL-BPPT-ECO-EN-150Kw.tok.txt", "PANL-BPPT-ECO-ID-150Kw.tok.txt")
+        _REPO + "PANL-BPPT-ECO-EN-ID-150Kw.tgz",
+        ("PANL-BPPT-ECO-EN-150Kw.txt", "PANL-BPPT-ECO-ID-150Kw.txt")
     ],
     [
-        _REPO + "PANL-BPPT-INT-EN-ID-150Kw.tok.tgz",
-        ("PANL-BPPT-INT-EN-150Kw.tok.txt", "PANL-BPPT-INT-ID-150Kw.tok.txt")
+        _REPO + "PANL-BPPT-INT-EN-ID-150Kw.tgz",
+        ("PANL-BPPT-INT-EN-150Kw.txt", "PANL-BPPT-INT-ID-150Kw.txt")
     ],
     [
-        _REPO + "PANL-BPPT-SCI-EN-ID-100Kw.tok.tgz",
-        ("PANL-BPPT-SCI-EN-100Kw.tok.txt", "PANL-BPPT-SCI-ID-100Kw.tok.txt")
+        _REPO + "PANL-BPPT-SCI-EN-ID-100Kw.tgz",
+        ("PANL-BPPT-SCI-EN-100Kw.txt", "PANL-BPPT-SCI-ID-100Kw.txt")
     ],
     [
-        _REPO + "PANL-BPPT-SPO-EN-ID-100Kw.tok.tgz",
-        ("PANL-BPPT-SPO-EN-100Kw.tok.txt", "PANL-BPPT-SPO-ID-100Kw.tok.txt")
+        _REPO + "PANL-BPPT-SPO-EN-ID-100Kw.tgz",
+        ("PANL-BPPT-SPO-EN-100Kw.txt", "PANL-BPPT-SPO-ID-100Kw.txt")
     ],
 ]
 
@@ -65,9 +65,9 @@
 # https://wit3.fbk.eu/mt.php?release=2017-01-more
 _ENID_TEST_DATASETS = [
     [
-        _REPO + "IWSLT17.TED.tst2017plus.en-id.tok.tgz",
-        ("IWSLT17.TED.tst2017plus.en-id.tok.en",
-         "IWSLT17.TED.tst2017plus.en-id.tok.id")
+        _REPO + "IWSLT17.TED.tst2017plus.en-id.tgz",
+        ("IWSLT17.TED.tst2017plus.en-id.en",
+         "IWSLT17.TED.tst2017plus.en-id.id")
     ]
 ]
 

From 259cbd016d9cac8b903cc4dfe3205bc2eac8d37d Mon Sep 17 00:00:00 2001
From: Gabe Grand <gabrieljgrand@gmail.com>
Date: Thu, 21 Nov 2019 14:18:51 -0500
Subject: [PATCH 2576/2720] Create Text2RealProblem class for regression-based
 problems (#1748)

* Create Text2RealProblem class for regression-based problems

* Global style fixes

* pylint style fixes

* Remove whitespace

* More whitespace
---
 .../data_generators/text_problems.py          | 104 ++++++++++++++++++
 .../data_generators/text_problems_test.py     |  17 +++
 tensor2tensor/models/transformer.py           |   6 +-
 tensor2tensor/utils/t2t_model.py              |  12 +-
 4 files changed, 133 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 4067a65aa..f807e6f51 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -17,6 +17,7 @@
 
 * Text2TextProblem: input=text, target=text.
 * Text2ClassProblem: input=text, target=class.
+* Text2RealProblem: input=text, target=float.
 * Text2SelfProblem (for language modeling): target=text
 * QuestionAndContext2TextProblem: input=text, context=text, target=text.
 
@@ -605,6 +606,94 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
       yield {"inputs": inputs, "targets": [label]}
 
 
+class Text2RealProblem(Text2TextProblem):
+  """Base class for text regression problems with one or more tasks.
+    Suitable for text-based problems where targets are continuous, real values.
+    When ntasks = 1, each text example is mapped to a single scalar value. When
+    ntasks > 1, each text example is mapped to a 1-d vector of length ntasks.
+  """
+
+  @property
+  def ntasks(self):
+    """Set to n > 1 for multitask regression."""
+    return 1
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    """Generate samples of text and real-valued target pairs.
+    Each yielded dict will be a single example. The inputs should be raw text.
+    The target should be a list containing ntasks floats.
+    Args:
+      data_dir: final data directory. Typically only used in this method to copy
+        over user-supplied vocab files (for example, if vocab_type ==
+        VocabType.TOKEN).
+      tmp_dir: temporary directory that you can use for downloading and scratch.
+      dataset_split: problem.DatasetSplit, which data split to generate samples
+        for (for example, training and evaluation).
+    Yields:
+      {"inputs": text, "targets": [x1, x2, ..., xN]} where N is ntasks
+    """
+    raise NotImplementedError()
+
+  def generate_text_for_vocab(self, data_dir, tmp_dir):
+    for i, sample in enumerate(
+        self.generate_samples(data_dir, tmp_dir, problem.DatasetSplit.TRAIN)):
+      yield sample["inputs"]
+      if self.max_samples_for_vocab and (i + 1) >= self.max_samples_for_vocab:
+        break
+
+  def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
+    generator = self.generate_samples(data_dir, tmp_dir, dataset_split)
+    encoder = self.get_or_create_vocab(data_dir, tmp_dir)
+    for sample in generator:
+      inputs = encoder.encode(sample["inputs"])
+      inputs.append(text_encoder.EOS_ID)
+      yield {"inputs": inputs, "targets": sample["targets"]}
+
+  def feature_encoders(self, data_dir):
+    encoder = self.get_or_create_vocab(data_dir, None, force_get=True)
+
+    return {
+        "inputs": encoder,
+        "targets": text_encoder.RealEncoder(),
+    }
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    p.modality = {
+        "inputs": modalities.ModalityType.SYMBOL,
+        "targets": modalities.ModalityType.REAL_L2_LOSS,
+    }
+    p.vocab_size = {
+        "inputs": self._encoders["inputs"].vocab_size,
+        "targets": self.ntasks
+    }
+    p.target_space_id = problem.SpaceID.REAL
+    p.add_hparam("regression_targets", True)
+
+  def max_length(self, model_hparams):
+    return model_hparams.batch_size * self.ntasks
+
+  def preprocess_example(self, example, unused_mode, unused_hparams):
+    example = problem.preprocess_example_common(example, unused_mode,
+                                                unused_hparams)
+    example["targets"] = tf.reshape(example["targets"], [1, 1, self.ntasks])
+    return example
+
+  def example_reading_spec(self):
+    data_fields = {
+        "inputs": tf.VarLenFeature(tf.int64),
+        "targets": tf.FixedLenFeature([self.ntasks], tf.float32),
+    }
+    data_items_to_decoders = None
+    return (data_fields, data_items_to_decoders)
+
+  def eval_metrics(self):
+    metrics_list = [metrics.Metrics.RMSE]
+    if self.ntasks == 1:
+      metrics_list.append(metrics.Metrics.PEARSON)
+    return metrics_list
+
+
 def txt_line_iterator(txt_path):
   """Iterate through lines of file."""
   with tf.gfile.Open(txt_path) as f:
@@ -692,6 +781,21 @@ def text2class_txt_iterator(source_txt_path, label_txt_path, class_strs=None):
     yield {"inputs": inputs, "label": label}
 
 
+def text2real_txt_iterator(source_txt_path, target_txt_path):
+  """Yield dicts for Text2RealProblem.generate_samples from lines of files.
+  Args:
+    source_txt_path: txt file with record per line.
+    target_txt_path: txt file with float (or space-separated float list for
+      multitask) per line.
+  Yields:
+    {"inputs": inputs, "targets": targets}
+  """
+  for inputs, targets in zip(
+      txt_line_iterator(source_txt_path), txt_line_iterator(target_txt_path)):
+    targets = [float(x) for x in targets.split(" ")]
+    yield {"inputs": inputs, "targets": targets}
+
+
 def text2text_txt_tab_iterator(txt_path):
   """Yield dicts for Text2TextProblem.generate_samples from lines of txt_path.
 
diff --git a/tensor2tensor/data_generators/text_problems_test.py b/tensor2tensor/data_generators/text_problems_test.py
index 51f948fcf..720f1ba68 100644
--- a/tensor2tensor/data_generators/text_problems_test.py
+++ b/tensor2tensor/data_generators/text_problems_test.py
@@ -94,6 +94,13 @@ def setUpClass(cls):
     tf.gfile.Copy(cls.targets_file, os.path.join(cls.tmp_dir,
                                                  "targets.eval.txt"))
 
+    cls.targets_regr = [[1.23, 2.34], [4.56, 5.67]]
+    cls.targets_regr_file = os.path.join(cls.tmp_dir, "targets_regr.train.txt")
+    with tf.gfile.Open(cls.targets_regr_file, "w") as f:
+      for targets in cls.targets_regr:
+        f.write(" ".join([str(x) for x in targets]) + "\n")
+
+
   def testTxtLineIterator(self):
     lines = [line for line in text_problems.txt_line_iterator(self.inputs_file)]
     self.assertEqual(lines, self.inputs)
@@ -136,6 +143,16 @@ def testText2ClassTxtIteratorWithStrs(self):
     self.assertEqual(inputs, self.inputs)
     self.assertEqual(labels, self.labels)
 
+  def testText2RealTxtIterator(self):
+    inputs = []
+    targets = []
+    for entry in text_problems.text2real_txt_iterator(self.inputs_file,
+                                                      self.targets_regr_file):
+      inputs.append(entry["inputs"])
+      targets.append(entry["targets"])
+    self.assertEqual(inputs, self.inputs)
+    self.assertEqual(targets, self.targets_regr)
+
   def testText2TextTxtTabIterator(self):
     inputs = []
     targets = []
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index eef24dba1..c1c75f121 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -462,7 +462,8 @@ def _fast_decode_tpu(self,
 
     if self.has_input:
       inputs_shape = common_layers.shape_list(features["inputs"])
-      if target_modality == modalities.ModalityType.CLASS_LABEL:
+      if (target_modality == modalities.ModalityType.CLASS_LABEL or
+          self._problem_hparams.get("regression_targets")):
         decode_length = 1
       else:
         decode_length = (
@@ -704,7 +705,8 @@ def _fast_decode(self,
           " of the dataset when decoding.")
     if self.has_input:
       inputs_shape = common_layers.shape_list(features["inputs"])
-      if target_modality == modalities.ModalityType.CLASS_LABEL:
+      if (target_modality == modalities.ModalityType.CLASS_LABEL or
+          self._problem_hparams.get("regression_targets")):
         decode_length = 1
       else:
         decode_length = (
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index fe33e6315..2b9e419ef 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -806,8 +806,10 @@ def infer(self,
 
       if self._problem_hparams:
         target_modality = self._problem_hparams.modality["targets"]
-        if target_modality == modalities.ModalityType.CLASS_LABEL:
-          beam_size = 1  # No use to run beam-search for a single class.
+      if (target_modality == modalities.ModalityType.CLASS_LABEL or
+          self._problem_hparams.get("regression_targets")):
+        # No use to run beam-search for classification or regression.
+        beam_size = 1
       if beam_size == 1:
         log_info("Greedy Decoding")
         results = self._greedy_infer(features, decode_length, use_tpu)
@@ -1064,7 +1066,8 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
     initial_output = tf.slice(initial_output, [0, 0, 0, 0],
                               common_layers.shape_list(initial_output))
     target_modality = self._problem_hparams.modality["targets"]
-    if target_modality == modalities.ModalityType.CLASS_LABEL:
+    if (target_modality == modalities.ModalityType.CLASS_LABEL or
+        self._problem_hparams.get("regression_targets")):
       decode_length = 1
     else:
       if "partial_targets" in features:
@@ -1243,7 +1246,8 @@ def infer_step(recent_output, recent_logits, unused_loss):
     initial_output = tf.slice(initial_output, [0, 0, 0, 0],
                               common_layers.shape_list(initial_output))
     target_modality = self._problem_hparams.modality["targets"]
-    if target_modality == modalities.ModalityType.CLASS_LABEL:
+    if (target_modality == modalities.ModalityType.CLASS_LABEL or
+        self._problem_hparams.get("regression_targets")):
       decode_length = 1
     else:
       if "partial_targets" in features:

From c57ccbb93c692a9d1cc36daf525a45055fe3fbc4 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 21 Nov 2019 12:10:10 -0800
Subject: [PATCH 2577/2720] Revert "Revert "Fix decoding in prepend mode
 (#1726)" (#1749)" (#1750)

This reverts commit b84d5e8f88c7465e6cbf6a86ade049b9011f98fd.
---
 .../models/research/universal_transformer.py       |  2 --
 tensor2tensor/models/transformer.py                | 13 ++++++++-----
 tensor2tensor/utils/decoding.py                    | 14 ++++++++++++++
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 28040584e..a5a65504e 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -458,7 +458,6 @@ def universal_transformer_base():
 @registry.register_hparams
 def universal_transformer_base_tpu():
   hparams = universal_transformer_base()
-  hparams = update_hparams_for_universal_transformer(hparams)
   transformer.update_hparams_for_tpu(hparams)
   hparams.add_step_timing_signal = False
   return hparams
@@ -467,7 +466,6 @@ def universal_transformer_base_tpu():
 @registry.register_hparams
 def universal_transformer_big():
   hparams = universal_transformer_base()
-  hparams = update_hparams_for_universal_transformer(hparams)
   hparams.hidden_size = 2048
   hparams.filter_size = 8192
   return hparams
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index c1c75f121..d2dda17ec 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -865,9 +865,15 @@ def symbols_to_logits_fn(ids, i, cache):
         vocab_size = tf.shape(ret)[1]
 
         def forced_logits():
+          # Workaround for: tf.one_hot(
+          #               tf.repeat(partial_targets[:, i], [beam_size]), vocab_size, 0.0,
+          #               -1e9)
+          # Can be replaced by the above in future versions (from tf 1.15).
           return tf.one_hot(
-              tf.tile(partial_targets[:, i], [beam_size]), vocab_size, 0.0,
-              -1e9)
+              tf.reshape(tf.tile(
+                  tf.reshape(partial_targets[:, i], [-1, 1]),
+                  [1, beam_size]), [-1]),
+              vocab_size, 0.0, -1e9)
 
         ret = tf.cond(
             tf.less(i, partial_targets_length), forced_logits, lambda: ret)
@@ -1170,9 +1176,6 @@ def fast_decode(encoder_output,
           "scores": decoding log probs from the beam search,
               None if using greedy decoding (beam_size=1)
       }
-
-    Raises:
-      NotImplementedError: If beam size > 1 with partial targets.
   """
   if encoder_output is not None:
     batch_size = common_layers.shape_list(encoder_output)[0]
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 656edfdfd..a63e83d5b 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -927,6 +927,13 @@ def _interactive_input_tensor_to_features_dict(feature_map, hparams):
   features["decode_length"] = (
       IMAGE_DECODE_LENGTH if input_is_image else inputs[1])
   features["inputs"] = x
+  # Save inputs to "partial_targets" when prepending inputs to targets. Also
+  # keep "inputs" as some models crash if they don't exist.
+  if getattr(hparams, "prepend_mode", "none") != "none":
+    shape = tf.shape(x)
+    partial_targets = tf.reshape(x, [shape[0], shape[1]])
+    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
+    features["partial_targets"] = partial_targets
   return features
 
 
@@ -957,6 +964,13 @@ def _decode_input_tensor_to_features_dict(feature_map, hparams):
   features["decode_length"] = (
       IMAGE_DECODE_LENGTH if input_is_image else tf.shape(x)[1] + 50)
   features["inputs"] = x
+  # Save inputs to "partial_targets" when prepending inputs to targets. Also
+  # keep "inputs" as some models crash if they don't exist.
+  if getattr(hparams, "prepend_mode", "none") != "none":
+    shape = tf.shape(x)
+    partial_targets = tf.reshape(x, [shape[0], shape[1]])
+    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
+    features["partial_targets"] = partial_targets
   return features
 
 

From 1f7cbd16471c5cd73e8030b7acfa56286e0a1a6f Mon Sep 17 00:00:00 2001
From: Sepehr Sameni <Sepehr.Sameni@gmail.com>
Date: Thu, 21 Nov 2019 09:59:23 -0800
Subject: [PATCH 2578/2720] Merge of PR #1724

PiperOrigin-RevId: 281772522
---
 .../data_generators/text_problems.py          | 104 ------------------
 .../data_generators/text_problems_test.py     |  17 ---
 .../models/image_transformer_2d_test.py       |   4 +-
 .../models/research/universal_transformer.py  |   2 +
 tensor2tensor/models/transformer.py           |  19 ++--
 tensor2tensor/utils/decoding.py               |  14 ---
 tensor2tensor/utils/t2t_model.py              |  12 +-
 7 files changed, 15 insertions(+), 157 deletions(-)

diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index f807e6f51..4067a65aa 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -17,7 +17,6 @@
 
 * Text2TextProblem: input=text, target=text.
 * Text2ClassProblem: input=text, target=class.
-* Text2RealProblem: input=text, target=float.
 * Text2SelfProblem (for language modeling): target=text
 * QuestionAndContext2TextProblem: input=text, context=text, target=text.
 
@@ -606,94 +605,6 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
       yield {"inputs": inputs, "targets": [label]}
 
 
-class Text2RealProblem(Text2TextProblem):
-  """Base class for text regression problems with one or more tasks.
-    Suitable for text-based problems where targets are continuous, real values.
-    When ntasks = 1, each text example is mapped to a single scalar value. When
-    ntasks > 1, each text example is mapped to a 1-d vector of length ntasks.
-  """
-
-  @property
-  def ntasks(self):
-    """Set to n > 1 for multitask regression."""
-    return 1
-
-  def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    """Generate samples of text and real-valued target pairs.
-    Each yielded dict will be a single example. The inputs should be raw text.
-    The target should be a list containing ntasks floats.
-    Args:
-      data_dir: final data directory. Typically only used in this method to copy
-        over user-supplied vocab files (for example, if vocab_type ==
-        VocabType.TOKEN).
-      tmp_dir: temporary directory that you can use for downloading and scratch.
-      dataset_split: problem.DatasetSplit, which data split to generate samples
-        for (for example, training and evaluation).
-    Yields:
-      {"inputs": text, "targets": [x1, x2, ..., xN]} where N is ntasks
-    """
-    raise NotImplementedError()
-
-  def generate_text_for_vocab(self, data_dir, tmp_dir):
-    for i, sample in enumerate(
-        self.generate_samples(data_dir, tmp_dir, problem.DatasetSplit.TRAIN)):
-      yield sample["inputs"]
-      if self.max_samples_for_vocab and (i + 1) >= self.max_samples_for_vocab:
-        break
-
-  def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
-    generator = self.generate_samples(data_dir, tmp_dir, dataset_split)
-    encoder = self.get_or_create_vocab(data_dir, tmp_dir)
-    for sample in generator:
-      inputs = encoder.encode(sample["inputs"])
-      inputs.append(text_encoder.EOS_ID)
-      yield {"inputs": inputs, "targets": sample["targets"]}
-
-  def feature_encoders(self, data_dir):
-    encoder = self.get_or_create_vocab(data_dir, None, force_get=True)
-
-    return {
-        "inputs": encoder,
-        "targets": text_encoder.RealEncoder(),
-    }
-
-  def hparams(self, defaults, unused_model_hparams):
-    p = defaults
-    p.modality = {
-        "inputs": modalities.ModalityType.SYMBOL,
-        "targets": modalities.ModalityType.REAL_L2_LOSS,
-    }
-    p.vocab_size = {
-        "inputs": self._encoders["inputs"].vocab_size,
-        "targets": self.ntasks
-    }
-    p.target_space_id = problem.SpaceID.REAL
-    p.add_hparam("regression_targets", True)
-
-  def max_length(self, model_hparams):
-    return model_hparams.batch_size * self.ntasks
-
-  def preprocess_example(self, example, unused_mode, unused_hparams):
-    example = problem.preprocess_example_common(example, unused_mode,
-                                                unused_hparams)
-    example["targets"] = tf.reshape(example["targets"], [1, 1, self.ntasks])
-    return example
-
-  def example_reading_spec(self):
-    data_fields = {
-        "inputs": tf.VarLenFeature(tf.int64),
-        "targets": tf.FixedLenFeature([self.ntasks], tf.float32),
-    }
-    data_items_to_decoders = None
-    return (data_fields, data_items_to_decoders)
-
-  def eval_metrics(self):
-    metrics_list = [metrics.Metrics.RMSE]
-    if self.ntasks == 1:
-      metrics_list.append(metrics.Metrics.PEARSON)
-    return metrics_list
-
-
 def txt_line_iterator(txt_path):
   """Iterate through lines of file."""
   with tf.gfile.Open(txt_path) as f:
@@ -781,21 +692,6 @@ def text2class_txt_iterator(source_txt_path, label_txt_path, class_strs=None):
     yield {"inputs": inputs, "label": label}
 
 
-def text2real_txt_iterator(source_txt_path, target_txt_path):
-  """Yield dicts for Text2RealProblem.generate_samples from lines of files.
-  Args:
-    source_txt_path: txt file with record per line.
-    target_txt_path: txt file with float (or space-separated float list for
-      multitask) per line.
-  Yields:
-    {"inputs": inputs, "targets": targets}
-  """
-  for inputs, targets in zip(
-      txt_line_iterator(source_txt_path), txt_line_iterator(target_txt_path)):
-    targets = [float(x) for x in targets.split(" ")]
-    yield {"inputs": inputs, "targets": targets}
-
-
 def text2text_txt_tab_iterator(txt_path):
   """Yield dicts for Text2TextProblem.generate_samples from lines of txt_path.
 
diff --git a/tensor2tensor/data_generators/text_problems_test.py b/tensor2tensor/data_generators/text_problems_test.py
index 720f1ba68..51f948fcf 100644
--- a/tensor2tensor/data_generators/text_problems_test.py
+++ b/tensor2tensor/data_generators/text_problems_test.py
@@ -94,13 +94,6 @@ def setUpClass(cls):
     tf.gfile.Copy(cls.targets_file, os.path.join(cls.tmp_dir,
                                                  "targets.eval.txt"))
 
-    cls.targets_regr = [[1.23, 2.34], [4.56, 5.67]]
-    cls.targets_regr_file = os.path.join(cls.tmp_dir, "targets_regr.train.txt")
-    with tf.gfile.Open(cls.targets_regr_file, "w") as f:
-      for targets in cls.targets_regr:
-        f.write(" ".join([str(x) for x in targets]) + "\n")
-
-
   def testTxtLineIterator(self):
     lines = [line for line in text_problems.txt_line_iterator(self.inputs_file)]
     self.assertEqual(lines, self.inputs)
@@ -143,16 +136,6 @@ def testText2ClassTxtIteratorWithStrs(self):
     self.assertEqual(inputs, self.inputs)
     self.assertEqual(labels, self.labels)
 
-  def testText2RealTxtIterator(self):
-    inputs = []
-    targets = []
-    for entry in text_problems.text2real_txt_iterator(self.inputs_file,
-                                                      self.targets_regr_file):
-      inputs.append(entry["inputs"])
-      targets.append(entry["targets"])
-    self.assertEqual(inputs, self.inputs)
-    self.assertEqual(targets, self.targets_regr)
-
   def testText2TextTxtTabIterator(self):
     inputs = []
     targets = []
diff --git a/tensor2tensor/models/image_transformer_2d_test.py b/tensor2tensor/models/image_transformer_2d_test.py
index 7f903fb15..7deddc870 100644
--- a/tensor2tensor/models/image_transformer_2d_test.py
+++ b/tensor2tensor/models/image_transformer_2d_test.py
@@ -35,8 +35,8 @@ def _test_img2img_transformer(self, net):
     hparams = image_transformer_2d.img2img_transformer2d_tiny()
     hparams.data_dir = ""
     p_hparams = registry.problem("image_celeba").get_hparams(hparams)
-    inputs = np.random.randint(256, size=(3, 4, 4, 3))
-    targets = np.random.randint(256, size=(3, 8, 8, 3))
+    inputs = np.random.randint(256, size=(batch_size, 4, 4, 3))
+    targets = np.random.randint(256, size=(batch_size, 8, 8, 3))
     with self.test_session() as session:
       features = {
           "inputs": tf.constant(inputs, dtype=tf.int32),
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index a5a65504e..28040584e 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -458,6 +458,7 @@ def universal_transformer_base():
 @registry.register_hparams
 def universal_transformer_base_tpu():
   hparams = universal_transformer_base()
+  hparams = update_hparams_for_universal_transformer(hparams)
   transformer.update_hparams_for_tpu(hparams)
   hparams.add_step_timing_signal = False
   return hparams
@@ -466,6 +467,7 @@ def universal_transformer_base_tpu():
 @registry.register_hparams
 def universal_transformer_big():
   hparams = universal_transformer_base()
+  hparams = update_hparams_for_universal_transformer(hparams)
   hparams.hidden_size = 2048
   hparams.filter_size = 8192
   return hparams
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index d2dda17ec..eef24dba1 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -462,8 +462,7 @@ def _fast_decode_tpu(self,
 
     if self.has_input:
       inputs_shape = common_layers.shape_list(features["inputs"])
-      if (target_modality == modalities.ModalityType.CLASS_LABEL or
-          self._problem_hparams.get("regression_targets")):
+      if target_modality == modalities.ModalityType.CLASS_LABEL:
         decode_length = 1
       else:
         decode_length = (
@@ -705,8 +704,7 @@ def _fast_decode(self,
           " of the dataset when decoding.")
     if self.has_input:
       inputs_shape = common_layers.shape_list(features["inputs"])
-      if (target_modality == modalities.ModalityType.CLASS_LABEL or
-          self._problem_hparams.get("regression_targets")):
+      if target_modality == modalities.ModalityType.CLASS_LABEL:
         decode_length = 1
       else:
         decode_length = (
@@ -865,15 +863,9 @@ def symbols_to_logits_fn(ids, i, cache):
         vocab_size = tf.shape(ret)[1]
 
         def forced_logits():
-          # Workaround for: tf.one_hot(
-          #               tf.repeat(partial_targets[:, i], [beam_size]), vocab_size, 0.0,
-          #               -1e9)
-          # Can be replaced by the above in future versions (from tf 1.15).
           return tf.one_hot(
-              tf.reshape(tf.tile(
-                  tf.reshape(partial_targets[:, i], [-1, 1]),
-                  [1, beam_size]), [-1]),
-              vocab_size, 0.0, -1e9)
+              tf.tile(partial_targets[:, i], [beam_size]), vocab_size, 0.0,
+              -1e9)
 
         ret = tf.cond(
             tf.less(i, partial_targets_length), forced_logits, lambda: ret)
@@ -1176,6 +1168,9 @@ def fast_decode(encoder_output,
           "scores": decoding log probs from the beam search,
               None if using greedy decoding (beam_size=1)
       }
+
+    Raises:
+      NotImplementedError: If beam size > 1 with partial targets.
   """
   if encoder_output is not None:
     batch_size = common_layers.shape_list(encoder_output)[0]
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index a63e83d5b..656edfdfd 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -927,13 +927,6 @@ def _interactive_input_tensor_to_features_dict(feature_map, hparams):
   features["decode_length"] = (
       IMAGE_DECODE_LENGTH if input_is_image else inputs[1])
   features["inputs"] = x
-  # Save inputs to "partial_targets" when prepending inputs to targets. Also
-  # keep "inputs" as some models crash if they don't exist.
-  if getattr(hparams, "prepend_mode", "none") != "none":
-    shape = tf.shape(x)
-    partial_targets = tf.reshape(x, [shape[0], shape[1]])
-    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
-    features["partial_targets"] = partial_targets
   return features
 
 
@@ -964,13 +957,6 @@ def _decode_input_tensor_to_features_dict(feature_map, hparams):
   features["decode_length"] = (
       IMAGE_DECODE_LENGTH if input_is_image else tf.shape(x)[1] + 50)
   features["inputs"] = x
-  # Save inputs to "partial_targets" when prepending inputs to targets. Also
-  # keep "inputs" as some models crash if they don't exist.
-  if getattr(hparams, "prepend_mode", "none") != "none":
-    shape = tf.shape(x)
-    partial_targets = tf.reshape(x, [shape[0], shape[1]])
-    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
-    features["partial_targets"] = partial_targets
   return features
 
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 2b9e419ef..fe33e6315 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -806,10 +806,8 @@ def infer(self,
 
       if self._problem_hparams:
         target_modality = self._problem_hparams.modality["targets"]
-      if (target_modality == modalities.ModalityType.CLASS_LABEL or
-          self._problem_hparams.get("regression_targets")):
-        # No use to run beam-search for classification or regression.
-        beam_size = 1
+        if target_modality == modalities.ModalityType.CLASS_LABEL:
+          beam_size = 1  # No use to run beam-search for a single class.
       if beam_size == 1:
         log_info("Greedy Decoding")
         results = self._greedy_infer(features, decode_length, use_tpu)
@@ -1066,8 +1064,7 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
     initial_output = tf.slice(initial_output, [0, 0, 0, 0],
                               common_layers.shape_list(initial_output))
     target_modality = self._problem_hparams.modality["targets"]
-    if (target_modality == modalities.ModalityType.CLASS_LABEL or
-        self._problem_hparams.get("regression_targets")):
+    if target_modality == modalities.ModalityType.CLASS_LABEL:
       decode_length = 1
     else:
       if "partial_targets" in features:
@@ -1246,8 +1243,7 @@ def infer_step(recent_output, recent_logits, unused_loss):
     initial_output = tf.slice(initial_output, [0, 0, 0, 0],
                               common_layers.shape_list(initial_output))
     target_modality = self._problem_hparams.modality["targets"]
-    if (target_modality == modalities.ModalityType.CLASS_LABEL or
-        self._problem_hparams.get("regression_targets")):
+    if target_modality == modalities.ModalityType.CLASS_LABEL:
       decode_length = 1
     else:
       if "partial_targets" in features:

From a4a0cf0af376dba9b9d66ddb43c4e08579df16cc Mon Sep 17 00:00:00 2001
From: Gabe Grand <gabrieljgrand@gmail.com>
Date: Thu, 21 Nov 2019 12:01:53 -0800
Subject: [PATCH 2579/2720] Merge of PR #1748

PiperOrigin-RevId: 281802220
---
 .../data_generators/text_problems.py          | 107 ++++++++++++++++++
 .../data_generators/text_problems_test.py     |  16 +++
 tensor2tensor/models/transformer.py           |   6 +-
 tensor2tensor/utils/t2t_model.py              |  12 +-
 4 files changed, 135 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 4067a65aa..a039c8a13 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -17,6 +17,7 @@
 
 * Text2TextProblem: input=text, target=text.
 * Text2ClassProblem: input=text, target=class.
+* Text2RealProblem: input=text, target=float.
 * Text2SelfProblem (for language modeling): target=text
 * QuestionAndContext2TextProblem: input=text, context=text, target=text.
 
@@ -605,6 +606,96 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
       yield {"inputs": inputs, "targets": [label]}
 
 
+class Text2RealProblem(Text2TextProblem):
+  """Base class for text regression problems with one or more tasks.
+
+    Suitable for text-based problems where targets are continuous, real values.
+    When ntasks = 1, each text example is mapped to a single scalar value. When
+    ntasks > 1, each text example is mapped to a 1-d vector of length ntasks.
+  """
+
+  @property
+  def ntasks(self):
+    """Set to n > 1 for multitask regression."""
+    return 1
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    """Generate samples of text and real-valued target pairs.
+
+    Each yielded dict will be a single example. The inputs should be raw text.
+    The target should be a list containing ntasks floats.
+    Args:
+      data_dir: final data directory. Typically only used in this method to copy
+        over user-supplied vocab files (for example, if vocab_type ==
+        VocabType.TOKEN).
+      tmp_dir: temporary directory that you can use for downloading and scratch.
+      dataset_split: problem.DatasetSplit, which data split to generate samples
+        for (for example, training and evaluation).
+    Yields:
+      {"inputs": text, "targets": [x1, x2, ..., xN]} where N is ntasks
+    """
+    raise NotImplementedError()
+
+  def generate_text_for_vocab(self, data_dir, tmp_dir):
+    for i, sample in enumerate(
+        self.generate_samples(data_dir, tmp_dir, problem.DatasetSplit.TRAIN)):
+      yield sample["inputs"]
+      if self.max_samples_for_vocab and (i + 1) >= self.max_samples_for_vocab:
+        break
+
+  def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
+    generator = self.generate_samples(data_dir, tmp_dir, dataset_split)
+    encoder = self.get_or_create_vocab(data_dir, tmp_dir)
+    for sample in generator:
+      inputs = encoder.encode(sample["inputs"])
+      inputs.append(text_encoder.EOS_ID)
+      yield {"inputs": inputs, "targets": sample["targets"]}
+
+  def feature_encoders(self, data_dir):
+    encoder = self.get_or_create_vocab(data_dir, None, force_get=True)
+
+    return {
+        "inputs": encoder,
+        "targets": text_encoder.RealEncoder(),
+    }
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    p.modality = {
+        "inputs": modalities.ModalityType.SYMBOL,
+        "targets": modalities.ModalityType.REAL_L2_LOSS,
+    }
+    p.vocab_size = {
+        "inputs": self._encoders["inputs"].vocab_size,
+        "targets": self.ntasks
+    }
+    p.target_space_id = problem.SpaceID.REAL
+    p.add_hparam("regression_targets", True)
+
+  def max_length(self, model_hparams):
+    return model_hparams.batch_size * self.ntasks
+
+  def preprocess_example(self, example, unused_mode, unused_hparams):
+    example = problem.preprocess_example_common(example, unused_mode,
+                                                unused_hparams)
+    example["targets"] = tf.reshape(example["targets"], [1, 1, self.ntasks])
+    return example
+
+  def example_reading_spec(self):
+    data_fields = {
+        "inputs": tf.VarLenFeature(tf.int64),
+        "targets": tf.FixedLenFeature([self.ntasks], tf.float32),
+    }
+    data_items_to_decoders = None
+    return (data_fields, data_items_to_decoders)
+
+  def eval_metrics(self):
+    metrics_list = [metrics.Metrics.RMSE]
+    if self.ntasks == 1:
+      metrics_list.append(metrics.Metrics.PEARSON)
+    return metrics_list
+
+
 def txt_line_iterator(txt_path):
   """Iterate through lines of file."""
   with tf.gfile.Open(txt_path) as f:
@@ -692,6 +783,22 @@ def text2class_txt_iterator(source_txt_path, label_txt_path, class_strs=None):
     yield {"inputs": inputs, "label": label}
 
 
+def text2real_txt_iterator(source_txt_path, target_txt_path):
+  """Yield dicts for Text2RealProblem.generate_samples from lines of files.
+
+  Args:
+    source_txt_path: txt file with record per line.
+    target_txt_path: txt file with float (or space-separated float list for
+      multitask) per line.
+  Yields:
+    {"inputs": inputs, "targets": targets}
+  """
+  for inputs, targets in zip(
+      txt_line_iterator(source_txt_path), txt_line_iterator(target_txt_path)):
+    targets = [float(x) for x in targets.split(" ")]
+    yield {"inputs": inputs, "targets": targets}
+
+
 def text2text_txt_tab_iterator(txt_path):
   """Yield dicts for Text2TextProblem.generate_samples from lines of txt_path.
 
diff --git a/tensor2tensor/data_generators/text_problems_test.py b/tensor2tensor/data_generators/text_problems_test.py
index 51f948fcf..b253421b3 100644
--- a/tensor2tensor/data_generators/text_problems_test.py
+++ b/tensor2tensor/data_generators/text_problems_test.py
@@ -94,6 +94,12 @@ def setUpClass(cls):
     tf.gfile.Copy(cls.targets_file, os.path.join(cls.tmp_dir,
                                                  "targets.eval.txt"))
 
+    cls.targets_regr = [[1.23, 2.34], [4.56, 5.67]]
+    cls.targets_regr_file = os.path.join(cls.tmp_dir, "targets_regr.train.txt")
+    with tf.gfile.Open(cls.targets_regr_file, "w") as f:
+      for targets in cls.targets_regr:
+        f.write(" ".join([str(x) for x in targets]) + "\n")
+
   def testTxtLineIterator(self):
     lines = [line for line in text_problems.txt_line_iterator(self.inputs_file)]
     self.assertEqual(lines, self.inputs)
@@ -136,6 +142,16 @@ def testText2ClassTxtIteratorWithStrs(self):
     self.assertEqual(inputs, self.inputs)
     self.assertEqual(labels, self.labels)
 
+  def testText2RealTxtIterator(self):
+    inputs = []
+    targets = []
+    for entry in text_problems.text2real_txt_iterator(self.inputs_file,
+                                                      self.targets_regr_file):
+      inputs.append(entry["inputs"])
+      targets.append(entry["targets"])
+    self.assertEqual(inputs, self.inputs)
+    self.assertEqual(targets, self.targets_regr)
+
   def testText2TextTxtTabIterator(self):
     inputs = []
     targets = []
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index eef24dba1..c1c75f121 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -462,7 +462,8 @@ def _fast_decode_tpu(self,
 
     if self.has_input:
       inputs_shape = common_layers.shape_list(features["inputs"])
-      if target_modality == modalities.ModalityType.CLASS_LABEL:
+      if (target_modality == modalities.ModalityType.CLASS_LABEL or
+          self._problem_hparams.get("regression_targets")):
         decode_length = 1
       else:
         decode_length = (
@@ -704,7 +705,8 @@ def _fast_decode(self,
           " of the dataset when decoding.")
     if self.has_input:
       inputs_shape = common_layers.shape_list(features["inputs"])
-      if target_modality == modalities.ModalityType.CLASS_LABEL:
+      if (target_modality == modalities.ModalityType.CLASS_LABEL or
+          self._problem_hparams.get("regression_targets")):
         decode_length = 1
       else:
         decode_length = (
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index fe33e6315..2b9e419ef 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -806,8 +806,10 @@ def infer(self,
 
       if self._problem_hparams:
         target_modality = self._problem_hparams.modality["targets"]
-        if target_modality == modalities.ModalityType.CLASS_LABEL:
-          beam_size = 1  # No use to run beam-search for a single class.
+      if (target_modality == modalities.ModalityType.CLASS_LABEL or
+          self._problem_hparams.get("regression_targets")):
+        # No use to run beam-search for classification or regression.
+        beam_size = 1
       if beam_size == 1:
         log_info("Greedy Decoding")
         results = self._greedy_infer(features, decode_length, use_tpu)
@@ -1064,7 +1066,8 @@ def infer_step(i, recent_output, recent_logits, unused_loss):
     initial_output = tf.slice(initial_output, [0, 0, 0, 0],
                               common_layers.shape_list(initial_output))
     target_modality = self._problem_hparams.modality["targets"]
-    if target_modality == modalities.ModalityType.CLASS_LABEL:
+    if (target_modality == modalities.ModalityType.CLASS_LABEL or
+        self._problem_hparams.get("regression_targets")):
       decode_length = 1
     else:
       if "partial_targets" in features:
@@ -1243,7 +1246,8 @@ def infer_step(recent_output, recent_logits, unused_loss):
     initial_output = tf.slice(initial_output, [0, 0, 0, 0],
                               common_layers.shape_list(initial_output))
     target_modality = self._problem_hparams.modality["targets"]
-    if target_modality == modalities.ModalityType.CLASS_LABEL:
+    if (target_modality == modalities.ModalityType.CLASS_LABEL or
+        self._problem_hparams.get("regression_targets")):
       decode_length = 1
     else:
       if "partial_targets" in features:

From 2efd8aa349d4022a7eb056ffd53c25ffc2afbaa4 Mon Sep 17 00:00:00 2001
From: Seppo Enarvi <seppo.git@marjaniemi.com>
Date: Thu, 21 Nov 2019 12:23:15 -0800
Subject: [PATCH 2580/2720] Merge of PR #1726

PiperOrigin-RevId: 281806525
---
 .../models/research/universal_transformer.py       |  2 --
 tensor2tensor/models/transformer.py                |  3 ---
 tensor2tensor/utils/decoding.py                    | 14 ++++++++++++++
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 28040584e..a5a65504e 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -458,7 +458,6 @@ def universal_transformer_base():
 @registry.register_hparams
 def universal_transformer_base_tpu():
   hparams = universal_transformer_base()
-  hparams = update_hparams_for_universal_transformer(hparams)
   transformer.update_hparams_for_tpu(hparams)
   hparams.add_step_timing_signal = False
   return hparams
@@ -467,7 +466,6 @@ def universal_transformer_base_tpu():
 @registry.register_hparams
 def universal_transformer_big():
   hparams = universal_transformer_base()
-  hparams = update_hparams_for_universal_transformer(hparams)
   hparams.hidden_size = 2048
   hparams.filter_size = 8192
   return hparams
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index c1c75f121..aeb4d63e9 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1170,9 +1170,6 @@ def fast_decode(encoder_output,
           "scores": decoding log probs from the beam search,
               None if using greedy decoding (beam_size=1)
       }
-
-    Raises:
-      NotImplementedError: If beam size > 1 with partial targets.
   """
   if encoder_output is not None:
     batch_size = common_layers.shape_list(encoder_output)[0]
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 656edfdfd..a63e83d5b 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -927,6 +927,13 @@ def _interactive_input_tensor_to_features_dict(feature_map, hparams):
   features["decode_length"] = (
       IMAGE_DECODE_LENGTH if input_is_image else inputs[1])
   features["inputs"] = x
+  # Save inputs to "partial_targets" when prepending inputs to targets. Also
+  # keep "inputs" as some models crash if they don't exist.
+  if getattr(hparams, "prepend_mode", "none") != "none":
+    shape = tf.shape(x)
+    partial_targets = tf.reshape(x, [shape[0], shape[1]])
+    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
+    features["partial_targets"] = partial_targets
   return features
 
 
@@ -957,6 +964,13 @@ def _decode_input_tensor_to_features_dict(feature_map, hparams):
   features["decode_length"] = (
       IMAGE_DECODE_LENGTH if input_is_image else tf.shape(x)[1] + 50)
   features["inputs"] = x
+  # Save inputs to "partial_targets" when prepending inputs to targets. Also
+  # keep "inputs" as some models crash if they don't exist.
+  if getattr(hparams, "prepend_mode", "none") != "none":
+    shape = tf.shape(x)
+    partial_targets = tf.reshape(x, [shape[0], shape[1]])
+    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
+    features["partial_targets"] = partial_targets
   return features
 
 

From 2565fbec2291320b4e6eb58fb8cab7527f89f39f Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 21 Nov 2019 17:08:31 -0800
Subject: [PATCH 2581/2720] Skip
 `CommonLayersTest::testFactoredTensorImplicitConversion` and
 `ModalityTest::testSymbolModalityTargetsFactored` since this has changed
 slightly internally and externally.

PiperOrigin-RevId: 281866052
---
 oss_scripts/oss_tests.sh | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index ede0fa1f8..f85b21917 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -108,6 +108,8 @@ pytest --disable-warnings \
   tensor2tensor/layers/ngram_test.py \
   tensor2tensor/utils/t2t_model_test.py \
   tensor2tensor/utils/test_utils_test.py \
+  --deselect=tensor2tensor/layers/common_layers_test.py::CommonLayersTest::testFactoredTensorImplicitConversion \
+  --deselect=tensor2tensor/layers/modalities_test.py::ModalityTest::testSymbolModalityTargetsFactored \
   --deselect=tensor2tensor/layers/common_video_test.py::CommonVideoTest::testGifSummary
 set_status
 

From 1317424c1c6b073bcb9d997bd5468b5e537f5a1f Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 21 Nov 2019 17:27:24 -0800
Subject: [PATCH 2582/2720] Move from tf flags to absl flags, since they seem
 to be gone in tf 1.15

PiperOrigin-RevId: 281869018
---
 setup.py                               | 2 +-
 tensor2tensor/data_generators/audio.py | 8 +++-----
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/setup.py b/setup.py
index de1a41645..3248053c7 100644
--- a/setup.py
+++ b/setup.py
@@ -33,6 +33,7 @@
         'tensor2tensor/bin/t2t-translate-all',
     ],
     install_requires=[
+        'absl-py',
         'bz2file',
         'dopamine-rl',
         'flask',
@@ -64,7 +65,6 @@
         'tensorflow': ['tensorflow>=1.15.0'],
         'tensorflow-hub': ['tensorflow-hub>=0.1.1'],
         'tests': [
-            'absl-py',
             # Needed to fix a Travis pytest error.
             # https://github.com/Julian/jsonschema/issues/449#issuecomment-411406525
             'attrs>=17.4.0',
diff --git a/tensor2tensor/data_generators/audio.py b/tensor2tensor/data_generators/audio.py
index 7a3480f6f..7a5d5d64f 100644
--- a/tensor2tensor/data_generators/audio.py
+++ b/tensor2tensor/data_generators/audio.py
@@ -19,14 +19,12 @@
 from __future__ import print_function
 
 import os
-from subprocess import call
+import subprocess
 import tarfile
 import wave
-# from tensor2tensor.data_generators import generator_utils
-
+from absl import flags
 import tensorflow as tf
 
-flags = tf.flags
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string("timit_paths", "",
@@ -78,7 +76,7 @@ def _get_audio_data(filepath):
   # Construct a true .wav file.
   out_filepath = filepath.strip(".WAV") + ".wav"
   # Assumes sox is installed on system. Sox converts from NIST SPHERE to WAV.
-  call(["sox", filepath, out_filepath])
+  subprocess.call(["sox", filepath, out_filepath])
   wav_file = wave.open(open(out_filepath))
   frame_count = wav_file.getnframes()
   byte_array = wav_file.readframes(frame_count)

From d979d9f11a767a0c6131defbce28b8f7ba966a37 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 21 Nov 2019 20:53:30 -0800
Subject: [PATCH 2583/2720] pip installing tensorflow serving seems to be
 uninstalling the installed version of tensorflow (1.15.0) and replacing with
 2.0, leading to errors of missing contrib -- so install a pinned version.

Make an error in all_problems.py log earlier.

Reverting the earlier change about tf flags not being shipped with TF, because of this TF 2.0 subterfuge.

PiperOrigin-RevId: 281893644
---
 oss_scripts/oss_integration_test.sh           | 2 +-
 setup.py                                      | 2 +-
 tensor2tensor/data_generators/all_problems.py | 4 ++--
 tensor2tensor/data_generators/audio.py        | 3 ++-
 4 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/oss_scripts/oss_integration_test.sh b/oss_scripts/oss_integration_test.sh
index f5bd63e3d..46ad0fe06 100755
--- a/oss_scripts/oss_integration_test.sh
+++ b/oss_scripts/oss_integration_test.sh
@@ -40,7 +40,7 @@ then
   sleep 10
 
   # Query
-  pip install tensorflow-serving-api
+  pip install tensorflow-serving-api=="$TF_VERSION"
   t2t-query-server \
       --server=localhost:$server_port \
       --servable_name=$model_name \
diff --git a/setup.py b/setup.py
index 3248053c7..de1a41645 100644
--- a/setup.py
+++ b/setup.py
@@ -33,7 +33,6 @@
         'tensor2tensor/bin/t2t-translate-all',
     ],
     install_requires=[
-        'absl-py',
         'bz2file',
         'dopamine-rl',
         'flask',
@@ -65,6 +64,7 @@
         'tensorflow': ['tensorflow>=1.15.0'],
         'tensorflow-hub': ['tensorflow-hub>=0.1.1'],
         'tests': [
+            'absl-py',
             # Needed to fix a Travis pytest error.
             # https://github.com/Julian/jsonschema/issues/449#issuecomment-411406525
             'attrs>=17.4.0',
diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 3efdc9cd3..f39addd71 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -126,11 +126,11 @@ def _handle_errors(errors):
   print(err_msg.format(num_missing=len(errors)))
   for module, err in errors:
     err_str = str(err)
+    if log_all:
+      print("Did not import module: %s; Cause: %s" % (module, err_str))
     if not _is_import_err_msg(err_str, module):
       print("From module %s" % module)
       raise err
-    if log_all:
-      print("Did not import module: %s; Cause: %s" % (module, err_str))
 
 
 def import_modules(modules):
diff --git a/tensor2tensor/data_generators/audio.py b/tensor2tensor/data_generators/audio.py
index 7a5d5d64f..4f7a65a14 100644
--- a/tensor2tensor/data_generators/audio.py
+++ b/tensor2tensor/data_generators/audio.py
@@ -22,9 +22,10 @@
 import subprocess
 import tarfile
 import wave
-from absl import flags
+
 import tensorflow as tf
 
+flags = tf.flags
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string("timit_paths", "",

From 259406973a597a7630784320509559fb60595184 Mon Sep 17 00:00:00 2001
From: Dero Gharibian <dero@google.com>
Date: Fri, 22 Nov 2019 14:30:57 -0800
Subject: [PATCH 2584/2720] Migrate from std::string to tensorflow::tstring.

Note that during the transition period tstring is typedef'ed to std::string.

See: https://github.com/tensorflow/community/pull/91
PiperOrigin-RevId: 282043204
---
 tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc b/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
index d0ba6ec8f..9e8959f96 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
@@ -14,6 +14,7 @@ using ::tensorflow::OpKernelContext;
 using ::tensorflow::Status;
 using ::tensorflow::Tensor;
 using ::tensorflow::TensorShape;
+using ::tensorflow::tstring;
 using ::tensorflow::shape_inference::InferenceContext;
 
 REGISTER_OP("SubwordTextEncoderEncode")
@@ -36,7 +37,7 @@ class SubwordTextEncoderEncodeOp : public OpKernel {
 
   void Compute(OpKernelContext* ctx) override {
     // Get input string and deserialize into ArticleExample proto.
-    const string& s = ctx->input(0).scalar<string>()();
+    absl::string_view s = ctx->input(0).scalar<tstring>()();
 
     // Construct encoded output tensors.
     std::vector<int> encoded_ids;

From 313fdfcf94062280ebde5902a0dd771431dd5c47 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 22 Nov 2019 15:39:47 -0800
Subject: [PATCH 2585/2720] Move from tf flags to absl flags, since they seem
 to be gone in tf 1.15 NOTE: This is only for libraries (and not binaries) in
 data_generators/

Also bump up version to 1.15.1

Also add a description to show on PyPI

PiperOrigin-RevId: 282056584
---
 setup.py                                            | 12 ++++++++++--
 tensor2tensor/data_generators/audio.py              |  3 +--
 tensor2tensor/data_generators/dialog_personachat.py |  2 --
 tensor2tensor/data_generators/tokenizer_test.py     |  1 -
 tensor2tensor/data_generators/translate.py          |  2 --
 tensor2tensor/data_generators/translate_encs.py     |  3 ---
 tensor2tensor/data_generators/translate_enfr.py     |  3 ---
 tensor2tensor/data_generators/translate_enzh.py     |  1 -
 tensor2tensor/data_generators/video_utils.py        |  6 +++---
 tensor2tensor/data_generators/wiki_revision.py      |  2 +-
 tensor2tensor/data_generators/wsj_parsing.py        |  5 +++--
 11 files changed, 18 insertions(+), 22 deletions(-)

diff --git a/setup.py b/setup.py
index de1a41645..0c6df6f73 100644
--- a/setup.py
+++ b/setup.py
@@ -5,8 +5,16 @@
 
 setup(
     name='tensor2tensor',
-    version='1.15.0',
+    version='1.15.1',
     description='Tensor2Tensor',
+    long_description=(
+        'Tensor2Tensor, or T2T for short, is a library of '
+        'deep learning models and datasets designed to make deep '
+        'learning more accessible and accelerate ML research. '
+        'T2T was developed by researchers and engineers in the Google '
+        'Brain team and a community of users. It is now in maintenance '
+        'mode -- we keep it running and welcome bug-fixes, but encourage '
+        'users to use the successor library Trax.'),
     author='Google Inc.',
     author_email='no-reply@google.com',
     url='http://github.com/tensorflow/tensor2tensor',
@@ -33,6 +41,7 @@
         'tensor2tensor/bin/t2t-translate-all',
     ],
     install_requires=[
+        'absl-py',
         'bz2file',
         'dopamine-rl',
         'flask',
@@ -64,7 +73,6 @@
         'tensorflow': ['tensorflow>=1.15.0'],
         'tensorflow-hub': ['tensorflow-hub>=0.1.1'],
         'tests': [
-            'absl-py',
             # Needed to fix a Travis pytest error.
             # https://github.com/Julian/jsonschema/issues/449#issuecomment-411406525
             'attrs>=17.4.0',
diff --git a/tensor2tensor/data_generators/audio.py b/tensor2tensor/data_generators/audio.py
index 4f7a65a14..7a5d5d64f 100644
--- a/tensor2tensor/data_generators/audio.py
+++ b/tensor2tensor/data_generators/audio.py
@@ -22,10 +22,9 @@
 import subprocess
 import tarfile
 import wave
-
+from absl import flags
 import tensorflow as tf
 
-flags = tf.flags
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string("timit_paths", "",
diff --git a/tensor2tensor/data_generators/dialog_personachat.py b/tensor2tensor/data_generators/dialog_personachat.py
index 0d848772c..ed085badc 100644
--- a/tensor2tensor/data_generators/dialog_personachat.py
+++ b/tensor2tensor/data_generators/dialog_personachat.py
@@ -27,9 +27,7 @@
 from tensor2tensor.data_generators import dialog_abstract
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.utils import registry
-import tensorflow as tf
 
-FLAGS = tf.flags.FLAGS
 
 # End-of-sentence marker.
 EOS = text_encoder.EOS_ID
diff --git a/tensor2tensor/data_generators/tokenizer_test.py b/tensor2tensor/data_generators/tokenizer_test.py
index 932529d90..a5bc99657 100644
--- a/tensor2tensor/data_generators/tokenizer_test.py
+++ b/tensor2tensor/data_generators/tokenizer_test.py
@@ -27,7 +27,6 @@
 from tensor2tensor.data_generators import tokenizer
 import tensorflow as tf
 
-FLAGS = tf.flags.FLAGS
 
 pkg_dir, _ = os.path.split(__file__)
 _TESTDATA = os.path.join(pkg_dir, "test_data")
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index b74c61ea5..8ac2f2575 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -34,8 +34,6 @@
 import tensorflow as tf
 import tf_slim as slim
 
-FLAGS = tf.flags.FLAGS
-
 
 class TranslateProblem(text_problems.Text2TextProblem):
   """Base class for translation problems."""
diff --git a/tensor2tensor/data_generators/translate_encs.py b/tensor2tensor/data_generators/translate_encs.py
index 5fd7a321b..7bd0fffc0 100644
--- a/tensor2tensor/data_generators/translate_encs.py
+++ b/tensor2tensor/data_generators/translate_encs.py
@@ -24,9 +24,6 @@
 from tensor2tensor.data_generators import translate
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
-
-FLAGS = tf.flags.FLAGS
 
 # End-of-sentence marker.
 EOS = text_encoder.EOS_ID
diff --git a/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index 4484b5c4e..81fd5e4b1 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -28,9 +28,6 @@
 from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
-
-FLAGS = tf.flags.FLAGS
 
 # End-of-sentence marker.
 EOS = text_encoder.EOS_ID
diff --git a/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index e5ed9cddb..c3a7370f8 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -29,7 +29,6 @@
 
 import tensorflow as tf
 
-FLAGS = tf.flags.FLAGS
 
 # End-of-sentence marker.
 EOS = text_encoder.EOS_ID
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 98683c0e6..4d3cec63c 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -21,9 +21,10 @@
 
 import functools
 import os
+
+from absl import flags
 import numpy as np
 import six
-
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.data_generators import problem
@@ -33,11 +34,10 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import video_metrics
-
 import tensorflow as tf
 import tf_slim as slim
 
-flags = tf.flags
+
 FLAGS = flags.FLAGS
 
 flags.DEFINE_bool(
diff --git a/tensor2tensor/data_generators/wiki_revision.py b/tensor2tensor/data_generators/wiki_revision.py
index 089c941c1..f41da7556 100644
--- a/tensor2tensor/data_generators/wiki_revision.py
+++ b/tensor2tensor/data_generators/wiki_revision.py
@@ -31,6 +31,7 @@
 import math
 import random
 
+from absl import flags
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
@@ -40,7 +41,6 @@
 
 import tensorflow as tf
 
-flags = tf.flags
 FLAGS = flags.FLAGS
 
 flags.DEFINE_integer("wiki_revision_num_train_shards", 50,
diff --git a/tensor2tensor/data_generators/wsj_parsing.py b/tensor2tensor/data_generators/wsj_parsing.py
index ae42e73d2..40225e43b 100644
--- a/tensor2tensor/data_generators/wsj_parsing.py
+++ b/tensor2tensor/data_generators/wsj_parsing.py
@@ -21,15 +21,16 @@
 
 import os
 
+from absl import flags
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 import tensorflow as tf
 
-tf.flags.DEFINE_string("parsing_path", "", "Path to parsing files in tmp_dir.")
+flags.DEFINE_string("parsing_path", "", "Path to parsing files in tmp_dir.")
 
 
-FLAGS = tf.flags.FLAGS
+FLAGS = flags.FLAGS
 
 
 @registry.register_problem

From 8107e53e8d8896c521522e9335a266b4b7510c37 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 23 Nov 2019 00:14:17 -0800
Subject: [PATCH 2586/2720] Changes needed for T2T's problems to be imported on
 TF 2.0 -- this doesn't mean we are using these problems.

This is only so that we can `import trax` on these environments.

PiperOrigin-RevId: 282108974
---
 setup.py                                      |  3 +--
 tensor2tensor/data_generators/allen_brain.py  |  3 +--
 .../data_generators/bair_robot_pushing.py     |  7 +++----
 tensor2tensor/data_generators/fsns.py         |  3 +--
 tensor2tensor/data_generators/gym_env.py      |  3 +--
 tensor2tensor/data_generators/image_utils.py  |  7 +++----
 tensor2tensor/data_generators/moving_mnist.py |  3 +--
 tensor2tensor/data_generators/problem.py      | 19 ++++++++++++-------
 .../data_generators/style_transfer.py         | 12 ++++--------
 tensor2tensor/data_generators/translate.py    |  5 ++---
 .../data_generators/video_generated.py        |  3 +--
 tensor2tensor/data_generators/video_utils.py  |  7 +++----
 tensor2tensor/data_generators/vqa.py          | 17 ++++++++---------
 tensor2tensor/envs/env_problem.py             |  3 +--
 tensor2tensor/envs/rendered_env_problem.py    |  3 +--
 tensor2tensor/envs/trajectory.py              |  8 ++++----
 tensor2tensor/layers/common_attention.py      |  3 ++-
 tensor2tensor/layers/common_layers.py         |  9 +++++++--
 tensor2tensor/layers/common_video.py          |  9 ++++++---
 tensor2tensor/utils/data_reader.py            |  3 +--
 tensor2tensor/utils/metrics.py                |  4 +++-
 tensor2tensor/utils/trainer_lib.py            |  2 +-
 22 files changed, 67 insertions(+), 69 deletions(-)

diff --git a/setup.py b/setup.py
index 0c6df6f73..b68e4ce6a 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.15.1',
+    version='1.15.2',
     description='Tensor2Tensor',
     long_description=(
         'Tensor2Tensor, or T2T for short, is a library of '
@@ -66,7 +66,6 @@
         'tensorflow-datasets',
         'tensorflow-gan',
         'tensorflow-probability==0.7.0',
-        'tf_slim',
         'tqdm',
     ],
     extras_require={
diff --git a/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index 2a738885a..2f58de3e9 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -44,7 +44,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 _BASE_EXAMPLE_IMAGE_SIZE = 64
 
@@ -351,7 +350,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "targets":
-            slim.tfexample_decoder.Image(
+            tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index b4e094dfb..eadd3275f 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -36,7 +36,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 DATA_URL = (
     "http://rail.eecs.berkeley.edu/datasets/bair_robot_pushing_dataset_v0.tar")
@@ -103,7 +102,7 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": slim.tfexample_decoder.Tensor(
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
     }
     return data_fields, decoders
@@ -188,9 +187,9 @@ def extra_reading_spec(self):
         "action": tf.FixedLenFeature([4], tf.float32),
     }
     decoders = {
-        "frame_number": slim.tfexample_decoder.Tensor(
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
-        "action": slim.tfexample_decoder.Tensor(tensor_key="action"),
+        "action": tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="action"),
     }
     return data_fields, decoders
 
diff --git a/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
index 4792263d7..8dff69547 100644
--- a/tensor2tensor/data_generators/fsns.py
+++ b/tensor2tensor/data_generators/fsns.py
@@ -28,7 +28,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 @registry.register_problem
@@ -77,5 +76,5 @@ def example_reading_spec(self):
         super(ImageFSNS, self).example_reading_spec())
     data_fields[label_key] = tf.VarLenFeature(tf.int64)
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index c325772cf..67aedef9b 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -36,7 +36,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 Frame = collections.namedtuple(
@@ -378,7 +377,7 @@ def extra_reading_spec(self):
         name: tf.FixedLenFeature([1], tf.int64) for name in field_names
     }
     decoders = {
-        name: slim.tfexample_decoder.Tensor(tensor_key=name)
+        name: tf.contrib.slim.tfexample_decoder.Tensor(tensor_key=name)
         for name in field_names
     }
     return (data_fields, decoders)
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 23c38c543..61766a89e 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -31,7 +31,6 @@
 from tensor2tensor.utils import metrics
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 def matplotlib_pyplot():
@@ -173,7 +172,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "inputs":
-            slim.tfexample_decoder.Image(
+            tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
@@ -240,7 +239,7 @@ def example_reading_spec(self):
     data_fields[label_key] = tf.FixedLenFeature((1,), tf.int64)
 
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
 
   def hparams(self, defaults, unused_model_hparams):
@@ -344,7 +343,7 @@ def example_reading_spec(self):
         super(Image2TextProblem, self).example_reading_spec())
     data_fields[label_key] = tf.VarLenFeature(tf.int64)
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
 
   def feature_encoders(self, data_dir):
diff --git a/tensor2tensor/data_generators/moving_mnist.py b/tensor2tensor/data_generators/moving_mnist.py
index d8841c741..507207623 100644
--- a/tensor2tensor/data_generators/moving_mnist.py
+++ b/tensor2tensor/data_generators/moving_mnist.py
@@ -37,7 +37,6 @@
 import tensorflow as tf
 import tensorflow_datasets as tfds
 from tensorflow_datasets.video import moving_sequence
-import tf_slim as slim
 
 
 DATA_URL = (
@@ -95,7 +94,7 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": slim.tfexample_decoder.Tensor(
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
     }
     return data_fields, decoders
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index c1ff4ef4f..aca91f31f 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -32,8 +32,13 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
-import tf_slim as slim
-from tensorflow.contrib.tpu.python.tpu import tpu_config
+# pylint: disable=g-import-not-at-top
+try:
+  from tensorflow.contrib.tpu.python.tpu import tpu_config
+except ImportError:
+  # TF 2.0 doesn't ship with contrib.
+  tpu_config = None
+# pylint: enable=g-import-not-at-top
 
 
 
@@ -199,7 +204,7 @@ class Problem(object):
         - Mutate defaults as needed
     * example_reading_spec
         - Specify the names and types of the features on disk.
-        - Specify slim.tfexample_decoder
+        - Specify tf.contrib.slim.tfexample_decoder
     * preprocess_example(example, mode, hparams)
         - Preprocess the example feature dict from feature name to Tensor or
           SparseTensor.
@@ -643,7 +648,7 @@ def dataset(self,
 
     data_filepattern = self.filepattern(data_dir, dataset_split, shard=shard)
     tf.logging.info("Reading data files from %s", data_filepattern)
-    data_files = sorted(slim.parallel_reader.get_data_files(
+    data_files = sorted(tf.contrib.slim.parallel_reader.get_data_files(
         data_filepattern))
 
     # Functions used in dataset transforms below. `filenames` can be either a
@@ -706,12 +711,12 @@ def decode_example(self, serialized_example):
     data_fields["batch_prediction_key"] = tf.FixedLenFeature([1], tf.int64, 0)
     if data_items_to_decoders is None:
       data_items_to_decoders = {
-          field: slim.tfexample_decoder.Tensor(field)
+          field: tf.contrib.slim.tfexample_decoder.Tensor(field)
           for field in data_fields
       }
 
-    decoder = slim.tfexample_decoder.TFExampleDecoder(data_fields,
-                                                      data_items_to_decoders)
+    decoder = tf.contrib.slim.tfexample_decoder.TFExampleDecoder(
+        data_fields, data_items_to_decoders)
 
     decode_items = list(sorted(data_items_to_decoders))
     decoded = decoder.decode(serialized_example, items=decode_items)
diff --git a/tensor2tensor/data_generators/style_transfer.py b/tensor2tensor/data_generators/style_transfer.py
index cdd9ceacc..146793666 100644
--- a/tensor2tensor/data_generators/style_transfer.py
+++ b/tensor2tensor/data_generators/style_transfer.py
@@ -32,16 +32,12 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
 
-logger = tf.logging
+# Modern-Shakespeare corpus is consisted of:
+# - 18,395 parallel sentences for training (train set),
+# - 1,218 parallel sentences for evaluation (dev set),
+# - 1,462 parallel sentence for testing (test set).
 
-"""
-Modern-Shakespeare corpus is consisted of:
-- 18,395 parallel sentences for training (train set),
-- 1,218 parallel sentences for evaluation (dev set),
-- 1,462 parallel sentence for testing (test set).
-"""
 
 _SHAKESPEARE_MODERN_TRAIN_DATASET = [[
     "https://github.com/tlatkowski/st/raw/master/shakespeare.train.tgz",
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 8ac2f2575..05e07d7c3 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -32,7 +32,6 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 class TranslateProblem(text_problems.Text2TextProblem):
@@ -277,8 +276,8 @@ def example_reading_spec(self):
 
     # hack: ignoring true targets and putting dist_targets in targets
     data_items_to_decoders = {
-        "inputs": slim.tfexample_decoder.Tensor("inputs"),
-        "targets": slim.tfexample_decoder.Tensor("dist_targets"),
+        "inputs": tf.contrib.slim.tfexample_decoder.Tensor("inputs"),
+        "targets": tf.contrib.slim.tfexample_decoder.Tensor("dist_targets"),
     }
 
     return (data_fields, data_items_to_decoders)
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index e4f2e5661..aa0e5bcf3 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -28,7 +28,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 try:
   import matplotlib  # pylint: disable=g-import-not-at-top
@@ -86,7 +85,7 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": slim.tfexample_decoder.Tensor(
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
     }
     return data_fields, decoders
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 4d3cec63c..ae2d2b313 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -35,7 +35,6 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import video_metrics
 import tensorflow as tf
-import tf_slim as slim
 
 
 FLAGS = flags.FLAGS
@@ -385,7 +384,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "frame":
-            slim.tfexample_decoder.Image(
+            tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 shape=[self.frame_height, self.frame_width, self.num_channels],
@@ -677,7 +676,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "inputs":
-            slim.tfexample_decoder.Image(
+            tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
@@ -767,7 +766,7 @@ def example_reading_spec(self):
         super(Video2ClassProblem, self).example_reading_spec())
     data_fields[label_key] = tf.FixedLenFeature((1,), tf.int64)
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
 
   def hparams(self, defaults, unused_model_hparams):
diff --git a/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
index e0c75d388..8a39f85f6 100644
--- a/tensor2tensor/data_generators/vqa.py
+++ b/tensor2tensor/data_generators/vqa.py
@@ -39,7 +39,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 def _get_vqa_v2_annotations(directory,
@@ -218,10 +217,10 @@ def example_reading_spec(self):
         (), tf.int64, allow_missing=True)
 
     data_items_to_decoders[
-        "question"] = slim.tfexample_decoder.Tensor(
+        "question"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/question")
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/answer")
     return data_fields, data_items_to_decoders
 
@@ -339,23 +338,23 @@ def example_reading_spec(self):
         (), tf.int64, allow_missing=True)
 
     data_items_to_decoders[
-        "inputs"] = slim.tfexample_decoder.Tensor(
+        "inputs"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/feature")
     data_items_to_decoders[
-        "question_id"] = slim.tfexample_decoder.Tensor(
+        "question_id"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/question_id")
     data_items_to_decoders[
-        "image_id"] = slim.tfexample_decoder.Tensor(
+        "image_id"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/image_id")
 
     data_items_to_decoders[
-        "spatial_feature"] = slim.tfexample_decoder.Tensor(
+        "spatial_feature"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/spatial_feature")
     data_items_to_decoders[
-        "question"] = slim.tfexample_decoder.Tensor(
+        "question"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/question")
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/answer")
 
     return data_fields, data_items_to_decoders
diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 7c02f43fd..43afb1a75 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -33,7 +33,6 @@
 from tensor2tensor.envs import trajectory
 from tensor2tensor.layers import modalities
 import tensorflow as tf
-import tf_slim as slim
 
 # Names for data fields in stored tf.Examples.
 TIMESTEP_FIELD = "timestep"
@@ -477,7 +476,7 @@ def example_reading_spec(self):
     }
 
     data_items_to_decoders = {
-        field: slim.tfexample_decoder.Tensor(field)
+        field: tf.contrib.slim.tfexample_decoder.Tensor(field)
         for field in data_fields
     }
 
diff --git a/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
index 82eb2fa71..92a7a7e54 100644
--- a/tensor2tensor/envs/rendered_env_problem.py
+++ b/tensor2tensor/envs/rendered_env_problem.py
@@ -25,7 +25,6 @@
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import gym_env_problem
 import tensorflow as tf
-import tf_slim as slim
 
 _IMAGE_ENCODED_FIELD = "image/encoded"
 _IMAGE_FORMAT_FIELD = "image/format"
@@ -81,7 +80,7 @@ def example_reading_spec(self):
     # Add frame number spec and decoder.
     env_fields[_FRAME_NUMBER_FIELD] = tf.FixedLenFeature((1,), tf.int64)
     env_decoders[
-        _FRAME_NUMBER_FIELD] = slim.tfexample_decoder.Tensor(
+        _FRAME_NUMBER_FIELD] = tf.contrib.slim.tfexample_decoder.Tensor(
             _FRAME_NUMBER_FIELD)
 
     # Add video fields and decoders
diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index cf4705a1c..e980db4d5 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -31,7 +31,7 @@
 import cloudpickle
 import numpy as np
 from tensor2tensor.envs import time_step
-from tensorflow.io import gfile
+import tensorflow as tf
 
 TRAJECTORY_FILE_FORMAT = r"trajectory_epoch_{epoch}_env_id_{env_id}_temperature_{temperature}_r_{r}.pkl"
 
@@ -528,7 +528,7 @@ def load_from_directory(trajectory_dir,
         r="*",
     )
 
-    trajectory_files = gfile.glob(
+    trajectory_files = tf.io.gfile.glob(
         os.path.join(trajectory_dir, trajectory_file_glob))
 
     if n_trajectories:
@@ -543,7 +543,7 @@ def load_from_directory(trajectory_dir,
         time.sleep(sleep_time_secs)
         max_tries -= 1
         sleep_time_secs = min(10.0, sleep_time_secs * 2)
-        trajectory_files = gfile.glob(
+        trajectory_files = tf.io.gfile.glob(
             os.path.join(trajectory_dir, trajectory_file_glob))
 
       # We can't get the required number of files and we can't up-sample either.
@@ -557,7 +557,7 @@ def load_from_directory(trajectory_dir,
     # We read and load all the files, revisit if this becomes a problem.
     trajectories_buffer = []
     for trajectory_file in trajectory_files:
-      with gfile.GFile(trajectory_file, "rb") as f:
+      with tf.io.gfile.GFile(trajectory_file, "rb") as f:
         trajectory = get_pickle_module().load(f)
         assert isinstance(trajectory, Trajectory)
         trajectories_buffer.append(trajectory)
diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index f3c43443e..a1589ed5d 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -951,7 +951,8 @@ def attention_bias_ignore_padding(memory_padding):
 
 
 @expert_utils.add_name_scope()
-def attention_bias_to_padding(attention_bias, cast_fn=tf.to_float):
+def attention_bias_to_padding(attention_bias,
+                              cast_fn=(lambda x: tf.cast(x, tf.float32))):
   """Inverse of attention_bias_ignore_padding().
 
   Args:
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 8eb6a8d19..aae4a6b4c 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -23,6 +23,7 @@
 import functools
 import math
 
+from absl import logging
 import numpy as np
 from six.moves import range  # pylint: disable=redefined-builtin
 
@@ -44,11 +45,15 @@ def layers():
   global _cached_layers
   if _cached_layers is not None:
     return _cached_layers
-  layers_module = tf.layers
+  layers_module = None
+  try:
+    layers_module = tf.layers
+  except AttributeError:
+    logging.info("Cannot access tf.layers, trying TF2 layers.")
   try:
     from tensorflow.python import tf2  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
     if tf2.enabled():
-      tf.logging.info("Running in V2 mode, using Keras layers.")
+      logging.info("Running in V2 mode, using Keras layers.")
       layers_module = tf.keras.layers
   except ImportError:
     pass
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 0d2a1af4c..d564db7fe 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -33,9 +33,12 @@
 except ImportError:
   distribute_summary_op_util = summary_op_util
 
-
-tfl = tf.layers
-tfcl = tf.contrib.layers
+tfl = common_layers.layers()
+tfcl = None
+try:
+  tfcl = tf.contrib.layers
+except AttributeError:
+  pass
 
 
 def swap_time_and_batch_axes(inputs):
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index e80aa6dc4..c65878163 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -28,7 +28,6 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 def cast_ints_to_int32(features):
@@ -380,7 +379,7 @@ def define_shapes(example):
     dataset = dataset.repeat()
 
   if is_training and skip_random_fraction_when_training:
-    data_files = slim.parallel_reader.get_data_files(filepattern)
+    data_files = tf.contrib.slim.parallel_reader.get_data_files(filepattern)
     #  In continuous_train_and_eval when switching between train and
     #  eval, this input_fn method gets called multiple times and it
     #  would give you the exact same samples from the last call
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index fcdf1aab2..59a54133b 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -29,7 +29,6 @@
 
 import tensorflow as tf
 
-from tensorflow.contrib.eager.python import tfe
 from tensorflow.python.util import tf_inspect as inspect
 
 
@@ -786,6 +785,9 @@ def create_eager_metrics_internal(metric_fns,
     (accum_fn(predictions, targets) => None,
      result_fn() => dict<str metric_name, float avg_val>
   """
+
+  from tensorflow.contrib.eager.python import tfe  # pylint: disable=g-import-not-at-top
+
   tfe_metrics = {}
 
   for name in metric_fns:
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 5a108f53c..36ffe363b 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -36,7 +36,6 @@
 
 import tensorflow as tf
 
-from tensorflow.contrib.tpu.python.tpu import tpu_estimator
 from tensorflow.core.protobuf import rewriter_config_pb2
 from tensorflow.python import debug
 
@@ -297,6 +296,7 @@ def create_estimator(model_name,
 
   del use_xla
   if use_tpu or use_tpu_estimator:
+    from tensorflow.contrib.tpu.python.tpu import tpu_estimator  # pylint: disable=g-import-not-at-top
     problem = hparams.problem
     batch_size = (
         problem.tpu_batch_size_per_shard(hparams) *

From ae3d13b3ec8c51ff47da144dcc3e2b01bc1f1d6d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 25 Nov 2019 13:01:13 -0800
Subject: [PATCH 2587/2720] Add a neural_assistant_base hparam config

PiperOrigin-RevId: 282416354
---
 tensor2tensor/models/neural_assistant.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/tensor2tensor/models/neural_assistant.py b/tensor2tensor/models/neural_assistant.py
index 454f467c6..3c619f47f 100644
--- a/tensor2tensor/models/neural_assistant.py
+++ b/tensor2tensor/models/neural_assistant.py
@@ -425,6 +425,25 @@ def compute_summary_embedding(input_embeddings, input_lengths, hparams):
     return compute_last_embedding(input_embeddings, input_lengths, hparams)
 
 
+@registry.register_hparams
+def neural_assistant_base():
+  """HParams for a base neural_assistant model."""
+  hparams = transformer.transformer_base()
+  hparams.add_hparam("pos_weight", 1.0)  # weight for positive triples
+  hparams.add_hparam("similarity_fuction",
+                     "bilinear")  # dot_product or bilinear
+  hparams.add_hparam("pool_technique", "average")  # avg or max pool or last
+  hparams.add_hparam("last_k", 1)  # number of last indices for averaging
+  hparams.add_hparam("max_triple_length", 30)  # max length of every triple
+  hparams.add_hparam("train_triple_num",
+                     5000)  # max number of triples during training
+  hparams.add_hparam("attend_kb", True)  # if False, it's a transformer model
+  hparams.add_hparam("kb_loss_weight", 0.0)  # weight for distant supervision
+  hparams.add_hparam("test_triple_num",
+                     28483)  # max triples of KB
+  return hparams
+
+
 @registry.register_hparams
 def neural_assistant_tiny():
   """HParams for tiny neural_assistant model."""

From 66e7ba8e96eb43707af603984b5db55aea6432f4 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 25 Nov 2019 16:08:28 -0800
Subject: [PATCH 2588/2720] T2T data reader was padding knowledge features to
 max_length unnecessarily because their shape is unknown. Fixed it by
 explicitly setting shape in preprocessing and removed the unnecessary slices
 in T2T model code.

PiperOrigin-RevId: 282453938
---
 tensor2tensor/models/neural_assistant.py | 27 +++++++++---------------
 1 file changed, 10 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/models/neural_assistant.py b/tensor2tensor/models/neural_assistant.py
index 3c619f47f..77695c0fd 100644
--- a/tensor2tensor/models/neural_assistant.py
+++ b/tensor2tensor/models/neural_assistant.py
@@ -18,7 +18,6 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
 import six
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
@@ -83,27 +82,22 @@ def model_fn(self, features):
 
   def encode_knowledge_bottom(self, features):
     tf.logging.info("Encoding knowledge " + str(self.triple_num))
-    hparams = self._hparams
-
     # Make sure this is embeddings for triples
-    # <tf.float32>[batch_size, max_triple_num*max_triple_length, 1, emb_dim]
+    # <tf.float32>[batch_size, triple_num*max_triple_length, 1, emb_dim]
     fact_embedding = features["encoded_triples"]
-    # [batch_size, max_triple_num*max_triple_length, emb_dim]
+    # [batch_size, triple_num*max_triple_length, emb_dim]
     fact_embedding = tf.squeeze(fact_embedding, 2)
 
     kb_shape = common_layers.shape_list(fact_embedding)
     batch_size = kb_shape[0]
     embed_dim = kb_shape[2]
-    max_triple_length = hparams.max_triple_length
-    fact_embedding = fact_embedding[:, :self.triple_num * max_triple_length, :]
-    # <tf.float32>[batch_size*max_triple_num, max_triple_length, emb_dim]
+    # <tf.float32>[batch_size*triple_num, max_triple_length, emb_dim]
     re_fact_embedding = tf.reshape(
         fact_embedding, [batch_size * self.triple_num, -1, embed_dim],
         name="reshape_fact_embedding")
 
-    # <tf.int64>[batch_size, max_triple_num]
+    # <tf.int64>[batch_size, triple_num]
     input_fact_lengths = features["triple_lens"]
-    input_fact_lengths = input_fact_lengths[:, :self.triple_num]
     # Stack the fact lengths.
     # <tf.int64>[batch_size*max_triple_num]
     re_fact_lengths = tf.reshape(
@@ -119,9 +113,9 @@ def compute_knowledge_selection_and_loss(self, features, encoder_output,
     Args:
       features: features.
       encoder_output: <tf.float32>[batch_size, input_length, hidden_dim]
-      fact_embedding: <tf.float32>[batch_size*max_triple_num, max_triple_length,
+      fact_embedding: <tf.float32>[batch_size*triple_num, max_triple_length,
         emb_dim]
-      fact_lengths: # <tf.int32>[batch_size*max_triple_num]
+      fact_lengths: # <tf.int32>[batch_size*triple_num]
 
     Returns:
       knowledge_weights:
@@ -156,16 +150,16 @@ def compute_knowledge_selection_and_loss(self, features, encoder_output,
                     tf.expand_dims(context_vector_summary, 2)), -1)
     elif hparams.similarity_fuction == "bilinear":
       # Tile the context vector summary.
-      # <tf.float32>[batch_size, max_triple_num*hidden_dim]
+      # <tf.float32>[batch_size, triple_num*hidden_dim]
       tiled_context_vector = tf.tile(context_vector_summary,
                                      [1, self.triple_num])
-      # <tf.float32>[batch_size, max_triple_num, hidden_dim]
+      # <tf.float32>[batch_size, triple_num, hidden_dim]
       context_vector = tf.reshape(tiled_context_vector,
                                   [-1, self.triple_num, encoder_hidden_dim])
       # compute outer product
       context_vector = tf.expand_dims(context_vector, -1)
       knowledge_encoder_output = tf.expand_dims(knowledge_encoder_output, 2)
-      # <tf.float32>[batch_size, max_triple_num, hidden_dim, hidden_dim]
+      # <tf.float32>[batch_size, triple_num, hidden_dim, hidden_dim]
       outer_product = tf.matmul(context_vector, knowledge_encoder_output)
       outer_product = tf.reshape(
           outer_product,
@@ -175,7 +169,6 @@ def compute_knowledge_selection_and_loss(self, features, encoder_output,
 
     avg_triple_loss = 0.0
     triple_labels = features["triple_labels"]
-    triple_labels = triple_labels[:, :self.triple_num]
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
       triple_losses = tf.nn.weighted_cross_entropy_with_logits(
           labels=triple_labels,
@@ -213,7 +206,7 @@ def body(self, features):
     with tf.variable_scope("knowledge"):
       with tf.name_scope("knowledge_encoding"):
         # Encode knowledge.
-        # <tf.float32>[batch_size, max_triple_num, emb_dim]
+        # <tf.float32>[batch_size, triple_num, emb_dim]
         fact_embedding, fact_lengths = self.encode_knowledge_bottom(features)
         tf.logging.info("Encoded knowledge")
 

From aa23c240b6bbb351bf1c4f4bd0d634e10c4e4c29 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 25 Nov 2019 19:00:35 -0800
Subject: [PATCH 2589/2720] small audio_encoder.py refactors

PiperOrigin-RevId: 282478845
---
 tensor2tensor/data_generators/audio_encoder.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/audio_encoder.py b/tensor2tensor/data_generators/audio_encoder.py
index 26fc01889..68d942299 100644
--- a/tensor2tensor/data_generators/audio_encoder.py
+++ b/tensor2tensor/data_generators/audio_encoder.py
@@ -42,20 +42,24 @@ def encode(self, s):
     Returns:
       samples: list of int16s
     """
+    def convert_to_wav(in_path, out_path, extra_args=None):
+      if not os.path.exists(out_path):
+        # TODO(dliebling) On Linux, check if libsox-fmt-mp3 is installed.
+        args = ["sox", "--rate", "16k", "--bits", "16", "--channel", "1"]
+        if extra_args:
+          args += extra_args
+        call(args + [in_path, out_path])
+
     # Make sure that the data is a single channel, 16bit, 16kHz wave.
     # TODO(chorowski): the directory may not be writable, this should fallback
     # to a temp path, and provide instructions for installing sox.
     if s.endswith(".mp3"):
-      # TODO(dliebling) On Linux, check if libsox-fmt-mp3 is installed.
       out_filepath = s[:-4] + ".wav"
-      call([
-          "sox", "--guard", s, "-r", "16k", "-b", "16", "-c", "1", out_filepath
-      ])
+      convert_to_wav(s, out_filepath, ["--guard"])
       s = out_filepath
     elif not s.endswith(".wav"):
       out_filepath = s + ".wav"
-      if not os.path.exists(out_filepath):
-        call(["sox", "-r", "16k", "-b", "16", "-c", "1", s, out_filepath])
+      convert_to_wav(s, out_filepath)
       s = out_filepath
     rate, data = wavfile.read(s)
     assert rate == self._sample_rate

From 838f1a99e24a9391a8faf6603e90d476444110a0 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 5 Dec 2019 07:30:12 -0800
Subject: [PATCH 2590/2720] Remove unused 'using' declarations.

PiperOrigin-RevId: 283969379
---
 tensor2tensor/data_generators/ops/subword_text_encoder.cc   | 1 -
 .../data_generators/ops/subword_text_encoder_test.cc        | 6 ------
 2 files changed, 7 deletions(-)

diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder.cc b/tensor2tensor/data_generators/ops/subword_text_encoder.cc
index 9199e2a83..bc3d22134 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder.cc
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder.cc
@@ -12,7 +12,6 @@ namespace tensor2tensor {
 namespace {
 
 using ::tensorflow::Env;
-using ::tensorflow::Tensor;
 
 // End of Sequence token ID to insert at end of encoded text.
 constexpr int64 kEosTokenId = 1;
diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder_test.cc b/tensor2tensor/data_generators/ops/subword_text_encoder_test.cc
index baef07c5f..9bad1523d 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder_test.cc
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder_test.cc
@@ -7,12 +7,6 @@
 namespace tensor2tensor {
 namespace {
 
-using ::tensorflow::DT_INT64;
-using ::tensorflow::Tensor;
-using ::tensorflow::TensorShape;
-using ::tensorflow::test::AsTensor;
-using ::tensorflow::test::ExpectTensorEqual;
-
 TEST(SubwordTextEncoderTest, EncodesSubTokens) {
   SubwordTextEncoder encoder("third_party/py/tensor2tensor/"
                              "data_generators/ops/testdata/subwords");

From d8433981b2dfcc3858ba7b8ccc8a37028991d04a Mon Sep 17 00:00:00 2001
From: James Thomin <james.thomin@gmail.com>
Date: Thu, 12 Dec 2019 02:20:18 -0600
Subject: [PATCH 2591/2720] Add support for training on V100 GPUs on Cloud ML
 (#1766)

Updated the runtime version to 1.14 and added support for
using V100 GPUs on Cloud ML.
---
 tensor2tensor/utils/cloud_mlengine.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 407a963cf..c8cf08670 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -35,7 +35,7 @@
 FLAGS = tf.flags.FLAGS
 
 CONSOLE_URL = "https://console.cloud.google.com/mlengine/jobs/"
-RUNTIME_VERSION = "1.13"
+RUNTIME_VERSION = "1.14"
 LIST_VM = "gcloud compute instances list"
 DEFAULT_PROJECT = "gcloud config get-value project"
 DEFAULT_REGION = "gcloud config get-value compute/region"
@@ -310,12 +310,15 @@ def validate_flags():
     if FLAGS.worker_gpu:
       if FLAGS.worker_gpu == 1:
         assert FLAGS.cloud_mlengine_master_type in ["standard_gpu",
-                                                    "standard_p100"]
+                                                    "standard_p100",
+                                                    "standard_v100"]
       elif FLAGS.worker_gpu == 4:
         assert FLAGS.cloud_mlengine_master_type in ["complex_model_m_gpu",
-                                                    "complex_model_m_p100"]
+                                                    "complex_model_m_p100",
+                                                    "complex_model_m_v100"]
       else:
-        assert FLAGS.cloud_mlengine_master_type == "complex_model_l_gpu"
+        assert FLAGS.cloud_mlengine_master_type in ["complex_model_l_gpu",
+                                                    "complex_model_l_v100"]
     else:
       assert FLAGS.cloud_mlengine_master_type in ["standard", "large_model",
                                                   "complex_model_s",

From fb067f82c89945b9a5ff73c318774b71c363131f Mon Sep 17 00:00:00 2001
From: Hamza Abbad <hamza.abbad@gmail.com>
Date: Thu, 12 Dec 2019 16:20:31 +0800
Subject: [PATCH 2592/2720] Fixed a mistake in a doc comment (#1760)

In T2TModel.infer, the comment mentioned that the returned outputs tensor has a shape of `[batch_size, <= decode_length] if beam_size == 1`, but this is wrong. It should be: `if top_beams == 1`. It is a simple mistake but it may really confuse a beginner (like me).
---
 tensor2tensor/utils/t2t_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 2b9e419ef..599381888 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -783,7 +783,7 @@ def infer(self,
     Returns:
       A dict of decoding results {
           "outputs": integer `Tensor` of decoded ids of shape
-              [batch_size, <= decode_length] if beam_size == 1 or
+              [batch_size, <= decode_length] if top_beams == 1 or
               [batch_size, top_beams, <= decode_length]
           "scores": decoding log probs from the beam search,
               None if using greedy decoding (beam_size=1)

From e3e54900d377940c2930285104d6b05bf7cf0330 Mon Sep 17 00:00:00 2001
From: geek_fly <real.szf@gmail.com>
Date: Thu, 12 Dec 2019 16:20:46 +0800
Subject: [PATCH 2593/2720] inaccurate calculation of first token in pos_emb
 (#1758)

since,
1. scaled_time is calculated by `scaled_time = tf.expand_dims(position, 1) * tf.expand_dims(inv_timescales, 0)`, while position_idx in position will **start by 0**.
2. scaled_time will be used to calculate position embedding by `signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)`

Thus, position_embedding of first token will always be combined with half zeros and half ones.

But the purpose of addding timing pos_embed should be helping model to learn relative and absolute position information between words by triangle relations of these words, i.e. sin(\alpha+\beta)=sin(\alpha)cos(\beta)+sin(\beta)cos(\alpha).
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index a1589ed5d..a2a062fce 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -439,7 +439,7 @@ def get_timing_signal_1d(length,
   Returns:
     a Tensor of timing signals [1, length, channels]
   """
-  position = tf.to_float(tf.range(length) + start_index)
+  position = tf.to_float(tf.range(1, 1+length) + start_index)
   num_timescales = channels // 2
   log_timescale_increment = (
       math.log(float(max_timescale) / float(min_timescale)) /

From d8338effc13b27f9747dc180359d7f983f9c2526 Mon Sep 17 00:00:00 2001
From: Zhanghao Wu <zhanghao.wu@outlook.com>
Date: Thu, 12 Dec 2019 03:20:57 -0500
Subject: [PATCH 2594/2720] Fix the decode length for decoding (#1752)

---
 tensor2tensor/utils/decoding.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index a63e83d5b..b8e77fc01 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -453,7 +453,7 @@ def input_fn():
           task_id=decode_hp.multiproblem_task_id, has_input=has_input)
       gen_fn = make_input_fn_from_generator(input_gen)
       example = gen_fn()
-      return _decode_input_tensor_to_features_dict(example, hparams)
+      return _decode_input_tensor_to_features_dict(example, hparams, decode_hp)
   decodes = []
   result_iter = estimator.predict(input_fn, checkpoint_path=checkpoint_path)
 
@@ -937,12 +937,13 @@ def _interactive_input_tensor_to_features_dict(feature_map, hparams):
   return features
 
 
-def _decode_input_tensor_to_features_dict(feature_map, hparams):
+def _decode_input_tensor_to_features_dict(feature_map, hparams, decode_hp):
   """Convert the interactive input format (see above) to a dictionary.
 
   Args:
     feature_map: dict with inputs.
     hparams: model hyperparameters
+    decode_hp: decode hyperparameters
 
   Returns:
     a features dictionary, as expected by the decoder.
@@ -962,7 +963,7 @@ def _decode_input_tensor_to_features_dict(feature_map, hparams):
   features["input_space_id"] = input_space_id
   features["target_space_id"] = target_space_id
   features["decode_length"] = (
-      IMAGE_DECODE_LENGTH if input_is_image else tf.shape(x)[1] + 50)
+      IMAGE_DECODE_LENGTH if input_is_image else tf.constant(decode_hp.extra_length))
   features["inputs"] = x
   # Save inputs to "partial_targets" when prepending inputs to targets. Also
   # keep "inputs" as some models crash if they don't exist.

From c8e5e9c0335d903957833fb8b41498b7562948b8 Mon Sep 17 00:00:00 2001
From: Zhanghao Wu <zhanghao.wu@outlook.com>
Date: Thu, 12 Dec 2019 00:21:32 -0800
Subject: [PATCH 2595/2720] Merge of PR #1752

PiperOrigin-RevId: 285138096
---
 tensor2tensor/layers/common_attention.py |  2 +-
 tensor2tensor/utils/cloud_mlengine.py    | 11 ++++-------
 tensor2tensor/utils/decoding.py          |  3 ++-
 tensor2tensor/utils/t2t_model.py         |  2 +-
 4 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index a2a062fce..a1589ed5d 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -439,7 +439,7 @@ def get_timing_signal_1d(length,
   Returns:
     a Tensor of timing signals [1, length, channels]
   """
-  position = tf.to_float(tf.range(1, 1+length) + start_index)
+  position = tf.to_float(tf.range(length) + start_index)
   num_timescales = channels // 2
   log_timescale_increment = (
       math.log(float(max_timescale) / float(min_timescale)) /
diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index c8cf08670..407a963cf 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -35,7 +35,7 @@
 FLAGS = tf.flags.FLAGS
 
 CONSOLE_URL = "https://console.cloud.google.com/mlengine/jobs/"
-RUNTIME_VERSION = "1.14"
+RUNTIME_VERSION = "1.13"
 LIST_VM = "gcloud compute instances list"
 DEFAULT_PROJECT = "gcloud config get-value project"
 DEFAULT_REGION = "gcloud config get-value compute/region"
@@ -310,15 +310,12 @@ def validate_flags():
     if FLAGS.worker_gpu:
       if FLAGS.worker_gpu == 1:
         assert FLAGS.cloud_mlengine_master_type in ["standard_gpu",
-                                                    "standard_p100",
-                                                    "standard_v100"]
+                                                    "standard_p100"]
       elif FLAGS.worker_gpu == 4:
         assert FLAGS.cloud_mlengine_master_type in ["complex_model_m_gpu",
-                                                    "complex_model_m_p100",
-                                                    "complex_model_m_v100"]
+                                                    "complex_model_m_p100"]
       else:
-        assert FLAGS.cloud_mlengine_master_type in ["complex_model_l_gpu",
-                                                    "complex_model_l_v100"]
+        assert FLAGS.cloud_mlengine_master_type == "complex_model_l_gpu"
     else:
       assert FLAGS.cloud_mlengine_master_type in ["standard", "large_model",
                                                   "complex_model_s",
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index b8e77fc01..d3ac9b2fe 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -963,7 +963,8 @@ def _decode_input_tensor_to_features_dict(feature_map, hparams, decode_hp):
   features["input_space_id"] = input_space_id
   features["target_space_id"] = target_space_id
   features["decode_length"] = (
-      IMAGE_DECODE_LENGTH if input_is_image else tf.constant(decode_hp.extra_length))
+      IMAGE_DECODE_LENGTH if input_is_image else
+      tf.constant(decode_hp.extra_length))
   features["inputs"] = x
   # Save inputs to "partial_targets" when prepending inputs to targets. Also
   # keep "inputs" as some models crash if they don't exist.
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 599381888..2b9e419ef 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -783,7 +783,7 @@ def infer(self,
     Returns:
       A dict of decoding results {
           "outputs": integer `Tensor` of decoded ids of shape
-              [batch_size, <= decode_length] if top_beams == 1 or
+              [batch_size, <= decode_length] if beam_size == 1 or
               [batch_size, top_beams, <= decode_length]
           "scores": decoding log probs from the beam search,
               None if using greedy decoding (beam_size=1)

From c4f1a89086de4678aa4b4a92b5cc17ccff602193 Mon Sep 17 00:00:00 2001
From: geek_fly <real.szf@gmail.com>
Date: Thu, 12 Dec 2019 00:21:50 -0800
Subject: [PATCH 2596/2720] Merge of PR #1758

PiperOrigin-RevId: 285138123
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index a1589ed5d..a2a062fce 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -439,7 +439,7 @@ def get_timing_signal_1d(length,
   Returns:
     a Tensor of timing signals [1, length, channels]
   """
-  position = tf.to_float(tf.range(length) + start_index)
+  position = tf.to_float(tf.range(1, 1+length) + start_index)
   num_timescales = channels // 2
   log_timescale_increment = (
       math.log(float(max_timescale) / float(min_timescale)) /

From c6ec6fcf7db52811f7e0a47e72f41ec24bd3f032 Mon Sep 17 00:00:00 2001
From: Hamza Abbad <hamza.abbad@gmail.com>
Date: Thu, 12 Dec 2019 00:21:59 -0800
Subject: [PATCH 2597/2720] Merge of PR #1760

PiperOrigin-RevId: 285138130
---
 tensor2tensor/utils/t2t_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 2b9e419ef..599381888 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -783,7 +783,7 @@ def infer(self,
     Returns:
       A dict of decoding results {
           "outputs": integer `Tensor` of decoded ids of shape
-              [batch_size, <= decode_length] if beam_size == 1 or
+              [batch_size, <= decode_length] if top_beams == 1 or
               [batch_size, top_beams, <= decode_length]
           "scores": decoding log probs from the beam search,
               None if using greedy decoding (beam_size=1)

From 7878b54b1e7b14594a921944f23efa26aefedfd5 Mon Sep 17 00:00:00 2001
From: James Thomin <james.thomin@gmail.com>
Date: Thu, 12 Dec 2019 00:22:09 -0800
Subject: [PATCH 2598/2720] Merge of PR #1766

PiperOrigin-RevId: 285138142
---
 tensor2tensor/utils/cloud_mlengine.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 407a963cf..c8cf08670 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -35,7 +35,7 @@
 FLAGS = tf.flags.FLAGS
 
 CONSOLE_URL = "https://console.cloud.google.com/mlengine/jobs/"
-RUNTIME_VERSION = "1.13"
+RUNTIME_VERSION = "1.14"
 LIST_VM = "gcloud compute instances list"
 DEFAULT_PROJECT = "gcloud config get-value project"
 DEFAULT_REGION = "gcloud config get-value compute/region"
@@ -310,12 +310,15 @@ def validate_flags():
     if FLAGS.worker_gpu:
       if FLAGS.worker_gpu == 1:
         assert FLAGS.cloud_mlengine_master_type in ["standard_gpu",
-                                                    "standard_p100"]
+                                                    "standard_p100",
+                                                    "standard_v100"]
       elif FLAGS.worker_gpu == 4:
         assert FLAGS.cloud_mlengine_master_type in ["complex_model_m_gpu",
-                                                    "complex_model_m_p100"]
+                                                    "complex_model_m_p100",
+                                                    "complex_model_m_v100"]
       else:
-        assert FLAGS.cloud_mlengine_master_type == "complex_model_l_gpu"
+        assert FLAGS.cloud_mlengine_master_type in ["complex_model_l_gpu",
+                                                    "complex_model_l_v100"]
     else:
       assert FLAGS.cloud_mlengine_master_type in ["standard", "large_model",
                                                   "complex_model_s",

From fc49d33e47475b6402a08039becaa7e39288a49c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 16 Dec 2019 14:34:51 -0800
Subject: [PATCH 2599/2720] Fix AudioEncoder comment

PiperOrigin-RevId: 285852472
---
 tensor2tensor/data_generators/audio_encoder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/audio_encoder.py b/tensor2tensor/data_generators/audio_encoder.py
index 68d942299..9ad4494eb 100644
--- a/tensor2tensor/data_generators/audio_encoder.py
+++ b/tensor2tensor/data_generators/audio_encoder.py
@@ -85,7 +85,7 @@ def decode(self, ids):
     return tmp_file_path
 
   def decode_list(self, ids):
-    """Transform a sequence of int ids into an image file.
+    """Transform a sequence of int ids into a wavform file.
 
     Args:
       ids: list of integers to be converted.

From a3031165a7b20e5b2698c4319d9abe51c2b94e2e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Em=C4=ABls=20Ozoli=C5=86=C5=A1?= <ozolinsemils@gmail.com>
Date: Tue, 17 Dec 2019 03:08:46 +0200
Subject: [PATCH 2600/2720] Neural Shuffle-Exchange Seq2Seq model (#1763)

* Neural Shuffle-Exchange model for Seq2Seq problems

* Format Shuffle-Exchange model code and add comments

* Fix Python 2.7 compatibility errors

* Fix lint errors in Shuffle-Exchange network

* Format docstrings for Neural Shuffle-Exchange network
---
 tensor2tensor/models/__init__.py              |  10 +-
 .../models/research/shuffle_network.py        | 413 ++++++++++++++++++
 2 files changed, 417 insertions(+), 6 deletions(-)
 create mode 100644 tensor2tensor/models/research/shuffle_network.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index d07545379..5f949dee5 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -20,8 +20,6 @@
 
 import six
 
-# pylint: disable=unused-import
-
 from tensor2tensor.layers import modalities  # pylint: disable=g-import-not-at-top
 from tensor2tensor.models import basic
 from tensor2tensor.models import bytenet
@@ -44,9 +42,7 @@
 from tensor2tensor.models import transformer
 from tensor2tensor.models import vanilla_gan
 from tensor2tensor.models import xception
-
 from tensor2tensor.models.neural_architecture_search import nas_model
-
 from tensor2tensor.models.research import adafactor_experiments
 from tensor2tensor.models.research import aligned
 from tensor2tensor.models.research import attention_lm
@@ -60,6 +56,7 @@
 from tensor2tensor.models.research import multiquery_paper
 from tensor2tensor.models.research import neural_stack
 from tensor2tensor.models.research import rl
+from tensor2tensor.models.research import shuffle_network
 from tensor2tensor.models.research import similarity_transformer
 from tensor2tensor.models.research import super_lm
 from tensor2tensor.models.research import transformer_moe
@@ -73,7 +70,6 @@
 from tensor2tensor.models.research import vqa_attention
 from tensor2tensor.models.research import vqa_recurrent_self_attention
 from tensor2tensor.models.research import vqa_self_attention
-
 from tensor2tensor.models.video import basic_deterministic
 from tensor2tensor.models.video import basic_recurrent
 from tensor2tensor.models.video import basic_stochastic
@@ -82,9 +78,11 @@
 from tensor2tensor.models.video import next_frame_glow
 from tensor2tensor.models.video import savp
 from tensor2tensor.models.video import sv2p
-
 from tensor2tensor.utils import registry
 
+
+# pylint: disable=unused-import
+
 # pylint: enable=unused-import
 
 
diff --git a/tensor2tensor/models/research/shuffle_network.py b/tensor2tensor/models/research/shuffle_network.py
new file mode 100644
index 000000000..a3307ce44
--- /dev/null
+++ b/tensor2tensor/models/research/shuffle_network.py
@@ -0,0 +1,413 @@
+"""Neural Shuffle-Exchange Network.
+
+Implementation of
+"Neural Shuffle-Exchange Networks - Sequence Processing in O(n log n) Time"
+paper by K.Freivalds, E.Ozolins, A.Sostaks.
+
+Paper: https://papers.nips.cc/paper/
+8889-neural-shuffle-exchange-networks-sequence-processing-in-on-log-n-time.pdf
+
+Original code: https://github.com/LUMII-Syslab/shuffle-exchange
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_hparams
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
+import tensorflow as tf
+
+
+def ror(x, n, p=1):
+  """Bitwise right rotation.
+
+  Args:
+    x: Input tensor
+    n: Bit count to represent x
+    p: Bit positions to shift
+
+  Returns:
+    tf.Tensor: x shifted by p positions in n bits
+  """
+
+  a = tf.bitwise.right_shift(x, p)
+  b = tf.bitwise.left_shift(1, p) - 1
+  c = tf.bitwise.bitwise_and(x, b)
+  d = tf.bitwise.left_shift(c, n - p)
+
+  return a + d
+
+
+def rol(x, n, p=1):
+  """Bitwise left rotation.
+
+  Args:
+    x: Input tensor
+    n: Bit count to represent x
+    p: Bit positions to shift
+
+  Returns:
+    tf.Tensor: x shifted by p positions in n bits
+  """
+  a = tf.bitwise.left_shift(x, p)
+  b = tf.bitwise.left_shift(1, n) - 1
+  c = tf.bitwise.bitwise_and(a, b)
+  d = tf.bitwise.right_shift(x, n - p)
+
+  return tf.bitwise.bitwise_or(c, d)
+
+
+def shuffle_layer(inputs, shuffle_fn=rol):
+  """Shuffles the elements according to bitwise left or right rotation.
+
+  Args:
+    inputs: Tensor input from previous layer
+    shuffle_fn: Shift function rol or ror
+
+  Returns:
+    tf.Tensor: Inputs shifted according to shuffle_fn
+  """
+
+  length = tf.shape(inputs)[1]
+  n_bits = tf.log(tf.cast(length - 1, tf.float32)) / tf.log(2.0)
+  n_bits = tf.cast(n_bits, tf.int32) + 1
+
+  indices = tf.range(0, length)
+  rev_indices = shuffle_fn(indices, n_bits)
+  return tf.gather(inputs, rev_indices, axis=1)
+
+
+def reverse_shuffle_layer(inputs):
+  """Reverse shuffle of inputs. Used in the second half of Benes block.
+
+  Args:
+    inputs: Inputs that should be shuffled
+
+  Returns:
+    tf.Tensor: Inputs shuffled according to bitwise right rotation
+  """
+
+  return shuffle_layer(inputs, ror)
+
+
+def conv_linear_map(inputs, nin, nout, bias_start, prefix):
+  """Convolutional liner map. Maps 3D tensor by last dimension.
+
+  Args:
+    inputs: Inputs that should be shuffled
+    nin: Input feature map count
+    nout: Output feature map count
+    bias_start: Bias start value
+    prefix: Name prefix
+
+  Returns:
+    tf.Tensor: Inputs with applied convolution
+  """
+
+  with tf.variable_scope(prefix):
+    inp_shape = tf.shape(inputs)
+
+    initializer = tf.variance_scaling_initializer(scale=1.0,
+                                                  mode="fan_avg",
+                                                  distribution="uniform")
+    kernel = tf.get_variable("CvK", [nin, nout], initializer=initializer)
+    bias_term = tf.get_variable("CvB", [nout],
+                                initializer=tf.constant_initializer(0.0))
+
+    mul_shape = [inp_shape[0] * inp_shape[1], nin]
+    res = tf.matmul(tf.reshape(inputs, mul_shape), kernel)
+    res = tf.reshape(res, [inp_shape[0], inp_shape[1], nout])
+    return res + bias_start + bias_term
+
+
+# pylint: disable=useless-object-inheritance
+class SwitchLayer(object):
+  """Switch layer of Neural Shuffle-Exchange network.
+  """
+
+  def __init__(self, prefix, dropout, mode):
+    """Initialize switch layer.
+
+    Args:
+      prefix: Name prefix for switch layer
+      dropout: Dropout rate
+      mode: Training mode
+    """
+
+    self.prefix = prefix
+    self.dropout = dropout
+    self.mode = mode
+    self.batch_size = None
+    self.length = None
+    self.num_units = None
+    self.n_bits = None
+
+  def linear_map(self, inputs, suffix, bias_start, in_units, out_units):
+    """2 input to 2 output linear map.
+
+    Args:
+      inputs: Input tensor
+      suffix: Linear map name suffix
+      bias_start: Bias start value
+      in_units: Size of input tensor feature map count
+      out_units: Size of output tensor feature map count
+
+    Return:
+      tf.Tensor: Convolution apply to input tensor
+    """
+    in_shape = [self.batch_size, self.length // 2, in_units * 2]
+    inputs = tf.reshape(inputs, in_shape)
+    res = conv_linear_map(inputs, in_units * 2, out_units * 2,
+                          bias_start, self.prefix + "/" + suffix)
+    return tf.reshape(res, [self.batch_size, self.length, out_units])
+
+  def gated_linear_map(self, inputs, suffix, bias_start_reset,
+                       in_units, out_units):
+    """Linear mapping with two reset gates.
+
+    Args:
+      inputs: Input tensor
+      suffix: Linear map name suffix
+      bias_start_reset: Bias start value for reset gate
+      in_units: Size of input tensor feature map count
+      out_units: Size of output tensor feature map count
+
+    Return:
+      tf.Tensor: Convolution apply to input tensor
+    """
+
+    def reset_gate(name):
+      prefix = self.prefix + name + suffix
+      reset = conv_linear_map(inputs, in_units * 2, in_units * 2,
+                              bias_start_reset, prefix)
+      return tf.nn.sigmoid(reset)
+
+    in_shape = [self.batch_size, self.length // 2, in_units * 2]
+    inputs = tf.reshape(inputs, in_shape)
+
+    reset1 = reset_gate("/reset1/")
+    reset2 = reset_gate("/reset2/")
+    res1 = conv_linear_map(inputs * reset1, in_units * 2,
+                           out_units, 0.0, self.prefix + "/cand1/" + suffix)
+    res2 = conv_linear_map(inputs * reset2, in_units * 2,
+                           out_units, 0.0, self.prefix + "/cand2/" + suffix)
+
+    res = tf.concat([res1, res2], axis=2)
+    res = tf.reshape(res, [self.batch_size, self.length, out_units])
+    return tf.nn.tanh(res)
+
+  def __call__(self, inputs, residual_inputs):
+    """Apply SwitchLayer to inputs.
+
+    Args:
+      inputs: Input tensor
+      residual_inputs: Residual connections from previous block
+
+    Returns:
+      tf.Tensor: New candidate value
+    """
+    input_shape = tf.shape(inputs)
+    self.batch_size = input_shape[0]
+    self.length = input_shape[1]
+    self.num_units = inputs.shape.as_list()[2]
+
+    self.n_bits = tf.log(tf.cast(self.length - 1, tf.float32)) / tf.log(2.0)
+    self.n_bits = tf.floor(self.n_bits) + 1
+
+    initializer = tf.constant_initializer(0.5)
+    residual_scale = tf.get_variable(self.prefix + "/residual_scale",
+                                     [self.num_units], initializer=initializer)
+
+    shuffled_input = self.swap_halves(inputs)
+    mem_all = inputs + residual_inputs * residual_scale
+
+    # calculate the new value
+    candidate = self.gated_linear_map(mem_all, "c", 0.5,
+                                      self.num_units, self.num_units)
+    gate = tf.nn.sigmoid(self.linear_map(mem_all, "g", 0.5,
+                                         self.num_units, self.num_units))
+
+    candidate = gate * shuffled_input + (1 - gate) * candidate
+
+    if self.dropout > 0:
+      candidate = tf.nn.dropout(candidate, rate=self.dropout / self.n_bits)
+    if self.dropout != 0.0 and self.mode == tf.estimator.ModeKeys.TRAIN:
+      noise = tf.random_normal(tf.shape(candidate), mean=1.0, stddev=0.001)
+      candidate = candidate * noise
+
+    return candidate
+
+  def swap_halves(self, inputs):
+    """Split inputs in half and then shuffle them as described in paper.
+
+    Args:
+      inputs: ShuffleLayer inputs
+
+    Return:
+      tf.Tensor: Inputs with swapped halves
+    """
+    x = tf.range(0, self.length)
+    xor_indices = tf.bitwise.bitwise_xor(x, 1)
+    input_xor = tf.gather(inputs[:, :, :self.num_units // 2],
+                          xor_indices, axis=1)
+    return tf.concat([input_xor, inputs[:, :, self.num_units // 2:]], axis=2)
+
+
+def shuffle_network(inputs, hparams):
+  """Neural Shuffle-Network with skip connections between blocks.
+
+  Args:
+    inputs: inputs to the Shuffle-Exchange network.
+    Should be in length of power of 2.
+    hparams: Model configuration
+
+  Returns:
+    tf.Tensor: Outputs of the Shuffle-Exchange last layer
+  """
+
+  def forward_step(state, layer_nr):
+    with tf.variable_scope("forward"):
+      last_state, residuals = state
+      prev = residuals[layer_nr, :, :, :]
+      switch = SwitchLayer("switch", hparams.dropout, hparams.mode)
+      cur = switch(last_state, prev)
+      return shuffle_layer(cur), residuals
+
+  def reverse_step(state, layer_nr):
+    with tf.variable_scope("reverse"):
+      last_state, residuals = state
+      prev = residuals[layer_nr, :, :, :]
+      switch = SwitchLayer("reverse_switch", hparams.dropout, hparams.mode)
+      cur = switch(last_state, prev)
+      return reverse_shuffle_layer(cur), residuals
+
+  input_shape = tf.shape(inputs)
+  n_bits = tf.log(tf.cast(input_shape[1] - 1, tf.float32)) / tf.log(2.0)
+  n_bits = tf.cast(n_bits, tf.int32) + 1
+
+  queue_shape = [n_bits * 2, input_shape[0], input_shape[1], input_shape[2]]
+  residuals_queue = tf.zeros(queue_shape)
+  block_out = tf.tanh(inputs)
+
+  for k in range(hparams.num_hidden_layers):
+    with tf.variable_scope("benes_block_" + str(k), reuse=tf.AUTO_REUSE):
+      forward_outputs, _ = tf.scan(forward_step,
+                                   tf.range(0, n_bits),
+                                   initializer=(block_out, residuals_queue),
+                                   parallel_iterations=1,
+                                   swap_memory=True)
+
+      forward_tensors = [tf.expand_dims(block_out, axis=0), forward_outputs]
+      forward_outputs = tf.concat(forward_tensors, axis=0)
+      forward_last = forward_outputs[-1, :, :, :]
+
+      reverse_outputs, _ = tf.scan(reverse_step,
+                                   tf.range(n_bits, n_bits * 2),
+                                   initializer=(forward_last, residuals_queue),
+                                   parallel_iterations=1,
+                                   swap_memory=True)
+
+      block_out = reverse_outputs[-1, :, :, :]
+      residuals_queue = tf.concat([forward_outputs, reverse_outputs], axis=0)
+
+  last_layer = SwitchLayer("last_layer", hparams.dropout, hparams.mode)
+  return last_layer(block_out, residuals_queue[n_bits * 2, :, :, :])
+
+
+@registry.register_model
+class ShuffleNetwork(t2t_model.T2TModel):
+  """Seq2Seq model for sequence processing in O(n log n) time.
+  """
+
+  def bottom(self, features):
+    """We add padding to the input and output so they are the same.
+    Length of input and output should be power of 2.
+
+    Args:
+      features: Dictionary of inputs and targets
+
+    Returns:
+      dictionary: Inputs and targets padded with 0 to the length of power of 2.
+      Both are same length.
+    """
+    inputs = features["inputs"]
+    targets = features["targets"]
+    inputs_length = tf.shape(inputs)[1]
+    targets_length = tf.shape(targets)[1]
+
+    length = tf.maximum(inputs_length, targets_length)
+    p = tf.log(tf.cast(length, tf.float32)) / tf.log(2.0)
+    p = tf.cast(tf.ceil(p), tf.int32)
+    pad_len = tf.pow(2, p)
+
+    input_padding = [[0, 0], [0, pad_len - inputs_length], [0, 0], [0, 0]]
+    features["inputs"] = tf.pad(inputs, input_padding)
+    target_padding = [[0, 0], [0, pad_len - targets_length], [0, 0], [0, 0]]
+    features["targets"] = tf.pad(targets, target_padding)
+    return super(ShuffleNetwork, self).bottom(features)
+
+  def loss(self, logits, features):
+    """Loss function for Neural Shuffle-Exchange network.
+
+    We use custom loss function as default loss function doesn't
+    use padding for calculating loss. We assume that output string is same
+    length as the input. If you need other type of output please feel
+    free to modify this.
+
+    Args:
+      logits: Logits from model
+      features: Features, not in one-hot format
+
+    Returns:
+       tf.Tensor: Loss value
+    """
+
+    onehot_labels = tf.one_hot(features["targets"],
+                               self._problem_hparams.vocab_size["targets"])
+    cost_vector = tf.nn.softmax_cross_entropy_with_logits_v2(
+      logits=logits,
+      labels=onehot_labels)
+    return tf.reduce_mean(cost_vector)
+
+  def body(self, features):
+    """Body of Neural Shuffle-Exchange network.
+
+    Args:
+      features: dictionary of inputs and targets
+    """
+
+    inputs = tf.squeeze(features["inputs"], axis=2)
+    logits = shuffle_network(inputs, self._hparams)
+    return tf.expand_dims(logits, axis=2)
+
+
+@registry.register_hparams
+def shuffle_network_baseline():
+  """Large Shuffle-Exchange configuration.
+
+  Returns:
+    dict: Neural Shuffle-Exchange configuration
+  """
+
+  hparams = common_hparams.basic_params1()
+  hparams.hidden_size = 48 * 8  # feature maps
+  hparams.num_hidden_layers = 2  # block count
+
+  hparams.clip_grad_norm = 0.  # no gradient clipping
+
+  hparams.optimizer = "adam"
+  hparams.optimizer_adam_epsilon = 1e-5
+  hparams.learning_rate_schedule = "legacy"
+  hparams.learning_rate_decay_scheme = "noam"
+  hparams.learning_rate = 0.1
+  hparams.initializer_gain = 1.0
+  hparams.initializer = "uniform_unit_scaling"
+  hparams.optimizer_adam_beta1 = 0.9
+  hparams.optimizer_adam_beta2 = 0.999
+
+  hparams.dropout = 0.1
+  hparams.label_smoothing = 0.
+  hparams.weight_decay = 0.
+
+  return hparams

From 0561ceb653893572e484ad0720f670598923486c Mon Sep 17 00:00:00 2001
From: Jonni Kanerva <jonni@google.com>
Date: Mon, 16 Dec 2019 14:48:39 -0800
Subject: [PATCH 2601/2720] Internal clean-up.

PiperOrigin-RevId: 285855523
---
 tensor2tensor/models/__init__.py              |  10 +-
 .../models/research/shuffle_network.py        | 413 ------------------
 2 files changed, 6 insertions(+), 417 deletions(-)
 delete mode 100644 tensor2tensor/models/research/shuffle_network.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 5f949dee5..d07545379 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -20,6 +20,8 @@
 
 import six
 
+# pylint: disable=unused-import
+
 from tensor2tensor.layers import modalities  # pylint: disable=g-import-not-at-top
 from tensor2tensor.models import basic
 from tensor2tensor.models import bytenet
@@ -42,7 +44,9 @@
 from tensor2tensor.models import transformer
 from tensor2tensor.models import vanilla_gan
 from tensor2tensor.models import xception
+
 from tensor2tensor.models.neural_architecture_search import nas_model
+
 from tensor2tensor.models.research import adafactor_experiments
 from tensor2tensor.models.research import aligned
 from tensor2tensor.models.research import attention_lm
@@ -56,7 +60,6 @@
 from tensor2tensor.models.research import multiquery_paper
 from tensor2tensor.models.research import neural_stack
 from tensor2tensor.models.research import rl
-from tensor2tensor.models.research import shuffle_network
 from tensor2tensor.models.research import similarity_transformer
 from tensor2tensor.models.research import super_lm
 from tensor2tensor.models.research import transformer_moe
@@ -70,6 +73,7 @@
 from tensor2tensor.models.research import vqa_attention
 from tensor2tensor.models.research import vqa_recurrent_self_attention
 from tensor2tensor.models.research import vqa_self_attention
+
 from tensor2tensor.models.video import basic_deterministic
 from tensor2tensor.models.video import basic_recurrent
 from tensor2tensor.models.video import basic_stochastic
@@ -78,10 +82,8 @@
 from tensor2tensor.models.video import next_frame_glow
 from tensor2tensor.models.video import savp
 from tensor2tensor.models.video import sv2p
-from tensor2tensor.utils import registry
-
 
-# pylint: disable=unused-import
+from tensor2tensor.utils import registry
 
 # pylint: enable=unused-import
 
diff --git a/tensor2tensor/models/research/shuffle_network.py b/tensor2tensor/models/research/shuffle_network.py
deleted file mode 100644
index a3307ce44..000000000
--- a/tensor2tensor/models/research/shuffle_network.py
+++ /dev/null
@@ -1,413 +0,0 @@
-"""Neural Shuffle-Exchange Network.
-
-Implementation of
-"Neural Shuffle-Exchange Networks - Sequence Processing in O(n log n) Time"
-paper by K.Freivalds, E.Ozolins, A.Sostaks.
-
-Paper: https://papers.nips.cc/paper/
-8889-neural-shuffle-exchange-networks-sequence-processing-in-on-log-n-time.pdf
-
-Original code: https://github.com/LUMII-Syslab/shuffle-exchange
-"""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensor2tensor.layers import common_hparams
-from tensor2tensor.utils import registry
-from tensor2tensor.utils import t2t_model
-import tensorflow as tf
-
-
-def ror(x, n, p=1):
-  """Bitwise right rotation.
-
-  Args:
-    x: Input tensor
-    n: Bit count to represent x
-    p: Bit positions to shift
-
-  Returns:
-    tf.Tensor: x shifted by p positions in n bits
-  """
-
-  a = tf.bitwise.right_shift(x, p)
-  b = tf.bitwise.left_shift(1, p) - 1
-  c = tf.bitwise.bitwise_and(x, b)
-  d = tf.bitwise.left_shift(c, n - p)
-
-  return a + d
-
-
-def rol(x, n, p=1):
-  """Bitwise left rotation.
-
-  Args:
-    x: Input tensor
-    n: Bit count to represent x
-    p: Bit positions to shift
-
-  Returns:
-    tf.Tensor: x shifted by p positions in n bits
-  """
-  a = tf.bitwise.left_shift(x, p)
-  b = tf.bitwise.left_shift(1, n) - 1
-  c = tf.bitwise.bitwise_and(a, b)
-  d = tf.bitwise.right_shift(x, n - p)
-
-  return tf.bitwise.bitwise_or(c, d)
-
-
-def shuffle_layer(inputs, shuffle_fn=rol):
-  """Shuffles the elements according to bitwise left or right rotation.
-
-  Args:
-    inputs: Tensor input from previous layer
-    shuffle_fn: Shift function rol or ror
-
-  Returns:
-    tf.Tensor: Inputs shifted according to shuffle_fn
-  """
-
-  length = tf.shape(inputs)[1]
-  n_bits = tf.log(tf.cast(length - 1, tf.float32)) / tf.log(2.0)
-  n_bits = tf.cast(n_bits, tf.int32) + 1
-
-  indices = tf.range(0, length)
-  rev_indices = shuffle_fn(indices, n_bits)
-  return tf.gather(inputs, rev_indices, axis=1)
-
-
-def reverse_shuffle_layer(inputs):
-  """Reverse shuffle of inputs. Used in the second half of Benes block.
-
-  Args:
-    inputs: Inputs that should be shuffled
-
-  Returns:
-    tf.Tensor: Inputs shuffled according to bitwise right rotation
-  """
-
-  return shuffle_layer(inputs, ror)
-
-
-def conv_linear_map(inputs, nin, nout, bias_start, prefix):
-  """Convolutional liner map. Maps 3D tensor by last dimension.
-
-  Args:
-    inputs: Inputs that should be shuffled
-    nin: Input feature map count
-    nout: Output feature map count
-    bias_start: Bias start value
-    prefix: Name prefix
-
-  Returns:
-    tf.Tensor: Inputs with applied convolution
-  """
-
-  with tf.variable_scope(prefix):
-    inp_shape = tf.shape(inputs)
-
-    initializer = tf.variance_scaling_initializer(scale=1.0,
-                                                  mode="fan_avg",
-                                                  distribution="uniform")
-    kernel = tf.get_variable("CvK", [nin, nout], initializer=initializer)
-    bias_term = tf.get_variable("CvB", [nout],
-                                initializer=tf.constant_initializer(0.0))
-
-    mul_shape = [inp_shape[0] * inp_shape[1], nin]
-    res = tf.matmul(tf.reshape(inputs, mul_shape), kernel)
-    res = tf.reshape(res, [inp_shape[0], inp_shape[1], nout])
-    return res + bias_start + bias_term
-
-
-# pylint: disable=useless-object-inheritance
-class SwitchLayer(object):
-  """Switch layer of Neural Shuffle-Exchange network.
-  """
-
-  def __init__(self, prefix, dropout, mode):
-    """Initialize switch layer.
-
-    Args:
-      prefix: Name prefix for switch layer
-      dropout: Dropout rate
-      mode: Training mode
-    """
-
-    self.prefix = prefix
-    self.dropout = dropout
-    self.mode = mode
-    self.batch_size = None
-    self.length = None
-    self.num_units = None
-    self.n_bits = None
-
-  def linear_map(self, inputs, suffix, bias_start, in_units, out_units):
-    """2 input to 2 output linear map.
-
-    Args:
-      inputs: Input tensor
-      suffix: Linear map name suffix
-      bias_start: Bias start value
-      in_units: Size of input tensor feature map count
-      out_units: Size of output tensor feature map count
-
-    Return:
-      tf.Tensor: Convolution apply to input tensor
-    """
-    in_shape = [self.batch_size, self.length // 2, in_units * 2]
-    inputs = tf.reshape(inputs, in_shape)
-    res = conv_linear_map(inputs, in_units * 2, out_units * 2,
-                          bias_start, self.prefix + "/" + suffix)
-    return tf.reshape(res, [self.batch_size, self.length, out_units])
-
-  def gated_linear_map(self, inputs, suffix, bias_start_reset,
-                       in_units, out_units):
-    """Linear mapping with two reset gates.
-
-    Args:
-      inputs: Input tensor
-      suffix: Linear map name suffix
-      bias_start_reset: Bias start value for reset gate
-      in_units: Size of input tensor feature map count
-      out_units: Size of output tensor feature map count
-
-    Return:
-      tf.Tensor: Convolution apply to input tensor
-    """
-
-    def reset_gate(name):
-      prefix = self.prefix + name + suffix
-      reset = conv_linear_map(inputs, in_units * 2, in_units * 2,
-                              bias_start_reset, prefix)
-      return tf.nn.sigmoid(reset)
-
-    in_shape = [self.batch_size, self.length // 2, in_units * 2]
-    inputs = tf.reshape(inputs, in_shape)
-
-    reset1 = reset_gate("/reset1/")
-    reset2 = reset_gate("/reset2/")
-    res1 = conv_linear_map(inputs * reset1, in_units * 2,
-                           out_units, 0.0, self.prefix + "/cand1/" + suffix)
-    res2 = conv_linear_map(inputs * reset2, in_units * 2,
-                           out_units, 0.0, self.prefix + "/cand2/" + suffix)
-
-    res = tf.concat([res1, res2], axis=2)
-    res = tf.reshape(res, [self.batch_size, self.length, out_units])
-    return tf.nn.tanh(res)
-
-  def __call__(self, inputs, residual_inputs):
-    """Apply SwitchLayer to inputs.
-
-    Args:
-      inputs: Input tensor
-      residual_inputs: Residual connections from previous block
-
-    Returns:
-      tf.Tensor: New candidate value
-    """
-    input_shape = tf.shape(inputs)
-    self.batch_size = input_shape[0]
-    self.length = input_shape[1]
-    self.num_units = inputs.shape.as_list()[2]
-
-    self.n_bits = tf.log(tf.cast(self.length - 1, tf.float32)) / tf.log(2.0)
-    self.n_bits = tf.floor(self.n_bits) + 1
-
-    initializer = tf.constant_initializer(0.5)
-    residual_scale = tf.get_variable(self.prefix + "/residual_scale",
-                                     [self.num_units], initializer=initializer)
-
-    shuffled_input = self.swap_halves(inputs)
-    mem_all = inputs + residual_inputs * residual_scale
-
-    # calculate the new value
-    candidate = self.gated_linear_map(mem_all, "c", 0.5,
-                                      self.num_units, self.num_units)
-    gate = tf.nn.sigmoid(self.linear_map(mem_all, "g", 0.5,
-                                         self.num_units, self.num_units))
-
-    candidate = gate * shuffled_input + (1 - gate) * candidate
-
-    if self.dropout > 0:
-      candidate = tf.nn.dropout(candidate, rate=self.dropout / self.n_bits)
-    if self.dropout != 0.0 and self.mode == tf.estimator.ModeKeys.TRAIN:
-      noise = tf.random_normal(tf.shape(candidate), mean=1.0, stddev=0.001)
-      candidate = candidate * noise
-
-    return candidate
-
-  def swap_halves(self, inputs):
-    """Split inputs in half and then shuffle them as described in paper.
-
-    Args:
-      inputs: ShuffleLayer inputs
-
-    Return:
-      tf.Tensor: Inputs with swapped halves
-    """
-    x = tf.range(0, self.length)
-    xor_indices = tf.bitwise.bitwise_xor(x, 1)
-    input_xor = tf.gather(inputs[:, :, :self.num_units // 2],
-                          xor_indices, axis=1)
-    return tf.concat([input_xor, inputs[:, :, self.num_units // 2:]], axis=2)
-
-
-def shuffle_network(inputs, hparams):
-  """Neural Shuffle-Network with skip connections between blocks.
-
-  Args:
-    inputs: inputs to the Shuffle-Exchange network.
-    Should be in length of power of 2.
-    hparams: Model configuration
-
-  Returns:
-    tf.Tensor: Outputs of the Shuffle-Exchange last layer
-  """
-
-  def forward_step(state, layer_nr):
-    with tf.variable_scope("forward"):
-      last_state, residuals = state
-      prev = residuals[layer_nr, :, :, :]
-      switch = SwitchLayer("switch", hparams.dropout, hparams.mode)
-      cur = switch(last_state, prev)
-      return shuffle_layer(cur), residuals
-
-  def reverse_step(state, layer_nr):
-    with tf.variable_scope("reverse"):
-      last_state, residuals = state
-      prev = residuals[layer_nr, :, :, :]
-      switch = SwitchLayer("reverse_switch", hparams.dropout, hparams.mode)
-      cur = switch(last_state, prev)
-      return reverse_shuffle_layer(cur), residuals
-
-  input_shape = tf.shape(inputs)
-  n_bits = tf.log(tf.cast(input_shape[1] - 1, tf.float32)) / tf.log(2.0)
-  n_bits = tf.cast(n_bits, tf.int32) + 1
-
-  queue_shape = [n_bits * 2, input_shape[0], input_shape[1], input_shape[2]]
-  residuals_queue = tf.zeros(queue_shape)
-  block_out = tf.tanh(inputs)
-
-  for k in range(hparams.num_hidden_layers):
-    with tf.variable_scope("benes_block_" + str(k), reuse=tf.AUTO_REUSE):
-      forward_outputs, _ = tf.scan(forward_step,
-                                   tf.range(0, n_bits),
-                                   initializer=(block_out, residuals_queue),
-                                   parallel_iterations=1,
-                                   swap_memory=True)
-
-      forward_tensors = [tf.expand_dims(block_out, axis=0), forward_outputs]
-      forward_outputs = tf.concat(forward_tensors, axis=0)
-      forward_last = forward_outputs[-1, :, :, :]
-
-      reverse_outputs, _ = tf.scan(reverse_step,
-                                   tf.range(n_bits, n_bits * 2),
-                                   initializer=(forward_last, residuals_queue),
-                                   parallel_iterations=1,
-                                   swap_memory=True)
-
-      block_out = reverse_outputs[-1, :, :, :]
-      residuals_queue = tf.concat([forward_outputs, reverse_outputs], axis=0)
-
-  last_layer = SwitchLayer("last_layer", hparams.dropout, hparams.mode)
-  return last_layer(block_out, residuals_queue[n_bits * 2, :, :, :])
-
-
-@registry.register_model
-class ShuffleNetwork(t2t_model.T2TModel):
-  """Seq2Seq model for sequence processing in O(n log n) time.
-  """
-
-  def bottom(self, features):
-    """We add padding to the input and output so they are the same.
-    Length of input and output should be power of 2.
-
-    Args:
-      features: Dictionary of inputs and targets
-
-    Returns:
-      dictionary: Inputs and targets padded with 0 to the length of power of 2.
-      Both are same length.
-    """
-    inputs = features["inputs"]
-    targets = features["targets"]
-    inputs_length = tf.shape(inputs)[1]
-    targets_length = tf.shape(targets)[1]
-
-    length = tf.maximum(inputs_length, targets_length)
-    p = tf.log(tf.cast(length, tf.float32)) / tf.log(2.0)
-    p = tf.cast(tf.ceil(p), tf.int32)
-    pad_len = tf.pow(2, p)
-
-    input_padding = [[0, 0], [0, pad_len - inputs_length], [0, 0], [0, 0]]
-    features["inputs"] = tf.pad(inputs, input_padding)
-    target_padding = [[0, 0], [0, pad_len - targets_length], [0, 0], [0, 0]]
-    features["targets"] = tf.pad(targets, target_padding)
-    return super(ShuffleNetwork, self).bottom(features)
-
-  def loss(self, logits, features):
-    """Loss function for Neural Shuffle-Exchange network.
-
-    We use custom loss function as default loss function doesn't
-    use padding for calculating loss. We assume that output string is same
-    length as the input. If you need other type of output please feel
-    free to modify this.
-
-    Args:
-      logits: Logits from model
-      features: Features, not in one-hot format
-
-    Returns:
-       tf.Tensor: Loss value
-    """
-
-    onehot_labels = tf.one_hot(features["targets"],
-                               self._problem_hparams.vocab_size["targets"])
-    cost_vector = tf.nn.softmax_cross_entropy_with_logits_v2(
-      logits=logits,
-      labels=onehot_labels)
-    return tf.reduce_mean(cost_vector)
-
-  def body(self, features):
-    """Body of Neural Shuffle-Exchange network.
-
-    Args:
-      features: dictionary of inputs and targets
-    """
-
-    inputs = tf.squeeze(features["inputs"], axis=2)
-    logits = shuffle_network(inputs, self._hparams)
-    return tf.expand_dims(logits, axis=2)
-
-
-@registry.register_hparams
-def shuffle_network_baseline():
-  """Large Shuffle-Exchange configuration.
-
-  Returns:
-    dict: Neural Shuffle-Exchange configuration
-  """
-
-  hparams = common_hparams.basic_params1()
-  hparams.hidden_size = 48 * 8  # feature maps
-  hparams.num_hidden_layers = 2  # block count
-
-  hparams.clip_grad_norm = 0.  # no gradient clipping
-
-  hparams.optimizer = "adam"
-  hparams.optimizer_adam_epsilon = 1e-5
-  hparams.learning_rate_schedule = "legacy"
-  hparams.learning_rate_decay_scheme = "noam"
-  hparams.learning_rate = 0.1
-  hparams.initializer_gain = 1.0
-  hparams.initializer = "uniform_unit_scaling"
-  hparams.optimizer_adam_beta1 = 0.9
-  hparams.optimizer_adam_beta2 = 0.999
-
-  hparams.dropout = 0.1
-  hparams.label_smoothing = 0.
-  hparams.weight_decay = 0.
-
-  return hparams

From 9514fb045efbc77a05aa640f89691ec03461c79f Mon Sep 17 00:00:00 2001
From: Em?ls Ozoli?? <ozolinsemils@gmail.com>
Date: Wed, 18 Dec 2019 16:30:21 -0800
Subject: [PATCH 2602/2720] Merge of PR #1763

PiperOrigin-RevId: 286289397
---
 tensor2tensor/models/__init__.py              |  10 +-
 .../models/research/shuffle_network.py        | 429 ++++++++++++++++++
 2 files changed, 433 insertions(+), 6 deletions(-)
 create mode 100644 tensor2tensor/models/research/shuffle_network.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index d07545379..5f949dee5 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -20,8 +20,6 @@
 
 import six
 
-# pylint: disable=unused-import
-
 from tensor2tensor.layers import modalities  # pylint: disable=g-import-not-at-top
 from tensor2tensor.models import basic
 from tensor2tensor.models import bytenet
@@ -44,9 +42,7 @@
 from tensor2tensor.models import transformer
 from tensor2tensor.models import vanilla_gan
 from tensor2tensor.models import xception
-
 from tensor2tensor.models.neural_architecture_search import nas_model
-
 from tensor2tensor.models.research import adafactor_experiments
 from tensor2tensor.models.research import aligned
 from tensor2tensor.models.research import attention_lm
@@ -60,6 +56,7 @@
 from tensor2tensor.models.research import multiquery_paper
 from tensor2tensor.models.research import neural_stack
 from tensor2tensor.models.research import rl
+from tensor2tensor.models.research import shuffle_network
 from tensor2tensor.models.research import similarity_transformer
 from tensor2tensor.models.research import super_lm
 from tensor2tensor.models.research import transformer_moe
@@ -73,7 +70,6 @@
 from tensor2tensor.models.research import vqa_attention
 from tensor2tensor.models.research import vqa_recurrent_self_attention
 from tensor2tensor.models.research import vqa_self_attention
-
 from tensor2tensor.models.video import basic_deterministic
 from tensor2tensor.models.video import basic_recurrent
 from tensor2tensor.models.video import basic_stochastic
@@ -82,9 +78,11 @@
 from tensor2tensor.models.video import next_frame_glow
 from tensor2tensor.models.video import savp
 from tensor2tensor.models.video import sv2p
-
 from tensor2tensor.utils import registry
 
+
+# pylint: disable=unused-import
+
 # pylint: enable=unused-import
 
 
diff --git a/tensor2tensor/models/research/shuffle_network.py b/tensor2tensor/models/research/shuffle_network.py
new file mode 100644
index 000000000..ca0f0d156
--- /dev/null
+++ b/tensor2tensor/models/research/shuffle_network.py
@@ -0,0 +1,429 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Neural Shuffle-Exchange Network.
+
+Implementation of
+"Neural Shuffle-Exchange Networks - Sequence Processing in O(n log n) Time"
+paper by K.Freivalds, E.Ozolins, A.Sostaks.
+
+Paper: https://papers.nips.cc/paper/
+8889-neural-shuffle-exchange-networks-sequence-processing-in-on-log-n-time.pdf
+
+Original code: https://github.com/LUMII-Syslab/shuffle-exchange
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_hparams
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
+import tensorflow as tf
+
+
+def ror(x, n, p=1):
+  """Bitwise right rotation.
+
+  Args:
+    x: Input tensor
+    n: Bit count to represent x
+    p: Bit positions to shift
+
+  Returns:
+    tf.Tensor: x shifted by p positions in n bits
+  """
+
+  a = tf.bitwise.right_shift(x, p)
+  b = tf.bitwise.left_shift(1, p) - 1
+  c = tf.bitwise.bitwise_and(x, b)
+  d = tf.bitwise.left_shift(c, n - p)
+
+  return a + d
+
+
+def rol(x, n, p=1):
+  """Bitwise left rotation.
+
+  Args:
+    x: Input tensor
+    n: Bit count to represent x
+    p: Bit positions to shift
+
+  Returns:
+    tf.Tensor: x shifted by p positions in n bits
+  """
+  a = tf.bitwise.left_shift(x, p)
+  b = tf.bitwise.left_shift(1, n) - 1
+  c = tf.bitwise.bitwise_and(a, b)
+  d = tf.bitwise.right_shift(x, n - p)
+
+  return tf.bitwise.bitwise_or(c, d)
+
+
+def shuffle_layer(inputs, shuffle_fn=rol):
+  """Shuffles the elements according to bitwise left or right rotation.
+
+  Args:
+    inputs: Tensor input from previous layer
+    shuffle_fn: Shift function rol or ror
+
+  Returns:
+    tf.Tensor: Inputs shifted according to shuffle_fn
+  """
+
+  length = tf.shape(inputs)[1]
+  n_bits = tf.log(tf.cast(length - 1, tf.float32)) / tf.log(2.0)
+  n_bits = tf.cast(n_bits, tf.int32) + 1
+
+  indices = tf.range(0, length)
+  rev_indices = shuffle_fn(indices, n_bits)
+  return tf.gather(inputs, rev_indices, axis=1)
+
+
+def reverse_shuffle_layer(inputs):
+  """Reverse shuffle of inputs.
+
+  Used in the second half of Benes block.
+
+  Args:
+    inputs: Inputs that should be shuffled
+
+  Returns:
+    tf.Tensor: Inputs shuffled according to bitwise right rotation
+  """
+
+  return shuffle_layer(inputs, ror)
+
+
+def conv_linear_map(inputs, nin, nout, bias_start, prefix):
+  """Convolutional liner map.
+
+  Maps 3D tensor by last dimension.
+
+  Args:
+    inputs: Inputs that should be shuffled
+    nin: Input feature map count
+    nout: Output feature map count
+    bias_start: Bias start value
+    prefix: Name prefix
+
+  Returns:
+    tf.Tensor: Inputs with applied convolution
+  """
+
+  with tf.variable_scope(prefix):
+    inp_shape = tf.shape(inputs)
+
+    initializer = tf.variance_scaling_initializer(
+        scale=1.0, mode="fan_avg", distribution="uniform")
+    kernel = tf.get_variable("CvK", [nin, nout], initializer=initializer)
+    bias_term = tf.get_variable(
+        "CvB", [nout], initializer=tf.constant_initializer(0.0))
+
+    mul_shape = [inp_shape[0] * inp_shape[1], nin]
+    res = tf.matmul(tf.reshape(inputs, mul_shape), kernel)
+    res = tf.reshape(res, [inp_shape[0], inp_shape[1], nout])
+    return res + bias_start + bias_term
+
+
+# pylint: disable=useless-object-inheritance
+class SwitchLayer(object):
+  """Switch layer of Neural Shuffle-Exchange network."""
+
+  def __init__(self, prefix, dropout, mode):
+    """Initialize switch layer.
+
+    Args:
+      prefix: Name prefix for switch layer
+      dropout: Dropout rate
+      mode: Training mode
+    """
+
+    self.prefix = prefix
+    self.dropout = dropout
+    self.mode = mode
+    self.batch_size = None
+    self.length = None
+    self.num_units = None
+    self.n_bits = None
+
+  def linear_map(self, inputs, suffix, bias_start, in_units, out_units):
+    """2 input to 2 output linear map.
+
+    Args:
+      inputs: Input tensor
+      suffix: Linear map name suffix
+      bias_start: Bias start value
+      in_units: Size of input tensor feature map count
+      out_units: Size of output tensor feature map count
+    Return:
+      tf.Tensor: Convolution apply to input tensor
+    """
+    in_shape = [self.batch_size, self.length // 2, in_units * 2]
+    inputs = tf.reshape(inputs, in_shape)
+    res = conv_linear_map(inputs, in_units * 2, out_units * 2, bias_start,
+                          self.prefix + "/" + suffix)
+    return tf.reshape(res, [self.batch_size, self.length, out_units])
+
+  def gated_linear_map(self, inputs, suffix, bias_start_reset, in_units,
+                       out_units):
+    """Linear mapping with two reset gates.
+
+    Args:
+      inputs: Input tensor
+      suffix: Linear map name suffix
+      bias_start_reset: Bias start value for reset gate
+      in_units: Size of input tensor feature map count
+      out_units: Size of output tensor feature map count
+    Return:
+      tf.Tensor: Convolution apply to input tensor
+    """
+
+    def reset_gate(name):
+      prefix = self.prefix + name + suffix
+      reset = conv_linear_map(inputs, in_units * 2, in_units * 2,
+                              bias_start_reset, prefix)
+      return tf.nn.sigmoid(reset)
+
+    in_shape = [self.batch_size, self.length // 2, in_units * 2]
+    inputs = tf.reshape(inputs, in_shape)
+
+    reset1 = reset_gate("/reset1/")
+    reset2 = reset_gate("/reset2/")
+    res1 = conv_linear_map(inputs * reset1, in_units * 2, out_units, 0.0,
+                           self.prefix + "/cand1/" + suffix)
+    res2 = conv_linear_map(inputs * reset2, in_units * 2, out_units, 0.0,
+                           self.prefix + "/cand2/" + suffix)
+
+    res = tf.concat([res1, res2], axis=2)
+    res = tf.reshape(res, [self.batch_size, self.length, out_units])
+    return tf.nn.tanh(res)
+
+  def __call__(self, inputs, residual_inputs):
+    """Apply SwitchLayer to inputs.
+
+    Args:
+      inputs: Input tensor
+      residual_inputs: Residual connections from previous block
+
+    Returns:
+      tf.Tensor: New candidate value
+    """
+    input_shape = tf.shape(inputs)
+    self.batch_size = input_shape[0]
+    self.length = input_shape[1]
+    self.num_units = inputs.shape.as_list()[2]
+
+    self.n_bits = tf.log(tf.cast(self.length - 1, tf.float32)) / tf.log(2.0)
+    self.n_bits = tf.floor(self.n_bits) + 1
+
+    initializer = tf.constant_initializer(0.5)
+    residual_scale = tf.get_variable(
+        self.prefix + "/residual_scale", [self.num_units],
+        initializer=initializer)
+
+    shuffled_input = self.swap_halves(inputs)
+    mem_all = inputs + residual_inputs * residual_scale
+
+    # calculate the new value
+    candidate = self.gated_linear_map(mem_all, "c", 0.5, self.num_units,
+                                      self.num_units)
+    gate = tf.nn.sigmoid(
+        self.linear_map(mem_all, "g", 0.5, self.num_units, self.num_units))
+
+    candidate = gate * shuffled_input + (1 - gate) * candidate
+
+    if self.dropout > 0:
+      candidate = tf.nn.dropout(candidate, rate=self.dropout / self.n_bits)
+    if self.dropout != 0.0 and self.mode == tf.estimator.ModeKeys.TRAIN:
+      noise = tf.random_normal(tf.shape(candidate), mean=1.0, stddev=0.001)
+      candidate = candidate * noise
+
+    return candidate
+
+  def swap_halves(self, inputs):
+    """Split inputs in half and then shuffle them as described in paper.
+
+    Args:
+      inputs: ShuffleLayer inputs
+    Return:
+      tf.Tensor: Inputs with swapped halves
+    """
+    x = tf.range(0, self.length)
+    xor_indices = tf.bitwise.bitwise_xor(x, 1)
+    input_xor = tf.gather(
+        inputs[:, :, :self.num_units // 2], xor_indices, axis=1)
+    return tf.concat([input_xor, inputs[:, :, self.num_units // 2:]], axis=2)
+
+
+def shuffle_network(inputs, hparams):
+  """Neural Shuffle-Network with skip connections between blocks.
+
+  Args:
+    inputs: inputs to the Shuffle-Exchange network. Should be in length of power
+      of 2.
+    hparams: Model configuration
+
+  Returns:
+    tf.Tensor: Outputs of the Shuffle-Exchange last layer
+  """
+
+  def forward_step(state, layer_nr):
+    with tf.variable_scope("forward"):
+      last_state, residuals = state
+      prev = residuals[layer_nr, :, :, :]
+      switch = SwitchLayer("switch", hparams.dropout, hparams.mode)
+      cur = switch(last_state, prev)
+      return shuffle_layer(cur), residuals
+
+  def reverse_step(state, layer_nr):
+    with tf.variable_scope("reverse"):
+      last_state, residuals = state
+      prev = residuals[layer_nr, :, :, :]
+      switch = SwitchLayer("reverse_switch", hparams.dropout, hparams.mode)
+      cur = switch(last_state, prev)
+      return reverse_shuffle_layer(cur), residuals
+
+  input_shape = tf.shape(inputs)
+  n_bits = tf.log(tf.cast(input_shape[1] - 1, tf.float32)) / tf.log(2.0)
+  n_bits = tf.cast(n_bits, tf.int32) + 1
+
+  queue_shape = [n_bits * 2, input_shape[0], input_shape[1], input_shape[2]]
+  residuals_queue = tf.zeros(queue_shape)
+  block_out = tf.tanh(inputs)
+
+  for k in range(hparams.num_hidden_layers):
+    with tf.variable_scope("benes_block_" + str(k), reuse=tf.AUTO_REUSE):
+      forward_outputs, _ = tf.scan(
+          forward_step,
+          tf.range(0, n_bits),
+          initializer=(block_out, residuals_queue),
+          parallel_iterations=1,
+          swap_memory=True)
+
+      forward_tensors = [tf.expand_dims(block_out, axis=0), forward_outputs]
+      forward_outputs = tf.concat(forward_tensors, axis=0)
+      forward_last = forward_outputs[-1, :, :, :]
+
+      reverse_outputs, _ = tf.scan(
+          reverse_step,
+          tf.range(n_bits, n_bits * 2),
+          initializer=(forward_last, residuals_queue),
+          parallel_iterations=1,
+          swap_memory=True)
+
+      block_out = reverse_outputs[-1, :, :, :]
+      residuals_queue = tf.concat([forward_outputs, reverse_outputs], axis=0)
+
+  last_layer = SwitchLayer("last_layer", hparams.dropout, hparams.mode)
+  return last_layer(block_out, residuals_queue[n_bits * 2, :, :, :])
+
+
+@registry.register_model
+class ShuffleNetwork(t2t_model.T2TModel):
+  """Seq2Seq model for sequence processing in O(n log n) time."""
+
+  def bottom(self, features):
+    """We add padding to the input and output so they are the same.
+
+    Length of input and output should be power of 2.
+
+    Args:
+      features: Dictionary of inputs and targets
+
+    Returns:
+      dictionary: Inputs and targets padded with 0 to the length of power of 2.
+      Both are same length.
+    """
+    inputs = features["inputs"]
+    targets = features["targets"]
+    inputs_length = tf.shape(inputs)[1]
+    targets_length = tf.shape(targets)[1]
+
+    length = tf.maximum(inputs_length, targets_length)
+    p = tf.log(tf.cast(length, tf.float32)) / tf.log(2.0)
+    p = tf.cast(tf.ceil(p), tf.int32)
+    pad_len = tf.pow(2, p)
+
+    input_padding = [[0, 0], [0, pad_len - inputs_length], [0, 0], [0, 0]]
+    features["inputs"] = tf.pad(inputs, input_padding)
+    target_padding = [[0, 0], [0, pad_len - targets_length], [0, 0], [0, 0]]
+    features["targets"] = tf.pad(targets, target_padding)
+    return super(ShuffleNetwork, self).bottom(features)
+
+  def loss(self, logits, features):
+    """Loss function for Neural Shuffle-Exchange network.
+
+    We use custom loss function as default loss function doesn't
+    use padding for calculating loss. We assume that output string is same
+    length as the input. If you need other type of output please feel
+    free to modify this.
+
+    Args:
+      logits: Logits from model
+      features: Features, not in one-hot format
+
+    Returns:
+       tf.Tensor: Loss value
+    """
+
+    onehot_labels = tf.one_hot(features["targets"],
+                               self._problem_hparams.vocab_size["targets"])
+    cost_vector = tf.nn.softmax_cross_entropy_with_logits_v2(
+        logits=logits, labels=onehot_labels)
+    return tf.reduce_mean(cost_vector)
+
+  def body(self, features):
+    """Body of Neural Shuffle-Exchange network.
+
+    Args:
+      features: dictionary of inputs and targets
+    """
+
+    inputs = tf.squeeze(features["inputs"], axis=2)
+    logits = shuffle_network(inputs, self._hparams)
+    return tf.expand_dims(logits, axis=2)
+
+
+@registry.register_hparams
+def shuffle_network_baseline():
+  """Large Shuffle-Exchange configuration.
+
+  Returns:
+    dict: Neural Shuffle-Exchange configuration
+  """
+
+  hparams = common_hparams.basic_params1()
+  hparams.hidden_size = 48 * 8  # feature maps
+  hparams.num_hidden_layers = 2  # block count
+
+  hparams.clip_grad_norm = 0.  # no gradient clipping
+
+  hparams.optimizer = "adam"
+  hparams.optimizer_adam_epsilon = 1e-5
+  hparams.learning_rate_schedule = "legacy"
+  hparams.learning_rate_decay_scheme = "noam"
+  hparams.learning_rate = 0.1
+  hparams.initializer_gain = 1.0
+  hparams.initializer = "uniform_unit_scaling"
+  hparams.optimizer_adam_beta1 = 0.9
+  hparams.optimizer_adam_beta2 = 0.999
+
+  hparams.dropout = 0.1
+  hparams.label_smoothing = 0.
+  hparams.weight_decay = 0.
+
+  return hparams

From 47e6390046f50d5254a209db9c2b9fc6b4f1882a Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 18 Dec 2019 16:48:02 -0800
Subject: [PATCH 2603/2720] Internal

PiperOrigin-RevId: 286292678
---
 tensor2tensor/models/__init__.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 5f949dee5..bf4a2e3d7 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -20,6 +20,8 @@
 
 import six
 
+# pylint: disable=unused-import
+
 from tensor2tensor.layers import modalities  # pylint: disable=g-import-not-at-top
 from tensor2tensor.models import basic
 from tensor2tensor.models import bytenet

From 4fa18266710e76160103d6f45de0e4d13d558415 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 19 Dec 2019 14:58:39 -0800
Subject: [PATCH 2604/2720] Specify a version for t2t notebooks intended for
 use in Colab.

Colab will soon update the default version of tensorflow to 2.1.0. In order for this notebook to continue to work, I'm adding a line magic that will ensure this notebook continues to use tensorflow 1.x and execute without errors.

PiperOrigin-RevId: 286469731
---
 .../notebooks/Transformer_translate.ipynb     |   1 +
 tensor2tensor/notebooks/asr_transformer.ipynb |   1 +
 tensor2tensor/notebooks/hello_t2t.ipynb       | 814 ++++++++----------
 tensor2tensor/notebooks/t2t_problem.ipynb     |   1 +
 4 files changed, 367 insertions(+), 450 deletions(-)

diff --git a/tensor2tensor/notebooks/Transformer_translate.ipynb b/tensor2tensor/notebooks/Transformer_translate.ipynb
index 07c350351..f32668f8c 100644
--- a/tensor2tensor/notebooks/Transformer_translate.ipynb
+++ b/tensor2tensor/notebooks/Transformer_translate.ipynb
@@ -108,6 +108,7 @@
         "colab": {}
       },
       "source": [
+        "%tensorflow_version 1.x\n",
         "import tensorflow as tf\n",
         "import os\n",
         "\n",
diff --git a/tensor2tensor/notebooks/asr_transformer.ipynb b/tensor2tensor/notebooks/asr_transformer.ipynb
index 82a0728a8..5c8b103a3 100644
--- a/tensor2tensor/notebooks/asr_transformer.ipynb
+++ b/tensor2tensor/notebooks/asr_transformer.ipynb
@@ -70,6 +70,7 @@
       },
       "outputs": [],
       "source": [
+        "%tensorflow_version 1.x\n",
         "import tensorflow as tf\n",
         "import matplotlib.pyplot as plt\n",
         "import numpy as np\n",
diff --git a/tensor2tensor/notebooks/hello_t2t.ipynb b/tensor2tensor/notebooks/hello_t2t.ipynb
index acd9c3e15..851fecba9 100644
--- a/tensor2tensor/notebooks/hello_t2t.ipynb
+++ b/tensor2tensor/notebooks/hello_t2t.ipynb
@@ -1,32 +1,11 @@
 {
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "name": "Tensor2Tensor Intro",
-      "version": "0.3.2",
-      "views": {},
-      "default_view": {},
-      "provenance": [
-        {
-          "file_id": "1-VScmaLkMqWiSbqgUCFWefzisSREd8l1",
-          "timestamp": 1512175750497
-        }
-      ],
-      "collapsed_sections": []
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    }
-  },
   "cells": [
     {
+      "cell_type": "markdown",
       "metadata": {
-        "id": "odi2vIMHC3Rm",
-        "colab_type": "text"
+        "colab_type": "text",
+        "id": "odi2vIMHC3Rm"
       },
-      "cell_type": "markdown",
       "source": [
         "# Welcome to the [Tensor2Tensor](https://github.com/tensorflow/tensor2tensor) Colab\n",
         "\n",
@@ -34,18 +13,15 @@
       ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "s19ucTii_wYb",
+        "cellView": "both",
+        "colab": {},
         "colab_type": "code",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        },
-        "cellView": "form"
+        "id": "s19ucTii_wYb"
       },
-      "cell_type": "code",
+      "outputs": [],
       "source": [
         "#@title\n",
         "# Copyright 2018 Google LLC.\n",
@@ -61,44 +37,35 @@
         "# 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."
-      ],
-      "execution_count": 0,
-      "outputs": []
+      ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "OPGni6fuvoTj",
+        "colab": {},
         "colab_type": "code",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        }
+        "id": "OPGni6fuvoTj"
       },
-      "cell_type": "code",
+      "outputs": [],
       "source": [
         "# Install deps\n",
         "!pip install -q -U tensor2tensor\n",
         "!pip install -q tensorflow matplotlib\n"
-      ],
-      "execution_count": 0,
-      "outputs": []
+      ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "oILRLCWN_16u",
+        "colab": {},
         "colab_type": "code",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        }
+        "id": "oILRLCWN_16u"
       },
-      "cell_type": "code",
+      "outputs": [],
       "source": [
         "# Imports we need.\n",
+        "%tensorflow_version 1.x\n",
         "import tensorflow as tf\n",
         "import matplotlib.pyplot as plt\n",
         "import numpy as np\n",
@@ -131,58 +98,43 @@
         "tf.gfile.MakeDirs(checkpoint_dir)\n",
         "gs_data_dir = \"gs://tensor2tensor-data\"\n",
         "gs_ckpt_dir = \"gs://tensor2tensor-checkpoints/\""
-      ],
-      "execution_count": 0,
-      "outputs": []
+      ]
     },
     {
+      "cell_type": "markdown",
       "metadata": {
-        "id": "0a69r1KDiZDe",
-        "colab_type": "text"
+        "colab_type": "text",
+        "id": "0a69r1KDiZDe"
       },
-      "cell_type": "markdown",
       "source": [
         "# Download MNIST and inspect it"
       ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "RYDMO4zArgkz",
-        "colab_type": "code",
         "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          },
           "base_uri": "https://localhost:8080/",
           "height": 1241
         },
-        "outputId": "f0f13103-a437-4b95-ac9d-38f2b57a5f4c",
+        "colab_type": "code",
         "executionInfo": {
+          "elapsed": 505,
           "status": "ok",
           "timestamp": 1512371452348,
-          "user_tz": 480,
-          "elapsed": 505,
           "user": {
             "displayName": "Lukasz Kaiser",
             "photoUrl": "//lh3.googleusercontent.com/-CbWIwcQ_VsA/AAAAAAAAAAI/AAAAAAAAAB8/jloHVR1qOhg/s50-c-k-no/photo.jpg",
             "userId": "109750154298538986950"
-          }
-        }
+          },
+          "user_tz": 480
+        },
+        "id": "RYDMO4zArgkz",
+        "outputId": "f0f13103-a437-4b95-ac9d-38f2b57a5f4c"
       },
-      "cell_type": "code",
-      "source": [
-        "# A Problem is a dataset together with some fixed pre-processing.\n",
-        "# It could be a translation dataset with a specific tokenization,\n",
-        "# or an image dataset with a specific resolution.\n",
-        "#\n",
-        "# There are many problems available in Tensor2Tensor\n",
-        "problems.available()"
-      ],
-      "execution_count": 0,
       "outputs": [
         {
-          "output_type": "execute_result",
           "data": {
             "text/plain": [
               "['algorithmic_addition_binary40',\n",
@@ -259,49 +211,48 @@
               " 'translate_enzh_wmt8k']"
             ]
           },
+          "execution_count": 4,
           "metadata": {
             "tags": []
           },
-          "execution_count": 4
+          "output_type": "execute_result"
         }
+      ],
+      "source": [
+        "# A Problem is a dataset together with some fixed pre-processing.\n",
+        "# It could be a translation dataset with a specific tokenization,\n",
+        "# or an image dataset with a specific resolution.\n",
+        "#\n",
+        "# There are many problems available in Tensor2Tensor\n",
+        "problems.available()"
       ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "JKc2uSk6WX5e",
-        "colab_type": "code",
         "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          },
           "base_uri": "https://localhost:8080/",
           "height": 306
         },
-        "outputId": "7e0cafb5-d035-49a7-9ff4-7f4150c905c7",
+        "colab_type": "code",
         "executionInfo": {
+          "elapsed": 21361,
           "status": "ok",
           "timestamp": 1512371478309,
-          "user_tz": 480,
-          "elapsed": 21361,
           "user": {
             "displayName": "Lukasz Kaiser",
             "photoUrl": "//lh3.googleusercontent.com/-CbWIwcQ_VsA/AAAAAAAAAAI/AAAAAAAAAB8/jloHVR1qOhg/s50-c-k-no/photo.jpg",
             "userId": "109750154298538986950"
-          }
-        }
+          },
+          "user_tz": 480
+        },
+        "id": "JKc2uSk6WX5e",
+        "outputId": "7e0cafb5-d035-49a7-9ff4-7f4150c905c7"
       },
-      "cell_type": "code",
-      "source": [
-        "# Fetch the MNIST problem\n",
-        "mnist_problem = problems.problem(\"image_mnist\")\n",
-        "# The generate_data method of a problem will download data and process it into\n",
-        "# a standard format ready for training and evaluation.\n",
-        "mnist_problem.generate_data(data_dir, tmp_dir)"
-      ],
-      "execution_count": 0,
       "outputs": [
         {
+          "name": "stdout",
           "output_type": "stream",
           "text": [
             "INFO:tensorflow:Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to /content/t2t/tmp/train-images-idx3-ubyte.gz\n",
@@ -321,106 +272,122 @@
             "INFO:tensorflow:Not downloading, file already found: /content/t2t/tmp/t10k-images-idx3-ubyte.gz\n",
             "INFO:tensorflow:Not downloading, file already found: /content/t2t/tmp/t10k-labels-idx1-ubyte.gz\n",
             "INFO:tensorflow:Shuffling data...\n"
-          ],
-          "name": "stdout"
+          ]
         }
+      ],
+      "source": [
+        "# Fetch the MNIST problem\n",
+        "mnist_problem = problems.problem(\"image_mnist\")\n",
+        "# The generate_data method of a problem will download data and process it into\n",
+        "# a standard format ready for training and evaluation.\n",
+        "mnist_problem.generate_data(data_dir, tmp_dir)"
       ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "VW6HCRANFPYV",
-        "colab_type": "code",
         "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          },
           "base_uri": "https://localhost:8080/",
           "height": 381
         },
-        "outputId": "3b33057c-5082-4377-ec83-79f67e5a8e84",
+        "colab_type": "code",
         "executionInfo": {
+          "elapsed": 471,
           "status": "ok",
           "timestamp": 1512371501917,
-          "user_tz": 480,
-          "elapsed": 471,
           "user": {
             "displayName": "Lukasz Kaiser",
             "photoUrl": "//lh3.googleusercontent.com/-CbWIwcQ_VsA/AAAAAAAAAAI/AAAAAAAAAB8/jloHVR1qOhg/s50-c-k-no/photo.jpg",
             "userId": "109750154298538986950"
-          }
-        }
+          },
+          "user_tz": 480
+        },
+        "id": "VW6HCRANFPYV",
+        "outputId": "3b33057c-5082-4377-ec83-79f67e5a8e84"
       },
-      "cell_type": "code",
-      "source": [
-        "# Now let's see the training MNIST data as Tensors.\n",
-        "mnist_example = tfe.Iterator(mnist_problem.dataset(Modes.TRAIN, data_dir)).next()\n",
-        "image = mnist_example[\"inputs\"]\n",
-        "label = mnist_example[\"targets\"]\n",
-        "\n",
-        "plt.imshow(image.numpy()[:, :, 0].astype(np.float32), cmap=plt.get_cmap('gray'))\n",
-        "print(\"Label: %d\" % label.numpy())"
-      ],
-      "execution_count": 0,
       "outputs": [
         {
+          "name": "stdout",
           "output_type": "stream",
           "text": [
             "INFO:tensorflow:Reading data files from /content/t2t/data/image_mnist-train*\n",
             "Label: 7\n"
-          ],
-          "name": "stdout"
+          ]
         },
         {
-          "output_type": "display_data",
           "data": {
             "image/png": "iVBORw0KGgoAAAANSUhEUgAAAUsAAAFKCAYAAACU6307AAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMS4wLCBo\ndHRwOi8vbWF0cGxvdGxpYi5vcmcvpW3flQAAEhNJREFUeJzt3V1IlPn7x/HP/J2VGir85arQYrtL\nGCtpBwuFGj1YEriwlNHDJiULHRRLkVmESA8LQZa5Rm4HqT0crCzMNkcdBErEQrQ6sR6EemJ1UCKt\naUkl2W7J/A9+/GTbHfVympn7nun9Ag+85+s918V3+nQ/zHfGEwqFQgIATOn/nC4AABIBYQkABoQl\nABgQlgBgQFgCgAFhCQAWoTiQFPanu7t70scS9ScZe0rWvugpcX7i1ddUPPF4n6XH4wm7PRQKTfpY\nokrGnqTk7IueEke8+poqDr2R7vTkyZO6e/euPB6PampqtHTp0kh3BQCuF1FY3rlzRw8fPpTf79eD\nBw9UU1Mjv98f7doAwDUiusHT0dGhkpISSdKiRYv0/PlzjY6ORrUwAHCTiI4sh4eHtWTJkonf58+f\nr6GhIc2ZMyfs+O7ubuXl5YV9LA6XTOMuGXuSkrMvekocTvcV8TXLv5uuifz8/En/LtkuRidjT1Jy\n9kVPicMNN3giOg3PzMzU8PDwxO9PnjxRRkZGJLsCgIQQUViuWLFCbW1tkqTe3l5lZmZOegoOAMkg\notPwL7/8UkuWLNE333wjj8ej48ePR7suAHAV3pQeZcnYk5ScfdFT4kjYa5YA8KEhLAHAgLAEAAPC\nEgAMCEsAMCAsAcCAsAQAA8ISAAwISwAwICwBwICwBAADwhIADAhLADAgLAHAgLAEAAPCEgAMCEsA\nMCAsAcCAsAQAA8ISAAwISwAwICwBwICwBAADwhIADAhLADAgLAHAgLAEAAPCEgAMCEsAMCAsAcCA\nsAQAA8ISAAwISwAwICwBwICwBAADwhIADAhLADAgLAHAgLAEAAPCEgAMCEsAMCAsAcCAsAQAA8IS\nAAy8kfxRMBjU/v37lZOTI0lavHixjh49GtXCAMBNIgpLSVq+fLkaGxujWQsAuBan4QBgEHFY3r9/\nX3v27NH27dt1+/btaNYEAK7jCYVCoZn+0eDgoLq6ulRaWqr+/n5VVFSovb1dqampYcf39PQoLy/v\nvYsFAKdEFJb/tHnzZp09e1bZ2dnhn8TjCbs9FApN+liiSsaepOTsi54SR7z6mioOIzoNv3btmi5d\nuiRJGhoa0tOnT5WVlRVZdQCQACI6shwdHdWhQ4f04sULvXnzRnv37tXq1asnfxKOLBNeMvZFT4nD\nDUeWUTkNnw5hmfiSsS96ShxuCMuI32cJJIrJrqWHe8z6zo6p9vlPDQ0N5rEHDx40j0V88T5LADAg\nLAHAgLAEAAPCEgAMCEsAMCAsAcCAsAQAA8ISAAwISwAwICwBwIDljkh6Uy1h/OdjM1nGaBUMBqO+\nT8QfR5YAYEBYAoABYQkABoQlABgQlgBgQFgCgAFhCQAGhCUAGBCWAGDAtztGWTL2JMWvL+sKGr/f\nb95nYWFhpOVMqqOjwzy2qKgo6s8/GV5/7/88k+HIEgAMCEsAMCAsAcCAsAQAA8ISAAwISwAwICwB\nwICwBAADwhIADAhLADDgC8vwjq1bt5oe++STT8z73LJli3lsLJYmzkR/f79pXDyXMMIdOLIEAAPC\nEgAMCEsAMCAsAcCAsAQAA8ISAAwISwAwICwBwICwBAADwhIADFju6DK//fabaZwTywJn8o2Kierq\n1atOlwCXMh1Z9vX1qaSkRK2trZKkx48fa+fOnSovL9f+/fv1119/xbRIAHDatGH56tUrnThx4p0j\nmcbGRpWXl+vnn3/Wp59+qkAgENMiAcBp04ZlamqqWlpalJmZObEtGAxq3bp1kqTi4uIZfeE8ACSi\naa9Zer1eeb3vDhsbG1NqaqokKT09XUNDQ7GpDgBc4r1v8IRCoWnHdHd3Ky8vL+K/TzTJ2NOHoqqq\nKqrjnJCsrz+n+4ooLH0+n16/fq1Zs2ZpcHDwnVP0cPLz88NuD4VC8ng8kZTgWu/bk5vvhn8IGhoa\nTOMOHjwY40oik4z/pqT49TVVIEf0PsuioiK1tbVJktrb27Vy5crIKgOABDHtkWVPT49Onz6tgYEB\neb1etbW1qb6+XtXV1fL7/VqwYIE2btwYj1oBwDHThmVeXp5++umnf22/cuVKTAoCADdiBY/LOH0t\ncrIv7MrOzn7nsbNnz5r3OTAwYB77yy+/mMbF6mJ/MBiMyX6R+FgbDgAGhCUAGBCWAGBAWAKAAWEJ\nAAaEJQAYEJYAYEBYAoABYQkABoQlABh4QnH4kLjJPlopGT9O6n17+uGHH0zjZrIsz7qEcCpum6tY\nvWwXLlxoGjfZslCnuW2eoiVhP6INAD40hCUAGBCWAGBAWAKAAWEJAAaEJQAYEJYAYEBYAoABYQkA\nBoQlABiw3DHKkrEnyX19xepl66YeI+G2eYoWljsCQIIgLAHAgLAEAAPCEgAMCEsAMCAsAcCAsAQA\nA8ISAAwISwAwICwBwICwBAADwhIADAhLADAgLAHAgLAEAAPCEgAMCEsAMCAsAcCAsAQAA8ISAAy8\nThcAJKKCggLz2IGBAfPY/v7+SMpBHHBkCQAGprDs6+tTSUmJWltbJUnV1dX6+uuvtXPnTu3cuVO/\n/vprLGsEAMdNexr+6tUrnThxQoWFhe9sr6qqUnFxccwKAwA3mfbIMjU1VS0tLcrMzIxHPQDgStMe\nWXq9Xnm9/x7W2tqqK1euKD09XUePHtX8+fMn3Ud3d7fy8vLCPhYKhWZQbmJIxp6k5O3r75Khx2To\nIRyn+4robviGDRuUlpam3NxcNTc36/z58zp27Nik4/Pz88NuD4VC8ng8kZTgWsnYk+S+vmL1D8fa\no1vvhrttnqIlXn1N9bqK6G54YWGhcnNzJUlr165VX19fZJUBQIKIKCz37ds38T9gMBhUTk5OVIsC\nALeZ9jS8p6dHp0+f1sDAgLxer9ra2rRjxw5VVlZq9uzZ8vl8qq2tjUetAOAYTygOV00nu9aQjNdX\nkrEnyX19cc0yPLfNU7S44Zolyx0RkUePHpnHZmdnm8devXo1knKiJhYhPJOetm7dGvXnR3Sw3BEA\nDAhLADAgLAHAgLAEAAPCEgAMCEsAMCAsAcCAsAQAA8ISAAwISwAwYLkjXGXLli1Ol2AykyWMBw8e\njGEliBeOLAHAgLAEAAPCEgAMCEsAMCAsAcCAsAQAA8ISAAwISwAwICwBwIBvd4yyZOxJ+ndfM/kS\nss2bN0e9noaGhqjvU5IWLlxoGve+38IYKx/K6y+WzzMZjiwBwICwBAADwhIADAhLADAgLAHAgLAE\nAAPCEgAMCEsAMCAsAcCAsAQAA5Y7Rlky9iS5r69YvWzd1GMk3DZP0cJyRwBIEIQlABgQlgBgQFgC\ngAFhCQAGhCUAGBCWAGBAWAKAAWEJAAaEJQAYeJ0uAHCTgoIC07jOzs4YVwK3MYVlXV2durq69Pbt\nW+3evVv5+fk6fPiwxsfHlZGRoTNnzig1NTXWtQKAY6YNy87OTt27d09+v18jIyMqKytTYWGhysvL\nVVpaqoaGBgUCAZWXl8ejXgBwxLTXLJctW6Zz585JkubNm6exsTEFg0GtW7dOklRcXKyOjo7YVgkA\nDps2LFNSUuTz+SRJgUBAq1at0tjY2MRpd3p6uoaGhmJbJQA4zHyD58aNGwoEArp8+bLWr18/sd3y\nuYLd3d3Ky8sL+1gcPk4z7pKxJyl5+/q7ZDhLStZ5crovU1jeunVLFy5c0MWLFzV37lz5fD69fv1a\ns2bN0uDgoDIzM6f8+/z8/LDbk/GDSpOxJ8l9fcXqH05hYaFpnFvvhrttnqIlIT789+XLl6qrq1NT\nU5PS0tIkSUVFRWpra5Mktbe3a+XKlVEqFQDcadojy+vXr2tkZESVlZUT206dOqUjR47I7/drwYIF\n2rhxY0yLBACn8R08UZaMPUnu64vT8PDcNk/R4obTcFbwICH19/ebx2ZnZ5vHDgwMRFIOPgCsDQcA\nA8ISAAwISwAwICwBwICwBAADwhIADAhLADAgLAHAgLAEAAPCEgAMWO6IhDSTtdkzWe5oXRs+k+WW\nSA4cWQKAAWEJAAaEJQAYEJYAYEBYAoABYQkABoQlABgQlgBgQFgCgAFhCQAGhCUAGBCWAGBAWAKA\nAWEJAAaEJQAYEJYAYEBYAoABYQkABoQlABgQlgBgQFgCgAFhCQAGhCUAGBCWAGBAWAKAAWEJAAaE\nJQAYEJYAYEBYAoABYQkABoQlABh4nS4AiEQgEDCPLSgoCLs9Oztb/f3972x79OjRe9WF5GUKy7q6\nOnV1dent27favXu3bt68qd7eXqWlpUmSdu3apTVr1sSyTgBw1LRh2dnZqXv37snv92tkZERlZWUq\nKChQVVWViouL41EjADhu2rBctmyZli5dKkmaN2+exsbGND4+HvPCAMBNpr3Bk5KSIp/PJ+m/14lW\nrVqllJQUtba2qqKiQgcOHNCzZ89iXigAOMkTCoVCloE3btxQU1OTLl++rJ6eHqWlpSk3N1fNzc36\n448/dOzYsUn/tqenR3l5eVErGgDizRSWt27d0rlz53Tx4sWJmzr/c//+fX3//fdqbW2d/Ek8nrDb\nQ6HQpI8lqmTsSXJfX1u3bjWPra+vD7s93N1w6347OzvNzx9PbpunaIlXX1PF4bSn4S9fvlRdXZ2a\nmpomgnLfvn0TL7JgMKicnJwolQoA7jTtDZ7r169rZGRElZWVE9s2bdqkyspKzZ49Wz6fT7W1tTEt\nEgCcNm1Ybtu2Tdu2bfvX9rKyspgUBABuxHJHADAw3w1/ryfhBk/CS8a+6ClxJMQNHgAAYQkAJoQl\nABgQlgBgQFgCgAFhCQAGhCUAGBCWAGBAWAKAAWEJAAaEJQAYEJYAYEBYAoABYQkABoQlABgQlgBg\nQFgCgAFhCQAGhCUAGBCWAGAQly8sA4BEx5ElABgQlgBgQFgCgAFhCQAGhCUAGBCWAGDgdeJJT548\nqbt378rj8aimpkZLly51ooyoCgaD2r9/v3JyciRJixcv1tGjRx2uKnJ9fX367rvv9O2332rHjh16\n/PixDh8+rPHxcWVkZOjMmTNKTU11uswZ+WdP1dXV6u3tVVpamiRp165dWrNmjbNFzlBdXZ26urr0\n9u1b7d69W/n5+Qk/T9K/+7p586bjcxX3sLxz544ePnwov9+vBw8eqKamRn6/P95lxMTy5cvV2Njo\ndBnv7dWrVzpx4oQKCwsntjU2Nqq8vFylpaVqaGhQIBBQeXm5g1XOTLieJKmqqkrFxcUOVfV+Ojs7\nde/ePfn9fo2MjKisrEyFhYUJPU9S+L4KCgocn6u4n4Z3dHSopKREkrRo0SI9f/5co6Oj8S4DU0hN\nTVVLS4syMzMntgWDQa1bt06SVFxcrI6ODqfKi0i4nhLdsmXLdO7cOUnSvHnzNDY2lvDzJIXva3x8\n3OGqHAjL4eFh/ec//5n4ff78+RoaGop3GTFx//597dmzR9u3b9ft27edLidiXq9Xs2bNemfb2NjY\nxOlcenp6ws1ZuJ4kqbW1VRUVFTpw4ICePXvmQGWRS0lJkc/nkyQFAgGtWrUq4edJCt9XSkqK43Pl\nyDXLv0uW1ZafffaZ9u7dq9LSUvX396uiokLt7e0Jeb1oOskyZxs2bFBaWppyc3PV3Nys8+fP69ix\nY06XNWM3btxQIBDQ5cuXtX79+ontiT5Pf++rp6fH8bmK+5FlZmamhoeHJ35/8uSJMjIy4l1G1GVl\nZemrr76Sx+PRwoUL9fHHH2twcNDpsqLG5/Pp9evXkqTBwcGkOJ0tLCxUbm6uJGnt2rXq6+tzuKKZ\nu3Xrli5cuKCWlhbNnTs3aebpn325Ya7iHpYrVqxQW1ubJKm3t1eZmZmaM2dOvMuIumvXrunSpUuS\npKGhIT19+lRZWVkOVxU9RUVFE/PW3t6ulStXOlzR+9u3b5/6+/sl/fea7P/eyZAoXr58qbq6OjU1\nNU3cJU6GeQrXlxvmypFPHaqvr9fvv/8uj8ej48eP64svvoh3CVE3OjqqQ4cO6cWLF3rz5o327t2r\n1atXO11WRHp6enT69GkNDAzI6/UqKytL9fX1qq6u1p9//qkFCxaotrZWH330kdOlmoXraceOHWpu\nbtbs2bPl8/lUW1ur9PR0p0s18/v9+vHHH/X5559PbDt16pSOHDmSsPMkhe9r06ZNam1tdXSu+Ig2\nADBgBQ8AGBCWAGBAWAKAAWEJAAaEJQAYEJYAYEBYAoABYQkABv8PicrBdxpy97QAAAAASUVORK5C\nYII=\n",
             "text/plain": [
-              "<matplotlib.figure.Figure at 0x7f9a730a8210>"
+              "\u003cmatplotlib.figure.Figure at 0x7f9a730a8210\u003e"
             ]
           },
           "metadata": {
             "tags": []
-          }
+          },
+          "output_type": "display_data"
         }
+      ],
+      "source": [
+        "# Now let's see the training MNIST data as Tensors.\n",
+        "mnist_example = tfe.Iterator(mnist_problem.dataset(Modes.TRAIN, data_dir)).next()\n",
+        "image = mnist_example[\"inputs\"]\n",
+        "label = mnist_example[\"targets\"]\n",
+        "\n",
+        "plt.imshow(image.numpy()[:, :, 0].astype(np.float32), cmap=plt.get_cmap('gray'))\n",
+        "print(\"Label: %d\" % label.numpy())"
       ]
     },
     {
+      "cell_type": "markdown",
       "metadata": {
-        "id": "gXL7_bVH49Kl",
-        "colab_type": "text"
+        "colab_type": "text",
+        "id": "gXL7_bVH49Kl"
       },
-      "cell_type": "markdown",
       "source": [
         "# Translate from English to German with a pre-trained model"
       ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "EB4MP7_y_SuQ",
-        "colab_type": "code",
         "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          },
           "base_uri": "https://localhost:8080/",
           "height": 170
         },
-        "outputId": "8fbdcd05-a8b6-45e5-88b2-ce6fdfec0351",
+        "colab_type": "code",
         "executionInfo": {
+          "elapsed": 2843,
           "status": "ok",
           "timestamp": 1512371509946,
-          "user_tz": 480,
-          "elapsed": 2843,
           "user": {
             "displayName": "Lukasz Kaiser",
             "photoUrl": "//lh3.googleusercontent.com/-CbWIwcQ_VsA/AAAAAAAAAAI/AAAAAAAAAB8/jloHVR1qOhg/s50-c-k-no/photo.jpg",
             "userId": "109750154298538986950"
-          }
-        }
+          },
+          "user_tz": 480
+        },
+        "id": "EB4MP7_y_SuQ",
+        "outputId": "8fbdcd05-a8b6-45e5-88b2-ce6fdfec0351"
       },
-      "cell_type": "code",
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\r\n",
+            "\r\n",
+            "Updates are available for some Cloud SDK components.  To install them,\r\n",
+            "please run:\r\n",
+            "  $ gcloud components update\r\n",
+            "\n",
+            "Copying gs://tensor2tensor-data/vocab.translate_ende_wmt32k.32768.subwords...\n",
+            "/ [1 files][316.4 KiB/316.4 KiB]                                                \n",
+            "Operation completed over 1 objects/316.4 KiB.                                    \n"
+          ]
+        }
+      ],
       "source": [
         "# Fetch the problem\n",
         "ende_problem = problems.problem(\"translate_ende_wmt32k\")\n",
@@ -447,38 +414,17 @@
         "  if 1 in integers:\n",
         "    integers = integers[:integers.index(1)]\n",
         "  return encoders[\"inputs\"].decode(np.squeeze(integers))"
-      ],
-      "execution_count": 0,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "text": [
-            "\r\n",
-            "\r\n",
-            "Updates are available for some Cloud SDK components.  To install them,\r\n",
-            "please run:\r\n",
-            "  $ gcloud components update\r\n",
-            "\n",
-            "Copying gs://tensor2tensor-data/vocab.translate_ende_wmt32k.32768.subwords...\n",
-            "/ [1 files][316.4 KiB/316.4 KiB]                                                \n",
-            "Operation completed over 1 objects/316.4 KiB.                                    \n"
-          ],
-          "name": "stdout"
-        }
       ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "g2aQW7Z6TOEu",
+        "colab": {},
         "colab_type": "code",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        }
+        "id": "g2aQW7Z6TOEu"
       },
-      "cell_type": "code",
+      "outputs": [],
       "source": [
         "# # Generate and view the data\n",
         "# # This cell is commented out because WMT data generation can take hours\n",
@@ -502,44 +448,33 @@
         "# # Example targets as a sentence.\n",
         "# print(\"Targets, decoded:\")\n",
         "# print(decode(targets))"
-      ],
-      "execution_count": 0,
-      "outputs": []
+      ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "WkFUEs7ZOA79",
-        "colab_type": "code",
         "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          },
           "base_uri": "https://localhost:8080/",
           "height": 408
         },
-        "outputId": "f8be52a4-e85c-4daf-9f77-24d75eea3ab0",
+        "colab_type": "code",
         "executionInfo": {
+          "elapsed": 496,
           "status": "ok",
           "timestamp": 1512371515918,
-          "user_tz": 480,
-          "elapsed": 496,
           "user": {
             "displayName": "Lukasz Kaiser",
             "photoUrl": "//lh3.googleusercontent.com/-CbWIwcQ_VsA/AAAAAAAAAAI/AAAAAAAAAB8/jloHVR1qOhg/s50-c-k-no/photo.jpg",
             "userId": "109750154298538986950"
-          }
-        }
+          },
+          "user_tz": 480
+        },
+        "id": "WkFUEs7ZOA79",
+        "outputId": "f8be52a4-e85c-4daf-9f77-24d75eea3ab0"
       },
-      "cell_type": "code",
-      "source": [
-        "# There are many models available in Tensor2Tensor\n",
-        "registry.list_models()"
-      ],
-      "execution_count": 0,
       "outputs": [
         {
-          "output_type": "execute_result",
           "data": {
             "text/plain": [
               "['resnet50',\n",
@@ -567,25 +502,27 @@
               " 'blue_net']"
             ]
           },
+          "execution_count": 9,
           "metadata": {
             "tags": []
           },
-          "execution_count": 9
+          "output_type": "execute_result"
         }
+      ],
+      "source": [
+        "# There are many models available in Tensor2Tensor\n",
+        "registry.list_models()"
       ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "9l6hDQbrRUYV",
+        "colab": {},
         "colab_type": "code",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        }
+        "id": "9l6hDQbrRUYV"
       },
-      "cell_type": "code",
+      "outputs": [],
       "source": [
         "# Create hparams and the model\n",
         "model_name = \"transformer\"\n",
@@ -597,86 +534,88 @@
         "# Layer and so subsequent instantiations will have different variable scopes\n",
         "# that will not match the checkpoint.\n",
         "translate_model = registry.model(model_name)(hparams, Modes.EVAL)"
-      ],
-      "execution_count": 0,
-      "outputs": []
+      ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "FEwNUVlMYOJi",
-        "colab_type": "code",
         "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          },
           "base_uri": "https://localhost:8080/",
           "height": 34
         },
-        "outputId": "86747a09-e83d-4a5f-d938-2fef25e4ce2f",
+        "colab_type": "code",
         "executionInfo": {
+          "elapsed": 13020,
           "status": "ok",
           "timestamp": 1512371536282,
-          "user_tz": 480,
-          "elapsed": 13020,
           "user": {
             "displayName": "Lukasz Kaiser",
             "photoUrl": "//lh3.googleusercontent.com/-CbWIwcQ_VsA/AAAAAAAAAAI/AAAAAAAAAB8/jloHVR1qOhg/s50-c-k-no/photo.jpg",
             "userId": "109750154298538986950"
-          }
-        }
+          },
+          "user_tz": 480
+        },
+        "id": "FEwNUVlMYOJi",
+        "outputId": "86747a09-e83d-4a5f-d938-2fef25e4ce2f"
       },
-      "cell_type": "code",
-      "source": [
-        "# Copy the pretrained checkpoint locally\n",
-        "ckpt_name = \"transformer_ende_test\"\n",
-        "gs_ckpt = os.path.join(gs_ckpt_dir, ckpt_name)\n",
-        "!gsutil -q cp -R {gs_ckpt} {checkpoint_dir}\n",
-        "ckpt_path = tf.train.latest_checkpoint(os.path.join(checkpoint_dir, ckpt_name))\n",
-        "ckpt_path"
-      ],
-      "execution_count": 0,
       "outputs": [
         {
-          "output_type": "execute_result",
           "data": {
             "text/plain": [
               "u'/content/t2t/checkpoints/transformer_ende_test/model.ckpt-350855'"
             ]
           },
+          "execution_count": 11,
           "metadata": {
             "tags": []
           },
-          "execution_count": 11
+          "output_type": "execute_result"
         }
+      ],
+      "source": [
+        "# Copy the pretrained checkpoint locally\n",
+        "ckpt_name = \"transformer_ende_test\"\n",
+        "gs_ckpt = os.path.join(gs_ckpt_dir, ckpt_name)\n",
+        "!gsutil -q cp -R {gs_ckpt} {checkpoint_dir}\n",
+        "ckpt_path = tf.train.latest_checkpoint(os.path.join(checkpoint_dir, ckpt_name))\n",
+        "ckpt_path"
       ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "3O-8E9d6TtuJ",
-        "colab_type": "code",
         "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          },
           "base_uri": "https://localhost:8080/",
           "height": 68
         },
-        "outputId": "cee729b7-8237-45bb-ac6f-dfadce9916b4",
+        "colab_type": "code",
         "executionInfo": {
+          "elapsed": 11397,
           "status": "ok",
           "timestamp": 1512371578480,
-          "user_tz": 480,
-          "elapsed": 11397,
           "user": {
             "displayName": "Lukasz Kaiser",
             "photoUrl": "//lh3.googleusercontent.com/-CbWIwcQ_VsA/AAAAAAAAAAI/AAAAAAAAAB8/jloHVR1qOhg/s50-c-k-no/photo.jpg",
             "userId": "109750154298538986950"
-          }
-        }
+          },
+          "user_tz": 480
+        },
+        "id": "3O-8E9d6TtuJ",
+        "outputId": "cee729b7-8237-45bb-ac6f-dfadce9916b4"
       },
-      "cell_type": "code",
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "INFO:tensorflow:Greedy Decoding\n",
+            "Inputs: The animal didn't cross the street because it was too tired\n",
+            "Outputs: Das Tier überquerte die Straße nicht, weil es zu müde war, weil es zu müde war.\n"
+          ]
+        }
+      ],
       "source": [
         "# Restore and translate!\n",
         "def translate(inputs):\n",
@@ -690,42 +629,27 @@
         "\n",
         "print(\"Inputs: %s\" % inputs)\n",
         "print(\"Outputs: %s\" % outputs)"
-      ],
-      "execution_count": 0,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "text": [
-            "INFO:tensorflow:Greedy Decoding\n",
-            "Inputs: The animal didn't cross the street because it was too tired\n",
-            "Outputs: Das Tier überquerte die Straße nicht, weil es zu müde war, weil es zu müde war.\n"
-          ],
-          "name": "stdout"
-        }
       ]
     },
     {
+      "cell_type": "markdown",
       "metadata": {
-        "id": "X3mkIEcbfiTP",
-        "colab_type": "text"
+        "colab_type": "text",
+        "id": "X3mkIEcbfiTP"
       },
-      "cell_type": "markdown",
       "source": [
         "## Attention Viz Utils"
       ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "r6GPPFy1fL2N",
+        "colab": {},
         "colab_type": "code",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        }
+        "id": "r6GPPFy1fL2N"
       },
-      "cell_type": "code",
+      "outputs": [],
       "source": [
         "from tensor2tensor.visualization import attention\n",
         "from tensor2tensor.data_generators import text_encoder\n",
@@ -771,36 +695,31 @@
         "  tokens = []\n",
         "  for _id in ids:\n",
         "    if _id == 0:\n",
-        "      tokens.append('<PAD>')\n",
+        "      tokens.append('\u003cPAD\u003e')\n",
         "    elif _id == 1:\n",
-        "      tokens.append('<EOS>')\n",
+        "      tokens.append('\u003cEOS\u003e')\n",
         "    elif _id == -1:\n",
-        "      tokens.append('<NULL>')\n",
+        "      tokens.append('\u003cNULL\u003e')\n",
         "    else:\n",
         "        tokens.append(subtokenizer._subtoken_id_to_subtoken_string(_id))\n",
         "  return tokens"
-      ],
-      "execution_count": 0,
-      "outputs": []
+      ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "wfF8_cW-OXPN",
+        "colab": {},
         "colab_type": "code",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        }
+        "id": "wfF8_cW-OXPN"
       },
-      "cell_type": "code",
+      "outputs": [],
       "source": [
         "def call_html():\n",
         "  import IPython\n",
         "  display(IPython.core.display.HTML('''\n",
-        "        <script src=\"/static/components/requirejs/require.js\"></script>\n",
-        "        <script>\n",
+        "        \u003cscript src=\"/static/components/requirejs/require.js\"\u003e\u003c/script\u003e\n",
+        "        \u003cscript\u003e\n",
         "          requirejs.config({\n",
         "            paths: {\n",
         "              base: '/static/base',\n",
@@ -808,81 +727,60 @@
         "              jquery: '//ajax.googleapis.com/ajax/libs/jquery/2.0.0/jquery.min',\n",
         "            },\n",
         "          });\n",
-        "        </script>\n",
+        "        \u003c/script\u003e\n",
         "        '''))"
-      ],
-      "execution_count": 0,
-      "outputs": []
+      ]
     },
     {
+      "cell_type": "markdown",
       "metadata": {
-        "id": "T7UJzFf6fmhp",
-        "colab_type": "text"
+        "colab_type": "text",
+        "id": "T7UJzFf6fmhp"
       },
-      "cell_type": "markdown",
       "source": [
         "## Display Attention"
       ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "OJKU36QAfqOC",
-        "colab_type": "code",
         "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          },
+          "base_uri": "https://localhost:8080/",
+          "height": 2006,
           "resources": {
             "http://localhost:8080/static/components/requirejs/require.js": {
               "data": "LyoqIHZpbTogZXQ6dHM9NDpzdz00OnN0cz00CiAqIEBsaWNlbnNlIFJlcXVpcmVKUyAyLjEuMjIgQ29weXJpZ2h0IChjKSAyMDEwLTIwMTUsIFRoZSBEb2pvIEZvdW5kYXRpb24gQWxsIFJpZ2h0cyBSZXNlcnZlZC4KICogQXZhaWxhYmxlIHZpYSB0aGUgTUlUIG9yIG5ldyBCU0QgbGljZW5zZS4KICogc2VlOiBodHRwOi8vZ2l0aHViLmNvbS9qcmJ1cmtlL3JlcXVpcmVqcyBmb3IgZGV0YWlscwogKi8KLy9Ob3QgdXNpbmcgc3RyaWN0OiB1bmV2ZW4gc3RyaWN0IHN1cHBvcnQgaW4gYnJvd3NlcnMsICMzOTIsIGFuZCBjYXVzZXMKLy9wcm9ibGVtcyB3aXRoIHJlcXVpcmVqcy5leGVjKCkvdHJhbnNwaWxlciBwbHVnaW5zIHRoYXQgbWF5IG5vdCBiZSBzdHJpY3QuCi8qanNsaW50IHJlZ2V4cDogdHJ1ZSwgbm9tZW46IHRydWUsIHNsb3BweTogdHJ1ZSAqLwovKmdsb2JhbCB3aW5kb3csIG5hdmlnYXRvciwgZG9jdW1lbnQsIGltcG9ydFNjcmlwdHMsIHNldFRpbWVvdXQsIG9wZXJhICovCgp2YXIgcmVxdWlyZWpzLCByZXF1aXJlLCBkZWZpbmU7CihmdW5jdGlvbiAoZ2xvYmFsKSB7CiAgICB2YXIgcmVxLCBzLCBoZWFkLCBiYXNlRWxlbWVudCwgZGF0YU1haW4sIHNyYywKICAgICAgICBpbnRlcmFjdGl2ZVNjcmlwdCwgY3VycmVudGx5QWRkaW5nU2NyaXB0LCBtYWluU2NyaXB0LCBzdWJQYXRoLAogICAgICAgIHZlcnNpb24gPSAnMi4xLjIyJywKICAgICAgICBjb21tZW50UmVnRXhwID0gLyhcL1wqKFtcc1xTXSo/KVwqXC98KFteOl18XilcL1wvKC4qKSQpL21nLAogICAgICAgIGNqc1JlcXVpcmVSZWdFeHAgPSAvW14uXVxzKnJlcXVpcmVccypcKFxzKlsiJ10oW14nIlxzXSspWyInXVxzKlwpL2csCiAgICAgICAganNTdWZmaXhSZWdFeHAgPSAvXC5qcyQvLAogICAgICAgIGN1cnJEaXJSZWdFeHAgPSAvXlwuXC8vLAogICAgICAgIG9wID0gT2JqZWN0LnByb3RvdHlwZSwKICAgICAgICBvc3RyaW5nID0gb3AudG9TdHJpbmcsCiAgICAgICAgaGFzT3duID0gb3AuaGFzT3duUHJvcGVydHksCiAgICAgICAgYXAgPSBBcnJheS5wcm90b3R5cGUsCiAgICAgICAgaXNCcm93c2VyID0gISEodHlwZW9mIHdpbmRvdyAhPT0gJ3VuZGVmaW5lZCcgJiYgdHlwZW9mIG5hdmlnYXRvciAhPT0gJ3VuZGVmaW5lZCcgJiYgd2luZG93LmRvY3VtZW50KSwKICAgICAgICBpc1dlYldvcmtlciA9ICFpc0Jyb3dzZXIgJiYgdHlwZW9mIGltcG9ydFNjcmlwdHMgIT09ICd1bmRlZmluZWQnLAogICAgICAgIC8vUFMzIGluZGljYXRlcyBsb2FkZWQgYW5kIGNvbXBsZXRlLCBidXQgbmVlZCB0byB3YWl0IGZvciBjb21wbGV0ZQogICAgICAgIC8vc3BlY2lmaWNhbGx5LiBTZXF1ZW5jZSBpcyAnbG9hZGluZycsICdsb2FkZWQnLCBleGVjdXRpb24sCiAgICAgICAgLy8gdGhlbiAnY29tcGxldGUnLiBUaGUgVUEgY2hlY2sgaXMgdW5mb3J0dW5hdGUsIGJ1dCBub3Qgc3VyZSBob3cKICAgICAgICAvL3RvIGZlYXR1cmUgdGVzdCB3L28gY2F1c2luZyBwZXJmIGlzc3Vlcy4KICAgICAgICByZWFkeVJlZ0V4cCA9IGlzQnJvd3NlciAmJiBuYXZpZ2F0b3IucGxhdGZvcm0gPT09ICdQTEFZU1RBVElPTiAzJyA/CiAgICAgICAgICAgICAgICAgICAgICAvXmNvbXBsZXRlJC8gOiAvXihjb21wbGV0ZXxsb2FkZWQpJC8sCiAgICAgICAgZGVmQ29udGV4dE5hbWUgPSAnXycsCiAgICAgICAgLy9PaCB0aGUgdHJhZ2VkeSwgZGV0ZWN0aW5nIG9wZXJhLiBTZWUgdGhlIHVzYWdlIG9mIGlzT3BlcmEgZm9yIHJlYXNvbi4KICAgICAgICBpc09wZXJhID0gdHlwZW9mIG9wZXJhICE9PSAndW5kZWZpbmVkJyAmJiBvcGVyYS50b1N0cmluZygpID09PSAnW29iamVjdCBPcGVyYV0nLAogICAgICAgIGNvbnRleHRzID0ge30sCiAgICAgICAgY2ZnID0ge30sCiAgICAgICAgZ2xvYmFsRGVmUXVldWUgPSBbXSwKICAgICAgICB1c2VJbnRlcmFjdGl2ZSA9IGZhbHNlOwoKICAgIGZ1bmN0aW9uIGlzRnVuY3Rpb24oaXQpIHsKICAgICAgICByZXR1cm4gb3N0cmluZy5jYWxsKGl0KSA9PT0gJ1tvYmplY3QgRnVuY3Rpb25dJzsKICAgIH0KCiAgICBmdW5jdGlvbiBpc0FycmF5KGl0KSB7CiAgICAgICAgcmV0dXJuIG9zdHJpbmcuY2FsbChpdCkgPT09ICdbb2JqZWN0IEFycmF5XSc7CiAgICB9CgogICAgLyoqCiAgICAgKiBIZWxwZXIgZnVuY3Rpb24gZm9yIGl0ZXJhdGluZyBvdmVyIGFuIGFycmF5LiBJZiB0aGUgZnVuYyByZXR1cm5zCiAgICAgKiBhIHRydWUgdmFsdWUsIGl0IHdpbGwgYnJlYWsgb3V0IG9mIHRoZSBsb29wLgogICAgICovCiAgICBmdW5jdGlvbiBlYWNoKGFyeSwgZnVuYykgewogICAgICAgIGlmIChhcnkpIHsKICAgICAgICAgICAgdmFyIGk7CiAgICAgICAgICAgIGZvciAoaSA9IDA7IGkgPCBhcnkubGVuZ3RoOyBpICs9IDEpIHsKICAgICAgICAgICAgICAgIGlmIChhcnlbaV0gJiYgZnVuYyhhcnlbaV0sIGksIGFyeSkpIHsKICAgICAgICAgICAgICAgICAgICBicmVhazsKICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgfQogICAgICAgIH0KICAgIH0KCiAgICAvKioKICAgICAqIEhlbHBlciBmdW5jdGlvbiBmb3IgaXRlcmF0aW5nIG92ZXIgYW4gYXJyYXkgYmFja3dhcmRzLiBJZiB0aGUgZnVuYwogICAgICogcmV0dXJucyBhIHRydWUgdmFsdWUsIGl0IHdpbGwgYnJlYWsgb3V0IG9mIHRoZSBsb29wLgogICAgICovCiAgICBmdW5jdGlvbiBlYWNoUmV2ZXJzZShhcnksIGZ1bmMpIHsKICAgICAgICBpZiAoYXJ5KSB7CiAgICAgICAgICAgIHZhciBpOwogICAgICAgICAgICBmb3IgKGkgPSBhcnkubGVuZ3RoIC0gMTsgaSA+IC0xOyBpIC09IDEpIHsKICAgICAgICAgICAgICAgIGlmIChhcnlbaV0gJiYgZnVuYyhhcnlbaV0sIGksIGFyeSkpIHsKICAgICAgICAgICAgICAgICAgICBicmVhazsKICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgfQogICAgICAgIH0KICAgIH0KCiAgICBmdW5jdGlvbiBoYXNQcm9wKG9iaiwgcHJvcCkgewogICAgICAgIHJldHVybiBoYXNPd24uY2FsbChvYmosIHByb3ApOwogICAgfQoKICAgIGZ1bmN0aW9uIGdldE93bihvYmosIHByb3ApIHsKICAgICAgICByZXR1cm4gaGFzUHJvcChvYmosIHByb3ApICYmIG9ialtwcm9wXTsKICAgIH0KCiAgICAvKioKICAgICAqIEN5Y2xlcyBvdmVyIHByb3BlcnRpZXMgaW4gYW4gb2JqZWN0IGFuZCBjYWxscyBhIGZ1bmN0aW9uIGZvciBlYWNoCiAgICAgKiBwcm9wZXJ0eSB2YWx1ZS4gSWYgdGhlIGZ1bmN0aW9uIHJldHVybnMgYSB0cnV0aHkgdmFsdWUsIHRoZW4gdGhlCiAgICAgKiBpdGVyYXRpb24gaXMgc3RvcHBlZC4KICAgICAqLwogICAgZnVuY3Rpb24gZWFjaFByb3Aob2JqLCBmdW5jKSB7CiAgICAgICAgdmFyIHByb3A7CiAgICAgICAgZm9yIChwcm9wIGluIG9iaikgewogICAgICAgICAgICBpZiAoaGFzUHJvcChvYmosIHByb3ApKSB7CiAgICAgICAgICAgICAgICBpZiAoZnVuYyhvYmpbcHJvcF0sIHByb3ApKSB7CiAgICAgICAgICAgICAgICAgICAgYnJlYWs7CiAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgIH0KICAgICAgICB9CiAgICB9CgogICAgLyoqCiAgICAgKiBTaW1wbGUgZnVuY3Rpb24gdG8gbWl4IGluIHByb3BlcnRpZXMgZnJvbSBzb3VyY2UgaW50byB0YXJnZXQsCiAgICAgKiBidXQgb25seSBpZiB0YXJnZXQgZG9lcyBub3QgYWxyZWFkeSBoYXZlIGEgcHJvcGVydHkgb2YgdGhlIHNhbWUgbmFtZS4KICAgICAqLwogICAgZnVuY3Rpb24gbWl4aW4odGFyZ2V0LCBzb3VyY2UsIGZvcmNlLCBkZWVwU3RyaW5nTWl4aW4pIHsKICAgICAgICBpZiAoc291cmNlKSB7CiAgICAgICAgICAgIGVhY2hQcm9wKHNvdXJjZSwgZnVuY3Rpb24gKHZhbHVlLCBwcm9wKSB7CiAgICAgICAgICAgICAgICBpZiAoZm9yY2UgfHwgIWhhc1Byb3AodGFyZ2V0LCBwcm9wKSkgewogICAgICAgICAgICAgICAgICAgIGlmIChkZWVwU3RyaW5nTWl4aW4gJiYgdHlwZW9mIHZhbHVlID09PSAnb2JqZWN0JyAmJiB2YWx1ZSAmJgogICAgICAgICAgICAgICAgICAgICAgICAhaXNBcnJheSh2YWx1ZSkgJiYgIWlzRnVuY3Rpb24odmFsdWUpICYmCiAgICAgICAgICAgICAgICAgICAgICAgICEodmFsdWUgaW5zdGFuY2VvZiBSZWdFeHApKSB7CgogICAgICAgICAgICAgICAgICAgICAgICBpZiAoIXRhcmdldFtwcm9wXSkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgdGFyZ2V0W3Byb3BdID0ge307CiAgICAgICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICAgICAgbWl4aW4odGFyZ2V0W3Byb3BdLCB2YWx1ZSwgZm9yY2UsIGRlZXBTdHJpbmdNaXhpbik7CiAgICAgICAgICAgICAgICAgICAgfSBlbHNlIHsKICAgICAgICAgICAgICAgICAgICAgICAgdGFyZ2V0W3Byb3BdID0gdmFsdWU7CiAgICAgICAgICAgICAgICAgICAgfQogICAgICAgICAgICAgICAgfQogICAgICAgICAgICB9KTsKICAgICAgICB9CiAgICAgICAgcmV0dXJuIHRhcmdldDsKICAgIH0KCiAgICAvL1NpbWlsYXIgdG8gRnVuY3Rpb24ucHJvdG90eXBlLmJpbmQsIGJ1dCB0aGUgJ3RoaXMnIG9iamVjdCBpcyBzcGVjaWZpZWQKICAgIC8vZmlyc3QsIHNpbmNlIGl0IGlzIGVhc2llciB0byByZWFkL2ZpZ3VyZSBvdXQgd2hhdCAndGhpcycgd2lsbCBiZS4KICAgIGZ1bmN0aW9uIGJpbmQob2JqLCBmbikgewogICAgICAgIHJldHVybiBmdW5jdGlvbiAoKSB7CiAgICAgICAgICAgIHJldHVybiBmbi5hcHBseShvYmosIGFyZ3VtZW50cyk7CiAgICAgICAgfTsKICAgIH0KCiAgICBmdW5jdGlvbiBzY3JpcHRzKCkgewogICAgICAgIHJldHVybiBkb2N1bWVudC5nZXRFbGVtZW50c0J5VGFnTmFtZSgnc2NyaXB0Jyk7CiAgICB9CgogICAgZnVuY3Rpb24gZGVmYXVsdE9uRXJyb3IoZXJyKSB7CiAgICAgICAgdGhyb3cgZXJyOwogICAgfQoKICAgIC8vQWxsb3cgZ2V0dGluZyBhIGdsb2JhbCB0aGF0IGlzIGV4cHJlc3NlZCBpbgogICAgLy9kb3Qgbm90YXRpb24sIGxpa2UgJ2EuYi5jJy4KICAgIGZ1bmN0aW9uIGdldEdsb2JhbCh2YWx1ZSkgewogICAgICAgIGlmICghdmFsdWUpIHsKICAgICAgICAgICAgcmV0dXJuIHZhbHVlOwogICAgICAgIH0KICAgICAgICB2YXIgZyA9IGdsb2JhbDsKICAgICAgICBlYWNoKHZhbHVlLnNwbGl0KCcuJyksIGZ1bmN0aW9uIChwYXJ0KSB7CiAgICAgICAgICAgIGcgPSBnW3BhcnRdOwogICAgICAgIH0pOwogICAgICAgIHJldHVybiBnOwogICAgfQoKICAgIC8qKgogICAgICogQ29uc3RydWN0cyBhbiBlcnJvciB3aXRoIGEgcG9pbnRlciB0byBhbiBVUkwgd2l0aCBtb3JlIGluZm9ybWF0aW9uLgogICAgICogQHBhcmFtIHtTdHJpbmd9IGlkIHRoZSBlcnJvciBJRCB0aGF0IG1hcHMgdG8gYW4gSUQgb24gYSB3ZWIgcGFnZS4KICAgICAqIEBwYXJhbSB7U3RyaW5nfSBtZXNzYWdlIGh1bWFuIHJlYWRhYmxlIGVycm9yLgogICAgICogQHBhcmFtIHtFcnJvcn0gW2Vycl0gdGhlIG9yaWdpbmFsIGVycm9yLCBpZiB0aGVyZSBpcyBvbmUuCiAgICAgKgogICAgICogQHJldHVybnMge0Vycm9yfQogICAgICovCiAgICBmdW5jdGlvbiBtYWtlRXJyb3IoaWQsIG1zZywgZXJyLCByZXF1aXJlTW9kdWxlcykgewogICAgICAgIHZhciBlID0gbmV3IEVycm9yKG1zZyArICdcbmh0dHA6Ly9yZXF1aXJlanMub3JnL2RvY3MvZXJyb3JzLmh0bWwjJyArIGlkKTsKICAgICAgICBlLnJlcXVpcmVUeXBlID0gaWQ7CiAgICAgICAgZS5yZXF1aXJlTW9kdWxlcyA9IHJlcXVpcmVNb2R1bGVzOwogICAgICAgIGlmIChlcnIpIHsKICAgICAgICAgICAgZS5vcmlnaW5hbEVycm9yID0gZXJyOwogICAgICAgIH0KICAgICAgICByZXR1cm4gZTsKICAgIH0KCiAgICBpZiAodHlwZW9mIGRlZmluZSAhPT0gJ3VuZGVmaW5lZCcpIHsKICAgICAgICAvL0lmIGEgZGVmaW5lIGlzIGFscmVhZHkgaW4gcGxheSB2aWEgYW5vdGhlciBBTUQgbG9hZGVyLAogICAgICAgIC8vZG8gbm90IG92ZXJ3cml0ZS4KICAgICAgICByZXR1cm47CiAgICB9CgogICAgaWYgKHR5cGVvZiByZXF1aXJlanMgIT09ICd1bmRlZmluZWQnKSB7CiAgICAgICAgaWYgKGlzRnVuY3Rpb24ocmVxdWlyZWpzKSkgewogICAgICAgICAgICAvL0RvIG5vdCBvdmVyd3JpdGUgYW4gZXhpc3RpbmcgcmVxdWlyZWpzIGluc3RhbmNlLgogICAgICAgICAgICByZXR1cm47CiAgICAgICAgfQogICAgICAgIGNmZyA9IHJlcXVpcmVqczsKICAgICAgICByZXF1aXJlanMgPSB1bmRlZmluZWQ7CiAgICB9CgogICAgLy9BbGxvdyBmb3IgYSByZXF1aXJlIGNvbmZpZyBvYmplY3QKICAgIGlmICh0eXBlb2YgcmVxdWlyZSAhPT0gJ3VuZGVmaW5lZCcgJiYgIWlzRnVuY3Rpb24ocmVxdWlyZSkpIHsKICAgICAgICAvL2Fzc3VtZSBpdCBpcyBhIGNvbmZpZyBvYmplY3QuCiAgICAgICAgY2ZnID0gcmVxdWlyZTsKICAgICAgICByZXF1aXJlID0gdW5kZWZpbmVkOwogICAgfQoKICAgIGZ1bmN0aW9uIG5ld0NvbnRleHQoY29udGV4dE5hbWUpIHsKICAgICAgICB2YXIgaW5DaGVja0xvYWRlZCwgTW9kdWxlLCBjb250ZXh0LCBoYW5kbGVycywKICAgICAgICAgICAgY2hlY2tMb2FkZWRUaW1lb3V0SWQsCiAgICAgICAgICAgIGNvbmZpZyA9IHsKICAgICAgICAgICAgICAgIC8vRGVmYXVsdHMuIERvIG5vdCBzZXQgYSBkZWZhdWx0IGZvciBtYXAKICAgICAgICAgICAgICAgIC8vY29uZmlnIHRvIHNwZWVkIHVwIG5vcm1hbGl6ZSgpLCB3aGljaAogICAgICAgICAgICAgICAgLy93aWxsIHJ1biBmYXN0ZXIgaWYgdGhlcmUgaXMgbm8gZGVmYXVsdC4KICAgICAgICAgICAgICAgIHdhaXRTZWNvbmRzOiA3LAogICAgICAgICAgICAgICAgYmFzZVVybDogJy4vJywKICAgICAgICAgICAgICAgIHBhdGhzOiB7fSwKICAgICAgICAgICAgICAgIGJ1bmRsZXM6IHt9LAogICAgICAgICAgICAgICAgcGtnczoge30sCiAgICAgICAgICAgICAgICBzaGltOiB7fSwKICAgICAgICAgICAgICAgIGNvbmZpZzoge30KICAgICAgICAgICAgfSwKICAgICAgICAgICAgcmVnaXN0cnkgPSB7fSwKICAgICAgICAgICAgLy9yZWdpc3RyeSBvZiBqdXN0IGVuYWJsZWQgbW9kdWxlcywgdG8gc3BlZWQKICAgICAgICAgICAgLy9jeWNsZSBicmVha2luZyBjb2RlIHdoZW4gbG90cyBvZiBtb2R1bGVzCiAgICAgICAgICAgIC8vYXJlIHJlZ2lzdGVyZWQsIGJ1dCBub3QgYWN0aXZhdGVkLgogICAgICAgICAgICBlbmFibGVkUmVnaXN0cnkgPSB7fSwKICAgICAgICAgICAgdW5kZWZFdmVudHMgPSB7fSwKICAgICAgICAgICAgZGVmUXVldWUgPSBbXSwKICAgICAgICAgICAgZGVmaW5lZCA9IHt9LAogICAgICAgICAgICB1cmxGZXRjaGVkID0ge30sCiAgICAgICAgICAgIGJ1bmRsZXNNYXAgPSB7fSwKICAgICAgICAgICAgcmVxdWlyZUNvdW50ZXIgPSAxLAogICAgICAgICAgICB1bm5vcm1hbGl6ZWRDb3VudGVyID0gMTsKCiAgICAgICAgLyoqCiAgICAgICAgICogVHJpbXMgdGhlIC4gYW5kIC4uIGZyb20gYW4gYXJyYXkgb2YgcGF0aCBzZWdtZW50cy4KICAgICAgICAgKiBJdCB3aWxsIGtlZXAgYSBsZWFkaW5nIHBhdGggc2VnbWVudCBpZiBhIC4uIHdpbGwgYmVjb21lCiAgICAgICAgICogdGhlIGZpcnN0IHBhdGggc2VnbWVudCwgdG8gaGVscCB3aXRoIG1vZHVsZSBuYW1lIGxvb2t1cHMsCiAgICAgICAgICogd2hpY2ggYWN0IGxpa2UgcGF0aHMsIGJ1dCBjYW4gYmUgcmVtYXBwZWQuIEJ1dCB0aGUgZW5kIHJlc3VsdCwKICAgICAgICAgKiBhbGwgcGF0aHMgdGhhdCB1c2UgdGhpcyBmdW5jdGlvbiBzaG91bGQgbG9vayBub3JtYWxpemVkLgogICAgICAgICAqIE5PVEU6IHRoaXMgbWV0aG9kIE1PRElGSUVTIHRoZSBpbnB1dCBhcnJheS4KICAgICAgICAgKiBAcGFyYW0ge0FycmF5fSBhcnkgdGhlIGFycmF5IG9mIHBhdGggc2VnbWVudHMuCiAgICAgICAgICovCiAgICAgICAgZnVuY3Rpb24gdHJpbURvdHMoYXJ5KSB7CiAgICAgICAgICAgIHZhciBpLCBwYXJ0OwogICAgICAgICAgICBmb3IgKGkgPSAwOyBpIDwgYXJ5Lmxlbmd0aDsgaSsrKSB7CiAgICAgICAgICAgICAgICBwYXJ0ID0gYXJ5W2ldOwogICAgICAgICAgICAgICAgaWYgKHBhcnQgPT09ICcuJykgewogICAgICAgICAgICAgICAgICAgIGFyeS5zcGxpY2UoaSwgMSk7CiAgICAgICAgICAgICAgICAgICAgaSAtPSAxOwogICAgICAgICAgICAgICAgfSBlbHNlIGlmIChwYXJ0ID09PSAnLi4nKSB7CiAgICAgICAgICAgICAgICAgICAgLy8gSWYgYXQgdGhlIHN0YXJ0LCBvciBwcmV2aW91cyB2YWx1ZSBpcyBzdGlsbCAuLiwKICAgICAgICAgICAgICAgICAgICAvLyBrZWVwIHRoZW0gc28gdGhhdCB3aGVuIGNvbnZlcnRlZCB0byBhIHBhdGggaXQgbWF5CiAgICAgICAgICAgICAgICAgICAgLy8gc3RpbGwgd29yayB3aGVuIGNvbnZlcnRlZCB0byBhIHBhdGgsIGV2ZW4gdGhvdWdoCiAgICAgICAgICAgICAgICAgICAgLy8gYXMgYW4gSUQgaXQgaXMgbGVzcyB0aGFuIGlkZWFsLiBJbiBsYXJnZXIgcG9pbnQKICAgICAgICAgICAgICAgICAgICAvLyByZWxlYXNlcywgbWF5IGJlIGJldHRlciB0byBqdXN0IGtpY2sgb3V0IGFuIGVycm9yLgogICAgICAgICAgICAgICAgICAgIGlmIChpID09PSAwIHx8IChpID09PSAxICYmIGFyeVsyXSA9PT0gJy4uJykgfHwgYXJ5W2kgLSAxXSA9PT0gJy4uJykgewogICAgICAgICAgICAgICAgICAgICAgICBjb250aW51ZTsKICAgICAgICAgICAgICAgICAgICB9IGVsc2UgaWYgKGkgPiAwKSB7CiAgICAgICAgICAgICAgICAgICAgICAgIGFyeS5zcGxpY2UoaSAtIDEsIDIpOwogICAgICAgICAgICAgICAgICAgICAgICBpIC09IDI7CiAgICAgICAgICAgICAgICAgICAgfQogICAgICAgICAgICAgICAgfQogICAgICAgICAgICB9CiAgICAgICAgfQoKICAgICAgICAvKioKICAgICAgICAgKiBHaXZlbiBhIHJlbGF0aXZlIG1vZHVsZSBuYW1lLCBsaWtlIC4vc29tZXRoaW5nLCBub3JtYWxpemUgaXQgdG8KICAgICAgICAgKiBhIHJlYWwgbmFtZSB0aGF0IGNhbiBiZSBtYXBwZWQgdG8gYSBwYXRoLgogICAgICAgICAqIEBwYXJhbSB7U3RyaW5nfSBuYW1lIHRoZSByZWxhdGl2ZSBuYW1lCiAgICAgICAgICogQHBhcmFtIHtTdHJpbmd9IGJhc2VOYW1lIGEgcmVhbCBuYW1lIHRoYXQgdGhlIG5hbWUgYXJnIGlzIHJlbGF0aXZlCiAgICAgICAgICogdG8uCiAgICAgICAgICogQHBhcmFtIHtCb29sZWFufSBhcHBseU1hcCBhcHBseSB0aGUgbWFwIGNvbmZpZyB0byB0aGUgdmFsdWUuIFNob3VsZAogICAgICAgICAqIG9ubHkgYmUgZG9uZSBpZiB0aGlzIG5vcm1hbGl6YXRpb24gaXMgZm9yIGEgZGVwZW5kZW5jeSBJRC4KICAgICAgICAgKiBAcmV0dXJucyB7U3RyaW5nfSBub3JtYWxpemVkIG5hbWUKICAgICAgICAgKi8KICAgICAgICBmdW5jdGlvbiBub3JtYWxpemUobmFtZSwgYmFzZU5hbWUsIGFwcGx5TWFwKSB7CiAgICAgICAgICAgIHZhciBwa2dNYWluLCBtYXBWYWx1ZSwgbmFtZVBhcnRzLCBpLCBqLCBuYW1lU2VnbWVudCwgbGFzdEluZGV4LAogICAgICAgICAgICAgICAgZm91bmRNYXAsIGZvdW5kSSwgZm91bmRTdGFyTWFwLCBzdGFySSwgbm9ybWFsaXplZEJhc2VQYXJ0cywKICAgICAgICAgICAgICAgIGJhc2VQYXJ0cyA9IChiYXNlTmFtZSAmJiBiYXNlTmFtZS5zcGxpdCgnLycpKSwKICAgICAgICAgICAgICAgIG1hcCA9IGNvbmZpZy5tYXAsCiAgICAgICAgICAgICAgICBzdGFyTWFwID0gbWFwICYmIG1hcFsnKiddOwoKICAgICAgICAgICAgLy9BZGp1c3QgYW55IHJlbGF0aXZlIHBhdGhzLgogICAgICAgICAgICBpZiAobmFtZSkgewogICAgICAgICAgICAgICAgbmFtZSA9IG5hbWUuc3BsaXQoJy8nKTsKICAgICAgICAgICAgICAgIGxhc3RJbmRleCA9IG5hbWUubGVuZ3RoIC0gMTsKCiAgICAgICAgICAgICAgICAvLyBJZiB3YW50aW5nIG5vZGUgSUQgY29tcGF0aWJpbGl0eSwgc3RyaXAgLmpzIGZyb20gZW5kCiAgICAgICAgICAgICAgICAvLyBvZiBJRHMuIEhhdmUgdG8gZG8gdGhpcyBoZXJlLCBhbmQgbm90IGluIG5hbWVUb1VybAogICAgICAgICAgICAgICAgLy8gYmVjYXVzZSBub2RlIGFsbG93cyBlaXRoZXIgLmpzIG9yIG5vbiAuanMgdG8gbWFwCiAgICAgICAgICAgICAgICAvLyB0byBzYW1lIGZpbGUuCiAgICAgICAgICAgICAgICBpZiAoY29uZmlnLm5vZGVJZENvbXBhdCAmJiBqc1N1ZmZpeFJlZ0V4cC50ZXN0KG5hbWVbbGFzdEluZGV4XSkpIHsKICAgICAgICAgICAgICAgICAgICBuYW1lW2xhc3RJbmRleF0gPSBuYW1lW2xhc3RJbmRleF0ucmVwbGFjZShqc1N1ZmZpeFJlZ0V4cCwgJycpOwogICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgIC8vIFN0YXJ0cyB3aXRoIGEgJy4nIHNvIG5lZWQgdGhlIGJhc2VOYW1lCiAgICAgICAgICAgICAgICBpZiAobmFtZVswXS5jaGFyQXQoMCkgPT09ICcuJyAmJiBiYXNlUGFydHMpIHsKICAgICAgICAgICAgICAgICAgICAvL0NvbnZlcnQgYmFzZU5hbWUgdG8gYXJyYXksIGFuZCBsb3Agb2ZmIHRoZSBsYXN0IHBhcnQsCiAgICAgICAgICAgICAgICAgICAgLy9zbyB0aGF0IC4gbWF0Y2hlcyB0aGF0ICdkaXJlY3RvcnknIGFuZCBub3QgbmFtZSBvZiB0aGUgYmFzZU5hbWUncwogICAgICAgICAgICAgICAgICAgIC8vbW9kdWxlLiBGb3IgaW5zdGFuY2UsIGJhc2VOYW1lIG9mICdvbmUvdHdvL3RocmVlJywgbWFwcyB0bwogICAgICAgICAgICAgICAgICAgIC8vJ29uZS90d28vdGhyZWUuanMnLCBidXQgd2Ugd2FudCB0aGUgZGlyZWN0b3J5LCAnb25lL3R3bycgZm9yCiAgICAgICAgICAgICAgICAgICAgLy90aGlzIG5vcm1hbGl6YXRpb24uCiAgICAgICAgICAgICAgICAgICAgbm9ybWFsaXplZEJhc2VQYXJ0cyA9IGJhc2VQYXJ0cy5zbGljZSgwLCBiYXNlUGFydHMubGVuZ3RoIC0gMSk7CiAgICAgICAgICAgICAgICAgICAgbmFtZSA9IG5vcm1hbGl6ZWRCYXNlUGFydHMuY29uY2F0KG5hbWUpOwogICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgIHRyaW1Eb3RzKG5hbWUpOwogICAgICAgICAgICAgICAgbmFtZSA9IG5hbWUuam9pbignLycpOwogICAgICAgICAgICB9CgogICAgICAgICAgICAvL0FwcGx5IG1hcCBjb25maWcgaWYgYXZhaWxhYmxlLgogICAgICAgICAgICBpZiAoYXBwbHlNYXAgJiYgbWFwICYmIChiYXNlUGFydHMgfHwgc3Rhck1hcCkpIHsKICAgICAgICAgICAgICAgIG5hbWVQYXJ0cyA9IG5hbWUuc3BsaXQoJy8nKTsKCiAgICAgICAgICAgICAgICBvdXRlckxvb3A6IGZvciAoaSA9IG5hbWVQYXJ0cy5sZW5ndGg7IGkgPiAwOyBpIC09IDEpIHsKICAgICAgICAgICAgICAgICAgICBuYW1lU2VnbWVudCA9IG5hbWVQYXJ0cy5zbGljZSgwLCBpKS5qb2luKCcvJyk7CgogICAgICAgICAgICAgICAgICAgIGlmIChiYXNlUGFydHMpIHsKICAgICAgICAgICAgICAgICAgICAgICAgLy9GaW5kIHRoZSBsb25nZXN0IGJhc2VOYW1lIHNlZ21lbnQgbWF0Y2ggaW4gdGhlIGNvbmZpZy4KICAgICAgICAgICAgICAgICAgICAgICAgLy9TbywgZG8gam9pbnMgb24gdGhlIGJpZ2dlc3QgdG8gc21hbGxlc3QgbGVuZ3RocyBvZiBiYXNlUGFydHMuCiAgICAgICAgICAgICAgICAgICAgICAgIGZvciAoaiA9IGJhc2VQYXJ0cy5sZW5ndGg7IGogPiAwOyBqIC09IDEpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIG1hcFZhbHVlID0gZ2V0T3duKG1hcCwgYmFzZVBhcnRzLnNsaWNlKDAsIGopLmpvaW4oJy8nKSk7CgogICAgICAgICAgICAgICAgICAgICAgICAgICAgLy9iYXNlTmFtZSBzZWdtZW50IGhhcyBjb25maWcsIGZpbmQgaWYgaXQgaGFzIG9uZSBmb3IKICAgICAgICAgICAgICAgICAgICAgICAgICAgIC8vdGhpcyBuYW1lLgogICAgICAgICAgICAgICAgICAgICAgICAgICAgaWYgKG1hcFZhbHVlKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgbWFwVmFsdWUgPSBnZXRPd24obWFwVmFsdWUsIG5hbWVTZWdtZW50KTsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBpZiAobWFwVmFsdWUpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgLy9NYXRjaCwgdXBkYXRlIG5hbWUgdG8gdGhlIG5ldyB2YWx1ZS4KICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgZm91bmRNYXAgPSBtYXBWYWx1ZTsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgZm91bmRJID0gaTsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgYnJlYWsgb3V0ZXJMb29wOwogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICAgICAgfQogICAgICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICAgICAgLy9DaGVjayBmb3IgYSBzdGFyIG1hcCBtYXRjaCwgYnV0IGp1c3QgaG9sZCBvbiB0byBpdCwKICAgICAgICAgICAgICAgICAgICAvL2lmIHRoZXJlIGlzIGEgc2hvcnRlciBzZWdtZW50IG1hdGNoIGxhdGVyIGluIGEgbWF0Y2hpbmcKICAgICAgICAgICAgICAgICAgICAvL2NvbmZpZywgdGhlbiBmYXZvciBvdmVyIHRoaXMgc3RhciBtYXAuCiAgICAgICAgICAgICAgICAgICAgaWYgKCFmb3VuZFN0YXJNYXAgJiYgc3Rhck1hcCAmJiBnZXRPd24oc3Rhck1hcCwgbmFtZVNlZ21lbnQpKSB7CiAgICAgICAgICAgICAgICAgICAgICAgIGZvdW5kU3Rhck1hcCA9IGdldE93bihzdGFyTWFwLCBuYW1lU2VnbWVudCk7CiAgICAgICAgICAgICAgICAgICAgICAgIHN0YXJJID0gaTsKICAgICAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgaWYgKCFmb3VuZE1hcCAmJiBmb3VuZFN0YXJNYXApIHsKICAgICAgICAgICAgICAgICAgICBmb3VuZE1hcCA9IGZvdW5kU3Rhck1hcDsKICAgICAgICAgICAgICAgICAgICBmb3VuZEkgPSBzdGFySTsKICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICBpZiAoZm91bmRNYXApIHsKICAgICAgICAgICAgICAgICAgICBuYW1lUGFydHMuc3BsaWNlKDAsIGZvdW5kSSwgZm91bmRNYXApOwogICAgICAgICAgICAgICAgICAgIG5hbWUgPSBuYW1lUGFydHMuam9pbignLycpOwogICAgICAgICAgICAgICAgfQogICAgICAgICAgICB9CgogICAgICAgICAgICAvLyBJZiB0aGUgbmFtZSBwb2ludHMgdG8gYSBwYWNrYWdlJ3MgbmFtZSwgdXNlCiAgICAgICAgICAgIC8vIHRoZSBwYWNrYWdlIG1haW4gaW5zdGVhZC4KICAgICAgICAgICAgcGtnTWFpbiA9IGdldE93bihjb25maWcucGtncywgbmFtZSk7CgogICAgICAgICAgICByZXR1cm4gcGtnTWFpbiA/IHBrZ01haW4gOiBuYW1lOwogICAgICAgIH0KCiAgICAgICAgZnVuY3Rpb24gcmVtb3ZlU2NyaXB0KG5hbWUpIHsKICAgICAgICAgICAgaWYgKGlzQnJvd3NlcikgewogICAgICAgICAgICAgICAgZWFjaChzY3JpcHRzKCksIGZ1bmN0aW9uIChzY3JpcHROb2RlKSB7CiAgICAgICAgICAgICAgICAgICAgaWYgKHNjcmlwdE5vZGUuZ2V0QXR0cmlidXRlKCdkYXRhLXJlcXVpcmVtb2R1bGUnKSA9PT0gbmFtZSAmJgogICAgICAgICAgICAgICAgICAgICAgICAgICAgc2NyaXB0Tm9kZS5nZXRBdHRyaWJ1dGUoJ2RhdGEtcmVxdWlyZWNvbnRleHQnKSA9PT0gY29udGV4dC5jb250ZXh0TmFtZSkgewogICAgICAgICAgICAgICAgICAgICAgICBzY3JpcHROb2RlLnBhcmVudE5vZGUucmVtb3ZlQ2hpbGQoc2NyaXB0Tm9kZSk7CiAgICAgICAgICAgICAgICAgICAgICAgIHJldHVybiB0cnVlOwogICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgIH0pOwogICAgICAgICAgICB9CiAgICAgICAgfQoKICAgICAgICBmdW5jdGlvbiBoYXNQYXRoRmFsbGJhY2soaWQpIHsKICAgICAgICAgICAgdmFyIHBhdGhDb25maWcgPSBnZXRPd24oY29uZmlnLnBhdGhzLCBpZCk7CiAgICAgICAgICAgIGlmIChwYXRoQ29uZmlnICYmIGlzQXJyYXkocGF0aENvbmZpZykgJiYgcGF0aENvbmZpZy5sZW5ndGggPiAxKSB7CiAgICAgICAgICAgICAgICAvL1BvcCBvZmYgdGhlIGZpcnN0IGFycmF5IHZhbHVlLCBzaW5jZSBpdCBmYWlsZWQsIGFuZAogICAgICAgICAgICAgICAgLy9yZXRyeQogICAgICAgICAgICAgICAgcGF0aENvbmZpZy5zaGlmdCgpOwogICAgICAgICAgICAgICAgY29udGV4dC5yZXF1aXJlLnVuZGVmKGlkKTsKCiAgICAgICAgICAgICAgICAvL0N1c3RvbSByZXF1aXJlIHRoYXQgZG9lcyBub3QgZG8gbWFwIHRyYW5zbGF0aW9uLCBzaW5jZQogICAgICAgICAgICAgICAgLy9JRCBpcyAiYWJzb2x1dGUiLCBhbHJlYWR5IG1hcHBlZC9yZXNvbHZlZC4KICAgICAgICAgICAgICAgIGNvbnRleHQubWFrZVJlcXVpcmUobnVsbCwgewogICAgICAgICAgICAgICAgICAgIHNraXBNYXA6IHRydWUKICAgICAgICAgICAgICAgIH0pKFtpZF0pOwoKICAgICAgICAgICAgICAgIHJldHVybiB0cnVlOwogICAgICAgICAgICB9CiAgICAgICAgfQoKICAgICAgICAvL1R1cm5zIGEgcGx1Z2luIXJlc291cmNlIHRvIFtwbHVnaW4sIHJlc291cmNlXQogICAgICAgIC8vd2l0aCB0aGUgcGx1Z2luIGJlaW5nIHVuZGVmaW5lZCBpZiB0aGUgbmFtZQogICAgICAgIC8vZGlkIG5vdCBoYXZlIGEgcGx1Z2luIHByZWZpeC4KICAgICAgICBmdW5jdGlvbiBzcGxpdFByZWZpeChuYW1lKSB7CiAgICAgICAgICAgIHZhciBwcmVmaXgsCiAgICAgICAgICAgICAgICBpbmRleCA9IG5hbWUgPyBuYW1lLmluZGV4T2YoJyEnKSA6IC0xOwogICAgICAgICAgICBpZiAoaW5kZXggPiAtMSkgewogICAgICAgICAgICAgICAgcHJlZml4ID0gbmFtZS5zdWJzdHJpbmcoMCwgaW5kZXgpOwogICAgICAgICAgICAgICAgbmFtZSA9IG5hbWUuc3Vic3RyaW5nKGluZGV4ICsgMSwgbmFtZS5sZW5ndGgpOwogICAgICAgICAgICB9CiAgICAgICAgICAgIHJldHVybiBbcHJlZml4LCBuYW1lXTsKICAgICAgICB9CgogICAgICAgIC8qKgogICAgICAgICAqIENyZWF0ZXMgYSBtb2R1bGUgbWFwcGluZyB0aGF0IGluY2x1ZGVzIHBsdWdpbiBwcmVmaXgsIG1vZHVsZQogICAgICAgICAqIG5hbWUsIGFuZCBwYXRoLiBJZiBwYXJlbnRNb2R1bGVNYXAgaXMgcHJvdmlkZWQgaXQgd2lsbAogICAgICAgICAqIGFsc28gbm9ybWFsaXplIHRoZSBuYW1lIHZpYSByZXF1aXJlLm5vcm1hbGl6ZSgpCiAgICAgICAgICoKICAgICAgICAgKiBAcGFyYW0ge1N0cmluZ30gbmFtZSB0aGUgbW9kdWxlIG5hbWUKICAgICAgICAgKiBAcGFyYW0ge1N0cmluZ30gW3BhcmVudE1vZHVsZU1hcF0gcGFyZW50IG1vZHVsZSBtYXAKICAgICAgICAgKiBmb3IgdGhlIG1vZHVsZSBuYW1lLCB1c2VkIHRvIHJlc29sdmUgcmVsYXRpdmUgbmFtZXMuCiAgICAgICAgICogQHBhcmFtIHtCb29sZWFufSBpc05vcm1hbGl6ZWQ6IGlzIHRoZSBJRCBhbHJlYWR5IG5vcm1hbGl6ZWQuCiAgICAgICAgICogVGhpcyBpcyB0cnVlIGlmIHRoaXMgY2FsbCBpcyBkb25lIGZvciBhIGRlZmluZSgpIG1vZHVsZSBJRC4KICAgICAgICAgKiBAcGFyYW0ge0Jvb2xlYW59IGFwcGx5TWFwOiBhcHBseSB0aGUgbWFwIGNvbmZpZyB0byB0aGUgSUQuCiAgICAgICAgICogU2hvdWxkIG9ubHkgYmUgdHJ1ZSBpZiB0aGlzIG1hcCBpcyBmb3IgYSBkZXBlbmRlbmN5LgogICAgICAgICAqCiAgICAgICAgICogQHJldHVybnMge09iamVjdH0KICAgICAgICAgKi8KICAgICAgICBmdW5jdGlvbiBtYWtlTW9kdWxlTWFwKG5hbWUsIHBhcmVudE1vZHVsZU1hcCwgaXNOb3JtYWxpemVkLCBhcHBseU1hcCkgewogICAgICAgICAgICB2YXIgdXJsLCBwbHVnaW5Nb2R1bGUsIHN1ZmZpeCwgbmFtZVBhcnRzLAogICAgICAgICAgICAgICAgcHJlZml4ID0gbnVsbCwKICAgICAgICAgICAgICAgIHBhcmVudE5hbWUgPSBwYXJlbnRNb2R1bGVNYXAgPyBwYXJlbnRNb2R1bGVNYXAubmFtZSA6IG51bGwsCiAgICAgICAgICAgICAgICBvcmlnaW5hbE5hbWUgPSBuYW1lLAogICAgICAgICAgICAgICAgaXNEZWZpbmUgPSB0cnVlLAogICAgICAgICAgICAgICAgbm9ybWFsaXplZE5hbWUgPSAnJzsKCiAgICAgICAgICAgIC8vSWYgbm8gbmFtZSwgdGhlbiBpdCBtZWFucyBpdCBpcyBhIHJlcXVpcmUgY2FsbCwgZ2VuZXJhdGUgYW4KICAgICAgICAgICAgLy9pbnRlcm5hbCBuYW1lLgogICAgICAgICAgICBpZiAoIW5hbWUpIHsKICAgICAgICAgICAgICAgIGlzRGVmaW5lID0gZmFsc2U7CiAgICAgICAgICAgICAgICBuYW1lID0gJ19AcicgKyAocmVxdWlyZUNvdW50ZXIgKz0gMSk7CiAgICAgICAgICAgIH0KCiAgICAgICAgICAgIG5hbWVQYXJ0cyA9IHNwbGl0UHJlZml4KG5hbWUpOwogICAgICAgICAgICBwcmVmaXggPSBuYW1lUGFydHNbMF07CiAgICAgICAgICAgIG5hbWUgPSBuYW1lUGFydHNbMV07CgogICAgICAgICAgICBpZiAocHJlZml4KSB7CiAgICAgICAgICAgICAgICBwcmVmaXggPSBub3JtYWxpemUocHJlZml4LCBwYXJlbnROYW1lLCBhcHBseU1hcCk7CiAgICAgICAgICAgICAgICBwbHVnaW5Nb2R1bGUgPSBnZXRPd24oZGVmaW5lZCwgcHJlZml4KTsKICAgICAgICAgICAgfQoKICAgICAgICAgICAgLy9BY2NvdW50IGZvciByZWxhdGl2ZSBwYXRocyBpZiB0aGVyZSBpcyBhIGJhc2UgbmFtZS4KICAgICAgICAgICAgaWYgKG5hbWUpIHsKICAgICAgICAgICAgICAgIGlmIChwcmVmaXgpIHsKICAgICAgICAgICAgICAgICAgICBpZiAocGx1Z2luTW9kdWxlICYmIHBsdWdpbk1vZHVsZS5ub3JtYWxpemUpIHsKICAgICAgICAgICAgICAgICAgICAgICAgLy9QbHVnaW4gaXMgbG9hZGVkLCB1c2UgaXRzIG5vcm1hbGl6ZSBtZXRob2QuCiAgICAgICAgICAgICAgICAgICAgICAgIG5vcm1hbGl6ZWROYW1lID0gcGx1Z2luTW9kdWxlLm5vcm1hbGl6ZShuYW1lLCBmdW5jdGlvbiAobmFtZSkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgcmV0dXJuIG5vcm1hbGl6ZShuYW1lLCBwYXJlbnROYW1lLCBhcHBseU1hcCk7CiAgICAgICAgICAgICAgICAgICAgICAgIH0pOwogICAgICAgICAgICAgICAgICAgIH0gZWxzZSB7CiAgICAgICAgICAgICAgICAgICAgICAgIC8vIElmIG5lc3RlZCBwbHVnaW4gcmVmZXJlbmNlcywgdGhlbiBkbyBub3QgdHJ5IHRvCiAgICAgICAgICAgICAgICAgICAgICAgIC8vIG5vcm1hbGl6ZSwgYXMgaXQgd2lsbCBub3Qgbm9ybWFsaXplIGNvcnJlY3RseS4gVGhpcwogICAgICAgICAgICAgICAgICAgICAgICAvLyBwbGFjZXMgYSByZXN0cmljdGlvbiBvbiByZXNvdXJjZUlkcywgYW5kIHRoZSBsb25nZXIKICAgICAgICAgICAgICAgICAgICAgICAgLy8gdGVybSBzb2x1dGlvbiBpcyBub3QgdG8gbm9ybWFsaXplIHVudGlsIHBsdWdpbnMgYXJlCiAgICAgICAgICAgICAgICAgICAgICAgIC8vIGxvYWRlZCBhbmQgYWxsIG5vcm1hbGl6YXRpb25zIHRvIGFsbG93IGZvciBhc3luYwogICAgICAgICAgICAgICAgICAgICAgICAvLyBsb2FkaW5nIG9mIGEgbG9hZGVyIHBsdWdpbi4gQnV0IGZvciBub3csIGZpeGVzIHRoZQogICAgICAgICAgICAgICAgICAgICAgICAvLyBjb21tb24gdXNlcy4gRGV0YWlscyBpbiAjMTEzMQogICAgICAgICAgICAgICAgICAgICAgICBub3JtYWxpemVkTmFtZSA9IG5hbWUuaW5kZXhPZignIScpID09PSAtMSA/CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgbm9ybWFsaXplKG5hbWUsIHBhcmVudE5hbWUsIGFwcGx5TWFwKSA6CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgbmFtZTsKICAgICAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgICAgICB9IGVsc2UgewogICAgICAgICAgICAgICAgICAgIC8vQSByZWd1bGFyIG1vZHVsZS4KICAgICAgICAgICAgICAgICAgICBub3JtYWxpemVkTmFtZSA9IG5vcm1hbGl6ZShuYW1lLCBwYXJlbnROYW1lLCBhcHBseU1hcCk7CgogICAgICAgICAgICAgICAgICAgIC8vTm9ybWFsaXplZCBuYW1lIG1heSBiZSBhIHBsdWdpbiBJRCBkdWUgdG8gbWFwIGNvbmZpZwogICAgICAgICAgICAgICAgICAgIC8vYXBwbGljYXRpb24gaW4gbm9ybWFsaXplLiBUaGUgbWFwIGNvbmZpZyB2YWx1ZXMgbXVzdAogICAgICAgICAgICAgICAgICAgIC8vYWxyZWFkeSBiZSBub3JtYWxpemVkLCBzbyBkbyBub3QgbmVlZCB0byByZWRvIHRoYXQgcGFydC4KICAgICAgICAgICAgICAgICAgICBuYW1lUGFydHMgPSBzcGxpdFByZWZpeChub3JtYWxpemVkTmFtZSk7CiAgICAgICAgICAgICAgICAgICAgcHJlZml4ID0gbmFtZVBhcnRzWzBdOwogICAgICAgICAgICAgICAgICAgIG5vcm1hbGl6ZWROYW1lID0gbmFtZVBhcnRzWzFdOwogICAgICAgICAgICAgICAgICAgIGlzTm9ybWFsaXplZCA9IHRydWU7CgogICAgICAgICAgICAgICAgICAgIHVybCA9IGNvbnRleHQubmFtZVRvVXJsKG5vcm1hbGl6ZWROYW1lKTsKICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgfQoKICAgICAgICAgICAgLy9JZiB0aGUgaWQgaXMgYSBwbHVnaW4gaWQgdGhhdCBjYW5ub3QgYmUgZGV0ZXJtaW5lZCBpZiBpdCBuZWVkcwogICAgICAgICAgICAvL25vcm1hbGl6YXRpb24sIHN0YW1wIGl0IHdpdGggYSB1bmlxdWUgSUQgc28gdHdvIG1hdGNoaW5nIHJlbGF0aXZlCiAgICAgICAgICAgIC8vaWRzIHRoYXQgbWF5IGNvbmZsaWN0IGNhbiBiZSBzZXBhcmF0ZS4KICAgICAgICAgICAgc3VmZml4ID0gcHJlZml4ICYmICFwbHVnaW5Nb2R1bGUgJiYgIWlzTm9ybWFsaXplZCA/CiAgICAgICAgICAgICAgICAgICAgICdfdW5ub3JtYWxpemVkJyArICh1bm5vcm1hbGl6ZWRDb3VudGVyICs9IDEpIDoKICAgICAgICAgICAgICAgICAgICAgJyc7CgogICAgICAgICAgICByZXR1cm4gewogICAgICAgICAgICAgICAgcHJlZml4OiBwcmVmaXgsCiAgICAgICAgICAgICAgICBuYW1lOiBub3JtYWxpemVkTmFtZSwKICAgICAgICAgICAgICAgIHBhcmVudE1hcDogcGFyZW50TW9kdWxlTWFwLAogICAgICAgICAgICAgICAgdW5ub3JtYWxpemVkOiAhIXN1ZmZpeCwKICAgICAgICAgICAgICAgIHVybDogdXJsLAogICAgICAgICAgICAgICAgb3JpZ2luYWxOYW1lOiBvcmlnaW5hbE5hbWUsCiAgICAgICAgICAgICAgICBpc0RlZmluZTogaXNEZWZpbmUsCiAgICAgICAgICAgICAgICBpZDogKHByZWZpeCA/CiAgICAgICAgICAgICAgICAgICAgICAgIHByZWZpeCArICchJyArIG5vcm1hbGl6ZWROYW1lIDoKICAgICAgICAgICAgICAgICAgICAgICAgbm9ybWFsaXplZE5hbWUpICsgc3VmZml4CiAgICAgICAgICAgIH07CiAgICAgICAgfQoKICAgICAgICBmdW5jdGlvbiBnZXRNb2R1bGUoZGVwTWFwKSB7CiAgICAgICAgICAgIHZhciBpZCA9IGRlcE1hcC5pZCwKICAgICAgICAgICAgICAgIG1vZCA9IGdldE93bihyZWdpc3RyeSwgaWQpOwoKICAgICAgICAgICAgaWYgKCFtb2QpIHsKICAgICAgICAgICAgICAgIG1vZCA9IHJlZ2lzdHJ5W2lkXSA9IG5ldyBjb250ZXh0Lk1vZHVsZShkZXBNYXApOwogICAgICAgICAgICB9CgogICAgICAgICAgICByZXR1cm4gbW9kOwogICAgICAgIH0KCiAgICAgICAgZnVuY3Rpb24gb24oZGVwTWFwLCBuYW1lLCBmbikgewogICAgICAgICAgICB2YXIgaWQgPSBkZXBNYXAuaWQsCiAgICAgICAgICAgICAgICBtb2QgPSBnZXRPd24ocmVnaXN0cnksIGlkKTsKCiAgICAgICAgICAgIGlmIChoYXNQcm9wKGRlZmluZWQsIGlkKSAmJgogICAgICAgICAgICAgICAgICAgICghbW9kIHx8IG1vZC5kZWZpbmVFbWl0Q29tcGxldGUpKSB7CiAgICAgICAgICAgICAgICBpZiAobmFtZSA9PT0gJ2RlZmluZWQnKSB7CiAgICAgICAgICAgICAgICAgICAgZm4oZGVmaW5lZFtpZF0pOwogICAgICAgICAgICAgICAgfQogICAgICAgICAgICB9IGVsc2UgewogICAgICAgICAgICAgICAgbW9kID0gZ2V0TW9kdWxlKGRlcE1hcCk7CiAgICAgICAgICAgICAgICBpZiAobW9kLmVycm9yICYmIG5hbWUgPT09ICdlcnJvcicpIHsKICAgICAgICAgICAgICAgICAgICBmbihtb2QuZXJyb3IpOwogICAgICAgICAgICAgICAgfSBlbHNlIHsKICAgICAgICAgICAgICAgICAgICBtb2Qub24obmFtZSwgZm4pOwogICAgICAgICAgICAgICAgfQogICAgICAgICAgICB9CiAgICAgICAgfQoKICAgICAgICBmdW5jdGlvbiBvbkVycm9yKGVyciwgZXJyYmFjaykgewogICAgICAgICAgICB2YXIgaWRzID0gZXJyLnJlcXVpcmVNb2R1bGVzLAogICAgICAgICAgICAgICAgbm90aWZpZWQgPSBmYWxzZTsKCiAgICAgICAgICAgIGlmIChlcnJiYWNrKSB7CiAgICAgICAgICAgICAgICBlcnJiYWNrKGVycik7CiAgICAgICAgICAgIH0gZWxzZSB7CiAgICAgICAgICAgICAgICBlYWNoKGlkcywgZnVuY3Rpb24gKGlkKSB7CiAgICAgICAgICAgICAgICAgICAgdmFyIG1vZCA9IGdldE93bihyZWdpc3RyeSwgaWQpOwogICAgICAgICAgICAgICAgICAgIGlmIChtb2QpIHsKICAgICAgICAgICAgICAgICAgICAgICAgLy9TZXQgZXJyb3Igb24gbW9kdWxlLCBzbyBpdCBza2lwcyB0aW1lb3V0IGNoZWNrcy4KICAgICAgICAgICAgICAgICAgICAgICAgbW9kLmVycm9yID0gZXJyOwogICAgICAgICAgICAgICAgICAgICAgICBpZiAobW9kLmV2ZW50cy5lcnJvcikgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgbm90aWZpZWQgPSB0cnVlOwogICAgICAgICAgICAgICAgICAgICAgICAgICAgbW9kLmVtaXQoJ2Vycm9yJywgZXJyKTsKICAgICAgICAgICAgICAgICAgICAgICAgfQogICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgIH0pOwoKICAgICAgICAgICAgICAgIGlmICghbm90aWZpZWQpIHsKICAgICAgICAgICAgICAgICAgICByZXEub25FcnJvcihlcnIpOwogICAgICAgICAgICAgICAgfQogICAgICAgICAgICB9CiAgICAgICAgfQoKICAgICAgICAvKioKICAgICAgICAgKiBJbnRlcm5hbCBtZXRob2QgdG8gdHJhbnNmZXIgZ2xvYmFsUXVldWUgaXRlbXMgdG8gdGhpcyBjb250ZXh0J3MKICAgICAgICAgKiBkZWZRdWV1ZS4KICAgICAgICAgKi8KICAgICAgICBmdW5jdGlvbiB0YWtlR2xvYmFsUXVldWUoKSB7CiAgICAgICAgICAgIC8vUHVzaCBhbGwgdGhlIGdsb2JhbERlZlF1ZXVlIGl0ZW1zIGludG8gdGhlIGNvbnRleHQncyBkZWZRdWV1ZQogICAgICAgICAgICBpZiAoZ2xvYmFsRGVmUXVldWUubGVuZ3RoKSB7CiAgICAgICAgICAgICAgICBlYWNoKGdsb2JhbERlZlF1ZXVlLCBmdW5jdGlvbihxdWV1ZUl0ZW0pIHsKICAgICAgICAgICAgICAgICAgICB2YXIgaWQgPSBxdWV1ZUl0ZW1bMF07CiAgICAgICAgICAgICAgICAgICAgaWYgKHR5cGVvZiBpZCA9PT0gJ3N0cmluZycpIHsKICAgICAgICAgICAgICAgICAgICAgICAgY29udGV4dC5kZWZRdWV1ZU1hcFtpZF0gPSB0cnVlOwogICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICBkZWZRdWV1ZS5wdXNoKHF1ZXVlSXRlbSk7CiAgICAgICAgICAgICAgICB9KTsKICAgICAgICAgICAgICAgIGdsb2JhbERlZlF1ZXVlID0gW107CiAgICAgICAgICAgIH0KICAgICAgICB9CgogICAgICAgIGhhbmRsZXJzID0gewogICAgICAgICAgICAncmVxdWlyZSc6IGZ1bmN0aW9uIChtb2QpIHsKICAgICAgICAgICAgICAgIGlmIChtb2QucmVxdWlyZSkgewogICAgICAgICAgICAgICAgICAgIHJldHVybiBtb2QucmVxdWlyZTsKICAgICAgICAgICAgICAgIH0gZWxzZSB7CiAgICAgICAgICAgICAgICAgICAgcmV0dXJuIChtb2QucmVxdWlyZSA9IGNvbnRleHQubWFrZVJlcXVpcmUobW9kLm1hcCkpOwogICAgICAgICAgICAgICAgfQogICAgICAgICAgICB9LAogICAgICAgICAgICAnZXhwb3J0cyc6IGZ1bmN0aW9uIChtb2QpIHsKICAgICAgICAgICAgICAgIG1vZC51c2luZ0V4cG9ydHMgPSB0cnVlOwogICAgICAgICAgICAgICAgaWYgKG1vZC5tYXAuaXNEZWZpbmUpIHsKICAgICAgICAgICAgICAgICAgICBpZiAobW9kLmV4cG9ydHMpIHsKICAgICAgICAgICAgICAgICAgICAgICAgcmV0dXJuIChkZWZpbmVkW21vZC5tYXAuaWRdID0gbW9kLmV4cG9ydHMpOwogICAgICAgICAgICAgICAgICAgIH0gZWxzZSB7CiAgICAgICAgICAgICAgICAgICAgICAgIHJldHVybiAobW9kLmV4cG9ydHMgPSBkZWZpbmVkW21vZC5tYXAuaWRdID0ge30pOwogICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgfSwKICAgICAgICAgICAgJ21vZHVsZSc6IGZ1bmN0aW9uIChtb2QpIHsKICAgICAgICAgICAgICAgIGlmIChtb2QubW9kdWxlKSB7CiAgICAgICAgICAgICAgICAgICAgcmV0dXJuIG1vZC5tb2R1bGU7CiAgICAgICAgICAgICAgICB9IGVsc2UgewogICAgICAgICAgICAgICAgICAgIHJldHVybiAobW9kLm1vZHVsZSA9IHsKICAgICAgICAgICAgICAgICAgICAgICAgaWQ6IG1vZC5tYXAuaWQsCiAgICAgICAgICAgICAgICAgICAgICAgIHVyaTogbW9kLm1hcC51cmwsCiAgICAgICAgICAgICAgICAgICAgICAgIGNvbmZpZzogZnVuY3Rpb24gKCkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgcmV0dXJuIGdldE93bihjb25maWcuY29uZmlnLCBtb2QubWFwLmlkKSB8fCB7fTsKICAgICAgICAgICAgICAgICAgICAgICAgfSwKICAgICAgICAgICAgICAgICAgICAgICAgZXhwb3J0czogbW9kLmV4cG9ydHMgfHwgKG1vZC5leHBvcnRzID0ge30pCiAgICAgICAgICAgICAgICAgICAgfSk7CiAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgIH0KICAgICAgICB9OwoKICAgICAgICBmdW5jdGlvbiBjbGVhblJlZ2lzdHJ5KGlkKSB7CiAgICAgICAgICAgIC8vQ2xlYW4gdXAgbWFjaGluZXJ5IHVzZWQgZm9yIHdhaXRpbmcgbW9kdWxlcy4KICAgICAgICAgICAgZGVsZXRlIHJlZ2lzdHJ5W2lkXTsKICAgICAgICAgICAgZGVsZXRlIGVuYWJsZWRSZWdpc3RyeVtpZF07CiAgICAgICAgfQoKICAgICAgICBmdW5jdGlvbiBicmVha0N5Y2xlKG1vZCwgdHJhY2VkLCBwcm9jZXNzZWQpIHsKICAgICAgICAgICAgdmFyIGlkID0gbW9kLm1hcC5pZDsKCiAgICAgICAgICAgIGlmIChtb2QuZXJyb3IpIHsKICAgICAgICAgICAgICAgIG1vZC5lbWl0KCdlcnJvcicsIG1vZC5lcnJvcik7CiAgICAgICAgICAgIH0gZWxzZSB7CiAgICAgICAgICAgICAgICB0cmFjZWRbaWRdID0gdHJ1ZTsKICAgICAgICAgICAgICAgIGVhY2gobW9kLmRlcE1hcHMsIGZ1bmN0aW9uIChkZXBNYXAsIGkpIHsKICAgICAgICAgICAgICAgICAgICB2YXIgZGVwSWQgPSBkZXBNYXAuaWQsCiAgICAgICAgICAgICAgICAgICAgICAgIGRlcCA9IGdldE93bihyZWdpc3RyeSwgZGVwSWQpOwoKICAgICAgICAgICAgICAgICAgICAvL09ubHkgZm9yY2UgdGhpbmdzIHRoYXQgaGF2ZSBub3QgY29tcGxldGVkCiAgICAgICAgICAgICAgICAgICAgLy9iZWluZyBkZWZpbmVkLCBzbyBzdGlsbCBpbiB0aGUgcmVnaXN0cnksCiAgICAgICAgICAgICAgICAgICAgLy9hbmQgb25seSBpZiBpdCBoYXMgbm90IGJlZW4gbWF0Y2hlZCB1cAogICAgICAgICAgICAgICAgICAgIC8vaW4gdGhlIG1vZHVsZSBhbHJlYWR5LgogICAgICAgICAgICAgICAgICAgIGlmIChkZXAgJiYgIW1vZC5kZXBNYXRjaGVkW2ldICYmICFwcm9jZXNzZWRbZGVwSWRdKSB7CiAgICAgICAgICAgICAgICAgICAgICAgIGlmIChnZXRPd24odHJhY2VkLCBkZXBJZCkpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIG1vZC5kZWZpbmVEZXAoaSwgZGVmaW5lZFtkZXBJZF0pOwogICAgICAgICAgICAgICAgICAgICAgICAgICAgbW9kLmNoZWNrKCk7IC8vcGFzcyBmYWxzZT8KICAgICAgICAgICAgICAgICAgICAgICAgfSBlbHNlIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIGJyZWFrQ3ljbGUoZGVwLCB0cmFjZWQsIHByb2Nlc3NlZCk7CiAgICAgICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgICAgICB9KTsKICAgICAgICAgICAgICAgIHByb2Nlc3NlZFtpZF0gPSB0cnVlOwogICAgICAgICAgICB9CiAgICAgICAgfQoKICAgICAgICBmdW5jdGlvbiBjaGVja0xvYWRlZCgpIHsKICAgICAgICAgICAgdmFyIGVyciwgdXNpbmdQYXRoRmFsbGJhY2ssCiAgICAgICAgICAgICAgICB3YWl0SW50ZXJ2YWwgPSBjb25maWcud2FpdFNlY29uZHMgKiAxMDAwLAogICAgICAgICAgICAgICAgLy9JdCBpcyBwb3NzaWJsZSB0byBkaXNhYmxlIHRoZSB3YWl0IGludGVydmFsIGJ5IHVzaW5nIHdhaXRTZWNvbmRzIG9mIDAuCiAgICAgICAgICAgICAgICBleHBpcmVkID0gd2FpdEludGVydmFsICYmIChjb250ZXh0LnN0YXJ0VGltZSArIHdhaXRJbnRlcnZhbCkgPCBuZXcgRGF0ZSgpLmdldFRpbWUoKSwKICAgICAgICAgICAgICAgIG5vTG9hZHMgPSBbXSwKICAgICAgICAgICAgICAgIHJlcUNhbGxzID0gW10sCiAgICAgICAgICAgICAgICBzdGlsbExvYWRpbmcgPSBmYWxzZSwKICAgICAgICAgICAgICAgIG5lZWRDeWNsZUNoZWNrID0gdHJ1ZTsKCiAgICAgICAgICAgIC8vRG8gbm90IGJvdGhlciBpZiB0aGlzIGNhbGwgd2FzIGEgcmVzdWx0IG9mIGEgY3ljbGUgYnJlYWsuCiAgICAgICAgICAgIGlmIChpbkNoZWNrTG9hZGVkKSB7CiAgICAgICAgICAgICAgICByZXR1cm47CiAgICAgICAgICAgIH0KCiAgICAgICAgICAgIGluQ2hlY2tMb2FkZWQgPSB0cnVlOwoKICAgICAgICAgICAgLy9GaWd1cmUgb3V0IHRoZSBzdGF0ZSBvZiBhbGwgdGhlIG1vZHVsZXMuCiAgICAgICAgICAgIGVhY2hQcm9wKGVuYWJsZWRSZWdpc3RyeSwgZnVuY3Rpb24gKG1vZCkgewogICAgICAgICAgICAgICAgdmFyIG1hcCA9IG1vZC5tYXAsCiAgICAgICAgICAgICAgICAgICAgbW9kSWQgPSBtYXAuaWQ7CgogICAgICAgICAgICAgICAgLy9Ta2lwIHRoaW5ncyB0aGF0IGFyZSBub3QgZW5hYmxlZCBvciBpbiBlcnJvciBzdGF0ZS4KICAgICAgICAgICAgICAgIGlmICghbW9kLmVuYWJsZWQpIHsKICAgICAgICAgICAgICAgICAgICByZXR1cm47CiAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgaWYgKCFtYXAuaXNEZWZpbmUpIHsKICAgICAgICAgICAgICAgICAgICByZXFDYWxscy5wdXNoKG1vZCk7CiAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgaWYgKCFtb2QuZXJyb3IpIHsKICAgICAgICAgICAgICAgICAgICAvL0lmIHRoZSBtb2R1bGUgc2hvdWxkIGJlIGV4ZWN1dGVkLCBhbmQgaXQgaGFzIG5vdAogICAgICAgICAgICAgICAgICAgIC8vYmVlbiBpbml0ZWQgYW5kIHRpbWUgaXMgdXAsIHJlbWVtYmVyIGl0LgogICAgICAgICAgICAgICAgICAgIGlmICghbW9kLmluaXRlZCAmJiBleHBpcmVkKSB7CiAgICAgICAgICAgICAgICAgICAgICAgIGlmIChoYXNQYXRoRmFsbGJhY2sobW9kSWQpKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICB1c2luZ1BhdGhGYWxsYmFjayA9IHRydWU7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICBzdGlsbExvYWRpbmcgPSB0cnVlOwogICAgICAgICAgICAgICAgICAgICAgICB9IGVsc2UgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgbm9Mb2Fkcy5wdXNoKG1vZElkKTsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIHJlbW92ZVNjcmlwdChtb2RJZCk7CiAgICAgICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICB9IGVsc2UgaWYgKCFtb2QuaW5pdGVkICYmIG1vZC5mZXRjaGVkICYmIG1hcC5pc0RlZmluZSkgewogICAgICAgICAgICAgICAgICAgICAgICBzdGlsbExvYWRpbmcgPSB0cnVlOwogICAgICAgICAgICAgICAgICAgICAgICBpZiAoIW1hcC5wcmVmaXgpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIC8vTm8gcmVhc29uIHRvIGtlZXAgbG9va2luZyBmb3IgdW5maW5pc2hlZAogICAgICAgICAgICAgICAgICAgICAgICAgICAgLy9sb2FkaW5nLiBJZiB0aGUgb25seSBzdGlsbExvYWRpbmcgaXMgYQogICAgICAgICAgICAgICAgICAgICAgICAgICAgLy9wbHVnaW4gcmVzb3VyY2UgdGhvdWdoLCBrZWVwIGdvaW5nLAogICAgICAgICAgICAgICAgICAgICAgICAgICAgLy9iZWNhdXNlIGl0IG1heSBiZSB0aGF0IGEgcGx1Z2luIHJlc291cmNlCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAvL2lzIHdhaXRpbmcgb24gYSBub24tcGx1Z2luIGN5Y2xlLgogICAgICAgICAgICAgICAgICAgICAgICAgICAgcmV0dXJuIChuZWVkQ3ljbGVDaGVjayA9IGZhbHNlKTsKICAgICAgICAgICAgICAgICAgICAgICAgfQogICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgfSk7CgogICAgICAgICAgICBpZiAoZXhwaXJlZCAmJiBub0xvYWRzLmxlbmd0aCkgewogICAgICAgICAgICAgICAgLy9JZiB3YWl0IHRpbWUgZXhwaXJlZCwgdGhyb3cgZXJyb3Igb2YgdW5sb2FkZWQgbW9kdWxlcy4KICAgICAgICAgICAgICAgIGVyciA9IG1ha2VFcnJvcigndGltZW91dCcsICdMb2FkIHRpbWVvdXQgZm9yIG1vZHVsZXM6ICcgKyBub0xvYWRzLCBudWxsLCBub0xvYWRzKTsKICAgICAgICAgICAgICAgIGVyci5jb250ZXh0TmFtZSA9IGNvbnRleHQuY29udGV4dE5hbWU7CiAgICAgICAgICAgICAgICByZXR1cm4gb25FcnJvcihlcnIpOwogICAgICAgICAgICB9CgogICAgICAgICAgICAvL05vdCBleHBpcmVkLCBjaGVjayBmb3IgYSBjeWNsZS4KICAgICAgICAgICAgaWYgKG5lZWRDeWNsZUNoZWNrKSB7CiAgICAgICAgICAgICAgICBlYWNoKHJlcUNhbGxzLCBmdW5jdGlvbiAobW9kKSB7CiAgICAgICAgICAgICAgICAgICAgYnJlYWtDeWNsZShtb2QsIHt9LCB7fSk7CiAgICAgICAgICAgICAgICB9KTsKICAgICAgICAgICAgfQoKICAgICAgICAgICAgLy9JZiBzdGlsbCB3YWl0aW5nIG9uIGxvYWRzLCBhbmQgdGhlIHdhaXRpbmcgbG9hZCBpcyBzb21ldGhpbmcKICAgICAgICAgICAgLy9vdGhlciB0aGFuIGEgcGx1Z2luIHJlc291cmNlLCBvciB0aGVyZSBhcmUgc3RpbGwgb3V0c3RhbmRpbmcKICAgICAgICAgICAgLy9zY3JpcHRzLCB0aGVuIGp1c3QgdHJ5IGJhY2sgbGF0ZXIuCiAgICAgICAgICAgIGlmICgoIWV4cGlyZWQgfHwgdXNpbmdQYXRoRmFsbGJhY2spICYmIHN0aWxsTG9hZGluZykgewogICAgICAgICAgICAgICAgLy9Tb21ldGhpbmcgaXMgc3RpbGwgd2FpdGluZyB0byBsb2FkLiBXYWl0IGZvciBpdCwgYnV0IG9ubHkKICAgICAgICAgICAgICAgIC8vaWYgYSB0aW1lb3V0IGlzIG5vdCBhbHJlYWR5IGluIGVmZmVjdC4KICAgICAgICAgICAgICAgIGlmICgoaXNCcm93c2VyIHx8IGlzV2ViV29ya2VyKSAmJiAhY2hlY2tMb2FkZWRUaW1lb3V0SWQpIHsKICAgICAgICAgICAgICAgICAgICBjaGVja0xvYWRlZFRpbWVvdXRJZCA9IHNldFRpbWVvdXQoZnVuY3Rpb24gKCkgewogICAgICAgICAgICAgICAgICAgICAgICBjaGVja0xvYWRlZFRpbWVvdXRJZCA9IDA7CiAgICAgICAgICAgICAgICAgICAgICAgIGNoZWNrTG9hZGVkKCk7CiAgICAgICAgICAgICAgICAgICAgfSwgNTApOwogICAgICAgICAgICAgICAgfQogICAgICAgICAgICB9CgogICAgICAgICAgICBpbkNoZWNrTG9hZGVkID0gZmFsc2U7CiAgICAgICAgfQoKICAgICAgICBNb2R1bGUgPSBmdW5jdGlvbiAobWFwKSB7CiAgICAgICAgICAgIHRoaXMuZXZlbnRzID0gZ2V0T3duKHVuZGVmRXZlbnRzLCBtYXAuaWQpIHx8IHt9OwogICAgICAgICAgICB0aGlzLm1hcCA9IG1hcDsKICAgICAgICAgICAgdGhpcy5zaGltID0gZ2V0T3duKGNvbmZpZy5zaGltLCBtYXAuaWQpOwogICAgICAgICAgICB0aGlzLmRlcEV4cG9ydHMgPSBbXTsKICAgICAgICAgICAgdGhpcy5kZXBNYXBzID0gW107CiAgICAgICAgICAgIHRoaXMuZGVwTWF0Y2hlZCA9IFtdOwogICAgICAgICAgICB0aGlzLnBsdWdpbk1hcHMgPSB7fTsKICAgICAgICAgICAgdGhpcy5kZXBDb3VudCA9IDA7CgogICAgICAgICAgICAvKiB0aGlzLmV4cG9ydHMgdGhpcy5mYWN0b3J5CiAgICAgICAgICAgICAgIHRoaXMuZGVwTWFwcyA9IFtdLAogICAgICAgICAgICAgICB0aGlzLmVuYWJsZWQsIHRoaXMuZmV0Y2hlZAogICAgICAgICAgICAqLwogICAgICAgIH07CgogICAgICAgIE1vZHVsZS5wcm90b3R5cGUgPSB7CiAgICAgICAgICAgIGluaXQ6IGZ1bmN0aW9uIChkZXBNYXBzLCBmYWN0b3J5LCBlcnJiYWNrLCBvcHRpb25zKSB7CiAgICAgICAgICAgICAgICBvcHRpb25zID0gb3B0aW9ucyB8fCB7fTsKCiAgICAgICAgICAgICAgICAvL0RvIG5vdCBkbyBtb3JlIGluaXRzIGlmIGFscmVhZHkgZG9uZS4gQ2FuIGhhcHBlbiBpZiB0aGVyZQogICAgICAgICAgICAgICAgLy9hcmUgbXVsdGlwbGUgZGVmaW5lIGNhbGxzIGZvciB0aGUgc2FtZSBtb2R1bGUuIFRoYXQgaXMgbm90CiAgICAgICAgICAgICAgICAvL2Egbm9ybWFsLCBjb21tb24gY2FzZSwgYnV0IGl0IGlzIGFsc28gbm90IHVuZXhwZWN0ZWQuCiAgICAgICAgICAgICAgICBpZiAodGhpcy5pbml0ZWQpIHsKICAgICAgICAgICAgICAgICAgICByZXR1cm47CiAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgdGhpcy5mYWN0b3J5ID0gZmFjdG9yeTsKCiAgICAgICAgICAgICAgICBpZiAoZXJyYmFjaykgewogICAgICAgICAgICAgICAgICAgIC8vUmVnaXN0ZXIgZm9yIGVycm9ycyBvbiB0aGlzIG1vZHVsZS4KICAgICAgICAgICAgICAgICAgICB0aGlzLm9uKCdlcnJvcicsIGVycmJhY2spOwogICAgICAgICAgICAgICAgfSBlbHNlIGlmICh0aGlzLmV2ZW50cy5lcnJvcikgewogICAgICAgICAgICAgICAgICAgIC8vSWYgbm8gZXJyYmFjayBhbHJlYWR5LCBidXQgdGhlcmUgYXJlIGVycm9yIGxpc3RlbmVycwogICAgICAgICAgICAgICAgICAgIC8vb24gdGhpcyBtb2R1bGUsIHNldCB1cCBhbiBlcnJiYWNrIHRvIHBhc3MgdG8gdGhlIGRlcHMuCiAgICAgICAgICAgICAgICAgICAgZXJyYmFjayA9IGJpbmQodGhpcywgZnVuY3Rpb24gKGVycikgewogICAgICAgICAgICAgICAgICAgICAgICB0aGlzLmVtaXQoJ2Vycm9yJywgZXJyKTsKICAgICAgICAgICAgICAgICAgICB9KTsKICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICAvL0RvIGEgY29weSBvZiB0aGUgZGVwZW5kZW5jeSBhcnJheSwgc28gdGhhdAogICAgICAgICAgICAgICAgLy9zb3VyY2UgaW5wdXRzIGFyZSBub3QgbW9kaWZpZWQuIEZvciBleGFtcGxlCiAgICAgICAgICAgICAgICAvLyJzaGltIiBkZXBzIGFyZSBwYXNzZWQgaW4gaGVyZSBkaXJlY3RseSwgYW5kCiAgICAgICAgICAgICAgICAvL2RvaW5nIGEgZGlyZWN0IG1vZGlmaWNhdGlvbiBvZiB0aGUgZGVwTWFwcyBhcnJheQogICAgICAgICAgICAgICAgLy93b3VsZCBhZmZlY3QgdGhhdCBjb25maWcuCiAgICAgICAgICAgICAgICB0aGlzLmRlcE1hcHMgPSBkZXBNYXBzICYmIGRlcE1hcHMuc2xpY2UoMCk7CgogICAgICAgICAgICAgICAgdGhpcy5lcnJiYWNrID0gZXJyYmFjazsKCiAgICAgICAgICAgICAgICAvL0luZGljYXRlIHRoaXMgbW9kdWxlIGhhcyBiZSBpbml0aWFsaXplZAogICAgICAgICAgICAgICAgdGhpcy5pbml0ZWQgPSB0cnVlOwoKICAgICAgICAgICAgICAgIHRoaXMuaWdub3JlID0gb3B0aW9ucy5pZ25vcmU7CgogICAgICAgICAgICAgICAgLy9Db3VsZCBoYXZlIG9wdGlvbiB0byBpbml0IHRoaXMgbW9kdWxlIGluIGVuYWJsZWQgbW9kZSwKICAgICAgICAgICAgICAgIC8vb3IgY291bGQgaGF2ZSBiZWVuIHByZXZpb3VzbHkgbWFya2VkIGFzIGVuYWJsZWQuIEhvd2V2ZXIsCiAgICAgICAgICAgICAgICAvL3RoZSBkZXBlbmRlbmNpZXMgYXJlIG5vdCBrbm93biB1bnRpbCBpbml0IGlzIGNhbGxlZC4gU28KICAgICAgICAgICAgICAgIC8vaWYgZW5hYmxlZCBwcmV2aW91c2x5LCBub3cgdHJpZ2dlciBkZXBlbmRlbmNpZXMgYXMgZW5hYmxlZC4KICAgICAgICAgICAgICAgIGlmIChvcHRpb25zLmVuYWJsZWQgfHwgdGhpcy5lbmFibGVkKSB7CiAgICAgICAgICAgICAgICAgICAgLy9FbmFibGUgdGhpcyBtb2R1bGUgYW5kIGRlcGVuZGVuY2llcy4KICAgICAgICAgICAgICAgICAgICAvL1dpbGwgY2FsbCB0aGlzLmNoZWNrKCkKICAgICAgICAgICAgICAgICAgICB0aGlzLmVuYWJsZSgpOwogICAgICAgICAgICAgICAgfSBlbHNlIHsKICAgICAgICAgICAgICAgICAgICB0aGlzLmNoZWNrKCk7CiAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgIH0sCgogICAgICAgICAgICBkZWZpbmVEZXA6IGZ1bmN0aW9uIChpLCBkZXBFeHBvcnRzKSB7CiAgICAgICAgICAgICAgICAvL0JlY2F1c2Ugb2YgY3ljbGVzLCBkZWZpbmVkIGNhbGxiYWNrIGZvciBhIGdpdmVuCiAgICAgICAgICAgICAgICAvL2V4cG9ydCBjYW4gYmUgY2FsbGVkIG1vcmUgdGhhbiBvbmNlLgogICAgICAgICAgICAgICAgaWYgKCF0aGlzLmRlcE1hdGNoZWRbaV0pIHsKICAgICAgICAgICAgICAgICAgICB0aGlzLmRlcE1hdGNoZWRbaV0gPSB0cnVlOwogICAgICAgICAgICAgICAgICAgIHRoaXMuZGVwQ291bnQgLT0gMTsKICAgICAgICAgICAgICAgICAgICB0aGlzLmRlcEV4cG9ydHNbaV0gPSBkZXBFeHBvcnRzOwogICAgICAgICAgICAgICAgfQogICAgICAgICAgICB9LAoKICAgICAgICAgICAgZmV0Y2g6IGZ1bmN0aW9uICgpIHsKICAgICAgICAgICAgICAgIGlmICh0aGlzLmZldGNoZWQpIHsKICAgICAgICAgICAgICAgICAgICByZXR1cm47CiAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgICAgICB0aGlzLmZldGNoZWQgPSB0cnVlOwoKICAgICAgICAgICAgICAgIGNvbnRleHQuc3RhcnRUaW1lID0gKG5ldyBEYXRlKCkpLmdldFRpbWUoKTsKCiAgICAgICAgICAgICAgICB2YXIgbWFwID0gdGhpcy5tYXA7CgogICAgICAgICAgICAgICAgLy9JZiB0aGUgbWFuYWdlciBpcyBmb3IgYSBwbHVnaW4gbWFuYWdlZCByZXNvdXJjZSwKICAgICAgICAgICAgICAgIC8vYXNrIHRoZSBwbHVnaW4gdG8gbG9hZCBpdCBub3cuCiAgICAgICAgICAgICAgICBpZiAodGhpcy5zaGltKSB7CiAgICAgICAgICAgICAgICAgICAgY29udGV4dC5tYWtlUmVxdWlyZSh0aGlzLm1hcCwgewogICAgICAgICAgICAgICAgICAgICAgICBlbmFibGVCdWlsZENhbGxiYWNrOiB0cnVlCiAgICAgICAgICAgICAgICAgICAgfSkodGhpcy5zaGltLmRlcHMgfHwgW10sIGJpbmQodGhpcywgZnVuY3Rpb24gKCkgewogICAgICAgICAgICAgICAgICAgICAgICByZXR1cm4gbWFwLnByZWZpeCA/IHRoaXMuY2FsbFBsdWdpbigpIDogdGhpcy5sb2FkKCk7CiAgICAgICAgICAgICAgICAgICAgfSkpOwogICAgICAgICAgICAgICAgfSBlbHNlIHsKICAgICAgICAgICAgICAgICAgICAvL1JlZ3VsYXIgZGVwZW5kZW5jeS4KICAgICAgICAgICAgICAgICAgICByZXR1cm4gbWFwLnByZWZpeCA/IHRoaXMuY2FsbFBsdWdpbigpIDogdGhpcy5sb2FkKCk7CiAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgIH0sCgogICAgICAgICAgICBsb2FkOiBmdW5jdGlvbiAoKSB7CiAgICAgICAgICAgICAgICB2YXIgdXJsID0gdGhpcy5tYXAudXJsOwoKICAgICAgICAgICAgICAgIC8vUmVndWxhciBkZXBlbmRlbmN5LgogICAgICAgICAgICAgICAgaWYgKCF1cmxGZXRjaGVkW3VybF0pIHsKICAgICAgICAgICAgICAgICAgICB1cmxGZXRjaGVkW3VybF0gPSB0cnVlOwogICAgICAgICAgICAgICAgICAgIGNvbnRleHQubG9hZCh0aGlzLm1hcC5pZCwgdXJsKTsKICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgfSwKCiAgICAgICAgICAgIC8qKgogICAgICAgICAgICAgKiBDaGVja3MgaWYgdGhlIG1vZHVsZSBpcyByZWFkeSB0byBkZWZpbmUgaXRzZWxmLCBhbmQgaWYgc28sCiAgICAgICAgICAgICAqIGRlZmluZSBpdC4KICAgICAgICAgICAgICovCiAgICAgICAgICAgIGNoZWNrOiBmdW5jdGlvbiAoKSB7CiAgICAgICAgICAgICAgICBpZiAoIXRoaXMuZW5hYmxlZCB8fCB0aGlzLmVuYWJsaW5nKSB7CiAgICAgICAgICAgICAgICAgICAgcmV0dXJuOwogICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgIHZhciBlcnIsIGNqc01vZHVsZSwKICAgICAgICAgICAgICAgICAgICBpZCA9IHRoaXMubWFwLmlkLAogICAgICAgICAgICAgICAgICAgIGRlcEV4cG9ydHMgPSB0aGlzLmRlcEV4cG9ydHMsCiAgICAgICAgICAgICAgICAgICAgZXhwb3J0cyA9IHRoaXMuZXhwb3J0cywKICAgICAgICAgICAgICAgICAgICBmYWN0b3J5ID0gdGhpcy5mYWN0b3J5OwoKICAgICAgICAgICAgICAgIGlmICghdGhpcy5pbml0ZWQpIHsKICAgICAgICAgICAgICAgICAgICAvLyBPbmx5IGZldGNoIGlmIG5vdCBhbHJlYWR5IGluIHRoZSBkZWZRdWV1ZS4KICAgICAgICAgICAgICAgICAgICBpZiAoIWhhc1Byb3AoY29udGV4dC5kZWZRdWV1ZU1hcCwgaWQpKSB7CiAgICAgICAgICAgICAgICAgICAgICAgIHRoaXMuZmV0Y2goKTsKICAgICAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgICAgICB9IGVsc2UgaWYgKHRoaXMuZXJyb3IpIHsKICAgICAgICAgICAgICAgICAgICB0aGlzLmVtaXQoJ2Vycm9yJywgdGhpcy5lcnJvcik7CiAgICAgICAgICAgICAgICB9IGVsc2UgaWYgKCF0aGlzLmRlZmluaW5nKSB7CiAgICAgICAgICAgICAgICAgICAgLy9UaGUgZmFjdG9yeSBjb3VsZCB0cmlnZ2VyIGFub3RoZXIgcmVxdWlyZSBjYWxsCiAgICAgICAgICAgICAgICAgICAgLy90aGF0IHdvdWxkIHJlc3VsdCBpbiBjaGVja2luZyB0aGlzIG1vZHVsZSB0bwogICAgICAgICAgICAgICAgICAgIC8vZGVmaW5lIGl0c2VsZiBhZ2Fpbi4gSWYgYWxyZWFkeSBpbiB0aGUgcHJvY2VzcwogICAgICAgICAgICAgICAgICAgIC8vb2YgZG9pbmcgdGhhdCwgc2tpcCB0aGlzIHdvcmsuCiAgICAgICAgICAgICAgICAgICAgdGhpcy5kZWZpbmluZyA9IHRydWU7CgogICAgICAgICAgICAgICAgICAgIGlmICh0aGlzLmRlcENvdW50IDwgMSAmJiAhdGhpcy5kZWZpbmVkKSB7CiAgICAgICAgICAgICAgICAgICAgICAgIGlmIChpc0Z1bmN0aW9uKGZhY3RvcnkpKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICB0cnkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIGV4cG9ydHMgPSBjb250ZXh0LmV4ZWNDYihpZCwgZmFjdG9yeSwgZGVwRXhwb3J0cywgZXhwb3J0cyk7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICB9IGNhdGNoIChlKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgZXJyID0gZTsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAvLyBGYXZvciByZXR1cm4gdmFsdWUgb3ZlciBleHBvcnRzLiBJZiBub2RlL2NqcyBpbiBwbGF5LAogICAgICAgICAgICAgICAgICAgICAgICAgICAgLy8gdGhlbiB3aWxsIG5vdCBoYXZlIGEgcmV0dXJuIHZhbHVlIGFueXdheS4gRmF2b3IKICAgICAgICAgICAgICAgICAgICAgICAgICAgIC8vIG1vZHVsZS5leHBvcnRzIGFzc2lnbm1lbnQgb3ZlciBleHBvcnRzIG9iamVjdC4KICAgICAgICAgICAgICAgICAgICAgICAgICAgIGlmICh0aGlzLm1hcC5pc0RlZmluZSAmJiBleHBvcnRzID09PSB1bmRlZmluZWQpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBjanNNb2R1bGUgPSB0aGlzLm1vZHVsZTsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBpZiAoY2pzTW9kdWxlKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIGV4cG9ydHMgPSBjanNNb2R1bGUuZXhwb3J0czsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICB9IGVsc2UgaWYgKHRoaXMudXNpbmdFeHBvcnRzKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIC8vZXhwb3J0cyBhbHJlYWR5IHNldCB0aGUgZGVmaW5lZCB2YWx1ZS4KICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgZXhwb3J0cyA9IHRoaXMuZXhwb3J0czsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgICAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgICAgICAgICAgICAgaWYgKGVycikgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIC8vIElmIHRoZXJlIGlzIGFuIGVycm9yIGxpc3RlbmVyLCBmYXZvciBwYXNzaW5nCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgLy8gdG8gdGhhdCBpbnN0ZWFkIG9mIHRocm93aW5nIGFuIGVycm9yLiBIb3dldmVyLAogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIC8vIG9ubHkgZG8gaXQgZm9yIGRlZmluZSgpJ2QgIG1vZHVsZXMuIHJlcXVpcmUKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAvLyBlcnJiYWNrcyBzaG91bGQgbm90IGJlIGNhbGxlZCBmb3IgZmFpbHVyZXMgaW4KICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAvLyB0aGVpciBjYWxsYmFja3MgKCM2OTkpLiBIb3dldmVyIGlmIGEgZ2xvYmFsCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgLy8gb25FcnJvciBpcyBzZXQsIHVzZSB0aGF0LgogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIGlmICgodGhpcy5ldmVudHMuZXJyb3IgJiYgdGhpcy5tYXAuaXNEZWZpbmUpIHx8CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIHJlcS5vbkVycm9yICE9PSBkZWZhdWx0T25FcnJvcikgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBlcnIucmVxdWlyZU1hcCA9IHRoaXMubWFwOwogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBlcnIucmVxdWlyZU1vZHVsZXMgPSB0aGlzLm1hcC5pc0RlZmluZSA/IFt0aGlzLm1hcC5pZF0gOiBudWxsOwogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBlcnIucmVxdWlyZVR5cGUgPSB0aGlzLm1hcC5pc0RlZmluZSA/ICdkZWZpbmUnIDogJ3JlcXVpcmUnOwogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICByZXR1cm4gb25FcnJvcigodGhpcy5lcnJvciA9IGVycikpOwogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIH0gZWxzZSBpZiAodHlwZW9mIGNvbnNvbGUgIT09ICd1bmRlZmluZWQnICYmCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBjb25zb2xlLmVycm9yKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIC8vIExvZyB0aGUgZXJyb3IgZm9yIGRlYnVnZ2luZy4gSWYgcHJvbWlzZXMgY291bGQgYmUKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgLy8gdXNlZCwgdGhpcyB3b3VsZCBiZSBkaWZmZXJlbnQsIGJ1dCBtYWtpbmcgZG8uCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIGNvbnNvbGUuZXJyb3IoZXJyKTsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICB9IGVsc2UgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAvLyBEbyBub3Qgd2FudCB0byBjb21wbGV0ZWx5IGxvc2UgdGhlIGVycm9yLiBXaGlsZSB0aGlzCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIC8vIHdpbGwgbWVzcyB1cCBwcm9jZXNzaW5nIGFuZCBsZWFkIHRvIHNpbWlsYXIgcmVzdWx0cwogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAvLyBhcyBidWcgMTQ0MCwgaXQgYXQgbGVhc3Qgc3VyZmFjZXMgdGhlIGVycm9yLgogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICByZXEub25FcnJvcihlcnIpOwogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICAgICAgfSBlbHNlIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIC8vSnVzdCBhIGxpdGVyYWwgdmFsdWUKICAgICAgICAgICAgICAgICAgICAgICAgICAgIGV4cG9ydHMgPSBmYWN0b3J5OwogICAgICAgICAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgICAgICAgICB0aGlzLmV4cG9ydHMgPSBleHBvcnRzOwoKICAgICAgICAgICAgICAgICAgICAgICAgaWYgKHRoaXMubWFwLmlzRGVmaW5lICYmICF0aGlzLmlnbm9yZSkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgZGVmaW5lZFtpZF0gPSBleHBvcnRzOwoKICAgICAgICAgICAgICAgICAgICAgICAgICAgIGlmIChyZXEub25SZXNvdXJjZUxvYWQpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICB2YXIgcmVzTG9hZE1hcHMgPSBbXTsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBlYWNoKHRoaXMuZGVwTWFwcywgZnVuY3Rpb24gKGRlcE1hcCkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICByZXNMb2FkTWFwcy5wdXNoKGRlcE1hcC5ub3JtYWxpemVkTWFwIHx8IGRlcE1hcCk7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgfSk7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgcmVxLm9uUmVzb3VyY2VMb2FkKGNvbnRleHQsIHRoaXMubWFwLCByZXNMb2FkTWFwcyk7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICAgICAgICAgIC8vQ2xlYW4gdXAKICAgICAgICAgICAgICAgICAgICAgICAgY2xlYW5SZWdpc3RyeShpZCk7CgogICAgICAgICAgICAgICAgICAgICAgICB0aGlzLmRlZmluZWQgPSB0cnVlOwogICAgICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICAgICAgLy9GaW5pc2hlZCB0aGUgZGVmaW5lIHN0YWdlLiBBbGxvdyBjYWxsaW5nIGNoZWNrIGFnYWluCiAgICAgICAgICAgICAgICAgICAgLy90byBhbGxvdyBkZWZpbmUgbm90aWZpY2F0aW9ucyBiZWxvdyBpbiB0aGUgY2FzZSBvZiBhCiAgICAgICAgICAgICAgICAgICAgLy9jeWNsZS4KICAgICAgICAgICAgICAgICAgICB0aGlzLmRlZmluaW5nID0gZmFsc2U7CgogICAgICAgICAgICAgICAgICAgIGlmICh0aGlzLmRlZmluZWQgJiYgIXRoaXMuZGVmaW5lRW1pdHRlZCkgewogICAgICAgICAgICAgICAgICAgICAgICB0aGlzLmRlZmluZUVtaXR0ZWQgPSB0cnVlOwogICAgICAgICAgICAgICAgICAgICAgICB0aGlzLmVtaXQoJ2RlZmluZWQnLCB0aGlzLmV4cG9ydHMpOwogICAgICAgICAgICAgICAgICAgICAgICB0aGlzLmRlZmluZUVtaXRDb21wbGV0ZSA9IHRydWU7CiAgICAgICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgfSwKCiAgICAgICAgICAgIGNhbGxQbHVnaW46IGZ1bmN0aW9uICgpIHsKICAgICAgICAgICAgICAgIHZhciBtYXAgPSB0aGlzLm1hcCwKICAgICAgICAgICAgICAgICAgICBpZCA9IG1hcC5pZCwKICAgICAgICAgICAgICAgICAgICAvL01hcCBhbHJlYWR5IG5vcm1hbGl6ZWQgdGhlIHByZWZpeC4KICAgICAgICAgICAgICAgICAgICBwbHVnaW5NYXAgPSBtYWtlTW9kdWxlTWFwKG1hcC5wcmVmaXgpOwoKICAgICAgICAgICAgICAgIC8vTWFyayB0aGlzIGFzIGEgZGVwZW5kZW5jeSBmb3IgdGhpcyBwbHVnaW4sIHNvIGl0CiAgICAgICAgICAgICAgICAvL2NhbiBiZSB0cmFjZWQgZm9yIGN5Y2xlcy4KICAgICAgICAgICAgICAgIHRoaXMuZGVwTWFwcy5wdXNoKHBsdWdpbk1hcCk7CgogICAgICAgICAgICAgICAgb24ocGx1Z2luTWFwLCAnZGVmaW5lZCcsIGJpbmQodGhpcywgZnVuY3Rpb24gKHBsdWdpbikgewogICAgICAgICAgICAgICAgICAgIHZhciBsb2FkLCBub3JtYWxpemVkTWFwLCBub3JtYWxpemVkTW9kLAogICAgICAgICAgICAgICAgICAgICAgICBidW5kbGVJZCA9IGdldE93bihidW5kbGVzTWFwLCB0aGlzLm1hcC5pZCksCiAgICAgICAgICAgICAgICAgICAgICAgIG5hbWUgPSB0aGlzLm1hcC5uYW1lLAogICAgICAgICAgICAgICAgICAgICAgICBwYXJlbnROYW1lID0gdGhpcy5tYXAucGFyZW50TWFwID8gdGhpcy5tYXAucGFyZW50TWFwLm5hbWUgOiBudWxsLAogICAgICAgICAgICAgICAgICAgICAgICBsb2NhbFJlcXVpcmUgPSBjb250ZXh0Lm1ha2VSZXF1aXJlKG1hcC5wYXJlbnRNYXAsIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIGVuYWJsZUJ1aWxkQ2FsbGJhY2s6IHRydWUKICAgICAgICAgICAgICAgICAgICAgICAgfSk7CgogICAgICAgICAgICAgICAgICAgIC8vSWYgY3VycmVudCBtYXAgaXMgbm90IG5vcm1hbGl6ZWQsIHdhaXQgZm9yIHRoYXQKICAgICAgICAgICAgICAgICAgICAvL25vcm1hbGl6ZWQgbmFtZSB0byBsb2FkIGluc3RlYWQgb2YgY29udGludWluZy4KICAgICAgICAgICAgICAgICAgICBpZiAodGhpcy5tYXAudW5ub3JtYWxpemVkKSB7CiAgICAgICAgICAgICAgICAgICAgICAgIC8vTm9ybWFsaXplIHRoZSBJRCBpZiB0aGUgcGx1Z2luIGFsbG93cyBpdC4KICAgICAgICAgICAgICAgICAgICAgICAgaWYgKHBsdWdpbi5ub3JtYWxpemUpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIG5hbWUgPSBwbHVnaW4ubm9ybWFsaXplKG5hbWUsIGZ1bmN0aW9uIChuYW1lKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgcmV0dXJuIG5vcm1hbGl6ZShuYW1lLCBwYXJlbnROYW1lLCB0cnVlKTsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIH0pIHx8ICcnOwogICAgICAgICAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgICAgICAgICAvL3ByZWZpeCBhbmQgbmFtZSBzaG91bGQgYWxyZWFkeSBiZSBub3JtYWxpemVkLCBubyBuZWVkCiAgICAgICAgICAgICAgICAgICAgICAgIC8vZm9yIGFwcGx5aW5nIG1hcCBjb25maWcgYWdhaW4gZWl0aGVyLgogICAgICAgICAgICAgICAgICAgICAgICBub3JtYWxpemVkTWFwID0gbWFrZU1vZHVsZU1hcChtYXAucHJlZml4ICsgJyEnICsgbmFtZSwKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgdGhpcy5tYXAucGFyZW50TWFwKTsKICAgICAgICAgICAgICAgICAgICAgICAgb24obm9ybWFsaXplZE1hcCwKICAgICAgICAgICAgICAgICAgICAgICAgICAgICdkZWZpbmVkJywgYmluZCh0aGlzLCBmdW5jdGlvbiAodmFsdWUpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICB0aGlzLm1hcC5ub3JtYWxpemVkTWFwID0gbm9ybWFsaXplZE1hcDsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICB0aGlzLmluaXQoW10sIGZ1bmN0aW9uICgpIHsgcmV0dXJuIHZhbHVlOyB9LCBudWxsLCB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIGVuYWJsZWQ6IHRydWUsCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIGlnbm9yZTogdHJ1ZQogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIH0pOwogICAgICAgICAgICAgICAgICAgICAgICAgICAgfSkpOwoKICAgICAgICAgICAgICAgICAgICAgICAgbm9ybWFsaXplZE1vZCA9IGdldE93bihyZWdpc3RyeSwgbm9ybWFsaXplZE1hcC5pZCk7CiAgICAgICAgICAgICAgICAgICAgICAgIGlmIChub3JtYWxpemVkTW9kKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAvL01hcmsgdGhpcyBhcyBhIGRlcGVuZGVuY3kgZm9yIHRoaXMgcGx1Z2luLCBzbyBpdAogICAgICAgICAgICAgICAgICAgICAgICAgICAgLy9jYW4gYmUgdHJhY2VkIGZvciBjeWNsZXMuCiAgICAgICAgICAgICAgICAgICAgICAgICAgICB0aGlzLmRlcE1hcHMucHVzaChub3JtYWxpemVkTWFwKTsKCiAgICAgICAgICAgICAgICAgICAgICAgICAgICBpZiAodGhpcy5ldmVudHMuZXJyb3IpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBub3JtYWxpemVkTW9kLm9uKCdlcnJvcicsIGJpbmQodGhpcywgZnVuY3Rpb24gKGVycikgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICB0aGlzLmVtaXQoJ2Vycm9yJywgZXJyKTsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICB9KSk7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgICAgICAgICAgICAgICAgICBub3JtYWxpemVkTW9kLmVuYWJsZSgpOwogICAgICAgICAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgICAgICAgICByZXR1cm47CiAgICAgICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgICAgICAvL0lmIGEgcGF0aHMgY29uZmlnLCB0aGVuIGp1c3QgbG9hZCB0aGF0IGZpbGUgaW5zdGVhZCB0bwogICAgICAgICAgICAgICAgICAgIC8vcmVzb2x2ZSB0aGUgcGx1Z2luLCBhcyBpdCBpcyBidWlsdCBpbnRvIHRoYXQgcGF0aHMgbGF5ZXIuCiAgICAgICAgICAgICAgICAgICAgaWYgKGJ1bmRsZUlkKSB7CiAgICAgICAgICAgICAgICAgICAgICAgIHRoaXMubWFwLnVybCA9IGNvbnRleHQubmFtZVRvVXJsKGJ1bmRsZUlkKTsKICAgICAgICAgICAgICAgICAgICAgICAgdGhpcy5sb2FkKCk7CiAgICAgICAgICAgICAgICAgICAgICAgIHJldHVybjsKICAgICAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgICAgIGxvYWQgPSBiaW5kKHRoaXMsIGZ1bmN0aW9uICh2YWx1ZSkgewogICAgICAgICAgICAgICAgICAgICAgICB0aGlzLmluaXQoW10sIGZ1bmN0aW9uICgpIHsgcmV0dXJuIHZhbHVlOyB9LCBudWxsLCB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICBlbmFibGVkOiB0cnVlCiAgICAgICAgICAgICAgICAgICAgICAgIH0pOwogICAgICAgICAgICAgICAgICAgIH0pOwoKICAgICAgICAgICAgICAgICAgICBsb2FkLmVycm9yID0gYmluZCh0aGlzLCBmdW5jdGlvbiAoZXJyKSB7CiAgICAgICAgICAgICAgICAgICAgICAgIHRoaXMuaW5pdGVkID0gdHJ1ZTsKICAgICAgICAgICAgICAgICAgICAgICAgdGhpcy5lcnJvciA9IGVycjsKICAgICAgICAgICAgICAgICAgICAgICAgZXJyLnJlcXVpcmVNb2R1bGVzID0gW2lkXTsKCiAgICAgICAgICAgICAgICAgICAgICAgIC8vUmVtb3ZlIHRlbXAgdW5ub3JtYWxpemVkIG1vZHVsZXMgZm9yIHRoaXMgbW9kdWxlLAogICAgICAgICAgICAgICAgICAgICAgICAvL3NpbmNlIHRoZXkgd2lsbCBuZXZlciBiZSByZXNvbHZlZCBvdGhlcndpc2Ugbm93LgogICAgICAgICAgICAgICAgICAgICAgICBlYWNoUHJvcChyZWdpc3RyeSwgZnVuY3Rpb24gKG1vZCkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgaWYgKG1vZC5tYXAuaWQuaW5kZXhPZihpZCArICdfdW5ub3JtYWxpemVkJykgPT09IDApIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBjbGVhblJlZ2lzdHJ5KG1vZC5tYXAuaWQpOwogICAgICAgICAgICAgICAgICAgICAgICAgICAgfQogICAgICAgICAgICAgICAgICAgICAgICB9KTsKCiAgICAgICAgICAgICAgICAgICAgICAgIG9uRXJyb3IoZXJyKTsKICAgICAgICAgICAgICAgICAgICB9KTsKCiAgICAgICAgICAgICAgICAgICAgLy9BbGxvdyBwbHVnaW5zIHRvIGxvYWQgb3RoZXIgY29kZSB3aXRob3V0IGhhdmluZyB0byBrbm93IHRoZQogICAgICAgICAgICAgICAgICAgIC8vY29udGV4dCBvciBob3cgdG8gJ2NvbXBsZXRlJyB0aGUgbG9hZC4KICAgICAgICAgICAgICAgICAgICBsb2FkLmZyb21UZXh0ID0gYmluZCh0aGlzLCBmdW5jdGlvbiAodGV4dCwgdGV4dEFsdCkgewogICAgICAgICAgICAgICAgICAgICAgICAvKmpzbGludCBldmlsOiB0cnVlICovCiAgICAgICAgICAgICAgICAgICAgICAgIHZhciBtb2R1bGVOYW1lID0gbWFwLm5hbWUsCiAgICAgICAgICAgICAgICAgICAgICAgICAgICBtb2R1bGVNYXAgPSBtYWtlTW9kdWxlTWFwKG1vZHVsZU5hbWUpLAogICAgICAgICAgICAgICAgICAgICAgICAgICAgaGFzSW50ZXJhY3RpdmUgPSB1c2VJbnRlcmFjdGl2ZTsKCiAgICAgICAgICAgICAgICAgICAgICAgIC8vQXMgb2YgMi4xLjAsIHN1cHBvcnQganVzdCBwYXNzaW5nIHRoZSB0ZXh0LCB0byByZWluZm9yY2UKICAgICAgICAgICAgICAgICAgICAgICAgLy9mcm9tVGV4dCBvbmx5IGJlaW5nIGNhbGxlZCBvbmNlIHBlciByZXNvdXJjZS4gU3RpbGwKICAgICAgICAgICAgICAgICAgICAgICAgLy9zdXBwb3J0IG9sZCBzdHlsZSBvZiBwYXNzaW5nIG1vZHVsZU5hbWUgYnV0IGRpc2NhcmQKICAgICAgICAgICAgICAgICAgICAgICAgLy90aGF0IG1vZHVsZU5hbWUgaW4gZmF2b3Igb2YgdGhlIGludGVybmFsIHJlZi4KICAgICAgICAgICAgICAgICAgICAgICAgaWYgKHRleHRBbHQpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIHRleHQgPSB0ZXh0QWx0OwogICAgICAgICAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgICAgICAgICAvL1R1cm4gb2ZmIGludGVyYWN0aXZlIHNjcmlwdCBtYXRjaGluZyBmb3IgSUUgZm9yIGFueSBkZWZpbmUKICAgICAgICAgICAgICAgICAgICAgICAgLy9jYWxscyBpbiB0aGUgdGV4dCwgdGhlbiB0dXJuIGl0IGJhY2sgb24gYXQgdGhlIGVuZC4KICAgICAgICAgICAgICAgICAgICAgICAgaWYgKGhhc0ludGVyYWN0aXZlKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICB1c2VJbnRlcmFjdGl2ZSA9IGZhbHNlOwogICAgICAgICAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgICAgICAgICAvL1ByaW1lIHRoZSBzeXN0ZW0gYnkgY3JlYXRpbmcgYSBtb2R1bGUgaW5zdGFuY2UgZm9yCiAgICAgICAgICAgICAgICAgICAgICAgIC8vaXQuCiAgICAgICAgICAgICAgICAgICAgICAgIGdldE1vZHVsZShtb2R1bGVNYXApOwoKICAgICAgICAgICAgICAgICAgICAgICAgLy9UcmFuc2ZlciBhbnkgY29uZmlnIHRvIHRoaXMgb3RoZXIgbW9kdWxlLgogICAgICAgICAgICAgICAgICAgICAgICBpZiAoaGFzUHJvcChjb25maWcuY29uZmlnLCBpZCkpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIGNvbmZpZy5jb25maWdbbW9kdWxlTmFtZV0gPSBjb25maWcuY29uZmlnW2lkXTsKICAgICAgICAgICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgICAgICAgICAgdHJ5IHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIHJlcS5leGVjKHRleHQpOwogICAgICAgICAgICAgICAgICAgICAgICB9IGNhdGNoIChlKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICByZXR1cm4gb25FcnJvcihtYWtlRXJyb3IoJ2Zyb210ZXh0ZXZhbCcsCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICdmcm9tVGV4dCBldmFsIGZvciAnICsgaWQgKwogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICcgZmFpbGVkOiAnICsgZSwKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgZSwKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgW2lkXSkpOwogICAgICAgICAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgICAgICAgICBpZiAoaGFzSW50ZXJhY3RpdmUpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIHVzZUludGVyYWN0aXZlID0gdHJ1ZTsKICAgICAgICAgICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgICAgICAgICAgLy9NYXJrIHRoaXMgYXMgYSBkZXBlbmRlbmN5IGZvciB0aGUgcGx1Z2luCiAgICAgICAgICAgICAgICAgICAgICAgIC8vcmVzb3VyY2UKICAgICAgICAgICAgICAgICAgICAgICAgdGhpcy5kZXBNYXBzLnB1c2gobW9kdWxlTWFwKTsKCiAgICAgICAgICAgICAgICAgICAgICAgIC8vU3VwcG9ydCBhbm9ueW1vdXMgbW9kdWxlcy4KICAgICAgICAgICAgICAgICAgICAgICAgY29udGV4dC5jb21wbGV0ZUxvYWQobW9kdWxlTmFtZSk7CgogICAgICAgICAgICAgICAgICAgICAgICAvL0JpbmQgdGhlIHZhbHVlIG9mIHRoYXQgbW9kdWxlIHRvIHRoZSB2YWx1ZSBmb3IgdGhpcwogICAgICAgICAgICAgICAgICAgICAgICAvL3Jlc291cmNlIElELgogICAgICAgICAgICAgICAgICAgICAgICBsb2NhbFJlcXVpcmUoW21vZHVsZU5hbWVdLCBsb2FkKTsKICAgICAgICAgICAgICAgICAgICB9KTsKCiAgICAgICAgICAgICAgICAgICAgLy9Vc2UgcGFyZW50TmFtZSBoZXJlIHNpbmNlIHRoZSBwbHVnaW4ncyBuYW1lIGlzIG5vdCByZWxpYWJsZSwKICAgICAgICAgICAgICAgICAgICAvL2NvdWxkIGJlIHNvbWUgd2VpcmQgc3RyaW5nIHdpdGggbm8gcGF0aCB0aGF0IGFjdHVhbGx5IHdhbnRzIHRvCiAgICAgICAgICAgICAgICAgICAgLy9yZWZlcmVuY2UgdGhlIHBhcmVudE5hbWUncyBwYXRoLgogICAgICAgICAgICAgICAgICAgIHBsdWdpbi5sb2FkKG1hcC5uYW1lLCBsb2NhbFJlcXVpcmUsIGxvYWQsIGNvbmZpZyk7CiAgICAgICAgICAgICAgICB9KSk7CgogICAgICAgICAgICAgICAgY29udGV4dC5lbmFibGUocGx1Z2luTWFwLCB0aGlzKTsKICAgICAgICAgICAgICAgIHRoaXMucGx1Z2luTWFwc1twbHVnaW5NYXAuaWRdID0gcGx1Z2luTWFwOwogICAgICAgICAgICB9LAoKICAgICAgICAgICAgZW5hYmxlOiBmdW5jdGlvbiAoKSB7CiAgICAgICAgICAgICAgICBlbmFibGVkUmVnaXN0cnlbdGhpcy5tYXAuaWRdID0gdGhpczsKICAgICAgICAgICAgICAgIHRoaXMuZW5hYmxlZCA9IHRydWU7CgogICAgICAgICAgICAgICAgLy9TZXQgZmxhZyBtZW50aW9uaW5nIHRoYXQgdGhlIG1vZHVsZSBpcyBlbmFibGluZywKICAgICAgICAgICAgICAgIC8vc28gdGhhdCBpbW1lZGlhdGUgY2FsbHMgdG8gdGhlIGRlZmluZWQgY2FsbGJhY2tzCiAgICAgICAgICAgICAgICAvL2ZvciBkZXBlbmRlbmNpZXMgZG8gbm90IHRyaWdnZXIgaW5hZHZlcnRlbnQgbG9hZAogICAgICAgICAgICAgICAgLy93aXRoIHRoZSBkZXBDb3VudCBzdGlsbCBiZWluZyB6ZXJvLgogICAgICAgICAgICAgICAgdGhpcy5lbmFibGluZyA9IHRydWU7CgogICAgICAgICAgICAgICAgLy9FbmFibGUgZWFjaCBkZXBlbmRlbmN5CiAgICAgICAgICAgICAgICBlYWNoKHRoaXMuZGVwTWFwcywgYmluZCh0aGlzLCBmdW5jdGlvbiAoZGVwTWFwLCBpKSB7CiAgICAgICAgICAgICAgICAgICAgdmFyIGlkLCBtb2QsIGhhbmRsZXI7CgogICAgICAgICAgICAgICAgICAgIGlmICh0eXBlb2YgZGVwTWFwID09PSAnc3RyaW5nJykgewogICAgICAgICAgICAgICAgICAgICAgICAvL0RlcGVuZGVuY3kgbmVlZHMgdG8gYmUgY29udmVydGVkIHRvIGEgZGVwTWFwCiAgICAgICAgICAgICAgICAgICAgICAgIC8vYW5kIHdpcmVkIHVwIHRvIHRoaXMgbW9kdWxlLgogICAgICAgICAgICAgICAgICAgICAgICBkZXBNYXAgPSBtYWtlTW9kdWxlTWFwKGRlcE1hcCwKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAodGhpcy5tYXAuaXNEZWZpbmUgPyB0aGlzLm1hcCA6IHRoaXMubWFwLnBhcmVudE1hcCksCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgZmFsc2UsCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIXRoaXMuc2tpcE1hcCk7CiAgICAgICAgICAgICAgICAgICAgICAgIHRoaXMuZGVwTWFwc1tpXSA9IGRlcE1hcDsKCiAgICAgICAgICAgICAgICAgICAgICAgIGhhbmRsZXIgPSBnZXRPd24oaGFuZGxlcnMsIGRlcE1hcC5pZCk7CgogICAgICAgICAgICAgICAgICAgICAgICBpZiAoaGFuZGxlcikgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgdGhpcy5kZXBFeHBvcnRzW2ldID0gaGFuZGxlcih0aGlzKTsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIHJldHVybjsKICAgICAgICAgICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgICAgICAgICAgdGhpcy5kZXBDb3VudCArPSAxOwoKICAgICAgICAgICAgICAgICAgICAgICAgb24oZGVwTWFwLCAnZGVmaW5lZCcsIGJpbmQodGhpcywgZnVuY3Rpb24gKGRlcEV4cG9ydHMpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIGlmICh0aGlzLnVuZGVmZWQpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICByZXR1cm47CiAgICAgICAgICAgICAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgICAgICAgICAgICAgICAgICB0aGlzLmRlZmluZURlcChpLCBkZXBFeHBvcnRzKTsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIHRoaXMuY2hlY2soKTsKICAgICAgICAgICAgICAgICAgICAgICAgfSkpOwoKICAgICAgICAgICAgICAgICAgICAgICAgaWYgKHRoaXMuZXJyYmFjaykgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgb24oZGVwTWFwLCAnZXJyb3InLCBiaW5kKHRoaXMsIHRoaXMuZXJyYmFjaykpOwogICAgICAgICAgICAgICAgICAgICAgICB9IGVsc2UgaWYgKHRoaXMuZXZlbnRzLmVycm9yKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAvLyBObyBkaXJlY3QgZXJyYmFjayBvbiB0aGlzIG1vZHVsZSwgYnV0IHNvbWV0aGluZwogICAgICAgICAgICAgICAgICAgICAgICAgICAgLy8gZWxzZSBpcyBsaXN0ZW5pbmcgZm9yIGVycm9ycywgc28gYmUgc3VyZSB0bwogICAgICAgICAgICAgICAgICAgICAgICAgICAgLy8gcHJvcGFnYXRlIHRoZSBlcnJvciBjb3JyZWN0bHkuCiAgICAgICAgICAgICAgICAgICAgICAgICAgICBvbihkZXBNYXAsICdlcnJvcicsIGJpbmQodGhpcywgZnVuY3Rpb24oZXJyKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgdGhpcy5lbWl0KCdlcnJvcicsIGVycik7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICB9KSk7CiAgICAgICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgICAgIGlkID0gZGVwTWFwLmlkOwogICAgICAgICAgICAgICAgICAgIG1vZCA9IHJlZ2lzdHJ5W2lkXTsKCiAgICAgICAgICAgICAgICAgICAgLy9Ta2lwIHNwZWNpYWwgbW9kdWxlcyBsaWtlICdyZXF1aXJlJywgJ2V4cG9ydHMnLCAnbW9kdWxlJwogICAgICAgICAgICAgICAgICAgIC8vQWxzbywgZG9uJ3QgY2FsbCBlbmFibGUgaWYgaXQgaXMgYWxyZWFkeSBlbmFibGVkLAogICAgICAgICAgICAgICAgICAgIC8vaW1wb3J0YW50IGluIGNpcmN1bGFyIGRlcGVuZGVuY3kgY2FzZXMuCiAgICAgICAgICAgICAgICAgICAgaWYgKCFoYXNQcm9wKGhhbmRsZXJzLCBpZCkgJiYgbW9kICYmICFtb2QuZW5hYmxlZCkgewogICAgICAgICAgICAgICAgICAgICAgICBjb250ZXh0LmVuYWJsZShkZXBNYXAsIHRoaXMpOwogICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgIH0pKTsKCiAgICAgICAgICAgICAgICAvL0VuYWJsZSBlYWNoIHBsdWdpbiB0aGF0IGlzIHVzZWQgaW4KICAgICAgICAgICAgICAgIC8vYSBkZXBlbmRlbmN5CiAgICAgICAgICAgICAgICBlYWNoUHJvcCh0aGlzLnBsdWdpbk1hcHMsIGJpbmQodGhpcywgZnVuY3Rpb24gKHBsdWdpbk1hcCkgewogICAgICAgICAgICAgICAgICAgIHZhciBtb2QgPSBnZXRPd24ocmVnaXN0cnksIHBsdWdpbk1hcC5pZCk7CiAgICAgICAgICAgICAgICAgICAgaWYgKG1vZCAmJiAhbW9kLmVuYWJsZWQpIHsKICAgICAgICAgICAgICAgICAgICAgICAgY29udGV4dC5lbmFibGUocGx1Z2luTWFwLCB0aGlzKTsKICAgICAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgICAgICB9KSk7CgogICAgICAgICAgICAgICAgdGhpcy5lbmFibGluZyA9IGZhbHNlOwoKICAgICAgICAgICAgICAgIHRoaXMuY2hlY2soKTsKICAgICAgICAgICAgfSwKCiAgICAgICAgICAgIG9uOiBmdW5jdGlvbiAobmFtZSwgY2IpIHsKICAgICAgICAgICAgICAgIHZhciBjYnMgPSB0aGlzLmV2ZW50c1tuYW1lXTsKICAgICAgICAgICAgICAgIGlmICghY2JzKSB7CiAgICAgICAgICAgICAgICAgICAgY2JzID0gdGhpcy5ldmVudHNbbmFtZV0gPSBbXTsKICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgIGNicy5wdXNoKGNiKTsKICAgICAgICAgICAgfSwKCiAgICAgICAgICAgIGVtaXQ6IGZ1bmN0aW9uIChuYW1lLCBldnQpIHsKICAgICAgICAgICAgICAgIGVhY2godGhpcy5ldmVudHNbbmFtZV0sIGZ1bmN0aW9uIChjYikgewogICAgICAgICAgICAgICAgICAgIGNiKGV2dCk7CiAgICAgICAgICAgICAgICB9KTsKICAgICAgICAgICAgICAgIGlmIChuYW1lID09PSAnZXJyb3InKSB7CiAgICAgICAgICAgICAgICAgICAgLy9Ob3cgdGhhdCB0aGUgZXJyb3IgaGFuZGxlciB3YXMgdHJpZ2dlcmVkLCByZW1vdmUKICAgICAgICAgICAgICAgICAgICAvL3RoZSBsaXN0ZW5lcnMsIHNpbmNlIHRoaXMgYnJva2VuIE1vZHVsZSBpbnN0YW5jZQogICAgICAgICAgICAgICAgICAgIC8vY2FuIHN0YXkgYXJvdW5kIGZvciBhIHdoaWxlIGluIHRoZSByZWdpc3RyeS4KICAgICAgICAgICAgICAgICAgICBkZWxldGUgdGhpcy5ldmVudHNbbmFtZV07CiAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgIH0KICAgICAgICB9OwoKICAgICAgICBmdW5jdGlvbiBjYWxsR2V0TW9kdWxlKGFyZ3MpIHsKICAgICAgICAgICAgLy9Ta2lwIG1vZHVsZXMgYWxyZWFkeSBkZWZpbmVkLgogICAgICAgICAgICBpZiAoIWhhc1Byb3AoZGVmaW5lZCwgYXJnc1swXSkpIHsKICAgICAgICAgICAgICAgIGdldE1vZHVsZShtYWtlTW9kdWxlTWFwKGFyZ3NbMF0sIG51bGwsIHRydWUpKS5pbml0KGFyZ3NbMV0sIGFyZ3NbMl0pOwogICAgICAgICAgICB9CiAgICAgICAgfQoKICAgICAgICBmdW5jdGlvbiByZW1vdmVMaXN0ZW5lcihub2RlLCBmdW5jLCBuYW1lLCBpZU5hbWUpIHsKICAgICAgICAgICAgLy9GYXZvciBkZXRhY2hFdmVudCBiZWNhdXNlIG9mIElFOQogICAgICAgICAgICAvL2lzc3VlLCBzZWUgYXR0YWNoRXZlbnQvYWRkRXZlbnRMaXN0ZW5lciBjb21tZW50IGVsc2V3aGVyZQogICAgICAgICAgICAvL2luIHRoaXMgZmlsZS4KICAgICAgICAgICAgaWYgKG5vZGUuZGV0YWNoRXZlbnQgJiYgIWlzT3BlcmEpIHsKICAgICAgICAgICAgICAgIC8vUHJvYmFibHkgSUUuIElmIG5vdCBpdCB3aWxsIHRocm93IGFuIGVycm9yLCB3aGljaCB3aWxsIGJlCiAgICAgICAgICAgICAgICAvL3VzZWZ1bCB0byBrbm93LgogICAgICAgICAgICAgICAgaWYgKGllTmFtZSkgewogICAgICAgICAgICAgICAgICAgIG5vZGUuZGV0YWNoRXZlbnQoaWVOYW1lLCBmdW5jKTsKICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgfSBlbHNlIHsKICAgICAgICAgICAgICAgIG5vZGUucmVtb3ZlRXZlbnRMaXN0ZW5lcihuYW1lLCBmdW5jLCBmYWxzZSk7CiAgICAgICAgICAgIH0KICAgICAgICB9CgogICAgICAgIC8qKgogICAgICAgICAqIEdpdmVuIGFuIGV2ZW50IGZyb20gYSBzY3JpcHQgbm9kZSwgZ2V0IHRoZSByZXF1aXJlanMgaW5mbyBmcm9tIGl0LAogICAgICAgICAqIGFuZCB0aGVuIHJlbW92ZXMgdGhlIGV2ZW50IGxpc3RlbmVycyBvbiB0aGUgbm9kZS4KICAgICAgICAgKiBAcGFyYW0ge0V2ZW50fSBldnQKICAgICAgICAgKiBAcmV0dXJucyB7T2JqZWN0fQogICAgICAgICAqLwogICAgICAgIGZ1bmN0aW9uIGdldFNjcmlwdERhdGEoZXZ0KSB7CiAgICAgICAgICAgIC8vVXNpbmcgY3VycmVudFRhcmdldCBpbnN0ZWFkIG9mIHRhcmdldCBmb3IgRmlyZWZveCAyLjAncyBzYWtlLiBOb3QKICAgICAgICAgICAgLy9hbGwgb2xkIGJyb3dzZXJzIHdpbGwgYmUgc3VwcG9ydGVkLCBidXQgdGhpcyBvbmUgd2FzIGVhc3kgZW5vdWdoCiAgICAgICAgICAgIC8vdG8gc3VwcG9ydCBhbmQgc3RpbGwgbWFrZXMgc2Vuc2UuCiAgICAgICAgICAgIHZhciBub2RlID0gZXZ0LmN1cnJlbnRUYXJnZXQgfHwgZXZ0LnNyY0VsZW1lbnQ7CgogICAgICAgICAgICAvL1JlbW92ZSB0aGUgbGlzdGVuZXJzIG9uY2UgaGVyZS4KICAgICAgICAgICAgcmVtb3ZlTGlzdGVuZXIobm9kZSwgY29udGV4dC5vblNjcmlwdExvYWQsICdsb2FkJywgJ29ucmVhZHlzdGF0ZWNoYW5nZScpOwogICAgICAgICAgICByZW1vdmVMaXN0ZW5lcihub2RlLCBjb250ZXh0Lm9uU2NyaXB0RXJyb3IsICdlcnJvcicpOwoKICAgICAgICAgICAgcmV0dXJuIHsKICAgICAgICAgICAgICAgIG5vZGU6IG5vZGUsCiAgICAgICAgICAgICAgICBpZDogbm9kZSAmJiBub2RlLmdldEF0dHJpYnV0ZSgnZGF0YS1yZXF1aXJlbW9kdWxlJykKICAgICAgICAgICAgfTsKICAgICAgICB9CgogICAgICAgIGZ1bmN0aW9uIGludGFrZURlZmluZXMoKSB7CiAgICAgICAgICAgIHZhciBhcmdzOwoKICAgICAgICAgICAgLy9BbnkgZGVmaW5lZCBtb2R1bGVzIGluIHRoZSBnbG9iYWwgcXVldWUsIGludGFrZSB0aGVtIG5vdy4KICAgICAgICAgICAgdGFrZUdsb2JhbFF1ZXVlKCk7CgogICAgICAgICAgICAvL01ha2Ugc3VyZSBhbnkgcmVtYWluaW5nIGRlZlF1ZXVlIGl0ZW1zIGdldCBwcm9wZXJseSBwcm9jZXNzZWQuCiAgICAgICAgICAgIHdoaWxlIChkZWZRdWV1ZS5sZW5ndGgpIHsKICAgICAgICAgICAgICAgIGFyZ3MgPSBkZWZRdWV1ZS5zaGlmdCgpOwogICAgICAgICAgICAgICAgaWYgKGFyZ3NbMF0gPT09IG51bGwpIHsKICAgICAgICAgICAgICAgICAgICByZXR1cm4gb25FcnJvcihtYWtlRXJyb3IoJ21pc21hdGNoJywgJ01pc21hdGNoZWQgYW5vbnltb3VzIGRlZmluZSgpIG1vZHVsZTogJyArCiAgICAgICAgICAgICAgICAgICAgICAgIGFyZ3NbYXJncy5sZW5ndGggLSAxXSkpOwogICAgICAgICAgICAgICAgfSBlbHNlIHsKICAgICAgICAgICAgICAgICAgICAvL2FyZ3MgYXJlIGlkLCBkZXBzLCBmYWN0b3J5LiBTaG91bGQgYmUgbm9ybWFsaXplZCBieSB0aGUKICAgICAgICAgICAgICAgICAgICAvL2RlZmluZSgpIGZ1bmN0aW9uLgogICAgICAgICAgICAgICAgICAgIGNhbGxHZXRNb2R1bGUoYXJncyk7CiAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgIH0KICAgICAgICAgICAgY29udGV4dC5kZWZRdWV1ZU1hcCA9IHt9OwogICAgICAgIH0KCiAgICAgICAgY29udGV4dCA9IHsKICAgICAgICAgICAgY29uZmlnOiBjb25maWcsCiAgICAgICAgICAgIGNvbnRleHROYW1lOiBjb250ZXh0TmFtZSwKICAgICAgICAgICAgcmVnaXN0cnk6IHJlZ2lzdHJ5LAogICAgICAgICAgICBkZWZpbmVkOiBkZWZpbmVkLAogICAgICAgICAgICB1cmxGZXRjaGVkOiB1cmxGZXRjaGVkLAogICAgICAgICAgICBkZWZRdWV1ZTogZGVmUXVldWUsCiAgICAgICAgICAgIGRlZlF1ZXVlTWFwOiB7fSwKICAgICAgICAgICAgTW9kdWxlOiBNb2R1bGUsCiAgICAgICAgICAgIG1ha2VNb2R1bGVNYXA6IG1ha2VNb2R1bGVNYXAsCiAgICAgICAgICAgIG5leHRUaWNrOiByZXEubmV4dFRpY2ssCiAgICAgICAgICAgIG9uRXJyb3I6IG9uRXJyb3IsCgogICAgICAgICAgICAvKioKICAgICAgICAgICAgICogU2V0IGEgY29uZmlndXJhdGlvbiBmb3IgdGhlIGNvbnRleHQuCiAgICAgICAgICAgICAqIEBwYXJhbSB7T2JqZWN0fSBjZmcgY29uZmlnIG9iamVjdCB0byBpbnRlZ3JhdGUuCiAgICAgICAgICAgICAqLwogICAgICAgICAgICBjb25maWd1cmU6IGZ1bmN0aW9uIChjZmcpIHsKICAgICAgICAgICAgICAgIC8vTWFrZSBzdXJlIHRoZSBiYXNlVXJsIGVuZHMgaW4gYSBzbGFzaC4KICAgICAgICAgICAgICAgIGlmIChjZmcuYmFzZVVybCkgewogICAgICAgICAgICAgICAgICAgIGlmIChjZmcuYmFzZVVybC5jaGFyQXQoY2ZnLmJhc2VVcmwubGVuZ3RoIC0gMSkgIT09ICcvJykgewogICAgICAgICAgICAgICAgICAgICAgICBjZmcuYmFzZVVybCArPSAnLyc7CiAgICAgICAgICAgICAgICAgICAgfQogICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgIC8vU2F2ZSBvZmYgdGhlIHBhdGhzIHNpbmNlIHRoZXkgcmVxdWlyZSBzcGVjaWFsIHByb2Nlc3NpbmcsCiAgICAgICAgICAgICAgICAvL3RoZXkgYXJlIGFkZGl0aXZlLgogICAgICAgICAgICAgICAgdmFyIHNoaW0gPSBjb25maWcuc2hpbSwKICAgICAgICAgICAgICAgICAgICBvYmpzID0gewogICAgICAgICAgICAgICAgICAgICAgICBwYXRoczogdHJ1ZSwKICAgICAgICAgICAgICAgICAgICAgICAgYnVuZGxlczogdHJ1ZSwKICAgICAgICAgICAgICAgICAgICAgICAgY29uZmlnOiB0cnVlLAogICAgICAgICAgICAgICAgICAgICAgICBtYXA6IHRydWUKICAgICAgICAgICAgICAgICAgICB9OwoKICAgICAgICAgICAgICAgIGVhY2hQcm9wKGNmZywgZnVuY3Rpb24gKHZhbHVlLCBwcm9wKSB7CiAgICAgICAgICAgICAgICAgICAgaWYgKG9ianNbcHJvcF0pIHsKICAgICAgICAgICAgICAgICAgICAgICAgaWYgKCFjb25maWdbcHJvcF0pIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIGNvbmZpZ1twcm9wXSA9IHt9OwogICAgICAgICAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgICAgICAgICAgICAgIG1peGluKGNvbmZpZ1twcm9wXSwgdmFsdWUsIHRydWUsIHRydWUpOwogICAgICAgICAgICAgICAgICAgIH0gZWxzZSB7CiAgICAgICAgICAgICAgICAgICAgICAgIGNvbmZpZ1twcm9wXSA9IHZhbHVlOwogICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgIH0pOwoKICAgICAgICAgICAgICAgIC8vUmV2ZXJzZSBtYXAgdGhlIGJ1bmRsZXMKICAgICAgICAgICAgICAgIGlmIChjZmcuYnVuZGxlcykgewogICAgICAgICAgICAgICAgICAgIGVhY2hQcm9wKGNmZy5idW5kbGVzLCBmdW5jdGlvbiAodmFsdWUsIHByb3ApIHsKICAgICAgICAgICAgICAgICAgICAgICAgZWFjaCh2YWx1ZSwgZnVuY3Rpb24gKHYpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIGlmICh2ICE9PSBwcm9wKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgYnVuZGxlc01hcFt2XSA9IHByb3A7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgICAgICAgICAgICAgIH0pOwogICAgICAgICAgICAgICAgICAgIH0pOwogICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgIC8vTWVyZ2Ugc2hpbQogICAgICAgICAgICAgICAgaWYgKGNmZy5zaGltKSB7CiAgICAgICAgICAgICAgICAgICAgZWFjaFByb3AoY2ZnLnNoaW0sIGZ1bmN0aW9uICh2YWx1ZSwgaWQpIHsKICAgICAgICAgICAgICAgICAgICAgICAgLy9Ob3JtYWxpemUgdGhlIHN0cnVjdHVyZQogICAgICAgICAgICAgICAgICAgICAgICBpZiAoaXNBcnJheSh2YWx1ZSkpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIHZhbHVlID0gewogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIGRlcHM6IHZhbHVlCiAgICAgICAgICAgICAgICAgICAgICAgICAgICB9OwogICAgICAgICAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgICAgICAgICAgICAgIGlmICgodmFsdWUuZXhwb3J0cyB8fCB2YWx1ZS5pbml0KSAmJiAhdmFsdWUuZXhwb3J0c0ZuKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICB2YWx1ZS5leHBvcnRzRm4gPSBjb250ZXh0Lm1ha2VTaGltRXhwb3J0cyh2YWx1ZSk7CiAgICAgICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICAgICAgc2hpbVtpZF0gPSB2YWx1ZTsKICAgICAgICAgICAgICAgICAgICB9KTsKICAgICAgICAgICAgICAgICAgICBjb25maWcuc2hpbSA9IHNoaW07CiAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgLy9BZGp1c3QgcGFja2FnZXMgaWYgbmVjZXNzYXJ5LgogICAgICAgICAgICAgICAgaWYgKGNmZy5wYWNrYWdlcykgewogICAgICAgICAgICAgICAgICAgIGVhY2goY2ZnLnBhY2thZ2VzLCBmdW5jdGlvbiAocGtnT2JqKSB7CiAgICAgICAgICAgICAgICAgICAgICAgIHZhciBsb2NhdGlvbiwgbmFtZTsKCiAgICAgICAgICAgICAgICAgICAgICAgIHBrZ09iaiA9IHR5cGVvZiBwa2dPYmogPT09ICdzdHJpbmcnID8ge25hbWU6IHBrZ09ian0gOiBwa2dPYmo7CgogICAgICAgICAgICAgICAgICAgICAgICBuYW1lID0gcGtnT2JqLm5hbWU7CiAgICAgICAgICAgICAgICAgICAgICAgIGxvY2F0aW9uID0gcGtnT2JqLmxvY2F0aW9uOwogICAgICAgICAgICAgICAgICAgICAgICBpZiAobG9jYXRpb24pIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIGNvbmZpZy5wYXRoc1tuYW1lXSA9IHBrZ09iai5sb2NhdGlvbjsKICAgICAgICAgICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgICAgICAgICAgLy9TYXZlIHBvaW50ZXIgdG8gbWFpbiBtb2R1bGUgSUQgZm9yIHBrZyBuYW1lLgogICAgICAgICAgICAgICAgICAgICAgICAvL1JlbW92ZSBsZWFkaW5nIGRvdCBpbiBtYWluLCBzbyBtYWluIHBhdGhzIGFyZSBub3JtYWxpemVkLAogICAgICAgICAgICAgICAgICAgICAgICAvL2FuZCByZW1vdmUgYW55IHRyYWlsaW5nIC5qcywgc2luY2UgZGlmZmVyZW50IHBhY2thZ2UKICAgICAgICAgICAgICAgICAgICAgICAgLy9lbnZzIGhhdmUgZGlmZmVyZW50IGNvbnZlbnRpb25zOiBzb21lIHVzZSBhIG1vZHVsZSBuYW1lLAogICAgICAgICAgICAgICAgICAgICAgICAvL3NvbWUgdXNlIGEgZmlsZSBuYW1lLgogICAgICAgICAgICAgICAgICAgICAgICBjb25maWcucGtnc1tuYW1lXSA9IHBrZ09iai5uYW1lICsgJy8nICsgKHBrZ09iai5tYWluIHx8ICdtYWluJykKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIC5yZXBsYWNlKGN1cnJEaXJSZWdFeHAsICcnKQogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgLnJlcGxhY2UoanNTdWZmaXhSZWdFeHAsICcnKTsKICAgICAgICAgICAgICAgICAgICB9KTsKICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICAvL0lmIHRoZXJlIGFyZSBhbnkgIndhaXRpbmcgdG8gZXhlY3V0ZSIgbW9kdWxlcyBpbiB0aGUgcmVnaXN0cnksCiAgICAgICAgICAgICAgICAvL3VwZGF0ZSB0aGUgbWFwcyBmb3IgdGhlbSwgc2luY2UgdGhlaXIgaW5mbywgbGlrZSBVUkxzIHRvIGxvYWQsCiAgICAgICAgICAgICAgICAvL21heSBoYXZlIGNoYW5nZWQuCiAgICAgICAgICAgICAgICBlYWNoUHJvcChyZWdpc3RyeSwgZnVuY3Rpb24gKG1vZCwgaWQpIHsKICAgICAgICAgICAgICAgICAgICAvL0lmIG1vZHVsZSBhbHJlYWR5IGhhcyBpbml0IGNhbGxlZCwgc2luY2UgaXQgaXMgdG9vCiAgICAgICAgICAgICAgICAgICAgLy9sYXRlIHRvIG1vZGlmeSB0aGVtLCBhbmQgaWdub3JlIHVubm9ybWFsaXplZCBvbmVzCiAgICAgICAgICAgICAgICAgICAgLy9zaW5jZSB0aGV5IGFyZSB0cmFuc2llbnQuCiAgICAgICAgICAgICAgICAgICAgaWYgKCFtb2QuaW5pdGVkICYmICFtb2QubWFwLnVubm9ybWFsaXplZCkgewogICAgICAgICAgICAgICAgICAgICAgICBtb2QubWFwID0gbWFrZU1vZHVsZU1hcChpZCwgbnVsbCwgdHJ1ZSk7CiAgICAgICAgICAgICAgICAgICAgfQogICAgICAgICAgICAgICAgfSk7CgogICAgICAgICAgICAgICAgLy9JZiBhIGRlcHMgYXJyYXkgb3IgYSBjb25maWcgY2FsbGJhY2sgaXMgc3BlY2lmaWVkLCB0aGVuIGNhbGwKICAgICAgICAgICAgICAgIC8vcmVxdWlyZSB3aXRoIHRob3NlIGFyZ3MuIFRoaXMgaXMgdXNlZnVsIHdoZW4gcmVxdWlyZSBpcyBkZWZpbmVkIGFzIGEKICAgICAgICAgICAgICAgIC8vY29uZmlnIG9iamVjdCBiZWZvcmUgcmVxdWlyZS5qcyBpcyBsb2FkZWQuCiAgICAgICAgICAgICAgICBpZiAoY2ZnLmRlcHMgfHwgY2ZnLmNhbGxiYWNrKSB7CiAgICAgICAgICAgICAgICAgICAgY29udGV4dC5yZXF1aXJlKGNmZy5kZXBzIHx8IFtdLCBjZmcuY2FsbGJhY2spOwogICAgICAgICAgICAgICAgfQogICAgICAgICAgICB9LAoKICAgICAgICAgICAgbWFrZVNoaW1FeHBvcnRzOiBmdW5jdGlvbiAodmFsdWUpIHsKICAgICAgICAgICAgICAgIGZ1bmN0aW9uIGZuKCkgewogICAgICAgICAgICAgICAgICAgIHZhciByZXQ7CiAgICAgICAgICAgICAgICAgICAgaWYgKHZhbHVlLmluaXQpIHsKICAgICAgICAgICAgICAgICAgICAgICAgcmV0ID0gdmFsdWUuaW5pdC5hcHBseShnbG9iYWwsIGFyZ3VtZW50cyk7CiAgICAgICAgICAgICAgICAgICAgfQogICAgICAgICAgICAgICAgICAgIHJldHVybiByZXQgfHwgKHZhbHVlLmV4cG9ydHMgJiYgZ2V0R2xvYmFsKHZhbHVlLmV4cG9ydHMpKTsKICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgIHJldHVybiBmbjsKICAgICAgICAgICAgfSwKCiAgICAgICAgICAgIG1ha2VSZXF1aXJlOiBmdW5jdGlvbiAocmVsTWFwLCBvcHRpb25zKSB7CiAgICAgICAgICAgICAgICBvcHRpb25zID0gb3B0aW9ucyB8fCB7fTsKCiAgICAgICAgICAgICAgICBmdW5jdGlvbiBsb2NhbFJlcXVpcmUoZGVwcywgY2FsbGJhY2ssIGVycmJhY2spIHsKICAgICAgICAgICAgICAgICAgICB2YXIgaWQsIG1hcCwgcmVxdWlyZU1vZDsKCiAgICAgICAgICAgICAgICAgICAgaWYgKG9wdGlvbnMuZW5hYmxlQnVpbGRDYWxsYmFjayAmJiBjYWxsYmFjayAmJiBpc0Z1bmN0aW9uKGNhbGxiYWNrKSkgewogICAgICAgICAgICAgICAgICAgICAgICBjYWxsYmFjay5fX3JlcXVpcmVKc0J1aWxkID0gdHJ1ZTsKICAgICAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgICAgIGlmICh0eXBlb2YgZGVwcyA9PT0gJ3N0cmluZycpIHsKICAgICAgICAgICAgICAgICAgICAgICAgaWYgKGlzRnVuY3Rpb24oY2FsbGJhY2spKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAvL0ludmFsaWQgY2FsbAogICAgICAgICAgICAgICAgICAgICAgICAgICAgcmV0dXJuIG9uRXJyb3IobWFrZUVycm9yKCdyZXF1aXJlYXJncycsICdJbnZhbGlkIHJlcXVpcmUgY2FsbCcpLCBlcnJiYWNrKTsKICAgICAgICAgICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgICAgICAgICAgLy9JZiByZXF1aXJlfGV4cG9ydHN8bW9kdWxlIGFyZSByZXF1ZXN0ZWQsIGdldCB0aGUKICAgICAgICAgICAgICAgICAgICAgICAgLy92YWx1ZSBmb3IgdGhlbSBmcm9tIHRoZSBzcGVjaWFsIGhhbmRsZXJzLiBDYXZlYXQ6CiAgICAgICAgICAgICAgICAgICAgICAgIC8vdGhpcyBvbmx5IHdvcmtzIHdoaWxlIG1vZHVsZSBpcyBiZWluZyBkZWZpbmVkLgogICAgICAgICAgICAgICAgICAgICAgICBpZiAocmVsTWFwICYmIGhhc1Byb3AoaGFuZGxlcnMsIGRlcHMpKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICByZXR1cm4gaGFuZGxlcnNbZGVwc10ocmVnaXN0cnlbcmVsTWFwLmlkXSk7CiAgICAgICAgICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICAgICAgICAgIC8vU3luY2hyb25vdXMgYWNjZXNzIHRvIG9uZSBtb2R1bGUuIElmIHJlcXVpcmUuZ2V0IGlzCiAgICAgICAgICAgICAgICAgICAgICAgIC8vYXZhaWxhYmxlIChhcyBpbiB0aGUgTm9kZSBhZGFwdGVyKSwgcHJlZmVyIHRoYXQuCiAgICAgICAgICAgICAgICAgICAgICAgIGlmIChyZXEuZ2V0KSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICByZXR1cm4gcmVxLmdldChjb250ZXh0LCBkZXBzLCByZWxNYXAsIGxvY2FsUmVxdWlyZSk7CiAgICAgICAgICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICAgICAgICAgIC8vTm9ybWFsaXplIG1vZHVsZSBuYW1lLCBpZiBpdCBjb250YWlucyAuIG9yIC4uCiAgICAgICAgICAgICAgICAgICAgICAgIG1hcCA9IG1ha2VNb2R1bGVNYXAoZGVwcywgcmVsTWFwLCBmYWxzZSwgdHJ1ZSk7CiAgICAgICAgICAgICAgICAgICAgICAgIGlkID0gbWFwLmlkOwoKICAgICAgICAgICAgICAgICAgICAgICAgaWYgKCFoYXNQcm9wKGRlZmluZWQsIGlkKSkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgcmV0dXJuIG9uRXJyb3IobWFrZUVycm9yKCdub3Rsb2FkZWQnLCAnTW9kdWxlIG5hbWUgIicgKwogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgaWQgKwogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgJyIgaGFzIG5vdCBiZWVuIGxvYWRlZCB5ZXQgZm9yIGNvbnRleHQ6ICcgKwogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgY29udGV4dE5hbWUgKwogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgKHJlbE1hcCA/ICcnIDogJy4gVXNlIHJlcXVpcmUoW10pJykpKTsKICAgICAgICAgICAgICAgICAgICAgICAgfQogICAgICAgICAgICAgICAgICAgICAgICByZXR1cm4gZGVmaW5lZFtpZF07CiAgICAgICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgICAgICAvL0dyYWIgZGVmaW5lcyB3YWl0aW5nIGluIHRoZSBnbG9iYWwgcXVldWUuCiAgICAgICAgICAgICAgICAgICAgaW50YWtlRGVmaW5lcygpOwoKICAgICAgICAgICAgICAgICAgICAvL01hcmsgYWxsIHRoZSBkZXBlbmRlbmNpZXMgYXMgbmVlZGluZyB0byBiZSBsb2FkZWQuCiAgICAgICAgICAgICAgICAgICAgY29udGV4dC5uZXh0VGljayhmdW5jdGlvbiAoKSB7CiAgICAgICAgICAgICAgICAgICAgICAgIC8vU29tZSBkZWZpbmVzIGNvdWxkIGhhdmUgYmVlbiBhZGRlZCBzaW5jZSB0aGUKICAgICAgICAgICAgICAgICAgICAgICAgLy9yZXF1aXJlIGNhbGwsIGNvbGxlY3QgdGhlbS4KICAgICAgICAgICAgICAgICAgICAgICAgaW50YWtlRGVmaW5lcygpOwoKICAgICAgICAgICAgICAgICAgICAgICAgcmVxdWlyZU1vZCA9IGdldE1vZHVsZShtYWtlTW9kdWxlTWFwKG51bGwsIHJlbE1hcCkpOwoKICAgICAgICAgICAgICAgICAgICAgICAgLy9TdG9yZSBpZiBtYXAgY29uZmlnIHNob3VsZCBiZSBhcHBsaWVkIHRvIHRoaXMgcmVxdWlyZQogICAgICAgICAgICAgICAgICAgICAgICAvL2NhbGwgZm9yIGRlcGVuZGVuY2llcy4KICAgICAgICAgICAgICAgICAgICAgICAgcmVxdWlyZU1vZC5za2lwTWFwID0gb3B0aW9ucy5za2lwTWFwOwoKICAgICAgICAgICAgICAgICAgICAgICAgcmVxdWlyZU1vZC5pbml0KGRlcHMsIGNhbGxiYWNrLCBlcnJiYWNrLCB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICBlbmFibGVkOiB0cnVlCiAgICAgICAgICAgICAgICAgICAgICAgIH0pOwoKICAgICAgICAgICAgICAgICAgICAgICAgY2hlY2tMb2FkZWQoKTsKICAgICAgICAgICAgICAgICAgICB9KTsKCiAgICAgICAgICAgICAgICAgICAgcmV0dXJuIGxvY2FsUmVxdWlyZTsKICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICBtaXhpbihsb2NhbFJlcXVpcmUsIHsKICAgICAgICAgICAgICAgICAgICBpc0Jyb3dzZXI6IGlzQnJvd3NlciwKCiAgICAgICAgICAgICAgICAgICAgLyoqCiAgICAgICAgICAgICAgICAgICAgICogQ29udmVydHMgYSBtb2R1bGUgbmFtZSArIC5leHRlbnNpb24gaW50byBhbiBVUkwgcGF0aC4KICAgICAgICAgICAgICAgICAgICAgKiAqUmVxdWlyZXMqIHRoZSB1c2Ugb2YgYSBtb2R1bGUgbmFtZS4gSXQgZG9lcyBub3Qgc3VwcG9ydCB1c2luZwogICAgICAgICAgICAgICAgICAgICAqIHBsYWluIFVSTHMgbGlrZSBuYW1lVG9VcmwuCiAgICAgICAgICAgICAgICAgICAgICovCiAgICAgICAgICAgICAgICAgICAgdG9Vcmw6IGZ1bmN0aW9uIChtb2R1bGVOYW1lUGx1c0V4dCkgewogICAgICAgICAgICAgICAgICAgICAgICB2YXIgZXh0LAogICAgICAgICAgICAgICAgICAgICAgICAgICAgaW5kZXggPSBtb2R1bGVOYW1lUGx1c0V4dC5sYXN0SW5kZXhPZignLicpLAogICAgICAgICAgICAgICAgICAgICAgICAgICAgc2VnbWVudCA9IG1vZHVsZU5hbWVQbHVzRXh0LnNwbGl0KCcvJylbMF0sCiAgICAgICAgICAgICAgICAgICAgICAgICAgICBpc1JlbGF0aXZlID0gc2VnbWVudCA9PT0gJy4nIHx8IHNlZ21lbnQgPT09ICcuLic7CgogICAgICAgICAgICAgICAgICAgICAgICAvL0hhdmUgYSBmaWxlIGV4dGVuc2lvbiBhbGlhcywgYW5kIGl0IGlzIG5vdCB0aGUKICAgICAgICAgICAgICAgICAgICAgICAgLy9kb3RzIGZyb20gYSByZWxhdGl2ZSBwYXRoLgogICAgICAgICAgICAgICAgICAgICAgICBpZiAoaW5kZXggIT09IC0xICYmICghaXNSZWxhdGl2ZSB8fCBpbmRleCA+IDEpKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICBleHQgPSBtb2R1bGVOYW1lUGx1c0V4dC5zdWJzdHJpbmcoaW5kZXgsIG1vZHVsZU5hbWVQbHVzRXh0Lmxlbmd0aCk7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICBtb2R1bGVOYW1lUGx1c0V4dCA9IG1vZHVsZU5hbWVQbHVzRXh0LnN1YnN0cmluZygwLCBpbmRleCk7CiAgICAgICAgICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICAgICAgICAgIHJldHVybiBjb250ZXh0Lm5hbWVUb1VybChub3JtYWxpemUobW9kdWxlTmFtZVBsdXNFeHQsCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIHJlbE1hcCAmJiByZWxNYXAuaWQsIHRydWUpLCBleHQsICB0cnVlKTsKICAgICAgICAgICAgICAgICAgICB9LAoKICAgICAgICAgICAgICAgICAgICBkZWZpbmVkOiBmdW5jdGlvbiAoaWQpIHsKICAgICAgICAgICAgICAgICAgICAgICAgcmV0dXJuIGhhc1Byb3AoZGVmaW5lZCwgbWFrZU1vZHVsZU1hcChpZCwgcmVsTWFwLCBmYWxzZSwgdHJ1ZSkuaWQpOwogICAgICAgICAgICAgICAgICAgIH0sCgogICAgICAgICAgICAgICAgICAgIHNwZWNpZmllZDogZnVuY3Rpb24gKGlkKSB7CiAgICAgICAgICAgICAgICAgICAgICAgIGlkID0gbWFrZU1vZHVsZU1hcChpZCwgcmVsTWFwLCBmYWxzZSwgdHJ1ZSkuaWQ7CiAgICAgICAgICAgICAgICAgICAgICAgIHJldHVybiBoYXNQcm9wKGRlZmluZWQsIGlkKSB8fCBoYXNQcm9wKHJlZ2lzdHJ5LCBpZCk7CiAgICAgICAgICAgICAgICAgICAgfQogICAgICAgICAgICAgICAgfSk7CgogICAgICAgICAgICAgICAgLy9Pbmx5IGFsbG93IHVuZGVmIG9uIHRvcCBsZXZlbCByZXF1aXJlIGNhbGxzCiAgICAgICAgICAgICAgICBpZiAoIXJlbE1hcCkgewogICAgICAgICAgICAgICAgICAgIGxvY2FsUmVxdWlyZS51bmRlZiA9IGZ1bmN0aW9uIChpZCkgewogICAgICAgICAgICAgICAgICAgICAgICAvL0JpbmQgYW55IHdhaXRpbmcgZGVmaW5lKCkgY2FsbHMgdG8gdGhpcyBjb250ZXh0LAogICAgICAgICAgICAgICAgICAgICAgICAvL2ZpeCBmb3IgIzQwOAogICAgICAgICAgICAgICAgICAgICAgICB0YWtlR2xvYmFsUXVldWUoKTsKCiAgICAgICAgICAgICAgICAgICAgICAgIHZhciBtYXAgPSBtYWtlTW9kdWxlTWFwKGlkLCByZWxNYXAsIHRydWUpLAogICAgICAgICAgICAgICAgICAgICAgICAgICAgbW9kID0gZ2V0T3duKHJlZ2lzdHJ5LCBpZCk7CgogICAgICAgICAgICAgICAgICAgICAgICBtb2QudW5kZWZlZCA9IHRydWU7CiAgICAgICAgICAgICAgICAgICAgICAgIHJlbW92ZVNjcmlwdChpZCk7CgogICAgICAgICAgICAgICAgICAgICAgICBkZWxldGUgZGVmaW5lZFtpZF07CiAgICAgICAgICAgICAgICAgICAgICAgIGRlbGV0ZSB1cmxGZXRjaGVkW21hcC51cmxdOwogICAgICAgICAgICAgICAgICAgICAgICBkZWxldGUgdW5kZWZFdmVudHNbaWRdOwoKICAgICAgICAgICAgICAgICAgICAgICAgLy9DbGVhbiBxdWV1ZWQgZGVmaW5lcyB0b28uIEdvIGJhY2t3YXJkcwogICAgICAgICAgICAgICAgICAgICAgICAvL2luIGFycmF5IHNvIHRoYXQgdGhlIHNwbGljZXMgZG8gbm90CiAgICAgICAgICAgICAgICAgICAgICAgIC8vbWVzcyB1cCB0aGUgaXRlcmF0aW9uLgogICAgICAgICAgICAgICAgICAgICAgICBlYWNoUmV2ZXJzZShkZWZRdWV1ZSwgZnVuY3Rpb24oYXJncywgaSkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgaWYgKGFyZ3NbMF0gPT09IGlkKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgZGVmUXVldWUuc3BsaWNlKGksIDEpOwogICAgICAgICAgICAgICAgICAgICAgICAgICAgfQogICAgICAgICAgICAgICAgICAgICAgICB9KTsKICAgICAgICAgICAgICAgICAgICAgICAgZGVsZXRlIGNvbnRleHQuZGVmUXVldWVNYXBbaWRdOwoKICAgICAgICAgICAgICAgICAgICAgICAgaWYgKG1vZCkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgLy9Ib2xkIG9uIHRvIGxpc3RlbmVycyBpbiBjYXNlIHRoZQogICAgICAgICAgICAgICAgICAgICAgICAgICAgLy9tb2R1bGUgd2lsbCBiZSBhdHRlbXB0ZWQgdG8gYmUgcmVsb2FkZWQKICAgICAgICAgICAgICAgICAgICAgICAgICAgIC8vdXNpbmcgYSBkaWZmZXJlbnQgY29uZmlnLgogICAgICAgICAgICAgICAgICAgICAgICAgICAgaWYgKG1vZC5ldmVudHMuZGVmaW5lZCkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIHVuZGVmRXZlbnRzW2lkXSA9IG1vZC5ldmVudHM7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgICAgICAgICAgICAgY2xlYW5SZWdpc3RyeShpZCk7CiAgICAgICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICB9OwogICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgIHJldHVybiBsb2NhbFJlcXVpcmU7CiAgICAgICAgICAgIH0sCgogICAgICAgICAgICAvKioKICAgICAgICAgICAgICogQ2FsbGVkIHRvIGVuYWJsZSBhIG1vZHVsZSBpZiBpdCBpcyBzdGlsbCBpbiB0aGUgcmVnaXN0cnkKICAgICAgICAgICAgICogYXdhaXRpbmcgZW5hYmxlbWVudC4gQSBzZWNvbmQgYXJnLCBwYXJlbnQsIHRoZSBwYXJlbnQgbW9kdWxlLAogICAgICAgICAgICAgKiBpcyBwYXNzZWQgaW4gZm9yIGNvbnRleHQsIHdoZW4gdGhpcyBtZXRob2QgaXMgb3ZlcnJpZGRlbiBieQogICAgICAgICAgICAgKiB0aGUgb3B0aW1pemVyLiBOb3Qgc2hvd24gaGVyZSB0byBrZWVwIGNvZGUgY29tcGFjdC4KICAgICAgICAgICAgICovCiAgICAgICAgICAgIGVuYWJsZTogZnVuY3Rpb24gKGRlcE1hcCkgewogICAgICAgICAgICAgICAgdmFyIG1vZCA9IGdldE93bihyZWdpc3RyeSwgZGVwTWFwLmlkKTsKICAgICAgICAgICAgICAgIGlmIChtb2QpIHsKICAgICAgICAgICAgICAgICAgICBnZXRNb2R1bGUoZGVwTWFwKS5lbmFibGUoKTsKICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgfSwKCiAgICAgICAgICAgIC8qKgogICAgICAgICAgICAgKiBJbnRlcm5hbCBtZXRob2QgdXNlZCBieSBlbnZpcm9ubWVudCBhZGFwdGVycyB0byBjb21wbGV0ZSBhIGxvYWQgZXZlbnQuCiAgICAgICAgICAgICAqIEEgbG9hZCBldmVudCBjb3VsZCBiZSBhIHNjcmlwdCBsb2FkIG9yIGp1c3QgYSBsb2FkIHBhc3MgZnJvbSBhIHN5bmNocm9ub3VzCiAgICAgICAgICAgICAqIGxvYWQgY2FsbC4KICAgICAgICAgICAgICogQHBhcmFtIHtTdHJpbmd9IG1vZHVsZU5hbWUgdGhlIG5hbWUgb2YgdGhlIG1vZHVsZSB0byBwb3RlbnRpYWxseSBjb21wbGV0ZS4KICAgICAgICAgICAgICovCiAgICAgICAgICAgIGNvbXBsZXRlTG9hZDogZnVuY3Rpb24gKG1vZHVsZU5hbWUpIHsKICAgICAgICAgICAgICAgIHZhciBmb3VuZCwgYXJncywgbW9kLAogICAgICAgICAgICAgICAgICAgIHNoaW0gPSBnZXRPd24oY29uZmlnLnNoaW0sIG1vZHVsZU5hbWUpIHx8IHt9LAogICAgICAgICAgICAgICAgICAgIHNoRXhwb3J0cyA9IHNoaW0uZXhwb3J0czsKCiAgICAgICAgICAgICAgICB0YWtlR2xvYmFsUXVldWUoKTsKCiAgICAgICAgICAgICAgICB3aGlsZSAoZGVmUXVldWUubGVuZ3RoKSB7CiAgICAgICAgICAgICAgICAgICAgYXJncyA9IGRlZlF1ZXVlLnNoaWZ0KCk7CiAgICAgICAgICAgICAgICAgICAgaWYgKGFyZ3NbMF0gPT09IG51bGwpIHsKICAgICAgICAgICAgICAgICAgICAgICAgYXJnc1swXSA9IG1vZHVsZU5hbWU7CiAgICAgICAgICAgICAgICAgICAgICAgIC8vSWYgYWxyZWFkeSBmb3VuZCBhbiBhbm9ueW1vdXMgbW9kdWxlIGFuZCBib3VuZCBpdAogICAgICAgICAgICAgICAgICAgICAgICAvL3RvIHRoaXMgbmFtZSwgdGhlbiB0aGlzIGlzIHNvbWUgb3RoZXIgYW5vbiBtb2R1bGUKICAgICAgICAgICAgICAgICAgICAgICAgLy93YWl0aW5nIGZvciBpdHMgY29tcGxldGVMb2FkIHRvIGZpcmUuCiAgICAgICAgICAgICAgICAgICAgICAgIGlmIChmb3VuZCkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgYnJlYWs7CiAgICAgICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICAgICAgZm91bmQgPSB0cnVlOwogICAgICAgICAgICAgICAgICAgIH0gZWxzZSBpZiAoYXJnc1swXSA9PT0gbW9kdWxlTmFtZSkgewogICAgICAgICAgICAgICAgICAgICAgICAvL0ZvdW5kIG1hdGNoaW5nIGRlZmluZSBjYWxsIGZvciB0aGlzIHNjcmlwdCEKICAgICAgICAgICAgICAgICAgICAgICAgZm91bmQgPSB0cnVlOwogICAgICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICAgICAgY2FsbEdldE1vZHVsZShhcmdzKTsKICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgIGNvbnRleHQuZGVmUXVldWVNYXAgPSB7fTsKCiAgICAgICAgICAgICAgICAvL0RvIHRoaXMgYWZ0ZXIgdGhlIGN5Y2xlIG9mIGNhbGxHZXRNb2R1bGUgaW4gY2FzZSB0aGUgcmVzdWx0CiAgICAgICAgICAgICAgICAvL29mIHRob3NlIGNhbGxzL2luaXQgY2FsbHMgY2hhbmdlcyB0aGUgcmVnaXN0cnkuCiAgICAgICAgICAgICAgICBtb2QgPSBnZXRPd24ocmVnaXN0cnksIG1vZHVsZU5hbWUpOwoKICAgICAgICAgICAgICAgIGlmICghZm91bmQgJiYgIWhhc1Byb3AoZGVmaW5lZCwgbW9kdWxlTmFtZSkgJiYgbW9kICYmICFtb2QuaW5pdGVkKSB7CiAgICAgICAgICAgICAgICAgICAgaWYgKGNvbmZpZy5lbmZvcmNlRGVmaW5lICYmICghc2hFeHBvcnRzIHx8ICFnZXRHbG9iYWwoc2hFeHBvcnRzKSkpIHsKICAgICAgICAgICAgICAgICAgICAgICAgaWYgKGhhc1BhdGhGYWxsYmFjayhtb2R1bGVOYW1lKSkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgcmV0dXJuOwogICAgICAgICAgICAgICAgICAgICAgICB9IGVsc2UgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgcmV0dXJuIG9uRXJyb3IobWFrZUVycm9yKCdub2RlZmluZScsCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICdObyBkZWZpbmUgY2FsbCBmb3IgJyArIG1vZHVsZU5hbWUsCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIG51bGwsCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIFttb2R1bGVOYW1lXSkpOwogICAgICAgICAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgICAgICAgICAgfSBlbHNlIHsKICAgICAgICAgICAgICAgICAgICAgICAgLy9BIHNjcmlwdCB0aGF0IGRvZXMgbm90IGNhbGwgZGVmaW5lKCksIHNvIGp1c3Qgc2ltdWxhdGUKICAgICAgICAgICAgICAgICAgICAgICAgLy90aGUgY2FsbCBmb3IgaXQuCiAgICAgICAgICAgICAgICAgICAgICAgIGNhbGxHZXRNb2R1bGUoW21vZHVsZU5hbWUsIChzaGltLmRlcHMgfHwgW10pLCBzaGltLmV4cG9ydHNGbl0pOwogICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICBjaGVja0xvYWRlZCgpOwogICAgICAgICAgICB9LAoKICAgICAgICAgICAgLyoqCiAgICAgICAgICAgICAqIENvbnZlcnRzIGEgbW9kdWxlIG5hbWUgdG8gYSBmaWxlIHBhdGguIFN1cHBvcnRzIGNhc2VzIHdoZXJlCiAgICAgICAgICAgICAqIG1vZHVsZU5hbWUgbWF5IGFjdHVhbGx5IGJlIGp1c3QgYW4gVVJMLgogICAgICAgICAgICAgKiBOb3RlIHRoYXQgaXQgKipkb2VzIG5vdCoqIGNhbGwgbm9ybWFsaXplIG9uIHRoZSBtb2R1bGVOYW1lLAogICAgICAgICAgICAgKiBpdCBpcyBhc3N1bWVkIHRvIGhhdmUgYWxyZWFkeSBiZWVuIG5vcm1hbGl6ZWQuIFRoaXMgaXMgYW4KICAgICAgICAgICAgICogaW50ZXJuYWwgQVBJLCBub3QgYSBwdWJsaWMgb25lLiBVc2UgdG9VcmwgZm9yIHRoZSBwdWJsaWMgQVBJLgogICAgICAgICAgICAgKi8KICAgICAgICAgICAgbmFtZVRvVXJsOiBmdW5jdGlvbiAobW9kdWxlTmFtZSwgZXh0LCBza2lwRXh0KSB7CiAgICAgICAgICAgICAgICB2YXIgcGF0aHMsIHN5bXMsIGksIHBhcmVudE1vZHVsZSwgdXJsLAogICAgICAgICAgICAgICAgICAgIHBhcmVudFBhdGgsIGJ1bmRsZUlkLAogICAgICAgICAgICAgICAgICAgIHBrZ01haW4gPSBnZXRPd24oY29uZmlnLnBrZ3MsIG1vZHVsZU5hbWUpOwoKICAgICAgICAgICAgICAgIGlmIChwa2dNYWluKSB7CiAgICAgICAgICAgICAgICAgICAgbW9kdWxlTmFtZSA9IHBrZ01haW47CiAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgYnVuZGxlSWQgPSBnZXRPd24oYnVuZGxlc01hcCwgbW9kdWxlTmFtZSk7CgogICAgICAgICAgICAgICAgaWYgKGJ1bmRsZUlkKSB7CiAgICAgICAgICAgICAgICAgICAgcmV0dXJuIGNvbnRleHQubmFtZVRvVXJsKGJ1bmRsZUlkLCBleHQsIHNraXBFeHQpOwogICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgIC8vSWYgYSBjb2xvbiBpcyBpbiB0aGUgVVJMLCBpdCBpbmRpY2F0ZXMgYSBwcm90b2NvbCBpcyB1c2VkIGFuZCBpdCBpcyBqdXN0CiAgICAgICAgICAgICAgICAvL2FuIFVSTCB0byBhIGZpbGUsIG9yIGlmIGl0IHN0YXJ0cyB3aXRoIGEgc2xhc2gsIGNvbnRhaW5zIGEgcXVlcnkgYXJnIChpLmUuID8pCiAgICAgICAgICAgICAgICAvL29yIGVuZHMgd2l0aCAuanMsIHRoZW4gYXNzdW1lIHRoZSB1c2VyIG1lYW50IHRvIHVzZSBhbiB1cmwgYW5kIG5vdCBhIG1vZHVsZSBpZC4KICAgICAgICAgICAgICAgIC8vVGhlIHNsYXNoIGlzIGltcG9ydGFudCBmb3IgcHJvdG9jb2wtbGVzcyBVUkxzIGFzIHdlbGwgYXMgZnVsbCBwYXRocy4KICAgICAgICAgICAgICAgIGlmIChyZXEuanNFeHRSZWdFeHAudGVzdChtb2R1bGVOYW1lKSkgewogICAgICAgICAgICAgICAgICAgIC8vSnVzdCBhIHBsYWluIHBhdGgsIG5vdCBtb2R1bGUgbmFtZSBsb29rdXAsIHNvIGp1c3QgcmV0dXJuIGl0LgogICAgICAgICAgICAgICAgICAgIC8vQWRkIGV4dGVuc2lvbiBpZiBpdCBpcyBpbmNsdWRlZC4gVGhpcyBpcyBhIGJpdCB3b25reSwgb25seSBub24tLmpzIHRoaW5ncyBwYXNzCiAgICAgICAgICAgICAgICAgICAgLy9hbiBleHRlbnNpb24sIHRoaXMgbWV0aG9kIHByb2JhYmx5IG5lZWRzIHRvIGJlIHJld29ya2VkLgogICAgICAgICAgICAgICAgICAgIHVybCA9IG1vZHVsZU5hbWUgKyAoZXh0IHx8ICcnKTsKICAgICAgICAgICAgICAgIH0gZWxzZSB7CiAgICAgICAgICAgICAgICAgICAgLy9BIG1vZHVsZSB0aGF0IG5lZWRzIHRvIGJlIGNvbnZlcnRlZCB0byBhIHBhdGguCiAgICAgICAgICAgICAgICAgICAgcGF0aHMgPSBjb25maWcucGF0aHM7CgogICAgICAgICAgICAgICAgICAgIHN5bXMgPSBtb2R1bGVOYW1lLnNwbGl0KCcvJyk7CiAgICAgICAgICAgICAgICAgICAgLy9Gb3IgZWFjaCBtb2R1bGUgbmFtZSBzZWdtZW50LCBzZWUgaWYgdGhlcmUgaXMgYSBwYXRoCiAgICAgICAgICAgICAgICAgICAgLy9yZWdpc3RlcmVkIGZvciBpdC4gU3RhcnQgd2l0aCBtb3N0IHNwZWNpZmljIG5hbWUKICAgICAgICAgICAgICAgICAgICAvL2FuZCB3b3JrIHVwIGZyb20gaXQuCiAgICAgICAgICAgICAgICAgICAgZm9yIChpID0gc3ltcy5sZW5ndGg7IGkgPiAwOyBpIC09IDEpIHsKICAgICAgICAgICAgICAgICAgICAgICAgcGFyZW50TW9kdWxlID0gc3ltcy5zbGljZSgwLCBpKS5qb2luKCcvJyk7CgogICAgICAgICAgICAgICAgICAgICAgICBwYXJlbnRQYXRoID0gZ2V0T3duKHBhdGhzLCBwYXJlbnRNb2R1bGUpOwogICAgICAgICAgICAgICAgICAgICAgICBpZiAocGFyZW50UGF0aCkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgLy9JZiBhbiBhcnJheSwgaXQgbWVhbnMgdGhlcmUgYXJlIGEgZmV3IGNob2ljZXMsCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAvL0Nob29zZSB0aGUgb25lIHRoYXQgaXMgZGVzaXJlZAogICAgICAgICAgICAgICAgICAgICAgICAgICAgaWYgKGlzQXJyYXkocGFyZW50UGF0aCkpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBwYXJlbnRQYXRoID0gcGFyZW50UGF0aFswXTsKICAgICAgICAgICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICAgICAgICAgIHN5bXMuc3BsaWNlKDAsIGksIHBhcmVudFBhdGgpOwogICAgICAgICAgICAgICAgICAgICAgICAgICAgYnJlYWs7CiAgICAgICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICB9CgogICAgICAgICAgICAgICAgICAgIC8vSm9pbiB0aGUgcGF0aCBwYXJ0cyB0b2dldGhlciwgdGhlbiBmaWd1cmUgb3V0IGlmIGJhc2VVcmwgaXMgbmVlZGVkLgogICAgICAgICAgICAgICAgICAgIHVybCA9IHN5bXMuam9pbignLycpOwogICAgICAgICAgICAgICAgICAgIHVybCArPSAoZXh0IHx8ICgvXmRhdGFcOnxcPy8udGVzdCh1cmwpIHx8IHNraXBFeHQgPyAnJyA6ICcuanMnKSk7CiAgICAgICAgICAgICAgICAgICAgdXJsID0gKHVybC5jaGFyQXQoMCkgPT09ICcvJyB8fCB1cmwubWF0Y2goL15bXHdcK1wuXC1dKzovKSA/ICcnIDogY29uZmlnLmJhc2VVcmwpICsgdXJsOwogICAgICAgICAgICAgICAgfQoKICAgICAgICAgICAgICAgIHJldHVybiBjb25maWcudXJsQXJncyA/IHVybCArCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAoKHVybC5pbmRleE9mKCc/JykgPT09IC0xID8gJz8nIDogJyYnKSArCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgY29uZmlnLnVybEFyZ3MpIDogdXJsOwogICAgICAgICAgICB9LAoKICAgICAgICAgICAgLy9EZWxlZ2F0ZXMgdG8gcmVxLmxvYWQuIEJyb2tlbiBvdXQgYXMgYSBzZXBhcmF0ZSBmdW5jdGlvbiB0bwogICAgICAgICAgICAvL2FsbG93IG92ZXJyaWRpbmcgaW4gdGhlIG9wdGltaXplci4KICAgICAgICAgICAgbG9hZDogZnVuY3Rpb24gKGlkLCB1cmwpIHsKICAgICAgICAgICAgICAgIHJlcS5sb2FkKGNvbnRleHQsIGlkLCB1cmwpOwogICAgICAgICAgICB9LAoKICAgICAgICAgICAgLyoqCiAgICAgICAgICAgICAqIEV4ZWN1dGVzIGEgbW9kdWxlIGNhbGxiYWNrIGZ1bmN0aW9uLiBCcm9rZW4gb3V0IGFzIGEgc2VwYXJhdGUgZnVuY3Rpb24KICAgICAgICAgICAgICogc29sZWx5IHRvIGFsbG93IHRoZSBidWlsZCBzeXN0ZW0gdG8gc2VxdWVuY2UgdGhlIGZpbGVzIGluIHRoZSBidWlsdAogICAgICAgICAgICAgKiBsYXllciBpbiB0aGUgcmlnaHQgc2VxdWVuY2UuCiAgICAgICAgICAgICAqCiAgICAgICAgICAgICAqIEBwcml2YXRlCiAgICAgICAgICAgICAqLwogICAgICAgICAgICBleGVjQ2I6IGZ1bmN0aW9uIChuYW1lLCBjYWxsYmFjaywgYXJncywgZXhwb3J0cykgewogICAgICAgICAgICAgICAgcmV0dXJuIGNhbGxiYWNrLmFwcGx5KGV4cG9ydHMsIGFyZ3MpOwogICAgICAgICAgICB9LAoKICAgICAgICAgICAgLyoqCiAgICAgICAgICAgICAqIGNhbGxiYWNrIGZvciBzY3JpcHQgbG9hZHMsIHVzZWQgdG8gY2hlY2sgc3RhdHVzIG9mIGxvYWRpbmcuCiAgICAgICAgICAgICAqCiAgICAgICAgICAgICAqIEBwYXJhbSB7RXZlbnR9IGV2dCB0aGUgZXZlbnQgZnJvbSB0aGUgYnJvd3NlciBmb3IgdGhlIHNjcmlwdAogICAgICAgICAgICAgKiB0aGF0IHdhcyBsb2FkZWQuCiAgICAgICAgICAgICAqLwogICAgICAgICAgICBvblNjcmlwdExvYWQ6IGZ1bmN0aW9uIChldnQpIHsKICAgICAgICAgICAgICAgIC8vVXNpbmcgY3VycmVudFRhcmdldCBpbnN0ZWFkIG9mIHRhcmdldCBmb3IgRmlyZWZveCAyLjAncyBzYWtlLiBOb3QKICAgICAgICAgICAgICAgIC8vYWxsIG9sZCBicm93c2VycyB3aWxsIGJlIHN1cHBvcnRlZCwgYnV0IHRoaXMgb25lIHdhcyBlYXN5IGVub3VnaAogICAgICAgICAgICAgICAgLy90byBzdXBwb3J0IGFuZCBzdGlsbCBtYWtlcyBzZW5zZS4KICAgICAgICAgICAgICAgIGlmIChldnQudHlwZSA9PT0gJ2xvYWQnIHx8CiAgICAgICAgICAgICAgICAgICAgICAgIChyZWFkeVJlZ0V4cC50ZXN0KChldnQuY3VycmVudFRhcmdldCB8fCBldnQuc3JjRWxlbWVudCkucmVhZHlTdGF0ZSkpKSB7CiAgICAgICAgICAgICAgICAgICAgLy9SZXNldCBpbnRlcmFjdGl2ZSBzY3JpcHQgc28gYSBzY3JpcHQgbm9kZSBpcyBub3QgaGVsZCBvbnRvIGZvcgogICAgICAgICAgICAgICAgICAgIC8vdG8gbG9uZy4KICAgICAgICAgICAgICAgICAgICBpbnRlcmFjdGl2ZVNjcmlwdCA9IG51bGw7CgogICAgICAgICAgICAgICAgICAgIC8vUHVsbCBvdXQgdGhlIG5hbWUgb2YgdGhlIG1vZHVsZSBhbmQgdGhlIGNvbnRleHQuCiAgICAgICAgICAgICAgICAgICAgdmFyIGRhdGEgPSBnZXRTY3JpcHREYXRhKGV2dCk7CiAgICAgICAgICAgICAgICAgICAgY29udGV4dC5jb21wbGV0ZUxvYWQoZGF0YS5pZCk7CiAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgIH0sCgogICAgICAgICAgICAvKioKICAgICAgICAgICAgICogQ2FsbGJhY2sgZm9yIHNjcmlwdCBlcnJvcnMuCiAgICAgICAgICAgICAqLwogICAgICAgICAgICBvblNjcmlwdEVycm9yOiBmdW5jdGlvbiAoZXZ0KSB7CiAgICAgICAgICAgICAgICB2YXIgZGF0YSA9IGdldFNjcmlwdERhdGEoZXZ0KTsKICAgICAgICAgICAgICAgIGlmICghaGFzUGF0aEZhbGxiYWNrKGRhdGEuaWQpKSB7CiAgICAgICAgICAgICAgICAgICAgdmFyIHBhcmVudHMgPSBbXTsKICAgICAgICAgICAgICAgICAgICBlYWNoUHJvcChyZWdpc3RyeSwgZnVuY3Rpb24odmFsdWUsIGtleSkgewogICAgICAgICAgICAgICAgICAgICAgICBpZiAoa2V5LmluZGV4T2YoJ19AcicpICE9PSAwKSB7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICBlYWNoKHZhbHVlLmRlcE1hcHMsIGZ1bmN0aW9uKGRlcE1hcCkgewogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIGlmIChkZXBNYXAuaWQgPT09IGRhdGEuaWQpIHsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgcGFyZW50cy5wdXNoKGtleSk7CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgfQogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIHJldHVybiB0cnVlOwogICAgICAgICAgICAgICAgICAgICAgICAgICAgfSk7CiAgICAgICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgICAgICB9KTsKICAgICAgICAgICAgICAgICAgICByZXR1cm4gb25FcnJvcihtYWtlRXJyb3IoJ3NjcmlwdGVycm9yJywgJ1NjcmlwdCBlcnJvciBmb3IgIicgKyBkYXRhLmlkICsKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgKHBhcmVudHMubGVuZ3RoID8KICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgJyIsIG5lZWRlZCBieTogJyArIHBhcmVudHMuam9pbignLCAnKSA6CiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICciJyksIGV2dCwgW2RhdGEuaWRdKSk7CiAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgIH0KICAgICAgICB9OwoKICAgICAgICBjb250ZXh0LnJlcXVpcmUgPSBjb250ZXh0Lm1ha2VSZXF1aXJlKCk7CiAgICAgICAgcmV0dXJuIGNvbnRleHQ7CiAgICB9CgogICAgLyoqCiAgICAgKiBNYWluIGVudHJ5IHBvaW50LgogICAgICoKICAgICAqIElmIHRoZSBvbmx5IGFyZ3VtZW50IHRvIHJlcXVpcmUgaXMgYSBzdHJpbmcsIHRoZW4gdGhlIG1vZHVsZSB0aGF0CiAgICAgKiBpcyByZXByZXNlbnRlZCBieSB0aGF0IHN0cmluZyBpcyBmZXRjaGVkIGZvciB0aGUgYXBwcm9wcmlhdGUgY29udGV4dC4KICAgICAqCiAgICAgKiBJZiB0aGUgZmlyc3QgYXJndW1lbnQgaXMgYW4gYXJyYXksIHRoZW4gaXQgd2lsbCBiZSB0cmVhdGVkIGFzIGFuIGFycmF5CiAgICAgKiBvZiBkZXBlbmRlbmN5IHN0cmluZyBuYW1lcyB0byBmZXRjaC4gQW4gb3B0aW9uYWwgZnVuY3Rpb24gY2FsbGJhY2sgY2FuCiAgICAgKiBiZSBzcGVjaWZpZWQgdG8gZXhlY3V0ZSB3aGVuIGFsbCBvZiB0aG9zZSBkZXBlbmRlbmNpZXMgYXJlIGF2YWlsYWJsZS4KICAgICAqCiAgICAgKiBNYWtlIGEgbG9jYWwgcmVxIHZhcmlhYmxlIHRvIGhlbHAgQ2FqYSBjb21wbGlhbmNlIChpdCBhc3N1bWVzIHRoaW5ncwogICAgICogb24gYSByZXF1aXJlIHRoYXQgYXJlIG5vdCBzdGFuZGFyZGl6ZWQpLCBhbmQgdG8gZ2l2ZSBhIHNob3J0CiAgICAgKiBuYW1lIGZvciBtaW5pZmljYXRpb24vbG9jYWwgc2NvcGUgdXNlLgogICAgICovCiAgICByZXEgPSByZXF1aXJlanMgPSBmdW5jdGlvbiAoZGVwcywgY2FsbGJhY2ssIGVycmJhY2ssIG9wdGlvbmFsKSB7CgogICAgICAgIC8vRmluZCB0aGUgcmlnaHQgY29udGV4dCwgdXNlIGRlZmF1bHQKICAgICAgICB2YXIgY29udGV4dCwgY29uZmlnLAogICAgICAgICAgICBjb250ZXh0TmFtZSA9IGRlZkNvbnRleHROYW1lOwoKICAgICAgICAvLyBEZXRlcm1pbmUgaWYgaGF2ZSBjb25maWcgb2JqZWN0IGluIHRoZSBjYWxsLgogICAgICAgIGlmICghaXNBcnJheShkZXBzKSAmJiB0eXBlb2YgZGVwcyAhPT0gJ3N0cmluZycpIHsKICAgICAgICAgICAgLy8gZGVwcyBpcyBhIGNvbmZpZyBvYmplY3QKICAgICAgICAgICAgY29uZmlnID0gZGVwczsKICAgICAgICAgICAgaWYgKGlzQXJyYXkoY2FsbGJhY2spKSB7CiAgICAgICAgICAgICAgICAvLyBBZGp1c3QgYXJncyBpZiB0aGVyZSBhcmUgZGVwZW5kZW5jaWVzCiAgICAgICAgICAgICAgICBkZXBzID0gY2FsbGJhY2s7CiAgICAgICAgICAgICAgICBjYWxsYmFjayA9IGVycmJhY2s7CiAgICAgICAgICAgICAgICBlcnJiYWNrID0gb3B0aW9uYWw7CiAgICAgICAgICAgIH0gZWxzZSB7CiAgICAgICAgICAgICAgICBkZXBzID0gW107CiAgICAgICAgICAgIH0KICAgICAgICB9CgogICAgICAgIGlmIChjb25maWcgJiYgY29uZmlnLmNvbnRleHQpIHsKICAgICAgICAgICAgY29udGV4dE5hbWUgPSBjb25maWcuY29udGV4dDsKICAgICAgICB9CgogICAgICAgIGNvbnRleHQgPSBnZXRPd24oY29udGV4dHMsIGNvbnRleHROYW1lKTsKICAgICAgICBpZiAoIWNvbnRleHQpIHsKICAgICAgICAgICAgY29udGV4dCA9IGNvbnRleHRzW2NvbnRleHROYW1lXSA9IHJlcS5zLm5ld0NvbnRleHQoY29udGV4dE5hbWUpOwogICAgICAgIH0KCiAgICAgICAgaWYgKGNvbmZpZykgewogICAgICAgICAgICBjb250ZXh0LmNvbmZpZ3VyZShjb25maWcpOwogICAgICAgIH0KCiAgICAgICAgcmV0dXJuIGNvbnRleHQucmVxdWlyZShkZXBzLCBjYWxsYmFjaywgZXJyYmFjayk7CiAgICB9OwoKICAgIC8qKgogICAgICogU3VwcG9ydCByZXF1aXJlLmNvbmZpZygpIHRvIG1ha2UgaXQgZWFzaWVyIHRvIGNvb3BlcmF0ZSB3aXRoIG90aGVyCiAgICAgKiBBTUQgbG9hZGVycyBvbiBnbG9iYWxseSBhZ3JlZWQgbmFtZXMuCiAgICAgKi8KICAgIHJlcS5jb25maWcgPSBmdW5jdGlvbiAoY29uZmlnKSB7CiAgICAgICAgcmV0dXJuIHJlcShjb25maWcpOwogICAgfTsKCiAgICAvKioKICAgICAqIEV4ZWN1dGUgc29tZXRoaW5nIGFmdGVyIHRoZSBjdXJyZW50IHRpY2sKICAgICAqIG9mIHRoZSBldmVudCBsb29wLiBPdmVycmlkZSBmb3Igb3RoZXIgZW52cwogICAgICogdGhhdCBoYXZlIGEgYmV0dGVyIHNvbHV0aW9uIHRoYW4gc2V0VGltZW91dC4KICAgICAqIEBwYXJhbSAge0Z1bmN0aW9ufSBmbiBmdW5jdGlvbiB0byBleGVjdXRlIGxhdGVyLgogICAgICovCiAgICByZXEubmV4dFRpY2sgPSB0eXBlb2Ygc2V0VGltZW91dCAhPT0gJ3VuZGVmaW5lZCcgPyBmdW5jdGlvbiAoZm4pIHsKICAgICAgICBzZXRUaW1lb3V0KGZuLCA0KTsKICAgIH0gOiBmdW5jdGlvbiAoZm4pIHsgZm4oKTsgfTsKCiAgICAvKioKICAgICAqIEV4cG9ydCByZXF1aXJlIGFzIGEgZ2xvYmFsLCBidXQgb25seSBpZiBpdCBkb2VzIG5vdCBhbHJlYWR5IGV4aXN0LgogICAgICovCiAgICBpZiAoIXJlcXVpcmUpIHsKICAgICAgICByZXF1aXJlID0gcmVxOwogICAgfQoKICAgIHJlcS52ZXJzaW9uID0gdmVyc2lvbjsKCiAgICAvL1VzZWQgdG8gZmlsdGVyIG91dCBkZXBlbmRlbmNpZXMgdGhhdCBhcmUgYWxyZWFkeSBwYXRocy4KICAgIHJlcS5qc0V4dFJlZ0V4cCA9IC9eXC98OnxcP3xcLmpzJC87CiAgICByZXEuaXNCcm93c2VyID0gaXNCcm93c2VyOwogICAgcyA9IHJlcS5zID0gewogICAgICAgIGNvbnRleHRzOiBjb250ZXh0cywKICAgICAgICBuZXdDb250ZXh0OiBuZXdDb250ZXh0CiAgICB9OwoKICAgIC8vQ3JlYXRlIGRlZmF1bHQgY29udGV4dC4KICAgIHJlcSh7fSk7CgogICAgLy9FeHBvcnRzIHNvbWUgY29udGV4dC1zZW5zaXRpdmUgbWV0aG9kcyBvbiBnbG9iYWwgcmVxdWlyZS4KICAgIGVhY2goWwogICAgICAgICd0b1VybCcsCiAgICAgICAgJ3VuZGVmJywKICAgICAgICAnZGVmaW5lZCcsCiAgICAgICAgJ3NwZWNpZmllZCcKICAgIF0sIGZ1bmN0aW9uIChwcm9wKSB7CiAgICAgICAgLy9SZWZlcmVuY2UgZnJvbSBjb250ZXh0cyBpbnN0ZWFkIG9mIGVhcmx5IGJpbmRpbmcgdG8gZGVmYXVsdCBjb250ZXh0LAogICAgICAgIC8vc28gdGhhdCBkdXJpbmcgYnVpbGRzLCB0aGUgbGF0ZXN0IGluc3RhbmNlIG9mIHRoZSBkZWZhdWx0IGNvbnRleHQKICAgICAgICAvL3dpdGggaXRzIGNvbmZpZyBnZXRzIHVzZWQuCiAgICAgICAgcmVxW3Byb3BdID0gZnVuY3Rpb24gKCkgewogICAgICAgICAgICB2YXIgY3R4ID0gY29udGV4dHNbZGVmQ29udGV4dE5hbWVdOwogICAgICAgICAgICByZXR1cm4gY3R4LnJlcXVpcmVbcHJvcF0uYXBwbHkoY3R4LCBhcmd1bWVudHMpOwogICAgICAgIH07CiAgICB9KTsKCiAgICBpZiAoaXNCcm93c2VyKSB7CiAgICAgICAgaGVhZCA9IHMuaGVhZCA9IGRvY3VtZW50LmdldEVsZW1lbnRzQnlUYWdOYW1lKCdoZWFkJylbMF07CiAgICAgICAgLy9JZiBCQVNFIHRhZyBpcyBpbiBwbGF5LCB1c2luZyBhcHBlbmRDaGlsZCBpcyBhIHByb2JsZW0gZm9yIElFNi4KICAgICAgICAvL1doZW4gdGhhdCBicm93c2VyIGRpZXMsIHRoaXMgY2FuIGJlIHJlbW92ZWQuIERldGFpbHMgaW4gdGhpcyBqUXVlcnkgYnVnOgogICAgICAgIC8vaHR0cDovL2Rldi5qcXVlcnkuY29tL3RpY2tldC8yNzA5CiAgICAgICAgYmFzZUVsZW1lbnQgPSBkb2N1bWVudC5nZXRFbGVtZW50c0J5VGFnTmFtZSgnYmFzZScpWzBdOwogICAgICAgIGlmIChiYXNlRWxlbWVudCkgewogICAgICAgICAgICBoZWFkID0gcy5oZWFkID0gYmFzZUVsZW1lbnQucGFyZW50Tm9kZTsKICAgICAgICB9CiAgICB9CgogICAgLyoqCiAgICAgKiBBbnkgZXJyb3JzIHRoYXQgcmVxdWlyZSBleHBsaWNpdGx5IGdlbmVyYXRlcyB3aWxsIGJlIHBhc3NlZCB0byB0aGlzCiAgICAgKiBmdW5jdGlvbi4gSW50ZXJjZXB0L292ZXJyaWRlIGl0IGlmIHlvdSB3YW50IGN1c3RvbSBlcnJvciBoYW5kbGluZy4KICAgICAqIEBwYXJhbSB7RXJyb3J9IGVyciB0aGUgZXJyb3Igb2JqZWN0LgogICAgICovCiAgICByZXEub25FcnJvciA9IGRlZmF1bHRPbkVycm9yOwoKICAgIC8qKgogICAgICogQ3JlYXRlcyB0aGUgbm9kZSBmb3IgdGhlIGxvYWQgY29tbWFuZC4gT25seSB1c2VkIGluIGJyb3dzZXIgZW52cy4KICAgICAqLwogICAgcmVxLmNyZWF0ZU5vZGUgPSBmdW5jdGlvbiAoY29uZmlnLCBtb2R1bGVOYW1lLCB1cmwpIHsKICAgICAgICB2YXIgbm9kZSA9IGNvbmZpZy54aHRtbCA/CiAgICAgICAgICAgICAgICBkb2N1bWVudC5jcmVhdGVFbGVtZW50TlMoJ2h0dHA6Ly93d3cudzMub3JnLzE5OTkveGh0bWwnLCAnaHRtbDpzY3JpcHQnKSA6CiAgICAgICAgICAgICAgICBkb2N1bWVudC5jcmVhdGVFbGVtZW50KCdzY3JpcHQnKTsKICAgICAgICBub2RlLnR5cGUgPSBjb25maWcuc2NyaXB0VHlwZSB8fCAndGV4dC9qYXZhc2NyaXB0JzsKICAgICAgICBub2RlLmNoYXJzZXQgPSAndXRmLTgnOwogICAgICAgIG5vZGUuYXN5bmMgPSB0cnVlOwogICAgICAgIHJldHVybiBub2RlOwogICAgfTsKCiAgICAvKioKICAgICAqIERvZXMgdGhlIHJlcXVlc3QgdG8gbG9hZCBhIG1vZHVsZSBmb3IgdGhlIGJyb3dzZXIgY2FzZS4KICAgICAqIE1ha2UgdGhpcyBhIHNlcGFyYXRlIGZ1bmN0aW9uIHRvIGFsbG93IG90aGVyIGVudmlyb25tZW50cwogICAgICogdG8gb3ZlcnJpZGUgaXQuCiAgICAgKgogICAgICogQHBhcmFtIHtPYmplY3R9IGNvbnRleHQgdGhlIHJlcXVpcmUgY29udGV4dCB0byBmaW5kIHN0YXRlLgogICAgICogQHBhcmFtIHtTdHJpbmd9IG1vZHVsZU5hbWUgdGhlIG5hbWUgb2YgdGhlIG1vZHVsZS4KICAgICAqIEBwYXJhbSB7T2JqZWN0fSB1cmwgdGhlIFVSTCB0byB0aGUgbW9kdWxlLgogICAgICovCiAgICByZXEubG9hZCA9IGZ1bmN0aW9uIChjb250ZXh0LCBtb2R1bGVOYW1lLCB1cmwpIHsKICAgICAgICB2YXIgY29uZmlnID0gKGNvbnRleHQgJiYgY29udGV4dC5jb25maWcpIHx8IHt9LAogICAgICAgICAgICBub2RlOwogICAgICAgIGlmIChpc0Jyb3dzZXIpIHsKICAgICAgICAgICAgLy9JbiB0aGUgYnJvd3NlciBzbyB1c2UgYSBzY3JpcHQgdGFnCiAgICAgICAgICAgIG5vZGUgPSByZXEuY3JlYXRlTm9kZShjb25maWcsIG1vZHVsZU5hbWUsIHVybCk7CiAgICAgICAgICAgIGlmIChjb25maWcub25Ob2RlQ3JlYXRlZCkgewogICAgICAgICAgICAgICAgY29uZmlnLm9uTm9kZUNyZWF0ZWQobm9kZSwgY29uZmlnLCBtb2R1bGVOYW1lLCB1cmwpOwogICAgICAgICAgICB9CgogICAgICAgICAgICBub2RlLnNldEF0dHJpYnV0ZSgnZGF0YS1yZXF1aXJlY29udGV4dCcsIGNvbnRleHQuY29udGV4dE5hbWUpOwogICAgICAgICAgICBub2RlLnNldEF0dHJpYnV0ZSgnZGF0YS1yZXF1aXJlbW9kdWxlJywgbW9kdWxlTmFtZSk7CgogICAgICAgICAgICAvL1NldCB1cCBsb2FkIGxpc3RlbmVyLiBUZXN0IGF0dGFjaEV2ZW50IGZpcnN0IGJlY2F1c2UgSUU5IGhhcwogICAgICAgICAgICAvL2Egc3VidGxlIGlzc3VlIGluIGl0cyBhZGRFdmVudExpc3RlbmVyIGFuZCBzY3JpcHQgb25sb2FkIGZpcmluZ3MKICAgICAgICAgICAgLy90aGF0IGRvIG5vdCBtYXRjaCB0aGUgYmVoYXZpb3Igb2YgYWxsIG90aGVyIGJyb3dzZXJzIHdpdGgKICAgICAgICAgICAgLy9hZGRFdmVudExpc3RlbmVyIHN1cHBvcnQsIHdoaWNoIGZpcmUgdGhlIG9ubG9hZCBldmVudCBmb3IgYQogICAgICAgICAgICAvL3NjcmlwdCByaWdodCBhZnRlciB0aGUgc2NyaXB0IGV4ZWN1dGlvbi4gU2VlOgogICAgICAgICAgICAvL2h0dHBzOi8vY29ubmVjdC5taWNyb3NvZnQuY29tL0lFL2ZlZWRiYWNrL2RldGFpbHMvNjQ4MDU3L3NjcmlwdC1vbmxvYWQtZXZlbnQtaXMtbm90LWZpcmVkLWltbWVkaWF0ZWx5LWFmdGVyLXNjcmlwdC1leGVjdXRpb24KICAgICAgICAgICAgLy9VTkZPUlRVTkFURUxZIE9wZXJhIGltcGxlbWVudHMgYXR0YWNoRXZlbnQgYnV0IGRvZXMgbm90IGZvbGxvdyB0aGUgc2NyaXB0CiAgICAgICAgICAgIC8vc2NyaXB0IGV4ZWN1dGlvbiBtb2RlLgogICAgICAgICAgICBpZiAobm9kZS5hdHRhY2hFdmVudCAmJgogICAgICAgICAgICAgICAgICAgIC8vQ2hlY2sgaWYgbm9kZS5hdHRhY2hFdmVudCBpcyBhcnRpZmljaWFsbHkgYWRkZWQgYnkgY3VzdG9tIHNjcmlwdCBvcgogICAgICAgICAgICAgICAgICAgIC8vbmF0aXZlbHkgc3VwcG9ydGVkIGJ5IGJyb3dzZXIKICAgICAgICAgICAgICAgICAgICAvL3JlYWQgaHR0cHM6Ly9naXRodWIuY29tL2pyYnVya2UvcmVxdWlyZWpzL2lzc3Vlcy8xODcKICAgICAgICAgICAgICAgICAgICAvL2lmIHdlIGNhbiBOT1QgZmluZCBbbmF0aXZlIGNvZGVdIHRoZW4gaXQgbXVzdCBOT1QgbmF0aXZlbHkgc3VwcG9ydGVkLgogICAgICAgICAgICAgICAgICAgIC8vaW4gSUU4LCBub2RlLmF0dGFjaEV2ZW50IGRvZXMgbm90IGhhdmUgdG9TdHJpbmcoKQogICAgICAgICAgICAgICAgICAgIC8vTm90ZSB0aGUgdGVzdCBmb3IgIltuYXRpdmUgY29kZSIgd2l0aCBubyBjbG9zaW5nIGJyYWNlLCBzZWU6CiAgICAgICAgICAgICAgICAgICAgLy9odHRwczovL2dpdGh1Yi5jb20vanJidXJrZS9yZXF1aXJlanMvaXNzdWVzLzI3MwogICAgICAgICAgICAgICAgICAgICEobm9kZS5hdHRhY2hFdmVudC50b1N0cmluZyAmJiBub2RlLmF0dGFjaEV2ZW50LnRvU3RyaW5nKCkuaW5kZXhPZignW25hdGl2ZSBjb2RlJykgPCAwKSAmJgogICAgICAgICAgICAgICAgICAgICFpc09wZXJhKSB7CiAgICAgICAgICAgICAgICAvL1Byb2JhYmx5IElFLiBJRSAoYXQgbGVhc3QgNi04KSBkbyBub3QgZmlyZQogICAgICAgICAgICAgICAgLy9zY3JpcHQgb25sb2FkIHJpZ2h0IGFmdGVyIGV4ZWN1dGluZyB0aGUgc2NyaXB0LCBzbwogICAgICAgICAgICAgICAgLy93ZSBjYW5ub3QgdGllIHRoZSBhbm9ueW1vdXMgZGVmaW5lIGNhbGwgdG8gYSBuYW1lLgogICAgICAgICAgICAgICAgLy9Ib3dldmVyLCBJRSByZXBvcnRzIHRoZSBzY3JpcHQgYXMgYmVpbmcgaW4gJ2ludGVyYWN0aXZlJwogICAgICAgICAgICAgICAgLy9yZWFkeVN0YXRlIGF0IHRoZSB0aW1lIG9mIHRoZSBkZWZpbmUgY2FsbC4KICAgICAgICAgICAgICAgIHVzZUludGVyYWN0aXZlID0gdHJ1ZTsKCiAgICAgICAgICAgICAgICBub2RlLmF0dGFjaEV2ZW50KCdvbnJlYWR5c3RhdGVjaGFuZ2UnLCBjb250ZXh0Lm9uU2NyaXB0TG9hZCk7CiAgICAgICAgICAgICAgICAvL0l0IHdvdWxkIGJlIGdyZWF0IHRvIGFkZCBhbiBlcnJvciBoYW5kbGVyIGhlcmUgdG8gY2F0Y2gKICAgICAgICAgICAgICAgIC8vNDA0cyBpbiBJRTkrLiBIb3dldmVyLCBvbnJlYWR5c3RhdGVjaGFuZ2Ugd2lsbCBmaXJlIGJlZm9yZQogICAgICAgICAgICAgICAgLy90aGUgZXJyb3IgaGFuZGxlciwgc28gdGhhdCBkb2VzIG5vdCBoZWxwLiBJZiBhZGRFdmVudExpc3RlbmVyCiAgICAgICAgICAgICAgICAvL2lzIHVzZWQsIHRoZW4gSUUgd2lsbCBmaXJlIGVycm9yIGJlZm9yZSBsb2FkLCBidXQgd2UgY2Fubm90CiAgICAgICAgICAgICAgICAvL3VzZSB0aGF0IHBhdGh3YXkgZ2l2ZW4gdGhlIGNvbm5lY3QubWljcm9zb2Z0LmNvbSBpc3N1ZQogICAgICAgICAgICAgICAgLy9tZW50aW9uZWQgYWJvdmUgYWJvdXQgbm90IGRvaW5nIHRoZSAnc2NyaXB0IGV4ZWN1dGUsCiAgICAgICAgICAgICAgICAvL3RoZW4gZmlyZSB0aGUgc2NyaXB0IGxvYWQgZXZlbnQgbGlzdGVuZXIgYmVmb3JlIGV4ZWN1dGUKICAgICAgICAgICAgICAgIC8vbmV4dCBzY3JpcHQnIHRoYXQgb3RoZXIgYnJvd3NlcnMgZG8uCiAgICAgICAgICAgICAgICAvL0Jlc3QgaG9wZTogSUUxMCBmaXhlcyB0aGUgaXNzdWVzLAogICAgICAgICAgICAgICAgLy9hbmQgdGhlbiBkZXN0cm95cyBhbGwgaW5zdGFsbHMgb2YgSUUgNi05LgogICAgICAgICAgICAgICAgLy9ub2RlLmF0dGFjaEV2ZW50KCdvbmVycm9yJywgY29udGV4dC5vblNjcmlwdEVycm9yKTsKICAgICAgICAgICAgfSBlbHNlIHsKICAgICAgICAgICAgICAgIG5vZGUuYWRkRXZlbnRMaXN0ZW5lcignbG9hZCcsIGNvbnRleHQub25TY3JpcHRMb2FkLCBmYWxzZSk7CiAgICAgICAgICAgICAgICBub2RlLmFkZEV2ZW50TGlzdGVuZXIoJ2Vycm9yJywgY29udGV4dC5vblNjcmlwdEVycm9yLCBmYWxzZSk7CiAgICAgICAgICAgIH0KICAgICAgICAgICAgbm9kZS5zcmMgPSB1cmw7CgogICAgICAgICAgICAvL0ZvciBzb21lIGNhY2hlIGNhc2VzIGluIElFIDYtOCwgdGhlIHNjcmlwdCBleGVjdXRlcyBiZWZvcmUgdGhlIGVuZAogICAgICAgICAgICAvL29mIHRoZSBhcHBlbmRDaGlsZCBleGVjdXRpb24sIHNvIHRvIHRpZSBhbiBhbm9ueW1vdXMgZGVmaW5lCiAgICAgICAgICAgIC8vY2FsbCB0byB0aGUgbW9kdWxlIG5hbWUgKHdoaWNoIGlzIHN0b3JlZCBvbiB0aGUgbm9kZSksIGhvbGQgb24KICAgICAgICAgICAgLy90byBhIHJlZmVyZW5jZSB0byB0aGlzIG5vZGUsIGJ1dCBjbGVhciBhZnRlciB0aGUgRE9NIGluc2VydGlvbi4KICAgICAgICAgICAgY3VycmVudGx5QWRkaW5nU2NyaXB0ID0gbm9kZTsKICAgICAgICAgICAgaWYgKGJhc2VFbGVtZW50KSB7CiAgICAgICAgICAgICAgICBoZWFkLmluc2VydEJlZm9yZShub2RlLCBiYXNlRWxlbWVudCk7CiAgICAgICAgICAgIH0gZWxzZSB7CiAgICAgICAgICAgICAgICBoZWFkLmFwcGVuZENoaWxkKG5vZGUpOwogICAgICAgICAgICB9CiAgICAgICAgICAgIGN1cnJlbnRseUFkZGluZ1NjcmlwdCA9IG51bGw7CgogICAgICAgICAgICByZXR1cm4gbm9kZTsKICAgICAgICB9IGVsc2UgaWYgKGlzV2ViV29ya2VyKSB7CiAgICAgICAgICAgIHRyeSB7CiAgICAgICAgICAgICAgICAvL0luIGEgd2ViIHdvcmtlciwgdXNlIGltcG9ydFNjcmlwdHMuIFRoaXMgaXMgbm90IGEgdmVyeQogICAgICAgICAgICAgICAgLy9lZmZpY2llbnQgdXNlIG9mIGltcG9ydFNjcmlwdHMsIGltcG9ydFNjcmlwdHMgd2lsbCBibG9jayB1bnRpbAogICAgICAgICAgICAgICAgLy9pdHMgc2NyaXB0IGlzIGRvd25sb2FkZWQgYW5kIGV2YWx1YXRlZC4gSG93ZXZlciwgaWYgd2ViIHdvcmtlcnMKICAgICAgICAgICAgICAgIC8vYXJlIGluIHBsYXksIHRoZSBleHBlY3RhdGlvbiBpcyB0aGF0IGEgYnVpbGQgaGFzIGJlZW4gZG9uZSBzbwogICAgICAgICAgICAgICAgLy90aGF0IG9ubHkgb25lIHNjcmlwdCBuZWVkcyB0byBiZSBsb2FkZWQgYW55d2F5LiBUaGlzIG1heSBuZWVkCiAgICAgICAgICAgICAgICAvL3RvIGJlIHJlZXZhbHVhdGVkIGlmIG90aGVyIHVzZSBjYXNlcyBiZWNvbWUgY29tbW9uLgogICAgICAgICAgICAgICAgaW1wb3J0U2NyaXB0cyh1cmwpOwoKICAgICAgICAgICAgICAgIC8vQWNjb3VudCBmb3IgYW5vbnltb3VzIG1vZHVsZXMKICAgICAgICAgICAgICAgIGNvbnRleHQuY29tcGxldGVMb2FkKG1vZHVsZU5hbWUpOwogICAgICAgICAgICB9IGNhdGNoIChlKSB7CiAgICAgICAgICAgICAgICBjb250ZXh0Lm9uRXJyb3IobWFrZUVycm9yKCdpbXBvcnRzY3JpcHRzJywKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAnaW1wb3J0U2NyaXB0cyBmYWlsZWQgZm9yICcgKwogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBtb2R1bGVOYW1lICsgJyBhdCAnICsgdXJsLAogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIGUsCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgW21vZHVsZU5hbWVdKSk7CiAgICAgICAgICAgIH0KICAgICAgICB9CiAgICB9OwoKICAgIGZ1bmN0aW9uIGdldEludGVyYWN0aXZlU2NyaXB0KCkgewogICAgICAgIGlmIChpbnRlcmFjdGl2ZVNjcmlwdCAmJiBpbnRlcmFjdGl2ZVNjcmlwdC5yZWFkeVN0YXRlID09PSAnaW50ZXJhY3RpdmUnKSB7CiAgICAgICAgICAgIHJldHVybiBpbnRlcmFjdGl2ZVNjcmlwdDsKICAgICAgICB9CgogICAgICAgIGVhY2hSZXZlcnNlKHNjcmlwdHMoKSwgZnVuY3Rpb24gKHNjcmlwdCkgewogICAgICAgICAgICBpZiAoc2NyaXB0LnJlYWR5U3RhdGUgPT09ICdpbnRlcmFjdGl2ZScpIHsKICAgICAgICAgICAgICAgIHJldHVybiAoaW50ZXJhY3RpdmVTY3JpcHQgPSBzY3JpcHQpOwogICAgICAgICAgICB9CiAgICAgICAgfSk7CiAgICAgICAgcmV0dXJuIGludGVyYWN0aXZlU2NyaXB0OwogICAgfQoKICAgIC8vTG9vayBmb3IgYSBkYXRhLW1haW4gc2NyaXB0IGF0dHJpYnV0ZSwgd2hpY2ggY291bGQgYWxzbyBhZGp1c3QgdGhlIGJhc2VVcmwuCiAgICBpZiAoaXNCcm93c2VyICYmICFjZmcuc2tpcERhdGFNYWluKSB7CiAgICAgICAgLy9GaWd1cmUgb3V0IGJhc2VVcmwuIEdldCBpdCBmcm9tIHRoZSBzY3JpcHQgdGFnIHdpdGggcmVxdWlyZS5qcyBpbiBpdC4KICAgICAgICBlYWNoUmV2ZXJzZShzY3JpcHRzKCksIGZ1bmN0aW9uIChzY3JpcHQpIHsKICAgICAgICAgICAgLy9TZXQgdGhlICdoZWFkJyB3aGVyZSB3ZSBjYW4gYXBwZW5kIGNoaWxkcmVuIGJ5CiAgICAgICAgICAgIC8vdXNpbmcgdGhlIHNjcmlwdCdzIHBhcmVudC4KICAgICAgICAgICAgaWYgKCFoZWFkKSB7CiAgICAgICAgICAgICAgICBoZWFkID0gc2NyaXB0LnBhcmVudE5vZGU7CiAgICAgICAgICAgIH0KCiAgICAgICAgICAgIC8vTG9vayBmb3IgYSBkYXRhLW1haW4gYXR0cmlidXRlIHRvIHNldCBtYWluIHNjcmlwdCBmb3IgdGhlIHBhZ2UKICAgICAgICAgICAgLy90byBsb2FkLiBJZiBpdCBpcyB0aGVyZSwgdGhlIHBhdGggdG8gZGF0YSBtYWluIGJlY29tZXMgdGhlCiAgICAgICAgICAgIC8vYmFzZVVybCwgaWYgaXQgaXMgbm90IGFscmVhZHkgc2V0LgogICAgICAgICAgICBkYXRhTWFpbiA9IHNjcmlwdC5nZXRBdHRyaWJ1dGUoJ2RhdGEtbWFpbicpOwogICAgICAgICAgICBpZiAoZGF0YU1haW4pIHsKICAgICAgICAgICAgICAgIC8vUHJlc2VydmUgZGF0YU1haW4gaW4gY2FzZSBpdCBpcyBhIHBhdGggKGkuZS4gY29udGFpbnMgJz8nKQogICAgICAgICAgICAgICAgbWFpblNjcmlwdCA9IGRhdGFNYWluOwoKICAgICAgICAgICAgICAgIC8vU2V0IGZpbmFsIGJhc2VVcmwgaWYgdGhlcmUgaXMgbm90IGFscmVhZHkgYW4gZXhwbGljaXQgb25lLgogICAgICAgICAgICAgICAgaWYgKCFjZmcuYmFzZVVybCkgewogICAgICAgICAgICAgICAgICAgIC8vUHVsbCBvZmYgdGhlIGRpcmVjdG9yeSBvZiBkYXRhLW1haW4gZm9yIHVzZSBhcyB0aGUKICAgICAgICAgICAgICAgICAgICAvL2Jhc2VVcmwuCiAgICAgICAgICAgICAgICAgICAgc3JjID0gbWFpblNjcmlwdC5zcGxpdCgnLycpOwogICAgICAgICAgICAgICAgICAgIG1haW5TY3JpcHQgPSBzcmMucG9wKCk7CiAgICAgICAgICAgICAgICAgICAgc3ViUGF0aCA9IHNyYy5sZW5ndGggPyBzcmMuam9pbignLycpICArICcvJyA6ICcuLyc7CgogICAgICAgICAgICAgICAgICAgIGNmZy5iYXNlVXJsID0gc3ViUGF0aDsKICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICAvL1N0cmlwIG9mZiBhbnkgdHJhaWxpbmcgLmpzIHNpbmNlIG1haW5TY3JpcHQgaXMgbm93CiAgICAgICAgICAgICAgICAvL2xpa2UgYSBtb2R1bGUgbmFtZS4KICAgICAgICAgICAgICAgIG1haW5TY3JpcHQgPSBtYWluU2NyaXB0LnJlcGxhY2UoanNTdWZmaXhSZWdFeHAsICcnKTsKCiAgICAgICAgICAgICAgICAvL0lmIG1haW5TY3JpcHQgaXMgc3RpbGwgYSBwYXRoLCBmYWxsIGJhY2sgdG8gZGF0YU1haW4KICAgICAgICAgICAgICAgIGlmIChyZXEuanNFeHRSZWdFeHAudGVzdChtYWluU2NyaXB0KSkgewogICAgICAgICAgICAgICAgICAgIG1haW5TY3JpcHQgPSBkYXRhTWFpbjsKICAgICAgICAgICAgICAgIH0KCiAgICAgICAgICAgICAgICAvL1B1dCB0aGUgZGF0YS1tYWluIHNjcmlwdCBpbiB0aGUgZmlsZXMgdG8gbG9hZC4KICAgICAgICAgICAgICAgIGNmZy5kZXBzID0gY2ZnLmRlcHMgPyBjZmcuZGVwcy5jb25jYXQobWFpblNjcmlwdCkgOiBbbWFpblNjcmlwdF07CgogICAgICAgICAgICAgICAgcmV0dXJuIHRydWU7CiAgICAgICAgICAgIH0KICAgICAgICB9KTsKICAgIH0KCiAgICAvKioKICAgICAqIFRoZSBmdW5jdGlvbiB0aGF0IGhhbmRsZXMgZGVmaW5pdGlvbnMgb2YgbW9kdWxlcy4gRGlmZmVycyBmcm9tCiAgICAgKiByZXF1aXJlKCkgaW4gdGhhdCBhIHN0cmluZyBmb3IgdGhlIG1vZHVsZSBzaG91bGQgYmUgdGhlIGZpcnN0IGFyZ3VtZW50LAogICAgICogYW5kIHRoZSBmdW5jdGlvbiB0byBleGVjdXRlIGFmdGVyIGRlcGVuZGVuY2llcyBhcmUgbG9hZGVkIHNob3VsZAogICAgICogcmV0dXJuIGEgdmFsdWUgdG8gZGVmaW5lIHRoZSBtb2R1bGUgY29ycmVzcG9uZGluZyB0byB0aGUgZmlyc3QgYXJndW1lbnQncwogICAgICogbmFtZS4KICAgICAqLwogICAgZGVmaW5lID0gZnVuY3Rpb24gKG5hbWUsIGRlcHMsIGNhbGxiYWNrKSB7CiAgICAgICAgdmFyIG5vZGUsIGNvbnRleHQ7CgogICAgICAgIC8vQWxsb3cgZm9yIGFub255bW91cyBtb2R1bGVzCiAgICAgICAgaWYgKHR5cGVvZiBuYW1lICE9PSAnc3RyaW5nJykgewogICAgICAgICAgICAvL0FkanVzdCBhcmdzIGFwcHJvcHJpYXRlbHkKICAgICAgICAgICAgY2FsbGJhY2sgPSBkZXBzOwogICAgICAgICAgICBkZXBzID0gbmFtZTsKICAgICAgICAgICAgbmFtZSA9IG51bGw7CiAgICAgICAgfQoKICAgICAgICAvL1RoaXMgbW9kdWxlIG1heSBub3QgaGF2ZSBkZXBlbmRlbmNpZXMKICAgICAgICBpZiAoIWlzQXJyYXkoZGVwcykpIHsKICAgICAgICAgICAgY2FsbGJhY2sgPSBkZXBzOwogICAgICAgICAgICBkZXBzID0gbnVsbDsKICAgICAgICB9CgogICAgICAgIC8vSWYgbm8gbmFtZSwgYW5kIGNhbGxiYWNrIGlzIGEgZnVuY3Rpb24sIHRoZW4gZmlndXJlIG91dCBpZiBpdCBhCiAgICAgICAgLy9Db21tb25KUyB0aGluZyB3aXRoIGRlcGVuZGVuY2llcy4KICAgICAgICBpZiAoIWRlcHMgJiYgaXNGdW5jdGlvbihjYWxsYmFjaykpIHsKICAgICAgICAgICAgZGVwcyA9IFtdOwogICAgICAgICAgICAvL1JlbW92ZSBjb21tZW50cyBmcm9tIHRoZSBjYWxsYmFjayBzdHJpbmcsCiAgICAgICAgICAgIC8vbG9vayBmb3IgcmVxdWlyZSBjYWxscywgYW5kIHB1bGwgdGhlbSBpbnRvIHRoZSBkZXBlbmRlbmNpZXMsCiAgICAgICAgICAgIC8vYnV0IG9ubHkgaWYgdGhlcmUgYXJlIGZ1bmN0aW9uIGFyZ3MuCiAgICAgICAgICAgIGlmIChjYWxsYmFjay5sZW5ndGgpIHsKICAgICAgICAgICAgICAgIGNhbGxiYWNrCiAgICAgICAgICAgICAgICAgICAgLnRvU3RyaW5nKCkKICAgICAgICAgICAgICAgICAgICAucmVwbGFjZShjb21tZW50UmVnRXhwLCAnJykKICAgICAgICAgICAgICAgICAgICAucmVwbGFjZShjanNSZXF1aXJlUmVnRXhwLCBmdW5jdGlvbiAobWF0Y2gsIGRlcCkgewogICAgICAgICAgICAgICAgICAgICAgICBkZXBzLnB1c2goZGVwKTsKICAgICAgICAgICAgICAgICAgICB9KTsKCiAgICAgICAgICAgICAgICAvL01heSBiZSBhIENvbW1vbkpTIHRoaW5nIGV2ZW4gd2l0aG91dCByZXF1aXJlIGNhbGxzLCBidXQgc3RpbGwKICAgICAgICAgICAgICAgIC8vY291bGQgdXNlIGV4cG9ydHMsIGFuZCBtb2R1bGUuIEF2b2lkIGRvaW5nIGV4cG9ydHMgYW5kIG1vZHVsZQogICAgICAgICAgICAgICAgLy93b3JrIHRob3VnaCBpZiBpdCBqdXN0IG5lZWRzIHJlcXVpcmUuCiAgICAgICAgICAgICAgICAvL1JFUVVJUkVTIHRoZSBmdW5jdGlvbiB0byBleHBlY3QgdGhlIENvbW1vbkpTIHZhcmlhYmxlcyBpbiB0aGUKICAgICAgICAgICAgICAgIC8vb3JkZXIgbGlzdGVkIGJlbG93LgogICAgICAgICAgICAgICAgZGVwcyA9IChjYWxsYmFjay5sZW5ndGggPT09IDEgPyBbJ3JlcXVpcmUnXSA6IFsncmVxdWlyZScsICdleHBvcnRzJywgJ21vZHVsZSddKS5jb25jYXQoZGVwcyk7CiAgICAgICAgICAgIH0KICAgICAgICB9CgogICAgICAgIC8vSWYgaW4gSUUgNi04IGFuZCBoaXQgYW4gYW5vbnltb3VzIGRlZmluZSgpIGNhbGwsIGRvIHRoZSBpbnRlcmFjdGl2ZQogICAgICAgIC8vd29yay4KICAgICAgICBpZiAodXNlSW50ZXJhY3RpdmUpIHsKICAgICAgICAgICAgbm9kZSA9IGN1cnJlbnRseUFkZGluZ1NjcmlwdCB8fCBnZXRJbnRlcmFjdGl2ZVNjcmlwdCgpOwogICAgICAgICAgICBpZiAobm9kZSkgewogICAgICAgICAgICAgICAgaWYgKCFuYW1lKSB7CiAgICAgICAgICAgICAgICAgICAgbmFtZSA9IG5vZGUuZ2V0QXR0cmlidXRlKCdkYXRhLXJlcXVpcmVtb2R1bGUnKTsKICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgIGNvbnRleHQgPSBjb250ZXh0c1tub2RlLmdldEF0dHJpYnV0ZSgnZGF0YS1yZXF1aXJlY29udGV4dCcpXTsKICAgICAgICAgICAgfQogICAgICAgIH0KCiAgICAgICAgLy9BbHdheXMgc2F2ZSBvZmYgZXZhbHVhdGluZyB0aGUgZGVmIGNhbGwgdW50aWwgdGhlIHNjcmlwdCBvbmxvYWQgaGFuZGxlci4KICAgICAgICAvL1RoaXMgYWxsb3dzIG11bHRpcGxlIG1vZHVsZXMgdG8gYmUgaW4gYSBmaWxlIHdpdGhvdXQgcHJlbWF0dXJlbHkKICAgICAgICAvL3RyYWNpbmcgZGVwZW5kZW5jaWVzLCBhbmQgYWxsb3dzIGZvciBhbm9ueW1vdXMgbW9kdWxlIHN1cHBvcnQsCiAgICAgICAgLy93aGVyZSB0aGUgbW9kdWxlIG5hbWUgaXMgbm90IGtub3duIHVudGlsIHRoZSBzY3JpcHQgb25sb2FkIGV2ZW50CiAgICAgICAgLy9vY2N1cnMuIElmIG5vIGNvbnRleHQsIHVzZSB0aGUgZ2xvYmFsIHF1ZXVlLCBhbmQgZ2V0IGl0IHByb2Nlc3NlZAogICAgICAgIC8vaW4gdGhlIG9uc2NyaXB0IGxvYWQgY2FsbGJhY2suCiAgICAgICAgaWYgKGNvbnRleHQpIHsKICAgICAgICAgICAgY29udGV4dC5kZWZRdWV1ZS5wdXNoKFtuYW1lLCBkZXBzLCBjYWxsYmFja10pOwogICAgICAgICAgICBjb250ZXh0LmRlZlF1ZXVlTWFwW25hbWVdID0gdHJ1ZTsKICAgICAgICB9IGVsc2UgewogICAgICAgICAgICBnbG9iYWxEZWZRdWV1ZS5wdXNoKFtuYW1lLCBkZXBzLCBjYWxsYmFja10pOwogICAgICAgIH0KICAgIH07CgogICAgZGVmaW5lLmFtZCA9IHsKICAgICAgICBqUXVlcnk6IHRydWUKICAgIH07CgogICAgLyoqCiAgICAgKiBFeGVjdXRlcyB0aGUgdGV4dC4gTm9ybWFsbHkganVzdCB1c2VzIGV2YWwsIGJ1dCBjYW4gYmUgbW9kaWZpZWQKICAgICAqIHRvIHVzZSBhIGJldHRlciwgZW52aXJvbm1lbnQtc3BlY2lmaWMgY2FsbC4gT25seSB1c2VkIGZvciB0cmFuc3BpbGluZwogICAgICogbG9hZGVyIHBsdWdpbnMsIG5vdCBmb3IgcGxhaW4gSlMgbW9kdWxlcy4KICAgICAqIEBwYXJhbSB7U3RyaW5nfSB0ZXh0IHRoZSB0ZXh0IHRvIGV4ZWN1dGUvZXZhbHVhdGUuCiAgICAgKi8KICAgIHJlcS5leGVjID0gZnVuY3Rpb24gKHRleHQpIHsKICAgICAgICAvKmpzbGludCBldmlsOiB0cnVlICovCiAgICAgICAgcmV0dXJuIGV2YWwodGV4dCk7CiAgICB9OwoKICAgIC8vU2V0IHVwIHdpdGggY29uZmlnIGluZm8uCiAgICByZXEoY2ZnKTsKfSh0aGlzKSk7Cg==",
-              "ok": true,
               "headers": [
                 [
                   "content-type",
                   "text/javascript"
                 ]
               ],
+              "ok": true,
               "status": 200,
               "status_text": ""
             }
-          },
-          "base_uri": "https://localhost:8080/",
-          "height": 2006
+          }
         },
-        "outputId": "0b3f497f-040f-41ef-8a32-70b4adf7d7d0",
+        "colab_type": "code",
         "executionInfo": {
+          "elapsed": 4242,
           "status": "ok",
           "timestamp": 1512371597785,
-          "user_tz": 480,
-          "elapsed": 4242,
           "user": {
             "displayName": "Lukasz Kaiser",
             "photoUrl": "//lh3.googleusercontent.com/-CbWIwcQ_VsA/AAAAAAAAAAI/AAAAAAAAAB8/jloHVR1qOhg/s50-c-k-no/photo.jpg",
             "userId": "109750154298538986950"
-          }
-        }
+          },
+          "user_tz": 480
+        },
+        "id": "OJKU36QAfqOC",
+        "outputId": "0b3f497f-040f-41ef-8a32-70b4adf7d7d0"
       },
-      "cell_type": "code",
-      "source": [
-        "# Convert inputs and outputs to subwords\n",
-        "inp_text = to_tokens(encoders[\"inputs\"].encode(inputs))\n",
-        "out_text = to_tokens(encoders[\"inputs\"].encode(outputs))\n",
-        "\n",
-        "# Run eval to collect attention weights\n",
-        "example = encode_eval(inputs, outputs)\n",
-        "with tfe.restore_variables_on_create(tf.train.latest_checkpoint(checkpoint_dir)):\n",
-        "  translate_model.set_mode(Modes.EVAL)\n",
-        "  translate_model(example)\n",
-        "# Get normalized attention weights for each layer\n",
-        "enc_atts, dec_atts, encdec_atts = get_att_mats()\n",
-        "\n",
-        "call_html()\n",
-        "attention.show(inp_text, out_text, enc_atts, dec_atts, encdec_atts)"
-      ],
-      "execution_count": 0,
       "outputs": [
         {
+          "name": "stdout",
           "output_type": "stream",
           "text": [
             "WARNING:tensorflow:From /usr/local/lib/python2.7/dist-packages/tensor2tensor/layers/common_layers.py:1671: softmax_cross_entropy_with_logits (from tensorflow.python.ops.nn_ops) is deprecated and will be removed in a future version.\n",
@@ -893,19 +791,14 @@
             "\n",
             "See tf.nn.softmax_cross_entropy_with_logits_v2.\n",
             "\n"
-          ],
-          "name": "stdout"
+          ]
         },
         {
-          "output_type": "display_data",
           "data": {
-            "text/plain": [
-              "<IPython.core.display.HTML object>"
-            ],
             "text/html": [
               "\n",
-              "        <script src=\"/static/components/requirejs/require.js\"></script>\n",
-              "        <script>\n",
+              "        \u003cscript src=\"/static/components/requirejs/require.js\"\u003e\u003c/script\u003e\n",
+              "        \u003cscript\u003e\n",
               "          requirejs.config({\n",
               "            paths: {\n",
               "              base: '/static/base',\n",
@@ -913,54 +806,57 @@
               "              jquery: '//ajax.googleapis.com/ajax/libs/jquery/2.0.0/jquery.min',\n",
               "            },\n",
               "          });\n",
-              "        </script>\n",
+              "        \u003c/script\u003e\n",
               "        "
+            ],
+            "text/plain": [
+              "\u003cIPython.core.display.HTML object\u003e"
             ]
           },
           "metadata": {
             "tags": []
-          }
+          },
+          "output_type": "display_data"
         },
         {
-          "output_type": "display_data",
           "data": {
-            "text/plain": [
-              "<IPython.core.display.HTML object>"
-            ],
             "text/html": [
               "\n",
-              "  <span style=\"user-select:none\">\n",
-              "    Layer: <select id=\"layer\"></select>\n",
-              "    Attention: <select id=\"att_type\">\n",
-              "      <option value=\"all\">All</option>\n",
-              "      <option value=\"inp_inp\">Input - Input</option>\n",
-              "      <option value=\"inp_out\">Input - Output</option>\n",
-              "      <option value=\"out_out\">Output - Output</option>\n",
-              "    </select>\n",
-              "  </span>\n",
-              "  <div id='vis'></div>\n"
+              "  \u003cspan style=\"user-select:none\"\u003e\n",
+              "    Layer: \u003cselect id=\"layer\"\u003e\u003c/select\u003e\n",
+              "    Attention: \u003cselect id=\"att_type\"\u003e\n",
+              "      \u003coption value=\"all\"\u003eAll\u003c/option\u003e\n",
+              "      \u003coption value=\"inp_inp\"\u003eInput - Input\u003c/option\u003e\n",
+              "      \u003coption value=\"inp_out\"\u003eInput - Output\u003c/option\u003e\n",
+              "      \u003coption value=\"out_out\"\u003eOutput - Output\u003c/option\u003e\n",
+              "    \u003c/select\u003e\n",
+              "  \u003c/span\u003e\n",
+              "  \u003cdiv id='vis'\u003e\u003c/div\u003e\n"
+            ],
+            "text/plain": [
+              "\u003cIPython.core.display.HTML object\u003e"
             ]
           },
           "metadata": {
             "tags": []
-          }
+          },
+          "output_type": "display_data"
         },
         {
-          "output_type": "display_data",
           "data": {
             "application/javascript": [
               "window.attention = {\"inp_out\": {\"top_text\": [\"The_\", \"animal_\", \"didn_\", \"'_\", \"t_\", \"cross_\", \"the_\", \"street_\", \"because_\", \"it_\", \"was_\", \"too_\", \"tire\", \"d_\"], \"att\": [[[[0.01107952743768692, 0.002038179198279977, 0.02572617679834366, 0.043437324464321136, 0.026865433901548386, 0.008821134455502033, 0.05896050110459328, 0.006038360297679901, 0.05802087485790253, 0.05262080207467079, 0.021981995552778244, 0.01655607670545578, 0.007265332620590925, 0.017941446974873543, 0.19668635725975037], [0.4201550781726837, 0.0003083523770328611, 0.003427971852943301, 0.027074502781033516, 0.0025770263746380806, 0.0006525526405312121, 0.0672224909067154, 0.0006329934694804251, 0.002376251621171832, 0.007315297145396471, 0.0018543159822002053, 0.0002170451043639332, 5.486799182108371e-06, 8.465739665552974e-05, 0.018722370266914368], [6.826388562330976e-05, 0.41254693269729614, 8.318798791151494e-05, 0.00021303755056578666, 2.6623651137924753e-05, 1.3030116861045826e-06, 3.3524677292007254e-06, 9.95700816019962e-07, 0.00025696202646940947, 0.00021154701244086027, 4.0387480112258345e-05, 7.382633339148015e-05, 0.0001871670683613047, 0.0001393109851051122, 0.00044668230111710727], [0.0012913167010992765, 0.46178945899009705, 0.0011929792817682028, 0.0014885100536048412, 0.001382660586386919, 0.00010778238356579095, 4.841455302084796e-05, 4.8626650823280215e-05, 0.0007912410655990243, 0.0019299217965453863, 0.0002972490037791431, 0.0004315593687351793, 0.013707359321415424, 0.0025058358442038298, 0.00208207662217319], [0.0008573953527957201, 5.803010481031379e-06, 0.0034995940513908863, 0.007113253697752953, 4.1040249925572425e-05, 0.48505696654319763, 0.0009781911503523588, 2.57480514846975e-05, 0.0006811833591200411, 0.011991027742624283, 0.013829604722559452, 0.02649468183517456, 0.018967876210808754, 0.008940043859183788, 0.0023627132177352905], [3.2793446735013276e-05, 4.91645641886862e-06, 0.0003670089063234627, 0.0005689052632078528, 0.0004337447171565145, 0.6979628205299377, 0.00025133590679615736, 1.3211038094596006e-05, 0.001040837960317731, 0.0008422345272265375, 0.00011131400242447853, 0.0007033413276076317, 0.00044049491407349706, 0.0004404923238325864, 0.00032976132933981717], [0.002877118531614542, 0.0015123215271160007, 0.21683953702449799, 0.042356427758932114, 0.09360139071941376, 0.7325531840324402, 0.007687804754823446, 0.0004983373219147325, 0.0008397439960390329, 0.018263472244143486, 0.01633409783244133, 0.06572946161031723, 0.029279880225658417, 0.13710656762123108, 0.013406738638877869], [0.09384340792894363, 0.002295592101290822, 0.05245966836810112, 0.10398446023464203, 0.13232196867465973, 0.2621823251247406, 0.7299563884735107, 0.01621837355196476, 0.008298774249851704, 0.019108427688479424, 0.013038183562457561, 0.008606976829469204, 0.0014156820252537727, 0.008462491445243359, 0.08448491245508194], [7.994164479896426e-05, 9.660106115916278e-06, 1.3390360436460469e-05, 0.0009496311540715396, 7.498388185922522e-06, 0.0023292596451938152, 0.0033705621026456356, 0.45610299706459045, 0.00048403104301542044, 0.0003956609289161861, 6.013430538587272e-05, 1.5610943592037074e-05, 4.899038231087616e-06, 1.0044974260381423e-05, 0.0011326958192512393], [0.0021254755556583405, 0.025354469195008278, 0.0505821667611599, 0.04718977212905884, 0.3544465899467468, 0.27984359860420227, 0.10468283295631409, 0.03827415779232979, 0.0065247067250311375, 0.003615353489294648, 0.001024437602609396, 0.02404061146080494, 0.00031744904117658734, 0.011979974806308746, 0.06911104917526245], [0.06793052703142166, 0.04423084855079651, 0.009074175730347633, 0.010606715455651283, 0.023761747404932976, 0.06765440851449966, 0.048715878278017044, 0.13498826324939728, 0.15846557915210724, 0.01835249364376068, 0.0033974519465118647, 0.011923078447580338, 0.0035463334061205387, 0.036997705698013306, 0.15195232629776], [0.00013637961819767952, 0.00010623007256072015, 0.00015417735266964883, 0.00014589299098588526, 0.0007127521676011384, 0.0008950252668000758, 0.00038585966103710234, 0.002901369472965598, 0.34460243582725525, 0.00040915730642154813, 0.00017379666678607464, 9.334777860203758e-05, 0.0002283527428517118, 0.0001650981866987422, 0.0021401161793619394], [0.03951041400432587, 0.015644539147615433, 0.002765331417322159, 0.020979223772883415, 0.001914863707497716, 0.049360573291778564, 0.010446744039654732, 0.06006397679448128, 0.18512527644634247, 0.5769777894020081, 0.07455664873123169, 0.016840822994709015, 0.21517987549304962, 0.030672460794448853, 0.04319411888718605], [0.0012064727488905191, 0.0013226938899606466, 0.002064700936898589, 0.008003294467926025, 0.002116014016792178, 0.0028530799318104982, 0.006337625440210104, 0.0002913604548666626, 0.0004794643900822848, 0.0026383439544588327, 0.0038926906418055296, 0.3737375736236572, 0.002772320294752717, 0.007620541378855705, 0.003997606225311756], [1.0432314411445986e-05, 4.745730166177964e-06, 1.672162215982098e-05, 2.360623693675734e-05, 4.496370820561424e-06, 1.767691173881758e-06, 4.21794857174973e-06, 1.7029789205480483e-06, 2.8430429665604606e-05, 7.409282261505723e-05, 0.00010478614422027022, 0.00017224416660610586, 0.480630487203598, 0.017292670905590057, 3.8113743357826024e-05], [0.00031966043752618134, 7.799067680025473e-05, 0.0005293181748129427, 0.0002383182873018086, 6.09634407737758e-05, 1.622732997930143e-05, 0.0001254813396371901, 4.548055585473776e-05, 0.0002202334435423836, 0.0014038329245522618, 0.008373874239623547, 0.0005300238262861967, 0.8584288358688354, 0.0721927285194397, 0.0012385909212753177], [0.008336205966770649, 0.000929497298784554, 0.060522519052028656, 0.02858084999024868, 0.004865946713835001, 0.19429318606853485, 0.006222299765795469, 0.00020022530225105584, 0.03241097182035446, 0.2199898362159729, 0.40489089488983154, 0.12284909188747406, 0.04783688485622406, 0.16652296483516693, 0.03165041282773018], [0.06735408306121826, 0.02395833097398281, 0.022876637056469917, 0.059418935328722, 0.020556019619107246, 0.006657767109572887, 0.01686989888548851, 0.03750348463654518, 0.0929105281829834, 0.11066772043704987, 0.07383746653795242, 0.04306775704026222, 0.1764260083436966, 0.2488536387681961, 0.14264866709709167], [0.00023218609567265958, 9.724824485601857e-05, 0.00017837552877608687, 0.000249945733230561, 0.00043016509152948856, 0.0002728255931288004, 0.0002596308768261224, 0.0021448382176458836, 0.33870813250541687, 0.0012523159384727478, 0.0004828754754271358, 7.525486580561846e-05, 0.001232807757332921, 0.00022845527564641088, 0.0029908884316682816], [0.044313203543424606, 0.014693659730255604, 0.001713237608782947, 0.01787775754928589, 0.001054717693477869, 0.03111616149544716, 0.005932849366217852, 0.035437386482954025, 0.10908837616443634, 0.6214090585708618, 0.11623460799455643, 0.018710769712924957, 0.26884767413139343, 0.036007944494485855, 0.04555344209074974], [0.0014647350180894136, 0.0016486160457134247, 0.001705971430055797, 0.008203698322176933, 0.0011827786220237613, 0.001036314177326858, 0.004107706248760223, 0.00018337460642214864, 0.0005908485618419945, 0.004427316598594189, 0.0075510423630476, 0.37528446316719055, 0.0045065670274198055, 0.01084148045629263, 0.0047609396278858185], [1.1546462701517157e-05, 6.3197094277711585e-06, 1.3665205187862739e-05, 2.3049220544635318e-05, 3.1024922009237343e-06, 9.712728115118807e-07, 4.2468768697290216e-06, 1.4032799526830786e-06, 2.1501631636056118e-05, 0.00011254433775320649, 0.00014821428339928389, 0.00021640797785948962, 0.4815296530723572, 0.022970588877797127, 4.596232975018211e-05], [0.0004618540406227112, 0.00011890243331436068, 0.0008028792799450457, 0.0003817373653873801, 7.645944424439222e-05, 2.0059787857462652e-05, 0.00017321997438557446, 3.885024489136413e-05, 0.00016429855895694345, 0.0017073642229661345, 0.011983372271060944, 0.0008083870052359998, 0.8495219349861145, 0.07573292404413223, 0.0017974229995161295], [0.00848880223929882, 0.0010204557329416275, 0.06384890526533127, 0.030244439840316772, 0.004545390605926514, 0.2111765593290329, 0.007047791499644518, 0.00020413362653926015, 0.03285042569041252, 0.2096482813358307, 0.40160003304481506, 0.12425301223993301, 0.05433715134859085, 0.2013336718082428, 0.03489448130130768], [0.018106432631611824, 0.01663283444941044, 0.006966447923332453, 0.06288447231054306, 0.008926548063755035, 0.0005806194385513663, 0.004527462646365166, 0.00047311693197116256, 0.010450053960084915, 0.008817908354103565, 0.02498125471174717, 0.02475220151245594, 0.006219316273927689, 0.034688226878643036, 0.15510374307632446]], [[0.011485431343317032, 0.057214245200157166, 0.11445975303649902, 0.035292237997055054, 0.17235025763511658, 0.21079879999160767, 0.08683252334594727, 0.33144259452819824, 0.2781406342983246, 0.07864350080490112, 0.10017280280590057, 0.0828540250658989, 0.17722147703170776, 0.21101748943328857, 0.15805292129516602], [0.041519034653902054, 0.11474552005529404, 0.04909001290798187, 0.1299373209476471, 0.06295691430568695, 0.0239214189350605, 0.22038953006267548, 0.6809458136558533, 0.03295678645372391, 0.34942832589149475, 0.1847512274980545, 0.22206875681877136, 0.13646042346954346, 0.277276873588562, 0.1334262192249298], [0.0764331966638565, 0.004937899298965931, 0.049346037209033966, 0.05165911093354225, 0.051789041608572006, 0.11632981896400452, 0.3382570743560791, 0.21805666387081146, 0.5269062519073486, 0.05627245828509331, 0.1284114420413971, 0.3053610324859619, 0.058564696460962296, 0.14431920647621155, 0.19175130128860474], [0.08274618536233902, 0.009897814132273197, 0.07511309534311295, 0.03663979470729828, 0.16369661688804626, 0.04579350724816322, 0.04420214146375656, 0.06866282969713211, 0.17000554502010345, 0.09549596160650253, 0.07313749194145203, 0.06223462149500847, 0.11603321135044098, 0.07143211364746094, 0.2059532254934311], [0.41769060492515564, 0.07210511714220047, 0.40716952085494995, 0.22363832592964172, 0.48781970143318176, 0.015007800422608852, 0.4504202902317047, 0.4675638973712921, 0.24936619400978088, 0.5447031855583191, 0.4296078681945801, 0.07025930285453796, 0.1902965009212494, 0.3567025065422058, 0.12464861571788788], [0.3858333230018616, 0.06937354803085327, 0.5601253509521484, 0.30969470739364624, 0.36272186040878296, 0.005774383433163166, 0.16290897130966187, 0.16338182985782623, 0.1734752655029297, 0.10127251595258713, 0.6812319159507751, 0.35078492760658264, 0.26554787158966064, 0.3089393675327301, 0.12310608476400375], [0.047016799449920654, 0.04388514533638954, 0.010725832544267178, 0.029561294242739677, 0.04913409426808357, 0.007112162187695503, 0.045616600662469864, 0.09563170373439789, 0.021758677437901497, 0.05606407672166824, 0.023780539631843567, 0.2586848735809326, 0.1317795366048813, 0.13214319944381714, 0.18490085005760193], [0.024271933361887932, 0.10952932387590408, 0.01092300284653902, 0.005798409227281809, 0.03478696197271347, 0.015390553511679173, 0.005925341974943876, 0.04537563398480415, 0.00714160455390811, 0.005484140943735838, 0.00704369880259037, 0.04858299717307091, 0.06617175042629242, 0.13874217867851257, 0.17208275198936462], [0.1448126882314682, 0.16020630300045013, 0.02696153335273266, 0.06902630627155304, 0.03837759047746658, 0.07682601362466812, 0.15773272514343262, 0.005734406877309084, 0.16041570901870728, 0.10849703103303909, 0.08964504301548004, 0.4313186705112457, 0.12084108591079712, 0.20548132061958313, 0.1913137137889862], [0.03147122263908386, 0.06498080492019653, 0.03835386037826538, 0.021906379610300064, 0.004580754786729813, 0.08777225762605667, 0.06548282504081726, 0.0501156747341156, 0.09960248321294785, 0.05812418833374977, 0.04425663501024246, 0.12932318449020386, 0.040425609797239304, 0.10523593425750732, 0.20731014013290405], [0.03185653313994408, 0.014990762807428837, 0.012671640142798424, 0.014554454945027828, 0.005096337758004665, 0.025306345894932747, 0.015522593632340431, 0.012109486386179924, 0.014945329166948795, 0.0111803337931633, 0.010501275770366192, 0.010505528189241886, 0.013426732271909714, 0.01895906589925289, 0.16498495638370514], [0.05249502509832382, 0.3800218403339386, 0.048091597855091095, 0.01820666529238224, 0.10161028057336807, 0.18240275979042053, 0.03954629600048065, 0.08666953444480896, 0.00239415536634624, 0.05545663461089134, 0.11899324506521225, 0.03552442044019699, 0.037884730845689774, 0.08727249503135681, 0.23120805621147156], [0.06818026304244995, 0.06384387612342834, 0.013627037405967712, 0.017488455399870872, 0.04112459346652031, 0.37204819917678833, 0.2269488275051117, 0.050778258591890335, 0.07564288377761841, 0.002337054116651416, 0.03256889060139656, 0.017944803461432457, 0.02268233709037304, 0.05458826571702957, 0.17415940761566162], [0.3350563049316406, 0.14807114005088806, 0.16856855154037476, 0.0634150505065918, 0.6115131974220276, 0.8617944717407227, 0.4784194529056549, 0.271447092294693, 0.44727417826652527, 0.03638387843966484, 0.0791390910744667, 0.0010650564217939973, 0.10882135480642319, 0.07249648869037628, 0.16217634081840515], [0.6229478120803833, 0.11473710834980011, 0.9313594102859497, 0.6977004408836365, 0.7760463953018188, 0.5547962784767151, 0.2850213646888733, 0.12024195492267609, 0.6867435574531555, 0.3715392053127289, 0.5383524894714355, 0.04410971701145172, 0.001209885231219232, 0.03505939990282059, 0.07057712972164154], [0.12039526551961899, 0.15183398127555847, 0.23466746509075165, 0.07534174621105194, 0.09489727020263672, 0.12723755836486816, 0.06088049337267876, 0.06659132242202759, 0.24534910917282104, 0.08624531328678131, 0.05703657865524292, 0.031156441196799278, 0.0026320687029510736, 0.016870809718966484, 0.16136524081230164], [0.024926312267780304, 0.055538877844810486, 0.0035579875111579895, 0.006728078704327345, 0.10179352015256882, 0.12386216968297958, 0.08368373662233353, 0.17138876020908356, 0.13290183246135712, 0.025975322350859642, 0.0007942751399241388, 0.08679928630590439, 0.006940893363207579, 0.006668384652584791, 0.2167840152978897], [0.03079223819077015, 0.008776835165917873, 0.025623725727200508, 0.02996702678501606, 0.076390340924263, 0.11722294241189957, 0.03722265735268593, 0.06894396245479584, 0.023492204025387764, 0.02721765637397766, 0.02432498149573803, 0.009946721605956554, 0.02367306686937809, 0.02709045261144638, 0.15603508055210114], [0.050754088908433914, 0.38707080483436584, 0.056088101118803024, 0.022330837324261665, 0.19594413042068481, 0.356031596660614, 0.05540256202220917, 0.17031489312648773, 0.002592364326119423, 0.0904960110783577, 0.17009596526622772, 0.02688765898346901, 0.05266827344894409, 0.09536514431238174, 0.2306852787733078], [0.052731066942214966, 0.07647765427827835, 0.009669344872236252, 0.013631273992359638, 0.037963252514600754, 0.40968915820121765, 0.1877974420785904, 0.06287717074155807, 0.06925270706415176, 0.0021469732746481895, 0.03106895461678505, 0.02147551439702511, 0.022071314975619316, 0.058794401586055756, 0.17150944471359253], [0.2993965446949005, 0.1887350082397461, 0.17583680152893066, 0.06075390800833702, 0.6836855411529541, 0.8825634121894836, 0.44942814111709595, 0.3110062777996063, 0.6245057582855225, 0.04149743914604187, 0.08928828686475754, 0.0010537458583712578, 0.13885420560836792, 0.09175378829240799, 0.16601231694221497], [0.6222140192985535, 0.13893182575702667, 0.9335290789604187, 0.7374492883682251, 0.8253674507141113, 0.5633905529975891, 0.4091120660305023, 0.12903769314289093, 0.8090996742248535, 0.490604043006897, 0.6206711530685425, 0.06171489879488945, 0.0013746770564466715, 0.055387232452631, 0.07617512345314026], [0.1216169223189354, 0.17628714442253113, 0.21903447806835175, 0.08471400290727615, 0.12100206315517426, 0.12684285640716553, 0.060168445110321045, 0.05725802481174469, 0.204857736825943, 0.07119028270244598, 0.04997517541050911, 0.046147700399160385, 0.002665548352524638, 0.01769380457699299, 0.1595369428396225], [0.02323095127940178, 0.05151251330971718, 0.002836216241121292, 0.007343180477619171, 0.11471041291952133, 0.09745588153600693, 0.08793136477470398, 0.19987791776657104, 0.2081962525844574, 0.026029428467154503, 0.0006721516838297248, 0.15218332409858704, 0.008676346391439438, 0.009503011591732502, 0.20713838934898376], [0.07751920074224472, 0.05964339151978493, 0.026831025257706642, 0.018057459965348244, 0.1489739865064621, 0.27560925483703613, 0.15271086990833282, 0.29336896538734436, 0.2548864185810089, 0.015449506230652332, 0.02643660455942154, 0.05839552357792854, 0.06659974157810211, 0.1841144859790802, 0.1324990689754486]], [[0.006645309738814831, 0.043047573417425156, 0.04108792915940285, 0.028674451634287834, 0.10265154391527176, 0.03326163440942764, 0.05858607590198517, 0.06312219053506851, 0.013714859262108803, 0.017589740455150604, 0.02732386440038681, 0.11026919633150101, 0.028857730329036713, 0.054291173815727234, 0.19011041522026062], [0.006623337976634502, 0.06184479594230652, 0.014693422242999077, 0.03981047496199608, 0.08752858638763428, 0.01962500624358654, 0.06706372648477554, 0.011501927860081196, 0.0061228955164551735, 0.013949333690106869, 0.018435969948768616, 0.03678559139370918, 0.022487374022603035, 0.0660797506570816, 0.28934401273727417], [0.04245300590991974, 0.10349805653095245, 0.03407163918018341, 0.007511724252253771, 0.011565770022571087, 0.010817471891641617, 0.05971734598278999, 0.00459411833435297, 0.00350962788797915, 0.021488210186362267, 0.02298545651137829, 0.06376963108778, 0.036461468786001205, 0.1865386664867401, 0.16962040960788727], [0.014149562455713749, 0.03299444913864136, 0.007003516890108585, 0.004260434303432703, 0.018919609487056732, 0.008522795513272285, 0.018369171768426895, 0.015471882186830044, 0.0008095644298009574, 0.012402600608766079, 0.0075600892305374146, 0.03885417431592941, 0.05682341009378433, 0.0525624044239521, 0.22132590413093567], [0.01582285761833191, 0.013434984721243382, 0.0299182441085577, 0.03647983819246292, 0.009840411134064198, 0.06101881340146065, 0.04943924769759178, 0.3809337913990021, 0.027872184291481972, 0.07177315652370453, 0.06987256556749344, 0.014244881458580494, 0.18650749325752258, 0.16280896961688995, 0.16209137439727783], [0.018014581874012947, 0.11459828168153763, 0.013770120218396187, 0.021584663540124893, 0.02155740186572075, 0.03133949637413025, 0.03938381373882294, 0.28105995059013367, 0.02592163160443306, 0.026603924110531807, 0.010026685893535614, 0.009953479282557964, 0.004658891819417477, 0.014652709476649761, 0.16460371017456055], [0.001359884045086801, 0.029354294762015343, 0.0013457777677103877, 0.0026418184861540794, 0.008543581701815128, 0.003654624568298459, 0.0034977763425558805, 0.039957791566848755, 0.00108401442412287, 0.0005604945472441614, 0.0003877367707900703, 0.0033066808246076107, 0.007358025759458542, 0.007617549039423466, 0.20286646485328674], [0.015068605542182922, 0.027786174789071083, 0.015096615999937057, 0.048349082469940186, 0.03296791389584541, 0.0033369800075888634, 0.004459223244339228, 0.01348987128585577, 0.0010384898632764816, 0.013556106016039848, 0.015940798446536064, 0.042712315917015076, 0.02055070362985134, 0.042082786560058594, 0.17761820554733276], [0.09032934159040451, 0.007927155122160912, 0.08835490047931671, 0.21186837553977966, 0.05379607528448105, 0.23637458682060242, 0.16646702587604523, 0.022663533687591553, 0.024165447801351547, 0.08468358218669891, 0.07286331057548523, 0.016201749444007874, 0.031014403328299522, 0.026781529188156128, 0.21159759163856506], [0.014649872668087482, 0.032003261148929596, 0.1914098560810089, 0.17710277438163757, 0.07542474567890167, 0.05287592485547066, 0.14732114970684052, 0.08320016413927078, 0.025441674515604973, 0.02800501137971878, 0.0780739113688469, 0.04154554009437561, 0.017996925860643387, 0.08907850831747055, 0.17056028544902802], [0.29397615790367126, 0.03400568664073944, 0.3242063522338867, 0.3681035339832306, 0.48163339495658875, 0.025333818048238754, 0.20042747259140015, 0.06051841378211975, 0.2913966476917267, 0.19229580461978912, 0.12739360332489014, 0.07057002186775208, 0.012750222347676754, 0.053084854036569595, 0.09877952188253403], [0.2290111482143402, 0.04351853206753731, 0.4067046046257019, 0.12047477811574936, 0.3140789866447449, 0.03630740940570831, 0.1768438071012497, 0.13207398355007172, 0.0676346942782402, 0.07621245086193085, 0.1797569841146469, 0.24804529547691345, 0.009716867469251156, 0.01671340875327587, 0.15996301174163818], [0.0448942668735981, 0.015721717849373817, 0.04864601418375969, 0.03494936227798462, 0.016112152487039566, 0.06668571382761002, 0.05302642658352852, 0.07182876765727997, 0.006946365814656019, 0.011091585271060467, 0.1120418831706047, 0.008756275288760662, 0.055249348282814026, 0.03253563493490219, 0.187040314078331], [0.3104230761528015, 0.04545353353023529, 0.3986057937145233, 0.6762936115264893, 0.03838818892836571, 0.03300129249691963, 0.27034318447113037, 0.21517230570316315, 0.008858010172843933, 0.2650390863418579, 0.2720700800418854, 0.005442188587039709, 0.06764175742864609, 0.053534120321273804, 0.18754751980304718], [0.011383982375264168, 0.11127021163702011, 0.0030386100988835096, 0.0067845494486391544, 0.013927198015153408, 0.08719860762357712, 0.03287587687373161, 0.5690041184425354, 0.03855481743812561, 0.020931608974933624, 0.01293823029845953, 0.047187648713588715, 0.021772168576717377, 0.1471272110939026, 0.18776896595954895], [0.005892250686883926, 0.03474593162536621, 0.023128867149353027, 0.002957691205665469, 0.03212961554527283, 0.015600761398673058, 0.0076070488430559635, 0.04006163775920868, 0.012522950768470764, 0.00397108681499958, 0.004476191475987434, 0.01931026391685009, 0.006290406920015812, 0.014653924852609634, 0.17843826115131378], [0.030382098630070686, 0.14396639168262482, 0.0023552696220576763, 0.003069670405238867, 0.03293609246611595, 0.010766614228487015, 0.04698408767580986, 0.0892328992486, 0.010764017701148987, 0.01645551063120365, 0.0007101192022673786, 0.14693684875965118, 0.10194381326436996, 0.06734117865562439, 0.21650707721710205], [0.11579495668411255, 0.04704239219427109, 0.08932461589574814, 0.10469675809144974, 0.3945455551147461, 0.10528933256864548, 0.15413445234298706, 0.13012593984603882, 0.37207290530204773, 0.07726370543241501, 0.08641648292541504, 0.07665102183818817, 0.02378079853951931, 0.06452124565839767, 0.12331708520650864], [0.20921318233013153, 0.07137931883335114, 0.3537597060203552, 0.1065746620297432, 0.30610421299934387, 0.07002534717321396, 0.22329437732696533, 0.23702743649482727, 0.06014438346028328, 0.05975072830915451, 0.17522762715816498, 0.3013332188129425, 0.02163097821176052, 0.016774384304881096, 0.15580035746097565], [0.037447404116392136, 0.022215796634554863, 0.033449236303567886, 0.026462113484740257, 0.01563168875873089, 0.07434160262346268, 0.05695066228508949, 0.11209315806627274, 0.007291351445019245, 0.008904322981834412, 0.08964232355356216, 0.01435061078518629, 0.07215401530265808, 0.030404584482312202, 0.17889626324176788], [0.35028940439224243, 0.06261257082223892, 0.400876522064209, 0.6601436138153076, 0.0364767424762249, 0.0348673090338707, 0.3584212362766266, 0.3042086958885193, 0.012779565528035164, 0.3784087598323822, 0.29859334230422974, 0.00785628892481327, 0.11913719773292542, 0.06971576809883118, 0.17937220633029938], [0.014627714641392231, 0.1739588975906372, 0.0033204040955752134, 0.007496224716305733, 0.011711684986948967, 0.10170583426952362, 0.050673384219408035, 0.6495208740234375, 0.040652137249708176, 0.03492900729179382, 0.01829371228814125, 0.07074988633394241, 0.02588740922510624, 0.18312060832977295, 0.1794223189353943], [0.006626310292631388, 0.049714479595422745, 0.02355029061436653, 0.0033578642178326845, 0.02970620058476925, 0.020507775247097015, 0.008351391181349754, 0.03789898753166199, 0.008593969978392124, 0.004206442274153233, 0.004605707712471485, 0.02678176388144493, 0.006028715055435896, 0.012980426661670208, 0.1725957691669464], [0.029822910204529762, 0.18419219553470612, 0.002088941168040037, 0.00302593014203012, 0.028257815167307854, 0.012486547231674194, 0.051940228790044785, 0.10161811858415604, 0.01137576438486576, 0.02022942155599594, 0.0007436276064254344, 0.2113851010799408, 0.1359580010175705, 0.08821411430835724, 0.2053057849407196], [0.016353517770767212, 0.03170220926403999, 0.014149405062198639, 0.013441388495266438, 0.037340469658374786, 0.010170645080506802, 0.0053974115289747715, 0.025274941697716713, 0.017184404656291008, 0.0020940443500876427, 0.006704597268253565, 0.009430822916328907, 0.030376460403203964, 0.024553189054131508, 0.15533798933029175]], [[0.005564282648265362, 0.001319661969318986, 0.028383644297719002, 0.01146539393812418, 0.028919272124767303, 0.012663042172789574, 0.023019153624773026, 0.0018097365973517299, 0.0143426563590765, 0.021044740453362465, 0.015969598665833473, 0.03200899809598923, 0.013908782042562962, 0.03448842838406563, 0.20206299424171448], [0.3364894986152649, 0.00033270660787820816, 0.017299778759479523, 0.02505551464855671, 0.00914769060909748, 0.0018482855521142483, 0.040363892912864685, 0.0008854345069266856, 0.020481230691075325, 0.022734129801392555, 0.016724254935979843, 0.0011141380527988076, 5.783090819022618e-05, 0.0005799515638500452, 0.07228588312864304], [0.0004661931307055056, 0.4122284948825836, 0.0022180580999702215, 0.00018468582129571587, 0.00030452435021288693, 5.825214248034172e-05, 0.0012309255544096231, 0.0017770789563655853, 1.19774986160337e-05, 0.0001907332189148292, 0.0007099026697687805, 0.0006694658659398556, 1.216385771840578e-05, 0.00011785236711148173, 0.00036971797817386687], [0.04950903728604317, 0.2967310845851898, 0.021222729235887527, 0.01289455872029066, 0.009955117478966713, 0.008917939849197865, 0.011312013491988182, 0.01272521447390318, 0.0006359940161928535, 0.011413054540753365, 0.006479735020548105, 0.0053005279041826725, 0.001741865067742765, 0.0027997863944619894, 0.08213357627391815], [0.020872987806797028, 3.087984805461019e-05, 0.009670623578131199, 0.0253498163074255, 0.010817835107445717, 0.4320962131023407, 0.017970044165849686, 0.0021109851077198982, 0.0003069202939514071, 0.008261006325483322, 0.006166533567011356, 0.7898750901222229, 0.11304597556591034, 0.12737329304218292, 0.011856237426400185], [0.06067817285656929, 0.005839335732161999, 0.025896329432725906, 0.03351203724741936, 0.025002295151352882, 0.25514867901802063, 0.4275963008403778, 0.0194717925041914, 0.0888834074139595, 0.04690318927168846, 0.03570560738444328, 0.0850825086236, 0.0388353131711483, 0.24394167959690094, 0.10019046813249588], [0.014415884390473366, 0.001141559099778533, 0.0678224116563797, 0.024646559730172157, 0.08796916157007217, 0.022639306262135506, 0.07784608006477356, 0.02605922892689705, 0.014093886129558086, 0.0286162830889225, 0.09674176573753357, 0.04692256450653076, 0.03519048914313316, 0.20982496440410614, 0.1800668090581894], [0.02086471952497959, 0.0008324789232574403, 0.01815967448055744, 0.002886975882574916, 0.0020961007103323936, 0.004472428001463413, 0.033020272850990295, 0.0047500282526016235, 0.012928733602166176, 0.014328529126942158, 0.015946470201015472, 0.06593997031450272, 0.00855537410825491, 0.07526978105306625, 0.1768130511045456], [0.0009654826717451215, 0.000225315525312908, 0.0006124225910753012, 0.0007836261647753417, 0.0007428302778862417, 0.003282200777903199, 0.008662715554237366, 0.45239004492759705, 4.857195381191559e-05, 0.0006357804522849619, 0.0010122592793777585, 0.0006606358801946044, 0.00025698603712953627, 0.0011707579251378775, 0.0028539940249174833], [0.0025523374788463116, 0.0009212270379066467, 0.09748471528291702, 0.057154957205057144, 0.4982932209968567, 0.000552327954210341, 0.02918482944369316, 0.0039253802970051765, 0.00450148293748498, 0.0014971394557505846, 0.009822547435760498, 0.0017059196252375841, 0.001570553402416408, 0.005804183427244425, 0.00957300141453743], [0.016401896253228188, 0.00043752315104939044, 0.0039018490351736546, 0.005885160993784666, 0.0023499932140111923, 0.0031332974322140217, 0.055512603372335434, 0.003903925186023116, 0.10197419673204422, 0.009071548469364643, 0.023729920387268066, 0.002627716166898608, 0.01914973370730877, 0.02837507426738739, 0.1623656302690506], [0.0004865071678068489, 2.4051656509982422e-05, 0.00020084556308574975, 0.0003736558719538152, 0.000646126689389348, 9.209318523062393e-05, 0.009753170423209667, 9.854567178990692e-05, 0.34485483169555664, 0.00047165394062176347, 0.0012700805673375726, 0.000479432987049222, 0.0015819557011127472, 0.0008011643076315522, 0.0017131956992670894], [0.03442303463816643, 0.014513631351292133, 0.003174385754391551, 0.00478995218873024, 0.0017101461999118328, 0.003900717245414853, 0.05713852494955063, 0.013628470711410046, 0.0976317971944809, 0.28217896819114685, 0.01894235610961914, 0.009533336386084557, 0.003816690994426608, 0.005922130309045315, 0.12864208221435547], [0.01004086248576641, 0.01997406780719757, 0.005450551863759756, 0.006583535112440586, 0.0027623113710433245, 0.002903316868469119, 0.03531726077198982, 0.008635452017188072, 0.029197845607995987, 0.02162068709731102, 0.013219092041254044, 0.2711889445781708, 0.00537630682811141, 0.006846235599368811, 0.06079954653978348], [0.00031272557680495083, 8.196506314561702e-06, 4.237617031321861e-05, 0.00043677922803908587, 0.00024717405904084444, 0.022641032934188843, 0.002573953475803137, 0.0004433683061506599, 0.0013428670354187489, 0.00034036010038107634, 0.0007929583080112934, 0.0033021108247339725, 0.4761846959590912, 0.05593165382742882, 0.00081905338447541], [0.00267792004160583, 4.751862070406787e-05, 0.014043050818145275, 0.02037942036986351, 0.04410611465573311, 0.04370833560824394, 0.06117184832692146, 0.01571183279156685, 0.11117196083068848, 0.006906491704285145, 0.0029646854382008314, 0.15407170355319977, 0.010935205966234207, 0.03797803074121475, 0.16977860033512115], [0.011722833849489689, 0.005004812031984329, 0.007801789790391922, 0.0020204312168061733, 0.004946417640894651, 0.000467105332063511, 0.11018845438957214, 0.016256244853138924, 0.05208335816860199, 0.08122430741786957, 0.4447634816169739, 0.0032620911952108145, 0.0036480925045907497, 0.02699565887451172, 0.038189876824617386], [0.024071840569376945, 0.0004321316082496196, 0.023504342883825302, 0.020648522302508354, 0.021508874371647835, 0.012214796617627144, 0.024360070005059242, 0.0013747027842327952, 0.0815734788775444, 0.08039785921573639, 0.06951787322759628, 0.017521949484944344, 0.04566040262579918, 0.08389204740524292, 0.15396325290203094], [0.0014979105908423662, 4.0405931940767914e-05, 0.0008743218495510519, 0.001329930848442018, 0.0032007889822125435, 0.0002464030694682151, 0.015361684374511242, 0.00014017200737725943, 0.3369258642196655, 0.0015512423124164343, 0.003011554479598999, 0.0010034784208983183, 0.0037561107892543077, 0.0018123533809557557, 0.0037892721593379974], [0.03386643901467323, 0.015328249894082546, 0.002211565151810646, 0.003828595858067274, 0.0012934240512549877, 0.004837968852370977, 0.04463785141706467, 0.014559985138475895, 0.04106945917010307, 0.26340487599372864, 0.017707379534840584, 0.01015215553343296, 0.0033097255509346724, 0.0058202859945595264, 0.13427288830280304], [0.011043943464756012, 0.029788998886942863, 0.004548549186438322, 0.006417197175323963, 0.0019613932818174362, 0.0028304944280534983, 0.02768276073038578, 0.006805655546486378, 0.02553243562579155, 0.0314837321639061, 0.015709027647972107, 0.2568790316581726, 0.008081428706645966, 0.009137820452451706, 0.06746803224086761], [0.0003306480939500034, 1.1417017958592623e-05, 3.816767639364116e-05, 0.000435528316302225, 0.00020690191013272852, 0.02179853804409504, 0.002864222740754485, 0.0005160043947398663, 0.001080053043551743, 0.0004847492673434317, 0.0009861867874860764, 0.003908392507582903, 0.47703394293785095, 0.07113853842020035, 0.000873323529958725], [0.0030808241572231054, 6.38188939774409e-05, 0.011707174591720104, 0.023645061999559402, 0.038246914744377136, 0.047200631350278854, 0.04958858713507652, 0.012573646381497383, 0.04961754009127617, 0.005252092145383358, 0.002489157486706972, 0.17429526150226593, 0.008030706085264683, 0.02717452496290207, 0.1679786741733551], [0.01455691922456026, 0.008012487553060055, 0.006938801147043705, 0.00259140832349658, 0.004911262542009354, 0.0004763725446537137, 0.10579084604978561, 0.021042171865701675, 0.03971559554338455, 0.07511086016893387, 0.43185338377952576, 0.0035418386105448008, 0.004437423776835203, 0.03184036538004875, 0.04226255044341087], [0.055085837841033936, 0.014846320264041424, 0.06939522176980972, 0.036867137998342514, 0.13156765699386597, 0.04343622922897339, 0.18117153644561768, 0.04244613274931908, 0.04596249759197235, 0.13158053159713745, 0.047130946069955826, 0.549620509147644, 0.24813801050186157, 0.3232562243938446, 0.11823604255914688]], [[0.7448275089263916, 0.00023065913410391659, 0.0003700565139297396, 0.0002745355886872858, 0.0005768057890236378, 1.0151054993912112e-05, 1.3715341992792673e-05, 7.643950084457174e-06, 0.0004341531603131443, 5.2913601393811405e-05, 5.353476808522828e-05, 8.812115265754983e-05, 1.1566834245968494e-06, 5.744800546381157e-06, 5.576572584686801e-05], [8.114575030049309e-05, 0.06691394746303558, 0.04036417603492737, 0.022258125245571136, 0.055233534425497055, 0.050445422530174255, 0.048324622213840485, 0.00889397319406271, 0.1270352452993393, 0.04156908392906189, 0.20929713547229767, 0.21122632920742035, 0.414194792509079, 0.12628954648971558, 0.25567519664764404], [0.0012628535041585565, 0.0008597301202826202, 0.036364536732435226, 0.0971999391913414, 0.04217860475182533, 0.10421664267778397, 0.16082510352134705, 0.03283625468611717, 0.09032318741083145, 0.09653837233781815, 0.21890851855278015, 0.06589526683092117, 0.47985169291496277, 0.21388037502765656, 0.21010825037956238], [0.0002990703214891255, 0.001862871926277876, 0.010526847094297409, 0.01025421917438507, 0.05592086538672447, 0.02697981521487236, 0.01570008136332035, 0.02568165771663189, 0.010194454342126846, 0.048093631863594055, 0.04421652480959892, 0.02353351190686226, 0.21245922148227692, 0.0448865108191967, 0.23352482914924622], [0.00015855174569878727, 0.013162538409233093, 0.006567019037902355, 0.004201928153634071, 0.006268346216529608, 0.00024757537175901234, 0.012954139150679111, 0.003747382666915655, 0.03740423545241356, 0.007960616610944271, 0.013323514722287655, 0.06273993849754333, 0.048431456089019775, 0.13987915217876434, 0.20342004299163818], [0.013553211465477943, 0.03824196010828018, 0.02278091199696064, 0.09299258887767792, 0.0559159517288208, 0.00022306715254671872, 0.031003709882497787, 0.010444254614412785, 0.16168788075447083, 0.03666102886199951, 0.00852662418037653, 0.4432809352874756, 0.009321487508714199, 0.024379035457968712, 0.17351986467838287], [0.00026768012321554124, 0.015254812315106392, 0.007090381346642971, 0.006173381581902504, 0.006773150525987148, 0.0008773274021223187, 0.00638232659548521, 0.016591282561421394, 0.004996343981474638, 0.009327422827482224, 0.008862738497555256, 0.05876166746020317, 0.009527520276606083, 0.00578573253005743, 0.20356230437755585], [0.0008312691352330148, 0.012717761099338531, 0.013986560516059399, 0.007093494758009911, 0.004876464139670134, 0.0027259632479399443, 0.0033886858727782965, 0.01589561626315117, 0.00876854918897152, 0.005017295014113188, 0.023178039118647575, 0.05755693465471268, 0.05451130494475365, 0.06928746402263641, 0.1796484887599945], [0.00016753048112150282, 0.011822681874036789, 0.005686081480234861, 0.011659285984933376, 0.004307762254029512, 0.0031254058703780174, 0.009316416457295418, 0.0016170619055628777, 0.012603488750755787, 0.0245236624032259, 0.01756892167031765, 0.011099276132881641, 0.11892349272966385, 0.02075323462486267, 0.2549600899219513], [0.00017647366621531546, 0.053185176104307175, 0.007304554805159569, 0.004834755789488554, 0.000954066461417824, 0.025718921795487404, 0.02985404059290886, 0.09960591793060303, 0.010695043951272964, 0.016483109444379807, 0.018774237483739853, 0.05090473219752312, 0.01008983701467514, 0.028674444183707237, 0.22871088981628418], [0.0008755451999604702, 0.020039640367031097, 0.003969491925090551, 0.007670485880225897, 0.006173306610435247, 0.012295764870941639, 0.0076020946726202965, 0.012137084268033504, 0.010956642217934132, 0.010541083291172981, 0.018125493079423904, 0.03226908668875694, 0.02587633579969406, 0.016216130927205086, 0.1660052388906479], [5.4335410823114216e-05, 0.03367479890584946, 0.004507457371801138, 0.004544241353869438, 0.00623831432312727, 0.002192543353885412, 0.004128816071897745, 0.021106822416186333, 0.0003909784718416631, 0.00830051489174366, 0.018183842301368713, 0.009683135896921158, 0.0325237475335598, 0.00792472343891859, 0.25227075815200806], [0.0006012204103171825, 0.01188816037029028, 0.023532994091510773, 0.00770517997443676, 0.007410787045955658, 0.007087987381964922, 0.021027186885476112, 0.013456426560878754, 0.03266710042953491, 0.001251929672434926, 0.09021235257387161, 0.024440091103315353, 0.024299103766679764, 0.02338516153395176, 0.1967199146747589], [0.0009616355528123677, 0.059039004147052765, 0.04997482895851135, 0.013552234508097172, 0.03981975466012955, 0.020335622131824493, 0.014380398206412792, 0.07606764137744904, 0.07161007821559906, 0.024130970239639282, 0.06891870498657227, 0.0008635766571387649, 0.023193923756480217, 0.02981526218354702, 0.21020111441612244], [0.0013424595817923546, 0.0746709555387497, 0.011544802226126194, 0.027912717312574387, 0.0729047879576683, 0.10483764857053757, 0.07119728624820709, 0.010606798343360424, 0.044552259147167206, 0.05723145231604576, 0.034647323191165924, 0.38214871287345886, 0.003923356998711824, 0.08778946846723557, 0.19581711292266846], [0.0016638260567560792, 0.01581355184316635, 0.08943041414022446, 0.02092832513153553, 0.021133122965693474, 0.012408973649144173, 0.01347691286355257, 0.00275444146245718, 0.027862150222063065, 0.01225491613149643, 0.018322426825761795, 0.008929668925702572, 0.00015579524915665388, 0.0014782899525016546, 0.18181975185871124], [0.0008640239248052239, 0.06174946948885918, 0.004653214477002621, 0.002717669354751706, 0.015129820443689823, 0.00935456808656454, 0.016078660264611244, 0.08089328557252884, 0.017857585102319717, 0.0025031790137290955, 0.00012101473839720711, 0.013123439624905586, 0.005499868653714657, 0.001559562049806118, 0.22764776647090912], [0.0008687095833010972, 0.025285501033067703, 0.01658034697175026, 0.02363765239715576, 0.02393241412937641, 0.0657346174120903, 0.015298763290047646, 0.01792113669216633, 0.021707117557525635, 0.018967296928167343, 0.037634264677762985, 0.013209421187639236, 0.02256513573229313, 0.007774183992296457, 0.15961462259292603], [0.0001073219973477535, 0.04253393039107323, 0.010077103972434998, 0.007349912542849779, 0.00879223458468914, 0.004757148679345846, 0.008167163468897343, 0.03753674402832985, 0.00042728587868623435, 0.014237778261303902, 0.029898250475525856, 0.006872681900858879, 0.045794516801834106, 0.007500257343053818, 0.2562271058559418], [0.0005320480559021235, 0.010701313614845276, 0.020972738042473793, 0.007364482618868351, 0.006165153346955776, 0.00950621161609888, 0.022682208567857742, 0.018515970557928085, 0.03319491446018219, 0.00125269521959126, 0.07773777842521667, 0.022826068103313446, 0.02051766775548458, 0.020874740555882454, 0.1872510462999344], [0.0008804904646240175, 0.05573932081460953, 0.06578188389539719, 0.01897181011736393, 0.043492771685123444, 0.026308609172701836, 0.016426166519522667, 0.09104844927787781, 0.12495335191488266, 0.04637341946363449, 0.0944451242685318, 0.0008321930072270334, 0.03243781998753548, 0.03530845418572426, 0.2013196051120758], [0.001610875129699707, 0.08435038477182388, 0.014167247340083122, 0.03493078798055649, 0.07050123810768127, 0.10772886872291565, 0.09850788861513138, 0.013066386803984642, 0.05027954652905464, 0.10465669631958008, 0.04533415287733078, 0.47037968039512634, 0.004505114629864693, 0.12196572870016098, 0.18816377222537994], [0.0018758929800242186, 0.019657986238598824, 0.1020394116640091, 0.033738646656274796, 0.024869924411177635, 0.012215637601912022, 0.015038376674056053, 0.002843664726242423, 0.02175789885222912, 0.01636381261050701, 0.01989913359284401, 0.01190999522805214, 0.00020280842727515846, 0.0016855570720508695, 0.17570628225803375], [0.0009206020040437579, 0.08179444819688797, 0.00436751963570714, 0.003652991494163871, 0.019383452832698822, 0.008280212059617043, 0.016885409131646156, 0.10377784073352814, 0.023152435198426247, 0.0037028237711638212, 0.0001251623034477234, 0.018928401172161102, 0.009926089085638523, 0.002465219935402274, 0.21539123356342316], [0.0005496710073202848, 0.039492249488830566, 0.016358638182282448, 0.007983607240021229, 0.006420070305466652, 0.0012171968119218946, 0.003928476013243198, 0.005028040148317814, 0.010722441598773003, 0.0025004756171256304, 0.015696601942181587, 0.006085758097469807, 0.0033880609553307295, 0.0056163351982831955, 0.1572248637676239]], [[0.09555985033512115, 0.6603901982307434, 0.4109249413013458, 0.6857163310050964, 0.16377028822898865, 0.1341286301612854, 0.19969937205314636, 0.28269705176353455, 0.14764364063739777, 0.41980865597724915, 0.4319525361061096, 0.3789142668247223, 0.49345141649246216, 0.26345306634902954, 0.00909768883138895], [0.1460653841495514, 0.2758752405643463, 0.2826583981513977, 0.551855206489563, 0.05612415447831154, 0.19304026663303375, 0.0849798247218132, 0.038316093385219574, 0.02312053181231022, 0.46154478192329407, 0.36433619260787964, 0.35877159237861633, 0.1596277803182602, 0.0554661750793457, 6.483463948825374e-05], [3.716628270922229e-05, 1.9402585849093157e-07, 1.0113188182003796e-05, 6.318590021692216e-05, 6.053787728887983e-07, 2.5790013751247898e-06, 0.00022986173280514777, 1.074662236533186e-06, 6.082240361138247e-06, 3.35614299729059e-06, 2.225729804194998e-05, 7.863033715693746e-06, 1.555537892272696e-06, 3.881560041918419e-05, 0.23657216131687164], [0.6150763630867004, 0.041665952652692795, 0.4174444377422333, 0.4949702024459839, 0.20794649422168732, 0.3307763934135437, 0.8098993897438049, 0.2721010744571686, 0.7274996042251587, 0.4779607057571411, 0.6233283281326294, 0.7560765147209167, 0.3628612458705902, 0.7672091722488403, 5.392584171204362e-06], [5.640763447445352e-06, 2.5884469323500525e-07, 1.2724142379738623e-06, 8.170181899913587e-06, 1.2345621769327408e-07, 1.310836523771286e-07, 1.02673438959755e-05, 9.661080184741877e-07, 6.520539272969472e-07, 7.602448022225872e-07, 2.058099425994442e-06, 6.885502301656743e-08, 1.0175665465794737e-06, 1.7383708836860023e-05, 0.20754273235797882], [9.27566077280062e-07, 5.395870630309219e-07, 1.8455818917573197e-07, 1.2775643654094893e-06, 2.105696061960316e-08, 3.1680112755338996e-08, 6.263408067752607e-06, 4.3284012463118415e-07, 1.918825773827848e-06, 1.694104128091567e-07, 3.363936968980852e-07, 9.135120215830739e-09, 4.4058825920956224e-08, 7.840970965844463e-07, 0.18219269812107086], [0.7144812345504761, 0.6739043593406677, 0.2952970862388611, 0.49478814005851746, 0.17151717841625214, 0.06989942491054535, 0.5132517218589783, 0.30886489152908325, 0.5621734261512756, 0.5728412866592407, 0.576314389705658, 0.34687095880508423, 0.25617536902427673, 0.29690253734588623, 7.371841547865188e-06], [0.6291437745094299, 0.5982875823974609, 0.4885888695716858, 0.5792520046234131, 0.2514877915382385, 0.5298613905906677, 0.11972777545452118, 0.6076628565788269, 0.04243328422307968, 0.5940482020378113, 0.6775911450386047, 0.3496588468551636, 0.4937344789505005, 0.40163323283195496, 2.9517783332266845e-05], [0.6414378881454468, 0.20530864596366882, 0.8448930978775024, 0.5841984748840332, 0.48009997606277466, 0.48003992438316345, 0.4468145966529846, 0.036266062408685684, 0.3466547429561615, 0.521195650100708, 0.7532409429550171, 0.14529024064540863, 0.3844791650772095, 0.46825459599494934, 2.1059213395346887e-05], [0.7977450489997864, 0.5162288546562195, 0.513008177280426, 0.6203657984733582, 0.04621165990829468, 0.2237500697374344, 0.10730908066034317, 0.17203836143016815, 0.028481170535087585, 0.5342445969581604, 0.7256113290786743, 0.5827998518943787, 0.755642294883728, 0.511749804019928, 0.00015279543003998697], [0.5001324415206909, 0.7283154129981995, 0.6225411295890808, 0.5096700191497803, 0.4470505714416504, 0.6475648880004883, 0.4919697046279907, 0.42729777097702026, 0.22966071963310242, 0.4533919394016266, 0.5539101958274841, 0.2698501944541931, 0.3532210886478424, 0.2643750309944153, 2.9741322578047402e-05], [0.42266348004341125, 0.20205438137054443, 0.42841264605522156, 0.6724829077720642, 0.29094210267066956, 0.4464052617549896, 0.24126748740673065, 0.22405968606472015, 0.21308888494968414, 0.3085091710090637, 0.4672502279281616, 0.14604215323925018, 0.09687051922082901, 0.12085973471403122, 2.7047781259170733e-05], [0.5077533721923828, 0.4866065979003906, 0.8742184638977051, 0.805268406867981, 0.8406472206115723, 0.45863693952560425, 0.3596036732196808, 0.36316972970962524, 0.38783764839172363, 0.03767421096563339, 0.43841618299484253, 0.3401361405849457, 0.3197961747646332, 0.20812755823135376, 7.5720936365542e-06], [0.12348711490631104, 0.49926623702049255, 0.1342328041791916, 0.07936512678861618, 0.11133208125829697, 0.032334309071302414, 0.028592387214303017, 0.036310840398073196, 0.036252155900001526, 0.10585709661245346, 0.19267472624778748, 0.34429997205734253, 0.16909800469875336, 0.2464863359928131, 3.1697504709882196e-06], [4.5035082507638435e-07, 4.8253248507990065e-08, 2.1990938847693542e-08, 4.3766593194050074e-07, 1.1283042766763174e-07, 2.4235429663121977e-08, 4.6985369408503175e-06, 1.5805973418991925e-07, 1.1619090578562918e-08, 1.9516033233912822e-08, 1.8456361772223318e-07, 2.2261544074808626e-07, 2.278205402106437e-09, 7.143006541809882e-07, 0.21044957637786865], [0.71169513463974, 0.2780396640300751, 0.44078493118286133, 0.7963916063308716, 0.6933308839797974, 0.5056049823760986, 0.7329073548316956, 0.810703694820404, 0.551677942276001, 0.6459015607833862, 0.6943050622940063, 0.2817550301551819, 0.10247289389371872, 0.7378624677658081, 8.274764695670456e-06], [0.723514199256897, 0.08602748066186905, 0.6093902587890625, 0.8655006289482117, 0.42677831649780273, 0.03823491558432579, 0.30262306332588196, 0.036271825432777405, 0.12300263345241547, 0.2776595950126648, 0.07632125169038773, 0.06917709112167358, 0.14498986303806305, 0.06881040334701538, 2.5871422622003593e-06], [0.7111753225326538, 0.8019941449165344, 0.7984396815299988, 0.6959745287895203, 0.34880974888801575, 0.5955101251602173, 0.6658092141151428, 0.5378626585006714, 0.35595381259918213, 0.5855972766876221, 0.5757258534431458, 0.133575439453125, 0.3884122669696808, 0.11617641150951385, 8.579120731155854e-06], [0.43439850211143494, 0.1714652180671692, 0.4214288294315338, 0.6560039520263672, 0.15961043536663055, 0.25604698061943054, 0.26937225461006165, 0.1702796220779419, 0.22940081357955933, 0.327440470457077, 0.3977930247783661, 0.08873222768306732, 0.13160161674022675, 0.07058954238891602, 2.3103428247850388e-05], [0.48717519640922546, 0.4504354000091553, 0.9026078581809998, 0.8262973427772522, 0.8697957992553711, 0.4322546720504761, 0.47440072894096375, 0.40584686398506165, 0.6554202437400818, 0.04447361081838608, 0.5114831924438477, 0.4020007252693176, 0.3586147725582123, 0.19603849947452545, 5.424046776170144e-06], [0.09346597641706467, 0.41046077013015747, 0.13097965717315674, 0.06711046397686005, 0.09538185596466064, 0.021688319742679596, 0.027864748612046242, 0.029869627207517624, 0.07506763935089111, 0.13717295229434967, 0.21322546899318695, 0.3559926152229309, 0.19059841334819794, 0.24045485258102417, 2.0756003777933074e-06], [4.6634454520244617e-07, 5.573102512812511e-08, 2.3018172257138758e-08, 3.889360016273713e-07, 9.709493298259986e-08, 2.4796046105279856e-08, 7.192591056082165e-06, 1.7916640615567303e-07, 1.8580767147113875e-08, 3.5935642017648206e-08, 2.774728216081712e-07, 3.801677337378351e-07, 2.8816848907098347e-09, 9.808413778955583e-07, 0.2028982788324356], [0.6667957305908203, 0.327456533908844, 0.4202725291252136, 0.7458598613739014, 0.6837785840034485, 0.5435037612915039, 0.7794858813285828, 0.849186360836029, 0.6942030787467957, 0.7531007528305054, 0.7604266405105591, 0.4857816696166992, 0.12311270833015442, 0.7958275079727173, 7.400509275612421e-06], [0.704485297203064, 0.08825523406267166, 0.5944071412086487, 0.8510531783103943, 0.4262540936470032, 0.04518446326255798, 0.38849392533302307, 0.055145543068647385, 0.277063250541687, 0.40566664934158325, 0.09198901802301407, 0.13750647008419037, 0.24822941422462463, 0.1165834292769432, 3.5331499930180144e-06], [0.5231692790985107, 0.6706213355064392, 0.7785398364067078, 0.7122241258621216, 0.34260621666908264, 0.579698920249939, 0.5863306522369385, 0.4822496175765991, 0.5804131031036377, 0.7801564335823059, 0.7983464002609253, 0.22512593865394592, 0.4790371060371399, 0.2274763584136963, 1.8860177078749985e-05]], [[0.12044757604598999, 0.22699733078479767, 0.3625817894935608, 0.18942511081695557, 0.468371719121933, 0.5971034169197083, 0.5581120252609253, 0.29680517315864563, 0.4773823618888855, 0.4035939574241638, 0.3702273666858673, 0.3751682937145233, 0.267861545085907, 0.4069889783859253, 0.040672045201063156], [0.0243044663220644, 0.4273812174797058, 0.5286219716072083, 0.05566978082060814, 0.4582313597202301, 0.5064847469329834, 0.09591992199420929, 0.1787465512752533, 0.7349562644958496, 0.00692495983093977, 0.04355573281645775, 0.04027868062257767, 0.03415951877832413, 0.02788657508790493, 0.03653726726770401], [0.1999487727880478, 0.02213704027235508, 0.750217854976654, 0.5677059292793274, 0.8556592464447021, 0.6869031190872192, 0.2201639711856842, 0.6947058439254761, 0.2711787521839142, 0.21462410688400269, 0.3783731162548065, 0.39328378438949585, 0.3796219229698181, 0.27560317516326904, 0.052095912396907806], [0.17733721435070038, 0.1195838525891304, 0.4294462502002716, 0.41039443016052246, 0.45686641335487366, 0.5433338284492493, 0.08341590315103531, 0.5749803781509399, 0.0773383378982544, 0.2876206338405609, 0.19534848630428314, 0.10015372186899185, 0.2102438062429428, 0.04678432643413544, 0.044711172580718994], [0.4523387849330902, 0.8917949795722961, 0.4903220534324646, 0.5869925022125244, 0.47626572847366333, 0.006232858635485172, 0.41125378012657166, 0.13404546678066254, 0.6460333466529846, 0.32553666830062866, 0.3429105877876282, 0.031081799417734146, 0.42998504638671875, 0.16709895431995392, 0.08821719139814377], [0.49767979979515076, 0.7566660642623901, 0.25263193249702454, 0.4967457056045532, 0.47193706035614014, 0.006824302952736616, 0.2858791947364807, 0.18135732412338257, 0.4390898644924164, 0.7668571472167969, 0.15391138195991516, 0.08414287865161896, 0.5640745759010315, 0.35628020763397217, 0.09142898768186569], [0.18697474896907806, 0.23196713626384735, 0.23554784059524536, 0.34321168065071106, 0.5325552225112915, 0.15430577099323273, 0.2887123227119446, 0.4957616627216339, 0.36584702134132385, 0.2891024053096771, 0.08069057762622833, 0.18119029700756073, 0.4536079466342926, 0.16425864398479462, 0.03777371346950531], [0.17079660296440125, 0.16765500605106354, 0.28291502594947815, 0.16039209067821503, 0.2695491909980774, 0.16163654625415802, 0.08897912502288818, 0.28747832775115967, 0.8989478349685669, 0.26775097846984863, 0.17184530198574066, 0.3264879584312439, 0.31386569142341614, 0.1549917310476303, 0.05264737084507942], [0.04084352031350136, 0.5361505150794983, 0.018223807215690613, 0.03828004375100136, 0.3140276074409485, 0.08277524262666702, 0.07094793766736984, 0.012667819857597351, 0.3304368853569031, 0.10053964704275131, 0.03868165612220764, 0.31755131483078003, 0.22644393146038055, 0.07613880187273026, 0.12961620092391968], [0.07373615354299545, 0.19122207164764404, 0.06966950744390488, 0.01624569669365883, 0.017842771485447884, 0.2144099771976471, 0.24285149574279785, 0.3761756718158722, 0.8141085505485535, 0.27487871050834656, 0.09974052757024765, 0.10127317160367966, 0.16323235630989075, 0.21032299101352692, 0.10343435406684875], [0.06651142984628677, 0.1456020176410675, 0.01741747185587883, 0.07566884905099869, 0.018790215253829956, 0.20801369845867157, 0.16892337799072266, 0.33592528104782104, 0.1834612786769867, 0.29906225204467773, 0.2579277753829956, 0.5998365879058838, 0.5642448663711548, 0.572043240070343, 0.0891154333949089], [0.03234146162867546, 0.1962265521287918, 0.0277019701898098, 0.06972747296094894, 0.10650040954351425, 0.07791601866483688, 0.38205334544181824, 0.4892197549343109, 0.003444283502176404, 0.414199560880661, 0.16890743374824524, 0.4916560649871826, 0.8149713277816772, 0.7298122048377991, 0.14976243674755096], [0.07799918204545975, 0.2381461262702942, 0.01647050306200981, 0.08363308757543564, 0.05209676921367645, 0.02968973107635975, 0.11220219731330872, 0.32446831464767456, 0.1546868085861206, 0.06510066986083984, 0.1935844123363495, 0.5264057517051697, 0.34881067276000977, 0.6311980485916138, 0.09822507947683334], [0.1688770204782486, 0.13700607419013977, 0.20374003052711487, 0.12288741022348404, 0.15864238142967224, 0.039533428847789764, 0.12642242014408112, 0.35126128792762756, 0.365562379360199, 0.48467183113098145, 0.3247453570365906, 0.003142370842397213, 0.5969579219818115, 0.5533550977706909, 0.1647837609052658], [0.3052995800971985, 0.6539703607559204, 0.022321274504065514, 0.1902511715888977, 0.05963977798819542, 0.17083951830863953, 0.5218495726585388, 0.2573777139186859, 0.17107829451560974, 0.46426069736480713, 0.3389802873134613, 0.4338558316230774, 0.014936042949557304, 0.6202957630157471, 0.13899832963943481], [0.12219581007957458, 0.5012378692626953, 0.06702763587236404, 0.06399006396532059, 0.07401375472545624, 0.24048954248428345, 0.08739905059337616, 0.050457850098609924, 0.030934542417526245, 0.1506662517786026, 0.1536494344472885, 0.49837279319763184, 0.018043117597699165, 0.11216632276773453, 0.12939369678497314], [0.11525271832942963, 0.521948516368866, 0.007329752668738365, 0.008543604053556919, 0.05213259160518646, 0.04235774278640747, 0.2166471928358078, 0.528154194355011, 0.42159566283226013, 0.22446103394031525, 0.0032521234825253487, 0.5035390257835388, 0.365617960691452, 0.44961339235305786, 0.15735329687595367], [0.03232282027602196, 0.08449342846870422, 0.004147443920373917, 0.050799064338207245, 0.037334948778152466, 0.08206064254045486, 0.07099173963069916, 0.19771835207939148, 0.021330662071704865, 0.08051090687513351, 0.1005825400352478, 0.700605034828186, 0.3027697801589966, 0.4364767074584961, 0.10480254143476486], [0.034268103539943695, 0.16091260313987732, 0.0168391652405262, 0.06967493146657944, 0.0915973111987114, 0.051104262471199036, 0.2385529726743698, 0.3295409679412842, 0.0004638703539967537, 0.22104156017303467, 0.13362999260425568, 0.5110065937042236, 0.7347238063812256, 0.7763577103614807, 0.15897347033023834], [0.08530293405056, 0.1988343894481659, 0.010091865435242653, 0.07736483961343765, 0.030177433043718338, 0.023718634620308876, 0.06320804357528687, 0.20902810990810394, 0.020835628733038902, 0.026085397228598595, 0.10371798276901245, 0.427949994802475, 0.2465561032295227, 0.6410334706306458, 0.12414435297250748], [0.17881684005260468, 0.09949745982885361, 0.17292529344558716, 0.14197823405265808, 0.0994792953133583, 0.022899990901350975, 0.07621151208877563, 0.20277591049671173, 0.059071850031614304, 0.23252709209918976, 0.2142648547887802, 0.0016634195344522595, 0.4786902368068695, 0.5105896592140198, 0.1802191287279129], [0.29184988141059875, 0.5299537181854248, 0.01714717224240303, 0.1581006944179535, 0.034420810639858246, 0.1480618417263031, 0.35555243492126465, 0.16130897402763367, 0.0352683924138546, 0.2384539395570755, 0.22334522008895874, 0.274210661649704, 0.008749962784349918, 0.5107676982879639, 0.16247788071632385], [0.1536586880683899, 0.39876002073287964, 0.060627128928899765, 0.08434724807739258, 0.06138864532113075, 0.18170806765556335, 0.0558285117149353, 0.026850836351513863, 0.004648242145776749, 0.05450701341032982, 0.08679821342229843, 0.24500715732574463, 0.009806739166378975, 0.06359081715345383, 0.14997224509716034], [0.1216418668627739, 0.4058372378349304, 0.00597163662314415, 0.009731672704219818, 0.04685758054256439, 0.030955728143453598, 0.14503908157348633, 0.4122965633869171, 0.13539999723434448, 0.08889995515346527, 0.0017191163497045636, 0.24694381654262543, 0.23039060831069946, 0.2996818721294403, 0.1837962418794632], [0.2966727912425995, 0.1567845344543457, 0.07310101389884949, 0.14124755561351776, 0.2961083948612213, 0.07968501001596451, 0.06122228875756264, 0.14724984765052795, 0.06047076731920242, 0.055829375982284546, 0.06430483609437943, 0.11614347994327545, 0.15107537806034088, 0.15706941485404968, 0.12527146935462952]], [[0.004390498157590628, 0.00876205787062645, 0.016465701162815094, 0.005714573431760073, 0.036494653671979904, 0.0032131776679307222, 0.01477664802223444, 0.018077310174703598, 0.010320773348212242, 0.006645719520747662, 0.03231831267476082, 0.004141036421060562, 0.011432528495788574, 0.011813640594482422, 0.20326180756092072], [0.024762088432908058, 0.05259820073843002, 0.06384432315826416, 0.1483391523361206, 0.26820069551467896, 0.20398226380348206, 0.37573596835136414, 0.08007726073265076, 0.052950888872146606, 0.09653404355049133, 0.1610451638698578, 0.12953783571720123, 0.2330068051815033, 0.4463363587856293, 0.19394421577453613], [0.679330587387085, 0.043791741132736206, 0.12768849730491638, 0.27546241879463196, 0.03847555071115494, 0.08167082816362381, 0.21957245469093323, 0.04802798852324486, 0.10780715942382812, 0.6106712222099304, 0.2505488693714142, 0.1709391176700592, 0.04529926925897598, 0.17936259508132935, 0.13903558254241943], [0.05959116667509079, 0.03547457605600357, 0.03805014118552208, 0.02909783646464348, 0.08531224727630615, 0.035567909479141235, 0.017052877694368362, 0.03032829985022545, 0.012725351378321648, 0.06508343666791916, 0.04963213950395584, 0.013415418565273285, 0.026129938662052155, 0.011819864623248577, 0.21026377379894257], [0.0922531858086586, 0.009465531446039677, 0.05285167694091797, 0.11621613800525665, 0.008946871384978294, 0.0003396931570023298, 0.056973982602357864, 0.011571673676371574, 0.03833528608083725, 0.02977353148162365, 0.12428728491067886, 0.005304301157593727, 0.012764646671712399, 0.03717968612909317, 0.1998610943555832], [0.024207258597016335, 0.015275360085070133, 0.12442810088396072, 0.044900182634592056, 0.06243159621953964, 0.002727220067754388, 0.05297050252556801, 0.34427115321159363, 0.10989916324615479, 0.020859790965914726, 0.11048608273267746, 0.02605186030268669, 0.1171213760972023, 0.05136575922369957, 0.16462838649749756], [0.03260662034153938, 0.00298042013309896, 0.16533112525939941, 0.056620776653289795, 0.049906134605407715, 0.008958332240581512, 0.05700542405247688, 0.016634995117783546, 0.029206881299614906, 0.025224529206752777, 0.19688823819160461, 0.03853357210755348, 0.07708126306533813, 0.04636078327894211, 0.17741571366786957], [0.04517968371510506, 0.08089613169431686, 0.11787059158086777, 0.09224344044923782, 0.27191361784935, 0.020393863320350647, 0.01454318780452013, 0.009129227139055729, 0.020442765206098557, 0.08070629835128784, 0.07541637122631073, 0.10045406222343445, 0.04119513928890228, 0.10953037440776825, 0.15667563676834106], [0.08136362582445145, 0.07834970951080322, 0.015254710800945759, 0.0832342654466629, 0.10864067077636719, 0.11524737626314163, 0.1366880238056183, 0.012557982467114925, 0.1251911222934723, 0.15952906012535095, 0.026927798986434937, 0.07786250859498978, 0.11803606152534485, 0.2014097422361374, 0.2085045427083969], [0.07754338532686234, 0.11610410362482071, 0.032187070697546005, 0.05519983917474747, 0.0022462301421910524, 0.11507689952850342, 0.2733137607574463, 0.17666463553905487, 0.010644900612533092, 0.08315187692642212, 0.02269633859395981, 0.06840697675943375, 0.010724963620305061, 0.0371541827917099, 0.21114735305309296], [0.022315502166748047, 0.012378118932247162, 0.0062178960070014, 0.0078407758846879, 0.015144318342208862, 0.010697844438254833, 0.011326298117637634, 0.013119788840413094, 0.009139686822891235, 0.006104558240622282, 0.005014281254261732, 0.002417754614725709, 0.007784656248986721, 0.009948876686394215, 0.16676713526248932], [0.2628116309642792, 0.1443735957145691, 0.08422664552927017, 0.11404431611299515, 0.17927099764347076, 0.25378888845443726, 0.1460212618112564, 0.04387032985687256, 0.023589681833982468, 0.13644081354141235, 0.045464351773262024, 0.06847606599330902, 0.006222521886229515, 0.036451175808906555, 0.20291540026664734], [0.22663825750350952, 0.15363532304763794, 0.01756531558930874, 0.025186356157064438, 0.038983430713415146, 0.01259024627506733, 0.15960636734962463, 0.10260611027479172, 0.059462085366249084, 0.02338782697916031, 0.039677273482084274, 0.055942799896001816, 0.010165784507989883, 0.013570738956332207, 0.1720115691423416], [0.04994741827249527, 0.08986728638410568, 0.03736276924610138, 0.029899757355451584, 0.03542618826031685, 0.007244490087032318, 0.040187276899814606, 0.040814109146595, 0.04076588898897171, 0.05965813249349594, 0.045340292155742645, 0.0002602309104986489, 0.026138437911868095, 0.02984587848186493, 0.21049101650714874], [0.058702513575553894, 0.04533839225769043, 0.03167680650949478, 0.07689032703638077, 0.07722999900579453, 0.05968516319990158, 0.08647314459085464, 0.04232413321733475, 0.05769982933998108, 0.08562258630990982, 0.07418374717235565, 0.08922348916530609, 0.0013435373548418283, 0.0365031398832798, 0.1955317258834839], [0.035160183906555176, 0.01820351555943489, 0.1303882896900177, 0.019772829487919807, 0.040328264236450195, 0.05493366718292236, 0.03643186390399933, 0.013673724606633186, 0.020261095836758614, 0.09265058487653732, 0.06087178364396095, 0.005874141119420528, 0.0010416797595098615, 0.00679743243381381, 0.17795756459236145], [0.0850016176700592, 0.12483492493629456, 0.30438917875289917, 0.08283902704715729, 0.36141735315322876, 0.5806636810302734, 0.21757252514362335, 0.0776025652885437, 0.2093839943408966, 0.1517311930656433, 0.0691467672586441, 0.05431315675377846, 0.323522686958313, 0.21248842775821686, 0.11186490952968597], [0.017619943246245384, 0.008017263375222683, 0.019503258168697357, 0.014857600443065166, 0.07692210376262665, 0.015309707261621952, 0.015313221141695976, 0.008549719117581844, 0.03095930442214012, 0.019377540796995163, 0.031960610300302505, 0.0054225618951022625, 0.016712497919797897, 0.015215321443974972, 0.15961019694805145], [0.2695287764072418, 0.16650046408176422, 0.14075446128845215, 0.1364857405424118, 0.23432065546512604, 0.261515349149704, 0.18958930671215057, 0.053015366196632385, 0.031337250024080276, 0.28422990441322327, 0.08986067771911621, 0.06408891826868057, 0.008591849356889725, 0.031372129917144775, 0.19151051342487335], [0.2586316764354706, 0.21131351590156555, 0.019284198060631752, 0.02717362530529499, 0.037918541580438614, 0.014535612426698208, 0.14439015090465546, 0.14164134860038757, 0.06384728103876114, 0.03232301026582718, 0.05240772292017937, 0.08253412693738937, 0.007928711362183094, 0.011026060208678246, 0.1583670824766159], [0.0646420493721962, 0.15151722729206085, 0.04734531044960022, 0.03642117232084274, 0.03833956643939018, 0.007805521599948406, 0.03985777497291565, 0.05410199984908104, 0.07749858498573303, 0.1281091719865799, 0.06692291796207428, 0.0004382343322504312, 0.02769407443702221, 0.03219819441437721, 0.20084568858146667], [0.06935474276542664, 0.07278740406036377, 0.0317843034863472, 0.061563972383737564, 0.057788632810115814, 0.05731336027383804, 0.08327846229076385, 0.046548519283533096, 0.06359860301017761, 0.13075897097587585, 0.09122113883495331, 0.1188196912407875, 0.0009191188146360219, 0.03464866429567337, 0.18994329869747162], [0.04588386043906212, 0.027941085398197174, 0.16196617484092712, 0.023955674842000008, 0.04093120992183685, 0.06800121814012527, 0.031365618109703064, 0.013349683955311775, 0.016157155856490135, 0.09367228299379349, 0.06382262706756592, 0.009268027730286121, 0.0006308736628852785, 0.005314440466463566, 0.17240527272224426], [0.09685268998146057, 0.17937548458576202, 0.31954076886177063, 0.09235721081495285, 0.3550800085067749, 0.5939842462539673, 0.19687135517597198, 0.10603781044483185, 0.27224627137184143, 0.17071248590946198, 0.0712975338101387, 0.10525800287723541, 0.3080449402332306, 0.250378280878067, 0.11120767891407013], [0.012543261051177979, 0.010277148336172104, 0.014658409170806408, 0.007294217124581337, 0.028056686744093895, 0.009602113626897335, 0.004711315967142582, 0.003909323364496231, 0.019910220056772232, 0.0035717461723834276, 0.016398703679442406, 0.01044577918946743, 0.015165981836616993, 0.04322582483291626, 0.1563079059123993]]], [[[0.017177388072013855, 0.0003127168456558138, 0.004294774029403925, 0.0025685238651931286, 0.0020048224832862616, 0.0018501998856663704, 0.004262528382241726, 0.00010045748058473691, 0.004143967293202877, 0.0026836262550204992, 0.0008790316642262042, 0.0012905423063784838, 8.68891947902739e-05, 0.00021419797849375755, 0.16245633363723755], [0.12795236706733704, 0.00371668953448534, 0.02831968478858471, 0.025539351627230644, 0.0009935664711520076, 0.0005314573645591736, 0.0308157317340374, 4.653090945794247e-05, 0.004544692113995552, 0.02307700179517269, 0.014357739128172398, 0.0017676070565357804, 1.5830510164960288e-05, 0.0005655316635966301, 0.23366259038448334], [0.0012442924780771136, 0.6349257826805115, 1.560185046400875e-05, 0.0005892697954550385, 2.671209358595661e-06, 1.747990245348774e-05, 0.00010909549746429548, 9.000968930195086e-06, 1.720580803521443e-05, 0.0008049540338106453, 0.00025925427326001227, 4.468534825718962e-06, 5.9764097386505455e-06, 7.895294402260333e-05, 0.00020540088007692248], [0.014811321161687374, 0.6550174951553345, 5.4754978918936104e-05, 0.0013682727003470063, 7.1730828494764864e-06, 3.513193587423302e-05, 0.00030579010490328074, 4.0161107790481765e-06, 8.621193410363048e-05, 0.0020331761334091425, 0.00018049145000986755, 1.5370842447737232e-05, 2.3058303213474574e-06, 3.803792060352862e-05, 0.0004018820764031261], [0.0038746336940675974, 0.000324725842801854, 0.0051879663951694965, 0.009153621271252632, 0.0008864403935149312, 0.6781038641929626, 0.057408660650253296, 0.0010902854846790433, 0.00043091498082503676, 0.000930881651584059, 0.00047575533972121775, 0.0024355631321668625, 0.0005705857765860856, 0.0003382607828825712, 0.0010924984235316515], [3.359095899213571e-06, 1.5333833403019526e-07, 3.112653939751908e-05, 0.00013510043208952993, 6.284327810135437e-06, 0.7821753025054932, 0.0016732696676626801, 2.949555346276611e-05, 1.1825303545265342e-06, 2.2443591660703532e-06, 4.938602842230466e-07, 8.253279020209447e-07, 2.1931487026449759e-07, 9.422030302630446e-07, 3.409375494811684e-06], [0.00014056767395231873, 5.100669682178705e-07, 0.0031089531257748604, 0.006296438630670309, 0.00044245802564546466, 0.5631491541862488, 0.006006886251270771, 0.00015836386592127383, 1.0129460861207917e-05, 9.741926623973995e-05, 8.02019567345269e-05, 2.8800504878745414e-05, 2.2740101485396735e-05, 9.966635116143152e-05, 5.9340749430703e-05], [0.07201159745454788, 9.12444302230142e-05, 0.07167930901050568, 0.07350550591945648, 0.008381813764572144, 0.32997292280197144, 0.32325229048728943, 0.006826527416706085, 0.005964158568531275, 0.01031426526606083, 0.0041834041476249695, 0.0003298712254036218, 2.8659975214395672e-05, 0.00019656911899801344, 0.02016262151300907], [0.0011574724921956658, 3.413460092360765e-07, 0.00010100962390424684, 0.0058910842053592205, 3.088227913394803e-06, 0.01394782867282629, 0.16852441430091858, 0.6476468443870544, 4.158269439358264e-05, 0.002217742381617427, 3.1430703529622406e-05, 8.318846812471747e-05, 7.552150123046886e-07, 2.136993316526059e-06, 0.00013183141709305346], [0.056869976222515106, 0.00018767332949209958, 0.07251239567995071, 0.21200358867645264, 0.5404223799705505, 0.01658189669251442, 0.03565289452672005, 0.0015120785683393478, 0.002293382305651903, 0.005935561377555132, 0.012055100873112679, 0.005193157121539116, 0.003556813346222043, 0.007320231292396784, 0.018532630056142807], [0.37012216448783875, 0.0030506134498864412, 0.585090160369873, 0.3774729073047638, 0.6362679600715637, 0.12865976989269257, 0.340728759765625, 0.01963443122804165, 0.11373940855264664, 0.0405576266348362, 0.04042620584368706, 0.006893007550388575, 0.0011100739939138293, 0.004035779275000095, 0.12706774473190308], [0.01695789396762848, 0.00023016006161924452, 0.013878279365599155, 0.04998883232474327, 0.0032932739704847336, 8.226843783631921e-05, 0.014781651087105274, 0.00017401285003870726, 0.4112556278705597, 0.007095593959093094, 0.01393651869148016, 0.000858593441080302, 0.0009966455399990082, 0.006141065154224634, 0.004614917561411858], [0.023780474439263344, 4.510316648520529e-05, 0.013797261752188206, 0.087004654109478, 0.0004407854867167771, 0.0013536562910303473, 0.04187630116939545, 0.0028901200275868177, 0.06213926523923874, 0.3483656048774719, 0.03705320879817009, 0.005524389911442995, 0.0004139445663895458, 0.0025706440210342407, 0.012163926847279072], [0.017730457708239555, 8.937691018218175e-05, 0.00767871318385005, 0.02321789041161537, 0.00010702417785068974, 0.004407694097608328, 0.0538853257894516, 0.011079255491495132, 0.003184565110132098, 0.026336153969168663, 0.005110009107738733, 0.3480301797389984, 0.002053677337244153, 0.01653059385716915, 0.00945478305220604], [0.00016590843733865768, 4.410037217894569e-05, 0.0031412369571626186, 0.0015988551313057542, 0.002399750053882599, 0.0004506838449742645, 0.001152031123638153, 0.00021803524577990174, 0.00054850586457178, 0.0001300607982557267, 0.001143390079960227, 0.0023531741462647915, 0.6484718322753906, 0.061944324523210526, 1.8855764210456982e-05], [5.492825607689156e-07, 1.991102926979238e-08, 2.3713612335996004e-06, 1.7095164366764948e-05, 8.657893886265811e-07, 3.6805211323098774e-08, 1.598790731804911e-06, 2.0731313554733788e-07, 4.274500042811269e-07, 5.490248440764844e-06, 0.00014167907647788525, 5.53526615476585e-06, 0.5851997137069702, 0.22563536465168, 1.0684430407081891e-07], [0.01633528247475624, 0.0006067559006623924, 0.047781698405742645, 0.1674666851758957, 0.0008243213524110615, 0.0007217283127829432, 0.005900595337152481, 0.0001012250068015419, 0.006910703144967556, 0.1343279927968979, 0.5695670247077942, 0.0034049933310598135, 0.008110514841973782, 0.0796104148030281, 0.00713667506352067], [0.02614973485469818, 0.001497315475717187, 0.11498566716909409, 0.08699594438076019, 0.006599655374884605, 0.0011878651566803455, 0.009639720432460308, 0.0002812722814269364, 0.014351817779242992, 0.06119270250201225, 0.19180962443351746, 0.06391202658414841, 0.4759237766265869, 0.44549837708473206, 0.058810409158468246], [0.041024841368198395, 0.0016396299470216036, 0.05072889104485512, 0.1323171705007553, 0.0024413676001131535, 0.00023246044293045998, 0.02059599943459034, 0.00033336327760480344, 0.7358176708221436, 0.04226389154791832, 0.0658484548330307, 0.002587914001196623, 0.013076293282210827, 0.0423613116145134, 0.051219869405031204], [0.025904469192028046, 0.00014531973283737898, 0.014812517911195755, 0.11958510428667068, 0.0003183217777404934, 0.0012536202557384968, 0.031174438074231148, 0.0025010022800415754, 0.045685503631830215, 0.4334242641925812, 0.057037968188524246, 0.005963113158941269, 0.0007164725102484226, 0.00356480129994452, 0.02565825544297695], [0.04193783551454544, 0.0005606984486803412, 0.01569434627890587, 0.058890990912914276, 0.00016686622984707355, 0.0032934362534433603, 0.10695304721593857, 0.011062747798860073, 0.008127261884510517, 0.04922156408429146, 0.01035262644290924, 0.3408533036708832, 0.003045044606551528, 0.019185535609722137, 0.046415992081165314], [0.00012501348101068288, 4.870840712101199e-05, 0.0024386774748563766, 0.001847597537562251, 0.0017206922639161348, 0.0002501157287042588, 0.0009360458934679627, 0.00021343374100979418, 0.0004799730086233467, 0.00017777700850274414, 0.0013057318283244967, 0.0019216074142605066, 0.7016423344612122, 0.059743087738752365, 1.6802117897896096e-05], [1.7574552657606546e-06, 9.272354617451128e-08, 1.001089003693778e-05, 5.891482942388393e-05, 3.3656547202554066e-06, 1.2065736143540562e-07, 6.7727110035775695e-06, 6.411150366147922e-07, 1.3192883443480241e-06, 1.1707085832313169e-05, 0.00026830541901290417, 1.0283902156515978e-05, 0.6812964081764221, 0.27208930253982544, 4.838558993469633e-07], [0.01900503970682621, 0.0008953948272392154, 0.09836827963590622, 0.2858547866344452, 0.0013939865166321397, 0.0011423979885876179, 0.011685764417052269, 0.00014273256238084286, 0.010754182003438473, 0.15914513170719147, 0.6438553333282471, 0.002441136632114649, 0.008362390100955963, 0.07132171094417572, 0.011131932027637959], [0.12417581677436829, 0.0153038389980793, 0.12986266613006592, 0.6406017541885376, 0.009386910125613213, 0.057520631700754166, 0.09723392128944397, 0.0041757188737392426, 0.030985616147518158, 0.12765046954154968, 0.052563395351171494, 0.09427980333566666, 0.010530965402722359, 0.01615813747048378, 0.110444575548172]], [[0.05668458715081215, 0.013551714830100536, 0.3300224542617798, 0.22417771816253662, 0.24923239648342133, 0.16107039153575897, 0.07639153301715851, 0.036736860871315, 0.044193096458911896, 0.14611276984214783, 0.15061600506305695, 0.035221245139837265, 0.0397845022380352, 0.06225845590233803, 0.12414046376943588], [0.29422780871391296, 0.3258638381958008, 0.027477310970425606, 0.10906420648097992, 0.003920723684132099, 0.020042676478624344, 0.05157224088907242, 0.0009247793932445347, 0.005282218102365732, 0.1744423359632492, 0.0761384516954422, 0.0033416510559618473, 0.0003361533163115382, 0.0012587645323947072, 0.013668928295373917], [0.19355924427509308, 0.1259031891822815, 0.004604514688253403, 0.04003702849149704, 0.0129036083817482, 0.019794460386037827, 0.06589072942733765, 0.0014933310449123383, 0.012753497809171677, 0.06252782791852951, 0.0361945815384388, 0.011655895970761776, 0.01012047752737999, 0.02639157697558403, 0.16549569368362427], [0.4293937385082245, 0.07181306928396225, 0.003158864099532366, 0.04697505012154579, 0.01354672759771347, 0.09221473336219788, 0.24058710038661957, 0.0037424738984555006, 0.07543525844812393, 0.0656844824552536, 0.01989266835153103, 0.06512395292520523, 0.01137665193527937, 0.029709961265325546, 0.18951866030693054], [0.052543047815561295, 0.03695955500006676, 0.100065678358078, 0.07546547800302505, 0.053252771496772766, 0.11382242292165756, 0.28551623225212097, 0.14051520824432373, 0.12815484404563904, 0.15533913671970367, 0.11139650642871857, 0.09512985497713089, 0.017796501517295837, 0.04266834259033203, 0.1351824700832367], [0.002040643012151122, 0.005490712355822325, 0.024769198149442673, 0.007002650294452906, 0.0020249236840754747, 0.03913044556975365, 0.01487613096833229, 0.09424738585948944, 0.010089649818837643, 0.05513475462794304, 0.0488949678838253, 0.007691625505685806, 0.002344577107578516, 0.012510538101196289, 0.20307941734790802], [0.04981796815991402, 0.13342007994651794, 0.4189896881580353, 0.06767702847719193, 0.007763800676912069, 0.11641503125429153, 0.029343493282794952, 0.11072052270174026, 0.06700066477060318, 0.1429358571767807, 0.3406253457069397, 0.00571059063076973, 0.0006326772854663432, 0.004126383922994137, 0.17491626739501953], [0.008032058365643024, 0.009898788295686245, 0.0165096465498209, 0.015990890562534332, 0.001612947671674192, 0.07025154680013657, 0.1309722512960434, 0.45684561133384705, 0.020022952929139137, 0.014566164463758469, 0.01627122238278389, 0.001012062537483871, 0.003352430183440447, 0.006583840120583773, 0.0849505066871643], [0.027854006737470627, 0.008844887837767601, 0.011581032536923885, 0.014227867126464844, 0.0022522227372974157, 0.6803511381149292, 0.24682462215423584, 0.11913055926561356, 0.0028406307101249695, 0.006190288811922073, 0.00574448611587286, 0.0012344244169071317, 0.010572707280516624, 0.00985674187541008, 0.11121391505002975], [0.11111988872289658, 0.0035893325693905354, 0.4007861316204071, 0.2033512443304062, 0.1986382007598877, 0.15137647092342377, 0.12109687924385071, 0.007575488183647394, 0.021906785666942596, 0.03087061457335949, 0.08533017337322235, 0.07086688280105591, 0.06729871034622192, 0.045789312571287155, 0.1673528403043747], [0.06468851119279861, 0.006587199401110411, 0.23617494106292725, 0.19800357520580292, 0.15495024621486664, 0.06172433868050575, 0.05180465057492256, 0.01833559013903141, 0.016546709463000298, 0.05746111273765564, 0.0824536681175232, 0.007550883572548628, 0.007943101227283478, 0.011712267994880676, 0.33849596977233887], [0.09414701163768768, 0.10295354574918747, 0.0844656303524971, 0.06548816710710526, 0.08529236167669296, 0.06227908656001091, 0.030192906036973, 0.010874724946916103, 0.025562399998307228, 0.005146168638020754, 0.014559037052094936, 0.013559900224208832, 0.06781303137540817, 0.05153109133243561, 0.33232951164245605], [0.314544141292572, 0.6832185983657837, 0.07794945687055588, 0.042061515152454376, 0.015504884533584118, 0.1916494369506836, 0.006379975005984306, 0.0006176759488880634, 0.0012508369982242584, 0.01929312013089657, 0.022219885140657425, 0.0019787217024713755, 0.01769268326461315, 0.008809820748865604, 0.08711312711238861], [0.027118999511003494, 0.07309459149837494, 0.04486501216888428, 0.012266037985682487, 0.024303032085299492, 0.030924739316105843, 0.021004648879170418, 0.003694491693750024, 0.01517508551478386, 0.025275954976677895, 0.0075909653678536415, 0.24021397531032562, 0.04135901853442192, 0.07603362947702408, 0.11061857640743256], [0.025165440514683723, 0.019109023734927177, 0.008520743809640408, 0.015198510140180588, 0.007751345168799162, 0.005125374533236027, 0.008160223253071308, 0.0017721926560625434, 0.08641061931848526, 0.07765892893075943, 0.017936453223228455, 0.020675569772720337, 0.0024341135285794735, 0.023971976712346077, 0.16557703912258148], [0.22320780158042908, 0.05348529666662216, 0.01734296977519989, 0.1172923669219017, 0.004340981598943472, 0.003372892737388611, 0.033841460943222046, 0.024162178859114647, 0.05216863751411438, 0.3090120553970337, 0.2295515090227127, 0.014075365848839283, 0.020010780543088913, 0.20773397386074066, 0.12411301583051682], [0.1383964717388153, 0.05579448863863945, 0.1563209742307663, 0.09128513187170029, 0.039257608354091644, 0.009886945597827435, 0.006391164381057024, 0.0007081980584189296, 0.006523598916828632, 0.16335614025592804, 0.02935076504945755, 0.023180969059467316, 0.19186609983444214, 0.2336183488368988, 0.16814255714416504], [0.1625337302684784, 0.007939358241856098, 0.11928629875183105, 0.1341797411441803, 0.005670298356562853, 0.0033473502844572067, 0.022544465959072113, 0.005534132476896048, 0.007299710530787706, 0.08667418360710144, 0.07403960824012756, 0.004230144899338484, 0.002401313977316022, 0.005503634922206402, 0.20701391994953156], [0.08204744011163712, 0.04882703348994255, 0.048393696546554565, 0.02867632359266281, 0.012730585411190987, 0.02805456519126892, 0.014470821246504784, 0.008571655489504337, 0.011637779884040356, 0.011116313748061657, 0.015620187856256962, 0.00444953003898263, 0.038398172706365585, 0.021771300584077835, 0.25556278228759766], [0.3818233609199524, 0.6690115928649902, 0.07648678869009018, 0.0345233753323555, 0.011518634855747223, 0.1436365395784378, 0.005264819134026766, 0.000502048700582236, 0.0017500953981652856, 0.03918173909187317, 0.04129163548350334, 0.0023984990548342466, 0.020183494314551353, 0.008427903987467289, 0.09516369551420212], [0.02332407608628273, 0.06938373297452927, 0.035716570913791656, 0.008126936852931976, 0.012537641450762749, 0.0137803228572011, 0.01513306051492691, 0.00204691500402987, 0.029820755124092102, 0.05474912002682686, 0.016170548275113106, 0.22342036664485931, 0.05026429146528244, 0.06863567978143692, 0.11948796361684799], [0.020166568458080292, 0.015762973576784134, 0.006330324336886406, 0.008625769056379795, 0.005781465210020542, 0.00451312493532896, 0.007413441780954599, 0.0018466140609234571, 0.14846709370613098, 0.1376892477273941, 0.02431248314678669, 0.03153817355632782, 0.0025850962847471237, 0.026987632736563683, 0.15984071791172028], [0.11904438585042953, 0.03637225553393364, 0.013324074447154999, 0.04586002975702286, 0.00359557312913239, 0.002297254279255867, 0.02453085221350193, 0.019205793738365173, 0.07615289092063904, 0.3510056436061859, 0.24748629331588745, 0.0179043747484684, 0.015299135819077492, 0.16336295008659363, 0.13914434611797333], [0.0598345547914505, 0.028141267597675323, 0.11996681243181229, 0.04193190485239029, 0.03001757152378559, 0.006633914541453123, 0.005910022184252739, 0.0007469199481420219, 0.010509159415960312, 0.18832749128341675, 0.032145459204912186, 0.022126449272036552, 0.16793787479400635, 0.1917877346277237, 0.16885708272457123], [0.30011340975761414, 0.029496116563677788, 0.21246175467967987, 0.11388618499040604, 0.019265230745077133, 0.011386800557374954, 0.02386542037129402, 0.0049255480989813805, 0.002113579073920846, 0.2235003262758255, 0.1410367637872696, 0.022971738129854202, 0.009332037530839443, 0.01034344732761383, 0.12311729788780212]], [[0.03517069295048714, 0.03549245744943619, 0.004381549544632435, 0.008797217160463333, 0.007323419209569693, 0.042320944368839264, 0.004849699325859547, 0.003679578425362706, 0.011580413207411766, 0.009367180056869984, 0.006541883572936058, 0.022973380982875824, 0.023761657997965813, 0.02892483025789261, 0.1581033319234848], [0.01528994832187891, 0.20408181846141815, 0.11101088672876358, 0.08111120015382767, 0.07986893504858017, 0.010126215405762196, 0.020366966724395752, 0.1417536586523056, 0.04787333309650421, 0.04340335354208946, 0.2409791648387909, 0.04442436248064041, 0.005909040104597807, 0.014603852294385433, 0.18931475281715393], [0.21622280776500702, 0.09626477211713791, 0.10110790282487869, 0.31975099444389343, 0.2572377920150757, 0.630383312702179, 0.1336757242679596, 0.17725828289985657, 0.02378956414759159, 0.22253809869289398, 0.13939163088798523, 0.30914127826690674, 0.35968318581581116, 0.48164138197898865, 0.09301326423883438], [0.168080672621727, 0.1516411453485489, 0.07150255143642426, 0.32225823402404785, 0.2490793913602829, 0.30686429142951965, 0.032337237149477005, 0.16698232293128967, 0.04405289515852928, 0.2310783565044403, 0.10561788827180862, 0.2769646644592285, 0.19830158352851868, 0.1653461754322052, 0.09653043746948242], [0.04038669914007187, 0.16624715924263, 0.3317047655582428, 0.3851986229419708, 0.42305275797843933, 0.008450526744127274, 0.09501849114894867, 0.24002836644649506, 0.4256587326526642, 0.15410973131656647, 0.19127053022384644, 0.04389801248908043, 0.030224177986383438, 0.05971052870154381, 0.11478950828313828], [0.04527302458882332, 0.15370813012123108, 0.46266382932662964, 0.06791326403617859, 0.6029869914054871, 0.018879592418670654, 0.07514301687479019, 0.07948564738035202, 0.6243545413017273, 0.11254889518022537, 0.24916931986808777, 0.08612842112779617, 0.07598677277565002, 0.13317255675792694, 0.04299912229180336], [0.03695433586835861, 0.028389452025294304, 0.2721908688545227, 0.07653216272592545, 0.6730886697769165, 0.004614274017512798, 0.004165990743786097, 0.01533985324203968, 0.28992146253585815, 0.028840038925409317, 0.055076081305742264, 0.024787841364741325, 0.0010191021719947457, 0.0022868094965815544, 0.030124979093670845], [0.005083801224827766, 0.09139324724674225, 0.28116321563720703, 0.08195066452026367, 0.6340349316596985, 0.012272918596863747, 0.0005934475339017808, 0.010692326352000237, 0.1514793336391449, 0.016046250239014626, 0.04672969505190849, 0.014393122866749763, 0.002580928150564432, 0.007409923244267702, 0.12582267820835114], [0.00605103699490428, 0.11548061668872833, 0.2870264947414398, 0.061026521027088165, 0.8064441084861755, 0.2189176380634308, 0.020241523161530495, 0.07779920846223831, 0.08952271938323975, 0.0073190852999687195, 0.02372264862060547, 0.038144610822200775, 0.07446137070655823, 0.09413070231676102, 0.030171062797307968], [0.08316895365715027, 0.6715664267539978, 0.04549514129757881, 0.17856287956237793, 0.018127189949154854, 0.38010329008102417, 0.16956135630607605, 0.5726994872093201, 0.1473512202501297, 0.13756032288074493, 0.044131502509117126, 0.03872460126876831, 0.13646697998046875, 0.07963203638792038, 0.10255669057369232], [0.0817432552576065, 0.2031053900718689, 0.02472570165991783, 0.02598942257463932, 0.05427335575222969, 0.43315476179122925, 0.06398319453001022, 0.14792829751968384, 0.18555517494678497, 0.020227503031492233, 0.03572608157992363, 0.008726409636437893, 0.33127138018608093, 0.0956021174788475, 0.032814960926771164], [0.36652442812919617, 0.4977355897426605, 0.09286413341760635, 0.21385566890239716, 0.18058304488658905, 0.4562758207321167, 0.4738945960998535, 0.2067655473947525, 0.17124009132385254, 0.035114847123622894, 0.05785587430000305, 0.03289380669593811, 0.3892229497432709, 0.2459530532360077, 0.0885753259062767], [0.3338637053966522, 0.241106316447258, 0.10183558613061905, 0.16975384950637817, 0.22215212881565094, 0.1208982765674591, 0.12069278955459595, 0.027770178392529488, 0.12589573860168457, 0.018161755055189133, 0.05639319866895676, 0.024462532252073288, 0.08646970242261887, 0.18506868183612823, 0.2994369864463806], [0.24999171495437622, 0.7484717965126038, 0.1908620148897171, 0.6611655354499817, 0.24442408978939056, 0.0825357735157013, 0.5622089505195618, 0.4391622543334961, 0.045715928077697754, 0.2250336855649948, 0.3067566156387329, 0.014471310190856457, 0.06388252228498459, 0.21674634516239166, 0.13583892583847046], [0.05097173899412155, 0.16686855256557465, 0.15120531618595123, 0.3698476254940033, 0.35846272110939026, 0.6895467042922974, 0.8159933686256409, 0.843620777130127, 0.6904561519622803, 0.307870090007782, 0.450530469417572, 0.6275950074195862, 0.15986312925815582, 0.5293903350830078, 0.07888244837522507], [0.3532100319862366, 0.1141892597079277, 0.06207668036222458, 0.23437273502349854, 0.13035829365253448, 0.16457295417785645, 0.6610441207885742, 0.6354422569274902, 0.6703211069107056, 0.18266227841377258, 0.16635818779468536, 0.1048990935087204, 0.1468038111925125, 0.17976891994476318, 0.0709633082151413], [0.18437133729457855, 0.20806346833705902, 0.06752406805753708, 0.15831130743026733, 0.3405534625053406, 0.0627271831035614, 0.3717433214187622, 0.3913803696632385, 0.5862330794334412, 0.29396724700927734, 0.02299528755247593, 0.060014016926288605, 0.08232607692480087, 0.15418194234371185, 0.15275102853775024], [0.07671413570642471, 0.17070698738098145, 0.13325846195220947, 0.07402658462524414, 0.6503690481185913, 0.1330946981906891, 0.165133535861969, 0.2397843301296234, 0.6370089054107666, 0.09848601371049881, 0.09929761290550232, 0.10903115570545197, 0.14141131937503815, 0.14783106744289398, 0.08112896233797073], [0.1416744738817215, 0.274202436208725, 0.13295260071754456, 0.20105819404125214, 0.3945937156677246, 0.333781898021698, 0.3556738793849945, 0.2839928865432739, 0.10343024134635925, 0.07706140726804733, 0.054361648857593536, 0.05752982571721077, 0.2817353904247284, 0.27278265357017517, 0.13429909944534302], [0.22879131138324738, 0.1777554452419281, 0.09183042496442795, 0.14726729691028595, 0.1873711347579956, 0.05672184377908707, 0.08326486498117447, 0.01781904511153698, 0.0835406556725502, 0.02614605240523815, 0.06876543164253235, 0.03439611196517944, 0.0621294341981411, 0.16512615978717804, 0.26481878757476807], [0.1532706916332245, 0.5982866883277893, 0.18050755560398102, 0.5800401568412781, 0.22030943632125854, 0.025230426341295242, 0.3744361996650696, 0.265155166387558, 0.03173244372010231, 0.2068646252155304, 0.27338433265686035, 0.012270096689462662, 0.05047086998820305, 0.14277896285057068, 0.15170519053936005], [0.04688200727105141, 0.12437571585178375, 0.1870293915271759, 0.4533093273639679, 0.3565751910209656, 0.5648568868637085, 0.7852934002876282, 0.7657470703125, 0.5417794585227966, 0.4419334828853607, 0.632922887802124, 0.7103447914123535, 0.15686877071857452, 0.6169639825820923, 0.08483293652534485], [0.2884610891342163, 0.10604135692119598, 0.07176870107650757, 0.2240629643201828, 0.12294583767652512, 0.10159854590892792, 0.6051279902458191, 0.5541971921920776, 0.5623130798339844, 0.16405576467514038, 0.18055777251720428, 0.13399486243724823, 0.12637703120708466, 0.18360036611557007, 0.09598042815923691], [0.10626664012670517, 0.1478983461856842, 0.07806308567523956, 0.11814259737730026, 0.31690794229507446, 0.03372211009263992, 0.30042603611946106, 0.29277828335762024, 0.44479742646217346, 0.216581329703331, 0.023049354553222656, 0.0511498898267746, 0.08494822680950165, 0.14207273721694946, 0.16419102251529694], [0.048457998782396317, 0.0638582855463028, 0.20956584811210632, 0.021124709397554398, 0.09014897048473358, 0.11662621796131134, 0.3483109474182129, 0.4503737986087799, 0.17136822640895844, 0.02997676283121109, 0.21708470582962036, 0.05856599286198616, 0.2859736979007721, 0.41663405299186707, 0.12262307107448578]], [[0.01622859761118889, 0.0033176897559314966, 0.006228303536772728, 0.003451053285971284, 0.011415286920964718, 0.016942020505666733, 0.0027556640561670065, 0.001647507306188345, 0.0010015909792855382, 0.0013629572931677103, 0.004746851045638323, 0.009338179603219032, 0.00885467603802681, 0.006604180671274662, 0.16180677711963654], [0.17455320060253143, 0.026163265109062195, 0.2041780799627304, 0.027548620477318764, 0.4711945950984955, 0.5480062365531921, 0.10718726366758347, 0.032194506376981735, 0.08035919070243835, 0.010791448876261711, 0.11821587383747101, 0.04372825473546982, 0.5788823962211609, 0.10199426859617233, 0.06844703108072281], [0.023936308920383453, 0.03560526669025421, 0.007881848141551018, 0.022994371131062508, 0.003501775674521923, 0.000663262908346951, 0.0027445319574326277, 0.0008202926255762577, 0.002215484855696559, 0.014335977844893932, 0.06139073148369789, 0.0039900378324091434, 0.004902976099401712, 0.006251698825508356, 0.21882350742816925], [0.01501577626913786, 0.026870740577578545, 0.007700353395193815, 0.02517320215702057, 0.005199552513659, 0.0040618558414280415, 0.0018289085710421205, 0.0005822794046252966, 0.008953371085226536, 0.004845716059207916, 0.02605423890054226, 0.010851072147488594, 0.011600007303059101, 0.011058725416660309, 0.2679094076156616], [0.05198093131184578, 0.026691097766160965, 0.04745011776685715, 0.02099662832915783, 0.007765383925288916, 0.0017653746763244271, 0.002459246199578047, 0.0005052239284850657, 0.0007161727407947183, 0.00449666241183877, 0.00950489193201065, 0.002728741616010666, 0.007593079470098019, 0.0031749741174280643, 0.1993207037448883], [0.0031879025045782328, 0.001219254801981151, 0.007273980416357517, 0.0029734931886196136, 9.794573998078704e-05, 0.0006066279602237046, 0.000905939843505621, 0.0002116545947501436, 0.00022416051069740206, 0.001432110439054668, 0.00046862047747708857, 0.0008043517009355128, 0.00010411434050183743, 0.0003457288257777691, 0.22099417448043823], [0.020157048478722572, 0.026601465418934822, 0.04540588706731796, 0.04344630241394043, 0.0022944926749914885, 0.0010618591913953424, 0.00406603142619133, 0.0029086798895150423, 0.0019963555969297886, 0.010005260817706585, 0.0020353682339191437, 0.0019374215044081211, 0.0013613863848149776, 0.001661884132772684, 0.34173521399497986], [0.09776000678539276, 0.012011643499135971, 0.12930582463741302, 0.019725820049643517, 0.03450663015246391, 0.44516250491142273, 0.09379248321056366, 0.011904217302799225, 0.012111036106944084, 0.007218031212687492, 0.028761520981788635, 0.011232447810471058, 0.17035166919231415, 0.022308414801955223, 0.055901553481817245], [0.0270126610994339, 0.0034831874072551727, 0.03977394104003906, 0.025583824142813683, 0.0007700100541114807, 0.002870001830160618, 0.0027750579174607992, 0.0016644555144011974, 0.0016086471732705832, 0.001177149242721498, 0.00746855279430747, 0.002065857872366905, 0.0016993783647194505, 0.0015537800500169396, 0.32808277010917664], [0.16020068526268005, 0.019860466942191124, 0.3786206543445587, 0.04546584561467171, 0.22538548707962036, 0.035959187895059586, 0.022749971598386765, 0.0223965086042881, 0.010994979180395603, 0.013655508868396282, 0.08095952123403549, 0.07914181798696518, 0.5184871554374695, 0.24710357189178467, 0.059729527682065964], [0.002354596508666873, 0.013563946820795536, 0.0012282072566449642, 0.0011236226418986917, 0.004269973374903202, 0.05393142253160477, 0.010044331662356853, 0.012847290374338627, 0.23206481337547302, 0.0042032524943351746, 0.002388538094237447, 0.005051162093877792, 0.004106870852410793, 0.003583247307687998, 0.0021634430158883333], [0.1318124532699585, 0.006612265948206186, 0.026151085272431374, 0.15551267564296722, 0.006537565030157566, 0.045402105897665024, 0.08115606755018234, 0.020273711532354355, 0.2617640495300293, 0.03846455365419388, 0.42425140738487244, 0.0063036843203008175, 0.045534029603004456, 0.06594183295965195, 0.0061628553085029125], [0.0171976238489151, 0.0023818486370146275, 0.036466922610998154, 0.011855212040245533, 0.019672302529215813, 0.007386004086583853, 0.02982362173497677, 0.0045198979787528515, 0.02385052479803562, 0.25256073474884033, 0.2446560561656952, 0.0453505739569664, 0.08819476515054703, 0.09139581024646759, 0.0022182920947670937], [0.023948049172759056, 0.006307430099695921, 0.014840157702565193, 0.01758965104818344, 0.0009477039566263556, 0.00178795016836375, 0.005927308928221464, 0.0026511158794164658, 0.00012311375758145005, 0.04321818798780441, 0.0496363490819931, 0.3416200280189514, 0.001097637927159667, 0.007029203698039055, 0.007338459137827158], [0.1633826345205307, 0.005062526557594538, 0.04231903329491615, 0.24309031665325165, 0.0009563505300320685, 0.0008045694557949901, 0.004994159564375877, 0.0011061460245400667, 0.0013372766552492976, 0.023061903193593025, 0.044598180800676346, 0.0017028035363182425, 2.3589664124301635e-05, 0.0003540365141816437, 0.16737498342990875], [0.1106855720281601, 0.005593962036073208, 0.014953872188925743, 0.19064223766326904, 0.0008905718568712473, 0.002549833618104458, 0.019427485764026642, 0.019940704107284546, 0.0020017458591610193, 0.029780413955450058, 0.01774613931775093, 0.00061158457538113, 0.0022336822003126144, 0.007989613339304924, 0.2558586895465851], [0.07112060487270355, 0.029737049713730812, 0.09336916357278824, 0.07307538390159607, 0.023197662085294724, 0.022866347804665565, 0.060328319668769836, 0.04474486783146858, 0.0006379868718795478, 0.027103934437036514, 0.2942929267883301, 0.011375843547284603, 0.07746338844299316, 0.09051978588104248, 0.11258094012737274], [0.15941812098026276, 0.02997875213623047, 0.08360203355550766, 0.10365118086338043, 0.03050130233168602, 0.39312028884887695, 0.3065427839756012, 0.2912093997001648, 0.135236918926239, 0.18899840116500854, 0.13724294304847717, 0.1948302835226059, 0.07353706657886505, 0.12220755219459534, 0.10422825068235397], [0.24064786732196808, 0.0051915524527430534, 0.09652373939752579, 0.2287912219762802, 0.019215410575270653, 0.13947954773902893, 0.15343742072582245, 0.07055477797985077, 0.05467608571052551, 0.10673969984054565, 0.5659986138343811, 0.014077076688408852, 0.1709020584821701, 0.23944324254989624, 0.026877261698246002], [0.019817974418401718, 0.002034382661804557, 0.04978875443339348, 0.009913384914398193, 0.033772312104701996, 0.0069160182029008865, 0.027356693521142006, 0.004301261156797409, 0.005268980748951435, 0.24062182009220123, 0.2975090742111206, 0.09841412305831909, 0.13523375988006592, 0.1965852826833725, 0.004198803100734949], [0.017094334587454796, 0.005556214600801468, 0.011722622439265251, 0.009952181950211525, 0.0008346029790118337, 0.0009373819339089096, 0.006794091779738665, 0.0019291864009574056, 4.7701923904241994e-05, 0.0364256277680397, 0.035398196429014206, 0.3890627920627594, 0.0013647697633132339, 0.008012092672288418, 0.013173048384487629], [0.12328237295150757, 0.0036286553367972374, 0.03202027454972267, 0.16562366485595703, 0.0006255045300349593, 0.00061140360776335, 0.00499368691816926, 0.0010923785157501698, 0.0008833102765493095, 0.03177933022379875, 0.04344986379146576, 0.00255553494207561, 2.260845576529391e-05, 0.0005036385264247656, 0.16160868108272552], [0.050196755677461624, 0.002699600299820304, 0.009293685667216778, 0.06999042630195618, 0.0006182404467836022, 0.0013977399794384837, 0.014421526342630386, 0.010930507443845272, 0.0008620836888439953, 0.015927143394947052, 0.008692404255270958, 0.0006625624373555183, 0.0011245491914451122, 0.0053406055085361, 0.2061784416437149], [0.04101766273379326, 0.020672734826803207, 0.08772061765193939, 0.04009746387600899, 0.01892852783203125, 0.017910925671458244, 0.057973578572273254, 0.03737492114305496, 0.00047206622548401356, 0.021084431558847427, 0.21054430305957794, 0.013546224683523178, 0.08985017240047455, 0.10610225051641464, 0.1389981210231781], [0.018278781324625015, 0.03789714351296425, 0.00408195098862052, 0.005283118225634098, 0.009515376761555672, 0.11360906809568405, 0.008760524913668633, 0.006613489706069231, 0.018946174532175064, 0.008831392042338848, 0.015675490722060204, 0.021136337891221046, 0.13481837511062622, 0.08728663623332977, 0.15406787395477295]], [[0.05651351809501648, 0.11774645000696182, 0.026926513761281967, 0.04848615080118179, 0.10334916412830353, 0.4247743785381317, 0.21147629618644714, 0.6254463195800781, 0.10587190836668015, 0.08194849640130997, 0.04674661532044411, 0.35135090351104736, 0.35409873723983765, 0.43208518624305725, 0.11939813196659088], [0.05609016492962837, 0.06931670010089874, 0.1576625108718872, 0.27308744192123413, 0.04202406853437424, 0.2399596869945526, 0.3320065140724182, 0.6272499561309814, 0.09423039108514786, 0.144412100315094, 0.2769482433795929, 0.05643320456147194, 0.11388154327869415, 0.32551372051239014, 0.13187405467033386], [0.1798395812511444, 0.02382134646177292, 0.024498937651515007, 0.28730508685112, 0.19651466608047485, 0.13693250715732574, 0.34929007291793823, 0.1055094301700592, 0.08990196883678436, 0.5189381837844849, 0.3313819468021393, 0.34343984723091125, 0.21719343960285187, 0.21188895404338837, 0.15588119626045227], [0.26584357023239136, 0.03035559318959713, 0.026536965742707253, 0.20298171043395996, 0.23938016593456268, 0.24181482195854187, 0.31930428743362427, 0.10626629739999771, 0.13103167712688446, 0.4636806845664978, 0.393515944480896, 0.3422740399837494, 0.342117577791214, 0.5495904088020325, 0.14030353724956512], [0.30834218859672546, 0.3875667452812195, 0.32842832803726196, 0.16462059319019318, 0.416511207818985, 0.03730625659227371, 0.23662680387496948, 0.5092235207557678, 0.08549848943948746, 0.3278381824493408, 0.507111668586731, 0.0415511280298233, 0.5590415596961975, 0.6185146570205688, 0.0664283037185669], [0.0765935555100441, 0.29552146792411804, 0.05705742537975311, 0.01913047581911087, 0.15779250860214233, 0.030224651098251343, 0.08988720178604126, 0.3389361500740051, 0.08153010904788971, 0.05811480060219765, 0.09408371150493622, 0.19600677490234375, 0.6126919388771057, 0.623294472694397, 0.13969288766384125], [0.4304950535297394, 0.5688965320587158, 0.09143517911434174, 0.09618712961673737, 0.13307496905326843, 0.014428870752453804, 0.040250685065984726, 0.15830516815185547, 0.10923942923545837, 0.23653797805309296, 0.3180045783519745, 0.5594316720962524, 0.5058388710021973, 0.3866141140460968, 0.14058275520801544], [0.31169822812080383, 0.7707167863845825, 0.30778199434280396, 0.10994993895292282, 0.18047340214252472, 0.01769133098423481, 0.014783667400479317, 0.009741406887769699, 0.1340220719575882, 0.11223828792572021, 0.46960482001304626, 0.360332190990448, 0.56731116771698, 0.5470200181007385, 0.18929171562194824], [0.2397254854440689, 0.361926406621933, 0.24345533549785614, 0.18179422616958618, 0.10373111069202423, 0.014045567251741886, 0.08654272556304932, 0.018043776974081993, 0.02193235233426094, 0.07134812325239182, 0.19312754273414612, 0.6192790865898132, 0.6039608716964722, 0.673239529132843, 0.15608295798301697], [0.32110491394996643, 0.2706402838230133, 0.034645695239305496, 0.029830342158675194, 0.00933478306978941, 0.25964564085006714, 0.17791348695755005, 0.11580535769462585, 0.07073061913251877, 0.10197918862104416, 0.06440304219722748, 0.2378954440355301, 0.09358810633420944, 0.24307624995708466, 0.22625915706157684], [0.18688960373401642, 0.6521251797676086, 0.05505351349711418, 0.05518023297190666, 0.07190049439668655, 0.15721110999584198, 0.11867944896221161, 0.2974295914173126, 0.018550140783190727, 0.1645369827747345, 0.09910324215888977, 0.499615877866745, 0.34706613421440125, 0.5406060218811035, 0.24014075100421906], [0.24844318628311157, 0.24823600053787231, 0.41713690757751465, 0.05438315495848656, 0.5823535323143005, 0.1801777333021164, 0.13823869824409485, 0.16278210282325745, 0.035736992955207825, 0.017554355785250664, 0.03778500482439995, 0.09959819167852402, 0.18642207980155945, 0.26950401067733765, 0.24913227558135986], [0.21744470298290253, 0.04392259195446968, 0.5108200907707214, 0.27167755365371704, 0.5572997331619263, 0.30860280990600586, 0.5083038210868835, 0.6815038919448853, 0.3754148483276367, 0.01992654800415039, 0.0589066781103611, 0.07934294641017914, 0.15649113059043884, 0.3772245943546295, 0.25267744064331055], [0.11088164150714874, 0.06568774580955505, 0.49295517802238464, 0.06175035238265991, 0.3928946256637573, 0.306259423494339, 0.1265336275100708, 0.29877781867980957, 0.061930101364851, 0.053618840873241425, 0.02546272985637188, 0.011733881197869778, 0.4200928509235382, 0.25557151436805725, 0.12701815366744995], [0.06005493924021721, 0.46575742959976196, 0.4922090172767639, 0.06956527382135391, 0.3788193464279175, 0.21330630779266357, 0.06565267592668533, 0.10461793839931488, 0.1200915202498436, 0.07597928494215012, 0.08451344817876816, 0.06952610611915588, 0.03487509861588478, 0.12158560007810593, 0.14820002019405365], [0.11028759926557541, 0.4027779996395111, 0.8237467408180237, 0.1328621804714203, 0.7811888456344604, 0.5416622757911682, 0.16887041926383972, 0.2001309096813202, 0.08848496526479721, 0.05607001483440399, 0.13165172934532166, 0.10739479213953018, 0.052385441958904266, 0.05461856350302696, 0.16259506344795227], [0.12960980832576752, 0.21605639159679413, 0.13754284381866455, 0.0687912181019783, 0.2001095861196518, 0.7652902007102966, 0.3308810591697693, 0.3389359712600708, 0.07430214434862137, 0.036511119455099106, 0.010612682439386845, 0.005050503648817539, 0.1584991067647934, 0.036481909453868866, 0.18724960088729858], [0.16838932037353516, 0.47491130232810974, 0.21776747703552246, 0.05912807583808899, 0.16565343737602234, 0.34125030040740967, 0.2414778620004654, 0.28169524669647217, 0.03973108157515526, 0.03921183571219444, 0.02238578163087368, 0.02449338510632515, 0.05498792976140976, 0.03159895911812782, 0.17659053206443787], [0.14295107126235962, 0.27777984738349915, 0.30436068773269653, 0.03198731318116188, 0.38494178652763367, 0.27411460876464844, 0.18790900707244873, 0.29966217279434204, 0.029011890292167664, 0.012050352990627289, 0.008839968591928482, 0.009298003278672695, 0.09229473769664764, 0.05935056507587433, 0.2074589878320694], [0.185210719704628, 0.0802093893289566, 0.4863169491291046, 0.24164138734340668, 0.5185936689376831, 0.381059467792511, 0.5372542142868042, 0.6922534108161926, 0.40473121404647827, 0.015452258288860321, 0.03550630062818527, 0.023993153125047684, 0.09803077578544617, 0.14391310513019562, 0.25199130177497864], [0.08245678246021271, 0.1390499472618103, 0.5461503863334656, 0.060220371931791306, 0.43899697065353394, 0.5144884586334229, 0.22183947265148163, 0.5088672041893005, 0.09321429580450058, 0.05354699492454529, 0.02214067056775093, 0.004303250927478075, 0.39110496640205383, 0.12463895231485367, 0.1568218618631363], [0.043030936270952225, 0.498334676027298, 0.5084810853004456, 0.06107298657298088, 0.3904430866241455, 0.35258427262306213, 0.08483341336250305, 0.17738159000873566, 0.1815967708826065, 0.09597334265708923, 0.08432064205408096, 0.040181081742048264, 0.02593160979449749, 0.08670566976070404, 0.14764654636383057], [0.0785449668765068, 0.4015392065048218, 0.8182658553123474, 0.10243776440620422, 0.7659414410591125, 0.5735372304916382, 0.16621330380439758, 0.21339072287082672, 0.12523002922534943, 0.05685745179653168, 0.1081186980009079, 0.07184037566184998, 0.02847907319664955, 0.031456008553504944, 0.15293413400650024], [0.07311940938234329, 0.15430475771427155, 0.1386927217245102, 0.04823235049843788, 0.20945730805397034, 0.8191487193107605, 0.33371293544769287, 0.3618466258049011, 0.1152336597442627, 0.031010858714580536, 0.008395140990614891, 0.002998974174261093, 0.13362915813922882, 0.02411211095750332, 0.1613900512456894], [0.2622520923614502, 0.7386532425880432, 0.41215938329696655, 0.08539438247680664, 0.7665934562683105, 0.5218235850334167, 0.42940571904182434, 0.4037780165672302, 0.7456067204475403, 0.07961834967136383, 0.02781907096505165, 0.02608557976782322, 0.15701159834861755, 0.05025498941540718, 0.11428551375865936]], [[0.5009713768959045, 0.11806200444698334, 0.543484628200531, 0.29247328639030457, 0.5261343717575073, 0.23446989059448242, 0.5474087595939636, 0.062012095004320145, 0.8189043998718262, 0.538780152797699, 0.6200674176216125, 0.43515679240226746, 0.24830776453018188, 0.341129869222641, 0.04290800169110298], [0.018064359202980995, 0.030848585069179535, 0.08071158826351166, 0.0676560178399086, 0.13447926938533783, 0.11551786214113235, 0.17043589055538177, 0.10128363966941833, 0.6618390679359436, 0.2855142652988434, 0.0971621423959732, 0.23388729989528656, 0.21859601140022278, 0.46025529503822327, 0.182326078414917], [0.04308566823601723, 0.03711610287427902, 0.06502576172351837, 0.10632220655679703, 0.09326566010713577, 0.08777783066034317, 0.3412204086780548, 0.6204424500465393, 0.8231819868087769, 0.09377399832010269, 0.1541169434785843, 0.21222646534442902, 0.11298450827598572, 0.15309588611125946, 0.11645805835723877], [0.07351326197385788, 0.05497964471578598, 0.07563240081071854, 0.32393333315849304, 0.057468246668577194, 0.2634526193141937, 0.3780488967895508, 0.7154850363731384, 0.7017503976821899, 0.20895157754421234, 0.29085400700569153, 0.06311048567295074, 0.03268700838088989, 0.14748480916023254, 0.03694311901926994], [0.15202973783016205, 0.07260382175445557, 0.07307075709104538, 0.01561899296939373, 0.03831832483410835, 0.04392734169960022, 0.07259247452020645, 0.03668325021862984, 0.315115749835968, 0.14016768336296082, 0.147903710603714, 0.09513753652572632, 0.08079177141189575, 0.04876280575990677, 0.1678115576505661], [0.20334205031394958, 0.03987862542271614, 0.2323523759841919, 0.08299659937620163, 0.11007620394229889, 0.049821991473436356, 0.05303451418876648, 0.020633194595575333, 0.20804192125797272, 0.621069610118866, 0.6013453006744385, 0.6998922824859619, 0.30664384365081787, 0.1810489445924759, 0.12484823167324066], [0.33830341696739197, 0.10967365652322769, 0.03348035365343094, 0.09579410403966904, 0.07735400646924973, 0.09874830394983292, 0.15181724727153778, 0.11190870404243469, 0.4600948095321655, 0.5270871520042419, 0.27297794818878174, 0.3748718500137329, 0.4609748125076294, 0.5019738078117371, 0.0790465772151947], [0.18835663795471191, 0.05185278132557869, 0.06106729805469513, 0.04512745887041092, 0.04466439411044121, 0.025852244347333908, 0.031750425696372986, 0.022515133023262024, 0.5077425837516785, 0.6734393835067749, 0.37964752316474915, 0.35936975479125977, 0.19831591844558716, 0.216437429189682, 0.2985125184059143], [0.5560556054115295, 0.47877317667007446, 0.15116584300994873, 0.40482252836227417, 0.04176756739616394, 0.04773563891649246, 0.13619393110275269, 0.07804162055253983, 0.07037016749382019, 0.5527278780937195, 0.486864298582077, 0.22204715013504028, 0.2625967860221863, 0.19855597615242004, 0.060070205479860306], [0.21585102379322052, 0.028776921331882477, 0.056070148944854736, 0.3207121789455414, 0.0078024002723395824, 0.016524065285921097, 0.3710367977619171, 0.14693383872509003, 0.12693363428115845, 0.6266815662384033, 0.6993157863616943, 0.5497558116912842, 0.14310741424560547, 0.3664083480834961, 0.047443971037864685], [0.28475576639175415, 0.10818006843328476, 0.08735410869121552, 0.329417884349823, 0.02252645045518875, 0.04752267897129059, 0.3733118176460266, 0.39454737305641174, 0.029050499200820923, 0.6059318780899048, 0.7311877012252808, 0.44807982444763184, 0.29598307609558105, 0.33838847279548645, 0.16424106061458588], [0.08968453854322433, 0.11453098803758621, 0.20413988828659058, 0.368092805147171, 0.07694120705127716, 0.048818718641996384, 0.12943927943706512, 0.036333490163087845, 0.04509947448968887, 0.25635746121406555, 0.2806471586227417, 0.5608395338058472, 0.1390012502670288, 0.28897786140441895, 0.04701472818851471], [0.05315335839986801, 0.017116300761699677, 0.1720367670059204, 0.3916313052177429, 0.05510414391756058, 0.2876152992248535, 0.22692401707172394, 0.14989952743053436, 0.3368622660636902, 0.0913245752453804, 0.3484038710594177, 0.3637443780899048, 0.007217096630483866, 0.103476881980896, 0.036375418305397034], [0.5125223994255066, 0.07351671159267426, 0.21591535210609436, 0.21059465408325195, 0.3288169205188751, 0.5466507077217102, 0.21618640422821045, 0.15017350018024445, 0.8681062459945679, 0.2442341297864914, 0.06865198910236359, 0.019835328683257103, 0.10077274590730667, 0.12228173017501831, 0.1682003289461136], [0.4846254289150238, 0.17620818316936493, 0.23995715379714966, 0.09631974995136261, 0.22585628926753998, 0.04512355476617813, 0.06700992584228516, 0.01503949984908104, 0.07369402050971985, 0.03452376648783684, 0.04930250719189644, 0.1451164036989212, 0.010093613527715206, 0.020862746983766556, 0.16003692150115967], [0.12189289927482605, 0.3658526837825775, 0.06606122851371765, 0.1638106107711792, 0.07819290459156036, 0.27624964714050293, 0.09599297493696213, 0.08126427978277206, 0.14055852591991425, 0.02327289618551731, 0.03783821687102318, 0.2963305115699768, 0.13405835628509521, 0.09205315262079239, 0.12166540324687958], [0.278896301984787, 0.1438806802034378, 0.46959513425827026, 0.3356979489326477, 0.3651174008846283, 0.1071292906999588, 0.18117688596248627, 0.20183299481868744, 0.29131460189819336, 0.13872042298316956, 0.021824011579155922, 0.06362087279558182, 0.34404000639915466, 0.13715140521526337, 0.1120462715625763], [0.2151702344417572, 0.2682046890258789, 0.2758127450942993, 0.20445802807807922, 0.06759822368621826, 0.058143485337495804, 0.21948587894439697, 0.1328936666250229, 0.04737214744091034, 0.09880322962999344, 0.06969184428453445, 0.0649414211511612, 0.09957331418991089, 0.08072139322757721, 0.15442174673080444], [0.10625648498535156, 0.3580685555934906, 0.2235240340232849, 0.2717205584049225, 0.14765356481075287, 0.1302592158317566, 0.182493656873703, 0.07402253895998001, 0.044094108045101166, 0.28373098373413086, 0.09141446650028229, 0.13240621984004974, 0.1622740924358368, 0.2716645896434784, 0.09359043836593628], [0.08181191235780716, 0.05183182656764984, 0.18780435621738434, 0.39972010254859924, 0.11086275428533554, 0.3443254232406616, 0.26716044545173645, 0.2157517671585083, 0.3917877972126007, 0.09846898168325424, 0.25891563296318054, 0.25942671298980713, 0.008535100147128105, 0.11220833659172058, 0.06895694881677628], [0.4507053792476654, 0.10277862101793289, 0.16431982815265656, 0.2027788907289505, 0.318918377161026, 0.4106469452381134, 0.24116744101047516, 0.1587350070476532, 0.8309358358383179, 0.2625651955604553, 0.047453198581933975, 0.009295494295656681, 0.07160880416631699, 0.07481760531663895, 0.19364440441131592], [0.5336673855781555, 0.18865860998630524, 0.19927646219730377, 0.10614699125289917, 0.21258802711963654, 0.035614922642707825, 0.07572873681783676, 0.021095039322972298, 0.08985494822263718, 0.061252057552337646, 0.05201297253370285, 0.10173538327217102, 0.008337927050888538, 0.017984798178076744, 0.15578274428844452], [0.11776354163885117, 0.337507039308548, 0.055947914719581604, 0.144154354929924, 0.09536269307136536, 0.2646341919898987, 0.10820504277944565, 0.0982295498251915, 0.1891198456287384, 0.027041049674153328, 0.03162495046854019, 0.2652260959148407, 0.10165920853614807, 0.07911970466375351, 0.1373925358057022], [0.20648452639579773, 0.10074114054441452, 0.42538517713546753, 0.26027214527130127, 0.3658106029033661, 0.09280957281589508, 0.23363487422466278, 0.27985435724258423, 0.3744349181652069, 0.1453229784965515, 0.02015594393014908, 0.05169985443353653, 0.3284047245979309, 0.12707991898059845, 0.12262601405382156], [0.019576620310544968, 0.03319034352898598, 0.0111849969252944, 0.010870445519685745, 0.03222370147705078, 0.13807591795921326, 0.0675833523273468, 0.0615379698574543, 0.013822048902511597, 0.008804764598608017, 0.004974161274731159, 0.01815059222280979, 0.1774466335773468, 0.06282598525285721, 0.15396134555339813]], [[0.07712388038635254, 0.042244281619787216, 0.004363007377833128, 0.0015959119191393256, 0.019252488389611244, 0.02118455246090889, 0.001846740604378283, 0.0012080060550943017, 0.0007866616360843182, 0.001261864323168993, 0.002815018408000469, 0.017323212698101997, 0.00286104716360569, 0.004067797679454088, 0.15733002126216888], [0.176344633102417, 0.3271441161632538, 0.08498391509056091, 0.04002806171774864, 0.06676299124956131, 0.008946515619754791, 0.012590638361871243, 0.0061616976745426655, 0.010515754111111164, 0.042563267052173615, 0.024306243285536766, 0.009260479360818863, 0.0002838150830939412, 0.0009972971165552735, 0.0829070582985878], [0.3345734477043152, 0.016792800277471542, 0.785018265247345, 0.16747814416885376, 0.3955724537372589, 0.09289640188217163, 0.041390396654605865, 0.004024161957204342, 0.04094661772251129, 0.023736434057354927, 0.20348279178142548, 0.041674140840768814, 0.012969214469194412, 0.03994787111878395, 0.04405270516872406], [0.027460135519504547, 0.0009503767942078412, 0.8045902252197266, 0.05251304432749748, 0.4111766219139099, 0.08071836084127426, 0.01928381621837616, 0.0005491983611136675, 0.029575586318969727, 0.001678029540926218, 0.033282194286584854, 0.007144003175199032, 0.012064780108630657, 0.008930332958698273, 0.0033295771572738886], [0.18455208837985992, 0.0566692017018795, 0.08522135764360428, 0.2798183560371399, 0.013304274529218674, 0.0006802850402891636, 0.09522412717342377, 0.0060977875255048275, 0.002369458321481943, 0.017453324049711227, 0.0036190226674079895, 2.9809654733981006e-05, 0.0002128492487827316, 0.0002820969675667584, 0.18610867857933044], [0.6536933779716492, 0.3485175371170044, 0.2007695585489273, 0.8106443881988525, 0.12433423846960068, 0.008092332631349564, 0.6807736158370972, 0.40895989537239075, 0.04516575112938881, 0.1387551873922348, 0.004862201400101185, 0.0003120531910099089, 0.00022667655139230192, 0.00031860917806625366, 0.07640787214040756], [0.08564082533121109, 0.05155009403824806, 0.10021068900823593, 0.5880905985832214, 0.0823356956243515, 0.0626063123345375, 0.7381499409675598, 0.566346287727356, 0.04188016802072525, 0.02469027414917946, 0.004355741199105978, 0.00042968738125637174, 2.4299803044414148e-05, 2.7212277927901596e-05, 0.001896930974908173], [0.03975995257496834, 0.012421448715031147, 0.08890707790851593, 0.605818510055542, 0.05048904940485954, 0.017510779201984406, 0.24702893197536469, 0.39587050676345825, 0.06098005548119545, 0.052625395357608795, 0.013424866832792759, 0.0005194320692680776, 0.000250102486461401, 0.0003063087642658502, 0.0010793216060847044], [0.11902385950088501, 0.011114073917269707, 0.22151720523834229, 0.2006509006023407, 0.03878694027662277, 0.01363028772175312, 0.3268369734287262, 0.04311302676796913, 0.8067907094955444, 0.34777864813804626, 0.25920552015304565, 0.09021251648664474, 0.035271789878606796, 0.0031717135570943356, 0.004271878860890865], [0.006270309444516897, 0.0001492560259066522, 0.00045137249981053174, 0.0007612273329868913, 7.476524478988722e-05, 0.013270817697048187, 0.04344405606389046, 0.014117085374891758, 0.6041488647460938, 0.07304701954126358, 0.010559855960309505, 0.0026350386906415224, 0.02638809196650982, 0.002994539914652705, 0.00020572090579662472], [0.002078789984807372, 0.000502656155731529, 0.00018232718866784126, 0.0008548289188183844, 0.0009249084978364408, 0.02029070071876049, 0.012032798491418362, 0.024348178878426552, 0.2300865352153778, 0.10343841463327408, 0.007660495117306709, 0.0012821657583117485, 0.0114271380007267, 0.0009412667131982744, 7.524124521296471e-05], [0.022463228553533554, 0.0013134862529113889, 0.00013891702110413462, 0.002816978842020035, 0.0011811865260824561, 0.0014538302784785628, 0.0005458829691633582, 0.0004073161107953638, 0.000992793939076364, 0.626685380935669, 0.1310541182756424, 0.1785772740840912, 0.1327074021100998, 0.014590581879019737, 3.459410072537139e-05], [0.004299411084502935, 0.00014757749158889055, 0.0013493087608367205, 0.003552102018147707, 0.004041418433189392, 0.004232631530612707, 0.00022051982523407787, 5.3625211876351386e-05, 0.008671559393405914, 0.2003454566001892, 0.2010745257139206, 0.20048564672470093, 0.327506959438324, 0.12215141952037811, 7.573522452730685e-05], [0.011497906409204006, 0.0014132088981568813, 0.002270179335027933, 0.006387166678905487, 5.5530636018374935e-05, 0.0020248510409146547, 0.0021348590962588787, 0.001147052156738937, 0.0024277162738144398, 0.3687064051628113, 0.5298402905464172, 0.006611559074372053, 0.3372868299484253, 0.2915361225605011, 0.0002606022753752768], [0.043351031839847565, 0.015730101615190506, 0.006545424461364746, 0.11301398277282715, 0.001535893650725484, 0.0002994980022776872, 0.002417969051748514, 0.0027875620871782303, 0.007663458585739136, 0.4366588592529297, 0.29866132140159607, 0.03879629448056221, 0.0005757116014137864, 0.10755223035812378, 0.15693426132202148], [0.05824243649840355, 0.00918568018823862, 0.004823020659387112, 0.12202360481023788, 0.001364732626825571, 0.009540650062263012, 0.017077280208468437, 0.02250218391418457, 0.031557418406009674, 0.39489659667015076, 0.4118596911430359, 0.4739699363708496, 0.04330656677484512, 0.22410848736763, 0.009354491718113422], [0.10114194452762604, 0.055991608649492264, 0.0056193675845861435, 0.044799599796533585, 0.005612906999886036, 0.0018076150445267558, 0.0035521595273166895, 0.003050913568586111, 0.014126029796898365, 0.18568304181098938, 0.044660091400146484, 0.8178999423980713, 0.12312521040439606, 0.22830259799957275, 0.0015339198289439082], [0.17329555749893188, 0.022842630743980408, 0.03050464764237404, 0.3040459156036377, 0.023058682680130005, 0.05675753578543663, 0.012084487825632095, 0.018060212954878807, 0.012510768137872219, 0.4205268621444702, 0.403047114610672, 0.5196431279182434, 0.14466160535812378, 0.15726853907108307, 0.003281315555796027], [0.21814380586147308, 0.013853680342435837, 0.0011839027283713222, 0.02006133459508419, 0.0059941732324659824, 0.004335244186222553, 0.0006587213138118386, 0.0008069095201790333, 6.766151636838913e-05, 0.4439576268196106, 0.16648612916469574, 0.7347545623779297, 0.19459886848926544, 0.05657987296581268, 0.0006026092451065779], [0.034262340515851974, 0.0017182001611217856, 0.005656392779201269, 0.017169898375868797, 0.0156857930123806, 0.01468763966113329, 0.0007699507405050099, 0.00017933807976078242, 0.002019587904214859, 0.09474337100982666, 0.21286551654338837, 0.39837440848350525, 0.44769343733787537, 0.30061447620391846, 0.0009720441303215921], [0.1974877417087555, 0.05350746586918831, 0.02080627717077732, 0.07140190154314041, 0.0007820951868779957, 0.021851971745491028, 0.023295408114790916, 0.011020028032362461, 0.0015720969531685114, 0.3204348385334015, 0.5890824198722839, 0.011122598312795162, 0.40923523902893066, 0.5521805882453918, 0.009284045547246933], [0.04384012520313263, 0.020103074610233307, 0.00601673498749733, 0.10121199488639832, 0.0015372235793620348, 0.00047879578778520226, 0.0028034253045916557, 0.0035304632037878036, 0.0019347126362845302, 0.15543726086616516, 0.10060140490531921, 0.012154079042375088, 0.00020098914683330804, 0.049742307513952255, 0.15931616723537445], [0.33183732628822327, 0.07794758677482605, 0.02364480309188366, 0.3878714144229889, 0.007764760870486498, 0.055411770939826965, 0.07855504751205444, 0.09397301822900772, 0.02721172571182251, 0.38145557045936584, 0.42047446966171265, 0.5078706741333008, 0.03859835863113403, 0.25985077023506165, 0.0625251829624176], [0.4473247230052948, 0.3730325996875763, 0.029895052313804626, 0.15908104181289673, 0.02762797847390175, 0.008889964781701565, 0.016516737639904022, 0.012883803807199001, 0.01523641124367714, 0.22003965079784393, 0.05771813541650772, 0.8456536531448364, 0.1770154982805252, 0.31127816438674927, 0.007925343699753284], [0.2188224196434021, 0.06026163697242737, 0.01674255169928074, 0.1205059364438057, 0.017392028123140335, 0.033714599907398224, 0.013199009001255035, 0.035441260784864426, 0.006878681946545839, 0.5097362399101257, 0.5390803217887878, 0.7098195552825928, 0.20610427856445312, 0.34404870867729187, 0.06464894115924835]], [[0.24012988805770874, 0.6692726612091064, 0.08029869198799133, 0.41845017671585083, 0.08128808438777924, 0.09738753736019135, 0.15100885927677155, 0.2691691815853119, 0.013517879880964756, 0.21848294138908386, 0.16758716106414795, 0.12734578549861908, 0.32224464416503906, 0.12471552193164825, 0.07385692000389099], [0.13747748732566833, 0.012865100987255573, 0.3056560158729553, 0.3759651184082031, 0.20075583457946777, 0.056869279593229294, 0.27502477169036865, 0.09038521349430084, 0.09535539150238037, 0.27579623460769653, 0.15189220011234283, 0.6071571111679077, 0.0820951759815216, 0.09481122344732285, 0.09779953956604004], [0.007538634352385998, 0.02957071363925934, 0.011847163550555706, 0.055522944778203964, 0.04100131243467331, 0.031534671783447266, 0.06567902117967606, 0.09044305235147476, 0.007193693891167641, 0.06334451586008072, 0.07378207892179489, 0.07786792516708374, 0.28214019536972046, 0.08070375770330429, 0.20607011020183563], [0.005881547927856445, 0.008371960371732712, 0.010823756456375122, 0.024797217920422554, 0.024142105132341385, 0.01083815935999155, 0.008304014801979065, 0.006388344801962376, 0.009114595130085945, 0.022048065438866615, 0.1306026130914688, 0.23451638221740723, 0.3918500244617462, 0.08784151822328568, 0.2650633752346039], [0.20629070699214935, 0.2529377341270447, 0.028870999813079834, 0.049127642065286636, 0.04690879210829735, 0.11594393104314804, 0.15515393018722534, 0.06585636734962463, 0.0420556403696537, 0.1996643990278244, 0.028717953711748123, 0.7190893292427063, 0.30376943945884705, 0.22654840350151062, 0.12926629185676575], [0.01586613617837429, 0.15566423535346985, 0.015082520432770252, 0.009204044006764889, 0.002680863719433546, 0.07106906920671463, 0.08370621502399445, 0.05749649554491043, 0.03059268370270729, 0.012942377477884293, 0.0011753733269870281, 0.00916373822838068, 0.0020018015056848526, 0.049308281391859055, 0.19197486340999603], [0.03849078342318535, 0.08146823942661285, 0.03517843410372734, 0.025976145640015602, 0.02364599145948887, 0.1389763057231903, 0.02619975060224533, 0.034312427043914795, 0.02985706366598606, 0.029806064441800117, 0.00684476038441062, 0.03280223533511162, 0.030126189813017845, 0.10321015119552612, 0.23163792490959167], [0.2772977352142334, 0.05161405727267265, 0.04358568787574768, 0.047931231558322906, 0.04583681374788284, 0.08128579705953598, 0.15782645344734192, 0.0856042429804802, 0.10767779499292374, 0.11355230212211609, 0.041377030313014984, 0.252811074256897, 0.05780917406082153, 0.19973745942115784, 0.22427907586097717], [0.023119861260056496, 0.02037731558084488, 0.0453791618347168, 0.1060030460357666, 0.006244942545890808, 0.0085020512342453, 0.012060720473527908, 0.014560479670763016, 0.00689319521188736, 0.011241135187447071, 0.023835573345422745, 0.02693312056362629, 0.011436404660344124, 0.019489392638206482, 0.30997538566589355], [0.045414164662361145, 0.005229660775512457, 0.011418518610298634, 0.009312640875577927, 0.0002147085906472057, 0.12653864920139313, 0.05854451283812523, 0.11896014213562012, 0.0156405046582222, 0.010270207189023495, 0.0032450463622808456, 0.015787174925208092, 0.011106730438768864, 0.007675709668546915, 0.3779195249080658], [0.007367350626736879, 0.012884993106126785, 0.01019106525927782, 0.011957473121583462, 0.054886650294065475, 0.09750530868768692, 0.029414953663945198, 0.08492925018072128, 0.17440666258335114, 0.003643231000751257, 0.00105402956251055, 0.02280060388147831, 0.0010922637302428484, 0.005130939185619354, 0.09500079602003098], [0.02996714971959591, 0.028387926518917084, 0.16122521460056305, 0.0898616760969162, 0.06381779164075851, 0.20551051199436188, 0.13175098598003387, 0.562389075756073, 0.04834860563278198, 0.013581722043454647, 0.03991095721721649, 0.10736902058124542, 0.03830268979072571, 0.05736052244901657, 0.27213579416275024], [0.03571658954024315, 0.012061648070812225, 0.08574458211660385, 0.022463832050561905, 0.12578466534614563, 0.07826194912195206, 0.06577891856431961, 0.13274507224559784, 0.06591502577066422, 0.05002211779356003, 0.03129255399107933, 0.27911075949668884, 0.31601372361183167, 0.10930214822292328, 0.30993908643722534], [0.04630875587463379, 0.03141915798187256, 0.03061339072883129, 0.007028677500784397, 0.008451082743704319, 0.02540888637304306, 0.012118873186409473, 0.09331455826759338, 0.0033372503239661455, 0.01357665192335844, 0.0069510783068835735, 0.017483821138739586, 0.033454760909080505, 0.014270796440541744, 0.44127020239830017], [0.1722828894853592, 0.15122008323669434, 0.056102070957422256, 0.09136570990085602, 0.02421834133565426, 0.045343294739723206, 0.034619707614183426, 0.030837759375572205, 0.019798463210463524, 0.04411705583333969, 0.05331422761082649, 0.09423463046550751, 0.1436629444360733, 0.13433872163295746, 0.1229754090309143], [0.022473091259598732, 0.0489150770008564, 0.010993139818310738, 0.03897916153073311, 0.003662768052890897, 0.002051829593256116, 0.0037445707712322474, 0.016557298600673676, 0.014907213859260082, 0.004300208762288094, 0.004852794576436281, 0.0027131394017487764, 0.016001524403691292, 0.008091894909739494, 0.25544992089271545], [0.08012817800045013, 0.2898695766925812, 0.022246699780225754, 0.06057273969054222, 0.025327028706669807, 0.02957070618867874, 0.04002644121646881, 0.019245512783527374, 0.01995179057121277, 0.020330116152763367, 0.006697094067931175, 0.015452835708856583, 0.014569609425961971, 0.04013357311487198, 0.2585589587688446], [0.01832924410700798, 0.023918962106108665, 0.024782713502645493, 0.033514510840177536, 0.050549402832984924, 0.013098560273647308, 0.023091215640306473, 0.030541786924004555, 0.1064886748790741, 0.006106832530349493, 0.0024854408111423254, 0.018918434157967567, 0.0075035663321614265, 0.009370497427880764, 0.21452490985393524], [0.027254067361354828, 0.020437292754650116, 0.14233240485191345, 0.08538791537284851, 0.03242940828204155, 0.0897425189614296, 0.08476056158542633, 0.2620556950569153, 0.02126460149884224, 0.023079702630639076, 0.03143052011728287, 0.04489685967564583, 0.046720463782548904, 0.03604652360081673, 0.23038896918296814], [0.042377930134534836, 0.017293933779001236, 0.08730384707450867, 0.030179454013705254, 0.12187745422124863, 0.05139933153986931, 0.047754548490047455, 0.066692054271698, 0.06521614640951157, 0.05196157470345497, 0.028108397498726845, 0.17703385651111603, 0.22747749090194702, 0.06955988705158234, 0.28824013471603394], [0.03372317552566528, 0.030876630917191505, 0.025082340463995934, 0.008588657714426517, 0.007454049773514271, 0.009771045297384262, 0.010381288826465607, 0.041183773428201675, 0.004549690056592226, 0.01619204692542553, 0.0060179769061505795, 0.009672058746218681, 0.022905999794602394, 0.009750566445291042, 0.30946746468544006], [0.18900562822818756, 0.14908763766288757, 0.05840699374675751, 0.10216160118579865, 0.03072887472808361, 0.04109037667512894, 0.03799780085682869, 0.02909342385828495, 0.03500371053814888, 0.0757574513554573, 0.061073921620845795, 0.09956928342580795, 0.10441071540117264, 0.14136889576911926, 0.13095542788505554], [0.014150185510516167, 0.03789284825325012, 0.007744992151856422, 0.02556411363184452, 0.0037681234534829855, 0.001123085618019104, 0.002939486177638173, 0.010072565637528896, 0.019109029322862625, 0.003645692951977253, 0.0027771664317697287, 0.002490789396688342, 0.007166225463151932, 0.005180294159799814, 0.2058444321155548], [0.0469474196434021, 0.1743137687444687, 0.021908296272158623, 0.046387769281864166, 0.02985612489283085, 0.019742406904697418, 0.040140021592378616, 0.01437240932136774, 0.02856219932436943, 0.018488112837076187, 0.004136314615607262, 0.01038376335054636, 0.009851893410086632, 0.026245350018143654, 0.22488054633140564], [0.00832295510917902, 0.021339448168873787, 0.00394090311601758, 0.002333499025553465, 0.05547437444329262, 0.007243151310831308, 0.011641105636954308, 0.0331541933119297, 0.010278979316353798, 0.011881710961461067, 0.001766148954629898, 0.04899042472243309, 0.01878243498504162, 0.01244808267802, 0.15685127675533295]]], [[[0.04773104563355446, 0.01963546872138977, 0.16452182829380035, 0.04063690826296806, 0.1849776655435562, 0.08088860660791397, 0.11659693717956543, 0.038044340908527374, 0.2744975686073303, 0.003083554795011878, 0.019721103832125664, 0.08137688785791397, 0.0169991385191679, 0.03939461708068848, 0.14168404042720795], [0.09676018357276917, 0.018249453976750374, 0.657112717628479, 0.5890088677406311, 0.5712416768074036, 0.2744671702384949, 0.48642322421073914, 0.26345524191856384, 0.23708243668079376, 0.03475205600261688, 0.15204745531082153, 0.0676480308175087, 0.050043635070323944, 0.0665324404835701, 0.036993421614170074], [0.04065309092402458, 0.0025235058274120092, 0.11838234961032867, 0.27863210439682007, 0.37560757994651794, 0.7046668529510498, 0.12516380846500397, 0.1912177950143814, 0.14992743730545044, 0.05949303135275841, 0.056387268006801605, 0.04353337734937668, 0.17471297085285187, 0.07017815858125687, 0.12025584280490875], [0.015422305092215538, 0.000844803755171597, 0.015767300501465797, 0.11098357290029526, 0.273564875125885, 0.3235251009464264, 0.14805495738983154, 0.17132841050624847, 0.25568780303001404, 0.034506767988204956, 0.046862825751304626, 0.03818853572010994, 0.025031423196196556, 0.027911247685551643, 0.009120252914726734], [0.01866327039897442, 0.11290711164474487, 0.007440958172082901, 0.031009642407298088, 0.059622399508953094, 0.035299621522426605, 0.012064317241311073, 0.17540854215621948, 0.06399405747652054, 0.010346408933401108, 0.023967623710632324, 0.006549614481627941, 0.015476463362574577, 0.017944032326340675, 0.15624091029167175], [0.115133136510849, 0.5564319491386414, 0.0024013265501707792, 0.014839398674666882, 0.027623601257801056, 0.003712957026436925, 0.11139625310897827, 0.4320802688598633, 0.18111301958560944, 0.025198934599757195, 0.05914938822388649, 0.029404014348983765, 0.1131783202290535, 0.1630096137523651, 0.14384765923023224], [0.047323077917099, 0.01987922191619873, 0.021367410197854042, 0.0816798061132431, 0.11104802042245865, 0.01310664601624012, 0.37855657935142517, 0.16697411239147186, 0.31461480259895325, 0.04616151005029678, 0.27547621726989746, 0.04939346760511398, 0.02232075110077858, 0.15515512228012085, 0.01579722762107849], [0.13229456543922424, 0.031869739294052124, 0.26943540573120117, 0.2586674690246582, 0.3796730637550354, 0.127562016248703, 0.20277942717075348, 0.05910756066441536, 0.14354895055294037, 0.08293455094099045, 0.2214740365743637, 0.23150987923145294, 0.18035069108009338, 0.2860051393508911, 0.07895194739103317], [0.09224988520145416, 0.07457923144102097, 0.05282874405384064, 0.09438028931617737, 0.06849074363708496, 0.012997711077332497, 0.007214613724499941, 0.004257954657077789, 0.2309093326330185, 0.38276976346969604, 0.5917518734931946, 0.7830951809883118, 0.8438952565193176, 0.7586230039596558, 0.04145537316799164], [0.014161140657961369, 0.027171263471245766, 0.0029068312142044306, 0.020549731329083443, 0.0005743438960053027, 0.00417140731588006, 0.003657599212601781, 0.00956815481185913, 0.34446486830711365, 0.5171273946762085, 0.39057764410972595, 0.2845093309879303, 0.1669711321592331, 0.5306525230407715, 0.015455210581421852], [0.02566671371459961, 0.00907080341130495, 0.0006065603229217231, 0.03001752682030201, 0.00023783017240930349, 0.0005533608491532505, 0.013808660209178925, 0.003767948364838958, 0.06461481004953384, 0.1359771490097046, 0.08153439313173294, 0.572087287902832, 0.36045318841934204, 0.44234389066696167, 0.0030113777611404657], [0.03087739646434784, 0.012099061161279678, 0.004942088853567839, 0.038267359137535095, 0.0023591304197907448, 0.0037323227152228355, 0.04966888204216957, 0.012427400797605515, 0.16158415377140045, 0.020882699638605118, 0.05600592866539955, 0.367767333984375, 0.24262923002243042, 0.38281354308128357, 0.00973587203770876], [0.04249054566025734, 0.0069285486824810505, 0.006088858004659414, 0.044397544115781784, 0.05390672758221626, 0.006144464481621981, 0.018320903182029724, 0.01545354351401329, 0.05193139612674713, 0.03221629932522774, 0.02379259280860424, 0.27246853709220886, 0.22103002667427063, 0.23179520666599274, 0.005589436274021864], [0.04184036701917648, 0.03700190782546997, 0.008264865726232529, 0.02439146116375923, 0.00799429602921009, 0.12502151727676392, 0.05032283812761307, 0.18101848661899567, 0.07329469919204712, 0.08409427851438522, 0.10790428519248962, 0.011960207484662533, 0.20496119558811188, 0.19276422262191772, 0.0069670299999415874], [0.06364590674638748, 0.06483624875545502, 0.015260975807905197, 0.1278582364320755, 0.006228389218449593, 0.02756887674331665, 0.020600903779268265, 0.015440343879163265, 0.018087223172187805, 0.017098410055041313, 0.025406692177057266, 0.0007098353235051036, 0.00014885497512295842, 0.0013503700029104948, 0.15608660876750946], [0.6220619678497314, 0.6306124329566956, 0.6737340092658997, 0.49940165877342224, 0.1517823040485382, 0.8503586649894714, 0.705633282661438, 0.6629571914672852, 0.11157920956611633, 0.39899003505706787, 0.3173867464065552, 0.027327625080943108, 0.014980590902268887, 0.009274562820792198, 0.08523338288068771], [0.15005189180374146, 0.04609784111380577, 0.17501141130924225, 0.21113994717597961, 0.26919078826904297, 0.6422000527381897, 0.7493206858634949, 0.2162598967552185, 0.010351919569075108, 0.09728528559207916, 0.09688232094049454, 0.028558582067489624, 0.10305432975292206, 0.05914681404829025, 0.11260810494422913], [0.09041088819503784, 0.052050016820430756, 0.08856991678476334, 0.2977358102798462, 0.04025371000170708, 0.3506464660167694, 0.6434463858604431, 0.25059518218040466, 0.01933867670595646, 0.04819375276565552, 0.07508239895105362, 0.04970608279109001, 0.02890131063759327, 0.02355407178401947, 0.12558245658874512], [0.18765486776828766, 0.021713200956583023, 0.21844394505023956, 0.3042432367801666, 0.17823228240013123, 0.1673380434513092, 0.8088975548744202, 0.46762967109680176, 0.05706785246729851, 0.009645337238907814, 0.0322297103703022, 0.09777479618787766, 0.08048812299966812, 0.10106904059648514, 0.17228879034519196], [0.4792143702507019, 0.09839366376399994, 0.1882246881723404, 0.4093988239765167, 0.7147246599197388, 0.24897223711013794, 0.4705742597579956, 0.4205995500087738, 0.01958448253571987, 0.026842152699828148, 0.02239188365638256, 0.15106931328773499, 0.08969185501337051, 0.10003618896007538, 0.1635625958442688], [0.40625429153442383, 0.3796224594116211, 0.2515096962451935, 0.36165565252304077, 0.24774380028247833, 0.8824228644371033, 0.8048573136329651, 0.857955813407898, 0.058371078222990036, 0.07109472155570984, 0.11402199417352676, 0.0021524245385080576, 0.019929109141230583, 0.030590593814849854, 0.11712031066417694], [0.04390633478760719, 0.032843075692653656, 0.010515165515244007, 0.11869800090789795, 0.005461697466671467, 0.023131608963012695, 0.01705162413418293, 0.008547519333660603, 0.003713170997798443, 0.008410640992224216, 0.009457322768867016, 0.00015943740436341614, 3.361727431183681e-05, 0.0002994383394252509, 0.1532706469297409], [0.6348351836204529, 0.5127235651016235, 0.5931673645973206, 0.5543242692947388, 0.12377271056175232, 0.8264753222465515, 0.6941898465156555, 0.5687963962554932, 0.03150533139705658, 0.12843358516693115, 0.11884576827287674, 0.005231617949903011, 0.0018767286092042923, 0.0011644444894045591, 0.11210005730390549], [0.10790421068668365, 0.016916295513510704, 0.09771728515625, 0.22749783098697662, 0.26325535774230957, 0.49138790369033813, 0.6275916695594788, 0.08931886404752731, 0.0033968419302254915, 0.024402111768722534, 0.018104346469044685, 0.003288157982751727, 0.010537534020841122, 0.006979967001825571, 0.12102893739938736], [0.028179557994008064, 0.011468129232525826, 0.016789404675364494, 0.00803140178322792, 0.00952040497213602, 0.02960360422730446, 0.24957160651683807, 0.03544437885284424, 0.005487674381583929, 0.0028927521780133247, 0.005656986031681299, 0.0040698484517633915, 0.04730471968650818, 0.0667993351817131, 0.1372966766357422]], [[0.11859580129384995, 0.07486707717180252, 0.21083025634288788, 0.32276296615600586, 0.08426652103662491, 0.03581860288977623, 0.24113436043262482, 0.608397364616394, 0.13584911823272705, 0.45509204268455505, 0.594833254814148, 0.30372148752212524, 0.8448506593704224, 0.7470672726631165, 0.09252076596021652], [0.04140070080757141, 0.00858838576823473, 0.11639615148305893, 0.1280786097049713, 0.2722368836402893, 0.21025919914245605, 0.4195333421230316, 0.631318211555481, 0.6560773253440857, 0.29341432452201843, 0.6862512230873108, 0.7675639986991882, 0.8915717005729675, 0.8601328730583191, 0.23356862366199493], [0.23441848158836365, 0.1666196584701538, 0.16664288938045502, 0.25857093930244446, 0.13334479928016663, 0.17917701601982117, 0.8257887363433838, 0.7395779490470886, 0.6802234053611755, 0.8125103712081909, 0.671615719795227, 0.8831866383552551, 0.6773648858070374, 0.7102506160736084, 0.08689045161008835], [0.24967892467975616, 0.48421844840049744, 0.036505091935396194, 0.17128480970859528, 0.01777578890323639, 0.09479225426912308, 0.36135032773017883, 0.0868472084403038, 0.16740600764751434, 0.523710310459137, 0.24439233541488647, 0.42307958006858826, 0.6259368062019348, 0.3662186563014984, 0.20058651268482208], [0.28931790590286255, 0.4439229369163513, 0.24370647966861725, 0.6020305752754211, 0.17363131046295166, 0.338454008102417, 0.5701692700386047, 0.33999428153038025, 0.68463534116745, 0.8701388239860535, 0.7831944823265076, 0.9611375331878662, 0.9679895043373108, 0.9072677493095398, 0.0468842089176178], [0.1225743219256401, 0.062406159937381744, 0.03387807682156563, 0.02868799865245819, 0.01787530817091465, 0.04143121838569641, 0.5920179486274719, 0.08798510581254959, 0.2968905568122864, 0.7129084467887878, 0.4609105885028839, 0.29060137271881104, 0.7909923791885376, 0.5701599717140198, 0.13614380359649658], [0.0705394446849823, 0.02209068462252617, 0.0211530439555645, 0.008882923051714897, 0.0033682750072330236, 0.08319123089313507, 0.11070933192968369, 0.0025125632528215647, 0.10380591452121735, 0.17744502425193787, 0.10391969978809357, 0.12427430599927902, 0.5562515258789062, 0.49710196256637573, 0.3223192095756531], [0.15847322344779968, 0.015464702621102333, 0.13866224884986877, 0.053395166993141174, 0.03494010120630264, 0.13738934695720673, 0.02684560976922512, 0.03214175999164581, 0.5759801864624023, 0.1755424290895462, 0.13409779965877533, 0.035038210451602936, 0.6489107012748718, 0.4460716247558594, 0.4074119031429291], [0.00857736449688673, 0.012718217447400093, 0.01174219325184822, 0.012934550642967224, 0.006551709491759539, 0.24597492814064026, 0.030029013752937317, 0.05923602730035782, 0.04650798439979553, 0.02447274886071682, 0.019859377294778824, 0.003505804343149066, 0.04937520623207092, 0.05625420808792114, 0.28037816286087036], [0.0015372766647487879, 0.015295127406716347, 0.018696704879403114, 0.004789609462022781, 0.19481690227985382, 0.04769033566117287, 0.01355075929313898, 0.02196505106985569, 0.08700259774923325, 0.020393503829836845, 0.02400771528482437, 0.18789233267307281, 0.15418098866939545, 0.08713112771511078, 0.19334079325199127], [0.04759770259261131, 0.04375501722097397, 0.02714523859322071, 0.05194481834769249, 0.05246514454483986, 0.14355513453483582, 0.17152011394500732, 0.14246520400047302, 0.1098044142127037, 0.013531663455069065, 0.008927365764975548, 0.03807468339800835, 0.10050502419471741, 0.02236531302332878, 0.3381733298301697], [0.10647730529308319, 0.04246760904788971, 0.08123224973678589, 0.13003453612327576, 0.07854175567626953, 0.24148082733154297, 0.6790831685066223, 0.7492273449897766, 0.28685522079467773, 0.03681188449263573, 0.15954196453094482, 0.2672117054462433, 0.11099980026483536, 0.04468434303998947, 0.4826459586620331], [0.2962004542350769, 0.47284576296806335, 0.11245852708816528, 0.23689918220043182, 0.10807513445615768, 0.8532499074935913, 0.5788733959197998, 0.6375027894973755, 0.33168625831604004, 0.06381742656230927, 0.004373080097138882, 0.015940984711050987, 0.3371734917163849, 0.06828418374061584, 0.21185840666294098], [0.3828115463256836, 0.12613584101200104, 0.47516295313835144, 0.4473835527896881, 0.17031393945217133, 0.6938255429267883, 0.7945614457130432, 0.34594833850860596, 0.5323623418807983, 0.34808266162872314, 0.11382761597633362, 0.1349307745695114, 0.013382190838456154, 0.0600610226392746, 0.30783677101135254], [0.7362364530563354, 0.8323087096214294, 0.9336822032928467, 0.7739728689193726, 0.8897883296012878, 0.9609381556510925, 0.9334329962730408, 0.9553548693656921, 0.7747710943222046, 0.4005538523197174, 0.5586770176887512, 0.25099167227745056, 0.4200068712234497, 0.1631680577993393, 0.06528117507696152], [0.07449624687433243, 0.061402805149555206, 0.09389828145503998, 0.048646457493305206, 0.024208296090364456, 0.10819891840219498, 0.10563155263662338, 0.1243496686220169, 0.048523951321840286, 0.14693649113178253, 0.06614942103624344, 0.0066792843863368034, 0.2858017086982727, 0.04383772611618042, 0.15409637987613678], [0.02467108517885208, 0.049052223563194275, 0.08135215938091278, 0.013768618926405907, 0.01176412496715784, 0.15210841596126556, 0.004693970084190369, 0.0041237217374145985, 0.018837640061974525, 0.03490369766950607, 0.036496780812740326, 0.0011750683188438416, 0.018557026982307434, 0.02382473833858967, 0.22122804820537567], [0.012043171562254429, 0.03080524504184723, 0.02248452790081501, 0.008785543963313103, 0.00550604984164238, 0.05614035204052925, 0.015958979725837708, 0.01727765053510666, 0.03423915058374405, 0.017799094319343567, 0.029912255704402924, 0.01144923735409975, 0.09533664584159851, 0.02436906285583973, 0.20283196866512299], [0.01959865354001522, 0.003073114436119795, 0.06498773396015167, 0.027286570519208908, 0.019540993496775627, 0.052237618714571, 0.08713454008102417, 0.28957968950271606, 0.3906492590904236, 0.044482238590717316, 0.17143161594867706, 0.1301742047071457, 0.10445850342512131, 0.03699616342782974, 0.2442801147699356], [0.11208802461624146, 0.11668127030134201, 0.09828943759202957, 0.10754654556512833, 0.015885351225733757, 0.38998937606811523, 0.183034285902977, 0.3230077624320984, 0.20506803691387177, 0.08733018487691879, 0.007069121580570936, 0.010435528121888638, 0.30221423506736755, 0.047303054481744766, 0.19994190335273743], [0.1682588905096054, 0.051582805812358856, 0.4415716230869293, 0.2735750675201416, 0.07878735661506653, 0.06776249408721924, 0.15038572251796722, 0.03211068734526634, 0.6709542274475098, 0.37688353657722473, 0.1879340261220932, 0.04096703231334686, 0.011627858504652977, 0.03471425548195839, 0.19384095072746277], [0.8205305933952332, 0.9214023947715759, 0.9559677839279175, 0.7988566160202026, 0.9105063080787659, 0.9672437906265259, 0.9506043195724487, 0.9735420346260071, 0.9064961075782776, 0.6156813502311707, 0.6370130777359009, 0.18943972885608673, 0.3681671619415283, 0.1194160059094429, 0.08283783495426178], [0.10534824430942535, 0.08027994632720947, 0.1381307989358902, 0.07063161581754684, 0.01806548424065113, 0.10409632325172424, 0.12885765731334686, 0.2072904407978058, 0.09267445653676987, 0.23836983740329742, 0.11645739525556564, 0.006059943698346615, 0.1595546454191208, 0.017974214628338814, 0.14464683830738068], [0.026579611003398895, 0.02949470281600952, 0.04954056441783905, 0.017031243070960045, 0.008355016820132732, 0.09075918793678284, 0.0036468924954533577, 0.0022332987282425165, 0.050134338438510895, 0.049380820244550705, 0.028885982930660248, 0.0007559077348560095, 0.015549316070973873, 0.013319555670022964, 0.1734825074672699], [0.05047497898340225, 0.027197130024433136, 0.11470095813274384, 0.007973222993314266, 0.12679167091846466, 0.4866730570793152, 0.17132264375686646, 0.15032453835010529, 0.14889459311962128, 0.01696154847741127, 0.0735161080956459, 0.0034290377516299486, 0.05194668471813202, 0.06144191324710846, 0.13309471309185028]], [[0.005987181328237057, 0.0011158415582031012, 0.0026756690349429846, 0.0011391430161893368, 0.0021053741220384836, 0.0005449134623631835, 0.0017384873935952783, 0.000736464629881084, 0.00014482461847364902, 0.0008784460369497538, 0.0008941806154325604, 0.0009559267782606184, 0.00015614555741194636, 0.00044419756159186363, 0.16329224407672882], [0.3448674976825714, 0.07203025370836258, 0.011963781900703907, 0.012941744178533554, 0.011539866216480732, 0.003333584638312459, 0.005511423572897911, 0.0016478801844641566, 0.003020848147571087, 0.006189296022057533, 0.0020935258362442255, 0.00048376841004937887, 8.994764357339591e-05, 0.00040787423495203257, 0.2113737165927887], [0.44219815731048584, 0.8124432563781738, 0.1900549679994583, 0.3808274269104004, 0.045300956815481186, 0.024617541581392288, 0.0172295980155468, 0.03488133102655411, 0.004235385917127132, 0.05999733507633209, 0.03787413239479065, 0.0011567235924303532, 0.0017442036187276244, 0.008845857344567776, 0.004224383272230625], [0.07874103635549545, 0.02866651676595211, 0.3287397623062134, 0.27984437346458435, 0.10563887655735016, 0.003691220423206687, 0.005916049238294363, 0.0007406381191685796, 0.0005066083394922316, 0.0481056272983551, 0.029072491452097893, 0.000652547983918339, 0.0003529583918862045, 0.0009863339364528656, 0.002192106796428561], [0.030638281255960464, 0.02597089111804962, 0.6577842831611633, 0.16596756875514984, 0.48041173815727234, 0.6114144921302795, 0.028207998722791672, 0.053615398705005646, 0.1417267620563507, 0.03454216569662094, 0.023575417697429657, 0.004873087164014578, 0.0009616028983145952, 0.00223313900642097, 0.0011337294708937407], [0.29477018117904663, 0.14754106104373932, 0.8534399271011353, 0.9182198643684387, 0.6083860993385315, 0.9389832019805908, 0.12579986453056335, 0.03590020909905434, 0.012173496186733246, 0.16479530930519104, 0.15366923809051514, 0.0035958383232355118, 0.002988115418702364, 0.026292480528354645, 0.0003885648038703948], [0.2897806465625763, 0.01695333980023861, 0.6714832782745361, 0.4471692144870758, 0.24303969740867615, 0.15563154220581055, 0.008645682595670223, 0.0004950988804921508, 0.0001695932005532086, 0.13566477596759796, 0.030448369681835175, 0.00021736785129178315, 9.297585347667336e-05, 0.0014399208594113588, 5.083655923954211e-05], [0.1102917492389679, 0.0027466323226690292, 0.13646264374256134, 0.07094646990299225, 0.17040857672691345, 0.6033481955528259, 0.41631338000297546, 0.013031017035245895, 0.00012492973473854363, 0.005976412910968065, 0.0002816450723912567, 4.682707003667019e-05, 0.00021861463028471917, 0.00019605428678914905, 0.001022772048600018], [0.7042187452316284, 0.49455204606056213, 0.43194010853767395, 0.7080989480018616, 0.382207989692688, 0.06800723820924759, 0.48792970180511475, 0.12651333212852478, 0.0012585417134687304, 0.07895761728286743, 0.01729964278638363, 0.0006471746601164341, 0.00013743228919338435, 0.00039039706462062895, 0.00010207234299741685], [0.5233215093612671, 0.7835124135017395, 0.3596530258655548, 0.5502080917358398, 0.589034378528595, 0.24138878285884857, 0.4714515507221222, 0.13250088691711426, 0.08884716778993607, 0.06473898142576218, 0.12478159368038177, 0.001717525301501155, 0.01358798798173666, 0.004862584639340639, 0.0004225081647746265], [0.0975094586610794, 0.14095744490623474, 0.009511731564998627, 0.03128954395651817, 0.01951521448791027, 0.0017430862644687295, 0.033708807080984116, 0.009512575343251228, 0.3042309582233429, 0.0025639990344643593, 0.0006334132049232721, 2.5987004846683703e-05, 0.0001574041525600478, 1.1997842193522956e-05, 1.5690195141360164e-05], [0.536220133304596, 0.12877297401428223, 0.013534938916563988, 0.13534405827522278, 0.015604051761329174, 0.0035537974908947945, 0.02344023622572422, 0.008398037403821945, 0.2580391466617584, 0.2587551474571228, 0.014949243515729904, 0.0010696486569941044, 0.00046315763029269874, 0.0013398011215031147, 8.422375685768202e-05], [0.028944578021764755, 0.013114584609866142, 0.0438210591673851, 0.05079193785786629, 0.03694206848740578, 0.0008442872785963118, 0.0030779552180320024, 0.002579997293651104, 0.01023491844534874, 0.21445545554161072, 0.2806929349899292, 0.00855539832264185, 0.03333647921681404, 0.06091907247900963, 1.9560096916393377e-05], [0.0058769844472408295, 0.06350620836019516, 0.003568005282431841, 0.0076079596765339375, 0.0037217612843960524, 0.004286385141313076, 0.03584115207195282, 0.14617407321929932, 0.0030082303564995527, 0.12143123894929886, 0.0793885663151741, 0.1555183082818985, 0.14442139863967896, 0.29275521636009216, 7.129996811272576e-05], [0.034930020570755005, 0.09419079124927521, 0.0127689428627491, 0.008763227611780167, 0.0065171802416443825, 0.008632887154817581, 0.02612082101404667, 0.02043459191918373, 0.0836663544178009, 0.5329904556274414, 0.3228733241558075, 0.7184357047080994, 0.5793755650520325, 0.783859133720398, 0.0001531920424895361], [0.0009532110998407006, 0.0024861039128154516, 7.189704774646088e-05, 0.00014637503772974014, 2.8552024105010787e-06, 3.0342853278853e-05, 0.0007709002820774913, 0.0005337693146429956, 6.919851330167148e-06, 0.02619163505733013, 0.02381032705307007, 0.008668542839586735, 0.39639002084732056, 0.7824769616127014, 1.1539431170604075e-06], [0.02785377763211727, 0.15845024585723877, 0.19323119521141052, 0.06543393433094025, 0.014044036157429218, 0.040286585688591, 0.07583826035261154, 0.6567350029945374, 0.004159754142165184, 0.35265031456947327, 0.6287637948989868, 0.12951745092868805, 0.32439297437667847, 0.653313934803009, 0.0008144593448378146], [0.02927210181951523, 0.04805546626448631, 0.295967698097229, 0.060625556856393814, 0.014990724623203278, 0.10397231578826904, 0.12186732143163681, 0.5237559080123901, 0.0203724168241024, 0.43874940276145935, 0.4409005343914032, 0.09095493704080582, 0.5531511306762695, 0.5263633728027344, 0.0002321143983863294], [0.5664732456207275, 0.02422192506492138, 0.3148367702960968, 0.37531769275665283, 0.06290365755558014, 0.02708868682384491, 0.03764869272708893, 0.06476183980703354, 0.09221415221691132, 0.3172641098499298, 0.088014617562294, 0.02202794700860977, 0.004314645659178495, 0.0619816817343235, 0.0017959593096747994], [0.04828598350286484, 0.01127469539642334, 0.1758044958114624, 0.0725238099694252, 0.01880812831223011, 0.003422890789806843, 0.0039800796657800674, 0.008112750947475433, 0.0007020575576461852, 0.0960424467921257, 0.3098883628845215, 0.03193678706884384, 0.03351299837231636, 0.2577627897262573, 0.0005041947006247938], [0.008833246305584908, 0.03231082111597061, 0.009648996405303478, 0.01135926228016615, 0.004257569555193186, 0.002696139505133033, 0.026390861719846725, 0.07894735038280487, 0.0002903220884036273, 0.05877671018242836, 0.0971919596195221, 0.32856324315071106, 0.08294347673654556, 0.6861463785171509, 0.00047716210247017443], [0.020260397344827652, 0.03928471356630325, 0.012783887796103954, 0.0091601787135005, 0.005565040744841099, 0.007968534715473652, 0.020862603560090065, 0.012279938906431198, 0.01832268387079239, 0.3204420506954193, 0.28696081042289734, 0.7937509417533875, 0.6314787864685059, 0.8277974724769592, 0.00014348741387948394], [0.00497927563264966, 0.011739314533770084, 0.0009416648535989225, 0.0009133343119174242, 2.0598932678694837e-05, 0.00024278588534798473, 0.00463896244764328, 0.0027787971775978804, 1.9694551156135276e-05, 0.026842234656214714, 0.05824153125286102, 0.023767979815602303, 0.7019069194793701, 0.8979114294052124, 1.5536308637820184e-05], [0.06832221150398254, 0.18812543153762817, 0.5426309108734131, 0.237625390291214, 0.041615329682826996, 0.11611851304769516, 0.16301436722278595, 0.827357828617096, 0.011619587428867817, 0.35340800881385803, 0.8248108625411987, 0.22083298861980438, 0.4978465139865875, 0.8379470109939575, 0.008811386302113533], [0.7676634788513184, 0.8615484237670898, 0.768317461013794, 0.9594964981079102, 0.36958935856819153, 0.4649639129638672, 0.5634418725967407, 0.8043064475059509, 0.6601962447166443, 0.9397303462028503, 0.8348119258880615, 0.9867405295372009, 0.7646960020065308, 0.8154686689376831, 0.03640103340148926]], [[0.5194346308708191, 0.08715501427650452, 0.09860441088676453, 0.08100719004869461, 0.11848669499158859, 0.14280925691127777, 0.19592297077178955, 0.1196337640285492, 0.2793996334075928, 0.0691760703921318, 0.09539081901311874, 0.05545644089579582, 0.02620256133377552, 0.03735822066664696, 0.09928011149168015], [0.002687783446162939, 0.2585922181606293, 0.004556892905384302, 0.0005560630816034973, 0.0013625096762552857, 0.000865808455273509, 2.095674426527694e-05, 0.013363445177674294, 1.4331720194604713e-05, 0.00023233501997310668, 0.013212678954005241, 0.00027388104354031384, 2.99917119264137e-05, 5.10126119479537e-05, 0.0653858631849289], [0.010489544831216335, 0.001751396106556058, 0.2775154411792755, 0.0030420231632888317, 0.08156438916921616, 0.0006471106316894293, 1.7804295566747896e-05, 0.00014657371502835304, 0.00035265504266135395, 0.00129506376106292, 0.018553601577878, 0.0019669390749186277, 0.009056665003299713, 0.05091148242354393, 0.1541917622089386], [0.0025869093369692564, 0.008571458049118519, 0.38431695103645325, 0.030530055984854698, 0.03365315869450569, 0.005854337941855192, 0.00010941662185359746, 4.1041937947738916e-05, 0.000364075880497694, 0.0011989381164312363, 0.014197473414242268, 0.0010815636487677693, 0.0004893331206403673, 0.0013785242335870862, 0.011478900909423828], [0.20589935779571533, 0.03613102436065674, 0.009011336602270603, 0.09399610757827759, 0.042497485876083374, 0.000576009857468307, 0.0040712482295930386, 0.00162220629863441, 0.00015305644774343818, 0.0034409475047141314, 0.025435233488678932, 2.175084773625713e-05, 1.0188268788624555e-05, 5.634217450278811e-05, 0.160919189453125], [0.00994176883250475, 0.015379102900624275, 0.000435269670560956, 0.004355194512754679, 0.002023787936195731, 4.86412636746536e-06, 0.0007220985717140138, 0.0004895065212622285, 0.0005591813242062926, 0.009127096273005009, 0.023014724254608154, 0.0003639610658865422, 3.1703839340480044e-05, 0.00036040451959706843, 0.1469942033290863], [0.31647789478302, 0.5689504742622375, 0.010991040617227554, 0.29046669602394104, 0.008814695291221142, 0.008600234054028988, 0.094898521900177, 0.02089405618607998, 0.005384301766753197, 0.1224634200334549, 0.2525540888309479, 0.011421876028180122, 9.89354812190868e-05, 0.00020726426737383008, 0.3419104218482971], [0.006757077760994434, 0.1354868859052658, 0.002759847091510892, 0.009205225855112076, 0.0038083188701421022, 0.0014255000278353691, 0.0007299972930923104, 0.2051592320203781, 0.00020230394147802144, 0.001623967313207686, 0.006681961473077536, 0.0021689198911190033, 5.557909025810659e-05, 0.000162289768923074, 0.20840437710285187], [0.010027364827692509, 0.02789497748017311, 0.0041139991953969, 0.012661347165703773, 0.0013435317669063807, 0.0034407242201268673, 0.0064836894161999226, 0.007366063538938761, 0.29601985216140747, 0.053567804396152496, 0.040060218423604965, 0.004607491660863161, 0.00018677859043236822, 3.186250978615135e-05, 0.10952453315258026], [0.19971387088298798, 0.012958711944520473, 0.001638519112020731, 0.17775660753250122, 0.0022716999519616365, 0.03685721755027771, 0.06948257982730865, 0.005452410783618689, 0.037147630006074905, 0.19678887724876404, 0.21911752223968506, 0.02466990426182747, 0.0004891769494861364, 6.33890085737221e-05, 0.21250228583812714], [0.05692211166024208, 0.036700569093227386, 0.0015533106634393334, 0.01848980039358139, 0.002404581755399704, 0.008354752324521542, 0.023693444207310677, 0.02836945652961731, 0.29948922991752625, 0.005321406293660402, 0.0022319734562188387, 0.0005214664852246642, 0.00019869217067025602, 5.8369230828247964e-05, 0.008838840760290623], [0.011123275384306908, 0.003955129534006119, 0.0015235289465636015, 0.011223106645047665, 0.002481319010257721, 0.000903434120118618, 0.0006720115779899061, 0.00024289102293550968, 0.010115177370607853, 0.26232361793518066, 0.014199022203683853, 0.0005582758458331227, 0.0001542939426144585, 5.357913687475957e-05, 0.050008371472358704], [0.025191567838191986, 0.009952094405889511, 0.015023785643279552, 0.0893990620970726, 0.006299919448792934, 0.0077370950020849705, 0.0004422276106197387, 0.00010742250742623582, 0.001807618304155767, 0.052116382867097855, 0.33116668462753296, 0.0029348258394747972, 0.004942082799971104, 0.0017646296182647347, 0.009777115657925606], [0.12133541703224182, 0.0033125760965049267, 0.008441481739282608, 0.0257105715572834, 0.005432062782347202, 0.020603680983185768, 0.0008238950395025313, 0.00019463927310425788, 0.0001117472565965727, 0.011082900688052177, 0.4118730425834656, 0.0024717452470213175, 0.21560189127922058, 0.015253315679728985, 0.03452993184328079], [0.00568122835829854, 0.003583817044273019, 0.0009402501164004207, 0.0034319525584578514, 0.014700439758598804, 0.00014027200813870877, 5.928567406954244e-05, 0.0005310353590175509, 0.001004774123430252, 0.00433507701382041, 0.003991644363850355, 0.0015378128737211227, 6.231402221601456e-05, 0.02625701017677784, 0.15481357276439667], [0.00503728911280632, 0.004739185329526663, 0.021364033222198486, 0.04603096470236778, 0.004565324168652296, 0.021244995296001434, 0.07592181116342545, 0.027910754084587097, 0.008603491820394993, 0.004941265098750591, 0.03103908710181713, 0.035909827798604965, 0.01818632334470749, 0.04406380280852318, 0.17931725084781647], [0.21416018903255463, 0.005411786492913961, 0.02111194096505642, 0.07001130282878876, 0.04736214876174927, 0.09187527745962143, 0.1399366855621338, 0.030981194227933884, 0.02342112548649311, 0.07424263656139374, 0.02716991677880287, 0.5710572600364685, 0.007255392149090767, 0.005560784600675106, 0.054831843823194504], [0.3339015245437622, 0.03176174685359001, 0.25991618633270264, 0.31748515367507935, 0.17923809587955475, 0.2977932095527649, 0.14185847342014313, 0.09826549887657166, 0.4168005883693695, 0.09961694478988647, 0.1390676498413086, 0.191667839884758, 0.0443519689142704, 0.10075851529836655, 0.08045557886362076], [0.018510108813643456, 0.0015040059806779027, 0.011199833825230598, 0.021222928538918495, 0.02421635016798973, 0.004175371024757624, 0.0007807075162418187, 0.0005349562270566821, 0.0038052168674767017, 0.3727143108844757, 0.022828511893749237, 0.01009275484830141, 0.0012628438416868448, 0.0009096930734813213, 0.10904579609632492], [0.05896773934364319, 0.023542853072285652, 0.0776505172252655, 0.15385140478610992, 0.011508575640618801, 0.0939982458949089, 0.0018089915392920375, 0.0003290986060164869, 0.0005636389250867069, 0.029514340683817863, 0.35146546363830566, 0.007090898230671883, 0.012099701911211014, 0.006742698606103659, 0.052738532423973083], [0.18205131590366364, 0.00472951028496027, 0.03192766383290291, 0.059333182871341705, 0.028221452608704567, 0.033883631229400635, 0.00131422549020499, 0.0001085989861167036, 5.632251122733578e-05, 0.004554648417979479, 0.2950275242328644, 0.0014449548907577991, 0.2329740822315216, 0.0520821250975132, 0.1361607313156128], [0.0063572716899216175, 0.002779513830319047, 0.0009721479145810008, 0.0035897656343877316, 0.019835324957966805, 0.00021187934908084571, 8.435463678324595e-05, 0.00043589723645709455, 0.0004945950931869447, 0.004414541646838188, 0.0027602717746049166, 0.0008482423145323992, 5.171148222871125e-05, 0.021799515932798386, 0.15211130678653717], [0.005286877974867821, 0.008391096256673336, 0.025823507457971573, 0.030178312212228775, 0.00857502967119217, 0.042816706001758575, 0.07608389109373093, 0.03679429367184639, 0.0067360359244048595, 0.0038807345554232597, 0.03710461035370827, 0.037315309047698975, 0.018847206607460976, 0.0415174663066864, 0.15352587401866913], [0.2992006242275238, 0.008802352473139763, 0.027079692110419273, 0.08564624935388565, 0.11560814827680588, 0.22971339523792267, 0.1826445311307907, 0.033842965960502625, 0.06175734102725983, 0.11205370724201202, 0.04016120731830597, 0.5851526856422424, 0.016921253874897957, 0.011652404442429543, 0.08951538056135178], [0.12446854263544083, 0.0009617851465009153, 0.004788657650351524, 0.0008746102685108781, 0.16037316620349884, 0.003065474098548293, 0.0056405095383524895, 0.005250739399343729, 0.05696318671107292, 0.013819074258208275, 0.028642717748880386, 0.0011808956041932106, 0.08446037769317627, 0.03008313849568367, 0.13710428774356842]], [[0.005261753685772419, 0.005328452680259943, 0.1075906753540039, 0.007504252251237631, 0.18196941912174225, 0.2677178680896759, 0.18533208966255188, 0.041308093816041946, 0.04052837938070297, 0.0018225060775876045, 0.004738607443869114, 0.028365809470415115, 0.07867489755153656, 0.032602421939373016, 0.14697469770908356], [0.024903474375605583, 0.2637169063091278, 0.01148936152458191, 0.01806865818798542, 0.010384032502770424, 0.05497525632381439, 0.01011874619871378, 6.159161421237513e-05, 0.03404803201556206, 0.01315199863165617, 0.004086918197572231, 0.033981483429670334, 0.0007253359071910381, 0.0010365481721237302, 0.023150891065597534], [0.03176039457321167, 0.002004105830565095, 0.011469452641904354, 0.003235333366319537, 0.011606591753661633, 0.01332010142505169, 0.007885226979851723, 0.0010319099528715014, 0.0026684575714170933, 0.003885145066305995, 0.002207087352871895, 0.010414022952318192, 0.015553043223917484, 0.01973811537027359, 0.1639232188463211], [0.24842531979084015, 0.031220050528645515, 0.028132880106568336, 0.029530569911003113, 0.01766534335911274, 0.36354437470436096, 0.06892471760511398, 0.02528339996933937, 0.01102821622043848, 0.15825842320919037, 0.13755246996879578, 0.07390110194683075, 0.19022952020168304, 0.1824880689382553, 0.1432848572731018], [0.0013664831640198827, 0.001714985934086144, 0.0013615208445116878, 0.0015855998499318957, 0.0011547008762136102, 0.007221538573503494, 0.01537399459630251, 0.020302001386880875, 0.0011185031617060304, 0.001242821803316474, 0.0004577837826218456, 0.0013307477347552776, 6.100967220845632e-05, 3.943840420106426e-05, 0.16435295343399048], [0.0006725311395712197, 0.000846685899887234, 0.001614874112419784, 0.000348375499015674, 0.0019150535808876157, 0.01370947528630495, 0.026421356946229935, 0.08118636161088943, 0.0008913385099731386, 0.0004401778569445014, 0.0003709472657646984, 0.0007744845934212208, 0.002328733913600445, 0.0003664834948722273, 0.14579549431800842], [0.011207095347344875, 0.029191432520747185, 0.015348215587437153, 0.012354064732789993, 0.002485303906723857, 0.7150441408157349, 0.0764552503824234, 0.14450958371162415, 0.0016117440536618233, 0.008765846490859985, 0.011787951923906803, 0.002862851833924651, 0.022502094507217407, 0.007210019044578075, 0.007054056040942669], [0.006926322355866432, 0.0050496323965489864, 0.010020078159868717, 0.021360181272029877, 0.0027102867607027292, 0.028520535677671432, 0.05918040871620178, 0.23060235381126404, 0.019199691712856293, 0.09477535635232925, 0.013206732459366322, 0.0014817069750279188, 0.0153219448402524, 0.01803957298398018, 0.07950127124786377], [0.009242678992450237, 0.05580667033791542, 0.014326682314276695, 0.04630666971206665, 0.010674487799406052, 0.5850453972816467, 0.4108324944972992, 0.4116209149360657, 0.007144990377128124, 0.20661039650440216, 0.037308260798454285, 0.054067905992269516, 0.037599414587020874, 0.03113422356545925, 0.22261686623096466], [0.0023711349349468946, 0.019731320440769196, 0.027566438540816307, 0.03758935630321503, 0.022646954283118248, 0.06538618355989456, 0.01152126956731081, 0.014797273091971874, 0.003413880243897438, 0.024214325472712517, 0.019466044381260872, 0.007235943805426359, 0.0008611958473920822, 0.0011126803001388907, 0.268255352973938], [0.08772679418325424, 0.02003292553126812, 0.09465871006250381, 0.41126132011413574, 0.07995565980672836, 0.5143890976905823, 0.1155472919344902, 0.01320470031350851, 0.02149844542145729, 0.06702866405248642, 0.6884661316871643, 0.09638151526451111, 0.35587188601493835, 0.2170087993144989, 0.019593046978116035], [0.01343127153813839, 0.0019279895350337029, 0.01925632171332836, 0.04226915165781975, 0.005290344823151827, 0.5555825233459473, 0.06846548616886139, 0.006453313864767551, 0.019162334501743317, 0.0017575293313711882, 0.2967261075973511, 0.11721283942461014, 0.4438721835613251, 0.1899448037147522, 0.007863422855734825], [0.12789316475391388, 0.004323228262364864, 0.03538274019956589, 0.05581461265683174, 0.020947236567735672, 0.09860846400260925, 0.11394336074590683, 0.010361305437982082, 0.011101406998932362, 0.33580121397972107, 0.13689599931240082, 0.038663506507873535, 0.19725953042507172, 0.10533706098794937, 0.008538279682397842], [0.007053391542285681, 0.012331487610936165, 0.008611395955085754, 0.031008008867502213, 0.004283395130187273, 0.0029549654573202133, 0.00849387887865305, 0.008564120158553123, 0.02629040740430355, 0.009985123760998249, 0.00761940935626626, 0.003499145619571209, 0.0015691317385062575, 0.005600257311016321, 0.5214234590530396], [0.0007030746201053262, 0.0001308645587414503, 0.0001913319865707308, 0.00016671058256179094, 0.000299752748105675, 0.0001608166057849303, 0.004501530434936285, 0.0010771069210022688, 0.003937124740332365, 0.001599485520273447, 0.0007339937728829682, 0.0030779645312577486, 3.4502605558373034e-05, 9.700484952190891e-05, 0.15641583502292633], [0.027913473546504974, 0.10055015236139297, 0.005828284192830324, 0.007361504249274731, 0.0010143647668883204, 0.000654859293717891, 0.0101061025634408, 0.029607031494379044, 0.04485415667295456, 0.09235014766454697, 0.05163425952196121, 0.03075464628636837, 0.027050884440541267, 0.021472401916980743, 0.18064866960048676], [0.0011193754617124796, 0.03864011913537979, 0.0033454783260822296, 0.0006957795703783631, 0.001480268081650138, 0.0012079592561349273, 0.00020605533791240305, 0.0011212154058739543, 0.0015670693246647716, 0.0014121911954134703, 0.0012700740480795503, 0.0019415348069742322, 0.001359732006676495, 0.0011440571397542953, 0.23876120150089264], [0.012943120673298836, 0.020876264199614525, 0.04825761169195175, 0.03707631304860115, 0.015636419877409935, 0.11923719942569733, 0.021652603521943092, 0.026653259992599487, 0.020431919023394585, 0.03287035599350929, 0.10921605676412582, 0.11103712767362595, 0.08490956574678421, 0.05352960154414177, 0.1791488379240036], [0.010143280029296875, 0.0011783033842220902, 0.07699523866176605, 0.04151652753353119, 0.013031265698373318, 0.6595657467842102, 0.04001229628920555, 0.015414847061038017, 0.05828738585114479, 0.00582495890557766, 0.39538952708244324, 0.3540988564491272, 0.5535411834716797, 0.14920510351657867, 0.05510678142309189], [0.10365689545869827, 0.011393263004720211, 0.09083462506532669, 0.05552159622311592, 0.021694108843803406, 0.23093751072883606, 0.12655670940876007, 0.02638416364789009, 0.016898566856980324, 0.4334920644760132, 0.1302367001771927, 0.07987051457166672, 0.26015403866767883, 0.07882147282361984, 0.06412448734045029], [0.0009046280756592751, 0.006186267826706171, 0.001710598124191165, 0.0040000369772315025, 0.0010556421475484967, 0.00010012275743065402, 0.000467440317152068, 0.00034073027200065553, 0.012450831942260265, 0.001776019111275673, 0.0016348852077499032, 0.0004490323772188276, 0.00023723821504972875, 0.0005369102582335472, 0.2610536217689514], [0.00040706052095629275, 5.995776882627979e-05, 0.00011266738147241995, 0.00010974665929097682, 0.00022393744438886642, 7.468188414350152e-05, 0.00239625689573586, 0.0004222780407872051, 0.002755024004727602, 0.0011263962369412184, 0.0004159261588938534, 0.0013214137870818377, 1.3015362128498964e-05, 3.146446033497341e-05, 0.15343648195266724], [0.02487853355705738, 0.06922142952680588, 0.005931189749389887, 0.005149703938513994, 0.0007503133383579552, 0.00046759017277508974, 0.004864065907895565, 0.010271446779370308, 0.03885169327259064, 0.0494176521897316, 0.032662954181432724, 0.015474021434783936, 0.005468437913805246, 0.0031831569503992796, 0.16160887479782104], [0.0006016235565766692, 0.010655699297785759, 0.0012552555417641997, 0.0004406629304867238, 0.0006771506741642952, 0.0004804672207683325, 8.584682655055076e-05, 0.00018533790716901422, 0.0020008538849651814, 0.0008522755815647542, 0.0005471827462315559, 0.0006654397584497929, 0.0003326669684611261, 0.00020969027536921203, 0.18202657997608185], [0.0006660889484919608, 0.0011989487102255225, 0.006168409250676632, 0.0007392434636130929, 0.002072105184197426, 0.0013732375809922814, 0.001215140800923109, 8.942947169998661e-05, 0.0032219376880675554, 0.00034276655060239136, 0.0006051870877854526, 0.0004003554640803486, 0.0006330502219498158, 9.228585986420512e-05, 0.13989190757274628]], [[0.17597882449626923, 0.03865775838494301, 0.04927876219153404, 0.19269852340221405, 0.07631995528936386, 0.03202155977487564, 0.04315444082021713, 0.0381813645362854, 0.14437337219715118, 0.14268529415130615, 0.12548406422138214, 0.22065725922584534, 0.007455701474100351, 0.012540786527097225, 0.13194040954113007], [0.12168548256158829, 0.12690430879592896, 0.03319493681192398, 0.044549524784088135, 0.022643521428108215, 0.12293753027915955, 0.012858373112976551, 0.056580886244773865, 0.0409478023648262, 0.5390252470970154, 0.04499629884958267, 0.010665545240044594, 0.0012580851325765252, 0.0006077282596379519, 0.16003872454166412], [0.004976227879524231, 0.0016218257369473577, 0.10218203067779541, 0.005807417444884777, 0.025330372154712677, 0.00805770605802536, 0.0010953968157991767, 0.007808555383235216, 0.03332183510065079, 0.01014297641813755, 0.0378553569316864, 0.0012688467977568507, 0.0070253219455480576, 0.006525768432766199, 0.1611432433128357], [0.018298039212822914, 0.043392445892095566, 0.026758581399917603, 0.06685060262680054, 0.007846164517104626, 0.0070086256600916386, 0.0011090404586866498, 0.0016357558779418468, 0.015295942313969135, 0.022091375663876534, 0.08676162362098694, 0.0013220091350376606, 0.0007799563463777304, 0.0005145008908584714, 0.5814905166625977], [0.16791731119155884, 0.01838838867843151, 0.03170344606041908, 0.04746389389038086, 0.024931352585554123, 0.002624210435897112, 0.3320338726043701, 0.32248422503471375, 0.021048149093985558, 0.02857070416212082, 0.11922428011894226, 4.079664358869195e-05, 0.0002566495386417955, 0.0005197013379074633, 0.1538068950176239], [0.03376027196645737, 0.001082546659745276, 0.003266592975705862, 0.006257645785808563, 0.023632841184735298, 0.00021245618700049818, 0.033721838146448135, 0.15340450406074524, 0.009442711248993874, 0.006162047851830721, 0.09923229366540909, 0.0001386175281368196, 0.0008165750186890364, 0.0010916005121544003, 0.14602994918823242], [0.04221357777714729, 0.03857824206352234, 0.004161412362009287, 0.06419923156499863, 0.010648604482412338, 0.008165394887328148, 0.04070910066366196, 0.34736329317092896, 0.0012154168216511607, 0.1630050241947174, 0.07001504302024841, 0.0033116117119789124, 0.00023883172252681106, 0.00045473958016373217, 0.2740376889705658], [0.007271567825227976, 0.0015110730892047286, 0.0014769553672522306, 0.0053740208968520164, 0.0038654205854982138, 0.0024983601178973913, 0.049697574228048325, 0.27208074927330017, 0.0006182760698720813, 0.014045008458197117, 0.00131281279027462, 0.00040628391434438527, 0.00037906834040768445, 0.0001199298130813986, 0.006693295668810606], [0.08829134702682495, 0.11286511272192001, 0.004967967513948679, 0.006996258161962032, 0.0014454894699156284, 0.006397548597306013, 0.01389994379132986, 0.27431485056877136, 0.0018983082845807076, 0.09154568612575531, 0.022492842748761177, 0.0017391144065186381, 0.000634143827483058, 4.5783879613736644e-05, 0.318096399307251], [0.02142007276415825, 0.007001234218478203, 0.00761477230116725, 0.018849696964025497, 0.010492328554391861, 0.01844215951859951, 0.008208145387470722, 0.01109394058585167, 0.006335548125207424, 0.01884968765079975, 0.01652243174612522, 0.016355833038687706, 0.0014795949682593346, 0.0011322565842419863, 0.27169719338417053], [0.17013461887836456, 0.14343884587287903, 0.017679741606116295, 0.10850679129362106, 0.01231957133859396, 0.010847942903637886, 0.04900640249252319, 0.023357992991805077, 0.014735743403434753, 0.014097570441663265, 0.012582896277308464, 0.0010529988212510943, 0.00046457236749120057, 0.0006211225991137326, 0.5663455724716187], [0.1586649864912033, 0.08337923884391785, 0.0181503314524889, 0.22676831483840942, 0.016727542504668236, 0.015186772681772709, 0.0050455182790756226, 0.00688449339941144, 0.025511443614959717, 0.20239992439746857, 0.024231791496276855, 0.0023393011651933193, 0.0011192933889105916, 0.0005647524958476424, 0.390881210565567], [0.3443087935447693, 0.28029316663742065, 0.23536846041679382, 0.34415915608406067, 0.11761639267206192, 0.006012732163071632, 0.008058828301727772, 0.005314267706125975, 0.013309409841895103, 0.09906232357025146, 0.10091385245323181, 0.018941059708595276, 0.025248508900403976, 0.014945760369300842, 0.7436007857322693], [0.0022638223599642515, 0.004991845227777958, 0.004655482713133097, 0.0007185174035839736, 0.0013901105849072337, 0.011776956729590893, 0.0005479936371557415, 0.00022604972764384001, 0.00024645475787110627, 0.009541304782032967, 0.011744895949959755, 0.0007132806931622326, 0.27867355942726135, 0.02834550105035305, 0.007979176938533783], [0.024570701643824577, 0.00167787482496351, 0.004072254989296198, 0.00223688711412251, 0.007143567781895399, 0.00014352552534546703, 0.0004634522774722427, 0.0016921478090807796, 0.003620122792199254, 0.007754941936582327, 0.011850811541080475, 0.0027722271624952555, 9.3724018370267e-05, 0.02145184949040413, 0.15506701171398163], [0.01723022572696209, 0.08018677681684494, 0.007713299244642258, 0.004271229729056358, 0.0005464836140163243, 0.00456921337172389, 0.0031762931030243635, 0.009469777345657349, 0.000385247083613649, 0.01870143786072731, 0.033109456300735474, 0.004042719956487417, 0.004976211115717888, 0.005646048113703728, 0.19230251014232635], [0.016216034069657326, 0.04777013510465622, 0.01620146818459034, 0.010810854844748974, 0.16034351289272308, 0.006931359879672527, 0.0032006967812776566, 0.032106515020132065, 0.0003033989341929555, 0.015325331129133701, 0.006036583799868822, 0.12791146337985992, 0.19952742755413055, 0.023708127439022064, 0.18307197093963623], [0.014499284327030182, 0.035677529871463776, 0.009275808930397034, 0.01653297245502472, 0.006223962642252445, 0.0020693510305136442, 0.007680083625018597, 0.013822571374475956, 0.00040966575033962727, 0.0038025544490665197, 0.013774569146335125, 0.006069935858249664, 0.004488381557166576, 0.005977130029350519, 0.217429518699646], [0.03237156197428703, 0.013441890478134155, 0.0194883793592453, 0.09343220293521881, 0.05379915237426758, 0.004893247038125992, 0.0011929833563044667, 0.009432576596736908, 0.015330814756453037, 0.14898745715618134, 0.018398255109786987, 0.01228779274970293, 0.00492482166737318, 0.0038985873106867075, 0.2601524889469147], [0.08357361704111099, 0.18220724165439606, 0.10462122410535812, 0.08245989680290222, 0.03124452568590641, 0.002170282183215022, 0.0020384257659316063, 0.004550496581941843, 0.003485089400783181, 0.036062099039554596, 0.0278666652739048, 0.011443988420069218, 0.01760544627904892, 0.013599698431789875, 0.3874043822288513], [0.001995340920984745, 0.011527596041560173, 0.005334027577191591, 0.0006887424970045686, 0.0023407095577567816, 0.00276917009614408, 0.00029977987287566066, 0.00012230046559125185, 0.00026578022516332567, 0.008239910937845707, 0.009819538332521915, 0.000393931899452582, 0.605858564376831, 0.08989311754703522, 0.011135715991258621], [0.021298440173268318, 0.001658836961723864, 0.004600299056619406, 0.0025729055050760508, 0.015332063660025597, 0.00017298871534876525, 0.0005721640191040933, 0.00186175387352705, 0.0037871075328439474, 0.009124312549829483, 0.01116581168025732, 0.0031747270841151476, 0.00012207991676405072, 0.029056062921881676, 0.15163807570934296], [0.020229021087288857, 0.11621151119470596, 0.015550180338323116, 0.006284819450229406, 0.0013723199954256415, 0.013658476993441582, 0.005685316864401102, 0.02063058130443096, 0.001440295367501676, 0.022225895896553993, 0.07092871516942978, 0.007373427972197533, 0.00771017000079155, 0.006927240639925003, 0.16024509072303772], [0.014029471203684807, 0.02389930933713913, 0.011611595749855042, 0.012217668816447258, 0.2477317750453949, 0.006976675242185593, 0.0035841658245772123, 0.022232146933674812, 0.0018886715406551957, 0.01750483363866806, 0.005654812324792147, 0.10889071226119995, 0.19916927814483643, 0.022882532328367233, 0.16074435412883759], [0.0032621105201542377, 0.006088452413678169, 0.012619324028491974, 0.008848619647324085, 0.17461968958377838, 8.660123421577737e-05, 0.0006109846872277558, 0.0007747155614197254, 0.003163054818287492, 0.017787659540772438, 0.029563669115304947, 0.0032195982057601213, 0.013336165808141232, 0.013171130791306496, 0.1387031376361847]], [[0.09661699831485748, 0.7619754076004028, 0.05676787346601486, 0.020180072635412216, 0.10883769392967224, 0.42711278796195984, 0.09064477682113647, 0.10612691193819046, 0.04782179743051529, 0.06935178488492966, 0.027948519214987755, 0.00755169615149498, 0.007339869160205126, 0.025803416967391968, 0.09292053431272507], [0.042798254638910294, 0.23223945498466492, 0.062359996140003204, 0.01933804154396057, 0.04838808253407478, 0.30189236998558044, 0.0354127362370491, 0.019764740020036697, 0.00920741818845272, 0.0097093116492033, 0.0160877276211977, 0.0032758424058556557, 0.005296806804835796, 0.011010169051587582, 0.02110680378973484], [0.02002989500761032, 0.001048662350513041, 0.03834937512874603, 0.030392715707421303, 0.09750902652740479, 0.056120067834854126, 0.008173296228051186, 0.006944228895008564, 0.004440560005605221, 0.005061029922217131, 0.007118762470781803, 0.008411978371441364, 0.023608768358826637, 0.04182775691151619, 0.16016238927841187], [0.041295986622571945, 0.19780276715755463, 0.03777160495519638, 0.1712082475423813, 0.20935285091400146, 0.158755823969841, 0.3937656581401825, 0.684601902961731, 0.2584594190120697, 0.11237194389104843, 0.1112959012389183, 0.09882687777280807, 0.05429066717624664, 0.24210131168365479, 0.016339490190148354], [0.26312491297721863, 0.2720799446105957, 0.005703570321202278, 0.0481516495347023, 0.027902500703930855, 0.0034437666181474924, 0.03425572067499161, 0.03555849939584732, 0.028000997379422188, 0.0429554246366024, 0.002753790933638811, 0.0017769382102414966, 0.002218457870185375, 0.003535473719239235, 0.1597488671541214], [0.22248251736164093, 0.03185709938406944, 0.000688861298840493, 0.005810217931866646, 0.007679672911763191, 0.0008787074475549161, 0.07858764380216599, 0.14273476600646973, 0.07306984066963196, 0.02433006465435028, 0.011720307171344757, 0.013396549038589, 0.017704129219055176, 0.034836068749427795, 0.1453055441379547], [0.1531120240688324, 0.15391655266284943, 0.006810865830630064, 0.07720811665058136, 0.008951452560722828, 0.01149735413491726, 0.2822602391242981, 0.30408379435539246, 0.48283058404922485, 0.33028021454811096, 0.16095426678657532, 0.031167738139629364, 0.03355513513088226, 0.13962571322917938, 0.012790725566446781], [0.03593587130308151, 0.03233448788523674, 0.22662676870822906, 0.405829519033432, 0.014032814651727676, 0.02822977490723133, 0.09231841564178467, 0.1225365549325943, 0.20093639194965363, 0.2508411109447479, 0.5826555490493774, 0.037383783608675, 0.07952429354190826, 0.10720134526491165, 0.15212680399417877], [0.037364520132541656, 0.04119153320789337, 0.0012645104434341192, 0.021537767723202705, 0.000536995125003159, 0.0011436643544584513, 0.019049961119890213, 0.06139632686972618, 0.385105162858963, 0.13276730477809906, 0.24771228432655334, 0.04952799528837204, 0.04911990836262703, 0.11973114311695099, 0.021608887240290642], [0.004867227748036385, 0.009626063518226147, 0.0003137234307359904, 0.0026314754504710436, 0.00027048110496252775, 0.000934475683607161, 0.007251756265759468, 0.03575620427727699, 0.40781450271606445, 0.05584407597780228, 0.040446195751428604, 0.005334825720638037, 0.007708138320595026, 0.06401336193084717, 0.010240204632282257], [0.19358457624912262, 0.2328234314918518, 0.0017398587660863996, 0.10100623220205307, 0.0019695234950631857, 0.1674531251192093, 0.4513051509857178, 0.6547151803970337, 0.030009860172867775, 0.7025956511497498, 0.1685936599969864, 0.03178222477436066, 0.13270388543605804, 0.23426049947738647, 0.010277668945491314], [0.09463346004486084, 0.5257620811462402, 0.0045187450014054775, 0.07222570478916168, 0.0025188177824020386, 0.1410406231880188, 0.06597349792718887, 0.0719805508852005, 0.09957849979400635, 0.17567123472690582, 0.18618373572826385, 0.02195402979850769, 0.042485080659389496, 0.12470933794975281, 0.00617468124255538], [0.027796348556876183, 0.06599752604961395, 0.002643989399075508, 0.029425768181681633, 0.008861851878464222, 0.013279970735311508, 0.25377023220062256, 0.2656356692314148, 0.055540941655635834, 0.027583830058574677, 0.004816746339201927, 0.3890189528465271, 0.12020140886306763, 0.33882811665534973, 0.0040408894419670105], [0.4147956669330597, 0.5514373779296875, 0.09636387228965759, 0.29775112867355347, 0.03436855599284172, 0.08799602836370468, 0.07023341208696365, 0.10276275128126144, 0.25543972849845886, 0.10302554070949554, 0.05857125297188759, 0.029829595237970352, 0.114840567111969, 0.33078575134277344, 0.07371985912322998], [0.07031518220901489, 0.001305539975874126, 0.0025430582463741302, 0.010662226937711239, 0.0007357596186921, 0.000663888524286449, 0.0014398572966456413, 0.0005107407923787832, 0.005960140842944384, 0.0030986208003014326, 0.0017578504048287868, 0.00018377922242507339, 1.743367283779662e-05, 4.847845411859453e-05, 0.15638960897922516], [0.24421003460884094, 0.03331591188907623, 0.07573812454938889, 0.33240795135498047, 0.006838400848209858, 0.008697851561009884, 0.06428743898868561, 0.06466686725616455, 0.006176145281642675, 0.06394235789775848, 0.09260299056768417, 0.19959890842437744, 0.02154124155640602, 0.021672323346138, 0.15025706589221954], [0.5462155342102051, 0.545982301235199, 0.3341628611087799, 0.5788259506225586, 0.08809857815504074, 0.06356553733348846, 0.022417092695832253, 0.0164126455783844, 0.00386660173535347, 0.10154324769973755, 0.14015790820121765, 0.0864240974187851, 0.34186482429504395, 0.22899740934371948, 0.05407746881246567], [0.48888036608695984, 0.6578190326690674, 0.030819885432720184, 0.2205304652452469, 0.004883326590061188, 0.0656682699918747, 0.04461565986275673, 0.05094402655959129, 0.0005314986919984221, 0.15455113351345062, 0.10763049870729446, 0.1186080202460289, 0.14419804513454437, 0.1328149437904358, 0.09490374475717545], [0.15812784433364868, 0.9118645191192627, 0.022590545937418938, 0.05952226370573044, 0.00360964541323483, 0.07875056564807892, 0.013187792152166367, 0.02020449750125408, 0.0020393244922161102, 0.033818699419498444, 0.0449705570936203, 0.02132066898047924, 0.0717315599322319, 0.12101268768310547, 0.06353376060724258], [0.07771441340446472, 0.4748976230621338, 0.012594498693943024, 0.043653786182403564, 0.006564431358128786, 0.024485116824507713, 0.20463299751281738, 0.1550481915473938, 0.0016144687542691827, 0.005543926265090704, 0.0017496985383331776, 0.3491710126399994, 0.23835937678813934, 0.3316482901573181, 0.08539295196533203], [0.22228576242923737, 0.3581831455230713, 0.10504736006259918, 0.2062736451625824, 0.015430409461259842, 0.007369442842900753, 0.009848481975495815, 0.0027359407395124435, 0.003257193835452199, 0.004766176920384169, 0.0058546122163534164, 0.0040231142193078995, 0.032162997871637344, 0.05548902228474617, 0.22239458560943604], [0.040305208414793015, 0.0008039010572247207, 0.001399470493197441, 0.006614126265048981, 0.0003286598657723516, 0.0002559607964940369, 0.0005696980515494943, 0.00010972175368806347, 0.0006102611077949405, 0.0009710662416182458, 0.0004746906051877886, 5.0628168537514284e-05, 6.201828455232317e-06, 1.1841932064271532e-05, 0.15342259407043457], [0.18667390942573547, 0.05485990643501282, 0.06146723031997681, 0.2094709873199463, 0.003188095986843109, 0.005957009736448526, 0.04363764822483063, 0.02604665607213974, 0.0011390803847461939, 0.022857926785945892, 0.035827361047267914, 0.07732249796390533, 0.00673074834048748, 0.004807854071259499, 0.15350142121315002], [0.46625471115112305, 0.6644052863121033, 0.19963930547237396, 0.36004284024238586, 0.06144074350595474, 0.06362717598676682, 0.016601700335741043, 0.006137203890830278, 0.0020489897578954697, 0.041981395334005356, 0.042364589869976044, 0.04546959325671196, 0.25786423683166504, 0.1048446074128151, 0.10812478512525558], [0.01868601329624653, 0.08739857375621796, 0.016145089641213417, 0.000850466953124851, 0.0035631621722131968, 0.013478883542120457, 0.0006747889565303922, 0.0010685214074328542, 0.013735192827880383, 0.0029910006560385227, 0.017663421109318733, 0.0005569100612774491, 0.0335303470492363, 0.010939561761915684, 0.13854636251926422]], [[0.03039383515715599, 0.011264979839324951, 0.30973049998283386, 0.33407092094421387, 0.24303670227527618, 0.013086382299661636, 0.12547586858272552, 0.047571711242198944, 0.07738520950078964, 0.2579103410243988, 0.13098950684070587, 0.3019145727157593, 0.018321001902222633, 0.10478901118040085, 0.1313871294260025], [0.32489657402038574, 0.01967906951904297, 0.10292623937129974, 0.18745845556259155, 0.06220339238643646, 0.03126899152994156, 0.030121171846985817, 0.013807957991957664, 0.01960192248225212, 0.10352540761232376, 0.08122410625219345, 0.11610747873783112, 0.05098450556397438, 0.06022121384739876, 0.24838198721408844], [0.21547414362430573, 0.011987588368356228, 0.09540344774723053, 0.03949207067489624, 0.22973625361919403, 0.013393656350672245, 0.014646085910499096, 0.018391601741313934, 0.12483032047748566, 0.04761500656604767, 0.16838808357715607, 0.0500614158809185, 0.09093409031629562, 0.09172232449054718, 0.14920873939990997], [0.3455514907836914, 0.20528344810009003, 0.14200778305530548, 0.1397678107023239, 0.3345029056072235, 0.04282815381884575, 0.020769812166690826, 0.02952164225280285, 0.29125186800956726, 0.09975660592317581, 0.3298649489879608, 0.36294782161712646, 0.10288939625024796, 0.1784013956785202, 0.03550736606121063], [0.023072484880685806, 0.08888474851846695, 0.04328835755586624, 0.009794876910746098, 0.18984860181808472, 0.0009663040982559323, 0.0038235578685998917, 0.05101485177874565, 0.059323158115148544, 0.00876270979642868, 0.021391507238149643, 0.02426949329674244, 0.013026251457631588, 0.06840420514345169, 0.15691325068473816], [0.20066522061824799, 0.18445545434951782, 0.10427504032850266, 0.02148139849305153, 0.3108636438846588, 0.0010669901967048645, 0.031332992017269135, 0.06621930748224258, 0.42585986852645874, 0.05703788995742798, 0.1919325739145279, 0.6617251038551331, 0.07196007668972015, 0.2038833349943161, 0.13549473881721497], [0.06934618204832077, 0.15043997764587402, 0.24868465960025787, 0.0180400051176548, 0.61164391040802, 0.0047634197399020195, 0.0077652581967413425, 0.01316747348755598, 0.09036756306886673, 0.016214115545153618, 0.09484434872865677, 0.7773507833480835, 0.3649398386478424, 0.19880527257919312, 0.026039909571409225], [0.5420496463775635, 0.775536835193634, 0.21455605328083038, 0.17522192001342773, 0.3905614912509918, 0.07102629542350769, 0.15213513374328613, 0.06534071266651154, 0.05938922241330147, 0.3742612600326538, 0.040289394557476044, 0.6919643878936768, 0.07523911446332932, 0.14220400154590607, 0.06588775664567947], [0.05002814158797264, 0.18039211630821228, 0.4788157641887665, 0.0970841720700264, 0.5287489891052246, 0.07699278742074966, 0.024560611695051193, 0.055294524878263474, 0.031155720353126526, 0.029308732599020004, 0.023515479639172554, 0.10280930250883102, 0.01905171573162079, 0.033789344131946564, 0.006217750255018473], [0.2326076328754425, 0.12470381706953049, 0.5816100239753723, 0.187625452876091, 0.17989297211170197, 0.58512943983078, 0.4148763120174408, 0.7688660621643066, 0.02497384324669838, 0.10204316675662994, 0.16508084535598755, 0.4722842574119568, 0.654721736907959, 0.31103214621543884, 0.02808636985719204], [0.32085803151130676, 0.3732209801673889, 0.8471049070358276, 0.2474840134382248, 0.8311324715614319, 0.1531035155057907, 0.14141014218330383, 0.12460694462060928, 0.15561653673648834, 0.05888388305902481, 0.03703024983406067, 0.2600737512111664, 0.049645353108644485, 0.08333000540733337, 0.053744472563266754], [0.048572178930044174, 0.20163586735725403, 0.8568418025970459, 0.3438677489757538, 0.8764770030975342, 0.038519736379384995, 0.10765119642019272, 0.14438603818416595, 0.13915397226810455, 0.04139794409275055, 0.24816225469112396, 0.22188685834407806, 0.1582770049571991, 0.255889892578125, 0.05260627716779709], [0.10717450082302094, 0.14654512703418732, 0.5492125749588013, 0.149112731218338, 0.6473506689071655, 0.014123019762337208, 0.023513145744800568, 0.06304500997066498, 0.5243880152702332, 0.17494699358940125, 0.11734810471534729, 0.2534768283367157, 0.06080847606062889, 0.1781260073184967, 0.01657547615468502], [0.024022793397307396, 0.20128284394741058, 0.39493197202682495, 0.16542883217334747, 0.7724959254264832, 0.05353498458862305, 0.039175428450107574, 0.21511156857013702, 0.10924636572599411, 0.3127569556236267, 0.20907098054885864, 0.6610769033432007, 0.026550091803073883, 0.07443477213382721, 0.04747246578335762], [0.0639173686504364, 0.0019661476835608482, 0.03054100275039673, 0.07290788739919662, 0.07458660751581192, 0.0017515828367322683, 0.01338117104023695, 0.0049591753631830215, 0.10895326733589172, 0.03256915882229805, 0.07470867037773132, 0.022291045635938644, 0.00026081688702106476, 0.003768018214032054, 0.15579301118850708], [0.00809751357883215, 0.08670660853385925, 0.12165205925703049, 0.06173386052250862, 0.8110419511795044, 0.006245153024792671, 0.03447260707616806, 0.08050490915775299, 0.779870867729187, 0.2479465901851654, 0.38426774740219116, 0.6870184540748596, 0.2310730367898941, 0.07155610620975494, 0.05814361199736595], [0.01971210353076458, 0.10859540849924088, 0.17558348178863525, 0.04931360110640526, 0.4077165424823761, 0.001824796199798584, 0.004386546555906534, 0.0422598272562027, 0.9374924302101135, 0.3226373493671417, 0.06322266161441803, 0.05341457948088646, 0.0039883931167423725, 0.004304073750972748, 0.13460686802864075], [0.018049566075205803, 0.12295468151569366, 0.24470828473567963, 0.04122815281152725, 0.7332677245140076, 0.004472800530493259, 0.0029204280581325293, 0.018685931339859962, 0.4878760874271393, 0.20441682636737823, 0.08441592752933502, 0.4205068051815033, 0.04466289281845093, 0.13263334333896637, 0.0994158536195755], [0.007120466325432062, 0.02300306409597397, 0.2714575231075287, 0.07745856046676636, 0.6446666717529297, 0.0059507740661501884, 0.011145476251840591, 0.13244189321994781, 0.38060593605041504, 0.06726288050413132, 0.22673718631267548, 0.3522229492664337, 0.17927831411361694, 0.524927020072937, 0.09379637986421585], [0.03649899363517761, 0.08160936087369919, 0.2519805133342743, 0.07504414021968842, 0.1795702874660492, 0.006024391856044531, 0.0073743402026593685, 0.061968039721250534, 0.7520835995674133, 0.28517279028892517, 0.1493321657180786, 0.3589819371700287, 0.04636238142848015, 0.16408585011959076, 0.046330999583005905], [0.009416425600647926, 0.1558573991060257, 0.15325002372264862, 0.08311447501182556, 0.6221630573272705, 0.0029961667023599148, 0.006436231546103954, 0.027678541839122772, 0.2543543577194214, 0.47390833497047424, 0.28851544857025146, 0.6220062375068665, 0.014266690239310265, 0.05054754391312599, 0.0578170008957386], [0.04693470522761345, 0.0011674511479213834, 0.01364858541637659, 0.06039872020483017, 0.0427468940615654, 0.0009404723532497883, 0.007858873344957829, 0.0028007859364151955, 0.06382106244564056, 0.03982963413000107, 0.05175205320119858, 0.011254650540649891, 0.0001272865483770147, 0.001588277518749237, 0.15313954651355743], [0.017768997699022293, 0.1465732455253601, 0.15898801386356354, 0.12304693460464478, 0.8442554473876953, 0.006285809446126223, 0.04204265773296356, 0.12739135324954987, 0.8276333808898926, 0.5079721808433533, 0.5299316644668579, 0.8274551630020142, 0.09790517389774323, 0.02651425078511238, 0.11435628682374954], [0.017107579857110977, 0.05770094692707062, 0.07052541524171829, 0.059498131275177, 0.2613165080547333, 0.0009367912425659597, 0.0028308003675192595, 0.01869240775704384, 0.8671534061431885, 0.40041688084602356, 0.03947103023529053, 0.0349445715546608, 0.00177917187102139, 0.002164072822779417, 0.1562660187482834], [0.006599111016839743, 0.004138579126447439, 0.06047067046165466, 0.013185898773372173, 0.15347044169902802, 0.000755132467020303, 0.007522573694586754, 0.002741254400461912, 0.10833818465471268, 0.005474736914038658, 0.009540018625557423, 0.00040286476723849773, 0.004092549905180931, 0.002003892557695508, 0.13896189630031586]]], [[[0.010830877348780632, 0.011870973743498325, 0.10922139137983322, 0.013140714727342129, 0.060979437083005905, 0.24213501811027527, 0.056873127818107605, 0.0565403513610363, 0.1606917381286621, 0.004471848253160715, 0.04391508549451828, 0.16444265842437744, 0.14521700143814087, 0.12183647602796555, 0.18165212869644165], [0.1442122757434845, 0.026047294959425926, 0.4262431859970093, 0.3211715519428253, 0.7946609258651733, 0.48857852816581726, 0.31943926215171814, 0.3322535455226898, 0.8442224860191345, 0.37700119614601135, 0.4491288661956787, 0.725179135799408, 0.5425247550010681, 0.7077597379684448, 0.47353750467300415], [0.004308484960347414, 0.0038143862038850784, 0.01376394834369421, 0.007213444449007511, 0.0352218858897686, 0.009065943770110607, 0.00796457938849926, 0.009648038074374199, 0.012818497605621815, 0.005304576829075813, 0.00578665267676115, 0.025514552369713783, 0.003588201943784952, 0.005116589833050966, 0.1385156214237213], [0.37350767850875854, 0.33144617080688477, 0.1264321357011795, 0.21400198340415955, 0.32627996802330017, 0.09132378548383713, 0.05067773535847664, 0.05911920592188835, 0.47554144263267517, 0.5285797715187073, 0.055136121809482574, 0.07909779250621796, 0.0048016151413321495, 0.023815851658582687, 0.05086187273263931], [0.026979738846421242, 0.17144815623760223, 0.016802728176116943, 0.011190843768417835, 0.05719228833913803, 0.006600439548492432, 0.02541169337928295, 0.056367360055446625, 0.2566111385822296, 0.13847731053829193, 0.02390860766172409, 0.10821771621704102, 0.004193281754851341, 0.024024199694395065, 0.1485961675643921], [0.010539665818214417, 0.02736317366361618, 0.020729688927531242, 0.012272891588509083, 0.037458207458257675, 0.020133765414357185, 0.006475721951574087, 0.0135318823158741, 0.14018985629081726, 0.043190933763980865, 0.014518915675580502, 0.06027117371559143, 0.013409063220024109, 0.008036705665290356, 0.12864065170288086], [0.06693296134471893, 0.05517994612455368, 0.31718623638153076, 0.09396946430206299, 0.13595829904079437, 0.09244473278522491, 0.0043823812156915665, 0.004134675953537226, 0.9252469539642334, 0.10048755258321762, 0.12945091724395752, 0.21572811901569366, 0.034586720168590546, 0.0726432204246521, 0.04207848384976387], [0.07686225324869156, 0.019675375893712044, 0.2417416274547577, 0.08641211688518524, 0.27890217304229736, 0.038729339838027954, 0.01047417800873518, 0.015033761039376259, 0.4832261800765991, 0.05870191380381584, 0.2969569265842438, 0.6193534731864929, 0.12871475517749786, 0.22289764881134033, 0.5152896642684937], [0.27357029914855957, 0.46676310896873474, 0.3964380621910095, 0.19407758116722107, 0.11257106065750122, 0.014855606481432915, 0.047355495393276215, 0.03237777575850487, 0.3466991186141968, 0.3347361087799072, 0.40522828698158264, 0.5460160970687866, 0.16927282512187958, 0.30020883679389954, 0.04839835315942764], [0.03550037741661072, 0.12907657027244568, 0.07532694190740585, 0.016156595200300217, 0.003630127990618348, 0.01967703178524971, 0.04095811769366264, 0.0179570484906435, 0.39472800493240356, 0.07661326229572296, 0.4370958209037781, 0.4819755256175995, 0.022724222391843796, 0.033822834491729736, 0.04362141340970993], [0.021909046918153763, 0.030848275870084763, 0.046106528490781784, 0.06202828511595726, 0.0325893796980381, 0.03412875533103943, 0.03159455209970474, 0.053456224501132965, 0.16627800464630127, 0.058593228459358215, 0.13071225583553314, 0.20816291868686676, 0.06561117619276047, 0.04416830837726593, 0.03868245705962181], [0.012810717336833477, 0.0013835412682965398, 0.03224228695034981, 0.08643268793821335, 0.03331959247589111, 0.030278367921710014, 0.07819522172212601, 0.03789946064352989, 0.1521843820810318, 0.04584735259413719, 0.022775838151574135, 0.3594759702682495, 0.37505412101745605, 0.4203481376171112, 0.0833948627114296], [0.12084313482046127, 0.009313090704381466, 0.17649081349372864, 0.125856414437294, 0.03634244203567505, 0.028733352199196815, 0.006864639464765787, 0.002353896852582693, 0.16829386353492737, 0.1124483197927475, 0.061692144721746445, 0.19240431487560272, 0.09329058974981308, 0.18641597032546997, 0.018957242369651794], [0.026597192510962486, 0.005893908906728029, 0.12369649112224579, 0.06400194019079208, 0.07115989178419113, 0.0058293454349040985, 0.008344992063939571, 0.00957680307328701, 0.04244829714298248, 0.036994293332099915, 0.07189996540546417, 0.04466360807418823, 0.12661096453666687, 0.2742233872413635, 0.042464204132556915], [0.0012156351003795862, 0.0009695529006421566, 0.021633058786392212, 0.003243132960051298, 0.017804604023694992, 0.006560572423040867, 0.00960883591324091, 0.043045539408922195, 0.008467147126793861, 0.0006170565611682832, 0.0028031598776578903, 0.004630656447261572, 1.7895566998049617e-05, 0.00023196694382932037, 0.14134538173675537], [0.3736850321292877, 0.29077818989753723, 0.43184730410575867, 0.4823248088359833, 0.7379603385925293, 0.5093098282814026, 0.5006043910980225, 0.3135696351528168, 0.5183887481689453, 0.13794882595539093, 0.04961319640278816, 0.12779268622398376, 0.1589212864637375, 0.22346213459968567, 0.1422436237335205], [0.15325459837913513, 0.1614270806312561, 0.4186149537563324, 0.16462315618991852, 0.44647181034088135, 0.7114150524139404, 0.12785741686820984, 0.04132780805230141, 0.047578196972608566, 0.12349404394626617, 0.3133608400821686, 0.35326144099235535, 0.30924320220947266, 0.31196898221969604, 0.028064150363206863], [0.06399086862802505, 0.06306004524230957, 0.1948489397764206, 0.12845031917095184, 0.26295408606529236, 0.38098499178886414, 0.0839061513543129, 0.02110268920660019, 0.07144157588481903, 0.01679118163883686, 0.14834797382354736, 0.479995995759964, 0.24741992354393005, 0.2288939356803894, 0.04729384183883667], [0.041305530816316605, 0.00217662681825459, 0.29091107845306396, 0.12698692083358765, 0.3031243085861206, 0.1103614866733551, 0.14891935884952545, 0.018863126635551453, 0.033797744661569595, 0.008303376846015453, 0.009713392704725266, 0.31765925884246826, 0.4755025804042816, 0.4005468487739563, 0.10761724412441254], [0.4954506754875183, 0.04642331227660179, 0.603453516960144, 0.26468321681022644, 0.3210473358631134, 0.15078485012054443, 0.027168329805135727, 0.004181328695267439, 0.10826757550239563, 0.10845811665058136, 0.053085505962371826, 0.20335085690021515, 0.12072784453630447, 0.17107200622558594, 0.059424202889204025], [0.21408557891845703, 0.03960772231221199, 0.43507251143455505, 0.10961537808179855, 0.42240580916404724, 0.06637464463710785, 0.08428787440061569, 0.03856734186410904, 0.0027873425278812647, 0.012926235795021057, 0.019708000123500824, 0.017574653029441833, 0.10679914057254791, 0.20499441027641296, 0.14648839831352234], [0.002137779025360942, 0.0005492505733855069, 0.03787382319569588, 0.004300523083657026, 0.03090864233672619, 0.003432363970205188, 0.010591491125524044, 0.028211969882249832, 0.003533262060955167, 0.0003883022291120142, 0.0014010752784088254, 0.0010855919681489468, 8.133743904181756e-06, 7.628504681633785e-05, 0.13786831498146057], [0.39364972710609436, 0.15414100885391235, 0.5289453864097595, 0.2158767729997635, 0.8369554877281189, 0.5879349708557129, 0.29191306233406067, 0.1240038275718689, 0.0375535674393177, 0.006134674418717623, 0.003127586329355836, 0.02892274223268032, 0.023530103266239166, 0.026029296219348907, 0.16074688732624054], [0.2684386968612671, 0.29252222180366516, 0.6921796798706055, 0.1771971732378006, 0.6445736885070801, 0.7333542704582214, 0.14767038822174072, 0.04686985909938812, 0.030383678153157234, 0.06000908464193344, 0.1879548877477646, 0.5258318781852722, 0.3533342778682709, 0.3370157778263092, 0.05586722865700722], [0.0015460141003131866, 0.010688474401831627, 0.09971211850643158, 0.017146917060017586, 0.1899741291999817, 0.03437719866633415, 0.022833971306681633, 0.015900788828730583, 0.05731913447380066, 0.0008445536368526518, 0.0073861475102603436, 0.06343144923448563, 0.11084617674350739, 0.11975067108869553, 0.13715405762195587]], [[0.021257108077406883, 0.04756314679980278, 0.05559564009308815, 0.030912479385733604, 0.2625647187232971, 0.138688862323761, 0.027820995077490807, 0.05787678435444832, 0.3002224862575531, 0.018701573833823204, 0.027547171339392662, 0.19844435155391693, 0.1917300671339035, 0.07151354849338531, 0.16648255288600922], [0.4235764741897583, 0.10086580365896225, 0.07221788167953491, 0.13654322922229767, 0.04923773929476738, 0.06516944617033005, 0.07642015814781189, 0.147566020488739, 0.013325832784175873, 0.07923475652933121, 0.03588176146149635, 0.02368854358792305, 0.12847480177879333, 0.04384613409638405, 0.18713882565498352], [0.8895729184150696, 0.7431688904762268, 0.3041851818561554, 0.5492796897888184, 0.7013789415359497, 0.2035668045282364, 0.4541507959365845, 0.17740322649478912, 0.37418368458747864, 0.7257221937179565, 0.3302299678325653, 0.32646968960762024, 0.4535413682460785, 0.2710181474685669, 0.06444819271564484], [0.18918083608150482, 0.07354198396205902, 0.03709281235933304, 0.039312511682510376, 0.2119109183549881, 0.32255253195762634, 0.06547961384057999, 0.022612132132053375, 0.0069438498467206955, 0.04682554677128792, 0.04775600507855415, 0.10260774195194244, 0.060122229158878326, 0.07651683688163757, 0.11037445813417435], [0.05778415873646736, 0.1888784021139145, 0.12087801843881607, 0.08340981602668762, 0.2725185453891754, 0.956253707408905, 0.6455949544906616, 0.6532288789749146, 0.3585406243801117, 0.18532338738441467, 0.18782632052898407, 0.09142936766147614, 0.8097347617149353, 0.3558001220226288, 0.037162330001592636], [0.04896414652466774, 0.25620371103286743, 0.11985385417938232, 0.0157163105905056, 0.14219185709953308, 0.22957918047904968, 0.36173656582832336, 0.07001917064189911, 0.3676673173904419, 0.12105175852775574, 0.22853095829486847, 0.07480601221323013, 0.5630075335502625, 0.8219463229179382, 0.12425509095191956], [0.04714362695813179, 0.01630709134042263, 0.04501143842935562, 0.03696214035153389, 0.036871057003736496, 0.14248797297477722, 0.08399422466754913, 0.03027486614882946, 0.0030259382911026478, 0.019033554941415787, 0.2224818617105484, 0.033125121146440506, 0.02079186774790287, 0.04913722351193428, 0.46250322461128235], [0.033912286162376404, 0.0072718155570328236, 0.013269636780023575, 0.010754123330116272, 0.003932052757591009, 0.022333307191729546, 0.05135813727974892, 0.17082874476909637, 0.004249163903295994, 0.009168761782348156, 0.00692910747602582, 0.00042953240335918963, 0.008801857940852642, 0.008872170932590961, 0.02866899035871029], [0.026226887479424477, 0.006219716742634773, 0.016528652980923653, 0.019500089809298515, 0.009756595827639103, 0.01771577261388302, 0.10877248644828796, 0.07924166321754456, 0.026382839307188988, 0.007807224057614803, 0.018975039944052696, 0.009491248056292534, 0.042680755257606506, 0.025040525943040848, 0.31068748235702515], [0.0181743074208498, 0.0022439020685851574, 0.027739310637116432, 0.07926302403211594, 0.007397042121738195, 0.01831221394240856, 0.057637136429548264, 0.025927647948265076, 0.03431807458400726, 0.03189869597554207, 0.20874466001987457, 0.006929311901330948, 0.08810199052095413, 0.09789149463176727, 0.25120988488197327], [0.0006848929915577173, 0.00015734595945104957, 0.0022563491947948933, 0.00281638465821743, 0.00390908308327198, 0.012311742641031742, 0.006667551584541798, 0.010898235253989697, 0.18826207518577576, 0.0010989188449457288, 0.003811799455434084, 0.0007082286756485701, 0.0025871950201690197, 0.0005297476891428232, 0.004719105549156666], [0.008918036706745625, 0.01932302489876747, 0.1743663251399994, 0.04276113957166672, 0.17357498407363892, 0.05217360332608223, 0.01903947815299034, 0.006896412931382656, 0.02532179281115532, 0.019349897280335426, 0.14434273540973663, 0.2454780638217926, 0.06247624009847641, 0.03444024175405502, 0.2827233076095581], [0.014348846860229969, 0.006216275505721569, 0.06011093780398369, 0.05047134682536125, 0.013856974430382252, 0.08402124047279358, 0.0029483914840966463, 0.0018935499247163534, 0.004232283215969801, 0.022591279819607735, 0.34387707710266113, 0.06330335885286331, 0.20501238107681274, 0.1859048306941986, 0.0244001317769289], [0.016000788658857346, 0.003648907644674182, 0.07618206739425659, 0.26581478118896484, 0.00828572828322649, 0.01491115428507328, 0.006984202191233635, 0.00572665361687541, 0.007784067187458277, 0.03336494415998459, 0.19996345043182373, 0.0026567107997834682, 0.14645317196846008, 0.1677580624818802, 0.0739188864827156], [0.033913157880306244, 0.5720782279968262, 0.09820353239774704, 0.06329890340566635, 0.10058190673589706, 0.8026418685913086, 0.08380495011806488, 0.37448471784591675, 0.04885341227054596, 0.01422097533941269, 0.32552391290664673, 0.701602578163147, 0.9988673329353333, 0.9602208137512207, 0.015194611623883247], [0.01701497472822666, 0.004510161932557821, 0.04222021996974945, 0.131240576505661, 0.007172171492129564, 0.0009335885988548398, 0.0025300730485469103, 0.0012859954731538892, 0.013300590217113495, 0.05520036071538925, 0.2908037602901459, 0.0021335158962756395, 0.11976832151412964, 0.046004947274923325, 0.029495948925614357], [0.0007848403765819967, 0.002563882153481245, 0.003471110016107559, 0.009534057229757309, 0.012083875946700573, 0.006908607203513384, 0.0028729254845529795, 0.0018324146512895823, 0.009593485854566097, 0.008395246230065823, 0.009609236381947994, 0.05064208433032036, 0.00595981115475297, 0.002902570180594921, 0.2071433663368225], [0.008253121748566628, 0.01393465232104063, 0.03316362947225571, 0.045629892498254776, 0.015712177380919456, 0.15894818305969238, 0.02510240487754345, 0.013996893540024757, 0.6886083483695984, 0.014645315706729889, 0.04062162712216377, 0.02812274731695652, 0.10265076905488968, 0.10770027339458466, 0.07716524600982666], [0.0017006727866828442, 0.008613905869424343, 0.08540165424346924, 0.014788517728447914, 0.11802737414836884, 0.058780014514923096, 0.008085138164460659, 0.003584004705771804, 0.06396479159593582, 0.006658769678324461, 0.02042919024825096, 0.3806440234184265, 0.01375669613480568, 0.01512871216982603, 0.1676391214132309], [0.017164628952741623, 0.028738657012581825, 0.06823595613241196, 0.08604145050048828, 0.04855107143521309, 0.24198594689369202, 0.008688676171004772, 0.003311790293082595, 0.059665460139513016, 0.08214288204908371, 0.34741461277008057, 0.15404720604419708, 0.18822570145130157, 0.19501997530460358, 0.062469229102134705], [0.04490135982632637, 0.02318926900625229, 0.15967297554016113, 0.36984479427337646, 0.027114713564515114, 0.1867561787366867, 0.04668368771672249, 0.02171866036951542, 0.05653616786003113, 0.08818016946315765, 0.14142879843711853, 0.002535451203584671, 0.06232175603508949, 0.12099058926105499, 0.16113655269145966], [0.07898441702127457, 0.817236065864563, 0.29267793893814087, 0.16063392162322998, 0.31295838952064514, 0.9265751838684082, 0.1967003047466278, 0.5436303615570068, 0.2332589328289032, 0.04864489659667015, 0.5440958142280579, 0.8931991457939148, 0.9993566870689392, 0.9798612594604492, 0.03687797114253044], [0.051174335181713104, 0.009388554841279984, 0.15813162922859192, 0.3707107603549957, 0.02142486348748207, 0.01361497025936842, 0.01679075136780739, 0.00489152641966939, 0.08238242566585541, 0.07653495669364929, 0.14888693392276764, 0.003932347521185875, 0.1416105329990387, 0.05760091543197632, 0.13266737759113312], [0.00042274355655536056, 0.0019217034569010139, 0.0013128711143508554, 0.004135955590754747, 0.004101510625332594, 0.004091422073543072, 0.0013299065176397562, 0.0007323773461394012, 0.006002569571137428, 0.003528070170432329, 0.004258603788912296, 0.04385730251669884, 0.006557406857609749, 0.0025679266545921564, 0.1728060394525528], [0.0034927180968225002, 0.014745223335921764, 0.025302981957793236, 0.04650698974728584, 0.0658985823392868, 0.10278132557868958, 0.009682145901024342, 0.010841106064617634, 0.1757735013961792, 0.03157021477818489, 0.006062814965844154, 0.2611170709133148, 0.3153221011161804, 0.08490109443664551, 0.13624651730060577]], [[0.01888529770076275, 0.5547894835472107, 0.0062187607400119305, 0.02304725907742977, 0.007431741803884506, 0.05333258956670761, 0.13557927310466766, 0.09608769416809082, 0.011193820275366306, 0.006900292821228504, 0.007560353726148605, 0.018807610496878624, 0.018169475719332695, 0.07717052102088928, 0.1439915895462036], [0.045791856944561005, 0.14471176266670227, 0.057932548224925995, 0.15441685914993286, 0.011981116607785225, 0.030152589082717896, 0.13976308703422546, 0.003811573376879096, 0.010053272359073162, 0.1557283103466034, 0.05080341920256615, 0.00967743806540966, 0.003085661679506302, 0.003445286303758621, 0.08783376961946487], [0.010936958715319633, 0.0031021125614643097, 0.009866965003311634, 0.09017129242420197, 0.02775183692574501, 0.0016267865430563688, 0.01958146132528782, 0.003049993421882391, 0.009465858340263367, 0.022049162536859512, 0.013875926844775677, 0.002902107546105981, 0.0008567434852011502, 0.0034160439390689135, 0.13799139857292175], [0.10994840413331985, 0.15032780170440674, 0.0035718681756407022, 0.1491042822599411, 0.020450405776500702, 0.013510379940271378, 0.47067153453826904, 0.6447877883911133, 0.18023402988910675, 0.1876010298728943, 0.011866661719977856, 0.006677938625216484, 0.0005242988117970526, 0.004238110035657883, 0.29615819454193115], [0.06992093473672867, 0.2791251242160797, 0.006900451611727476, 0.053067900240421295, 0.010168666951358318, 0.0023874202743172646, 0.05137968435883522, 0.06462283432483673, 0.11192043125629425, 0.10690896213054657, 0.009735661558806896, 0.04335656389594078, 0.0031411510426551104, 0.011707558296620846, 0.14929862320423126], [0.24040630459785461, 0.43853774666786194, 0.0175826046615839, 0.06282828748226166, 0.03055599145591259, 0.20223812758922577, 0.5439046025276184, 0.8139520287513733, 0.30283859372138977, 0.4911571145057678, 0.09772597998380661, 0.1337594985961914, 0.08667796850204468, 0.03606351464986801, 0.12256386131048203], [0.03999294713139534, 0.1864590346813202, 0.003897173795849085, 0.04184543341398239, 0.0012414547381922603, 0.025941016152501106, 0.05348599702119827, 0.5434274673461914, 0.012460692785680294, 0.31306707859039307, 0.06930337846279144, 0.0021947044879198074, 0.023592861369252205, 0.04260588437318802, 0.01969532109797001], [0.053744781762361526, 0.006899113766849041, 0.0563664473593235, 0.12695427238941193, 0.012777185067534447, 0.08455551415681839, 0.11441048979759216, 0.13062608242034912, 0.19371363520622253, 0.6254263520240784, 0.24294114112854004, 0.020724456757307053, 0.019838949665427208, 0.022365091368556023, 0.1131007969379425], [0.11661048978567123, 0.35882315039634705, 0.03118491731584072, 0.06881216168403625, 0.014698721468448639, 0.0038598491810262203, 0.1485612690448761, 0.39066970348358154, 0.07792866975069046, 0.22571811079978943, 0.040231697261333466, 0.265895277261734, 0.2000368982553482, 0.1125464141368866, 0.24931347370147705], [0.03291217237710953, 0.23853188753128052, 0.04644821211695671, 0.031600918620824814, 0.045192934572696686, 0.0019951597787439823, 0.11113008856773376, 0.36339887976646423, 0.010439107194542885, 0.20188210904598236, 0.027288423851132393, 0.21054767072200775, 0.04143378138542175, 0.0853629931807518, 0.2336580902338028], [0.07334253191947937, 0.14656193554401398, 0.004660916980355978, 0.03353964164853096, 0.00998624786734581, 0.00235390174202621, 0.04832129552960396, 0.031250230967998505, 0.0017524310387670994, 0.10710166394710541, 0.04863408952951431, 0.11276239901781082, 0.00949337612837553, 0.024303043261170387, 0.5020502805709839], [0.15921767055988312, 0.18694822490215302, 0.011401425115764141, 0.15920288860797882, 0.0017978762043640018, 0.00600996520370245, 0.1401643455028534, 0.08585444837808609, 0.05989503860473633, 0.2726706564426422, 0.041456613689661026, 0.0019109381828457117, 0.0026012342423200607, 0.00675933575257659, 0.05683350935578346], [0.6248686909675598, 0.8166397213935852, 0.05456394702196121, 0.3034517765045166, 0.0032548136077821255, 0.03656908869743347, 0.3933179974555969, 0.635881781578064, 0.4090532660484314, 0.6309216618537903, 0.09238837659358978, 0.01225167978554964, 0.0038302247412502766, 0.05015851929783821, 0.4316881597042084], [0.6506885886192322, 0.26984432339668274, 0.19192098081111908, 0.45030322670936584, 0.018604522570967674, 0.06438936293125153, 0.16284945607185364, 0.46218666434288025, 0.2198290228843689, 0.6063108444213867, 0.13934792578220367, 0.19822801649570465, 0.009406321682035923, 0.07906869053840637, 0.39550670981407166], [0.6516265273094177, 0.3494286835193634, 0.13445304334163666, 0.40472084283828735, 0.05377691984176636, 0.043724507093429565, 0.6220480799674988, 0.09338771551847458, 0.1620686650276184, 0.8232020139694214, 0.17699383199214935, 0.03535428270697594, 4.775904380949214e-05, 0.000580178399104625, 0.13870029151439667], [0.40970566868782043, 0.3527304232120514, 0.004458754323422909, 0.09938450157642365, 0.006175781134516001, 0.014084810391068459, 0.22543573379516602, 0.4835565686225891, 0.025563040748238564, 0.39703506231307983, 0.00602720445021987, 0.0051488312892615795, 0.0008810341823846102, 0.0033910071942955256, 0.2277533859014511], [0.19487805664539337, 0.1991150975227356, 0.010765495710074902, 0.08231080323457718, 0.014791524969041348, 0.005413876846432686, 0.2905171811580658, 0.06453394889831543, 0.003980779554694891, 0.08378233760595322, 0.012941073626279831, 0.009292078204452991, 0.0008543379371985793, 0.002103410428389907, 0.1794004589319229], [0.12092277407646179, 0.17967110872268677, 0.0018819703254848719, 0.04615653306245804, 0.002711376640945673, 0.0007180452230386436, 0.10793514549732208, 0.09669310599565506, 0.0005949889309704304, 0.15432700514793396, 0.015202132984995842, 0.003636009059846401, 0.00047353014815598726, 0.0022874167189002037, 0.22825637459754944], [0.14498451352119446, 0.2535317540168762, 0.027076847851276398, 0.14632807672023773, 0.0057570356875658035, 0.011071202345192432, 0.31473973393440247, 0.2956455647945404, 0.07720959931612015, 0.1944134682416916, 0.008117430843412876, 0.0006636073812842369, 0.0008167477208189666, 0.0018315445631742477, 0.15913215279579163], [0.22215187549591064, 0.47823596000671387, 0.018273456022143364, 0.13293205201625824, 0.0049734353087842464, 0.0265207476913929, 0.27213141322135925, 0.33180302381515503, 0.1344960778951645, 0.335622638463974, 0.010143149644136429, 0.0012862810399383307, 0.00035499766818247736, 0.0037611438892781734, 0.27220219373703003], [0.3673586845397949, 0.057844266295433044, 0.06040150299668312, 0.09888742864131927, 0.023171812295913696, 0.05270017683506012, 0.11794743686914444, 0.1507657766342163, 0.008498218841850758, 0.09498187899589539, 0.003615680383518338, 0.010834122076630592, 0.00024780313833616674, 0.0017297717276960611, 0.20351538062095642], [0.6060628294944763, 0.1373525857925415, 0.13755829632282257, 0.4113396406173706, 0.07285188883543015, 0.014519162476062775, 0.5372579097747803, 0.0630655512213707, 0.14564833045005798, 0.695697009563446, 0.06662726402282715, 0.006644518580287695, 1.2849791346525308e-05, 0.00011718441965058446, 0.13694217801094055], [0.16518473625183105, 0.10184229910373688, 0.002064367523416877, 0.05309450253844261, 0.004080682527273893, 0.012669779360294342, 0.18988992273807526, 0.5354599356651306, 0.004024976398795843, 0.07357845455408096, 0.00022774768876843154, 0.00034433722612448037, 4.428778629517183e-05, 0.00011935137445107102, 0.17481543123722076], [0.060375016182661057, 0.09738604724407196, 0.004719918128103018, 0.05357348173856735, 0.007510221563279629, 0.002087255474179983, 0.1777726411819458, 0.04658319056034088, 0.0022654803469777107, 0.02657914347946644, 0.002838509390130639, 0.0023206211626529694, 0.00029234393150545657, 0.0006460589938797057, 0.15720529854297638], [0.006292517296969891, 0.056422796100378036, 0.003871192689985037, 0.016857203096151352, 0.0060961381532251835, 0.01021772250533104, 0.02558758109807968, 0.004345982801169157, 0.003136568469926715, 0.011386821046471596, 0.0007550015579909086, 0.014218548312783241, 0.002899263286963105, 0.00665974011644721, 0.1386014223098755]], [[0.19101674854755402, 0.0880991518497467, 0.25550922751426697, 0.3376496732234955, 0.25425824522972107, 0.2177356481552124, 0.35922226309776306, 0.13405567407608032, 0.2859460711479187, 0.47983312606811523, 0.235154390335083, 0.26708394289016724, 0.2646999657154083, 0.4890832304954529, 0.0349225178360939], [0.12788966298103333, 0.14897412061691284, 0.18708589673042297, 0.1539590060710907, 0.06750026345252991, 0.06459501385688782, 0.24742794036865234, 0.0008040289394557476, 0.08417094498872757, 0.08338519930839539, 0.09756942838430405, 0.05163748189806938, 0.06044981628656387, 0.1204136312007904, 0.005185095127671957], [0.00823432207107544, 0.006774595472961664, 0.011488616466522217, 0.031759701669216156, 0.014620696194469929, 0.015192853286862373, 0.015498323366045952, 0.001623230637051165, 0.04214249551296234, 0.022796856239438057, 0.0813785269856453, 0.058821164071559906, 0.018185952678322792, 0.030505431815981865, 0.13797427713871002], [0.07304069399833679, 0.17316529154777527, 0.0638275146484375, 0.06216027960181236, 0.10879980027675629, 0.2286580353975296, 0.12489848583936691, 0.06798849999904633, 0.12340370565652847, 0.11364749073982239, 0.33209869265556335, 0.7156579494476318, 0.917570948600769, 0.8780012726783752, 0.004697424825280905], [0.04041377454996109, 0.06032548099756241, 0.013153426349163055, 0.12010756880044937, 0.032379359006881714, 0.02533758245408535, 0.03651244193315506, 0.05168384686112404, 0.05184069648385048, 0.20407944917678833, 0.10554968565702438, 0.5571502447128296, 0.039276935160160065, 0.10380254685878754, 0.1458612084388733], [0.025283029302954674, 0.14580176770687103, 0.0262577123939991, 0.01834816485643387, 0.02426275424659252, 0.5010125637054443, 0.025797395035624504, 0.08120379596948624, 0.10846563428640366, 0.05807282403111458, 0.047331083565950394, 0.01890925131738186, 0.041984543204307556, 0.021773895248770714, 0.12734822928905487], [0.11099886894226074, 0.272359162569046, 0.07267793267965317, 0.02685651369392872, 0.04662291333079338, 0.6599292755126953, 0.15850403904914856, 0.1944371908903122, 0.02196124941110611, 0.18415939807891846, 0.2094753533601761, 0.11699666827917099, 0.8625363111495972, 0.6611498594284058, 0.034588079899549484], [0.10045554488897324, 0.003808635985478759, 0.012772331945598125, 0.008206314407289028, 0.016907531768083572, 0.2308196723461151, 0.04502535238862038, 0.16794730722904205, 0.14683513343334198, 0.07804886251688004, 0.12962646782398224, 0.03242946416139603, 0.45433515310287476, 0.3931583762168884, 0.023861808702349663], [0.020261207595467567, 0.011864200234413147, 0.013516101986169815, 0.00783876795321703, 0.006360001862049103, 0.5825139880180359, 0.27136117219924927, 0.28645893931388855, 0.002775657456368208, 0.05587191879749298, 0.01021821890026331, 0.03437367081642151, 0.37942126393318176, 0.11788230389356613, 0.047214996069669724], [0.3444993495941162, 0.4299255907535553, 0.3897337317466736, 0.11608962714672089, 0.07001375406980515, 0.1826992928981781, 0.3195875883102417, 0.1513850837945938, 0.014436168596148491, 0.25265297293663025, 0.18822813034057617, 0.20145024359226227, 0.648497998714447, 0.6856710314750671, 0.13566814363002777], [0.37375974655151367, 0.2605052888393402, 0.636468231678009, 0.14340142905712128, 0.5107957124710083, 0.683059811592102, 0.3617965579032898, 0.3775153160095215, 0.0734284520149231, 0.5245854258537292, 0.5329803228378296, 0.541839063167572, 0.8546188473701477, 0.8892531991004944, 0.08003345131874084], [0.1478864699602127, 0.26107946038246155, 0.2706110179424286, 0.022070137783885002, 0.08394861966371536, 0.7104908227920532, 0.22173403203487396, 0.18465854227542877, 0.3481738865375519, 0.02706378884613514, 0.14399166405200958, 0.24452990293502808, 0.3432118594646454, 0.3138853907585144, 0.0603480227291584], [0.03315366804599762, 0.109662726521492, 0.165960431098938, 0.03089676797389984, 0.00589095801115036, 0.7119044065475464, 0.04612211138010025, 0.03627030551433563, 0.019800378009676933, 0.02169116772711277, 0.07954178750514984, 0.014483828097581863, 0.3210127055644989, 0.25073835253715515, 0.021559905260801315], [0.1801593005657196, 0.7095129489898682, 0.41699883341789246, 0.14223065972328186, 0.03218872845172882, 0.8857168555259705, 0.325775682926178, 0.46090880036354065, 0.31827157735824585, 0.19596631824970245, 0.36584827303886414, 0.568932831287384, 0.05918605625629425, 0.12899020314216614, 0.03239220380783081], [0.15587098896503448, 0.007851594127714634, 0.38951343297958374, 0.26023998856544495, 0.2678505480289459, 0.04164084047079086, 0.060063086450099945, 0.06729273498058319, 0.019880756735801697, 0.0442759171128273, 0.10040930658578873, 0.1083277016878128, 0.0003995952138211578, 0.001039322349242866, 0.14095477759838104], [0.08899319916963577, 0.2356371134519577, 0.40766164660453796, 0.08200893551111221, 0.14033742249011993, 0.12043434381484985, 0.050508081912994385, 0.04391980916261673, 0.2084629088640213, 0.07807423919439316, 0.06514080613851547, 0.6571899652481079, 0.6522034406661987, 0.4899447560310364, 0.0237458273768425], [0.3269592225551605, 0.23715397715568542, 0.21103474497795105, 0.29856637120246887, 0.031984660774469376, 0.019636303186416626, 0.2648169696331024, 0.0041971527971327305, 0.6909844875335693, 0.5414000153541565, 0.4092715382575989, 0.02185220457613468, 0.006548420060425997, 0.013211028650403023, 0.06752441078424454], [0.40959432721138, 0.2696213126182556, 0.4055677354335785, 0.265968382358551, 0.12281941622495651, 0.10883577167987823, 0.16766701638698578, 0.053767129778862, 0.028326192870736122, 0.5353591442108154, 0.3247348368167877, 0.03339260071516037, 0.1199125200510025, 0.14055927097797394, 0.07849014550447464], [0.0703776553273201, 0.17115768790245056, 0.14820680022239685, 0.014450321905314922, 0.036940984427928925, 0.4336852431297302, 0.18269671499729156, 0.1382565200328827, 0.5314536690711975, 0.05019254609942436, 0.11642822623252869, 0.17526941001415253, 0.3684784173965454, 0.3591882586479187, 0.09016428142786026], [0.020959746092557907, 0.2473447471857071, 0.04995026811957359, 0.032434724271297455, 0.004538285546004772, 0.38885483145713806, 0.04268676042556763, 0.035024866461753845, 0.14864443242549896, 0.14174208045005798, 0.13687251508235931, 0.021197974681854248, 0.4566997289657593, 0.37854352593421936, 0.051512595266103745], [0.11558277904987335, 0.8023946285247803, 0.11340320110321045, 0.07801315933465958, 0.012690390460193157, 0.363363116979599, 0.22989940643310547, 0.28700947761535645, 0.3164795935153961, 0.28987860679626465, 0.20186272263526917, 0.5113669037818909, 0.04614659398794174, 0.13675883412361145, 0.05756649002432823], [0.13439694046974182, 0.004173143766820431, 0.22800596058368683, 0.19857077300548553, 0.1396344006061554, 0.007145485375076532, 0.03306930512189865, 0.026599518954753876, 0.02599666267633438, 0.04890456795692444, 0.0713912844657898, 0.040079280734062195, 0.00020046728604938835, 0.0004629320465028286, 0.13767622411251068], [0.21178027987480164, 0.5613860487937927, 0.18598653376102448, 0.13814353942871094, 0.06437420845031738, 0.1469835489988327, 0.09205848723649979, 0.07043211162090302, 0.3314816355705261, 0.1618121713399887, 0.0553976409137249, 0.7871544361114502, 0.7398563027381897, 0.533365786075592, 0.06109875440597534], [0.308572918176651, 0.1810312271118164, 0.10904403775930405, 0.38784971833229065, 0.013434378430247307, 0.011286276392638683, 0.26633715629577637, 0.0027595413848757744, 0.7609409689903259, 0.7608016729354858, 0.6143397688865662, 0.036307673901319504, 0.013564765453338623, 0.02826162986457348, 0.07738469541072845], [0.1500416249036789, 0.027276279404759407, 0.32022449374198914, 0.45847558975219727, 0.23693141341209412, 0.1596660166978836, 0.2821829915046692, 0.005833256058394909, 0.32143598794937134, 0.14477354288101196, 0.029714325442910194, 0.15291856229305267, 0.007731991354376078, 0.029727784916758537, 0.12283544987440109]], [[0.2602275013923645, 0.0514441579580307, 0.4731021821498871, 0.5077798962593079, 0.22717851400375366, 0.04740440100431442, 0.27564913034439087, 0.24302659928798676, 0.05887439846992493, 0.3509802222251892, 0.6124410033226013, 0.11394976824522018, 0.0489780493080616, 0.04593530669808388, 0.01042554248124361], [0.032066281884908676, 0.1349876970052719, 0.04647025838494301, 0.02243492752313614, 0.02574889175593853, 0.03298051655292511, 0.026965852826833725, 0.3248708248138428, 0.005728535819798708, 0.08351098001003265, 0.1499667763710022, 0.16844461858272552, 0.05473209172487259, 0.05656114220619202, 0.10718395560979843], [0.005181984044611454, 0.0008690498070791364, 0.00864254217594862, 0.00306740403175354, 0.10709173232316971, 0.0007182863773778081, 0.004329775460064411, 0.010956686921417713, 0.06760676205158234, 0.010445973835885525, 0.012115269899368286, 0.06696799397468567, 0.0054829977452754974, 0.025371035560965538, 0.13854098320007324], [0.03556624799966812, 0.11754146218299866, 0.010577056556940079, 0.008073115721344948, 0.06965696066617966, 0.0032990325707942247, 0.011276635341346264, 0.09485359489917755, 0.10517128556966782, 0.0125450249761343, 0.007751243654638529, 0.0650070384144783, 0.0006160335033200681, 0.002038064645603299, 0.4774436056613922], [0.13858208060264587, 0.06875398755073547, 0.01532802265137434, 0.10744626820087433, 0.18273182213306427, 0.002165634883567691, 0.069672591984272, 0.11672408878803253, 0.005795653443783522, 0.0880894884467125, 0.05771886929869652, 0.025581423193216324, 0.03904194384813309, 0.07354751974344254, 0.14365413784980774], [0.16291819512844086, 0.050931405276060104, 0.14806726574897766, 0.2683573365211487, 0.2810481786727905, 0.002092417562380433, 0.012745368294417858, 0.01212888304144144, 0.014305775985121727, 0.17753903567790985, 0.1299620419740677, 0.10299177467823029, 0.21836693584918976, 0.06576120108366013, 0.12406044453382492], [0.12156791239976883, 0.39120492339134216, 0.1209033653140068, 0.08395244181156158, 0.29989197850227356, 0.044024936854839325, 0.023133939132094383, 0.05934688448905945, 0.02561376802623272, 0.024757277220487595, 0.04535222053527832, 0.11912120133638382, 0.02126661129295826, 0.03811139240860939, 0.248785600066185], [0.106705442070961, 0.8169862627983093, 0.1967339813709259, 0.01375850010663271, 0.13418887555599213, 0.16134029626846313, 0.005958847235888243, 0.09247319400310516, 0.04806499928236008, 0.025876127183437347, 0.08311128616333008, 0.22926460206508636, 0.05653654783964157, 0.04726153612136841, 0.20836575329303741], [0.04722486063838005, 0.04722658172249794, 0.05176655203104019, 0.00462702801451087, 0.20528024435043335, 0.0011717488523572683, 0.004415996838361025, 0.014451048336923122, 0.028127426281571388, 0.007240481209009886, 0.004411954898387194, 0.10081291943788528, 0.07703132927417755, 0.033158108592033386, 0.21852079033851624], [0.032722555100917816, 0.027063244953751564, 0.014943713322281837, 0.0013555125333368778, 0.016471203416585922, 0.005467826500535011, 0.02999643050134182, 0.014794600196182728, 0.03837134689092636, 0.004397213459014893, 0.01024235412478447, 0.04855721816420555, 0.05723624676465988, 0.051476139575242996, 0.2643129825592041], [0.052069392055273056, 0.003948261961340904, 0.01313212513923645, 0.010319330729544163, 0.04011767730116844, 0.00066552241332829, 0.01502715889364481, 0.007099903654307127, 0.16779832541942596, 0.03226454555988312, 0.052614975720644, 0.014822165481746197, 0.002071568975225091, 0.001763610984198749, 0.05304422974586487], [0.022045070305466652, 0.036587294191122055, 0.06798984855413437, 0.040110163390636444, 0.5405737161636353, 0.015278805047273636, 0.02948732301592827, 0.034845639020204544, 0.27487096190452576, 0.008005083538591862, 0.012681123800575733, 0.10707750916481018, 0.02124345488846302, 0.00868641585111618, 0.4183328449726105], [0.07479816675186157, 0.018890362232923508, 0.2873721718788147, 0.028116360306739807, 0.7967413067817688, 0.008446138352155685, 0.020726248621940613, 0.018564706668257713, 0.33813604712486267, 0.003492887830361724, 0.010393181815743446, 0.18903475999832153, 0.00443642633035779, 0.0231452826410532, 0.42231008410453796], [0.07108656316995621, 0.0021144712809473276, 0.0671088695526123, 0.03148089721798897, 0.7113023400306702, 0.006737539079040289, 0.2500847280025482, 0.023258471861481667, 0.23158760368824005, 0.011219021864235401, 0.04227704927325249, 0.03650788217782974, 0.15078191459178925, 0.09633734077215195, 0.15066072344779968], [0.04487757384777069, 0.009540342725813389, 0.2420971691608429, 0.01275626104325056, 0.3918483257293701, 0.0218670591711998, 0.022137846797704697, 0.08132637292146683, 0.11900310963392258, 0.000993919325992465, 0.03630243241786957, 0.087126724421978, 0.0003738462692126632, 0.02454514056444168, 0.14072805643081665], [0.0048965876922011375, 0.019337626174092293, 0.002879639156162739, 0.0027576948050409555, 0.04260760545730591, 0.003218113211914897, 0.003307115286588669, 0.026640478521585464, 0.011750566773116589, 0.0005104524316266179, 9.575913281878456e-05, 0.057879798114299774, 0.004244217649102211, 0.00609983503818512, 0.28528884053230286], [0.0335795059800148, 0.030716734007000923, 0.023829646408557892, 0.03415534272789955, 0.08875380456447601, 0.0019310596399009228, 0.017619425430893898, 0.012105603702366352, 0.002468202030286193, 0.010380377061665058, 0.01267782598733902, 0.10606792569160461, 0.0014069904573261738, 0.0004161447286605835, 0.19442977011203766], [0.17404082417488098, 0.05758971348404884, 0.12847737967967987, 0.07598815858364105, 0.49957963824272156, 0.003085564589127898, 0.05114232748746872, 0.011464038863778114, 0.06926580518484116, 0.06844814121723175, 0.06813240051269531, 0.08604259043931961, 0.004740274045616388, 0.009239559061825275, 0.19994765520095825], [0.011875619180500507, 0.026503771543502808, 0.054018229246139526, 0.01668175496160984, 0.3499281406402588, 0.01803278550505638, 0.01878167688846588, 0.01221490278840065, 0.15005004405975342, 0.0046301730908453465, 0.005843435879796743, 0.032064031809568405, 0.010490885935723782, 0.00555034726858139, 0.27147379517555237], [0.0646943747997284, 0.047236885875463486, 0.11903148144483566, 0.02203843556344509, 0.4764179587364197, 0.008550588972866535, 0.013687309809029102, 0.008890991099178791, 0.32491248846054077, 0.011557912454009056, 0.009869826957583427, 0.0921611338853836, 0.0031256151851266623, 0.016340140253305435, 0.3438139855861664], [0.17560914158821106, 0.007353567518293858, 0.056802812963724136, 0.032415200024843216, 0.4015137553215027, 0.02137722261250019, 0.35710790753364563, 0.018633568659424782, 0.05862341821193695, 0.02506905421614647, 0.018169963732361794, 0.009134531952440739, 0.07779684662818909, 0.07867905497550964, 0.1750962883234024], [0.05210466682910919, 0.006375414319336414, 0.22638031840324402, 0.012961659580469131, 0.3225522041320801, 0.012402641586959362, 0.024030247703194618, 0.056293144822120667, 0.11919546872377396, 0.0012290689628571272, 0.027758106589317322, 0.025181178003549576, 0.00022994892788119614, 0.012616506777703762, 0.1375768631696701], [0.005459210369735956, 0.03143180534243584, 0.0014205367770045996, 0.0012642937945201993, 0.01687682792544365, 0.007108580321073532, 0.004234722815454006, 0.017920657992362976, 0.003724986221641302, 0.0002761750074569136, 2.4563792976550758e-05, 0.011889445595443249, 0.0013067404506728053, 0.002636768389493227, 0.19040453433990479], [0.031027475371956825, 0.05656901001930237, 0.0113890515640378, 0.024300340563058853, 0.03550150617957115, 0.0024159413296729326, 0.02035972848534584, 0.01581081561744213, 0.002032301388680935, 0.009238713420927525, 0.01651322841644287, 0.11367840319871902, 0.003108791308477521, 0.00086622079834342, 0.16520220041275024], [0.7154905796051025, 0.15825338661670685, 0.49722805619239807, 0.38231807947158813, 0.39668020606040955, 0.051081933081150055, 0.4188354015350342, 0.3623049259185791, 0.3077245056629181, 0.4494604766368866, 0.7933229804039001, 0.20231026411056519, 0.27286192774772644, 0.2623305022716522, 0.06808917224407196]], [[0.437301367521286, 0.15179137885570526, 0.09085877984762192, 0.06997784972190857, 0.17732757329940796, 0.23180970549583435, 0.11514479666948318, 0.32073739171028137, 0.15501314401626587, 0.1294255405664444, 0.06762269139289856, 0.21488851308822632, 0.2614101469516754, 0.12734454870224, 0.049641113728284836], [0.028495818376541138, 0.1544514149427414, 0.06366834789514542, 0.016971074044704437, 0.02302762120962143, 0.054101087152957916, 0.012630121782422066, 0.018889501690864563, 0.004939573351293802, 0.01251249760389328, 0.1164683923125267, 0.009905983693897724, 0.01818472519516945, 0.01017050538212061, 0.04256897792220116], [0.007633751258254051, 0.002589557319879532, 0.02251260355114937, 0.05040144920349121, 0.032673582434654236, 0.0022981506772339344, 0.00627527991309762, 0.0006094649434089661, 0.01362280547618866, 0.006205975078046322, 0.006417383905500174, 0.0010467394022271037, 0.0010408272501081228, 0.007578521966934204, 0.13823428750038147], [0.0074798669666051865, 0.011802621185779572, 0.3115181624889374, 0.22458955645561218, 0.10706131160259247, 0.016402821987867355, 0.046956516802310944, 0.004200803115963936, 0.01468481682240963, 0.014471452683210373, 0.27619558572769165, 0.0038709931541234255, 0.00034889893140643835, 0.0020716534927487373, 0.01783183217048645], [0.015254770405590534, 0.01172303594648838, 0.002065492793917656, 0.005149758420884609, 0.013159574940800667, 0.001197350095026195, 0.018971139565110207, 0.004385960288345814, 0.06813318282365799, 0.021520443260669708, 0.005575989838689566, 0.001505104242824018, 0.0019181625684723258, 0.005167691968381405, 0.15193934738636017], [0.026872141286730766, 0.003412047168239951, 0.03895608335733414, 0.03612855076789856, 0.02536499686539173, 0.03102046251296997, 0.004315483849495649, 0.0027427596505731344, 0.03512648865580559, 0.022632958367466927, 0.05171700567007065, 0.0026941397227346897, 0.0031264815479516983, 0.024213580414652824, 0.12838274240493774], [0.0600903183221817, 0.002928798785433173, 0.0064612883143126965, 0.05414368212223053, 0.029363246634602547, 0.006244697142392397, 0.397325724363327, 0.040878646075725555, 0.005305922590196133, 0.27715954184532166, 0.04618077725172043, 0.008418801240622997, 0.01155431941151619, 0.05281350389122963, 0.025860372930765152], [0.0013151391176506877, 0.002262294292449951, 0.0012738551013171673, 0.0034272209741175175, 0.0030726443510502577, 0.04279911145567894, 0.008567760698497295, 0.17885291576385498, 0.00929640606045723, 0.001624501310288906, 0.02533317357301712, 0.005113683640956879, 0.027247918769717216, 0.07258909195661545, 0.014188846573233604], [0.3408622145652771, 0.07445694506168365, 0.03113507851958275, 0.0754152163863182, 0.014415460638701916, 0.002693483140319586, 0.09953030943870544, 0.11086118221282959, 0.5124953985214233, 0.329039990901947, 0.5092117786407471, 0.027396254241466522, 0.055544231086969376, 0.4057520925998688, 0.09588415175676346], [0.09238530695438385, 0.007053247652947903, 0.0017291916301473975, 0.005093103274703026, 0.0007437380263581872, 0.0014228186337277293, 0.02520381473004818, 0.019087698310613632, 0.47848576307296753, 0.29748132824897766, 0.057576071470975876, 0.01139640249311924, 0.004621520172804594, 0.02937469258904457, 0.015335291624069214], [0.0720675140619278, 0.012255199253559113, 0.04221949726343155, 0.09128241240978241, 0.009349699132144451, 0.008273615501821041, 0.014371694065630436, 0.01100369542837143, 0.1737149953842163, 0.16746114194393158, 0.1696900725364685, 0.014558696188032627, 0.01365632750093937, 0.0269284937530756, 0.016150163486599922], [0.052127860486507416, 0.0038822691421955824, 0.01307338010519743, 0.12611117959022522, 0.013002983294427395, 0.054914653301239014, 0.022843925282359123, 0.0017219025176018476, 0.025739489123225212, 0.3090609014034271, 0.10414470732212067, 0.006550551857799292, 0.006861968897283077, 0.010005415417253971, 0.011784915812313557], [0.074305959045887, 0.010457544587552547, 0.07050318270921707, 0.4022633135318756, 0.04945780336856842, 0.04771194979548454, 0.4660364091396332, 0.07594453543424606, 0.018491366878151894, 0.1513216346502304, 0.09796185791492462, 0.23858080804347992, 0.011272062547504902, 0.09385059028863907, 0.06640274822711945], [0.025815313681960106, 0.0033349080476909876, 0.00924734864383936, 0.012487816624343395, 0.03726305067539215, 0.016575457528233528, 0.23753590881824493, 0.025156090036034584, 0.11919926106929779, 0.04390435293316841, 0.0095932362601161, 0.04137176275253296, 0.08216788619756699, 0.1757660061120987, 0.30195334553718567], [0.05659867450594902, 0.020075146108865738, 0.01205957867205143, 0.004331792704761028, 0.052221644669771194, 0.0230423454195261, 0.0683140978217125, 0.09752152115106583, 0.2100839763879776, 0.0003861601871903986, 0.0032946986611932516, 0.0004593236662913114, 5.027504084864631e-05, 0.0022022551856935024, 0.14128009974956512], [0.08638240396976471, 0.0710444375872612, 0.06771891564130783, 0.17398057878017426, 0.05179189518094063, 0.34193578362464905, 0.2095513492822647, 0.09331211447715759, 0.052257001399993896, 0.006232596468180418, 0.002646914916113019, 0.06318453699350357, 0.019070196896791458, 0.02972061187028885, 0.2659039795398712], [0.26895081996917725, 0.1478959172964096, 0.3258365988731384, 0.404258131980896, 0.3733697533607483, 0.19055484235286713, 0.19857566058635712, 0.01781378500163555, 0.07512970268726349, 0.11693259328603745, 0.1175057590007782, 0.24425068497657776, 0.20241285860538483, 0.2411348670721054, 0.06638508290052414], [0.17850612103939056, 0.12822727859020233, 0.17801056802272797, 0.28459492325782776, 0.058830633759498596, 0.03884930908679962, 0.3513718843460083, 0.061017971485853195, 0.06718380004167557, 0.071348175406456, 0.23821549117565155, 0.03658399358391762, 0.03897847980260849, 0.20709341764450073, 0.13892877101898193], [0.4637373983860016, 0.04377487301826477, 0.15646661818027496, 0.36986854672431946, 0.09056738018989563, 0.23626187443733215, 0.11398540437221527, 0.0026716177817434072, 0.006399102043360472, 0.2626173198223114, 0.20860937237739563, 0.01349638868123293, 0.014208723790943623, 0.042171213775873184, 0.08208009600639343], [0.13806220889091492, 0.04062362387776375, 0.09515099227428436, 0.37904345989227295, 0.10653041303157806, 0.052835192531347275, 0.5728973150253296, 0.03487204387784004, 0.0029783223289996386, 0.07966885715723038, 0.03475099802017212, 0.13843636214733124, 0.006917618680745363, 0.06183210015296936, 0.1688811033964157], [0.02612869068980217, 0.003477374091744423, 0.007765303365886211, 0.0023155075032263994, 0.018893033266067505, 0.022398637607693672, 0.09549611806869507, 0.004012360703200102, 0.0013466936070472002, 0.0021441734861582518, 0.0004924506065435708, 0.006835760548710823, 0.011635211296379566, 0.023846328258514404, 0.22376547753810883], [0.08347997069358826, 0.014491320587694645, 0.015744350850582123, 0.0043899440206587315, 0.05038629099726677, 0.008546282537281513, 0.06458569318056107, 0.03869106248021126, 0.0615551732480526, 0.0002168803766835481, 0.0014501431724056602, 0.00013847390073351562, 1.5032101146061905e-05, 0.0007368824444711208, 0.13783538341522217], [0.072405144572258, 0.036094967275857925, 0.060353852808475494, 0.1382489949464798, 0.03810955956578255, 0.1803218573331833, 0.3716851472854614, 0.04992733895778656, 0.002898369450122118, 0.0008571037324145436, 0.00035707451752386987, 0.02692999318242073, 0.003073085332289338, 0.009645520709455013, 0.17640869319438934], [0.30767515301704407, 0.17313888669013977, 0.17682777345180511, 0.3453424274921417, 0.2732711434364319, 0.18888972699642181, 0.2821650207042694, 0.011036374606192112, 0.013345124199986458, 0.030917862430214882, 0.037141598761081696, 0.14430613815784454, 0.09504004567861557, 0.16429893672466278, 0.0962204858660698], [0.038221023976802826, 0.4632723033428192, 0.022520000115036964, 0.005303966347128153, 0.07163825631141663, 0.030774233862757683, 0.006099082063883543, 0.008936556056141853, 0.02098681591451168, 0.004558844491839409, 0.0029896388296037912, 0.018592750653624535, 0.20478543639183044, 0.08578886091709137, 0.1358346790075302]], [[0.04784957319498062, 0.004609245341271162, 0.006819143425673246, 0.0166594497859478, 0.006965316366404295, 0.000989345251582563, 0.006434451788663864, 0.005414100829511881, 0.027048002928495407, 0.008730669505894184, 0.003844247665256262, 0.0032386775128543377, 0.00916406698524952, 0.02474893629550934, 0.20862001180648804], [0.07474544644355774, 0.14463284611701965, 0.06348620355129242, 0.11649901419878006, 0.010943777859210968, 0.05790672451257706, 0.023460205644369125, 0.09132371097803116, 0.013804412446916103, 0.11923354864120483, 0.04609918221831322, 0.0031168698333203793, 0.02482042834162712, 0.018085025250911713, 0.06715727597475052], [0.07159372419118881, 0.23599489033222198, 0.6269188523292542, 0.2670744061470032, 0.07840307801961899, 0.7659233808517456, 0.4897821247577667, 0.7919513583183289, 0.47275444865226746, 0.20698092877864838, 0.5493778586387634, 0.516223669052124, 0.5164197683334351, 0.6560667753219604, 0.10535097867250443], [0.030506769195199013, 0.030577607452869415, 0.37364113330841064, 0.17907775938510895, 0.011576596647500992, 0.0018289608415216208, 0.0013806972419843078, 0.0006740305689163506, 0.006688407156616449, 0.02554805763065815, 0.1984224021434784, 0.0020999175030738115, 0.0001219362675328739, 0.0009508132934570312, 0.00851912796497345], [0.6425503492355347, 0.21330313384532928, 0.8213226199150085, 0.6104346513748169, 0.4307103455066681, 0.005470798350870609, 0.1284545361995697, 0.017213305458426476, 0.14068865776062012, 0.2507726550102234, 0.6069697737693787, 0.17266355454921722, 0.10257546603679657, 0.4255537688732147, 0.07138645648956299], [0.4833258390426636, 0.07765677571296692, 0.6261626482009888, 0.5845412611961365, 0.457427054643631, 0.012895571999251842, 0.037013884633779526, 0.0045295762829482555, 0.030468540266156197, 0.08583686500787735, 0.4300892949104309, 0.6064226627349854, 0.07339996099472046, 0.02218388393521309, 0.11548874527215958], [0.47047996520996094, 0.06838852912187576, 0.42273014783859253, 0.6319702863693237, 0.4177776277065277, 0.0021309976000338793, 0.00800495408475399, 0.0009326375438831747, 0.00536699453368783, 0.07440605759620667, 0.2710660994052887, 0.5013447999954224, 0.021646764129400253, 0.07749785482883453, 0.039263706654310226], [0.5323148965835571, 0.13256511092185974, 0.352451890707016, 0.6556484699249268, 0.4897412359714508, 0.22345507144927979, 0.17913641035556793, 0.12689323723316193, 0.025374194607138634, 0.169284388422966, 0.17072416841983795, 0.08815333992242813, 0.10821512341499329, 0.18704712390899658, 0.05398408696055412], [0.14081209897994995, 0.02785991132259369, 0.37397870421409607, 0.3742114305496216, 0.4757237732410431, 0.0011322007048875093, 0.0019287536852061749, 0.00011125820310553536, 0.00032575102522969246, 0.0042410544119775295, 0.007025705184787512, 0.007957610301673412, 0.0022035131696611643, 0.0008391661685891449, 0.0013405061326920986], [0.17781563103199005, 0.10205524414777756, 0.04494810104370117, 0.011432765983045101, 0.0031803075689822435, 0.6873405575752258, 0.1935015618801117, 0.2538544535636902, 0.0006125010550022125, 0.0012519293231889606, 0.0009674279135651886, 0.0007319907890632749, 0.006560447160154581, 0.0005926102166995406, 0.045413821935653687], [0.24551935493946075, 0.010881111957132816, 0.16116493940353394, 0.28567203879356384, 0.017490731552243233, 0.03198051080107689, 0.25225502252578735, 0.04009091481566429, 0.1379493623971939, 0.030329206958413124, 0.00725751556456089, 0.0005535308737307787, 0.0001769027003319934, 0.0002177381538785994, 0.11288075149059296], [0.2663186192512512, 0.0841110497713089, 0.39283427596092224, 0.3631373345851898, 0.12446267902851105, 0.0023146900348365307, 0.05166012421250343, 0.025394057855010033, 0.09723125398159027, 0.2633029520511627, 0.09458169341087341, 0.0066002910025417805, 0.0024958536960184574, 0.0033851033076643944, 0.0521465502679348], [0.032533496618270874, 0.005542360246181488, 0.14801643788814545, 0.028237437829375267, 0.09192534536123276, 0.002004631096497178, 0.0014868990983814, 0.0018816014053300023, 0.026168106123805046, 0.03666744753718376, 0.2621643543243408, 0.27366670966148376, 0.011460919864475727, 0.012693443335592747, 0.006134080700576305], [0.028670914471149445, 0.004855436272919178, 0.1069486141204834, 0.02764085866510868, 0.11977140605449677, 0.002686614403501153, 0.007388734724372625, 0.00704799173399806, 0.05677136406302452, 0.0688808336853981, 0.16234178841114044, 0.10548661649227142, 0.1935848444700241, 0.06036479026079178, 0.0025575226172804832], [0.04708265885710716, 0.030478408560156822, 0.0932990089058876, 0.24881142377853394, 0.1139858141541481, 0.03301549330353737, 0.12353643029928207, 0.18121947348117828, 0.3742617964744568, 0.11242274194955826, 0.2673158049583435, 0.05749531090259552, 0.00021243211813271046, 0.005648713558912277, 0.14063234627246857], [0.0034641579259186983, 0.015587975271046162, 0.04098831117153168, 0.025328122079372406, 0.012870541773736477, 0.002695741830393672, 0.0012444279855117202, 0.005834754556417465, 0.005115050356835127, 0.10742342472076416, 0.29450723528862, 0.004624508786946535, 0.028462348505854607, 0.09151851385831833, 0.02349407598376274], [0.00187075010035187, 0.017386021092534065, 0.0033179710153490305, 0.00216178921982646, 0.0006196821923367679, 0.0036519868299365044, 0.020315727218985558, 0.0735914558172226, 0.011879049241542816, 0.05418893322348595, 0.04255518689751625, 0.006776698864996433, 0.007105604745447636, 0.005562894977629185, 0.20312508940696716], [0.018124327063560486, 0.011053304187953472, 0.041496749967336655, 0.08067373931407928, 0.008039752952754498, 0.27361106872558594, 0.12004023045301437, 0.14489491283893585, 0.05115145817399025, 0.09850911796092987, 0.102595254778862, 0.03553636744618416, 0.03690872713923454, 0.062350839376449585, 0.18180564045906067], [0.12148405611515045, 0.0812632218003273, 0.2165963500738144, 0.1931358426809311, 0.08697410672903061, 0.006551810074597597, 0.06685828417539597, 0.03445844352245331, 0.0957593098282814, 0.40685340762138367, 0.14669549465179443, 0.05295614153146744, 0.013317806646227837, 0.016840115189552307, 0.07654187083244324], [0.00987213384360075, 0.006524993572384119, 0.026135168969631195, 0.011839349754154682, 0.033334147185087204, 0.0041054473258554935, 0.0015945311170071363, 0.0032734640408307314, 0.04142798110842705, 0.08157128095626831, 0.26105597615242004, 0.34578391909599304, 0.018666768446564674, 0.02866668626666069, 0.00917118415236473], [0.024172252044081688, 0.01827125810086727, 0.0764245018362999, 0.024589890614151955, 0.045055974274873734, 0.08366040140390396, 0.049236495047807693, 0.16330885887145996, 0.05235174670815468, 0.18916647136211395, 0.2596777379512787, 0.12284716963768005, 0.3776375353336334, 0.3416304290294647, 0.00993264652788639], [0.03498423844575882, 0.015507807955145836, 0.05400218814611435, 0.2035217136144638, 0.06879755109548569, 0.01839861460030079, 0.1265679895877838, 0.19229170680046082, 0.28682830929756165, 0.19846217334270477, 0.19391797482967377, 0.03128731623291969, 0.00016305393364746124, 0.003939830232411623, 0.1374405473470688], [0.013754391111433506, 0.07632532715797424, 0.05588589236140251, 0.060033075511455536, 0.015113652683794498, 0.024528013542294502, 0.0056539555080235004, 0.025407979264855385, 0.0030256062746047974, 0.3076882064342499, 0.2846599221229553, 0.01613902486860752, 0.07589408755302429, 0.25697121024131775, 0.08533195406198502], [0.0015476603293791413, 0.017548631876707077, 0.0017550711054354906, 0.0017123925499618053, 0.0004861274501308799, 0.0013240363914519548, 0.007671059109270573, 0.03281305357813835, 0.0013763409806415439, 0.060824256390333176, 0.04298469424247742, 0.011416267603635788, 0.012759965844452381, 0.012971585616469383, 0.16966485977172852], [0.005211545154452324, 0.0055291797034442425, 0.0040288688614964485, 0.011110500432550907, 0.002710954286158085, 0.0645279660820961, 0.01716793328523636, 0.025083528831601143, 0.010282285511493683, 0.009002536535263062, 0.0011292833369225264, 0.0045064822770655155, 0.007478337734937668, 0.004868943244218826, 0.13875910639762878]], [[0.01263146661221981, 0.08983241021633148, 0.002674827352166176, 0.0008326905663125217, 0.0032944290433079004, 0.06790440529584885, 0.02327594719827175, 0.08626140654087067, 0.0010102109517902136, 0.0009567838278599083, 0.001915089669637382, 0.019144434481859207, 0.060631223022937775, 0.04236740246415138, 0.2042645514011383], [0.12322216480970383, 0.14532910287380219, 0.08289580047130585, 0.07800436019897461, 0.016899574548006058, 0.20651613175868988, 0.15389330685138702, 0.08048079907894135, 0.023754820227622986, 0.08939354121685028, 0.05408218502998352, 0.0083498889580369, 0.16772767901420593, 0.03971855714917183, 0.029394451528787613], [0.002537816995754838, 0.0036866364534944296, 0.0026212686207145452, 0.0010326605988666415, 0.0028582154773175716, 0.0016078348271548748, 0.0024177017621695995, 0.004757970105856657, 0.007405414246022701, 0.0004943490494042635, 0.0008183143800124526, 0.0020540759433060884, 0.0008841927628964186, 0.0009274804615415633, 0.13894422352313995], [0.18076959252357483, 0.11159703880548477, 0.07333940267562866, 0.12368053197860718, 0.1442640721797943, 0.3224244713783264, 0.2286587655544281, 0.10576390475034714, 0.0873323604464531, 0.0707816481590271, 0.07077325880527496, 0.024980774149298668, 0.015894055366516113, 0.01236753724515438, 0.034113459289073944], [0.008514223620295525, 0.006442691199481487, 0.003549255197867751, 0.00919315591454506, 0.0011393448803573847, 0.0005870977183803916, 0.02400296926498413, 0.03577389195561409, 0.006469632964581251, 0.004828252829611301, 0.0027150637470185757, 9.597353346180171e-05, 0.00011822552187368274, 0.000396552961319685, 0.1521017998456955], [0.0016907083336263895, 9.336868970422074e-05, 0.0023900996893644333, 0.0018071996746584773, 0.001690928009338677, 0.0010278637055307627, 0.008010926656425, 0.0018918663263320923, 0.0009378245449624956, 0.0005185406771488488, 0.00012474792310968041, 0.00014544214354828, 2.7525844416231848e-05, 2.095987474604044e-05, 0.12926018238067627], [0.08279342949390411, 0.00717265997081995, 0.01113244891166687, 0.030300047248601913, 0.03227340802550316, 0.02679654024541378, 0.2711687386035919, 0.12656770646572113, 0.0010184150887653232, 0.0069296094588935375, 0.006689318455755711, 0.00307065830565989, 0.004024384077638388, 0.006041096989065409, 0.12722525000572205], [0.09468965977430344, 0.010531323030591011, 0.1253902167081833, 0.09483902901411057, 0.060478318482637405, 0.1959676593542099, 0.5850688219070435, 0.11734473705291748, 0.08924026787281036, 0.031869061291217804, 0.04437774419784546, 0.004531644284725189, 0.19630968570709229, 0.04580901935696602, 0.04253998026251793], [0.03443194553256035, 0.006786322686821222, 0.08545193076133728, 0.2555176913738251, 0.16119416058063507, 0.3760574460029602, 0.3180745542049408, 0.0858285129070282, 0.0052651395089924335, 0.035345133394002914, 0.0046972003765404224, 0.00805696938186884, 0.0738091915845871, 0.004572577308863401, 0.028640231117606163], [0.26599034667015076, 0.06405031681060791, 0.39913085103034973, 0.7390084862709045, 0.8533709049224854, 0.0830850899219513, 0.22198519110679626, 0.15359464287757874, 0.0286090150475502, 0.1338224709033966, 0.06985709816217422, 0.03841168060898781, 0.1308237761259079, 0.01580808497965336, 0.010780439712107182], [0.16064751148223877, 0.5348425507545471, 0.09399141371250153, 0.3709404170513153, 0.3757614493370056, 0.2272261530160904, 0.2699662148952484, 0.46868544816970825, 0.09081633388996124, 0.07856583595275879, 0.054298948496580124, 0.10659310221672058, 0.05178465321660042, 0.012835889123380184, 0.19243957102298737], [0.33067551255226135, 0.40668511390686035, 0.03748138248920441, 0.16017457842826843, 0.02931954525411129, 0.1285390406847, 0.43687552213668823, 0.6227295398712158, 0.016583241522312164, 0.054699335247278214, 0.43602558970451355, 0.028376825153827667, 0.1860552728176117, 0.202489972114563, 0.03443598374724388], [0.025147954002022743, 0.023277895525097847, 0.036982107907533646, 0.030706623569130898, 0.00253032217733562, 0.08060919493436813, 0.062497250735759735, 0.22720953822135925, 0.015824737027287483, 0.020865583792328835, 0.051981136202812195, 0.016274577006697655, 0.3496847152709961, 0.19709302484989166, 0.00854758732020855], [0.0009813109645619988, 0.0007951235747896135, 0.007896890863776207, 0.006039812229573727, 0.001424357295036316, 0.003153599100187421, 0.0010362794855609536, 0.006138501223176718, 0.00410880520939827, 0.003359388094395399, 0.008728301152586937, 0.0021525975316762924, 0.2318088710308075, 0.017491629347205162, 0.0005464124260470271], [0.008814784698188305, 0.009578033350408077, 0.008741176687180996, 0.002597709419205785, 0.0019302073633298278, 0.02750723622739315, 0.010486552491784096, 0.061721935868263245, 0.05738110467791557, 0.0038812088314443827, 0.08735688030719757, 0.00500333309173584, 3.085857315454632e-05, 0.005531619768589735, 0.14116442203521729], [0.015857994556427002, 0.010374038480222225, 0.002225207630544901, 0.002974742790684104, 0.0010843537747859955, 0.007387869525700808, 0.006818806286901236, 0.0318806953728199, 0.1651621013879776, 0.21757511794567108, 0.2911650240421295, 0.08204617351293564, 0.016449127346277237, 0.10985822230577469, 0.0020742996130138636], [0.01972219906747341, 0.20374125242233276, 0.0031293979845941067, 0.004390338435769081, 0.031924858689308167, 0.06048818305134773, 0.0774247944355011, 0.7845978140830994, 0.15838612616062164, 0.06142642721533775, 0.0820784792304039, 0.20785683393478394, 0.46646884083747864, 0.42270010709762573, 0.053927596658468246], [0.026567673310637474, 0.2768426239490509, 0.016553064808249474, 0.07253812253475189, 0.029352964833378792, 0.034967049956321716, 0.09283487498760223, 0.5970632433891296, 0.02342795394361019, 0.04057195410132408, 0.06215028092265129, 0.2966896891593933, 0.4489157795906067, 0.24187524616718292, 0.048112284392118454], [0.14453455805778503, 0.4129781723022461, 0.021322425454854965, 0.11776001751422882, 0.008680691011250019, 0.12525556981563568, 0.1459336131811142, 0.4943058490753174, 0.041365865617990494, 0.06633096933364868, 0.48416346311569214, 0.027247071266174316, 0.10342812538146973, 0.15874288976192474, 0.04535134881734848], [0.03164434805512428, 0.10487183183431625, 0.019769076257944107, 0.0709872916340828, 0.0046073514968156815, 0.12636253237724304, 0.06114564463496208, 0.5786424875259399, 0.17960773408412933, 0.15923625230789185, 0.14680741727352142, 0.04373620077967644, 0.20528176426887512, 0.14476445317268372, 0.03252548724412918], [0.03216148540377617, 0.04786192253232002, 0.0904572606086731, 0.284318745136261, 0.04915444552898407, 0.20336958765983582, 0.019341057166457176, 0.31598398089408875, 0.503376841545105, 0.2976534068584442, 0.3550446927547455, 0.318871408700943, 0.31741514801979065, 0.09137054532766342, 0.022498751059174538], [0.00784912146627903, 0.004314524121582508, 0.007757026236504316, 0.004281783476471901, 0.001910648075863719, 0.00898022297769785, 0.007197065278887749, 0.05121663585305214, 0.12398385256528854, 0.006457128562033176, 0.09335841238498688, 0.0023844544775784016, 1.3785818737233058e-05, 0.0021891386713832617, 0.13778245449066162], [0.0865921899676323, 0.029389984905719757, 0.007211814168840647, 0.022628001868724823, 0.003064699238166213, 0.026838112622499466, 0.02777392417192459, 0.17195671796798706, 0.5349084734916687, 0.37311822175979614, 0.5073185563087463, 0.12468769401311874, 0.014684900641441345, 0.11363118886947632, 0.01852630451321602], [0.021940317004919052, 0.17988227307796478, 0.0027716639451682568, 0.0058884406462311745, 0.02112143486738205, 0.056551095098257065, 0.09669405966997147, 0.8433947563171387, 0.1836535632610321, 0.048101164400577545, 0.0939687192440033, 0.12228170782327652, 0.5153423547744751, 0.4533718526363373, 0.10564926266670227], [0.07970402389764786, 0.263812392950058, 0.027112353593111038, 0.06228066235780716, 0.03007029928267002, 0.5465735197067261, 0.2176109254360199, 0.5667538046836853, 0.10334119945764542, 0.3484029769897461, 0.1586397886276245, 0.28290486335754395, 0.07807470858097076, 0.405972421169281, 0.12247955799102783]]], [[[0.02659090794622898, 0.049626123160123825, 0.04500019550323486, 0.012677792459726334, 0.33557751774787903, 0.02776678465306759, 0.02675992250442505, 0.09967876970767975, 0.04216820374131203, 0.009756066836416721, 0.0133897690102458, 0.12886802852153778, 0.03152704983949661, 0.046163998544216156, 0.21004843711853027], [0.05978302285075188, 0.18161648511886597, 0.038620203733444214, 0.022025080397725105, 0.09790226072072983, 0.04398013651371002, 0.00788698997348547, 0.04135579988360405, 0.0068543110974133015, 0.03809167072176933, 0.03150040656328201, 0.0462106354534626, 0.024762138724327087, 0.011792140081524849, 0.015839271247386932], [0.005166883580386639, 0.0005590450600720942, 0.007114546839147806, 0.0015656572068110108, 0.02179996483027935, 0.0010864944197237492, 0.0051814797334373, 0.0011148365447297692, 0.00816393457353115, 0.0019027285743504763, 0.005033016670495272, 0.010743028484284878, 0.0006906923954375088, 0.0011143455049023032, 0.16189540922641754], [0.17136499285697937, 0.002046054694801569, 0.4725193679332733, 0.24347566068172455, 0.1026763990521431, 0.00369152519851923, 0.013768541626632214, 0.003912978805601597, 0.022358577698469162, 0.06323882192373276, 0.28539538383483887, 0.009778834879398346, 0.0043070269748568535, 0.020384330302476883, 0.006856778170913458], [0.18433871865272522, 0.013500750064849854, 0.42166435718536377, 0.1935500204563141, 0.3502363860607147, 0.0009389789775013924, 0.0472395233809948, 0.015336934477090836, 0.07204270362854004, 0.07276465743780136, 0.4023721218109131, 0.016390468925237656, 0.00493515282869339, 0.01088448241353035, 0.18081046640872955], [0.01929071731865406, 3.154709338559769e-05, 0.04895680397748947, 0.04499320685863495, 0.03726757690310478, 0.0012487026397138834, 0.06078735366463661, 0.0025376947596669197, 0.023622047156095505, 0.008605116978287697, 0.05601886287331581, 0.011475598439574242, 0.0013240767875686288, 0.009706309996545315, 0.13962702453136444], [0.032548993825912476, 0.0047013829462230206, 0.08043498545885086, 0.08197268843650818, 0.43236956000328064, 0.013080407865345478, 0.006017346400767565, 0.05529334023594856, 0.01970849372446537, 0.004050384275615215, 0.0073967562057077885, 0.005829385481774807, 0.0008975209202617407, 0.0025361862499266863, 0.011671289801597595], [0.046304989606142044, 0.026358718052506447, 0.20277923345565796, 0.3021180331707001, 0.6281617879867554, 0.19840610027313232, 0.12000668793916702, 0.21165543794631958, 0.0507807619869709, 0.10083203762769699, 0.17539183795452118, 0.08392243832349777, 0.036049142479896545, 0.06088141351938248, 0.024198466911911964], [0.016816509887576103, 0.003118144813925028, 0.035858120769262314, 0.02315649762749672, 0.2957051992416382, 0.0033856350928545, 0.008419573307037354, 0.013085800223052502, 0.0065522813238203526, 0.004261805210262537, 0.0022621729876846075, 0.0015856586396694183, 0.00012999074533581734, 0.00036330719012767076, 0.004947974346578121], [0.13966688513755798, 0.051315873861312866, 0.16794879734516144, 0.17204447090625763, 0.02530861273407936, 0.1971883773803711, 0.6035643219947815, 0.35590535402297974, 0.01904589682817459, 0.14328262209892273, 0.05827813595533371, 0.12283631414175034, 0.08582676202058792, 0.021607764065265656, 0.09174748510122299], [0.07622234523296356, 0.021088531240820885, 0.13214311003684998, 0.1876712292432785, 0.09946685284376144, 0.0739995539188385, 0.16667790710926056, 0.06527374684810638, 0.2691768705844879, 0.1298666000366211, 0.20347969233989716, 0.28972044587135315, 0.16063560545444489, 0.23408198356628418, 0.02879655919969082], [0.04186922311782837, 0.028065834194421768, 0.2365874946117401, 0.22718128561973572, 0.717268168926239, 0.0283160749822855, 0.047574929893016815, 0.22635598480701447, 0.046485841274261475, 0.11764083057641983, 0.11684223264455795, 0.600357711315155, 0.07936308532953262, 0.1614740490913391, 0.02326863817870617], [0.002160860225558281, 0.00041385856457054615, 0.0032894921023398638, 0.004175879992544651, 0.09230346977710724, 0.00037096597952768207, 0.00036027038004249334, 0.000777967507019639, 0.0010948613053187728, 0.006351495627313852, 0.00803811103105545, 0.2546491026878357, 0.005140772555023432, 0.0052158161997795105, 0.0018242541700601578], [0.01453752163797617, 0.0016249779146164656, 0.07837095856666565, 0.046283330768346786, 0.5220571756362915, 0.00571427633985877, 0.011274048127233982, 0.0005770810530520976, 0.06172677502036095, 0.028573052957654, 0.1375623345375061, 0.2926015257835388, 0.17741695046424866, 0.13592077791690826, 0.025488857179880142], [0.0018199050100520253, 1.759366932674311e-05, 0.005607981700450182, 0.029583722352981567, 0.009902501478791237, 0.00240499060600996, 0.016255119815468788, 0.008434450253844261, 0.0070381201803684235, 0.006882159970700741, 0.008103356696665287, 0.009371891617774963, 3.180988642270677e-05, 0.0005422193789854646, 0.14323127269744873], [0.04913536086678505, 0.005111359525471926, 0.3943053185939789, 0.16504207253456116, 0.1333204060792923, 0.007373967207968235, 0.00649205781519413, 0.005781218875199556, 0.0696163922548294, 0.17078818380832672, 0.43588367104530334, 0.2441176176071167, 0.044073574244976044, 0.13962700963020325, 0.0038013174198567867], [0.02972331829369068, 0.032405998557806015, 0.13676248490810394, 0.2985995411872864, 0.6838041543960571, 0.17950911819934845, 0.02566559985280037, 0.299430251121521, 0.06906868517398834, 0.09219349920749664, 0.14271143078804016, 0.15384355187416077, 0.31184810400009155, 0.37699857354164124, 0.11869719624519348], [0.035901740193367004, 0.049252428114414215, 0.13651704788208008, 0.3431343734264374, 0.4621880352497101, 0.07741573452949524, 0.035817742347717285, 0.1879495084285736, 0.09167803823947906, 0.15167558193206787, 0.20264029502868652, 0.22310277819633484, 0.27972275018692017, 0.27912822365760803, 0.1079779863357544], [0.03869367763400078, 0.07609386742115021, 0.09811960905790329, 0.19582945108413696, 0.7770717144012451, 0.05828123167157173, 0.03398818522691727, 0.4334997236728668, 0.06648975610733032, 0.07675088942050934, 0.06197739765048027, 0.7435874938964844, 0.14106591045856476, 0.2445826381444931, 0.04634908586740494], [0.0033209763932973146, 0.0013802923494949937, 0.007923663593828678, 0.01537866611033678, 0.27329060435295105, 0.0012711664894595742, 0.000925537955481559, 0.0031033798586577177, 0.00518713379278779, 0.008014743216335773, 0.01865261048078537, 0.32840412855148315, 0.015081376768648624, 0.0187647957354784, 0.007287481799721718], [0.012120293453335762, 0.00801909901201725, 0.05887366458773613, 0.08173726499080658, 0.42918333411216736, 0.0074272770434618, 0.018144551664590836, 0.002390465000644326, 0.19959968328475952, 0.01595914363861084, 0.19477497041225433, 0.24081164598464966, 0.32190656661987305, 0.2620943486690521, 0.06223426014184952], [0.001324097509495914, 1.9873512428603135e-05, 0.0026336663868278265, 0.025088831782341003, 0.006480309646576643, 0.0015246026450768113, 0.009156930260360241, 0.006450172513723373, 0.006447002291679382, 0.003797400277107954, 0.0037222199607640505, 0.006030225194990635, 1.9453302229521796e-05, 0.0003723614208865911, 0.13770580291748047], [0.23361828923225403, 0.06709202378988266, 0.7719610333442688, 0.734594464302063, 0.7922726273536682, 0.049216482788324356, 0.04663456231355667, 0.060855433344841, 0.40224209427833557, 0.20935069024562836, 0.5060975551605225, 0.5454070568084717, 0.2919921875, 0.420108824968338, 0.08753460645675659], [0.01675574854016304, 0.0394110269844532, 0.07827049493789673, 0.20941881835460663, 0.5690934658050537, 0.13831959664821625, 0.015872817486524582, 0.2790753245353699, 0.07380014657974243, 0.05484941974282265, 0.11329877376556396, 0.046586740761995316, 0.27540746331214905, 0.3769146502017975, 0.12728242576122284], [0.13399043679237366, 0.38312259316444397, 0.21414920687675476, 0.1335369348526001, 0.883351743221283, 0.17629003524780273, 0.21391625702381134, 0.35840436816215515, 0.7405950427055359, 0.11166028678417206, 0.2222289741039276, 0.2562817633152008, 0.20710349082946777, 0.2988908290863037, 0.10401280969381332]], [[0.169734388589859, 0.018695855513215065, 0.1739528477191925, 0.1591939628124237, 0.2628772258758545, 0.10412096232175827, 0.10786166787147522, 0.024563027545809746, 0.26776236295700073, 0.15710414946079254, 0.04751116409897804, 0.10171505063772202, 0.02745870314538479, 0.022933470085263252, 0.11237789690494537], [0.04881957918405533, 0.17062845826148987, 0.0187830850481987, 0.030382977798581123, 0.08311481773853302, 0.03788991644978523, 0.005156277678906918, 0.026916639879345894, 0.06639944016933441, 0.03180782124400139, 0.02173716016113758, 0.05343012511730194, 0.01850084401667118, 0.0033381145913153887, 0.04681381955742836], [0.11046597361564636, 0.13029024004936218, 0.30802851915359497, 0.31618139147758484, 0.21513698995113373, 0.08858107775449753, 0.07770872116088867, 0.030179373919963837, 0.2956576347351074, 0.19506438076496124, 0.06668522953987122, 0.15814362466335297, 0.07954283803701401, 0.09008871018886566, 0.11347464472055435], [0.14630576968193054, 0.10272074490785599, 0.06626180559396744, 0.39613619446754456, 0.5213132500648499, 0.09462913125753403, 0.19745559990406036, 0.14176879823207855, 0.45916420221328735, 0.2814978361129761, 0.19076579809188843, 0.7478294968605042, 0.15201923251152039, 0.4428024888038635, 0.11204658448696136], [0.17077980935573578, 0.372023344039917, 0.03066021017730236, 0.20403380692005157, 0.25160810351371765, 0.047236956655979156, 0.19034826755523682, 0.09997845441102982, 0.22249065339565277, 0.14956896007061005, 0.12211201339960098, 0.43811750411987305, 0.32559871673583984, 0.4463178217411041, 0.1688702404499054], [0.001587467617355287, 0.0028523027431219816, 0.001275891438126564, 0.007771230302751064, 0.06833823025226593, 0.016362184658646584, 0.01554875634610653, 0.0395360104739666, 0.020186755806207657, 0.02848842740058899, 0.006796931382268667, 0.08043718338012695, 0.1258731484413147, 0.048048797994852066, 0.14538481831550598], [0.19441094994544983, 0.026329312473535538, 0.03907056525349617, 0.5187185406684875, 0.06508557498455048, 0.04464683309197426, 0.23734036087989807, 0.10510969161987305, 0.23671847581863403, 0.2550508677959442, 0.2969563603401184, 0.31371036171913147, 0.023362383246421814, 0.04756302013993263, 0.09379850327968597], [0.009693926200270653, 0.06855454295873642, 0.04046608507633209, 0.021632034331560135, 0.07003092765808105, 0.1099655032157898, 0.02166297659277916, 0.14673617482185364, 0.08559776097536087, 0.021444879472255707, 0.06376301497220993, 0.07838241755962372, 0.2981177270412445, 0.05645254626870155, 0.11510419100522995], [0.1475960612297058, 0.11415769904851913, 0.09677327424287796, 0.22716772556304932, 0.05128113925457001, 0.0685737207531929, 0.17258046567440033, 0.05221087113022804, 0.2985250651836395, 0.36185649037361145, 0.6199293732643127, 0.5016448497772217, 0.08136574923992157, 0.06544326990842819, 0.09482244402170181], [0.16866622865200043, 0.03890697658061981, 0.038960762321949005, 0.045146964490413666, 0.003443084890022874, 0.025941072031855583, 0.02535194903612137, 0.01214737631380558, 0.39030662178993225, 0.11890958994626999, 0.2736153304576874, 0.3244759440422058, 0.00968784186989069, 0.014615286141633987, 0.03826850652694702], [0.08395736664533615, 0.10560688376426697, 0.29490047693252563, 0.15838190913200378, 0.20854075253009796, 0.047574300318956375, 0.025914132595062256, 0.0076736449263989925, 0.23083198070526123, 0.11239635199308395, 0.08150741457939148, 0.3915822207927704, 0.126749187707901, 0.08327525854110718, 0.07453686743974686], [0.08537011593580246, 0.01334940642118454, 0.026223814114928246, 0.09485415369272232, 0.04081009700894356, 0.021519087255001068, 0.04835912212729454, 0.008561250753700733, 0.1425430029630661, 0.15310505032539368, 0.12245412170886993, 0.15674236416816711, 0.03265313804149628, 0.020860055461525917, 0.1338454782962799], [0.009048069827258587, 0.008220783434808254, 0.0010462020291015506, 0.0073586152866482735, 0.01628630980849266, 0.0030796914361417294, 0.0014804736711084843, 0.0016866090008988976, 0.021953675895929337, 0.024090107530355453, 0.02321471832692623, 0.2417944222688675, 0.00791110284626484, 0.012413977645337582, 0.02231968566775322], [0.02412300556898117, 0.02128133550286293, 0.018482450395822525, 0.016898121684789658, 0.07439899444580078, 0.03563898429274559, 0.04473365843296051, 0.0026737016160041094, 0.06965204328298569, 0.10727399587631226, 0.046027760952711105, 0.33166152238845825, 0.12371443957090378, 0.07036767154932022, 0.15801618993282318], [0.007644897326827049, 0.000292555516352877, 0.08444877713918686, 0.17402730882167816, 0.16615508496761322, 0.013423392549157143, 0.054235123097896576, 0.007257240824401379, 0.08712441474199295, 0.012547464109957218, 0.0328214131295681, 0.2736492455005646, 0.0037261026445776224, 0.09982366114854813, 0.13941559195518494], [0.07466596364974976, 0.11066461354494095, 0.02582395263016224, 0.1052846685051918, 0.0988694354891777, 0.13372771441936493, 0.10285167396068573, 0.04043884575366974, 0.12614820897579193, 0.00874736811965704, 0.006169801577925682, 0.3642371892929077, 0.13258321583271027, 0.14621633291244507, 0.16873647272586823], [0.23522600531578064, 0.0398484542965889, 0.3737937808036804, 0.288825660943985, 0.10485613346099854, 0.11366727948188782, 0.29695606231689453, 0.06251946091651917, 0.35146233439445496, 0.04921486973762512, 0.25325968861579895, 0.33112239837646484, 0.06967249512672424, 0.050063006579875946, 0.0896972194314003], [0.1151093989610672, 0.085483118891716, 0.1238018348813057, 0.10984596610069275, 0.07372570037841797, 0.07080911099910736, 0.04283013194799423, 0.011434272862970829, 0.6184931993484497, 0.031299810856580734, 0.1232943907380104, 0.4399086534976959, 0.16973690688610077, 0.18915507197380066, 0.06319096684455872], [0.23179487884044647, 0.03441762179136276, 0.058240070939064026, 0.17834095656871796, 0.049968671053647995, 0.038375332951545715, 0.05405527353286743, 0.00672679441049695, 0.09475977718830109, 0.0764862671494484, 0.1440851390361786, 0.11337311565876007, 0.06998162716627121, 0.031302694231271744, 0.13650138676166534], [0.037197839468717575, 0.022889001294970512, 0.00443503400310874, 0.02830665186047554, 0.056754183024168015, 0.011282439343631268, 0.008815057575702667, 0.005641489755362272, 0.03366301208734512, 0.01200089417397976, 0.022881681099534035, 0.24835483729839325, 0.020306341350078583, 0.028865927830338478, 0.09140723943710327], [0.019821494817733765, 0.0461096465587616, 0.009799499064683914, 0.008886821568012238, 0.03164605051279068, 0.03408728539943695, 0.06531291455030441, 0.004583337344229221, 0.015776870772242546, 0.0067581660114228725, 0.005247185938060284, 0.0803409293293953, 0.12878651916980743, 0.033680036664009094, 0.15540239214897156], [0.006374652031809092, 0.0003620072384364903, 0.05079201981425285, 0.10443739593029022, 0.13200052082538605, 0.007841442711651325, 0.04038690775632858, 0.005943085998296738, 0.04502689838409424, 0.005707652773708105, 0.010736361145973206, 0.17095635831356049, 0.0034604808315634727, 0.08947119116783142, 0.1356668770313263], [0.05784226581454277, 0.06101800128817558, 0.011293647810816765, 0.030310506001114845, 0.02692366950213909, 0.10355494171380997, 0.1643158346414566, 0.02146345190703869, 0.10686127096414566, 0.0006235101609490812, 0.001034505432471633, 0.12770172953605652, 0.08152752369642258, 0.06569667905569077, 0.13584844768047333], [0.24130187928676605, 0.04057329148054123, 0.37395209074020386, 0.32695549726486206, 0.18701796233654022, 0.1542418897151947, 0.4307348132133484, 0.07850468903779984, 0.24226921796798706, 0.027551302686333656, 0.17328326404094696, 0.256756991147995, 0.1007629856467247, 0.0746576264500618, 0.1026487648487091], [0.18065117299556732, 0.0850963443517685, 0.37481072545051575, 0.36960142850875854, 0.042269542813301086, 0.04689870774745941, 0.10553675144910812, 0.031215613707900047, 0.03850337490439415, 0.055640675127506256, 0.11964564025402069, 0.20274300873279572, 0.22541530430316925, 0.07314471900463104, 0.12492100149393082]], [[0.2626786530017853, 0.0849713385105133, 0.11954734474420547, 0.09299539029598236, 0.12019845843315125, 0.1675114780664444, 0.12060416489839554, 0.1292921006679535, 0.33819568157196045, 0.3146125078201294, 0.20831438899040222, 0.39596518874168396, 0.2145393043756485, 0.2666572332382202, 0.05294949933886528], [0.1368129849433899, 0.16135744750499725, 0.15528292953968048, 0.24771884083747864, 0.1416730433702469, 0.05803852900862694, 0.07394444942474365, 0.10563277453184128, 0.033661823719739914, 0.18054474890232086, 0.1985052525997162, 0.05316935107111931, 0.05009648948907852, 0.043446026742458344, 0.03412564843893051], [0.0030849967151880264, 0.0006440586876124144, 0.016017315909266472, 0.0037563794758170843, 0.009170617908239365, 0.0008218333241529763, 0.0032779525499790907, 0.0006974118296056986, 0.12044321000576019, 0.005983977112919092, 0.011704917997121811, 0.023849062621593475, 0.0031650178134441376, 0.01169323269277811, 0.16145823895931244], [0.02798222377896309, 0.012448069639503956, 0.018199993297457695, 0.0069459048099815845, 0.042531996965408325, 0.009718443267047405, 0.013791781850159168, 0.04370715469121933, 0.21814176440238953, 0.024645699188113213, 0.0633857473731041, 0.0802498310804367, 0.006771658081561327, 0.040147896856069565, 0.4109969139099121], [0.02001010812819004, 0.02580004744231701, 0.006869276985526085, 0.007543967105448246, 0.017537932842969894, 0.00023914838675409555, 0.006739956792443991, 0.008227680809795856, 0.05446772649884224, 0.03320171311497688, 0.022232946008443832, 0.01063306163996458, 0.0007752752280794084, 0.0028256638906896114, 0.2078467756509781], [0.0034786108881235123, 0.00011826713307527825, 0.002407492371276021, 0.005452741403132677, 0.002847136929631233, 0.003419033018872142, 0.013516861945390701, 0.002940082224085927, 0.002004653448238969, 0.006652397103607655, 0.004079414997249842, 0.0028307989705353975, 0.0006369714974425733, 0.002542868722230196, 0.1463778167963028], [0.0762338638305664, 0.11778479814529419, 0.03105221875011921, 0.006415408570319414, 0.0190818402916193, 0.027191398665308952, 0.005222225561738014, 0.0170834269374609, 0.05309534817934036, 0.00936796236783266, 0.03816217556595802, 0.17940494418144226, 0.020440110936760902, 0.13513173162937164, 0.3000544309616089], [0.16228125989437103, 0.35454851388931274, 0.04026315361261368, 0.03822629526257515, 0.023396998643875122, 0.30800631642341614, 0.24136781692504883, 0.15176478028297424, 0.0788438618183136, 0.07347536832094193, 0.030298085883259773, 0.007365733850747347, 0.1061745211482048, 0.2841038405895233, 0.07787416130304337], [0.05645078793168068, 0.023840615525841713, 0.013567867688834667, 0.00750470208004117, 0.07643276453018188, 0.08809614926576614, 0.06102507561445236, 0.021034346893429756, 0.039108242839574814, 0.02081543207168579, 0.011458326131105423, 0.20520520210266113, 0.027348484843969345, 0.06299317628145218, 0.2514360249042511], [0.016126127913594246, 0.01087501272559166, 0.01213990617543459, 0.004450921434909105, 0.014690833166241646, 0.30525338649749756, 0.02716207131743431, 0.09981174021959305, 0.027048761025071144, 0.01336466334760189, 0.006663064938038588, 0.0520603246986866, 0.042623523622751236, 0.018071996048092842, 0.1948687732219696], [0.04185086488723755, 0.034399643540382385, 0.041276611387729645, 0.0584070086479187, 0.019824109971523285, 0.00856409315019846, 0.08867836743593216, 0.10337970405817032, 0.09468665719032288, 0.02033121883869171, 0.018058426678180695, 0.059728462249040604, 0.09321711957454681, 0.20168805122375488, 0.1941128522157669], [0.01436887588351965, 0.027922889217734337, 0.046481672674417496, 0.010071231983602047, 0.026127830147743225, 0.06003356724977493, 0.022118212655186653, 0.08160483092069626, 0.07784195244312286, 0.010694753378629684, 0.017130734398961067, 0.05340806022286415, 0.041410259902477264, 0.035884104669094086, 0.2491855025291443], [0.053393200039863586, 0.04828185588121414, 0.03453819081187248, 0.013636122457683086, 0.25098806619644165, 0.12313847243785858, 0.02266266942024231, 0.017618268728256226, 0.019785437732934952, 0.005274764262139797, 0.021053072065114975, 0.20679616928100586, 0.021523641422390938, 0.03855947405099869, 0.1109846979379654], [0.12851715087890625, 0.12400124222040176, 0.2637093663215637, 0.02439347468316555, 0.07038086652755737, 0.12665364146232605, 0.04898465424776077, 0.03412041813135147, 0.0263816025108099, 0.023226425051689148, 0.11513664573431015, 0.09503531455993652, 0.1215861439704895, 0.11158601939678192, 0.14799171686172485], [0.0010214513167738914, 0.004835289902985096, 0.0042709591798484325, 0.0026378841139376163, 0.005866974592208862, 0.008331544697284698, 0.006240549497306347, 0.01365274004638195, 0.1720106601715088, 0.0005307683604769409, 0.0007543729152530432, 0.004353509750217199, 0.0002490385086275637, 0.0017186965560540557, 0.14317919313907623], [0.07205050438642502, 0.12816517055034637, 0.23753608763217926, 0.08243206143379211, 0.5041552186012268, 0.11970840394496918, 0.04837331175804138, 0.034129947423934937, 0.16484025120735168, 0.011070297099649906, 0.05054215341806412, 0.039082955569028854, 0.09205758571624756, 0.1322212517261505, 0.16203875839710236], [0.014979850500822067, 0.03769220784306526, 0.04367470741271973, 0.009415187872946262, 0.019922776147723198, 0.11522040516138077, 0.014906312339007854, 0.04722318425774574, 0.06570684164762497, 0.008925273083150387, 0.019600573927164078, 0.0472339391708374, 0.005348374601453543, 0.0017698986921459436, 0.1612817794084549], [0.023198002949357033, 0.06148262694478035, 0.046858664602041245, 0.013079512864351273, 0.08762317895889282, 0.00949429627507925, 0.0484880767762661, 0.025388503447175026, 0.04432932287454605, 0.006038118619471788, 0.010164186358451843, 0.08949221670627594, 0.06122652441263199, 0.11895263940095901, 0.16355113685131073], [0.009917332790791988, 0.01408212911337614, 0.047434139996767044, 0.005388779100030661, 0.023170381784439087, 0.034844160079956055, 0.009820640087127686, 0.03569778800010681, 0.05789060518145561, 0.0037882563192397356, 0.013808010146021843, 0.04879388585686684, 0.03114072047173977, 0.0507131889462471, 0.18661679327487946], [0.0652787834405899, 0.04612350836396217, 0.04522763565182686, 0.014745297841727734, 0.27657532691955566, 0.16156227886676788, 0.025164838880300522, 0.017732013016939163, 0.023105354979634285, 0.005499221384525299, 0.020183373242616653, 0.19132839143276215, 0.020515967160463333, 0.056384406983852386, 0.14304831624031067], [0.14539514482021332, 0.21388974785804749, 0.34906452894210815, 0.031415559351444244, 0.062017399817705154, 0.08485611528158188, 0.03913363441824913, 0.03569692373275757, 0.023448940366506577, 0.020669998601078987, 0.1622902750968933, 0.1315622329711914, 0.09182734042406082, 0.1796703040599823, 0.13702963292598724], [0.0009059146977961063, 0.004442692268639803, 0.002850044285878539, 0.0024173678830266, 0.006019651889801025, 0.004450949374586344, 0.003768310882151127, 0.009272964671254158, 0.19643637537956238, 0.0004391498805489391, 0.0004852984275203198, 0.005083973053842783, 0.000164541692356579, 0.001456208759918809, 0.13767127692699432], [0.03601038455963135, 0.08602340519428253, 0.042799800634384155, 0.007577326148748398, 0.12637566030025482, 0.07399067282676697, 0.02205651067197323, 0.01475659292191267, 0.14170114696025848, 0.004405674524605274, 0.013175459578633308, 0.03142356127500534, 0.06839168816804886, 0.09161193668842316, 0.1376270353794098], [0.014056011103093624, 0.020953036844730377, 0.03237491473555565, 0.0042424313724040985, 0.017438247799873352, 0.08849667757749557, 0.005714876111596823, 0.025588830932974815, 0.08735965192317963, 0.009712125174701214, 0.02371004782617092, 0.06271149963140488, 0.00425978796556592, 0.0027238703332841396, 0.14272134006023407], [0.15719948709011078, 0.03286461904644966, 0.12916648387908936, 0.10299614071846008, 0.014032969251275063, 0.011700707487761974, 0.06680437922477722, 0.016068298369646072, 0.04505150765180588, 0.056866806000471115, 0.07287567108869553, 0.09101171046495438, 0.06734755635261536, 0.17371943593025208, 0.1297563910484314]], [[0.010018138214945793, 0.02516627125442028, 0.027397310361266136, 0.005101055838167667, 0.025938771665096283, 0.13529063761234283, 0.02690303698182106, 0.11719205975532532, 0.027814749628305435, 0.019565219059586525, 0.07996311038732529, 0.0991574078798294, 0.16288702189922333, 0.1113416850566864, 0.22370746731758118], [0.05219842493534088, 0.1440066546201706, 0.27922260761260986, 0.2058621197938919, 0.11230742931365967, 0.6016822457313538, 0.20846855640411377, 0.04777589067816734, 0.20611444115638733, 0.15481434762477875, 0.11950203776359558, 0.02679699845612049, 0.0639302060008049, 0.047183193266391754, 0.04897741973400116], [0.01555164996534586, 0.0014379153726622462, 0.01706753298640251, 0.003720618085935712, 0.10093016922473907, 0.027928827330470085, 0.015380543656647205, 0.0025812943931668997, 0.020822137594223022, 0.014309070073068142, 0.017923271283507347, 0.0120958611369133, 0.014481468126177788, 0.009491728618741035, 0.15904544293880463], [0.11612647771835327, 0.0010205605067312717, 0.020188286900520325, 0.027076182886958122, 0.09822120517492294, 0.3221674859523773, 0.1250218003988266, 0.002691123867407441, 0.005359187722206116, 0.04976291581988335, 0.023232540115714073, 0.04237976670265198, 0.028708819299936295, 0.049411751329898834, 0.005618311930447817], [0.0470837838947773, 0.007497857324779034, 0.004583081230521202, 0.022991856560111046, 0.0278051495552063, 0.00051211251411587, 0.0627230703830719, 0.011764267459511757, 0.010903585702180862, 0.07272983342409134, 0.011678352952003479, 0.09392477571964264, 0.01558940764516592, 0.03351595252752304, 0.2068868726491928], [0.0024584962520748377, 8.163625898305327e-05, 0.00016154914919752628, 0.0002508168399799615, 0.0019916424062103033, 0.0004536219348665327, 0.0036078437697142363, 0.0008641426684334874, 0.00021941671730019152, 0.0014423344982787967, 0.0004360634775366634, 0.004383172374218702, 0.0009428760386072099, 0.0009436326217837632, 0.14683274924755096], [0.02989446185529232, 0.007703323382884264, 0.12996061146259308, 0.025068828836083412, 0.2812304198741913, 0.0071953474543988705, 0.0021352169569581747, 0.0025125211104750633, 0.0014658492291346192, 0.007028855849057436, 0.0448734275996685, 0.09462164342403412, 0.0503704659640789, 0.11768583953380585, 0.12974096834659576], [0.16756094992160797, 0.028098214417696, 0.20756086707115173, 0.2207580953836441, 0.10928753018379211, 0.13773545622825623, 0.2233184576034546, 0.1774815022945404, 0.13830231130123138, 0.20932619273662567, 0.18267595767974854, 0.05961548537015915, 0.07697918266057968, 0.18739080429077148, 0.06796090304851532], [0.017068415880203247, 0.00098085415083915, 0.010854640044271946, 0.006490680854767561, 0.29060667753219604, 0.006710599176585674, 0.0118483304977417, 0.0008181483135558665, 0.00011296885350020602, 0.0034601599909365177, 0.005098147317767143, 0.010750477202236652, 0.010399019345641136, 0.009376241825520992, 0.017405353486537933], [0.1331326961517334, 0.019769106060266495, 0.01612294837832451, 0.028521019965410233, 0.007509702816605568, 0.2665199935436249, 0.19958320260047913, 0.1385747790336609, 0.0059373765252530575, 0.08046255260705948, 0.052418529987335205, 0.004961848258972168, 0.10941796749830246, 0.06705309450626373, 0.17611992359161377], [0.019668979570269585, 0.0081618782132864, 0.12552350759506226, 0.0802406370639801, 0.07089362293481827, 0.18871739506721497, 0.12778939306735992, 0.04829992726445198, 0.04307088255882263, 0.02314154990017414, 0.14194107055664062, 0.05861861631274223, 0.19650596380233765, 0.11930099874734879, 0.18420156836509705], [0.00538466265425086, 0.0270208939909935, 0.18066750466823578, 0.06076826527714729, 0.035171061754226685, 0.411039799451828, 0.09634009003639221, 0.26394954323768616, 0.1915867179632187, 0.03318370133638382, 0.3213040828704834, 0.10995125770568848, 0.5320225954055786, 0.4394112527370453, 0.15243512392044067], [0.0030147582292556763, 0.00625306461006403, 0.017102748155593872, 0.008551767095923424, 0.0727200135588646, 0.015153692103922367, 0.0023096217773854733, 0.011201570741832256, 0.002435098635032773, 0.006847116630524397, 0.016829995438456535, 0.12519565224647522, 0.3878204822540283, 0.13249750435352325, 0.028183329850435257], [0.066617950797081, 0.006649812217801809, 0.04142908379435539, 0.13957993686199188, 0.025706114247441292, 0.08231058716773987, 0.08377126604318619, 0.02330365777015686, 0.04652002453804016, 0.11060080677270889, 0.09014575183391571, 0.07117310166358948, 0.15938407182693481, 0.1624550223350525, 0.05356656014919281], [0.004379222169518471, 0.0002637936850078404, 0.0022587613202631474, 0.006711117923259735, 0.0006837267428636551, 0.007989797741174698, 0.02997850626707077, 0.045127563178539276, 0.008224103599786758, 0.0034686585422605276, 0.0038658890407532454, 0.00034815416438505054, 7.646608719369397e-05, 0.00017854337056633085, 0.14325816929340363], [0.25216665863990784, 0.1422366499900818, 0.10172943770885468, 0.3735504150390625, 0.0612066313624382, 0.06238102167844772, 0.11154207587242126, 0.031159698963165283, 0.011768986470997334, 0.4107469618320465, 0.1557808816432953, 0.07179611176252365, 0.186580628156662, 0.18789765238761902, 0.099563829600811], [0.0073658498004078865, 0.1486257165670395, 0.03456511348485947, 0.0081891855224967, 0.009660922922194004, 0.09341325610876083, 0.010183881968259811, 0.09390538185834885, 0.005950886756181717, 0.019719628617167473, 0.060451164841651917, 0.021925343200564384, 0.19991156458854675, 0.17004182934761047, 0.15761280059814453], [0.0057948376052081585, 0.023180164396762848, 0.018019115552306175, 0.008233858272433281, 0.005580522585660219, 0.09526203572750092, 0.025384269654750824, 0.05396068096160889, 0.022398412227630615, 0.010895788669586182, 0.02884012460708618, 0.008390026167035103, 0.1754663735628128, 0.0998048186302185, 0.1692073941230774], [0.0038264640606939793, 0.023839879781007767, 0.12264026701450348, 0.02543032169342041, 0.01467527449131012, 0.22457416355609894, 0.02885078825056553, 0.18430863320827484, 0.08557040989398956, 0.016987022012472153, 0.3513573110103607, 0.04023189842700958, 0.40384334325790405, 0.4235673248767853, 0.16652488708496094], [0.006266402080655098, 0.015031179413199425, 0.02853887900710106, 0.010518345981836319, 0.09044987708330154, 0.021657679229974747, 0.0031435268465429544, 0.020945381373167038, 0.004824943374842405, 0.0127853499725461, 0.04820985347032547, 0.12459135800600052, 0.5573670268058777, 0.2566193640232086, 0.05160163715481758], [0.3002758324146271, 0.08866846561431885, 0.06544900685548782, 0.25531354546546936, 0.028160221874713898, 0.12210531532764435, 0.16810676455497742, 0.0764283761382103, 0.17981933057308197, 0.3050864636898041, 0.2806880474090576, 0.13050490617752075, 0.19047558307647705, 0.3216065764427185, 0.07704814523458481], [0.005926316604018211, 0.0003559965989552438, 0.0015365411527454853, 0.005924532189965248, 0.0005743101937696338, 0.007415232714265585, 0.024156678467988968, 0.045611582696437836, 0.009969166480004787, 0.003380746114999056, 0.003106702584773302, 0.0003880919248331338, 4.0538176108384505e-05, 0.00014580521383322775, 0.13770556449890137], [0.1617586314678192, 0.29556339979171753, 0.028325924649834633, 0.059843577444553375, 0.009868957102298737, 0.03965649753808975, 0.07811643928289413, 0.06809397041797638, 0.009963614866137505, 0.11740529537200928, 0.08369920402765274, 0.039758261293172836, 0.13982373476028442, 0.1197674348950386, 0.13220268487930298], [0.012153265066444874, 0.16048333048820496, 0.041802890598773956, 0.00796045083552599, 0.018259191885590553, 0.10963782668113708, 0.009757153689861298, 0.07023902982473373, 0.01128031499683857, 0.030125515535473824, 0.0943576917052269, 0.02206866256892681, 0.1321137398481369, 0.19507774710655212, 0.1400403380393982], [0.005033975467085838, 0.01824766956269741, 0.015512547455728054, 0.006673634983599186, 0.005676268134266138, 0.04240407794713974, 0.023996027186512947, 0.1038113459944725, 0.02023463323712349, 0.0080516142770648, 0.052543867379426956, 0.1188565045595169, 0.05977800861001015, 0.05786403268575668, 0.13343320786952972]], [[0.1022859737277031, 0.17571765184402466, 0.1416551172733307, 0.11749783158302307, 0.09062699973583221, 0.07838433235883713, 0.09344526380300522, 0.3238999545574188, 0.11371968686580658, 0.10100032389163971, 0.09302259236574173, 0.0389624647796154, 0.16697892546653748, 0.1419355273246765, 0.1285012662410736], [0.24028724431991577, 0.14351274073123932, 0.051798444241285324, 0.16382630169391632, 0.04226303845643997, 0.020662518218159676, 0.11527843773365021, 0.29321926832199097, 0.02218940667808056, 0.0878078043460846, 0.10535410046577454, 0.011972848325967789, 0.07032275199890137, 0.04715458303689957, 0.0739566907286644], [0.2799055874347687, 0.11053244769573212, 0.1936434954404831, 0.029654914513230324, 0.3583168685436249, 0.552708625793457, 0.34459343552589417, 0.33612802624702454, 0.17023301124572754, 0.19969996809959412, 0.18768110871315002, 0.6793866157531738, 0.791401207447052, 0.7463385462760925, 0.09094473719596863], [0.1572730988264084, 0.12077052146196365, 0.0489557608962059, 0.1575693041086197, 0.05669395253062248, 0.21311312913894653, 0.07387427985668182, 0.12006285786628723, 0.06427917629480362, 0.05486075580120087, 0.09722346067428589, 0.0672946497797966, 0.519307017326355, 0.15919242799282074, 0.07895061373710632], [0.056666091084480286, 0.13304737210273743, 0.023897293955087662, 0.04679059237241745, 0.045941345393657684, 0.32384783029556274, 0.44531556963920593, 0.533463716506958, 0.08588721603155136, 0.10118058323860168, 0.027683693915605545, 0.15270595252513885, 0.45412689447402954, 0.19033603370189667, 0.009601723402738571], [0.026866083964705467, 0.01856745034456253, 0.00889106560498476, 0.023431263864040375, 0.014423922635614872, 0.06721587479114532, 0.30465173721313477, 0.5084072351455688, 0.06748852878808975, 0.09416066110134125, 0.028160765767097473, 0.08301042765378952, 0.13479003310203552, 0.08470122516155243, 0.14269311726093292], [0.07283831387758255, 0.02513016201555729, 0.513066828250885, 0.1692790985107422, 0.12089971452951431, 0.05420007184147835, 0.019427694380283356, 0.038392528891563416, 0.31973040103912354, 0.29048243165016174, 0.4046151340007782, 0.10607112944126129, 0.0885496586561203, 0.07017665356397629, 0.1372956782579422], [0.27857187390327454, 0.3617483973503113, 0.2938012182712555, 0.22770966589450836, 0.06824903935194016, 0.055705904960632324, 0.2735913395881653, 0.10727421194314957, 0.15245027840137482, 0.12983311712741852, 0.2781352400779724, 0.010307536460459232, 0.09433942288160324, 0.07780664414167404, 0.13000918924808502], [0.09918209165334702, 0.053455647081136703, 0.645177960395813, 0.40746453404426575, 0.08205579966306686, 0.11053493618965149, 0.09200509637594223, 0.0519426129758358, 0.15867555141448975, 0.14363400638103485, 0.08945868164300919, 0.009240956045687199, 0.05626320466399193, 0.024817338213324547, 0.10628006607294083], [0.21029417216777802, 0.16975507140159607, 0.4791514277458191, 0.5080997347831726, 0.14877668023109436, 0.04306463524699211, 0.02225780300796032, 0.027854960411787033, 0.09907854348421097, 0.17716829478740692, 0.027767561376094818, 0.04010230675339699, 0.1045137569308281, 0.07445494085550308, 0.1349247545003891], [0.05318222567439079, 0.11344952136278152, 0.09562063962221146, 0.10165436565876007, 0.11442670226097107, 0.07387696951627731, 0.04448265954852104, 0.12469986081123352, 0.10296554863452911, 0.029610879719257355, 0.006854650564491749, 0.06481806933879852, 0.038151390850543976, 0.029200172051787376, 0.19021393358707428], [0.024841444566845894, 0.16249340772628784, 0.20643305778503418, 0.09402812272310257, 0.0850510448217392, 0.023708872497081757, 0.027868179604411125, 0.16653721034526825, 0.2575382590293884, 0.07176022976636887, 0.04638299718499184, 0.019721999764442444, 0.08340867608785629, 0.04306621477007866, 0.19255293905735016], [0.24242781102657318, 0.4547469913959503, 0.7904132008552551, 0.7443370819091797, 0.4808639585971832, 0.2640213668346405, 0.06001711264252663, 0.24681034684181213, 0.5675581097602844, 0.2725449204444885, 0.247804656624794, 0.029579274356365204, 0.19247104227542877, 0.09198179841041565, 0.18542104959487915], [0.10456986725330353, 0.23679938912391663, 0.29603201150894165, 0.2020668387413025, 0.14429134130477905, 0.4285147190093994, 0.3221139907836914, 0.592944860458374, 0.47945162653923035, 0.273953914642334, 0.2270997315645218, 0.05125115066766739, 0.15167200565338135, 0.14498752355575562, 0.03565559163689613], [0.005393329542130232, 0.004602347034960985, 0.02125353366136551, 0.017772456631064415, 0.029431374743580818, 0.06670433282852173, 0.07382840663194656, 0.05640842020511627, 0.2022721767425537, 0.02110537886619568, 0.006757265422493219, 0.0065305884927511215, 0.00012849831546191126, 0.0015581984771415591, 0.14312443137168884], [0.03693488612771034, 0.3099628686904907, 0.02452116832137108, 0.038606833666563034, 0.04603191837668419, 0.056979674845933914, 0.014461892656981945, 0.021202413365244865, 0.4372372031211853, 0.02073492854833603, 0.005594322457909584, 0.11605570465326309, 0.05724794790148735, 0.01605997234582901, 0.1753198802471161], [0.17487157881259918, 0.2829012870788574, 0.22657853364944458, 0.2227388322353363, 0.09278897941112518, 0.05522100254893303, 0.023270972073078156, 0.031554628163576126, 0.32194823026657104, 0.13948096334934235, 0.09803083539009094, 0.2809208631515503, 0.14969345927238464, 0.03018103539943695, 0.10283161699771881], [0.06711219251155853, 0.13971862196922302, 0.10573939234018326, 0.08062157034873962, 0.22173365950584412, 0.04757346957921982, 0.02002648264169693, 0.06195787340402603, 0.09553409367799759, 0.04351034387946129, 0.015184497460722923, 0.17841440439224243, 0.07658158242702484, 0.04646967723965645, 0.1461518555879593], [0.015694430097937584, 0.09081663191318512, 0.2731003761291504, 0.09780610352754593, 0.06437630951404572, 0.024092676118016243, 0.017730340361595154, 0.09997125715017319, 0.24317535758018494, 0.06615940481424332, 0.05322461575269699, 0.013002216815948486, 0.10308460891246796, 0.03947872668504715, 0.16966252028942108], [0.19514591991901398, 0.2590837776660919, 0.7111572027206421, 0.6245842576026917, 0.2279123067855835, 0.21324849128723145, 0.0465325303375721, 0.16129039227962494, 0.5552195906639099, 0.24888396263122559, 0.16995932161808014, 0.017819084227085114, 0.13601525127887726, 0.04923256114125252, 0.1924036145210266], [0.11466818302869797, 0.23749157786369324, 0.22078867256641388, 0.21260471642017365, 0.1054922342300415, 0.38443663716316223, 0.35735341906547546, 0.3432110548019409, 0.45766645669937134, 0.30316272377967834, 0.15794025361537933, 0.23222389817237854, 0.18522031605243683, 0.12369272857904434, 0.062224190682172775], [0.004928229842334986, 0.004764902405440807, 0.014567935839295387, 0.014073353260755539, 0.020878629758954048, 0.04901519790291786, 0.05124438554048538, 0.042454566806554794, 0.19801755249500275, 0.018003307282924652, 0.004736864008009434, 0.006620202213525772, 0.00011398878996260464, 0.001381832524202764, 0.13761556148529053], [0.013776288367807865, 0.25124475359916687, 0.00789756141602993, 0.00910337083041668, 0.005072988104075193, 0.015830766409635544, 0.005818341393023729, 0.011153762228786945, 0.14152461290359497, 0.008211367763578892, 0.002360414480790496, 0.06666377186775208, 0.057822320610284805, 0.009000283665955067, 0.13980405032634735], [0.25532495975494385, 0.3110601603984833, 0.28066542744636536, 0.29941898584365845, 0.09561395645141602, 0.06004221364855766, 0.0257351566106081, 0.04446575790643692, 0.3475395441055298, 0.2538500130176544, 0.25107017159461975, 0.4736424386501312, 0.29699820280075073, 0.06975124776363373, 0.11745814979076385], [0.06876020133495331, 0.07319146394729614, 0.08357107639312744, 0.06905727088451385, 0.010884120129048824, 0.012632370926439762, 0.04344229772686958, 0.06033884361386299, 0.05559740215539932, 0.048808641731739044, 0.06204793229699135, 0.017201891168951988, 0.028970519080758095, 0.021960163488984108, 0.13179059326648712]], [[0.1855485588312149, 0.4779467284679413, 0.0886944904923439, 0.027812138199806213, 0.051930978894233704, 0.20570456981658936, 0.13285183906555176, 0.12479114532470703, 0.03275279700756073, 0.13280591368675232, 0.10831113904714584, 0.13358037173748016, 0.31709861755371094, 0.18639257550239563, 0.0658930093050003], [0.04738391190767288, 0.17884546518325806, 0.030679181218147278, 0.09374479204416275, 0.015219364315271378, 0.004209337756037712, 0.011544613167643547, 0.014519347809255123, 0.0008998611010611057, 0.03714418038725853, 0.02808041125535965, 0.0015275280456990004, 0.014074422419071198, 0.01773718185722828, 0.02865048497915268], [0.4282352328300476, 0.07421883940696716, 0.37614062428474426, 0.6016114950180054, 0.16448479890823364, 0.10949403792619705, 0.43647968769073486, 0.17394804954528809, 0.2346193641424179, 0.5131813287734985, 0.6543169021606445, 0.06318124383687973, 0.059741634875535965, 0.08049911260604858, 0.08155221492052078], [0.04248558357357979, 0.005498564336448908, 0.015051363967359066, 0.021896474063396454, 0.031015703454613686, 0.23631463944911957, 0.5231030583381653, 0.1651564985513687, 0.010708797723054886, 0.0702022984623909, 0.015817642211914062, 0.01968570239841938, 0.2309122085571289, 0.11954572051763535, 0.04909561946988106], [0.019823409616947174, 0.02119731903076172, 0.0447932668030262, 0.04950243979692459, 0.11350910365581512, 0.3172611892223358, 0.1175147220492363, 0.16474604606628418, 0.025614900514483452, 0.11684545129537582, 0.027774598449468613, 0.03366768732666969, 0.1657668650150299, 0.20241110026836395, 0.02058284729719162], [0.024027986451983452, 0.07085671275854111, 0.014559593982994556, 0.003951122052967548, 0.5812088251113892, 0.07389754801988602, 0.10464153438806534, 0.06822511553764343, 0.1849648803472519, 0.02429678477346897, 0.014226456172764301, 0.2123226672410965, 0.1049809455871582, 0.17609325051307678, 0.13661964237689972], [0.20496347546577454, 0.09403666108846664, 0.02112487144768238, 0.025338320061564445, 0.008130905218422413, 0.1783977895975113, 0.3754851818084717, 0.0950397253036499, 0.0030220954213291407, 0.08205359429121017, 0.011042395606637001, 0.018588367849588394, 0.1888807862997055, 0.10302136838436127, 0.14473272860050201], [0.037373751401901245, 0.07382072508335114, 0.08205787092447281, 0.10832883417606354, 0.02859049290418625, 0.1663966327905655, 0.058918725699186325, 0.17053310573101044, 0.011018002405762672, 0.15213745832443237, 0.027154715731739998, 0.0019660431426018476, 0.22162862122058868, 0.11411792784929276, 0.08493959158658981], [0.015705576166510582, 0.016172299161553383, 0.006149389781057835, 0.0038101596292108297, 0.007736767642199993, 0.20371977984905243, 0.12438680231571198, 0.06649734079837799, 0.004926482681185007, 0.004153827205300331, 0.0012289183214306831, 0.003863752353936434, 0.0550994910299778, 0.04052891582250595, 0.36571574211120605], [0.008730506524443626, 0.002757954876869917, 0.0122150257229805, 0.006305738352239132, 0.004681416787207127, 0.06460410356521606, 0.008150112815201283, 0.010960009880363941, 0.004299533553421497, 0.004670997615903616, 0.0034528695978224277, 0.0024545302148908377, 0.005013267509639263, 0.008545692078769207, 0.23703089356422424], [0.09499987959861755, 0.010673395358026028, 0.007046178914606571, 0.020993953570723534, 0.010670008137822151, 0.07466354966163635, 0.06417079269886017, 0.023990478366613388, 0.17728924751281738, 0.15624059736728668, 0.004560643341392279, 0.010690598748624325, 0.03727814555168152, 0.017693333327770233, 0.14084658026695251], [0.688500165939331, 0.16286028921604156, 0.04583478718996048, 0.22473743557929993, 0.025797681882977486, 0.04771623760461807, 0.5437547564506531, 0.0642164871096611, 0.01443459838628769, 0.2519066631793976, 0.017869845032691956, 0.003991205245256424, 0.04630482196807861, 0.029587149620056152, 0.049375567585229874], [0.14772717654705048, 0.11627800017595291, 0.034884992986917496, 0.02596234902739525, 0.031621210277080536, 0.39286479353904724, 0.6627658009529114, 0.20747745037078857, 0.019052494317293167, 0.06071586161851883, 0.014515946619212627, 0.03545556217432022, 0.1622975915670395, 0.05619712546467781, 0.4560142755508423], [0.3253695070743561, 0.18678773939609528, 0.23196454346179962, 0.43925735354423523, 0.09974130243062973, 0.1577768325805664, 0.26045241951942444, 0.07323815673589706, 0.005399893503636122, 0.23951157927513123, 0.04431937262415886, 0.013187061063945293, 0.0749824121594429, 0.025474021211266518, 0.2768867611885071], [0.049311667680740356, 0.10222040861845016, 0.30249276757240295, 0.11109475791454315, 0.4333159327507019, 0.4476950168609619, 0.14919614791870117, 0.45436185598373413, 0.10977044701576233, 0.101465605199337, 0.28612539172172546, 0.15904487669467926, 0.4858849048614502, 0.19411928951740265, 0.08273273706436157], [0.08865676820278168, 0.0832996591925621, 0.0360012948513031, 0.026901112869381905, 0.0488949753344059, 0.5697077512741089, 0.2118675261735916, 0.21166029572486877, 0.009457184933125973, 0.042189937084913254, 0.010147118009626865, 0.027016732841730118, 0.1966082751750946, 0.18848717212677002, 0.17412608861923218], [0.09455566853284836, 0.047932155430316925, 0.06032469496130943, 0.027359262108802795, 0.004525639116764069, 0.19231697916984558, 0.29536089301109314, 0.10446369647979736, 0.004957688972353935, 0.22148354351520538, 0.017980555072426796, 0.016062501817941666, 0.01227590162307024, 0.007468203082680702, 0.14047065377235413], [0.18475790321826935, 0.03305341675877571, 0.022945405915379524, 0.02499788999557495, 0.016275716945528984, 0.44049808382987976, 0.3255404233932495, 0.03656867519021034, 0.008760510943830013, 0.28132569789886475, 0.00872495025396347, 0.02103549800813198, 0.09103824943304062, 0.045535117387771606, 0.1431308537721634], [0.5226730704307556, 0.08511564135551453, 0.13128292560577393, 0.22977954149246216, 0.025636736303567886, 0.14430683851242065, 0.697600245475769, 0.08303582668304443, 0.03326253592967987, 0.30183717608451843, 0.04944504052400589, 0.004384536296129227, 0.07144975662231445, 0.05258011445403099, 0.06879302859306335], [0.06703877449035645, 0.049393996596336365, 0.041539933532476425, 0.021373772993683815, 0.02868128940463066, 0.32991066575050354, 0.488584041595459, 0.0702073872089386, 0.0075523643754422665, 0.038572411984205246, 0.012813442386686802, 0.04136957228183746, 0.06929102540016174, 0.03757195174694061, 0.23515936732292175], [0.15618596971035004, 0.12941822409629822, 0.2654253840446472, 0.28590527176856995, 0.31243884563446045, 0.1085575670003891, 0.15852880477905273, 0.026613548398017883, 0.004155577160418034, 0.15324708819389343, 0.037679530680179596, 0.09416285902261734, 0.02134908176958561, 0.010629331693053246, 0.17846201360225677], [0.058257974684238434, 0.12017454952001572, 0.32657214999198914, 0.12284700572490692, 0.5568311810493469, 0.41536086797714233, 0.16300946474075317, 0.49100223183631897, 0.15462136268615723, 0.11520260572433472, 0.260068416595459, 0.28476831316947937, 0.501883327960968, 0.21151991188526154, 0.09330709278583527], [0.04007576033473015, 0.04011448100209236, 0.02015572600066662, 0.006723308004438877, 0.01584162376821041, 0.6745935082435608, 0.14270515739917755, 0.05812964215874672, 0.0018657244509086013, 0.018765496090054512, 0.004551106132566929, 0.05217724293470383, 0.21886952221393585, 0.13090433180332184, 0.13149680197238922], [0.051524627953767776, 0.037071868777275085, 0.09267362952232361, 0.03285788744688034, 0.006808253470808268, 0.2584725618362427, 0.21142001450061798, 0.06556515395641327, 0.003410812932997942, 0.18829914927482605, 0.028329605236649513, 0.02864006720483303, 0.014232979156076908, 0.014326054602861404, 0.12804241478443146], [0.13503411412239075, 0.06798373907804489, 0.08072269707918167, 0.04104887321591377, 0.027653640136122704, 0.5933560132980347, 0.15723249316215515, 0.044575583189725876, 0.017590617761015892, 0.04771400988101959, 0.07117579132318497, 0.10345834493637085, 0.10624422132968903, 0.027206260710954666, 0.1271171271800995]], [[0.04247138649225235, 0.01728098653256893, 0.06617120653390884, 0.009399485774338245, 0.0730140432715416, 0.14221039414405823, 0.11889991164207458, 0.10651882737874985, 0.10687308758497238, 0.0351867638528347, 0.09164245426654816, 0.06160420924425125, 0.04699656739830971, 0.14884592592716217, 0.20088525116443634], [0.35919252038002014, 0.017007382586598396, 0.3711448311805725, 0.05260182172060013, 0.23237934708595276, 0.17189942300319672, 0.06846722215414047, 0.25480321049690247, 0.4269619286060333, 0.141769677400589, 0.19745108485221863, 0.3101239502429962, 0.12419883906841278, 0.061588384211063385, 0.3489930033683777], [0.1570073962211609, 0.6818748116493225, 0.08056136965751648, 0.04282544180750847, 0.09609510749578476, 0.21831035614013672, 0.11452964693307877, 0.4344905614852905, 0.09872471541166306, 0.06769980490207672, 0.054214250296354294, 0.015440859831869602, 0.04572026804089546, 0.05267196521162987, 0.06955287605524063], [0.1362180858850479, 0.01786869764328003, 0.3548091650009155, 0.13650378584861755, 0.07479218393564224, 0.08773932605981827, 0.007214170414954424, 0.020996512845158577, 0.09793394804000854, 0.26323461532592773, 0.31718939542770386, 0.004400049336254597, 0.01118874829262495, 0.016452480107545853, 0.0059462906792759895], [0.13787487149238586, 0.02221597172319889, 0.46063661575317383, 0.42787930369377136, 0.16819633543491364, 0.30927538871765137, 0.10940644890069962, 0.14741046726703644, 0.3708270192146301, 0.08424455672502518, 0.34931957721710205, 0.015041538514196873, 0.02219252847135067, 0.0637117251753807, 0.001682900357991457], [0.09526984393596649, 0.013222168199717999, 0.9035038352012634, 0.8715099692344666, 0.20107677578926086, 0.7829492688179016, 0.28305909037590027, 0.141366645693779, 0.15355023741722107, 0.11376345157623291, 0.804192841053009, 0.012117957696318626, 0.3312073349952698, 0.4514775276184082, 0.016239164397120476], [0.34537556767463684, 0.010514522902667522, 0.04824088513851166, 0.12771852314472198, 0.005308120045810938, 0.17857761681079865, 0.2263273000717163, 0.26537755131721497, 0.3297313451766968, 0.3104889690876007, 0.11654951423406601, 0.08535956591367722, 0.02363554947078228, 0.031254567205905914, 0.10634612292051315], [0.2808375656604767, 0.07436379790306091, 0.11235158890485764, 0.07017786800861359, 0.034851111471652985, 0.01653558947145939, 0.025893066078424454, 0.02911091037094593, 0.23654304444789886, 0.2646749019622803, 0.20617236196994781, 0.25081631541252136, 0.013157923705875874, 0.04621773213148117, 0.2354249358177185], [0.5487799644470215, 0.03728892654180527, 0.05227963626384735, 0.18957917392253876, 0.014632479287683964, 0.19499987363815308, 0.29326584935188293, 0.6778355836868286, 0.45779454708099365, 0.33408117294311523, 0.11356081813573837, 0.01941866986453533, 0.010207045823335648, 0.013884961605072021, 0.09069465100765228], [0.09531711786985397, 0.03595840558409691, 0.017401238903403282, 0.061305541545152664, 0.1627957820892334, 0.050434935837984085, 0.05516263470053673, 0.23917846381664276, 0.3637218177318573, 0.09729932248592377, 0.03891580551862717, 0.19205324351787567, 0.041229162365198135, 0.046046942472457886, 0.03756402060389519], [0.08811857551336288, 0.010963470675051212, 0.2593647241592407, 0.26678594946861267, 0.42746680974960327, 0.41530901193618774, 0.07491520792245865, 0.18910719454288483, 0.04928334057331085, 0.04599721357226372, 0.4843277335166931, 0.07717985659837723, 0.09353034198284149, 0.07800954580307007, 0.08156391978263855], [0.04596662148833275, 0.005170373246073723, 0.12165658175945282, 0.15079215168952942, 0.04554709792137146, 0.08856093138456345, 0.04626012593507767, 0.020681705325841904, 0.17637456953525543, 0.26189061999320984, 0.13335715234279633, 0.046832337975502014, 0.018430203199386597, 0.01621258072555065, 0.10917440801858902], [0.5138411521911621, 0.0654044821858406, 0.1128465011715889, 0.18054738640785217, 0.038166921585798264, 0.13531430065631866, 0.12295213341712952, 0.28065726161003113, 0.2875981628894806, 0.5909985899925232, 0.601227879524231, 0.03077608533203602, 0.04096299037337303, 0.09236451238393784, 0.1495288461446762], [0.07072688639163971, 0.012152088806033134, 0.021357353776693344, 0.04663744568824768, 0.020319821313023567, 0.05489102751016617, 0.07223928719758987, 0.23148301243782043, 0.18188072741031647, 0.10590049624443054, 0.10450157523155212, 0.03876996785402298, 0.13536545634269714, 0.10362161695957184, 0.12556865811347961], [0.07390952110290527, 0.023819932714104652, 0.4992673993110657, 0.293674498796463, 0.18016116321086884, 0.3294305205345154, 0.5326097011566162, 0.20817913115024567, 0.231731578707695, 0.17336609959602356, 0.4696378707885742, 0.3560185134410858, 0.5055418610572815, 0.687153697013855, 0.06569264829158783], [0.19887569546699524, 0.009285598993301392, 0.17495201528072357, 0.1799449920654297, 0.0410592183470726, 0.0050115324556827545, 0.025978662073612213, 0.011312133632600307, 0.04069671407341957, 0.23767657577991486, 0.3294059634208679, 0.09899688512086868, 0.03285939246416092, 0.08387716114521027, 0.04885585233569145], [0.054675761610269547, 0.04458622261881828, 0.0536046139895916, 0.016943499445915222, 0.02146792784333229, 0.1686052531003952, 0.036354243755340576, 0.08614800870418549, 0.1611979901790619, 0.170720174908638, 0.163726344704628, 0.09202460944652557, 0.016866492107510567, 0.019021833315491676, 0.13082824647426605], [0.254617303609848, 0.09600356966257095, 0.5283652544021606, 0.35948434472084045, 0.11690203100442886, 0.22449535131454468, 0.07030754536390305, 0.14074397087097168, 0.11056768894195557, 0.2017645388841629, 0.5897989273071289, 0.032950446009635925, 0.0850306898355484, 0.16881772875785828, 0.07667817175388336], [0.06611059606075287, 0.009380446746945381, 0.1600489318370819, 0.18714633584022522, 0.028496628627181053, 0.28509950637817383, 0.06793918460607529, 0.036412376910448074, 0.3864555358886719, 0.38031718134880066, 0.19321800768375397, 0.03279240429401398, 0.024823389947414398, 0.02684853971004486, 0.10572600364685059], [0.5806823372840881, 0.09046274423599243, 0.1468239277601242, 0.2587219774723053, 0.018666794523596764, 0.17986845970153809, 0.1758078932762146, 0.26734092831611633, 0.30597683787345886, 0.6407824158668518, 0.6427304148674011, 0.011203133501112461, 0.017842967063188553, 0.05609212443232536, 0.1528221219778061], [0.09578646719455719, 0.04883359372615814, 0.014442636631429195, 0.07719788700342178, 0.013871591538190842, 0.24272511899471283, 0.11848346889019012, 0.48695430159568787, 0.10090471804141998, 0.15632015466690063, 0.12246286869049072, 0.056596189737319946, 0.051980338990688324, 0.03806659206748009, 0.1369783878326416], [0.12923087179660797, 0.04506811499595642, 0.5631698966026306, 0.4945719838142395, 0.16776354610919952, 0.4656532406806946, 0.6344242095947266, 0.28209388256073, 0.297488808631897, 0.3520771265029907, 0.6463941931724548, 0.3803158104419708, 0.4924411177635193, 0.6891878843307495, 0.08469904214143753], [0.3177553117275238, 0.027823492884635925, 0.11541304737329483, 0.1464630663394928, 0.010460668243467808, 0.028609508648514748, 0.14352867007255554, 0.043905869126319885, 0.18215790390968323, 0.6030426025390625, 0.38763877749443054, 0.1293274313211441, 0.07180552184581757, 0.1464845985174179, 0.10971048474311829], [0.03459807112812996, 0.05000016465783119, 0.02839210256934166, 0.008521324954926968, 0.009519261308014393, 0.12168280780315399, 0.03372196480631828, 0.07665831595659256, 0.21765880286693573, 0.11945746093988419, 0.0821232944726944, 0.058310747146606445, 0.011853469535708427, 0.02031784877181053, 0.13586042821407318], [0.02964477799832821, 0.1353258490562439, 0.017653465270996094, 0.011115004308521748, 0.008141545578837395, 0.05911250412464142, 0.01831989735364914, 0.05519499629735947, 0.03573962301015854, 0.02204814739525318, 0.05097896233201027, 0.08341387659311295, 0.08060181885957718, 0.10490117967128754, 0.13247323036193848]], [[0.20067201554775238, 0.150595024228096, 0.3375815153121948, 0.5753223896026611, 0.03983612731099129, 0.13901081681251526, 0.37267425656318665, 0.07406412810087204, 0.07071352750062943, 0.22996902465820312, 0.35784539580345154, 0.0401473231613636, 0.03251379355788231, 0.07572956383228302, 0.005637211725115776], [0.055522263050079346, 0.0030253075528889894, 0.054468654096126556, 0.18383808434009552, 0.2751407325267792, 0.06163792684674263, 0.5092534422874451, 0.21577699482440948, 0.23691882193088531, 0.32801976799964905, 0.29786956310272217, 0.4967685043811798, 0.6341143250465393, 0.7677603363990784, 0.40264371037483215], [0.0005822544917464256, 0.0004425827646628022, 0.0014265297213569283, 0.0006841197027824819, 0.03406556695699692, 0.0010687633184716105, 0.0028485425282269716, 0.020860498771071434, 0.05133597180247307, 0.002158694202080369, 0.002441320102661848, 0.037159714847803116, 0.005256796721369028, 0.008102376013994217, 0.16207638382911682], [0.20224374532699585, 0.7376267313957214, 0.004014236852526665, 0.0103965038433671, 0.07275543361902237, 0.03262623772025108, 0.04577071964740753, 0.5017040371894836, 0.12205435335636139, 0.19255708158016205, 0.006990006659179926, 0.028381695970892906, 0.046785227954387665, 0.15206293761730194, 0.330488920211792], [0.3634231686592102, 0.404717355966568, 0.00689590023830533, 0.04770800471305847, 0.0251657422631979, 0.0006883289897814393, 0.02071242779493332, 0.019072405993938446, 0.15776626765727997, 0.3694642186164856, 0.036826737225055695, 0.23951902985572815, 0.011015082709491253, 0.04999716952443123, 0.2037181556224823], [0.8270207643508911, 0.8942698836326599, 0.020243747159838676, 0.04263966530561447, 0.09284591674804688, 0.054453812539577484, 0.21418678760528564, 0.23612302541732788, 0.5479635000228882, 0.7225908041000366, 0.08608872443437576, 0.5934221148490906, 0.30024465918540955, 0.22648638486862183, 0.12622572481632233], [0.043734412640333176, 0.7137998342514038, 0.1370490938425064, 0.045488547533750534, 0.06789389997720718, 0.49671053886413574, 0.1280447244644165, 0.4211912155151367, 0.03652801364660263, 0.041476957499980927, 0.08040425181388855, 0.19641457498073578, 0.603863537311554, 0.49263066053390503, 0.07636027038097382], [0.017375759780406952, 0.012506993487477303, 0.020720014348626137, 0.011049210093915462, 0.03743210807442665, 0.0072485157288610935, 0.03524084761738777, 0.005443913396447897, 0.24646395444869995, 0.048276107758283615, 0.03640883043408394, 0.507624089717865, 0.15355341136455536, 0.1730290949344635, 0.2644885182380676], [0.09840062260627747, 0.7509858012199402, 0.13933908939361572, 0.13482652604579926, 0.18154919147491455, 0.32397931814193726, 0.23646889626979828, 0.11657525599002838, 0.03430478647351265, 0.1277371644973755, 0.15700362622737885, 0.24829043447971344, 0.7591869831085205, 0.7825927138328552, 0.06869770586490631], [0.22806629538536072, 0.6706615686416626, 0.2560598850250244, 0.17412559688091278, 0.6327939033508301, 0.04699348285794258, 0.058767881244421005, 0.11556732654571533, 0.09056147933006287, 0.3648419678211212, 0.5388886332511902, 0.261055588722229, 0.6016876697540283, 0.7496042847633362, 0.0894755870103836], [0.5419997572898865, 0.6956567168235779, 0.044124722480773926, 0.12586495280265808, 0.048711128532886505, 0.11729516834020615, 0.4073715806007385, 0.43757542967796326, 0.032695479691028595, 0.4824156165122986, 0.05927032604813576, 0.04766178876161575, 0.25393223762512207, 0.23675066232681274, 0.10572775453329086], [0.09369882941246033, 0.5731168985366821, 0.13611510396003723, 0.13756731152534485, 0.024227088317275047, 0.31910547614097595, 0.16772453486919403, 0.1680929958820343, 0.09319504350423813, 0.0998181626200676, 0.22465890645980835, 0.00899507012218237, 0.16640731692314148, 0.25350457429885864, 0.09016240388154984], [0.02838694490492344, 0.30040091276168823, 0.005878766532987356, 0.015430719591677189, 0.017050068825483322, 0.06605669111013412, 0.12745192646980286, 0.23377051949501038, 0.08052214235067368, 0.033177152276039124, 0.06731567531824112, 0.07575374841690063, 0.18187224864959717, 0.570769727230072, 0.04572387412190437], [0.2655380666255951, 0.4107033908367157, 0.04865417629480362, 0.08488347381353378, 0.04310445114970207, 0.10849997401237488, 0.15643075108528137, 0.04165918007493019, 0.12898734211921692, 0.11095981299877167, 0.23520684242248535, 0.10632039606571198, 0.055878568440675735, 0.24558725953102112, 0.17682571709156036], [0.8565200567245483, 0.8639481067657471, 0.0803997814655304, 0.36449819803237915, 0.17448320984840393, 0.12402030825614929, 0.13765643537044525, 0.2065785825252533, 0.18182852864265442, 0.6806339025497437, 0.1919344812631607, 0.19068314135074615, 0.004361266735941172, 0.01490570418536663, 0.13936595618724823], [0.22751423716545105, 0.21127405762672424, 0.005130667705088854, 0.028237944468855858, 0.06646221876144409, 0.045109983533620834, 0.478432834148407, 0.6443154215812683, 0.140235036611557, 0.0980456992983818, 0.006476161070168018, 0.038696710020303726, 0.25798937678337097, 0.10561345517635345, 0.16755780577659607], [0.3886019289493561, 0.36600789427757263, 0.07069597393274307, 0.12792876362800598, 0.0629734918475151, 0.0820467472076416, 0.2973020672798157, 0.27475541830062866, 0.019707435742020607, 0.2982620298862457, 0.24423947930335999, 0.05686682090163231, 0.23438367247581482, 0.3444555997848511, 0.09858046472072601], [0.31350865960121155, 0.5118260383605957, 0.01775331422686577, 0.060602445155382156, 0.015971101820468903, 0.03445184975862503, 0.4316053092479706, 0.4819965064525604, 0.008238772861659527, 0.27349013090133667, 0.02135261707007885, 0.006705985404551029, 0.06119696795940399, 0.05213680863380432, 0.13011163473129272], [0.11128952354192734, 0.6662537455558777, 0.10913366079330444, 0.08027850091457367, 0.016604425385594368, 0.1904260814189911, 0.09001538157463074, 0.12034764140844345, 0.032395973801612854, 0.07767382264137268, 0.13288450241088867, 0.0038343279156833887, 0.15461067855358124, 0.13092683255672455, 0.1198263093829155], [0.045069050043821335, 0.5156355500221252, 0.014353718608617783, 0.026371080428361893, 0.027669712901115417, 0.08119883388280869, 0.2510265111923218, 0.45373910665512085, 0.0644708126783371, 0.03346102684736252, 0.06456929445266724, 0.036929432302713394, 0.1635800451040268, 0.4964689314365387, 0.12627021968364716], [0.15574656426906586, 0.22756966948509216, 0.016156630590558052, 0.0469389408826828, 0.01719032973051071, 0.01580459624528885, 0.07493647187948227, 0.02412206307053566, 0.018628407269716263, 0.03879624605178833, 0.03891688585281372, 0.03379734605550766, 0.008454171009361744, 0.03055991418659687, 0.1906210333108902], [0.7930518984794617, 0.8248118162155151, 0.03787774592638016, 0.2306395173072815, 0.10945193469524384, 0.048738475888967514, 0.07385316491127014, 0.1171715259552002, 0.09199279546737671, 0.5013920664787292, 0.07074998319149017, 0.14583703875541687, 0.0018764830892905593, 0.00646476075053215, 0.13562877476215363], [0.139163076877594, 0.17112046480178833, 0.0021531793754547834, 0.0053843106143176556, 0.013183848932385445, 0.014547600410878658, 0.39682450890541077, 0.7216413021087646, 0.013683686964213848, 0.038195278495550156, 0.0014429710572585464, 0.0075409854762256145, 0.06976743042469025, 0.016425929963588715, 0.1257757991552353], [0.37428542971611023, 0.3404470980167389, 0.07186836749315262, 0.11062464118003845, 0.09624961018562317, 0.06910651177167892, 0.26704323291778564, 0.35990291833877563, 0.016681469976902008, 0.31615501642227173, 0.23382727801799774, 0.051282789558172226, 0.1643712818622589, 0.24623094499111176, 0.1059461385011673], [0.2896858751773834, 0.2041676938533783, 0.0844137892127037, 0.26597079634666443, 0.007990201003849506, 0.057605594396591187, 0.37075188755989075, 0.33039090037345886, 0.04668770357966423, 0.6492098569869995, 0.34850311279296875, 0.12703292071819305, 0.22453922033309937, 0.2423134297132492, 0.11649563163518906]]], [[[0.12698857486248016, 0.15100647509098053, 0.08910781890153885, 0.09401589632034302, 0.14288602769374847, 0.07712502032518387, 0.1496707946062088, 0.23784373700618744, 0.024656152352690697, 0.07261883467435837, 0.11269068717956543, 0.10889188945293427, 0.23155105113983154, 0.10633593797683716, 0.14060717821121216], [0.33520859479904175, 0.17541100084781647, 0.043081097304821014, 0.07071122527122498, 0.031066332012414932, 0.05302952229976654, 0.13712948560714722, 0.0819549486041069, 0.010218805633485317, 0.05350261554121971, 0.03376028686761856, 0.016291575506329536, 0.04384060204029083, 0.016914406791329384, 0.06937505304813385], [0.2972787618637085, 0.14542943239212036, 0.2801832854747772, 0.6946116089820862, 0.3750338852405548, 0.09368664771318436, 0.11078806221485138, 0.124379463493824, 0.028408339247107506, 0.3442523181438446, 0.15075638890266418, 0.08511755615472794, 0.32891392707824707, 0.12337944656610489, 0.05913665145635605], [0.06821048259735107, 0.007578656077384949, 0.033511072397232056, 0.039627932012081146, 0.016393400728702545, 0.20925503969192505, 0.15704192221164703, 0.024064799770712852, 0.005696912761777639, 0.01698312722146511, 0.15042142570018768, 0.0017041407991200686, 0.016995420679450035, 0.005758653394877911, 0.015053601935505867], [0.05268644914031029, 0.018480738624930382, 0.006206580437719822, 0.01908770017325878, 0.009213676676154137, 0.012446015141904354, 0.2606332302093506, 0.15275397896766663, 0.004711512941867113, 0.01064901053905487, 0.00940486416220665, 0.00429189158603549, 0.014810611493885517, 0.012880465015769005, 0.15466143190860748], [0.017502065747976303, 0.09008979797363281, 0.045234303921461105, 0.04321402683854103, 0.014162504114210606, 0.2841097414493561, 0.10382679849863052, 0.4497845470905304, 0.042821191251277924, 0.03918898105621338, 0.06416238099336624, 0.04602029174566269, 0.2197093665599823, 0.07547488063573837, 0.13285692036151886], [0.02909473329782486, 0.05293780937790871, 0.025932423770427704, 0.061369478702545166, 0.12287095934152603, 0.12207728624343872, 0.20267462730407715, 0.3647293746471405, 0.036313559859991074, 0.028358493000268936, 0.054471470415592194, 0.007501897402107716, 0.10796680301427841, 0.05851392075419426, 0.12157665193080902], [0.02889016829431057, 0.05256107077002525, 0.05110660940408707, 0.09513585269451141, 0.049980901181697845, 0.07343146204948425, 0.21190620958805084, 0.10279127210378647, 0.1787082403898239, 0.022944355383515358, 0.03947293758392334, 0.008258121088147163, 0.09723227471113205, 0.030062679201364517, 0.14898137748241425], [0.027054987847805023, 0.06796294450759888, 0.02347770519554615, 0.04540639370679855, 0.13579830527305603, 0.1935206949710846, 0.09281998127698898, 0.22921815514564514, 0.012567882426083088, 0.02752627059817314, 0.05939676612615585, 0.00633750855922699, 0.24427738785743713, 0.10302533209323883, 0.18246731162071228], [0.13923436403274536, 0.07431720942258835, 0.06541924923658371, 0.14132679998874664, 0.10506866127252579, 0.06156519800424576, 0.21440355479717255, 0.06509862840175629, 0.02759510651230812, 0.10144857317209244, 0.13265900313854218, 0.048845868557691574, 0.16166719794273376, 0.1116088330745697, 0.15105699002742767], [0.14352908730506897, 0.10288456827402115, 0.05261845886707306, 0.1541282832622528, 0.05661991983652115, 0.12065587192773819, 0.10697692632675171, 0.15951323509216309, 0.1055477038025856, 0.14385449886322021, 0.23090383410453796, 0.08539394289255142, 0.09938428550958633, 0.08322764188051224, 0.11896289885044098], [0.24387870728969574, 0.11191204935312271, 0.06428070366382599, 0.3038298189640045, 0.14750736951828003, 0.1200045570731163, 0.46686112880706787, 0.3116493225097656, 0.10273779183626175, 0.10795925557613373, 0.1416371762752533, 0.09460661560297012, 0.27618303894996643, 0.09149192273616791, 0.10828596353530884], [0.1039203479886055, 0.05052376165986061, 0.051659513264894485, 0.18036356568336487, 0.11265069991350174, 0.047071922570466995, 0.3453211784362793, 0.29340654611587524, 0.007079527713358402, 0.06730296462774277, 0.08055143058300018, 0.02563900128006935, 0.19650228321552277, 0.060815099626779556, 0.13184599578380585], [0.1947154402732849, 0.003113611601293087, 0.028957238420844078, 0.026910793036222458, 0.017121652141213417, 0.08169777691364288, 0.32467299699783325, 0.05661681666970253, 0.007502032909542322, 0.02869880571961403, 0.020577264949679375, 0.0070375413633883, 0.16551434993743896, 0.06083058565855026, 0.06852211803197861], [0.018467016518115997, 0.004791167099028826, 0.015553582459688187, 0.021664531901478767, 0.025298617780208588, 0.1971224695444107, 0.13395515084266663, 0.1881190687417984, 0.05309745669364929, 0.018728721886873245, 0.018886514008045197, 0.023248562589287758, 0.008927382528781891, 0.03253133222460747, 0.130488321185112], [0.4018593430519104, 0.09619066119194031, 0.047895513474941254, 0.0887020081281662, 0.04670756310224533, 0.17605426907539368, 0.21604543924331665, 0.1403813511133194, 0.0010993692558258772, 0.07762767374515533, 0.0958188846707344, 0.1024225577712059, 0.06565871089696884, 0.04857100546360016, 0.1717240959405899], [0.31909966468811035, 0.26355716586112976, 0.16833621263504028, 0.334572434425354, 0.18670302629470825, 0.11206400394439697, 0.46585598587989807, 0.15377958118915558, 0.014857469126582146, 0.07049962878227234, 0.1590365469455719, 0.09933225810527802, 0.23580892384052277, 0.09940709918737411, 0.11795931309461594], [0.3361136317253113, 0.18450267612934113, 0.10482683777809143, 0.3672127425670624, 0.09347432106733322, 0.06302808225154877, 0.17493662238121033, 0.11965186893939972, 0.06742112338542938, 0.13331438601016998, 0.26999813318252563, 0.03264465183019638, 0.07908355444669724, 0.09376725554466248, 0.11511774361133575], [0.271436870098114, 0.16103556752204895, 0.09723401814699173, 0.3494490087032318, 0.1582973301410675, 0.11393263936042786, 0.41371721029281616, 0.2938876152038574, 0.08068472146987915, 0.08301044255495071, 0.11968915909528732, 0.07779402285814285, 0.24559125304222107, 0.07589462399482727, 0.1087639182806015], [0.1091129332780838, 0.08970999717712402, 0.08557470142841339, 0.23009367287158966, 0.13180004060268402, 0.0638015940785408, 0.31095248460769653, 0.2814267873764038, 0.0075759077444672585, 0.039292845875024796, 0.06780961900949478, 0.013560868799686432, 0.15987654030323029, 0.04180291295051575, 0.12740370631217957], [0.4568881392478943, 0.01152532733976841, 0.12744615972042084, 0.16633041203022003, 0.05682089552283287, 0.22013583779335022, 0.46718865633010864, 0.06831676512956619, 0.011846139095723629, 0.051503561437129974, 0.07631707936525345, 0.017341753467917442, 0.16032609343528748, 0.06682911515235901, 0.06364742666482925], [0.0270079392939806, 0.003701634705066681, 0.024473953992128372, 0.035727839916944504, 0.031186459586024284, 0.22590965032577515, 0.1764952838420868, 0.1725662350654602, 0.06108492240309715, 0.017804577946662903, 0.01644762232899666, 0.018474329262971878, 0.0059660994447767735, 0.026993868872523308, 0.12890712916851044], [0.32686647772789, 0.10561588406562805, 0.10599718242883682, 0.08397059142589569, 0.05158340185880661, 0.22573474049568176, 0.19403943419456482, 0.08219113945960999, 0.0007591660832986236, 0.028280239552259445, 0.06139420345425606, 0.03943438082933426, 0.025857241824269295, 0.027251310646533966, 0.1435350626707077], [0.21139562129974365, 0.21867576241493225, 0.17973701655864716, 0.29884445667266846, 0.19560806453227997, 0.11132223159074783, 0.28179141879081726, 0.10507592558860779, 0.014165982604026794, 0.04481332749128342, 0.1297360062599182, 0.07738039642572403, 0.2323194295167923, 0.09134778380393982, 0.12234959006309509], [0.2484172284603119, 0.2714419662952423, 0.13623963296413422, 0.33317360281944275, 0.14056812226772308, 0.16453251242637634, 0.23482279479503632, 0.2797185182571411, 0.08398787677288055, 0.13855448365211487, 0.19988903403282166, 0.12159004807472229, 0.21263501048088074, 0.1342880129814148, 0.11613592505455017]], [[0.1659475415945053, 0.1821746528148651, 0.2680368423461914, 0.3257308900356293, 0.2135642170906067, 0.10952500998973846, 0.23729652166366577, 0.15246635675430298, 0.09328519552946091, 0.22413431107997894, 0.22322525084018707, 0.11237151175737381, 0.18681256473064423, 0.1572018712759018, 0.06837792694568634], [0.14290380477905273, 0.026570750400424004, 0.14845344424247742, 0.26635152101516724, 0.12476544827222824, 0.1522083431482315, 0.287058562040329, 0.16522644460201263, 0.21008911728858948, 0.3761942982673645, 0.12840349972248077, 0.0757022351026535, 0.39944273233413696, 0.379029244184494, 0.1911974847316742], [0.00885845348238945, 0.005625984165817499, 0.0020030708983540535, 0.005766861606389284, 0.001782223698683083, 0.004346099682152271, 0.014438317157328129, 0.010037342086434364, 0.0175970196723938, 0.0067982920445501804, 0.003056151093915105, 0.005088370759040117, 0.0035549686290323734, 0.002117584692314267, 0.17935973405838013], [0.04871530085802078, 0.2322341799736023, 0.043161727488040924, 0.046935759484767914, 0.04166096821427345, 0.048159919679164886, 0.2838554382324219, 0.5679410696029663, 0.17445935308933258, 0.05776107683777809, 0.14550535380840302, 0.04300517588853836, 0.2332015484571457, 0.28196635842323303, 0.4675023853778839], [0.03277377411723137, 0.28776609897613525, 0.0018310850718989968, 0.006392122711986303, 0.0034063432831317186, 0.0006021481240168214, 0.02006486989557743, 0.09552518278360367, 0.02804744802415371, 0.060428690165281296, 0.004742977675050497, 0.018782831728458405, 0.016696294769644737, 0.023774143308401108, 0.16262513399124146], [0.006045958958566189, 0.0958699956536293, 0.007954242639243603, 0.011606856249272823, 0.004544504452496767, 0.010406642220914364, 0.011899203062057495, 0.07300186902284622, 0.002370428293943405, 0.012239865958690643, 0.020374998450279236, 0.012496876530349255, 0.024265890941023827, 0.0274967048317194, 0.1423870474100113], [0.008809137158095837, 0.13565093278884888, 0.03191651031374931, 0.0483417883515358, 0.028707973659038544, 0.039296794682741165, 0.018359076231718063, 0.07145766168832779, 0.13921810686588287, 0.01646633818745613, 0.06145479157567024, 0.028490308672189713, 0.056069642305374146, 0.13838331401348114, 0.19134177267551422], [0.39272594451904297, 0.39728477597236633, 0.32111606001853943, 0.41796234250068665, 0.15293559432029724, 0.04586965963244438, 0.16940170526504517, 0.022719532251358032, 0.14239482581615448, 0.5121501088142395, 0.19016578793525696, 0.06530822068452835, 0.29211705923080444, 0.14742477238178253, 0.11553633958101273], [0.009060109965503216, 0.08736205101013184, 0.03623565658926964, 0.046393588185310364, 0.04293924570083618, 0.049119193106889725, 0.018734706565737724, 0.10957584530115128, 0.04821338504552841, 0.02008068934082985, 0.029284991323947906, 0.015971768647432327, 0.05779576674103737, 0.21830672025680542, 0.21264111995697021], [0.02833615615963936, 0.24966742098331451, 0.06237170845270157, 0.03993965685367584, 0.10454770177602768, 0.019859671592712402, 0.03772445023059845, 0.19178973138332367, 0.012827831320464611, 0.03533304110169411, 0.024230163544416428, 0.054630037397146225, 0.032379381358623505, 0.08906079828739166, 0.17152637243270874], [0.015255320817232132, 0.21888743340969086, 0.1253896951675415, 0.08362822234630585, 0.12500159442424774, 0.02890017069876194, 0.03405824303627014, 0.07477163523435593, 0.0229325033724308, 0.01863025315105915, 0.044950928539037704, 0.0560457706451416, 0.04699615016579628, 0.08650227636098862, 0.1548503190279007], [0.011826024390757084, 0.10608652234077454, 0.04723645746707916, 0.057715099304914474, 0.03395959734916687, 0.028910892084240913, 0.011586843058466911, 0.050380002707242966, 0.030421555042266846, 0.00583301018923521, 0.015118762850761414, 0.014350258745253086, 0.01606619358062744, 0.025515934452414513, 0.18496018648147583], [0.015032858587801456, 0.5077551603317261, 0.07541441917419434, 0.08020945638418198, 0.10545077919960022, 0.2137133628129959, 0.01040775515139103, 0.09528981149196625, 0.09038985520601273, 0.012094871141016483, 0.025733938440680504, 0.06706724315881729, 0.03145073354244232, 0.09538157284259796, 0.34148263931274414], [0.32250380516052246, 0.7984310388565063, 0.3962976634502411, 0.40014326572418213, 0.3554738759994507, 0.47898975014686584, 0.10853014886379242, 0.20243746042251587, 0.127571240067482, 0.2699570655822754, 0.16473528742790222, 0.08001074939966202, 0.03713205084204674, 0.14643853902816772, 0.4229389429092407], [0.023898553103208542, 0.03448064997792244, 0.007101188413798809, 0.020377272740006447, 0.09085186570882797, 0.008504875935614109, 0.01689869724214077, 0.021393392235040665, 0.03013733960688114, 0.004040753003209829, 0.000672544410917908, 0.0007860396872274578, 0.0003324192948639393, 0.0003073772240895778, 0.13160185515880585], [0.025859396904706955, 0.29733914136886597, 0.09033425897359848, 0.06196272000670433, 0.10889838635921478, 0.14661002159118652, 0.034964289516210556, 0.07059973478317261, 0.007527152542024851, 0.007617437280714512, 0.006072000600397587, 0.0492180734872818, 0.0069811418652534485, 0.011496509425342083, 0.22706106305122375], [0.014849718660116196, 0.1462036818265915, 0.11065799742937088, 0.06219353526830673, 0.08005399256944656, 0.016894571483135223, 0.010269397869706154, 0.02562439627945423, 0.009192260913550854, 0.009821194224059582, 0.015785057097673416, 0.019254932180047035, 0.01222837995737791, 0.011684795841574669, 0.16154925525188446], [0.01973692700266838, 0.11480830609798431, 0.07148479670286179, 0.05237298831343651, 0.0777522474527359, 0.019268590956926346, 0.01592963933944702, 0.01235677395015955, 0.06519288569688797, 0.019938096404075623, 0.03185376524925232, 0.0271891038864851, 0.01742159202694893, 0.040164995938539505, 0.1837940812110901], [0.006014276295900345, 0.07228019088506699, 0.029915854334831238, 0.031709808856248856, 0.01963544264435768, 0.01660715602338314, 0.00532315531745553, 0.03606380149722099, 0.029185649007558823, 0.0046777487732470036, 0.01710142381489277, 0.013257446698844433, 0.01389795821160078, 0.02201540581882, 0.16183340549468994], [0.008549164049327374, 0.34144893288612366, 0.03957316279411316, 0.03764811158180237, 0.04039980471134186, 0.07271253317594528, 0.00613941578194499, 0.04612124711275101, 0.0911136344075203, 0.008750539273023605, 0.01715807057917118, 0.03749352693557739, 0.024577608332037926, 0.06848984956741333, 0.2503378689289093], [0.1472499966621399, 0.4703251123428345, 0.2558133602142334, 0.283985435962677, 0.21470209956169128, 0.17662864923477173, 0.07007063925266266, 0.06038873642683029, 0.20766907930374146, 0.26984694600105286, 0.16889145970344543, 0.27114859223365784, 0.03473396599292755, 0.13903996348381042, 0.2962591350078583], [0.020655758678913116, 0.020222418010234833, 0.006879583932459354, 0.019070995971560478, 0.07609020173549652, 0.006032301113009453, 0.015974652022123337, 0.01717195473611355, 0.05267442390322685, 0.004277344327419996, 0.0005684247589670122, 0.0007490122807212174, 0.0002994663082063198, 0.0002370573638472706, 0.12958088517189026], [0.009374987334012985, 0.23445867002010345, 0.05258592590689659, 0.020285839214920998, 0.024131227284669876, 0.0535256564617157, 0.01552440132945776, 0.032435644418001175, 0.006646827794611454, 0.005740212742239237, 0.005195626523345709, 0.07125341892242432, 0.0043562185019254684, 0.01014760322868824, 0.17807012796401978], [0.018758203834295273, 0.11843696236610413, 0.09101122617721558, 0.0610043928027153, 0.06165887042880058, 0.012400476261973381, 0.011786350980401039, 0.021215293556451797, 0.014211799949407578, 0.011016220785677433, 0.02130991406738758, 0.02418670989573002, 0.015627985820174217, 0.013993974775075912, 0.14536960422992706], [0.03985379636287689, 0.12957410514354706, 0.13386031985282898, 0.10592924803495407, 0.09455320239067078, 0.03913174197077751, 0.052976641803979874, 0.03812992200255394, 0.11070051789283752, 0.042073190212249756, 0.05433963984251022, 0.058929286897182465, 0.03380222246050835, 0.05054538697004318, 0.1317562311887741]], [[0.038382355123758316, 0.16509199142456055, 0.03795319423079491, 0.018471574410796165, 0.017937200143933296, 0.20822547376155853, 0.036850690841674805, 0.07025959342718124, 0.026183662936091423, 0.008891633711755276, 0.011525453999638557, 0.06559614092111588, 0.10240377485752106, 0.05705304443836212, 0.19186913967132568], [0.18736660480499268, 0.12802250683307648, 0.06000450998544693, 0.07085607945919037, 0.02492770366370678, 0.13308653235435486, 0.01379183866083622, 0.01460492704063654, 0.018005041405558586, 0.18972568213939667, 0.18918126821517944, 0.05261359363794327, 0.08419474214315414, 0.039842329919338226, 0.12843605875968933], [0.003212069161236286, 0.04924406483769417, 0.010131219401955605, 0.0015629208646714687, 0.009065762162208557, 0.04507109895348549, 0.003221129300072789, 0.07382506877183914, 0.0011923180427402258, 0.004047631751745939, 0.006328214425593615, 0.012952281162142754, 0.0641837865114212, 0.02541324496269226, 0.1715373396873474], [0.002438034862279892, 0.0007996301865205169, 0.10929557681083679, 0.030698396265506744, 0.007961505092680454, 0.21520712971687317, 0.0018748894799500704, 0.0015670642023906112, 0.00039643081254325807, 0.0017966092564165592, 0.010619523003697395, 0.0026792865246534348, 0.0035868084523826838, 0.001077426946721971, 0.003137440187856555], [0.04913554713129997, 0.023452362045645714, 0.16805477440357208, 0.2746557891368866, 0.369334876537323, 0.025402046740055084, 0.03595297038555145, 0.27975642681121826, 0.005478397477418184, 0.044800374656915665, 0.028408128768205643, 0.025396348908543587, 0.1202942430973053, 0.22760754823684692, 0.12602998316287994], [0.0008230121457017958, 0.006709535606205463, 0.005090394522994757, 0.005009432788938284, 0.0009200142812915146, 0.002589132636785507, 0.003276216797530651, 0.011904137209057808, 0.0009605096420273185, 0.0016532291192561388, 0.001647727913223207, 0.0010296034161001444, 0.00474548852071166, 0.004530362784862518, 0.14385877549648285], [0.011407818645238876, 0.11073090881109238, 0.11066732555627823, 0.07063236832618713, 0.2326628416776657, 0.057718440890312195, 0.005228970665484667, 0.12933272123336792, 0.010014788247644901, 0.0034599530044943094, 0.015450170263648033, 0.004393222741782665, 0.010258005000650883, 0.00790967233479023, 0.16524673998355865], [0.024886149913072586, 0.019822845235466957, 0.050577834248542786, 0.042761147022247314, 0.013624369166791439, 0.03171548992395401, 0.03447520360350609, 0.057101696729660034, 0.018126925453543663, 0.012612801045179367, 0.056599393486976624, 0.005686976481229067, 0.022324958816170692, 0.021004129201173782, 0.18438492715358734], [0.012148641981184483, 0.047028496861457825, 0.07792042940855026, 0.1455426812171936, 0.3985011875629425, 0.08270914107561111, 0.0031603944953531027, 0.07123681157827377, 0.020226983353495598, 0.005742877256125212, 0.009367674589157104, 0.007002389058470726, 0.013849785551428795, 0.006732230074703693, 0.14449873566627502], [0.029934342950582504, 0.04287242144346237, 0.10493571311235428, 0.10647397488355637, 0.01039193756878376, 0.1410648375749588, 0.06155749782919884, 0.08983614295721054, 0.05490254610776901, 0.038721270859241486, 0.021267540752887726, 0.05536682903766632, 0.019229264929890633, 0.008436290547251701, 0.15105655789375305], [0.009979508817195892, 0.08308109641075134, 0.026161497458815575, 0.023276647552847862, 0.0017319537000730634, 0.056630972772836685, 0.012614267878234386, 0.041058339178562164, 0.026752248406410217, 0.01169703807681799, 0.011314285919070244, 0.007283498533070087, 0.05053415521979332, 0.019243547692894936, 0.16277745366096497], [0.04712976887822151, 0.24274323880672455, 0.053717970848083496, 0.06948067992925644, 0.009206406772136688, 0.0471884086728096, 0.010105792433023453, 0.05801715701818466, 0.01891178824007511, 0.07684698700904846, 0.07729421555995941, 0.042662668973207474, 0.10241091996431351, 0.038032110780477524, 0.15563422441482544], [0.009955390356481075, 0.06358544528484344, 0.028598172590136528, 0.04170457646250725, 0.01363537646830082, 0.011423949152231216, 0.003101062262430787, 0.04170127958059311, 0.01145926769822836, 0.01274544931948185, 0.020664334297180176, 0.15329574048519135, 0.20515742897987366, 0.07666952162981033, 0.13521607220172882], [0.006747167091816664, 0.006801524665206671, 0.007903891615569592, 0.00237295706756413, 0.0009535709978081286, 0.0006887177005410194, 0.0011137888068333268, 0.0005580680444836617, 0.004365934059023857, 0.0043631866574287415, 0.004836279433220625, 0.0014166004257276654, 0.1882382482290268, 0.04424351081252098, 0.006875277496874332], [0.0040101236663758755, 0.00047035442548803985, 0.0008357138140127063, 0.009736553765833378, 0.00025759977870620787, 2.9679033104912378e-05, 0.008525178767740726, 0.0036214631982147694, 0.0009930779924616218, 0.0008531230851076543, 0.0029921825043857098, 7.93160234024981e-06, 6.746472354279831e-05, 0.0017078705132007599, 0.13162609934806824], [0.021027032285928726, 0.04388788715004921, 0.07337366044521332, 0.13240061700344086, 0.005691900383681059, 0.08179081231355667, 0.010154702700674534, 0.019539857283234596, 0.013572044670581818, 0.03972425311803818, 0.14196330308914185, 0.0491810142993927, 0.029326222836971283, 0.024830663576722145, 0.1775946319103241], [0.020570920780301094, 0.07008225470781326, 0.05771828070282936, 0.10093566030263901, 0.0037175160832703114, 0.10588520765304565, 0.008791210129857063, 0.07720224559307098, 0.037850137799978256, 0.016810759902000427, 0.0763774886727333, 0.06772230565547943, 0.10185997188091278, 0.02133399061858654, 0.1501101702451706], [0.027059482410550117, 0.22707954049110413, 0.13379518687725067, 0.08346803486347198, 0.011664706282317638, 0.1994924694299698, 0.013729198835790157, 0.07924864441156387, 0.10303384810686111, 0.02253318764269352, 0.06352351605892181, 0.13561668992042542, 0.3492315113544464, 0.13069112598896027, 0.12187084555625916], [0.038929592818021774, 0.2334582358598709, 0.12089657783508301, 0.17347271740436554, 0.023068996146321297, 0.04853734001517296, 0.008499456569552422, 0.0867975577712059, 0.02351396717131138, 0.04524386301636696, 0.12492679059505463, 0.06575564295053482, 0.10587428510189056, 0.055128976702690125, 0.1414995789527893], [0.011872883886098862, 0.08469298481941223, 0.054403409361839294, 0.08831894397735596, 0.02684788778424263, 0.021699469536542892, 0.0027920349966734648, 0.05190650746226311, 0.006984782870858908, 0.008844600059092045, 0.02751598134636879, 0.22613400220870972, 0.15431185066699982, 0.06476734578609467, 0.1412026435136795], [0.015115483663976192, 0.08628259599208832, 0.023322032764554024, 0.012461238540709019, 0.0028755213133990765, 0.010226217098534107, 0.0010302395094186068, 0.002081838669255376, 0.003762529231607914, 0.013111302629113197, 0.0290949996560812, 0.013309521600604057, 0.22778895497322083, 0.05992528051137924, 0.00796937569975853], [0.0057023135013878345, 0.0003758604871109128, 0.0009645622340030968, 0.01432577334344387, 0.00027227052487432957, 3.7724938010796905e-05, 0.007459490094333887, 0.0037525389343500137, 0.001061747083440423, 0.0008801367366686463, 0.0023195864632725716, 8.150678695528768e-06, 4.0667833673069254e-05, 0.001007204526104033, 0.12961283326148987], [0.017900969833135605, 0.026770949363708496, 0.15903817117214203, 0.31877970695495605, 0.014844128862023354, 0.10845804959535599, 0.00868347566574812, 0.015460771508514881, 0.008762474171817303, 0.01190071552991867, 0.07999671250581741, 0.053750935941934586, 0.013735906220972538, 0.020958656445145607, 0.15606556832790375], [0.022256335243582726, 0.07135839015245438, 0.07359576225280762, 0.12423767894506454, 0.006224590353667736, 0.13500085473060608, 0.008429165929555893, 0.08156562596559525, 0.02983916364610195, 0.013062523677945137, 0.10225346684455872, 0.04065772891044617, 0.06899033486843109, 0.012502058409154415, 0.13831046223640442], [0.016071150079369545, 0.06728275120258331, 0.025518205016851425, 0.023689931258559227, 0.0069392030127346516, 0.04150809720158577, 0.00898416806012392, 0.016712933778762817, 0.005143268499523401, 0.020111138001084328, 0.03020956739783287, 0.01359627302736044, 0.018198341131210327, 0.01637156493961811, 0.1379418522119522]], [[0.029921628534793854, 0.09876842796802521, 0.1324968934059143, 0.09236511588096619, 0.02831152267754078, 0.08077768236398697, 0.03118293546140194, 0.1750149130821228, 0.015778981149196625, 0.07032441347837448, 0.22269371151924133, 0.07579661160707474, 0.029184984043240547, 0.053061336278915405, 0.18562854826450348], [0.07805982232093811, 0.05365234240889549, 0.2842547595500946, 0.2606758773326874, 0.21293140947818756, 0.02651267871260643, 0.08033362030982971, 0.07913534343242645, 0.17101624608039856, 0.12522375583648682, 0.14315897226333618, 0.16815446317195892, 0.0695369690656662, 0.13316825032234192, 0.19111928343772888], [0.11272483319044113, 0.11636882275342941, 0.45685258507728577, 0.0910579040646553, 0.3091263473033905, 0.12632955610752106, 0.1822080761194229, 0.18498732149600983, 0.6353387832641602, 0.08394157886505127, 0.3285849094390869, 0.4818887710571289, 0.08592816442251205, 0.3495768904685974, 0.07449600845575333], [0.2834128737449646, 0.1102365031838417, 0.1840669959783554, 0.5708534121513367, 0.3157653212547302, 0.041008107364177704, 0.038309745490550995, 0.03211268410086632, 0.6102551817893982, 0.20786605775356293, 0.21116787195205688, 0.10018377006053925, 0.04653669148683548, 0.17929011583328247, 0.11314841359853745], [0.5993789434432983, 0.0908532664179802, 0.49218761920928955, 0.41100576519966125, 0.18825526535511017, 0.4342217445373535, 0.12116678059101105, 0.10673660039901733, 0.822167158126831, 0.4385586380958557, 0.6995345950126648, 0.18085956573486328, 0.1357179582118988, 0.2864921987056732, 0.034255724400281906], [0.858432412147522, 0.34460219740867615, 0.7778953909873962, 0.7743141651153564, 0.4405529797077179, 0.4761039614677429, 0.6155950427055359, 0.06873662024736404, 0.7323919534683228, 0.7086790204048157, 0.6720118522644043, 0.45794978737831116, 0.1628962755203247, 0.4249861538410187, 0.040913816541433334], [0.04546767473220825, 0.0383436344563961, 0.10268200188875198, 0.20100316405296326, 0.185649111866951, 0.08432896435260773, 0.060354892164468765, 0.07717668265104294, 0.3201402723789215, 0.04503992572426796, 0.088813915848732, 0.3990366756916046, 0.1564548909664154, 0.08066049963235855, 0.11440145969390869], [0.21178147196769714, 0.043018583208322525, 0.1065564677119255, 0.10858221352100372, 0.05675008147954941, 0.06700197607278824, 0.12675313651561737, 0.058651700615882874, 0.18508696556091309, 0.05493801832199097, 0.037313126027584076, 0.19010567665100098, 0.07823225855827332, 0.034572359174489975, 0.16783590614795685], [0.053469568490982056, 0.03894811123609543, 0.06651152670383453, 0.10646583139896393, 0.08985435962677002, 0.07578439265489578, 0.03395741805434227, 0.09802807122468948, 0.190333291888237, 0.07748086005449295, 0.07400990277528763, 0.6643930077552795, 0.07830479741096497, 0.07947986572980881, 0.11464671790599823], [0.1680978536605835, 0.06724530458450317, 0.16071708500385284, 0.2987021803855896, 0.11997595429420471, 0.007637033239006996, 0.05953739956021309, 0.06456195563077927, 0.07405640929937363, 0.11493658274412155, 0.07269633561372757, 0.12183233350515366, 0.019239120185375214, 0.0931614562869072, 0.15387272834777832], [0.09433168172836304, 0.05311369523406029, 0.44581180810928345, 0.2857709527015686, 0.11141614615917206, 0.04973546415567398, 0.10592624545097351, 0.0732862576842308, 0.26435965299606323, 0.07302475720643997, 0.17637307941913605, 0.06760746240615845, 0.052111051976680756, 0.29667070508003235, 0.11431443691253662], [0.07687122374773026, 0.10929025709629059, 0.4687592387199402, 0.20397132635116577, 0.26744040846824646, 0.03514130413532257, 0.033296968787908554, 0.08783485740423203, 0.22074763476848602, 0.08713625371456146, 0.12920482456684113, 0.05166565254330635, 0.07679110020399094, 0.17419996857643127, 0.1387287825345993], [0.061203911900520325, 0.12594261765480042, 0.353413462638855, 0.22131817042827606, 0.41015592217445374, 0.11432977020740509, 0.010031531564891338, 0.048355478793382645, 0.27572426199913025, 0.07773520797491074, 0.2322542816400528, 0.1527126431465149, 0.05797232687473297, 0.09810248017311096, 0.16366761922836304], [0.10230414569377899, 0.03857935592532158, 0.05230129137635231, 0.14396332204341888, 0.09251677989959717, 0.03541665896773338, 0.005624003708362579, 0.014271721243858337, 0.042375415563583374, 0.13543996214866638, 0.061749108135700226, 0.00788076315075159, 0.1602918803691864, 0.07564403861761093, 0.09375559538602829], [0.705120861530304, 0.026186510920524597, 0.8528315424919128, 0.8252069354057312, 0.24319231510162354, 0.07270172983407974, 0.09487330913543701, 0.07207771390676498, 0.4722364544868469, 0.7067926526069641, 0.8624283075332642, 0.07399676740169525, 0.0075901346281170845, 0.016478050500154495, 0.12560917437076569], [0.27840110659599304, 0.06363435834646225, 0.3689763844013214, 0.33064448833465576, 0.25749024748802185, 0.1453908383846283, 0.03645810857415199, 0.00836147554218769, 0.3977815508842468, 0.41805213689804077, 0.17756043374538422, 0.05318059027194977, 0.011340576224029064, 0.020938394591212273, 0.05934957042336464], [0.17816129326820374, 0.10609658807516098, 0.17893879115581512, 0.28182876110076904, 0.15060719847679138, 0.03372456133365631, 0.04276707395911217, 0.050946421921253204, 0.04137968271970749, 0.16634012758731842, 0.16395889222621918, 0.24548840522766113, 0.05229371041059494, 0.09448723495006561, 0.12793652713298798], [0.14424489438533783, 0.0705854520201683, 0.24214811623096466, 0.24549053609371185, 0.19939330220222473, 0.02639644220471382, 0.021373553201556206, 0.024115193635225296, 0.08405331522226334, 0.14685925841331482, 0.15661610662937164, 0.06219787895679474, 0.032059792429208755, 0.09036684036254883, 0.15146715939044952], [0.06650430709123611, 0.10705426335334778, 0.3146411180496216, 0.1647443175315857, 0.23945462703704834, 0.035643309354782104, 0.026562364771962166, 0.09605439007282257, 0.19827118515968323, 0.1037423387169838, 0.14283734560012817, 0.08165161311626434, 0.07012972235679626, 0.11072988063097, 0.13417953252792358], [0.06460674107074738, 0.10897383838891983, 0.18354696035385132, 0.20187535881996155, 0.38844820857048035, 0.04722803831100464, 0.010622762143611908, 0.04332485795021057, 0.31279584765434265, 0.11892355233430862, 0.20366235077381134, 0.1460915356874466, 0.041410893201828, 0.060890424996614456, 0.16885291039943695], [0.08445128798484802, 0.07278266549110413, 0.017734743654727936, 0.12906457483768463, 0.17354236543178558, 0.01439378596842289, 0.0032682251185178757, 0.009051240049302578, 0.02403325028717518, 0.17859239876270294, 0.05114053934812546, 0.026160510256886482, 0.17188863456249237, 0.059929899871349335, 0.12745818495750427], [0.6940725445747375, 0.016104217618703842, 0.8427497148513794, 0.8075915575027466, 0.2572270333766937, 0.04667792096734047, 0.07690176367759705, 0.06650352478027344, 0.4641934931278229, 0.7403572797775269, 0.892522931098938, 0.08286882191896439, 0.00509345019236207, 0.009769911877810955, 0.1252693384885788], [0.47638654708862305, 0.08160793781280518, 0.2188907116651535, 0.3983159363269806, 0.3041192293167114, 0.0773146003484726, 0.041229549795389175, 0.00785501953214407, 0.20719125866889954, 0.6323855519294739, 0.1790589690208435, 0.15920953452587128, 0.005728188902139664, 0.011172757484018803, 0.10331764072179794], [0.3162515461444855, 0.12029282748699188, 0.1898643672466278, 0.3138664960861206, 0.22235795855522156, 0.03812789171934128, 0.07994988560676575, 0.07006566971540451, 0.06856126338243484, 0.2470276951789856, 0.2142392098903656, 0.4667101502418518, 0.07071195542812347, 0.09391427785158157, 0.11791101843118668], [0.15722334384918213, 0.11492010205984116, 0.22595097124576569, 0.17283931374549866, 0.11246844381093979, 0.07424511015415192, 0.1308857947587967, 0.1509532928466797, 0.12219540029764175, 0.14498494565486908, 0.13763099908828735, 0.16327989101409912, 0.12245305627584457, 0.21428720653057098, 0.12265608459711075]], [[0.03995227441191673, 0.02612248808145523, 0.09039098769426346, 0.04685363546013832, 0.14171013236045837, 0.3046724796295166, 0.08713044226169586, 0.11726538836956024, 0.3945818245410919, 0.03867875412106514, 0.060879118740558624, 0.3211958110332489, 0.1562168449163437, 0.1954476237297058, 0.12928469479084015], [0.138319730758667, 0.1925395429134369, 0.06914161890745163, 0.1830926090478897, 0.22252067923545837, 0.24239645898342133, 0.2738734483718872, 0.3115195333957672, 0.287569522857666, 0.12556934356689453, 0.047479670494794846, 0.1859251707792282, 0.015966184437274933, 0.050888173282146454, 0.04287213087081909], [0.059622667729854584, 0.19761067628860474, 0.019807182252407074, 0.02911451645195484, 0.11472073942422867, 0.03754669055342674, 0.08183436095714569, 0.09122617542743683, 0.10595303028821945, 0.094895139336586, 0.022252719849348068, 0.087751105427742, 0.015402892604470253, 0.02668953314423561, 0.15029701590538025], [0.4440009295940399, 0.5055950880050659, 0.14072291553020477, 0.20776981115341187, 0.24339812994003296, 0.01946749910712242, 0.1477651447057724, 0.24892206490039825, 0.13990418612957, 0.5277839303016663, 0.22113053500652313, 0.7815175652503967, 0.04741470143198967, 0.31336119771003723, 0.318754643201828], [0.003975332248955965, 0.09357346594333649, 0.000580776366405189, 0.001556370290927589, 0.0040078358724713326, 0.00020105167641304433, 0.005314813926815987, 0.0463886484503746, 0.0025405578780919313, 0.008098164573311806, 0.0004367573419585824, 0.0955028310418129, 0.0013312119990587234, 0.008472515270113945, 0.16612127423286438], [0.00713347876444459, 0.11304348707199097, 0.007166451308876276, 0.017305465415120125, 0.01892760582268238, 0.004294875077903271, 0.013284130021929741, 0.05641845986247063, 0.006293897051364183, 0.008091668598353863, 0.004229044076055288, 0.03852742537856102, 0.036073870956897736, 0.030675750225782394, 0.1423715502023697], [0.112990602850914, 0.20299020409584045, 0.29141831398010254, 0.1917479783296585, 0.25626659393310547, 0.40023526549339294, 0.045914653688669205, 0.05403761938214302, 0.3577503561973572, 0.11164049804210663, 0.20054538547992706, 0.23382915556430817, 0.3541012704372406, 0.39880213141441345, 0.05442150682210922], [0.11769542098045349, 0.22490660846233368, 0.16446754336357117, 0.17726869881153107, 0.24409359693527222, 0.16966795921325684, 0.06426751613616943, 0.1868649125099182, 0.17593497037887573, 0.10732528567314148, 0.1210716962814331, 0.18835949897766113, 0.07820838689804077, 0.12172650545835495, 0.0815061554312706], [0.08801974356174469, 0.2964327037334442, 0.17140379548072815, 0.1086457222700119, 0.1790848970413208, 0.042561717331409454, 0.02568918652832508, 0.12736740708351135, 0.4644424617290497, 0.09952269494533539, 0.1403166949748993, 0.12085206061601639, 0.2499331831932068, 0.14905890822410583, 0.04691213369369507], [0.28339406847953796, 0.25363603234291077, 0.49371209740638733, 0.28714650869369507, 0.42171764373779297, 0.03586414083838463, 0.140908345580101, 0.27345338463783264, 0.06897412985563278, 0.24740128219127655, 0.5061832070350647, 0.4192107915878296, 0.43851029872894287, 0.29079654812812805, 0.10071542859077454], [0.049345988780260086, 0.1473262906074524, 0.10952533781528473, 0.16707968711853027, 0.25493475794792175, 0.03866606950759888, 0.046480532735586166, 0.16288119554519653, 0.06614720076322556, 0.0629507377743721, 0.07218940556049347, 0.3448391556739807, 0.06943795084953308, 0.058807674795389175, 0.135455921292305], [0.05557708069682121, 0.024377070367336273, 0.171014666557312, 0.1548214852809906, 0.21205416321754456, 0.29049578309059143, 0.08155391365289688, 0.2053205668926239, 0.09979691356420517, 0.11640740185976028, 0.23155182600021362, 0.4772811830043793, 0.2134055644273758, 0.3209300637245178, 0.0739695355296135], [0.046621087938547134, 0.02855776995420456, 0.11975010484457016, 0.2049850970506668, 0.16244490444660187, 0.14614170789718628, 0.03785347566008568, 0.2537410259246826, 0.3719625771045685, 0.1159287542104721, 0.23734091222286224, 0.26474830508232117, 0.04938332363963127, 0.17566856741905212, 0.034675102680921555], [0.08535599708557129, 0.01230260543525219, 0.28460273146629333, 0.3323705196380615, 0.13364574313163757, 0.14216013252735138, 0.16550986468791962, 0.36634352803230286, 0.3233327269554138, 0.13755354285240173, 0.6341029405593872, 0.1276889443397522, 0.0818048045039177, 0.2633805274963379, 0.10007897019386292], [0.014263293705880642, 0.07173046469688416, 0.01932992786169052, 0.01909404993057251, 0.16755935549736023, 0.2271488904953003, 0.1093294620513916, 0.14342457056045532, 0.0580194853246212, 0.01671113632619381, 0.03395597264170647, 0.0692841187119484, 0.07175575196743011, 0.04972841590642929, 0.12856654822826385], [0.06590985506772995, 0.1636172980070114, 0.09935098141431808, 0.20126965641975403, 0.4101002812385559, 0.21936923265457153, 0.26084569096565247, 0.3593950569629669, 0.014820259064435959, 0.05201014503836632, 0.03426084294915199, 0.38774317502975464, 0.1401163786649704, 0.3782513439655304, 0.13036324083805084], [0.05128908529877663, 0.11090300232172012, 0.24501535296440125, 0.07115167379379272, 0.3950805068016052, 0.2010982632637024, 0.08927696198225021, 0.2923780679702759, 0.11195118725299835, 0.05971711874008179, 0.14540457725524902, 0.4000069797039032, 0.2374461144208908, 0.47139719128608704, 0.10731440782546997], [0.014083221554756165, 0.029302498325705528, 0.019839908927679062, 0.019802037626504898, 0.11310776323080063, 0.014347831718623638, 0.013065088540315628, 0.0404186025261879, 0.14103254675865173, 0.01056672353297472, 0.02028844505548477, 0.4335528016090393, 0.019943613559007645, 0.08491621166467667, 0.15365199744701385], [0.04251990094780922, 0.025738505646586418, 0.19788101315498352, 0.08900192379951477, 0.20504283905029297, 0.36725619435310364, 0.05852765589952469, 0.12635937333106995, 0.07596885412931442, 0.055006030946969986, 0.1975020170211792, 0.39253395795822144, 0.2602497935295105, 0.3791850209236145, 0.11310473829507828], [0.06150972843170166, 0.049163203686475754, 0.14174170792102814, 0.13322500884532928, 0.16170991957187653, 0.21354396641254425, 0.04667104035615921, 0.26311540603637695, 0.32218027114868164, 0.0809161439538002, 0.18361496925354004, 0.23948682844638824, 0.09133663028478622, 0.25973111391067505, 0.07212682068347931], [0.12382826954126358, 0.035204268991947174, 0.3469122052192688, 0.27821084856987, 0.12485836446285248, 0.1130678728222847, 0.12963837385177612, 0.3451126217842102, 0.16417652368545532, 0.12570835649967194, 0.5000419616699219, 0.09880878776311874, 0.042446259409189224, 0.2635292708873749, 0.16834798455238342], [0.010800065472722054, 0.04851265624165535, 0.01629789173603058, 0.013155121356248856, 0.14412836730480194, 0.10944324731826782, 0.08000180870294571, 0.10409139841794968, 0.054843056946992874, 0.011575616896152496, 0.02017728053033352, 0.044063322246074677, 0.04816943034529686, 0.03936787694692612, 0.1280953288078308], [0.03501533716917038, 0.12365423142910004, 0.058643028140068054, 0.026187611743807793, 0.2106953263282776, 0.09627192467451096, 0.1373300403356552, 0.209503173828125, 0.00544273667037487, 0.010177833028137684, 0.00795654021203518, 0.17826952040195465, 0.06280092895030975, 0.2785777747631073, 0.15446779131889343], [0.055331505835056305, 0.14680130779743195, 0.22850985825061798, 0.040600359439849854, 0.2299574315547943, 0.21366852521896362, 0.10291176289319992, 0.2649042010307312, 0.07482050359249115, 0.04207760840654373, 0.11352740973234177, 0.22353075444698334, 0.2551318407058716, 0.4900997579097748, 0.11985023319721222], [0.04223596677184105, 0.14613933861255646, 0.08112313598394394, 0.04192597419023514, 0.11981905251741409, 0.18680673837661743, 0.07695262134075165, 0.14058402180671692, 0.1875196099281311, 0.05864474177360535, 0.0581248439848423, 0.23554684221744537, 0.21983209252357483, 0.1619952768087387, 0.12595340609550476]], [[0.24939602613449097, 0.0921018123626709, 0.20195554196834564, 0.25931593775749207, 0.24976609647274017, 0.08025927096605301, 0.10602997988462448, 0.08455296605825424, 0.038250602781772614, 0.34039628505706787, 0.2528480887413025, 0.17168891429901123, 0.12038858979940414, 0.16591216623783112, 0.05973837152123451], [0.04881530627608299, 0.07757209986448288, 0.080610491335392, 0.047049663960933685, 0.2744564712047577, 0.18291208148002625, 0.11781244724988937, 0.130965456366539, 0.16412131488323212, 0.049904536455869675, 0.10192018002271652, 0.46385079622268677, 0.23078110814094543, 0.23192283511161804, 0.17445482313632965], [0.11153621971607208, 0.27696484327316284, 0.0350787453353405, 0.011731116101145744, 0.08945441246032715, 0.2750371992588043, 0.07341955602169037, 0.12011690437793732, 0.026965567842125893, 0.023494159802794456, 0.015654105693101883, 0.05704642832279205, 0.11022293567657471, 0.0463077574968338, 0.1307818740606308], [0.06216026097536087, 0.123567596077919, 0.044055916368961334, 0.012494971975684166, 0.045035671442747116, 0.18137943744659424, 0.1501520872116089, 0.0996006652712822, 0.05310875549912453, 0.11289763450622559, 0.05045852065086365, 0.055306825786828995, 0.3424266576766968, 0.1600506752729416, 0.04121629521250725], [0.03470996022224426, 0.38486456871032715, 0.007671448867768049, 0.014272118918597698, 0.01295357197523117, 0.001353065250441432, 0.035229261964559555, 0.10929086059331894, 0.03641098737716675, 0.08741087466478348, 0.01870635710656643, 0.10011491179466248, 0.03142678365111351, 0.12343490868806839, 0.15971165895462036], [0.03053746558725834, 0.24113330245018005, 0.009466315619647503, 0.01980357989668846, 0.04114365205168724, 0.05523357167840004, 0.027042368426918983, 0.10979101061820984, 0.004461985547095537, 0.04689180105924606, 0.04529552906751633, 0.1364448219537735, 0.054305437952280045, 0.06579019129276276, 0.13895106315612793], [0.3289671242237091, 0.3443813920021057, 0.38217487931251526, 0.32642021775245667, 0.12515123188495636, 0.04144418612122536, 0.06740343570709229, 0.024584289640188217, 0.007359183859080076, 0.39375364780426025, 0.38123685121536255, 0.3035361170768738, 0.18788036704063416, 0.13260427117347717, 0.09976762533187866], [0.1711268573999405, 0.1900682896375656, 0.20778892934322357, 0.08847668021917343, 0.39589688181877136, 0.3955995440483093, 0.3348483741283417, 0.11133389919996262, 0.10861264914274216, 0.14033687114715576, 0.26926568150520325, 0.4846358299255371, 0.23405344784259796, 0.4343181252479553, 0.08998383581638336], [0.4154844284057617, 0.4073733687400818, 0.5541329383850098, 0.43809109926223755, 0.11503908038139343, 0.02849700301885605, 0.025097709149122238, 0.014711813069880009, 0.006424109451472759, 0.39197838306427, 0.4694826304912567, 0.17039237916469574, 0.16142874956130981, 0.19919125735759735, 0.054951149970293045], [0.24498042464256287, 0.277620404958725, 0.060333866626024246, 0.030503980815410614, 0.04090564325451851, 0.4659561812877655, 0.2110646367073059, 0.11101182550191879, 0.028219982981681824, 0.10508411377668381, 0.025386929512023926, 0.0648839995265007, 0.13676653802394867, 0.07622335106134415, 0.09164498746395111], [0.4220424294471741, 0.21296784281730652, 0.10483475774526596, 0.11319100856781006, 0.14396990835666656, 0.1309618502855301, 0.13656088709831238, 0.2097199261188507, 0.1397993415594101, 0.263439804315567, 0.10735370218753815, 0.27457332611083984, 0.26051631569862366, 0.18891198933124542, 0.10100831091403961], [0.12607140839099884, 0.08847615122795105, 0.09191321581602097, 0.06030821427702904, 0.21649383008480072, 0.10438336431980133, 0.07331530004739761, 0.1330888420343399, 0.04176999628543854, 0.06727378815412521, 0.06257567554712296, 0.21110908687114716, 0.09018781781196594, 0.09389244765043259, 0.13621515035629272], [0.062066610902547836, 0.07845254987478256, 0.24838510155677795, 0.16541223227977753, 0.16867581009864807, 0.019677892327308655, 0.021460779011249542, 0.018530650064349174, 0.023010587319731712, 0.10349667817354202, 0.16099916398525238, 0.3089703619480133, 0.08426959812641144, 0.16459643840789795, 0.06073381006717682], [0.11642084270715714, 0.11190053075551987, 0.12368596345186234, 0.04549993947148323, 0.3567850887775421, 0.06569506227970123, 0.07286660373210907, 0.03259556367993355, 0.09530685096979141, 0.19273261725902557, 0.06463074684143066, 0.7640278339385986, 0.06371455639600754, 0.1593337506055832, 0.2193848341703415], [0.11034999042749405, 0.03210863843560219, 0.010996339842677116, 0.026450032368302345, 0.051475513726472855, 0.02743532694876194, 0.3610350787639618, 0.20538736879825592, 0.017281753942370415, 0.05300014466047287, 0.012052728794515133, 0.08001075685024261, 0.0069017065688967705, 0.010893179103732109, 0.13085691630840302], [0.07615644484758377, 0.1536630541086197, 0.1253354847431183, 0.048576656728982925, 0.05276811867952347, 0.1611642986536026, 0.12317243963479996, 0.32385867834091187, 0.012925365939736366, 0.0864856168627739, 0.08918802440166473, 0.23886144161224365, 0.20351386070251465, 0.20744860172271729, 0.13318131864070892], [0.051417503505945206, 0.1600690335035324, 0.08639511466026306, 0.02997625432908535, 0.08503448963165283, 0.32695260643959045, 0.06822863221168518, 0.16364485025405884, 0.06138167902827263, 0.07786902785301208, 0.04443247988820076, 0.0585777647793293, 0.1263807862997055, 0.10769001394510269, 0.13808733224868774], [0.1321558654308319, 0.24967153370380402, 0.0761917233467102, 0.044561922550201416, 0.12028387933969498, 0.19908402860164642, 0.04708404839038849, 0.10076720267534256, 0.09921064227819443, 0.18345412611961365, 0.09404058009386063, 0.21650025248527527, 0.11625839024782181, 0.1530369222164154, 0.12011245638132095], [0.10757170617580414, 0.1042957603931427, 0.13590699434280396, 0.06331591308116913, 0.24158470332622528, 0.09161848574876785, 0.0633605495095253, 0.13977625966072083, 0.03925082087516785, 0.07121878862380981, 0.1023484393954277, 0.26378345489501953, 0.10990181565284729, 0.12030858546495438, 0.1261080652475357], [0.06512168049812317, 0.13837532699108124, 0.3250073194503784, 0.16753129661083221, 0.21647527813911438, 0.04118574038147926, 0.03336784988641739, 0.029927842319011688, 0.03334499150514603, 0.08782976865768433, 0.17631417512893677, 0.3171449303627014, 0.10520178824663162, 0.15139654278755188, 0.0914224162697792], [0.06382797658443451, 0.2566763758659363, 0.11056842654943466, 0.028001734986901283, 0.2813059389591217, 0.24806144833564758, 0.07807287573814392, 0.05373501405119896, 0.21183612942695618, 0.09658068418502808, 0.05084875971078873, 0.501965343952179, 0.06208595260977745, 0.10913741588592529, 0.26912179589271545], [0.08548272401094437, 0.017544403672218323, 0.011271107010543346, 0.022962557151913643, 0.05241750180721283, 0.02648325450718403, 0.3057800531387329, 0.19772306084632874, 0.025625178590416908, 0.03652432560920715, 0.006945622619241476, 0.05576859414577484, 0.00584550853818655, 0.008180957287549973, 0.12917736172676086], [0.03209112584590912, 0.1926622986793518, 0.09989916533231735, 0.02044818177819252, 0.04127199947834015, 0.22930434346199036, 0.09912838786840439, 0.3779822289943695, 0.007566491607576609, 0.046152934432029724, 0.04734500125050545, 0.35250937938690186, 0.10047939419746399, 0.16575956344604492, 0.13635975122451782], [0.05301084369421005, 0.1661737710237503, 0.08216799795627594, 0.025789698585867882, 0.07900767773389816, 0.3054123520851135, 0.08738221228122711, 0.17720931768417358, 0.06289011240005493, 0.06967967748641968, 0.05491774156689644, 0.02886299602687359, 0.10253670811653137, 0.09415244311094284, 0.129754438996315], [0.1895110011100769, 0.09308972954750061, 0.1887637972831726, 0.14927715063095093, 0.3653167188167572, 0.1686658412218094, 0.1126369759440422, 0.17013703286647797, 0.0685301423072815, 0.15278968214988708, 0.19327588379383087, 0.18825437128543854, 0.143904447555542, 0.143670454621315, 0.1203024610877037]], [[0.20045556128025055, 0.06346653401851654, 0.1246497705578804, 0.132145956158638, 0.18068760633468628, 0.0611145943403244, 0.3011611998081207, 0.09648064523935318, 0.3848741054534912, 0.20776434242725372, 0.09024091809988022, 0.10095226764678955, 0.05726093426346779, 0.17784324288368225, 0.06983170658349991], [0.06639314442873001, 0.03837187588214874, 0.306266725063324, 0.09758531302213669, 0.10875808447599411, 0.20901371538639069, 0.0894559919834137, 0.21620051562786102, 0.13805773854255676, 0.07912127673625946, 0.3521624505519867, 0.036526914685964584, 0.1551785171031952, 0.14622288942337036, 0.19236178696155548], [0.03379146009683609, 0.11666905134916306, 0.02791847102344036, 0.04754703491926193, 0.02039634808897972, 0.23185299336910248, 0.07985613495111465, 0.3240954875946045, 0.04561735317111015, 0.061520081013441086, 0.18156962096691132, 0.10860903561115265, 0.3409081995487213, 0.3218340575695038, 0.13103368878364563], [0.06278766691684723, 0.001863734913058579, 0.30563783645629883, 0.056017640978097916, 0.245498925447464, 0.11060530692338943, 0.09064232558012009, 0.004372697789222002, 0.007118886336684227, 0.06251134723424911, 0.17941752076148987, 0.004394095856696367, 0.11450538039207458, 0.046043287962675095, 0.021101655438542366], [0.11553236097097397, 0.0885467380285263, 0.2750205993652344, 0.21104735136032104, 0.3459762930870056, 0.07976578176021576, 0.218110129237175, 0.05760955810546875, 0.09680842608213425, 0.2662138342857361, 0.21090076863765717, 0.41520535945892334, 0.21548694372177124, 0.2248467653989792, 0.10481394827365875], [0.03112325258553028, 0.08175794035196304, 0.035110849887132645, 0.038375336676836014, 0.2468937784433365, 0.060934457927942276, 0.0843387246131897, 0.03423367813229561, 0.02026834897696972, 0.07970783859491348, 0.08959806710481644, 0.1693299561738968, 0.16057033836841583, 0.21660663187503815, 0.13329552114009857], [0.09539461880922318, 0.058681365102529526, 0.01674766093492508, 0.02866855263710022, 0.012030106969177723, 0.21465063095092773, 0.034089475870132446, 0.04479566961526871, 0.014019637368619442, 0.035355255007743835, 0.1569557934999466, 0.01038492750376463, 0.06631091982126236, 0.1547483503818512, 0.19284123182296753], [0.04954487085342407, 0.07065968960523605, 0.07275094836950302, 0.040997497737407684, 0.07946129143238068, 0.17300859093666077, 0.03222974017262459, 0.02469809167087078, 0.18557047843933105, 0.13542628288269043, 0.26776814460754395, 0.056715987622737885, 0.15973475575447083, 0.19029632210731506, 0.17610958218574524], [0.047577280551195145, 0.02606579288840294, 0.0165295097976923, 0.04137043654918671, 0.013305035419762135, 0.32835593819618225, 0.026565413922071457, 0.06772360950708389, 0.010228256694972515, 0.041277337819337845, 0.1336892545223236, 0.008326719515025616, 0.10322394222021103, 0.1976388841867447, 0.21077491343021393], [0.043893925845623016, 0.021177353337407112, 0.028366681188344955, 0.07016126066446304, 0.07573862373828888, 0.22699910402297974, 0.055615294724702835, 0.07980518788099289, 0.009269739501178265, 0.09460800141096115, 0.16427507996559143, 0.20832805335521698, 0.1427353024482727, 0.2680304944515228, 0.13907650113105774], [0.03411688283085823, 0.056632235646247864, 0.07365043461322784, 0.10934542864561081, 0.09185239672660828, 0.5077250003814697, 0.05141168087720871, 0.047258101403713226, 0.053326722234487534, 0.13365329802036285, 0.28296661376953125, 0.041020717471838, 0.08861301094293594, 0.13371184468269348, 0.11519401520490646], [0.04096442833542824, 0.07374820858240128, 0.07300861179828644, 0.10121195018291473, 0.051522452384233475, 0.3508135676383972, 0.03948133811354637, 0.047985587269067764, 0.06340529769659042, 0.06765846908092499, 0.281475692987442, 0.05536516010761261, 0.1822110116481781, 0.22272904217243195, 0.13150985538959503], [0.07982534170150757, 0.06016559898853302, 0.03820561617612839, 0.02410227432847023, 0.006901262793689966, 0.42442968487739563, 0.02364957146346569, 0.07835549116134644, 0.027230771258473396, 0.12123586237430573, 0.15446297824382782, 0.018115278333425522, 0.21087171137332916, 0.29417684674263, 0.08362340182065964], [0.05696694925427437, 0.014171368442475796, 0.06200120970606804, 0.021368764340877533, 0.012162269093096256, 0.0841592326760292, 0.03827953711152077, 0.07895056158304214, 0.01159723848104477, 0.05937046930193901, 0.023348387330770493, 0.008824712596833706, 0.13521961867809296, 0.23698511719703674, 0.03196632117033005], [0.11678174138069153, 0.8205142617225647, 0.01038320455700159, 0.023903295397758484, 0.21764065325260162, 0.2580764889717102, 0.20165181159973145, 0.2900886535644531, 0.03504627197980881, 0.10256802290678024, 0.03713424876332283, 0.7063723206520081, 0.8779962062835693, 0.8367014527320862, 0.0919082760810852], [0.038494985550642014, 0.05109047889709473, 0.07501792907714844, 0.04001014679670334, 0.021166233345866203, 0.03079657442867756, 0.01494709774851799, 0.010983827523887157, 0.0029027159325778484, 0.0995086133480072, 0.350593626499176, 0.02021479234099388, 0.34575650095939636, 0.21952421963214874, 0.05450797453522682], [0.028108511120080948, 0.08174566179513931, 0.03328564018011093, 0.03230520337820053, 0.012646276503801346, 0.1872790902853012, 0.025206655263900757, 0.06737280637025833, 0.033121660351753235, 0.08641302585601807, 0.2848047614097595, 0.059273794293403625, 0.18425194919109344, 0.15244826674461365, 0.1352420449256897], [0.07509021461009979, 0.05027765780687332, 0.23718997836112976, 0.11438266932964325, 0.11051909625530243, 0.431958943605423, 0.046987809240818024, 0.021854011341929436, 0.15366314351558685, 0.1928708851337433, 0.2900879681110382, 0.052021902054548264, 0.11538787186145782, 0.25173547863960266, 0.10233873873949051], [0.03257948160171509, 0.08023553341627121, 0.06238585337996483, 0.06856023520231247, 0.02927098423242569, 0.2968010902404785, 0.03317389637231827, 0.04758336395025253, 0.07943073660135269, 0.053982626646757126, 0.21416282653808594, 0.05025764927268028, 0.14347779750823975, 0.19969123601913452, 0.13921964168548584], [0.07817428559064865, 0.11046875268220901, 0.040724072605371475, 0.024797527119517326, 0.004808576311916113, 0.5141928791999817, 0.024754824116826057, 0.080713652074337, 0.03179122135043144, 0.12244449555873871, 0.22665926814079285, 0.013305582106113434, 0.23485711216926575, 0.323343425989151, 0.10171245783567429], [0.03765244409441948, 0.0463164821267128, 0.06456112116575241, 0.05319739878177643, 0.010156691074371338, 0.1155625581741333, 0.02458079345524311, 0.07648347318172455, 0.019683409482240677, 0.06488858163356781, 0.09342794120311737, 0.059032924473285675, 0.15581923723220825, 0.2894386053085327, 0.04157077521085739], [0.14924734830856323, 0.8862696886062622, 0.013125438243150711, 0.033269379287958145, 0.22599543631076813, 0.33975404500961304, 0.25561264157295227, 0.36481109261512756, 0.05327271297574043, 0.09902165085077286, 0.03598061203956604, 0.754990816116333, 0.9104278087615967, 0.8631682395935059, 0.10125402361154556], [0.03672042489051819, 0.12888115644454956, 0.1578092873096466, 0.056865133345127106, 0.03288109228014946, 0.1379515379667282, 0.021150214597582817, 0.013284055516123772, 0.003249341854825616, 0.08646353334188461, 0.5471532940864563, 0.0361909456551075, 0.5093809366226196, 0.39931434392929077, 0.07520455867052078], [0.03492635861039162, 0.09938696771860123, 0.028945090249180794, 0.03084651380777359, 0.012707062065601349, 0.15071596205234528, 0.029011720791459084, 0.05455483868718147, 0.03256314992904663, 0.07100401073694229, 0.2587825059890747, 0.05546442046761513, 0.17298617959022522, 0.15517692267894745, 0.13362783193588257], [0.050736088305711746, 0.10139954090118408, 0.08949553966522217, 0.0938185378909111, 0.06053004041314125, 0.18139560520648956, 0.0767659917473793, 0.11340610682964325, 0.19499026238918304, 0.11419404298067093, 0.23666803538799286, 0.05730360746383667, 0.07293370366096497, 0.11558260023593903, 0.12613430619239807]], [[0.1489560306072235, 0.2212677150964737, 0.055408962070941925, 0.03110104240477085, 0.02513720653951168, 0.07830048352479935, 0.05067736655473709, 0.06611648201942444, 0.02238955721259117, 0.03719142824411392, 0.025896798819303513, 0.04350690543651581, 0.11618120968341827, 0.08714473247528076, 0.15466241538524628], [0.002932992298156023, 0.307859867811203, 0.008187332190573215, 0.003677746979519725, 0.0005738585605286062, 0.0008406178676523268, 0.0005446207360364497, 0.00039283244404941797, 0.0009221792570315301, 0.000758469570428133, 0.003933709114789963, 0.0009352274937555194, 0.001059120986610651, 0.0020118390675634146, 0.010183396749198437], [0.37297555804252625, 0.09208715707063675, 0.16802547872066498, 0.11860792338848114, 0.08042033761739731, 0.18612971901893616, 0.45423436164855957, 0.07133221626281738, 0.13892753422260284, 0.3810507357120514, 0.291797935962677, 0.16154640913009644, 0.050885219126939774, 0.10468144714832306, 0.10335776954889297], [0.028274476528167725, 0.018124615773558617, 0.13954800367355347, 0.03560209274291992, 0.08428613841533661, 0.17491763830184937, 0.13035845756530762, 0.0214189775288105, 0.009060325101017952, 0.012400318868458271, 0.031279344111680984, 0.011209131218492985, 0.19533281028270721, 0.012452301569283009, 0.020085560157895088], [0.11180772632360458, 0.012462746351957321, 0.04844700172543526, 0.06198285147547722, 0.06685204058885574, 0.44600817561149597, 0.30352795124053955, 0.1519387811422348, 0.003835479263216257, 0.08384031802415848, 0.027865614742040634, 0.159846231341362, 0.46423590183258057, 0.09249147027730942, 0.09178084880113602], [0.04840230569243431, 0.026793736964464188, 0.1120820939540863, 0.09037120640277863, 0.2328549474477768, 0.1063276007771492, 0.14073747396469116, 0.19612964987754822, 0.1904316544532776, 0.10354755818843842, 0.10268037766218185, 0.13820117712020874, 0.3374333083629608, 0.15443934500217438, 0.12536528706550598], [0.36786824464797974, 0.056283749639987946, 0.03846094757318497, 0.07181648164987564, 0.03666122257709503, 0.04024837538599968, 0.5659748911857605, 0.2338860183954239, 0.11518415063619614, 0.3659259080886841, 0.04107162728905678, 0.012827688828110695, 0.0609581284224987, 0.02837788313627243, 0.060403015464544296], [0.0033490851055830717, 0.001678164815530181, 0.02563566155731678, 0.028815647587180138, 0.007257265504449606, 0.04370535537600517, 0.026118090376257896, 0.435838907957077, 0.005564961116760969, 0.014266176149249077, 0.018343305215239525, 0.0009297388605773449, 0.03809681162238121, 0.020595146343111992, 0.03566184639930725], [0.34718528389930725, 0.028826624155044556, 0.05378839746117592, 0.0680842474102974, 0.0254778191447258, 0.1994519978761673, 0.7739751935005188, 0.28213825821876526, 0.24756361544132233, 0.3363908529281616, 0.08445209264755249, 0.0067241075448691845, 0.09118638187646866, 0.04656682163476944, 0.0331079363822937], [0.06212884560227394, 0.013463910669088364, 0.024143628776073456, 0.025745615363121033, 0.12165382504463196, 0.04105379059910774, 0.21918880939483643, 0.12444313615560532, 0.7241542935371399, 0.2624671459197998, 0.05330171436071396, 0.026902005076408386, 0.04947282373905182, 0.06268218904733658, 0.04105047509074211], [0.23139908909797668, 0.12510670721530914, 0.062008026987314224, 0.06357982009649277, 0.21447335183620453, 0.06672460585832596, 0.5059712529182434, 0.23151132464408875, 0.3211345672607422, 0.29274967312812805, 0.07394816726446152, 0.12323616445064545, 0.33240705728530884, 0.13292434811592102, 0.0974365845322609], [0.3976813554763794, 0.24336650967597961, 0.030069073662161827, 0.04866141080856323, 0.061815883964300156, 0.023062149062752724, 0.2837987542152405, 0.10572359710931778, 0.42220908403396606, 0.47088485956192017, 0.06114182993769646, 0.05295940861105919, 0.04274435341358185, 0.033208493143320084, 0.07069624215364456], [0.6213744282722473, 0.08501708507537842, 0.08457361906766891, 0.0819045826792717, 0.02008524350821972, 0.02321169711649418, 0.5481746196746826, 0.17061969637870789, 0.19314314424991608, 0.48946020007133484, 0.08799289166927338, 0.009451461024582386, 0.1643926501274109, 0.03458939492702484, 0.0487554594874382], [0.11498570442199707, 0.014700047671794891, 0.04425002261996269, 0.027370423078536987, 0.031341005116701126, 0.11119254678487778, 0.2834031581878662, 0.24822625517845154, 0.387948602437973, 0.17188440263271332, 0.026020031422376633, 0.003112945705652237, 0.1680845320224762, 0.013143973425030708, 0.05647796019911766], [0.00710845272988081, 0.009718026034533978, 0.08296849578619003, 0.05356726795434952, 0.20372402667999268, 0.20898059010505676, 0.07373131066560745, 0.07588774710893631, 0.33318811655044556, 0.09730548411607742, 0.031877510249614716, 0.04629351943731308, 0.026428943499922752, 0.05165233090519905, 0.12934288382530212], [0.092291921377182, 0.13057716190814972, 0.11971572786569595, 0.09643372148275375, 0.0971774011850357, 0.03882397338747978, 0.30341219902038574, 0.06688009947538376, 0.5493715405464172, 0.21897412836551666, 0.10454282909631729, 0.09917838126420975, 0.19730664789676666, 0.0889393612742424, 0.0462181456387043], [0.3365032970905304, 0.06134270504117012, 0.11965256929397583, 0.08703643828630447, 0.08615697175264359, 0.01610170491039753, 0.289604127407074, 0.16905160248279572, 0.690265953540802, 0.5125291347503662, 0.11020015180110931, 0.05034353584051132, 0.04973014071583748, 0.04155145213007927, 0.06180096045136452], [0.25151577591896057, 0.0737723708152771, 0.11452356725931168, 0.07270905375480652, 0.27380475401878357, 0.046423640102148056, 0.6668940782546997, 0.60158771276474, 0.286392480134964, 0.2904633581638336, 0.07359147071838379, 0.040276750922203064, 0.2706137001514435, 0.15532110631465912, 0.051646988838911057], [0.4344438314437866, 0.2159019559621811, 0.0411386713385582, 0.059745997190475464, 0.08364511281251907, 0.02960371784865856, 0.3908357322216034, 0.17347759008407593, 0.4736940562725067, 0.5831181406974792, 0.08143209666013718, 0.05496616289019585, 0.0508774034678936, 0.03704635798931122, 0.07529113441705704], [0.6010525822639465, 0.07716702669858932, 0.12942874431610107, 0.11651009321212769, 0.029510293155908585, 0.025635747238993645, 0.564699649810791, 0.20346374809741974, 0.1942133754491806, 0.5329980254173279, 0.09726559370756149, 0.006782675161957741, 0.1884276419878006, 0.02957840822637081, 0.046941183507442474], [0.07098641246557236, 0.02088714949786663, 0.0536419078707695, 0.04874833673238754, 0.1357380896806717, 0.10192368179559708, 0.22615019977092743, 0.3848302960395813, 0.3569928705692291, 0.19976821541786194, 0.030237246304750443, 0.012232640758156776, 0.14491091668605804, 0.01217038556933403, 0.025625383481383324], [0.007031308952718973, 0.007269172929227352, 0.08423776179552078, 0.053896792232990265, 0.21268267929553986, 0.2456619292497635, 0.0817742720246315, 0.07338020205497742, 0.2872445285320282, 0.08955906331539154, 0.02503780461847782, 0.043076977133750916, 0.024157537147402763, 0.05127491056919098, 0.1281031221151352], [0.06564409285783768, 0.10634885728359222, 0.14713656902313232, 0.07514703273773193, 0.3204736113548279, 0.07143916934728622, 0.4829144775867462, 0.2612879276275635, 0.7603816986083984, 0.17889906466007233, 0.07189968973398209, 0.10938191413879395, 0.2776612341403961, 0.08681799471378326, 0.052979547530412674], [0.28806957602500916, 0.05887402966618538, 0.12616868317127228, 0.10481040924787521, 0.19247829914093018, 0.033351678401231766, 0.39873749017715454, 0.22540906071662903, 0.7029480338096619, 0.5013188719749451, 0.10523373633623123, 0.08320688456296921, 0.0816955640912056, 0.04881281033158302, 0.09282685816287994], [0.2559513747692108, 0.07615252584218979, 0.11904845386743546, 0.07934627681970596, 0.09980516135692596, 0.14371442794799805, 0.3059750497341156, 0.09035829454660416, 0.22693291306495667, 0.32864776253700256, 0.08986205607652664, 0.1614997386932373, 0.17624114453792572, 0.16325940191745758, 0.119119793176651]]]], \"bot_text\": [\"Das_\", \"Tier\", \"_\", \"\\u00fcber\", \"quer\", \"te_\", \"die_\", \"Stra\\u00dfe_\", \"nicht_\", \", _\", \"weil_\", \"es_\", \"zu_\", \"m\\u00fc\", \"de_\", \"war_\", \", _\", \"weil_\", \"es_\", \"zu_\", \"m\\u00fc\", \"de_\", \"war_\", \"._\"]}, \"all\": {\"top_text\": [\"The_\", \"animal_\", \"didn_\", \"'_\", \"t_\", \"cross_\", \"the_\", \"street_\", \"because_\", \"it_\", \"was_\", \"too_\", \"tire\", \"d_\", \"Das_\", \"Tier\", \"_\", \"\\u00fcber\", \"quer\", \"te_\", \"die_\", \"Stra\\u00dfe_\", \"nicht_\", \", _\", \"weil_\", \"es_\", \"zu_\", \"m\\u00fc\", \"de_\", \"war_\", \", _\", \"weil_\", \"es_\", \"zu_\", \"m\\u00fc\", \"de_\", \"war_\", \"._\"], \"att\": [[[[0.04540494084358215, 0.009098929353058338, 0.06841860711574554, 0.050027038902044296, 0.1867244392633438, 0.20893266797065735, 0.15536439418792725, 0.2501838803291321, 0.03253718465566635, 0.045193806290626526, 0.01405471283942461, 0.15126678347587585, 0.5554144382476807, 0.07120772451162338, 0.21479088068008423, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.010880604386329651, 0.008569094352424145, 0.3644530475139618, 0.032524824142456055, 0.15862980484962463, 0.2895345985889435, 0.007411073427647352, 0.03074379824101925, 0.23678991198539734, 0.04092710092663765, 0.21633881330490112, 0.10217994451522827, 0.5741018652915955, 0.08794906735420227, 0.15811748802661896, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1548197716474533, 0.04407857358455658, 0.04267416149377823, 0.14390510320663452, 0.39150071144104004, 0.10470721870660782, 0.21010224521160126, 0.37398451566696167, 0.24677534401416779, 0.3071460425853729, 0.12511251866817474, 0.37053829431533813, 0.34731435775756836, 0.21468856930732727, 0.22426171600818634, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01666487753391266, 0.070415198802948, 0.13558338582515717, 0.030082950368523598, 0.17114414274692535, 0.20995233952999115, 0.018852930516004562, 0.2688913345336914, 0.024380644783377647, 0.01614876091480255, 0.058318838477134705, 0.003357462352141738, 0.22233186662197113, 0.08606056123971939, 0.08522026240825653, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.26702794432640076, 0.10013092309236526, 0.15535299479961395, 0.01822819747030735, 0.19259323179721832, 0.1620739996433258, 0.06925511360168457, 0.14121465384960175, 0.30160874128341675, 0.138941690325737, 0.14571446180343628, 0.1845642775297165, 0.3172887861728668, 0.1378965824842453, 0.15321676433086395, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05774107202887535, 0.08979255706071854, 0.15777261555194855, 0.0986839085817337, 0.04042482376098633, 0.02364284358918667, 0.006265458185225725, 0.20312650501728058, 0.04589210823178291, 0.2705432176589966, 0.29482388496398926, 0.25277185440063477, 0.21941334009170532, 0.09023746848106384, 0.12374064326286316, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10808208584785461, 0.08377770334482193, 0.3031982481479645, 0.08575166761875153, 0.1659224033355713, 0.02410510927438736, 0.024052061140537262, 0.06346622854471207, 0.012278172187507153, 0.033475130796432495, 0.02865537814795971, 0.2309909611940384, 0.5272806286811829, 0.058207638561725616, 0.12589795887470245, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2848440408706665, 0.04557379335165024, 0.07043055444955826, 0.13887976109981537, 0.25104182958602905, 0.08729252219200134, 0.03900376707315445, 0.06159999966621399, 0.07028467953205109, 0.1360185593366623, 0.12163159996271133, 0.4339398145675659, 0.18035274744033813, 0.13636742532253265, 0.35040098428726196, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03364454582333565, 0.06385143101215363, 0.4650610089302063, 0.13847006857395172, 0.12132523953914642, 0.23606915771961212, 0.02828356996178627, 0.17786316573619843, 0.0068073878064751625, 0.0032905752304941416, 0.04716186597943306, 0.060036350041627884, 0.5867005586624146, 0.23594366014003754, 0.05739189311861992, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04961356148123741, 0.4571499228477478, 0.32633671164512634, 0.044803813099861145, 0.12193554639816284, 0.15620054304599762, 0.031114954501390457, 0.37925899028778076, 0.023853085935115814, 0.007363635115325451, 0.0625552162528038, 0.04359081760048866, 0.12771400809288025, 0.10945692658424377, 0.03218715265393257, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.054336514323949814, 0.12682472169399261, 0.28572455048561096, 0.7098703384399414, 0.04356186464428902, 0.036012813448905945, 0.12616953253746033, 0.12438997626304626, 0.06097114831209183, 0.011340769939124584, 0.00453603221103549, 0.02511424943804741, 0.15918391942977905, 0.004009802360087633, 0.1337292641401291, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.029656492173671722, 0.11861541867256165, 0.25968441367149353, 0.6952800154685974, 0.06073199212551117, 0.3734285235404968, 0.030824951827526093, 0.09641394764184952, 0.0529148206114769, 0.01715172454714775, 0.01323915645480156, 0.055627286434173584, 0.11593649536371231, 0.04441850632429123, 0.04630020260810852, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10554661601781845, 0.6362442970275879, 0.6959939002990723, 0.018170323222875595, 0.40134888887405396, 0.15823723375797272, 0.1629355400800705, 0.11358990520238876, 0.24731940031051636, 0.23558683693408966, 0.07505767047405243, 0.03725680336356163, 0.014009351842105389, 0.03713200241327286, 0.09585387259721756, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4055319130420685, 0.2534714341163635, 0.44874629378318787, 0.14194901287555695, 0.3008168041706085, 0.20029903948307037, 0.07248799502849579, 0.26174047589302063, 0.1826024055480957, 0.0982341319322586, 0.09884719550609589, 0.22728654742240906, 0.04277953878045082, 0.06280668079853058, 0.09454112499952316, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.025013893842697144, 0.013348683714866638, 0.22353146970272064, 0.0037027201615273952, 0.14888618886470795, 0.22346094250679016, 0.021921563893556595, 0.6342950463294983, 0.03356323391199112, 0.06236502528190613, 0.03522828221321106, 0.17797930538654327, 0.04731723666191101, 0.06786928325891495, 0.042550042271614075, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01107952743768692, 0.002038179198279977, 0.02572617679834366, 0.043437324464321136, 0.026865433901548386, 0.008821134455502033, 0.05896050110459328, 0.006038360297679901, 0.05802087485790253, 0.05262080207467079, 0.021981995552778244, 0.01655607670545578, 0.007265332620590925, 0.017941446974873543, 0.19668635725975037, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4201550781726837, 0.0003083523770328611, 0.003427971852943301, 0.027074502781033516, 0.0025770263746380806, 0.0006525526405312121, 0.0672224909067154, 0.0006329934694804251, 0.002376251621171832, 0.007315297145396471, 0.0018543159822002053, 0.0002170451043639332, 5.486799182108371e-06, 8.465739665552974e-05, 0.018722370266914368, 0.33067038655281067, 0.02820705994963646, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [6.826388562330976e-05, 0.41254693269729614, 8.318798791151494e-05, 0.00021303755056578666, 2.6623651137924753e-05, 1.3030116861045826e-06, 3.3524677292007254e-06, 9.95700816019962e-07, 0.00025696202646940947, 0.00021154701244086027, 4.0387480112258345e-05, 7.382633339148015e-05, 0.0001871670683613047, 0.0001393109851051122, 0.00044668230111710727, 0.43891066312789917, 0.3106566071510315, 0.006947982590645552, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0012913167010992765, 0.46178945899009705, 0.0011929792817682028, 0.0014885100536048412, 0.001382660586386919, 0.00010778238356579095, 4.841455302084796e-05, 4.8626650823280215e-05, 0.0007912410655990243, 0.0019299217965453863, 0.0002972490037791431, 0.0004315593687351793, 0.013707359321415424, 0.0025058358442038298, 0.00208207662217319, 0.8740342259407043, 0.6547167897224426, 0.0062981778755784035, 0.46666401624679565, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0008573953527957201, 5.803010481031379e-06, 0.0034995940513908863, 0.007113253697752953, 4.1040249925572425e-05, 0.48505696654319763, 0.0009781911503523588, 2.57480514846975e-05, 0.0006811833591200411, 0.011991027742624283, 0.013829604722559452, 0.02649468183517456, 0.018967876210808754, 0.008940043859183788, 0.0023627132177352905, 0.009682492353022099, 0.17458303272724152, 0.7120969891548157, 0.10496775060892105, 0.0038010317366570234, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [3.2793446735013276e-05, 4.91645641886862e-06, 0.0003670089063234627, 0.0005689052632078528, 0.0004337447171565145, 0.6979628205299377, 0.00025133590679615736, 1.3211038094596006e-05, 0.001040837960317731, 0.0008422345272265375, 0.00011131400242447853, 0.0007033413276076317, 0.00044049491407349706, 0.0004404923238325864, 0.00032976132933981717, 0.31054121255874634, 0.41146165132522583, 0.4573209881782532, 0.639615535736084, 0.038498248904943466, 0.06232544779777527, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.002877118531614542, 0.0015123215271160007, 0.21683953702449799, 0.042356427758932114, 0.09360139071941376, 0.7325531840324402, 0.007687804754823446, 0.0004983373219147325, 0.0008397439960390329, 0.018263472244143486, 0.01633409783244133, 0.06572946161031723, 0.029279880225658417, 0.13710656762123108, 0.013406738638877869, 0.2996446192264557, 0.18095439672470093, 0.8072441220283508, 0.6008384227752686, 0.045412980020046234, 0.09029265493154526, 0.15878555178642273, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09384340792894363, 0.002295592101290822, 0.05245966836810112, 0.10398446023464203, 0.13232196867465973, 0.2621823251247406, 0.7299563884735107, 0.01621837355196476, 0.008298774249851704, 0.019108427688479424, 0.013038183562457561, 0.008606976829469204, 0.0014156820252537727, 0.008462491445243359, 0.08448491245508194, 0.07671086490154266, 0.13175785541534424, 0.032809216529130936, 0.06887537240982056, 0.32570284605026245, 0.22846734523773193, 0.06983717530965805, 0.07415641844272614, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [7.994164479896426e-05, 9.660106115916278e-06, 1.3390360436460469e-05, 0.0009496311540715396, 7.498388185922522e-06, 0.0023292596451938152, 0.0033705621026456356, 0.45610299706459045, 0.00048403104301542044, 0.0003956609289161861, 6.013430538587272e-05, 1.5610943592037074e-05, 4.899038231087616e-06, 1.0044974260381423e-05, 0.0011326958192512393, 0.4443431496620178, 0.2924090623855591, 0.09237049520015717, 0.07077033072710037, 0.05661908909678459, 0.1886560618877411, 0.5792031288146973, 0.23326165974140167, 0.024399278685450554, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0021254755556583405, 0.025354469195008278, 0.0505821667611599, 0.04718977212905884, 0.3544465899467468, 0.27984359860420227, 0.10468283295631409, 0.03827415779232979, 0.0065247067250311375, 0.003615353489294648, 0.001024437602609396, 0.02404061146080494, 0.00031744904117658734, 0.011979974806308746, 0.06911104917526245, 0.0045473226346075535, 0.015263181179761887, 0.11153102666139603, 0.01091472152620554, 0.07137833535671234, 0.14599360525608063, 0.24649137258529663, 0.2676219940185547, 0.14942915737628937, 0.03359955921769142, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06793052703142166, 0.04423084855079651, 0.009074175730347633, 0.010606715455651283, 0.023761747404932976, 0.06765440851449966, 0.048715878278017044, 0.13498826324939728, 0.15846557915210724, 0.01835249364376068, 0.0033974519465118647, 0.011923078447580338, 0.0035463334061205387, 0.036997705698013306, 0.15195232629776, 0.0021246292162686586, 0.019146723672747612, 0.0190261360257864, 0.004887872841209173, 0.032842181622982025, 0.009469296783208847, 0.015122202225029469, 0.056959331035614014, 0.014146327041089535, 0.2864534854888916, 0.028167642652988434, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00013637961819767952, 0.00010623007256072015, 0.00015417735266964883, 0.00014589299098588526, 0.0007127521676011384, 0.0008950252668000758, 0.00038585966103710234, 0.002901369472965598, 0.34460243582725525, 0.00040915730642154813, 0.00017379666678607464, 9.334777860203758e-05, 0.0002283527428517118, 0.0001650981866987422, 0.0021401161793619394, 0.007321672048419714, 0.06949152052402496, 0.18409577012062073, 0.05168240889906883, 0.5332358479499817, 0.12983477115631104, 0.020923368632793427, 0.015086837112903595, 0.05491120368242264, 0.38865622878074646, 0.036598365753889084, 0.02645716816186905, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03951041400432587, 0.015644539147615433, 0.002765331417322159, 0.020979223772883415, 0.001914863707497716, 0.049360573291778564, 0.010446744039654732, 0.06006397679448128, 0.18512527644634247, 0.5769777894020081, 0.07455664873123169, 0.016840822994709015, 0.21517987549304962, 0.030672460794448853, 0.04319411888718605, 0.004608431365340948, 0.07759333401918411, 0.05611182749271393, 0.031112710013985634, 0.06043193116784096, 0.023203425109386444, 0.01299421489238739, 0.011212858371436596, 0.2615091800689697, 0.5089370608329773, 0.22289350628852844, 0.10276756435632706, 0.03959360718727112, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0012064727488905191, 0.0013226938899606466, 0.002064700936898589, 0.008003294467926025, 0.002116014016792178, 0.0028530799318104982, 0.006337625440210104, 0.0002913604548666626, 0.0004794643900822848, 0.0026383439544588327, 0.0038926906418055296, 0.3737375736236572, 0.002772320294752717, 0.007620541378855705, 0.003997606225311756, 0.012221934273838997, 0.040381401777267456, 0.0694599524140358, 0.0800129845738411, 0.023234205320477486, 0.003881127340719104, 0.03062801994383335, 0.024260450154542923, 0.012832778505980968, 0.01656900905072689, 0.2333584874868393, 0.3572527766227722, 0.0072386497631669044, 0.014752739109098911, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [1.0432314411445986e-05, 4.745730166177964e-06, 1.672162215982098e-05, 2.360623693675734e-05, 4.496370820561424e-06, 1.767691173881758e-06, 4.21794857174973e-06, 1.7029789205480483e-06, 2.8430429665604606e-05, 7.409282261505723e-05, 0.00010478614422027022, 0.00017224416660610586, 0.480630487203598, 0.017292670905590057, 3.8113743357826024e-05, 0.09144259989261627, 0.1256924569606781, 0.6557105779647827, 0.1641494482755661, 0.04417502135038376, 0.42902442812919617, 0.377028226852417, 0.1956152766942978, 0.27481555938720703, 0.37677863240242004, 0.4323487877845764, 0.6219720244407654, 0.3997260332107544, 0.1145903542637825, 0.041462015360593796, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00031966043752618134, 7.799067680025473e-05, 0.0005293181748129427, 0.0002383182873018086, 6.09634407737758e-05, 1.622732997930143e-05, 0.0001254813396371901, 4.548055585473776e-05, 0.0002202334435423836, 0.0014038329245522618, 0.008373874239623547, 0.0005300238262861967, 0.8584288358688354, 0.0721927285194397, 0.0012385909212753177, 0.5997433662414551, 0.1045081838965416, 0.10960735380649567, 0.047688476741313934, 0.31575047969818115, 0.1532202959060669, 0.4197675585746765, 0.16546213626861572, 0.31973955035209656, 0.23332525789737701, 0.15541672706604004, 0.05988143011927605, 0.5733460187911987, 0.8565582036972046, 0.009604076854884624, 0.030047349631786346, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.008336205966770649, 0.000929497298784554, 0.060522519052028656, 0.02858084999024868, 0.004865946713835001, 0.19429318606853485, 0.006222299765795469, 0.00020022530225105584, 0.03241097182035446, 0.2199898362159729, 0.40489089488983154, 0.12284909188747406, 0.04783688485622406, 0.16652296483516693, 0.03165041282773018, 0.02339007519185543, 0.01581897959113121, 0.02374129369854927, 0.02252129279077053, 0.08995510637760162, 0.0626068115234375, 0.27313846349716187, 0.036778680980205536, 0.22608895599842072, 0.06801939755678177, 0.035735905170440674, 0.022851483896374702, 0.06078701093792915, 0.42404335737228394, 0.41984546184539795, 0.08353053033351898, 0.058427464216947556, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06735408306121826, 0.02395833097398281, 0.022876637056469917, 0.059418935328722, 0.020556019619107246, 0.006657767109572887, 0.01686989888548851, 0.03750348463654518, 0.0929105281829834, 0.11066772043704987, 0.07383746653795242, 0.04306775704026222, 0.1764260083436966, 0.2488536387681961, 0.14264866709709167, 0.034203190356492996, 0.23458202183246613, 0.15632590651512146, 0.02520577609539032, 0.26413342356681824, 0.06292548030614853, 0.06378099322319031, 0.08676797896623611, 0.02988903410732746, 0.3430734872817993, 0.007843950763344765, 0.03405369073152542, 0.01887335814535618, 0.39618176221847534, 0.2528276741504669, 0.10531513392925262, 0.12583006918430328, 0.09389571845531464, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00023218609567265958, 9.724824485601857e-05, 0.00017837552877608687, 0.000249945733230561, 0.00043016509152948856, 0.0002728255931288004, 0.0002596308768261224, 0.0021448382176458836, 0.33870813250541687, 0.0012523159384727478, 0.0004828754754271358, 7.525486580561846e-05, 0.001232807757332921, 0.00022845527564641088, 0.0029908884316682816, 0.009769688360393047, 0.056299567222595215, 0.11172951757907867, 0.02802591770887375, 0.3647110164165497, 0.09813904017210007, 0.016619421541690826, 0.006417513824999332, 0.016537560150027275, 0.15495160222053528, 0.023067951202392578, 0.011397394351661205, 0.029141509905457497, 0.0527399443089962, 0.2784731984138489, 0.059669919312000275, 0.5969582796096802, 0.09549567103385925, 0.03235183656215668, NaN, NaN, NaN, NaN, NaN, NaN], [0.044313203543424606, 0.014693659730255604, 0.001713237608782947, 0.01787775754928589, 0.001054717693477869, 0.03111616149544716, 0.005932849366217852, 0.035437386482954025, 0.10908837616443634, 0.6214090585708618, 0.11623460799455643, 0.018710769712924957, 0.26884767413139343, 0.036007944494485855, 0.04555344209074974, 0.00987912341952324, 0.12349259853363037, 0.037169262766838074, 0.01944275200366974, 0.06324917078018188, 0.02598830871284008, 0.020618943497538567, 0.009103300981223583, 0.1360517293214798, 0.09789924323558807, 0.06809242814779282, 0.12332575768232346, 0.034675393253564835, 0.16954950988292694, 0.010956126265227795, 0.11111389100551605, 0.1871008574962616, 0.2434563934803009, 0.10274684429168701, 0.0379486046731472, NaN, NaN, NaN, NaN, NaN], [0.0014647350180894136, 0.0016486160457134247, 0.001705971430055797, 0.008203698322176933, 0.0011827786220237613, 0.001036314177326858, 0.004107706248760223, 0.00018337460642214864, 0.0005908485618419945, 0.004427316598594189, 0.0075510423630476, 0.37528446316719055, 0.0045065670274198055, 0.01084148045629263, 0.0047609396278858185, 0.010987702757120132, 0.03791751340031624, 0.03792046010494232, 0.0400051474571228, 0.008841714821755886, 0.002161285374313593, 0.031619150191545486, 0.01907121017575264, 0.0057282340712845325, 0.002385619329288602, 0.03308374434709549, 0.11032091826200485, 0.0044158026576042175, 0.05701944977045059, 0.0651637390255928, 0.027267253026366234, 0.3151875138282776, 0.17881636321544647, 0.3164456784725189, 0.005250148009508848, 0.011875288560986519, NaN, NaN, NaN, NaN], [1.1546462701517157e-05, 6.3197094277711585e-06, 1.3665205187862739e-05, 2.3049220544635318e-05, 3.1024922009237343e-06, 9.712728115118807e-07, 4.2468768697290216e-06, 1.4032799526830786e-06, 2.1501631636056118e-05, 0.00011254433775320649, 0.00014821428339928389, 0.00021640797785948962, 0.4815296530723572, 0.022970588877797127, 4.596232975018211e-05, 0.08034691959619522, 0.1792650669813156, 0.6813479661941528, 0.11697664856910706, 0.022037051618099213, 0.4362119436264038, 0.3332834541797638, 0.16648675501346588, 0.3133866786956787, 0.21180157363414764, 0.22306133806705475, 0.5634312033653259, 0.2539531886577606, 0.28583550453186035, 0.0421890914440155, 0.24185270071029663, 0.9185315370559692, 0.5444227457046509, 0.7130873799324036, 0.36675870418548584, 0.1082441657781601, 0.02894955314695835, NaN, NaN, NaN], [0.0004618540406227112, 0.00011890243331436068, 0.0008028792799450457, 0.0003817373653873801, 7.645944424439222e-05, 2.0059787857462652e-05, 0.00017321997438557446, 3.885024489136413e-05, 0.00016429855895694345, 0.0017073642229661345, 0.011983372271060944, 0.0008083870052359998, 0.8495219349861145, 0.07573292404413223, 0.0017974229995161295, 0.3316553831100464, 0.07297243922948837, 0.18084223568439484, 0.0543624572455883, 0.141310915350914, 0.15985439717769623, 0.22593949735164642, 0.09976530820131302, 0.2670679986476898, 0.12590403854846954, 0.10189743340015411, 0.06066418066620827, 0.14688965678215027, 0.6279550790786743, 0.004891595803201199, 0.013660040684044361, 0.19539086520671844, 0.13336770236492157, 0.11226529628038406, 0.4554508626461029, 0.7914823293685913, 0.007615156006067991, 0.015521766617894173, NaN, NaN], [0.00848880223929882, 0.0010204557329416275, 0.06384890526533127, 0.030244439840316772, 0.004545390605926514, 0.2111765593290329, 0.007047791499644518, 0.00020413362653926015, 0.03285042569041252, 0.2096482813358307, 0.40160003304481506, 0.12425301223993301, 0.05433715134859085, 0.2013336718082428, 0.03489448130130768, 0.010082974098622799, 0.009416572749614716, 0.026376336812973022, 0.021534079685807228, 0.041008636355400085, 0.028814975172281265, 0.09862472116947174, 0.019531887024641037, 0.1915404349565506, 0.055525705218315125, 0.03489372506737709, 0.035597167909145355, 0.017297467216849327, 0.13875839114189148, 0.18795406818389893, 0.13025526702404022, 0.03705297037959099, 0.016517892479896545, 0.028779756277799606, 0.02632485330104828, 0.36631691455841064, 0.4771501123905182, 0.10461407899856567, 0.07566797733306885, NaN], [0.018106432631611824, 0.01663283444941044, 0.006966447923332453, 0.06288447231054306, 0.008926548063755035, 0.0005806194385513663, 0.004527462646365166, 0.00047311693197116256, 0.010450053960084915, 0.008817908354103565, 0.02498125471174717, 0.02475220151245594, 0.006219316273927689, 0.034688226878643036, 0.15510374307632446, 0.00671275844797492, 0.019956005737185478, 0.15321078896522522, 0.00987993273884058, 0.1430601179599762, 0.02432059310376644, 0.007838046178221703, 0.016839532181620598, 0.017622128129005432, 0.03075602278113365, 0.01907699555158615, 0.30206096172332764, 0.010013632476329803, 0.06018203869462013, 0.19546428322792053, 0.020215312018990517, 0.04091925173997879, 0.022548291832208633, 0.26572445034980774, 0.010653333738446236, 0.1212434321641922, 0.3668496906757355, 0.1586136817932129, 0.14579400420188904, 0.04911552369594574]], [[0.1577349603176117, 0.09554319828748703, 0.02016325853765011, 0.08440300822257996, 0.33925309777259827, 0.35353752970695496, 0.49755600094795227, 0.2782062292098999, 0.2544572949409485, 0.6230229735374451, 0.04059281200170517, 0.12019311636686325, 0.2659685015678406, 0.3508304953575134, 0.10784413665533066, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.053030457347631454, 0.00926118716597557, 0.08361255377531052, 0.1587543487548828, 0.42493122816085815, 0.0713140144944191, 0.05032603442668915, 0.790120005607605, 0.4618776738643646, 0.3647898733615875, 0.20375682413578033, 0.2847990393638611, 0.20242592692375183, 0.33538198471069336, 0.174686461687088, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08703262358903885, 0.32554149627685547, 0.013934381306171417, 0.05831753462553024, 0.13550086319446564, 0.24707834422588348, 0.10738440603017807, 0.2015978991985321, 0.20393061637878418, 0.3176687955856323, 0.11071985214948654, 0.18533341586589813, 0.23293758928775787, 0.34885379672050476, 0.5850104689598083, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10977373272180557, 0.1966770738363266, 0.08552326261997223, 0.3559982180595398, 0.025181425735354424, 0.05637436732649803, 0.04466243088245392, 0.30799123644828796, 0.24855823814868927, 0.13041310012340546, 0.16531962156295776, 0.11238406598567963, 0.33737656474113464, 0.08863592892885208, 0.043888676911592484, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5166918635368347, 0.35558366775512695, 0.01755080744624138, 0.011931763030588627, 0.556053638458252, 0.21828243136405945, 0.17387567460536957, 0.11686032265424728, 0.22141756117343903, 0.6036979556083679, 0.3235246241092682, 0.21816273033618927, 0.20258961617946625, 0.7225815653800964, 0.3817636966705322, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.34899845719337463, 0.35567307472229004, 0.2643766403198242, 0.12664493918418884, 0.18397535383701324, 0.012551958672702312, 0.056629326194524765, 0.06369142234325409, 0.252005010843277, 0.3601645529270172, 0.3771168887615204, 0.4479873776435852, 0.13717319071292877, 0.6667386293411255, 0.1451762467622757, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5782451629638672, 0.6189379096031189, 0.11758852005004883, 0.3125992715358734, 0.3504111170768738, 0.10631152987480164, 0.16217094659805298, 0.04177623987197876, 0.10916820168495178, 0.3274877965450287, 0.10721725970506668, 0.11595069617033005, 0.11270644515752792, 0.32787472009658813, 0.13412055373191833, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2553749084472656, 0.5479037165641785, 0.3395489752292633, 0.13140854239463806, 0.07771788537502289, 0.06743729114532471, 0.04718935862183571, 0.022107038646936417, 0.2706955075263977, 0.06462319940328598, 0.20574931800365448, 0.08401398360729218, 0.11249610781669617, 0.20925462245941162, 0.07354141771793365, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.15992610156536102, 0.4297313988208771, 0.11996463686227798, 0.29957810044288635, 0.19940054416656494, 0.6192947030067444, 0.07005859166383743, 0.4058174192905426, 0.0451255701482296, 0.02480492927134037, 0.052432600408792496, 0.13078351318836212, 0.14195236563682556, 0.12686756253242493, 0.10959619283676147, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.13202522695064545, 0.3311104476451874, 0.12707853317260742, 0.06901858001947403, 0.13186469674110413, 0.37057942152023315, 0.1482420712709427, 0.21941475570201874, 0.1949346363544464, 0.11534072458744049, 0.011536079458892345, 0.018882060423493385, 0.16279305517673492, 0.07962523400783539, 0.11737312376499176, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0604790523648262, 0.5140921473503113, 0.37517040967941284, 0.060462601482868195, 0.14644990861415863, 0.49839717149734497, 0.08009912073612213, 0.3367377519607544, 0.0785842090845108, 0.043956201523542404, 0.0826396569609642, 0.015624956227838993, 0.10417986661195755, 0.07971351593732834, 0.018050679937005043, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10509271919727325, 0.5468136072158813, 0.2136838436126709, 0.13898353278636932, 0.11654751002788544, 0.1982421725988388, 0.03731672093272209, 0.5618436336517334, 0.37511539459228516, 0.015668287873268127, 0.07859797775745392, 0.026544239372015, 0.11879771202802658, 0.051024846732616425, 0.03191406652331352, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2583395540714264, 0.306291788816452, 0.15283380448818207, 0.48663485050201416, 0.24239543080329895, 0.6472541093826294, 0.11895711719989777, 0.7050262093544006, 0.43789902329444885, 0.07257331907749176, 0.1529301553964615, 0.07237879186868668, 0.029207568615674973, 0.031136667355895042, 0.04320577159523964, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.37997886538505554, 0.3090342879295349, 0.09529577195644379, 0.06091787666082382, 0.5611693859100342, 0.5351426005363464, 0.5250707268714905, 0.4058402180671692, 0.08284364640712738, 0.7192233204841614, 0.12988585233688354, 0.24924960732460022, 0.016598563641309738, 0.6531801819801331, 0.22117754817008972, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.31734058260917664, 0.02799793891608715, 0.08435621112585068, 0.4273812472820282, 0.37900310754776, 0.1551857888698578, 0.12445898354053497, 0.02975497953593731, 0.13922178745269775, 0.25836795568466187, 0.3142063617706299, 0.5329877138137817, 0.020000692456960678, 0.19246473908424377, 0.34441179037094116, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.011485431343317032, 0.057214245200157166, 0.11445975303649902, 0.035292237997055054, 0.17235025763511658, 0.21079879999160767, 0.08683252334594727, 0.33144259452819824, 0.2781406342983246, 0.07864350080490112, 0.10017280280590057, 0.0828540250658989, 0.17722147703170776, 0.21101748943328857, 0.15805292129516602, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.041519034653902054, 0.11474552005529404, 0.04909001290798187, 0.1299373209476471, 0.06295691430568695, 0.0239214189350605, 0.22038953006267548, 0.6809458136558533, 0.03295678645372391, 0.34942832589149475, 0.1847512274980545, 0.22206875681877136, 0.13646042346954346, 0.277276873588562, 0.1334262192249298, 0.00017037145153153688, 0.1837475299835205, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0764331966638565, 0.004937899298965931, 0.049346037209033966, 0.05165911093354225, 0.051789041608572006, 0.11632981896400452, 0.3382570743560791, 0.21805666387081146, 0.5269062519073486, 0.05627245828509331, 0.1284114420413971, 0.3053610324859619, 0.058564696460962296, 0.14431920647621155, 0.19175130128860474, 4.619961600837996e-06, 0.00011092388740507886, 0.19595862925052643, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08274618536233902, 0.009897814132273197, 0.07511309534311295, 0.03663979470729828, 0.16369661688804626, 0.04579350724816322, 0.04420214146375656, 0.06866282969713211, 0.17000554502010345, 0.09549596160650253, 0.07313749194145203, 0.06223462149500847, 0.11603321135044098, 0.07143211364746094, 0.2059532254934311, 7.402049959637225e-07, 0.0014410031726583838, 0.15330694615840912, 0.0009438465931452811, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.41769060492515564, 0.07210511714220047, 0.40716952085494995, 0.22363832592964172, 0.48781970143318176, 0.015007800422608852, 0.4504202902317047, 0.4675638973712921, 0.24936619400978088, 0.5447031855583191, 0.4296078681945801, 0.07025930285453796, 0.1902965009212494, 0.3567025065422058, 0.12464861571788788, 6.564930572494632e-07, 1.2471617083065212e-05, 0.0012651559663936496, 1.2094314115529414e-05, 0.2683168947696686, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3858333230018616, 0.06937354803085327, 0.5601253509521484, 0.30969470739364624, 0.36272186040878296, 0.005774383433163166, 0.16290897130966187, 0.16338182985782623, 0.1734752655029297, 0.10127251595258713, 0.6812319159507751, 0.35078492760658264, 0.26554787158966064, 0.3089393675327301, 0.12310608476400375, 3.960849710438197e-07, 2.835777740983758e-05, 0.0015905762556940317, 5.72201497561764e-05, 0.20671997964382172, 0.03618929535150528, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.047016799449920654, 0.04388514533638954, 0.010725832544267178, 0.029561294242739677, 0.04913409426808357, 0.007112162187695503, 0.045616600662469864, 0.09563170373439789, 0.021758677437901497, 0.05606407672166824, 0.023780539631843567, 0.2586848735809326, 0.1317795366048813, 0.13214319944381714, 0.18490085005760193, 3.613545777625404e-05, 4.069158967467956e-05, 0.0019799659494310617, 4.598083614837378e-05, 0.28016433119773865, 0.1021510660648346, 0.0019787675701081753, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.024271933361887932, 0.10952932387590408, 0.01092300284653902, 0.005798409227281809, 0.03478696197271347, 0.015390553511679173, 0.005925341974943876, 0.04537563398480415, 0.00714160455390811, 0.005484140943735838, 0.00704369880259037, 0.04858299717307091, 0.06617175042629242, 0.13874217867851257, 0.17208275198936462, 0.03414154052734375, 0.018152736127376556, 0.002861178945749998, 0.0031036457512527704, 0.2743661403656006, 0.08905426412820816, 0.058365415781736374, 0.2834230065345764, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1448126882314682, 0.16020630300045013, 0.02696153335273266, 0.06902630627155304, 0.03837759047746658, 0.07682601362466812, 0.15773272514343262, 0.005734406877309084, 0.16041570901870728, 0.10849703103303909, 0.08964504301548004, 0.4313186705112457, 0.12084108591079712, 0.20548132061958313, 0.1913137137889862, 0.0001288916973862797, 0.0019113116431981325, 0.0011359998025000095, 2.5460678443778306e-05, 0.0018093753606081009, 0.008086470887064934, 0.005666371434926987, 0.0014489549212157726, 0.27176737785339355, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03147122263908386, 0.06498080492019653, 0.03835386037826538, 0.021906379610300064, 0.004580754786729813, 0.08777225762605667, 0.06548282504081726, 0.0501156747341156, 0.09960248321294785, 0.05812418833374977, 0.04425663501024246, 0.12932318449020386, 0.040425609797239304, 0.10523593425750732, 0.20731014013290405, 0.0013363973703235388, 0.015213730745017529, 0.019847076386213303, 0.0016770424554124475, 0.6085457801818848, 0.051846977323293686, 0.06904839724302292, 0.023163089528679848, 0.0024616841692477465, 0.4075135886669159, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03185653313994408, 0.014990762807428837, 0.012671640142798424, 0.014554454945027828, 0.005096337758004665, 0.025306345894932747, 0.015522593632340431, 0.012109486386179924, 0.014945329166948795, 0.0111803337931633, 0.010501275770366192, 0.010505528189241886, 0.013426732271909714, 0.01895906589925289, 0.16498495638370514, 1.5705205441918224e-05, 0.00011942459968850017, 3.308789018774405e-05, 0.00047703171730972826, 1.5581523257424124e-05, 3.566192026482895e-05, 0.000621139828581363, 0.002513762330636382, 0.0013953398447483778, 0.001656065694987774, 0.6708395481109619, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05249502509832382, 0.3800218403339386, 0.048091597855091095, 0.01820666529238224, 0.10161028057336807, 0.18240275979042053, 0.03954629600048065, 0.08666953444480896, 0.00239415536634624, 0.05545663461089134, 0.11899324506521225, 0.03552442044019699, 0.037884730845689774, 0.08727249503135681, 0.23120805621147156, 0.0009777048835530877, 0.006719581317156553, 0.017090875655412674, 0.007835427299141884, 0.0003081739123445004, 0.0027951891534030437, 0.0031432590913027525, 0.011542102321982384, 0.01903962530195713, 0.032312098890542984, 0.23448777198791504, 0.18604722619056702, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06818026304244995, 0.06384387612342834, 0.013627037405967712, 0.017488455399870872, 0.04112459346652031, 0.37204819917678833, 0.2269488275051117, 0.050778258591890335, 0.07564288377761841, 0.002337054116651416, 0.03256889060139656, 0.017944803461432457, 0.02268233709037304, 0.05458826571702957, 0.17415940761566162, 0.0010771078523248434, 0.00013067253166809678, 0.0004810431564692408, 0.0005832655006088316, 0.27172601222991943, 0.023587899282574654, 0.0011203349567949772, 0.0001570776366861537, 3.2636336982250214e-05, 0.008125105872750282, 0.3860749900341034, 0.011222672648727894, 0.4488545358181, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3350563049316406, 0.14807114005088806, 0.16856855154037476, 0.0634150505065918, 0.6115131974220276, 0.8617944717407227, 0.4784194529056549, 0.271447092294693, 0.44727417826652527, 0.03638387843966484, 0.0791390910744667, 0.0010650564217939973, 0.10882135480642319, 0.07249648869037628, 0.16217634081840515, 0.0018897228874266148, 0.00010004806244978681, 0.040837980806827545, 0.0009045379119925201, 0.4036760926246643, 0.033945482224226, 0.0009020724683068693, 2.477952148183249e-05, 0.0006147518288344145, 2.3498352675233036e-05, 0.0003015661786776036, 0.00019162058015353978, 0.0013656887458637357, 0.9207848906517029, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.6229478120803833, 0.11473710834980011, 0.9313594102859497, 0.6977004408836365, 0.7760463953018188, 0.5547962784767151, 0.2850213646888733, 0.12024195492267609, 0.6867435574531555, 0.3715392053127289, 0.5383524894714355, 0.04410971701145172, 0.001209885231219232, 0.03505939990282059, 0.07057712972164154, 3.0049262932152487e-05, 0.00032340767211280763, 0.0004620190302375704, 1.456133759347722e-05, 0.4214256703853607, 0.00038119935197755694, 2.2086916942498647e-05, 5.437946310848929e-05, 0.0005922063137404621, 0.0002251591213280335, 4.171442924416624e-05, 0.0011568808695301414, 6.667344860034063e-05, 0.004539569839835167, 0.07099039107561111, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12039526551961899, 0.15183398127555847, 0.23466746509075165, 0.07534174621105194, 0.09489727020263672, 0.12723755836486816, 0.06088049337267876, 0.06659132242202759, 0.24534910917282104, 0.08624531328678131, 0.05703657865524292, 0.031156441196799278, 0.0026320687029510736, 0.016870809718966484, 0.16136524081230164, 0.0001142411565524526, 0.001007341779768467, 0.5582761764526367, 0.0006983705679886043, 0.04208780825138092, 0.07311324775218964, 0.011010478250682354, 0.00018356108921580017, 0.11227726191282272, 1.5535662896581925e-05, 7.865564111853018e-05, 8.497068483848125e-05, 0.007107958197593689, 0.04726947844028473, 0.03816111385822296, 0.7400538921356201, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.024926312267780304, 0.055538877844810486, 0.0035579875111579895, 0.006728078704327345, 0.10179352015256882, 0.12386216968297958, 0.08368373662233353, 0.17138876020908356, 0.13290183246135712, 0.025975322350859642, 0.0007942751399241388, 0.08679928630590439, 0.006940893363207579, 0.006668384652584791, 0.2167840152978897, 9.270196460420266e-05, 0.00014002913667354733, 0.006266205105930567, 8.287983655463904e-05, 0.029540851712226868, 0.019505193457007408, 0.0002005908900173381, 0.0002361711667617783, 0.002089217072352767, 0.0007247799658216536, 0.0003387654141988605, 3.3522373996675014e-05, 0.00015295531193260103, 0.005682599265128374, 0.01914886385202408, 0.006167547311633825, 0.6065680980682373, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03079223819077015, 0.008776835165917873, 0.025623725727200508, 0.02996702678501606, 0.076390340924263, 0.11722294241189957, 0.03722265735268593, 0.06894396245479584, 0.023492204025387764, 0.02721765637397766, 0.02432498149573803, 0.009946721605956554, 0.02367306686937809, 0.02709045261144638, 0.15603508055210114, 0.017243418842554092, 0.0717378556728363, 0.015470567159354687, 0.14577892422676086, 0.003815611358731985, 0.01656431145966053, 0.21609994769096375, 0.24452562630176544, 0.07360902428627014, 0.020440302789211273, 0.9522358775138855, 0.0012982342159375548, 0.00034142163349315524, 4.905217429040931e-05, 0.0002677988959476352, 0.0020047405268996954, 0.013444142416119576, 0.5238149166107178, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.050754088908433914, 0.38707080483436584, 0.056088101118803024, 0.022330837324261665, 0.19594413042068481, 0.356031596660614, 0.05540256202220917, 0.17031489312648773, 0.002592364326119423, 0.0904960110783577, 0.17009596526622772, 0.02688765898346901, 0.05266827344894409, 0.09536514431238174, 0.2306852787733078, 0.006589227356016636, 0.025933612138032913, 0.05151839554309845, 0.019538801163434982, 0.000567624403629452, 0.011064885184168816, 0.018599001690745354, 0.0389220230281353, 0.03263486549258232, 0.03920944407582283, 0.309482604265213, 0.18455958366394043, 0.0028949796687811613, 0.0009189100819639862, 0.01304793544113636, 0.01903691701591015, 0.0013186958385631442, 0.1459255963563919, 0.2617945969104767, NaN, NaN, NaN, NaN, NaN, NaN], [0.052731066942214966, 0.07647765427827835, 0.009669344872236252, 0.013631273992359638, 0.037963252514600754, 0.40968915820121765, 0.1877974420785904, 0.06287717074155807, 0.06925270706415176, 0.0021469732746481895, 0.03106895461678505, 0.02147551439702511, 0.022071314975619316, 0.058794401586055756, 0.17150944471359253, 0.000940846570301801, 6.996696902206168e-05, 0.0001185448418254964, 0.00013115631008986384, 0.04620806872844696, 0.009408986195921898, 0.0010798430303111672, 0.00010642426059348509, 1.4586596989829559e-05, 0.0008147742482833564, 0.049950405955314636, 0.0020658469293266535, 0.020368386059999466, 0.0015965981874614954, 0.0005227082292549312, 8.089001494226977e-05, 0.42970454692840576, 0.3893451988697052, 0.006195466499775648, 0.2630486488342285, NaN, NaN, NaN, NaN, NaN], [0.2993965446949005, 0.1887350082397461, 0.17583680152893066, 0.06075390800833702, 0.6836855411529541, 0.8825634121894836, 0.44942814111709595, 0.3110062777996063, 0.6245057582855225, 0.04149743914604187, 0.08928828686475754, 0.0010537458583712578, 0.13885420560836792, 0.09175378829240799, 0.16601231694221497, 0.0015646422980353236, 5.644361226586625e-05, 0.015588155947625637, 0.0004337269929237664, 0.061090677976608276, 0.015012362040579319, 0.0009935805574059486, 3.2441483199363574e-05, 0.0006383971776813269, 7.901599929027725e-06, 0.00011085882579209283, 2.031324947893154e-05, 0.0001886440732050687, 0.1558367908000946, 2.918860081990715e-05, 0.00031420652521774173, 3.769064642256126e-05, 0.000311522075207904, 8.488001913065091e-05, 0.001447036280296743, 0.9016569256782532, NaN, NaN, NaN, NaN], [0.6222140192985535, 0.13893182575702667, 0.9335290789604187, 0.7374492883682251, 0.8253674507141113, 0.5633905529975891, 0.4091120660305023, 0.12903769314289093, 0.8090996742248535, 0.490604043006897, 0.6206711530685425, 0.06171489879488945, 0.0013746770564466715, 0.055387232452631, 0.07617512345314026, 6.329882307909429e-05, 0.0007932570297271013, 0.0008974742377176881, 3.545067738741636e-05, 0.41645264625549316, 0.0012166639789938927, 5.162824527360499e-05, 0.00016062096983660012, 0.0028807471971958876, 0.0007734368555247784, 0.0001738688733894378, 0.0017386887921020389, 8.449772576568648e-05, 0.008313576690852642, 0.04833607003092766, 5.605717160506174e-05, 0.000497612461913377, 0.00019103533122688532, 0.0018799308454617858, 0.000193181011127308, 0.010939341969788074, 0.11687301844358444, NaN, NaN, NaN], [0.1216169223189354, 0.17628714442253113, 0.21903447806835175, 0.08471400290727615, 0.12100206315517426, 0.12684285640716553, 0.060168445110321045, 0.05725802481174469, 0.204857736825943, 0.07119028270244598, 0.04997517541050911, 0.046147700399160385, 0.002665548352524638, 0.01769380457699299, 0.1595369428396225, 2.7039888664148748e-05, 0.0002653435221873224, 0.3520841896533966, 0.0011641159653663635, 0.017258664593100548, 0.13898366689682007, 0.004804374184459448, 0.0001136215214501135, 0.10132589936256409, 1.9021857951884158e-05, 0.00018713112513069063, 5.577637057285756e-05, 0.0021825090516358614, 0.016621561720967293, 0.003813497256487608, 0.05257569998502731, 7.136658678064123e-05, 0.00013083907833788544, 8.304342918563634e-05, 0.009517401456832886, 0.07102376222610474, 0.0242641419172287, 0.791592538356781, NaN, NaN], [0.02323095127940178, 0.05151251330971718, 0.002836216241121292, 0.007343180477619171, 0.11471041291952133, 0.09745588153600693, 0.08793136477470398, 0.19987791776657104, 0.2081962525844574, 0.026029428467154503, 0.0006721516838297248, 0.15218332409858704, 0.008676346391439438, 0.009503011591732502, 0.20713838934898376, 1.8426982933306135e-05, 6.735812348779291e-05, 0.005383457988500595, 0.0002568464260548353, 0.03709089383482933, 0.05173188075423241, 0.00015440442075487226, 0.00026214553508907557, 0.0031172526068985462, 0.0018413036596029997, 0.001364374067634344, 0.0001026472236844711, 0.00015940713637974113, 0.00464483629912138, 0.007250420283526182, 0.006640422623604536, 0.10042263567447662, 0.00037284562131389976, 5.502302519744262e-05, 0.00017516437219455838, 0.013823487795889378, 0.028728578239679337, 0.014491567388176918, 0.5602642297744751, NaN], [0.07751920074224472, 0.05964339151978493, 0.026831025257706642, 0.018057459965348244, 0.1489739865064621, 0.27560925483703613, 0.15271086990833282, 0.29336896538734436, 0.2548864185810089, 0.015449506230652332, 0.02643660455942154, 0.05839552357792854, 0.06659974157810211, 0.1841144859790802, 0.1324990689754486, 1.3810687960358337e-05, 0.0002572945086285472, 0.008041280321776867, 0.00040080497274175286, 0.00010326507617719471, 0.0013340600999072194, 0.00019016038277186453, 0.00019489554688334465, 0.0007417663000524044, 0.0012533330591395497, 0.0032668926287442446, 0.001072657760232687, 5.286548912408762e-05, 4.225512952871213e-07, 1.0035311788669787e-05, 2.1279807697283104e-05, 0.0006032216479070485, 0.00048016011714935303, 0.00037273563793860376, 3.447151175350882e-05, 9.715819260236458e-07, 2.8930742701049894e-05, 0.0003854547976516187, 0.005018792115151882, 0.4505775570869446]], [[0.022252710536122322, 0.017558962106704712, 0.12289869785308838, 0.01514213066548109, 0.04983796179294586, 0.160098597407341, 0.09159664064645767, 0.03634485974907875, 0.27353572845458984, 0.14908282458782196, 0.8423851132392883, 0.33708906173706055, 0.03012021631002426, 0.05972116440534592, 0.2686574459075928, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.13637107610702515, 0.02899317629635334, 0.09026061743497849, 0.22582301497459412, 0.09117049723863602, 0.19661013782024384, 0.30083417892456055, 0.13528303802013397, 0.1352328211069107, 0.18504901230335236, 0.3621358573436737, 0.504258930683136, 0.10044156759977341, 0.37106865644454956, 0.36433035135269165, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10935092717409134, 0.06271693855524063, 0.044740546494722366, 0.1709805577993393, 0.22382155060768127, 0.2615796625614166, 0.3429900109767914, 0.02677186205983162, 0.39723172783851624, 0.1559167355298996, 0.6381150484085083, 0.34350308775901794, 0.14388519525527954, 0.322640985250473, 0.07209958881139755, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.11123806983232498, 0.14550834894180298, 0.12841136753559113, 0.013620064593851566, 0.006130752619355917, 0.025231752544641495, 0.11538708955049515, 0.09429272264242172, 0.3855685293674469, 0.016912028193473816, 0.3869503438472748, 0.1961694061756134, 0.15352581441402435, 0.019190048798918724, 0.4291467070579529, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1283823847770691, 0.33987957239151, 0.06837885081768036, 0.03946131095290184, 0.03139644116163254, 0.11983324587345123, 0.12062173336744308, 0.46404916048049927, 0.24212448298931122, 0.1594262570142746, 0.4298713207244873, 0.5236353278160095, 0.2188095897436142, 0.049411591142416, 0.10146455466747284, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.010564678348600864, 0.32722386717796326, 0.19864077866077423, 0.015389330685138702, 0.0028029000386595726, 0.007416849955916405, 0.003262599464505911, 0.23795713484287262, 0.05000551417469978, 0.075996033847332, 0.049679387360811234, 0.21265098452568054, 0.2097157984972, 0.01007634773850441, 0.03895873948931694, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10390599817037582, 0.04329453781247139, 0.42168325185775757, 0.06385642290115356, 0.04340887442231178, 0.029213739559054375, 0.036663200706243515, 0.0028809772338718176, 0.19718152284622192, 0.16335125267505646, 0.6605148315429688, 0.17834524810314178, 0.08135847747325897, 0.05741032958030701, 0.24636343121528625, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.010566278360784054, 0.32608217000961304, 0.34194469451904297, 0.08201102167367935, 0.036688148975372314, 0.12155891954898834, 0.015490439720451832, 0.05858473479747772, 0.1731383204460144, 0.12207219004631042, 0.0636284351348877, 0.2239474654197693, 0.2988812327384949, 0.033257871866226196, 0.04593053460121155, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.26241976022720337, 0.0378817655146122, 0.10770448297262192, 0.11944369971752167, 0.367754727602005, 0.041288651525974274, 0.25914207100868225, 0.061461515724658966, 0.061867646872997284, 0.08977923542261124, 0.03797370195388794, 0.2101898193359375, 0.035329420119524, 0.38835543394088745, 0.3324989080429077, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3753410875797272, 0.031615160405635834, 0.1074504628777504, 0.07966858148574829, 0.16393397748470306, 0.01204571221023798, 0.36072632670402527, 0.026240641251206398, 0.09493876993656158, 0.12203314155340195, 0.0640302300453186, 0.13458214700222015, 0.19451306760311127, 0.3176366686820984, 0.19878560304641724, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.19523903727531433, 0.1090913861989975, 0.11059779673814774, 0.03402426466345787, 0.4491459131240845, 0.1729225516319275, 0.3482173979282379, 0.01764478161931038, 0.14307594299316406, 0.22771455347537994, 0.04787566140294075, 0.14714154601097107, 0.028272001072764397, 0.23823784291744232, 0.19700175523757935, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1428564339876175, 0.03585843741893768, 0.023294193670153618, 0.1143055409193039, 0.07461919635534286, 0.13578416407108307, 0.4153969883918762, 0.03374828025698662, 0.10746961832046509, 0.17216910421848297, 0.02314077876508236, 0.02450137585401535, 0.06497504562139511, 0.381274551153183, 0.14229674637317657, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5444629788398743, 0.049506742507219315, 0.09827632457017899, 0.29229700565338135, 0.06650383025407791, 0.11397240310907364, 0.597455620765686, 0.1362738311290741, 0.15222173929214478, 0.2562837302684784, 0.13646292686462402, 0.38294121623039246, 0.030382927507162094, 0.038297515362501144, 0.465526819229126, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.12950241565704346, 0.2834409177303314, 0.40745216608047485, 0.040315985679626465, 0.09126543253660202, 0.16738829016685486, 0.24838824570178986, 0.2707839906215668, 0.5177856087684631, 0.1416875720024109, 0.6573355793952942, 0.4225574731826782, 0.02239617332816124, 0.07502269744873047, 0.07588320225477219, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00751910824328661, 0.5024122595787048, 0.38239815831184387, 0.016937274485826492, 0.039716992527246475, 0.11479316651821136, 0.004478333052247763, 0.02017248421907425, 0.011771232821047306, 0.0035600941628217697, 0.03807784244418144, 0.07125832885503769, 0.1964063048362732, 0.0026467873249202967, 0.00302477041259408, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.006645309738814831, 0.043047573417425156, 0.04108792915940285, 0.028674451634287834, 0.10265154391527176, 0.03326163440942764, 0.05858607590198517, 0.06312219053506851, 0.013714859262108803, 0.017589740455150604, 0.02732386440038681, 0.11026919633150101, 0.028857730329036713, 0.054291173815727234, 0.19011041522026062, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.006623337976634502, 0.06184479594230652, 0.014693422242999077, 0.03981047496199608, 0.08752858638763428, 0.01962500624358654, 0.06706372648477554, 0.011501927860081196, 0.0061228955164551735, 0.013949333690106869, 0.018435969948768616, 0.03678559139370918, 0.022487374022603035, 0.0660797506570816, 0.28934401273727417, 4.347301455709385e-06, 0.18382565677165985, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04245300590991974, 0.10349805653095245, 0.03407163918018341, 0.007511724252253771, 0.011565770022571087, 0.010817471891641617, 0.05971734598278999, 0.00459411833435297, 0.00350962788797915, 0.021488210186362267, 0.02298545651137829, 0.06376963108778, 0.036461468786001205, 0.1865386664867401, 0.16962040960788727, 0.0001576173526700586, 0.00605444610118866, 0.19315025210380554, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.014149562455713749, 0.03299444913864136, 0.007003516890108585, 0.004260434303432703, 0.018919609487056732, 0.008522795513272285, 0.018369171768426895, 0.015471882186830044, 0.0008095644298009574, 0.012402600608766079, 0.0075600892305374146, 0.03885417431592941, 0.05682341009378433, 0.0525624044239521, 0.22132590413093567, 0.0015271879965439439, 0.2696094512939453, 0.0976908802986145, 0.19172586500644684, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01582285761833191, 0.013434984721243382, 0.0299182441085577, 0.03647983819246292, 0.009840411134064198, 0.06101881340146065, 0.04943924769759178, 0.3809337913990021, 0.027872184291481972, 0.07177315652370453, 0.06987256556749344, 0.014244881458580494, 0.18650749325752258, 0.16280896961688995, 0.16209137439727783, 0.018620789051055908, 0.1513659805059433, 0.1261996626853943, 0.04123798385262489, 0.18324223160743713, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.018014581874012947, 0.11459828168153763, 0.013770120218396187, 0.021584663540124893, 0.02155740186572075, 0.03133949637413025, 0.03938381373882294, 0.28105995059013367, 0.02592163160443306, 0.026603924110531807, 0.010026685893535614, 0.009953479282557964, 0.004658891819417477, 0.014652709476649761, 0.16460371017456055, 7.739824650343508e-05, 0.0007302183075807989, 0.0020413347519934177, 0.0010007238015532494, 0.20195050537586212, 0.04546361416578293, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.001359884045086801, 0.029354294762015343, 0.0013457777677103877, 0.0026418184861540794, 0.008543581701815128, 0.003654624568298459, 0.0034977763425558805, 0.039957791566848755, 0.00108401442412287, 0.0005604945472441614, 0.0003877367707900703, 0.0033066808246076107, 0.007358025759458542, 0.007617549039423466, 0.20286646485328674, 0.0007431988487951458, 0.330532044172287, 0.08558935672044754, 0.06556878238916397, 0.10690004378557205, 0.1145712360739708, 0.06475446373224258, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.015068605542182922, 0.027786174789071083, 0.015096615999937057, 0.048349082469940186, 0.03296791389584541, 0.0033369800075888634, 0.004459223244339228, 0.01348987128585577, 0.0010384898632764816, 0.013556106016039848, 0.015940798446536064, 0.042712315917015076, 0.02055070362985134, 0.042082786560058594, 0.17761820554733276, 0.015635214745998383, 0.050190601497888565, 0.02352251298725605, 0.24284599721431732, 0.06325101107358932, 0.02171560376882553, 0.015677697956562042, 0.4775830805301666, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09032934159040451, 0.007927155122160912, 0.08835490047931671, 0.21186837553977966, 0.05379607528448105, 0.23637458682060242, 0.16646702587604523, 0.022663533687591553, 0.024165447801351547, 0.08468358218669891, 0.07286331057548523, 0.016201749444007874, 0.031014403328299522, 0.026781529188156128, 0.21159759163856506, 0.03602181747555733, 0.2262161672115326, 0.11374488472938538, 0.22297167778015137, 0.018925879150629044, 0.2400040328502655, 0.13629396259784698, 0.14897051453590393, 0.11721047759056091, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.014649872668087482, 0.032003261148929596, 0.1914098560810089, 0.17710277438163757, 0.07542474567890167, 0.05287592485547066, 0.14732114970684052, 0.08320016413927078, 0.025441674515604973, 0.02800501137971878, 0.0780739113688469, 0.04154554009437561, 0.017996925860643387, 0.08907850831747055, 0.17056028544902802, 0.001669732853770256, 0.0008830919396132231, 0.007873992435634136, 0.004793200176209211, 0.032567575573921204, 0.019068563356995583, 0.01167156733572483, 0.006520072463899851, 0.001765590044669807, 0.479371041059494, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.29397615790367126, 0.03400568664073944, 0.3242063522338867, 0.3681035339832306, 0.48163339495658875, 0.025333818048238754, 0.20042747259140015, 0.06051841378211975, 0.2913966476917267, 0.19229580461978912, 0.12739360332489014, 0.07057002186775208, 0.012750222347676754, 0.053084854036569595, 0.09877952188253403, 0.04264334216713905, 0.01628556102514267, 0.012549073435366154, 0.1270730197429657, 0.09553729742765427, 0.12904676795005798, 0.28088441491127014, 0.08353402465581894, 0.19219043850898743, 0.1467161476612091, 0.04815742373466492, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2290111482143402, 0.04351853206753731, 0.4067046046257019, 0.12047477811574936, 0.3140789866447449, 0.03630740940570831, 0.1768438071012497, 0.13207398355007172, 0.0676346942782402, 0.07621245086193085, 0.1797569841146469, 0.24804529547691345, 0.009716867469251156, 0.01671340875327587, 0.15996301174163818, 0.006975929252803326, 0.05510300025343895, 0.007132354192435741, 0.0349782258272171, 0.02191060781478882, 0.018211986869573593, 0.026551326736807823, 0.03648876026272774, 0.06464254856109619, 0.049987878650426865, 0.05908217281103134, 0.5448521375656128, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0448942668735981, 0.015721717849373817, 0.04864601418375969, 0.03494936227798462, 0.016112152487039566, 0.06668571382761002, 0.05302642658352852, 0.07182876765727997, 0.006946365814656019, 0.011091585271060467, 0.1120418831706047, 0.008756275288760662, 0.055249348282814026, 0.03253563493490219, 0.187040314078331, 0.000807860866189003, 0.00374230626039207, 0.004482839722186327, 0.005506760906428099, 0.000447272410383448, 0.003816538956016302, 0.03234753757715225, 0.014306235127151012, 0.01718331128358841, 0.04840204864740372, 0.06595310568809509, 0.18900929391384125, 0.0723472312092781, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3104230761528015, 0.04545353353023529, 0.3986057937145233, 0.6762936115264893, 0.03838818892836571, 0.03300129249691963, 0.27034318447113037, 0.21517230570316315, 0.008858010172843933, 0.2650390863418579, 0.2720700800418854, 0.005442188587039709, 0.06764175742864609, 0.053534120321273804, 0.18754751980304718, 0.00447529973462224, 0.019966747611761093, 0.03737834841012955, 0.3797287940979004, 0.010614297352731228, 0.05463654175400734, 0.32780376076698303, 0.0739898681640625, 0.25606051087379456, 0.8621841073036194, 0.2645638585090637, 0.25103500485420227, 0.016027942299842834, 0.004609693773090839, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.011383982375264168, 0.11127021163702011, 0.0030386100988835096, 0.0067845494486391544, 0.013927198015153408, 0.08719860762357712, 0.03287587687373161, 0.5690041184425354, 0.03855481743812561, 0.020931608974933624, 0.01293823029845953, 0.047187648713588715, 0.021772168576717377, 0.1471272110939026, 0.18776896595954895, 0.0010164460400119424, 0.011448963545262814, 0.03378765657544136, 0.02785181999206543, 0.056788451969623566, 0.07099426537752151, 0.008927138522267342, 0.01755385287106037, 0.039185769855976105, 0.09313513338565826, 0.027632856741547585, 0.12282836437225342, 0.017955774441361427, 0.02453978732228279, 0.267269104719162, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.005892250686883926, 0.03474593162536621, 0.023128867149353027, 0.002957691205665469, 0.03212961554527283, 0.015600761398673058, 0.0076070488430559635, 0.04006163775920868, 0.012522950768470764, 0.00397108681499958, 0.004476191475987434, 0.01931026391685009, 0.006290406920015812, 0.014653924852609634, 0.17843826115131378, 0.09903331845998764, 0.854941725730896, 0.020280463621020317, 0.8786925673484802, 0.37992238998413086, 0.20425425469875336, 0.32038459181785583, 0.8171603083610535, 0.2503354549407959, 0.7644308805465698, 0.7474347949028015, 0.935006856918335, 0.36836859583854675, 0.03383934497833252, 0.0021248040720820427, 0.21007098257541656, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.030382098630070686, 0.14396639168262482, 0.0023552696220576763, 0.003069670405238867, 0.03293609246611595, 0.010766614228487015, 0.04698408767580986, 0.0892328992486, 0.010764017701148987, 0.01645551063120365, 0.0007101192022673786, 0.14693684875965118, 0.10194381326436996, 0.06734117865562439, 0.21650707721710205, 0.09584157168865204, 0.00421579135581851, 0.0017077650409191847, 0.0670090913772583, 0.10943465679883957, 0.05715145170688629, 0.03694647178053856, 0.04514404758810997, 0.04956913739442825, 0.07195062190294266, 0.4566742479801178, 0.20942343771457672, 0.1548582911491394, 0.3906869888305664, 0.03925589844584465, 0.005858495831489563, 0.23115697503089905, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11579495668411255, 0.04704239219427109, 0.08932461589574814, 0.10469675809144974, 0.3945455551147461, 0.10528933256864548, 0.15413445234298706, 0.13012593984603882, 0.37207290530204773, 0.07726370543241501, 0.08641648292541504, 0.07665102183818817, 0.02378079853951931, 0.06452124565839767, 0.12331708520650864, 0.10393274575471878, 0.03258725255727768, 0.01998279243707657, 0.13928532600402832, 0.08602269738912582, 0.139993816614151, 0.2561682462692261, 0.08122693002223969, 0.28790318965911865, 0.34215468168258667, 0.023110536858439445, 0.8003224730491638, 0.11519370973110199, 0.5406965613365173, 0.2252652645111084, 0.07071924954652786, 0.03988110274076462, 0.09249765425920486, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20921318233013153, 0.07137931883335114, 0.3537597060203552, 0.1065746620297432, 0.30610421299934387, 0.07002534717321396, 0.22329437732696533, 0.23702743649482727, 0.06014438346028328, 0.05975072830915451, 0.17522762715816498, 0.3013332188129425, 0.02163097821176052, 0.016774384304881096, 0.15580035746097565, 0.006400381214916706, 0.03668399527668953, 0.006957556586712599, 0.024804070591926575, 0.013962345197796822, 0.010118995793163776, 0.014814852736890316, 0.02360437996685505, 0.038752347230911255, 0.10996780544519424, 0.24877001345157623, 0.7050904035568237, 0.103914275765419, 0.0656881257891655, 0.03925013542175293, 0.0268316138535738, 0.009403076022863388, 0.042995911091566086, 0.38370969891548157, NaN, NaN, NaN, NaN, NaN, NaN], [0.037447404116392136, 0.022215796634554863, 0.033449236303567886, 0.026462113484740257, 0.01563168875873089, 0.07434160262346268, 0.05695066228508949, 0.11209315806627274, 0.007291351445019245, 0.008904322981834412, 0.08964232355356216, 0.01435061078518629, 0.07215401530265808, 0.030404584482312202, 0.17889626324176788, 0.0005728903925046325, 0.0018518416909500957, 0.003297911025583744, 0.002339646453037858, 0.0003125199000351131, 0.0013706001918762922, 0.011640608310699463, 0.005699110683053732, 0.00646078959107399, 0.029403753578662872, 0.09435103088617325, 0.4532504379749298, 0.1454003006219864, 0.08155784755945206, 0.1478416919708252, 0.06988534331321716, 0.07031917572021484, 0.08092489838600159, 0.16178953647613525, 0.09959835559129715, NaN, NaN, NaN, NaN, NaN], [0.35028940439224243, 0.06261257082223892, 0.400876522064209, 0.6601436138153076, 0.0364767424762249, 0.0348673090338707, 0.3584212362766266, 0.3042086958885193, 0.012779565528035164, 0.3784087598323822, 0.29859334230422974, 0.00785628892481327, 0.11913719773292542, 0.06971576809883118, 0.17937220633029938, 0.007587960455566645, 0.01947515644133091, 0.06775914877653122, 0.37032291293144226, 0.014833947643637657, 0.04509717598557472, 0.2979332506656647, 0.08052700757980347, 0.2017516791820526, 0.8817963004112244, 0.3514429032802582, 0.3636293411254883, 0.14158478379249573, 0.09958238899707794, 0.13573585450649261, 0.27771836519241333, 0.47418463230133057, 0.36210212111473083, 0.2140081375837326, 0.022566867992281914, 0.004614678677171469, NaN, NaN, NaN, NaN], [0.014627714641392231, 0.1739588975906372, 0.0033204040955752134, 0.007496224716305733, 0.011711684986948967, 0.10170583426952362, 0.050673384219408035, 0.6495208740234375, 0.040652137249708176, 0.03492900729179382, 0.01829371228814125, 0.07074988633394241, 0.02588740922510624, 0.18312060832977295, 0.1794223189353943, 0.0009141381597146392, 0.00906511303037405, 0.026196878403425217, 0.011460180394351482, 0.03924085199832916, 0.05833837762475014, 0.004696658346801996, 0.009781464003026485, 0.029306253418326378, 0.06398104876279831, 0.017127037048339844, 0.0922316163778305, 0.03436172753572464, 0.12105685472488403, 0.475220263004303, 0.20121201872825623, 0.0066191148944199085, 0.018271028995513916, 0.05732923001050949, 0.018915977329015732, 0.019877590239048004, 0.23682713508605957, NaN, NaN, NaN], [0.006626310292631388, 0.049714479595422745, 0.02355029061436653, 0.0033578642178326845, 0.02970620058476925, 0.020507775247097015, 0.008351391181349754, 0.03789898753166199, 0.008593969978392124, 0.004206442274153233, 0.004605707712471485, 0.02678176388144493, 0.006028715055435896, 0.012980426661670208, 0.1725957691669464, 0.14320576190948486, 0.892350971698761, 0.030759859830141068, 0.8051734566688538, 0.7149769067764282, 0.4937312602996826, 0.3181091248989105, 0.8743517994880676, 0.3442763686180115, 0.8711729049682617, 0.7545801997184753, 0.9297782182693481, 0.6998263001441956, 0.17287810146808624, 0.008261360228061676, 0.9148194789886475, 0.7390273213386536, 0.743715763092041, 0.8801547288894653, 0.47275617718696594, 0.02699747122824192, 0.002916275057941675, 0.1803632229566574, NaN, NaN], [0.029822910204529762, 0.18419219553470612, 0.002088941168040037, 0.00302593014203012, 0.028257815167307854, 0.012486547231674194, 0.051940228790044785, 0.10161811858415604, 0.01137576438486576, 0.02022942155599594, 0.0007436276064254344, 0.2113851010799408, 0.1359580010175705, 0.08821411430835724, 0.2053057849407196, 0.0431031733751297, 0.0034584910608828068, 0.0008681766339577734, 0.032780423760414124, 0.11873625963926315, 0.03893061354756355, 0.019801655784249306, 0.03132590278983116, 0.05763043835759163, 0.06388700753450394, 0.3317660689353943, 0.16543246805667877, 0.10311393439769745, 0.4146954417228699, 0.09686555713415146, 0.06189668923616409, 0.5733434557914734, 0.2515217959880829, 0.17396190762519836, 0.13145960867404938, 0.40639445185661316, 0.07709264755249023, 0.007335619535297155, 0.2446187138557434, NaN], [0.016353517770767212, 0.03170220926403999, 0.014149405062198639, 0.013441388495266438, 0.037340469658374786, 0.010170645080506802, 0.0053974115289747715, 0.025274941697716713, 0.017184404656291008, 0.0020940443500876427, 0.006704597268253565, 0.009430822916328907, 0.030376460403203964, 0.024553189054131508, 0.15533798933029175, 0.046706411987543106, 0.31744489073753357, 0.6429179310798645, 0.4889025092124939, 0.43930482864379883, 0.3055577576160431, 0.6935683488845825, 0.25992196798324585, 0.7758384346961975, 0.2076689600944519, 0.8320663571357727, 0.39907822012901306, 0.8469056487083435, 0.5997118353843689, 0.31635957956314087, 0.36650604009628296, 0.2247273474931717, 0.7608639597892761, 0.37947097420692444, 0.8680096864700317, 0.5816919803619385, 0.19056683778762817, 0.27210569381713867, 0.06685535609722137, 0.040061503648757935]], [[0.06952784210443497, 0.0770183801651001, 0.23747292160987854, 0.022874178364872932, 0.14143598079681396, 0.08435114473104477, 0.0795491486787796, 0.054600730538368225, 0.015159118920564651, 0.06120437756180763, 0.02771361917257309, 0.06765643507242203, 0.013518131338059902, 0.15485556423664093, 0.21279898285865784, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2531612813472748, 0.03241151198744774, 0.04793045297265053, 0.13835468888282776, 0.05921119078993797, 0.20751594007015228, 0.5453532934188843, 0.021712571382522583, 0.07093679159879684, 0.2689567506313324, 0.13515745103359222, 0.05570060759782791, 0.04099860414862633, 0.03517309948801994, 0.11268090456724167, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.35043928027153015, 0.18572849035263062, 0.0481790192425251, 0.19426384568214417, 0.018465382978320122, 0.2676069438457489, 0.3000488579273224, 0.2726097106933594, 0.08134563267230988, 0.10164237022399902, 0.05787196010351181, 0.03694695979356766, 0.21335498988628387, 0.0815601795911789, 0.051584985107183456, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10967924445867538, 0.047143928706645966, 0.06498727947473526, 0.0161599051207304, 0.08311080187559128, 0.25361040234565735, 0.2589581310749054, 0.0646943673491478, 0.11701063811779022, 0.7398742437362671, 0.11236728727817535, 0.4240334630012512, 0.09019055217504501, 0.1980810910463333, 0.08526580780744553, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0050394656136631966, 0.005000656470656395, 0.01952306181192398, 0.4184519350528717, 0.012662295252084732, 0.015614073723554611, 0.006089636590331793, 0.027387546375393867, 0.007885311730206013, 0.009227052330970764, 0.015002718195319176, 0.002679894445464015, 0.040426015853881836, 0.023895790800452232, 0.031263262033462524, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1104135811328888, 0.16341662406921387, 0.10040471702814102, 0.15014782547950745, 0.22085179388523102, 0.07417210936546326, 0.08140900731086731, 0.21936744451522827, 0.12380684167146683, 0.030364450067281723, 0.008148477412760258, 0.040405042469501495, 0.016740301623940468, 0.05651557818055153, 0.03777482733130455, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.021739037707448006, 0.025255737826228142, 0.041796568781137466, 0.028582973405718803, 0.06361079961061478, 0.10603900998830795, 0.04079660773277283, 0.23573672771453857, 0.031395647674798965, 0.17699679732322693, 0.11518478393554688, 0.12758946418762207, 0.029195530340075493, 0.19761133193969727, 0.24158287048339844, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1121676117181778, 0.056780170649290085, 0.05766424164175987, 0.4753672778606415, 0.17093990743160248, 0.055545274168252945, 0.23774300515651703, 0.047642335295677185, 0.2396271675825119, 0.07084424793720245, 0.05071293190121651, 0.15200014412403107, 0.17973174154758453, 0.16349640488624573, 0.16329222917556763, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08155515789985657, 0.04415197670459747, 0.09395420551300049, 0.06736686080694199, 0.009449290111660957, 0.007789341267198324, 0.08313233405351639, 0.018231436610221863, 0.2736586928367615, 0.12516330182552338, 0.14283257722854614, 0.03993181511759758, 0.11735112965106964, 0.037545330822467804, 0.095799021422863, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07989984005689621, 0.019307896494865417, 0.05061032995581627, 0.29983657598495483, 0.009587445296347141, 0.23453857004642487, 0.06259765475988388, 0.014452173374593258, 0.026213111355900764, 0.03952796012163162, 0.12968890368938446, 0.019515926018357277, 0.23016268014907837, 0.18980233371257782, 0.14884653687477112, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.042069002985954285, 0.007410319056361914, 0.027750220149755478, 0.14348776638507843, 0.190275177359581, 0.0696464255452156, 0.09576459228992462, 0.08924749493598938, 0.16830699145793915, 0.14098002016544342, 0.2945949137210846, 0.08460760116577148, 0.11812892556190491, 0.2108343094587326, 0.28860458731651306, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.509858250617981, 0.07021021842956543, 0.044154465198516846, 0.005825423635542393, 0.5241404175758362, 0.030089300125837326, 0.19222509860992432, 0.02549084462225437, 0.1939508020877838, 0.09437919408082962, 0.10883274674415588, 0.13631868362426758, 0.08004569262266159, 0.04784407094120979, 0.14005501568317413, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.029798628762364388, 0.0011461747344583273, 0.00650657806545496, 0.02902117185294628, 0.007348767947405577, 0.012432223185896873, 0.018553903326392174, 0.006125486921519041, 0.008405826054513454, 0.057926055043935776, 0.04542696848511696, 0.21123111248016357, 0.05352021008729935, 0.2931033968925476, 0.1833699345588684, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01627730205655098, 0.0057758791372179985, 0.013731835409998894, 0.6289489269256592, 0.011782719753682613, 0.006108477246016264, 0.005309773609042168, 0.023312430828809738, 0.012817217037081718, 0.00939176045358181, 0.04320970177650452, 0.012798959389328957, 0.1585281491279602, 0.11795029044151306, 0.13285225629806519, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.39748579263687134, 0.10528232902288437, 0.006042438093572855, 0.07306646555662155, 0.020484283566474915, 0.09288878738880157, 0.6331413388252258, 0.03478514030575752, 0.016230005770921707, 0.039869412779808044, 0.10224607586860657, 0.005181388463824987, 0.007975003682076931, 0.01008305512368679, 0.026732152327895164, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.005564282648265362, 0.001319661969318986, 0.028383644297719002, 0.01146539393812418, 0.028919272124767303, 0.012663042172789574, 0.023019153624773026, 0.0018097365973517299, 0.0143426563590765, 0.021044740453362465, 0.015969598665833473, 0.03200899809598923, 0.013908782042562962, 0.03448842838406563, 0.20206299424171448, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3364894986152649, 0.00033270660787820816, 0.017299778759479523, 0.02505551464855671, 0.00914769060909748, 0.0018482855521142483, 0.040363892912864685, 0.0008854345069266856, 0.020481230691075325, 0.022734129801392555, 0.016724254935979843, 0.0011141380527988076, 5.783090819022618e-05, 0.0005799515638500452, 0.07228588312864304, 0.17503570020198822, 0.10145211219787598, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0004661931307055056, 0.4122284948825836, 0.0022180580999702215, 0.00018468582129571587, 0.00030452435021288693, 5.825214248034172e-05, 0.0012309255544096231, 0.0017770789563655853, 1.19774986160337e-05, 0.0001907332189148292, 0.0007099026697687805, 0.0006694658659398556, 1.216385771840578e-05, 0.00011785236711148173, 0.00036971797817386687, 0.002467370592057705, 0.014373218640685081, 0.18901397287845612, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04950903728604317, 0.2967310845851898, 0.021222729235887527, 0.01289455872029066, 0.009955117478966713, 0.008917939849197865, 0.011312013491988182, 0.01272521447390318, 0.0006359940161928535, 0.011413054540753365, 0.006479735020548105, 0.0053005279041826725, 0.001741865067742765, 0.0027997863944619894, 0.08213357627391815, 4.782021278515458e-05, 0.0002036100922850892, 0.15351639688014984, 0.001678619533777237, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.020872987806797028, 3.087984805461019e-05, 0.009670623578131199, 0.0253498163074255, 0.010817835107445717, 0.4320962131023407, 0.017970044165849686, 0.0021109851077198982, 0.0003069202939514071, 0.008261006325483322, 0.006166533567011356, 0.7898750901222229, 0.11304597556591034, 0.12737329304218292, 0.011856237426400185, 0.015930648893117905, 0.006582066882401705, 0.10560829937458038, 0.3465193808078766, 0.012144939973950386, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06067817285656929, 0.005839335732161999, 0.025896329432725906, 0.03351203724741936, 0.025002295151352882, 0.25514867901802063, 0.4275963008403778, 0.0194717925041914, 0.0888834074139595, 0.04690318927168846, 0.03570560738444328, 0.0850825086236, 0.0388353131711483, 0.24394167959690094, 0.10019046813249588, 0.010950141586363316, 0.003185260808095336, 0.03380253165960312, 0.13516294956207275, 0.16374172270298004, 0.0833682045340538, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.014415884390473366, 0.001141559099778533, 0.0678224116563797, 0.024646559730172157, 0.08796916157007217, 0.022639306262135506, 0.07784608006477356, 0.02605922892689705, 0.014093886129558086, 0.0286162830889225, 0.09674176573753357, 0.04692256450653076, 0.03519048914313316, 0.20982496440410614, 0.1800668090581894, 4.016391176264733e-05, 0.0003202538937330246, 0.0050767818465828896, 1.7212016246048734e-05, 0.5176156759262085, 0.003749872324988246, 0.00026106167933903635, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02086471952497959, 0.0008324789232574403, 0.01815967448055744, 0.002886975882574916, 0.0020961007103323936, 0.004472428001463413, 0.033020272850990295, 0.0047500282526016235, 0.012928733602166176, 0.014328529126942158, 0.015946470201015472, 0.06593997031450272, 0.00855537410825491, 0.07526978105306625, 0.1768130511045456, 0.13457109034061432, 0.07774609327316284, 0.006220821291208267, 0.0008077693055383861, 0.2509746253490448, 0.17662860453128815, 0.13796226680278778, 0.053514063358306885, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0009654826717451215, 0.000225315525312908, 0.0006124225910753012, 0.0007836261647753417, 0.0007428302778862417, 0.003282200777903199, 0.008662715554237366, 0.45239004492759705, 4.857195381191559e-05, 0.0006357804522849619, 0.0010122592793777585, 0.0006606358801946044, 0.00025698603712953627, 0.0011707579251378775, 0.0028539940249174833, 0.06553670763969421, 0.09473168104887009, 0.013516419567167759, 0.0013789478689432144, 0.03089364431798458, 0.0676402598619461, 0.03963227570056915, 0.17151857912540436, 0.1338733434677124, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0025523374788463116, 0.0009212270379066467, 0.09748471528291702, 0.057154957205057144, 0.4982932209968567, 0.000552327954210341, 0.02918482944369316, 0.0039253802970051765, 0.00450148293748498, 0.0014971394557505846, 0.009822547435760498, 0.0017059196252375841, 0.001570553402416408, 0.005804183427244425, 0.00957300141453743, 0.07379595190286636, 0.1714182198047638, 0.13684017956256866, 0.00734432740136981, 0.0039545828476548195, 0.09408346563577652, 0.0452522449195385, 0.2525797188282013, 0.15314188599586487, 0.008748584426939487, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.016401896253228188, 0.00043752315104939044, 0.0039018490351736546, 0.005885160993784666, 0.0023499932140111923, 0.0031332974322140217, 0.055512603372335434, 0.003903925186023116, 0.10197419673204422, 0.009071548469364643, 0.023729920387268066, 0.002627716166898608, 0.01914973370730877, 0.02837507426738739, 0.1623656302690506, 0.006909683812409639, 0.034793343394994736, 0.13824458420276642, 0.0004423256032168865, 0.38493895530700684, 0.12702688574790955, 0.0007700703572481871, 0.005257567390799522, 0.3978818655014038, 0.028774550184607506, 0.016022928059101105, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0004865071678068489, 2.4051656509982422e-05, 0.00020084556308574975, 0.0003736558719538152, 0.000646126689389348, 9.209318523062393e-05, 0.009753170423209667, 9.854567178990692e-05, 0.34485483169555664, 0.00047165394062176347, 0.0012700805673375726, 0.000479432987049222, 0.0015819557011127472, 0.0008011643076315522, 0.0017131956992670894, 0.15589091181755066, 0.059809040278196335, 0.2019805759191513, 0.006274765357375145, 0.053891621530056, 0.38889890909194946, 0.024021193385124207, 0.016828669235110283, 0.09206627309322357, 0.15270450711250305, 0.10960505902767181, 0.14381197094917297, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03442303463816643, 0.014513631351292133, 0.003174385754391551, 0.00478995218873024, 0.0017101461999118328, 0.003900717245414853, 0.05713852494955063, 0.013628470711410046, 0.0976317971944809, 0.28217896819114685, 0.01894235610961914, 0.009533336386084557, 0.003816690994426608, 0.005922130309045315, 0.12864208221435547, 0.0011966965394094586, 0.0013769377255812287, 0.0006101150647737086, 4.0936538425739855e-05, 0.008213219232857227, 0.03395655378699303, 0.0003392287762835622, 0.00015790743054822087, 0.000944053172133863, 0.0007261222926899791, 0.011664116755127907, 0.22049497067928314, 0.0034024016931653023, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01004086248576641, 0.01997406780719757, 0.005450551863759756, 0.006583535112440586, 0.0027623113710433245, 0.002903316868469119, 0.03531726077198982, 0.008635452017188072, 0.029197845607995987, 0.02162068709731102, 0.013219092041254044, 0.2711889445781708, 0.00537630682811141, 0.006846235599368811, 0.06079954653978348, 0.2470119595527649, 0.22662757337093353, 0.086290642619133, 0.0011605313047766685, 0.20862528681755066, 0.31339770555496216, 0.007298772688955069, 0.00864456407725811, 0.010568802244961262, 0.01924213580787182, 0.034804634749889374, 0.16789764165878296, 0.11296499520540237, 0.017940307036042213, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00031272557680495083, 8.196506314561702e-06, 4.237617031321861e-05, 0.00043677922803908587, 0.00024717405904084444, 0.022641032934188843, 0.002573953475803137, 0.0004433683061506599, 0.0013428670354187489, 0.00034036010038107634, 0.0007929583080112934, 0.0033021108247339725, 0.4761846959590912, 0.05593165382742882, 0.00081905338447541, 0.3800778388977051, 0.4679488241672516, 0.19362112879753113, 0.18464821577072144, 0.046723559498786926, 0.160307839512825, 0.24654103815555573, 0.2610638439655304, 0.07595612108707428, 0.1325986683368683, 0.022732526063919067, 0.1294456422328949, 0.2688123285770416, 0.12097980827093124, 0.12297553569078445, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00267792004160583, 4.751862070406787e-05, 0.014043050818145275, 0.02037942036986351, 0.04410611465573311, 0.04370833560824394, 0.06117184832692146, 0.01571183279156685, 0.11117196083068848, 0.006906491704285145, 0.0029646854382008314, 0.15407170355319977, 0.010935205966234207, 0.03797803074121475, 0.16977860033512115, 0.005153980106115341, 0.0002073257346637547, 0.12819816172122955, 0.00011319551413180307, 0.08506736904382706, 0.013190183788537979, 0.0028314462397247553, 0.00016588614380452782, 0.009067418053746223, 0.0008525841985829175, 0.00018506577180232853, 0.0002737078757490963, 0.0002474631182849407, 0.04919072240591049, 0.1850043386220932, 0.0018668848788365722, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.011722833849489689, 0.005004812031984329, 0.007801789790391922, 0.0020204312168061733, 0.004946417640894651, 0.000467105332063511, 0.11018845438957214, 0.016256244853138924, 0.05208335816860199, 0.08122430741786957, 0.4447634816169739, 0.0032620911952108145, 0.0036480925045907497, 0.02699565887451172, 0.038189876824617386, 0.4235798418521881, 0.8363600969314575, 0.13292381167411804, 0.03160996362566948, 0.6294970512390137, 0.3827916085720062, 0.01768689975142479, 0.031598031520843506, 0.05291707068681717, 0.004268768709152937, 0.01666090451180935, 0.0017059938982129097, 0.03961870074272156, 0.006749838124960661, 0.2787548303604126, 0.12898604571819305, 0.00984524842351675, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.024071840569376945, 0.0004321316082496196, 0.023504342883825302, 0.020648522302508354, 0.021508874371647835, 0.012214796617627144, 0.024360070005059242, 0.0013747027842327952, 0.0815734788775444, 0.08039785921573639, 0.06951787322759628, 0.017521949484944344, 0.04566040262579918, 0.08389204740524292, 0.15396325290203094, 0.001200420199893415, 0.004923743661493063, 0.03312471881508827, 7.996988279046491e-05, 0.2118730992078781, 0.0288531631231308, 0.00010192030458711088, 0.0002958755649160594, 0.007303019054234028, 0.00011155433458043262, 2.6572593014861923e-06, 0.00035481253871694207, 2.4723947262828005e-06, 2.6933960270980606e-06, 0.017764916643500328, 0.0003658832865767181, 0.25218549370765686, 0.002238432876765728, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0014979105908423662, 4.0405931940767914e-05, 0.0008743218495510519, 0.001329930848442018, 0.0032007889822125435, 0.0002464030694682151, 0.015361684374511242, 0.00014017200737725943, 0.3369258642196655, 0.0015512423124164343, 0.003011554479598999, 0.0010034784208983183, 0.0037561107892543077, 0.0018123533809557557, 0.0037892721593379974, 0.16854390501976013, 0.046801913529634476, 0.18834064900875092, 0.005545254796743393, 0.10321269929409027, 0.3906272351741791, 0.03742265701293945, 0.024458711966872215, 0.05521516501903534, 0.07171308994293213, 0.021107476204633713, 0.025199010968208313, 0.0027974944096058607, 0.0025010560639202595, 0.02306896261870861, 0.15930885076522827, 0.06242140382528305, 0.11754277348518372, 0.21403564512729645, NaN, NaN, NaN, NaN, NaN, NaN], [0.03386643901467323, 0.015328249894082546, 0.002211565151810646, 0.003828595858067274, 0.0012934240512549877, 0.004837968852370977, 0.04463785141706467, 0.014559985138475895, 0.04106945917010307, 0.26340487599372864, 0.017707379534840584, 0.01015215553343296, 0.0033097255509346724, 0.0058202859945595264, 0.13427288830280304, 0.0004002669302280992, 0.00040952101699076593, 0.00012874403910245746, 8.880775567376986e-06, 0.005201425869017839, 0.007163480389863253, 0.0002137795090675354, 0.00012960725871380419, 0.0005550362984649837, 0.0001244707527803257, 0.0006415210082195699, 0.03161805495619774, 4.1008814150700346e-05, 0.000599265971686691, 0.00399716105312109, 5.7038221711991355e-05, 0.0033261284697800875, 0.006950944196432829, 0.22392861545085907, 0.0028074102010577917, NaN, NaN, NaN, NaN, NaN], [0.011043943464756012, 0.029788998886942863, 0.004548549186438322, 0.006417197175323963, 0.0019613932818174362, 0.0028304944280534983, 0.02768276073038578, 0.006805655546486378, 0.02553243562579155, 0.0314837321639061, 0.015709027647972107, 0.2568790316581726, 0.008081428706645966, 0.009137820452451706, 0.06746803224086761, 0.22722585499286652, 0.18426381051540375, 0.07697561383247375, 0.0012757674558088183, 0.23254786431789398, 0.14769063889980316, 0.013780240900814533, 0.02735842764377594, 0.04001649469137192, 0.031179115176200867, 0.015889445319771767, 0.062248069792985916, 0.013498637825250626, 0.0052745710127055645, 0.2219674438238144, 0.0031969451811164618, 0.0037056237924844027, 0.028058722615242004, 0.22486938536167145, 0.09661445021629333, 0.02616964653134346, NaN, NaN, NaN, NaN], [0.0003306480939500034, 1.1417017958592623e-05, 3.816767639364116e-05, 0.000435528316302225, 0.00020690191013272852, 0.02179853804409504, 0.002864222740754485, 0.0005160043947398663, 0.001080053043551743, 0.0004847492673434317, 0.0009861867874860764, 0.003908392507582903, 0.47703394293785095, 0.07113853842020035, 0.000873323529958725, 0.27366653084754944, 0.354305237531662, 0.16368547081947327, 0.1598840057849884, 0.02900015190243721, 0.10581760108470917, 0.21902981400489807, 0.27043354511260986, 0.19813168048858643, 0.2514232099056244, 0.025616073980927467, 0.12471329420804977, 0.09682969748973846, 0.07310353219509125, 0.02883375994861126, 0.09285400807857513, 0.013515813276171684, 0.021914459764957428, 0.14159631729125977, 0.3238908648490906, 0.1783936321735382, 0.11570748686790466, NaN, NaN, NaN], [0.0030808241572231054, 6.38188939774409e-05, 0.011707174591720104, 0.023645061999559402, 0.038246914744377136, 0.047200631350278854, 0.04958858713507652, 0.012573646381497383, 0.04961754009127617, 0.005252092145383358, 0.002489157486706972, 0.17429526150226593, 0.008030706085264683, 0.02717452496290207, 0.1679786741733551, 0.0030968550126999617, 7.297070260392502e-05, 0.1371629387140274, 0.00018204482330475003, 0.04798782989382744, 0.01213640347123146, 0.0023585439193993807, 0.00011540603009052575, 0.016970379278063774, 0.0015150568215176463, 0.0003718302759807557, 0.00044133648043498397, 0.00012143531785113737, 0.021671650931239128, 0.023021340370178223, 0.00010860650218091905, 0.0005334930610843003, 0.000257489358773455, 0.0005856966599822044, 0.00045311596477404237, 0.09709983319044113, 0.18528476357460022, 0.0029071324970573187, NaN, NaN], [0.01455691922456026, 0.008012487553060055, 0.006938801147043705, 0.00259140832349658, 0.004911262542009354, 0.0004763725446537137, 0.10579084604978561, 0.021042171865701675, 0.03971559554338455, 0.07511086016893387, 0.43185338377952576, 0.0035418386105448008, 0.004437423776835203, 0.03184036538004875, 0.04226255044341087, 0.49188995361328125, 0.918917715549469, 0.2054058462381363, 0.08403602242469788, 0.6967929005622864, 0.5653088688850403, 0.03772272169589996, 0.04957969859242439, 0.18319177627563477, 0.012161915190517902, 0.07060753554105759, 0.009896048344671726, 0.1126827672123909, 0.010653471574187279, 0.1938174068927765, 0.1352803260087967, 0.0021707522682845592, 0.030638370662927628, 0.003963022027164698, 0.03303877264261246, 0.004082953091710806, 0.20578816533088684, 0.11854958534240723, 0.02041587606072426, NaN], [0.055085837841033936, 0.014846320264041424, 0.06939522176980972, 0.036867137998342514, 0.13156765699386597, 0.04343622922897339, 0.18117153644561768, 0.04244613274931908, 0.04596249759197235, 0.13158053159713745, 0.047130946069955826, 0.549620509147644, 0.24813801050186157, 0.3232562243938446, 0.11823604255914688, 0.001465475419536233, 0.00045102695003151894, 0.017218099907040596, 0.00030212500132620335, 0.11662620306015015, 0.017841650173068047, 0.00014393724268302321, 0.0003088460653088987, 0.006560556124895811, 0.0005491081974469125, 5.78465114813298e-05, 0.0019656207878142595, 0.00016285650781355798, 0.0002489366161171347, 0.011378495953977108, 0.0017521223053336143, 0.00787137821316719, 8.434856863459572e-05, 0.0012881350703537464, 7.287580228876323e-05, 0.00021561238099820912, 0.020317554473876953, 0.04195580258965492, 0.24219898879528046, 0.0017395684262737632]], [[0.2484879046678543, 0.12593188881874084, 0.11472177505493164, 0.6318025588989258, 0.009745504707098007, 0.030495919287204742, 0.054615989327430725, 0.004801109898835421, 0.23875823616981506, 0.011562658473849297, 0.02087206020951271, 0.059635717421770096, 0.011483770795166492, 0.07716090232133865, 0.041850361973047256, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3294946551322937, 0.17723912000656128, 0.041080135852098465, 0.30134642124176025, 0.0073102316819131374, 0.049291279166936874, 0.0495959147810936, 0.0037847748026251793, 0.014987694099545479, 0.07676513493061066, 0.039059415459632874, 0.006041571032255888, 0.011380840092897415, 0.011979957111179829, 0.02782473713159561, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.008675806224346161, 0.016726570203900337, 0.19906938076019287, 0.3167073726654053, 0.022006884217262268, 0.014510865323245525, 0.00237266905605793, 0.00938868336379528, 0.004848333541303873, 0.00305117666721344, 0.042285457253456116, 0.0026737553998827934, 0.017337674275040627, 0.0016427191440016031, 0.0027906473260372877, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.06292864680290222, 0.010060630738735199, 0.07846219092607498, 0.3009726405143738, 0.09911586344242096, 0.3769649565219879, 0.290684312582016, 0.048859626054763794, 0.015964722260832787, 0.02972962148487568, 0.25837212800979614, 0.050403933972120285, 0.052831199020147324, 0.44793814420700073, 0.12096201628446579, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0647541731595993, 0.06744952499866486, 0.010754776187241077, 0.15598785877227783, 0.08916914463043213, 0.4045051634311676, 0.5958212018013, 0.10594789683818817, 0.12025819718837738, 0.04822946712374687, 0.02913811057806015, 0.014846491627395153, 0.17111137509346008, 0.049513354897499084, 0.14188753068447113, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07069405168294907, 0.0006015333347022533, 0.0017680496675893664, 0.0010985832195729017, 0.0012869784841313958, 0.22278346121311188, 0.4465882480144501, 0.06128238886594772, 0.02642727456986904, 0.03756114840507507, 0.002607540925964713, 0.0018699204083532095, 0.0059012919664382935, 0.020283877849578857, 0.03355809301137924, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0861939862370491, 0.03346291184425354, 0.009915103204548359, 0.35010838508605957, 0.03437130153179169, 0.18394741415977478, 0.5006390810012817, 0.0633198693394661, 0.36160194873809814, 0.07578127831220627, 0.038500167429447174, 0.08213403075933456, 0.026455186307430267, 0.12013117223978043, 0.1146865040063858, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2484544962644577, 0.00790119543671608, 0.004407763481140137, 0.02700735628604889, 0.015422074124217033, 0.015295883640646935, 0.40846768021583557, 0.10706920176744461, 0.06367217004299164, 0.22094424068927765, 0.21221157908439636, 0.006999517325311899, 0.054566796869039536, 0.124799944460392, 0.09114839136600494, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1237153485417366, 0.029043834656476974, 0.07521974295377731, 0.04068650305271149, 0.002623512176796794, 0.008706655353307724, 0.03832445293664932, 0.14616532623767853, 0.1701044738292694, 0.20599642395973206, 0.11677426844835281, 0.2341107875108719, 0.06235762685537338, 0.003964806441217661, 0.15731573104858398, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.034962959587574005, 0.023077068850398064, 0.034600574523210526, 0.14041800796985626, 0.0021679585333913565, 0.009290770627558231, 0.07274696230888367, 0.014187950640916824, 0.1371506154537201, 0.39440277218818665, 0.2198760211467743, 0.19940708577632904, 0.11203428357839584, 0.08552268147468567, 0.11737436801195145, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.015330069698393345, 0.007386082783341408, 0.017500948160886765, 0.01906486414372921, 0.010120063088834286, 0.05364372953772545, 0.043298348784446716, 0.12658876180648804, 0.06039673835039139, 0.02238147333264351, 0.16429400444030762, 0.06984445452690125, 0.3043651580810547, 0.055543575435876846, 0.11423089355230331, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.09644094854593277, 0.0058854687958955765, 0.03721459209918976, 0.0025620406959205866, 0.062300242483615875, 0.003563062520697713, 0.07219880819320679, 0.03924282267689705, 0.025451356545090675, 0.06598387658596039, 0.026776403188705444, 0.07250863313674927, 0.45021528005599976, 0.08199745416641235, 0.4220075309276581, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01460834126919508, 0.0005662022740580142, 0.0013911814894527197, 0.05315173417329788, 0.008028149604797363, 0.016604119911789894, 0.011740745045244694, 0.008678588084876537, 0.0025609249714761972, 0.01638207584619522, 0.018210044130682945, 0.014119945466518402, 0.06550943106412888, 0.34254926443099976, 0.04794229939579964, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05372002348303795, 0.14061135053634644, 0.018787089735269547, 0.0958278551697731, 0.0019092779839411378, 0.03348369151353836, 0.13957257568836212, 0.031220966950058937, 0.19735871255397797, 0.017847368493676186, 0.0589337982237339, 0.01900595612823963, 0.1276925951242447, 0.04769464209675789, 0.4384888708591461, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08416850119829178, 0.1088641807436943, 0.0573052242398262, 0.27551695704460144, 0.030813831835985184, 0.18022866547107697, 0.10468263924121857, 0.09972096234560013, 0.31189021468162537, 0.3315774202346802, 0.2321816384792328, 0.034622836858034134, 0.14143656194210052, 0.04640315845608711, 0.09621720016002655, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7448275089263916, 0.00023065913410391659, 0.0003700565139297396, 0.0002745355886872858, 0.0005768057890236378, 1.0151054993912112e-05, 1.3715341992792673e-05, 7.643950084457174e-06, 0.0004341531603131443, 5.2913601393811405e-05, 5.353476808522828e-05, 8.812115265754983e-05, 1.1566834245968494e-06, 5.744800546381157e-06, 5.576572584686801e-05, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [8.114575030049309e-05, 0.06691394746303558, 0.04036417603492737, 0.022258125245571136, 0.055233534425497055, 0.050445422530174255, 0.048324622213840485, 0.00889397319406271, 0.1270352452993393, 0.04156908392906189, 0.20929713547229767, 0.21122632920742035, 0.414194792509079, 0.12628954648971558, 0.25567519664764404, 0.39058852195739746, 8.28505744721042e-06, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0012628535041585565, 0.0008597301202826202, 0.036364536732435226, 0.0971999391913414, 0.04217860475182533, 0.10421664267778397, 0.16082510352134705, 0.03283625468611717, 0.09032318741083145, 0.09653837233781815, 0.21890851855278015, 0.06589526683092117, 0.47985169291496277, 0.21388037502765656, 0.21010825037956238, 2.7811127438326366e-05, 0.4158080220222473, 0.0005852450849488378, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0002990703214891255, 0.001862871926277876, 0.010526847094297409, 0.01025421917438507, 0.05592086538672447, 0.02697981521487236, 0.01570008136332035, 0.02568165771663189, 0.010194454342126846, 0.048093631863594055, 0.04421652480959892, 0.02353351190686226, 0.21245922148227692, 0.0448865108191967, 0.23352482914924622, 9.039229868085252e-13, 4.1926887206500396e-05, 0.15358270704746246, 0.00044542484101839364, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00015855174569878727, 0.013162538409233093, 0.006567019037902355, 0.004201928153634071, 0.006268346216529608, 0.00024757537175901234, 0.012954139150679111, 0.003747382666915655, 0.03740423545241356, 0.007960616610944271, 0.013323514722287655, 0.06273993849754333, 0.048431456089019775, 0.13987915217876434, 0.20342004299163818, 1.9216391628896996e-16, 4.9363904963684035e-08, 0.0004218998074065894, 0.40449434518814087, 4.695959432865493e-06, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.013553211465477943, 0.03824196010828018, 0.02278091199696064, 0.09299258887767792, 0.0559159517288208, 0.00022306715254671872, 0.031003709882497787, 0.010444254614412785, 0.16168788075447083, 0.03666102886199951, 0.00852662418037653, 0.4432809352874756, 0.009321487508714199, 0.024379035457968712, 0.17351986467838287, 1.7349648803667746e-14, 5.141012060505545e-09, 3.7822364902240224e-06, 0.0002717413299251348, 0.22465285658836365, 2.698016260183067e-06, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00026768012321554124, 0.015254812315106392, 0.007090381346642971, 0.006173381581902504, 0.006773150525987148, 0.0008773274021223187, 0.00638232659548521, 0.016591282561421394, 0.004996343981474638, 0.009327422827482224, 0.008862738497555256, 0.05876166746020317, 0.009527520276606083, 0.00578573253005743, 0.20356230437755585, 3.6696812255598843e-09, 2.368522711293508e-09, 3.1902116006676806e-06, 9.520445587440918e-08, 9.990107355406508e-05, 0.2170185148715973, 0.019131841138005257, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0008312691352330148, 0.012717761099338531, 0.013986560516059399, 0.007093494758009911, 0.004876464139670134, 0.0027259632479399443, 0.0033886858727782965, 0.01589561626315117, 0.00876854918897152, 0.005017295014113188, 0.023178039118647575, 0.05755693465471268, 0.05451130494475365, 0.06928746402263641, 0.1796484887599945, 2.292660354896725e-07, 1.4062491449085002e-10, 1.0373556180720556e-11, 2.945570870549474e-11, 1.3987125901948616e-09, 1.1205498822164373e-06, 0.3382871150970459, 0.0008390913717448711, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00016753048112150282, 0.011822681874036789, 0.005686081480234861, 0.011659285984933376, 0.004307762254029512, 0.0031254058703780174, 0.009316416457295418, 0.0016170619055628777, 0.012603488750755787, 0.0245236624032259, 0.01756892167031765, 0.011099276132881641, 0.11892349272966385, 0.02075323462486267, 0.2549600899219513, 2.3133984541345853e-06, 0.00017511146143078804, 1.441240442545677e-06, 3.064446918443764e-09, 3.097617096159411e-08, 7.23518027712089e-08, 0.0017295092111453414, 0.39626115560531616, 0.00019915253506042063, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00017647366621531546, 0.053185176104307175, 0.007304554805159569, 0.004834755789488554, 0.000954066461417824, 0.025718921795487404, 0.02985404059290886, 0.09960591793060303, 0.010695043951272964, 0.016483109444379807, 0.018774237483739853, 0.05090473219752312, 0.01008983701467514, 0.028674444183707237, 0.22871088981628418, 8.689644937311981e-15, 2.8357308110571466e-06, 5.0946681540153804e-08, 2.0269605438549831e-10, 1.289949813632063e-10, 3.375676821404383e-11, 8.602300205495794e-09, 4.5097981455910485e-06, 0.29888245463371277, 6.641173968091607e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0008755451999604702, 0.020039640367031097, 0.003969491925090551, 0.007670485880225897, 0.006173306610435247, 0.012295764870941639, 0.0076020946726202965, 0.012137084268033504, 0.010956642217934132, 0.010541083291172981, 0.018125493079423904, 0.03226908668875694, 0.02587633579969406, 0.016216130927205086, 0.1660052388906479, 2.8127108337250475e-18, 1.3557467148928026e-08, 7.431774662336466e-08, 2.301476165200711e-08, 1.1707952315975767e-11, 7.274678689300762e-12, 7.034611066401852e-13, 5.257664963120856e-13, 3.4044413041556254e-05, 0.32336506247520447, 4.600838292390108e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [5.4335410823114216e-05, 0.03367479890584946, 0.004507457371801138, 0.004544241353869438, 0.00623831432312727, 0.002192543353885412, 0.004128816071897745, 0.021106822416186333, 0.0003909784718416631, 0.00830051489174366, 0.018183842301368713, 0.009683135896921158, 0.0325237475335598, 0.00792472343891859, 0.25227075815200806, 6.300134025583048e-13, 5.676838910062543e-08, 1.822371018533886e-06, 2.3448223146260716e-05, 2.5415656068616954e-07, 3.417801153204891e-08, 5.353474885616549e-10, 2.141239963115993e-11, 3.762530198514469e-08, 6.24434178462252e-05, 0.33693620562553406, 3.183486114721745e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0006012204103171825, 0.01188816037029028, 0.023532994091510773, 0.00770517997443676, 0.007410787045955658, 0.007087987381964922, 0.021027186885476112, 0.013456426560878754, 0.03266710042953491, 0.001251929672434926, 0.09021235257387161, 0.024440091103315353, 0.024299103766679764, 0.02338516153395176, 0.1967199146747589, 1.5877897954763576e-12, 1.2288996487086479e-09, 3.458522428445576e-07, 9.462546586291865e-06, 7.457422907464206e-05, 0.0005706463125534356, 1.4425116212635203e-08, 4.5430816769144455e-13, 2.616490357709722e-12, 3.545688542772041e-08, 0.00016559385403525084, 0.22770871222019196, 0.0009294600458815694, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0009616355528123677, 0.059039004147052765, 0.04997482895851135, 0.013552234508097172, 0.03981975466012955, 0.020335622131824493, 0.014380398206412792, 0.07606764137744904, 0.07161007821559906, 0.024130970239639282, 0.06891870498657227, 0.0008635766571387649, 0.023193923756480217, 0.02981526218354702, 0.21020111441612244, 2.579016999959549e-10, 1.5412886245069757e-10, 5.557828156033118e-11, 1.2367832313842086e-09, 3.3751638284229557e-07, 4.776334208145272e-07, 1.75399406998622e-07, 9.608910021829953e-12, 7.499024594652057e-14, 2.8573548556528813e-14, 3.2670008191793e-12, 4.494925178732956e-06, 0.37381958961486816, 3.638648195192218e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0013424595817923546, 0.0746709555387497, 0.011544802226126194, 0.027912717312574387, 0.0729047879576683, 0.10483764857053757, 0.07119728624820709, 0.010606798343360424, 0.044552259147167206, 0.05723145231604576, 0.034647323191165924, 0.38214871287345886, 0.003923356998711824, 0.08778946846723557, 0.19581711292266846, 3.090227983193472e-05, 8.430293382843956e-05, 4.32313208875712e-05, 1.6493000885020592e-06, 8.794136192591395e-06, 0.0005616153357550502, 0.0013158570509403944, 0.0005267951055429876, 3.675571861094795e-05, 2.42239195813454e-07, 8.356466074666002e-10, 2.3424906885338714e-06, 0.0012797197559848428, 0.6210904717445374, 0.0014036636566743255, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0016638260567560792, 0.01581355184316635, 0.08943041414022446, 0.02092832513153553, 0.021133122965693474, 0.012408973649144173, 0.01347691286355257, 0.00275444146245718, 0.027862150222063065, 0.01225491613149643, 0.018322426825761795, 0.008929668925702572, 0.00015579524915665388, 0.0014782899525016546, 0.18181975185871124, 7.67247776423119e-09, 2.954437938740284e-08, 8.54147774731473e-09, 2.011255162415182e-09, 5.265776792384713e-08, 1.4630668898618637e-09, 2.2913241082278546e-06, 3.266295323101076e-08, 1.6124132571349037e-06, 1.13081211061683e-11, 2.6358108895513247e-15, 7.728456763445024e-11, 2.3767283696685126e-09, 2.1271845980663784e-05, 0.19462287425994873, 6.456446044467157e-06, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0008640239248052239, 0.06174946948885918, 0.004653214477002621, 0.002717669354751706, 0.015129820443689823, 0.00935456808656454, 0.016078660264611244, 0.08089328557252884, 0.017857585102319717, 0.0025031790137290955, 0.00012101473839720711, 0.013123439624905586, 0.005499868653714657, 0.001559562049806118, 0.22764776647090912, 4.312543703220706e-13, 2.1705271535665815e-07, 1.1365986551936658e-07, 1.9739390211270802e-07, 7.690645453806155e-09, 4.219609994748907e-09, 9.716764060030414e-10, 3.915795687703394e-08, 3.0873563900968293e-06, 5.5168204227129536e-08, 1.0056843552375128e-10, 6.254387632798064e-12, 4.318517331930449e-12, 1.5618051990573534e-11, 6.033264071447775e-05, 0.4116440713405609, 1.8908482161350548e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0008687095833010972, 0.025285501033067703, 0.01658034697175026, 0.02363765239715576, 0.02393241412937641, 0.0657346174120903, 0.015298763290047646, 0.01792113669216633, 0.021707117557525635, 0.018967296928167343, 0.037634264677762985, 0.013209421187639236, 0.02256513573229313, 0.007774183992296457, 0.15961462259292603, 1.797858697974407e-17, 3.5553746058347713e-10, 1.0377114723070235e-09, 5.157609006545272e-09, 5.5740526777592336e-11, 3.675403037473046e-11, 3.015720268992328e-12, 1.2632186895361434e-14, 3.2584634990229233e-09, 2.7093712162695738e-08, 2.733851353305984e-15, 2.0347772078377346e-10, 7.802066534575867e-16, 1.702402683943053e-16, 1.8298086656987067e-10, 6.30185184036236e-08, 0.2592085301876068, 3.469779585429933e-06, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0001073219973477535, 0.04253393039107323, 0.010077103972434998, 0.007349912542849779, 0.00879223458468914, 0.004757148679345846, 0.008167163468897343, 0.03753674402832985, 0.00042728587868623435, 0.014237778261303902, 0.029898250475525856, 0.006872681900858879, 0.045794516801834106, 0.007500257343053818, 0.2562271058559418, 3.386366187463352e-10, 1.5587464474720036e-07, 5.430682108453766e-07, 1.926859113154933e-05, 2.7584928830037825e-06, 5.553058031182445e-07, 6.554741815989473e-08, 7.146391256540596e-10, 4.225638150501254e-08, 2.0539353045023745e-06, 0.00010312868107575923, 2.5505174860995794e-08, 1.3659710695890226e-08, 4.206753695390475e-11, 5.200286035123014e-11, 3.842067428649898e-07, 1.4282905794971157e-05, 0.31164512038230896, 0.00011869923037011176, NaN, NaN, NaN, NaN, NaN, NaN], [0.0005320480559021235, 0.010701313614845276, 0.020972738042473793, 0.007364482618868351, 0.006165153346955776, 0.00950621161609888, 0.022682208567857742, 0.018515970557928085, 0.03319491446018219, 0.00125269521959126, 0.07773777842521667, 0.022826068103313446, 0.02051766775548458, 0.020874740555882454, 0.1872510462999344, 3.098006018387167e-10, 3.2388165482899467e-09, 1.8609943808201024e-08, 5.099297482047405e-07, 4.603737033903599e-05, 0.00016448901442345232, 1.6998721719119203e-07, 1.7718410072475876e-11, 2.5886336477154437e-11, 9.218055652127077e-09, 1.2046231745443947e-07, 7.304957398446277e-05, 2.3164133111652774e-10, 2.8952129582648922e-09, 2.9085676575557606e-11, 8.895827650901023e-12, 8.14965606110718e-09, 8.762691868469119e-05, 0.2280847281217575, 0.0004104141262359917, NaN, NaN, NaN, NaN, NaN], [0.0008804904646240175, 0.05573932081460953, 0.06578188389539719, 0.01897181011736393, 0.043492771685123444, 0.026308609172701836, 0.016426166519522667, 0.09104844927787781, 0.12495335191488266, 0.04637341946363449, 0.0944451242685318, 0.0008321930072270334, 0.03243781998753548, 0.03530845418572426, 0.2013196051120758, 1.3149543676149733e-09, 1.080373679407387e-09, 5.5150013028582023e-11, 7.800748935693491e-10, 1.7859061074432248e-07, 2.183157299384675e-08, 2.5236221290469985e-07, 2.35878039323012e-10, 9.060349692724401e-12, 1.4339956088890715e-12, 1.7799637631876752e-12, 2.9941787715870305e-08, 6.0217857935640495e-06, 3.1683756313016787e-11, 4.5713120788715145e-11, 3.4124135808721867e-13, 3.591858459424911e-15, 1.3559961530365539e-12, 3.119595021416899e-06, 0.35679423809051514, 3.964137067669071e-05, NaN, NaN, NaN, NaN], [0.001610875129699707, 0.08435038477182388, 0.014167247340083122, 0.03493078798055649, 0.07050123810768127, 0.10772886872291565, 0.09850788861513138, 0.013066386803984642, 0.05027954652905464, 0.10465669631958008, 0.04533415287733078, 0.47037968039512634, 0.004505114629864693, 0.12196572870016098, 0.18816377222537994, 4.326914222474443e-06, 0.00023807807883713394, 0.00026310785324312747, 8.714396244613454e-06, 1.617559973965399e-05, 0.0001319001312367618, 0.0005945482989773154, 0.000823884445708245, 0.0008506007143296301, 1.7805428797146305e-05, 2.734714854568665e-08, 2.8855724849563558e-06, 4.891938442597166e-05, 0.0011682395124807954, 8.529372053089901e-07, 0.00017029111040756106, 1.0359013202787537e-07, 7.06834313302096e-10, 1.0861956525332062e-06, 0.0008713650749996305, 0.596385657787323, 0.0009257638594135642, NaN, NaN, NaN], [0.0018758929800242186, 0.019657986238598824, 0.1020394116640091, 0.033738646656274796, 0.024869924411177635, 0.012215637601912022, 0.015038376674056053, 0.002843664726242423, 0.02175789885222912, 0.01636381261050701, 0.01989913359284401, 0.01190999522805214, 0.00020280842727515846, 0.0016855570720508695, 0.17570628225803375, 1.4773272882795396e-10, 2.3448599506536993e-08, 6.434380566133768e-07, 3.8027360460546333e-07, 2.454226432746509e-06, 5.541529457531169e-09, 3.5226184991188347e-06, 2.5443886997322807e-08, 1.7749154721968807e-05, 1.8393259137994278e-09, 4.026108439691978e-12, 6.382850692432385e-09, 1.7809153263215194e-08, 8.996512974590587e-07, 0.00010512088192626834, 1.1464897607671443e-11, 2.794342757184154e-09, 2.4549680847631107e-15, 9.933188299671158e-11, 7.3009864820505754e-09, 8.105817687464878e-05, 0.2077004611492157, 2.0097606466151774e-05, NaN, NaN], [0.0009206020040437579, 0.08179444819688797, 0.00436751963570714, 0.003652991494163871, 0.019383452832698822, 0.008280212059617043, 0.016885409131646156, 0.10377784073352814, 0.023152435198426247, 0.0037028237711638212, 0.0001251623034477234, 0.018928401172161102, 0.009926089085638523, 0.002465219935402274, 0.21539123356342316, 1.1257004341538607e-14, 1.3137036347643516e-08, 4.6611327775281097e-07, 3.0405328743654536e-06, 1.5423474053477548e-07, 2.520166120234535e-08, 3.4643394819511286e-09, 1.1558090484697914e-08, 1.417677253812144e-06, 9.112129362165433e-08, 4.2694305868451465e-09, 3.7723260626343347e-10, 4.1450526344632976e-10, 2.7357388923676673e-11, 6.112880441833113e-07, 3.9687514799879864e-05, 8.382351063263016e-11, 8.293656039715103e-11, 4.97465783844131e-12, 4.144883221368634e-12, 1.4191136113450575e-11, 2.5566061594872735e-05, 0.4056495428085327, 4.4409513066057116e-05, NaN], [0.0005496710073202848, 0.039492249488830566, 0.016358638182282448, 0.007983607240021229, 0.006420070305466652, 0.0012171968119218946, 0.003928476013243198, 0.005028040148317814, 0.010722441598773003, 0.0025004756171256304, 0.015696601942181587, 0.006085758097469807, 0.0033880609553307295, 0.0056163351982831955, 0.1572248637676239, 9.215334861117716e-19, 2.6557794852166694e-10, 5.799645919069008e-07, 1.003176621633406e-11, 7.217926736302616e-07, 4.876178394397357e-08, 8.254863459455919e-11, 1.424103456687531e-12, 1.1857503423584603e-08, 1.3074058502482444e-09, 8.580362115262474e-12, 5.829819293978744e-09, 1.8017319407259702e-12, 9.234832950427707e-14, 3.576115098491428e-11, 1.9265784523270213e-09, 1.8997316146851517e-06, 1.949248054633479e-11, 8.860704392432694e-10, 2.8198800851872777e-14, 5.674391451236226e-15, 1.0258181110112119e-10, 6.93914080329705e-06, 0.25534507632255554, 2.742740150551981e-07]], [[0.130781888961792, 0.31469303369522095, 0.10550640523433685, 0.05234318599104881, 0.073336161673069, 0.022349786013364792, 0.04807984083890915, 0.1931842416524887, 0.06399697810411453, 0.042083337903022766, 0.026750531047582626, 0.11997608095407486, 0.008983415551483631, 0.03431839123368263, 0.019280044361948967, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1582711637020111, 0.14862558245658875, 0.20016248524188995, 0.08876624703407288, 0.11006557196378708, 0.14632253348827362, 0.04025046527385712, 0.010204354301095009, 0.017868297174572945, 0.059372395277023315, 0.02111685276031494, 0.04181571304798126, 0.025184988975524902, 0.09681157767772675, 0.11611668020486832, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.23875439167022705, 0.3084685802459717, 0.14188633859157562, 0.026331612840294838, 0.0149313323199749, 0.09176106750965118, 0.03131069242954254, 0.10051372647285461, 0.03149634972214699, 0.11085867136716843, 0.014410188421607018, 0.02796255424618721, 0.034816499799489975, 0.025807565078139305, 0.01846306212246418, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3404518961906433, 0.24260303378105164, 0.15383434295654297, 0.17020593583583832, 0.011800014413893223, 0.014385397545993328, 0.09441643208265305, 0.12204645574092865, 0.13843503594398499, 0.045293405652046204, 0.010667533613741398, 0.19693949818611145, 0.10281307995319366, 0.01422606036067009, 0.06984427571296692, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.002873742487281561, 0.008706165477633476, 0.35573768615722656, 0.0015586970839649439, 0.015496796928346157, 0.003392455168068409, 0.01149011217057705, 0.01891980692744255, 0.016394488513469696, 0.003960000351071358, 0.0035995631478726864, 0.008501716889441013, 0.018164046108722687, 0.004727588500827551, 0.013562880456447601, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.044807154685258865, 0.02788197249174118, 0.03947468474507332, 0.1271299421787262, 0.17640650272369385, 0.25110092759132385, 0.08349309861660004, 0.02069718949496746, 0.45751577615737915, 0.039922621101140976, 0.1781769096851349, 0.002931024879217148, 0.16567888855934143, 0.1177627220749855, 0.5156693458557129, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.005990047473460436, 0.04782475531101227, 0.01399919856339693, 0.010489771142601967, 0.06132129579782486, 0.030459748581051826, 0.010153756476938725, 0.3387801945209503, 0.06446883827447891, 0.007243711035698652, 0.00693717272952199, 0.020023254677653313, 0.007285784464329481, 0.009139767847955227, 0.0044054011814296246, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.020405659452080727, 0.00729386368766427, 0.06661678105592728, 0.08295443654060364, 0.20373474061489105, 0.3448184132575989, 0.04295210912823677, 0.20947468280792236, 0.03081577830016613, 0.010805373080074787, 0.17521467804908752, 0.06567652523517609, 0.012400656938552856, 0.10652147233486176, 0.07385163754224777, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.21573591232299805, 0.13175059854984283, 0.04085814207792282, 0.04119405150413513, 0.03551999852061272, 0.023009058088064194, 0.2751774191856384, 0.047030266374349594, 0.14272502064704895, 0.20153193175792694, 0.09575672447681427, 0.11327007412910461, 0.008532780222594738, 0.053245026618242264, 0.08952803909778595, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2778390347957611, 0.11423225700855255, 0.3034791946411133, 0.34643107652664185, 0.5395972728729248, 0.06785042583942413, 0.13029156625270844, 0.18737749755382538, 0.029348008334636688, 0.16667678952217102, 0.021040884777903557, 0.008728248998522758, 0.037633832544088364, 0.02033349499106407, 0.03947347402572632, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4898838996887207, 0.08082167059183121, 0.07362432777881622, 0.02171795442700386, 0.1333591789007187, 0.09000474214553833, 0.13501934707164764, 0.03979193791747093, 0.19113953411579132, 0.13522492349147797, 0.16557832062244415, 0.16255514323711395, 0.07687958329916, 0.15948235988616943, 0.09843874722719193, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.045906297862529755, 0.18602333962917328, 0.4082620143890381, 0.010370302945375443, 0.04507172852754593, 0.19693265855312347, 0.04021843150258064, 0.027866821736097336, 0.1546991914510727, 0.33766424655914307, 0.09260500222444534, 0.05066358670592308, 0.05655887722969055, 0.13157807290554047, 0.06850539147853851, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.020344020798802376, 0.0030158585868775845, 0.004445259924978018, 0.022628312930464745, 0.030150510370731354, 0.027700912207365036, 0.026311388239264488, 0.012862108647823334, 0.07009940594434738, 0.24656175076961517, 0.10596039146184921, 0.1143152266740799, 0.3679012656211853, 0.0068145813420414925, 0.04171491786837578, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.004749340936541557, 0.00182742765173316, 0.0021293568424880505, 0.00394084258005023, 0.004750867374241352, 5.3125138947507367e-05, 0.0026011874433606863, 0.000718552153557539, 0.002356230979785323, 0.00125187449157238, 0.0021339249797165394, 0.00044074622564949095, 0.2141493707895279, 0.0029175111558288336, 0.00477015832439065, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.12991508841514587, 0.06724811345338821, 0.06397818773984909, 0.15923364460468292, 0.2566852867603302, 0.07963784784078598, 0.09182894974946976, 0.040824584662914276, 0.21298912167549133, 0.2517295181751251, 0.2285410314798355, 0.11115844547748566, 0.1010512113571167, 0.3968040943145752, 0.1870165765285492, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.09555985033512115, 0.6603901982307434, 0.4109249413013458, 0.6857163310050964, 0.16377028822898865, 0.1341286301612854, 0.19969937205314636, 0.28269705176353455, 0.14764364063739777, 0.41980865597724915, 0.4319525361061096, 0.3789142668247223, 0.49345141649246216, 0.26345306634902954, 0.00909768883138895, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1460653841495514, 0.2758752405643463, 0.2826583981513977, 0.551855206489563, 0.05612415447831154, 0.19304026663303375, 0.0849798247218132, 0.038316093385219574, 0.02312053181231022, 0.46154478192329407, 0.36433619260787964, 0.35877159237861633, 0.1596277803182602, 0.0554661750793457, 6.483463948825374e-05, 0.0002614231198094785, 0.183704674243927, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [3.716628270922229e-05, 1.9402585849093157e-07, 1.0113188182003796e-05, 6.318590021692216e-05, 6.053787728887983e-07, 2.5790013751247898e-06, 0.00022986173280514777, 1.074662236533186e-06, 6.082240361138247e-06, 3.35614299729059e-06, 2.225729804194998e-05, 7.863033715693746e-06, 1.555537892272696e-06, 3.881560041918419e-05, 0.23657216131687164, 1.3331101555991154e-08, 0.003119559260085225, 0.19454506039619446, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.6150763630867004, 0.041665952652692795, 0.4174444377422333, 0.4949702024459839, 0.20794649422168732, 0.3307763934135437, 0.8098993897438049, 0.2721010744571686, 0.7274996042251587, 0.4779607057571411, 0.6233283281326294, 0.7560765147209167, 0.3628612458705902, 0.7672091722488403, 5.392584171204362e-06, 1.1244888353800775e-09, 0.0005117341643199325, 0.15345418453216553, 0.0018621939234435558, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [5.640763447445352e-06, 2.5884469323500525e-07, 1.2724142379738623e-06, 8.170181899913587e-06, 1.2345621769327408e-07, 1.310836523771286e-07, 1.02673438959755e-05, 9.661080184741877e-07, 6.520539272969472e-07, 7.602448022225872e-07, 2.058099425994442e-06, 6.885502301656743e-08, 1.0175665465794737e-06, 1.7383708836860023e-05, 0.20754273235797882, 2.882708471929618e-08, 0.0006895777769386768, 0.008299488574266434, 0.004234161227941513, 0.26378652453422546, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [9.27566077280062e-07, 5.395870630309219e-07, 1.8455818917573197e-07, 1.2775643654094893e-06, 2.105696061960316e-08, 3.1680112755338996e-08, 6.263408067752607e-06, 4.3284012463118415e-07, 1.918825773827848e-06, 1.694104128091567e-07, 3.363936968980852e-07, 9.135120215830739e-09, 4.4058825920956224e-08, 7.840970965844463e-07, 0.18219269812107086, 6.507164653157815e-05, 0.0030905166640877724, 0.269605815410614, 0.06594818085432053, 0.07055308669805527, 0.24370616674423218, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.7144812345504761, 0.6739043593406677, 0.2952970862388611, 0.49478814005851746, 0.17151717841625214, 0.06989942491054535, 0.5132517218589783, 0.30886489152908325, 0.5621734261512756, 0.5728412866592407, 0.576314389705658, 0.34687095880508423, 0.25617536902427673, 0.29690253734588623, 7.371841547865188e-06, 5.806248736917041e-05, 0.0008924558642320335, 0.00047033390728756785, 0.003593915607780218, 0.044251326471567154, 0.18547922372817993, 0.19724349677562714, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.6291437745094299, 0.5982875823974609, 0.4885888695716858, 0.5792520046234131, 0.2514877915382385, 0.5298613905906677, 0.11972777545452118, 0.6076628565788269, 0.04243328422307968, 0.5940482020378113, 0.6775911450386047, 0.3496588468551636, 0.4937344789505005, 0.40163323283195496, 2.9517783332266845e-05, 0.03321969881653786, 0.1786998063325882, 0.0021111152600497007, 0.00015362887643277645, 0.0013223892310634255, 0.01674751006066799, 0.27181917428970337, 0.0704144611954689, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.6414378881454468, 0.20530864596366882, 0.8448930978775024, 0.5841984748840332, 0.48009997606277466, 0.48003992438316345, 0.4468145966529846, 0.036266062408685684, 0.3466547429561615, 0.521195650100708, 0.7532409429550171, 0.14529024064540863, 0.3844791650772095, 0.46825459599494934, 2.1059213395346887e-05, 0.0005316429305821657, 0.0021434861700981855, 0.0005638045258820057, 2.0347550162114203e-05, 8.372889715246856e-05, 0.0012170294066891074, 0.0006328476592898369, 0.0015302025713026524, 0.2731996476650238, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.7977450489997864, 0.5162288546562195, 0.513008177280426, 0.6203657984733582, 0.04621165990829468, 0.2237500697374344, 0.10730908066034317, 0.17203836143016815, 0.028481170535087585, 0.5342445969581604, 0.7256113290786743, 0.5827998518943787, 0.755642294883728, 0.511749804019928, 0.00015279543003998697, 3.384976253073546e-06, 0.0032942681573331356, 0.003179847961291671, 0.0003072107210755348, 3.0923787562642246e-05, 0.0003082206822000444, 0.0026841319631785154, 0.011449099518358707, 0.2928124964237213, 0.0015787724405527115, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5001324415206909, 0.7283154129981995, 0.6225411295890808, 0.5096700191497803, 0.4470505714416504, 0.6475648880004883, 0.4919697046279907, 0.42729777097702026, 0.22966071963310242, 0.4533919394016266, 0.5539101958274841, 0.2698501944541931, 0.3532210886478424, 0.2643750309944153, 2.9741322578047402e-05, 4.910896677756682e-05, 0.01189705915749073, 0.0036808690056204796, 0.006090851966291666, 0.0029882052913308144, 0.006760776974260807, 0.0002592294185888022, 0.0001972121826838702, 0.15788163244724274, 0.14973512291908264, 0.14614373445510864, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.42266348004341125, 0.20205438137054443, 0.42841264605522156, 0.6724829077720642, 0.29094210267066956, 0.4464052617549896, 0.24126748740673065, 0.22405968606472015, 0.21308888494968414, 0.3085091710090637, 0.4672502279281616, 0.14604215323925018, 0.09687051922082901, 0.12085973471403122, 2.7047781259170733e-05, 7.539001671830192e-05, 0.036947283893823624, 0.01112621370702982, 0.04119950905442238, 0.06979847699403763, 0.01383589580655098, 0.008948443457484245, 9.020609286380932e-05, 0.0005221512983553112, 0.34183818101882935, 0.12104173004627228, 0.027292484417557716, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5077533721923828, 0.4866065979003906, 0.8742184638977051, 0.805268406867981, 0.8406472206115723, 0.45863693952560425, 0.3596036732196808, 0.36316972970962524, 0.38783764839172363, 0.03767421096563339, 0.43841618299484253, 0.3401361405849457, 0.3197961747646332, 0.20812755823135376, 7.5720936365542e-06, 5.4811065638205037e-05, 0.015359039418399334, 0.005874635651707649, 0.024854328483343124, 0.16572602093219757, 0.13195344805717468, 0.08553953468799591, 0.00124072446487844, 0.0008515206864103675, 0.0025517549365758896, 0.03817262500524521, 0.1957935392856598, 0.020919298753142357, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12348711490631104, 0.49926623702049255, 0.1342328041791916, 0.07936512678861618, 0.11133208125829697, 0.032334309071302414, 0.028592387214303017, 0.036310840398073196, 0.036252155900001526, 0.10585709661245346, 0.19267472624778748, 0.34429997205734253, 0.16909800469875336, 0.2464863359928131, 3.1697504709882196e-06, 3.401398498681374e-05, 0.0008079431718215346, 0.00045223115012049675, 0.00013304724416229874, 0.0006849576020613313, 0.009534466080367565, 0.010466179810464382, 0.00030334663460962474, 0.00033610902028158307, 2.1021634893259034e-05, 6.891421071486548e-05, 0.0028196852654218674, 0.3685440421104431, 0.0008976467652246356, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [4.5035082507638435e-07, 4.8253248507990065e-08, 2.1990938847693542e-08, 4.3766593194050074e-07, 1.1283042766763174e-07, 2.4235429663121977e-08, 4.6985369408503175e-06, 1.5805973418991925e-07, 1.1619090578562918e-08, 1.9516033233912822e-08, 1.8456361772223318e-07, 2.2261544074808626e-07, 2.278205402106437e-09, 7.143006541809882e-07, 0.21044957637786865, 0.0012722803512588143, 0.07485485821962357, 0.004568059463053942, 0.008557068184018135, 0.04491077736020088, 0.010689688846468925, 0.010801602154970169, 0.015439217910170555, 0.001288879313506186, 0.032191790640354156, 9.430324280401692e-05, 0.0010071481810882688, 0.03593403846025467, 0.015365669503808022, 0.28865233063697815, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.71169513463974, 0.2780396640300751, 0.44078493118286133, 0.7963916063308716, 0.6933308839797974, 0.5056049823760986, 0.7329073548316956, 0.810703694820404, 0.551677942276001, 0.6459015607833862, 0.6943050622940063, 0.2817550301551819, 0.10247289389371872, 0.7378624677658081, 8.274764695670456e-06, 0.0003195737663190812, 0.0016381103778257966, 0.001899963477626443, 0.000450764549896121, 0.0029568641912192106, 0.0004077073244843632, 0.006739944685250521, 5.316005626809783e-05, 0.000977654941380024, 0.00033480822457931936, 1.5544836060144007e-05, 5.177688763069455e-06, 0.000280524865956977, 8.569184137741104e-05, 0.19435854256153107, 0.0009946423815563321, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.723514199256897, 0.08602748066186905, 0.6093902587890625, 0.8655006289482117, 0.42677831649780273, 0.03823491558432579, 0.30262306332588196, 0.036271825432777405, 0.12300263345241547, 0.2776595950126648, 0.07632125169038773, 0.06917709112167358, 0.14498986303806305, 0.06881040334701538, 2.5871422622003593e-06, 0.0004552309401333332, 0.00916277151554823, 0.2859989106655121, 0.028668222948908806, 0.004703177139163017, 0.013283651322126389, 0.011935138143599033, 0.00041849465924315155, 0.021506765857338905, 0.0005354905733838677, 2.3408898414345458e-05, 5.557515123655321e-06, 4.006853941973532e-06, 0.000782388960942626, 0.032734211534261703, 0.33600685000419617, 0.05645810067653656, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.7111753225326538, 0.8019941449165344, 0.7984396815299988, 0.6959745287895203, 0.34880974888801575, 0.5955101251602173, 0.6658092141151428, 0.5378626585006714, 0.35595381259918213, 0.5855972766876221, 0.5757258534431458, 0.133575439453125, 0.3884122669696808, 0.11617641150951385, 8.579120731155854e-06, 0.001615832676179707, 0.0592908076941967, 0.004439341835677624, 0.0221478920429945, 0.05761101841926575, 0.08599329739809036, 0.009327156469225883, 0.0014337823959067464, 0.22479815781116486, 0.007599419914186001, 0.00010282513540005311, 0.003995772451162338, 0.0007532926392741501, 0.0001985877170227468, 0.042725738137960434, 0.609107255935669, 0.032340146601200104, 0.2600889503955841, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.43439850211143494, 0.1714652180671692, 0.4214288294315338, 0.6560039520263672, 0.15961043536663055, 0.25604698061943054, 0.26937225461006165, 0.1702796220779419, 0.22940081357955933, 0.327440470457077, 0.3977930247783661, 0.08873222768306732, 0.13160161674022675, 0.07058954238891602, 2.3103428247850388e-05, 0.0007894318550825119, 0.08912800997495651, 0.00870462041348219, 0.062210533767938614, 0.21669252216815948, 0.04955689236521721, 0.12036743760108948, 0.001276280265301466, 0.002290783217176795, 0.4637441337108612, 0.041003014892339706, 0.007595454342663288, 0.0049859327264130116, 0.030789200216531754, 0.01441932376474142, 0.02666427381336689, 0.013092019595205784, 0.22824719548225403, 0.07290598005056381, NaN, NaN, NaN, NaN, NaN, NaN], [0.48717519640922546, 0.4504354000091553, 0.9026078581809998, 0.8262973427772522, 0.8697957992553711, 0.4322546720504761, 0.47440072894096375, 0.40584686398506165, 0.6554202437400818, 0.04447361081838608, 0.5114831924438477, 0.4020007252693176, 0.3586147725582123, 0.19603849947452545, 5.424046776170144e-06, 4.2991967347916216e-05, 0.006631283089518547, 0.0006027332856319845, 0.004053125157952309, 0.03894652798771858, 0.031787656247615814, 0.10168109834194183, 0.004267984535545111, 0.002045443281531334, 0.0010633694473654032, 0.005091637372970581, 0.031351421028375626, 6.663963722530752e-05, 0.09428737312555313, 0.0008465268765576184, 0.00024849644978530705, 0.002269570017233491, 0.01905866153538227, 0.2164839655160904, 0.010082208551466465, NaN, NaN, NaN, NaN, NaN], [0.09346597641706467, 0.41046077013015747, 0.13097965717315674, 0.06711046397686005, 0.09538185596466064, 0.021688319742679596, 0.027864748612046242, 0.029869627207517624, 0.07506763935089111, 0.13717295229434967, 0.21322546899318695, 0.3559926152229309, 0.19059841334819794, 0.24045485258102417, 2.0756003777933074e-06, 1.1191940757271368e-05, 0.0006002296577207744, 0.0002709901600610465, 9.913583926390857e-05, 0.0001758227008394897, 0.0029332106932997704, 0.008675863035023212, 0.0011328428518027067, 0.0023299665190279484, 6.693489558529109e-05, 0.00013525204849429429, 0.0013442488852888346, 0.022858861833810806, 2.321010106243193e-05, 0.0010626229923218489, 2.5993340386776254e-05, 3.972689592046663e-05, 5.326797690941021e-05, 0.0033412689808756113, 0.35271701216697693, 0.0008956229430623353, NaN, NaN, NaN, NaN], [4.6634454520244617e-07, 5.573102512812511e-08, 2.3018172257138758e-08, 3.889360016273713e-07, 9.709493298259986e-08, 2.4796046105279856e-08, 7.192591056082165e-06, 1.7916640615567303e-07, 1.8580767147113875e-08, 3.5935642017648206e-08, 2.774728216081712e-07, 3.801677337378351e-07, 2.8816848907098347e-09, 9.808413778955583e-07, 0.2028982788324356, 0.00036489564809016883, 0.07616367936134338, 0.00673737283796072, 0.011110173538327217, 0.021392904222011566, 0.010494116693735123, 0.006134945899248123, 0.015969248488545418, 0.005187375005334616, 0.12039955705404282, 0.0005341891082935035, 0.0022901638876646757, 0.027128320187330246, 0.005907480139285326, 0.033119603991508484, 0.002176248235628009, 0.0003625153622124344, 6.369769835146144e-05, 0.0007003483478911221, 0.03456505015492439, 0.01570759527385235, 0.28412890434265137, NaN, NaN, NaN], [0.6667957305908203, 0.327456533908844, 0.4202725291252136, 0.7458598613739014, 0.6837785840034485, 0.5435037612915039, 0.7794858813285828, 0.849186360836029, 0.6942030787467957, 0.7531007528305054, 0.7604266405105591, 0.4857816696166992, 0.12311270833015442, 0.7958275079727173, 7.400509275612421e-06, 3.192616713931784e-05, 0.00035208670306019485, 0.002478531561791897, 0.0006564928335137665, 0.0008886585710570216, 0.0005662215990014374, 0.0016915983287617564, 1.3900444173486903e-05, 0.0009738726075738668, 0.00042995362309738994, 8.639829320600256e-05, 1.4000924238644075e-05, 0.00033226466621272266, 2.9785558581352234e-05, 0.00921203475445509, 3.390025085536763e-06, 5.1574592362158e-05, 2.3835823412809987e-06, 1.9022172637050971e-06, 0.00016878120368346572, 9.063100151252002e-05, 0.20696188509464264, 0.001649125711992383, NaN, NaN], [0.704485297203064, 0.08825523406267166, 0.5944071412086487, 0.8510531783103943, 0.4262540936470032, 0.04518446326255798, 0.38849392533302307, 0.055145543068647385, 0.277063250541687, 0.40566664934158325, 0.09198901802301407, 0.13750647008419037, 0.24822941422462463, 0.1165834292769432, 3.5331499930180144e-06, 0.00019471753330435604, 0.003537738462910056, 0.2800489366054535, 0.036592625081539154, 0.002127013634890318, 0.024595409631729126, 0.008275463245809078, 0.00023266732750926167, 0.021680369973182678, 0.0005173377576284111, 7.175304199336097e-05, 2.6857771445065737e-05, 1.6371919627999887e-05, 0.0012281013187021017, 0.011112956330180168, 0.058813560754060745, 0.0009629606502130628, 1.1531898962857667e-05, 4.947432444168953e-06, 2.475359451636905e-06, 0.0005685617215931416, 0.0267820842564106, 0.3296748399734497, 0.06147307902574539, NaN], [0.5231692790985107, 0.6706213355064392, 0.7785398364067078, 0.7122241258621216, 0.34260621666908264, 0.579698920249939, 0.5863306522369385, 0.4822496175765991, 0.5804131031036377, 0.7801564335823059, 0.7983464002609253, 0.22512593865394592, 0.4790371060371399, 0.2274763584136963, 1.8860177078749985e-05, 3.20236104300875e-08, 0.00013383101031649858, 0.00029007354169152677, 0.002788462908938527, 0.0014709108509123325, 0.0009710633894428611, 0.0001290659129153937, 2.0881772798020393e-05, 7.236683813971467e-06, 3.12792144541163e-05, 7.099155482137576e-05, 3.213396485080011e-05, 3.9666349039180204e-05, 0.00022854047711007297, 0.0037343965377658606, 1.487573445047019e-05, 0.00019343644089531153, 8.10168421594426e-05, 1.1448363693489227e-05, 3.5921341350331204e-06, 2.216967368440237e-05, 0.0017730530817061663, 0.0001526248233858496, 0.009769736789166927, 0.4419056475162506]], [[0.06147387623786926, 0.0657946914434433, 0.22564710676670074, 0.1299343705177307, 0.021580645814538002, 0.08992400765419006, 0.025479430332779884, 0.04823821783065796, 0.05891237407922745, 0.016958819702267647, 0.0021926285699009895, 0.017513686791062355, 0.09859969466924667, 0.16368542611598969, 0.038398925215005875, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.029852252453565598, 0.26626214385032654, 0.14803646504878998, 0.038784727454185486, 0.07803148031234741, 0.006210723891854286, 0.0026457132771611214, 0.006018034182488918, 0.05453306809067726, 0.002730109030380845, 0.015730326995253563, 0.0017557059181854129, 0.034912969917058945, 0.03208531066775322, 0.03983413055539131, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01053018867969513, 0.02744918502867222, 0.2530466914176941, 0.05846027657389641, 0.1744728684425354, 0.011957419104874134, 0.003304906887933612, 0.00205883732996881, 0.00874510407447815, 0.0014524421421810985, 0.0009729861048981547, 0.0026561047416180372, 0.0023208027705550194, 0.0038251704536378384, 0.005045189522206783, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.016039762645959854, 0.05755838379263878, 0.10756286233663559, 0.03799062967300415, 0.5738711953163147, 0.061907339841127396, 0.128611221909523, 0.01847657933831215, 0.06501789391040802, 0.015564735978841782, 0.0016139671206474304, 0.014343881979584694, 0.020734043791890144, 0.14008449018001556, 0.13515408337116241, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.005847899243235588, 0.11914067715406418, 0.01715121790766716, 0.3517457842826843, 0.0661543607711792, 0.07493122667074203, 0.012425812892615795, 0.11745280772447586, 0.08440648764371872, 0.020029406994581223, 0.05165768414735794, 0.04094480350613594, 0.024548601359128952, 0.005826729815453291, 0.13841456174850464, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.015926362946629524, 0.007578620687127113, 0.1226087138056755, 0.030128292739391327, 0.03851892054080963, 0.3367418944835663, 0.01694057136774063, 0.09829536825418472, 0.0361555740237236, 0.10537439584732056, 0.007450005039572716, 0.029753634706139565, 0.22920416295528412, 0.01793695241212845, 0.05258304625749588, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01326388493180275, 0.05337870866060257, 0.047661036252975464, 0.08615607023239136, 0.12425915151834488, 0.4180251955986023, 0.04702466353774071, 0.0717325434088707, 0.05138256773352623, 0.06877672672271729, 0.0152205191552639, 0.0719875767827034, 0.1666427105665207, 0.13322126865386963, 0.053655143827199936, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.026802292093634605, 0.003955241292715073, 0.0206829272210598, 0.02742936834692955, 0.06016179919242859, 0.15127348899841309, 0.06774158030748367, 0.2981398105621338, 0.05239749699831009, 0.09365928173065186, 0.035629644989967346, 0.020771589130163193, 0.13655303418636322, 0.012941722758114338, 0.05640798062086105, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.06469012051820755, 0.1851334124803543, 0.08788572251796722, 0.19977343082427979, 0.00846380740404129, 0.03702360764145851, 0.0876760184764862, 0.046302031725645065, 0.11564433574676514, 0.05180440843105316, 0.49518024921417236, 0.1649368405342102, 0.030481798574328423, 0.10461966693401337, 0.07739346474409103, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.020106524229049683, 0.01925482228398323, 0.006043681409209967, 0.01652396097779274, 0.001572003006003797, 0.005779887083917856, 0.015335858799517155, 0.03537710756063461, 0.009967570193111897, 0.09144406765699387, 0.43651703000068665, 0.2613205015659332, 0.0483890138566494, 0.06553913652896881, 0.055434126406908035, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07980967313051224, 0.14815203845500946, 0.09271827340126038, 0.004086778499186039, 0.010790406726300716, 0.0747552439570427, 0.10995902121067047, 0.04728228971362114, 0.1809520274400711, 0.025821411982178688, 0.06657237559556961, 0.1431768387556076, 0.19449584186077118, 0.20780201256275177, 0.10148976743221283, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05537823587656021, 0.008725662715733051, 0.0058344281278550625, 0.029011448845267296, 0.048424966633319855, 0.047911662608385086, 0.16901308298110962, 0.17019973695278168, 0.011648884043097496, 0.08953043073415756, 0.5360274910926819, 0.10330803692340851, 0.078437939286232, 0.12202966213226318, 0.11905822902917862, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01546903420239687, 0.0005347061669453979, 0.0015839362749829888, 0.053056132048368454, 0.23614321649074554, 0.013318118639290333, 0.051473915576934814, 0.011966699734330177, 0.007302975282073021, 0.09275621920824051, 0.06646261364221573, 0.010813506320118904, 0.13289499282836914, 0.22826357185840607, 0.04386172071099281, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.009458722546696663, 0.0058342707343399525, 0.012789146974682808, 0.005895438138395548, 0.026010286062955856, 0.057482823729515076, 0.005663284566253424, 0.005727604031562805, 0.0033144087065011263, 0.011671853251755238, 0.00424896739423275, 0.056589994579553604, 0.20401620864868164, 0.03777612745761871, 0.03114682249724865, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0012354525970295072, 0.034024473279714584, 0.10020612925291061, 0.02267461270093918, 0.08676987141370773, 0.14216794073581696, 0.0033775768242776394, 0.07320579141378403, 0.07390473037958145, 0.0168889332562685, 0.00386308366432786, 0.02569040097296238, 0.24664165079593658, 0.2674221694469452, 0.014589445665478706, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.12044757604598999, 0.22699733078479767, 0.3625817894935608, 0.18942511081695557, 0.468371719121933, 0.5971034169197083, 0.5581120252609253, 0.29680517315864563, 0.4773823618888855, 0.4035939574241638, 0.3702273666858673, 0.3751682937145233, 0.267861545085907, 0.4069889783859253, 0.040672045201063156, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0243044663220644, 0.4273812174797058, 0.5286219716072083, 0.05566978082060814, 0.4582313597202301, 0.5064847469329834, 0.09591992199420929, 0.1787465512752533, 0.7349562644958496, 0.00692495983093977, 0.04355573281645775, 0.04027868062257767, 0.03415951877832413, 0.02788657508790493, 0.03653726726770401, 0.07662782073020935, 0.14776498079299927, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1999487727880478, 0.02213704027235508, 0.750217854976654, 0.5677059292793274, 0.8556592464447021, 0.6869031190872192, 0.2201639711856842, 0.6947058439254761, 0.2711787521839142, 0.21462410688400269, 0.3783731162548065, 0.39328378438949585, 0.3796219229698181, 0.27560317516326904, 0.052095912396907806, 0.0006832284270785749, 0.003495789598673582, 0.19430121779441833, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17733721435070038, 0.1195838525891304, 0.4294462502002716, 0.41039443016052246, 0.45686641335487366, 0.5433338284492493, 0.08341590315103531, 0.5749803781509399, 0.0773383378982544, 0.2876206338405609, 0.19534848630428314, 0.10015372186899185, 0.2102438062429428, 0.04678432643413544, 0.044711172580718994, 0.00020953372586518526, 0.007476589176803827, 0.1521030217409134, 0.003494996577501297, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4523387849330902, 0.8917949795722961, 0.4903220534324646, 0.5869925022125244, 0.47626572847366333, 0.006232858635485172, 0.41125378012657166, 0.13404546678066254, 0.6460333466529846, 0.32553666830062866, 0.3429105877876282, 0.031081799417734146, 0.42998504638671875, 0.16709895431995392, 0.08821719139814377, 0.00048688906827010214, 0.0011088894680142403, 0.0024602855555713177, 0.0005520267877727747, 0.26744863390922546, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.49767979979515076, 0.7566660642623901, 0.25263193249702454, 0.4967457056045532, 0.47193706035614014, 0.006824302952736616, 0.2858791947364807, 0.18135732412338257, 0.4390898644924164, 0.7668571472167969, 0.15391138195991516, 0.08414287865161896, 0.5640745759010315, 0.35628020763397217, 0.09142898768186569, 0.0004194685607217252, 0.0005068383179605007, 0.026896899566054344, 0.0004147894505877048, 0.006156287621706724, 0.4387049376964569, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18697474896907806, 0.23196713626384735, 0.23554784059524536, 0.34321168065071106, 0.5325552225112915, 0.15430577099323273, 0.2887123227119446, 0.4957616627216339, 0.36584702134132385, 0.2891024053096771, 0.08069057762622833, 0.18119029700756073, 0.4536079466342926, 0.16425864398479462, 0.03777371346950531, 1.0518371709622443e-05, 5.5142045312095433e-05, 0.016997506842017174, 3.693701364682056e-05, 0.0006244040559977293, 0.21657241880893707, 0.01345360092818737, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17079660296440125, 0.16765500605106354, 0.28291502594947815, 0.16039209067821503, 0.2695491909980774, 0.16163654625415802, 0.08897912502288818, 0.28747832775115967, 0.8989478349685669, 0.26775097846984863, 0.17184530198574066, 0.3264879584312439, 0.31386569142341614, 0.1549917310476303, 0.05264737084507942, 0.3619365394115448, 0.25655418634414673, 0.3611752688884735, 0.14710570871829987, 0.018539972603321075, 0.21814967691898346, 0.09323819726705551, 0.01780291646718979, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04084352031350136, 0.5361505150794983, 0.018223807215690613, 0.03828004375100136, 0.3140276074409485, 0.08277524262666702, 0.07094793766736984, 0.012667819857597351, 0.3304368853569031, 0.10053964704275131, 0.03868165612220764, 0.31755131483078003, 0.22644393146038055, 0.07613880187273026, 0.12961620092391968, 0.004012200981378555, 0.004658036399632692, 0.017421945929527283, 0.0026806569658219814, 0.590861439704895, 0.051964171230793, 0.007618917152285576, 0.0007336572161875665, 0.12340892106294632, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07373615354299545, 0.19122207164764404, 0.06966950744390488, 0.01624569669365883, 0.017842771485447884, 0.2144099771976471, 0.24285149574279785, 0.3761756718158722, 0.8141085505485535, 0.27487871050834656, 0.09974052757024765, 0.10127317160367966, 0.16323235630989075, 0.21032299101352692, 0.10343435406684875, 0.44725751876831055, 0.6053639054298401, 0.07041247189044952, 0.07085516303777695, 0.003138674655929208, 0.2879992425441742, 0.049135204404592514, 0.14297868311405182, 0.06008363142609596, 0.06304289400577545, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06651142984628677, 0.1456020176410675, 0.01741747185587883, 0.07566884905099869, 0.018790215253829956, 0.20801369845867157, 0.16892337799072266, 0.33592528104782104, 0.1834612786769867, 0.29906225204467773, 0.2579277753829956, 0.5998365879058838, 0.5642448663711548, 0.572043240070343, 0.0891154333949089, 0.7072809338569641, 0.7582566142082214, 0.16150887310504913, 0.18586905300617218, 0.015776842832565308, 0.08385244756937027, 0.32581770420074463, 0.5540359020233154, 0.13379113376140594, 0.0028463751077651978, 0.051922835409641266, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03234146162867546, 0.1962265521287918, 0.0277019701898098, 0.06972747296094894, 0.10650040954351425, 0.07791601866483688, 0.38205334544181824, 0.4892197549343109, 0.003444283502176404, 0.414199560880661, 0.16890743374824524, 0.4916560649871826, 0.8149713277816772, 0.7298122048377991, 0.14976243674755096, 0.4378974437713623, 0.10523661971092224, 0.014314417727291584, 0.30093127489089966, 0.06324318051338196, 0.08432605862617493, 0.2594241797924042, 0.6188808083534241, 0.3929617404937744, 0.00827555637806654, 0.07725780457258224, 0.06407154351472855, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07799918204545975, 0.2381461262702942, 0.01647050306200981, 0.08363308757543564, 0.05209676921367645, 0.02968973107635975, 0.11220219731330872, 0.32446831464767456, 0.1546868085861206, 0.06510066986083984, 0.1935844123363495, 0.5264057517051697, 0.34881067276000977, 0.6311980485916138, 0.09822507947683334, 0.2013174593448639, 0.5200937390327454, 0.3190821707248688, 0.5249915719032288, 0.18779213726520538, 0.1779765784740448, 0.29882070422172546, 0.5049118399620056, 0.06443758308887482, 0.007539320737123489, 0.16998757421970367, 0.031686559319496155, 0.3610091209411621, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1688770204782486, 0.13700607419013977, 0.20374003052711487, 0.12288741022348404, 0.15864238142967224, 0.039533428847789764, 0.12642242014408112, 0.35126128792762756, 0.365562379360199, 0.48467183113098145, 0.3247453570365906, 0.003142370842397213, 0.5969579219818115, 0.5533550977706909, 0.1647837609052658, 0.5546301603317261, 0.5397829413414001, 0.43089261651039124, 0.08987504988908768, 0.3114354610443115, 0.4812281131744385, 0.11215226352214813, 0.17198431491851807, 0.5790820121765137, 0.03648975491523743, 0.0541677288711071, 0.04165489599108696, 0.07749651372432709, 0.030232839286327362, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3052995800971985, 0.6539703607559204, 0.022321274504065514, 0.1902511715888977, 0.05963977798819542, 0.17083951830863953, 0.5218495726585388, 0.2573777139186859, 0.17107829451560974, 0.46426069736480713, 0.3389802873134613, 0.4338558316230774, 0.014936042949557304, 0.6202957630157471, 0.13899832963943481, 0.005376005079597235, 0.010858614929020405, 0.02991071715950966, 0.029742157086730003, 0.04020260274410248, 0.1695990264415741, 0.0604972317814827, 0.10318762809038162, 0.48727869987487793, 0.07163358479738235, 0.025501595810055733, 0.05125340074300766, 0.22269804775714874, 0.08394679427146912, 0.19870582222938538, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12219581007957458, 0.5012378692626953, 0.06702763587236404, 0.06399006396532059, 0.07401375472545624, 0.24048954248428345, 0.08739905059337616, 0.050457850098609924, 0.030934542417526245, 0.1506662517786026, 0.1536494344472885, 0.49837279319763184, 0.018043117597699165, 0.11216632276773453, 0.12939369678497314, 0.0006954512791708112, 0.0002132337394868955, 0.037006676197052, 0.0018452922813594341, 0.16118928790092468, 0.5505160689353943, 0.028353480622172356, 0.0021746368147432804, 0.027092093601822853, 0.0001434519508620724, 0.0029707583598792553, 4.2726576793938875e-05, 0.0012847317848354578, 0.0010433235438540578, 0.18891005218029022, 0.014656933024525642, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11525271832942963, 0.521948516368866, 0.007329752668738365, 0.008543604053556919, 0.05213259160518646, 0.04235774278640747, 0.2166471928358078, 0.528154194355011, 0.42159566283226013, 0.22446103394031525, 0.0032521234825253487, 0.5035390257835388, 0.365617960691452, 0.44961339235305786, 0.15735329687595367, 0.013874622993171215, 0.0695175901055336, 0.005752294324338436, 0.005697373300790787, 0.0021822804119437933, 0.02415846660733223, 0.00723307253792882, 0.3120453357696533, 0.016472192481160164, 0.004319194238632917, 0.041901107877492905, 0.7052133083343506, 0.0035930864978581667, 0.020578961819410324, 0.0021869041956961155, 0.0003597450559027493, 0.0005889505264349282, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03232282027602196, 0.08449342846870422, 0.004147443920373917, 0.050799064338207245, 0.037334948778152466, 0.08206064254045486, 0.07099173963069916, 0.19771835207939148, 0.021330662071704865, 0.08051090687513351, 0.1005825400352478, 0.700605034828186, 0.3027697801589966, 0.4364767074584961, 0.10480254143476486, 0.29724666476249695, 0.30918487906455994, 0.0693497508764267, 0.04026606306433678, 0.00593132060021162, 0.04497085511684418, 0.07199602574110031, 0.16270284354686737, 0.058071933686733246, 0.0005904879071749747, 0.0013724194141104817, 0.013050474226474762, 0.002609569113701582, 0.013482913374900818, 0.089314766228199, 0.03341012820601463, 0.21929660439491272, 0.006776490714401007, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.034268103539943695, 0.16091260313987732, 0.0168391652405262, 0.06967493146657944, 0.0915973111987114, 0.051104262471199036, 0.2385529726743698, 0.3295409679412842, 0.0004638703539967537, 0.22104156017303467, 0.13362999260425568, 0.5110065937042236, 0.7347238063812256, 0.7763577103614807, 0.15897347033023834, 0.3422777056694031, 0.07256462424993515, 0.012822822667658329, 0.21187257766723633, 0.060081083327531815, 0.09390594810247421, 0.19744858145713806, 0.5327264666557312, 0.3024030029773712, 0.013231869786977768, 0.1601967215538025, 0.04191795364022255, 0.5788960456848145, 0.791706383228302, 0.2698511779308319, 0.26516515016555786, 0.2890409529209137, 0.032140959054231644, 0.02436642162501812, NaN, NaN, NaN, NaN, NaN, NaN], [0.08530293405056, 0.1988343894481659, 0.010091865435242653, 0.07736483961343765, 0.030177433043718338, 0.023718634620308876, 0.06320804357528687, 0.20902810990810394, 0.020835628733038902, 0.026085397228598595, 0.10371798276901245, 0.427949994802475, 0.2465561032295227, 0.6410334706306458, 0.12414435297250748, 0.15722303092479706, 0.44676893949508667, 0.24300073087215424, 0.3980245292186737, 0.29666030406951904, 0.21130049228668213, 0.31708449125289917, 0.45276522636413574, 0.04954151436686516, 0.006070373114198446, 0.23888874053955078, 0.06321726739406586, 0.48237892985343933, 0.09136107563972473, 0.571183979511261, 0.36026179790496826, 0.0799446776509285, 0.1583012342453003, 0.025381257757544518, 0.5154083371162415, NaN, NaN, NaN, NaN, NaN], [0.17881684005260468, 0.09949745982885361, 0.17292529344558716, 0.14197823405265808, 0.0994792953133583, 0.022899990901350975, 0.07621151208877563, 0.20277591049671173, 0.059071850031614304, 0.23252709209918976, 0.2142648547887802, 0.0016634195344522595, 0.4786902368068695, 0.5105896592140198, 0.1802191287279129, 0.6566299200057983, 0.6752134561538696, 0.5489535927772522, 0.1520741730928421, 0.6433172821998596, 0.7151104211807251, 0.290630042552948, 0.3418242335319519, 0.686417818069458, 0.046654678881168365, 0.09611856192350388, 0.0634889155626297, 0.4891318380832672, 0.46607306599617004, 0.5581225156784058, 0.4337400496006012, 0.06152508407831192, 0.08386452496051788, 0.0397774837911129, 0.11068917065858841, 0.04009125009179115, NaN, NaN, NaN, NaN], [0.29184988141059875, 0.5299537181854248, 0.01714717224240303, 0.1581006944179535, 0.034420810639858246, 0.1480618417263031, 0.35555243492126465, 0.16130897402763367, 0.0352683924138546, 0.2384539395570755, 0.22334522008895874, 0.274210661649704, 0.008749962784349918, 0.5107676982879639, 0.16247788071632385, 0.0024060788564383984, 0.006098441779613495, 0.013975032605230808, 0.014695755206048489, 0.022452646866440773, 0.10514718294143677, 0.04751533642411232, 0.0609392412006855, 0.31799331307411194, 0.04427095875144005, 0.01951766200363636, 0.04202713817358017, 0.3371936082839966, 0.2731744647026062, 0.3478449583053589, 0.03363266587257385, 0.011759405955672264, 0.01767517626285553, 0.024101490154862404, 0.19511322677135468, 0.05518092215061188, 0.2097322940826416, NaN, NaN, NaN], [0.1536586880683899, 0.39876002073287964, 0.060627128928899765, 0.08434724807739258, 0.06138864532113075, 0.18170806765556335, 0.0558285117149353, 0.026850836351513863, 0.004648242145776749, 0.05450701341032982, 0.08679821342229843, 0.24500715732574463, 0.009806739166378975, 0.06359081715345383, 0.14997224509716034, 0.000109505133877974, 2.9198725314927287e-05, 0.01053665205836296, 0.0007290886132977903, 0.055462777614593506, 0.18011406064033508, 0.013305839151144028, 0.0007181179826147854, 0.008689867332577705, 4.760328374686651e-05, 0.0016827695071697235, 2.2867327061248943e-05, 0.000821226101834327, 0.0012459746794775128, 0.2353316843509674, 0.004575389437377453, 0.003901307238265872, 0.0009429306373931468, 1.1980442650383338e-05, 0.0003497266152407974, 0.00027309934375807643, 0.1965111494064331, 0.005757085047662258, NaN, NaN], [0.1216418668627739, 0.4058372378349304, 0.00597163662314415, 0.009731672704219818, 0.04685758054256439, 0.030955728143453598, 0.14503908157348633, 0.4122965633869171, 0.13539999723434448, 0.08889995515346527, 0.0017191163497045636, 0.24694381654262543, 0.23039060831069946, 0.2996818721294403, 0.1837962418794632, 0.0017744784709066153, 0.012578981928527355, 0.0015974465059116483, 0.002320722443982959, 0.0008557687979191542, 0.004459704738110304, 0.00322481500916183, 0.13683773577213287, 0.010506929829716682, 0.0027294831816107035, 0.03936534747481346, 0.7146239876747131, 0.0021277000196278095, 0.014929071068763733, 0.003117389976978302, 0.0010002683848142624, 0.0005979579291306436, 0.037009548395872116, 0.6984097361564636, 0.0021584301721304655, 0.012162267230451107, 0.002483450109139085, 0.00014705986541230232, 0.0003713203768711537, NaN], [0.2966727912425995, 0.1567845344543457, 0.07310101389884949, 0.14124755561351776, 0.2961083948612213, 0.07968501001596451, 0.06122228875756264, 0.14724984765052795, 0.06047076731920242, 0.055829375982284546, 0.06430483609437943, 0.11614347994327545, 0.15107537806034088, 0.15706941485404968, 0.12527146935462952, 0.10933294892311096, 0.0594157911837101, 0.01442565955221653, 0.027944112196564674, 0.24928514659404755, 0.3314722180366516, 0.036283038556575775, 0.01824975199997425, 0.03247179090976715, 0.02741291932761669, 0.0011664694175124168, 0.03365480154752731, 0.10097742080688477, 0.021067792549729347, 0.42791858315467834, 0.11242418736219406, 0.11434369534254074, 0.000791618600487709, 0.02291581965982914, 0.07201644033193588, 0.02081850729882717, 0.39859694242477417, 0.2763477563858032, 0.13874487578868866, 0.003258609212934971]], [[0.2643359303474426, 0.2943609654903412, 0.10517127066850662, 0.013473477214574814, 0.17808614671230316, 0.05031028389930725, 0.0477585569024086, 0.13444076478481293, 0.0626431554555893, 0.05089121311903, 0.025438696146011353, 0.12666909396648407, 0.015911895781755447, 0.08822031319141388, 0.09637932479381561, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.02893858775496483, 0.3286381959915161, 0.024464154615998268, 0.015645690262317657, 0.07065004110336304, 0.03320073336362839, 0.0035833900328725576, 0.002133443485945463, 0.0077736834064126015, 0.0014096481027081609, 0.006704544182866812, 0.0034484381321817636, 0.010553284548223019, 0.029550330713391304, 0.0064092278480529785, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0403970405459404, 0.029290249571204185, 0.2564694881439209, 0.03103366494178772, 0.01930038072168827, 0.0007984130643308163, 0.0024861868005245924, 0.013074777089059353, 0.025626862421631813, 0.0022637112997472286, 0.010511897504329681, 0.03038576804101467, 0.00803295336663723, 0.000980974524281919, 0.040744345635175705, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.23322375118732452, 0.23003342747688293, 0.24563531577587128, 0.07496963441371918, 0.029645830392837524, 0.0015733843902125955, 0.048427432775497437, 0.07474764436483383, 0.005064227152615786, 0.006064139772206545, 0.00639030896127224, 0.0023683567997068167, 0.0201968252658844, 0.0057837339118123055, 0.030518243089318275, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.009382463060319424, 0.004108777269721031, 0.355550616979599, 0.0026344929356127977, 0.036474164575338364, 0.0013674235669896007, 0.010420771315693855, 0.008167937397956848, 0.005904712714254856, 0.0164882093667984, 0.0014915319625288248, 0.00666471105068922, 0.007061991840600967, 0.006146776955574751, 0.03842667490243912, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.340854674577713, 0.027831802144646645, 0.11495380103588104, 0.4507772624492645, 0.33573275804519653, 0.07158998399972916, 0.3054116368293762, 0.09558256715536118, 0.008191889151930809, 0.08007357269525528, 0.08199689537286758, 0.011630101129412651, 0.016172919422388077, 0.020448284223675728, 0.05253906920552254, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0825798362493515, 0.09406770020723343, 0.044158000499010086, 0.06245531886816025, 0.15669509768486023, 0.1018981784582138, 0.17849969863891602, 0.1823071539402008, 0.1725231111049652, 0.14688736200332642, 0.027769910171628, 0.1729786992073059, 0.04907526820898056, 0.09640378504991531, 0.07928813993930817, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04138464853167534, 0.0045098732225596905, 0.098704032599926, 0.034942083060741425, 0.1842936873435974, 0.1567782759666443, 0.14141200482845306, 0.1953822374343872, 0.09936889261007309, 0.281032919883728, 0.13522183895111084, 0.012650868855416775, 0.02501768246293068, 0.2133605033159256, 0.14542686939239502, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05831298604607582, 0.07845572382211685, 0.00935202743858099, 0.09348727762699127, 0.2554629147052765, 0.026818757876753807, 0.15820558369159698, 0.09712891280651093, 0.18406683206558228, 0.297629177570343, 0.011888068169355392, 0.04674078896641731, 0.01729435659945011, 0.04945852607488632, 0.08047669380903244, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.030211733654141426, 0.004252443555742502, 0.044400423765182495, 0.0032993308268487453, 0.029341043904423714, 0.14371474087238312, 0.17894455790519714, 0.12369092553853989, 0.48359414935112, 0.06321088969707489, 0.05475561320781708, 0.3139732778072357, 0.086760014295578, 0.13208359479904175, 0.2905256450176239, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.06285266578197479, 0.0062216646037995815, 0.016913438215851784, 0.007285475265234709, 0.01629750058054924, 0.004617355298250914, 0.06147269159555435, 0.21831700205802917, 0.11657348275184631, 0.39258062839508057, 0.17390909790992737, 0.3519352376461029, 0.014494672417640686, 0.04437657818198204, 0.04845427721738815, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.014810703694820404, 0.027867808938026428, 0.00787208043038845, 0.003661711234599352, 0.06816401332616806, 0.014048570767045021, 0.04280591011047363, 0.04519394412636757, 0.07874996215105057, 0.2074531614780426, 0.12078044563531876, 0.53052818775177, 0.035032909363508224, 0.1398327797651291, 0.02986292913556099, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.011430865153670311, 0.002694258699193597, 0.03896895423531532, 0.04504057392477989, 0.00808126013725996, 0.01048098411411047, 0.012571780942380428, 0.0054772221483290195, 0.07419075071811676, 0.02193005569279194, 0.3994891941547394, 0.15694338083267212, 0.3065741956233978, 0.022703034803271294, 0.07852455973625183, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0007813395350240171, 4.470362910069525e-06, 0.0010683261789381504, 0.022204171866178513, 0.0022952572908252478, 4.198186070425436e-05, 0.0009061718010343611, 0.0006557627930305898, 0.0009219115017913282, 0.0006920882733538747, 0.005404994357377291, 0.012070748023688793, 0.21383939683437347, 0.0026518681552261114, 0.0011399114737287164, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03732156753540039, 0.14082211256027222, 0.08218222856521606, 0.02148711122572422, 0.037640467286109924, 0.011636778712272644, 0.01611051708459854, 0.06724098324775696, 0.20042963325977325, 0.035641491413116455, 0.045655738562345505, 0.041121501475572586, 0.23917138576507568, 0.01630677469074726, 0.2854580283164978, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.004390498157590628, 0.00876205787062645, 0.016465701162815094, 0.005714573431760073, 0.036494653671979904, 0.0032131776679307222, 0.01477664802223444, 0.018077310174703598, 0.010320773348212242, 0.006645719520747662, 0.03231831267476082, 0.004141036421060562, 0.011432528495788574, 0.011813640594482422, 0.20326180756092072, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.024762088432908058, 0.05259820073843002, 0.06384432315826416, 0.1483391523361206, 0.26820069551467896, 0.20398226380348206, 0.37573596835136414, 0.08007726073265076, 0.052950888872146606, 0.09653404355049133, 0.1610451638698578, 0.12953783571720123, 0.2330068051815033, 0.4463363587856293, 0.19394421577453613, 0.026641450822353363, 0.17128966748714447, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.679330587387085, 0.043791741132736206, 0.12768849730491638, 0.27546241879463196, 0.03847555071115494, 0.08167082816362381, 0.21957245469093323, 0.04802798852324486, 0.10780715942382812, 0.6106712222099304, 0.2505488693714142, 0.1709391176700592, 0.04529926925897598, 0.17936259508132935, 0.13903558254241943, 0.5577486157417297, 0.24638143181800842, 0.025497647002339363, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05959116667509079, 0.03547457605600357, 0.03805014118552208, 0.02909783646464348, 0.08531224727630615, 0.035567909479141235, 0.017052877694368362, 0.03032829985022545, 0.012725351378321648, 0.06508343666791916, 0.04963213950395584, 0.013415418565273285, 0.026129938662052155, 0.011819864623248577, 0.21026377379894257, 0.1241803988814354, 0.06599891930818558, 0.13004763424396515, 0.33318501710891724, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0922531858086586, 0.009465531446039677, 0.05285167694091797, 0.11621613800525665, 0.008946871384978294, 0.0003396931570023298, 0.056973982602357864, 0.011571673676371574, 0.03833528608083725, 0.02977353148162365, 0.12428728491067886, 0.005304301157593727, 0.012764646671712399, 0.03717968612909317, 0.1998610943555832, 0.9552784562110901, 0.6656578779220581, 0.04364815354347229, 0.097982257604599, 0.0012550450628623366, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.024207258597016335, 0.015275360085070133, 0.12442810088396072, 0.044900182634592056, 0.06243159621953964, 0.002727220067754388, 0.05297050252556801, 0.34427115321159363, 0.10989916324615479, 0.020859790965914726, 0.11048608273267746, 0.02605186030268669, 0.1171213760972023, 0.05136575922369957, 0.16462838649749756, 0.6779462695121765, 0.5809971690177917, 0.2087380737066269, 0.15752893686294556, 0.08772724121809006, 0.09023962169885635, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03260662034153938, 0.00298042013309896, 0.16533112525939941, 0.056620776653289795, 0.049906134605407715, 0.008958332240581512, 0.05700542405247688, 0.016634995117783546, 0.029206881299614906, 0.025224529206752777, 0.19688823819160461, 0.03853357210755348, 0.07708126306533813, 0.04636078327894211, 0.17741571366786957, 0.6994673609733582, 0.48720496892929077, 0.08263873308897018, 0.3298986256122589, 0.0049313209019601345, 0.07016509026288986, 0.5443912744522095, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04517968371510506, 0.08089613169431686, 0.11787059158086777, 0.09224344044923782, 0.27191361784935, 0.020393863320350647, 0.01454318780452013, 0.009129227139055729, 0.020442765206098557, 0.08070629835128784, 0.07541637122631073, 0.10045406222343445, 0.04119513928890228, 0.10953037440776825, 0.15667563676834106, 0.3437848389148712, 0.28689879179000854, 0.5712999105453491, 0.5371078252792358, 0.06584293395280838, 0.2492358684539795, 0.014812931418418884, 0.02226697839796543, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08136362582445145, 0.07834970951080322, 0.015254710800945759, 0.0832342654466629, 0.10864067077636719, 0.11524737626314163, 0.1366880238056183, 0.012557982467114925, 0.1251911222934723, 0.15952906012535095, 0.026927798986434937, 0.07786250859498978, 0.11803606152534485, 0.2014097422361374, 0.2085045427083969, 0.44942334294319153, 0.3777551054954529, 0.7612449526786804, 0.7021526098251343, 0.30080679059028625, 0.4424319267272949, 0.22922295331954956, 0.04627525433897972, 0.055941756814718246, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07754338532686234, 0.11610410362482071, 0.032187070697546005, 0.05519983917474747, 0.0022462301421910524, 0.11507689952850342, 0.2733137607574463, 0.17666463553905487, 0.010644900612533092, 0.08315187692642212, 0.02269633859395981, 0.06840697675943375, 0.010724963620305061, 0.0371541827917099, 0.21114735305309296, 0.47138965129852295, 0.18856076896190643, 0.6503154039382935, 0.9041082859039307, 0.2803841233253479, 0.4006999135017395, 0.5757170915603638, 0.295682817697525, 0.04142303764820099, 0.006079117301851511, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.022315502166748047, 0.012378118932247162, 0.0062178960070014, 0.0078407758846879, 0.015144318342208862, 0.010697844438254833, 0.011326298117637634, 0.013119788840413094, 0.009139686822891235, 0.006104558240622282, 0.005014281254261732, 0.002417754614725709, 0.007784656248986721, 0.009948876686394215, 0.16676713526248932, 0.24097655713558197, 0.15950126945972443, 0.6649572849273682, 0.6751598119735718, 0.46790093183517456, 0.6438081860542297, 0.3765251934528351, 0.2975021302700043, 0.10267924517393112, 0.060453154146671295, 0.03869982063770294, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2628116309642792, 0.1443735957145691, 0.08422664552927017, 0.11404431611299515, 0.17927099764347076, 0.25378888845443726, 0.1460212618112564, 0.04387032985687256, 0.023589681833982468, 0.13644081354141235, 0.045464351773262024, 0.06847606599330902, 0.006222521886229515, 0.036451175808906555, 0.20291540026664734, 0.39086097478866577, 0.6666929125785828, 0.5642580389976501, 0.557075023651123, 0.25761184096336365, 0.3620971143245697, 0.656988263130188, 0.301082581281662, 0.3758563995361328, 0.026163028553128242, 0.024990877136588097, 0.0074356794357299805, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.22663825750350952, 0.15363532304763794, 0.01756531558930874, 0.025186356157064438, 0.038983430713415146, 0.01259024627506733, 0.15960636734962463, 0.10260611027479172, 0.059462085366249084, 0.02338782697916031, 0.039677273482084274, 0.055942799896001816, 0.010165784507989883, 0.013570738956332207, 0.1720115691423416, 0.7909376621246338, 0.3817039430141449, 0.6133569478988647, 0.41290101408958435, 0.30558884143829346, 0.6049348711967468, 0.5688384175300598, 0.4680134057998657, 0.6550416946411133, 0.42371857166290283, 0.10508850961923599, 0.021316751837730408, 0.05294431000947952, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04994741827249527, 0.08986728638410568, 0.03736276924610138, 0.029899757355451584, 0.03542618826031685, 0.007244490087032318, 0.040187276899814606, 0.040814109146595, 0.04076588898897171, 0.05965813249349594, 0.045340292155742645, 0.0002602309104986489, 0.026138437911868095, 0.02984587848186493, 0.21049101650714874, 0.17973686754703522, 0.17233335971832275, 0.334688276052475, 0.4481850564479828, 0.04172942414879799, 0.10337609797716141, 0.5107487440109253, 0.7207926511764526, 0.1405051052570343, 0.0654703825712204, 0.41273486614227295, 0.17914383113384247, 0.042542651295661926, 0.010745447129011154, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.058702513575553894, 0.04533839225769043, 0.03167680650949478, 0.07689032703638077, 0.07722999900579453, 0.05968516319990158, 0.08647314459085464, 0.04232413321733475, 0.05769982933998108, 0.08562258630990982, 0.07418374717235565, 0.08922348916530609, 0.0013435373548418283, 0.0365031398832798, 0.1955317258834839, 0.5207539200782776, 0.308788537979126, 0.08189663290977478, 0.5850351452827454, 0.3457651734352112, 0.15844188630580902, 0.2948668897151947, 0.4065589904785156, 0.12084604799747467, 0.29343682527542114, 0.49164822697639465, 0.07233413308858871, 0.0535273477435112, 0.014947501011192799, 0.008541097864508629, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.035160183906555176, 0.01820351555943489, 0.1303882896900177, 0.019772829487919807, 0.040328264236450195, 0.05493366718292236, 0.03643186390399933, 0.013673724606633186, 0.020261095836758614, 0.09265058487653732, 0.06087178364396095, 0.005874141119420528, 0.0010416797595098615, 0.00679743243381381, 0.17795756459236145, 0.2949400544166565, 0.03748409450054169, 0.14473117887973785, 0.0705113336443901, 0.013025683350861073, 0.005298166535794735, 0.21091029047966003, 0.014800299890339375, 0.2805088758468628, 0.000897476973477751, 0.0938984826207161, 0.004705057479441166, 0.04936474934220314, 0.011992034502327442, 0.18721424043178558, 0.00230285432189703, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0850016176700592, 0.12483492493629456, 0.30438917875289917, 0.08283902704715729, 0.36141735315322876, 0.5806636810302734, 0.21757252514362335, 0.0776025652885437, 0.2093839943408966, 0.1517311930656433, 0.0691467672586441, 0.05431315675377846, 0.323522686958313, 0.21248842775821686, 0.11186490952968597, 0.44276589155197144, 0.06478449702262878, 0.543609619140625, 0.8444110155105591, 0.13468694686889648, 0.4405028522014618, 0.6528593897819519, 0.5737791061401367, 0.6313535571098328, 0.8501816987991333, 0.4486657381057739, 0.06076665595173836, 0.7409859299659729, 0.15147589147090912, 0.20801351964473724, 0.027446726337075233, 0.036936238408088684, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.017619943246245384, 0.008017263375222683, 0.019503258168697357, 0.014857600443065166, 0.07692210376262665, 0.015309707261621952, 0.015313221141695976, 0.008549719117581844, 0.03095930442214012, 0.019377540796995163, 0.031960610300302505, 0.0054225618951022625, 0.016712497919797897, 0.015215321443974972, 0.15961019694805145, 0.5445577502250671, 0.2876933515071869, 0.7013069987297058, 0.627236008644104, 0.37061285972595215, 0.6206991076469421, 0.38252583146095276, 0.4230470061302185, 0.31842562556266785, 0.28603002429008484, 0.015331648290157318, 0.14692452549934387, 0.8622261881828308, 0.049388445913791656, 0.37183380126953125, 0.17907747626304626, 0.05781394988298416, 0.020684318616986275, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2695287764072418, 0.16650046408176422, 0.14075446128845215, 0.1364857405424118, 0.23432065546512604, 0.261515349149704, 0.18958930671215057, 0.053015366196632385, 0.031337250024080276, 0.28422990441322327, 0.08986067771911621, 0.06408891826868057, 0.008591849356889725, 0.031372129917144775, 0.19151051342487335, 0.4656296670436859, 0.6725881099700928, 0.6199259161949158, 0.6479836702346802, 0.24076998233795166, 0.34658652544021606, 0.5947279930114746, 0.37259459495544434, 0.5521662831306458, 0.14718003571033478, 0.19626900553703308, 0.024240192025899887, 0.27736979722976685, 0.05565635487437248, 0.3618892729282379, 0.44332295656204224, 0.027751203626394272, 0.0260067880153656, 0.010717106983065605, NaN, NaN, NaN, NaN, NaN, NaN], [0.2586316764354706, 0.21131351590156555, 0.019284198060631752, 0.02717362530529499, 0.037918541580438614, 0.014535612426698208, 0.14439015090465546, 0.14164134860038757, 0.06384728103876114, 0.03232301026582718, 0.05240772292017937, 0.08253412693738937, 0.007928711362183094, 0.011026060208678246, 0.1583670824766159, 0.830940842628479, 0.42077580094337463, 0.7156820893287659, 0.57599937915802, 0.5493759512901306, 0.7128159999847412, 0.5476810932159424, 0.527928352355957, 0.8053308725357056, 0.8646240234375, 0.542984127998352, 0.2950981855392456, 0.3170693516731262, 0.5610483884811401, 0.26465174555778503, 0.45835256576538086, 0.22733505070209503, 0.10187508910894394, 0.03538959100842476, 0.07069608569145203, NaN, NaN, NaN, NaN, NaN], [0.0646420493721962, 0.15151722729206085, 0.04734531044960022, 0.03642117232084274, 0.03833956643939018, 0.007805521599948406, 0.03985777497291565, 0.05410199984908104, 0.07749858498573303, 0.1281091719865799, 0.06692291796207428, 0.0004382343322504312, 0.02769407443702221, 0.03219819441437721, 0.20084568858146667, 0.09599269181489944, 0.08247342705726624, 0.25253206491470337, 0.4357891380786896, 0.039192523807287216, 0.0719948410987854, 0.3563676178455353, 0.5300538539886475, 0.06311739236116409, 0.037909455597400665, 0.5032193064689636, 0.39894816279411316, 0.3283153772354126, 0.21619060635566711, 0.017918655648827553, 0.2577371895313263, 0.14531975984573364, 0.346793532371521, 0.2014700472354889, 0.0539211668074131, 0.0146569162607193, NaN, NaN, NaN, NaN], [0.06935474276542664, 0.07278740406036377, 0.0317843034863472, 0.061563972383737564, 0.057788632810115814, 0.05731336027383804, 0.08327846229076385, 0.046548519283533096, 0.06359860301017761, 0.13075897097587585, 0.09122113883495331, 0.1188196912407875, 0.0009191188146360219, 0.03464866429567337, 0.18994329869747162, 0.6422337889671326, 0.3740711212158203, 0.10689651221036911, 0.6858291029930115, 0.4494076073169708, 0.2826421856880188, 0.3886936604976654, 0.475405216217041, 0.13226336240768433, 0.3073323965072632, 0.7139697670936584, 0.17356495559215546, 0.25040003657341003, 0.23144030570983887, 0.024455448612570763, 0.4280460476875305, 0.048713963478803635, 0.3974619209766388, 0.06130422651767731, 0.05969162657856941, 0.015271119773387909, 0.00685582309961319, NaN, NaN, NaN], [0.04588386043906212, 0.027941085398197174, 0.16196617484092712, 0.023955674842000008, 0.04093120992183685, 0.06800121814012527, 0.031365618109703064, 0.013349683955311775, 0.016157155856490135, 0.09367228299379349, 0.06382262706756592, 0.009268027730286121, 0.0006308736628852785, 0.005314440466463566, 0.17240527272224426, 0.5218734741210938, 0.03395698964595795, 0.2861349880695343, 0.13773199915885925, 0.02211177349090576, 0.014614011161029339, 0.43378758430480957, 0.02492188662290573, 0.26067787408828735, 0.0009113854030147195, 0.1411941796541214, 0.009023642167448997, 0.14982649683952332, 0.15959703922271729, 0.7153633832931519, 0.014257365837693214, 0.06102409213781357, 0.12158294767141342, 0.006897313520312309, 0.06130388379096985, 0.012951835058629513, 0.16874605417251587, 0.002189028775319457, NaN, NaN], [0.09685268998146057, 0.17937548458576202, 0.31954076886177063, 0.09235721081495285, 0.3550800085067749, 0.5939842462539673, 0.19687135517597198, 0.10603781044483185, 0.27224627137184143, 0.17071248590946198, 0.0712975338101387, 0.10525800287723541, 0.3080449402332306, 0.250378280878067, 0.11120767891407013, 0.45293620228767395, 0.05202305316925049, 0.4803192913532257, 0.8224762082099915, 0.10338833183050156, 0.2861584722995758, 0.8321961760520935, 0.7622299790382385, 0.5323314070701599, 0.8633370995521545, 0.5219312310218811, 0.07432084530591965, 0.7646023631095886, 0.4150907099246979, 0.4998815357685089, 0.606073796749115, 0.2854492664337158, 0.6639280319213867, 0.09482558071613312, 0.806840717792511, 0.19665148854255676, 0.18194931745529175, 0.01953776553273201, 0.037144362926483154, NaN], [0.012543261051177979, 0.010277148336172104, 0.014658409170806408, 0.007294217124581337, 0.028056686744093895, 0.009602113626897335, 0.004711315967142582, 0.003909323364496231, 0.019910220056772232, 0.0035717461723834276, 0.016398703679442406, 0.01044577918946743, 0.015165981836616993, 0.04322582483291626, 0.1563079059123993, 0.8357685804367065, 0.6023411154747009, 0.16389556229114532, 0.4697819948196411, 0.05014880374073982, 0.3185025751590729, 0.2618474066257477, 0.7044641375541687, 0.16675803065299988, 0.7323283553123474, 0.14429442584514618, 0.2621355652809143, 0.041847843676805496, 0.3185603618621826, 0.04513467848300934, 0.49906620383262634, 0.611339807510376, 0.21515053510665894, 0.3302164673805237, 0.04920952767133713, 0.2760073244571686, 0.0218669306486845, 0.25043201446533203, 0.13627314567565918, 0.01334126852452755]]], [[[0.00028402332100085914, 1.9304454923485537e-08, 1.5483598847509938e-09, 7.885660006923256e-12, 2.7246130684943637e-08, 2.9440096113830805e-05, 4.3406546978985716e-07, 3.7434634236888087e-07, 3.9264233464564313e-07, 1.911867819615054e-08, 6.894639170695882e-08, 1.9322192201798316e-06, 1.594805780769093e-06, 1.097217136702966e-06, 0.25163131952285767, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8221166729927063, 0.0031213052570819855, 7.842657214496285e-05, 5.977510153520882e-10, 6.043178735204435e-10, 7.336016096815001e-07, 0.0001510237343609333, 0.000765863514970988, 0.0003504687047097832, 5.704807790607447e-07, 3.8402351520971933e-08, 3.7901799032624695e-07, 1.534954208182171e-05, 4.934078606311232e-05, 0.00023439944197889417, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0023944040294736624, 0.796754002571106, 0.004422985017299652, 9.068900226338883e-07, 5.795331436964091e-10, 1.0343059742012883e-08, 4.4964113499190717e-07, 0.0014743957435712218, 0.00028717826353386045, 7.994436600711197e-05, 3.3569827451174206e-07, 1.215876466176269e-07, 7.940250839055807e-07, 4.835407253267476e-06, 2.585098854979151e-07, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [4.3931080995207594e-11, 0.0005229745293036103, 0.5791732668876648, 0.0002632129180710763, 3.316774765949049e-08, 1.7754019825469425e-12, 1.4596207272357664e-14, 1.5350217763554497e-09, 1.2882580335826788e-07, 7.457471838279162e-06, 1.2410231420290074e-06, 2.736720361440348e-08, 3.621486097116211e-11, 3.919724787804224e-12, 2.306477925317907e-12, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [3.994035801418473e-14, 1.3595737036187217e-10, 5.270875135465758e-06, 0.5513067841529846, 0.00020578903786372393, 1.9226330039145978e-07, 1.181193272532799e-12, 2.80986930771554e-13, 9.120337812881449e-14, 1.37843805814164e-10, 7.154308718781976e-07, 1.5133276747292257e-06, 7.425698944629744e-10, 2.2010659354171347e-13, 1.8997327582565005e-12, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [2.3444651168352815e-12, 2.1774425253313912e-13, 1.857566878094019e-09, 0.00030468025943264365, 0.9472002983093262, 0.00010681805724743754, 2.00606624645161e-08, 5.2167251502746245e-14, 1.354494091723496e-15, 5.737065011425513e-13, 8.729777456473187e-10, 3.2425006793346256e-05, 7.676636641917867e-07, 1.870739785303499e-09, 2.3914221713994266e-09, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [3.644098217625569e-11, 3.867062572937563e-11, 4.1057553190615437e-11, 1.5412249254609378e-09, 0.018834512680768967, 0.505605936050415, 0.0010763276368379593, 5.434728933551014e-08, 2.6194791127864825e-11, 6.074670846504876e-15, 3.814499497517554e-12, 1.2291486939375318e-07, 9.572526323609054e-06, 4.437842653715052e-05, 7.18067713023629e-06, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [5.002242687623948e-05, 2.445471238843311e-07, 7.217475506138271e-09, 2.943958878759423e-12, 1.391844648424012e-07, 0.0035048718564212322, 0.755942702293396, 0.0011242764303460717, 1.4866960555082187e-05, 9.753278740198823e-11, 3.792431321238132e-13, 1.6398679289486573e-11, 1.3850768709744443e-07, 0.0002873632765840739, 2.565975592005998e-05, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [7.748224284398475e-09, 3.667011867491965e-07, 1.7906526261768363e-09, 1.001209222569038e-16, 4.707358499311462e-15, 2.921879960204876e-10, 4.77575849799905e-06, 0.9355171918869019, 1.7088919776142575e-05, 1.5246609308405823e-08, 1.546373502880373e-14, 1.9256968477537417e-16, 2.8356877952137637e-15, 6.199032398512827e-10, 3.679770266273863e-09, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [6.04271771509346e-11, 2.349539499846287e-06, 6.254656170767703e-08, 2.0915530592191534e-12, 3.303753013789688e-16, 1.0466700578893717e-14, 7.288482968201282e-13, 0.0006303040427155793, 0.47335511445999146, 8.928982424549758e-05, 1.5872458902776998e-08, 1.3611594998645584e-14, 1.3777586457132233e-16, 1.589055302510104e-15, 8.100658338561217e-11, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [3.812023474658588e-10, 1.421315573679749e-06, 2.2867025109007955e-06, 2.6682736020688935e-08, 3.632111755455525e-12, 1.6831340872913367e-14, 3.240909670081289e-14, 1.4920277635610546e-07, 0.0005182845052331686, 0.39297640323638916, 0.0007259719423018396, 1.2580667174688642e-08, 3.7229049595736974e-13, 2.157145159519631e-15, 1.0612778433838344e-09, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [6.84109713322556e-10, 1.9775532322796607e-08, 5.041609938416514e-07, 0.00017906920402310789, 1.631619738873269e-06, 2.0158734681530177e-09, 9.65507530290054e-15, 4.2181228128435055e-12, 8.564649545128589e-10, 0.00023218656133394688, 0.6439363956451416, 0.000818322179839015, 1.3831699163802114e-07, 2.1358659198916774e-12, 5.4572883101400294e-08, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.4084274191361601e-08, 2.1930364191291574e-09, 7.004614666072939e-09, 2.0828078959311824e-06, 6.64705439703539e-05, 3.6118690331932157e-06, 4.0857584676645686e-11, 1.0090924406833124e-12, 5.430448080009356e-15, 6.815135122906213e-09, 0.0007384128402918577, 0.9033229351043701, 0.0037223652470856905, 5.428325380307797e-07, 5.097080588711833e-07, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [3.370899046006848e-11, 1.5044922772877722e-12, 1.903236411786996e-13, 5.2399131041103164e-12, 5.3600892613303586e-09, 3.287689196440624e-07, 1.293990137263279e-09, 3.2395277866498207e-13, 8.98320316581696e-19, 7.591717251043266e-18, 2.4333673097343134e-12, 7.08575316821225e-05, 0.3025490641593933, 0.00011370918218744919, 1.7842703314840946e-08, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0009491983219049871, 3.734114216058515e-05, 0.00010643315181368962, 4.299266220186837e-05, 0.0019948105327785015, 0.012520392425358295, 0.0005770812276750803, 0.00013455892622005194, 0.0002518744731787592, 0.0005399127840064466, 0.0017743584467098117, 0.004756112117320299, 0.00398082984611392, 0.002925803419202566, 0.1746407300233841, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.017177388072013855, 0.0003127168456558138, 0.004294774029403925, 0.0025685238651931286, 0.0020048224832862616, 0.0018501998856663704, 0.004262528382241726, 0.00010045748058473691, 0.004143967293202877, 0.0026836262550204992, 0.0008790316642262042, 0.0012905423063784838, 8.68891947902739e-05, 0.00021419797849375755, 0.16245633363723755, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12795236706733704, 0.00371668953448534, 0.02831968478858471, 0.025539351627230644, 0.0009935664711520076, 0.0005314573645591736, 0.0308157317340374, 4.653090945794247e-05, 0.004544692113995552, 0.02307700179517269, 0.014357739128172398, 0.0017676070565357804, 1.5830510164960288e-05, 0.0005655316635966301, 0.23366259038448334, 0.13569742441177368, 0.0376364141702652, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0012442924780771136, 0.6349257826805115, 1.560185046400875e-05, 0.0005892697954550385, 2.671209358595661e-06, 1.747990245348774e-05, 0.00010909549746429548, 9.000968930195086e-06, 1.720580803521443e-05, 0.0008049540338106453, 0.00025925427326001227, 4.468534825718962e-06, 5.9764097386505455e-06, 7.895294402260333e-05, 0.00020540088007692248, 0.05053132027387619, 0.5417848825454712, 0.07814626395702362, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.014811321161687374, 0.6550174951553345, 5.4754978918936104e-05, 0.0013682727003470063, 7.1730828494764864e-06, 3.513193587423302e-05, 0.00030579010490328074, 4.0161107790481765e-06, 8.621193410363048e-05, 0.0020331761334091425, 0.00018049145000986755, 1.5370842447737232e-05, 2.3058303213474574e-06, 3.803792060352862e-05, 0.0004018820764031261, 0.03762863576412201, 0.4749486744403839, 0.013701170682907104, 0.053301598876714706, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0038746336940675974, 0.000324725842801854, 0.0051879663951694965, 0.009153621271252632, 0.0008864403935149312, 0.6781038641929626, 0.057408660650253296, 0.0010902854846790433, 0.00043091498082503676, 0.000930881651584059, 0.00047575533972121775, 0.0024355631321668625, 0.0005705857765860856, 0.0003382607828825712, 0.0010924984235316515, 0.10598134994506836, 0.16776065528392792, 0.11929589509963989, 0.16846179962158203, 0.40715572237968445, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [3.359095899213571e-06, 1.5333833403019526e-07, 3.112653939751908e-05, 0.00013510043208952993, 6.284327810135437e-06, 0.7821753025054932, 0.0016732696676626801, 2.949555346276611e-05, 1.1825303545265342e-06, 2.2443591660703532e-06, 4.938602842230466e-07, 8.253279020209447e-07, 2.1931487026449759e-07, 9.422030302630446e-07, 3.409375494811684e-06, 0.05147748813033104, 0.203742116689682, 0.11462464928627014, 0.46246808767318726, 0.01836300455033779, 0.02458924613893032, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00014056767395231873, 5.100669682178705e-07, 0.0031089531257748604, 0.006296438630670309, 0.00044245802564546466, 0.5631491541862488, 0.006006886251270771, 0.00015836386592127383, 1.0129460861207917e-05, 9.741926623973995e-05, 8.02019567345269e-05, 2.8800504878745414e-05, 2.2740101485396735e-05, 9.966635116143152e-05, 5.9340749430703e-05, 0.17594558000564575, 0.17753779888153076, 0.024665912613272667, 0.19817322492599487, 0.008797828108072281, 0.022263213992118835, 0.29173722863197327, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07201159745454788, 9.12444302230142e-05, 0.07167930901050568, 0.07350550591945648, 0.008381813764572144, 0.32997292280197144, 0.32325229048728943, 0.006826527416706085, 0.005964158568531275, 0.01031426526606083, 0.0041834041476249695, 0.0003298712254036218, 2.8659975214395672e-05, 0.00019656911899801344, 0.02016262151300907, 0.016114797443151474, 0.0061007170006632805, 0.028504224494099617, 0.017245782539248466, 0.08753485232591629, 0.11264273524284363, 0.6154332160949707, 0.029144972562789917, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0011574724921956658, 3.413460092360765e-07, 0.00010100962390424684, 0.0058910842053592205, 3.088227913394803e-06, 0.01394782867282629, 0.16852441430091858, 0.6476468443870544, 4.158269439358264e-05, 0.002217742381617427, 3.1430703529622406e-05, 8.318846812471747e-05, 7.552150123046886e-07, 2.136993316526059e-06, 0.00013183141709305346, 0.027042992413043976, 0.032212790101766586, 0.019619816914200783, 0.014702342450618744, 0.06721275299787521, 0.2560867667198181, 0.5545244216918945, 0.40561506152153015, 0.037922732532024384, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.056869976222515106, 0.00018767332949209958, 0.07251239567995071, 0.21200358867645264, 0.5404223799705505, 0.01658189669251442, 0.03565289452672005, 0.0015120785683393478, 0.002293382305651903, 0.005935561377555132, 0.012055100873112679, 0.005193157121539116, 0.003556813346222043, 0.007320231292396784, 0.018532630056142807, 0.1654873937368393, 0.013622531667351723, 0.0656571239233017, 0.09179358184337616, 0.03440919890999794, 0.08533406257629395, 0.16269220411777496, 0.1151970624923706, 0.09265416115522385, 0.028269361704587936, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.37012216448783875, 0.0030506134498864412, 0.585090160369873, 0.3774729073047638, 0.6362679600715637, 0.12865976989269257, 0.340728759765625, 0.01963443122804165, 0.11373940855264664, 0.0405576266348362, 0.04042620584368706, 0.006893007550388575, 0.0011100739939138293, 0.004035779275000095, 0.12706774473190308, 0.2598540484905243, 0.010173649527132511, 0.004170349799096584, 0.003479698905721307, 0.0014636714477092028, 0.0011101020500063896, 0.001677120802924037, 0.034040722995996475, 0.0041177538223564625, 0.024958845227956772, 0.016315795481204987, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01695789396762848, 0.00023016006161924452, 0.013878279365599155, 0.04998883232474327, 0.0032932739704847336, 8.226843783631921e-05, 0.014781651087105274, 0.00017401285003870726, 0.4112556278705597, 0.007095593959093094, 0.01393651869148016, 0.000858593441080302, 0.0009966455399990082, 0.006141065154224634, 0.004614917561411858, 0.17492477595806122, 0.010013026185333729, 0.005800239276140928, 0.0069971769116818905, 0.0036480696871876717, 0.001016399241052568, 0.0060493675991892815, 0.0034581662621349096, 0.00659980857744813, 0.0047594537027180195, 0.3941299021244049, 0.2407994568347931, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.023780474439263344, 4.510316648520529e-05, 0.013797261752188206, 0.087004654109478, 0.0004407854867167771, 0.0013536562910303473, 0.04187630116939545, 0.0028901200275868177, 0.06213926523923874, 0.3483656048774719, 0.03705320879817009, 0.005524389911442995, 0.0004139445663895458, 0.0025706440210342407, 0.012163926847279072, 0.06559828668832779, 0.005602334160357714, 0.0005807551206089556, 0.0005322807701304555, 0.004617360420525074, 0.00354054500348866, 0.005599506665021181, 0.011434626765549183, 0.006905066315084696, 0.009602343663573265, 0.11027393490076065, 0.36931946873664856, 0.06368503719568253, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.017730457708239555, 8.937691018218175e-05, 0.00767871318385005, 0.02321789041161537, 0.00010702417785068974, 0.004407694097608328, 0.0538853257894516, 0.011079255491495132, 0.003184565110132098, 0.026336153969168663, 0.005110009107738733, 0.3480301797389984, 0.002053677337244153, 0.01653059385716915, 0.00945478305220604, 0.015983520075678825, 0.012168757617473602, 0.0015684146201238036, 0.0005484889261424541, 0.00233695306815207, 0.0038106110878288746, 0.005947766825556755, 0.04194773733615875, 0.014443459920585155, 0.06465759128332138, 0.14989611506462097, 0.5095774531364441, 0.1882752925157547, 0.02387852594256401, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00016590843733865768, 4.410037217894569e-05, 0.0031412369571626186, 0.0015988551313057542, 0.002399750053882599, 0.0004506838449742645, 0.001152031123638153, 0.00021803524577990174, 0.00054850586457178, 0.0001300607982557267, 0.001143390079960227, 0.0023531741462647915, 0.6484718322753906, 0.061944324523210526, 1.8855764210456982e-05, 0.11159919947385788, 0.06036144495010376, 0.06681493669748306, 0.0798669382929802, 0.03668922558426857, 0.018710536882281303, 0.029976846650242805, 0.0675768032670021, 0.03372039645910263, 0.057603828608989716, 0.14515243470668793, 0.25060775876045227, 0.23181115090847015, 0.14262832701206207, 0.33286023139953613, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [5.492825607689156e-07, 1.991102926979238e-08, 2.3713612335996004e-06, 1.7095164366764948e-05, 8.657893886265811e-07, 3.6805211323098774e-08, 1.598790731804911e-06, 2.0731313554733788e-07, 4.274500042811269e-07, 5.490248440764844e-06, 0.00014167907647788525, 5.53526615476585e-06, 0.5851997137069702, 0.22563536465168, 1.0684430407081891e-07, 0.018035059794783592, 0.02341379225254059, 0.0019442361081019044, 0.004369894042611122, 0.00136191223282367, 0.00017434914479963481, 0.0011034610215574503, 0.06787250190973282, 0.060198791325092316, 0.12004764378070831, 0.11878902465105057, 0.2063554972410202, 0.28332868218421936, 0.35319504141807556, 0.008158767595887184, 0.26057863235473633, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01633528247475624, 0.0006067559006623924, 0.047781698405742645, 0.1674666851758957, 0.0008243213524110615, 0.0007217283127829432, 0.005900595337152481, 0.0001012250068015419, 0.006910703144967556, 0.1343279927968979, 0.5695670247077942, 0.0034049933310598135, 0.008110514841973782, 0.0796104148030281, 0.00713667506352067, 0.17278411984443665, 0.007028562016785145, 0.010641193017363548, 0.013809186406433582, 0.0005732428980991244, 0.001056239241734147, 0.0005258666351437569, 0.03639528155326843, 0.02256075292825699, 0.01660884916782379, 0.1527748554944992, 0.1477358043193817, 0.2577149271965027, 0.03867224231362343, 0.04304511100053787, 0.11759469658136368, 0.0762997567653656, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02614973485469818, 0.001497315475717187, 0.11498566716909409, 0.08699594438076019, 0.006599655374884605, 0.0011878651566803455, 0.009639720432460308, 0.0002812722814269364, 0.014351817779242992, 0.06119270250201225, 0.19180962443351746, 0.06391202658414841, 0.4759237766265869, 0.44549837708473206, 0.058810409158468246, 0.38573285937309265, 0.0028330886270850897, 0.0014278099406510592, 0.0009824484586715698, 9.371336636831984e-05, 0.00015483389142900705, 6.760591350030154e-05, 0.0035791138652712107, 0.0002520910056773573, 0.0005180046427994967, 0.00024238335026893765, 0.011901103891432285, 0.011019378900527954, 0.006276060827076435, 0.0026990415062755346, 0.016820058226585388, 0.03330027312040329, 0.047877803444862366, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.041024841368198395, 0.0016396299470216036, 0.05072889104485512, 0.1323171705007553, 0.0024413676001131535, 0.00023246044293045998, 0.02059599943459034, 0.00033336327760480344, 0.7358176708221436, 0.04226389154791832, 0.0658484548330307, 0.002587914001196623, 0.013076293282210827, 0.0423613116145134, 0.051219869405031204, 0.21399648487567902, 0.008264300413429737, 0.0051351506263017654, 0.005111425183713436, 0.0020249083172529936, 0.00047485672985203564, 0.0018332998733967543, 0.0008904117858037353, 0.0017731828847900033, 0.000539442349690944, 0.03944296017289162, 0.039767228066921234, 0.00580678740516305, 0.004312179517000914, 0.003937484696507454, 0.00913114845752716, 0.006211036816239357, 0.3553882837295532, 0.3024981617927551, NaN, NaN, NaN, NaN, NaN, NaN], [0.025904469192028046, 0.00014531973283737898, 0.014812517911195755, 0.11958510428667068, 0.0003183217777404934, 0.0012536202557384968, 0.031174438074231148, 0.0025010022800415754, 0.045685503631830215, 0.4334242641925812, 0.057037968188524246, 0.005963113158941269, 0.0007164725102484226, 0.00356480129994452, 0.02565825544297695, 0.05261809378862381, 0.004144520964473486, 0.00047606538282707334, 0.0003396419051568955, 0.002880769083276391, 0.0015178520698100328, 0.0018901955336332321, 0.0029504895210266113, 0.0017174717504531145, 0.0006908842478878796, 0.0046035549603402615, 0.09042679518461227, 0.0032755613792687654, 0.007712012622505426, 0.032594844698905945, 0.02268057130277157, 0.033856723457574844, 0.07955116033554077, 0.4074561595916748, 0.07153668999671936, NaN, NaN, NaN, NaN, NaN], [0.04193783551454544, 0.0005606984486803412, 0.01569434627890587, 0.058890990912914276, 0.00016686622984707355, 0.0032934362534433603, 0.10695304721593857, 0.011062747798860073, 0.008127261884510517, 0.04922156408429146, 0.01035262644290924, 0.3408533036708832, 0.003045044606551528, 0.019185535609722137, 0.046415992081165314, 0.019381573423743248, 0.012705344706773758, 0.0019882190972566605, 0.0005741973291151226, 0.0020475401543080807, 0.0023934554774314165, 0.004172713495790958, 0.021013854071497917, 0.005879250820726156, 0.006729640066623688, 0.00632414361461997, 0.09735815972089767, 0.01909361220896244, 0.00100265524815768, 0.003452989971265197, 0.008203250356018543, 0.05971603840589523, 0.11904174834489822, 0.5188009142875671, 0.2541559338569641, 0.029506316408514977, NaN, NaN, NaN, NaN], [0.00012501348101068288, 4.870840712101199e-05, 0.0024386774748563766, 0.001847597537562251, 0.0017206922639161348, 0.0002501157287042588, 0.0009360458934679627, 0.00021343374100979418, 0.0004799730086233467, 0.00017777700850274414, 0.0013057318283244967, 0.0019216074142605066, 0.7016423344612122, 0.059743087738752365, 1.6802117897896096e-05, 0.10572486370801926, 0.04525948688387871, 0.055838145315647125, 0.050681136548519135, 0.027844024822115898, 0.014026278629899025, 0.025656970217823982, 0.0361209474503994, 0.017075760290026665, 0.01003955863416195, 0.016965145245194435, 0.04991300031542778, 0.01522271428257227, 0.007584442384541035, 0.03757705166935921, 0.03609456866979599, 0.10922907292842865, 0.19329114258289337, 0.2903786897659302, 0.29551932215690613, 0.1564989984035492, 0.3518115282058716, NaN, NaN, NaN], [1.7574552657606546e-06, 9.272354617451128e-08, 1.001089003693778e-05, 5.891482942388393e-05, 3.3656547202554066e-06, 1.2065736143540562e-07, 6.7727110035775695e-06, 6.411150366147922e-07, 1.3192883443480241e-06, 1.1707085832313169e-05, 0.00026830541901290417, 1.0283902156515978e-05, 0.6812964081764221, 0.27208930253982544, 4.838558993469633e-07, 0.017342884093523026, 0.024629754945635796, 0.0017386168474331498, 0.003977979999035597, 0.0011948446044698358, 0.0001711023651296273, 0.0019097719341516495, 0.050265345722436905, 0.048485398292541504, 0.025773482397198677, 0.011941587552428246, 0.02582539990544319, 0.014500979334115982, 0.011088544502854347, 0.0004536270862445235, 0.001346826204098761, 0.09912228584289551, 0.03899921476840973, 0.19399496912956238, 0.33165985345840454, 0.3351045250892639, 0.007158405613154173, 0.26822295784950256, NaN, NaN], [0.01900503970682621, 0.0008953948272392154, 0.09836827963590622, 0.2858547866344452, 0.0013939865166321397, 0.0011423979885876179, 0.011685764417052269, 0.00014273256238084286, 0.010754182003438473, 0.15914513170719147, 0.6438553333282471, 0.002441136632114649, 0.008362390100955963, 0.07132171094417572, 0.011131932027637959, 0.15815527737140656, 0.009173951111733913, 0.012453499250113964, 0.01756284572184086, 0.0007500716019421816, 0.0020462200045585632, 0.00166225153952837, 0.05335438624024391, 0.037105023860931396, 0.009711050428450108, 0.05516523867845535, 0.04893142729997635, 0.03887411952018738, 0.002221355913206935, 0.004346344619989395, 0.004376854281872511, 0.001785764587111771, 0.09844812005758286, 0.14674220979213715, 0.34636548161506653, 0.04763580113649368, 0.057022612541913986, 0.12166893482208252, 0.13556897640228271, NaN], [0.12417581677436829, 0.0153038389980793, 0.12986266613006592, 0.6406017541885376, 0.009386910125613213, 0.057520631700754166, 0.09723392128944397, 0.0041757188737392426, 0.030985616147518158, 0.12765046954154968, 0.052563395351171494, 0.09427980333566666, 0.010530965402722359, 0.01615813747048378, 0.110444575548172, 0.16895240545272827, 0.0006144722574390471, 0.0027162963524460793, 0.0007400937611237168, 0.0007253509247675538, 0.0007097159395925701, 0.000199983871425502, 0.0005034026107750833, 0.0002540702698752284, 0.0002154638059437275, 0.0004817947919946164, 0.0019994170870631933, 0.0003459753352217376, 6.575404404429719e-05, 0.004540599416941404, 0.00010029276745626703, 0.0005050064064562321, 0.003569946391507983, 0.008527955040335655, 0.003213587449863553, 0.0022120880894362926, 0.11142478138208389, 0.01313241571187973, 0.055687084794044495, 0.21235007047653198]], [[0.1577264666557312, 0.03251823037862778, 0.4939506947994232, 0.8334789872169495, 0.6927971243858337, 0.3147047460079193, 0.7604361176490784, 0.11822030693292618, 0.7022377848625183, 0.6516091823577881, 0.14691989123821259, 0.2232232689857483, 0.14339210093021393, 0.3761228322982788, 0.014605461619794369, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.028655482456088066, 0.14083503186702728, 0.08485368639230728, 0.8299343585968018, 0.8304422497749329, 0.5664599537849426, 0.834579586982727, 0.7438958287239075, 0.8452481031417847, 0.8614712953567505, 0.3640905022621155, 0.805733323097229, 0.3481642007827759, 0.795884370803833, 0.05269646272063255, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.02106422185897827, 0.010846637189388275, 0.073356993496418, 0.017661061137914658, 0.8741048574447632, 0.5687165856361389, 0.5249210000038147, 0.5693489909172058, 0.5103186368942261, 0.5253384709358215, 0.6472406387329102, 0.4561024308204651, 0.1524587720632553, 0.45141565799713135, 0.034538887441158295, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2203565090894699, 0.02154199220240116, 0.007279311306774616, 0.003464027540758252, 0.18461424112319946, 0.07773485034704208, 0.7297388315200806, 0.2260110229253769, 0.6848539113998413, 0.2328294813632965, 0.22646839916706085, 0.3173597455024719, 0.10388152301311493, 0.06158056855201721, 0.11330780386924744, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1574045568704605, 0.12516136467456818, 0.04707150533795357, 0.0032313871197402477, 0.19444315135478973, 0.046962298452854156, 0.48863229155540466, 0.8290899991989136, 0.892469584941864, 0.6836395859718323, 0.83636474609375, 0.47956424951553345, 0.034452617168426514, 0.38761135935783386, 0.055785421282052994, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4389230012893677, 0.6133158802986145, 0.4783843159675598, 0.11230780929327011, 0.006951127201318741, 0.0644199401140213, 0.03406795859336853, 0.33251792192459106, 0.9552598595619202, 0.8827710747718811, 0.9276224970817566, 0.8325800895690918, 0.737617552280426, 0.745059609413147, 0.05149900168180466, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3395847976207733, 0.09897124767303467, 0.16763220727443695, 0.1671983003616333, 0.049412358552217484, 0.007114487700164318, 0.3340696394443512, 0.018166696652770042, 0.7235669493675232, 0.9639523029327393, 0.851059079170227, 0.7306914925575256, 0.5801126956939697, 0.8017169237136841, 0.08099871873855591, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.44394704699516296, 0.6082286238670349, 0.37166181206703186, 0.3715074956417084, 0.35315781831741333, 0.10853563994169235, 0.013190319761633873, 0.07092351466417313, 0.03435605764389038, 0.25131845474243164, 0.921750545501709, 0.8745512366294861, 0.7473158240318298, 0.834020733833313, 0.1216883435845375, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.18251584470272064, 0.8759727478027344, 0.1439245641231537, 0.06640342622995377, 0.060579828917980194, 0.2710072100162506, 0.011089610867202282, 0.034396518021821976, 0.1700025051832199, 0.043876904994249344, 0.14450228214263916, 0.9449294805526733, 0.9689385294914246, 0.939329981803894, 0.07954179495573044, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.32071176171302795, 0.7452729344367981, 0.11999625712633133, 0.08053360879421234, 0.3748469650745392, 0.31863275170326233, 0.028054066002368927, 0.2197551280260086, 0.01771731488406658, 0.23943577706813812, 0.01906767673790455, 0.8113164901733398, 0.9739595055580139, 0.9691897630691528, 0.21732129156589508, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6261264085769653, 0.6649302244186401, 0.5194191336631775, 0.6324451565742493, 0.6771988272666931, 0.7814968228340149, 0.4118405878543854, 0.3728334903717041, 0.03296521306037903, 0.008678224869072437, 0.6047253012657166, 0.11251461505889893, 0.21560458838939667, 0.9244948625564575, 0.10127653181552887, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3176693320274353, 0.5172579884529114, 0.1793123036623001, 0.37762320041656494, 0.23678036034107208, 0.5621929168701172, 0.08773050457239151, 0.24525783956050873, 0.010828782804310322, 0.025829488411545753, 0.0057976157404482365, 0.08708162605762482, 0.04166324809193611, 0.5714256167411804, 0.16898052394390106, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6460146307945251, 0.8194199800491333, 0.48921409249305725, 0.6910595297813416, 0.5259124636650085, 0.6389046311378479, 0.3241840600967407, 0.7817367911338806, 0.17853572964668274, 0.1606196016073227, 0.06383053213357925, 0.007355134002864361, 0.02128707617521286, 0.02206379547715187, 0.23354344069957733, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5992116332054138, 0.6358246803283691, 0.47243836522102356, 0.5617506504058838, 0.6971379518508911, 0.6431114673614502, 0.39991113543510437, 0.8182389140129089, 0.2704472243785858, 0.20400457084178925, 0.059529319405555725, 0.06732083112001419, 0.008503233082592487, 0.06121496111154556, 0.2071741670370102, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2342938333749771, 0.5683650374412537, 0.6037701964378357, 0.7331977486610413, 0.7349027395248413, 0.6651985049247742, 0.23853524029254913, 0.2293619066476822, 0.48426058888435364, 0.7077944874763489, 0.5918195843696594, 0.8169012665748596, 0.7005065679550171, 0.4784330725669861, 0.015931207686662674, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05668458715081215, 0.013551714830100536, 0.3300224542617798, 0.22417771816253662, 0.24923239648342133, 0.16107039153575897, 0.07639153301715851, 0.036736860871315, 0.044193096458911896, 0.14611276984214783, 0.15061600506305695, 0.035221245139837265, 0.0397845022380352, 0.06225845590233803, 0.12414046376943588, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.29422780871391296, 0.3258638381958008, 0.027477310970425606, 0.10906420648097992, 0.003920723684132099, 0.020042676478624344, 0.05157224088907242, 0.0009247793932445347, 0.005282218102365732, 0.1744423359632492, 0.0761384516954422, 0.0033416510559618473, 0.0003361533163115382, 0.0012587645323947072, 0.013668928295373917, 0.13440807163715363, 0.048166193068027496, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19355924427509308, 0.1259031891822815, 0.004604514688253403, 0.04003702849149704, 0.0129036083817482, 0.019794460386037827, 0.06589072942733765, 0.0014933310449123383, 0.012753497809171677, 0.06252782791852951, 0.0361945815384388, 0.011655895970761776, 0.01012047752737999, 0.02639157697558403, 0.16549569368362427, 0.14904144406318665, 0.03273539990186691, 0.03615117073059082, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4293937385082245, 0.07181306928396225, 0.003158864099532366, 0.04697505012154579, 0.01354672759771347, 0.09221473336219788, 0.24058710038661957, 0.0037424738984555006, 0.07543525844812393, 0.0656844824552536, 0.01989266835153103, 0.06512395292520523, 0.01137665193527937, 0.029709961265325546, 0.18951866030693054, 0.17614386975765228, 0.0854690745472908, 0.038236960768699646, 0.12011754512786865, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.052543047815561295, 0.03695955500006676, 0.100065678358078, 0.07546547800302505, 0.053252771496772766, 0.11382242292165756, 0.28551623225212097, 0.14051520824432373, 0.12815484404563904, 0.15533913671970367, 0.11139650642871857, 0.09512985497713089, 0.017796501517295837, 0.04266834259033203, 0.1351824700832367, 0.14069411158561707, 0.1466522365808487, 0.07941046357154846, 0.06070372834801674, 0.045592159032821655, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.002040643012151122, 0.005490712355822325, 0.024769198149442673, 0.007002650294452906, 0.0020249236840754747, 0.03913044556975365, 0.01487613096833229, 0.09424738585948944, 0.010089649818837643, 0.05513475462794304, 0.0488949678838253, 0.007691625505685806, 0.002344577107578516, 0.012510538101196289, 0.20307941734790802, 0.15778480470180511, 0.11167039722204208, 0.20017755031585693, 0.10082826018333435, 0.013994856737554073, 0.07346371561288834, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04981796815991402, 0.13342007994651794, 0.4189896881580353, 0.06767702847719193, 0.007763800676912069, 0.11641503125429153, 0.029343493282794952, 0.11072052270174026, 0.06700066477060318, 0.1429358571767807, 0.3406253457069397, 0.00571059063076973, 0.0006326772854663432, 0.004126383922994137, 0.17491626739501953, 0.15305520594120026, 0.26692208647727966, 0.1222626119852066, 0.14178596436977386, 0.012799645774066448, 0.019025815650820732, 0.14782781898975372, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.008032058365643024, 0.009898788295686245, 0.0165096465498209, 0.015990890562534332, 0.001612947671674192, 0.07025154680013657, 0.1309722512960434, 0.45684561133384705, 0.020022952929139137, 0.014566164463758469, 0.01627122238278389, 0.001012062537483871, 0.003352430183440447, 0.006583840120583773, 0.0849505066871643, 0.050227321684360504, 0.49922510981559753, 0.2564227879047394, 0.37594476342201233, 0.05222875997424126, 0.019398091360926628, 0.07475102692842484, 0.13636687397956848, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.027854006737470627, 0.008844887837767601, 0.011581032536923885, 0.014227867126464844, 0.0022522227372974157, 0.6803511381149292, 0.24682462215423584, 0.11913055926561356, 0.0028406307101249695, 0.006190288811922073, 0.00574448611587286, 0.0012344244169071317, 0.010572707280516624, 0.00985674187541008, 0.11121391505002975, 0.1278427243232727, 0.4489462971687317, 0.09382158517837524, 0.09914611279964447, 0.11451858282089233, 0.14035384356975555, 0.0858180820941925, 0.1395546793937683, 0.05027398467063904, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11111988872289658, 0.0035893325693905354, 0.4007861316204071, 0.2033512443304062, 0.1986382007598877, 0.15137647092342377, 0.12109687924385071, 0.007575488183647394, 0.021906785666942596, 0.03087061457335949, 0.08533017337322235, 0.07086688280105591, 0.06729871034622192, 0.045789312571287155, 0.1673528403043747, 0.06907324492931366, 0.44302117824554443, 0.21607427299022675, 0.21861647069454193, 0.14559195935726166, 0.12854896485805511, 0.21420170366764069, 0.5056769251823425, 0.05036870762705803, 0.14160890877246857, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06468851119279861, 0.006587199401110411, 0.23617494106292725, 0.19800357520580292, 0.15495024621486664, 0.06172433868050575, 0.05180465057492256, 0.01833559013903141, 0.016546709463000298, 0.05746111273765564, 0.0824536681175232, 0.007550883572548628, 0.007943101227283478, 0.011712267994880676, 0.33849596977233887, 0.08832916617393494, 0.4917650520801544, 0.16961733996868134, 0.21240676939487457, 0.17275941371917725, 0.13381528854370117, 0.1763075888156891, 0.3443826735019684, 0.022638684138655663, 0.14659351110458374, 0.05034468695521355, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09414701163768768, 0.10295354574918747, 0.0844656303524971, 0.06548816710710526, 0.08529236167669296, 0.06227908656001091, 0.030192906036973, 0.010874724946916103, 0.025562399998307228, 0.005146168638020754, 0.014559037052094936, 0.013559900224208832, 0.06781303137540817, 0.05153109133243561, 0.33232951164245605, 0.10765255987644196, 0.1569133847951889, 0.14696621894836426, 0.12414205074310303, 0.1321374922990799, 0.32589367032051086, 0.09939466416835785, 0.15668180584907532, 0.035531532019376755, 0.18526552617549896, 0.100669264793396, 0.1766001582145691, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.314544141292572, 0.6832185983657837, 0.07794945687055588, 0.042061515152454376, 0.015504884533584118, 0.1916494369506836, 0.006379975005984306, 0.0006176759488880634, 0.0012508369982242584, 0.01929312013089657, 0.022219885140657425, 0.0019787217024713755, 0.01769268326461315, 0.008809820748865604, 0.08711312711238861, 0.0920143872499466, 0.03631591796875, 0.10338561236858368, 0.13865944743156433, 0.14365890622138977, 0.19164490699768066, 0.08302215486764908, 0.17053648829460144, 0.20418454706668854, 0.4243081212043762, 0.23730118572711945, 0.11353020370006561, 0.062482837587594986, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.027118999511003494, 0.07309459149837494, 0.04486501216888428, 0.012266037985682487, 0.024303032085299492, 0.030924739316105843, 0.021004648879170418, 0.003694491693750024, 0.01517508551478386, 0.025275954976677895, 0.0075909653678536415, 0.24021397531032562, 0.04135901853442192, 0.07603362947702408, 0.11061857640743256, 0.14247462153434753, 0.10275112092494965, 0.08782284706830978, 0.07633533328771591, 0.09427531808614731, 0.2382509559392929, 0.11237408220767975, 0.1274290829896927, 0.09234490990638733, 0.29983192682266235, 0.19681134819984436, 0.09119200706481934, 0.1394888311624527, 0.02876400761306286, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.025165440514683723, 0.019109023734927177, 0.008520743809640408, 0.015198510140180588, 0.007751345168799162, 0.005125374533236027, 0.008160223253071308, 0.0017721926560625434, 0.08641061931848526, 0.07765892893075943, 0.017936453223228455, 0.020675569772720337, 0.0024341135285794735, 0.023971976712346077, 0.16557703912258148, 0.14126147329807281, 0.06271495670080185, 0.09029032289981842, 0.10313913226127625, 0.08530516922473907, 0.05194256827235222, 0.09853952378034592, 0.05407971888780594, 0.10021005570888519, 0.14394013583660126, 0.19472479820251465, 0.17138735949993134, 0.055624835193157196, 0.022259291261434555, 0.010825252160429955, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.22320780158042908, 0.05348529666662216, 0.01734296977519989, 0.1172923669219017, 0.004340981598943472, 0.003372892737388611, 0.033841460943222046, 0.024162178859114647, 0.05216863751411438, 0.3090120553970337, 0.2295515090227127, 0.014075365848839283, 0.020010780543088913, 0.20773397386074066, 0.12411301583051682, 0.15579406917095184, 0.5571659207344055, 0.09220181405544281, 0.09424383193254471, 0.2893342971801758, 0.14449337124824524, 0.08881417661905289, 0.09621196240186691, 0.05768556892871857, 0.34467604756355286, 0.16894927620887756, 0.32070621848106384, 0.32385867834091187, 0.08616255223751068, 0.0030245021916925907, 0.011462957598268986, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1383964717388153, 0.05579448863863945, 0.1563209742307663, 0.09128513187170029, 0.039257608354091644, 0.009886945597827435, 0.006391164381057024, 0.0007081980584189296, 0.006523598916828632, 0.16335614025592804, 0.02935076504945755, 0.023180969059467316, 0.19186609983444214, 0.2336183488368988, 0.16814255714416504, 0.06543286889791489, 0.3303832709789276, 0.1981877088546753, 0.17906354367733002, 0.08578304201364517, 0.12075137346982956, 0.09918820112943649, 0.14948950707912445, 0.0696079283952713, 0.2870473861694336, 0.2037079930305481, 0.20505982637405396, 0.415317177772522, 0.18504147231578827, 0.05944397673010826, 0.03780561313033104, 0.06350213289260864, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1625337302684784, 0.007939358241856098, 0.11928629875183105, 0.1341797411441803, 0.005670298356562853, 0.0033473502844572067, 0.022544465959072113, 0.005534132476896048, 0.007299710530787706, 0.08667418360710144, 0.07403960824012756, 0.004230144899338484, 0.002401313977316022, 0.005503634922206402, 0.20701391994953156, 0.08806300163269043, 0.5073549151420593, 0.15216797590255737, 0.1779468059539795, 0.08599209040403366, 0.038353316485881805, 0.05095306783914566, 0.13815101981163025, 0.05531492829322815, 0.3680262565612793, 0.045964885503053665, 0.5803228616714478, 0.2365681380033493, 0.10053237527608871, 0.016326427459716797, 0.011199035681784153, 0.02849578857421875, 0.09785498678684235, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08204744011163712, 0.04882703348994255, 0.048393696546554565, 0.02867632359266281, 0.012730585411190987, 0.02805456519126892, 0.014470821246504784, 0.008571655489504337, 0.011637779884040356, 0.011116313748061657, 0.015620187856256962, 0.00444953003898263, 0.038398172706365585, 0.021771300584077835, 0.25556278228759766, 0.10047968477010727, 0.17735490202903748, 0.1303417980670929, 0.1233980730175972, 0.11124629527330399, 0.27208706736564636, 0.09057758748531342, 0.20949512720108032, 0.0595981664955616, 0.32820063829421997, 0.19304482638835907, 0.3008245825767517, 0.24370267987251282, 0.0977335274219513, 0.0604717954993248, 0.08826017379760742, 0.05976974964141846, 0.11658596247434616, 0.26095637679100037, NaN, NaN, NaN, NaN, NaN, NaN], [0.3818233609199524, 0.6690115928649902, 0.07648678869009018, 0.0345233753323555, 0.011518634855747223, 0.1436365395784378, 0.005264819134026766, 0.000502048700582236, 0.0017500953981652856, 0.03918173909187317, 0.04129163548350334, 0.0023984990548342466, 0.020183494314551353, 0.008427903987467289, 0.09516369551420212, 0.08956606686115265, 0.03296149522066116, 0.07127847522497177, 0.10275094956159592, 0.12852256000041962, 0.15250688791275024, 0.05763629823923111, 0.13953621685504913, 0.2147330343723297, 0.3297017514705658, 0.25630685687065125, 0.3529660999774933, 0.05266188457608223, 0.19866161048412323, 0.08034973591566086, 0.16050152480602264, 0.12120798975229263, 0.21796129643917084, 0.13665789365768433, 0.05867582932114601, NaN, NaN, NaN, NaN, NaN], [0.02332407608628273, 0.06938373297452927, 0.035716570913791656, 0.008126936852931976, 0.012537641450762749, 0.0137803228572011, 0.01513306051492691, 0.00204691500402987, 0.029820755124092102, 0.05474912002682686, 0.016170548275113106, 0.22342036664485931, 0.05026429146528244, 0.06863567978143692, 0.11948796361684799, 0.16931524872779846, 0.06866136193275452, 0.058377113193273544, 0.054153572767972946, 0.06997817754745483, 0.17294903099536896, 0.06504172086715698, 0.09800923615694046, 0.07601338624954224, 0.22323867678642273, 0.17471107840538025, 0.20914696156978607, 0.32561469078063965, 0.04201642796397209, 0.014874166809022427, 0.043757203966379166, 0.11901038885116577, 0.15924809873104095, 0.08216992020606995, 0.13305248320102692, 0.031323518604040146, NaN, NaN, NaN, NaN], [0.020166568458080292, 0.015762973576784134, 0.006330324336886406, 0.008625769056379795, 0.005781465210020542, 0.00451312493532896, 0.007413441780954599, 0.0018466140609234571, 0.14846709370613098, 0.1376892477273941, 0.02431248314678669, 0.03153817355632782, 0.0025850962847471237, 0.026987632736563683, 0.15984071791172028, 0.14597494900226593, 0.05063166096806526, 0.07245789468288422, 0.08537694066762924, 0.07253167033195496, 0.03945168852806091, 0.07488631457090378, 0.04114159941673279, 0.09447583556175232, 0.11984950304031372, 0.21245841681957245, 0.24130037426948547, 0.053050536662340164, 0.036372195929288864, 0.012788524851202965, 0.05413965508341789, 0.17548364400863647, 0.18113258481025696, 0.17045176029205322, 0.056165628135204315, 0.023532675579190254, 0.007599800359457731, NaN, NaN, NaN], [0.11904438585042953, 0.03637225553393364, 0.013324074447154999, 0.04586002975702286, 0.00359557312913239, 0.002297254279255867, 0.02453085221350193, 0.019205793738365173, 0.07615289092063904, 0.3510056436061859, 0.24748629331588745, 0.0179043747484684, 0.015299135819077492, 0.16336295008659363, 0.13914434611797333, 0.20880575478076935, 0.4742221236228943, 0.0684090405702591, 0.07499475032091141, 0.22897963225841522, 0.11411925405263901, 0.06380540132522583, 0.06602712720632553, 0.04886250197887421, 0.25098055601119995, 0.16695836186408997, 0.41882073879241943, 0.45364588499069214, 0.19780457019805908, 0.004864717833697796, 0.007611281704157591, 0.23698794841766357, 0.08390159159898758, 0.28844529390335083, 0.28151822090148926, 0.0680297240614891, 0.0018790157046169043, 0.008693840354681015, NaN, NaN], [0.0598345547914505, 0.028141267597675323, 0.11996681243181229, 0.04193190485239029, 0.03001757152378559, 0.006633914541453123, 0.005910022184252739, 0.0007469199481420219, 0.010509159415960312, 0.18832749128341675, 0.032145459204912186, 0.022126449272036552, 0.16793787479400635, 0.1917877346277237, 0.16885708272457123, 0.06649312376976013, 0.2272576093673706, 0.15548978745937347, 0.13675269484519958, 0.06747769564390182, 0.09888236224651337, 0.07679145783185959, 0.09811051189899445, 0.059132058173418045, 0.16564641892910004, 0.1534833461046219, 0.21299242973327637, 0.46317315101623535, 0.18783308565616608, 0.06707606464624405, 0.07066023349761963, 0.038238298147916794, 0.13390158116817474, 0.1738123893737793, 0.3894510865211487, 0.199345201253891, 0.05267143249511719, 0.03450411930680275, 0.0674150139093399, NaN], [0.30011340975761414, 0.029496116563677788, 0.21246175467967987, 0.11388618499040604, 0.019265230745077133, 0.011386800557374954, 0.02386542037129402, 0.0049255480989813805, 0.002113579073920846, 0.2235003262758255, 0.1410367637872696, 0.022971738129854202, 0.009332037530839443, 0.01034344732761383, 0.12311729788780212, 0.13068987429141998, 0.5177554488182068, 0.21822108328342438, 0.17411521077156067, 0.11371950805187225, 0.10282127559185028, 0.14754493534564972, 0.10529720038175583, 0.04059072583913803, 0.1422514021396637, 0.16688787937164307, 0.3468432128429413, 0.07328897714614868, 0.033892080187797546, 0.005811289418488741, 0.006848806049674749, 0.033459149301052094, 0.08608346432447433, 0.29348817467689514, 0.07146795839071274, 0.05563248693943024, 0.008248405531048775, 0.00942459236830473, 0.03898181766271591, 0.13983668386936188]], [[0.04383472725749016, 0.02773081697523594, 0.016415273770689964, 0.024880478158593178, 0.005487722344696522, 0.14834517240524292, 0.010061212815344334, 0.013310510665178299, 0.03559315577149391, 0.022788431495428085, 0.016539618372917175, 0.022621937096118927, 0.3853665292263031, 0.02895752713084221, 0.21785423159599304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.02212689444422722, 0.0360226184129715, 0.0007962794625200331, 0.005733562167733908, 0.0017349227564409375, 0.011109595187008381, 0.02015179581940174, 0.048344310373067856, 0.003794114338234067, 0.016348786652088165, 0.0018908409401774406, 0.010183308273553848, 0.04822028428316116, 0.011540568433701992, 0.21287554502487183, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.19621919095516205, 0.02568935602903366, 0.012553256005048752, 0.05958101898431778, 0.0049527534283697605, 0.009129918180406094, 0.035662900656461716, 0.006033026147633791, 0.01979534700512886, 0.016174430027604103, 0.025959551334381104, 0.017891131341457367, 0.21532145142555237, 0.010915487073361874, 0.2776879370212555, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.22681212425231934, 0.26364389061927795, 0.1368870735168457, 0.07472710311412811, 0.004966794513165951, 0.17209400236606598, 0.07595591247081757, 0.10330677032470703, 0.009879215620458126, 0.30214887857437134, 0.027453631162643433, 0.07928238064050674, 0.6068928837776184, 0.0009245484252460301, 0.41711828112602234, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03220081329345703, 0.07110226154327393, 0.19687172770500183, 0.32465922832489014, 0.06123804301023483, 0.009123058058321476, 0.008925903588533401, 0.001694322214461863, 0.009767607785761356, 0.012425252236425877, 0.021234901621937752, 0.006749649532139301, 0.022427640855312347, 0.00419656652957201, 0.11337225884199142, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1499132513999939, 0.1588381826877594, 0.006192722357809544, 0.06905046850442886, 0.021936854347586632, 0.04223879054188728, 0.01654554158449173, 0.012800824828445911, 0.001194898271933198, 0.011350413784384727, 0.0011690479004755616, 0.03650015965104103, 0.0330234132707119, 0.032408226281404495, 0.30060991644859314, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10197536647319794, 0.32784661650657654, 0.22266407310962677, 0.37194594740867615, 0.4840903878211975, 0.2562866806983948, 0.20682689547538757, 0.01685171388089657, 0.02662164717912674, 0.01744299754500389, 0.07043293118476868, 0.06053447723388672, 0.13449640572071075, 0.0437617152929306, 0.15905345976352692, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04155902937054634, 0.02725875750184059, 0.06621034443378448, 0.15740959346294403, 0.22226983308792114, 0.11737026274204254, 0.021176597103476524, 0.037896860390901566, 0.001983239781111479, 0.07737525552511215, 0.040612466633319855, 0.036445699632167816, 0.04206009954214096, 0.005294053349643946, 0.22695806622505188, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3731417655944824, 0.020610323175787926, 0.04687204957008362, 0.19942151010036469, 0.0219199787825346, 0.023319954052567482, 0.607546865940094, 0.0038317576982080936, 0.05746426433324814, 0.0039819530211389065, 0.0020286834333091974, 0.023514816537499428, 0.0007224131841212511, 0.0017132725333794951, 0.31377115845680237, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.007707278709858656, 0.04994801804423332, 0.0602150596678257, 0.1843070536851883, 0.023052150383591652, 0.00867108628153801, 0.0030793596524745226, 0.008175634779036045, 0.3707427382469177, 0.032583341002464294, 0.030614105984568596, 0.003414844162762165, 0.0027733321767300367, 0.00039667857345193624, 0.06665757298469543, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.06275568902492523, 0.15385569632053375, 0.07121506333351135, 0.04657430946826935, 0.08974524587392807, 0.017753345891833305, 0.09537442773580551, 0.08409535884857178, 0.4617481529712677, 0.05371565744280815, 0.051210206001996994, 0.014556556940078735, 0.0261379461735487, 0.0015151489060372114, 0.25993233919143677, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.037524934858083725, 0.08964382112026215, 0.11503562331199646, 0.2385229468345642, 0.14595970511436462, 0.01507873460650444, 0.07354842126369476, 0.014194677583873272, 0.01029899064451456, 0.3145633935928345, 0.08443433046340942, 0.02799280546605587, 0.006364578381180763, 0.0011598452692851424, 0.25597554445266724, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03498825803399086, 0.003427299438044429, 0.012860815972089767, 0.00960747804492712, 0.0073430403135716915, 0.002194140339270234, 0.020218953490257263, 0.04016692563891411, 0.0035721054300665855, 0.11439335346221924, 0.03179614990949631, 0.0055262502282857895, 0.08811097592115402, 0.0019241927657276392, 0.31578439474105835, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0003122057532891631, 0.0005657155998051167, 0.0003099576279055327, 0.018182117491960526, 8.608390635345131e-05, 0.00029685357003472745, 0.00030423246789723635, 0.0039575002156198025, 0.00041145391878671944, 0.0009832053910940886, 0.0007515411707572639, 0.006357411853969097, 0.3007054328918457, 0.00010537439811741933, 0.00161165336612612, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.052370160818099976, 0.019386928528547287, 0.0404941625893116, 0.16087706387043, 0.14014431834220886, 0.0561581589281559, 0.1907973736524582, 0.027806226164102554, 0.022970959544181824, 0.05846026912331581, 0.09902504831552505, 0.038958851248025894, 0.016928229480981827, 0.04114920645952225, 0.14461401104927063, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03517069295048714, 0.03549245744943619, 0.004381549544632435, 0.008797217160463333, 0.007323419209569693, 0.042320944368839264, 0.004849699325859547, 0.003679578425362706, 0.011580413207411766, 0.009367180056869984, 0.006541883572936058, 0.022973380982875824, 0.023761657997965813, 0.02892483025789261, 0.1581033319234848, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01528994832187891, 0.20408181846141815, 0.11101088672876358, 0.08111120015382767, 0.07986893504858017, 0.010126215405762196, 0.020366966724395752, 0.1417536586523056, 0.04787333309650421, 0.04340335354208946, 0.2409791648387909, 0.04442436248064041, 0.005909040104597807, 0.014603852294385433, 0.18931475281715393, 0.13037645816802979, 0.08109150826931, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21622280776500702, 0.09626477211713791, 0.10110790282487869, 0.31975099444389343, 0.2572377920150757, 0.630383312702179, 0.1336757242679596, 0.17725828289985657, 0.02378956414759159, 0.22253809869289398, 0.13939163088798523, 0.30914127826690674, 0.35968318581581116, 0.48164138197898865, 0.09301326423883438, 0.14859925210475922, 0.02925589494407177, 0.0505123995244503, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.168080672621727, 0.1516411453485489, 0.07150255143642426, 0.32225823402404785, 0.2490793913602829, 0.30686429142951965, 0.032337237149477005, 0.16698232293128967, 0.04405289515852928, 0.2310783565044403, 0.10561788827180862, 0.2769646644592285, 0.19830158352851868, 0.1653461754322052, 0.09653043746948242, 0.21387919783592224, 0.03206360712647438, 0.012896520085632801, 0.06630519032478333, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04038669914007187, 0.16624715924263, 0.3317047655582428, 0.3851986229419708, 0.42305275797843933, 0.008450526744127274, 0.09501849114894867, 0.24002836644649506, 0.4256587326526642, 0.15410973131656647, 0.19127053022384644, 0.04389801248908043, 0.030224177986383438, 0.05971052870154381, 0.11478950828313828, 0.15968731045722961, 0.046736959367990494, 0.014681101776659489, 0.01418250147253275, 0.011044399812817574, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04527302458882332, 0.15370813012123108, 0.46266382932662964, 0.06791326403617859, 0.6029869914054871, 0.018879592418670654, 0.07514301687479019, 0.07948564738035202, 0.6243545413017273, 0.11254889518022537, 0.24916931986808777, 0.08612842112779617, 0.07598677277565002, 0.13317255675792694, 0.04299912229180336, 0.22570300102233887, 0.051045093685388565, 0.020206425338983536, 0.021926334127783775, 0.008406145498156548, 0.0702541247010231, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03695433586835861, 0.028389452025294304, 0.2721908688545227, 0.07653216272592545, 0.6730886697769165, 0.004614274017512798, 0.004165990743786097, 0.01533985324203968, 0.28992146253585815, 0.028840038925409317, 0.055076081305742264, 0.024787841364741325, 0.0010191021719947457, 0.0022868094965815544, 0.030124979093670845, 0.28555917739868164, 0.03329295665025711, 0.036049578338861465, 0.038853298872709274, 0.007190736476331949, 0.006643606815487146, 0.08228380233049393, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.005083801224827766, 0.09139324724674225, 0.28116321563720703, 0.08195066452026367, 0.6340349316596985, 0.012272918596863747, 0.0005934475339017808, 0.010692326352000237, 0.1514793336391449, 0.016046250239014626, 0.04672969505190849, 0.014393122866749763, 0.002580928150564432, 0.007409923244267702, 0.12582267820835114, 0.2511760890483856, 0.07463249564170837, 0.04988643527030945, 0.0701586976647377, 0.028143733739852905, 0.007391677238047123, 0.02261284738779068, 0.0737045407295227, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00605103699490428, 0.11548061668872833, 0.2870264947414398, 0.061026521027088165, 0.8064441084861755, 0.2189176380634308, 0.020241523161530495, 0.07779920846223831, 0.08952271938323975, 0.0073190852999687195, 0.02372264862060547, 0.038144610822200775, 0.07446137070655823, 0.09413070231676102, 0.030171062797307968, 0.15217745304107666, 0.19177564978599548, 0.125013530254364, 0.1473270058631897, 0.20325084030628204, 0.10669662803411484, 0.07946557551622391, 0.027662983164191246, 0.09494684636592865, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08316895365715027, 0.6715664267539978, 0.04549514129757881, 0.17856287956237793, 0.018127189949154854, 0.38010329008102417, 0.16956135630607605, 0.5726994872093201, 0.1473512202501297, 0.13756032288074493, 0.044131502509117126, 0.03872460126876831, 0.13646697998046875, 0.07963203638792038, 0.10255669057369232, 0.13806378841400146, 0.2514709234237671, 0.17176732420921326, 0.21858137845993042, 0.17882317304611206, 0.16198168694972992, 0.20351995527744293, 0.07158615440130234, 0.0266498401761055, 0.23213928937911987, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0817432552576065, 0.2031053900718689, 0.02472570165991783, 0.02598942257463932, 0.05427335575222969, 0.43315476179122925, 0.06398319453001022, 0.14792829751968384, 0.18555517494678497, 0.020227503031492233, 0.03572608157992363, 0.008726409636437893, 0.33127138018608093, 0.0956021174788475, 0.032814960926771164, 0.17152094841003418, 0.15314172208309174, 0.15820659697055817, 0.19208288192749023, 0.19640566408634186, 0.061033159494400024, 0.12321671098470688, 0.07748300582170486, 0.07906179875135422, 0.032524362206459045, 0.08073069155216217, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.36652442812919617, 0.4977355897426605, 0.09286413341760635, 0.21385566890239716, 0.18058304488658905, 0.4562758207321167, 0.4738945960998535, 0.2067655473947525, 0.17124009132385254, 0.035114847123622894, 0.05785587430000305, 0.03289380669593811, 0.3892229497432709, 0.2459530532360077, 0.0885753259062767, 0.11935991793870926, 0.25889015197753906, 0.181893989443779, 0.2521744966506958, 0.2510518431663513, 0.1320696324110031, 0.17421388626098633, 0.10352174937725067, 0.13144756853580475, 0.06071629375219345, 0.07381404936313629, 0.11898738145828247, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3338637053966522, 0.241106316447258, 0.10183558613061905, 0.16975384950637817, 0.22215212881565094, 0.1208982765674591, 0.12069278955459595, 0.027770178392529488, 0.12589573860168457, 0.018161755055189133, 0.05639319866895676, 0.024462532252073288, 0.08646970242261887, 0.18506868183612823, 0.2994369864463806, 0.11384479701519012, 0.12307179719209671, 0.17695116996765137, 0.21105043590068817, 0.2652710974216461, 0.1994313895702362, 0.5530626177787781, 0.33474239706993103, 0.11353342235088348, 0.20157715678215027, 0.12058570981025696, 0.02405776083469391, 0.20302970707416534, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24999171495437622, 0.7484717965126038, 0.1908620148897171, 0.6611655354499817, 0.24442408978939056, 0.0825357735157013, 0.5622089505195618, 0.4391622543334961, 0.045715928077697754, 0.2250336855649948, 0.3067566156387329, 0.014471310190856457, 0.06388252228498459, 0.21674634516239166, 0.13583892583847046, 0.1661912202835083, 0.3088836967945099, 0.3049609959125519, 0.34614017605781555, 0.3287224769592285, 0.19484750926494598, 0.49978625774383545, 0.2471936047077179, 0.14924246072769165, 0.2264283001422882, 0.11719675362110138, 0.028577886521816254, 0.03125511854887009, 0.04683076590299606, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05097173899412155, 0.16686855256557465, 0.15120531618595123, 0.3698476254940033, 0.35846272110939026, 0.6895467042922974, 0.8159933686256409, 0.843620777130127, 0.6904561519622803, 0.307870090007782, 0.450530469417572, 0.6275950074195862, 0.15986312925815582, 0.5293903350830078, 0.07888244837522507, 0.1382068395614624, 0.14312644302845, 0.15027517080307007, 0.2806132137775421, 0.10704077035188675, 0.15715429186820984, 0.3545873463153839, 0.2772214114665985, 0.11900671571493149, 0.16433128714561462, 0.08395379036664963, 0.0337035246193409, 0.08286106586456299, 0.029390821233391762, 0.07092607021331787, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3532100319862366, 0.1141892597079277, 0.06207668036222458, 0.23437273502349854, 0.13035829365253448, 0.16457295417785645, 0.6610441207885742, 0.6354422569274902, 0.6703211069107056, 0.18266227841377258, 0.16635818779468536, 0.1048990935087204, 0.1468038111925125, 0.17976891994476318, 0.0709633082151413, 0.31265145540237427, 0.17018769681453705, 0.42172688245773315, 0.3373875319957733, 0.26503118872642517, 0.3668123483657837, 0.6080453991889954, 0.3421963155269623, 0.29850897192955017, 0.22005639970302582, 0.08626232296228409, 0.05660916119813919, 0.04967416450381279, 0.020023291930556297, 0.01626538299024105, 0.03365384787321091, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18437133729457855, 0.20806346833705902, 0.06752406805753708, 0.15831130743026733, 0.3405534625053406, 0.0627271831035614, 0.3717433214187622, 0.3913803696632385, 0.5862330794334412, 0.29396724700927734, 0.02299528755247593, 0.060014016926288605, 0.08232607692480087, 0.15418194234371185, 0.15275102853775024, 0.11847452819347382, 0.5065410137176514, 0.4161456227302551, 0.44356557726860046, 0.358999639749527, 0.34202155470848083, 0.6410406231880188, 0.5693260431289673, 0.3344528377056122, 0.3382241725921631, 0.16963228583335876, 0.12081613391637802, 0.09492655098438263, 0.06781262904405594, 0.059771545231342316, 0.013083304278552532, 0.15846344828605652, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07671413570642471, 0.17070698738098145, 0.13325846195220947, 0.07402658462524414, 0.6503690481185913, 0.1330946981906891, 0.165133535861969, 0.2397843301296234, 0.6370089054107666, 0.09848601371049881, 0.09929761290550232, 0.10903115570545197, 0.14141131937503815, 0.14783106744289398, 0.08112896233797073, 0.14143924415111542, 0.33810776472091675, 0.4273369610309601, 0.4442084729671478, 0.4867575168609619, 0.40271657705307007, 0.7919159531593323, 0.5796146988868713, 0.41502290964126587, 0.19611117243766785, 0.2659074366092682, 0.0590454526245594, 0.09533000737428665, 0.06579555571079254, 0.049002423882484436, 0.011413656175136566, 0.05989237129688263, 0.0694013461470604, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1416744738817215, 0.274202436208725, 0.13295260071754456, 0.20105819404125214, 0.3945937156677246, 0.333781898021698, 0.3556738793849945, 0.2839928865432739, 0.10343024134635925, 0.07706140726804733, 0.054361648857593536, 0.05752982571721077, 0.2817353904247284, 0.27278265357017517, 0.13429909944534302, 0.06363721936941147, 0.3402014374732971, 0.30108359456062317, 0.3598821461200714, 0.356340229511261, 0.2955020070075989, 0.3913557827472687, 0.34592464566230774, 0.3881937265396118, 0.23078370094299316, 0.49122318625450134, 0.3432621657848358, 0.1563359946012497, 0.12668228149414062, 0.1534397453069687, 0.06296171993017197, 0.07472987473011017, 0.07419107109308243, 0.08810260146856308, NaN, NaN, NaN, NaN, NaN, NaN], [0.22879131138324738, 0.1777554452419281, 0.09183042496442795, 0.14726729691028595, 0.1873711347579956, 0.05672184377908707, 0.08326486498117447, 0.01781904511153698, 0.0835406556725502, 0.02614605240523815, 0.06876543164253235, 0.03439611196517944, 0.0621294341981411, 0.16512615978717804, 0.26481878757476807, 0.06025628373026848, 0.1445734202861786, 0.2208743691444397, 0.22917300462722778, 0.34805941581726074, 0.30598515272140503, 0.6932811141014099, 0.6030279994010925, 0.2491629421710968, 0.46458470821380615, 0.5228609442710876, 0.2136632800102234, 0.610046923160553, 0.25265923142433167, 0.14038830995559692, 0.07342293113470078, 0.22653138637542725, 0.10003089159727097, 0.02225746400654316, 0.14559555053710938, NaN, NaN, NaN, NaN, NaN], [0.1532706916332245, 0.5982866883277893, 0.18050755560398102, 0.5800401568412781, 0.22030943632125854, 0.025230426341295242, 0.3744361996650696, 0.265155166387558, 0.03173244372010231, 0.2068646252155304, 0.27338433265686035, 0.012270096689462662, 0.05047086998820305, 0.14277896285057068, 0.15170519053936005, 0.0902293398976326, 0.5066702961921692, 0.45472872257232666, 0.45485398173332214, 0.5058757662773132, 0.3594079613685608, 0.7028806209564209, 0.5180745720863342, 0.25713953375816345, 0.5372852683067322, 0.6213670372962952, 0.2659974694252014, 0.3181111812591553, 0.5259383916854858, 0.33730512857437134, 0.13441412150859833, 0.36266574263572693, 0.10496268421411514, 0.02362431399524212, 0.020191077142953873, 0.04590708762407303, NaN, NaN, NaN, NaN], [0.04688200727105141, 0.12437571585178375, 0.1870293915271759, 0.4533093273639679, 0.3565751910209656, 0.5648568868637085, 0.7852934002876282, 0.7657470703125, 0.5417794585227966, 0.4419334828853607, 0.632922887802124, 0.7103447914123535, 0.15686877071857452, 0.6169639825820923, 0.08483293652534485, 0.1059701219201088, 0.2303982675075531, 0.21762119233608246, 0.3580361306667328, 0.17096057534217834, 0.24843183159828186, 0.5131583213806152, 0.47260501980781555, 0.21650557219982147, 0.38561707735061646, 0.416827529668808, 0.1716565638780594, 0.3172723054885864, 0.29216328263282776, 0.47280052304267883, 0.38235870003700256, 0.1798420399427414, 0.1762932986021042, 0.04000748321413994, 0.08066289126873016, 0.03975420445203781, 0.08505715429782867, NaN, NaN, NaN], [0.2884610891342163, 0.10604135692119598, 0.07176870107650757, 0.2240629643201828, 0.12294583767652512, 0.10159854590892792, 0.6051279902458191, 0.5541971921920776, 0.5623130798339844, 0.16405576467514038, 0.18055777251720428, 0.13399486243724823, 0.12637703120708466, 0.18360036611557007, 0.09598042815923691, 0.2317487895488739, 0.2560827136039734, 0.5102789998054504, 0.4199059009552002, 0.44283756613731384, 0.5258800983428955, 0.732390284538269, 0.4491574466228485, 0.4244932234287262, 0.5298821926116943, 0.43037980794906616, 0.2800268232822418, 0.3093121647834778, 0.4250229299068451, 0.19317308068275452, 0.2640416920185089, 0.38813653588294983, 0.11181202530860901, 0.054203763604164124, 0.037284549325704575, 0.018739882856607437, 0.014264266937971115, 0.035236652940511703, NaN, NaN], [0.10626664012670517, 0.1478983461856842, 0.07806308567523956, 0.11814259737730026, 0.31690794229507446, 0.03372211009263992, 0.30042603611946106, 0.29277828335762024, 0.44479742646217346, 0.216581329703331, 0.023049354553222656, 0.0511498898267746, 0.08494822680950165, 0.14207273721694946, 0.16419102251529694, 0.08032029122114182, 0.6358892321586609, 0.5042787194252014, 0.5074477195739746, 0.5223307013511658, 0.5343775749206543, 0.703619122505188, 0.6657658815383911, 0.45647403597831726, 0.602655827999115, 0.5387927889823914, 0.39006462693214417, 0.39567169547080994, 0.43596506118774414, 0.41000646352767944, 0.269907683134079, 0.5412885546684265, 0.2038634866476059, 0.10306636989116669, 0.05501747503876686, 0.04515310004353523, 0.04695969074964523, 0.008877278305590153, 0.09985174983739853, NaN], [0.048457998782396317, 0.0638582855463028, 0.20956584811210632, 0.021124709397554398, 0.09014897048473358, 0.11662621796131134, 0.3483109474182129, 0.4503737986087799, 0.17136822640895844, 0.02997676283121109, 0.21708470582962036, 0.05856599286198616, 0.2859736979007721, 0.41663405299186707, 0.12262307107448578, 0.03129265457391739, 0.2636677324771881, 0.3672870099544525, 0.438161164522171, 0.7497870922088623, 0.43876102566719055, 0.6747432947158813, 0.5918557643890381, 0.5535795092582703, 0.7133825421333313, 0.7440239787101746, 0.3780657947063446, 0.4423457384109497, 0.6450315713882446, 0.5939705967903137, 0.7279283404350281, 0.4253756105899811, 0.4950290024280548, 0.13756991922855377, 0.08432447165250778, 0.11775307357311249, 0.12791647017002106, 0.07922011613845825, 0.04417572543025017, 0.3473970592021942]], [[0.1774463951587677, 0.26868411898612976, 0.03527391701936722, 0.01705012284219265, 0.00047759010340087116, 0.006241941824555397, 0.0031507122330367565, 0.2944689095020294, 0.038735195994377136, 0.003944840747863054, 0.004385389853268862, 0.004225992131978273, 0.03986744210124016, 0.00549504067748785, 0.07870971411466599, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00027908835909329355, 0.005506355315446854, 0.001626787707209587, 0.13775548338890076, 0.0008261757320724428, 0.00028156363987363875, 0.0002459189563523978, 0.0025131029542535543, 0.0009445812902413309, 0.001017659087665379, 0.002250042976811528, 0.0015115974238142371, 0.0017954352078959346, 0.0006745054270140827, 0.21780018508434296, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.021244889125227928, 0.1178143173456192, 0.008956437930464745, 0.14321640133857727, 0.023635229095816612, 0.3068733811378479, 0.15845780074596405, 0.3092327415943146, 0.0024783278349786997, 0.06481246650218964, 0.008965774439275265, 0.019083118066191673, 0.04005150496959686, 0.01112168189138174, 0.19139143824577332, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00042023108107969165, 0.0008873279439285398, 0.0019056870369240642, 0.007766622584313154, 0.23140135407447815, 0.5036463141441345, 0.015440672636032104, 0.008361338637769222, 0.001879698014818132, 0.0006688520661555231, 0.01133010908961296, 0.09722423553466797, 0.03314661607146263, 0.006971372757107019, 0.02285030484199524, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.002678314223885536, 0.004764833487570286, 0.0003137744788546115, 0.0006636036559939384, 0.07552827149629593, 0.36051952838897705, 0.21059149503707886, 0.11911091953516006, 0.00013829045929014683, 0.00018005385936703533, 0.00021675217431038618, 0.007453517522662878, 0.004449300933629274, 0.03708551451563835, 0.13281597197055817, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.008487393148243427, 0.014329447411000729, 0.005103611387312412, 0.0017902699764817953, 0.00018748251022771, 0.07080603390932083, 0.1865091174840927, 0.03389747440814972, 0.0026728338561952114, 0.00012369015894364566, 0.0001717496052151546, 0.0016556874616071582, 0.0035823825746774673, 0.018341869115829468, 0.2051384449005127, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0016413311241194606, 0.0038119314704090357, 0.0005628983490169048, 6.117233715485781e-05, 0.00011399950017221272, 0.0007454796577803791, 0.054881561547517776, 0.30246245861053467, 0.15667226910591125, 0.0004453254514373839, 0.0002609542279969901, 0.0001120980887208134, 0.0006856885738670826, 0.00573006272315979, 0.011146760545670986, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.001007524086162448, 0.0022212164476513863, 0.00036003260174766183, 2.8946307793376036e-05, 1.0167077562073246e-05, 0.00012231878645252436, 0.00022786400222685188, 0.03619853034615517, 0.005354967433959246, 0.003357505425810814, 0.0005030903848819435, 5.3131421736907214e-05, 4.2532476072665304e-05, 0.00010396525613032281, 0.2518664300441742, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.004948427900671959, 0.0037361346185207367, 0.0040338728576898575, 0.0015943445032462478, 3.9753424061927944e-05, 0.00016846440848894417, 0.00017597683472558856, 0.003258961718529463, 0.06328149139881134, 0.43567389249801636, 0.03252503648400307, 0.006277996581047773, 3.634384847828187e-05, 2.672040500328876e-05, 0.030029548332095146, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00322673749178648, 0.017767680808901787, 0.0033617434091866016, 0.029219835996627808, 0.0009114073473028839, 0.002889687195420265, 0.00012576105655170977, 0.01574547402560711, 0.0018639388727024198, 0.6032934188842773, 0.1301620751619339, 0.04121570661664009, 0.0035096178762614727, 0.00032833084696903825, 0.3004224896430969, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.033899419009685516, 0.07324357330799103, 0.00985381193459034, 0.017461512237787247, 0.019165849313139915, 0.07006029784679413, 0.01799222268164158, 0.013579626567661762, 0.00021177329472266138, 0.026033537462353706, 0.13102787733078003, 0.2077469676733017, 0.7029638886451721, 0.029135672375559807, 0.05414650961756706, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0015424743760377169, 0.007544125430285931, 0.010602829977869987, 0.0016127177514135838, 0.006006686482578516, 0.08514653891324997, 0.003129118587821722, 0.0036380700767040253, 1.298951519856928e-05, 6.919799488969147e-05, 0.0003367147874087095, 0.031529009342193604, 0.36636054515838623, 0.21289798617362976, 0.04463290795683861, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.005653384607285261, 0.005221519153565168, 0.010438429191708565, 0.0023121859412640333, 0.0034771040081977844, 0.01156994141638279, 0.006321457680314779, 0.006196276750415564, 2.671167931111995e-05, 0.00012823205906897783, 0.00023895784397609532, 0.0015353390481323004, 0.06888392567634583, 0.3010466396808624, 0.05789510905742645, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0025978884659707546, 0.0011408268474042416, 0.0005907863960601389, 0.0073682027868926525, 5.514698841579957e-06, 0.0001586068101460114, 0.0016139426734298468, 0.002635698765516281, 2.2516995159094222e-05, 7.803570952091832e-06, 4.170422926108586e-06, 4.799172893399373e-05, 8.148160122800618e-05, 0.006126015912741423, 0.363029420375824, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.018444720655679703, 0.036891017109155655, 0.08301377296447754, 0.04485299810767174, 0.0371856652200222, 0.0472157783806324, 0.022677546367049217, 0.017107300460338593, 0.03217196837067604, 0.03369837626814842, 0.021089907735586166, 0.018274538218975067, 0.020997297018766403, 0.034321803599596024, 0.1648317128419876, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01622859761118889, 0.0033176897559314966, 0.006228303536772728, 0.003451053285971284, 0.011415286920964718, 0.016942020505666733, 0.0027556640561670065, 0.001647507306188345, 0.0010015909792855382, 0.0013629572931677103, 0.004746851045638323, 0.009338179603219032, 0.00885467603802681, 0.006604180671274662, 0.16180677711963654, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17455320060253143, 0.026163265109062195, 0.2041780799627304, 0.027548620477318764, 0.4711945950984955, 0.5480062365531921, 0.10718726366758347, 0.032194506376981735, 0.08035919070243835, 0.010791448876261711, 0.11821587383747101, 0.04372825473546982, 0.5788823962211609, 0.10199426859617233, 0.06844703108072281, 0.13398022949695587, 0.051660239696502686, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.023936308920383453, 0.03560526669025421, 0.007881848141551018, 0.022994371131062508, 0.003501775674521923, 0.000663262908346951, 0.0027445319574326277, 0.0008202926255762577, 0.002215484855696559, 0.014335977844893932, 0.06139073148369789, 0.0039900378324091434, 0.004902976099401712, 0.006251698825508356, 0.21882350742816925, 0.14254364371299744, 0.023038247600197792, 0.14531654119491577, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01501577626913786, 0.026870740577578545, 0.007700353395193815, 0.02517320215702057, 0.005199552513659, 0.0040618558414280415, 0.0018289085710421205, 0.0005822794046252966, 0.008953371085226536, 0.004845716059207916, 0.02605423890054226, 0.010851072147488594, 0.011600007303059101, 0.011058725416660309, 0.2679094076156616, 0.17795929312705994, 0.024941343814134598, 0.06730933487415314, 0.21388311684131622, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05198093131184578, 0.026691097766160965, 0.04745011776685715, 0.02099662832915783, 0.007765383925288916, 0.0017653746763244271, 0.002459246199578047, 0.0005052239284850657, 0.0007161727407947183, 0.00449666241183877, 0.00950489193201065, 0.002728741616010666, 0.007593079470098019, 0.0031749741174280643, 0.1993207037448883, 0.09399491548538208, 0.3603954315185547, 0.2704434394836426, 0.1475897580385208, 0.18568314611911774, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0031879025045782328, 0.001219254801981151, 0.007273980416357517, 0.0029734931886196136, 9.794573998078704e-05, 0.0006066279602237046, 0.000905939843505621, 0.0002116545947501436, 0.00022416051069740206, 0.001432110439054668, 0.00046862047747708857, 0.0008043517009355128, 0.00010411434050183743, 0.0003457288257777691, 0.22099417448043823, 0.14775781333446503, 0.19919507205486298, 0.14170727133750916, 0.05924544855952263, 0.05067846551537514, 0.45942243933677673, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.020157048478722572, 0.026601465418934822, 0.04540588706731796, 0.04344630241394043, 0.0022944926749914885, 0.0010618591913953424, 0.00406603142619133, 0.0029086798895150423, 0.0019963555969297886, 0.010005260817706585, 0.0020353682339191437, 0.0019374215044081211, 0.0013613863848149776, 0.001661884132772684, 0.34173521399497986, 0.14211317896842957, 0.055850330740213394, 0.31645503640174866, 0.16900919377803802, 0.038168299943208694, 0.07897188514471054, 0.2625669240951538, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09776000678539276, 0.012011643499135971, 0.12930582463741302, 0.019725820049643517, 0.03450663015246391, 0.44516250491142273, 0.09379248321056366, 0.011904217302799225, 0.012111036106944084, 0.007218031212687492, 0.028761520981788635, 0.011232447810471058, 0.17035166919231415, 0.022308414801955223, 0.055901553481817245, 0.08848852664232254, 0.1616290658712387, 0.37575462460517883, 0.24721546471118927, 0.16591095924377441, 0.06889674067497253, 0.052010323852300644, 0.12634019553661346, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0270126610994339, 0.0034831874072551727, 0.03977394104003906, 0.025583824142813683, 0.0007700100541114807, 0.002870001830160618, 0.0027750579174607992, 0.0016644555144011974, 0.0016086471732705832, 0.001177149242721498, 0.00746855279430747, 0.002065857872366905, 0.0016993783647194505, 0.0015537800500169396, 0.32808277010917664, 0.0747382640838623, 0.14914710819721222, 0.6135430335998535, 0.5929751992225647, 0.35069379210472107, 0.2108047604560852, 0.11502823978662491, 0.02365955151617527, 0.17759312689304352, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16020068526268005, 0.019860466942191124, 0.3786206543445587, 0.04546584561467171, 0.22538548707962036, 0.035959187895059586, 0.022749971598386765, 0.0223965086042881, 0.010994979180395603, 0.013655508868396282, 0.08095952123403549, 0.07914181798696518, 0.5184871554374695, 0.24710357189178467, 0.059729527682065964, 0.02855301834642887, 0.21659326553344727, 0.4310435652732849, 0.40604472160339355, 0.3670090436935425, 0.48140615224838257, 0.27167943120002747, 0.09097199141979218, 0.1627163589000702, 0.1288144737482071, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.002354596508666873, 0.013563946820795536, 0.0012282072566449642, 0.0011236226418986917, 0.004269973374903202, 0.05393142253160477, 0.010044331662356853, 0.012847290374338627, 0.23206481337547302, 0.0042032524943351746, 0.002388538094237447, 0.005051162093877792, 0.004106870852410793, 0.003583247307687998, 0.0021634430158883333, 0.03365316241979599, 0.14809295535087585, 0.3644290566444397, 0.4046455919742584, 0.26744210720062256, 0.32108214497566223, 0.1678413599729538, 0.190241739153862, 0.22121649980545044, 0.03444775566458702, 0.46765974164009094, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1318124532699585, 0.006612265948206186, 0.026151085272431374, 0.15551267564296722, 0.006537565030157566, 0.045402105897665024, 0.08115606755018234, 0.020273711532354355, 0.2617640495300293, 0.03846455365419388, 0.42425140738487244, 0.0063036843203008175, 0.045534029603004456, 0.06594183295965195, 0.0061628553085029125, 0.038216885179281235, 0.2552680969238281, 0.4071650505065918, 0.3936895430088043, 0.4416206479072571, 0.38015541434288025, 0.1657901555299759, 0.15260477364063263, 0.22771137952804565, 0.10614379495382309, 0.0724361315369606, 0.1760038137435913, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0171976238489151, 0.0023818486370146275, 0.036466922610998154, 0.011855212040245533, 0.019672302529215813, 0.007386004086583853, 0.02982362173497677, 0.0045198979787528515, 0.02385052479803562, 0.25256073474884033, 0.2446560561656952, 0.0453505739569664, 0.08819476515054703, 0.09139581024646759, 0.0022182920947670937, 0.07068492472171783, 0.07818713039159775, 0.3302493095397949, 0.299561083316803, 0.46339741349220276, 0.48102065920829773, 0.15714748203754425, 0.27301517128944397, 0.38065311312675476, 0.19789563119411469, 0.11113718152046204, 0.05171056091785431, 0.13386131823062897, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.023948049172759056, 0.006307430099695921, 0.014840157702565193, 0.01758965104818344, 0.0009477039566263556, 0.00178795016836375, 0.005927308928221464, 0.0026511158794164658, 0.00012311375758145005, 0.04321818798780441, 0.0496363490819931, 0.3416200280189514, 0.001097637927159667, 0.007029203698039055, 0.007338459137827158, 0.05115865543484688, 0.44867002964019775, 0.49208834767341614, 0.477664977312088, 0.4642978608608246, 0.46059542894363403, 0.25649622082710266, 0.406831830739975, 0.27858051657676697, 0.2405669242143631, 0.11958811432123184, 0.1450459510087967, 0.0628136694431305, 0.09898709505796432, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1633826345205307, 0.005062526557594538, 0.04231903329491615, 0.24309031665325165, 0.0009563505300320685, 0.0008045694557949901, 0.004994159564375877, 0.0011061460245400667, 0.0013372766552492976, 0.023061903193593025, 0.044598180800676346, 0.0017028035363182425, 2.3589664124301635e-05, 0.0003540365141816437, 0.16737498342990875, 0.04031704366207123, 0.6707005500793457, 0.529548704624176, 0.4586588144302368, 0.3106471002101898, 0.6713098287582397, 0.4458201229572296, 0.5507155060768127, 0.6255134344100952, 0.5032600164413452, 0.18919125199317932, 0.2968505918979645, 0.3902440667152405, 0.16804949939250946, 0.088200144469738, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1106855720281601, 0.005593962036073208, 0.014953872188925743, 0.19064223766326904, 0.0008905718568712473, 0.002549833618104458, 0.019427485764026642, 0.019940704107284546, 0.0020017458591610193, 0.029780413955450058, 0.01774613931775093, 0.00061158457538113, 0.0022336822003126144, 0.007989613339304924, 0.2558586895465851, 0.13188821077346802, 0.1971314549446106, 0.3902590274810791, 0.4961083233356476, 0.37017205357551575, 0.46889960765838623, 0.2874276340007782, 0.1815745085477829, 0.39618349075317383, 0.17909032106399536, 0.26052209734916687, 0.13463276624679565, 0.11223814636468887, 0.05094114691019058, 0.030694767832756042, 0.23131275177001953, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07112060487270355, 0.029737049713730812, 0.09336916357278824, 0.07307538390159607, 0.023197662085294724, 0.022866347804665565, 0.060328319668769836, 0.04474486783146858, 0.0006379868718795478, 0.027103934437036514, 0.2942929267883301, 0.011375843547284603, 0.07746338844299316, 0.09051978588104248, 0.11258094012737274, 0.029627619311213493, 0.0727827325463295, 0.2382729947566986, 0.16726669669151306, 0.3644602298736572, 0.47072863578796387, 0.2034798413515091, 0.1723088026046753, 0.43477845191955566, 0.18565386533737183, 0.3540991544723511, 0.2379947453737259, 0.07713616639375687, 0.19858470559120178, 0.17015229165554047, 0.0891638696193695, 0.22899208962917328, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15941812098026276, 0.02997875213623047, 0.08360203355550766, 0.10365118086338043, 0.03050130233168602, 0.39312028884887695, 0.3065427839756012, 0.2912093997001648, 0.135236918926239, 0.18899840116500854, 0.13724294304847717, 0.1948302835226059, 0.07353706657886505, 0.12220755219459534, 0.10422825068235397, 0.01839388906955719, 0.10223808884620667, 0.244280606508255, 0.22035017609596252, 0.2828108072280884, 0.41914066672325134, 0.09010869264602661, 0.14338640868663788, 0.35142722725868225, 0.12073972821235657, 0.6723650693893433, 0.17433631420135498, 0.20010362565517426, 0.17566151916980743, 0.17214345932006836, 0.06743419170379639, 0.08234895765781403, 0.4274884760379791, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24064786732196808, 0.0051915524527430534, 0.09652373939752579, 0.2287912219762802, 0.019215410575270653, 0.13947954773902893, 0.15343742072582245, 0.07055477797985077, 0.05467608571052551, 0.10673969984054565, 0.5659986138343811, 0.014077076688408852, 0.1709020584821701, 0.23944324254989624, 0.026877261698246002, 0.02117752842605114, 0.17625343799591064, 0.2448491007089615, 0.23410049080848694, 0.3357784152030945, 0.2992798388004303, 0.09099920094013214, 0.1110134869813919, 0.20308172702789307, 0.1763213574886322, 0.1646280288696289, 0.23259523510932922, 0.3615821301937103, 0.32664546370506287, 0.296549916267395, 0.2726198732852936, 0.07387500256299973, 0.07587912678718567, 0.14093360304832458, NaN, NaN, NaN, NaN, NaN, NaN], [0.019817974418401718, 0.002034382661804557, 0.04978875443339348, 0.009913384914398193, 0.033772312104701996, 0.0069160182029008865, 0.027356693521142006, 0.004301261156797409, 0.005268980748951435, 0.24062182009220123, 0.2975090742111206, 0.09841412305831909, 0.13523375988006592, 0.1965852826833725, 0.004198803100734949, 0.05486638844013214, 0.06597498804330826, 0.2194771021604538, 0.1927901804447174, 0.37433308362960815, 0.412477970123291, 0.07100911438465118, 0.1499587744474411, 0.3056679368019104, 0.16932857036590576, 0.15193165838718414, 0.19111526012420654, 0.291239857673645, 0.37710845470428467, 0.510109543800354, 0.47089657187461853, 0.17204606533050537, 0.09759342670440674, 0.05198577418923378, 0.1557197868824005, NaN, NaN, NaN, NaN, NaN], [0.017094334587454796, 0.005556214600801468, 0.011722622439265251, 0.009952181950211525, 0.0008346029790118337, 0.0009373819339089096, 0.006794091779738665, 0.0019291864009574056, 4.7701923904241994e-05, 0.0364256277680397, 0.035398196429014206, 0.3890627920627594, 0.0013647697633132339, 0.008012092672288418, 0.013173048384487629, 0.03942986950278282, 0.2940163016319275, 0.3192412853240967, 0.3550935387611389, 0.28974649310112, 0.35144588351249695, 0.111830934882164, 0.2212614268064499, 0.1942923218011856, 0.16557106375694275, 0.12293191254138947, 0.3516637980937958, 0.22679129242897034, 0.3504909574985504, 0.4427362084388733, 0.6422855854034424, 0.29741936922073364, 0.17250965535640717, 0.13341550529003143, 0.05469499155879021, 0.0792233869433403, NaN, NaN, NaN, NaN], [0.12328237295150757, 0.0036286553367972374, 0.03202027454972267, 0.16562366485595703, 0.0006255045300349593, 0.00061140360776335, 0.00499368691816926, 0.0010923785157501698, 0.0008833102765493095, 0.03177933022379875, 0.04344986379146576, 0.00255553494207561, 2.260845576529391e-05, 0.0005036385264247656, 0.16160868108272552, 0.03949292004108429, 0.6095755696296692, 0.4376317858695984, 0.4024345874786377, 0.24819140136241913, 0.555855929851532, 0.2881583273410797, 0.40402302145957947, 0.5775710940361023, 0.42070186138153076, 0.22824901342391968, 0.4547353982925415, 0.567461371421814, 0.5762937664985657, 0.33163049817085266, 0.41951635479927063, 0.37286072969436646, 0.25620296597480774, 0.25266289710998535, 0.3395143151283264, 0.13239842653274536, 0.07333662360906601, NaN, NaN, NaN], [0.050196755677461624, 0.002699600299820304, 0.009293685667216778, 0.06999042630195618, 0.0006182404467836022, 0.0013977399794384837, 0.014421526342630386, 0.010930507443845272, 0.0008620836888439953, 0.015927143394947052, 0.008692404255270958, 0.0006625624373555183, 0.0011245491914451122, 0.0053406055085361, 0.2061784416437149, 0.11607979983091354, 0.18507249653339386, 0.30528268218040466, 0.41669708490371704, 0.22673273086547852, 0.3321194052696228, 0.17922396957874298, 0.1181870847940445, 0.299829363822937, 0.11785572022199631, 0.23005077242851257, 0.1731709986925125, 0.17971253395080566, 0.2448451966047287, 0.15796169638633728, 0.701153576374054, 0.1659945547580719, 0.4861533045768738, 0.20215842127799988, 0.13506482541561127, 0.058445703238248825, 0.03114200383424759, 0.21790345013141632, NaN, NaN], [0.04101766273379326, 0.020672734826803207, 0.08772061765193939, 0.04009746387600899, 0.01892852783203125, 0.017910925671458244, 0.057973578572273254, 0.03737492114305496, 0.00047206622548401356, 0.021084431558847427, 0.21054430305957794, 0.013546224683523178, 0.08985017240047455, 0.10610225051641464, 0.1389981210231781, 0.017429474741220474, 0.04190561920404434, 0.14842365682125092, 0.09654705971479416, 0.16489917039871216, 0.24686570465564728, 0.09686223417520523, 0.09368213266134262, 0.2918589413166046, 0.08991989493370056, 0.18521137535572052, 0.19666530191898346, 0.06316249072551727, 0.222347229719162, 0.3215444087982178, 0.3288835287094116, 0.38603323698043823, 0.4142700135707855, 0.25910744071006775, 0.0714699923992157, 0.2130158245563507, 0.1895158588886261, 0.07420682162046432, 0.2235250473022461, NaN], [0.018278781324625015, 0.03789714351296425, 0.00408195098862052, 0.005283118225634098, 0.009515376761555672, 0.11360906809568405, 0.008760524913668633, 0.006613489706069231, 0.018946174532175064, 0.008831392042338848, 0.015675490722060204, 0.021136337891221046, 0.13481837511062622, 0.08728663623332977, 0.15406787395477295, 0.011625233106315136, 0.13701221346855164, 0.3079974055290222, 0.17742200195789337, 0.10538481175899506, 0.17213597893714905, 0.08605048805475235, 0.13507568836212158, 0.2275547832250595, 0.07923908531665802, 0.07705283164978027, 0.2479921281337738, 0.3453103303909302, 0.2883259654045105, 0.36409828066825867, 0.18068012595176697, 0.4896908700466156, 0.399289608001709, 0.5261627435684204, 0.6339481472969055, 0.6382991671562195, 0.5417840480804443, 0.2542280852794647, 0.330732524394989, 0.21995915472507477]], [[0.2133164256811142, 0.025492815300822258, 0.20653849840164185, 0.07043907791376114, 0.10411863774061203, 0.3043566346168518, 0.06760577112436295, 0.5064103603363037, 0.08081910014152527, 0.27507925033569336, 0.5432406663894653, 0.27881479263305664, 0.16320040822029114, 0.2653813064098358, 0.11116068065166473, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.015402763150632381, 0.2444494515657425, 0.0030522451270371675, 0.00048490799963474274, 0.0026600188575685024, 0.06905494630336761, 0.012269481085240841, 0.014592616818845272, 0.004205085337162018, 0.0039128707721829414, 0.0037959537003189325, 0.012499181553721428, 0.02713301219046116, 0.00563135975971818, 0.19437076151371002, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04805738478899002, 0.007929358631372452, 0.4969516396522522, 0.08109094947576523, 0.008613435551524162, 0.06128339096903801, 0.020970679819583893, 0.014624540694057941, 0.001800250494852662, 0.04372387006878853, 0.036881472915410995, 0.022519467398524284, 0.032134752720594406, 0.17586740851402283, 0.15428785979747772, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.021660206839442253, 0.06483402103185654, 0.07990853488445282, 0.8655576705932617, 0.10770212858915329, 0.042777951806783676, 0.004243527539074421, 0.04141073673963547, 0.0011197980493307114, 0.0010354480473324656, 0.007620980031788349, 0.009411019273102283, 0.023886993527412415, 0.8532692193984985, 0.009252375923097134, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03802541270852089, 0.5626884698867798, 0.3869370222091675, 0.012873617932200432, 0.11968709528446198, 0.014900745823979378, 0.02957817167043686, 0.018288375809788704, 0.005979553796350956, 0.03379013389348984, 0.016338851302862167, 0.01766209304332733, 0.8086205720901489, 0.08052025735378265, 0.13067808747291565, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0663566142320633, 0.02082742564380169, 0.009716741740703583, 0.003548208624124527, 0.0008020728128030896, 0.4547119140625, 0.03523911535739899, 0.0031006578356027603, 0.006736437324434519, 0.0009184986702166498, 0.0011584048625081778, 0.04212343320250511, 0.019468490034341812, 0.001240313402377069, 0.20631356537342072, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.004470710642635822, 0.02006937935948372, 0.020011691376566887, 0.019766854122281075, 0.12330501526594162, 0.15558527410030365, 0.04160740226507187, 0.1780312955379486, 0.014384130015969276, 0.005233153235167265, 0.004123131278902292, 0.05227937176823616, 0.013469746336340904, 0.022578507661819458, 0.07922197878360748, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.17898443341255188, 0.006772744003683329, 0.041487641632556915, 0.009575014933943748, 0.016729410737752914, 0.2668032944202423, 0.12321095168590546, 0.6781973838806152, 0.0025635806377977133, 0.01087682880461216, 0.002732365159317851, 0.020299792289733887, 0.0031363710295408964, 0.0008204782498069108, 0.05180227383971214, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.12461799383163452, 0.013122161850333214, 0.02311752177774906, 0.0762406587600708, 0.09383975714445114, 0.007501720450818539, 0.07133012264966965, 0.008159258402884007, 0.13900579512119293, 0.006521029397845268, 0.021471921354532242, 0.012502939440310001, 0.0014349960256367922, 0.011674328707158566, 0.3848530650138855, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.014992507174611092, 0.010756749659776688, 0.10129547864198685, 0.15213072299957275, 0.1363232582807541, 0.16603931784629822, 0.0040587568655610085, 0.505429208278656, 0.0025213102344423532, 0.05678342655301094, 0.20746274292469025, 0.04314066469669342, 0.0019582516979426146, 0.01985819824039936, 0.18090446293354034, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.11427638679742813, 0.0123747568577528, 0.020808644592761993, 0.1336503028869629, 0.008563186042010784, 0.09643486887216568, 0.15193390846252441, 0.050255559384822845, 0.0023536821827292442, 0.3208443820476532, 0.021319447085261345, 0.003293143818154931, 0.027340535074472427, 0.01197835523635149, 0.09007034450769424, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.15923485159873962, 0.11477550864219666, 0.21969333291053772, 0.09681756794452667, 0.07061057537794113, 0.1670638769865036, 0.1398637294769287, 0.059452954679727554, 0.00850652251392603, 0.062244825065135956, 0.03212086483836174, 0.10482167452573776, 0.05658517777919769, 0.03675027936697006, 0.24718202650547028, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.004966236650943756, 0.001515651005320251, 0.002549123717471957, 0.006106496322900057, 0.00036676786839962006, 0.0014838402858003974, 0.008350875228643417, 0.003760475432500243, 9.004020830616355e-05, 0.003012964967638254, 0.000879374798387289, 0.0023141989950090647, 0.5349817276000977, 0.00013737898552790284, 0.18041089177131653, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [3.0577066354453564e-05, 0.00011073229688918218, 0.0002722943318076432, 0.00012968607188668102, 3.925479541067034e-05, 9.284611587645486e-05, 1.1375399481039494e-05, 0.00013649655738845468, 2.160583608201705e-05, 3.872126853821101e-06, 4.776401965500554e-06, 5.892393892281689e-05, 0.3018791675567627, 0.0016873051645234227, 0.00020723984926007688, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0053407615050673485, 0.002270790981128812, 0.015077341347932816, 0.008943013846874237, 0.01947944425046444, 0.013856526464223862, 0.021029049530625343, 0.011522401124238968, 0.019980257377028465, 0.021877266466617584, 0.03018842823803425, 0.06539047509431839, 0.04945596680045128, 0.008784771896898746, 0.1688213050365448, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05651351809501648, 0.11774645000696182, 0.026926513761281967, 0.04848615080118179, 0.10334916412830353, 0.4247743785381317, 0.21147629618644714, 0.6254463195800781, 0.10587190836668015, 0.08194849640130997, 0.04674661532044411, 0.35135090351104736, 0.35409873723983765, 0.43208518624305725, 0.11939813196659088, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05609016492962837, 0.06931670010089874, 0.1576625108718872, 0.27308744192123413, 0.04202406853437424, 0.2399596869945526, 0.3320065140724182, 0.6272499561309814, 0.09423039108514786, 0.144412100315094, 0.2769482433795929, 0.05643320456147194, 0.11388154327869415, 0.32551372051239014, 0.13187405467033386, 0.04915444552898407, 0.7444152235984802, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1798395812511444, 0.02382134646177292, 0.024498937651515007, 0.28730508685112, 0.19651466608047485, 0.13693250715732574, 0.34929007291793823, 0.1055094301700592, 0.08990196883678436, 0.5189381837844849, 0.3313819468021393, 0.34343984723091125, 0.21719343960285187, 0.21188895404338837, 0.15588119626045227, 0.10270431637763977, 0.20103313028812408, 0.23083212971687317, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.26584357023239136, 0.03035559318959713, 0.026536965742707253, 0.20298171043395996, 0.23938016593456268, 0.24181482195854187, 0.31930428743362427, 0.10626629739999771, 0.13103167712688446, 0.4636806845664978, 0.393515944480896, 0.3422740399837494, 0.342117577791214, 0.5495904088020325, 0.14030353724956512, 0.1558120846748352, 0.09243088960647583, 0.02280065417289734, 0.32627996802330017, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.30834218859672546, 0.3875667452812195, 0.32842832803726196, 0.16462059319019318, 0.416511207818985, 0.03730625659227371, 0.23662680387496948, 0.5092235207557678, 0.08549848943948746, 0.3278381824493408, 0.507111668586731, 0.0415511280298233, 0.5590415596961975, 0.6185146570205688, 0.0664283037185669, 0.1265193670988083, 0.1639627069234848, 0.12297425419092178, 0.08557231724262238, 0.1833999902009964, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0765935555100441, 0.29552146792411804, 0.05705742537975311, 0.01913047581911087, 0.15779250860214233, 0.030224651098251343, 0.08988720178604126, 0.3389361500740051, 0.08153010904788971, 0.05811480060219765, 0.09408371150493622, 0.19600677490234375, 0.6126919388771057, 0.623294472694397, 0.13969288766384125, 0.11118379235267639, 0.23907560110092163, 0.16732671856880188, 0.1982172429561615, 0.02825341187417507, 0.15412425994873047, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4304950535297394, 0.5688965320587158, 0.09143517911434174, 0.09618712961673737, 0.13307496905326843, 0.014428870752453804, 0.040250685065984726, 0.15830516815185547, 0.10923942923545837, 0.23653797805309296, 0.3180045783519745, 0.5594316720962524, 0.5058388710021973, 0.3866141140460968, 0.14058275520801544, 0.06564534455537796, 0.4107542335987091, 0.09891282767057419, 0.3507450222969055, 0.0021941487211734056, 0.004341787192970514, 0.11288701742887497, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.31169822812080383, 0.7707167863845825, 0.30778199434280396, 0.10994993895292282, 0.18047340214252472, 0.01769133098423481, 0.014783667400479317, 0.009741406887769699, 0.1340220719575882, 0.11223828792572021, 0.46960482001304626, 0.360332190990448, 0.56731116771698, 0.5470200181007385, 0.18929171562194824, 0.09254656732082367, 0.17870496213436127, 0.11882538348436356, 0.2565489113330841, 0.06709786504507065, 0.020701991394162178, 0.05621851608157158, 0.571487307548523, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2397254854440689, 0.361926406621933, 0.24345533549785614, 0.18179422616958618, 0.10373111069202423, 0.014045567251741886, 0.08654272556304932, 0.018043776974081993, 0.02193235233426094, 0.07134812325239182, 0.19312754273414612, 0.6192790865898132, 0.6039608716964722, 0.673239529132843, 0.15608295798301697, 0.12130707502365112, 0.06869146227836609, 0.052872415632009506, 0.07373122870922089, 0.03967232629656792, 0.019552208483219147, 0.024196362122893333, 0.1570335328578949, 0.3329051434993744, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.32110491394996643, 0.2706402838230133, 0.034645695239305496, 0.029830342158675194, 0.00933478306978941, 0.25964564085006714, 0.17791348695755005, 0.11580535769462585, 0.07073061913251877, 0.10197918862104416, 0.06440304219722748, 0.2378954440355301, 0.09358810633420944, 0.24307624995708466, 0.22625915706157684, 0.12370187789201736, 0.027735348790884018, 0.007442266680300236, 0.018701551482081413, 0.04923407360911369, 0.022976329550147057, 0.06834850460290909, 0.13354788720607758, 0.13089321553707123, 0.41554775834083557, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18688960373401642, 0.6521251797676086, 0.05505351349711418, 0.05518023297190666, 0.07190049439668655, 0.15721110999584198, 0.11867944896221161, 0.2974295914173126, 0.018550140783190727, 0.1645369827747345, 0.09910324215888977, 0.499615877866745, 0.34706613421440125, 0.5406060218811035, 0.24014075100421906, 0.08012630045413971, 0.020899765193462372, 0.032236725091934204, 0.011631320230662823, 0.1322554349899292, 0.13739252090454102, 0.3272823691368103, 0.10228703171014786, 0.16136890649795532, 0.12631160020828247, 0.3315902352333069, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24844318628311157, 0.24823600053787231, 0.41713690757751465, 0.05438315495848656, 0.5823535323143005, 0.1801777333021164, 0.13823869824409485, 0.16278210282325745, 0.035736992955207825, 0.017554355785250664, 0.03778500482439995, 0.09959819167852402, 0.18642207980155945, 0.26950401067733765, 0.24913227558135986, 0.07002493739128113, 0.03239390626549721, 0.05209453031420708, 0.033656563609838486, 0.10301846265792847, 0.08080227673053741, 0.10908480733633041, 0.10694557428359985, 0.2992934286594391, 0.26628223061561584, 0.1579413264989853, 0.18216297030448914, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21744470298290253, 0.04392259195446968, 0.5108200907707214, 0.27167755365371704, 0.5572997331619263, 0.30860280990600586, 0.5083038210868835, 0.6815038919448853, 0.3754148483276367, 0.01992654800415039, 0.0589066781103611, 0.07934294641017914, 0.15649113059043884, 0.3772245943546295, 0.25267744064331055, 0.23901967704296112, 0.02059122547507286, 0.03393668681383133, 0.04736512154340744, 0.05927135422825813, 0.02361929975450039, 0.006761881057173014, 0.05556455999612808, 0.1379650980234146, 0.12424714863300323, 0.191926509141922, 0.01547694206237793, 0.05743350088596344, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11088164150714874, 0.06568774580955505, 0.49295517802238464, 0.06175035238265991, 0.3928946256637573, 0.306259423494339, 0.1265336275100708, 0.29877781867980957, 0.061930101364851, 0.053618840873241425, 0.02546272985637188, 0.011733881197869778, 0.4200928509235382, 0.25557151436805725, 0.12701815366744995, 0.0662187710404396, 0.02669837884604931, 0.008789082989096642, 0.004751283209770918, 0.0528719425201416, 0.011242655105888844, 0.018989307805895805, 0.07620660215616226, 0.012969521805644035, 0.039284493774175644, 0.22954939305782318, 0.04563957825303078, 0.029234008863568306, 0.7488549947738647, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06005493924021721, 0.46575742959976196, 0.4922090172767639, 0.06956527382135391, 0.3788193464279175, 0.21330630779266357, 0.06565267592668533, 0.10461793839931488, 0.1200915202498436, 0.07597928494215012, 0.08451344817876816, 0.06952610611915588, 0.03487509861588478, 0.12158560007810593, 0.14820002019405365, 0.10826153308153152, 0.014460555277764797, 0.0725417360663414, 0.03217141702771187, 0.06698039174079895, 0.08051858842372894, 0.05872708931565285, 0.022866755723953247, 0.06705553829669952, 0.07034263759851456, 0.3507814407348633, 0.05356235057115555, 0.08709309250116348, 0.23604632914066315, 0.324868768453598, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11028759926557541, 0.4027779996395111, 0.8237467408180237, 0.1328621804714203, 0.7811888456344604, 0.5416622757911682, 0.16887041926383972, 0.2001309096813202, 0.08848496526479721, 0.05607001483440399, 0.13165172934532166, 0.10739479213953018, 0.052385441958904266, 0.05461856350302696, 0.16259506344795227, 0.13878783583641052, 0.02536645717918873, 0.06943535804748535, 0.05891912057995796, 0.006977759767323732, 0.003910682164132595, 0.004916978534311056, 0.04463541880249977, 0.07985055446624756, 0.07872368395328522, 0.291103333234787, 0.21302121877670288, 0.16995804011821747, 0.19893744587898254, 0.01890285685658455, 0.3838881254196167, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12960980832576752, 0.21605639159679413, 0.13754284381866455, 0.0687912181019783, 0.2001095861196518, 0.7652902007102966, 0.3308810591697693, 0.3389359712600708, 0.07430214434862137, 0.036511119455099106, 0.010612682439386845, 0.005050503648817539, 0.1584991067647934, 0.036481909453868866, 0.18724960088729858, 0.04579493775963783, 0.04550570994615555, 0.013287660665810108, 0.023886512964963913, 0.024052713066339493, 0.017023656517267227, 0.04836693033576012, 0.030526861548423767, 0.017645621672272682, 0.03170713782310486, 0.09266000241041183, 0.23106807470321655, 0.03557471185922623, 0.12432269752025604, 0.10334902256727219, 0.3233395516872406, 0.3770029842853546, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16838932037353516, 0.47491130232810974, 0.21776747703552246, 0.05912807583808899, 0.16565343737602234, 0.34125030040740967, 0.2414778620004654, 0.28169524669647217, 0.03973108157515526, 0.03921183571219444, 0.02238578163087368, 0.02449338510632515, 0.05498792976140976, 0.03159895911812782, 0.17659053206443787, 0.0394071489572525, 0.011173942126333714, 0.019201254472136497, 0.012027204036712646, 0.1043756976723671, 0.09629304707050323, 0.044260744005441666, 0.010774374939501286, 0.027033720165491104, 0.01529898401349783, 0.004158060997724533, 0.03471178933978081, 0.3574643135070801, 0.04469288885593414, 0.27014297246932983, 0.10925178974866867, 0.34427598118782043, 0.2875407040119171, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14295107126235962, 0.27777984738349915, 0.30436068773269653, 0.03198731318116188, 0.38494178652763367, 0.27411460876464844, 0.18790900707244873, 0.29966217279434204, 0.029011890292167664, 0.012050352990627289, 0.008839968591928482, 0.009298003278672695, 0.09229473769664764, 0.05935056507587433, 0.2074589878320694, 0.08343059569597244, 0.043180350214242935, 0.0767669752240181, 0.06360654532909393, 0.1271795630455017, 0.0800960585474968, 0.06889919936656952, 0.05648425221443176, 0.1521727591753006, 0.09240606427192688, 0.03566697984933853, 0.03560119867324829, 0.1492718607187271, 0.18653850257396698, 0.3474813401699066, 0.3278762698173523, 0.10706853121519089, 0.127774178981781, 0.1299499273300171, NaN, NaN, NaN, NaN, NaN, NaN], [0.185210719704628, 0.0802093893289566, 0.4863169491291046, 0.24164138734340668, 0.5185936689376831, 0.381059467792511, 0.5372542142868042, 0.6922534108161926, 0.40473121404647827, 0.015452258288860321, 0.03550630062818527, 0.023993153125047684, 0.09803077578544617, 0.14391310513019562, 0.25199130177497864, 0.23721955716609955, 0.02343675307929516, 0.03610215708613396, 0.05973569303750992, 0.07488072663545609, 0.026813305914402008, 0.0050082337111234665, 0.03149579092860222, 0.06251367926597595, 0.02305557392537594, 0.025774041190743446, 0.007636546157300472, 0.004965651780366898, 0.09922869503498077, 0.133448526263237, 0.1956746131181717, 0.04676169902086258, 0.27956491708755493, 0.021136147901415825, 0.057313986122608185, NaN, NaN, NaN, NaN, NaN], [0.08245678246021271, 0.1390499472618103, 0.5461503863334656, 0.060220371931791306, 0.43899697065353394, 0.5144884586334229, 0.22183947265148163, 0.5088672041893005, 0.09321429580450058, 0.05354699492454529, 0.02214067056775093, 0.004303250927478075, 0.39110496640205383, 0.12463895231485367, 0.1568218618631363, 0.0697786882519722, 0.028010839596390724, 0.012634677812457085, 0.007894599810242653, 0.0697624459862709, 0.015741104260087013, 0.01737123914062977, 0.05471426621079445, 0.0063003492541611195, 0.009287585504353046, 0.02825707383453846, 0.016440505161881447, 0.0038715004920959473, 0.07019948214292526, 0.02518516778945923, 0.041359793394804, 0.06545242667198181, 0.29174378514289856, 0.05010553449392319, 0.020036837086081505, 0.7549301981925964, NaN, NaN, NaN, NaN], [0.043030936270952225, 0.498334676027298, 0.5084810853004456, 0.06107298657298088, 0.3904430866241455, 0.35258427262306213, 0.08483341336250305, 0.17738159000873566, 0.1815967708826065, 0.09597334265708923, 0.08432064205408096, 0.040181081742048264, 0.02593160979449749, 0.08670566976070404, 0.14764654636383057, 0.12042609602212906, 0.016146911308169365, 0.09666067361831665, 0.04101520776748657, 0.09386932849884033, 0.11830881983041763, 0.08227012306451797, 0.02001151442527771, 0.0443122573196888, 0.028465820476412773, 0.11253371834754944, 0.02299223281443119, 0.013287386856973171, 0.043506089597940445, 0.09705191105604172, 0.08899306505918503, 0.14267200231552124, 0.1414598524570465, 0.04555709660053253, 0.08242949843406677, 0.2358742356300354, 0.30384859442710876, NaN, NaN, NaN], [0.0785449668765068, 0.4015392065048218, 0.8182658553123474, 0.10243776440620422, 0.7659414410591125, 0.5735372304916382, 0.16621330380439758, 0.21339072287082672, 0.12523002922534943, 0.05685745179653168, 0.1081186980009079, 0.07184037566184998, 0.02847907319664955, 0.031456008553504944, 0.15293413400650024, 0.14026813209056854, 0.02709769457578659, 0.07936792075634003, 0.07383942604064941, 0.01026969589293003, 0.007506935391575098, 0.01013263501226902, 0.043357811868190765, 0.054843299090862274, 0.032377004623413086, 0.07885654270648956, 0.05951513722538948, 0.021026868373155594, 0.029062975198030472, 0.004067933652549982, 0.00896876398473978, 0.031901001930236816, 0.2457016408443451, 0.1949184089899063, 0.16180625557899475, 0.23649972677230835, 0.020314330235123634, 0.390868216753006, NaN, NaN], [0.07311940938234329, 0.15430475771427155, 0.1386927217245102, 0.04823235049843788, 0.20945730805397034, 0.8191487193107605, 0.33371293544769287, 0.3618466258049011, 0.1152336597442627, 0.031010858714580536, 0.008395140990614891, 0.002998974174261093, 0.13362915813922882, 0.02411211095750332, 0.1613900512456894, 0.036581799387931824, 0.048626694828271866, 0.015552042052149773, 0.027681825682520866, 0.03610476478934288, 0.033903565257787704, 0.10816461592912674, 0.038128215819597244, 0.015381437726318836, 0.020138615742325783, 0.04596110060811043, 0.12391334027051926, 0.008882056921720505, 0.017164889723062515, 0.019657107070088387, 0.039318498224020004, 0.012226631864905357, 0.12883862853050232, 0.2578184902667999, 0.03228205814957619, 0.13855229318141937, 0.08962707966566086, 0.32015570998191833, 0.32621434330940247, NaN], [0.2622520923614502, 0.7386532425880432, 0.41215938329696655, 0.08539438247680664, 0.7665934562683105, 0.5218235850334167, 0.42940571904182434, 0.4037780165672302, 0.7456067204475403, 0.07961834967136383, 0.02781907096505165, 0.02608557976782322, 0.15701159834861755, 0.05025498941540718, 0.11428551375865936, 0.16620944440364838, 0.03880922496318817, 0.027515552937984467, 0.018877340480685234, 0.019147777929902077, 0.2389368712902069, 0.02623477764427662, 0.012871777638792992, 0.013969821855425835, 0.021991701796650887, 0.0026013199239969254, 0.00741098215803504, 0.01774594374001026, 0.003101027337834239, 0.007316285278648138, 0.009464021772146225, 0.007634901907294989, 0.005969886668026447, 0.011287253350019455, 0.04429420828819275, 0.016200777143239975, 0.03440575301647186, 0.14183124899864197, 0.1436305195093155, 0.03402799740433693]], [[0.09667091816663742, 0.08969368785619736, 0.16646768152713776, 0.01428181305527687, 0.1262292116880417, 0.03015410713851452, 0.00857650488615036, 0.013287652283906937, 0.013465571217238903, 0.009945754893124104, 0.03584994748234749, 0.07976501435041428, 0.013894102536141872, 0.07191513478755951, 0.16682514548301697, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00307486648671329, 0.2169581949710846, 0.015313946641981602, 0.005070009268820286, 0.13766343891620636, 0.036365993320941925, 0.013734312728047371, 0.012890451587736607, 0.00037508379318751395, 0.002069024136289954, 0.0038654597010463476, 0.007793853525072336, 0.006365353707224131, 0.02897111512720585, 0.19472798705101013, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.013033762574195862, 0.0016745100729167461, 0.09789733588695526, 0.11557573825120926, 0.070904940366745, 0.039959780871868134, 0.06112189590930939, 0.005926545709371567, 0.05931684747338295, 0.06562750041484833, 0.015556245110929012, 0.2949027419090271, 0.09280899167060852, 0.18960142135620117, 0.2321171909570694, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0009253448224626482, 0.0011463494738563895, 0.0022407870274037123, 0.022192178294062614, 0.18083734810352325, 0.18906380236148834, 0.06340676546096802, 0.5556718111038208, 0.008876022882759571, 0.00195835973136127, 0.009641225449740887, 0.13488754630088806, 0.03692271187901497, 0.0069083282724022865, 0.19416382908821106, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.020195724442601204, 0.0026999269612133503, 0.0047158133238554, 0.017117822542786598, 0.22690622508525848, 0.009801734238862991, 0.18513473868370056, 0.000916039280127734, 0.006044555455446243, 0.006021710112690926, 0.010346228256821632, 0.04500352963805199, 0.008295656181871891, 0.1122727021574974, 0.4271945357322693, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.02983868308365345, 0.03651329129934311, 0.005064305383712053, 0.00043434457620605826, 0.001774297677911818, 0.10316617041826248, 0.10274261981248856, 0.570116400718689, 0.0018607155652716756, 0.004884766880422831, 0.0001192242925753817, 0.01004798710346222, 0.011760696768760681, 0.020220324397087097, 0.036799319088459015, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.020830435678362846, 0.04066089913249016, 0.01340602245181799, 0.0007146665593609214, 0.05329689383506775, 0.010700137354433537, 0.06310626864433289, 0.1416247934103012, 0.059007443487644196, 0.009734428487718105, 0.023192377761006355, 0.030464952811598778, 0.011454294435679913, 0.06458231806755066, 0.29838618636131287, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04047420993447304, 0.05575861781835556, 0.0035385461524128914, 0.00047053993330337107, 0.010776028037071228, 0.0002634078555274755, 0.006466362159699202, 0.09768779575824738, 0.011305907741189003, 0.6455902457237244, 0.005685864482074976, 0.009437574073672295, 0.0014128481270745397, 0.0036261524073779583, 0.1994941532611847, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.001968077849596739, 0.00013096239126753062, 0.014192181639373302, 0.0025808673817664385, 1.1752749742299784e-05, 7.090794679243118e-05, 8.489128958899528e-05, 7.501097570639104e-05, 0.005588378757238388, 0.00024033378576859832, 0.7911840081214905, 0.0006417080294340849, 0.00012212486763019115, 0.0026151463389396667, 0.024830428883433342, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.007711799815297127, 0.006852409336715937, 0.005409319419413805, 0.029324712231755257, 0.0012151957489550114, 0.0014427780406549573, 0.0002848623844329268, 0.0011284908978268504, 0.00042831210885196924, 0.0035933239851146936, 0.2853389084339142, 0.04352247342467308, 0.0011324246879667044, 0.0015205255476757884, 0.05924868583679199, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.06333743035793304, 0.004831443540751934, 0.017261236906051636, 0.05893971398472786, 0.005950291641056538, 0.002105317311361432, 0.003185122972354293, 0.0028415010310709476, 0.004572128411382437, 0.007815520279109478, 0.07613655924797058, 0.10669270157814026, 0.027066918089985847, 0.03207901865243912, 0.4743220806121826, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10327208787202835, 0.004544916562736034, 0.05445469170808792, 0.010814311914145947, 0.026858847588300705, 0.011217474937438965, 0.07071709632873535, 0.05960191786289215, 0.0010665962472558022, 0.025403864681720734, 0.006131312809884548, 0.5720618963241577, 0.029676837846636772, 0.17520834505558014, 0.23297326266765594, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.011414228938519955, 0.002735550981014967, 0.015156290493905544, 0.0027777000796049833, 0.009832575917243958, 0.015552453696727753, 0.017305195331573486, 0.004722784738987684, 4.7792200348339975e-05, 0.0034479873720556498, 0.0004017044266220182, 0.0011886333813890815, 0.18307994306087494, 0.2786843478679657, 0.04159880056977272, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0032662157900631428, 0.004168938845396042, 0.0016457620076835155, 0.0005059303948655725, 0.0003206630062777549, 0.000853654695674777, 0.010604765266180038, 0.005784912034869194, 0.00014833646127954125, 0.0001704594906186685, 5.580573997576721e-05, 0.0004662217397708446, 0.0009024841128848493, 0.025914611294865608, 0.3543371260166168, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.057395875453948975, 0.01834016665816307, 0.017516011372208595, 0.011936328373849392, 0.010095582343637943, 0.018046732991933823, 0.24530914425849915, 0.01257838774472475, 0.014466731809079647, 0.027552323415875435, 0.054997242987155914, 0.013960911892354488, 0.0074861980974674225, 0.03251070901751518, 0.14566579461097717, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5009713768959045, 0.11806200444698334, 0.543484628200531, 0.29247328639030457, 0.5261343717575073, 0.23446989059448242, 0.5474087595939636, 0.062012095004320145, 0.8189043998718262, 0.538780152797699, 0.6200674176216125, 0.43515679240226746, 0.24830776453018188, 0.341129869222641, 0.04290800169110298, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.018064359202980995, 0.030848585069179535, 0.08071158826351166, 0.0676560178399086, 0.13447926938533783, 0.11551786214113235, 0.17043589055538177, 0.10128363966941833, 0.6618390679359436, 0.2855142652988434, 0.0971621423959732, 0.23388729989528656, 0.21859601140022278, 0.46025529503822327, 0.182326078414917, 0.13823550939559937, 0.01690824329853058, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04308566823601723, 0.03711610287427902, 0.06502576172351837, 0.10632220655679703, 0.09326566010713577, 0.08777783066034317, 0.3412204086780548, 0.6204424500465393, 0.8231819868087769, 0.09377399832010269, 0.1541169434785843, 0.21222646534442902, 0.11298450827598572, 0.15309588611125946, 0.11645805835723877, 0.1366243064403534, 0.10029595345258713, 0.03309698402881622, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07351326197385788, 0.05497964471578598, 0.07563240081071854, 0.32393333315849304, 0.057468246668577194, 0.2634526193141937, 0.3780488967895508, 0.7154850363731384, 0.7017503976821899, 0.20895157754421234, 0.29085400700569153, 0.06311048567295074, 0.03268700838088989, 0.14748480916023254, 0.03694311901926994, 0.14204008877277374, 0.17578311264514923, 0.058153361082077026, 0.03275991603732109, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15202973783016205, 0.07260382175445557, 0.07307075709104538, 0.01561899296939373, 0.03831832483410835, 0.04392734169960022, 0.07259247452020645, 0.03668325021862984, 0.315115749835968, 0.14016768336296082, 0.147903710603714, 0.09513753652572632, 0.08079177141189575, 0.04876280575990677, 0.1678115576505661, 0.15378697216510773, 0.06811928749084473, 0.031730279326438904, 0.02174059860408306, 0.06419884413480759, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20334205031394958, 0.03987862542271614, 0.2323523759841919, 0.08299659937620163, 0.11007620394229889, 0.049821991473436356, 0.05303451418876648, 0.020633194595575333, 0.20804192125797272, 0.621069610118866, 0.6013453006744385, 0.6998922824859619, 0.30664384365081787, 0.1810489445924759, 0.12484823167324066, 0.2336570769548416, 0.05475717782974243, 0.004165933933109045, 0.0025384188629686832, 0.005177688784897327, 0.12858138978481293, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.33830341696739197, 0.10967365652322769, 0.03348035365343094, 0.09579410403966904, 0.07735400646924973, 0.09874830394983292, 0.15181724727153778, 0.11190870404243469, 0.4600948095321655, 0.5270871520042419, 0.27297794818878174, 0.3748718500137329, 0.4609748125076294, 0.5019738078117371, 0.0790465772151947, 0.1292651742696762, 0.01662198081612587, 0.01174056064337492, 0.002378111705183983, 0.04036910459399223, 0.6038607358932495, 0.053664252161979675, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18835663795471191, 0.05185278132557869, 0.06106729805469513, 0.04512745887041092, 0.04466439411044121, 0.025852244347333908, 0.031750425696372986, 0.022515133023262024, 0.5077425837516785, 0.6734393835067749, 0.37964752316474915, 0.35936975479125977, 0.19831591844558716, 0.216437429189682, 0.2985125184059143, 0.13257111608982086, 0.0015173845458775759, 0.11979293078184128, 0.025075461715459824, 0.17128729820251465, 0.38108551502227783, 0.04533570259809494, 0.02173132263123989, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5560556054115295, 0.47877317667007446, 0.15116584300994873, 0.40482252836227417, 0.04176756739616394, 0.04773563891649246, 0.13619393110275269, 0.07804162055253983, 0.07037016749382019, 0.5527278780937195, 0.486864298582077, 0.22204715013504028, 0.2625967860221863, 0.19855597615242004, 0.060070205479860306, 0.12533389031887054, 0.01691550202667713, 0.03341663256287575, 0.04296481981873512, 0.13898836076259613, 0.21484552323818207, 0.09921174496412277, 0.178620383143425, 0.08540544658899307, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21585102379322052, 0.028776921331882477, 0.056070148944854736, 0.3207121789455414, 0.0078024002723395824, 0.016524065285921097, 0.3710367977619171, 0.14693383872509003, 0.12693363428115845, 0.6266815662384033, 0.6993157863616943, 0.5497558116912842, 0.14310741424560547, 0.3664083480834961, 0.047443971037864685, 0.19628551602363586, 0.0262758769094944, 0.06177970767021179, 0.020167797803878784, 0.21508394181728363, 0.05243970826268196, 0.05236654728651047, 0.019688904285430908, 0.04470491781830788, 0.03636182099580765, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.28475576639175415, 0.10818006843328476, 0.08735410869121552, 0.329417884349823, 0.02252645045518875, 0.04752267897129059, 0.3733118176460266, 0.39454737305641174, 0.029050499200820923, 0.6059318780899048, 0.7311877012252808, 0.44807982444763184, 0.29598307609558105, 0.33838847279548645, 0.16424106061458588, 0.10685201734304428, 0.1520930975675583, 0.22691352665424347, 0.1206204891204834, 0.20647111535072327, 0.3387817144393921, 0.17652125656604767, 0.14866295456886292, 0.058651361614465714, 0.13512541353702545, 0.029732942581176758, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08968453854322433, 0.11453098803758621, 0.20413988828659058, 0.368092805147171, 0.07694120705127716, 0.048818718641996384, 0.12943927943706512, 0.036333490163087845, 0.04509947448968887, 0.25635746121406555, 0.2806471586227417, 0.5608395338058472, 0.1390012502670288, 0.28897786140441895, 0.04701472818851471, 0.14931687712669373, 0.17397953569889069, 0.045104723423719406, 0.029273295775055885, 0.009919327683746815, 0.05321130529046059, 0.40632039308547974, 0.053491849452257156, 0.10154163092374802, 0.08916116505861282, 0.038379959762096405, 0.050926242023706436, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05315335839986801, 0.017116300761699677, 0.1720367670059204, 0.3916313052177429, 0.05510414391756058, 0.2876152992248535, 0.22692401707172394, 0.14989952743053436, 0.3368622660636902, 0.0913245752453804, 0.3484038710594177, 0.3637443780899048, 0.007217096630483866, 0.103476881980896, 0.036375418305397034, 0.1467411071062088, 0.6613936424255371, 0.30691561102867126, 0.27473992109298706, 0.05103013291954994, 0.09803401678800583, 0.18992389738559723, 0.012332501821219921, 0.08918186277151108, 0.009687116369605064, 0.01925584301352501, 0.0046735359355807304, 0.006799460854381323, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5125223994255066, 0.07351671159267426, 0.21591535210609436, 0.21059465408325195, 0.3288169205188751, 0.5466507077217102, 0.21618640422821045, 0.15017350018024445, 0.8681062459945679, 0.2442341297864914, 0.06865198910236359, 0.019835328683257103, 0.10077274590730667, 0.12228173017501831, 0.1682003289461136, 0.23535212874412537, 0.03722311928868294, 0.0383867472410202, 0.06886720657348633, 0.040591221302747726, 0.07368911802768707, 0.09838991612195969, 0.052333034574985504, 0.3684787154197693, 0.05692664161324501, 0.030762571841478348, 0.0074586388655006886, 0.017855344340205193, 0.004115242511034012, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4846254289150238, 0.17620818316936493, 0.23995715379714966, 0.09631974995136261, 0.22585628926753998, 0.04512355476617813, 0.06700992584228516, 0.01503949984908104, 0.07369402050971985, 0.03452376648783684, 0.04930250719189644, 0.1451164036989212, 0.010093613527715206, 0.020862746983766556, 0.16003692150115967, 0.17482686042785645, 0.020169643685221672, 0.038628242909908295, 0.03409411385655403, 0.011309999041259289, 0.013418656773865223, 0.010934274643659592, 0.0036632094997912645, 0.017374617978930473, 0.023464469239115715, 0.0031370571814477444, 0.004764250945299864, 0.022831382229924202, 0.0012565170181915164, 0.01132481824606657, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12189289927482605, 0.3658526837825775, 0.06606122851371765, 0.1638106107711792, 0.07819290459156036, 0.27624964714050293, 0.09599297493696213, 0.08126427978277206, 0.14055852591991425, 0.02327289618551731, 0.03783821687102318, 0.2963305115699768, 0.13405835628509521, 0.09205315262079239, 0.12166540324687958, 0.2204812914133072, 0.0262058824300766, 0.011961801908910275, 0.00864139012992382, 0.033310361206531525, 0.014301336370408535, 0.009627565741539001, 0.26419174671173096, 0.09070254862308502, 0.04369048774242401, 0.05080936849117279, 0.022543352097272873, 0.012377972714602947, 0.030277462676167488, 0.2341402769088745, 0.01971697248518467, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.278896301984787, 0.1438806802034378, 0.46959513425827026, 0.3356979489326477, 0.3651174008846283, 0.1071292906999588, 0.18117688596248627, 0.20183299481868744, 0.29131460189819336, 0.13872042298316956, 0.021824011579155922, 0.06362087279558182, 0.34404000639915466, 0.13715140521526337, 0.1120462715625763, 0.253863126039505, 0.004828702192753553, 0.05376851186156273, 0.11550138890743256, 0.1064227893948555, 0.03894256055355072, 0.006152869202196598, 0.03161965310573578, 0.06215812265872955, 0.10950783640146255, 0.01032247580587864, 0.005066303536295891, 0.011880352161824703, 0.09494113177061081, 0.06700112670660019, 0.10617008060216904, 0.020382743328809738, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2151702344417572, 0.2682046890258789, 0.2758127450942993, 0.20445802807807922, 0.06759822368621826, 0.058143485337495804, 0.21948587894439697, 0.1328936666250229, 0.04737214744091034, 0.09880322962999344, 0.06969184428453445, 0.0649414211511612, 0.09957331418991089, 0.08072139322757721, 0.15442174673080444, 0.04813924431800842, 0.008662978187203407, 0.10469061881303787, 0.06787187606096268, 0.02962217852473259, 0.04144993796944618, 0.019078848883509636, 0.10597121715545654, 0.0923849567770958, 0.24696239829063416, 0.010940729640424252, 0.060362689197063446, 0.059540145099163055, 0.36283043026924133, 0.1817280501127243, 0.2542697787284851, 0.10456714779138565, 0.017782384529709816, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10625648498535156, 0.3580685555934906, 0.2235240340232849, 0.2717205584049225, 0.14765356481075287, 0.1302592158317566, 0.182493656873703, 0.07402253895998001, 0.044094108045101166, 0.28373098373413086, 0.09141446650028229, 0.13240621984004974, 0.1622740924358368, 0.2716645896434784, 0.09359043836593628, 0.10143542289733887, 0.13917230069637299, 0.040259018540382385, 0.030723553150892258, 0.006155712995678186, 0.031952716410160065, 0.3338092863559723, 0.06915750354528427, 0.1324792504310608, 0.11542332917451859, 0.05764009431004524, 0.04023035988211632, 0.03596781566739082, 0.1495574563741684, 0.02840258926153183, 0.049019940197467804, 0.4096885919570923, 0.03150010108947754, 0.02953496389091015, NaN, NaN, NaN, NaN, NaN, NaN], [0.08181191235780716, 0.05183182656764984, 0.18780435621738434, 0.39972010254859924, 0.11086275428533554, 0.3443254232406616, 0.26716044545173645, 0.2157517671585083, 0.3917877972126007, 0.09846898168325424, 0.25891563296318054, 0.25942671298980713, 0.008535100147128105, 0.11220833659172058, 0.06895694881677628, 0.1521255224943161, 0.6490614414215088, 0.39427587389945984, 0.3861289620399475, 0.05361294746398926, 0.09808307886123657, 0.16810499131679535, 0.014004985801875591, 0.1451900601387024, 0.008040589280426502, 0.022555561736226082, 0.013471563346683979, 0.006859058979898691, 0.05312783271074295, 0.04058152437210083, 0.023753749206662178, 0.3811529278755188, 0.052651502192020416, 0.007359141018241644, 0.007947265170514584, NaN, NaN, NaN, NaN, NaN], [0.4507053792476654, 0.10277862101793289, 0.16431982815265656, 0.2027788907289505, 0.318918377161026, 0.4106469452381134, 0.24116744101047516, 0.1587350070476532, 0.8309358358383179, 0.2625651955604553, 0.047453198581933975, 0.009295494295656681, 0.07160880416631699, 0.07481760531663895, 0.19364440441131592, 0.2650813162326813, 0.032561566680669785, 0.05222610384225845, 0.09714324027299881, 0.038093939423561096, 0.08016244322061539, 0.09171951562166214, 0.056265611201524734, 0.42980653047561646, 0.0462084598839283, 0.03524700179696083, 0.017182864248752594, 0.04137876257300377, 0.007372017949819565, 0.08077534288167953, 0.07507885992527008, 0.050101280212402344, 0.02560576982796192, 0.006666052620857954, 0.016142593696713448, 0.003943128511309624, NaN, NaN, NaN, NaN], [0.5336673855781555, 0.18865860998630524, 0.19927646219730377, 0.10614699125289917, 0.21258802711963654, 0.035614922642707825, 0.07572873681783676, 0.021095039322972298, 0.08985494822263718, 0.061252057552337646, 0.05201297253370285, 0.10173538327217102, 0.008337927050888538, 0.017984798178076744, 0.15578274428844452, 0.186274453997612, 0.02024305984377861, 0.052268851548433304, 0.04830823838710785, 0.011142827570438385, 0.015970220789313316, 0.01383616030216217, 0.004258061293512583, 0.024750858545303345, 0.02320612221956253, 0.004944193176925182, 0.006908308248966932, 0.022138824686408043, 0.002315782941877842, 0.022694725543260574, 0.010753386653959751, 0.0032616793178021908, 0.0013332129456102848, 0.0031688748858869076, 0.015737321227788925, 0.00092066585784778, 0.009911282919347286, NaN, NaN, NaN], [0.11776354163885117, 0.337507039308548, 0.055947914719581604, 0.144154354929924, 0.09536269307136536, 0.2646341919898987, 0.10820504277944565, 0.0982295498251915, 0.1891198456287384, 0.027041049674153328, 0.03162495046854019, 0.2652260959148407, 0.10165920853614807, 0.07911970466375351, 0.1373925358057022, 0.2620354890823364, 0.032388050109148026, 0.01473915670067072, 0.01008685864508152, 0.03682388737797737, 0.017798764631152153, 0.012407293543219566, 0.2692665457725525, 0.10958822816610336, 0.03793380409479141, 0.07735131680965424, 0.03087974339723587, 0.01817244663834572, 0.0740593820810318, 0.5664002895355225, 0.01639901101589203, 0.07361851632595062, 0.02498074807226658, 0.01953950524330139, 0.011185318231582642, 0.024920325726270676, 0.19407986104488373, 0.01722806692123413, NaN, NaN], [0.20648452639579773, 0.10074114054441452, 0.42538517713546753, 0.26027214527130127, 0.3658106029033661, 0.09280957281589508, 0.23363487422466278, 0.27985435724258423, 0.3744349181652069, 0.1453229784965515, 0.02015594393014908, 0.05169985443353653, 0.3284047245979309, 0.12707991898059845, 0.12262601405382156, 0.27593934535980225, 0.005811678245663643, 0.07111961394548416, 0.13982559740543365, 0.1345955729484558, 0.06462955474853516, 0.009384723380208015, 0.03974011912941933, 0.0818282812833786, 0.09768332540988922, 0.015042337588965893, 0.006764655001461506, 0.01590757444500923, 0.11177312582731247, 0.1289886087179184, 0.2743605673313141, 0.018859822303056717, 0.01428449247032404, 0.0072670611552894115, 0.013756940141320229, 0.08787993341684341, 0.08323681354522705, 0.09635237604379654, 0.025643613189458847, NaN], [0.019576620310544968, 0.03319034352898598, 0.0111849969252944, 0.010870445519685745, 0.03222370147705078, 0.13807591795921326, 0.0675833523273468, 0.0615379698574543, 0.013822048902511597, 0.008804764598608017, 0.004974161274731159, 0.01815059222280979, 0.1774466335773468, 0.06282598525285721, 0.15396134555339813, 0.17263205349445343, 0.01194645743817091, 0.02866498939692974, 0.16296441853046417, 0.0019488729303702712, 0.034664519131183624, 0.05397665500640869, 0.1285821497440338, 0.10828299820423126, 0.02950196899473667, 0.008275950327515602, 0.008977574296295643, 0.09588290750980377, 0.01758315972983837, 0.00981396809220314, 0.06520896404981613, 0.03634792938828468, 0.007794357370585203, 0.007516053505241871, 0.0633511170744896, 0.016588596627116203, 0.008872142061591148, 0.04887184873223305, 0.025813041254878044, 0.0022019031457602978]], [[0.3107149600982666, 0.049285680055618286, 0.08128133416175842, 0.03986956924200058, 0.07088969647884369, 0.1961679309606552, 0.15016919374465942, 0.05429982393980026, 0.1291487067937851, 0.03663256764411926, 0.25306442379951477, 0.3913470208644867, 0.2542778253555298, 0.252127081155777, 0.15921251475811005, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10834414511919022, 0.3508348762989044, 0.02124197781085968, 0.019397908821702003, 0.026673240587115288, 0.3167271912097931, 0.11886779963970184, 0.17699773609638214, 0.14507175981998444, 0.115145742893219, 0.6241064667701721, 0.1622784435749054, 0.5683063268661499, 0.15724869072437286, 0.12728430330753326, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6979861855506897, 0.039286430925130844, 0.3014020621776581, 0.003208757843822241, 0.01772892102599144, 0.014036925509572029, 0.19886529445648193, 0.09335973858833313, 0.4060034155845642, 0.28424081206321716, 0.26539483666419983, 0.1895008385181427, 0.4672236740589142, 0.16107353568077087, 0.10992881655693054, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5298255681991577, 0.6474234461784363, 0.19260530173778534, 0.026028962805867195, 0.013013242743909359, 0.01466711051762104, 0.11121421307325363, 0.06523838639259338, 0.29339125752449036, 0.46135157346725464, 0.7174844145774841, 0.3618351221084595, 0.19526919722557068, 0.0703459233045578, 0.24330592155456543, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7494951486587524, 0.23358309268951416, 0.3640848398208618, 0.09014757722616196, 0.32190942764282227, 0.0021980239544063807, 0.07713330537080765, 0.030900368466973305, 0.08560045808553696, 0.26394325494766235, 0.11549779027700424, 0.44356539845466614, 0.12175428122282028, 0.3783136308193207, 0.14015373587608337, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3064809739589691, 0.15617568790912628, 0.4955383241176605, 0.8125641942024231, 0.02114781178534031, 0.2633197009563446, 0.014569958671927452, 0.04754461348056793, 0.03227522596716881, 0.09995166957378387, 0.0697590634226799, 0.0770602896809578, 0.19454655051231384, 0.18272873759269714, 0.19963966310024261, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5314973592758179, 0.5086395144462585, 0.5757231116294861, 0.44031307101249695, 0.2709468603134155, 0.0639616996049881, 0.2984015941619873, 0.0039451331831514835, 0.0197422094643116, 0.0031917106825858355, 0.05093149095773697, 0.12591752409934998, 0.25977155566215515, 0.0615861676633358, 0.3711840510368347, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2939777970314026, 0.2997593581676483, 0.5167340040206909, 0.46100836992263794, 0.39705657958984375, 0.5034002065658569, 0.07978513836860657, 0.0779491513967514, 0.012053987942636013, 0.01132633350789547, 0.028715649619698524, 0.059212565422058105, 0.20603224635124207, 0.15584728121757507, 0.14816488325595856, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3128078877925873, 0.0864272266626358, 0.7678588032722473, 0.6537591814994812, 0.8236088752746582, 0.6979317665100098, 0.30976778268814087, 0.014760972931981087, 0.5645584464073181, 0.004590533208101988, 0.008271697908639908, 0.012132997624576092, 0.028745530173182487, 0.04464057460427284, 0.1669740080833435, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6456499099731445, 0.1693999022245407, 0.7097220420837402, 0.5244839191436768, 0.46365103125572205, 0.5023244023323059, 0.9643971920013428, 0.24913577735424042, 0.13337120413780212, 0.06419410556554794, 0.012416149489581585, 0.0573885552585125, 0.016666844487190247, 0.008706454187631607, 0.1754455268383026, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.09960467368364334, 0.0907629206776619, 0.36143985390663147, 0.11092879623174667, 0.19937658309936523, 0.03214935213327408, 0.3196737766265869, 0.4763943552970886, 0.497630774974823, 0.1899363249540329, 0.1145005002617836, 0.004749455489218235, 0.0008605146431364119, 0.0007969819707795978, 0.02025206945836544, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3807562589645386, 0.26623356342315674, 0.4209006428718567, 0.27443018555641174, 0.5137820839881897, 0.1592678278684616, 0.6250110864639282, 0.6178545951843262, 0.9692861437797546, 0.5716569423675537, 0.22724294662475586, 0.17567582428455353, 0.008769324980676174, 0.002557128667831421, 0.05025441572070122, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2969632148742676, 0.16767999529838562, 0.46978121995925903, 0.28813451528549194, 0.45300158858299255, 0.33029136061668396, 0.6236194968223572, 0.1634167730808258, 0.8177276253700256, 0.718397855758667, 0.9021148681640625, 0.07875741273164749, 0.09992827475070953, 0.004932410083711147, 0.1707668900489807, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3945808410644531, 0.3581867516040802, 0.5247420072555542, 0.4120633900165558, 0.3024104833602905, 0.35548633337020874, 0.5872392654418945, 0.15815261006355286, 0.7289484143257141, 0.7948301434516907, 0.9396543502807617, 0.9256777167320251, 0.08537369966506958, 0.03166399896144867, 0.03224433213472366, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.004588960204273462, 0.041907694190740585, 0.17755450308322906, 0.039724841713905334, 0.047663237899541855, 0.09274838864803314, 0.010110240429639816, 0.014862497337162495, 0.11161036789417267, 0.0490046888589859, 0.18517035245895386, 0.029471391811966896, 0.05094437301158905, 0.002971563721075654, 0.16300250589847565, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07712388038635254, 0.042244281619787216, 0.004363007377833128, 0.0015959119191393256, 0.019252488389611244, 0.02118455246090889, 0.001846740604378283, 0.0012080060550943017, 0.0007866616360843182, 0.001261864323168993, 0.002815018408000469, 0.017323212698101997, 0.00286104716360569, 0.004067797679454088, 0.15733002126216888, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.176344633102417, 0.3271441161632538, 0.08498391509056091, 0.04002806171774864, 0.06676299124956131, 0.008946515619754791, 0.012590638361871243, 0.0061616976745426655, 0.010515754111111164, 0.042563267052173615, 0.024306243285536766, 0.009260479360818863, 0.0002838150830939412, 0.0009972971165552735, 0.0829070582985878, 0.13826748728752136, 0.016647184267640114, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3345734477043152, 0.016792800277471542, 0.785018265247345, 0.16747814416885376, 0.3955724537372589, 0.09289640188217163, 0.041390396654605865, 0.004024161957204342, 0.04094661772251129, 0.023736434057354927, 0.20348279178142548, 0.041674140840768814, 0.012969214469194412, 0.03994787111878395, 0.04405270516872406, 0.12115656584501266, 0.053111400455236435, 0.35221540927886963, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.027460135519504547, 0.0009503767942078412, 0.8045902252197266, 0.05251304432749748, 0.4111766219139099, 0.08071836084127426, 0.01928381621837616, 0.0005491983611136675, 0.029575586318969727, 0.001678029540926218, 0.033282194286584854, 0.007144003175199032, 0.012064780108630657, 0.008930332958698273, 0.0033295771572738886, 0.06620940566062927, 0.0874415934085846, 0.3174281120300293, 0.09698687493801117, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18455208837985992, 0.0566692017018795, 0.08522135764360428, 0.2798183560371399, 0.013304274529218674, 0.0006802850402891636, 0.09522412717342377, 0.0060977875255048275, 0.002369458321481943, 0.017453324049711227, 0.0036190226674079895, 2.9809654733981006e-05, 0.0002128492487827316, 0.0002820969675667584, 0.18610867857933044, 0.05510773882269859, 0.045387670397758484, 0.35701045393943787, 0.5011870265007019, 0.0787656381726265, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.6536933779716492, 0.3485175371170044, 0.2007695585489273, 0.8106443881988525, 0.12433423846960068, 0.008092332631349564, 0.6807736158370972, 0.40895989537239075, 0.04516575112938881, 0.1387551873922348, 0.004862201400101185, 0.0003120531910099089, 0.00022667655139230192, 0.00031860917806625366, 0.07640787214040756, 0.05231153964996338, 0.1393265277147293, 0.34751832485198975, 0.15474379062652588, 0.1892920285463333, 0.06652400642633438, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08564082533121109, 0.05155009403824806, 0.10021068900823593, 0.5880905985832214, 0.0823356956243515, 0.0626063123345375, 0.7381499409675598, 0.566346287727356, 0.04188016802072525, 0.02469027414917946, 0.004355741199105978, 0.00042968738125637174, 2.4299803044414148e-05, 2.7212277927901596e-05, 0.001896930974908173, 0.04669328033924103, 0.038986966013908386, 0.38860636949539185, 0.09904015064239502, 0.3339899182319641, 0.027963249012827873, 0.04134462773799896, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03975995257496834, 0.012421448715031147, 0.08890707790851593, 0.605818510055542, 0.05048904940485954, 0.017510779201984406, 0.24702893197536469, 0.39587050676345825, 0.06098005548119545, 0.052625395357608795, 0.013424866832792759, 0.0005194320692680776, 0.000250102486461401, 0.0003063087642658502, 0.0010793216060847044, 0.20758312940597534, 0.07789289951324463, 0.047907259315252304, 0.006299893371760845, 0.2608397901058197, 0.044556185603141785, 0.061705876141786575, 0.034865181893110275, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11902385950088501, 0.011114073917269707, 0.22151720523834229, 0.2006509006023407, 0.03878694027662277, 0.01363028772175312, 0.3268369734287262, 0.04311302676796913, 0.8067907094955444, 0.34777864813804626, 0.25920552015304565, 0.09021251648664474, 0.035271789878606796, 0.0031717135570943356, 0.004271878860890865, 0.18052776157855988, 0.08179321140050888, 0.059846919029951096, 0.02793782763183117, 0.062999427318573, 0.04310278594493866, 0.024987775832414627, 0.015387488529086113, 0.132792130112648, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.006270309444516897, 0.0001492560259066522, 0.00045137249981053174, 0.0007612273329868913, 7.476524478988722e-05, 0.013270817697048187, 0.04344405606389046, 0.014117085374891758, 0.6041488647460938, 0.07304701954126358, 0.010559855960309505, 0.0026350386906415224, 0.02638809196650982, 0.002994539914652705, 0.00020572090579662472, 0.03587701544165611, 0.020078828558325768, 0.04571571201086044, 0.02593454346060753, 0.007220670115202665, 0.03280382603406906, 0.012364541180431843, 0.04736338183283806, 0.48638036847114563, 0.015403805300593376, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.002078789984807372, 0.000502656155731529, 0.00018232718866784126, 0.0008548289188183844, 0.0009249084978364408, 0.02029070071876049, 0.012032798491418362, 0.024348178878426552, 0.2300865352153778, 0.10343841463327408, 0.007660495117306709, 0.0012821657583117485, 0.0114271380007267, 0.0009412667131982744, 7.524124521296471e-05, 0.010417330078780651, 0.019508572295308113, 0.03964173421263695, 0.041229844093322754, 0.021899865940213203, 0.0029071751050651073, 0.010124437510967255, 0.08508285880088806, 0.40291228890419006, 0.4734281599521637, 0.015163381583988667, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.022463228553533554, 0.0013134862529113889, 0.00013891702110413462, 0.002816978842020035, 0.0011811865260824561, 0.0014538302784785628, 0.0005458829691633582, 0.0004073161107953638, 0.000992793939076364, 0.626685380935669, 0.1310541182756424, 0.1785772740840912, 0.1327074021100998, 0.014590581879019737, 3.459410072537139e-05, 0.08744391798973083, 0.1107466071844101, 0.15557123720645905, 0.13837403059005737, 0.05803389474749565, 0.026755833998322487, 0.03754325956106186, 0.4220706820487976, 0.16102783381938934, 0.2859216034412384, 0.1457504779100418, 0.03281670808792114, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.004299411084502935, 0.00014757749158889055, 0.0013493087608367205, 0.003552102018147707, 0.004041418433189392, 0.004232631530612707, 0.00022051982523407787, 5.3625211876351386e-05, 0.008671559393405914, 0.2003454566001892, 0.2010745257139206, 0.20048564672470093, 0.327506959438324, 0.12215141952037811, 7.573522452730685e-05, 0.21633882820606232, 0.07441287487745285, 0.04740259423851967, 0.026924576610326767, 0.012407396920025349, 0.002398786135017872, 0.0038467273116111755, 0.13835540413856506, 0.06710492819547653, 0.026295386254787445, 0.17057135701179504, 0.013244924135506153, 0.46883779764175415, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.011497906409204006, 0.0014132088981568813, 0.002270179335027933, 0.006387166678905487, 5.5530636018374935e-05, 0.0020248510409146547, 0.0021348590962588787, 0.001147052156738937, 0.0024277162738144398, 0.3687064051628113, 0.5298402905464172, 0.006611559074372053, 0.3372868299484253, 0.2915361225605011, 0.0002606022753752768, 0.027107199653983116, 0.05742119997739792, 0.06533583253622055, 0.024222400039434433, 0.014050583355128765, 0.013653005473315716, 0.0030738371424376965, 0.04425956308841705, 0.06826918572187424, 0.011929179541766644, 0.14959540963172913, 0.16161218285560608, 0.5212987065315247, 0.041249219328165054, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.043351031839847565, 0.015730101615190506, 0.006545424461364746, 0.11301398277282715, 0.001535893650725484, 0.0002994980022776872, 0.002417969051748514, 0.0027875620871782303, 0.007663458585739136, 0.4366588592529297, 0.29866132140159607, 0.03879629448056221, 0.0005757116014137864, 0.10755223035812378, 0.15693426132202148, 0.12232528626918793, 0.02327316626906395, 0.043996360152959824, 0.010462167672812939, 0.05786772817373276, 0.006097386125475168, 0.001271827262826264, 0.022651376202702522, 0.03627351298928261, 0.030646052211523056, 0.03145253658294678, 0.18536151945590973, 0.10030946880578995, 0.3235938847064972, 0.09760642796754837, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05824243649840355, 0.00918568018823862, 0.004823020659387112, 0.12202360481023788, 0.001364732626825571, 0.009540650062263012, 0.017077280208468437, 0.02250218391418457, 0.031557418406009674, 0.39489659667015076, 0.4118596911430359, 0.4739699363708496, 0.04330656677484512, 0.22410848736763, 0.009354491718113422, 0.01696004532277584, 0.0005225083441473544, 0.012039890512824059, 0.0003213977033738047, 0.024568837136030197, 0.0005492557538673282, 6.035636397427879e-05, 0.0032521369867026806, 0.016784805804491043, 0.013033770024776459, 0.023488081991672516, 0.04594254866242409, 0.04732683673501015, 0.2366781234741211, 0.2578820288181305, 0.02447950839996338, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10114194452762604, 0.055991608649492264, 0.0056193675845861435, 0.044799599796533585, 0.005612906999886036, 0.0018076150445267558, 0.0035521595273166895, 0.003050913568586111, 0.014126029796898365, 0.18568304181098938, 0.044660091400146484, 0.8178999423980713, 0.12312521040439606, 0.22830259799957275, 0.0015339198289439082, 0.016271475702524185, 0.026037830859422684, 0.05988215655088425, 0.04065781086683273, 0.0548781082034111, 0.0059303357265889645, 0.000490839418489486, 0.009792556054890156, 0.05564826726913452, 0.029693011194467545, 0.015783851966261864, 0.050408631563186646, 0.10483089834451675, 0.18894171714782715, 0.4590488076210022, 0.24355939030647278, 0.03408684581518173, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17329555749893188, 0.022842630743980408, 0.03050464764237404, 0.3040459156036377, 0.023058682680130005, 0.05675753578543663, 0.012084487825632095, 0.018060212954878807, 0.012510768137872219, 0.4205268621444702, 0.403047114610672, 0.5196431279182434, 0.14466160535812378, 0.15726853907108307, 0.003281315555796027, 0.011992339976131916, 0.02786487340927124, 0.025577154010534286, 0.02912752889096737, 0.009845648892223835, 0.0007121131638996303, 0.001387864351272583, 0.015649031847715378, 0.05334821715950966, 0.05039743706583977, 0.0003855754912365228, 0.07798124849796295, 0.03745294734835625, 0.16697214543819427, 0.29521557688713074, 0.2776513993740082, 0.29445046186447144, 0.031993161886930466, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21814380586147308, 0.013853680342435837, 0.0011839027283713222, 0.02006133459508419, 0.0059941732324659824, 0.004335244186222553, 0.0006587213138118386, 0.0008069095201790333, 6.766151636838913e-05, 0.4439576268196106, 0.16648612916469574, 0.7347545623779297, 0.19459886848926544, 0.05657987296581268, 0.0006026092451065779, 0.11517049372196198, 0.11416894942522049, 0.19162771105766296, 0.14611610770225525, 0.060761958360672, 0.02055470645427704, 0.021888524293899536, 0.20655019581317902, 0.047658227384090424, 0.055987950414419174, 0.01683689095079899, 0.005808014422655106, 0.045862384140491486, 0.09340663254261017, 0.10908356308937073, 0.18944555521011353, 0.26804569363594055, 0.20485185086727142, 0.037772081792354584, NaN, NaN, NaN, NaN, NaN, NaN], [0.034262340515851974, 0.0017182001611217856, 0.005656392779201269, 0.017169898375868797, 0.0156857930123806, 0.01468763966113329, 0.0007699507405050099, 0.00017933807976078242, 0.002019587904214859, 0.09474337100982666, 0.21286551654338837, 0.39837440848350525, 0.44769343733787537, 0.30061447620391846, 0.0009720441303215921, 0.24184046685695648, 0.07921410351991653, 0.056290365755558014, 0.026794791221618652, 0.016941547393798828, 0.0021516080014407635, 0.0023830668069422245, 0.05685606598854065, 0.02070370689034462, 0.003236053278669715, 0.01165463775396347, 0.004370343871414661, 0.030780060216784477, 0.00907946564257145, 0.06188458576798439, 0.04407832771539688, 0.006142587400972843, 0.14762946963310242, 0.013672620058059692, 0.4999893307685852, NaN, NaN, NaN, NaN, NaN], [0.1974877417087555, 0.05350746586918831, 0.02080627717077732, 0.07140190154314041, 0.0007820951868779957, 0.021851971745491028, 0.023295408114790916, 0.011020028032362461, 0.0015720969531685114, 0.3204348385334015, 0.5890824198722839, 0.011122598312795162, 0.40923523902893066, 0.5521805882453918, 0.009284045547246933, 0.03566991165280342, 0.0538097508251667, 0.09943600744009018, 0.028607800602912903, 0.020965654402971268, 0.013461945578455925, 0.002478980924934149, 0.02911236882209778, 0.02446376532316208, 0.0022762087173759937, 0.010774179361760616, 0.04047773778438568, 0.06471210718154907, 0.0026813328731805086, 0.07523855566978455, 0.030470186844468117, 0.0345987044274807, 0.1238497719168663, 0.17781274020671844, 0.4970780611038208, 0.04515520855784416, NaN, NaN, NaN, NaN], [0.04384012520313263, 0.020103074610233307, 0.00601673498749733, 0.10121199488639832, 0.0015372235793620348, 0.00047879578778520226, 0.0028034253045916557, 0.0035304632037878036, 0.0019347126362845302, 0.15543726086616516, 0.10060140490531921, 0.012154079042375088, 0.00020098914683330804, 0.049742307513952255, 0.15931616723537445, 0.12716706097126007, 0.02434932254254818, 0.05787394568324089, 0.013031681068241596, 0.06681805849075317, 0.007088592275977135, 0.0018475945107638836, 0.021072670817375183, 0.024636711925268173, 0.010089303366839886, 0.0076353950425982475, 0.05158482864499092, 0.009980393573641777, 0.034229546785354614, 0.01627102866768837, 0.008032353594899178, 0.013575052842497826, 0.04940066114068031, 0.19428585469722748, 0.10819438844919205, 0.2976790964603424, 0.08516447991132736, NaN, NaN, NaN], [0.33183732628822327, 0.07794758677482605, 0.02364480309188366, 0.3878714144229889, 0.007764760870486498, 0.055411770939826965, 0.07855504751205444, 0.09397301822900772, 0.02721172571182251, 0.38145557045936584, 0.42047446966171265, 0.5078706741333008, 0.03859835863113403, 0.25985077023506165, 0.0625251829624176, 0.01713084802031517, 0.000499976216815412, 0.019638467580080032, 0.00048709739348851144, 0.03356647491455078, 0.0008144291932694614, 0.00011953162174904719, 0.003664336632937193, 0.013800683431327343, 0.004805452190339565, 0.004433726891875267, 0.011711561121046543, 0.003556638490408659, 0.01588965393602848, 0.025807680562138557, 0.00022126971452962607, 0.004036479629576206, 0.00837762001901865, 0.04655361920595169, 0.04086336866021156, 0.22630761563777924, 0.2765483856201172, 0.02425519935786724, NaN, NaN], [0.4473247230052948, 0.3730325996875763, 0.029895052313804626, 0.15908104181289673, 0.02762797847390175, 0.008889964781701565, 0.016516737639904022, 0.012883803807199001, 0.01523641124367714, 0.22003965079784393, 0.05771813541650772, 0.8456536531448364, 0.1770154982805252, 0.31127816438674927, 0.007925343699753284, 0.010901566594839096, 0.020337969064712524, 0.07802019268274307, 0.0504593625664711, 0.06312800198793411, 0.009868033230304718, 0.000861799344420433, 0.010114955715835094, 0.052247028797864914, 0.012602821923792362, 0.005399123765528202, 0.01934058591723442, 0.013776490464806557, 0.010564911179244518, 0.04300173744559288, 0.008748980239033699, 0.0006391598144546151, 0.006108305882662535, 0.05087457224726677, 0.09035929292440414, 0.18751013278961182, 0.4462290108203888, 0.28552356362342834, 0.05451636388897896, NaN], [0.2188224196434021, 0.06026163697242737, 0.01674255169928074, 0.1205059364438057, 0.017392028123140335, 0.033714599907398224, 0.013199009001255035, 0.035441260784864426, 0.006878681946545839, 0.5097362399101257, 0.5390803217887878, 0.7098195552825928, 0.20610427856445312, 0.34404870867729187, 0.06464894115924835, 0.1367119550704956, 0.02979014255106449, 0.04602046683430672, 0.022530242800712585, 0.009278235025703907, 0.01184787880629301, 0.010125648230314255, 0.02445557340979576, 0.052750833332538605, 0.013119504787027836, 0.0006633299053646624, 0.007243738044053316, 0.02398994006216526, 0.00908573716878891, 0.013761860318481922, 0.007176807615906, 0.00677318312227726, 0.0021949538495391607, 0.01309704128652811, 0.09677710384130478, 0.12711098790168762, 0.1613820642232895, 0.37058699131011963, 0.3504316806793213, 0.02586444839835167]], [[6.113462859502761e-06, 0.5065946578979492, 7.261813152581453e-05, 5.1066386498122354e-14, 1.0490246824277965e-15, 1.4956003015903496e-12, 2.5734427609724886e-13, 2.1143946469237562e-06, 9.544867651811728e-08, 4.2543565892394497e-10, 6.215519418595328e-12, 1.687761909396901e-11, 1.6993320528513323e-08, 1.0583119935958507e-09, 9.857150189418462e-07, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [4.727198188447801e-08, 0.002272214274853468, 0.8730366826057434, 0.0016238681273534894, 9.849362297975617e-11, 6.310171162720105e-14, 1.3311845115798748e-12, 1.350557283785747e-07, 1.07800769910682e-05, 3.4101576602552086e-05, 7.529693561991735e-07, 3.7022258592145363e-09, 3.1551092294357375e-10, 8.851498527195911e-12, 1.024629546009237e-05, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [6.003397223786067e-10, 5.335852165444521e-06, 0.00445933174341917, 0.5796651840209961, 5.976808097329922e-05, 2.377180230439535e-09, 1.7792844021063958e-12, 1.2140626282075573e-09, 6.417224529542409e-09, 2.601910637167748e-06, 1.1842810181406094e-06, 1.8266834445057611e-07, 1.3081095096012518e-09, 1.5776791765370612e-12, 4.7676843678345904e-05, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [2.4071971206038626e-15, 2.3560551770727793e-14, 9.98394700246763e-11, 1.7167060661904543e-07, 0.2774648666381836, 1.6012703781598248e-05, 9.760837530760607e-15, 4.654387315338889e-18, 8.039692137064508e-20, 2.1508527635127157e-16, 1.789740057545064e-11, 2.4233797191186568e-08, 2.7592322870972907e-10, 4.956549239646573e-15, 1.5411848153235042e-06, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.9919477308935618e-13, 5.266535346254387e-16, 1.2917133013982517e-14, 7.221083175856791e-10, 8.195231930585578e-05, 0.5564944744110107, 4.117699063499458e-06, 5.438900198273533e-13, 2.4172004338169554e-20, 9.57835365503234e-22, 9.376302678036402e-17, 3.235451073724249e-10, 6.101883442966027e-09, 9.971044129253315e-11, 1.6162671201414014e-08, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [9.771466125130246e-08, 3.17872256294649e-11, 3.1429036890379125e-13, 5.901367481980172e-16, 4.2342058748090494e-09, 0.0012305855052545667, 0.6103256940841675, 2.2161180822877213e-05, 7.972257402844019e-12, 6.481494664823834e-19, 5.35928561114305e-19, 7.863773244772346e-14, 1.1593314752644801e-07, 8.808668212623161e-07, 1.1730364235518209e-07, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [2.6939844799400703e-10, 3.892770337188267e-07, 2.2438891023046637e-10, 2.095593632707407e-18, 1.8655412772298346e-14, 2.206185598652155e-07, 3.0316745323943906e-05, 0.33891788125038147, 5.437008439912461e-06, 1.3213468337612382e-14, 2.5347562276209975e-18, 1.0659246862729562e-18, 2.6392999114346893e-13, 9.868956762915104e-10, 1.6170986327779246e-06, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.3015508670832787e-09, 4.1474245904282725e-07, 7.619819371029735e-06, 9.079691751061325e-13, 5.725895077835787e-16, 1.0568446176517903e-14, 8.978999488373773e-11, 2.253716047562193e-05, 0.9323674440383911, 0.0001553743495605886, 1.1094852814252931e-10, 4.251380123255501e-17, 3.4548606558270072e-18, 1.563022274271835e-14, 1.7832363141678798e-07, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.2218349942916262e-10, 4.9370779464652514e-08, 1.0212672805209877e-06, 3.802215486903293e-11, 4.1323817879847246e-16, 3.8503187577578586e-16, 6.2032051316354e-15, 3.2203126920649083e-07, 8.202762546716258e-05, 0.5051153898239136, 1.6483796571264975e-05, 2.317061202194298e-13, 9.134085045449695e-19, 4.959048342554486e-21, 1.9839136555788173e-08, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [3.5615963439117673e-14, 6.311461336200308e-12, 7.572167781688677e-09, 7.864790063649707e-08, 5.871175941252194e-13, 4.399392566282849e-15, 3.6105855357745724e-20, 8.408651243829376e-14, 2.915925279012299e-09, 2.7294316168990918e-05, 0.31493836641311646, 1.4271394093157141e-06, 7.57530499374999e-14, 1.0444343699767344e-21, 5.65783730976932e-09, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.619628042792698e-10, 6.862534152052291e-11, 7.238428190170509e-10, 5.1994692995549485e-08, 8.193378420173758e-08, 6.734891755399985e-09, 1.47457238341411e-14, 5.793711288450045e-15, 1.5065480465795492e-14, 1.167909147170576e-08, 0.0003541565383784473, 0.5504465699195862, 2.5677532903500833e-05, 4.9321430864142715e-14, 1.3459792569392448e-07, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [8.003913504195381e-11, 5.626729984720136e-12, 4.9737857062137625e-12, 1.4365373474101162e-11, 1.165467935493325e-07, 3.263785401941277e-05, 9.4434834951862e-11, 2.6144878938953817e-15, 6.540743544149476e-19, 2.5930401594030658e-17, 1.8366722587259687e-09, 1.8794700736179948e-05, 0.49058014154434204, 8.066950840657228e-07, 1.3585024589701788e-06, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.0801989728040362e-12, 2.2359935084037552e-13, 1.1691597126203823e-12, 1.0214807062303036e-16, 2.4270561688882752e-12, 4.4484740890915475e-10, 1.1468358207533669e-10, 1.5131759777478604e-13, 3.7208958865722007e-20, 6.888861115537483e-21, 1.5888746801787275e-18, 3.2241334168431335e-12, 5.685043561243219e-06, 0.3912107050418854, 3.0407140694244106e-10, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [5.397048425948014e-07, 2.3629811494174646e-06, 8.614414923613367e-07, 8.006720286779512e-13, 4.92412575016192e-14, 2.066644277931573e-08, 0.00031528103863820434, 0.011093947105109692, 3.7555511767095595e-07, 1.151808547627739e-13, 5.505821095062543e-16, 1.6971218267519683e-12, 5.383023108151974e-06, 0.8731740117073059, 0.04139598086476326, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6266164779663086, 0.3128010928630829, 0.06246759742498398, 0.00042505442979745567, 0.008534153923392296, 0.09425555169582367, 0.2709643542766571, 0.686626672744751, 0.3142872750759125, 0.10107265412807465, 0.015935143455863, 0.012286541052162647, 0.14970052242279053, 0.3989029824733734, 0.022492708638310432, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.24012988805770874, 0.6692726612091064, 0.08029869198799133, 0.41845017671585083, 0.08128808438777924, 0.09738753736019135, 0.15100885927677155, 0.2691691815853119, 0.013517879880964756, 0.21848294138908386, 0.16758716106414795, 0.12734578549861908, 0.32224464416503906, 0.12471552193164825, 0.07385692000389099, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13747748732566833, 0.012865100987255573, 0.3056560158729553, 0.3759651184082031, 0.20075583457946777, 0.056869279593229294, 0.27502477169036865, 0.09038521349430084, 0.09535539150238037, 0.27579623460769653, 0.15189220011234283, 0.6071571111679077, 0.0820951759815216, 0.09481122344732285, 0.09779953956604004, 0.13988038897514343, 0.003474950324743986, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.007538634352385998, 0.02957071363925934, 0.011847163550555706, 0.055522944778203964, 0.04100131243467331, 0.031534671783447266, 0.06567902117967606, 0.09044305235147476, 0.007193693891167641, 0.06334451586008072, 0.07378207892179489, 0.07786792516708374, 0.28214019536972046, 0.08070375770330429, 0.20607011020183563, 0.14879919588565826, 0.018745053559541702, 0.07372914999723434, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.005881547927856445, 0.008371960371732712, 0.010823756456375122, 0.024797217920422554, 0.024142105132341385, 0.01083815935999155, 0.008304014801979065, 0.006388344801962376, 0.009114595130085945, 0.022048065438866615, 0.1306026130914688, 0.23451638221740723, 0.3918500244617462, 0.08784151822328568, 0.2650633752346039, 0.030327370390295982, 0.02692173607647419, 0.46947386860847473, 0.09036581218242645, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20629070699214935, 0.2529377341270447, 0.028870999813079834, 0.049127642065286636, 0.04690879210829735, 0.11594393104314804, 0.15515393018722534, 0.06585636734962463, 0.0420556403696537, 0.1996643990278244, 0.028717953711748123, 0.7190893292427063, 0.30376943945884705, 0.22654840350151062, 0.12926629185676575, 0.164228156208992, 0.0009850627975538373, 0.0044541023671627045, 0.0005622706958092749, 0.024160074070096016, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01586613617837429, 0.15566423535346985, 0.015082520432770252, 0.009204044006764889, 0.002680863719433546, 0.07106906920671463, 0.08370621502399445, 0.05749649554491043, 0.03059268370270729, 0.012942377477884293, 0.0011753733269870281, 0.00916373822838068, 0.0020018015056848526, 0.049308281391859055, 0.19197486340999603, 0.020124448463320732, 0.0011880549136549234, 0.0042731426656246185, 3.242780803702772e-05, 0.6858344078063965, 0.023040860891342163, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03849078342318535, 0.08146823942661285, 0.03517843410372734, 0.025976145640015602, 0.02364599145948887, 0.1389763057231903, 0.02619975060224533, 0.034312427043914795, 0.02985706366598606, 0.029806064441800117, 0.00684476038441062, 0.03280223533511162, 0.030126189813017845, 0.10321015119552612, 0.23163792490959167, 0.0017230550292879343, 3.356653905939311e-05, 0.001307086437009275, 1.4968540199333802e-05, 0.5564903616905212, 0.236929789185524, 0.007688341196626425, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2772977352142334, 0.05161405727267265, 0.04358568787574768, 0.047931231558322906, 0.04583681374788284, 0.08128579705953598, 0.15782645344734192, 0.0856042429804802, 0.10767779499292374, 0.11355230212211609, 0.041377030313014984, 0.252811074256897, 0.05780917406082153, 0.19973745942115784, 0.22427907586097717, 0.1612924486398697, 0.00029754414572380483, 0.0029063820838928223, 0.0015110797248780727, 0.16695675253868103, 0.3453270196914673, 0.07193248718976974, 0.006359610706567764, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.023119861260056496, 0.02037731558084488, 0.0453791618347168, 0.1060030460357666, 0.006244942545890808, 0.0085020512342453, 0.012060720473527908, 0.014560479670763016, 0.00689319521188736, 0.011241135187447071, 0.023835573345422745, 0.02693312056362629, 0.011436404660344124, 0.019489392638206482, 0.30997538566589355, 0.1910298615694046, 0.01051796693354845, 0.0018660163041204214, 0.0012154864380136132, 0.022663934156298637, 0.008557457476854324, 0.016767704859375954, 0.05246622860431671, 0.08816055208444595, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.045414164662361145, 0.005229660775512457, 0.011418518610298634, 0.009312640875577927, 0.0002147085906472057, 0.12653864920139313, 0.05854451283812523, 0.11896014213562012, 0.0156405046582222, 0.010270207189023495, 0.0032450463622808456, 0.015787174925208092, 0.011106730438768864, 0.007675709668546915, 0.3779195249080658, 0.24295811355113983, 0.0012021175352856517, 0.0005200211890041828, 0.00015996988804545254, 0.002627951791509986, 0.03450923040509224, 0.014827161096036434, 0.015967652201652527, 0.005632439162582159, 0.001854590023867786, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.007367350626736879, 0.012884993106126785, 0.01019106525927782, 0.011957473121583462, 0.054886650294065475, 0.09750530868768692, 0.029414953663945198, 0.08492925018072128, 0.17440666258335114, 0.003643231000751257, 0.00105402956251055, 0.02280060388147831, 0.0010922637302428484, 0.005130939185619354, 0.09500079602003098, 0.2492469847202301, 0.004325273912400007, 0.004784590099006891, 0.013903478160500526, 0.0013026667293161154, 0.003877879586070776, 0.017029188573360443, 0.01781909167766571, 0.05003270506858826, 0.026610376313328743, 0.008462576195597649, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02996714971959591, 0.028387926518917084, 0.16122521460056305, 0.0898616760969162, 0.06381779164075851, 0.20551051199436188, 0.13175098598003387, 0.562389075756073, 0.04834860563278198, 0.013581722043454647, 0.03991095721721649, 0.10736902058124542, 0.03830268979072571, 0.05736052244901657, 0.27213579416275024, 0.25306010246276855, 0.0017952719936147332, 0.005404005758464336, 0.021692873910069466, 0.0005702165653929114, 9.544018394080922e-05, 0.001603480544872582, 0.001225438085384667, 0.036846794188022614, 0.001749897957779467, 0.016878794878721237, 0.021703237667679787, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03571658954024315, 0.012061648070812225, 0.08574458211660385, 0.022463832050561905, 0.12578466534614563, 0.07826194912195206, 0.06577891856431961, 0.13274507224559784, 0.06591502577066422, 0.05002211779356003, 0.03129255399107933, 0.27911075949668884, 0.31601372361183167, 0.10930214822292328, 0.30993908643722534, 0.055758021771907806, 0.000425096252001822, 0.0005783061496913433, 0.0011671994579955935, 0.00034630659501999617, 0.00031045774812810123, 0.0006358043756335974, 0.004018810577690601, 0.0004720573779195547, 0.006387148518115282, 0.038948215544223785, 0.40798652172088623, 0.0038703898899257183, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04630875587463379, 0.03141915798187256, 0.03061339072883129, 0.007028677500784397, 0.008451082743704319, 0.02540888637304306, 0.012118873186409473, 0.09331455826759338, 0.0033372503239661455, 0.01357665192335844, 0.0069510783068835735, 0.017483821138739586, 0.033454760909080505, 0.014270796440541744, 0.44127020239830017, 0.29551389813423157, 0.006183725781738758, 0.0010477532632648945, 0.001470124931074679, 0.0028535614255815744, 0.003910644445568323, 0.004942604340612888, 0.003798475954681635, 0.01567114144563675, 0.060374900698661804, 0.006600319407880306, 0.010896215215325356, 0.009779008105397224, 0.007320093456655741, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1722828894853592, 0.15122008323669434, 0.056102070957422256, 0.09136570990085602, 0.02421834133565426, 0.045343294739723206, 0.034619707614183426, 0.030837759375572205, 0.019798463210463524, 0.04411705583333969, 0.05331422761082649, 0.09423463046550751, 0.1436629444360733, 0.13433872163295746, 0.1229754090309143, 0.1632017195224762, 0.00519327400252223, 0.00790441408753395, 0.0009941658936440945, 0.3241596221923828, 0.0008480648975819349, 0.0001429034018656239, 0.0012253100285306573, 0.0008457236108370125, 0.006411578040570021, 0.0016067628748714924, 0.003762597683817148, 0.029224932193756104, 0.07677540183067322, 0.06338826566934586, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.022473091259598732, 0.0489150770008564, 0.010993139818310738, 0.03897916153073311, 0.003662768052890897, 0.002051829593256116, 0.0037445707712322474, 0.016557298600673676, 0.014907213859260082, 0.004300208762288094, 0.004852794576436281, 0.0027131394017487764, 0.016001524403691292, 0.008091894909739494, 0.25544992089271545, 0.005401996895670891, 6.3005199990584515e-06, 0.0004310416697990149, 8.47076989884954e-06, 0.009243682958185673, 0.0008590375073254108, 4.37394373875577e-06, 6.523932825075462e-05, 8.531090134056285e-05, 0.0006816720124334097, 7.644478318979964e-05, 0.00018924157484434545, 0.0012375408550724387, 0.023784970864653587, 0.4309314787387848, 0.034907225519418716, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08012817800045013, 0.2898695766925812, 0.022246699780225754, 0.06057273969054222, 0.025327028706669807, 0.02957070618867874, 0.04002644121646881, 0.019245512783527374, 0.01995179057121277, 0.020330116152763367, 0.006697094067931175, 0.015452835708856583, 0.014569609425961971, 0.04013357311487198, 0.2585589587688446, 0.29775136709213257, 0.006892140489071608, 0.009814155288040638, 0.016249310225248337, 0.004830268211662769, 0.0035455955658107996, 0.0007549467263743281, 0.000541276705916971, 0.0031480982434004545, 0.001557780895382166, 0.0010192448971793056, 0.0018504501786082983, 0.002619183622300625, 0.1016833484172821, 0.03818811476230621, 0.06928347051143646, 0.0412699431180954, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01832924410700798, 0.023918962106108665, 0.024782713502645493, 0.033514510840177536, 0.050549402832984924, 0.013098560273647308, 0.023091215640306473, 0.030541786924004555, 0.1064886748790741, 0.006106832530349493, 0.0024854408111423254, 0.018918434157967567, 0.0075035663321614265, 0.009370497427880764, 0.21452490985393524, 0.26683223247528076, 0.0017643374158069491, 0.02531762421131134, 0.047485485672950745, 0.0005023732082918286, 0.0011795219033956528, 0.002227108459919691, 0.0028741960413753986, 0.005215880926698446, 0.001946018310263753, 3.592624852899462e-05, 0.001338632428087294, 0.0025214410852640867, 0.07723907381296158, 0.012742026709020138, 0.25196006894111633, 0.052669085562229156, 0.020061112940311432, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.027254067361354828, 0.020437292754650116, 0.14233240485191345, 0.08538791537284851, 0.03242940828204155, 0.0897425189614296, 0.08476056158542633, 0.2620556950569153, 0.02126460149884224, 0.023079702630639076, 0.03143052011728287, 0.04489685967564583, 0.046720463782548904, 0.03604652360081673, 0.23038896918296814, 0.3006725609302521, 0.0014043879928067327, 0.009936605580151081, 0.037061650305986404, 0.0005129858036525548, 5.274279828881845e-05, 0.0006371501949615777, 0.00048446646542288363, 0.015043019317090511, 0.0003374778898432851, 0.0015171451959758997, 0.001911269617266953, 0.0014702629996463656, 0.015123972669243813, 0.0006335150101222098, 0.0006853189552202821, 0.0006114236894063652, 0.013829384930431843, 0.010252222418785095, NaN, NaN, NaN, NaN, NaN, NaN], [0.042377930134534836, 0.017293933779001236, 0.08730384707450867, 0.030179454013705254, 0.12187745422124863, 0.05139933153986931, 0.047754548490047455, 0.066692054271698, 0.06521614640951157, 0.05196157470345497, 0.028108397498726845, 0.17703385651111603, 0.22747749090194702, 0.06955988705158234, 0.28824013471603394, 0.11150761693716049, 0.0006332705961540341, 0.0012255925685167313, 0.0022868558298796415, 0.0007688697660341859, 0.00046408100752159953, 0.0006869957433082163, 0.0021696356125175953, 0.0003113164857495576, 0.0013619231758639216, 0.004312699660658836, 0.1263500303030014, 0.0001710234791971743, 0.0024227115791291, 0.0006429344066418707, 0.008991677314043045, 0.01230061985552311, 0.025017380714416504, 0.33947470784187317, 0.0032216052059084177, NaN, NaN, NaN, NaN, NaN], [0.03372317552566528, 0.030876630917191505, 0.025082340463995934, 0.008588657714426517, 0.007454049773514271, 0.009771045297384262, 0.010381288826465607, 0.041183773428201675, 0.004549690056592226, 0.01619204692542553, 0.0060179769061505795, 0.009672058746218681, 0.022905999794602394, 0.009750566445291042, 0.30946746468544006, 0.31111404299736023, 0.0035644923336803913, 0.0013678895775228739, 0.0016790243098512292, 0.0035299588926136494, 0.004438228905200958, 0.004504224751144648, 0.0015486004995182157, 0.006104794796556234, 0.009403211995959282, 0.00038756802678108215, 0.001732571516185999, 0.00042684219079092145, 0.00029873420135118067, 0.02043243870139122, 0.02443091571331024, 0.011036018840968609, 0.0030384601559489965, 0.007405058480799198, 0.004648045636713505, 0.010011163540184498, NaN, NaN, NaN, NaN], [0.18900562822818756, 0.14908763766288757, 0.05840699374675751, 0.10216160118579865, 0.03072887472808361, 0.04109037667512894, 0.03799780085682869, 0.02909342385828495, 0.03500371053814888, 0.0757574513554573, 0.061073921620845795, 0.09956928342580795, 0.10441071540117264, 0.14136889576911926, 0.13095542788505554, 0.16896948218345642, 0.0033956619445234537, 0.009647470898926258, 0.0011160745052620769, 0.30864211916923523, 0.0008666384965181351, 0.0001862353819888085, 0.0007671809289604425, 0.0006719603552483022, 0.002030742121860385, 0.00038655498065054417, 0.0009093419066630304, 0.0015865613240748644, 0.007534818258136511, 0.009185722097754478, 0.00011195908882655203, 0.003075815038755536, 0.000886340974830091, 0.0034873690456151962, 0.021776562556624413, 0.11334169656038284, 0.0832705944776535, NaN, NaN, NaN], [0.014150185510516167, 0.03789284825325012, 0.007744992151856422, 0.02556411363184452, 0.0037681234534829855, 0.001123085618019104, 0.002939486177638173, 0.010072565637528896, 0.019109029322862625, 0.003645692951977253, 0.0027771664317697287, 0.002490789396688342, 0.007166225463151932, 0.005180294159799814, 0.2058444321155548, 0.006588279269635677, 7.165617716964334e-06, 0.0005450915195979178, 1.0953889614029322e-05, 0.01959507167339325, 0.001590097788721323, 1.1096496564277913e-05, 7.439414184773341e-05, 9.72584675764665e-05, 0.00039174238918349147, 2.7912905352422968e-05, 4.964227991877124e-05, 7.256279786815867e-05, 0.00222678086720407, 0.04727102443575859, 0.0002576226834207773, 0.00020273383415769786, 7.391278631985188e-05, 0.00018598776659928262, 0.000617648009210825, 0.03195251524448395, 0.45461374521255493, 0.037591490894556046, NaN, NaN], [0.0469474196434021, 0.1743137687444687, 0.021908296272158623, 0.046387769281864166, 0.02985612489283085, 0.019742406904697418, 0.040140021592378616, 0.01437240932136774, 0.02856219932436943, 0.018488112837076187, 0.004136314615607262, 0.01038376335054636, 0.009851893410086632, 0.026245350018143654, 0.22488054633140564, 0.35417911410331726, 0.010997277684509754, 0.014662563800811768, 0.023722819983959198, 0.01071385107934475, 0.009427045471966267, 0.002653747797012329, 0.0011037624208256602, 0.005973298568278551, 0.0016420705942437053, 0.0009447215707041323, 0.001327668083831668, 0.0005524749867618084, 0.012130306102335453, 0.005379356909543276, 0.0037436189595609903, 0.0009285339619964361, 0.0002853046462405473, 0.0013114019529893994, 0.0012977200094610453, 0.08090774714946747, 0.034737478941679, 0.058711227029561996, 0.0672648623585701, NaN], [0.00832295510917902, 0.021339448168873787, 0.00394090311601758, 0.002333499025553465, 0.05547437444329262, 0.007243151310831308, 0.011641105636954308, 0.0331541933119297, 0.010278979316353798, 0.011881710961461067, 0.001766148954629898, 0.04899042472243309, 0.01878243498504162, 0.01244808267802, 0.15685127675533295, 0.18188641965389252, 0.00040442554745823145, 0.0015771333128213882, 0.005189571529626846, 8.387575689994264e-06, 0.0001226859458256513, 0.0011242604814469814, 0.0013583728577941656, 0.0030172227416187525, 0.00029841059586033225, 1.2829146726289764e-05, 0.001467264024540782, 0.001090237987227738, 0.002914785873144865, 0.0006871690275147557, 0.002592542441561818, 0.00021328746515791863, 6.871169898658991e-05, 0.002350796014070511, 0.0026233955286443233, 0.02620280720293522, 0.005966363474726677, 0.08270465582609177, 0.010547555983066559, 0.018362630158662796]]], [[[0.1393769532442093, 0.0735321119427681, 0.701509952545166, 0.10650816559791565, 0.05110495164990425, 0.021589145064353943, 0.0033319133799523115, 0.0014166238252073526, 0.01486207265406847, 0.006584684830158949, 0.002582702785730362, 0.0004108685825485736, 0.010701421648263931, 0.009390643797814846, 0.06290604919195175, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0030957262497395277, 0.0237117987126112, 0.7945073246955872, 0.09792613238096237, 0.2614360749721527, 0.179405078291893, 0.011310527101159096, 0.009954328648746014, 0.009489532560110092, 0.0005609119543805718, 0.000751268700696528, 0.0001462608779547736, 0.004604416899383068, 0.004964352585375309, 0.019775664433836937, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.002461136318743229, 0.024594180285930634, 0.009559455327689648, 0.055053047835826874, 0.30010533332824707, 0.4690517783164978, 0.03334644436836243, 0.0075769852846860886, 0.007821744307875633, 0.004109389614313841, 0.0022267017047852278, 0.000916018383577466, 0.0037954216822981834, 0.0007741246954537928, 0.004415341652929783, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0019876149017363787, 0.0012237336486577988, 0.00015556006110273302, 0.0003553472051862627, 0.4419420659542084, 0.6252713799476624, 0.02062046155333519, 0.0028509902767837048, 0.00548406969755888, 0.0003452444798313081, 0.0001962203241419047, 0.0008938669925555587, 0.0009214308229275048, 1.2216354662086815e-05, 0.0019377138232812285, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00020824302919209003, 0.00021322975226212293, 4.6913473852328025e-06, 0.00017657040734775364, 0.0005752452998422086, 0.5289100408554077, 0.1970362812280655, 0.12947966158390045, 0.0005265067447908223, 0.000227929005632177, 6.233566091395915e-05, 0.0001991882745642215, 0.00032238851417787373, 0.0003627484547905624, 0.0016414258861914277, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0010278578847646713, 0.0029486939311027527, 0.00014835220645181835, 0.00036925319000147283, 0.00742883887141943, 0.03272741660475731, 0.8576475977897644, 0.03500620648264885, 0.2982224225997925, 0.0003585784579627216, 5.663683623424731e-05, 0.0011889662127941847, 0.00576341338455677, 0.003998933359980583, 0.03130826726555824, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.002113666385412216, 0.004151111003011465, 0.002428078791126609, 0.002119476906955242, 0.001100956811569631, 0.003687644377350807, 0.13543397188186646, 0.11922256648540497, 0.7567945718765259, 0.2570010721683502, 0.004903816152364016, 0.0001005519661703147, 0.000830159813631326, 0.001259618904441595, 0.14076685905456543, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0010344160255044699, 0.00660368800163269, 0.0025270660407841206, 0.00023567670723423362, 0.0004021638887934387, 0.0030120171140879393, 0.0016376315616071224, 0.0524386465549469, 0.7797302007675171, 0.1269131302833557, 0.004214781802147627, 0.0002750723797362298, 0.002267329953610897, 0.001067862962372601, 0.16698867082595825, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0009750229655764997, 0.0120720649138093, 0.0038384809158742428, 0.0036232813727110624, 0.004431525245308876, 0.0007613649941049516, 5.662842158926651e-05, 0.01338160876184702, 0.041878536343574524, 0.7091978788375854, 0.2535402476787567, 0.13969287276268005, 0.026510832831263542, 0.0006678565987385809, 0.015569130890071392, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0002093962684739381, 0.00030164673808030784, 0.00010105424007633701, 5.030819465901004e-06, 0.001411793869920075, 0.003664590884000063, 0.00017403968377038836, 0.0011218853760510683, 0.011106000281870365, 0.003924186807125807, 0.07315385341644287, 0.3008219599723816, 0.36353737115859985, 0.025737306103110313, 0.0060785748064517975, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0001716838014544919, 0.0008840822265483439, 4.3183892557863146e-05, 3.6494086543825688e-06, 0.0005770743009634316, 0.010045445524156094, 0.00010205945727648214, 6.57988857710734e-05, 0.0006949909729883075, 0.004452799912542105, 0.009000658988952637, 0.49080607295036316, 0.17717383801937103, 0.11174798011779785, 0.021669577807188034, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.019416164606809616, 0.0014941463014110923, 0.001027028076350689, 0.001502541359513998, 0.0085412273183465, 0.12493651360273361, 0.0035243057645857334, 0.0026196581311523914, 0.0008317703031934798, 0.0015569254755973816, 0.060888972133398056, 0.06929422169923782, 0.3396435081958771, 0.387500524520874, 0.017253199592232704, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04994890093803406, 0.15025374293327332, 0.024391163140535355, 0.00227133696898818, 0.012616162188351154, 0.2894521951675415, 0.4185648262500763, 0.19089959561824799, 0.027421748265624046, 0.001001756638288498, 0.0036985764745622873, 0.06802930682897568, 0.02484762854874134, 0.057649459689855576, 0.1606004238128662, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03736208751797676, 0.11793919652700424, 0.0180205088108778, 0.0001436693564755842, 0.0030756669584661722, 0.08228655159473419, 0.12110688537359238, 0.09650447964668274, 0.015347721055150032, 0.0004259537090547383, 0.00022625335259363055, 0.001013986300677061, 0.0784289613366127, 0.2240448147058487, 0.18707746267318726, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7529165148735046, 0.7075774073600769, 0.6068683862686157, 0.3852986991405487, 0.6197313666343689, 0.6735447645187378, 0.6598724722862244, 0.7226093411445618, 0.31395286321640015, 0.2518909275531769, 0.07010441273450851, 0.21793116629123688, 0.4325476884841919, 0.7029338479042053, 0.06848814338445663, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04773104563355446, 0.01963546872138977, 0.16452182829380035, 0.04063690826296806, 0.1849776655435562, 0.08088860660791397, 0.11659693717956543, 0.038044340908527374, 0.2744975686073303, 0.003083554795011878, 0.019721103832125664, 0.08137688785791397, 0.0169991385191679, 0.03939461708068848, 0.14168404042720795, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09676018357276917, 0.018249453976750374, 0.657112717628479, 0.5890088677406311, 0.5712416768074036, 0.2744671702384949, 0.48642322421073914, 0.26345524191856384, 0.23708243668079376, 0.03475205600261688, 0.15204745531082153, 0.0676480308175087, 0.050043635070323944, 0.0665324404835701, 0.036993421614170074, 0.13007116317749023, 0.035988736897706985, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04065309092402458, 0.0025235058274120092, 0.11838234961032867, 0.27863210439682007, 0.37560757994651794, 0.7046668529510498, 0.12516380846500397, 0.1912177950143814, 0.14992743730545044, 0.05949303135275841, 0.056387268006801605, 0.04353337734937668, 0.17471297085285187, 0.07017815858125687, 0.12025584280490875, 0.17991511523723602, 0.05124381557106972, 0.013642107136547565, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.015422305092215538, 0.000844803755171597, 0.015767300501465797, 0.11098357290029526, 0.273564875125885, 0.3235251009464264, 0.14805495738983154, 0.17132841050624847, 0.25568780303001404, 0.034506767988204956, 0.046862825751304626, 0.03818853572010994, 0.025031423196196556, 0.027911247685551643, 0.009120252914726734, 0.16831281781196594, 0.043814778327941895, 0.0950295478105545, 0.07350433617830276, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01866327039897442, 0.11290711164474487, 0.007440958172082901, 0.031009642407298088, 0.059622399508953094, 0.035299621522426605, 0.012064317241311073, 0.17540854215621948, 0.06399405747652054, 0.010346408933401108, 0.023967623710632324, 0.006549614481627941, 0.015476463362574577, 0.017944032326340675, 0.15624091029167175, 0.13759823143482208, 0.14112484455108643, 0.20577600598335266, 0.13910864293575287, 0.034107428044080734, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.115133136510849, 0.5564319491386414, 0.0024013265501707792, 0.014839398674666882, 0.027623601257801056, 0.003712957026436925, 0.11139625310897827, 0.4320802688598633, 0.18111301958560944, 0.025198934599757195, 0.05914938822388649, 0.029404014348983765, 0.1131783202290535, 0.1630096137523651, 0.14384765923023224, 0.11619941890239716, 0.038306448608636856, 0.06045802682638168, 0.03494013100862503, 0.374624639749527, 0.22046393156051636, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.047323077917099, 0.01987922191619873, 0.021367410197854042, 0.0816798061132431, 0.11104802042245865, 0.01310664601624012, 0.37855657935142517, 0.16697411239147186, 0.31461480259895325, 0.04616151005029678, 0.27547621726989746, 0.04939346760511398, 0.02232075110077858, 0.15515512228012085, 0.01579722762107849, 0.08332619816064835, 0.009484739042818546, 0.012810231186449528, 0.0027760458178818226, 0.3268325924873352, 0.26342087984085083, 0.17634892463684082, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13229456543922424, 0.031869739294052124, 0.26943540573120117, 0.2586674690246582, 0.3796730637550354, 0.127562016248703, 0.20277942717075348, 0.05910756066441536, 0.14354895055294037, 0.08293455094099045, 0.2214740365743637, 0.23150987923145294, 0.18035069108009338, 0.2860051393508911, 0.07895194739103317, 0.057563915848731995, 0.01992173306643963, 0.03713805601000786, 0.014863312244415283, 0.25726908445358276, 0.14832180738449097, 0.402090460062027, 0.06479739397764206, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09224988520145416, 0.07457923144102097, 0.05282874405384064, 0.09438028931617737, 0.06849074363708496, 0.012997711077332497, 0.007214613724499941, 0.004257954657077789, 0.2309093326330185, 0.38276976346969604, 0.5917518734931946, 0.7830951809883118, 0.8438952565193176, 0.7586230039596558, 0.04145537316799164, 0.21478669345378876, 0.15359601378440857, 0.26770198345184326, 0.12653663754463196, 0.09151764959096909, 0.07003500312566757, 0.19363711774349213, 0.014233908616006374, 0.023967349901795387, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.014161140657961369, 0.027171263471245766, 0.0029068312142044306, 0.020549731329083443, 0.0005743438960053027, 0.00417140731588006, 0.003657599212601781, 0.00956815481185913, 0.34446486830711365, 0.5171273946762085, 0.39057764410972595, 0.2845093309879303, 0.1669711321592331, 0.5306525230407715, 0.015455210581421852, 0.2834857702255249, 0.07559704780578613, 0.07655511796474457, 0.16202391684055328, 0.08316012471914291, 0.11911017447710037, 0.0204884335398674, 0.011816238984465599, 0.13204774260520935, 0.039266277104616165, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02566671371459961, 0.00907080341130495, 0.0006065603229217231, 0.03001752682030201, 0.00023783017240930349, 0.0005533608491532505, 0.013808660209178925, 0.003767948364838958, 0.06461481004953384, 0.1359771490097046, 0.08153439313173294, 0.572087287902832, 0.36045318841934204, 0.44234389066696167, 0.0030113777611404657, 0.23006244003772736, 0.03933367133140564, 0.07187695801258087, 0.04476522281765938, 0.01073860377073288, 0.0032203071750700474, 0.00176758982706815, 0.018770985305309296, 0.12121162563562393, 0.18536020815372467, 0.01582610420882702, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03087739646434784, 0.012099061161279678, 0.004942088853567839, 0.038267359137535095, 0.0023591304197907448, 0.0037323227152228355, 0.04966888204216957, 0.012427400797605515, 0.16158415377140045, 0.020882699638605118, 0.05600592866539955, 0.367767333984375, 0.24262923002243042, 0.38281354308128357, 0.00973587203770876, 0.18067117035388947, 0.009833509102463722, 0.03744787722826004, 0.016920698806643486, 0.05744745582342148, 0.04540643468499184, 0.008024180307984352, 0.012110988609492779, 0.09370782226324081, 0.08820194005966187, 0.06259123980998993, 0.025030089542269707, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04249054566025734, 0.0069285486824810505, 0.006088858004659414, 0.044397544115781784, 0.05390672758221626, 0.006144464481621981, 0.018320903182029724, 0.01545354351401329, 0.05193139612674713, 0.03221629932522774, 0.02379259280860424, 0.27246853709220886, 0.22103002667427063, 0.23179520666599274, 0.005589436274021864, 0.11523616313934326, 0.03200709819793701, 0.050564926117658615, 0.010618647560477257, 0.09430865943431854, 0.018685024231672287, 0.022438397631049156, 0.017720744013786316, 0.1592920571565628, 0.21717989444732666, 0.2463550567626953, 0.2194516956806183, 0.0009421245777048171, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04184036701917648, 0.03700190782546997, 0.008264865726232529, 0.02439146116375923, 0.00799429602921009, 0.12502151727676392, 0.05032283812761307, 0.18101848661899567, 0.07329469919204712, 0.08409427851438522, 0.10790428519248962, 0.011960207484662533, 0.20496119558811188, 0.19276422262191772, 0.0069670299999415874, 0.09747911244630814, 0.1645127683877945, 0.1875433474779129, 0.09478750824928284, 0.08721300214529037, 0.02294742316007614, 0.02039182186126709, 0.07351931929588318, 0.1815827339887619, 0.5564144849777222, 0.41975197196006775, 0.2698606848716736, 0.05650324374437332, 0.05821085348725319, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06364590674638748, 0.06483624875545502, 0.015260975807905197, 0.1278582364320755, 0.006228389218449593, 0.02756887674331665, 0.020600903779268265, 0.015440343879163265, 0.018087223172187805, 0.017098410055041313, 0.025406692177057266, 0.0007098353235051036, 0.00014885497512295842, 0.0013503700029104948, 0.15608660876750946, 0.14833268523216248, 0.1209164559841156, 0.08990822732448578, 0.0656033307313919, 0.23720099031925201, 0.11782333254814148, 0.04633651673793793, 0.16808320581912994, 0.06126163899898529, 0.43528908491134644, 0.3754012882709503, 0.13757933676242828, 0.05596579611301422, 0.16984672844409943, 0.002737722359597683, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.6220619678497314, 0.6306124329566956, 0.6737340092658997, 0.49940165877342224, 0.1517823040485382, 0.8503586649894714, 0.705633282661438, 0.6629571914672852, 0.11157920956611633, 0.39899003505706787, 0.3173867464065552, 0.027327625080943108, 0.014980590902268887, 0.009274562820792198, 0.08523338288068771, 0.19258342683315277, 0.05838138237595558, 0.04652376100420952, 0.017318567261099815, 0.23482391238212585, 0.16333334147930145, 0.02100907638669014, 0.048424359411001205, 0.06841404736042023, 0.3133482038974762, 0.07921069860458374, 0.021035969257354736, 0.03291412815451622, 0.18175286054611206, 0.1566929817199707, 0.053215935826301575, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15005189180374146, 0.04609784111380577, 0.17501141130924225, 0.21113994717597961, 0.26919078826904297, 0.6422000527381897, 0.7493206858634949, 0.2162598967552185, 0.010351919569075108, 0.09728528559207916, 0.09688232094049454, 0.028558582067489624, 0.10305432975292206, 0.05914681404829025, 0.11260810494422913, 0.17641158401966095, 0.15294750034809113, 0.15352487564086914, 0.10843643546104431, 0.08260629326105118, 0.016529222950339317, 0.012650150805711746, 0.07893627882003784, 0.1388573795557022, 0.19094663858413696, 0.03751035034656525, 0.05650494620203972, 0.2426995038986206, 0.16961677372455597, 0.07263431698083878, 0.152814581990242, 0.018521834164857864, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09041088819503784, 0.052050016820430756, 0.08856991678476334, 0.2977358102798462, 0.04025371000170708, 0.3506464660167694, 0.6434463858604431, 0.25059518218040466, 0.01933867670595646, 0.04819375276565552, 0.07508239895105362, 0.04970608279109001, 0.02890131063759327, 0.02355407178401947, 0.12558245658874512, 0.25574439764022827, 0.04364950954914093, 0.05707173049449921, 0.02453112043440342, 0.016254547983407974, 0.0026636396069079638, 0.0035282839089632034, 0.015699811279773712, 0.03404982015490532, 0.04375504329800606, 0.001423283712938428, 0.05359426140785217, 0.1740386039018631, 0.10691730678081512, 0.03620539605617523, 0.04950953647494316, 0.022295303642749786, 0.025807255879044533, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18765486776828766, 0.021713200956583023, 0.21844394505023956, 0.3042432367801666, 0.17823228240013123, 0.1673380434513092, 0.8088975548744202, 0.46762967109680176, 0.05706785246729851, 0.009645337238907814, 0.0322297103703022, 0.09777479618787766, 0.08048812299966812, 0.10106904059648514, 0.17228879034519196, 0.216966450214386, 0.016096990555524826, 0.08351551741361618, 0.02645382098853588, 0.05811392888426781, 0.04091750830411911, 0.014506897889077663, 0.015038754791021347, 0.07221462577581406, 0.08585365861654282, 0.059816163033246994, 0.04502185434103012, 0.00397779606282711, 0.041175276041030884, 0.04448581859469414, 0.10983181744813919, 0.01911303587257862, 0.07987141609191895, 0.062483180314302444, NaN, NaN, NaN, NaN, NaN, NaN], [0.4792143702507019, 0.09839366376399994, 0.1882246881723404, 0.4093988239765167, 0.7147246599197388, 0.24897223711013794, 0.4705742597579956, 0.4205995500087738, 0.01958448253571987, 0.026842152699828148, 0.02239188365638256, 0.15106931328773499, 0.08969185501337051, 0.10003618896007538, 0.1635625958442688, 0.11257521063089371, 0.027663733810186386, 0.023284420371055603, 0.0038690094370394945, 0.053685132414102554, 0.008445030078291893, 0.014706910587847233, 0.009755544364452362, 0.06406830251216888, 0.10475295782089233, 0.08554040640592575, 0.16072620451450348, 0.00029980239924043417, 0.03509804978966713, 0.03031017631292343, 0.04435117170214653, 0.06420817226171494, 0.2780051827430725, 0.2271702140569687, 0.0013584558619186282, NaN, NaN, NaN, NaN, NaN], [0.40625429153442383, 0.3796224594116211, 0.2515096962451935, 0.36165565252304077, 0.24774380028247833, 0.8824228644371033, 0.8048573136329651, 0.857955813407898, 0.058371078222990036, 0.07109472155570984, 0.11402199417352676, 0.0021524245385080576, 0.019929109141230583, 0.030590593814849854, 0.11712031066417694, 0.10895614326000214, 0.15509657561779022, 0.19682957231998444, 0.07681374996900558, 0.06229116767644882, 0.016663551330566406, 0.015513443388044834, 0.04232686012983322, 0.0986364334821701, 0.35070890188217163, 0.19941051304340363, 0.163076713681221, 0.026361489668488503, 0.018140846863389015, 0.016411108896136284, 0.03203867748379707, 0.053678009659051895, 0.19773079454898834, 0.3572796881198883, 0.059515852481126785, 0.04298213869333267, NaN, NaN, NaN, NaN], [0.04390633478760719, 0.032843075692653656, 0.010515165515244007, 0.11869800090789795, 0.005461697466671467, 0.023131608963012695, 0.01705162413418293, 0.008547519333660603, 0.003713170997798443, 0.008410640992224216, 0.009457322768867016, 0.00015943740436341614, 3.361727431183681e-05, 0.0002994383394252509, 0.1532706469297409, 0.15568822622299194, 0.11876019835472107, 0.09203660488128662, 0.059780094772577286, 0.24089980125427246, 0.06525673717260361, 0.029934749007225037, 0.11168782413005829, 0.03211824223399162, 0.30118685960769653, 0.22822384536266327, 0.08190999180078506, 0.018841415643692017, 0.1366286426782608, 0.0017427116399630904, 0.02601366490125656, 0.09386949241161346, 0.19522085785865784, 0.1546826809644699, 0.06491755694150925, 0.19679579138755798, 0.0025137634947896004, NaN, NaN, NaN], [0.6348351836204529, 0.5127235651016235, 0.5931673645973206, 0.5543242692947388, 0.12377271056175232, 0.8264753222465515, 0.6941898465156555, 0.5687963962554932, 0.03150533139705658, 0.12843358516693115, 0.11884576827287674, 0.005231617949903011, 0.0018767286092042923, 0.0011644444894045591, 0.11210005730390549, 0.26271528005599976, 0.07045364379882812, 0.0520184300839901, 0.023400958627462387, 0.11433269083499908, 0.07895253598690033, 0.012276851572096348, 0.023823700845241547, 0.04200353845953941, 0.16687022149562836, 0.05654531344771385, 0.038080912083387375, 0.012698299251496792, 0.10473722219467163, 0.0643644630908966, 0.015445034019649029, 0.014234953559935093, 0.06144930049777031, 0.05821693688631058, 0.0568128302693367, 0.1767931431531906, 0.1402994990348816, 0.07714083790779114, NaN, NaN], [0.10790421068668365, 0.016916295513510704, 0.09771728515625, 0.22749783098697662, 0.26325535774230957, 0.49138790369033813, 0.6275916695594788, 0.08931886404752731, 0.0033968419302254915, 0.024402111768722534, 0.018104346469044685, 0.003288157982751727, 0.010537534020841122, 0.006979967001825571, 0.12102893739938736, 0.1969611942768097, 0.16093717515468597, 0.1609625220298767, 0.11138524115085602, 0.026131147518754005, 0.00619129091501236, 0.005407778546214104, 0.04104578495025635, 0.06517186760902405, 0.06833471357822418, 0.020616043359041214, 0.03467438742518425, 0.095084547996521, 0.06247802451252937, 0.022057469934225082, 0.06569864600896835, 0.0052108620293438435, 0.03032413311302662, 0.0838729590177536, 0.3427644968032837, 0.19215865433216095, 0.08116735517978668, 0.14785417914390564, 0.015012684278190136, NaN], [0.028179557994008064, 0.011468129232525826, 0.016789404675364494, 0.00803140178322792, 0.00952040497213602, 0.02960360422730446, 0.24957160651683807, 0.03544437885284424, 0.005487674381583929, 0.0028927521780133247, 0.005656986031681299, 0.0040698484517633915, 0.04730471968650818, 0.0667993351817131, 0.1372966766357422, 0.1272672563791275, 0.008308093063533306, 0.030398543924093246, 0.02721896767616272, 0.016537277027964592, 0.021588556468486786, 0.002818688517436385, 0.010970782488584518, 0.01434051152318716, 0.012293173000216484, 0.04184769093990326, 0.03683166950941086, 0.023453323170542717, 0.020430248230695724, 0.03333409130573273, 0.068024642765522, 0.02648366242647171, 0.1640448421239853, 0.109919473528862, 0.1576652079820633, 0.14138163626194, 0.16884489357471466, 0.30372628569602966, 0.2283693552017212, 0.17022481560707092]], [[0.0006553527782671154, 0.5631614327430725, 0.0008777088369242847, 0.00020331511041149497, 0.0014234310947358608, 0.013944034464657307, 9.958680493582506e-06, 0.01898920349776745, 0.00014103656576480716, 1.4779416233068332e-06, 1.1701366275929104e-07, 1.195983372781484e-06, 0.00012817273091059178, 3.365538941579871e-05, 0.00028557839686982334, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00638999929651618, 0.7093943953514099, 0.004974186420440674, 0.06159398332238197, 0.003979360219091177, 0.06536109745502472, 0.005324128083884716, 0.02885170467197895, 0.0003847253101412207, 0.0002721542550716549, 4.3882369936909527e-05, 0.00024302180099766701, 0.00612376956269145, 0.006710950285196304, 0.0343138724565506, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.109707772731781, 0.1680740863084793, 0.05170662701129913, 0.04158816486597061, 0.026700180023908615, 0.23248757421970367, 0.5156019330024719, 0.3799504041671753, 0.02909121848642826, 0.009008231572806835, 0.0013055672170594335, 0.0032788640819489956, 0.0791734829545021, 0.010587821714580059, 0.06850002706050873, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04004191607236862, 0.02257939800620079, 0.01325287576764822, 0.14834734797477722, 0.0700073167681694, 0.12831416726112366, 0.47980472445487976, 0.3121630549430847, 0.05984592065215111, 0.015101294964551926, 0.002668763743713498, 0.0007187540177255869, 0.04004915803670883, 0.0007627750164829195, 0.05523831769824028, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0007188548916019499, 0.006864115130156279, 0.00033292395528405905, 0.000431404507253319, 0.0152564262971282, 0.2775210440158844, 0.03714991733431816, 0.7278205156326294, 0.004819776862859726, 0.00047404138604179025, 0.0003997469611931592, 0.0001266899926122278, 0.0201359074562788, 0.0027800032403320074, 0.042311206459999084, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00020999301341362298, 0.0025689874310046434, 3.502765650864603e-07, 6.610702985199168e-05, 0.00024143110204022378, 0.018905406817793846, 0.033397458493709564, 0.4650881290435791, 0.004783111158758402, 0.00013528004637919366, 5.751344360760413e-06, 7.93816871009767e-05, 0.0039043116848915815, 0.0005016719806008041, 0.07914639264345169, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00019393693946767598, 0.07456899434328079, 1.429513213224709e-05, 4.6383509470615536e-05, 6.820548151154071e-05, 0.004400796256959438, 0.0021800962276756763, 0.45963534712791443, 0.00143687822856009, 0.0008175616967491806, 6.983020284678787e-05, 3.49152869603131e-05, 0.0030698180198669434, 0.0006545006763190031, 0.001625033444724977, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.004301158711314201, 0.013502174988389015, 4.788395017385483e-05, 0.00021532995742745697, 7.713190279901028e-05, 0.001439842046238482, 0.005622516851872206, 0.121849425137043, 0.006593172438442707, 0.006624745205044746, 0.0006814572843722999, 0.0002721978526096791, 0.0009267745190300047, 0.0016606011195108294, 0.2357456088066101, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0064394231885671616, 0.03409593552350998, 0.0025135872419923544, 0.0008376456098631024, 0.0004409599641803652, 0.0026055865455418825, 0.005634414032101631, 0.014003962278366089, 0.2343187928199768, 0.08099395036697388, 0.23927520215511322, 0.01715606264770031, 0.10332414507865906, 0.021894987672567368, 0.1941189020872116, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0004975660121999681, 0.0015548047376796603, 6.826691333117196e-06, 1.0557592986515374e-06, 2.731301538005937e-05, 0.0005447702133096755, 0.00042012380436062813, 0.0503113828599453, 0.0053693996742367744, 0.0012762928381562233, 0.0017790982965379953, 0.019809026271104813, 0.47653263807296753, 0.008869247511029243, 0.017010610550642014, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00012974163109902292, 0.005610004533082247, 2.3442629753844813e-05, 1.8520654521125834e-06, 3.9678394387010485e-05, 0.0016583451069891453, 0.00029088594601489604, 0.004530484322458506, 0.0021493860986083746, 0.00029196502873674035, 0.0005848451401107013, 0.0028240433894097805, 0.4590959846973419, 0.22978197038173676, 0.0020738127641379833, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00021855060185771435, 0.005491270218044519, 1.9927349057979882e-05, 7.633860150235705e-06, 0.0004071943403687328, 0.008836714550852776, 7.301902951439843e-05, 0.011723233386874199, 1.7278060113312677e-05, 0.0001269245840376243, 0.00022235361393541098, 0.016586007550358772, 0.41012606024742126, 0.37776312232017517, 0.0024871949572116137, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.02619638666510582, 0.18392468988895416, 0.0003054745029658079, 0.00016413358389399946, 0.0015171386767178774, 0.004799532704055309, 0.004810427315533161, 0.058836404234170914, 0.0003794554795604199, 0.0017285931389778852, 0.000568193441722542, 0.003299211384728551, 0.6178385019302368, 0.5079926252365112, 0.05467592179775238, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03445081040263176, 0.14193737506866455, 0.0007241201237775385, 0.0002892682678066194, 0.0003202178922947496, 0.003702279180288315, 0.01134149543941021, 0.12129464000463486, 0.0006569268880411983, 0.0008894759230315685, 8.523569704266265e-05, 0.00030898841214366257, 0.7088924646377563, 0.10790188610553741, 0.05374660715460777, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04547691345214844, 0.010678221471607685, 0.0016328264027833939, 0.024403419345617294, 0.012795579619705677, 0.004323439672589302, 0.06414945423603058, 0.014008321799337864, 0.011475995182991028, 0.00871653389185667, 0.012156924232840538, 0.0147528275847435, 0.009472412057220936, 0.0331418551504612, 0.1366012692451477, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.11859580129384995, 0.07486707717180252, 0.21083025634288788, 0.32276296615600586, 0.08426652103662491, 0.03581860288977623, 0.24113436043262482, 0.608397364616394, 0.13584911823272705, 0.45509204268455505, 0.594833254814148, 0.30372148752212524, 0.8448506593704224, 0.7470672726631165, 0.09252076596021652, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04140070080757141, 0.00858838576823473, 0.11639615148305893, 0.1280786097049713, 0.2722368836402893, 0.21025919914245605, 0.4195333421230316, 0.631318211555481, 0.6560773253440857, 0.29341432452201843, 0.6862512230873108, 0.7675639986991882, 0.8915717005729675, 0.8601328730583191, 0.23356862366199493, 0.12451039254665375, 0.1335938721895218, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.23441848158836365, 0.1666196584701538, 0.16664288938045502, 0.25857093930244446, 0.13334479928016663, 0.17917701601982117, 0.8257887363433838, 0.7395779490470886, 0.6802234053611755, 0.8125103712081909, 0.671615719795227, 0.8831866383552551, 0.6773648858070374, 0.7102506160736084, 0.08689045161008835, 0.18396444618701935, 0.017508728429675102, 0.02471269853413105, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24967892467975616, 0.48421844840049744, 0.036505091935396194, 0.17128480970859528, 0.01777578890323639, 0.09479225426912308, 0.36135032773017883, 0.0868472084403038, 0.16740600764751434, 0.523710310459137, 0.24439233541488647, 0.42307958006858826, 0.6259368062019348, 0.3662186563014984, 0.20058651268482208, 0.18453162908554077, 0.038695670664310455, 0.04155581444501877, 0.05072518810629845, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.28931790590286255, 0.4439229369163513, 0.24370647966861725, 0.6020305752754211, 0.17363131046295166, 0.338454008102417, 0.5701692700386047, 0.33999428153038025, 0.68463534116745, 0.8701388239860535, 0.7831944823265076, 0.9611375331878662, 0.9679895043373108, 0.9072677493095398, 0.0468842089176178, 0.14826133847236633, 0.04252630099654198, 0.08689215034246445, 0.08308856934309006, 0.015247097238898277, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1225743219256401, 0.062406159937381744, 0.03387807682156563, 0.02868799865245819, 0.01787530817091465, 0.04143121838569641, 0.5920179486274719, 0.08798510581254959, 0.2968905568122864, 0.7129084467887878, 0.4609105885028839, 0.29060137271881104, 0.7909923791885376, 0.5701599717140198, 0.13614380359649658, 0.1348571479320526, 0.07033194601535797, 0.10030655562877655, 0.13752251863479614, 0.030713800340890884, 0.1331333965063095, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0705394446849823, 0.02209068462252617, 0.0211530439555645, 0.008882923051714897, 0.0033682750072330236, 0.08319123089313507, 0.11070933192968369, 0.0025125632528215647, 0.10380591452121735, 0.17744502425193787, 0.10391969978809357, 0.12427430599927902, 0.5562515258789062, 0.49710196256637573, 0.3223192095756531, 0.20671042799949646, 0.05809834972023964, 0.1630101054906845, 0.06033356115221977, 0.07501133531332016, 0.017328333109617233, 0.028450097888708115, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15847322344779968, 0.015464702621102333, 0.13866224884986877, 0.053395166993141174, 0.03494010120630264, 0.13738934695720673, 0.02684560976922512, 0.03214175999164581, 0.5759801864624023, 0.1755424290895462, 0.13409779965877533, 0.035038210451602936, 0.6489107012748718, 0.4460716247558594, 0.4074119031429291, 0.15813153982162476, 0.14090144634246826, 0.26030233502388, 0.10773709416389465, 0.16133210062980652, 0.04816069453954697, 0.01304988656193018, 0.13335363566875458, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00857736449688673, 0.012718217447400093, 0.01174219325184822, 0.012934550642967224, 0.006551709491759539, 0.24597492814064026, 0.030029013752937317, 0.05923602730035782, 0.04650798439979553, 0.02447274886071682, 0.019859377294778824, 0.003505804343149066, 0.04937520623207092, 0.05625420808792114, 0.28037816286087036, 0.3033713400363922, 0.22469042241573334, 0.4264413118362427, 0.3422197103500366, 0.14910078048706055, 0.06983038783073425, 0.023690486326813698, 0.010566752403974533, 0.05880258232355118, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0015372766647487879, 0.015295127406716347, 0.018696704879403114, 0.004789609462022781, 0.19481690227985382, 0.04769033566117287, 0.01355075929313898, 0.02196505106985569, 0.08700259774923325, 0.020393503829836845, 0.02400771528482437, 0.18789233267307281, 0.15418098866939545, 0.08713112771511078, 0.19334079325199127, 0.25368839502334595, 0.33459752798080444, 0.3829180896282196, 0.2782860994338989, 0.2427205741405487, 0.08768615871667862, 0.031752120703458786, 0.02143564634025097, 0.03798065707087517, 0.07379034906625748, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04759770259261131, 0.04375501722097397, 0.02714523859322071, 0.05194481834769249, 0.05246514454483986, 0.14355513453483582, 0.17152011394500732, 0.14246520400047302, 0.1098044142127037, 0.013531663455069065, 0.008927365764975548, 0.03807468339800835, 0.10050502419471741, 0.02236531302332878, 0.3381733298301697, 0.14200474321842194, 0.2391311228275299, 0.18728229403495789, 0.11236919462680817, 0.20923744142055511, 0.13365258276462555, 0.052715059369802475, 0.134474515914917, 0.14480768144130707, 0.06683899462223053, 0.104619100689888, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10647730529308319, 0.04246760904788971, 0.08123224973678589, 0.13003453612327576, 0.07854175567626953, 0.24148082733154297, 0.6790831685066223, 0.7492273449897766, 0.28685522079467773, 0.03681188449263573, 0.15954196453094482, 0.2672117054462433, 0.11099980026483536, 0.04468434303998947, 0.4826459586620331, 0.09595079720020294, 0.2752297520637512, 0.21842314302921295, 0.13660691678524017, 0.35477691888809204, 0.37130749225616455, 0.20556269586086273, 0.35276445746421814, 0.31008264422416687, 0.11074709892272949, 0.19841141998767853, 0.07199764251708984, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2962004542350769, 0.47284576296806335, 0.11245852708816528, 0.23689918220043182, 0.10807513445615768, 0.8532499074935913, 0.5788733959197998, 0.6375027894973755, 0.33168625831604004, 0.06381742656230927, 0.004373080097138882, 0.015940984711050987, 0.3371734917163849, 0.06828418374061584, 0.21185840666294098, 0.15323933959007263, 0.4611065983772278, 0.07869336754083633, 0.03600241616368294, 0.47375282645225525, 0.7350273132324219, 0.297486275434494, 0.6052883863449097, 0.4953201115131378, 0.144621342420578, 0.3493393063545227, 0.04881289228796959, 0.10520726442337036, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3828115463256836, 0.12613584101200104, 0.47516295313835144, 0.4473835527896881, 0.17031393945217133, 0.6938255429267883, 0.7945614457130432, 0.34594833850860596, 0.5323623418807983, 0.34808266162872314, 0.11382761597633362, 0.1349307745695114, 0.013382190838456154, 0.0600610226392746, 0.30783677101135254, 0.12003841996192932, 0.2704387903213501, 0.20063650608062744, 0.23778890073299408, 0.36254584789276123, 0.5319709777832031, 0.4483972191810608, 0.15058189630508423, 0.11134153604507446, 0.09426670521497726, 0.21241672337055206, 0.10488338023424149, 0.049764484167099, 0.15823495388031006, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.7362364530563354, 0.8323087096214294, 0.9336822032928467, 0.7739728689193726, 0.8897883296012878, 0.9609381556510925, 0.9334329962730408, 0.9553548693656921, 0.7747710943222046, 0.4005538523197174, 0.5586770176887512, 0.25099167227745056, 0.4200068712234497, 0.1631680577993393, 0.06528117507696152, 0.15233570337295532, 0.21891875565052032, 0.13215333223342896, 0.2837490439414978, 0.08042775094509125, 0.43866410851478577, 0.2773631513118744, 0.12773916125297546, 0.3155127763748169, 0.07932031899690628, 0.1219707503914833, 0.11212008446455002, 0.1944955438375473, 0.07170752435922623, 0.004313962999731302, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07449624687433243, 0.061402805149555206, 0.09389828145503998, 0.048646457493305206, 0.024208296090364456, 0.10819891840219498, 0.10563155263662338, 0.1243496686220169, 0.048523951321840286, 0.14693649113178253, 0.06614942103624344, 0.0066792843863368034, 0.2858017086982727, 0.04383772611618042, 0.15409637987613678, 0.2607015371322632, 0.3645761013031006, 0.37828943133354187, 0.3385462462902069, 0.2960833013057709, 0.5598280429840088, 0.544554591178894, 0.47054967284202576, 0.3477361798286438, 0.13701467216014862, 0.14822737872600555, 0.030188634991645813, 0.05528556555509567, 0.058441486209630966, 0.03410256654024124, 0.17273126542568207, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02467108517885208, 0.049052223563194275, 0.08135215938091278, 0.013768618926405907, 0.01176412496715784, 0.15210841596126556, 0.004693970084190369, 0.0041237217374145985, 0.018837640061974525, 0.03490369766950607, 0.036496780812740326, 0.0011750683188438416, 0.018557026982307434, 0.02382473833858967, 0.22122804820537567, 0.1872977614402771, 0.29805198311805725, 0.5206820368766785, 0.33024296164512634, 0.6395015716552734, 0.7210167050361633, 0.353913813829422, 0.406305193901062, 0.5096184015274048, 0.26257815957069397, 0.07301049679517746, 0.03464117646217346, 0.0787002444267273, 0.10916904360055923, 0.3557807505130768, 0.08364078402519226, 0.08538500964641571, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.012043171562254429, 0.03080524504184723, 0.02248452790081501, 0.008785543963313103, 0.00550604984164238, 0.05614035204052925, 0.015958979725837708, 0.01727765053510666, 0.03423915058374405, 0.017799094319343567, 0.029912255704402924, 0.01144923735409975, 0.09533664584159851, 0.02436906285583973, 0.20283196866512299, 0.13269101083278656, 0.2835436165332794, 0.47488275170326233, 0.24851854145526886, 0.694171130657196, 0.6760384440422058, 0.2759343385696411, 0.29058361053466797, 0.7136873602867126, 0.20711864531040192, 0.04295802861452103, 0.07691331952810287, 0.11943909525871277, 0.1323360651731491, 0.20847304165363312, 0.05967296287417412, 0.12062160670757294, 0.09502720832824707, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01959865354001522, 0.003073114436119795, 0.06498773396015167, 0.027286570519208908, 0.019540993496775627, 0.052237618714571, 0.08713454008102417, 0.28957968950271606, 0.3906492590904236, 0.044482238590717316, 0.17143161594867706, 0.1301742047071457, 0.10445850342512131, 0.03699616342782974, 0.2442801147699356, 0.058743223547935486, 0.276242733001709, 0.29826071858406067, 0.20218241214752197, 0.4631478488445282, 0.48415693640708923, 0.2865871787071228, 0.3694051504135132, 0.4054408073425293, 0.19627220928668976, 0.2907293438911438, 0.09057808667421341, 0.11348091810941696, 0.21781016886234283, 0.38082650303840637, 0.3570795953273773, 0.22612451016902924, 0.09323522448539734, 0.03618632256984711, NaN, NaN, NaN, NaN, NaN, NaN], [0.11208802461624146, 0.11668127030134201, 0.09828943759202957, 0.10754654556512833, 0.015885351225733757, 0.38998937606811523, 0.183034285902977, 0.3230077624320984, 0.20506803691387177, 0.08733018487691879, 0.007069121580570936, 0.010435528121888638, 0.30221423506736755, 0.047303054481744766, 0.19994190335273743, 0.07694489508867264, 0.41184449195861816, 0.038429711014032364, 0.018668875098228455, 0.5307568907737732, 0.7476497888565063, 0.4137455224990845, 0.6917499303817749, 0.6703397035598755, 0.3623183071613312, 0.579600989818573, 0.12613137066364288, 0.20100651681423187, 0.40998968482017517, 0.46115902066230774, 0.575211763381958, 0.35096046328544617, 0.163946270942688, 0.021770814433693886, 0.09986086189746857, NaN, NaN, NaN, NaN, NaN], [0.1682588905096054, 0.051582805812358856, 0.4415716230869293, 0.2735750675201416, 0.07878735661506653, 0.06776249408721924, 0.15038572251796722, 0.03211068734526634, 0.6709542274475098, 0.37688353657722473, 0.1879340261220932, 0.04096703231334686, 0.011627858504652977, 0.03471425548195839, 0.19384095072746277, 0.0834016501903534, 0.33346420526504517, 0.238715261220932, 0.28079062700271606, 0.5652539134025574, 0.6881173849105835, 0.5534363985061646, 0.22000034153461456, 0.1979052871465683, 0.3127084970474243, 0.4257359504699707, 0.18722867965698242, 0.1397658735513687, 0.3447277843952179, 0.13513657450675964, 0.31811001896858215, 0.32070791721343994, 0.12404847145080566, 0.05496959760785103, 0.04215753450989723, 0.16014836728572845, NaN, NaN, NaN, NaN], [0.8205305933952332, 0.9214023947715759, 0.9559677839279175, 0.7988566160202026, 0.9105063080787659, 0.9672437906265259, 0.9506043195724487, 0.9735420346260071, 0.9064961075782776, 0.6156813502311707, 0.6370130777359009, 0.18943972885608673, 0.3681671619415283, 0.1194160059094429, 0.08283783495426178, 0.13260646164417267, 0.29362690448760986, 0.18431688845157623, 0.38109344244003296, 0.20342527329921722, 0.5946046113967896, 0.4558189809322357, 0.26072001457214355, 0.5455912351608276, 0.2635512351989746, 0.31394094228744507, 0.23975242674350739, 0.36583349108695984, 0.2753828167915344, 0.01127256266772747, 0.41475725173950195, 0.29836422204971313, 0.2503683567047119, 0.10983213782310486, 0.21767295897006989, 0.0692884549498558, 0.003035380970686674, NaN, NaN, NaN], [0.10534824430942535, 0.08027994632720947, 0.1381307989358902, 0.07063161581754684, 0.01806548424065113, 0.10409632325172424, 0.12885765731334686, 0.2072904407978058, 0.09267445653676987, 0.23836983740329742, 0.11645739525556564, 0.006059943698346615, 0.1595546454191208, 0.017974214628338814, 0.14464683830738068, 0.2068602293729782, 0.4467880427837372, 0.4564751386642456, 0.4485791325569153, 0.45999279618263245, 0.6740500330924988, 0.7906107902526855, 0.6832103133201599, 0.5420533418655396, 0.4096798300743103, 0.3950984477996826, 0.13646338880062103, 0.10497336834669113, 0.17230592668056488, 0.07012390345335007, 0.27583980560302734, 0.3079235553741455, 0.1555996537208557, 0.038740403950214386, 0.05588690564036369, 0.03859011456370354, 0.02352789230644703, 0.12950412929058075, NaN, NaN], [0.026579611003398895, 0.02949470281600952, 0.04954056441783905, 0.017031243070960045, 0.008355016820132732, 0.09075918793678284, 0.0036468924954533577, 0.0022332987282425165, 0.050134338438510895, 0.049380820244550705, 0.028885982930660248, 0.0007559077348560095, 0.015549316070973873, 0.013319555670022964, 0.1734825074672699, 0.16561447083950043, 0.3958832919597626, 0.5531814098358154, 0.4040684700012207, 0.7809365391731262, 0.8175305128097534, 0.5712264180183411, 0.6113651394844055, 0.6668697595596313, 0.4850655198097229, 0.18787693977355957, 0.08608534932136536, 0.19115354120731354, 0.2498423308134079, 0.6246696710586548, 0.31422460079193115, 0.373276948928833, 0.049351077526807785, 0.046956032514572144, 0.08076699078083038, 0.09392194449901581, 0.3349837362766266, 0.062239501625299454, 0.10001940280199051, NaN], [0.05047497898340225, 0.027197130024433136, 0.11470095813274384, 0.007973222993314266, 0.12679167091846466, 0.4866730570793152, 0.17132264375686646, 0.15032453835010529, 0.14889459311962128, 0.01696154847741127, 0.0735161080956459, 0.0034290377516299486, 0.05194668471813202, 0.06144191324710846, 0.13309471309185028, 0.06568613648414612, 0.36780038475990295, 0.6246912479400635, 0.7116879820823669, 0.754679262638092, 0.7714072465896606, 0.7616819739341736, 0.5837911367416382, 0.9111838936805725, 0.8262851238250732, 0.6737059354782104, 0.5146453380584717, 0.7674095630645752, 0.7359525561332703, 0.5679676532745361, 0.7213301062583923, 0.6703079342842102, 0.5636342167854309, 0.38883939385414124, 0.5560528635978699, 0.518941342830658, 0.3739706873893738, 0.32013192772865295, 0.3743935525417328, 0.3977084755897522]], [[0.3143080472946167, 0.014564945362508297, 0.07743841409683228, 0.19665417075157166, 0.23130221664905548, 0.03274351730942726, 0.23599109053611755, 0.04763320833444595, 0.20168107748031616, 0.7521476149559021, 0.7922006249427795, 0.840878427028656, 0.6463541388511658, 0.6008138656616211, 0.0070990691892802715, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05880431830883026, 0.004086965229362249, 0.06557433307170868, 0.4476080536842346, 0.32179930806159973, 0.2046266496181488, 0.5952353477478027, 0.20483972132205963, 0.7834360599517822, 0.27592822909355164, 0.5900363922119141, 0.6986290812492371, 0.3548848032951355, 0.36629796028137207, 0.07452832907438278, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4484235942363739, 0.0712433010339737, 0.09740526974201202, 0.49982836842536926, 0.18807044625282288, 0.007537430617958307, 0.2073078453540802, 0.015238385647535324, 0.18028782308101654, 0.6095888018608093, 0.4225178062915802, 0.6769288778305054, 0.3957397937774658, 0.7102670669555664, 0.05611870437860489, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4341801106929779, 0.05481646955013275, 0.17834456264972687, 0.2579769194126129, 0.326920747756958, 0.0030261597130447626, 0.03147314488887787, 0.003279186552390456, 0.09941483289003372, 0.5679370760917664, 0.8480010032653809, 0.8133074045181274, 0.4710683822631836, 0.9189481139183044, 0.04321537911891937, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.559230387210846, 0.08983521163463593, 0.16111011803150177, 0.14667965471744537, 0.32596829533576965, 0.008685072883963585, 0.1111784353852272, 0.02690659649670124, 0.06770152598619461, 0.18340016901493073, 0.4614297151565552, 0.502476155757904, 0.42325475811958313, 0.5992166996002197, 0.05437220633029938, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.367906779050827, 0.21432256698608398, 0.3548191487789154, 0.2603428363800049, 0.22096140682697296, 0.0013341127196326852, 0.021726170554757118, 0.005543001927435398, 0.5389296412467957, 0.818263828754425, 0.919593095779419, 0.8187286257743835, 0.4823090434074402, 0.4897681474685669, 0.07018090784549713, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7116888761520386, 0.17206020653247833, 0.6874114871025085, 0.19288089871406555, 0.20990870893001556, 0.011273512616753578, 0.2026582807302475, 0.004371582996100187, 0.10976968705654144, 0.4432500898838043, 0.7022042274475098, 0.8704607486724854, 0.721519947052002, 0.7422701716423035, 0.025589054450392723, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7674684524536133, 0.20032620429992676, 0.42808812856674194, 0.11714937537908554, 0.32732346653938293, 0.009955272078514099, 0.05444686487317085, 0.0040375906974077225, 0.12078685313463211, 0.6266691088676453, 0.5163981914520264, 0.8307003378868103, 0.32096055150032043, 0.24524804949760437, 0.04717922583222389, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7549813389778137, 0.15439504384994507, 0.33331331610679626, 0.24930144846439362, 0.2927357852458954, 0.04936225712299347, 0.44933974742889404, 0.06466211378574371, 0.09519664198160172, 0.08716140687465668, 0.058296240866184235, 0.09990595281124115, 0.5117565989494324, 0.1508449912071228, 0.039490822702646255, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.654628574848175, 0.3205694854259491, 0.5841068029403687, 0.21299651265144348, 0.365792840719223, 0.0401315838098526, 0.18686936795711517, 0.05883712321519852, 0.05069931596517563, 0.33667507767677307, 0.3354107439517975, 0.22027519345283508, 0.05277648940682411, 0.09031395614147186, 0.015531455166637897, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3366456627845764, 0.1530359387397766, 0.41866233944892883, 0.39775165915489197, 0.7769761681556702, 0.06979230791330338, 0.41583842039108276, 0.02130916155874729, 0.14617334306240082, 0.25815388560295105, 0.1423572301864624, 0.18894770741462708, 0.041056301444768906, 0.026175418868660927, 0.03888533264398575, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.24913249909877777, 0.0818726196885109, 0.5426726341247559, 0.1687711775302887, 0.8305720090866089, 0.26261457800865173, 0.39635857939720154, 0.1712585836648941, 0.1158638522028923, 0.17366157472133636, 0.12521226704120636, 0.5298976302146912, 0.041029125452041626, 0.02415779046714306, 0.1170416921377182, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3567614257335663, 0.035316068679094315, 0.3819185495376587, 0.10469090938568115, 0.3454773426055908, 0.09596268832683563, 0.3821227550506592, 0.17425164580345154, 0.40528857707977295, 0.1745157092809677, 0.10956539213657379, 0.5078453421592712, 0.0026470222510397434, 0.016186503693461418, 0.08932095021009445, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.330766886472702, 0.039845019578933716, 0.6981685757637024, 0.09713104367256165, 0.8411048650741577, 0.16356231272220612, 0.3630223274230957, 0.1627381145954132, 0.6954487562179565, 0.17326875030994415, 0.1752558946609497, 0.24479816854000092, 0.026946308091282845, 0.016200177371501923, 0.06702017039060593, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07683827728033066, 0.07034450024366379, 0.21707428991794586, 0.2902449369430542, 0.1834353357553482, 0.01726321130990982, 0.13144701719284058, 0.005189047660678625, 0.150242418050766, 0.1182665303349495, 0.4041094183921814, 0.12062898278236389, 0.05959685891866684, 0.1186181977391243, 0.1283060759305954, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.005987181328237057, 0.0011158415582031012, 0.0026756690349429846, 0.0011391430161893368, 0.0021053741220384836, 0.0005449134623631835, 0.0017384873935952783, 0.000736464629881084, 0.00014482461847364902, 0.0008784460369497538, 0.0008941806154325604, 0.0009559267782606184, 0.00015614555741194636, 0.00044419756159186363, 0.16329224407672882, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3448674976825714, 0.07203025370836258, 0.011963781900703907, 0.012941744178533554, 0.011539866216480732, 0.003333584638312459, 0.005511423572897911, 0.0016478801844641566, 0.003020848147571087, 0.006189296022057533, 0.0020935258362442255, 0.00048376841004937887, 8.994764357339591e-05, 0.00040787423495203257, 0.2113737165927887, 0.1305680274963379, 0.02726716920733452, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.44219815731048584, 0.8124432563781738, 0.1900549679994583, 0.3808274269104004, 0.045300956815481186, 0.024617541581392288, 0.0172295980155468, 0.03488133102655411, 0.004235385917127132, 0.05999733507633209, 0.03787413239479065, 0.0011567235924303532, 0.0017442036187276244, 0.008845857344567776, 0.004224383272230625, 0.002169837476685643, 0.0032534021884202957, 0.5694547891616821, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07874103635549545, 0.02866651676595211, 0.3287397623062134, 0.27984437346458435, 0.10563887655735016, 0.003691220423206687, 0.005916049238294363, 0.0007406381191685796, 0.0005066083394922316, 0.0481056272983551, 0.029072491452097893, 0.000652547983918339, 0.0003529583918862045, 0.0009863339364528656, 0.002192106796428561, 0.1568225622177124, 0.12336109578609467, 0.028200775384902954, 0.03890102356672287, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.030638281255960464, 0.02597089111804962, 0.6577842831611633, 0.16596756875514984, 0.48041173815727234, 0.6114144921302795, 0.028207998722791672, 0.053615398705005646, 0.1417267620563507, 0.03454216569662094, 0.023575417697429657, 0.004873087164014578, 0.0009616028983145952, 0.00223313900642097, 0.0011337294708937407, 0.008017625659704208, 0.013223886489868164, 0.04581261798739433, 0.017950134351849556, 0.8790656328201294, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.29477018117904663, 0.14754106104373932, 0.8534399271011353, 0.9182198643684387, 0.6083860993385315, 0.9389832019805908, 0.12579986453056335, 0.03590020909905434, 0.012173496186733246, 0.16479530930519104, 0.15366923809051514, 0.0035958383232355118, 0.002988115418702364, 0.026292480528354645, 0.0003885648038703948, 0.08130903542041779, 0.2643316090106964, 0.5756329894065857, 0.29882851243019104, 0.31516125798225403, 0.09644471108913422, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2897806465625763, 0.01695333980023861, 0.6714832782745361, 0.4471692144870758, 0.24303969740867615, 0.15563154220581055, 0.008645682595670223, 0.0004950988804921508, 0.0001695932005532086, 0.13566477596759796, 0.030448369681835175, 0.00021736785129178315, 9.297585347667336e-05, 0.0014399208594113588, 5.083655923954211e-05, 0.20484277606010437, 0.3443664610385895, 0.0019387316424399614, 0.017399819567799568, 0.0004214652581140399, 0.00013534165918827057, 0.01563790813088417, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1102917492389679, 0.0027466323226690292, 0.13646264374256134, 0.07094646990299225, 0.17040857672691345, 0.6033481955528259, 0.41631338000297546, 0.013031017035245895, 0.00012492973473854363, 0.005976412910968065, 0.0002816450723912567, 4.682707003667019e-05, 0.00021861463028471917, 0.00019605428678914905, 0.001022772048600018, 0.1571786254644394, 0.5643889307975769, 0.13441002368927002, 0.09036820381879807, 0.02947377972304821, 0.015878956764936447, 0.022048691287636757, 0.14189693331718445, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.7042187452316284, 0.49455204606056213, 0.43194010853767395, 0.7080989480018616, 0.382207989692688, 0.06800723820924759, 0.48792970180511475, 0.12651333212852478, 0.0012585417134687304, 0.07895761728286743, 0.01729964278638363, 0.0006471746601164341, 0.00013743228919338435, 0.00039039706462062895, 0.00010207234299741685, 0.005826869048178196, 0.13292454183101654, 0.00521356426179409, 0.005004087463021278, 0.10703893005847931, 0.26877719163894653, 0.1785666048526764, 0.23197543621063232, 0.007970587350428104, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5233215093612671, 0.7835124135017395, 0.3596530258655548, 0.5502080917358398, 0.589034378528595, 0.24138878285884857, 0.4714515507221222, 0.13250088691711426, 0.08884716778993607, 0.06473898142576218, 0.12478159368038177, 0.001717525301501155, 0.01358798798173666, 0.004862584639340639, 0.0004225081647746265, 0.03136341646313667, 0.08873608708381653, 0.009185479953885078, 0.03043411858379841, 0.3010490834712982, 0.36070317029953003, 0.178965762257576, 0.21872122585773468, 0.005464768502861261, 0.06020791083574295, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0975094586610794, 0.14095744490623474, 0.009511731564998627, 0.03128954395651817, 0.01951521448791027, 0.0017430862644687295, 0.033708807080984116, 0.009512575343251228, 0.3042309582233429, 0.0025639990344643593, 0.0006334132049232721, 2.5987004846683703e-05, 0.0001574041525600478, 1.1997842193522956e-05, 1.5690195141360164e-05, 0.07854610681533813, 0.03772095590829849, 0.016643106937408447, 0.02832828275859356, 0.0785825327038765, 0.09336084127426147, 0.24177083373069763, 0.2718014717102051, 0.12932275235652924, 0.08437053114175797, 0.24188947677612305, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.536220133304596, 0.12877297401428223, 0.013534938916563988, 0.13534405827522278, 0.015604051761329174, 0.0035537974908947945, 0.02344023622572422, 0.008398037403821945, 0.2580391466617584, 0.2587551474571228, 0.014949243515729904, 0.0010696486569941044, 0.00046315763029269874, 0.0013398011215031147, 8.422375685768202e-05, 0.17239268124103546, 0.029533302411437035, 0.030515655875205994, 0.026403654366731644, 0.05037287250161171, 0.13986584544181824, 0.11416076123714447, 0.08228978514671326, 0.26975753903388977, 0.020502708852291107, 0.030797043815255165, 0.006723156664520502, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.028944578021764755, 0.013114584609866142, 0.0438210591673851, 0.05079193785786629, 0.03694206848740578, 0.0008442872785963118, 0.0030779552180320024, 0.002579997293651104, 0.01023491844534874, 0.21445545554161072, 0.2806929349899292, 0.00855539832264185, 0.03333647921681404, 0.06091907247900963, 1.9560096916393377e-05, 0.35662412643432617, 0.005917226430028677, 0.00044432797585614026, 0.00022813511895947158, 0.0073361690156161785, 0.0027237480971962214, 0.007987208664417267, 0.021625559777021408, 0.010472757741808891, 0.0008755659800954163, 0.012584702111780643, 0.000526397256180644, 0.01033733133226633, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0058769844472408295, 0.06350620836019516, 0.003568005282431841, 0.0076079596765339375, 0.0037217612843960524, 0.004286385141313076, 0.03584115207195282, 0.14617407321929932, 0.0030082303564995527, 0.12143123894929886, 0.0793885663151741, 0.1555183082818985, 0.14442139863967896, 0.29275521636009216, 7.129996811272576e-05, 0.189227893948555, 0.01606086827814579, 0.0030457540415227413, 0.005861388053745031, 0.04963670298457146, 0.004091562703251839, 0.01225967425853014, 0.037419673055410385, 0.01020084973424673, 0.003108290024101734, 0.01512740459293127, 0.006679146084934473, 0.014098022133111954, 0.03816642239689827, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.034930020570755005, 0.09419079124927521, 0.0127689428627491, 0.008763227611780167, 0.0065171802416443825, 0.008632887154817581, 0.02612082101404667, 0.02043459191918373, 0.0836663544178009, 0.5329904556274414, 0.3228733241558075, 0.7184357047080994, 0.5793755650520325, 0.783859133720398, 0.0001531920424895361, 0.00965302623808384, 0.0035168000031262636, 0.03902876377105713, 0.0158648993819952, 0.32648226618766785, 0.0038036927580833435, 0.002248003613203764, 0.002372291637584567, 0.014672092162072659, 0.007728067692369223, 0.022481968626379967, 0.028911879286170006, 0.044244468212127686, 0.021532919257879257, 0.6417658925056458, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0009532110998407006, 0.0024861039128154516, 7.189704774646088e-05, 0.00014637503772974014, 2.8552024105010787e-06, 3.0342853278853e-05, 0.0007709002820774913, 0.0005337693146429956, 6.919851330167148e-06, 0.02619163505733013, 0.02381032705307007, 0.008668542839586735, 0.39639002084732056, 0.7824769616127014, 1.1539431170604075e-06, 0.037641312927007675, 0.005557402968406677, 0.0006393054500222206, 0.006437606643885374, 0.007460788358002901, 0.0009530181414447725, 0.0016025539953261614, 0.0067516821436584, 0.02322007343173027, 0.018459537997841835, 0.011051125824451447, 0.006488891318440437, 0.04039585590362549, 0.18200218677520752, 0.0006002468289807439, 0.6243939995765686, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02785377763211727, 0.15845024585723877, 0.19323119521141052, 0.06543393433094025, 0.014044036157429218, 0.040286585688591, 0.07583826035261154, 0.6567350029945374, 0.004159754142165184, 0.35265031456947327, 0.6287637948989868, 0.12951745092868805, 0.32439297437667847, 0.653313934803009, 0.0008144593448378146, 0.01615065336227417, 0.01699231006205082, 0.00012957912986166775, 0.016060354188084602, 0.0006264564581215382, 0.0012908404460176826, 0.002684527076780796, 0.027531128376722336, 0.015566377900540829, 0.003692139405757189, 0.5753727555274963, 0.5145941376686096, 0.03750383481383324, 0.009545800276100636, 0.0034461882896721363, 0.005381980445235968, 0.00046628122800029814, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02927210181951523, 0.04805546626448631, 0.295967698097229, 0.060625556856393814, 0.014990724623203278, 0.10397231578826904, 0.12186732143163681, 0.5237559080123901, 0.0203724168241024, 0.43874940276145935, 0.4409005343914032, 0.09095493704080582, 0.5531511306762695, 0.5263633728027344, 0.0002321143983863294, 0.021861553192138672, 0.01695878431200981, 0.0018149337265640497, 0.015764223411679268, 0.007719711866229773, 0.0034752548672258854, 0.007653116714209318, 0.03472340479493141, 0.038436826318502426, 0.014262136071920395, 0.8426622748374939, 0.36256304383277893, 0.21876515448093414, 0.019672129303216934, 0.020847154781222343, 0.00781619269400835, 0.005409067030996084, 0.16073459386825562, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5664732456207275, 0.02422192506492138, 0.3148367702960968, 0.37531769275665283, 0.06290365755558014, 0.02708868682384491, 0.03764869272708893, 0.06476183980703354, 0.09221415221691132, 0.3172641098499298, 0.088014617562294, 0.02202794700860977, 0.004314645659178495, 0.0619816817343235, 0.0017959593096747994, 0.18507197499275208, 0.027911728248000145, 0.014699580147862434, 0.025536103174090385, 0.014524195343255997, 0.045023027807474136, 0.031167738139629364, 0.07539253681898117, 0.22652071714401245, 0.011904416605830193, 0.08752688765525818, 0.03955431655049324, 0.2908211648464203, 0.03612781688570976, 0.00514488760381937, 0.017019467428326607, 0.07116629183292389, 0.03509910777211189, 0.02026083506643772, NaN, NaN, NaN, NaN, NaN, NaN], [0.04828598350286484, 0.01127469539642334, 0.1758044958114624, 0.0725238099694252, 0.01880812831223011, 0.003422890789806843, 0.0039800796657800674, 0.008112750947475433, 0.0007020575576461852, 0.0960424467921257, 0.3098883628845215, 0.03193678706884384, 0.03351299837231636, 0.2577627897262573, 0.0005041947006247938, 0.40259334444999695, 0.005078054964542389, 0.00017122419376391917, 9.21270766411908e-05, 0.002624903805553913, 0.0009363252320326865, 0.00360113475471735, 0.01331485528498888, 0.008243494667112827, 0.0007176694343797863, 0.019634194672107697, 0.002027983544394374, 0.02349759265780449, 0.030203014612197876, 0.000993669149465859, 0.0008422310347668827, 0.013102295808494091, 0.025159381330013275, 0.0006507099606096745, 0.018182074651122093, NaN, NaN, NaN, NaN, NaN], [0.008833246305584908, 0.03231082111597061, 0.009648996405303478, 0.01135926228016615, 0.004257569555193186, 0.002696139505133033, 0.026390861719846725, 0.07894735038280487, 0.0002903220884036273, 0.05877671018242836, 0.0971919596195221, 0.32856324315071106, 0.08294347673654556, 0.6861463785171509, 0.00047716210247017443, 0.2579963207244873, 0.021157346665859222, 0.002921733073890209, 0.006211739499121904, 0.031850416213274, 0.0022005264181643724, 0.0070661455392837524, 0.036871425807476044, 0.012320333160459995, 0.005331193562597036, 0.033889420330524445, 0.020235266536474228, 0.07458563148975372, 0.1398555487394333, 0.008059950545430183, 0.0405682735145092, 0.03368399292230606, 0.012085597030818462, 0.010676471516489983, 0.03411625698208809, 0.08152885735034943, NaN, NaN, NaN, NaN], [0.020260397344827652, 0.03928471356630325, 0.012783887796103954, 0.0091601787135005, 0.005565040744841099, 0.007968534715473652, 0.020862603560090065, 0.012279938906431198, 0.01832268387079239, 0.3204420506954193, 0.28696081042289734, 0.7937509417533875, 0.6314787864685059, 0.8277974724769592, 0.00014348741387948394, 0.005019576288759708, 0.001437423750758171, 0.014701779931783676, 0.005876661743968725, 0.15098156034946442, 0.001037455745972693, 0.0006782425916753709, 0.0010664333822205663, 0.006170186679810286, 0.004750464111566544, 0.015587885864078999, 0.020612932741642, 0.024904461577534676, 0.027292385697364807, 0.6522603631019592, 0.02780178189277649, 0.009980881586670876, 0.010863273404538631, 0.016993993893265724, 0.026612548157572746, 0.013426730409264565, 0.6643192768096924, NaN, NaN, NaN], [0.00497927563264966, 0.011739314533770084, 0.0009416648535989225, 0.0009133343119174242, 2.0598932678694837e-05, 0.00024278588534798473, 0.00463896244764328, 0.0027787971775978804, 1.9694551156135276e-05, 0.026842234656214714, 0.05824153125286102, 0.023767979815602303, 0.7019069194793701, 0.8979114294052124, 1.5536308637820184e-05, 0.023952102288603783, 0.0025056565646082163, 0.0002975048264488578, 0.0031560298521071672, 0.002087814500555396, 0.00019765450269915164, 0.00028781042783521116, 0.0023521913681179285, 0.009429593570530415, 0.010675383731722832, 0.013774069957435131, 0.012372920289635658, 0.030660077929496765, 0.3810364305973053, 0.0006224916432984173, 0.6039706468582153, 0.2701583206653595, 0.012816790491342545, 0.005745226051658392, 0.052403513342142105, 0.18411211669445038, 0.00043697847286239266, 0.6234135627746582, NaN, NaN], [0.06832221150398254, 0.18812543153762817, 0.5426309108734131, 0.237625390291214, 0.041615329682826996, 0.11611851304769516, 0.16301436722278595, 0.827357828617096, 0.011619587428867817, 0.35340800881385803, 0.8248108625411987, 0.22083298861980438, 0.4978465139865875, 0.8379470109939575, 0.008811386302113533, 0.007988094352185726, 0.006256349850445986, 4.065780740347691e-05, 0.006692530121654272, 0.00010113247117260471, 0.0002641561150085181, 0.0006015493418090045, 0.009669815190136433, 0.00486318813636899, 0.0012557843001559377, 0.43231210112571716, 0.35852983593940735, 0.01959061808884144, 0.007567983586341143, 0.0019125458784401417, 0.00857639778405428, 0.0005027590086683631, 0.41286540031433105, 0.4292365312576294, 0.01753525249660015, 0.005813234485685825, 0.00216498039662838, 0.003382693277671933, 0.00027526391204446554, NaN], [0.7676634788513184, 0.8615484237670898, 0.768317461013794, 0.9594964981079102, 0.36958935856819153, 0.4649639129638672, 0.5634418725967407, 0.8043064475059509, 0.6601962447166443, 0.9397303462028503, 0.8348119258880615, 0.9867405295372009, 0.7646960020065308, 0.8154686689376831, 0.03640103340148926, 0.1387476772069931, 0.027318276464939117, 0.00785337295383215, 0.019197843968868256, 0.013794281519949436, 0.020801816135644913, 0.013009469024837017, 0.07068510353565216, 0.020734209567308426, 0.024748992174863815, 0.04673967882990837, 0.025586238130927086, 0.01648368127644062, 0.06557000428438187, 0.022920427843928337, 0.013843921944499016, 0.04100487753748894, 0.0375630147755146, 0.023956134915351868, 0.018727701157331467, 0.05957711860537529, 0.020177751779556274, 0.007389482576400042, 0.027843382209539413, 0.025224220007658005]], [[0.06827192008495331, 0.0036808219738304615, 0.005701950751245022, 0.005157816223800182, 0.003777393838390708, 0.024757172912359238, 0.0020165019668638706, 0.010267351754009724, 0.013163687661290169, 0.001690453034825623, 0.00837681908160448, 0.00522418599575758, 0.061038240790367126, 0.015438525006175041, 0.325132817029953, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7422951459884644, 0.028774140402674675, 0.06394203752279282, 0.00887901522219181, 0.04345611855387688, 0.027670713141560555, 0.0295904241502285, 0.01398912351578474, 0.025535697117447853, 0.02094031311571598, 0.022182827815413475, 0.009663421660661697, 0.049684178084135056, 0.026225639507174492, 0.13834334909915924, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.20897099375724792, 0.21868035197257996, 0.23815643787384033, 0.005872054491192102, 0.0010661164997145534, 0.0017293300479650497, 0.00042713910806924105, 0.002609806600958109, 0.016046296805143356, 0.009100147522985935, 0.014420107938349247, 0.0022624030243605375, 0.010553905740380287, 0.007111164275556803, 0.25332581996917725, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2508500814437866, 0.20390872657299042, 0.7329782247543335, 0.07117453217506409, 0.016424261033535004, 0.021444672718644142, 0.001510130357928574, 0.004098558332771063, 0.0484151765704155, 0.02061472274363041, 0.001126835006289184, 0.0022107160184532404, 0.007578131277114153, 0.004504901356995106, 0.1403624713420868, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.27370113134384155, 0.8174626231193542, 0.7193068861961365, 0.7076587677001953, 0.07771007716655731, 0.01620337925851345, 0.004001453518867493, 0.004182097036391497, 0.03681829199194908, 0.09453201293945312, 0.026799198240041733, 0.006044679321348667, 0.03725922852754593, 0.016391301527619362, 0.04474738612771034, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3889567255973816, 0.4487122893333435, 0.5870586037635803, 0.6609426140785217, 0.6319714188575745, 0.10676700621843338, 0.009257740341126919, 0.0017087672604247928, 0.027955975383520126, 0.07590407133102417, 0.006841681431978941, 0.08621303737163544, 0.05063363164663315, 0.016846608370542526, 0.05719457566738129, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00991373136639595, 0.0983041524887085, 0.15667210519313812, 0.19277995824813843, 0.5809133052825928, 0.7996482253074646, 0.06316149979829788, 0.004939877428114414, 0.023352928459644318, 0.010926214046776295, 0.008795071393251419, 0.006998055148869753, 0.0765714943408966, 0.006783204153180122, 0.05886436253786087, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07887525111436844, 0.017153050750494003, 0.2216421663761139, 0.13068468868732452, 0.5295770764350891, 0.35302138328552246, 0.8493326902389526, 0.04265422001481056, 0.052519019693136215, 0.027357611805200577, 0.01357424259185791, 0.004279646556824446, 0.026089098304510117, 0.04089489206671715, 0.014124121516942978, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03465811163187027, 0.15351061522960663, 0.2825109362602234, 0.08174889534711838, 0.19755861163139343, 0.5825939774513245, 0.37084007263183594, 0.7892780900001526, 0.1287456750869751, 0.006381133571267128, 0.001940184272825718, 0.00047384126810356975, 0.011903955601155758, 0.003972942009568214, 0.06710142642259598, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.013788340613245964, 0.006632686126977205, 0.02207767777144909, 0.0785517543554306, 0.014113685116171837, 0.048156753182411194, 0.1944313496351242, 0.22155866026878357, 0.49656373262405396, 0.009422117844223976, 0.004702835343778133, 0.0007582302205264568, 0.00014129001647233963, 0.00033574484405107796, 0.23994654417037964, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00469209672883153, 0.015491061843931675, 0.035103749483823776, 0.009631682187318802, 0.008573818951845169, 0.051444172859191895, 0.04315423220396042, 0.05495374649763107, 0.6859460473060608, 0.5370080471038818, 0.06784479320049286, 0.004556083586066961, 0.001035997993312776, 0.0006345660076476634, 0.13974453508853912, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.02668480947613716, 0.016245348379015923, 0.01112398225814104, 0.008507933467626572, 0.02067524567246437, 0.17763113975524902, 0.05662769451737404, 0.04544723033905029, 0.7948054671287537, 0.7384940385818481, 0.5224500298500061, 0.1060851439833641, 0.014122114516794682, 0.0019289307529106736, 0.08371670544147491, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.02394592948257923, 0.04371663182973862, 0.028385786339640617, 0.007640721742063761, 0.014576996676623821, 0.08887659758329391, 0.017377078533172607, 0.020801657810807228, 0.187345951795578, 0.5047414302825928, 0.6342922449111938, 0.3672487437725067, 0.04719087854027748, 0.10966072231531143, 0.08543073385953903, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.009629062376916409, 0.020042795687913895, 0.006009343545883894, 0.001406975439749658, 0.0026742229238152504, 0.006072318647056818, 0.006495587062090635, 0.0032924923580139875, 0.034326668828725815, 0.5998041033744812, 0.7456773519515991, 0.7204623818397522, 0.012111457996070385, 0.018825965002179146, 0.008305574767291546, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08114123344421387, 0.05478224158287048, 0.11802507936954498, 0.1980995535850525, 0.15338915586471558, 0.11414031684398651, 0.06528255343437195, 0.04494854062795639, 0.26375874876976013, 0.30061599612236023, 0.26960447430610657, 0.5329554677009583, 0.4288364350795746, 0.12292250245809555, 0.12395624816417694, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5194346308708191, 0.08715501427650452, 0.09860441088676453, 0.08100719004869461, 0.11848669499158859, 0.14280925691127777, 0.19592297077178955, 0.1196337640285492, 0.2793996334075928, 0.0691760703921318, 0.09539081901311874, 0.05545644089579582, 0.02620256133377552, 0.03735822066664696, 0.09928011149168015, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.002687783446162939, 0.2585922181606293, 0.004556892905384302, 0.0005560630816034973, 0.0013625096762552857, 0.000865808455273509, 2.095674426527694e-05, 0.013363445177674294, 1.4331720194604713e-05, 0.00023233501997310668, 0.013212678954005241, 0.00027388104354031384, 2.99917119264137e-05, 5.10126119479537e-05, 0.0653858631849289, 0.1319446712732315, 0.003103907685726881, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.010489544831216335, 0.001751396106556058, 0.2775154411792755, 0.0030420231632888317, 0.08156438916921616, 0.0006471106316894293, 1.7804295566747896e-05, 0.00014657371502835304, 0.00035265504266135395, 0.00129506376106292, 0.018553601577878, 0.0019669390749186277, 0.009056665003299713, 0.05091148242354393, 0.1541917622089386, 0.004627853631973267, 0.8189921975135803, 0.006355744786560535, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0025869093369692564, 0.008571458049118519, 0.38431695103645325, 0.030530055984854698, 0.03365315869450569, 0.005854337941855192, 0.00010941662185359746, 4.1041937947738916e-05, 0.000364075880497694, 0.0011989381164312363, 0.014197473414242268, 0.0010815636487677693, 0.0004893331206403673, 0.0013785242335870862, 0.011478900909423828, 0.0004822930786758661, 0.5574855208396912, 0.0058120423927903175, 0.014268792234361172, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20589935779571533, 0.03613102436065674, 0.009011336602270603, 0.09399610757827759, 0.042497485876083374, 0.000576009857468307, 0.0040712482295930386, 0.00162220629863441, 0.00015305644774343818, 0.0034409475047141314, 0.025435233488678932, 2.175084773625713e-05, 1.0188268788624555e-05, 5.634217450278811e-05, 0.160919189453125, 0.15055440366268158, 0.0014966451562941074, 0.1733904629945755, 0.05038055405020714, 0.0057296124286949635, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00994176883250475, 0.015379102900624275, 0.000435269670560956, 0.004355194512754679, 0.002023787936195731, 4.86412636746536e-06, 0.0007220985717140138, 0.0004895065212622285, 0.0005591813242062926, 0.009127096273005009, 0.023014724254608154, 0.0003639610658865422, 3.1703839340480044e-05, 0.00036040451959706843, 0.1469942033290863, 0.1304439753293991, 0.00022060537594370544, 0.03428095951676369, 0.0157721396535635, 0.20856629312038422, 0.2746620774269104, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.31647789478302, 0.5689504742622375, 0.010991040617227554, 0.29046669602394104, 0.008814695291221142, 0.008600234054028988, 0.094898521900177, 0.02089405618607998, 0.005384301766753197, 0.1224634200334549, 0.2525540888309479, 0.011421876028180122, 9.89354812190868e-05, 0.00020726426737383008, 0.3419104218482971, 0.017820989713072777, 1.0936159014818259e-05, 0.0006241680239327252, 4.3406893382780254e-05, 0.2565733790397644, 0.5255003571510315, 0.040596142411231995, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.006757077760994434, 0.1354868859052658, 0.002759847091510892, 0.009205225855112076, 0.0038083188701421022, 0.0014255000278353691, 0.0007299972930923104, 0.2051592320203781, 0.00020230394147802144, 0.001623967313207686, 0.006681961473077536, 0.0021689198911190033, 5.557909025810659e-05, 0.000162289768923074, 0.20840437710285187, 0.2143511176109314, 3.818454570136964e-05, 0.0006476931739598513, 0.00012842394062317908, 0.007853559218347073, 0.008102592080831528, 0.0005345920799300075, 0.00793861411511898, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.010027364827692509, 0.02789497748017311, 0.0041139991953969, 0.012661347165703773, 0.0013435317669063807, 0.0034407242201268673, 0.0064836894161999226, 0.007366063538938761, 0.29601985216140747, 0.053567804396152496, 0.040060218423604965, 0.004607491660863161, 0.00018677859043236822, 3.186250978615135e-05, 0.10952453315258026, 0.00014670012751594186, 7.536429620813578e-06, 0.0001294321846216917, 0.00024457855033688247, 0.00022483686916530132, 0.001284220488741994, 0.0014163334853947163, 0.5552030801773071, 0.006061996798962355, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19971387088298798, 0.012958711944520473, 0.001638519112020731, 0.17775660753250122, 0.0022716999519616365, 0.03685721755027771, 0.06948257982730865, 0.005452410783618689, 0.037147630006074905, 0.19678887724876404, 0.21911752223968506, 0.02466990426182747, 0.0004891769494861364, 6.33890085737221e-05, 0.21250228583812714, 0.09223808348178864, 0.004348577931523323, 0.013163902796804905, 0.018216131255030632, 0.035016678273677826, 0.11075899004936218, 0.1728493720293045, 0.19621391594409943, 0.029301786795258522, 0.46166056394577026, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05692211166024208, 0.036700569093227386, 0.0015533106634393334, 0.01848980039358139, 0.002404581755399704, 0.008354752324521542, 0.023693444207310677, 0.02836945652961731, 0.29948922991752625, 0.005321406293660402, 0.0022319734562188387, 0.0005214664852246642, 0.00019869217067025602, 5.8369230828247964e-05, 0.008838840760290623, 0.11309938877820969, 0.004489036742597818, 0.0485633909702301, 0.021462395787239075, 0.4192940890789032, 0.26214849948883057, 0.22032421827316284, 0.0067114257253706455, 0.010406548157334328, 0.11692964285612106, 0.23004111647605896, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.011123275384306908, 0.003955129534006119, 0.0015235289465636015, 0.011223106645047665, 0.002481319010257721, 0.000903434120118618, 0.0006720115779899061, 0.00024289102293550968, 0.010115177370607853, 0.26232361793518066, 0.014199022203683853, 0.0005582758458331227, 0.0001542939426144585, 5.357913687475957e-05, 0.050008371472358704, 0.14281870424747467, 0.000545236689504236, 0.003893920686095953, 0.0005153689999133348, 0.01790653169155121, 0.004868220537900925, 0.0031487985979765654, 0.0011714915744960308, 0.0043698386289179325, 0.020373020321130753, 0.02358497679233551, 0.2682037353515625, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.025191567838191986, 0.009952094405889511, 0.015023785643279552, 0.0893990620970726, 0.006299919448792934, 0.0077370950020849705, 0.0004422276106197387, 0.00010742250742623582, 0.001807618304155767, 0.052116382867097855, 0.33116668462753296, 0.0029348258394747972, 0.004942082799971104, 0.0017646296182647347, 0.009777115657925606, 0.09794370085000992, 0.0018320194212719798, 0.000285644200630486, 3.260145604144782e-05, 0.00041393720312044024, 0.0043053096160292625, 0.002047628629952669, 0.0003047001373488456, 0.002447759034112096, 0.0016152235912159085, 0.024524936452507973, 0.29461416602134705, 0.014563476666808128, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12133541703224182, 0.0033125760965049267, 0.008441481739282608, 0.0257105715572834, 0.005432062782347202, 0.020603680983185768, 0.0008238950395025313, 0.00019463927310425788, 0.0001117472565965727, 0.011082900688052177, 0.4118730425834656, 0.0024717452470213175, 0.21560189127922058, 0.015253315679728985, 0.03452993184328079, 0.13817672431468964, 0.0034516772720962763, 0.002911344636231661, 0.0003800573176704347, 0.001462712767533958, 0.001961951842531562, 0.0040230052545666695, 0.0023086154833436012, 0.002483226591721177, 0.028553131967782974, 0.014239847660064697, 0.18359807133674622, 0.09542248398065567, 0.2067933827638626, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00568122835829854, 0.003583817044273019, 0.0009402501164004207, 0.0034319525584578514, 0.014700439758598804, 0.00014027200813870877, 5.928567406954244e-05, 0.0005310353590175509, 0.001004774123430252, 0.00433507701382041, 0.003991644363850355, 0.0015378128737211227, 6.231402221601456e-05, 0.02625701017677784, 0.15481357276439667, 0.14011409878730774, 0.01466476172208786, 0.09487155824899673, 0.03769487887620926, 0.062972791492939, 0.003495296463370323, 0.0004466120735742152, 0.0044098952785134315, 0.056031279265880585, 0.12585759162902832, 0.04736572876572609, 0.02727479301393032, 0.06542934477329254, 0.563940703868866, 0.024195805191993713, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00503728911280632, 0.004739185329526663, 0.021364033222198486, 0.04603096470236778, 0.004565324168652296, 0.021244995296001434, 0.07592181116342545, 0.027910754084587097, 0.008603491820394993, 0.004941265098750591, 0.03103908710181713, 0.035909827798604965, 0.01818632334470749, 0.04406380280852318, 0.17931725084781647, 0.05395817384123802, 6.747527368133888e-05, 0.0018676340114325285, 0.0002809480356518179, 0.03275269269943237, 0.005758063402026892, 9.199039777740836e-05, 0.00011598093260545284, 0.0015754709020256996, 0.026104740798473358, 0.009686414152383804, 0.001081737456843257, 0.0017741151386871934, 0.49180474877357483, 0.007121484261006117, 0.013531914912164211, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21416018903255463, 0.005411786492913961, 0.02111194096505642, 0.07001130282878876, 0.04736214876174927, 0.09187527745962143, 0.1399366855621338, 0.030981194227933884, 0.02342112548649311, 0.07424263656139374, 0.02716991677880287, 0.5710572600364685, 0.007255392149090767, 0.005560784600675106, 0.054831843823194504, 0.03839295729994774, 0.0002068357716780156, 0.006204192526638508, 0.0054313126020133495, 0.011207946576178074, 0.0013116636546328664, 0.008276019245386124, 0.002269806107506156, 0.004080863669514656, 0.01488969475030899, 0.0006726597202941775, 0.009391524828970432, 0.039596475660800934, 0.19840312004089355, 0.043704546988010406, 0.31202515959739685, 0.23529505729675293, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3339015245437622, 0.03176174685359001, 0.25991618633270264, 0.31748515367507935, 0.17923809587955475, 0.2977932095527649, 0.14185847342014313, 0.09826549887657166, 0.4168005883693695, 0.09961694478988647, 0.1390676498413086, 0.191667839884758, 0.0443519689142704, 0.10075851529836655, 0.08045557886362076, 0.07469534128904343, 0.001304430770687759, 0.0239309910684824, 0.008060658350586891, 0.021029237657785416, 0.015191669575870037, 0.006979105528444052, 0.0016427322989329696, 0.002132130553945899, 0.015241370536386967, 0.0018563566263765097, 0.035101406276226044, 0.06515936553478241, 0.27313047647476196, 0.10352547466754913, 0.2570805549621582, 0.45083746314048767, 0.1295340657234192, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.018510108813643456, 0.0015040059806779027, 0.011199833825230598, 0.021222928538918495, 0.02421635016798973, 0.004175371024757624, 0.0007807075162418187, 0.0005349562270566821, 0.0038052168674767017, 0.3727143108844757, 0.022828511893749237, 0.01009275484830141, 0.0012628438416868448, 0.0009096930734813213, 0.10904579609632492, 0.19253067672252655, 0.0008209676598198712, 0.004669400863349438, 0.00047802351764403284, 0.013135433197021484, 0.0034620855003595352, 0.0016354827675968409, 0.0008273401763290167, 0.0018895546672865748, 0.009773027151823044, 0.006215384230017662, 0.2356690764427185, 0.01036232803016901, 0.06144833192229271, 0.008870624005794525, 0.024212215095758438, 0.008509873412549496, 0.01347219105809927, 0.35532569885253906, NaN, NaN, NaN, NaN, NaN, NaN], [0.05896773934364319, 0.023542853072285652, 0.0776505172252655, 0.15385140478610992, 0.011508575640618801, 0.0939982458949089, 0.0018089915392920375, 0.0003290986060164869, 0.0005636389250867069, 0.029514340683817863, 0.35146546363830566, 0.007090898230671883, 0.012099701911211014, 0.006742698606103659, 0.052738532423973083, 0.10910779982805252, 0.002221200615167618, 0.0001436042075511068, 1.1848528629343491e-05, 0.0001887700636871159, 0.0020721519831568003, 0.0009632316650822759, 0.00014056939107831568, 0.0007320817094296217, 0.0006829273188486695, 0.007395991589874029, 0.2889891564846039, 0.007074101362377405, 0.0002627878566272557, 0.004363438580185175, 0.0018575063440948725, 0.00557676050812006, 0.012322820723056793, 0.31134024262428284, 0.027276715263724327, NaN, NaN, NaN, NaN, NaN], [0.18205131590366364, 0.00472951028496027, 0.03192766383290291, 0.059333182871341705, 0.028221452608704567, 0.033883631229400635, 0.00131422549020499, 0.0001085989861167036, 5.632251122733578e-05, 0.004554648417979479, 0.2950275242328644, 0.0014449548907577991, 0.2329740822315216, 0.0520821250975132, 0.1361607313156128, 0.18170765042304993, 0.003209297079592943, 0.0023912524338811636, 0.00020479358499869704, 0.0009326079743914306, 0.0013757160631939769, 0.0021110770758241415, 0.0008730489062145352, 0.000792569131590426, 0.01825624145567417, 0.0059272306971251965, 0.11984144151210785, 0.05654650926589966, 0.08423373848199844, 0.024963613599538803, 0.027966396883130074, 0.1777324080467224, 0.005578523967415094, 0.14623191952705383, 0.11331525444984436, 0.2157108038663864, NaN, NaN, NaN, NaN], [0.0063572716899216175, 0.002779513830319047, 0.0009721479145810008, 0.0035897656343877316, 0.019835324957966805, 0.00021187934908084571, 8.435463678324595e-05, 0.00043589723645709455, 0.0004945950931869447, 0.004414541646838188, 0.0027602717746049166, 0.0008482423145323992, 5.171148222871125e-05, 0.021799515932798386, 0.15211130678653717, 0.1515214741230011, 0.008395697921514511, 0.0657893642783165, 0.019086696207523346, 0.05097401514649391, 0.0016111076110973954, 0.00021851839846931398, 0.002003778237849474, 0.01669292151927948, 0.06321260333061218, 0.015100682154297829, 0.010209205560386181, 0.015906400978565216, 0.30131736397743225, 0.012282183393836021, 0.09666845202445984, 0.00808996893465519, 0.03798958286643028, 0.013879657723009586, 0.047733187675476074, 0.5371345281600952, 0.020763304084539413, NaN, NaN, NaN], [0.005286877974867821, 0.008391096256673336, 0.025823507457971573, 0.030178312212228775, 0.00857502967119217, 0.042816706001758575, 0.07608389109373093, 0.03679429367184639, 0.0067360359244048595, 0.0038807345554232597, 0.03710461035370827, 0.037315309047698975, 0.018847206607460976, 0.0415174663066864, 0.15352587401866913, 0.07945924997329712, 4.7485355025855824e-05, 0.0020416006445884705, 0.00022757358965463936, 0.013386114500463009, 0.001981395063921809, 3.6917605029884726e-05, 2.620528539409861e-05, 0.0003202208608854562, 0.009042860940098763, 0.0030785591807216406, 0.0011855574557557702, 0.0005728560499846935, 0.20002734661102295, 0.00213914574123919, 0.002927121240645647, 0.004968173801898956, 0.0065933396108448505, 0.002585601294413209, 0.002817549044266343, 0.547335147857666, 0.006171087268739939, 0.018697692081332207, NaN, NaN], [0.2992006242275238, 0.008802352473139763, 0.027079692110419273, 0.08564624935388565, 0.11560814827680588, 0.22971339523792267, 0.1826445311307907, 0.033842965960502625, 0.06175734102725983, 0.11205370724201202, 0.04016120731830597, 0.5851526856422424, 0.016921253874897957, 0.011652404442429543, 0.08951538056135178, 0.059381648898124695, 0.00026094831991940737, 0.007586375344544649, 0.006061093881726265, 0.0039266073144972324, 0.0004965912085026503, 0.003665223019197583, 0.0008195870905183256, 0.0014654117403551936, 0.0045553394593298435, 0.00032001128420233727, 0.004615657962858677, 0.017150992527604103, 0.07922492176294327, 0.012805018573999405, 0.1320599913597107, 0.09461667388677597, 0.003555287839844823, 0.019601207226514816, 0.047796737402677536, 0.29085052013397217, 0.04383813217282295, 0.32529252767562866, 0.24933147430419922, NaN], [0.12446854263544083, 0.0009617851465009153, 0.004788657650351524, 0.0008746102685108781, 0.16037316620349884, 0.003065474098548293, 0.0056405095383524895, 0.005250739399343729, 0.05696318671107292, 0.013819074258208275, 0.028642717748880386, 0.0011808956041932106, 0.08446037769317627, 0.03008313849568367, 0.13710428774356842, 0.13618361949920654, 0.0007103006355464458, 0.025071904063224792, 0.004419561009854078, 0.001962232170626521, 0.0023795748129487038, 0.002366183791309595, 0.0003890783409588039, 0.00022811641974840313, 0.0010611300822347403, 0.001608739490620792, 0.028126444667577744, 0.005591525696218014, 0.0024579197634011507, 0.004123267717659473, 0.0409882515668869, 0.010364435613155365, 0.010518459603190422, 0.09771004319190979, 0.037823982536792755, 0.019979961216449738, 0.018303534016013145, 0.22492042183876038, 0.09256016463041306, 0.005498841404914856]], [[0.09139528125524521, 0.1232069656252861, 0.06926427036523819, 0.03596228361129761, 0.08677947521209717, 0.3523865342140198, 0.17220446467399597, 0.3048216700553894, 0.24129998683929443, 0.008230631239712238, 0.012852879241108894, 0.0024019270204007626, 0.003931952640414238, 0.002576343482360244, 0.13348431885242462, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.005495021585375071, 0.009821278043091297, 0.006606503389775753, 0.0009270968730561435, 0.022634856402873993, 0.02637101709842682, 0.03666122257709503, 0.003247066168114543, 0.03138025477528572, 0.0023785934317857027, 0.007012520916759968, 0.0027185468934476376, 0.001623710268177092, 0.009003029204905033, 0.24841202795505524, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.004891206510365009, 0.01856830157339573, 0.01660238206386566, 0.05400720611214638, 0.2678459584712982, 0.21548990905284882, 0.0901486948132515, 0.14165979623794556, 0.4387242794036865, 0.0060303402133286, 0.03774549812078476, 0.022296983748674393, 0.014843892306089401, 0.003844154067337513, 0.0701230987906456, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.009136357344686985, 0.005524215288460255, 0.002000550739467144, 0.004360574297606945, 0.06230698525905609, 0.032116882503032684, 0.14447683095932007, 0.11250873655080795, 0.12456412613391876, 0.017903752624988556, 0.03641437739133835, 0.030236193910241127, 0.03817100450396538, 0.0020203718449920416, 0.24235397577285767, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.011458649300038815, 0.0028747334145009518, 0.0048751854337751865, 0.0034302298445254564, 0.032581884413957596, 0.009492963552474976, 0.29646721482276917, 0.024549754336476326, 0.5199102163314819, 0.07497825473546982, 0.039336495101451874, 0.23366358876228333, 0.2855432629585266, 0.0047793262638151646, 0.131587415933609, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0048281243070960045, 0.014400148764252663, 0.00021499136346392334, 0.00015902110317256302, 0.0008502291166223586, 0.005816742777824402, 0.03721616789698601, 0.31765323877334595, 0.006985681131482124, 9.90723492577672e-05, 0.0015535155544057488, 0.002471775049343705, 0.00966054666787386, 0.002636645222082734, 0.15553238987922668, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01824354939162731, 0.02838711440563202, 0.0006440957658924162, 0.00040316785452887416, 0.00041587575105950236, 0.0021029487252235413, 0.07766012847423553, 0.3384210765361786, 0.005884509067982435, 0.02229108288884163, 0.02292727865278721, 0.00326070049777627, 0.002748187631368637, 0.004811563994735479, 0.08466839045286179, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0009052195237018168, 0.00028935770387761295, 0.00010135041520697996, 4.4237076508579776e-05, 9.765469440026209e-05, 0.0003226006228942424, 0.0006174442823976278, 0.003764552064239979, 0.001191335148178041, 0.0005841490346938372, 0.001988127361983061, 0.0019700597040355206, 0.0006354944198392332, 0.0011416736524552107, 0.25631290674209595, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.007226317655295134, 0.015471585094928741, 0.027516253292560577, 0.0063530029729008675, 0.015222059562802315, 0.004327190574258566, 0.010739101096987724, 0.0023785619996488094, 0.053105201572179794, 0.0674574077129364, 0.31870341300964355, 0.4986713230609894, 0.027042971923947334, 0.0736011192202568, 0.116986483335495, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.015794623643159866, 0.009404269978404045, 0.017993446439504623, 0.003823975333943963, 0.004969433881342411, 0.03679484874010086, 0.04242165759205818, 0.017222637310624123, 0.1201641708612442, 0.016131659969687462, 0.3518509864807129, 0.3061373829841614, 0.0458594486117363, 0.15943044424057007, 0.17968055605888367, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.006380036938935518, 0.028477374464273453, 0.006851766724139452, 0.005024573765695095, 0.02579522877931595, 0.052536945790052414, 0.0111169358715415, 0.0038714397232979536, 0.008046599105000496, 0.008921324275434017, 0.011395278386771679, 0.10255969315767288, 0.21638940274715424, 0.44467252492904663, 0.05895284563302994, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.010142950341105461, 0.001643709372729063, 0.002422438468784094, 0.0009472724632360041, 0.0033483330626040697, 0.003415578044950962, 0.03889569267630577, 0.005287462379783392, 0.00042015319922938943, 0.0010667687747627497, 0.00740370387211442, 0.00895014964044094, 0.0067735291086137295, 0.017782215029001236, 0.26753443479537964, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.11724554747343063, 0.0023070531897246838, 0.004510094877332449, 0.0014967885799705982, 0.007825762964785099, 0.00018500315491110086, 0.013543304987251759, 0.0012864026939496398, 0.0007778326398693025, 0.00044295378029346466, 0.001640060218051076, 0.0014512997586280107, 0.002360806567594409, 0.2112705558538437, 0.19457924365997314, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.09882069379091263, 0.014871560037136078, 0.005077258683741093, 0.0014827846316620708, 0.005620975513011217, 0.0024449406191706657, 0.07368315756320953, 0.06950978189706802, 0.0017206794582307339, 0.00039900749106891453, 0.0006052122334949672, 0.0005968212499283254, 0.004762541502714157, 0.0232950821518898, 0.2500154376029968, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.001020739320665598, 0.001402992638759315, 0.0006185534875839949, 0.0003395593084860593, 0.0013021298218518496, 0.0008022591937333345, 0.003452729433774948, 0.0026675688568502665, 0.0021077031269669533, 0.0008018113439902663, 0.0017594166565686464, 0.0005115982494316995, 0.0007778447470627725, 0.0008368113776668906, 0.13888627290725708, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.005261753685772419, 0.005328452680259943, 0.1075906753540039, 0.007504252251237631, 0.18196941912174225, 0.2677178680896759, 0.18533208966255188, 0.041308093816041946, 0.04052837938070297, 0.0018225060775876045, 0.004738607443869114, 0.028365809470415115, 0.07867489755153656, 0.032602421939373016, 0.14697469770908356, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.024903474375605583, 0.2637169063091278, 0.01148936152458191, 0.01806865818798542, 0.010384032502770424, 0.05497525632381439, 0.01011874619871378, 6.159161421237513e-05, 0.03404803201556206, 0.01315199863165617, 0.004086918197572231, 0.033981483429670334, 0.0007253359071910381, 0.0010365481721237302, 0.023150891065597534, 0.11621169000864029, 0.2792567312717438, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03176039457321167, 0.002004105830565095, 0.011469452641904354, 0.003235333366319537, 0.011606591753661633, 0.01332010142505169, 0.007885226979851723, 0.0010319099528715014, 0.0026684575714170933, 0.003885145066305995, 0.002207087352871895, 0.010414022952318192, 0.015553043223917484, 0.01973811537027359, 0.1639232188463211, 0.16788142919540405, 0.08717074245214462, 0.024576181545853615, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24842531979084015, 0.031220050528645515, 0.028132880106568336, 0.029530569911003113, 0.01766534335911274, 0.36354437470436096, 0.06892471760511398, 0.02528339996933937, 0.01102821622043848, 0.15825842320919037, 0.13755246996879578, 0.07390110194683075, 0.19022952020168304, 0.1824880689382553, 0.1432848572731018, 0.14762163162231445, 0.09094145894050598, 0.023598572239279747, 0.2273045778274536, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0013664831640198827, 0.001714985934086144, 0.0013615208445116878, 0.0015855998499318957, 0.0011547008762136102, 0.007221538573503494, 0.01537399459630251, 0.020302001386880875, 0.0011185031617060304, 0.001242821803316474, 0.0004577837826218456, 0.0013307477347552776, 6.100967220845632e-05, 3.943840420106426e-05, 0.16435295343399048, 0.10424397885799408, 0.7145561575889587, 0.21233327686786652, 0.5272893309593201, 0.04291817173361778, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0006725311395712197, 0.000846685899887234, 0.001614874112419784, 0.000348375499015674, 0.0019150535808876157, 0.01370947528630495, 0.026421356946229935, 0.08118636161088943, 0.0008913385099731386, 0.0004401778569445014, 0.0003709472657646984, 0.0007744845934212208, 0.002328733913600445, 0.0003664834948722273, 0.14579549431800842, 0.11001076549291611, 0.4734446108341217, 0.06134912371635437, 0.2925608456134796, 0.02150837518274784, 0.19962187111377716, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.011207095347344875, 0.029191432520747185, 0.015348215587437153, 0.012354064732789993, 0.002485303906723857, 0.7150441408157349, 0.0764552503824234, 0.14450958371162415, 0.0016117440536618233, 0.008765846490859985, 0.011787951923906803, 0.002862851833924651, 0.022502094507217407, 0.007210019044578075, 0.007054056040942669, 0.17212024331092834, 0.1419786959886551, 0.05631781369447708, 0.2185172289609909, 0.002532752463594079, 0.0032626313623040915, 0.18381445109844208, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.006926322355866432, 0.0050496323965489864, 0.010020078159868717, 0.021360181272029877, 0.0027102867607027292, 0.028520535677671432, 0.05918040871620178, 0.23060235381126404, 0.019199691712856293, 0.09477535635232925, 0.013206732459366322, 0.0014817069750279188, 0.0153219448402524, 0.01803957298398018, 0.07950127124786377, 0.09107878059148788, 0.12160263955593109, 0.2150201052427292, 0.3705081045627594, 0.07164584845304489, 0.05021890252828598, 0.14392021298408508, 0.39638784527778625, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.009242678992450237, 0.05580667033791542, 0.014326682314276695, 0.04630666971206665, 0.010674487799406052, 0.5850453972816467, 0.4108324944972992, 0.4116209149360657, 0.007144990377128124, 0.20661039650440216, 0.037308260798454285, 0.054067905992269516, 0.037599414587020874, 0.03113422356545925, 0.22261686623096466, 0.2121918499469757, 0.20806513726711273, 0.15205760300159454, 0.38131871819496155, 0.1009124368429184, 0.09936784207820892, 0.07077471911907196, 0.05006752535700798, 0.14871110022068024, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0023711349349468946, 0.019731320440769196, 0.027566438540816307, 0.03758935630321503, 0.022646954283118248, 0.06538618355989456, 0.01152126956731081, 0.014797273091971874, 0.003413880243897438, 0.024214325472712517, 0.019466044381260872, 0.007235943805426359, 0.0008611958473920822, 0.0011126803001388907, 0.268255352973938, 0.21685828268527985, 0.23333710432052612, 0.06609098613262177, 0.12803798913955688, 0.1004808098077774, 0.025170300155878067, 0.04069148004055023, 0.10828333348035812, 0.10351972281932831, 0.29450517892837524, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08772679418325424, 0.02003292553126812, 0.09465871006250381, 0.41126132011413574, 0.07995565980672836, 0.5143890976905823, 0.1155472919344902, 0.01320470031350851, 0.02149844542145729, 0.06702866405248642, 0.6884661316871643, 0.09638151526451111, 0.35587188601493835, 0.2170087993144989, 0.019593046978116035, 0.05205162987112999, 0.22306090593338013, 0.049221184104681015, 0.061203524470329285, 0.09776578843593597, 0.06183243915438652, 0.17444021999835968, 0.321644127368927, 0.054029058665037155, 0.2629997134208679, 0.2757931053638458, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01343127153813839, 0.0019279895350337029, 0.01925632171332836, 0.04226915165781975, 0.005290344823151827, 0.5555825233459473, 0.06846548616886139, 0.006453313864767551, 0.019162334501743317, 0.0017575293313711882, 0.2967261075973511, 0.11721283942461014, 0.4438721835613251, 0.1899448037147522, 0.007863422855734825, 0.05800137668848038, 0.32540804147720337, 0.13333332538604736, 0.05756821855902672, 0.12640602886676788, 0.11846329271793365, 0.2918737828731537, 0.3632459342479706, 0.18816226720809937, 0.6433262228965759, 0.3291742205619812, 0.12170911580324173, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12789316475391388, 0.004323228262364864, 0.03538274019956589, 0.05581461265683174, 0.020947236567735672, 0.09860846400260925, 0.11394336074590683, 0.010361305437982082, 0.011101406998932362, 0.33580121397972107, 0.13689599931240082, 0.038663506507873535, 0.19725953042507172, 0.10533706098794937, 0.008538279682397842, 0.11078674346208572, 0.40781712532043457, 0.06261185556650162, 0.05779192969202995, 0.18194560706615448, 0.1120922714471817, 0.5645142793655396, 0.33037880063056946, 0.18058234453201294, 0.6155731678009033, 0.21430827677249908, 0.044265877455472946, 0.20548948645591736, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.007053391542285681, 0.012331487610936165, 0.008611395955085754, 0.031008008867502213, 0.004283395130187273, 0.0029549654573202133, 0.00849387887865305, 0.008564120158553123, 0.02629040740430355, 0.009985123760998249, 0.00761940935626626, 0.003499145619571209, 0.0015691317385062575, 0.005600257311016321, 0.5214234590530396, 0.08288691937923431, 0.2962968051433563, 0.2819015085697174, 0.19574381411075592, 0.1136796846985817, 0.07755676656961441, 0.20596812665462494, 0.3330870270729065, 0.21944326162338257, 0.22804425656795502, 0.1688224822282791, 0.2872299253940582, 0.13759873807430267, 0.09907422959804535, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0007030746201053262, 0.0001308645587414503, 0.0001913319865707308, 0.00016671058256179094, 0.000299752748105675, 0.0001608166057849303, 0.004501530434936285, 0.0010771069210022688, 0.003937124740332365, 0.001599485520273447, 0.0007339937728829682, 0.0030779645312577486, 3.4502605558373034e-05, 9.700484952190891e-05, 0.15641583502292633, 0.11118441820144653, 0.6110438108444214, 0.6292654871940613, 0.5805363655090332, 0.22765980660915375, 0.4274957776069641, 0.6573506593704224, 0.6816673278808594, 0.5361799597740173, 0.320940226316452, 0.3845328688621521, 0.6242536306381226, 0.41633498668670654, 0.12922972440719604, 0.01991792768239975, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.027913473546504974, 0.10055015236139297, 0.005828284192830324, 0.007361504249274731, 0.0010143647668883204, 0.000654859293717891, 0.0101061025634408, 0.029607031494379044, 0.04485415667295456, 0.09235014766454697, 0.05163425952196121, 0.03075464628636837, 0.027050884440541267, 0.021472401916980743, 0.18064866960048676, 0.10675505548715591, 0.1912444829940796, 0.23975566029548645, 0.32351911067962646, 0.046362437307834625, 0.08004549145698547, 0.3363644778728485, 0.2706483006477356, 0.26792168617248535, 0.2952979505062103, 0.4496033787727356, 0.1126319095492363, 0.5116660594940186, 0.015820369124412537, 0.030236991122364998, 0.03603934869170189, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0011193754617124796, 0.03864011913537979, 0.0033454783260822296, 0.0006957795703783631, 0.001480268081650138, 0.0012079592561349273, 0.00020605533791240305, 0.0011212154058739543, 0.0015670693246647716, 0.0014121911954134703, 0.0012700740480795503, 0.0019415348069742322, 0.001359732006676495, 0.0011440571397542953, 0.23876120150089264, 0.2233639359474182, 0.0911012589931488, 0.12918633222579956, 0.17958812415599823, 0.037158817052841187, 0.06043876335024834, 0.43303725123405457, 0.3349981904029846, 0.09061599522829056, 0.23225362598896027, 0.1514965295791626, 0.09056703746318817, 0.2480165809392929, 0.056160230189561844, 0.015552842989563942, 0.007365798112004995, 0.17054231464862823, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.012943120673298836, 0.020876264199614525, 0.04825761169195175, 0.03707631304860115, 0.015636419877409935, 0.11923719942569733, 0.021652603521943092, 0.026653259992599487, 0.020431919023394585, 0.03287035599350929, 0.10921605676412582, 0.11103712767362595, 0.08490956574678421, 0.05352960154414177, 0.1791488379240036, 0.09585364907979965, 0.22669152915477753, 0.08040254563093185, 0.0638674795627594, 0.15364862978458405, 0.13237975537776947, 0.3887532651424408, 0.5357696413993835, 0.07155110687017441, 0.4139500856399536, 0.05426981300115585, 0.1238613948225975, 0.07816720753908157, 0.14353296160697937, 0.021915707737207413, 0.02897939831018448, 0.22262324392795563, 0.4835837185382843, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.010143280029296875, 0.0011783033842220902, 0.07699523866176605, 0.04151652753353119, 0.013031265698373318, 0.6595657467842102, 0.04001229628920555, 0.015414847061038017, 0.05828738585114479, 0.00582495890557766, 0.39538952708244324, 0.3540988564491272, 0.5535411834716797, 0.14920510351657867, 0.05510678142309189, 0.05190133675932884, 0.3522363007068634, 0.14802464842796326, 0.07656959444284439, 0.12417534738779068, 0.17628712952136993, 0.33604755997657776, 0.38481405377388, 0.20552395284175873, 0.5797679424285889, 0.3262830972671509, 0.19466114044189453, 0.045280374586582184, 0.2712458372116089, 0.041196610778570175, 0.08666794002056122, 0.3327068090438843, 0.1922111064195633, 0.10969121754169464, NaN, NaN, NaN, NaN, NaN, NaN], [0.10365689545869827, 0.011393263004720211, 0.09083462506532669, 0.05552159622311592, 0.021694108843803406, 0.23093751072883606, 0.12655670940876007, 0.02638416364789009, 0.016898566856980324, 0.4334920644760132, 0.1302367001771927, 0.07987051457166672, 0.26015403866767883, 0.07882147282361984, 0.06412448734045029, 0.10818891227245331, 0.3937702178955078, 0.030490810051560402, 0.030189264565706253, 0.11243001371622086, 0.07142115384340286, 0.3648340702056885, 0.2467786818742752, 0.13009557127952576, 0.5037410855293274, 0.18716548383235931, 0.08825942128896713, 0.23451530933380127, 0.24434491991996765, 0.03496113047003746, 0.04431905224919319, 0.3934983015060425, 0.31427451968193054, 0.05462265387177467, 0.2524711489677429, NaN, NaN, NaN, NaN, NaN], [0.0009046280756592751, 0.006186267826706171, 0.001710598124191165, 0.0040000369772315025, 0.0010556421475484967, 0.00010012275743065402, 0.000467440317152068, 0.00034073027200065553, 0.012450831942260265, 0.001776019111275673, 0.0016348852077499032, 0.0004490323772188276, 0.00023723821504972875, 0.0005369102582335472, 0.2610536217689514, 0.06088699772953987, 0.23725801706314087, 0.2046121060848236, 0.14171433448791504, 0.06688592582941055, 0.06064169481396675, 0.14286598563194275, 0.21723276376724243, 0.13491223752498627, 0.2083195000886917, 0.15285742282867432, 0.34066644310951233, 0.18166381120681763, 0.10532425343990326, 0.06318715214729309, 0.052211396396160126, 0.20970472693443298, 0.20715771615505219, 0.28281068801879883, 0.13935938477516174, 0.11923542618751526, NaN, NaN, NaN, NaN], [0.00040706052095629275, 5.995776882627979e-05, 0.00011266738147241995, 0.00010974665929097682, 0.00022393744438886642, 7.468188414350152e-05, 0.00239625689573586, 0.0004222780407872051, 0.002755024004727602, 0.0011263962369412184, 0.0004159261588938534, 0.0013214137870818377, 1.3015362128498964e-05, 3.146446033497341e-05, 0.15343648195266724, 0.09884612262248993, 0.5530695915222168, 0.6301063299179077, 0.5187459588050842, 0.28427499532699585, 0.33059176802635193, 0.49595603346824646, 0.6107674241065979, 0.387560099363327, 0.3283739984035492, 0.3905918300151825, 0.5949583053588867, 0.2912430167198181, 0.19163259863853455, 0.03091937117278576, 0.3911139667034149, 0.3233675956726074, 0.421701043844223, 0.6310504674911499, 0.4068542718887329, 0.13317596912384033, 0.02126597985625267, NaN, NaN, NaN], [0.02487853355705738, 0.06922142952680588, 0.005931189749389887, 0.005149703938513994, 0.0007503133383579552, 0.00046759017277508974, 0.004864065907895565, 0.010271446779370308, 0.03885169327259064, 0.0494176521897316, 0.032662954181432724, 0.015474021434783936, 0.005468437913805246, 0.0031831569503992796, 0.16160887479782104, 0.07192745804786682, 0.09934075176715851, 0.15662430226802826, 0.18248029053211212, 0.021172231063246727, 0.037516966462135315, 0.12766626477241516, 0.09711621701717377, 0.09662153571844101, 0.1303528994321823, 0.3114719092845917, 0.1600099802017212, 0.265144020318985, 0.011710498481988907, 0.02471126988530159, 0.012725233100354671, 0.12533646821975708, 0.446529746055603, 0.11092787981033325, 0.45893827080726624, 0.011159577406942844, 0.028070949018001556, 0.024378135800361633, NaN, NaN], [0.0006016235565766692, 0.010655699297785759, 0.0012552555417641997, 0.0004406629304867238, 0.0006771506741642952, 0.0004804672207683325, 8.584682655055076e-05, 0.00018533790716901422, 0.0020008538849651814, 0.0008522755815647542, 0.0005471827462315559, 0.0006654397584497929, 0.0003326669684611261, 0.00020969027536921203, 0.18202657997608185, 0.21178482472896576, 0.0713806003332138, 0.12116114795207977, 0.16551871597766876, 0.025692136958241463, 0.03932836279273033, 0.255863755941391, 0.20887790620326996, 0.05500240623950958, 0.14075487852096558, 0.158308207988739, 0.10016348958015442, 0.22940821945667267, 0.06542190909385681, 0.016673747450113297, 0.011679067276418209, 0.21266934275627136, 0.27460965514183044, 0.08977667987346649, 0.1985965520143509, 0.05640871822834015, 0.014301197603344917, 0.004748867359012365, 0.1251523643732071, NaN], [0.0006660889484919608, 0.0011989487102255225, 0.006168409250676632, 0.0007392434636130929, 0.002072105184197426, 0.0013732375809922814, 0.001215140800923109, 8.942947169998661e-05, 0.0032219376880675554, 0.00034276655060239136, 0.0006051870877854526, 0.0004003554640803486, 0.0006330502219498158, 9.228585986420512e-05, 0.13989190757274628, 0.11377177387475967, 0.4656391441822052, 0.26672884821891785, 0.20802536606788635, 0.1860857605934143, 0.16829806566238403, 0.19711202383041382, 0.3023360073566437, 0.035885076969861984, 0.11114621162414551, 0.21048156917095184, 0.27827921509742737, 0.11178875714540482, 0.13154125213623047, 0.3096882104873657, 0.09530708193778992, 0.2201821655035019, 0.1989239901304245, 0.27841058373451233, 0.15223632752895355, 0.2206900417804718, 0.34536775946617126, 0.09229245036840439, 0.24595825374126434, 0.2865155339241028]], [[0.04622220993041992, 0.12740419805049896, 0.05372706800699234, 0.5582705140113831, 0.030120277777314186, 0.3703221380710602, 0.020304178819060326, 0.3357560634613037, 0.11819478869438171, 0.0765489861369133, 0.09261158853769302, 0.03858334198594093, 0.13079233467578888, 0.0447748564183712, 0.11706516146659851, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0919138491153717, 0.05798470228910446, 0.02827676385641098, 0.34965166449546814, 0.05504997447133064, 0.1526506543159485, 0.09941896051168442, 0.4367760419845581, 0.061004042625427246, 0.5390062928199768, 0.28723591566085815, 0.15840129554271698, 0.2018149495124817, 0.11561664938926697, 0.1249081939458847, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.032068803906440735, 0.0549696609377861, 0.018587671220302582, 0.2202640324831009, 0.0011182812741026282, 0.03810814768075943, 0.027008401229977608, 0.3763306438922882, 0.11146998405456543, 0.16719762980937958, 0.13283231854438782, 0.014421377331018448, 0.07254088670015335, 0.007401765324175358, 0.20662666857242584, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10753453522920609, 0.479284405708313, 0.009764611721038818, 0.0431443527340889, 0.0008862981921993196, 0.03188035264611244, 0.00600279588252306, 0.43093177676200867, 0.08460848033428192, 0.18502341210842133, 0.038902610540390015, 0.030237559229135513, 0.1820157915353775, 0.03367093205451965, 0.14427724480628967, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.013928310945630074, 0.032752107828855515, 0.0024797581136226654, 0.10617181658744812, 0.0002726189268287271, 0.011333486996591091, 0.005626056343317032, 0.05421115458011627, 0.020341530442237854, 0.0548044852912426, 0.027503041550517082, 0.005752534605562687, 0.033552803099155426, 0.008454940281808376, 0.388910174369812, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.15046736598014832, 0.296213299036026, 0.044096194207668304, 0.05168119817972183, 0.02727358601987362, 0.04717152938246727, 0.0016543868696317077, 0.035376399755477905, 0.027143586426973343, 0.0870317667722702, 0.05812281742691994, 0.06705813109874725, 0.3147181272506714, 0.39039844274520874, 0.23394177854061127, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.14644725620746613, 0.5605929493904114, 0.11812092363834381, 0.5902084112167358, 0.021858595311641693, 0.10718227922916412, 0.007383488584309816, 0.019886687397956848, 0.06570647656917572, 0.10820640623569489, 0.1357717514038086, 0.025582531467080116, 0.077891044318676, 0.061965201050043106, 0.164744034409523, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.049012791365385056, 0.35138410329818726, 0.26388463377952576, 0.7301797866821289, 0.014552393928170204, 0.24720129370689392, 0.0041521950624883175, 0.07795857638120651, 0.014070906676352024, 0.04667593538761139, 0.1480453461408615, 0.010990227572619915, 0.20039354264736176, 0.17517414689064026, 0.0717916414141655, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.09980960935354233, 0.4834202826023102, 0.20237547159194946, 0.5161312222480774, 0.2011035680770874, 0.31254804134368896, 0.023049525916576385, 0.09284620732069016, 0.030714770779013634, 0.009841320104897022, 0.03625232353806496, 0.02249438874423504, 0.030981028452515602, 0.01249231118708849, 0.19809871912002563, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2242409735918045, 0.5898000001907349, 0.2996082305908203, 0.6961580514907837, 0.3950251638889313, 0.824604332447052, 0.0551396869122982, 0.5436567068099976, 0.06683327257633209, 0.03568824753165245, 0.060814060270786285, 0.00592254800722003, 0.012778226286172867, 0.017990900203585625, 0.1082865446805954, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03427329286932945, 0.7018846869468689, 0.18350760638713837, 0.5559015274047852, 0.03810380771756172, 0.7226935029029846, 0.05184842646121979, 0.881024181842804, 0.06315085291862488, 0.03384441137313843, 0.014913397841155529, 0.002015632577240467, 0.008405282162129879, 0.0011906703002750874, 0.2768104076385498, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.022437993437051773, 0.7336767315864563, 0.2893984615802765, 0.7315550446510315, 0.021726222708821297, 0.3247562646865845, 0.05117126554250717, 0.7097986340522766, 0.03149837628006935, 0.017582548782229424, 0.017906883731484413, 0.004864181391894817, 0.0014982494758442044, 0.0005988480988889933, 0.17147301137447357, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.279982328414917, 0.427709698677063, 0.4798988997936249, 0.811837911605835, 0.5607104301452637, 0.3233453035354614, 0.03364620357751846, 0.48738226294517517, 0.20507316291332245, 0.2806957960128784, 0.20560167729854584, 0.021487781777977943, 0.0051806773990392685, 0.018182942643761635, 0.10378202050924301, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.15081651508808136, 0.5779510736465454, 0.21354816854000092, 0.8126901984214783, 0.041816346347332, 0.5376638174057007, 0.02729017473757267, 0.45972490310668945, 0.1708957701921463, 0.17148789763450623, 0.06268936395645142, 0.0045938147231936455, 0.0036332160234451294, 0.0009066996863111854, 0.10311751067638397, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.009540104307234287, 0.03889232128858566, 0.016071060672402382, 0.08366316556930542, 0.004574422258883715, 0.029401082545518875, 0.00834547821432352, 0.0893266350030899, 0.14732055366039276, 0.09065960347652435, 0.14173488318920135, 0.042114999145269394, 0.004022075328975916, 0.003513866104185581, 0.1347859650850296, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.17597882449626923, 0.03865775838494301, 0.04927876219153404, 0.19269852340221405, 0.07631995528936386, 0.03202155977487564, 0.04315444082021713, 0.0381813645362854, 0.14437337219715118, 0.14268529415130615, 0.12548406422138214, 0.22065725922584534, 0.007455701474100351, 0.012540786527097225, 0.13194040954113007, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12168548256158829, 0.12690430879592896, 0.03319493681192398, 0.044549524784088135, 0.022643521428108215, 0.12293753027915955, 0.012858373112976551, 0.056580886244773865, 0.0409478023648262, 0.5390252470970154, 0.04499629884958267, 0.010665545240044594, 0.0012580851325765252, 0.0006077282596379519, 0.16003872454166412, 0.13124778866767883, 0.015335792675614357, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.004976227879524231, 0.0016218257369473577, 0.10218203067779541, 0.005807417444884777, 0.025330372154712677, 0.00805770605802536, 0.0010953968157991767, 0.007808555383235216, 0.03332183510065079, 0.01014297641813755, 0.0378553569316864, 0.0012688467977568507, 0.0070253219455480576, 0.006525768432766199, 0.1611432433128357, 0.19323189556598663, 0.005229663103818893, 0.005805561784654856, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.018298039212822914, 0.043392445892095566, 0.026758581399917603, 0.06685060262680054, 0.007846164517104626, 0.0070086256600916386, 0.0011090404586866498, 0.0016357558779418468, 0.015295942313969135, 0.022091375663876534, 0.08676162362098694, 0.0013220091350376606, 0.0007799563463777304, 0.0005145008908584714, 0.5814905166625977, 0.06695510447025299, 0.08997365087270737, 0.32878753542900085, 0.35321861505508423, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16791731119155884, 0.01838838867843151, 0.03170344606041908, 0.04746389389038086, 0.024931352585554123, 0.002624210435897112, 0.3320338726043701, 0.32248422503471375, 0.021048149093985558, 0.02857070416212082, 0.11922428011894226, 4.079664358869195e-05, 0.0002566495386417955, 0.0005197013379074633, 0.1538068950176239, 0.1452476531267166, 0.07996584475040436, 0.2002653181552887, 0.13149262964725494, 0.005022347904741764, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03376027196645737, 0.001082546659745276, 0.003266592975705862, 0.006257645785808563, 0.023632841184735298, 0.00021245618700049818, 0.033721838146448135, 0.15340450406074524, 0.009442711248993874, 0.006162047851830721, 0.09923229366540909, 0.0001386175281368196, 0.0008165750186890364, 0.0010916005121544003, 0.14602994918823242, 0.1274433135986328, 0.13577045500278473, 0.16066212952136993, 0.1959238052368164, 0.04180024936795235, 0.06788772344589233, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04221357777714729, 0.03857824206352234, 0.004161412362009287, 0.06419923156499863, 0.010648604482412338, 0.008165394887328148, 0.04070910066366196, 0.34736329317092896, 0.0012154168216511607, 0.1630050241947174, 0.07001504302024841, 0.0033116117119789124, 0.00023883172252681106, 0.00045473958016373217, 0.2740376889705658, 0.14809708297252655, 0.29017606377601624, 0.22457490861415863, 0.17088554799556732, 0.041788797825574875, 0.013634788803756237, 0.02984887920320034, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.007271567825227976, 0.0015110730892047286, 0.0014769553672522306, 0.0053740208968520164, 0.0038654205854982138, 0.0024983601178973913, 0.049697574228048325, 0.27208074927330017, 0.0006182760698720813, 0.014045008458197117, 0.00131281279027462, 0.00040628391434438527, 0.00037906834040768445, 0.0001199298130813986, 0.006693295668810606, 0.21402230858802795, 0.012405444867908955, 0.0014808804262429476, 0.0009161182679235935, 0.0035427443217486143, 0.0017166208708658814, 0.001927618752233684, 0.015056394040584564, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08829134702682495, 0.11286511272192001, 0.004967967513948679, 0.006996258161962032, 0.0014454894699156284, 0.006397548597306013, 0.01389994379132986, 0.27431485056877136, 0.0018983082845807076, 0.09154568612575531, 0.022492842748761177, 0.0017391144065186381, 0.000634143827483058, 4.5783879613736644e-05, 0.318096399307251, 0.10794443637132645, 0.13477572798728943, 0.046750620007514954, 0.03419584408402443, 0.30604344606399536, 0.11879221349954605, 0.08022946119308472, 0.11745522916316986, 0.21712547540664673, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02142007276415825, 0.007001234218478203, 0.00761477230116725, 0.018849696964025497, 0.010492328554391861, 0.01844215951859951, 0.008208145387470722, 0.01109394058585167, 0.006335548125207424, 0.01884968765079975, 0.01652243174612522, 0.016355833038687706, 0.0014795949682593346, 0.0011322565842419863, 0.27169719338417053, 0.06259628385305405, 0.21873348951339722, 0.248628169298172, 0.2344663441181183, 0.09133727103471756, 0.05752522125840187, 0.03945200890302658, 0.39403918385505676, 0.15040725469589233, 0.009099425747990608, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17013461887836456, 0.14343884587287903, 0.017679741606116295, 0.10850679129362106, 0.01231957133859396, 0.010847942903637886, 0.04900640249252319, 0.023357992991805077, 0.014735743403434753, 0.014097570441663265, 0.012582896277308464, 0.0010529988212510943, 0.00046457236749120057, 0.0006211225991137326, 0.5663455724716187, 0.06400181353092194, 0.3208324611186981, 0.5040323138237, 0.6282902359962463, 0.04389061778783798, 0.08030739426612854, 0.10539824515581131, 0.1485716998577118, 0.08085520565509796, 0.13963551819324493, 0.0947280004620552, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1586649864912033, 0.08337923884391785, 0.0181503314524889, 0.22676831483840942, 0.016727542504668236, 0.015186772681772709, 0.0050455182790756226, 0.00688449339941144, 0.025511443614959717, 0.20239992439746857, 0.024231791496276855, 0.0023393011651933193, 0.0011192933889105916, 0.0005647524958476424, 0.390881210565567, 0.0935494601726532, 0.3055664598941803, 0.46751275658607483, 0.6914730072021484, 0.12860655784606934, 0.15726737678050995, 0.2987912595272064, 0.1529359668493271, 0.062232255935668945, 0.041881486773490906, 0.03399288281798363, 0.026789270341396332, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3443087935447693, 0.28029316663742065, 0.23536846041679382, 0.34415915608406067, 0.11761639267206192, 0.006012732163071632, 0.008058828301727772, 0.005314267706125975, 0.013309409841895103, 0.09906232357025146, 0.10091385245323181, 0.018941059708595276, 0.025248508900403976, 0.014945760369300842, 0.7436007857322693, 0.012478480115532875, 0.051689472049474716, 0.7194163799285889, 0.8485123515129089, 0.006671697832643986, 0.03636787086725235, 0.05433559790253639, 0.01463489979505539, 0.0011851346353068948, 0.0010049004340544343, 0.012586181983351707, 0.0039429632015526295, 0.0029262336902320385, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0022638223599642515, 0.004991845227777958, 0.004655482713133097, 0.0007185174035839736, 0.0013901105849072337, 0.011776956729590893, 0.0005479936371557415, 0.00022604972764384001, 0.00024645475787110627, 0.009541304782032967, 0.011744895949959755, 0.0007132806931622326, 0.27867355942726135, 0.02834550105035305, 0.007979176938533783, 0.16095376014709473, 0.10161679983139038, 0.15561290085315704, 0.27214428782463074, 0.06339859217405319, 0.047669682651758194, 0.16775988042354584, 0.30333516001701355, 0.29585903882980347, 0.026492541655898094, 0.03390856087207794, 0.020966142416000366, 0.027538424357771873, 0.040642742067575455, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.024570701643824577, 0.00167787482496351, 0.004072254989296198, 0.00223688711412251, 0.007143567781895399, 0.00014352552534546703, 0.0004634522774722427, 0.0016921478090807796, 0.003620122792199254, 0.007754941936582327, 0.011850811541080475, 0.0027722271624952555, 9.3724018370267e-05, 0.02145184949040413, 0.15506701171398163, 0.1701768934726715, 0.015393235720694065, 0.0020776872988790274, 0.011533004231750965, 0.013215321116149426, 0.004845780786126852, 0.011772604659199715, 0.006262979004532099, 0.00390799343585968, 0.007256041280925274, 0.0014780729543417692, 0.007152961101382971, 0.1450572907924652, 0.009833375923335552, 0.004788131918758154, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01723022572696209, 0.08018677681684494, 0.007713299244642258, 0.004271229729056358, 0.0005464836140163243, 0.00456921337172389, 0.0031762931030243635, 0.009469777345657349, 0.000385247083613649, 0.01870143786072731, 0.033109456300735474, 0.004042719956487417, 0.004976211115717888, 0.005646048113703728, 0.19230251014232635, 0.27953270077705383, 0.3106633424758911, 0.3078516721725464, 0.2835734188556671, 0.23220741748809814, 0.10028243064880371, 0.059542566537857056, 0.10900203883647919, 0.24247398972511292, 0.19294817745685577, 0.04455278813838959, 0.032558612525463104, 0.2623904049396515, 0.04071282595396042, 0.07101175934076309, 0.01397540420293808, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.016216034069657326, 0.04777013510465622, 0.01620146818459034, 0.010810854844748974, 0.16034351289272308, 0.006931359879672527, 0.0032006967812776566, 0.032106515020132065, 0.0003033989341929555, 0.015325331129133701, 0.006036583799868822, 0.12791146337985992, 0.19952742755413055, 0.023708127439022064, 0.18307197093963623, 0.15828359127044678, 0.26215362548828125, 0.1828027367591858, 0.3383132517337799, 0.14976613223552704, 0.17187725007534027, 0.16098640859127045, 0.10713529586791992, 0.2253616452217102, 0.27887699007987976, 0.0991593673825264, 0.1987481713294983, 0.2010713517665863, 0.24892166256904602, 0.09143882989883423, 0.028894133865833282, 0.0226773452013731, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.014499284327030182, 0.035677529871463776, 0.009275808930397034, 0.01653297245502472, 0.006223962642252445, 0.0020693510305136442, 0.007680083625018597, 0.013822571374475956, 0.00040966575033962727, 0.0038025544490665197, 0.013774569146335125, 0.006069935858249664, 0.004488381557166576, 0.005977130029350519, 0.217429518699646, 0.08621957898139954, 0.39239373803138733, 0.32060059905052185, 0.6169360876083374, 0.04211895540356636, 0.07954877614974976, 0.28241875767707825, 0.1073535904288292, 0.10431969910860062, 0.28138864040374756, 0.05428503826260567, 0.29005417227745056, 0.2829020619392395, 0.1771886944770813, 0.12728992104530334, 0.029228007420897484, 0.09527892619371414, 0.030012397095561028, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03237156197428703, 0.013441890478134155, 0.0194883793592453, 0.09343220293521881, 0.05379915237426758, 0.004893247038125992, 0.0011929833563044667, 0.009432576596736908, 0.015330814756453037, 0.14898745715618134, 0.018398255109786987, 0.01228779274970293, 0.00492482166737318, 0.0038985873106867075, 0.2601524889469147, 0.10387677699327469, 0.28899070620536804, 0.34778735041618347, 0.5978891849517822, 0.08856049180030823, 0.11093756556510925, 0.2773001492023468, 0.1387036144733429, 0.05535874143242836, 0.040542375296354294, 0.057020239531993866, 0.08593740314245224, 0.3575255870819092, 0.1780063509941101, 0.03115975111722946, 0.05683879926800728, 0.20087137818336487, 0.022991398349404335, 0.024780578911304474, NaN, NaN, NaN, NaN, NaN, NaN], [0.08357361704111099, 0.18220724165439606, 0.10462122410535812, 0.08245989680290222, 0.03124452568590641, 0.002170282183215022, 0.0020384257659316063, 0.004550496581941843, 0.003485089400783181, 0.036062099039554596, 0.0278666652739048, 0.011443988420069218, 0.01760544627904892, 0.013599698431789875, 0.3874043822288513, 0.027872784063220024, 0.11975038051605225, 0.8484699726104736, 0.9221431016921997, 0.010032964870333672, 0.05817321315407753, 0.14408904314041138, 0.03149182349443436, 0.0027255630120635033, 0.003546576714143157, 0.054592132568359375, 0.03846639767289162, 0.0179138146340847, 0.04004756733775139, 0.0025625908747315407, 0.006073353346437216, 0.017890095710754395, 0.006128084380179644, 0.0035659971181303263, 0.005842072889208794, NaN, NaN, NaN, NaN, NaN], [0.001995340920984745, 0.011527596041560173, 0.005334027577191591, 0.0006887424970045686, 0.0023407095577567816, 0.00276917009614408, 0.00029977987287566066, 0.00012230046559125185, 0.00026578022516332567, 0.008239910937845707, 0.009819538332521915, 0.000393931899452582, 0.605858564376831, 0.08989311754703522, 0.011135715991258621, 0.21095024049282074, 0.16082847118377686, 0.2551726996898651, 0.40046265721321106, 0.07841236889362335, 0.05558479577302933, 0.20925307273864746, 0.4381427764892578, 0.47918838262557983, 0.07096414268016815, 0.11106863617897034, 0.09138666838407516, 0.1393880993127823, 0.1506565660238266, 0.07743309438228607, 0.06943798065185547, 0.09801105409860611, 0.017720624804496765, 0.015859564766287804, 0.029157793149352074, 0.0392736941576004, NaN, NaN, NaN, NaN], [0.021298440173268318, 0.001658836961723864, 0.004600299056619406, 0.0025729055050760508, 0.015332063660025597, 0.00017298871534876525, 0.0005721640191040933, 0.00186175387352705, 0.0037871075328439474, 0.009124312549829483, 0.01116581168025732, 0.0031747270841151476, 0.00012207991676405072, 0.029056062921881676, 0.15163807570934296, 0.17935752868652344, 0.014263968914747238, 0.0022281131241470575, 0.011617614887654781, 0.022433524951338768, 0.0047986325807869434, 0.013686214573681355, 0.007696506567299366, 0.004939754959195852, 0.012488129548728466, 0.002878576284274459, 0.013457567431032658, 0.23303280770778656, 0.030022362247109413, 0.013181640766561031, 0.027029545977711678, 0.010247751139104366, 0.0006795030203647912, 0.0032072996255010366, 0.1104368045926094, 0.006663828622549772, 0.003364446572959423, NaN, NaN, NaN], [0.020229021087288857, 0.11621151119470596, 0.015550180338323116, 0.006284819450229406, 0.0013723199954256415, 0.013658476993441582, 0.005685316864401102, 0.02063058130443096, 0.001440295367501676, 0.022225895896553993, 0.07092871516942978, 0.007373427972197533, 0.00771017000079155, 0.006927240639925003, 0.16024509072303772, 0.3113161623477936, 0.29550519585609436, 0.2834082841873169, 0.292662650346756, 0.1380799263715744, 0.055221766233444214, 0.0487985797226429, 0.10219268500804901, 0.25612032413482666, 0.2569950222969055, 0.10279092192649841, 0.16084249317646027, 0.5340818166732788, 0.10305190831422806, 0.16831228137016296, 0.03310799598693848, 0.10521702468395233, 0.008185362443327904, 0.02029210887849331, 0.2447529286146164, 0.0189062412828207, 0.051586367189884186, 0.011271311901509762, NaN, NaN], [0.014029471203684807, 0.02389930933713913, 0.011611595749855042, 0.012217668816447258, 0.2477317750453949, 0.006976675242185593, 0.0035841658245772123, 0.022232146933674812, 0.0018886715406551957, 0.01750483363866806, 0.005654812324792147, 0.10889071226119995, 0.19916927814483643, 0.022882532328367233, 0.16074435412883759, 0.21913117170333862, 0.2667233347892761, 0.15068072080612183, 0.2934513986110687, 0.11010763049125671, 0.11770202964544296, 0.1548316478729248, 0.10880382359027863, 0.19848009943962097, 0.2926469147205353, 0.17939361929893494, 0.38748762011528015, 0.38622626662254333, 0.4369211196899414, 0.14473943412303925, 0.11290202289819717, 0.11878126114606857, 0.013051117770373821, 0.18458649516105652, 0.15622372925281525, 0.14840805530548096, 0.06742489337921143, 0.01624887064099312, 0.028317920863628387, NaN], [0.0032621105201542377, 0.006088452413678169, 0.012619324028491974, 0.008848619647324085, 0.17461968958377838, 8.660123421577737e-05, 0.0006109846872277558, 0.0007747155614197254, 0.003163054818287492, 0.017787659540772438, 0.029563669115304947, 0.0032195982057601213, 0.013336165808141232, 0.013171130791306496, 0.1387031376361847, 0.13670727610588074, 0.11102687567472458, 0.008893890306353569, 0.008979070000350475, 0.01785319298505783, 0.008134939707815647, 0.02043774165213108, 0.030145585536956787, 0.014907605946063995, 0.021436721086502075, 0.020207075402140617, 0.10284662246704102, 0.06823904067277908, 0.04208305850625038, 0.03810393810272217, 0.04656955599784851, 0.025087369605898857, 0.005296032875776291, 0.07358870655298233, 0.057817310094833374, 0.033472564071416855, 0.02220221422612667, 0.01758744567632675, 0.012124869041144848, 0.052647966891527176]], [[0.009570755064487457, 0.005546795669943094, 0.006825579330325127, 0.033384330570697784, 0.3769712448120117, 0.15916845202445984, 0.5290282368659973, 0.24695992469787598, 0.2377869039773941, 0.0913546234369278, 0.07570143043994904, 0.06522544473409653, 0.12397455424070358, 0.2645682692527771, 0.1787039041519165, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0061562443152070045, 0.040286894887685776, 0.0029807272367179394, 0.016133036464452744, 0.1151214987039566, 0.07519882172346115, 0.10128971189260483, 0.046498823910951614, 0.04111110791563988, 0.11845260113477707, 0.08915312588214874, 0.10556784272193909, 0.16933780908584595, 0.3531811535358429, 0.21578538417816162, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.14712950587272644, 0.04435151070356369, 0.015454337000846863, 0.01427951455116272, 0.08342041075229645, 0.005383625626564026, 0.10468690097332001, 0.05861024558544159, 0.08666124939918518, 0.15304753184318542, 0.23543620109558105, 0.2374279797077179, 0.10751555860042572, 0.10399115085601807, 0.23440681397914886, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0859314426779747, 0.15731151401996613, 0.005385389551520348, 0.04620514437556267, 0.010708490386605263, 0.006711416877806187, 0.012445325031876564, 0.056288186460733414, 0.097142793238163, 0.07020799815654755, 0.02479076385498047, 0.0890590250492096, 0.22972674667835236, 0.034618109464645386, 0.28529092669487, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07441635429859161, 0.018118128180503845, 0.016377849504351616, 0.003080169903114438, 0.20936372876167297, 0.0007255859090946615, 0.03578657656908035, 0.00550744216889143, 0.1172742024064064, 0.5684130191802979, 0.3980042636394501, 0.15252694487571716, 0.10817506164312363, 0.23486874997615814, 0.2619861364364624, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05188249424099922, 0.0069924332201480865, 0.0009591103880666196, 0.0061192926950752735, 0.002253405749797821, 0.006572761107236147, 0.004667140077799559, 0.11107926070690155, 0.03415685519576073, 0.010113962925970554, 0.006655086297541857, 0.010832482948899269, 0.03651394695043564, 0.040573474019765854, 0.2686486840248108, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08095332235097885, 0.02014574408531189, 0.011188640259206295, 0.0037319576367735863, 0.024485761299729347, 0.0018746056593954563, 0.04114176332950592, 0.034570205956697464, 0.009728988632559776, 0.07755846530199051, 0.09898480027914047, 0.0613434873521328, 0.09528356045484543, 0.1511603444814682, 0.2821846306324005, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04335615411400795, 0.026033984497189522, 0.03572213277220726, 0.017578190192580223, 0.05956277251243591, 0.01715734601020813, 0.011929154396057129, 0.28936532139778137, 0.0027683174703270197, 0.061091482639312744, 0.23734883964061737, 0.10397756844758987, 0.16337142884731293, 0.37352773547172546, 0.18409839272499084, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.06077902019023895, 0.031166722998023033, 0.11759120225906372, 0.1409873068332672, 0.24215947091579437, 0.009796793572604656, 0.10265856236219406, 0.01014934666454792, 0.2757207751274109, 0.023714441806077957, 0.038815632462501526, 0.15303847193717957, 0.14991649985313416, 0.6824791431427002, 0.13190437853336334, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.06505369395017624, 0.006089756730943918, 0.036541152745485306, 0.005829536356031895, 0.20233574509620667, 0.029401954263448715, 0.49993017315864563, 0.030510973185300827, 0.01976127363741398, 0.07993583381175995, 0.017815636470913887, 0.04079095646739006, 0.022992853075265884, 0.6425142288208008, 0.26567763090133667, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6054520010948181, 0.07051455229520798, 0.2702813744544983, 0.029061302542686462, 0.13962645828723907, 0.07908772677183151, 0.4563634395599365, 0.02414957620203495, 0.02722080610692501, 0.03215296193957329, 0.015534932725131512, 0.009437407366931438, 0.0218642745167017, 0.08506882190704346, 0.4000338017940521, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3943043351173401, 0.11258544027805328, 0.12088752537965775, 0.0732470229268074, 0.030587676912546158, 0.056065596640110016, 0.2533946633338928, 0.04020307958126068, 0.03702285513281822, 0.018525324761867523, 0.009753274731338024, 0.01584538072347641, 0.006842197384685278, 0.013304048217833042, 0.2415902465581894, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.09087645262479782, 0.0733630359172821, 0.03259122744202614, 0.05433432757854462, 0.028730718418955803, 0.026890264824032784, 0.0992540791630745, 0.042951032519340515, 0.1659460812807083, 0.017093859612941742, 0.006921885069459677, 0.0007972968742251396, 0.010357401333749294, 0.037234287708997726, 0.1852690428495407, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2766205668449402, 0.06249983608722687, 0.03302843123674393, 0.08374682813882828, 0.07296875864267349, 0.016804786399006844, 0.2612326145172119, 0.06074067950248718, 0.06402052938938141, 0.021471360698342323, 0.00216249143704772, 0.001582604949362576, 0.0037338242400437593, 0.005314995069056749, 0.23526467382907867, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.005338736344128847, 0.013486125506460667, 0.016210375353693962, 0.00714905746281147, 0.01115293800830841, 0.008639699779450893, 0.009605110622942448, 0.01017976924777031, 0.008433598093688488, 0.06244685873389244, 0.040223702788352966, 0.009117859415709972, 0.005228321999311447, 0.0028589563444256783, 0.13790398836135864, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.09661699831485748, 0.7619754076004028, 0.05676787346601486, 0.020180072635412216, 0.10883769392967224, 0.42711278796195984, 0.09064477682113647, 0.10612691193819046, 0.04782179743051529, 0.06935178488492966, 0.027948519214987755, 0.00755169615149498, 0.007339869160205126, 0.025803416967391968, 0.09292053431272507, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.042798254638910294, 0.23223945498466492, 0.062359996140003204, 0.01933804154396057, 0.04838808253407478, 0.30189236998558044, 0.0354127362370491, 0.019764740020036697, 0.00920741818845272, 0.0097093116492033, 0.0160877276211977, 0.0032758424058556557, 0.005296806804835796, 0.011010169051587582, 0.02110680378973484, 0.1301431953907013, 0.0347244068980217, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02002989500761032, 0.001048662350513041, 0.03834937512874603, 0.030392715707421303, 0.09750902652740479, 0.056120067834854126, 0.008173296228051186, 0.006944228895008564, 0.004440560005605221, 0.005061029922217131, 0.007118762470781803, 0.008411978371441364, 0.023608768358826637, 0.04182775691151619, 0.16016238927841187, 0.19350707530975342, 0.0006586865638382733, 0.008110460825264454, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.041295986622571945, 0.19780276715755463, 0.03777160495519638, 0.1712082475423813, 0.20935285091400146, 0.158755823969841, 0.3937656581401825, 0.684601902961731, 0.2584594190120697, 0.11237194389104843, 0.1112959012389183, 0.09882687777280807, 0.05429066717624664, 0.24210131168365479, 0.016339490190148354, 0.07742509245872498, 0.025898784399032593, 0.46813124418258667, 0.21566073596477509, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.26312491297721863, 0.2720799446105957, 0.005703570321202278, 0.0481516495347023, 0.027902500703930855, 0.0034437666181474924, 0.03425572067499161, 0.03555849939584732, 0.028000997379422188, 0.0429554246366024, 0.002753790933638811, 0.0017769382102414966, 0.002218457870185375, 0.003535473719239235, 0.1597488671541214, 0.15508510172367096, 0.002848779782652855, 0.006727630738168955, 0.01290579792112112, 0.0019038956379517913, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.22248251736164093, 0.03185709938406944, 0.000688861298840493, 0.005810217931866646, 0.007679672911763191, 0.0008787074475549161, 0.07858764380216599, 0.14273476600646973, 0.07306984066963196, 0.02433006465435028, 0.011720307171344757, 0.013396549038589, 0.017704129219055176, 0.034836068749427795, 0.1453055441379547, 0.1506490558385849, 0.0018329949816688895, 0.0011812039883807302, 0.010563074611127377, 0.0007367127691395581, 0.0007524989196099341, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1531120240688324, 0.15391655266284943, 0.006810865830630064, 0.07720811665058136, 0.008951452560722828, 0.01149735413491726, 0.2822602391242981, 0.30408379435539246, 0.48283058404922485, 0.33028021454811096, 0.16095426678657532, 0.031167738139629364, 0.03355513513088226, 0.13962571322917938, 0.012790725566446781, 0.0463392436504364, 0.0861721858382225, 0.5342088341712952, 0.5262086987495422, 0.252642959356308, 0.014757110737264156, 0.02778990939259529, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03593587130308151, 0.03233448788523674, 0.22662676870822906, 0.405829519033432, 0.014032814651727676, 0.02822977490723133, 0.09231841564178467, 0.1225365549325943, 0.20093639194965363, 0.2508411109447479, 0.5826555490493774, 0.037383783608675, 0.07952429354190826, 0.10720134526491165, 0.15212680399417877, 0.08082517981529236, 0.10121051222085953, 0.3481808602809906, 0.41374534368515015, 0.38359278440475464, 0.07890304177999496, 0.1096968874335289, 0.1685827672481537, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.037364520132541656, 0.04119153320789337, 0.0012645104434341192, 0.021537767723202705, 0.000536995125003159, 0.0011436643544584513, 0.019049961119890213, 0.06139632686972618, 0.385105162858963, 0.13276730477809906, 0.24771228432655334, 0.04952799528837204, 0.04911990836262703, 0.11973114311695099, 0.021608887240290642, 0.1433362513780594, 0.13670213520526886, 0.10138670355081558, 0.1093992069363594, 0.236768901348114, 0.09415888041257858, 0.011134332977235317, 0.019298367202281952, 0.5348934531211853, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.004867227748036385, 0.009626063518226147, 0.0003137234307359904, 0.0026314754504710436, 0.00027048110496252775, 0.000934475683607161, 0.007251756265759468, 0.03575620427727699, 0.40781450271606445, 0.05584407597780228, 0.040446195751428604, 0.005334825720638037, 0.007708138320595026, 0.06401336193084717, 0.010240204632282257, 0.024931270629167557, 0.02871265634894371, 0.20136752724647522, 0.1457405984401703, 0.13753218948841095, 0.13171687722206116, 0.07031083852052689, 0.04771474376320839, 0.5403124690055847, 0.04482616111636162, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19358457624912262, 0.2328234314918518, 0.0017398587660863996, 0.10100623220205307, 0.0019695234950631857, 0.1674531251192093, 0.4513051509857178, 0.6547151803970337, 0.030009860172867775, 0.7025956511497498, 0.1685936599969864, 0.03178222477436066, 0.13270388543605804, 0.23426049947738647, 0.010277668945491314, 0.026511939242482185, 0.12058579176664352, 0.09381356090307236, 0.09726550430059433, 0.13490843772888184, 0.36408668756484985, 0.19949088990688324, 0.09435784071683884, 0.45831772685050964, 0.1274537742137909, 0.014095090329647064, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09463346004486084, 0.5257620811462402, 0.0045187450014054775, 0.07222570478916168, 0.0025188177824020386, 0.1410406231880188, 0.06597349792718887, 0.0719805508852005, 0.09957849979400635, 0.17567123472690582, 0.18618373572826385, 0.02195402979850769, 0.042485080659389496, 0.12470933794975281, 0.00617468124255538, 0.12624163925647736, 0.03293433412909508, 0.07055910676717758, 0.06304988265037537, 0.23899653553962708, 0.15645378828048706, 0.07000429183244705, 0.02516351453959942, 0.06797400116920471, 0.07094329595565796, 0.1311238706111908, 0.21208471059799194, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.027796348556876183, 0.06599752604961395, 0.002643989399075508, 0.029425768181681633, 0.008861851878464222, 0.013279970735311508, 0.25377023220062256, 0.2656356692314148, 0.055540941655635834, 0.027583830058574677, 0.004816746339201927, 0.3890189528465271, 0.12020140886306763, 0.33882811665534973, 0.0040408894419670105, 0.1118171289563179, 0.015469676814973354, 0.08768722414970398, 0.046650953590869904, 0.23542486131191254, 0.09032069146633148, 0.05012429133057594, 0.004171812906861305, 0.15006321668624878, 0.017805932089686394, 0.049085501581430435, 0.035517167299985886, 0.6428134441375732, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4147956669330597, 0.5514373779296875, 0.09636387228965759, 0.29775112867355347, 0.03436855599284172, 0.08799602836370468, 0.07023341208696365, 0.10276275128126144, 0.25543972849845886, 0.10302554070949554, 0.05857125297188759, 0.029829595237970352, 0.114840567111969, 0.33078575134277344, 0.07371985912322998, 0.09301143884658813, 0.13257478177547455, 0.1489255279302597, 0.18642880022525787, 0.318376362323761, 0.31357452273368835, 0.1382697969675064, 0.07457731664180756, 0.17392435669898987, 0.00920780934393406, 0.020603884011507034, 0.049020376056432724, 0.322329580783844, 0.3050764203071594, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07031518220901489, 0.001305539975874126, 0.0025430582463741302, 0.010662226937711239, 0.0007357596186921, 0.000663888524286449, 0.0014398572966456413, 0.0005107407923787832, 0.005960140842944384, 0.0030986208003014326, 0.0017578504048287868, 0.00018377922242507339, 1.743367283779662e-05, 4.847845411859453e-05, 0.15638960897922516, 0.17444664239883423, 0.0007958812057040632, 5.6854176364140585e-05, 0.0004179355164524168, 0.00013179269444663078, 0.00024977640714496374, 0.0001107741700252518, 7.639485556865111e-05, 0.0008396806661039591, 0.00030287212575785816, 0.00023763117496855557, 0.003834246192127466, 0.003433886216953397, 0.00015348535089287907, 0.00014843019016552716, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24421003460884094, 0.03331591188907623, 0.07573812454938889, 0.33240795135498047, 0.006838400848209858, 0.008697851561009884, 0.06428743898868561, 0.06466686725616455, 0.006176145281642675, 0.06394235789775848, 0.09260299056768417, 0.19959890842437744, 0.02154124155640602, 0.021672323346138, 0.15025706589221954, 0.00841783918440342, 0.03505324944853783, 0.02469123899936676, 0.026689309626817703, 0.1500382125377655, 0.08861804753541946, 0.006530162878334522, 0.060150377452373505, 0.04669034481048584, 0.007807246409356594, 0.02131708152592182, 0.012364925816655159, 0.041818197816610336, 0.02841370552778244, 0.6981374621391296, 0.06836962699890137, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5462155342102051, 0.545982301235199, 0.3341628611087799, 0.5788259506225586, 0.08809857815504074, 0.06356553733348846, 0.022417092695832253, 0.0164126455783844, 0.00386660173535347, 0.10154324769973755, 0.14015790820121765, 0.0864240974187851, 0.34186482429504395, 0.22899740934371948, 0.05407746881246567, 0.0009672276792116463, 0.0037913541309535503, 0.00524782482534647, 0.006044968497008085, 0.07807419449090958, 0.026950905099511147, 0.0024354930501431227, 0.005482541862875223, 0.013836389407515526, 0.002816400956362486, 0.0006559633184224367, 0.002845867071300745, 0.018497759476304054, 0.19704575836658478, 0.41393977403640747, 0.4024144113063812, 0.00308317132294178, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.48888036608695984, 0.6578190326690674, 0.030819885432720184, 0.2205304652452469, 0.004883326590061188, 0.0656682699918747, 0.04461565986275673, 0.05094402655959129, 0.0005314986919984221, 0.15455113351345062, 0.10763049870729446, 0.1186080202460289, 0.14419804513454437, 0.1328149437904358, 0.09490374475717545, 0.0023347423411905766, 0.018236415460705757, 0.011423468589782715, 0.014267664402723312, 0.06272618472576141, 0.09006785601377487, 0.023437032476067543, 0.008957883343100548, 0.03532397374510765, 0.006200278177857399, 0.0002018583327298984, 0.016960909590125084, 0.04933774098753929, 0.1362536996603012, 0.47770828008651733, 0.5670948624610901, 0.06992122530937195, 0.03068283386528492, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15812784433364868, 0.9118645191192627, 0.022590545937418938, 0.05952226370573044, 0.00360964541323483, 0.07875056564807892, 0.013187792152166367, 0.02020449750125408, 0.0020393244922161102, 0.033818699419498444, 0.0449705570936203, 0.02132066898047924, 0.0717315599322319, 0.12101268768310547, 0.06353376060724258, 0.0730348452925682, 0.024321116507053375, 0.06646358221769333, 0.0630527138710022, 0.23201428353786469, 0.1378810703754425, 0.04738042131066322, 0.010255109518766403, 0.0316733755171299, 0.07226394861936569, 0.06345586478710175, 0.13366159796714783, 0.1651405692100525, 0.1875276118516922, 0.475235253572464, 0.34701114892959595, 0.106105737388134, 0.17074023187160492, 0.14835108816623688, NaN, NaN, NaN, NaN, NaN, NaN], [0.07771441340446472, 0.4748976230621338, 0.012594498693943024, 0.043653786182403564, 0.006564431358128786, 0.024485116824507713, 0.20463299751281738, 0.1550481915473938, 0.0016144687542691827, 0.005543926265090704, 0.0017496985383331776, 0.3491710126399994, 0.23835937678813934, 0.3316482901573181, 0.08539295196533203, 0.1317213624715805, 0.02603350207209587, 0.05892709270119667, 0.02498493157327175, 0.2902502715587616, 0.11121267080307007, 0.057563167065382004, 0.004654969088733196, 0.12363925576210022, 0.02343585342168808, 0.03682887554168701, 0.054189957678318024, 0.5043657422065735, 0.23388440907001495, 0.46154457330703735, 0.32561513781547546, 0.055846668779850006, 0.06476935744285583, 0.026345595717430115, 0.5623452067375183, NaN, NaN, NaN, NaN, NaN], [0.22228576242923737, 0.3581831455230713, 0.10504736006259918, 0.2062736451625824, 0.015430409461259842, 0.007369442842900753, 0.009848481975495815, 0.0027359407395124435, 0.003257193835452199, 0.004766176920384169, 0.0058546122163534164, 0.0040231142193078995, 0.032162997871637344, 0.05548902228474617, 0.22239458560943604, 0.037178635597229004, 0.08259578794240952, 0.0920928493142128, 0.09107104688882828, 0.19359135627746582, 0.17535823583602905, 0.06819135695695877, 0.03716395050287247, 0.07458745688199997, 0.0064619481563568115, 0.009060872718691826, 0.02094256319105625, 0.1461041122674942, 0.11104261875152588, 0.6685899496078491, 0.4500047266483307, 0.029085516929626465, 0.03437849134206772, 0.03590574488043785, 0.20188003778457642, 0.23542997241020203, NaN, NaN, NaN, NaN], [0.040305208414793015, 0.0008039010572247207, 0.001399470493197441, 0.006614126265048981, 0.0003286598657723516, 0.0002559607964940369, 0.0005696980515494943, 0.00010972175368806347, 0.0006102611077949405, 0.0009710662416182458, 0.0004746906051877886, 5.0628168537514284e-05, 6.201828455232317e-06, 1.1841932064271532e-05, 0.15342259407043457, 0.18516498804092407, 0.0009336460498161614, 7.266629108926281e-05, 0.00041225351742468774, 0.00023152375069912523, 0.0002865330025088042, 0.00012637366307899356, 8.909442112781107e-05, 0.0006568549433723092, 0.0003727772564161569, 0.00021836791711393744, 0.0030449857003986835, 0.002062517451122403, 0.0001740154402796179, 0.00019746039470192045, 0.0010639599058777094, 3.738106170203537e-05, 0.00018948569777421653, 0.0017019548686221242, 0.0021623496431857347, 7.414143328787759e-05, 0.00010166682477574795, NaN, NaN, NaN], [0.18667390942573547, 0.05485990643501282, 0.06146723031997681, 0.2094709873199463, 0.003188095986843109, 0.005957009736448526, 0.04363764822483063, 0.02604665607213974, 0.0011390803847461939, 0.022857926785945892, 0.035827361047267914, 0.07732249796390533, 0.00673074834048748, 0.004807854071259499, 0.15350142121315002, 0.014717604033648968, 0.07327108085155487, 0.049021750688552856, 0.04824157431721687, 0.2509053647518158, 0.1518847495317459, 0.011399514973163605, 0.08240412920713425, 0.052963949739933014, 0.012185328640043736, 0.03166860342025757, 0.029948236420750618, 0.0332757867872715, 0.026646502315998077, 0.6691258549690247, 0.05157328397035599, 0.010373775847256184, 0.027277877554297447, 0.022091276943683624, 0.06386284530162811, 0.02213944122195244, 0.7486419677734375, 0.1026511937379837, NaN, NaN], [0.46625471115112305, 0.6644052863121033, 0.19963930547237396, 0.36004284024238586, 0.06144074350595474, 0.06362717598676682, 0.016601700335741043, 0.006137203890830278, 0.0020489897578954697, 0.041981395334005356, 0.042364589869976044, 0.04546959325671196, 0.25786423683166504, 0.1048446074128151, 0.10812478512525558, 0.0010381464380770922, 0.0033105257898569107, 0.005275417119264603, 0.005129440221935511, 0.05292869359254837, 0.018404772505164146, 0.0016328096389770508, 0.0039754449389874935, 0.007563540246337652, 0.0015294092008844018, 0.00038045260589569807, 0.0016144785331562161, 0.00974529329687357, 0.09415796399116516, 0.176291361451149, 0.35064396262168884, 0.0026081653777509928, 0.0026635529939085245, 0.004589376971125603, 0.028667066246271133, 0.20089752972126007, 0.45412325859069824, 0.4352543354034424, 0.005037708207964897, NaN], [0.01868601329624653, 0.08739857375621796, 0.016145089641213417, 0.000850466953124851, 0.0035631621722131968, 0.013478883542120457, 0.0006747889565303922, 0.0010685214074328542, 0.013735192827880383, 0.0029910006560385227, 0.017663421109318733, 0.0005569100612774491, 0.0335303470492363, 0.010939561761915684, 0.13854636251926422, 0.1408424973487854, 0.01142195239663124, 0.027654578909277916, 0.018255943432450294, 0.00871819257736206, 0.007302883546799421, 0.002508251927793026, 0.0010894191218540072, 0.002539109904319048, 0.0016572934109717607, 0.002274427330121398, 0.00915378425270319, 0.004932411015033722, 0.000505969044752419, 0.0064278775826096535, 0.013472460210323334, 0.0009905033512040973, 0.004150861874222755, 0.015419019386172295, 0.013300818391144276, 0.00147106999065727, 0.01399929728358984, 0.03311459720134735, 0.0035406623501330614, 0.008275571279227734]], [[0.3301994204521179, 0.08890271931886673, 0.08465498685836792, 0.06385943293571472, 0.21852104365825653, 0.02508896216750145, 0.03711355850100517, 0.034155964851379395, 0.1728704422712326, 0.06344152241945267, 0.01567375846207142, 0.047274719923734665, 0.023079151287674904, 0.06240373104810715, 0.17532315850257874, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08584976941347122, 0.12593986093997955, 0.03313801810145378, 0.017280908301472664, 0.17652282118797302, 0.268716037273407, 0.12116961926221848, 0.2558431923389435, 0.04765854403376579, 0.04246087744832039, 0.0035840249620378017, 0.02463056705892086, 0.2119264155626297, 0.11800020188093185, 0.14393316209316254, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.046346988528966904, 0.39951857924461365, 0.5525277853012085, 0.10910754650831223, 0.13167327642440796, 0.030212268233299255, 0.021472660824656487, 0.018023721873760223, 0.1298973113298416, 0.04191790521144867, 0.1535157859325409, 0.04246748238801956, 0.3158371150493622, 0.15602277219295502, 0.1064835637807846, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0703379437327385, 0.07535148411989212, 0.05811825022101402, 0.428435742855072, 0.07080380618572235, 0.15123498439788818, 0.3036666214466095, 0.07787945121526718, 0.48052453994750977, 0.12286645174026489, 0.04789941385388374, 0.033336445689201355, 0.030469346791505814, 0.005462532863020897, 0.08732402324676514, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0663379579782486, 0.03187985718250275, 0.09551261365413666, 0.0323714055120945, 0.33827176690101624, 0.1471284031867981, 0.3127540946006775, 0.02734280750155449, 0.23260797560214996, 0.02317011170089245, 0.046465177088975906, 0.0992102101445198, 0.09175661206245422, 0.13314616680145264, 0.07444406300783157, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.034720633178949356, 0.01384154986590147, 0.012703170999884605, 0.020319687202572823, 0.10901976376771927, 0.7807050347328186, 0.03443336486816406, 0.028544975444674492, 0.061822760850191116, 0.00809338316321373, 0.007171421777456999, 0.01342758722603321, 0.09649696201086044, 0.05527613312005997, 0.10404697060585022, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.030445659533143044, 0.041789710521698, 0.023520270362496376, 0.01782963052392006, 0.16124852001667023, 0.06983006745576859, 0.4703807234764099, 0.01895260065793991, 0.027326058596372604, 0.07994905114173889, 0.026343191042542458, 0.032219063490629196, 0.022085823118686676, 0.031095484271645546, 0.24155765771865845, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.055046502500772476, 0.3847074508666992, 0.04798666015267372, 0.003912709187716246, 0.06840738654136658, 0.36789029836654663, 0.07226144522428513, 0.4079316258430481, 0.022340288385748863, 0.10408379882574081, 0.07774890959262848, 0.04753485694527626, 0.285355806350708, 0.16128498315811157, 0.02375940792262554, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03513112664222717, 0.11586778610944748, 0.03034079447388649, 0.001017131027765572, 0.04634808376431465, 0.03800477832555771, 0.03768199309706688, 0.013300161808729172, 0.14031966030597687, 0.015252463519573212, 0.053176701068878174, 0.06856708973646164, 0.13856393098831177, 0.054046642035245895, 0.2367301732301712, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.025786809623241425, 0.06564735621213913, 0.039564721286296844, 0.0026341548655182123, 0.016324089840054512, 0.016701271757483482, 0.020613567903637886, 0.0767805427312851, 0.22950275242328644, 0.51694655418396, 0.1544727236032486, 0.1054847463965416, 0.025381706655025482, 0.05480813980102539, 0.1677880734205246, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.012255452573299408, 0.02410232275724411, 0.08552651852369308, 0.002623841166496277, 0.010307574644684792, 0.0127415731549263, 0.021285703405737877, 0.010095748119056225, 0.06661782413721085, 0.12517453730106354, 0.7383688688278198, 0.19885332882404327, 0.07497892528772354, 0.10072800517082214, 0.06182975694537163, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2776626944541931, 0.046990759670734406, 0.032447993755340576, 0.015461347065865993, 0.08414210379123688, 0.04174359515309334, 0.19995476305484772, 0.013662091456353664, 0.019540153443813324, 0.048985805362463, 0.25616249442100525, 0.2484772503376007, 0.1799653023481369, 0.17696446180343628, 0.09890354424715042, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05504303798079491, 0.08340897411108017, 0.04799877479672432, 0.017563870176672935, 0.028545444831252098, 0.1704884171485901, 0.030681313946843147, 0.02359093725681305, 0.007767115719616413, 0.019779905676841736, 0.03771185874938965, 0.029841119423508644, 0.28957709670066833, 0.04182300344109535, 0.12634176015853882, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.06153338775038719, 0.02491314895451069, 0.02542346529662609, 0.0031092099379748106, 0.03241894021630287, 0.1874629557132721, 0.1358277052640915, 0.02619485929608345, 0.017582973465323448, 0.03225348889827728, 0.01329810544848442, 0.026643214747309685, 0.1614912450313568, 0.6035103797912598, 0.09545250982046127, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.027727488428354263, 0.10283610969781876, 0.02349940501153469, 0.010801603086292744, 0.0136191351339221, 0.1518852412700653, 0.05784522369503975, 0.11107083410024643, 0.10270816832780838, 0.1666017472743988, 0.06030665338039398, 0.06198698654770851, 0.05951831862330437, 0.015173939988017082, 0.1310720145702362, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03039383515715599, 0.011264979839324951, 0.30973049998283386, 0.33407092094421387, 0.24303670227527618, 0.013086382299661636, 0.12547586858272552, 0.047571711242198944, 0.07738520950078964, 0.2579103410243988, 0.13098950684070587, 0.3019145727157593, 0.018321001902222633, 0.10478901118040085, 0.1313871294260025, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.32489657402038574, 0.01967906951904297, 0.10292623937129974, 0.18745845556259155, 0.06220339238643646, 0.03126899152994156, 0.030121171846985817, 0.013807957991957664, 0.01960192248225212, 0.10352540761232376, 0.08122410625219345, 0.11610747873783112, 0.05098450556397438, 0.06022121384739876, 0.24838198721408844, 0.10530310869216919, 0.47072935104370117, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21547414362430573, 0.011987588368356228, 0.09540344774723053, 0.03949207067489624, 0.22973625361919403, 0.013393656350672245, 0.014646085910499096, 0.018391601741313934, 0.12483032047748566, 0.04761500656604767, 0.16838808357715607, 0.0500614158809185, 0.09093409031629562, 0.09172232449054718, 0.14920873939990997, 0.07470229268074036, 0.01594272069633007, 0.3473423421382904, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3455514907836914, 0.20528344810009003, 0.14200778305530548, 0.1397678107023239, 0.3345029056072235, 0.04282815381884575, 0.020769812166690826, 0.02952164225280285, 0.29125186800956726, 0.09975660592317581, 0.3298649489879608, 0.36294782161712646, 0.10288939625024796, 0.1784013956785202, 0.03550736606121063, 0.19784890115261078, 0.02982909232378006, 0.008884507231414318, 0.026416730135679245, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.023072484880685806, 0.08888474851846695, 0.04328835755586624, 0.009794876910746098, 0.18984860181808472, 0.0009663040982559323, 0.0038235578685998917, 0.05101485177874565, 0.059323158115148544, 0.00876270979642868, 0.021391507238149643, 0.02426949329674244, 0.013026251457631588, 0.06840420514345169, 0.15691325068473816, 0.15099161863327026, 0.004257611930370331, 0.06880252063274384, 0.03778434172272682, 0.016005711629986763, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20066522061824799, 0.18445545434951782, 0.10427504032850266, 0.02148139849305153, 0.3108636438846588, 0.0010669901967048645, 0.031332992017269135, 0.06621930748224258, 0.42585986852645874, 0.05703788995742798, 0.1919325739145279, 0.6617251038551331, 0.07196007668972015, 0.2038833349943161, 0.13549473881721497, 0.14908726513385773, 0.01576131209731102, 0.006129090208560228, 0.013888919726014137, 0.006888655014336109, 0.007033796049654484, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06934618204832077, 0.15043997764587402, 0.24868465960025787, 0.0180400051176548, 0.61164391040802, 0.0047634197399020195, 0.0077652581967413425, 0.01316747348755598, 0.09036756306886673, 0.016214115545153618, 0.09484434872865677, 0.7773507833480835, 0.3649398386478424, 0.19880527257919312, 0.026039909571409225, 0.1207430437207222, 0.0697125568985939, 0.0065151299349963665, 0.0038357542362064123, 0.04419673979282379, 0.16196060180664062, 0.49751368165016174, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5420496463775635, 0.775536835193634, 0.21455605328083038, 0.17522192001342773, 0.3905614912509918, 0.07102629542350769, 0.15213513374328613, 0.06534071266651154, 0.05938922241330147, 0.3742612600326538, 0.040289394557476044, 0.6919643878936768, 0.07523911446332932, 0.14220400154590607, 0.06588775664567947, 0.02684849314391613, 0.03953110799193382, 0.00281998747959733, 0.001733462675474584, 0.08529012650251389, 0.6486974358558655, 0.306731641292572, 0.07198647409677505, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05002814158797264, 0.18039211630821228, 0.4788157641887665, 0.0970841720700264, 0.5287489891052246, 0.07699278742074966, 0.024560611695051193, 0.055294524878263474, 0.031155720353126526, 0.029308732599020004, 0.023515479639172554, 0.10280930250883102, 0.01905171573162079, 0.033789344131946564, 0.006217750255018473, 0.012395885773003101, 0.009238478727638721, 0.0003186498652212322, 0.0010813054395839572, 0.008392964489758015, 0.2777543067932129, 0.44055092334747314, 0.0011997584952041507, 0.00246741552837193, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2326076328754425, 0.12470381706953049, 0.5816100239753723, 0.187625452876091, 0.17989297211170197, 0.58512943983078, 0.4148763120174408, 0.7688660621643066, 0.02497384324669838, 0.10204316675662994, 0.16508084535598755, 0.4722842574119568, 0.654721736907959, 0.31103214621543884, 0.02808636985719204, 0.034838397055864334, 0.015937600284814835, 0.002090656431391835, 0.002794815693050623, 0.008703295141458511, 0.10732896625995636, 0.4454900026321411, 0.001775766140781343, 0.0009654808673076332, 0.016644174233078957, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.32085803151130676, 0.3732209801673889, 0.8471049070358276, 0.2474840134382248, 0.8311324715614319, 0.1531035155057907, 0.14141014218330383, 0.12460694462060928, 0.15561653673648834, 0.05888388305902481, 0.03703024983406067, 0.2600737512111664, 0.049645353108644485, 0.08333000540733337, 0.053744472563266754, 0.293722003698349, 0.0148458918556571, 0.02856721729040146, 0.006315621547400951, 0.005582483485341072, 0.0013911855639889836, 0.004092940129339695, 0.0036679452750831842, 0.0010494120651856065, 0.016411608085036278, 0.023008037358522415, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.048572178930044174, 0.20163586735725403, 0.8568418025970459, 0.3438677489757538, 0.8764770030975342, 0.038519736379384995, 0.10765119642019272, 0.14438603818416595, 0.13915397226810455, 0.04139794409275055, 0.24816225469112396, 0.22188685834407806, 0.1582770049571991, 0.255889892578125, 0.05260627716779709, 0.13037414848804474, 0.020949387922883034, 0.03831411898136139, 0.007462172769010067, 0.02548721246421337, 0.006367610301822424, 0.008434200659394264, 0.010317808948457241, 0.003713584039360285, 0.00402417778968811, 0.19032441079616547, 0.26746228337287903, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10717450082302094, 0.14654512703418732, 0.5492125749588013, 0.149112731218338, 0.6473506689071655, 0.014123019762337208, 0.023513145744800568, 0.06304500997066498, 0.5243880152702332, 0.17494699358940125, 0.11734810471534729, 0.2534768283367157, 0.06080847606062889, 0.1781260073184967, 0.01657547615468502, 0.041874390095472336, 0.024160701781511307, 0.00029624058515764773, 0.00016299582784995437, 0.00014630405348725617, 0.0004776908899657428, 0.0010664566652849317, 0.005874973721802235, 0.000636687153019011, 0.0013240330154076219, 0.0912160873413086, 0.35286882519721985, 0.01772063784301281, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.024022793397307396, 0.20128284394741058, 0.39493197202682495, 0.16542883217334747, 0.7724959254264832, 0.05353498458862305, 0.039175428450107574, 0.21511156857013702, 0.10924636572599411, 0.3127569556236267, 0.20907098054885864, 0.6610769033432007, 0.026550091803073883, 0.07443477213382721, 0.04747246578335762, 0.11822566390037537, 0.015047432854771614, 0.019423136487603188, 0.00686526857316494, 0.0036870460025966167, 0.00022719512344338, 0.002930518239736557, 0.025171050801873207, 0.005165010690689087, 0.05391281098127365, 0.11512911319732666, 0.07776232063770294, 0.2967449426651001, 0.09380093216896057, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0639173686504364, 0.0019661476835608482, 0.03054100275039673, 0.07290788739919662, 0.07458660751581192, 0.0017515828367322683, 0.01338117104023695, 0.0049591753631830215, 0.10895326733589172, 0.03256915882229805, 0.07470867037773132, 0.022291045635938644, 0.00026081688702106476, 0.003768018214032054, 0.15579301118850708, 0.09375648200511932, 0.01475021056830883, 0.012638024985790253, 0.0046005831100046635, 0.051909249275922775, 0.0036223391070961952, 0.004371740389615297, 0.009388775564730167, 0.01159447617828846, 0.023305783048272133, 0.046531662344932556, 0.058873143047094345, 0.07503876090049744, 0.0337555818259716, 0.30213212966918945, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00809751357883215, 0.08670660853385925, 0.12165205925703049, 0.06173386052250862, 0.8110419511795044, 0.006245153024792671, 0.03447260707616806, 0.08050490915775299, 0.779870867729187, 0.2479465901851654, 0.38426774740219116, 0.6870184540748596, 0.2310730367898941, 0.07155610620975494, 0.05814361199736595, 0.060409948229789734, 0.03445665165781975, 0.000381257850676775, 0.0036348046269267797, 0.0002713070425670594, 0.0011815812904387712, 0.03030458651483059, 0.03435760363936424, 0.0019682012498378754, 0.00901943538337946, 0.2363511621952057, 0.7836493253707886, 0.05375572293996811, 0.0010517562041059136, 0.002096510259434581, 0.017742546275258064, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01971210353076458, 0.10859540849924088, 0.17558348178863525, 0.04931360110640526, 0.4077165424823761, 0.001824796199798584, 0.004386546555906534, 0.0422598272562027, 0.9374924302101135, 0.3226373493671417, 0.06322266161441803, 0.05341457948088646, 0.0039883931167423725, 0.004304073750972748, 0.13460686802864075, 0.19913224875926971, 0.17475517094135284, 0.0022224360145628452, 0.015882516279816628, 0.001058473251760006, 0.0005846276762895286, 0.02601638250052929, 0.037341512739658356, 0.002062901621684432, 0.01394632738083601, 0.062121838331222534, 0.09270716458559036, 0.13391432166099548, 0.011137665249407291, 0.003502808278426528, 0.007463122718036175, 0.4640289545059204, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.018049566075205803, 0.12295468151569366, 0.24470828473567963, 0.04122815281152725, 0.7332677245140076, 0.004472800530493259, 0.0029204280581325293, 0.018685931339859962, 0.4878760874271393, 0.20441682636737823, 0.08441592752933502, 0.4205068051815033, 0.04466289281845093, 0.13263334333896637, 0.0994158536195755, 0.33059969544410706, 0.017222048714756966, 0.029873082414269447, 0.008054245263338089, 0.002331576542928815, 0.0006345488945953548, 0.011296147480607033, 0.005269323009997606, 0.0004991231253370643, 0.01808379590511322, 0.0023433570750057697, 0.0409514382481575, 0.01219080574810505, 0.010968736372888088, 0.004035044461488724, 0.000618473335634917, 0.01301309373229742, 0.04461785778403282, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.007120466325432062, 0.02300306409597397, 0.2714575231075287, 0.07745856046676636, 0.6446666717529297, 0.0059507740661501884, 0.011145476251840591, 0.13244189321994781, 0.38060593605041504, 0.06726288050413132, 0.22673718631267548, 0.3522229492664337, 0.17927831411361694, 0.524927020072937, 0.09379637986421585, 0.11787470430135727, 0.013379373587667942, 0.03657921776175499, 0.007838133722543716, 0.006328434217721224, 0.0013346761697903275, 0.005374525673687458, 0.005563441663980484, 0.0013783610193058848, 0.003622437361627817, 0.10895299166440964, 0.17491653561592102, 0.013411260209977627, 0.006658618804067373, 0.013080593198537827, 0.0013389869127422571, 0.03540230169892311, 0.3923792839050293, 0.2429211437702179, NaN, NaN, NaN, NaN, NaN, NaN], [0.03649899363517761, 0.08160936087369919, 0.2519805133342743, 0.07504414021968842, 0.1795702874660492, 0.006024391856044531, 0.0073743402026593685, 0.061968039721250534, 0.7520835995674133, 0.28517279028892517, 0.1493321657180786, 0.3589819371700287, 0.04636238142848015, 0.16408585011959076, 0.046330999583005905, 0.03099578432738781, 0.01363852247595787, 8.312943100463599e-05, 4.0873743273550645e-05, 3.1056373700266704e-05, 8.971957868197933e-05, 0.0004970009904354811, 0.0021136843133717775, 0.00015606316446792334, 0.0008045462891459465, 0.029241982847452164, 0.24120952188968658, 0.011327153071761131, 0.006169632077217102, 0.004105421248823404, 0.0017298789462074637, 0.09891722351312637, 0.13539430499076843, 0.3545337915420532, 0.03266340494155884, NaN, NaN, NaN, NaN, NaN], [0.009416425600647926, 0.1558573991060257, 0.15325002372264862, 0.08311447501182556, 0.6221630573272705, 0.0029961667023599148, 0.006436231546103954, 0.027678541839122772, 0.2543543577194214, 0.47390833497047424, 0.28851544857025146, 0.6220062375068665, 0.014266690239310265, 0.05054754391312599, 0.0578170008957386, 0.05892227217555046, 0.006390280555933714, 0.00726453959941864, 0.002730957930907607, 0.0007821861072443426, 5.8160956541541964e-05, 0.0015625637024641037, 0.007388831116259098, 0.0016573512693867087, 0.027249574661254883, 0.062049947679042816, 0.056622181087732315, 0.2355845421552658, 0.04601869359612465, 0.006218506023287773, 0.00966239720582962, 0.07739637047052383, 0.4012998342514038, 0.09626632183790207, 0.38049787282943726, 0.10569068044424057, NaN, NaN, NaN, NaN], [0.04693470522761345, 0.0011674511479213834, 0.01364858541637659, 0.06039872020483017, 0.0427468940615654, 0.0009404723532497883, 0.007858873344957829, 0.0028007859364151955, 0.06382106244564056, 0.03982963413000107, 0.05175205320119858, 0.011254650540649891, 0.0001272865483770147, 0.001588277518749237, 0.15313954651355743, 0.09179559350013733, 0.00951253343373537, 0.010748236440122128, 0.0033872865606099367, 0.04677930101752281, 0.0018132117111235857, 0.0035809800028800964, 0.005968866869807243, 0.0062707834877073765, 0.02606387436389923, 0.033457815647125244, 0.03605461120605469, 0.04817588999867439, 0.03754975646734238, 0.2781437933444977, 0.015551367774605751, 0.2560427486896515, 0.08298799395561218, 0.06865174323320389, 0.12361031025648117, 0.04344068095088005, 0.28463616967201233, NaN, NaN, NaN], [0.017768997699022293, 0.1465732455253601, 0.15898801386356354, 0.12304693460464478, 0.8442554473876953, 0.006285809446126223, 0.04204265773296356, 0.12739135324954987, 0.8276333808898926, 0.5079721808433533, 0.5299316644668579, 0.8274551630020142, 0.09790517389774323, 0.02651425078511238, 0.11435628682374954, 0.02905191108584404, 0.012088212184607983, 0.00011298860044917092, 0.0012518719304352999, 4.317293132771738e-05, 0.0001948956778505817, 0.008923283778131008, 0.008874665014445782, 0.00048750368296168745, 0.0041984752751886845, 0.08557221293449402, 0.46109655499458313, 0.018593793734908104, 0.0004841866611968726, 0.0006005582981742918, 0.004410868044942617, 0.1617877185344696, 0.2815479040145874, 0.7414005398750305, 0.06452517956495285, 0.0009642028599046171, 0.0012653517769649625, 0.012943175621330738, NaN, NaN], [0.017107579857110977, 0.05770094692707062, 0.07052541524171829, 0.059498131275177, 0.2613165080547333, 0.0009367912425659597, 0.0028308003675192595, 0.01869240775704384, 0.8671534061431885, 0.40041688084602356, 0.03947103023529053, 0.0349445715546608, 0.00177917187102139, 0.002164072822779417, 0.1562660187482834, 0.1381005197763443, 0.0952477678656578, 0.0011117071844637394, 0.007693122606724501, 0.0001761779421940446, 8.233776316046715e-05, 0.0067709037102758884, 0.015442474745213985, 0.0005836034542880952, 0.005857429001480341, 0.020792629569768906, 0.02682901732623577, 0.05164036154747009, 0.0043857707642018795, 0.0008507486782036722, 0.004215322434902191, 0.19233396649360657, 0.21357974410057068, 0.14138071238994598, 0.12764914333820343, 0.011541306972503662, 0.001996394479647279, 0.004979089833796024, 0.4768531322479248, NaN], [0.006599111016839743, 0.004138579126447439, 0.06047067046165466, 0.013185898773372173, 0.15347044169902802, 0.000755132467020303, 0.007522573694586754, 0.002741254400461912, 0.10833818465471268, 0.005474736914038658, 0.009540018625557423, 0.00040286476723849773, 0.004092549905180931, 0.002003892557695508, 0.13896189630031586, 0.14079369604587555, 0.0077750058844685555, 0.008707624860107899, 0.002215370535850525, 0.0003697987995110452, 8.685041393619031e-05, 6.568676326423883e-05, 0.0005928067839704454, 0.00018151948461309075, 0.0013713521184399724, 0.003134837606921792, 0.004530616104602814, 0.0021016064565628767, 0.0014590725768357515, 0.01743447594344616, 0.0004639088874682784, 0.00557903666049242, 0.015868593007326126, 0.012156624346971512, 0.006375743541866541, 0.004486390855163336, 0.037133798003196716, 0.0008373309392482042, 0.015209782868623734, 0.053904592990875244]]], [[[0.042950913310050964, 0.0007196685182861984, 0.027302199974656105, 0.006393556483089924, 0.09642192721366882, 0.01637418009340763, 0.0023990001063793898, 0.0024961719755083323, 0.0020593979861587286, 0.0015603104839101434, 0.03318732604384422, 0.35782966017723083, 0.0989728793501854, 0.061845745891332626, 0.203965961933136, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10955026745796204, 0.02388770505785942, 0.04351670667529106, 0.023162608966231346, 0.012142845429480076, 0.035775765776634216, 0.03457501530647278, 0.11992064118385315, 0.01240380760282278, 0.007506475783884525, 0.05337386205792427, 0.6535924673080444, 0.5536571145057678, 0.19680790603160858, 0.140446737408638, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.005947283003479242, 0.0010204642312601209, 0.18009734153747559, 0.006447697523981333, 0.012463629245758057, 7.613956404384226e-05, 7.241032290039584e-05, 0.00011841111700050533, 0.0034185522235929966, 0.0034766956232488155, 0.002135018352419138, 0.005925178527832031, 0.003751354990527034, 0.0019247139571234584, 0.28479355573654175, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.014483454637229443, 0.022866876795887947, 0.32726621627807617, 0.007662326563149691, 0.09431912004947662, 0.0004296264669392258, 0.0011131323408335447, 0.0014158609556034207, 0.018019702285528183, 0.01865016296505928, 0.0020740600302815437, 0.0029411758296191692, 0.0016890126280486584, 0.0063899424858391285, 0.12852828204631805, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.030419446527957916, 0.058438073843717575, 0.3924228250980377, 0.035587672144174576, 0.08137891441583633, 0.010925069451332092, 0.001356365391984582, 0.0012006007600575686, 0.053269751369953156, 0.0027948038186877966, 0.04010261595249176, 0.01993635483086109, 0.004820133093744516, 0.004111820366233587, 0.21765674650669098, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07767480611801147, 0.006269918289035559, 0.09326869994401932, 0.6196063756942749, 0.11043263971805573, 0.052975643426179886, 0.02037718892097473, 0.0008919782703742385, 0.008360025472939014, 0.002104781800881028, 0.0179440937936306, 0.10498880594968796, 0.011864815838634968, 0.002359954407438636, 0.24602332711219788, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00026913435431197286, 8.159392746165395e-05, 0.007915529422461987, 0.05068095400929451, 0.6570689678192139, 0.32081079483032227, 0.05758208408951759, 0.0006442792946472764, 0.0015821922570466995, 6.469202344305813e-05, 0.003034515306353569, 0.0310077928006649, 0.025656316429376602, 0.0025228438898921013, 0.023106882348656654, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0005435149651020765, 0.0005490019102580845, 0.034476928412914276, 0.01287262886762619, 0.25229769945144653, 0.4536571502685547, 0.10281822830438614, 0.012222280725836754, 0.016108570620417595, 0.00031008716905489564, 0.0026372161228209734, 0.0034134499728679657, 0.0248859953135252, 0.017225822433829308, 0.02475895546376705, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.000726195692550391, 0.00036735343746840954, 0.007114858832210302, 0.0026034389156848192, 0.01250846590846777, 0.009484091773629189, 0.0354158952832222, 0.0016834242269396782, 0.19215336441993713, 0.007594457361847162, 0.003938279580324888, 2.8376112823025323e-05, 0.001137340790592134, 0.00011368053674232215, 0.29228782653808594, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0005387092242017388, 0.0003453432582318783, 0.015091696754097939, 0.06184916943311691, 0.003162123030051589, 0.014056581072509289, 0.012467358261346817, 0.009164737537503242, 0.05548334866762161, 0.008076494559645653, 0.005971547681838274, 0.001972777536138892, 0.006774900481104851, 0.001264052465558052, 0.2362799048423767, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0025044670328497887, 0.0023456772323697805, 0.07385681569576263, 0.006188494618982077, 0.021690815687179565, 0.0007893598522059619, 0.002135526854544878, 0.006048245821148157, 0.25190338492393494, 0.09442908316850662, 0.19532348215579987, 0.031008923426270485, 0.009561427868902683, 0.0021240306086838245, 0.21234139800071716, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.015501828864216805, 0.0072255814447999, 0.006012998055666685, 0.008203291334211826, 0.0171041339635849, 0.001770812552422285, 0.00655776634812355, 0.002186145167797804, 0.15154685080051422, 0.5713958144187927, 0.05368567630648613, 0.051326390355825424, 0.01612916588783264, 0.0019418209558352828, 0.18746227025985718, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05876695737242699, 0.005032649263739586, 0.05515526235103607, 0.012789947912096977, 0.017388533800840378, 0.00580496434122324, 0.015462081879377365, 0.009339934214949608, 0.0222479198127985, 0.03960718587040901, 0.14906688034534454, 0.2817051410675049, 0.14850065112113953, 0.09505022317171097, 0.10619710385799408, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.012425977736711502, 0.0006452641100622714, 0.00298808584921062, 0.001349467202089727, 0.014642779715359211, 0.0010115096811205149, 0.0033098396379500628, 0.00038259345456026495, 0.0035037249326705933, 0.008293021470308304, 0.03801131248474121, 0.8317341208457947, 0.018821584060788155, 0.057542454451322556, 0.011905365623533726, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04682805389165878, 0.01908799074590206, 0.10485747456550598, 0.060083843767642975, 0.15075230598449707, 0.029059063643217087, 0.04093548655509949, 0.03368941321969032, 0.017014725133776665, 0.011203174479305744, 0.0391479916870594, 0.24882012605667114, 0.37940239906311035, 0.12485622614622116, 0.12782400846481323, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.010830877348780632, 0.011870973743498325, 0.10922139137983322, 0.013140714727342129, 0.060979437083005905, 0.24213501811027527, 0.056873127818107605, 0.0565403513610363, 0.1606917381286621, 0.004471848253160715, 0.04391508549451828, 0.16444265842437744, 0.14521700143814087, 0.12183647602796555, 0.18165212869644165, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1442122757434845, 0.026047294959425926, 0.4262431859970093, 0.3211715519428253, 0.7946609258651733, 0.48857852816581726, 0.31943926215171814, 0.3322535455226898, 0.8442224860191345, 0.37700119614601135, 0.4491288661956787, 0.725179135799408, 0.5425247550010681, 0.7077597379684448, 0.47353750467300415, 0.12363631278276443, 0.14845161139965057, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.004308484960347414, 0.0038143862038850784, 0.01376394834369421, 0.007213444449007511, 0.0352218858897686, 0.009065943770110607, 0.00796457938849926, 0.009648038074374199, 0.012818497605621815, 0.005304576829075813, 0.00578665267676115, 0.025514552369713783, 0.003588201943784952, 0.005116589833050966, 0.1385156214237213, 0.14363405108451843, 0.021847352385520935, 0.10135873407125473, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.37350767850875854, 0.33144617080688477, 0.1264321357011795, 0.21400198340415955, 0.32627996802330017, 0.09132378548383713, 0.05067773535847664, 0.05911920592188835, 0.47554144263267517, 0.5285797715187073, 0.055136121809482574, 0.07909779250621796, 0.0048016151413321495, 0.023815851658582687, 0.05086187273263931, 0.13959342241287231, 0.059129536151885986, 0.04632453992962837, 0.0506979376077652, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.026979738846421242, 0.17144815623760223, 0.016802728176116943, 0.011190843768417835, 0.05719228833913803, 0.006600439548492432, 0.02541169337928295, 0.056367360055446625, 0.2566111385822296, 0.13847731053829193, 0.02390860766172409, 0.10821771621704102, 0.004193281754851341, 0.024024199694395065, 0.1485961675643921, 0.1401052325963974, 0.20328059792518616, 0.08711162209510803, 0.021569250151515007, 0.06437158584594727, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.010539665818214417, 0.02736317366361618, 0.020729688927531242, 0.012272891588509083, 0.037458207458257675, 0.020133765414357185, 0.006475721951574087, 0.0135318823158741, 0.14018985629081726, 0.043190933763980865, 0.014518915675580502, 0.06027117371559143, 0.013409063220024109, 0.008036705665290356, 0.12864065170288086, 0.14849096536636353, 0.24162742495536804, 0.13733072578907013, 0.023916935548186302, 0.4261094033718109, 0.034874048084020615, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06693296134471893, 0.05517994612455368, 0.31718623638153076, 0.09396946430206299, 0.13595829904079437, 0.09244473278522491, 0.0043823812156915665, 0.004134675953537226, 0.9252469539642334, 0.10048755258321762, 0.12945091724395752, 0.21572811901569366, 0.034586720168590546, 0.0726432204246521, 0.04207848384976387, 0.1122843325138092, 0.27548718452453613, 0.3164171576499939, 0.11597670614719391, 0.521038293838501, 0.1305568367242813, 0.04802507162094116, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07686225324869156, 0.019675375893712044, 0.2417416274547577, 0.08641211688518524, 0.27890217304229736, 0.038729339838027954, 0.01047417800873518, 0.015033761039376259, 0.4832261800765991, 0.05870191380381584, 0.2969569265842438, 0.6193534731864929, 0.12871475517749786, 0.22289764881134033, 0.5152896642684937, 0.13016629219055176, 0.2326299250125885, 0.3132029175758362, 0.32591310143470764, 0.1516764611005783, 0.09795279055833817, 0.02053435519337654, 0.1865263283252716, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.27357029914855957, 0.46676310896873474, 0.3964380621910095, 0.19407758116722107, 0.11257106065750122, 0.014855606481432915, 0.047355495393276215, 0.03237777575850487, 0.3466991186141968, 0.3347361087799072, 0.40522828698158264, 0.5460160970687866, 0.16927282512187958, 0.30020883679389954, 0.04839835315942764, 0.121080182492733, 0.4840172827243805, 0.47487083077430725, 0.3000609576702118, 0.5299880504608154, 0.09183567762374878, 0.057097259908914566, 0.12967270612716675, 0.04215369373559952, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03550037741661072, 0.12907657027244568, 0.07532694190740585, 0.016156595200300217, 0.003630127990618348, 0.01967703178524971, 0.04095811769366264, 0.0179570484906435, 0.39472800493240356, 0.07661326229572296, 0.4370958209037781, 0.4819755256175995, 0.022724222391843796, 0.033822834491729736, 0.04362141340970993, 0.08035996556282043, 0.5049515962600708, 0.21779249608516693, 0.22551923990249634, 0.48642098903656006, 0.17451445758342743, 0.14853931963443756, 0.2973877787590027, 0.02990546263754368, 0.12922555208206177, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.021909046918153763, 0.030848275870084763, 0.046106528490781784, 0.06202828511595726, 0.0325893796980381, 0.03412875533103943, 0.03159455209970474, 0.053456224501132965, 0.16627800464630127, 0.058593228459358215, 0.13071225583553314, 0.20816291868686676, 0.06561117619276047, 0.04416830837726593, 0.03868245705962181, 0.15412510931491852, 0.24815845489501953, 0.21706829965114594, 0.15909965336322784, 0.3919820487499237, 0.2097313106060028, 0.05961627885699272, 0.10788830369710922, 0.04644578695297241, 0.008778278715908527, 0.1666601300239563, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.012810717336833477, 0.0013835412682965398, 0.03224228695034981, 0.08643268793821335, 0.03331959247589111, 0.030278367921710014, 0.07819522172212601, 0.03789946064352989, 0.1521843820810318, 0.04584735259413719, 0.022775838151574135, 0.3594759702682495, 0.37505412101745605, 0.4203481376171112, 0.0833948627114296, 0.1319347769021988, 0.07332690805196762, 0.3709748387336731, 0.10343886911869049, 0.2416648119688034, 0.273651659488678, 0.142499178647995, 0.032821010798215866, 0.08169299364089966, 0.04221141338348389, 0.04960552975535393, 0.14849121868610382, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12084313482046127, 0.009313090704381466, 0.17649081349372864, 0.125856414437294, 0.03634244203567505, 0.028733352199196815, 0.006864639464765787, 0.002353896852582693, 0.16829386353492737, 0.1124483197927475, 0.061692144721746445, 0.19240431487560272, 0.09329058974981308, 0.18641597032546997, 0.018957242369651794, 0.15117543935775757, 0.09085448831319809, 0.23665060102939606, 0.09974268078804016, 0.5293540358543396, 0.2969721853733063, 0.0923411101102829, 0.04701923578977585, 0.47750627994537354, 0.31436240673065186, 0.11817371100187302, 0.08098391443490982, 0.05702001228928566, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.026597192510962486, 0.005893908906728029, 0.12369649112224579, 0.06400194019079208, 0.07115989178419113, 0.0058293454349040985, 0.008344992063939571, 0.00957680307328701, 0.04244829714298248, 0.036994293332099915, 0.07189996540546417, 0.04466360807418823, 0.12661096453666687, 0.2742233872413635, 0.042464204132556915, 0.2022491842508316, 0.0666579008102417, 0.032761361449956894, 0.03407268971204758, 0.3113752603530884, 0.5905517935752869, 0.21839523315429688, 0.043745849281549454, 0.02789805829524994, 0.042396336793899536, 0.08724991232156754, 0.07408890873193741, 0.010044119320809841, 0.12108539044857025, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0012156351003795862, 0.0009695529006421566, 0.021633058786392212, 0.003243132960051298, 0.017804604023694992, 0.006560572423040867, 0.00960883591324091, 0.043045539408922195, 0.008467147126793861, 0.0006170565611682832, 0.0028031598776578903, 0.004630656447261572, 1.7895566998049617e-05, 0.00023196694382932037, 0.14134538173675537, 0.14857184886932373, 0.38842764496803284, 0.16100677847862244, 0.1839173436164856, 0.03719957172870636, 0.5251989364624023, 0.25831982493400574, 0.06345110386610031, 0.01966739259660244, 0.013820506632328033, 0.10135386884212494, 0.06285497546195984, 0.037499457597732544, 0.09235794097185135, 0.06518241763114929, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3736850321292877, 0.29077818989753723, 0.43184730410575867, 0.4823248088359833, 0.7379603385925293, 0.5093098282814026, 0.5006043910980225, 0.3135696351528168, 0.5183887481689453, 0.13794882595539093, 0.04961319640278816, 0.12779268622398376, 0.1589212864637375, 0.22346213459968567, 0.1422436237335205, 0.15810954570770264, 0.08897967636585236, 0.2754043936729431, 0.11542505025863647, 0.7166418433189392, 0.6856120824813843, 0.15602687001228333, 0.03588242083787918, 0.10233978182077408, 0.06907100230455399, 0.13906386494636536, 0.06064911186695099, 0.02474391460418701, 0.09316151589155197, 0.5409220457077026, 0.18577302992343903, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15325459837913513, 0.1614270806312561, 0.4186149537563324, 0.16462315618991852, 0.44647181034088135, 0.7114150524139404, 0.12785741686820984, 0.04132780805230141, 0.047578196972608566, 0.12349404394626617, 0.3133608400821686, 0.35326144099235535, 0.30924320220947266, 0.31196898221969604, 0.028064150363206863, 0.07972963899374008, 0.06995329260826111, 0.2565014958381653, 0.11985079944133759, 0.5429201126098633, 0.3072132468223572, 0.04467121511697769, 0.06233014911413193, 0.06391221284866333, 0.06306523084640503, 0.04008801653981209, 0.16940940916538239, 0.21208623051643372, 0.3237960636615753, 0.4987465739250183, 0.14530567824840546, 0.42085787653923035, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06399086862802505, 0.06306004524230957, 0.1948489397764206, 0.12845031917095184, 0.26295408606529236, 0.38098499178886414, 0.0839061513543129, 0.02110268920660019, 0.07144157588481903, 0.01679118163883686, 0.14834797382354736, 0.479995995759964, 0.24741992354393005, 0.2288939356803894, 0.04729384183883667, 0.057688161730766296, 0.05957844480872154, 0.09227755665779114, 0.06308872997760773, 0.6051628589630127, 0.41719216108322144, 0.06513097882270813, 0.11441777646541595, 0.2576654255390167, 0.039566945284605026, 0.04989808052778244, 0.41204503178596497, 0.6269510388374329, 0.0653882622718811, 0.2309982180595398, 0.05030554160475731, 0.12162061780691147, 0.2016562819480896, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.041305530816316605, 0.00217662681825459, 0.29091107845306396, 0.12698692083358765, 0.3031243085861206, 0.1103614866733551, 0.14891935884952545, 0.018863126635551453, 0.033797744661569595, 0.008303376846015453, 0.009713392704725266, 0.31765925884246826, 0.4755025804042816, 0.4005468487739563, 0.10761724412441254, 0.08513950556516647, 0.05776134505867958, 0.44855204224586487, 0.15441171824932098, 0.37962910532951355, 0.43142464756965637, 0.21386101841926575, 0.07478547096252441, 0.22071515023708344, 0.1727379858493805, 0.06471506506204605, 0.1414414495229721, 0.20356127619743347, 0.23849359154701233, 0.28116941452026367, 0.22387196123600006, 0.24124523997306824, 0.10411572456359863, 0.14086224138736725, NaN, NaN, NaN, NaN, NaN, NaN], [0.4954506754875183, 0.04642331227660179, 0.603453516960144, 0.26468321681022644, 0.3210473358631134, 0.15078485012054443, 0.027168329805135727, 0.004181328695267439, 0.10826757550239563, 0.10845811665058136, 0.053085505962371826, 0.20335085690021515, 0.12072784453630447, 0.17107200622558594, 0.059424202889204025, 0.09857918322086334, 0.08268877118825912, 0.17155912518501282, 0.08326277136802673, 0.3910389840602875, 0.23102693259716034, 0.0706368237733841, 0.04062340036034584, 0.34264665842056274, 0.40400993824005127, 0.14310938119888306, 0.07597656548023224, 0.059025220572948456, 0.46083009243011475, 0.6441643834114075, 0.8002472519874573, 0.34466618299484253, 0.10859531164169312, 0.04317509010434151, 0.042760394513607025, NaN, NaN, NaN, NaN, NaN], [0.21408557891845703, 0.03960772231221199, 0.43507251143455505, 0.10961537808179855, 0.42240580916404724, 0.06637464463710785, 0.08428787440061569, 0.03856734186410904, 0.0027873425278812647, 0.012926235795021057, 0.019708000123500824, 0.017574653029441833, 0.10679914057254791, 0.20499441027641296, 0.14648839831352234, 0.07982634007930756, 0.027687683701515198, 0.01305405143648386, 0.01568622700870037, 0.15395750105381012, 0.36470726132392883, 0.09429053217172623, 0.02618592418730259, 0.00988653302192688, 0.03718657046556473, 0.057223062962293625, 0.036843542009592056, 0.008861655369400978, 0.039983998984098434, 0.5628355145454407, 0.5858935713768005, 0.11540589481592178, 0.07112369686365128, 0.022479010745882988, 0.0049066911451518536, 0.07443748414516449, NaN, NaN, NaN, NaN], [0.002137779025360942, 0.0005492505733855069, 0.03787382319569588, 0.004300523083657026, 0.03090864233672619, 0.003432363970205188, 0.010591491125524044, 0.028211969882249832, 0.003533262060955167, 0.0003883022291120142, 0.0014010752784088254, 0.0010855919681489468, 8.133743904181756e-06, 7.628504681633785e-05, 0.13786831498146057, 0.13230623304843903, 0.39635705947875977, 0.12619565427303314, 0.23844560980796814, 0.04749276116490364, 0.5552228093147278, 0.304650217294693, 0.16151569783687592, 0.05923860892653465, 0.03940735384821892, 0.37161606550216675, 0.13852664828300476, 0.1098584458231926, 0.421970933675766, 0.059641290456056595, 0.35413044691085815, 0.2336989790201187, 0.21869167685508728, 0.04408164322376251, 0.03093402087688446, 0.08392708003520966, 0.038801465183496475, NaN, NaN, NaN], [0.39364972710609436, 0.15414100885391235, 0.5289453864097595, 0.2158767729997635, 0.8369554877281189, 0.5879349708557129, 0.29191306233406067, 0.1240038275718689, 0.0375535674393177, 0.006134674418717623, 0.003127586329355836, 0.02892274223268032, 0.023530103266239166, 0.026029296219348907, 0.16074688732624054, 0.06938444077968597, 0.08034616708755493, 0.1555827558040619, 0.07347460091114044, 0.4763748347759247, 0.40589335560798645, 0.07265187799930573, 0.022002995014190674, 0.0527057945728302, 0.07314148545265198, 0.11090734601020813, 0.03504399210214615, 0.0172868762165308, 0.14030121266841888, 0.3467526137828827, 0.21038202941417694, 0.6312639117240906, 0.1208876520395279, 0.020520374178886414, 0.014591614715754986, 0.03736459091305733, 0.22129306197166443, 0.05682671070098877, NaN, NaN], [0.2684386968612671, 0.29252222180366516, 0.6921796798706055, 0.1771971732378006, 0.6445736885070801, 0.7333542704582214, 0.14767038822174072, 0.04686985909938812, 0.030383678153157234, 0.06000908464193344, 0.1879548877477646, 0.5258318781852722, 0.3533342778682709, 0.3370157778263092, 0.05586722865700722, 0.08218587934970856, 0.08353152126073837, 0.244074746966362, 0.15340235829353333, 0.5709766745567322, 0.4268343448638916, 0.06391507387161255, 0.13458560407161713, 0.14046461880207062, 0.13024689257144928, 0.043825987726449966, 0.1802380084991455, 0.2593124508857727, 0.4235299825668335, 0.23401854932308197, 0.23376718163490295, 0.4458163380622864, 0.1644086241722107, 0.22351105511188507, 0.25077733397483826, 0.28149890899658203, 0.3320602774620056, 0.05098887160420418, 0.4388013482093811, NaN], [0.0015460141003131866, 0.010688474401831627, 0.09971211850643158, 0.017146917060017586, 0.1899741291999817, 0.03437719866633415, 0.022833971306681633, 0.015900788828730583, 0.05731913447380066, 0.0008445536368526518, 0.0073861475102603436, 0.06343144923448563, 0.11084617674350739, 0.11975067108869553, 0.13715405762195587, 0.13887250423431396, 0.1972966492176056, 0.3352757692337036, 0.30585116147994995, 0.6380553841590881, 0.5158089995384216, 0.3850407004356384, 0.3912012279033661, 0.2877788245677948, 0.30187875032424927, 0.20025724172592163, 0.34020906686782837, 0.47167572379112244, 0.3815076947212219, 0.5385518074035645, 0.20663535594940186, 0.37741178274154663, 0.29376763105392456, 0.3577961027622223, 0.21765607595443726, 0.14290691912174225, 0.3544510304927826, 0.07646653801202774, 0.1391337811946869, 0.019570577889680862]], [[0.010500228963792324, 0.7224081754684448, 0.030353030189871788, 0.00683749420568347, 0.007232841569930315, 0.018554184585809708, 0.0004432629211805761, 0.02719983458518982, 0.0006519495509564877, 0.0012597806053236127, 0.006804677192121744, 0.0011734187137335539, 0.003679303452372551, 0.010371293872594833, 0.019012004137039185, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0004097823693882674, 0.007568135391920805, 0.05432860180735588, 0.08570658415555954, 0.005480978172272444, 0.0009473124518990517, 0.000799189496319741, 0.0012391285272315145, 0.00044785221689380705, 0.0009745006100274622, 0.013956908136606216, 0.00011593959061428905, 0.004404959734529257, 0.0031790253706276417, 0.20507724583148956, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.022728245705366135, 0.0194535069167614, 0.024020839482545853, 0.023168254643678665, 0.45748311281204224, 0.5855799913406372, 0.21754446625709534, 0.1001717820763588, 0.0221620611846447, 0.0033511894289404154, 0.03508710116147995, 0.20201759040355682, 0.2973189353942871, 0.04947788640856743, 0.0494859553873539, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.010499863885343075, 0.004784405697137117, 0.0035181313287466764, 0.007238015066832304, 0.4155227243900299, 0.8333501219749451, 0.07475034892559052, 0.20445603132247925, 0.005854693241417408, 0.001852003508247435, 0.02841898612678051, 0.243921160697937, 0.10275343060493469, 0.13816815614700317, 0.07406751066446304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00768234534189105, 0.012151399627327919, 0.0006104251369833946, 0.0018971813842654228, 0.08389636874198914, 0.7291921973228455, 0.2573831081390381, 0.13359335064888, 0.0011000150116160512, 0.0005446228897199035, 0.036390628665685654, 0.06110000237822533, 0.1527252048254013, 0.14593005180358887, 0.05624886974692345, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0037335127126425505, 0.004452059045433998, 0.00018280810036230832, 0.016856878995895386, 0.0016014263965189457, 0.05306785926222801, 0.5318921208381653, 0.2889253497123718, 0.0004385874199215323, 0.007465890143066645, 0.0005691659171134233, 0.008836256340146065, 0.00793292187154293, 0.0033322598319500685, 0.1706118881702423, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00023320072796195745, 0.0486629419028759, 0.0005405444535426795, 0.005952970590442419, 0.0009982762858271599, 0.004001363180577755, 0.009125707671046257, 0.6945337057113647, 0.006549985148012638, 0.007807720452547073, 0.003924727905541658, 0.004149672109633684, 0.003537258366122842, 0.001676861196756363, 0.11541670560836792, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0021667596884071827, 0.0005287157837301493, 0.009149480611085892, 0.024324318394064903, 0.0018866003956645727, 0.0003624066011980176, 0.0004668526817113161, 0.0064473398961126804, 0.0217228215187788, 0.0031395854894071817, 0.0052951243706047535, 0.004629157949239016, 0.003511544084176421, 0.0017145106103271246, 0.2705381214618683, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0036477160174399614, 0.018601393327116966, 0.00400471780449152, 0.016223786398768425, 0.015442389994859695, 0.030637366697192192, 0.04816145822405815, 0.009263478219509125, 0.08580432087182999, 0.07024423778057098, 0.17587034404277802, 0.2670482397079468, 0.10741393268108368, 0.11723090708255768, 0.197556272149086, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0067135002464056015, 0.005400336813181639, 0.002429268090054393, 0.0005210567032918334, 0.0009090648964047432, 0.056922394782304764, 0.006305574905127287, 0.02051912061870098, 0.009087055921554565, 0.0029723523184657097, 0.5903128385543823, 0.4623943269252777, 0.5148944854736328, 0.10147220641374588, 0.10177940130233765, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.016283290460705757, 0.004236595239490271, 0.00024049253261182457, 0.00013081195356789976, 0.004825976211577654, 0.03370611369609833, 0.030076656490564346, 0.006495397537946701, 0.015585500746965408, 0.0006116450531408191, 0.009124655276536942, 0.7220618724822998, 0.5160555839538574, 0.16948190331459045, 0.04205150157213211, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04056651145219803, 0.05449386313557625, 0.007923644036054611, 0.00034379694261588156, 0.0072999089024960995, 0.005707062315195799, 0.018278487026691437, 0.00924981851130724, 0.0004191468469798565, 0.0015566512010991573, 0.0019580996595323086, 0.06517467647790909, 0.4938390851020813, 0.1360015720129013, 0.14540629088878632, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.02595147117972374, 0.0358305424451828, 0.021912503987550735, 0.01559682097285986, 0.0029425774700939655, 0.008820675313472748, 0.259022980928421, 0.24083182215690613, 0.0008326273527927697, 0.009937180206179619, 0.008380424231290817, 0.0008840225636959076, 0.11912944912910461, 0.5976794362068176, 0.17433230578899384, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.024576334282755852, 0.01131413970142603, 0.0036256120074540377, 0.007047882303595543, 0.015460383147001266, 0.007877636700868607, 0.035456594079732895, 0.017273712903261185, 0.0020541276317089796, 0.005268692504614592, 0.003138576401397586, 0.0058868261985480785, 0.09279357641935349, 0.45485755801200867, 0.2460370808839798, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.02016485668718815, 0.03839857131242752, 0.0345035195350647, 0.005700604524463415, 0.03111962042748928, 0.03698137030005455, 0.056010663509368896, 0.043163470923900604, 0.004449993837624788, 0.000997284660115838, 0.006035848520696163, 0.0027079761493951082, 0.009604639373719692, 0.02099894918501377, 0.13394789397716522, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.021257108077406883, 0.04756314679980278, 0.05559564009308815, 0.030912479385733604, 0.2625647187232971, 0.138688862323761, 0.027820995077490807, 0.05787678435444832, 0.3002224862575531, 0.018701573833823204, 0.027547171339392662, 0.19844435155391693, 0.1917300671339035, 0.07151354849338531, 0.16648255288600922, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4235764741897583, 0.10086580365896225, 0.07221788167953491, 0.13654322922229767, 0.04923773929476738, 0.06516944617033005, 0.07642015814781189, 0.147566020488739, 0.013325832784175873, 0.07923475652933121, 0.03588176146149635, 0.02368854358792305, 0.12847480177879333, 0.04384613409638405, 0.18713882565498352, 0.10658828914165497, 0.44162610173225403, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.8895729184150696, 0.7431688904762268, 0.3041851818561554, 0.5492796897888184, 0.7013789415359497, 0.2035668045282364, 0.4541507959365845, 0.17740322649478912, 0.37418368458747864, 0.7257221937179565, 0.3302299678325653, 0.32646968960762024, 0.4535413682460785, 0.2710181474685669, 0.06444819271564484, 0.14346696436405182, 0.1105659008026123, 0.04705679044127464, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18918083608150482, 0.07354198396205902, 0.03709281235933304, 0.039312511682510376, 0.2119109183549881, 0.32255253195762634, 0.06547961384057999, 0.022612132132053375, 0.0069438498467206955, 0.04682554677128792, 0.04775600507855415, 0.10260774195194244, 0.060122229158878326, 0.07651683688163757, 0.11037445813417435, 0.14569434523582458, 0.006359750870615244, 0.06321832537651062, 0.009962446056306362, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05778415873646736, 0.1888784021139145, 0.12087801843881607, 0.08340981602668762, 0.2725185453891754, 0.956253707408905, 0.6455949544906616, 0.6532288789749146, 0.3585406243801117, 0.18532338738441467, 0.18782632052898407, 0.09142936766147614, 0.8097347617149353, 0.3558001220226288, 0.037162330001592636, 0.14614860713481903, 0.0770370289683342, 0.14572308957576752, 0.11918944120407104, 0.003047030884772539, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04896414652466774, 0.25620371103286743, 0.11985385417938232, 0.0157163105905056, 0.14219185709953308, 0.22957918047904968, 0.36173656582832336, 0.07001917064189911, 0.3676673173904419, 0.12105175852775574, 0.22853095829486847, 0.07480601221323013, 0.5630075335502625, 0.8219463229179382, 0.12425509095191956, 0.16211360692977905, 0.1199408695101738, 0.008137544617056847, 0.026895001530647278, 0.022997038438916206, 0.0004772362008225173, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04714362695813179, 0.01630709134042263, 0.04501143842935562, 0.03696214035153389, 0.036871057003736496, 0.14248797297477722, 0.08399422466754913, 0.03027486614882946, 0.0030259382911026478, 0.019033554941415787, 0.2224818617105484, 0.033125121146440506, 0.02079186774790287, 0.04913722351193428, 0.46250322461128235, 0.1276824176311493, 0.05415544658899307, 0.008876973763108253, 0.006533092353492975, 0.16286829113960266, 0.4191088378429413, 0.11241274327039719, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.033912286162376404, 0.0072718155570328236, 0.013269636780023575, 0.010754123330116272, 0.003932052757591009, 0.022333307191729546, 0.05135813727974892, 0.17082874476909637, 0.004249163903295994, 0.009168761782348156, 0.00692910747602582, 0.00042953240335918963, 0.008801857940852642, 0.008872170932590961, 0.02866899035871029, 0.1310766041278839, 0.09720440953969955, 0.005617472343146801, 0.018550021573901176, 0.07474999874830246, 0.03211009502410889, 0.01561786886304617, 0.5897646546363831, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.026226887479424477, 0.006219716742634773, 0.016528652980923653, 0.019500089809298515, 0.009756595827639103, 0.01771577261388302, 0.10877248644828796, 0.07924166321754456, 0.026382839307188988, 0.007807224057614803, 0.018975039944052696, 0.009491248056292534, 0.042680755257606506, 0.025040525943040848, 0.31068748235702515, 0.07142644375562668, 0.019657818600535393, 0.044225241988897324, 0.006672952324151993, 0.015112369321286678, 0.03715437650680542, 0.012035970576107502, 0.08684496581554413, 0.5578015446662903, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0181743074208498, 0.0022439020685851574, 0.027739310637116432, 0.07926302403211594, 0.007397042121738195, 0.01831221394240856, 0.057637136429548264, 0.025927647948265076, 0.03431807458400726, 0.03189869597554207, 0.20874466001987457, 0.006929311901330948, 0.08810199052095413, 0.09789149463176727, 0.25120988488197327, 0.06384367495775223, 0.009399783797562122, 0.06692944467067719, 0.013825987465679646, 0.01438650768250227, 0.11814092099666595, 0.025182364508509636, 0.04756484180688858, 0.4922580420970917, 0.010614832863211632, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0006848929915577173, 0.00015734595945104957, 0.0022563491947948933, 0.00281638465821743, 0.00390908308327198, 0.012311742641031742, 0.006667551584541798, 0.010898235253989697, 0.18826207518577576, 0.0010989188449457288, 0.003811799455434084, 0.0007082286756485701, 0.0025871950201690197, 0.0005297476891428232, 0.004719105549156666, 0.21570175886154175, 0.004600263200700283, 0.0039491499774158, 0.0010213260538876057, 0.00511409854516387, 0.00780195789411664, 0.0035460677463561296, 0.06005942076444626, 0.002209970960393548, 0.0011990047059953213, 0.010184505954384804, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.008918036706745625, 0.01932302489876747, 0.1743663251399994, 0.04276113957166672, 0.17357498407363892, 0.05217360332608223, 0.01903947815299034, 0.006896412931382656, 0.02532179281115532, 0.019349897280335426, 0.14434273540973663, 0.2454780638217926, 0.06247624009847641, 0.03444024175405502, 0.2827233076095581, 0.15804870426654816, 0.10358668118715286, 0.018792977556586266, 0.0036350360605865717, 0.02226737141609192, 0.007843486964702606, 0.002713214373216033, 0.3624168336391449, 0.00397031893953681, 0.013842551037669182, 0.05391863361001015, 0.040338534861803055, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.014348846860229969, 0.006216275505721569, 0.06011093780398369, 0.05047134682536125, 0.013856974430382252, 0.08402124047279358, 0.0029483914840966463, 0.0018935499247163534, 0.004232283215969801, 0.022591279819607735, 0.34387707710266113, 0.06330335885286331, 0.20501238107681274, 0.1859048306941986, 0.0244001317769289, 0.0703621581196785, 0.01676221750676632, 0.03283774480223656, 0.005265639629215002, 0.016811830922961235, 0.008307189680635929, 0.0008217993890866637, 0.06662888079881668, 0.006444453727453947, 0.0015952866524457932, 0.03341786190867424, 0.28674793243408203, 0.09830270707607269, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.016000788658857346, 0.003648907644674182, 0.07618206739425659, 0.26581478118896484, 0.00828572828322649, 0.01491115428507328, 0.006984202191233635, 0.00572665361687541, 0.007784067187458277, 0.03336494415998459, 0.19996345043182373, 0.0026567107997834682, 0.14645317196846008, 0.1677580624818802, 0.0739188864827156, 0.00274313404224813, 0.01220498327165842, 0.001565106911584735, 0.014617281965911388, 0.0015394951915368438, 0.00014163085143081844, 0.0032730719540268183, 0.04253724217414856, 0.01929563470184803, 0.0011092370841652155, 0.008900013752281666, 0.14250728487968445, 0.44352540373802185, 0.012739983387291431, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.033913157880306244, 0.5720782279968262, 0.09820353239774704, 0.06329890340566635, 0.10058190673589706, 0.8026418685913086, 0.08380495011806488, 0.37448471784591675, 0.04885341227054596, 0.01422097533941269, 0.32552391290664673, 0.701602578163147, 0.9988673329353333, 0.9602208137512207, 0.015194611623883247, 0.12441921979188919, 0.09727630764245987, 0.031539320945739746, 0.0390433706343174, 0.004017204977571964, 0.003718326799571514, 0.06902258098125458, 0.21229486167430878, 0.1692674309015274, 0.507585346698761, 0.24224399030208588, 0.4713107943534851, 0.22175242006778717, 0.1071210727095604, 0.001354279462248087, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01701497472822666, 0.004510161932557821, 0.04222021996974945, 0.131240576505661, 0.007172171492129564, 0.0009335885988548398, 0.0025300730485469103, 0.0012859954731538892, 0.013300590217113495, 0.05520036071538925, 0.2908037602901459, 0.0021335158962756395, 0.11976832151412964, 0.046004947274923325, 0.029495948925614357, 0.11131177842617035, 0.045754965394735336, 0.13187335431575775, 0.021390099078416824, 0.2008819729089737, 0.1753949522972107, 0.029810786247253418, 0.1191062182188034, 0.0330519825220108, 0.021209293976426125, 0.007793682627379894, 0.004569755867123604, 0.21031485497951508, 0.08390634506940842, 0.11696453392505646, 0.2920413017272949, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0007848403765819967, 0.002563882153481245, 0.003471110016107559, 0.009534057229757309, 0.012083875946700573, 0.006908607203513384, 0.0028729254845529795, 0.0018324146512895823, 0.009593485854566097, 0.008395246230065823, 0.009609236381947994, 0.05064208433032036, 0.00595981115475297, 0.002902570180594921, 0.2071433663368225, 0.28942060470581055, 0.004874760750681162, 0.02575746178627014, 0.03629674017429352, 0.0339069589972496, 0.06067432835698128, 0.06949229538440704, 0.17600718140602112, 0.04042575880885124, 0.0021073101088404655, 0.002125136088579893, 0.0013297069817781448, 0.013164625503122807, 0.019647862762212753, 0.0625171884894371, 0.003036472015082836, 0.15673543512821198, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.008253121748566628, 0.01393465232104063, 0.03316362947225571, 0.045629892498254776, 0.015712177380919456, 0.15894818305969238, 0.02510240487754345, 0.013996893540024757, 0.6886083483695984, 0.014645315706729889, 0.04062162712216377, 0.02812274731695652, 0.10265076905488968, 0.10770027339458466, 0.07716524600982666, 0.29843398928642273, 0.006499151699244976, 0.002175502711907029, 0.00474061444401741, 0.012194045819342136, 0.024305779486894608, 0.05332900583744049, 0.20892387628555298, 0.06725459545850754, 0.0056669809855520725, 0.023831704631447792, 0.0038352743722498417, 0.008001168258488178, 0.00692057004198432, 0.006051996257156134, 0.0008782879449427128, 0.0244371946901083, 0.05294432491064072, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0017006727866828442, 0.008613905869424343, 0.08540165424346924, 0.014788517728447914, 0.11802737414836884, 0.058780014514923096, 0.008085138164460659, 0.003584004705771804, 0.06396479159593582, 0.006658769678324461, 0.02042919024825096, 0.3806440234184265, 0.01375669613480568, 0.01512871216982603, 0.1676391214132309, 0.19362471997737885, 0.05030333995819092, 0.012831996195018291, 0.0028119448106735945, 0.011659904383122921, 0.0070129260420799255, 0.002673238283023238, 0.1857692450284958, 0.0015845311572775245, 0.003893241984769702, 0.009055504575371742, 0.013083641417324543, 0.009338575415313244, 0.007860029116272926, 0.009482803754508495, 0.019751103594899178, 0.03845033049583435, 0.03947525471448898, 0.03009573556482792, NaN, NaN, NaN, NaN, NaN, NaN], [0.017164628952741623, 0.028738657012581825, 0.06823595613241196, 0.08604145050048828, 0.04855107143521309, 0.24198594689369202, 0.008688676171004772, 0.003311790293082595, 0.059665460139513016, 0.08214288204908371, 0.34741461277008057, 0.15404720604419708, 0.18822570145130157, 0.19501997530460358, 0.062469229102134705, 0.08181142061948776, 0.013090993277728558, 0.025600923225283623, 0.0045991819351911545, 0.007844633422791958, 0.0066622160375118256, 0.0006054755649529397, 0.01805841363966465, 0.0025927021633833647, 0.0006796378293074667, 0.012531430460512638, 0.18806973099708557, 0.04688132554292679, 0.005460845306515694, 0.053047653287649155, 0.013497358188033104, 0.040136244148015976, 0.022071214392781258, 0.31691932678222656, 0.07654344290494919, NaN, NaN, NaN, NaN, NaN], [0.04490135982632637, 0.02318926900625229, 0.15967297554016113, 0.36984479427337646, 0.027114713564515114, 0.1867561787366867, 0.04668368771672249, 0.02171866036951542, 0.05653616786003113, 0.08818016946315765, 0.14142879843711853, 0.002535451203584671, 0.06232175603508949, 0.12099058926105499, 0.16113655269145966, 0.003571689361706376, 0.007330529857426882, 0.0009176949388347566, 0.011351491324603558, 0.0005700239562429488, 0.0001114286933443509, 0.0023790227714926004, 0.011217805556952953, 0.004490875173360109, 0.00038650527130812407, 0.0025467458181083202, 0.048559535294771194, 0.22723886370658875, 0.0019670024048537016, 0.0002542402071412653, 0.027445662766695023, 0.015111691318452358, 0.029036840423941612, 0.2144545316696167, 0.4208240211009979, 0.013829981908202171, NaN, NaN, NaN, NaN], [0.07898441702127457, 0.817236065864563, 0.29267793893814087, 0.16063392162322998, 0.31295838952064514, 0.9265751838684082, 0.1967003047466278, 0.5436303615570068, 0.2332589328289032, 0.04864489659667015, 0.5440958142280579, 0.8931991457939148, 0.9993566870689392, 0.9798612594604492, 0.03687797114253044, 0.11162849515676498, 0.06633912026882172, 0.017337389290332794, 0.030477523803710938, 0.0024834000505506992, 0.001867939718067646, 0.03932232782244682, 0.1628599613904953, 0.14192035794258118, 0.2944621741771698, 0.21811458468437195, 0.42557209730148315, 0.2638176381587982, 0.14630424976348877, 0.0005040403339080513, 0.32521945238113403, 0.2411627173423767, 0.28287336230278015, 0.40539565682411194, 0.1682160645723343, 0.08244442939758301, 0.001218001707457006, NaN, NaN, NaN], [0.051174335181713104, 0.009388554841279984, 0.15813162922859192, 0.3707107603549957, 0.02142486348748207, 0.01361497025936842, 0.01679075136780739, 0.00489152641966939, 0.08238242566585541, 0.07653495669364929, 0.14888693392276764, 0.003932347521185875, 0.1416105329990387, 0.05760091543197632, 0.13266737759113312, 0.20973265171051025, 0.07712213695049286, 0.20427735149860382, 0.025535617023706436, 0.4053865373134613, 0.41131824254989624, 0.030548784881830215, 0.060146916657686234, 0.012079673819243908, 0.01592317223548889, 0.0048461491242051125, 0.0021770852617919445, 0.09957096725702286, 0.1170588806271553, 0.13386258482933044, 0.16141492128372192, 0.004613581579178572, 0.015190798789262772, 0.003683852730318904, 0.1389266699552536, 0.07006954401731491, 0.1815212517976761, 0.17825333774089813, NaN, NaN], [0.00042274355655536056, 0.0019217034569010139, 0.0013128711143508554, 0.004135955590754747, 0.004101510625332594, 0.004091422073543072, 0.0013299065176397562, 0.0007323773461394012, 0.006002569571137428, 0.003528070170432329, 0.004258603788912296, 0.04385730251669884, 0.006557406857609749, 0.0025679266545921564, 0.1728060394525528, 0.3360293209552765, 0.0046190484426915646, 0.024437543004751205, 0.03736568242311478, 0.023848971351981163, 0.05927197262644768, 0.0542423352599144, 0.09209144860506058, 0.023972967639565468, 0.000766670098528266, 0.0006589474505744874, 0.0007115502958185971, 0.00637162895873189, 0.012912634760141373, 0.014624576084315777, 0.0019432539120316505, 0.05897590517997742, 0.0038116518408060074, 0.0016802565660327673, 0.011611220426857471, 0.025170182809233665, 0.04455949738621712, 0.0020357028115540743, 0.14134161174297333, NaN], [0.0034927180968225002, 0.014745223335921764, 0.025302981957793236, 0.04650698974728584, 0.0658985823392868, 0.10278132557868958, 0.009682145901024342, 0.010841106064617634, 0.1757735013961792, 0.03157021477818489, 0.006062814965844154, 0.2611170709133148, 0.3153221011161804, 0.08490109443664551, 0.13624651730060577, 0.187117338180542, 0.005916869733482599, 0.020901108160614967, 0.0559980571269989, 0.0324174202978611, 0.008547084406018257, 0.044511571526527405, 0.04880741238594055, 0.05289075896143913, 0.038245368748903275, 0.003611604683101177, 0.002279189880937338, 0.01790045015513897, 0.008863909170031548, 0.01127588003873825, 0.005861865822225809, 0.17173975706100464, 0.009364882484078407, 0.005221609957516193, 0.012455414980649948, 0.007264893501996994, 0.016177698969841003, 0.008824422955513, 0.18642237782478333, 0.0006185321253724396]], [[0.11855445802211761, 0.018203705549240112, 0.014699782244861126, 0.005997231230139732, 0.012317956425249577, 0.005482070613652468, 0.020501872524619102, 0.04173066467046738, 0.028033137321472168, 0.007907108403742313, 0.13633504509925842, 0.11779958009719849, 0.02402079664170742, 0.08686818182468414, 0.19919154047966003, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.015789268538355827, 0.07802969217300415, 0.024552250280976295, 0.007203033193945885, 0.015197299420833588, 0.0086579704657197, 0.005928180180490017, 0.015956610441207886, 0.019966211169958115, 0.002508557867258787, 0.048071712255477905, 0.0452260747551918, 0.027286410331726074, 0.034357864409685135, 0.19209280610084534, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7560696601867676, 0.09646204113960266, 0.24264514446258545, 0.03150765225291252, 0.15196740627288818, 0.027980739250779152, 0.025865402072668076, 0.037002913653850555, 0.02429634891450405, 0.014392002485692501, 0.11331582069396973, 0.2883520722389221, 0.24113057553768158, 0.5529852509498596, 0.13967400789260864, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6593953371047974, 0.14735713601112366, 0.007992099039256573, 0.03938791900873184, 0.047611087560653687, 0.002478603972122073, 0.00756214139983058, 0.01120123453438282, 0.017771385610103607, 0.011085578240454197, 0.01766165718436241, 0.07185176759958267, 0.01590064913034439, 0.05699647217988968, 0.22524236142635345, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8214750289916992, 0.5506035089492798, 0.04117008298635483, 0.00517136137932539, 0.5628769993782043, 0.013714980334043503, 0.018153639510273933, 0.019494647160172462, 0.02796507254242897, 0.003693098435178399, 0.052905939519405365, 0.024033749476075172, 0.017759546637535095, 0.154443621635437, 0.2181331366300583, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.47579920291900635, 0.4996025860309601, 0.02201933227479458, 0.032786499708890915, 0.003352785250172019, 0.402157723903656, 0.028392860665917397, 0.03425603359937668, 0.017302367836236954, 0.007774383760988712, 0.03628184646368027, 0.015436487272381783, 0.09682580828666687, 0.09163853526115417, 0.1807471215724945, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6324970722198486, 0.5132108926773071, 0.14723047614097595, 0.10531618446111679, 0.14770705997943878, 0.01965152472257614, 0.16446776688098907, 0.023718399927020073, 0.014144167304039001, 0.003392518265172839, 0.03989372402429581, 0.048702552914619446, 0.05385157838463783, 0.06003360450267792, 0.2021118402481079, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2804942727088928, 0.4447323679924011, 0.40719398856163025, 0.15280602872371674, 0.5485119223594666, 0.006256175693124533, 0.005905789323151112, 0.0894087627530098, 0.014159541577100754, 0.0037697115913033485, 0.08780182898044586, 0.04568948596715927, 0.08344046771526337, 0.08309336006641388, 0.1791403889656067, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.38668709993362427, 0.3767029941082001, 0.5765653848648071, 0.14457443356513977, 0.830109715461731, 0.558448314666748, 0.2105703204870224, 0.015437009744346142, 0.0802588015794754, 0.0035789015237241983, 0.009509528055787086, 0.011719968169927597, 0.04601259157061577, 0.015442220494151115, 0.02989899180829525, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.42374563217163086, 0.4557475447654724, 0.5995064973831177, 0.22240440547466278, 0.8298278450965881, 0.26192477345466614, 0.5618261694908142, 0.2755923569202423, 0.03321446478366852, 0.014314521104097366, 0.030895033851265907, 0.0061126528307795525, 0.0033166268840432167, 0.0021476708352565765, 0.12580153346061707, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4742293357849121, 0.32335561513900757, 0.5931060910224915, 0.0772920548915863, 0.3757626712322235, 0.211185023188591, 0.42018893361091614, 0.37329575419425964, 0.26276469230651855, 0.012583179399371147, 0.3317490220069885, 0.002885210793465376, 0.011435287073254585, 0.00757939275354147, 0.1435183733701706, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.21439705789089203, 0.17853425443172455, 0.32548797130584717, 0.06489395350217819, 0.64824378490448, 0.1159982681274414, 0.19616922736167908, 0.27417391538619995, 0.6047332286834717, 0.1810707151889801, 0.034782104194164276, 0.10310898721218109, 0.0316632017493248, 0.025309519842267036, 0.09833981841802597, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.19860051572322845, 0.10174965113401413, 0.08606765419244766, 0.053267233073711395, 0.11251617968082428, 0.2378872036933899, 0.16651752591133118, 0.1490997076034546, 0.4605393707752228, 0.18029887974262238, 0.1883857697248459, 0.007075145840644836, 0.25310245156288147, 0.08171047270298004, 0.15088772773742676, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2976968586444855, 0.21286718547344208, 0.04716610535979271, 0.025928588584065437, 0.1317281424999237, 0.12927810847759247, 0.2939497232437134, 0.23276808857917786, 0.5986261367797852, 0.05386120826005936, 0.05668044835329056, 0.025143466889858246, 0.007965278811752796, 0.03647890314459801, 0.16275253891944885, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.34472423791885376, 0.33325105905532837, 0.5841152667999268, 0.8456752300262451, 0.4377557933330536, 0.4159393310546875, 0.33224907517433167, 0.1488359123468399, 0.2203720510005951, 0.7425854206085205, 0.7086009383201599, 0.5293036699295044, 0.2777566909790039, 0.22530661523342133, 0.09936152398586273, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01888529770076275, 0.5547894835472107, 0.0062187607400119305, 0.02304725907742977, 0.007431741803884506, 0.05333258956670761, 0.13557927310466766, 0.09608769416809082, 0.011193820275366306, 0.006900292821228504, 0.007560353726148605, 0.018807610496878624, 0.018169475719332695, 0.07717052102088928, 0.1439915895462036, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.045791856944561005, 0.14471176266670227, 0.057932548224925995, 0.15441685914993286, 0.011981116607785225, 0.030152589082717896, 0.13976308703422546, 0.003811573376879096, 0.010053272359073162, 0.1557283103466034, 0.05080341920256615, 0.00967743806540966, 0.003085661679506302, 0.003445286303758621, 0.08783376961946487, 0.12484697252511978, 0.1276315450668335, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.010936958715319633, 0.0031021125614643097, 0.009866965003311634, 0.09017129242420197, 0.02775183692574501, 0.0016267865430563688, 0.01958146132528782, 0.003049993421882391, 0.009465858340263367, 0.022049162536859512, 0.013875926844775677, 0.002902107546105981, 0.0008567434852011502, 0.0034160439390689135, 0.13799139857292175, 0.15841424465179443, 0.03031034581363201, 0.02654799446463585, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10994840413331985, 0.15032780170440674, 0.0035718681756407022, 0.1491042822599411, 0.020450405776500702, 0.013510379940271378, 0.47067153453826904, 0.6447877883911133, 0.18023402988910675, 0.1876010298728943, 0.011866661719977856, 0.006677938625216484, 0.0005242988117970526, 0.004238110035657883, 0.29615819454193115, 0.13769303262233734, 0.09575259685516357, 0.025977646932005882, 0.052591271698474884, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06992093473672867, 0.2791251242160797, 0.006900451611727476, 0.053067900240421295, 0.010168666951358318, 0.0023874202743172646, 0.05137968435883522, 0.06462283432483673, 0.11192043125629425, 0.10690896213054657, 0.009735661558806896, 0.04335656389594078, 0.0031411510426551104, 0.011707558296620846, 0.14929862320423126, 0.15085087716579437, 0.15096567571163177, 0.09222358465194702, 0.028469638898968697, 0.0012114758137613535, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24040630459785461, 0.43853774666786194, 0.0175826046615839, 0.06282828748226166, 0.03055599145591259, 0.20223812758922577, 0.5439046025276184, 0.8139520287513733, 0.30283859372138977, 0.4911571145057678, 0.09772597998380661, 0.1337594985961914, 0.08667796850204468, 0.03606351464986801, 0.12256386131048203, 0.16431185603141785, 0.07204771786928177, 0.05053501948714256, 0.012478960677981377, 0.05114812031388283, 0.00039714027661830187, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03999294713139534, 0.1864590346813202, 0.003897173795849085, 0.04184543341398239, 0.0012414547381922603, 0.025941016152501106, 0.05348599702119827, 0.5434274673461914, 0.012460692785680294, 0.31306707859039307, 0.06930337846279144, 0.0021947044879198074, 0.023592861369252205, 0.04260588437318802, 0.01969532109797001, 0.1666734665632248, 0.06891340762376785, 0.013632094487547874, 0.018171580508351326, 0.002599227475002408, 0.0009873181115835905, 0.0006481229793280363, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.053744781762361526, 0.006899113766849041, 0.0563664473593235, 0.12695427238941193, 0.012777185067534447, 0.08455551415681839, 0.11441048979759216, 0.13062608242034912, 0.19371363520622253, 0.6254263520240784, 0.24294114112854004, 0.020724456757307053, 0.019838949665427208, 0.022365091368556023, 0.1131007969379425, 0.14423918724060059, 0.12251336872577667, 0.10176724940538406, 0.33380815386772156, 0.1583750993013382, 0.023372141644358635, 0.026839546859264374, 0.06730155646800995, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11661048978567123, 0.35882315039634705, 0.03118491731584072, 0.06881216168403625, 0.014698721468448639, 0.0038598491810262203, 0.1485612690448761, 0.39066970348358154, 0.07792866975069046, 0.22571811079978943, 0.040231697261333466, 0.265895277261734, 0.2000368982553482, 0.1125464141368866, 0.24931347370147705, 0.2790219187736511, 0.15446610748767853, 0.015893638134002686, 0.03619629144668579, 0.003051391802728176, 0.00038247412885539234, 0.0007123185787349939, 0.010222047567367554, 0.0010863485513255, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03291217237710953, 0.23853188753128052, 0.04644821211695671, 0.031600918620824814, 0.045192934572696686, 0.0019951597787439823, 0.11113008856773376, 0.36339887976646423, 0.010439107194542885, 0.20188210904598236, 0.027288423851132393, 0.21054767072200775, 0.04143378138542175, 0.0853629931807518, 0.2336580902338028, 0.26870372891426086, 0.10405707359313965, 0.00916238222271204, 0.058617573231458664, 0.0049601029604673386, 0.0005682760966010392, 0.004407011903822422, 0.03309918940067291, 0.0036104319151490927, 0.12174393236637115, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07334253191947937, 0.14656193554401398, 0.004660916980355978, 0.03353964164853096, 0.00998624786734581, 0.00235390174202621, 0.04832129552960396, 0.031250230967998505, 0.0017524310387670994, 0.10710166394710541, 0.04863408952951431, 0.11276239901781082, 0.00949337612837553, 0.024303043261170387, 0.5020502805709839, 0.05985519662499428, 0.14893494546413422, 0.09544339030981064, 0.18974637985229492, 0.1120084673166275, 0.28269606828689575, 0.4275827407836914, 0.12184610962867737, 0.40095797181129456, 0.08120625466108322, 0.27448615431785583, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15921767055988312, 0.18694822490215302, 0.011401425115764141, 0.15920288860797882, 0.0017978762043640018, 0.00600996520370245, 0.1401643455028534, 0.08585444837808609, 0.05989503860473633, 0.2726706564426422, 0.041456613689661026, 0.0019109381828457117, 0.0026012342423200607, 0.00675933575257659, 0.05683350935578346, 0.06809581816196442, 0.09586934000253677, 0.10229554027318954, 0.057183876633644104, 0.25635847449302673, 0.19582371413707733, 0.4237477481365204, 0.37648820877075195, 0.48733898997306824, 0.20777222514152527, 0.24944597482681274, 0.45371755957603455, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.6248686909675598, 0.8166397213935852, 0.05456394702196121, 0.3034517765045166, 0.0032548136077821255, 0.03656908869743347, 0.3933179974555969, 0.635881781578064, 0.4090532660484314, 0.6309216618537903, 0.09238837659358978, 0.01225167978554964, 0.0038302247412502766, 0.05015851929783821, 0.4316881597042084, 0.05513762682676315, 0.16880887746810913, 0.02300925739109516, 0.03029457852244377, 0.032050080597400665, 0.0745139941573143, 0.08332593739032745, 0.5048279166221619, 0.051856089383363724, 0.16889351606369019, 0.22218117117881775, 0.29087209701538086, 0.03443009778857231, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.6506885886192322, 0.26984432339668274, 0.19192098081111908, 0.45030322670936584, 0.018604522570967674, 0.06438936293125153, 0.16284945607185364, 0.46218666434288025, 0.2198290228843689, 0.6063108444213867, 0.13934792578220367, 0.19822801649570465, 0.009406321682035923, 0.07906869053840637, 0.39550670981407166, 0.07503295689821243, 0.22708888351917267, 0.011672623455524445, 0.03240634873509407, 0.051372844725847244, 0.0555996336042881, 0.1055832952260971, 0.27455389499664307, 0.019383858889341354, 0.29115474224090576, 0.25329896807670593, 0.3762655258178711, 0.06596359610557556, 0.027243560180068016, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.6516265273094177, 0.3494286835193634, 0.13445304334163666, 0.40472084283828735, 0.05377691984176636, 0.043724507093429565, 0.6220480799674988, 0.09338771551847458, 0.1620686650276184, 0.8232020139694214, 0.17699383199214935, 0.03535428270697594, 4.775904380949214e-05, 0.000580178399104625, 0.13870029151439667, 0.15851522982120514, 0.22386471927165985, 0.13473065197467804, 0.10273782163858414, 0.539568305015564, 0.23089595139026642, 0.2947250008583069, 0.2566256523132324, 0.08758009225130081, 0.04963833838701248, 0.026406293734908104, 0.02359875850379467, 0.06999926269054413, 0.014701825566589832, 0.008440684527158737, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.40970566868782043, 0.3527304232120514, 0.004458754323422909, 0.09938450157642365, 0.006175781134516001, 0.014084810391068459, 0.22543573379516602, 0.4835565686225891, 0.025563040748238564, 0.39703506231307983, 0.00602720445021987, 0.0051488312892615795, 0.0008810341823846102, 0.0033910071942955256, 0.2277533859014511, 0.1888987272977829, 0.22277534008026123, 0.06621028482913971, 0.04940320923924446, 0.013609242625534534, 0.012980671599507332, 0.0275713000446558, 0.5000426769256592, 0.025658253580331802, 0.28077542781829834, 0.21061377227306366, 0.1005047932267189, 0.0123829934746027, 0.005874408408999443, 0.04495157673954964, 0.007559731602668762, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19487805664539337, 0.1991150975227356, 0.010765495710074902, 0.08231080323457718, 0.014791524969041348, 0.005413876846432686, 0.2905171811580658, 0.06453394889831543, 0.003980779554694891, 0.08378233760595322, 0.012941073626279831, 0.009292078204452991, 0.0008543379371985793, 0.002103410428389907, 0.1794004589319229, 0.10630622506141663, 0.1130438968539238, 0.04711592569947243, 0.14829613268375397, 0.0012987125664949417, 0.0009870391804724932, 0.002409427659586072, 0.10731083154678345, 0.010861101560294628, 0.02266101725399494, 0.22295407950878143, 0.37738272547721863, 0.21324896812438965, 0.09625840187072754, 0.01478838175535202, 0.004724964965134859, 0.13376930356025696, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12092277407646179, 0.17967110872268677, 0.0018819703254848719, 0.04615653306245804, 0.002711376640945673, 0.0007180452230386436, 0.10793514549732208, 0.09669310599565506, 0.0005949889309704304, 0.15432700514793396, 0.015202132984995842, 0.003636009059846401, 0.00047353014815598726, 0.0022874167189002037, 0.22825637459754944, 0.0042772903107106686, 0.006450775545090437, 0.00791113544255495, 0.01871791109442711, 0.02349945716559887, 0.036059893667697906, 0.09560179710388184, 0.01157363597303629, 0.020316841080784798, 0.002858342370018363, 0.0015840751584619284, 0.03869258984923363, 0.04008479043841362, 0.0456826388835907, 0.061234306544065475, 0.32812535762786865, 0.4548730254173279, 0.048923686146736145, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14498451352119446, 0.2535317540168762, 0.027076847851276398, 0.14632807672023773, 0.0057570356875658035, 0.011071202345192432, 0.31473973393440247, 0.2956455647945404, 0.07720959931612015, 0.1944134682416916, 0.008117430843412876, 0.0006636073812842369, 0.0008167477208189666, 0.0018315445631742477, 0.15913215279579163, 0.034464891999959946, 0.04304976761341095, 0.0730237364768982, 0.07959159463644028, 0.156441330909729, 0.14927342534065247, 0.37836754322052, 0.2500280439853668, 0.265838086605072, 0.038285933434963226, 0.0458042174577713, 0.2175784856081009, 0.055615901947021484, 0.32925114035606384, 0.23017114400863647, 0.5254709720611572, 0.3807608187198639, 0.4477500319480896, 0.3941081464290619, NaN, NaN, NaN, NaN, NaN, NaN], [0.22215187549591064, 0.47823596000671387, 0.018273456022143364, 0.13293205201625824, 0.0049734353087842464, 0.0265207476913929, 0.27213141322135925, 0.33180302381515503, 0.1344960778951645, 0.335622638463974, 0.010143149644136429, 0.0012862810399383307, 0.00035499766818247736, 0.0037611438892781734, 0.27220219373703003, 0.024431752040982246, 0.057854264974594116, 0.009785568341612816, 0.015689833089709282, 0.010099711827933788, 0.022971261292696, 0.026158222928643227, 0.08270542323589325, 0.00771379703655839, 0.023359954357147217, 0.06216609850525856, 0.1452798992395401, 0.010090651921927929, 0.13497084379196167, 0.023736534640192986, 0.06422590464353561, 0.2799428105354309, 0.34307411313056946, 0.27198341488838196, 0.018816450610756874, NaN, NaN, NaN, NaN, NaN], [0.3673586845397949, 0.057844266295433044, 0.06040150299668312, 0.09888742864131927, 0.023171812295913696, 0.05270017683506012, 0.11794743686914444, 0.1507657766342163, 0.008498218841850758, 0.09498187899589539, 0.003615680383518338, 0.010834122076630592, 0.00024780313833616674, 0.0017297717276960611, 0.20351538062095642, 0.032250434160232544, 0.07008427381515503, 0.003495490411296487, 0.011726448312401772, 0.013232100754976273, 0.021211393177509308, 0.02240551821887493, 0.050749149173498154, 0.0020511853508651257, 0.034987252205610275, 0.05167752131819725, 0.10231753438711166, 0.017492327839136124, 0.0036121474113315344, 0.0030979528091847897, 0.14347726106643677, 0.4107814431190491, 0.18759746849536896, 0.28042495250701904, 0.02327493391931057, 0.023935986682772636, NaN, NaN, NaN, NaN], [0.6060628294944763, 0.1373525857925415, 0.13755829632282257, 0.4113396406173706, 0.07285188883543015, 0.014519162476062775, 0.5372579097747803, 0.0630655512213707, 0.14564833045005798, 0.695697009563446, 0.06662726402282715, 0.006644518580287695, 1.2849791346525308e-05, 0.00011718441965058446, 0.13694217801094055, 0.17385193705558777, 0.24280618131160736, 0.0901411697268486, 0.1509939581155777, 0.5964542627334595, 0.18189039826393127, 0.25377142429351807, 0.39126867055892944, 0.11990400403738022, 0.04869762808084488, 0.06967514008283615, 0.0491257943212986, 0.1536286324262619, 0.04553663358092308, 0.006321897264569998, 0.008409527130424976, 0.01950901933014393, 0.028066763654351234, 0.039955586194992065, 0.08575458079576492, 0.02489100769162178, 0.0107131227850914, NaN, NaN, NaN], [0.16518473625183105, 0.10184229910373688, 0.002064367523416877, 0.05309450253844261, 0.004080682527273893, 0.012669779360294342, 0.18988992273807526, 0.5354599356651306, 0.004024976398795843, 0.07357845455408096, 0.00022774768876843154, 0.00034433722612448037, 4.428778629517183e-05, 0.00011935137445107102, 0.17481543123722076, 0.18693126738071442, 0.25040745735168457, 0.07803116738796234, 0.06071358174085617, 0.018153348937630653, 0.012512190267443657, 0.012858238071203232, 0.18478038907051086, 0.008756724186241627, 0.14063727855682373, 0.16963867843151093, 0.06472224742174149, 0.008233368396759033, 0.010625114664435387, 0.04533438757061958, 0.004584541078656912, 0.04685693234205246, 0.3269248306751251, 0.13935554027557373, 0.022706659510731697, 0.015514994971454144, 0.09856907278299332, 0.009564985521137714, NaN, NaN], [0.060375016182661057, 0.09738604724407196, 0.004719918128103018, 0.05357348173856735, 0.007510221563279629, 0.002087255474179983, 0.1777726411819458, 0.04658319056034088, 0.0022654803469777107, 0.02657914347946644, 0.002838509390130639, 0.0023206211626529694, 0.00029234393150545657, 0.0006460589938797057, 0.15720529854297638, 0.10220125317573547, 0.06584151834249496, 0.046970706433057785, 0.16499453783035278, 0.0008504274883307517, 0.000721337681170553, 0.0015187861863523722, 0.050142802298069, 0.005332621280103922, 0.005509581416845322, 0.0572623535990715, 0.172898530960083, 0.12213093042373657, 0.0640687644481659, 0.004657925106585026, 0.002522988012060523, 0.028443191200494766, 0.29674383997917175, 0.3544806241989136, 0.20916549861431122, 0.09151047468185425, 0.014975211583077908, 0.0019209993770346045, 0.07398010790348053, NaN], [0.006292517296969891, 0.056422796100378036, 0.003871192689985037, 0.016857203096151352, 0.0060961381532251835, 0.01021772250533104, 0.02558758109807968, 0.004345982801169157, 0.003136568469926715, 0.011386821046471596, 0.0007550015579909086, 0.014218548312783241, 0.002899263286963105, 0.00665974011644721, 0.1386014223098755, 0.014319260604679585, 0.019726725295186043, 0.010809341445565224, 0.06728478521108627, 0.024899542331695557, 0.06927011907100677, 0.2726534307003021, 0.06849226355552673, 0.06274150311946869, 0.0032663261517882347, 0.007571991998702288, 0.011041088029742241, 0.0653790682554245, 0.06552072614431381, 0.10165777057409286, 0.05923810228705406, 0.20752549171447754, 0.1128133162856102, 0.041725482791662216, 0.12833572924137115, 0.10405165702104568, 0.2233171910047531, 0.10715138167142868, 0.3742898404598236, 0.43902406096458435]], [[0.3582096993923187, 0.12323450297117233, 0.41414904594421387, 0.12697191536426544, 0.2567327618598938, 0.12921607494354248, 0.303745299577713, 0.26060354709625244, 0.2067556530237198, 0.0739586353302002, 0.038356974720954895, 0.018690073862671852, 0.019858568906784058, 0.03828525170683861, 0.09448481351137161, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.034560851752758026, 0.06147807836532593, 0.09719342738389969, 0.03090484067797661, 0.05040246620774269, 0.10769589245319366, 0.28225648403167725, 0.03959896042943001, 0.04561477154493332, 0.015998149290680885, 0.010396423749625683, 0.0027313604950904846, 0.02088637463748455, 0.02540828473865986, 0.1729334592819214, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.031599532812833786, 0.03154325857758522, 0.01938430592417717, 0.10300880670547485, 0.07719798386096954, 0.3211115002632141, 0.5488157868385315, 0.6110779047012329, 0.03511836752295494, 0.03874386474490166, 0.02549627609550953, 0.08684590458869934, 0.1071673184633255, 0.10855282843112946, 0.09071482717990875, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05947110056877136, 0.046990834176540375, 0.001917339744977653, 0.019972380250692368, 0.14856000244617462, 0.10937333106994629, 0.7613639235496521, 0.43800127506256104, 0.038890283554792404, 0.0702563002705574, 0.052807219326496124, 0.20175476372241974, 0.09827514737844467, 0.19838720560073853, 0.1799801141023636, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.010548654943704605, 0.056933727115392685, 0.0004277318366803229, 0.0005220972234383225, 0.03427216783165932, 0.15697234869003296, 0.44382861256599426, 0.28639304637908936, 0.1278306096792221, 0.0589531809091568, 0.07240739464759827, 0.21584689617156982, 0.623681902885437, 0.39177897572517395, 0.053747572004795074, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.012333033606410027, 0.11936485022306442, 0.0015480549773201346, 0.05167163908481598, 0.003915506415069103, 0.05033823475241661, 0.18770258128643036, 0.5247471332550049, 0.13492631912231445, 0.0999734029173851, 0.02801361307501793, 0.04943297058343887, 0.067798912525177, 0.02220618724822998, 0.04863249137997627, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.023225123062729836, 0.03936318680644035, 0.0654693990945816, 0.0780135840177536, 0.03190883249044418, 0.007237496320158243, 0.3230750560760498, 0.11266676336526871, 0.3152024447917938, 0.12503208220005035, 0.08215073496103287, 0.20814812183380127, 0.054794978350400925, 0.014369799755513668, 0.31165388226509094, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.021642545238137245, 0.05032852664589882, 0.10916808992624283, 0.14173567295074463, 0.025796422734856606, 0.002176823327317834, 0.004212724044919014, 0.11230720579624176, 0.2761599123477936, 0.18545517325401306, 0.30032697319984436, 0.18456220626831055, 0.1202857494354248, 0.02383211813867092, 0.22383396327495575, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.014165909960865974, 0.030938388779759407, 0.019327908754348755, 0.025021186098456383, 0.018685894086956978, 0.058899857103824615, 0.05705944076180458, 0.013411193154752254, 0.27564239501953125, 0.14192135632038116, 0.4484158754348755, 0.49174171686172485, 0.42328834533691406, 0.5148258805274963, 0.024227913469076157, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.030343737453222275, 0.035576362162828445, 0.011198173277080059, 0.0029289661906659603, 0.004656192846596241, 0.19044476747512817, 0.14425727725028992, 0.14593322575092316, 0.02429576776921749, 0.03922351822257042, 0.03158531337976456, 0.3954472541809082, 0.18761666119098663, 0.829915463924408, 0.05755764618515968, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07378673553466797, 0.08269044756889343, 0.008506381884217262, 0.004565858747810125, 0.0033621611073613167, 0.47163471579551697, 0.3437289595603943, 0.16293375194072723, 0.0103234788402915, 0.006828381214290857, 0.025515833869576454, 0.13491219282150269, 0.23380780220031738, 0.7675665616989136, 0.06853343546390533, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.19539110362529755, 0.20751968026161194, 0.012997383251786232, 0.004634191282093525, 0.004486567340791225, 0.10301963984966278, 0.2361651211977005, 0.10510270297527313, 0.007245894055813551, 0.02498149685561657, 0.005201807711273432, 0.12586773931980133, 0.2985144853591919, 0.741521954536438, 0.061252206563949585, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3654796779155731, 0.656768798828125, 0.02389511466026306, 0.057929087430238724, 0.025417884811758995, 0.2985052168369293, 0.29244741797447205, 0.15614598989486694, 0.02199239283800125, 0.027919312939047813, 0.024499662220478058, 0.0015409317566081882, 0.18344998359680176, 0.05587974563241005, 0.11099682748317719, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.24996283650398254, 0.30432745814323425, 0.08651068061590195, 0.27794384956359863, 0.10948572307825089, 0.32318809628486633, 0.40224379301071167, 0.24700750410556793, 0.016620514914393425, 0.03902489319443703, 0.01563531532883644, 0.008603462018072605, 0.029363060370087624, 0.20380347967147827, 0.1635625809431076, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08184575289487839, 0.05559774115681648, 0.012900986708700657, 0.004766350146383047, 0.02465618960559368, 0.0658264234662056, 0.16982027888298035, 0.09995799511671066, 0.1946410834789276, 0.03345171734690666, 0.026332948356866837, 0.010880211368203163, 0.01684177853167057, 0.011932285502552986, 0.13059602677822113, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.19101674854755402, 0.0880991518497467, 0.25550922751426697, 0.3376496732234955, 0.25425824522972107, 0.2177356481552124, 0.35922226309776306, 0.13405567407608032, 0.2859460711479187, 0.47983312606811523, 0.235154390335083, 0.26708394289016724, 0.2646999657154083, 0.4890832304954529, 0.0349225178360939, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12788966298103333, 0.14897412061691284, 0.18708589673042297, 0.1539590060710907, 0.06750026345252991, 0.06459501385688782, 0.24742794036865234, 0.0008040289394557476, 0.08417094498872757, 0.08338519930839539, 0.09756942838430405, 0.05163748189806938, 0.06044981628656387, 0.1204136312007904, 0.005185095127671957, 0.12878015637397766, 0.05999259278178215, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00823432207107544, 0.006774595472961664, 0.011488616466522217, 0.031759701669216156, 0.014620696194469929, 0.015192853286862373, 0.015498323366045952, 0.001623230637051165, 0.04214249551296234, 0.022796856239438057, 0.0813785269856453, 0.058821164071559906, 0.018185952678322792, 0.030505431815981865, 0.13797427713871002, 0.16734670102596283, 0.0018487111665308475, 0.002184537472203374, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07304069399833679, 0.17316529154777527, 0.0638275146484375, 0.06216027960181236, 0.10879980027675629, 0.2286580353975296, 0.12489848583936691, 0.06798849999904633, 0.12340370565652847, 0.11364749073982239, 0.33209869265556335, 0.7156579494476318, 0.917570948600769, 0.8780012726783752, 0.004697424825280905, 0.06620991975069046, 0.4480140209197998, 0.42379117012023926, 0.3748236298561096, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04041377454996109, 0.06032548099756241, 0.013153426349163055, 0.12010756880044937, 0.032379359006881714, 0.02533758245408535, 0.03651244193315506, 0.05168384686112404, 0.05184069648385048, 0.20407944917678833, 0.10554968565702438, 0.5571502447128296, 0.039276935160160065, 0.10380254685878754, 0.1458612084388733, 0.1498516947031021, 0.091057188808918, 0.11073686927556992, 0.05954570695757866, 0.00012444167805369943, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.025283029302954674, 0.14580176770687103, 0.0262577123939991, 0.01834816485643387, 0.02426275424659252, 0.5010125637054443, 0.025797395035624504, 0.08120379596948624, 0.10846563428640366, 0.05807282403111458, 0.047331083565950394, 0.01890925131738186, 0.041984543204307556, 0.021773895248770714, 0.12734822928905487, 0.15789009630680084, 0.05178086459636688, 0.2272004932165146, 0.05532779544591904, 0.002530630910769105, 0.00011625503975665197, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11099886894226074, 0.272359162569046, 0.07267793267965317, 0.02685651369392872, 0.04662291333079338, 0.6599292755126953, 0.15850403904914856, 0.1944371908903122, 0.02196124941110611, 0.18415939807891846, 0.2094753533601761, 0.11699666827917099, 0.8625363111495972, 0.6611498594284058, 0.034588079899549484, 0.05158510431647301, 0.42307329177856445, 0.4962795376777649, 0.6637455821037292, 0.11636865884065628, 0.027691489085555077, 0.059323750436306, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10045554488897324, 0.003808635985478759, 0.012772331945598125, 0.008206314407289028, 0.016907531768083572, 0.2308196723461151, 0.04502535238862038, 0.16794730722904205, 0.14683513343334198, 0.07804886251688004, 0.12962646782398224, 0.03242946416139603, 0.45433515310287476, 0.3931583762168884, 0.023861808702349663, 0.1440366506576538, 0.37752795219421387, 0.42684903740882874, 0.13104133307933807, 0.0449170246720314, 0.0360451340675354, 0.007316120434552431, 0.03281773626804352, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.020261207595467567, 0.011864200234413147, 0.013516101986169815, 0.00783876795321703, 0.006360001862049103, 0.5825139880180359, 0.27136117219924927, 0.28645893931388855, 0.002775657456368208, 0.05587191879749298, 0.01021821890026331, 0.03437367081642151, 0.37942126393318176, 0.11788230389356613, 0.047214996069669724, 0.018571142107248306, 0.11001976579427719, 0.16728174686431885, 0.33147770166397095, 0.29621925950050354, 0.11174014210700989, 0.46736985445022583, 0.18467408418655396, 0.05186863988637924, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3444993495941162, 0.4299255907535553, 0.3897337317466736, 0.11608962714672089, 0.07001375406980515, 0.1826992928981781, 0.3195875883102417, 0.1513850837945938, 0.014436168596148491, 0.25265297293663025, 0.18822813034057617, 0.20145024359226227, 0.648497998714447, 0.6856710314750671, 0.13566814363002777, 0.0193540807813406, 0.11997552216053009, 0.4339123070240021, 0.4291674792766571, 0.22741732001304626, 0.21840345859527588, 0.4310562014579773, 0.16546283662319183, 0.05634206160902977, 0.03477246314287186, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.37375974655151367, 0.2605052888393402, 0.636468231678009, 0.14340142905712128, 0.5107957124710083, 0.683059811592102, 0.3617965579032898, 0.3775153160095215, 0.0734284520149231, 0.5245854258537292, 0.5329803228378296, 0.541839063167572, 0.8546188473701477, 0.8892531991004944, 0.08003345131874084, 0.07166115939617157, 0.34385329484939575, 0.5272834300994873, 0.4769807457923889, 0.34829023480415344, 0.19288644194602966, 0.1752767115831375, 0.3240547180175781, 0.026788396760821342, 0.09653788805007935, 0.14339366555213928, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1478864699602127, 0.26107946038246155, 0.2706110179424286, 0.022070137783885002, 0.08394861966371536, 0.7104908227920532, 0.22173403203487396, 0.18465854227542877, 0.3481738865375519, 0.02706378884613514, 0.14399166405200958, 0.24452990293502808, 0.3432118594646454, 0.3138853907585144, 0.0603480227291584, 0.09568949043750763, 0.2010803371667862, 0.1452081948518753, 0.13633964955806732, 0.13264110684394836, 0.11369673907756805, 0.18754418194293976, 0.10573749244213104, 0.12209529429674149, 0.3772747814655304, 0.4260762333869934, 0.1448964774608612, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03315366804599762, 0.109662726521492, 0.165960431098938, 0.03089676797389984, 0.00589095801115036, 0.7119044065475464, 0.04612211138010025, 0.03627030551433563, 0.019800378009676933, 0.02169116772711277, 0.07954178750514984, 0.014483828097581863, 0.3210127055644989, 0.25073835253715515, 0.021559905260801315, 0.1600937843322754, 0.32966408133506775, 0.46643200516700745, 0.2761552929878235, 0.1128716766834259, 0.16030451655387878, 0.13808301091194153, 0.12019707262516022, 0.08980843424797058, 0.23569302260875702, 0.18699060380458832, 0.06252679228782654, 0.02190866880118847, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1801593005657196, 0.7095129489898682, 0.41699883341789246, 0.14223065972328186, 0.03218872845172882, 0.8857168555259705, 0.325775682926178, 0.46090880036354065, 0.31827157735824585, 0.19596631824970245, 0.36584827303886414, 0.568932831287384, 0.05918605625629425, 0.12899020314216614, 0.03239220380783081, 0.09671676903963089, 0.3181785047054291, 0.5044789910316467, 0.5311775803565979, 0.43058764934539795, 0.24623769521713257, 0.546705424785614, 0.20948244631290436, 0.5971428155899048, 0.15125280618667603, 0.21692372858524323, 0.08393274247646332, 0.0805632621049881, 0.11463441699743271, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15587098896503448, 0.007851594127714634, 0.38951343297958374, 0.26023998856544495, 0.2678505480289459, 0.04164084047079086, 0.060063086450099945, 0.06729273498058319, 0.019880756735801697, 0.0442759171128273, 0.10040930658578873, 0.1083277016878128, 0.0003995952138211578, 0.001039322349242866, 0.14095477759838104, 0.17538371682167053, 0.005170984659343958, 0.01562126912176609, 0.012803001329302788, 0.0004321248270571232, 0.003303500125184655, 0.010391591116786003, 0.0083633316680789, 0.001453742035664618, 0.0005911564221605659, 0.001968160504475236, 0.018067756667733192, 0.0012553221313282847, 0.0006174716982059181, 0.0014710418181493878, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08899319916963577, 0.2356371134519577, 0.40766164660453796, 0.08200893551111221, 0.14033742249011993, 0.12043434381484985, 0.050508081912994385, 0.04391980916261673, 0.2084629088640213, 0.07807423919439316, 0.06514080613851547, 0.6571899652481079, 0.6522034406661987, 0.4899447560310364, 0.0237458273768425, 0.00964878499507904, 0.07296860218048096, 0.1732037365436554, 0.2482636272907257, 0.018695944920182228, 0.04061494395136833, 0.019565006718039513, 0.048743683844804764, 0.15582872927188873, 0.0506676621735096, 0.08059392869472504, 0.2691291868686676, 0.4701274335384369, 0.05269847437739372, 0.15863555669784546, 0.011098350398242474, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3269592225551605, 0.23715397715568542, 0.21103474497795105, 0.29856637120246887, 0.031984660774469376, 0.019636303186416626, 0.2648169696331024, 0.0041971527971327305, 0.6909844875335693, 0.5414000153541565, 0.4092715382575989, 0.02185220457613468, 0.006548420060425997, 0.013211028650403023, 0.06752441078424454, 0.023792432621121407, 0.42975902557373047, 0.3812340199947357, 0.23295366764068604, 0.2699258625507355, 0.32472288608551025, 0.04527096822857857, 0.2556793987751007, 0.5905154347419739, 0.8116171360015869, 0.684613823890686, 0.13916483521461487, 0.05671815946698189, 0.0401710644364357, 0.30002903938293457, 0.014873968437314034, 0.1109585389494896, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.40959432721138, 0.2696213126182556, 0.4055677354335785, 0.265968382358551, 0.12281941622495651, 0.10883577167987823, 0.16766701638698578, 0.053767129778862, 0.028326192870736122, 0.5353591442108154, 0.3247348368167877, 0.03339260071516037, 0.1199125200510025, 0.14055927097797394, 0.07849014550447464, 0.07327478379011154, 0.42313894629478455, 0.7821765542030334, 0.6752634048461914, 0.18926696479320526, 0.27897483110427856, 0.1972714066505432, 0.26650866866111755, 0.21928414702415466, 0.6610813736915588, 0.8023169040679932, 0.32853400707244873, 0.043605707585811615, 0.04177317023277283, 0.5147100687026978, 0.014965414069592953, 0.041893746703863144, 0.10476090759038925, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0703776553273201, 0.17115768790245056, 0.14820680022239685, 0.014450321905314922, 0.036940984427928925, 0.4336852431297302, 0.18269671499729156, 0.1382565200328827, 0.5314536690711975, 0.05019254609942436, 0.11642822623252869, 0.17526941001415253, 0.3684784173965454, 0.3591882586479187, 0.09016428142786026, 0.09543995559215546, 0.1369307041168213, 0.1906978189945221, 0.1367466300725937, 0.17180036008358002, 0.12260185182094574, 0.13847540318965912, 0.1559406965970993, 0.13510896265506744, 0.4644373655319214, 0.6843520402908325, 0.2938932180404663, 0.08134166151285172, 0.16692468523979187, 0.35020914673805237, 0.0983358696103096, 0.26928237080574036, 0.11322443932294846, 0.14002281427383423, NaN, NaN, NaN, NaN, NaN, NaN], [0.020959746092557907, 0.2473447471857071, 0.04995026811957359, 0.032434724271297455, 0.004538285546004772, 0.38885483145713806, 0.04268676042556763, 0.035024866461753845, 0.14864443242549896, 0.14174208045005798, 0.13687251508235931, 0.021197974681854248, 0.4566997289657593, 0.37854352593421936, 0.051512595266103745, 0.17294523119926453, 0.44891712069511414, 0.5596615076065063, 0.3151743412017822, 0.15508009493350983, 0.20398668944835663, 0.18162229657173157, 0.14380685985088348, 0.09279182553291321, 0.25614914298057556, 0.37145668268203735, 0.2047339379787445, 0.05775143578648567, 0.06389063596725464, 0.19947569072246552, 0.07508620619773865, 0.162083700299263, 0.036575064063072205, 0.05963924527168274, 0.02704720012843609, NaN, NaN, NaN, NaN, NaN], [0.11558277904987335, 0.8023946285247803, 0.11340320110321045, 0.07801315933465958, 0.012690390460193157, 0.363363116979599, 0.22989940643310547, 0.28700947761535645, 0.3164795935153961, 0.28987860679626465, 0.20186272263526917, 0.5113669037818909, 0.04614659398794174, 0.13675883412361145, 0.05756649002432823, 0.09450869262218475, 0.5263407230377197, 0.5685468316078186, 0.6246378421783447, 0.5457862615585327, 0.4288109838962555, 0.7265884876251221, 0.4213257133960724, 0.7441360354423523, 0.37028953433036804, 0.4906199276447296, 0.24940308928489685, 0.2854059636592865, 0.25606390833854675, 0.06486664712429047, 0.03651905804872513, 0.215606689453125, 0.16494624316692352, 0.07126681506633759, 0.0978088453412056, 0.18553400039672852, NaN, NaN, NaN, NaN], [0.13439694046974182, 0.004173143766820431, 0.22800596058368683, 0.19857077300548553, 0.1396344006061554, 0.007145485375076532, 0.03306930512189865, 0.026599518954753876, 0.02599666267633438, 0.04890456795692444, 0.0713912844657898, 0.040079280734062195, 0.00020046728604938835, 0.0004629320465028286, 0.13767622411251068, 0.19233128428459167, 0.0069253402762115, 0.019198253750801086, 0.024288823828101158, 0.0006626379326917231, 0.0032825330272316933, 0.012745865620672703, 0.02121213637292385, 0.004573441576212645, 0.001344278221949935, 0.010449343360960484, 0.07998955249786377, 0.008849495090544224, 0.005957764107733965, 0.00281895836815238, 0.0006993816932663321, 0.0011300387559458613, 0.0034355262760072947, 0.006048144306987524, 0.0007683978183194995, 0.00029024321702308953, 0.0009215899626724422, NaN, NaN, NaN], [0.21178027987480164, 0.5613860487937927, 0.18598653376102448, 0.13814353942871094, 0.06437420845031738, 0.1469835489988327, 0.09205848723649979, 0.07043211162090302, 0.3314816355705261, 0.1618121713399887, 0.0553976409137249, 0.7871544361114502, 0.7398563027381897, 0.533365786075592, 0.06109875440597534, 0.00490582175552845, 0.09978753328323364, 0.17523892223834991, 0.18201382458209991, 0.025161702185869217, 0.0351867638528347, 0.008898423984646797, 0.033712878823280334, 0.06612548977136612, 0.044598400592803955, 0.0818907842040062, 0.31783777475357056, 0.6522275805473328, 0.26521986722946167, 0.31609129905700684, 0.0543142631649971, 0.07028744369745255, 0.06436092406511307, 0.12702754139900208, 0.4257008731365204, 0.05356784537434578, 0.20406562089920044, 0.022904740646481514, NaN, NaN], [0.308572918176651, 0.1810312271118164, 0.10904403775930405, 0.38784971833229065, 0.013434378430247307, 0.011286276392638683, 0.26633715629577637, 0.0027595413848757744, 0.7609409689903259, 0.7608016729354858, 0.6143397688865662, 0.036307673901319504, 0.013564765453338623, 0.02826162986457348, 0.07738469541072845, 0.02933959849178791, 0.5456263422966003, 0.4945109188556671, 0.26123103499412537, 0.3237256109714508, 0.3705388903617859, 0.04209306091070175, 0.3351372182369232, 0.658141016960144, 0.8126230239868164, 0.8673186898231506, 0.28273773193359375, 0.11254162341356277, 0.17348313331604004, 0.7003386616706848, 0.1474425047636032, 0.36997753381729126, 0.41849759221076965, 0.091117262840271, 0.03724836930632591, 0.036747273057699203, 0.47380825877189636, 0.017722588032484055, 0.0920308530330658, NaN], [0.1500416249036789, 0.027276279404759407, 0.32022449374198914, 0.45847558975219727, 0.23693141341209412, 0.1596660166978836, 0.2821829915046692, 0.005833256058394909, 0.32143598794937134, 0.14477354288101196, 0.029714325442910194, 0.15291856229305267, 0.007731991354376078, 0.029727784916758537, 0.12283544987440109, 0.1429738998413086, 0.11406568437814713, 0.30407312512397766, 0.04420004412531853, 0.050888776779174805, 0.009020227938890457, 0.026264725252985954, 0.20154790580272675, 0.284900963306427, 0.16813665628433228, 0.6384625434875488, 0.35198092460632324, 0.0041788192465901375, 0.017796171829104424, 0.06702794879674911, 0.017356209456920624, 0.11703062057495117, 0.363391250371933, 0.08829980343580246, 0.0006652214215137064, 0.002063008025288582, 0.01232101023197174, 0.0010344748152419925, 0.005295889917761087, 0.10532692819833755]], [[0.06378140300512314, 0.013955923728644848, 0.058693334460258484, 0.014864355325698853, 0.02882157638669014, 0.02533077634871006, 0.013877282850444317, 0.02919653430581093, 0.029733512550592422, 0.010929838754236698, 0.2184230536222458, 0.404588907957077, 0.5044611692428589, 0.4171900451183319, 0.18600669503211975, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.09787620604038239, 0.3741878271102905, 0.1718531847000122, 0.22170154750347137, 0.11211875081062317, 0.06884550303220749, 0.023903023451566696, 0.00765330670401454, 0.043831951916217804, 0.04742401838302612, 0.08705892413854599, 0.19904442131519318, 0.1439688503742218, 0.08975595235824585, 0.124632827937603, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.024405136704444885, 0.006321595516055822, 0.03571266308426857, 0.0050111510790884495, 0.01807553507387638, 6.11300565651618e-05, 0.0022184934932738543, 0.002461126074194908, 0.00987271312624216, 0.03944821655750275, 0.02587837167084217, 0.009154303930699825, 0.018459370359778404, 0.07083768397569656, 0.2838045060634613, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.02829434722661972, 0.05303699150681496, 0.03342747688293457, 0.026768406853079796, 0.06776657700538635, 0.0015663451049476862, 0.0066550131887197495, 0.028257621452212334, 0.02201445959508419, 0.024995435029268265, 0.014314326457679272, 0.019762825220823288, 0.019060753285884857, 0.09995586425065994, 0.2721303105354309, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.011709636077284813, 0.13082386553287506, 0.3091292977333069, 0.012390679679811, 0.06598176062107086, 0.0025066242087632418, 0.008877930231392384, 0.03396160528063774, 0.01681593246757984, 0.01466491911560297, 0.12272557616233826, 0.010357965715229511, 0.009066522121429443, 0.12291242927312851, 0.3062548041343689, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05738264322280884, 0.12342102825641632, 0.7862259149551392, 0.20355252921581268, 0.007363088894635439, 0.0717976987361908, 0.032159313559532166, 0.018495721742510796, 0.0034321516286581755, 0.0013732254737988114, 0.006710591726005077, 0.0023603499867022038, 0.007563347462564707, 0.05948156490921974, 0.12037239223718643, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.015277753584086895, 0.006394209805876017, 0.6686000227928162, 0.29117655754089355, 0.06745831668376923, 0.2462725043296814, 0.06154515966773033, 0.015117062255740166, 0.004134421236813068, 0.0023558081593364477, 0.08952713012695312, 0.04650713875889778, 0.023702487349510193, 0.01321239210665226, 0.09701406955718994, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.028385812416672707, 0.012191490270197392, 0.27066752314567566, 0.18411272764205933, 0.040896836668252945, 0.48173367977142334, 0.02650352008640766, 0.07071101665496826, 0.007758310064673424, 0.001958101289346814, 0.01839292421936989, 0.023066602647304535, 0.03435399383306503, 0.03657263144850731, 0.029525745660066605, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04876675456762314, 0.422792911529541, 0.22041767835617065, 0.2559551000595093, 0.08884847164154053, 0.01230597123503685, 0.025672338902950287, 0.003895203350111842, 0.022659877315163612, 0.0043840305879712105, 0.007982935756444931, 0.010924039408564568, 0.06971067935228348, 0.0061518345028162, 0.21563398838043213, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.015657104551792145, 0.02366352081298828, 0.07373688369989395, 0.10379613190889359, 0.013535204343497753, 0.07323776930570602, 0.048540983349084854, 0.008235346525907516, 0.01638718694448471, 0.012322558090090752, 0.073370561003685, 0.03809332847595215, 0.021602218970656395, 0.003090204205363989, 0.23272792994976044, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.018198516219854355, 0.011175387538969517, 0.02189311571419239, 0.012938260100781918, 0.09454065561294556, 0.010837653651833534, 0.04214898869395256, 0.03231353685259819, 0.2788335978984833, 0.02807164192199707, 0.0381515808403492, 0.013884211890399456, 0.014051362872123718, 0.00934662390500307, 0.24102351069450378, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01114112138748169, 0.11382883787155151, 0.017900465056300163, 0.008639826439321041, 0.024639632552862167, 0.020821422338485718, 0.022935912013053894, 0.04321465268731117, 0.055257730185985565, 0.0561254657804966, 0.006350866984575987, 0.034159135073423386, 0.001170721254311502, 0.00040716465446166694, 0.2438717484474182, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01806582696735859, 0.014762195758521557, 0.02654433250427246, 0.025726040825247765, 0.03240499645471573, 0.020733002573251724, 0.04244884103536606, 0.02047092467546463, 0.13412125408649445, 0.512605607509613, 0.5156171321868896, 0.023306455463171005, 0.0489252470433712, 0.06594526767730713, 0.173824280500412, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.018763704225420952, 0.010509289801120758, 0.06387435644865036, 0.02487548068165779, 0.10975509881973267, 0.01984621025621891, 0.06460897624492645, 0.03137337416410446, 0.1802622228860855, 0.7354047894477844, 0.7864400148391724, 0.1003832221031189, 0.007522855885326862, 0.14785504341125488, 0.08187610656023026, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.02117479033768177, 0.061044495552778244, 0.02157888375222683, 0.021421663463115692, 0.04618487507104874, 0.05167240649461746, 0.01054168026894331, 0.009977741166949272, 0.0295058935880661, 0.008349624462425709, 0.02268156036734581, 0.026699911803007126, 0.020697196945548058, 0.013632250018417835, 0.13365623354911804, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2602275013923645, 0.0514441579580307, 0.4731021821498871, 0.5077798962593079, 0.22717851400375366, 0.04740440100431442, 0.27564913034439087, 0.24302659928798676, 0.05887439846992493, 0.3509802222251892, 0.6124410033226013, 0.11394976824522018, 0.0489780493080616, 0.04593530669808388, 0.01042554248124361, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.032066281884908676, 0.1349876970052719, 0.04647025838494301, 0.02243492752313614, 0.02574889175593853, 0.03298051655292511, 0.026965852826833725, 0.3248708248138428, 0.005728535819798708, 0.08351098001003265, 0.1499667763710022, 0.16844461858272552, 0.05473209172487259, 0.05656114220619202, 0.10718395560979843, 0.1283751130104065, 0.06695841252803802, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.005181984044611454, 0.0008690498070791364, 0.00864254217594862, 0.00306740403175354, 0.10709173232316971, 0.0007182863773778081, 0.004329775460064411, 0.010956686921417713, 0.06760676205158234, 0.010445973835885525, 0.012115269899368286, 0.06696799397468567, 0.0054829977452754974, 0.025371035560965538, 0.13854098320007324, 5.319380943547003e-05, 9.114345448324457e-05, 0.7905611991882324, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03556624799966812, 0.11754146218299866, 0.010577056556940079, 0.008073115721344948, 0.06965696066617966, 0.0032990325707942247, 0.011276635341346264, 0.09485359489917755, 0.10517128556966782, 0.0125450249761343, 0.007751243654638529, 0.0650070384144783, 0.0006160335033200681, 0.002038064645603299, 0.4774436056613922, 0.10777772217988968, 0.19019582867622375, 0.12566408514976501, 0.295462429523468, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13858208060264587, 0.06875398755073547, 0.01532802265137434, 0.10744626820087433, 0.18273182213306427, 0.002165634883567691, 0.069672591984272, 0.11672408878803253, 0.005795653443783522, 0.0880894884467125, 0.05771886929869652, 0.025581423193216324, 0.03904194384813309, 0.07354751974344254, 0.14365413784980774, 2.4899240088416263e-05, 2.9243250537547283e-05, 0.0014855118934065104, 3.888772698701359e-05, 0.9169090986251831, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16291819512844086, 0.050931405276060104, 0.14806726574897766, 0.2683573365211487, 0.2810481786727905, 0.002092417562380433, 0.012745368294417858, 0.01212888304144144, 0.014305775985121727, 0.17753903567790985, 0.1299620419740677, 0.10299177467823029, 0.21836693584918976, 0.06576120108366013, 0.12406044453382492, 3.5349924587535497e-07, 4.689470642915694e-06, 0.02691131830215454, 1.3325815416465048e-05, 0.19568589329719543, 0.956480085849762, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12156791239976883, 0.39120492339134216, 0.1209033653140068, 0.08395244181156158, 0.29989197850227356, 0.044024936854839325, 0.023133939132094383, 0.05934688448905945, 0.02561376802623272, 0.024757277220487595, 0.04535222053527832, 0.11912120133638382, 0.02126661129295826, 0.03811139240860939, 0.248785600066185, 0.08490768820047379, 0.04920955002307892, 0.012384464032948017, 0.04339546710252762, 0.010612337850034237, 0.05702771991491318, 0.7263003587722778, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.106705442070961, 0.8169862627983093, 0.1967339813709259, 0.01375850010663271, 0.13418887555599213, 0.16134029626846313, 0.005958847235888243, 0.09247319400310516, 0.04806499928236008, 0.025876127183437347, 0.08311128616333008, 0.22926460206508636, 0.05653654783964157, 0.04726153612136841, 0.20836575329303741, 0.16491760313510895, 0.04815620183944702, 0.0007595600909553468, 0.006606678944081068, 0.0006115635624155402, 0.0007167417788878083, 0.0015418223338201642, 0.0024032427463680506, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04722486063838005, 0.04722658172249794, 0.05176655203104019, 0.00462702801451087, 0.20528024435043335, 0.0011717488523572683, 0.004415996838361025, 0.014451048336923122, 0.028127426281571388, 0.007240481209009886, 0.004411954898387194, 0.10081291943788528, 0.07703132927417755, 0.033158108592033386, 0.21852079033851624, 0.012053201906383038, 0.18336322903633118, 0.0033893296495079994, 0.22584111988544464, 0.004534169565886259, 0.003455487545579672, 0.30805450677871704, 0.5499533414840698, 0.13390673696994781, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.032722555100917816, 0.027063244953751564, 0.014943713322281837, 0.0013555125333368778, 0.016471203416585922, 0.005467826500535011, 0.02999643050134182, 0.014794600196182728, 0.03837134689092636, 0.004397213459014893, 0.01024235412478447, 0.04855721816420555, 0.05723624676465988, 0.051476139575242996, 0.2643129825592041, 0.02224119007587433, 0.09969844669103622, 0.01827961951494217, 0.1828235685825348, 0.009660250507295132, 0.005268027540296316, 0.13511976599693298, 0.39505934715270996, 0.1772008240222931, 0.6222725510597229, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.052069392055273056, 0.003948261961340904, 0.01313212513923645, 0.010319330729544163, 0.04011767730116844, 0.00066552241332829, 0.01502715889364481, 0.007099903654307127, 0.16779832541942596, 0.03226454555988312, 0.052614975720644, 0.014822165481746197, 0.002071568975225091, 0.001763610984198749, 0.05304422974586487, 0.19008594751358032, 0.025696618482470512, 0.004118501208722591, 0.03605509176850319, 0.002144730417057872, 0.0023362801875919104, 0.16961191594600677, 0.015426162630319595, 0.016875047236680984, 0.017404966056346893, 0.032629188150167465, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.022045070305466652, 0.036587294191122055, 0.06798984855413437, 0.040110163390636444, 0.5405737161636353, 0.015278805047273636, 0.02948732301592827, 0.034845639020204544, 0.27487096190452576, 0.008005083538591862, 0.012681123800575733, 0.10707750916481018, 0.02124345488846302, 0.00868641585111618, 0.4183328449726105, 0.1594686657190323, 0.03835373371839523, 0.021387629210948944, 0.028402678668498993, 0.12163796275854111, 0.1348690688610077, 0.027878204360604286, 0.016979072242975235, 0.009301519952714443, 0.047045812010765076, 0.103324294090271, 0.0978349894285202, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07479816675186157, 0.018890362232923508, 0.2873721718788147, 0.028116360306739807, 0.7967413067817688, 0.008446138352155685, 0.020726248621940613, 0.018564706668257713, 0.33813604712486267, 0.003492887830361724, 0.010393181815743446, 0.18903475999832153, 0.00443642633035779, 0.0231452826410532, 0.42231008410453796, 0.08206925541162491, 0.0482555516064167, 0.03066202998161316, 0.14434732496738434, 0.10149279236793518, 0.1536794900894165, 0.16425268352031708, 0.00592045346274972, 0.002011190867051482, 0.030538976192474365, 0.015422381460666656, 0.0400862954556942, 0.6933969259262085, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07108656316995621, 0.0021144712809473276, 0.0671088695526123, 0.03148089721798897, 0.7113023400306702, 0.006737539079040289, 0.2500847280025482, 0.023258471861481667, 0.23158760368824005, 0.011219021864235401, 0.04227704927325249, 0.03650788217782974, 0.15078191459178925, 0.09633734077215195, 0.15066072344779968, 0.11962933838367462, 0.08867897093296051, 0.023231033235788345, 0.019267449155449867, 0.06578893214464188, 0.01314490009099245, 0.028238458558917046, 0.2009190320968628, 0.005505711771547794, 0.024347275495529175, 0.005847027525305748, 0.13606473803520203, 0.11386173218488693, 0.6883828639984131, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04487757384777069, 0.009540342725813389, 0.2420971691608429, 0.01275626104325056, 0.3918483257293701, 0.0218670591711998, 0.022137846797704697, 0.08132637292146683, 0.11900310963392258, 0.000993919325992465, 0.03630243241786957, 0.087126724421978, 0.0003738462692126632, 0.02454514056444168, 0.14072805643081665, 0.004133098293095827, 0.007605875376611948, 0.380069762468338, 0.01569206453859806, 0.3162667751312256, 0.06185031309723854, 0.003268925240263343, 0.007663627155125141, 0.00711404625326395, 0.0016827658982947469, 0.002885768422856927, 0.009058460593223572, 0.0104479705914855, 0.0013903286308050156, 0.9176042079925537, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0048965876922011375, 0.019337626174092293, 0.002879639156162739, 0.0027576948050409555, 0.04260760545730591, 0.003218113211914897, 0.003307115286588669, 0.026640478521585464, 0.011750566773116589, 0.0005104524316266179, 9.575913281878456e-05, 0.057879798114299774, 0.004244217649102211, 0.00609983503818512, 0.28528884053230286, 0.19946889579296112, 0.004915847908705473, 0.0015343156410381198, 0.012221671640872955, 0.003153382334858179, 0.0001576353097334504, 0.0020530277397483587, 0.003957398701459169, 0.010446527041494846, 0.012547693215310574, 0.03473197668790817, 0.06650777161121368, 0.014228541404008865, 0.02601468935608864, 0.0018418998224660754, 0.08826413750648499, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0335795059800148, 0.030716734007000923, 0.023829646408557892, 0.03415534272789955, 0.08875380456447601, 0.0019310596399009228, 0.017619425430893898, 0.012105603702366352, 0.002468202030286193, 0.010380377061665058, 0.01267782598733902, 0.10606792569160461, 0.0014069904573261738, 0.0004161447286605835, 0.19442977011203766, 0.14040440320968628, 0.29221969842910767, 0.09665771573781967, 0.2947876751422882, 0.00611721258610487, 0.012681002728641033, 0.7610099911689758, 0.27993685007095337, 0.19895455241203308, 0.07963719218969345, 0.025141140446066856, 0.30299919843673706, 0.4374280273914337, 0.12315846234560013, 0.011889583431184292, 0.00027308438438922167, 0.03226177766919136, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17404082417488098, 0.05758971348404884, 0.12847737967967987, 0.07598815858364105, 0.49957963824272156, 0.003085564589127898, 0.05114232748746872, 0.011464038863778114, 0.06926580518484116, 0.06844814121723175, 0.06813240051269531, 0.08604259043931961, 0.004740274045616388, 0.009239559061825275, 0.19994765520095825, 0.22362156212329865, 0.19648011028766632, 0.02122899703681469, 0.12822405993938446, 0.013841216452419758, 0.009505078196525574, 0.4746513366699219, 0.1753886640071869, 0.09167484194040298, 0.038334570825099945, 0.04122844338417053, 0.14653263986110687, 0.17874038219451904, 0.023550381883978844, 0.014212163165211678, 0.001423373818397522, 0.0059451088309288025, 0.09707646816968918, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.011875619180500507, 0.026503771543502808, 0.054018229246139526, 0.01668175496160984, 0.3499281406402588, 0.01803278550505638, 0.01878167688846588, 0.01221490278840065, 0.15005004405975342, 0.0046301730908453465, 0.005843435879796743, 0.032064031809568405, 0.010490885935723782, 0.00555034726858139, 0.27147379517555237, 0.167328879237175, 0.06208498775959015, 0.010482249781489372, 0.03574186563491821, 0.0675959512591362, 0.06477286666631699, 0.04995346441864967, 0.05412250757217407, 0.009984727017581463, 0.03347667679190636, 0.11074735969305038, 0.16135196387767792, 0.07774785906076431, 0.01735900156199932, 0.007863441482186317, 0.019525114446878433, 0.005842071026563644, 0.1275986284017563, 0.0955328494310379, NaN, NaN, NaN, NaN, NaN, NaN], [0.0646943747997284, 0.047236885875463486, 0.11903148144483566, 0.02203843556344509, 0.4764179587364197, 0.008550588972866535, 0.013687309809029102, 0.008890991099178791, 0.32491248846054077, 0.011557912454009056, 0.009869826957583427, 0.0921611338853836, 0.0031256151851266623, 0.016340140253305435, 0.3438139855861664, 0.05032582953572273, 0.03989394009113312, 0.02223959006369114, 0.07248460501432419, 0.04305185005068779, 0.04872481897473335, 0.09144517779350281, 0.0032577940728515387, 0.000561918190214783, 0.015125684440135956, 0.018474824726581573, 0.0519116036593914, 0.7149417400360107, 0.023930398747324944, 0.005549557972699404, 0.0027118371799588203, 0.08418004959821701, 0.22684048116207123, 0.052481237798929214, 0.7548789381980896, NaN, NaN, NaN, NaN, NaN], [0.17560914158821106, 0.007353567518293858, 0.056802812963724136, 0.032415200024843216, 0.4015137553215027, 0.02137722261250019, 0.35710790753364563, 0.018633568659424782, 0.05862341821193695, 0.02506905421614647, 0.018169963732361794, 0.009134531952440739, 0.07779684662818909, 0.07867905497550964, 0.1750962883234024, 0.14971917867660522, 0.12296220660209656, 0.03256092593073845, 0.015910452231764793, 0.08324312418699265, 0.010959222912788391, 0.03249981626868248, 0.2630986273288727, 0.0023772413842380047, 0.021863164380192757, 0.014683729968965054, 0.3797665238380432, 0.26638853549957275, 0.6724205613136292, 0.015757206827402115, 0.01569446735084057, 0.01732691004872322, 0.06738004088401794, 0.17602917551994324, 0.12501026690006256, 0.6636221408843994, NaN, NaN, NaN, NaN], [0.05210466682910919, 0.006375414319336414, 0.22638031840324402, 0.012961659580469131, 0.3225522041320801, 0.012402641586959362, 0.024030247703194618, 0.056293144822120667, 0.11919546872377396, 0.0012290689628571272, 0.027758106589317322, 0.025181178003549576, 0.00022994892788119614, 0.012616506777703762, 0.1375768631696701, 0.0045495470985770226, 0.007598123978823423, 0.48235079646110535, 0.017675379291176796, 0.30638325214385986, 0.03773635998368263, 0.0025513810105621815, 0.013349749147891998, 0.011474208906292915, 0.002688285429030657, 0.009704438969492912, 0.024301802739501, 0.030528949573636055, 0.006023744586855173, 0.9289764761924744, 0.008095184341073036, 0.015121471136808395, 0.003912394400686026, 0.005678378511220217, 0.005922055337578058, 0.0012866485631093383, 0.9431078433990479, NaN, NaN, NaN], [0.005459210369735956, 0.03143180534243584, 0.0014205367770045996, 0.0012642937945201993, 0.01687682792544365, 0.007108580321073532, 0.004234722815454006, 0.017920657992362976, 0.003724986221641302, 0.0002761750074569136, 2.4563792976550758e-05, 0.011889445595443249, 0.0013067404506728053, 0.002636768389493227, 0.19040453433990479, 0.25144028663635254, 0.013477480970323086, 0.004043558146804571, 0.02197866141796112, 0.005731666926294565, 0.00035365403164178133, 0.0028230457101017237, 0.003569219959899783, 0.00616231607273221, 0.023324957117438316, 0.07691453397274017, 0.11847300082445145, 0.025281671434640884, 0.05239935964345932, 0.002384425140917301, 0.16120819747447968, 0.011955172754824162, 0.09212952852249146, 0.03993848338723183, 0.017148757353425026, 0.01459744293242693, 0.0018050760263577104, 0.08139479160308838, NaN, NaN], [0.031027475371956825, 0.05656901001930237, 0.0113890515640378, 0.024300340563058853, 0.03550150617957115, 0.0024159413296729326, 0.02035972848534584, 0.01581081561744213, 0.002032301388680935, 0.009238713420927525, 0.01651322841644287, 0.11367840319871902, 0.003108791308477521, 0.00086622079834342, 0.16520220041275024, 0.08713241666555405, 0.22884246706962585, 0.12139283120632172, 0.21789073944091797, 0.00419022049754858, 0.011025986634194851, 0.8093750476837158, 0.24520863592624664, 0.11868450790643692, 0.037659380584955215, 0.014297883957624435, 0.35379931330680847, 0.4382935166358948, 0.17632676661014557, 0.006937071681022644, 0.0007303177262656391, 0.027538392692804337, 0.0690605565905571, 0.3237524628639221, 0.41753751039505005, 0.09520361572504044, 0.013310365378856659, 0.0003602981742005795, 0.032565031200647354, NaN], [0.7154905796051025, 0.15825338661670685, 0.49722805619239807, 0.38231807947158813, 0.39668020606040955, 0.051081933081150055, 0.4188354015350342, 0.3623049259185791, 0.3077245056629181, 0.4494604766368866, 0.7933229804039001, 0.20231026411056519, 0.27286192774772644, 0.2623305022716522, 0.06808917224407196, 0.01268855668604374, 0.009620537050068378, 0.0011078648967668414, 0.01395372860133648, 0.00034480926115065813, 0.0002369812864344567, 0.14032205939292908, 0.12187758088111877, 0.004498081747442484, 6.632315489696339e-05, 0.01873306930065155, 0.07693066447973251, 0.06357964873313904, 0.012718681246042252, 0.02489433065056801, 0.4312428832054138, 0.013737366534769535, 0.0326746366918087, 0.34456172585487366, 0.0668448805809021, 0.006646350026130676, 0.04233057424426079, 0.4123155176639557, 0.007851892150938511, 0.43338367342948914]], [[4.754594192490913e-05, 2.1380438752771624e-08, 2.918067565360616e-08, 2.8621201408896013e-08, 2.499384379461844e-07, 0.0002631827082950622, 5.21495513439163e-10, 2.490414274802788e-08, 1.4592379216082918e-07, 4.660217989282955e-09, 1.3478041793746343e-08, 1.530838318331007e-07, 4.6195887989597395e-05, 8.429636181972455e-06, 0.2157532423734665, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6645432114601135, 0.00044607618474401534, 8.70102576300269e-06, 1.056492124007491e-06, 4.43653931370136e-07, 3.5252294310339494e-06, 0.013106754049658775, 0.0008970960625447333, 5.719662112824153e-07, 3.2791810156140855e-08, 1.0544068729245737e-08, 3.57371057191358e-08, 0.00012361648259684443, 0.0008665899513289332, 0.00011794524471042678, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [5.6636022236489225e-06, 0.771808385848999, 0.2603715658187866, 7.618767995154485e-05, 2.6443340175319463e-05, 1.448297037853763e-08, 1.7459943213449236e-10, 0.0005545829189941287, 1.3129211993145873e-06, 0.0003596498572733253, 1.3187416243454209e-06, 1.2532552773336647e-08, 5.7067543821176514e-05, 1.4676837054139469e-05, 8.822963764032465e-07, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [7.866851170490463e-09, 0.0015575109282508492, 0.5911858677864075, 0.005255529191344976, 0.00012560673349071294, 1.2381517144888221e-08, 1.3975322635251253e-12, 4.631081083061872e-06, 1.8297629367225454e-06, 0.043241821229457855, 0.00025465109501965344, 1.6550380621538352e-07, 1.5873881693551084e-06, 1.3629888329091955e-08, 2.2046858560997862e-08, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.6020940130090366e-10, 3.2446525892737554e-06, 0.1964423805475235, 0.9067507982254028, 4.244087540428154e-05, 3.027215825568419e-05, 6.154020626425449e-10, 3.570748958736658e-07, 2.493328743469192e-08, 1.327106815551815e-07, 5.116170723340474e-05, 7.67620722541551e-09, 6.538175512105227e-07, 1.6885725528936746e-07, 1.9495971503857845e-09, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [4.057985947270026e-09, 1.6926858803500977e-09, 0.00014235911658033729, 0.0026504932902753353, 0.8634750843048096, 1.9555229300749488e-05, 1.294085109293519e-06, 2.6649362894204387e-07, 3.0507638082433175e-10, 5.069419550807197e-09, 1.108148239836737e-07, 1.7377595213474706e-05, 9.726352800498717e-06, 1.823265733946755e-06, 5.869507617717318e-07, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.9094309466893833e-12, 2.4682887027685507e-13, 6.382604444965523e-10, 6.302604549368596e-10, 1.4692274817207363e-05, 0.3734012544155121, 3.483030241113738e-06, 1.1820202594492457e-08, 1.9522692351614523e-09, 1.394072303342181e-13, 1.7670450172535546e-11, 1.716609077107023e-09, 3.7749509829154704e-06, 2.593782255644328e-06, 3.855710133393586e-07, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [8.508453674949124e-08, 1.863478038544031e-09, 1.257351167627263e-10, 5.331373190142763e-11, 3.337832410466035e-08, 1.777973557182122e-05, 0.8244234323501587, 8.755041926633567e-05, 1.7572835409040977e-09, 1.3142270258170718e-11, 7.735358035533546e-13, 4.927841815161038e-11, 5.296478775562719e-07, 0.000259329448454082, 1.8429471282388477e-08, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.2582735964272729e-09, 2.3675827378610848e-06, 5.770066309196409e-07, 5.0431950282536775e-11, 2.6034334410507398e-11, 1.7287857190240175e-07, 9.084228622668888e-06, 0.8877476453781128, 0.0008898449596017599, 7.2106473680833e-08, 1.9634756043274137e-08, 4.930736808433922e-13, 3.217972377456135e-08, 1.2906410120194778e-05, 9.568290160189008e-09, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [2.8039692789860737e-09, 1.3000158105569426e-06, 4.493769978353157e-08, 2.493898698663344e-10, 7.932443764346875e-12, 1.7288407150317653e-08, 2.642636942606913e-10, 3.576151357265189e-05, 0.8324669599533081, 5.240505197434686e-05, 8.11301958947297e-07, 9.422521651814009e-10, 4.6924657937097436e-08, 2.8963553333483105e-08, 6.33739318800508e-08, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [2.873091320410026e-09, 7.32139524188824e-05, 1.393846559949452e-05, 2.2707215663331226e-08, 3.602095333121724e-08, 7.893682235637911e-12, 1.2799745258921386e-13, 1.2971109697446082e-07, 4.534097752184607e-05, 0.7187873721122742, 0.0028858170844614506, 4.860597982769832e-06, 3.316463335067965e-06, 6.64895694058032e-08, 4.189383506769673e-09, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [3.5802516507033033e-10, 3.3775189312024168e-09, 1.689890041234321e-06, 2.72409181434341e-07, 2.3650377656281307e-08, 3.1582386705863996e-10, 4.773196676235644e-14, 6.179980832632381e-11, 1.0790042637154329e-07, 0.00019566719129215926, 0.8666706681251526, 0.00033315850305370986, 7.101260734998505e-07, 3.226231015673875e-08, 6.780910499770698e-09, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [7.800644574729176e-09, 1.700809604265885e-09, 9.215954577257435e-08, 4.046364665555302e-07, 0.00011374137102393433, 5.132134901941754e-06, 5.991689921991394e-10, 9.107053305923429e-11, 5.105777606262407e-11, 3.3974476565390432e-09, 3.904122058884241e-05, 0.65162193775177, 0.00035754009149968624, 6.446759653044865e-05, 8.575011065659055e-07, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [5.410449865905775e-10, 1.9016622998524468e-10, 1.651180719930423e-10, 9.184660809680167e-10, 4.749936000081334e-09, 6.8993631430203095e-06, 9.186856830822876e-10, 1.2120262259107673e-11, 1.0679299241797557e-12, 7.136916383397585e-13, 1.9098522763272285e-10, 9.612936082703527e-06, 0.7662882208824158, 0.00778515450656414, 3.0943773765557125e-08, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0058370670303702354, 0.00017831011791713536, 6.727457275701454e-06, 4.542615897662472e-06, 0.0008248149533756077, 0.04996809363365173, 0.010534689761698246, 8.931134652812034e-05, 2.4081384708551923e-07, 6.080232139993313e-08, 3.077615701840841e-06, 0.00041306819184683263, 0.062034472823143005, 0.37576472759246826, 0.1323644071817398, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.437301367521286, 0.15179137885570526, 0.09085877984762192, 0.06997784972190857, 0.17732757329940796, 0.23180970549583435, 0.11514479666948318, 0.32073739171028137, 0.15501314401626587, 0.1294255405664444, 0.06762269139289856, 0.21488851308822632, 0.2614101469516754, 0.12734454870224, 0.049641113728284836, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.028495818376541138, 0.1544514149427414, 0.06366834789514542, 0.016971074044704437, 0.02302762120962143, 0.054101087152957916, 0.012630121782422066, 0.018889501690864563, 0.004939573351293802, 0.01251249760389328, 0.1164683923125267, 0.009905983693897724, 0.01818472519516945, 0.01017050538212061, 0.04256897792220116, 0.13150663673877716, 0.013105388730764389, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.007633751258254051, 0.002589557319879532, 0.02251260355114937, 0.05040144920349121, 0.032673582434654236, 0.0022981506772339344, 0.00627527991309762, 0.0006094649434089661, 0.01362280547618866, 0.006205975078046322, 0.006417383905500174, 0.0010467394022271037, 0.0010408272501081228, 0.007578521966934204, 0.13823428750038147, 0.16704899072647095, 0.0014066778821870685, 0.003860085504129529, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0074798669666051865, 0.011802621185779572, 0.3115181624889374, 0.22458955645561218, 0.10706131160259247, 0.016402821987867355, 0.046956516802310944, 0.004200803115963936, 0.01468481682240963, 0.014471452683210373, 0.27619558572769165, 0.0038709931541234255, 0.00034889893140643835, 0.0020716534927487373, 0.01783183217048645, 0.14769184589385986, 0.005059333052486181, 0.0053715878166258335, 0.026609797030687332, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.015254770405590534, 0.01172303594648838, 0.002065492793917656, 0.005149758420884609, 0.013159574940800667, 0.001197350095026195, 0.018971139565110207, 0.004385960288345814, 0.06813318282365799, 0.021520443260669708, 0.005575989838689566, 0.001505104242824018, 0.0019181625684723258, 0.005167691968381405, 0.15193934738636017, 0.15381431579589844, 0.05056624114513397, 0.015615872107446194, 0.004382571205496788, 0.00015187788812909275, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.026872141286730766, 0.003412047168239951, 0.03895608335733414, 0.03612855076789856, 0.02536499686539173, 0.03102046251296997, 0.004315483849495649, 0.0027427596505731344, 0.03512648865580559, 0.022632958367466927, 0.05171700567007065, 0.0026941397227346897, 0.0031264815479516983, 0.024213580414652824, 0.12838274240493774, 0.16606314480304718, 0.03878505155444145, 0.01631396822631359, 0.011268166825175285, 0.00036908386391587555, 0.00010962320084217936, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0600903183221817, 0.002928798785433173, 0.0064612883143126965, 0.05414368212223053, 0.029363246634602547, 0.006244697142392397, 0.397325724363327, 0.040878646075725555, 0.005305922590196133, 0.27715954184532166, 0.04618077725172043, 0.008418801240622997, 0.01155431941151619, 0.05281350389122963, 0.025860372930765152, 0.16556474566459656, 0.059035927057266235, 0.018687130883336067, 0.020593103021383286, 0.0006985706277191639, 0.0006753651541657746, 0.01174053642898798, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0013151391176506877, 0.002262294292449951, 0.0012738551013171673, 0.0034272209741175175, 0.0030726443510502577, 0.04279911145567894, 0.008567760698497295, 0.17885291576385498, 0.00929640606045723, 0.001624501310288906, 0.02533317357301712, 0.005113683640956879, 0.027247918769717216, 0.07258909195661545, 0.014188846573233604, 0.16100119054317474, 0.03705580160021782, 0.08672276139259338, 0.05696912482380867, 0.00507472176104784, 0.006951047107577324, 0.0023692583199590445, 0.004235508386045694, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3408622145652771, 0.07445694506168365, 0.03113507851958275, 0.0754152163863182, 0.014415460638701916, 0.002693483140319586, 0.09953030943870544, 0.11086118221282959, 0.5124953985214233, 0.329039990901947, 0.5092117786407471, 0.027396254241466522, 0.055544231086969376, 0.4057520925998688, 0.09588415175676346, 0.288095086812973, 0.011840847320854664, 0.005622565280646086, 0.00535928551107645, 0.0008760345517657697, 0.0004899614141322672, 0.001179057639092207, 0.0010409504175186157, 0.0012723063118755817, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09238530695438385, 0.007053247652947903, 0.0017291916301473975, 0.005093103274703026, 0.0007437380263581872, 0.0014228186337277293, 0.02520381473004818, 0.019087698310613632, 0.47848576307296753, 0.29748132824897766, 0.057576071470975876, 0.01139640249311924, 0.004621520172804594, 0.02937469258904457, 0.015335291624069214, 0.2984195351600647, 0.024577315896749496, 0.008883590810000896, 0.0237559974193573, 0.001871026586741209, 0.002048116410151124, 0.00452006608247757, 0.0067189703695476055, 0.002311990363523364, 0.0035932722967118025, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0720675140619278, 0.012255199253559113, 0.04221949726343155, 0.09128241240978241, 0.009349699132144451, 0.008273615501821041, 0.014371694065630436, 0.01100369542837143, 0.1737149953842163, 0.16746114194393158, 0.1696900725364685, 0.014558696188032627, 0.01365632750093937, 0.0269284937530756, 0.016150163486599922, 0.19755195081233978, 0.08605571836233139, 0.04371126368641853, 0.045333728194236755, 0.005393510684370995, 0.006479238625615835, 0.018500106409192085, 0.012994848191738129, 0.011254888959228992, 0.03004884347319603, 0.011813223361968994, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.052127860486507416, 0.0038822691421955824, 0.01307338010519743, 0.12611117959022522, 0.013002983294427395, 0.054914653301239014, 0.022843925282359123, 0.0017219025176018476, 0.025739489123225212, 0.3090609014034271, 0.10414470732212067, 0.006550551857799292, 0.006861968897283077, 0.010005415417253971, 0.011784915812313557, 0.05165635421872139, 0.44527125358581543, 0.31059694290161133, 0.6649516224861145, 0.027770839631557465, 0.02873762883245945, 0.17512862384319305, 0.06940869987010956, 0.1633579134941101, 0.028000785037875175, 0.003091411432251334, 0.016245586797595024, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.074305959045887, 0.010457544587552547, 0.07050318270921707, 0.4022633135318756, 0.04945780336856842, 0.04771194979548454, 0.4660364091396332, 0.07594453543424606, 0.018491366878151894, 0.1513216346502304, 0.09796185791492462, 0.23858080804347992, 0.011272062547504902, 0.09385059028863907, 0.06640274822711945, 0.19151811301708221, 0.1383962333202362, 0.13229386508464813, 0.35712042450904846, 0.18756243586540222, 0.2871147096157074, 0.5138459801673889, 0.22405852377414703, 0.28785935044288635, 0.04021993279457092, 0.0012617700267583132, 0.004019713494926691, 0.003964945673942566, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.025815313681960106, 0.0033349080476909876, 0.00924734864383936, 0.012487816624343395, 0.03726305067539215, 0.016575457528233528, 0.23753590881824493, 0.025156090036034584, 0.11919926106929779, 0.04390435293316841, 0.0095932362601161, 0.04137176275253296, 0.08216788619756699, 0.1757660061120987, 0.30195334553718567, 0.24189773201942444, 0.08955204486846924, 0.32067012786865234, 0.20245005190372467, 0.11740265786647797, 0.08460556715726852, 0.044664137065410614, 0.025831788778305054, 0.07413194328546524, 0.0068964180536568165, 0.002961511956527829, 0.005619046278297901, 0.0014741680352017283, 0.00546230049803853, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05659867450594902, 0.020075146108865738, 0.01205957867205143, 0.004331792704761028, 0.052221644669771194, 0.0230423454195261, 0.0683140978217125, 0.09752152115106583, 0.2100839763879776, 0.0003861601871903986, 0.0032946986611932516, 0.0004593236662913114, 5.027504084864631e-05, 0.0022022551856935024, 0.14128009974956512, 0.1724659651517868, 0.13219435513019562, 0.15014058351516724, 0.12075512856245041, 0.0006761215627193451, 0.10174072533845901, 0.19516822695732117, 0.009559075348079205, 0.057678524404764175, 0.08239483833312988, 0.0039215064607560635, 0.0027616096194833517, 0.013109313324093819, 0.002305442001670599, 0.00021083203318994492, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08638240396976471, 0.0710444375872612, 0.06771891564130783, 0.17398057878017426, 0.05179189518094063, 0.34193578362464905, 0.2095513492822647, 0.09331211447715759, 0.052257001399993896, 0.006232596468180418, 0.002646914916113019, 0.06318453699350357, 0.019070196896791458, 0.02972061187028885, 0.2659039795398712, 0.19843007624149323, 0.15979865193367004, 0.14398488402366638, 0.41609427332878113, 0.010126790963113308, 0.04840107262134552, 0.7232485413551331, 0.22829605638980865, 0.34322667121887207, 0.08224418759346008, 0.03167981281876564, 0.020198417827486992, 0.013381149619817734, 0.0009459191933274269, 0.006438484415411949, 0.008794432505965233, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.26895081996917725, 0.1478959172964096, 0.3258365988731384, 0.404258131980896, 0.3733697533607483, 0.19055484235286713, 0.19857566058635712, 0.01781378500163555, 0.07512970268726349, 0.11693259328603745, 0.1175057590007782, 0.24425068497657776, 0.20241285860538483, 0.2411348670721054, 0.06638508290052414, 0.30347728729248047, 0.04726674035191536, 0.010849116370081902, 0.12094812840223312, 0.0013257962418720126, 0.0025908409152179956, 0.0014983253786340356, 0.03437754884362221, 0.009621781297028065, 0.006184253375977278, 0.00671237800270319, 0.0018636187305673957, 0.01123903226107359, 0.0035993149504065514, 0.0012990115210413933, 0.00021464838937390596, 0.001025065197609365, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17850612103939056, 0.12822727859020233, 0.17801056802272797, 0.28459492325782776, 0.058830633759498596, 0.03884930908679962, 0.3513718843460083, 0.061017971485853195, 0.06718380004167557, 0.071348175406456, 0.23821549117565155, 0.03658399358391762, 0.03897847980260849, 0.20709341764450073, 0.13892877101898193, 0.2792417109012604, 0.26782968640327454, 0.03489779308438301, 0.07551994919776917, 0.018111348152160645, 0.04002813994884491, 0.03850500285625458, 0.11152958869934082, 0.21995633840560913, 0.07949108630418777, 0.0037619988434016705, 0.03436713665723801, 0.020695386454463005, 0.017524488270282745, 0.010141805745661259, 0.003556826151907444, 0.0020958345849066973, 0.0058519174344837666, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4637373983860016, 0.04377487301826477, 0.15646661818027496, 0.36986854672431946, 0.09056738018989563, 0.23626187443733215, 0.11398540437221527, 0.0026716177817434072, 0.006399102043360472, 0.2626173198223114, 0.20860937237739563, 0.01349638868123293, 0.014208723790943623, 0.042171213775873184, 0.08208009600639343, 0.05386974662542343, 0.6086578965187073, 0.22683310508728027, 0.5828835964202881, 0.02668178826570511, 0.03663201630115509, 0.14977867901325226, 0.2173178791999817, 0.2744499444961548, 0.08338183909654617, 0.008825525641441345, 0.06588608771562576, 0.5592238306999207, 0.17532478272914886, 0.006846817210316658, 0.028904464095830917, 0.01721598580479622, 0.006393561605364084, 0.010461881756782532, NaN, NaN, NaN, NaN, NaN, NaN], [0.13806220889091492, 0.04062362387776375, 0.09515099227428436, 0.37904345989227295, 0.10653041303157806, 0.052835192531347275, 0.5728973150253296, 0.03487204387784004, 0.0029783223289996386, 0.07966885715723038, 0.03475099802017212, 0.13843636214733124, 0.006917618680745363, 0.06183210015296936, 0.1688811033964157, 0.24167264997959137, 0.2504684031009674, 0.15247754752635956, 0.4417489171028137, 0.37691444158554077, 0.47509273886680603, 0.6227271556854248, 0.6949021220207214, 0.5199849605560303, 0.14203055202960968, 0.006932773161679506, 0.02713918127119541, 0.026524275541305542, 0.28478434681892395, 0.05304509028792381, 0.03063105419278145, 0.007391192018985748, 0.001299944007769227, 0.0022179351653903723, 0.0017378581687808037, NaN, NaN, NaN, NaN, NaN], [0.02612869068980217, 0.003477374091744423, 0.007765303365886211, 0.0023155075032263994, 0.018893033266067505, 0.022398637607693672, 0.09549611806869507, 0.004012360703200102, 0.0013466936070472002, 0.0021441734861582518, 0.0004924506065435708, 0.006835760548710823, 0.011635211296379566, 0.023846328258514404, 0.22376547753810883, 0.3587647080421448, 0.13152657449245453, 0.3170546591281891, 0.1872878074645996, 0.17338471114635468, 0.16099165380001068, 0.050314128398895264, 0.07316549867391586, 0.1506616473197937, 0.027928102761507034, 0.013985591009259224, 0.03077181987464428, 0.00928373821079731, 0.01458327379077673, 0.34401679039001465, 0.1675042062997818, 0.008024912327528, 0.00340651860460639, 0.001158604514785111, 0.0004595925274770707, 0.0022153020836412907, NaN, NaN, NaN, NaN], [0.08347997069358826, 0.014491320587694645, 0.015744350850582123, 0.0043899440206587315, 0.05038629099726677, 0.008546282537281513, 0.06458569318056107, 0.03869106248021126, 0.0615551732480526, 0.0002168803766835481, 0.0014501431724056602, 0.00013847390073351562, 1.5032101146061905e-05, 0.0007368824444711208, 0.13783538341522217, 0.18021628260612488, 0.21554027497768402, 0.22428971529006958, 0.28362634778022766, 0.0019759181886911392, 0.19364571571350098, 0.3129161596298218, 0.05571373924612999, 0.43670228123664856, 0.5364305973052979, 0.045233964920043945, 0.02291695959866047, 0.15668357908725739, 0.03788933902978897, 0.0009749932214617729, 0.15011590719223022, 0.009233620017766953, 0.023490505293011665, 0.0018092861864715815, 0.01433361042290926, 0.002351803006604314, 0.00025271173217333853, NaN, NaN, NaN], [0.072405144572258, 0.036094967275857925, 0.060353852808475494, 0.1382489949464798, 0.03810955956578255, 0.1803218573331833, 0.3716851472854614, 0.04992733895778656, 0.002898369450122118, 0.0008571037324145436, 0.00035707451752386987, 0.02692999318242073, 0.003073085332289338, 0.009645520709455013, 0.17640869319438934, 0.18984580039978027, 0.30305740237236023, 0.22004783153533936, 0.5488721132278442, 0.023633448407053947, 0.10360189527273178, 0.8517335653305054, 0.6748489141464233, 0.77315753698349, 0.4876308739185333, 0.2048063576221466, 0.14540305733680725, 0.08473058044910431, 0.012403973378241062, 0.06795734912157059, 0.17164894938468933, 0.18992502987384796, 0.12247806042432785, 0.011528578586876392, 0.009636401198804379, 0.0008312705904245377, 0.013430905528366566, 0.011612125672399998, NaN, NaN], [0.30767515301704407, 0.17313888669013977, 0.17682777345180511, 0.3453424274921417, 0.2732711434364319, 0.18888972699642181, 0.2821650207042694, 0.011036374606192112, 0.013345124199986458, 0.030917862430214882, 0.037141598761081696, 0.14430613815784454, 0.09504004567861557, 0.16429893672466278, 0.0962204858660698, 0.3384567201137543, 0.062264904379844666, 0.014819102361798286, 0.14853152632713318, 0.0019540644716471434, 0.003596463706344366, 0.001872691442258656, 0.11878995597362518, 0.02639206312596798, 0.009769541211426258, 0.011811794713139534, 0.006684192456305027, 0.045877717435359955, 0.019279729574918747, 0.005480214022099972, 0.003932234365493059, 0.006437724456191063, 0.0240105502307415, 0.0011211916571483016, 0.004233745392411947, 0.001469226786866784, 0.0013713098596781492, 0.00014342667418532073, 0.0008160521974787116, NaN], [0.038221023976802826, 0.4632723033428192, 0.022520000115036964, 0.005303966347128153, 0.07163825631141663, 0.030774233862757683, 0.006099082063883543, 0.008936556056141853, 0.02098681591451168, 0.004558844491839409, 0.0029896388296037912, 0.018592750653624535, 0.20478543639183044, 0.08578886091709137, 0.1358346790075302, 0.1837155818939209, 0.5941455364227295, 0.2251758873462677, 0.3662757873535156, 0.039659783244132996, 0.3226933479309082, 0.014135366305708885, 0.028798755258321762, 0.10863638669252396, 0.34925851225852966, 0.03930900990962982, 0.08864527195692062, 0.10118203610181808, 0.05801505595445633, 0.11320658773183823, 0.05595846846699715, 0.0026757779996842146, 0.007132661063224077, 0.010286321863532066, 0.015962811186909676, 0.004528969060629606, 0.01888921484351158, 0.004036444239318371, 0.00027040645363740623, 0.0002387895801803097]], [[0.278582364320755, 0.012074317783117294, 0.4035726487636566, 0.05818924307823181, 0.5308449864387512, 0.7759386301040649, 0.6032847166061401, 0.04120228812098503, 0.6623223423957825, 0.4034832715988159, 0.2541539669036865, 0.023309720680117607, 0.054716046899557114, 0.3570294678211212, 0.004749305546283722, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03977029398083687, 0.025161603465676308, 0.4579423666000366, 0.3708552420139313, 0.767479419708252, 0.5835962295532227, 0.5609359741210938, 0.14304085075855255, 0.8166816234588623, 0.848468542098999, 0.5771627426147461, 0.07112090289592743, 0.12416274100542068, 0.618628740310669, 0.06885465234518051, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.004083612468093634, 0.0006101519684307277, 0.12011494487524033, 0.04229450225830078, 0.17203551530838013, 0.013333754613995552, 0.01874622330069542, 0.021773431450128555, 0.8914079666137695, 0.25239333510398865, 0.2674473226070404, 0.0986163467168808, 0.10968483239412308, 0.05420238524675369, 0.020816486328840256, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00974054355174303, 0.009372939355671406, 0.016473596915602684, 0.12944141030311584, 0.06805374473333359, 0.019993484020233154, 0.038472987711429596, 0.21791628003120422, 0.8550615310668945, 0.2646826505661011, 0.7350810766220093, 0.17277619242668152, 0.36265626549720764, 0.3741258382797241, 0.06228891760110855, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0007183643756434321, 0.0016902177594602108, 0.0015671673463657498, 0.000663107552099973, 0.015286565758287907, 0.000776923552621156, 0.007700319401919842, 0.11482121050357819, 0.7658083438873291, 0.5443719625473022, 0.22170989215373993, 0.027013972401618958, 0.025342080742120743, 0.049981117248535156, 0.0074298488907516, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.011776593513786793, 0.00668947771191597, 0.05204532667994499, 0.026732588186860085, 0.007738037500530481, 0.19347773492336273, 0.08661007881164551, 0.02065080776810646, 0.8265263438224792, 0.77967369556427, 0.8155033588409424, 0.7568296194076538, 0.6889008283615112, 0.7797287106513977, 0.04647013917565346, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03701920434832573, 0.011276619508862495, 0.026248518377542496, 0.01771446317434311, 0.046063318848609924, 0.020064320415258408, 0.23005641996860504, 0.032302577048540115, 0.6365551948547363, 0.6746889352798462, 0.6497765183448792, 0.5260909199714661, 0.6955898404121399, 0.8770567178726196, 0.04424796253442764, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3583561182022095, 0.034818924963474274, 0.1010005921125412, 0.08171684294939041, 0.0902533084154129, 0.0273053590208292, 0.029195906594395638, 0.10516665875911713, 0.5163984894752502, 0.7107389569282532, 0.5390304327011108, 0.6552954316139221, 0.648922324180603, 0.8148984909057617, 0.13771982491016388, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04790134355425835, 0.016352321952581406, 0.004838719964027405, 0.039540428668260574, 0.004614146891981363, 0.10033231228590012, 0.05411757901310921, 0.012187371961772442, 0.25466611981391907, 0.4822390675544739, 0.22996564209461212, 0.2013523131608963, 0.3018202781677246, 0.325538694858551, 0.10763657093048096, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.18817435204982758, 0.007200991734862328, 0.0915139690041542, 0.00800582580268383, 0.007660675328224897, 0.27090781927108765, 0.08786749839782715, 0.014442713931202888, 0.017244037240743637, 0.8212726712226868, 0.22018176317214966, 0.05063365772366524, 0.16457810997962952, 0.059498634189367294, 0.11578860878944397, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1423795521259308, 0.008703344501554966, 0.2208349108695984, 0.02527845837175846, 0.027401143684983253, 0.09980836510658264, 0.024800043553113937, 0.009310302324593067, 0.11915526539087296, 0.048824433237314224, 0.23738479614257812, 0.04641610383987427, 0.11649724096059799, 0.03864651918411255, 0.200869619846344, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.19247660040855408, 0.028833042830228806, 0.1872357279062271, 0.03232081979513168, 0.031028537079691887, 0.3644941747188568, 0.11239293217658997, 0.0803447812795639, 0.13423573970794678, 0.07468846440315247, 0.009079186245799065, 0.19545331597328186, 0.09625646471977234, 0.07526607811450958, 0.1802312582731247, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1263553649187088, 0.009648445062339306, 0.47829046845436096, 0.22347994148731232, 0.2749265432357788, 0.23197446763515472, 0.05249631777405739, 0.01617230661213398, 0.3326357305049896, 0.1497221142053604, 0.04782721772789955, 0.011572148650884628, 0.1354474574327469, 0.0791783407330513, 0.15636207163333893, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.166306734085083, 0.04561271890997887, 0.48400574922561646, 0.31743937730789185, 0.4171416163444519, 0.1806352734565735, 0.04328177124261856, 0.022486848756670952, 0.1779668778181076, 0.03957689553499222, 0.009708160534501076, 0.01422630064189434, 0.013467496261000633, 0.06257133930921555, 0.22838094830513, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.39438390731811523, 0.20185884833335876, 0.19486168026924133, 0.053202297538518906, 0.29429352283477783, 0.31667405366897583, 0.3313867747783661, 0.37864530086517334, 0.4971301257610321, 0.178373321890831, 0.16689708828926086, 0.16029801964759827, 0.22925321757793427, 0.22496484220027924, 0.11296840012073517, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04784957319498062, 0.004609245341271162, 0.006819143425673246, 0.0166594497859478, 0.006965316366404295, 0.000989345251582563, 0.006434451788663864, 0.005414100829511881, 0.027048002928495407, 0.008730669505894184, 0.003844247665256262, 0.0032386775128543377, 0.00916406698524952, 0.02474893629550934, 0.20862001180648804, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07474544644355774, 0.14463284611701965, 0.06348620355129242, 0.11649901419878006, 0.010943777859210968, 0.05790672451257706, 0.023460205644369125, 0.09132371097803116, 0.013804412446916103, 0.11923354864120483, 0.04609918221831322, 0.0031168698333203793, 0.02482042834162712, 0.018085025250911713, 0.06715727597475052, 0.12851747870445251, 0.06451001763343811, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07159372419118881, 0.23599489033222198, 0.6269188523292542, 0.2670744061470032, 0.07840307801961899, 0.7659233808517456, 0.4897821247577667, 0.7919513583183289, 0.47275444865226746, 0.20698092877864838, 0.5493778586387634, 0.516223669052124, 0.5164197683334351, 0.6560667753219604, 0.10535097867250443, 0.16148854792118073, 0.04709945246577263, 0.0016553826862946153, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.030506769195199013, 0.030577607452869415, 0.37364113330841064, 0.17907775938510895, 0.011576596647500992, 0.0018289608415216208, 0.0013806972419843078, 0.0006740305689163506, 0.006688407156616449, 0.02554805763065815, 0.1984224021434784, 0.0020999175030738115, 0.0001219362675328739, 0.0009508132934570312, 0.00851912796497345, 0.12575848400592804, 0.13552792370319366, 0.1085570901632309, 0.11512085795402527, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.6425503492355347, 0.21330313384532928, 0.8213226199150085, 0.6104346513748169, 0.4307103455066681, 0.005470798350870609, 0.1284545361995697, 0.017213305458426476, 0.14068865776062012, 0.2507726550102234, 0.6069697737693787, 0.17266355454921722, 0.10257546603679657, 0.4255537688732147, 0.07138645648956299, 0.14333586394786835, 0.24668441712856293, 0.19262480735778809, 0.13920731842517853, 0.0020065978169441223, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4833258390426636, 0.07765677571296692, 0.6261626482009888, 0.5845412611961365, 0.457427054643631, 0.012895571999251842, 0.037013884633779526, 0.0045295762829482555, 0.030468540266156197, 0.08583686500787735, 0.4300892949104309, 0.6064226627349854, 0.07339996099472046, 0.02218388393521309, 0.11548874527215958, 0.1578390896320343, 0.19358907639980316, 0.02251395769417286, 0.04702039062976837, 0.018520673736929893, 0.0005939522525295615, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.47047996520996094, 0.06838852912187576, 0.42273014783859253, 0.6319702863693237, 0.4177776277065277, 0.0021309976000338793, 0.00800495408475399, 0.0009326375438831747, 0.00536699453368783, 0.07440605759620667, 0.2710660994052887, 0.5013447999954224, 0.021646764129400253, 0.07749785482883453, 0.039263706654310226, 0.14088943600654602, 0.05360155552625656, 0.043673839420080185, 0.0087194312363863, 0.14876413345336914, 0.3311525881290436, 0.029076436534523964, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5323148965835571, 0.13256511092185974, 0.352451890707016, 0.6556484699249268, 0.4897412359714508, 0.22345507144927979, 0.17913641035556793, 0.12689323723316193, 0.025374194607138634, 0.169284388422966, 0.17072416841983795, 0.08815333992242813, 0.10821512341499329, 0.18704712390899658, 0.05398408696055412, 0.11886978894472122, 0.08032860606908798, 0.053777631372213364, 0.06359982490539551, 0.49348562955856323, 0.7690801620483398, 0.032007213681936264, 0.00921344943344593, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14081209897994995, 0.02785991132259369, 0.37397870421409607, 0.3742114305496216, 0.4757237732410431, 0.0011322007048875093, 0.0019287536852061749, 0.00011125820310553536, 0.00032575102522969246, 0.0042410544119775295, 0.007025705184787512, 0.007957610301673412, 0.0022035131696611643, 0.0008391661685891449, 0.0013405061326920986, 0.013988303020596504, 0.031309448182582855, 0.021422432735562325, 0.015959911048412323, 0.13852538168430328, 0.7482463121414185, 0.1306946873664856, 0.0026366086676716805, 0.006285007111728191, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17781563103199005, 0.10205524414777756, 0.04494810104370117, 0.011432765983045101, 0.0031803075689822435, 0.6873405575752258, 0.1935015618801117, 0.2538544535636902, 0.0006125010550022125, 0.0012519293231889606, 0.0009674279135651886, 0.0007319907890632749, 0.006560447160154581, 0.0005926102166995406, 0.045413821935653687, 0.02759428508579731, 0.1341203898191452, 0.1143924742937088, 0.04895513132214546, 0.2507959306240082, 0.47495928406715393, 0.24884849786758423, 0.04048554226756096, 0.06435439735651016, 0.02207104302942753, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24551935493946075, 0.010881111957132816, 0.16116493940353394, 0.28567203879356384, 0.017490731552243233, 0.03198051080107689, 0.25225502252578735, 0.04009091481566429, 0.1379493623971939, 0.030329206958413124, 0.00725751556456089, 0.0005535308737307787, 0.0001769027003319934, 0.0002177381538785994, 0.11288075149059296, 0.08376637101173401, 0.08644555509090424, 0.08414626121520996, 0.08246676623821259, 0.09393073618412018, 0.2536129355430603, 0.09570588916540146, 0.057335685938596725, 0.27625876665115356, 0.23640654981136322, 0.22554923593997955, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2663186192512512, 0.0841110497713089, 0.39283427596092224, 0.3631373345851898, 0.12446267902851105, 0.0023146900348365307, 0.05166012421250343, 0.025394057855010033, 0.09723125398159027, 0.2633029520511627, 0.09458169341087341, 0.0066002910025417805, 0.0024958536960184574, 0.0033851033076643944, 0.0521465502679348, 0.16592197120189667, 0.037314873188734055, 0.020350072532892227, 0.005164262373000383, 0.009123047813773155, 0.005826999898999929, 0.003451529424637556, 0.017567342147231102, 0.055315494537353516, 0.2317170798778534, 0.05933540314435959, 0.06010079011321068, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.032533496618270874, 0.005542360246181488, 0.14801643788814545, 0.028237437829375267, 0.09192534536123276, 0.002004631096497178, 0.0014868990983814, 0.0018816014053300023, 0.026168106123805046, 0.03666744753718376, 0.2621643543243408, 0.27366670966148376, 0.011460919864475727, 0.012693443335592747, 0.006134080700576305, 0.07053745537996292, 0.19491763412952423, 0.06705262511968613, 0.08265279233455658, 0.006405644118785858, 0.0031596925109624863, 0.005410268437117338, 0.030676638707518578, 0.08307406306266785, 0.20774710178375244, 0.4213918149471283, 0.23337899148464203, 0.08583765476942062, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.028670914471149445, 0.004855436272919178, 0.1069486141204834, 0.02764085866510868, 0.11977140605449677, 0.002686614403501153, 0.007388734724372625, 0.00704799173399806, 0.05677136406302452, 0.0688808336853981, 0.16234178841114044, 0.10548661649227142, 0.1935848444700241, 0.06036479026079178, 0.0025575226172804832, 0.13580749928951263, 0.17484943568706512, 0.09017936140298843, 0.11502011120319366, 0.015199831686913967, 0.008567527867853642, 0.04639086127281189, 0.16773870587348938, 0.16907723248004913, 0.43436557054519653, 0.2870768904685974, 0.10786425322294235, 0.08931463956832886, 0.011009148322045803, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04708265885710716, 0.030478408560156822, 0.0932990089058876, 0.24881142377853394, 0.1139858141541481, 0.03301549330353737, 0.12353643029928207, 0.18121947348117828, 0.3742617964744568, 0.11242274194955826, 0.2673158049583435, 0.05749531090259552, 0.00021243211813271046, 0.005648713558912277, 0.14063234627246857, 0.1727631837129593, 0.039101891219615936, 0.0065339612774550915, 0.0278339721262455, 0.004674504045397043, 0.014613990671932697, 0.03457005321979523, 0.04850766807794571, 0.02412491664290428, 0.009369020350277424, 0.022906647995114326, 0.04899173229932785, 0.01023520715534687, 0.0022774694953113794, 7.664388976991177e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0034641579259186983, 0.015587975271046162, 0.04098831117153168, 0.025328122079372406, 0.012870541773736477, 0.002695741830393672, 0.0012444279855117202, 0.005834754556417465, 0.005115050356835127, 0.10742342472076416, 0.29450723528862, 0.004624508786946535, 0.028462348505854607, 0.09151851385831833, 0.02349407598376274, 0.08213489502668381, 0.3905046880245209, 0.07204636186361313, 0.08312273025512695, 0.02625700645148754, 0.02937941811978817, 0.04131421819329262, 0.05289716273546219, 0.16493423283100128, 0.290347158908844, 0.47713640332221985, 0.44352003931999207, 0.11574649810791016, 0.0847686156630516, 0.047198787331581116, 0.1300322264432907, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00187075010035187, 0.017386021092534065, 0.0033179710153490305, 0.00216178921982646, 0.0006196821923367679, 0.0036519868299365044, 0.020315727218985558, 0.0735914558172226, 0.011879049241542816, 0.05418893322348595, 0.04255518689751625, 0.006776698864996433, 0.007105604745447636, 0.005562894977629185, 0.20312508940696716, 0.056048911064863205, 0.04177262261509895, 0.18134142458438873, 0.04556399583816528, 0.1435631662607193, 0.2900937497615814, 0.07549438625574112, 0.08105770498514175, 0.08377190679311752, 0.011481991037726402, 0.017289845272898674, 0.006863615941256285, 0.013694294728338718, 0.13657283782958984, 0.0735873132944107, 0.3659329116344452, 0.0919225886464119, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.018124327063560486, 0.011053304187953472, 0.041496749967336655, 0.08067373931407928, 0.008039752952754498, 0.27361106872558594, 0.12004023045301437, 0.14489491283893585, 0.05115145817399025, 0.09850911796092987, 0.102595254778862, 0.03553636744618416, 0.03690872713923454, 0.062350839376449585, 0.18180564045906067, 0.06230737641453743, 0.038521286100149155, 0.05914388969540596, 0.03398321941494942, 0.13657090067863464, 0.19265799224376678, 0.07424072921276093, 0.08660972863435745, 0.10718739032745361, 0.16533604264259338, 0.0767570361495018, 0.03204379230737686, 0.028188396245241165, 0.21943823993206024, 0.11997849494218826, 0.2698959410190582, 0.12308003753423691, 0.45223531126976013, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12148405611515045, 0.0812632218003273, 0.2165963500738144, 0.1931358426809311, 0.08697410672903061, 0.006551810074597597, 0.06685828417539597, 0.03445844352245331, 0.0957593098282814, 0.40685340762138367, 0.14669549465179443, 0.05295614153146744, 0.013317806646227837, 0.016840115189552307, 0.07654187083244324, 0.18667352199554443, 0.0350969135761261, 0.030425790697336197, 0.0065561928786337376, 0.028277983888983727, 0.010725672356784344, 0.005219776649028063, 0.03378060460090637, 0.04241056367754936, 0.18939200043678284, 0.06338198482990265, 0.08136797696352005, 0.004227515775710344, 0.024540461599826813, 0.057830944657325745, 0.038525767624378204, 0.0177453625947237, 0.06933332234621048, 0.08866386860609055, NaN, NaN, NaN, NaN, NaN, NaN], [0.00987213384360075, 0.006524993572384119, 0.026135168969631195, 0.011839349754154682, 0.033334147185087204, 0.0041054473258554935, 0.0015945311170071363, 0.0032734640408307314, 0.04142798110842705, 0.08157128095626831, 0.26105597615242004, 0.34578391909599304, 0.018666768446564674, 0.02866668626666069, 0.00917118415236473, 0.04736897721886635, 0.0950922816991806, 0.05233628675341606, 0.0639958381652832, 0.009022187441587448, 0.002768130972981453, 0.005348078906536102, 0.016458049416542053, 0.03350484371185303, 0.1584910899400711, 0.3849281072616577, 0.30566492676734924, 0.08282434195280075, 0.02534077689051628, 0.01897522434592247, 0.013481524772942066, 0.08136109262704849, 0.25969398021698, 0.2513872981071472, 0.07361149042844772, NaN, NaN, NaN, NaN, NaN], [0.024172252044081688, 0.01827125810086727, 0.0764245018362999, 0.024589890614151955, 0.045055974274873734, 0.08366040140390396, 0.049236495047807693, 0.16330885887145996, 0.05235174670815468, 0.18916647136211395, 0.2596777379512787, 0.12284716963768005, 0.3776375353336334, 0.3416304290294647, 0.00993264652788639, 0.15279658138751984, 0.09928575158119202, 0.0573631152510643, 0.10790141671895981, 0.026906443759799004, 0.012519991025328636, 0.06774256378412247, 0.1448669582605362, 0.07826853543519974, 0.4991803467273712, 0.34429702162742615, 0.12145370990037918, 0.10719165205955505, 0.008088642731308937, 0.007662023417651653, 0.013441860675811768, 0.13362208008766174, 0.34251537919044495, 0.10342243313789368, 0.07045409828424454, 0.010391364805400372, NaN, NaN, NaN, NaN], [0.03498423844575882, 0.015507807955145836, 0.05400218814611435, 0.2035217136144638, 0.06879755109548569, 0.01839861460030079, 0.1265679895877838, 0.19229170680046082, 0.28682830929756165, 0.19846217334270477, 0.19391797482967377, 0.03128731623291969, 0.00016305393364746124, 0.003939830232411623, 0.1374405473470688, 0.1865139603614807, 0.02971193566918373, 0.005512321833521128, 0.039164237678050995, 0.007472363766282797, 0.012969624251127243, 0.03476016968488693, 0.0836154893040657, 0.050758667290210724, 0.017821883782744408, 0.08676476776599884, 0.13045690953731537, 0.03245873004198074, 0.009119128808379173, 7.800521416356787e-05, 0.0006276130443438888, 0.0024839011020958424, 0.06682475656270981, 0.06347990781068802, 0.009879485704004765, 0.0017003080574795604, 6.444661266868934e-05, NaN, NaN, NaN], [0.013754391111433506, 0.07632532715797424, 0.05588589236140251, 0.060033075511455536, 0.015113652683794498, 0.024528013542294502, 0.0056539555080235004, 0.025407979264855385, 0.0030256062746047974, 0.3076882064342499, 0.2846599221229553, 0.01613902486860752, 0.07589408755302429, 0.25697121024131775, 0.08533195406198502, 0.029208103194832802, 0.15452517569065094, 0.02615012601017952, 0.034968301653862, 0.030517179518938065, 0.023491270840168, 0.02012590691447258, 0.01683984510600567, 0.047155413776636124, 0.1569623053073883, 0.34555378556251526, 0.29876279830932617, 0.06633269041776657, 0.090775266289711, 0.05117363482713699, 0.14964616298675537, 0.024973956868052483, 0.22028914093971252, 0.5953715443611145, 0.10930891335010529, 0.05826140195131302, 0.08348876982927322, 0.2024080604314804, NaN, NaN], [0.0015476603293791413, 0.017548631876707077, 0.0017550711054354906, 0.0017123925499618053, 0.0004861274501308799, 0.0013240363914519548, 0.007671059109270573, 0.03281305357813835, 0.0013763409806415439, 0.060824256390333176, 0.04298469424247742, 0.011416267603635788, 0.012759965844452381, 0.012971585616469383, 0.16966485977172852, 0.023966457694768906, 0.008770916610956192, 0.0534873865544796, 0.015555462799966335, 0.07408829033374786, 0.12750747799873352, 0.026930494233965874, 0.023400133475661278, 0.02665247581899166, 0.00316479685716331, 0.004739005118608475, 0.002742160577327013, 0.006070322822779417, 0.09564805775880814, 0.029174519702792168, 0.5144217014312744, 0.05911846086382866, 0.020064763724803925, 0.0023497287184000015, 0.004584830719977617, 0.10225256532430649, 0.05520752817392349, 0.4466201066970825, 0.09660884737968445, NaN], [0.005211545154452324, 0.0055291797034442425, 0.0040288688614964485, 0.011110500432550907, 0.002710954286158085, 0.0645279660820961, 0.01716793328523636, 0.025083528831601143, 0.010282285511493683, 0.009002536535263062, 0.0011292833369225264, 0.0045064822770655155, 0.007478337734937668, 0.004868943244218826, 0.13875910639762878, 0.18986307084560394, 0.036011889576911926, 0.08335232734680176, 0.12826237082481384, 0.08758756518363953, 0.027860891073942184, 0.10198243707418442, 0.0981309786438942, 0.17985263466835022, 0.11864234507083893, 0.08274368196725845, 0.1066904067993164, 0.051979877054691315, 0.06548189371824265, 0.03337343409657478, 0.0824524462223053, 0.012718076817691326, 0.0349668525159359, 0.03024965338408947, 0.01082769688218832, 0.0127665214240551, 0.014164488762617111, 0.01925024762749672, 0.0028478982858359814, 0.0007362329051829875]], [[0.12737327814102173, 0.10940374433994293, 0.05123003572225571, 0.7807462215423584, 0.0676276683807373, 0.02884089946746826, 0.05574861168861389, 0.5975708961486816, 0.07044392824172974, 0.5009010434150696, 0.31273892521858215, 0.07660850137472153, 0.29424503445625305, 0.028401609510183334, 0.07683643698692322, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03750006482005119, 0.429240882396698, 0.15060469508171082, 0.2604650557041168, 0.037177786231040955, 0.1944778561592102, 0.07849539071321487, 0.6716934442520142, 0.06105323135852814, 0.07711976766586304, 0.20997941493988037, 0.028168758377432823, 0.12550987303256989, 0.030995607376098633, 0.0958443135023117, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.15516091883182526, 0.07278051972389221, 0.11765316128730774, 0.7884857058525085, 0.11075033247470856, 0.051856692880392075, 0.18673725426197052, 0.2268398553133011, 0.013722711242735386, 0.6478350162506104, 0.5306386947631836, 0.3090885877609253, 0.22243055701255798, 0.16200464963912964, 0.13070979714393616, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.21811531484127045, 0.7140333652496338, 0.018219277262687683, 0.764274001121521, 0.15804116427898407, 0.03280843421816826, 0.11008237302303314, 0.09874711185693741, 0.0423860140144825, 0.5652360320091248, 0.14938808977603912, 0.2869919240474701, 0.39966318011283875, 0.1259765923023224, 0.0577625073492527, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.11744663864374161, 0.1893559694290161, 0.05823011323809624, 0.03701714053750038, 0.15626470744609833, 0.08588159829378128, 0.26269999146461487, 0.41053518652915955, 0.007210245821624994, 0.3749772906303406, 0.4537068009376526, 0.6417111158370972, 0.1666039228439331, 0.13084180653095245, 0.14052902162075043, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3613002598285675, 0.240200012922287, 0.044567547738552094, 0.04614294692873955, 0.0021214759908616543, 0.17616558074951172, 0.11286458373069763, 0.11203286051750183, 0.009014172479510307, 0.10163455456495285, 0.0949772298336029, 0.06209810823202133, 0.11910365521907806, 0.04125094786286354, 0.1871420443058014, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2914785146713257, 0.381010502576828, 0.08399549126625061, 0.4511452913284302, 0.048780620098114014, 0.008560722693800926, 0.1541443020105362, 0.12101723253726959, 0.02183164842426777, 0.18665823340415955, 0.13169258832931519, 0.13539372384548187, 0.14286382496356964, 0.031125182285904884, 0.2064482420682907, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3084108829498291, 0.4568510055541992, 0.068343386054039, 0.40243175625801086, 0.04035715013742447, 0.028490515425801277, 0.006473515648394823, 0.6036491990089417, 0.14769236743450165, 0.09462843090295792, 0.04651549458503723, 0.08334364742040634, 0.08459941297769547, 0.022403797134757042, 0.13448290526866913, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4981050491333008, 0.13424238562583923, 0.16773013770580292, 0.5160816311836243, 0.029790958389639854, 0.22989192605018616, 0.568993866443634, 0.056374672800302505, 0.08792523294687271, 0.2900378406047821, 0.12431738525629044, 0.017185388132929802, 0.05061684548854828, 0.020683959126472473, 0.13275840878486633, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.33482691645622253, 0.4720645546913147, 0.20652346312999725, 0.6004944443702698, 0.1402488797903061, 0.13250590860843658, 0.13873517513275146, 0.5260767936706543, 0.01182119082659483, 0.1017654612660408, 0.047682080417871475, 0.04534589499235153, 0.10121697187423706, 0.0026118881069123745, 0.13006491959095, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.27261805534362793, 0.5674196481704712, 0.08154824376106262, 0.8736060261726379, 0.4724165201187134, 0.1720387041568756, 0.13692085444927216, 0.40960294008255005, 0.06138879805803299, 0.0898643285036087, 0.15986473858356476, 0.04882661625742912, 0.09858791530132294, 0.005254920106381178, 0.09166211634874344, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.33052578568458557, 0.40956470370292664, 0.44244009256362915, 0.8809638619422913, 0.26719745993614197, 0.38818857073783875, 0.40750059485435486, 0.4857279658317566, 0.04656125605106354, 0.08998580276966095, 0.02227160707116127, 0.42457664012908936, 0.06242617964744568, 0.019552020356059074, 0.08343644440174103, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.20678018033504486, 0.17620769143104553, 0.3081345558166504, 0.6112105250358582, 0.534289538860321, 0.19626931846141815, 0.17160479724407196, 0.4079393148422241, 0.027630727738142014, 0.07990976423025131, 0.0661839172244072, 0.022294294089078903, 0.11108729988336563, 0.024492109194397926, 0.12739884853363037, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2302674651145935, 0.4147239625453949, 0.3118293881416321, 0.3454154133796692, 0.20178626477718353, 0.3381562829017639, 0.1571493148803711, 0.4487079083919525, 0.02096635475754738, 0.11857040971517563, 0.09038619697093964, 0.01401298213750124, 0.06377796083688736, 0.029106009751558304, 0.10548537224531174, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0850413590669632, 0.2905830442905426, 0.047175440937280655, 0.009145522490143776, 0.014412813819944859, 0.03387918695807457, 0.04852135106921196, 0.2856408655643463, 0.03688584640622139, 0.02503933012485504, 0.030300520360469818, 0.020876996219158173, 0.004409631714224815, 0.0025441893376410007, 0.1292814165353775, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01263146661221981, 0.08983241021633148, 0.002674827352166176, 0.0008326905663125217, 0.0032944290433079004, 0.06790440529584885, 0.02327594719827175, 0.08626140654087067, 0.0010102109517902136, 0.0009567838278599083, 0.001915089669637382, 0.019144434481859207, 0.060631223022937775, 0.04236740246415138, 0.2042645514011383, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12322216480970383, 0.14532910287380219, 0.08289580047130585, 0.07800436019897461, 0.016899574548006058, 0.20651613175868988, 0.15389330685138702, 0.08048079907894135, 0.023754820227622986, 0.08939354121685028, 0.05408218502998352, 0.0083498889580369, 0.16772767901420593, 0.03971855714917183, 0.029394451528787613, 0.12774905562400818, 0.07772441953420639, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.002537816995754838, 0.0036866364534944296, 0.0026212686207145452, 0.0010326605988666415, 0.0028582154773175716, 0.0016078348271548748, 0.0024177017621695995, 0.004757970105856657, 0.007405414246022701, 0.0004943490494042635, 0.0008183143800124526, 0.0020540759433060884, 0.0008841927628964186, 0.0009274804615415633, 0.13894422352313995, 0.058547187596559525, 0.7868303656578064, 0.02677525207400322, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18076959252357483, 0.11159703880548477, 0.07333940267562866, 0.12368053197860718, 0.1442640721797943, 0.3224244713783264, 0.2286587655544281, 0.10576390475034714, 0.0873323604464531, 0.0707816481590271, 0.07077325880527496, 0.024980774149298668, 0.015894055366516113, 0.01236753724515438, 0.034113459289073944, 0.12958122789859772, 0.05996095389127731, 0.20109553635120392, 0.07473170012235641, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.008514223620295525, 0.006442691199481487, 0.003549255197867751, 0.00919315591454506, 0.0011393448803573847, 0.0005870977183803916, 0.02400296926498413, 0.03577389195561409, 0.006469632964581251, 0.004828252829611301, 0.0027150637470185757, 9.597353346180171e-05, 0.00011822552187368274, 0.000396552961319685, 0.1521017998456955, 0.11586850136518478, 0.18037959933280945, 0.354478657245636, 0.6275972127914429, 0.01217791810631752, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0016907083336263895, 9.336868970422074e-05, 0.0023900996893644333, 0.0018071996746584773, 0.001690928009338677, 0.0010278637055307627, 0.008010926656425, 0.0018918663263320923, 0.0009378245449624956, 0.0005185406771488488, 0.00012474792310968041, 0.00014544214354828, 2.7525844416231848e-05, 2.095987474604044e-05, 0.12926018238067627, 0.04329086095094681, 0.2822243273258209, 0.5110569596290588, 0.8230794668197632, 0.28263914585113525, 0.006951561663299799, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08279342949390411, 0.00717265997081995, 0.01113244891166687, 0.030300047248601913, 0.03227340802550316, 0.02679654024541378, 0.2711687386035919, 0.12656770646572113, 0.0010184150887653232, 0.0069296094588935375, 0.006689318455755711, 0.00307065830565989, 0.004024384077638388, 0.006041096989065409, 0.12722525000572205, 0.15041278302669525, 0.01652364432811737, 0.09004879742860794, 0.1228649914264679, 0.03705046698451042, 0.03279988467693329, 0.012472960166633129, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09468965977430344, 0.010531323030591011, 0.1253902167081833, 0.09483902901411057, 0.060478318482637405, 0.1959676593542099, 0.5850688219070435, 0.11734473705291748, 0.08924026787281036, 0.031869061291217804, 0.04437774419784546, 0.004531644284725189, 0.19630968570709229, 0.04580901935696602, 0.04253998026251793, 0.005692727863788605, 0.004583822097629309, 0.011303454637527466, 0.06351188570261002, 0.07110948860645294, 0.03377191722393036, 0.8937738537788391, 0.1077374666929245, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03443194553256035, 0.006786322686821222, 0.08545193076133728, 0.2555176913738251, 0.16119416058063507, 0.3760574460029602, 0.3180745542049408, 0.0858285129070282, 0.0052651395089924335, 0.035345133394002914, 0.0046972003765404224, 0.00805696938186884, 0.0738091915845871, 0.004572577308863401, 0.028640231117606163, 0.1957636922597885, 0.00532554043456912, 0.2672942280769348, 0.07843183726072311, 0.01169322058558464, 0.006695515010505915, 0.022856300696730614, 0.03495524823665619, 0.2056257426738739, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.26599034667015076, 0.06405031681060791, 0.39913085103034973, 0.7390084862709045, 0.8533709049224854, 0.0830850899219513, 0.22198519110679626, 0.15359464287757874, 0.0286090150475502, 0.1338224709033966, 0.06985709816217422, 0.03841168060898781, 0.1308237761259079, 0.01580808497965336, 0.010780439712107182, 0.21948350965976715, 0.003219911362975836, 0.13064762949943542, 0.017335020005702972, 0.004487968049943447, 0.006097455509006977, 0.0023269150406122208, 0.014221499674022198, 0.1740167737007141, 0.05570632219314575, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16064751148223877, 0.5348425507545471, 0.09399141371250153, 0.3709404170513153, 0.3757614493370056, 0.2272261530160904, 0.2699662148952484, 0.46868544816970825, 0.09081633388996124, 0.07856583595275879, 0.054298948496580124, 0.10659310221672058, 0.05178465321660042, 0.012835889123380184, 0.19243957102298737, 0.027252521365880966, 0.05625513195991516, 0.024279700592160225, 0.009296371601521969, 0.04113621264696121, 0.04445572942495346, 0.05016031116247177, 0.300394743680954, 0.219209223985672, 0.5284181833267212, 0.13528388738632202, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.33067551255226135, 0.40668511390686035, 0.03748138248920441, 0.16017457842826843, 0.02931954525411129, 0.1285390406847, 0.43687552213668823, 0.6227295398712158, 0.016583241522312164, 0.054699335247278214, 0.43602558970451355, 0.028376825153827667, 0.1860552728176117, 0.202489972114563, 0.03443598374724388, 0.16918426752090454, 0.005196947604417801, 0.010393726639449596, 0.0008839815272949636, 0.18853645026683807, 0.23955073952674866, 0.03703731670975685, 0.018581384792923927, 0.07692746073007584, 0.05213537812232971, 0.05520249530673027, 0.03837481513619423, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.025147954002022743, 0.023277895525097847, 0.036982107907533646, 0.030706623569130898, 0.00253032217733562, 0.08060919493436813, 0.062497250735759735, 0.22720953822135925, 0.015824737027287483, 0.020865583792328835, 0.051981136202812195, 0.016274577006697655, 0.3496847152709961, 0.19709302484989166, 0.00854758732020855, 0.21910618245601654, 0.012340836226940155, 0.011061819270253181, 0.004421355202794075, 0.01345156505703926, 0.015948239713907242, 0.001919197733514011, 0.0006712953327223659, 0.0014401280786842108, 0.0009498890140093863, 0.0011606297921389341, 0.0013843519845977426, 0.005138876382261515, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0009813109645619988, 0.0007951235747896135, 0.007896890863776207, 0.006039812229573727, 0.001424357295036316, 0.003153599100187421, 0.0010362794855609536, 0.006138501223176718, 0.00410880520939827, 0.003359388094395399, 0.008728301152586937, 0.0021525975316762924, 0.2318088710308075, 0.017491629347205162, 0.0005464124260470271, 0.12592341005802155, 0.022789308801293373, 0.01544136367738247, 0.05098855495452881, 0.006733328104019165, 0.0011512627825140953, 0.0067494111135602, 0.03519098460674286, 0.08756479620933533, 0.04847756400704384, 0.13774195313453674, 0.07365753501653671, 0.19525301456451416, 0.019442297518253326, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.008814784698188305, 0.009578033350408077, 0.008741176687180996, 0.002597709419205785, 0.0019302073633298278, 0.02750723622739315, 0.010486552491784096, 0.061721935868263245, 0.05738110467791557, 0.0038812088314443827, 0.08735688030719757, 0.00500333309173584, 3.085857315454632e-05, 0.005531619768589735, 0.14116442203521729, 0.04374772310256958, 0.10635814815759659, 0.1203576922416687, 0.4972172677516937, 0.09716533124446869, 0.05867829546332359, 0.13453392684459686, 0.39353471994400024, 0.6331138610839844, 0.33491814136505127, 0.5983138680458069, 0.3633559048175812, 0.6357010006904602, 0.7792285084724426, 0.005659972317516804, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.015857994556427002, 0.010374038480222225, 0.002225207630544901, 0.002974742790684104, 0.0010843537747859955, 0.007387869525700808, 0.006818806286901236, 0.0318806953728199, 0.1651621013879776, 0.21757511794567108, 0.2911650240421295, 0.08204617351293564, 0.016449127346277237, 0.10985822230577469, 0.0020742996130138636, 0.05199728533625603, 0.014302223920822144, 0.13574257493019104, 0.05407930538058281, 0.010633953846991062, 0.007459194865077734, 0.0004102779785171151, 0.01107444055378437, 0.16451390087604523, 0.19313758611679077, 0.018386593088507652, 0.03492085263133049, 0.1390746384859085, 0.6526300311088562, 0.08304706960916519, 0.27643677592277527, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01972219906747341, 0.20374125242233276, 0.0031293979845941067, 0.004390338435769081, 0.031924858689308167, 0.06048818305134773, 0.0774247944355011, 0.7845978140830994, 0.15838612616062164, 0.06142642721533775, 0.0820784792304039, 0.20785683393478394, 0.46646884083747864, 0.42270010709762573, 0.053927596658468246, 0.0008206118363887072, 0.0011099595576524734, 0.0005428412696346641, 0.0013029578840360045, 0.0009422241128049791, 0.001036918954923749, 0.00015340711979661137, 0.003300317795947194, 0.0019372785463929176, 0.003245894331485033, 0.0010756017873063684, 0.0009867959888651967, 0.04242069274187088, 0.25679609179496765, 0.03714281693100929, 0.46563825011253357, 0.052469443529844284, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.026567673310637474, 0.2768426239490509, 0.016553064808249474, 0.07253812253475189, 0.029352964833378792, 0.034967049956321716, 0.09283487498760223, 0.5970632433891296, 0.02342795394361019, 0.04057195410132408, 0.06215028092265129, 0.2966896891593933, 0.4489157795906067, 0.24187524616718292, 0.048112284392118454, 0.0011551693314686418, 0.0015016108518466353, 0.00018865184392780066, 0.0004620797117240727, 0.001353209256194532, 0.001276124152354896, 0.001269699539989233, 0.02504812367260456, 0.016660472378134727, 0.007664685603231192, 0.000621759332716465, 0.0039494638331234455, 0.05373308062553406, 0.5797222256660461, 0.04267296567559242, 0.3308492600917816, 0.22605444490909576, 0.03655111417174339, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14453455805778503, 0.4129781723022461, 0.021322425454854965, 0.11776001751422882, 0.008680691011250019, 0.12525556981563568, 0.1459336131811142, 0.4943058490753174, 0.041365865617990494, 0.06633096933364868, 0.48416346311569214, 0.027247071266174316, 0.10342812538146973, 0.15874288976192474, 0.04535134881734848, 0.18345873057842255, 0.006115049123764038, 0.007153322920203209, 0.00125643250066787, 0.15791349112987518, 0.17755654454231262, 0.06167090684175491, 0.028255566954612732, 0.04990806803107262, 0.014394938945770264, 0.013118196278810501, 0.02539716847240925, 0.00894339382648468, 0.04024626687169075, 0.05642623454332352, 0.04561464861035347, 0.029457826167345047, 0.09210912138223648, 0.1002524197101593, NaN, NaN, NaN, NaN, NaN, NaN], [0.03164434805512428, 0.10487183183431625, 0.019769076257944107, 0.0709872916340828, 0.0046073514968156815, 0.12636253237724304, 0.06114564463496208, 0.5786424875259399, 0.17960773408412933, 0.15923625230789185, 0.14680741727352142, 0.04373620077967644, 0.20528176426887512, 0.14476445317268372, 0.03252548724412918, 0.2828649580478668, 0.011994204483926296, 0.006339475512504578, 0.0030444697476923466, 0.006948052905499935, 0.008767204359173775, 0.0014567734906449914, 0.00018795454525388777, 0.00020330831466708332, 0.0001539710647193715, 0.0004007722018286586, 0.0012242270167917013, 0.001961026806384325, 0.0007920600473880768, 0.002005743095651269, 0.00011892847396666184, 0.00023868663993198425, 0.0018499011639505625, 0.002196513582020998, 0.004604275804013014, NaN, NaN, NaN, NaN, NaN], [0.03216148540377617, 0.04786192253232002, 0.0904572606086731, 0.284318745136261, 0.04915444552898407, 0.20336958765983582, 0.019341057166457176, 0.31598398089408875, 0.503376841545105, 0.2976534068584442, 0.3550446927547455, 0.318871408700943, 0.31741514801979065, 0.09137054532766342, 0.022498751059174538, 0.128562331199646, 0.014782274141907692, 0.007007280830293894, 0.02549830637872219, 0.0029198189731687307, 0.0006880113505758345, 0.0037798655685037374, 0.009390356950461864, 0.008127862587571144, 0.00817851535975933, 0.024966517463326454, 0.0308842696249485, 0.07813727855682373, 0.003280356992036104, 0.001509596244432032, 0.010023933835327625, 0.08412036299705505, 0.1339937299489975, 0.13076454401016235, 0.2572615444660187, 0.02603374607861042, NaN, NaN, NaN, NaN], [0.00784912146627903, 0.004314524121582508, 0.007757026236504316, 0.004281783476471901, 0.001910648075863719, 0.00898022297769785, 0.007197065278887749, 0.05121663585305214, 0.12398385256528854, 0.006457128562033176, 0.09335841238498688, 0.0023844544775784016, 1.3785818737233058e-05, 0.0021891386713832617, 0.13778245449066162, 0.018602287396788597, 0.034721970558166504, 0.034974802285432816, 0.21532808244228363, 0.037075310945510864, 0.013384592719376087, 0.039282385259866714, 0.11046459525823593, 0.17542847990989685, 0.05914776027202606, 0.1884417086839676, 0.12911023199558258, 0.24417443573474884, 0.327198326587677, 0.0006843891460448503, 0.1527024656534195, 0.4776603579521179, 0.37270504236221313, 0.4335513412952423, 0.6841917634010315, 0.8031085133552551, 0.004920803010463715, NaN, NaN, NaN], [0.0865921899676323, 0.029389984905719757, 0.007211814168840647, 0.022628001868724823, 0.003064699238166213, 0.026838112622499466, 0.02777392417192459, 0.17195671796798706, 0.5349084734916687, 0.37311822175979614, 0.5073185563087463, 0.12468769401311874, 0.014684900641441345, 0.11363118886947632, 0.01852630451321602, 0.05855157971382141, 0.021276630461215973, 0.13662834465503693, 0.05244326964020729, 0.015041220933198929, 0.007642571348696947, 0.00036013865610584617, 0.004098850768059492, 0.033856965601444244, 0.05778159946203232, 0.005442364141345024, 0.017580043524503708, 0.04633626714348793, 0.3112163841724396, 0.03644357994198799, 0.0868009626865387, 0.020123973488807678, 0.03773906081914902, 0.06257405877113342, 0.2619801461696625, 0.7497928738594055, 0.19582624733448029, 0.4370352327823639, NaN, NaN], [0.021940317004919052, 0.17988227307796478, 0.0027716639451682568, 0.0058884406462311745, 0.02112143486738205, 0.056551095098257065, 0.09669405966997147, 0.8433947563171387, 0.1836535632610321, 0.048101164400577545, 0.0939687192440033, 0.12228170782327652, 0.5153423547744751, 0.4533718526363373, 0.10564926266670227, 0.0006882869056425989, 0.0005033394554629922, 0.00030677669565193355, 0.001028614118695259, 0.00036578672006726265, 0.0005035633221268654, 5.2447539928834885e-05, 0.0006442382582463324, 0.0003597578906919807, 0.0002600657753646374, 8.536354289390147e-05, 0.00018848010222427547, 0.00940172839909792, 0.03475101292133331, 0.004768407437950373, 0.09523987770080566, 0.0036924693267792463, 0.0034024319611489773, 0.001987446565181017, 0.06484154611825943, 0.36614781618118286, 0.06470755487680435, 0.48020803928375244, 0.12385622411966324, NaN], [0.07970402389764786, 0.263812392950058, 0.027112353593111038, 0.06228066235780716, 0.03007029928267002, 0.5465735197067261, 0.2176109254360199, 0.5667538046836853, 0.10334119945764542, 0.3484029769897461, 0.1586397886276245, 0.28290486335754395, 0.07807470858097076, 0.405972421169281, 0.12247955799102783, 0.13044977188110352, 0.023216107860207558, 0.019304566085338593, 0.018173998221755028, 0.12614674866199493, 0.04656239226460457, 0.015089727938175201, 0.04114385321736336, 0.018700774759054184, 0.020505733788013458, 0.009310846216976643, 0.02222343534231186, 0.22412429749965668, 0.3900958001613617, 0.1100122332572937, 0.14125461876392365, 0.09716113656759262, 0.14588865637779236, 0.12185929715633392, 0.5472521185874939, 0.7197717428207397, 0.31834876537323, 0.37092098593711853, 0.2838878929615021, 0.0011011400492861867]]], [[[0.00039591442327946424, 4.3682277464540675e-05, 1.7448855942348018e-05, 4.859234650211874e-06, 1.1413659422032651e-06, 1.0625568393152207e-05, 1.9137923246148603e-08, 5.615326585939329e-07, 5.487099315359956e-06, 2.1910665282121045e-07, 2.532970881929941e-07, 7.501878940274764e-07, 1.657212578720646e-06, 1.0862070212169783e-06, 0.18717002868652344, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6005652546882629, 0.09179380536079407, 0.017407523468136787, 0.009556752629578114, 0.001977206440642476, 0.02417689561843872, 0.001285116421058774, 0.0015866898465901613, 0.0007265046588145196, 0.0008927723974920809, 0.008914382196962833, 0.0016361800953745842, 0.1313493698835373, 0.006872364319860935, 0.052507203072309494, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00456381356343627, 0.8302816152572632, 0.11558636277914047, 0.010320104658603668, 0.00024428890901617706, 9.749805758474395e-05, 7.678471774852369e-06, 0.0030259541235864162, 3.9539358112961054e-05, 7.781033491482958e-05, 0.0003711417084559798, 9.1652873379644e-06, 0.0006458949064835906, 0.00023330377007368952, 0.00865631178021431, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0011992683866992593, 0.008629350923001766, 0.6251504421234131, 0.015135818161070347, 0.001978840446099639, 0.000745285302400589, 5.708653407054953e-05, 0.00043479635496623814, 0.0005481417756527662, 0.0016355890547856688, 0.0002436988870613277, 5.164237336430233e-06, 4.976044510840438e-05, 3.400173591217026e-05, 0.00024351823958568275, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.006698334589600563, 0.006304558366537094, 0.34660738706588745, 0.7217360138893127, 0.06864907592535019, 0.0027605369687080383, 0.0006927561480551958, 0.00010832686530193314, 0.0002978279662784189, 0.007849807851016521, 0.0023863124661147594, 8.873132173903286e-06, 2.0952818886144087e-05, 4.62439584225649e-06, 0.000559441396035254, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0006861803703941405, 0.036174044013023376, 0.4128260612487793, 0.09897080808877945, 0.6376775503158569, 0.19431157410144806, 0.0007082957308739424, 0.05852581560611725, 0.0003548018867149949, 0.00026609119959175587, 0.0006576658925041556, 0.0007862210040912032, 0.027955245226621628, 0.006076914723962545, 0.0010327105410397053, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.7293352305713938e-09, 1.4693102912133327e-06, 3.0192679332685657e-05, 1.0152590220968705e-05, 0.005660888738930225, 0.5108420252799988, 0.0005426039570011199, 0.0008102089632302523, 3.168102921335958e-06, 6.12798771726375e-08, 2.5310575324510864e-07, 5.088519174023531e-06, 0.00021843344438821077, 2.5946601454052143e-06, 2.594279294498847e-06, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [7.755387923680246e-05, 3.5259185096947476e-05, 0.0012139425380155444, 0.00035162578569725156, 0.00505053298547864, 0.4696201980113983, 0.5859625339508057, 0.009771172888576984, 0.0005853781476616859, 3.0261137453635456e-06, 1.2206013707327656e-05, 2.2465645088232122e-05, 0.013555033132433891, 0.0011026648571714759, 7.656160596525297e-05, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [3.390625025190275e-08, 5.7732322602532804e-05, 3.19563605444273e-06, 2.0829493507790175e-07, 5.039521965954918e-06, 0.00017657184798736125, 0.000729007413610816, 0.8331114649772644, 0.0037640428636223078, 1.5948112377373036e-06, 5.8014775277115405e-06, 4.528372699041938e-07, 0.00020723954366985708, 0.00025866259238682687, 1.95706252270611e-06, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [2.7739795882553153e-07, 2.501485141692683e-05, 4.778147285833256e-06, 3.7190903867667657e-07, 9.610201523457818e-09, 1.1292572708043735e-06, 1.2355405942798825e-07, 3.984562499681488e-05, 0.6202287077903748, 0.0002610959345474839, 0.00017016819037962705, 9.242457963409834e-07, 2.799387630147976e-06, 3.2760857493485673e-07, 1.038134087139042e-06, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.2775580216839444e-05, 0.0010497755138203502, 6.564326031366363e-05, 4.172011358605232e-06, 4.676745959386608e-07, 3.6489967669695034e-07, 8.09820832614605e-08, 5.78842673348845e-06, 0.0015375507064163685, 0.7445451617240906, 0.026254041120409966, 8.213486580643803e-05, 1.1159563655382954e-05, 3.0355058697750792e-05, 2.6809220798895694e-06, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.3068409316474572e-05, 0.00010775982809718698, 0.00024633039720356464, 3.3576598070794716e-05, 4.556980275083333e-05, 1.0597023702985098e-07, 9.86238859468358e-08, 2.1072135041322326e-06, 0.0013669389300048351, 0.5916010141372681, 0.4436832368373871, 0.0013138806680217385, 4.73510908705066e-06, 6.116700660641072e-06, 2.961193558803643e-06, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [4.950460061081685e-05, 0.0011237917933613062, 0.017257435247302055, 0.0011414129985496402, 0.025087760761380196, 0.00036485170130617917, 3.213326635886915e-05, 5.293267349770758e-06, 4.4593522034119815e-05, 0.001686945091933012, 0.00823597889393568, 0.8047888278961182, 0.014818375930190086, 0.006413417402654886, 2.281446177221369e-05, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.000998240546323359, 0.1768636256456375, 0.0663335844874382, 0.02716292440891266, 0.03197554498910904, 0.001621886040084064, 0.00012482069723773748, 7.020989141892642e-05, 0.08078382909297943, 0.1701173484325409, 0.08303841948509216, 0.5506232380867004, 0.06293172389268875, 0.03332124650478363, 0.0033543158788233995, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.021357281133532524, 0.0013016555458307266, 0.00422634556889534, 0.00104909623041749, 0.012563652358949184, 0.07401228696107864, 0.007866809144616127, 0.0024991247337311506, 0.0011657974682748318, 5.4276370065053925e-06, 0.0024851916823536158, 0.0298884529620409, 0.4522511959075928, 0.2182934284210205, 0.14462554454803467, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.02659090794622898, 0.049626123160123825, 0.04500019550323486, 0.012677792459726334, 0.33557751774787903, 0.02776678465306759, 0.02675992250442505, 0.09967876970767975, 0.04216820374131203, 0.009756066836416721, 0.0133897690102458, 0.12886802852153778, 0.03152704983949661, 0.046163998544216156, 0.21004843711853027, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05978302285075188, 0.18161648511886597, 0.038620203733444214, 0.022025080397725105, 0.09790226072072983, 0.04398013651371002, 0.00788698997348547, 0.04135579988360405, 0.0068543110974133015, 0.03809167072176933, 0.03150040656328201, 0.0462106354534626, 0.024762138724327087, 0.011792140081524849, 0.015839271247386932, 0.16810710728168488, 0.017288343980908394, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.005166883580386639, 0.0005590450600720942, 0.007114546839147806, 0.0015656572068110108, 0.02179996483027935, 0.0010864944197237492, 0.0051814797334373, 0.0011148365447297692, 0.00816393457353115, 0.0019027285743504763, 0.005033016670495272, 0.010743028484284878, 0.0006906923954375088, 0.0011143455049023032, 0.16189540922641754, 0.12647151947021484, 0.25301796197891235, 0.03169602155685425, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17136499285697937, 0.002046054694801569, 0.4725193679332733, 0.24347566068172455, 0.1026763990521431, 0.00369152519851923, 0.013768541626632214, 0.003912978805601597, 0.022358577698469162, 0.06323882192373276, 0.28539538383483887, 0.009778834879398346, 0.0043070269748568535, 0.020384330302476883, 0.006856778170913458, 0.15976493060588837, 0.03159531578421593, 0.05609510838985443, 0.007400199305266142, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18433871865272522, 0.013500750064849854, 0.42166435718536377, 0.1935500204563141, 0.3502363860607147, 0.0009389789775013924, 0.0472395233809948, 0.015336934477090836, 0.07204270362854004, 0.07276465743780136, 0.4023721218109131, 0.016390468925237656, 0.00493515282869339, 0.01088448241353035, 0.18081046640872955, 0.16021955013275146, 0.26433131098747253, 0.07329617440700531, 0.11257290840148926, 0.001577433431521058, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01929071731865406, 3.154709338559769e-05, 0.04895680397748947, 0.04499320685863495, 0.03726757690310478, 0.0012487026397138834, 0.06078735366463661, 0.0025376947596669197, 0.023622047156095505, 0.008605116978287697, 0.05601886287331581, 0.011475598439574242, 0.0013240767875686288, 0.009706309996545315, 0.13962702453136444, 0.22870834171772003, 0.043985288590192795, 0.04075293987989426, 0.0035545979626476765, 0.0075324228964746, 0.00014864112017676234, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.032548993825912476, 0.0047013829462230206, 0.08043498545885086, 0.08197268843650818, 0.43236956000328064, 0.013080407865345478, 0.006017346400767565, 0.05529334023594856, 0.01970849372446537, 0.004050384275615215, 0.0073967562057077885, 0.005829385481774807, 0.0008975209202617407, 0.0025361862499266863, 0.011671289801597595, 0.047688793390989304, 0.14664201438426971, 0.03658692538738251, 0.6408759355545044, 0.43873438239097595, 0.20478755235671997, 0.00511742290109396, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.046304989606142044, 0.026358718052506447, 0.20277923345565796, 0.3021180331707001, 0.6281617879867554, 0.19840610027313232, 0.12000668793916702, 0.21165543794631958, 0.0507807619869709, 0.10083203762769699, 0.17539183795452118, 0.08392243832349777, 0.036049142479896545, 0.06088141351938248, 0.024198466911911964, 0.07761336117982864, 0.07061085104942322, 0.041570939123630524, 0.1916733682155609, 0.159084752202034, 0.3477410674095154, 0.5968326330184937, 0.004175147507339716, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.016816509887576103, 0.003118144813925028, 0.035858120769262314, 0.02315649762749672, 0.2957051992416382, 0.0033856350928545, 0.008419573307037354, 0.013085800223052502, 0.0065522813238203526, 0.004261805210262537, 0.0022621729876846075, 0.0015856586396694183, 0.00012999074533581734, 0.00036330719012767076, 0.004947974346578121, 0.07191380113363266, 0.05497179180383682, 0.3517811894416809, 0.9035707116127014, 0.14233137667179108, 0.1767667979001999, 0.04289708659052849, 0.00892895832657814, 0.001834895578213036, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13966688513755798, 0.051315873861312866, 0.16794879734516144, 0.17204447090625763, 0.02530861273407936, 0.1971883773803711, 0.6035643219947815, 0.35590535402297974, 0.01904589682817459, 0.14328262209892273, 0.05827813595533371, 0.12283631414175034, 0.08582676202058792, 0.021607764065265656, 0.09174748510122299, 0.21536989510059357, 0.19956108927726746, 0.3517906069755554, 0.458966463804245, 0.09842110425233841, 0.08277469873428345, 0.03296331316232681, 0.04812879115343094, 0.009344152174890041, 0.006280441302806139, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07622234523296356, 0.021088531240820885, 0.13214311003684998, 0.1876712292432785, 0.09946685284376144, 0.0739995539188385, 0.16667790710926056, 0.06527374684810638, 0.2691768705844879, 0.1298666000366211, 0.20347969233989716, 0.28972044587135315, 0.16063560545444489, 0.23408198356628418, 0.02879655919969082, 0.24051256477832794, 0.10134825110435486, 0.04672827199101448, 0.021085558459162712, 0.02245912328362465, 0.026835136115550995, 0.005604758393019438, 0.028772464022040367, 0.01708872988820076, 0.008745603263378143, 0.02540087327361107, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04186922311782837, 0.028065834194421768, 0.2365874946117401, 0.22718128561973572, 0.717268168926239, 0.0283160749822855, 0.047574929893016815, 0.22635598480701447, 0.046485841274261475, 0.11764083057641983, 0.11684223264455795, 0.600357711315155, 0.07936308532953262, 0.1614740490913391, 0.02326863817870617, 0.18141932785511017, 0.024432087317109108, 0.0408032201230526, 0.004596539307385683, 0.0778040885925293, 0.025828123092651367, 0.04467899724841118, 0.0885351300239563, 0.026468785479664803, 0.030213410034775734, 0.16925157606601715, 0.003915028180927038, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.002160860225558281, 0.00041385856457054615, 0.0032894921023398638, 0.004175879992544651, 0.09230346977710724, 0.00037096597952768207, 0.00036027038004249334, 0.000777967507019639, 0.0010948613053187728, 0.006351495627313852, 0.00803811103105545, 0.2546491026878357, 0.005140772555023432, 0.0052158161997795105, 0.0018242541700601578, 0.0821177139878273, 0.0264634620398283, 0.01841210387647152, 0.010007970035076141, 0.006691556889563799, 0.0167625043541193, 0.0005595253896899521, 0.020632673054933548, 0.0021230748388916254, 0.10790054500102997, 0.5654488801956177, 0.3003200888633728, 0.01571945659816265, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01453752163797617, 0.0016249779146164656, 0.07837095856666565, 0.046283330768346786, 0.5220571756362915, 0.00571427633985877, 0.011274048127233982, 0.0005770810530520976, 0.06172677502036095, 0.028573052957654, 0.1375623345375061, 0.2926015257835388, 0.17741695046424866, 0.13592077791690826, 0.025488857179880142, 0.0726943239569664, 0.09770844131708145, 0.050709616392850876, 0.04594658315181732, 0.009083828888833523, 0.024983327835798264, 0.021837929263710976, 0.11926575750112534, 0.11382617056369781, 0.22249171137809753, 0.3826439678668976, 0.22458447515964508, 0.24531354010105133, 0.05176876112818718, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0018199050100520253, 1.759366932674311e-05, 0.005607981700450182, 0.029583722352981567, 0.009902501478791237, 0.00240499060600996, 0.016255119815468788, 0.008434450253844261, 0.0070381201803684235, 0.006882159970700741, 0.008103356696665287, 0.009371891617774963, 3.180988642270677e-05, 0.0005422193789854646, 0.14323127269744873, 0.28158777952194214, 0.045097555965185165, 0.02117414027452469, 0.05809389799833298, 0.0014524150174111128, 0.006964406464248896, 0.010582090355455875, 0.011965163983404636, 0.02265000529587269, 0.020484870299696922, 0.019729144871234894, 0.028731632977724075, 0.004907289054244757, 0.0051048253662884235, 0.00039794077747501433, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04913536086678505, 0.005111359525471926, 0.3943053185939789, 0.16504207253456116, 0.1333204060792923, 0.007373967207968235, 0.00649205781519413, 0.005781218875199556, 0.0696163922548294, 0.17078818380832672, 0.43588367104530334, 0.2441176176071167, 0.044073574244976044, 0.13962700963020325, 0.0038013174198567867, 0.18024474382400513, 0.03336771950125694, 0.025161737576127052, 0.03788529708981514, 0.010167604312300682, 0.0039537386037409306, 3.701886089402251e-05, 0.046124417334795, 0.08654022216796875, 0.06664562225341797, 0.11276466399431229, 0.09791301190853119, 0.08758807182312012, 0.277656227350235, 0.5478507876396179, 0.06896418333053589, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02972331829369068, 0.032405998557806015, 0.13676248490810394, 0.2985995411872864, 0.6838041543960571, 0.17950911819934845, 0.02566559985280037, 0.299430251121521, 0.06906868517398834, 0.09219349920749664, 0.14271143078804016, 0.15384355187416077, 0.31184810400009155, 0.37699857354164124, 0.11869719624519348, 0.10793236643075943, 0.04864804446697235, 0.0019557650666683912, 0.14817607402801514, 0.0378977507352829, 0.049347102642059326, 0.0036467635072767735, 0.0038541490212082863, 0.0034904496278613806, 0.0012115711579099298, 0.047197386622428894, 0.05697714909911156, 0.11328870058059692, 0.8784908056259155, 0.019691603258252144, 0.23420120775699615, 0.004765921737998724, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.035901740193367004, 0.049252428114414215, 0.13651704788208008, 0.3431343734264374, 0.4621880352497101, 0.07741573452949524, 0.035817742347717285, 0.1879495084285736, 0.09167803823947906, 0.15167558193206787, 0.20264029502868652, 0.22310277819633484, 0.27972275018692017, 0.27912822365760803, 0.1079779863357544, 0.1524984985589981, 0.08107080310583115, 0.005865868646651506, 0.00971321389079094, 0.007243088912218809, 0.011549782939255238, 0.00268083019182086, 0.03457775339484215, 0.0031127233523875475, 0.000510410696733743, 0.009807620197534561, 0.008875550702214241, 0.023541534319519997, 0.527433454990387, 0.015368063934147358, 0.16288210451602936, 0.20708848536014557, 0.014573587104678154, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03869367763400078, 0.07609386742115021, 0.09811960905790329, 0.19582945108413696, 0.7770717144012451, 0.05828123167157173, 0.03398818522691727, 0.4334997236728668, 0.06648975610733032, 0.07675088942050934, 0.06197739765048027, 0.7435874938964844, 0.14106591045856476, 0.2445826381444931, 0.04634908586740494, 0.16305263340473175, 0.020936982706189156, 0.020989498123526573, 0.007437185384333134, 0.034894589334726334, 0.016221558675169945, 0.04928300529718399, 0.02460765466094017, 0.006940784398466349, 0.010303718037903309, 0.11923910677433014, 0.002430608496069908, 0.020191287621855736, 0.019723495468497276, 0.015607062727212906, 0.14493703842163086, 0.29023703932762146, 0.2954525649547577, 0.024419967085123062, NaN, NaN, NaN, NaN, NaN, NaN], [0.0033209763932973146, 0.0013802923494949937, 0.007923663593828678, 0.01537866611033678, 0.27329060435295105, 0.0012711664894595742, 0.000925537955481559, 0.0031033798586577177, 0.00518713379278779, 0.008014743216335773, 0.01865261048078537, 0.32840412855148315, 0.015081376768648624, 0.0187647957354784, 0.007287481799721718, 0.04235544800758362, 0.014461617916822433, 0.006770138628780842, 0.009241613559424877, 0.002999901305884123, 0.0037356300745159388, 0.00043396188993938267, 0.005936506669968367, 0.00027135247364640236, 0.00836905650794506, 0.38652852177619934, 0.1805782914161682, 0.00859912484884262, 0.13720881938934326, 0.026457296684384346, 0.044793374836444855, 0.41905051469802856, 0.48846107721328735, 0.271888792514801, 0.02787640690803528, NaN, NaN, NaN, NaN, NaN], [0.012120293453335762, 0.00801909901201725, 0.05887366458773613, 0.08173726499080658, 0.42918333411216736, 0.0074272770434618, 0.018144551664590836, 0.002390465000644326, 0.19959968328475952, 0.01595914363861084, 0.19477497041225433, 0.24081164598464966, 0.32190656661987305, 0.2620943486690521, 0.06223426014184952, 0.03824670985341072, 0.05110237002372742, 0.016365332528948784, 0.027689939364790916, 0.004054062534123659, 0.0016762956511229277, 0.0059990487061440945, 0.061629924923181534, 0.02193543128669262, 0.004144957754760981, 0.11336920410394669, 0.0855039581656456, 0.16943661868572235, 0.007511935196816921, 0.0029296777211129665, 0.005633122753351927, 0.04470856487751007, 0.19621509313583374, 0.1449754536151886, 0.4407651424407959, 0.012849990278482437, NaN, NaN, NaN, NaN], [0.001324097509495914, 1.9873512428603135e-05, 0.0026336663868278265, 0.025088831782341003, 0.006480309646576643, 0.0015246026450768113, 0.009156930260360241, 0.006450172513723373, 0.006447002291679382, 0.003797400277107954, 0.0037222199607640505, 0.006030225194990635, 1.9453302229521796e-05, 0.0003723614208865911, 0.13770580291748047, 0.29710885882377625, 0.04157622903585434, 0.022785142064094543, 0.06820578873157501, 0.0019051277777180076, 0.004196317866444588, 0.012664434500038624, 0.010533612221479416, 0.00958634540438652, 0.006948783528059721, 0.024731770157814026, 0.04424457997083664, 0.0092665059491992, 0.008317369967699051, 0.00025302590802311897, 0.03921425715088844, 0.024433301761746407, 0.005475904326885939, 0.02041386440396309, 0.005526822991669178, 0.006030899006873369, 0.000147900907904841, NaN, NaN, NaN], [0.23361828923225403, 0.06709202378988266, 0.7719610333442688, 0.734594464302063, 0.7922726273536682, 0.049216482788324356, 0.04663456231355667, 0.060855433344841, 0.40224209427833557, 0.20935069024562836, 0.5060975551605225, 0.5454070568084717, 0.2919921875, 0.420108824968338, 0.08753460645675659, 0.15116539597511292, 0.029300624504685402, 0.014213098213076591, 0.04858435317873955, 0.008192096836864948, 0.0029929669108241796, 0.00010039177868748084, 0.02851700410246849, 0.014845605008304119, 0.01335279829800129, 0.07330357283353806, 0.08230004459619522, 0.06801280379295349, 0.12962418794631958, 0.38807213306427, 0.021973537281155586, 0.0005578201962634921, 0.13413770496845245, 0.18835364282131195, 0.15109674632549286, 0.5815849900245667, 0.6008182764053345, 0.10515720397233963, NaN, NaN], [0.01675574854016304, 0.0394110269844532, 0.07827049493789673, 0.20941881835460663, 0.5690934658050537, 0.13831959664821625, 0.015872817486524582, 0.2790753245353699, 0.07380014657974243, 0.05484941974282265, 0.11329877376556396, 0.046586740761995316, 0.27540746331214905, 0.3769146502017975, 0.12728242576122284, 0.05911188945174217, 0.013889956288039684, 0.00048160224105231464, 0.10393460839986801, 0.009916743263602257, 0.013972792774438858, 0.0005543273873627186, 0.0008135904208756983, 0.0005866698920726776, 0.00012856724788434803, 0.016669562086462975, 0.022332170978188515, 0.03126570209860802, 0.39481881260871887, 0.0021035531535744667, 0.09696949273347855, 0.0003469766234047711, 0.012058700434863567, 0.1351245492696762, 0.1276140809059143, 0.8529128432273865, 0.013427066616714, 0.3029053509235382, 0.0016288348706439137, NaN], [0.13399043679237366, 0.38312259316444397, 0.21414920687675476, 0.1335369348526001, 0.883351743221283, 0.17629003524780273, 0.21391625702381134, 0.35840436816215515, 0.7405950427055359, 0.11166028678417206, 0.2222289741039276, 0.2562817633152008, 0.20710349082946777, 0.2988908290863037, 0.10401280969381332, 0.22241219878196716, 0.00997188687324524, 0.004307668190449476, 0.0318865031003952, 0.026490027084946632, 0.04937301576137543, 0.016565896570682526, 0.0013930558925494552, 0.01958940364420414, 0.015218929387629032, 0.1830211728811264, 0.11458480358123779, 0.1729872077703476, 0.047152113169431686, 0.017883911728858948, 0.118315190076828, 0.07728181034326553, 0.31889867782592773, 0.1497264951467514, 0.2596881091594696, 0.15263305604457855, 0.024473916739225388, 0.19167250394821167, 0.12363447993993759, 0.010316992178559303]], [[0.03249572962522507, 0.01680905371904373, 0.01368993055075407, 0.005182549823075533, 0.0014828554121777415, 0.0045396420173347, 0.0006250899168662727, 0.01684878207743168, 0.005824672989547253, 0.007428525947034359, 0.009805276058614254, 0.003550198394805193, 0.007900950498878956, 0.009690256789326668, 0.18011362850666046, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.11159665137529373, 0.10346578061580658, 0.414338618516922, 0.08694489300251007, 0.2136271595954895, 0.10264819115400314, 0.023593097925186157, 0.0335584320127964, 0.0575689822435379, 0.06024341657757759, 0.1307218372821808, 0.13801440596580505, 0.1756829470396042, 0.14866231381893158, 0.1320090889930725, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1948547214269638, 0.038279034197330475, 0.07790879160165787, 0.04177340865135193, 0.004589961376041174, 0.0009778933599591255, 0.002051346004009247, 0.006739486940205097, 0.009280361235141754, 0.0007642557029612362, 0.0012637393083423376, 0.00433916924521327, 0.00236115837469697, 0.008354227058589458, 0.2381056696176529, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07799407094717026, 0.10201291739940643, 0.037178199738264084, 0.03369736298918724, 0.035083431750535965, 0.003606606973335147, 0.0009816481033340096, 0.010917055420577526, 0.019562464207410812, 0.004011118784546852, 0.0029224867466837168, 0.0011325542582198977, 0.00486336974427104, 0.007979645393788815, 0.2784355580806732, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.11467810720205307, 0.4025481641292572, 0.4041208028793335, 0.13489782810211182, 0.520052433013916, 0.013409112580120564, 0.0056337821297347546, 0.04408307746052742, 0.06485209614038467, 0.0023049998562783003, 0.0050890627317130566, 0.004091872368007898, 0.006159461103379726, 0.0242836382240057, 0.07189745455980301, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1516697108745575, 0.2241159826517105, 0.5074643492698669, 0.3874017000198364, 0.2519407868385315, 0.032381314784288406, 0.015091626904904842, 0.006451433524489403, 0.09749187529087067, 0.007731522433459759, 0.00912014115601778, 0.029297562316060066, 0.05765664204955101, 0.059585090726614, 0.023513801395893097, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01171550527215004, 0.10137046873569489, 0.870269238948822, 0.5154522657394409, 0.6626715660095215, 0.08923148363828659, 0.047533176839351654, 0.015608957968652248, 0.11948943883180618, 0.008091520518064499, 0.008133050054311752, 0.012773845344781876, 0.051611315459012985, 0.01502595841884613, 0.00961183663457632, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01722140610218048, 0.036506716161966324, 0.7147647738456726, 0.20675897598266602, 0.8291797637939453, 0.31030455231666565, 0.11803850531578064, 0.03327609598636627, 0.4245462417602539, 0.013293992727994919, 0.008976193144917488, 0.054750751703977585, 0.1754072904586792, 0.04528210312128067, 0.012820743955671787, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01982569508254528, 0.15988187491893768, 0.12975367903709412, 0.1326102912425995, 0.6299260258674622, 0.28946900367736816, 0.34108322858810425, 0.11804011464118958, 0.16752222180366516, 0.01777276024222374, 0.0021109972149133682, 0.0006076672580093145, 0.0030632279813289642, 0.00126487051602453, 0.1333881914615631, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.005461913999170065, 0.03046412020921707, 0.008993657305836678, 0.005659051705151796, 0.004244270734488964, 0.02773391455411911, 0.042834386229515076, 0.13534432649612427, 0.27069228887557983, 0.04962563514709473, 0.015227400697767735, 0.0016283531440421939, 0.0014969720505177975, 0.0027089377399533987, 0.17130999267101288, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01672529987990856, 0.10339350253343582, 0.009749630466103554, 0.02030925825238228, 0.017326004803180695, 0.03957638517022133, 0.030999623239040375, 0.10308665037155151, 0.5008098483085632, 0.09767498821020126, 0.09780175238847733, 0.025981366634368896, 0.003117683343589306, 0.00962040200829506, 0.1932818591594696, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.026731140911579132, 0.05838552862405777, 0.07611822336912155, 0.05796685442328453, 0.5904980301856995, 0.010755263268947601, 0.0517524816095829, 0.055663660168647766, 0.29654714465141296, 0.1307908594608307, 0.1585402488708496, 0.03976760059595108, 0.07525579631328583, 0.16488958895206451, 0.1035238653421402, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.024593327194452286, 0.12932555377483368, 0.13568159937858582, 0.16021546721458435, 0.3227141201496124, 0.029398979619145393, 0.01611196994781494, 0.016819216310977936, 0.2378186136484146, 0.5602607131004333, 0.7615779638290405, 0.08417549729347229, 0.10783103108406067, 0.2013072967529297, 0.06744378060102463, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.018169090151786804, 0.26050350069999695, 0.078061044216156, 0.023439347743988037, 0.05254700779914856, 0.0014709478709846735, 0.002907117595896125, 0.009980114176869392, 0.1381266713142395, 0.5626046061515808, 0.5405392646789551, 0.11909772455692291, 0.008021530695259571, 0.06359856575727463, 0.009888176806271076, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08646434545516968, 0.009946366772055626, 0.041608210653066635, 0.009163393639028072, 0.12723588943481445, 0.17822976410388947, 0.01437843032181263, 0.0057503837160766125, 0.008486853912472725, 0.002935740165412426, 0.019836073741316795, 0.07525425404310226, 0.02854214422404766, 0.0230310820043087, 0.1518138200044632, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.169734388589859, 0.018695855513215065, 0.1739528477191925, 0.1591939628124237, 0.2628772258758545, 0.10412096232175827, 0.10786166787147522, 0.024563027545809746, 0.26776236295700073, 0.15710414946079254, 0.04751116409897804, 0.10171505063772202, 0.02745870314538479, 0.022933470085263252, 0.11237789690494537, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04881957918405533, 0.17062845826148987, 0.0187830850481987, 0.030382977798581123, 0.08311481773853302, 0.03788991644978523, 0.005156277678906918, 0.026916639879345894, 0.06639944016933441, 0.03180782124400139, 0.02173716016113758, 0.05343012511730194, 0.01850084401667118, 0.0033381145913153887, 0.04681381955742836, 0.12855423986911774, 0.11611904203891754, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11046597361564636, 0.13029024004936218, 0.30802851915359497, 0.31618139147758484, 0.21513698995113373, 0.08858107775449753, 0.07770872116088867, 0.030179373919963837, 0.2956576347351074, 0.19506438076496124, 0.06668522953987122, 0.15814362466335297, 0.07954283803701401, 0.09008871018886566, 0.11347464472055435, 0.1812644749879837, 0.04049589857459068, 0.04480821266770363, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14630576968193054, 0.10272074490785599, 0.06626180559396744, 0.39613619446754456, 0.5213132500648499, 0.09462913125753403, 0.19745559990406036, 0.14176879823207855, 0.45916420221328735, 0.2814978361129761, 0.19076579809188843, 0.7478294968605042, 0.15201923251152039, 0.4428024888038635, 0.11204658448696136, 0.14001408219337463, 0.11702272295951843, 0.5616602897644043, 0.021032487973570824, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17077980935573578, 0.372023344039917, 0.03066021017730236, 0.20403380692005157, 0.25160810351371765, 0.047236956655979156, 0.19034826755523682, 0.09997845441102982, 0.22249065339565277, 0.14956896007061005, 0.12211201339960098, 0.43811750411987305, 0.32559871673583984, 0.4463178217411041, 0.1688702404499054, 0.17309650778770447, 0.011261633597314358, 0.0023054813500493765, 0.0014516497030854225, 0.17103753983974457, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.001587467617355287, 0.0028523027431219816, 0.001275891438126564, 0.007771230302751064, 0.06833823025226593, 0.016362184658646584, 0.01554875634610653, 0.0395360104739666, 0.020186755806207657, 0.02848842740058899, 0.006796931382268667, 0.08043718338012695, 0.1258731484413147, 0.048048797994852066, 0.14538481831550598, 0.21775518357753754, 0.1599237471818924, 0.031671781092882156, 0.0027859890833497047, 0.1030324175953865, 0.009803196415305138, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19441094994544983, 0.026329312473535538, 0.03907056525349617, 0.5187185406684875, 0.06508557498455048, 0.04464683309197426, 0.23734036087989807, 0.10510969161987305, 0.23671847581863403, 0.2550508677959442, 0.2969563603401184, 0.31371036171913147, 0.023362383246421814, 0.04756302013993263, 0.09379850327968597, 0.1265520304441452, 0.2245447188615799, 0.3357183039188385, 0.19591355323791504, 0.030100535601377487, 0.11038237810134888, 0.012957160361111164, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.009693926200270653, 0.06855454295873642, 0.04046608507633209, 0.021632034331560135, 0.07003092765808105, 0.1099655032157898, 0.02166297659277916, 0.14673617482185364, 0.08559776097536087, 0.021444879472255707, 0.06376301497220993, 0.07838241755962372, 0.2981177270412445, 0.05645254626870155, 0.11510419100522995, 0.12113019824028015, 0.07331034541130066, 0.073086217045784, 0.038516201078891754, 0.16168329119682312, 0.12152494490146637, 0.1929183006286621, 0.11648087203502655, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1475960612297058, 0.11415769904851913, 0.09677327424287796, 0.22716772556304932, 0.05128113925457001, 0.0685737207531929, 0.17258046567440033, 0.05221087113022804, 0.2985250651836395, 0.36185649037361145, 0.6199293732643127, 0.5016448497772217, 0.08136574923992157, 0.06544326990842819, 0.09482244402170181, 0.15162895619869232, 0.16000056266784668, 0.47010278701782227, 0.008242717012763023, 0.016423694789409637, 0.19619418680667877, 0.014187236316502094, 0.2187093049287796, 0.3917299807071686, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16866622865200043, 0.03890697658061981, 0.038960762321949005, 0.045146964490413666, 0.003443084890022874, 0.025941072031855583, 0.02535194903612137, 0.01214737631380558, 0.39030662178993225, 0.11890958994626999, 0.2736153304576874, 0.3244759440422058, 0.00968784186989069, 0.014615286141633987, 0.03826850652694702, 0.1371021270751953, 0.24055053293704987, 0.39826682209968567, 0.0653936043381691, 0.06886317580938339, 0.1729464828968048, 0.02453671395778656, 0.2748231589794159, 0.23215962946414948, 0.03306089714169502, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08395736664533615, 0.10560688376426697, 0.29490047693252563, 0.15838190913200378, 0.20854075253009796, 0.047574300318956375, 0.025914132595062256, 0.0076736449263989925, 0.23083198070526123, 0.11239635199308395, 0.08150741457939148, 0.3915822207927704, 0.126749187707901, 0.08327525854110718, 0.07453686743974686, 0.05615014582872391, 0.17226241528987885, 0.4426397681236267, 0.534454345703125, 0.0034056571312248707, 0.0038566330913454294, 0.24011781811714172, 0.31882721185684204, 0.4456172287464142, 0.1489524245262146, 0.03087311051785946, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08537011593580246, 0.01334940642118454, 0.026223814114928246, 0.09485415369272232, 0.04081009700894356, 0.021519087255001068, 0.04835912212729454, 0.008561250753700733, 0.1425430029630661, 0.15310505032539368, 0.12245412170886993, 0.15674236416816711, 0.03265313804149628, 0.020860055461525917, 0.1338454782962799, 0.037336766719818115, 0.065662682056427, 0.18869149684906006, 0.795316219329834, 0.14649540185928345, 0.021824514493346214, 0.13452036678791046, 0.026823654770851135, 0.35548609495162964, 0.18523786962032318, 0.020790524780750275, 0.09485815465450287, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.009048069827258587, 0.008220783434808254, 0.0010462020291015506, 0.0073586152866482735, 0.01628630980849266, 0.0030796914361417294, 0.0014804736711084843, 0.0016866090008988976, 0.021953675895929337, 0.024090107530355453, 0.02321471832692623, 0.2417944222688675, 0.00791110284626484, 0.012413977645337582, 0.02231968566775322, 0.17983746528625488, 0.09746579825878143, 0.46259593963623047, 0.706605851650238, 0.09193093329668045, 0.2823830544948578, 0.007526541594415903, 0.10234087705612183, 0.24847157299518585, 0.2038285881280899, 0.012590465135872364, 0.002493936335667968, 0.04428662359714508, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02412300556898117, 0.02128133550286293, 0.018482450395822525, 0.016898121684789658, 0.07439899444580078, 0.03563898429274559, 0.04473365843296051, 0.0026737016160041094, 0.06965204328298569, 0.10727399587631226, 0.046027760952711105, 0.33166152238845825, 0.12371443957090378, 0.07036767154932022, 0.15801618993282318, 0.1421777307987213, 0.23310348391532898, 0.2705342471599579, 0.5351002812385559, 0.02795390971004963, 0.06031421944499016, 0.012775074690580368, 0.20022329688072205, 0.6570897698402405, 0.2668534517288208, 0.033325545489788055, 0.023841219022870064, 0.1455993354320526, 0.03172359615564346, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.007644897326827049, 0.000292555516352877, 0.08444877713918686, 0.17402730882167816, 0.16615508496761322, 0.013423392549157143, 0.054235123097896576, 0.007257240824401379, 0.08712441474199295, 0.012547464109957218, 0.0328214131295681, 0.2736492455005646, 0.0037261026445776224, 0.09982366114854813, 0.13941559195518494, 0.11665362864732742, 0.1886645257472992, 0.03897944837808609, 0.07137740403413773, 0.15634050965309143, 0.15400150418281555, 0.13745756447315216, 0.05537642911076546, 0.2729690372943878, 0.04749782383441925, 0.05948880687355995, 0.014797642827033997, 0.11365658044815063, 0.002582019427791238, 0.20324750244617462, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07466596364974976, 0.11066461354494095, 0.02582395263016224, 0.1052846685051918, 0.0988694354891777, 0.13372771441936493, 0.10285167396068573, 0.04043884575366974, 0.12614820897579193, 0.00874736811965704, 0.006169801577925682, 0.3642371892929077, 0.13258321583271027, 0.14621633291244507, 0.16873647272586823, 0.29635345935821533, 0.04781435802578926, 0.41243496537208557, 0.03004680573940277, 0.13952067494392395, 0.045467544347047806, 4.634694050764665e-05, 0.20948387682437897, 0.002634957665577531, 0.005124728661030531, 0.0019075855379924178, 0.0009838729165494442, 0.0013485344825312495, 0.004148871172219515, 0.03574635088443756, 0.23113909363746643, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.23522600531578064, 0.0398484542965889, 0.3737937808036804, 0.288825660943985, 0.10485613346099854, 0.11366727948188782, 0.29695606231689453, 0.06251946091651917, 0.35146233439445496, 0.04921486973762512, 0.25325968861579895, 0.33112239837646484, 0.06967249512672424, 0.050063006579875946, 0.0896972194314003, 0.22071197628974915, 0.019423967227339745, 0.06694509834051132, 0.2386176735162735, 0.015943216159939766, 0.14270655810832977, 0.039743710309267044, 0.014324809424579144, 0.581375777721405, 0.040944233536720276, 0.011615565046668053, 0.02482481673359871, 0.06486763060092926, 0.002298883395269513, 0.009274494834244251, 0.012798607349395752, 0.009606687352061272, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1151093989610672, 0.085483118891716, 0.1238018348813057, 0.10984596610069275, 0.07372570037841797, 0.07080911099910736, 0.04283013194799423, 0.011434272862970829, 0.6184931993484497, 0.031299810856580734, 0.1232943907380104, 0.4399086534976959, 0.16973690688610077, 0.18915507197380066, 0.06319096684455872, 0.04979729279875755, 0.005993144121021032, 0.05621323734521866, 0.3196869492530823, 0.0036542851012200117, 0.006608159281313419, 0.07202935218811035, 0.023804083466529846, 0.08581908792257309, 0.002907529706135392, 0.0022882334887981415, 0.155064657330513, 0.6752456426620483, 0.19066885113716125, 0.033486951142549515, 0.1545412391424179, 0.3257397711277008, 0.07836033403873444, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.23179487884044647, 0.03441762179136276, 0.058240070939064026, 0.17834095656871796, 0.049968671053647995, 0.038375332951545715, 0.05405527353286743, 0.00672679441049695, 0.09475977718830109, 0.0764862671494484, 0.1440851390361786, 0.11337311565876007, 0.06998162716627121, 0.031302694231271744, 0.13650138676166534, 0.02027127519249916, 0.036089565604925156, 0.0908525288105011, 0.6094546914100647, 0.035198476165533066, 0.01578100211918354, 0.08828305453062057, 0.00740778585895896, 0.08938029408454895, 0.055872198194265366, 0.01406459603458643, 0.05842210724949837, 0.7085317969322205, 0.04043729975819588, 0.00861792266368866, 0.05839632451534271, 0.306302547454834, 0.11257344484329224, 0.09490343183279037, NaN, NaN, NaN, NaN, NaN, NaN], [0.037197839468717575, 0.022889001294970512, 0.00443503400310874, 0.02830665186047554, 0.056754183024168015, 0.011282439343631268, 0.008815057575702667, 0.005641489755362272, 0.03366301208734512, 0.01200089417397976, 0.022881681099534035, 0.24835483729839325, 0.020306341350078583, 0.028865927830338478, 0.09140723943710327, 0.2219613641500473, 0.0726998969912529, 0.3657586872577667, 0.6172192692756653, 0.07194076478481293, 0.17607101798057556, 0.009873087517917156, 0.09032700955867767, 0.1240842267870903, 0.06592906266450882, 0.021971723064780235, 0.004476875066757202, 0.04292584955692291, 0.013240871019661427, 0.03868407383561134, 0.0364602766931057, 0.007298360578715801, 0.02817610278725624, 0.0009550384129397571, 0.033005379140377045, NaN, NaN, NaN, NaN, NaN], [0.019821494817733765, 0.0461096465587616, 0.009799499064683914, 0.008886821568012238, 0.03164605051279068, 0.03408728539943695, 0.06531291455030441, 0.004583337344229221, 0.015776870772242546, 0.0067581660114228725, 0.005247185938060284, 0.0803409293293953, 0.12878651916980743, 0.033680036664009094, 0.15540239214897156, 0.2832254469394684, 0.40537261962890625, 0.25111812353134155, 0.4335843026638031, 0.05173255130648613, 0.02949104830622673, 0.00834138598293066, 0.5043417811393738, 0.45271721482276917, 0.10732957720756531, 0.08741836994886398, 0.06616821885108948, 0.1252485066652298, 0.04288535565137863, 0.0027607728261500597, 0.11496254801750183, 0.007436650805175304, 0.04789961501955986, 0.014611729420721531, 0.05419020354747772, 0.013982507400214672, NaN, NaN, NaN, NaN], [0.006374652031809092, 0.0003620072384364903, 0.05079201981425285, 0.10443739593029022, 0.13200052082538605, 0.007841442711651325, 0.04038690775632858, 0.005943085998296738, 0.04502689838409424, 0.005707652773708105, 0.010736361145973206, 0.17095635831356049, 0.0034604808315634727, 0.08947119116783142, 0.1356668770313263, 0.1133793368935585, 0.2190774381160736, 0.04727642610669136, 0.08785698562860489, 0.22799502313137054, 0.1395695060491562, 0.17899513244628906, 0.05776361748576164, 0.19579172134399414, 0.03426501154899597, 0.08577524870634079, 0.027239171788096428, 0.22711482644081116, 0.005856664851307869, 0.3394412696361542, 0.03666312247514725, 0.053877539932727814, 0.02460121363401413, 0.02095765992999077, 0.08733106404542923, 0.0007995758787728846, 0.19509249925613403, NaN, NaN, NaN], [0.05784226581454277, 0.06101800128817558, 0.011293647810816765, 0.030310506001114845, 0.02692366950213909, 0.10355494171380997, 0.1643158346414566, 0.02146345190703869, 0.10686127096414566, 0.0006235101609490812, 0.001034505432471633, 0.12770172953605652, 0.08152752369642258, 0.06569667905569077, 0.13584844768047333, 0.32134389877319336, 0.08582156896591187, 0.36053547263145447, 0.06279635429382324, 0.1449708491563797, 0.041098933666944504, 0.0002254477294627577, 0.3326246738433838, 0.0031729326583445072, 0.011426791548728943, 0.00305219367146492, 0.0021134610287845135, 0.0029090954922139645, 0.0035086346324533224, 0.0884322077035904, 0.7275413274765015, 4.6366836613742635e-05, 0.004567307885736227, 0.00048746803076937795, 0.0006845259922556579, 0.00036436106893233955, 0.0336419902741909, 0.19370199739933014, NaN, NaN], [0.24130187928676605, 0.04057329148054123, 0.37395209074020386, 0.32695549726486206, 0.18701796233654022, 0.1542418897151947, 0.4307348132133484, 0.07850468903779984, 0.24226921796798706, 0.027551302686333656, 0.17328326404094696, 0.256756991147995, 0.1007629856467247, 0.0746576264500618, 0.1026487648487091, 0.2431764006614685, 0.00993723887950182, 0.023469794541597366, 0.12711890041828156, 0.013049022294580936, 0.09880916029214859, 0.014819139614701271, 0.015189954079687595, 0.19677633047103882, 0.012298321351408958, 0.006653454154729843, 0.017306946218013763, 0.044382814317941666, 0.005554118659347296, 0.008197239600121975, 0.025704391300678253, 0.01238576602190733, 0.005520223639905453, 0.018611198291182518, 0.07344726473093033, 0.00026948421145789325, 0.012129159644246101, 0.01222553662955761, 0.005697384011000395, NaN], [0.18065117299556732, 0.0850963443517685, 0.37481072545051575, 0.36960142850875854, 0.042269542813301086, 0.04689870774745941, 0.10553675144910812, 0.031215613707900047, 0.03850337490439415, 0.055640675127506256, 0.11964564025402069, 0.20274300873279572, 0.22541530430316925, 0.07314471900463104, 0.12492100149393082, 0.018590128049254417, 0.012204503640532494, 0.0029425490647554398, 0.01610950194299221, 0.024503106251358986, 0.04006015509366989, 0.018976394087076187, 0.006591797806322575, 0.002320006489753723, 0.001339062349870801, 0.028667215257883072, 0.03959575667977333, 0.00960585381835699, 0.009797154925763607, 0.022796805948019028, 0.1637655347585678, 0.20084494352340698, 0.05620957538485527, 0.12549559772014618, 0.022888751700520515, 0.037492163479328156, 0.04711981862783432, 0.44462573528289795, 0.3949664235115051, 0.3300856053829193]], [[0.7472922801971436, 0.06644202023744583, 0.12477048486471176, 0.07691145688295364, 0.17426471412181854, 0.17453429102897644, 0.8713244795799255, 0.22852616012096405, 0.7413471937179565, 0.5253387689590454, 0.16250024735927582, 0.19445888698101044, 0.10716042667627335, 0.2310180366039276, 0.05536508187651634, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.13811203837394714, 0.40626850724220276, 0.2430061399936676, 0.22277961671352386, 0.18414726853370667, 0.21574343740940094, 0.8225958943367004, 0.5822084546089172, 0.41659367084503174, 0.35776287317276, 0.4909748136997223, 0.39181941747665405, 0.34554892778396606, 0.6003718972206116, 0.043436333537101746, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03130434453487396, 0.0024298657663166523, 0.43690061569213867, 0.5043830275535583, 0.07530603557825089, 0.015139158815145493, 0.03498073294758797, 0.012510559521615505, 0.6034607291221619, 0.7801509499549866, 0.8402397036552429, 0.5008089542388916, 0.17657218873500824, 0.11879491806030273, 0.05205746740102768, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.09661327302455902, 0.049034956842660904, 0.05331439897418022, 0.7222777009010315, 0.25703296065330505, 0.020087046548724174, 0.06235986202955246, 0.0651831179857254, 0.32113927602767944, 0.5460676550865173, 0.7442458271980286, 0.5571728348731995, 0.08091285824775696, 0.059992171823978424, 0.029936296865344048, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00972762517631054, 0.007879518903791904, 0.02767527848482132, 0.019306808710098267, 0.22303025424480438, 0.007516835816204548, 0.007440114859491587, 0.022099999710917473, 0.29848337173461914, 0.9075287580490112, 0.5192471742630005, 0.8959035873413086, 0.055479276925325394, 0.04288056865334511, 0.021558567881584167, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03836950287222862, 0.05839527025818825, 0.005887853913009167, 0.08494037389755249, 0.012977076694369316, 0.5726994872093201, 0.09935679286718369, 0.13719113171100616, 0.448569655418396, 0.5218547582626343, 0.13800226151943207, 0.1732572466135025, 0.4354798197746277, 0.4542965292930603, 0.12337890267372131, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.17566490173339844, 0.03925755247473717, 0.01956782303750515, 0.04187121242284775, 0.02149910107254982, 0.049183186143636703, 0.5663522481918335, 0.045388396829366684, 0.45039302110671997, 0.19015204906463623, 0.22913624346256256, 0.10953018814325333, 0.21400360763072968, 0.572381854057312, 0.1667298972606659, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2136794924736023, 0.20810233056545258, 0.08830246329307556, 0.27903637290000916, 0.02317022904753685, 0.10591837763786316, 0.15087167918682098, 0.5299598574638367, 0.3452024757862091, 0.15965056419372559, 0.2765912711620331, 0.516273021697998, 0.2846863567829132, 0.3888777792453766, 0.0719258189201355, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07398565858602524, 0.04620325192809105, 0.3374384939670563, 0.19415578246116638, 0.025615269318223, 0.010194968432188034, 0.018451105803251266, 0.0005573831731453538, 0.5073301196098328, 0.25312942266464233, 0.15244188904762268, 0.143111914396286, 0.051979612559080124, 0.04884689673781395, 0.12363318353891373, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5805832147598267, 0.09438126534223557, 0.24455930292606354, 0.06023820489645004, 0.03943831846117973, 0.021930387243628502, 0.026398053392767906, 0.012488989159464836, 0.011794325895607471, 0.767930269241333, 0.4412824809551239, 0.07896611094474792, 0.01228941697627306, 0.018458310514688492, 0.10866446793079376, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1145540103316307, 0.05171298235654831, 0.7072227597236633, 0.4839639961719513, 0.11294537037611008, 0.06211492419242859, 0.021921994164586067, 0.0025394419208168983, 0.0033554628025740385, 0.07357389479875565, 0.7795555591583252, 0.05686911940574646, 0.022035235539078712, 0.034172482788562775, 0.07262071967124939, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08121224492788315, 0.025126218795776367, 0.4891066551208496, 0.29065003991127014, 0.20622830092906952, 0.36699986457824707, 0.07864820212125778, 0.014422299340367317, 0.016684990376234055, 0.0649130716919899, 0.07936163991689682, 0.6605017185211182, 0.18783104419708252, 0.08294262737035751, 0.03477967903017998, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0700722336769104, 0.1311686784029007, 0.5332850813865662, 0.1558467000722885, 0.36321985721588135, 0.7912644743919373, 0.32202765345573425, 0.1934671401977539, 0.031114375218749046, 0.09986341744661331, 0.08630139380693436, 0.055017780512571335, 0.44781896471977234, 0.42446693778038025, 0.1060790941119194, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08875010907649994, 0.06247853487730026, 0.4616371989250183, 0.12711729109287262, 0.3074216842651367, 0.19363558292388916, 0.2020244151353836, 0.0779867023229599, 0.019831692799925804, 0.03570472076535225, 0.07392378151416779, 0.04282142594456673, 0.0921483263373375, 0.3143211603164673, 0.22281906008720398, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5682113766670227, 0.1249876543879509, 0.7342633008956909, 0.902918815612793, 0.7035764455795288, 0.3718622326850891, 0.6157594919204712, 0.15625660121440887, 0.8438207507133484, 0.9341241121292114, 0.8159937858581543, 0.6624717712402344, 0.3264457583427429, 0.5970154404640198, 0.003644895739853382, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2626786530017853, 0.0849713385105133, 0.11954734474420547, 0.09299539029598236, 0.12019845843315125, 0.1675114780664444, 0.12060416489839554, 0.1292921006679535, 0.33819568157196045, 0.3146125078201294, 0.20831438899040222, 0.39596518874168396, 0.2145393043756485, 0.2666572332382202, 0.05294949933886528, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1368129849433899, 0.16135744750499725, 0.15528292953968048, 0.24771884083747864, 0.1416730433702469, 0.05803852900862694, 0.07394444942474365, 0.10563277453184128, 0.033661823719739914, 0.18054474890232086, 0.1985052525997162, 0.05316935107111931, 0.05009648948907852, 0.043446026742458344, 0.03412564843893051, 0.16815106570720673, 0.017178548499941826, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0030849967151880264, 0.0006440586876124144, 0.016017315909266472, 0.0037563794758170843, 0.009170617908239365, 0.0008218333241529763, 0.0032779525499790907, 0.0006974118296056986, 0.12044321000576019, 0.005983977112919092, 0.011704917997121811, 0.023849062621593475, 0.0031650178134441376, 0.01169323269277811, 0.16145823895931244, 0.2022658735513687, 0.005017802584916353, 0.01763225719332695, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02798222377896309, 0.012448069639503956, 0.018199993297457695, 0.0069459048099815845, 0.042531996965408325, 0.009718443267047405, 0.013791781850159168, 0.04370715469121933, 0.21814176440238953, 0.024645699188113213, 0.0633857473731041, 0.0802498310804367, 0.006771658081561327, 0.040147896856069565, 0.4109969139099121, 0.16166983544826508, 0.033678483217954636, 0.014520054683089256, 0.003462842432782054, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02001010812819004, 0.02580004744231701, 0.006869276985526085, 0.007543967105448246, 0.017537932842969894, 0.00023914838675409555, 0.006739956792443991, 0.008227680809795856, 0.05446772649884224, 0.03320171311497688, 0.022232946008443832, 0.01063306163996458, 0.0007752752280794084, 0.0028256638906896114, 0.2078467756509781, 0.10712886601686478, 0.3422684967517853, 0.05748933553695679, 0.2768969237804413, 0.004922540858387947, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0034786108881235123, 0.00011826713307527825, 0.002407492371276021, 0.005452741403132677, 0.002847136929631233, 0.003419033018872142, 0.013516861945390701, 0.002940082224085927, 0.002004653448238969, 0.006652397103607655, 0.004079414997249842, 0.0028307989705353975, 0.0006369714974425733, 0.002542868722230196, 0.1463778167963028, 0.047501806169748306, 0.48201972246170044, 0.4827657639980316, 0.48466482758522034, 0.022285524755716324, 0.00022009640815667808, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0762338638305664, 0.11778479814529419, 0.03105221875011921, 0.006415408570319414, 0.0190818402916193, 0.027191398665308952, 0.005222225561738014, 0.0170834269374609, 0.05309534817934036, 0.00936796236783266, 0.03816217556595802, 0.17940494418144226, 0.020440110936760902, 0.13513173162937164, 0.3000544309616089, 0.1517350822687149, 0.04445230960845947, 0.09343461692333221, 0.05873756855726242, 0.07171032577753067, 0.22849556803703308, 0.05614512786269188, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16228125989437103, 0.35454851388931274, 0.04026315361261368, 0.03822629526257515, 0.023396998643875122, 0.30800631642341614, 0.24136781692504883, 0.15176478028297424, 0.0788438618183136, 0.07347536832094193, 0.030298085883259773, 0.007365733850747347, 0.1061745211482048, 0.2841038405895233, 0.07787416130304337, 0.25680339336395264, 0.00010820403986144811, 0.0123103903606534, 0.007049524690955877, 0.001952940714545548, 0.027401963248848915, 0.0028134624008089304, 0.00041907382546924055, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05645078793168068, 0.023840615525841713, 0.013567867688834667, 0.00750470208004117, 0.07643276453018188, 0.08809614926576614, 0.06102507561445236, 0.021034346893429756, 0.039108242839574814, 0.02081543207168579, 0.011458326131105423, 0.20520520210266113, 0.027348484843969345, 0.06299317628145218, 0.2514360249042511, 0.005559808574616909, 0.007462772540748119, 0.013313480652868748, 0.017376750707626343, 0.0038542840629816055, 0.006728595122694969, 0.5333897471427917, 0.03155524656176567, 0.15571120381355286, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.016126127913594246, 0.01087501272559166, 0.01213990617543459, 0.004450921434909105, 0.014690833166241646, 0.30525338649749756, 0.02716207131743431, 0.09981174021959305, 0.027048761025071144, 0.01336466334760189, 0.006663064938038588, 0.0520603246986866, 0.042623523622751236, 0.018071996048092842, 0.1948687732219696, 0.004124458413571119, 0.004751718603074551, 0.016015900298953056, 0.01742120459675789, 0.032125748693943024, 0.010460411198437214, 0.45809611678123474, 0.07138781994581223, 0.5171095728874207, 0.17626723647117615, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04185086488723755, 0.034399643540382385, 0.041276611387729645, 0.0584070086479187, 0.019824109971523285, 0.00856409315019846, 0.08867836743593216, 0.10337970405817032, 0.09468665719032288, 0.02033121883869171, 0.018058426678180695, 0.059728462249040604, 0.09321711957454681, 0.20168805122375488, 0.1941128522157669, 0.24881334602832794, 0.005821824539452791, 0.031170587986707687, 0.009853766299784184, 0.027254868298768997, 0.01885347068309784, 0.02900754101574421, 0.013663586229085922, 0.012090054340660572, 0.0009272377355955541, 0.0030740045476704836, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01436887588351965, 0.027922889217734337, 0.046481672674417496, 0.010071231983602047, 0.026127830147743225, 0.06003356724977493, 0.022118212655186653, 0.08160483092069626, 0.07784195244312286, 0.010694753378629684, 0.017130734398961067, 0.05340806022286415, 0.041410259902477264, 0.035884104669094086, 0.2491855025291443, 0.19627800583839417, 0.054823894053697586, 0.1886557787656784, 0.00739922234788537, 0.09451853483915329, 0.01572227105498314, 0.0010023268405348063, 0.0061036646366119385, 0.0014733865391463041, 0.0003654434985946864, 0.006776102818548679, 0.0027319795917719603, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.053393200039863586, 0.04828185588121414, 0.03453819081187248, 0.013636122457683086, 0.25098806619644165, 0.12313847243785858, 0.02266266942024231, 0.017618268728256226, 0.019785437732934952, 0.005274764262139797, 0.021053072065114975, 0.20679616928100586, 0.021523641422390938, 0.03855947405099869, 0.1109846979379654, 0.07900664210319519, 0.04510375112295151, 0.002657376928254962, 0.0032053724862635136, 0.0027717212215065956, 0.008140889927744865, 0.0011833005119115114, 0.04105996713042259, 0.0017470002640038729, 0.008194361813366413, 0.019470002502202988, 0.3834601640701294, 0.013146632350981236, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12851715087890625, 0.12400124222040176, 0.2637093663215637, 0.02439347468316555, 0.07038086652755737, 0.12665364146232605, 0.04898465424776077, 0.03412041813135147, 0.0263816025108099, 0.023226425051689148, 0.11513664573431015, 0.09503531455993652, 0.1215861439704895, 0.11158601939678192, 0.14799171686172485, 0.06578069925308228, 0.08975866436958313, 0.022234706208109856, 0.015388325788080692, 0.006578383035957813, 0.011582762002944946, 0.014906905591487885, 0.04645423963665962, 0.008417387492954731, 0.0318351611495018, 0.024524353444576263, 0.5050408244132996, 0.1078883558511734, 0.09876319766044617, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0010214513167738914, 0.004835289902985096, 0.0042709591798484325, 0.0026378841139376163, 0.005866974592208862, 0.008331544697284698, 0.006240549497306347, 0.01365274004638195, 0.1720106601715088, 0.0005307683604769409, 0.0007543729152530432, 0.004353509750217199, 0.0002490385086275637, 0.0017186965560540557, 0.14317919313907623, 0.010224410332739353, 0.16048979759216309, 0.09242240339517593, 0.259725958108902, 0.06779038906097412, 0.007232773117721081, 0.09601377695798874, 0.28109633922576904, 0.2723717987537384, 0.1275584101676941, 0.06318827718496323, 0.25179460644721985, 0.2496732771396637, 0.6837621927261353, 0.0018262360244989395, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07205050438642502, 0.12816517055034637, 0.23753608763217926, 0.08243206143379211, 0.5041552186012268, 0.11970840394496918, 0.04837331175804138, 0.034129947423934937, 0.16484025120735168, 0.011070297099649906, 0.05054215341806412, 0.039082955569028854, 0.09205758571624756, 0.1322212517261505, 0.16203875839710236, 0.04991341754794121, 0.05319196358323097, 0.14821480214595795, 0.020963814109563828, 0.03095317631959915, 0.024693654850125313, 0.008621936663985252, 0.14259999990463257, 0.042305052280426025, 0.09002435952425003, 0.005839803721755743, 0.061309609562158585, 0.23589004576206207, 0.30903181433677673, 0.18008928000926971, 0.49815359711647034, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.014979850500822067, 0.03769220784306526, 0.04367470741271973, 0.009415187872946262, 0.019922776147723198, 0.11522040516138077, 0.014906312339007854, 0.04722318425774574, 0.06570684164762497, 0.008925273083150387, 0.019600573927164078, 0.0472339391708374, 0.005348374601453543, 0.0017698986921459436, 0.1612817794084549, 0.015294999815523624, 0.03185835853219032, 0.0202027577906847, 0.03976168856024742, 0.0711589902639389, 0.13473857939243317, 0.0059967683628201485, 0.0031582280062139034, 0.003374348394572735, 0.002362155122682452, 0.015532899647951126, 0.038825590163469315, 0.08611883223056793, 0.03844507411122322, 0.009673628956079483, 0.7068554162979126, 0.013729983940720558, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.023198002949357033, 0.06148262694478035, 0.046858664602041245, 0.013079512864351273, 0.08762317895889282, 0.00949429627507925, 0.0484880767762661, 0.025388503447175026, 0.04432932287454605, 0.006038118619471788, 0.010164186358451843, 0.08949221670627594, 0.06122652441263199, 0.11895263940095901, 0.16355113685131073, 0.2531464695930481, 0.013071080669760704, 0.035546887665987015, 0.020458703860640526, 0.01740572415292263, 0.009577612392604351, 0.014396607875823975, 0.05952044576406479, 0.013841827400028706, 0.0003843819722533226, 0.0024746267590671778, 0.007157978601753712, 0.013787134550511837, 0.033782534301280975, 0.003469215938821435, 0.007898973301053047, 0.05525756999850273, 0.003914556000381708, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.009917332790791988, 0.01408212911337614, 0.047434139996767044, 0.005388779100030661, 0.023170381784439087, 0.034844160079956055, 0.009820640087127686, 0.03569778800010681, 0.05789060518145561, 0.0037882563192397356, 0.013808010146021843, 0.04879388585686684, 0.03114072047173977, 0.0507131889462471, 0.18661679327487946, 0.20273520052433014, 0.05025332421064377, 0.2335304319858551, 0.009442931972444057, 0.13508503139019012, 0.0181263517588377, 0.0010557285277172923, 0.003822105238214135, 0.0018545370548963547, 0.0003744752029888332, 0.0046313730999827385, 0.0008518796530552208, 0.006319030188024044, 0.014203540980815887, 0.0018540708115324378, 0.003058186499401927, 0.002516325796023011, 0.001575352856889367, 0.0014869269216433167, NaN, NaN, NaN, NaN, NaN, NaN], [0.0652787834405899, 0.04612350836396217, 0.04522763565182686, 0.014745297841727734, 0.27657532691955566, 0.16156227886676788, 0.025164838880300522, 0.017732013016939163, 0.023105354979634285, 0.005499221384525299, 0.020183373242616653, 0.19132839143276215, 0.020515967160463333, 0.056384406983852386, 0.14304831624031067, 0.059709664434194565, 0.021975213661789894, 0.002582199638709426, 0.002308695577085018, 0.00240446999669075, 0.004605048336088657, 0.0013587460853159428, 0.04497997462749481, 0.0009150391560979187, 0.0030208472162485123, 0.016492530703544617, 0.2572183907032013, 0.006429646629840136, 0.013558420352637768, 0.06110598146915436, 0.03728436306118965, 0.019318275153636932, 0.03907725587487221, 0.4492114782333374, 0.01579420454800129, NaN, NaN, NaN, NaN, NaN], [0.14539514482021332, 0.21388974785804749, 0.34906452894210815, 0.031415559351444244, 0.062017399817705154, 0.08485611528158188, 0.03913363441824913, 0.03569692373275757, 0.023448940366506577, 0.020669998601078987, 0.1622902750968933, 0.1315622329711914, 0.09182734042406082, 0.1796703040599823, 0.13702963292598724, 0.025836847722530365, 0.04185229912400246, 0.017175624147057533, 0.005038154777139425, 0.006518983747810125, 0.0043221269734203815, 0.004393702372908592, 0.03134007006883621, 0.002082354621961713, 0.00246719503775239, 0.00855192355811596, 0.28023120760917664, 0.0558621920645237, 0.020582975819706917, 0.00264686718583107, 0.052114877849817276, 0.01051351334899664, 0.0282430537045002, 0.640393853187561, 0.11605942994356155, 0.042242906987667084, NaN, NaN, NaN, NaN], [0.0009059146977961063, 0.004442692268639803, 0.002850044285878539, 0.0024173678830266, 0.006019651889801025, 0.004450949374586344, 0.003768310882151127, 0.009272964671254158, 0.19643637537956238, 0.0004391498805489391, 0.0004852984275203198, 0.005083973053842783, 0.000164541692356579, 0.001456208759918809, 0.13767127692699432, 0.00790853425860405, 0.07249781489372253, 0.09275110065937042, 0.13612288236618042, 0.0654025748372078, 0.0028184219263494015, 0.039562828838825226, 0.11378230899572372, 0.08281006664037704, 0.029445864260196686, 0.03387679159641266, 0.16786670684814453, 0.2288694977760315, 0.6801032423973083, 0.0008468713494949043, 0.32477572560310364, 0.20243169367313385, 0.04291461780667305, 0.2565927505493164, 0.2435160130262375, 0.8255255222320557, 0.0008029205491766334, NaN, NaN, NaN], [0.03601038455963135, 0.08602340519428253, 0.042799800634384155, 0.007577326148748398, 0.12637566030025482, 0.07399067282676697, 0.02205651067197323, 0.01475659292191267, 0.14170114696025848, 0.004405674524605274, 0.013175459578633308, 0.03142356127500534, 0.06839168816804886, 0.09161193668842316, 0.1376270353794098, 0.06791312247514725, 0.034157127141952515, 0.26634278893470764, 0.01933334954082966, 0.08246968686580658, 0.03419587388634682, 0.019395295530557632, 0.1259232461452484, 0.02923283353447914, 0.07644251734018326, 0.00482177222147584, 0.03381035849452019, 0.2429695725440979, 0.4201262295246124, 0.21319957077503204, 0.1469077318906784, 0.005101305432617664, 0.05322602018713951, 0.08754345029592514, 0.4596864581108093, 0.32625797390937805, 0.2286616712808609, 0.6285872459411621, NaN, NaN], [0.014056011103093624, 0.020953036844730377, 0.03237491473555565, 0.0042424313724040985, 0.017438247799873352, 0.08849667757749557, 0.005714876111596823, 0.025588830932974815, 0.08735965192317963, 0.009712125174701214, 0.02371004782617092, 0.06271149963140488, 0.00425978796556592, 0.0027238703332841396, 0.14272134006023407, 0.0236026793718338, 0.032931454479694366, 0.018642868846654892, 0.052601076662540436, 0.09147398918867111, 0.11555580049753189, 0.00512799434363842, 0.006684163119643927, 0.005264784675091505, 0.0023014512844383717, 0.005628940649330616, 0.03778252378106117, 0.09737572073936462, 0.12753169238567352, 0.00698094442486763, 0.6853439807891846, 0.02319822832942009, 0.018658116459846497, 0.08199534565210342, 0.18709556758403778, 0.07321563363075256, 0.027500100433826447, 0.6534799337387085, 0.01572287082672119, NaN], [0.15719948709011078, 0.03286461904644966, 0.12916648387908936, 0.10299614071846008, 0.014032969251275063, 0.011700707487761974, 0.06680437922477722, 0.016068298369646072, 0.04505150765180588, 0.056866806000471115, 0.07287567108869553, 0.09101171046495438, 0.06734755635261536, 0.17371943593025208, 0.1297563910484314, 0.24674107134342194, 0.007728901691734791, 0.010779940523207188, 0.01413859985768795, 0.08573849499225616, 0.014258946292102337, 0.014431791380047798, 0.00199147523380816, 0.006254997570067644, 0.003036148613318801, 0.015209752134978771, 0.015118316747248173, 0.05811062082648277, 0.01987045258283615, 0.012226228602230549, 0.021392136812210083, 0.08141177892684937, 0.016042163595557213, 0.01565614528954029, 0.05352389067411423, 0.01607833430171013, 0.014641694724559784, 0.020306598395109177, 0.06722531467676163, 0.005379782523959875]], [[0.0183254461735487, 0.00659788167104125, 0.046570390462875366, 0.04327844828367233, 0.10241857916116714, 0.5407979488372803, 0.0026681027375161648, 0.15349310636520386, 0.0016508381813764572, 0.010916458442807198, 0.036675866693258286, 0.15769276022911072, 0.4073828458786011, 0.04228133708238602, 0.15622197091579437, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07985992729663849, 0.06383417546749115, 0.024972105398774147, 0.18746882677078247, 0.11770728975534439, 0.13333363831043243, 0.006719768047332764, 0.04288880154490471, 0.001412510173395276, 0.058754052966833115, 0.14280158281326294, 0.13529875874519348, 0.08268098533153534, 0.02367851696908474, 0.1494951695203781, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01403640117496252, 0.014278309419751167, 0.1034439280629158, 0.022417087107896805, 0.10706920921802521, 0.018271848559379578, 0.046350300312042236, 0.04233889281749725, 0.037542134523391724, 0.0005760823260061443, 0.004724643658846617, 0.233056902885437, 0.2574465572834015, 0.1892177164554596, 0.21611936390399933, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.032590243965387344, 0.14464972913265228, 0.1993260532617569, 0.12327495217323303, 0.27639931440353394, 0.011173157021403313, 0.012838426046073437, 0.0802190750837326, 0.0400678850710392, 0.013469994999468327, 0.025247203186154366, 0.30583158135414124, 0.6397863626480103, 0.258308470249176, 0.08317234367132187, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.007401467300951481, 0.04209339618682861, 0.1104009672999382, 0.04737341031432152, 0.06253770738840103, 0.0023836863692849874, 0.05026397854089737, 0.01439946424216032, 0.006556188687682152, 0.001721409265883267, 0.01908556930720806, 0.022761031985282898, 0.01600046642124653, 0.22344018518924713, 0.2855986952781677, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00031611474696546793, 0.010241325944662094, 0.005327185150235891, 0.007503898814320564, 0.009216651320457458, 0.08986854553222656, 0.0022410263773053885, 0.04830501973628998, 0.013246790505945683, 0.0036830154713243246, 0.001605262397788465, 0.004246865399181843, 0.005818811245262623, 0.00778583250939846, 0.2319662719964981, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00028042105259373784, 0.004604758229106665, 0.008834331296384335, 0.010530425235629082, 0.04934454336762428, 0.3239482641220093, 0.02964387647807598, 0.041019540280103683, 0.028070107102394104, 0.002580034313723445, 0.0034616885241121054, 0.006594499107450247, 0.07731658220291138, 0.01784621551632881, 0.10414844751358032, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.002352550160139799, 0.00811008270829916, 0.007519579492509365, 0.09616736322641373, 0.00784054771065712, 0.06404154002666473, 0.025837063789367676, 0.06720300018787384, 0.008001329377293587, 0.016075177118182182, 0.0036620565224438906, 0.031110821291804314, 0.1529460847377777, 0.03003939613699913, 0.19531111419200897, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.014062762260437012, 0.03979215770959854, 0.0070105125196278095, 0.010145032778382301, 0.023933248594403267, 0.08613994717597961, 0.027301009744405746, 0.007488427218049765, 0.04610109701752663, 0.00706111453473568, 0.005716769024729729, 0.008516461588442326, 0.04168170318007469, 0.004054774064570665, 0.3198099434375763, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0027477010153234005, 0.009237049147486687, 0.005884162615984678, 0.004349177703261375, 0.039300523698329926, 0.06504905968904495, 0.005921225529164076, 0.05048412084579468, 0.004538795445114374, 0.019958311691880226, 0.08035917580127716, 0.1339075267314911, 0.45191076397895813, 0.1108468547463417, 0.15996994078159332, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0004566281568259001, 0.0044615683145821095, 0.008062957786023617, 0.0003266451822128147, 0.032452184706926346, 0.004190187435597181, 0.0009983428753912449, 0.0015420016134157777, 0.025539150461554527, 0.0009114624699577689, 0.001308016013354063, 0.11249691247940063, 0.5262115597724915, 0.16036535799503326, 0.02284345217049122, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.006384413689374924, 0.006966868881136179, 0.013256898149847984, 0.008146845735609531, 0.005910678766667843, 0.005924733821302652, 0.0029809526167809963, 0.004338744096457958, 0.0021091948729008436, 0.02691148780286312, 0.09123647958040237, 0.0904775932431221, 0.10420377552509308, 0.019918829202651978, 0.21981710195541382, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.004395737312734127, 0.0342060811817646, 0.08344801515340805, 0.012639162130653858, 0.07537969946861267, 0.00383414002135396, 0.007808698806911707, 0.007516762241721153, 0.0023650380317121744, 0.055798787623643875, 0.025632014498114586, 0.040716953575611115, 0.16482838988304138, 0.13848447799682617, 0.17180821299552917, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0016022673808038235, 0.013307235203683376, 0.012306403368711472, 0.0029055906925350428, 0.06092625483870506, 0.01653674617409706, 0.008309547789394855, 0.00395687622949481, 0.002493055537343025, 0.0038927635177969933, 0.009680269286036491, 0.23031921684741974, 0.35693949460983276, 0.1708209365606308, 0.050492819398641586, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.009627100080251694, 0.006502249743789434, 0.0023533182684332132, 0.0021814347710460424, 0.007286426145583391, 0.024909881874918938, 0.01453662570565939, 0.010449647903442383, 0.0028000103775411844, 0.001988302916288376, 0.001580765936523676, 0.013102496974170208, 0.001836722600273788, 0.0008430163725279272, 0.15720587968826294, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.010018138214945793, 0.02516627125442028, 0.027397310361266136, 0.005101055838167667, 0.025938771665096283, 0.13529063761234283, 0.02690303698182106, 0.11719205975532532, 0.027814749628305435, 0.019565219059586525, 0.07996311038732529, 0.0991574078798294, 0.16288702189922333, 0.1113416850566864, 0.22370746731758118, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05219842493534088, 0.1440066546201706, 0.27922260761260986, 0.2058621197938919, 0.11230742931365967, 0.6016822457313538, 0.20846855640411377, 0.04777589067816734, 0.20611444115638733, 0.15481434762477875, 0.11950203776359558, 0.02679699845612049, 0.0639302060008049, 0.047183193266391754, 0.04897741973400116, 0.147435262799263, 0.06894105672836304, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01555164996534586, 0.0014379153726622462, 0.01706753298640251, 0.003720618085935712, 0.10093016922473907, 0.027928827330470085, 0.015380543656647205, 0.0025812943931668997, 0.020822137594223022, 0.014309070073068142, 0.017923271283507347, 0.0120958611369133, 0.014481468126177788, 0.009491728618741035, 0.15904544293880463, 0.18660759925842285, 0.013697005808353424, 0.050341442227363586, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11612647771835327, 0.0010205605067312717, 0.020188286900520325, 0.027076182886958122, 0.09822120517492294, 0.3221674859523773, 0.1250218003988266, 0.002691123867407441, 0.005359187722206116, 0.04976291581988335, 0.023232540115714073, 0.04237976670265198, 0.028708819299936295, 0.049411751329898834, 0.005618311930447817, 0.14907698333263397, 0.12682567536830902, 0.14014844596385956, 0.024977339431643486, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0470837838947773, 0.007497857324779034, 0.004583081230521202, 0.022991856560111046, 0.0278051495552063, 0.00051211251411587, 0.0627230703830719, 0.011764267459511757, 0.010903585702180862, 0.07272983342409134, 0.011678352952003479, 0.09392477571964264, 0.01558940764516592, 0.03351595252752304, 0.2068868726491928, 0.20074230432510376, 0.11179281026124954, 0.012457489967346191, 0.01455892063677311, 0.011106430552899837, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0024584962520748377, 8.163625898305327e-05, 0.00016154914919752628, 0.0002508168399799615, 0.0019916424062103033, 0.0004536219348665327, 0.0036078437697142363, 0.0008641426684334874, 0.00021941671730019152, 0.0014423344982787967, 0.0004360634775366634, 0.004383172374218702, 0.0009428760386072099, 0.0009436326217837632, 0.14683274924755096, 0.20768699049949646, 0.16985096037387848, 0.19526726007461548, 0.016829432919621468, 0.05647609382867813, 0.022808711975812912, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02989446185529232, 0.007703323382884264, 0.12996061146259308, 0.025068828836083412, 0.2812304198741913, 0.0071953474543988705, 0.0021352169569581747, 0.0025125211104750633, 0.0014658492291346192, 0.007028855849057436, 0.0448734275996685, 0.09462164342403412, 0.0503704659640789, 0.11768583953380585, 0.12974096834659576, 0.14349573850631714, 0.41078659892082214, 0.5100967288017273, 0.04046756774187088, 0.2924310266971588, 0.07987978309392929, 0.007180717773735523, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16756094992160797, 0.028098214417696, 0.20756086707115173, 0.2207580953836441, 0.10928753018379211, 0.13773545622825623, 0.2233184576034546, 0.1774815022945404, 0.13830231130123138, 0.20932619273662567, 0.18267595767974854, 0.05961548537015915, 0.07697918266057968, 0.18739080429077148, 0.06796090304851532, 0.11146429926156998, 0.3579395115375519, 0.7730652093887329, 0.5723751783370972, 0.2817910611629486, 0.25461745262145996, 0.060240793973207474, 0.08399515599012375, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.017068415880203247, 0.00098085415083915, 0.010854640044271946, 0.006490680854767561, 0.29060667753219604, 0.006710599176585674, 0.0118483304977417, 0.0008181483135558665, 0.00011296885350020602, 0.0034601599909365177, 0.005098147317767143, 0.010750477202236652, 0.010399019345641136, 0.009376241825520992, 0.017405353486537933, 0.13904383778572083, 0.44345301389694214, 0.1345542073249817, 0.05706587806344032, 0.7818705439567566, 0.04436418041586876, 0.015915511175990105, 0.31926584243774414, 0.26167550683021545, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1331326961517334, 0.019769106060266495, 0.01612294837832451, 0.028521019965410233, 0.007509702816605568, 0.2665199935436249, 0.19958320260047913, 0.1385747790336609, 0.0059373765252530575, 0.08046255260705948, 0.052418529987335205, 0.004961848258972168, 0.10941796749830246, 0.06705309450626373, 0.17611992359161377, 0.12236351519823074, 0.40148651599884033, 0.12099923938512802, 0.38539087772369385, 0.6352627873420715, 0.0574735552072525, 0.027495326474308968, 0.25199854373931885, 0.07788273692131042, 0.1824284791946411, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.019668979570269585, 0.0081618782132864, 0.12552350759506226, 0.0802406370639801, 0.07089362293481827, 0.18871739506721497, 0.12778939306735992, 0.04829992726445198, 0.04307088255882263, 0.02314154990017414, 0.14194107055664062, 0.05861861631274223, 0.19650596380233765, 0.11930099874734879, 0.18420156836509705, 0.0776049941778183, 0.26076433062553406, 0.12800094485282898, 0.15216867625713348, 0.36678510904312134, 0.31404268741607666, 0.13151897490024567, 0.1709745228290558, 0.2591820955276489, 0.18929390609264374, 0.08235450834035873, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00538466265425086, 0.0270208939909935, 0.18066750466823578, 0.06076826527714729, 0.035171061754226685, 0.411039799451828, 0.09634009003639221, 0.26394954323768616, 0.1915867179632187, 0.03318370133638382, 0.3213040828704834, 0.10995125770568848, 0.5320225954055786, 0.4394112527370453, 0.15243512392044067, 0.08287283033132553, 0.26698997616767883, 0.29562729597091675, 0.13922370970249176, 0.3693794012069702, 0.22139106690883636, 0.612119734287262, 0.1618482619524002, 0.40734153985977173, 0.10604425519704819, 0.2217203825712204, 0.14197519421577454, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0030147582292556763, 0.00625306461006403, 0.017102748155593872, 0.008551767095923424, 0.0727200135588646, 0.015153692103922367, 0.0023096217773854733, 0.011201570741832256, 0.002435098635032773, 0.006847116630524397, 0.016829995438456535, 0.12519565224647522, 0.3878204822540283, 0.13249750435352325, 0.028183329850435257, 0.0676846131682396, 0.5803259611129761, 0.47128230333328247, 0.2430339902639389, 0.43893957138061523, 0.5822793245315552, 0.9563859105110168, 0.5092246532440186, 0.7397804260253906, 0.6675750613212585, 0.2242172360420227, 0.046741336584091187, 0.09371624141931534, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.066617950797081, 0.006649812217801809, 0.04142908379435539, 0.13957993686199188, 0.025706114247441292, 0.08231058716773987, 0.08377126604318619, 0.02330365777015686, 0.04652002453804016, 0.11060080677270889, 0.09014575183391571, 0.07117310166358948, 0.15938407182693481, 0.1624550223350525, 0.05356656014919281, 0.16273218393325806, 0.4245251417160034, 0.44257473945617676, 0.1064363345503807, 0.22264361381530762, 0.638583779335022, 0.7456080913543701, 0.17856015264987946, 0.09681503474712372, 0.3901955187320709, 0.4154786765575409, 0.10903800278902054, 0.0281606987118721, 0.027353502810001373, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.004379222169518471, 0.0002637936850078404, 0.0022587613202631474, 0.006711117923259735, 0.0006837267428636551, 0.007989797741174698, 0.02997850626707077, 0.045127563178539276, 0.008224103599786758, 0.0034686585422605276, 0.0038658890407532454, 0.00034815416438505054, 7.646608719369397e-05, 0.00017854337056633085, 0.14325816929340363, 0.2541956901550293, 0.2554672658443451, 0.13483673334121704, 0.33163735270500183, 0.11067650467157364, 0.3400806486606598, 0.4272999167442322, 0.2955835163593292, 0.293487548828125, 0.2820315957069397, 0.17141510546207428, 0.08369391411542892, 0.012903732247650623, 0.010530934669077396, 0.015047149732708931, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.25216665863990784, 0.1422366499900818, 0.10172943770885468, 0.3735504150390625, 0.0612066313624382, 0.06238102167844772, 0.11154207587242126, 0.031159698963165283, 0.011768986470997334, 0.4107469618320465, 0.1557808816432953, 0.07179611176252365, 0.186580628156662, 0.18789765238761902, 0.099563829600811, 0.07456009835004807, 0.09125705808401108, 0.20381297171115875, 0.09053967893123627, 0.6734579801559448, 0.8927901983261108, 0.9854956865310669, 0.19160649180412292, 0.848483681678772, 0.3795100748538971, 0.0351644828915596, 0.06069617718458176, 0.0190274715423584, 0.13319239020347595, 0.1618155688047409, 0.029784632846713066, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0073658498004078865, 0.1486257165670395, 0.03456511348485947, 0.0081891855224967, 0.009660922922194004, 0.09341325610876083, 0.010183881968259811, 0.09390538185834885, 0.005950886756181717, 0.019719628617167473, 0.060451164841651917, 0.021925343200564384, 0.19991156458854675, 0.17004182934761047, 0.15761280059814453, 0.13663174211978912, 0.5250937938690186, 0.20416004955768585, 0.37758082151412964, 0.7281314134597778, 0.24714940786361694, 0.006291824858635664, 0.029336191713809967, 0.258807897567749, 0.17944614589214325, 0.2768983840942383, 0.49996671080589294, 0.6760725975036621, 0.0684136375784874, 0.9500845074653625, 0.04427658021450043, 0.027829600498080254, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0057948376052081585, 0.023180164396762848, 0.018019115552306175, 0.008233858272433281, 0.005580522585660219, 0.09526203572750092, 0.025384269654750824, 0.05396068096160889, 0.022398412227630615, 0.010895788669586182, 0.02884012460708618, 0.008390026167035103, 0.1754663735628128, 0.0998048186302185, 0.1692073941230774, 0.05520259216427803, 0.4062710404396057, 0.11698392778635025, 0.09814880043268204, 0.8328142166137695, 0.46247926354408264, 0.07190129905939102, 0.3418641984462738, 0.14486591517925262, 0.025201991200447083, 0.042143724858760834, 0.4074908196926117, 0.1494714319705963, 0.17342594265937805, 0.908286988735199, 0.5950636863708496, 0.14296366274356842, 0.20851416885852814, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0038264640606939793, 0.023839879781007767, 0.12264026701450348, 0.02543032169342041, 0.01467527449131012, 0.22457416355609894, 0.02885078825056553, 0.18430863320827484, 0.08557040989398956, 0.016987022012472153, 0.3513573110103607, 0.04023189842700958, 0.40384334325790405, 0.4235673248767853, 0.16652488708496094, 0.08497714251279831, 0.5087416172027588, 0.4508724510669708, 0.33144411444664, 0.600685715675354, 0.523800790309906, 0.4743403494358063, 0.10964386910200119, 0.6009643077850342, 0.29714730381965637, 0.1661888062953949, 0.10026849061250687, 0.19036318361759186, 0.07889659702777863, 0.29447081685066223, 0.5917950868606567, 0.05482999235391617, 0.0994495078921318, 0.08629819005727768, NaN, NaN, NaN, NaN, NaN, NaN], [0.006266402080655098, 0.015031179413199425, 0.02853887900710106, 0.010518345981836319, 0.09044987708330154, 0.021657679229974747, 0.0031435268465429544, 0.020945381373167038, 0.004824943374842405, 0.0127853499725461, 0.04820985347032547, 0.12459135800600052, 0.5573670268058777, 0.2566193640232086, 0.05160163715481758, 0.04716389998793602, 0.6635201573371887, 0.5744545459747314, 0.33429521322250366, 0.755266010761261, 0.7800281643867493, 0.9541771411895752, 0.5776658058166504, 0.8714791536331177, 0.9158549308776855, 0.2818737030029297, 0.06938906759023666, 0.10379814356565475, 0.3064776659011841, 0.7474142909049988, 0.7715258002281189, 0.37782159447669983, 0.057383324950933456, 0.013433223590254784, 0.03400390222668648, NaN, NaN, NaN, NaN, NaN], [0.3002758324146271, 0.08866846561431885, 0.06544900685548782, 0.25531354546546936, 0.028160221874713898, 0.12210531532764435, 0.16810676455497742, 0.0764283761382103, 0.17981933057308197, 0.3050864636898041, 0.2806880474090576, 0.13050490617752075, 0.19047558307647705, 0.3216065764427185, 0.07704814523458481, 0.1486319750547409, 0.22267495095729828, 0.42902871966362, 0.07982667535543442, 0.5459871888160706, 0.9060689210891724, 0.8350642919540405, 0.10920917987823486, 0.4773065447807312, 0.7826967239379883, 0.5733710527420044, 0.26356616616249084, 0.040332335978746414, 0.031653065234422684, 0.8572309613227844, 0.5636150240898132, 0.07464684545993805, 0.03465104475617409, 0.03009859099984169, 0.008700854144990444, 0.005375253036618233, NaN, NaN, NaN, NaN], [0.005926316604018211, 0.0003559965989552438, 0.0015365411527454853, 0.005924532189965248, 0.0005743101937696338, 0.007415232714265585, 0.024156678467988968, 0.045611582696437836, 0.009969166480004787, 0.003380746114999056, 0.003106702584773302, 0.0003880919248331338, 4.0538176108384505e-05, 0.00014580521383322775, 0.13770556449890137, 0.25873932242393494, 0.5196211338043213, 0.3300914764404297, 0.5837901830673218, 0.4101006090641022, 0.7175306677818298, 0.6572118401527405, 0.6919461488723755, 0.6594171524047852, 0.7066829204559326, 0.46555259823799133, 0.3380126953125, 0.05317035689949989, 0.053740378469228745, 0.031323984265327454, 0.30507126450538635, 0.1422475129365921, 0.03319966048002243, 0.08714800328016281, 0.01252773217856884, 0.006611488293856382, 0.007115270011126995, NaN, NaN, NaN], [0.1617586314678192, 0.29556339979171753, 0.028325924649834633, 0.059843577444553375, 0.009868957102298737, 0.03965649753808975, 0.07811643928289413, 0.06809397041797638, 0.009963614866137505, 0.11740529537200928, 0.08369920402765274, 0.039758261293172836, 0.13982373476028442, 0.1197674348950386, 0.13220268487930298, 0.011579165235161781, 0.05381239950656891, 0.044945720583200455, 0.035533830523490906, 0.6624263525009155, 0.8997865319252014, 0.9679857492446899, 0.17051655054092407, 0.940772533416748, 0.6132625341415405, 0.01721411757171154, 0.04632151871919632, 0.010550450533628464, 0.08354383707046509, 0.12839946150779724, 0.02755529060959816, 0.44050073623657227, 0.04286862909793854, 0.01342833787202835, 0.003870438551530242, 0.026607532054185867, 0.02663758397102356, 0.005111980251967907, NaN, NaN], [0.012153265066444874, 0.16048333048820496, 0.041802890598773956, 0.00796045083552599, 0.018259191885590553, 0.10963782668113708, 0.009757153689861298, 0.07023902982473373, 0.01128031499683857, 0.030125515535473824, 0.0943576917052269, 0.02206866256892681, 0.1321137398481369, 0.19507774710655212, 0.1400403380393982, 0.13300661742687225, 0.5851269960403442, 0.20284885168075562, 0.5700805187225342, 0.7479174137115479, 0.39722636342048645, 0.004733124747872353, 0.0698152482509613, 0.6515945196151733, 0.5409151315689087, 0.25820717215538025, 0.4583084285259247, 0.6744768619537354, 0.3421478569507599, 0.9633424878120422, 0.1852269172668457, 0.04996338114142418, 0.5482219457626343, 0.296283096075058, 0.48366567492485046, 0.06441208720207214, 0.9149421453475952, 0.02780383825302124, 0.0073219588957726955, NaN], [0.005033975467085838, 0.01824766956269741, 0.015512547455728054, 0.006673634983599186, 0.005676268134266138, 0.04240407794713974, 0.023996027186512947, 0.1038113459944725, 0.02023463323712349, 0.0080516142770648, 0.052543867379426956, 0.1188565045595169, 0.05977800861001015, 0.05786403268575668, 0.13343320786952972, 0.14593175053596497, 0.2687321603298187, 0.04604685679078102, 0.30660173296928406, 0.3806478679180145, 0.38105660676956177, 0.15303322672843933, 0.014211257919669151, 0.05383581668138504, 0.20604565739631653, 0.2462100237607956, 0.5718756914138794, 0.5113963484764099, 0.21981710195541382, 0.4276719391345978, 0.5577609539031982, 0.4118191599845886, 0.31598320603370667, 0.5468451976776123, 0.4359907805919647, 0.2059280127286911, 0.3916337192058563, 0.2548142671585083, 0.2198532670736313, 0.026425611227750778]], [[0.060514166951179504, 0.09119007736444473, 0.5136731863021851, 0.024349171668291092, 0.41056114435195923, 0.043175265192985535, 0.016160618513822556, 0.12711943686008453, 0.029147693887352943, 0.01592664048075676, 0.04504424333572388, 0.03736018016934395, 0.026280265301465988, 0.042564861476421356, 0.13562467694282532, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.009338664822280407, 0.09596994519233704, 0.12376897037029266, 0.01794583536684513, 0.059337858110666275, 0.04990454390645027, 0.003890786785632372, 0.07171432673931122, 0.0057785604149103165, 0.005389686673879623, 0.009663187898695469, 0.014342015609145164, 0.020640142261981964, 0.04060304909944534, 0.16408833861351013, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07689530402421951, 0.027863014489412308, 0.15549975633621216, 0.2693096697330475, 0.73520827293396, 0.03749871999025345, 0.3640631139278412, 0.14002074301242828, 0.16656053066253662, 0.02643253095448017, 0.0061660525389015675, 0.054253485053777695, 0.14240022003650665, 0.14975441992282867, 0.13701564073562622, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.21953634917736053, 0.22122228145599365, 0.04846278205513954, 0.07968296110630035, 0.3619323670864105, 0.03181222453713417, 0.6669740080833435, 0.3975786566734314, 0.11174946278333664, 0.15518029034137726, 0.004886193200945854, 0.010736972093582153, 0.07725195586681366, 0.09191425889730453, 0.1523013859987259, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0740056112408638, 0.054083533585071564, 0.027193741872906685, 0.014972379431128502, 0.04523617774248123, 0.012482533231377602, 0.4212614595890045, 0.25695085525512695, 0.3699147403240204, 0.013461914844810963, 0.08041262626647949, 0.015268572606146336, 0.627507209777832, 0.13811761140823364, 0.19850368797779083, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.029503263533115387, 0.09333665668964386, 0.016309864819049835, 0.1364656686782837, 0.03873518481850624, 0.019083604216575623, 0.758955180644989, 0.6250144243240356, 0.10551930963993073, 0.0059091635048389435, 0.001959211425855756, 0.004587537609040737, 0.0029548059683293104, 0.011073557659983635, 0.10497581213712692, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0038599083200097084, 0.03815716505050659, 0.004112291149795055, 0.0037336996756494045, 0.02896580658853054, 0.003606554586440325, 0.2724342346191406, 0.5795999765396118, 0.041377726942300797, 0.01812332309782505, 0.006642999593168497, 0.006629596464335918, 0.018780261278152466, 0.00801254715770483, 0.11063171178102493, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.023342538625001907, 0.1589166522026062, 0.01254882663488388, 0.01894153468310833, 0.04743911698460579, 0.015340029262006283, 0.06989605724811554, 0.22605817019939423, 0.016811540350317955, 0.014681086875498295, 0.0061398339457809925, 0.02630683407187462, 0.032653048634529114, 0.05358496680855751, 0.18197578191757202, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01728241890668869, 0.12100599706172943, 0.003952578641474247, 0.038103699684143066, 0.00803869217634201, 0.017839567735791206, 0.040644098073244095, 0.014622771181166172, 0.07288665324449539, 0.4550913870334625, 0.18886235356330872, 0.2150641530752182, 0.487347275018692, 0.42817094922065735, 0.12942945957183838, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.011775199323892593, 0.1349712610244751, 0.005470172502100468, 0.003098055487498641, 0.028361253440380096, 0.03303566575050354, 0.007174484897404909, 0.015601159073412418, 0.006606224924325943, 0.08859884738922119, 0.18040567636489868, 0.31761303544044495, 0.2462366670370102, 0.4818485677242279, 0.12394269555807114, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05270439758896828, 0.1637289971113205, 0.009510326199233532, 0.008013473823666573, 0.14090411365032196, 0.011389089748263359, 0.013123652897775173, 0.023534703999757767, 0.009078129194676876, 0.02855684608221054, 0.026650836691260338, 0.39132389426231384, 0.16291603446006775, 0.25967708230018616, 0.10212607681751251, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.19571052491664886, 0.10246216505765915, 0.02142595686018467, 0.012254489585757256, 0.00365867605432868, 0.007110960781574249, 0.020346596837043762, 0.03192196041345596, 0.00833944883197546, 0.07423693686723709, 0.09786227345466614, 0.08075869083404541, 0.1330210417509079, 0.26891645789146423, 0.17930860817432404, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.11616674810647964, 0.175978422164917, 0.00425378605723381, 0.017427049577236176, 0.011484457179903984, 0.030517226085066795, 0.08637198060750961, 0.1500588357448578, 0.0009573447750881314, 0.044167183339595795, 0.005869577638804913, 0.0011607500491663814, 0.014711305499076843, 0.027834221720695496, 0.18594378232955933, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.11675343662500381, 0.17556257545948029, 0.016423039138317108, 0.02097608894109726, 0.06606884300708771, 0.06371303647756577, 0.09760221093893051, 0.2481643557548523, 0.0015754855703562498, 0.03009907715022564, 0.03618617355823517, 0.012020162306725979, 0.17486301064491272, 0.22630257904529572, 0.2108311653137207, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.004961065016686916, 0.011551961302757263, 0.006318831816315651, 0.002851473866030574, 0.003461753251031041, 0.011111320927739143, 0.004611799493432045, 0.004697122145444155, 0.0026004482060670853, 0.0010426584631204605, 0.0060967751778662205, 0.01239971723407507, 0.004622939508408308, 0.002610035240650177, 0.15716104209423065, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1022859737277031, 0.17571765184402466, 0.1416551172733307, 0.11749783158302307, 0.09062699973583221, 0.07838433235883713, 0.09344526380300522, 0.3238999545574188, 0.11371968686580658, 0.10100032389163971, 0.09302259236574173, 0.0389624647796154, 0.16697892546653748, 0.1419355273246765, 0.1285012662410736, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24028724431991577, 0.14351274073123932, 0.051798444241285324, 0.16382630169391632, 0.04226303845643997, 0.020662518218159676, 0.11527843773365021, 0.29321926832199097, 0.02218940667808056, 0.0878078043460846, 0.10535410046577454, 0.011972848325967789, 0.07032275199890137, 0.04715458303689957, 0.0739566907286644, 0.1684475541114807, 0.01643766649067402, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2799055874347687, 0.11053244769573212, 0.1936434954404831, 0.029654914513230324, 0.3583168685436249, 0.552708625793457, 0.34459343552589417, 0.33612802624702454, 0.17023301124572754, 0.19969996809959412, 0.18768110871315002, 0.6793866157531738, 0.791401207447052, 0.7463385462760925, 0.09094473719596863, 0.20323613286018372, 0.02236698381602764, 0.0030780781526118517, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1572730988264084, 0.12077052146196365, 0.0489557608962059, 0.1575693041086197, 0.05669395253062248, 0.21311312913894653, 0.07387427985668182, 0.12006285786628723, 0.06427917629480362, 0.05486075580120087, 0.09722346067428589, 0.0672946497797966, 0.519307017326355, 0.15919242799282074, 0.07895061373710632, 0.15523119270801544, 0.029148569330573082, 0.04869325831532478, 0.027081435546278954, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.056666091084480286, 0.13304737210273743, 0.023897293955087662, 0.04679059237241745, 0.045941345393657684, 0.32384783029556274, 0.44531556963920593, 0.533463716506958, 0.08588721603155136, 0.10118058323860168, 0.027683693915605545, 0.15270595252513885, 0.45412689447402954, 0.19033603370189667, 0.009601723402738571, 0.20906439423561096, 0.016835892572999, 0.005647255107760429, 0.004844226874411106, 0.00019458922906778753, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.026866083964705467, 0.01856745034456253, 0.00889106560498476, 0.023431263864040375, 0.014423922635614872, 0.06721587479114532, 0.30465173721313477, 0.5084072351455688, 0.06748852878808975, 0.09416066110134125, 0.028160765767097473, 0.08301042765378952, 0.13479003310203552, 0.08470122516155243, 0.14269311726093292, 0.19736447930335999, 0.01826038584113121, 0.012854915112257004, 0.09684289991855621, 0.0006958578014746308, 4.3345058656996116e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07283831387758255, 0.02513016201555729, 0.513066828250885, 0.1692790985107422, 0.12089971452951431, 0.05420007184147835, 0.019427694380283356, 0.038392528891563416, 0.31973040103912354, 0.29048243165016174, 0.4046151340007782, 0.10607112944126129, 0.0885496586561203, 0.07017665356397629, 0.1372956782579422, 0.16369424760341644, 0.023256592452526093, 0.01855486072599888, 0.06154748797416687, 0.06098903343081474, 0.10795246064662933, 0.023746412247419357, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.27857187390327454, 0.3617483973503113, 0.2938012182712555, 0.22770966589450836, 0.06824903935194016, 0.055705904960632324, 0.2735913395881653, 0.10727421194314957, 0.15245027840137482, 0.12983311712741852, 0.2781352400779724, 0.010307536460459232, 0.09433942288160324, 0.07780664414167404, 0.13000918924808502, 0.19143380224704742, 0.11398851871490479, 0.03716170787811279, 0.07628969103097916, 0.38886839151382446, 0.24263328313827515, 0.13712459802627563, 0.02201412245631218, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09918209165334702, 0.053455647081136703, 0.645177960395813, 0.40746453404426575, 0.08205579966306686, 0.11053493618965149, 0.09200509637594223, 0.0519426129758358, 0.15867555141448975, 0.14363400638103485, 0.08945868164300919, 0.009240956045687199, 0.05626320466399193, 0.024817338213324547, 0.10628006607294083, 0.2130274772644043, 0.007986752316355705, 0.02235114760696888, 0.0019427334191277623, 0.005593507084995508, 0.012699572369456291, 0.006745419930666685, 0.06126464158296585, 0.14077326655387878, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21029417216777802, 0.16975507140159607, 0.4791514277458191, 0.5080997347831726, 0.14877668023109436, 0.04306463524699211, 0.02225780300796032, 0.027854960411787033, 0.09907854348421097, 0.17716829478740692, 0.027767561376094818, 0.04010230675339699, 0.1045137569308281, 0.07445494085550308, 0.1349247545003891, 0.22579564154148102, 0.013292824849486351, 0.10215212404727936, 0.005943832919001579, 0.013894540257751942, 0.01404587086290121, 0.02319374494254589, 0.10344905406236649, 0.1325504034757614, 0.008661924861371517, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05318222567439079, 0.11344952136278152, 0.09562063962221146, 0.10165436565876007, 0.11442670226097107, 0.07387696951627731, 0.04448265954852104, 0.12469986081123352, 0.10296554863452911, 0.029610879719257355, 0.006854650564491749, 0.06481806933879852, 0.038151390850543976, 0.029200172051787376, 0.19021393358707428, 0.1733061671257019, 0.07715445756912231, 0.2302267998456955, 0.05804288014769554, 0.07560069113969803, 0.23177897930145264, 0.2901765704154968, 0.042333029210567474, 0.08450006693601608, 0.04456959664821625, 0.015471314080059528, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.024841444566845894, 0.16249340772628784, 0.20643305778503418, 0.09402812272310257, 0.0850510448217392, 0.023708872497081757, 0.027868179604411125, 0.16653721034526825, 0.2575382590293884, 0.07176022976636887, 0.04638299718499184, 0.019721999764442444, 0.08340867608785629, 0.04306621477007866, 0.19255293905735016, 0.16428759694099426, 0.01361166127026081, 0.2167942076921463, 0.03707392141222954, 0.09917350113391876, 0.2872558534145355, 0.08793877810239792, 0.03127053380012512, 0.051127880811691284, 0.02603980340063572, 0.12251178920269012, 0.06466985493898392, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24242781102657318, 0.4547469913959503, 0.7904132008552551, 0.7443370819091797, 0.4808639585971832, 0.2640213668346405, 0.06001711264252663, 0.24681034684181213, 0.5675581097602844, 0.2725449204444885, 0.247804656624794, 0.029579274356365204, 0.19247104227542877, 0.09198179841041565, 0.18542104959487915, 0.2214493751525879, 0.0034381633158773184, 0.025536755099892616, 0.005642351228743792, 0.0024517737329006195, 0.00733930105343461, 0.0003064426709897816, 0.024970028549432755, 0.0009503457695245743, 0.0013023557839915156, 0.012362079694867134, 0.002213133964687586, 0.0037243058905005455, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10456986725330353, 0.23679938912391663, 0.29603201150894165, 0.2020668387413025, 0.14429134130477905, 0.4285147190093994, 0.3221139907836914, 0.592944860458374, 0.47945162653923035, 0.273953914642334, 0.2270997315645218, 0.05125115066766739, 0.15167200565338135, 0.14498752355575562, 0.03565559163689613, 0.21803884208202362, 0.044672977179288864, 0.15033316612243652, 0.24480289220809937, 0.0010314357932657003, 0.006885815411806107, 0.017953861504793167, 0.09280995279550552, 0.09214792400598526, 0.01309943851083517, 0.026278402656316757, 0.029330603778362274, 0.10137840360403061, 0.0009828503243625164, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.005393329542130232, 0.004602347034960985, 0.02125353366136551, 0.017772456631064415, 0.029431374743580818, 0.06670433282852173, 0.07382840663194656, 0.05640842020511627, 0.2022721767425537, 0.02110537886619568, 0.006757265422493219, 0.0065305884927511215, 0.00012849831546191126, 0.0015581984771415591, 0.14312443137168884, 0.28474918007850647, 0.005827821791172028, 0.0010850036051124334, 0.005180059466511011, 0.00018831032502930611, 0.002925402717664838, 0.0029562395066022873, 0.005281978752464056, 0.002952893264591694, 0.013548285700380802, 0.01663871854543686, 0.02234998345375061, 0.001472283387556672, 0.00024227210087701678, 9.911999950418249e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03693488612771034, 0.3099628686904907, 0.02452116832137108, 0.038606833666563034, 0.04603191837668419, 0.056979674845933914, 0.014461892656981945, 0.021202413365244865, 0.4372372031211853, 0.02073492854833603, 0.005594322457909584, 0.11605570465326309, 0.05724794790148735, 0.01605997234582901, 0.1753198802471161, 0.11472342163324356, 0.017006950452923775, 0.03429265320301056, 0.05351921543478966, 0.010289198718965054, 0.02545105293393135, 0.002036151010543108, 0.08590202778577805, 0.007977829314768314, 0.008050770498812199, 0.02079172432422638, 0.07815419882535934, 0.25072064995765686, 0.11726108938455582, 0.04080193489789963, 0.020839283242821693, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17487157881259918, 0.2829012870788574, 0.22657853364944458, 0.2227388322353363, 0.09278897941112518, 0.05522100254893303, 0.023270972073078156, 0.031554628163576126, 0.32194823026657104, 0.13948096334934235, 0.09803083539009094, 0.2809208631515503, 0.14969345927238464, 0.03018103539943695, 0.10283161699771881, 0.25351014733314514, 0.018978603184223175, 0.013279697857797146, 0.14657457172870636, 0.0005683518829755485, 0.003044809214770794, 0.0003673452010843903, 0.0009085922501981258, 0.00026260188315063715, 6.703466351609677e-05, 0.00393629027530551, 0.0411190427839756, 0.014572926796972752, 0.0009043514728546143, 0.001453216653317213, 0.001335341832600534, 0.0036634530406445265, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06711219251155853, 0.13971862196922302, 0.10573939234018326, 0.08062157034873962, 0.22173365950584412, 0.04757346957921982, 0.02002648264169693, 0.06195787340402603, 0.09553409367799759, 0.04351034387946129, 0.015184497460722923, 0.17841440439224243, 0.07658158242702484, 0.04646967723965645, 0.1461518555879593, 0.2249869406223297, 0.0773954764008522, 0.10561174154281616, 0.3267342746257782, 0.011780736967921257, 0.03227663040161133, 0.09185110032558441, 0.03840579837560654, 0.01289159432053566, 0.002641883445903659, 0.03386297821998596, 0.16820214688777924, 0.06345225125551224, 0.027306171134114265, 0.007737002335488796, 0.018253128975629807, 0.0508209764957428, 0.015562118031084538, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.015694430097937584, 0.09081663191318512, 0.2731003761291504, 0.09780610352754593, 0.06437630951404572, 0.024092676118016243, 0.017730340361595154, 0.09997125715017319, 0.24317535758018494, 0.06615940481424332, 0.05322461575269699, 0.013002216815948486, 0.10308460891246796, 0.03947872668504715, 0.16966252028942108, 0.17073971033096313, 0.01119090337306261, 0.07090220600366592, 0.026190776377916336, 0.04357914999127388, 0.10384812206029892, 0.05681576952338219, 0.008270802907645702, 0.011212479323148727, 0.016114890575408936, 0.1306251734495163, 0.04437248408794403, 0.022720789536833763, 0.0017881430685520172, 0.005742507986724377, 0.03271590173244476, 0.12170897424221039, 0.18442584574222565, 0.07238933444023132, NaN, NaN, NaN, NaN, NaN, NaN], [0.19514591991901398, 0.2590837776660919, 0.7111572027206421, 0.6245842576026917, 0.2279123067855835, 0.21324849128723145, 0.0465325303375721, 0.16129039227962494, 0.5552195906639099, 0.24888396263122559, 0.16995932161808014, 0.017819084227085114, 0.13601525127887726, 0.04923256114125252, 0.1924036145210266, 0.2460513859987259, 0.004599481821060181, 0.030415518209338188, 0.006707339081913233, 0.001940727117471397, 0.0018293699249625206, 0.0002438600640743971, 0.021702459082007408, 0.00019114103633910418, 0.0004616644873749465, 0.02795419655740261, 0.007376548834145069, 0.009364028461277485, 0.0008695388678461313, 0.027626920491456985, 0.002984545426443219, 0.0021758046932518482, 0.005276597570627928, 0.0015223525697365403, 0.0046029179356992245, NaN, NaN, NaN, NaN, NaN], [0.11466818302869797, 0.23749157786369324, 0.22078867256641388, 0.21260471642017365, 0.1054922342300415, 0.38443663716316223, 0.35735341906547546, 0.3432110548019409, 0.45766645669937134, 0.30316272377967834, 0.15794025361537933, 0.23222389817237854, 0.18522031605243683, 0.12369272857904434, 0.062224190682172775, 0.1682240217924118, 0.15532228350639343, 0.17499232292175293, 0.31528380513191223, 0.0016938054468482733, 0.0013859918108209968, 0.0071086762472987175, 0.08609996736049652, 0.02145048975944519, 0.00334079097956419, 0.08546027541160583, 0.16909679770469666, 0.5000762343406677, 0.012536582536995411, 0.0033327846322208643, 0.01681024581193924, 0.01291667390614748, 0.11205089092254639, 0.06917328387498856, 0.24062496423721313, 0.003104837378486991, NaN, NaN, NaN, NaN], [0.004928229842334986, 0.004764902405440807, 0.014567935839295387, 0.014073353260755539, 0.020878629758954048, 0.04901519790291786, 0.05124438554048538, 0.042454566806554794, 0.19801755249500275, 0.018003307282924652, 0.004736864008009434, 0.006620202213525772, 0.00011398878996260464, 0.001381832524202764, 0.13761556148529053, 0.30163663625717163, 0.008585775271058083, 0.0018221536884084344, 0.004949942696839571, 0.0002661931503098458, 0.0017199779395014048, 0.00286088977009058, 0.004591777920722961, 0.0013412131229415536, 0.009152509272098541, 0.029603971168398857, 0.059182800352573395, 0.004352512303739786, 0.0009281163802370429, 0.00013420419418253005, 0.0015637356555089355, 0.004895435180515051, 0.0020298720337450504, 0.016267914324998856, 0.0014363413210958242, 0.00015049855574034154, 4.989441003999673e-05, NaN, NaN, NaN], [0.013776288367807865, 0.25124475359916687, 0.00789756141602993, 0.00910337083041668, 0.005072988104075193, 0.015830766409635544, 0.005818341393023729, 0.011153762228786945, 0.14152461290359497, 0.008211367763578892, 0.002360414480790496, 0.06666377186775208, 0.057822320610284805, 0.009000283665955067, 0.13980405032634735, 0.1420876681804657, 0.030559053644537926, 0.035777460783720016, 0.0549585185945034, 0.010907668620347977, 0.018195953220129013, 0.005288956221193075, 0.07946551591157913, 0.003352995030581951, 0.00945360492914915, 0.03057919070124626, 0.20277532935142517, 0.5438944697380066, 0.2487112432718277, 0.11027072370052338, 0.03672702983021736, 0.009589559398591518, 0.03681262582540512, 0.12653782963752747, 0.3100517988204956, 0.04488144814968109, 0.07299992442131042, 0.024292031303048134, NaN, NaN], [0.25532495975494385, 0.3110601603984833, 0.28066542744636536, 0.29941898584365845, 0.09561395645141602, 0.06004221364855766, 0.0257351566106081, 0.04446575790643692, 0.3475395441055298, 0.2538500130176544, 0.25107017159461975, 0.4736424386501312, 0.29699820280075073, 0.06975124776363373, 0.11745814979076385, 0.2571920156478882, 0.012253361754119396, 0.00982633139938116, 0.09085621684789658, 0.00026428516139276326, 0.001174133620224893, 0.00010905979434028268, 0.0006958161829970777, 9.435929678147659e-05, 1.889842314994894e-05, 0.0019355103140696883, 0.03233037516474724, 0.014144179411232471, 0.0034062752965837717, 0.0014896523207426071, 0.0032966958824545145, 0.0043079969473183155, 0.002425077836960554, 0.0237245112657547, 0.017915409058332443, 0.0004631538176909089, 0.0033925946336239576, 0.0019653798080980778, 0.0010656031081452966, NaN], [0.06876020133495331, 0.07319146394729614, 0.08357107639312744, 0.06905727088451385, 0.010884120129048824, 0.012632370926439762, 0.04344229772686958, 0.06033884361386299, 0.05559740215539932, 0.048808641731739044, 0.06204793229699135, 0.017201891168951988, 0.028970519080758095, 0.021960163488984108, 0.13179059326648712, 0.25252944231033325, 0.012149164453148842, 0.019892947748303413, 0.013666713610291481, 0.05940697342157364, 0.04882493242621422, 0.025430571287870407, 0.00045668394886888564, 0.0054928152821958065, 0.005623141769319773, 0.004253733437508345, 0.014798035845160484, 0.012909402139484882, 0.011927488259971142, 0.007018915377557278, 0.021986471489071846, 0.016502689570188522, 0.002887164242565632, 0.006932961288839579, 0.007926056161522865, 0.015145027078688145, 0.005945136770606041, 0.016453862190246582, 0.011257275938987732, 0.0009747393196448684]], [[0.027552247047424316, 0.013821233063936234, 0.004237555433064699, 0.0007387229125015438, 0.0009859473211690784, 0.001997306477278471, 0.002160864183679223, 0.009250090457499027, 0.0009738927474245429, 0.0009403586154803634, 0.003406830132007599, 0.0010056114988401532, 0.008306043222546577, 0.06191018968820572, 0.18169914186000824, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0056476471945643425, 0.0617278628051281, 0.026225095614790916, 0.009516767226159573, 0.019543437287211418, 0.011766157113015652, 0.0015307252760976553, 0.004000868182629347, 0.006223553325980902, 0.02180931344628334, 0.02397397719323635, 0.025289250537753105, 0.01872297003865242, 0.05591608211398125, 0.17309869825839996, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5742589831352234, 0.02769068442285061, 0.03131784498691559, 0.008496972732245922, 0.005279624368995428, 0.0009009581408463418, 0.013010378926992416, 0.009255914948880672, 0.08095329999923706, 0.0017015798948705196, 0.0027918636333197355, 0.01474103331565857, 0.07241056859493256, 0.2960302531719208, 0.1991364061832428, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3870091140270233, 0.24428580701351166, 0.004871743265539408, 0.01251932606101036, 0.004600874613970518, 0.007045479491353035, 0.011942178010940552, 0.06100638955831528, 0.06223933771252632, 0.00421120086684823, 0.0017708303639665246, 0.010406754910945892, 0.016386834904551506, 0.038040366023778915, 0.25559180974960327, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6136646866798401, 0.2692064642906189, 0.043582458049058914, 0.00652115186676383, 0.05291604623198509, 0.006654517259448767, 0.03398957848548889, 0.03886384516954422, 0.13169772922992706, 0.002106831641867757, 0.005907678045332432, 0.01888049766421318, 0.04876947030425072, 0.2226717472076416, 0.22327177226543427, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.685612678527832, 0.0861489400267601, 0.03236214071512222, 0.16196951270103455, 0.03394145518541336, 0.05551951378583908, 0.027528556063771248, 0.06770895421504974, 0.19389298558235168, 0.03780713677406311, 0.0038191182538866997, 0.05989958345890045, 0.13479465246200562, 0.24111053347587585, 0.15613426268100739, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6876600384712219, 0.0606975182890892, 0.05783677101135254, 0.05387236177921295, 0.11914167553186417, 0.004756046459078789, 0.031782086938619614, 0.011465699411928654, 0.1448838710784912, 0.09538520872592926, 0.007872258313000202, 0.033316925168037415, 0.09786565601825714, 0.08940181881189346, 0.23629719018936157, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5363585352897644, 0.11579979956150055, 0.10718797892332077, 0.21453110873699188, 0.030864767730236053, 0.026318436488509178, 0.03807519003748894, 0.12262200564146042, 0.08015674352645874, 0.06537020206451416, 0.004594390746206045, 0.015254726633429527, 0.06485987454652786, 0.039039257913827896, 0.16586215794086456, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6220377087593079, 0.17304541170597076, 0.23731492459774017, 0.32412996888160706, 0.2203587144613266, 0.09306959062814713, 0.2822628319263458, 0.008407875895500183, 0.14113475382328033, 0.022416740655899048, 0.005183607805520296, 0.0005837879725731909, 0.00799399521201849, 0.006284625735133886, 0.12005029618740082, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.18509520590305328, 0.21334251761436462, 0.12845394015312195, 0.3693835139274597, 0.41559898853302, 0.19613976776599884, 0.7053389549255371, 0.3886314332485199, 0.06599769741296768, 0.04325481504201889, 0.029052795842289925, 0.001557054347358644, 0.0018087843200191855, 0.0036887156311422586, 0.18107539415359497, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.612794041633606, 0.24153079092502594, 0.076973557472229, 0.17341682314872742, 0.06242084503173828, 0.2242424041032791, 0.8304246068000793, 0.5655775666236877, 0.4262824058532715, 0.00936043355613947, 0.03881426528096199, 0.0046007027849555016, 0.005786797031760216, 0.020520325750112534, 0.226027712225914, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.21637925505638123, 0.22487440705299377, 0.19202512502670288, 0.3957260847091675, 0.15970049798488617, 0.16693006455898285, 0.3690066933631897, 0.5193001627922058, 0.6459834575653076, 0.047006867825984955, 0.06868032366037369, 0.043628890067338943, 0.02405296452343464, 0.05333276465535164, 0.08607933670282364, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5923737287521362, 0.3536633849143982, 0.08390633016824722, 0.2980528473854065, 0.042989592999219894, 0.026934657245874405, 0.1647067815065384, 0.1620720773935318, 0.6647022366523743, 0.13678880035877228, 0.10115252435207367, 0.012052871286869049, 0.2444845736026764, 0.1799331158399582, 0.10357851535081863, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3260110914707184, 0.10825559496879578, 0.040669191628694534, 0.08903322368860245, 0.055108752101659775, 0.014200238510966301, 0.06877616047859192, 0.07561883330345154, 0.7116665244102478, 0.08518233895301819, 0.13964912295341492, 0.01787719503045082, 0.027594367042183876, 0.0709126889705658, 0.09409899264574051, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.26070404052734375, 0.8011303544044495, 0.17980173230171204, 0.0725909024477005, 0.12434736639261246, 0.28980228304862976, 0.3281027674674988, 0.7843722701072693, 0.12677432596683502, 0.054726697504520416, 0.13370326161384583, 0.19018130004405975, 0.1707623451948166, 0.14939220249652863, 0.07447532564401627, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1855485588312149, 0.4779467284679413, 0.0886944904923439, 0.027812138199806213, 0.051930978894233704, 0.20570456981658936, 0.13285183906555176, 0.12479114532470703, 0.03275279700756073, 0.13280591368675232, 0.10831113904714584, 0.13358037173748016, 0.31709861755371094, 0.18639257550239563, 0.0658930093050003, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04738391190767288, 0.17884546518325806, 0.030679181218147278, 0.09374479204416275, 0.015219364315271378, 0.004209337756037712, 0.011544613167643547, 0.014519347809255123, 0.0008998611010611057, 0.03714418038725853, 0.02808041125535965, 0.0015275280456990004, 0.014074422419071198, 0.01773718185722828, 0.02865048497915268, 0.14568212628364563, 0.073321633040905, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4282352328300476, 0.07421883940696716, 0.37614062428474426, 0.6016114950180054, 0.16448479890823364, 0.10949403792619705, 0.43647968769073486, 0.17394804954528809, 0.2346193641424179, 0.5131813287734985, 0.6543169021606445, 0.06318124383687973, 0.059741634875535965, 0.08049911260604858, 0.08155221492052078, 0.07740449905395508, 0.019538799300789833, 0.31676185131073, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04248558357357979, 0.005498564336448908, 0.015051363967359066, 0.021896474063396454, 0.031015703454613686, 0.23631463944911957, 0.5231030583381653, 0.1651564985513687, 0.010708797723054886, 0.0702022984623909, 0.015817642211914062, 0.01968570239841938, 0.2309122085571289, 0.11954572051763535, 0.04909561946988106, 0.11254165321588516, 0.04977253079414368, 0.12113941460847855, 0.18998825550079346, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.019823409616947174, 0.02119731903076172, 0.0447932668030262, 0.04950243979692459, 0.11350910365581512, 0.3172611892223358, 0.1175147220492363, 0.16474604606628418, 0.025614900514483452, 0.11684545129537582, 0.027774598449468613, 0.03366768732666969, 0.1657668650150299, 0.20241110026836395, 0.02058284729719162, 0.09693466126918793, 0.12094055861234665, 0.48810020089149475, 0.07605772465467453, 0.10663138329982758, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.024027986451983452, 0.07085671275854111, 0.014559593982994556, 0.003951122052967548, 0.5812088251113892, 0.07389754801988602, 0.10464153438806534, 0.06822511553764343, 0.1849648803472519, 0.02429678477346897, 0.014226456172764301, 0.2123226672410965, 0.1049809455871582, 0.17609325051307678, 0.13661964237689972, 0.002718105213716626, 0.037000641226768494, 0.1506986916065216, 0.012303436174988747, 0.09212689101696014, 0.5217995047569275, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20496347546577454, 0.09403666108846664, 0.02112487144768238, 0.025338320061564445, 0.008130905218422413, 0.1783977895975113, 0.3754851818084717, 0.0950397253036499, 0.0030220954213291407, 0.08205359429121017, 0.011042395606637001, 0.018588367849588394, 0.1888807862997055, 0.10302136838436127, 0.14473272860050201, 0.17887507379055023, 0.10589989274740219, 0.004075651057064533, 0.0014342612121254206, 0.00521382549777627, 0.031908128410577774, 0.003124895039945841, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.037373751401901245, 0.07382072508335114, 0.08205787092447281, 0.10832883417606354, 0.02859049290418625, 0.1663966327905655, 0.058918725699186325, 0.17053310573101044, 0.011018002405762672, 0.15213745832443237, 0.027154715731739998, 0.0019660431426018476, 0.22162862122058868, 0.11411792784929276, 0.08493959158658981, 0.23519471287727356, 0.3653021454811096, 0.05512593686580658, 0.10675911605358124, 0.0014886436983942986, 0.001230676076374948, 0.003634560154750943, 0.00975269265472889, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.015705576166510582, 0.016172299161553383, 0.006149389781057835, 0.0038101596292108297, 0.007736767642199993, 0.20371977984905243, 0.12438680231571198, 0.06649734079837799, 0.004926482681185007, 0.004153827205300331, 0.0012289183214306831, 0.003863752353936434, 0.0550994910299778, 0.04052891582250595, 0.36571574211120605, 0.19171930849552155, 0.3204987347126007, 0.0060858046635985374, 0.010409774258732796, 0.003722283523529768, 0.0010954621247947216, 0.0028676562942564487, 0.35306307673454285, 0.01622932404279709, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.008730506524443626, 0.002757954876869917, 0.0122150257229805, 0.006305738352239132, 0.004681416787207127, 0.06460410356521606, 0.008150112815201283, 0.010960009880363941, 0.004299533553421497, 0.004670997615903616, 0.0034528695978224277, 0.0024545302148908377, 0.005013267509639263, 0.008545692078769207, 0.23703089356422424, 0.25555557012557983, 0.13076956570148468, 0.003832729533314705, 0.0447237528860569, 0.014599477872252464, 0.0024878191761672497, 0.0016443775966763496, 0.20187559723854065, 0.0005508072790689766, 0.0029457835480570793, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09499987959861755, 0.010673395358026028, 0.007046178914606571, 0.020993953570723534, 0.010670008137822151, 0.07466354966163635, 0.06417079269886017, 0.023990478366613388, 0.17728924751281738, 0.15624059736728668, 0.004560643341392279, 0.010690598748624325, 0.03727814555168152, 0.017693333327770233, 0.14084658026695251, 0.13948844373226166, 0.2463626265525818, 0.09502393007278442, 0.197096586227417, 0.47678983211517334, 0.3142886161804199, 0.09103813022375107, 0.10499368607997894, 0.07698603719472885, 0.026083102449774742, 0.3110981583595276, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.688500165939331, 0.16286028921604156, 0.04583478718996048, 0.22473743557929993, 0.025797681882977486, 0.04771623760461807, 0.5437547564506531, 0.0642164871096611, 0.01443459838628769, 0.2519066631793976, 0.017869845032691956, 0.003991205245256424, 0.04630482196807861, 0.029587149620056152, 0.049375567585229874, 0.1511228382587433, 0.027682308107614517, 0.014322453178465366, 0.0030328254215419292, 0.04723867028951645, 0.30981165170669556, 0.025852922350168228, 0.018514074385166168, 0.01515920553356409, 0.009253463707864285, 0.10175863653421402, 0.16996310651302338, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14772717654705048, 0.11627800017595291, 0.034884992986917496, 0.02596234902739525, 0.031621210277080536, 0.39286479353904724, 0.6627658009529114, 0.20747745037078857, 0.019052494317293167, 0.06071586161851883, 0.014515946619212627, 0.03545556217432022, 0.1622975915670395, 0.05619712546467781, 0.4560142755508423, 0.1847103387117386, 0.05052594095468521, 0.005765186157077551, 0.018545929342508316, 0.00881477165967226, 0.0375242680311203, 0.027162199839949608, 0.09025334566831589, 0.0028228689916431904, 0.0033718899358063936, 0.1103500947356224, 0.0837099552154541, 0.0044236015528440475, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3253695070743561, 0.18678773939609528, 0.23196454346179962, 0.43925735354423523, 0.09974130243062973, 0.1577768325805664, 0.26045241951942444, 0.07323815673589706, 0.005399893503636122, 0.23951157927513123, 0.04431937262415886, 0.013187061063945293, 0.0749824121594429, 0.025474021211266518, 0.2768867611885071, 0.27341794967651367, 0.03427007421851158, 0.008004172705113888, 0.009254892356693745, 0.005621441174298525, 0.00972525030374527, 0.005248658824712038, 0.02184745855629444, 0.0006181569187901914, 0.0005494534852914512, 0.06994801014661789, 0.02213645726442337, 0.004287416115403175, 0.0008399627404287457, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.049311667680740356, 0.10222040861845016, 0.30249276757240295, 0.11109475791454315, 0.4333159327507019, 0.4476950168609619, 0.14919614791870117, 0.45436185598373413, 0.10977044701576233, 0.101465605199337, 0.28612539172172546, 0.15904487669467926, 0.4858849048614502, 0.19411928951740265, 0.08273273706436157, 0.008804291486740112, 0.07617928832769394, 0.47516930103302, 0.07513945549726486, 0.5241973400115967, 0.4384346902370453, 0.06213618069887161, 0.06345370411872864, 0.0682281106710434, 0.15877418220043182, 0.023486817255616188, 0.026526909321546555, 0.0028373831883072853, 0.001617963775061071, 0.37629759311676025, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08865676820278168, 0.0832996591925621, 0.0360012948513031, 0.026901112869381905, 0.0488949753344059, 0.5697077512741089, 0.2118675261735916, 0.21166029572486877, 0.009457184933125973, 0.042189937084913254, 0.010147118009626865, 0.027016732841730118, 0.1966082751750946, 0.18848717212677002, 0.17412608861923218, 0.26533833146095276, 0.10994716733694077, 0.010266831144690514, 0.037150826305150986, 0.009969023987650871, 0.00030588259687647223, 8.988264016807079e-05, 0.07940464466810226, 0.00027601365582086146, 0.0013282618019729853, 0.009904097765684128, 0.03278518095612526, 0.0630892813205719, 0.10911130160093307, 0.016624033451080322, 0.011541539803147316, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09455566853284836, 0.047932155430316925, 0.06032469496130943, 0.027359262108802795, 0.004525639116764069, 0.19231697916984558, 0.29536089301109314, 0.10446369647979736, 0.004957688972353935, 0.22148354351520538, 0.017980555072426796, 0.016062501817941666, 0.01227590162307024, 0.007468203082680702, 0.14047065377235413, 0.2451263964176178, 0.014867580495774746, 0.0005470102187246084, 0.0054298522882163525, 0.0004450916312634945, 0.0006575370789505541, 3.8741818570997566e-05, 0.0010275153908878565, 0.0013172366889193654, 0.0019110681023448706, 0.13600468635559082, 0.29138538241386414, 0.011091821826994419, 0.0002334356977371499, 0.0002162840828532353, 0.0001727231137920171, 0.004782650154083967, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18475790321826935, 0.03305341675877571, 0.022945405915379524, 0.02499788999557495, 0.016275716945528984, 0.44049808382987976, 0.3255404233932495, 0.03656867519021034, 0.008760510943830013, 0.28132569789886475, 0.00872495025396347, 0.02103549800813198, 0.09103824943304062, 0.045535117387771606, 0.1431308537721634, 0.18341027200222015, 0.31211209297180176, 0.08544175326824188, 0.17215219140052795, 0.07786234468221664, 0.033002957701683044, 0.028957894071936607, 0.08467604964971542, 0.018818018957972527, 0.0016417433507740498, 0.15075404942035675, 0.1522863805294037, 0.03350237384438515, 0.006119633559137583, 0.022573737427592278, 0.03810621052980423, 0.13675758242607117, 0.1992093175649643, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5226730704307556, 0.08511564135551453, 0.13128292560577393, 0.22977954149246216, 0.025636736303567886, 0.14430683851242065, 0.697600245475769, 0.08303582668304443, 0.03326253592967987, 0.30183717608451843, 0.04944504052400589, 0.004384536296129227, 0.07144975662231445, 0.05258011445403099, 0.06879302859306335, 0.1540856957435608, 0.05453011393547058, 0.023697303608059883, 0.003979950677603483, 0.014029269106686115, 0.1104540005326271, 0.019629694521427155, 0.011429534293711185, 0.010672842152416706, 0.00807265006005764, 0.1843080371618271, 0.19234825670719147, 0.0017768212128430605, 0.006891301833093166, 0.08265318721532822, 0.014878016896545887, 0.09550431370735168, 0.1691773235797882, 0.20674942433834076, NaN, NaN, NaN, NaN, NaN, NaN], [0.06703877449035645, 0.049393996596336365, 0.041539933532476425, 0.021373772993683815, 0.02868128940463066, 0.32991066575050354, 0.488584041595459, 0.0702073872089386, 0.0075523643754422665, 0.038572411984205246, 0.012813442386686802, 0.04136957228183746, 0.06929102540016174, 0.03757195174694061, 0.23515936732292175, 0.21139073371887207, 0.06409671157598495, 0.007977590896189213, 0.017582383006811142, 0.004139575641602278, 0.008497070521116257, 0.024324562400579453, 0.12332659959793091, 0.0006915424601174891, 0.0006991134723648429, 0.09821731597185135, 0.18821127712726593, 0.009975801222026348, 0.024784373119473457, 0.009686794131994247, 0.0016004297649487853, 0.006526788230985403, 0.04246864095330238, 0.05479469522833824, 0.004482009913772345, NaN, NaN, NaN, NaN, NaN], [0.15618596971035004, 0.12941822409629822, 0.2654253840446472, 0.28590527176856995, 0.31243884563446045, 0.1085575670003891, 0.15852880477905273, 0.026613548398017883, 0.004155577160418034, 0.15324708819389343, 0.037679530680179596, 0.09416285902261734, 0.02134908176958561, 0.010629331693053246, 0.17846201360225677, 0.33224669098854065, 0.07294216006994247, 0.01592269167304039, 0.006994656287133694, 0.003661615075543523, 0.0007586313877254725, 0.0006907262722961605, 0.022764746099710464, 0.000276167003903538, 9.849678463069722e-05, 0.08613532781600952, 0.07070992141962051, 0.03258151933550835, 0.002256957348436117, 0.00035050295991823077, 0.002809839555993676, 0.005992868449538946, 0.14088936150074005, 0.024111032485961914, 0.015468394383788109, 0.000736193498596549, NaN, NaN, NaN, NaN], [0.058257974684238434, 0.12017454952001572, 0.32657214999198914, 0.12284700572490692, 0.5568311810493469, 0.41536086797714233, 0.16300946474075317, 0.49100223183631897, 0.15462136268615723, 0.11520260572433472, 0.260068416595459, 0.28476831316947937, 0.501883327960968, 0.21151991188526154, 0.09330709278583527, 0.00368693470954895, 0.0603332445025444, 0.389295369386673, 0.03955860063433647, 0.26089394092559814, 0.125760018825531, 0.029167605563998222, 0.03710402920842171, 0.03377004712820053, 0.08135493099689484, 0.01946301944553852, 0.033920928835868835, 0.00409010099247098, 0.0020981510169804096, 0.4028157889842987, 0.01821253076195717, 0.03254074230790138, 0.005954912398010492, 0.016414301469922066, 0.0033934058155864477, 0.0012025205651298165, 0.37666910886764526, NaN, NaN, NaN], [0.04007576033473015, 0.04011448100209236, 0.02015572600066662, 0.006723308004438877, 0.01584162376821041, 0.6745935082435608, 0.14270515739917755, 0.05812964215874672, 0.0018657244509086013, 0.018765496090054512, 0.004551106132566929, 0.05217724293470383, 0.21886952221393585, 0.13090433180332184, 0.13149680197238922, 0.30478137731552124, 0.23805196583271027, 0.009743728674948215, 0.02953244559466839, 0.005627358797937632, 0.00013927526015322655, 0.00016958850028458983, 0.09182754158973694, 0.00019882968626916409, 0.0018803260754793882, 0.01743759773671627, 0.09691343456506729, 0.09625609964132309, 0.0949849784374237, 0.057061683386564255, 0.028116967529058456, 0.00013736996334046125, 0.022905906662344933, 0.02515738271176815, 0.029101604595780373, 0.01233749371021986, 0.027021989226341248, 0.012159456498920918, NaN, NaN], [0.051524627953767776, 0.037071868777275085, 0.09267362952232361, 0.03285788744688034, 0.006808253470808268, 0.2584725618362427, 0.21142001450061798, 0.06556515395641327, 0.003410812932997942, 0.18829914927482605, 0.028329605236649513, 0.02864006720483303, 0.014232979156076908, 0.014326054602861404, 0.12804241478443146, 0.2508227825164795, 0.013127491809427738, 0.0004774215049110353, 0.005875048227608204, 0.00014762053615413606, 0.0003128673997707665, 1.7799626220948994e-05, 0.0017815351020544767, 0.0009225650574080646, 0.0009481729357503355, 0.09391504526138306, 0.24316561222076416, 0.008820290677249432, 0.0015348505694419146, 0.0002856143401004374, 0.00038499117363244295, 0.010248353704810143, 0.0923430323600769, 0.1539699137210846, 0.0089821582660079, 0.00013843990745954216, 0.0004539538058452308, 6.709429726470262e-05, 0.0014084051363170147, NaN], [0.13503411412239075, 0.06798373907804489, 0.08072269707918167, 0.04104887321591377, 0.027653640136122704, 0.5933560132980347, 0.15723249316215515, 0.044575583189725876, 0.017590617761015892, 0.04771400988101959, 0.07117579132318497, 0.10345834493637085, 0.10624422132968903, 0.027206260710954666, 0.1271171271800995, 0.06230561435222626, 0.051613274961709976, 0.02077883668243885, 0.04204944148659706, 0.07247611880302429, 0.11675790697336197, 0.004215644672513008, 0.00555834174156189, 0.008976897224783897, 0.017200933769345284, 0.007355507928878069, 0.06492317467927933, 0.04215962812304497, 0.02968345396220684, 0.23223130404949188, 0.03253115341067314, 0.08794146776199341, 0.025323374196887016, 0.08459514379501343, 0.05644838511943817, 0.04970480501651764, 0.3588789105415344, 0.028869707137346268, 0.11940079927444458, 0.27181047201156616]], [[0.10194799304008484, 0.042179130017757416, 0.27587375044822693, 0.8387316465377808, 0.3051532208919525, 0.225641667842865, 0.10655678808689117, 0.4426303505897522, 0.21958006918430328, 0.4376780688762665, 0.7421585917472839, 0.6036965250968933, 0.4420715570449829, 0.6119644045829773, 0.08460802584886551, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.052479684352874756, 0.018692737445235252, 0.13130725920200348, 0.4463008642196655, 0.4007475674152374, 0.4465942680835724, 0.13863760232925415, 0.26287177205085754, 0.5015351176261902, 0.48749616742134094, 0.19089040160179138, 0.2783986032009125, 0.20843097567558289, 0.11412637680768967, 0.11901978403329849, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.09998084604740143, 0.05760321766138077, 0.06884635984897614, 0.1367950737476349, 0.03696327656507492, 0.02052011340856552, 0.23966658115386963, 0.6639524102210999, 0.08913422375917435, 0.1896458864212036, 0.14239966869354248, 0.18587030470371246, 0.2512775659561157, 0.1800404042005539, 0.13985422253608704, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.17776982486248016, 0.2164098620414734, 0.03016561083495617, 0.006355184596031904, 0.04318562150001526, 0.004709928296506405, 0.02340516820549965, 0.07859960943460464, 0.3921053409576416, 0.27134451270103455, 0.2182498425245285, 0.1118401437997818, 0.13378913700580597, 0.4978374242782593, 0.18931511044502258, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.16739480197429657, 0.20097726583480835, 0.038037389516830444, 0.05488090589642525, 0.020769814029335976, 0.044557277113199234, 0.32692524790763855, 0.5529306530952454, 0.06495681405067444, 0.061963245272636414, 0.3602059483528137, 0.040287844836711884, 0.11072657257318497, 0.3166219890117645, 0.19249440729618073, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07948607206344604, 0.4389178156852722, 0.019072405993938446, 0.11389600485563278, 0.015004596672952175, 0.0008035529754124582, 0.00560334138572216, 0.007579134311527014, 0.12602436542510986, 0.4041804373264313, 0.8435949087142944, 0.7255359292030334, 0.3334953784942627, 0.21919409930706024, 0.13174442946910858, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.11827840656042099, 0.43549492955207825, 0.035650141537189484, 0.3500109016895294, 0.10479609668254852, 0.0029047641437500715, 0.016262628138065338, 0.008920608088374138, 0.1923075020313263, 0.6588289737701416, 0.7271849513053894, 0.8207041025161743, 0.5342087149620056, 0.29674431681632996, 0.16698533296585083, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.19771254062652588, 0.43774574995040894, 0.057631127536296844, 0.15638697147369385, 0.05497771501541138, 0.0015852008946239948, 0.004800108727067709, 0.0038221883587539196, 0.11230877041816711, 0.6780416369438171, 0.6535694003105164, 0.33372464776039124, 0.2617355287075043, 0.4378974735736847, 0.15096917748451233, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2510830760002136, 0.455088347196579, 0.2769528925418854, 0.28598156571388245, 0.08308438956737518, 0.495423823595047, 0.2878262400627136, 0.017540372908115387, 0.036487918347120285, 0.07030303031206131, 0.04537871107459068, 0.017587929964065552, 0.15749330818653107, 0.15622387826442719, 0.134229376912117, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2108728438615799, 0.12734071910381317, 0.6047671437263489, 0.5566261410713196, 0.4727993309497833, 0.6295000314712524, 0.20963285863399506, 0.3828260004520416, 0.01981351152062416, 0.02910005673766136, 0.17932364344596863, 0.029557999223470688, 0.02868420071899891, 0.05513756722211838, 0.1339428722858429, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2013130933046341, 0.35711804032325745, 0.18803814053535461, 0.31239861249923706, 0.6328845024108887, 0.6068195104598999, 0.09879770874977112, 0.295420378446579, 0.033300116658210754, 0.04495004564523697, 0.027333615347743034, 0.034196678549051285, 0.011724627576768398, 0.023517103865742683, 0.3543241322040558, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.27807915210723877, 0.07025524973869324, 0.15421687066555023, 0.23079168796539307, 0.0323871448636055, 0.4182601273059845, 0.43312954902648926, 0.3330070972442627, 0.027521615847945213, 0.03977188467979431, 0.03152378648519516, 0.00340716983191669, 0.005408053286373615, 0.0057552107609808445, 0.23170912265777588, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.15765754878520966, 0.07761365175247192, 0.1382310688495636, 0.33822664618492126, 0.15857987105846405, 0.11602839827537537, 0.3749851584434509, 0.3412497341632843, 0.06253337115049362, 0.09931040555238724, 0.010201470926404, 0.0010190334869548678, 0.0007929145358502865, 0.0016151106683537364, 0.1723894327878952, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.39988550543785095, 0.09145350754261017, 0.3013111352920532, 0.5813722610473633, 0.4042908251285553, 0.2935561537742615, 0.4903331696987152, 0.4357178807258606, 0.04456466808915138, 0.10430204123258591, 0.10590728372335434, 0.007762597873806953, 0.0026525144930928946, 0.0052152471616864204, 0.24974997341632843, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03366217389702797, 0.03653215244412422, 0.027766529470682144, 0.007369572762399912, 0.014929202385246754, 0.04527684673666954, 0.00940654892474413, 0.023517949506640434, 0.010960820131003857, 0.0019369145156815648, 0.01981637440621853, 0.00444602407515049, 0.014915830455720425, 0.007271313574165106, 0.15384840965270996, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04247138649225235, 0.01728098653256893, 0.06617120653390884, 0.009399485774338245, 0.0730140432715416, 0.14221039414405823, 0.11889991164207458, 0.10651882737874985, 0.10687308758497238, 0.0351867638528347, 0.09164245426654816, 0.06160420924425125, 0.04699656739830971, 0.14884592592716217, 0.20088525116443634, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.35919252038002014, 0.017007382586598396, 0.3711448311805725, 0.05260182172060013, 0.23237934708595276, 0.17189942300319672, 0.06846722215414047, 0.25480321049690247, 0.4269619286060333, 0.141769677400589, 0.19745108485221863, 0.3101239502429962, 0.12419883906841278, 0.061588384211063385, 0.3489930033683777, 0.04884753376245499, 0.31528204679489136, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1570073962211609, 0.6818748116493225, 0.08056136965751648, 0.04282544180750847, 0.09609510749578476, 0.21831035614013672, 0.11452964693307877, 0.4344905614852905, 0.09872471541166306, 0.06769980490207672, 0.054214250296354294, 0.015440859831869602, 0.04572026804089546, 0.05267196521162987, 0.06955287605524063, 7.444373295584228e-06, 4.17321571148932e-05, 0.5221405029296875, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1362180858850479, 0.01786869764328003, 0.3548091650009155, 0.13650378584861755, 0.07479218393564224, 0.08773932605981827, 0.007214170414954424, 0.020996512845158577, 0.09793394804000854, 0.26323461532592773, 0.31718939542770386, 0.004400049336254597, 0.01118874829262495, 0.016452480107545853, 0.0059462906792759895, 0.09023705869913101, 0.59262615442276, 0.038057319819927216, 0.1896824985742569, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13787487149238586, 0.02221597172319889, 0.46063661575317383, 0.42787930369377136, 0.16819633543491364, 0.30927538871765137, 0.10940644890069962, 0.14741046726703644, 0.3708270192146301, 0.08424455672502518, 0.34931957721710205, 0.015041538514196873, 0.02219252847135067, 0.0637117251753807, 0.001682900357991457, 0.0001943353418027982, 0.004992108792066574, 0.35714879631996155, 0.028785984963178635, 0.7041940689086914, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09526984393596649, 0.013222168199717999, 0.9035038352012634, 0.8715099692344666, 0.20107677578926086, 0.7829492688179016, 0.28305909037590027, 0.141366645693779, 0.15355023741722107, 0.11376345157623291, 0.804192841053009, 0.012117957696318626, 0.3312073349952698, 0.4514775276184082, 0.016239164397120476, 1.0879062756430358e-05, 5.022298137191683e-05, 0.0836932584643364, 0.0041815838776528835, 0.7177854776382446, 0.4451410174369812, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.34537556767463684, 0.010514522902667522, 0.04824088513851166, 0.12771852314472198, 0.005308120045810938, 0.17857761681079865, 0.2263273000717163, 0.26537755131721497, 0.3297313451766968, 0.3104889690876007, 0.11654951423406601, 0.08535956591367722, 0.02363554947078228, 0.031254567205905914, 0.10634612292051315, 0.003986984025686979, 0.03902542591094971, 0.00027279910864308476, 0.00016326647892128676, 0.09999275952577591, 0.23601794242858887, 0.8888784646987915, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2808375656604767, 0.07436379790306091, 0.11235158890485764, 0.07017786800861359, 0.034851111471652985, 0.01653558947145939, 0.025893066078424454, 0.02911091037094593, 0.23654304444789886, 0.2646749019622803, 0.20617236196994781, 0.25081631541252136, 0.013157923705875874, 0.04621773213148117, 0.2354249358177185, 0.0004483810334932059, 0.01581367664039135, 0.00053547159768641, 0.005416989792138338, 0.0004931549192406237, 1.743426764733158e-06, 0.0002464183489792049, 0.38669928908348083, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5487799644470215, 0.03728892654180527, 0.05227963626384735, 0.18957917392253876, 0.014632479287683964, 0.19499987363815308, 0.29326584935188293, 0.6778355836868286, 0.45779454708099365, 0.33408117294311523, 0.11356081813573837, 0.01941866986453533, 0.010207045823335648, 0.013884961605072021, 0.09069465100765228, 0.0014915558276697993, 0.0036082565784454346, 0.0005674233543686569, 0.0010717788245528936, 0.04321836307644844, 0.5446166396141052, 0.38359156250953674, 0.006869717035442591, 0.0028910271357744932, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09531711786985397, 0.03595840558409691, 0.017401238903403282, 0.061305541545152664, 0.1627957820892334, 0.050434935837984085, 0.05516263470053673, 0.23917846381664276, 0.3637218177318573, 0.09729932248592377, 0.03891580551862717, 0.19205324351787567, 0.041229162365198135, 0.046046942472457886, 0.03756402060389519, 8.035104838199914e-05, 0.005924052093178034, 0.005847892723977566, 0.020417997613549232, 0.11436353623867035, 0.6555760502815247, 0.4247216582298279, 0.04553407058119774, 0.00039129320066422224, 0.013846640475094318, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08811857551336288, 0.010963470675051212, 0.2593647241592407, 0.26678594946861267, 0.42746680974960327, 0.41530901193618774, 0.07491520792245865, 0.18910719454288483, 0.04928334057331085, 0.04599721357226372, 0.4843277335166931, 0.07717985659837723, 0.09353034198284149, 0.07800954580307007, 0.08156391978263855, 0.0012459981953725219, 0.12171746790409088, 0.022806251421570778, 0.021380947902798653, 0.018195364624261856, 0.08835338801145554, 0.20732422173023224, 0.30439698696136475, 0.09951408952474594, 0.2512991428375244, 0.4290468692779541, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04596662148833275, 0.005170373246073723, 0.12165658175945282, 0.15079215168952942, 0.04554709792137146, 0.08856093138456345, 0.04626012593507767, 0.020681705325841904, 0.17637456953525543, 0.26189061999320984, 0.13335715234279633, 0.046832337975502014, 0.018430203199386597, 0.01621258072555065, 0.10917440801858902, 0.007976139895617962, 0.03435874730348587, 0.026849543675780296, 0.002102706115692854, 0.13315419852733612, 0.1177494078874588, 0.08904305100440979, 0.576798677444458, 0.140389084815979, 0.6266443729400635, 0.32779327034950256, 0.5110495090484619, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5138411521911621, 0.0654044821858406, 0.1128465011715889, 0.18054738640785217, 0.038166921585798264, 0.13531430065631866, 0.12295213341712952, 0.28065726161003113, 0.2875981628894806, 0.5909985899925232, 0.601227879524231, 0.03077608533203602, 0.04096299037337303, 0.09236451238393784, 0.1495288461446762, 0.0015641784993931651, 0.09294694662094116, 0.006881145294755697, 0.0020365919917821884, 0.4301930069923401, 0.06383264064788818, 0.0045266724191606045, 0.17422647774219513, 0.00404678238555789, 0.006469257641583681, 0.052995309233665466, 0.1725381463766098, 0.668171763420105, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07072688639163971, 0.012152088806033134, 0.021357353776693344, 0.04663744568824768, 0.020319821313023567, 0.05489102751016617, 0.07223928719758987, 0.23148301243782043, 0.18188072741031647, 0.10590049624443054, 0.10450157523155212, 0.03876996785402298, 0.13536545634269714, 0.10362161695957184, 0.12556865811347961, 0.004304439760744572, 0.05993141233921051, 0.054169829934835434, 0.025809768587350845, 0.7262899279594421, 0.2466905415058136, 0.15344326198101044, 0.33606013655662537, 0.02952432446181774, 0.07010773569345474, 0.008777104318141937, 0.03394261747598648, 0.032566726207733154, 0.6152393221855164, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07390952110290527, 0.023819932714104652, 0.4992673993110657, 0.293674498796463, 0.18016116321086884, 0.3294305205345154, 0.5326097011566162, 0.20817913115024567, 0.231731578707695, 0.17336609959602356, 0.4696378707885742, 0.3560185134410858, 0.5055418610572815, 0.687153697013855, 0.06569264829158783, 1.0540320545260329e-05, 0.0013190202880650759, 0.20101842284202576, 0.004686327185481787, 0.13271625339984894, 0.04526880756020546, 0.0007031870190985501, 0.0011485026916489005, 0.002882149303331971, 0.0005991549696773291, 0.0030197217129170895, 0.004800362046808004, 0.004403174854815006, 0.002436757553368807, 0.4002683460712433, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19887569546699524, 0.009285598993301392, 0.17495201528072357, 0.1799449920654297, 0.0410592183470726, 0.0050115324556827545, 0.025978662073612213, 0.011312133632600307, 0.04069671407341957, 0.23767657577991486, 0.3294059634208679, 0.09899688512086868, 0.03285939246416092, 0.08387716114521027, 0.04885585233569145, 0.0003210107679478824, 0.5876501798629761, 0.16318874061107635, 0.7096263766288757, 0.11595475673675537, 0.007003267295658588, 0.001205803593620658, 0.1902448534965515, 0.011727835983037949, 0.44888344407081604, 0.8117052912712097, 0.45698752999305725, 0.023960944265127182, 0.010929742828011513, 0.005293603055179119, 0.00987145397812128, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.054675761610269547, 0.04458622261881828, 0.0536046139895916, 0.016943499445915222, 0.02146792784333229, 0.1686052531003952, 0.036354243755340576, 0.08614800870418549, 0.1611979901790619, 0.170720174908638, 0.163726344704628, 0.09202460944652557, 0.016866492107510567, 0.019021833315491676, 0.13082824647426605, 0.020372437313199043, 0.3410835862159729, 0.6929088234901428, 0.04383905977010727, 0.1458517462015152, 0.4223538339138031, 0.9439106583595276, 0.9473816156387329, 0.15120889246463776, 0.7730743288993835, 0.5082507133483887, 0.0460858978331089, 0.032336097210645676, 0.011211436241865158, 0.009573124349117279, 0.0003536108124535531, 0.06564418971538544, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.254617303609848, 0.09600356966257095, 0.5283652544021606, 0.35948434472084045, 0.11690203100442886, 0.22449535131454468, 0.07030754536390305, 0.14074397087097168, 0.11056768894195557, 0.2017645388841629, 0.5897989273071289, 0.032950446009635925, 0.0850306898355484, 0.16881772875785828, 0.07667817175388336, 0.020423829555511475, 0.09150233864784241, 0.593336284160614, 0.050333935767412186, 0.04262891411781311, 0.44151586294174194, 0.7098277807235718, 0.36869171261787415, 0.7183430194854736, 0.3146522641181946, 0.5934929251670837, 0.08962199836969376, 0.01141325756907463, 0.0268073882907629, 0.008290876634418964, 0.022364463657140732, 0.0520397312939167, 0.3134966492652893, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06611059606075287, 0.009380446746945381, 0.1600489318370819, 0.18714633584022522, 0.028496628627181053, 0.28509950637817383, 0.06793918460607529, 0.036412376910448074, 0.3864555358886719, 0.38031718134880066, 0.19321800768375397, 0.03279240429401398, 0.024823389947414398, 0.02684853971004486, 0.10572600364685059, 0.008604546077549458, 0.07562410086393356, 0.10463645309209824, 0.003217896446585655, 0.1296835094690323, 0.21162182092666626, 0.30799001455307007, 0.7962209582328796, 0.27782267332077026, 0.5974112749099731, 0.3643631041049957, 0.5975222587585449, 0.032379183918237686, 0.8344925045967102, 0.5903766751289368, 0.1521190106868744, 0.10492946952581406, 0.10503242909908295, 0.5022279620170593, NaN, NaN, NaN, NaN, NaN, NaN], [0.5806823372840881, 0.09046274423599243, 0.1468239277601242, 0.2587219774723053, 0.018666794523596764, 0.17986845970153809, 0.1758078932762146, 0.26734092831611633, 0.30597683787345886, 0.6407824158668518, 0.6427304148674011, 0.011203133501112461, 0.017842967063188553, 0.05609212443232536, 0.1528221219778061, 0.0010157334618270397, 0.08574047684669495, 0.010654903016984463, 0.003869200125336647, 0.15051355957984924, 0.02434478886425495, 0.005829520523548126, 0.10341739654541016, 0.0023463659454137087, 0.00469975033774972, 0.1621563881635666, 0.27765417098999023, 0.6246147155761719, 0.44377410411834717, 0.0757245346903801, 0.08620554953813553, 0.08146335929632187, 0.32109129428863525, 0.1958039551973343, 0.5327519178390503, NaN, NaN, NaN, NaN, NaN], [0.09578646719455719, 0.04883359372615814, 0.014442636631429195, 0.07719788700342178, 0.013871591538190842, 0.24272511899471283, 0.11848346889019012, 0.48695430159568787, 0.10090471804141998, 0.15632015466690063, 0.12246286869049072, 0.056596189737319946, 0.051980338990688324, 0.03806659206748009, 0.1369783878326416, 0.0009064326295629144, 0.04867112636566162, 0.09537991136312485, 0.12993541359901428, 0.38632717728614807, 0.056282784789800644, 0.13602504134178162, 0.18383464217185974, 0.024170320481061935, 0.09972675889730453, 0.022063996642827988, 0.042059145867824554, 0.01842264086008072, 0.8592916131019592, 0.1306053251028061, 0.06485681235790253, 0.048735883086919785, 0.037178389728069305, 0.017466288059949875, 0.006924192421138287, 0.8764364123344421, NaN, NaN, NaN, NaN], [0.12923087179660797, 0.04506811499595642, 0.5631698966026306, 0.4945719838142395, 0.16776354610919952, 0.4656532406806946, 0.6344242095947266, 0.28209388256073, 0.297488808631897, 0.3520771265029907, 0.6463941931724548, 0.3803158104419708, 0.4924411177635193, 0.6891878843307495, 0.08469904214143753, 1.2418378219081205e-06, 0.0003037750138901174, 0.10264009237289429, 0.0010840333998203278, 0.03004724159836769, 0.00720690144225955, 0.00017297905287705362, 0.00021026108879595995, 0.0005732537247240543, 0.00013229742762632668, 0.0014890850288793445, 0.0027206502854824066, 0.0022100789938122034, 0.0018764312844723463, 0.22427155077457428, 0.0012303950497880578, 0.0001426686649210751, 0.0015814924845471978, 0.00487141590565443, 0.0029599322006106377, 0.003610847517848015, 0.41901907324790955, NaN, NaN, NaN], [0.3177553117275238, 0.027823492884635925, 0.11541304737329483, 0.1464630663394928, 0.010460668243467808, 0.028609508648514748, 0.14352867007255554, 0.043905869126319885, 0.18215790390968323, 0.6030426025390625, 0.38763877749443054, 0.1293274313211441, 0.07180552184581757, 0.1464845985174179, 0.10971048474311829, 0.00015546051145065576, 0.5271192193031311, 0.2684091329574585, 0.7487277388572693, 0.0846778005361557, 0.003557654097676277, 0.0064069912768900394, 0.16770148277282715, 0.008421340025961399, 0.27412623167037964, 0.8534677624702454, 0.5243650078773499, 0.02665238454937935, 0.01776440255343914, 0.013793676160275936, 0.00868560466915369, 0.08064579218626022, 0.69512540102005, 0.49261555075645447, 0.010526523925364017, 0.0028473760467022657, 0.008281596936285496, 0.007198471110314131, NaN, NaN], [0.03459807112812996, 0.05000016465783119, 0.02839210256934166, 0.008521324954926968, 0.009519261308014393, 0.12168280780315399, 0.03372196480631828, 0.07665831595659256, 0.21765880286693573, 0.11945746093988419, 0.0821232944726944, 0.058310747146606445, 0.011853469535708427, 0.02031784877181053, 0.13586042821407318, 0.03285643830895424, 0.3327244818210602, 0.7442528605461121, 0.049526505172252655, 0.13722854852676392, 0.37294694781303406, 0.9746374487876892, 0.9050161242485046, 0.144730344414711, 0.44314900040626526, 0.6168692708015442, 0.18840178847312927, 0.12898683547973633, 0.1250022053718567, 0.01759251020848751, 0.0030696040485054255, 0.6704888939857483, 0.3205258250236511, 0.28675025701522827, 0.09770815074443817, 0.0085873082280159, 0.028106005862355232, 0.0015327840810641646, 0.12156207114458084, NaN], [0.02964477799832821, 0.1353258490562439, 0.017653465270996094, 0.011115004308521748, 0.008141545578837395, 0.05911250412464142, 0.01831989735364914, 0.05519499629735947, 0.03573962301015854, 0.02204814739525318, 0.05097896233201027, 0.08341387659311295, 0.08060181885957718, 0.10490117967128754, 0.13247323036193848, 0.027913866564631462, 0.6360336542129517, 0.8947576880455017, 0.5603421926498413, 0.3501611351966858, 0.3494046926498413, 0.7655782103538513, 0.9696423411369324, 0.8922762274742126, 0.42980051040649414, 0.4555767774581909, 0.17016178369522095, 0.1410100758075714, 0.652664303779602, 0.2781027853488922, 0.07839874923229218, 0.11400053650140762, 0.10023999214172363, 0.04957454651594162, 0.07193805277347565, 0.5185664892196655, 0.15356925129890442, 0.02747632935643196, 0.046240244060754776, 0.017650051042437553]], [[0.011476250365376472, 0.7629169225692749, 0.02116730809211731, 0.010803135111927986, 0.005132503807544708, 0.009303245693445206, 0.0005040443502366543, 0.022131631150841713, 0.001470191520638764, 0.0017710012616589665, 0.0004086543631274253, 0.0022351557854562998, 0.000896299781743437, 0.0005698543391190469, 0.019197434186935425, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0024000771809369326, 0.158247172832489, 0.01897430047392845, 0.019486481323838234, 0.0029122373089194298, 0.015832845121622086, 0.0017470666207373142, 0.00117065932136029, 0.01016113068908453, 0.007651789113879204, 0.0020597530528903008, 0.015201352536678314, 0.016943661496043205, 0.009769451804459095, 0.16634535789489746, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.00410552928224206, 0.0015743908006697893, 0.01049637421965599, 0.006504607852548361, 0.035339318215847015, 0.9065937995910645, 0.2998698651790619, 0.12215600907802582, 0.013029203750193119, 0.000650988076813519, 0.002043183660134673, 0.006920983083546162, 0.09688588231801987, 0.057574767619371414, 0.009054930880665779, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.007287806831300259, 0.01375514268875122, 0.001530585577711463, 0.007056740578263998, 0.01978658139705658, 0.9208202958106995, 0.2214416116476059, 0.30606138706207275, 0.052588097751140594, 0.004079628270119429, 0.0024339878000319004, 0.0028739250265061855, 0.04695972800254822, 0.045893676578998566, 0.0110039496794343, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.006429406348615885, 0.016907041892409325, 0.0023819799534976482, 0.0003115522558800876, 0.006808500271290541, 0.9102355241775513, 0.15379303693771362, 0.07056371122598648, 0.06324119120836258, 0.0030630400869995356, 0.007665702607482672, 0.002797773340716958, 0.13533660769462585, 0.03197972849011421, 0.006115978583693504, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.014356410130858421, 0.0526699461042881, 0.0007501932559534907, 0.008851941674947739, 0.0005067299935035408, 0.035332534462213516, 0.09051518887281418, 0.049224019050598145, 0.014900125563144684, 0.01856788620352745, 0.0012414768571034074, 0.002389064058661461, 0.0018446464091539383, 0.000877396494615823, 0.22725383937358856, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0025407460052520037, 0.32041609287261963, 0.0036992463283240795, 0.02451898716390133, 0.007920290343463421, 0.015527674928307533, 0.03544912114739418, 0.29718661308288574, 0.02347515895962715, 0.026838794350624084, 0.01756858080625534, 0.010445725172758102, 0.005995406303554773, 0.0005847325082868338, 0.2055930197238922, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.009255345910787582, 0.034783441573381424, 0.010831266641616821, 0.02782595343887806, 0.001477425335906446, 0.006871670484542847, 0.006518858019262552, 0.0072874827310442924, 0.012387615628540516, 0.05288432911038399, 0.04645476117730141, 0.02255677618086338, 0.014156763441860676, 0.00417641457170248, 0.22105874121189117, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0017225841293111444, 0.0049251834861934185, 0.007573804818093777, 0.014873476698994637, 0.00903867557644844, 0.0076865823939442635, 0.0017025101697072387, 0.00023153165238909423, 0.024773191660642624, 0.1742238849401474, 0.6002998948097229, 0.6145275831222534, 0.25023365020751953, 0.35489538311958313, 0.039457567036151886, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0034636815544217825, 0.39023807644844055, 0.0018667654367163777, 0.0006454490358009934, 0.00025732445647008717, 0.026610050350427628, 0.0026998629327863455, 0.014584111049771309, 0.00032847325201146305, 0.0012709795264527202, 0.07417861372232437, 0.43676891922950745, 0.25757044553756714, 0.32731080055236816, 0.12109360098838806, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0014396773185580969, 0.07700426131486893, 0.0003769460890907794, 0.0015669490676373243, 0.0010665652807801962, 0.05166712775826454, 0.003733921330422163, 0.00829349085688591, 9.729996236274019e-05, 0.0004270579374860972, 0.0022819112055003643, 0.3744491934776306, 0.2681969404220581, 0.4920969009399414, 0.028773367404937744, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.19549021124839783, 0.5118218064308167, 0.053603943437337875, 0.004430307075381279, 0.0015711480518803, 0.024018822237849236, 0.0441354438662529, 0.04134393110871315, 0.0014472270850092173, 0.024767767637968063, 0.029112013056874275, 0.08014442026615143, 0.4702226519584656, 0.40423843264579773, 0.14477935433387756, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.034691162407398224, 0.09692039340734482, 0.003936667460948229, 0.0164506658911705, 0.0005446859868243337, 0.0016573348548263311, 0.02795562334358692, 0.12881094217300415, 0.0004645287699531764, 0.0021237744949758053, 0.0010291342623531818, 0.001068241661414504, 0.00471450574696064, 0.019945403560996056, 0.19273433089256287, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04783029109239578, 0.11157537996768951, 0.02325829118490219, 0.12799327075481415, 0.0216610599309206, 0.41526544094085693, 0.129922553896904, 0.14850500226020813, 0.0009580283658578992, 0.008097043260931969, 0.01107556838542223, 0.019478609785437584, 0.2748490571975708, 0.11550750583410263, 0.15876543521881104, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.015012643299996853, 0.00804762914776802, 0.00366173661313951, 0.0018753333715721965, 0.0065993256866931915, 0.00479541253298521, 0.005337378475815058, 0.012457020580768585, 0.0033909485209733248, 0.0032401280477643013, 0.00048777347547002137, 0.012255984358489513, 0.0006230318685993552, 0.001543535152450204, 0.1572250872850418, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.20067201554775238, 0.150595024228096, 0.3375815153121948, 0.5753223896026611, 0.03983612731099129, 0.13901081681251526, 0.37267425656318665, 0.07406412810087204, 0.07071352750062943, 0.22996902465820312, 0.35784539580345154, 0.0401473231613636, 0.03251379355788231, 0.07572956383228302, 0.005637211725115776, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.055522263050079346, 0.0030253075528889894, 0.054468654096126556, 0.18383808434009552, 0.2751407325267792, 0.06163792684674263, 0.5092534422874451, 0.21577699482440948, 0.23691882193088531, 0.32801976799964905, 0.29786956310272217, 0.4967685043811798, 0.6341143250465393, 0.7677603363990784, 0.40264371037483215, 0.02477514185011387, 0.37543168663978577, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0005822544917464256, 0.0004425827646628022, 0.0014265297213569283, 0.0006841197027824819, 0.03406556695699692, 0.0010687633184716105, 0.0028485425282269716, 0.020860498771071434, 0.05133597180247307, 0.002158694202080369, 0.002441320102661848, 0.037159714847803116, 0.005256796721369028, 0.008102376013994217, 0.16207638382911682, 0.02274254709482193, 0.6458237767219543, 0.013541627675294876, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20224374532699585, 0.7376267313957214, 0.004014236852526665, 0.0103965038433671, 0.07275543361902237, 0.03262623772025108, 0.04577071964740753, 0.5017040371894836, 0.12205435335636139, 0.19255708158016205, 0.006990006659179926, 0.028381695970892906, 0.046785227954387665, 0.15206293761730194, 0.330488920211792, 0.03146426007151604, 0.019330549985170364, 0.019686071202158928, 0.5363749265670776, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3634231686592102, 0.404717355966568, 0.00689590023830533, 0.04770800471305847, 0.0251657422631979, 0.0006883289897814393, 0.02071242779493332, 0.019072405993938446, 0.15776626765727997, 0.3694642186164856, 0.036826737225055695, 0.23951902985572815, 0.011015082709491253, 0.04999716952443123, 0.2037181556224823, 0.05261930450797081, 0.12757715582847595, 0.003555318573489785, 0.48483166098594666, 0.00033596818684600294, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.8270207643508911, 0.8942698836326599, 0.020243747159838676, 0.04263966530561447, 0.09284591674804688, 0.054453812539577484, 0.21418678760528564, 0.23612302541732788, 0.5479635000228882, 0.7225908041000366, 0.08608872443437576, 0.5934221148490906, 0.30024465918540955, 0.22648638486862183, 0.12622572481632233, 0.09825422614812851, 0.08890903741121292, 0.0022953739389777184, 0.3788372278213501, 6.525879871333018e-05, 3.547202504705638e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.043734412640333176, 0.7137998342514038, 0.1370490938425064, 0.045488547533750534, 0.06789389997720718, 0.49671053886413574, 0.1280447244644165, 0.4211912155151367, 0.03652801364660263, 0.041476957499980927, 0.08040425181388855, 0.19641457498073578, 0.603863537311554, 0.49263066053390503, 0.07636027038097382, 0.1839720457792282, 0.005392392631620169, 0.0012601928319782019, 0.000860364583786577, 0.0008281354093924165, 0.0005760629428550601, 0.002849774667993188, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.017375759780406952, 0.012506993487477303, 0.020720014348626137, 0.011049210093915462, 0.03743210807442665, 0.0072485157288610935, 0.03524084761738777, 0.005443913396447897, 0.24646395444869995, 0.048276107758283615, 0.03640883043408394, 0.507624089717865, 0.15355341136455536, 0.1730290949344635, 0.2644885182380676, 0.005911883432418108, 0.0029267233330756426, 0.007144090253859758, 0.001919957809150219, 0.004637785721570253, 0.004848909098654985, 0.006189228966832161, 0.3764636814594269, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09840062260627747, 0.7509858012199402, 0.13933908939361572, 0.13482652604579926, 0.18154919147491455, 0.32397931814193726, 0.23646889626979828, 0.11657525599002838, 0.03430478647351265, 0.1277371644973755, 0.15700362622737885, 0.24829043447971344, 0.7591869831085205, 0.7825927138328552, 0.06869770586490631, 0.2256152480840683, 0.0020181250292807817, 0.0012439934071153402, 0.00031968209077604115, 0.0029859780333936214, 0.017534615471959114, 0.0004058087943121791, 0.00034323628642596304, 0.029154805466532707, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.22806629538536072, 0.6706615686416626, 0.2560598850250244, 0.17412559688091278, 0.6327939033508301, 0.04699348285794258, 0.058767881244421005, 0.11556732654571533, 0.09056147933006287, 0.3648419678211212, 0.5388886332511902, 0.261055588722229, 0.6016876697540283, 0.7496042847633362, 0.0894755870103836, 0.03960844501852989, 0.0036635666619986296, 0.00109457119833678, 0.0017422186210751534, 0.022469639778137207, 0.004235065542161465, 0.007348764222115278, 0.00280297570861876, 0.030011437833309174, 0.576508641242981, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5419997572898865, 0.6956567168235779, 0.044124722480773926, 0.12586495280265808, 0.048711128532886505, 0.11729516834020615, 0.4073715806007385, 0.43757542967796326, 0.032695479691028595, 0.4824156165122986, 0.05927032604813576, 0.04766178876161575, 0.25393223762512207, 0.23675066232681274, 0.10572775453329086, 0.0628783106803894, 0.014568633399903774, 0.003403500886633992, 0.005917230620980263, 0.009509358555078506, 0.0019911406561732292, 0.005211993586272001, 0.01603839360177517, 0.00502167409285903, 0.3301290273666382, 0.10268117487430573, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09369882941246033, 0.5731168985366821, 0.13611510396003723, 0.13756731152534485, 0.024227088317275047, 0.31910547614097595, 0.16772453486919403, 0.1680929958820343, 0.09319504350423813, 0.0998181626200676, 0.22465890645980835, 0.00899507012218237, 0.16640731692314148, 0.25350457429885864, 0.09016240388154984, 0.178706556558609, 0.5124386548995972, 0.028256116434931755, 0.011254883371293545, 0.03223628178238869, 0.0004171380714979023, 0.004843876231461763, 0.09010603278875351, 0.0025540743954479694, 0.016201328486204147, 0.029397757723927498, 0.010837158188223839, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02838694490492344, 0.30040091276168823, 0.005878766532987356, 0.015430719591677189, 0.017050068825483322, 0.06605669111013412, 0.12745192646980286, 0.23377051949501038, 0.08052214235067368, 0.033177152276039124, 0.06731567531824112, 0.07575374841690063, 0.18187224864959717, 0.570769727230072, 0.04572387412190437, 0.18362975120544434, 0.10373001545667648, 0.006869313772767782, 0.010921900160610676, 0.01820673979818821, 0.0017379705095663667, 0.002349345711991191, 0.03729201853275299, 5.792165029561147e-05, 0.0013579311780631542, 0.0025659396778792143, 0.008523254655301571, 0.1568114459514618, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2655380666255951, 0.4107033908367157, 0.04865417629480362, 0.08488347381353378, 0.04310445114970207, 0.10849997401237488, 0.15643075108528137, 0.04165918007493019, 0.12898734211921692, 0.11095981299877167, 0.23520684242248535, 0.10632039606571198, 0.055878568440675735, 0.24558725953102112, 0.17682571709156036, 0.060853905975818634, 0.016029829159379005, 0.001439533894881606, 0.017260756343603134, 0.0007974627078510821, 0.0012342276750132442, 0.028226196765899658, 0.0047790613025426865, 0.0015612602001056075, 0.004867547657340765, 0.039023980498313904, 0.05208572745323181, 0.33480554819107056, 0.17332881689071655, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.8565200567245483, 0.8639481067657471, 0.0803997814655304, 0.36449819803237915, 0.17448320984840393, 0.12402030825614929, 0.13765643537044525, 0.2065785825252533, 0.18182852864265442, 0.6806339025497437, 0.1919344812631607, 0.19068314135074615, 0.004361266735941172, 0.01490570418536663, 0.13936595618724823, 0.043774526566267014, 0.2669547498226166, 0.035314492881298065, 0.1941595822572708, 0.006638282909989357, 0.005091785918921232, 0.2628510892391205, 0.2860943675041199, 0.06445851922035217, 0.34950578212738037, 0.6430334448814392, 0.5673049688339233, 0.6101463437080383, 0.29372307658195496, 0.0028161092195659876, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.22751423716545105, 0.21127405762672424, 0.005130667705088854, 0.028237944468855858, 0.06646221876144409, 0.045109983533620834, 0.478432834148407, 0.6443154215812683, 0.140235036611557, 0.0980456992983818, 0.006476161070168018, 0.038696710020303726, 0.25798937678337097, 0.10561345517635345, 0.16755780577659607, 0.018545497208833694, 0.059764593839645386, 0.0026272537652403116, 0.020267995074391365, 0.009687644429504871, 0.00033462722785770893, 0.0024671528954058886, 0.054633729159832, 5.4464391723740846e-05, 0.00043273900519125164, 0.0019224031129851937, 0.21117039024829865, 0.3183750510215759, 0.03866858780384064, 0.011778384447097778, 0.1297062188386917, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3886019289493561, 0.36600789427757263, 0.07069597393274307, 0.12792876362800598, 0.0629734918475151, 0.0820467472076416, 0.2973020672798157, 0.27475541830062866, 0.019707435742020607, 0.2982620298862457, 0.24423947930335999, 0.05686682090163231, 0.23438367247581482, 0.3444555997848511, 0.09858046472072601, 0.0004199208051431924, 4.603992783813737e-05, 8.09443406524224e-07, 2.029701317951549e-05, 3.386533080629306e-06, 2.203315261795069e-06, 4.220597020321293e-06, 8.901660294213798e-06, 0.00016298270202241838, 0.000983458710834384, 0.0005640776362270117, 0.0008154786773957312, 0.001651398022659123, 2.400618996034609e-06, 3.3168395020766184e-05, 6.549440058734035e-06, 0.8699775338172913, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.31350865960121155, 0.5118260383605957, 0.01775331422686577, 0.060602445155382156, 0.015971101820468903, 0.03445184975862503, 0.4316053092479706, 0.4819965064525604, 0.008238772861659527, 0.27349013090133667, 0.02135261707007885, 0.006705985404551029, 0.06119696795940399, 0.05213680863380432, 0.13011163473129272, 0.06053417548537254, 0.012584012933075428, 0.0010002547642216086, 0.0027718576602637768, 0.006610550452023745, 0.0029896856285631657, 0.008355176076292992, 0.048459943383932114, 0.002307809190824628, 0.65205979347229, 0.1651758849620819, 0.011300449259579182, 0.029586348682641983, 0.014456091448664665, 0.0007872084970586002, 0.0008902085828594863, 0.029332326725125313, 0.16636918485164642, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11128952354192734, 0.6662537455558777, 0.10913366079330444, 0.08027850091457367, 0.016604425385594368, 0.1904260814189911, 0.09001538157463074, 0.12034764140844345, 0.032395973801612854, 0.07767382264137268, 0.13288450241088867, 0.0038343279156833887, 0.15461067855358124, 0.13092683255672455, 0.1198263093829155, 0.19553376734256744, 0.2426333725452423, 0.004519153386354446, 0.00883245188742876, 0.006844275165349245, 0.00014635240950156003, 0.00260242260992527, 0.03859727829694748, 0.0011520206462591887, 0.014703472144901752, 0.016579829156398773, 0.003783928230404854, 0.01771795004606247, 0.0035672299563884735, 0.000677697011269629, 0.002100451150909066, 0.023971345275640488, 0.03231354430317879, 0.011524699628353119, NaN, NaN, NaN, NaN, NaN, NaN], [0.045069050043821335, 0.5156355500221252, 0.014353718608617783, 0.026371080428361893, 0.027669712901115417, 0.08119883388280869, 0.2510265111923218, 0.45373910665512085, 0.0644708126783371, 0.03346102684736252, 0.06456929445266724, 0.036929432302713394, 0.1635800451040268, 0.4964689314365387, 0.12627021968364716, 0.17035169899463654, 0.07290639728307724, 0.0013864204520359635, 0.008776376023888588, 0.010795027948915958, 0.0008890280150808394, 0.00375909055583179, 0.03264426812529564, 2.1074760297778994e-05, 0.0009656226029619575, 0.004805654752999544, 0.015095297247171402, 0.19429266452789307, 0.060086220502853394, 0.013300183229148388, 0.019145654514431953, 0.08634541183710098, 0.018065713346004486, 0.012390428222715855, 0.3474832773208618, NaN, NaN, NaN, NaN, NaN], [0.15574656426906586, 0.22756966948509216, 0.016156630590558052, 0.0469389408826828, 0.01719032973051071, 0.01580459624528885, 0.07493647187948227, 0.02412206307053566, 0.018628407269716263, 0.03879624605178833, 0.03891688585281372, 0.03379734605550766, 0.008454171009361744, 0.03055991418659687, 0.1906210333108902, 0.002681915881112218, 0.0020622191950678825, 1.740588413667865e-05, 0.001647116499952972, 2.462047996232286e-05, 1.4256034774007276e-05, 0.0023770714178681374, 0.0007797144935466349, 6.146806117612869e-05, 0.00019536878971848637, 0.023629816249012947, 0.022664623335003853, 0.058040015399456024, 0.02328144572675228, 0.00014305225340649486, 0.1791975051164627, 0.7950490117073059, 0.40287262201309204, 0.05916967615485191, 0.11726692318916321, 0.045271970331668854, NaN, NaN, NaN, NaN], [0.7930518984794617, 0.8248118162155151, 0.03787774592638016, 0.2306395173072815, 0.10945193469524384, 0.048738475888967514, 0.07385316491127014, 0.1171715259552002, 0.09199279546737671, 0.5013920664787292, 0.07074998319149017, 0.14583703875541687, 0.0018764830892905593, 0.00646476075053215, 0.13562877476215363, 0.017539121210575104, 0.07800457626581192, 0.013338283635675907, 0.07843150943517685, 0.003389358287677169, 0.0011982140131294727, 0.07936429977416992, 0.08406823873519897, 0.016710255295038223, 0.13201765716075897, 0.339507520198822, 0.3268124461174011, 0.4709261357784271, 0.24707961082458496, 0.0009133804705925286, 0.27326905727386475, 0.539431095123291, 0.8842423558235168, 0.5773340463638306, 0.643308699131012, 0.15606866776943207, 0.0011033734772354364, NaN, NaN, NaN], [0.139163076877594, 0.17112046480178833, 0.0021531793754547834, 0.0053843106143176556, 0.013183848932385445, 0.014547600410878658, 0.39682450890541077, 0.7216413021087646, 0.013683686964213848, 0.038195278495550156, 0.0014429710572585464, 0.0075409854762256145, 0.06976743042469025, 0.016425929963588715, 0.1257757991552353, 0.0009739195229485631, 0.0011780881322920322, 3.265493069193326e-05, 0.0005334040033631027, 0.0007281061843968928, 3.2774634746601805e-05, 0.0004276044783182442, 0.00342408730648458, 2.9227990125946235e-06, 5.522280844161287e-05, 0.00012372780474834144, 0.011400841176509857, 0.008755120448768139, 0.0017365129897370934, 0.0007705622701905668, 0.0024924452882260084, 0.4634210169315338, 0.010356471873819828, 0.06587640196084976, 0.03498200699687004, 0.005118835251778364, 0.0019369632937014103, 0.023791478946805, NaN, NaN], [0.37428542971611023, 0.3404470980167389, 0.07186836749315262, 0.11062464118003845, 0.09624961018562317, 0.06910651177167892, 0.26704323291778564, 0.35990291833877563, 0.016681469976902008, 0.31615501642227173, 0.23382727801799774, 0.051282789558172226, 0.1643712818622589, 0.24623094499111176, 0.1059461385011673, 0.00023119446996133775, 9.065014637599234e-06, 3.0932378081161005e-07, 7.128239758458221e-06, 2.417179757685517e-06, 1.9917408735636855e-06, 1.0686825362427044e-06, 3.5747166293731425e-06, 3.038432441826444e-05, 0.00024045849568210542, 0.00012102597975172102, 0.0003720777458511293, 0.0005474414792843163, 4.2138731259910855e-06, 8.004362825886346e-06, 4.010584234492853e-06, 0.22906039655208588, 0.00024706448311917484, 0.003541025100275874, 0.0035716970451176167, 1.1338630656609894e-06, 4.888530747848563e-05, 2.00755093828775e-05, 0.8455927968025208, NaN], [0.2896858751773834, 0.2041676938533783, 0.0844137892127037, 0.26597079634666443, 0.007990201003849506, 0.057605594396591187, 0.37075188755989075, 0.33039090037345886, 0.04668770357966423, 0.6492098569869995, 0.34850311279296875, 0.12703292071819305, 0.22453922033309937, 0.2423134297132492, 0.11649563163518906, 0.023575956001877785, 0.001566409133374691, 0.0004935376346111298, 0.015205318108201027, 0.0005761805805377662, 0.00026375881861895323, 0.0017682479228824377, 0.00015503005124628544, 0.011253873817622662, 0.321735680103302, 0.05970581993460655, 0.008942467160522938, 0.051820773631334305, 0.009087985381484032, 0.002068085130304098, 0.00584985688328743, 0.01019755844026804, 0.16441591084003448, 0.021173937246203423, 0.09159599989652634, 0.004452125634998083, 0.0037374526727944613, 0.01578103005886078, 0.01742226630449295, 0.3373567461967468]]], [[[0.016101790592074394, 0.0050575402565300465, 0.008322462439537048, 0.006855499465018511, 0.003766664071008563, 0.0032708626240491867, 0.008669405244290829, 0.016983401030302048, 0.023632090538740158, 0.0007983215618878603, 0.006762287113815546, 0.019076332449913025, 0.0018054646207019687, 0.011848386377096176, 0.23875673115253448, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03118298575282097, 0.022700916975736618, 0.01820814236998558, 0.011041272431612015, 0.013735579326748848, 0.003388292621821165, 0.014374880120158195, 0.0029534229543060064, 0.06276529282331467, 0.0010488847037777305, 0.005698299501091242, 0.018068330362439156, 0.009247002191841602, 0.010645000264048576, 0.2274351567029953, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10749327391386032, 0.01361121516674757, 0.01930609717965126, 0.025707745924592018, 0.010174103081226349, 0.0019352196250110865, 0.006933925207704306, 0.026056114584207535, 0.003662128932774067, 0.006897854618728161, 0.0015213300939649343, 0.006132383830845356, 0.0028239174280315638, 0.013304864056408405, 0.22739072144031525, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.25010421872138977, 0.005582309328019619, 0.006115755997598171, 0.08664196729660034, 0.005224197171628475, 0.005311913322657347, 0.03281412273645401, 0.024678068235516548, 0.018595430999994278, 0.0819764956831932, 0.005479714833199978, 0.008821909315884113, 0.02042486146092415, 0.03525637462735176, 0.19444485008716583, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1781134456396103, 0.021083489060401917, 0.038613177835941315, 0.16417931020259857, 0.0029645320028066635, 0.00899361353367567, 0.009076704271137714, 0.01357053779065609, 0.01101364754140377, 0.04086701199412346, 0.014270029030740261, 0.011464214883744717, 0.011689195409417152, 0.0706799253821373, 0.3730076551437378, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3090042769908905, 0.031162124127149582, 0.033009856939315796, 0.14512063562870026, 0.00411824369803071, 0.07382509857416153, 0.02702517993748188, 0.07667822390794754, 0.021658627316355705, 0.01615101285278797, 0.0066233747638762, 0.008623828180134296, 0.0008525048615410924, 0.011195158585906029, 0.2578849792480469, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3291372060775757, 0.0561586357653141, 0.4192807674407959, 0.4571635127067566, 0.057550910860300064, 0.04359428584575653, 0.005270917434245348, 0.03804505616426468, 0.03733760863542557, 0.20409555733203888, 0.04554562643170357, 0.024629684165120125, 0.018161950632929802, 0.04353561997413635, 0.145583838224411, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3828665316104889, 0.019200418144464493, 0.34599530696868896, 0.4376910328865051, 0.07537391781806946, 0.036528222262859344, 0.04610925167798996, 0.04538694769144058, 0.1663823127746582, 0.04690397158265114, 0.05553056299686432, 0.021811597049236298, 0.012554574757814407, 0.03599526360630989, 0.1534716635942459, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08861738443374634, 0.06363938748836517, 0.7135313749313354, 0.146565243601799, 0.3346884250640869, 0.3544132113456726, 0.12204702943563461, 0.028818881139159203, 0.04564356431365013, 0.03288809210062027, 0.06753166019916534, 0.12387087196111679, 0.029650555923581123, 0.014753012917935848, 0.04379607364535332, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03655187785625458, 0.006058508530259132, 0.04018249735236168, 0.08900216966867447, 0.027111714705824852, 0.006408872082829475, 0.03783104568719864, 0.010064247064292431, 0.2550305724143982, 0.008420061320066452, 0.012097015976905823, 0.017737949267029762, 0.0012783813290297985, 0.0026436946354806423, 0.172612726688385, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1163061186671257, 0.04424217715859413, 0.014033653773367405, 0.03590161353349686, 0.06527962535619736, 0.00195779325440526, 0.027195196598768234, 0.1581626534461975, 0.30849722027778625, 0.1652299016714096, 0.04234298691153526, 0.05585171654820442, 0.016547594219446182, 0.04909297078847885, 0.08752257376909256, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1013311892747879, 0.06866802275180817, 0.06425411254167557, 0.4572087228298187, 0.04987834766507149, 0.005650981329381466, 0.053177352994680405, 0.04739876464009285, 0.2551265060901642, 0.06654207408428192, 0.20209699869155884, 0.04737241193652153, 0.042119286954402924, 0.22778292000293732, 0.10508881509304047, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.24632138013839722, 0.045121580362319946, 0.12561434507369995, 0.43826135993003845, 0.07532560080289841, 0.002372375223785639, 0.0398109070956707, 0.026653334498405457, 0.5938559174537659, 0.12655052542686462, 0.04707850515842438, 0.018195422366261482, 0.010826833546161652, 0.023274976760149002, 0.14916135370731354, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.12666325271129608, 0.047387395054101944, 0.04497509077191353, 0.23918962478637695, 0.016611548140645027, 0.009305250830948353, 0.02713325433433056, 0.030590379610657692, 0.4573454260826111, 0.17728003859519958, 0.08635216951370239, 0.05938294902443886, 0.008936652913689613, 0.028742672875523567, 0.15077541768550873, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03701020032167435, 0.037774376571178436, 0.1161394715309143, 0.09335700422525406, 0.015312368050217628, 0.026739761233329773, 0.013009096495807171, 0.005902147851884365, 0.07189750671386719, 0.00625182269141078, 0.056744903326034546, 0.06423129141330719, 0.06661844998598099, 0.02100159414112568, 0.2252311259508133, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.12698857486248016, 0.15100647509098053, 0.08910781890153885, 0.09401589632034302, 0.14288602769374847, 0.07712502032518387, 0.1496707946062088, 0.23784373700618744, 0.024656152352690697, 0.07261883467435837, 0.11269068717956543, 0.10889188945293427, 0.23155105113983154, 0.10633593797683716, 0.14060717821121216, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.33520859479904175, 0.17541100084781647, 0.043081097304821014, 0.07071122527122498, 0.031066332012414932, 0.05302952229976654, 0.13712948560714722, 0.0819549486041069, 0.010218805633485317, 0.05350261554121971, 0.03376028686761856, 0.016291575506329536, 0.04384060204029083, 0.016914406791329384, 0.06937505304813385, 0.1729947179555893, 0.014742943458259106, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2972787618637085, 0.14542943239212036, 0.2801832854747772, 0.6946116089820862, 0.3750338852405548, 0.09368664771318436, 0.11078806221485138, 0.124379463493824, 0.028408339247107506, 0.3442523181438446, 0.15075638890266418, 0.08511755615472794, 0.32891392707824707, 0.12337944656610489, 0.05913665145635605, 0.11518532782793045, 0.28854820132255554, 0.0005498379468917847, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06821048259735107, 0.007578656077384949, 0.033511072397232056, 0.039627932012081146, 0.016393400728702545, 0.20925503969192505, 0.15704192221164703, 0.024064799770712852, 0.005696912761777639, 0.01698312722146511, 0.15042142570018768, 0.0017041407991200686, 0.016995420679450035, 0.005758653394877911, 0.015053601935505867, 0.12768876552581787, 0.007979520596563816, 0.05741023272275925, 0.14377589523792267, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05268644914031029, 0.018480738624930382, 0.006206580437719822, 0.01908770017325878, 0.009213676676154137, 0.012446015141904354, 0.2606332302093506, 0.15275397896766663, 0.004711512941867113, 0.01064901053905487, 0.00940486416220665, 0.00429189158603549, 0.014810611493885517, 0.012880465015769005, 0.15466143190860748, 0.25598737597465515, 0.03471918776631355, 0.08263758569955826, 0.03616967797279358, 0.0012629067059606314, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.017502065747976303, 0.09008979797363281, 0.045234303921461105, 0.04321402683854103, 0.014162504114210606, 0.2841097414493561, 0.10382679849863052, 0.4497845470905304, 0.042821191251277924, 0.03918898105621338, 0.06416238099336624, 0.04602029174566269, 0.2197093665599823, 0.07547488063573837, 0.13285692036151886, 0.29742351174354553, 0.10481993854045868, 0.07552393525838852, 0.008401650935411453, 0.3407011330127716, 0.028353586792945862, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02909473329782486, 0.05293780937790871, 0.025932423770427704, 0.061369478702545166, 0.12287095934152603, 0.12207728624343872, 0.20267462730407715, 0.3647293746471405, 0.036313559859991074, 0.028358493000268936, 0.054471470415592194, 0.007501897402107716, 0.10796680301427841, 0.05851392075419426, 0.12157665193080902, 0.17861823737621307, 0.07256677001714706, 0.1795390099287033, 0.04586997628211975, 0.27750420570373535, 0.0032322825863957405, 0.09472999721765518, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02889016829431057, 0.05256107077002525, 0.05110660940408707, 0.09513585269451141, 0.049980901181697845, 0.07343146204948425, 0.21190620958805084, 0.10279127210378647, 0.1787082403898239, 0.022944355383515358, 0.03947293758392334, 0.008258121088147163, 0.09723227471113205, 0.030062679201364517, 0.14898137748241425, 0.1281835287809372, 0.008169662207365036, 0.10209551453590393, 0.22781534492969513, 0.13339588046073914, 0.022249281406402588, 0.2580547630786896, 0.0071509419940412045, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.027054987847805023, 0.06796294450759888, 0.02347770519554615, 0.04540639370679855, 0.13579830527305603, 0.1935206949710846, 0.09281998127698898, 0.22921815514564514, 0.012567882426083088, 0.02752627059817314, 0.05939676612615585, 0.00633750855922699, 0.24427738785743713, 0.10302533209323883, 0.18246731162071228, 0.19490991532802582, 0.0105251120403409, 0.07082764059305191, 0.07746586948633194, 0.10047772526741028, 0.007984980009496212, 0.045915842056274414, 0.030714787542819977, 0.09154831618070602, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13923436403274536, 0.07431720942258835, 0.06541924923658371, 0.14132679998874664, 0.10506866127252579, 0.06156519800424576, 0.21440355479717255, 0.06509862840175629, 0.02759510651230812, 0.10144857317209244, 0.13265900313854218, 0.048845868557691574, 0.16166719794273376, 0.1116088330745697, 0.15105699002742767, 0.2116595059633255, 0.006228659767657518, 0.09237925708293915, 0.33000993728637695, 0.06037600710988045, 0.06468494236469269, 0.028822004795074463, 0.015993207693099976, 0.023504862561821938, 0.014777855016291142, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14352908730506897, 0.10288456827402115, 0.05261845886707306, 0.1541282832622528, 0.05661991983652115, 0.12065587192773819, 0.10697692632675171, 0.15951323509216309, 0.1055477038025856, 0.14385449886322021, 0.23090383410453796, 0.08539394289255142, 0.09938428550958633, 0.08322764188051224, 0.11896289885044098, 0.11546289920806885, 0.0627092570066452, 0.1015198826789856, 0.17440570890903473, 0.11644574254751205, 0.15138378739356995, 0.17151175439357758, 0.07174428552389145, 0.1994275599718094, 0.20994937419891357, 0.08254047483205795, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24387870728969574, 0.11191204935312271, 0.06428070366382599, 0.3038298189640045, 0.14750736951828003, 0.1200045570731163, 0.46686112880706787, 0.3116493225097656, 0.10273779183626175, 0.10795925557613373, 0.1416371762752533, 0.09460661560297012, 0.27618303894996643, 0.09149192273616791, 0.10828596353530884, 0.13584046065807343, 0.09117304533720016, 0.15590398013591766, 0.10968183726072311, 0.5585501790046692, 0.07535546272993088, 0.2762793302536011, 0.32588398456573486, 0.3246583938598633, 0.41251155734062195, 0.043567951768636703, 0.0185235645622015, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1039203479886055, 0.05052376165986061, 0.051659513264894485, 0.18036356568336487, 0.11265069991350174, 0.047071922570466995, 0.3453211784362793, 0.29340654611587524, 0.007079527713358402, 0.06730296462774277, 0.08055143058300018, 0.02563900128006935, 0.19650228321552277, 0.060815099626779556, 0.13184599578380585, 0.1674133688211441, 0.12648360431194305, 0.27492284774780273, 0.24355122447013855, 0.8769406676292419, 0.6096609234809875, 0.4704851806163788, 0.055198147892951965, 0.6140321493148804, 0.2705269455909729, 0.07450747489929199, 0.04471021145582199, 0.05369797348976135, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1947154402732849, 0.003113611601293087, 0.028957238420844078, 0.026910793036222458, 0.017121652141213417, 0.08169777691364288, 0.32467299699783325, 0.05661681666970253, 0.007502032909542322, 0.02869880571961403, 0.020577264949679375, 0.0070375413633883, 0.16551434993743896, 0.06083058565855026, 0.06852211803197861, 0.035074394196271896, 0.012203776277601719, 0.2713678479194641, 0.27628132700920105, 0.5399907231330872, 0.3242804706096649, 0.5765586495399475, 0.02925838902592659, 0.3159044086933136, 0.11935708671808243, 0.16010764241218567, 0.31936678290367126, 0.22831447422504425, 0.09149928390979767, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.018467016518115997, 0.004791167099028826, 0.015553582459688187, 0.021664531901478767, 0.025298617780208588, 0.1971224695444107, 0.13395515084266663, 0.1881190687417984, 0.05309745669364929, 0.018728721886873245, 0.018886514008045197, 0.023248562589287758, 0.008927382528781891, 0.03253133222460747, 0.130488321185112, 0.1354324370622635, 0.08839684724807739, 0.010535157285630703, 0.3809414505958557, 0.006101538427174091, 0.04204240441322327, 0.6714356541633606, 0.02054513990879059, 0.44751474261283875, 0.5217893123626709, 0.16833685338497162, 0.4138224124908447, 0.5945862531661987, 0.14406909048557281, 0.000551112403627485, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4018593430519104, 0.09619066119194031, 0.047895513474941254, 0.0887020081281662, 0.04670756310224533, 0.17605426907539368, 0.21604543924331665, 0.1403813511133194, 0.0010993692558258772, 0.07762767374515533, 0.0958188846707344, 0.1024225577712059, 0.06565871089696884, 0.04857100546360016, 0.1717240959405899, 0.26645413041114807, 0.038747917860746384, 0.15441381931304932, 0.6166976094245911, 0.04416924715042114, 0.07849516719579697, 0.41569313406944275, 0.018940549343824387, 0.18770581483840942, 0.11268321424722672, 0.0962471142411232, 0.028718965128064156, 0.019747000187635422, 0.011864973232150078, 0.07090434432029724, 0.02976600080728531, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.31909966468811035, 0.26355716586112976, 0.16833621263504028, 0.334572434425354, 0.18670302629470825, 0.11206400394439697, 0.46585598587989807, 0.15377958118915558, 0.014857469126582146, 0.07049962878227234, 0.1590365469455719, 0.09933225810527802, 0.23580892384052277, 0.09940709918737411, 0.11795931309461594, 0.26584282517433167, 0.03641113266348839, 0.24681606888771057, 0.03326011076569557, 0.5612249970436096, 0.11044078320264816, 0.038705065846443176, 0.07638699561357498, 0.20042885839939117, 0.41367095708847046, 0.16446417570114136, 0.05500950291752815, 0.0458536334335804, 0.038293108344078064, 0.05886702984571457, 0.005421455018222332, 0.03447017818689346, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3361136317253113, 0.18450267612934113, 0.10482683777809143, 0.3672127425670624, 0.09347432106733322, 0.06302808225154877, 0.17493662238121033, 0.11965186893939972, 0.06742112338542938, 0.13331438601016998, 0.26999813318252563, 0.03264465183019638, 0.07908355444669724, 0.09376725554466248, 0.11511774361133575, 0.052208781242370605, 0.10399425774812698, 0.2661847770214081, 0.06582632660865784, 0.5218088626861572, 0.41107869148254395, 0.18652401864528656, 0.10915308445692062, 0.2499890774488449, 0.21385571360588074, 0.11996328830718994, 0.2169666439294815, 0.17541900277137756, 0.34852319955825806, 0.29904353618621826, 0.3583068549633026, 0.0660485103726387, 0.0772518739104271, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.271436870098114, 0.16103556752204895, 0.09723401814699173, 0.3494490087032318, 0.1582973301410675, 0.11393263936042786, 0.41371721029281616, 0.2938876152038574, 0.08068472146987915, 0.08301044255495071, 0.11968915909528732, 0.07779402285814285, 0.24559125304222107, 0.07589462399482727, 0.1087639182806015, 0.1452419012784958, 0.08285138756036758, 0.20162978768348694, 0.10332676023244858, 0.7324197292327881, 0.1815183311700821, 0.27558720111846924, 0.41944485902786255, 0.4614993929862976, 0.7035390734672546, 0.14779764413833618, 0.07484183460474014, 0.09274464100599289, 0.1956741362810135, 0.4027537703514099, 0.17018413543701172, 0.15845544636249542, 0.03217604011297226, 0.027846908196806908, NaN, NaN, NaN, NaN, NaN, NaN], [0.1091129332780838, 0.08970999717712402, 0.08557470142841339, 0.23009367287158966, 0.13180004060268402, 0.0638015940785408, 0.31095248460769653, 0.2814267873764038, 0.0075759077444672585, 0.039292845875024796, 0.06780961900949478, 0.013560868799686432, 0.15987654030323029, 0.04180291295051575, 0.12740370631217957, 0.06803880631923676, 0.0777740478515625, 0.3149954080581665, 0.17862020432949066, 0.9274848103523254, 0.6797788739204407, 0.28538215160369873, 0.04841757193207741, 0.524702250957489, 0.33268001675605774, 0.06556227803230286, 0.08207366615533829, 0.08443650603294373, 0.19301387667655945, 0.68314129114151, 0.7843886613845825, 0.24039600789546967, 0.0983721911907196, 0.035574402660131454, 0.04086223617196083, NaN, NaN, NaN, NaN, NaN], [0.4568881392478943, 0.01152532733976841, 0.12744615972042084, 0.16633041203022003, 0.05682089552283287, 0.22013583779335022, 0.46718865633010864, 0.06831676512956619, 0.011846139095723629, 0.051503561437129974, 0.07631707936525345, 0.017341753467917442, 0.16032609343528748, 0.06682911515235901, 0.06364742666482925, 0.004222579766064882, 0.012189013883471489, 0.38177239894866943, 0.23501808941364288, 0.3822557032108307, 0.273560494184494, 0.28252631425857544, 0.039307549595832825, 0.41269388794898987, 0.3037600517272949, 0.1617780327796936, 0.33094146847724915, 0.37525615096092224, 0.1388353556394577, 0.8142803907394409, 0.5916069149971008, 0.18943282961845398, 0.08566068857908249, 0.11778654158115387, 0.1818830519914627, 0.04465563967823982, NaN, NaN, NaN, NaN], [0.0270079392939806, 0.003701634705066681, 0.024473953992128372, 0.035727839916944504, 0.031186459586024284, 0.22590965032577515, 0.1764952838420868, 0.1725662350654602, 0.06108492240309715, 0.017804577946662903, 0.01644762232899666, 0.018474329262971878, 0.0059660994447767735, 0.026993868872523308, 0.12890712916851044, 0.0780838280916214, 0.07355974614620209, 0.01093215774744749, 0.22770193219184875, 0.008550305850803852, 0.06503485888242722, 0.5060688257217407, 0.02145100012421608, 0.43843212723731995, 0.6872871518135071, 0.1969044953584671, 0.45010682940483093, 0.7415768504142761, 0.3103433847427368, 0.001054091495461762, 0.20113487541675568, 0.21400661766529083, 0.41673052310943604, 0.3260871469974518, 0.620118260383606, 0.12724098563194275, 0.0004952864837832749, NaN, NaN, NaN], [0.32686647772789, 0.10561588406562805, 0.10599718242883682, 0.08397059142589569, 0.05158340185880661, 0.22573474049568176, 0.19403943419456482, 0.08219113945960999, 0.0007591660832986236, 0.028280239552259445, 0.06139420345425606, 0.03943438082933426, 0.025857241824269295, 0.027251310646533966, 0.1435350626707077, 0.3314567506313324, 0.06341477483510971, 0.5618032217025757, 0.642646074295044, 0.27415919303894043, 0.23788774013519287, 0.38833677768707275, 0.08984735608100891, 0.42147237062454224, 0.6564009785652161, 0.2928015887737274, 0.1047874391078949, 0.1023104265332222, 0.06365151703357697, 0.39097070693969727, 0.14560170471668243, 0.23420175909996033, 0.08592629432678223, 0.02493405155837536, 0.011453422717750072, 0.006046658381819725, 0.1451905518770218, 0.005812718998640776, NaN, NaN], [0.21139562129974365, 0.21867576241493225, 0.17973701655864716, 0.29884445667266846, 0.19560806453227997, 0.11132223159074783, 0.28179141879081726, 0.10507592558860779, 0.014165982604026794, 0.04481332749128342, 0.1297360062599182, 0.07738039642572403, 0.2323194295167923, 0.09134778380393982, 0.12234959006309509, 0.21756824851036072, 0.03937938064336777, 0.3266570568084717, 0.05877631530165672, 0.5281912088394165, 0.11102446913719177, 0.03890432044863701, 0.10487684607505798, 0.2815292179584503, 0.4750865697860718, 0.3058159351348877, 0.11602579057216644, 0.12021853774785995, 0.06692790240049362, 0.1190272718667984, 0.019106050953269005, 0.21307361125946045, 0.15337608754634857, 0.06824280321598053, 0.040861621499061584, 0.032932352274656296, 0.052440475672483444, 0.005818615201860666, 0.0524408333003521, NaN], [0.2484172284603119, 0.2714419662952423, 0.13623963296413422, 0.33317360281944275, 0.14056812226772308, 0.16453251242637634, 0.23482279479503632, 0.2797185182571411, 0.08398787677288055, 0.13855448365211487, 0.19988903403282166, 0.12159004807472229, 0.21263501048088074, 0.1342880129814148, 0.11613592505455017, 0.21100056171417236, 0.13406150043010712, 0.10563220083713531, 0.15389345586299896, 0.10192565619945526, 0.07836726307868958, 0.22881029546260834, 0.05055452138185501, 0.24765580892562866, 0.48160815238952637, 0.2201593518257141, 0.1761431246995926, 0.21236160397529602, 0.20979638397693634, 0.10962515324354172, 0.09009265154600143, 0.0623038187623024, 0.17415094375610352, 0.13285446166992188, 0.11576873064041138, 0.10801524668931961, 0.0743527039885521, 0.03413216769695282, 0.027520645409822464, 0.06626196205615997]], [[0.0034671342000365257, 0.05013812705874443, 0.16192083060741425, 0.3595426082611084, 0.20735634863376617, 0.08139260113239288, 0.009979248046875, 0.05037669837474823, 0.0023427342530339956, 6.08037480560597e-05, 0.003484810469672084, 0.023961462080478668, 0.38460296392440796, 0.24992075562477112, 0.13989195227622986, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6699675917625427, 0.09382463991641998, 0.2939082980155945, 0.17940783500671387, 0.06414232403039932, 0.05161670595407486, 0.09315118193626404, 0.0025183490943163633, 0.0024716362822800875, 0.00784118939191103, 0.06077995523810387, 0.010742363519966602, 0.027031319215893745, 0.033606547862291336, 0.020909229293465614, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2646949589252472, 0.029353437945246696, 0.21451972424983978, 0.10881441831588745, 0.06597915291786194, 0.0030848400201648474, 0.011694483458995819, 0.021679535508155823, 0.002872215351089835, 0.013158812187612057, 0.002100167330354452, 6.679360376438126e-05, 0.004520595073699951, 0.019191764295101166, 0.15631338953971863, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.040224652737379074, 0.02035309188067913, 0.3179875612258911, 0.11730892956256866, 0.5032125115394592, 0.4173433780670166, 0.2045394331216812, 0.3468436896800995, 0.0142394183203578, 0.034110911190509796, 0.0166803989559412, 0.0005183254834264517, 0.014372344128787518, 0.013749183155596256, 0.07609989494085312, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0153636634349823, 0.002009550342336297, 0.5970484614372253, 0.5668097734451294, 0.03708057850599289, 0.030387206003069878, 0.003990367520600557, 0.00021067907800897956, 0.0006718098884448409, 0.004241611808538437, 0.01157804112881422, 0.0002699779870454222, 0.0015558624872937799, 0.0029094237834215164, 0.04601351544260979, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03574535250663757, 0.009626551531255245, 0.4402237832546234, 0.2294078767299652, 0.26443710923194885, 0.01504121907055378, 0.016090886667370796, 0.007329131942242384, 0.002309221774339676, 0.0030864060390740633, 0.0026519321836531162, 0.0004272839578334242, 0.0011082548880949616, 0.01614256016910076, 0.03275791555643082, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [6.553631828865036e-05, 0.000357702374458313, 0.08750326931476593, 0.01436514500528574, 0.006815748754888773, 0.6623476147651672, 0.0034670215100049973, 0.0015547194052487612, 0.00029766204534098506, 1.8653441657079384e-05, 0.0003687080170493573, 0.00015007570618763566, 0.0009929342195391655, 0.00030579339363612235, 0.0016504023224115372, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0004548979632090777, 7.145033305278048e-05, 0.025678247213363647, 0.00989772193133831, 0.007979623042047024, 0.6904858946800232, 0.04177143797278404, 0.0005172804230824113, 0.00045151059748604894, 9.678980859462172e-05, 0.0003766386944334954, 0.00020437331113498658, 0.0009936039568856359, 0.0004823105991818011, 0.001104293274693191, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.02770741656422615, 0.15481999516487122, 0.0164713803678751, 0.029219333082437515, 0.01727348566055298, 0.0033895254600793123, 0.08395758271217346, 0.08886045962572098, 0.06561290472745895, 0.23454923927783966, 0.01131775975227356, 0.00014876923523843288, 0.021633606404066086, 0.032435301691293716, 0.2441566288471222, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0002423129917588085, 0.0011915951035916805, 0.0022339578717947006, 0.006169029977172613, 0.0026169228367507458, 0.006970150861889124, 0.0023872333113104105, 0.020186979323625565, 0.5034035444259644, 0.061859097331762314, 0.01802009530365467, 0.08541904389858246, 0.11395227909088135, 0.12879255414009094, 0.06123032420873642, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0016445622313767672, 0.0006882954621687531, 0.0003155411686748266, 0.0014561355346813798, 0.0007120753289200366, 0.00010650769399944693, 0.0005508221802301705, 0.004306118004024029, 0.4519909620285034, 0.2298276424407959, 0.04858560487627983, 0.008956322446465492, 0.005770590156316757, 0.011063157580792904, 0.0306133683770895, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0032223593443632126, 0.0006265831179916859, 0.002176017500460148, 0.010606854222714901, 0.0010762742022052407, 6.259929068619385e-05, 0.0013370343949645758, 0.0014808439882472157, 0.030783534049987793, 0.7491747736930847, 0.34058046340942383, 0.00350938574410975, 0.02303031086921692, 0.0742756798863411, 0.006112673785537481, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.010601752437651157, 0.009935700334608555, 0.0694134384393692, 0.14514312148094177, 0.01701076701283455, 0.0001025431411108002, 0.003628269536420703, 0.007610301487147808, 0.1447119563817978, 0.2691461443901062, 0.7685887217521667, 0.06739932298660278, 0.05600086599588394, 0.567065417766571, 0.01997430995106697, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0020818221382796764, 0.006225256249308586, 0.007747206371277571, 0.02054281160235405, 0.00644321832805872, 0.00019787036580964923, 0.0007576930802315474, 0.0013290452770888805, 0.1748982071876526, 0.20870953798294067, 0.6057864427566528, 0.2165842056274414, 0.10265108197927475, 0.12960675358772278, 0.026959752663969994, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0929064005613327, 0.3412420153617859, 0.13197122514247894, 0.20421825349330902, 0.6308890581130981, 0.08085004985332489, 0.35388287901878357, 0.3416491150856018, 0.024628864601254463, 0.013967287726700306, 0.0762757882475853, 0.26007020473480225, 0.3328040838241577, 0.09019435197114944, 0.014360385946929455, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1659475415945053, 0.1821746528148651, 0.2680368423461914, 0.3257308900356293, 0.2135642170906067, 0.10952500998973846, 0.23729652166366577, 0.15246635675430298, 0.09328519552946091, 0.22413431107997894, 0.22322525084018707, 0.11237151175737381, 0.18681256473064423, 0.1572018712759018, 0.06837792694568634, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14290380477905273, 0.026570750400424004, 0.14845344424247742, 0.26635152101516724, 0.12476544827222824, 0.1522083431482315, 0.287058562040329, 0.16522644460201263, 0.21008911728858948, 0.3761942982673645, 0.12840349972248077, 0.0757022351026535, 0.39944273233413696, 0.379029244184494, 0.1911974847316742, 0.0702696219086647, 0.2507307231426239, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00885845348238945, 0.005625984165817499, 0.0020030708983540535, 0.005766861606389284, 0.001782223698683083, 0.004346099682152271, 0.014438317157328129, 0.010037342086434364, 0.0175970196723938, 0.0067982920445501804, 0.003056151093915105, 0.005088370759040117, 0.0035549686290323734, 0.002117584692314267, 0.17935973405838013, 0.028418319299817085, 0.003963488154113293, 0.4144974946975708, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04871530085802078, 0.2322341799736023, 0.043161727488040924, 0.046935759484767914, 0.04166096821427345, 0.048159919679164886, 0.2838554382324219, 0.5679410696029663, 0.17445935308933258, 0.05776107683777809, 0.14550535380840302, 0.04300517588853836, 0.2332015484571457, 0.28196635842323303, 0.4675023853778839, 0.13786309957504272, 0.03506092354655266, 0.02415982447564602, 0.10726116597652435, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03277377411723137, 0.28776609897613525, 0.0018310850718989968, 0.006392122711986303, 0.0034063432831317186, 0.0006021481240168214, 0.02006486989557743, 0.09552518278360367, 0.02804744802415371, 0.060428690165281296, 0.004742977675050497, 0.018782831728458405, 0.016696294769644737, 0.023774143308401108, 0.16262513399124146, 0.011229841969907284, 0.008138949982821941, 0.04613415151834488, 0.2518063187599182, 0.013397655449807644, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.006045958958566189, 0.0958699956536293, 0.007954242639243603, 0.011606856249272823, 0.004544504452496767, 0.010406642220914364, 0.011899203062057495, 0.07300186902284622, 0.002370428293943405, 0.012239865958690643, 0.020374998450279236, 0.012496876530349255, 0.024265890941023827, 0.0274967048317194, 0.1423870474100113, 0.0016812672838568687, 0.012760624289512634, 0.002261990448459983, 0.2769384980201721, 0.03090759925544262, 0.0014064738061279058, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.008809137158095837, 0.13565093278884888, 0.03191651031374931, 0.0483417883515358, 0.028707973659038544, 0.039296794682741165, 0.018359076231718063, 0.07145766168832779, 0.13921810686588287, 0.01646633818745613, 0.06145479157567024, 0.028490308672189713, 0.056069642305374146, 0.13838331401348114, 0.19134177267551422, 0.11822758615016937, 0.07095540314912796, 0.030966516584157944, 0.03516996279358864, 0.2070395052433014, 0.02684318646788597, 0.2317354679107666, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.39272594451904297, 0.39728477597236633, 0.32111606001853943, 0.41796234250068665, 0.15293559432029724, 0.04586965963244438, 0.16940170526504517, 0.022719532251358032, 0.14239482581615448, 0.5121501088142395, 0.19016578793525696, 0.06530822068452835, 0.29211705923080444, 0.14742477238178253, 0.11553633958101273, 0.23311708867549896, 0.026411496102809906, 0.011159970425069332, 0.03808103874325752, 0.017219573259353638, 0.006694006733596325, 0.001702688867226243, 0.009211051277816296, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.009060109965503216, 0.08736205101013184, 0.03623565658926964, 0.046393588185310364, 0.04293924570083618, 0.049119193106889725, 0.018734706565737724, 0.10957584530115128, 0.04821338504552841, 0.02008068934082985, 0.029284991323947906, 0.015971768647432327, 0.05779576674103737, 0.21830672025680542, 0.21264111995697021, 0.1427604705095291, 0.06787170469760895, 0.04101337492465973, 0.04024908319115639, 0.2669386863708496, 0.04579312726855278, 0.07587221264839172, 0.10059545934200287, 0.18715938925743103, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02833615615963936, 0.24966742098331451, 0.06237170845270157, 0.03993965685367584, 0.10454770177602768, 0.019859671592712402, 0.03772445023059845, 0.19178973138332367, 0.012827831320464611, 0.03533304110169411, 0.024230163544416428, 0.054630037397146225, 0.032379381358623505, 0.08906079828739166, 0.17152637243270874, 0.059837497770786285, 0.10673120617866516, 0.06554628908634186, 0.047321293503046036, 0.26084935665130615, 0.05379262939095497, 0.09055614471435547, 0.09319713711738586, 0.334230899810791, 0.23545128107070923, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.015255320817232132, 0.21888743340969086, 0.1253896951675415, 0.08362822234630585, 0.12500159442424774, 0.02890017069876194, 0.03405824303627014, 0.07477163523435593, 0.0229325033724308, 0.01863025315105915, 0.044950928539037704, 0.0560457706451416, 0.04699615016579628, 0.08650227636098862, 0.1548503190279007, 0.06699422001838684, 0.48348554968833923, 0.10470042377710342, 0.2643885016441345, 0.49639153480529785, 0.11732041090726852, 0.061902400106191635, 0.1530170738697052, 0.11711295694112778, 0.23237623274326324, 0.09402092546224594, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.011826024390757084, 0.10608652234077454, 0.04723645746707916, 0.057715099304914474, 0.03395959734916687, 0.028910892084240913, 0.011586843058466911, 0.050380002707242966, 0.030421555042266846, 0.00583301018923521, 0.015118762850761414, 0.014350258745253086, 0.01606619358062744, 0.025515934452414513, 0.18496018648147583, 0.050390250980854034, 0.2627623975276947, 0.057036180049180984, 0.10587681084871292, 0.22481703758239746, 0.07078704982995987, 0.028480585664510727, 0.47086307406425476, 0.03990349546074867, 0.16108965873718262, 0.02393723465502262, 0.06960758566856384, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.015032858587801456, 0.5077551603317261, 0.07541441917419434, 0.08020945638418198, 0.10545077919960022, 0.2137133628129959, 0.01040775515139103, 0.09528981149196625, 0.09038985520601273, 0.012094871141016483, 0.025733938440680504, 0.06706724315881729, 0.03145073354244232, 0.09538157284259796, 0.34148263931274414, 0.29633763432502747, 0.1570599228143692, 0.07358378916978836, 0.08321648091077805, 0.01657349243760109, 0.02100137248635292, 0.019902318716049194, 0.5162196755409241, 0.03987365961074829, 0.018146652728319168, 0.026169516146183014, 0.00614600395783782, 0.07103840261697769, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.32250380516052246, 0.7984310388565063, 0.3962976634502411, 0.40014326572418213, 0.3554738759994507, 0.47898975014686584, 0.10853014886379242, 0.20243746042251587, 0.127571240067482, 0.2699570655822754, 0.16473528742790222, 0.08001074939966202, 0.03713205084204674, 0.14643853902816772, 0.4229389429092407, 0.1833065152168274, 0.0826280415058136, 0.06509751826524734, 0.017351830378174782, 0.08598462492227554, 0.028223805129528046, 0.03195580840110779, 0.045467328280210495, 0.041934747248888016, 0.016390223056077957, 0.05298775061964989, 0.05077003315091133, 0.2718433141708374, 0.04039132222533226, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.023898553103208542, 0.03448064997792244, 0.007101188413798809, 0.020377272740006447, 0.09085186570882797, 0.008504875935614109, 0.01689869724214077, 0.021393392235040665, 0.03013733960688114, 0.004040753003209829, 0.000672544410917908, 0.0007860396872274578, 0.0003324192948639393, 0.0003073772240895778, 0.13160185515880585, 0.09722712635993958, 0.09857381135225296, 0.2290657013654709, 0.162257120013237, 0.3208743929862976, 0.7083525657653809, 0.08285251259803772, 0.05820265784859657, 0.14296579360961914, 0.06442547589540482, 0.3963678479194641, 0.1963234394788742, 0.13509824872016907, 0.0551372766494751, 0.1773844212293625, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.025859396904706955, 0.29733914136886597, 0.09033425897359848, 0.06196272000670433, 0.10889838635921478, 0.14661002159118652, 0.034964289516210556, 0.07059973478317261, 0.007527152542024851, 0.007617437280714512, 0.006072000600397587, 0.0492180734872818, 0.0069811418652534485, 0.011496509425342083, 0.22706106305122375, 0.1786596029996872, 0.03035295568406582, 0.011360704898834229, 0.0041356864385306835, 0.02253635786473751, 0.032254207879304886, 0.05765725299715996, 0.06512543559074402, 0.26075252890586853, 0.14487245678901672, 0.06064848601818085, 0.02561355009675026, 0.06785233318805695, 0.08367668837308884, 0.11658230423927307, 0.21664968132972717, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.014849718660116196, 0.1462036818265915, 0.11065799742937088, 0.06219353526830673, 0.08005399256944656, 0.016894571483135223, 0.010269397869706154, 0.02562439627945423, 0.009192260913550854, 0.009821194224059582, 0.015785057097673416, 0.019254932180047035, 0.01222837995737791, 0.011684795841574669, 0.16154925525188446, 0.02336198277771473, 0.027563903480768204, 0.02503703534603119, 0.002219978952780366, 0.024155667051672935, 0.005802824627608061, 0.011775066144764423, 0.03527237847447395, 0.0438326895236969, 0.16127318143844604, 0.07829897105693817, 0.04636809974908829, 0.16168944537639618, 0.17395752668380737, 0.5116502642631531, 0.11367138475179672, 0.24585914611816406, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01973692700266838, 0.11480830609798431, 0.07148479670286179, 0.05237298831343651, 0.0777522474527359, 0.019268590956926346, 0.01592963933944702, 0.01235677395015955, 0.06519288569688797, 0.019938096404075623, 0.03185376524925232, 0.0271891038864851, 0.01742159202694893, 0.040164995938539505, 0.1837940812110901, 0.14312313497066498, 0.6151867508888245, 0.2511911392211914, 0.34089455008506775, 0.21357816457748413, 0.06974375993013382, 0.04017443582415581, 0.4436698257923126, 0.0627409890294075, 0.029346130788326263, 0.06214871257543564, 0.07426106929779053, 0.37162381410598755, 0.1908751130104065, 0.2730017304420471, 0.09601876139640808, 0.07787502557039261, 0.1985486000776291, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.006014276295900345, 0.07228019088506699, 0.029915854334831238, 0.031709808856248856, 0.01963544264435768, 0.01660715602338314, 0.00532315531745553, 0.03606380149722099, 0.029185649007558823, 0.0046777487732470036, 0.01710142381489277, 0.013257446698844433, 0.01389795821160078, 0.02201540581882, 0.16183340549468994, 0.05929486081004143, 0.1356429159641266, 0.08288607001304626, 0.1716676652431488, 0.17707081139087677, 0.11502664536237717, 0.023076828569173813, 0.41179341077804565, 0.03153251111507416, 0.08080360293388367, 0.03793509677052498, 0.0956316813826561, 0.40457794070243835, 0.3355584144592285, 0.2116643786430359, 0.2117510586977005, 0.0911363810300827, 0.13469243049621582, 0.08244834095239639, NaN, NaN, NaN, NaN, NaN, NaN], [0.008549164049327374, 0.34144893288612366, 0.03957316279411316, 0.03764811158180237, 0.04039980471134186, 0.07271253317594528, 0.00613941578194499, 0.04612124711275101, 0.0911136344075203, 0.008750539273023605, 0.01715807057917118, 0.03749352693557739, 0.024577608332037926, 0.06848984956741333, 0.2503378689289093, 0.34530380368232727, 0.14280815422534943, 0.08469259738922119, 0.20386184751987457, 0.018106382340192795, 0.025206930935382843, 0.03376462310552597, 0.665645956993103, 0.06945709139108658, 0.030968131497502327, 0.031062953174114227, 0.015101979486644268, 0.10170532017946243, 0.03453005850315094, 0.05652596056461334, 0.028510402888059616, 0.036133769899606705, 0.04489430412650108, 0.010548176243901253, 0.07425779104232788, NaN, NaN, NaN, NaN, NaN], [0.1472499966621399, 0.4703251123428345, 0.2558133602142334, 0.283985435962677, 0.21470209956169128, 0.17662864923477173, 0.07007063925266266, 0.06038873642683029, 0.20766907930374146, 0.26984694600105286, 0.16889145970344543, 0.27114859223365784, 0.03473396599292755, 0.13903996348381042, 0.2962591350078583, 0.21361097693443298, 0.09641434252262115, 0.0472431480884552, 0.030436551198363304, 0.12823571264743805, 0.024378983303904533, 0.03781319037079811, 0.04478050768375397, 0.04302188381552696, 0.031242409721016884, 0.06916327774524689, 0.08240062743425369, 0.2609483301639557, 0.04106062278151512, 0.01303931511938572, 0.014160559512674809, 0.011109860613942146, 0.034855347126722336, 0.10407929867506027, 0.21024775505065918, 0.08525354415178299, NaN, NaN, NaN, NaN], [0.020655758678913116, 0.020222418010234833, 0.006879583932459354, 0.019070995971560478, 0.07609020173549652, 0.006032301113009453, 0.015974652022123337, 0.01717195473611355, 0.05267442390322685, 0.004277344327419996, 0.0005684247589670122, 0.0007490122807212174, 0.0002994663082063198, 0.0002370573638472706, 0.12958088517189026, 0.056013792753219604, 0.04104574769735336, 0.13420559465885162, 0.14404895901679993, 0.30753612518310547, 0.5552563667297363, 0.06356479972600937, 0.02527950517833233, 0.09324341267347336, 0.03306487947702408, 0.2522013187408447, 0.14255186915397644, 0.09901494532823563, 0.06439376622438431, 0.10042564570903778, 0.43083739280700684, 0.20968028903007507, 0.35324180126190186, 0.2700602114200592, 0.23262809216976166, 0.11776822060346603, 0.14138048887252808, NaN, NaN, NaN], [0.009374987334012985, 0.23445867002010345, 0.05258592590689659, 0.020285839214920998, 0.024131227284669876, 0.0535256564617157, 0.01552440132945776, 0.032435644418001175, 0.006646827794611454, 0.005740212742239237, 0.005195626523345709, 0.07125341892242432, 0.0043562185019254684, 0.01014760322868824, 0.17807012796401978, 0.1699744164943695, 0.02438814751803875, 0.00377153092995286, 0.0020952692721039057, 0.017941365018486977, 0.009907160885632038, 0.04197421669960022, 0.08005423098802567, 0.16825814545154572, 0.08759146183729172, 0.037892259657382965, 0.02378804422914982, 0.12696562707424164, 0.21072204411029816, 0.039158232510089874, 0.12900760769844055, 0.018357207998633385, 0.09957201033830643, 0.024237502366304398, 0.12091250717639923, 0.2524404227733612, 0.044468626379966736, 0.19958341121673584, NaN, NaN], [0.018758203834295273, 0.11843696236610413, 0.09101122617721558, 0.0610043928027153, 0.06165887042880058, 0.012400476261973381, 0.011786350980401039, 0.021215293556451797, 0.014211799949407578, 0.011016220785677433, 0.02130991406738758, 0.02418670989573002, 0.015627985820174217, 0.013993974775075912, 0.14536960422992706, 0.016944430768489838, 0.011726072989404202, 0.017351148650050163, 0.0028529188130050898, 0.013441222719848156, 0.005811003036797047, 0.010734970681369305, 0.020825698971748352, 0.04144507274031639, 0.0777476355433464, 0.07330787181854248, 0.0589311420917511, 0.1305314600467682, 0.09686601907014847, 0.49986732006073, 0.09861493855714798, 0.24486178159713745, 0.2709232568740845, 0.08328418433666229, 0.1665872186422348, 0.2741791903972626, 0.5570544600486755, 0.09308093041181564, 0.18428745865821838, NaN], [0.03985379636287689, 0.12957410514354706, 0.13386031985282898, 0.10592924803495407, 0.09455320239067078, 0.03913174197077751, 0.052976641803979874, 0.03812992200255394, 0.11070051789283752, 0.042073190212249756, 0.05433963984251022, 0.058929286897182465, 0.03380222246050835, 0.05054538697004318, 0.1317562311887741, 0.043635401874780655, 0.027883753180503845, 0.11735352873802185, 0.09225393831729889, 0.11462916433811188, 0.1478782296180725, 0.04645288363099098, 0.049018505960702896, 0.08540874719619751, 0.16189652681350708, 0.081883005797863, 0.13365384936332703, 0.17616337537765503, 0.16547891497612, 0.3400772511959076, 0.14388780295848846, 0.2768324613571167, 0.1609276533126831, 0.18515954911708832, 0.2950800061225891, 0.32982173562049866, 0.4366631507873535, 0.3681013882160187, 0.34051525592803955, 0.05319627374410629]], [[0.014275058172643185, 0.006687531713396311, 0.3026585280895233, 0.06917963922023773, 0.2396276444196701, 0.6229325532913208, 0.15904799103736877, 0.13992713391780853, 0.10272591561079025, 0.6685669422149658, 0.22624024748802185, 0.09492585808038712, 0.40837499499320984, 0.2735627591609955, 0.011893448419868946, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.021194536238908768, 0.020265106111764908, 0.1736137419939041, 0.08712188154459, 0.3174395263195038, 0.3545694649219513, 0.3640749752521515, 0.11553992331027985, 0.3069344758987427, 0.7487083673477173, 0.45964598655700684, 0.41950592398643494, 0.6157799363136292, 0.47228363156318665, 0.04039919748902321, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.008898869156837463, 0.002019912237301469, 0.021509699523448944, 0.0182319525629282, 0.07474909722805023, 0.02385670319199562, 0.013716273009777069, 0.008799813687801361, 0.3437807857990265, 0.008914400823414326, 0.012629772536456585, 0.10342472046613693, 0.0370708666741848, 0.023541903123259544, 0.18654775619506836, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01223641075193882, 0.003142833709716797, 0.006001354195177555, 0.003996475599706173, 0.0579916350543499, 0.01896491087973118, 0.01948327198624611, 0.013184066861867905, 0.30560916662216187, 0.015957718715071678, 0.016950437799096107, 0.06207568570971489, 0.044481322169303894, 0.01894378289580345, 0.19150091707706451, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.003971019294112921, 0.0012432326329872012, 0.005908531602472067, 0.0021760377567261457, 0.002044213702902198, 0.01004379615187645, 0.01574278064072132, 0.026324355974793434, 0.4105670154094696, 0.05117517337203026, 0.02775881439447403, 0.023424910381436348, 0.009920927695930004, 0.011210974305868149, 0.16597995162010193, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.007421860471367836, 0.006305157672613859, 0.011464249342679977, 0.020268600434064865, 0.025753991678357124, 0.031131377443671227, 0.03418951481580734, 0.0052986773662269115, 0.5788748264312744, 0.46168622374534607, 0.07252157479524612, 0.06022901460528374, 0.017210712656378746, 0.04054110497236252, 0.15131165087223053, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.001541785546578467, 0.0008907613810151815, 0.004846525378525257, 0.001811343478038907, 0.0069520194083452225, 0.008084121160209179, 0.021458715200424194, 0.02802192233502865, 0.3832707405090332, 0.25552085041999817, 0.014592574909329414, 0.01065820176154375, 0.012523604556918144, 0.010731800459325314, 0.22416816651821136, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.004116748925298452, 0.0016883857315406203, 0.014749680645763874, 0.00869818776845932, 0.01003838051110506, 0.007631313521414995, 0.02068890631198883, 0.027104953303933144, 0.13497500121593475, 0.6378710865974426, 0.10288828611373901, 0.0942029282450676, 0.028772620484232903, 0.05935161933302879, 0.21764545142650604, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.06222981959581375, 0.01881357654929161, 0.00486758491024375, 0.015509632416069508, 0.0009378677350468934, 0.004574655555188656, 0.005093523766845465, 0.0076056248508393764, 0.02507362887263298, 0.02107030339539051, 0.007815904915332794, 0.010442771948873997, 0.011698074638843536, 0.006942160427570343, 0.31572407484054565, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.01727244071662426, 0.009210732765495777, 0.005953751504421234, 0.0013454181607812643, 0.005081892944872379, 0.04435739293694496, 0.006434922106564045, 0.0007962443050928414, 0.0007702711154706776, 0.16453301906585693, 0.5625144839286804, 0.34227296710014343, 0.6355522871017456, 0.6161591410636902, 0.02771596610546112, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.12786830961704254, 0.008172453381121159, 0.0017843057867139578, 0.004017683211714029, 0.007877650670707226, 0.0018398476531729102, 0.01566770300269127, 0.0026914728805422783, 0.0035052604507654905, 0.0037441153544932604, 0.011492998339235783, 0.10472051054239273, 0.01954079605638981, 0.025050928816199303, 0.24727097153663635, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1465907245874405, 0.037033673375844955, 0.013877127319574356, 0.00413108617067337, 0.00966043584048748, 0.02326187677681446, 0.04576379433274269, 0.010370912030339241, 0.05009477958083153, 0.002161832293495536, 0.012562266550958157, 0.08835282921791077, 0.018735390156507492, 0.07781965285539627, 0.21298982203006744, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.018177246674895287, 0.009594686329364777, 0.010616189800202847, 0.003939185757189989, 0.020018288865685463, 0.006944165099412203, 0.014553648419678211, 0.014575640670955181, 0.031773608177900314, 0.0201406329870224, 0.008282337337732315, 0.02822018228471279, 0.008926213718950748, 0.030271533876657486, 0.18345791101455688, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.029857823625206947, 0.018949948251247406, 0.0061294399201869965, 0.002908851485699415, 0.00919707678258419, 0.00952958408743143, 0.01205661240965128, 0.00758303003385663, 0.05086279660463333, 0.007759919855743647, 0.006360263098031282, 0.02717713639140129, 0.006157578434795141, 0.027468249201774597, 0.21562480926513672, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.035946138203144073, 0.021175134927034378, 0.025809520855545998, 0.0228139478713274, 0.02454732172191143, 0.008901212364435196, 0.01817207969725132, 0.024075007066130638, 0.042662542313337326, 0.10151555389165878, 0.03429628908634186, 0.025050567463040352, 0.015684176236391068, 0.028640326112508774, 0.23519039154052734, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.038382355123758316, 0.16509199142456055, 0.03795319423079491, 0.018471574410796165, 0.017937200143933296, 0.20822547376155853, 0.036850690841674805, 0.07025959342718124, 0.026183662936091423, 0.008891633711755276, 0.011525453999638557, 0.06559614092111588, 0.10240377485752106, 0.05705304443836212, 0.19186913967132568, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18736660480499268, 0.12802250683307648, 0.06000450998544693, 0.07085607945919037, 0.02492770366370678, 0.13308653235435486, 0.01379183866083622, 0.01460492704063654, 0.018005041405558586, 0.18972568213939667, 0.18918126821517944, 0.05261359363794327, 0.08419474214315414, 0.039842329919338226, 0.12843605875968933, 0.1755252629518509, 0.00892956368625164, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.003212069161236286, 0.04924406483769417, 0.010131219401955605, 0.0015629208646714687, 0.009065762162208557, 0.04507109895348549, 0.003221129300072789, 0.07382506877183914, 0.0011923180427402258, 0.004047631751745939, 0.006328214425593615, 0.012952281162142754, 0.0641837865114212, 0.02541324496269226, 0.1715373396873474, 0.18403629958629608, 0.12486936897039413, 0.01289399154484272, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.002438034862279892, 0.0007996301865205169, 0.10929557681083679, 0.030698396265506744, 0.007961505092680454, 0.21520712971687317, 0.0018748894799500704, 0.0015670642023906112, 0.00039643081254325807, 0.0017966092564165592, 0.010619523003697395, 0.0026792865246534348, 0.0035868084523826838, 0.001077426946721971, 0.003137440187856555, 0.07995349168777466, 0.1140136644244194, 0.16089488565921783, 0.271826833486557, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04913554713129997, 0.023452362045645714, 0.16805477440357208, 0.2746557891368866, 0.369334876537323, 0.025402046740055084, 0.03595297038555145, 0.27975642681121826, 0.005478397477418184, 0.044800374656915665, 0.028408128768205643, 0.025396348908543587, 0.1202942430973053, 0.22760754823684692, 0.12602998316287994, 0.19368642568588257, 0.20833823084831238, 0.38513559103012085, 0.0724099725484848, 0.026710418984293938, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0008230121457017958, 0.006709535606205463, 0.005090394522994757, 0.005009432788938284, 0.0009200142812915146, 0.002589132636785507, 0.003276216797530651, 0.011904137209057808, 0.0009605096420273185, 0.0016532291192561388, 0.001647727913223207, 0.0010296034161001444, 0.00474548852071166, 0.004530362784862518, 0.14385877549648285, 0.2920932173728943, 0.20408804714679718, 0.47836723923683167, 0.009784400463104248, 0.41401228308677673, 0.0022880665492266417, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.011407818645238876, 0.11073090881109238, 0.11066732555627823, 0.07063236832618713, 0.2326628416776657, 0.057718440890312195, 0.005228970665484667, 0.12933272123336792, 0.010014788247644901, 0.0034599530044943094, 0.015450170263648033, 0.004393222741782665, 0.010258005000650883, 0.00790967233479023, 0.16524673998355865, 0.2459677904844284, 0.013399376533925533, 0.165635347366333, 0.0016970435390248895, 0.00861914549022913, 0.0019094902090728283, 0.006659353617578745, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.024886149913072586, 0.019822845235466957, 0.050577834248542786, 0.042761147022247314, 0.013624369166791439, 0.03171548992395401, 0.03447520360350609, 0.057101696729660034, 0.018126925453543663, 0.012612801045179367, 0.056599393486976624, 0.005686976481229067, 0.022324958816170692, 0.021004129201173782, 0.18438492715358734, 0.1659669429063797, 0.3024148941040039, 0.4638516902923584, 0.19814886152744293, 0.06386706978082657, 0.37022748589515686, 0.096834197640419, 0.004976118449121714, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.012148641981184483, 0.047028496861457825, 0.07792042940855026, 0.1455426812171936, 0.3985011875629425, 0.08270914107561111, 0.0031603944953531027, 0.07123681157827377, 0.020226983353495598, 0.005742877256125212, 0.009367674589157104, 0.007002389058470726, 0.013849785551428795, 0.006732230074703693, 0.14449873566627502, 0.23605915904045105, 0.015010624192655087, 0.29689958691596985, 0.002272083656862378, 0.02557971514761448, 0.04829570651054382, 0.03933914750814438, 0.012097989208996296, 0.005491157062351704, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.029934342950582504, 0.04287242144346237, 0.10493571311235428, 0.10647397488355637, 0.01039193756878376, 0.1410648375749588, 0.06155749782919884, 0.08983614295721054, 0.05490254610776901, 0.038721270859241486, 0.021267540752887726, 0.05536682903766632, 0.019229264929890633, 0.008436290547251701, 0.15105655789375305, 0.2229652851819992, 0.011020033620297909, 0.07613904774188995, 0.00492003234103322, 0.11613531410694122, 0.12462546676397324, 0.03799906745553017, 0.029671484604477882, 0.022334527224302292, 0.003809461137279868, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.009979508817195892, 0.08308109641075134, 0.026161497458815575, 0.023276647552847862, 0.0017319537000730634, 0.056630972772836685, 0.012614267878234386, 0.041058339178562164, 0.026752248406410217, 0.01169703807681799, 0.011314285919070244, 0.007283498533070087, 0.05053415521979332, 0.019243547692894936, 0.16277745366096497, 0.30055463314056396, 0.03860635682940483, 0.08235271275043488, 0.12519411742687225, 0.07496307790279388, 0.24307869374752045, 0.02970520593225956, 0.043270040303468704, 0.01804984174668789, 0.008444367907941341, 0.04573319852352142, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04712976887822151, 0.24274323880672455, 0.053717970848083496, 0.06948067992925644, 0.009206406772136688, 0.0471884086728096, 0.010105792433023453, 0.05801715701818466, 0.01891178824007511, 0.07684698700904846, 0.07729421555995941, 0.042662668973207474, 0.10241091996431351, 0.038032110780477524, 0.15563422441482544, 0.361846923828125, 0.0072926427237689495, 0.07028269022703171, 0.038334887474775314, 0.02117738127708435, 0.035939738154411316, 0.03011121228337288, 0.01985063962638378, 0.03699057549238205, 0.0448327511548996, 0.07655268162488937, 0.03217002749443054, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.009955390356481075, 0.06358544528484344, 0.028598172590136528, 0.04170457646250725, 0.01363537646830082, 0.011423949152231216, 0.003101062262430787, 0.04170127958059311, 0.01145926769822836, 0.01274544931948185, 0.020664334297180176, 0.15329574048519135, 0.20515742897987366, 0.07666952162981033, 0.13521607220172882, 0.18510019779205322, 0.0857149139046669, 0.2959531545639038, 0.10870446264743805, 0.034602705389261246, 0.04019882157444954, 0.02403290942311287, 0.05409723520278931, 0.04566982761025429, 0.19149497151374817, 0.23549742996692657, 0.074503093957901, 0.01255789864808321, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.006747167091816664, 0.006801524665206671, 0.007903891615569592, 0.00237295706756413, 0.0009535709978081286, 0.0006887177005410194, 0.0011137888068333268, 0.0005580680444836617, 0.004365934059023857, 0.0043631866574287415, 0.004836279433220625, 0.0014166004257276654, 0.1882382482290268, 0.04424351081252098, 0.006875277496874332, 0.03710656613111496, 0.054964251816272736, 0.037898506969213486, 0.3724515438079834, 0.058691613376140594, 0.03363177552819252, 0.06933214515447617, 0.05247700959444046, 0.15643684566020966, 0.589249849319458, 0.349843829870224, 0.29659491777420044, 0.2287619560956955, 0.05358140170574188, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0040101236663758755, 0.00047035442548803985, 0.0008357138140127063, 0.009736553765833378, 0.00025759977870620787, 2.9679033104912378e-05, 0.008525178767740726, 0.0036214631982147694, 0.0009930779924616218, 0.0008531230851076543, 0.0029921825043857098, 7.93160234024981e-06, 6.746472354279831e-05, 0.0017078705132007599, 0.13162609934806824, 0.2688547670841217, 0.1434442549943924, 0.18350595235824585, 0.07485228031873703, 0.0647219642996788, 0.04773847386240959, 0.14254990220069885, 0.03905782103538513, 0.2126167118549347, 0.24802155792713165, 0.30339401960372925, 0.17472584545612335, 0.03891041502356529, 0.02338952198624611, 0.026767900213599205, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.021027032285928726, 0.04388788715004921, 0.07337366044521332, 0.13240061700344086, 0.005691900383681059, 0.08179081231355667, 0.010154702700674534, 0.019539857283234596, 0.013572044670581818, 0.03972425311803818, 0.14196330308914185, 0.0491810142993927, 0.029326222836971283, 0.024830663576722145, 0.1775946319103241, 0.1340402513742447, 0.12347351759672165, 0.42842522263526917, 0.0631304681301117, 0.06392616778612137, 0.1770109236240387, 0.11116458475589752, 0.04706185683608055, 0.09571156650781631, 0.3872493505477905, 0.5415271520614624, 0.14801958203315735, 0.013348261825740337, 0.016769861802458763, 0.019784821197390556, 0.012107723392546177, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.020570920780301094, 0.07008225470781326, 0.05771828070282936, 0.10093566030263901, 0.0037175160832703114, 0.10588520765304565, 0.008791210129857063, 0.07720224559307098, 0.037850137799978256, 0.016810759902000427, 0.0763774886727333, 0.06772230565547943, 0.10185997188091278, 0.02133399061858654, 0.1501101702451706, 0.3128407299518585, 0.02314484678208828, 0.20690661668777466, 0.0038596922531723976, 0.10119188576936722, 0.375572144985199, 0.077932208776474, 0.16011959314346313, 0.07805528491735458, 0.020400837063789368, 0.2237216979265213, 0.1006372720003128, 0.022764090448617935, 0.005061473231762648, 0.0205483790487051, 0.0018506759079173207, 0.001139476546086371, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.027059482410550117, 0.22707954049110413, 0.13379518687725067, 0.08346803486347198, 0.011664706282317638, 0.1994924694299698, 0.013729198835790157, 0.07924864441156387, 0.10303384810686111, 0.02253318764269352, 0.06352351605892181, 0.13561668992042542, 0.3492315113544464, 0.13069112598896027, 0.12187084555625916, 0.5802629590034485, 0.17577120661735535, 0.22907592356204987, 0.3224048614501953, 0.21584153175354004, 0.3719359040260315, 0.08852899819612503, 0.18978306651115417, 0.06894023716449738, 0.008546161465346813, 0.34136468172073364, 0.44251179695129395, 0.07915834337472916, 0.27557075023651123, 0.0915302038192749, 0.0036887326277792454, 0.0038842300418764353, 0.015524323098361492, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.038929592818021774, 0.2334582358598709, 0.12089657783508301, 0.17347271740436554, 0.023068996146321297, 0.04853734001517296, 0.008499456569552422, 0.0867975577712059, 0.02351396717131138, 0.04524386301636696, 0.12492679059505463, 0.06575564295053482, 0.10587428510189056, 0.055128976702690125, 0.1414995789527893, 0.5194967985153198, 0.010316978208720684, 0.10247951745986938, 0.03023943491280079, 0.02351299114525318, 0.05376119539141655, 0.03751303628087044, 0.02858700230717659, 0.03933052346110344, 0.026450933888554573, 0.16396890580654144, 0.08825679868459702, 0.01957540772855282, 0.02957809716463089, 0.0652899444103241, 0.003373907646164298, 0.007670924998819828, 0.004321575630456209, 0.024295708164572716, NaN, NaN, NaN, NaN, NaN, NaN], [0.011872883886098862, 0.08469298481941223, 0.054403409361839294, 0.08831894397735596, 0.02684788778424263, 0.021699469536542892, 0.0027920349966734648, 0.05190650746226311, 0.006984782870858908, 0.008844600059092045, 0.02751598134636879, 0.22613400220870972, 0.15431185066699982, 0.06476734578609467, 0.1412026435136795, 0.2508450150489807, 0.1962328553199768, 0.3596697747707367, 0.1504865288734436, 0.029224414378404617, 0.0663013905286789, 0.043777331709861755, 0.06269483268260956, 0.06556038558483124, 0.2250475436449051, 0.35171735286712646, 0.22191122174263, 0.018188640475273132, 0.026326660066843033, 0.017122289165854454, 0.0037187051493674517, 0.024730468168854713, 0.035062648355960846, 0.09351257234811783, 0.011442800983786583, NaN, NaN, NaN, NaN, NaN], [0.015115483663976192, 0.08628259599208832, 0.023322032764554024, 0.012461238540709019, 0.0028755213133990765, 0.010226217098534107, 0.0010302395094186068, 0.002081838669255376, 0.003762529231607914, 0.013111302629113197, 0.0290949996560812, 0.013309521600604057, 0.22778895497322083, 0.05992528051137924, 0.00796937569975853, 0.007168593350797892, 0.033368390053510666, 0.00873665139079094, 0.16062632203102112, 0.028196215629577637, 0.02527499757707119, 0.06866460293531418, 0.0198657363653183, 0.1544157713651657, 0.2752910256385803, 0.14698350429534912, 0.1242247000336647, 0.13061578571796417, 0.010920656844973564, 0.0055906628258526325, 0.006986986380070448, 0.030699225142598152, 0.36674854159355164, 0.2189747393131256, 0.2510429620742798, 0.04264682158827782, NaN, NaN, NaN, NaN], [0.0057023135013878345, 0.0003758604871109128, 0.0009645622340030968, 0.01432577334344387, 0.00027227052487432957, 3.7724938010796905e-05, 0.007459490094333887, 0.0037525389343500137, 0.001061747083440423, 0.0008801367366686463, 0.0023195864632725716, 8.150678695528768e-06, 4.0667833673069254e-05, 0.001007204526104033, 0.12961283326148987, 0.317547470331192, 0.16016888618469238, 0.1976199448108673, 0.10644932836294174, 0.09830258786678314, 0.07801979035139084, 0.301817923784256, 0.05034731701016426, 0.32512444257736206, 0.2241876721382141, 0.4657731354236603, 0.2891538441181183, 0.08093820512294769, 0.06031876429915428, 0.06730521470308304, 0.14267991483211517, 0.289673775434494, 0.1076083853840828, 0.2949788272380829, 0.0365237332880497, 0.015645001083612442, 0.03993191570043564, NaN, NaN, NaN], [0.017900969833135605, 0.026770949363708496, 0.15903817117214203, 0.31877970695495605, 0.014844128862023354, 0.10845804959535599, 0.00868347566574812, 0.015460771508514881, 0.008762474171817303, 0.01190071552991867, 0.07999671250581741, 0.053750935941934586, 0.013735906220972538, 0.020958656445145607, 0.15606556832790375, 0.17233391106128693, 0.22507980465888977, 0.300968736410141, 0.03457535058259964, 0.06539295613765717, 0.2556630074977875, 0.12555503845214844, 0.08745130896568298, 0.10011813044548035, 0.13041436672210693, 0.501103937625885, 0.14929187297821045, 0.03132137656211853, 0.02265048772096634, 0.03383776918053627, 0.006481703836470842, 0.011523596942424774, 0.35894638299942017, 0.1662973165512085, 0.034177642315626144, 0.02702290564775467, 0.036704160273075104, 0.014952532015740871, NaN, NaN], [0.022256335243582726, 0.07135839015245438, 0.07359576225280762, 0.12423767894506454, 0.006224590353667736, 0.13500085473060608, 0.008429165929555893, 0.08156562596559525, 0.02983916364610195, 0.013062523677945137, 0.10225346684455872, 0.04065772891044617, 0.06899033486843109, 0.012502058409154415, 0.13831046223640442, 0.4115316569805145, 0.042032964527606964, 0.21366682648658752, 0.010602481663227081, 0.11737099289894104, 0.5779745578765869, 0.13523340225219727, 0.2636784315109253, 0.170937180519104, 0.020469455048441887, 0.3112620711326599, 0.17165400087833405, 0.044973500072956085, 0.006653682328760624, 0.053596071898937225, 0.008654352277517319, 0.002382548525929451, 0.02675137296319008, 0.09427332878112793, 0.01890433207154274, 0.002222384326159954, 0.018390605226159096, 0.0013299400452524424, 0.0009657714981585741, NaN], [0.016071150079369545, 0.06728275120258331, 0.025518205016851425, 0.023689931258559227, 0.0069392030127346516, 0.04150809720158577, 0.00898416806012392, 0.016712933778762817, 0.005143268499523401, 0.020111138001084328, 0.03020956739783287, 0.01359627302736044, 0.018198341131210327, 0.01637156493961811, 0.1379418522119522, 0.38502925634384155, 0.1563987135887146, 0.13578397035598755, 0.1404726654291153, 0.14828255772590637, 0.28480827808380127, 0.15350891649723053, 0.09994281083345413, 0.06321649998426437, 0.030282480642199516, 0.13266463577747345, 0.1722954362630844, 0.07113035768270493, 0.024887708947062492, 0.016665330156683922, 0.03949398547410965, 0.020136239007115364, 0.01368448045104742, 0.09379612654447556, 0.030771953985095024, 0.011002926155924797, 0.007083212956786156, 0.009242233820259571, 0.007993990555405617, 0.018528543412685394]], [[0.29903000593185425, 0.5539957880973816, 0.06723504513502121, 0.06922264397144318, 0.12363186478614807, 0.04431891441345215, 0.10694187879562378, 0.08094406872987747, 0.15170463919639587, 0.05897890776395798, 0.026665056124329567, 0.04277891665697098, 0.011532573029398918, 0.016366619616746902, 0.08233406394720078, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.030788322910666466, 0.06814564764499664, 0.1441766321659088, 0.42568475008010864, 0.23481200635433197, 0.09723259508609772, 0.20801249146461487, 0.2833361029624939, 0.12989479303359985, 0.09075285494327545, 0.02217184565961361, 0.10632100701332092, 0.07123817503452301, 0.18399499356746674, 0.11842577904462814, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.21215111017227173, 0.2570435404777527, 0.03298918902873993, 0.11753708124160767, 0.2531988024711609, 0.2834656238555908, 0.13087181746959686, 0.14389817416667938, 0.06408312171697617, 0.023736948147416115, 0.043677639216184616, 0.007582403719425201, 0.08098249137401581, 0.042930904775857925, 0.09848955273628235, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.24232596158981323, 0.4370230436325073, 0.27921250462532043, 0.32216426730155945, 0.14763100445270538, 0.1446210741996765, 0.041608523577451706, 0.05782362446188927, 0.03667302429676056, 0.015881532803177834, 0.09886573255062103, 0.0007486737449653447, 0.022804880514740944, 0.01436265092343092, 0.04328664019703865, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0417991504073143, 0.06808368116617203, 0.22980956733226776, 0.06044253334403038, 0.09120408445596695, 0.3664403557777405, 0.01738058589398861, 0.026107804849743843, 0.16878005862236023, 0.007388730999082327, 0.6907519698143005, 0.00283504044637084, 0.004864559043198824, 0.017621232196688652, 0.04920867085456848, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07025078684091568, 0.08007846027612686, 0.18737106025218964, 0.08649075031280518, 0.14398247003555298, 0.03926409035921097, 0.10999412834644318, 0.10028164088726044, 0.2733333110809326, 0.07497494667768478, 0.6277027726173401, 0.03760387748479843, 0.07242996245622635, 0.04469411447644234, 0.0635850802063942, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.18292218446731567, 0.29889917373657227, 0.16216641664505005, 0.041324593126773834, 0.08738134056329727, 0.03374062106013298, 0.10780933499336243, 0.1685270518064499, 0.3661736249923706, 0.13795819878578186, 0.7607439160346985, 0.022037923336029053, 0.11896573007106781, 0.017960727214813232, 0.09792909026145935, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.29104405641555786, 0.7119240164756775, 0.16990531980991364, 0.02345188707113266, 0.15646961331367493, 0.008449066430330276, 0.06418811529874802, 0.018176060169935226, 0.3091927766799927, 0.08911041170358658, 0.3005200922489166, 0.04236089810729027, 0.2996547222137451, 0.08733220398426056, 0.07523740082979202, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.046947941184043884, 0.14375551044940948, 0.004344047512859106, 0.0067795743234455585, 0.02948000282049179, 0.08397668600082397, 0.06400846689939499, 0.18865461647510529, 0.023663662374019623, 0.08527978509664536, 0.02815503440797329, 0.04117048531770706, 0.5833349823951721, 0.0677085593342781, 0.23153413832187653, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08349642902612686, 0.4532567262649536, 0.004409583285450935, 0.009004302322864532, 0.007938031107187271, 0.13749390840530396, 0.1858609914779663, 0.31525370478630066, 0.018453413620591164, 0.12712040543556213, 0.04680929332971573, 0.12408707290887833, 0.13737666606903076, 0.12311573326587677, 0.142713725566864, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05042501911520958, 0.07026762515306473, 0.0020696106366813183, 0.010109566152095795, 0.07710029184818268, 0.05610239878296852, 0.05948542803525925, 0.19247274100780487, 0.001940111513249576, 0.05155838653445244, 0.04620450362563133, 0.20989066362380981, 0.485702246427536, 0.4166657328605652, 0.18102103471755981, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.09080760926008224, 0.09187275916337967, 0.012195594608783722, 0.021634280681610107, 0.019499676302075386, 0.09054076671600342, 0.11008334904909134, 0.23214302957057953, 0.0423310361802578, 0.034868963062763214, 0.06751228123903275, 0.049237679690122604, 0.03915484994649887, 0.08995199203491211, 0.1941523253917694, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0706457570195198, 0.10473088920116425, 0.039385173469781876, 0.02697153575718403, 0.04372800514101982, 0.06655491143465042, 0.23491710424423218, 0.19935868680477142, 0.036273516714572906, 0.06345809996128082, 0.020782677456736565, 0.12393849343061447, 0.05726756155490875, 0.041495081037282944, 0.15982753038406372, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.039186086505651474, 0.11076691001653671, 0.03891725465655327, 0.009549588896334171, 0.01825849525630474, 0.051163915544748306, 0.1146436408162117, 0.1649821698665619, 0.03586947172880173, 0.06679365783929825, 0.09092967957258224, 0.14827685058116913, 0.10948126018047333, 0.10746686905622482, 0.1515202671289444, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.14541134238243103, 0.05313154682517052, 0.01991144008934498, 0.08764121681451797, 0.014597749337553978, 0.03937898576259613, 0.04872390255331993, 0.04689335823059082, 0.04558950290083885, 0.051970891654491425, 0.02520112879574299, 0.022838978096842766, 0.00921469647437334, 0.00801294855773449, 0.21471147239208221, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.029921628534793854, 0.09876842796802521, 0.1324968934059143, 0.09236511588096619, 0.02831152267754078, 0.08077768236398697, 0.03118293546140194, 0.1750149130821228, 0.015778981149196625, 0.07032441347837448, 0.22269371151924133, 0.07579661160707474, 0.029184984043240547, 0.053061336278915405, 0.18562854826450348, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07805982232093811, 0.05365234240889549, 0.2842547595500946, 0.2606758773326874, 0.21293140947818756, 0.02651267871260643, 0.08033362030982971, 0.07913534343242645, 0.17101624608039856, 0.12522375583648682, 0.14315897226333618, 0.16815446317195892, 0.0695369690656662, 0.13316825032234192, 0.19111928343772888, 0.17860974371433258, 0.0018437139224261045, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11272483319044113, 0.11636882275342941, 0.45685258507728577, 0.0910579040646553, 0.3091263473033905, 0.12632955610752106, 0.1822080761194229, 0.18498732149600983, 0.6353387832641602, 0.08394157886505127, 0.3285849094390869, 0.4818887710571289, 0.08592816442251205, 0.3495768904685974, 0.07449600845575333, 0.20284786820411682, 0.0034877806901931763, 0.08334594964981079, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2834128737449646, 0.1102365031838417, 0.1840669959783554, 0.5708534121513367, 0.3157653212547302, 0.041008107364177704, 0.038309745490550995, 0.03211268410086632, 0.6102551817893982, 0.20786605775356293, 0.21116787195205688, 0.10018377006053925, 0.04653669148683548, 0.17929011583328247, 0.11314841359853745, 0.1494244486093521, 0.3379342555999756, 0.0649241954088211, 0.006597604602575302, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5993789434432983, 0.0908532664179802, 0.49218761920928955, 0.41100576519966125, 0.18825526535511017, 0.4342217445373535, 0.12116678059101105, 0.10673660039901733, 0.822167158126831, 0.4385586380958557, 0.6995345950126648, 0.18085956573486328, 0.1357179582118988, 0.2864921987056732, 0.034255724400281906, 0.2969810962677002, 0.005403619725257158, 0.054099179804325104, 0.0006044544279575348, 0.009600944817066193, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.858432412147522, 0.34460219740867615, 0.7778953909873962, 0.7743141651153564, 0.4405529797077179, 0.4761039614677429, 0.6155950427055359, 0.06873662024736404, 0.7323919534683228, 0.7086790204048157, 0.6720118522644043, 0.45794978737831116, 0.1628962755203247, 0.4249861538410187, 0.040913816541433334, 0.32280662655830383, 0.01735025830566883, 0.15535852313041687, 0.00028658873634412885, 0.016427762806415558, 0.001579301548190415, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04546767473220825, 0.0383436344563961, 0.10268200188875198, 0.20100316405296326, 0.185649111866951, 0.08432896435260773, 0.060354892164468765, 0.07717668265104294, 0.3201402723789215, 0.04503992572426796, 0.088813915848732, 0.3990366756916046, 0.1564548909664154, 0.08066049963235855, 0.11440145969390869, 0.016787199303507805, 0.10643576830625534, 0.24800433218479156, 0.4802894592285156, 0.03762362524867058, 0.06816797703504562, 0.10676699876785278, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21178147196769714, 0.043018583208322525, 0.1065564677119255, 0.10858221352100372, 0.05675008147954941, 0.06700197607278824, 0.12675313651561737, 0.058651700615882874, 0.18508696556091309, 0.05493801832199097, 0.037313126027584076, 0.19010567665100098, 0.07823225855827332, 0.034572359174489975, 0.16783590614795685, 0.22070105373859406, 0.03063296526670456, 0.12860903143882751, 0.04803713783621788, 0.06528759002685547, 0.3172104060649872, 0.012414618395268917, 0.008628717623651028, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.053469568490982056, 0.03894811123609543, 0.06651152670383453, 0.10646583139896393, 0.08985435962677002, 0.07578439265489578, 0.03395741805434227, 0.09802807122468948, 0.190333291888237, 0.07748086005449295, 0.07400990277528763, 0.6643930077552795, 0.07830479741096497, 0.07947986572980881, 0.11464671790599823, 0.0170818492770195, 0.2921580374240875, 0.24774892628192902, 0.2979756295681, 0.16657015681266785, 0.03825104981660843, 0.39123743772506714, 0.0541624091565609, 0.01715947687625885, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1680978536605835, 0.06724530458450317, 0.16071708500385284, 0.2987021803855896, 0.11997595429420471, 0.007637033239006996, 0.05953739956021309, 0.06456195563077927, 0.07405640929937363, 0.11493658274412155, 0.07269633561372757, 0.12183233350515366, 0.019239120185375214, 0.0931614562869072, 0.15387272834777832, 0.06952934712171555, 0.09443160146474838, 0.3155873417854309, 0.2511345446109772, 0.20146684348583221, 0.17959536612033844, 0.500001072883606, 0.3407229483127594, 0.15127938985824585, 0.026401039212942123, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09433168172836304, 0.05311369523406029, 0.44581180810928345, 0.2857709527015686, 0.11141614615917206, 0.04973546415567398, 0.10592624545097351, 0.0732862576842308, 0.26435965299606323, 0.07302475720643997, 0.17637307941913605, 0.06760746240615845, 0.052111051976680756, 0.29667070508003235, 0.11431443691253662, 0.12491581588983536, 0.08139167726039886, 0.045777399092912674, 0.07585746794939041, 0.05243801325559616, 0.09790124744176865, 0.17415514588356018, 0.44996151328086853, 0.13761505484580994, 0.06580806523561478, 0.1016187071800232, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07687122374773026, 0.10929025709629059, 0.4687592387199402, 0.20397132635116577, 0.26744040846824646, 0.03514130413532257, 0.033296968787908554, 0.08783485740423203, 0.22074763476848602, 0.08713625371456146, 0.12920482456684113, 0.05166565254330635, 0.07679110020399094, 0.17419996857643127, 0.1387287825345993, 0.03772348165512085, 0.0006561332265846431, 0.04040418565273285, 0.23337695002555847, 0.0037602160591632128, 0.1251135915517807, 0.07994246482849121, 0.0032252452801913023, 0.044697076082229614, 0.05314825102686882, 0.16676445305347443, 0.42838534712791443, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.061203911900520325, 0.12594261765480042, 0.353413462638855, 0.22131817042827606, 0.41015592217445374, 0.11432977020740509, 0.010031531564891338, 0.048355478793382645, 0.27572426199913025, 0.07773520797491074, 0.2322542816400528, 0.1527126431465149, 0.05797232687473297, 0.09810248017311096, 0.16366761922836304, 0.008380687795579433, 0.11938491463661194, 0.03761400282382965, 0.10612092912197113, 0.004111893475055695, 0.07536520808935165, 0.06150262430310249, 0.010061400011181831, 0.01712355576455593, 0.026476707309484482, 0.05440329760313034, 0.37643373012542725, 0.12204637378454208, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10230414569377899, 0.03857935592532158, 0.05230129137635231, 0.14396332204341888, 0.09251677989959717, 0.03541665896773338, 0.005624003708362579, 0.014271721243858337, 0.042375415563583374, 0.13543996214866638, 0.061749108135700226, 0.00788076315075159, 0.1602918803691864, 0.07564403861761093, 0.09375559538602829, 0.0973815768957138, 0.1330094188451767, 0.2356250286102295, 0.23801013827323914, 0.16962124407291412, 0.3808935284614563, 0.19062454998493195, 0.12487400323152542, 0.4241224527359009, 0.1858355700969696, 0.1843334436416626, 0.17186462879180908, 0.1674181967973709, 0.03679514676332474, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.705120861530304, 0.026186510920524597, 0.8528315424919128, 0.8252069354057312, 0.24319231510162354, 0.07270172983407974, 0.09487330913543701, 0.07207771390676498, 0.4722364544868469, 0.7067926526069641, 0.8624283075332642, 0.07399676740169525, 0.0075901346281170845, 0.016478050500154495, 0.12560917437076569, 0.28161293268203735, 0.39586660265922546, 0.35408592224121094, 0.26687130331993103, 0.036089953035116196, 0.12106626480817795, 0.05175312981009483, 0.6374836564064026, 0.06537415832281113, 0.01867927983403206, 0.03261437267065048, 0.05161871388554573, 0.026679201051592827, 0.0063977655954658985, 0.0581950880587101, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.27840110659599304, 0.06363435834646225, 0.3689763844013214, 0.33064448833465576, 0.25749024748802185, 0.1453908383846283, 0.03645810857415199, 0.00836147554218769, 0.3977815508842468, 0.41805213689804077, 0.17756043374538422, 0.05318059027194977, 0.011340576224029064, 0.020938394591212273, 0.05934957042336464, 0.052721865475177765, 0.30848002433776855, 0.24953237175941467, 0.2790854275226593, 0.7654650807380676, 0.6871634125709534, 0.13210926949977875, 0.673875629901886, 0.04467727988958359, 0.018614191561937332, 0.08283445239067078, 0.0906965509057045, 0.06073237210512161, 0.12131030112504959, 0.06997358053922653, 0.3489122688770294, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17816129326820374, 0.10609658807516098, 0.17893879115581512, 0.28182876110076904, 0.15060719847679138, 0.03372456133365631, 0.04276707395911217, 0.050946421921253204, 0.04137968271970749, 0.16634012758731842, 0.16395889222621918, 0.24548840522766113, 0.05229371041059494, 0.09448723495006561, 0.12793652713298798, 0.03943483531475067, 0.28613966703414917, 0.07243800908327103, 0.8744964599609375, 0.029915155842900276, 0.331167072057724, 0.4079437255859375, 0.5431530475616455, 0.3259604275226593, 0.1150238886475563, 0.3324905335903168, 0.44221389293670654, 0.2450132817029953, 0.12577538192272186, 0.11014749854803085, 0.1900990903377533, 0.042790502309799194, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14424489438533783, 0.0705854520201683, 0.24214811623096466, 0.24549053609371185, 0.19939330220222473, 0.02639644220471382, 0.021373553201556206, 0.024115193635225296, 0.08405331522226334, 0.14685925841331482, 0.15661610662937164, 0.06219787895679474, 0.032059792429208755, 0.09036684036254883, 0.15146715939044952, 0.06558705866336823, 0.020870981737971306, 0.007642277050763369, 0.028054187074303627, 0.010532653890550137, 0.10334379225969315, 0.12033270299434662, 0.1911371499300003, 0.30930495262145996, 0.04741071164608002, 0.06516209989786148, 0.09313901513814926, 0.24243950843811035, 0.15116305649280548, 0.09231718629598618, 0.47254911065101624, 0.053373783826828, 0.18162642419338226, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06650430709123611, 0.10705426335334778, 0.3146411180496216, 0.1647443175315857, 0.23945462703704834, 0.035643309354782104, 0.026562364771962166, 0.09605439007282257, 0.19827118515968323, 0.1037423387169838, 0.14283734560012817, 0.08165161311626434, 0.07012972235679626, 0.11072988063097, 0.13417953252792358, 0.017124762758612633, 0.00014164860476739705, 0.01482362300157547, 0.13952724635601044, 0.0008921221597120166, 0.07150562852621078, 0.037848807871341705, 0.0009583857608959079, 0.0160027127712965, 0.01657933183014393, 0.09754330664873123, 0.3402610719203949, 0.02766183763742447, 0.011668790131807327, 0.019427720457315445, 0.01879642903804779, 0.06977814435958862, 0.23379765450954437, 0.41046860814094543, NaN, NaN, NaN, NaN, NaN, NaN], [0.06460674107074738, 0.10897383838891983, 0.18354696035385132, 0.20187535881996155, 0.38844820857048035, 0.04722803831100464, 0.010622762143611908, 0.04332485795021057, 0.31279584765434265, 0.11892355233430862, 0.20366235077381134, 0.1460915356874466, 0.041410893201828, 0.060890424996614456, 0.16885291039943695, 0.0033047832548618317, 0.043024010956287384, 0.009507044218480587, 0.05758155509829521, 0.0012058177962899208, 0.04777836054563522, 0.038867104798555374, 0.0027761561796069145, 0.008453112095594406, 0.011027430184185505, 0.021058345213532448, 0.3453521430492401, 0.05058252438902855, 0.004837945103645325, 0.0014179014833644032, 0.06873936206102371, 0.10687354952096939, 0.21186815202236176, 0.44615596532821655, 0.10872229933738708, NaN, NaN, NaN, NaN, NaN], [0.08445128798484802, 0.07278266549110413, 0.017734743654727936, 0.12906457483768463, 0.17354236543178558, 0.01439378596842289, 0.0032682251185178757, 0.009051240049302578, 0.02403325028717518, 0.17859239876270294, 0.05114053934812546, 0.026160510256886482, 0.17188863456249237, 0.059929899871349335, 0.12745818495750427, 0.05260666832327843, 0.09784732013940811, 0.08957145363092422, 0.40504154562950134, 0.2393025904893875, 0.37446328997612, 0.33926665782928467, 0.06915906071662903, 0.28494811058044434, 0.18951286375522614, 0.21801336109638214, 0.2963850796222687, 0.09700386226177216, 0.02254888415336609, 0.016780056059360504, 0.3380737006664276, 0.17247304320335388, 0.15711140632629395, 0.27414536476135254, 0.12462585419416428, 0.05461693927645683, NaN, NaN, NaN, NaN], [0.6940725445747375, 0.016104217618703842, 0.8427497148513794, 0.8075915575027466, 0.2572270333766937, 0.04667792096734047, 0.07690176367759705, 0.06650352478027344, 0.4641934931278229, 0.7403572797775269, 0.892522931098938, 0.08286882191896439, 0.00509345019236207, 0.009769911877810955, 0.1252693384885788, 0.4168609082698822, 0.5786882042884827, 0.4795728027820587, 0.4880480170249939, 0.07741907238960266, 0.22295767068862915, 0.10229793190956116, 0.7397969365119934, 0.09120289236307144, 0.02111845649778843, 0.040493883192539215, 0.06478337198495865, 0.029333919286727905, 0.01266437117010355, 0.08807221800088882, 0.12442159652709961, 0.019878262653946877, 0.02248454838991165, 0.045759230852127075, 0.02396523579955101, 0.002620323793962598, 0.04143214225769043, NaN, NaN, NaN], [0.47638654708862305, 0.08160793781280518, 0.2188907116651535, 0.3983159363269806, 0.3041192293167114, 0.0773146003484726, 0.041229549795389175, 0.00785501953214407, 0.20719125866889954, 0.6323855519294739, 0.1790589690208435, 0.15920953452587128, 0.005728188902139664, 0.011172757484018803, 0.10331764072179794, 0.05813424289226532, 0.29987069964408875, 0.06046860292553902, 0.2948205769062042, 0.6036045551300049, 0.4684220552444458, 0.10851431638002396, 0.5970842242240906, 0.03630568087100983, 0.009022231213748455, 0.034897517412900925, 0.044963937252759933, 0.06918716430664062, 0.06464210897684097, 0.027029458433389664, 0.39741793274879456, 0.1858920007944107, 0.0860959067940712, 0.03553689271211624, 0.03651457652449608, 0.07401836663484573, 0.02850046567618847, 0.457316130399704, NaN, NaN], [0.3162515461444855, 0.12029282748699188, 0.1898643672466278, 0.3138664960861206, 0.22235795855522156, 0.03812789171934128, 0.07994988560676575, 0.07006566971540451, 0.06856126338243484, 0.2470276951789856, 0.2142392098903656, 0.4667101502418518, 0.07071195542812347, 0.09391427785158157, 0.11791101843118668, 0.011862307786941528, 0.06274299323558807, 0.019264375790953636, 0.7077140212059021, 0.009838010184466839, 0.08938813954591751, 0.2665976285934448, 0.21134285628795624, 0.19931168854236603, 0.029879093170166016, 0.11873869597911835, 0.2187809944152832, 0.10740162432193756, 0.03893040865659714, 0.02778119407594204, 0.17118902504444122, 0.03705315291881561, 0.41107529401779175, 0.3035467863082886, 0.1782693862915039, 0.062172479927539825, 0.04369974508881569, 0.43116021156311035, 0.04090215638279915, NaN], [0.15722334384918213, 0.11492010205984116, 0.22595097124576569, 0.17283931374549866, 0.11246844381093979, 0.07424511015415192, 0.1308857947587967, 0.1509532928466797, 0.12219540029764175, 0.14498494565486908, 0.13763099908828735, 0.16327989101409912, 0.12245305627584457, 0.21428720653057098, 0.12265608459711075, 0.13294808566570282, 0.07747184485197067, 0.06700501590967178, 0.24500344693660736, 0.07035010308027267, 0.06088097393512726, 0.15465889871120453, 0.22422827780246735, 0.20946520566940308, 0.06346394866704941, 0.1416163444519043, 0.10671631991863251, 0.07756247371435165, 0.14874279499053955, 0.2551397681236267, 0.18877547979354858, 0.07302238047122955, 0.24805422127246857, 0.1228112131357193, 0.08095405995845795, 0.12022056430578232, 0.20888803899288177, 0.1654488444328308, 0.07207347452640533, 0.12261014431715012]], [[0.009874092414975166, 0.0475393682718277, 0.0700187012553215, 0.05995699018239975, 0.023110831156373024, 0.04304451867938042, 0.02397323027253151, 0.09104450792074203, 0.13320927321910858, 0.0718994140625, 0.16378211975097656, 0.06306017935276031, 0.03516274318099022, 0.06407153606414795, 0.1927335411310196, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.007679122034460306, 0.008519956842064857, 0.023641018196940422, 0.036320336163043976, 0.005810021422803402, 0.002834178740158677, 0.01027101743966341, 0.005131446290761232, 0.05288401618599892, 0.022729018703103065, 0.02885960415005684, 0.007142365910112858, 0.005423326510936022, 0.00592823838815093, 0.23125353455543518, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.17363575100898743, 0.08529574424028397, 0.018747013062238693, 0.09323837608098984, 0.07366655766963959, 0.2784116566181183, 0.6226999759674072, 0.6422466039657593, 0.18433590233325958, 0.44911590218544006, 0.07703087478876114, 0.23628254234790802, 0.37835898995399475, 0.3362680971622467, 0.10061702132225037, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.039354946464300156, 0.028671007603406906, 0.0009692042949609458, 0.010166235268115997, 0.003592043649405241, 0.024686597287654877, 0.0576656274497509, 0.10543617606163025, 0.069565050303936, 0.23999209702014923, 0.0370241142809391, 0.07099387794733047, 0.08031197637319565, 0.0629396140575409, 0.19831009209156036, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.07821620255708694, 0.07413192838430405, 0.008470119908452034, 0.005837618373334408, 0.016890503466129303, 0.34118980169296265, 0.6424257159233093, 0.5736639499664307, 0.18751046061515808, 0.08286380022764206, 0.013973995111882687, 0.16452431678771973, 0.6265572905540466, 0.24633896350860596, 0.03771306574344635, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08601168543100357, 0.11519530415534973, 0.00501672737300396, 0.0384475477039814, 0.0009856059914454818, 0.020220156759023666, 0.4602939486503601, 0.41334664821624756, 0.011432202532887459, 0.039776530116796494, 0.004202698357403278, 0.012451107613742352, 0.012797003611922264, 0.0109980758279562, 0.22371669113636017, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05821564793586731, 0.2493630200624466, 0.017187682911753654, 0.007334073074162006, 0.002277297666296363, 0.012770043686032295, 0.014771709218621254, 0.06810285151004791, 0.008148171938955784, 0.093966543674469, 0.03078475221991539, 0.016961626708507538, 0.009818210266530514, 0.005369590129703283, 0.2805846929550171, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0315314382314682, 0.006441309116780758, 0.005187691655009985, 0.0023020647931843996, 0.001103160553611815, 0.0010285694152116776, 0.0036586276255548, 0.0034369472414255142, 0.02540425956249237, 0.018933216109871864, 0.011261656880378723, 0.014689027331769466, 0.0047272746451199055, 0.003173592034727335, 0.27608010172843933, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.052501752972602844, 0.03902341425418854, 0.022159013897180557, 0.15980832278728485, 0.04565480723977089, 0.04961955174803734, 0.10487794876098633, 0.03556728735566139, 0.011893571354448795, 0.350600004196167, 0.8153157234191895, 0.696418821811676, 0.19642634689807892, 0.7945331335067749, 0.025074943900108337, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.008775658905506134, 0.0231929961591959, 0.001974506536498666, 0.02221933752298355, 0.002016729209572077, 0.03464629501104355, 0.020560195669531822, 0.015741808339953423, 0.024821357801556587, 0.03194829449057579, 0.062133170664310455, 0.009445058181881905, 0.008440939709544182, 0.031038939952850342, 0.24359388649463654, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.15448324382305145, 0.15535393357276917, 0.0009195139864459634, 0.02347545325756073, 0.010745828039944172, 0.05933469906449318, 0.0886014774441719, 0.09891750663518906, 0.008176282048225403, 0.17814745008945465, 0.04613054543733597, 0.10348650068044662, 0.06132601201534271, 0.10257216542959213, 0.2144334316253662, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1637454628944397, 0.3587695062160492, 0.013175190426409245, 0.027070751413702965, 0.009701711125671864, 0.027045298367738724, 0.06057014688849449, 0.08674251288175583, 0.018084047362208366, 0.012978773564100266, 0.04984384402632713, 0.0746963769197464, 0.21545591950416565, 0.18275731801986694, 0.18403297662734985, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04016833007335663, 0.03071952983736992, 0.0073937661945819855, 0.044594794511795044, 0.005693770945072174, 0.007929249666631222, 0.19023852050304413, 0.12198647856712341, 0.00967123731970787, 0.05747445672750473, 0.006795276887714863, 0.006636326666921377, 0.014849998988211155, 0.02297961339354515, 0.1823122203350067, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08359953761100769, 0.14515268802642822, 0.009139984846115112, 0.10055579245090485, 0.007817201316356659, 0.06191832944750786, 0.24591712653636932, 0.26670339703559875, 0.008127851411700249, 0.05132465437054634, 0.011226493865251541, 0.020721180364489555, 0.025672290474176407, 0.06137499585747719, 0.19538666307926178, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.004038439132273197, 0.01158715970814228, 0.012492671608924866, 0.008604439906775951, 0.0044732466340065, 0.001471644383855164, 0.003622728632763028, 0.005392232909798622, 0.024040954187512398, 0.002572751836851239, 0.011896335519850254, 0.00655994052067399, 0.004419950768351555, 0.0023605322930961847, 0.2578853368759155, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03995227441191673, 0.02612248808145523, 0.09039098769426346, 0.04685363546013832, 0.14171013236045837, 0.3046724796295166, 0.08713044226169586, 0.11726538836956024, 0.3945818245410919, 0.03867875412106514, 0.060879118740558624, 0.3211958110332489, 0.1562168449163437, 0.1954476237297058, 0.12928469479084015, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.138319730758667, 0.1925395429134369, 0.06914161890745163, 0.1830926090478897, 0.22252067923545837, 0.24239645898342133, 0.2738734483718872, 0.3115195333957672, 0.287569522857666, 0.12556934356689453, 0.047479670494794846, 0.1859251707792282, 0.015966184437274933, 0.050888173282146454, 0.04287213087081909, 0.04818185046315193, 0.30147239565849304, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.059622667729854584, 0.19761067628860474, 0.019807182252407074, 0.02911451645195484, 0.11472073942422867, 0.03754669055342674, 0.08183436095714569, 0.09122617542743683, 0.10595303028821945, 0.094895139336586, 0.022252719849348068, 0.087751105427742, 0.015402892604470253, 0.02668953314423561, 0.15029701590538025, 0.000490668579004705, 0.5364181399345398, 0.0016803600592538714, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4440009295940399, 0.5055950880050659, 0.14072291553020477, 0.20776981115341187, 0.24339812994003296, 0.01946749910712242, 0.1477651447057724, 0.24892206490039825, 0.13990418612957, 0.5277839303016663, 0.22113053500652313, 0.7815175652503967, 0.04741470143198967, 0.31336119771003723, 0.318754643201828, 0.17249688506126404, 0.003960400819778442, 1.1815190191555303e-05, 0.00205309153534472, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.003975332248955965, 0.09357346594333649, 0.000580776366405189, 0.001556370290927589, 0.0040078358724713326, 0.00020105167641304433, 0.005314813926815987, 0.0463886484503746, 0.0025405578780919313, 0.008098164573311806, 0.0004367573419585824, 0.0955028310418129, 0.0013312119990587234, 0.008472515270113945, 0.16612127423286438, 0.08659190684556961, 0.2260276973247528, 0.018877657130360603, 0.019257033243775368, 0.9179584980010986, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00713347876444459, 0.11304348707199097, 0.007166451308876276, 0.017305465415120125, 0.01892760582268238, 0.004294875077903271, 0.013284130021929741, 0.05641845986247063, 0.006293897051364183, 0.008091668598353863, 0.004229044076055288, 0.03852742537856102, 0.036073870956897736, 0.030675750225782394, 0.1423715502023697, 2.1155383365112357e-05, 0.00016346832853741944, 0.0004644138098228723, 9.852640505414456e-05, 0.009302367456257343, 0.8758521676063538, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.112990602850914, 0.20299020409584045, 0.29141831398010254, 0.1917479783296585, 0.25626659393310547, 0.40023526549339294, 0.045914653688669205, 0.05403761938214302, 0.3577503561973572, 0.11164049804210663, 0.20054538547992706, 0.23382915556430817, 0.3541012704372406, 0.39880213141441345, 0.05442150682210922, 0.0038963633123785257, 0.11578002572059631, 0.06833135336637497, 0.2930091321468353, 0.06728219240903854, 0.588379442691803, 0.190787211060524, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11769542098045349, 0.22490660846233368, 0.16446754336357117, 0.17726869881153107, 0.24409359693527222, 0.16966795921325684, 0.06426751613616943, 0.1868649125099182, 0.17593497037887573, 0.10732528567314148, 0.1210716962814331, 0.18835949897766113, 0.07820838689804077, 0.12172650545835495, 0.0815061554312706, 0.04113525524735451, 0.03917931765317917, 0.013817446306347847, 0.06874216347932816, 0.027753230184316635, 0.04752122610807419, 0.17637789249420166, 0.2964049279689789, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08801974356174469, 0.2964327037334442, 0.17140379548072815, 0.1086457222700119, 0.1790848970413208, 0.042561717331409454, 0.02568918652832508, 0.12736740708351135, 0.4644424617290497, 0.09952269494533539, 0.1403166949748993, 0.12085206061601639, 0.2499331831932068, 0.14905890822410583, 0.04691213369369507, 0.006397286430001259, 0.008155078627169132, 0.02385183423757553, 0.08218340575695038, 0.09733399748802185, 0.7216709852218628, 0.11420661956071854, 0.028804002329707146, 0.49512770771980286, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.28339406847953796, 0.25363603234291077, 0.49371209740638733, 0.28714650869369507, 0.42171764373779297, 0.03586414083838463, 0.140908345580101, 0.27345338463783264, 0.06897412985563278, 0.24740128219127655, 0.5061832070350647, 0.4192107915878296, 0.43851029872894287, 0.29079654812812805, 0.10071542859077454, 0.007080267183482647, 0.010165071114897728, 0.007166726514697075, 0.04547898843884468, 0.014898931607604027, 0.06153866648674011, 0.05960511788725853, 0.025653565302491188, 0.05574938654899597, 0.5054050087928772, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.049345988780260086, 0.1473262906074524, 0.10952533781528473, 0.16707968711853027, 0.25493475794792175, 0.03866606950759888, 0.046480532735586166, 0.16288119554519653, 0.06614720076322556, 0.0629507377743721, 0.07218940556049347, 0.3448391556739807, 0.06943795084953308, 0.058807674795389175, 0.135455921292305, 0.12821261584758759, 0.09823491424322128, 0.2407415509223938, 0.03722868487238884, 0.07500484585762024, 0.23719841241836548, 0.08696958422660828, 0.10033686459064484, 0.08637046813964844, 0.05946339666843414, 0.17889682948589325, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05557708069682121, 0.024377070367336273, 0.171014666557312, 0.1548214852809906, 0.21205416321754456, 0.29049578309059143, 0.08155391365289688, 0.2053205668926239, 0.09979691356420517, 0.11640740185976028, 0.23155182600021362, 0.4772811830043793, 0.2134055644273758, 0.3209300637245178, 0.0739695355296135, 0.018611561506986618, 0.530681848526001, 0.37442806363105774, 0.09326046705245972, 0.039934538304805756, 0.607749342918396, 0.1011725440621376, 0.041957128793001175, 0.061673425137996674, 0.012941170483827591, 0.012897199019789696, 0.02531522512435913, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.046621087938547134, 0.02855776995420456, 0.11975010484457016, 0.2049850970506668, 0.16244490444660187, 0.14614170789718628, 0.03785347566008568, 0.2537410259246826, 0.3719625771045685, 0.1159287542104721, 0.23734091222286224, 0.26474830508232117, 0.04938332363963127, 0.17566856741905212, 0.034675102680921555, 0.025258230045437813, 0.013820141553878784, 0.020238902419805527, 0.20186173915863037, 0.008764497935771942, 0.044081512838602066, 0.11685895919799805, 0.12131167203187943, 0.03466574102640152, 0.0033257410395890474, 0.009427645243704319, 0.00932170171290636, 0.6215367317199707, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08535599708557129, 0.01230260543525219, 0.28460273146629333, 0.3323705196380615, 0.13364574313163757, 0.14216013252735138, 0.16550986468791962, 0.36634352803230286, 0.3233327269554138, 0.13755354285240173, 0.6341029405593872, 0.1276889443397522, 0.0818048045039177, 0.2633805274963379, 0.10007897019386292, 0.0027034373488277197, 0.008653531782329082, 0.0021412167698144913, 0.02395743690431118, 0.06537352502346039, 0.05110874027013779, 0.050060901790857315, 0.023448945954442024, 0.0059632728807628155, 0.0016337132547050714, 0.0060929651372134686, 0.00957516860216856, 0.05008334666490555, 0.696637749671936, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.014263293705880642, 0.07173046469688416, 0.01932992786169052, 0.01909404993057251, 0.16755935549736023, 0.2271488904953003, 0.1093294620513916, 0.14342457056045532, 0.0580194853246212, 0.01671113632619381, 0.03395597264170647, 0.0692841187119484, 0.07175575196743011, 0.04972841590642929, 0.12856654822826385, 5.63129390229733e-07, 0.00027805642457678914, 1.7160025890916586e-05, 5.958595011179568e-06, 0.00078710971865803, 1.2566613349918043e-06, 9.03528507478768e-06, 2.1993335394654423e-05, 4.528845238382928e-06, 1.0594538935038145e-06, 2.375837993895402e-06, 1.0765622391772922e-05, 0.00012861557479482144, 0.000270194374024868, 0.4203896224498749, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06590985506772995, 0.1636172980070114, 0.09935098141431808, 0.20126965641975403, 0.4101002812385559, 0.21936923265457153, 0.26084569096565247, 0.3593950569629669, 0.014820259064435959, 0.05201014503836632, 0.03426084294915199, 0.38774317502975464, 0.1401163786649704, 0.3782513439655304, 0.13036324083805084, 0.19651824235916138, 0.009276115335524082, 0.0007576652569696307, 0.02043321169912815, 0.000937489268835634, 0.0014158851699903607, 0.02691410481929779, 0.025149332359433174, 0.015754513442516327, 0.002638434525579214, 0.03568584471940994, 0.28478676080703735, 0.08937329053878784, 0.04057440906763077, 0.41798362135887146, 0.02812151424586773, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05128908529877663, 0.11090300232172012, 0.24501535296440125, 0.07115167379379272, 0.3950805068016052, 0.2010982632637024, 0.08927696198225021, 0.2923780679702759, 0.11195118725299835, 0.05971711874008179, 0.14540457725524902, 0.4000069797039032, 0.2374461144208908, 0.47139719128608704, 0.10731440782546997, 0.0009883381426334381, 0.005475975573062897, 0.017872320488095284, 0.0038598645478487015, 0.01383217889815569, 0.1060260757803917, 0.010558119975030422, 0.0004280287539586425, 0.011488020420074463, 0.004323506727814674, 0.015877770259976387, 0.025533713400363922, 0.06758329272270203, 0.005362953990697861, 0.03033292666077614, 0.3987913429737091, 0.22715723514556885, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.014083221554756165, 0.029302498325705528, 0.019839908927679062, 0.019802037626504898, 0.11310776323080063, 0.014347831718623638, 0.013065088540315628, 0.0404186025261879, 0.14103254675865173, 0.01056672353297472, 0.02028844505548477, 0.4335528016090393, 0.019943613559007645, 0.08491621166467667, 0.15365199744701385, 0.025437461212277412, 0.027387555688619614, 0.0211916733533144, 0.0013409400125965476, 0.0016278955154120922, 0.0205780491232872, 0.006606978829950094, 0.005105526186525822, 0.008417481556534767, 0.008475488983094692, 0.016475802287459373, 0.021865585818886757, 0.04041945934295654, 0.001965513452887535, 0.030297037214040756, 0.018051480874419212, 0.2940014600753784, 0.09546513855457306, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04251990094780922, 0.025738505646586418, 0.19788101315498352, 0.08900192379951477, 0.20504283905029297, 0.36725619435310364, 0.05852765589952469, 0.12635937333106995, 0.07596885412931442, 0.055006030946969986, 0.1975020170211792, 0.39253395795822144, 0.2602497935295105, 0.3791850209236145, 0.11310473829507828, 0.014116446487605572, 0.6685785055160522, 0.40577325224876404, 0.09365412592887878, 0.008716625161468983, 0.504762589931488, 0.11037815362215042, 0.03693895787000656, 0.066362664103508, 0.025546396151185036, 0.030971869826316833, 0.07333581149578094, 0.21910515427589417, 0.03128749132156372, 0.013437384739518166, 0.06674141436815262, 0.055549826472997665, 0.02615067921578884, 0.05289305001497269, NaN, NaN, NaN, NaN, NaN, NaN], [0.06150972843170166, 0.049163203686475754, 0.14174170792102814, 0.13322500884532928, 0.16170991957187653, 0.21354396641254425, 0.04667104035615921, 0.26311540603637695, 0.32218027114868164, 0.0809161439538002, 0.18361496925354004, 0.23948682844638824, 0.09133663028478622, 0.25973111391067505, 0.07212682068347931, 0.01752244122326374, 0.013681006617844105, 0.015325021930038929, 0.15400148928165436, 0.0017620606813579798, 0.03783759847283363, 0.07285356521606445, 0.042190372943878174, 0.019725583493709564, 0.004497688263654709, 0.010335608385503292, 0.023485884070396423, 0.5969190001487732, 0.22785267233848572, 0.05655405670404434, 0.05765213817358017, 0.006416310556232929, 0.029401889070868492, 0.022928474470973015, 0.6468356251716614, NaN, NaN, NaN, NaN, NaN], [0.12382826954126358, 0.035204268991947174, 0.3469122052192688, 0.27821084856987, 0.12485836446285248, 0.1130678728222847, 0.12963837385177612, 0.3451126217842102, 0.16417652368545532, 0.12570835649967194, 0.5000419616699219, 0.09880878776311874, 0.042446259409189224, 0.2635292708873749, 0.16834798455238342, 0.003705248236656189, 0.09392052888870239, 0.0011726000811904669, 0.042238909751176834, 0.07787514477968216, 0.11800158768892288, 0.09318403154611588, 0.018972182646393776, 0.022339271381497383, 0.02290215529501438, 0.009648749604821205, 0.020298194140195847, 0.09632600843906403, 0.6665039658546448, 0.01913357712328434, 0.016501925885677338, 0.01550414226949215, 0.014767719432711601, 0.035943012684583664, 0.1298983097076416, 0.7307590246200562, NaN, NaN, NaN, NaN], [0.010800065472722054, 0.04851265624165535, 0.01629789173603058, 0.013155121356248856, 0.14412836730480194, 0.10944324731826782, 0.08000180870294571, 0.10409139841794968, 0.054843056946992874, 0.011575616896152496, 0.02017728053033352, 0.044063322246074677, 0.04816943034529686, 0.03936787694692612, 0.1280953288078308, 3.2450822118335054e-07, 0.0001958437787834555, 1.195628647110425e-05, 3.192948497598991e-06, 0.00034392892848700285, 1.3818779507346335e-06, 6.319523890851997e-06, 9.25252061279025e-06, 3.2897685287025524e-06, 1.041492623699014e-06, 2.450263082209858e-06, 1.1291336704744026e-05, 9.216016042046249e-05, 0.00025747373001649976, 0.3770022690296173, 7.494814053643495e-05, 0.00011931787594221532, 5.454379424918443e-05, 3.481862586340867e-05, 0.0001493972522439435, 6.532184488605708e-05, 0.4379080533981323, NaN, NaN, NaN], [0.03501533716917038, 0.12365423142910004, 0.058643028140068054, 0.026187611743807793, 0.2106953263282776, 0.09627192467451096, 0.1373300403356552, 0.209503173828125, 0.00544273667037487, 0.010177833028137684, 0.00795654021203518, 0.17826952040195465, 0.06280092895030975, 0.2785777747631073, 0.15446779131889343, 0.11172444373369217, 0.00812594499439001, 0.000803561822976917, 0.011673782020807266, 0.00013412271800916642, 0.002435607835650444, 0.021002406254410744, 0.009926681406795979, 0.014218374155461788, 0.0044799866154789925, 0.03462693840265274, 0.49634605646133423, 0.1610735058784485, 0.03537029027938843, 0.3717024624347687, 0.0470024012029171, 0.0025306264869868755, 0.08426976948976517, 0.5137573480606079, 0.047759927809238434, 0.008752438239753246, 0.5270217657089233, 0.020567137748003006, NaN, NaN], [0.055331505835056305, 0.14680130779743195, 0.22850985825061798, 0.040600359439849854, 0.2299574315547943, 0.21366852521896362, 0.10291176289319992, 0.2649042010307312, 0.07482050359249115, 0.04207760840654373, 0.11352740973234177, 0.22353075444698334, 0.2551318407058716, 0.4900997579097748, 0.11985023319721222, 0.00039373920299112797, 0.00142151047475636, 0.016346368938684464, 0.0038184949662536383, 0.00426360173150897, 0.10012070834636688, 0.007060237228870392, 0.00022489627008326352, 0.006389277055859566, 0.0014407823327928782, 0.01344740204513073, 0.019176417961716652, 0.04953484237194061, 0.003102741902694106, 0.017501499503850937, 0.25968801975250244, 0.12805432081222534, 0.03450275957584381, 0.03214799612760544, 0.06495527178049088, 0.007038496434688568, 0.018200475722551346, 0.2228115350008011, 0.24082934856414795, NaN], [0.04223596677184105, 0.14613933861255646, 0.08112313598394394, 0.04192597419023514, 0.11981905251741409, 0.18680673837661743, 0.07695262134075165, 0.14058402180671692, 0.1875196099281311, 0.05864474177360535, 0.0581248439848423, 0.23554684221744537, 0.21983209252357483, 0.1619952768087387, 0.12595340609550476, 0.004585978575050831, 0.008592751808464527, 0.20804427564144135, 0.003501898143440485, 0.01809401623904705, 0.0088487658649683, 0.01839679665863514, 0.009930659085512161, 0.019693726673722267, 0.015943868085741997, 0.06719032675027847, 0.03678698092699051, 0.03292753919959068, 0.02313893660902977, 0.023240724578499794, 0.03294161707162857, 0.24390928447246552, 0.10472099483013153, 0.0623757429420948, 0.06489475816488266, 0.03424002602696419, 0.03615953400731087, 0.05666068568825722, 0.29077935218811035, 0.20903274416923523]], [[0.020951254293322563, 0.19576001167297363, 0.05422525107860565, 0.000516751199029386, 0.0576050765812397, 0.039616964757442474, 0.0011584623716771603, 0.06260760873556137, 0.05524995177984238, 5.760174462920986e-05, 0.0005486492882482708, 0.01856253668665886, 0.008022493682801723, 0.0032547120936214924, 0.1980074942111969, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.15878187119960785, 0.5755441188812256, 0.073322594165802, 0.006848999299108982, 0.04221894592046738, 0.057610929012298584, 0.01498481910675764, 0.15564584732055664, 0.02557745948433876, 0.010493909008800983, 0.04444737732410431, 0.10564734041690826, 0.04703369736671448, 0.007807346060872078, 0.10371111333370209, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0667557343840599, 0.5756934881210327, 0.02783285267651081, 0.001271323417313397, 0.13096383213996887, 0.007863562554121017, 0.0004880728665739298, 0.00786207988858223, 0.030193913727998734, 0.0004458925104700029, 0.0008183285826817155, 0.003005507169291377, 0.008833326399326324, 0.014566708356142044, 0.09050195664167404, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.006902126595377922, 0.22582471370697021, 0.027240794152021408, 0.000252248632023111, 0.08146748691797256, 0.008376134559512138, 0.0017193618696182966, 0.010283069685101509, 0.09191752970218658, 1.873078872449696e-05, 0.0001427968527423218, 0.0006295929779298604, 0.016630304977297783, 0.005029548890888691, 0.17517179250717163, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.46813952922821045, 0.7474208474159241, 0.04419572278857231, 0.039987821131944656, 0.07900705188512802, 0.010286353528499603, 0.008277984336018562, 0.21022778749465942, 0.018339863047003746, 0.003122991183772683, 0.0047759185545146465, 0.0031952662393450737, 0.0037801233120262623, 0.005526377819478512, 0.11187370121479034, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.08057912439107895, 0.09254536032676697, 0.26037144660949707, 0.04459136351943016, 0.19053104519844055, 0.18187369406223297, 0.04494835063815117, 0.08866222947835922, 0.05515718460083008, 0.011219717562198639, 0.041749756783246994, 0.13417255878448486, 0.43527963757514954, 0.4240920841693878, 0.05903848633170128, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.005677447654306889, 0.1104632169008255, 0.17886187136173248, 0.06816153228282928, 0.31320425868034363, 0.08580746501684189, 0.044242095202207565, 0.4031389355659485, 0.13310441374778748, 8.991359209176153e-05, 0.00051962147699669, 0.017516016960144043, 0.02517649158835411, 0.02827705629169941, 0.13873830437660217, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.009441166184842587, 0.04568161070346832, 0.08503290265798569, 0.055850934237241745, 0.15800173580646515, 0.09921947866678238, 0.2719998359680176, 0.7131122350692749, 0.12690743803977966, 0.0015569856623187661, 0.019959524273872375, 0.06398878246545792, 0.1124982088804245, 0.07506788522005081, 0.06075114384293556, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1778930425643921, 0.41812169551849365, 0.05459700897336006, 0.015388981439173222, 0.296997606754303, 0.041353121399879456, 0.1696915328502655, 0.1226804181933403, 0.3453136682510376, 0.006036087870597839, 0.008416525088250637, 0.004891113843768835, 0.003974124789237976, 0.0023401544895023108, 0.04184575751423836, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0018550200620666146, 0.2628808617591858, 0.0018376001389697194, 9.925621998263523e-05, 0.008250601589679718, 0.11965687572956085, 0.011913565918803215, 0.3649533987045288, 0.12527383863925934, 0.0011617891723290086, 0.002173396060243249, 0.011088940314948559, 0.02579125389456749, 0.004398738034069538, 0.18079015612602234, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0033212341368198395, 0.4786561131477356, 0.00019389556837268174, 4.100392834516242e-05, 0.03255903348326683, 0.004482456482946873, 0.0018638258334249258, 0.04032744839787483, 0.151435986161232, 0.0011174781247973442, 0.0008650964009575546, 0.049343932420015335, 0.013284855522215366, 0.009702197276055813, 0.17111515998840332, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.015286837704479694, 0.17760051786899567, 0.012107143178582191, 0.004069492220878601, 0.40114596486091614, 0.005856915842741728, 0.025313973426818848, 0.23595470190048218, 0.5599475502967834, 0.019674712792038918, 0.01789786107838154, 0.0449712835252285, 0.024323459714651108, 0.008310162462294102, 0.10516723990440369, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.013816175982356071, 0.10832668840885162, 0.014126134105026722, 0.0044770012609660625, 0.18972823023796082, 0.04144473373889923, 0.013167506083846092, 0.0398833267390728, 0.08117146790027618, 0.03379456326365471, 0.04336484149098396, 0.6766878366470337, 0.6025072932243347, 0.24042664468288422, 0.05677386373281479, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.010657100938260555, 0.1729527860879898, 0.006031150463968515, 0.006062258500605822, 0.10042858123779297, 0.007653414737433195, 0.0031583579257130623, 0.014785557985305786, 0.13275322318077087, 0.05689838156104088, 0.04302775487303734, 0.36964303255081177, 0.3870774507522583, 0.31299954652786255, 0.07590257376432419, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.014769526198506355, 0.05199434980750084, 0.11582475155591965, 0.14804258942604065, 0.05702318996191025, 0.3275434374809265, 0.3759170472621918, 0.3329218327999115, 0.027774346992373466, 0.12548163533210754, 0.13219930231571198, 0.029332099482417107, 0.2028164267539978, 0.518939197063446, 4.3280975660309196e-05, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.24939602613449097, 0.0921018123626709, 0.20195554196834564, 0.25931593775749207, 0.24976609647274017, 0.08025927096605301, 0.10602997988462448, 0.08455296605825424, 0.038250602781772614, 0.34039628505706787, 0.2528480887413025, 0.17168891429901123, 0.12038858979940414, 0.16591216623783112, 0.05973837152123451, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04881530627608299, 0.07757209986448288, 0.080610491335392, 0.047049663960933685, 0.2744564712047577, 0.18291208148002625, 0.11781244724988937, 0.130965456366539, 0.16412131488323212, 0.049904536455869675, 0.10192018002271652, 0.46385079622268677, 0.23078110814094543, 0.23192283511161804, 0.17445482313632965, 0.15880486369132996, 0.04734092205762863, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11153621971607208, 0.27696484327316284, 0.0350787453353405, 0.011731116101145744, 0.08945441246032715, 0.2750371992588043, 0.07341955602169037, 0.12011690437793732, 0.026965567842125893, 0.023494159802794456, 0.015654105693101883, 0.05704642832279205, 0.11022293567657471, 0.0463077574968338, 0.1307818740606308, 0.22883240878582, 0.015307039953768253, 0.023610780015587807, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06216026097536087, 0.123567596077919, 0.044055916368961334, 0.012494971975684166, 0.045035671442747116, 0.18137943744659424, 0.1501520872116089, 0.0996006652712822, 0.05310875549912453, 0.11289763450622559, 0.05045852065086365, 0.055306825786828995, 0.3424266576766968, 0.1600506752729416, 0.04121629521250725, 0.15376803278923035, 0.17623378336429596, 0.16427822411060333, 0.018553992733359337, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03470996022224426, 0.38486456871032715, 0.007671448867768049, 0.014272118918597698, 0.01295357197523117, 0.001353065250441432, 0.035229261964559555, 0.10929086059331894, 0.03641098737716675, 0.08741087466478348, 0.01870635710656643, 0.10011491179466248, 0.03142678365111351, 0.12343490868806839, 0.15971165895462036, 0.12576976418495178, 0.44071146845817566, 0.38860467076301575, 0.12043511122465134, 0.027116619050502777, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03053746558725834, 0.24113330245018005, 0.009466315619647503, 0.01980357989668846, 0.04114365205168724, 0.05523357167840004, 0.027042368426918983, 0.10979101061820984, 0.004461985547095537, 0.04689180105924606, 0.04529552906751633, 0.1364448219537735, 0.054305437952280045, 0.06579019129276276, 0.13895106315612793, 0.03928220644593239, 0.42239660024642944, 0.2546820342540741, 0.22367709875106812, 0.1215892881155014, 0.001983387628570199, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3289671242237091, 0.3443813920021057, 0.38217487931251526, 0.32642021775245667, 0.12515123188495636, 0.04144418612122536, 0.06740343570709229, 0.024584289640188217, 0.007359183859080076, 0.39375364780426025, 0.38123685121536255, 0.3035361170768738, 0.18788036704063416, 0.13260427117347717, 0.09976762533187866, 0.17152060568332672, 0.49365419149398804, 0.08085957914590836, 0.02207508496940136, 0.19231174886226654, 0.008304901421070099, 0.03878962993621826, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1711268573999405, 0.1900682896375656, 0.20778892934322357, 0.08847668021917343, 0.39589688181877136, 0.3955995440483093, 0.3348483741283417, 0.11133389919996262, 0.10861264914274216, 0.14033687114715576, 0.26926568150520325, 0.4846358299255371, 0.23405344784259796, 0.4343181252479553, 0.08998383581638336, 0.13843253254890442, 0.07047099620103836, 0.2525072991847992, 0.13487939536571503, 0.27911728620529175, 0.11727599054574966, 0.022392159327864647, 0.1764850914478302, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4154844284057617, 0.4073733687400818, 0.5541329383850098, 0.43809109926223755, 0.11503908038139343, 0.02849700301885605, 0.025097709149122238, 0.014711813069880009, 0.006424109451472759, 0.39197838306427, 0.4694826304912567, 0.17039237916469574, 0.16142874956130981, 0.19919125735759735, 0.054951149970293045, 0.10915631055831909, 0.30942168831825256, 0.19657404720783234, 0.031007295474410057, 0.23716343939304352, 0.05435822904109955, 0.08149112015962601, 0.6613667011260986, 0.11670006066560745, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24498042464256287, 0.277620404958725, 0.060333866626024246, 0.030503980815410614, 0.04090564325451851, 0.4659561812877655, 0.2110646367073059, 0.11101182550191879, 0.028219982981681824, 0.10508411377668381, 0.025386929512023926, 0.0648839995265007, 0.13676653802394867, 0.07622335106134415, 0.09164498746395111, 0.0640818402171135, 0.41535088419914246, 0.29784247279167175, 0.05657188221812248, 0.036311421543359756, 0.08192699402570724, 0.16688455641269684, 0.10144203901290894, 0.346017450094223, 0.15466110408306122, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4220424294471741, 0.21296784281730652, 0.10483475774526596, 0.11319100856781006, 0.14396990835666656, 0.1309618502855301, 0.13656088709831238, 0.2097199261188507, 0.1397993415594101, 0.263439804315567, 0.10735370218753815, 0.27457332611083984, 0.26051631569862366, 0.18891198933124542, 0.10100831091403961, 0.04877842590212822, 0.16450235247612, 0.23761717975139618, 0.0720985159277916, 0.12954245507717133, 0.08035153150558472, 0.18124118447303772, 0.05973014980554581, 0.26483285427093506, 0.39028850197792053, 0.05098416656255722, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12607140839099884, 0.08847615122795105, 0.09191321581602097, 0.06030821427702904, 0.21649383008480072, 0.10438336431980133, 0.07331530004739761, 0.1330888420343399, 0.04176999628543854, 0.06727378815412521, 0.06257567554712296, 0.21110908687114716, 0.09018781781196594, 0.09389244765043259, 0.13621515035629272, 0.11044558137655258, 0.08550350368022919, 0.2513507902622223, 0.28401821851730347, 0.12441904842853546, 0.05029991641640663, 0.42405593395233154, 0.08374682813882828, 0.43869927525520325, 0.14253327250480652, 0.10876792669296265, 0.09369473904371262, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.062066610902547836, 0.07845254987478256, 0.24838510155677795, 0.16541223227977753, 0.16867581009864807, 0.019677892327308655, 0.021460779011249542, 0.018530650064349174, 0.023010587319731712, 0.10349667817354202, 0.16099916398525238, 0.3089703619480133, 0.08426959812641144, 0.16459643840789795, 0.06073381006717682, 0.08764015138149261, 0.46941375732421875, 0.23278135061264038, 0.11763583868741989, 0.0354606918990612, 0.16624747216701508, 0.2793619632720947, 0.1965668648481369, 0.23052528500556946, 0.3914787769317627, 0.08669382333755493, 0.10678009688854218, 0.08708767592906952, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11642084270715714, 0.11190053075551987, 0.12368596345186234, 0.04549993947148323, 0.3567850887775421, 0.06569506227970123, 0.07286660373210907, 0.03259556367993355, 0.09530685096979141, 0.19273261725902557, 0.06463074684143066, 0.7640278339385986, 0.06371455639600754, 0.1593337506055832, 0.2193848341703415, 0.2116944044828415, 0.06720030307769775, 0.29984304308891296, 0.010844358243048191, 0.051072586327791214, 0.15023349225521088, 0.04554526135325432, 0.1560167670249939, 0.03609438240528107, 0.026584016159176826, 0.14512087404727936, 0.05890262499451637, 0.015816861763596535, 0.07422769069671631, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11034999042749405, 0.03210863843560219, 0.010996339842677116, 0.026450032368302345, 0.051475513726472855, 0.02743532694876194, 0.3610350787639618, 0.20538736879825592, 0.017281753942370415, 0.05300014466047287, 0.012052728794515133, 0.08001075685024261, 0.0069017065688967705, 0.010893179103732109, 0.13085691630840302, 0.056502565741539, 0.15541820228099823, 0.07158821076154709, 0.00490804947912693, 0.015012365765869617, 0.06302572786808014, 0.01116714347153902, 0.22065599262714386, 0.021468764171004295, 0.01365464273840189, 0.022816751152276993, 0.019708380103111267, 0.0059420084580779076, 0.0700121819972992, 0.287899911403656, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07615644484758377, 0.1536630541086197, 0.1253354847431183, 0.048576656728982925, 0.05276811867952347, 0.1611642986536026, 0.12317243963479996, 0.32385867834091187, 0.012925365939736366, 0.0864856168627739, 0.08918802440166473, 0.23886144161224365, 0.20351386070251465, 0.20744860172271729, 0.13318131864070892, 0.058403778821229935, 0.0693131536245346, 0.04999461770057678, 0.004054869059473276, 0.0624610111117363, 0.018093721941113472, 0.07961009442806244, 0.1545858234167099, 0.3008257746696472, 0.14455094933509827, 0.09800520539283752, 0.09531621634960175, 0.27401015162467957, 0.4782770574092865, 0.11211755871772766, 0.01358953770250082, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.051417503505945206, 0.1600690335035324, 0.08639511466026306, 0.02997625432908535, 0.08503448963165283, 0.32695260643959045, 0.06822863221168518, 0.16364485025405884, 0.06138167902827263, 0.07786902785301208, 0.04443247988820076, 0.0585777647793293, 0.1263807862997055, 0.10769001394510269, 0.13808733224868774, 0.1399688720703125, 0.5559014678001404, 0.20350231230258942, 0.042011573910713196, 0.020507201552391052, 0.03915366902947426, 0.4243565797805786, 0.11376935243606567, 0.31140708923339844, 0.051479678601026535, 0.07416504621505737, 0.2654426097869873, 0.3960915207862854, 0.5790604948997498, 0.18063338100910187, 0.1939544379711151, 0.04191381484270096, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1321558654308319, 0.24967153370380402, 0.0761917233467102, 0.044561922550201416, 0.12028387933969498, 0.19908402860164642, 0.04708404839038849, 0.10076720267534256, 0.09921064227819443, 0.18345412611961365, 0.09404058009386063, 0.21650025248527527, 0.11625839024782181, 0.1530369222164154, 0.12011245638132095, 0.027515297755599022, 0.0486784465610981, 0.06845460832118988, 0.023408811539411545, 0.008863206952810287, 0.008533195592463017, 0.24178741872310638, 0.01229054294526577, 0.25817692279815674, 0.6869812607765198, 0.049950506538152695, 0.12178820371627808, 0.0564231351017952, 0.02026011236011982, 0.004908477421849966, 0.03562311828136444, 0.12746450304985046, 0.0016219470417127013, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10757170617580414, 0.1042957603931427, 0.13590699434280396, 0.06331591308116913, 0.24158470332622528, 0.09161848574876785, 0.0633605495095253, 0.13977625966072083, 0.03925082087516785, 0.07121878862380981, 0.1023484393954277, 0.26378345489501953, 0.10990181565284729, 0.12030858546495438, 0.1261080652475357, 0.11620164662599564, 0.09937138110399246, 0.17538107931613922, 0.40406307578086853, 0.043817292898893356, 0.05759625509381294, 0.49306368827819824, 0.09120260924100876, 0.36450278759002686, 0.08042807132005692, 0.1856311559677124, 0.1376025527715683, 0.1998283714056015, 0.3654005527496338, 0.15910619497299194, 0.4969707429409027, 0.08565060794353485, 0.02514367550611496, 0.090617336332798, NaN, NaN, NaN, NaN, NaN, NaN], [0.06512168049812317, 0.13837532699108124, 0.3250073194503784, 0.16753129661083221, 0.21647527813911438, 0.04118574038147926, 0.03336784988641739, 0.029927842319011688, 0.03334499150514603, 0.08782976865768433, 0.17631417512893677, 0.3171449303627014, 0.10520178824663162, 0.15139654278755188, 0.0914224162697792, 0.0739481970667839, 0.5182103514671326, 0.19721719622612, 0.21118015050888062, 0.015751224011182785, 0.12249443680047989, 0.5174803733825684, 0.17075838148593903, 0.30025264620780945, 0.29246312379837036, 0.0875946432352066, 0.2326347827911377, 0.13986286520957947, 0.511695921421051, 0.12602318823337555, 0.03662485629320145, 0.1263200044631958, 0.0166145209223032, 0.19702456891536713, 0.09621746093034744, NaN, NaN, NaN, NaN, NaN], [0.06382797658443451, 0.2566763758659363, 0.11056842654943466, 0.028001734986901283, 0.2813059389591217, 0.24806144833564758, 0.07807287573814392, 0.05373501405119896, 0.21183612942695618, 0.09658068418502808, 0.05084875971078873, 0.501965343952179, 0.06208595260977745, 0.10913741588592529, 0.26912179589271545, 0.3052336871623993, 0.37224864959716797, 0.45515015721321106, 0.04986808821558952, 0.05332064628601074, 0.13846120238304138, 0.15990367531776428, 0.20659208297729492, 0.06640873104333878, 0.035323526710271835, 0.30340465903282166, 0.10174556821584702, 0.02102985605597496, 0.11508277803659439, 0.09203195571899414, 0.0029288395307958126, 0.023838462308049202, 0.004605103749781847, 0.052648112177848816, 0.006431906949728727, 0.026736242696642876, NaN, NaN, NaN, NaN], [0.08548272401094437, 0.017544403672218323, 0.011271107010543346, 0.022962557151913643, 0.05241750180721283, 0.02648325450718403, 0.3057800531387329, 0.19772306084632874, 0.025625178590416908, 0.03652432560920715, 0.006945622619241476, 0.05576859414577484, 0.00584550853818655, 0.008180957287549973, 0.12917736172676086, 0.047024402767419815, 0.1257133185863495, 0.052377521991729736, 0.009844984859228134, 0.015597687102854252, 0.06965665519237518, 0.01849394477903843, 0.1603521853685379, 0.02587857097387314, 0.00957732368260622, 0.023523790761828423, 0.020081259310245514, 0.008425970561802387, 0.10955916345119476, 0.35300737619400024, 0.023505402728915215, 0.00786643661558628, 0.007557017263025045, 0.013908758759498596, 0.004675114993005991, 0.035296451300382614, 0.3261549174785614, NaN, NaN, NaN], [0.03209112584590912, 0.1926622986793518, 0.09989916533231735, 0.02044818177819252, 0.04127199947834015, 0.22930434346199036, 0.09912838786840439, 0.3779822289943695, 0.007566491607576609, 0.046152934432029724, 0.04734500125050545, 0.35250937938690186, 0.10047939419746399, 0.16575956344604492, 0.13635975122451782, 0.11014947295188904, 0.08461853116750717, 0.02981843426823616, 0.004099451471120119, 0.009237504564225674, 0.011130756698548794, 0.132149338722229, 0.11619938164949417, 0.22203940153121948, 0.02292616292834282, 0.06793706119060516, 0.07227552682161331, 0.3262397348880768, 0.40601006150245667, 0.08270477503538132, 0.013506797142326832, 0.03135772421956062, 0.07034049183130264, 0.09623772650957108, 0.20842698216438293, 0.2752794623374939, 0.1234828308224678, 0.04129752516746521, NaN, NaN], [0.05301084369421005, 0.1661737710237503, 0.08216799795627594, 0.025789698585867882, 0.07900767773389816, 0.3054123520851135, 0.08738221228122711, 0.17720931768417358, 0.06289011240005493, 0.06967967748641968, 0.05491774156689644, 0.02886299602687359, 0.10253670811653137, 0.09415244311094284, 0.129754438996315, 0.1182219609618187, 0.7384620308876038, 0.11492461711168289, 0.09884578734636307, 0.012010940350592136, 0.038200050592422485, 0.4905328154563904, 0.23439669609069824, 0.2528713345527649, 0.015177865512669086, 0.07817362248897552, 0.33532261848449707, 0.4971323609352112, 0.7384514212608337, 0.2383432686328888, 0.2306600660085678, 0.025716517120599747, 0.023198120296001434, 0.3352215886116028, 0.4797173738479614, 0.5688640475273132, 0.2555003762245178, 0.1890360713005066, 0.06237812712788582, NaN], [0.1895110011100769, 0.09308972954750061, 0.1887637972831726, 0.14927715063095093, 0.3653167188167572, 0.1686658412218094, 0.1126369759440422, 0.17013703286647797, 0.0685301423072815, 0.15278968214988708, 0.19327588379383087, 0.18825437128543854, 0.143904447555542, 0.143670454621315, 0.1203024610877037, 0.13153354823589325, 0.5476850867271423, 0.27465543150901794, 0.27658137679100037, 0.5121651291847229, 0.3939417600631714, 0.2527337968349457, 0.41937416791915894, 0.2437492311000824, 0.1485103964805603, 0.10651403665542603, 0.241710364818573, 0.34289923310279846, 0.3691290616989136, 0.108230821788311, 0.32214298844337463, 0.08876177668571472, 0.03369928151369095, 0.23942533135414124, 0.302080899477005, 0.3531237244606018, 0.09724070131778717, 0.19267186522483826, 0.06874143332242966, 0.052875734865665436]], [[0.5917359590530396, 0.12410512566566467, 0.24872945249080658, 0.20040015876293182, 0.21720361709594727, 0.11561702191829681, 0.58521568775177, 0.41413450241088867, 0.22558750212192535, 0.117314413189888, 0.3378458619117737, 0.10710897296667099, 0.0625920221209526, 0.24034489691257477, 0.0060951621271669865, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03933318331837654, 0.17479471862316132, 0.1999012678861618, 0.1507989913225174, 0.2344110906124115, 0.41628938913345337, 0.19733835756778717, 0.42009472846984863, 0.32125937938690186, 0.09302358329296112, 0.29758843779563904, 0.2500022351741791, 0.15192696452140808, 0.19621950387954712, 0.06078135594725609, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03998054191470146, 0.02165106125175953, 0.5779209733009338, 0.4094802737236023, 0.3219829499721527, 0.23359909653663635, 0.15223096311092377, 0.0776560828089714, 0.11850404739379883, 0.1752316802740097, 0.7765606641769409, 0.15624035894870758, 0.19448350369930267, 0.3389243483543396, 0.015656093135476112, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2606712579727173, 0.23122362792491913, 0.33188652992248535, 0.327752023935318, 0.0930425301194191, 0.13157396018505096, 0.5079332590103149, 0.15524731576442719, 0.2039693295955658, 0.336448073387146, 0.7406277656555176, 0.11173539608716965, 0.03980698063969612, 0.2757716476917267, 0.009055807255208492, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.03992704302072525, 0.03562299162149429, 0.05761631205677986, 0.04593607783317566, 0.747100830078125, 0.13848423957824707, 0.25807130336761475, 0.11098858714103699, 0.025020861998200417, 0.027831630781292915, 0.07712040096521378, 0.5344594120979309, 0.28488224744796753, 0.37143638730049133, 0.060307834297418594, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.146702840924263, 0.5779150128364563, 0.04704871401190758, 0.12512727081775665, 0.05839477851986885, 0.5817644596099854, 0.2541782557964325, 0.167904794216156, 0.020014837384223938, 0.0557471327483654, 0.1778557300567627, 0.29983726143836975, 0.34978994727134705, 0.3759990334510803, 0.07532685250043869, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.14372284710407257, 0.20398879051208496, 0.060162752866744995, 0.022449441254138947, 0.15882903337478638, 0.12907396256923676, 0.7781419157981873, 0.20689332485198975, 0.023098474368453026, 0.02567201852798462, 0.04225016012787819, 0.05647281929850578, 0.5644452571868896, 0.8062969446182251, 0.0037398021668195724, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.09274263679981232, 0.19406189024448395, 0.18035270273685455, 0.18292436003684998, 0.2674761116504669, 0.1057504341006279, 0.5214765071868896, 0.1765710562467575, 0.15375129878520966, 0.08563723415136337, 0.35003283619880676, 0.12250327318906784, 0.4574505388736725, 0.6043637990951538, 0.046846963465213776, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3136129081249237, 0.10648278146982193, 0.02492944709956646, 0.07937752455472946, 0.16382691264152527, 0.40212482213974, 0.2148500233888626, 0.5046796798706055, 0.25625455379486084, 0.10382789373397827, 0.027611082419753075, 0.07138189673423767, 0.1265101283788681, 0.05298655480146408, 0.01642199046909809, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7252353429794312, 0.23862500488758087, 0.17466871440410614, 0.2584758698940277, 0.15821219980716705, 0.41019105911254883, 0.4795793294906616, 0.2558479905128479, 0.061036378145217896, 0.5831483006477356, 0.23237691819667816, 0.36767491698265076, 0.07294586300849915, 0.0734395682811737, 0.006080146878957748, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.18402060866355896, 0.2199273407459259, 0.10670217871665955, 0.36498934030532837, 0.37264159321784973, 0.5975290536880493, 0.641157865524292, 0.4798426032066345, 0.07047704607248306, 0.30389490723609924, 0.6835307478904724, 0.29959914088249207, 0.32009243965148926, 0.2076108753681183, 0.015385132282972336, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.18547095358371735, 0.1046445369720459, 0.17664410173892975, 0.031107882037758827, 0.4872691333293915, 0.6876094937324524, 0.29805243015289307, 0.2697339355945587, 0.03289056569337845, 0.04577193781733513, 0.2390383929014206, 0.650258481502533, 0.6253164410591125, 0.2719551920890808, 0.042574722319841385, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.06026101112365723, 0.4596063494682312, 0.11362233757972717, 0.050736263394355774, 0.47900232672691345, 0.8146356344223022, 0.23428170382976532, 0.5258204936981201, 0.07407079637050629, 0.24087238311767578, 0.04631686583161354, 0.04097185283899307, 0.24002470076084137, 0.051092784851789474, 0.10185284167528152, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.05915316566824913, 0.3385859429836273, 0.23845957219600677, 0.13520635664463043, 0.49372056126594543, 0.8321547508239746, 0.47351959347724915, 0.4942004382610321, 0.11661165207624435, 0.273796945810318, 0.09639480710029602, 0.07113680988550186, 0.3545372784137726, 0.3069557547569275, 0.026768943294882774, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6326229572296143, 0.28129494190216064, 0.2424720972776413, 0.23961131274700165, 0.1532977670431137, 0.03248026221990585, 0.07237446308135986, 0.03991716355085373, 0.058106135576963425, 0.6791825294494629, 0.4868316352367401, 0.4841252863407135, 0.1838759332895279, 0.16229771077632904, 0.03779346123337746, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.20045556128025055, 0.06346653401851654, 0.1246497705578804, 0.132145956158638, 0.18068760633468628, 0.0611145943403244, 0.3011611998081207, 0.09648064523935318, 0.3848741054534912, 0.20776434242725372, 0.09024091809988022, 0.10095226764678955, 0.05726093426346779, 0.17784324288368225, 0.06983170658349991, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06639314442873001, 0.03837187588214874, 0.306266725063324, 0.09758531302213669, 0.10875808447599411, 0.20901371538639069, 0.0894559919834137, 0.21620051562786102, 0.13805773854255676, 0.07912127673625946, 0.3521624505519867, 0.036526914685964584, 0.1551785171031952, 0.14622288942337036, 0.19236178696155548, 0.03290099650621414, 0.3365767002105713, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03379146009683609, 0.11666905134916306, 0.02791847102344036, 0.04754703491926193, 0.02039634808897972, 0.23185299336910248, 0.07985613495111465, 0.3240954875946045, 0.04561735317111015, 0.061520081013441086, 0.18156962096691132, 0.10860903561115265, 0.3409081995487213, 0.3218340575695038, 0.13103368878364563, 0.003547579748556018, 0.004082763101905584, 0.4616691768169403, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06278766691684723, 0.001863734913058579, 0.30563783645629883, 0.056017640978097916, 0.245498925447464, 0.11060530692338943, 0.09064232558012009, 0.004372697789222002, 0.007118886336684227, 0.06251134723424911, 0.17941752076148987, 0.004394095856696367, 0.11450538039207458, 0.046043287962675095, 0.021101655438542366, 0.03595791012048721, 0.1313885897397995, 0.007101066876202822, 0.42131781578063965, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11553236097097397, 0.0885467380285263, 0.2750205993652344, 0.21104735136032104, 0.3459762930870056, 0.07976578176021576, 0.218110129237175, 0.05760955810546875, 0.09680842608213425, 0.2662138342857361, 0.21090076863765717, 0.41520535945892334, 0.21548694372177124, 0.2248467653989792, 0.10481394827365875, 0.007601147051900625, 0.014137630350887775, 0.01938864029943943, 0.2572920322418213, 0.0011994435917586088, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03112325258553028, 0.08175794035196304, 0.035110849887132645, 0.038375336676836014, 0.2468937784433365, 0.060934457927942276, 0.0843387246131897, 0.03423367813229561, 0.02026834897696972, 0.07970783859491348, 0.08959806710481644, 0.1693299561738968, 0.16057033836841583, 0.21660663187503815, 0.13329552114009857, 0.00011468974116723984, 0.0032473355531692505, 0.00037737423554062843, 0.2793608605861664, 0.003465541172772646, 5.061212868895382e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09539461880922318, 0.058681365102529526, 0.01674766093492508, 0.02866855263710022, 0.012030106969177723, 0.21465063095092773, 0.034089475870132446, 0.04479566961526871, 0.014019637368619442, 0.035355255007743835, 0.1569557934999466, 0.01038492750376463, 0.06631091982126236, 0.1547483503818512, 0.19284123182296753, 0.21311266720294952, 0.10434294492006302, 0.011484598740935326, 0.0013334749964997172, 0.03845251351594925, 0.028238367289304733, 0.05654546618461609, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04954487085342407, 0.07065968960523605, 0.07275094836950302, 0.040997497737407684, 0.07946129143238068, 0.17300859093666077, 0.03222974017262459, 0.02469809167087078, 0.18557047843933105, 0.13542628288269043, 0.26776814460754395, 0.056715987622737885, 0.15973475575447083, 0.19029632210731506, 0.17610958218574524, 0.052184704691171646, 0.499632865190506, 0.005138374865055084, 0.10169705748558044, 0.09997230768203735, 0.036990027874708176, 0.07566682249307632, 0.32418423891067505, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.047577280551195145, 0.02606579288840294, 0.0165295097976923, 0.04137043654918671, 0.013305035419762135, 0.32835593819618225, 0.026565413922071457, 0.06772360950708389, 0.010228256694972515, 0.041277337819337845, 0.1336892545223236, 0.008326719515025616, 0.10322394222021103, 0.1976388841867447, 0.21077491343021393, 0.23645982146263123, 0.016864946112036705, 0.013305210508406162, 0.0007752762176096439, 0.017555342987179756, 0.03100133314728737, 0.04085567593574524, 0.029846351593732834, 0.010373883880674839, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.043893925845623016, 0.021177353337407112, 0.028366681188344955, 0.07016126066446304, 0.07573862373828888, 0.22699910402297974, 0.055615294724702835, 0.07980518788099289, 0.009269739501178265, 0.09460800141096115, 0.16427507996559143, 0.20832805335521698, 0.1427353024482727, 0.2680304944515228, 0.13907650113105774, 0.18805328011512756, 0.046367619186639786, 0.10314629226922989, 0.018223291262984276, 0.27720585465431213, 0.3798944056034088, 0.09291481226682663, 0.09293034672737122, 0.04290880635380745, 0.03370373696088791, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03411688283085823, 0.056632235646247864, 0.07365043461322784, 0.10934542864561081, 0.09185239672660828, 0.5077250003814697, 0.05141168087720871, 0.047258101403713226, 0.053326722234487534, 0.13365329802036285, 0.28296661376953125, 0.041020717471838, 0.08861301094293594, 0.13371184468269348, 0.11519401520490646, 0.028641005977988243, 0.03295213729143143, 0.0065453751012682915, 0.16686026751995087, 0.028714975342154503, 0.015397193841636181, 0.02003423683345318, 0.019093815237283707, 0.020523719489574432, 0.016172079369425774, 0.3490104377269745, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04096442833542824, 0.07374820858240128, 0.07300861179828644, 0.10121195018291473, 0.051522452384233475, 0.3508135676383972, 0.03948133811354637, 0.047985587269067764, 0.06340529769659042, 0.06765846908092499, 0.281475692987442, 0.05536516010761261, 0.1822110116481781, 0.22272904217243195, 0.13150985538959503, 0.10839971899986267, 0.004465002100914717, 0.016082070767879486, 0.035488102585077286, 0.015600458718836308, 0.012030484154820442, 0.015872180461883545, 0.01552913524210453, 0.03533920273184776, 0.11401902139186859, 0.31523072719573975, 0.20448055863380432, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07982534170150757, 0.06016559898853302, 0.03820561617612839, 0.02410227432847023, 0.006901262793689966, 0.42442968487739563, 0.02364957146346569, 0.07835549116134644, 0.027230771258473396, 0.12123586237430573, 0.15446297824382782, 0.018115278333425522, 0.21087171137332916, 0.29417684674263, 0.08362340182065964, 0.18776558339595795, 0.0060520414263010025, 0.017473671585321426, 0.005528539884835482, 0.0027145782951265574, 0.012176988646388054, 0.0031525399535894394, 0.004637573380023241, 0.011988476850092411, 0.06979440897703171, 0.38327983021736145, 0.020156072452664375, 0.010166948661208153, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05696694925427437, 0.014171368442475796, 0.06200120970606804, 0.021368764340877533, 0.012162269093096256, 0.0841592326760292, 0.03827953711152077, 0.07895056158304214, 0.01159723848104477, 0.05937046930193901, 0.023348387330770493, 0.008824712596833706, 0.13521961867809296, 0.23698511719703674, 0.03196632117033005, 0.3064975440502167, 0.004262991715222597, 0.009997943416237831, 0.00034317225799895823, 0.013912403024733067, 0.02852706052362919, 0.004078225698322058, 0.001928618410602212, 0.006367305759340525, 0.035507142543792725, 0.050674788653850555, 0.007057875394821167, 0.0049485149793326855, 0.0049379738047719, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11678174138069153, 0.8205142617225647, 0.01038320455700159, 0.023903295397758484, 0.21764065325260162, 0.2580764889717102, 0.20165181159973145, 0.2900886535644531, 0.03504627197980881, 0.10256802290678024, 0.03713424876332283, 0.7063723206520081, 0.8779962062835693, 0.8367014527320862, 0.0919082760810852, 0.14988604187965393, 0.015584584325551987, 0.137997567653656, 0.0031439096201211214, 0.5546696782112122, 0.01658078096807003, 0.0025873971171677113, 0.0010246702004224062, 0.019667595624923706, 0.012580120004713535, 0.015491531230509281, 0.029023459181189537, 0.021588340401649475, 0.25595030188560486, 0.02325037308037281, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.038494985550642014, 0.05109047889709473, 0.07501792907714844, 0.04001014679670334, 0.021166233345866203, 0.03079657442867756, 0.01494709774851799, 0.010983827523887157, 0.0029027159325778484, 0.0995086133480072, 0.350593626499176, 0.02021479234099388, 0.34575650095939636, 0.21952421963214874, 0.05450797453522682, 0.07357528805732727, 0.007756352424621582, 0.002724927617236972, 0.001402079127728939, 0.0004431438574101776, 0.00010925461538136005, 0.0029409730341285467, 0.005563507787883282, 0.012139370664954185, 0.03890732303261757, 0.05558362230658531, 0.03318313509225845, 0.4270496368408203, 0.07112571597099304, 0.15036046504974365, 0.020786603912711143, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.028108511120080948, 0.08174566179513931, 0.03328564018011093, 0.03230520337820053, 0.012646276503801346, 0.1872790902853012, 0.025206655263900757, 0.06737280637025833, 0.033121660351753235, 0.08641302585601807, 0.2848047614097595, 0.059273794293403625, 0.18425194919109344, 0.15244826674461365, 0.1352420449256897, 0.012120572850108147, 0.0003307444858364761, 0.009640182368457317, 0.00017808230768423527, 0.0021490382496267557, 0.0008148089982569218, 0.0008481521508656442, 0.0019973982125520706, 0.005024890415370464, 0.01719486527144909, 0.044799502938985825, 0.006444229744374752, 0.018026985228061676, 0.0067391968332231045, 0.061299871653318405, 0.01281613577157259, 0.3084925711154938, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07509021461009979, 0.05027765780687332, 0.23718997836112976, 0.11438266932964325, 0.11051909625530243, 0.431958943605423, 0.046987809240818024, 0.021854011341929436, 0.15366314351558685, 0.1928708851337433, 0.2900879681110382, 0.052021902054548264, 0.11538787186145782, 0.25173547863960266, 0.10233873873949051, 0.011204708367586136, 0.0033799665980041027, 0.008117830380797386, 0.1567971557378769, 0.012545537203550339, 0.002854604972526431, 0.0037395430263131857, 0.0003391341888345778, 0.002928558737039566, 0.004266565665602684, 0.28180748224258423, 0.005543314386159182, 0.0059068226255476475, 0.004401014186441898, 0.09436267614364624, 0.003524675266817212, 0.09697568416595459, 0.3818984925746918, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03257948160171509, 0.08023553341627121, 0.06238585337996483, 0.06856023520231247, 0.02927098423242569, 0.2968010902404785, 0.03317389637231827, 0.04758336395025253, 0.07943073660135269, 0.053982626646757126, 0.21416282653808594, 0.05025764927268028, 0.14347779750823975, 0.19969123601913452, 0.13921964168548584, 0.1085091158747673, 0.0013132937019690871, 0.011304548010230064, 0.014309195801615715, 0.009265521541237831, 0.00682368129491806, 0.01179590355604887, 0.005223054438829422, 0.01697726733982563, 0.05782441794872284, 0.2522926330566406, 0.16053971648216248, 0.020927468314766884, 0.02051178365945816, 0.1114674061536789, 0.014847181737422943, 0.40623563528060913, 0.12017090618610382, 0.2281613051891327, NaN, NaN, NaN, NaN, NaN, NaN], [0.07817428559064865, 0.11046875268220901, 0.040724072605371475, 0.024797527119517326, 0.004808576311916113, 0.5141928791999817, 0.024754824116826057, 0.080713652074337, 0.03179122135043144, 0.12244449555873871, 0.22665926814079285, 0.013305582106113434, 0.23485711216926575, 0.323343425989151, 0.10171245783567429, 0.23926517367362976, 0.007461922243237495, 0.015478387475013733, 0.02120528556406498, 0.0046339076943695545, 0.01287792343646288, 0.005305645987391472, 0.0037130024284124374, 0.011430526152253151, 0.10132863372564316, 0.42019084095954895, 0.03134358674287796, 0.006659360136836767, 0.0015345009742304683, 0.05340040102601051, 0.0021821516565978527, 0.15366847813129425, 0.09343723207712173, 0.04055917635560036, 0.009410854429006577, NaN, NaN, NaN, NaN, NaN], [0.03765244409441948, 0.0463164821267128, 0.06456112116575241, 0.05319739878177643, 0.010156691074371338, 0.1155625581741333, 0.02458079345524311, 0.07648347318172455, 0.019683409482240677, 0.06488858163356781, 0.09342794120311737, 0.059032924473285675, 0.15581923723220825, 0.2894386053085327, 0.04157077521085739, 0.3882482349872589, 0.012203006073832512, 0.008404962718486786, 0.0008633172838017344, 0.07213836163282394, 0.03903299570083618, 0.006879106629639864, 0.0025245456490665674, 0.011604986153542995, 0.1302306056022644, 0.05970751494169235, 0.005057368893176317, 0.0025832061655819416, 0.003548768814653158, 0.03821956738829613, 0.0041786422953009605, 0.029319334775209427, 0.009258194826543331, 0.010013489983975887, 0.0024901984725147486, 0.009316755458712578, NaN, NaN, NaN, NaN], [0.14924734830856323, 0.8862696886062622, 0.013125438243150711, 0.033269379287958145, 0.22599543631076813, 0.33975404500961304, 0.25561264157295227, 0.36481109261512756, 0.05327271297574043, 0.09902165085077286, 0.03598061203956604, 0.754990816116333, 0.9104278087615967, 0.8631682395935059, 0.10125402361154556, 0.08333727717399597, 0.009125825949013233, 0.12352871894836426, 0.0034849271178245544, 0.49194949865341187, 0.008760062977671623, 0.002427457133308053, 0.0004761714953929186, 0.014378424733877182, 0.007653949782252312, 0.010163314640522003, 0.018072640523314476, 0.014914281666278839, 0.33540958166122437, 0.012212751433253288, 0.050671979784965515, 0.08942927420139313, 0.0058481828309595585, 0.02088618278503418, 0.013520943000912666, 0.3026564419269562, 0.011637967079877853, NaN, NaN, NaN], [0.03672042489051819, 0.12888115644454956, 0.1578092873096466, 0.056865133345127106, 0.03288109228014946, 0.1379515379667282, 0.021150214597582817, 0.013284055516123772, 0.003249341854825616, 0.08646353334188461, 0.5471532940864563, 0.0361909456551075, 0.5093809366226196, 0.39931434392929077, 0.07520455867052078, 0.019913960248231888, 0.003490668721497059, 0.00020567848696373403, 0.00036819992237724364, 0.00019341551524121314, 3.8652269722661003e-05, 0.0008544524316675961, 0.002890991745516658, 0.001110991695895791, 0.005157719366252422, 0.008338885381817818, 0.0030357406940311193, 0.14557099342346191, 0.021602485328912735, 0.04367346689105034, 0.0015647107502445579, 0.009655454196035862, 0.14827704429626465, 0.008163533173501492, 0.49237948656082153, 0.06938102096319199, 0.08394628763198853, 0.049248531460762024, NaN, NaN], [0.03492635861039162, 0.09938696771860123, 0.028945090249180794, 0.03084651380777359, 0.012707062065601349, 0.15071596205234528, 0.029011720791459084, 0.05455483868718147, 0.03256314992904663, 0.07100401073694229, 0.2587825059890747, 0.05546442046761513, 0.17298617959022522, 0.15517692267894745, 0.13362783193588257, 0.010580360889434814, 0.00023049254377838224, 0.00745873898267746, 0.00016025979130063206, 0.002226235345005989, 0.0004258991975802928, 0.000578688399400562, 0.0014760587364435196, 0.002039685845375061, 0.0048048608005046844, 0.019996320828795433, 0.0029125709552317858, 0.006709430366754532, 0.0017099445685744286, 0.02097223326563835, 0.0024284888058900833, 0.10361000150442123, 0.022238893434405327, 0.009704988449811935, 0.017071064561605453, 0.011506098322570324, 0.0406200997531414, 0.0063119689002633095, 0.36112311482429504, NaN], [0.050736088305711746, 0.10139954090118408, 0.08949553966522217, 0.0938185378909111, 0.06053004041314125, 0.18139560520648956, 0.0767659917473793, 0.11340610682964325, 0.19499026238918304, 0.11419404298067093, 0.23666803538799286, 0.05730360746383667, 0.07293370366096497, 0.11558260023593903, 0.12613430619239807, 0.07011571526527405, 0.029766615480184555, 0.05616272985935211, 0.02569880336523056, 0.02553572878241539, 0.010698755271732807, 0.02022577077150345, 0.01824677176773548, 0.03918607532978058, 0.034657131880521774, 0.11515442281961441, 0.05569382756948471, 0.035370998084545135, 0.047812946140766144, 0.1140216588973999, 0.018943075090646744, 0.09709078818559647, 0.08172454684972763, 0.04602199047803879, 0.02941049635410309, 0.031383853405714035, 0.10708537697792053, 0.012693268246948719, 0.07050468772649765, 0.25427982211112976]], [[0.04456469416618347, 0.016716457903385162, 0.08688971400260925, 0.23432573676109314, 0.12769784033298492, 0.0498066172003746, 0.10501405596733093, 0.14398211240768433, 0.3055479824542999, 0.0823235884308815, 0.23467087745666504, 0.6305257678031921, 0.08790664374828339, 0.14063040912151337, 0.13028757274150848, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.04107241332530975, 0.03620494529604912, 0.07322828471660614, 0.1027759537100792, 0.08743055909872055, 0.016458408907055855, 0.09779228270053864, 0.014780157245695591, 0.09821301698684692, 0.025402111932635307, 0.0808086097240448, 0.08257035166025162, 0.07231960445642471, 0.0895148441195488, 0.19708459079265594, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1263897716999054, 0.01533158216625452, 0.08717449009418488, 0.22571881115436554, 0.06928549706935883, 0.16778334975242615, 0.06136450543999672, 0.07180161774158478, 0.2525678873062134, 0.32249853014945984, 0.08566119521856308, 0.48726531863212585, 0.2929263114929199, 0.21127133071422577, 0.12448348850011826, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1481804996728897, 0.04817945510149002, 0.03058626689016819, 0.13171793520450592, 0.10783855617046356, 0.24912205338478088, 0.1342363804578781, 0.28650397062301636, 0.25943103432655334, 0.2756144404411316, 0.08422903716564178, 0.7444766163825989, 0.7611673474311829, 0.5739472508430481, 0.11213001608848572, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1744699776172638, 0.050404343754053116, 0.018338145688176155, 0.11463086307048798, 0.02370826154947281, 0.09417468309402466, 0.04503462836146355, 0.0389062762260437, 0.1780962496995926, 0.7825090885162354, 0.15977078676223755, 0.2598268687725067, 0.05674973130226135, 0.2742767333984375, 0.15589554607868195, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.26428407430648804, 0.0871720165014267, 0.015494171530008316, 0.31054598093032837, 0.31179672479629517, 0.05687993764877319, 0.05327969416975975, 0.14049863815307617, 0.03721972927451134, 0.33735793828964233, 0.06669215857982635, 0.44665512442588806, 0.1105320155620575, 0.07633788883686066, 0.13637836277484894, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.27871736884117126, 0.07987862080335617, 0.06999076902866364, 0.3873903453350067, 0.3669894337654114, 0.0245819091796875, 0.02483827993273735, 0.08571609854698181, 0.04856930300593376, 0.2826782464981079, 0.10519464313983917, 0.8515737056732178, 0.24991582334041595, 0.08752243965864182, 0.1076057106256485, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.18780259788036346, 0.02093103528022766, 0.1730981320142746, 0.27918383479118347, 0.32355740666389465, 0.05090703070163727, 0.030107326805591583, 0.015694553032517433, 0.08293543756008148, 0.11989035457372665, 0.1594303995370865, 0.6402391195297241, 0.08334839344024658, 0.13423335552215576, 0.16886292397975922, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.23048973083496094, 0.05534357205033302, 0.15910016000270844, 0.5473513603210449, 0.11114095151424408, 0.060548413544893265, 0.23547381162643433, 0.0231330469250679, 0.22654443979263306, 0.16574865579605103, 0.03383632004261017, 0.05167527496814728, 0.026772163808345795, 0.028301218524575233, 0.08144620060920715, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.126570925116539, 0.0055835917592048645, 0.7687394022941589, 0.6136845350265503, 0.7887718677520752, 0.24027548730373383, 0.25543272495269775, 0.017155619338154793, 0.01121050026267767, 0.02180907502770424, 0.06387564539909363, 0.04227403923869133, 0.004662328865379095, 0.0204116590321064, 0.16526305675506592, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3619309663772583, 0.022692076861858368, 0.8739812970161438, 0.5600091814994812, 0.4330839216709137, 0.27864721417427063, 0.1654776781797409, 0.02327956072986126, 0.003977042157202959, 0.0664801374077797, 0.12084753066301346, 0.16815124452114105, 0.07773539423942566, 0.17824198305606842, 0.05263833701610565, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.29354482889175415, 0.16078433394432068, 0.705570638179779, 0.44417092204093933, 0.02176845259964466, 0.15997210144996643, 0.4057019054889679, 0.11617531627416611, 0.010741903446614742, 0.06882698833942413, 0.07046788930892944, 0.041601523756980896, 0.011864392086863518, 0.06714706867933273, 0.14988133311271667, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5400083065032959, 0.2319646179676056, 0.6198285818099976, 0.2858767509460449, 0.1694929450750351, 0.06001640111207962, 0.26940232515335083, 0.06411167979240417, 0.02847147174179554, 0.18856319785118103, 0.05879069119691849, 0.03795049339532852, 0.009596540592610836, 0.023393897339701653, 0.14663995802402496, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6488012075424194, 0.15997910499572754, 0.6486002802848816, 0.4859846830368042, 0.34752336144447327, 0.028076842427253723, 0.12281371653079987, 0.019826101139187813, 0.023531395941972733, 0.15743687748908997, 0.059922393411397934, 0.08707788586616516, 0.005486410576850176, 0.025385212153196335, 0.15706156194210052, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.037294961512088776, 0.2018004208803177, 0.33537882566452026, 0.19571122527122498, 0.0998593419790268, 0.48263466358184814, 0.11429780721664429, 0.20324908196926117, 0.7053001523017883, 0.01905757561326027, 0.1765546351671219, 0.10779165476560593, 0.18456625938415527, 0.16855330765247345, 0.014784654602408409, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.1489560306072235, 0.2212677150964737, 0.055408962070941925, 0.03110104240477085, 0.02513720653951168, 0.07830048352479935, 0.05067736655473709, 0.06611648201942444, 0.02238955721259117, 0.03719142824411392, 0.025896798819303513, 0.04350690543651581, 0.11618120968341827, 0.08714473247528076, 0.15466241538524628, 0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.002932992298156023, 0.307859867811203, 0.008187332190573215, 0.003677746979519725, 0.0005738585605286062, 0.0008406178676523268, 0.0005446207360364497, 0.00039283244404941797, 0.0009221792570315301, 0.000758469570428133, 0.003933709114789963, 0.0009352274937555194, 0.001059120986610651, 0.0020118390675634146, 0.010183396749198437, 0.1627129465341568, 0.03836298733949661, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.37297555804252625, 0.09208715707063675, 0.16802547872066498, 0.11860792338848114, 0.08042033761739731, 0.18612971901893616, 0.45423436164855957, 0.07133221626281738, 0.13892753422260284, 0.3810507357120514, 0.291797935962677, 0.16154640913009644, 0.050885219126939774, 0.10468144714832306, 0.10335776954889297, 0.23664157092571259, 0.02332315407693386, 0.0017523575806990266, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.028274476528167725, 0.018124615773558617, 0.13954800367355347, 0.03560209274291992, 0.08428613841533661, 0.17491763830184937, 0.13035845756530762, 0.0214189775288105, 0.009060325101017952, 0.012400318868458271, 0.031279344111680984, 0.011209131218492985, 0.19533281028270721, 0.012452301569283009, 0.020085560157895088, 0.14284735918045044, 0.19342879951000214, 0.5212197303771973, 0.028613613918423653, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11180772632360458, 0.012462746351957321, 0.04844700172543526, 0.06198285147547722, 0.06685204058885574, 0.44600817561149597, 0.30352795124053955, 0.1519387811422348, 0.003835479263216257, 0.08384031802415848, 0.027865614742040634, 0.159846231341362, 0.46423590183258057, 0.09249147027730942, 0.09178084880113602, 0.022152410820126534, 0.06252314150333405, 0.005122532602399588, 0.24202540516853333, 0.0027534610126167536, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04840230569243431, 0.026793736964464188, 0.1120820939540863, 0.09037120640277863, 0.2328549474477768, 0.1063276007771492, 0.14073747396469116, 0.19612964987754822, 0.1904316544532776, 0.10354755818843842, 0.10268037766218185, 0.13820117712020874, 0.3374333083629608, 0.15443934500217438, 0.12536528706550598, 0.04657726734876633, 0.23517371714115143, 0.03296450525522232, 0.2014523595571518, 0.06359406560659409, 0.0884864553809166, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.36786824464797974, 0.056283749639987946, 0.03846094757318497, 0.07181648164987564, 0.03666122257709503, 0.04024837538599968, 0.5659748911857605, 0.2338860183954239, 0.11518415063619614, 0.3659259080886841, 0.04107162728905678, 0.012827688828110695, 0.0609581284224987, 0.02837788313627243, 0.060403015464544296, 0.05186963453888893, 0.02286554127931595, 0.21517929434776306, 0.12055587023496628, 0.1711670458316803, 0.27492430806159973, 0.27398592233657837, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0033490851055830717, 0.001678164815530181, 0.02563566155731678, 0.028815647587180138, 0.007257265504449606, 0.04370535537600517, 0.026118090376257896, 0.435838907957077, 0.005564961116760969, 0.014266176149249077, 0.018343305215239525, 0.0009297388605773449, 0.03809681162238121, 0.020595146343111992, 0.03566184639930725, 0.020278872922062874, 0.02308776043355465, 0.022820638492703438, 0.18259893357753754, 0.3133871257305145, 0.08183155953884125, 0.35655686259269714, 0.17295894026756287, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.34718528389930725, 0.028826624155044556, 0.05378839746117592, 0.0680842474102974, 0.0254778191447258, 0.1994519978761673, 0.7739751935005188, 0.28213825821876526, 0.24756361544132233, 0.3363908529281616, 0.08445209264755249, 0.0067241075448691845, 0.09118638187646866, 0.04656682163476944, 0.0331079363822937, 0.057175230234861374, 0.2799927890300751, 0.10977934300899506, 0.4680712819099426, 0.08838099986314774, 0.05264464393258095, 0.21108192205429077, 0.08241217583417892, 0.0764400064945221, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06212884560227394, 0.013463910669088364, 0.024143628776073456, 0.025745615363121033, 0.12165382504463196, 0.04105379059910774, 0.21918880939483643, 0.12444313615560532, 0.7241542935371399, 0.2624671459197998, 0.05330171436071396, 0.026902005076408386, 0.04947282373905182, 0.06268218904733658, 0.04105047509074211, 0.17679302394390106, 0.30970489978790283, 0.042192552238702774, 0.2463400512933731, 0.032756272703409195, 0.05394153669476509, 0.02321716584265232, 0.30038926005363464, 0.023974716663360596, 0.0257905051112175, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.23139908909797668, 0.12510670721530914, 0.062008026987314224, 0.06357982009649277, 0.21447335183620453, 0.06672460585832596, 0.5059712529182434, 0.23151132464408875, 0.3211345672607422, 0.29274967312812805, 0.07394816726446152, 0.12323616445064545, 0.33240705728530884, 0.13292434811592102, 0.0974365845322609, 0.1864403486251831, 0.03811780363321304, 0.18074536323547363, 0.08396673202514648, 0.026499373838305473, 0.05736878141760826, 0.274480402469635, 0.10284627228975296, 0.15606749057769775, 0.017497936263680458, 0.09719526022672653, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3976813554763794, 0.24336650967597961, 0.030069073662161827, 0.04866141080856323, 0.061815883964300156, 0.023062149062752724, 0.2837987542152405, 0.10572359710931778, 0.42220908403396606, 0.47088485956192017, 0.06114182993769646, 0.05295940861105919, 0.04274435341358185, 0.033208493143320084, 0.07069624215364456, 0.1767420768737793, 0.017465414479374886, 0.034512054175138474, 0.0999627411365509, 0.011741198599338531, 0.022724410519003868, 0.04408577084541321, 0.03894393891096115, 0.018038587644696236, 0.058924250304698944, 0.2522818148136139, 0.12782295048236847, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.6213744282722473, 0.08501708507537842, 0.08457361906766891, 0.0819045826792717, 0.02008524350821972, 0.02321169711649418, 0.5481746196746826, 0.17061969637870789, 0.19314314424991608, 0.48946020007133484, 0.08799289166927338, 0.009451461024582386, 0.1643926501274109, 0.03458939492702484, 0.0487554594874382, 0.042104240506887436, 0.022070694714784622, 0.04743226245045662, 0.13338083028793335, 0.020831480622291565, 0.031267598271369934, 0.024703562259674072, 0.041907425969839096, 0.006121364887803793, 0.02875565178692341, 0.13002096116542816, 0.36194902658462524, 0.021867850795388222, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11498570442199707, 0.014700047671794891, 0.04425002261996269, 0.027370423078536987, 0.031341005116701126, 0.11119254678487778, 0.2834031581878662, 0.24822625517845154, 0.387948602437973, 0.17188440263271332, 0.026020031422376633, 0.003112945705652237, 0.1680845320224762, 0.013143973425030708, 0.05647796019911766, 0.12623563408851624, 0.6370776891708374, 0.07802888005971909, 0.06076015904545784, 0.015353387221693993, 0.0031011439859867096, 0.031844403594732285, 0.5665289163589478, 0.013176449574530125, 0.025442441925406456, 0.05083877220749855, 0.08586791157722473, 0.03281332179903984, 0.0019294946687296033, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00710845272988081, 0.009718026034533978, 0.08296849578619003, 0.05356726795434952, 0.20372402667999268, 0.20898059010505676, 0.07373131066560745, 0.07588774710893631, 0.33318811655044556, 0.09730548411607742, 0.031877510249614716, 0.04629351943731308, 0.026428943499922752, 0.05165233090519905, 0.12934288382530212, 0.010483458638191223, 0.10243765264749527, 0.013204336166381836, 0.1070198118686676, 0.001742976950481534, 0.0011925535509362817, 0.03764529153704643, 0.023008054122328758, 0.09038762003183365, 0.1208486333489418, 0.06097627431154251, 0.11476689577102661, 0.17706690728664398, 0.4447736442089081, 0.005561552010476589, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.092291921377182, 0.13057716190814972, 0.11971572786569595, 0.09643372148275375, 0.0971774011850357, 0.03882397338747978, 0.30341219902038574, 0.06688009947538376, 0.5493715405464172, 0.21897412836551666, 0.10454282909631729, 0.09917838126420975, 0.19730664789676666, 0.0889393612742424, 0.0462181456387043, 0.03962688520550728, 0.412600040435791, 0.1027907133102417, 0.011060677468776703, 0.04006139934062958, 0.005457504652440548, 0.17391063272953033, 0.009697728790342808, 0.08243320137262344, 0.1504840850830078, 0.029468167573213577, 0.29366523027420044, 0.04788699373602867, 0.17640100419521332, 0.04229334741830826, 0.3300667107105255, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3365032970905304, 0.06134270504117012, 0.11965256929397583, 0.08703643828630447, 0.08615697175264359, 0.01610170491039753, 0.289604127407074, 0.16905160248279572, 0.690265953540802, 0.5125291347503662, 0.11020015180110931, 0.05034353584051132, 0.04973014071583748, 0.04155145213007927, 0.06180096045136452, 0.20544184744358063, 0.06503231078386307, 0.21778742969036102, 0.04011436551809311, 0.2470238208770752, 0.03102266602218151, 0.027881061658263206, 0.06887322664260864, 0.023802783340215683, 0.2166331559419632, 0.06618232280015945, 0.058350641280412674, 0.04297764599323273, 0.06574989855289459, 0.02652076631784439, 0.08339553326368332, 0.09817715734243393, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.25151577591896057, 0.0737723708152771, 0.11452356725931168, 0.07270905375480652, 0.27380475401878357, 0.046423640102148056, 0.6668940782546997, 0.60158771276474, 0.286392480134964, 0.2904633581638336, 0.07359147071838379, 0.040276750922203064, 0.2706137001514435, 0.15532110631465912, 0.051646988838911057, 0.09466058760881424, 0.0047309016808867455, 0.1481417566537857, 0.06127317249774933, 0.015202163718640804, 0.011932089924812317, 0.31230586767196655, 0.04852164536714554, 0.039501819759607315, 0.001117925625294447, 0.06312739849090576, 0.023924386128783226, 0.02860989049077034, 0.007241260260343552, 0.11453913897275925, 0.012237192131578922, 0.2803768217563629, 0.0480632521212101, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4344438314437866, 0.2159019559621811, 0.0411386713385582, 0.059745997190475464, 0.08364511281251907, 0.02960371784865856, 0.3908357322216034, 0.17347759008407593, 0.4736940562725067, 0.5831181406974792, 0.08143209666013718, 0.05496616289019585, 0.0508774034678936, 0.03704635798931122, 0.07529113441705704, 0.02001449465751648, 0.0017837424529716372, 0.005722085013985634, 0.04321253299713135, 0.00430489843711257, 0.009005578234791756, 0.010736249387264252, 0.0058517144061625, 0.003792154835537076, 0.008828205987811089, 0.0838593989610672, 0.029530486091971397, 0.015579215250909328, 0.010320665314793587, 0.016853220760822296, 0.017335176467895508, 0.12552303075790405, 0.42354699969291687, 0.08326870948076248, NaN, NaN, NaN, NaN, NaN, NaN], [0.6010525822639465, 0.07716702669858932, 0.12942874431610107, 0.11651009321212769, 0.029510293155908585, 0.025635747238993645, 0.564699649810791, 0.20346374809741974, 0.1942133754491806, 0.5329980254173279, 0.09726559370756149, 0.006782675161957741, 0.1884276419878006, 0.02957840822637081, 0.046941183507442474, 0.001771818962879479, 0.000807587115559727, 0.0031146325636655092, 0.023062998428940773, 0.0018312688916921616, 0.007724495604634285, 0.002569216303527355, 0.003803644794970751, 0.00041838324978016317, 0.001987496856600046, 0.012477965094149113, 0.04809670150279999, 0.0016458284808322787, 0.00020838514319621027, 0.005814890842884779, 0.018183711916208267, 0.30546146631240845, 0.4703490138053894, 0.15369661152362823, 0.012250960804522038, NaN, NaN, NaN, NaN, NaN], [0.07098641246557236, 0.02088714949786663, 0.0536419078707695, 0.04874833673238754, 0.1357380896806717, 0.10192368179559708, 0.22615019977092743, 0.3848302960395813, 0.3569928705692291, 0.19976821541786194, 0.030237246304750443, 0.012232640758156776, 0.14491091668605804, 0.01217038556933403, 0.025625383481383324, 0.02520398050546646, 0.2818087637424469, 0.007948609068989754, 0.07590723037719727, 0.01867567002773285, 0.006826441269367933, 0.011762343347072601, 0.5987983345985413, 0.0045673479326069355, 0.01173742488026619, 0.03130093589425087, 0.03894692659378052, 0.016236862167716026, 0.0014989122282713652, 0.0009245824767276645, 0.025562506169080734, 0.5276230573654175, 0.32699310779571533, 0.1864093542098999, 0.0933799296617508, 0.0060149896889925, NaN, NaN, NaN, NaN], [0.007031308952718973, 0.007269172929227352, 0.08423776179552078, 0.053896792232990265, 0.21268267929553986, 0.2456619292497635, 0.0817742720246315, 0.07338020205497742, 0.2872445285320282, 0.08955906331539154, 0.02503780461847782, 0.043076977133750916, 0.024157537147402763, 0.05127491056919098, 0.1281031221151352, 0.0011320068733766675, 0.011502433568239212, 0.0017513524508103728, 0.020418671891093254, 0.0003008104977197945, 0.00031320590642280877, 0.0053228470496833324, 0.0022876623552292585, 0.011736828833818436, 0.017109515145421028, 0.010937619023025036, 0.015238909050822258, 0.025703608989715576, 0.10705357789993286, 0.0009204442030750215, 0.02667400799691677, 0.16934601962566376, 0.08647502958774567, 0.028284918516874313, 0.06841914355754852, 0.39870724081993103, 0.0010592876933515072, NaN, NaN, NaN], [0.06564409285783768, 0.10634885728359222, 0.14713656902313232, 0.07514703273773193, 0.3204736113548279, 0.07143916934728622, 0.4829144775867462, 0.2612879276275635, 0.7603816986083984, 0.17889906466007233, 0.07189968973398209, 0.10938191413879395, 0.2776612341403961, 0.08681799471378326, 0.052979547530412674, 0.02631283551454544, 0.29101136326789856, 0.042160265147686005, 0.009721376933157444, 0.02933679334819317, 0.014515053480863571, 0.18161341547966003, 0.016545770689845085, 0.03647695854306221, 0.0840071588754654, 0.02240183763206005, 0.1055113896727562, 0.037331126630306244, 0.17535105347633362, 0.010923052206635475, 0.2594170868396759, 0.5064816474914551, 0.06657205522060394, 0.130835622549057, 0.0483754500746727, 0.2870587110519409, 0.010685333050787449, 0.21122200787067413, NaN, NaN], [0.28806957602500916, 0.05887402966618538, 0.12616868317127228, 0.10481040924787521, 0.19247829914093018, 0.033351678401231766, 0.39873749017715454, 0.22540906071662903, 0.7029480338096619, 0.5013188719749451, 0.10523373633623123, 0.08320688456296921, 0.0816955640912056, 0.04881281033158302, 0.09282685816287994, 0.21289733052253723, 0.10400458425283432, 0.2843308448791504, 0.11722961068153381, 0.31265783309936523, 0.07705509662628174, 0.050357937812805176, 0.1631784737110138, 0.04547655209898949, 0.37539371848106384, 0.07925810664892197, 0.07719646394252777, 0.043498191982507706, 0.04735783487558365, 0.022911155596375465, 0.20965908467769623, 0.2452480047941208, 0.05793433263897896, 0.07357832789421082, 0.03363368287682533, 0.041085004806518555, 0.014093895442783833, 0.05045074224472046, 0.0570731945335865, NaN], [0.2559513747692108, 0.07615252584218979, 0.11904845386743546, 0.07934627681970596, 0.09980516135692596, 0.14371442794799805, 0.3059750497341156, 0.09035829454660416, 0.22693291306495667, 0.32864776253700256, 0.08986205607652664, 0.1614997386932373, 0.17624114453792572, 0.16325940191745758, 0.119119793176651, 0.02115148864686489, 0.018139760941267014, 0.03536282852292061, 0.06259438395500183, 0.00901759136468172, 0.014575985260307789, 0.12521256506443024, 0.12870429456233978, 0.09162478893995285, 0.06363746523857117, 0.1348179280757904, 0.07700010389089584, 0.05158444121479988, 0.01101324986666441, 0.03299920633435249, 0.163722425699234, 0.13794326782226562, 0.18303781747817993, 0.117555633187294, 0.08103907853364944, 0.012191864661872387, 0.032527241855859756, 0.16104964911937714, 0.12187117338180542, 0.22321484982967377]]]], \"bot_text\": [\"The_\", \"animal_\", \"didn_\", \"'_\", \"t_\", \"cross_\", \"the_\", \"street_\", \"because_\", \"it_\", \"was_\", \"too_\", \"tire\", \"d_\", \"Das_\", \"Tier\", \"_\", \"\\u00fcber\", \"quer\", \"te_\", \"die_\", \"Stra\\u00dfe_\", \"nicht_\", \", _\", \"weil_\", \"es_\", \"zu_\", \"m\\u00fc\", \"de_\", \"war_\", \", _\", \"weil_\", \"es_\", \"zu_\", \"m\\u00fc\", \"de_\", \"war_\", \"._\"]}, \"out_out\": {\"top_text\": [\"Das_\", \"Tier\", \"_\", \"\\u00fcber\", \"quer\", \"te_\", \"die_\", \"Stra\\u00dfe_\", \"nicht_\", \", _\", \"weil_\", \"es_\", \"zu_\", \"m\\u00fc\", \"de_\", \"war_\", \", _\", \"weil_\", \"es_\", \"zu_\", \"m\\u00fc\", \"de_\", \"war_\", \"._\"], \"att\": [[[[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.33067038655281067, 0.02820705994963646, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.43891066312789917, 0.3106566071510315, 0.006947982590645552, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.8740342259407043, 0.6547167897224426, 0.0062981778755784035, 0.46666401624679565, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.009682492353022099, 0.17458303272724152, 0.7120969891548157, 0.10496775060892105, 0.0038010317366570234, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.31054121255874634, 0.41146165132522583, 0.4573209881782532, 0.639615535736084, 0.038498248904943466, 0.06232544779777527, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2996446192264557, 0.18095439672470093, 0.8072441220283508, 0.6008384227752686, 0.045412980020046234, 0.09029265493154526, 0.15878555178642273, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07671086490154266, 0.13175785541534424, 0.032809216529130936, 0.06887537240982056, 0.32570284605026245, 0.22846734523773193, 0.06983717530965805, 0.07415641844272614, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4443431496620178, 0.2924090623855591, 0.09237049520015717, 0.07077033072710037, 0.05661908909678459, 0.1886560618877411, 0.5792031288146973, 0.23326165974140167, 0.024399278685450554, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0045473226346075535, 0.015263181179761887, 0.11153102666139603, 0.01091472152620554, 0.07137833535671234, 0.14599360525608063, 0.24649137258529663, 0.2676219940185547, 0.14942915737628937, 0.03359955921769142, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0021246292162686586, 0.019146723672747612, 0.0190261360257864, 0.004887872841209173, 0.032842181622982025, 0.009469296783208847, 0.015122202225029469, 0.056959331035614014, 0.014146327041089535, 0.2864534854888916, 0.028167642652988434, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.007321672048419714, 0.06949152052402496, 0.18409577012062073, 0.05168240889906883, 0.5332358479499817, 0.12983477115631104, 0.020923368632793427, 0.015086837112903595, 0.05491120368242264, 0.38865622878074646, 0.036598365753889084, 0.02645716816186905, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.004608431365340948, 0.07759333401918411, 0.05611182749271393, 0.031112710013985634, 0.06043193116784096, 0.023203425109386444, 0.01299421489238739, 0.011212858371436596, 0.2615091800689697, 0.5089370608329773, 0.22289350628852844, 0.10276756435632706, 0.03959360718727112, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.012221934273838997, 0.040381401777267456, 0.0694599524140358, 0.0800129845738411, 0.023234205320477486, 0.003881127340719104, 0.03062801994383335, 0.024260450154542923, 0.012832778505980968, 0.01656900905072689, 0.2333584874868393, 0.3572527766227722, 0.0072386497631669044, 0.014752739109098911, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09144259989261627, 0.1256924569606781, 0.6557105779647827, 0.1641494482755661, 0.04417502135038376, 0.42902442812919617, 0.377028226852417, 0.1956152766942978, 0.27481555938720703, 0.37677863240242004, 0.4323487877845764, 0.6219720244407654, 0.3997260332107544, 0.1145903542637825, 0.041462015360593796, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5997433662414551, 0.1045081838965416, 0.10960735380649567, 0.047688476741313934, 0.31575047969818115, 0.1532202959060669, 0.4197675585746765, 0.16546213626861572, 0.31973955035209656, 0.23332525789737701, 0.15541672706604004, 0.05988143011927605, 0.5733460187911987, 0.8565582036972046, 0.009604076854884624, 0.030047349631786346, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02339007519185543, 0.01581897959113121, 0.02374129369854927, 0.02252129279077053, 0.08995510637760162, 0.0626068115234375, 0.27313846349716187, 0.036778680980205536, 0.22608895599842072, 0.06801939755678177, 0.035735905170440674, 0.022851483896374702, 0.06078701093792915, 0.42404335737228394, 0.41984546184539795, 0.08353053033351898, 0.058427464216947556, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.034203190356492996, 0.23458202183246613, 0.15632590651512146, 0.02520577609539032, 0.26413342356681824, 0.06292548030614853, 0.06378099322319031, 0.08676797896623611, 0.02988903410732746, 0.3430734872817993, 0.007843950763344765, 0.03405369073152542, 0.01887335814535618, 0.39618176221847534, 0.2528276741504669, 0.10531513392925262, 0.12583006918430328, 0.09389571845531464, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.009769688360393047, 0.056299567222595215, 0.11172951757907867, 0.02802591770887375, 0.3647110164165497, 0.09813904017210007, 0.016619421541690826, 0.006417513824999332, 0.016537560150027275, 0.15495160222053528, 0.023067951202392578, 0.011397394351661205, 0.029141509905457497, 0.0527399443089962, 0.2784731984138489, 0.059669919312000275, 0.5969582796096802, 0.09549567103385925, 0.03235183656215668, NaN, NaN, NaN, NaN, NaN, NaN], [0.00987912341952324, 0.12349259853363037, 0.037169262766838074, 0.01944275200366974, 0.06324917078018188, 0.02598830871284008, 0.020618943497538567, 0.009103300981223583, 0.1360517293214798, 0.09789924323558807, 0.06809242814779282, 0.12332575768232346, 0.034675393253564835, 0.16954950988292694, 0.010956126265227795, 0.11111389100551605, 0.1871008574962616, 0.2434563934803009, 0.10274684429168701, 0.0379486046731472, NaN, NaN, NaN, NaN, NaN], [0.010987702757120132, 0.03791751340031624, 0.03792046010494232, 0.0400051474571228, 0.008841714821755886, 0.002161285374313593, 0.031619150191545486, 0.01907121017575264, 0.0057282340712845325, 0.002385619329288602, 0.03308374434709549, 0.11032091826200485, 0.0044158026576042175, 0.05701944977045059, 0.0651637390255928, 0.027267253026366234, 0.3151875138282776, 0.17881636321544647, 0.3164456784725189, 0.005250148009508848, 0.011875288560986519, NaN, NaN, NaN, NaN], [0.08034691959619522, 0.1792650669813156, 0.6813479661941528, 0.11697664856910706, 0.022037051618099213, 0.4362119436264038, 0.3332834541797638, 0.16648675501346588, 0.3133866786956787, 0.21180157363414764, 0.22306133806705475, 0.5634312033653259, 0.2539531886577606, 0.28583550453186035, 0.0421890914440155, 0.24185270071029663, 0.9185315370559692, 0.5444227457046509, 0.7130873799324036, 0.36675870418548584, 0.1082441657781601, 0.02894955314695835, NaN, NaN, NaN], [0.3316553831100464, 0.07297243922948837, 0.18084223568439484, 0.0543624572455883, 0.141310915350914, 0.15985439717769623, 0.22593949735164642, 0.09976530820131302, 0.2670679986476898, 0.12590403854846954, 0.10189743340015411, 0.06066418066620827, 0.14688965678215027, 0.6279550790786743, 0.004891595803201199, 0.013660040684044361, 0.19539086520671844, 0.13336770236492157, 0.11226529628038406, 0.4554508626461029, 0.7914823293685913, 0.007615156006067991, 0.015521766617894173, NaN, NaN], [0.010082974098622799, 0.009416572749614716, 0.026376336812973022, 0.021534079685807228, 0.041008636355400085, 0.028814975172281265, 0.09862472116947174, 0.019531887024641037, 0.1915404349565506, 0.055525705218315125, 0.03489372506737709, 0.035597167909145355, 0.017297467216849327, 0.13875839114189148, 0.18795406818389893, 0.13025526702404022, 0.03705297037959099, 0.016517892479896545, 0.028779756277799606, 0.02632485330104828, 0.36631691455841064, 0.4771501123905182, 0.10461407899856567, 0.07566797733306885, NaN], [0.00671275844797492, 0.019956005737185478, 0.15321078896522522, 0.00987993273884058, 0.1430601179599762, 0.02432059310376644, 0.007838046178221703, 0.016839532181620598, 0.017622128129005432, 0.03075602278113365, 0.01907699555158615, 0.30206096172332764, 0.010013632476329803, 0.06018203869462013, 0.19546428322792053, 0.020215312018990517, 0.04091925173997879, 0.022548291832208633, 0.26572445034980774, 0.010653333738446236, 0.1212434321641922, 0.3668496906757355, 0.1586136817932129, 0.14579400420188904, 0.04911552369594574]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00017037145153153688, 0.1837475299835205, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [4.619961600837996e-06, 0.00011092388740507886, 0.19595862925052643, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [7.402049959637225e-07, 0.0014410031726583838, 0.15330694615840912, 0.0009438465931452811, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [6.564930572494632e-07, 1.2471617083065212e-05, 0.0012651559663936496, 1.2094314115529414e-05, 0.2683168947696686, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [3.960849710438197e-07, 2.835777740983758e-05, 0.0015905762556940317, 5.72201497561764e-05, 0.20671997964382172, 0.03618929535150528, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [3.613545777625404e-05, 4.069158967467956e-05, 0.0019799659494310617, 4.598083614837378e-05, 0.28016433119773865, 0.1021510660648346, 0.0019787675701081753, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03414154052734375, 0.018152736127376556, 0.002861178945749998, 0.0031036457512527704, 0.2743661403656006, 0.08905426412820816, 0.058365415781736374, 0.2834230065345764, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0001288916973862797, 0.0019113116431981325, 0.0011359998025000095, 2.5460678443778306e-05, 0.0018093753606081009, 0.008086470887064934, 0.005666371434926987, 0.0014489549212157726, 0.27176737785339355, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0013363973703235388, 0.015213730745017529, 0.019847076386213303, 0.0016770424554124475, 0.6085457801818848, 0.051846977323293686, 0.06904839724302292, 0.023163089528679848, 0.0024616841692477465, 0.4075135886669159, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [1.5705205441918224e-05, 0.00011942459968850017, 3.308789018774405e-05, 0.00047703171730972826, 1.5581523257424124e-05, 3.566192026482895e-05, 0.000621139828581363, 0.002513762330636382, 0.0013953398447483778, 0.001656065694987774, 0.6708395481109619, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0009777048835530877, 0.006719581317156553, 0.017090875655412674, 0.007835427299141884, 0.0003081739123445004, 0.0027951891534030437, 0.0031432590913027525, 0.011542102321982384, 0.01903962530195713, 0.032312098890542984, 0.23448777198791504, 0.18604722619056702, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0010771078523248434, 0.00013067253166809678, 0.0004810431564692408, 0.0005832655006088316, 0.27172601222991943, 0.023587899282574654, 0.0011203349567949772, 0.0001570776366861537, 3.2636336982250214e-05, 0.008125105872750282, 0.3860749900341034, 0.011222672648727894, 0.4488545358181, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0018897228874266148, 0.00010004806244978681, 0.040837980806827545, 0.0009045379119925201, 0.4036760926246643, 0.033945482224226, 0.0009020724683068693, 2.477952148183249e-05, 0.0006147518288344145, 2.3498352675233036e-05, 0.0003015661786776036, 0.00019162058015353978, 0.0013656887458637357, 0.9207848906517029, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [3.0049262932152487e-05, 0.00032340767211280763, 0.0004620190302375704, 1.456133759347722e-05, 0.4214256703853607, 0.00038119935197755694, 2.2086916942498647e-05, 5.437946310848929e-05, 0.0005922063137404621, 0.0002251591213280335, 4.171442924416624e-05, 0.0011568808695301414, 6.667344860034063e-05, 0.004539569839835167, 0.07099039107561111, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0001142411565524526, 0.001007341779768467, 0.5582761764526367, 0.0006983705679886043, 0.04208780825138092, 0.07311324775218964, 0.011010478250682354, 0.00018356108921580017, 0.11227726191282272, 1.5535662896581925e-05, 7.865564111853018e-05, 8.497068483848125e-05, 0.007107958197593689, 0.04726947844028473, 0.03816111385822296, 0.7400538921356201, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [9.270196460420266e-05, 0.00014002913667354733, 0.006266205105930567, 8.287983655463904e-05, 0.029540851712226868, 0.019505193457007408, 0.0002005908900173381, 0.0002361711667617783, 0.002089217072352767, 0.0007247799658216536, 0.0003387654141988605, 3.3522373996675014e-05, 0.00015295531193260103, 0.005682599265128374, 0.01914886385202408, 0.006167547311633825, 0.6065680980682373, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.017243418842554092, 0.0717378556728363, 0.015470567159354687, 0.14577892422676086, 0.003815611358731985, 0.01656431145966053, 0.21609994769096375, 0.24452562630176544, 0.07360902428627014, 0.020440302789211273, 0.9522358775138855, 0.0012982342159375548, 0.00034142163349315524, 4.905217429040931e-05, 0.0002677988959476352, 0.0020047405268996954, 0.013444142416119576, 0.5238149166107178, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.006589227356016636, 0.025933612138032913, 0.05151839554309845, 0.019538801163434982, 0.000567624403629452, 0.011064885184168816, 0.018599001690745354, 0.0389220230281353, 0.03263486549258232, 0.03920944407582283, 0.309482604265213, 0.18455958366394043, 0.0028949796687811613, 0.0009189100819639862, 0.01304793544113636, 0.01903691701591015, 0.0013186958385631442, 0.1459255963563919, 0.2617945969104767, NaN, NaN, NaN, NaN, NaN, NaN], [0.000940846570301801, 6.996696902206168e-05, 0.0001185448418254964, 0.00013115631008986384, 0.04620806872844696, 0.009408986195921898, 0.0010798430303111672, 0.00010642426059348509, 1.4586596989829559e-05, 0.0008147742482833564, 0.049950405955314636, 0.0020658469293266535, 0.020368386059999466, 0.0015965981874614954, 0.0005227082292549312, 8.089001494226977e-05, 0.42970454692840576, 0.3893451988697052, 0.006195466499775648, 0.2630486488342285, NaN, NaN, NaN, NaN, NaN], [0.0015646422980353236, 5.644361226586625e-05, 0.015588155947625637, 0.0004337269929237664, 0.061090677976608276, 0.015012362040579319, 0.0009935805574059486, 3.2441483199363574e-05, 0.0006383971776813269, 7.901599929027725e-06, 0.00011085882579209283, 2.031324947893154e-05, 0.0001886440732050687, 0.1558367908000946, 2.918860081990715e-05, 0.00031420652521774173, 3.769064642256126e-05, 0.000311522075207904, 8.488001913065091e-05, 0.001447036280296743, 0.9016569256782532, NaN, NaN, NaN, NaN], [6.329882307909429e-05, 0.0007932570297271013, 0.0008974742377176881, 3.545067738741636e-05, 0.41645264625549316, 0.0012166639789938927, 5.162824527360499e-05, 0.00016062096983660012, 0.0028807471971958876, 0.0007734368555247784, 0.0001738688733894378, 0.0017386887921020389, 8.449772576568648e-05, 0.008313576690852642, 0.04833607003092766, 5.605717160506174e-05, 0.000497612461913377, 0.00019103533122688532, 0.0018799308454617858, 0.000193181011127308, 0.010939341969788074, 0.11687301844358444, NaN, NaN, NaN], [2.7039888664148748e-05, 0.0002653435221873224, 0.3520841896533966, 0.0011641159653663635, 0.017258664593100548, 0.13898366689682007, 0.004804374184459448, 0.0001136215214501135, 0.10132589936256409, 1.9021857951884158e-05, 0.00018713112513069063, 5.577637057285756e-05, 0.0021825090516358614, 0.016621561720967293, 0.003813497256487608, 0.05257569998502731, 7.136658678064123e-05, 0.00013083907833788544, 8.304342918563634e-05, 0.009517401456832886, 0.07102376222610474, 0.0242641419172287, 0.791592538356781, NaN, NaN], [1.8426982933306135e-05, 6.735812348779291e-05, 0.005383457988500595, 0.0002568464260548353, 0.03709089383482933, 0.05173188075423241, 0.00015440442075487226, 0.00026214553508907557, 0.0031172526068985462, 0.0018413036596029997, 0.001364374067634344, 0.0001026472236844711, 0.00015940713637974113, 0.00464483629912138, 0.007250420283526182, 0.006640422623604536, 0.10042263567447662, 0.00037284562131389976, 5.502302519744262e-05, 0.00017516437219455838, 0.013823487795889378, 0.028728578239679337, 0.014491567388176918, 0.5602642297744751, NaN], [1.3810687960358337e-05, 0.0002572945086285472, 0.008041280321776867, 0.00040080497274175286, 0.00010326507617719471, 0.0013340600999072194, 0.00019016038277186453, 0.00019489554688334465, 0.0007417663000524044, 0.0012533330591395497, 0.0032668926287442446, 0.001072657760232687, 5.286548912408762e-05, 4.225512952871213e-07, 1.0035311788669787e-05, 2.1279807697283104e-05, 0.0006032216479070485, 0.00048016011714935303, 0.00037273563793860376, 3.447151175350882e-05, 9.715819260236458e-07, 2.8930742701049894e-05, 0.0003854547976516187, 0.005018792115151882, 0.4505775570869446]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [4.347301455709385e-06, 0.18382565677165985, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0001576173526700586, 0.00605444610118866, 0.19315025210380554, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0015271879965439439, 0.2696094512939453, 0.0976908802986145, 0.19172586500644684, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.018620789051055908, 0.1513659805059433, 0.1261996626853943, 0.04123798385262489, 0.18324223160743713, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [7.739824650343508e-05, 0.0007302183075807989, 0.0020413347519934177, 0.0010007238015532494, 0.20195050537586212, 0.04546361416578293, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0007431988487951458, 0.330532044172287, 0.08558935672044754, 0.06556878238916397, 0.10690004378557205, 0.1145712360739708, 0.06475446373224258, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.015635214745998383, 0.050190601497888565, 0.02352251298725605, 0.24284599721431732, 0.06325101107358932, 0.02171560376882553, 0.015677697956562042, 0.4775830805301666, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03602181747555733, 0.2262161672115326, 0.11374488472938538, 0.22297167778015137, 0.018925879150629044, 0.2400040328502655, 0.13629396259784698, 0.14897051453590393, 0.11721047759056091, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.001669732853770256, 0.0008830919396132231, 0.007873992435634136, 0.004793200176209211, 0.032567575573921204, 0.019068563356995583, 0.01167156733572483, 0.006520072463899851, 0.001765590044669807, 0.479371041059494, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04264334216713905, 0.01628556102514267, 0.012549073435366154, 0.1270730197429657, 0.09553729742765427, 0.12904676795005798, 0.28088441491127014, 0.08353402465581894, 0.19219043850898743, 0.1467161476612091, 0.04815742373466492, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.006975929252803326, 0.05510300025343895, 0.007132354192435741, 0.0349782258272171, 0.02191060781478882, 0.018211986869573593, 0.026551326736807823, 0.03648876026272774, 0.06464254856109619, 0.049987878650426865, 0.05908217281103134, 0.5448521375656128, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.000807860866189003, 0.00374230626039207, 0.004482839722186327, 0.005506760906428099, 0.000447272410383448, 0.003816538956016302, 0.03234753757715225, 0.014306235127151012, 0.01718331128358841, 0.04840204864740372, 0.06595310568809509, 0.18900929391384125, 0.0723472312092781, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00447529973462224, 0.019966747611761093, 0.03737834841012955, 0.3797287940979004, 0.010614297352731228, 0.05463654175400734, 0.32780376076698303, 0.0739898681640625, 0.25606051087379456, 0.8621841073036194, 0.2645638585090637, 0.25103500485420227, 0.016027942299842834, 0.004609693773090839, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0010164460400119424, 0.011448963545262814, 0.03378765657544136, 0.02785181999206543, 0.056788451969623566, 0.07099426537752151, 0.008927138522267342, 0.01755385287106037, 0.039185769855976105, 0.09313513338565826, 0.027632856741547585, 0.12282836437225342, 0.017955774441361427, 0.02453978732228279, 0.267269104719162, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09903331845998764, 0.854941725730896, 0.020280463621020317, 0.8786925673484802, 0.37992238998413086, 0.20425425469875336, 0.32038459181785583, 0.8171603083610535, 0.2503354549407959, 0.7644308805465698, 0.7474347949028015, 0.935006856918335, 0.36836859583854675, 0.03383934497833252, 0.0021248040720820427, 0.21007098257541656, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09584157168865204, 0.00421579135581851, 0.0017077650409191847, 0.0670090913772583, 0.10943465679883957, 0.05715145170688629, 0.03694647178053856, 0.04514404758810997, 0.04956913739442825, 0.07195062190294266, 0.4566742479801178, 0.20942343771457672, 0.1548582911491394, 0.3906869888305664, 0.03925589844584465, 0.005858495831489563, 0.23115697503089905, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10393274575471878, 0.03258725255727768, 0.01998279243707657, 0.13928532600402832, 0.08602269738912582, 0.139993816614151, 0.2561682462692261, 0.08122693002223969, 0.28790318965911865, 0.34215468168258667, 0.023110536858439445, 0.8003224730491638, 0.11519370973110199, 0.5406965613365173, 0.2252652645111084, 0.07071924954652786, 0.03988110274076462, 0.09249765425920486, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.006400381214916706, 0.03668399527668953, 0.006957556586712599, 0.024804070591926575, 0.013962345197796822, 0.010118995793163776, 0.014814852736890316, 0.02360437996685505, 0.038752347230911255, 0.10996780544519424, 0.24877001345157623, 0.7050904035568237, 0.103914275765419, 0.0656881257891655, 0.03925013542175293, 0.0268316138535738, 0.009403076022863388, 0.042995911091566086, 0.38370969891548157, NaN, NaN, NaN, NaN, NaN, NaN], [0.0005728903925046325, 0.0018518416909500957, 0.003297911025583744, 0.002339646453037858, 0.0003125199000351131, 0.0013706001918762922, 0.011640608310699463, 0.005699110683053732, 0.00646078959107399, 0.029403753578662872, 0.09435103088617325, 0.4532504379749298, 0.1454003006219864, 0.08155784755945206, 0.1478416919708252, 0.06988534331321716, 0.07031917572021484, 0.08092489838600159, 0.16178953647613525, 0.09959835559129715, NaN, NaN, NaN, NaN, NaN], [0.007587960455566645, 0.01947515644133091, 0.06775914877653122, 0.37032291293144226, 0.014833947643637657, 0.04509717598557472, 0.2979332506656647, 0.08052700757980347, 0.2017516791820526, 0.8817963004112244, 0.3514429032802582, 0.3636293411254883, 0.14158478379249573, 0.09958238899707794, 0.13573585450649261, 0.27771836519241333, 0.47418463230133057, 0.36210212111473083, 0.2140081375837326, 0.022566867992281914, 0.004614678677171469, NaN, NaN, NaN, NaN], [0.0009141381597146392, 0.00906511303037405, 0.026196878403425217, 0.011460180394351482, 0.03924085199832916, 0.05833837762475014, 0.004696658346801996, 0.009781464003026485, 0.029306253418326378, 0.06398104876279831, 0.017127037048339844, 0.0922316163778305, 0.03436172753572464, 0.12105685472488403, 0.475220263004303, 0.20121201872825623, 0.0066191148944199085, 0.018271028995513916, 0.05732923001050949, 0.018915977329015732, 0.019877590239048004, 0.23682713508605957, NaN, NaN, NaN], [0.14320576190948486, 0.892350971698761, 0.030759859830141068, 0.8051734566688538, 0.7149769067764282, 0.4937312602996826, 0.3181091248989105, 0.8743517994880676, 0.3442763686180115, 0.8711729049682617, 0.7545801997184753, 0.9297782182693481, 0.6998263001441956, 0.17287810146808624, 0.008261360228061676, 0.9148194789886475, 0.7390273213386536, 0.743715763092041, 0.8801547288894653, 0.47275617718696594, 0.02699747122824192, 0.002916275057941675, 0.1803632229566574, NaN, NaN], [0.0431031733751297, 0.0034584910608828068, 0.0008681766339577734, 0.032780423760414124, 0.11873625963926315, 0.03893061354756355, 0.019801655784249306, 0.03132590278983116, 0.05763043835759163, 0.06388700753450394, 0.3317660689353943, 0.16543246805667877, 0.10311393439769745, 0.4146954417228699, 0.09686555713415146, 0.06189668923616409, 0.5733434557914734, 0.2515217959880829, 0.17396190762519836, 0.13145960867404938, 0.40639445185661316, 0.07709264755249023, 0.007335619535297155, 0.2446187138557434, NaN], [0.046706411987543106, 0.31744489073753357, 0.6429179310798645, 0.4889025092124939, 0.43930482864379883, 0.3055577576160431, 0.6935683488845825, 0.25992196798324585, 0.7758384346961975, 0.2076689600944519, 0.8320663571357727, 0.39907822012901306, 0.8469056487083435, 0.5997118353843689, 0.31635957956314087, 0.36650604009628296, 0.2247273474931717, 0.7608639597892761, 0.37947097420692444, 0.8680096864700317, 0.5816919803619385, 0.19056683778762817, 0.27210569381713867, 0.06685535609722137, 0.040061503648757935]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17503570020198822, 0.10145211219787598, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.002467370592057705, 0.014373218640685081, 0.18901397287845612, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [4.782021278515458e-05, 0.0002036100922850892, 0.15351639688014984, 0.001678619533777237, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.015930648893117905, 0.006582066882401705, 0.10560829937458038, 0.3465193808078766, 0.012144939973950386, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.010950141586363316, 0.003185260808095336, 0.03380253165960312, 0.13516294956207275, 0.16374172270298004, 0.0833682045340538, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [4.016391176264733e-05, 0.0003202538937330246, 0.0050767818465828896, 1.7212016246048734e-05, 0.5176156759262085, 0.003749872324988246, 0.00026106167933903635, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13457109034061432, 0.07774609327316284, 0.006220821291208267, 0.0008077693055383861, 0.2509746253490448, 0.17662860453128815, 0.13796226680278778, 0.053514063358306885, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06553670763969421, 0.09473168104887009, 0.013516419567167759, 0.0013789478689432144, 0.03089364431798458, 0.0676402598619461, 0.03963227570056915, 0.17151857912540436, 0.1338733434677124, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07379595190286636, 0.1714182198047638, 0.13684017956256866, 0.00734432740136981, 0.0039545828476548195, 0.09408346563577652, 0.0452522449195385, 0.2525797188282013, 0.15314188599586487, 0.008748584426939487, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.006909683812409639, 0.034793343394994736, 0.13824458420276642, 0.0004423256032168865, 0.38493895530700684, 0.12702688574790955, 0.0007700703572481871, 0.005257567390799522, 0.3978818655014038, 0.028774550184607506, 0.016022928059101105, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15589091181755066, 0.059809040278196335, 0.2019805759191513, 0.006274765357375145, 0.053891621530056, 0.38889890909194946, 0.024021193385124207, 0.016828669235110283, 0.09206627309322357, 0.15270450711250305, 0.10960505902767181, 0.14381197094917297, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0011966965394094586, 0.0013769377255812287, 0.0006101150647737086, 4.0936538425739855e-05, 0.008213219232857227, 0.03395655378699303, 0.0003392287762835622, 0.00015790743054822087, 0.000944053172133863, 0.0007261222926899791, 0.011664116755127907, 0.22049497067928314, 0.0034024016931653023, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2470119595527649, 0.22662757337093353, 0.086290642619133, 0.0011605313047766685, 0.20862528681755066, 0.31339770555496216, 0.007298772688955069, 0.00864456407725811, 0.010568802244961262, 0.01924213580787182, 0.034804634749889374, 0.16789764165878296, 0.11296499520540237, 0.017940307036042213, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3800778388977051, 0.4679488241672516, 0.19362112879753113, 0.18464821577072144, 0.046723559498786926, 0.160307839512825, 0.24654103815555573, 0.2610638439655304, 0.07595612108707428, 0.1325986683368683, 0.022732526063919067, 0.1294456422328949, 0.2688123285770416, 0.12097980827093124, 0.12297553569078445, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.005153980106115341, 0.0002073257346637547, 0.12819816172122955, 0.00011319551413180307, 0.08506736904382706, 0.013190183788537979, 0.0028314462397247553, 0.00016588614380452782, 0.009067418053746223, 0.0008525841985829175, 0.00018506577180232853, 0.0002737078757490963, 0.0002474631182849407, 0.04919072240591049, 0.1850043386220932, 0.0018668848788365722, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4235798418521881, 0.8363600969314575, 0.13292381167411804, 0.03160996362566948, 0.6294970512390137, 0.3827916085720062, 0.01768689975142479, 0.031598031520843506, 0.05291707068681717, 0.004268768709152937, 0.01666090451180935, 0.0017059938982129097, 0.03961870074272156, 0.006749838124960661, 0.2787548303604126, 0.12898604571819305, 0.00984524842351675, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.001200420199893415, 0.004923743661493063, 0.03312471881508827, 7.996988279046491e-05, 0.2118730992078781, 0.0288531631231308, 0.00010192030458711088, 0.0002958755649160594, 0.007303019054234028, 0.00011155433458043262, 2.6572593014861923e-06, 0.00035481253871694207, 2.4723947262828005e-06, 2.6933960270980606e-06, 0.017764916643500328, 0.0003658832865767181, 0.25218549370765686, 0.002238432876765728, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16854390501976013, 0.046801913529634476, 0.18834064900875092, 0.005545254796743393, 0.10321269929409027, 0.3906272351741791, 0.03742265701293945, 0.024458711966872215, 0.05521516501903534, 0.07171308994293213, 0.021107476204633713, 0.025199010968208313, 0.0027974944096058607, 0.0025010560639202595, 0.02306896261870861, 0.15930885076522827, 0.06242140382528305, 0.11754277348518372, 0.21403564512729645, NaN, NaN, NaN, NaN, NaN, NaN], [0.0004002669302280992, 0.00040952101699076593, 0.00012874403910245746, 8.880775567376986e-06, 0.005201425869017839, 0.007163480389863253, 0.0002137795090675354, 0.00012960725871380419, 0.0005550362984649837, 0.0001244707527803257, 0.0006415210082195699, 0.03161805495619774, 4.1008814150700346e-05, 0.000599265971686691, 0.00399716105312109, 5.7038221711991355e-05, 0.0033261284697800875, 0.006950944196432829, 0.22392861545085907, 0.0028074102010577917, NaN, NaN, NaN, NaN, NaN], [0.22722585499286652, 0.18426381051540375, 0.07697561383247375, 0.0012757674558088183, 0.23254786431789398, 0.14769063889980316, 0.013780240900814533, 0.02735842764377594, 0.04001649469137192, 0.031179115176200867, 0.015889445319771767, 0.062248069792985916, 0.013498637825250626, 0.0052745710127055645, 0.2219674438238144, 0.0031969451811164618, 0.0037056237924844027, 0.028058722615242004, 0.22486938536167145, 0.09661445021629333, 0.02616964653134346, NaN, NaN, NaN, NaN], [0.27366653084754944, 0.354305237531662, 0.16368547081947327, 0.1598840057849884, 0.02900015190243721, 0.10581760108470917, 0.21902981400489807, 0.27043354511260986, 0.19813168048858643, 0.2514232099056244, 0.025616073980927467, 0.12471329420804977, 0.09682969748973846, 0.07310353219509125, 0.02883375994861126, 0.09285400807857513, 0.013515813276171684, 0.021914459764957428, 0.14159631729125977, 0.3238908648490906, 0.1783936321735382, 0.11570748686790466, NaN, NaN, NaN], [0.0030968550126999617, 7.297070260392502e-05, 0.1371629387140274, 0.00018204482330475003, 0.04798782989382744, 0.01213640347123146, 0.0023585439193993807, 0.00011540603009052575, 0.016970379278063774, 0.0015150568215176463, 0.0003718302759807557, 0.00044133648043498397, 0.00012143531785113737, 0.021671650931239128, 0.023021340370178223, 0.00010860650218091905, 0.0005334930610843003, 0.000257489358773455, 0.0005856966599822044, 0.00045311596477404237, 0.09709983319044113, 0.18528476357460022, 0.0029071324970573187, NaN, NaN], [0.49188995361328125, 0.918917715549469, 0.2054058462381363, 0.08403602242469788, 0.6967929005622864, 0.5653088688850403, 0.03772272169589996, 0.04957969859242439, 0.18319177627563477, 0.012161915190517902, 0.07060753554105759, 0.009896048344671726, 0.1126827672123909, 0.010653471574187279, 0.1938174068927765, 0.1352803260087967, 0.0021707522682845592, 0.030638370662927628, 0.003963022027164698, 0.03303877264261246, 0.004082953091710806, 0.20578816533088684, 0.11854958534240723, 0.02041587606072426, NaN], [0.001465475419536233, 0.00045102695003151894, 0.017218099907040596, 0.00030212500132620335, 0.11662620306015015, 0.017841650173068047, 0.00014393724268302321, 0.0003088460653088987, 0.006560556124895811, 0.0005491081974469125, 5.78465114813298e-05, 0.0019656207878142595, 0.00016285650781355798, 0.0002489366161171347, 0.011378495953977108, 0.0017521223053336143, 0.00787137821316719, 8.434856863459572e-05, 0.0012881350703537464, 7.287580228876323e-05, 0.00021561238099820912, 0.020317554473876953, 0.04195580258965492, 0.24219898879528046, 0.0017395684262737632]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.39058852195739746, 8.28505744721042e-06, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [2.7811127438326366e-05, 0.4158080220222473, 0.0005852450849488378, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [9.039229868085252e-13, 4.1926887206500396e-05, 0.15358270704746246, 0.00044542484101839364, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [1.9216391628896996e-16, 4.9363904963684035e-08, 0.0004218998074065894, 0.40449434518814087, 4.695959432865493e-06, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [1.7349648803667746e-14, 5.141012060505545e-09, 3.7822364902240224e-06, 0.0002717413299251348, 0.22465285658836365, 2.698016260183067e-06, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [3.6696812255598843e-09, 2.368522711293508e-09, 3.1902116006676806e-06, 9.520445587440918e-08, 9.990107355406508e-05, 0.2170185148715973, 0.019131841138005257, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [2.292660354896725e-07, 1.4062491449085002e-10, 1.0373556180720556e-11, 2.945570870549474e-11, 1.3987125901948616e-09, 1.1205498822164373e-06, 0.3382871150970459, 0.0008390913717448711, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [2.3133984541345853e-06, 0.00017511146143078804, 1.441240442545677e-06, 3.064446918443764e-09, 3.097617096159411e-08, 7.23518027712089e-08, 0.0017295092111453414, 0.39626115560531616, 0.00019915253506042063, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [8.689644937311981e-15, 2.8357308110571466e-06, 5.0946681540153804e-08, 2.0269605438549831e-10, 1.289949813632063e-10, 3.375676821404383e-11, 8.602300205495794e-09, 4.5097981455910485e-06, 0.29888245463371277, 6.641173968091607e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [2.8127108337250475e-18, 1.3557467148928026e-08, 7.431774662336466e-08, 2.301476165200711e-08, 1.1707952315975767e-11, 7.274678689300762e-12, 7.034611066401852e-13, 5.257664963120856e-13, 3.4044413041556254e-05, 0.32336506247520447, 4.600838292390108e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [6.300134025583048e-13, 5.676838910062543e-08, 1.822371018533886e-06, 2.3448223146260716e-05, 2.5415656068616954e-07, 3.417801153204891e-08, 5.353474885616549e-10, 2.141239963115993e-11, 3.762530198514469e-08, 6.24434178462252e-05, 0.33693620562553406, 3.183486114721745e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [1.5877897954763576e-12, 1.2288996487086479e-09, 3.458522428445576e-07, 9.462546586291865e-06, 7.457422907464206e-05, 0.0005706463125534356, 1.4425116212635203e-08, 4.5430816769144455e-13, 2.616490357709722e-12, 3.545688542772041e-08, 0.00016559385403525084, 0.22770871222019196, 0.0009294600458815694, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [2.579016999959549e-10, 1.5412886245069757e-10, 5.557828156033118e-11, 1.2367832313842086e-09, 3.3751638284229557e-07, 4.776334208145272e-07, 1.75399406998622e-07, 9.608910021829953e-12, 7.499024594652057e-14, 2.8573548556528813e-14, 3.2670008191793e-12, 4.494925178732956e-06, 0.37381958961486816, 3.638648195192218e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [3.090227983193472e-05, 8.430293382843956e-05, 4.32313208875712e-05, 1.6493000885020592e-06, 8.794136192591395e-06, 0.0005616153357550502, 0.0013158570509403944, 0.0005267951055429876, 3.675571861094795e-05, 2.42239195813454e-07, 8.356466074666002e-10, 2.3424906885338714e-06, 0.0012797197559848428, 0.6210904717445374, 0.0014036636566743255, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [7.67247776423119e-09, 2.954437938740284e-08, 8.54147774731473e-09, 2.011255162415182e-09, 5.265776792384713e-08, 1.4630668898618637e-09, 2.2913241082278546e-06, 3.266295323101076e-08, 1.6124132571349037e-06, 1.13081211061683e-11, 2.6358108895513247e-15, 7.728456763445024e-11, 2.3767283696685126e-09, 2.1271845980663784e-05, 0.19462287425994873, 6.456446044467157e-06, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [4.312543703220706e-13, 2.1705271535665815e-07, 1.1365986551936658e-07, 1.9739390211270802e-07, 7.690645453806155e-09, 4.219609994748907e-09, 9.716764060030414e-10, 3.915795687703394e-08, 3.0873563900968293e-06, 5.5168204227129536e-08, 1.0056843552375128e-10, 6.254387632798064e-12, 4.318517331930449e-12, 1.5618051990573534e-11, 6.033264071447775e-05, 0.4116440713405609, 1.8908482161350548e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [1.797858697974407e-17, 3.5553746058347713e-10, 1.0377114723070235e-09, 5.157609006545272e-09, 5.5740526777592336e-11, 3.675403037473046e-11, 3.015720268992328e-12, 1.2632186895361434e-14, 3.2584634990229233e-09, 2.7093712162695738e-08, 2.733851353305984e-15, 2.0347772078377346e-10, 7.802066534575867e-16, 1.702402683943053e-16, 1.8298086656987067e-10, 6.30185184036236e-08, 0.2592085301876068, 3.469779585429933e-06, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [3.386366187463352e-10, 1.5587464474720036e-07, 5.430682108453766e-07, 1.926859113154933e-05, 2.7584928830037825e-06, 5.553058031182445e-07, 6.554741815989473e-08, 7.146391256540596e-10, 4.225638150501254e-08, 2.0539353045023745e-06, 0.00010312868107575923, 2.5505174860995794e-08, 1.3659710695890226e-08, 4.206753695390475e-11, 5.200286035123014e-11, 3.842067428649898e-07, 1.4282905794971157e-05, 0.31164512038230896, 0.00011869923037011176, NaN, NaN, NaN, NaN, NaN, NaN], [3.098006018387167e-10, 3.2388165482899467e-09, 1.8609943808201024e-08, 5.099297482047405e-07, 4.603737033903599e-05, 0.00016448901442345232, 1.6998721719119203e-07, 1.7718410072475876e-11, 2.5886336477154437e-11, 9.218055652127077e-09, 1.2046231745443947e-07, 7.304957398446277e-05, 2.3164133111652774e-10, 2.8952129582648922e-09, 2.9085676575557606e-11, 8.895827650901023e-12, 8.14965606110718e-09, 8.762691868469119e-05, 0.2280847281217575, 0.0004104141262359917, NaN, NaN, NaN, NaN, NaN], [1.3149543676149733e-09, 1.080373679407387e-09, 5.5150013028582023e-11, 7.800748935693491e-10, 1.7859061074432248e-07, 2.183157299384675e-08, 2.5236221290469985e-07, 2.35878039323012e-10, 9.060349692724401e-12, 1.4339956088890715e-12, 1.7799637631876752e-12, 2.9941787715870305e-08, 6.0217857935640495e-06, 3.1683756313016787e-11, 4.5713120788715145e-11, 3.4124135808721867e-13, 3.591858459424911e-15, 1.3559961530365539e-12, 3.119595021416899e-06, 0.35679423809051514, 3.964137067669071e-05, NaN, NaN, NaN, NaN], [4.326914222474443e-06, 0.00023807807883713394, 0.00026310785324312747, 8.714396244613454e-06, 1.617559973965399e-05, 0.0001319001312367618, 0.0005945482989773154, 0.000823884445708245, 0.0008506007143296301, 1.7805428797146305e-05, 2.734714854568665e-08, 2.8855724849563558e-06, 4.891938442597166e-05, 0.0011682395124807954, 8.529372053089901e-07, 0.00017029111040756106, 1.0359013202787537e-07, 7.06834313302096e-10, 1.0861956525332062e-06, 0.0008713650749996305, 0.596385657787323, 0.0009257638594135642, NaN, NaN, NaN], [1.4773272882795396e-10, 2.3448599506536993e-08, 6.434380566133768e-07, 3.8027360460546333e-07, 2.454226432746509e-06, 5.541529457531169e-09, 3.5226184991188347e-06, 2.5443886997322807e-08, 1.7749154721968807e-05, 1.8393259137994278e-09, 4.026108439691978e-12, 6.382850692432385e-09, 1.7809153263215194e-08, 8.996512974590587e-07, 0.00010512088192626834, 1.1464897607671443e-11, 2.794342757184154e-09, 2.4549680847631107e-15, 9.933188299671158e-11, 7.3009864820505754e-09, 8.105817687464878e-05, 0.2077004611492157, 2.0097606466151774e-05, NaN, NaN], [1.1257004341538607e-14, 1.3137036347643516e-08, 4.6611327775281097e-07, 3.0405328743654536e-06, 1.5423474053477548e-07, 2.520166120234535e-08, 3.4643394819511286e-09, 1.1558090484697914e-08, 1.417677253812144e-06, 9.112129362165433e-08, 4.2694305868451465e-09, 3.7723260626343347e-10, 4.1450526344632976e-10, 2.7357388923676673e-11, 6.112880441833113e-07, 3.9687514799879864e-05, 8.382351063263016e-11, 8.293656039715103e-11, 4.97465783844131e-12, 4.144883221368634e-12, 1.4191136113450575e-11, 2.5566061594872735e-05, 0.4056495428085327, 4.4409513066057116e-05, NaN], [9.215334861117716e-19, 2.6557794852166694e-10, 5.799645919069008e-07, 1.003176621633406e-11, 7.217926736302616e-07, 4.876178394397357e-08, 8.254863459455919e-11, 1.424103456687531e-12, 1.1857503423584603e-08, 1.3074058502482444e-09, 8.580362115262474e-12, 5.829819293978744e-09, 1.8017319407259702e-12, 9.234832950427707e-14, 3.576115098491428e-11, 1.9265784523270213e-09, 1.8997316146851517e-06, 1.949248054633479e-11, 8.860704392432694e-10, 2.8198800851872777e-14, 5.674391451236226e-15, 1.0258181110112119e-10, 6.93914080329705e-06, 0.25534507632255554, 2.742740150551981e-07]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0002614231198094785, 0.183704674243927, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [1.3331101555991154e-08, 0.003119559260085225, 0.19454506039619446, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [1.1244888353800775e-09, 0.0005117341643199325, 0.15345418453216553, 0.0018621939234435558, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [2.882708471929618e-08, 0.0006895777769386768, 0.008299488574266434, 0.004234161227941513, 0.26378652453422546, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [6.507164653157815e-05, 0.0030905166640877724, 0.269605815410614, 0.06594818085432053, 0.07055308669805527, 0.24370616674423218, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [5.806248736917041e-05, 0.0008924558642320335, 0.00047033390728756785, 0.003593915607780218, 0.044251326471567154, 0.18547922372817993, 0.19724349677562714, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03321969881653786, 0.1786998063325882, 0.0021111152600497007, 0.00015362887643277645, 0.0013223892310634255, 0.01674751006066799, 0.27181917428970337, 0.0704144611954689, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0005316429305821657, 0.0021434861700981855, 0.0005638045258820057, 2.0347550162114203e-05, 8.372889715246856e-05, 0.0012170294066891074, 0.0006328476592898369, 0.0015302025713026524, 0.2731996476650238, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [3.384976253073546e-06, 0.0032942681573331356, 0.003179847961291671, 0.0003072107210755348, 3.0923787562642246e-05, 0.0003082206822000444, 0.0026841319631785154, 0.011449099518358707, 0.2928124964237213, 0.0015787724405527115, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [4.910896677756682e-05, 0.01189705915749073, 0.0036808690056204796, 0.006090851966291666, 0.0029882052913308144, 0.006760776974260807, 0.0002592294185888022, 0.0001972121826838702, 0.15788163244724274, 0.14973512291908264, 0.14614373445510864, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [7.539001671830192e-05, 0.036947283893823624, 0.01112621370702982, 0.04119950905442238, 0.06979847699403763, 0.01383589580655098, 0.008948443457484245, 9.020609286380932e-05, 0.0005221512983553112, 0.34183818101882935, 0.12104173004627228, 0.027292484417557716, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [5.4811065638205037e-05, 0.015359039418399334, 0.005874635651707649, 0.024854328483343124, 0.16572602093219757, 0.13195344805717468, 0.08553953468799591, 0.00124072446487844, 0.0008515206864103675, 0.0025517549365758896, 0.03817262500524521, 0.1957935392856598, 0.020919298753142357, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [3.401398498681374e-05, 0.0008079431718215346, 0.00045223115012049675, 0.00013304724416229874, 0.0006849576020613313, 0.009534466080367565, 0.010466179810464382, 0.00030334663460962474, 0.00033610902028158307, 2.1021634893259034e-05, 6.891421071486548e-05, 0.0028196852654218674, 0.3685440421104431, 0.0008976467652246356, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0012722803512588143, 0.07485485821962357, 0.004568059463053942, 0.008557068184018135, 0.04491077736020088, 0.010689688846468925, 0.010801602154970169, 0.015439217910170555, 0.001288879313506186, 0.032191790640354156, 9.430324280401692e-05, 0.0010071481810882688, 0.03593403846025467, 0.015365669503808022, 0.28865233063697815, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0003195737663190812, 0.0016381103778257966, 0.001899963477626443, 0.000450764549896121, 0.0029568641912192106, 0.0004077073244843632, 0.006739944685250521, 5.316005626809783e-05, 0.000977654941380024, 0.00033480822457931936, 1.5544836060144007e-05, 5.177688763069455e-06, 0.000280524865956977, 8.569184137741104e-05, 0.19435854256153107, 0.0009946423815563321, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0004552309401333332, 0.00916277151554823, 0.2859989106655121, 0.028668222948908806, 0.004703177139163017, 0.013283651322126389, 0.011935138143599033, 0.00041849465924315155, 0.021506765857338905, 0.0005354905733838677, 2.3408898414345458e-05, 5.557515123655321e-06, 4.006853941973532e-06, 0.000782388960942626, 0.032734211534261703, 0.33600685000419617, 0.05645810067653656, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.001615832676179707, 0.0592908076941967, 0.004439341835677624, 0.0221478920429945, 0.05761101841926575, 0.08599329739809036, 0.009327156469225883, 0.0014337823959067464, 0.22479815781116486, 0.007599419914186001, 0.00010282513540005311, 0.003995772451162338, 0.0007532926392741501, 0.0001985877170227468, 0.042725738137960434, 0.609107255935669, 0.032340146601200104, 0.2600889503955841, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0007894318550825119, 0.08912800997495651, 0.00870462041348219, 0.062210533767938614, 0.21669252216815948, 0.04955689236521721, 0.12036743760108948, 0.001276280265301466, 0.002290783217176795, 0.4637441337108612, 0.041003014892339706, 0.007595454342663288, 0.0049859327264130116, 0.030789200216531754, 0.01441932376474142, 0.02666427381336689, 0.013092019595205784, 0.22824719548225403, 0.07290598005056381, NaN, NaN, NaN, NaN, NaN, NaN], [4.2991967347916216e-05, 0.006631283089518547, 0.0006027332856319845, 0.004053125157952309, 0.03894652798771858, 0.031787656247615814, 0.10168109834194183, 0.004267984535545111, 0.002045443281531334, 0.0010633694473654032, 0.005091637372970581, 0.031351421028375626, 6.663963722530752e-05, 0.09428737312555313, 0.0008465268765576184, 0.00024849644978530705, 0.002269570017233491, 0.01905866153538227, 0.2164839655160904, 0.010082208551466465, NaN, NaN, NaN, NaN, NaN], [1.1191940757271368e-05, 0.0006002296577207744, 0.0002709901600610465, 9.913583926390857e-05, 0.0001758227008394897, 0.0029332106932997704, 0.008675863035023212, 0.0011328428518027067, 0.0023299665190279484, 6.693489558529109e-05, 0.00013525204849429429, 0.0013442488852888346, 0.022858861833810806, 2.321010106243193e-05, 0.0010626229923218489, 2.5993340386776254e-05, 3.972689592046663e-05, 5.326797690941021e-05, 0.0033412689808756113, 0.35271701216697693, 0.0008956229430623353, NaN, NaN, NaN, NaN], [0.00036489564809016883, 0.07616367936134338, 0.00673737283796072, 0.011110173538327217, 0.021392904222011566, 0.010494116693735123, 0.006134945899248123, 0.015969248488545418, 0.005187375005334616, 0.12039955705404282, 0.0005341891082935035, 0.0022901638876646757, 0.027128320187330246, 0.005907480139285326, 0.033119603991508484, 0.002176248235628009, 0.0003625153622124344, 6.369769835146144e-05, 0.0007003483478911221, 0.03456505015492439, 0.01570759527385235, 0.28412890434265137, NaN, NaN, NaN], [3.192616713931784e-05, 0.00035208670306019485, 0.002478531561791897, 0.0006564928335137665, 0.0008886585710570216, 0.0005662215990014374, 0.0016915983287617564, 1.3900444173486903e-05, 0.0009738726075738668, 0.00042995362309738994, 8.639829320600256e-05, 1.4000924238644075e-05, 0.00033226466621272266, 2.9785558581352234e-05, 0.00921203475445509, 3.390025085536763e-06, 5.1574592362158e-05, 2.3835823412809987e-06, 1.9022172637050971e-06, 0.00016878120368346572, 9.063100151252002e-05, 0.20696188509464264, 0.001649125711992383, NaN, NaN], [0.00019471753330435604, 0.003537738462910056, 0.2800489366054535, 0.036592625081539154, 0.002127013634890318, 0.024595409631729126, 0.008275463245809078, 0.00023266732750926167, 0.021680369973182678, 0.0005173377576284111, 7.175304199336097e-05, 2.6857771445065737e-05, 1.6371919627999887e-05, 0.0012281013187021017, 0.011112956330180168, 0.058813560754060745, 0.0009629606502130628, 1.1531898962857667e-05, 4.947432444168953e-06, 2.475359451636905e-06, 0.0005685617215931416, 0.0267820842564106, 0.3296748399734497, 0.06147307902574539, NaN], [3.20236104300875e-08, 0.00013383101031649858, 0.00029007354169152677, 0.002788462908938527, 0.0014709108509123325, 0.0009710633894428611, 0.0001290659129153937, 2.0881772798020393e-05, 7.236683813971467e-06, 3.12792144541163e-05, 7.099155482137576e-05, 3.213396485080011e-05, 3.9666349039180204e-05, 0.00022854047711007297, 0.0037343965377658606, 1.487573445047019e-05, 0.00019343644089531153, 8.10168421594426e-05, 1.1448363693489227e-05, 3.5921341350331204e-06, 2.216967368440237e-05, 0.0017730530817061663, 0.0001526248233858496, 0.009769736789166927, 0.4419056475162506]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07662782073020935, 0.14776498079299927, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0006832284270785749, 0.003495789598673582, 0.19430121779441833, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00020953372586518526, 0.007476589176803827, 0.1521030217409134, 0.003494996577501297, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00048688906827010214, 0.0011088894680142403, 0.0024602855555713177, 0.0005520267877727747, 0.26744863390922546, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0004194685607217252, 0.0005068383179605007, 0.026896899566054344, 0.0004147894505877048, 0.006156287621706724, 0.4387049376964569, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [1.0518371709622443e-05, 5.5142045312095433e-05, 0.016997506842017174, 3.693701364682056e-05, 0.0006244040559977293, 0.21657241880893707, 0.01345360092818737, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3619365394115448, 0.25655418634414673, 0.3611752688884735, 0.14710570871829987, 0.018539972603321075, 0.21814967691898346, 0.09323819726705551, 0.01780291646718979, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.004012200981378555, 0.004658036399632692, 0.017421945929527283, 0.0026806569658219814, 0.590861439704895, 0.051964171230793, 0.007618917152285576, 0.0007336572161875665, 0.12340892106294632, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.44725751876831055, 0.6053639054298401, 0.07041247189044952, 0.07085516303777695, 0.003138674655929208, 0.2879992425441742, 0.049135204404592514, 0.14297868311405182, 0.06008363142609596, 0.06304289400577545, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.7072809338569641, 0.7582566142082214, 0.16150887310504913, 0.18586905300617218, 0.015776842832565308, 0.08385244756937027, 0.32581770420074463, 0.5540359020233154, 0.13379113376140594, 0.0028463751077651978, 0.051922835409641266, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4378974437713623, 0.10523661971092224, 0.014314417727291584, 0.30093127489089966, 0.06324318051338196, 0.08432605862617493, 0.2594241797924042, 0.6188808083534241, 0.3929617404937744, 0.00827555637806654, 0.07725780457258224, 0.06407154351472855, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2013174593448639, 0.5200937390327454, 0.3190821707248688, 0.5249915719032288, 0.18779213726520538, 0.1779765784740448, 0.29882070422172546, 0.5049118399620056, 0.06443758308887482, 0.007539320737123489, 0.16998757421970367, 0.031686559319496155, 0.3610091209411621, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5546301603317261, 0.5397829413414001, 0.43089261651039124, 0.08987504988908768, 0.3114354610443115, 0.4812281131744385, 0.11215226352214813, 0.17198431491851807, 0.5790820121765137, 0.03648975491523743, 0.0541677288711071, 0.04165489599108696, 0.07749651372432709, 0.030232839286327362, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.005376005079597235, 0.010858614929020405, 0.02991071715950966, 0.029742157086730003, 0.04020260274410248, 0.1695990264415741, 0.0604972317814827, 0.10318762809038162, 0.48727869987487793, 0.07163358479738235, 0.025501595810055733, 0.05125340074300766, 0.22269804775714874, 0.08394679427146912, 0.19870582222938538, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0006954512791708112, 0.0002132337394868955, 0.037006676197052, 0.0018452922813594341, 0.16118928790092468, 0.5505160689353943, 0.028353480622172356, 0.0021746368147432804, 0.027092093601822853, 0.0001434519508620724, 0.0029707583598792553, 4.2726576793938875e-05, 0.0012847317848354578, 0.0010433235438540578, 0.18891005218029022, 0.014656933024525642, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.013874622993171215, 0.0695175901055336, 0.005752294324338436, 0.005697373300790787, 0.0021822804119437933, 0.02415846660733223, 0.00723307253792882, 0.3120453357696533, 0.016472192481160164, 0.004319194238632917, 0.041901107877492905, 0.7052133083343506, 0.0035930864978581667, 0.020578961819410324, 0.0021869041956961155, 0.0003597450559027493, 0.0005889505264349282, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.29724666476249695, 0.30918487906455994, 0.0693497508764267, 0.04026606306433678, 0.00593132060021162, 0.04497085511684418, 0.07199602574110031, 0.16270284354686737, 0.058071933686733246, 0.0005904879071749747, 0.0013724194141104817, 0.013050474226474762, 0.002609569113701582, 0.013482913374900818, 0.089314766228199, 0.03341012820601463, 0.21929660439491272, 0.006776490714401007, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3422777056694031, 0.07256462424993515, 0.012822822667658329, 0.21187257766723633, 0.060081083327531815, 0.09390594810247421, 0.19744858145713806, 0.5327264666557312, 0.3024030029773712, 0.013231869786977768, 0.1601967215538025, 0.04191795364022255, 0.5788960456848145, 0.791706383228302, 0.2698511779308319, 0.26516515016555786, 0.2890409529209137, 0.032140959054231644, 0.02436642162501812, NaN, NaN, NaN, NaN, NaN, NaN], [0.15722303092479706, 0.44676893949508667, 0.24300073087215424, 0.3980245292186737, 0.29666030406951904, 0.21130049228668213, 0.31708449125289917, 0.45276522636413574, 0.04954151436686516, 0.006070373114198446, 0.23888874053955078, 0.06321726739406586, 0.48237892985343933, 0.09136107563972473, 0.571183979511261, 0.36026179790496826, 0.0799446776509285, 0.1583012342453003, 0.025381257757544518, 0.5154083371162415, NaN, NaN, NaN, NaN, NaN], [0.6566299200057983, 0.6752134561538696, 0.5489535927772522, 0.1520741730928421, 0.6433172821998596, 0.7151104211807251, 0.290630042552948, 0.3418242335319519, 0.686417818069458, 0.046654678881168365, 0.09611856192350388, 0.0634889155626297, 0.4891318380832672, 0.46607306599617004, 0.5581225156784058, 0.4337400496006012, 0.06152508407831192, 0.08386452496051788, 0.0397774837911129, 0.11068917065858841, 0.04009125009179115, NaN, NaN, NaN, NaN], [0.0024060788564383984, 0.006098441779613495, 0.013975032605230808, 0.014695755206048489, 0.022452646866440773, 0.10514718294143677, 0.04751533642411232, 0.0609392412006855, 0.31799331307411194, 0.04427095875144005, 0.01951766200363636, 0.04202713817358017, 0.3371936082839966, 0.2731744647026062, 0.3478449583053589, 0.03363266587257385, 0.011759405955672264, 0.01767517626285553, 0.024101490154862404, 0.19511322677135468, 0.05518092215061188, 0.2097322940826416, NaN, NaN, NaN], [0.000109505133877974, 2.9198725314927287e-05, 0.01053665205836296, 0.0007290886132977903, 0.055462777614593506, 0.18011406064033508, 0.013305839151144028, 0.0007181179826147854, 0.008689867332577705, 4.760328374686651e-05, 0.0016827695071697235, 2.2867327061248943e-05, 0.000821226101834327, 0.0012459746794775128, 0.2353316843509674, 0.004575389437377453, 0.003901307238265872, 0.0009429306373931468, 1.1980442650383338e-05, 0.0003497266152407974, 0.00027309934375807643, 0.1965111494064331, 0.005757085047662258, NaN, NaN], [0.0017744784709066153, 0.012578981928527355, 0.0015974465059116483, 0.002320722443982959, 0.0008557687979191542, 0.004459704738110304, 0.00322481500916183, 0.13683773577213287, 0.010506929829716682, 0.0027294831816107035, 0.03936534747481346, 0.7146239876747131, 0.0021277000196278095, 0.014929071068763733, 0.003117389976978302, 0.0010002683848142624, 0.0005979579291306436, 0.037009548395872116, 0.6984097361564636, 0.0021584301721304655, 0.012162267230451107, 0.002483450109139085, 0.00014705986541230232, 0.0003713203768711537, NaN], [0.10933294892311096, 0.0594157911837101, 0.01442565955221653, 0.027944112196564674, 0.24928514659404755, 0.3314722180366516, 0.036283038556575775, 0.01824975199997425, 0.03247179090976715, 0.02741291932761669, 0.0011664694175124168, 0.03365480154752731, 0.10097742080688477, 0.021067792549729347, 0.42791858315467834, 0.11242418736219406, 0.11434369534254074, 0.000791618600487709, 0.02291581965982914, 0.07201644033193588, 0.02081850729882717, 0.39859694242477417, 0.2763477563858032, 0.13874487578868866, 0.003258609212934971]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.026641450822353363, 0.17128966748714447, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5577486157417297, 0.24638143181800842, 0.025497647002339363, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1241803988814354, 0.06599891930818558, 0.13004763424396515, 0.33318501710891724, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.9552784562110901, 0.6656578779220581, 0.04364815354347229, 0.097982257604599, 0.0012550450628623366, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.6779462695121765, 0.5809971690177917, 0.2087380737066269, 0.15752893686294556, 0.08772724121809006, 0.09023962169885635, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.6994673609733582, 0.48720496892929077, 0.08263873308897018, 0.3298986256122589, 0.0049313209019601345, 0.07016509026288986, 0.5443912744522095, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3437848389148712, 0.28689879179000854, 0.5712999105453491, 0.5371078252792358, 0.06584293395280838, 0.2492358684539795, 0.014812931418418884, 0.02226697839796543, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.44942334294319153, 0.3777551054954529, 0.7612449526786804, 0.7021526098251343, 0.30080679059028625, 0.4424319267272949, 0.22922295331954956, 0.04627525433897972, 0.055941756814718246, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.47138965129852295, 0.18856076896190643, 0.6503154039382935, 0.9041082859039307, 0.2803841233253479, 0.4006999135017395, 0.5757170915603638, 0.295682817697525, 0.04142303764820099, 0.006079117301851511, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24097655713558197, 0.15950126945972443, 0.6649572849273682, 0.6751598119735718, 0.46790093183517456, 0.6438081860542297, 0.3765251934528351, 0.2975021302700043, 0.10267924517393112, 0.060453154146671295, 0.03869982063770294, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.39086097478866577, 0.6666929125785828, 0.5642580389976501, 0.557075023651123, 0.25761184096336365, 0.3620971143245697, 0.656988263130188, 0.301082581281662, 0.3758563995361328, 0.026163028553128242, 0.024990877136588097, 0.0074356794357299805, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.7909376621246338, 0.3817039430141449, 0.6133569478988647, 0.41290101408958435, 0.30558884143829346, 0.6049348711967468, 0.5688384175300598, 0.4680134057998657, 0.6550416946411133, 0.42371857166290283, 0.10508850961923599, 0.021316751837730408, 0.05294431000947952, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17973686754703522, 0.17233335971832275, 0.334688276052475, 0.4481850564479828, 0.04172942414879799, 0.10337609797716141, 0.5107487440109253, 0.7207926511764526, 0.1405051052570343, 0.0654703825712204, 0.41273486614227295, 0.17914383113384247, 0.042542651295661926, 0.010745447129011154, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5207539200782776, 0.308788537979126, 0.08189663290977478, 0.5850351452827454, 0.3457651734352112, 0.15844188630580902, 0.2948668897151947, 0.4065589904785156, 0.12084604799747467, 0.29343682527542114, 0.49164822697639465, 0.07233413308858871, 0.0535273477435112, 0.014947501011192799, 0.008541097864508629, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2949400544166565, 0.03748409450054169, 0.14473117887973785, 0.0705113336443901, 0.013025683350861073, 0.005298166535794735, 0.21091029047966003, 0.014800299890339375, 0.2805088758468628, 0.000897476973477751, 0.0938984826207161, 0.004705057479441166, 0.04936474934220314, 0.011992034502327442, 0.18721424043178558, 0.00230285432189703, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.44276589155197144, 0.06478449702262878, 0.543609619140625, 0.8444110155105591, 0.13468694686889648, 0.4405028522014618, 0.6528593897819519, 0.5737791061401367, 0.6313535571098328, 0.8501816987991333, 0.4486657381057739, 0.06076665595173836, 0.7409859299659729, 0.15147589147090912, 0.20801351964473724, 0.027446726337075233, 0.036936238408088684, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5445577502250671, 0.2876933515071869, 0.7013069987297058, 0.627236008644104, 0.37061285972595215, 0.6206991076469421, 0.38252583146095276, 0.4230470061302185, 0.31842562556266785, 0.28603002429008484, 0.015331648290157318, 0.14692452549934387, 0.8622261881828308, 0.049388445913791656, 0.37183380126953125, 0.17907747626304626, 0.05781394988298416, 0.020684318616986275, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.4656296670436859, 0.6725881099700928, 0.6199259161949158, 0.6479836702346802, 0.24076998233795166, 0.34658652544021606, 0.5947279930114746, 0.37259459495544434, 0.5521662831306458, 0.14718003571033478, 0.19626900553703308, 0.024240192025899887, 0.27736979722976685, 0.05565635487437248, 0.3618892729282379, 0.44332295656204224, 0.027751203626394272, 0.0260067880153656, 0.010717106983065605, NaN, NaN, NaN, NaN, NaN, NaN], [0.830940842628479, 0.42077580094337463, 0.7156820893287659, 0.57599937915802, 0.5493759512901306, 0.7128159999847412, 0.5476810932159424, 0.527928352355957, 0.8053308725357056, 0.8646240234375, 0.542984127998352, 0.2950981855392456, 0.3170693516731262, 0.5610483884811401, 0.26465174555778503, 0.45835256576538086, 0.22733505070209503, 0.10187508910894394, 0.03538959100842476, 0.07069608569145203, NaN, NaN, NaN, NaN, NaN], [0.09599269181489944, 0.08247342705726624, 0.25253206491470337, 0.4357891380786896, 0.039192523807287216, 0.0719948410987854, 0.3563676178455353, 0.5300538539886475, 0.06311739236116409, 0.037909455597400665, 0.5032193064689636, 0.39894816279411316, 0.3283153772354126, 0.21619060635566711, 0.017918655648827553, 0.2577371895313263, 0.14531975984573364, 0.346793532371521, 0.2014700472354889, 0.0539211668074131, 0.0146569162607193, NaN, NaN, NaN, NaN], [0.6422337889671326, 0.3740711212158203, 0.10689651221036911, 0.6858291029930115, 0.4494076073169708, 0.2826421856880188, 0.3886936604976654, 0.475405216217041, 0.13226336240768433, 0.3073323965072632, 0.7139697670936584, 0.17356495559215546, 0.25040003657341003, 0.23144030570983887, 0.024455448612570763, 0.4280460476875305, 0.048713963478803635, 0.3974619209766388, 0.06130422651767731, 0.05969162657856941, 0.015271119773387909, 0.00685582309961319, NaN, NaN, NaN], [0.5218734741210938, 0.03395698964595795, 0.2861349880695343, 0.13773199915885925, 0.02211177349090576, 0.014614011161029339, 0.43378758430480957, 0.02492188662290573, 0.26067787408828735, 0.0009113854030147195, 0.1411941796541214, 0.009023642167448997, 0.14982649683952332, 0.15959703922271729, 0.7153633832931519, 0.014257365837693214, 0.06102409213781357, 0.12158294767141342, 0.006897313520312309, 0.06130388379096985, 0.012951835058629513, 0.16874605417251587, 0.002189028775319457, NaN, NaN], [0.45293620228767395, 0.05202305316925049, 0.4803192913532257, 0.8224762082099915, 0.10338833183050156, 0.2861584722995758, 0.8321961760520935, 0.7622299790382385, 0.5323314070701599, 0.8633370995521545, 0.5219312310218811, 0.07432084530591965, 0.7646023631095886, 0.4150907099246979, 0.4998815357685089, 0.606073796749115, 0.2854492664337158, 0.6639280319213867, 0.09482558071613312, 0.806840717792511, 0.19665148854255676, 0.18194931745529175, 0.01953776553273201, 0.037144362926483154, NaN], [0.8357685804367065, 0.6023411154747009, 0.16389556229114532, 0.4697819948196411, 0.05014880374073982, 0.3185025751590729, 0.2618474066257477, 0.7044641375541687, 0.16675803065299988, 0.7323283553123474, 0.14429442584514618, 0.2621355652809143, 0.041847843676805496, 0.3185603618621826, 0.04513467848300934, 0.49906620383262634, 0.611339807510376, 0.21515053510665894, 0.3302164673805237, 0.04920952767133713, 0.2760073244571686, 0.0218669306486845, 0.25043201446533203, 0.13627314567565918, 0.01334126852452755]]], [[[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13569742441177368, 0.0376364141702652, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05053132027387619, 0.5417848825454712, 0.07814626395702362, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03762863576412201, 0.4749486744403839, 0.013701170682907104, 0.053301598876714706, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10598134994506836, 0.16776065528392792, 0.11929589509963989, 0.16846179962158203, 0.40715572237968445, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05147748813033104, 0.203742116689682, 0.11462464928627014, 0.46246808767318726, 0.01836300455033779, 0.02458924613893032, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17594558000564575, 0.17753779888153076, 0.024665912613272667, 0.19817322492599487, 0.008797828108072281, 0.022263213992118835, 0.29173722863197327, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.016114797443151474, 0.0061007170006632805, 0.028504224494099617, 0.017245782539248466, 0.08753485232591629, 0.11264273524284363, 0.6154332160949707, 0.029144972562789917, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.027042992413043976, 0.032212790101766586, 0.019619816914200783, 0.014702342450618744, 0.06721275299787521, 0.2560867667198181, 0.5545244216918945, 0.40561506152153015, 0.037922732532024384, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1654873937368393, 0.013622531667351723, 0.0656571239233017, 0.09179358184337616, 0.03440919890999794, 0.08533406257629395, 0.16269220411777496, 0.1151970624923706, 0.09265416115522385, 0.028269361704587936, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2598540484905243, 0.010173649527132511, 0.004170349799096584, 0.003479698905721307, 0.0014636714477092028, 0.0011101020500063896, 0.001677120802924037, 0.034040722995996475, 0.0041177538223564625, 0.024958845227956772, 0.016315795481204987, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17492477595806122, 0.010013026185333729, 0.005800239276140928, 0.0069971769116818905, 0.0036480696871876717, 0.001016399241052568, 0.0060493675991892815, 0.0034581662621349096, 0.00659980857744813, 0.0047594537027180195, 0.3941299021244049, 0.2407994568347931, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06559828668832779, 0.005602334160357714, 0.0005807551206089556, 0.0005322807701304555, 0.004617360420525074, 0.00354054500348866, 0.005599506665021181, 0.011434626765549183, 0.006905066315084696, 0.009602343663573265, 0.11027393490076065, 0.36931946873664856, 0.06368503719568253, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.015983520075678825, 0.012168757617473602, 0.0015684146201238036, 0.0005484889261424541, 0.00233695306815207, 0.0038106110878288746, 0.005947766825556755, 0.04194773733615875, 0.014443459920585155, 0.06465759128332138, 0.14989611506462097, 0.5095774531364441, 0.1882752925157547, 0.02387852594256401, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11159919947385788, 0.06036144495010376, 0.06681493669748306, 0.0798669382929802, 0.03668922558426857, 0.018710536882281303, 0.029976846650242805, 0.0675768032670021, 0.03372039645910263, 0.057603828608989716, 0.14515243470668793, 0.25060775876045227, 0.23181115090847015, 0.14262832701206207, 0.33286023139953613, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.018035059794783592, 0.02341379225254059, 0.0019442361081019044, 0.004369894042611122, 0.00136191223282367, 0.00017434914479963481, 0.0011034610215574503, 0.06787250190973282, 0.060198791325092316, 0.12004764378070831, 0.11878902465105057, 0.2063554972410202, 0.28332868218421936, 0.35319504141807556, 0.008158767595887184, 0.26057863235473633, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17278411984443665, 0.007028562016785145, 0.010641193017363548, 0.013809186406433582, 0.0005732428980991244, 0.001056239241734147, 0.0005258666351437569, 0.03639528155326843, 0.02256075292825699, 0.01660884916782379, 0.1527748554944992, 0.1477358043193817, 0.2577149271965027, 0.03867224231362343, 0.04304511100053787, 0.11759469658136368, 0.0762997567653656, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.38573285937309265, 0.0028330886270850897, 0.0014278099406510592, 0.0009824484586715698, 9.371336636831984e-05, 0.00015483389142900705, 6.760591350030154e-05, 0.0035791138652712107, 0.0002520910056773573, 0.0005180046427994967, 0.00024238335026893765, 0.011901103891432285, 0.011019378900527954, 0.006276060827076435, 0.0026990415062755346, 0.016820058226585388, 0.03330027312040329, 0.047877803444862366, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21399648487567902, 0.008264300413429737, 0.0051351506263017654, 0.005111425183713436, 0.0020249083172529936, 0.00047485672985203564, 0.0018332998733967543, 0.0008904117858037353, 0.0017731828847900033, 0.000539442349690944, 0.03944296017289162, 0.039767228066921234, 0.00580678740516305, 0.004312179517000914, 0.003937484696507454, 0.00913114845752716, 0.006211036816239357, 0.3553882837295532, 0.3024981617927551, NaN, NaN, NaN, NaN, NaN, NaN], [0.05261809378862381, 0.004144520964473486, 0.00047606538282707334, 0.0003396419051568955, 0.002880769083276391, 0.0015178520698100328, 0.0018901955336332321, 0.0029504895210266113, 0.0017174717504531145, 0.0006908842478878796, 0.0046035549603402615, 0.09042679518461227, 0.0032755613792687654, 0.007712012622505426, 0.032594844698905945, 0.02268057130277157, 0.033856723457574844, 0.07955116033554077, 0.4074561595916748, 0.07153668999671936, NaN, NaN, NaN, NaN, NaN], [0.019381573423743248, 0.012705344706773758, 0.0019882190972566605, 0.0005741973291151226, 0.0020475401543080807, 0.0023934554774314165, 0.004172713495790958, 0.021013854071497917, 0.005879250820726156, 0.006729640066623688, 0.00632414361461997, 0.09735815972089767, 0.01909361220896244, 0.00100265524815768, 0.003452989971265197, 0.008203250356018543, 0.05971603840589523, 0.11904174834489822, 0.5188009142875671, 0.2541559338569641, 0.029506316408514977, NaN, NaN, NaN, NaN], [0.10572486370801926, 0.04525948688387871, 0.055838145315647125, 0.050681136548519135, 0.027844024822115898, 0.014026278629899025, 0.025656970217823982, 0.0361209474503994, 0.017075760290026665, 0.01003955863416195, 0.016965145245194435, 0.04991300031542778, 0.01522271428257227, 0.007584442384541035, 0.03757705166935921, 0.03609456866979599, 0.10922907292842865, 0.19329114258289337, 0.2903786897659302, 0.29551932215690613, 0.1564989984035492, 0.3518115282058716, NaN, NaN, NaN], [0.017342884093523026, 0.024629754945635796, 0.0017386168474331498, 0.003977979999035597, 0.0011948446044698358, 0.0001711023651296273, 0.0019097719341516495, 0.050265345722436905, 0.048485398292541504, 0.025773482397198677, 0.011941587552428246, 0.02582539990544319, 0.014500979334115982, 0.011088544502854347, 0.0004536270862445235, 0.001346826204098761, 0.09912228584289551, 0.03899921476840973, 0.19399496912956238, 0.33165985345840454, 0.3351045250892639, 0.007158405613154173, 0.26822295784950256, NaN, NaN], [0.15815527737140656, 0.009173951111733913, 0.012453499250113964, 0.01756284572184086, 0.0007500716019421816, 0.0020462200045585632, 0.00166225153952837, 0.05335438624024391, 0.037105023860931396, 0.009711050428450108, 0.05516523867845535, 0.04893142729997635, 0.03887411952018738, 0.002221355913206935, 0.004346344619989395, 0.004376854281872511, 0.001785764587111771, 0.09844812005758286, 0.14674220979213715, 0.34636548161506653, 0.04763580113649368, 0.057022612541913986, 0.12166893482208252, 0.13556897640228271, NaN], [0.16895240545272827, 0.0006144722574390471, 0.0027162963524460793, 0.0007400937611237168, 0.0007253509247675538, 0.0007097159395925701, 0.000199983871425502, 0.0005034026107750833, 0.0002540702698752284, 0.0002154638059437275, 0.0004817947919946164, 0.0019994170870631933, 0.0003459753352217376, 6.575404404429719e-05, 0.004540599416941404, 0.00010029276745626703, 0.0005050064064562321, 0.003569946391507983, 0.008527955040335655, 0.003213587449863553, 0.0022120880894362926, 0.11142478138208389, 0.01313241571187973, 0.055687084794044495, 0.21235007047653198]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13440807163715363, 0.048166193068027496, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14904144406318665, 0.03273539990186691, 0.03615117073059082, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17614386975765228, 0.0854690745472908, 0.038236960768699646, 0.12011754512786865, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14069411158561707, 0.1466522365808487, 0.07941046357154846, 0.06070372834801674, 0.045592159032821655, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15778480470180511, 0.11167039722204208, 0.20017755031585693, 0.10082826018333435, 0.013994856737554073, 0.07346371561288834, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15305520594120026, 0.26692208647727966, 0.1222626119852066, 0.14178596436977386, 0.012799645774066448, 0.019025815650820732, 0.14782781898975372, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.050227321684360504, 0.49922510981559753, 0.2564227879047394, 0.37594476342201233, 0.05222875997424126, 0.019398091360926628, 0.07475102692842484, 0.13636687397956848, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1278427243232727, 0.4489462971687317, 0.09382158517837524, 0.09914611279964447, 0.11451858282089233, 0.14035384356975555, 0.0858180820941925, 0.1395546793937683, 0.05027398467063904, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06907324492931366, 0.44302117824554443, 0.21607427299022675, 0.21861647069454193, 0.14559195935726166, 0.12854896485805511, 0.21420170366764069, 0.5056769251823425, 0.05036870762705803, 0.14160890877246857, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08832916617393494, 0.4917650520801544, 0.16961733996868134, 0.21240676939487457, 0.17275941371917725, 0.13381528854370117, 0.1763075888156891, 0.3443826735019684, 0.022638684138655663, 0.14659351110458374, 0.05034468695521355, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10765255987644196, 0.1569133847951889, 0.14696621894836426, 0.12414205074310303, 0.1321374922990799, 0.32589367032051086, 0.09939466416835785, 0.15668180584907532, 0.035531532019376755, 0.18526552617549896, 0.100669264793396, 0.1766001582145691, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0920143872499466, 0.03631591796875, 0.10338561236858368, 0.13865944743156433, 0.14365890622138977, 0.19164490699768066, 0.08302215486764908, 0.17053648829460144, 0.20418454706668854, 0.4243081212043762, 0.23730118572711945, 0.11353020370006561, 0.062482837587594986, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14247462153434753, 0.10275112092494965, 0.08782284706830978, 0.07633533328771591, 0.09427531808614731, 0.2382509559392929, 0.11237408220767975, 0.1274290829896927, 0.09234490990638733, 0.29983192682266235, 0.19681134819984436, 0.09119200706481934, 0.1394888311624527, 0.02876400761306286, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14126147329807281, 0.06271495670080185, 0.09029032289981842, 0.10313913226127625, 0.08530516922473907, 0.05194256827235222, 0.09853952378034592, 0.05407971888780594, 0.10021005570888519, 0.14394013583660126, 0.19472479820251465, 0.17138735949993134, 0.055624835193157196, 0.022259291261434555, 0.010825252160429955, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15579406917095184, 0.5571659207344055, 0.09220181405544281, 0.09424383193254471, 0.2893342971801758, 0.14449337124824524, 0.08881417661905289, 0.09621196240186691, 0.05768556892871857, 0.34467604756355286, 0.16894927620887756, 0.32070621848106384, 0.32385867834091187, 0.08616255223751068, 0.0030245021916925907, 0.011462957598268986, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06543286889791489, 0.3303832709789276, 0.1981877088546753, 0.17906354367733002, 0.08578304201364517, 0.12075137346982956, 0.09918820112943649, 0.14948950707912445, 0.0696079283952713, 0.2870473861694336, 0.2037079930305481, 0.20505982637405396, 0.415317177772522, 0.18504147231578827, 0.05944397673010826, 0.03780561313033104, 0.06350213289260864, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08806300163269043, 0.5073549151420593, 0.15216797590255737, 0.1779468059539795, 0.08599209040403366, 0.038353316485881805, 0.05095306783914566, 0.13815101981163025, 0.05531492829322815, 0.3680262565612793, 0.045964885503053665, 0.5803228616714478, 0.2365681380033493, 0.10053237527608871, 0.016326427459716797, 0.011199035681784153, 0.02849578857421875, 0.09785498678684235, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10047968477010727, 0.17735490202903748, 0.1303417980670929, 0.1233980730175972, 0.11124629527330399, 0.27208706736564636, 0.09057758748531342, 0.20949512720108032, 0.0595981664955616, 0.32820063829421997, 0.19304482638835907, 0.3008245825767517, 0.24370267987251282, 0.0977335274219513, 0.0604717954993248, 0.08826017379760742, 0.05976974964141846, 0.11658596247434616, 0.26095637679100037, NaN, NaN, NaN, NaN, NaN, NaN], [0.08956606686115265, 0.03296149522066116, 0.07127847522497177, 0.10275094956159592, 0.12852256000041962, 0.15250688791275024, 0.05763629823923111, 0.13953621685504913, 0.2147330343723297, 0.3297017514705658, 0.25630685687065125, 0.3529660999774933, 0.05266188457608223, 0.19866161048412323, 0.08034973591566086, 0.16050152480602264, 0.12120798975229263, 0.21796129643917084, 0.13665789365768433, 0.05867582932114601, NaN, NaN, NaN, NaN, NaN], [0.16931524872779846, 0.06866136193275452, 0.058377113193273544, 0.054153572767972946, 0.06997817754745483, 0.17294903099536896, 0.06504172086715698, 0.09800923615694046, 0.07601338624954224, 0.22323867678642273, 0.17471107840538025, 0.20914696156978607, 0.32561469078063965, 0.04201642796397209, 0.014874166809022427, 0.043757203966379166, 0.11901038885116577, 0.15924809873104095, 0.08216992020606995, 0.13305248320102692, 0.031323518604040146, NaN, NaN, NaN, NaN], [0.14597494900226593, 0.05063166096806526, 0.07245789468288422, 0.08537694066762924, 0.07253167033195496, 0.03945168852806091, 0.07488631457090378, 0.04114159941673279, 0.09447583556175232, 0.11984950304031372, 0.21245841681957245, 0.24130037426948547, 0.053050536662340164, 0.036372195929288864, 0.012788524851202965, 0.05413965508341789, 0.17548364400863647, 0.18113258481025696, 0.17045176029205322, 0.056165628135204315, 0.023532675579190254, 0.007599800359457731, NaN, NaN, NaN], [0.20880575478076935, 0.4742221236228943, 0.0684090405702591, 0.07499475032091141, 0.22897963225841522, 0.11411925405263901, 0.06380540132522583, 0.06602712720632553, 0.04886250197887421, 0.25098055601119995, 0.16695836186408997, 0.41882073879241943, 0.45364588499069214, 0.19780457019805908, 0.004864717833697796, 0.007611281704157591, 0.23698794841766357, 0.08390159159898758, 0.28844529390335083, 0.28151822090148926, 0.0680297240614891, 0.0018790157046169043, 0.008693840354681015, NaN, NaN], [0.06649312376976013, 0.2272576093673706, 0.15548978745937347, 0.13675269484519958, 0.06747769564390182, 0.09888236224651337, 0.07679145783185959, 0.09811051189899445, 0.059132058173418045, 0.16564641892910004, 0.1534833461046219, 0.21299242973327637, 0.46317315101623535, 0.18783308565616608, 0.06707606464624405, 0.07066023349761963, 0.038238298147916794, 0.13390158116817474, 0.1738123893737793, 0.3894510865211487, 0.199345201253891, 0.05267143249511719, 0.03450411930680275, 0.0674150139093399, NaN], [0.13068987429141998, 0.5177554488182068, 0.21822108328342438, 0.17411521077156067, 0.11371950805187225, 0.10282127559185028, 0.14754493534564972, 0.10529720038175583, 0.04059072583913803, 0.1422514021396637, 0.16688787937164307, 0.3468432128429413, 0.07328897714614868, 0.033892080187797546, 0.005811289418488741, 0.006848806049674749, 0.033459149301052094, 0.08608346432447433, 0.29348817467689514, 0.07146795839071274, 0.05563248693943024, 0.008248405531048775, 0.00942459236830473, 0.03898181766271591, 0.13983668386936188]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13037645816802979, 0.08109150826931, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14859925210475922, 0.02925589494407177, 0.0505123995244503, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21387919783592224, 0.03206360712647438, 0.012896520085632801, 0.06630519032478333, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15968731045722961, 0.046736959367990494, 0.014681101776659489, 0.01418250147253275, 0.011044399812817574, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.22570300102233887, 0.051045093685388565, 0.020206425338983536, 0.021926334127783775, 0.008406145498156548, 0.0702541247010231, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.28555917739868164, 0.03329295665025711, 0.036049578338861465, 0.038853298872709274, 0.007190736476331949, 0.006643606815487146, 0.08228380233049393, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2511760890483856, 0.07463249564170837, 0.04988643527030945, 0.0701586976647377, 0.028143733739852905, 0.007391677238047123, 0.02261284738779068, 0.0737045407295227, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15217745304107666, 0.19177564978599548, 0.125013530254364, 0.1473270058631897, 0.20325084030628204, 0.10669662803411484, 0.07946557551622391, 0.027662983164191246, 0.09494684636592865, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13806378841400146, 0.2514709234237671, 0.17176732420921326, 0.21858137845993042, 0.17882317304611206, 0.16198168694972992, 0.20351995527744293, 0.07158615440130234, 0.0266498401761055, 0.23213928937911987, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17152094841003418, 0.15314172208309174, 0.15820659697055817, 0.19208288192749023, 0.19640566408634186, 0.061033159494400024, 0.12321671098470688, 0.07748300582170486, 0.07906179875135422, 0.032524362206459045, 0.08073069155216217, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11935991793870926, 0.25889015197753906, 0.181893989443779, 0.2521744966506958, 0.2510518431663513, 0.1320696324110031, 0.17421388626098633, 0.10352174937725067, 0.13144756853580475, 0.06071629375219345, 0.07381404936313629, 0.11898738145828247, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11384479701519012, 0.12307179719209671, 0.17695116996765137, 0.21105043590068817, 0.2652710974216461, 0.1994313895702362, 0.5530626177787781, 0.33474239706993103, 0.11353342235088348, 0.20157715678215027, 0.12058570981025696, 0.02405776083469391, 0.20302970707416534, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1661912202835083, 0.3088836967945099, 0.3049609959125519, 0.34614017605781555, 0.3287224769592285, 0.19484750926494598, 0.49978625774383545, 0.2471936047077179, 0.14924246072769165, 0.2264283001422882, 0.11719675362110138, 0.028577886521816254, 0.03125511854887009, 0.04683076590299606, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1382068395614624, 0.14312644302845, 0.15027517080307007, 0.2806132137775421, 0.10704077035188675, 0.15715429186820984, 0.3545873463153839, 0.2772214114665985, 0.11900671571493149, 0.16433128714561462, 0.08395379036664963, 0.0337035246193409, 0.08286106586456299, 0.029390821233391762, 0.07092607021331787, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.31265145540237427, 0.17018769681453705, 0.42172688245773315, 0.3373875319957733, 0.26503118872642517, 0.3668123483657837, 0.6080453991889954, 0.3421963155269623, 0.29850897192955017, 0.22005639970302582, 0.08626232296228409, 0.05660916119813919, 0.04967416450381279, 0.020023291930556297, 0.01626538299024105, 0.03365384787321091, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11847452819347382, 0.5065410137176514, 0.4161456227302551, 0.44356557726860046, 0.358999639749527, 0.34202155470848083, 0.6410406231880188, 0.5693260431289673, 0.3344528377056122, 0.3382241725921631, 0.16963228583335876, 0.12081613391637802, 0.09492655098438263, 0.06781262904405594, 0.059771545231342316, 0.013083304278552532, 0.15846344828605652, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14143924415111542, 0.33810776472091675, 0.4273369610309601, 0.4442084729671478, 0.4867575168609619, 0.40271657705307007, 0.7919159531593323, 0.5796146988868713, 0.41502290964126587, 0.19611117243766785, 0.2659074366092682, 0.0590454526245594, 0.09533000737428665, 0.06579555571079254, 0.049002423882484436, 0.011413656175136566, 0.05989237129688263, 0.0694013461470604, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06363721936941147, 0.3402014374732971, 0.30108359456062317, 0.3598821461200714, 0.356340229511261, 0.2955020070075989, 0.3913557827472687, 0.34592464566230774, 0.3881937265396118, 0.23078370094299316, 0.49122318625450134, 0.3432621657848358, 0.1563359946012497, 0.12668228149414062, 0.1534397453069687, 0.06296171993017197, 0.07472987473011017, 0.07419107109308243, 0.08810260146856308, NaN, NaN, NaN, NaN, NaN, NaN], [0.06025628373026848, 0.1445734202861786, 0.2208743691444397, 0.22917300462722778, 0.34805941581726074, 0.30598515272140503, 0.6932811141014099, 0.6030279994010925, 0.2491629421710968, 0.46458470821380615, 0.5228609442710876, 0.2136632800102234, 0.610046923160553, 0.25265923142433167, 0.14038830995559692, 0.07342293113470078, 0.22653138637542725, 0.10003089159727097, 0.02225746400654316, 0.14559555053710938, NaN, NaN, NaN, NaN, NaN], [0.0902293398976326, 0.5066702961921692, 0.45472872257232666, 0.45485398173332214, 0.5058757662773132, 0.3594079613685608, 0.7028806209564209, 0.5180745720863342, 0.25713953375816345, 0.5372852683067322, 0.6213670372962952, 0.2659974694252014, 0.3181111812591553, 0.5259383916854858, 0.33730512857437134, 0.13441412150859833, 0.36266574263572693, 0.10496268421411514, 0.02362431399524212, 0.020191077142953873, 0.04590708762407303, NaN, NaN, NaN, NaN], [0.1059701219201088, 0.2303982675075531, 0.21762119233608246, 0.3580361306667328, 0.17096057534217834, 0.24843183159828186, 0.5131583213806152, 0.47260501980781555, 0.21650557219982147, 0.38561707735061646, 0.416827529668808, 0.1716565638780594, 0.3172723054885864, 0.29216328263282776, 0.47280052304267883, 0.38235870003700256, 0.1798420399427414, 0.1762932986021042, 0.04000748321413994, 0.08066289126873016, 0.03975420445203781, 0.08505715429782867, NaN, NaN, NaN], [0.2317487895488739, 0.2560827136039734, 0.5102789998054504, 0.4199059009552002, 0.44283756613731384, 0.5258800983428955, 0.732390284538269, 0.4491574466228485, 0.4244932234287262, 0.5298821926116943, 0.43037980794906616, 0.2800268232822418, 0.3093121647834778, 0.4250229299068451, 0.19317308068275452, 0.2640416920185089, 0.38813653588294983, 0.11181202530860901, 0.054203763604164124, 0.037284549325704575, 0.018739882856607437, 0.014264266937971115, 0.035236652940511703, NaN, NaN], [0.08032029122114182, 0.6358892321586609, 0.5042787194252014, 0.5074477195739746, 0.5223307013511658, 0.5343775749206543, 0.703619122505188, 0.6657658815383911, 0.45647403597831726, 0.602655827999115, 0.5387927889823914, 0.39006462693214417, 0.39567169547080994, 0.43596506118774414, 0.41000646352767944, 0.269907683134079, 0.5412885546684265, 0.2038634866476059, 0.10306636989116669, 0.05501747503876686, 0.04515310004353523, 0.04695969074964523, 0.008877278305590153, 0.09985174983739853, NaN], [0.03129265457391739, 0.2636677324771881, 0.3672870099544525, 0.438161164522171, 0.7497870922088623, 0.43876102566719055, 0.6747432947158813, 0.5918557643890381, 0.5535795092582703, 0.7133825421333313, 0.7440239787101746, 0.3780657947063446, 0.4423457384109497, 0.6450315713882446, 0.5939705967903137, 0.7279283404350281, 0.4253756105899811, 0.4950290024280548, 0.13756991922855377, 0.08432447165250778, 0.11775307357311249, 0.12791647017002106, 0.07922011613845825, 0.04417572543025017, 0.3473970592021942]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13398022949695587, 0.051660239696502686, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14254364371299744, 0.023038247600197792, 0.14531654119491577, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17795929312705994, 0.024941343814134598, 0.06730933487415314, 0.21388311684131622, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09399491548538208, 0.3603954315185547, 0.2704434394836426, 0.1475897580385208, 0.18568314611911774, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14775781333446503, 0.19919507205486298, 0.14170727133750916, 0.05924544855952263, 0.05067846551537514, 0.45942243933677673, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14211317896842957, 0.055850330740213394, 0.31645503640174866, 0.16900919377803802, 0.038168299943208694, 0.07897188514471054, 0.2625669240951538, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08848852664232254, 0.1616290658712387, 0.37575462460517883, 0.24721546471118927, 0.16591095924377441, 0.06889674067497253, 0.052010323852300644, 0.12634019553661346, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0747382640838623, 0.14914710819721222, 0.6135430335998535, 0.5929751992225647, 0.35069379210472107, 0.2108047604560852, 0.11502823978662491, 0.02365955151617527, 0.17759312689304352, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02855301834642887, 0.21659326553344727, 0.4310435652732849, 0.40604472160339355, 0.3670090436935425, 0.48140615224838257, 0.27167943120002747, 0.09097199141979218, 0.1627163589000702, 0.1288144737482071, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03365316241979599, 0.14809295535087585, 0.3644290566444397, 0.4046455919742584, 0.26744210720062256, 0.32108214497566223, 0.1678413599729538, 0.190241739153862, 0.22121649980545044, 0.03444775566458702, 0.46765974164009094, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.038216885179281235, 0.2552680969238281, 0.4071650505065918, 0.3936895430088043, 0.4416206479072571, 0.38015541434288025, 0.1657901555299759, 0.15260477364063263, 0.22771137952804565, 0.10614379495382309, 0.0724361315369606, 0.1760038137435913, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07068492472171783, 0.07818713039159775, 0.3302493095397949, 0.299561083316803, 0.46339741349220276, 0.48102065920829773, 0.15714748203754425, 0.27301517128944397, 0.38065311312675476, 0.19789563119411469, 0.11113718152046204, 0.05171056091785431, 0.13386131823062897, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05115865543484688, 0.44867002964019775, 0.49208834767341614, 0.477664977312088, 0.4642978608608246, 0.46059542894363403, 0.25649622082710266, 0.406831830739975, 0.27858051657676697, 0.2405669242143631, 0.11958811432123184, 0.1450459510087967, 0.0628136694431305, 0.09898709505796432, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04031704366207123, 0.6707005500793457, 0.529548704624176, 0.4586588144302368, 0.3106471002101898, 0.6713098287582397, 0.4458201229572296, 0.5507155060768127, 0.6255134344100952, 0.5032600164413452, 0.18919125199317932, 0.2968505918979645, 0.3902440667152405, 0.16804949939250946, 0.088200144469738, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13188821077346802, 0.1971314549446106, 0.3902590274810791, 0.4961083233356476, 0.37017205357551575, 0.46889960765838623, 0.2874276340007782, 0.1815745085477829, 0.39618349075317383, 0.17909032106399536, 0.26052209734916687, 0.13463276624679565, 0.11223814636468887, 0.05094114691019058, 0.030694767832756042, 0.23131275177001953, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.029627619311213493, 0.0727827325463295, 0.2382729947566986, 0.16726669669151306, 0.3644602298736572, 0.47072863578796387, 0.2034798413515091, 0.1723088026046753, 0.43477845191955566, 0.18565386533737183, 0.3540991544723511, 0.2379947453737259, 0.07713616639375687, 0.19858470559120178, 0.17015229165554047, 0.0891638696193695, 0.22899208962917328, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01839388906955719, 0.10223808884620667, 0.244280606508255, 0.22035017609596252, 0.2828108072280884, 0.41914066672325134, 0.09010869264602661, 0.14338640868663788, 0.35142722725868225, 0.12073972821235657, 0.6723650693893433, 0.17433631420135498, 0.20010362565517426, 0.17566151916980743, 0.17214345932006836, 0.06743419170379639, 0.08234895765781403, 0.4274884760379791, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02117752842605114, 0.17625343799591064, 0.2448491007089615, 0.23410049080848694, 0.3357784152030945, 0.2992798388004303, 0.09099920094013214, 0.1110134869813919, 0.20308172702789307, 0.1763213574886322, 0.1646280288696289, 0.23259523510932922, 0.3615821301937103, 0.32664546370506287, 0.296549916267395, 0.2726198732852936, 0.07387500256299973, 0.07587912678718567, 0.14093360304832458, NaN, NaN, NaN, NaN, NaN, NaN], [0.05486638844013214, 0.06597498804330826, 0.2194771021604538, 0.1927901804447174, 0.37433308362960815, 0.412477970123291, 0.07100911438465118, 0.1499587744474411, 0.3056679368019104, 0.16932857036590576, 0.15193165838718414, 0.19111526012420654, 0.291239857673645, 0.37710845470428467, 0.510109543800354, 0.47089657187461853, 0.17204606533050537, 0.09759342670440674, 0.05198577418923378, 0.1557197868824005, NaN, NaN, NaN, NaN, NaN], [0.03942986950278282, 0.2940163016319275, 0.3192412853240967, 0.3550935387611389, 0.28974649310112, 0.35144588351249695, 0.111830934882164, 0.2212614268064499, 0.1942923218011856, 0.16557106375694275, 0.12293191254138947, 0.3516637980937958, 0.22679129242897034, 0.3504909574985504, 0.4427362084388733, 0.6422855854034424, 0.29741936922073364, 0.17250965535640717, 0.13341550529003143, 0.05469499155879021, 0.0792233869433403, NaN, NaN, NaN, NaN], [0.03949292004108429, 0.6095755696296692, 0.4376317858695984, 0.4024345874786377, 0.24819140136241913, 0.555855929851532, 0.2881583273410797, 0.40402302145957947, 0.5775710940361023, 0.42070186138153076, 0.22824901342391968, 0.4547353982925415, 0.567461371421814, 0.5762937664985657, 0.33163049817085266, 0.41951635479927063, 0.37286072969436646, 0.25620296597480774, 0.25266289710998535, 0.3395143151283264, 0.13239842653274536, 0.07333662360906601, NaN, NaN, NaN], [0.11607979983091354, 0.18507249653339386, 0.30528268218040466, 0.41669708490371704, 0.22673273086547852, 0.3321194052696228, 0.17922396957874298, 0.1181870847940445, 0.299829363822937, 0.11785572022199631, 0.23005077242851257, 0.1731709986925125, 0.17971253395080566, 0.2448451966047287, 0.15796169638633728, 0.701153576374054, 0.1659945547580719, 0.4861533045768738, 0.20215842127799988, 0.13506482541561127, 0.058445703238248825, 0.03114200383424759, 0.21790345013141632, NaN, NaN], [0.017429474741220474, 0.04190561920404434, 0.14842365682125092, 0.09654705971479416, 0.16489917039871216, 0.24686570465564728, 0.09686223417520523, 0.09368213266134262, 0.2918589413166046, 0.08991989493370056, 0.18521137535572052, 0.19666530191898346, 0.06316249072551727, 0.222347229719162, 0.3215444087982178, 0.3288835287094116, 0.38603323698043823, 0.4142700135707855, 0.25910744071006775, 0.0714699923992157, 0.2130158245563507, 0.1895158588886261, 0.07420682162046432, 0.2235250473022461, NaN], [0.011625233106315136, 0.13701221346855164, 0.3079974055290222, 0.17742200195789337, 0.10538481175899506, 0.17213597893714905, 0.08605048805475235, 0.13507568836212158, 0.2275547832250595, 0.07923908531665802, 0.07705283164978027, 0.2479921281337738, 0.3453103303909302, 0.2883259654045105, 0.36409828066825867, 0.18068012595176697, 0.4896908700466156, 0.399289608001709, 0.5261627435684204, 0.6339481472969055, 0.6382991671562195, 0.5417840480804443, 0.2542280852794647, 0.330732524394989, 0.21995915472507477]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04915444552898407, 0.7444152235984802, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10270431637763977, 0.20103313028812408, 0.23083212971687317, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1558120846748352, 0.09243088960647583, 0.02280065417289734, 0.32627996802330017, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1265193670988083, 0.1639627069234848, 0.12297425419092178, 0.08557231724262238, 0.1833999902009964, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11118379235267639, 0.23907560110092163, 0.16732671856880188, 0.1982172429561615, 0.02825341187417507, 0.15412425994873047, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06564534455537796, 0.4107542335987091, 0.09891282767057419, 0.3507450222969055, 0.0021941487211734056, 0.004341787192970514, 0.11288701742887497, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09254656732082367, 0.17870496213436127, 0.11882538348436356, 0.2565489113330841, 0.06709786504507065, 0.020701991394162178, 0.05621851608157158, 0.571487307548523, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12130707502365112, 0.06869146227836609, 0.052872415632009506, 0.07373122870922089, 0.03967232629656792, 0.019552208483219147, 0.024196362122893333, 0.1570335328578949, 0.3329051434993744, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12370187789201736, 0.027735348790884018, 0.007442266680300236, 0.018701551482081413, 0.04923407360911369, 0.022976329550147057, 0.06834850460290909, 0.13354788720607758, 0.13089321553707123, 0.41554775834083557, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08012630045413971, 0.020899765193462372, 0.032236725091934204, 0.011631320230662823, 0.1322554349899292, 0.13739252090454102, 0.3272823691368103, 0.10228703171014786, 0.16136890649795532, 0.12631160020828247, 0.3315902352333069, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07002493739128113, 0.03239390626549721, 0.05209453031420708, 0.033656563609838486, 0.10301846265792847, 0.08080227673053741, 0.10908480733633041, 0.10694557428359985, 0.2992934286594391, 0.26628223061561584, 0.1579413264989853, 0.18216297030448914, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.23901967704296112, 0.02059122547507286, 0.03393668681383133, 0.04736512154340744, 0.05927135422825813, 0.02361929975450039, 0.006761881057173014, 0.05556455999612808, 0.1379650980234146, 0.12424714863300323, 0.191926509141922, 0.01547694206237793, 0.05743350088596344, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0662187710404396, 0.02669837884604931, 0.008789082989096642, 0.004751283209770918, 0.0528719425201416, 0.011242655105888844, 0.018989307805895805, 0.07620660215616226, 0.012969521805644035, 0.039284493774175644, 0.22954939305782318, 0.04563957825303078, 0.029234008863568306, 0.7488549947738647, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10826153308153152, 0.014460555277764797, 0.0725417360663414, 0.03217141702771187, 0.06698039174079895, 0.08051858842372894, 0.05872708931565285, 0.022866755723953247, 0.06705553829669952, 0.07034263759851456, 0.3507814407348633, 0.05356235057115555, 0.08709309250116348, 0.23604632914066315, 0.324868768453598, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13878783583641052, 0.02536645717918873, 0.06943535804748535, 0.05891912057995796, 0.006977759767323732, 0.003910682164132595, 0.004916978534311056, 0.04463541880249977, 0.07985055446624756, 0.07872368395328522, 0.291103333234787, 0.21302121877670288, 0.16995804011821747, 0.19893744587898254, 0.01890285685658455, 0.3838881254196167, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04579493775963783, 0.04550570994615555, 0.013287660665810108, 0.023886512964963913, 0.024052713066339493, 0.017023656517267227, 0.04836693033576012, 0.030526861548423767, 0.017645621672272682, 0.03170713782310486, 0.09266000241041183, 0.23106807470321655, 0.03557471185922623, 0.12432269752025604, 0.10334902256727219, 0.3233395516872406, 0.3770029842853546, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0394071489572525, 0.011173942126333714, 0.019201254472136497, 0.012027204036712646, 0.1043756976723671, 0.09629304707050323, 0.044260744005441666, 0.010774374939501286, 0.027033720165491104, 0.01529898401349783, 0.004158060997724533, 0.03471178933978081, 0.3574643135070801, 0.04469288885593414, 0.27014297246932983, 0.10925178974866867, 0.34427598118782043, 0.2875407040119171, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08343059569597244, 0.043180350214242935, 0.0767669752240181, 0.06360654532909393, 0.1271795630455017, 0.0800960585474968, 0.06889919936656952, 0.05648425221443176, 0.1521727591753006, 0.09240606427192688, 0.03566697984933853, 0.03560119867324829, 0.1492718607187271, 0.18653850257396698, 0.3474813401699066, 0.3278762698173523, 0.10706853121519089, 0.127774178981781, 0.1299499273300171, NaN, NaN, NaN, NaN, NaN, NaN], [0.23721955716609955, 0.02343675307929516, 0.03610215708613396, 0.05973569303750992, 0.07488072663545609, 0.026813305914402008, 0.0050082337111234665, 0.03149579092860222, 0.06251367926597595, 0.02305557392537594, 0.025774041190743446, 0.007636546157300472, 0.004965651780366898, 0.09922869503498077, 0.133448526263237, 0.1956746131181717, 0.04676169902086258, 0.27956491708755493, 0.021136147901415825, 0.057313986122608185, NaN, NaN, NaN, NaN, NaN], [0.0697786882519722, 0.028010839596390724, 0.012634677812457085, 0.007894599810242653, 0.0697624459862709, 0.015741104260087013, 0.01737123914062977, 0.05471426621079445, 0.0063003492541611195, 0.009287585504353046, 0.02825707383453846, 0.016440505161881447, 0.0038715004920959473, 0.07019948214292526, 0.02518516778945923, 0.041359793394804, 0.06545242667198181, 0.29174378514289856, 0.05010553449392319, 0.020036837086081505, 0.7549301981925964, NaN, NaN, NaN, NaN], [0.12042609602212906, 0.016146911308169365, 0.09666067361831665, 0.04101520776748657, 0.09386932849884033, 0.11830881983041763, 0.08227012306451797, 0.02001151442527771, 0.0443122573196888, 0.028465820476412773, 0.11253371834754944, 0.02299223281443119, 0.013287386856973171, 0.043506089597940445, 0.09705191105604172, 0.08899306505918503, 0.14267200231552124, 0.1414598524570465, 0.04555709660053253, 0.08242949843406677, 0.2358742356300354, 0.30384859442710876, NaN, NaN, NaN], [0.14026813209056854, 0.02709769457578659, 0.07936792075634003, 0.07383942604064941, 0.01026969589293003, 0.007506935391575098, 0.01013263501226902, 0.043357811868190765, 0.054843299090862274, 0.032377004623413086, 0.07885654270648956, 0.05951513722538948, 0.021026868373155594, 0.029062975198030472, 0.004067933652549982, 0.00896876398473978, 0.031901001930236816, 0.2457016408443451, 0.1949184089899063, 0.16180625557899475, 0.23649972677230835, 0.020314330235123634, 0.390868216753006, NaN, NaN], [0.036581799387931824, 0.048626694828271866, 0.015552042052149773, 0.027681825682520866, 0.03610476478934288, 0.033903565257787704, 0.10816461592912674, 0.038128215819597244, 0.015381437726318836, 0.020138615742325783, 0.04596110060811043, 0.12391334027051926, 0.008882056921720505, 0.017164889723062515, 0.019657107070088387, 0.039318498224020004, 0.012226631864905357, 0.12883862853050232, 0.2578184902667999, 0.03228205814957619, 0.13855229318141937, 0.08962707966566086, 0.32015570998191833, 0.32621434330940247, NaN], [0.16620944440364838, 0.03880922496318817, 0.027515552937984467, 0.018877340480685234, 0.019147777929902077, 0.2389368712902069, 0.02623477764427662, 0.012871777638792992, 0.013969821855425835, 0.021991701796650887, 0.0026013199239969254, 0.00741098215803504, 0.01774594374001026, 0.003101027337834239, 0.007316285278648138, 0.009464021772146225, 0.007634901907294989, 0.005969886668026447, 0.011287253350019455, 0.04429420828819275, 0.016200777143239975, 0.03440575301647186, 0.14183124899864197, 0.1436305195093155, 0.03402799740433693]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13823550939559937, 0.01690824329853058, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1366243064403534, 0.10029595345258713, 0.03309698402881622, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14204008877277374, 0.17578311264514923, 0.058153361082077026, 0.03275991603732109, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15378697216510773, 0.06811928749084473, 0.031730279326438904, 0.02174059860408306, 0.06419884413480759, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2336570769548416, 0.05475717782974243, 0.004165933933109045, 0.0025384188629686832, 0.005177688784897327, 0.12858138978481293, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1292651742696762, 0.01662198081612587, 0.01174056064337492, 0.002378111705183983, 0.04036910459399223, 0.6038607358932495, 0.053664252161979675, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13257111608982086, 0.0015173845458775759, 0.11979293078184128, 0.025075461715459824, 0.17128729820251465, 0.38108551502227783, 0.04533570259809494, 0.02173132263123989, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12533389031887054, 0.01691550202667713, 0.03341663256287575, 0.04296481981873512, 0.13898836076259613, 0.21484552323818207, 0.09921174496412277, 0.178620383143425, 0.08540544658899307, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19628551602363586, 0.0262758769094944, 0.06177970767021179, 0.020167797803878784, 0.21508394181728363, 0.05243970826268196, 0.05236654728651047, 0.019688904285430908, 0.04470491781830788, 0.03636182099580765, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10685201734304428, 0.1520930975675583, 0.22691352665424347, 0.1206204891204834, 0.20647111535072327, 0.3387817144393921, 0.17652125656604767, 0.14866295456886292, 0.058651361614465714, 0.13512541353702545, 0.029732942581176758, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14931687712669373, 0.17397953569889069, 0.045104723423719406, 0.029273295775055885, 0.009919327683746815, 0.05321130529046059, 0.40632039308547974, 0.053491849452257156, 0.10154163092374802, 0.08916116505861282, 0.038379959762096405, 0.050926242023706436, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1467411071062088, 0.6613936424255371, 0.30691561102867126, 0.27473992109298706, 0.05103013291954994, 0.09803401678800583, 0.18992389738559723, 0.012332501821219921, 0.08918186277151108, 0.009687116369605064, 0.01925584301352501, 0.0046735359355807304, 0.006799460854381323, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.23535212874412537, 0.03722311928868294, 0.0383867472410202, 0.06886720657348633, 0.040591221302747726, 0.07368911802768707, 0.09838991612195969, 0.052333034574985504, 0.3684787154197693, 0.05692664161324501, 0.030762571841478348, 0.0074586388655006886, 0.017855344340205193, 0.004115242511034012, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17482686042785645, 0.020169643685221672, 0.038628242909908295, 0.03409411385655403, 0.011309999041259289, 0.013418656773865223, 0.010934274643659592, 0.0036632094997912645, 0.017374617978930473, 0.023464469239115715, 0.0031370571814477444, 0.004764250945299864, 0.022831382229924202, 0.0012565170181915164, 0.01132481824606657, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2204812914133072, 0.0262058824300766, 0.011961801908910275, 0.00864139012992382, 0.033310361206531525, 0.014301336370408535, 0.009627565741539001, 0.26419174671173096, 0.09070254862308502, 0.04369048774242401, 0.05080936849117279, 0.022543352097272873, 0.012377972714602947, 0.030277462676167488, 0.2341402769088745, 0.01971697248518467, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.253863126039505, 0.004828702192753553, 0.05376851186156273, 0.11550138890743256, 0.1064227893948555, 0.03894256055355072, 0.006152869202196598, 0.03161965310573578, 0.06215812265872955, 0.10950783640146255, 0.01032247580587864, 0.005066303536295891, 0.011880352161824703, 0.09494113177061081, 0.06700112670660019, 0.10617008060216904, 0.020382743328809738, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04813924431800842, 0.008662978187203407, 0.10469061881303787, 0.06787187606096268, 0.02962217852473259, 0.04144993796944618, 0.019078848883509636, 0.10597121715545654, 0.0923849567770958, 0.24696239829063416, 0.010940729640424252, 0.060362689197063446, 0.059540145099163055, 0.36283043026924133, 0.1817280501127243, 0.2542697787284851, 0.10456714779138565, 0.017782384529709816, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10143542289733887, 0.13917230069637299, 0.040259018540382385, 0.030723553150892258, 0.006155712995678186, 0.031952716410160065, 0.3338092863559723, 0.06915750354528427, 0.1324792504310608, 0.11542332917451859, 0.05764009431004524, 0.04023035988211632, 0.03596781566739082, 0.1495574563741684, 0.02840258926153183, 0.049019940197467804, 0.4096885919570923, 0.03150010108947754, 0.02953496389091015, NaN, NaN, NaN, NaN, NaN, NaN], [0.1521255224943161, 0.6490614414215088, 0.39427587389945984, 0.3861289620399475, 0.05361294746398926, 0.09808307886123657, 0.16810499131679535, 0.014004985801875591, 0.1451900601387024, 0.008040589280426502, 0.022555561736226082, 0.013471563346683979, 0.006859058979898691, 0.05312783271074295, 0.04058152437210083, 0.023753749206662178, 0.3811529278755188, 0.052651502192020416, 0.007359141018241644, 0.007947265170514584, NaN, NaN, NaN, NaN, NaN], [0.2650813162326813, 0.032561566680669785, 0.05222610384225845, 0.09714324027299881, 0.038093939423561096, 0.08016244322061539, 0.09171951562166214, 0.056265611201524734, 0.42980653047561646, 0.0462084598839283, 0.03524700179696083, 0.017182864248752594, 0.04137876257300377, 0.007372017949819565, 0.08077534288167953, 0.07507885992527008, 0.050101280212402344, 0.02560576982796192, 0.006666052620857954, 0.016142593696713448, 0.003943128511309624, NaN, NaN, NaN, NaN], [0.186274453997612, 0.02024305984377861, 0.052268851548433304, 0.04830823838710785, 0.011142827570438385, 0.015970220789313316, 0.01383616030216217, 0.004258061293512583, 0.024750858545303345, 0.02320612221956253, 0.004944193176925182, 0.006908308248966932, 0.022138824686408043, 0.002315782941877842, 0.022694725543260574, 0.010753386653959751, 0.0032616793178021908, 0.0013332129456102848, 0.0031688748858869076, 0.015737321227788925, 0.00092066585784778, 0.009911282919347286, NaN, NaN, NaN], [0.2620354890823364, 0.032388050109148026, 0.01473915670067072, 0.01008685864508152, 0.03682388737797737, 0.017798764631152153, 0.012407293543219566, 0.2692665457725525, 0.10958822816610336, 0.03793380409479141, 0.07735131680965424, 0.03087974339723587, 0.01817244663834572, 0.0740593820810318, 0.5664002895355225, 0.01639901101589203, 0.07361851632595062, 0.02498074807226658, 0.01953950524330139, 0.011185318231582642, 0.024920325726270676, 0.19407986104488373, 0.01722806692123413, NaN, NaN], [0.27593934535980225, 0.005811678245663643, 0.07111961394548416, 0.13982559740543365, 0.1345955729484558, 0.06462955474853516, 0.009384723380208015, 0.03974011912941933, 0.0818282812833786, 0.09768332540988922, 0.015042337588965893, 0.006764655001461506, 0.01590757444500923, 0.11177312582731247, 0.1289886087179184, 0.2743605673313141, 0.018859822303056717, 0.01428449247032404, 0.0072670611552894115, 0.013756940141320229, 0.08787993341684341, 0.08323681354522705, 0.09635237604379654, 0.025643613189458847, NaN], [0.17263205349445343, 0.01194645743817091, 0.02866498939692974, 0.16296441853046417, 0.0019488729303702712, 0.034664519131183624, 0.05397665500640869, 0.1285821497440338, 0.10828299820423126, 0.02950196899473667, 0.008275950327515602, 0.008977574296295643, 0.09588290750980377, 0.01758315972983837, 0.00981396809220314, 0.06520896404981613, 0.03634792938828468, 0.007794357370585203, 0.007516053505241871, 0.0633511170744896, 0.016588596627116203, 0.008872142061591148, 0.04887184873223305, 0.025813041254878044, 0.0022019031457602978]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13826748728752136, 0.016647184267640114, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12115656584501266, 0.053111400455236435, 0.35221540927886963, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06620940566062927, 0.0874415934085846, 0.3174281120300293, 0.09698687493801117, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05510773882269859, 0.045387670397758484, 0.35701045393943787, 0.5011870265007019, 0.0787656381726265, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05231153964996338, 0.1393265277147293, 0.34751832485198975, 0.15474379062652588, 0.1892920285463333, 0.06652400642633438, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04669328033924103, 0.038986966013908386, 0.38860636949539185, 0.09904015064239502, 0.3339899182319641, 0.027963249012827873, 0.04134462773799896, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20758312940597534, 0.07789289951324463, 0.047907259315252304, 0.006299893371760845, 0.2608397901058197, 0.044556185603141785, 0.061705876141786575, 0.034865181893110275, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18052776157855988, 0.08179321140050888, 0.059846919029951096, 0.02793782763183117, 0.062999427318573, 0.04310278594493866, 0.024987775832414627, 0.015387488529086113, 0.132792130112648, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03587701544165611, 0.020078828558325768, 0.04571571201086044, 0.02593454346060753, 0.007220670115202665, 0.03280382603406906, 0.012364541180431843, 0.04736338183283806, 0.48638036847114563, 0.015403805300593376, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.010417330078780651, 0.019508572295308113, 0.03964173421263695, 0.041229844093322754, 0.021899865940213203, 0.0029071751050651073, 0.010124437510967255, 0.08508285880088806, 0.40291228890419006, 0.4734281599521637, 0.015163381583988667, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08744391798973083, 0.1107466071844101, 0.15557123720645905, 0.13837403059005737, 0.05803389474749565, 0.026755833998322487, 0.03754325956106186, 0.4220706820487976, 0.16102783381938934, 0.2859216034412384, 0.1457504779100418, 0.03281670808792114, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21633882820606232, 0.07441287487745285, 0.04740259423851967, 0.026924576610326767, 0.012407396920025349, 0.002398786135017872, 0.0038467273116111755, 0.13835540413856506, 0.06710492819547653, 0.026295386254787445, 0.17057135701179504, 0.013244924135506153, 0.46883779764175415, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.027107199653983116, 0.05742119997739792, 0.06533583253622055, 0.024222400039434433, 0.014050583355128765, 0.013653005473315716, 0.0030738371424376965, 0.04425956308841705, 0.06826918572187424, 0.011929179541766644, 0.14959540963172913, 0.16161218285560608, 0.5212987065315247, 0.041249219328165054, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12232528626918793, 0.02327316626906395, 0.043996360152959824, 0.010462167672812939, 0.05786772817373276, 0.006097386125475168, 0.001271827262826264, 0.022651376202702522, 0.03627351298928261, 0.030646052211523056, 0.03145253658294678, 0.18536151945590973, 0.10030946880578995, 0.3235938847064972, 0.09760642796754837, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01696004532277584, 0.0005225083441473544, 0.012039890512824059, 0.0003213977033738047, 0.024568837136030197, 0.0005492557538673282, 6.035636397427879e-05, 0.0032521369867026806, 0.016784805804491043, 0.013033770024776459, 0.023488081991672516, 0.04594254866242409, 0.04732683673501015, 0.2366781234741211, 0.2578820288181305, 0.02447950839996338, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.016271475702524185, 0.026037830859422684, 0.05988215655088425, 0.04065781086683273, 0.0548781082034111, 0.0059303357265889645, 0.000490839418489486, 0.009792556054890156, 0.05564826726913452, 0.029693011194467545, 0.015783851966261864, 0.050408631563186646, 0.10483089834451675, 0.18894171714782715, 0.4590488076210022, 0.24355939030647278, 0.03408684581518173, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.011992339976131916, 0.02786487340927124, 0.025577154010534286, 0.02912752889096737, 0.009845648892223835, 0.0007121131638996303, 0.001387864351272583, 0.015649031847715378, 0.05334821715950966, 0.05039743706583977, 0.0003855754912365228, 0.07798124849796295, 0.03745294734835625, 0.16697214543819427, 0.29521557688713074, 0.2776513993740082, 0.29445046186447144, 0.031993161886930466, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11517049372196198, 0.11416894942522049, 0.19162771105766296, 0.14611610770225525, 0.060761958360672, 0.02055470645427704, 0.021888524293899536, 0.20655019581317902, 0.047658227384090424, 0.055987950414419174, 0.01683689095079899, 0.005808014422655106, 0.045862384140491486, 0.09340663254261017, 0.10908356308937073, 0.18944555521011353, 0.26804569363594055, 0.20485185086727142, 0.037772081792354584, NaN, NaN, NaN, NaN, NaN, NaN], [0.24184046685695648, 0.07921410351991653, 0.056290365755558014, 0.026794791221618652, 0.016941547393798828, 0.0021516080014407635, 0.0023830668069422245, 0.05685606598854065, 0.02070370689034462, 0.003236053278669715, 0.01165463775396347, 0.004370343871414661, 0.030780060216784477, 0.00907946564257145, 0.06188458576798439, 0.04407832771539688, 0.006142587400972843, 0.14762946963310242, 0.013672620058059692, 0.4999893307685852, NaN, NaN, NaN, NaN, NaN], [0.03566991165280342, 0.0538097508251667, 0.09943600744009018, 0.028607800602912903, 0.020965654402971268, 0.013461945578455925, 0.002478980924934149, 0.02911236882209778, 0.02446376532316208, 0.0022762087173759937, 0.010774179361760616, 0.04047773778438568, 0.06471210718154907, 0.0026813328731805086, 0.07523855566978455, 0.030470186844468117, 0.0345987044274807, 0.1238497719168663, 0.17781274020671844, 0.4970780611038208, 0.04515520855784416, NaN, NaN, NaN, NaN], [0.12716706097126007, 0.02434932254254818, 0.05787394568324089, 0.013031681068241596, 0.06681805849075317, 0.007088592275977135, 0.0018475945107638836, 0.021072670817375183, 0.024636711925268173, 0.010089303366839886, 0.0076353950425982475, 0.05158482864499092, 0.009980393573641777, 0.034229546785354614, 0.01627102866768837, 0.008032353594899178, 0.013575052842497826, 0.04940066114068031, 0.19428585469722748, 0.10819438844919205, 0.2976790964603424, 0.08516447991132736, NaN, NaN, NaN], [0.01713084802031517, 0.000499976216815412, 0.019638467580080032, 0.00048709739348851144, 0.03356647491455078, 0.0008144291932694614, 0.00011953162174904719, 0.003664336632937193, 0.013800683431327343, 0.004805452190339565, 0.004433726891875267, 0.011711561121046543, 0.003556638490408659, 0.01588965393602848, 0.025807680562138557, 0.00022126971452962607, 0.004036479629576206, 0.00837762001901865, 0.04655361920595169, 0.04086336866021156, 0.22630761563777924, 0.2765483856201172, 0.02425519935786724, NaN, NaN], [0.010901566594839096, 0.020337969064712524, 0.07802019268274307, 0.0504593625664711, 0.06312800198793411, 0.009868033230304718, 0.000861799344420433, 0.010114955715835094, 0.052247028797864914, 0.012602821923792362, 0.005399123765528202, 0.01934058591723442, 0.013776490464806557, 0.010564911179244518, 0.04300173744559288, 0.008748980239033699, 0.0006391598144546151, 0.006108305882662535, 0.05087457224726677, 0.09035929292440414, 0.18751013278961182, 0.4462290108203888, 0.28552356362342834, 0.05451636388897896, NaN], [0.1367119550704956, 0.02979014255106449, 0.04602046683430672, 0.022530242800712585, 0.009278235025703907, 0.01184787880629301, 0.010125648230314255, 0.02445557340979576, 0.052750833332538605, 0.013119504787027836, 0.0006633299053646624, 0.007243738044053316, 0.02398994006216526, 0.00908573716878891, 0.013761860318481922, 0.007176807615906, 0.00677318312227726, 0.0021949538495391607, 0.01309704128652811, 0.09677710384130478, 0.12711098790168762, 0.1613820642232895, 0.37058699131011963, 0.3504316806793213, 0.02586444839835167]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13988038897514343, 0.003474950324743986, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14879919588565826, 0.018745053559541702, 0.07372914999723434, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.030327370390295982, 0.02692173607647419, 0.46947386860847473, 0.09036581218242645, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.164228156208992, 0.0009850627975538373, 0.0044541023671627045, 0.0005622706958092749, 0.024160074070096016, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.020124448463320732, 0.0011880549136549234, 0.0042731426656246185, 3.242780803702772e-05, 0.6858344078063965, 0.023040860891342163, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0017230550292879343, 3.356653905939311e-05, 0.001307086437009275, 1.4968540199333802e-05, 0.5564903616905212, 0.236929789185524, 0.007688341196626425, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1612924486398697, 0.00029754414572380483, 0.0029063820838928223, 0.0015110797248780727, 0.16695675253868103, 0.3453270196914673, 0.07193248718976974, 0.006359610706567764, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1910298615694046, 0.01051796693354845, 0.0018660163041204214, 0.0012154864380136132, 0.022663934156298637, 0.008557457476854324, 0.016767704859375954, 0.05246622860431671, 0.08816055208444595, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24295811355113983, 0.0012021175352856517, 0.0005200211890041828, 0.00015996988804545254, 0.002627951791509986, 0.03450923040509224, 0.014827161096036434, 0.015967652201652527, 0.005632439162582159, 0.001854590023867786, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2492469847202301, 0.004325273912400007, 0.004784590099006891, 0.013903478160500526, 0.0013026667293161154, 0.003877879586070776, 0.017029188573360443, 0.01781909167766571, 0.05003270506858826, 0.026610376313328743, 0.008462576195597649, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.25306010246276855, 0.0017952719936147332, 0.005404005758464336, 0.021692873910069466, 0.0005702165653929114, 9.544018394080922e-05, 0.001603480544872582, 0.001225438085384667, 0.036846794188022614, 0.001749897957779467, 0.016878794878721237, 0.021703237667679787, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.055758021771907806, 0.000425096252001822, 0.0005783061496913433, 0.0011671994579955935, 0.00034630659501999617, 0.00031045774812810123, 0.0006358043756335974, 0.004018810577690601, 0.0004720573779195547, 0.006387148518115282, 0.038948215544223785, 0.40798652172088623, 0.0038703898899257183, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.29551389813423157, 0.006183725781738758, 0.0010477532632648945, 0.001470124931074679, 0.0028535614255815744, 0.003910644445568323, 0.004942604340612888, 0.003798475954681635, 0.01567114144563675, 0.060374900698661804, 0.006600319407880306, 0.010896215215325356, 0.009779008105397224, 0.007320093456655741, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1632017195224762, 0.00519327400252223, 0.00790441408753395, 0.0009941658936440945, 0.3241596221923828, 0.0008480648975819349, 0.0001429034018656239, 0.0012253100285306573, 0.0008457236108370125, 0.006411578040570021, 0.0016067628748714924, 0.003762597683817148, 0.029224932193756104, 0.07677540183067322, 0.06338826566934586, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.005401996895670891, 6.3005199990584515e-06, 0.0004310416697990149, 8.47076989884954e-06, 0.009243682958185673, 0.0008590375073254108, 4.37394373875577e-06, 6.523932825075462e-05, 8.531090134056285e-05, 0.0006816720124334097, 7.644478318979964e-05, 0.00018924157484434545, 0.0012375408550724387, 0.023784970864653587, 0.4309314787387848, 0.034907225519418716, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.29775136709213257, 0.006892140489071608, 0.009814155288040638, 0.016249310225248337, 0.004830268211662769, 0.0035455955658107996, 0.0007549467263743281, 0.000541276705916971, 0.0031480982434004545, 0.001557780895382166, 0.0010192448971793056, 0.0018504501786082983, 0.002619183622300625, 0.1016833484172821, 0.03818811476230621, 0.06928347051143646, 0.0412699431180954, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.26683223247528076, 0.0017643374158069491, 0.02531762421131134, 0.047485485672950745, 0.0005023732082918286, 0.0011795219033956528, 0.002227108459919691, 0.0028741960413753986, 0.005215880926698446, 0.001946018310263753, 3.592624852899462e-05, 0.001338632428087294, 0.0025214410852640867, 0.07723907381296158, 0.012742026709020138, 0.25196006894111633, 0.052669085562229156, 0.020061112940311432, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3006725609302521, 0.0014043879928067327, 0.009936605580151081, 0.037061650305986404, 0.0005129858036525548, 5.274279828881845e-05, 0.0006371501949615777, 0.00048446646542288363, 0.015043019317090511, 0.0003374778898432851, 0.0015171451959758997, 0.001911269617266953, 0.0014702629996463656, 0.015123972669243813, 0.0006335150101222098, 0.0006853189552202821, 0.0006114236894063652, 0.013829384930431843, 0.010252222418785095, NaN, NaN, NaN, NaN, NaN, NaN], [0.11150761693716049, 0.0006332705961540341, 0.0012255925685167313, 0.0022868558298796415, 0.0007688697660341859, 0.00046408100752159953, 0.0006869957433082163, 0.0021696356125175953, 0.0003113164857495576, 0.0013619231758639216, 0.004312699660658836, 0.1263500303030014, 0.0001710234791971743, 0.0024227115791291, 0.0006429344066418707, 0.008991677314043045, 0.01230061985552311, 0.025017380714416504, 0.33947470784187317, 0.0032216052059084177, NaN, NaN, NaN, NaN, NaN], [0.31111404299736023, 0.0035644923336803913, 0.0013678895775228739, 0.0016790243098512292, 0.0035299588926136494, 0.004438228905200958, 0.004504224751144648, 0.0015486004995182157, 0.006104794796556234, 0.009403211995959282, 0.00038756802678108215, 0.001732571516185999, 0.00042684219079092145, 0.00029873420135118067, 0.02043243870139122, 0.02443091571331024, 0.011036018840968609, 0.0030384601559489965, 0.007405058480799198, 0.004648045636713505, 0.010011163540184498, NaN, NaN, NaN, NaN], [0.16896948218345642, 0.0033956619445234537, 0.009647470898926258, 0.0011160745052620769, 0.30864211916923523, 0.0008666384965181351, 0.0001862353819888085, 0.0007671809289604425, 0.0006719603552483022, 0.002030742121860385, 0.00038655498065054417, 0.0009093419066630304, 0.0015865613240748644, 0.007534818258136511, 0.009185722097754478, 0.00011195908882655203, 0.003075815038755536, 0.000886340974830091, 0.0034873690456151962, 0.021776562556624413, 0.11334169656038284, 0.0832705944776535, NaN, NaN, NaN], [0.006588279269635677, 7.165617716964334e-06, 0.0005450915195979178, 1.0953889614029322e-05, 0.01959507167339325, 0.001590097788721323, 1.1096496564277913e-05, 7.439414184773341e-05, 9.72584675764665e-05, 0.00039174238918349147, 2.7912905352422968e-05, 4.964227991877124e-05, 7.256279786815867e-05, 0.00222678086720407, 0.04727102443575859, 0.0002576226834207773, 0.00020273383415769786, 7.391278631985188e-05, 0.00018598776659928262, 0.000617648009210825, 0.03195251524448395, 0.45461374521255493, 0.037591490894556046, NaN, NaN], [0.35417911410331726, 0.010997277684509754, 0.014662563800811768, 0.023722819983959198, 0.01071385107934475, 0.009427045471966267, 0.002653747797012329, 0.0011037624208256602, 0.005973298568278551, 0.0016420705942437053, 0.0009447215707041323, 0.001327668083831668, 0.0005524749867618084, 0.012130306102335453, 0.005379356909543276, 0.0037436189595609903, 0.0009285339619964361, 0.0002853046462405473, 0.0013114019529893994, 0.0012977200094610453, 0.08090774714946747, 0.034737478941679, 0.058711227029561996, 0.0672648623585701, NaN], [0.18188641965389252, 0.00040442554745823145, 0.0015771333128213882, 0.005189571529626846, 8.387575689994264e-06, 0.0001226859458256513, 0.0011242604814469814, 0.0013583728577941656, 0.0030172227416187525, 0.00029841059586033225, 1.2829146726289764e-05, 0.001467264024540782, 0.001090237987227738, 0.002914785873144865, 0.0006871690275147557, 0.002592542441561818, 0.00021328746515791863, 6.871169898658991e-05, 0.002350796014070511, 0.0026233955286443233, 0.02620280720293522, 0.005966363474726677, 0.08270465582609177, 0.010547555983066559, 0.018362630158662796]]], [[[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13007116317749023, 0.035988736897706985, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17991511523723602, 0.05124381557106972, 0.013642107136547565, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16831281781196594, 0.043814778327941895, 0.0950295478105545, 0.07350433617830276, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13759823143482208, 0.14112484455108643, 0.20577600598335266, 0.13910864293575287, 0.034107428044080734, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11619941890239716, 0.038306448608636856, 0.06045802682638168, 0.03494013100862503, 0.374624639749527, 0.22046393156051636, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08332619816064835, 0.009484739042818546, 0.012810231186449528, 0.0027760458178818226, 0.3268325924873352, 0.26342087984085083, 0.17634892463684082, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.057563915848731995, 0.01992173306643963, 0.03713805601000786, 0.014863312244415283, 0.25726908445358276, 0.14832180738449097, 0.402090460062027, 0.06479739397764206, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21478669345378876, 0.15359601378440857, 0.26770198345184326, 0.12653663754463196, 0.09151764959096909, 0.07003500312566757, 0.19363711774349213, 0.014233908616006374, 0.023967349901795387, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2834857702255249, 0.07559704780578613, 0.07655511796474457, 0.16202391684055328, 0.08316012471914291, 0.11911017447710037, 0.0204884335398674, 0.011816238984465599, 0.13204774260520935, 0.039266277104616165, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.23006244003772736, 0.03933367133140564, 0.07187695801258087, 0.04476522281765938, 0.01073860377073288, 0.0032203071750700474, 0.00176758982706815, 0.018770985305309296, 0.12121162563562393, 0.18536020815372467, 0.01582610420882702, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18067117035388947, 0.009833509102463722, 0.03744787722826004, 0.016920698806643486, 0.05744745582342148, 0.04540643468499184, 0.008024180307984352, 0.012110988609492779, 0.09370782226324081, 0.08820194005966187, 0.06259123980998993, 0.025030089542269707, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11523616313934326, 0.03200709819793701, 0.050564926117658615, 0.010618647560477257, 0.09430865943431854, 0.018685024231672287, 0.022438397631049156, 0.017720744013786316, 0.1592920571565628, 0.21717989444732666, 0.2463550567626953, 0.2194516956806183, 0.0009421245777048171, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09747911244630814, 0.1645127683877945, 0.1875433474779129, 0.09478750824928284, 0.08721300214529037, 0.02294742316007614, 0.02039182186126709, 0.07351931929588318, 0.1815827339887619, 0.5564144849777222, 0.41975197196006775, 0.2698606848716736, 0.05650324374437332, 0.05821085348725319, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14833268523216248, 0.1209164559841156, 0.08990822732448578, 0.0656033307313919, 0.23720099031925201, 0.11782333254814148, 0.04633651673793793, 0.16808320581912994, 0.06126163899898529, 0.43528908491134644, 0.3754012882709503, 0.13757933676242828, 0.05596579611301422, 0.16984672844409943, 0.002737722359597683, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19258342683315277, 0.05838138237595558, 0.04652376100420952, 0.017318567261099815, 0.23482391238212585, 0.16333334147930145, 0.02100907638669014, 0.048424359411001205, 0.06841404736042023, 0.3133482038974762, 0.07921069860458374, 0.021035969257354736, 0.03291412815451622, 0.18175286054611206, 0.1566929817199707, 0.053215935826301575, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17641158401966095, 0.15294750034809113, 0.15352487564086914, 0.10843643546104431, 0.08260629326105118, 0.016529222950339317, 0.012650150805711746, 0.07893627882003784, 0.1388573795557022, 0.19094663858413696, 0.03751035034656525, 0.05650494620203972, 0.2426995038986206, 0.16961677372455597, 0.07263431698083878, 0.152814581990242, 0.018521834164857864, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.25574439764022827, 0.04364950954914093, 0.05707173049449921, 0.02453112043440342, 0.016254547983407974, 0.0026636396069079638, 0.0035282839089632034, 0.015699811279773712, 0.03404982015490532, 0.04375504329800606, 0.001423283712938428, 0.05359426140785217, 0.1740386039018631, 0.10691730678081512, 0.03620539605617523, 0.04950953647494316, 0.022295303642749786, 0.025807255879044533, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.216966450214386, 0.016096990555524826, 0.08351551741361618, 0.02645382098853588, 0.05811392888426781, 0.04091750830411911, 0.014506897889077663, 0.015038754791021347, 0.07221462577581406, 0.08585365861654282, 0.059816163033246994, 0.04502185434103012, 0.00397779606282711, 0.041175276041030884, 0.04448581859469414, 0.10983181744813919, 0.01911303587257862, 0.07987141609191895, 0.062483180314302444, NaN, NaN, NaN, NaN, NaN, NaN], [0.11257521063089371, 0.027663733810186386, 0.023284420371055603, 0.0038690094370394945, 0.053685132414102554, 0.008445030078291893, 0.014706910587847233, 0.009755544364452362, 0.06406830251216888, 0.10475295782089233, 0.08554040640592575, 0.16072620451450348, 0.00029980239924043417, 0.03509804978966713, 0.03031017631292343, 0.04435117170214653, 0.06420817226171494, 0.2780051827430725, 0.2271702140569687, 0.0013584558619186282, NaN, NaN, NaN, NaN, NaN], [0.10895614326000214, 0.15509657561779022, 0.19682957231998444, 0.07681374996900558, 0.06229116767644882, 0.016663551330566406, 0.015513443388044834, 0.04232686012983322, 0.0986364334821701, 0.35070890188217163, 0.19941051304340363, 0.163076713681221, 0.026361489668488503, 0.018140846863389015, 0.016411108896136284, 0.03203867748379707, 0.053678009659051895, 0.19773079454898834, 0.3572796881198883, 0.059515852481126785, 0.04298213869333267, NaN, NaN, NaN, NaN], [0.15568822622299194, 0.11876019835472107, 0.09203660488128662, 0.059780094772577286, 0.24089980125427246, 0.06525673717260361, 0.029934749007225037, 0.11168782413005829, 0.03211824223399162, 0.30118685960769653, 0.22822384536266327, 0.08190999180078506, 0.018841415643692017, 0.1366286426782608, 0.0017427116399630904, 0.02601366490125656, 0.09386949241161346, 0.19522085785865784, 0.1546826809644699, 0.06491755694150925, 0.19679579138755798, 0.0025137634947896004, NaN, NaN, NaN], [0.26271528005599976, 0.07045364379882812, 0.0520184300839901, 0.023400958627462387, 0.11433269083499908, 0.07895253598690033, 0.012276851572096348, 0.023823700845241547, 0.04200353845953941, 0.16687022149562836, 0.05654531344771385, 0.038080912083387375, 0.012698299251496792, 0.10473722219467163, 0.0643644630908966, 0.015445034019649029, 0.014234953559935093, 0.06144930049777031, 0.05821693688631058, 0.0568128302693367, 0.1767931431531906, 0.1402994990348816, 0.07714083790779114, NaN, NaN], [0.1969611942768097, 0.16093717515468597, 0.1609625220298767, 0.11138524115085602, 0.026131147518754005, 0.00619129091501236, 0.005407778546214104, 0.04104578495025635, 0.06517186760902405, 0.06833471357822418, 0.020616043359041214, 0.03467438742518425, 0.095084547996521, 0.06247802451252937, 0.022057469934225082, 0.06569864600896835, 0.0052108620293438435, 0.03032413311302662, 0.0838729590177536, 0.3427644968032837, 0.19215865433216095, 0.08116735517978668, 0.14785417914390564, 0.015012684278190136, NaN], [0.1272672563791275, 0.008308093063533306, 0.030398543924093246, 0.02721896767616272, 0.016537277027964592, 0.021588556468486786, 0.002818688517436385, 0.010970782488584518, 0.01434051152318716, 0.012293173000216484, 0.04184769093990326, 0.03683166950941086, 0.023453323170542717, 0.020430248230695724, 0.03333409130573273, 0.068024642765522, 0.02648366242647171, 0.1640448421239853, 0.109919473528862, 0.1576652079820633, 0.14138163626194, 0.16884489357471466, 0.30372628569602966, 0.2283693552017212, 0.17022481560707092]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12451039254665375, 0.1335938721895218, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18396444618701935, 0.017508728429675102, 0.02471269853413105, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18453162908554077, 0.038695670664310455, 0.04155581444501877, 0.05072518810629845, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14826133847236633, 0.04252630099654198, 0.08689215034246445, 0.08308856934309006, 0.015247097238898277, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1348571479320526, 0.07033194601535797, 0.10030655562877655, 0.13752251863479614, 0.030713800340890884, 0.1331333965063095, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20671042799949646, 0.05809834972023964, 0.1630101054906845, 0.06033356115221977, 0.07501133531332016, 0.017328333109617233, 0.028450097888708115, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15813153982162476, 0.14090144634246826, 0.26030233502388, 0.10773709416389465, 0.16133210062980652, 0.04816069453954697, 0.01304988656193018, 0.13335363566875458, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3033713400363922, 0.22469042241573334, 0.4264413118362427, 0.3422197103500366, 0.14910078048706055, 0.06983038783073425, 0.023690486326813698, 0.010566752403974533, 0.05880258232355118, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.25368839502334595, 0.33459752798080444, 0.3829180896282196, 0.2782860994338989, 0.2427205741405487, 0.08768615871667862, 0.031752120703458786, 0.02143564634025097, 0.03798065707087517, 0.07379034906625748, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14200474321842194, 0.2391311228275299, 0.18728229403495789, 0.11236919462680817, 0.20923744142055511, 0.13365258276462555, 0.052715059369802475, 0.134474515914917, 0.14480768144130707, 0.06683899462223053, 0.104619100689888, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09595079720020294, 0.2752297520637512, 0.21842314302921295, 0.13660691678524017, 0.35477691888809204, 0.37130749225616455, 0.20556269586086273, 0.35276445746421814, 0.31008264422416687, 0.11074709892272949, 0.19841141998767853, 0.07199764251708984, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15323933959007263, 0.4611065983772278, 0.07869336754083633, 0.03600241616368294, 0.47375282645225525, 0.7350273132324219, 0.297486275434494, 0.6052883863449097, 0.4953201115131378, 0.144621342420578, 0.3493393063545227, 0.04881289228796959, 0.10520726442337036, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12003841996192932, 0.2704387903213501, 0.20063650608062744, 0.23778890073299408, 0.36254584789276123, 0.5319709777832031, 0.4483972191810608, 0.15058189630508423, 0.11134153604507446, 0.09426670521497726, 0.21241672337055206, 0.10488338023424149, 0.049764484167099, 0.15823495388031006, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15233570337295532, 0.21891875565052032, 0.13215333223342896, 0.2837490439414978, 0.08042775094509125, 0.43866410851478577, 0.2773631513118744, 0.12773916125297546, 0.3155127763748169, 0.07932031899690628, 0.1219707503914833, 0.11212008446455002, 0.1944955438375473, 0.07170752435922623, 0.004313962999731302, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2607015371322632, 0.3645761013031006, 0.37828943133354187, 0.3385462462902069, 0.2960833013057709, 0.5598280429840088, 0.544554591178894, 0.47054967284202576, 0.3477361798286438, 0.13701467216014862, 0.14822737872600555, 0.030188634991645813, 0.05528556555509567, 0.058441486209630966, 0.03410256654024124, 0.17273126542568207, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1872977614402771, 0.29805198311805725, 0.5206820368766785, 0.33024296164512634, 0.6395015716552734, 0.7210167050361633, 0.353913813829422, 0.406305193901062, 0.5096184015274048, 0.26257815957069397, 0.07301049679517746, 0.03464117646217346, 0.0787002444267273, 0.10916904360055923, 0.3557807505130768, 0.08364078402519226, 0.08538500964641571, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13269101083278656, 0.2835436165332794, 0.47488275170326233, 0.24851854145526886, 0.694171130657196, 0.6760384440422058, 0.2759343385696411, 0.29058361053466797, 0.7136873602867126, 0.20711864531040192, 0.04295802861452103, 0.07691331952810287, 0.11943909525871277, 0.1323360651731491, 0.20847304165363312, 0.05967296287417412, 0.12062160670757294, 0.09502720832824707, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.058743223547935486, 0.276242733001709, 0.29826071858406067, 0.20218241214752197, 0.4631478488445282, 0.48415693640708923, 0.2865871787071228, 0.3694051504135132, 0.4054408073425293, 0.19627220928668976, 0.2907293438911438, 0.09057808667421341, 0.11348091810941696, 0.21781016886234283, 0.38082650303840637, 0.3570795953273773, 0.22612451016902924, 0.09323522448539734, 0.03618632256984711, NaN, NaN, NaN, NaN, NaN, NaN], [0.07694489508867264, 0.41184449195861816, 0.038429711014032364, 0.018668875098228455, 0.5307568907737732, 0.7476497888565063, 0.4137455224990845, 0.6917499303817749, 0.6703397035598755, 0.3623183071613312, 0.579600989818573, 0.12613137066364288, 0.20100651681423187, 0.40998968482017517, 0.46115902066230774, 0.575211763381958, 0.35096046328544617, 0.163946270942688, 0.021770814433693886, 0.09986086189746857, NaN, NaN, NaN, NaN, NaN], [0.0834016501903534, 0.33346420526504517, 0.238715261220932, 0.28079062700271606, 0.5652539134025574, 0.6881173849105835, 0.5534363985061646, 0.22000034153461456, 0.1979052871465683, 0.3127084970474243, 0.4257359504699707, 0.18722867965698242, 0.1397658735513687, 0.3447277843952179, 0.13513657450675964, 0.31811001896858215, 0.32070791721343994, 0.12404847145080566, 0.05496959760785103, 0.04215753450989723, 0.16014836728572845, NaN, NaN, NaN, NaN], [0.13260646164417267, 0.29362690448760986, 0.18431688845157623, 0.38109344244003296, 0.20342527329921722, 0.5946046113967896, 0.4558189809322357, 0.26072001457214355, 0.5455912351608276, 0.2635512351989746, 0.31394094228744507, 0.23975242674350739, 0.36583349108695984, 0.2753828167915344, 0.01127256266772747, 0.41475725173950195, 0.29836422204971313, 0.2503683567047119, 0.10983213782310486, 0.21767295897006989, 0.0692884549498558, 0.003035380970686674, NaN, NaN, NaN], [0.2068602293729782, 0.4467880427837372, 0.4564751386642456, 0.4485791325569153, 0.45999279618263245, 0.6740500330924988, 0.7906107902526855, 0.6832103133201599, 0.5420533418655396, 0.4096798300743103, 0.3950984477996826, 0.13646338880062103, 0.10497336834669113, 0.17230592668056488, 0.07012390345335007, 0.27583980560302734, 0.3079235553741455, 0.1555996537208557, 0.038740403950214386, 0.05588690564036369, 0.03859011456370354, 0.02352789230644703, 0.12950412929058075, NaN, NaN], [0.16561447083950043, 0.3958832919597626, 0.5531814098358154, 0.4040684700012207, 0.7809365391731262, 0.8175305128097534, 0.5712264180183411, 0.6113651394844055, 0.6668697595596313, 0.4850655198097229, 0.18787693977355957, 0.08608534932136536, 0.19115354120731354, 0.2498423308134079, 0.6246696710586548, 0.31422460079193115, 0.373276948928833, 0.049351077526807785, 0.046956032514572144, 0.08076699078083038, 0.09392194449901581, 0.3349837362766266, 0.062239501625299454, 0.10001940280199051, NaN], [0.06568613648414612, 0.36780038475990295, 0.6246912479400635, 0.7116879820823669, 0.754679262638092, 0.7714072465896606, 0.7616819739341736, 0.5837911367416382, 0.9111838936805725, 0.8262851238250732, 0.6737059354782104, 0.5146453380584717, 0.7674095630645752, 0.7359525561332703, 0.5679676532745361, 0.7213301062583923, 0.6703079342842102, 0.5636342167854309, 0.38883939385414124, 0.5560528635978699, 0.518941342830658, 0.3739706873893738, 0.32013192772865295, 0.3743935525417328, 0.3977084755897522]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1305680274963379, 0.02726716920733452, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.002169837476685643, 0.0032534021884202957, 0.5694547891616821, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1568225622177124, 0.12336109578609467, 0.028200775384902954, 0.03890102356672287, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.008017625659704208, 0.013223886489868164, 0.04581261798739433, 0.017950134351849556, 0.8790656328201294, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08130903542041779, 0.2643316090106964, 0.5756329894065857, 0.29882851243019104, 0.31516125798225403, 0.09644471108913422, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20484277606010437, 0.3443664610385895, 0.0019387316424399614, 0.017399819567799568, 0.0004214652581140399, 0.00013534165918827057, 0.01563790813088417, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1571786254644394, 0.5643889307975769, 0.13441002368927002, 0.09036820381879807, 0.02947377972304821, 0.015878956764936447, 0.022048691287636757, 0.14189693331718445, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.005826869048178196, 0.13292454183101654, 0.00521356426179409, 0.005004087463021278, 0.10703893005847931, 0.26877719163894653, 0.1785666048526764, 0.23197543621063232, 0.007970587350428104, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03136341646313667, 0.08873608708381653, 0.009185479953885078, 0.03043411858379841, 0.3010490834712982, 0.36070317029953003, 0.178965762257576, 0.21872122585773468, 0.005464768502861261, 0.06020791083574295, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07854610681533813, 0.03772095590829849, 0.016643106937408447, 0.02832828275859356, 0.0785825327038765, 0.09336084127426147, 0.24177083373069763, 0.2718014717102051, 0.12932275235652924, 0.08437053114175797, 0.24188947677612305, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17239268124103546, 0.029533302411437035, 0.030515655875205994, 0.026403654366731644, 0.05037287250161171, 0.13986584544181824, 0.11416076123714447, 0.08228978514671326, 0.26975753903388977, 0.020502708852291107, 0.030797043815255165, 0.006723156664520502, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.35662412643432617, 0.005917226430028677, 0.00044432797585614026, 0.00022813511895947158, 0.0073361690156161785, 0.0027237480971962214, 0.007987208664417267, 0.021625559777021408, 0.010472757741808891, 0.0008755659800954163, 0.012584702111780643, 0.000526397256180644, 0.01033733133226633, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.189227893948555, 0.01606086827814579, 0.0030457540415227413, 0.005861388053745031, 0.04963670298457146, 0.004091562703251839, 0.01225967425853014, 0.037419673055410385, 0.01020084973424673, 0.003108290024101734, 0.01512740459293127, 0.006679146084934473, 0.014098022133111954, 0.03816642239689827, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00965302623808384, 0.0035168000031262636, 0.03902876377105713, 0.0158648993819952, 0.32648226618766785, 0.0038036927580833435, 0.002248003613203764, 0.002372291637584567, 0.014672092162072659, 0.007728067692369223, 0.022481968626379967, 0.028911879286170006, 0.044244468212127686, 0.021532919257879257, 0.6417658925056458, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.037641312927007675, 0.005557402968406677, 0.0006393054500222206, 0.006437606643885374, 0.007460788358002901, 0.0009530181414447725, 0.0016025539953261614, 0.0067516821436584, 0.02322007343173027, 0.018459537997841835, 0.011051125824451447, 0.006488891318440437, 0.04039585590362549, 0.18200218677520752, 0.0006002468289807439, 0.6243939995765686, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.01615065336227417, 0.01699231006205082, 0.00012957912986166775, 0.016060354188084602, 0.0006264564581215382, 0.0012908404460176826, 0.002684527076780796, 0.027531128376722336, 0.015566377900540829, 0.003692139405757189, 0.5753727555274963, 0.5145941376686096, 0.03750383481383324, 0.009545800276100636, 0.0034461882896721363, 0.005381980445235968, 0.00046628122800029814, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.021861553192138672, 0.01695878431200981, 0.0018149337265640497, 0.015764223411679268, 0.007719711866229773, 0.0034752548672258854, 0.007653116714209318, 0.03472340479493141, 0.038436826318502426, 0.014262136071920395, 0.8426622748374939, 0.36256304383277893, 0.21876515448093414, 0.019672129303216934, 0.020847154781222343, 0.00781619269400835, 0.005409067030996084, 0.16073459386825562, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18507197499275208, 0.027911728248000145, 0.014699580147862434, 0.025536103174090385, 0.014524195343255997, 0.045023027807474136, 0.031167738139629364, 0.07539253681898117, 0.22652071714401245, 0.011904416605830193, 0.08752688765525818, 0.03955431655049324, 0.2908211648464203, 0.03612781688570976, 0.00514488760381937, 0.017019467428326607, 0.07116629183292389, 0.03509910777211189, 0.02026083506643772, NaN, NaN, NaN, NaN, NaN, NaN], [0.40259334444999695, 0.005078054964542389, 0.00017122419376391917, 9.21270766411908e-05, 0.002624903805553913, 0.0009363252320326865, 0.00360113475471735, 0.01331485528498888, 0.008243494667112827, 0.0007176694343797863, 0.019634194672107697, 0.002027983544394374, 0.02349759265780449, 0.030203014612197876, 0.000993669149465859, 0.0008422310347668827, 0.013102295808494091, 0.025159381330013275, 0.0006507099606096745, 0.018182074651122093, NaN, NaN, NaN, NaN, NaN], [0.2579963207244873, 0.021157346665859222, 0.002921733073890209, 0.006211739499121904, 0.031850416213274, 0.0022005264181643724, 0.0070661455392837524, 0.036871425807476044, 0.012320333160459995, 0.005331193562597036, 0.033889420330524445, 0.020235266536474228, 0.07458563148975372, 0.1398555487394333, 0.008059950545430183, 0.0405682735145092, 0.03368399292230606, 0.012085597030818462, 0.010676471516489983, 0.03411625698208809, 0.08152885735034943, NaN, NaN, NaN, NaN], [0.005019576288759708, 0.001437423750758171, 0.014701779931783676, 0.005876661743968725, 0.15098156034946442, 0.001037455745972693, 0.0006782425916753709, 0.0010664333822205663, 0.006170186679810286, 0.004750464111566544, 0.015587885864078999, 0.020612932741642, 0.024904461577534676, 0.027292385697364807, 0.6522603631019592, 0.02780178189277649, 0.009980881586670876, 0.010863273404538631, 0.016993993893265724, 0.026612548157572746, 0.013426730409264565, 0.6643192768096924, NaN, NaN, NaN], [0.023952102288603783, 0.0025056565646082163, 0.0002975048264488578, 0.0031560298521071672, 0.002087814500555396, 0.00019765450269915164, 0.00028781042783521116, 0.0023521913681179285, 0.009429593570530415, 0.010675383731722832, 0.013774069957435131, 0.012372920289635658, 0.030660077929496765, 0.3810364305973053, 0.0006224916432984173, 0.6039706468582153, 0.2701583206653595, 0.012816790491342545, 0.005745226051658392, 0.052403513342142105, 0.18411211669445038, 0.00043697847286239266, 0.6234135627746582, NaN, NaN], [0.007988094352185726, 0.006256349850445986, 4.065780740347691e-05, 0.006692530121654272, 0.00010113247117260471, 0.0002641561150085181, 0.0006015493418090045, 0.009669815190136433, 0.00486318813636899, 0.0012557843001559377, 0.43231210112571716, 0.35852983593940735, 0.01959061808884144, 0.007567983586341143, 0.0019125458784401417, 0.00857639778405428, 0.0005027590086683631, 0.41286540031433105, 0.4292365312576294, 0.01753525249660015, 0.005813234485685825, 0.00216498039662838, 0.003382693277671933, 0.00027526391204446554, NaN], [0.1387476772069931, 0.027318276464939117, 0.00785337295383215, 0.019197843968868256, 0.013794281519949436, 0.020801816135644913, 0.013009469024837017, 0.07068510353565216, 0.020734209567308426, 0.024748992174863815, 0.04673967882990837, 0.025586238130927086, 0.01648368127644062, 0.06557000428438187, 0.022920427843928337, 0.013843921944499016, 0.04100487753748894, 0.0375630147755146, 0.023956134915351868, 0.018727701157331467, 0.05957711860537529, 0.020177751779556274, 0.007389482576400042, 0.027843382209539413, 0.025224220007658005]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1319446712732315, 0.003103907685726881, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.004627853631973267, 0.8189921975135803, 0.006355744786560535, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0004822930786758661, 0.5574855208396912, 0.0058120423927903175, 0.014268792234361172, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15055440366268158, 0.0014966451562941074, 0.1733904629945755, 0.05038055405020714, 0.0057296124286949635, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1304439753293991, 0.00022060537594370544, 0.03428095951676369, 0.0157721396535635, 0.20856629312038422, 0.2746620774269104, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.017820989713072777, 1.0936159014818259e-05, 0.0006241680239327252, 4.3406893382780254e-05, 0.2565733790397644, 0.5255003571510315, 0.040596142411231995, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2143511176109314, 3.818454570136964e-05, 0.0006476931739598513, 0.00012842394062317908, 0.007853559218347073, 0.008102592080831528, 0.0005345920799300075, 0.00793861411511898, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00014670012751594186, 7.536429620813578e-06, 0.0001294321846216917, 0.00024457855033688247, 0.00022483686916530132, 0.001284220488741994, 0.0014163334853947163, 0.5552030801773071, 0.006061996798962355, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09223808348178864, 0.004348577931523323, 0.013163902796804905, 0.018216131255030632, 0.035016678273677826, 0.11075899004936218, 0.1728493720293045, 0.19621391594409943, 0.029301786795258522, 0.46166056394577026, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11309938877820969, 0.004489036742597818, 0.0485633909702301, 0.021462395787239075, 0.4192940890789032, 0.26214849948883057, 0.22032421827316284, 0.0067114257253706455, 0.010406548157334328, 0.11692964285612106, 0.23004111647605896, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14281870424747467, 0.000545236689504236, 0.003893920686095953, 0.0005153689999133348, 0.01790653169155121, 0.004868220537900925, 0.0031487985979765654, 0.0011714915744960308, 0.0043698386289179325, 0.020373020321130753, 0.02358497679233551, 0.2682037353515625, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09794370085000992, 0.0018320194212719798, 0.000285644200630486, 3.260145604144782e-05, 0.00041393720312044024, 0.0043053096160292625, 0.002047628629952669, 0.0003047001373488456, 0.002447759034112096, 0.0016152235912159085, 0.024524936452507973, 0.29461416602134705, 0.014563476666808128, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13817672431468964, 0.0034516772720962763, 0.002911344636231661, 0.0003800573176704347, 0.001462712767533958, 0.001961951842531562, 0.0040230052545666695, 0.0023086154833436012, 0.002483226591721177, 0.028553131967782974, 0.014239847660064697, 0.18359807133674622, 0.09542248398065567, 0.2067933827638626, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14011409878730774, 0.01466476172208786, 0.09487155824899673, 0.03769487887620926, 0.062972791492939, 0.003495296463370323, 0.0004466120735742152, 0.0044098952785134315, 0.056031279265880585, 0.12585759162902832, 0.04736572876572609, 0.02727479301393032, 0.06542934477329254, 0.563940703868866, 0.024195805191993713, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05395817384123802, 6.747527368133888e-05, 0.0018676340114325285, 0.0002809480356518179, 0.03275269269943237, 0.005758063402026892, 9.199039777740836e-05, 0.00011598093260545284, 0.0015754709020256996, 0.026104740798473358, 0.009686414152383804, 0.001081737456843257, 0.0017741151386871934, 0.49180474877357483, 0.007121484261006117, 0.013531914912164211, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03839295729994774, 0.0002068357716780156, 0.006204192526638508, 0.0054313126020133495, 0.011207946576178074, 0.0013116636546328664, 0.008276019245386124, 0.002269806107506156, 0.004080863669514656, 0.01488969475030899, 0.0006726597202941775, 0.009391524828970432, 0.039596475660800934, 0.19840312004089355, 0.043704546988010406, 0.31202515959739685, 0.23529505729675293, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07469534128904343, 0.001304430770687759, 0.0239309910684824, 0.008060658350586891, 0.021029237657785416, 0.015191669575870037, 0.006979105528444052, 0.0016427322989329696, 0.002132130553945899, 0.015241370536386967, 0.0018563566263765097, 0.035101406276226044, 0.06515936553478241, 0.27313047647476196, 0.10352547466754913, 0.2570805549621582, 0.45083746314048767, 0.1295340657234192, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19253067672252655, 0.0008209676598198712, 0.004669400863349438, 0.00047802351764403284, 0.013135433197021484, 0.0034620855003595352, 0.0016354827675968409, 0.0008273401763290167, 0.0018895546672865748, 0.009773027151823044, 0.006215384230017662, 0.2356690764427185, 0.01036232803016901, 0.06144833192229271, 0.008870624005794525, 0.024212215095758438, 0.008509873412549496, 0.01347219105809927, 0.35532569885253906, NaN, NaN, NaN, NaN, NaN, NaN], [0.10910779982805252, 0.002221200615167618, 0.0001436042075511068, 1.1848528629343491e-05, 0.0001887700636871159, 0.0020721519831568003, 0.0009632316650822759, 0.00014056939107831568, 0.0007320817094296217, 0.0006829273188486695, 0.007395991589874029, 0.2889891564846039, 0.007074101362377405, 0.0002627878566272557, 0.004363438580185175, 0.0018575063440948725, 0.00557676050812006, 0.012322820723056793, 0.31134024262428284, 0.027276715263724327, NaN, NaN, NaN, NaN, NaN], [0.18170765042304993, 0.003209297079592943, 0.0023912524338811636, 0.00020479358499869704, 0.0009326079743914306, 0.0013757160631939769, 0.0021110770758241415, 0.0008730489062145352, 0.000792569131590426, 0.01825624145567417, 0.0059272306971251965, 0.11984144151210785, 0.05654650926589966, 0.08423373848199844, 0.024963613599538803, 0.027966396883130074, 0.1777324080467224, 0.005578523967415094, 0.14623191952705383, 0.11331525444984436, 0.2157108038663864, NaN, NaN, NaN, NaN], [0.1515214741230011, 0.008395697921514511, 0.0657893642783165, 0.019086696207523346, 0.05097401514649391, 0.0016111076110973954, 0.00021851839846931398, 0.002003778237849474, 0.01669292151927948, 0.06321260333061218, 0.015100682154297829, 0.010209205560386181, 0.015906400978565216, 0.30131736397743225, 0.012282183393836021, 0.09666845202445984, 0.00808996893465519, 0.03798958286643028, 0.013879657723009586, 0.047733187675476074, 0.5371345281600952, 0.020763304084539413, NaN, NaN, NaN], [0.07945924997329712, 4.7485355025855824e-05, 0.0020416006445884705, 0.00022757358965463936, 0.013386114500463009, 0.001981395063921809, 3.6917605029884726e-05, 2.620528539409861e-05, 0.0003202208608854562, 0.009042860940098763, 0.0030785591807216406, 0.0011855574557557702, 0.0005728560499846935, 0.20002734661102295, 0.00213914574123919, 0.002927121240645647, 0.004968173801898956, 0.0065933396108448505, 0.002585601294413209, 0.002817549044266343, 0.547335147857666, 0.006171087268739939, 0.018697692081332207, NaN, NaN], [0.059381648898124695, 0.00026094831991940737, 0.007586375344544649, 0.006061093881726265, 0.0039266073144972324, 0.0004965912085026503, 0.003665223019197583, 0.0008195870905183256, 0.0014654117403551936, 0.0045553394593298435, 0.00032001128420233727, 0.004615657962858677, 0.017150992527604103, 0.07922492176294327, 0.012805018573999405, 0.1320599913597107, 0.09461667388677597, 0.003555287839844823, 0.019601207226514816, 0.047796737402677536, 0.29085052013397217, 0.04383813217282295, 0.32529252767562866, 0.24933147430419922, NaN], [0.13618361949920654, 0.0007103006355464458, 0.025071904063224792, 0.004419561009854078, 0.001962232170626521, 0.0023795748129487038, 0.002366183791309595, 0.0003890783409588039, 0.00022811641974840313, 0.0010611300822347403, 0.001608739490620792, 0.028126444667577744, 0.005591525696218014, 0.0024579197634011507, 0.004123267717659473, 0.0409882515668869, 0.010364435613155365, 0.010518459603190422, 0.09771004319190979, 0.037823982536792755, 0.019979961216449738, 0.018303534016013145, 0.22492042183876038, 0.09256016463041306, 0.005498841404914856]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11621169000864029, 0.2792567312717438, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16788142919540405, 0.08717074245214462, 0.024576181545853615, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14762163162231445, 0.09094145894050598, 0.023598572239279747, 0.2273045778274536, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10424397885799408, 0.7145561575889587, 0.21233327686786652, 0.5272893309593201, 0.04291817173361778, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11001076549291611, 0.4734446108341217, 0.06134912371635437, 0.2925608456134796, 0.02150837518274784, 0.19962187111377716, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17212024331092834, 0.1419786959886551, 0.05631781369447708, 0.2185172289609909, 0.002532752463594079, 0.0032626313623040915, 0.18381445109844208, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09107878059148788, 0.12160263955593109, 0.2150201052427292, 0.3705081045627594, 0.07164584845304489, 0.05021890252828598, 0.14392021298408508, 0.39638784527778625, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2121918499469757, 0.20806513726711273, 0.15205760300159454, 0.38131871819496155, 0.1009124368429184, 0.09936784207820892, 0.07077471911907196, 0.05006752535700798, 0.14871110022068024, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21685828268527985, 0.23333710432052612, 0.06609098613262177, 0.12803798913955688, 0.1004808098077774, 0.025170300155878067, 0.04069148004055023, 0.10828333348035812, 0.10351972281932831, 0.29450517892837524, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05205162987112999, 0.22306090593338013, 0.049221184104681015, 0.061203524470329285, 0.09776578843593597, 0.06183243915438652, 0.17444021999835968, 0.321644127368927, 0.054029058665037155, 0.2629997134208679, 0.2757931053638458, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05800137668848038, 0.32540804147720337, 0.13333332538604736, 0.05756821855902672, 0.12640602886676788, 0.11846329271793365, 0.2918737828731537, 0.3632459342479706, 0.18816226720809937, 0.6433262228965759, 0.3291742205619812, 0.12170911580324173, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11078674346208572, 0.40781712532043457, 0.06261185556650162, 0.05779192969202995, 0.18194560706615448, 0.1120922714471817, 0.5645142793655396, 0.33037880063056946, 0.18058234453201294, 0.6155731678009033, 0.21430827677249908, 0.044265877455472946, 0.20548948645591736, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08288691937923431, 0.2962968051433563, 0.2819015085697174, 0.19574381411075592, 0.1136796846985817, 0.07755676656961441, 0.20596812665462494, 0.3330870270729065, 0.21944326162338257, 0.22804425656795502, 0.1688224822282791, 0.2872299253940582, 0.13759873807430267, 0.09907422959804535, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11118441820144653, 0.6110438108444214, 0.6292654871940613, 0.5805363655090332, 0.22765980660915375, 0.4274957776069641, 0.6573506593704224, 0.6816673278808594, 0.5361799597740173, 0.320940226316452, 0.3845328688621521, 0.6242536306381226, 0.41633498668670654, 0.12922972440719604, 0.01991792768239975, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10675505548715591, 0.1912444829940796, 0.23975566029548645, 0.32351911067962646, 0.046362437307834625, 0.08004549145698547, 0.3363644778728485, 0.2706483006477356, 0.26792168617248535, 0.2952979505062103, 0.4496033787727356, 0.1126319095492363, 0.5116660594940186, 0.015820369124412537, 0.030236991122364998, 0.03603934869170189, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2233639359474182, 0.0911012589931488, 0.12918633222579956, 0.17958812415599823, 0.037158817052841187, 0.06043876335024834, 0.43303725123405457, 0.3349981904029846, 0.09061599522829056, 0.23225362598896027, 0.1514965295791626, 0.09056703746318817, 0.2480165809392929, 0.056160230189561844, 0.015552842989563942, 0.007365798112004995, 0.17054231464862823, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09585364907979965, 0.22669152915477753, 0.08040254563093185, 0.0638674795627594, 0.15364862978458405, 0.13237975537776947, 0.3887532651424408, 0.5357696413993835, 0.07155110687017441, 0.4139500856399536, 0.05426981300115585, 0.1238613948225975, 0.07816720753908157, 0.14353296160697937, 0.021915707737207413, 0.02897939831018448, 0.22262324392795563, 0.4835837185382843, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05190133675932884, 0.3522363007068634, 0.14802464842796326, 0.07656959444284439, 0.12417534738779068, 0.17628712952136993, 0.33604755997657776, 0.38481405377388, 0.20552395284175873, 0.5797679424285889, 0.3262830972671509, 0.19466114044189453, 0.045280374586582184, 0.2712458372116089, 0.041196610778570175, 0.08666794002056122, 0.3327068090438843, 0.1922111064195633, 0.10969121754169464, NaN, NaN, NaN, NaN, NaN, NaN], [0.10818891227245331, 0.3937702178955078, 0.030490810051560402, 0.030189264565706253, 0.11243001371622086, 0.07142115384340286, 0.3648340702056885, 0.2467786818742752, 0.13009557127952576, 0.5037410855293274, 0.18716548383235931, 0.08825942128896713, 0.23451530933380127, 0.24434491991996765, 0.03496113047003746, 0.04431905224919319, 0.3934983015060425, 0.31427451968193054, 0.05462265387177467, 0.2524711489677429, NaN, NaN, NaN, NaN, NaN], [0.06088699772953987, 0.23725801706314087, 0.2046121060848236, 0.14171433448791504, 0.06688592582941055, 0.06064169481396675, 0.14286598563194275, 0.21723276376724243, 0.13491223752498627, 0.2083195000886917, 0.15285742282867432, 0.34066644310951233, 0.18166381120681763, 0.10532425343990326, 0.06318715214729309, 0.052211396396160126, 0.20970472693443298, 0.20715771615505219, 0.28281068801879883, 0.13935938477516174, 0.11923542618751526, NaN, NaN, NaN, NaN], [0.09884612262248993, 0.5530695915222168, 0.6301063299179077, 0.5187459588050842, 0.28427499532699585, 0.33059176802635193, 0.49595603346824646, 0.6107674241065979, 0.387560099363327, 0.3283739984035492, 0.3905918300151825, 0.5949583053588867, 0.2912430167198181, 0.19163259863853455, 0.03091937117278576, 0.3911139667034149, 0.3233675956726074, 0.421701043844223, 0.6310504674911499, 0.4068542718887329, 0.13317596912384033, 0.02126597985625267, NaN, NaN, NaN], [0.07192745804786682, 0.09934075176715851, 0.15662430226802826, 0.18248029053211212, 0.021172231063246727, 0.037516966462135315, 0.12766626477241516, 0.09711621701717377, 0.09662153571844101, 0.1303528994321823, 0.3114719092845917, 0.1600099802017212, 0.265144020318985, 0.011710498481988907, 0.02471126988530159, 0.012725233100354671, 0.12533646821975708, 0.446529746055603, 0.11092787981033325, 0.45893827080726624, 0.011159577406942844, 0.028070949018001556, 0.024378135800361633, NaN, NaN], [0.21178482472896576, 0.0713806003332138, 0.12116114795207977, 0.16551871597766876, 0.025692136958241463, 0.03932836279273033, 0.255863755941391, 0.20887790620326996, 0.05500240623950958, 0.14075487852096558, 0.158308207988739, 0.10016348958015442, 0.22940821945667267, 0.06542190909385681, 0.016673747450113297, 0.011679067276418209, 0.21266934275627136, 0.27460965514183044, 0.08977667987346649, 0.1985965520143509, 0.05640871822834015, 0.014301197603344917, 0.004748867359012365, 0.1251523643732071, NaN], [0.11377177387475967, 0.4656391441822052, 0.26672884821891785, 0.20802536606788635, 0.1860857605934143, 0.16829806566238403, 0.19711202383041382, 0.3023360073566437, 0.035885076969861984, 0.11114621162414551, 0.21048156917095184, 0.27827921509742737, 0.11178875714540482, 0.13154125213623047, 0.3096882104873657, 0.09530708193778992, 0.2201821655035019, 0.1989239901304245, 0.27841058373451233, 0.15223632752895355, 0.2206900417804718, 0.34536775946617126, 0.09229245036840439, 0.24595825374126434, 0.2865155339241028]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13124778866767883, 0.015335792675614357, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19323189556598663, 0.005229663103818893, 0.005805561784654856, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06695510447025299, 0.08997365087270737, 0.32878753542900085, 0.35321861505508423, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1452476531267166, 0.07996584475040436, 0.2002653181552887, 0.13149262964725494, 0.005022347904741764, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1274433135986328, 0.13577045500278473, 0.16066212952136993, 0.1959238052368164, 0.04180024936795235, 0.06788772344589233, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14809708297252655, 0.29017606377601624, 0.22457490861415863, 0.17088554799556732, 0.041788797825574875, 0.013634788803756237, 0.02984887920320034, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21402230858802795, 0.012405444867908955, 0.0014808804262429476, 0.0009161182679235935, 0.0035427443217486143, 0.0017166208708658814, 0.001927618752233684, 0.015056394040584564, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10794443637132645, 0.13477572798728943, 0.046750620007514954, 0.03419584408402443, 0.30604344606399536, 0.11879221349954605, 0.08022946119308472, 0.11745522916316986, 0.21712547540664673, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06259628385305405, 0.21873348951339722, 0.248628169298172, 0.2344663441181183, 0.09133727103471756, 0.05752522125840187, 0.03945200890302658, 0.39403918385505676, 0.15040725469589233, 0.009099425747990608, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06400181353092194, 0.3208324611186981, 0.5040323138237, 0.6282902359962463, 0.04389061778783798, 0.08030739426612854, 0.10539824515581131, 0.1485716998577118, 0.08085520565509796, 0.13963551819324493, 0.0947280004620552, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0935494601726532, 0.3055664598941803, 0.46751275658607483, 0.6914730072021484, 0.12860655784606934, 0.15726737678050995, 0.2987912595272064, 0.1529359668493271, 0.062232255935668945, 0.041881486773490906, 0.03399288281798363, 0.026789270341396332, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.012478480115532875, 0.051689472049474716, 0.7194163799285889, 0.8485123515129089, 0.006671697832643986, 0.03636787086725235, 0.05433559790253639, 0.01463489979505539, 0.0011851346353068948, 0.0010049004340544343, 0.012586181983351707, 0.0039429632015526295, 0.0029262336902320385, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16095376014709473, 0.10161679983139038, 0.15561290085315704, 0.27214428782463074, 0.06339859217405319, 0.047669682651758194, 0.16775988042354584, 0.30333516001701355, 0.29585903882980347, 0.026492541655898094, 0.03390856087207794, 0.020966142416000366, 0.027538424357771873, 0.040642742067575455, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1701768934726715, 0.015393235720694065, 0.0020776872988790274, 0.011533004231750965, 0.013215321116149426, 0.004845780786126852, 0.011772604659199715, 0.006262979004532099, 0.00390799343585968, 0.007256041280925274, 0.0014780729543417692, 0.007152961101382971, 0.1450572907924652, 0.009833375923335552, 0.004788131918758154, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.27953270077705383, 0.3106633424758911, 0.3078516721725464, 0.2835734188556671, 0.23220741748809814, 0.10028243064880371, 0.059542566537857056, 0.10900203883647919, 0.24247398972511292, 0.19294817745685577, 0.04455278813838959, 0.032558612525463104, 0.2623904049396515, 0.04071282595396042, 0.07101175934076309, 0.01397540420293808, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15828359127044678, 0.26215362548828125, 0.1828027367591858, 0.3383132517337799, 0.14976613223552704, 0.17187725007534027, 0.16098640859127045, 0.10713529586791992, 0.2253616452217102, 0.27887699007987976, 0.0991593673825264, 0.1987481713294983, 0.2010713517665863, 0.24892166256904602, 0.09143882989883423, 0.028894133865833282, 0.0226773452013731, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08621957898139954, 0.39239373803138733, 0.32060059905052185, 0.6169360876083374, 0.04211895540356636, 0.07954877614974976, 0.28241875767707825, 0.1073535904288292, 0.10431969910860062, 0.28138864040374756, 0.05428503826260567, 0.29005417227745056, 0.2829020619392395, 0.1771886944770813, 0.12728992104530334, 0.029228007420897484, 0.09527892619371414, 0.030012397095561028, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10387677699327469, 0.28899070620536804, 0.34778735041618347, 0.5978891849517822, 0.08856049180030823, 0.11093756556510925, 0.2773001492023468, 0.1387036144733429, 0.05535874143242836, 0.040542375296354294, 0.057020239531993866, 0.08593740314245224, 0.3575255870819092, 0.1780063509941101, 0.03115975111722946, 0.05683879926800728, 0.20087137818336487, 0.022991398349404335, 0.024780578911304474, NaN, NaN, NaN, NaN, NaN, NaN], [0.027872784063220024, 0.11975038051605225, 0.8484699726104736, 0.9221431016921997, 0.010032964870333672, 0.05817321315407753, 0.14408904314041138, 0.03149182349443436, 0.0027255630120635033, 0.003546576714143157, 0.054592132568359375, 0.03846639767289162, 0.0179138146340847, 0.04004756733775139, 0.0025625908747315407, 0.006073353346437216, 0.017890095710754395, 0.006128084380179644, 0.0035659971181303263, 0.005842072889208794, NaN, NaN, NaN, NaN, NaN], [0.21095024049282074, 0.16082847118377686, 0.2551726996898651, 0.40046265721321106, 0.07841236889362335, 0.05558479577302933, 0.20925307273864746, 0.4381427764892578, 0.47918838262557983, 0.07096414268016815, 0.11106863617897034, 0.09138666838407516, 0.1393880993127823, 0.1506565660238266, 0.07743309438228607, 0.06943798065185547, 0.09801105409860611, 0.017720624804496765, 0.015859564766287804, 0.029157793149352074, 0.0392736941576004, NaN, NaN, NaN, NaN], [0.17935752868652344, 0.014263968914747238, 0.0022281131241470575, 0.011617614887654781, 0.022433524951338768, 0.0047986325807869434, 0.013686214573681355, 0.007696506567299366, 0.004939754959195852, 0.012488129548728466, 0.002878576284274459, 0.013457567431032658, 0.23303280770778656, 0.030022362247109413, 0.013181640766561031, 0.027029545977711678, 0.010247751139104366, 0.0006795030203647912, 0.0032072996255010366, 0.1104368045926094, 0.006663828622549772, 0.003364446572959423, NaN, NaN, NaN], [0.3113161623477936, 0.29550519585609436, 0.2834082841873169, 0.292662650346756, 0.1380799263715744, 0.055221766233444214, 0.0487985797226429, 0.10219268500804901, 0.25612032413482666, 0.2569950222969055, 0.10279092192649841, 0.16084249317646027, 0.5340818166732788, 0.10305190831422806, 0.16831228137016296, 0.03310799598693848, 0.10521702468395233, 0.008185362443327904, 0.02029210887849331, 0.2447529286146164, 0.0189062412828207, 0.051586367189884186, 0.011271311901509762, NaN, NaN], [0.21913117170333862, 0.2667233347892761, 0.15068072080612183, 0.2934513986110687, 0.11010763049125671, 0.11770202964544296, 0.1548316478729248, 0.10880382359027863, 0.19848009943962097, 0.2926469147205353, 0.17939361929893494, 0.38748762011528015, 0.38622626662254333, 0.4369211196899414, 0.14473943412303925, 0.11290202289819717, 0.11878126114606857, 0.013051117770373821, 0.18458649516105652, 0.15622372925281525, 0.14840805530548096, 0.06742489337921143, 0.01624887064099312, 0.028317920863628387, NaN], [0.13670727610588074, 0.11102687567472458, 0.008893890306353569, 0.008979070000350475, 0.01785319298505783, 0.008134939707815647, 0.02043774165213108, 0.030145585536956787, 0.014907605946063995, 0.021436721086502075, 0.020207075402140617, 0.10284662246704102, 0.06823904067277908, 0.04208305850625038, 0.03810393810272217, 0.04656955599784851, 0.025087369605898857, 0.005296032875776291, 0.07358870655298233, 0.057817310094833374, 0.033472564071416855, 0.02220221422612667, 0.01758744567632675, 0.012124869041144848, 0.052647966891527176]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1301431953907013, 0.0347244068980217, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19350707530975342, 0.0006586865638382733, 0.008110460825264454, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07742509245872498, 0.025898784399032593, 0.46813124418258667, 0.21566073596477509, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15508510172367096, 0.002848779782652855, 0.006727630738168955, 0.01290579792112112, 0.0019038956379517913, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1506490558385849, 0.0018329949816688895, 0.0011812039883807302, 0.010563074611127377, 0.0007367127691395581, 0.0007524989196099341, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0463392436504364, 0.0861721858382225, 0.5342088341712952, 0.5262086987495422, 0.252642959356308, 0.014757110737264156, 0.02778990939259529, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08082517981529236, 0.10121051222085953, 0.3481808602809906, 0.41374534368515015, 0.38359278440475464, 0.07890304177999496, 0.1096968874335289, 0.1685827672481537, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1433362513780594, 0.13670213520526886, 0.10138670355081558, 0.1093992069363594, 0.236768901348114, 0.09415888041257858, 0.011134332977235317, 0.019298367202281952, 0.5348934531211853, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.024931270629167557, 0.02871265634894371, 0.20136752724647522, 0.1457405984401703, 0.13753218948841095, 0.13171687722206116, 0.07031083852052689, 0.04771474376320839, 0.5403124690055847, 0.04482616111636162, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.026511939242482185, 0.12058579176664352, 0.09381356090307236, 0.09726550430059433, 0.13490843772888184, 0.36408668756484985, 0.19949088990688324, 0.09435784071683884, 0.45831772685050964, 0.1274537742137909, 0.014095090329647064, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12624163925647736, 0.03293433412909508, 0.07055910676717758, 0.06304988265037537, 0.23899653553962708, 0.15645378828048706, 0.07000429183244705, 0.02516351453959942, 0.06797400116920471, 0.07094329595565796, 0.1311238706111908, 0.21208471059799194, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1118171289563179, 0.015469676814973354, 0.08768722414970398, 0.046650953590869904, 0.23542486131191254, 0.09032069146633148, 0.05012429133057594, 0.004171812906861305, 0.15006321668624878, 0.017805932089686394, 0.049085501581430435, 0.035517167299985886, 0.6428134441375732, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09301143884658813, 0.13257478177547455, 0.1489255279302597, 0.18642880022525787, 0.318376362323761, 0.31357452273368835, 0.1382697969675064, 0.07457731664180756, 0.17392435669898987, 0.00920780934393406, 0.020603884011507034, 0.049020376056432724, 0.322329580783844, 0.3050764203071594, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17444664239883423, 0.0007958812057040632, 5.6854176364140585e-05, 0.0004179355164524168, 0.00013179269444663078, 0.00024977640714496374, 0.0001107741700252518, 7.639485556865111e-05, 0.0008396806661039591, 0.00030287212575785816, 0.00023763117496855557, 0.003834246192127466, 0.003433886216953397, 0.00015348535089287907, 0.00014843019016552716, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00841783918440342, 0.03505324944853783, 0.02469123899936676, 0.026689309626817703, 0.1500382125377655, 0.08861804753541946, 0.006530162878334522, 0.060150377452373505, 0.04669034481048584, 0.007807246409356594, 0.02131708152592182, 0.012364925816655159, 0.041818197816610336, 0.02841370552778244, 0.6981374621391296, 0.06836962699890137, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0009672276792116463, 0.0037913541309535503, 0.00524782482534647, 0.006044968497008085, 0.07807419449090958, 0.026950905099511147, 0.0024354930501431227, 0.005482541862875223, 0.013836389407515526, 0.002816400956362486, 0.0006559633184224367, 0.002845867071300745, 0.018497759476304054, 0.19704575836658478, 0.41393977403640747, 0.4024144113063812, 0.00308317132294178, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0023347423411905766, 0.018236415460705757, 0.011423468589782715, 0.014267664402723312, 0.06272618472576141, 0.09006785601377487, 0.023437032476067543, 0.008957883343100548, 0.03532397374510765, 0.006200278177857399, 0.0002018583327298984, 0.016960909590125084, 0.04933774098753929, 0.1362536996603012, 0.47770828008651733, 0.5670948624610901, 0.06992122530937195, 0.03068283386528492, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0730348452925682, 0.024321116507053375, 0.06646358221769333, 0.0630527138710022, 0.23201428353786469, 0.1378810703754425, 0.04738042131066322, 0.010255109518766403, 0.0316733755171299, 0.07226394861936569, 0.06345586478710175, 0.13366159796714783, 0.1651405692100525, 0.1875276118516922, 0.475235253572464, 0.34701114892959595, 0.106105737388134, 0.17074023187160492, 0.14835108816623688, NaN, NaN, NaN, NaN, NaN, NaN], [0.1317213624715805, 0.02603350207209587, 0.05892709270119667, 0.02498493157327175, 0.2902502715587616, 0.11121267080307007, 0.057563167065382004, 0.004654969088733196, 0.12363925576210022, 0.02343585342168808, 0.03682887554168701, 0.054189957678318024, 0.5043657422065735, 0.23388440907001495, 0.46154457330703735, 0.32561513781547546, 0.055846668779850006, 0.06476935744285583, 0.026345595717430115, 0.5623452067375183, NaN, NaN, NaN, NaN, NaN], [0.037178635597229004, 0.08259578794240952, 0.0920928493142128, 0.09107104688882828, 0.19359135627746582, 0.17535823583602905, 0.06819135695695877, 0.03716395050287247, 0.07458745688199997, 0.0064619481563568115, 0.009060872718691826, 0.02094256319105625, 0.1461041122674942, 0.11104261875152588, 0.6685899496078491, 0.4500047266483307, 0.029085516929626465, 0.03437849134206772, 0.03590574488043785, 0.20188003778457642, 0.23542997241020203, NaN, NaN, NaN, NaN], [0.18516498804092407, 0.0009336460498161614, 7.266629108926281e-05, 0.00041225351742468774, 0.00023152375069912523, 0.0002865330025088042, 0.00012637366307899356, 8.909442112781107e-05, 0.0006568549433723092, 0.0003727772564161569, 0.00021836791711393744, 0.0030449857003986835, 0.002062517451122403, 0.0001740154402796179, 0.00019746039470192045, 0.0010639599058777094, 3.738106170203537e-05, 0.00018948569777421653, 0.0017019548686221242, 0.0021623496431857347, 7.414143328787759e-05, 0.00010166682477574795, NaN, NaN, NaN], [0.014717604033648968, 0.07327108085155487, 0.049021750688552856, 0.04824157431721687, 0.2509053647518158, 0.1518847495317459, 0.011399514973163605, 0.08240412920713425, 0.052963949739933014, 0.012185328640043736, 0.03166860342025757, 0.029948236420750618, 0.0332757867872715, 0.026646502315998077, 0.6691258549690247, 0.05157328397035599, 0.010373775847256184, 0.027277877554297447, 0.022091276943683624, 0.06386284530162811, 0.02213944122195244, 0.7486419677734375, 0.1026511937379837, NaN, NaN], [0.0010381464380770922, 0.0033105257898569107, 0.005275417119264603, 0.005129440221935511, 0.05292869359254837, 0.018404772505164146, 0.0016328096389770508, 0.0039754449389874935, 0.007563540246337652, 0.0015294092008844018, 0.00038045260589569807, 0.0016144785331562161, 0.00974529329687357, 0.09415796399116516, 0.176291361451149, 0.35064396262168884, 0.0026081653777509928, 0.0026635529939085245, 0.004589376971125603, 0.028667066246271133, 0.20089752972126007, 0.45412325859069824, 0.4352543354034424, 0.005037708207964897, NaN], [0.1408424973487854, 0.01142195239663124, 0.027654578909277916, 0.018255943432450294, 0.00871819257736206, 0.007302883546799421, 0.002508251927793026, 0.0010894191218540072, 0.002539109904319048, 0.0016572934109717607, 0.002274427330121398, 0.00915378425270319, 0.004932411015033722, 0.000505969044752419, 0.0064278775826096535, 0.013472460210323334, 0.0009905033512040973, 0.004150861874222755, 0.015419019386172295, 0.013300818391144276, 0.00147106999065727, 0.01399929728358984, 0.03311459720134735, 0.0035406623501330614, 0.008275571279227734]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10530310869216919, 0.47072935104370117, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07470229268074036, 0.01594272069633007, 0.3473423421382904, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19784890115261078, 0.02982909232378006, 0.008884507231414318, 0.026416730135679245, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15099161863327026, 0.004257611930370331, 0.06880252063274384, 0.03778434172272682, 0.016005711629986763, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14908726513385773, 0.01576131209731102, 0.006129090208560228, 0.013888919726014137, 0.006888655014336109, 0.007033796049654484, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1207430437207222, 0.0697125568985939, 0.0065151299349963665, 0.0038357542362064123, 0.04419673979282379, 0.16196060180664062, 0.49751368165016174, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02684849314391613, 0.03953110799193382, 0.00281998747959733, 0.001733462675474584, 0.08529012650251389, 0.6486974358558655, 0.306731641292572, 0.07198647409677505, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.012395885773003101, 0.009238478727638721, 0.0003186498652212322, 0.0010813054395839572, 0.008392964489758015, 0.2777543067932129, 0.44055092334747314, 0.0011997584952041507, 0.00246741552837193, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.034838397055864334, 0.015937600284814835, 0.002090656431391835, 0.002794815693050623, 0.008703295141458511, 0.10732896625995636, 0.4454900026321411, 0.001775766140781343, 0.0009654808673076332, 0.016644174233078957, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.293722003698349, 0.0148458918556571, 0.02856721729040146, 0.006315621547400951, 0.005582483485341072, 0.0013911855639889836, 0.004092940129339695, 0.0036679452750831842, 0.0010494120651856065, 0.016411608085036278, 0.023008037358522415, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13037414848804474, 0.020949387922883034, 0.03831411898136139, 0.007462172769010067, 0.02548721246421337, 0.006367610301822424, 0.008434200659394264, 0.010317808948457241, 0.003713584039360285, 0.00402417778968811, 0.19032441079616547, 0.26746228337287903, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.041874390095472336, 0.024160701781511307, 0.00029624058515764773, 0.00016299582784995437, 0.00014630405348725617, 0.0004776908899657428, 0.0010664566652849317, 0.005874973721802235, 0.000636687153019011, 0.0013240330154076219, 0.0912160873413086, 0.35286882519721985, 0.01772063784301281, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11822566390037537, 0.015047432854771614, 0.019423136487603188, 0.00686526857316494, 0.0036870460025966167, 0.00022719512344338, 0.002930518239736557, 0.025171050801873207, 0.005165010690689087, 0.05391281098127365, 0.11512911319732666, 0.07776232063770294, 0.2967449426651001, 0.09380093216896057, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09375648200511932, 0.01475021056830883, 0.012638024985790253, 0.0046005831100046635, 0.051909249275922775, 0.0036223391070961952, 0.004371740389615297, 0.009388775564730167, 0.01159447617828846, 0.023305783048272133, 0.046531662344932556, 0.058873143047094345, 0.07503876090049744, 0.0337555818259716, 0.30213212966918945, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.060409948229789734, 0.03445665165781975, 0.000381257850676775, 0.0036348046269267797, 0.0002713070425670594, 0.0011815812904387712, 0.03030458651483059, 0.03435760363936424, 0.0019682012498378754, 0.00901943538337946, 0.2363511621952057, 0.7836493253707886, 0.05375572293996811, 0.0010517562041059136, 0.002096510259434581, 0.017742546275258064, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19913224875926971, 0.17475517094135284, 0.0022224360145628452, 0.015882516279816628, 0.001058473251760006, 0.0005846276762895286, 0.02601638250052929, 0.037341512739658356, 0.002062901621684432, 0.01394632738083601, 0.062121838331222534, 0.09270716458559036, 0.13391432166099548, 0.011137665249407291, 0.003502808278426528, 0.007463122718036175, 0.4640289545059204, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.33059969544410706, 0.017222048714756966, 0.029873082414269447, 0.008054245263338089, 0.002331576542928815, 0.0006345488945953548, 0.011296147480607033, 0.005269323009997606, 0.0004991231253370643, 0.01808379590511322, 0.0023433570750057697, 0.0409514382481575, 0.01219080574810505, 0.010968736372888088, 0.004035044461488724, 0.000618473335634917, 0.01301309373229742, 0.04461785778403282, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11787470430135727, 0.013379373587667942, 0.03657921776175499, 0.007838133722543716, 0.006328434217721224, 0.0013346761697903275, 0.005374525673687458, 0.005563441663980484, 0.0013783610193058848, 0.003622437361627817, 0.10895299166440964, 0.17491653561592102, 0.013411260209977627, 0.006658618804067373, 0.013080593198537827, 0.0013389869127422571, 0.03540230169892311, 0.3923792839050293, 0.2429211437702179, NaN, NaN, NaN, NaN, NaN, NaN], [0.03099578432738781, 0.01363852247595787, 8.312943100463599e-05, 4.0873743273550645e-05, 3.1056373700266704e-05, 8.971957868197933e-05, 0.0004970009904354811, 0.0021136843133717775, 0.00015606316446792334, 0.0008045462891459465, 0.029241982847452164, 0.24120952188968658, 0.011327153071761131, 0.006169632077217102, 0.004105421248823404, 0.0017298789462074637, 0.09891722351312637, 0.13539430499076843, 0.3545337915420532, 0.03266340494155884, NaN, NaN, NaN, NaN, NaN], [0.05892227217555046, 0.006390280555933714, 0.00726453959941864, 0.002730957930907607, 0.0007821861072443426, 5.8160956541541964e-05, 0.0015625637024641037, 0.007388831116259098, 0.0016573512693867087, 0.027249574661254883, 0.062049947679042816, 0.056622181087732315, 0.2355845421552658, 0.04601869359612465, 0.006218506023287773, 0.00966239720582962, 0.07739637047052383, 0.4012998342514038, 0.09626632183790207, 0.38049787282943726, 0.10569068044424057, NaN, NaN, NaN, NaN], [0.09179559350013733, 0.00951253343373537, 0.010748236440122128, 0.0033872865606099367, 0.04677930101752281, 0.0018132117111235857, 0.0035809800028800964, 0.005968866869807243, 0.0062707834877073765, 0.02606387436389923, 0.033457815647125244, 0.03605461120605469, 0.04817588999867439, 0.03754975646734238, 0.2781437933444977, 0.015551367774605751, 0.2560427486896515, 0.08298799395561218, 0.06865174323320389, 0.12361031025648117, 0.04344068095088005, 0.28463616967201233, NaN, NaN, NaN], [0.02905191108584404, 0.012088212184607983, 0.00011298860044917092, 0.0012518719304352999, 4.317293132771738e-05, 0.0001948956778505817, 0.008923283778131008, 0.008874665014445782, 0.00048750368296168745, 0.0041984752751886845, 0.08557221293449402, 0.46109655499458313, 0.018593793734908104, 0.0004841866611968726, 0.0006005582981742918, 0.004410868044942617, 0.1617877185344696, 0.2815479040145874, 0.7414005398750305, 0.06452517956495285, 0.0009642028599046171, 0.0012653517769649625, 0.012943175621330738, NaN, NaN], [0.1381005197763443, 0.0952477678656578, 0.0011117071844637394, 0.007693122606724501, 0.0001761779421940446, 8.233776316046715e-05, 0.0067709037102758884, 0.015442474745213985, 0.0005836034542880952, 0.005857429001480341, 0.020792629569768906, 0.02682901732623577, 0.05164036154747009, 0.0043857707642018795, 0.0008507486782036722, 0.004215322434902191, 0.19233396649360657, 0.21357974410057068, 0.14138071238994598, 0.12764914333820343, 0.011541306972503662, 0.001996394479647279, 0.004979089833796024, 0.4768531322479248, NaN], [0.14079369604587555, 0.0077750058844685555, 0.008707624860107899, 0.002215370535850525, 0.0003697987995110452, 8.685041393619031e-05, 6.568676326423883e-05, 0.0005928067839704454, 0.00018151948461309075, 0.0013713521184399724, 0.003134837606921792, 0.004530616104602814, 0.0021016064565628767, 0.0014590725768357515, 0.01743447594344616, 0.0004639088874682784, 0.00557903666049242, 0.015868593007326126, 0.012156624346971512, 0.006375743541866541, 0.004486390855163336, 0.037133798003196716, 0.0008373309392482042, 0.015209782868623734, 0.053904592990875244]]], [[[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12363631278276443, 0.14845161139965057, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14363405108451843, 0.021847352385520935, 0.10135873407125473, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13959342241287231, 0.059129536151885986, 0.04632453992962837, 0.0506979376077652, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1401052325963974, 0.20328059792518616, 0.08711162209510803, 0.021569250151515007, 0.06437158584594727, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14849096536636353, 0.24162742495536804, 0.13733072578907013, 0.023916935548186302, 0.4261094033718109, 0.034874048084020615, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1122843325138092, 0.27548718452453613, 0.3164171576499939, 0.11597670614719391, 0.521038293838501, 0.1305568367242813, 0.04802507162094116, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13016629219055176, 0.2326299250125885, 0.3132029175758362, 0.32591310143470764, 0.1516764611005783, 0.09795279055833817, 0.02053435519337654, 0.1865263283252716, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.121080182492733, 0.4840172827243805, 0.47487083077430725, 0.3000609576702118, 0.5299880504608154, 0.09183567762374878, 0.057097259908914566, 0.12967270612716675, 0.04215369373559952, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08035996556282043, 0.5049515962600708, 0.21779249608516693, 0.22551923990249634, 0.48642098903656006, 0.17451445758342743, 0.14853931963443756, 0.2973877787590027, 0.02990546263754368, 0.12922555208206177, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15412510931491852, 0.24815845489501953, 0.21706829965114594, 0.15909965336322784, 0.3919820487499237, 0.2097313106060028, 0.05961627885699272, 0.10788830369710922, 0.04644578695297241, 0.008778278715908527, 0.1666601300239563, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1319347769021988, 0.07332690805196762, 0.3709748387336731, 0.10343886911869049, 0.2416648119688034, 0.273651659488678, 0.142499178647995, 0.032821010798215866, 0.08169299364089966, 0.04221141338348389, 0.04960552975535393, 0.14849121868610382, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15117543935775757, 0.09085448831319809, 0.23665060102939606, 0.09974268078804016, 0.5293540358543396, 0.2969721853733063, 0.0923411101102829, 0.04701923578977585, 0.47750627994537354, 0.31436240673065186, 0.11817371100187302, 0.08098391443490982, 0.05702001228928566, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2022491842508316, 0.0666579008102417, 0.032761361449956894, 0.03407268971204758, 0.3113752603530884, 0.5905517935752869, 0.21839523315429688, 0.043745849281549454, 0.02789805829524994, 0.042396336793899536, 0.08724991232156754, 0.07408890873193741, 0.010044119320809841, 0.12108539044857025, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14857184886932373, 0.38842764496803284, 0.16100677847862244, 0.1839173436164856, 0.03719957172870636, 0.5251989364624023, 0.25831982493400574, 0.06345110386610031, 0.01966739259660244, 0.013820506632328033, 0.10135386884212494, 0.06285497546195984, 0.037499457597732544, 0.09235794097185135, 0.06518241763114929, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15810954570770264, 0.08897967636585236, 0.2754043936729431, 0.11542505025863647, 0.7166418433189392, 0.6856120824813843, 0.15602687001228333, 0.03588242083787918, 0.10233978182077408, 0.06907100230455399, 0.13906386494636536, 0.06064911186695099, 0.02474391460418701, 0.09316151589155197, 0.5409220457077026, 0.18577302992343903, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07972963899374008, 0.06995329260826111, 0.2565014958381653, 0.11985079944133759, 0.5429201126098633, 0.3072132468223572, 0.04467121511697769, 0.06233014911413193, 0.06391221284866333, 0.06306523084640503, 0.04008801653981209, 0.16940940916538239, 0.21208623051643372, 0.3237960636615753, 0.4987465739250183, 0.14530567824840546, 0.42085787653923035, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.057688161730766296, 0.05957844480872154, 0.09227755665779114, 0.06308872997760773, 0.6051628589630127, 0.41719216108322144, 0.06513097882270813, 0.11441777646541595, 0.2576654255390167, 0.039566945284605026, 0.04989808052778244, 0.41204503178596497, 0.6269510388374329, 0.0653882622718811, 0.2309982180595398, 0.05030554160475731, 0.12162061780691147, 0.2016562819480896, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08513950556516647, 0.05776134505867958, 0.44855204224586487, 0.15441171824932098, 0.37962910532951355, 0.43142464756965637, 0.21386101841926575, 0.07478547096252441, 0.22071515023708344, 0.1727379858493805, 0.06471506506204605, 0.1414414495229721, 0.20356127619743347, 0.23849359154701233, 0.28116941452026367, 0.22387196123600006, 0.24124523997306824, 0.10411572456359863, 0.14086224138736725, NaN, NaN, NaN, NaN, NaN, NaN], [0.09857918322086334, 0.08268877118825912, 0.17155912518501282, 0.08326277136802673, 0.3910389840602875, 0.23102693259716034, 0.0706368237733841, 0.04062340036034584, 0.34264665842056274, 0.40400993824005127, 0.14310938119888306, 0.07597656548023224, 0.059025220572948456, 0.46083009243011475, 0.6441643834114075, 0.8002472519874573, 0.34466618299484253, 0.10859531164169312, 0.04317509010434151, 0.042760394513607025, NaN, NaN, NaN, NaN, NaN], [0.07982634007930756, 0.027687683701515198, 0.01305405143648386, 0.01568622700870037, 0.15395750105381012, 0.36470726132392883, 0.09429053217172623, 0.02618592418730259, 0.00988653302192688, 0.03718657046556473, 0.057223062962293625, 0.036843542009592056, 0.008861655369400978, 0.039983998984098434, 0.5628355145454407, 0.5858935713768005, 0.11540589481592178, 0.07112369686365128, 0.022479010745882988, 0.0049066911451518536, 0.07443748414516449, NaN, NaN, NaN, NaN], [0.13230623304843903, 0.39635705947875977, 0.12619565427303314, 0.23844560980796814, 0.04749276116490364, 0.5552228093147278, 0.304650217294693, 0.16151569783687592, 0.05923860892653465, 0.03940735384821892, 0.37161606550216675, 0.13852664828300476, 0.1098584458231926, 0.421970933675766, 0.059641290456056595, 0.35413044691085815, 0.2336989790201187, 0.21869167685508728, 0.04408164322376251, 0.03093402087688446, 0.08392708003520966, 0.038801465183496475, NaN, NaN, NaN], [0.06938444077968597, 0.08034616708755493, 0.1555827558040619, 0.07347460091114044, 0.4763748347759247, 0.40589335560798645, 0.07265187799930573, 0.022002995014190674, 0.0527057945728302, 0.07314148545265198, 0.11090734601020813, 0.03504399210214615, 0.0172868762165308, 0.14030121266841888, 0.3467526137828827, 0.21038202941417694, 0.6312639117240906, 0.1208876520395279, 0.020520374178886414, 0.014591614715754986, 0.03736459091305733, 0.22129306197166443, 0.05682671070098877, NaN, NaN], [0.08218587934970856, 0.08353152126073837, 0.244074746966362, 0.15340235829353333, 0.5709766745567322, 0.4268343448638916, 0.06391507387161255, 0.13458560407161713, 0.14046461880207062, 0.13024689257144928, 0.043825987726449966, 0.1802380084991455, 0.2593124508857727, 0.4235299825668335, 0.23401854932308197, 0.23376718163490295, 0.4458163380622864, 0.1644086241722107, 0.22351105511188507, 0.25077733397483826, 0.28149890899658203, 0.3320602774620056, 0.05098887160420418, 0.4388013482093811, NaN], [0.13887250423431396, 0.1972966492176056, 0.3352757692337036, 0.30585116147994995, 0.6380553841590881, 0.5158089995384216, 0.3850407004356384, 0.3912012279033661, 0.2877788245677948, 0.30187875032424927, 0.20025724172592163, 0.34020906686782837, 0.47167572379112244, 0.3815076947212219, 0.5385518074035645, 0.20663535594940186, 0.37741178274154663, 0.29376763105392456, 0.3577961027622223, 0.21765607595443726, 0.14290691912174225, 0.3544510304927826, 0.07646653801202774, 0.1391337811946869, 0.019570577889680862]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10658828914165497, 0.44162610173225403, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14346696436405182, 0.1105659008026123, 0.04705679044127464, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14569434523582458, 0.006359750870615244, 0.06321832537651062, 0.009962446056306362, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14614860713481903, 0.0770370289683342, 0.14572308957576752, 0.11918944120407104, 0.003047030884772539, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16211360692977905, 0.1199408695101738, 0.008137544617056847, 0.026895001530647278, 0.022997038438916206, 0.0004772362008225173, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1276824176311493, 0.05415544658899307, 0.008876973763108253, 0.006533092353492975, 0.16286829113960266, 0.4191088378429413, 0.11241274327039719, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1310766041278839, 0.09720440953969955, 0.005617472343146801, 0.018550021573901176, 0.07474999874830246, 0.03211009502410889, 0.01561786886304617, 0.5897646546363831, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07142644375562668, 0.019657818600535393, 0.044225241988897324, 0.006672952324151993, 0.015112369321286678, 0.03715437650680542, 0.012035970576107502, 0.08684496581554413, 0.5578015446662903, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06384367495775223, 0.009399783797562122, 0.06692944467067719, 0.013825987465679646, 0.01438650768250227, 0.11814092099666595, 0.025182364508509636, 0.04756484180688858, 0.4922580420970917, 0.010614832863211632, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21570175886154175, 0.004600263200700283, 0.0039491499774158, 0.0010213260538876057, 0.00511409854516387, 0.00780195789411664, 0.0035460677463561296, 0.06005942076444626, 0.002209970960393548, 0.0011990047059953213, 0.010184505954384804, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15804870426654816, 0.10358668118715286, 0.018792977556586266, 0.0036350360605865717, 0.02226737141609192, 0.007843486964702606, 0.002713214373216033, 0.3624168336391449, 0.00397031893953681, 0.013842551037669182, 0.05391863361001015, 0.040338534861803055, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0703621581196785, 0.01676221750676632, 0.03283774480223656, 0.005265639629215002, 0.016811830922961235, 0.008307189680635929, 0.0008217993890866637, 0.06662888079881668, 0.006444453727453947, 0.0015952866524457932, 0.03341786190867424, 0.28674793243408203, 0.09830270707607269, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00274313404224813, 0.01220498327165842, 0.001565106911584735, 0.014617281965911388, 0.0015394951915368438, 0.00014163085143081844, 0.0032730719540268183, 0.04253724217414856, 0.01929563470184803, 0.0011092370841652155, 0.008900013752281666, 0.14250728487968445, 0.44352540373802185, 0.012739983387291431, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12441921979188919, 0.09727630764245987, 0.031539320945739746, 0.0390433706343174, 0.004017204977571964, 0.003718326799571514, 0.06902258098125458, 0.21229486167430878, 0.1692674309015274, 0.507585346698761, 0.24224399030208588, 0.4713107943534851, 0.22175242006778717, 0.1071210727095604, 0.001354279462248087, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11131177842617035, 0.045754965394735336, 0.13187335431575775, 0.021390099078416824, 0.2008819729089737, 0.1753949522972107, 0.029810786247253418, 0.1191062182188034, 0.0330519825220108, 0.021209293976426125, 0.007793682627379894, 0.004569755867123604, 0.21031485497951508, 0.08390634506940842, 0.11696453392505646, 0.2920413017272949, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.28942060470581055, 0.004874760750681162, 0.02575746178627014, 0.03629674017429352, 0.0339069589972496, 0.06067432835698128, 0.06949229538440704, 0.17600718140602112, 0.04042575880885124, 0.0021073101088404655, 0.002125136088579893, 0.0013297069817781448, 0.013164625503122807, 0.019647862762212753, 0.0625171884894371, 0.003036472015082836, 0.15673543512821198, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.29843398928642273, 0.006499151699244976, 0.002175502711907029, 0.00474061444401741, 0.012194045819342136, 0.024305779486894608, 0.05332900583744049, 0.20892387628555298, 0.06725459545850754, 0.0056669809855520725, 0.023831704631447792, 0.0038352743722498417, 0.008001168258488178, 0.00692057004198432, 0.006051996257156134, 0.0008782879449427128, 0.0244371946901083, 0.05294432491064072, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19362471997737885, 0.05030333995819092, 0.012831996195018291, 0.0028119448106735945, 0.011659904383122921, 0.0070129260420799255, 0.002673238283023238, 0.1857692450284958, 0.0015845311572775245, 0.003893241984769702, 0.009055504575371742, 0.013083641417324543, 0.009338575415313244, 0.007860029116272926, 0.009482803754508495, 0.019751103594899178, 0.03845033049583435, 0.03947525471448898, 0.03009573556482792, NaN, NaN, NaN, NaN, NaN, NaN], [0.08181142061948776, 0.013090993277728558, 0.025600923225283623, 0.0045991819351911545, 0.007844633422791958, 0.0066622160375118256, 0.0006054755649529397, 0.01805841363966465, 0.0025927021633833647, 0.0006796378293074667, 0.012531430460512638, 0.18806973099708557, 0.04688132554292679, 0.005460845306515694, 0.053047653287649155, 0.013497358188033104, 0.040136244148015976, 0.022071214392781258, 0.31691932678222656, 0.07654344290494919, NaN, NaN, NaN, NaN, NaN], [0.003571689361706376, 0.007330529857426882, 0.0009176949388347566, 0.011351491324603558, 0.0005700239562429488, 0.0001114286933443509, 0.0023790227714926004, 0.011217805556952953, 0.004490875173360109, 0.00038650527130812407, 0.0025467458181083202, 0.048559535294771194, 0.22723886370658875, 0.0019670024048537016, 0.0002542402071412653, 0.027445662766695023, 0.015111691318452358, 0.029036840423941612, 0.2144545316696167, 0.4208240211009979, 0.013829981908202171, NaN, NaN, NaN, NaN], [0.11162849515676498, 0.06633912026882172, 0.017337389290332794, 0.030477523803710938, 0.0024834000505506992, 0.001867939718067646, 0.03932232782244682, 0.1628599613904953, 0.14192035794258118, 0.2944621741771698, 0.21811458468437195, 0.42557209730148315, 0.2638176381587982, 0.14630424976348877, 0.0005040403339080513, 0.32521945238113403, 0.2411627173423767, 0.28287336230278015, 0.40539565682411194, 0.1682160645723343, 0.08244442939758301, 0.001218001707457006, NaN, NaN, NaN], [0.20973265171051025, 0.07712213695049286, 0.20427735149860382, 0.025535617023706436, 0.4053865373134613, 0.41131824254989624, 0.030548784881830215, 0.060146916657686234, 0.012079673819243908, 0.01592317223548889, 0.0048461491242051125, 0.0021770852617919445, 0.09957096725702286, 0.1170588806271553, 0.13386258482933044, 0.16141492128372192, 0.004613581579178572, 0.015190798789262772, 0.003683852730318904, 0.1389266699552536, 0.07006954401731491, 0.1815212517976761, 0.17825333774089813, NaN, NaN], [0.3360293209552765, 0.0046190484426915646, 0.024437543004751205, 0.03736568242311478, 0.023848971351981163, 0.05927197262644768, 0.0542423352599144, 0.09209144860506058, 0.023972967639565468, 0.000766670098528266, 0.0006589474505744874, 0.0007115502958185971, 0.00637162895873189, 0.012912634760141373, 0.014624576084315777, 0.0019432539120316505, 0.05897590517997742, 0.0038116518408060074, 0.0016802565660327673, 0.011611220426857471, 0.025170182809233665, 0.04455949738621712, 0.0020357028115540743, 0.14134161174297333, NaN], [0.187117338180542, 0.005916869733482599, 0.020901108160614967, 0.0559980571269989, 0.0324174202978611, 0.008547084406018257, 0.044511571526527405, 0.04880741238594055, 0.05289075896143913, 0.038245368748903275, 0.003611604683101177, 0.002279189880937338, 0.01790045015513897, 0.008863909170031548, 0.01127588003873825, 0.005861865822225809, 0.17173975706100464, 0.009364882484078407, 0.005221609957516193, 0.012455414980649948, 0.007264893501996994, 0.016177698969841003, 0.008824422955513, 0.18642237782478333, 0.0006185321253724396]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12484697252511978, 0.1276315450668335, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15841424465179443, 0.03031034581363201, 0.02654799446463585, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13769303262233734, 0.09575259685516357, 0.025977646932005882, 0.052591271698474884, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15085087716579437, 0.15096567571163177, 0.09222358465194702, 0.028469638898968697, 0.0012114758137613535, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16431185603141785, 0.07204771786928177, 0.05053501948714256, 0.012478960677981377, 0.05114812031388283, 0.00039714027661830187, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1666734665632248, 0.06891340762376785, 0.013632094487547874, 0.018171580508351326, 0.002599227475002408, 0.0009873181115835905, 0.0006481229793280363, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14423918724060059, 0.12251336872577667, 0.10176724940538406, 0.33380815386772156, 0.1583750993013382, 0.023372141644358635, 0.026839546859264374, 0.06730155646800995, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2790219187736511, 0.15446610748767853, 0.015893638134002686, 0.03619629144668579, 0.003051391802728176, 0.00038247412885539234, 0.0007123185787349939, 0.010222047567367554, 0.0010863485513255, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.26870372891426086, 0.10405707359313965, 0.00916238222271204, 0.058617573231458664, 0.0049601029604673386, 0.0005682760966010392, 0.004407011903822422, 0.03309918940067291, 0.0036104319151490927, 0.12174393236637115, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05985519662499428, 0.14893494546413422, 0.09544339030981064, 0.18974637985229492, 0.1120084673166275, 0.28269606828689575, 0.4275827407836914, 0.12184610962867737, 0.40095797181129456, 0.08120625466108322, 0.27448615431785583, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06809581816196442, 0.09586934000253677, 0.10229554027318954, 0.057183876633644104, 0.25635847449302673, 0.19582371413707733, 0.4237477481365204, 0.37648820877075195, 0.48733898997306824, 0.20777222514152527, 0.24944597482681274, 0.45371755957603455, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05513762682676315, 0.16880887746810913, 0.02300925739109516, 0.03029457852244377, 0.032050080597400665, 0.0745139941573143, 0.08332593739032745, 0.5048279166221619, 0.051856089383363724, 0.16889351606369019, 0.22218117117881775, 0.29087209701538086, 0.03443009778857231, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07503295689821243, 0.22708888351917267, 0.011672623455524445, 0.03240634873509407, 0.051372844725847244, 0.0555996336042881, 0.1055832952260971, 0.27455389499664307, 0.019383858889341354, 0.29115474224090576, 0.25329896807670593, 0.3762655258178711, 0.06596359610557556, 0.027243560180068016, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15851522982120514, 0.22386471927165985, 0.13473065197467804, 0.10273782163858414, 0.539568305015564, 0.23089595139026642, 0.2947250008583069, 0.2566256523132324, 0.08758009225130081, 0.04963833838701248, 0.026406293734908104, 0.02359875850379467, 0.06999926269054413, 0.014701825566589832, 0.008440684527158737, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1888987272977829, 0.22277534008026123, 0.06621028482913971, 0.04940320923924446, 0.013609242625534534, 0.012980671599507332, 0.0275713000446558, 0.5000426769256592, 0.025658253580331802, 0.28077542781829834, 0.21061377227306366, 0.1005047932267189, 0.0123829934746027, 0.005874408408999443, 0.04495157673954964, 0.007559731602668762, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10630622506141663, 0.1130438968539238, 0.04711592569947243, 0.14829613268375397, 0.0012987125664949417, 0.0009870391804724932, 0.002409427659586072, 0.10731083154678345, 0.010861101560294628, 0.02266101725399494, 0.22295407950878143, 0.37738272547721863, 0.21324896812438965, 0.09625840187072754, 0.01478838175535202, 0.004724964965134859, 0.13376930356025696, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0042772903107106686, 0.006450775545090437, 0.00791113544255495, 0.01871791109442711, 0.02349945716559887, 0.036059893667697906, 0.09560179710388184, 0.01157363597303629, 0.020316841080784798, 0.002858342370018363, 0.0015840751584619284, 0.03869258984923363, 0.04008479043841362, 0.0456826388835907, 0.061234306544065475, 0.32812535762786865, 0.4548730254173279, 0.048923686146736145, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.034464891999959946, 0.04304976761341095, 0.0730237364768982, 0.07959159463644028, 0.156441330909729, 0.14927342534065247, 0.37836754322052, 0.2500280439853668, 0.265838086605072, 0.038285933434963226, 0.0458042174577713, 0.2175784856081009, 0.055615901947021484, 0.32925114035606384, 0.23017114400863647, 0.5254709720611572, 0.3807608187198639, 0.4477500319480896, 0.3941081464290619, NaN, NaN, NaN, NaN, NaN, NaN], [0.024431752040982246, 0.057854264974594116, 0.009785568341612816, 0.015689833089709282, 0.010099711827933788, 0.022971261292696, 0.026158222928643227, 0.08270542323589325, 0.00771379703655839, 0.023359954357147217, 0.06216609850525856, 0.1452798992395401, 0.010090651921927929, 0.13497084379196167, 0.023736534640192986, 0.06422590464353561, 0.2799428105354309, 0.34307411313056946, 0.27198341488838196, 0.018816450610756874, NaN, NaN, NaN, NaN, NaN], [0.032250434160232544, 0.07008427381515503, 0.003495490411296487, 0.011726448312401772, 0.013232100754976273, 0.021211393177509308, 0.02240551821887493, 0.050749149173498154, 0.0020511853508651257, 0.034987252205610275, 0.05167752131819725, 0.10231753438711166, 0.017492327839136124, 0.0036121474113315344, 0.0030979528091847897, 0.14347726106643677, 0.4107814431190491, 0.18759746849536896, 0.28042495250701904, 0.02327493391931057, 0.023935986682772636, NaN, NaN, NaN, NaN], [0.17385193705558777, 0.24280618131160736, 0.0901411697268486, 0.1509939581155777, 0.5964542627334595, 0.18189039826393127, 0.25377142429351807, 0.39126867055892944, 0.11990400403738022, 0.04869762808084488, 0.06967514008283615, 0.0491257943212986, 0.1536286324262619, 0.04553663358092308, 0.006321897264569998, 0.008409527130424976, 0.01950901933014393, 0.028066763654351234, 0.039955586194992065, 0.08575458079576492, 0.02489100769162178, 0.0107131227850914, NaN, NaN, NaN], [0.18693126738071442, 0.25040745735168457, 0.07803116738796234, 0.06071358174085617, 0.018153348937630653, 0.012512190267443657, 0.012858238071203232, 0.18478038907051086, 0.008756724186241627, 0.14063727855682373, 0.16963867843151093, 0.06472224742174149, 0.008233368396759033, 0.010625114664435387, 0.04533438757061958, 0.004584541078656912, 0.04685693234205246, 0.3269248306751251, 0.13935554027557373, 0.022706659510731697, 0.015514994971454144, 0.09856907278299332, 0.009564985521137714, NaN, NaN], [0.10220125317573547, 0.06584151834249496, 0.046970706433057785, 0.16499453783035278, 0.0008504274883307517, 0.000721337681170553, 0.0015187861863523722, 0.050142802298069, 0.005332621280103922, 0.005509581416845322, 0.0572623535990715, 0.172898530960083, 0.12213093042373657, 0.0640687644481659, 0.004657925106585026, 0.002522988012060523, 0.028443191200494766, 0.29674383997917175, 0.3544806241989136, 0.20916549861431122, 0.09151047468185425, 0.014975211583077908, 0.0019209993770346045, 0.07398010790348053, NaN], [0.014319260604679585, 0.019726725295186043, 0.010809341445565224, 0.06728478521108627, 0.024899542331695557, 0.06927011907100677, 0.2726534307003021, 0.06849226355552673, 0.06274150311946869, 0.0032663261517882347, 0.007571991998702288, 0.011041088029742241, 0.0653790682554245, 0.06552072614431381, 0.10165777057409286, 0.05923810228705406, 0.20752549171447754, 0.1128133162856102, 0.041725482791662216, 0.12833572924137115, 0.10405165702104568, 0.2233171910047531, 0.10715138167142868, 0.3742898404598236, 0.43902406096458435]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12878015637397766, 0.05999259278178215, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16734670102596283, 0.0018487111665308475, 0.002184537472203374, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06620991975069046, 0.4480140209197998, 0.42379117012023926, 0.3748236298561096, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1498516947031021, 0.091057188808918, 0.11073686927556992, 0.05954570695757866, 0.00012444167805369943, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15789009630680084, 0.05178086459636688, 0.2272004932165146, 0.05532779544591904, 0.002530630910769105, 0.00011625503975665197, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05158510431647301, 0.42307329177856445, 0.4962795376777649, 0.6637455821037292, 0.11636865884065628, 0.027691489085555077, 0.059323750436306, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1440366506576538, 0.37752795219421387, 0.42684903740882874, 0.13104133307933807, 0.0449170246720314, 0.0360451340675354, 0.007316120434552431, 0.03281773626804352, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.018571142107248306, 0.11001976579427719, 0.16728174686431885, 0.33147770166397095, 0.29621925950050354, 0.11174014210700989, 0.46736985445022583, 0.18467408418655396, 0.05186863988637924, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0193540807813406, 0.11997552216053009, 0.4339123070240021, 0.4291674792766571, 0.22741732001304626, 0.21840345859527588, 0.4310562014579773, 0.16546283662319183, 0.05634206160902977, 0.03477246314287186, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07166115939617157, 0.34385329484939575, 0.5272834300994873, 0.4769807457923889, 0.34829023480415344, 0.19288644194602966, 0.1752767115831375, 0.3240547180175781, 0.026788396760821342, 0.09653788805007935, 0.14339366555213928, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09568949043750763, 0.2010803371667862, 0.1452081948518753, 0.13633964955806732, 0.13264110684394836, 0.11369673907756805, 0.18754418194293976, 0.10573749244213104, 0.12209529429674149, 0.3772747814655304, 0.4260762333869934, 0.1448964774608612, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1600937843322754, 0.32966408133506775, 0.46643200516700745, 0.2761552929878235, 0.1128716766834259, 0.16030451655387878, 0.13808301091194153, 0.12019707262516022, 0.08980843424797058, 0.23569302260875702, 0.18699060380458832, 0.06252679228782654, 0.02190866880118847, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09671676903963089, 0.3181785047054291, 0.5044789910316467, 0.5311775803565979, 0.43058764934539795, 0.24623769521713257, 0.546705424785614, 0.20948244631290436, 0.5971428155899048, 0.15125280618667603, 0.21692372858524323, 0.08393274247646332, 0.0805632621049881, 0.11463441699743271, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17538371682167053, 0.005170984659343958, 0.01562126912176609, 0.012803001329302788, 0.0004321248270571232, 0.003303500125184655, 0.010391591116786003, 0.0083633316680789, 0.001453742035664618, 0.0005911564221605659, 0.001968160504475236, 0.018067756667733192, 0.0012553221313282847, 0.0006174716982059181, 0.0014710418181493878, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00964878499507904, 0.07296860218048096, 0.1732037365436554, 0.2482636272907257, 0.018695944920182228, 0.04061494395136833, 0.019565006718039513, 0.048743683844804764, 0.15582872927188873, 0.0506676621735096, 0.08059392869472504, 0.2691291868686676, 0.4701274335384369, 0.05269847437739372, 0.15863555669784546, 0.011098350398242474, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.023792432621121407, 0.42975902557373047, 0.3812340199947357, 0.23295366764068604, 0.2699258625507355, 0.32472288608551025, 0.04527096822857857, 0.2556793987751007, 0.5905154347419739, 0.8116171360015869, 0.684613823890686, 0.13916483521461487, 0.05671815946698189, 0.0401710644364357, 0.30002903938293457, 0.014873968437314034, 0.1109585389494896, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07327478379011154, 0.42313894629478455, 0.7821765542030334, 0.6752634048461914, 0.18926696479320526, 0.27897483110427856, 0.1972714066505432, 0.26650866866111755, 0.21928414702415466, 0.6610813736915588, 0.8023169040679932, 0.32853400707244873, 0.043605707585811615, 0.04177317023277283, 0.5147100687026978, 0.014965414069592953, 0.041893746703863144, 0.10476090759038925, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09543995559215546, 0.1369307041168213, 0.1906978189945221, 0.1367466300725937, 0.17180036008358002, 0.12260185182094574, 0.13847540318965912, 0.1559406965970993, 0.13510896265506744, 0.4644373655319214, 0.6843520402908325, 0.2938932180404663, 0.08134166151285172, 0.16692468523979187, 0.35020914673805237, 0.0983358696103096, 0.26928237080574036, 0.11322443932294846, 0.14002281427383423, NaN, NaN, NaN, NaN, NaN, NaN], [0.17294523119926453, 0.44891712069511414, 0.5596615076065063, 0.3151743412017822, 0.15508009493350983, 0.20398668944835663, 0.18162229657173157, 0.14380685985088348, 0.09279182553291321, 0.25614914298057556, 0.37145668268203735, 0.2047339379787445, 0.05775143578648567, 0.06389063596725464, 0.19947569072246552, 0.07508620619773865, 0.162083700299263, 0.036575064063072205, 0.05963924527168274, 0.02704720012843609, NaN, NaN, NaN, NaN, NaN], [0.09450869262218475, 0.5263407230377197, 0.5685468316078186, 0.6246378421783447, 0.5457862615585327, 0.4288109838962555, 0.7265884876251221, 0.4213257133960724, 0.7441360354423523, 0.37028953433036804, 0.4906199276447296, 0.24940308928489685, 0.2854059636592865, 0.25606390833854675, 0.06486664712429047, 0.03651905804872513, 0.215606689453125, 0.16494624316692352, 0.07126681506633759, 0.0978088453412056, 0.18553400039672852, NaN, NaN, NaN, NaN], [0.19233128428459167, 0.0069253402762115, 0.019198253750801086, 0.024288823828101158, 0.0006626379326917231, 0.0032825330272316933, 0.012745865620672703, 0.02121213637292385, 0.004573441576212645, 0.001344278221949935, 0.010449343360960484, 0.07998955249786377, 0.008849495090544224, 0.005957764107733965, 0.00281895836815238, 0.0006993816932663321, 0.0011300387559458613, 0.0034355262760072947, 0.006048144306987524, 0.0007683978183194995, 0.00029024321702308953, 0.0009215899626724422, NaN, NaN, NaN], [0.00490582175552845, 0.09978753328323364, 0.17523892223834991, 0.18201382458209991, 0.025161702185869217, 0.0351867638528347, 0.008898423984646797, 0.033712878823280334, 0.06612548977136612, 0.044598400592803955, 0.0818907842040062, 0.31783777475357056, 0.6522275805473328, 0.26521986722946167, 0.31609129905700684, 0.0543142631649971, 0.07028744369745255, 0.06436092406511307, 0.12702754139900208, 0.4257008731365204, 0.05356784537434578, 0.20406562089920044, 0.022904740646481514, NaN, NaN], [0.02933959849178791, 0.5456263422966003, 0.4945109188556671, 0.26123103499412537, 0.3237256109714508, 0.3705388903617859, 0.04209306091070175, 0.3351372182369232, 0.658141016960144, 0.8126230239868164, 0.8673186898231506, 0.28273773193359375, 0.11254162341356277, 0.17348313331604004, 0.7003386616706848, 0.1474425047636032, 0.36997753381729126, 0.41849759221076965, 0.091117262840271, 0.03724836930632591, 0.036747273057699203, 0.47380825877189636, 0.017722588032484055, 0.0920308530330658, NaN], [0.1429738998413086, 0.11406568437814713, 0.30407312512397766, 0.04420004412531853, 0.050888776779174805, 0.009020227938890457, 0.026264725252985954, 0.20154790580272675, 0.284900963306427, 0.16813665628433228, 0.6384625434875488, 0.35198092460632324, 0.0041788192465901375, 0.017796171829104424, 0.06702794879674911, 0.017356209456920624, 0.11703062057495117, 0.363391250371933, 0.08829980343580246, 0.0006652214215137064, 0.002063008025288582, 0.01232101023197174, 0.0010344748152419925, 0.005295889917761087, 0.10532692819833755]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1283751130104065, 0.06695841252803802, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [5.319380943547003e-05, 9.114345448324457e-05, 0.7905611991882324, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10777772217988968, 0.19019582867622375, 0.12566408514976501, 0.295462429523468, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [2.4899240088416263e-05, 2.9243250537547283e-05, 0.0014855118934065104, 3.888772698701359e-05, 0.9169090986251831, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [3.5349924587535497e-07, 4.689470642915694e-06, 0.02691131830215454, 1.3325815416465048e-05, 0.19568589329719543, 0.956480085849762, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08490768820047379, 0.04920955002307892, 0.012384464032948017, 0.04339546710252762, 0.010612337850034237, 0.05702771991491318, 0.7263003587722778, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16491760313510895, 0.04815620183944702, 0.0007595600909553468, 0.006606678944081068, 0.0006115635624155402, 0.0007167417788878083, 0.0015418223338201642, 0.0024032427463680506, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.012053201906383038, 0.18336322903633118, 0.0033893296495079994, 0.22584111988544464, 0.004534169565886259, 0.003455487545579672, 0.30805450677871704, 0.5499533414840698, 0.13390673696994781, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02224119007587433, 0.09969844669103622, 0.01827961951494217, 0.1828235685825348, 0.009660250507295132, 0.005268027540296316, 0.13511976599693298, 0.39505934715270996, 0.1772008240222931, 0.6222725510597229, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19008594751358032, 0.025696618482470512, 0.004118501208722591, 0.03605509176850319, 0.002144730417057872, 0.0023362801875919104, 0.16961191594600677, 0.015426162630319595, 0.016875047236680984, 0.017404966056346893, 0.032629188150167465, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1594686657190323, 0.03835373371839523, 0.021387629210948944, 0.028402678668498993, 0.12163796275854111, 0.1348690688610077, 0.027878204360604286, 0.016979072242975235, 0.009301519952714443, 0.047045812010765076, 0.103324294090271, 0.0978349894285202, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08206925541162491, 0.0482555516064167, 0.03066202998161316, 0.14434732496738434, 0.10149279236793518, 0.1536794900894165, 0.16425268352031708, 0.00592045346274972, 0.002011190867051482, 0.030538976192474365, 0.015422381460666656, 0.0400862954556942, 0.6933969259262085, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11962933838367462, 0.08867897093296051, 0.023231033235788345, 0.019267449155449867, 0.06578893214464188, 0.01314490009099245, 0.028238458558917046, 0.2009190320968628, 0.005505711771547794, 0.024347275495529175, 0.005847027525305748, 0.13606473803520203, 0.11386173218488693, 0.6883828639984131, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.004133098293095827, 0.007605875376611948, 0.380069762468338, 0.01569206453859806, 0.3162667751312256, 0.06185031309723854, 0.003268925240263343, 0.007663627155125141, 0.00711404625326395, 0.0016827658982947469, 0.002885768422856927, 0.009058460593223572, 0.0104479705914855, 0.0013903286308050156, 0.9176042079925537, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19946889579296112, 0.004915847908705473, 0.0015343156410381198, 0.012221671640872955, 0.003153382334858179, 0.0001576353097334504, 0.0020530277397483587, 0.003957398701459169, 0.010446527041494846, 0.012547693215310574, 0.03473197668790817, 0.06650777161121368, 0.014228541404008865, 0.02601468935608864, 0.0018418998224660754, 0.08826413750648499, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14040440320968628, 0.29221969842910767, 0.09665771573781967, 0.2947876751422882, 0.00611721258610487, 0.012681002728641033, 0.7610099911689758, 0.27993685007095337, 0.19895455241203308, 0.07963719218969345, 0.025141140446066856, 0.30299919843673706, 0.4374280273914337, 0.12315846234560013, 0.011889583431184292, 0.00027308438438922167, 0.03226177766919136, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.22362156212329865, 0.19648011028766632, 0.02122899703681469, 0.12822405993938446, 0.013841216452419758, 0.009505078196525574, 0.4746513366699219, 0.1753886640071869, 0.09167484194040298, 0.038334570825099945, 0.04122844338417053, 0.14653263986110687, 0.17874038219451904, 0.023550381883978844, 0.014212163165211678, 0.001423373818397522, 0.0059451088309288025, 0.09707646816968918, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.167328879237175, 0.06208498775959015, 0.010482249781489372, 0.03574186563491821, 0.0675959512591362, 0.06477286666631699, 0.04995346441864967, 0.05412250757217407, 0.009984727017581463, 0.03347667679190636, 0.11074735969305038, 0.16135196387767792, 0.07774785906076431, 0.01735900156199932, 0.007863441482186317, 0.019525114446878433, 0.005842071026563644, 0.1275986284017563, 0.0955328494310379, NaN, NaN, NaN, NaN, NaN, NaN], [0.05032582953572273, 0.03989394009113312, 0.02223959006369114, 0.07248460501432419, 0.04305185005068779, 0.04872481897473335, 0.09144517779350281, 0.0032577940728515387, 0.000561918190214783, 0.015125684440135956, 0.018474824726581573, 0.0519116036593914, 0.7149417400360107, 0.023930398747324944, 0.005549557972699404, 0.0027118371799588203, 0.08418004959821701, 0.22684048116207123, 0.052481237798929214, 0.7548789381980896, NaN, NaN, NaN, NaN, NaN], [0.14971917867660522, 0.12296220660209656, 0.03256092593073845, 0.015910452231764793, 0.08324312418699265, 0.010959222912788391, 0.03249981626868248, 0.2630986273288727, 0.0023772413842380047, 0.021863164380192757, 0.014683729968965054, 0.3797665238380432, 0.26638853549957275, 0.6724205613136292, 0.015757206827402115, 0.01569446735084057, 0.01732691004872322, 0.06738004088401794, 0.17602917551994324, 0.12501026690006256, 0.6636221408843994, NaN, NaN, NaN, NaN], [0.0045495470985770226, 0.007598123978823423, 0.48235079646110535, 0.017675379291176796, 0.30638325214385986, 0.03773635998368263, 0.0025513810105621815, 0.013349749147891998, 0.011474208906292915, 0.002688285429030657, 0.009704438969492912, 0.024301802739501, 0.030528949573636055, 0.006023744586855173, 0.9289764761924744, 0.008095184341073036, 0.015121471136808395, 0.003912394400686026, 0.005678378511220217, 0.005922055337578058, 0.0012866485631093383, 0.9431078433990479, NaN, NaN, NaN], [0.25144028663635254, 0.013477480970323086, 0.004043558146804571, 0.02197866141796112, 0.005731666926294565, 0.00035365403164178133, 0.0028230457101017237, 0.003569219959899783, 0.00616231607273221, 0.023324957117438316, 0.07691453397274017, 0.11847300082445145, 0.025281671434640884, 0.05239935964345932, 0.002384425140917301, 0.16120819747447968, 0.011955172754824162, 0.09212952852249146, 0.03993848338723183, 0.017148757353425026, 0.01459744293242693, 0.0018050760263577104, 0.08139479160308838, NaN, NaN], [0.08713241666555405, 0.22884246706962585, 0.12139283120632172, 0.21789073944091797, 0.00419022049754858, 0.011025986634194851, 0.8093750476837158, 0.24520863592624664, 0.11868450790643692, 0.037659380584955215, 0.014297883957624435, 0.35379931330680847, 0.4382935166358948, 0.17632676661014557, 0.006937071681022644, 0.0007303177262656391, 0.027538392692804337, 0.0690605565905571, 0.3237524628639221, 0.41753751039505005, 0.09520361572504044, 0.013310365378856659, 0.0003602981742005795, 0.032565031200647354, NaN], [0.01268855668604374, 0.009620537050068378, 0.0011078648967668414, 0.01395372860133648, 0.00034480926115065813, 0.0002369812864344567, 0.14032205939292908, 0.12187758088111877, 0.004498081747442484, 6.632315489696339e-05, 0.01873306930065155, 0.07693066447973251, 0.06357964873313904, 0.012718681246042252, 0.02489433065056801, 0.4312428832054138, 0.013737366534769535, 0.0326746366918087, 0.34456172585487366, 0.0668448805809021, 0.006646350026130676, 0.04233057424426079, 0.4123155176639557, 0.007851892150938511, 0.43338367342948914]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13150663673877716, 0.013105388730764389, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16704899072647095, 0.0014066778821870685, 0.003860085504129529, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14769184589385986, 0.005059333052486181, 0.0053715878166258335, 0.026609797030687332, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15381431579589844, 0.05056624114513397, 0.015615872107446194, 0.004382571205496788, 0.00015187788812909275, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16606314480304718, 0.03878505155444145, 0.01631396822631359, 0.011268166825175285, 0.00036908386391587555, 0.00010962320084217936, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16556474566459656, 0.059035927057266235, 0.018687130883336067, 0.020593103021383286, 0.0006985706277191639, 0.0006753651541657746, 0.01174053642898798, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16100119054317474, 0.03705580160021782, 0.08672276139259338, 0.05696912482380867, 0.00507472176104784, 0.006951047107577324, 0.0023692583199590445, 0.004235508386045694, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.288095086812973, 0.011840847320854664, 0.005622565280646086, 0.00535928551107645, 0.0008760345517657697, 0.0004899614141322672, 0.001179057639092207, 0.0010409504175186157, 0.0012723063118755817, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2984195351600647, 0.024577315896749496, 0.008883590810000896, 0.0237559974193573, 0.001871026586741209, 0.002048116410151124, 0.00452006608247757, 0.0067189703695476055, 0.002311990363523364, 0.0035932722967118025, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19755195081233978, 0.08605571836233139, 0.04371126368641853, 0.045333728194236755, 0.005393510684370995, 0.006479238625615835, 0.018500106409192085, 0.012994848191738129, 0.011254888959228992, 0.03004884347319603, 0.011813223361968994, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05165635421872139, 0.44527125358581543, 0.31059694290161133, 0.6649516224861145, 0.027770839631557465, 0.02873762883245945, 0.17512862384319305, 0.06940869987010956, 0.1633579134941101, 0.028000785037875175, 0.003091411432251334, 0.016245586797595024, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19151811301708221, 0.1383962333202362, 0.13229386508464813, 0.35712042450904846, 0.18756243586540222, 0.2871147096157074, 0.5138459801673889, 0.22405852377414703, 0.28785935044288635, 0.04021993279457092, 0.0012617700267583132, 0.004019713494926691, 0.003964945673942566, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24189773201942444, 0.08955204486846924, 0.32067012786865234, 0.20245005190372467, 0.11740265786647797, 0.08460556715726852, 0.044664137065410614, 0.025831788778305054, 0.07413194328546524, 0.0068964180536568165, 0.002961511956527829, 0.005619046278297901, 0.0014741680352017283, 0.00546230049803853, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1724659651517868, 0.13219435513019562, 0.15014058351516724, 0.12075512856245041, 0.0006761215627193451, 0.10174072533845901, 0.19516822695732117, 0.009559075348079205, 0.057678524404764175, 0.08239483833312988, 0.0039215064607560635, 0.0027616096194833517, 0.013109313324093819, 0.002305442001670599, 0.00021083203318994492, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19843007624149323, 0.15979865193367004, 0.14398488402366638, 0.41609427332878113, 0.010126790963113308, 0.04840107262134552, 0.7232485413551331, 0.22829605638980865, 0.34322667121887207, 0.08224418759346008, 0.03167981281876564, 0.020198417827486992, 0.013381149619817734, 0.0009459191933274269, 0.006438484415411949, 0.008794432505965233, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.30347728729248047, 0.04726674035191536, 0.010849116370081902, 0.12094812840223312, 0.0013257962418720126, 0.0025908409152179956, 0.0014983253786340356, 0.03437754884362221, 0.009621781297028065, 0.006184253375977278, 0.00671237800270319, 0.0018636187305673957, 0.01123903226107359, 0.0035993149504065514, 0.0012990115210413933, 0.00021464838937390596, 0.001025065197609365, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2792417109012604, 0.26782968640327454, 0.03489779308438301, 0.07551994919776917, 0.018111348152160645, 0.04002813994884491, 0.03850500285625458, 0.11152958869934082, 0.21995633840560913, 0.07949108630418777, 0.0037619988434016705, 0.03436713665723801, 0.020695386454463005, 0.017524488270282745, 0.010141805745661259, 0.003556826151907444, 0.0020958345849066973, 0.0058519174344837666, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05386974662542343, 0.6086578965187073, 0.22683310508728027, 0.5828835964202881, 0.02668178826570511, 0.03663201630115509, 0.14977867901325226, 0.2173178791999817, 0.2744499444961548, 0.08338183909654617, 0.008825525641441345, 0.06588608771562576, 0.5592238306999207, 0.17532478272914886, 0.006846817210316658, 0.028904464095830917, 0.01721598580479622, 0.006393561605364084, 0.010461881756782532, NaN, NaN, NaN, NaN, NaN, NaN], [0.24167264997959137, 0.2504684031009674, 0.15247754752635956, 0.4417489171028137, 0.37691444158554077, 0.47509273886680603, 0.6227271556854248, 0.6949021220207214, 0.5199849605560303, 0.14203055202960968, 0.006932773161679506, 0.02713918127119541, 0.026524275541305542, 0.28478434681892395, 0.05304509028792381, 0.03063105419278145, 0.007391192018985748, 0.001299944007769227, 0.0022179351653903723, 0.0017378581687808037, NaN, NaN, NaN, NaN, NaN], [0.3587647080421448, 0.13152657449245453, 0.3170546591281891, 0.1872878074645996, 0.17338471114635468, 0.16099165380001068, 0.050314128398895264, 0.07316549867391586, 0.1506616473197937, 0.027928102761507034, 0.013985591009259224, 0.03077181987464428, 0.00928373821079731, 0.01458327379077673, 0.34401679039001465, 0.1675042062997818, 0.008024912327528, 0.00340651860460639, 0.001158604514785111, 0.0004595925274770707, 0.0022153020836412907, NaN, NaN, NaN, NaN], [0.18021628260612488, 0.21554027497768402, 0.22428971529006958, 0.28362634778022766, 0.0019759181886911392, 0.19364571571350098, 0.3129161596298218, 0.05571373924612999, 0.43670228123664856, 0.5364305973052979, 0.045233964920043945, 0.02291695959866047, 0.15668357908725739, 0.03788933902978897, 0.0009749932214617729, 0.15011590719223022, 0.009233620017766953, 0.023490505293011665, 0.0018092861864715815, 0.01433361042290926, 0.002351803006604314, 0.00025271173217333853, NaN, NaN, NaN], [0.18984580039978027, 0.30305740237236023, 0.22004783153533936, 0.5488721132278442, 0.023633448407053947, 0.10360189527273178, 0.8517335653305054, 0.6748489141464233, 0.77315753698349, 0.4876308739185333, 0.2048063576221466, 0.14540305733680725, 0.08473058044910431, 0.012403973378241062, 0.06795734912157059, 0.17164894938468933, 0.18992502987384796, 0.12247806042432785, 0.011528578586876392, 0.009636401198804379, 0.0008312705904245377, 0.013430905528366566, 0.011612125672399998, NaN, NaN], [0.3384567201137543, 0.062264904379844666, 0.014819102361798286, 0.14853152632713318, 0.0019540644716471434, 0.003596463706344366, 0.001872691442258656, 0.11878995597362518, 0.02639206312596798, 0.009769541211426258, 0.011811794713139534, 0.006684192456305027, 0.045877717435359955, 0.019279729574918747, 0.005480214022099972, 0.003932234365493059, 0.006437724456191063, 0.0240105502307415, 0.0011211916571483016, 0.004233745392411947, 0.001469226786866784, 0.0013713098596781492, 0.00014342667418532073, 0.0008160521974787116, NaN], [0.1837155818939209, 0.5941455364227295, 0.2251758873462677, 0.3662757873535156, 0.039659783244132996, 0.3226933479309082, 0.014135366305708885, 0.028798755258321762, 0.10863638669252396, 0.34925851225852966, 0.03930900990962982, 0.08864527195692062, 0.10118203610181808, 0.05801505595445633, 0.11320658773183823, 0.05595846846699715, 0.0026757779996842146, 0.007132661063224077, 0.010286321863532066, 0.015962811186909676, 0.004528969060629606, 0.01888921484351158, 0.004036444239318371, 0.00027040645363740623, 0.0002387895801803097]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12851747870445251, 0.06451001763343811, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16148854792118073, 0.04709945246577263, 0.0016553826862946153, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12575848400592804, 0.13552792370319366, 0.1085570901632309, 0.11512085795402527, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14333586394786835, 0.24668441712856293, 0.19262480735778809, 0.13920731842517853, 0.0020065978169441223, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1578390896320343, 0.19358907639980316, 0.02251395769417286, 0.04702039062976837, 0.018520673736929893, 0.0005939522525295615, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14088943600654602, 0.05360155552625656, 0.043673839420080185, 0.0087194312363863, 0.14876413345336914, 0.3311525881290436, 0.029076436534523964, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11886978894472122, 0.08032860606908798, 0.053777631372213364, 0.06359982490539551, 0.49348562955856323, 0.7690801620483398, 0.032007213681936264, 0.00921344943344593, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.013988303020596504, 0.031309448182582855, 0.021422432735562325, 0.015959911048412323, 0.13852538168430328, 0.7482463121414185, 0.1306946873664856, 0.0026366086676716805, 0.006285007111728191, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02759428508579731, 0.1341203898191452, 0.1143924742937088, 0.04895513132214546, 0.2507959306240082, 0.47495928406715393, 0.24884849786758423, 0.04048554226756096, 0.06435439735651016, 0.02207104302942753, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08376637101173401, 0.08644555509090424, 0.08414626121520996, 0.08246676623821259, 0.09393073618412018, 0.2536129355430603, 0.09570588916540146, 0.057335685938596725, 0.27625876665115356, 0.23640654981136322, 0.22554923593997955, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16592197120189667, 0.037314873188734055, 0.020350072532892227, 0.005164262373000383, 0.009123047813773155, 0.005826999898999929, 0.003451529424637556, 0.017567342147231102, 0.055315494537353516, 0.2317170798778534, 0.05933540314435959, 0.06010079011321068, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07053745537996292, 0.19491763412952423, 0.06705262511968613, 0.08265279233455658, 0.006405644118785858, 0.0031596925109624863, 0.005410268437117338, 0.030676638707518578, 0.08307406306266785, 0.20774710178375244, 0.4213918149471283, 0.23337899148464203, 0.08583765476942062, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13580749928951263, 0.17484943568706512, 0.09017936140298843, 0.11502011120319366, 0.015199831686913967, 0.008567527867853642, 0.04639086127281189, 0.16773870587348938, 0.16907723248004913, 0.43436557054519653, 0.2870768904685974, 0.10786425322294235, 0.08931463956832886, 0.011009148322045803, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1727631837129593, 0.039101891219615936, 0.0065339612774550915, 0.0278339721262455, 0.004674504045397043, 0.014613990671932697, 0.03457005321979523, 0.04850766807794571, 0.02412491664290428, 0.009369020350277424, 0.022906647995114326, 0.04899173229932785, 0.01023520715534687, 0.0022774694953113794, 7.664388976991177e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08213489502668381, 0.3905046880245209, 0.07204636186361313, 0.08312273025512695, 0.02625700645148754, 0.02937941811978817, 0.04131421819329262, 0.05289716273546219, 0.16493423283100128, 0.290347158908844, 0.47713640332221985, 0.44352003931999207, 0.11574649810791016, 0.0847686156630516, 0.047198787331581116, 0.1300322264432907, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.056048911064863205, 0.04177262261509895, 0.18134142458438873, 0.04556399583816528, 0.1435631662607193, 0.2900937497615814, 0.07549438625574112, 0.08105770498514175, 0.08377190679311752, 0.011481991037726402, 0.017289845272898674, 0.006863615941256285, 0.013694294728338718, 0.13657283782958984, 0.0735873132944107, 0.3659329116344452, 0.0919225886464119, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06230737641453743, 0.038521286100149155, 0.05914388969540596, 0.03398321941494942, 0.13657090067863464, 0.19265799224376678, 0.07424072921276093, 0.08660972863435745, 0.10718739032745361, 0.16533604264259338, 0.0767570361495018, 0.03204379230737686, 0.028188396245241165, 0.21943823993206024, 0.11997849494218826, 0.2698959410190582, 0.12308003753423691, 0.45223531126976013, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18667352199554443, 0.0350969135761261, 0.030425790697336197, 0.0065561928786337376, 0.028277983888983727, 0.010725672356784344, 0.005219776649028063, 0.03378060460090637, 0.04241056367754936, 0.18939200043678284, 0.06338198482990265, 0.08136797696352005, 0.004227515775710344, 0.024540461599826813, 0.057830944657325745, 0.038525767624378204, 0.0177453625947237, 0.06933332234621048, 0.08866386860609055, NaN, NaN, NaN, NaN, NaN, NaN], [0.04736897721886635, 0.0950922816991806, 0.05233628675341606, 0.0639958381652832, 0.009022187441587448, 0.002768130972981453, 0.005348078906536102, 0.016458049416542053, 0.03350484371185303, 0.1584910899400711, 0.3849281072616577, 0.30566492676734924, 0.08282434195280075, 0.02534077689051628, 0.01897522434592247, 0.013481524772942066, 0.08136109262704849, 0.25969398021698, 0.2513872981071472, 0.07361149042844772, NaN, NaN, NaN, NaN, NaN], [0.15279658138751984, 0.09928575158119202, 0.0573631152510643, 0.10790141671895981, 0.026906443759799004, 0.012519991025328636, 0.06774256378412247, 0.1448669582605362, 0.07826853543519974, 0.4991803467273712, 0.34429702162742615, 0.12145370990037918, 0.10719165205955505, 0.008088642731308937, 0.007662023417651653, 0.013441860675811768, 0.13362208008766174, 0.34251537919044495, 0.10342243313789368, 0.07045409828424454, 0.010391364805400372, NaN, NaN, NaN, NaN], [0.1865139603614807, 0.02971193566918373, 0.005512321833521128, 0.039164237678050995, 0.007472363766282797, 0.012969624251127243, 0.03476016968488693, 0.0836154893040657, 0.050758667290210724, 0.017821883782744408, 0.08676476776599884, 0.13045690953731537, 0.03245873004198074, 0.009119128808379173, 7.800521416356787e-05, 0.0006276130443438888, 0.0024839011020958424, 0.06682475656270981, 0.06347990781068802, 0.009879485704004765, 0.0017003080574795604, 6.444661266868934e-05, NaN, NaN, NaN], [0.029208103194832802, 0.15452517569065094, 0.02615012601017952, 0.034968301653862, 0.030517179518938065, 0.023491270840168, 0.02012590691447258, 0.01683984510600567, 0.047155413776636124, 0.1569623053073883, 0.34555378556251526, 0.29876279830932617, 0.06633269041776657, 0.090775266289711, 0.05117363482713699, 0.14964616298675537, 0.024973956868052483, 0.22028914093971252, 0.5953715443611145, 0.10930891335010529, 0.05826140195131302, 0.08348876982927322, 0.2024080604314804, NaN, NaN], [0.023966457694768906, 0.008770916610956192, 0.0534873865544796, 0.015555462799966335, 0.07408829033374786, 0.12750747799873352, 0.026930494233965874, 0.023400133475661278, 0.02665247581899166, 0.00316479685716331, 0.004739005118608475, 0.002742160577327013, 0.006070322822779417, 0.09564805775880814, 0.029174519702792168, 0.5144217014312744, 0.05911846086382866, 0.020064763724803925, 0.0023497287184000015, 0.004584830719977617, 0.10225256532430649, 0.05520752817392349, 0.4466201066970825, 0.09660884737968445, NaN], [0.18986307084560394, 0.036011889576911926, 0.08335232734680176, 0.12826237082481384, 0.08758756518363953, 0.027860891073942184, 0.10198243707418442, 0.0981309786438942, 0.17985263466835022, 0.11864234507083893, 0.08274368196725845, 0.1066904067993164, 0.051979877054691315, 0.06548189371824265, 0.03337343409657478, 0.0824524462223053, 0.012718076817691326, 0.0349668525159359, 0.03024965338408947, 0.01082769688218832, 0.0127665214240551, 0.014164488762617111, 0.01925024762749672, 0.0028478982858359814, 0.0007362329051829875]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12774905562400818, 0.07772441953420639, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.058547187596559525, 0.7868303656578064, 0.02677525207400322, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12958122789859772, 0.05996095389127731, 0.20109553635120392, 0.07473170012235641, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11586850136518478, 0.18037959933280945, 0.354478657245636, 0.6275972127914429, 0.01217791810631752, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04329086095094681, 0.2822243273258209, 0.5110569596290588, 0.8230794668197632, 0.28263914585113525, 0.006951561663299799, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15041278302669525, 0.01652364432811737, 0.09004879742860794, 0.1228649914264679, 0.03705046698451042, 0.03279988467693329, 0.012472960166633129, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.005692727863788605, 0.004583822097629309, 0.011303454637527466, 0.06351188570261002, 0.07110948860645294, 0.03377191722393036, 0.8937738537788391, 0.1077374666929245, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1957636922597885, 0.00532554043456912, 0.2672942280769348, 0.07843183726072311, 0.01169322058558464, 0.006695515010505915, 0.022856300696730614, 0.03495524823665619, 0.2056257426738739, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21948350965976715, 0.003219911362975836, 0.13064762949943542, 0.017335020005702972, 0.004487968049943447, 0.006097455509006977, 0.0023269150406122208, 0.014221499674022198, 0.1740167737007141, 0.05570632219314575, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.027252521365880966, 0.05625513195991516, 0.024279700592160225, 0.009296371601521969, 0.04113621264696121, 0.04445572942495346, 0.05016031116247177, 0.300394743680954, 0.219209223985672, 0.5284181833267212, 0.13528388738632202, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16918426752090454, 0.005196947604417801, 0.010393726639449596, 0.0008839815272949636, 0.18853645026683807, 0.23955073952674866, 0.03703731670975685, 0.018581384792923927, 0.07692746073007584, 0.05213537812232971, 0.05520249530673027, 0.03837481513619423, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21910618245601654, 0.012340836226940155, 0.011061819270253181, 0.004421355202794075, 0.01345156505703926, 0.015948239713907242, 0.001919197733514011, 0.0006712953327223659, 0.0014401280786842108, 0.0009498890140093863, 0.0011606297921389341, 0.0013843519845977426, 0.005138876382261515, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12592341005802155, 0.022789308801293373, 0.01544136367738247, 0.05098855495452881, 0.006733328104019165, 0.0011512627825140953, 0.0067494111135602, 0.03519098460674286, 0.08756479620933533, 0.04847756400704384, 0.13774195313453674, 0.07365753501653671, 0.19525301456451416, 0.019442297518253326, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04374772310256958, 0.10635814815759659, 0.1203576922416687, 0.4972172677516937, 0.09716533124446869, 0.05867829546332359, 0.13453392684459686, 0.39353471994400024, 0.6331138610839844, 0.33491814136505127, 0.5983138680458069, 0.3633559048175812, 0.6357010006904602, 0.7792285084724426, 0.005659972317516804, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05199728533625603, 0.014302223920822144, 0.13574257493019104, 0.05407930538058281, 0.010633953846991062, 0.007459194865077734, 0.0004102779785171151, 0.01107444055378437, 0.16451390087604523, 0.19313758611679077, 0.018386593088507652, 0.03492085263133049, 0.1390746384859085, 0.6526300311088562, 0.08304706960916519, 0.27643677592277527, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0008206118363887072, 0.0011099595576524734, 0.0005428412696346641, 0.0013029578840360045, 0.0009422241128049791, 0.001036918954923749, 0.00015340711979661137, 0.003300317795947194, 0.0019372785463929176, 0.003245894331485033, 0.0010756017873063684, 0.0009867959888651967, 0.04242069274187088, 0.25679609179496765, 0.03714281693100929, 0.46563825011253357, 0.052469443529844284, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0011551693314686418, 0.0015016108518466353, 0.00018865184392780066, 0.0004620797117240727, 0.001353209256194532, 0.001276124152354896, 0.001269699539989233, 0.02504812367260456, 0.016660472378134727, 0.007664685603231192, 0.000621759332716465, 0.0039494638331234455, 0.05373308062553406, 0.5797222256660461, 0.04267296567559242, 0.3308492600917816, 0.22605444490909576, 0.03655111417174339, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18345873057842255, 0.006115049123764038, 0.007153322920203209, 0.00125643250066787, 0.15791349112987518, 0.17755654454231262, 0.06167090684175491, 0.028255566954612732, 0.04990806803107262, 0.014394938945770264, 0.013118196278810501, 0.02539716847240925, 0.00894339382648468, 0.04024626687169075, 0.05642623454332352, 0.04561464861035347, 0.029457826167345047, 0.09210912138223648, 0.1002524197101593, NaN, NaN, NaN, NaN, NaN, NaN], [0.2828649580478668, 0.011994204483926296, 0.006339475512504578, 0.0030444697476923466, 0.006948052905499935, 0.008767204359173775, 0.0014567734906449914, 0.00018795454525388777, 0.00020330831466708332, 0.0001539710647193715, 0.0004007722018286586, 0.0012242270167917013, 0.001961026806384325, 0.0007920600473880768, 0.002005743095651269, 0.00011892847396666184, 0.00023868663993198425, 0.0018499011639505625, 0.002196513582020998, 0.004604275804013014, NaN, NaN, NaN, NaN, NaN], [0.128562331199646, 0.014782274141907692, 0.007007280830293894, 0.02549830637872219, 0.0029198189731687307, 0.0006880113505758345, 0.0037798655685037374, 0.009390356950461864, 0.008127862587571144, 0.00817851535975933, 0.024966517463326454, 0.0308842696249485, 0.07813727855682373, 0.003280356992036104, 0.001509596244432032, 0.010023933835327625, 0.08412036299705505, 0.1339937299489975, 0.13076454401016235, 0.2572615444660187, 0.02603374607861042, NaN, NaN, NaN, NaN], [0.018602287396788597, 0.034721970558166504, 0.034974802285432816, 0.21532808244228363, 0.037075310945510864, 0.013384592719376087, 0.039282385259866714, 0.11046459525823593, 0.17542847990989685, 0.05914776027202606, 0.1884417086839676, 0.12911023199558258, 0.24417443573474884, 0.327198326587677, 0.0006843891460448503, 0.1527024656534195, 0.4776603579521179, 0.37270504236221313, 0.4335513412952423, 0.6841917634010315, 0.8031085133552551, 0.004920803010463715, NaN, NaN, NaN], [0.05855157971382141, 0.021276630461215973, 0.13662834465503693, 0.05244326964020729, 0.015041220933198929, 0.007642571348696947, 0.00036013865610584617, 0.004098850768059492, 0.033856965601444244, 0.05778159946203232, 0.005442364141345024, 0.017580043524503708, 0.04633626714348793, 0.3112163841724396, 0.03644357994198799, 0.0868009626865387, 0.020123973488807678, 0.03773906081914902, 0.06257405877113342, 0.2619801461696625, 0.7497928738594055, 0.19582624733448029, 0.4370352327823639, NaN, NaN], [0.0006882869056425989, 0.0005033394554629922, 0.00030677669565193355, 0.001028614118695259, 0.00036578672006726265, 0.0005035633221268654, 5.2447539928834885e-05, 0.0006442382582463324, 0.0003597578906919807, 0.0002600657753646374, 8.536354289390147e-05, 0.00018848010222427547, 0.00940172839909792, 0.03475101292133331, 0.004768407437950373, 0.09523987770080566, 0.0036924693267792463, 0.0034024319611489773, 0.001987446565181017, 0.06484154611825943, 0.36614781618118286, 0.06470755487680435, 0.48020803928375244, 0.12385622411966324, NaN], [0.13044977188110352, 0.023216107860207558, 0.019304566085338593, 0.018173998221755028, 0.12614674866199493, 0.04656239226460457, 0.015089727938175201, 0.04114385321736336, 0.018700774759054184, 0.020505733788013458, 0.009310846216976643, 0.02222343534231186, 0.22412429749965668, 0.3900958001613617, 0.1100122332572937, 0.14125461876392365, 0.09716113656759262, 0.14588865637779236, 0.12185929715633392, 0.5472521185874939, 0.7197717428207397, 0.31834876537323, 0.37092098593711853, 0.2838878929615021, 0.0011011400492861867]]], [[[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16810710728168488, 0.017288343980908394, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12647151947021484, 0.25301796197891235, 0.03169602155685425, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15976493060588837, 0.03159531578421593, 0.05609510838985443, 0.007400199305266142, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16021955013275146, 0.26433131098747253, 0.07329617440700531, 0.11257290840148926, 0.001577433431521058, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.22870834171772003, 0.043985288590192795, 0.04075293987989426, 0.0035545979626476765, 0.0075324228964746, 0.00014864112017676234, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.047688793390989304, 0.14664201438426971, 0.03658692538738251, 0.6408759355545044, 0.43873438239097595, 0.20478755235671997, 0.00511742290109396, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07761336117982864, 0.07061085104942322, 0.041570939123630524, 0.1916733682155609, 0.159084752202034, 0.3477410674095154, 0.5968326330184937, 0.004175147507339716, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07191380113363266, 0.05497179180383682, 0.3517811894416809, 0.9035707116127014, 0.14233137667179108, 0.1767667979001999, 0.04289708659052849, 0.00892895832657814, 0.001834895578213036, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21536989510059357, 0.19956108927726746, 0.3517906069755554, 0.458966463804245, 0.09842110425233841, 0.08277469873428345, 0.03296331316232681, 0.04812879115343094, 0.009344152174890041, 0.006280441302806139, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24051256477832794, 0.10134825110435486, 0.04672827199101448, 0.021085558459162712, 0.02245912328362465, 0.026835136115550995, 0.005604758393019438, 0.028772464022040367, 0.01708872988820076, 0.008745603263378143, 0.02540087327361107, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18141932785511017, 0.024432087317109108, 0.0408032201230526, 0.004596539307385683, 0.0778040885925293, 0.025828123092651367, 0.04467899724841118, 0.0885351300239563, 0.026468785479664803, 0.030213410034775734, 0.16925157606601715, 0.003915028180927038, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0821177139878273, 0.0264634620398283, 0.01841210387647152, 0.010007970035076141, 0.006691556889563799, 0.0167625043541193, 0.0005595253896899521, 0.020632673054933548, 0.0021230748388916254, 0.10790054500102997, 0.5654488801956177, 0.3003200888633728, 0.01571945659816265, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0726943239569664, 0.09770844131708145, 0.050709616392850876, 0.04594658315181732, 0.009083828888833523, 0.024983327835798264, 0.021837929263710976, 0.11926575750112534, 0.11382617056369781, 0.22249171137809753, 0.3826439678668976, 0.22458447515964508, 0.24531354010105133, 0.05176876112818718, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.28158777952194214, 0.045097555965185165, 0.02117414027452469, 0.05809389799833298, 0.0014524150174111128, 0.006964406464248896, 0.010582090355455875, 0.011965163983404636, 0.02265000529587269, 0.020484870299696922, 0.019729144871234894, 0.028731632977724075, 0.004907289054244757, 0.0051048253662884235, 0.00039794077747501433, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18024474382400513, 0.03336771950125694, 0.025161737576127052, 0.03788529708981514, 0.010167604312300682, 0.0039537386037409306, 3.701886089402251e-05, 0.046124417334795, 0.08654022216796875, 0.06664562225341797, 0.11276466399431229, 0.09791301190853119, 0.08758807182312012, 0.277656227350235, 0.5478507876396179, 0.06896418333053589, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10793236643075943, 0.04864804446697235, 0.0019557650666683912, 0.14817607402801514, 0.0378977507352829, 0.049347102642059326, 0.0036467635072767735, 0.0038541490212082863, 0.0034904496278613806, 0.0012115711579099298, 0.047197386622428894, 0.05697714909911156, 0.11328870058059692, 0.8784908056259155, 0.019691603258252144, 0.23420120775699615, 0.004765921737998724, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1524984985589981, 0.08107080310583115, 0.005865868646651506, 0.00971321389079094, 0.007243088912218809, 0.011549782939255238, 0.00268083019182086, 0.03457775339484215, 0.0031127233523875475, 0.000510410696733743, 0.009807620197534561, 0.008875550702214241, 0.023541534319519997, 0.527433454990387, 0.015368063934147358, 0.16288210451602936, 0.20708848536014557, 0.014573587104678154, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16305263340473175, 0.020936982706189156, 0.020989498123526573, 0.007437185384333134, 0.034894589334726334, 0.016221558675169945, 0.04928300529718399, 0.02460765466094017, 0.006940784398466349, 0.010303718037903309, 0.11923910677433014, 0.002430608496069908, 0.020191287621855736, 0.019723495468497276, 0.015607062727212906, 0.14493703842163086, 0.29023703932762146, 0.2954525649547577, 0.024419967085123062, NaN, NaN, NaN, NaN, NaN, NaN], [0.04235544800758362, 0.014461617916822433, 0.006770138628780842, 0.009241613559424877, 0.002999901305884123, 0.0037356300745159388, 0.00043396188993938267, 0.005936506669968367, 0.00027135247364640236, 0.00836905650794506, 0.38652852177619934, 0.1805782914161682, 0.00859912484884262, 0.13720881938934326, 0.026457296684384346, 0.044793374836444855, 0.41905051469802856, 0.48846107721328735, 0.271888792514801, 0.02787640690803528, NaN, NaN, NaN, NaN, NaN], [0.03824670985341072, 0.05110237002372742, 0.016365332528948784, 0.027689939364790916, 0.004054062534123659, 0.0016762956511229277, 0.0059990487061440945, 0.061629924923181534, 0.02193543128669262, 0.004144957754760981, 0.11336920410394669, 0.0855039581656456, 0.16943661868572235, 0.007511935196816921, 0.0029296777211129665, 0.005633122753351927, 0.04470856487751007, 0.19621509313583374, 0.1449754536151886, 0.4407651424407959, 0.012849990278482437, NaN, NaN, NaN, NaN], [0.29710885882377625, 0.04157622903585434, 0.022785142064094543, 0.06820578873157501, 0.0019051277777180076, 0.004196317866444588, 0.012664434500038624, 0.010533612221479416, 0.00958634540438652, 0.006948783528059721, 0.024731770157814026, 0.04424457997083664, 0.0092665059491992, 0.008317369967699051, 0.00025302590802311897, 0.03921425715088844, 0.024433301761746407, 0.005475904326885939, 0.02041386440396309, 0.005526822991669178, 0.006030899006873369, 0.000147900907904841, NaN, NaN, NaN], [0.15116539597511292, 0.029300624504685402, 0.014213098213076591, 0.04858435317873955, 0.008192096836864948, 0.0029929669108241796, 0.00010039177868748084, 0.02851700410246849, 0.014845605008304119, 0.01335279829800129, 0.07330357283353806, 0.08230004459619522, 0.06801280379295349, 0.12962418794631958, 0.38807213306427, 0.021973537281155586, 0.0005578201962634921, 0.13413770496845245, 0.18835364282131195, 0.15109674632549286, 0.5815849900245667, 0.6008182764053345, 0.10515720397233963, NaN, NaN], [0.05911188945174217, 0.013889956288039684, 0.00048160224105231464, 0.10393460839986801, 0.009916743263602257, 0.013972792774438858, 0.0005543273873627186, 0.0008135904208756983, 0.0005866698920726776, 0.00012856724788434803, 0.016669562086462975, 0.022332170978188515, 0.03126570209860802, 0.39481881260871887, 0.0021035531535744667, 0.09696949273347855, 0.0003469766234047711, 0.012058700434863567, 0.1351245492696762, 0.1276140809059143, 0.8529128432273865, 0.013427066616714, 0.3029053509235382, 0.0016288348706439137, NaN], [0.22241219878196716, 0.00997188687324524, 0.004307668190449476, 0.0318865031003952, 0.026490027084946632, 0.04937301576137543, 0.016565896570682526, 0.0013930558925494552, 0.01958940364420414, 0.015218929387629032, 0.1830211728811264, 0.11458480358123779, 0.1729872077703476, 0.047152113169431686, 0.017883911728858948, 0.118315190076828, 0.07728181034326553, 0.31889867782592773, 0.1497264951467514, 0.2596881091594696, 0.15263305604457855, 0.024473916739225388, 0.19167250394821167, 0.12363447993993759, 0.010316992178559303]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12855423986911774, 0.11611904203891754, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1812644749879837, 0.04049589857459068, 0.04480821266770363, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14001408219337463, 0.11702272295951843, 0.5616602897644043, 0.021032487973570824, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17309650778770447, 0.011261633597314358, 0.0023054813500493765, 0.0014516497030854225, 0.17103753983974457, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21775518357753754, 0.1599237471818924, 0.031671781092882156, 0.0027859890833497047, 0.1030324175953865, 0.009803196415305138, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1265520304441452, 0.2245447188615799, 0.3357183039188385, 0.19591355323791504, 0.030100535601377487, 0.11038237810134888, 0.012957160361111164, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12113019824028015, 0.07331034541130066, 0.073086217045784, 0.038516201078891754, 0.16168329119682312, 0.12152494490146637, 0.1929183006286621, 0.11648087203502655, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15162895619869232, 0.16000056266784668, 0.47010278701782227, 0.008242717012763023, 0.016423694789409637, 0.19619418680667877, 0.014187236316502094, 0.2187093049287796, 0.3917299807071686, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1371021270751953, 0.24055053293704987, 0.39826682209968567, 0.0653936043381691, 0.06886317580938339, 0.1729464828968048, 0.02453671395778656, 0.2748231589794159, 0.23215962946414948, 0.03306089714169502, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05615014582872391, 0.17226241528987885, 0.4426397681236267, 0.534454345703125, 0.0034056571312248707, 0.0038566330913454294, 0.24011781811714172, 0.31882721185684204, 0.4456172287464142, 0.1489524245262146, 0.03087311051785946, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.037336766719818115, 0.065662682056427, 0.18869149684906006, 0.795316219329834, 0.14649540185928345, 0.021824514493346214, 0.13452036678791046, 0.026823654770851135, 0.35548609495162964, 0.18523786962032318, 0.020790524780750275, 0.09485815465450287, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17983746528625488, 0.09746579825878143, 0.46259593963623047, 0.706605851650238, 0.09193093329668045, 0.2823830544948578, 0.007526541594415903, 0.10234087705612183, 0.24847157299518585, 0.2038285881280899, 0.012590465135872364, 0.002493936335667968, 0.04428662359714508, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1421777307987213, 0.23310348391532898, 0.2705342471599579, 0.5351002812385559, 0.02795390971004963, 0.06031421944499016, 0.012775074690580368, 0.20022329688072205, 0.6570897698402405, 0.2668534517288208, 0.033325545489788055, 0.023841219022870064, 0.1455993354320526, 0.03172359615564346, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11665362864732742, 0.1886645257472992, 0.03897944837808609, 0.07137740403413773, 0.15634050965309143, 0.15400150418281555, 0.13745756447315216, 0.05537642911076546, 0.2729690372943878, 0.04749782383441925, 0.05948880687355995, 0.014797642827033997, 0.11365658044815063, 0.002582019427791238, 0.20324750244617462, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.29635345935821533, 0.04781435802578926, 0.41243496537208557, 0.03004680573940277, 0.13952067494392395, 0.045467544347047806, 4.634694050764665e-05, 0.20948387682437897, 0.002634957665577531, 0.005124728661030531, 0.0019075855379924178, 0.0009838729165494442, 0.0013485344825312495, 0.004148871172219515, 0.03574635088443756, 0.23113909363746643, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.22071197628974915, 0.019423967227339745, 0.06694509834051132, 0.2386176735162735, 0.015943216159939766, 0.14270655810832977, 0.039743710309267044, 0.014324809424579144, 0.581375777721405, 0.040944233536720276, 0.011615565046668053, 0.02482481673359871, 0.06486763060092926, 0.002298883395269513, 0.009274494834244251, 0.012798607349395752, 0.009606687352061272, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04979729279875755, 0.005993144121021032, 0.05621323734521866, 0.3196869492530823, 0.0036542851012200117, 0.006608159281313419, 0.07202935218811035, 0.023804083466529846, 0.08581908792257309, 0.002907529706135392, 0.0022882334887981415, 0.155064657330513, 0.6752456426620483, 0.19066885113716125, 0.033486951142549515, 0.1545412391424179, 0.3257397711277008, 0.07836033403873444, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02027127519249916, 0.036089565604925156, 0.0908525288105011, 0.6094546914100647, 0.035198476165533066, 0.01578100211918354, 0.08828305453062057, 0.00740778585895896, 0.08938029408454895, 0.055872198194265366, 0.01406459603458643, 0.05842210724949837, 0.7085317969322205, 0.04043729975819588, 0.00861792266368866, 0.05839632451534271, 0.306302547454834, 0.11257344484329224, 0.09490343183279037, NaN, NaN, NaN, NaN, NaN, NaN], [0.2219613641500473, 0.0726998969912529, 0.3657586872577667, 0.6172192692756653, 0.07194076478481293, 0.17607101798057556, 0.009873087517917156, 0.09032700955867767, 0.1240842267870903, 0.06592906266450882, 0.021971723064780235, 0.004476875066757202, 0.04292584955692291, 0.013240871019661427, 0.03868407383561134, 0.0364602766931057, 0.007298360578715801, 0.02817610278725624, 0.0009550384129397571, 0.033005379140377045, NaN, NaN, NaN, NaN, NaN], [0.2832254469394684, 0.40537261962890625, 0.25111812353134155, 0.4335843026638031, 0.05173255130648613, 0.02949104830622673, 0.00834138598293066, 0.5043417811393738, 0.45271721482276917, 0.10732957720756531, 0.08741836994886398, 0.06616821885108948, 0.1252485066652298, 0.04288535565137863, 0.0027607728261500597, 0.11496254801750183, 0.007436650805175304, 0.04789961501955986, 0.014611729420721531, 0.05419020354747772, 0.013982507400214672, NaN, NaN, NaN, NaN], [0.1133793368935585, 0.2190774381160736, 0.04727642610669136, 0.08785698562860489, 0.22799502313137054, 0.1395695060491562, 0.17899513244628906, 0.05776361748576164, 0.19579172134399414, 0.03426501154899597, 0.08577524870634079, 0.027239171788096428, 0.22711482644081116, 0.005856664851307869, 0.3394412696361542, 0.03666312247514725, 0.053877539932727814, 0.02460121363401413, 0.02095765992999077, 0.08733106404542923, 0.0007995758787728846, 0.19509249925613403, NaN, NaN, NaN], [0.32134389877319336, 0.08582156896591187, 0.36053547263145447, 0.06279635429382324, 0.1449708491563797, 0.041098933666944504, 0.0002254477294627577, 0.3326246738433838, 0.0031729326583445072, 0.011426791548728943, 0.00305219367146492, 0.0021134610287845135, 0.0029090954922139645, 0.0035086346324533224, 0.0884322077035904, 0.7275413274765015, 4.6366836613742635e-05, 0.004567307885736227, 0.00048746803076937795, 0.0006845259922556579, 0.00036436106893233955, 0.0336419902741909, 0.19370199739933014, NaN, NaN], [0.2431764006614685, 0.00993723887950182, 0.023469794541597366, 0.12711890041828156, 0.013049022294580936, 0.09880916029214859, 0.014819139614701271, 0.015189954079687595, 0.19677633047103882, 0.012298321351408958, 0.006653454154729843, 0.017306946218013763, 0.044382814317941666, 0.005554118659347296, 0.008197239600121975, 0.025704391300678253, 0.01238576602190733, 0.005520223639905453, 0.018611198291182518, 0.07344726473093033, 0.00026948421145789325, 0.012129159644246101, 0.01222553662955761, 0.005697384011000395, NaN], [0.018590128049254417, 0.012204503640532494, 0.0029425490647554398, 0.01610950194299221, 0.024503106251358986, 0.04006015509366989, 0.018976394087076187, 0.006591797806322575, 0.002320006489753723, 0.001339062349870801, 0.028667215257883072, 0.03959575667977333, 0.00960585381835699, 0.009797154925763607, 0.022796805948019028, 0.1637655347585678, 0.20084494352340698, 0.05620957538485527, 0.12549559772014618, 0.022888751700520515, 0.037492163479328156, 0.04711981862783432, 0.44462573528289795, 0.3949664235115051, 0.3300856053829193]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16815106570720673, 0.017178548499941826, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2022658735513687, 0.005017802584916353, 0.01763225719332695, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16166983544826508, 0.033678483217954636, 0.014520054683089256, 0.003462842432782054, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10712886601686478, 0.3422684967517853, 0.05748933553695679, 0.2768969237804413, 0.004922540858387947, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.047501806169748306, 0.48201972246170044, 0.4827657639980316, 0.48466482758522034, 0.022285524755716324, 0.00022009640815667808, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1517350822687149, 0.04445230960845947, 0.09343461692333221, 0.05873756855726242, 0.07171032577753067, 0.22849556803703308, 0.05614512786269188, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.25680339336395264, 0.00010820403986144811, 0.0123103903606534, 0.007049524690955877, 0.001952940714545548, 0.027401963248848915, 0.0028134624008089304, 0.00041907382546924055, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.005559808574616909, 0.007462772540748119, 0.013313480652868748, 0.017376750707626343, 0.0038542840629816055, 0.006728595122694969, 0.5333897471427917, 0.03155524656176567, 0.15571120381355286, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.004124458413571119, 0.004751718603074551, 0.016015900298953056, 0.01742120459675789, 0.032125748693943024, 0.010460411198437214, 0.45809611678123474, 0.07138781994581223, 0.5171095728874207, 0.17626723647117615, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.24881334602832794, 0.005821824539452791, 0.031170587986707687, 0.009853766299784184, 0.027254868298768997, 0.01885347068309784, 0.02900754101574421, 0.013663586229085922, 0.012090054340660572, 0.0009272377355955541, 0.0030740045476704836, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19627800583839417, 0.054823894053697586, 0.1886557787656784, 0.00739922234788537, 0.09451853483915329, 0.01572227105498314, 0.0010023268405348063, 0.0061036646366119385, 0.0014733865391463041, 0.0003654434985946864, 0.006776102818548679, 0.0027319795917719603, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07900664210319519, 0.04510375112295151, 0.002657376928254962, 0.0032053724862635136, 0.0027717212215065956, 0.008140889927744865, 0.0011833005119115114, 0.04105996713042259, 0.0017470002640038729, 0.008194361813366413, 0.019470002502202988, 0.3834601640701294, 0.013146632350981236, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06578069925308228, 0.08975866436958313, 0.022234706208109856, 0.015388325788080692, 0.006578383035957813, 0.011582762002944946, 0.014906905591487885, 0.04645423963665962, 0.008417387492954731, 0.0318351611495018, 0.024524353444576263, 0.5050408244132996, 0.1078883558511734, 0.09876319766044617, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.010224410332739353, 0.16048979759216309, 0.09242240339517593, 0.259725958108902, 0.06779038906097412, 0.007232773117721081, 0.09601377695798874, 0.28109633922576904, 0.2723717987537384, 0.1275584101676941, 0.06318827718496323, 0.25179460644721985, 0.2496732771396637, 0.6837621927261353, 0.0018262360244989395, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04991341754794121, 0.05319196358323097, 0.14821480214595795, 0.020963814109563828, 0.03095317631959915, 0.024693654850125313, 0.008621936663985252, 0.14259999990463257, 0.042305052280426025, 0.09002435952425003, 0.005839803721755743, 0.061309609562158585, 0.23589004576206207, 0.30903181433677673, 0.18008928000926971, 0.49815359711647034, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.015294999815523624, 0.03185835853219032, 0.0202027577906847, 0.03976168856024742, 0.0711589902639389, 0.13473857939243317, 0.0059967683628201485, 0.0031582280062139034, 0.003374348394572735, 0.002362155122682452, 0.015532899647951126, 0.038825590163469315, 0.08611883223056793, 0.03844507411122322, 0.009673628956079483, 0.7068554162979126, 0.013729983940720558, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2531464695930481, 0.013071080669760704, 0.035546887665987015, 0.020458703860640526, 0.01740572415292263, 0.009577612392604351, 0.014396607875823975, 0.05952044576406479, 0.013841827400028706, 0.0003843819722533226, 0.0024746267590671778, 0.007157978601753712, 0.013787134550511837, 0.033782534301280975, 0.003469215938821435, 0.007898973301053047, 0.05525756999850273, 0.003914556000381708, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20273520052433014, 0.05025332421064377, 0.2335304319858551, 0.009442931972444057, 0.13508503139019012, 0.0181263517588377, 0.0010557285277172923, 0.003822105238214135, 0.0018545370548963547, 0.0003744752029888332, 0.0046313730999827385, 0.0008518796530552208, 0.006319030188024044, 0.014203540980815887, 0.0018540708115324378, 0.003058186499401927, 0.002516325796023011, 0.001575352856889367, 0.0014869269216433167, NaN, NaN, NaN, NaN, NaN, NaN], [0.059709664434194565, 0.021975213661789894, 0.002582199638709426, 0.002308695577085018, 0.00240446999669075, 0.004605048336088657, 0.0013587460853159428, 0.04497997462749481, 0.0009150391560979187, 0.0030208472162485123, 0.016492530703544617, 0.2572183907032013, 0.006429646629840136, 0.013558420352637768, 0.06110598146915436, 0.03728436306118965, 0.019318275153636932, 0.03907725587487221, 0.4492114782333374, 0.01579420454800129, NaN, NaN, NaN, NaN, NaN], [0.025836847722530365, 0.04185229912400246, 0.017175624147057533, 0.005038154777139425, 0.006518983747810125, 0.0043221269734203815, 0.004393702372908592, 0.03134007006883621, 0.002082354621961713, 0.00246719503775239, 0.00855192355811596, 0.28023120760917664, 0.0558621920645237, 0.020582975819706917, 0.00264686718583107, 0.052114877849817276, 0.01051351334899664, 0.0282430537045002, 0.640393853187561, 0.11605942994356155, 0.042242906987667084, NaN, NaN, NaN, NaN], [0.00790853425860405, 0.07249781489372253, 0.09275110065937042, 0.13612288236618042, 0.0654025748372078, 0.0028184219263494015, 0.039562828838825226, 0.11378230899572372, 0.08281006664037704, 0.029445864260196686, 0.03387679159641266, 0.16786670684814453, 0.2288694977760315, 0.6801032423973083, 0.0008468713494949043, 0.32477572560310364, 0.20243169367313385, 0.04291461780667305, 0.2565927505493164, 0.2435160130262375, 0.8255255222320557, 0.0008029205491766334, NaN, NaN, NaN], [0.06791312247514725, 0.034157127141952515, 0.26634278893470764, 0.01933334954082966, 0.08246968686580658, 0.03419587388634682, 0.019395295530557632, 0.1259232461452484, 0.02923283353447914, 0.07644251734018326, 0.00482177222147584, 0.03381035849452019, 0.2429695725440979, 0.4201262295246124, 0.21319957077503204, 0.1469077318906784, 0.005101305432617664, 0.05322602018713951, 0.08754345029592514, 0.4596864581108093, 0.32625797390937805, 0.2286616712808609, 0.6285872459411621, NaN, NaN], [0.0236026793718338, 0.032931454479694366, 0.018642868846654892, 0.052601076662540436, 0.09147398918867111, 0.11555580049753189, 0.00512799434363842, 0.006684163119643927, 0.005264784675091505, 0.0023014512844383717, 0.005628940649330616, 0.03778252378106117, 0.09737572073936462, 0.12753169238567352, 0.00698094442486763, 0.6853439807891846, 0.02319822832942009, 0.018658116459846497, 0.08199534565210342, 0.18709556758403778, 0.07321563363075256, 0.027500100433826447, 0.6534799337387085, 0.01572287082672119, NaN], [0.24674107134342194, 0.007728901691734791, 0.010779940523207188, 0.01413859985768795, 0.08573849499225616, 0.014258946292102337, 0.014431791380047798, 0.00199147523380816, 0.006254997570067644, 0.003036148613318801, 0.015209752134978771, 0.015118316747248173, 0.05811062082648277, 0.01987045258283615, 0.012226228602230549, 0.021392136812210083, 0.08141177892684937, 0.016042163595557213, 0.01565614528954029, 0.05352389067411423, 0.01607833430171013, 0.014641694724559784, 0.020306598395109177, 0.06722531467676163, 0.005379782523959875]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.147435262799263, 0.06894105672836304, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18660759925842285, 0.013697005808353424, 0.050341442227363586, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14907698333263397, 0.12682567536830902, 0.14014844596385956, 0.024977339431643486, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20074230432510376, 0.11179281026124954, 0.012457489967346191, 0.01455892063677311, 0.011106430552899837, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20768699049949646, 0.16985096037387848, 0.19526726007461548, 0.016829432919621468, 0.05647609382867813, 0.022808711975812912, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14349573850631714, 0.41078659892082214, 0.5100967288017273, 0.04046756774187088, 0.2924310266971588, 0.07987978309392929, 0.007180717773735523, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11146429926156998, 0.3579395115375519, 0.7730652093887329, 0.5723751783370972, 0.2817910611629486, 0.25461745262145996, 0.060240793973207474, 0.08399515599012375, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13904383778572083, 0.44345301389694214, 0.1345542073249817, 0.05706587806344032, 0.7818705439567566, 0.04436418041586876, 0.015915511175990105, 0.31926584243774414, 0.26167550683021545, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12236351519823074, 0.40148651599884033, 0.12099923938512802, 0.38539087772369385, 0.6352627873420715, 0.0574735552072525, 0.027495326474308968, 0.25199854373931885, 0.07788273692131042, 0.1824284791946411, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0776049941778183, 0.26076433062553406, 0.12800094485282898, 0.15216867625713348, 0.36678510904312134, 0.31404268741607666, 0.13151897490024567, 0.1709745228290558, 0.2591820955276489, 0.18929390609264374, 0.08235450834035873, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08287283033132553, 0.26698997616767883, 0.29562729597091675, 0.13922370970249176, 0.3693794012069702, 0.22139106690883636, 0.612119734287262, 0.1618482619524002, 0.40734153985977173, 0.10604425519704819, 0.2217203825712204, 0.14197519421577454, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0676846131682396, 0.5803259611129761, 0.47128230333328247, 0.2430339902639389, 0.43893957138061523, 0.5822793245315552, 0.9563859105110168, 0.5092246532440186, 0.7397804260253906, 0.6675750613212585, 0.2242172360420227, 0.046741336584091187, 0.09371624141931534, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16273218393325806, 0.4245251417160034, 0.44257473945617676, 0.1064363345503807, 0.22264361381530762, 0.638583779335022, 0.7456080913543701, 0.17856015264987946, 0.09681503474712372, 0.3901955187320709, 0.4154786765575409, 0.10903800278902054, 0.0281606987118721, 0.027353502810001373, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2541956901550293, 0.2554672658443451, 0.13483673334121704, 0.33163735270500183, 0.11067650467157364, 0.3400806486606598, 0.4272999167442322, 0.2955835163593292, 0.293487548828125, 0.2820315957069397, 0.17141510546207428, 0.08369391411542892, 0.012903732247650623, 0.010530934669077396, 0.015047149732708931, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07456009835004807, 0.09125705808401108, 0.20381297171115875, 0.09053967893123627, 0.6734579801559448, 0.8927901983261108, 0.9854956865310669, 0.19160649180412292, 0.848483681678772, 0.3795100748538971, 0.0351644828915596, 0.06069617718458176, 0.0190274715423584, 0.13319239020347595, 0.1618155688047409, 0.029784632846713066, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13663174211978912, 0.5250937938690186, 0.20416004955768585, 0.37758082151412964, 0.7281314134597778, 0.24714940786361694, 0.006291824858635664, 0.029336191713809967, 0.258807897567749, 0.17944614589214325, 0.2768983840942383, 0.49996671080589294, 0.6760725975036621, 0.0684136375784874, 0.9500845074653625, 0.04427658021450043, 0.027829600498080254, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05520259216427803, 0.4062710404396057, 0.11698392778635025, 0.09814880043268204, 0.8328142166137695, 0.46247926354408264, 0.07190129905939102, 0.3418641984462738, 0.14486591517925262, 0.025201991200447083, 0.042143724858760834, 0.4074908196926117, 0.1494714319705963, 0.17342594265937805, 0.908286988735199, 0.5950636863708496, 0.14296366274356842, 0.20851416885852814, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08497714251279831, 0.5087416172027588, 0.4508724510669708, 0.33144411444664, 0.600685715675354, 0.523800790309906, 0.4743403494358063, 0.10964386910200119, 0.6009643077850342, 0.29714730381965637, 0.1661888062953949, 0.10026849061250687, 0.19036318361759186, 0.07889659702777863, 0.29447081685066223, 0.5917950868606567, 0.05482999235391617, 0.0994495078921318, 0.08629819005727768, NaN, NaN, NaN, NaN, NaN, NaN], [0.04716389998793602, 0.6635201573371887, 0.5744545459747314, 0.33429521322250366, 0.755266010761261, 0.7800281643867493, 0.9541771411895752, 0.5776658058166504, 0.8714791536331177, 0.9158549308776855, 0.2818737030029297, 0.06938906759023666, 0.10379814356565475, 0.3064776659011841, 0.7474142909049988, 0.7715258002281189, 0.37782159447669983, 0.057383324950933456, 0.013433223590254784, 0.03400390222668648, NaN, NaN, NaN, NaN, NaN], [0.1486319750547409, 0.22267495095729828, 0.42902871966362, 0.07982667535543442, 0.5459871888160706, 0.9060689210891724, 0.8350642919540405, 0.10920917987823486, 0.4773065447807312, 0.7826967239379883, 0.5733710527420044, 0.26356616616249084, 0.040332335978746414, 0.031653065234422684, 0.8572309613227844, 0.5636150240898132, 0.07464684545993805, 0.03465104475617409, 0.03009859099984169, 0.008700854144990444, 0.005375253036618233, NaN, NaN, NaN, NaN], [0.25873932242393494, 0.5196211338043213, 0.3300914764404297, 0.5837901830673218, 0.4101006090641022, 0.7175306677818298, 0.6572118401527405, 0.6919461488723755, 0.6594171524047852, 0.7066829204559326, 0.46555259823799133, 0.3380126953125, 0.05317035689949989, 0.053740378469228745, 0.031323984265327454, 0.30507126450538635, 0.1422475129365921, 0.03319966048002243, 0.08714800328016281, 0.01252773217856884, 0.006611488293856382, 0.007115270011126995, NaN, NaN, NaN], [0.011579165235161781, 0.05381239950656891, 0.044945720583200455, 0.035533830523490906, 0.6624263525009155, 0.8997865319252014, 0.9679857492446899, 0.17051655054092407, 0.940772533416748, 0.6132625341415405, 0.01721411757171154, 0.04632151871919632, 0.010550450533628464, 0.08354383707046509, 0.12839946150779724, 0.02755529060959816, 0.44050073623657227, 0.04286862909793854, 0.01342833787202835, 0.003870438551530242, 0.026607532054185867, 0.02663758397102356, 0.005111980251967907, NaN, NaN], [0.13300661742687225, 0.5851269960403442, 0.20284885168075562, 0.5700805187225342, 0.7479174137115479, 0.39722636342048645, 0.004733124747872353, 0.0698152482509613, 0.6515945196151733, 0.5409151315689087, 0.25820717215538025, 0.4583084285259247, 0.6744768619537354, 0.3421478569507599, 0.9633424878120422, 0.1852269172668457, 0.04996338114142418, 0.5482219457626343, 0.296283096075058, 0.48366567492485046, 0.06441208720207214, 0.9149421453475952, 0.02780383825302124, 0.0073219588957726955, NaN], [0.14593175053596497, 0.2687321603298187, 0.04604685679078102, 0.30660173296928406, 0.3806478679180145, 0.38105660676956177, 0.15303322672843933, 0.014211257919669151, 0.05383581668138504, 0.20604565739631653, 0.2462100237607956, 0.5718756914138794, 0.5113963484764099, 0.21981710195541382, 0.4276719391345978, 0.5577609539031982, 0.4118191599845886, 0.31598320603370667, 0.5468451976776123, 0.4359907805919647, 0.2059280127286911, 0.3916337192058563, 0.2548142671585083, 0.2198532670736313, 0.026425611227750778]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1684475541114807, 0.01643766649067402, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20323613286018372, 0.02236698381602764, 0.0030780781526118517, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15523119270801544, 0.029148569330573082, 0.04869325831532478, 0.027081435546278954, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20906439423561096, 0.016835892572999, 0.005647255107760429, 0.004844226874411106, 0.00019458922906778753, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19736447930335999, 0.01826038584113121, 0.012854915112257004, 0.09684289991855621, 0.0006958578014746308, 4.3345058656996116e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16369424760341644, 0.023256592452526093, 0.01855486072599888, 0.06154748797416687, 0.06098903343081474, 0.10795246064662933, 0.023746412247419357, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19143380224704742, 0.11398851871490479, 0.03716170787811279, 0.07628969103097916, 0.38886839151382446, 0.24263328313827515, 0.13712459802627563, 0.02201412245631218, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2130274772644043, 0.007986752316355705, 0.02235114760696888, 0.0019427334191277623, 0.005593507084995508, 0.012699572369456291, 0.006745419930666685, 0.06126464158296585, 0.14077326655387878, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.22579564154148102, 0.013292824849486351, 0.10215212404727936, 0.005943832919001579, 0.013894540257751942, 0.01404587086290121, 0.02319374494254589, 0.10344905406236649, 0.1325504034757614, 0.008661924861371517, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1733061671257019, 0.07715445756912231, 0.2302267998456955, 0.05804288014769554, 0.07560069113969803, 0.23177897930145264, 0.2901765704154968, 0.042333029210567474, 0.08450006693601608, 0.04456959664821625, 0.015471314080059528, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.16428759694099426, 0.01361166127026081, 0.2167942076921463, 0.03707392141222954, 0.09917350113391876, 0.2872558534145355, 0.08793877810239792, 0.03127053380012512, 0.051127880811691284, 0.02603980340063572, 0.12251178920269012, 0.06466985493898392, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2214493751525879, 0.0034381633158773184, 0.025536755099892616, 0.005642351228743792, 0.0024517737329006195, 0.00733930105343461, 0.0003064426709897816, 0.024970028549432755, 0.0009503457695245743, 0.0013023557839915156, 0.012362079694867134, 0.002213133964687586, 0.0037243058905005455, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21803884208202362, 0.044672977179288864, 0.15033316612243652, 0.24480289220809937, 0.0010314357932657003, 0.006885815411806107, 0.017953861504793167, 0.09280995279550552, 0.09214792400598526, 0.01309943851083517, 0.026278402656316757, 0.029330603778362274, 0.10137840360403061, 0.0009828503243625164, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.28474918007850647, 0.005827821791172028, 0.0010850036051124334, 0.005180059466511011, 0.00018831032502930611, 0.002925402717664838, 0.0029562395066022873, 0.005281978752464056, 0.002952893264591694, 0.013548285700380802, 0.01663871854543686, 0.02234998345375061, 0.001472283387556672, 0.00024227210087701678, 9.911999950418249e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11472342163324356, 0.017006950452923775, 0.03429265320301056, 0.05351921543478966, 0.010289198718965054, 0.02545105293393135, 0.002036151010543108, 0.08590202778577805, 0.007977829314768314, 0.008050770498812199, 0.02079172432422638, 0.07815419882535934, 0.25072064995765686, 0.11726108938455582, 0.04080193489789963, 0.020839283242821693, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.25351014733314514, 0.018978603184223175, 0.013279697857797146, 0.14657457172870636, 0.0005683518829755485, 0.003044809214770794, 0.0003673452010843903, 0.0009085922501981258, 0.00026260188315063715, 6.703466351609677e-05, 0.00393629027530551, 0.0411190427839756, 0.014572926796972752, 0.0009043514728546143, 0.001453216653317213, 0.001335341832600534, 0.0036634530406445265, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2249869406223297, 0.0773954764008522, 0.10561174154281616, 0.3267342746257782, 0.011780736967921257, 0.03227663040161133, 0.09185110032558441, 0.03840579837560654, 0.01289159432053566, 0.002641883445903659, 0.03386297821998596, 0.16820214688777924, 0.06345225125551224, 0.027306171134114265, 0.007737002335488796, 0.018253128975629807, 0.0508209764957428, 0.015562118031084538, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17073971033096313, 0.01119090337306261, 0.07090220600366592, 0.026190776377916336, 0.04357914999127388, 0.10384812206029892, 0.05681576952338219, 0.008270802907645702, 0.011212479323148727, 0.016114890575408936, 0.1306251734495163, 0.04437248408794403, 0.022720789536833763, 0.0017881430685520172, 0.005742507986724377, 0.03271590173244476, 0.12170897424221039, 0.18442584574222565, 0.07238933444023132, NaN, NaN, NaN, NaN, NaN, NaN], [0.2460513859987259, 0.004599481821060181, 0.030415518209338188, 0.006707339081913233, 0.001940727117471397, 0.0018293699249625206, 0.0002438600640743971, 0.021702459082007408, 0.00019114103633910418, 0.0004616644873749465, 0.02795419655740261, 0.007376548834145069, 0.009364028461277485, 0.0008695388678461313, 0.027626920491456985, 0.002984545426443219, 0.0021758046932518482, 0.005276597570627928, 0.0015223525697365403, 0.0046029179356992245, NaN, NaN, NaN, NaN, NaN], [0.1682240217924118, 0.15532228350639343, 0.17499232292175293, 0.31528380513191223, 0.0016938054468482733, 0.0013859918108209968, 0.0071086762472987175, 0.08609996736049652, 0.02145048975944519, 0.00334079097956419, 0.08546027541160583, 0.16909679770469666, 0.5000762343406677, 0.012536582536995411, 0.0033327846322208643, 0.01681024581193924, 0.01291667390614748, 0.11205089092254639, 0.06917328387498856, 0.24062496423721313, 0.003104837378486991, NaN, NaN, NaN, NaN], [0.30163663625717163, 0.008585775271058083, 0.0018221536884084344, 0.004949942696839571, 0.0002661931503098458, 0.0017199779395014048, 0.00286088977009058, 0.004591777920722961, 0.0013412131229415536, 0.009152509272098541, 0.029603971168398857, 0.059182800352573395, 0.004352512303739786, 0.0009281163802370429, 0.00013420419418253005, 0.0015637356555089355, 0.004895435180515051, 0.0020298720337450504, 0.016267914324998856, 0.0014363413210958242, 0.00015049855574034154, 4.989441003999673e-05, NaN, NaN, NaN], [0.1420876681804657, 0.030559053644537926, 0.035777460783720016, 0.0549585185945034, 0.010907668620347977, 0.018195953220129013, 0.005288956221193075, 0.07946551591157913, 0.003352995030581951, 0.00945360492914915, 0.03057919070124626, 0.20277532935142517, 0.5438944697380066, 0.2487112432718277, 0.11027072370052338, 0.03672702983021736, 0.009589559398591518, 0.03681262582540512, 0.12653782963752747, 0.3100517988204956, 0.04488144814968109, 0.07299992442131042, 0.024292031303048134, NaN, NaN], [0.2571920156478882, 0.012253361754119396, 0.00982633139938116, 0.09085621684789658, 0.00026428516139276326, 0.001174133620224893, 0.00010905979434028268, 0.0006958161829970777, 9.435929678147659e-05, 1.889842314994894e-05, 0.0019355103140696883, 0.03233037516474724, 0.014144179411232471, 0.0034062752965837717, 0.0014896523207426071, 0.0032966958824545145, 0.0043079969473183155, 0.002425077836960554, 0.0237245112657547, 0.017915409058332443, 0.0004631538176909089, 0.0033925946336239576, 0.0019653798080980778, 0.0010656031081452966, NaN], [0.25252944231033325, 0.012149164453148842, 0.019892947748303413, 0.013666713610291481, 0.05940697342157364, 0.04882493242621422, 0.025430571287870407, 0.00045668394886888564, 0.0054928152821958065, 0.005623141769319773, 0.004253733437508345, 0.014798035845160484, 0.012909402139484882, 0.011927488259971142, 0.007018915377557278, 0.021986471489071846, 0.016502689570188522, 0.002887164242565632, 0.006932961288839579, 0.007926056161522865, 0.015145027078688145, 0.005945136770606041, 0.016453862190246582, 0.011257275938987732, 0.0009747393196448684]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14568212628364563, 0.073321633040905, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07740449905395508, 0.019538799300789833, 0.31676185131073, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11254165321588516, 0.04977253079414368, 0.12113941460847855, 0.18998825550079346, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09693466126918793, 0.12094055861234665, 0.48810020089149475, 0.07605772465467453, 0.10663138329982758, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.002718105213716626, 0.037000641226768494, 0.1506986916065216, 0.012303436174988747, 0.09212689101696014, 0.5217995047569275, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17887507379055023, 0.10589989274740219, 0.004075651057064533, 0.0014342612121254206, 0.00521382549777627, 0.031908128410577774, 0.003124895039945841, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.23519471287727356, 0.3653021454811096, 0.05512593686580658, 0.10675911605358124, 0.0014886436983942986, 0.001230676076374948, 0.003634560154750943, 0.00975269265472889, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19171930849552155, 0.3204987347126007, 0.0060858046635985374, 0.010409774258732796, 0.003722283523529768, 0.0010954621247947216, 0.0028676562942564487, 0.35306307673454285, 0.01622932404279709, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.25555557012557983, 0.13076956570148468, 0.003832729533314705, 0.0447237528860569, 0.014599477872252464, 0.0024878191761672497, 0.0016443775966763496, 0.20187559723854065, 0.0005508072790689766, 0.0029457835480570793, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13948844373226166, 0.2463626265525818, 0.09502393007278442, 0.197096586227417, 0.47678983211517334, 0.3142886161804199, 0.09103813022375107, 0.10499368607997894, 0.07698603719472885, 0.026083102449774742, 0.3110981583595276, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1511228382587433, 0.027682308107614517, 0.014322453178465366, 0.0030328254215419292, 0.04723867028951645, 0.30981165170669556, 0.025852922350168228, 0.018514074385166168, 0.01515920553356409, 0.009253463707864285, 0.10175863653421402, 0.16996310651302338, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1847103387117386, 0.05052594095468521, 0.005765186157077551, 0.018545929342508316, 0.00881477165967226, 0.0375242680311203, 0.027162199839949608, 0.09025334566831589, 0.0028228689916431904, 0.0033718899358063936, 0.1103500947356224, 0.0837099552154541, 0.0044236015528440475, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.27341794967651367, 0.03427007421851158, 0.008004172705113888, 0.009254892356693745, 0.005621441174298525, 0.00972525030374527, 0.005248658824712038, 0.02184745855629444, 0.0006181569187901914, 0.0005494534852914512, 0.06994801014661789, 0.02213645726442337, 0.004287416115403175, 0.0008399627404287457, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.008804291486740112, 0.07617928832769394, 0.47516930103302, 0.07513945549726486, 0.5241973400115967, 0.4384346902370453, 0.06213618069887161, 0.06345370411872864, 0.0682281106710434, 0.15877418220043182, 0.023486817255616188, 0.026526909321546555, 0.0028373831883072853, 0.001617963775061071, 0.37629759311676025, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.26533833146095276, 0.10994716733694077, 0.010266831144690514, 0.037150826305150986, 0.009969023987650871, 0.00030588259687647223, 8.988264016807079e-05, 0.07940464466810226, 0.00027601365582086146, 0.0013282618019729853, 0.009904097765684128, 0.03278518095612526, 0.0630892813205719, 0.10911130160093307, 0.016624033451080322, 0.011541539803147316, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2451263964176178, 0.014867580495774746, 0.0005470102187246084, 0.0054298522882163525, 0.0004450916312634945, 0.0006575370789505541, 3.8741818570997566e-05, 0.0010275153908878565, 0.0013172366889193654, 0.0019110681023448706, 0.13600468635559082, 0.29138538241386414, 0.011091821826994419, 0.0002334356977371499, 0.0002162840828532353, 0.0001727231137920171, 0.004782650154083967, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18341027200222015, 0.31211209297180176, 0.08544175326824188, 0.17215219140052795, 0.07786234468221664, 0.033002957701683044, 0.028957894071936607, 0.08467604964971542, 0.018818018957972527, 0.0016417433507740498, 0.15075404942035675, 0.1522863805294037, 0.03350237384438515, 0.006119633559137583, 0.022573737427592278, 0.03810621052980423, 0.13675758242607117, 0.1992093175649643, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1540856957435608, 0.05453011393547058, 0.023697303608059883, 0.003979950677603483, 0.014029269106686115, 0.1104540005326271, 0.019629694521427155, 0.011429534293711185, 0.010672842152416706, 0.00807265006005764, 0.1843080371618271, 0.19234825670719147, 0.0017768212128430605, 0.006891301833093166, 0.08265318721532822, 0.014878016896545887, 0.09550431370735168, 0.1691773235797882, 0.20674942433834076, NaN, NaN, NaN, NaN, NaN, NaN], [0.21139073371887207, 0.06409671157598495, 0.007977590896189213, 0.017582383006811142, 0.004139575641602278, 0.008497070521116257, 0.024324562400579453, 0.12332659959793091, 0.0006915424601174891, 0.0006991134723648429, 0.09821731597185135, 0.18821127712726593, 0.009975801222026348, 0.024784373119473457, 0.009686794131994247, 0.0016004297649487853, 0.006526788230985403, 0.04246864095330238, 0.05479469522833824, 0.004482009913772345, NaN, NaN, NaN, NaN, NaN], [0.33224669098854065, 0.07294216006994247, 0.01592269167304039, 0.006994656287133694, 0.003661615075543523, 0.0007586313877254725, 0.0006907262722961605, 0.022764746099710464, 0.000276167003903538, 9.849678463069722e-05, 0.08613532781600952, 0.07070992141962051, 0.03258151933550835, 0.002256957348436117, 0.00035050295991823077, 0.002809839555993676, 0.005992868449538946, 0.14088936150074005, 0.024111032485961914, 0.015468394383788109, 0.000736193498596549, NaN, NaN, NaN, NaN], [0.00368693470954895, 0.0603332445025444, 0.389295369386673, 0.03955860063433647, 0.26089394092559814, 0.125760018825531, 0.029167605563998222, 0.03710402920842171, 0.03377004712820053, 0.08135493099689484, 0.01946301944553852, 0.033920928835868835, 0.00409010099247098, 0.0020981510169804096, 0.4028157889842987, 0.01821253076195717, 0.03254074230790138, 0.005954912398010492, 0.016414301469922066, 0.0033934058155864477, 0.0012025205651298165, 0.37666910886764526, NaN, NaN, NaN], [0.30478137731552124, 0.23805196583271027, 0.009743728674948215, 0.02953244559466839, 0.005627358797937632, 0.00013927526015322655, 0.00016958850028458983, 0.09182754158973694, 0.00019882968626916409, 0.0018803260754793882, 0.01743759773671627, 0.09691343456506729, 0.09625609964132309, 0.0949849784374237, 0.057061683386564255, 0.028116967529058456, 0.00013736996334046125, 0.022905906662344933, 0.02515738271176815, 0.029101604595780373, 0.01233749371021986, 0.027021989226341248, 0.012159456498920918, NaN, NaN], [0.2508227825164795, 0.013127491809427738, 0.0004774215049110353, 0.005875048227608204, 0.00014762053615413606, 0.0003128673997707665, 1.7799626220948994e-05, 0.0017815351020544767, 0.0009225650574080646, 0.0009481729357503355, 0.09391504526138306, 0.24316561222076416, 0.008820290677249432, 0.0015348505694419146, 0.0002856143401004374, 0.00038499117363244295, 0.010248353704810143, 0.0923430323600769, 0.1539699137210846, 0.0089821582660079, 0.00013843990745954216, 0.0004539538058452308, 6.709429726470262e-05, 0.0014084051363170147, NaN], [0.06230561435222626, 0.051613274961709976, 0.02077883668243885, 0.04204944148659706, 0.07247611880302429, 0.11675790697336197, 0.004215644672513008, 0.00555834174156189, 0.008976897224783897, 0.017200933769345284, 0.007355507928878069, 0.06492317467927933, 0.04215962812304497, 0.02968345396220684, 0.23223130404949188, 0.03253115341067314, 0.08794146776199341, 0.025323374196887016, 0.08459514379501343, 0.05644838511943817, 0.04970480501651764, 0.3588789105415344, 0.028869707137346268, 0.11940079927444458, 0.27181047201156616]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04884753376245499, 0.31528204679489136, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [7.444373295584228e-06, 4.17321571148932e-05, 0.5221405029296875, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09023705869913101, 0.59262615442276, 0.038057319819927216, 0.1896824985742569, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0001943353418027982, 0.004992108792066574, 0.35714879631996155, 0.028785984963178635, 0.7041940689086914, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [1.0879062756430358e-05, 5.022298137191683e-05, 0.0836932584643364, 0.0041815838776528835, 0.7177854776382446, 0.4451410174369812, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.003986984025686979, 0.03902542591094971, 0.00027279910864308476, 0.00016326647892128676, 0.09999275952577591, 0.23601794242858887, 0.8888784646987915, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0004483810334932059, 0.01581367664039135, 0.00053547159768641, 0.005416989792138338, 0.0004931549192406237, 1.743426764733158e-06, 0.0002464183489792049, 0.38669928908348083, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0014915558276697993, 0.0036082565784454346, 0.0005674233543686569, 0.0010717788245528936, 0.04321836307644844, 0.5446166396141052, 0.38359156250953674, 0.006869717035442591, 0.0028910271357744932, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [8.035104838199914e-05, 0.005924052093178034, 0.005847892723977566, 0.020417997613549232, 0.11436353623867035, 0.6555760502815247, 0.4247216582298279, 0.04553407058119774, 0.00039129320066422224, 0.013846640475094318, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0012459981953725219, 0.12171746790409088, 0.022806251421570778, 0.021380947902798653, 0.018195364624261856, 0.08835338801145554, 0.20732422173023224, 0.30439698696136475, 0.09951408952474594, 0.2512991428375244, 0.4290468692779541, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.007976139895617962, 0.03435874730348587, 0.026849543675780296, 0.002102706115692854, 0.13315419852733612, 0.1177494078874588, 0.08904305100440979, 0.576798677444458, 0.140389084815979, 0.6266443729400635, 0.32779327034950256, 0.5110495090484619, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0015641784993931651, 0.09294694662094116, 0.006881145294755697, 0.0020365919917821884, 0.4301930069923401, 0.06383264064788818, 0.0045266724191606045, 0.17422647774219513, 0.00404678238555789, 0.006469257641583681, 0.052995309233665466, 0.1725381463766098, 0.668171763420105, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.004304439760744572, 0.05993141233921051, 0.054169829934835434, 0.025809768587350845, 0.7262899279594421, 0.2466905415058136, 0.15344326198101044, 0.33606013655662537, 0.02952432446181774, 0.07010773569345474, 0.008777104318141937, 0.03394261747598648, 0.032566726207733154, 0.6152393221855164, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [1.0540320545260329e-05, 0.0013190202880650759, 0.20101842284202576, 0.004686327185481787, 0.13271625339984894, 0.04526880756020546, 0.0007031870190985501, 0.0011485026916489005, 0.002882149303331971, 0.0005991549696773291, 0.0030197217129170895, 0.004800362046808004, 0.004403174854815006, 0.002436757553368807, 0.4002683460712433, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0003210107679478824, 0.5876501798629761, 0.16318874061107635, 0.7096263766288757, 0.11595475673675537, 0.007003267295658588, 0.001205803593620658, 0.1902448534965515, 0.011727835983037949, 0.44888344407081604, 0.8117052912712097, 0.45698752999305725, 0.023960944265127182, 0.010929742828011513, 0.005293603055179119, 0.00987145397812128, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.020372437313199043, 0.3410835862159729, 0.6929088234901428, 0.04383905977010727, 0.1458517462015152, 0.4223538339138031, 0.9439106583595276, 0.9473816156387329, 0.15120889246463776, 0.7730743288993835, 0.5082507133483887, 0.0460858978331089, 0.032336097210645676, 0.011211436241865158, 0.009573124349117279, 0.0003536108124535531, 0.06564418971538544, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.020423829555511475, 0.09150233864784241, 0.593336284160614, 0.050333935767412186, 0.04262891411781311, 0.44151586294174194, 0.7098277807235718, 0.36869171261787415, 0.7183430194854736, 0.3146522641181946, 0.5934929251670837, 0.08962199836969376, 0.01141325756907463, 0.0268073882907629, 0.008290876634418964, 0.022364463657140732, 0.0520397312939167, 0.3134966492652893, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.008604546077549458, 0.07562410086393356, 0.10463645309209824, 0.003217896446585655, 0.1296835094690323, 0.21162182092666626, 0.30799001455307007, 0.7962209582328796, 0.27782267332077026, 0.5974112749099731, 0.3643631041049957, 0.5975222587585449, 0.032379183918237686, 0.8344925045967102, 0.5903766751289368, 0.1521190106868744, 0.10492946952581406, 0.10503242909908295, 0.5022279620170593, NaN, NaN, NaN, NaN, NaN, NaN], [0.0010157334618270397, 0.08574047684669495, 0.010654903016984463, 0.003869200125336647, 0.15051355957984924, 0.02434478886425495, 0.005829520523548126, 0.10341739654541016, 0.0023463659454137087, 0.00469975033774972, 0.1621563881635666, 0.27765417098999023, 0.6246147155761719, 0.44377410411834717, 0.0757245346903801, 0.08620554953813553, 0.08146335929632187, 0.32109129428863525, 0.1958039551973343, 0.5327519178390503, NaN, NaN, NaN, NaN, NaN], [0.0009064326295629144, 0.04867112636566162, 0.09537991136312485, 0.12993541359901428, 0.38632717728614807, 0.056282784789800644, 0.13602504134178162, 0.18383464217185974, 0.024170320481061935, 0.09972675889730453, 0.022063996642827988, 0.042059145867824554, 0.01842264086008072, 0.8592916131019592, 0.1306053251028061, 0.06485681235790253, 0.048735883086919785, 0.037178389728069305, 0.017466288059949875, 0.006924192421138287, 0.8764364123344421, NaN, NaN, NaN, NaN], [1.2418378219081205e-06, 0.0003037750138901174, 0.10264009237289429, 0.0010840333998203278, 0.03004724159836769, 0.00720690144225955, 0.00017297905287705362, 0.00021026108879595995, 0.0005732537247240543, 0.00013229742762632668, 0.0014890850288793445, 0.0027206502854824066, 0.0022100789938122034, 0.0018764312844723463, 0.22427155077457428, 0.0012303950497880578, 0.0001426686649210751, 0.0015814924845471978, 0.00487141590565443, 0.0029599322006106377, 0.003610847517848015, 0.41901907324790955, NaN, NaN, NaN], [0.00015546051145065576, 0.5271192193031311, 0.2684091329574585, 0.7487277388572693, 0.0846778005361557, 0.003557654097676277, 0.0064069912768900394, 0.16770148277282715, 0.008421340025961399, 0.27412623167037964, 0.8534677624702454, 0.5243650078773499, 0.02665238454937935, 0.01776440255343914, 0.013793676160275936, 0.00868560466915369, 0.08064579218626022, 0.69512540102005, 0.49261555075645447, 0.010526523925364017, 0.0028473760467022657, 0.008281596936285496, 0.007198471110314131, NaN, NaN], [0.03285643830895424, 0.3327244818210602, 0.7442528605461121, 0.049526505172252655, 0.13722854852676392, 0.37294694781303406, 0.9746374487876892, 0.9050161242485046, 0.144730344414711, 0.44314900040626526, 0.6168692708015442, 0.18840178847312927, 0.12898683547973633, 0.1250022053718567, 0.01759251020848751, 0.0030696040485054255, 0.6704888939857483, 0.3205258250236511, 0.28675025701522827, 0.09770815074443817, 0.0085873082280159, 0.028106005862355232, 0.0015327840810641646, 0.12156207114458084, NaN], [0.027913866564631462, 0.6360336542129517, 0.8947576880455017, 0.5603421926498413, 0.3501611351966858, 0.3494046926498413, 0.7655782103538513, 0.9696423411369324, 0.8922762274742126, 0.42980051040649414, 0.4555767774581909, 0.17016178369522095, 0.1410100758075714, 0.652664303779602, 0.2781027853488922, 0.07839874923229218, 0.11400053650140762, 0.10023999214172363, 0.04957454651594162, 0.07193805277347565, 0.5185664892196655, 0.15356925129890442, 0.02747632935643196, 0.046240244060754776, 0.017650051042437553]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02477514185011387, 0.37543168663978577, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02274254709482193, 0.6458237767219543, 0.013541627675294876, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03146426007151604, 0.019330549985170364, 0.019686071202158928, 0.5363749265670776, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05261930450797081, 0.12757715582847595, 0.003555318573489785, 0.48483166098594666, 0.00033596818684600294, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09825422614812851, 0.08890903741121292, 0.0022953739389777184, 0.3788372278213501, 6.525879871333018e-05, 3.547202504705638e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1839720457792282, 0.005392392631620169, 0.0012601928319782019, 0.000860364583786577, 0.0008281354093924165, 0.0005760629428550601, 0.002849774667993188, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.005911883432418108, 0.0029267233330756426, 0.007144090253859758, 0.001919957809150219, 0.004637785721570253, 0.004848909098654985, 0.006189228966832161, 0.3764636814594269, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2256152480840683, 0.0020181250292807817, 0.0012439934071153402, 0.00031968209077604115, 0.0029859780333936214, 0.017534615471959114, 0.0004058087943121791, 0.00034323628642596304, 0.029154805466532707, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03960844501852989, 0.0036635666619986296, 0.00109457119833678, 0.0017422186210751534, 0.022469639778137207, 0.004235065542161465, 0.007348764222115278, 0.00280297570861876, 0.030011437833309174, 0.576508641242981, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0628783106803894, 0.014568633399903774, 0.003403500886633992, 0.005917230620980263, 0.009509358555078506, 0.0019911406561732292, 0.005211993586272001, 0.01603839360177517, 0.00502167409285903, 0.3301290273666382, 0.10268117487430573, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.178706556558609, 0.5124386548995972, 0.028256116434931755, 0.011254883371293545, 0.03223628178238869, 0.0004171380714979023, 0.004843876231461763, 0.09010603278875351, 0.0025540743954479694, 0.016201328486204147, 0.029397757723927498, 0.010837158188223839, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18362975120544434, 0.10373001545667648, 0.006869313772767782, 0.010921900160610676, 0.01820673979818821, 0.0017379705095663667, 0.002349345711991191, 0.03729201853275299, 5.792165029561147e-05, 0.0013579311780631542, 0.0025659396778792143, 0.008523254655301571, 0.1568114459514618, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.060853905975818634, 0.016029829159379005, 0.001439533894881606, 0.017260756343603134, 0.0007974627078510821, 0.0012342276750132442, 0.028226196765899658, 0.0047790613025426865, 0.0015612602001056075, 0.004867547657340765, 0.039023980498313904, 0.05208572745323181, 0.33480554819107056, 0.17332881689071655, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.043774526566267014, 0.2669547498226166, 0.035314492881298065, 0.1941595822572708, 0.006638282909989357, 0.005091785918921232, 0.2628510892391205, 0.2860943675041199, 0.06445851922035217, 0.34950578212738037, 0.6430334448814392, 0.5673049688339233, 0.6101463437080383, 0.29372307658195496, 0.0028161092195659876, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.018545497208833694, 0.059764593839645386, 0.0026272537652403116, 0.020267995074391365, 0.009687644429504871, 0.00033462722785770893, 0.0024671528954058886, 0.054633729159832, 5.4464391723740846e-05, 0.00043273900519125164, 0.0019224031129851937, 0.21117039024829865, 0.3183750510215759, 0.03866858780384064, 0.011778384447097778, 0.1297062188386917, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0004199208051431924, 4.603992783813737e-05, 8.09443406524224e-07, 2.029701317951549e-05, 3.386533080629306e-06, 2.203315261795069e-06, 4.220597020321293e-06, 8.901660294213798e-06, 0.00016298270202241838, 0.000983458710834384, 0.0005640776362270117, 0.0008154786773957312, 0.001651398022659123, 2.400618996034609e-06, 3.3168395020766184e-05, 6.549440058734035e-06, 0.8699775338172913, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06053417548537254, 0.012584012933075428, 0.0010002547642216086, 0.0027718576602637768, 0.006610550452023745, 0.0029896856285631657, 0.008355176076292992, 0.048459943383932114, 0.002307809190824628, 0.65205979347229, 0.1651758849620819, 0.011300449259579182, 0.029586348682641983, 0.014456091448664665, 0.0007872084970586002, 0.0008902085828594863, 0.029332326725125313, 0.16636918485164642, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19553376734256744, 0.2426333725452423, 0.004519153386354446, 0.00883245188742876, 0.006844275165349245, 0.00014635240950156003, 0.00260242260992527, 0.03859727829694748, 0.0011520206462591887, 0.014703472144901752, 0.016579829156398773, 0.003783928230404854, 0.01771795004606247, 0.0035672299563884735, 0.000677697011269629, 0.002100451150909066, 0.023971345275640488, 0.03231354430317879, 0.011524699628353119, NaN, NaN, NaN, NaN, NaN, NaN], [0.17035169899463654, 0.07290639728307724, 0.0013864204520359635, 0.008776376023888588, 0.010795027948915958, 0.0008890280150808394, 0.00375909055583179, 0.03264426812529564, 2.1074760297778994e-05, 0.0009656226029619575, 0.004805654752999544, 0.015095297247171402, 0.19429266452789307, 0.060086220502853394, 0.013300183229148388, 0.019145654514431953, 0.08634541183710098, 0.018065713346004486, 0.012390428222715855, 0.3474832773208618, NaN, NaN, NaN, NaN, NaN], [0.002681915881112218, 0.0020622191950678825, 1.740588413667865e-05, 0.001647116499952972, 2.462047996232286e-05, 1.4256034774007276e-05, 0.0023770714178681374, 0.0007797144935466349, 6.146806117612869e-05, 0.00019536878971848637, 0.023629816249012947, 0.022664623335003853, 0.058040015399456024, 0.02328144572675228, 0.00014305225340649486, 0.1791975051164627, 0.7950490117073059, 0.40287262201309204, 0.05916967615485191, 0.11726692318916321, 0.045271970331668854, NaN, NaN, NaN, NaN], [0.017539121210575104, 0.07800457626581192, 0.013338283635675907, 0.07843150943517685, 0.003389358287677169, 0.0011982140131294727, 0.07936429977416992, 0.08406823873519897, 0.016710255295038223, 0.13201765716075897, 0.339507520198822, 0.3268124461174011, 0.4709261357784271, 0.24707961082458496, 0.0009133804705925286, 0.27326905727386475, 0.539431095123291, 0.8842423558235168, 0.5773340463638306, 0.643308699131012, 0.15606866776943207, 0.0011033734772354364, NaN, NaN, NaN], [0.0009739195229485631, 0.0011780881322920322, 3.265493069193326e-05, 0.0005334040033631027, 0.0007281061843968928, 3.2774634746601805e-05, 0.0004276044783182442, 0.00342408730648458, 2.9227990125946235e-06, 5.522280844161287e-05, 0.00012372780474834144, 0.011400841176509857, 0.008755120448768139, 0.0017365129897370934, 0.0007705622701905668, 0.0024924452882260084, 0.4634210169315338, 0.010356471873819828, 0.06587640196084976, 0.03498200699687004, 0.005118835251778364, 0.0019369632937014103, 0.023791478946805, NaN, NaN], [0.00023119446996133775, 9.065014637599234e-06, 3.0932378081161005e-07, 7.128239758458221e-06, 2.417179757685517e-06, 1.9917408735636855e-06, 1.0686825362427044e-06, 3.5747166293731425e-06, 3.038432441826444e-05, 0.00024045849568210542, 0.00012102597975172102, 0.0003720777458511293, 0.0005474414792843163, 4.2138731259910855e-06, 8.004362825886346e-06, 4.010584234492853e-06, 0.22906039655208588, 0.00024706448311917484, 0.003541025100275874, 0.0035716970451176167, 1.1338630656609894e-06, 4.888530747848563e-05, 2.00755093828775e-05, 0.8455927968025208, NaN], [0.023575956001877785, 0.001566409133374691, 0.0004935376346111298, 0.015205318108201027, 0.0005761805805377662, 0.00026375881861895323, 0.0017682479228824377, 0.00015503005124628544, 0.011253873817622662, 0.321735680103302, 0.05970581993460655, 0.008942467160522938, 0.051820773631334305, 0.009087985381484032, 0.002068085130304098, 0.00584985688328743, 0.01019755844026804, 0.16441591084003448, 0.021173937246203423, 0.09159599989652634, 0.004452125634998083, 0.0037374526727944613, 0.01578103005886078, 0.01742226630449295, 0.3373567461967468]]], [[[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1729947179555893, 0.014742943458259106, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11518532782793045, 0.28854820132255554, 0.0005498379468917847, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12768876552581787, 0.007979520596563816, 0.05741023272275925, 0.14377589523792267, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.25598737597465515, 0.03471918776631355, 0.08263758569955826, 0.03616967797279358, 0.0012629067059606314, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.29742351174354553, 0.10481993854045868, 0.07552393525838852, 0.008401650935411453, 0.3407011330127716, 0.028353586792945862, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17861823737621307, 0.07256677001714706, 0.1795390099287033, 0.04586997628211975, 0.27750420570373535, 0.0032322825863957405, 0.09472999721765518, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1281835287809372, 0.008169662207365036, 0.10209551453590393, 0.22781534492969513, 0.13339588046073914, 0.022249281406402588, 0.2580547630786896, 0.0071509419940412045, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19490991532802582, 0.0105251120403409, 0.07082764059305191, 0.07746586948633194, 0.10047772526741028, 0.007984980009496212, 0.045915842056274414, 0.030714787542819977, 0.09154831618070602, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2116595059633255, 0.006228659767657518, 0.09237925708293915, 0.33000993728637695, 0.06037600710988045, 0.06468494236469269, 0.028822004795074463, 0.015993207693099976, 0.023504862561821938, 0.014777855016291142, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11546289920806885, 0.0627092570066452, 0.1015198826789856, 0.17440570890903473, 0.11644574254751205, 0.15138378739356995, 0.17151175439357758, 0.07174428552389145, 0.1994275599718094, 0.20994937419891357, 0.08254047483205795, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13584046065807343, 0.09117304533720016, 0.15590398013591766, 0.10968183726072311, 0.5585501790046692, 0.07535546272993088, 0.2762793302536011, 0.32588398456573486, 0.3246583938598633, 0.41251155734062195, 0.043567951768636703, 0.0185235645622015, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1674133688211441, 0.12648360431194305, 0.27492284774780273, 0.24355122447013855, 0.8769406676292419, 0.6096609234809875, 0.4704851806163788, 0.055198147892951965, 0.6140321493148804, 0.2705269455909729, 0.07450747489929199, 0.04471021145582199, 0.05369797348976135, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.035074394196271896, 0.012203776277601719, 0.2713678479194641, 0.27628132700920105, 0.5399907231330872, 0.3242804706096649, 0.5765586495399475, 0.02925838902592659, 0.3159044086933136, 0.11935708671808243, 0.16010764241218567, 0.31936678290367126, 0.22831447422504425, 0.09149928390979767, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1354324370622635, 0.08839684724807739, 0.010535157285630703, 0.3809414505958557, 0.006101538427174091, 0.04204240441322327, 0.6714356541633606, 0.02054513990879059, 0.44751474261283875, 0.5217893123626709, 0.16833685338497162, 0.4138224124908447, 0.5945862531661987, 0.14406909048557281, 0.000551112403627485, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.26645413041114807, 0.038747917860746384, 0.15441381931304932, 0.6166976094245911, 0.04416924715042114, 0.07849516719579697, 0.41569313406944275, 0.018940549343824387, 0.18770581483840942, 0.11268321424722672, 0.0962471142411232, 0.028718965128064156, 0.019747000187635422, 0.011864973232150078, 0.07090434432029724, 0.02976600080728531, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.26584282517433167, 0.03641113266348839, 0.24681606888771057, 0.03326011076569557, 0.5612249970436096, 0.11044078320264816, 0.038705065846443176, 0.07638699561357498, 0.20042885839939117, 0.41367095708847046, 0.16446417570114136, 0.05500950291752815, 0.0458536334335804, 0.038293108344078064, 0.05886702984571457, 0.005421455018222332, 0.03447017818689346, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.052208781242370605, 0.10399425774812698, 0.2661847770214081, 0.06582632660865784, 0.5218088626861572, 0.41107869148254395, 0.18652401864528656, 0.10915308445692062, 0.2499890774488449, 0.21385571360588074, 0.11996328830718994, 0.2169666439294815, 0.17541900277137756, 0.34852319955825806, 0.29904353618621826, 0.3583068549633026, 0.0660485103726387, 0.0772518739104271, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1452419012784958, 0.08285138756036758, 0.20162978768348694, 0.10332676023244858, 0.7324197292327881, 0.1815183311700821, 0.27558720111846924, 0.41944485902786255, 0.4614993929862976, 0.7035390734672546, 0.14779764413833618, 0.07484183460474014, 0.09274464100599289, 0.1956741362810135, 0.4027537703514099, 0.17018413543701172, 0.15845544636249542, 0.03217604011297226, 0.027846908196806908, NaN, NaN, NaN, NaN, NaN, NaN], [0.06803880631923676, 0.0777740478515625, 0.3149954080581665, 0.17862020432949066, 0.9274848103523254, 0.6797788739204407, 0.28538215160369873, 0.04841757193207741, 0.524702250957489, 0.33268001675605774, 0.06556227803230286, 0.08207366615533829, 0.08443650603294373, 0.19301387667655945, 0.68314129114151, 0.7843886613845825, 0.24039600789546967, 0.0983721911907196, 0.035574402660131454, 0.04086223617196083, NaN, NaN, NaN, NaN, NaN], [0.004222579766064882, 0.012189013883471489, 0.38177239894866943, 0.23501808941364288, 0.3822557032108307, 0.273560494184494, 0.28252631425857544, 0.039307549595832825, 0.41269388794898987, 0.3037600517272949, 0.1617780327796936, 0.33094146847724915, 0.37525615096092224, 0.1388353556394577, 0.8142803907394409, 0.5916069149971008, 0.18943282961845398, 0.08566068857908249, 0.11778654158115387, 0.1818830519914627, 0.04465563967823982, NaN, NaN, NaN, NaN], [0.0780838280916214, 0.07355974614620209, 0.01093215774744749, 0.22770193219184875, 0.008550305850803852, 0.06503485888242722, 0.5060688257217407, 0.02145100012421608, 0.43843212723731995, 0.6872871518135071, 0.1969044953584671, 0.45010682940483093, 0.7415768504142761, 0.3103433847427368, 0.001054091495461762, 0.20113487541675568, 0.21400661766529083, 0.41673052310943604, 0.3260871469974518, 0.620118260383606, 0.12724098563194275, 0.0004952864837832749, NaN, NaN, NaN], [0.3314567506313324, 0.06341477483510971, 0.5618032217025757, 0.642646074295044, 0.27415919303894043, 0.23788774013519287, 0.38833677768707275, 0.08984735608100891, 0.42147237062454224, 0.6564009785652161, 0.2928015887737274, 0.1047874391078949, 0.1023104265332222, 0.06365151703357697, 0.39097070693969727, 0.14560170471668243, 0.23420175909996033, 0.08592629432678223, 0.02493405155837536, 0.011453422717750072, 0.006046658381819725, 0.1451905518770218, 0.005812718998640776, NaN, NaN], [0.21756824851036072, 0.03937938064336777, 0.3266570568084717, 0.05877631530165672, 0.5281912088394165, 0.11102446913719177, 0.03890432044863701, 0.10487684607505798, 0.2815292179584503, 0.4750865697860718, 0.3058159351348877, 0.11602579057216644, 0.12021853774785995, 0.06692790240049362, 0.1190272718667984, 0.019106050953269005, 0.21307361125946045, 0.15337608754634857, 0.06824280321598053, 0.040861621499061584, 0.032932352274656296, 0.052440475672483444, 0.005818615201860666, 0.0524408333003521, NaN], [0.21100056171417236, 0.13406150043010712, 0.10563220083713531, 0.15389345586299896, 0.10192565619945526, 0.07836726307868958, 0.22881029546260834, 0.05055452138185501, 0.24765580892562866, 0.48160815238952637, 0.2201593518257141, 0.1761431246995926, 0.21236160397529602, 0.20979638397693634, 0.10962515324354172, 0.09009265154600143, 0.0623038187623024, 0.17415094375610352, 0.13285446166992188, 0.11576873064041138, 0.10801524668931961, 0.0743527039885521, 0.03413216769695282, 0.027520645409822464, 0.06626196205615997]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0702696219086647, 0.2507307231426239, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.028418319299817085, 0.003963488154113293, 0.4144974946975708, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13786309957504272, 0.03506092354655266, 0.02415982447564602, 0.10726116597652435, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.011229841969907284, 0.008138949982821941, 0.04613415151834488, 0.2518063187599182, 0.013397655449807644, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0016812672838568687, 0.012760624289512634, 0.002261990448459983, 0.2769384980201721, 0.03090759925544262, 0.0014064738061279058, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11822758615016937, 0.07095540314912796, 0.030966516584157944, 0.03516996279358864, 0.2070395052433014, 0.02684318646788597, 0.2317354679107666, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.23311708867549896, 0.026411496102809906, 0.011159970425069332, 0.03808103874325752, 0.017219573259353638, 0.006694006733596325, 0.001702688867226243, 0.009211051277816296, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1427604705095291, 0.06787170469760895, 0.04101337492465973, 0.04024908319115639, 0.2669386863708496, 0.04579312726855278, 0.07587221264839172, 0.10059545934200287, 0.18715938925743103, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.059837497770786285, 0.10673120617866516, 0.06554628908634186, 0.047321293503046036, 0.26084935665130615, 0.05379262939095497, 0.09055614471435547, 0.09319713711738586, 0.334230899810791, 0.23545128107070923, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06699422001838684, 0.48348554968833923, 0.10470042377710342, 0.2643885016441345, 0.49639153480529785, 0.11732041090726852, 0.061902400106191635, 0.1530170738697052, 0.11711295694112778, 0.23237623274326324, 0.09402092546224594, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.050390250980854034, 0.2627623975276947, 0.057036180049180984, 0.10587681084871292, 0.22481703758239746, 0.07078704982995987, 0.028480585664510727, 0.47086307406425476, 0.03990349546074867, 0.16108965873718262, 0.02393723465502262, 0.06960758566856384, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.29633763432502747, 0.1570599228143692, 0.07358378916978836, 0.08321648091077805, 0.01657349243760109, 0.02100137248635292, 0.019902318716049194, 0.5162196755409241, 0.03987365961074829, 0.018146652728319168, 0.026169516146183014, 0.00614600395783782, 0.07103840261697769, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1833065152168274, 0.0826280415058136, 0.06509751826524734, 0.017351830378174782, 0.08598462492227554, 0.028223805129528046, 0.03195580840110779, 0.045467328280210495, 0.041934747248888016, 0.016390223056077957, 0.05298775061964989, 0.05077003315091133, 0.2718433141708374, 0.04039132222533226, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09722712635993958, 0.09857381135225296, 0.2290657013654709, 0.162257120013237, 0.3208743929862976, 0.7083525657653809, 0.08285251259803772, 0.05820265784859657, 0.14296579360961914, 0.06442547589540482, 0.3963678479194641, 0.1963234394788742, 0.13509824872016907, 0.0551372766494751, 0.1773844212293625, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1786596029996872, 0.03035295568406582, 0.011360704898834229, 0.0041356864385306835, 0.02253635786473751, 0.032254207879304886, 0.05765725299715996, 0.06512543559074402, 0.26075252890586853, 0.14487245678901672, 0.06064848601818085, 0.02561355009675026, 0.06785233318805695, 0.08367668837308884, 0.11658230423927307, 0.21664968132972717, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02336198277771473, 0.027563903480768204, 0.02503703534603119, 0.002219978952780366, 0.024155667051672935, 0.005802824627608061, 0.011775066144764423, 0.03527237847447395, 0.0438326895236969, 0.16127318143844604, 0.07829897105693817, 0.04636809974908829, 0.16168944537639618, 0.17395752668380737, 0.5116502642631531, 0.11367138475179672, 0.24585914611816406, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14312313497066498, 0.6151867508888245, 0.2511911392211914, 0.34089455008506775, 0.21357816457748413, 0.06974375993013382, 0.04017443582415581, 0.4436698257923126, 0.0627409890294075, 0.029346130788326263, 0.06214871257543564, 0.07426106929779053, 0.37162381410598755, 0.1908751130104065, 0.2730017304420471, 0.09601876139640808, 0.07787502557039261, 0.1985486000776291, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05929486081004143, 0.1356429159641266, 0.08288607001304626, 0.1716676652431488, 0.17707081139087677, 0.11502664536237717, 0.023076828569173813, 0.41179341077804565, 0.03153251111507416, 0.08080360293388367, 0.03793509677052498, 0.0956316813826561, 0.40457794070243835, 0.3355584144592285, 0.2116643786430359, 0.2117510586977005, 0.0911363810300827, 0.13469243049621582, 0.08244834095239639, NaN, NaN, NaN, NaN, NaN, NaN], [0.34530380368232727, 0.14280815422534943, 0.08469259738922119, 0.20386184751987457, 0.018106382340192795, 0.025206930935382843, 0.03376462310552597, 0.665645956993103, 0.06945709139108658, 0.030968131497502327, 0.031062953174114227, 0.015101979486644268, 0.10170532017946243, 0.03453005850315094, 0.05652596056461334, 0.028510402888059616, 0.036133769899606705, 0.04489430412650108, 0.010548176243901253, 0.07425779104232788, NaN, NaN, NaN, NaN, NaN], [0.21361097693443298, 0.09641434252262115, 0.0472431480884552, 0.030436551198363304, 0.12823571264743805, 0.024378983303904533, 0.03781319037079811, 0.04478050768375397, 0.04302188381552696, 0.031242409721016884, 0.06916327774524689, 0.08240062743425369, 0.2609483301639557, 0.04106062278151512, 0.01303931511938572, 0.014160559512674809, 0.011109860613942146, 0.034855347126722336, 0.10407929867506027, 0.21024775505065918, 0.08525354415178299, NaN, NaN, NaN, NaN], [0.056013792753219604, 0.04104574769735336, 0.13420559465885162, 0.14404895901679993, 0.30753612518310547, 0.5552563667297363, 0.06356479972600937, 0.02527950517833233, 0.09324341267347336, 0.03306487947702408, 0.2522013187408447, 0.14255186915397644, 0.09901494532823563, 0.06439376622438431, 0.10042564570903778, 0.43083739280700684, 0.20968028903007507, 0.35324180126190186, 0.2700602114200592, 0.23262809216976166, 0.11776822060346603, 0.14138048887252808, NaN, NaN, NaN], [0.1699744164943695, 0.02438814751803875, 0.00377153092995286, 0.0020952692721039057, 0.017941365018486977, 0.009907160885632038, 0.04197421669960022, 0.08005423098802567, 0.16825814545154572, 0.08759146183729172, 0.037892259657382965, 0.02378804422914982, 0.12696562707424164, 0.21072204411029816, 0.039158232510089874, 0.12900760769844055, 0.018357207998633385, 0.09957201033830643, 0.024237502366304398, 0.12091250717639923, 0.2524404227733612, 0.044468626379966736, 0.19958341121673584, NaN, NaN], [0.016944430768489838, 0.011726072989404202, 0.017351148650050163, 0.0028529188130050898, 0.013441222719848156, 0.005811003036797047, 0.010734970681369305, 0.020825698971748352, 0.04144507274031639, 0.0777476355433464, 0.07330787181854248, 0.0589311420917511, 0.1305314600467682, 0.09686601907014847, 0.49986732006073, 0.09861493855714798, 0.24486178159713745, 0.2709232568740845, 0.08328418433666229, 0.1665872186422348, 0.2741791903972626, 0.5570544600486755, 0.09308093041181564, 0.18428745865821838, NaN], [0.043635401874780655, 0.027883753180503845, 0.11735352873802185, 0.09225393831729889, 0.11462916433811188, 0.1478782296180725, 0.04645288363099098, 0.049018505960702896, 0.08540874719619751, 0.16189652681350708, 0.081883005797863, 0.13365384936332703, 0.17616337537765503, 0.16547891497612, 0.3400772511959076, 0.14388780295848846, 0.2768324613571167, 0.1609276533126831, 0.18515954911708832, 0.2950800061225891, 0.32982173562049866, 0.4366631507873535, 0.3681013882160187, 0.34051525592803955, 0.05319627374410629]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1755252629518509, 0.00892956368625164, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18403629958629608, 0.12486936897039413, 0.01289399154484272, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07995349168777466, 0.1140136644244194, 0.16089488565921783, 0.271826833486557, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19368642568588257, 0.20833823084831238, 0.38513559103012085, 0.0724099725484848, 0.026710418984293938, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2920932173728943, 0.20408804714679718, 0.47836723923683167, 0.009784400463104248, 0.41401228308677673, 0.0022880665492266417, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2459677904844284, 0.013399376533925533, 0.165635347366333, 0.0016970435390248895, 0.00861914549022913, 0.0019094902090728283, 0.006659353617578745, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1659669429063797, 0.3024148941040039, 0.4638516902923584, 0.19814886152744293, 0.06386706978082657, 0.37022748589515686, 0.096834197640419, 0.004976118449121714, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.23605915904045105, 0.015010624192655087, 0.29689958691596985, 0.002272083656862378, 0.02557971514761448, 0.04829570651054382, 0.03933914750814438, 0.012097989208996296, 0.005491157062351704, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2229652851819992, 0.011020033620297909, 0.07613904774188995, 0.00492003234103322, 0.11613531410694122, 0.12462546676397324, 0.03799906745553017, 0.029671484604477882, 0.022334527224302292, 0.003809461137279868, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.30055463314056396, 0.03860635682940483, 0.08235271275043488, 0.12519411742687225, 0.07496307790279388, 0.24307869374752045, 0.02970520593225956, 0.043270040303468704, 0.01804984174668789, 0.008444367907941341, 0.04573319852352142, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.361846923828125, 0.0072926427237689495, 0.07028269022703171, 0.038334887474775314, 0.02117738127708435, 0.035939738154411316, 0.03011121228337288, 0.01985063962638378, 0.03699057549238205, 0.0448327511548996, 0.07655268162488937, 0.03217002749443054, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18510019779205322, 0.0857149139046669, 0.2959531545639038, 0.10870446264743805, 0.034602705389261246, 0.04019882157444954, 0.02403290942311287, 0.05409723520278931, 0.04566982761025429, 0.19149497151374817, 0.23549742996692657, 0.074503093957901, 0.01255789864808321, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03710656613111496, 0.054964251816272736, 0.037898506969213486, 0.3724515438079834, 0.058691613376140594, 0.03363177552819252, 0.06933214515447617, 0.05247700959444046, 0.15643684566020966, 0.589249849319458, 0.349843829870224, 0.29659491777420044, 0.2287619560956955, 0.05358140170574188, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2688547670841217, 0.1434442549943924, 0.18350595235824585, 0.07485228031873703, 0.0647219642996788, 0.04773847386240959, 0.14254990220069885, 0.03905782103538513, 0.2126167118549347, 0.24802155792713165, 0.30339401960372925, 0.17472584545612335, 0.03891041502356529, 0.02338952198624611, 0.026767900213599205, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1340402513742447, 0.12347351759672165, 0.42842522263526917, 0.0631304681301117, 0.06392616778612137, 0.1770109236240387, 0.11116458475589752, 0.04706185683608055, 0.09571156650781631, 0.3872493505477905, 0.5415271520614624, 0.14801958203315735, 0.013348261825740337, 0.016769861802458763, 0.019784821197390556, 0.012107723392546177, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3128407299518585, 0.02314484678208828, 0.20690661668777466, 0.0038596922531723976, 0.10119188576936722, 0.375572144985199, 0.077932208776474, 0.16011959314346313, 0.07805528491735458, 0.020400837063789368, 0.2237216979265213, 0.1006372720003128, 0.022764090448617935, 0.005061473231762648, 0.0205483790487051, 0.0018506759079173207, 0.001139476546086371, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5802629590034485, 0.17577120661735535, 0.22907592356204987, 0.3224048614501953, 0.21584153175354004, 0.3719359040260315, 0.08852899819612503, 0.18978306651115417, 0.06894023716449738, 0.008546161465346813, 0.34136468172073364, 0.44251179695129395, 0.07915834337472916, 0.27557075023651123, 0.0915302038192749, 0.0036887326277792454, 0.0038842300418764353, 0.015524323098361492, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.5194967985153198, 0.010316978208720684, 0.10247951745986938, 0.03023943491280079, 0.02351299114525318, 0.05376119539141655, 0.03751303628087044, 0.02858700230717659, 0.03933052346110344, 0.026450933888554573, 0.16396890580654144, 0.08825679868459702, 0.01957540772855282, 0.02957809716463089, 0.0652899444103241, 0.003373907646164298, 0.007670924998819828, 0.004321575630456209, 0.024295708164572716, NaN, NaN, NaN, NaN, NaN, NaN], [0.2508450150489807, 0.1962328553199768, 0.3596697747707367, 0.1504865288734436, 0.029224414378404617, 0.0663013905286789, 0.043777331709861755, 0.06269483268260956, 0.06556038558483124, 0.2250475436449051, 0.35171735286712646, 0.22191122174263, 0.018188640475273132, 0.026326660066843033, 0.017122289165854454, 0.0037187051493674517, 0.024730468168854713, 0.035062648355960846, 0.09351257234811783, 0.011442800983786583, NaN, NaN, NaN, NaN, NaN], [0.007168593350797892, 0.033368390053510666, 0.00873665139079094, 0.16062632203102112, 0.028196215629577637, 0.02527499757707119, 0.06866460293531418, 0.0198657363653183, 0.1544157713651657, 0.2752910256385803, 0.14698350429534912, 0.1242247000336647, 0.13061578571796417, 0.010920656844973564, 0.0055906628258526325, 0.006986986380070448, 0.030699225142598152, 0.36674854159355164, 0.2189747393131256, 0.2510429620742798, 0.04264682158827782, NaN, NaN, NaN, NaN], [0.317547470331192, 0.16016888618469238, 0.1976199448108673, 0.10644932836294174, 0.09830258786678314, 0.07801979035139084, 0.301817923784256, 0.05034731701016426, 0.32512444257736206, 0.2241876721382141, 0.4657731354236603, 0.2891538441181183, 0.08093820512294769, 0.06031876429915428, 0.06730521470308304, 0.14267991483211517, 0.289673775434494, 0.1076083853840828, 0.2949788272380829, 0.0365237332880497, 0.015645001083612442, 0.03993191570043564, NaN, NaN, NaN], [0.17233391106128693, 0.22507980465888977, 0.300968736410141, 0.03457535058259964, 0.06539295613765717, 0.2556630074977875, 0.12555503845214844, 0.08745130896568298, 0.10011813044548035, 0.13041436672210693, 0.501103937625885, 0.14929187297821045, 0.03132137656211853, 0.02265048772096634, 0.03383776918053627, 0.006481703836470842, 0.011523596942424774, 0.35894638299942017, 0.1662973165512085, 0.034177642315626144, 0.02702290564775467, 0.036704160273075104, 0.014952532015740871, NaN, NaN], [0.4115316569805145, 0.042032964527606964, 0.21366682648658752, 0.010602481663227081, 0.11737099289894104, 0.5779745578765869, 0.13523340225219727, 0.2636784315109253, 0.170937180519104, 0.020469455048441887, 0.3112620711326599, 0.17165400087833405, 0.044973500072956085, 0.006653682328760624, 0.053596071898937225, 0.008654352277517319, 0.002382548525929451, 0.02675137296319008, 0.09427332878112793, 0.01890433207154274, 0.002222384326159954, 0.018390605226159096, 0.0013299400452524424, 0.0009657714981585741, NaN], [0.38502925634384155, 0.1563987135887146, 0.13578397035598755, 0.1404726654291153, 0.14828255772590637, 0.28480827808380127, 0.15350891649723053, 0.09994281083345413, 0.06321649998426437, 0.030282480642199516, 0.13266463577747345, 0.1722954362630844, 0.07113035768270493, 0.024887708947062492, 0.016665330156683922, 0.03949398547410965, 0.020136239007115364, 0.01368448045104742, 0.09379612654447556, 0.030771953985095024, 0.011002926155924797, 0.007083212956786156, 0.009242233820259571, 0.007993990555405617, 0.018528543412685394]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17860974371433258, 0.0018437139224261045, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20284786820411682, 0.0034877806901931763, 0.08334594964981079, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1494244486093521, 0.3379342555999756, 0.0649241954088211, 0.006597604602575302, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2969810962677002, 0.005403619725257158, 0.054099179804325104, 0.0006044544279575348, 0.009600944817066193, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.32280662655830383, 0.01735025830566883, 0.15535852313041687, 0.00028658873634412885, 0.016427762806415558, 0.001579301548190415, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.016787199303507805, 0.10643576830625534, 0.24800433218479156, 0.4802894592285156, 0.03762362524867058, 0.06816797703504562, 0.10676699876785278, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.22070105373859406, 0.03063296526670456, 0.12860903143882751, 0.04803713783621788, 0.06528759002685547, 0.3172104060649872, 0.012414618395268917, 0.008628717623651028, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0170818492770195, 0.2921580374240875, 0.24774892628192902, 0.2979756295681, 0.16657015681266785, 0.03825104981660843, 0.39123743772506714, 0.0541624091565609, 0.01715947687625885, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06952934712171555, 0.09443160146474838, 0.3155873417854309, 0.2511345446109772, 0.20146684348583221, 0.17959536612033844, 0.500001072883606, 0.3407229483127594, 0.15127938985824585, 0.026401039212942123, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12491581588983536, 0.08139167726039886, 0.045777399092912674, 0.07585746794939041, 0.05243801325559616, 0.09790124744176865, 0.17415514588356018, 0.44996151328086853, 0.13761505484580994, 0.06580806523561478, 0.1016187071800232, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03772348165512085, 0.0006561332265846431, 0.04040418565273285, 0.23337695002555847, 0.0037602160591632128, 0.1251135915517807, 0.07994246482849121, 0.0032252452801913023, 0.044697076082229614, 0.05314825102686882, 0.16676445305347443, 0.42838534712791443, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.008380687795579433, 0.11938491463661194, 0.03761400282382965, 0.10612092912197113, 0.004111893475055695, 0.07536520808935165, 0.06150262430310249, 0.010061400011181831, 0.01712355576455593, 0.026476707309484482, 0.05440329760313034, 0.37643373012542725, 0.12204637378454208, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0973815768957138, 0.1330094188451767, 0.2356250286102295, 0.23801013827323914, 0.16962124407291412, 0.3808935284614563, 0.19062454998493195, 0.12487400323152542, 0.4241224527359009, 0.1858355700969696, 0.1843334436416626, 0.17186462879180908, 0.1674181967973709, 0.03679514676332474, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.28161293268203735, 0.39586660265922546, 0.35408592224121094, 0.26687130331993103, 0.036089953035116196, 0.12106626480817795, 0.05175312981009483, 0.6374836564064026, 0.06537415832281113, 0.01867927983403206, 0.03261437267065048, 0.05161871388554573, 0.026679201051592827, 0.0063977655954658985, 0.0581950880587101, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.052721865475177765, 0.30848002433776855, 0.24953237175941467, 0.2790854275226593, 0.7654650807380676, 0.6871634125709534, 0.13210926949977875, 0.673875629901886, 0.04467727988958359, 0.018614191561937332, 0.08283445239067078, 0.0906965509057045, 0.06073237210512161, 0.12131030112504959, 0.06997358053922653, 0.3489122688770294, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03943483531475067, 0.28613966703414917, 0.07243800908327103, 0.8744964599609375, 0.029915155842900276, 0.331167072057724, 0.4079437255859375, 0.5431530475616455, 0.3259604275226593, 0.1150238886475563, 0.3324905335903168, 0.44221389293670654, 0.2450132817029953, 0.12577538192272186, 0.11014749854803085, 0.1900990903377533, 0.042790502309799194, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.06558705866336823, 0.020870981737971306, 0.007642277050763369, 0.028054187074303627, 0.010532653890550137, 0.10334379225969315, 0.12033270299434662, 0.1911371499300003, 0.30930495262145996, 0.04741071164608002, 0.06516209989786148, 0.09313901513814926, 0.24243950843811035, 0.15116305649280548, 0.09231718629598618, 0.47254911065101624, 0.053373783826828, 0.18162642419338226, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.017124762758612633, 0.00014164860476739705, 0.01482362300157547, 0.13952724635601044, 0.0008921221597120166, 0.07150562852621078, 0.037848807871341705, 0.0009583857608959079, 0.0160027127712965, 0.01657933183014393, 0.09754330664873123, 0.3402610719203949, 0.02766183763742447, 0.011668790131807327, 0.019427720457315445, 0.01879642903804779, 0.06977814435958862, 0.23379765450954437, 0.41046860814094543, NaN, NaN, NaN, NaN, NaN, NaN], [0.0033047832548618317, 0.043024010956287384, 0.009507044218480587, 0.05758155509829521, 0.0012058177962899208, 0.04777836054563522, 0.038867104798555374, 0.0027761561796069145, 0.008453112095594406, 0.011027430184185505, 0.021058345213532448, 0.3453521430492401, 0.05058252438902855, 0.004837945103645325, 0.0014179014833644032, 0.06873936206102371, 0.10687354952096939, 0.21186815202236176, 0.44615596532821655, 0.10872229933738708, NaN, NaN, NaN, NaN, NaN], [0.05260666832327843, 0.09784732013940811, 0.08957145363092422, 0.40504154562950134, 0.2393025904893875, 0.37446328997612, 0.33926665782928467, 0.06915906071662903, 0.28494811058044434, 0.18951286375522614, 0.21801336109638214, 0.2963850796222687, 0.09700386226177216, 0.02254888415336609, 0.016780056059360504, 0.3380737006664276, 0.17247304320335388, 0.15711140632629395, 0.27414536476135254, 0.12462585419416428, 0.05461693927645683, NaN, NaN, NaN, NaN], [0.4168609082698822, 0.5786882042884827, 0.4795728027820587, 0.4880480170249939, 0.07741907238960266, 0.22295767068862915, 0.10229793190956116, 0.7397969365119934, 0.09120289236307144, 0.02111845649778843, 0.040493883192539215, 0.06478337198495865, 0.029333919286727905, 0.01266437117010355, 0.08807221800088882, 0.12442159652709961, 0.019878262653946877, 0.02248454838991165, 0.045759230852127075, 0.02396523579955101, 0.002620323793962598, 0.04143214225769043, NaN, NaN, NaN], [0.05813424289226532, 0.29987069964408875, 0.06046860292553902, 0.2948205769062042, 0.6036045551300049, 0.4684220552444458, 0.10851431638002396, 0.5970842242240906, 0.03630568087100983, 0.009022231213748455, 0.034897517412900925, 0.044963937252759933, 0.06918716430664062, 0.06464210897684097, 0.027029458433389664, 0.39741793274879456, 0.1858920007944107, 0.0860959067940712, 0.03553689271211624, 0.03651457652449608, 0.07401836663484573, 0.02850046567618847, 0.457316130399704, NaN, NaN], [0.011862307786941528, 0.06274299323558807, 0.019264375790953636, 0.7077140212059021, 0.009838010184466839, 0.08938813954591751, 0.2665976285934448, 0.21134285628795624, 0.19931168854236603, 0.029879093170166016, 0.11873869597911835, 0.2187809944152832, 0.10740162432193756, 0.03893040865659714, 0.02778119407594204, 0.17118902504444122, 0.03705315291881561, 0.41107529401779175, 0.3035467863082886, 0.1782693862915039, 0.062172479927539825, 0.04369974508881569, 0.43116021156311035, 0.04090215638279915, NaN], [0.13294808566570282, 0.07747184485197067, 0.06700501590967178, 0.24500344693660736, 0.07035010308027267, 0.06088097393512726, 0.15465889871120453, 0.22422827780246735, 0.20946520566940308, 0.06346394866704941, 0.1416163444519043, 0.10671631991863251, 0.07756247371435165, 0.14874279499053955, 0.2551397681236267, 0.18877547979354858, 0.07302238047122955, 0.24805422127246857, 0.1228112131357193, 0.08095405995845795, 0.12022056430578232, 0.20888803899288177, 0.1654488444328308, 0.07207347452640533, 0.12261014431715012]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04818185046315193, 0.30147239565849304, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.000490668579004705, 0.5364181399345398, 0.0016803600592538714, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17249688506126404, 0.003960400819778442, 1.1815190191555303e-05, 0.00205309153534472, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08659190684556961, 0.2260276973247528, 0.018877657130360603, 0.019257033243775368, 0.9179584980010986, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [2.1155383365112357e-05, 0.00016346832853741944, 0.0004644138098228723, 9.852640505414456e-05, 0.009302367456257343, 0.8758521676063538, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0038963633123785257, 0.11578002572059631, 0.06833135336637497, 0.2930091321468353, 0.06728219240903854, 0.588379442691803, 0.190787211060524, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04113525524735451, 0.03917931765317917, 0.013817446306347847, 0.06874216347932816, 0.027753230184316635, 0.04752122610807419, 0.17637789249420166, 0.2964049279689789, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.006397286430001259, 0.008155078627169132, 0.02385183423757553, 0.08218340575695038, 0.09733399748802185, 0.7216709852218628, 0.11420661956071854, 0.028804002329707146, 0.49512770771980286, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.007080267183482647, 0.010165071114897728, 0.007166726514697075, 0.04547898843884468, 0.014898931607604027, 0.06153866648674011, 0.05960511788725853, 0.025653565302491188, 0.05574938654899597, 0.5054050087928772, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12821261584758759, 0.09823491424322128, 0.2407415509223938, 0.03722868487238884, 0.07500484585762024, 0.23719841241836548, 0.08696958422660828, 0.10033686459064484, 0.08637046813964844, 0.05946339666843414, 0.17889682948589325, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.018611561506986618, 0.530681848526001, 0.37442806363105774, 0.09326046705245972, 0.039934538304805756, 0.607749342918396, 0.1011725440621376, 0.041957128793001175, 0.061673425137996674, 0.012941170483827591, 0.012897199019789696, 0.02531522512435913, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.025258230045437813, 0.013820141553878784, 0.020238902419805527, 0.20186173915863037, 0.008764497935771942, 0.044081512838602066, 0.11685895919799805, 0.12131167203187943, 0.03466574102640152, 0.0033257410395890474, 0.009427645243704319, 0.00932170171290636, 0.6215367317199707, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0027034373488277197, 0.008653531782329082, 0.0021412167698144913, 0.02395743690431118, 0.06537352502346039, 0.05110874027013779, 0.050060901790857315, 0.023448945954442024, 0.0059632728807628155, 0.0016337132547050714, 0.0060929651372134686, 0.00957516860216856, 0.05008334666490555, 0.696637749671936, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [5.63129390229733e-07, 0.00027805642457678914, 1.7160025890916586e-05, 5.958595011179568e-06, 0.00078710971865803, 1.2566613349918043e-06, 9.03528507478768e-06, 2.1993335394654423e-05, 4.528845238382928e-06, 1.0594538935038145e-06, 2.375837993895402e-06, 1.0765622391772922e-05, 0.00012861557479482144, 0.000270194374024868, 0.4203896224498749, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.19651824235916138, 0.009276115335524082, 0.0007576652569696307, 0.02043321169912815, 0.000937489268835634, 0.0014158851699903607, 0.02691410481929779, 0.025149332359433174, 0.015754513442516327, 0.002638434525579214, 0.03568584471940994, 0.28478676080703735, 0.08937329053878784, 0.04057440906763077, 0.41798362135887146, 0.02812151424586773, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0009883381426334381, 0.005475975573062897, 0.017872320488095284, 0.0038598645478487015, 0.01383217889815569, 0.1060260757803917, 0.010558119975030422, 0.0004280287539586425, 0.011488020420074463, 0.004323506727814674, 0.015877770259976387, 0.025533713400363922, 0.06758329272270203, 0.005362953990697861, 0.03033292666077614, 0.3987913429737091, 0.22715723514556885, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.025437461212277412, 0.027387555688619614, 0.0211916733533144, 0.0013409400125965476, 0.0016278955154120922, 0.0205780491232872, 0.006606978829950094, 0.005105526186525822, 0.008417481556534767, 0.008475488983094692, 0.016475802287459373, 0.021865585818886757, 0.04041945934295654, 0.001965513452887535, 0.030297037214040756, 0.018051480874419212, 0.2940014600753784, 0.09546513855457306, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.014116446487605572, 0.6685785055160522, 0.40577325224876404, 0.09365412592887878, 0.008716625161468983, 0.504762589931488, 0.11037815362215042, 0.03693895787000656, 0.066362664103508, 0.025546396151185036, 0.030971869826316833, 0.07333581149578094, 0.21910515427589417, 0.03128749132156372, 0.013437384739518166, 0.06674141436815262, 0.055549826472997665, 0.02615067921578884, 0.05289305001497269, NaN, NaN, NaN, NaN, NaN, NaN], [0.01752244122326374, 0.013681006617844105, 0.015325021930038929, 0.15400148928165436, 0.0017620606813579798, 0.03783759847283363, 0.07285356521606445, 0.042190372943878174, 0.019725583493709564, 0.004497688263654709, 0.010335608385503292, 0.023485884070396423, 0.5969190001487732, 0.22785267233848572, 0.05655405670404434, 0.05765213817358017, 0.006416310556232929, 0.029401889070868492, 0.022928474470973015, 0.6468356251716614, NaN, NaN, NaN, NaN, NaN], [0.003705248236656189, 0.09392052888870239, 0.0011726000811904669, 0.042238909751176834, 0.07787514477968216, 0.11800158768892288, 0.09318403154611588, 0.018972182646393776, 0.022339271381497383, 0.02290215529501438, 0.009648749604821205, 0.020298194140195847, 0.09632600843906403, 0.6665039658546448, 0.01913357712328434, 0.016501925885677338, 0.01550414226949215, 0.014767719432711601, 0.035943012684583664, 0.1298983097076416, 0.7307590246200562, NaN, NaN, NaN, NaN], [3.2450822118335054e-07, 0.0001958437787834555, 1.195628647110425e-05, 3.192948497598991e-06, 0.00034392892848700285, 1.3818779507346335e-06, 6.319523890851997e-06, 9.25252061279025e-06, 3.2897685287025524e-06, 1.041492623699014e-06, 2.450263082209858e-06, 1.1291336704744026e-05, 9.216016042046249e-05, 0.00025747373001649976, 0.3770022690296173, 7.494814053643495e-05, 0.00011931787594221532, 5.454379424918443e-05, 3.481862586340867e-05, 0.0001493972522439435, 6.532184488605708e-05, 0.4379080533981323, NaN, NaN, NaN], [0.11172444373369217, 0.00812594499439001, 0.000803561822976917, 0.011673782020807266, 0.00013412271800916642, 0.002435607835650444, 0.021002406254410744, 0.009926681406795979, 0.014218374155461788, 0.0044799866154789925, 0.03462693840265274, 0.49634605646133423, 0.1610735058784485, 0.03537029027938843, 0.3717024624347687, 0.0470024012029171, 0.0025306264869868755, 0.08426976948976517, 0.5137573480606079, 0.047759927809238434, 0.008752438239753246, 0.5270217657089233, 0.020567137748003006, NaN, NaN], [0.00039373920299112797, 0.00142151047475636, 0.016346368938684464, 0.0038184949662536383, 0.00426360173150897, 0.10012070834636688, 0.007060237228870392, 0.00022489627008326352, 0.006389277055859566, 0.0014407823327928782, 0.01344740204513073, 0.019176417961716652, 0.04953484237194061, 0.003102741902694106, 0.017501499503850937, 0.25968801975250244, 0.12805432081222534, 0.03450275957584381, 0.03214799612760544, 0.06495527178049088, 0.007038496434688568, 0.018200475722551346, 0.2228115350008011, 0.24082934856414795, NaN], [0.004585978575050831, 0.008592751808464527, 0.20804427564144135, 0.003501898143440485, 0.01809401623904705, 0.0088487658649683, 0.01839679665863514, 0.009930659085512161, 0.019693726673722267, 0.015943868085741997, 0.06719032675027847, 0.03678698092699051, 0.03292753919959068, 0.02313893660902977, 0.023240724578499794, 0.03294161707162857, 0.24390928447246552, 0.10472099483013153, 0.0623757429420948, 0.06489475816488266, 0.03424002602696419, 0.03615953400731087, 0.05666068568825722, 0.29077935218811035, 0.20903274416923523]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15880486369132996, 0.04734092205762863, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.22883240878582, 0.015307039953768253, 0.023610780015587807, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.15376803278923035, 0.17623378336429596, 0.16427822411060333, 0.018553992733359337, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12576976418495178, 0.44071146845817566, 0.38860467076301575, 0.12043511122465134, 0.027116619050502777, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03928220644593239, 0.42239660024642944, 0.2546820342540741, 0.22367709875106812, 0.1215892881155014, 0.001983387628570199, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17152060568332672, 0.49365419149398804, 0.08085957914590836, 0.02207508496940136, 0.19231174886226654, 0.008304901421070099, 0.03878962993621826, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.13843253254890442, 0.07047099620103836, 0.2525072991847992, 0.13487939536571503, 0.27911728620529175, 0.11727599054574966, 0.022392159327864647, 0.1764850914478302, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10915631055831909, 0.30942168831825256, 0.19657404720783234, 0.031007295474410057, 0.23716343939304352, 0.05435822904109955, 0.08149112015962601, 0.6613667011260986, 0.11670006066560745, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.0640818402171135, 0.41535088419914246, 0.29784247279167175, 0.05657188221812248, 0.036311421543359756, 0.08192699402570724, 0.16688455641269684, 0.10144203901290894, 0.346017450094223, 0.15466110408306122, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04877842590212822, 0.16450235247612, 0.23761717975139618, 0.0720985159277916, 0.12954245507717133, 0.08035153150558472, 0.18124118447303772, 0.05973014980554581, 0.26483285427093506, 0.39028850197792053, 0.05098416656255722, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11044558137655258, 0.08550350368022919, 0.2513507902622223, 0.28401821851730347, 0.12441904842853546, 0.05029991641640663, 0.42405593395233154, 0.08374682813882828, 0.43869927525520325, 0.14253327250480652, 0.10876792669296265, 0.09369473904371262, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.08764015138149261, 0.46941375732421875, 0.23278135061264038, 0.11763583868741989, 0.0354606918990612, 0.16624747216701508, 0.2793619632720947, 0.1965668648481369, 0.23052528500556946, 0.3914787769317627, 0.08669382333755493, 0.10678009688854218, 0.08708767592906952, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.2116944044828415, 0.06720030307769775, 0.29984304308891296, 0.010844358243048191, 0.051072586327791214, 0.15023349225521088, 0.04554526135325432, 0.1560167670249939, 0.03609438240528107, 0.026584016159176826, 0.14512087404727936, 0.05890262499451637, 0.015816861763596535, 0.07422769069671631, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.056502565741539, 0.15541820228099823, 0.07158821076154709, 0.00490804947912693, 0.015012365765869617, 0.06302572786808014, 0.01116714347153902, 0.22065599262714386, 0.021468764171004295, 0.01365464273840189, 0.022816751152276993, 0.019708380103111267, 0.0059420084580779076, 0.0700121819972992, 0.287899911403656, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.058403778821229935, 0.0693131536245346, 0.04999461770057678, 0.004054869059473276, 0.0624610111117363, 0.018093721941113472, 0.07961009442806244, 0.1545858234167099, 0.3008257746696472, 0.14455094933509827, 0.09800520539283752, 0.09531621634960175, 0.27401015162467957, 0.4782770574092865, 0.11211755871772766, 0.01358953770250082, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1399688720703125, 0.5559014678001404, 0.20350231230258942, 0.042011573910713196, 0.020507201552391052, 0.03915366902947426, 0.4243565797805786, 0.11376935243606567, 0.31140708923339844, 0.051479678601026535, 0.07416504621505737, 0.2654426097869873, 0.3960915207862854, 0.5790604948997498, 0.18063338100910187, 0.1939544379711151, 0.04191381484270096, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.027515297755599022, 0.0486784465610981, 0.06845460832118988, 0.023408811539411545, 0.008863206952810287, 0.008533195592463017, 0.24178741872310638, 0.01229054294526577, 0.25817692279815674, 0.6869812607765198, 0.049950506538152695, 0.12178820371627808, 0.0564231351017952, 0.02026011236011982, 0.004908477421849966, 0.03562311828136444, 0.12746450304985046, 0.0016219470417127013, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.11620164662599564, 0.09937138110399246, 0.17538107931613922, 0.40406307578086853, 0.043817292898893356, 0.05759625509381294, 0.49306368827819824, 0.09120260924100876, 0.36450278759002686, 0.08042807132005692, 0.1856311559677124, 0.1376025527715683, 0.1998283714056015, 0.3654005527496338, 0.15910619497299194, 0.4969707429409027, 0.08565060794353485, 0.02514367550611496, 0.090617336332798, NaN, NaN, NaN, NaN, NaN, NaN], [0.0739481970667839, 0.5182103514671326, 0.19721719622612, 0.21118015050888062, 0.015751224011182785, 0.12249443680047989, 0.5174803733825684, 0.17075838148593903, 0.30025264620780945, 0.29246312379837036, 0.0875946432352066, 0.2326347827911377, 0.13986286520957947, 0.511695921421051, 0.12602318823337555, 0.03662485629320145, 0.1263200044631958, 0.0166145209223032, 0.19702456891536713, 0.09621746093034744, NaN, NaN, NaN, NaN, NaN], [0.3052336871623993, 0.37224864959716797, 0.45515015721321106, 0.04986808821558952, 0.05332064628601074, 0.13846120238304138, 0.15990367531776428, 0.20659208297729492, 0.06640873104333878, 0.035323526710271835, 0.30340465903282166, 0.10174556821584702, 0.02102985605597496, 0.11508277803659439, 0.09203195571899414, 0.0029288395307958126, 0.023838462308049202, 0.004605103749781847, 0.052648112177848816, 0.006431906949728727, 0.026736242696642876, NaN, NaN, NaN, NaN], [0.047024402767419815, 0.1257133185863495, 0.052377521991729736, 0.009844984859228134, 0.015597687102854252, 0.06965665519237518, 0.01849394477903843, 0.1603521853685379, 0.02587857097387314, 0.00957732368260622, 0.023523790761828423, 0.020081259310245514, 0.008425970561802387, 0.10955916345119476, 0.35300737619400024, 0.023505402728915215, 0.00786643661558628, 0.007557017263025045, 0.013908758759498596, 0.004675114993005991, 0.035296451300382614, 0.3261549174785614, NaN, NaN, NaN], [0.11014947295188904, 0.08461853116750717, 0.02981843426823616, 0.004099451471120119, 0.009237504564225674, 0.011130756698548794, 0.132149338722229, 0.11619938164949417, 0.22203940153121948, 0.02292616292834282, 0.06793706119060516, 0.07227552682161331, 0.3262397348880768, 0.40601006150245667, 0.08270477503538132, 0.013506797142326832, 0.03135772421956062, 0.07034049183130264, 0.09623772650957108, 0.20842698216438293, 0.2752794623374939, 0.1234828308224678, 0.04129752516746521, NaN, NaN], [0.1182219609618187, 0.7384620308876038, 0.11492461711168289, 0.09884578734636307, 0.012010940350592136, 0.038200050592422485, 0.4905328154563904, 0.23439669609069824, 0.2528713345527649, 0.015177865512669086, 0.07817362248897552, 0.33532261848449707, 0.4971323609352112, 0.7384514212608337, 0.2383432686328888, 0.2306600660085678, 0.025716517120599747, 0.023198120296001434, 0.3352215886116028, 0.4797173738479614, 0.5688640475273132, 0.2555003762245178, 0.1890360713005066, 0.06237812712788582, NaN], [0.13153354823589325, 0.5476850867271423, 0.27465543150901794, 0.27658137679100037, 0.5121651291847229, 0.3939417600631714, 0.2527337968349457, 0.41937416791915894, 0.2437492311000824, 0.1485103964805603, 0.10651403665542603, 0.241710364818573, 0.34289923310279846, 0.3691290616989136, 0.108230821788311, 0.32214298844337463, 0.08876177668571472, 0.03369928151369095, 0.23942533135414124, 0.302080899477005, 0.3531237244606018, 0.09724070131778717, 0.19267186522483826, 0.06874143332242966, 0.052875734865665436]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03290099650621414, 0.3365767002105713, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.003547579748556018, 0.004082763101905584, 0.4616691768169403, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03595791012048721, 0.1313885897397995, 0.007101066876202822, 0.42131781578063965, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.007601147051900625, 0.014137630350887775, 0.01938864029943943, 0.2572920322418213, 0.0011994435917586088, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.00011468974116723984, 0.0032473355531692505, 0.00037737423554062843, 0.2793608605861664, 0.003465541172772646, 5.061212868895382e-05, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.21311266720294952, 0.10434294492006302, 0.011484598740935326, 0.0013334749964997172, 0.03845251351594925, 0.028238367289304733, 0.05654546618461609, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.052184704691171646, 0.499632865190506, 0.005138374865055084, 0.10169705748558044, 0.09997230768203735, 0.036990027874708176, 0.07566682249307632, 0.32418423891067505, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.23645982146263123, 0.016864946112036705, 0.013305210508406162, 0.0007752762176096439, 0.017555342987179756, 0.03100133314728737, 0.04085567593574524, 0.029846351593732834, 0.010373883880674839, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18805328011512756, 0.046367619186639786, 0.10314629226922989, 0.018223291262984276, 0.27720585465431213, 0.3798944056034088, 0.09291481226682663, 0.09293034672737122, 0.04290880635380745, 0.03370373696088791, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.028641005977988243, 0.03295213729143143, 0.0065453751012682915, 0.16686026751995087, 0.028714975342154503, 0.015397193841636181, 0.02003423683345318, 0.019093815237283707, 0.020523719489574432, 0.016172079369425774, 0.3490104377269745, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.10839971899986267, 0.004465002100914717, 0.016082070767879486, 0.035488102585077286, 0.015600458718836308, 0.012030484154820442, 0.015872180461883545, 0.01552913524210453, 0.03533920273184776, 0.11401902139186859, 0.31523072719573975, 0.20448055863380432, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.18776558339595795, 0.0060520414263010025, 0.017473671585321426, 0.005528539884835482, 0.0027145782951265574, 0.012176988646388054, 0.0031525399535894394, 0.004637573380023241, 0.011988476850092411, 0.06979440897703171, 0.38327983021736145, 0.020156072452664375, 0.010166948661208153, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.3064975440502167, 0.004262991715222597, 0.009997943416237831, 0.00034317225799895823, 0.013912403024733067, 0.02852706052362919, 0.004078225698322058, 0.001928618410602212, 0.006367305759340525, 0.035507142543792725, 0.050674788653850555, 0.007057875394821167, 0.0049485149793326855, 0.0049379738047719, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14988604187965393, 0.015584584325551987, 0.137997567653656, 0.0031439096201211214, 0.5546696782112122, 0.01658078096807003, 0.0025873971171677113, 0.0010246702004224062, 0.019667595624923706, 0.012580120004713535, 0.015491531230509281, 0.029023459181189537, 0.021588340401649475, 0.25595030188560486, 0.02325037308037281, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.07357528805732727, 0.007756352424621582, 0.002724927617236972, 0.001402079127728939, 0.0004431438574101776, 0.00010925461538136005, 0.0029409730341285467, 0.005563507787883282, 0.012139370664954185, 0.03890732303261757, 0.05558362230658531, 0.03318313509225845, 0.4270496368408203, 0.07112571597099304, 0.15036046504974365, 0.020786603912711143, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.012120572850108147, 0.0003307444858364761, 0.009640182368457317, 0.00017808230768423527, 0.0021490382496267557, 0.0008148089982569218, 0.0008481521508656442, 0.0019973982125520706, 0.005024890415370464, 0.01719486527144909, 0.044799502938985825, 0.006444229744374752, 0.018026985228061676, 0.0067391968332231045, 0.061299871653318405, 0.01281613577157259, 0.3084925711154938, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.011204708367586136, 0.0033799665980041027, 0.008117830380797386, 0.1567971557378769, 0.012545537203550339, 0.002854604972526431, 0.0037395430263131857, 0.0003391341888345778, 0.002928558737039566, 0.004266565665602684, 0.28180748224258423, 0.005543314386159182, 0.0059068226255476475, 0.004401014186441898, 0.09436267614364624, 0.003524675266817212, 0.09697568416595459, 0.3818984925746918, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1085091158747673, 0.0013132937019690871, 0.011304548010230064, 0.014309195801615715, 0.009265521541237831, 0.00682368129491806, 0.01179590355604887, 0.005223054438829422, 0.01697726733982563, 0.05782441794872284, 0.2522926330566406, 0.16053971648216248, 0.020927468314766884, 0.02051178365945816, 0.1114674061536789, 0.014847181737422943, 0.40623563528060913, 0.12017090618610382, 0.2281613051891327, NaN, NaN, NaN, NaN, NaN, NaN], [0.23926517367362976, 0.007461922243237495, 0.015478387475013733, 0.02120528556406498, 0.0046339076943695545, 0.01287792343646288, 0.005305645987391472, 0.0037130024284124374, 0.011430526152253151, 0.10132863372564316, 0.42019084095954895, 0.03134358674287796, 0.006659360136836767, 0.0015345009742304683, 0.05340040102601051, 0.0021821516565978527, 0.15366847813129425, 0.09343723207712173, 0.04055917635560036, 0.009410854429006577, NaN, NaN, NaN, NaN, NaN], [0.3882482349872589, 0.012203006073832512, 0.008404962718486786, 0.0008633172838017344, 0.07213836163282394, 0.03903299570083618, 0.006879106629639864, 0.0025245456490665674, 0.011604986153542995, 0.1302306056022644, 0.05970751494169235, 0.005057368893176317, 0.0025832061655819416, 0.003548768814653158, 0.03821956738829613, 0.0041786422953009605, 0.029319334775209427, 0.009258194826543331, 0.010013489983975887, 0.0024901984725147486, 0.009316755458712578, NaN, NaN, NaN, NaN], [0.08333727717399597, 0.009125825949013233, 0.12352871894836426, 0.0034849271178245544, 0.49194949865341187, 0.008760062977671623, 0.002427457133308053, 0.0004761714953929186, 0.014378424733877182, 0.007653949782252312, 0.010163314640522003, 0.018072640523314476, 0.014914281666278839, 0.33540958166122437, 0.012212751433253288, 0.050671979784965515, 0.08942927420139313, 0.0058481828309595585, 0.02088618278503418, 0.013520943000912666, 0.3026564419269562, 0.011637967079877853, NaN, NaN, NaN], [0.019913960248231888, 0.003490668721497059, 0.00020567848696373403, 0.00036819992237724364, 0.00019341551524121314, 3.8652269722661003e-05, 0.0008544524316675961, 0.002890991745516658, 0.001110991695895791, 0.005157719366252422, 0.008338885381817818, 0.0030357406940311193, 0.14557099342346191, 0.021602485328912735, 0.04367346689105034, 0.0015647107502445579, 0.009655454196035862, 0.14827704429626465, 0.008163533173501492, 0.49237948656082153, 0.06938102096319199, 0.08394628763198853, 0.049248531460762024, NaN, NaN], [0.010580360889434814, 0.00023049254377838224, 0.00745873898267746, 0.00016025979130063206, 0.002226235345005989, 0.0004258991975802928, 0.000578688399400562, 0.0014760587364435196, 0.002039685845375061, 0.0048048608005046844, 0.019996320828795433, 0.0029125709552317858, 0.006709430366754532, 0.0017099445685744286, 0.02097223326563835, 0.0024284888058900833, 0.10361000150442123, 0.022238893434405327, 0.009704988449811935, 0.017071064561605453, 0.011506098322570324, 0.0406200997531414, 0.0063119689002633095, 0.36112311482429504, NaN], [0.07011571526527405, 0.029766615480184555, 0.05616272985935211, 0.02569880336523056, 0.02553572878241539, 0.010698755271732807, 0.02022577077150345, 0.01824677176773548, 0.03918607532978058, 0.034657131880521774, 0.11515442281961441, 0.05569382756948471, 0.035370998084545135, 0.047812946140766144, 0.1140216588973999, 0.018943075090646744, 0.09709078818559647, 0.08172454684972763, 0.04602199047803879, 0.02941049635410309, 0.031383853405714035, 0.10708537697792053, 0.012693268246948719, 0.07050468772649765, 0.25427982211112976]], [[0.125, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1627129465341568, 0.03836298733949661, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.23664157092571259, 0.02332315407693386, 0.0017523575806990266, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.14284735918045044, 0.19342879951000214, 0.5212197303771973, 0.028613613918423653, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.022152410820126534, 0.06252314150333405, 0.005122532602399588, 0.24202540516853333, 0.0027534610126167536, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.04657726734876633, 0.23517371714115143, 0.03296450525522232, 0.2014523595571518, 0.06359406560659409, 0.0884864553809166, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.05186963453888893, 0.02286554127931595, 0.21517929434776306, 0.12055587023496628, 0.1711670458316803, 0.27492430806159973, 0.27398592233657837, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.020278872922062874, 0.02308776043355465, 0.022820638492703438, 0.18259893357753754, 0.3133871257305145, 0.08183155953884125, 0.35655686259269714, 0.17295894026756287, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.057175230234861374, 0.2799927890300751, 0.10977934300899506, 0.4680712819099426, 0.08838099986314774, 0.05264464393258095, 0.21108192205429077, 0.08241217583417892, 0.0764400064945221, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.17679302394390106, 0.30970489978790283, 0.042192552238702774, 0.2463400512933731, 0.032756272703409195, 0.05394153669476509, 0.02321716584265232, 0.30038926005363464, 0.023974716663360596, 0.0257905051112175, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1864403486251831, 0.03811780363321304, 0.18074536323547363, 0.08396673202514648, 0.026499373838305473, 0.05736878141760826, 0.274480402469635, 0.10284627228975296, 0.15606749057769775, 0.017497936263680458, 0.09719526022672653, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.1767420768737793, 0.017465414479374886, 0.034512054175138474, 0.0999627411365509, 0.011741198599338531, 0.022724410519003868, 0.04408577084541321, 0.03894393891096115, 0.018038587644696236, 0.058924250304698944, 0.2522818148136139, 0.12782295048236847, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.042104240506887436, 0.022070694714784622, 0.04743226245045662, 0.13338083028793335, 0.020831480622291565, 0.031267598271369934, 0.024703562259674072, 0.041907425969839096, 0.006121364887803793, 0.02875565178692341, 0.13002096116542816, 0.36194902658462524, 0.021867850795388222, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.12623563408851624, 0.6370776891708374, 0.07802888005971909, 0.06076015904545784, 0.015353387221693993, 0.0031011439859867096, 0.031844403594732285, 0.5665289163589478, 0.013176449574530125, 0.025442441925406456, 0.05083877220749855, 0.08586791157722473, 0.03281332179903984, 0.0019294946687296033, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.010483458638191223, 0.10243765264749527, 0.013204336166381836, 0.1070198118686676, 0.001742976950481534, 0.0011925535509362817, 0.03764529153704643, 0.023008054122328758, 0.09038762003183365, 0.1208486333489418, 0.06097627431154251, 0.11476689577102661, 0.17706690728664398, 0.4447736442089081, 0.005561552010476589, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.03962688520550728, 0.412600040435791, 0.1027907133102417, 0.011060677468776703, 0.04006139934062958, 0.005457504652440548, 0.17391063272953033, 0.009697728790342808, 0.08243320137262344, 0.1504840850830078, 0.029468167573213577, 0.29366523027420044, 0.04788699373602867, 0.17640100419521332, 0.04229334741830826, 0.3300667107105255, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.20544184744358063, 0.06503231078386307, 0.21778742969036102, 0.04011436551809311, 0.2470238208770752, 0.03102266602218151, 0.027881061658263206, 0.06887322664260864, 0.023802783340215683, 0.2166331559419632, 0.06618232280015945, 0.058350641280412674, 0.04297764599323273, 0.06574989855289459, 0.02652076631784439, 0.08339553326368332, 0.09817715734243393, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.09466058760881424, 0.0047309016808867455, 0.1481417566537857, 0.06127317249774933, 0.015202163718640804, 0.011932089924812317, 0.31230586767196655, 0.04852164536714554, 0.039501819759607315, 0.001117925625294447, 0.06312739849090576, 0.023924386128783226, 0.02860989049077034, 0.007241260260343552, 0.11453913897275925, 0.012237192131578922, 0.2803768217563629, 0.0480632521212101, NaN, NaN, NaN, NaN, NaN, NaN, NaN], [0.02001449465751648, 0.0017837424529716372, 0.005722085013985634, 0.04321253299713135, 0.00430489843711257, 0.009005578234791756, 0.010736249387264252, 0.0058517144061625, 0.003792154835537076, 0.008828205987811089, 0.0838593989610672, 0.029530486091971397, 0.015579215250909328, 0.010320665314793587, 0.016853220760822296, 0.017335176467895508, 0.12552303075790405, 0.42354699969291687, 0.08326870948076248, NaN, NaN, NaN, NaN, NaN, NaN], [0.001771818962879479, 0.000807587115559727, 0.0031146325636655092, 0.023062998428940773, 0.0018312688916921616, 0.007724495604634285, 0.002569216303527355, 0.003803644794970751, 0.00041838324978016317, 0.001987496856600046, 0.012477965094149113, 0.04809670150279999, 0.0016458284808322787, 0.00020838514319621027, 0.005814890842884779, 0.018183711916208267, 0.30546146631240845, 0.4703490138053894, 0.15369661152362823, 0.012250960804522038, NaN, NaN, NaN, NaN, NaN], [0.02520398050546646, 0.2818087637424469, 0.007948609068989754, 0.07590723037719727, 0.01867567002773285, 0.006826441269367933, 0.011762343347072601, 0.5987983345985413, 0.0045673479326069355, 0.01173742488026619, 0.03130093589425087, 0.03894692659378052, 0.016236862167716026, 0.0014989122282713652, 0.0009245824767276645, 0.025562506169080734, 0.5276230573654175, 0.32699310779571533, 0.1864093542098999, 0.0933799296617508, 0.0060149896889925, NaN, NaN, NaN, NaN], [0.0011320068733766675, 0.011502433568239212, 0.0017513524508103728, 0.020418671891093254, 0.0003008104977197945, 0.00031320590642280877, 0.0053228470496833324, 0.0022876623552292585, 0.011736828833818436, 0.017109515145421028, 0.010937619023025036, 0.015238909050822258, 0.025703608989715576, 0.10705357789993286, 0.0009204442030750215, 0.02667400799691677, 0.16934601962566376, 0.08647502958774567, 0.028284918516874313, 0.06841914355754852, 0.39870724081993103, 0.0010592876933515072, NaN, NaN, NaN], [0.02631283551454544, 0.29101136326789856, 0.042160265147686005, 0.009721376933157444, 0.02933679334819317, 0.014515053480863571, 0.18161341547966003, 0.016545770689845085, 0.03647695854306221, 0.0840071588754654, 0.02240183763206005, 0.1055113896727562, 0.037331126630306244, 0.17535105347633362, 0.010923052206635475, 0.2594170868396759, 0.5064816474914551, 0.06657205522060394, 0.130835622549057, 0.0483754500746727, 0.2870587110519409, 0.010685333050787449, 0.21122200787067413, NaN, NaN], [0.21289733052253723, 0.10400458425283432, 0.2843308448791504, 0.11722961068153381, 0.31265783309936523, 0.07705509662628174, 0.050357937812805176, 0.1631784737110138, 0.04547655209898949, 0.37539371848106384, 0.07925810664892197, 0.07719646394252777, 0.043498191982507706, 0.04735783487558365, 0.022911155596375465, 0.20965908467769623, 0.2452480047941208, 0.05793433263897896, 0.07357832789421082, 0.03363368287682533, 0.041085004806518555, 0.014093895442783833, 0.05045074224472046, 0.0570731945335865, NaN], [0.02115148864686489, 0.018139760941267014, 0.03536282852292061, 0.06259438395500183, 0.00901759136468172, 0.014575985260307789, 0.12521256506443024, 0.12870429456233978, 0.09162478893995285, 0.06363746523857117, 0.1348179280757904, 0.07700010389089584, 0.05158444121479988, 0.01101324986666441, 0.03299920633435249, 0.163722425699234, 0.13794326782226562, 0.18303781747817993, 0.117555633187294, 0.08103907853364944, 0.012191864661872387, 0.032527241855859756, 0.16104964911937714, 0.12187117338180542, 0.22321484982967377]]]], \"bot_text\": [\"Das_\", \"Tier\", \"_\", \"\\u00fcber\", \"quer\", \"te_\", \"die_\", \"Stra\\u00dfe_\", \"nicht_\", \", _\", \"weil_\", \"es_\", \"zu_\", \"m\\u00fc\", \"de_\", \"war_\", \", _\", \"weil_\", \"es_\", \"zu_\", \"m\\u00fc\", \"de_\", \"war_\", \"._\"]}, \"inp_inp\": {\"top_text\": [\"The_\", \"animal_\", \"didn_\", \"'_\", \"t_\", \"cross_\", \"the_\", \"street_\", \"because_\", \"it_\", \"was_\", \"too_\", \"tire\", \"d_\"], \"att\": [[[[0.04540494084358215, 0.009098929353058338, 0.06841860711574554, 0.050027038902044296, 0.1867244392633438, 0.20893266797065735, 0.15536439418792725, 0.2501838803291321, 0.03253718465566635, 0.045193806290626526, 0.01405471283942461, 0.15126678347587585, 0.5554144382476807, 0.07120772451162338, 0.21479088068008423], [0.010880604386329651, 0.008569094352424145, 0.3644530475139618, 0.032524824142456055, 0.15862980484962463, 0.2895345985889435, 0.007411073427647352, 0.03074379824101925, 0.23678991198539734, 0.04092710092663765, 0.21633881330490112, 0.10217994451522827, 0.5741018652915955, 0.08794906735420227, 0.15811748802661896], [0.1548197716474533, 0.04407857358455658, 0.04267416149377823, 0.14390510320663452, 0.39150071144104004, 0.10470721870660782, 0.21010224521160126, 0.37398451566696167, 0.24677534401416779, 0.3071460425853729, 0.12511251866817474, 0.37053829431533813, 0.34731435775756836, 0.21468856930732727, 0.22426171600818634], [0.01666487753391266, 0.070415198802948, 0.13558338582515717, 0.030082950368523598, 0.17114414274692535, 0.20995233952999115, 0.018852930516004562, 0.2688913345336914, 0.024380644783377647, 0.01614876091480255, 0.058318838477134705, 0.003357462352141738, 0.22233186662197113, 0.08606056123971939, 0.08522026240825653], [0.26702794432640076, 0.10013092309236526, 0.15535299479961395, 0.01822819747030735, 0.19259323179721832, 0.1620739996433258, 0.06925511360168457, 0.14121465384960175, 0.30160874128341675, 0.138941690325737, 0.14571446180343628, 0.1845642775297165, 0.3172887861728668, 0.1378965824842453, 0.15321676433086395], [0.05774107202887535, 0.08979255706071854, 0.15777261555194855, 0.0986839085817337, 0.04042482376098633, 0.02364284358918667, 0.006265458185225725, 0.20312650501728058, 0.04589210823178291, 0.2705432176589966, 0.29482388496398926, 0.25277185440063477, 0.21941334009170532, 0.09023746848106384, 0.12374064326286316], [0.10808208584785461, 0.08377770334482193, 0.3031982481479645, 0.08575166761875153, 0.1659224033355713, 0.02410510927438736, 0.024052061140537262, 0.06346622854471207, 0.012278172187507153, 0.033475130796432495, 0.02865537814795971, 0.2309909611940384, 0.5272806286811829, 0.058207638561725616, 0.12589795887470245], [0.2848440408706665, 0.04557379335165024, 0.07043055444955826, 0.13887976109981537, 0.25104182958602905, 0.08729252219200134, 0.03900376707315445, 0.06159999966621399, 0.07028467953205109, 0.1360185593366623, 0.12163159996271133, 0.4339398145675659, 0.18035274744033813, 0.13636742532253265, 0.35040098428726196], [0.03364454582333565, 0.06385143101215363, 0.4650610089302063, 0.13847006857395172, 0.12132523953914642, 0.23606915771961212, 0.02828356996178627, 0.17786316573619843, 0.0068073878064751625, 0.0032905752304941416, 0.04716186597943306, 0.060036350041627884, 0.5867005586624146, 0.23594366014003754, 0.05739189311861992], [0.04961356148123741, 0.4571499228477478, 0.32633671164512634, 0.044803813099861145, 0.12193554639816284, 0.15620054304599762, 0.031114954501390457, 0.37925899028778076, 0.023853085935115814, 0.007363635115325451, 0.0625552162528038, 0.04359081760048866, 0.12771400809288025, 0.10945692658424377, 0.03218715265393257], [0.054336514323949814, 0.12682472169399261, 0.28572455048561096, 0.7098703384399414, 0.04356186464428902, 0.036012813448905945, 0.12616953253746033, 0.12438997626304626, 0.06097114831209183, 0.011340769939124584, 0.00453603221103549, 0.02511424943804741, 0.15918391942977905, 0.004009802360087633, 0.1337292641401291], [0.029656492173671722, 0.11861541867256165, 0.25968441367149353, 0.6952800154685974, 0.06073199212551117, 0.3734285235404968, 0.030824951827526093, 0.09641394764184952, 0.0529148206114769, 0.01715172454714775, 0.01323915645480156, 0.055627286434173584, 0.11593649536371231, 0.04441850632429123, 0.04630020260810852], [0.10554661601781845, 0.6362442970275879, 0.6959939002990723, 0.018170323222875595, 0.40134888887405396, 0.15823723375797272, 0.1629355400800705, 0.11358990520238876, 0.24731940031051636, 0.23558683693408966, 0.07505767047405243, 0.03725680336356163, 0.014009351842105389, 0.03713200241327286, 0.09585387259721756], [0.4055319130420685, 0.2534714341163635, 0.44874629378318787, 0.14194901287555695, 0.3008168041706085, 0.20029903948307037, 0.07248799502849579, 0.26174047589302063, 0.1826024055480957, 0.0982341319322586, 0.09884719550609589, 0.22728654742240906, 0.04277953878045082, 0.06280668079853058, 0.09454112499952316], [0.025013893842697144, 0.013348683714866638, 0.22353146970272064, 0.0037027201615273952, 0.14888618886470795, 0.22346094250679016, 0.021921563893556595, 0.6342950463294983, 0.03356323391199112, 0.06236502528190613, 0.03522828221321106, 0.17797930538654327, 0.04731723666191101, 0.06786928325891495, 0.042550042271614075]], [[0.1577349603176117, 0.09554319828748703, 0.02016325853765011, 0.08440300822257996, 0.33925309777259827, 0.35353752970695496, 0.49755600094795227, 0.2782062292098999, 0.2544572949409485, 0.6230229735374451, 0.04059281200170517, 0.12019311636686325, 0.2659685015678406, 0.3508304953575134, 0.10784413665533066], [0.053030457347631454, 0.00926118716597557, 0.08361255377531052, 0.1587543487548828, 0.42493122816085815, 0.0713140144944191, 0.05032603442668915, 0.790120005607605, 0.4618776738643646, 0.3647898733615875, 0.20375682413578033, 0.2847990393638611, 0.20242592692375183, 0.33538198471069336, 0.174686461687088], [0.08703262358903885, 0.32554149627685547, 0.013934381306171417, 0.05831753462553024, 0.13550086319446564, 0.24707834422588348, 0.10738440603017807, 0.2015978991985321, 0.20393061637878418, 0.3176687955856323, 0.11071985214948654, 0.18533341586589813, 0.23293758928775787, 0.34885379672050476, 0.5850104689598083], [0.10977373272180557, 0.1966770738363266, 0.08552326261997223, 0.3559982180595398, 0.025181425735354424, 0.05637436732649803, 0.04466243088245392, 0.30799123644828796, 0.24855823814868927, 0.13041310012340546, 0.16531962156295776, 0.11238406598567963, 0.33737656474113464, 0.08863592892885208, 0.043888676911592484], [0.5166918635368347, 0.35558366775512695, 0.01755080744624138, 0.011931763030588627, 0.556053638458252, 0.21828243136405945, 0.17387567460536957, 0.11686032265424728, 0.22141756117343903, 0.6036979556083679, 0.3235246241092682, 0.21816273033618927, 0.20258961617946625, 0.7225815653800964, 0.3817636966705322], [0.34899845719337463, 0.35567307472229004, 0.2643766403198242, 0.12664493918418884, 0.18397535383701324, 0.012551958672702312, 0.056629326194524765, 0.06369142234325409, 0.252005010843277, 0.3601645529270172, 0.3771168887615204, 0.4479873776435852, 0.13717319071292877, 0.6667386293411255, 0.1451762467622757], [0.5782451629638672, 0.6189379096031189, 0.11758852005004883, 0.3125992715358734, 0.3504111170768738, 0.10631152987480164, 0.16217094659805298, 0.04177623987197876, 0.10916820168495178, 0.3274877965450287, 0.10721725970506668, 0.11595069617033005, 0.11270644515752792, 0.32787472009658813, 0.13412055373191833], [0.2553749084472656, 0.5479037165641785, 0.3395489752292633, 0.13140854239463806, 0.07771788537502289, 0.06743729114532471, 0.04718935862183571, 0.022107038646936417, 0.2706955075263977, 0.06462319940328598, 0.20574931800365448, 0.08401398360729218, 0.11249610781669617, 0.20925462245941162, 0.07354141771793365], [0.15992610156536102, 0.4297313988208771, 0.11996463686227798, 0.29957810044288635, 0.19940054416656494, 0.6192947030067444, 0.07005859166383743, 0.4058174192905426, 0.0451255701482296, 0.02480492927134037, 0.052432600408792496, 0.13078351318836212, 0.14195236563682556, 0.12686756253242493, 0.10959619283676147], [0.13202522695064545, 0.3311104476451874, 0.12707853317260742, 0.06901858001947403, 0.13186469674110413, 0.37057942152023315, 0.1482420712709427, 0.21941475570201874, 0.1949346363544464, 0.11534072458744049, 0.011536079458892345, 0.018882060423493385, 0.16279305517673492, 0.07962523400783539, 0.11737312376499176], [0.0604790523648262, 0.5140921473503113, 0.37517040967941284, 0.060462601482868195, 0.14644990861415863, 0.49839717149734497, 0.08009912073612213, 0.3367377519607544, 0.0785842090845108, 0.043956201523542404, 0.0826396569609642, 0.015624956227838993, 0.10417986661195755, 0.07971351593732834, 0.018050679937005043], [0.10509271919727325, 0.5468136072158813, 0.2136838436126709, 0.13898353278636932, 0.11654751002788544, 0.1982421725988388, 0.03731672093272209, 0.5618436336517334, 0.37511539459228516, 0.015668287873268127, 0.07859797775745392, 0.026544239372015, 0.11879771202802658, 0.051024846732616425, 0.03191406652331352], [0.2583395540714264, 0.306291788816452, 0.15283380448818207, 0.48663485050201416, 0.24239543080329895, 0.6472541093826294, 0.11895711719989777, 0.7050262093544006, 0.43789902329444885, 0.07257331907749176, 0.1529301553964615, 0.07237879186868668, 0.029207568615674973, 0.031136667355895042, 0.04320577159523964], [0.37997886538505554, 0.3090342879295349, 0.09529577195644379, 0.06091787666082382, 0.5611693859100342, 0.5351426005363464, 0.5250707268714905, 0.4058402180671692, 0.08284364640712738, 0.7192233204841614, 0.12988585233688354, 0.24924960732460022, 0.016598563641309738, 0.6531801819801331, 0.22117754817008972], [0.31734058260917664, 0.02799793891608715, 0.08435621112585068, 0.4273812472820282, 0.37900310754776, 0.1551857888698578, 0.12445898354053497, 0.02975497953593731, 0.13922178745269775, 0.25836795568466187, 0.3142063617706299, 0.5329877138137817, 0.020000692456960678, 0.19246473908424377, 0.34441179037094116]], [[0.022252710536122322, 0.017558962106704712, 0.12289869785308838, 0.01514213066548109, 0.04983796179294586, 0.160098597407341, 0.09159664064645767, 0.03634485974907875, 0.27353572845458984, 0.14908282458782196, 0.8423851132392883, 0.33708906173706055, 0.03012021631002426, 0.05972116440534592, 0.2686574459075928], [0.13637107610702515, 0.02899317629635334, 0.09026061743497849, 0.22582301497459412, 0.09117049723863602, 0.19661013782024384, 0.30083417892456055, 0.13528303802013397, 0.1352328211069107, 0.18504901230335236, 0.3621358573436737, 0.504258930683136, 0.10044156759977341, 0.37106865644454956, 0.36433035135269165], [0.10935092717409134, 0.06271693855524063, 0.044740546494722366, 0.1709805577993393, 0.22382155060768127, 0.2615796625614166, 0.3429900109767914, 0.02677186205983162, 0.39723172783851624, 0.1559167355298996, 0.6381150484085083, 0.34350308775901794, 0.14388519525527954, 0.322640985250473, 0.07209958881139755], [0.11123806983232498, 0.14550834894180298, 0.12841136753559113, 0.013620064593851566, 0.006130752619355917, 0.025231752544641495, 0.11538708955049515, 0.09429272264242172, 0.3855685293674469, 0.016912028193473816, 0.3869503438472748, 0.1961694061756134, 0.15352581441402435, 0.019190048798918724, 0.4291467070579529], [0.1283823847770691, 0.33987957239151, 0.06837885081768036, 0.03946131095290184, 0.03139644116163254, 0.11983324587345123, 0.12062173336744308, 0.46404916048049927, 0.24212448298931122, 0.1594262570142746, 0.4298713207244873, 0.5236353278160095, 0.2188095897436142, 0.049411591142416, 0.10146455466747284], [0.010564678348600864, 0.32722386717796326, 0.19864077866077423, 0.015389330685138702, 0.0028029000386595726, 0.007416849955916405, 0.003262599464505911, 0.23795713484287262, 0.05000551417469978, 0.075996033847332, 0.049679387360811234, 0.21265098452568054, 0.2097157984972, 0.01007634773850441, 0.03895873948931694], [0.10390599817037582, 0.04329453781247139, 0.42168325185775757, 0.06385642290115356, 0.04340887442231178, 0.029213739559054375, 0.036663200706243515, 0.0028809772338718176, 0.19718152284622192, 0.16335125267505646, 0.6605148315429688, 0.17834524810314178, 0.08135847747325897, 0.05741032958030701, 0.24636343121528625], [0.010566278360784054, 0.32608217000961304, 0.34194469451904297, 0.08201102167367935, 0.036688148975372314, 0.12155891954898834, 0.015490439720451832, 0.05858473479747772, 0.1731383204460144, 0.12207219004631042, 0.0636284351348877, 0.2239474654197693, 0.2988812327384949, 0.033257871866226196, 0.04593053460121155], [0.26241976022720337, 0.0378817655146122, 0.10770448297262192, 0.11944369971752167, 0.367754727602005, 0.041288651525974274, 0.25914207100868225, 0.061461515724658966, 0.061867646872997284, 0.08977923542261124, 0.03797370195388794, 0.2101898193359375, 0.035329420119524, 0.38835543394088745, 0.3324989080429077], [0.3753410875797272, 0.031615160405635834, 0.1074504628777504, 0.07966858148574829, 0.16393397748470306, 0.01204571221023798, 0.36072632670402527, 0.026240641251206398, 0.09493876993656158, 0.12203314155340195, 0.0640302300453186, 0.13458214700222015, 0.19451306760311127, 0.3176366686820984, 0.19878560304641724], [0.19523903727531433, 0.1090913861989975, 0.11059779673814774, 0.03402426466345787, 0.4491459131240845, 0.1729225516319275, 0.3482173979282379, 0.01764478161931038, 0.14307594299316406, 0.22771455347537994, 0.04787566140294075, 0.14714154601097107, 0.028272001072764397, 0.23823784291744232, 0.19700175523757935], [0.1428564339876175, 0.03585843741893768, 0.023294193670153618, 0.1143055409193039, 0.07461919635534286, 0.13578416407108307, 0.4153969883918762, 0.03374828025698662, 0.10746961832046509, 0.17216910421848297, 0.02314077876508236, 0.02450137585401535, 0.06497504562139511, 0.381274551153183, 0.14229674637317657], [0.5444629788398743, 0.049506742507219315, 0.09827632457017899, 0.29229700565338135, 0.06650383025407791, 0.11397240310907364, 0.597455620765686, 0.1362738311290741, 0.15222173929214478, 0.2562837302684784, 0.13646292686462402, 0.38294121623039246, 0.030382927507162094, 0.038297515362501144, 0.465526819229126], [0.12950241565704346, 0.2834409177303314, 0.40745216608047485, 0.040315985679626465, 0.09126543253660202, 0.16738829016685486, 0.24838824570178986, 0.2707839906215668, 0.5177856087684631, 0.1416875720024109, 0.6573355793952942, 0.4225574731826782, 0.02239617332816124, 0.07502269744873047, 0.07588320225477219], [0.00751910824328661, 0.5024122595787048, 0.38239815831184387, 0.016937274485826492, 0.039716992527246475, 0.11479316651821136, 0.004478333052247763, 0.02017248421907425, 0.011771232821047306, 0.0035600941628217697, 0.03807784244418144, 0.07125832885503769, 0.1964063048362732, 0.0026467873249202967, 0.00302477041259408]], [[0.06952784210443497, 0.0770183801651001, 0.23747292160987854, 0.022874178364872932, 0.14143598079681396, 0.08435114473104477, 0.0795491486787796, 0.054600730538368225, 0.015159118920564651, 0.06120437756180763, 0.02771361917257309, 0.06765643507242203, 0.013518131338059902, 0.15485556423664093, 0.21279898285865784], [0.2531612813472748, 0.03241151198744774, 0.04793045297265053, 0.13835468888282776, 0.05921119078993797, 0.20751594007015228, 0.5453532934188843, 0.021712571382522583, 0.07093679159879684, 0.2689567506313324, 0.13515745103359222, 0.05570060759782791, 0.04099860414862633, 0.03517309948801994, 0.11268090456724167], [0.35043928027153015, 0.18572849035263062, 0.0481790192425251, 0.19426384568214417, 0.018465382978320122, 0.2676069438457489, 0.3000488579273224, 0.2726097106933594, 0.08134563267230988, 0.10164237022399902, 0.05787196010351181, 0.03694695979356766, 0.21335498988628387, 0.0815601795911789, 0.051584985107183456], [0.10967924445867538, 0.047143928706645966, 0.06498727947473526, 0.0161599051207304, 0.08311080187559128, 0.25361040234565735, 0.2589581310749054, 0.0646943673491478, 0.11701063811779022, 0.7398742437362671, 0.11236728727817535, 0.4240334630012512, 0.09019055217504501, 0.1980810910463333, 0.08526580780744553], [0.0050394656136631966, 0.005000656470656395, 0.01952306181192398, 0.4184519350528717, 0.012662295252084732, 0.015614073723554611, 0.006089636590331793, 0.027387546375393867, 0.007885311730206013, 0.009227052330970764, 0.015002718195319176, 0.002679894445464015, 0.040426015853881836, 0.023895790800452232, 0.031263262033462524], [0.1104135811328888, 0.16341662406921387, 0.10040471702814102, 0.15014782547950745, 0.22085179388523102, 0.07417210936546326, 0.08140900731086731, 0.21936744451522827, 0.12380684167146683, 0.030364450067281723, 0.008148477412760258, 0.040405042469501495, 0.016740301623940468, 0.05651557818055153, 0.03777482733130455], [0.021739037707448006, 0.025255737826228142, 0.041796568781137466, 0.028582973405718803, 0.06361079961061478, 0.10603900998830795, 0.04079660773277283, 0.23573672771453857, 0.031395647674798965, 0.17699679732322693, 0.11518478393554688, 0.12758946418762207, 0.029195530340075493, 0.19761133193969727, 0.24158287048339844], [0.1121676117181778, 0.056780170649290085, 0.05766424164175987, 0.4753672778606415, 0.17093990743160248, 0.055545274168252945, 0.23774300515651703, 0.047642335295677185, 0.2396271675825119, 0.07084424793720245, 0.05071293190121651, 0.15200014412403107, 0.17973174154758453, 0.16349640488624573, 0.16329222917556763], [0.08155515789985657, 0.04415197670459747, 0.09395420551300049, 0.06736686080694199, 0.009449290111660957, 0.007789341267198324, 0.08313233405351639, 0.018231436610221863, 0.2736586928367615, 0.12516330182552338, 0.14283257722854614, 0.03993181511759758, 0.11735112965106964, 0.037545330822467804, 0.095799021422863], [0.07989984005689621, 0.019307896494865417, 0.05061032995581627, 0.29983657598495483, 0.009587445296347141, 0.23453857004642487, 0.06259765475988388, 0.014452173374593258, 0.026213111355900764, 0.03952796012163162, 0.12968890368938446, 0.019515926018357277, 0.23016268014907837, 0.18980233371257782, 0.14884653687477112], [0.042069002985954285, 0.007410319056361914, 0.027750220149755478, 0.14348776638507843, 0.190275177359581, 0.0696464255452156, 0.09576459228992462, 0.08924749493598938, 0.16830699145793915, 0.14098002016544342, 0.2945949137210846, 0.08460760116577148, 0.11812892556190491, 0.2108343094587326, 0.28860458731651306], [0.509858250617981, 0.07021021842956543, 0.044154465198516846, 0.005825423635542393, 0.5241404175758362, 0.030089300125837326, 0.19222509860992432, 0.02549084462225437, 0.1939508020877838, 0.09437919408082962, 0.10883274674415588, 0.13631868362426758, 0.08004569262266159, 0.04784407094120979, 0.14005501568317413], [0.029798628762364388, 0.0011461747344583273, 0.00650657806545496, 0.02902117185294628, 0.007348767947405577, 0.012432223185896873, 0.018553903326392174, 0.006125486921519041, 0.008405826054513454, 0.057926055043935776, 0.04542696848511696, 0.21123111248016357, 0.05352021008729935, 0.2931033968925476, 0.1833699345588684], [0.01627730205655098, 0.0057758791372179985, 0.013731835409998894, 0.6289489269256592, 0.011782719753682613, 0.006108477246016264, 0.005309773609042168, 0.023312430828809738, 0.012817217037081718, 0.00939176045358181, 0.04320970177650452, 0.012798959389328957, 0.1585281491279602, 0.11795029044151306, 0.13285225629806519], [0.39748579263687134, 0.10528232902288437, 0.006042438093572855, 0.07306646555662155, 0.020484283566474915, 0.09288878738880157, 0.6331413388252258, 0.03478514030575752, 0.016230005770921707, 0.039869412779808044, 0.10224607586860657, 0.005181388463824987, 0.007975003682076931, 0.01008305512368679, 0.026732152327895164]], [[0.2484879046678543, 0.12593188881874084, 0.11472177505493164, 0.6318025588989258, 0.009745504707098007, 0.030495919287204742, 0.054615989327430725, 0.004801109898835421, 0.23875823616981506, 0.011562658473849297, 0.02087206020951271, 0.059635717421770096, 0.011483770795166492, 0.07716090232133865, 0.041850361973047256], [0.3294946551322937, 0.17723912000656128, 0.041080135852098465, 0.30134642124176025, 0.0073102316819131374, 0.049291279166936874, 0.0495959147810936, 0.0037847748026251793, 0.014987694099545479, 0.07676513493061066, 0.039059415459632874, 0.006041571032255888, 0.011380840092897415, 0.011979957111179829, 0.02782473713159561], [0.008675806224346161, 0.016726570203900337, 0.19906938076019287, 0.3167073726654053, 0.022006884217262268, 0.014510865323245525, 0.00237266905605793, 0.00938868336379528, 0.004848333541303873, 0.00305117666721344, 0.042285457253456116, 0.0026737553998827934, 0.017337674275040627, 0.0016427191440016031, 0.0027906473260372877], [0.06292864680290222, 0.010060630738735199, 0.07846219092607498, 0.3009726405143738, 0.09911586344242096, 0.3769649565219879, 0.290684312582016, 0.048859626054763794, 0.015964722260832787, 0.02972962148487568, 0.25837212800979614, 0.050403933972120285, 0.052831199020147324, 0.44793814420700073, 0.12096201628446579], [0.0647541731595993, 0.06744952499866486, 0.010754776187241077, 0.15598785877227783, 0.08916914463043213, 0.4045051634311676, 0.5958212018013, 0.10594789683818817, 0.12025819718837738, 0.04822946712374687, 0.02913811057806015, 0.014846491627395153, 0.17111137509346008, 0.049513354897499084, 0.14188753068447113], [0.07069405168294907, 0.0006015333347022533, 0.0017680496675893664, 0.0010985832195729017, 0.0012869784841313958, 0.22278346121311188, 0.4465882480144501, 0.06128238886594772, 0.02642727456986904, 0.03756114840507507, 0.002607540925964713, 0.0018699204083532095, 0.0059012919664382935, 0.020283877849578857, 0.03355809301137924], [0.0861939862370491, 0.03346291184425354, 0.009915103204548359, 0.35010838508605957, 0.03437130153179169, 0.18394741415977478, 0.5006390810012817, 0.0633198693394661, 0.36160194873809814, 0.07578127831220627, 0.038500167429447174, 0.08213403075933456, 0.026455186307430267, 0.12013117223978043, 0.1146865040063858], [0.2484544962644577, 0.00790119543671608, 0.004407763481140137, 0.02700735628604889, 0.015422074124217033, 0.015295883640646935, 0.40846768021583557, 0.10706920176744461, 0.06367217004299164, 0.22094424068927765, 0.21221157908439636, 0.006999517325311899, 0.054566796869039536, 0.124799944460392, 0.09114839136600494], [0.1237153485417366, 0.029043834656476974, 0.07521974295377731, 0.04068650305271149, 0.002623512176796794, 0.008706655353307724, 0.03832445293664932, 0.14616532623767853, 0.1701044738292694, 0.20599642395973206, 0.11677426844835281, 0.2341107875108719, 0.06235762685537338, 0.003964806441217661, 0.15731573104858398], [0.034962959587574005, 0.023077068850398064, 0.034600574523210526, 0.14041800796985626, 0.0021679585333913565, 0.009290770627558231, 0.07274696230888367, 0.014187950640916824, 0.1371506154537201, 0.39440277218818665, 0.2198760211467743, 0.19940708577632904, 0.11203428357839584, 0.08552268147468567, 0.11737436801195145], [0.015330069698393345, 0.007386082783341408, 0.017500948160886765, 0.01906486414372921, 0.010120063088834286, 0.05364372953772545, 0.043298348784446716, 0.12658876180648804, 0.06039673835039139, 0.02238147333264351, 0.16429400444030762, 0.06984445452690125, 0.3043651580810547, 0.055543575435876846, 0.11423089355230331], [0.09644094854593277, 0.0058854687958955765, 0.03721459209918976, 0.0025620406959205866, 0.062300242483615875, 0.003563062520697713, 0.07219880819320679, 0.03924282267689705, 0.025451356545090675, 0.06598387658596039, 0.026776403188705444, 0.07250863313674927, 0.45021528005599976, 0.08199745416641235, 0.4220075309276581], [0.01460834126919508, 0.0005662022740580142, 0.0013911814894527197, 0.05315173417329788, 0.008028149604797363, 0.016604119911789894, 0.011740745045244694, 0.008678588084876537, 0.0025609249714761972, 0.01638207584619522, 0.018210044130682945, 0.014119945466518402, 0.06550943106412888, 0.34254926443099976, 0.04794229939579964], [0.05372002348303795, 0.14061135053634644, 0.018787089735269547, 0.0958278551697731, 0.0019092779839411378, 0.03348369151353836, 0.13957257568836212, 0.031220966950058937, 0.19735871255397797, 0.017847368493676186, 0.0589337982237339, 0.01900595612823963, 0.1276925951242447, 0.04769464209675789, 0.4384888708591461], [0.08416850119829178, 0.1088641807436943, 0.0573052242398262, 0.27551695704460144, 0.030813831835985184, 0.18022866547107697, 0.10468263924121857, 0.09972096234560013, 0.31189021468162537, 0.3315774202346802, 0.2321816384792328, 0.034622836858034134, 0.14143656194210052, 0.04640315845608711, 0.09621720016002655]], [[0.130781888961792, 0.31469303369522095, 0.10550640523433685, 0.05234318599104881, 0.073336161673069, 0.022349786013364792, 0.04807984083890915, 0.1931842416524887, 0.06399697810411453, 0.042083337903022766, 0.026750531047582626, 0.11997608095407486, 0.008983415551483631, 0.03431839123368263, 0.019280044361948967], [0.1582711637020111, 0.14862558245658875, 0.20016248524188995, 0.08876624703407288, 0.11006557196378708, 0.14632253348827362, 0.04025046527385712, 0.010204354301095009, 0.017868297174572945, 0.059372395277023315, 0.02111685276031494, 0.04181571304798126, 0.025184988975524902, 0.09681157767772675, 0.11611668020486832], [0.23875439167022705, 0.3084685802459717, 0.14188633859157562, 0.026331612840294838, 0.0149313323199749, 0.09176106750965118, 0.03131069242954254, 0.10051372647285461, 0.03149634972214699, 0.11085867136716843, 0.014410188421607018, 0.02796255424618721, 0.034816499799489975, 0.025807565078139305, 0.01846306212246418], [0.3404518961906433, 0.24260303378105164, 0.15383434295654297, 0.17020593583583832, 0.011800014413893223, 0.014385397545993328, 0.09441643208265305, 0.12204645574092865, 0.13843503594398499, 0.045293405652046204, 0.010667533613741398, 0.19693949818611145, 0.10281307995319366, 0.01422606036067009, 0.06984427571296692], [0.002873742487281561, 0.008706165477633476, 0.35573768615722656, 0.0015586970839649439, 0.015496796928346157, 0.003392455168068409, 0.01149011217057705, 0.01891980692744255, 0.016394488513469696, 0.003960000351071358, 0.0035995631478726864, 0.008501716889441013, 0.018164046108722687, 0.004727588500827551, 0.013562880456447601], [0.044807154685258865, 0.02788197249174118, 0.03947468474507332, 0.1271299421787262, 0.17640650272369385, 0.25110092759132385, 0.08349309861660004, 0.02069718949496746, 0.45751577615737915, 0.039922621101140976, 0.1781769096851349, 0.002931024879217148, 0.16567888855934143, 0.1177627220749855, 0.5156693458557129], [0.005990047473460436, 0.04782475531101227, 0.01399919856339693, 0.010489771142601967, 0.06132129579782486, 0.030459748581051826, 0.010153756476938725, 0.3387801945209503, 0.06446883827447891, 0.007243711035698652, 0.00693717272952199, 0.020023254677653313, 0.007285784464329481, 0.009139767847955227, 0.0044054011814296246], [0.020405659452080727, 0.00729386368766427, 0.06661678105592728, 0.08295443654060364, 0.20373474061489105, 0.3448184132575989, 0.04295210912823677, 0.20947468280792236, 0.03081577830016613, 0.010805373080074787, 0.17521467804908752, 0.06567652523517609, 0.012400656938552856, 0.10652147233486176, 0.07385163754224777], [0.21573591232299805, 0.13175059854984283, 0.04085814207792282, 0.04119405150413513, 0.03551999852061272, 0.023009058088064194, 0.2751774191856384, 0.047030266374349594, 0.14272502064704895, 0.20153193175792694, 0.09575672447681427, 0.11327007412910461, 0.008532780222594738, 0.053245026618242264, 0.08952803909778595], [0.2778390347957611, 0.11423225700855255, 0.3034791946411133, 0.34643107652664185, 0.5395972728729248, 0.06785042583942413, 0.13029156625270844, 0.18737749755382538, 0.029348008334636688, 0.16667678952217102, 0.021040884777903557, 0.008728248998522758, 0.037633832544088364, 0.02033349499106407, 0.03947347402572632], [0.4898838996887207, 0.08082167059183121, 0.07362432777881622, 0.02171795442700386, 0.1333591789007187, 0.09000474214553833, 0.13501934707164764, 0.03979193791747093, 0.19113953411579132, 0.13522492349147797, 0.16557832062244415, 0.16255514323711395, 0.07687958329916, 0.15948235988616943, 0.09843874722719193], [0.045906297862529755, 0.18602333962917328, 0.4082620143890381, 0.010370302945375443, 0.04507172852754593, 0.19693265855312347, 0.04021843150258064, 0.027866821736097336, 0.1546991914510727, 0.33766424655914307, 0.09260500222444534, 0.05066358670592308, 0.05655887722969055, 0.13157807290554047, 0.06850539147853851], [0.020344020798802376, 0.0030158585868775845, 0.004445259924978018, 0.022628312930464745, 0.030150510370731354, 0.027700912207365036, 0.026311388239264488, 0.012862108647823334, 0.07009940594434738, 0.24656175076961517, 0.10596039146184921, 0.1143152266740799, 0.3679012656211853, 0.0068145813420414925, 0.04171491786837578], [0.004749340936541557, 0.00182742765173316, 0.0021293568424880505, 0.00394084258005023, 0.004750867374241352, 5.3125138947507367e-05, 0.0026011874433606863, 0.000718552153557539, 0.002356230979785323, 0.00125187449157238, 0.0021339249797165394, 0.00044074622564949095, 0.2141493707895279, 0.0029175111558288336, 0.00477015832439065], [0.12991508841514587, 0.06724811345338821, 0.06397818773984909, 0.15923364460468292, 0.2566852867603302, 0.07963784784078598, 0.09182894974946976, 0.040824584662914276, 0.21298912167549133, 0.2517295181751251, 0.2285410314798355, 0.11115844547748566, 0.1010512113571167, 0.3968040943145752, 0.1870165765285492]], [[0.06147387623786926, 0.0657946914434433, 0.22564710676670074, 0.1299343705177307, 0.021580645814538002, 0.08992400765419006, 0.025479430332779884, 0.04823821783065796, 0.05891237407922745, 0.016958819702267647, 0.0021926285699009895, 0.017513686791062355, 0.09859969466924667, 0.16368542611598969, 0.038398925215005875], [0.029852252453565598, 0.26626214385032654, 0.14803646504878998, 0.038784727454185486, 0.07803148031234741, 0.006210723891854286, 0.0026457132771611214, 0.006018034182488918, 0.05453306809067726, 0.002730109030380845, 0.015730326995253563, 0.0017557059181854129, 0.034912969917058945, 0.03208531066775322, 0.03983413055539131], [0.01053018867969513, 0.02744918502867222, 0.2530466914176941, 0.05846027657389641, 0.1744728684425354, 0.011957419104874134, 0.003304906887933612, 0.00205883732996881, 0.00874510407447815, 0.0014524421421810985, 0.0009729861048981547, 0.0026561047416180372, 0.0023208027705550194, 0.0038251704536378384, 0.005045189522206783], [0.016039762645959854, 0.05755838379263878, 0.10756286233663559, 0.03799062967300415, 0.5738711953163147, 0.061907339841127396, 0.128611221909523, 0.01847657933831215, 0.06501789391040802, 0.015564735978841782, 0.0016139671206474304, 0.014343881979584694, 0.020734043791890144, 0.14008449018001556, 0.13515408337116241], [0.005847899243235588, 0.11914067715406418, 0.01715121790766716, 0.3517457842826843, 0.0661543607711792, 0.07493122667074203, 0.012425812892615795, 0.11745280772447586, 0.08440648764371872, 0.020029406994581223, 0.05165768414735794, 0.04094480350613594, 0.024548601359128952, 0.005826729815453291, 0.13841456174850464], [0.015926362946629524, 0.007578620687127113, 0.1226087138056755, 0.030128292739391327, 0.03851892054080963, 0.3367418944835663, 0.01694057136774063, 0.09829536825418472, 0.0361555740237236, 0.10537439584732056, 0.007450005039572716, 0.029753634706139565, 0.22920416295528412, 0.01793695241212845, 0.05258304625749588], [0.01326388493180275, 0.05337870866060257, 0.047661036252975464, 0.08615607023239136, 0.12425915151834488, 0.4180251955986023, 0.04702466353774071, 0.0717325434088707, 0.05138256773352623, 0.06877672672271729, 0.0152205191552639, 0.0719875767827034, 0.1666427105665207, 0.13322126865386963, 0.053655143827199936], [0.026802292093634605, 0.003955241292715073, 0.0206829272210598, 0.02742936834692955, 0.06016179919242859, 0.15127348899841309, 0.06774158030748367, 0.2981398105621338, 0.05239749699831009, 0.09365928173065186, 0.035629644989967346, 0.020771589130163193, 0.13655303418636322, 0.012941722758114338, 0.05640798062086105], [0.06469012051820755, 0.1851334124803543, 0.08788572251796722, 0.19977343082427979, 0.00846380740404129, 0.03702360764145851, 0.0876760184764862, 0.046302031725645065, 0.11564433574676514, 0.05180440843105316, 0.49518024921417236, 0.1649368405342102, 0.030481798574328423, 0.10461966693401337, 0.07739346474409103], [0.020106524229049683, 0.01925482228398323, 0.006043681409209967, 0.01652396097779274, 0.001572003006003797, 0.005779887083917856, 0.015335858799517155, 0.03537710756063461, 0.009967570193111897, 0.09144406765699387, 0.43651703000068665, 0.2613205015659332, 0.0483890138566494, 0.06553913652896881, 0.055434126406908035], [0.07980967313051224, 0.14815203845500946, 0.09271827340126038, 0.004086778499186039, 0.010790406726300716, 0.0747552439570427, 0.10995902121067047, 0.04728228971362114, 0.1809520274400711, 0.025821411982178688, 0.06657237559556961, 0.1431768387556076, 0.19449584186077118, 0.20780201256275177, 0.10148976743221283], [0.05537823587656021, 0.008725662715733051, 0.0058344281278550625, 0.029011448845267296, 0.048424966633319855, 0.047911662608385086, 0.16901308298110962, 0.17019973695278168, 0.011648884043097496, 0.08953043073415756, 0.5360274910926819, 0.10330803692340851, 0.078437939286232, 0.12202966213226318, 0.11905822902917862], [0.01546903420239687, 0.0005347061669453979, 0.0015839362749829888, 0.053056132048368454, 0.23614321649074554, 0.013318118639290333, 0.051473915576934814, 0.011966699734330177, 0.007302975282073021, 0.09275621920824051, 0.06646261364221573, 0.010813506320118904, 0.13289499282836914, 0.22826357185840607, 0.04386172071099281], [0.009458722546696663, 0.0058342707343399525, 0.012789146974682808, 0.005895438138395548, 0.026010286062955856, 0.057482823729515076, 0.005663284566253424, 0.005727604031562805, 0.0033144087065011263, 0.011671853251755238, 0.00424896739423275, 0.056589994579553604, 0.20401620864868164, 0.03777612745761871, 0.03114682249724865], [0.0012354525970295072, 0.034024473279714584, 0.10020612925291061, 0.02267461270093918, 0.08676987141370773, 0.14216794073581696, 0.0033775768242776394, 0.07320579141378403, 0.07390473037958145, 0.0168889332562685, 0.00386308366432786, 0.02569040097296238, 0.24664165079593658, 0.2674221694469452, 0.014589445665478706]], [[0.2643359303474426, 0.2943609654903412, 0.10517127066850662, 0.013473477214574814, 0.17808614671230316, 0.05031028389930725, 0.0477585569024086, 0.13444076478481293, 0.0626431554555893, 0.05089121311903, 0.025438696146011353, 0.12666909396648407, 0.015911895781755447, 0.08822031319141388, 0.09637932479381561], [0.02893858775496483, 0.3286381959915161, 0.024464154615998268, 0.015645690262317657, 0.07065004110336304, 0.03320073336362839, 0.0035833900328725576, 0.002133443485945463, 0.0077736834064126015, 0.0014096481027081609, 0.006704544182866812, 0.0034484381321817636, 0.010553284548223019, 0.029550330713391304, 0.0064092278480529785], [0.0403970405459404, 0.029290249571204185, 0.2564694881439209, 0.03103366494178772, 0.01930038072168827, 0.0007984130643308163, 0.0024861868005245924, 0.013074777089059353, 0.025626862421631813, 0.0022637112997472286, 0.010511897504329681, 0.03038576804101467, 0.00803295336663723, 0.000980974524281919, 0.040744345635175705], [0.23322375118732452, 0.23003342747688293, 0.24563531577587128, 0.07496963441371918, 0.029645830392837524, 0.0015733843902125955, 0.048427432775497437, 0.07474764436483383, 0.005064227152615786, 0.006064139772206545, 0.00639030896127224, 0.0023683567997068167, 0.0201968252658844, 0.0057837339118123055, 0.030518243089318275], [0.009382463060319424, 0.004108777269721031, 0.355550616979599, 0.0026344929356127977, 0.036474164575338364, 0.0013674235669896007, 0.010420771315693855, 0.008167937397956848, 0.005904712714254856, 0.0164882093667984, 0.0014915319625288248, 0.00666471105068922, 0.007061991840600967, 0.006146776955574751, 0.03842667490243912], [0.340854674577713, 0.027831802144646645, 0.11495380103588104, 0.4507772624492645, 0.33573275804519653, 0.07158998399972916, 0.3054116368293762, 0.09558256715536118, 0.008191889151930809, 0.08007357269525528, 0.08199689537286758, 0.011630101129412651, 0.016172919422388077, 0.020448284223675728, 0.05253906920552254], [0.0825798362493515, 0.09406770020723343, 0.044158000499010086, 0.06245531886816025, 0.15669509768486023, 0.1018981784582138, 0.17849969863891602, 0.1823071539402008, 0.1725231111049652, 0.14688736200332642, 0.027769910171628, 0.1729786992073059, 0.04907526820898056, 0.09640378504991531, 0.07928813993930817], [0.04138464853167534, 0.0045098732225596905, 0.098704032599926, 0.034942083060741425, 0.1842936873435974, 0.1567782759666443, 0.14141200482845306, 0.1953822374343872, 0.09936889261007309, 0.281032919883728, 0.13522183895111084, 0.012650868855416775, 0.02501768246293068, 0.2133605033159256, 0.14542686939239502], [0.05831298604607582, 0.07845572382211685, 0.00935202743858099, 0.09348727762699127, 0.2554629147052765, 0.026818757876753807, 0.15820558369159698, 0.09712891280651093, 0.18406683206558228, 0.297629177570343, 0.011888068169355392, 0.04674078896641731, 0.01729435659945011, 0.04945852607488632, 0.08047669380903244], [0.030211733654141426, 0.004252443555742502, 0.044400423765182495, 0.0032993308268487453, 0.029341043904423714, 0.14371474087238312, 0.17894455790519714, 0.12369092553853989, 0.48359414935112, 0.06321088969707489, 0.05475561320781708, 0.3139732778072357, 0.086760014295578, 0.13208359479904175, 0.2905256450176239], [0.06285266578197479, 0.0062216646037995815, 0.016913438215851784, 0.007285475265234709, 0.01629750058054924, 0.004617355298250914, 0.06147269159555435, 0.21831700205802917, 0.11657348275184631, 0.39258062839508057, 0.17390909790992737, 0.3519352376461029, 0.014494672417640686, 0.04437657818198204, 0.04845427721738815], [0.014810703694820404, 0.027867808938026428, 0.00787208043038845, 0.003661711234599352, 0.06816401332616806, 0.014048570767045021, 0.04280591011047363, 0.04519394412636757, 0.07874996215105057, 0.2074531614780426, 0.12078044563531876, 0.53052818775177, 0.035032909363508224, 0.1398327797651291, 0.02986292913556099], [0.011430865153670311, 0.002694258699193597, 0.03896895423531532, 0.04504057392477989, 0.00808126013725996, 0.01048098411411047, 0.012571780942380428, 0.0054772221483290195, 0.07419075071811676, 0.02193005569279194, 0.3994891941547394, 0.15694338083267212, 0.3065741956233978, 0.022703034803271294, 0.07852455973625183], [0.0007813395350240171, 4.470362910069525e-06, 0.0010683261789381504, 0.022204171866178513, 0.0022952572908252478, 4.198186070425436e-05, 0.0009061718010343611, 0.0006557627930305898, 0.0009219115017913282, 0.0006920882733538747, 0.005404994357377291, 0.012070748023688793, 0.21383939683437347, 0.0026518681552261114, 0.0011399114737287164], [0.03732156753540039, 0.14082211256027222, 0.08218222856521606, 0.02148711122572422, 0.037640467286109924, 0.011636778712272644, 0.01611051708459854, 0.06724098324775696, 0.20042963325977325, 0.035641491413116455, 0.045655738562345505, 0.041121501475572586, 0.23917138576507568, 0.01630677469074726, 0.2854580283164978]]], [[[0.00028402332100085914, 1.9304454923485537e-08, 1.5483598847509938e-09, 7.885660006923256e-12, 2.7246130684943637e-08, 2.9440096113830805e-05, 4.3406546978985716e-07, 3.7434634236888087e-07, 3.9264233464564313e-07, 1.911867819615054e-08, 6.894639170695882e-08, 1.9322192201798316e-06, 1.594805780769093e-06, 1.097217136702966e-06, 0.25163131952285767], [0.8221166729927063, 0.0031213052570819855, 7.842657214496285e-05, 5.977510153520882e-10, 6.043178735204435e-10, 7.336016096815001e-07, 0.0001510237343609333, 0.000765863514970988, 0.0003504687047097832, 5.704807790607447e-07, 3.8402351520971933e-08, 3.7901799032624695e-07, 1.534954208182171e-05, 4.934078606311232e-05, 0.00023439944197889417], [0.0023944040294736624, 0.796754002571106, 0.004422985017299652, 9.068900226338883e-07, 5.795331436964091e-10, 1.0343059742012883e-08, 4.4964113499190717e-07, 0.0014743957435712218, 0.00028717826353386045, 7.994436600711197e-05, 3.3569827451174206e-07, 1.215876466176269e-07, 7.940250839055807e-07, 4.835407253267476e-06, 2.585098854979151e-07], [4.3931080995207594e-11, 0.0005229745293036103, 0.5791732668876648, 0.0002632129180710763, 3.316774765949049e-08, 1.7754019825469425e-12, 1.4596207272357664e-14, 1.5350217763554497e-09, 1.2882580335826788e-07, 7.457471838279162e-06, 1.2410231420290074e-06, 2.736720361440348e-08, 3.621486097116211e-11, 3.919724787804224e-12, 2.306477925317907e-12], [3.994035801418473e-14, 1.3595737036187217e-10, 5.270875135465758e-06, 0.5513067841529846, 0.00020578903786372393, 1.9226330039145978e-07, 1.181193272532799e-12, 2.80986930771554e-13, 9.120337812881449e-14, 1.37843805814164e-10, 7.154308718781976e-07, 1.5133276747292257e-06, 7.425698944629744e-10, 2.2010659354171347e-13, 1.8997327582565005e-12], [2.3444651168352815e-12, 2.1774425253313912e-13, 1.857566878094019e-09, 0.00030468025943264365, 0.9472002983093262, 0.00010681805724743754, 2.00606624645161e-08, 5.2167251502746245e-14, 1.354494091723496e-15, 5.737065011425513e-13, 8.729777456473187e-10, 3.2425006793346256e-05, 7.676636641917867e-07, 1.870739785303499e-09, 2.3914221713994266e-09], [3.644098217625569e-11, 3.867062572937563e-11, 4.1057553190615437e-11, 1.5412249254609378e-09, 0.018834512680768967, 0.505605936050415, 0.0010763276368379593, 5.434728933551014e-08, 2.6194791127864825e-11, 6.074670846504876e-15, 3.814499497517554e-12, 1.2291486939375318e-07, 9.572526323609054e-06, 4.437842653715052e-05, 7.18067713023629e-06], [5.002242687623948e-05, 2.445471238843311e-07, 7.217475506138271e-09, 2.943958878759423e-12, 1.391844648424012e-07, 0.0035048718564212322, 0.755942702293396, 0.0011242764303460717, 1.4866960555082187e-05, 9.753278740198823e-11, 3.792431321238132e-13, 1.6398679289486573e-11, 1.3850768709744443e-07, 0.0002873632765840739, 2.565975592005998e-05], [7.748224284398475e-09, 3.667011867491965e-07, 1.7906526261768363e-09, 1.001209222569038e-16, 4.707358499311462e-15, 2.921879960204876e-10, 4.77575849799905e-06, 0.9355171918869019, 1.7088919776142575e-05, 1.5246609308405823e-08, 1.546373502880373e-14, 1.9256968477537417e-16, 2.8356877952137637e-15, 6.199032398512827e-10, 3.679770266273863e-09], [6.04271771509346e-11, 2.349539499846287e-06, 6.254656170767703e-08, 2.0915530592191534e-12, 3.303753013789688e-16, 1.0466700578893717e-14, 7.288482968201282e-13, 0.0006303040427155793, 0.47335511445999146, 8.928982424549758e-05, 1.5872458902776998e-08, 1.3611594998645584e-14, 1.3777586457132233e-16, 1.589055302510104e-15, 8.100658338561217e-11], [3.812023474658588e-10, 1.421315573679749e-06, 2.2867025109007955e-06, 2.6682736020688935e-08, 3.632111755455525e-12, 1.6831340872913367e-14, 3.240909670081289e-14, 1.4920277635610546e-07, 0.0005182845052331686, 0.39297640323638916, 0.0007259719423018396, 1.2580667174688642e-08, 3.7229049595736974e-13, 2.157145159519631e-15, 1.0612778433838344e-09], [6.84109713322556e-10, 1.9775532322796607e-08, 5.041609938416514e-07, 0.00017906920402310789, 1.631619738873269e-06, 2.0158734681530177e-09, 9.65507530290054e-15, 4.2181228128435055e-12, 8.564649545128589e-10, 0.00023218656133394688, 0.6439363956451416, 0.000818322179839015, 1.3831699163802114e-07, 2.1358659198916774e-12, 5.4572883101400294e-08], [1.4084274191361601e-08, 2.1930364191291574e-09, 7.004614666072939e-09, 2.0828078959311824e-06, 6.64705439703539e-05, 3.6118690331932157e-06, 4.0857584676645686e-11, 1.0090924406833124e-12, 5.430448080009356e-15, 6.815135122906213e-09, 0.0007384128402918577, 0.9033229351043701, 0.0037223652470856905, 5.428325380307797e-07, 5.097080588711833e-07], [3.370899046006848e-11, 1.5044922772877722e-12, 1.903236411786996e-13, 5.2399131041103164e-12, 5.3600892613303586e-09, 3.287689196440624e-07, 1.293990137263279e-09, 3.2395277866498207e-13, 8.98320316581696e-19, 7.591717251043266e-18, 2.4333673097343134e-12, 7.08575316821225e-05, 0.3025490641593933, 0.00011370918218744919, 1.7842703314840946e-08], [0.0009491983219049871, 3.734114216058515e-05, 0.00010643315181368962, 4.299266220186837e-05, 0.0019948105327785015, 0.012520392425358295, 0.0005770812276750803, 0.00013455892622005194, 0.0002518744731787592, 0.0005399127840064466, 0.0017743584467098117, 0.004756112117320299, 0.00398082984611392, 0.002925803419202566, 0.1746407300233841]], [[0.1577264666557312, 0.03251823037862778, 0.4939506947994232, 0.8334789872169495, 0.6927971243858337, 0.3147047460079193, 0.7604361176490784, 0.11822030693292618, 0.7022377848625183, 0.6516091823577881, 0.14691989123821259, 0.2232232689857483, 0.14339210093021393, 0.3761228322982788, 0.014605461619794369], [0.028655482456088066, 0.14083503186702728, 0.08485368639230728, 0.8299343585968018, 0.8304422497749329, 0.5664599537849426, 0.834579586982727, 0.7438958287239075, 0.8452481031417847, 0.8614712953567505, 0.3640905022621155, 0.805733323097229, 0.3481642007827759, 0.795884370803833, 0.05269646272063255], [0.02106422185897827, 0.010846637189388275, 0.073356993496418, 0.017661061137914658, 0.8741048574447632, 0.5687165856361389, 0.5249210000038147, 0.5693489909172058, 0.5103186368942261, 0.5253384709358215, 0.6472406387329102, 0.4561024308204651, 0.1524587720632553, 0.45141565799713135, 0.034538887441158295], [0.2203565090894699, 0.02154199220240116, 0.007279311306774616, 0.003464027540758252, 0.18461424112319946, 0.07773485034704208, 0.7297388315200806, 0.2260110229253769, 0.6848539113998413, 0.2328294813632965, 0.22646839916706085, 0.3173597455024719, 0.10388152301311493, 0.06158056855201721, 0.11330780386924744], [0.1574045568704605, 0.12516136467456818, 0.04707150533795357, 0.0032313871197402477, 0.19444315135478973, 0.046962298452854156, 0.48863229155540466, 0.8290899991989136, 0.892469584941864, 0.6836395859718323, 0.83636474609375, 0.47956424951553345, 0.034452617168426514, 0.38761135935783386, 0.055785421282052994], [0.4389230012893677, 0.6133158802986145, 0.4783843159675598, 0.11230780929327011, 0.006951127201318741, 0.0644199401140213, 0.03406795859336853, 0.33251792192459106, 0.9552598595619202, 0.8827710747718811, 0.9276224970817566, 0.8325800895690918, 0.737617552280426, 0.745059609413147, 0.05149900168180466], [0.3395847976207733, 0.09897124767303467, 0.16763220727443695, 0.1671983003616333, 0.049412358552217484, 0.007114487700164318, 0.3340696394443512, 0.018166696652770042, 0.7235669493675232, 0.9639523029327393, 0.851059079170227, 0.7306914925575256, 0.5801126956939697, 0.8017169237136841, 0.08099871873855591], [0.44394704699516296, 0.6082286238670349, 0.37166181206703186, 0.3715074956417084, 0.35315781831741333, 0.10853563994169235, 0.013190319761633873, 0.07092351466417313, 0.03435605764389038, 0.25131845474243164, 0.921750545501709, 0.8745512366294861, 0.7473158240318298, 0.834020733833313, 0.1216883435845375], [0.18251584470272064, 0.8759727478027344, 0.1439245641231537, 0.06640342622995377, 0.060579828917980194, 0.2710072100162506, 0.011089610867202282, 0.034396518021821976, 0.1700025051832199, 0.043876904994249344, 0.14450228214263916, 0.9449294805526733, 0.9689385294914246, 0.939329981803894, 0.07954179495573044], [0.32071176171302795, 0.7452729344367981, 0.11999625712633133, 0.08053360879421234, 0.3748469650745392, 0.31863275170326233, 0.028054066002368927, 0.2197551280260086, 0.01771731488406658, 0.23943577706813812, 0.01906767673790455, 0.8113164901733398, 0.9739595055580139, 0.9691897630691528, 0.21732129156589508], [0.6261264085769653, 0.6649302244186401, 0.5194191336631775, 0.6324451565742493, 0.6771988272666931, 0.7814968228340149, 0.4118405878543854, 0.3728334903717041, 0.03296521306037903, 0.008678224869072437, 0.6047253012657166, 0.11251461505889893, 0.21560458838939667, 0.9244948625564575, 0.10127653181552887], [0.3176693320274353, 0.5172579884529114, 0.1793123036623001, 0.37762320041656494, 0.23678036034107208, 0.5621929168701172, 0.08773050457239151, 0.24525783956050873, 0.010828782804310322, 0.025829488411545753, 0.0057976157404482365, 0.08708162605762482, 0.04166324809193611, 0.5714256167411804, 0.16898052394390106], [0.6460146307945251, 0.8194199800491333, 0.48921409249305725, 0.6910595297813416, 0.5259124636650085, 0.6389046311378479, 0.3241840600967407, 0.7817367911338806, 0.17853572964668274, 0.1606196016073227, 0.06383053213357925, 0.007355134002864361, 0.02128707617521286, 0.02206379547715187, 0.23354344069957733], [0.5992116332054138, 0.6358246803283691, 0.47243836522102356, 0.5617506504058838, 0.6971379518508911, 0.6431114673614502, 0.39991113543510437, 0.8182389140129089, 0.2704472243785858, 0.20400457084178925, 0.059529319405555725, 0.06732083112001419, 0.008503233082592487, 0.06121496111154556, 0.2071741670370102], [0.2342938333749771, 0.5683650374412537, 0.6037701964378357, 0.7331977486610413, 0.7349027395248413, 0.6651985049247742, 0.23853524029254913, 0.2293619066476822, 0.48426058888435364, 0.7077944874763489, 0.5918195843696594, 0.8169012665748596, 0.7005065679550171, 0.4784330725669861, 0.015931207686662674]], [[0.04383472725749016, 0.02773081697523594, 0.016415273770689964, 0.024880478158593178, 0.005487722344696522, 0.14834517240524292, 0.010061212815344334, 0.013310510665178299, 0.03559315577149391, 0.022788431495428085, 0.016539618372917175, 0.022621937096118927, 0.3853665292263031, 0.02895752713084221, 0.21785423159599304], [0.02212689444422722, 0.0360226184129715, 0.0007962794625200331, 0.005733562167733908, 0.0017349227564409375, 0.011109595187008381, 0.02015179581940174, 0.048344310373067856, 0.003794114338234067, 0.016348786652088165, 0.0018908409401774406, 0.010183308273553848, 0.04822028428316116, 0.011540568433701992, 0.21287554502487183], [0.19621919095516205, 0.02568935602903366, 0.012553256005048752, 0.05958101898431778, 0.0049527534283697605, 0.009129918180406094, 0.035662900656461716, 0.006033026147633791, 0.01979534700512886, 0.016174430027604103, 0.025959551334381104, 0.017891131341457367, 0.21532145142555237, 0.010915487073361874, 0.2776879370212555], [0.22681212425231934, 0.26364389061927795, 0.1368870735168457, 0.07472710311412811, 0.004966794513165951, 0.17209400236606598, 0.07595591247081757, 0.10330677032470703, 0.009879215620458126, 0.30214887857437134, 0.027453631162643433, 0.07928238064050674, 0.6068928837776184, 0.0009245484252460301, 0.41711828112602234], [0.03220081329345703, 0.07110226154327393, 0.19687172770500183, 0.32465922832489014, 0.06123804301023483, 0.009123058058321476, 0.008925903588533401, 0.001694322214461863, 0.009767607785761356, 0.012425252236425877, 0.021234901621937752, 0.006749649532139301, 0.022427640855312347, 0.00419656652957201, 0.11337225884199142], [0.1499132513999939, 0.1588381826877594, 0.006192722357809544, 0.06905046850442886, 0.021936854347586632, 0.04223879054188728, 0.01654554158449173, 0.012800824828445911, 0.001194898271933198, 0.011350413784384727, 0.0011690479004755616, 0.03650015965104103, 0.0330234132707119, 0.032408226281404495, 0.30060991644859314], [0.10197536647319794, 0.32784661650657654, 0.22266407310962677, 0.37194594740867615, 0.4840903878211975, 0.2562866806983948, 0.20682689547538757, 0.01685171388089657, 0.02662164717912674, 0.01744299754500389, 0.07043293118476868, 0.06053447723388672, 0.13449640572071075, 0.0437617152929306, 0.15905345976352692], [0.04155902937054634, 0.02725875750184059, 0.06621034443378448, 0.15740959346294403, 0.22226983308792114, 0.11737026274204254, 0.021176597103476524, 0.037896860390901566, 0.001983239781111479, 0.07737525552511215, 0.040612466633319855, 0.036445699632167816, 0.04206009954214096, 0.005294053349643946, 0.22695806622505188], [0.3731417655944824, 0.020610323175787926, 0.04687204957008362, 0.19942151010036469, 0.0219199787825346, 0.023319954052567482, 0.607546865940094, 0.0038317576982080936, 0.05746426433324814, 0.0039819530211389065, 0.0020286834333091974, 0.023514816537499428, 0.0007224131841212511, 0.0017132725333794951, 0.31377115845680237], [0.007707278709858656, 0.04994801804423332, 0.0602150596678257, 0.1843070536851883, 0.023052150383591652, 0.00867108628153801, 0.0030793596524745226, 0.008175634779036045, 0.3707427382469177, 0.032583341002464294, 0.030614105984568596, 0.003414844162762165, 0.0027733321767300367, 0.00039667857345193624, 0.06665757298469543], [0.06275568902492523, 0.15385569632053375, 0.07121506333351135, 0.04657430946826935, 0.08974524587392807, 0.017753345891833305, 0.09537442773580551, 0.08409535884857178, 0.4617481529712677, 0.05371565744280815, 0.051210206001996994, 0.014556556940078735, 0.0261379461735487, 0.0015151489060372114, 0.25993233919143677], [0.037524934858083725, 0.08964382112026215, 0.11503562331199646, 0.2385229468345642, 0.14595970511436462, 0.01507873460650444, 0.07354842126369476, 0.014194677583873272, 0.01029899064451456, 0.3145633935928345, 0.08443433046340942, 0.02799280546605587, 0.006364578381180763, 0.0011598452692851424, 0.25597554445266724], [0.03498825803399086, 0.003427299438044429, 0.012860815972089767, 0.00960747804492712, 0.0073430403135716915, 0.002194140339270234, 0.020218953490257263, 0.04016692563891411, 0.0035721054300665855, 0.11439335346221924, 0.03179614990949631, 0.0055262502282857895, 0.08811097592115402, 0.0019241927657276392, 0.31578439474105835], [0.0003122057532891631, 0.0005657155998051167, 0.0003099576279055327, 0.018182117491960526, 8.608390635345131e-05, 0.00029685357003472745, 0.00030423246789723635, 0.0039575002156198025, 0.00041145391878671944, 0.0009832053910940886, 0.0007515411707572639, 0.006357411853969097, 0.3007054328918457, 0.00010537439811741933, 0.00161165336612612], [0.052370160818099976, 0.019386928528547287, 0.0404941625893116, 0.16087706387043, 0.14014431834220886, 0.0561581589281559, 0.1907973736524582, 0.027806226164102554, 0.022970959544181824, 0.05846026912331581, 0.09902504831552505, 0.038958851248025894, 0.016928229480981827, 0.04114920645952225, 0.14461401104927063]], [[0.1774463951587677, 0.26868411898612976, 0.03527391701936722, 0.01705012284219265, 0.00047759010340087116, 0.006241941824555397, 0.0031507122330367565, 0.2944689095020294, 0.038735195994377136, 0.003944840747863054, 0.004385389853268862, 0.004225992131978273, 0.03986744210124016, 0.00549504067748785, 0.07870971411466599], [0.00027908835909329355, 0.005506355315446854, 0.001626787707209587, 0.13775548338890076, 0.0008261757320724428, 0.00028156363987363875, 0.0002459189563523978, 0.0025131029542535543, 0.0009445812902413309, 0.001017659087665379, 0.002250042976811528, 0.0015115974238142371, 0.0017954352078959346, 0.0006745054270140827, 0.21780018508434296], [0.021244889125227928, 0.1178143173456192, 0.008956437930464745, 0.14321640133857727, 0.023635229095816612, 0.3068733811378479, 0.15845780074596405, 0.3092327415943146, 0.0024783278349786997, 0.06481246650218964, 0.008965774439275265, 0.019083118066191673, 0.04005150496959686, 0.01112168189138174, 0.19139143824577332], [0.00042023108107969165, 0.0008873279439285398, 0.0019056870369240642, 0.007766622584313154, 0.23140135407447815, 0.5036463141441345, 0.015440672636032104, 0.008361338637769222, 0.001879698014818132, 0.0006688520661555231, 0.01133010908961296, 0.09722423553466797, 0.03314661607146263, 0.006971372757107019, 0.02285030484199524], [0.002678314223885536, 0.004764833487570286, 0.0003137744788546115, 0.0006636036559939384, 0.07552827149629593, 0.36051952838897705, 0.21059149503707886, 0.11911091953516006, 0.00013829045929014683, 0.00018005385936703533, 0.00021675217431038618, 0.007453517522662878, 0.004449300933629274, 0.03708551451563835, 0.13281597197055817], [0.008487393148243427, 0.014329447411000729, 0.005103611387312412, 0.0017902699764817953, 0.00018748251022771, 0.07080603390932083, 0.1865091174840927, 0.03389747440814972, 0.0026728338561952114, 0.00012369015894364566, 0.0001717496052151546, 0.0016556874616071582, 0.0035823825746774673, 0.018341869115829468, 0.2051384449005127], [0.0016413311241194606, 0.0038119314704090357, 0.0005628983490169048, 6.117233715485781e-05, 0.00011399950017221272, 0.0007454796577803791, 0.054881561547517776, 0.30246245861053467, 0.15667226910591125, 0.0004453254514373839, 0.0002609542279969901, 0.0001120980887208134, 0.0006856885738670826, 0.00573006272315979, 0.011146760545670986], [0.001007524086162448, 0.0022212164476513863, 0.00036003260174766183, 2.8946307793376036e-05, 1.0167077562073246e-05, 0.00012231878645252436, 0.00022786400222685188, 0.03619853034615517, 0.005354967433959246, 0.003357505425810814, 0.0005030903848819435, 5.3131421736907214e-05, 4.2532476072665304e-05, 0.00010396525613032281, 0.2518664300441742], [0.004948427900671959, 0.0037361346185207367, 0.0040338728576898575, 0.0015943445032462478, 3.9753424061927944e-05, 0.00016846440848894417, 0.00017597683472558856, 0.003258961718529463, 0.06328149139881134, 0.43567389249801636, 0.03252503648400307, 0.006277996581047773, 3.634384847828187e-05, 2.672040500328876e-05, 0.030029548332095146], [0.00322673749178648, 0.017767680808901787, 0.0033617434091866016, 0.029219835996627808, 0.0009114073473028839, 0.002889687195420265, 0.00012576105655170977, 0.01574547402560711, 0.0018639388727024198, 0.6032934188842773, 0.1301620751619339, 0.04121570661664009, 0.0035096178762614727, 0.00032833084696903825, 0.3004224896430969], [0.033899419009685516, 0.07324357330799103, 0.00985381193459034, 0.017461512237787247, 0.019165849313139915, 0.07006029784679413, 0.01799222268164158, 0.013579626567661762, 0.00021177329472266138, 0.026033537462353706, 0.13102787733078003, 0.2077469676733017, 0.7029638886451721, 0.029135672375559807, 0.05414650961756706], [0.0015424743760377169, 0.007544125430285931, 0.010602829977869987, 0.0016127177514135838, 0.006006686482578516, 0.08514653891324997, 0.003129118587821722, 0.0036380700767040253, 1.298951519856928e-05, 6.919799488969147e-05, 0.0003367147874087095, 0.031529009342193604, 0.36636054515838623, 0.21289798617362976, 0.04463290795683861], [0.005653384607285261, 0.005221519153565168, 0.010438429191708565, 0.0023121859412640333, 0.0034771040081977844, 0.01156994141638279, 0.006321457680314779, 0.006196276750415564, 2.671167931111995e-05, 0.00012823205906897783, 0.00023895784397609532, 0.0015353390481323004, 0.06888392567634583, 0.3010466396808624, 0.05789510905742645], [0.0025978884659707546, 0.0011408268474042416, 0.0005907863960601389, 0.0073682027868926525, 5.514698841579957e-06, 0.0001586068101460114, 0.0016139426734298468, 0.002635698765516281, 2.2516995159094222e-05, 7.803570952091832e-06, 4.170422926108586e-06, 4.799172893399373e-05, 8.148160122800618e-05, 0.006126015912741423, 0.363029420375824], [0.018444720655679703, 0.036891017109155655, 0.08301377296447754, 0.04485299810767174, 0.0371856652200222, 0.0472157783806324, 0.022677546367049217, 0.017107300460338593, 0.03217196837067604, 0.03369837626814842, 0.021089907735586166, 0.018274538218975067, 0.020997297018766403, 0.034321803599596024, 0.1648317128419876]], [[0.2133164256811142, 0.025492815300822258, 0.20653849840164185, 0.07043907791376114, 0.10411863774061203, 0.3043566346168518, 0.06760577112436295, 0.5064103603363037, 0.08081910014152527, 0.27507925033569336, 0.5432406663894653, 0.27881479263305664, 0.16320040822029114, 0.2653813064098358, 0.11116068065166473], [0.015402763150632381, 0.2444494515657425, 0.0030522451270371675, 0.00048490799963474274, 0.0026600188575685024, 0.06905494630336761, 0.012269481085240841, 0.014592616818845272, 0.004205085337162018, 0.0039128707721829414, 0.0037959537003189325, 0.012499181553721428, 0.02713301219046116, 0.00563135975971818, 0.19437076151371002], [0.04805738478899002, 0.007929358631372452, 0.4969516396522522, 0.08109094947576523, 0.008613435551524162, 0.06128339096903801, 0.020970679819583893, 0.014624540694057941, 0.001800250494852662, 0.04372387006878853, 0.036881472915410995, 0.022519467398524284, 0.032134752720594406, 0.17586740851402283, 0.15428785979747772], [0.021660206839442253, 0.06483402103185654, 0.07990853488445282, 0.8655576705932617, 0.10770212858915329, 0.042777951806783676, 0.004243527539074421, 0.04141073673963547, 0.0011197980493307114, 0.0010354480473324656, 0.007620980031788349, 0.009411019273102283, 0.023886993527412415, 0.8532692193984985, 0.009252375923097134], [0.03802541270852089, 0.5626884698867798, 0.3869370222091675, 0.012873617932200432, 0.11968709528446198, 0.014900745823979378, 0.02957817167043686, 0.018288375809788704, 0.005979553796350956, 0.03379013389348984, 0.016338851302862167, 0.01766209304332733, 0.8086205720901489, 0.08052025735378265, 0.13067808747291565], [0.0663566142320633, 0.02082742564380169, 0.009716741740703583, 0.003548208624124527, 0.0008020728128030896, 0.4547119140625, 0.03523911535739899, 0.0031006578356027603, 0.006736437324434519, 0.0009184986702166498, 0.0011584048625081778, 0.04212343320250511, 0.019468490034341812, 0.001240313402377069, 0.20631356537342072], [0.004470710642635822, 0.02006937935948372, 0.020011691376566887, 0.019766854122281075, 0.12330501526594162, 0.15558527410030365, 0.04160740226507187, 0.1780312955379486, 0.014384130015969276, 0.005233153235167265, 0.004123131278902292, 0.05227937176823616, 0.013469746336340904, 0.022578507661819458, 0.07922197878360748], [0.17898443341255188, 0.006772744003683329, 0.041487641632556915, 0.009575014933943748, 0.016729410737752914, 0.2668032944202423, 0.12321095168590546, 0.6781973838806152, 0.0025635806377977133, 0.01087682880461216, 0.002732365159317851, 0.020299792289733887, 0.0031363710295408964, 0.0008204782498069108, 0.05180227383971214], [0.12461799383163452, 0.013122161850333214, 0.02311752177774906, 0.0762406587600708, 0.09383975714445114, 0.007501720450818539, 0.07133012264966965, 0.008159258402884007, 0.13900579512119293, 0.006521029397845268, 0.021471921354532242, 0.012502939440310001, 0.0014349960256367922, 0.011674328707158566, 0.3848530650138855], [0.014992507174611092, 0.010756749659776688, 0.10129547864198685, 0.15213072299957275, 0.1363232582807541, 0.16603931784629822, 0.0040587568655610085, 0.505429208278656, 0.0025213102344423532, 0.05678342655301094, 0.20746274292469025, 0.04314066469669342, 0.0019582516979426146, 0.01985819824039936, 0.18090446293354034], [0.11427638679742813, 0.0123747568577528, 0.020808644592761993, 0.1336503028869629, 0.008563186042010784, 0.09643486887216568, 0.15193390846252441, 0.050255559384822845, 0.0023536821827292442, 0.3208443820476532, 0.021319447085261345, 0.003293143818154931, 0.027340535074472427, 0.01197835523635149, 0.09007034450769424], [0.15923485159873962, 0.11477550864219666, 0.21969333291053772, 0.09681756794452667, 0.07061057537794113, 0.1670638769865036, 0.1398637294769287, 0.059452954679727554, 0.00850652251392603, 0.062244825065135956, 0.03212086483836174, 0.10482167452573776, 0.05658517777919769, 0.03675027936697006, 0.24718202650547028], [0.004966236650943756, 0.001515651005320251, 0.002549123717471957, 0.006106496322900057, 0.00036676786839962006, 0.0014838402858003974, 0.008350875228643417, 0.003760475432500243, 9.004020830616355e-05, 0.003012964967638254, 0.000879374798387289, 0.0023141989950090647, 0.5349817276000977, 0.00013737898552790284, 0.18041089177131653], [3.0577066354453564e-05, 0.00011073229688918218, 0.0002722943318076432, 0.00012968607188668102, 3.925479541067034e-05, 9.284611587645486e-05, 1.1375399481039494e-05, 0.00013649655738845468, 2.160583608201705e-05, 3.872126853821101e-06, 4.776401965500554e-06, 5.892393892281689e-05, 0.3018791675567627, 0.0016873051645234227, 0.00020723984926007688], [0.0053407615050673485, 0.002270790981128812, 0.015077341347932816, 0.008943013846874237, 0.01947944425046444, 0.013856526464223862, 0.021029049530625343, 0.011522401124238968, 0.019980257377028465, 0.021877266466617584, 0.03018842823803425, 0.06539047509431839, 0.04945596680045128, 0.008784771896898746, 0.1688213050365448]], [[0.09667091816663742, 0.08969368785619736, 0.16646768152713776, 0.01428181305527687, 0.1262292116880417, 0.03015410713851452, 0.00857650488615036, 0.013287652283906937, 0.013465571217238903, 0.009945754893124104, 0.03584994748234749, 0.07976501435041428, 0.013894102536141872, 0.07191513478755951, 0.16682514548301697], [0.00307486648671329, 0.2169581949710846, 0.015313946641981602, 0.005070009268820286, 0.13766343891620636, 0.036365993320941925, 0.013734312728047371, 0.012890451587736607, 0.00037508379318751395, 0.002069024136289954, 0.0038654597010463476, 0.007793853525072336, 0.006365353707224131, 0.02897111512720585, 0.19472798705101013], [0.013033762574195862, 0.0016745100729167461, 0.09789733588695526, 0.11557573825120926, 0.070904940366745, 0.039959780871868134, 0.06112189590930939, 0.005926545709371567, 0.05931684747338295, 0.06562750041484833, 0.015556245110929012, 0.2949027419090271, 0.09280899167060852, 0.18960142135620117, 0.2321171909570694], [0.0009253448224626482, 0.0011463494738563895, 0.0022407870274037123, 0.022192178294062614, 0.18083734810352325, 0.18906380236148834, 0.06340676546096802, 0.5556718111038208, 0.008876022882759571, 0.00195835973136127, 0.009641225449740887, 0.13488754630088806, 0.03692271187901497, 0.0069083282724022865, 0.19416382908821106], [0.020195724442601204, 0.0026999269612133503, 0.0047158133238554, 0.017117822542786598, 0.22690622508525848, 0.009801734238862991, 0.18513473868370056, 0.000916039280127734, 0.006044555455446243, 0.006021710112690926, 0.010346228256821632, 0.04500352963805199, 0.008295656181871891, 0.1122727021574974, 0.4271945357322693], [0.02983868308365345, 0.03651329129934311, 0.005064305383712053, 0.00043434457620605826, 0.001774297677911818, 0.10316617041826248, 0.10274261981248856, 0.570116400718689, 0.0018607155652716756, 0.004884766880422831, 0.0001192242925753817, 0.01004798710346222, 0.011760696768760681, 0.020220324397087097, 0.036799319088459015], [0.020830435678362846, 0.04066089913249016, 0.01340602245181799, 0.0007146665593609214, 0.05329689383506775, 0.010700137354433537, 0.06310626864433289, 0.1416247934103012, 0.059007443487644196, 0.009734428487718105, 0.023192377761006355, 0.030464952811598778, 0.011454294435679913, 0.06458231806755066, 0.29838618636131287], [0.04047420993447304, 0.05575861781835556, 0.0035385461524128914, 0.00047053993330337107, 0.010776028037071228, 0.0002634078555274755, 0.006466362159699202, 0.09768779575824738, 0.011305907741189003, 0.6455902457237244, 0.005685864482074976, 0.009437574073672295, 0.0014128481270745397, 0.0036261524073779583, 0.1994941532611847], [0.001968077849596739, 0.00013096239126753062, 0.014192181639373302, 0.0025808673817664385, 1.1752749742299784e-05, 7.090794679243118e-05, 8.489128958899528e-05, 7.501097570639104e-05, 0.005588378757238388, 0.00024033378576859832, 0.7911840081214905, 0.0006417080294340849, 0.00012212486763019115, 0.0026151463389396667, 0.024830428883433342], [0.007711799815297127, 0.006852409336715937, 0.005409319419413805, 0.029324712231755257, 0.0012151957489550114, 0.0014427780406549573, 0.0002848623844329268, 0.0011284908978268504, 0.00042831210885196924, 0.0035933239851146936, 0.2853389084339142, 0.04352247342467308, 0.0011324246879667044, 0.0015205255476757884, 0.05924868583679199], [0.06333743035793304, 0.004831443540751934, 0.017261236906051636, 0.05893971398472786, 0.005950291641056538, 0.002105317311361432, 0.003185122972354293, 0.0028415010310709476, 0.004572128411382437, 0.007815520279109478, 0.07613655924797058, 0.10669270157814026, 0.027066918089985847, 0.03207901865243912, 0.4743220806121826], [0.10327208787202835, 0.004544916562736034, 0.05445469170808792, 0.010814311914145947, 0.026858847588300705, 0.011217474937438965, 0.07071709632873535, 0.05960191786289215, 0.0010665962472558022, 0.025403864681720734, 0.006131312809884548, 0.5720618963241577, 0.029676837846636772, 0.17520834505558014, 0.23297326266765594], [0.011414228938519955, 0.002735550981014967, 0.015156290493905544, 0.0027777000796049833, 0.009832575917243958, 0.015552453696727753, 0.017305195331573486, 0.004722784738987684, 4.7792200348339975e-05, 0.0034479873720556498, 0.0004017044266220182, 0.0011886333813890815, 0.18307994306087494, 0.2786843478679657, 0.04159880056977272], [0.0032662157900631428, 0.004168938845396042, 0.0016457620076835155, 0.0005059303948655725, 0.0003206630062777549, 0.000853654695674777, 0.010604765266180038, 0.005784912034869194, 0.00014833646127954125, 0.0001704594906186685, 5.580573997576721e-05, 0.0004662217397708446, 0.0009024841128848493, 0.025914611294865608, 0.3543371260166168], [0.057395875453948975, 0.01834016665816307, 0.017516011372208595, 0.011936328373849392, 0.010095582343637943, 0.018046732991933823, 0.24530914425849915, 0.01257838774472475, 0.014466731809079647, 0.027552323415875435, 0.054997242987155914, 0.013960911892354488, 0.0074861980974674225, 0.03251070901751518, 0.14566579461097717]], [[0.3107149600982666, 0.049285680055618286, 0.08128133416175842, 0.03986956924200058, 0.07088969647884369, 0.1961679309606552, 0.15016919374465942, 0.05429982393980026, 0.1291487067937851, 0.03663256764411926, 0.25306442379951477, 0.3913470208644867, 0.2542778253555298, 0.252127081155777, 0.15921251475811005], [0.10834414511919022, 0.3508348762989044, 0.02124197781085968, 0.019397908821702003, 0.026673240587115288, 0.3167271912097931, 0.11886779963970184, 0.17699773609638214, 0.14507175981998444, 0.115145742893219, 0.6241064667701721, 0.1622784435749054, 0.5683063268661499, 0.15724869072437286, 0.12728430330753326], [0.6979861855506897, 0.039286430925130844, 0.3014020621776581, 0.003208757843822241, 0.01772892102599144, 0.014036925509572029, 0.19886529445648193, 0.09335973858833313, 0.4060034155845642, 0.28424081206321716, 0.26539483666419983, 0.1895008385181427, 0.4672236740589142, 0.16107353568077087, 0.10992881655693054], [0.5298255681991577, 0.6474234461784363, 0.19260530173778534, 0.026028962805867195, 0.013013242743909359, 0.01466711051762104, 0.11121421307325363, 0.06523838639259338, 0.29339125752449036, 0.46135157346725464, 0.7174844145774841, 0.3618351221084595, 0.19526919722557068, 0.0703459233045578, 0.24330592155456543], [0.7494951486587524, 0.23358309268951416, 0.3640848398208618, 0.09014757722616196, 0.32190942764282227, 0.0021980239544063807, 0.07713330537080765, 0.030900368466973305, 0.08560045808553696, 0.26394325494766235, 0.11549779027700424, 0.44356539845466614, 0.12175428122282028, 0.3783136308193207, 0.14015373587608337], [0.3064809739589691, 0.15617568790912628, 0.4955383241176605, 0.8125641942024231, 0.02114781178534031, 0.2633197009563446, 0.014569958671927452, 0.04754461348056793, 0.03227522596716881, 0.09995166957378387, 0.0697590634226799, 0.0770602896809578, 0.19454655051231384, 0.18272873759269714, 0.19963966310024261], [0.5314973592758179, 0.5086395144462585, 0.5757231116294861, 0.44031307101249695, 0.2709468603134155, 0.0639616996049881, 0.2984015941619873, 0.0039451331831514835, 0.0197422094643116, 0.0031917106825858355, 0.05093149095773697, 0.12591752409934998, 0.25977155566215515, 0.0615861676633358, 0.3711840510368347], [0.2939777970314026, 0.2997593581676483, 0.5167340040206909, 0.46100836992263794, 0.39705657958984375, 0.5034002065658569, 0.07978513836860657, 0.0779491513967514, 0.012053987942636013, 0.01132633350789547, 0.028715649619698524, 0.059212565422058105, 0.20603224635124207, 0.15584728121757507, 0.14816488325595856], [0.3128078877925873, 0.0864272266626358, 0.7678588032722473, 0.6537591814994812, 0.8236088752746582, 0.6979317665100098, 0.30976778268814087, 0.014760972931981087, 0.5645584464073181, 0.004590533208101988, 0.008271697908639908, 0.012132997624576092, 0.028745530173182487, 0.04464057460427284, 0.1669740080833435], [0.6456499099731445, 0.1693999022245407, 0.7097220420837402, 0.5244839191436768, 0.46365103125572205, 0.5023244023323059, 0.9643971920013428, 0.24913577735424042, 0.13337120413780212, 0.06419410556554794, 0.012416149489581585, 0.0573885552585125, 0.016666844487190247, 0.008706454187631607, 0.1754455268383026], [0.09960467368364334, 0.0907629206776619, 0.36143985390663147, 0.11092879623174667, 0.19937658309936523, 0.03214935213327408, 0.3196737766265869, 0.4763943552970886, 0.497630774974823, 0.1899363249540329, 0.1145005002617836, 0.004749455489218235, 0.0008605146431364119, 0.0007969819707795978, 0.02025206945836544], [0.3807562589645386, 0.26623356342315674, 0.4209006428718567, 0.27443018555641174, 0.5137820839881897, 0.1592678278684616, 0.6250110864639282, 0.6178545951843262, 0.9692861437797546, 0.5716569423675537, 0.22724294662475586, 0.17567582428455353, 0.008769324980676174, 0.002557128667831421, 0.05025441572070122], [0.2969632148742676, 0.16767999529838562, 0.46978121995925903, 0.28813451528549194, 0.45300158858299255, 0.33029136061668396, 0.6236194968223572, 0.1634167730808258, 0.8177276253700256, 0.718397855758667, 0.9021148681640625, 0.07875741273164749, 0.09992827475070953, 0.004932410083711147, 0.1707668900489807], [0.3945808410644531, 0.3581867516040802, 0.5247420072555542, 0.4120633900165558, 0.3024104833602905, 0.35548633337020874, 0.5872392654418945, 0.15815261006355286, 0.7289484143257141, 0.7948301434516907, 0.9396543502807617, 0.9256777167320251, 0.08537369966506958, 0.03166399896144867, 0.03224433213472366], [0.004588960204273462, 0.041907694190740585, 0.17755450308322906, 0.039724841713905334, 0.047663237899541855, 0.09274838864803314, 0.010110240429639816, 0.014862497337162495, 0.11161036789417267, 0.0490046888589859, 0.18517035245895386, 0.029471391811966896, 0.05094437301158905, 0.002971563721075654, 0.16300250589847565]], [[6.113462859502761e-06, 0.5065946578979492, 7.261813152581453e-05, 5.1066386498122354e-14, 1.0490246824277965e-15, 1.4956003015903496e-12, 2.5734427609724886e-13, 2.1143946469237562e-06, 9.544867651811728e-08, 4.2543565892394497e-10, 6.215519418595328e-12, 1.687761909396901e-11, 1.6993320528513323e-08, 1.0583119935958507e-09, 9.857150189418462e-07], [4.727198188447801e-08, 0.002272214274853468, 0.8730366826057434, 0.0016238681273534894, 9.849362297975617e-11, 6.310171162720105e-14, 1.3311845115798748e-12, 1.350557283785747e-07, 1.07800769910682e-05, 3.4101576602552086e-05, 7.529693561991735e-07, 3.7022258592145363e-09, 3.1551092294357375e-10, 8.851498527195911e-12, 1.024629546009237e-05], [6.003397223786067e-10, 5.335852165444521e-06, 0.00445933174341917, 0.5796651840209961, 5.976808097329922e-05, 2.377180230439535e-09, 1.7792844021063958e-12, 1.2140626282075573e-09, 6.417224529542409e-09, 2.601910637167748e-06, 1.1842810181406094e-06, 1.8266834445057611e-07, 1.3081095096012518e-09, 1.5776791765370612e-12, 4.7676843678345904e-05], [2.4071971206038626e-15, 2.3560551770727793e-14, 9.98394700246763e-11, 1.7167060661904543e-07, 0.2774648666381836, 1.6012703781598248e-05, 9.760837530760607e-15, 4.654387315338889e-18, 8.039692137064508e-20, 2.1508527635127157e-16, 1.789740057545064e-11, 2.4233797191186568e-08, 2.7592322870972907e-10, 4.956549239646573e-15, 1.5411848153235042e-06], [1.9919477308935618e-13, 5.266535346254387e-16, 1.2917133013982517e-14, 7.221083175856791e-10, 8.195231930585578e-05, 0.5564944744110107, 4.117699063499458e-06, 5.438900198273533e-13, 2.4172004338169554e-20, 9.57835365503234e-22, 9.376302678036402e-17, 3.235451073724249e-10, 6.101883442966027e-09, 9.971044129253315e-11, 1.6162671201414014e-08], [9.771466125130246e-08, 3.17872256294649e-11, 3.1429036890379125e-13, 5.901367481980172e-16, 4.2342058748090494e-09, 0.0012305855052545667, 0.6103256940841675, 2.2161180822877213e-05, 7.972257402844019e-12, 6.481494664823834e-19, 5.35928561114305e-19, 7.863773244772346e-14, 1.1593314752644801e-07, 8.808668212623161e-07, 1.1730364235518209e-07], [2.6939844799400703e-10, 3.892770337188267e-07, 2.2438891023046637e-10, 2.095593632707407e-18, 1.8655412772298346e-14, 2.206185598652155e-07, 3.0316745323943906e-05, 0.33891788125038147, 5.437008439912461e-06, 1.3213468337612382e-14, 2.5347562276209975e-18, 1.0659246862729562e-18, 2.6392999114346893e-13, 9.868956762915104e-10, 1.6170986327779246e-06], [1.3015508670832787e-09, 4.1474245904282725e-07, 7.619819371029735e-06, 9.079691751061325e-13, 5.725895077835787e-16, 1.0568446176517903e-14, 8.978999488373773e-11, 2.253716047562193e-05, 0.9323674440383911, 0.0001553743495605886, 1.1094852814252931e-10, 4.251380123255501e-17, 3.4548606558270072e-18, 1.563022274271835e-14, 1.7832363141678798e-07], [1.2218349942916262e-10, 4.9370779464652514e-08, 1.0212672805209877e-06, 3.802215486903293e-11, 4.1323817879847246e-16, 3.8503187577578586e-16, 6.2032051316354e-15, 3.2203126920649083e-07, 8.202762546716258e-05, 0.5051153898239136, 1.6483796571264975e-05, 2.317061202194298e-13, 9.134085045449695e-19, 4.959048342554486e-21, 1.9839136555788173e-08], [3.5615963439117673e-14, 6.311461336200308e-12, 7.572167781688677e-09, 7.864790063649707e-08, 5.871175941252194e-13, 4.399392566282849e-15, 3.6105855357745724e-20, 8.408651243829376e-14, 2.915925279012299e-09, 2.7294316168990918e-05, 0.31493836641311646, 1.4271394093157141e-06, 7.57530499374999e-14, 1.0444343699767344e-21, 5.65783730976932e-09], [1.619628042792698e-10, 6.862534152052291e-11, 7.238428190170509e-10, 5.1994692995549485e-08, 8.193378420173758e-08, 6.734891755399985e-09, 1.47457238341411e-14, 5.793711288450045e-15, 1.5065480465795492e-14, 1.167909147170576e-08, 0.0003541565383784473, 0.5504465699195862, 2.5677532903500833e-05, 4.9321430864142715e-14, 1.3459792569392448e-07], [8.003913504195381e-11, 5.626729984720136e-12, 4.9737857062137625e-12, 1.4365373474101162e-11, 1.165467935493325e-07, 3.263785401941277e-05, 9.4434834951862e-11, 2.6144878938953817e-15, 6.540743544149476e-19, 2.5930401594030658e-17, 1.8366722587259687e-09, 1.8794700736179948e-05, 0.49058014154434204, 8.066950840657228e-07, 1.3585024589701788e-06], [1.0801989728040362e-12, 2.2359935084037552e-13, 1.1691597126203823e-12, 1.0214807062303036e-16, 2.4270561688882752e-12, 4.4484740890915475e-10, 1.1468358207533669e-10, 1.5131759777478604e-13, 3.7208958865722007e-20, 6.888861115537483e-21, 1.5888746801787275e-18, 3.2241334168431335e-12, 5.685043561243219e-06, 0.3912107050418854, 3.0407140694244106e-10], [5.397048425948014e-07, 2.3629811494174646e-06, 8.614414923613367e-07, 8.006720286779512e-13, 4.92412575016192e-14, 2.066644277931573e-08, 0.00031528103863820434, 0.011093947105109692, 3.7555511767095595e-07, 1.151808547627739e-13, 5.505821095062543e-16, 1.6971218267519683e-12, 5.383023108151974e-06, 0.8731740117073059, 0.04139598086476326], [0.6266164779663086, 0.3128010928630829, 0.06246759742498398, 0.00042505442979745567, 0.008534153923392296, 0.09425555169582367, 0.2709643542766571, 0.686626672744751, 0.3142872750759125, 0.10107265412807465, 0.015935143455863, 0.012286541052162647, 0.14970052242279053, 0.3989029824733734, 0.022492708638310432]]], [[[0.1393769532442093, 0.0735321119427681, 0.701509952545166, 0.10650816559791565, 0.05110495164990425, 0.021589145064353943, 0.0033319133799523115, 0.0014166238252073526, 0.01486207265406847, 0.006584684830158949, 0.002582702785730362, 0.0004108685825485736, 0.010701421648263931, 0.009390643797814846, 0.06290604919195175], [0.0030957262497395277, 0.0237117987126112, 0.7945073246955872, 0.09792613238096237, 0.2614360749721527, 0.179405078291893, 0.011310527101159096, 0.009954328648746014, 0.009489532560110092, 0.0005609119543805718, 0.000751268700696528, 0.0001462608779547736, 0.004604416899383068, 0.004964352585375309, 0.019775664433836937], [0.002461136318743229, 0.024594180285930634, 0.009559455327689648, 0.055053047835826874, 0.30010533332824707, 0.4690517783164978, 0.03334644436836243, 0.0075769852846860886, 0.007821744307875633, 0.004109389614313841, 0.0022267017047852278, 0.000916018383577466, 0.0037954216822981834, 0.0007741246954537928, 0.004415341652929783], [0.0019876149017363787, 0.0012237336486577988, 0.00015556006110273302, 0.0003553472051862627, 0.4419420659542084, 0.6252713799476624, 0.02062046155333519, 0.0028509902767837048, 0.00548406969755888, 0.0003452444798313081, 0.0001962203241419047, 0.0008938669925555587, 0.0009214308229275048, 1.2216354662086815e-05, 0.0019377138232812285], [0.00020824302919209003, 0.00021322975226212293, 4.6913473852328025e-06, 0.00017657040734775364, 0.0005752452998422086, 0.5289100408554077, 0.1970362812280655, 0.12947966158390045, 0.0005265067447908223, 0.000227929005632177, 6.233566091395915e-05, 0.0001991882745642215, 0.00032238851417787373, 0.0003627484547905624, 0.0016414258861914277], [0.0010278578847646713, 0.0029486939311027527, 0.00014835220645181835, 0.00036925319000147283, 0.00742883887141943, 0.03272741660475731, 0.8576475977897644, 0.03500620648264885, 0.2982224225997925, 0.0003585784579627216, 5.663683623424731e-05, 0.0011889662127941847, 0.00576341338455677, 0.003998933359980583, 0.03130826726555824], [0.002113666385412216, 0.004151111003011465, 0.002428078791126609, 0.002119476906955242, 0.001100956811569631, 0.003687644377350807, 0.13543397188186646, 0.11922256648540497, 0.7567945718765259, 0.2570010721683502, 0.004903816152364016, 0.0001005519661703147, 0.000830159813631326, 0.001259618904441595, 0.14076685905456543], [0.0010344160255044699, 0.00660368800163269, 0.0025270660407841206, 0.00023567670723423362, 0.0004021638887934387, 0.0030120171140879393, 0.0016376315616071224, 0.0524386465549469, 0.7797302007675171, 0.1269131302833557, 0.004214781802147627, 0.0002750723797362298, 0.002267329953610897, 0.001067862962372601, 0.16698867082595825], [0.0009750229655764997, 0.0120720649138093, 0.0038384809158742428, 0.0036232813727110624, 0.004431525245308876, 0.0007613649941049516, 5.662842158926651e-05, 0.01338160876184702, 0.041878536343574524, 0.7091978788375854, 0.2535402476787567, 0.13969287276268005, 0.026510832831263542, 0.0006678565987385809, 0.015569130890071392], [0.0002093962684739381, 0.00030164673808030784, 0.00010105424007633701, 5.030819465901004e-06, 0.001411793869920075, 0.003664590884000063, 0.00017403968377038836, 0.0011218853760510683, 0.011106000281870365, 0.003924186807125807, 0.07315385341644287, 0.3008219599723816, 0.36353737115859985, 0.025737306103110313, 0.0060785748064517975], [0.0001716838014544919, 0.0008840822265483439, 4.3183892557863146e-05, 3.6494086543825688e-06, 0.0005770743009634316, 0.010045445524156094, 0.00010205945727648214, 6.57988857710734e-05, 0.0006949909729883075, 0.004452799912542105, 0.009000658988952637, 0.49080607295036316, 0.17717383801937103, 0.11174798011779785, 0.021669577807188034], [0.019416164606809616, 0.0014941463014110923, 0.001027028076350689, 0.001502541359513998, 0.0085412273183465, 0.12493651360273361, 0.0035243057645857334, 0.0026196581311523914, 0.0008317703031934798, 0.0015569254755973816, 0.060888972133398056, 0.06929422169923782, 0.3396435081958771, 0.387500524520874, 0.017253199592232704], [0.04994890093803406, 0.15025374293327332, 0.024391163140535355, 0.00227133696898818, 0.012616162188351154, 0.2894521951675415, 0.4185648262500763, 0.19089959561824799, 0.027421748265624046, 0.001001756638288498, 0.0036985764745622873, 0.06802930682897568, 0.02484762854874134, 0.057649459689855576, 0.1606004238128662], [0.03736208751797676, 0.11793919652700424, 0.0180205088108778, 0.0001436693564755842, 0.0030756669584661722, 0.08228655159473419, 0.12110688537359238, 0.09650447964668274, 0.015347721055150032, 0.0004259537090547383, 0.00022625335259363055, 0.001013986300677061, 0.0784289613366127, 0.2240448147058487, 0.18707746267318726], [0.7529165148735046, 0.7075774073600769, 0.6068683862686157, 0.3852986991405487, 0.6197313666343689, 0.6735447645187378, 0.6598724722862244, 0.7226093411445618, 0.31395286321640015, 0.2518909275531769, 0.07010441273450851, 0.21793116629123688, 0.4325476884841919, 0.7029338479042053, 0.06848814338445663]], [[0.0006553527782671154, 0.5631614327430725, 0.0008777088369242847, 0.00020331511041149497, 0.0014234310947358608, 0.013944034464657307, 9.958680493582506e-06, 0.01898920349776745, 0.00014103656576480716, 1.4779416233068332e-06, 1.1701366275929104e-07, 1.195983372781484e-06, 0.00012817273091059178, 3.365538941579871e-05, 0.00028557839686982334], [0.00638999929651618, 0.7093943953514099, 0.004974186420440674, 0.06159398332238197, 0.003979360219091177, 0.06536109745502472, 0.005324128083884716, 0.02885170467197895, 0.0003847253101412207, 0.0002721542550716549, 4.3882369936909527e-05, 0.00024302180099766701, 0.00612376956269145, 0.006710950285196304, 0.0343138724565506], [0.109707772731781, 0.1680740863084793, 0.05170662701129913, 0.04158816486597061, 0.026700180023908615, 0.23248757421970367, 0.5156019330024719, 0.3799504041671753, 0.02909121848642826, 0.009008231572806835, 0.0013055672170594335, 0.0032788640819489956, 0.0791734829545021, 0.010587821714580059, 0.06850002706050873], [0.04004191607236862, 0.02257939800620079, 0.01325287576764822, 0.14834734797477722, 0.0700073167681694, 0.12831416726112366, 0.47980472445487976, 0.3121630549430847, 0.05984592065215111, 0.015101294964551926, 0.002668763743713498, 0.0007187540177255869, 0.04004915803670883, 0.0007627750164829195, 0.05523831769824028], [0.0007188548916019499, 0.006864115130156279, 0.00033292395528405905, 0.000431404507253319, 0.0152564262971282, 0.2775210440158844, 0.03714991733431816, 0.7278205156326294, 0.004819776862859726, 0.00047404138604179025, 0.0003997469611931592, 0.0001266899926122278, 0.0201359074562788, 0.0027800032403320074, 0.042311206459999084], [0.00020999301341362298, 0.0025689874310046434, 3.502765650864603e-07, 6.610702985199168e-05, 0.00024143110204022378, 0.018905406817793846, 0.033397458493709564, 0.4650881290435791, 0.004783111158758402, 0.00013528004637919366, 5.751344360760413e-06, 7.93816871009767e-05, 0.0039043116848915815, 0.0005016719806008041, 0.07914639264345169], [0.00019393693946767598, 0.07456899434328079, 1.429513213224709e-05, 4.6383509470615536e-05, 6.820548151154071e-05, 0.004400796256959438, 0.0021800962276756763, 0.45963534712791443, 0.00143687822856009, 0.0008175616967491806, 6.983020284678787e-05, 3.49152869603131e-05, 0.0030698180198669434, 0.0006545006763190031, 0.001625033444724977], [0.004301158711314201, 0.013502174988389015, 4.788395017385483e-05, 0.00021532995742745697, 7.713190279901028e-05, 0.001439842046238482, 0.005622516851872206, 0.121849425137043, 0.006593172438442707, 0.006624745205044746, 0.0006814572843722999, 0.0002721978526096791, 0.0009267745190300047, 0.0016606011195108294, 0.2357456088066101], [0.0064394231885671616, 0.03409593552350998, 0.0025135872419923544, 0.0008376456098631024, 0.0004409599641803652, 0.0026055865455418825, 0.005634414032101631, 0.014003962278366089, 0.2343187928199768, 0.08099395036697388, 0.23927520215511322, 0.01715606264770031, 0.10332414507865906, 0.021894987672567368, 0.1941189020872116], [0.0004975660121999681, 0.0015548047376796603, 6.826691333117196e-06, 1.0557592986515374e-06, 2.731301538005937e-05, 0.0005447702133096755, 0.00042012380436062813, 0.0503113828599453, 0.0053693996742367744, 0.0012762928381562233, 0.0017790982965379953, 0.019809026271104813, 0.47653263807296753, 0.008869247511029243, 0.017010610550642014], [0.00012974163109902292, 0.005610004533082247, 2.3442629753844813e-05, 1.8520654521125834e-06, 3.9678394387010485e-05, 0.0016583451069891453, 0.00029088594601489604, 0.004530484322458506, 0.0021493860986083746, 0.00029196502873674035, 0.0005848451401107013, 0.0028240433894097805, 0.4590959846973419, 0.22978197038173676, 0.0020738127641379833], [0.00021855060185771435, 0.005491270218044519, 1.9927349057979882e-05, 7.633860150235705e-06, 0.0004071943403687328, 0.008836714550852776, 7.301902951439843e-05, 0.011723233386874199, 1.7278060113312677e-05, 0.0001269245840376243, 0.00022235361393541098, 0.016586007550358772, 0.41012606024742126, 0.37776312232017517, 0.0024871949572116137], [0.02619638666510582, 0.18392468988895416, 0.0003054745029658079, 0.00016413358389399946, 0.0015171386767178774, 0.004799532704055309, 0.004810427315533161, 0.058836404234170914, 0.0003794554795604199, 0.0017285931389778852, 0.000568193441722542, 0.003299211384728551, 0.6178385019302368, 0.5079926252365112, 0.05467592179775238], [0.03445081040263176, 0.14193737506866455, 0.0007241201237775385, 0.0002892682678066194, 0.0003202178922947496, 0.003702279180288315, 0.01134149543941021, 0.12129464000463486, 0.0006569268880411983, 0.0008894759230315685, 8.523569704266265e-05, 0.00030898841214366257, 0.7088924646377563, 0.10790188610553741, 0.05374660715460777], [0.04547691345214844, 0.010678221471607685, 0.0016328264027833939, 0.024403419345617294, 0.012795579619705677, 0.004323439672589302, 0.06414945423603058, 0.014008321799337864, 0.011475995182991028, 0.00871653389185667, 0.012156924232840538, 0.0147528275847435, 0.009472412057220936, 0.0331418551504612, 0.1366012692451477]], [[0.3143080472946167, 0.014564945362508297, 0.07743841409683228, 0.19665417075157166, 0.23130221664905548, 0.03274351730942726, 0.23599109053611755, 0.04763320833444595, 0.20168107748031616, 0.7521476149559021, 0.7922006249427795, 0.840878427028656, 0.6463541388511658, 0.6008138656616211, 0.0070990691892802715], [0.05880431830883026, 0.004086965229362249, 0.06557433307170868, 0.4476080536842346, 0.32179930806159973, 0.2046266496181488, 0.5952353477478027, 0.20483972132205963, 0.7834360599517822, 0.27592822909355164, 0.5900363922119141, 0.6986290812492371, 0.3548848032951355, 0.36629796028137207, 0.07452832907438278], [0.4484235942363739, 0.0712433010339737, 0.09740526974201202, 0.49982836842536926, 0.18807044625282288, 0.007537430617958307, 0.2073078453540802, 0.015238385647535324, 0.18028782308101654, 0.6095888018608093, 0.4225178062915802, 0.6769288778305054, 0.3957397937774658, 0.7102670669555664, 0.05611870437860489], [0.4341801106929779, 0.05481646955013275, 0.17834456264972687, 0.2579769194126129, 0.326920747756958, 0.0030261597130447626, 0.03147314488887787, 0.003279186552390456, 0.09941483289003372, 0.5679370760917664, 0.8480010032653809, 0.8133074045181274, 0.4710683822631836, 0.9189481139183044, 0.04321537911891937], [0.559230387210846, 0.08983521163463593, 0.16111011803150177, 0.14667965471744537, 0.32596829533576965, 0.008685072883963585, 0.1111784353852272, 0.02690659649670124, 0.06770152598619461, 0.18340016901493073, 0.4614297151565552, 0.502476155757904, 0.42325475811958313, 0.5992166996002197, 0.05437220633029938], [0.367906779050827, 0.21432256698608398, 0.3548191487789154, 0.2603428363800049, 0.22096140682697296, 0.0013341127196326852, 0.021726170554757118, 0.005543001927435398, 0.5389296412467957, 0.818263828754425, 0.919593095779419, 0.8187286257743835, 0.4823090434074402, 0.4897681474685669, 0.07018090784549713], [0.7116888761520386, 0.17206020653247833, 0.6874114871025085, 0.19288089871406555, 0.20990870893001556, 0.011273512616753578, 0.2026582807302475, 0.004371582996100187, 0.10976968705654144, 0.4432500898838043, 0.7022042274475098, 0.8704607486724854, 0.721519947052002, 0.7422701716423035, 0.025589054450392723], [0.7674684524536133, 0.20032620429992676, 0.42808812856674194, 0.11714937537908554, 0.32732346653938293, 0.009955272078514099, 0.05444686487317085, 0.0040375906974077225, 0.12078685313463211, 0.6266691088676453, 0.5163981914520264, 0.8307003378868103, 0.32096055150032043, 0.24524804949760437, 0.04717922583222389], [0.7549813389778137, 0.15439504384994507, 0.33331331610679626, 0.24930144846439362, 0.2927357852458954, 0.04936225712299347, 0.44933974742889404, 0.06466211378574371, 0.09519664198160172, 0.08716140687465668, 0.058296240866184235, 0.09990595281124115, 0.5117565989494324, 0.1508449912071228, 0.039490822702646255], [0.654628574848175, 0.3205694854259491, 0.5841068029403687, 0.21299651265144348, 0.365792840719223, 0.0401315838098526, 0.18686936795711517, 0.05883712321519852, 0.05069931596517563, 0.33667507767677307, 0.3354107439517975, 0.22027519345283508, 0.05277648940682411, 0.09031395614147186, 0.015531455166637897], [0.3366456627845764, 0.1530359387397766, 0.41866233944892883, 0.39775165915489197, 0.7769761681556702, 0.06979230791330338, 0.41583842039108276, 0.02130916155874729, 0.14617334306240082, 0.25815388560295105, 0.1423572301864624, 0.18894770741462708, 0.041056301444768906, 0.026175418868660927, 0.03888533264398575], [0.24913249909877777, 0.0818726196885109, 0.5426726341247559, 0.1687711775302887, 0.8305720090866089, 0.26261457800865173, 0.39635857939720154, 0.1712585836648941, 0.1158638522028923, 0.17366157472133636, 0.12521226704120636, 0.5298976302146912, 0.041029125452041626, 0.02415779046714306, 0.1170416921377182], [0.3567614257335663, 0.035316068679094315, 0.3819185495376587, 0.10469090938568115, 0.3454773426055908, 0.09596268832683563, 0.3821227550506592, 0.17425164580345154, 0.40528857707977295, 0.1745157092809677, 0.10956539213657379, 0.5078453421592712, 0.0026470222510397434, 0.016186503693461418, 0.08932095021009445], [0.330766886472702, 0.039845019578933716, 0.6981685757637024, 0.09713104367256165, 0.8411048650741577, 0.16356231272220612, 0.3630223274230957, 0.1627381145954132, 0.6954487562179565, 0.17326875030994415, 0.1752558946609497, 0.24479816854000092, 0.026946308091282845, 0.016200177371501923, 0.06702017039060593], [0.07683827728033066, 0.07034450024366379, 0.21707428991794586, 0.2902449369430542, 0.1834353357553482, 0.01726321130990982, 0.13144701719284058, 0.005189047660678625, 0.150242418050766, 0.1182665303349495, 0.4041094183921814, 0.12062898278236389, 0.05959685891866684, 0.1186181977391243, 0.1283060759305954]], [[0.06827192008495331, 0.0036808219738304615, 0.005701950751245022, 0.005157816223800182, 0.003777393838390708, 0.024757172912359238, 0.0020165019668638706, 0.010267351754009724, 0.013163687661290169, 0.001690453034825623, 0.00837681908160448, 0.00522418599575758, 0.061038240790367126, 0.015438525006175041, 0.325132817029953], [0.7422951459884644, 0.028774140402674675, 0.06394203752279282, 0.00887901522219181, 0.04345611855387688, 0.027670713141560555, 0.0295904241502285, 0.01398912351578474, 0.025535697117447853, 0.02094031311571598, 0.022182827815413475, 0.009663421660661697, 0.049684178084135056, 0.026225639507174492, 0.13834334909915924], [0.20897099375724792, 0.21868035197257996, 0.23815643787384033, 0.005872054491192102, 0.0010661164997145534, 0.0017293300479650497, 0.00042713910806924105, 0.002609806600958109, 0.016046296805143356, 0.009100147522985935, 0.014420107938349247, 0.0022624030243605375, 0.010553905740380287, 0.007111164275556803, 0.25332581996917725], [0.2508500814437866, 0.20390872657299042, 0.7329782247543335, 0.07117453217506409, 0.016424261033535004, 0.021444672718644142, 0.001510130357928574, 0.004098558332771063, 0.0484151765704155, 0.02061472274363041, 0.001126835006289184, 0.0022107160184532404, 0.007578131277114153, 0.004504901356995106, 0.1403624713420868], [0.27370113134384155, 0.8174626231193542, 0.7193068861961365, 0.7076587677001953, 0.07771007716655731, 0.01620337925851345, 0.004001453518867493, 0.004182097036391497, 0.03681829199194908, 0.09453201293945312, 0.026799198240041733, 0.006044679321348667, 0.03725922852754593, 0.016391301527619362, 0.04474738612771034], [0.3889567255973816, 0.4487122893333435, 0.5870586037635803, 0.6609426140785217, 0.6319714188575745, 0.10676700621843338, 0.009257740341126919, 0.0017087672604247928, 0.027955975383520126, 0.07590407133102417, 0.006841681431978941, 0.08621303737163544, 0.05063363164663315, 0.016846608370542526, 0.05719457566738129], [0.00991373136639595, 0.0983041524887085, 0.15667210519313812, 0.19277995824813843, 0.5809133052825928, 0.7996482253074646, 0.06316149979829788, 0.004939877428114414, 0.023352928459644318, 0.010926214046776295, 0.008795071393251419, 0.006998055148869753, 0.0765714943408966, 0.006783204153180122, 0.05886436253786087], [0.07887525111436844, 0.017153050750494003, 0.2216421663761139, 0.13068468868732452, 0.5295770764350891, 0.35302138328552246, 0.8493326902389526, 0.04265422001481056, 0.052519019693136215, 0.027357611805200577, 0.01357424259185791, 0.004279646556824446, 0.026089098304510117, 0.04089489206671715, 0.014124121516942978], [0.03465811163187027, 0.15351061522960663, 0.2825109362602234, 0.08174889534711838, 0.19755861163139343, 0.5825939774513245, 0.37084007263183594, 0.7892780900001526, 0.1287456750869751, 0.006381133571267128, 0.001940184272825718, 0.00047384126810356975, 0.011903955601155758, 0.003972942009568214, 0.06710142642259598], [0.013788340613245964, 0.006632686126977205, 0.02207767777144909, 0.0785517543554306, 0.014113685116171837, 0.048156753182411194, 0.1944313496351242, 0.22155866026878357, 0.49656373262405396, 0.009422117844223976, 0.004702835343778133, 0.0007582302205264568, 0.00014129001647233963, 0.00033574484405107796, 0.23994654417037964], [0.00469209672883153, 0.015491061843931675, 0.035103749483823776, 0.009631682187318802, 0.008573818951845169, 0.051444172859191895, 0.04315423220396042, 0.05495374649763107, 0.6859460473060608, 0.5370080471038818, 0.06784479320049286, 0.004556083586066961, 0.001035997993312776, 0.0006345660076476634, 0.13974453508853912], [0.02668480947613716, 0.016245348379015923, 0.01112398225814104, 0.008507933467626572, 0.02067524567246437, 0.17763113975524902, 0.05662769451737404, 0.04544723033905029, 0.7948054671287537, 0.7384940385818481, 0.5224500298500061, 0.1060851439833641, 0.014122114516794682, 0.0019289307529106736, 0.08371670544147491], [0.02394592948257923, 0.04371663182973862, 0.028385786339640617, 0.007640721742063761, 0.014576996676623821, 0.08887659758329391, 0.017377078533172607, 0.020801657810807228, 0.187345951795578, 0.5047414302825928, 0.6342922449111938, 0.3672487437725067, 0.04719087854027748, 0.10966072231531143, 0.08543073385953903], [0.009629062376916409, 0.020042795687913895, 0.006009343545883894, 0.001406975439749658, 0.0026742229238152504, 0.006072318647056818, 0.006495587062090635, 0.0032924923580139875, 0.034326668828725815, 0.5998041033744812, 0.7456773519515991, 0.7204623818397522, 0.012111457996070385, 0.018825965002179146, 0.008305574767291546], [0.08114123344421387, 0.05478224158287048, 0.11802507936954498, 0.1980995535850525, 0.15338915586471558, 0.11414031684398651, 0.06528255343437195, 0.04494854062795639, 0.26375874876976013, 0.30061599612236023, 0.26960447430610657, 0.5329554677009583, 0.4288364350795746, 0.12292250245809555, 0.12395624816417694]], [[0.09139528125524521, 0.1232069656252861, 0.06926427036523819, 0.03596228361129761, 0.08677947521209717, 0.3523865342140198, 0.17220446467399597, 0.3048216700553894, 0.24129998683929443, 0.008230631239712238, 0.012852879241108894, 0.0024019270204007626, 0.003931952640414238, 0.002576343482360244, 0.13348431885242462], [0.005495021585375071, 0.009821278043091297, 0.006606503389775753, 0.0009270968730561435, 0.022634856402873993, 0.02637101709842682, 0.03666122257709503, 0.003247066168114543, 0.03138025477528572, 0.0023785934317857027, 0.007012520916759968, 0.0027185468934476376, 0.001623710268177092, 0.009003029204905033, 0.24841202795505524], [0.004891206510365009, 0.01856830157339573, 0.01660238206386566, 0.05400720611214638, 0.2678459584712982, 0.21548990905284882, 0.0901486948132515, 0.14165979623794556, 0.4387242794036865, 0.0060303402133286, 0.03774549812078476, 0.022296983748674393, 0.014843892306089401, 0.003844154067337513, 0.0701230987906456], [0.009136357344686985, 0.005524215288460255, 0.002000550739467144, 0.004360574297606945, 0.06230698525905609, 0.032116882503032684, 0.14447683095932007, 0.11250873655080795, 0.12456412613391876, 0.017903752624988556, 0.03641437739133835, 0.030236193910241127, 0.03817100450396538, 0.0020203718449920416, 0.24235397577285767], [0.011458649300038815, 0.0028747334145009518, 0.0048751854337751865, 0.0034302298445254564, 0.032581884413957596, 0.009492963552474976, 0.29646721482276917, 0.024549754336476326, 0.5199102163314819, 0.07497825473546982, 0.039336495101451874, 0.23366358876228333, 0.2855432629585266, 0.0047793262638151646, 0.131587415933609], [0.0048281243070960045, 0.014400148764252663, 0.00021499136346392334, 0.00015902110317256302, 0.0008502291166223586, 0.005816742777824402, 0.03721616789698601, 0.31765323877334595, 0.006985681131482124, 9.90723492577672e-05, 0.0015535155544057488, 0.002471775049343705, 0.00966054666787386, 0.002636645222082734, 0.15553238987922668], [0.01824354939162731, 0.02838711440563202, 0.0006440957658924162, 0.00040316785452887416, 0.00041587575105950236, 0.0021029487252235413, 0.07766012847423553, 0.3384210765361786, 0.005884509067982435, 0.02229108288884163, 0.02292727865278721, 0.00326070049777627, 0.002748187631368637, 0.004811563994735479, 0.08466839045286179], [0.0009052195237018168, 0.00028935770387761295, 0.00010135041520697996, 4.4237076508579776e-05, 9.765469440026209e-05, 0.0003226006228942424, 0.0006174442823976278, 0.003764552064239979, 0.001191335148178041, 0.0005841490346938372, 0.001988127361983061, 0.0019700597040355206, 0.0006354944198392332, 0.0011416736524552107, 0.25631290674209595], [0.007226317655295134, 0.015471585094928741, 0.027516253292560577, 0.0063530029729008675, 0.015222059562802315, 0.004327190574258566, 0.010739101096987724, 0.0023785619996488094, 0.053105201572179794, 0.0674574077129364, 0.31870341300964355, 0.4986713230609894, 0.027042971923947334, 0.0736011192202568, 0.116986483335495], [0.015794623643159866, 0.009404269978404045, 0.017993446439504623, 0.003823975333943963, 0.004969433881342411, 0.03679484874010086, 0.04242165759205818, 0.017222637310624123, 0.1201641708612442, 0.016131659969687462, 0.3518509864807129, 0.3061373829841614, 0.0458594486117363, 0.15943044424057007, 0.17968055605888367], [0.006380036938935518, 0.028477374464273453, 0.006851766724139452, 0.005024573765695095, 0.02579522877931595, 0.052536945790052414, 0.0111169358715415, 0.0038714397232979536, 0.008046599105000496, 0.008921324275434017, 0.011395278386771679, 0.10255969315767288, 0.21638940274715424, 0.44467252492904663, 0.05895284563302994], [0.010142950341105461, 0.001643709372729063, 0.002422438468784094, 0.0009472724632360041, 0.0033483330626040697, 0.003415578044950962, 0.03889569267630577, 0.005287462379783392, 0.00042015319922938943, 0.0010667687747627497, 0.00740370387211442, 0.00895014964044094, 0.0067735291086137295, 0.017782215029001236, 0.26753443479537964], [0.11724554747343063, 0.0023070531897246838, 0.004510094877332449, 0.0014967885799705982, 0.007825762964785099, 0.00018500315491110086, 0.013543304987251759, 0.0012864026939496398, 0.0007778326398693025, 0.00044295378029346466, 0.001640060218051076, 0.0014512997586280107, 0.002360806567594409, 0.2112705558538437, 0.19457924365997314], [0.09882069379091263, 0.014871560037136078, 0.005077258683741093, 0.0014827846316620708, 0.005620975513011217, 0.0024449406191706657, 0.07368315756320953, 0.06950978189706802, 0.0017206794582307339, 0.00039900749106891453, 0.0006052122334949672, 0.0005968212499283254, 0.004762541502714157, 0.0232950821518898, 0.2500154376029968], [0.001020739320665598, 0.001402992638759315, 0.0006185534875839949, 0.0003395593084860593, 0.0013021298218518496, 0.0008022591937333345, 0.003452729433774948, 0.0026675688568502665, 0.0021077031269669533, 0.0008018113439902663, 0.0017594166565686464, 0.0005115982494316995, 0.0007778447470627725, 0.0008368113776668906, 0.13888627290725708]], [[0.04622220993041992, 0.12740419805049896, 0.05372706800699234, 0.5582705140113831, 0.030120277777314186, 0.3703221380710602, 0.020304178819060326, 0.3357560634613037, 0.11819478869438171, 0.0765489861369133, 0.09261158853769302, 0.03858334198594093, 0.13079233467578888, 0.0447748564183712, 0.11706516146659851], [0.0919138491153717, 0.05798470228910446, 0.02827676385641098, 0.34965166449546814, 0.05504997447133064, 0.1526506543159485, 0.09941896051168442, 0.4367760419845581, 0.061004042625427246, 0.5390062928199768, 0.28723591566085815, 0.15840129554271698, 0.2018149495124817, 0.11561664938926697, 0.1249081939458847], [0.032068803906440735, 0.0549696609377861, 0.018587671220302582, 0.2202640324831009, 0.0011182812741026282, 0.03810814768075943, 0.027008401229977608, 0.3763306438922882, 0.11146998405456543, 0.16719762980937958, 0.13283231854438782, 0.014421377331018448, 0.07254088670015335, 0.007401765324175358, 0.20662666857242584], [0.10753453522920609, 0.479284405708313, 0.009764611721038818, 0.0431443527340889, 0.0008862981921993196, 0.03188035264611244, 0.00600279588252306, 0.43093177676200867, 0.08460848033428192, 0.18502341210842133, 0.038902610540390015, 0.030237559229135513, 0.1820157915353775, 0.03367093205451965, 0.14427724480628967], [0.013928310945630074, 0.032752107828855515, 0.0024797581136226654, 0.10617181658744812, 0.0002726189268287271, 0.011333486996591091, 0.005626056343317032, 0.05421115458011627, 0.020341530442237854, 0.0548044852912426, 0.027503041550517082, 0.005752534605562687, 0.033552803099155426, 0.008454940281808376, 0.388910174369812], [0.15046736598014832, 0.296213299036026, 0.044096194207668304, 0.05168119817972183, 0.02727358601987362, 0.04717152938246727, 0.0016543868696317077, 0.035376399755477905, 0.027143586426973343, 0.0870317667722702, 0.05812281742691994, 0.06705813109874725, 0.3147181272506714, 0.39039844274520874, 0.23394177854061127], [0.14644725620746613, 0.5605929493904114, 0.11812092363834381, 0.5902084112167358, 0.021858595311641693, 0.10718227922916412, 0.007383488584309816, 0.019886687397956848, 0.06570647656917572, 0.10820640623569489, 0.1357717514038086, 0.025582531467080116, 0.077891044318676, 0.061965201050043106, 0.164744034409523], [0.049012791365385056, 0.35138410329818726, 0.26388463377952576, 0.7301797866821289, 0.014552393928170204, 0.24720129370689392, 0.0041521950624883175, 0.07795857638120651, 0.014070906676352024, 0.04667593538761139, 0.1480453461408615, 0.010990227572619915, 0.20039354264736176, 0.17517414689064026, 0.0717916414141655], [0.09980960935354233, 0.4834202826023102, 0.20237547159194946, 0.5161312222480774, 0.2011035680770874, 0.31254804134368896, 0.023049525916576385, 0.09284620732069016, 0.030714770779013634, 0.009841320104897022, 0.03625232353806496, 0.02249438874423504, 0.030981028452515602, 0.01249231118708849, 0.19809871912002563], [0.2242409735918045, 0.5898000001907349, 0.2996082305908203, 0.6961580514907837, 0.3950251638889313, 0.824604332447052, 0.0551396869122982, 0.5436567068099976, 0.06683327257633209, 0.03568824753165245, 0.060814060270786285, 0.00592254800722003, 0.012778226286172867, 0.017990900203585625, 0.1082865446805954], [0.03427329286932945, 0.7018846869468689, 0.18350760638713837, 0.5559015274047852, 0.03810380771756172, 0.7226935029029846, 0.05184842646121979, 0.881024181842804, 0.06315085291862488, 0.03384441137313843, 0.014913397841155529, 0.002015632577240467, 0.008405282162129879, 0.0011906703002750874, 0.2768104076385498], [0.022437993437051773, 0.7336767315864563, 0.2893984615802765, 0.7315550446510315, 0.021726222708821297, 0.3247562646865845, 0.05117126554250717, 0.7097986340522766, 0.03149837628006935, 0.017582548782229424, 0.017906883731484413, 0.004864181391894817, 0.0014982494758442044, 0.0005988480988889933, 0.17147301137447357], [0.279982328414917, 0.427709698677063, 0.4798988997936249, 0.811837911605835, 0.5607104301452637, 0.3233453035354614, 0.03364620357751846, 0.48738226294517517, 0.20507316291332245, 0.2806957960128784, 0.20560167729854584, 0.021487781777977943, 0.0051806773990392685, 0.018182942643761635, 0.10378202050924301], [0.15081651508808136, 0.5779510736465454, 0.21354816854000092, 0.8126901984214783, 0.041816346347332, 0.5376638174057007, 0.02729017473757267, 0.45972490310668945, 0.1708957701921463, 0.17148789763450623, 0.06268936395645142, 0.0045938147231936455, 0.0036332160234451294, 0.0009066996863111854, 0.10311751067638397], [0.009540104307234287, 0.03889232128858566, 0.016071060672402382, 0.08366316556930542, 0.004574422258883715, 0.029401082545518875, 0.00834547821432352, 0.0893266350030899, 0.14732055366039276, 0.09065960347652435, 0.14173488318920135, 0.042114999145269394, 0.004022075328975916, 0.003513866104185581, 0.1347859650850296]], [[0.009570755064487457, 0.005546795669943094, 0.006825579330325127, 0.033384330570697784, 0.3769712448120117, 0.15916845202445984, 0.5290282368659973, 0.24695992469787598, 0.2377869039773941, 0.0913546234369278, 0.07570143043994904, 0.06522544473409653, 0.12397455424070358, 0.2645682692527771, 0.1787039041519165], [0.0061562443152070045, 0.040286894887685776, 0.0029807272367179394, 0.016133036464452744, 0.1151214987039566, 0.07519882172346115, 0.10128971189260483, 0.046498823910951614, 0.04111110791563988, 0.11845260113477707, 0.08915312588214874, 0.10556784272193909, 0.16933780908584595, 0.3531811535358429, 0.21578538417816162], [0.14712950587272644, 0.04435151070356369, 0.015454337000846863, 0.01427951455116272, 0.08342041075229645, 0.005383625626564026, 0.10468690097332001, 0.05861024558544159, 0.08666124939918518, 0.15304753184318542, 0.23543620109558105, 0.2374279797077179, 0.10751555860042572, 0.10399115085601807, 0.23440681397914886], [0.0859314426779747, 0.15731151401996613, 0.005385389551520348, 0.04620514437556267, 0.010708490386605263, 0.006711416877806187, 0.012445325031876564, 0.056288186460733414, 0.097142793238163, 0.07020799815654755, 0.02479076385498047, 0.0890590250492096, 0.22972674667835236, 0.034618109464645386, 0.28529092669487], [0.07441635429859161, 0.018118128180503845, 0.016377849504351616, 0.003080169903114438, 0.20936372876167297, 0.0007255859090946615, 0.03578657656908035, 0.00550744216889143, 0.1172742024064064, 0.5684130191802979, 0.3980042636394501, 0.15252694487571716, 0.10817506164312363, 0.23486874997615814, 0.2619861364364624], [0.05188249424099922, 0.0069924332201480865, 0.0009591103880666196, 0.0061192926950752735, 0.002253405749797821, 0.006572761107236147, 0.004667140077799559, 0.11107926070690155, 0.03415685519576073, 0.010113962925970554, 0.006655086297541857, 0.010832482948899269, 0.03651394695043564, 0.040573474019765854, 0.2686486840248108], [0.08095332235097885, 0.02014574408531189, 0.011188640259206295, 0.0037319576367735863, 0.024485761299729347, 0.0018746056593954563, 0.04114176332950592, 0.034570205956697464, 0.009728988632559776, 0.07755846530199051, 0.09898480027914047, 0.0613434873521328, 0.09528356045484543, 0.1511603444814682, 0.2821846306324005], [0.04335615411400795, 0.026033984497189522, 0.03572213277220726, 0.017578190192580223, 0.05956277251243591, 0.01715734601020813, 0.011929154396057129, 0.28936532139778137, 0.0027683174703270197, 0.061091482639312744, 0.23734883964061737, 0.10397756844758987, 0.16337142884731293, 0.37352773547172546, 0.18409839272499084], [0.06077902019023895, 0.031166722998023033, 0.11759120225906372, 0.1409873068332672, 0.24215947091579437, 0.009796793572604656, 0.10265856236219406, 0.01014934666454792, 0.2757207751274109, 0.023714441806077957, 0.038815632462501526, 0.15303847193717957, 0.14991649985313416, 0.6824791431427002, 0.13190437853336334], [0.06505369395017624, 0.006089756730943918, 0.036541152745485306, 0.005829536356031895, 0.20233574509620667, 0.029401954263448715, 0.49993017315864563, 0.030510973185300827, 0.01976127363741398, 0.07993583381175995, 0.017815636470913887, 0.04079095646739006, 0.022992853075265884, 0.6425142288208008, 0.26567763090133667], [0.6054520010948181, 0.07051455229520798, 0.2702813744544983, 0.029061302542686462, 0.13962645828723907, 0.07908772677183151, 0.4563634395599365, 0.02414957620203495, 0.02722080610692501, 0.03215296193957329, 0.015534932725131512, 0.009437407366931438, 0.0218642745167017, 0.08506882190704346, 0.4000338017940521], [0.3943043351173401, 0.11258544027805328, 0.12088752537965775, 0.0732470229268074, 0.030587676912546158, 0.056065596640110016, 0.2533946633338928, 0.04020307958126068, 0.03702285513281822, 0.018525324761867523, 0.009753274731338024, 0.01584538072347641, 0.006842197384685278, 0.013304048217833042, 0.2415902465581894], [0.09087645262479782, 0.0733630359172821, 0.03259122744202614, 0.05433432757854462, 0.028730718418955803, 0.026890264824032784, 0.0992540791630745, 0.042951032519340515, 0.1659460812807083, 0.017093859612941742, 0.006921885069459677, 0.0007972968742251396, 0.010357401333749294, 0.037234287708997726, 0.1852690428495407], [0.2766205668449402, 0.06249983608722687, 0.03302843123674393, 0.08374682813882828, 0.07296875864267349, 0.016804786399006844, 0.2612326145172119, 0.06074067950248718, 0.06402052938938141, 0.021471360698342323, 0.00216249143704772, 0.001582604949362576, 0.0037338242400437593, 0.005314995069056749, 0.23526467382907867], [0.005338736344128847, 0.013486125506460667, 0.016210375353693962, 0.00714905746281147, 0.01115293800830841, 0.008639699779450893, 0.009605110622942448, 0.01017976924777031, 0.008433598093688488, 0.06244685873389244, 0.040223702788352966, 0.009117859415709972, 0.005228321999311447, 0.0028589563444256783, 0.13790398836135864]], [[0.3301994204521179, 0.08890271931886673, 0.08465498685836792, 0.06385943293571472, 0.21852104365825653, 0.02508896216750145, 0.03711355850100517, 0.034155964851379395, 0.1728704422712326, 0.06344152241945267, 0.01567375846207142, 0.047274719923734665, 0.023079151287674904, 0.06240373104810715, 0.17532315850257874], [0.08584976941347122, 0.12593986093997955, 0.03313801810145378, 0.017280908301472664, 0.17652282118797302, 0.268716037273407, 0.12116961926221848, 0.2558431923389435, 0.04765854403376579, 0.04246087744832039, 0.0035840249620378017, 0.02463056705892086, 0.2119264155626297, 0.11800020188093185, 0.14393316209316254], [0.046346988528966904, 0.39951857924461365, 0.5525277853012085, 0.10910754650831223, 0.13167327642440796, 0.030212268233299255, 0.021472660824656487, 0.018023721873760223, 0.1298973113298416, 0.04191790521144867, 0.1535157859325409, 0.04246748238801956, 0.3158371150493622, 0.15602277219295502, 0.1064835637807846], [0.0703379437327385, 0.07535148411989212, 0.05811825022101402, 0.428435742855072, 0.07080380618572235, 0.15123498439788818, 0.3036666214466095, 0.07787945121526718, 0.48052453994750977, 0.12286645174026489, 0.04789941385388374, 0.033336445689201355, 0.030469346791505814, 0.005462532863020897, 0.08732402324676514], [0.0663379579782486, 0.03187985718250275, 0.09551261365413666, 0.0323714055120945, 0.33827176690101624, 0.1471284031867981, 0.3127540946006775, 0.02734280750155449, 0.23260797560214996, 0.02317011170089245, 0.046465177088975906, 0.0992102101445198, 0.09175661206245422, 0.13314616680145264, 0.07444406300783157], [0.034720633178949356, 0.01384154986590147, 0.012703170999884605, 0.020319687202572823, 0.10901976376771927, 0.7807050347328186, 0.03443336486816406, 0.028544975444674492, 0.061822760850191116, 0.00809338316321373, 0.007171421777456999, 0.01342758722603321, 0.09649696201086044, 0.05527613312005997, 0.10404697060585022], [0.030445659533143044, 0.041789710521698, 0.023520270362496376, 0.01782963052392006, 0.16124852001667023, 0.06983006745576859, 0.4703807234764099, 0.01895260065793991, 0.027326058596372604, 0.07994905114173889, 0.026343191042542458, 0.032219063490629196, 0.022085823118686676, 0.031095484271645546, 0.24155765771865845], [0.055046502500772476, 0.3847074508666992, 0.04798666015267372, 0.003912709187716246, 0.06840738654136658, 0.36789029836654663, 0.07226144522428513, 0.4079316258430481, 0.022340288385748863, 0.10408379882574081, 0.07774890959262848, 0.04753485694527626, 0.285355806350708, 0.16128498315811157, 0.02375940792262554], [0.03513112664222717, 0.11586778610944748, 0.03034079447388649, 0.001017131027765572, 0.04634808376431465, 0.03800477832555771, 0.03768199309706688, 0.013300161808729172, 0.14031966030597687, 0.015252463519573212, 0.053176701068878174, 0.06856708973646164, 0.13856393098831177, 0.054046642035245895, 0.2367301732301712], [0.025786809623241425, 0.06564735621213913, 0.039564721286296844, 0.0026341548655182123, 0.016324089840054512, 0.016701271757483482, 0.020613567903637886, 0.0767805427312851, 0.22950275242328644, 0.51694655418396, 0.1544727236032486, 0.1054847463965416, 0.025381706655025482, 0.05480813980102539, 0.1677880734205246], [0.012255452573299408, 0.02410232275724411, 0.08552651852369308, 0.002623841166496277, 0.010307574644684792, 0.0127415731549263, 0.021285703405737877, 0.010095748119056225, 0.06661782413721085, 0.12517453730106354, 0.7383688688278198, 0.19885332882404327, 0.07497892528772354, 0.10072800517082214, 0.06182975694537163], [0.2776626944541931, 0.046990759670734406, 0.032447993755340576, 0.015461347065865993, 0.08414210379123688, 0.04174359515309334, 0.19995476305484772, 0.013662091456353664, 0.019540153443813324, 0.048985805362463, 0.25616249442100525, 0.2484772503376007, 0.1799653023481369, 0.17696446180343628, 0.09890354424715042], [0.05504303798079491, 0.08340897411108017, 0.04799877479672432, 0.017563870176672935, 0.028545444831252098, 0.1704884171485901, 0.030681313946843147, 0.02359093725681305, 0.007767115719616413, 0.019779905676841736, 0.03771185874938965, 0.029841119423508644, 0.28957709670066833, 0.04182300344109535, 0.12634176015853882], [0.06153338775038719, 0.02491314895451069, 0.02542346529662609, 0.0031092099379748106, 0.03241894021630287, 0.1874629557132721, 0.1358277052640915, 0.02619485929608345, 0.017582973465323448, 0.03225348889827728, 0.01329810544848442, 0.026643214747309685, 0.1614912450313568, 0.6035103797912598, 0.09545250982046127], [0.027727488428354263, 0.10283610969781876, 0.02349940501153469, 0.010801603086292744, 0.0136191351339221, 0.1518852412700653, 0.05784522369503975, 0.11107083410024643, 0.10270816832780838, 0.1666017472743988, 0.06030665338039398, 0.06198698654770851, 0.05951831862330437, 0.015173939988017082, 0.1310720145702362]]], [[[0.042950913310050964, 0.0007196685182861984, 0.027302199974656105, 0.006393556483089924, 0.09642192721366882, 0.01637418009340763, 0.0023990001063793898, 0.0024961719755083323, 0.0020593979861587286, 0.0015603104839101434, 0.03318732604384422, 0.35782966017723083, 0.0989728793501854, 0.061845745891332626, 0.203965961933136], [0.10955026745796204, 0.02388770505785942, 0.04351670667529106, 0.023162608966231346, 0.012142845429480076, 0.035775765776634216, 0.03457501530647278, 0.11992064118385315, 0.01240380760282278, 0.007506475783884525, 0.05337386205792427, 0.6535924673080444, 0.5536571145057678, 0.19680790603160858, 0.140446737408638], [0.005947283003479242, 0.0010204642312601209, 0.18009734153747559, 0.006447697523981333, 0.012463629245758057, 7.613956404384226e-05, 7.241032290039584e-05, 0.00011841111700050533, 0.0034185522235929966, 0.0034766956232488155, 0.002135018352419138, 0.005925178527832031, 0.003751354990527034, 0.0019247139571234584, 0.28479355573654175], [0.014483454637229443, 0.022866876795887947, 0.32726621627807617, 0.007662326563149691, 0.09431912004947662, 0.0004296264669392258, 0.0011131323408335447, 0.0014158609556034207, 0.018019702285528183, 0.01865016296505928, 0.0020740600302815437, 0.0029411758296191692, 0.0016890126280486584, 0.0063899424858391285, 0.12852828204631805], [0.030419446527957916, 0.058438073843717575, 0.3924228250980377, 0.035587672144174576, 0.08137891441583633, 0.010925069451332092, 0.001356365391984582, 0.0012006007600575686, 0.053269751369953156, 0.0027948038186877966, 0.04010261595249176, 0.01993635483086109, 0.004820133093744516, 0.004111820366233587, 0.21765674650669098], [0.07767480611801147, 0.006269918289035559, 0.09326869994401932, 0.6196063756942749, 0.11043263971805573, 0.052975643426179886, 0.02037718892097473, 0.0008919782703742385, 0.008360025472939014, 0.002104781800881028, 0.0179440937936306, 0.10498880594968796, 0.011864815838634968, 0.002359954407438636, 0.24602332711219788], [0.00026913435431197286, 8.159392746165395e-05, 0.007915529422461987, 0.05068095400929451, 0.6570689678192139, 0.32081079483032227, 0.05758208408951759, 0.0006442792946472764, 0.0015821922570466995, 6.469202344305813e-05, 0.003034515306353569, 0.0310077928006649, 0.025656316429376602, 0.0025228438898921013, 0.023106882348656654], [0.0005435149651020765, 0.0005490019102580845, 0.034476928412914276, 0.01287262886762619, 0.25229769945144653, 0.4536571502685547, 0.10281822830438614, 0.012222280725836754, 0.016108570620417595, 0.00031008716905489564, 0.0026372161228209734, 0.0034134499728679657, 0.0248859953135252, 0.017225822433829308, 0.02475895546376705], [0.000726195692550391, 0.00036735343746840954, 0.007114858832210302, 0.0026034389156848192, 0.01250846590846777, 0.009484091773629189, 0.0354158952832222, 0.0016834242269396782, 0.19215336441993713, 0.007594457361847162, 0.003938279580324888, 2.8376112823025323e-05, 0.001137340790592134, 0.00011368053674232215, 0.29228782653808594], [0.0005387092242017388, 0.0003453432582318783, 0.015091696754097939, 0.06184916943311691, 0.003162123030051589, 0.014056581072509289, 0.012467358261346817, 0.009164737537503242, 0.05548334866762161, 0.008076494559645653, 0.005971547681838274, 0.001972777536138892, 0.006774900481104851, 0.001264052465558052, 0.2362799048423767], [0.0025044670328497887, 0.0023456772323697805, 0.07385681569576263, 0.006188494618982077, 0.021690815687179565, 0.0007893598522059619, 0.002135526854544878, 0.006048245821148157, 0.25190338492393494, 0.09442908316850662, 0.19532348215579987, 0.031008923426270485, 0.009561427868902683, 0.0021240306086838245, 0.21234139800071716], [0.015501828864216805, 0.0072255814447999, 0.006012998055666685, 0.008203291334211826, 0.0171041339635849, 0.001770812552422285, 0.00655776634812355, 0.002186145167797804, 0.15154685080051422, 0.5713958144187927, 0.05368567630648613, 0.051326390355825424, 0.01612916588783264, 0.0019418209558352828, 0.18746227025985718], [0.05876695737242699, 0.005032649263739586, 0.05515526235103607, 0.012789947912096977, 0.017388533800840378, 0.00580496434122324, 0.015462081879377365, 0.009339934214949608, 0.0222479198127985, 0.03960718587040901, 0.14906688034534454, 0.2817051410675049, 0.14850065112113953, 0.09505022317171097, 0.10619710385799408], [0.012425977736711502, 0.0006452641100622714, 0.00298808584921062, 0.001349467202089727, 0.014642779715359211, 0.0010115096811205149, 0.0033098396379500628, 0.00038259345456026495, 0.0035037249326705933, 0.008293021470308304, 0.03801131248474121, 0.8317341208457947, 0.018821584060788155, 0.057542454451322556, 0.011905365623533726], [0.04682805389165878, 0.01908799074590206, 0.10485747456550598, 0.060083843767642975, 0.15075230598449707, 0.029059063643217087, 0.04093548655509949, 0.03368941321969032, 0.017014725133776665, 0.011203174479305744, 0.0391479916870594, 0.24882012605667114, 0.37940239906311035, 0.12485622614622116, 0.12782400846481323]], [[0.010500228963792324, 0.7224081754684448, 0.030353030189871788, 0.00683749420568347, 0.007232841569930315, 0.018554184585809708, 0.0004432629211805761, 0.02719983458518982, 0.0006519495509564877, 0.0012597806053236127, 0.006804677192121744, 0.0011734187137335539, 0.003679303452372551, 0.010371293872594833, 0.019012004137039185], [0.0004097823693882674, 0.007568135391920805, 0.05432860180735588, 0.08570658415555954, 0.005480978172272444, 0.0009473124518990517, 0.000799189496319741, 0.0012391285272315145, 0.00044785221689380705, 0.0009745006100274622, 0.013956908136606216, 0.00011593959061428905, 0.004404959734529257, 0.0031790253706276417, 0.20507724583148956], [0.022728245705366135, 0.0194535069167614, 0.024020839482545853, 0.023168254643678665, 0.45748311281204224, 0.5855799913406372, 0.21754446625709534, 0.1001717820763588, 0.0221620611846447, 0.0033511894289404154, 0.03508710116147995, 0.20201759040355682, 0.2973189353942871, 0.04947788640856743, 0.0494859553873539], [0.010499863885343075, 0.004784405697137117, 0.0035181313287466764, 0.007238015066832304, 0.4155227243900299, 0.8333501219749451, 0.07475034892559052, 0.20445603132247925, 0.005854693241417408, 0.001852003508247435, 0.02841898612678051, 0.243921160697937, 0.10275343060493469, 0.13816815614700317, 0.07406751066446304], [0.00768234534189105, 0.012151399627327919, 0.0006104251369833946, 0.0018971813842654228, 0.08389636874198914, 0.7291921973228455, 0.2573831081390381, 0.13359335064888, 0.0011000150116160512, 0.0005446228897199035, 0.036390628665685654, 0.06110000237822533, 0.1527252048254013, 0.14593005180358887, 0.05624886974692345], [0.0037335127126425505, 0.004452059045433998, 0.00018280810036230832, 0.016856878995895386, 0.0016014263965189457, 0.05306785926222801, 0.5318921208381653, 0.2889253497123718, 0.0004385874199215323, 0.007465890143066645, 0.0005691659171134233, 0.008836256340146065, 0.00793292187154293, 0.0033322598319500685, 0.1706118881702423], [0.00023320072796195745, 0.0486629419028759, 0.0005405444535426795, 0.005952970590442419, 0.0009982762858271599, 0.004001363180577755, 0.009125707671046257, 0.6945337057113647, 0.006549985148012638, 0.007807720452547073, 0.003924727905541658, 0.004149672109633684, 0.003537258366122842, 0.001676861196756363, 0.11541670560836792], [0.0021667596884071827, 0.0005287157837301493, 0.009149480611085892, 0.024324318394064903, 0.0018866003956645727, 0.0003624066011980176, 0.0004668526817113161, 0.0064473398961126804, 0.0217228215187788, 0.0031395854894071817, 0.0052951243706047535, 0.004629157949239016, 0.003511544084176421, 0.0017145106103271246, 0.2705381214618683], [0.0036477160174399614, 0.018601393327116966, 0.00400471780449152, 0.016223786398768425, 0.015442389994859695, 0.030637366697192192, 0.04816145822405815, 0.009263478219509125, 0.08580432087182999, 0.07024423778057098, 0.17587034404277802, 0.2670482397079468, 0.10741393268108368, 0.11723090708255768, 0.197556272149086], [0.0067135002464056015, 0.005400336813181639, 0.002429268090054393, 0.0005210567032918334, 0.0009090648964047432, 0.056922394782304764, 0.006305574905127287, 0.02051912061870098, 0.009087055921554565, 0.0029723523184657097, 0.5903128385543823, 0.4623943269252777, 0.5148944854736328, 0.10147220641374588, 0.10177940130233765], [0.016283290460705757, 0.004236595239490271, 0.00024049253261182457, 0.00013081195356789976, 0.004825976211577654, 0.03370611369609833, 0.030076656490564346, 0.006495397537946701, 0.015585500746965408, 0.0006116450531408191, 0.009124655276536942, 0.7220618724822998, 0.5160555839538574, 0.16948190331459045, 0.04205150157213211], [0.04056651145219803, 0.05449386313557625, 0.007923644036054611, 0.00034379694261588156, 0.0072999089024960995, 0.005707062315195799, 0.018278487026691437, 0.00924981851130724, 0.0004191468469798565, 0.0015566512010991573, 0.0019580996595323086, 0.06517467647790909, 0.4938390851020813, 0.1360015720129013, 0.14540629088878632], [0.02595147117972374, 0.0358305424451828, 0.021912503987550735, 0.01559682097285986, 0.0029425774700939655, 0.008820675313472748, 0.259022980928421, 0.24083182215690613, 0.0008326273527927697, 0.009937180206179619, 0.008380424231290817, 0.0008840225636959076, 0.11912944912910461, 0.5976794362068176, 0.17433230578899384], [0.024576334282755852, 0.01131413970142603, 0.0036256120074540377, 0.007047882303595543, 0.015460383147001266, 0.007877636700868607, 0.035456594079732895, 0.017273712903261185, 0.0020541276317089796, 0.005268692504614592, 0.003138576401397586, 0.0058868261985480785, 0.09279357641935349, 0.45485755801200867, 0.2460370808839798], [0.02016485668718815, 0.03839857131242752, 0.0345035195350647, 0.005700604524463415, 0.03111962042748928, 0.03698137030005455, 0.056010663509368896, 0.043163470923900604, 0.004449993837624788, 0.000997284660115838, 0.006035848520696163, 0.0027079761493951082, 0.009604639373719692, 0.02099894918501377, 0.13394789397716522]], [[0.11855445802211761, 0.018203705549240112, 0.014699782244861126, 0.005997231230139732, 0.012317956425249577, 0.005482070613652468, 0.020501872524619102, 0.04173066467046738, 0.028033137321472168, 0.007907108403742313, 0.13633504509925842, 0.11779958009719849, 0.02402079664170742, 0.08686818182468414, 0.19919154047966003], [0.015789268538355827, 0.07802969217300415, 0.024552250280976295, 0.007203033193945885, 0.015197299420833588, 0.0086579704657197, 0.005928180180490017, 0.015956610441207886, 0.019966211169958115, 0.002508557867258787, 0.048071712255477905, 0.0452260747551918, 0.027286410331726074, 0.034357864409685135, 0.19209280610084534], [0.7560696601867676, 0.09646204113960266, 0.24264514446258545, 0.03150765225291252, 0.15196740627288818, 0.027980739250779152, 0.025865402072668076, 0.037002913653850555, 0.02429634891450405, 0.014392002485692501, 0.11331582069396973, 0.2883520722389221, 0.24113057553768158, 0.5529852509498596, 0.13967400789260864], [0.6593953371047974, 0.14735713601112366, 0.007992099039256573, 0.03938791900873184, 0.047611087560653687, 0.002478603972122073, 0.00756214139983058, 0.01120123453438282, 0.017771385610103607, 0.011085578240454197, 0.01766165718436241, 0.07185176759958267, 0.01590064913034439, 0.05699647217988968, 0.22524236142635345], [0.8214750289916992, 0.5506035089492798, 0.04117008298635483, 0.00517136137932539, 0.5628769993782043, 0.013714980334043503, 0.018153639510273933, 0.019494647160172462, 0.02796507254242897, 0.003693098435178399, 0.052905939519405365, 0.024033749476075172, 0.017759546637535095, 0.154443621635437, 0.2181331366300583], [0.47579920291900635, 0.4996025860309601, 0.02201933227479458, 0.032786499708890915, 0.003352785250172019, 0.402157723903656, 0.028392860665917397, 0.03425603359937668, 0.017302367836236954, 0.007774383760988712, 0.03628184646368027, 0.015436487272381783, 0.09682580828666687, 0.09163853526115417, 0.1807471215724945], [0.6324970722198486, 0.5132108926773071, 0.14723047614097595, 0.10531618446111679, 0.14770705997943878, 0.01965152472257614, 0.16446776688098907, 0.023718399927020073, 0.014144167304039001, 0.003392518265172839, 0.03989372402429581, 0.048702552914619446, 0.05385157838463783, 0.06003360450267792, 0.2021118402481079], [0.2804942727088928, 0.4447323679924011, 0.40719398856163025, 0.15280602872371674, 0.5485119223594666, 0.006256175693124533, 0.005905789323151112, 0.0894087627530098, 0.014159541577100754, 0.0037697115913033485, 0.08780182898044586, 0.04568948596715927, 0.08344046771526337, 0.08309336006641388, 0.1791403889656067], [0.38668709993362427, 0.3767029941082001, 0.5765653848648071, 0.14457443356513977, 0.830109715461731, 0.558448314666748, 0.2105703204870224, 0.015437009744346142, 0.0802588015794754, 0.0035789015237241983, 0.009509528055787086, 0.011719968169927597, 0.04601259157061577, 0.015442220494151115, 0.02989899180829525], [0.42374563217163086, 0.4557475447654724, 0.5995064973831177, 0.22240440547466278, 0.8298278450965881, 0.26192477345466614, 0.5618261694908142, 0.2755923569202423, 0.03321446478366852, 0.014314521104097366, 0.030895033851265907, 0.0061126528307795525, 0.0033166268840432167, 0.0021476708352565765, 0.12580153346061707], [0.4742293357849121, 0.32335561513900757, 0.5931060910224915, 0.0772920548915863, 0.3757626712322235, 0.211185023188591, 0.42018893361091614, 0.37329575419425964, 0.26276469230651855, 0.012583179399371147, 0.3317490220069885, 0.002885210793465376, 0.011435287073254585, 0.00757939275354147, 0.1435183733701706], [0.21439705789089203, 0.17853425443172455, 0.32548797130584717, 0.06489395350217819, 0.64824378490448, 0.1159982681274414, 0.19616922736167908, 0.27417391538619995, 0.6047332286834717, 0.1810707151889801, 0.034782104194164276, 0.10310898721218109, 0.0316632017493248, 0.025309519842267036, 0.09833981841802597], [0.19860051572322845, 0.10174965113401413, 0.08606765419244766, 0.053267233073711395, 0.11251617968082428, 0.2378872036933899, 0.16651752591133118, 0.1490997076034546, 0.4605393707752228, 0.18029887974262238, 0.1883857697248459, 0.007075145840644836, 0.25310245156288147, 0.08171047270298004, 0.15088772773742676], [0.2976968586444855, 0.21286718547344208, 0.04716610535979271, 0.025928588584065437, 0.1317281424999237, 0.12927810847759247, 0.2939497232437134, 0.23276808857917786, 0.5986261367797852, 0.05386120826005936, 0.05668044835329056, 0.025143466889858246, 0.007965278811752796, 0.03647890314459801, 0.16275253891944885], [0.34472423791885376, 0.33325105905532837, 0.5841152667999268, 0.8456752300262451, 0.4377557933330536, 0.4159393310546875, 0.33224907517433167, 0.1488359123468399, 0.2203720510005951, 0.7425854206085205, 0.7086009383201599, 0.5293036699295044, 0.2777566909790039, 0.22530661523342133, 0.09936152398586273]], [[0.3582096993923187, 0.12323450297117233, 0.41414904594421387, 0.12697191536426544, 0.2567327618598938, 0.12921607494354248, 0.303745299577713, 0.26060354709625244, 0.2067556530237198, 0.0739586353302002, 0.038356974720954895, 0.018690073862671852, 0.019858568906784058, 0.03828525170683861, 0.09448481351137161], [0.034560851752758026, 0.06147807836532593, 0.09719342738389969, 0.03090484067797661, 0.05040246620774269, 0.10769589245319366, 0.28225648403167725, 0.03959896042943001, 0.04561477154493332, 0.015998149290680885, 0.010396423749625683, 0.0027313604950904846, 0.02088637463748455, 0.02540828473865986, 0.1729334592819214], [0.031599532812833786, 0.03154325857758522, 0.01938430592417717, 0.10300880670547485, 0.07719798386096954, 0.3211115002632141, 0.5488157868385315, 0.6110779047012329, 0.03511836752295494, 0.03874386474490166, 0.02549627609550953, 0.08684590458869934, 0.1071673184633255, 0.10855282843112946, 0.09071482717990875], [0.05947110056877136, 0.046990834176540375, 0.001917339744977653, 0.019972380250692368, 0.14856000244617462, 0.10937333106994629, 0.7613639235496521, 0.43800127506256104, 0.038890283554792404, 0.0702563002705574, 0.052807219326496124, 0.20175476372241974, 0.09827514737844467, 0.19838720560073853, 0.1799801141023636], [0.010548654943704605, 0.056933727115392685, 0.0004277318366803229, 0.0005220972234383225, 0.03427216783165932, 0.15697234869003296, 0.44382861256599426, 0.28639304637908936, 0.1278306096792221, 0.0589531809091568, 0.07240739464759827, 0.21584689617156982, 0.623681902885437, 0.39177897572517395, 0.053747572004795074], [0.012333033606410027, 0.11936485022306442, 0.0015480549773201346, 0.05167163908481598, 0.003915506415069103, 0.05033823475241661, 0.18770258128643036, 0.5247471332550049, 0.13492631912231445, 0.0999734029173851, 0.02801361307501793, 0.04943297058343887, 0.067798912525177, 0.02220618724822998, 0.04863249137997627], [0.023225123062729836, 0.03936318680644035, 0.0654693990945816, 0.0780135840177536, 0.03190883249044418, 0.007237496320158243, 0.3230750560760498, 0.11266676336526871, 0.3152024447917938, 0.12503208220005035, 0.08215073496103287, 0.20814812183380127, 0.054794978350400925, 0.014369799755513668, 0.31165388226509094], [0.021642545238137245, 0.05032852664589882, 0.10916808992624283, 0.14173567295074463, 0.025796422734856606, 0.002176823327317834, 0.004212724044919014, 0.11230720579624176, 0.2761599123477936, 0.18545517325401306, 0.30032697319984436, 0.18456220626831055, 0.1202857494354248, 0.02383211813867092, 0.22383396327495575], [0.014165909960865974, 0.030938388779759407, 0.019327908754348755, 0.025021186098456383, 0.018685894086956978, 0.058899857103824615, 0.05705944076180458, 0.013411193154752254, 0.27564239501953125, 0.14192135632038116, 0.4484158754348755, 0.49174171686172485, 0.42328834533691406, 0.5148258805274963, 0.024227913469076157], [0.030343737453222275, 0.035576362162828445, 0.011198173277080059, 0.0029289661906659603, 0.004656192846596241, 0.19044476747512817, 0.14425727725028992, 0.14593322575092316, 0.02429576776921749, 0.03922351822257042, 0.03158531337976456, 0.3954472541809082, 0.18761666119098663, 0.829915463924408, 0.05755764618515968], [0.07378673553466797, 0.08269044756889343, 0.008506381884217262, 0.004565858747810125, 0.0033621611073613167, 0.47163471579551697, 0.3437289595603943, 0.16293375194072723, 0.0103234788402915, 0.006828381214290857, 0.025515833869576454, 0.13491219282150269, 0.23380780220031738, 0.7675665616989136, 0.06853343546390533], [0.19539110362529755, 0.20751968026161194, 0.012997383251786232, 0.004634191282093525, 0.004486567340791225, 0.10301963984966278, 0.2361651211977005, 0.10510270297527313, 0.007245894055813551, 0.02498149685561657, 0.005201807711273432, 0.12586773931980133, 0.2985144853591919, 0.741521954536438, 0.061252206563949585], [0.3654796779155731, 0.656768798828125, 0.02389511466026306, 0.057929087430238724, 0.025417884811758995, 0.2985052168369293, 0.29244741797447205, 0.15614598989486694, 0.02199239283800125, 0.027919312939047813, 0.024499662220478058, 0.0015409317566081882, 0.18344998359680176, 0.05587974563241005, 0.11099682748317719], [0.24996283650398254, 0.30432745814323425, 0.08651068061590195, 0.27794384956359863, 0.10948572307825089, 0.32318809628486633, 0.40224379301071167, 0.24700750410556793, 0.016620514914393425, 0.03902489319443703, 0.01563531532883644, 0.008603462018072605, 0.029363060370087624, 0.20380347967147827, 0.1635625809431076], [0.08184575289487839, 0.05559774115681648, 0.012900986708700657, 0.004766350146383047, 0.02465618960559368, 0.0658264234662056, 0.16982027888298035, 0.09995799511671066, 0.1946410834789276, 0.03345171734690666, 0.026332948356866837, 0.010880211368203163, 0.01684177853167057, 0.011932285502552986, 0.13059602677822113]], [[0.06378140300512314, 0.013955923728644848, 0.058693334460258484, 0.014864355325698853, 0.02882157638669014, 0.02533077634871006, 0.013877282850444317, 0.02919653430581093, 0.029733512550592422, 0.010929838754236698, 0.2184230536222458, 0.404588907957077, 0.5044611692428589, 0.4171900451183319, 0.18600669503211975], [0.09787620604038239, 0.3741878271102905, 0.1718531847000122, 0.22170154750347137, 0.11211875081062317, 0.06884550303220749, 0.023903023451566696, 0.00765330670401454, 0.043831951916217804, 0.04742401838302612, 0.08705892413854599, 0.19904442131519318, 0.1439688503742218, 0.08975595235824585, 0.124632827937603], [0.024405136704444885, 0.006321595516055822, 0.03571266308426857, 0.0050111510790884495, 0.01807553507387638, 6.11300565651618e-05, 0.0022184934932738543, 0.002461126074194908, 0.00987271312624216, 0.03944821655750275, 0.02587837167084217, 0.009154303930699825, 0.018459370359778404, 0.07083768397569656, 0.2838045060634613], [0.02829434722661972, 0.05303699150681496, 0.03342747688293457, 0.026768406853079796, 0.06776657700538635, 0.0015663451049476862, 0.0066550131887197495, 0.028257621452212334, 0.02201445959508419, 0.024995435029268265, 0.014314326457679272, 0.019762825220823288, 0.019060753285884857, 0.09995586425065994, 0.2721303105354309], [0.011709636077284813, 0.13082386553287506, 0.3091292977333069, 0.012390679679811, 0.06598176062107086, 0.0025066242087632418, 0.008877930231392384, 0.03396160528063774, 0.01681593246757984, 0.01466491911560297, 0.12272557616233826, 0.010357965715229511, 0.009066522121429443, 0.12291242927312851, 0.3062548041343689], [0.05738264322280884, 0.12342102825641632, 0.7862259149551392, 0.20355252921581268, 0.007363088894635439, 0.0717976987361908, 0.032159313559532166, 0.018495721742510796, 0.0034321516286581755, 0.0013732254737988114, 0.006710591726005077, 0.0023603499867022038, 0.007563347462564707, 0.05948156490921974, 0.12037239223718643], [0.015277753584086895, 0.006394209805876017, 0.6686000227928162, 0.29117655754089355, 0.06745831668376923, 0.2462725043296814, 0.06154515966773033, 0.015117062255740166, 0.004134421236813068, 0.0023558081593364477, 0.08952713012695312, 0.04650713875889778, 0.023702487349510193, 0.01321239210665226, 0.09701406955718994], [0.028385812416672707, 0.012191490270197392, 0.27066752314567566, 0.18411272764205933, 0.040896836668252945, 0.48173367977142334, 0.02650352008640766, 0.07071101665496826, 0.007758310064673424, 0.001958101289346814, 0.01839292421936989, 0.023066602647304535, 0.03435399383306503, 0.03657263144850731, 0.029525745660066605], [0.04876675456762314, 0.422792911529541, 0.22041767835617065, 0.2559551000595093, 0.08884847164154053, 0.01230597123503685, 0.025672338902950287, 0.003895203350111842, 0.022659877315163612, 0.0043840305879712105, 0.007982935756444931, 0.010924039408564568, 0.06971067935228348, 0.0061518345028162, 0.21563398838043213], [0.015657104551792145, 0.02366352081298828, 0.07373688369989395, 0.10379613190889359, 0.013535204343497753, 0.07323776930570602, 0.048540983349084854, 0.008235346525907516, 0.01638718694448471, 0.012322558090090752, 0.073370561003685, 0.03809332847595215, 0.021602218970656395, 0.003090204205363989, 0.23272792994976044], [0.018198516219854355, 0.011175387538969517, 0.02189311571419239, 0.012938260100781918, 0.09454065561294556, 0.010837653651833534, 0.04214898869395256, 0.03231353685259819, 0.2788335978984833, 0.02807164192199707, 0.0381515808403492, 0.013884211890399456, 0.014051362872123718, 0.00934662390500307, 0.24102351069450378], [0.01114112138748169, 0.11382883787155151, 0.017900465056300163, 0.008639826439321041, 0.024639632552862167, 0.020821422338485718, 0.022935912013053894, 0.04321465268731117, 0.055257730185985565, 0.0561254657804966, 0.006350866984575987, 0.034159135073423386, 0.001170721254311502, 0.00040716465446166694, 0.2438717484474182], [0.01806582696735859, 0.014762195758521557, 0.02654433250427246, 0.025726040825247765, 0.03240499645471573, 0.020733002573251724, 0.04244884103536606, 0.02047092467546463, 0.13412125408649445, 0.512605607509613, 0.5156171321868896, 0.023306455463171005, 0.0489252470433712, 0.06594526767730713, 0.173824280500412], [0.018763704225420952, 0.010509289801120758, 0.06387435644865036, 0.02487548068165779, 0.10975509881973267, 0.01984621025621891, 0.06460897624492645, 0.03137337416410446, 0.1802622228860855, 0.7354047894477844, 0.7864400148391724, 0.1003832221031189, 0.007522855885326862, 0.14785504341125488, 0.08187610656023026], [0.02117479033768177, 0.061044495552778244, 0.02157888375222683, 0.021421663463115692, 0.04618487507104874, 0.05167240649461746, 0.01054168026894331, 0.009977741166949272, 0.0295058935880661, 0.008349624462425709, 0.02268156036734581, 0.026699911803007126, 0.020697196945548058, 0.013632250018417835, 0.13365623354911804]], [[4.754594192490913e-05, 2.1380438752771624e-08, 2.918067565360616e-08, 2.8621201408896013e-08, 2.499384379461844e-07, 0.0002631827082950622, 5.21495513439163e-10, 2.490414274802788e-08, 1.4592379216082918e-07, 4.660217989282955e-09, 1.3478041793746343e-08, 1.530838318331007e-07, 4.6195887989597395e-05, 8.429636181972455e-06, 0.2157532423734665], [0.6645432114601135, 0.00044607618474401534, 8.70102576300269e-06, 1.056492124007491e-06, 4.43653931370136e-07, 3.5252294310339494e-06, 0.013106754049658775, 0.0008970960625447333, 5.719662112824153e-07, 3.2791810156140855e-08, 1.0544068729245737e-08, 3.57371057191358e-08, 0.00012361648259684443, 0.0008665899513289332, 0.00011794524471042678], [5.6636022236489225e-06, 0.771808385848999, 0.2603715658187866, 7.618767995154485e-05, 2.6443340175319463e-05, 1.448297037853763e-08, 1.7459943213449236e-10, 0.0005545829189941287, 1.3129211993145873e-06, 0.0003596498572733253, 1.3187416243454209e-06, 1.2532552773336647e-08, 5.7067543821176514e-05, 1.4676837054139469e-05, 8.822963764032465e-07], [7.866851170490463e-09, 0.0015575109282508492, 0.5911858677864075, 0.005255529191344976, 0.00012560673349071294, 1.2381517144888221e-08, 1.3975322635251253e-12, 4.631081083061872e-06, 1.8297629367225454e-06, 0.043241821229457855, 0.00025465109501965344, 1.6550380621538352e-07, 1.5873881693551084e-06, 1.3629888329091955e-08, 2.2046858560997862e-08], [1.6020940130090366e-10, 3.2446525892737554e-06, 0.1964423805475235, 0.9067507982254028, 4.244087540428154e-05, 3.027215825568419e-05, 6.154020626425449e-10, 3.570748958736658e-07, 2.493328743469192e-08, 1.327106815551815e-07, 5.116170723340474e-05, 7.67620722541551e-09, 6.538175512105227e-07, 1.6885725528936746e-07, 1.9495971503857845e-09], [4.057985947270026e-09, 1.6926858803500977e-09, 0.00014235911658033729, 0.0026504932902753353, 0.8634750843048096, 1.9555229300749488e-05, 1.294085109293519e-06, 2.6649362894204387e-07, 3.0507638082433175e-10, 5.069419550807197e-09, 1.108148239836737e-07, 1.7377595213474706e-05, 9.726352800498717e-06, 1.823265733946755e-06, 5.869507617717318e-07], [1.9094309466893833e-12, 2.4682887027685507e-13, 6.382604444965523e-10, 6.302604549368596e-10, 1.4692274817207363e-05, 0.3734012544155121, 3.483030241113738e-06, 1.1820202594492457e-08, 1.9522692351614523e-09, 1.394072303342181e-13, 1.7670450172535546e-11, 1.716609077107023e-09, 3.7749509829154704e-06, 2.593782255644328e-06, 3.855710133393586e-07], [8.508453674949124e-08, 1.863478038544031e-09, 1.257351167627263e-10, 5.331373190142763e-11, 3.337832410466035e-08, 1.777973557182122e-05, 0.8244234323501587, 8.755041926633567e-05, 1.7572835409040977e-09, 1.3142270258170718e-11, 7.735358035533546e-13, 4.927841815161038e-11, 5.296478775562719e-07, 0.000259329448454082, 1.8429471282388477e-08], [1.2582735964272729e-09, 2.3675827378610848e-06, 5.770066309196409e-07, 5.0431950282536775e-11, 2.6034334410507398e-11, 1.7287857190240175e-07, 9.084228622668888e-06, 0.8877476453781128, 0.0008898449596017599, 7.2106473680833e-08, 1.9634756043274137e-08, 4.930736808433922e-13, 3.217972377456135e-08, 1.2906410120194778e-05, 9.568290160189008e-09], [2.8039692789860737e-09, 1.3000158105569426e-06, 4.493769978353157e-08, 2.493898698663344e-10, 7.932443764346875e-12, 1.7288407150317653e-08, 2.642636942606913e-10, 3.576151357265189e-05, 0.8324669599533081, 5.240505197434686e-05, 8.11301958947297e-07, 9.422521651814009e-10, 4.6924657937097436e-08, 2.8963553333483105e-08, 6.33739318800508e-08], [2.873091320410026e-09, 7.32139524188824e-05, 1.393846559949452e-05, 2.2707215663331226e-08, 3.602095333121724e-08, 7.893682235637911e-12, 1.2799745258921386e-13, 1.2971109697446082e-07, 4.534097752184607e-05, 0.7187873721122742, 0.0028858170844614506, 4.860597982769832e-06, 3.316463335067965e-06, 6.64895694058032e-08, 4.189383506769673e-09], [3.5802516507033033e-10, 3.3775189312024168e-09, 1.689890041234321e-06, 2.72409181434341e-07, 2.3650377656281307e-08, 3.1582386705863996e-10, 4.773196676235644e-14, 6.179980832632381e-11, 1.0790042637154329e-07, 0.00019566719129215926, 0.8666706681251526, 0.00033315850305370986, 7.101260734998505e-07, 3.226231015673875e-08, 6.780910499770698e-09], [7.800644574729176e-09, 1.700809604265885e-09, 9.215954577257435e-08, 4.046364665555302e-07, 0.00011374137102393433, 5.132134901941754e-06, 5.991689921991394e-10, 9.107053305923429e-11, 5.105777606262407e-11, 3.3974476565390432e-09, 3.904122058884241e-05, 0.65162193775177, 0.00035754009149968624, 6.446759653044865e-05, 8.575011065659055e-07], [5.410449865905775e-10, 1.9016622998524468e-10, 1.651180719930423e-10, 9.184660809680167e-10, 4.749936000081334e-09, 6.8993631430203095e-06, 9.186856830822876e-10, 1.2120262259107673e-11, 1.0679299241797557e-12, 7.136916383397585e-13, 1.9098522763272285e-10, 9.612936082703527e-06, 0.7662882208824158, 0.00778515450656414, 3.0943773765557125e-08], [0.0058370670303702354, 0.00017831011791713536, 6.727457275701454e-06, 4.542615897662472e-06, 0.0008248149533756077, 0.04996809363365173, 0.010534689761698246, 8.931134652812034e-05, 2.4081384708551923e-07, 6.080232139993313e-08, 3.077615701840841e-06, 0.00041306819184683263, 0.062034472823143005, 0.37576472759246826, 0.1323644071817398]], [[0.278582364320755, 0.012074317783117294, 0.4035726487636566, 0.05818924307823181, 0.5308449864387512, 0.7759386301040649, 0.6032847166061401, 0.04120228812098503, 0.6623223423957825, 0.4034832715988159, 0.2541539669036865, 0.023309720680117607, 0.054716046899557114, 0.3570294678211212, 0.004749305546283722], [0.03977029398083687, 0.025161603465676308, 0.4579423666000366, 0.3708552420139313, 0.767479419708252, 0.5835962295532227, 0.5609359741210938, 0.14304085075855255, 0.8166816234588623, 0.848468542098999, 0.5771627426147461, 0.07112090289592743, 0.12416274100542068, 0.618628740310669, 0.06885465234518051], [0.004083612468093634, 0.0006101519684307277, 0.12011494487524033, 0.04229450225830078, 0.17203551530838013, 0.013333754613995552, 0.01874622330069542, 0.021773431450128555, 0.8914079666137695, 0.25239333510398865, 0.2674473226070404, 0.0986163467168808, 0.10968483239412308, 0.05420238524675369, 0.020816486328840256], [0.00974054355174303, 0.009372939355671406, 0.016473596915602684, 0.12944141030311584, 0.06805374473333359, 0.019993484020233154, 0.038472987711429596, 0.21791628003120422, 0.8550615310668945, 0.2646826505661011, 0.7350810766220093, 0.17277619242668152, 0.36265626549720764, 0.3741258382797241, 0.06228891760110855], [0.0007183643756434321, 0.0016902177594602108, 0.0015671673463657498, 0.000663107552099973, 0.015286565758287907, 0.000776923552621156, 0.007700319401919842, 0.11482121050357819, 0.7658083438873291, 0.5443719625473022, 0.22170989215373993, 0.027013972401618958, 0.025342080742120743, 0.049981117248535156, 0.0074298488907516], [0.011776593513786793, 0.00668947771191597, 0.05204532667994499, 0.026732588186860085, 0.007738037500530481, 0.19347773492336273, 0.08661007881164551, 0.02065080776810646, 0.8265263438224792, 0.77967369556427, 0.8155033588409424, 0.7568296194076538, 0.6889008283615112, 0.7797287106513977, 0.04647013917565346], [0.03701920434832573, 0.011276619508862495, 0.026248518377542496, 0.01771446317434311, 0.046063318848609924, 0.020064320415258408, 0.23005641996860504, 0.032302577048540115, 0.6365551948547363, 0.6746889352798462, 0.6497765183448792, 0.5260909199714661, 0.6955898404121399, 0.8770567178726196, 0.04424796253442764], [0.3583561182022095, 0.034818924963474274, 0.1010005921125412, 0.08171684294939041, 0.0902533084154129, 0.0273053590208292, 0.029195906594395638, 0.10516665875911713, 0.5163984894752502, 0.7107389569282532, 0.5390304327011108, 0.6552954316139221, 0.648922324180603, 0.8148984909057617, 0.13771982491016388], [0.04790134355425835, 0.016352321952581406, 0.004838719964027405, 0.039540428668260574, 0.004614146891981363, 0.10033231228590012, 0.05411757901310921, 0.012187371961772442, 0.25466611981391907, 0.4822390675544739, 0.22996564209461212, 0.2013523131608963, 0.3018202781677246, 0.325538694858551, 0.10763657093048096], [0.18817435204982758, 0.007200991734862328, 0.0915139690041542, 0.00800582580268383, 0.007660675328224897, 0.27090781927108765, 0.08786749839782715, 0.014442713931202888, 0.017244037240743637, 0.8212726712226868, 0.22018176317214966, 0.05063365772366524, 0.16457810997962952, 0.059498634189367294, 0.11578860878944397], [0.1423795521259308, 0.008703344501554966, 0.2208349108695984, 0.02527845837175846, 0.027401143684983253, 0.09980836510658264, 0.024800043553113937, 0.009310302324593067, 0.11915526539087296, 0.048824433237314224, 0.23738479614257812, 0.04641610383987427, 0.11649724096059799, 0.03864651918411255, 0.200869619846344], [0.19247660040855408, 0.028833042830228806, 0.1872357279062271, 0.03232081979513168, 0.031028537079691887, 0.3644941747188568, 0.11239293217658997, 0.0803447812795639, 0.13423573970794678, 0.07468846440315247, 0.009079186245799065, 0.19545331597328186, 0.09625646471977234, 0.07526607811450958, 0.1802312582731247], [0.1263553649187088, 0.009648445062339306, 0.47829046845436096, 0.22347994148731232, 0.2749265432357788, 0.23197446763515472, 0.05249631777405739, 0.01617230661213398, 0.3326357305049896, 0.1497221142053604, 0.04782721772789955, 0.011572148650884628, 0.1354474574327469, 0.0791783407330513, 0.15636207163333893], [0.166306734085083, 0.04561271890997887, 0.48400574922561646, 0.31743937730789185, 0.4171416163444519, 0.1806352734565735, 0.04328177124261856, 0.022486848756670952, 0.1779668778181076, 0.03957689553499222, 0.009708160534501076, 0.01422630064189434, 0.013467496261000633, 0.06257133930921555, 0.22838094830513], [0.39438390731811523, 0.20185884833335876, 0.19486168026924133, 0.053202297538518906, 0.29429352283477783, 0.31667405366897583, 0.3313867747783661, 0.37864530086517334, 0.4971301257610321, 0.178373321890831, 0.16689708828926086, 0.16029801964759827, 0.22925321757793427, 0.22496484220027924, 0.11296840012073517]], [[0.12737327814102173, 0.10940374433994293, 0.05123003572225571, 0.7807462215423584, 0.0676276683807373, 0.02884089946746826, 0.05574861168861389, 0.5975708961486816, 0.07044392824172974, 0.5009010434150696, 0.31273892521858215, 0.07660850137472153, 0.29424503445625305, 0.028401609510183334, 0.07683643698692322], [0.03750006482005119, 0.429240882396698, 0.15060469508171082, 0.2604650557041168, 0.037177786231040955, 0.1944778561592102, 0.07849539071321487, 0.6716934442520142, 0.06105323135852814, 0.07711976766586304, 0.20997941493988037, 0.028168758377432823, 0.12550987303256989, 0.030995607376098633, 0.0958443135023117], [0.15516091883182526, 0.07278051972389221, 0.11765316128730774, 0.7884857058525085, 0.11075033247470856, 0.051856692880392075, 0.18673725426197052, 0.2268398553133011, 0.013722711242735386, 0.6478350162506104, 0.5306386947631836, 0.3090885877609253, 0.22243055701255798, 0.16200464963912964, 0.13070979714393616], [0.21811531484127045, 0.7140333652496338, 0.018219277262687683, 0.764274001121521, 0.15804116427898407, 0.03280843421816826, 0.11008237302303314, 0.09874711185693741, 0.0423860140144825, 0.5652360320091248, 0.14938808977603912, 0.2869919240474701, 0.39966318011283875, 0.1259765923023224, 0.0577625073492527], [0.11744663864374161, 0.1893559694290161, 0.05823011323809624, 0.03701714053750038, 0.15626470744609833, 0.08588159829378128, 0.26269999146461487, 0.41053518652915955, 0.007210245821624994, 0.3749772906303406, 0.4537068009376526, 0.6417111158370972, 0.1666039228439331, 0.13084180653095245, 0.14052902162075043], [0.3613002598285675, 0.240200012922287, 0.044567547738552094, 0.04614294692873955, 0.0021214759908616543, 0.17616558074951172, 0.11286458373069763, 0.11203286051750183, 0.009014172479510307, 0.10163455456495285, 0.0949772298336029, 0.06209810823202133, 0.11910365521907806, 0.04125094786286354, 0.1871420443058014], [0.2914785146713257, 0.381010502576828, 0.08399549126625061, 0.4511452913284302, 0.048780620098114014, 0.008560722693800926, 0.1541443020105362, 0.12101723253726959, 0.02183164842426777, 0.18665823340415955, 0.13169258832931519, 0.13539372384548187, 0.14286382496356964, 0.031125182285904884, 0.2064482420682907], [0.3084108829498291, 0.4568510055541992, 0.068343386054039, 0.40243175625801086, 0.04035715013742447, 0.028490515425801277, 0.006473515648394823, 0.6036491990089417, 0.14769236743450165, 0.09462843090295792, 0.04651549458503723, 0.08334364742040634, 0.08459941297769547, 0.022403797134757042, 0.13448290526866913], [0.4981050491333008, 0.13424238562583923, 0.16773013770580292, 0.5160816311836243, 0.029790958389639854, 0.22989192605018616, 0.568993866443634, 0.056374672800302505, 0.08792523294687271, 0.2900378406047821, 0.12431738525629044, 0.017185388132929802, 0.05061684548854828, 0.020683959126472473, 0.13275840878486633], [0.33482691645622253, 0.4720645546913147, 0.20652346312999725, 0.6004944443702698, 0.1402488797903061, 0.13250590860843658, 0.13873517513275146, 0.5260767936706543, 0.01182119082659483, 0.1017654612660408, 0.047682080417871475, 0.04534589499235153, 0.10121697187423706, 0.0026118881069123745, 0.13006491959095], [0.27261805534362793, 0.5674196481704712, 0.08154824376106262, 0.8736060261726379, 0.4724165201187134, 0.1720387041568756, 0.13692085444927216, 0.40960294008255005, 0.06138879805803299, 0.0898643285036087, 0.15986473858356476, 0.04882661625742912, 0.09858791530132294, 0.005254920106381178, 0.09166211634874344], [0.33052578568458557, 0.40956470370292664, 0.44244009256362915, 0.8809638619422913, 0.26719745993614197, 0.38818857073783875, 0.40750059485435486, 0.4857279658317566, 0.04656125605106354, 0.08998580276966095, 0.02227160707116127, 0.42457664012908936, 0.06242617964744568, 0.019552020356059074, 0.08343644440174103], [0.20678018033504486, 0.17620769143104553, 0.3081345558166504, 0.6112105250358582, 0.534289538860321, 0.19626931846141815, 0.17160479724407196, 0.4079393148422241, 0.027630727738142014, 0.07990976423025131, 0.0661839172244072, 0.022294294089078903, 0.11108729988336563, 0.024492109194397926, 0.12739884853363037], [0.2302674651145935, 0.4147239625453949, 0.3118293881416321, 0.3454154133796692, 0.20178626477718353, 0.3381562829017639, 0.1571493148803711, 0.4487079083919525, 0.02096635475754738, 0.11857040971517563, 0.09038619697093964, 0.01401298213750124, 0.06377796083688736, 0.029106009751558304, 0.10548537224531174], [0.0850413590669632, 0.2905830442905426, 0.047175440937280655, 0.009145522490143776, 0.014412813819944859, 0.03387918695807457, 0.04852135106921196, 0.2856408655643463, 0.03688584640622139, 0.02503933012485504, 0.030300520360469818, 0.020876996219158173, 0.004409631714224815, 0.0025441893376410007, 0.1292814165353775]]], [[[0.00039591442327946424, 4.3682277464540675e-05, 1.7448855942348018e-05, 4.859234650211874e-06, 1.1413659422032651e-06, 1.0625568393152207e-05, 1.9137923246148603e-08, 5.615326585939329e-07, 5.487099315359956e-06, 2.1910665282121045e-07, 2.532970881929941e-07, 7.501878940274764e-07, 1.657212578720646e-06, 1.0862070212169783e-06, 0.18717002868652344], [0.6005652546882629, 0.09179380536079407, 0.017407523468136787, 0.009556752629578114, 0.001977206440642476, 0.02417689561843872, 0.001285116421058774, 0.0015866898465901613, 0.0007265046588145196, 0.0008927723974920809, 0.008914382196962833, 0.0016361800953745842, 0.1313493698835373, 0.006872364319860935, 0.052507203072309494], [0.00456381356343627, 0.8302816152572632, 0.11558636277914047, 0.010320104658603668, 0.00024428890901617706, 9.749805758474395e-05, 7.678471774852369e-06, 0.0030259541235864162, 3.9539358112961054e-05, 7.781033491482958e-05, 0.0003711417084559798, 9.1652873379644e-06, 0.0006458949064835906, 0.00023330377007368952, 0.00865631178021431], [0.0011992683866992593, 0.008629350923001766, 0.6251504421234131, 0.015135818161070347, 0.001978840446099639, 0.000745285302400589, 5.708653407054953e-05, 0.00043479635496623814, 0.0005481417756527662, 0.0016355890547856688, 0.0002436988870613277, 5.164237336430233e-06, 4.976044510840438e-05, 3.400173591217026e-05, 0.00024351823958568275], [0.006698334589600563, 0.006304558366537094, 0.34660738706588745, 0.7217360138893127, 0.06864907592535019, 0.0027605369687080383, 0.0006927561480551958, 0.00010832686530193314, 0.0002978279662784189, 0.007849807851016521, 0.0023863124661147594, 8.873132173903286e-06, 2.0952818886144087e-05, 4.62439584225649e-06, 0.000559441396035254], [0.0006861803703941405, 0.036174044013023376, 0.4128260612487793, 0.09897080808877945, 0.6376775503158569, 0.19431157410144806, 0.0007082957308739424, 0.05852581560611725, 0.0003548018867149949, 0.00026609119959175587, 0.0006576658925041556, 0.0007862210040912032, 0.027955245226621628, 0.006076914723962545, 0.0010327105410397053], [1.7293352305713938e-09, 1.4693102912133327e-06, 3.0192679332685657e-05, 1.0152590220968705e-05, 0.005660888738930225, 0.5108420252799988, 0.0005426039570011199, 0.0008102089632302523, 3.168102921335958e-06, 6.12798771726375e-08, 2.5310575324510864e-07, 5.088519174023531e-06, 0.00021843344438821077, 2.5946601454052143e-06, 2.594279294498847e-06], [7.755387923680246e-05, 3.5259185096947476e-05, 0.0012139425380155444, 0.00035162578569725156, 0.00505053298547864, 0.4696201980113983, 0.5859625339508057, 0.009771172888576984, 0.0005853781476616859, 3.0261137453635456e-06, 1.2206013707327656e-05, 2.2465645088232122e-05, 0.013555033132433891, 0.0011026648571714759, 7.656160596525297e-05], [3.390625025190275e-08, 5.7732322602532804e-05, 3.19563605444273e-06, 2.0829493507790175e-07, 5.039521965954918e-06, 0.00017657184798736125, 0.000729007413610816, 0.8331114649772644, 0.0037640428636223078, 1.5948112377373036e-06, 5.8014775277115405e-06, 4.528372699041938e-07, 0.00020723954366985708, 0.00025866259238682687, 1.95706252270611e-06], [2.7739795882553153e-07, 2.501485141692683e-05, 4.778147285833256e-06, 3.7190903867667657e-07, 9.610201523457818e-09, 1.1292572708043735e-06, 1.2355405942798825e-07, 3.984562499681488e-05, 0.6202287077903748, 0.0002610959345474839, 0.00017016819037962705, 9.242457963409834e-07, 2.799387630147976e-06, 3.2760857493485673e-07, 1.038134087139042e-06], [1.2775580216839444e-05, 0.0010497755138203502, 6.564326031366363e-05, 4.172011358605232e-06, 4.676745959386608e-07, 3.6489967669695034e-07, 8.09820832614605e-08, 5.78842673348845e-06, 0.0015375507064163685, 0.7445451617240906, 0.026254041120409966, 8.213486580643803e-05, 1.1159563655382954e-05, 3.0355058697750792e-05, 2.6809220798895694e-06], [1.3068409316474572e-05, 0.00010775982809718698, 0.00024633039720356464, 3.3576598070794716e-05, 4.556980275083333e-05, 1.0597023702985098e-07, 9.86238859468358e-08, 2.1072135041322326e-06, 0.0013669389300048351, 0.5916010141372681, 0.4436832368373871, 0.0013138806680217385, 4.73510908705066e-06, 6.116700660641072e-06, 2.961193558803643e-06], [4.950460061081685e-05, 0.0011237917933613062, 0.017257435247302055, 0.0011414129985496402, 0.025087760761380196, 0.00036485170130617917, 3.213326635886915e-05, 5.293267349770758e-06, 4.4593522034119815e-05, 0.001686945091933012, 0.00823597889393568, 0.8047888278961182, 0.014818375930190086, 0.006413417402654886, 2.281446177221369e-05], [0.000998240546323359, 0.1768636256456375, 0.0663335844874382, 0.02716292440891266, 0.03197554498910904, 0.001621886040084064, 0.00012482069723773748, 7.020989141892642e-05, 0.08078382909297943, 0.1701173484325409, 0.08303841948509216, 0.5506232380867004, 0.06293172389268875, 0.03332124650478363, 0.0033543158788233995], [0.021357281133532524, 0.0013016555458307266, 0.00422634556889534, 0.00104909623041749, 0.012563652358949184, 0.07401228696107864, 0.007866809144616127, 0.0024991247337311506, 0.0011657974682748318, 5.4276370065053925e-06, 0.0024851916823536158, 0.0298884529620409, 0.4522511959075928, 0.2182934284210205, 0.14462554454803467]], [[0.03249572962522507, 0.01680905371904373, 0.01368993055075407, 0.005182549823075533, 0.0014828554121777415, 0.0045396420173347, 0.0006250899168662727, 0.01684878207743168, 0.005824672989547253, 0.007428525947034359, 0.009805276058614254, 0.003550198394805193, 0.007900950498878956, 0.009690256789326668, 0.18011362850666046], [0.11159665137529373, 0.10346578061580658, 0.414338618516922, 0.08694489300251007, 0.2136271595954895, 0.10264819115400314, 0.023593097925186157, 0.0335584320127964, 0.0575689822435379, 0.06024341657757759, 0.1307218372821808, 0.13801440596580505, 0.1756829470396042, 0.14866231381893158, 0.1320090889930725], [0.1948547214269638, 0.038279034197330475, 0.07790879160165787, 0.04177340865135193, 0.004589961376041174, 0.0009778933599591255, 0.002051346004009247, 0.006739486940205097, 0.009280361235141754, 0.0007642557029612362, 0.0012637393083423376, 0.00433916924521327, 0.00236115837469697, 0.008354227058589458, 0.2381056696176529], [0.07799407094717026, 0.10201291739940643, 0.037178199738264084, 0.03369736298918724, 0.035083431750535965, 0.003606606973335147, 0.0009816481033340096, 0.010917055420577526, 0.019562464207410812, 0.004011118784546852, 0.0029224867466837168, 0.0011325542582198977, 0.00486336974427104, 0.007979645393788815, 0.2784355580806732], [0.11467810720205307, 0.4025481641292572, 0.4041208028793335, 0.13489782810211182, 0.520052433013916, 0.013409112580120564, 0.0056337821297347546, 0.04408307746052742, 0.06485209614038467, 0.0023049998562783003, 0.0050890627317130566, 0.004091872368007898, 0.006159461103379726, 0.0242836382240057, 0.07189745455980301], [0.1516697108745575, 0.2241159826517105, 0.5074643492698669, 0.3874017000198364, 0.2519407868385315, 0.032381314784288406, 0.015091626904904842, 0.006451433524489403, 0.09749187529087067, 0.007731522433459759, 0.00912014115601778, 0.029297562316060066, 0.05765664204955101, 0.059585090726614, 0.023513801395893097], [0.01171550527215004, 0.10137046873569489, 0.870269238948822, 0.5154522657394409, 0.6626715660095215, 0.08923148363828659, 0.047533176839351654, 0.015608957968652248, 0.11948943883180618, 0.008091520518064499, 0.008133050054311752, 0.012773845344781876, 0.051611315459012985, 0.01502595841884613, 0.00961183663457632], [0.01722140610218048, 0.036506716161966324, 0.7147647738456726, 0.20675897598266602, 0.8291797637939453, 0.31030455231666565, 0.11803850531578064, 0.03327609598636627, 0.4245462417602539, 0.013293992727994919, 0.008976193144917488, 0.054750751703977585, 0.1754072904586792, 0.04528210312128067, 0.012820743955671787], [0.01982569508254528, 0.15988187491893768, 0.12975367903709412, 0.1326102912425995, 0.6299260258674622, 0.28946900367736816, 0.34108322858810425, 0.11804011464118958, 0.16752222180366516, 0.01777276024222374, 0.0021109972149133682, 0.0006076672580093145, 0.0030632279813289642, 0.00126487051602453, 0.1333881914615631], [0.005461913999170065, 0.03046412020921707, 0.008993657305836678, 0.005659051705151796, 0.004244270734488964, 0.02773391455411911, 0.042834386229515076, 0.13534432649612427, 0.27069228887557983, 0.04962563514709473, 0.015227400697767735, 0.0016283531440421939, 0.0014969720505177975, 0.0027089377399533987, 0.17130999267101288], [0.01672529987990856, 0.10339350253343582, 0.009749630466103554, 0.02030925825238228, 0.017326004803180695, 0.03957638517022133, 0.030999623239040375, 0.10308665037155151, 0.5008098483085632, 0.09767498821020126, 0.09780175238847733, 0.025981366634368896, 0.003117683343589306, 0.00962040200829506, 0.1932818591594696], [0.026731140911579132, 0.05838552862405777, 0.07611822336912155, 0.05796685442328453, 0.5904980301856995, 0.010755263268947601, 0.0517524816095829, 0.055663660168647766, 0.29654714465141296, 0.1307908594608307, 0.1585402488708496, 0.03976760059595108, 0.07525579631328583, 0.16488958895206451, 0.1035238653421402], [0.024593327194452286, 0.12932555377483368, 0.13568159937858582, 0.16021546721458435, 0.3227141201496124, 0.029398979619145393, 0.01611196994781494, 0.016819216310977936, 0.2378186136484146, 0.5602607131004333, 0.7615779638290405, 0.08417549729347229, 0.10783103108406067, 0.2013072967529297, 0.06744378060102463], [0.018169090151786804, 0.26050350069999695, 0.078061044216156, 0.023439347743988037, 0.05254700779914856, 0.0014709478709846735, 0.002907117595896125, 0.009980114176869392, 0.1381266713142395, 0.5626046061515808, 0.5405392646789551, 0.11909772455692291, 0.008021530695259571, 0.06359856575727463, 0.009888176806271076], [0.08646434545516968, 0.009946366772055626, 0.041608210653066635, 0.009163393639028072, 0.12723588943481445, 0.17822976410388947, 0.01437843032181263, 0.0057503837160766125, 0.008486853912472725, 0.002935740165412426, 0.019836073741316795, 0.07525425404310226, 0.02854214422404766, 0.0230310820043087, 0.1518138200044632]], [[0.7472922801971436, 0.06644202023744583, 0.12477048486471176, 0.07691145688295364, 0.17426471412181854, 0.17453429102897644, 0.8713244795799255, 0.22852616012096405, 0.7413471937179565, 0.5253387689590454, 0.16250024735927582, 0.19445888698101044, 0.10716042667627335, 0.2310180366039276, 0.05536508187651634], [0.13811203837394714, 0.40626850724220276, 0.2430061399936676, 0.22277961671352386, 0.18414726853370667, 0.21574343740940094, 0.8225958943367004, 0.5822084546089172, 0.41659367084503174, 0.35776287317276, 0.4909748136997223, 0.39181941747665405, 0.34554892778396606, 0.6003718972206116, 0.043436333537101746], [0.03130434453487396, 0.0024298657663166523, 0.43690061569213867, 0.5043830275535583, 0.07530603557825089, 0.015139158815145493, 0.03498073294758797, 0.012510559521615505, 0.6034607291221619, 0.7801509499549866, 0.8402397036552429, 0.5008089542388916, 0.17657218873500824, 0.11879491806030273, 0.05205746740102768], [0.09661327302455902, 0.049034956842660904, 0.05331439897418022, 0.7222777009010315, 0.25703296065330505, 0.020087046548724174, 0.06235986202955246, 0.0651831179857254, 0.32113927602767944, 0.5460676550865173, 0.7442458271980286, 0.5571728348731995, 0.08091285824775696, 0.059992171823978424, 0.029936296865344048], [0.00972762517631054, 0.007879518903791904, 0.02767527848482132, 0.019306808710098267, 0.22303025424480438, 0.007516835816204548, 0.007440114859491587, 0.022099999710917473, 0.29848337173461914, 0.9075287580490112, 0.5192471742630005, 0.8959035873413086, 0.055479276925325394, 0.04288056865334511, 0.021558567881584167], [0.03836950287222862, 0.05839527025818825, 0.005887853913009167, 0.08494037389755249, 0.012977076694369316, 0.5726994872093201, 0.09935679286718369, 0.13719113171100616, 0.448569655418396, 0.5218547582626343, 0.13800226151943207, 0.1732572466135025, 0.4354798197746277, 0.4542965292930603, 0.12337890267372131], [0.17566490173339844, 0.03925755247473717, 0.01956782303750515, 0.04187121242284775, 0.02149910107254982, 0.049183186143636703, 0.5663522481918335, 0.045388396829366684, 0.45039302110671997, 0.19015204906463623, 0.22913624346256256, 0.10953018814325333, 0.21400360763072968, 0.572381854057312, 0.1667298972606659], [0.2136794924736023, 0.20810233056545258, 0.08830246329307556, 0.27903637290000916, 0.02317022904753685, 0.10591837763786316, 0.15087167918682098, 0.5299598574638367, 0.3452024757862091, 0.15965056419372559, 0.2765912711620331, 0.516273021697998, 0.2846863567829132, 0.3888777792453766, 0.0719258189201355], [0.07398565858602524, 0.04620325192809105, 0.3374384939670563, 0.19415578246116638, 0.025615269318223, 0.010194968432188034, 0.018451105803251266, 0.0005573831731453538, 0.5073301196098328, 0.25312942266464233, 0.15244188904762268, 0.143111914396286, 0.051979612559080124, 0.04884689673781395, 0.12363318353891373], [0.5805832147598267, 0.09438126534223557, 0.24455930292606354, 0.06023820489645004, 0.03943831846117973, 0.021930387243628502, 0.026398053392767906, 0.012488989159464836, 0.011794325895607471, 0.767930269241333, 0.4412824809551239, 0.07896611094474792, 0.01228941697627306, 0.018458310514688492, 0.10866446793079376], [0.1145540103316307, 0.05171298235654831, 0.7072227597236633, 0.4839639961719513, 0.11294537037611008, 0.06211492419242859, 0.021921994164586067, 0.0025394419208168983, 0.0033554628025740385, 0.07357389479875565, 0.7795555591583252, 0.05686911940574646, 0.022035235539078712, 0.034172482788562775, 0.07262071967124939], [0.08121224492788315, 0.025126218795776367, 0.4891066551208496, 0.29065003991127014, 0.20622830092906952, 0.36699986457824707, 0.07864820212125778, 0.014422299340367317, 0.016684990376234055, 0.0649130716919899, 0.07936163991689682, 0.6605017185211182, 0.18783104419708252, 0.08294262737035751, 0.03477967903017998], [0.0700722336769104, 0.1311686784029007, 0.5332850813865662, 0.1558467000722885, 0.36321985721588135, 0.7912644743919373, 0.32202765345573425, 0.1934671401977539, 0.031114375218749046, 0.09986341744661331, 0.08630139380693436, 0.055017780512571335, 0.44781896471977234, 0.42446693778038025, 0.1060790941119194], [0.08875010907649994, 0.06247853487730026, 0.4616371989250183, 0.12711729109287262, 0.3074216842651367, 0.19363558292388916, 0.2020244151353836, 0.0779867023229599, 0.019831692799925804, 0.03570472076535225, 0.07392378151416779, 0.04282142594456673, 0.0921483263373375, 0.3143211603164673, 0.22281906008720398], [0.5682113766670227, 0.1249876543879509, 0.7342633008956909, 0.902918815612793, 0.7035764455795288, 0.3718622326850891, 0.6157594919204712, 0.15625660121440887, 0.8438207507133484, 0.9341241121292114, 0.8159937858581543, 0.6624717712402344, 0.3264457583427429, 0.5970154404640198, 0.003644895739853382]], [[0.0183254461735487, 0.00659788167104125, 0.046570390462875366, 0.04327844828367233, 0.10241857916116714, 0.5407979488372803, 0.0026681027375161648, 0.15349310636520386, 0.0016508381813764572, 0.010916458442807198, 0.036675866693258286, 0.15769276022911072, 0.4073828458786011, 0.04228133708238602, 0.15622197091579437], [0.07985992729663849, 0.06383417546749115, 0.024972105398774147, 0.18746882677078247, 0.11770728975534439, 0.13333363831043243, 0.006719768047332764, 0.04288880154490471, 0.001412510173395276, 0.058754052966833115, 0.14280158281326294, 0.13529875874519348, 0.08268098533153534, 0.02367851696908474, 0.1494951695203781], [0.01403640117496252, 0.014278309419751167, 0.1034439280629158, 0.022417087107896805, 0.10706920921802521, 0.018271848559379578, 0.046350300312042236, 0.04233889281749725, 0.037542134523391724, 0.0005760823260061443, 0.004724643658846617, 0.233056902885437, 0.2574465572834015, 0.1892177164554596, 0.21611936390399933], [0.032590243965387344, 0.14464972913265228, 0.1993260532617569, 0.12327495217323303, 0.27639931440353394, 0.011173157021403313, 0.012838426046073437, 0.0802190750837326, 0.0400678850710392, 0.013469994999468327, 0.025247203186154366, 0.30583158135414124, 0.6397863626480103, 0.258308470249176, 0.08317234367132187], [0.007401467300951481, 0.04209339618682861, 0.1104009672999382, 0.04737341031432152, 0.06253770738840103, 0.0023836863692849874, 0.05026397854089737, 0.01439946424216032, 0.006556188687682152, 0.001721409265883267, 0.01908556930720806, 0.022761031985282898, 0.01600046642124653, 0.22344018518924713, 0.2855986952781677], [0.00031611474696546793, 0.010241325944662094, 0.005327185150235891, 0.007503898814320564, 0.009216651320457458, 0.08986854553222656, 0.0022410263773053885, 0.04830501973628998, 0.013246790505945683, 0.0036830154713243246, 0.001605262397788465, 0.004246865399181843, 0.005818811245262623, 0.00778583250939846, 0.2319662719964981], [0.00028042105259373784, 0.004604758229106665, 0.008834331296384335, 0.010530425235629082, 0.04934454336762428, 0.3239482641220093, 0.02964387647807598, 0.041019540280103683, 0.028070107102394104, 0.002580034313723445, 0.0034616885241121054, 0.006594499107450247, 0.07731658220291138, 0.01784621551632881, 0.10414844751358032], [0.002352550160139799, 0.00811008270829916, 0.007519579492509365, 0.09616736322641373, 0.00784054771065712, 0.06404154002666473, 0.025837063789367676, 0.06720300018787384, 0.008001329377293587, 0.016075177118182182, 0.0036620565224438906, 0.031110821291804314, 0.1529460847377777, 0.03003939613699913, 0.19531111419200897], [0.014062762260437012, 0.03979215770959854, 0.0070105125196278095, 0.010145032778382301, 0.023933248594403267, 0.08613994717597961, 0.027301009744405746, 0.007488427218049765, 0.04610109701752663, 0.00706111453473568, 0.005716769024729729, 0.008516461588442326, 0.04168170318007469, 0.004054774064570665, 0.3198099434375763], [0.0027477010153234005, 0.009237049147486687, 0.005884162615984678, 0.004349177703261375, 0.039300523698329926, 0.06504905968904495, 0.005921225529164076, 0.05048412084579468, 0.004538795445114374, 0.019958311691880226, 0.08035917580127716, 0.1339075267314911, 0.45191076397895813, 0.1108468547463417, 0.15996994078159332], [0.0004566281568259001, 0.0044615683145821095, 0.008062957786023617, 0.0003266451822128147, 0.032452184706926346, 0.004190187435597181, 0.0009983428753912449, 0.0015420016134157777, 0.025539150461554527, 0.0009114624699577689, 0.001308016013354063, 0.11249691247940063, 0.5262115597724915, 0.16036535799503326, 0.02284345217049122], [0.006384413689374924, 0.006966868881136179, 0.013256898149847984, 0.008146845735609531, 0.005910678766667843, 0.005924733821302652, 0.0029809526167809963, 0.004338744096457958, 0.0021091948729008436, 0.02691148780286312, 0.09123647958040237, 0.0904775932431221, 0.10420377552509308, 0.019918829202651978, 0.21981710195541382], [0.004395737312734127, 0.0342060811817646, 0.08344801515340805, 0.012639162130653858, 0.07537969946861267, 0.00383414002135396, 0.007808698806911707, 0.007516762241721153, 0.0023650380317121744, 0.055798787623643875, 0.025632014498114586, 0.040716953575611115, 0.16482838988304138, 0.13848447799682617, 0.17180821299552917], [0.0016022673808038235, 0.013307235203683376, 0.012306403368711472, 0.0029055906925350428, 0.06092625483870506, 0.01653674617409706, 0.008309547789394855, 0.00395687622949481, 0.002493055537343025, 0.0038927635177969933, 0.009680269286036491, 0.23031921684741974, 0.35693949460983276, 0.1708209365606308, 0.050492819398641586], [0.009627100080251694, 0.006502249743789434, 0.0023533182684332132, 0.0021814347710460424, 0.007286426145583391, 0.024909881874918938, 0.01453662570565939, 0.010449647903442383, 0.0028000103775411844, 0.001988302916288376, 0.001580765936523676, 0.013102496974170208, 0.001836722600273788, 0.0008430163725279272, 0.15720587968826294]], [[0.060514166951179504, 0.09119007736444473, 0.5136731863021851, 0.024349171668291092, 0.41056114435195923, 0.043175265192985535, 0.016160618513822556, 0.12711943686008453, 0.029147693887352943, 0.01592664048075676, 0.04504424333572388, 0.03736018016934395, 0.026280265301465988, 0.042564861476421356, 0.13562467694282532], [0.009338664822280407, 0.09596994519233704, 0.12376897037029266, 0.01794583536684513, 0.059337858110666275, 0.04990454390645027, 0.003890786785632372, 0.07171432673931122, 0.0057785604149103165, 0.005389686673879623, 0.009663187898695469, 0.014342015609145164, 0.020640142261981964, 0.04060304909944534, 0.16408833861351013], [0.07689530402421951, 0.027863014489412308, 0.15549975633621216, 0.2693096697330475, 0.73520827293396, 0.03749871999025345, 0.3640631139278412, 0.14002074301242828, 0.16656053066253662, 0.02643253095448017, 0.0061660525389015675, 0.054253485053777695, 0.14240022003650665, 0.14975441992282867, 0.13701564073562622], [0.21953634917736053, 0.22122228145599365, 0.04846278205513954, 0.07968296110630035, 0.3619323670864105, 0.03181222453713417, 0.6669740080833435, 0.3975786566734314, 0.11174946278333664, 0.15518029034137726, 0.004886193200945854, 0.010736972093582153, 0.07725195586681366, 0.09191425889730453, 0.1523013859987259], [0.0740056112408638, 0.054083533585071564, 0.027193741872906685, 0.014972379431128502, 0.04523617774248123, 0.012482533231377602, 0.4212614595890045, 0.25695085525512695, 0.3699147403240204, 0.013461914844810963, 0.08041262626647949, 0.015268572606146336, 0.627507209777832, 0.13811761140823364, 0.19850368797779083], [0.029503263533115387, 0.09333665668964386, 0.016309864819049835, 0.1364656686782837, 0.03873518481850624, 0.019083604216575623, 0.758955180644989, 0.6250144243240356, 0.10551930963993073, 0.0059091635048389435, 0.001959211425855756, 0.004587537609040737, 0.0029548059683293104, 0.011073557659983635, 0.10497581213712692], [0.0038599083200097084, 0.03815716505050659, 0.004112291149795055, 0.0037336996756494045, 0.02896580658853054, 0.003606554586440325, 0.2724342346191406, 0.5795999765396118, 0.041377726942300797, 0.01812332309782505, 0.006642999593168497, 0.006629596464335918, 0.018780261278152466, 0.00801254715770483, 0.11063171178102493], [0.023342538625001907, 0.1589166522026062, 0.01254882663488388, 0.01894153468310833, 0.04743911698460579, 0.015340029262006283, 0.06989605724811554, 0.22605817019939423, 0.016811540350317955, 0.014681086875498295, 0.0061398339457809925, 0.02630683407187462, 0.032653048634529114, 0.05358496680855751, 0.18197578191757202], [0.01728241890668869, 0.12100599706172943, 0.003952578641474247, 0.038103699684143066, 0.00803869217634201, 0.017839567735791206, 0.040644098073244095, 0.014622771181166172, 0.07288665324449539, 0.4550913870334625, 0.18886235356330872, 0.2150641530752182, 0.487347275018692, 0.42817094922065735, 0.12942945957183838], [0.011775199323892593, 0.1349712610244751, 0.005470172502100468, 0.003098055487498641, 0.028361253440380096, 0.03303566575050354, 0.007174484897404909, 0.015601159073412418, 0.006606224924325943, 0.08859884738922119, 0.18040567636489868, 0.31761303544044495, 0.2462366670370102, 0.4818485677242279, 0.12394269555807114], [0.05270439758896828, 0.1637289971113205, 0.009510326199233532, 0.008013473823666573, 0.14090411365032196, 0.011389089748263359, 0.013123652897775173, 0.023534703999757767, 0.009078129194676876, 0.02855684608221054, 0.026650836691260338, 0.39132389426231384, 0.16291603446006775, 0.25967708230018616, 0.10212607681751251], [0.19571052491664886, 0.10246216505765915, 0.02142595686018467, 0.012254489585757256, 0.00365867605432868, 0.007110960781574249, 0.020346596837043762, 0.03192196041345596, 0.00833944883197546, 0.07423693686723709, 0.09786227345466614, 0.08075869083404541, 0.1330210417509079, 0.26891645789146423, 0.17930860817432404], [0.11616674810647964, 0.175978422164917, 0.00425378605723381, 0.017427049577236176, 0.011484457179903984, 0.030517226085066795, 0.08637198060750961, 0.1500588357448578, 0.0009573447750881314, 0.044167183339595795, 0.005869577638804913, 0.0011607500491663814, 0.014711305499076843, 0.027834221720695496, 0.18594378232955933], [0.11675343662500381, 0.17556257545948029, 0.016423039138317108, 0.02097608894109726, 0.06606884300708771, 0.06371303647756577, 0.09760221093893051, 0.2481643557548523, 0.0015754855703562498, 0.03009907715022564, 0.03618617355823517, 0.012020162306725979, 0.17486301064491272, 0.22630257904529572, 0.2108311653137207], [0.004961065016686916, 0.011551961302757263, 0.006318831816315651, 0.002851473866030574, 0.003461753251031041, 0.011111320927739143, 0.004611799493432045, 0.004697122145444155, 0.0026004482060670853, 0.0010426584631204605, 0.0060967751778662205, 0.01239971723407507, 0.004622939508408308, 0.002610035240650177, 0.15716104209423065]], [[0.027552247047424316, 0.013821233063936234, 0.004237555433064699, 0.0007387229125015438, 0.0009859473211690784, 0.001997306477278471, 0.002160864183679223, 0.009250090457499027, 0.0009738927474245429, 0.0009403586154803634, 0.003406830132007599, 0.0010056114988401532, 0.008306043222546577, 0.06191018968820572, 0.18169914186000824], [0.0056476471945643425, 0.0617278628051281, 0.026225095614790916, 0.009516767226159573, 0.019543437287211418, 0.011766157113015652, 0.0015307252760976553, 0.004000868182629347, 0.006223553325980902, 0.02180931344628334, 0.02397397719323635, 0.025289250537753105, 0.01872297003865242, 0.05591608211398125, 0.17309869825839996], [0.5742589831352234, 0.02769068442285061, 0.03131784498691559, 0.008496972732245922, 0.005279624368995428, 0.0009009581408463418, 0.013010378926992416, 0.009255914948880672, 0.08095329999923706, 0.0017015798948705196, 0.0027918636333197355, 0.01474103331565857, 0.07241056859493256, 0.2960302531719208, 0.1991364061832428], [0.3870091140270233, 0.24428580701351166, 0.004871743265539408, 0.01251932606101036, 0.004600874613970518, 0.007045479491353035, 0.011942178010940552, 0.06100638955831528, 0.06223933771252632, 0.00421120086684823, 0.0017708303639665246, 0.010406754910945892, 0.016386834904551506, 0.038040366023778915, 0.25559180974960327], [0.6136646866798401, 0.2692064642906189, 0.043582458049058914, 0.00652115186676383, 0.05291604623198509, 0.006654517259448767, 0.03398957848548889, 0.03886384516954422, 0.13169772922992706, 0.002106831641867757, 0.005907678045332432, 0.01888049766421318, 0.04876947030425072, 0.2226717472076416, 0.22327177226543427], [0.685612678527832, 0.0861489400267601, 0.03236214071512222, 0.16196951270103455, 0.03394145518541336, 0.05551951378583908, 0.027528556063771248, 0.06770895421504974, 0.19389298558235168, 0.03780713677406311, 0.0038191182538866997, 0.05989958345890045, 0.13479465246200562, 0.24111053347587585, 0.15613426268100739], [0.6876600384712219, 0.0606975182890892, 0.05783677101135254, 0.05387236177921295, 0.11914167553186417, 0.004756046459078789, 0.031782086938619614, 0.011465699411928654, 0.1448838710784912, 0.09538520872592926, 0.007872258313000202, 0.033316925168037415, 0.09786565601825714, 0.08940181881189346, 0.23629719018936157], [0.5363585352897644, 0.11579979956150055, 0.10718797892332077, 0.21453110873699188, 0.030864767730236053, 0.026318436488509178, 0.03807519003748894, 0.12262200564146042, 0.08015674352645874, 0.06537020206451416, 0.004594390746206045, 0.015254726633429527, 0.06485987454652786, 0.039039257913827896, 0.16586215794086456], [0.6220377087593079, 0.17304541170597076, 0.23731492459774017, 0.32412996888160706, 0.2203587144613266, 0.09306959062814713, 0.2822628319263458, 0.008407875895500183, 0.14113475382328033, 0.022416740655899048, 0.005183607805520296, 0.0005837879725731909, 0.00799399521201849, 0.006284625735133886, 0.12005029618740082], [0.18509520590305328, 0.21334251761436462, 0.12845394015312195, 0.3693835139274597, 0.41559898853302, 0.19613976776599884, 0.7053389549255371, 0.3886314332485199, 0.06599769741296768, 0.04325481504201889, 0.029052795842289925, 0.001557054347358644, 0.0018087843200191855, 0.0036887156311422586, 0.18107539415359497], [0.612794041633606, 0.24153079092502594, 0.076973557472229, 0.17341682314872742, 0.06242084503173828, 0.2242424041032791, 0.8304246068000793, 0.5655775666236877, 0.4262824058532715, 0.00936043355613947, 0.03881426528096199, 0.0046007027849555016, 0.005786797031760216, 0.020520325750112534, 0.226027712225914], [0.21637925505638123, 0.22487440705299377, 0.19202512502670288, 0.3957260847091675, 0.15970049798488617, 0.16693006455898285, 0.3690066933631897, 0.5193001627922058, 0.6459834575653076, 0.047006867825984955, 0.06868032366037369, 0.043628890067338943, 0.02405296452343464, 0.05333276465535164, 0.08607933670282364], [0.5923737287521362, 0.3536633849143982, 0.08390633016824722, 0.2980528473854065, 0.042989592999219894, 0.026934657245874405, 0.1647067815065384, 0.1620720773935318, 0.6647022366523743, 0.13678880035877228, 0.10115252435207367, 0.012052871286869049, 0.2444845736026764, 0.1799331158399582, 0.10357851535081863], [0.3260110914707184, 0.10825559496879578, 0.040669191628694534, 0.08903322368860245, 0.055108752101659775, 0.014200238510966301, 0.06877616047859192, 0.07561883330345154, 0.7116665244102478, 0.08518233895301819, 0.13964912295341492, 0.01787719503045082, 0.027594367042183876, 0.0709126889705658, 0.09409899264574051], [0.26070404052734375, 0.8011303544044495, 0.17980173230171204, 0.0725909024477005, 0.12434736639261246, 0.28980228304862976, 0.3281027674674988, 0.7843722701072693, 0.12677432596683502, 0.054726697504520416, 0.13370326161384583, 0.19018130004405975, 0.1707623451948166, 0.14939220249652863, 0.07447532564401627]], [[0.10194799304008484, 0.042179130017757416, 0.27587375044822693, 0.8387316465377808, 0.3051532208919525, 0.225641667842865, 0.10655678808689117, 0.4426303505897522, 0.21958006918430328, 0.4376780688762665, 0.7421585917472839, 0.6036965250968933, 0.4420715570449829, 0.6119644045829773, 0.08460802584886551], [0.052479684352874756, 0.018692737445235252, 0.13130725920200348, 0.4463008642196655, 0.4007475674152374, 0.4465942680835724, 0.13863760232925415, 0.26287177205085754, 0.5015351176261902, 0.48749616742134094, 0.19089040160179138, 0.2783986032009125, 0.20843097567558289, 0.11412637680768967, 0.11901978403329849], [0.09998084604740143, 0.05760321766138077, 0.06884635984897614, 0.1367950737476349, 0.03696327656507492, 0.02052011340856552, 0.23966658115386963, 0.6639524102210999, 0.08913422375917435, 0.1896458864212036, 0.14239966869354248, 0.18587030470371246, 0.2512775659561157, 0.1800404042005539, 0.13985422253608704], [0.17776982486248016, 0.2164098620414734, 0.03016561083495617, 0.006355184596031904, 0.04318562150001526, 0.004709928296506405, 0.02340516820549965, 0.07859960943460464, 0.3921053409576416, 0.27134451270103455, 0.2182498425245285, 0.1118401437997818, 0.13378913700580597, 0.4978374242782593, 0.18931511044502258], [0.16739480197429657, 0.20097726583480835, 0.038037389516830444, 0.05488090589642525, 0.020769814029335976, 0.044557277113199234, 0.32692524790763855, 0.5529306530952454, 0.06495681405067444, 0.061963245272636414, 0.3602059483528137, 0.040287844836711884, 0.11072657257318497, 0.3166219890117645, 0.19249440729618073], [0.07948607206344604, 0.4389178156852722, 0.019072405993938446, 0.11389600485563278, 0.015004596672952175, 0.0008035529754124582, 0.00560334138572216, 0.007579134311527014, 0.12602436542510986, 0.4041804373264313, 0.8435949087142944, 0.7255359292030334, 0.3334953784942627, 0.21919409930706024, 0.13174442946910858], [0.11827840656042099, 0.43549492955207825, 0.035650141537189484, 0.3500109016895294, 0.10479609668254852, 0.0029047641437500715, 0.016262628138065338, 0.008920608088374138, 0.1923075020313263, 0.6588289737701416, 0.7271849513053894, 0.8207041025161743, 0.5342087149620056, 0.29674431681632996, 0.16698533296585083], [0.19771254062652588, 0.43774574995040894, 0.057631127536296844, 0.15638697147369385, 0.05497771501541138, 0.0015852008946239948, 0.004800108727067709, 0.0038221883587539196, 0.11230877041816711, 0.6780416369438171, 0.6535694003105164, 0.33372464776039124, 0.2617355287075043, 0.4378974735736847, 0.15096917748451233], [0.2510830760002136, 0.455088347196579, 0.2769528925418854, 0.28598156571388245, 0.08308438956737518, 0.495423823595047, 0.2878262400627136, 0.017540372908115387, 0.036487918347120285, 0.07030303031206131, 0.04537871107459068, 0.017587929964065552, 0.15749330818653107, 0.15622387826442719, 0.134229376912117], [0.2108728438615799, 0.12734071910381317, 0.6047671437263489, 0.5566261410713196, 0.4727993309497833, 0.6295000314712524, 0.20963285863399506, 0.3828260004520416, 0.01981351152062416, 0.02910005673766136, 0.17932364344596863, 0.029557999223470688, 0.02868420071899891, 0.05513756722211838, 0.1339428722858429], [0.2013130933046341, 0.35711804032325745, 0.18803814053535461, 0.31239861249923706, 0.6328845024108887, 0.6068195104598999, 0.09879770874977112, 0.295420378446579, 0.033300116658210754, 0.04495004564523697, 0.027333615347743034, 0.034196678549051285, 0.011724627576768398, 0.023517103865742683, 0.3543241322040558], [0.27807915210723877, 0.07025524973869324, 0.15421687066555023, 0.23079168796539307, 0.0323871448636055, 0.4182601273059845, 0.43312954902648926, 0.3330070972442627, 0.027521615847945213, 0.03977188467979431, 0.03152378648519516, 0.00340716983191669, 0.005408053286373615, 0.0057552107609808445, 0.23170912265777588], [0.15765754878520966, 0.07761365175247192, 0.1382310688495636, 0.33822664618492126, 0.15857987105846405, 0.11602839827537537, 0.3749851584434509, 0.3412497341632843, 0.06253337115049362, 0.09931040555238724, 0.010201470926404, 0.0010190334869548678, 0.0007929145358502865, 0.0016151106683537364, 0.1723894327878952], [0.39988550543785095, 0.09145350754261017, 0.3013111352920532, 0.5813722610473633, 0.4042908251285553, 0.2935561537742615, 0.4903331696987152, 0.4357178807258606, 0.04456466808915138, 0.10430204123258591, 0.10590728372335434, 0.007762597873806953, 0.0026525144930928946, 0.0052152471616864204, 0.24974997341632843], [0.03366217389702797, 0.03653215244412422, 0.027766529470682144, 0.007369572762399912, 0.014929202385246754, 0.04527684673666954, 0.00940654892474413, 0.023517949506640434, 0.010960820131003857, 0.0019369145156815648, 0.01981637440621853, 0.00444602407515049, 0.014915830455720425, 0.007271313574165106, 0.15384840965270996]], [[0.011476250365376472, 0.7629169225692749, 0.02116730809211731, 0.010803135111927986, 0.005132503807544708, 0.009303245693445206, 0.0005040443502366543, 0.022131631150841713, 0.001470191520638764, 0.0017710012616589665, 0.0004086543631274253, 0.0022351557854562998, 0.000896299781743437, 0.0005698543391190469, 0.019197434186935425], [0.0024000771809369326, 0.158247172832489, 0.01897430047392845, 0.019486481323838234, 0.0029122373089194298, 0.015832845121622086, 0.0017470666207373142, 0.00117065932136029, 0.01016113068908453, 0.007651789113879204, 0.0020597530528903008, 0.015201352536678314, 0.016943661496043205, 0.009769451804459095, 0.16634535789489746], [0.00410552928224206, 0.0015743908006697893, 0.01049637421965599, 0.006504607852548361, 0.035339318215847015, 0.9065937995910645, 0.2998698651790619, 0.12215600907802582, 0.013029203750193119, 0.000650988076813519, 0.002043183660134673, 0.006920983083546162, 0.09688588231801987, 0.057574767619371414, 0.009054930880665779], [0.007287806831300259, 0.01375514268875122, 0.001530585577711463, 0.007056740578263998, 0.01978658139705658, 0.9208202958106995, 0.2214416116476059, 0.30606138706207275, 0.052588097751140594, 0.004079628270119429, 0.0024339878000319004, 0.0028739250265061855, 0.04695972800254822, 0.045893676578998566, 0.0110039496794343], [0.006429406348615885, 0.016907041892409325, 0.0023819799534976482, 0.0003115522558800876, 0.006808500271290541, 0.9102355241775513, 0.15379303693771362, 0.07056371122598648, 0.06324119120836258, 0.0030630400869995356, 0.007665702607482672, 0.002797773340716958, 0.13533660769462585, 0.03197972849011421, 0.006115978583693504], [0.014356410130858421, 0.0526699461042881, 0.0007501932559534907, 0.008851941674947739, 0.0005067299935035408, 0.035332534462213516, 0.09051518887281418, 0.049224019050598145, 0.014900125563144684, 0.01856788620352745, 0.0012414768571034074, 0.002389064058661461, 0.0018446464091539383, 0.000877396494615823, 0.22725383937358856], [0.0025407460052520037, 0.32041609287261963, 0.0036992463283240795, 0.02451898716390133, 0.007920290343463421, 0.015527674928307533, 0.03544912114739418, 0.29718661308288574, 0.02347515895962715, 0.026838794350624084, 0.01756858080625534, 0.010445725172758102, 0.005995406303554773, 0.0005847325082868338, 0.2055930197238922], [0.009255345910787582, 0.034783441573381424, 0.010831266641616821, 0.02782595343887806, 0.001477425335906446, 0.006871670484542847, 0.006518858019262552, 0.0072874827310442924, 0.012387615628540516, 0.05288432911038399, 0.04645476117730141, 0.02255677618086338, 0.014156763441860676, 0.00417641457170248, 0.22105874121189117], [0.0017225841293111444, 0.0049251834861934185, 0.007573804818093777, 0.014873476698994637, 0.00903867557644844, 0.0076865823939442635, 0.0017025101697072387, 0.00023153165238909423, 0.024773191660642624, 0.1742238849401474, 0.6002998948097229, 0.6145275831222534, 0.25023365020751953, 0.35489538311958313, 0.039457567036151886], [0.0034636815544217825, 0.39023807644844055, 0.0018667654367163777, 0.0006454490358009934, 0.00025732445647008717, 0.026610050350427628, 0.0026998629327863455, 0.014584111049771309, 0.00032847325201146305, 0.0012709795264527202, 0.07417861372232437, 0.43676891922950745, 0.25757044553756714, 0.32731080055236816, 0.12109360098838806], [0.0014396773185580969, 0.07700426131486893, 0.0003769460890907794, 0.0015669490676373243, 0.0010665652807801962, 0.05166712775826454, 0.003733921330422163, 0.00829349085688591, 9.729996236274019e-05, 0.0004270579374860972, 0.0022819112055003643, 0.3744491934776306, 0.2681969404220581, 0.4920969009399414, 0.028773367404937744], [0.19549021124839783, 0.5118218064308167, 0.053603943437337875, 0.004430307075381279, 0.0015711480518803, 0.024018822237849236, 0.0441354438662529, 0.04134393110871315, 0.0014472270850092173, 0.024767767637968063, 0.029112013056874275, 0.08014442026615143, 0.4702226519584656, 0.40423843264579773, 0.14477935433387756], [0.034691162407398224, 0.09692039340734482, 0.003936667460948229, 0.0164506658911705, 0.0005446859868243337, 0.0016573348548263311, 0.02795562334358692, 0.12881094217300415, 0.0004645287699531764, 0.0021237744949758053, 0.0010291342623531818, 0.001068241661414504, 0.00471450574696064, 0.019945403560996056, 0.19273433089256287], [0.04783029109239578, 0.11157537996768951, 0.02325829118490219, 0.12799327075481415, 0.0216610599309206, 0.41526544094085693, 0.129922553896904, 0.14850500226020813, 0.0009580283658578992, 0.008097043260931969, 0.01107556838542223, 0.019478609785437584, 0.2748490571975708, 0.11550750583410263, 0.15876543521881104], [0.015012643299996853, 0.00804762914776802, 0.00366173661313951, 0.0018753333715721965, 0.0065993256866931915, 0.00479541253298521, 0.005337378475815058, 0.012457020580768585, 0.0033909485209733248, 0.0032401280477643013, 0.00048777347547002137, 0.012255984358489513, 0.0006230318685993552, 0.001543535152450204, 0.1572250872850418]]], [[[0.016101790592074394, 0.0050575402565300465, 0.008322462439537048, 0.006855499465018511, 0.003766664071008563, 0.0032708626240491867, 0.008669405244290829, 0.016983401030302048, 0.023632090538740158, 0.0007983215618878603, 0.006762287113815546, 0.019076332449913025, 0.0018054646207019687, 0.011848386377096176, 0.23875673115253448], [0.03118298575282097, 0.022700916975736618, 0.01820814236998558, 0.011041272431612015, 0.013735579326748848, 0.003388292621821165, 0.014374880120158195, 0.0029534229543060064, 0.06276529282331467, 0.0010488847037777305, 0.005698299501091242, 0.018068330362439156, 0.009247002191841602, 0.010645000264048576, 0.2274351567029953], [0.10749327391386032, 0.01361121516674757, 0.01930609717965126, 0.025707745924592018, 0.010174103081226349, 0.0019352196250110865, 0.006933925207704306, 0.026056114584207535, 0.003662128932774067, 0.006897854618728161, 0.0015213300939649343, 0.006132383830845356, 0.0028239174280315638, 0.013304864056408405, 0.22739072144031525], [0.25010421872138977, 0.005582309328019619, 0.006115755997598171, 0.08664196729660034, 0.005224197171628475, 0.005311913322657347, 0.03281412273645401, 0.024678068235516548, 0.018595430999994278, 0.0819764956831932, 0.005479714833199978, 0.008821909315884113, 0.02042486146092415, 0.03525637462735176, 0.19444485008716583], [0.1781134456396103, 0.021083489060401917, 0.038613177835941315, 0.16417931020259857, 0.0029645320028066635, 0.00899361353367567, 0.009076704271137714, 0.01357053779065609, 0.01101364754140377, 0.04086701199412346, 0.014270029030740261, 0.011464214883744717, 0.011689195409417152, 0.0706799253821373, 0.3730076551437378], [0.3090042769908905, 0.031162124127149582, 0.033009856939315796, 0.14512063562870026, 0.00411824369803071, 0.07382509857416153, 0.02702517993748188, 0.07667822390794754, 0.021658627316355705, 0.01615101285278797, 0.0066233747638762, 0.008623828180134296, 0.0008525048615410924, 0.011195158585906029, 0.2578849792480469], [0.3291372060775757, 0.0561586357653141, 0.4192807674407959, 0.4571635127067566, 0.057550910860300064, 0.04359428584575653, 0.005270917434245348, 0.03804505616426468, 0.03733760863542557, 0.20409555733203888, 0.04554562643170357, 0.024629684165120125, 0.018161950632929802, 0.04353561997413635, 0.145583838224411], [0.3828665316104889, 0.019200418144464493, 0.34599530696868896, 0.4376910328865051, 0.07537391781806946, 0.036528222262859344, 0.04610925167798996, 0.04538694769144058, 0.1663823127746582, 0.04690397158265114, 0.05553056299686432, 0.021811597049236298, 0.012554574757814407, 0.03599526360630989, 0.1534716635942459], [0.08861738443374634, 0.06363938748836517, 0.7135313749313354, 0.146565243601799, 0.3346884250640869, 0.3544132113456726, 0.12204702943563461, 0.028818881139159203, 0.04564356431365013, 0.03288809210062027, 0.06753166019916534, 0.12387087196111679, 0.029650555923581123, 0.014753012917935848, 0.04379607364535332], [0.03655187785625458, 0.006058508530259132, 0.04018249735236168, 0.08900216966867447, 0.027111714705824852, 0.006408872082829475, 0.03783104568719864, 0.010064247064292431, 0.2550305724143982, 0.008420061320066452, 0.012097015976905823, 0.017737949267029762, 0.0012783813290297985, 0.0026436946354806423, 0.172612726688385], [0.1163061186671257, 0.04424217715859413, 0.014033653773367405, 0.03590161353349686, 0.06527962535619736, 0.00195779325440526, 0.027195196598768234, 0.1581626534461975, 0.30849722027778625, 0.1652299016714096, 0.04234298691153526, 0.05585171654820442, 0.016547594219446182, 0.04909297078847885, 0.08752257376909256], [0.1013311892747879, 0.06866802275180817, 0.06425411254167557, 0.4572087228298187, 0.04987834766507149, 0.005650981329381466, 0.053177352994680405, 0.04739876464009285, 0.2551265060901642, 0.06654207408428192, 0.20209699869155884, 0.04737241193652153, 0.042119286954402924, 0.22778292000293732, 0.10508881509304047], [0.24632138013839722, 0.045121580362319946, 0.12561434507369995, 0.43826135993003845, 0.07532560080289841, 0.002372375223785639, 0.0398109070956707, 0.026653334498405457, 0.5938559174537659, 0.12655052542686462, 0.04707850515842438, 0.018195422366261482, 0.010826833546161652, 0.023274976760149002, 0.14916135370731354], [0.12666325271129608, 0.047387395054101944, 0.04497509077191353, 0.23918962478637695, 0.016611548140645027, 0.009305250830948353, 0.02713325433433056, 0.030590379610657692, 0.4573454260826111, 0.17728003859519958, 0.08635216951370239, 0.05938294902443886, 0.008936652913689613, 0.028742672875523567, 0.15077541768550873], [0.03701020032167435, 0.037774376571178436, 0.1161394715309143, 0.09335700422525406, 0.015312368050217628, 0.026739761233329773, 0.013009096495807171, 0.005902147851884365, 0.07189750671386719, 0.00625182269141078, 0.056744903326034546, 0.06423129141330719, 0.06661844998598099, 0.02100159414112568, 0.2252311259508133]], [[0.0034671342000365257, 0.05013812705874443, 0.16192083060741425, 0.3595426082611084, 0.20735634863376617, 0.08139260113239288, 0.009979248046875, 0.05037669837474823, 0.0023427342530339956, 6.08037480560597e-05, 0.003484810469672084, 0.023961462080478668, 0.38460296392440796, 0.24992075562477112, 0.13989195227622986], [0.6699675917625427, 0.09382463991641998, 0.2939082980155945, 0.17940783500671387, 0.06414232403039932, 0.05161670595407486, 0.09315118193626404, 0.0025183490943163633, 0.0024716362822800875, 0.00784118939191103, 0.06077995523810387, 0.010742363519966602, 0.027031319215893745, 0.033606547862291336, 0.020909229293465614], [0.2646949589252472, 0.029353437945246696, 0.21451972424983978, 0.10881441831588745, 0.06597915291786194, 0.0030848400201648474, 0.011694483458995819, 0.021679535508155823, 0.002872215351089835, 0.013158812187612057, 0.002100167330354452, 6.679360376438126e-05, 0.004520595073699951, 0.019191764295101166, 0.15631338953971863], [0.040224652737379074, 0.02035309188067913, 0.3179875612258911, 0.11730892956256866, 0.5032125115394592, 0.4173433780670166, 0.2045394331216812, 0.3468436896800995, 0.0142394183203578, 0.034110911190509796, 0.0166803989559412, 0.0005183254834264517, 0.014372344128787518, 0.013749183155596256, 0.07609989494085312], [0.0153636634349823, 0.002009550342336297, 0.5970484614372253, 0.5668097734451294, 0.03708057850599289, 0.030387206003069878, 0.003990367520600557, 0.00021067907800897956, 0.0006718098884448409, 0.004241611808538437, 0.01157804112881422, 0.0002699779870454222, 0.0015558624872937799, 0.0029094237834215164, 0.04601351544260979], [0.03574535250663757, 0.009626551531255245, 0.4402237832546234, 0.2294078767299652, 0.26443710923194885, 0.01504121907055378, 0.016090886667370796, 0.007329131942242384, 0.002309221774339676, 0.0030864060390740633, 0.0026519321836531162, 0.0004272839578334242, 0.0011082548880949616, 0.01614256016910076, 0.03275791555643082], [6.553631828865036e-05, 0.000357702374458313, 0.08750326931476593, 0.01436514500528574, 0.006815748754888773, 0.6623476147651672, 0.0034670215100049973, 0.0015547194052487612, 0.00029766204534098506, 1.8653441657079384e-05, 0.0003687080170493573, 0.00015007570618763566, 0.0009929342195391655, 0.00030579339363612235, 0.0016504023224115372], [0.0004548979632090777, 7.145033305278048e-05, 0.025678247213363647, 0.00989772193133831, 0.007979623042047024, 0.6904858946800232, 0.04177143797278404, 0.0005172804230824113, 0.00045151059748604894, 9.678980859462172e-05, 0.0003766386944334954, 0.00020437331113498658, 0.0009936039568856359, 0.0004823105991818011, 0.001104293274693191], [0.02770741656422615, 0.15481999516487122, 0.0164713803678751, 0.029219333082437515, 0.01727348566055298, 0.0033895254600793123, 0.08395758271217346, 0.08886045962572098, 0.06561290472745895, 0.23454923927783966, 0.01131775975227356, 0.00014876923523843288, 0.021633606404066086, 0.032435301691293716, 0.2441566288471222], [0.0002423129917588085, 0.0011915951035916805, 0.0022339578717947006, 0.006169029977172613, 0.0026169228367507458, 0.006970150861889124, 0.0023872333113104105, 0.020186979323625565, 0.5034035444259644, 0.061859097331762314, 0.01802009530365467, 0.08541904389858246, 0.11395227909088135, 0.12879255414009094, 0.06123032420873642], [0.0016445622313767672, 0.0006882954621687531, 0.0003155411686748266, 0.0014561355346813798, 0.0007120753289200366, 0.00010650769399944693, 0.0005508221802301705, 0.004306118004024029, 0.4519909620285034, 0.2298276424407959, 0.04858560487627983, 0.008956322446465492, 0.005770590156316757, 0.011063157580792904, 0.0306133683770895], [0.0032223593443632126, 0.0006265831179916859, 0.002176017500460148, 0.010606854222714901, 0.0010762742022052407, 6.259929068619385e-05, 0.0013370343949645758, 0.0014808439882472157, 0.030783534049987793, 0.7491747736930847, 0.34058046340942383, 0.00350938574410975, 0.02303031086921692, 0.0742756798863411, 0.006112673785537481], [0.010601752437651157, 0.009935700334608555, 0.0694134384393692, 0.14514312148094177, 0.01701076701283455, 0.0001025431411108002, 0.003628269536420703, 0.007610301487147808, 0.1447119563817978, 0.2691461443901062, 0.7685887217521667, 0.06739932298660278, 0.05600086599588394, 0.567065417766571, 0.01997430995106697], [0.0020818221382796764, 0.006225256249308586, 0.007747206371277571, 0.02054281160235405, 0.00644321832805872, 0.00019787036580964923, 0.0007576930802315474, 0.0013290452770888805, 0.1748982071876526, 0.20870953798294067, 0.6057864427566528, 0.2165842056274414, 0.10265108197927475, 0.12960675358772278, 0.026959752663969994], [0.0929064005613327, 0.3412420153617859, 0.13197122514247894, 0.20421825349330902, 0.6308890581130981, 0.08085004985332489, 0.35388287901878357, 0.3416491150856018, 0.024628864601254463, 0.013967287726700306, 0.0762757882475853, 0.26007020473480225, 0.3328040838241577, 0.09019435197114944, 0.014360385946929455]], [[0.014275058172643185, 0.006687531713396311, 0.3026585280895233, 0.06917963922023773, 0.2396276444196701, 0.6229325532913208, 0.15904799103736877, 0.13992713391780853, 0.10272591561079025, 0.6685669422149658, 0.22624024748802185, 0.09492585808038712, 0.40837499499320984, 0.2735627591609955, 0.011893448419868946], [0.021194536238908768, 0.020265106111764908, 0.1736137419939041, 0.08712188154459, 0.3174395263195038, 0.3545694649219513, 0.3640749752521515, 0.11553992331027985, 0.3069344758987427, 0.7487083673477173, 0.45964598655700684, 0.41950592398643494, 0.6157799363136292, 0.47228363156318665, 0.04039919748902321], [0.008898869156837463, 0.002019912237301469, 0.021509699523448944, 0.0182319525629282, 0.07474909722805023, 0.02385670319199562, 0.013716273009777069, 0.008799813687801361, 0.3437807857990265, 0.008914400823414326, 0.012629772536456585, 0.10342472046613693, 0.0370708666741848, 0.023541903123259544, 0.18654775619506836], [0.01223641075193882, 0.003142833709716797, 0.006001354195177555, 0.003996475599706173, 0.0579916350543499, 0.01896491087973118, 0.01948327198624611, 0.013184066861867905, 0.30560916662216187, 0.015957718715071678, 0.016950437799096107, 0.06207568570971489, 0.044481322169303894, 0.01894378289580345, 0.19150091707706451], [0.003971019294112921, 0.0012432326329872012, 0.005908531602472067, 0.0021760377567261457, 0.002044213702902198, 0.01004379615187645, 0.01574278064072132, 0.026324355974793434, 0.4105670154094696, 0.05117517337203026, 0.02775881439447403, 0.023424910381436348, 0.009920927695930004, 0.011210974305868149, 0.16597995162010193], [0.007421860471367836, 0.006305157672613859, 0.011464249342679977, 0.020268600434064865, 0.025753991678357124, 0.031131377443671227, 0.03418951481580734, 0.0052986773662269115, 0.5788748264312744, 0.46168622374534607, 0.07252157479524612, 0.06022901460528374, 0.017210712656378746, 0.04054110497236252, 0.15131165087223053], [0.001541785546578467, 0.0008907613810151815, 0.004846525378525257, 0.001811343478038907, 0.0069520194083452225, 0.008084121160209179, 0.021458715200424194, 0.02802192233502865, 0.3832707405090332, 0.25552085041999817, 0.014592574909329414, 0.01065820176154375, 0.012523604556918144, 0.010731800459325314, 0.22416816651821136], [0.004116748925298452, 0.0016883857315406203, 0.014749680645763874, 0.00869818776845932, 0.01003838051110506, 0.007631313521414995, 0.02068890631198883, 0.027104953303933144, 0.13497500121593475, 0.6378710865974426, 0.10288828611373901, 0.0942029282450676, 0.028772620484232903, 0.05935161933302879, 0.21764545142650604], [0.06222981959581375, 0.01881357654929161, 0.00486758491024375, 0.015509632416069508, 0.0009378677350468934, 0.004574655555188656, 0.005093523766845465, 0.0076056248508393764, 0.02507362887263298, 0.02107030339539051, 0.007815904915332794, 0.010442771948873997, 0.011698074638843536, 0.006942160427570343, 0.31572407484054565], [0.01727244071662426, 0.009210732765495777, 0.005953751504421234, 0.0013454181607812643, 0.005081892944872379, 0.04435739293694496, 0.006434922106564045, 0.0007962443050928414, 0.0007702711154706776, 0.16453301906585693, 0.5625144839286804, 0.34227296710014343, 0.6355522871017456, 0.6161591410636902, 0.02771596610546112], [0.12786830961704254, 0.008172453381121159, 0.0017843057867139578, 0.004017683211714029, 0.007877650670707226, 0.0018398476531729102, 0.01566770300269127, 0.0026914728805422783, 0.0035052604507654905, 0.0037441153544932604, 0.011492998339235783, 0.10472051054239273, 0.01954079605638981, 0.025050928816199303, 0.24727097153663635], [0.1465907245874405, 0.037033673375844955, 0.013877127319574356, 0.00413108617067337, 0.00966043584048748, 0.02326187677681446, 0.04576379433274269, 0.010370912030339241, 0.05009477958083153, 0.002161832293495536, 0.012562266550958157, 0.08835282921791077, 0.018735390156507492, 0.07781965285539627, 0.21298982203006744], [0.018177246674895287, 0.009594686329364777, 0.010616189800202847, 0.003939185757189989, 0.020018288865685463, 0.006944165099412203, 0.014553648419678211, 0.014575640670955181, 0.031773608177900314, 0.0201406329870224, 0.008282337337732315, 0.02822018228471279, 0.008926213718950748, 0.030271533876657486, 0.18345791101455688], [0.029857823625206947, 0.018949948251247406, 0.0061294399201869965, 0.002908851485699415, 0.00919707678258419, 0.00952958408743143, 0.01205661240965128, 0.00758303003385663, 0.05086279660463333, 0.007759919855743647, 0.006360263098031282, 0.02717713639140129, 0.006157578434795141, 0.027468249201774597, 0.21562480926513672], [0.035946138203144073, 0.021175134927034378, 0.025809520855545998, 0.0228139478713274, 0.02454732172191143, 0.008901212364435196, 0.01817207969725132, 0.024075007066130638, 0.042662542313337326, 0.10151555389165878, 0.03429628908634186, 0.025050567463040352, 0.015684176236391068, 0.028640326112508774, 0.23519039154052734]], [[0.29903000593185425, 0.5539957880973816, 0.06723504513502121, 0.06922264397144318, 0.12363186478614807, 0.04431891441345215, 0.10694187879562378, 0.08094406872987747, 0.15170463919639587, 0.05897890776395798, 0.026665056124329567, 0.04277891665697098, 0.011532573029398918, 0.016366619616746902, 0.08233406394720078], [0.030788322910666466, 0.06814564764499664, 0.1441766321659088, 0.42568475008010864, 0.23481200635433197, 0.09723259508609772, 0.20801249146461487, 0.2833361029624939, 0.12989479303359985, 0.09075285494327545, 0.02217184565961361, 0.10632100701332092, 0.07123817503452301, 0.18399499356746674, 0.11842577904462814], [0.21215111017227173, 0.2570435404777527, 0.03298918902873993, 0.11753708124160767, 0.2531988024711609, 0.2834656238555908, 0.13087181746959686, 0.14389817416667938, 0.06408312171697617, 0.023736948147416115, 0.043677639216184616, 0.007582403719425201, 0.08098249137401581, 0.042930904775857925, 0.09848955273628235], [0.24232596158981323, 0.4370230436325073, 0.27921250462532043, 0.32216426730155945, 0.14763100445270538, 0.1446210741996765, 0.041608523577451706, 0.05782362446188927, 0.03667302429676056, 0.015881532803177834, 0.09886573255062103, 0.0007486737449653447, 0.022804880514740944, 0.01436265092343092, 0.04328664019703865], [0.0417991504073143, 0.06808368116617203, 0.22980956733226776, 0.06044253334403038, 0.09120408445596695, 0.3664403557777405, 0.01738058589398861, 0.026107804849743843, 0.16878005862236023, 0.007388730999082327, 0.6907519698143005, 0.00283504044637084, 0.004864559043198824, 0.017621232196688652, 0.04920867085456848], [0.07025078684091568, 0.08007846027612686, 0.18737106025218964, 0.08649075031280518, 0.14398247003555298, 0.03926409035921097, 0.10999412834644318, 0.10028164088726044, 0.2733333110809326, 0.07497494667768478, 0.6277027726173401, 0.03760387748479843, 0.07242996245622635, 0.04469411447644234, 0.0635850802063942], [0.18292218446731567, 0.29889917373657227, 0.16216641664505005, 0.041324593126773834, 0.08738134056329727, 0.03374062106013298, 0.10780933499336243, 0.1685270518064499, 0.3661736249923706, 0.13795819878578186, 0.7607439160346985, 0.022037923336029053, 0.11896573007106781, 0.017960727214813232, 0.09792909026145935], [0.29104405641555786, 0.7119240164756775, 0.16990531980991364, 0.02345188707113266, 0.15646961331367493, 0.008449066430330276, 0.06418811529874802, 0.018176060169935226, 0.3091927766799927, 0.08911041170358658, 0.3005200922489166, 0.04236089810729027, 0.2996547222137451, 0.08733220398426056, 0.07523740082979202], [0.046947941184043884, 0.14375551044940948, 0.004344047512859106, 0.0067795743234455585, 0.02948000282049179, 0.08397668600082397, 0.06400846689939499, 0.18865461647510529, 0.023663662374019623, 0.08527978509664536, 0.02815503440797329, 0.04117048531770706, 0.5833349823951721, 0.0677085593342781, 0.23153413832187653], [0.08349642902612686, 0.4532567262649536, 0.004409583285450935, 0.009004302322864532, 0.007938031107187271, 0.13749390840530396, 0.1858609914779663, 0.31525370478630066, 0.018453413620591164, 0.12712040543556213, 0.04680929332971573, 0.12408707290887833, 0.13737666606903076, 0.12311573326587677, 0.142713725566864], [0.05042501911520958, 0.07026762515306473, 0.0020696106366813183, 0.010109566152095795, 0.07710029184818268, 0.05610239878296852, 0.05948542803525925, 0.19247274100780487, 0.001940111513249576, 0.05155838653445244, 0.04620450362563133, 0.20989066362380981, 0.485702246427536, 0.4166657328605652, 0.18102103471755981], [0.09080760926008224, 0.09187275916337967, 0.012195594608783722, 0.021634280681610107, 0.019499676302075386, 0.09054076671600342, 0.11008334904909134, 0.23214302957057953, 0.0423310361802578, 0.034868963062763214, 0.06751228123903275, 0.049237679690122604, 0.03915484994649887, 0.08995199203491211, 0.1941523253917694], [0.0706457570195198, 0.10473088920116425, 0.039385173469781876, 0.02697153575718403, 0.04372800514101982, 0.06655491143465042, 0.23491710424423218, 0.19935868680477142, 0.036273516714572906, 0.06345809996128082, 0.020782677456736565, 0.12393849343061447, 0.05726756155490875, 0.041495081037282944, 0.15982753038406372], [0.039186086505651474, 0.11076691001653671, 0.03891725465655327, 0.009549588896334171, 0.01825849525630474, 0.051163915544748306, 0.1146436408162117, 0.1649821698665619, 0.03586947172880173, 0.06679365783929825, 0.09092967957258224, 0.14827685058116913, 0.10948126018047333, 0.10746686905622482, 0.1515202671289444], [0.14541134238243103, 0.05313154682517052, 0.01991144008934498, 0.08764121681451797, 0.014597749337553978, 0.03937898576259613, 0.04872390255331993, 0.04689335823059082, 0.04558950290083885, 0.051970891654491425, 0.02520112879574299, 0.022838978096842766, 0.00921469647437334, 0.00801294855773449, 0.21471147239208221]], [[0.009874092414975166, 0.0475393682718277, 0.0700187012553215, 0.05995699018239975, 0.023110831156373024, 0.04304451867938042, 0.02397323027253151, 0.09104450792074203, 0.13320927321910858, 0.0718994140625, 0.16378211975097656, 0.06306017935276031, 0.03516274318099022, 0.06407153606414795, 0.1927335411310196], [0.007679122034460306, 0.008519956842064857, 0.023641018196940422, 0.036320336163043976, 0.005810021422803402, 0.002834178740158677, 0.01027101743966341, 0.005131446290761232, 0.05288401618599892, 0.022729018703103065, 0.02885960415005684, 0.007142365910112858, 0.005423326510936022, 0.00592823838815093, 0.23125353455543518], [0.17363575100898743, 0.08529574424028397, 0.018747013062238693, 0.09323837608098984, 0.07366655766963959, 0.2784116566181183, 0.6226999759674072, 0.6422466039657593, 0.18433590233325958, 0.44911590218544006, 0.07703087478876114, 0.23628254234790802, 0.37835898995399475, 0.3362680971622467, 0.10061702132225037], [0.039354946464300156, 0.028671007603406906, 0.0009692042949609458, 0.010166235268115997, 0.003592043649405241, 0.024686597287654877, 0.0576656274497509, 0.10543617606163025, 0.069565050303936, 0.23999209702014923, 0.0370241142809391, 0.07099387794733047, 0.08031197637319565, 0.0629396140575409, 0.19831009209156036], [0.07821620255708694, 0.07413192838430405, 0.008470119908452034, 0.005837618373334408, 0.016890503466129303, 0.34118980169296265, 0.6424257159233093, 0.5736639499664307, 0.18751046061515808, 0.08286380022764206, 0.013973995111882687, 0.16452431678771973, 0.6265572905540466, 0.24633896350860596, 0.03771306574344635], [0.08601168543100357, 0.11519530415534973, 0.00501672737300396, 0.0384475477039814, 0.0009856059914454818, 0.020220156759023666, 0.4602939486503601, 0.41334664821624756, 0.011432202532887459, 0.039776530116796494, 0.004202698357403278, 0.012451107613742352, 0.012797003611922264, 0.0109980758279562, 0.22371669113636017], [0.05821564793586731, 0.2493630200624466, 0.017187682911753654, 0.007334073074162006, 0.002277297666296363, 0.012770043686032295, 0.014771709218621254, 0.06810285151004791, 0.008148171938955784, 0.093966543674469, 0.03078475221991539, 0.016961626708507538, 0.009818210266530514, 0.005369590129703283, 0.2805846929550171], [0.0315314382314682, 0.006441309116780758, 0.005187691655009985, 0.0023020647931843996, 0.001103160553611815, 0.0010285694152116776, 0.0036586276255548, 0.0034369472414255142, 0.02540425956249237, 0.018933216109871864, 0.011261656880378723, 0.014689027331769466, 0.0047272746451199055, 0.003173592034727335, 0.27608010172843933], [0.052501752972602844, 0.03902341425418854, 0.022159013897180557, 0.15980832278728485, 0.04565480723977089, 0.04961955174803734, 0.10487794876098633, 0.03556728735566139, 0.011893571354448795, 0.350600004196167, 0.8153157234191895, 0.696418821811676, 0.19642634689807892, 0.7945331335067749, 0.025074943900108337], [0.008775658905506134, 0.0231929961591959, 0.001974506536498666, 0.02221933752298355, 0.002016729209572077, 0.03464629501104355, 0.020560195669531822, 0.015741808339953423, 0.024821357801556587, 0.03194829449057579, 0.062133170664310455, 0.009445058181881905, 0.008440939709544182, 0.031038939952850342, 0.24359388649463654], [0.15448324382305145, 0.15535393357276917, 0.0009195139864459634, 0.02347545325756073, 0.010745828039944172, 0.05933469906449318, 0.0886014774441719, 0.09891750663518906, 0.008176282048225403, 0.17814745008945465, 0.04613054543733597, 0.10348650068044662, 0.06132601201534271, 0.10257216542959213, 0.2144334316253662], [0.1637454628944397, 0.3587695062160492, 0.013175190426409245, 0.027070751413702965, 0.009701711125671864, 0.027045298367738724, 0.06057014688849449, 0.08674251288175583, 0.018084047362208366, 0.012978773564100266, 0.04984384402632713, 0.0746963769197464, 0.21545591950416565, 0.18275731801986694, 0.18403297662734985], [0.04016833007335663, 0.03071952983736992, 0.0073937661945819855, 0.044594794511795044, 0.005693770945072174, 0.007929249666631222, 0.19023852050304413, 0.12198647856712341, 0.00967123731970787, 0.05747445672750473, 0.006795276887714863, 0.006636326666921377, 0.014849998988211155, 0.02297961339354515, 0.1823122203350067], [0.08359953761100769, 0.14515268802642822, 0.009139984846115112, 0.10055579245090485, 0.007817201316356659, 0.06191832944750786, 0.24591712653636932, 0.26670339703559875, 0.008127851411700249, 0.05132465437054634, 0.011226493865251541, 0.020721180364489555, 0.025672290474176407, 0.06137499585747719, 0.19538666307926178], [0.004038439132273197, 0.01158715970814228, 0.012492671608924866, 0.008604439906775951, 0.0044732466340065, 0.001471644383855164, 0.003622728632763028, 0.005392232909798622, 0.024040954187512398, 0.002572751836851239, 0.011896335519850254, 0.00655994052067399, 0.004419950768351555, 0.0023605322930961847, 0.2578853368759155]], [[0.020951254293322563, 0.19576001167297363, 0.05422525107860565, 0.000516751199029386, 0.0576050765812397, 0.039616964757442474, 0.0011584623716771603, 0.06260760873556137, 0.05524995177984238, 5.760174462920986e-05, 0.0005486492882482708, 0.01856253668665886, 0.008022493682801723, 0.0032547120936214924, 0.1980074942111969], [0.15878187119960785, 0.5755441188812256, 0.073322594165802, 0.006848999299108982, 0.04221894592046738, 0.057610929012298584, 0.01498481910675764, 0.15564584732055664, 0.02557745948433876, 0.010493909008800983, 0.04444737732410431, 0.10564734041690826, 0.04703369736671448, 0.007807346060872078, 0.10371111333370209], [0.0667557343840599, 0.5756934881210327, 0.02783285267651081, 0.001271323417313397, 0.13096383213996887, 0.007863562554121017, 0.0004880728665739298, 0.00786207988858223, 0.030193913727998734, 0.0004458925104700029, 0.0008183285826817155, 0.003005507169291377, 0.008833326399326324, 0.014566708356142044, 0.09050195664167404], [0.006902126595377922, 0.22582471370697021, 0.027240794152021408, 0.000252248632023111, 0.08146748691797256, 0.008376134559512138, 0.0017193618696182966, 0.010283069685101509, 0.09191752970218658, 1.873078872449696e-05, 0.0001427968527423218, 0.0006295929779298604, 0.016630304977297783, 0.005029548890888691, 0.17517179250717163], [0.46813952922821045, 0.7474208474159241, 0.04419572278857231, 0.039987821131944656, 0.07900705188512802, 0.010286353528499603, 0.008277984336018562, 0.21022778749465942, 0.018339863047003746, 0.003122991183772683, 0.0047759185545146465, 0.0031952662393450737, 0.0037801233120262623, 0.005526377819478512, 0.11187370121479034], [0.08057912439107895, 0.09254536032676697, 0.26037144660949707, 0.04459136351943016, 0.19053104519844055, 0.18187369406223297, 0.04494835063815117, 0.08866222947835922, 0.05515718460083008, 0.011219717562198639, 0.041749756783246994, 0.13417255878448486, 0.43527963757514954, 0.4240920841693878, 0.05903848633170128], [0.005677447654306889, 0.1104632169008255, 0.17886187136173248, 0.06816153228282928, 0.31320425868034363, 0.08580746501684189, 0.044242095202207565, 0.4031389355659485, 0.13310441374778748, 8.991359209176153e-05, 0.00051962147699669, 0.017516016960144043, 0.02517649158835411, 0.02827705629169941, 0.13873830437660217], [0.009441166184842587, 0.04568161070346832, 0.08503290265798569, 0.055850934237241745, 0.15800173580646515, 0.09921947866678238, 0.2719998359680176, 0.7131122350692749, 0.12690743803977966, 0.0015569856623187661, 0.019959524273872375, 0.06398878246545792, 0.1124982088804245, 0.07506788522005081, 0.06075114384293556], [0.1778930425643921, 0.41812169551849365, 0.05459700897336006, 0.015388981439173222, 0.296997606754303, 0.041353121399879456, 0.1696915328502655, 0.1226804181933403, 0.3453136682510376, 0.006036087870597839, 0.008416525088250637, 0.004891113843768835, 0.003974124789237976, 0.0023401544895023108, 0.04184575751423836], [0.0018550200620666146, 0.2628808617591858, 0.0018376001389697194, 9.925621998263523e-05, 0.008250601589679718, 0.11965687572956085, 0.011913565918803215, 0.3649533987045288, 0.12527383863925934, 0.0011617891723290086, 0.002173396060243249, 0.011088940314948559, 0.02579125389456749, 0.004398738034069538, 0.18079015612602234], [0.0033212341368198395, 0.4786561131477356, 0.00019389556837268174, 4.100392834516242e-05, 0.03255903348326683, 0.004482456482946873, 0.0018638258334249258, 0.04032744839787483, 0.151435986161232, 0.0011174781247973442, 0.0008650964009575546, 0.049343932420015335, 0.013284855522215366, 0.009702197276055813, 0.17111515998840332], [0.015286837704479694, 0.17760051786899567, 0.012107143178582191, 0.004069492220878601, 0.40114596486091614, 0.005856915842741728, 0.025313973426818848, 0.23595470190048218, 0.5599475502967834, 0.019674712792038918, 0.01789786107838154, 0.0449712835252285, 0.024323459714651108, 0.008310162462294102, 0.10516723990440369], [0.013816175982356071, 0.10832668840885162, 0.014126134105026722, 0.0044770012609660625, 0.18972823023796082, 0.04144473373889923, 0.013167506083846092, 0.0398833267390728, 0.08117146790027618, 0.03379456326365471, 0.04336484149098396, 0.6766878366470337, 0.6025072932243347, 0.24042664468288422, 0.05677386373281479], [0.010657100938260555, 0.1729527860879898, 0.006031150463968515, 0.006062258500605822, 0.10042858123779297, 0.007653414737433195, 0.0031583579257130623, 0.014785557985305786, 0.13275322318077087, 0.05689838156104088, 0.04302775487303734, 0.36964303255081177, 0.3870774507522583, 0.31299954652786255, 0.07590257376432419], [0.014769526198506355, 0.05199434980750084, 0.11582475155591965, 0.14804258942604065, 0.05702318996191025, 0.3275434374809265, 0.3759170472621918, 0.3329218327999115, 0.027774346992373466, 0.12548163533210754, 0.13219930231571198, 0.029332099482417107, 0.2028164267539978, 0.518939197063446, 4.3280975660309196e-05]], [[0.5917359590530396, 0.12410512566566467, 0.24872945249080658, 0.20040015876293182, 0.21720361709594727, 0.11561702191829681, 0.58521568775177, 0.41413450241088867, 0.22558750212192535, 0.117314413189888, 0.3378458619117737, 0.10710897296667099, 0.0625920221209526, 0.24034489691257477, 0.0060951621271669865], [0.03933318331837654, 0.17479471862316132, 0.1999012678861618, 0.1507989913225174, 0.2344110906124115, 0.41628938913345337, 0.19733835756778717, 0.42009472846984863, 0.32125937938690186, 0.09302358329296112, 0.29758843779563904, 0.2500022351741791, 0.15192696452140808, 0.19621950387954712, 0.06078135594725609], [0.03998054191470146, 0.02165106125175953, 0.5779209733009338, 0.4094802737236023, 0.3219829499721527, 0.23359909653663635, 0.15223096311092377, 0.0776560828089714, 0.11850404739379883, 0.1752316802740097, 0.7765606641769409, 0.15624035894870758, 0.19448350369930267, 0.3389243483543396, 0.015656093135476112], [0.2606712579727173, 0.23122362792491913, 0.33188652992248535, 0.327752023935318, 0.0930425301194191, 0.13157396018505096, 0.5079332590103149, 0.15524731576442719, 0.2039693295955658, 0.336448073387146, 0.7406277656555176, 0.11173539608716965, 0.03980698063969612, 0.2757716476917267, 0.009055807255208492], [0.03992704302072525, 0.03562299162149429, 0.05761631205677986, 0.04593607783317566, 0.747100830078125, 0.13848423957824707, 0.25807130336761475, 0.11098858714103699, 0.025020861998200417, 0.027831630781292915, 0.07712040096521378, 0.5344594120979309, 0.28488224744796753, 0.37143638730049133, 0.060307834297418594], [0.146702840924263, 0.5779150128364563, 0.04704871401190758, 0.12512727081775665, 0.05839477851986885, 0.5817644596099854, 0.2541782557964325, 0.167904794216156, 0.020014837384223938, 0.0557471327483654, 0.1778557300567627, 0.29983726143836975, 0.34978994727134705, 0.3759990334510803, 0.07532685250043869], [0.14372284710407257, 0.20398879051208496, 0.060162752866744995, 0.022449441254138947, 0.15882903337478638, 0.12907396256923676, 0.7781419157981873, 0.20689332485198975, 0.023098474368453026, 0.02567201852798462, 0.04225016012787819, 0.05647281929850578, 0.5644452571868896, 0.8062969446182251, 0.0037398021668195724], [0.09274263679981232, 0.19406189024448395, 0.18035270273685455, 0.18292436003684998, 0.2674761116504669, 0.1057504341006279, 0.5214765071868896, 0.1765710562467575, 0.15375129878520966, 0.08563723415136337, 0.35003283619880676, 0.12250327318906784, 0.4574505388736725, 0.6043637990951538, 0.046846963465213776], [0.3136129081249237, 0.10648278146982193, 0.02492944709956646, 0.07937752455472946, 0.16382691264152527, 0.40212482213974, 0.2148500233888626, 0.5046796798706055, 0.25625455379486084, 0.10382789373397827, 0.027611082419753075, 0.07138189673423767, 0.1265101283788681, 0.05298655480146408, 0.01642199046909809], [0.7252353429794312, 0.23862500488758087, 0.17466871440410614, 0.2584758698940277, 0.15821219980716705, 0.41019105911254883, 0.4795793294906616, 0.2558479905128479, 0.061036378145217896, 0.5831483006477356, 0.23237691819667816, 0.36767491698265076, 0.07294586300849915, 0.0734395682811737, 0.006080146878957748], [0.18402060866355896, 0.2199273407459259, 0.10670217871665955, 0.36498934030532837, 0.37264159321784973, 0.5975290536880493, 0.641157865524292, 0.4798426032066345, 0.07047704607248306, 0.30389490723609924, 0.6835307478904724, 0.29959914088249207, 0.32009243965148926, 0.2076108753681183, 0.015385132282972336], [0.18547095358371735, 0.1046445369720459, 0.17664410173892975, 0.031107882037758827, 0.4872691333293915, 0.6876094937324524, 0.29805243015289307, 0.2697339355945587, 0.03289056569337845, 0.04577193781733513, 0.2390383929014206, 0.650258481502533, 0.6253164410591125, 0.2719551920890808, 0.042574722319841385], [0.06026101112365723, 0.4596063494682312, 0.11362233757972717, 0.050736263394355774, 0.47900232672691345, 0.8146356344223022, 0.23428170382976532, 0.5258204936981201, 0.07407079637050629, 0.24087238311767578, 0.04631686583161354, 0.04097185283899307, 0.24002470076084137, 0.051092784851789474, 0.10185284167528152], [0.05915316566824913, 0.3385859429836273, 0.23845957219600677, 0.13520635664463043, 0.49372056126594543, 0.8321547508239746, 0.47351959347724915, 0.4942004382610321, 0.11661165207624435, 0.273796945810318, 0.09639480710029602, 0.07113680988550186, 0.3545372784137726, 0.3069557547569275, 0.026768943294882774], [0.6326229572296143, 0.28129494190216064, 0.2424720972776413, 0.23961131274700165, 0.1532977670431137, 0.03248026221990585, 0.07237446308135986, 0.03991716355085373, 0.058106135576963425, 0.6791825294494629, 0.4868316352367401, 0.4841252863407135, 0.1838759332895279, 0.16229771077632904, 0.03779346123337746]], [[0.04456469416618347, 0.016716457903385162, 0.08688971400260925, 0.23432573676109314, 0.12769784033298492, 0.0498066172003746, 0.10501405596733093, 0.14398211240768433, 0.3055479824542999, 0.0823235884308815, 0.23467087745666504, 0.6305257678031921, 0.08790664374828339, 0.14063040912151337, 0.13028757274150848], [0.04107241332530975, 0.03620494529604912, 0.07322828471660614, 0.1027759537100792, 0.08743055909872055, 0.016458408907055855, 0.09779228270053864, 0.014780157245695591, 0.09821301698684692, 0.025402111932635307, 0.0808086097240448, 0.08257035166025162, 0.07231960445642471, 0.0895148441195488, 0.19708459079265594], [0.1263897716999054, 0.01533158216625452, 0.08717449009418488, 0.22571881115436554, 0.06928549706935883, 0.16778334975242615, 0.06136450543999672, 0.07180161774158478, 0.2525678873062134, 0.32249853014945984, 0.08566119521856308, 0.48726531863212585, 0.2929263114929199, 0.21127133071422577, 0.12448348850011826], [0.1481804996728897, 0.04817945510149002, 0.03058626689016819, 0.13171793520450592, 0.10783855617046356, 0.24912205338478088, 0.1342363804578781, 0.28650397062301636, 0.25943103432655334, 0.2756144404411316, 0.08422903716564178, 0.7444766163825989, 0.7611673474311829, 0.5739472508430481, 0.11213001608848572], [0.1744699776172638, 0.050404343754053116, 0.018338145688176155, 0.11463086307048798, 0.02370826154947281, 0.09417468309402466, 0.04503462836146355, 0.0389062762260437, 0.1780962496995926, 0.7825090885162354, 0.15977078676223755, 0.2598268687725067, 0.05674973130226135, 0.2742767333984375, 0.15589554607868195], [0.26428407430648804, 0.0871720165014267, 0.015494171530008316, 0.31054598093032837, 0.31179672479629517, 0.05687993764877319, 0.05327969416975975, 0.14049863815307617, 0.03721972927451134, 0.33735793828964233, 0.06669215857982635, 0.44665512442588806, 0.1105320155620575, 0.07633788883686066, 0.13637836277484894], [0.27871736884117126, 0.07987862080335617, 0.06999076902866364, 0.3873903453350067, 0.3669894337654114, 0.0245819091796875, 0.02483827993273735, 0.08571609854698181, 0.04856930300593376, 0.2826782464981079, 0.10519464313983917, 0.8515737056732178, 0.24991582334041595, 0.08752243965864182, 0.1076057106256485], [0.18780259788036346, 0.02093103528022766, 0.1730981320142746, 0.27918383479118347, 0.32355740666389465, 0.05090703070163727, 0.030107326805591583, 0.015694553032517433, 0.08293543756008148, 0.11989035457372665, 0.1594303995370865, 0.6402391195297241, 0.08334839344024658, 0.13423335552215576, 0.16886292397975922], [0.23048973083496094, 0.05534357205033302, 0.15910016000270844, 0.5473513603210449, 0.11114095151424408, 0.060548413544893265, 0.23547381162643433, 0.0231330469250679, 0.22654443979263306, 0.16574865579605103, 0.03383632004261017, 0.05167527496814728, 0.026772163808345795, 0.028301218524575233, 0.08144620060920715], [0.126570925116539, 0.0055835917592048645, 0.7687394022941589, 0.6136845350265503, 0.7887718677520752, 0.24027548730373383, 0.25543272495269775, 0.017155619338154793, 0.01121050026267767, 0.02180907502770424, 0.06387564539909363, 0.04227403923869133, 0.004662328865379095, 0.0204116590321064, 0.16526305675506592], [0.3619309663772583, 0.022692076861858368, 0.8739812970161438, 0.5600091814994812, 0.4330839216709137, 0.27864721417427063, 0.1654776781797409, 0.02327956072986126, 0.003977042157202959, 0.0664801374077797, 0.12084753066301346, 0.16815124452114105, 0.07773539423942566, 0.17824198305606842, 0.05263833701610565], [0.29354482889175415, 0.16078433394432068, 0.705570638179779, 0.44417092204093933, 0.02176845259964466, 0.15997210144996643, 0.4057019054889679, 0.11617531627416611, 0.010741903446614742, 0.06882698833942413, 0.07046788930892944, 0.041601523756980896, 0.011864392086863518, 0.06714706867933273, 0.14988133311271667], [0.5400083065032959, 0.2319646179676056, 0.6198285818099976, 0.2858767509460449, 0.1694929450750351, 0.06001640111207962, 0.26940232515335083, 0.06411167979240417, 0.02847147174179554, 0.18856319785118103, 0.05879069119691849, 0.03795049339532852, 0.009596540592610836, 0.023393897339701653, 0.14663995802402496], [0.6488012075424194, 0.15997910499572754, 0.6486002802848816, 0.4859846830368042, 0.34752336144447327, 0.028076842427253723, 0.12281371653079987, 0.019826101139187813, 0.023531395941972733, 0.15743687748908997, 0.059922393411397934, 0.08707788586616516, 0.005486410576850176, 0.025385212153196335, 0.15706156194210052], [0.037294961512088776, 0.2018004208803177, 0.33537882566452026, 0.19571122527122498, 0.0998593419790268, 0.48263466358184814, 0.11429780721664429, 0.20324908196926117, 0.7053001523017883, 0.01905757561326027, 0.1765546351671219, 0.10779165476560593, 0.18456625938415527, 0.16855330765247345, 0.014784654602408409]]]], \"bot_text\": [\"The_\", \"animal_\", \"didn_\", \"'_\", \"t_\", \"cross_\", \"the_\", \"street_\", \"because_\", \"it_\", \"was_\", \"too_\", \"tire\", \"d_\"]}}"
             ],
             "text/plain": [
-              "<IPython.core.display.Javascript object>"
+              "\u003cIPython.core.display.Javascript object\u003e"
             ]
           },
           "metadata": {
             "tags": []
-          }
+          },
+          "output_type": "display_data"
         },
         {
-          "output_type": "display_data",
           "data": {
             "application/javascript": [
               "/**\n",
@@ -1012,7 +908,7 @@
               "            .attr(\"height\", HEIGHT);\n",
               "\n",
               "  var att_data = [];\n",
-              "  for (var i=0; i < attention_heads.length; i++) {\n",
+              "  for (var i=0; i \u003c attention_heads.length; i++) {\n",
               "    var att_trans = transpose(attention_heads[i]);\n",
               "    att_data.push(zip(attention_heads[i], att_trans));\n",
               "  }\n",
@@ -1198,9 +1094,9 @@
               "function renderAttention(svg, attention_heads) {\n",
               "  var line_container = svg.selectAll(\".attention_heads\");\n",
               "  line_container.html(null);\n",
-              "  for(var h=0; h<attention_heads.length; h++) {\n",
-              "    for(var a=0; a<attention_heads[h].length; a++) {\n",
-              "      for(var s=0; s<attention_heads[h][a].length; s++) {\n",
+              "  for(var h=0; h\u003cattention_heads.length; h++) {\n",
+              "    for(var a=0; a\u003cattention_heads[h].length; a++) {\n",
+              "      for(var s=0; s\u003cattention_heads[h][a].length; s++) {\n",
               "        line_container.append(\"line\")\n",
               "        .attr(\"y1\", (s+1) * BOXHEIGHT + (BOXHEIGHT/2))\n",
               "        .attr(\"x1\", BOXWIDTH)\n",
@@ -1223,7 +1119,7 @@
               "// Checkboxes\n",
               "function box_offset(i) {\n",
               "  var num_head_above = config.head_vis.reduce(\n",
-              "      function(acc, val, cur) {return val && cur < i ? acc + 1: acc;}, 0);\n",
+              "      function(acc, val, cur) {return val \u0026\u0026 cur \u003c i ? acc + 1: acc;}, 0);\n",
               "  return num_head_above*(BOXWIDTH / active_heads());\n",
               "}\n",
               "\n",
@@ -1262,7 +1158,7 @@
               "  update_checkboxes();\n",
               "\n",
               "  checkbox.on(\"click\", function(d, i) {\n",
-              "    if (config.head_vis[i] && active_heads() == 1) return;\n",
+              "    if (config.head_vis[i] \u0026\u0026 active_heads() == 1) return;\n",
               "    config.head_vis[i] = !config.head_vis[i];\n",
               "    update_checkboxes();\n",
               "    renderAttention(svg, attention_heads);\n",
@@ -1270,7 +1166,7 @@
               "\n",
               "  checkbox.on(\"dblclick\", function(d, i) {\n",
               "    // If we double click on the only active head then reset\n",
-              "    if (config.head_vis[i] && active_heads() == 1) {\n",
+              "    if (config.head_vis[i] \u0026\u0026 active_heads() == 1) {\n",
               "      config.head_vis = new Array(config.num_heads).fill(true);\n",
               "    } else {\n",
               "      config.head_vis = new Array(config.num_heads).fill(false);\n",
@@ -1307,8 +1203,8 @@
               "}\n",
               "\n",
               "$(\"#layer\").empty();\n",
-              "for(var i=0; i<6; i++) {\n",
-              "  $(\"#layer\").append($(\"<option />\").val(i).text(i));\n",
+              "for(var i=0; i\u003c6; i++) {\n",
+              "  $(\"#layer\").append($(\"\u003coption /\u003e\").val(i).text(i));\n",
               "}\n",
               "\n",
               "$(\"#layer\").on('change', function(e) {\n",
@@ -1328,37 +1224,51 @@
               "});\n"
             ],
             "text/plain": [
-              "<IPython.core.display.Javascript object>"
+              "\u003cIPython.core.display.Javascript object\u003e"
             ]
           },
           "metadata": {
             "tags": []
-          }
+          },
+          "output_type": "display_data"
         }
+      ],
+      "source": [
+        "# Convert inputs and outputs to subwords\n",
+        "inp_text = to_tokens(encoders[\"inputs\"].encode(inputs))\n",
+        "out_text = to_tokens(encoders[\"inputs\"].encode(outputs))\n",
+        "\n",
+        "# Run eval to collect attention weights\n",
+        "example = encode_eval(inputs, outputs)\n",
+        "with tfe.restore_variables_on_create(tf.train.latest_checkpoint(checkpoint_dir)):\n",
+        "  translate_model.set_mode(Modes.EVAL)\n",
+        "  translate_model(example)\n",
+        "# Get normalized attention weights for each layer\n",
+        "enc_atts, dec_atts, encdec_atts = get_att_mats()\n",
+        "\n",
+        "call_html()\n",
+        "attention.show(inp_text, out_text, enc_atts, dec_atts, encdec_atts)"
       ]
     },
     {
+      "cell_type": "markdown",
       "metadata": {
-        "id": "i7BZuO7T5BB4",
-        "colab_type": "text"
+        "colab_type": "text",
+        "id": "i7BZuO7T5BB4"
       },
-      "cell_type": "markdown",
       "source": [
         "# Train a custom model on MNIST"
       ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "-H25oG91YQj3",
+        "colab": {},
         "colab_type": "code",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        }
+        "id": "-H25oG91YQj3"
       },
-      "cell_type": "code",
+      "outputs": [],
       "source": [
         "# Create your own model\n",
         "\n",
@@ -1377,36 +1287,40 @@
         "hparams = trainer_lib.create_hparams(\"basic_1\", data_dir=data_dir, problem_name=\"image_mnist\")\n",
         "hparams.hidden_size = 64\n",
         "model = MySimpleModel(hparams, Modes.TRAIN)"
-      ],
-      "execution_count": 0,
-      "outputs": []
+      ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "7GEmpYQ2ZMnB",
-        "colab_type": "code",
         "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          },
           "base_uri": "https://localhost:8080/",
           "height": 34
         },
-        "outputId": "a574a1a3-ce56-4715-9ad3-8289c61ade3b",
+        "colab_type": "code",
         "executionInfo": {
+          "elapsed": 625,
           "status": "ok",
           "timestamp": 1512369563515,
-          "user_tz": 480,
-          "elapsed": 625,
           "user": {
             "displayName": "Niki Parmar",
             "photoUrl": "//lh3.googleusercontent.com/-ReuwZvCmGE8/AAAAAAAAAAI/AAAAAAAAAIc/fcvytJVpitE/s50-c-k-no/photo.jpg",
             "userId": "115864460963462186442"
-          }
-        }
+          },
+          "user_tz": 480
+        },
+        "id": "7GEmpYQ2ZMnB",
+        "outputId": "a574a1a3-ce56-4715-9ad3-8289c61ade3b"
       },
-      "cell_type": "code",
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "INFO:tensorflow:Reading data files from /content/t2t/data/image_mnist-train*\n"
+          ]
+        }
+      ],
       "source": [
         "# Prepare for the training loop\n",
         "\n",
@@ -1423,61 +1337,34 @@
         "mnist_train_dataset = mnist_train_dataset.repeat(None).batch(BATCH_SIZE)\n",
         "\n",
         "optimizer = tf.train.AdamOptimizer()"
-      ],
-      "execution_count": 0,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "text": [
-            "INFO:tensorflow:Reading data files from /content/t2t/data/image_mnist-train*\n"
-          ],
-          "name": "stdout"
-        }
       ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "AWVd2I7PYz6H",
-        "colab_type": "code",
         "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          },
           "base_uri": "https://localhost:8080/",
           "height": 204
         },
-        "outputId": "504a7876-8bbb-4e5f-f303-f951c2e071b2",
+        "colab_type": "code",
         "executionInfo": {
+          "elapsed": 103766,
           "status": "ok",
           "timestamp": 1512369756046,
-          "user_tz": 480,
-          "elapsed": 103766,
           "user": {
             "displayName": "Niki Parmar",
             "photoUrl": "//lh3.googleusercontent.com/-ReuwZvCmGE8/AAAAAAAAAAI/AAAAAAAAAIc/fcvytJVpitE/s50-c-k-no/photo.jpg",
             "userId": "115864460963462186442"
-          }
-        }
+          },
+          "user_tz": 480
+        },
+        "id": "AWVd2I7PYz6H",
+        "outputId": "504a7876-8bbb-4e5f-f303-f951c2e071b2"
       },
-      "cell_type": "code",
-      "source": [
-        "# Train\n",
-        "NUM_STEPS = 500\n",
-        "\n",
-        "for count, example in enumerate(tfe.Iterator(mnist_train_dataset)):\n",
-        "  example[\"targets\"] = tf.reshape(example[\"targets\"], [BATCH_SIZE, 1, 1, 1])  # Make it 4D.\n",
-        "  loss, gv = loss_fn(example)\n",
-        "  optimizer.apply_gradients(gv)\n",
-        "\n",
-        "  if count % 50 == 0:\n",
-        "    print(\"Step: %d, Loss: %.3f\" % (count, loss.numpy()))\n",
-        "  if count >= NUM_STEPS:\n",
-        "    break"
-      ],
-      "execution_count": 0,
       "outputs": [
         {
+          "name": "stdout",
           "output_type": "stream",
           "text": [
             "Step: 0, Loss: 0.513\n",
@@ -1491,37 +1378,58 @@
             "Step: 400, Loss: 0.250\n",
             "Step: 450, Loss: 0.247\n",
             "Step: 500, Loss: 0.338\n"
-          ],
-          "name": "stdout"
+          ]
         }
+      ],
+      "source": [
+        "# Train\n",
+        "NUM_STEPS = 500\n",
+        "\n",
+        "for count, example in enumerate(tfe.Iterator(mnist_train_dataset)):\n",
+        "  example[\"targets\"] = tf.reshape(example[\"targets\"], [BATCH_SIZE, 1, 1, 1])  # Make it 4D.\n",
+        "  loss, gv = loss_fn(example)\n",
+        "  optimizer.apply_gradients(gv)\n",
+        "\n",
+        "  if count % 50 == 0:\n",
+        "    print(\"Step: %d, Loss: %.3f\" % (count, loss.numpy()))\n",
+        "  if count \u003e= NUM_STEPS:\n",
+        "    break"
       ]
     },
     {
+      "cell_type": "code",
+      "execution_count": 0,
       "metadata": {
-        "id": "CIFlkiVOd8jO",
-        "colab_type": "code",
         "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          },
           "base_uri": "https://localhost:8080/",
           "height": 68
         },
-        "outputId": "ef33057a-1a22-4ab8-ab7b-3c90d9f6a850",
+        "colab_type": "code",
         "executionInfo": {
+          "elapsed": 3833,
           "status": "ok",
           "timestamp": 1512369759917,
-          "user_tz": 480,
-          "elapsed": 3833,
           "user": {
             "displayName": "Niki Parmar",
             "photoUrl": "//lh3.googleusercontent.com/-ReuwZvCmGE8/AAAAAAAAAAI/AAAAAAAAAIc/fcvytJVpitE/s50-c-k-no/photo.jpg",
             "userId": "115864460963462186442"
-          }
-        }
+          },
+          "user_tz": 480
+        },
+        "id": "CIFlkiVOd8jO",
+        "outputId": "ef33057a-1a22-4ab8-ab7b-3c90d9f6a850"
       },
-      "cell_type": "code",
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "INFO:tensorflow:Reading data files from /content/t2t/data/image_mnist-dev*\n",
+            "accuracy_top5: 1.00\n",
+            "accuracy: 0.99\n"
+          ]
+        }
+      ],
       "source": [
         "model.set_mode(Modes.EVAL)\n",
         "mnist_eval_dataset = mnist_problem.dataset(Modes.EVAL, data_dir)\n",
@@ -1532,7 +1440,7 @@
         "    [metrics.Metrics.ACC, metrics.Metrics.ACC_TOP5])\n",
         "\n",
         "for count, example in enumerate(tfe.Iterator(mnist_eval_dataset)):\n",
-        "  if count >= 200:\n",
+        "  if count \u003e= 200:\n",
         "    break\n",
         "\n",
         "  # Make the inputs and targets 4D\n",
@@ -1548,19 +1456,25 @@
         "# Print out the averaged metric values on the eval data\n",
         "for name, val in metrics_result().items():\n",
         "  print(\"%s: %.2f\" % (name, val))"
-      ],
-      "execution_count": 0,
-      "outputs": [
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [],
+      "name": "Tensor2Tensor Intro",
+      "provenance": [
         {
-          "output_type": "stream",
-          "text": [
-            "INFO:tensorflow:Reading data files from /content/t2t/data/image_mnist-dev*\n",
-            "accuracy_top5: 1.00\n",
-            "accuracy: 0.99\n"
-          ],
-          "name": "stdout"
+          "file_id": "1-VScmaLkMqWiSbqgUCFWefzisSREd8l1",
+          "timestamp": 1512175750497
         }
       ]
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
     }
-  ]
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
 }
diff --git a/tensor2tensor/notebooks/t2t_problem.ipynb b/tensor2tensor/notebooks/t2t_problem.ipynb
index 592cbad39..a5d85cfa1 100644
--- a/tensor2tensor/notebooks/t2t_problem.ipynb
+++ b/tensor2tensor/notebooks/t2t_problem.ipynb
@@ -136,6 +136,7 @@
       "source": [
         "#@title Run this only once - Sets up TF Eager execution.\n",
         "\n",
+        "%tensorflow_version 1.x\n",
         "import tensorflow as tf\n",
         "\n",
         "# Enable Eager execution - useful for seeing the generated data.\n",

From e05b32dd725dfabdf1a026d2290d5022b04772f8 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 20 Dec 2019 12:39:11 -0800
Subject: [PATCH 2605/2720] Allow transformer models to have multiple
 user-defined timing signals.

PiperOrigin-RevId: 286624049
---
 tensor2tensor/layers/common_attention.py      | 71 +++++++++++++++++++
 tensor2tensor/layers/common_attention_test.py | 67 ++++++++++++++++-
 tensor2tensor/layers/transformer_layers.py    |  3 +
 tensor2tensor/models/transformer.py           | 12 ++++
 4 files changed, 151 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index a2a062fce..c00c99540 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -577,6 +577,77 @@ def add_layer_timing_signal_sinusoid_1d(x, layer, num_layers):
   return x + signal
 
 
+@expert_utils.add_name_scope()
+def add_timing_signals_given_positions(x,
+                                       positions,
+                                       min_timescale=1.0,
+                                       max_timescale=1.0e4):
+  """Adds sinusoids of diff frequencies to a Tensor, with timing positions given.
+
+  Args:
+    x: a Tensor with shape [batch, length, channels]
+    positions: a list of positions, each of which can either be a Tensor of
+      shape [batch, length] or None for a default of (0..length]
+    min_timescale: a float
+    max_timescale: a float
+
+  Returns:
+    a Tensor the same shape as x.
+  """
+  shape = common_layers.shape_list(x)
+  batch = shape[0]
+  length = shape[1]
+  channels = shape[2]
+  num_dims = len(positions)
+  num_timescales = channels // (num_dims * 2)
+  log_timescale_increment = (
+      math.log(float(max_timescale) / float(min_timescale)) /
+      (tf.to_float(num_timescales) - 1))
+  inv_timescales = min_timescale * tf.exp(
+      tf.to_float(tf.range(num_timescales)) * -log_timescale_increment)
+  for dim, position in enumerate(positions):
+    if position is None:
+      # Create a [batch, length] Tensor of incrementing positions 0..length-1.
+      position = tf.tile(
+          tf.transpose(tf.expand_dims(tf.range(0, length), axis=1)), [batch, 1])
+    scaled_time = (
+        tf.expand_dims(tf.to_float(position), 2) *
+        tf.expand_dims(tf.expand_dims(inv_timescales, 0), 0))
+    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=2)
+    prepad = dim * 2 * num_timescales
+    postpad = channels - (dim + 1) * 2 * num_timescales
+    signal = tf.pad(signal, [[0, 0], [0, 0], [prepad, postpad]])
+    signal = common_layers.cast_like(signal, x)
+    x += signal
+  return x
+
+
+@expert_utils.add_name_scope()
+def add_timing_signals_from_features(x,
+                                     features,
+                                     position_features,
+                                     min_timescale=1.0,
+                                     max_timescale=1.0e4):
+  """Adds timing signals from features named in `position_features`.
+
+  Args:
+    x: a Tensor with shape [batch, length, channels]
+    features: a features dictionary
+    position_features: a comma-delimited string where each item is either a
+      feature key or the empty string (which denotes a default position tensor
+      of [0..length])
+    min_timescale: a float
+    max_timescale: a float
+
+  Returns:
+    a Tensor the same shape as x.
+  """
+  return add_timing_signals_given_positions(x, [
+      features.get(position_feature)
+      for position_feature in position_features.split(",")
+  ], min_timescale, max_timescale)
+
+
 @expert_utils.add_name_scope()
 def add_timing_signal_1d_given_position(x,
                                         position,
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 250e0c864..3be2ba409 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -19,14 +19,14 @@
 from __future__ import division
 from __future__ import print_function
 
+import math
+
 from absl.testing import parameterized
 import kfac
 import numpy as np
-
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import test_utils
-
 import tensorflow as tf
 tf.compat.v1.enable_eager_execution()
 
@@ -102,6 +102,69 @@ def testAddPositionalEmbeddingNd(self, input_shape):
     res = self.evaluate(y)
     self.assertEqual(res.shape, input_shape)
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testAddTimingSignalsGivenPositions(self):
+    x_positions = tf.expand_dims(
+        tf.constant([0, 1, 2, 3], dtype=tf.float32), axis=0)
+    y_positions = tf.expand_dims(
+        tf.constant([4, 5, 6, 7], dtype=tf.float32), axis=0)
+    x = tf.zeros([1, 4, 8], dtype=tf.float32)
+    self.assertAllClose(
+        common_attention.add_timing_signals_given_positions(
+            x, [x_positions, y_positions]),
+        tf.constant([[
+            [
+                math.sin(0),
+                math.sin(0 * 1e-4),
+                math.cos(0),
+                math.cos(0 * 1e-4),
+                math.sin(4),
+                math.sin(4 * 1e-4),
+                math.cos(4),
+                math.cos(4 * 1e-4)
+            ],
+            [
+                math.sin(1),
+                math.sin(1 * 1e-4),
+                math.cos(1),
+                math.cos(1 * 1e-4),
+                math.sin(5),
+                math.sin(5 * 1e-4),
+                math.cos(5),
+                math.cos(5 * 1e-4)
+            ],
+            [
+                math.sin(2),
+                math.sin(2 * 1e-4),
+                math.cos(2),
+                math.cos(2 * 1e-4),
+                math.sin(6),
+                math.sin(6 * 1e-4),
+                math.cos(6),
+                math.cos(6 * 1e-4)
+            ],
+            [
+                math.sin(3),
+                math.sin(3 * 1e-4),
+                math.cos(3),
+                math.cos(3 * 1e-4),
+                math.sin(7),
+                math.sin(7 * 1e-4),
+                math.cos(7),
+                math.cos(7 * 1e-4)
+            ],
+        ]]))
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testAddTimingSignalsGivenPositionsEquivalent(self):
+    x = tf.zeros([1, 10, 128], dtype=tf.float32)
+    positions = tf.expand_dims(tf.range(0, 10, dtype=tf.float32), axis=0)
+    # The method add_timing_signal_1d_given_position could be replaced by
+    # add_timing_signals_given_positions:
+    tf.assert_equal(
+        common_attention.add_timing_signal_1d_given_position(x, positions),
+        common_attention.add_timing_signals_given_positions(x, [positions]))
+
   @test_utils.run_in_graph_and_eager_modes()
   def testDotProductAttention(self):
     x = np.random.rand(5, 7, 12, 32)
diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 236deb685..68ce8baa6 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -111,6 +111,9 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None,
           encoder_input, inputs_position)
     else:
       encoder_input = common_attention.add_timing_signal_1d(encoder_input)
+  elif hparams.pos == "timing_from_features":
+    encoder_input = common_attention.add_timing_signals_from_features(
+        encoder_input, features, hparams.position_features)
   elif hparams.pos == "emb":
     encoder_input = common_attention.add_positional_embedding(
         encoder_input, hparams.max_length, "inputs_positional_embedding",
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index aeb4d63e9..928063a37 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -503,6 +503,10 @@ def _fast_decode_tpu(self,
     if hparams.pos == "timing":
       positional_encoding = common_attention.get_timing_signal_1d(
           decode_length + 1, hparams.hidden_size)
+    elif hparams.pos == "timing_from_features":
+      positional_encoding = common_attention.add_timing_signals_from_features(
+          tf.zeros([1, decode_length + 1, hparams.hidden_size]), features,
+          hparams.position_features)
     elif hparams.pos == "emb":
       positional_encoding = common_attention.add_positional_embedding(
           tf.zeros([1, decode_length + 1, hparams.hidden_size]),
@@ -748,6 +752,10 @@ def _fast_decode(self,
     if hparams.pos == "timing":
       positional_encoding = common_attention.get_timing_signal_1d(
           decode_length + 1, hparams.hidden_size)
+    elif hparams.pos == "timing_from_features":
+      positional_encoding = common_attention.add_timing_signals_from_features(
+          tf.zeros([1, decode_length + 1, hparams.hidden_size]), features,
+          hparams.position_features)
     elif hparams.pos == "emb":
       positional_encoding = common_attention.add_positional_embedding(
           tf.zeros([1, decode_length, hparams.hidden_size]), hparams.max_length,
@@ -1402,6 +1410,9 @@ def transformer_prepare_decoder(targets, hparams, features=None, pad=None):
           decoder_input, targets_position)
     else:
       decoder_input = common_attention.add_timing_signal_1d(decoder_input)
+  elif hparams.pos == "timing_from_features":
+    decoder_input = common_attention.add_timing_signals_from_features(
+        decoder_input, features, hparams.position_features)
   elif hparams.pos == "emb":
     decoder_input = common_attention.add_positional_embedding(
         decoder_input, hparams.max_length, "targets_positional_embedding",
@@ -1768,6 +1779,7 @@ def transformer_base_v1():
   hparams.add_hparam("relu_dropout", 0.0)
   hparams.add_hparam("relu_dropout_broadcast_dims", "")
   hparams.add_hparam("pos", "timing")  # timing, none
+  hparams.add_hparam("position_features", "")
   hparams.add_hparam("nbr_decoder_problems", 1)
   hparams.add_hparam("proximity_bias", False)
   hparams.add_hparam("causal_decoder_self_attention", True)

From e4ea462e73388b6c8a2ea99a5d92aacd9c3237c5 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sat, 21 Dec 2019 20:59:49 -0800
Subject: [PATCH 2606/2720] *** Reason for rollback ***

This change will break existing checkpoints in a hard to debug, hard to notice manner. Existing checkpoints will still load and run, but produce entirely incorrect results.

*** Original change description ***

PR Title: inaccurate calculation of first_token's pos_embedding

PR Body:
since,
1. scaled_time is calculated by `scaled_time = tf.expand_dims(position, 1) * tf.expand_dims(inv_timescales, 0)`, while position_idx in position will **start by 0**.
2. scaled_time will be used to calculate position embedding by `signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)`

Thus, position_embedding of first token will always be combined with half zeros and half ones.

But the purpose of ad...

***

PiperOrigin-RevId: 286760574
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index c00c99540..d0062083e 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -439,7 +439,7 @@ def get_timing_signal_1d(length,
   Returns:
     a Tensor of timing signals [1, length, channels]
   """
-  position = tf.to_float(tf.range(1, 1+length) + start_index)
+  position = tf.to_float(tf.range(length) + start_index)
   num_timescales = channels // 2
   log_timescale_increment = (
       math.log(float(max_timescale) / float(min_timescale)) /

From 955c4f3c57b7ff67d4f5b742d27716291f55199b Mon Sep 17 00:00:00 2001
From: Yanhua Sun <yanhuasun@google.com>
Date: Mon, 23 Dec 2019 16:22:12 -0800
Subject: [PATCH 2607/2720] Explicitly replace "import tensorflow" with
 "tensorflow.compat.v1"

PiperOrigin-RevId: 286956992
---
 tensor2tensor/data_generators/algorithmic.py                    | 2 +-
 tensor2tensor/data_generators/algorithmic_math_deepmind.py      | 2 +-
 tensor2tensor/data_generators/algorithmic_math_test.py          | 2 +-
 tensor2tensor/data_generators/algorithmic_math_two_variables.py | 2 +-
 tensor2tensor/data_generators/algorithmic_test.py               | 2 +-
 tensor2tensor/data_generators/audio.py                          | 2 +-
 tensor2tensor/data_generators/audio_test.py                     | 2 +-
 tensor2tensor/data_generators/babi_qa.py                        | 2 +-
 tensor2tensor/data_generators/celeba.py                         | 2 +-
 tensor2tensor/data_generators/celeba_test.py                    | 2 +-
 tensor2tensor/data_generators/celebahq.py                       | 2 +-
 tensor2tensor/data_generators/cifar.py                          | 2 +-
 tensor2tensor/data_generators/cleaner_en_xx.py                  | 2 +-
 tensor2tensor/data_generators/cnn_dailymail.py                  | 2 +-
 tensor2tensor/data_generators/cola.py                           | 2 +-
 tensor2tensor/data_generators/common_voice.py                   | 2 +-
 tensor2tensor/data_generators/common_voice_test.py              | 2 +-
 tensor2tensor/data_generators/conll_ner.py                      | 2 +-
 tensor2tensor/data_generators/desc2code.py                      | 2 +-
 tensor2tensor/data_generators/desc2code_test.py                 | 2 +-
 tensor2tensor/data_generators/dialog_abstract.py                | 2 +-
 tensor2tensor/data_generators/dna_encoder_test.py               | 2 +-
 tensor2tensor/data_generators/enwik8.py                         | 2 +-
 tensor2tensor/data_generators/function_docstring.py             | 2 +-
 tensor2tensor/data_generators/gene_expression.py                | 2 +-
 tensor2tensor/data_generators/gene_expression_test.py           | 2 +-
 tensor2tensor/data_generators/generator_utils.py                | 2 +-
 tensor2tensor/data_generators/generator_utils_test.py           | 2 +-
 tensor2tensor/data_generators/google_robot_pushing.py           | 2 +-
 tensor2tensor/data_generators/gym_env_test.py                   | 2 +-
 tensor2tensor/data_generators/image_lsun.py                     | 2 +-
 tensor2tensor/data_generators/image_utils_test.py               | 2 +-
 tensor2tensor/data_generators/imagenet.py                       | 2 +-
 tensor2tensor/data_generators/imagenet_test.py                  | 2 +-
 tensor2tensor/data_generators/imdb.py                           | 2 +-
 tensor2tensor/data_generators/lambada.py                        | 2 +-
 tensor2tensor/data_generators/librispeech.py                    | 2 +-
 tensor2tensor/data_generators/lm1b.py                           | 2 +-
 tensor2tensor/data_generators/mnist.py                          | 2 +-
 tensor2tensor/data_generators/mrpc.py                           | 2 +-
 tensor2tensor/data_generators/mscoco.py                         | 2 +-
 tensor2tensor/data_generators/mscoco_test.py                    | 2 +-
 tensor2tensor/data_generators/multi_problem.py                  | 2 +-
 tensor2tensor/data_generators/multi_problem_v2.py               | 2 +-
 tensor2tensor/data_generators/multi_problem_v2_test.py          | 2 +-
 tensor2tensor/data_generators/multinli.py                       | 2 +-
 tensor2tensor/data_generators/ocr.py                            | 2 +-
 tensor2tensor/data_generators/paraphrase_ms_coco.py             | 2 +-
 tensor2tensor/data_generators/paraphrase_ms_coco_test.py        | 2 +-
 tensor2tensor/data_generators/pointer_generator_word.py         | 2 +-
 tensor2tensor/data_generators/problem_hparams.py                | 2 +-
 tensor2tensor/data_generators/problem_test.py                   | 2 +-
 tensor2tensor/data_generators/program_search.py                 | 2 +-
 tensor2tensor/data_generators/program_search_test.py            | 2 +-
 tensor2tensor/data_generators/ptb.py                            | 2 +-
 tensor2tensor/data_generators/qnli.py                           | 2 +-
 tensor2tensor/data_generators/quora_qpairs.py                   | 2 +-
 tensor2tensor/data_generators/rte.py                            | 2 +-
 tensor2tensor/data_generators/scitail.py                        | 2 +-
 tensor2tensor/data_generators/snli.py                           | 2 +-
 tensor2tensor/data_generators/speech_recognition.py             | 2 +-
 tensor2tensor/data_generators/squad.py                          | 2 +-
 tensor2tensor/data_generators/sst_binary.py                     | 2 +-
 tensor2tensor/data_generators/stanford_nli.py                   | 2 +-
 tensor2tensor/data_generators/style_transfer_test.py            | 2 +-
 tensor2tensor/data_generators/subject_verb_agreement.py         | 2 +-
 tensor2tensor/data_generators/text_encoder.py                   | 2 +-
 tensor2tensor/data_generators/text_encoder_test.py              | 2 +-
 tensor2tensor/data_generators/text_problems.py                  | 2 +-
 tensor2tensor/data_generators/text_problems_test.py             | 2 +-
 tensor2tensor/data_generators/timeseries.py                     | 2 +-
 tensor2tensor/data_generators/timeseries_data_generator_test.py | 2 +-
 tensor2tensor/data_generators/timeseries_test.py                | 2 +-
 tensor2tensor/data_generators/tokenizer.py                      | 2 +-
 tensor2tensor/data_generators/tokenizer_test.py                 | 2 +-
 tensor2tensor/data_generators/transduction_problems.py          | 2 +-
 tensor2tensor/data_generators/transduction_problems_test.py     | 2 +-
 tensor2tensor/data_generators/translate_ende_test.py            | 2 +-
 tensor2tensor/data_generators/translate_enzh.py                 | 2 +-
 tensor2tensor/data_generators/translate_test.py                 | 2 +-
 tensor2tensor/data_generators/video_utils_test.py               | 2 +-
 tensor2tensor/data_generators/vqa_utils.py                      | 2 +-
 tensor2tensor/data_generators/wiki.py                           | 2 +-
 tensor2tensor/data_generators/wiki_lm.py                        | 2 +-
 tensor2tensor/data_generators/wiki_revision.py                  | 2 +-
 tensor2tensor/data_generators/wiki_revision_utils.py            | 2 +-
 tensor2tensor/data_generators/wikitext103.py                    | 2 +-
 tensor2tensor/data_generators/wnli.py                           | 2 +-
 tensor2tensor/data_generators/wsj_parsing.py                    | 2 +-
 tensor2tensor/data_generators/yelp_full.py                      | 2 +-
 tensor2tensor/data_generators/yelp_polarity.py                  | 2 +-
 91 files changed, 91 insertions(+), 91 deletions(-)

diff --git a/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index f5458f2ec..fe4f8ee0e 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -28,7 +28,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class AlgorithmicProblem(problem.Problem):
diff --git a/tensor2tensor/data_generators/algorithmic_math_deepmind.py b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
index dce068130..1d849c075 100644
--- a/tensor2tensor/data_generators/algorithmic_math_deepmind.py
+++ b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
@@ -30,7 +30,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 _URL = "https://storage.cloud.google.com/mathematics-dataset/mathematics_dataset-v1.0.tar.gz"
diff --git a/tensor2tensor/data_generators/algorithmic_math_test.py b/tensor2tensor/data_generators/algorithmic_math_test.py
index 035ec16ce..006a64ded 100644
--- a/tensor2tensor/data_generators/algorithmic_math_test.py
+++ b/tensor2tensor/data_generators/algorithmic_math_test.py
@@ -23,7 +23,7 @@
 import sympy
 from tensor2tensor.data_generators import algorithmic_math
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class AlgorithmicMathTest(tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/algorithmic_math_two_variables.py b/tensor2tensor/data_generators/algorithmic_math_two_variables.py
index 2d40a437b..bad119cbf 100644
--- a/tensor2tensor/data_generators/algorithmic_math_two_variables.py
+++ b/tensor2tensor/data_generators/algorithmic_math_two_variables.py
@@ -50,7 +50,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 _URL = ("https://art.wangperawong.com/mathematical_language_understanding"
diff --git a/tensor2tensor/data_generators/algorithmic_test.py b/tensor2tensor/data_generators/algorithmic_test.py
index c15bc76d3..ab6448399 100644
--- a/tensor2tensor/data_generators/algorithmic_test.py
+++ b/tensor2tensor/data_generators/algorithmic_test.py
@@ -22,7 +22,7 @@
 
 from tensor2tensor.data_generators import algorithmic
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class AlgorithmicTest(tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/audio.py b/tensor2tensor/data_generators/audio.py
index 7a5d5d64f..b37a4d6e9 100644
--- a/tensor2tensor/data_generators/audio.py
+++ b/tensor2tensor/data_generators/audio.py
@@ -23,7 +23,7 @@
 import tarfile
 import wave
 from absl import flags
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 FLAGS = flags.FLAGS
 
diff --git a/tensor2tensor/data_generators/audio_test.py b/tensor2tensor/data_generators/audio_test.py
index 1af2ffd74..329c8d77b 100644
--- a/tensor2tensor/data_generators/audio_test.py
+++ b/tensor2tensor/data_generators/audio_test.py
@@ -23,7 +23,7 @@
 import os
 from tensor2tensor.data_generators import audio
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class AudioTest(tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index d312a8272..0d49655e3 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -45,7 +45,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 _DIR_NAME = "tasks_1-20_v1-2"
diff --git a/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index abfd0b784..c4d5776c5 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -26,7 +26,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/celeba_test.py b/tensor2tensor/data_generators/celeba_test.py
index a782bac16..deb3472c1 100644
--- a/tensor2tensor/data_generators/celeba_test.py
+++ b/tensor2tensor/data_generators/celeba_test.py
@@ -23,7 +23,7 @@
 from tensor2tensor.data_generators import celeba
 from tensor2tensor.utils import hparam
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class CelebaTest(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/celebahq.py b/tensor2tensor/data_generators/celebahq.py
index 259de8407..d170883ed 100644
--- a/tensor2tensor/data_generators/celebahq.py
+++ b/tensor2tensor/data_generators/celebahq.py
@@ -25,7 +25,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index 22bdfc831..e644ed550 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -34,7 +34,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # URLs and filenames for CIFAR data.
 _CIFAR10_URL = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
diff --git a/tensor2tensor/data_generators/cleaner_en_xx.py b/tensor2tensor/data_generators/cleaner_en_xx.py
index 15243d17b..56737701e 100644
--- a/tensor2tensor/data_generators/cleaner_en_xx.py
+++ b/tensor2tensor/data_generators/cleaner_en_xx.py
@@ -36,7 +36,7 @@
 
 from tensor2tensor.data_generators import text_encoder
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 _RE_GOOD_S_START = re.compile(r'^["“”]?[A-Z]')
diff --git a/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index af11c928b..c792d5b06 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -31,7 +31,7 @@
 from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # Links to data from http://cs.nyu.edu/~kcho/DMQA/
 _CNN_STORIES_DRIVE_URL = ("https://drive.google.com/uc?"
diff --git a/tensor2tensor/data_generators/cola.py b/tensor2tensor/data_generators/cola.py
index be5fa6da5..7b08178b0 100644
--- a/tensor2tensor/data_generators/cola.py
+++ b/tensor2tensor/data_generators/cola.py
@@ -26,7 +26,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 EOS = text_encoder.EOS
 
diff --git a/tensor2tensor/data_generators/common_voice.py b/tensor2tensor/data_generators/common_voice.py
index 2dda2ef99..cc4867150 100644
--- a/tensor2tensor/data_generators/common_voice.py
+++ b/tensor2tensor/data_generators/common_voice.py
@@ -30,7 +30,7 @@
 from tensor2tensor.data_generators import speech_recognition
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 _COMMONVOICE_URL = "https://common-voice-data-download.s3.amazonaws.com/cv_corpus_v1.tar.gz"  # pylint: disable=line-too-long
 
diff --git a/tensor2tensor/data_generators/common_voice_test.py b/tensor2tensor/data_generators/common_voice_test.py
index 337c4228c..005c4034a 100644
--- a/tensor2tensor/data_generators/common_voice_test.py
+++ b/tensor2tensor/data_generators/common_voice_test.py
@@ -22,7 +22,7 @@
 import os
 from tensor2tensor.data_generators import common_voice
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 pkg_dir, _ = os.path.split(__file__)
 _TESTDATA = os.path.join(pkg_dir, "test_data")
diff --git a/tensor2tensor/data_generators/conll_ner.py b/tensor2tensor/data_generators/conll_ner.py
index ff742b2c6..0e2b536b7 100644
--- a/tensor2tensor/data_generators/conll_ner.py
+++ b/tensor2tensor/data_generators/conll_ner.py
@@ -26,7 +26,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/desc2code.py b/tensor2tensor/data_generators/desc2code.py
index 11216a06b..1524765d4 100644
--- a/tensor2tensor/data_generators/desc2code.py
+++ b/tensor2tensor/data_generators/desc2code.py
@@ -30,7 +30,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 # End-of-sentence marker.
diff --git a/tensor2tensor/data_generators/desc2code_test.py b/tensor2tensor/data_generators/desc2code_test.py
index b8a2bffac..f33016ab9 100644
--- a/tensor2tensor/data_generators/desc2code_test.py
+++ b/tensor2tensor/data_generators/desc2code_test.py
@@ -21,7 +21,7 @@
 
 from tensor2tensor.data_generators import desc2code
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 CODE_CPP_IN = """
   #include <iostream>
diff --git a/tensor2tensor/data_generators/dialog_abstract.py b/tensor2tensor/data_generators/dialog_abstract.py
index 16e791e0a..9346910e4 100644
--- a/tensor2tensor/data_generators/dialog_abstract.py
+++ b/tensor2tensor/data_generators/dialog_abstract.py
@@ -32,7 +32,7 @@
 from tensor2tensor.data_generators.text_problems import VocabType
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # End-of-sentence marker.
 EOS = text_encoder.EOS_ID
diff --git a/tensor2tensor/data_generators/dna_encoder_test.py b/tensor2tensor/data_generators/dna_encoder_test.py
index 2ceb91c75..b7739fae7 100644
--- a/tensor2tensor/data_generators/dna_encoder_test.py
+++ b/tensor2tensor/data_generators/dna_encoder_test.py
@@ -19,7 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 from tensor2tensor.data_generators import dna_encoder
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class DnaEncoderTest(tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/enwik8.py b/tensor2tensor/data_generators/enwik8.py
index 8fa41a272..c752230ce 100644
--- a/tensor2tensor/data_generators/enwik8.py
+++ b/tensor2tensor/data_generators/enwik8.py
@@ -27,7 +27,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def _maybe_download_corpus(tmp_dir):
diff --git a/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
index 04ddbc8ad..8e17e7d32 100644
--- a/tensor2tensor/data_generators/function_docstring.py
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -20,7 +20,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/gene_expression.py b/tensor2tensor/data_generators/gene_expression.py
index ee51db912..ab84ff790 100644
--- a/tensor2tensor/data_generators/gene_expression.py
+++ b/tensor2tensor/data_generators/gene_expression.py
@@ -51,7 +51,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 MAX_CONCURRENT_PROCESSES = 10
 
diff --git a/tensor2tensor/data_generators/gene_expression_test.py b/tensor2tensor/data_generators/gene_expression_test.py
index 2461bbed9..c07221a03 100644
--- a/tensor2tensor/data_generators/gene_expression_test.py
+++ b/tensor2tensor/data_generators/gene_expression_test.py
@@ -22,7 +22,7 @@
 from tensor2tensor.data_generators import dna_encoder
 from tensor2tensor.data_generators import gene_expression
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class GeneticsTest(tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index b66517a83..1fe678e14 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -38,7 +38,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.utils import mlperf_log
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 UNSHUFFLED_SUFFIX = "-unshuffled"
 
diff --git a/tensor2tensor/data_generators/generator_utils_test.py b/tensor2tensor/data_generators/generator_utils_test.py
index 3b83a968d..b729d6076 100644
--- a/tensor2tensor/data_generators/generator_utils_test.py
+++ b/tensor2tensor/data_generators/generator_utils_test.py
@@ -27,7 +27,7 @@
 
 from tensor2tensor.data_generators import generator_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 INPUTS = (
diff --git a/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index f49bd46e4..d4bbe0793 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -35,7 +35,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 BASE_URL = "https://storage.googleapis.com/brain-robotics-data/push/"
 DATA_TRAIN = (264, "push_train/push_train.tfrecord-{:05d}-of-00264")
diff --git a/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 3efd54715..dd252f394 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -31,7 +31,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.rl.gym_utils import make_gym_env
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TestEnv(gym.Env):
diff --git a/tensor2tensor/data_generators/image_lsun.py b/tensor2tensor/data_generators/image_lsun.py
index c26e312cc..653fab2f1 100644
--- a/tensor2tensor/data_generators/image_lsun.py
+++ b/tensor2tensor/data_generators/image_lsun.py
@@ -25,7 +25,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 _LSUN_URL = "http://lsun.cs.princeton.edu/htbin/download.cgi?tag=latest&category=%s&set=%s"
diff --git a/tensor2tensor/data_generators/image_utils_test.py b/tensor2tensor/data_generators/image_utils_test.py
index 0dd020c24..330d88768 100644
--- a/tensor2tensor/data_generators/image_utils_test.py
+++ b/tensor2tensor/data_generators/image_utils_test.py
@@ -22,7 +22,7 @@
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.utils import decoding
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class ImageTest(tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 5e1b3d358..58a98062b 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -27,7 +27,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # URLs and filenames for IMAGENET 32x32 data from
 # https://arxiv.org/abs/1601.06759.
diff --git a/tensor2tensor/data_generators/imagenet_test.py b/tensor2tensor/data_generators/imagenet_test.py
index 9ea07a43c..6a9299589 100644
--- a/tensor2tensor/data_generators/imagenet_test.py
+++ b/tensor2tensor/data_generators/imagenet_test.py
@@ -23,7 +23,7 @@
 from tensor2tensor.data_generators import imagenet
 from tensor2tensor.utils import hparam
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class ImagenetTest(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/imdb.py b/tensor2tensor/data_generators/imdb.py
index 801b7e865..68f5111a3 100644
--- a/tensor2tensor/data_generators/imdb.py
+++ b/tensor2tensor/data_generators/imdb.py
@@ -26,7 +26,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/lambada.py b/tensor2tensor/data_generators/lambada.py
index 753767c94..605e4a1e3 100644
--- a/tensor2tensor/data_generators/lambada.py
+++ b/tensor2tensor/data_generators/lambada.py
@@ -43,7 +43,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 _UNK = "<UNK>"
diff --git a/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index 9d2841f39..97c0c6588 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -22,7 +22,7 @@
 from tensor2tensor.data_generators import speech_recognition
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 _LIBRISPEECH_TRAIN_DATASETS = [
     [
diff --git a/tensor2tensor/data_generators/lm1b.py b/tensor2tensor/data_generators/lm1b.py
index c44d82d35..a98545438 100644
--- a/tensor2tensor/data_generators/lm1b.py
+++ b/tensor2tensor/data_generators/lm1b.py
@@ -29,7 +29,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def _original_vocab(tmp_dir):
diff --git a/tensor2tensor/data_generators/mnist.py b/tensor2tensor/data_generators/mnist.py
index 40f03977b..67ca5b9e2 100644
--- a/tensor2tensor/data_generators/mnist.py
+++ b/tensor2tensor/data_generators/mnist.py
@@ -28,7 +28,7 @@
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # URLs and filenames for MNIST data.
 _MNIST_URL = "http://yann.lecun.com/exdb/mnist/"
diff --git a/tensor2tensor/data_generators/mrpc.py b/tensor2tensor/data_generators/mrpc.py
index 7848a98a4..2eb37c4df 100644
--- a/tensor2tensor/data_generators/mrpc.py
+++ b/tensor2tensor/data_generators/mrpc.py
@@ -25,7 +25,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 EOS = text_encoder.EOS
 
diff --git a/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index dcb029723..28b77c07a 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -32,7 +32,7 @@
 from tensor2tensor.data_generators import translate_ende
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # URLs and filenames for MSCOCO data.
 _MSCOCO_ROOT_URL = "http://msvocds.blob.core.windows.net/"
diff --git a/tensor2tensor/data_generators/mscoco_test.py b/tensor2tensor/data_generators/mscoco_test.py
index e2ffadd75..8443fa914 100644
--- a/tensor2tensor/data_generators/mscoco_test.py
+++ b/tensor2tensor/data_generators/mscoco_test.py
@@ -23,7 +23,7 @@
 from tensor2tensor.data_generators import mscoco
 from tensor2tensor.utils import hparam
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class MscocoTest(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index fe534f050..7d51dc002 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -25,7 +25,7 @@
 from tensor2tensor.layers import discretization
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class MixingSchedule(object):
diff --git a/tensor2tensor/data_generators/multi_problem_v2.py b/tensor2tensor/data_generators/multi_problem_v2.py
index 62bd8c4ed..7d5c8f9ef 100644
--- a/tensor2tensor/data_generators/multi_problem_v2.py
+++ b/tensor2tensor/data_generators/multi_problem_v2.py
@@ -61,7 +61,7 @@
 
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_problems
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class MultiProblemV2(problem.Problem):
diff --git a/tensor2tensor/data_generators/multi_problem_v2_test.py b/tensor2tensor/data_generators/multi_problem_v2_test.py
index e714850c9..56b3d67e5 100644
--- a/tensor2tensor/data_generators/multi_problem_v2_test.py
+++ b/tensor2tensor/data_generators/multi_problem_v2_test.py
@@ -24,7 +24,7 @@
 
 from tensor2tensor.data_generators import multi_problem_v2
 from tensor2tensor.data_generators import problem
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class MultiProblemV2Test(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index 29f236e44..8dfaab72c 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -28,7 +28,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 EOS = text_encoder.EOS
 
diff --git a/tensor2tensor/data_generators/ocr.py b/tensor2tensor/data_generators/ocr.py
index 8f68d9885..4d682c4fa 100644
--- a/tensor2tensor/data_generators/ocr.py
+++ b/tensor2tensor/data_generators/ocr.py
@@ -24,7 +24,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/paraphrase_ms_coco.py b/tensor2tensor/data_generators/paraphrase_ms_coco.py
index df7bd30a8..36fbde85b 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco.py
@@ -30,7 +30,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 _MS_COCO_DOWNLOAD_URL = "http://msvocds.blob.core.windows.net/annotations-1-0-3"
 _MS_COCO_ZIPPED_FILE = "captions_train-val2014.zip"
diff --git a/tensor2tensor/data_generators/paraphrase_ms_coco_test.py b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
index aa57bce3b..f9dfdf52c 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
@@ -23,7 +23,7 @@
 
 from tensor2tensor.data_generators import paraphrase_ms_coco
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class ParaphraseGenerationProblemTest(tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/pointer_generator_word.py b/tensor2tensor/data_generators/pointer_generator_word.py
index ea14ba1c3..093ea03cd 100644
--- a/tensor2tensor/data_generators/pointer_generator_word.py
+++ b/tensor2tensor/data_generators/pointer_generator_word.py
@@ -24,7 +24,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/problem_hparams.py b/tensor2tensor/data_generators/problem_hparams.py
index ad0b76d27..41bbb47e0 100644
--- a/tensor2tensor/data_generators/problem_hparams.py
+++ b/tensor2tensor/data_generators/problem_hparams.py
@@ -26,7 +26,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # TODO(rsepassi): Merge these problems with their data generators. Currently
 # they only implement the hparams.
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 51aeb4f34..6fdd94521 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -29,7 +29,7 @@
 from tensor2tensor.utils import hparam
 from tensor2tensor.utils import test_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 tf.compat.v1.enable_eager_execution()
 
 
diff --git a/tensor2tensor/data_generators/program_search.py b/tensor2tensor/data_generators/program_search.py
index 160949857..06f6a8d56 100644
--- a/tensor2tensor/data_generators/program_search.py
+++ b/tensor2tensor/data_generators/program_search.py
@@ -28,7 +28,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/program_search_test.py b/tensor2tensor/data_generators/program_search_test.py
index 41fcaa37d..14b1d3291 100644
--- a/tensor2tensor/data_generators/program_search_test.py
+++ b/tensor2tensor/data_generators/program_search_test.py
@@ -28,7 +28,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import program_search
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class ProgramSearchAlgolispStub(program_search.ProgramSearchAlgolisp):
diff --git a/tensor2tensor/data_generators/ptb.py b/tensor2tensor/data_generators/ptb.py
index 6f63077cc..6e9dfd5b8 100644
--- a/tensor2tensor/data_generators/ptb.py
+++ b/tensor2tensor/data_generators/ptb.py
@@ -29,7 +29,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 EOS = text_encoder.EOS
diff --git a/tensor2tensor/data_generators/qnli.py b/tensor2tensor/data_generators/qnli.py
index 732a46c18..9889ee3f6 100644
--- a/tensor2tensor/data_generators/qnli.py
+++ b/tensor2tensor/data_generators/qnli.py
@@ -26,7 +26,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 EOS = text_encoder.EOS
 
diff --git a/tensor2tensor/data_generators/quora_qpairs.py b/tensor2tensor/data_generators/quora_qpairs.py
index 003f34e28..81d6c1a13 100644
--- a/tensor2tensor/data_generators/quora_qpairs.py
+++ b/tensor2tensor/data_generators/quora_qpairs.py
@@ -26,7 +26,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 EOS = text_encoder.EOS
 
diff --git a/tensor2tensor/data_generators/rte.py b/tensor2tensor/data_generators/rte.py
index 892c40f32..e39f5c4ac 100644
--- a/tensor2tensor/data_generators/rte.py
+++ b/tensor2tensor/data_generators/rte.py
@@ -26,7 +26,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 EOS = text_encoder.EOS
 
diff --git a/tensor2tensor/data_generators/scitail.py b/tensor2tensor/data_generators/scitail.py
index a9d023a55..e044a3101 100644
--- a/tensor2tensor/data_generators/scitail.py
+++ b/tensor2tensor/data_generators/scitail.py
@@ -27,7 +27,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 EOS = text_encoder.EOS
 
diff --git a/tensor2tensor/data_generators/snli.py b/tensor2tensor/data_generators/snli.py
index 8bad0f699..70b1a2139 100644
--- a/tensor2tensor/data_generators/snli.py
+++ b/tensor2tensor/data_generators/snli.py
@@ -25,7 +25,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import tokenizer
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 _EOS = 1
 _SEP = 2
diff --git a/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 3201a92fd..b105bd55c 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -29,7 +29,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class ByteTextEncoderWithEos(text_encoder.ByteTextEncoder):
diff --git a/tensor2tensor/data_generators/squad.py b/tensor2tensor/data_generators/squad.py
index dc9117f6b..cc3eb5c81 100644
--- a/tensor2tensor/data_generators/squad.py
+++ b/tensor2tensor/data_generators/squad.py
@@ -28,7 +28,7 @@
 from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 _DEV_SET = "dev-v1.1.json"
diff --git a/tensor2tensor/data_generators/sst_binary.py b/tensor2tensor/data_generators/sst_binary.py
index c5682404c..3b61d5ef0 100644
--- a/tensor2tensor/data_generators/sst_binary.py
+++ b/tensor2tensor/data_generators/sst_binary.py
@@ -26,7 +26,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 EOS = text_encoder.EOS
 
diff --git a/tensor2tensor/data_generators/stanford_nli.py b/tensor2tensor/data_generators/stanford_nli.py
index 2ff90d0d6..b3f038651 100644
--- a/tensor2tensor/data_generators/stanford_nli.py
+++ b/tensor2tensor/data_generators/stanford_nli.py
@@ -28,7 +28,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.data_generators import wiki_lm
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 EOS = text_encoder.EOS
 
diff --git a/tensor2tensor/data_generators/style_transfer_test.py b/tensor2tensor/data_generators/style_transfer_test.py
index 534f6335a..4f0947556 100644
--- a/tensor2tensor/data_generators/style_transfer_test.py
+++ b/tensor2tensor/data_generators/style_transfer_test.py
@@ -21,7 +21,7 @@
 
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import style_transfer
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class StyleTransferProblemShakespeareTest(tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/subject_verb_agreement.py b/tensor2tensor/data_generators/subject_verb_agreement.py
index 57f0992d4..7f2630b59 100644
--- a/tensor2tensor/data_generators/subject_verb_agreement.py
+++ b/tensor2tensor/data_generators/subject_verb_agreement.py
@@ -39,7 +39,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 _FILE_NAME = 'agr_50_mostcommon_10K'
 _TAR = _FILE_NAME + '.tsv.gz'
diff --git a/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 1d241164d..a9f51f1be 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -35,7 +35,7 @@
 from six.moves import range  # pylint: disable=redefined-builtin
 from tensor2tensor.data_generators import tokenizer
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # Reserved tokens for things like padding and EOS symbols.
 PAD = "<pad>"
diff --git a/tensor2tensor/data_generators/text_encoder_test.py b/tensor2tensor/data_generators/text_encoder_test.py
index 29f7c808f..e37203d53 100644
--- a/tensor2tensor/data_generators/text_encoder_test.py
+++ b/tensor2tensor/data_generators/text_encoder_test.py
@@ -32,7 +32,7 @@
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.data_generators import text_encoder
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class NativeToUnicodeTest(tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index a039c8a13..d868c58b6 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -40,7 +40,7 @@
 from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class VocabType(object):
diff --git a/tensor2tensor/data_generators/text_problems_test.py b/tensor2tensor/data_generators/text_problems_test.py
index b253421b3..7cf895e6c 100644
--- a/tensor2tensor/data_generators/text_problems_test.py
+++ b/tensor2tensor/data_generators/text_problems_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class Test1(text_problems.Text2textTmpdir):
diff --git a/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index e8f01e511..e0f92fe8b 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -27,7 +27,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TimeseriesProblem(problem.Problem):
diff --git a/tensor2tensor/data_generators/timeseries_data_generator_test.py b/tensor2tensor/data_generators/timeseries_data_generator_test.py
index 1050ab45b..85356bc97 100644
--- a/tensor2tensor/data_generators/timeseries_data_generator_test.py
+++ b/tensor2tensor/data_generators/timeseries_data_generator_test.py
@@ -23,7 +23,7 @@
 
 from tensor2tensor.data_generators import timeseries_data_generator
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TimeseriesDataGeneratorTest(tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index e0947108a..41c149837 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -24,7 +24,7 @@
 
 from tensor2tensor.data_generators import timeseries
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TimeseriesTest(tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/tokenizer.py b/tensor2tensor/data_generators/tokenizer.py
index 0bce48fd3..f00b5845b 100644
--- a/tensor2tensor/data_generators/tokenizer.py
+++ b/tensor2tensor/data_generators/tokenizer.py
@@ -50,7 +50,7 @@
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
 from tensor2tensor.utils import mlperf_log
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # Conversion between Unicode and UTF-8, if required (on Python2)
 _native_to_unicode = (lambda s: s.decode("utf-8")) if six.PY2 else (lambda s: s)
diff --git a/tensor2tensor/data_generators/tokenizer_test.py b/tensor2tensor/data_generators/tokenizer_test.py
index a5bc99657..60ffd0823 100644
--- a/tensor2tensor/data_generators/tokenizer_test.py
+++ b/tensor2tensor/data_generators/tokenizer_test.py
@@ -25,7 +25,7 @@
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
 from tensor2tensor.data_generators import tokenizer
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 pkg_dir, _ = os.path.split(__file__)
diff --git a/tensor2tensor/data_generators/transduction_problems.py b/tensor2tensor/data_generators/transduction_problems.py
index 37082dcbf..a3aa3984b 100644
--- a/tensor2tensor/data_generators/transduction_problems.py
+++ b/tensor2tensor/data_generators/transduction_problems.py
@@ -40,7 +40,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 MAX_GENERATOR_ATTEMPTS = 100
diff --git a/tensor2tensor/data_generators/transduction_problems_test.py b/tensor2tensor/data_generators/transduction_problems_test.py
index 3678b24e4..033ca3922 100644
--- a/tensor2tensor/data_generators/transduction_problems_test.py
+++ b/tensor2tensor/data_generators/transduction_problems_test.py
@@ -29,7 +29,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import transduction_problems
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TransductionProblem(parameterized.TestCase):
diff --git a/tensor2tensor/data_generators/translate_ende_test.py b/tensor2tensor/data_generators/translate_ende_test.py
index b8b5e4550..bd9467af0 100644
--- a/tensor2tensor/data_generators/translate_ende_test.py
+++ b/tensor2tensor/data_generators/translate_ende_test.py
@@ -22,7 +22,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import translate_ende
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TranslateEndeTest(tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index c3a7370f8..02acc5381 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -27,7 +27,7 @@
 from tensor2tensor.data_generators import translate
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 # End-of-sentence marker.
diff --git a/tensor2tensor/data_generators/translate_test.py b/tensor2tensor/data_generators/translate_test.py
index f0bbdd6dc..949c5e161 100644
--- a/tensor2tensor/data_generators/translate_test.py
+++ b/tensor2tensor/data_generators/translate_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.data_generators import translate
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TranslateTest(tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index 7db91dade..0af54dad6 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class VideoUtilsTest(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/vqa_utils.py b/tensor2tensor/data_generators/vqa_utils.py
index 660d3c943..74aa8a5b6 100644
--- a/tensor2tensor/data_generators/vqa_utils.py
+++ b/tensor2tensor/data_generators/vqa_utils.py
@@ -20,7 +20,7 @@
 from __future__ import print_function
 
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.python.ops import control_flow_ops
 
diff --git a/tensor2tensor/data_generators/wiki.py b/tensor2tensor/data_generators/wiki.py
index 3ea3d9e39..17a7f6330 100644
--- a/tensor2tensor/data_generators/wiki.py
+++ b/tensor2tensor/data_generators/wiki.py
@@ -29,7 +29,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/wiki_lm.py b/tensor2tensor/data_generators/wiki_lm.py
index 1e6dd3ca9..c64fe9fd9 100644
--- a/tensor2tensor/data_generators/wiki_lm.py
+++ b/tensor2tensor/data_generators/wiki_lm.py
@@ -27,7 +27,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def concat_generator(filename, up_threshold, low_threshold=10):
diff --git a/tensor2tensor/data_generators/wiki_revision.py b/tensor2tensor/data_generators/wiki_revision.py
index f41da7556..2a925cc67 100644
--- a/tensor2tensor/data_generators/wiki_revision.py
+++ b/tensor2tensor/data_generators/wiki_revision.py
@@ -39,7 +39,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 FLAGS = flags.FLAGS
 
diff --git a/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index 22f86b966..eb027203c 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -30,7 +30,7 @@
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import text_encoder
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def include_revision(revision_num, skip_factor=1.1):
diff --git a/tensor2tensor/data_generators/wikitext103.py b/tensor2tensor/data_generators/wikitext103.py
index 1e53eefbc..968fc3fea 100644
--- a/tensor2tensor/data_generators/wikitext103.py
+++ b/tensor2tensor/data_generators/wikitext103.py
@@ -31,7 +31,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def _build_vocab(filename, vocab_dir, vocab_name):
diff --git a/tensor2tensor/data_generators/wnli.py b/tensor2tensor/data_generators/wnli.py
index 413ffbf88..c6267403e 100644
--- a/tensor2tensor/data_generators/wnli.py
+++ b/tensor2tensor/data_generators/wnli.py
@@ -26,7 +26,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 EOS = text_encoder.EOS
 
diff --git a/tensor2tensor/data_generators/wsj_parsing.py b/tensor2tensor/data_generators/wsj_parsing.py
index 40225e43b..228c80656 100644
--- a/tensor2tensor/data_generators/wsj_parsing.py
+++ b/tensor2tensor/data_generators/wsj_parsing.py
@@ -25,7 +25,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags.DEFINE_string("parsing_path", "", "Path to parsing files in tmp_dir.")
 
diff --git a/tensor2tensor/data_generators/yelp_full.py b/tensor2tensor/data_generators/yelp_full.py
index 1ce222482..88e61c840 100644
--- a/tensor2tensor/data_generators/yelp_full.py
+++ b/tensor2tensor/data_generators/yelp_full.py
@@ -26,7 +26,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/yelp_polarity.py b/tensor2tensor/data_generators/yelp_polarity.py
index ae1e9dce3..bf2faa099 100644
--- a/tensor2tensor/data_generators/yelp_polarity.py
+++ b/tensor2tensor/data_generators/yelp_polarity.py
@@ -26,7 +26,7 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_problem

From cfa6c98b30a531e0900d1e0916d5df3d00bb93c8 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 24 Dec 2019 06:54:03 -0800
Subject: [PATCH 2608/2720] Pretraining the knowledge triple embeddings for
 improved augmentation of external knowledge in goal-oriented dialogue
 generation.

PiperOrigin-RevId: 287021893
---
 tensor2tensor/models/neural_assistant.py | 105 +++++++++++++++++++++--
 1 file changed, 99 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/neural_assistant.py b/tensor2tensor/models/neural_assistant.py
index 77695c0fd..40e6788c4 100644
--- a/tensor2tensor/models/neural_assistant.py
+++ b/tensor2tensor/models/neural_assistant.py
@@ -63,15 +63,20 @@ def model_fn(self, features):
         logits = self.top(output, features)
         losses["training"] = 0.0
         cur_kb_loss = losses["kb_loss"]
+        cur_knowledge_training_loss = losses["transe_loss"]
         cur_kb_loss_weight = self._hparams.kb_loss_weight
+        kb_train_weight = self._hparams.kb_train_weight
         cur_lm_loss_weight = 1.0 - cur_kb_loss_weight
         # Finalize loss
         if (self._hparams.mode != tf.estimator.ModeKeys.PREDICT and
             self._hparams.mode != "attack"):
           lm_loss_num, lm_loss_denom = self.loss(logits, features)
-          total_loss = cur_kb_loss * cur_kb_loss_weight + (
-              lm_loss_num / lm_loss_denom) * cur_lm_loss_weight
+          total_loss = (kb_train_weight) * cur_knowledge_training_loss + (
+              1 - kb_train_weight) * (
+                  cur_kb_loss * cur_kb_loss_weight +
+                  (lm_loss_num / lm_loss_denom) * cur_lm_loss_weight)
           tf.summary.scalar("kb_loss", cur_kb_loss)
+          tf.summary.scalar("transe_loss", cur_knowledge_training_loss)
           tf.summary.scalar("lm_loss", (lm_loss_num / lm_loss_denom))
           tf.summary.scalar("cur_kb_loss_weight",
                             tf.reshape(cur_kb_loss_weight, []))
@@ -107,7 +112,8 @@ def encode_knowledge_bottom(self, features):
     return re_fact_embedding, re_fact_lengths
 
   def compute_knowledge_selection_and_loss(self, features, encoder_output,
-                                           fact_embedding, fact_lengths):
+                                           fact_embedding, fact_lengths, margin,
+                                           num_negative_samples):
     """Compute knowledge selection and loss.
 
     Args:
@@ -116,6 +122,9 @@ def compute_knowledge_selection_and_loss(self, features, encoder_output,
       fact_embedding: <tf.float32>[batch_size*triple_num, max_triple_length,
         emb_dim]
       fact_lengths: # <tf.int32>[batch_size*triple_num]
+      margin: integer value for max margin in TransE loss,
+      num_negative_samples: shuffle and sample multiple negative examples for
+      the TransE loss
 
     Returns:
       knowledge_weights:
@@ -169,6 +178,82 @@ def compute_knowledge_selection_and_loss(self, features, encoder_output,
 
     avg_triple_loss = 0.0
     triple_labels = features["triple_labels"]
+
+    subject_mask = tf.reshape(features["subject_mask"],
+                              [-1, self.triple_num, hparams.max_triple_length])
+    subject_mask = tf.reshape(subject_mask, [-1, hparams.max_triple_length])
+
+    predicate_mask = tf.reshape(
+        features["predicate_mask"],
+        [-1, self.triple_num, hparams.max_triple_length])
+    predicate_mask = tf.reshape(predicate_mask, [-1, hparams.max_triple_length])
+
+    object_mask = tf.reshape(features["object_mask"],
+                             [-1, self.triple_num, hparams.max_triple_length])
+    object_mask = tf.reshape(object_mask, [-1, hparams.max_triple_length])
+
+    # mask : [bs, max_seq_len, triple_num]
+    # the below operation will result in [bs*triple_num,emb_dim]
+    subject_length = tf.cast(
+        tf.expand_dims(tf.reduce_sum(subject_mask, -1), 1),
+        tf.float32)  # [bs*tn]
+    object_length = tf.cast(
+        tf.expand_dims(tf.reduce_sum(object_mask, -1), 1), tf.float32)
+    predicate_length = tf.cast(
+        tf.expand_dims(tf.reduce_sum(predicate_mask, -1), 1), tf.float32)
+
+    # expand dimension 2 to be able to broadcast
+    subject_mask = tf.cast(tf.expand_dims(subject_mask, 2), tf.float32)
+    predicate_mask = tf.cast(tf.expand_dims(predicate_mask, 2), tf.float32)
+    object_mask = tf.cast(tf.expand_dims(object_mask, 2), tf.float32)
+
+    subject_vect = tf.reduce_sum(tf.multiply(
+        fact_embedding, subject_mask), 1) / (
+            subject_length +
+            tf.broadcast_to(tf.constant([1e-5]), tf.shape(subject_length)))
+    object_vect = tf.reduce_sum(tf.multiply(fact_embedding, object_mask), 1) / (
+        object_length +
+        tf.broadcast_to(tf.constant([1e-5]), tf.shape(object_length)))
+    predicate_vect = tf.reduce_sum(
+        tf.multiply(fact_embedding, predicate_mask), 1) / (
+            predicate_length +
+            tf.broadcast_to(tf.constant([1e-5]), tf.shape(predicate_length)))
+
+    # Shuffled rows to generate adversarial samples
+    shuffled_subject_vect = []
+    shuffled_object_vect = []
+
+    for _ in range(num_negative_samples):
+      shuffled_subject_vect += [
+          tf.gather(subject_vect,
+                    tf.random.shuffle(tf.range(tf.shape(subject_vect)[0])))
+      ]  # [bs*tn,d]
+      shuffled_object_vect += [
+          tf.gather(object_vect,
+                    tf.random.shuffle(tf.range(tf.shape(object_vect)[0])))
+      ]  # [bs*tn,d]
+
+    # KB pretraining loss
+
+    positive_loss = tf.reduce_mean(
+        tf.squared_difference(subject_vect + predicate_vect, object_vect))
+    negative_loss = 0
+    for n_adv in range(num_negative_samples):
+      negative_loss += tf.reduce_mean(
+          tf.squared_difference(shuffled_subject_vect[n_adv] + predicate_vect,
+                                object_vect))
+      negative_loss += tf.reduce_mean(
+          tf.squared_difference(subject_vect + predicate_vect,
+                                shuffled_object_vect[n_adv]))
+
+    # TransE Loss
+
+    negative_loss = negative_loss / (2 * num_negative_samples)
+
+    transe_loss = tf.clip_by_value(
+        margin + positive_loss - negative_loss,
+        clip_value_min=0,
+        clip_value_max=100)
     if hparams.mode != tf.estimator.ModeKeys.PREDICT:
       triple_losses = tf.nn.weighted_cross_entropy_with_logits(
           labels=triple_labels,
@@ -177,7 +262,7 @@ def compute_knowledge_selection_and_loss(self, features, encoder_output,
       avg_triple_loss = tf.reduce_mean(triple_losses)
       tf.summary.scalar("triple_loss", avg_triple_loss)
 
-    return triple_logits, avg_triple_loss, original_knowledge_encoder_output
+    return triple_logits, avg_triple_loss, original_knowledge_encoder_output, transe_loss
 
   def body(self, features):
     """Transformer main model_fn.
@@ -212,9 +297,11 @@ def body(self, features):
 
       with tf.name_scope("knowledge_selection_and_loss"):
         # Compute knowledge selection and loss.
-        triple_logits, avg_triple_selection_loss, knowledge_encoder_output = self.compute_knowledge_selection_and_loss(
-            features, encoder_output, fact_embedding, fact_lengths)
+        triple_logits, avg_triple_selection_loss, knowledge_encoder_output, transe_loss = self.compute_knowledge_selection_and_loss(
+            features, encoder_output, fact_embedding, fact_lengths,
+            hparams.margin, hparams.num_negative_samples)
         losses["kb_loss"] = avg_triple_selection_loss
+        losses["transe_loss"] = transe_loss
 
     if hparams.attend_kb:
       tf.logging.info("ATTEND_KB is ACTIVE")
@@ -453,6 +540,12 @@ def neural_assistant_tiny():
   hparams.add_hparam("kb_loss_weight", 0.0)  # weight for distant supervision
   hparams.add_hparam("test_triple_num",
                      28483)  # max triples of KB
+  hparams.add_hparam("margin", 1.0)  # KB training max-margin loss
+  hparams.add_hparam(
+      "num_negative_samples",
+      1)  # Sampling number of different adversarial training examples
+  hparams.add_hparam("kb_train_weight", 0.0)
+  # KB_training loss weight which combines Language model and KB selection loss
   return hparams
 
 

From 6e1a183f181f197d7ccb09042b6b75acb7d79260 Mon Sep 17 00:00:00 2001
From: Yanhua Sun <yanhuasun@google.com>
Date: Fri, 27 Dec 2019 14:16:46 -0800
Subject: [PATCH 2609/2720] Explicitly replace "import tensorflow" with
 "tensorflow.compat.v1"

PiperOrigin-RevId: 287340268
---
 tensor2tensor/layers/area_attention.py                        | 2 +-
 tensor2tensor/layers/area_attention_test.py                   | 2 +-
 tensor2tensor/layers/common_audio.py                          | 2 +-
 tensor2tensor/layers/common_hparams.py                        | 2 +-
 tensor2tensor/layers/common_image_attention.py                | 2 +-
 tensor2tensor/layers/common_image_attention_test.py           | 2 +-
 tensor2tensor/layers/common_video_test.py                     | 2 +-
 tensor2tensor/layers/discretization.py                        | 2 +-
 tensor2tensor/layers/discretization_test.py                   | 2 +-
 tensor2tensor/layers/latent_layers.py                         | 2 +-
 tensor2tensor/layers/latent_layers_test.py                    | 2 +-
 tensor2tensor/layers/message_passing_attention.py             | 2 +-
 tensor2tensor/layers/modalities.py                            | 2 +-
 tensor2tensor/layers/modalities_test.py                       | 2 +-
 tensor2tensor/layers/ngram.py                                 | 2 +-
 tensor2tensor/layers/ngram_test.py                            | 2 +-
 tensor2tensor/layers/transformer_layers.py                    | 2 +-
 tensor2tensor/layers/transformer_memory.py                    | 2 +-
 tensor2tensor/layers/transformer_memory_test.py               | 2 +-
 tensor2tensor/layers/vq_discrete.py                           | 2 +-
 tensor2tensor/models/basic.py                                 | 2 +-
 tensor2tensor/models/basic_test.py                            | 2 +-
 tensor2tensor/models/bytenet.py                               | 2 +-
 tensor2tensor/models/bytenet_test.py                          | 2 +-
 tensor2tensor/models/distillation.py                          | 2 +-
 tensor2tensor/models/evolved_transformer.py                   | 2 +-
 tensor2tensor/models/evolved_transformer_test.py              | 2 +-
 tensor2tensor/models/image_transformer.py                     | 2 +-
 tensor2tensor/models/image_transformer_2d.py                  | 2 +-
 tensor2tensor/models/image_transformer_2d_test.py             | 2 +-
 tensor2tensor/models/image_transformer_test.py                | 2 +-
 tensor2tensor/models/lstm_test.py                             | 2 +-
 tensor2tensor/models/mtf_image_transformer.py                 | 2 +-
 tensor2tensor/models/mtf_image_transformer_test.py            | 2 +-
 tensor2tensor/models/mtf_resnet.py                            | 2 +-
 tensor2tensor/models/mtf_transformer.py                       | 2 +-
 tensor2tensor/models/mtf_transformer2.py                      | 2 +-
 tensor2tensor/models/mtf_transformer_test.py                  | 2 +-
 tensor2tensor/models/neural_assistant.py                      | 2 +-
 tensor2tensor/models/neural_gpu.py                            | 2 +-
 tensor2tensor/models/neural_gpu_test.py                       | 2 +-
 tensor2tensor/models/research/aligned.py                      | 2 +-
 tensor2tensor/models/research/attention_lm_moe.py             | 2 +-
 tensor2tensor/models/research/autoencoders.py                 | 2 +-
 tensor2tensor/models/research/autoencoders_test.py            | 2 +-
 tensor2tensor/models/research/cycle_gan.py                    | 2 +-
 tensor2tensor/models/research/gene_expression_test.py         | 2 +-
 tensor2tensor/models/research/glow_init_hook.py               | 2 +-
 tensor2tensor/models/research/glow_test.py                    | 2 +-
 tensor2tensor/models/research/moe.py                          | 2 +-
 tensor2tensor/models/research/rl.py                           | 2 +-
 tensor2tensor/models/research/similarity_transformer.py       | 2 +-
 tensor2tensor/models/research/super_lm.py                     | 2 +-
 tensor2tensor/models/research/transformer_aux.py              | 2 +-
 tensor2tensor/models/research/transformer_aux_test.py         | 2 +-
 tensor2tensor/models/research/transformer_moe.py              | 2 +-
 tensor2tensor/models/research/transformer_nat.py              | 2 +-
 tensor2tensor/models/research/transformer_parallel.py         | 2 +-
 tensor2tensor/models/research/transformer_revnet_test.py      | 2 +-
 tensor2tensor/models/research/transformer_sketch.py           | 2 +-
 tensor2tensor/models/research/transformer_symshard.py         | 2 +-
 tensor2tensor/models/research/transformer_vae_test.py         | 2 +-
 tensor2tensor/models/research/universal_transformer_test.py   | 2 +-
 tensor2tensor/models/research/vqa_attention_test.py           | 2 +-
 tensor2tensor/models/research/vqa_recurrent_self_attention.py | 2 +-
 tensor2tensor/models/research/vqa_self_attention.py           | 2 +-
 tensor2tensor/models/resnet.py                                | 2 +-
 tensor2tensor/models/resnet_test.py                           | 2 +-
 tensor2tensor/models/revnet_test.py                           | 2 +-
 tensor2tensor/models/shake_shake.py                           | 2 +-
 tensor2tensor/models/slicenet.py                              | 2 +-
 tensor2tensor/models/slicenet_test.py                         | 2 +-
 tensor2tensor/models/text_cnn.py                              | 2 +-
 tensor2tensor/models/transformer.py                           | 2 +-
 tensor2tensor/models/transformer_test.py                      | 2 +-
 tensor2tensor/models/vanilla_gan.py                           | 2 +-
 tensor2tensor/models/xception.py                              | 2 +-
 tensor2tensor/models/xception_test.py                         | 2 +-
 78 files changed, 78 insertions(+), 78 deletions(-)

diff --git a/tensor2tensor/layers/area_attention.py b/tensor2tensor/layers/area_attention.py
index 2635608a4..7e1338fbc 100644
--- a/tensor2tensor/layers/area_attention.py
+++ b/tensor2tensor/layers/area_attention.py
@@ -21,7 +21,7 @@
 import numpy as np
 from six.moves import range  # pylint: disable=redefined-builtin
 from tensor2tensor.layers import common_layers
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def lengths_to_area_mask(feature_length, length, max_area_size):
diff --git a/tensor2tensor/layers/area_attention_test.py b/tensor2tensor/layers/area_attention_test.py
index 1b464a6a0..2a2ebc9b4 100644
--- a/tensor2tensor/layers/area_attention_test.py
+++ b/tensor2tensor/layers/area_attention_test.py
@@ -22,7 +22,7 @@
 from absl.testing import parameterized
 import numpy as np
 from tensor2tensor.layers import area_attention
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class AreaAttentionTest(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/layers/common_audio.py b/tensor2tensor/layers/common_audio.py
index 545e831b7..5cb114495 100644
--- a/tensor2tensor/layers/common_audio.py
+++ b/tensor2tensor/layers/common_audio.py
@@ -22,7 +22,7 @@
 import functools
 import numpy as np
 import scipy.signal
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def add_delta_deltas(filterbanks, name=None):
diff --git a/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index eb5d687a0..7aec40c15 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -22,7 +22,7 @@
 from tensor2tensor.utils import hparam
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_hparams("basic_1")
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index c4af3bc1a..a64f71e50 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -26,7 +26,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import expert_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class AttentionType(object):
diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index 26f60c68f..c9e140e5a 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -24,7 +24,7 @@
 from tensor2tensor.layers import common_image_attention
 from tensor2tensor.utils import hparam
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class CommonImageAttentionTest(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index 50493069d..b7188cd34 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.layers import common_video
 from tensor2tensor.utils import test_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 tf.compat.v1.enable_eager_execution()
 
 
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 667b6e5a4..6bcdd7db5 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -24,7 +24,7 @@
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow_probability as tfp
 
 from tensorflow.python.training import moving_averages  # pylint: disable=g-direct-tensorflow-import
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index 1ba6fde3e..5d4eec265 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -23,7 +23,7 @@
 from tensor2tensor.layers import discretization
 from tensor2tensor.utils import test_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 tf.compat.v1.enable_eager_execution()
 
 
diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 77c51a0bb..f23ba00a7 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -26,7 +26,7 @@
 from tensor2tensor.layers import transformer_layers
 from tensor2tensor.utils import beam_search
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow_probability as tfp
 
 DO_SUMMARIES = True
diff --git a/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
index 94ef58801..097465117 100644
--- a/tensor2tensor/layers/latent_layers_test.py
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -27,7 +27,7 @@
 from tensor2tensor.models import transformer
 from tensor2tensor.utils import test_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 tf.compat.v1.enable_eager_execution()
 
 
diff --git a/tensor2tensor/layers/message_passing_attention.py b/tensor2tensor/layers/message_passing_attention.py
index a87304f73..887d88b78 100644
--- a/tensor2tensor/layers/message_passing_attention.py
+++ b/tensor2tensor/layers/message_passing_attention.py
@@ -22,7 +22,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import expert_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def multihead_graph_attention(query_antecedent,
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index eaf56b120..36a5afefa 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -32,7 +32,7 @@
 from tensor2tensor.layers import common_video
 from tensor2tensor.layers import discretization
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow_probability as tfp
 
 
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 69504bb8e..393c558aa 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.utils import expert_utils
 from tensor2tensor.utils import test_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 tf.compat.v1.enable_eager_execution()
 
 
diff --git a/tensor2tensor/layers/ngram.py b/tensor2tensor/layers/ngram.py
index 675a20448..8372f85c8 100644
--- a/tensor2tensor/layers/ngram.py
+++ b/tensor2tensor/layers/ngram.py
@@ -20,7 +20,7 @@
 from __future__ import print_function
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class NGram(tf.keras.layers.Layer):
diff --git a/tensor2tensor/layers/ngram_test.py b/tensor2tensor/layers/ngram_test.py
index 02ef6ba06..0233722e5 100644
--- a/tensor2tensor/layers/ngram_test.py
+++ b/tensor2tensor/layers/ngram_test.py
@@ -23,7 +23,7 @@
 
 from tensor2tensor.utils import test_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 tf.compat.v1.enable_eager_execution()
 
 
diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 68ce8baa6..e0e6a7f54 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -24,7 +24,7 @@
 from tensor2tensor.utils import expert_utils
 from tensor2tensor.utils import mlperf_log
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 # TODO(lukaszkaiser): remove this function when not needed any more.
diff --git a/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index be2410fbb..63e2b1d66 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -19,7 +19,7 @@
 from __future__ import print_function
 
 from tensor2tensor.layers import common_layers
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class RecurrentMemory(object):
diff --git a/tensor2tensor/layers/transformer_memory_test.py b/tensor2tensor/layers/transformer_memory_test.py
index cd86dae95..c98c9790c 100644
--- a/tensor2tensor/layers/transformer_memory_test.py
+++ b/tensor2tensor/layers/transformer_memory_test.py
@@ -21,7 +21,7 @@
 
 from absl.testing import parameterized
 from tensor2tensor.layers import transformer_memory
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TransformerMemoryTest(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index b593d7a87..7799469dd 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -19,7 +19,7 @@
 from __future__ import print_function
 from functools import partial
 from tensor2tensor.layers import common_layers
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from tensorflow.python.training import moving_averages
 
 
diff --git a/tensor2tensor/models/basic.py b/tensor2tensor/models/basic.py
index 155a73ec5..3840e5a59 100644
--- a/tensor2tensor/models/basic.py
+++ b/tensor2tensor/models/basic.py
@@ -24,7 +24,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/basic_test.py b/tensor2tensor/models/basic_test.py
index a89550da9..f3b29cad3 100644
--- a/tensor2tensor/models/basic_test.py
+++ b/tensor2tensor/models/basic_test.py
@@ -24,7 +24,7 @@
 from tensor2tensor.models import basic
 from tensor2tensor.utils import trainer_lib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class BasicTest(tf.test.TestCase):
diff --git a/tensor2tensor/models/bytenet.py b/tensor2tensor/models/bytenet.py
index de5e6c550..23f9c0962 100644
--- a/tensor2tensor/models/bytenet.py
+++ b/tensor2tensor/models/bytenet.py
@@ -25,7 +25,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def residual_dilated_conv(x, repeat, padding, name, hparams):
diff --git a/tensor2tensor/models/bytenet_test.py b/tensor2tensor/models/bytenet_test.py
index 934b6ead4..e65d2022b 100644
--- a/tensor2tensor/models/bytenet_test.py
+++ b/tensor2tensor/models/bytenet_test.py
@@ -23,7 +23,7 @@
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.models import bytenet
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class ByteNetTest(tf.test.TestCase):
diff --git a/tensor2tensor/models/distillation.py b/tensor2tensor/models/distillation.py
index 0a7f69908..4d90166bc 100644
--- a/tensor2tensor/models/distillation.py
+++ b/tensor2tensor/models/distillation.py
@@ -22,7 +22,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 5529fb03c..87253ce61 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -28,7 +28,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.ops import inplace_ops
diff --git a/tensor2tensor/models/evolved_transformer_test.py b/tensor2tensor/models/evolved_transformer_test.py
index a7d782b25..3e6db316d 100644
--- a/tensor2tensor/models/evolved_transformer_test.py
+++ b/tensor2tensor/models/evolved_transformer_test.py
@@ -23,7 +23,7 @@
 from tensor2tensor.models import evolved_transformer
 from tensor2tensor.models import transformer
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 BATCH_SIZE = 3
 INPUT_LENGTH = 5
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index d43f96c15..9bd066e67 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -32,7 +32,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 80ce13d53..783a44ce0 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -33,7 +33,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/image_transformer_2d_test.py b/tensor2tensor/models/image_transformer_2d_test.py
index 7deddc870..9499a24b0 100644
--- a/tensor2tensor/models/image_transformer_2d_test.py
+++ b/tensor2tensor/models/image_transformer_2d_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.models import image_transformer_2d
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class Img2imgTransformerTest(tf.test.TestCase):
diff --git a/tensor2tensor/models/image_transformer_test.py b/tensor2tensor/models/image_transformer_test.py
index 3b61cfb40..87485b9e5 100644
--- a/tensor2tensor/models/image_transformer_test.py
+++ b/tensor2tensor/models/image_transformer_test.py
@@ -26,7 +26,7 @@
 from tensor2tensor.layers import common_image_attention
 from tensor2tensor.models import image_transformer
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class ImagetransformerTest(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/models/lstm_test.py b/tensor2tensor/models/lstm_test.py
index 95b6fe4f0..a4dc22d0d 100644
--- a/tensor2tensor/models/lstm_test.py
+++ b/tensor2tensor/models/lstm_test.py
@@ -23,7 +23,7 @@
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.models import lstm
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class LSTMTest(tf.test.TestCase):
diff --git a/tensor2tensor/models/mtf_image_transformer.py b/tensor2tensor/models/mtf_image_transformer.py
index 67915f656..1a71f673c 100644
--- a/tensor2tensor/models/mtf_image_transformer.py
+++ b/tensor2tensor/models/mtf_image_transformer.py
@@ -32,7 +32,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import mtf_model
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/mtf_image_transformer_test.py b/tensor2tensor/models/mtf_image_transformer_test.py
index 8bcef7a45..81fca2c9d 100644
--- a/tensor2tensor/models/mtf_image_transformer_test.py
+++ b/tensor2tensor/models/mtf_image_transformer_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.models import mtf_image_transformer
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # Constants shared between all functions.
 BATCH_SIZE = 8
diff --git a/tensor2tensor/models/mtf_resnet.py b/tensor2tensor/models/mtf_resnet.py
index 794efc1b8..2c2baf98f 100644
--- a/tensor2tensor/models/mtf_resnet.py
+++ b/tensor2tensor/models/mtf_resnet.py
@@ -28,7 +28,7 @@
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.utils import mtf_model
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 BATCH_NORM_DECAY = 0.9
diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index b8d4c56d8..526565ec1 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -28,7 +28,7 @@
 from tensor2tensor.utils import mtf_model
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index f4b37752b..a4262a960 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -29,7 +29,7 @@
 from tensor2tensor.utils import mtf_model
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/mtf_transformer_test.py b/tensor2tensor/models/mtf_transformer_test.py
index b6b8b5334..c7153be72 100644
--- a/tensor2tensor/models/mtf_transformer_test.py
+++ b/tensor2tensor/models/mtf_transformer_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.models import mtf_transformer
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # Constants shared between all functions.
 BATCH_SIZE = 2
diff --git a/tensor2tensor/models/neural_assistant.py b/tensor2tensor/models/neural_assistant.py
index 40e6788c4..01d55bb13 100644
--- a/tensor2tensor/models/neural_assistant.py
+++ b/tensor2tensor/models/neural_assistant.py
@@ -23,7 +23,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/neural_gpu.py b/tensor2tensor/models/neural_gpu.py
index 1b824d349..1fb255455 100644
--- a/tensor2tensor/models/neural_gpu.py
+++ b/tensor2tensor/models/neural_gpu.py
@@ -25,7 +25,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def neural_gpu_body(inputs, hparams, name=None):
diff --git a/tensor2tensor/models/neural_gpu_test.py b/tensor2tensor/models/neural_gpu_test.py
index 8d3e7068b..e8961f287 100644
--- a/tensor2tensor/models/neural_gpu_test.py
+++ b/tensor2tensor/models/neural_gpu_test.py
@@ -24,7 +24,7 @@
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.models import neural_gpu
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class NeuralGPUTest(tf.test.TestCase):
diff --git a/tensor2tensor/models/research/aligned.py b/tensor2tensor/models/research/aligned.py
index edea1218a..625270ffd 100644
--- a/tensor2tensor/models/research/aligned.py
+++ b/tensor2tensor/models/research/aligned.py
@@ -33,7 +33,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 ModeKeys = tf.estimator.ModeKeys  # pylint: disable=invalid-name
 
diff --git a/tensor2tensor/models/research/attention_lm_moe.py b/tensor2tensor/models/research/attention_lm_moe.py
index 6aa7112e0..04a9cfa91 100644
--- a/tensor2tensor/models/research/attention_lm_moe.py
+++ b/tensor2tensor/models/research/attention_lm_moe.py
@@ -35,7 +35,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 ModeKeys = tf.estimator.ModeKeys  # pylint: disable=invalid-name
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index cf4305a30..1908548da 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -28,7 +28,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def reverse_gradient(x, lr=1.0):
diff --git a/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index 6b7ff62d8..54a13cfe6 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class AutoencoderTest(tf.test.TestCase):
diff --git a/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index 3ba2e64ab..78f2260c3 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -24,7 +24,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def discriminator(x, compress, hparams, name, reuse=None):
diff --git a/tensor2tensor/models/research/gene_expression_test.py b/tensor2tensor/models/research/gene_expression_test.py
index 142b4315b..b58883e36 100644
--- a/tensor2tensor/models/research/gene_expression_test.py
+++ b/tensor2tensor/models/research/gene_expression_test.py
@@ -23,7 +23,7 @@
 from tensor2tensor.layers import modalities  # pylint: disable=unused-import
 from tensor2tensor.models.research import gene_expression
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def gene_expression_conv_test():
diff --git a/tensor2tensor/models/research/glow_init_hook.py b/tensor2tensor/models/research/glow_init_hook.py
index 86009f3e4..51a679d9d 100644
--- a/tensor2tensor/models/research/glow_init_hook.py
+++ b/tensor2tensor/models/research/glow_init_hook.py
@@ -19,7 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class GlowInitHook(tf.train.SessionRunHook):
diff --git a/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
index e157f53cc..67f6309f9 100644
--- a/tensor2tensor/models/research/glow_test.py
+++ b/tensor2tensor/models/research/glow_test.py
@@ -28,7 +28,7 @@
 from tensor2tensor.data_generators import cifar  # pylint: disable=unused-import
 from tensor2tensor.models.research import glow
 from tensor2tensor.utils import registry  # pylint: disable=unused-import
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 MODES = tf.estimator.ModeKeys
 
diff --git a/tensor2tensor/models/research/moe.py b/tensor2tensor/models/research/moe.py
index f70b0eece..370d2e2eb 100644
--- a/tensor2tensor/models/research/moe.py
+++ b/tensor2tensor/models/research/moe.py
@@ -24,7 +24,7 @@
 from __future__ import print_function
 
 import mesh_tensorflow as mtf
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def transformer_moe_layer_v1(inputs, output_dim, hparams, train,
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 68ba22a10..a6f1b5c3a 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -39,7 +39,7 @@
 from tensor2tensor.utils import t2t_model
 from tensor2tensor.utils import trainer_lib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow_probability as tfp
 
 
diff --git a/tensor2tensor/models/research/similarity_transformer.py b/tensor2tensor/models/research/similarity_transformer.py
index 0da345b19..9b5d63f18 100644
--- a/tensor2tensor/models/research/similarity_transformer.py
+++ b/tensor2tensor/models/research/similarity_transformer.py
@@ -19,7 +19,7 @@
 from tensor2tensor.models import transformer
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/research/super_lm.py b/tensor2tensor/models/research/super_lm.py
index 0e7dd58c8..9bf69ebed 100644
--- a/tensor2tensor/models/research/super_lm.py
+++ b/tensor2tensor/models/research/super_lm.py
@@ -38,7 +38,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 ModeKeys = tf.estimator.ModeKeys  # pylint: disable=invalid-name
 
diff --git a/tensor2tensor/models/research/transformer_aux.py b/tensor2tensor/models/research/transformer_aux.py
index ff5aafacf..097cf6a1a 100644
--- a/tensor2tensor/models/research/transformer_aux.py
+++ b/tensor2tensor/models/research/transformer_aux.py
@@ -23,7 +23,7 @@
 from tensor2tensor.models import transformer
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def shift_and_pad(tensor, shift, axis=0):
diff --git a/tensor2tensor/models/research/transformer_aux_test.py b/tensor2tensor/models/research/transformer_aux_test.py
index 11ec146f0..a748f7ffb 100644
--- a/tensor2tensor/models/research/transformer_aux_test.py
+++ b/tensor2tensor/models/research/transformer_aux_test.py
@@ -24,7 +24,7 @@
 import numpy as np
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.models.research import transformer_aux
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TransformerAuxTest(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index 2a4899f04..418ddddb2 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -28,7 +28,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 # The transformer architecture can be defined using the layer_types hparams.
diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 1eb01c5cd..fddcaacc9 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -24,7 +24,7 @@
 from tensor2tensor.utils import beam_search
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from tensorflow.python.training import moving_averages  # pylint: disable=g-direct-tensorflow-import
 
 
diff --git a/tensor2tensor/models/research/transformer_parallel.py b/tensor2tensor/models/research/transformer_parallel.py
index 166c512e0..6bca08dfe 100644
--- a/tensor2tensor/models/research/transformer_parallel.py
+++ b/tensor2tensor/models/research/transformer_parallel.py
@@ -24,7 +24,7 @@
 from tensor2tensor.models import transformer
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/research/transformer_revnet_test.py b/tensor2tensor/models/research/transformer_revnet_test.py
index 37c3648ba..da0ddf258 100644
--- a/tensor2tensor/models/research/transformer_revnet_test.py
+++ b/tensor2tensor/models/research/transformer_revnet_test.py
@@ -23,7 +23,7 @@
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.models.research import transformer_revnet
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def transformer_revnet_test():
diff --git a/tensor2tensor/models/research/transformer_sketch.py b/tensor2tensor/models/research/transformer_sketch.py
index 9b3a0cf94..0677fe1b5 100644
--- a/tensor2tensor/models/research/transformer_sketch.py
+++ b/tensor2tensor/models/research/transformer_sketch.py
@@ -24,7 +24,7 @@
 from tensor2tensor.models import transformer
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/research/transformer_symshard.py b/tensor2tensor/models/research/transformer_symshard.py
index 7f5bf39dd..1e4e1bf17 100644
--- a/tensor2tensor/models/research/transformer_symshard.py
+++ b/tensor2tensor/models/research/transformer_symshard.py
@@ -53,7 +53,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/research/transformer_vae_test.py b/tensor2tensor/models/research/transformer_vae_test.py
index bb04f44bc..18c639f1a 100644
--- a/tensor2tensor/models/research/transformer_vae_test.py
+++ b/tensor2tensor/models/research/transformer_vae_test.py
@@ -20,7 +20,7 @@
 import numpy as np
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.models.research import transformer_vae
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TransformerVaeTest(tf.test.TestCase):
diff --git a/tensor2tensor/models/research/universal_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
index 947cecc6f..1b9d9f4cf 100644
--- a/tensor2tensor/models/research/universal_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.models.research import universal_transformer
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 BATCH_SIZE = 3
 INPUT_LENGTH = 5
diff --git a/tensor2tensor/models/research/vqa_attention_test.py b/tensor2tensor/models/research/vqa_attention_test.py
index da6fc0bc3..86922f7d0 100644
--- a/tensor2tensor/models/research/vqa_attention_test.py
+++ b/tensor2tensor/models/research/vqa_attention_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.models.research import vqa_attention
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class VqaAttentionBaselineTest(tf.test.TestCase):
diff --git a/tensor2tensor/models/research/vqa_recurrent_self_attention.py b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
index 6765bce16..ca6976e64 100644
--- a/tensor2tensor/models/research/vqa_recurrent_self_attention.py
+++ b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
@@ -30,7 +30,7 @@
 from tensor2tensor.utils import registry
 # from tensor2tensor.utils import restore_hook
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.contrib.layers.python.layers import utils
 
diff --git a/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
index c9ae04680..efe08ff38 100644
--- a/tensor2tensor/models/research/vqa_self_attention.py
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -29,7 +29,7 @@
 from tensor2tensor.utils import registry
 # from tensor2tensor.utils import restore_hook
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.contrib.layers.python.layers import utils
 
diff --git a/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index 3e33e4c02..1560c13fe 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -26,7 +26,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 BATCH_NORM_DECAY = 0.9
diff --git a/tensor2tensor/models/resnet_test.py b/tensor2tensor/models/resnet_test.py
index 17e930a79..cd2481d7c 100644
--- a/tensor2tensor/models/resnet_test.py
+++ b/tensor2tensor/models/resnet_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.models import resnet
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def resnet_tiny_cpu():
diff --git a/tensor2tensor/models/revnet_test.py b/tensor2tensor/models/revnet_test.py
index 8344f4090..7bad272ff 100644
--- a/tensor2tensor/models/revnet_test.py
+++ b/tensor2tensor/models/revnet_test.py
@@ -16,7 +16,7 @@
 """Tests for Revnet."""
 
 from tensor2tensor.models import revnet
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class RevnetTest(tf.test.TestCase):
diff --git a/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index 3594a9f56..ba87619c2 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -24,7 +24,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def shake_shake_skip_connection(x, output_filters, stride, is_training):
diff --git a/tensor2tensor/models/slicenet.py b/tensor2tensor/models/slicenet.py
index 55254b3a1..75e73243d 100644
--- a/tensor2tensor/models/slicenet.py
+++ b/tensor2tensor/models/slicenet.py
@@ -27,7 +27,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 # pylint: disable=unused-argument
diff --git a/tensor2tensor/models/slicenet_test.py b/tensor2tensor/models/slicenet_test.py
index 4faf71817..1e20fc00a 100644
--- a/tensor2tensor/models/slicenet_test.py
+++ b/tensor2tensor/models/slicenet_test.py
@@ -26,7 +26,7 @@
 from tensor2tensor.models import slicenet
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class SliceNetTest(tf.test.TestCase):
diff --git a/tensor2tensor/models/text_cnn.py b/tensor2tensor/models/text_cnn.py
index 0a6957f2f..7b01565d1 100644
--- a/tensor2tensor/models/text_cnn.py
+++ b/tensor2tensor/models/text_cnn.py
@@ -24,7 +24,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 928063a37..31222cf6d 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -41,7 +41,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.ops import inplace_ops
diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 678a25341..50e613d4e 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -24,7 +24,7 @@
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.models import transformer
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 BATCH_SIZE = 3
diff --git a/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index bf922df50..79fc4f37d 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -27,7 +27,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def lrelu(input_, leak=0.2, name="lrelu"):
diff --git a/tensor2tensor/models/xception.py b/tensor2tensor/models/xception.py
index 9361745e0..d89fad897 100644
--- a/tensor2tensor/models/xception.py
+++ b/tensor2tensor/models/xception.py
@@ -27,7 +27,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def residual_block(x, hparams):
diff --git a/tensor2tensor/models/xception_test.py b/tensor2tensor/models/xception_test.py
index 041081889..0ac84d5d0 100644
--- a/tensor2tensor/models/xception_test.py
+++ b/tensor2tensor/models/xception_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.models import xception
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class XceptionTest(tf.test.TestCase):

From ee6ffe88a13e20c7bd6d5faf8fa3984a9d939bdf Mon Sep 17 00:00:00 2001
From: Yanhua Sun <yanhuasun@google.com>
Date: Fri, 27 Dec 2019 14:17:38 -0800
Subject: [PATCH 2610/2720] Explicitly replace "import tensorflow" with
 "tensorflow.compat.v1"

PiperOrigin-RevId: 287340335
---
 tensor2tensor/data_generators/ops/pack_sequences_ops_test.py  | 2 +-
 .../data_generators/ops/subword_text_encoder_ops_test.py      | 2 +-
 tensor2tensor/data_generators/wikisum/utils.py                | 2 +-
 tensor2tensor/data_generators/wikisum/utils_test.py           | 2 +-
 tensor2tensor/data_generators/wikisum/wikisum.py              | 2 +-
 tensor2tensor/envs/env_problem_utils_test.py                  | 2 +-
 tensor2tensor/envs/gym_env_problem_test.py                    | 2 +-
 tensor2tensor/envs/gym_spaces_utils.py                        | 2 +-
 tensor2tensor/envs/gym_spaces_utils_test.py                   | 2 +-
 tensor2tensor/envs/mujoco_problems_test.py                    | 2 +-
 tensor2tensor/envs/rendered_env_problem_test.py               | 2 +-
 tensor2tensor/envs/tic_tac_toe_env_problem_test.py            | 2 +-
 tensor2tensor/envs/tic_tac_toe_env_test.py                    | 2 +-
 tensor2tensor/envs/time_step_test.py                          | 2 +-
 tensor2tensor/envs/trajectory.py                              | 2 +-
 tensor2tensor/envs/trajectory_test.py                         | 4 ++--
 tensor2tensor/metrics/video_conditional_fvd_test.py           | 2 +-
 tensor2tensor/models/neural_architecture_search/nas_layers.py | 2 +-
 .../models/neural_architecture_search/nas_layers_test.py      | 2 +-
 .../models/neural_architecture_search/nas_model_test.py       | 2 +-
 tensor2tensor/models/video/base.py                            | 2 +-
 tensor2tensor/models/video/base_vae.py                        | 2 +-
 tensor2tensor/models/video/basic_deterministic.py             | 2 +-
 tensor2tensor/models/video/basic_deterministic_test.py        | 2 +-
 tensor2tensor/models/video/basic_recurrent_test.py            | 2 +-
 tensor2tensor/models/video/basic_stochastic.py                | 2 +-
 tensor2tensor/models/video/basic_stochastic_test.py           | 2 +-
 tensor2tensor/models/video/emily_test.py                      | 2 +-
 tensor2tensor/models/video/nfg_conv3d_test.py                 | 2 +-
 tensor2tensor/models/video/nfg_conv_lstm_test.py              | 2 +-
 tensor2tensor/models/video/nfg_conv_test.py                   | 2 +-
 tensor2tensor/models/video/nfg_test_utils.py                  | 2 +-
 tensor2tensor/models/video/nfg_uncond_test.py                 | 2 +-
 tensor2tensor/models/video/savp_test.py                       | 2 +-
 tensor2tensor/models/video/sv2p_test.py                       | 2 +-
 tensor2tensor/models/video/tests_utils.py                     | 2 +-
 tensor2tensor/problems_test.py                                | 2 +-
 tensor2tensor/rl/batch_dqn_agent_test.py                      | 2 +-
 tensor2tensor/rl/batch_runner_test.py                         | 2 +-
 tensor2tensor/rl/dopamine_connector.py                        | 2 +-
 tensor2tensor/rl/envs/in_graph_batch_env.py                   | 2 +-
 tensor2tensor/rl/envs/py_func_batch_env.py                    | 2 +-
 tensor2tensor/rl/envs/simulated_batch_env.py                  | 2 +-
 tensor2tensor/rl/envs/simulated_batch_gym_env.py              | 2 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py                    | 2 +-
 tensor2tensor/rl/evaluator_test.py                            | 2 +-
 tensor2tensor/rl/gym_utils_test.py                            | 2 +-
 tensor2tensor/rl/ppo.py                                       | 2 +-
 tensor2tensor/rl/ppo_learner.py                               | 2 +-
 tensor2tensor/rl/restarter.py                                 | 2 +-
 tensor2tensor/rl/restarter_test.py                            | 2 +-
 tensor2tensor/rl/rl_utils.py                                  | 2 +-
 tensor2tensor/visualization/visualization.py                  | 2 +-
 tensor2tensor/visualization/visualization_test.py             | 2 +-
 54 files changed, 55 insertions(+), 55 deletions(-)

diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
index 3531d7631..c0c033635 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
@@ -20,7 +20,7 @@
 from __future__ import print_function
 
 from tensor2tensor.data_generators.ops import pack_sequences_ops
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class PackSequencesOpsTest(tf.test.TestCase):
diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py b/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
index 677f95de1..28bb135ac 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
@@ -20,7 +20,7 @@
 from __future__ import print_function
 
 from tensor2tensor.data_generators.ops import subword_text_encoder_ops
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 vocab_file = (
     "third_party/py/tensor2tensor/data_generators/ops/testdata/subwords")
diff --git a/tensor2tensor/data_generators/wikisum/utils.py b/tensor2tensor/data_generators/wikisum/utils.py
index 396172216..541d2c268 100644
--- a/tensor2tensor/data_generators/wikisum/utils.py
+++ b/tensor2tensor/data_generators/wikisum/utils.py
@@ -26,7 +26,7 @@
 import re
 import urllib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # pylint: disable=g-import-not-at-top
 # To maintain compatibility with Python 2 and 3
diff --git a/tensor2tensor/data_generators/wikisum/utils_test.py b/tensor2tensor/data_generators/wikisum/utils_test.py
index e713d1938..3bf6ef71b 100644
--- a/tensor2tensor/data_generators/wikisum/utils_test.py
+++ b/tensor2tensor/data_generators/wikisum/utils_test.py
@@ -22,7 +22,7 @@
 import os
 from tensor2tensor.data_generators.wikisum import utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 pkg_dir = os.path.abspath(__file__)
 pkg_dir, _ = os.path.split(pkg_dir)
diff --git a/tensor2tensor/data_generators/wikisum/wikisum.py b/tensor2tensor/data_generators/wikisum/wikisum.py
index 8e994ea9d..8be4a59a6 100644
--- a/tensor2tensor/data_generators/wikisum/wikisum.py
+++ b/tensor2tensor/data_generators/wikisum/wikisum.py
@@ -36,7 +36,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 PROCESS_FOLDER_PREFIX = "process"
 REF_SHARD_FILE_PREFIX = "references.tfrecords.gz"
diff --git a/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index 572294bca..07324f30e 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -26,7 +26,7 @@
 from tensor2tensor.envs import tic_tac_toe_env  # pylint: disable=unused-import
 from tensor2tensor.envs import tic_tac_toe_env_problem
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class EnvProblemUtilsTest(tf.test.TestCase):
diff --git a/tensor2tensor/envs/gym_env_problem_test.py b/tensor2tensor/envs/gym_env_problem_test.py
index d75e81645..ea4b94167 100644
--- a/tensor2tensor/envs/gym_env_problem_test.py
+++ b/tensor2tensor/envs/gym_env_problem_test.py
@@ -30,7 +30,7 @@
 from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.envs import gym_env_problem
 from tensor2tensor.layers import modalities
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class GymEnvProblemTest(tf.test.TestCase):
diff --git a/tensor2tensor/envs/gym_spaces_utils.py b/tensor2tensor/envs/gym_spaces_utils.py
index 9b96948de..66fb3825a 100644
--- a/tensor2tensor/envs/gym_spaces_utils.py
+++ b/tensor2tensor/envs/gym_spaces_utils.py
@@ -26,7 +26,7 @@
 from gym.spaces import Discrete
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def box_space_spec(box_space, tf_dtype):
diff --git a/tensor2tensor/envs/gym_spaces_utils_test.py b/tensor2tensor/envs/gym_spaces_utils_test.py
index af7b76112..26cd6f574 100644
--- a/tensor2tensor/envs/gym_spaces_utils_test.py
+++ b/tensor2tensor/envs/gym_spaces_utils_test.py
@@ -23,7 +23,7 @@
 from gym.spaces import Discrete
 import numpy as np
 from tensor2tensor.envs import gym_spaces_utils
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class GymSpacesUtilsTest(tf.test.TestCase):
diff --git a/tensor2tensor/envs/mujoco_problems_test.py b/tensor2tensor/envs/mujoco_problems_test.py
index 5a66f5cd1..4bad4590f 100644
--- a/tensor2tensor/envs/mujoco_problems_test.py
+++ b/tensor2tensor/envs/mujoco_problems_test.py
@@ -23,7 +23,7 @@
 from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.envs import mujoco_problems  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class ReacherEnvProblemTest(tf.test.TestCase):
diff --git a/tensor2tensor/envs/rendered_env_problem_test.py b/tensor2tensor/envs/rendered_env_problem_test.py
index d52dce618..d1ec67cc1 100644
--- a/tensor2tensor/envs/rendered_env_problem_test.py
+++ b/tensor2tensor/envs/rendered_env_problem_test.py
@@ -23,7 +23,7 @@
 from tensor2tensor.envs import env_problem_utils
 from tensor2tensor.envs import rendered_env_problem
 from tensor2tensor.envs.mujoco_problems import ReacherEnvProblem
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class RenderedEnvProblemTest(tf.test.TestCase):
diff --git a/tensor2tensor/envs/tic_tac_toe_env_problem_test.py b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
index ac5a549cf..89517e8ba 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
@@ -24,7 +24,7 @@
 from tensor2tensor.envs import tic_tac_toe_env  # pylint: disable=unused-import
 from tensor2tensor.envs import tic_tac_toe_env_problem  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TicTacToeEnvProblemTest(tf.test.TestCase):
diff --git a/tensor2tensor/envs/tic_tac_toe_env_test.py b/tensor2tensor/envs/tic_tac_toe_env_test.py
index 8925d1dbf..c00b53683 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_test.py
@@ -20,7 +20,7 @@
 from __future__ import print_function
 
 from tensor2tensor.envs import tic_tac_toe_env as ttt_env
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TicTacToeEnvTest(tf.test.TestCase):
diff --git a/tensor2tensor/envs/time_step_test.py b/tensor2tensor/envs/time_step_test.py
index 1ab8f26b7..ec75c32a7 100644
--- a/tensor2tensor/envs/time_step_test.py
+++ b/tensor2tensor/envs/time_step_test.py
@@ -21,7 +21,7 @@
 
 from tensor2tensor.envs import time_step
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TimeStepTest(tf.test.TestCase):
diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index e980db4d5..5eec464da 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -31,7 +31,7 @@
 import cloudpickle
 import numpy as np
 from tensor2tensor.envs import time_step
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 TRAJECTORY_FILE_FORMAT = r"trajectory_epoch_{epoch}_env_id_{env_id}_temperature_{temperature}_r_{r}.pkl"
 
diff --git a/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index 6bc3dc442..e98d33a22 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -23,8 +23,8 @@
 import numpy as np
 from tensor2tensor.envs import time_step
 from tensor2tensor.envs import trajectory
-import tensorflow as tf
-from tensorflow.io import gfile
+import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1.io import gfile
 
 
 class TrajectoryTest(tf.test.TestCase):
diff --git a/tensor2tensor/metrics/video_conditional_fvd_test.py b/tensor2tensor/metrics/video_conditional_fvd_test.py
index 84446325e..2d9b847f8 100644
--- a/tensor2tensor/metrics/video_conditional_fvd_test.py
+++ b/tensor2tensor/metrics/video_conditional_fvd_test.py
@@ -20,7 +20,7 @@
 from __future__ import print_function
 
 from tensor2tensor.metrics import video_conditional_fvd
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class VideoConditionalFvdTest(tf.test.TestCase):
diff --git a/tensor2tensor/models/neural_architecture_search/nas_layers.py b/tensor2tensor/models/neural_architecture_search/nas_layers.py
index f0220575c..f86b8e742 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_layers.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers.py
@@ -28,7 +28,7 @@
 
 from tensor2tensor.layers import common_attention
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # Registry layer keys.
 ATTEND_TO_ENCODER_REGISTRY_KEY = "attend_to_encoder"
diff --git a/tensor2tensor/models/neural_architecture_search/nas_layers_test.py b/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
index 1d1fc705e..657b09a31 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
@@ -28,7 +28,7 @@
 from tensor2tensor.models import transformer
 from tensor2tensor.models.neural_architecture_search import nas_layers as layers
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 _BATCH_SIZE = 32
 _TOTAL_SEQUENCE_LENGTH = 20
diff --git a/tensor2tensor/models/neural_architecture_search/nas_model_test.py b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
index c5fd3775c..86d706e9b 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model_test.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.models import transformer
 from tensor2tensor.models.neural_architecture_search import nas_layers as layers
 from tensor2tensor.models.neural_architecture_search import nas_model as translation_nas_net
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 _BATCH_SIZE = 5
 _INPUT_LENGTH = 5
diff --git a/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index bdebf9632..55ca3dcdd 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -30,7 +30,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def flat_lists(list_of_lists):
diff --git a/tensor2tensor/models/video/base_vae.py b/tensor2tensor/models/video/base_vae.py
index ca9620dae..069397034 100644
--- a/tensor2tensor/models/video/base_vae.py
+++ b/tensor2tensor/models/video/base_vae.py
@@ -22,7 +22,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class NextFrameBaseVae(object):
diff --git a/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index ff712f5e8..cfd592ae9 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -27,7 +27,7 @@
 from tensor2tensor.models.video import basic_deterministic_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/video/basic_deterministic_test.py b/tensor2tensor/models/video/basic_deterministic_test.py
index 0c16666a7..9f9d3a216 100644
--- a/tensor2tensor/models/video/basic_deterministic_test.py
+++ b/tensor2tensor/models/video/basic_deterministic_test.py
@@ -23,7 +23,7 @@
 from tensor2tensor.models.video import basic_deterministic_params
 from tensor2tensor.models.video import tests_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class NextFrameTest(tests_utils.BaseNextFrameTest):
diff --git a/tensor2tensor/models/video/basic_recurrent_test.py b/tensor2tensor/models/video/basic_recurrent_test.py
index dfcef5bd7..f19f3ac7a 100644
--- a/tensor2tensor/models/video/basic_recurrent_test.py
+++ b/tensor2tensor/models/video/basic_recurrent_test.py
@@ -22,7 +22,7 @@
 from tensor2tensor.models.video import basic_recurrent
 from tensor2tensor.models.video import tests_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class NextFrameTest(tests_utils.BaseNextFrameTest):
diff --git a/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 7fa519e1a..e46b142bf 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -30,7 +30,7 @@
 
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 tfl = tf.layers
 _MAX_BATCH = 128
diff --git a/tensor2tensor/models/video/basic_stochastic_test.py b/tensor2tensor/models/video/basic_stochastic_test.py
index 4eb339d64..c0e27b76d 100644
--- a/tensor2tensor/models/video/basic_stochastic_test.py
+++ b/tensor2tensor/models/video/basic_stochastic_test.py
@@ -22,7 +22,7 @@
 from tensor2tensor.models.video import basic_stochastic
 from tensor2tensor.models.video import tests_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class NextFrameTest(tests_utils.BaseNextFrameTest):
diff --git a/tensor2tensor/models/video/emily_test.py b/tensor2tensor/models/video/emily_test.py
index 9ccf34c79..6c35d5951 100644
--- a/tensor2tensor/models/video/emily_test.py
+++ b/tensor2tensor/models/video/emily_test.py
@@ -23,7 +23,7 @@
 from tensor2tensor.models.video import tests_utils
 
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class NextFrameTest(tests_utils.BaseNextFrameTest):
diff --git a/tensor2tensor/models/video/nfg_conv3d_test.py b/tensor2tensor/models/video/nfg_conv3d_test.py
index 9c7f0fab8..9050ea955 100644
--- a/tensor2tensor/models/video/nfg_conv3d_test.py
+++ b/tensor2tensor/models/video/nfg_conv3d_test.py
@@ -21,7 +21,7 @@
 
 from absl.testing import parameterized
 from tensor2tensor.models.video import nfg_test_utils
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 conv3d_net_hparams = (
     ("conv3d_net", 2, 2, "conv3d_net", "conditional", -1, 3),
diff --git a/tensor2tensor/models/video/nfg_conv_lstm_test.py b/tensor2tensor/models/video/nfg_conv_lstm_test.py
index 046cec72d..607b445e3 100644
--- a/tensor2tensor/models/video/nfg_conv_lstm_test.py
+++ b/tensor2tensor/models/video/nfg_conv_lstm_test.py
@@ -21,7 +21,7 @@
 
 from absl.testing import parameterized
 from tensor2tensor.models.video import nfg_test_utils
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 conv_lstm_hparams = (
     ("in_3_out_2_lstm", 2, 1, "conv_lstm", "conditional", -1),
diff --git a/tensor2tensor/models/video/nfg_conv_test.py b/tensor2tensor/models/video/nfg_conv_test.py
index 51e00ed1c..b155967db 100644
--- a/tensor2tensor/models/video/nfg_conv_test.py
+++ b/tensor2tensor/models/video/nfg_conv_test.py
@@ -17,7 +17,7 @@
 
 from absl.testing import parameterized
 from tensor2tensor.models.video import nfg_test_utils
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 conv_net_hparams = (
     ("in_3_out_2_conv", 3, 1, "conv_net", "conditional"),
diff --git a/tensor2tensor/models/video/nfg_test_utils.py b/tensor2tensor/models/video/nfg_test_utils.py
index ec48cf818..e135dd6ea 100644
--- a/tensor2tensor/models/video/nfg_test_utils.py
+++ b/tensor2tensor/models/video/nfg_test_utils.py
@@ -25,7 +25,7 @@
 from tensor2tensor.data_generators import video_generated  # pylint: disable=unused-import
 from tensor2tensor.models.video import next_frame_glow
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 MODES = tf.estimator.ModeKeys
 
 
diff --git a/tensor2tensor/models/video/nfg_uncond_test.py b/tensor2tensor/models/video/nfg_uncond_test.py
index 97fa984d1..a533a995c 100644
--- a/tensor2tensor/models/video/nfg_uncond_test.py
+++ b/tensor2tensor/models/video/nfg_uncond_test.py
@@ -21,7 +21,7 @@
 
 from absl.testing import parameterized
 from tensor2tensor.models.video import nfg_test_utils
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 uncond_hparams = (
     ("in_1_out_1", 1, 1, "pointwise", "conditional"),
diff --git a/tensor2tensor/models/video/savp_test.py b/tensor2tensor/models/video/savp_test.py
index 5ec184b4e..49d14d5a8 100644
--- a/tensor2tensor/models/video/savp_test.py
+++ b/tensor2tensor/models/video/savp_test.py
@@ -24,7 +24,7 @@
 from tensor2tensor.models.video import tests_utils
 
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class NextFrameTest(tests_utils.BaseNextFrameTest):
diff --git a/tensor2tensor/models/video/sv2p_test.py b/tensor2tensor/models/video/sv2p_test.py
index 3694af807..39815b5b1 100644
--- a/tensor2tensor/models/video/sv2p_test.py
+++ b/tensor2tensor/models/video/sv2p_test.py
@@ -23,7 +23,7 @@
 from tensor2tensor.models.video import sv2p_params
 from tensor2tensor.models.video import tests_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class NextFrameTest(tests_utils.BaseNextFrameTest):
diff --git a/tensor2tensor/models/video/tests_utils.py b/tensor2tensor/models/video/tests_utils.py
index e28936993..e9d2d843c 100644
--- a/tensor2tensor/models/video/tests_utils.py
+++ b/tensor2tensor/models/video/tests_utils.py
@@ -25,7 +25,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def fill_hparams(hparams, in_frames, out_frames):
diff --git a/tensor2tensor/problems_test.py b/tensor2tensor/problems_test.py
index c49611ed6..bab90968a 100644
--- a/tensor2tensor/problems_test.py
+++ b/tensor2tensor/problems_test.py
@@ -20,7 +20,7 @@
 
 from tensor2tensor import problems
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class ProblemsTest(tf.test.TestCase):
diff --git a/tensor2tensor/rl/batch_dqn_agent_test.py b/tensor2tensor/rl/batch_dqn_agent_test.py
index 87152f49d..7e2821788 100644
--- a/tensor2tensor/rl/batch_dqn_agent_test.py
+++ b/tensor2tensor/rl/batch_dqn_agent_test.py
@@ -28,7 +28,7 @@
 
 from tensor2tensor.rl import dopamine_connector
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 FLAGS = flags.FLAGS
 
diff --git a/tensor2tensor/rl/batch_runner_test.py b/tensor2tensor/rl/batch_runner_test.py
index 3b3ee21d4..a260d9d97 100644
--- a/tensor2tensor/rl/batch_runner_test.py
+++ b/tensor2tensor/rl/batch_runner_test.py
@@ -29,7 +29,7 @@
 
 from tensor2tensor.rl import dopamine_connector
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 48fd92c92..136d837c1 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -33,7 +33,7 @@
 import numpy as np
 
 from tensor2tensor.rl.policy_learner import PolicyLearner
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # pylint: disable=g-import-not-at-top
 # pylint: disable=ungrouped-imports
diff --git a/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index 1164f7872..f8ba6967c 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -24,7 +24,7 @@
 
 import gym
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class InGraphBatchEnv(object):
diff --git a/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index 50d3ee3dc..ca455b4ac 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -25,7 +25,7 @@
 import numpy as np
 
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class PyFuncBatchEnv(InGraphBatchEnv):
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index bfee48de0..bca298bcd 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -34,7 +34,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 # Lazy load PIL.Image
diff --git a/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index fd220837f..6a69a045f 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -25,7 +25,7 @@
 
 from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class FlatBatchEnv(Env):
diff --git a/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 5c801c9ba..03567cad4 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -23,7 +23,7 @@
 
 from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class WrapperBase(InGraphBatchEnv):
diff --git a/tensor2tensor/rl/evaluator_test.py b/tensor2tensor/rl/evaluator_test.py
index dea1a6666..e8adeeb30 100644
--- a/tensor2tensor/rl/evaluator_test.py
+++ b/tensor2tensor/rl/evaluator_test.py
@@ -21,7 +21,7 @@
 from tensor2tensor.rl import evaluator
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class EvalTest(tf.test.TestCase):
diff --git a/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index 5d3d46f16..b347ba906 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -23,7 +23,7 @@
 from gym import spaces
 import numpy as np
 from tensor2tensor.rl import gym_utils
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class SimpleEnv(gym.Env):
diff --git a/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index 92d65b35b..a9f476ca0 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -26,7 +26,7 @@
 from tensor2tensor.utils import learning_rate
 from tensor2tensor.utils import optimize
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow_probability as tfp
 
 
diff --git a/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 9ca4ffccd..04c053098 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -31,7 +31,7 @@
 from tensor2tensor.rl.restarter import Restarter
 from tensor2tensor.utils import trainer_lib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow_probability as tfp
 
 
diff --git a/tensor2tensor/rl/restarter.py b/tensor2tensor/rl/restarter.py
index e996984b9..72e8628e1 100644
--- a/tensor2tensor/rl/restarter.py
+++ b/tensor2tensor/rl/restarter.py
@@ -18,7 +18,7 @@
 import contextlib
 import os
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class Restarter(object):
diff --git a/tensor2tensor/rl/restarter_test.py b/tensor2tensor/rl/restarter_test.py
index ff0455ecd..d28da52b7 100644
--- a/tensor2tensor/rl/restarter_test.py
+++ b/tensor2tensor/rl/restarter_test.py
@@ -19,7 +19,7 @@
 
 from tensor2tensor.rl.restarter import Restarter
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 TEST_MODE_1 = "mode1"
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index ffdcb5ef0..aebabf2a1 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -39,7 +39,7 @@
 from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import trainer_lib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def compute_mean_reward(rollouts, clipped):
diff --git a/tensor2tensor/visualization/visualization.py b/tensor2tensor/visualization/visualization.py
index 3519751b9..a476f1ea3 100644
--- a/tensor2tensor/visualization/visualization.py
+++ b/tensor2tensor/visualization/visualization.py
@@ -26,7 +26,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 EOS_ID = 1
 
diff --git a/tensor2tensor/visualization/visualization_test.py b/tensor2tensor/visualization/visualization_test.py
index 8ad86d107..06167d34a 100644
--- a/tensor2tensor/visualization/visualization_test.py
+++ b/tensor2tensor/visualization/visualization_test.py
@@ -33,7 +33,7 @@
 
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.visualization import visualization
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def get_data_dir():

From 8a084a4d56839733612d494db46503885cfc3962 Mon Sep 17 00:00:00 2001
From: Yanhua Sun <yanhuasun@google.com>
Date: Fri, 27 Dec 2019 14:17:42 -0800
Subject: [PATCH 2611/2720] Explicitly replace "import tensorflow" with
 "tensorflow.compat.v1"

PiperOrigin-RevId: 287340342
---
 tensor2tensor/utils/adafactor.py                     | 2 +-
 tensor2tensor/utils/adv_attack_utils.py              | 2 +-
 tensor2tensor/utils/beam_search.py                   | 2 +-
 tensor2tensor/utils/beam_search_test.py              | 2 +-
 tensor2tensor/utils/bleu_hook_test.py                | 2 +-
 tensor2tensor/utils/checkpoint_compatibility_test.py | 2 +-
 tensor2tensor/utils/data_reader_test.py              | 2 +-
 tensor2tensor/utils/diet.py                          | 2 +-
 tensor2tensor/utils/diet_test.py                     | 2 +-
 tensor2tensor/utils/expert_utils.py                  | 2 +-
 tensor2tensor/utils/expert_utils_test.py             | 2 +-
 tensor2tensor/utils/hparam_test.py                   | 2 +-
 tensor2tensor/utils/hparams_lib.py                   | 2 +-
 tensor2tensor/utils/hparams_lib_test.py              | 2 +-
 tensor2tensor/utils/learning_rate.py                 | 2 +-
 tensor2tensor/utils/metrics_hook.py                  | 2 +-
 tensor2tensor/utils/metrics_hook_test.py             | 2 +-
 tensor2tensor/utils/metrics_test.py                  | 2 +-
 tensor2tensor/utils/misc_utils_test.py               | 2 +-
 tensor2tensor/utils/mtf_model.py                     | 2 +-
 tensor2tensor/utils/multistep_optimizer.py           | 2 +-
 tensor2tensor/utils/multistep_optimizer_test.py      | 2 +-
 tensor2tensor/utils/optimize_test.py                 | 2 +-
 tensor2tensor/utils/partial_checkpoint_load_hook.py  | 2 +-
 tensor2tensor/utils/pruning_utils.py                 | 2 +-
 tensor2tensor/utils/quantization.py                  | 2 +-
 tensor2tensor/utils/registry.py                      | 2 +-
 tensor2tensor/utils/registry_test.py                 | 2 +-
 tensor2tensor/utils/rouge.py                         | 2 +-
 tensor2tensor/utils/rouge_test.py                    | 2 +-
 tensor2tensor/utils/sari_hook.py                     | 2 +-
 tensor2tensor/utils/sari_hook_test.py                | 2 +-
 tensor2tensor/utils/scheduled_sampling.py            | 2 +-
 tensor2tensor/utils/t2t_model_test.py                | 2 +-
 tensor2tensor/utils/test_utils.py                    | 2 +-
 tensor2tensor/utils/test_utils_test.py               | 2 +-
 tensor2tensor/utils/trainer_lib_test.py              | 2 +-
 tensor2tensor/utils/update_ops_hook.py               | 2 +-
 tensor2tensor/utils/usr_dir.py                       | 2 +-
 tensor2tensor/utils/video_metrics.py                 | 2 +-
 tensor2tensor/utils/video_metrics_test.py            | 2 +-
 tensor2tensor/utils/yellowfin.py                     | 2 +-
 tensor2tensor/utils/yellowfin_test.py                | 2 +-
 43 files changed, 43 insertions(+), 43 deletions(-)

diff --git a/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index bbb56f0e8..58525a740 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -21,7 +21,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import quantization
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class AdafactorOptimizer(tf.compat.v1.train.Optimizer):
diff --git a/tensor2tensor/utils/adv_attack_utils.py b/tensor2tensor/utils/adv_attack_utils.py
index b1ec90d9d..ddad56f59 100644
--- a/tensor2tensor/utils/adv_attack_utils.py
+++ b/tensor2tensor/utils/adv_attack_utils.py
@@ -23,7 +23,7 @@
 
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_attack
diff --git a/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index 85a32f8a1..53a931ba6 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -23,7 +23,7 @@
 import numpy as np
 
 from tensor2tensor.layers import common_layers
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.python.ops import inplace_ops
 from tensorflow.python.util import nest
diff --git a/tensor2tensor/utils/beam_search_test.py b/tensor2tensor/utils/beam_search_test.py
index 47991405d..88bc9b43e 100644
--- a/tensor2tensor/utils/beam_search_test.py
+++ b/tensor2tensor/utils/beam_search_test.py
@@ -21,7 +21,7 @@
 import numpy as np
 from tensor2tensor.utils import beam_search
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class BeamSearchTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/bleu_hook_test.py b/tensor2tensor/utils/bleu_hook_test.py
index ae086342d..e7ee9c264 100644
--- a/tensor2tensor/utils/bleu_hook_test.py
+++ b/tensor2tensor/utils/bleu_hook_test.py
@@ -27,7 +27,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.utils import bleu_hook
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class BleuHookTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/checkpoint_compatibility_test.py b/tensor2tensor/utils/checkpoint_compatibility_test.py
index 759105b99..91e7ff363 100644
--- a/tensor2tensor/utils/checkpoint_compatibility_test.py
+++ b/tensor2tensor/utils/checkpoint_compatibility_test.py
@@ -37,7 +37,7 @@
 from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import trainer_lib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def get_data_dir():
diff --git a/tensor2tensor/utils/data_reader_test.py b/tensor2tensor/utils/data_reader_test.py
index ef4c1f314..4c60b556e 100644
--- a/tensor2tensor/utils/data_reader_test.py
+++ b/tensor2tensor/utils/data_reader_test.py
@@ -30,7 +30,7 @@
 from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_problem
diff --git a/tensor2tensor/utils/diet.py b/tensor2tensor/utils/diet.py
index fa06b9902..61a20e3aa 100644
--- a/tensor2tensor/utils/diet.py
+++ b/tensor2tensor/utils/diet.py
@@ -28,7 +28,7 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import hparam
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def diet_adam_optimizer_params():
diff --git a/tensor2tensor/utils/diet_test.py b/tensor2tensor/utils/diet_test.py
index 0a41e58f1..e523cfde0 100644
--- a/tensor2tensor/utils/diet_test.py
+++ b/tensor2tensor/utils/diet_test.py
@@ -20,7 +20,7 @@
 from __future__ import print_function
 from tensor2tensor.utils import diet
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class DietVarTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index f177f2ac9..f8880e4a0 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -32,7 +32,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers.vq_discrete import DiscreteBottleneck
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 DEFAULT_DEV_STRING = "existing_device"
 
diff --git a/tensor2tensor/utils/expert_utils_test.py b/tensor2tensor/utils/expert_utils_test.py
index 3c6feec2f..48ab95e1a 100644
--- a/tensor2tensor/utils/expert_utils_test.py
+++ b/tensor2tensor/utils/expert_utils_test.py
@@ -20,7 +20,7 @@
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.utils import expert_utils
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class ExpertUtilsTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/hparam_test.py b/tensor2tensor/utils/hparam_test.py
index bea9ee36f..637299ad5 100644
--- a/tensor2tensor/utils/hparam_test.py
+++ b/tensor2tensor/utils/hparam_test.py
@@ -21,7 +21,7 @@
 
 from tensor2tensor.utils import hparam
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class HParamsTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/hparams_lib.py b/tensor2tensor/utils/hparams_lib.py
index 07be4e4f6..891fee943 100644
--- a/tensor2tensor/utils/hparams_lib.py
+++ b/tensor2tensor/utils/hparams_lib.py
@@ -25,7 +25,7 @@
 from tensor2tensor.utils import hparam
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def copy_hparams(hparams):
diff --git a/tensor2tensor/utils/hparams_lib_test.py b/tensor2tensor/utils/hparams_lib_test.py
index 589caee1b..e98cff0aa 100644
--- a/tensor2tensor/utils/hparams_lib_test.py
+++ b/tensor2tensor/utils/hparams_lib_test.py
@@ -23,7 +23,7 @@
 
 from tensor2tensor.utils import hparams_lib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class HparamsLibTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index 26b397f8b..d8bcd6a2d 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -20,7 +20,7 @@
 import numpy as np
 
 from tensor2tensor.utils import mlperf_log
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def learning_rate_factor(name, step_num, hparams):
diff --git a/tensor2tensor/utils/metrics_hook.py b/tensor2tensor/utils/metrics_hook.py
index 6186b1df6..640904fc1 100644
--- a/tensor2tensor/utils/metrics_hook.py
+++ b/tensor2tensor/utils/metrics_hook.py
@@ -19,7 +19,7 @@
 from __future__ import print_function
 
 import os
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorboard.backend.event_processing import event_accumulator
 from tensorboard.backend.event_processing import event_multiplexer
diff --git a/tensor2tensor/utils/metrics_hook_test.py b/tensor2tensor/utils/metrics_hook_test.py
index 3069a8f8d..74739730f 100644
--- a/tensor2tensor/utils/metrics_hook_test.py
+++ b/tensor2tensor/utils/metrics_hook_test.py
@@ -24,7 +24,7 @@
 import shutil
 from tensor2tensor.utils import metrics_hook
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class DummyHook(metrics_hook.MetricsBasedHook):
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 05c62b6c2..1fd4c7719 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -21,7 +21,7 @@
 import numpy as np
 from tensor2tensor.utils import metrics
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class MetricsTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/misc_utils_test.py b/tensor2tensor/utils/misc_utils_test.py
index 6beedb098..282cef0c3 100644
--- a/tensor2tensor/utils/misc_utils_test.py
+++ b/tensor2tensor/utils/misc_utils_test.py
@@ -21,7 +21,7 @@
 
 from tensor2tensor.utils import hparam
 from tensor2tensor.utils import misc_utils
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class MiscUtilsTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/mtf_model.py b/tensor2tensor/utils/mtf_model.py
index d39fe9902..607ead673 100644
--- a/tensor2tensor/utils/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -28,7 +28,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.contrib.tpu.python.tpu import tpu_estimator
 
diff --git a/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
index a9195b22d..f745a5d40 100644
--- a/tensor2tensor/utils/multistep_optimizer.py
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -26,7 +26,7 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class MultistepAdamOptimizer(tf.compat.v1.train.AdamOptimizer):
diff --git a/tensor2tensor/utils/multistep_optimizer_test.py b/tensor2tensor/utils/multistep_optimizer_test.py
index c11fb21b6..45cc7b15e 100644
--- a/tensor2tensor/utils/multistep_optimizer_test.py
+++ b/tensor2tensor/utils/multistep_optimizer_test.py
@@ -20,7 +20,7 @@
 
 import numpy as np
 from tensor2tensor.utils import multistep_optimizer
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class MultistepAdamOptimizerTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/optimize_test.py b/tensor2tensor/utils/optimize_test.py
index 1a24a9d34..e882b4263 100644
--- a/tensor2tensor/utils/optimize_test.py
+++ b/tensor2tensor/utils/optimize_test.py
@@ -22,7 +22,7 @@
 from absl.testing import parameterized
 from tensor2tensor.utils import hparams_lib
 from tensor2tensor.utils import optimize
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class OptimizeTest(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/utils/partial_checkpoint_load_hook.py b/tensor2tensor/utils/partial_checkpoint_load_hook.py
index 70b6c566b..0e8f42e66 100644
--- a/tensor2tensor/utils/partial_checkpoint_load_hook.py
+++ b/tensor2tensor/utils/partial_checkpoint_load_hook.py
@@ -18,7 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class PartialCheckpointLoad(tf.train.SessionRunHook):
diff --git a/tensor2tensor/utils/pruning_utils.py b/tensor2tensor/utils/pruning_utils.py
index 1023ce904..835c727f8 100644
--- a/tensor2tensor/utils/pruning_utils.py
+++ b/tensor2tensor/utils/pruning_utils.py
@@ -20,7 +20,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_pruning_strategy
diff --git a/tensor2tensor/utils/quantization.py b/tensor2tensor/utils/quantization.py
index 06e84ec5c..95000fadb 100644
--- a/tensor2tensor/utils/quantization.py
+++ b/tensor2tensor/utils/quantization.py
@@ -17,7 +17,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.python.framework import function
 
diff --git a/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 7786b5ec7..603b73e0e 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -69,7 +69,7 @@ class MyModel(T2TModel):
 import collections
 
 from tensor2tensor.utils import misc_utils
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.python.util import tf_inspect as inspect  # pylint: disable=g-direct-tensorflow-import
 
diff --git a/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index adf4bfccd..f5f0b7d09 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -22,7 +22,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 # pylint: disable=unused-variable,unused-argument
diff --git a/tensor2tensor/utils/rouge.py b/tensor2tensor/utils/rouge.py
index c3e6c7ebc..d18504d42 100644
--- a/tensor2tensor/utils/rouge.py
+++ b/tensor2tensor/utils/rouge.py
@@ -27,7 +27,7 @@
 
 import numpy as np
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def _len_lcs(x, y):
diff --git a/tensor2tensor/utils/rouge_test.py b/tensor2tensor/utils/rouge_test.py
index 5dc6a42cf..08005984e 100644
--- a/tensor2tensor/utils/rouge_test.py
+++ b/tensor2tensor/utils/rouge_test.py
@@ -21,7 +21,7 @@
 import numpy as np
 from tensor2tensor.utils import rouge
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TestRouge2Metric(tf.test.TestCase):
diff --git a/tensor2tensor/utils/sari_hook.py b/tensor2tensor/utils/sari_hook.py
index fba2899ea..df611dec7 100644
--- a/tensor2tensor/utils/sari_hook.py
+++ b/tensor2tensor/utils/sari_hook.py
@@ -39,7 +39,7 @@
 import collections
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # The paper that intoduces the SARI score uses only the precision of the deleted
 # tokens (i.e. beta=0). To give more emphasis on recall, you may set, e.g.,
diff --git a/tensor2tensor/utils/sari_hook_test.py b/tensor2tensor/utils/sari_hook_test.py
index e3c9a3aa8..3be960534 100644
--- a/tensor2tensor/utils/sari_hook_test.py
+++ b/tensor2tensor/utils/sari_hook_test.py
@@ -23,7 +23,7 @@
 
 import numpy as np
 from tensor2tensor.utils import sari_hook
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class SariHookTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/scheduled_sampling.py b/tensor2tensor/utils/scheduled_sampling.py
index 80def28bd..a26d20d66 100644
--- a/tensor2tensor/utils/scheduled_sampling.py
+++ b/tensor2tensor/utils/scheduled_sampling.py
@@ -36,7 +36,7 @@
 import copy
 
 from tensor2tensor.layers import common_layers
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.python.ops import inplace_ops  # pylint: disable=g-direct-tensorflow-import
 
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index b914d6607..a204b7cc6 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -24,7 +24,7 @@
 from tensor2tensor.utils import t2t_model
 from tensor2tensor.utils import test_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 tf.compat.v1.enable_eager_execution()
 
 
diff --git a/tensor2tensor/utils/test_utils.py b/tensor2tensor/utils/test_utils.py
index 292949f6e..81fcd5a6f 100644
--- a/tensor2tensor/utils/test_utils.py
+++ b/tensor2tensor/utils/test_utils.py
@@ -19,7 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def run_in_graph_and_eager_modes(func=None,
diff --git a/tensor2tensor/utils/test_utils_test.py b/tensor2tensor/utils/test_utils_test.py
index f9f07502e..0b29391de 100644
--- a/tensor2tensor/utils/test_utils_test.py
+++ b/tensor2tensor/utils/test_utils_test.py
@@ -21,7 +21,7 @@
 
 from tensor2tensor.utils import test_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 tf.compat.v1.enable_eager_execution()
 
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index 0c7ffaeb4..faf6d8090 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TrainerLibTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/update_ops_hook.py b/tensor2tensor/utils/update_ops_hook.py
index f3252060b..46ac18c84 100644
--- a/tensor2tensor/utils/update_ops_hook.py
+++ b/tensor2tensor/utils/update_ops_hook.py
@@ -18,7 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class UpdateOpsHook(tf.train.SessionRunHook):
diff --git a/tensor2tensor/utils/usr_dir.py b/tensor2tensor/utils/usr_dir.py
index 6817e9831..c349c4f2c 100644
--- a/tensor2tensor/utils/usr_dir.py
+++ b/tensor2tensor/utils/usr_dir.py
@@ -21,7 +21,7 @@
 import importlib
 import os
 import sys
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 INTERNAL_USR_DIR_PACKAGE = "t2t_usr_dir_internal"
diff --git a/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index 4515bb808..fe4d8638a 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -24,7 +24,7 @@
 import six
 
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def load_image_map_function(filename, frame_shape):
diff --git a/tensor2tensor/utils/video_metrics_test.py b/tensor2tensor/utils/video_metrics_test.py
index b3f7b53fa..24da20c46 100644
--- a/tensor2tensor/utils/video_metrics_test.py
+++ b/tensor2tensor/utils/video_metrics_test.py
@@ -21,7 +21,7 @@
 
 import numpy as np
 from tensor2tensor.utils import video_metrics
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class VideoMetricsTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/yellowfin.py b/tensor2tensor/utils/yellowfin.py
index 8896801c2..eb50ad01b 100644
--- a/tensor2tensor/utils/yellowfin.py
+++ b/tensor2tensor/utils/yellowfin.py
@@ -18,7 +18,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 # Values for gate_gradients.
diff --git a/tensor2tensor/utils/yellowfin_test.py b/tensor2tensor/utils/yellowfin_test.py
index 0a59c3917..693394362 100644
--- a/tensor2tensor/utils/yellowfin_test.py
+++ b/tensor2tensor/utils/yellowfin_test.py
@@ -22,7 +22,7 @@
 
 from tensor2tensor.utils.yellowfin import YellowFinOptimizer
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 n_dim = 1000000

From 1834e5614221163a63656a9f9f21a55f13d5bee4 Mon Sep 17 00:00:00 2001
From: Manoj Kumar <mechcoder@google.com>
Date: Thu, 2 Jan 2020 02:41:58 -0800
Subject: [PATCH 2612/2720] internal

PiperOrigin-RevId: 287810815
---
 tensor2tensor/bin/t2t_trainer.py              |   3 +-
 tensor2tensor/data_generators/allen_brain.py  |   4 +-
 .../data_generators/allen_brain_test.py       |   4 +-
 .../data_generators/bair_robot_pushing.py     |  13 +-
 tensor2tensor/data_generators/fsns.py         |   5 +-
 tensor2tensor/data_generators/gym_env.py      |   3 +-
 tensor2tensor/data_generators/image_utils.py  |  13 +-
 tensor2tensor/data_generators/moving_mnist.py |   5 +-
 tensor2tensor/data_generators/problem.py      |   9 +-
 tensor2tensor/data_generators/translate.py    |   5 +-
 .../data_generators/video_generated.py        |   5 +-
 tensor2tensor/data_generators/video_utils.py  |   9 +-
 tensor2tensor/data_generators/vqa.py          |  45 +++--
 tensor2tensor/envs/env_problem.py             |   3 +-
 tensor2tensor/envs/rendered_env_problem.py    |   7 +-
 tensor2tensor/layers/common_attention.py      |   5 +-
 tensor2tensor/layers/common_attention_test.py |  18 +-
 tensor2tensor/layers/common_layers.py         |  19 ++-
 tensor2tensor/layers/common_layers_test.py    |   2 +-
 tensor2tensor/layers/common_video.py          |  17 +-
 tensor2tensor/layers/vqa_layers.py            |   5 +-
 tensor2tensor/models/lstm.py                  |  11 +-
 .../neural_architecture_search/nas_model.py   |   5 +-
 tensor2tensor/models/research/attention_lm.py |   8 +-
 .../models/research/gene_expression.py        |   3 +-
 tensor2tensor/models/research/glow.py         |   5 +-
 tensor2tensor/models/research/glow_ops.py     |   5 +-
 .../models/research/glow_ops_test.py          |   9 +-
 tensor2tensor/models/research/neural_stack.py |  12 +-
 .../models/research/neural_stack_test.py      |   3 +-
 .../models/research/transformer_revnet.py     |   5 +-
 .../models/research/transformer_vae.py        |   3 +-
 .../models/research/universal_transformer.py  |   5 +-
 .../research/universal_transformer_util.py    |  37 +++--
 .../models/research/vqa_attention.py          |   3 +-
 tensor2tensor/models/revnet.py                |   8 +-
 tensor2tensor/models/video/emily.py           |   3 +-
 tensor2tensor/models/video/epva.py            |  27 +--
 tensor2tensor/models/video/next_frame_glow.py |   4 +-
 tensor2tensor/models/video/nfg_interpolate.py |   3 +-
 tensor2tensor/models/video/savp.py            |   3 +-
 tensor2tensor/models/video/sv2p.py            |   3 +-
 tensor2tensor/serving/serving_utils.py        |   2 +
 tensor2tensor/utils/contrib.py                | 156 ++++++++++++++++++
 tensor2tensor/utils/data_reader.py            |   3 +-
 tensor2tensor/utils/decoding.py               |   7 +-
 tensor2tensor/utils/metrics.py                |   6 +-
 tensor2tensor/utils/optimize.py               |  15 +-
 tensor2tensor/utils/restore_hook.py           |   3 +-
 tensor2tensor/utils/t2t_model.py              |  21 +--
 tensor2tensor/utils/trainer_lib.py            |  24 +--
 51 files changed, 396 insertions(+), 205 deletions(-)
 create mode 100644 tensor2tensor/utils/contrib.py

diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 0d7c7529a..dc4ad34a6 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -26,6 +26,7 @@
 from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
 
 from tensor2tensor.utils import cloud_mlengine
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
 from tensor2tensor.utils import hparams_lib
@@ -304,7 +305,7 @@ def generate_data():
 @contextlib.contextmanager
 def profile_context():
   if FLAGS.profile:
-    with tf.contrib.tfprof.ProfileContext(
+    with contrib.tfprof().ProfileContext(
         "t2tprof", trace_steps=range(100), dump_steps=range(100)) as pctx:
       opts = tf.profiler.ProfileOptionBuilder.time_and_memory()
       pctx.add_auto_profiling("op", opts, range(100))
diff --git a/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index 2f58de3e9..67afce1dd 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -44,6 +44,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
+from tensorflow.contrib import slim as contrib_slim
 
 _BASE_EXAMPLE_IMAGE_SIZE = 64
 
@@ -350,7 +351,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "targets":
-            tf.contrib.slim.tfexample_decoder.Image(
+            contrib_slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
@@ -442,4 +443,3 @@ def output_dim(self):
   @property
   def inpaint_fraction(self):
     return 0.01
-
diff --git a/tensor2tensor/data_generators/allen_brain_test.py b/tensor2tensor/data_generators/allen_brain_test.py
index e2c875789..487517ab6 100644
--- a/tensor2tensor/data_generators/allen_brain_test.py
+++ b/tensor2tensor/data_generators/allen_brain_test.py
@@ -23,10 +23,12 @@
 
 from tensor2tensor.data_generators import allen_brain
 from tensor2tensor.models import image_transformer_2d
+from tensor2tensor.utils import contrib
 
 import tensorflow as tf
 
-tfe = tf.contrib.eager
+
+tfe = contrib.eager()
 tfe.enable_eager_execution()
 Modes = tf.estimator.ModeKeys  # pylint: disable=invalid-name
 
diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index eadd3275f..ac3ae4152 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -33,6 +33,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.layers import modalities
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -102,8 +103,8 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
-            tensor_key="frame_number"),
+        "frame_number":
+            contrib.slim().tfexample_decoder.Tensor(tensor_key="frame_number"),
     }
     return data_fields, decoders
 
@@ -187,9 +188,9 @@ def extra_reading_spec(self):
         "action": tf.FixedLenFeature([4], tf.float32),
     }
     decoders = {
-        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
-            tensor_key="frame_number"),
-        "action": tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="action"),
+        "frame_number":
+            contrib.slim().tfexample_decoder.Tensor(tensor_key="frame_number"),
+        "action":
+            contrib.slim().tfexample_decoder.Tensor(tensor_key="action"),
     }
     return data_fields, decoders
-
diff --git a/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
index 8dff69547..2ae9a2734 100644
--- a/tensor2tensor/data_generators/fsns.py
+++ b/tensor2tensor/data_generators/fsns.py
@@ -25,6 +25,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.layers import modalities
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -75,6 +76,6 @@ def example_reading_spec(self):
     data_fields, data_items_to_decoders = (
         super(ImageFSNS, self).example_reading_spec())
     data_fields[label_key] = tf.VarLenFeature(tf.int64)
-    data_items_to_decoders[
-        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
+    data_items_to_decoders["targets"] = contrib.slim().tfexample_decoder.Tensor(
+        label_key)
     return data_fields, data_items_to_decoders
diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 67aedef9b..fa6ad4d17 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -31,6 +31,7 @@
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.layers import modalities
 from tensor2tensor.rl import gym_utils
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import registry
@@ -377,7 +378,7 @@ def extra_reading_spec(self):
         name: tf.FixedLenFeature([1], tf.int64) for name in field_names
     }
     decoders = {
-        name: tf.contrib.slim.tfexample_decoder.Tensor(tensor_key=name)
+        name: contrib.slim().tfexample_decoder.Tensor(tensor_key=name)
         for name in field_names
     }
     return (data_fields, decoders)
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 61766a89e..eed55809c 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -28,6 +28,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import modalities
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import metrics
 
 import tensorflow as tf
@@ -172,7 +173,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "inputs":
-            tf.contrib.slim.tfexample_decoder.Image(
+            contrib.slim().tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
@@ -238,8 +239,8 @@ def example_reading_spec(self):
         super(Image2ClassProblem, self).example_reading_spec())
     data_fields[label_key] = tf.FixedLenFeature((1,), tf.int64)
 
-    data_items_to_decoders[
-        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
+    data_items_to_decoders["targets"] = contrib.slim().tfexample_decoder.Tensor(
+        label_key)
     return data_fields, data_items_to_decoders
 
   def hparams(self, defaults, unused_model_hparams):
@@ -342,8 +343,8 @@ def example_reading_spec(self):
     data_fields, data_items_to_decoders = (
         super(Image2TextProblem, self).example_reading_spec())
     data_fields[label_key] = tf.VarLenFeature(tf.int64)
-    data_items_to_decoders[
-        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
+    data_items_to_decoders["targets"] = contrib.slim().tfexample_decoder.Tensor(
+        label_key)
     return data_fields, data_items_to_decoders
 
   def feature_encoders(self, data_dir):
@@ -422,4 +423,4 @@ def random_shift(image, wsr=0.1, hsr=0.1):
   height_translations = tf.random_uniform((1,), -height_range, height_range)
   width_translations = tf.random_uniform((1,), -width_range, width_range)
   translations = tf.concat((height_translations, width_translations), axis=0)
-  return tf.contrib.image.translate(image, translations=translations)
+  return contrib.image().translate(image, translations=translations)
diff --git a/tensor2tensor/data_generators/moving_mnist.py b/tensor2tensor/data_generators/moving_mnist.py
index 507207623..1dfc82313 100644
--- a/tensor2tensor/data_generators/moving_mnist.py
+++ b/tensor2tensor/data_generators/moving_mnist.py
@@ -32,6 +32,7 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.layers import modalities
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -94,8 +95,8 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
-            tensor_key="frame_number"),
+        "frame_number":
+            contrib.slim().tfexample_decoder.Tensor(tensor_key="frame_number"),
     }
     return data_fields, decoders
 
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index aca91f31f..7732225ce 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -26,6 +26,7 @@
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import hparam
 from tensor2tensor.utils import metrics
@@ -648,8 +649,8 @@ def dataset(self,
 
     data_filepattern = self.filepattern(data_dir, dataset_split, shard=shard)
     tf.logging.info("Reading data files from %s", data_filepattern)
-    data_files = sorted(tf.contrib.slim.parallel_reader.get_data_files(
-        data_filepattern))
+    data_files = sorted(
+        contrib.slim().parallel_reader.get_data_files(data_filepattern))
 
     # Functions used in dataset transforms below. `filenames` can be either a
     # `tf.string` tensor or `tf.data.Dataset` containing one or more filenames.
@@ -711,11 +712,11 @@ def decode_example(self, serialized_example):
     data_fields["batch_prediction_key"] = tf.FixedLenFeature([1], tf.int64, 0)
     if data_items_to_decoders is None:
       data_items_to_decoders = {
-          field: tf.contrib.slim.tfexample_decoder.Tensor(field)
+          field: contrib.slim().tfexample_decoder.Tensor(field)
           for field in data_fields
       }
 
-    decoder = tf.contrib.slim.tfexample_decoder.TFExampleDecoder(
+    decoder = contrib.slim().tfexample_decoder.TFExampleDecoder(
         data_fields, data_items_to_decoders)
 
     decode_items = list(sorted(data_items_to_decoders))
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 05e07d7c3..487c6f143 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -29,6 +29,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import bleu_hook
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
@@ -276,8 +277,8 @@ def example_reading_spec(self):
 
     # hack: ignoring true targets and putting dist_targets in targets
     data_items_to_decoders = {
-        "inputs": tf.contrib.slim.tfexample_decoder.Tensor("inputs"),
-        "targets": tf.contrib.slim.tfexample_decoder.Tensor("dist_targets"),
+        "inputs": contrib.slim().tfexample_decoder.Tensor("inputs"),
+        "targets": contrib.slim().tfexample_decoder.Tensor("dist_targets"),
     }
 
     return (data_fields, data_items_to_decoders)
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index aa0e5bcf3..c7a5a6afa 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -25,6 +25,7 @@
 
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.layers import modalities
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -85,8 +86,8 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
-            tensor_key="frame_number"),
+        "frame_number":
+            contrib.slim().tfexample_decoder.Tensor(tensor_key="frame_number"),
     }
     return data_fields, decoders
 
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index ae2d2b313..b816ef8dc 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -35,6 +35,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import video_metrics
 import tensorflow as tf
+from tensorflow.contrib import slim as contrib_slim
 
 
 FLAGS = flags.FLAGS
@@ -384,7 +385,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "frame":
-            tf.contrib.slim.tfexample_decoder.Image(
+            contrib_slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 shape=[self.frame_height, self.frame_width, self.num_channels],
@@ -676,7 +677,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "inputs":
-            tf.contrib.slim.tfexample_decoder.Image(
+            contrib_slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
@@ -765,8 +766,8 @@ def example_reading_spec(self):
     data_fields, data_items_to_decoders = (
         super(Video2ClassProblem, self).example_reading_spec())
     data_fields[label_key] = tf.FixedLenFeature((1,), tf.int64)
-    data_items_to_decoders[
-        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
+    data_items_to_decoders["targets"] = contrib_slim.tfexample_decoder.Tensor(
+        label_key)
     return data_fields, data_items_to_decoders
 
   def hparams(self, defaults, unused_model_hparams):
diff --git a/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
index 8a39f85f6..783200f54 100644
--- a/tensor2tensor/data_generators/vqa.py
+++ b/tensor2tensor/data_generators/vqa.py
@@ -35,6 +35,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import vqa_utils
 from tensor2tensor.layers import modalities
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
@@ -216,12 +217,11 @@ def example_reading_spec(self):
     data_fields["image/answer"] = tf.FixedLenSequenceFeature(
         (), tf.int64, allow_missing=True)
 
-    data_items_to_decoders[
-        "question"] = tf.contrib.slim.tfexample_decoder.Tensor(
-            "image/question")
-    data_items_to_decoders[
-        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(
-            "image/answer")
+    slim = contrib.slim()
+    data_items_to_decoders["question"] = slim.tfexample_decoder.Tensor(
+        "image/question")
+    data_items_to_decoders["targets"] = slim.tfexample_decoder.Tensor(
+        "image/answer")
     return data_fields, data_items_to_decoders
 
   def preprocess_example(self, example, mode, hparams):
@@ -325,6 +325,7 @@ def preprocess_example(self, example, mode, hparams):
     return example
 
   def example_reading_spec(self):
+    slim = contrib.slim()
     data_fields, data_items_to_decoders = {}, {}
     data_fields["image/feature"] = tf.FixedLenSequenceFeature(
         (), tf.float32, allow_missing=True)
@@ -337,25 +338,19 @@ def example_reading_spec(self):
     data_fields["image/answer"] = tf.FixedLenSequenceFeature(
         (), tf.int64, allow_missing=True)
 
-    data_items_to_decoders[
-        "inputs"] = tf.contrib.slim.tfexample_decoder.Tensor(
-            "image/feature")
-    data_items_to_decoders[
-        "question_id"] = tf.contrib.slim.tfexample_decoder.Tensor(
-            "image/question_id")
-    data_items_to_decoders[
-        "image_id"] = tf.contrib.slim.tfexample_decoder.Tensor(
-            "image/image_id")
-
-    data_items_to_decoders[
-        "spatial_feature"] = tf.contrib.slim.tfexample_decoder.Tensor(
-            "image/spatial_feature")
-    data_items_to_decoders[
-        "question"] = tf.contrib.slim.tfexample_decoder.Tensor(
-            "image/question")
-    data_items_to_decoders[
-        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(
-            "image/answer")
+    data_items_to_decoders["inputs"] = slim.tfexample_decoder.Tensor(
+        "image/feature")
+    data_items_to_decoders["question_id"] = slim.tfexample_decoder.Tensor(
+        "image/question_id")
+    data_items_to_decoders["image_id"] = slim.tfexample_decoder.Tensor(
+        "image/image_id")
+
+    data_items_to_decoders["spatial_feature"] = slim.tfexample_decoder.Tensor(
+        "image/spatial_feature")
+    data_items_to_decoders["question"] = slim.tfexample_decoder.Tensor(
+        "image/question")
+    data_items_to_decoders["targets"] = slim.tfexample_decoder.Tensor(
+        "image/answer")
 
     return data_fields, data_items_to_decoders
 
diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 43afb1a75..f85fa1f48 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -32,6 +32,7 @@
 from tensor2tensor.envs import gym_spaces_utils
 from tensor2tensor.envs import trajectory
 from tensor2tensor.layers import modalities
+from tensor2tensor.utils import contrib
 import tensorflow as tf
 
 # Names for data fields in stored tf.Examples.
@@ -476,7 +477,7 @@ def example_reading_spec(self):
     }
 
     data_items_to_decoders = {
-        field: tf.contrib.slim.tfexample_decoder.Tensor(field)
+        field: contrib.slim().tfexample_decoder.Tensor(field)
         for field in data_fields
     }
 
diff --git a/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
index 92a7a7e54..34d03a357 100644
--- a/tensor2tensor/envs/rendered_env_problem.py
+++ b/tensor2tensor/envs/rendered_env_problem.py
@@ -24,6 +24,7 @@
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import gym_env_problem
+from tensor2tensor.utils import contrib
 import tensorflow as tf
 
 _IMAGE_ENCODED_FIELD = "image/encoded"
@@ -68,6 +69,7 @@ def initialize_environments(self,
 
   def example_reading_spec(self):
     """Return a mix of env and video data fields and decoders."""
+    slim = contrib.slim()
     video_fields, video_decoders = (
         video_utils.VideoProblem.example_reading_spec(self))
     env_fields, env_decoders = (
@@ -79,9 +81,8 @@ def example_reading_spec(self):
 
     # Add frame number spec and decoder.
     env_fields[_FRAME_NUMBER_FIELD] = tf.FixedLenFeature((1,), tf.int64)
-    env_decoders[
-        _FRAME_NUMBER_FIELD] = tf.contrib.slim.tfexample_decoder.Tensor(
-            _FRAME_NUMBER_FIELD)
+    env_decoders[_FRAME_NUMBER_FIELD] = slim.tfexample_decoder.Tensor(
+        _FRAME_NUMBER_FIELD)
 
     # Add video fields and decoders
     env_fields.update(video_fields)
diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index d0062083e..409b978c0 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -31,6 +31,7 @@
 
 from tensor2tensor.layers import area_attention
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import expert_utils
 
 import tensorflow as tf
@@ -2775,7 +2776,7 @@ def masked_within_block_local_attention_1d(q, k, v, block_length=64, name=None):
     batch, heads, length, depth_k = common_layers.shape_list(q)
     depth_v = common_layers.shape_list(v)[-1]
     if isinstance(block_length, tf.Tensor):
-      const = tf.contrib.util.constant_value(block_length)
+      const = contrib.util().constant_value(block_length)
       if const is not None:
         block_length = int(const)
 
@@ -2869,7 +2870,7 @@ def masked_local_attention_1d(q,
     batch, heads, length, depth_k = common_layers.shape_list(q)
     depth_v = common_layers.shape_list(v)[-1]
     if isinstance(block_length, tf.Tensor):
-      const = tf.contrib.util.constant_value(block_length)
+      const = contrib.util().constant_value(block_length)
       if const is not None:
         block_length = int(const)
     # If (length < 2 * block_length), then we use only one block.
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 3be2ba409..aab56b5bc 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -26,8 +26,14 @@
 import numpy as np
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import test_utils
+
 import tensorflow as tf
+
+
+tfe = contrib.tfe()
+# from tensorflow.contrib.eager.python import tfe as tfe
 tf.compat.v1.enable_eager_execution()
 
 
@@ -931,7 +937,7 @@ def testDotProductUnMaskedAttentionRelativeV2(self):
     res = self.evaluate(a)
     self.assertEqual(res.shape, (5, 7, 12, 32))
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @tfe.run_test_in_graph_and_eager_modes()
   def testExtractblocks(self):
 
     batch_size = 1
@@ -985,7 +991,7 @@ def python_get_2d_local_memory(self, t, batch_size, num_heads, height, width,
                                                  memory_w_index]
     return out
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @tfe.run_test_in_graph_and_eager_modes()
   def testGet2dLocalMemory(self):
     batch_size = 3
     num_heads = 3
@@ -1019,7 +1025,7 @@ def testGet2dLocalMemory(self):
 
     self.assertAllClose(res, out)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @tfe.run_test_in_graph_and_eager_modes()
   def testSplitAlongWidth(self):
     batch_size = 1
     num_heads = 3
@@ -1063,7 +1069,7 @@ def testSplitAlongWidth(self):
     self.assertAllClose(res_l, out_l)
     self.assertAllClose(res_r, out_r)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @tfe.run_test_in_graph_and_eager_modes()
   def testGetLeftRightBlocks(self):
     batch_size = 1
     num_heads = 3
@@ -1118,7 +1124,7 @@ def testGetLeftRightBlocks(self):
     self.assertAllClose(res_l, out_l)
     self.assertAllClose(res_r, out_r)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @tfe.run_test_in_graph_and_eager_modes()
   def testDotProductUnmaskedAttentionLocal2dTpu(self):
     batch_size = 1
     num_heads = 3
@@ -1205,7 +1211,7 @@ def testDotProductUnmaskedAttentionLocal2dTpu(self):
     out = out[:, :, :height, :width, :]
     self.assertAllClose(res, out)
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
+  @tfe.run_test_in_graph_and_eager_modes()
   def testDotProductUnmaskedAttentionLocal2dTpuSimple(self):
     batch_size = 1
     num_heads = 3
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index aae4a6b4c..ca6621635 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -27,6 +27,7 @@
 import numpy as np
 from six.moves import range  # pylint: disable=redefined-builtin
 
+from tensor2tensor.utils import contrib
 import tensorflow as tf
 import tensorflow_probability as tfp
 
@@ -2761,7 +2762,7 @@ def _fn_with_custom_grad(fn, inputs, grad_fn, use_global_vars=False):
 
   def custom_grad_fn(op, *dys):
     """Custom grad fn applying grad_fn for identity Defun."""
-    fn_inputs, fn_vars, fn_outputs = tf.contrib.framework.nest.pack_sequence_as(
+    fn_inputs, fn_vars, fn_outputs = contrib.framework().nest.pack_sequence_as(
         defun_inputs, list(op.inputs))
     dys = list(dys)
     assert len(fn_outputs) == len(outputs)
@@ -2785,10 +2786,10 @@ def custom_grad_fn(op, *dys):
       python_grad_func=custom_grad_fn,
       shape_func=lambda _: [t.get_shape() for t in outputs])
   def identity(*args):
-    _, _, outs = tf.contrib.framework.nest.pack_sequence_as(defun_inputs, args)
+    _, _, outs = contrib.framework().nest.pack_sequence_as(defun_inputs, args)
     return tuple([tf.identity(t) for t in outs])
 
-  flat_inputs = tf.contrib.framework.nest.flatten(defun_inputs)
+  flat_inputs = contrib.framework().nest.flatten(defun_inputs)
   id_out = identity(*flat_inputs)
   return id_out
 
@@ -2956,7 +2957,7 @@ def sample_with_temperature(logits, temperature, sampling_keep_top_k=-1):
 
       vocab_size = shape_list(logits)[1]
 
-      k_largest = tf.contrib.nn.nth_element(
+      k_largest = contrib.nn().nth_element(
           logits, n=sampling_keep_top_k, reverse=True)
       k_largest = tf.tile(tf.reshape(k_largest, [-1, 1]), [1, vocab_size])
 
@@ -3050,7 +3051,7 @@ def grad_fn(inputs, variables, outputs, output_grads):
     variables = [underlying_variable_ref(v) for v in variables]
     # Recompute outputs
     with tf.control_dependencies(output_grads):
-      with tf.contrib.framework.arg_scope(cached_arg_scope[0]):
+      with contrib.framework().arg_scope(cached_arg_scope[0]):
         with tf.variable_scope(cached_vs[0], reuse=True):
           outputs = fn(*inputs)
 
@@ -3070,7 +3071,7 @@ def grad_fn(inputs, variables, outputs, output_grads):
   @fn_with_custom_grad(grad_fn)
   def fn_with_recompute(*args):
     cached_vs.append(tf.compat.v1.get_variable_scope())
-    cached_arg_scope.append(tf.contrib.framework.current_arg_scope())
+    cached_arg_scope.append(contrib.framework().current_arg_scope())
     return fn(*args)
 
   return fn_with_recompute(*args)
@@ -3411,7 +3412,7 @@ def should_generate_summaries():
   Returns:
     a boolean
   """
-  name_scope = tf.contrib.framework.get_name_scope()
+  name_scope = contrib.framework().get_name_scope()
   if name_scope and "while/" in name_scope:
     # Summaries don't work well within tf.while_loop()
     return False
@@ -3826,7 +3827,7 @@ def weight_targeting(w, k):
   w = tf.reshape(w, [size, w_shape[-1]])
 
   transpose_w = tf.transpose(w)
-  thres = tf.contrib.framework.sort(tf.abs(transpose_w), axis=1)[:, k]
+  thres = contrib.framework().sort(tf.abs(transpose_w), axis=1)[:, k]
   mask = to_float(thres[None, :] >= tf.abs(w))
 
   return tf.reshape(mask, w_shape)
@@ -3840,7 +3841,7 @@ def unit_targeting(w, k):
   w = tf.reshape(w, [size, w_shape[-1]])
 
   norm = tf.norm(w, axis=0)
-  thres = tf.contrib.framework.sort(norm, axis=0)[k]
+  thres = contrib.framework().sort(norm, axis=0)[k]
   mask = to_float(thres >= norm)[None, :]
   mask = tf.tile(mask, [size, 1])
 
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index fe20c2956..c7e7c7a2f 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -848,7 +848,7 @@ def testRecompute(self):
 
     def layer(x, name=None):
       with tf.variable_scope(name, default_name="layer"):
-        x = tf.contrib.layers.layer_norm(x)
+        x = common_layers.layer_norm(x)
         x = tf.layers.conv1d(
             x,
             10,
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index d564db7fe..eb05ebb3c 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -22,6 +22,7 @@
 import numpy as np
 
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import contrib
 import tensorflow as tf
 
 from tensorflow.python.ops import summary_op_util  # pylint: disable=g-direct-tensorflow-import
@@ -34,11 +35,6 @@
   distribute_summary_op_util = summary_op_util
 
 tfl = common_layers.layers()
-tfcl = None
-try:
-  tfcl = tf.contrib.layers
-except AttributeError:
-  pass
 
 
 def swap_time_and_batch_axes(inputs):
@@ -119,9 +115,8 @@ def conv_lstm_2d(inputs, state, output_channels,
   else:
     input_shape = spatial_dims + [input_channels]
 
-  cell = tf.contrib.rnn.ConvLSTMCell(
-      2, input_shape, output_channels,
-      [kernel_size, kernel_size], name=name)
+  cell = contrib.rnn().ConvLSTMCell(
+      2, input_shape, output_channels, [kernel_size, kernel_size], name=name)
   if state is None:
     state = cell.zero_state(batch_size, tf.float32)
   outputs, new_state = cell(inputs, state)
@@ -556,14 +551,14 @@ def conv_latent_tower(images, time_axis, latent_channels=1, min_logvar=-5,
     x = common_layers.make_even_size(x)
     x = tfl.conv2d(x, conv_size[0], [3, 3], strides=(2, 2),
                    padding="SAME", activation=tf.nn.relu, name="latent_conv1")
-    x = tfcl.layer_norm(x)
+    x = contrib.layers().layer_norm(x)
     if not small_mode:
       x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="latent_conv2")
-      x = tfcl.layer_norm(x)
+      x = contrib.layers().layer_norm(x)
     x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(1, 1),
                    padding="SAME", activation=tf.nn.relu, name="latent_conv3")
-    x = tfcl.layer_norm(x)
+    x = contrib.layers().layer_norm(x)
 
     nc = latent_channels
     mean = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
diff --git a/tensor2tensor/layers/vqa_layers.py b/tensor2tensor/layers/vqa_layers.py
index f81c66709..341fbf463 100644
--- a/tensor2tensor/layers/vqa_layers.py
+++ b/tensor2tensor/layers/vqa_layers.py
@@ -21,6 +21,7 @@
 
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import contrib
 
 import tensorflow as tf
 
@@ -68,7 +69,7 @@ def image_embedding(images,
   }
 
   if trainable:
-    weights_regularizer = tf.contrib.layers.l2_regularizer(weight_decay)
+    weights_regularizer = contrib.layers().l2_regularizer(weight_decay)
   else:
     weights_regularizer = None
 
@@ -94,7 +95,7 @@ def image_embedding(images,
 
   if add_summaries:
     for v in end_points.values():
-      tf.contrib.layers.summaries.summarize_activation(v)
+      contrib.layers().summaries.summarize_activation(v)
 
   return net
 
diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index f0ef2aae6..246a5be54 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -23,6 +23,7 @@
 from tensor2tensor.layers import area_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -97,9 +98,9 @@ def lstm_attention_decoder(inputs, hparams, train, name, initial_state,
   layers = [_dropout_lstm_cell(hparams, train)
             for _ in range(hparams.num_hidden_layers)]
   if hparams.attention_mechanism == "luong":
-    attention_mechanism_class = tf.contrib.seq2seq.LuongAttention
+    attention_mechanism_class = contrib.seq2seq().LuongAttention
   elif hparams.attention_mechanism == "bahdanau":
-    attention_mechanism_class = tf.contrib.seq2seq.BahdanauAttention
+    attention_mechanism_class = contrib.seq2seq().BahdanauAttention
   else:
     raise ValueError("Unknown hparams.attention_mechanism = %s, must be "
                      "luong or bahdanau." % hparams.attention_mechanism)
@@ -143,10 +144,10 @@ def _area_prob_fn(score):
   else:
     attention_mechanism = attention_mechanism_class(hparams.hidden_size,
                                                     encoder_outputs)
-  cell = tf.contrib.seq2seq.AttentionWrapper(
+  cell = contrib.seq2seq().AttentionWrapper(
       tf.nn.rnn_cell.MultiRNNCell(layers),
-      [attention_mechanism]*hparams.num_heads,
-      attention_layer_size=[hparams.attention_layer_size]*hparams.num_heads,
+      [attention_mechanism] * hparams.num_heads,
+      attention_layer_size=[hparams.attention_layer_size] * hparams.num_heads,
       output_attention=(hparams.output_attention == 1))
 
   batch_size = common_layers.shape_list(inputs)[0]
diff --git a/tensor2tensor/models/neural_architecture_search/nas_model.py b/tensor2tensor/models/neural_architecture_search/nas_model.py
index 5c85390e1..dd2073004 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model.py
@@ -35,6 +35,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
 from tensor2tensor.models.neural_architecture_search import nas_layers as layers
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
@@ -444,12 +445,12 @@ def _tpu_estimator_spec_eval(self, features, logits, labels, loss,
       # For TPU, logits dict will be passed as keyword arguments to
       # eval_metrics_fn. Here we add the labels to those arguments.
       logits.update({"labels": labels})
-      return tf.contrib.tpu.TPUEstimatorSpec(
+      return contrib.tpu().TPUEstimatorSpec(
           tf.estimator.ModeKeys.EVAL,
           eval_metrics=(eval_metrics_fn, logits),
           loss=loss)
     else:
-      return tf.contrib.tpu.TPUEstimatorSpec(
+      return contrib.tpu().TPUEstimatorSpec(
           tf.estimator.ModeKeys.EVAL,
           eval_metrics=(eval_metrics_fn, [logits, labels]),
           loss=loss)
diff --git a/tensor2tensor/models/research/attention_lm.py b/tensor2tensor/models/research/attention_lm.py
index 93d6f5e1a..b44afdaea 100644
--- a/tensor2tensor/models/research/attention_lm.py
+++ b/tensor2tensor/models/research/attention_lm.py
@@ -31,15 +31,17 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
 import tensorflow as tf
 
+framework = contrib.framework(msg="warn")
 
-@tf.contrib.framework.deprecated(
-    "2018-09-15",
-    "Use Transformer, which supports decoder-only mode when "
+
+@framework.deprecated(
+    "2018-09-15", "Use Transformer, which supports decoder-only mode when "
     "Transformer.has_input=False.")
 @registry.register_model
 class AttentionLM(t2t_model.T2TModel):
diff --git a/tensor2tensor/models/research/gene_expression.py b/tensor2tensor/models/research/gene_expression.py
index 955980298..dbf843ebf 100644
--- a/tensor2tensor/models/research/gene_expression.py
+++ b/tensor2tensor/models/research/gene_expression.py
@@ -21,6 +21,7 @@
 
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -118,7 +119,7 @@ def fc_layer(x, num_out, dropout_rate, name="fc"):
   with tf.variable_scope(name):
     out = x
     out = tf.layers.dense(out, num_out)
-    out = tf.contrib.layers.layer_norm(out)
+    out = contrib.layers().layer_norm(out)
     out = tf.nn.relu(out)
     out = tf.layers.dropout(out, dropout_rate)
     return out
diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 9fa39d742..f795e4ef2 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -24,12 +24,13 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models.research import glow_init_hook
 from tensor2tensor.models.research import glow_ops
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 import tensorflow as tf
 
-arg_scope = tf.contrib.framework.arg_scope
-add_arg_scope = tf.contrib.framework.add_arg_scope
+arg_scope = contrib.framework().arg_scope
+add_arg_scope = contrib.framework().add_arg_scope
 
 GLOW_DECODE_HPARAMS = ("identity_output=True,log_results=False,"
                        "decode_in_memory=True,display_decoded_images=True")
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 84c44f6a3..ee21fa5e0 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -24,11 +24,12 @@
 import scipy
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
+from tensor2tensor.utils import contrib
 import tensorflow as tf
 import tensorflow_probability as tfp
 
-arg_scope = tf.contrib.framework.arg_scope
-add_arg_scope = tf.contrib.framework.add_arg_scope
+arg_scope = contrib.framework().arg_scope
+add_arg_scope = contrib.framework().add_arg_scope
 
 
 def linear_interpolate(tensor1, tensor2, coeffs):
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index c77215753..6fb2d4e9c 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -28,15 +28,12 @@
 from six.moves import zip
 from tensor2tensor.models.research import glow
 from tensor2tensor.models.research import glow_ops
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import hparam
 import tensorflow as tf
 
-arg_scope = tf.contrib.framework.arg_scope
-add_arg_scope = tf.contrib.framework.add_arg_scope
-
-
-arg_scope = tf.contrib.framework.arg_scope
-add_arg_scope = tf.contrib.framework.add_arg_scope
+arg_scope = contrib.framework().arg_scope
+add_arg_scope = contrib.framework().add_arg_scope
 
 
 class GlowOpsTest(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/models/research/neural_stack.py b/tensor2tensor/models/research/neural_stack.py
index 8909e9748..9b6942b68 100644
--- a/tensor2tensor/models/research/neural_stack.py
+++ b/tensor2tensor/models/research/neural_stack.py
@@ -31,6 +31,7 @@
 
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -213,7 +214,7 @@ def build_controller(self):
     """Create the RNN and output projections for controlling the stack.
     """
     with tf.name_scope("controller"):
-      self.rnn = tf.contrib.rnn.BasicRNNCell(self._num_units)
+      self.rnn = contrib.rnn().BasicRNNCell(self._num_units)
       self._input_proj = self.add_variable(
           "input_projection_weights",
           shape=[self._embedding_size * (self._num_read_heads + 1),
@@ -280,9 +281,10 @@ def call_controller(self, input_value, read_values, prev_state, batch_size):
       # Concatenate the current input value with the read values from the
       # previous timestep before feeding them into the controller.
       controller_inputs = tf.concat([
-          tf.contrib.layers.flatten(input_value),
-          tf.contrib.layers.flatten(read_values),
-      ], axis=1)
+          contrib.layers().flatten(input_value),
+          contrib.layers().flatten(read_values),
+      ],
+                                    axis=1)
 
       rnn_input = tf.tanh(tf.nn.bias_add(tf.matmul(
           controller_inputs, self._input_proj), self._input_bias))
@@ -523,7 +525,7 @@ def _rnn(self, inputs, name, initial_state=None, sequence_length=None):
               for layer_size in self._hparams.controller_layer_sizes]
     with tf.variable_scope(name):
       return tf.nn.dynamic_rnn(
-          tf.contrib.rnn.MultiRNNCell(layers),
+          contrib.rnn().MultiRNNCell(layers),
           inputs,
           initial_state=initial_state,
           sequence_length=sequence_length,
diff --git a/tensor2tensor/models/research/neural_stack_test.py b/tensor2tensor/models/research/neural_stack_test.py
index a43ba6e23..132315f5f 100644
--- a/tensor2tensor/models/research/neural_stack_test.py
+++ b/tensor2tensor/models/research/neural_stack_test.py
@@ -23,6 +23,7 @@
 
 from tensor2tensor.layers import modalities
 from tensor2tensor.models.research import neural_stack
+from tensor2tensor.utils import contrib
 
 import tensorflow as tf
 
@@ -399,7 +400,7 @@ def test_model_shapes(self):
     vocab_size = 128
 
     hparams = neural_stack.neural_stack()
-    problem_hparams = tf.contrib.training.HParams()
+    problem_hparams = contrib.training().HParams()
 
     problem_hparams.add_hparam("modality", {
         "inputs": modalities.ModalityType.SYMBOL,
diff --git a/tensor2tensor/models/research/transformer_revnet.py b/tensor2tensor/models/research/transformer_revnet.py
index 21a17788b..7c3ba2c14 100644
--- a/tensor2tensor/models/research/transformer_revnet.py
+++ b/tensor2tensor/models/research/transformer_revnet.py
@@ -21,6 +21,7 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -120,7 +121,7 @@ def g(x):
   x1, x2 = tf.split(encoder_input, 2, axis=-1)
 
   with tf.variable_scope(name):
-    y1, y2 = tf.contrib.layers.rev_block(
+    y1, y2 = contrib.layers().rev_block(
         x1,
         x2,
         f,
@@ -198,7 +199,7 @@ def g(x):
   x1, x2 = tf.split(decoder_input, 2, axis=-1)
 
   with tf.variable_scope(name):
-    y1, y2 = tf.contrib.layers.rev_block(
+    y1, y2 = contrib.layers().rev_block(
         x1,
         x2,
         f,
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index eca25e1a3..a50fe1313 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -32,6 +32,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.models import transformer
 from tensor2tensor.utils import beam_search
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import expert_utils
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
@@ -377,7 +378,7 @@ def ae_transformer_internal(inputs,
                                 minval=0,
                                 maxval=1 + hparams.word_shuffle)
       targets_idx += noise
-      permutation = tf.contrib.framework.argsort(targets_idx)
+      permutation = contrib.framework().argsort(targets_idx)
       targets_permuted = tf.gather(targets, indices=permutation, axis=1)
       targets = targets_permuted
     targets, _ = common_layers.pad_to_same_length(
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index a5a65504e..878828921 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -32,6 +32,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
 from tensor2tensor.models.research import universal_transformer_util
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -219,7 +220,7 @@ def body(self, features):
           hparams.act_loss_weight *
           tf.reduce_mean(dec_ponder_times + dec_remainders))
       act_loss = enc_act_loss + dec_act_loss
-      tf.contrib.summary.scalar("act_loss", act_loss)
+      contrib.summary().scalar("act_loss", act_loss)
       return decoder_output, {"act_loss": act_loss}
 
     return decoder_output
@@ -347,7 +348,7 @@ def body(self, features):
       ponder_times, remainders = enc_extra_output
       act_loss = hparams.act_loss_weight * tf.reduce_mean(ponder_times +
                                                           remainders)
-      tf.contrib.summary.scalar("act_loss", act_loss)
+      contrib.summary().scalar("act_loss", act_loss)
 
       return encoder_output, {"act_loss": act_loss}
     return encoder_output
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index b7c6ae25c..5e1c4cc69 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -54,6 +54,7 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.models import transformer
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import expert_utils
 
 import tensorflow as tf
@@ -671,11 +672,11 @@ def universal_transformer_highway(layer_inputs,
 
   new_state = state * carry_gate + transformed_state * transform_gate
 
-  tf.contrib.summary.scalar("highway_transform_gate_layer",
-                            tf.reduce_mean(transform_gate))
+  contrib.summary().scalar("highway_transform_gate_layer",
+                           tf.reduce_mean(transform_gate))
 
-  tf.contrib.summary.scalar("highway_carry_gate_layer",
-                            tf.reduce_mean(carry_gate))
+  contrib.summary().scalar("highway_carry_gate_layer",
+                           tf.reduce_mean(carry_gate))
 
   return new_state, inputs, memory
 
@@ -761,10 +762,10 @@ def universal_transformer_skip(layer_inputs,
         pad_remover=pad_remover,
         preprocess=True)
 
-  tf.contrib.summary.scalar("skip_transform_gate_layer",
-                            tf.reduce_mean(transform_gate))
+  contrib.summary().scalar("skip_transform_gate_layer",
+                           tf.reduce_mean(transform_gate))
 
-  tf.contrib.summary.scalar("skip_carry_gate_layer", tf.reduce_mean(carry_gate))
+  contrib.summary().scalar("skip_carry_gate_layer", tf.reduce_mean(carry_gate))
 
   new_state = inputs * carry_gate + transformed_state * transform_gate
   return new_state, inputs, memory
@@ -879,8 +880,8 @@ def universal_transformer_with_gru_as_transition_function(
         activation=tf.sigmoid,
         pad_remover=pad_remover)
 
-    tf.contrib.summary.scalar("gru_update_gate",
-                              tf.reduce_mean(transition_function_update_gate))
+    contrib.summary().scalar("gru_update_gate",
+                             tf.reduce_mean(transition_function_update_gate))
 
     # gru reset gate: r_t = sigmoid(W_r.x_t + U_r.h_{t-1})
     transition_function_reset_gate = _ffn_layer_multi_inputs(
@@ -891,8 +892,8 @@ def universal_transformer_with_gru_as_transition_function(
         activation=tf.sigmoid,
         pad_remover=pad_remover)
 
-    tf.contrib.summary.scalar("gru_reset_gate",
-                              tf.reduce_mean(transition_function_reset_gate))
+    contrib.summary().scalar("gru_reset_gate",
+                             tf.reduce_mean(transition_function_reset_gate))
     reset_state = transition_function_reset_gate * state
 
     # gru_candidate_activation: h' = tanh(W_{x_t} + U (r_t h_{t-1})
@@ -967,8 +968,8 @@ def universal_transformer_with_lstm_as_transition_function(
         activation=tf.sigmoid,
         pad_remover=pad_remover)
 
-    tf.contrib.summary.scalar("lstm_input_gate",
-                              tf.reduce_mean(transition_function_input_gate))
+    contrib.summary().scalar("lstm_input_gate",
+                             tf.reduce_mean(transition_function_input_gate))
 
     # lstm forget gate: f_t = sigmoid(W_f.x_t + U_f.h_{t-1})
     transition_function_forget_gate = _ffn_layer_multi_inputs(
@@ -982,8 +983,8 @@ def universal_transformer_with_lstm_as_transition_function(
     transition_function_forget_gate = tf.sigmoid(
         transition_function_forget_gate + forget_bias_tensor)
 
-    tf.contrib.summary.scalar("lstm_forget_gate",
-                              tf.reduce_mean(transition_function_forget_gate))
+    contrib.summary().scalar("lstm_forget_gate",
+                             tf.reduce_mean(transition_function_forget_gate))
 
     # lstm output gate: o_t = sigmoid(W_o.x_t + U_o.h_{t-1})
     transition_function_output_gate = _ffn_layer_multi_inputs(
@@ -994,8 +995,8 @@ def universal_transformer_with_lstm_as_transition_function(
         activation=tf.sigmoid,
         pad_remover=pad_remover)
 
-    tf.contrib.summary.scalar("lstm_output_gate",
-                              tf.reduce_mean(transition_function_output_gate))
+    contrib.summary().scalar("lstm_output_gate",
+                             tf.reduce_mean(transition_function_output_gate))
 
     # lstm input modulation
     transition_function_input_modulation = _ffn_layer_multi_inputs(
@@ -1189,7 +1190,7 @@ def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
   ponder_times = n_updates
   remainders = remainder
 
-  tf.contrib.summary.scalar("ponder_times", tf.reduce_mean(ponder_times))
+  contrib.summary().scalar("ponder_times", tf.reduce_mean(ponder_times))
 
   return new_state, (ponder_times, remainders)
 
diff --git a/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index 801ef577f..c09875c7f 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -28,6 +28,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow as tf
+from tensorflow.contrib import rnn as contrib_rnn
 
 # pylint: disable=unused-import
 from tensorflow.contrib.layers.python.layers import utils
@@ -238,7 +239,7 @@ def _get_rnn_cell(hparams):
   if hparams.rnn_type == "lstm":
     rnn_cell = tf.nn.rnn_cell.BasicLSTMCell
   elif hparams.rnn_type == "lstm_layernorm":
-    rnn_cell = tf.contrib.rnn.LayerNormBasicLSTMCell
+    rnn_cell = contrib_rnn.LayerNormBasicLSTMCell
   return tf.nn.rnn_cell.DropoutWrapper(
       rnn_cell(hparams.hidden_size),
       output_keep_prob=1.0-hparams.dropout)
diff --git a/tensor2tensor/models/revnet.py b/tensor2tensor/models/revnet.py
index 8a4f24542..193cbb3b2 100644
--- a/tensor2tensor/models/revnet.py
+++ b/tensor2tensor/models/revnet.py
@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 """Creates a RevNet with the bottleneck residual function.
 
 Implements the following equations described in the RevNet paper:
@@ -37,6 +36,7 @@
 
 import functools
 from tensor2tensor.layers import common_hparams
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -251,10 +251,8 @@ def unit(x1, x2, block_num, depth, num_layers, dim='2d',
 
     # Full block using memory-efficient rev_block implementation.
     with tf.variable_scope('full_block'):
-      x1, x2 = tf.contrib.layers.rev_block(x1, x2,
-                                           residual,
-                                           residual,
-                                           num_layers=num_layers)
+      x1, x2 = contrib.layers().rev_block(
+          x1, x2, residual, residual, num_layers=num_layers)
       return x1, x2
 
 
diff --git a/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index 58afe4890..58b990a70 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -32,12 +32,13 @@
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.video import sv2p
 from tensor2tensor.models.video import sv2p_params
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
 tfl = tf.layers
-tfcl = tf.contrib.layers
+tfcl = contrib.layers()
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index a1e4941b7..735002fff 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -34,6 +34,7 @@
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.video import epva_params  # pylint: disable=unused-import
 from tensor2tensor.models.video import sv2p
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
@@ -42,7 +43,7 @@
 from tensorflow.contrib.slim.python.slim.nets import vgg
 
 tfl = tf.layers
-tfcl = tf.contrib.layers
+tfcl = contrib.layers()
 
 IMG_WIDTH = 64
 IMG_HEIGHT = 64
@@ -71,12 +72,12 @@ def van_image_enc_2d(x, first_depth, reuse=False, hparams=None):
 
     enc = tf.layers.conv2d(
         x, first_depth, 3, padding='same', activation=tf.nn.relu, strides=1)
-    enc = tf.contrib.layers.layer_norm(enc)
+    enc = contrib.layers().layer_norm(enc)
     enc = tf.layers.conv2d(
         enc, first_depth, 3, padding='same', activation=tf.nn.relu, strides=1)
     enc = tf.nn.max_pool(enc, [1, 2, 2, 1], [1, 2, 2, 1], 'SAME')
     enc = tf.nn.dropout(enc, hparams.van_keep_prob)
-    enc = tf.contrib.layers.layer_norm(enc)
+    enc = contrib.layers().layer_norm(enc)
     enc_history.append(enc)
 
     enc = tf.layers.conv2d(
@@ -95,7 +96,7 @@ def van_image_enc_2d(x, first_depth, reuse=False, hparams=None):
         strides=1)
     enc = tf.nn.max_pool(enc, [1, 2, 2, 1], [1, 2, 2, 1], 'SAME')
     enc = tf.nn.dropout(enc, hparams.van_keep_prob)
-    enc = tf.contrib.layers.layer_norm(enc)
+    enc = contrib.layers().layer_norm(enc)
     enc_history.append(enc)
 
     enc = tf.layers.conv2d(
@@ -144,13 +145,13 @@ def van_enc_2d(x, first_depth, reuse=False):
     # a, b = 4,4
     enc = tf.nn.relu(x)
     enc = tf.layers.dense(enc, first_depth * a * b, tf.nn.relu)
-    enc = tf.contrib.layers.layer_norm(enc)
+    enc = contrib.layers().layer_norm(enc)
 
     enc = tf.reshape(enc, [-1, a, b, first_depth])
 
     enc = tf.layers.conv2d_transpose(
         enc, first_depth, 3, padding='same', activation=tf.nn.relu, strides=1)
-    enc = tf.contrib.layers.layer_norm(enc)
+    enc = contrib.layers().layer_norm(enc)
     enc = tf.layers.conv2d_transpose(
         enc,
         first_depth * 2,
@@ -167,7 +168,7 @@ def van_enc_2d(x, first_depth, reuse=False):
         padding='same',
         activation=tf.nn.relu,
         strides=1)
-    enc = tf.contrib.layers.layer_norm(enc)
+    enc = contrib.layers().layer_norm(enc)
     enc = tf.layers.conv2d_transpose(
         enc,
         first_depth * 4,
@@ -199,7 +200,7 @@ def van_dec_2d(x, skip_connections, output_shape, first_depth, hparams=None):
     dec = tf.layers.conv2d_transpose(
         x, first_depth * 4, 3, padding='same', activation=tf.nn.relu, strides=2)
     dec = tf.nn.dropout(dec, hparams.van_keep_prob)
-    dec = tf.contrib.layers.layer_norm(dec)
+    dec = contrib.layers().layer_norm(dec)
     dec = tf.layers.conv2d_transpose(
         dec,
         first_depth * 4,
@@ -216,7 +217,7 @@ def van_dec_2d(x, skip_connections, output_shape, first_depth, hparams=None):
         activation=tf.nn.relu,
         strides=1)
     dec = tf.nn.dropout(dec, hparams.van_keep_prob)
-    dec = tf.contrib.layers.layer_norm(dec)
+    dec = contrib.layers().layer_norm(dec)
 
     dec = tf.layers.conv2d_transpose(
         dec,
@@ -229,7 +230,7 @@ def van_dec_2d(x, skip_connections, output_shape, first_depth, hparams=None):
     dec = tf.layers.conv2d_transpose(
         dec, first_depth, 3, padding='same', activation=tf.nn.relu, strides=1)
     dec = tf.nn.dropout(dec, hparams.van_keep_prob)
-    dec = tf.contrib.layers.layer_norm(dec)
+    dec = contrib.layers().layer_norm(dec)
 
     dec = tf.layers.conv2d_transpose(
         dec,
@@ -281,7 +282,7 @@ def analogy_computation_2d(f_first_enc,
         padding='same',
         activation=tf.nn.relu,
         strides=1)
-    analogy = tf.contrib.layers.layer_norm(analogy)
+    analogy = contrib.layers().layer_norm(analogy)
     analogy = tf.layers.conv2d(
         analogy,
         first_depth * 4,
@@ -458,7 +459,7 @@ def predictor(enc_flat,
         use_peepholes=True,
         initializer=tf.truncated_normal_initializer(stddev=lstm_init_stddev),
         num_proj=initial_size)
-    part_pred = tf.contrib.layers.layer_norm(part_pred)
+    part_pred = contrib.layers().layer_norm(part_pred)
     pred = part_pred
 
     for pred_layer_num in range(1, pred_depth, 2):
@@ -478,7 +479,7 @@ def predictor(enc_flat,
           use_peepholes=True,
           initializer=tf.truncated_normal_initializer(stddev=lstm_init_stddev),
           num_proj=initial_size)
-      part_pred = tf.contrib.layers.layer_norm(part_pred)
+      part_pred = contrib.layers().layer_norm(part_pred)
       pred += part_pred
 
     pred = tf.layers.dense(
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index e82ed2876..e37044420 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -25,12 +25,12 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.models.research import glow
 from tensor2tensor.models.research import glow_ops
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 import tensorflow as tf
 import tensorflow_probability as tfp
 
-
-arg_scope = tf.contrib.framework.arg_scope
+arg_scope = contrib.framework().arg_scope
 
 
 @registry.register_hparams
diff --git a/tensor2tensor/models/video/nfg_interpolate.py b/tensor2tensor/models/video/nfg_interpolate.py
index db665853c..535a6c61c 100644
--- a/tensor2tensor/models/video/nfg_interpolate.py
+++ b/tensor2tensor/models/video/nfg_interpolate.py
@@ -26,6 +26,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import glow_ops
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import trainer_lib
 import tensorflow as tf
@@ -47,7 +48,7 @@
 FLAGS = flags.FLAGS
 
 
-arg_scope = tf.contrib.framework.arg_scope
+arg_scope = contrib.framework().arg_scope
 
 
 def decode_hparams(overrides=""):
diff --git a/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index 6e44e50e3..186378bb7 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -28,6 +28,7 @@
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.video import savp_params  # pylint: disable=unused-import
 from tensor2tensor.models.video import sv2p
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import update_ops_hook
 
@@ -82,7 +83,7 @@ def encoder(self, inputs, n_layers=3):
           padded = tf.pad(inputs, padding)
         convolved = tf.layers.conv2d(padded, filters=n_filters, kernel_size=4,
                                      strides=2, padding="VALID")
-        normalized = tf.contrib.layers.instance_norm(convolved)
+        normalized = contrib.layers().instance_norm(convolved)
         rectified = tf.nn.leaky_relu(normalized, alpha=0.2)
 
     # Mean pooling across all spatial dimensions.
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index ea8f26e15..2aacc8da7 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -31,12 +31,13 @@
 
 from tensor2tensor.models.video import base
 from tensor2tensor.models.video import base_vae
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
 tfl = tf.layers
-tfcl = tf.contrib.layers
+tfcl = contrib.layers()
 
 
 @registry.register_model
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index ba535ca88..9facc7eeb 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -28,7 +28,9 @@
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.utils import cloud_mlengine as cloud
+from tensor2tensor.utils import contrib
 import tensorflow as tf
+
 from tensorflow_serving.apis import predict_pb2
 from tensorflow_serving.apis import prediction_service_pb2_grpc
 
diff --git a/tensor2tensor/utils/contrib.py b/tensor2tensor/utils/contrib.py
new file mode 100644
index 000000000..e3b34e758
--- /dev/null
+++ b/tensor2tensor/utils/contrib.py
@@ -0,0 +1,156 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Wrappers around tf.contrib to dynamically import contrib packages.
+
+This makes sure that libraries depending on T2T and TF2, do not crash at import.
+"""
+
+from __future__ import absolute_import
+from __future__ import division  # Not necessary in a Python 3-only module
+from __future__ import print_function  # Not necessary in a Python 3-only module
+
+from absl import logging
+from tensorflow.python import tf2  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+is_tf2 = tf2.enabled()
+
+
+def err_if_tf2(msg='err'):
+  if is_tf2:
+    msg = 'contrib is unavailable in tf2.'
+    if msg == 'err':
+      raise ImportError(msg)
+    else:
+      logging.info(msg)
+
+
+def slim():
+  err_if_tf2()
+  from tensorflow.contrib import slim as contrib_slim  # pylint: disable=g-import-not-at-top
+  return contrib_slim
+
+
+def util():
+  err_if_tf2()
+  from tensorflow.contrib import util as contrib_util  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_util
+
+
+def tfe():
+  err_if_tf2(msg='warn')
+  from tensorflow.contrib.eager.python import tfe as contrib_eager  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_eager
+
+
+def framework(msg='err'):
+  err_if_tf2(msg=msg)
+  from tensorflow.contrib import framework as contrib_framework  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_framework
+
+
+def nn():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import nn as contrib_nn  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_nn
+
+
+def layers():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import layers as contrib_layers  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_layers
+
+
+def rnn():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import rnn as contrib_rnn  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_rnn
+
+
+def seq2seq():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import seq2seq as contrib_seq2seq  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_seq2seq
+
+
+def tpu():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import tpu as contrib_tpu  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_tpu
+
+
+def training():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import training as contrib_training  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_training
+
+
+def summary():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import summary as contrib_summary  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_summary
+
+
+def metrics():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import metrics as contrib_metrics  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_metrics
+
+
+def opt():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import opt as contrib_opt  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_opt
+
+
+def mixed_precision():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import mixed_precision as contrib_mixed_precision  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_mixed_precision
+
+
+def cluster_resolver():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import cluster_resolver as contrib_cluster_resolver  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_cluster_resolver
+
+
+def distribute():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import distribute as contrib_distribute  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_distribute
+
+
+def learn():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import learn as contrib_learn  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_learn
+
+
+def tf_prof():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import tfprof as contrib_tfprof  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_tfprof
+
+
+def eager():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib.eager.python import tfe as contrib_eager  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_eager
+
+
+def image():
+  err_if_tf2(msg='err')
+  from tensorflow.contrib import image as contrib_image  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+  return contrib_image
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index c65878163..a979d37d7 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -25,6 +25,7 @@
 import six
 from six.moves import range  # pylint: disable=redefined-builtin
 
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
@@ -379,7 +380,7 @@ def define_shapes(example):
     dataset = dataset.repeat()
 
   if is_training and skip_random_fraction_when_training:
-    data_files = tf.contrib.slim.parallel_reader.get_data_files(filepattern)
+    data_files = contrib.slim().parallel_reader.get_data_files(filepattern)
     #  In continuous_train_and_eval when switching between train and
     #  eval, this input_fn method gets called multiple times and it
     #  would give you the exact same samples from the last call
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index d3ac9b2fe..405cca8a6 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -33,6 +33,7 @@
 from tensor2tensor.data_generators import problem as problem_lib
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import text_problems
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import hparam
 from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import registry
@@ -601,7 +602,7 @@ def _decode_filename(base_filename, problem_name, decode_hp):
 def make_input_fn_from_generator(gen):
   """Use py_func to yield elements from the given generator."""
   first_ex = six.next(gen)
-  flattened = tf.contrib.framework.nest.flatten(first_ex)
+  flattened = contrib.framework().nest.flatten(first_ex)
   types = [t.dtype for t in flattened]
   shapes = [[None] * len(t.shape) for t in flattened]
   first_ex_list = [first_ex]
@@ -611,12 +612,12 @@ def py_func():
       example = first_ex_list.pop()
     else:
       example = six.next(gen)
-    return tf.contrib.framework.nest.flatten(example)
+    return contrib.framework().nest.flatten(example)
 
   def input_fn():
     flat_example = tf.py_func(py_func, [], types)
     _ = [t.set_shape(shape) for t, shape in zip(flat_example, shapes)]
-    example = tf.contrib.framework.nest.pack_sequence_as(first_ex, flat_example)
+    example = contrib.framework().nest.pack_sequence_as(first_ex, flat_example)
     return example
 
   return input_fn
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 59a54133b..edcf2cdc3 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -24,11 +24,11 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import bleu_hook
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import rouge
 from tensor2tensor.utils import sari_hook
 
 import tensorflow as tf
-
 from tensorflow.python.util import tf_inspect as inspect
 
 
@@ -880,8 +880,8 @@ def pearson_correlation_coefficient(predictions, labels, weights_fn=None):
     The pearson correlation coefficient.
   """
   del weights_fn
-  _, pearson = tf.contrib.metrics.streaming_pearson_correlation(predictions,
-                                                                labels)
+  _, pearson = contrib.metrics().streaming_pearson_correlation(
+      predictions, labels)
   return pearson, tf.constant(1.0)
 
 # Metrics are functions that take predictions and labels and return
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index cdfd8e1b4..b0d848ac9 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -21,6 +21,7 @@
 import numpy as np
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import adafactor as adafactor_lib
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import multistep_optimizer
@@ -68,7 +69,7 @@ def optimize(loss,
       diet_vars, "Diet Variables", verbose=hparams.summarize_vars)
   opt = ConditionalOptimizer(hparams.optimizer, learning_rate, hparams, use_tpu)
   if use_tpu:
-    opt = tf.contrib.tpu.CrossShardOptimizer(opt)
+    opt = contrib.tpu().CrossShardOptimizer(opt)
   if getattr(hparams, "gpu_automatic_mixed_precision", False):
     if use_tpu:
       raise RuntimeError("GPU auto mixed precision cannot be used with TPU")
@@ -95,7 +96,7 @@ def optimize(loss,
     tf.logging.info("Adding noise to gradients, noise scale: %0.5f",
                     hparams.grad_noise_scale)
 
-  train_op = tf.contrib.layers.optimize_loss(
+  train_op = contrib.layers().optimize_loss(
       name="training",
       loss=loss,
       global_step=tf.train.get_or_create_global_step(),
@@ -113,7 +114,7 @@ def optimize(loss,
 def adam(learning_rate, hparams):
   # We change the default epsilon for Adam.
   # Using LazyAdam as it's much faster for large vocabulary embeddings.
-  return tf.contrib.opt.LazyAdamOptimizer(
+  return contrib.opt().LazyAdamOptimizer(
       learning_rate,
       beta1=hparams.optimizer_adam_beta1,
       beta2=hparams.optimizer_adam_beta2,
@@ -156,7 +157,7 @@ def true_adam(learning_rate, hparams):
 
 @registry.register_optimizer
 def adam_w(learning_rate, hparams):
-  return tf.contrib.opt.AdamWOptimizer(
+  return contrib.opt().AdamWOptimizer(
       weight_decay=hparams.weight_decay,
       learning_rate=learning_rate,
       beta1=hparams.optimizer_adam_beta1,
@@ -179,7 +180,7 @@ def _register_base_optimizer(name, opt):
       lambda learning_rate, hparams: opt(learning_rate))
 
 
-for _name, _opt in tf.contrib.layers.OPTIMIZER_CLS_NAMES.items():
+for _name, _opt in contrib.layers().OPTIMIZER_CLS_NAMES.items():
   _register_base_optimizer(_name, _opt)
 
 
@@ -216,13 +217,13 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
             ("Using Exponential Update Loss Scaler with",
              "init loss scale of {}".format(
                  hparams.mixed_precision_optimizer_init_loss_scale)))
-        manager = tf.contrib.mixed_precision.ExponentialUpdateLossScaleManager(
+        manager = contrib.mixed_precision().ExponentialUpdateLossScaleManager(
             init_loss_scale=hparams.mixed_precision_optimizer_init_loss_scale,
             incr_every_n_steps=2000,
             decr_every_n_nan_or_inf=2,
             incr_ratio=2,
             decr_ratio=0.5)
-        self._opt = tf.contrib.mixed_precision.LossScaleOptimizer(
+        self._opt = contrib.mixed_precision().LossScaleOptimizer(
             self._opt, manager)
 
     self._zero_grads = hparams.optimizer_zero_grads
diff --git a/tensor2tensor/utils/restore_hook.py b/tensor2tensor/utils/restore_hook.py
index 6d649c209..5960b26da 100644
--- a/tensor2tensor/utils/restore_hook.py
+++ b/tensor2tensor/utils/restore_hook.py
@@ -21,6 +21,7 @@
 
 import six
 
+from tensor2tensor.utils import contrib
 import tensorflow as tf
 
 
@@ -44,7 +45,7 @@ def begin(self):
     match the old_model_scope and remove the suffix :0.
 
     """
-    variables_to_restore = tf.contrib.framework.get_variables_to_restore(
+    variables_to_restore = contrib.framework().get_variables_to_restore(
         include=self._include, exclude=self._exclude)
     # remove new_model_scope from variable name prefix
     assignment_map = {variable.name[len(self._new_model_scope):]: variable
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 599381888..40f4da8fa 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -34,6 +34,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.layers.common_attention import mixed_precision_is_enabled
 from tensor2tensor.utils import beam_search
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import expert_utils as eu
 from tensor2tensor.utils import hparams_lib
@@ -1578,7 +1579,7 @@ def scaffold_fn():
 
       remove_summaries()
 
-      return tf.contrib.tpu.TPUEstimatorSpec(
+      return contrib.tpu().TPUEstimatorSpec(
           tf.estimator.ModeKeys.TRAIN,
           loss=loss,
           train_op=train_op,
@@ -1635,7 +1636,7 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
       )
 
       eval_metrics_fn_flat_args = _flatten_dict(eval_metrics_fn_args)
-      return tf.contrib.tpu.TPUEstimatorSpec(
+      return contrib.tpu().TPUEstimatorSpec(
           tf.estimator.ModeKeys.EVAL,
           eval_metrics=(eval_metrics_fn, eval_metrics_fn_flat_args),
           host_call=host_call,
@@ -1760,7 +1761,7 @@ def estimator_spec_predict(self, features, use_tpu=False):
 
       remove_summaries()
 
-      return tf.contrib.tpu.TPUEstimatorSpec(
+      return contrib.tpu().TPUEstimatorSpec(
           tf.estimator.ModeKeys.PREDICT,
           predictions=predictions,
           host_call=host_call,
@@ -2171,19 +2172,19 @@ def host_call_fn(**kwargs):
       List of summary ops to run on the CPU host.
     """
     gs = tf.to_int64(kwargs.pop("global_step")[0])
-    with tf.contrib.summary.create_file_writer(model_dir).as_default():
-      with tf.contrib.summary.always_record_summaries():
+    with contrib.summary().create_file_writer(model_dir).as_default():
+      with contrib.summary().always_record_summaries():
         # We need to use tf.contrib.summary in order to feed the `step`.
         for name, value in sorted(six.iteritems(kwargs)):
           if name.startswith("ScalarSummary"):
             name = name[len("ScalarSummary"):]
-            tf.contrib.summary.scalar(
+            contrib.summary().scalar(
                 name, tf.reduce_mean(tf.to_float(value)), step=gs)
           elif name.startswith("ImageSummary"):
             name = name[len("ImageSummary"):]
-            tf.contrib.summary.image(name, value, step=gs)
+            contrib.summary().image(name, value, step=gs)
 
-        return tf.contrib.summary.all_summary_ops()
+        return contrib.summary().all_summary_ops()
 
   return (host_call_fn, summary_kwargs)
 
@@ -2332,9 +2333,9 @@ def initialize_from_ckpt(ckpt_dir, hparams):
     return
 
   tf.logging.info("Checkpoint dir: %s", ckpt_dir)
-  reader = tf.contrib.framework.load_checkpoint(ckpt_dir)
+  reader = contrib.framework().load_checkpoint(ckpt_dir)
   variable_map = {}
-  for var in tf.contrib.framework.get_trainable_variables():
+  for var in contrib.framework().get_trainable_variables():
     var_name = var.name.split(":")[0]
     if reader.has_tensor(var_name):
       tf.logging.info("Loading variable from checkpoint: %s", var_name)
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 36ffe363b..b71ddac8d 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -26,6 +26,7 @@
 import random
 import numpy as np
 
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import devices
 from tensor2tensor.utils import hparams_lib
@@ -59,7 +60,7 @@ def next_checkpoint(model_dir, timeout_mins=240):
   if timeout_mins != -1:
     timeout_secs = timeout_mins * 60
   while True:
-    last_ckpt = tf.contrib.training.wait_for_new_checkpoint(
+    last_ckpt = contrib.training().wait_for_new_checkpoint(
         model_dir, last_ckpt, seconds_to_sleep=60, timeout=timeout_secs)
 
     if last_ckpt is None:
@@ -76,7 +77,7 @@ def next_undecoded_checkpoint(model_dir, timeout_mins=240):
   last_step = 0
   while True:
     # Get the latest checkpoint.
-    last_ckpt = tf.contrib.training.wait_for_new_checkpoint(
+    last_ckpt = contrib.training().wait_for_new_checkpoint(
         model_dir, last_ckpt, seconds_to_sleep=60, timeout=60 * timeout_mins)
     # Get all the checkpoint from the model dir.
     ckpt_path = tf.train.get_checkpoint_state(model_dir)
@@ -202,7 +203,7 @@ def create_run_config(model_name,
   }
   if save_checkpoints_secs:
     del run_config_args["save_checkpoints_steps"]
-  run_config_cls = tf.contrib.learn.RunConfig
+  run_config_cls = contrib.learn().RunConfig
 
   if use_tpu or use_tpu_estimator:
     # If using TPUEstimator, use TPU RunConfig, add TPUConfig, and add
@@ -215,9 +216,8 @@ def create_run_config(model_name,
     }
     if tpu_config_extra_kwargs is not None:
       tpu_config_kwargs.update(tpu_config_extra_kwargs)
-    run_config_cls = tf.contrib.tpu.RunConfig
-    tpu_config = tf.contrib.tpu.TPUConfig(
-        **tpu_config_kwargs)
+    run_config_cls = contrib.tpu().RunConfig
+    tpu_config = contrib.tpu().TPUConfig(**tpu_config_kwargs)
     run_config_args["tpu_config"] = tpu_config
     if not master and "KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS" in os.environ:
       # If running on TPU but no master is set and the KUBE env var is present
@@ -228,7 +228,7 @@ def create_run_config(model_name,
     elif not master and cloud_tpu_name:
       # Update run_config to use cluster instead of master/evaluation_master
       # as we need the cluster spec to use Cloud Pods
-      tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
+      tpu_cluster_resolver = contrib.cluster_resolver().TPUClusterResolver(
           cloud_tpu_name)
       run_config_args["cluster"] = tpu_cluster_resolver
       del run_config_args["master"]
@@ -257,7 +257,7 @@ def create_run_config(model_name,
           "Configuring MirroredStrategy DistributionStrategy to replicate the "
           "model."
       )
-      distribution = tf.contrib.distribute.MirroredStrategy()
+      distribution = contrib.distribute().MirroredStrategy()
       config = config.replace(train_distribute=distribution)
       config.data_parallelism = None
     else:
@@ -367,7 +367,7 @@ def tpu_model_fn(features, labels, mode, params):
       estimator_model_fn = tpu_model_fn
     else:
       raise ValueError("Flag export_saved_model_api_version must be 1 or 2.")
-    estimator = tf.contrib.tpu.TPUEstimator(
+    estimator = contrib.tpu().TPUEstimator(
         model_fn=estimator_model_fn,
         model_dir=run_config.model_dir,
         config=run_config,
@@ -412,7 +412,7 @@ def create_hooks(use_tfdbg=False,
   if use_validation_monitor:
     tf.logging.info("Using ValidationMonitor")
     train_hooks.append(
-        tf.contrib.learn.monitors.ValidationMonitor(
+        contrib.learn().monitors.ValidationMonitor(
             hooks=eval_hooks, **validation_monitor_kwargs))
 
   if use_early_stopping:
@@ -812,9 +812,9 @@ def serving_input_receiver_fn(hparams, decode_hparams, use_tpu):
   if additional_eval_hooks:
     eval_hooks += additional_eval_hooks
 
-  train_hooks = tf.contrib.learn.monitors.replace_monitors_with_hooks(
+  train_hooks = contrib.learn().monitors.replace_monitors_with_hooks(
       train_hooks, estimator)
-  eval_hooks = tf.contrib.learn.monitors.replace_monitors_with_hooks(
+  eval_hooks = contrib.learn().monitors.replace_monitors_with_hooks(
       eval_hooks, estimator)
 
   train_spec = tf.estimator.TrainSpec(

From 670ddafd233daab956d0b541c6bb360a86a45c1b Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 6 Jan 2020 13:36:41 -0800
Subject: [PATCH 2613/2720] Remove extra scope as it's already v1.

PiperOrigin-RevId: 288365574
---
 tensor2tensor/utils/expert_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index f8880e4a0..b7231fe97 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -65,7 +65,7 @@ def decorated(*args, **kwargs):
 
 
 def add_var_scope(scope=None):
-  return add_scope(scope, scope_fn=tf.compat.v1.variable_scope)
+  return add_scope(scope, scope_fn=tf.variable_scope)
 
 
 def add_name_scope(scope=None):

From 1655763b017b945da29443156372756da12b6d24 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 7 Jan 2020 16:03:05 -0800
Subject: [PATCH 2614/2720] Internal change

PiperOrigin-RevId: 288592605
---
 tensor2tensor/data_generators/wikisum/produce_examples.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/data_generators/wikisum/produce_examples.py b/tensor2tensor/data_generators/wikisum/produce_examples.py
index 16b9a6c98..a31afe14e 100644
--- a/tensor2tensor/data_generators/wikisum/produce_examples.py
+++ b/tensor2tensor/data_generators/wikisum/produce_examples.py
@@ -20,6 +20,7 @@
 
 import os
 
+from six.moves import range
 from tensor2tensor.data_generators.wikisum import utils
 from tensor2tensor.data_generators.wikisum import wikisum
 

From 02962d18889ad64fb1155a723b668bc87f65c493 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 10 Jan 2020 10:25:59 -0800
Subject: [PATCH 2615/2720] Hide internal colab magic methods from OSS, they
 don't work with jupyter.

PiperOrigin-RevId: 289117201
---
 tensor2tensor/notebooks/Transformer_translate.ipynb | 3 +--
 tensor2tensor/notebooks/asr_transformer.ipynb       | 1 -
 tensor2tensor/notebooks/hello_t2t.ipynb             | 1 -
 tensor2tensor/notebooks/t2t_problem.ipynb           | 1 -
 4 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/tensor2tensor/notebooks/Transformer_translate.ipynb b/tensor2tensor/notebooks/Transformer_translate.ipynb
index f32668f8c..0ebe9d6bc 100644
--- a/tensor2tensor/notebooks/Transformer_translate.ipynb
+++ b/tensor2tensor/notebooks/Transformer_translate.ipynb
@@ -108,7 +108,6 @@
         "colab": {}
       },
       "source": [
-        "%tensorflow_version 1.x\n",
         "import tensorflow as tf\n",
         "import os\n",
         "\n",
@@ -1100,4 +1099,4 @@
       "outputs": []
     }
   ]
-}
\ No newline at end of file
+}
diff --git a/tensor2tensor/notebooks/asr_transformer.ipynb b/tensor2tensor/notebooks/asr_transformer.ipynb
index 5c8b103a3..82a0728a8 100644
--- a/tensor2tensor/notebooks/asr_transformer.ipynb
+++ b/tensor2tensor/notebooks/asr_transformer.ipynb
@@ -70,7 +70,6 @@
       },
       "outputs": [],
       "source": [
-        "%tensorflow_version 1.x\n",
         "import tensorflow as tf\n",
         "import matplotlib.pyplot as plt\n",
         "import numpy as np\n",
diff --git a/tensor2tensor/notebooks/hello_t2t.ipynb b/tensor2tensor/notebooks/hello_t2t.ipynb
index 851fecba9..219aa3d77 100644
--- a/tensor2tensor/notebooks/hello_t2t.ipynb
+++ b/tensor2tensor/notebooks/hello_t2t.ipynb
@@ -65,7 +65,6 @@
       "outputs": [],
       "source": [
         "# Imports we need.\n",
-        "%tensorflow_version 1.x\n",
         "import tensorflow as tf\n",
         "import matplotlib.pyplot as plt\n",
         "import numpy as np\n",
diff --git a/tensor2tensor/notebooks/t2t_problem.ipynb b/tensor2tensor/notebooks/t2t_problem.ipynb
index a5d85cfa1..592cbad39 100644
--- a/tensor2tensor/notebooks/t2t_problem.ipynb
+++ b/tensor2tensor/notebooks/t2t_problem.ipynb
@@ -136,7 +136,6 @@
       "source": [
         "#@title Run this only once - Sets up TF Eager execution.\n",
         "\n",
-        "%tensorflow_version 1.x\n",
         "import tensorflow as tf\n",
         "\n",
         "# Enable Eager execution - useful for seeing the generated data.\n",

From bc29d16c23745bcca5dfa22c336e73ea87a94aeb Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 10 Jan 2020 11:16:58 -0800
Subject: [PATCH 2616/2720] Fix implementation of TWO_CLASS_LOG_LIKELIHOOD.

PiperOrigin-RevId: 289128069
---
 tensor2tensor/utils/metrics.py      | 15 ++++++---------
 tensor2tensor/utils/metrics_test.py | 23 +++++++++++++++++++++++
 2 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index edcf2cdc3..32c23751d 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -199,16 +199,13 @@ def two_class_log_likelihood(predictions, labels, weights_fn=None):
     A pair, with the average log likelihood in the first component.
   """
   del weights_fn
-  float_labels = tf.cast(labels, dtype=tf.float64)
   float_predictions = tf.cast(tf.squeeze(predictions), dtype=tf.float64)
-  # likelihood should be just p for class 1, and 1 - p for class 0.
-  # signs is 1 for class 1, and -1 for class 0
-  signs = 2 * float_labels - tf.ones_like(float_labels)
-  # constant_term is 1 for class 0, and 0 for class 1.
-  constant_term = tf.ones_like(float_labels) - float_labels
-  likelihoods = constant_term + signs * float_predictions
-  log_likelihoods = tf.log(likelihoods)
-  avg_log_likelihood = tf.reduce_mean(log_likelihoods)
+  batch_probs = tf.stack([1. - float_predictions, float_predictions], axis=-1)
+  int_labels = tf.cast(tf.squeeze(labels), dtype=tf.int32)
+  onehot_targets = tf.cast(tf.one_hot(int_labels, 2), dtype=tf.float64)
+  chosen_probs = tf.einsum(
+      "ij,ij->i", batch_probs, onehot_targets, name="chosen_probs")
+  avg_log_likelihood = tf.reduce_mean(tf.log(chosen_probs))
   return avg_log_likelihood, tf.constant(1.0)
 
 
diff --git a/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 1fd4c7719..877590d1a 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -111,6 +111,29 @@ def testTwoClassLogLikelihood(self):
       actual = session.run(avg_log_likelihood)
     self.assertAlmostEqual(actual, expected)
 
+  def testTwoClassLogLikelihoodVersusOldImplementation(self):
+    def alt_two_class_log_likelihood_impl(predictions, labels):
+      float_labels = tf.cast(labels, dtype=tf.float64)
+      float_predictions = tf.cast(tf.squeeze(predictions), dtype=tf.float64)
+      # likelihood should be just p for class 1, and 1 - p for class 0.
+      # signs is 1 for class 1, and -1 for class 0
+      signs = 2 * float_labels - tf.ones_like(float_labels)
+      # constant_term is 1 for class 0, and 0 for class 1.
+      constant_term = tf.ones_like(float_labels) - float_labels
+      likelihoods = constant_term + signs * float_predictions
+      log_likelihoods = tf.log(likelihoods)
+      avg_log_likelihood = tf.reduce_mean(log_likelihoods)
+      return avg_log_likelihood
+    predictions = np.random.rand(1, 10, 1)
+    targets = np.random.randint(2, size=10)
+    with self.test_session() as session:
+      new_log_likelihood, _ = metrics.two_class_log_likelihood(
+          predictions, targets)
+      alt_log_likelihood = alt_two_class_log_likelihood_impl(
+          predictions, targets)
+      new_impl, alt_impl = session.run([new_log_likelihood, alt_log_likelihood])
+    self.assertAlmostEqual(new_impl, alt_impl)
+
   def testRMSEMetric(self):
     predictions = np.full((10, 1), 1)  # All 1's
     targets = np.full((10, 1), 3)  # All 3's

From 8b0e68906d312013e495ae1f8e138efff2827180 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 10 Jan 2020 11:42:09 -0800
Subject: [PATCH 2617/2720] Bump version to 1.15.3

PiperOrigin-RevId: 289133315
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index b68e4ce6a..39abd77e4 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.15.2',
+    version='1.15.3',
     description='Tensor2Tensor',
     long_description=(
         'Tensor2Tensor, or T2T for short, is a library of '

From e908bcf9e31a65c703046985f58c4d161c6faddd Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 10 Jan 2020 11:55:19 -0800
Subject: [PATCH 2618/2720] In files where we explicitly import tf.compat.v1,
 don't qualify access of tf.compat.v1 again.

PiperOrigin-RevId: 289135863
---
 tensor2tensor/data_generators/generator_utils.py | 8 ++++----
 tensor2tensor/data_generators/problem_test.py    | 2 +-
 tensor2tensor/layers/common_video_test.py        | 2 +-
 tensor2tensor/layers/discretization_test.py      | 2 +-
 tensor2tensor/layers/latent_layers_test.py       | 2 +-
 tensor2tensor/layers/modalities_test.py          | 2 +-
 tensor2tensor/layers/ngram_test.py               | 2 +-
 tensor2tensor/utils/adafactor.py                 | 2 +-
 tensor2tensor/utils/multistep_optimizer.py       | 2 +-
 tensor2tensor/utils/optimize.py                  | 2 +-
 tensor2tensor/utils/t2t_model_test.py            | 2 +-
 tensor2tensor/utils/test_utils.py                | 4 ++--
 tensor2tensor/utils/test_utils_test.py           | 3 +--
 tensor2tensor/utils/yellowfin.py                 | 6 +++---
 14 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 1fe678e14..0b532caf9 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -867,18 +867,18 @@ def dict_pack(example):
 
   def _standardize(self, dataset, keys):
     """Force dataset structure into a tuple of Tensors."""
-    shapes = tf.compat.v1.data.get_output_shapes(dataset)
+    shapes = tf.data.get_output_shapes(dataset)
 
     if isinstance(shapes, dict):
       keys = keys or tuple(shapes.keys())
       dataset = dataset.map(lambda x: tuple(x[k] for k in keys))
-      shapes = tf.compat.v1.data.get_output_shapes(dataset)
+      shapes = tf.data.get_output_shapes(dataset)
 
     if not all(isinstance(i, tf.TensorShape) for i in shapes):
       # Internally this class expects tuples of Tensors, even for the degenerate
       # case of a single sequence.
       dataset = dataset.map(lambda x: (x,))
-      shapes = tf.compat.v1.data.get_output_shapes(dataset)
+      shapes = tf.data.get_output_shapes(dataset)
 
     for s in shapes:
       if not s.is_compatible_with(tf.TensorShape([None])):
@@ -890,7 +890,7 @@ def _standardize(self, dataset, keys):
     if self._chop_long_sequences and len(shapes) != 1:
       raise ValueError("chop_long_sequences expects a single sequence dataset.")
 
-    token_types = tf.compat.v1.data.get_output_types(dataset)
+    token_types = tf.data.get_output_types(dataset)
     if len(set(token_types)) > 1:
       raise ValueError("Inconsistent dtypes: {}".format(token_types))
 
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 6fdd94521..f0f7b99f4 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -30,7 +30,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 def assert_tensors_equal(sess, t1, t2, n):
diff --git a/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index b7188cd34..90ce75b7d 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -26,7 +26,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class CommonVideoTest(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index 5d4eec265..d0454e2e1 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -24,7 +24,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class DiscretizationTest(tf.test.TestCase):
diff --git a/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
index 097465117..59e6ca506 100644
--- a/tensor2tensor/layers/latent_layers_test.py
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -28,7 +28,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 def imagetransformer_latent_tiny():
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 393c558aa..adbb86414 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -26,7 +26,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class ModalityTest(tf.test.TestCase):
diff --git a/tensor2tensor/layers/ngram_test.py b/tensor2tensor/layers/ngram_test.py
index 0233722e5..958788f23 100644
--- a/tensor2tensor/layers/ngram_test.py
+++ b/tensor2tensor/layers/ngram_test.py
@@ -24,7 +24,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class NGramTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index 58525a740..87617ecd9 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -24,7 +24,7 @@
 import tensorflow.compat.v1 as tf
 
 
-class AdafactorOptimizer(tf.compat.v1.train.Optimizer):
+class AdafactorOptimizer(tf.train.Optimizer):
   """Optimizer that implements the Adafactor algorithm.
 
   Adafactor is described in https://arxiv.org/abs/1804.04235.
diff --git a/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
index f745a5d40..2367c8437 100644
--- a/tensor2tensor/utils/multistep_optimizer.py
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -29,7 +29,7 @@
 import tensorflow.compat.v1 as tf
 
 
-class MultistepAdamOptimizer(tf.compat.v1.train.AdamOptimizer):
+class MultistepAdamOptimizer(tf.train.AdamOptimizer):
   """Adam with SGD updates every n steps with accumulated gradients."""
 
   def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8,
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index b0d848ac9..a88f31ad5 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -184,7 +184,7 @@ def _register_base_optimizer(name, opt):
   _register_base_optimizer(_name, _opt)
 
 
-class ConditionalOptimizer(tf.compat.v1.train.Optimizer):
+class ConditionalOptimizer(tf.train.Optimizer):
   """Conditional optimizer."""
 
   def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disable=super-init-not-called
diff --git a/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index a204b7cc6..7a8423996 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class T2TModelTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/test_utils.py b/tensor2tensor/utils/test_utils.py
index 81fcd5a6f..56912b439 100644
--- a/tensor2tensor/utils/test_utils.py
+++ b/tensor2tensor/utils/test_utils.py
@@ -39,7 +39,7 @@ def run_in_graph_and_eager_modes(func=None,
   For example, consider the following unittest:
 
   ```python
-  tf.compat.v1.enable_eager_execution()
+  tf.enable_eager_execution()
 
   class SomeTest(tf.test.TestCase):
 
@@ -120,5 +120,5 @@ def decorated(self, *args, **kwargs):
 
 
 def test_main():
-  tf.compat.v1.enable_eager_execution()
+  tf.enable_eager_execution()
   tf.test.main()
diff --git a/tensor2tensor/utils/test_utils_test.py b/tensor2tensor/utils/test_utils_test.py
index 0b29391de..f5f949701 100644
--- a/tensor2tensor/utils/test_utils_test.py
+++ b/tensor2tensor/utils/test_utils_test.py
@@ -22,8 +22,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
-
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class RunInGraphAndEagerTest(tf.test.TestCase):
diff --git a/tensor2tensor/utils/yellowfin.py b/tensor2tensor/utils/yellowfin.py
index eb50ad01b..d95820eae 100644
--- a/tensor2tensor/utils/yellowfin.py
+++ b/tensor2tensor/utils/yellowfin.py
@@ -22,9 +22,9 @@
 
 
 # Values for gate_gradients.
-GATE_NONE = tf.compat.v1.train.Optimizer.GATE_NONE
-GATE_OP = tf.compat.v1.train.Optimizer.GATE_OP
-GATE_GRAPH = tf.compat.v1.train.Optimizer.GATE_GRAPH
+GATE_NONE = tf.train.Optimizer.GATE_NONE
+GATE_OP = tf.train.Optimizer.GATE_OP
+GATE_GRAPH = tf.train.Optimizer.GATE_GRAPH
 
 
 class YellowFinOptimizer(object):

From 15e0be6e069c62f0be998863a7a5cade8173467c Mon Sep 17 00:00:00 2001
From: Sharan Narang <sharannarang@google.com>
Date: Fri, 10 Jan 2020 13:27:53 -0800
Subject: [PATCH 2619/2720] Update hparams for neural assistant base model.

PiperOrigin-RevId: 289154271
---
 tensor2tensor/models/neural_assistant.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/models/neural_assistant.py b/tensor2tensor/models/neural_assistant.py
index 01d55bb13..6cc357739 100644
--- a/tensor2tensor/models/neural_assistant.py
+++ b/tensor2tensor/models/neural_assistant.py
@@ -508,7 +508,7 @@ def compute_summary_embedding(input_embeddings, input_lengths, hparams):
 @registry.register_hparams
 def neural_assistant_base():
   """HParams for a base neural_assistant model."""
-  hparams = transformer.transformer_base()
+  hparams = transformer.transformer_tpu()
   hparams.add_hparam("pos_weight", 1.0)  # weight for positive triples
   hparams.add_hparam("similarity_fuction",
                      "bilinear")  # dot_product or bilinear
@@ -521,6 +521,12 @@ def neural_assistant_base():
   hparams.add_hparam("kb_loss_weight", 0.0)  # weight for distant supervision
   hparams.add_hparam("test_triple_num",
                      28483)  # max triples of KB
+  hparams.add_hparam("margin", 0.0)  # KB training max-margin loss
+  hparams.add_hparam(
+      "num_negative_samples",
+      1)  # Sampling number of different adversarial training examples
+  hparams.add_hparam("kb_train_weight", 0.0)
+  # KB_training loss weight which combines Language model and KB selection loss
   return hparams
 
 

From 0f6955b4b330e1b55fc8c1f8d5df225545d3136c Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 10 Jan 2020 14:51:17 -0800
Subject: [PATCH 2620/2720] Latest pip has some bug, so invoke as `python -m
 pip install X` rather than `pip install X`

PiperOrigin-RevId: 289170283
---
 oss_scripts/oss_release.sh | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/oss_scripts/oss_release.sh b/oss_scripts/oss_release.sh
index 6038bae21..ea0e1412c 100755
--- a/oss_scripts/oss_release.sh
+++ b/oss_scripts/oss_release.sh
@@ -14,7 +14,9 @@ git clone https://github.com/tensorflow/tensor2tensor.git
 cd tensor2tensor
 git checkout $GIT_COMMIT_ID
 
-pip install wheel twine pyopenssl
+# Without `python -m` we sometimes get module not callable error:
+# https://stackoverflow.com/questions/58451650/pip-no-longer-working-after-update-error-module-object-is-not-callable
+python -m pip install wheel twine pyopenssl
 
 # Build the distribution
 echo "Building distribution"

From 14e56317f85d50d9f28441d218e3ebbd58142d00 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 10 Jan 2020 14:56:23 -0800
Subject: [PATCH 2621/2720] Print tf version in oss_pip_install just to make
 sure.

PiperOrigin-RevId: 289171246
---
 oss_scripts/oss_pip_install.sh | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/oss_scripts/oss_pip_install.sh b/oss_scripts/oss_pip_install.sh
index ad5dd5e41..c86e619ea 100755
--- a/oss_scripts/oss_pip_install.sh
+++ b/oss_scripts/oss_pip_install.sh
@@ -14,6 +14,9 @@ pip install -q -U setuptools
 pip install -q -U numpy
 pip install -q "tensorflow==$TF_VERSION"
 
+# Just print the version again to make sure.
+python -c 'import tensorflow as tf; print(tf.__version__)'
+
 # First ensure that the base dependencies are sufficient for a full import
 pip install -q -e .
 t2t-trainer --registry_help 2>&1 >/dev/null

From c1b1478acbe698a61ece28b2b36cc4f170d22340 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 10 Jan 2020 15:00:09 -0800
Subject: [PATCH 2622/2720] Drop support for py 2.7, make tests that were
 running on 2.7 run on py3 instead.

PiperOrigin-RevId: 289171890
---
 .travis.yml                         |  1 -
 oss_scripts/oss_integration_test.sh |  2 +-
 oss_scripts/oss_tests.sh            | 25 ++++++++++++-------------
 3 files changed, 13 insertions(+), 15 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 3f9d78e3b..0bc70f988 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -7,7 +7,6 @@ git:
 services:
   - docker
 python:
-  - "2.7"
   - "3.6"
 env:
   global:
diff --git a/oss_scripts/oss_integration_test.sh b/oss_scripts/oss_integration_test.sh
index 46ad0fe06..1700cfae0 100755
--- a/oss_scripts/oss_integration_test.sh
+++ b/oss_scripts/oss_integration_test.sh
@@ -20,7 +20,7 @@ t2t-trainer --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR --model=transformer
 t2t-decoder --problem=$T2T_PROBLEM --data_dir=$T2T_DATA_DIR --model=transformer --hparams_set=transformer_tiny --output_dir=$T2T_TRAIN_DIR --decode_hparams='num_samples=10'
 
 # Test serving
-if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]
+if [[ "$TF_VERSION" == "$TF_LATEST"  ]]
 then
   # Export for serving
   pip install tensorflow_hub
diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
index f85b21917..512015ff6 100755
--- a/oss_scripts/oss_tests.sh
+++ b/oss_scripts/oss_tests.sh
@@ -158,19 +158,7 @@ set_status
 #  set_status
 #fi
 
-if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]
-then
-    # TODO(afrozm): Once we drop support for 1.10 we can get rid of this.
-    pytest --disable-warnings \
-      tensor2tensor/utils/beam_search_test.py::BeamSearchTest::testTPUBeam
-    set_status
-    # TODO(afrozm): Enable other tests in the RL directory.
-    # Can't add disable warning here since it parses flags.
-    pytest tensor2tensor/rl/trainer_model_based_test.py
-    set_status
-fi
-
-if [[ "$TRAVIS_PYTHON_VERSION" == "3.6" ]] && [[ "$TF_VERSION" == "$TF_LATEST"  ]]
+if [[ "$TF_VERSION" == "$TF_LATEST"  ]]
 then
     jupyter nbconvert --ExecutePreprocessor.kernel_name=python3 \
       --ExecutePreprocessor.timeout=600 --to notebook --execute \
@@ -181,6 +169,17 @@ then
       --ExecutePreprocessor.timeout=600 --to notebook --execute \
       tensor2tensor/notebooks/t2t_problem.ipynb;
     set_status
+
+    # TODO(afrozm): Once we drop support for 1.10 we can get rid of this.
+    pytest --disable-warnings \
+      tensor2tensor/utils/beam_search_test.py::BeamSearchTest::testTPUBeam
+    set_status
+
+    # TODO(afrozm): Enable other tests in the RL directory.
+    # Can't add disable warning here since it parses flags.
+    pytest tensor2tensor/rl/trainer_model_based_test.py
+    set_status
+
 fi
 
 # Test --t2t_usr_dir

From 9ee82545668f6d3dfe666625727cf675217ca7dd Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 10 Jan 2020 16:28:16 -0800
Subject: [PATCH 2623/2720] Flush out some more contrib/slim.

PiperOrigin-RevId: 289187366
---
 tensor2tensor/data_generators/allen_brain.py | 4 ++--
 tensor2tensor/data_generators/video_utils.py | 8 ++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index 67afce1dd..056f6a8f2 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -40,11 +40,11 @@
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.layers import modalities
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-from tensorflow.contrib import slim as contrib_slim
 
 _BASE_EXAMPLE_IMAGE_SIZE = 64
 
@@ -351,7 +351,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "targets":
-            contrib_slim.tfexample_decoder.Image(
+            contrib.slim().tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index b816ef8dc..4ea7e244c 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -32,10 +32,10 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 from tensor2tensor.layers import modalities
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import video_metrics
 import tensorflow as tf
-from tensorflow.contrib import slim as contrib_slim
 
 
 FLAGS = flags.FLAGS
@@ -385,7 +385,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "frame":
-            contrib_slim.tfexample_decoder.Image(
+            contrib.slim().tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 shape=[self.frame_height, self.frame_width, self.num_channels],
@@ -677,7 +677,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "inputs":
-            contrib_slim.tfexample_decoder.Image(
+            contrib.slim().tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
@@ -766,7 +766,7 @@ def example_reading_spec(self):
     data_fields, data_items_to_decoders = (
         super(Video2ClassProblem, self).example_reading_spec())
     data_fields[label_key] = tf.FixedLenFeature((1,), tf.int64)
-    data_items_to_decoders["targets"] = contrib_slim.tfexample_decoder.Tensor(
+    data_items_to_decoders["targets"] = contrib.slim().tfexample_decoder.Tensor(
         label_key)
     return data_fields, data_items_to_decoders
 

From a338743aa3b757c7f401ac47f6a8a8f2606d1c49 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 10 Jan 2020 16:45:23 -0800
Subject: [PATCH 2624/2720] Bump up version to 1.15.4

PiperOrigin-RevId: 289189869
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 39abd77e4..1482f9f7a 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.15.3',
+    version='1.15.4',
     description='Tensor2Tensor',
     long_description=(
         'Tensor2Tensor, or T2T for short, is a library of '

From 19a8e7e27bfac03bfe6bc5a9b31d24af63e2681b Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Mon, 13 Jan 2020 14:19:29 -0800
Subject: [PATCH 2625/2720] Move to python 3.

PiperOrigin-RevId: 289518000
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 409b978c0..bb90f1e9d 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -4396,7 +4396,7 @@ def compute_attention_component(antecedent,
           "KFAC implementation only supports filter_width=1 (actual: {}) and "
           "vars_3d_num_heads=0 (actual: {}).".format(
               filter_width, vars_3d_num_heads))
-  if vars_3d_num_heads > 0:
+  if vars_3d_num_heads is not None and vars_3d_num_heads > 0:
     assert filter_width == 1
     input_depth = antecedent.get_shape().as_list()[-1]
     depth_per_head = total_depth // vars_3d_num_heads

From 87e4531513aa5c50f5c3843a527d368eee30d50c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 14 Jan 2020 11:44:43 -0800
Subject: [PATCH 2626/2720] Hide internal colab magic methods from OSS, they
 don't work with jupyter.

PiperOrigin-RevId: 289697475
---
 tensor2tensor/notebooks/Transformer_translate.ipynb | 3 +++
 tensor2tensor/notebooks/asr_transformer.ipynb       | 3 +++
 tensor2tensor/notebooks/hello_t2t.ipynb             | 3 +++
 tensor2tensor/notebooks/t2t_problem.ipynb           | 3 +++
 4 files changed, 12 insertions(+)

diff --git a/tensor2tensor/notebooks/Transformer_translate.ipynb b/tensor2tensor/notebooks/Transformer_translate.ipynb
index 0ebe9d6bc..4e9925498 100644
--- a/tensor2tensor/notebooks/Transformer_translate.ipynb
+++ b/tensor2tensor/notebooks/Transformer_translate.ipynb
@@ -108,6 +108,9 @@
         "colab": {}
       },
       "source": [
+        "import sys\n",
+        "if 'google.colab' in sys.modules: # Colab-only TensorFlow version selector\n",
+        "  %tensorflow_version 1.x\n",
         "import tensorflow as tf\n",
         "import os\n",
         "\n",
diff --git a/tensor2tensor/notebooks/asr_transformer.ipynb b/tensor2tensor/notebooks/asr_transformer.ipynb
index 82a0728a8..71a8bf456 100644
--- a/tensor2tensor/notebooks/asr_transformer.ipynb
+++ b/tensor2tensor/notebooks/asr_transformer.ipynb
@@ -70,6 +70,9 @@
       },
       "outputs": [],
       "source": [
+        "import sys\n",
+        "if 'google.colab' in sys.modules: # Colab-only TensorFlow version selector\n",
+        "  %tensorflow_version 1.x\n",
         "import tensorflow as tf\n",
         "import matplotlib.pyplot as plt\n",
         "import numpy as np\n",
diff --git a/tensor2tensor/notebooks/hello_t2t.ipynb b/tensor2tensor/notebooks/hello_t2t.ipynb
index 219aa3d77..b69fd48c1 100644
--- a/tensor2tensor/notebooks/hello_t2t.ipynb
+++ b/tensor2tensor/notebooks/hello_t2t.ipynb
@@ -65,6 +65,9 @@
       "outputs": [],
       "source": [
         "# Imports we need.\n",
+        "import sys\n",
+        "if 'google.colab' in sys.modules: # Colab-only TensorFlow version selector\n",
+        "  %tensorflow_version 1.x\n",
         "import tensorflow as tf\n",
         "import matplotlib.pyplot as plt\n",
         "import numpy as np\n",
diff --git a/tensor2tensor/notebooks/t2t_problem.ipynb b/tensor2tensor/notebooks/t2t_problem.ipynb
index 592cbad39..91e5235a1 100644
--- a/tensor2tensor/notebooks/t2t_problem.ipynb
+++ b/tensor2tensor/notebooks/t2t_problem.ipynb
@@ -136,6 +136,9 @@
       "source": [
         "#@title Run this only once - Sets up TF Eager execution.\n",
         "\n",
+        "import sys\n",
+        "if 'google.colab' in sys.modules: # Colab-only TensorFlow version selector\n",
+        "  %tensorflow_version 1.x\n",
         "import tensorflow as tf\n",
         "\n",
         "# Enable Eager execution - useful for seeing the generated data.\n",

From 2cb949d19e41d32cc205c34bfa1068aa2de54f17 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 14 Jan 2020 13:46:20 -0800
Subject: [PATCH 2627/2720] Correct some more tf.contrib violations.

PiperOrigin-RevId: 289721730
---
 tensor2tensor/notebooks/t2t_problem.ipynb | 86 +++++------------------
 1 file changed, 16 insertions(+), 70 deletions(-)

diff --git a/tensor2tensor/notebooks/t2t_problem.ipynb b/tensor2tensor/notebooks/t2t_problem.ipynb
index 91e5235a1..6e1eb1311 100644
--- a/tensor2tensor/notebooks/t2t_problem.ipynb
+++ b/tensor2tensor/notebooks/t2t_problem.ipynb
@@ -21,12 +21,7 @@
       "execution_count": 0,
       "metadata": {
         "cellView": "form",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        },
+        "colab": {},
         "colab_type": "code",
         "id": "FesA0dakI2kh"
       },
@@ -100,13 +95,8 @@
       "cell_type": "code",
       "execution_count": 0,
       "metadata": {
-        "cellView": "form",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        },
+        "cellView": "both",
+        "colab": {},
         "colab_type": "code",
         "id": "IBWBeE39JYaR"
       },
@@ -122,13 +112,8 @@
       "cell_type": "code",
       "execution_count": 0,
       "metadata": {
-        "cellView": "form",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        },
+        "cellView": "both",
+        "colab": {},
         "colab_type": "code",
         "id": "sbTULiroLs2w"
       },
@@ -150,12 +135,7 @@
       "execution_count": 0,
       "metadata": {
         "cellView": "form",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        },
+        "colab": {},
         "colab_type": "code",
         "id": "A8JljOzDYF-Z"
       },
@@ -175,12 +155,7 @@
       "execution_count": 0,
       "metadata": {
         "cellView": "form",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        },
+        "colab": {},
         "colab_type": "code",
         "id": "ioW-V1qpqSCE"
       },
@@ -228,12 +203,7 @@
       "execution_count": 0,
       "metadata": {
         "cellView": "form",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        },
+        "colab": {},
         "colab_type": "code",
         "id": "pDDiPxqg9UF-"
       },
@@ -265,12 +235,7 @@
       "cell_type": "code",
       "execution_count": 0,
       "metadata": {
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        },
+        "colab": {},
         "colab_type": "code",
         "id": "KcT_x4ma-Uaq"
       },
@@ -447,12 +412,7 @@
       "cell_type": "code",
       "execution_count": 0,
       "metadata": {
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        },
+        "colab": {},
         "colab_type": "code",
         "id": "atYWRpM1FgaJ"
       },
@@ -479,24 +439,18 @@
       "cell_type": "code",
       "execution_count": 0,
       "metadata": {
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        },
+        "colab": {},
         "colab_type": "code",
         "id": "PZczDWnOQDp2"
       },
       "outputs": [],
       "source": [
-        "tfe = tf.contrib.eager\n",
-        "\n",
         "Modes = tf.estimator.ModeKeys\n",
         "\n",
         "# We can iterate over our examples by making an iterator and calling next on it.\n",
-        "eager_iterator = tfe.Iterator(sort_len_problem.dataset(Modes.EVAL, DATA_DIR))\n",
-        "example = eager_iterator.next()\n",
+        "sort_len_problem_dataset = sort_len_problem.dataset(Modes.EVAL, DATA_DIR)\n",
+        "eager_iterator = sort_len_problem_dataset.make_one_shot_iterator()\n",
+        "example = next(eager_iterator)\n",
         "\n",
         "input_tensor = example[\"inputs\"]\n",
         "target_tensor = example[\"targets\"]\n",
@@ -511,12 +465,7 @@
       "cell_type": "code",
       "execution_count": 0,
       "metadata": {
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        },
+        "colab": {},
         "colab_type": "code",
         "id": "1DtfzgqivAxl"
       },
@@ -552,16 +501,13 @@
   "metadata": {
     "colab": {
       "collapsed_sections": [],
-      "default_view": {},
       "name": "t2t_problem.ipynb",
       "provenance": [
         {
           "file_id": "1FwspR4PzEZAiQCGziob5oov-8DyEXSnw",
           "timestamp": 1533664607636
         }
-      ],
-      "version": "0.3.2",
-      "views": {}
+      ]
     },
     "kernelspec": {
       "display_name": "Python 3",

From 2cfa849963c9147fbe9985392b7270838614d74d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sat, 18 Jan 2020 00:48:24 -0800
Subject: [PATCH 2628/2720] Export placeholder for sampling temperature in
 model signature.

PiperOrigin-RevId: 290405441
---
 tensor2tensor/data_generators/problem.py | 11 ++++++++++-
 tensor2tensor/layers/common_layers.py    |  2 +-
 tensor2tensor/models/transformer.py      | 19 +++++++++++++++----
 tensor2tensor/utils/t2t_model.py         |  3 ++-
 4 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 7732225ce..9ad0ae410 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -908,6 +908,10 @@ def serving_input_fn(self, hparams, decode_hparams=None, use_tpu=False):
     mode = tf.estimator.ModeKeys.PREDICT
     serialized_example = tf.placeholder(
         dtype=tf.string, shape=[None], name="serialized_example")
+    sampling_temp = tf.placeholder_with_default(
+        tf.constant(getattr(hparams, "sampling_temp", 0.0), dtype=tf.float32),
+        shape=[],
+        name="sampling_temp")
     dataset = tf.data.Dataset.from_tensor_slices(serialized_example)
     dataset = dataset.map(self.decode_example)
     dataset = dataset.map(lambda ex: self.preprocess_example(ex, mode, hparams))
@@ -929,12 +933,17 @@ def serving_input_fn(self, hparams, decode_hparams=None, use_tpu=False):
 
     dataset = dataset.map(data_reader.standardize_shapes)
     features = tf.data.experimental.get_single_element(dataset)
+    features["sampling_temp"] = sampling_temp
 
     if self.has_inputs:
       features.pop("targets", None)
 
     return tf.estimator.export.ServingInputReceiver(
-        features=features, receiver_tensors=serialized_example)
+        features=features,
+        receiver_tensors={
+            "input": serialized_example,
+            "sampling_temp": sampling_temp
+        })
 
 
 class FeatureInfo(object):
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index ca6621635..4a4496d50 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2949,7 +2949,7 @@ def sample_with_temperature(logits, temperature, sampling_keep_top_k=-1):
     argmax = tf.argmax(tf.reshape(logits, [-1, logits_shape[-1]]), axis=1)
     return tf.reshape(argmax, logits_shape[:-1])
   else:
-    assert temperature > 0.0
+    tf.debugging.assert_greater(temperature, 0.0)
 
     if sampling_keep_top_k != -1:
       if sampling_keep_top_k <= 0:
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 31222cf6d..927a4e9bf 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -623,6 +623,9 @@ def forced_logits():
       return ret, cache
 
     eos_id = self.get_decode_end_id() or beam_search.EOS_ID
+    temperature = features.get("sampling_temp",
+                               getattr(hparams, "sampling_temp", 0.0))
+
     ret = fast_decode_tpu(
         encoder_output=encoder_output,
         encoder_decoder_attention_bias=encoder_decoder_attention_bias,
@@ -636,7 +639,8 @@ def forced_logits():
         alpha=alpha,
         batch_size=batch_size,
         force_decode_length=self._decode_hparams.force_decode_length,
-        eos_id=eos_id)
+        eos_id=eos_id,
+        sampling_temperature=temperature)
     if partial_targets is not None:
       if beam_size <= 1 or top_beams <= 1:
         ret["outputs"] = ret["outputs"][:, partial_targets_length:]
@@ -883,6 +887,8 @@ def forced_logits():
 
     sos_id = self.get_decode_start_id() or 0
     eos_id = self.get_decode_end_id() or beam_search.EOS_ID
+    temperature = features.get("sampling_temp",
+                               getattr(hparams, "sampling_temp", 0.0))
 
     ret = fast_decode(
         encoder_output=encoder_output,
@@ -899,6 +905,7 @@ def forced_logits():
         force_decode_length=self._decode_hparams.force_decode_length,
         sos_id=sos_id,
         eos_id=eos_id,
+        sampling_temperature=temperature,
         cache=att_cache)
     if partial_targets is not None:
       if beam_size <= 1 or top_beams <= 1:
@@ -986,7 +993,8 @@ def fast_decode_tpu(encoder_output,
                     batch_size=None,
                     force_decode_length=False,
                     scope_prefix="body/",
-                    use_top_k_with_unique=True):
+                    use_top_k_with_unique=True,
+                    sampling_temperature=0.0):
   """Given encoder output and a symbols to logits function, does fast decoding.
 
   Implements both greedy and beam search decoding for TPU, uses beam search iff
@@ -1014,6 +1022,7 @@ def fast_decode_tpu(encoder_output,
     scope_prefix: str, prefix for decoder layer variable scopes.
     use_top_k_with_unique: bool, whether to use a fast (but decreased precision)
       top_k during beam search.
+    sampling_temperature: scalar, temperature with which to sample.
 
   Returns:
     A dict of decoding results {
@@ -1071,7 +1080,7 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       """One step of greedy decoding."""
       logits, cache = symbols_to_logits_fn(next_id, i, cache)
       log_probs = common_layers.log_prob_from_logits(logits)
-      temperature = getattr(hparams, "sampling_temp", 0.0)
+      temperature = sampling_temperature
       keep_top = getattr(hparams, "sampling_keep_top_k", -1)
       if hparams.sampling_method == "argmax":
         temperature = 0.0
@@ -1142,6 +1151,7 @@ def fast_decode(encoder_output,
                 batch_size=None,
                 force_decode_length=False,
                 scope_prefix="body/",
+                sampling_temperature=0.0,
                 cache=None):
   """Given encoder output and a symbols to logits function, does fast decoding.
 
@@ -1168,6 +1178,7 @@ def fast_decode(encoder_output,
     force_decode_length: bool, whether to force the full decode length, or if
       False, stop when all beams hit eos_id.
     scope_prefix: str, prefix for decoder layer variable scopes.
+    sampling_temperature: scalar, temperature with which to sample.
     cache: cache dictionary for additional predictions.
 
   Returns:
@@ -1216,7 +1227,7 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       """One step of greedy decoding."""
       logits, cache = symbols_to_logits_fn(next_id, i, cache)
       log_probs = common_layers.log_prob_from_logits(logits)
-      temperature = getattr(hparams, "sampling_temp", 0.0)
+      temperature = sampling_temperature
       keep_top = getattr(hparams, "sampling_keep_top_k", -1)
       if hparams.sampling_method == "argmax":
         temperature = 0.0
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 40f4da8fa..d4bf75c93 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1357,7 +1357,8 @@ def multinomial_squeeze(logits, temperature=1.0):
         choices = tf.reshape(choices, logits_shape[:-1])
         return choices
 
-      samples = multinomial_squeeze(logits, self.hparams.sampling_temp)
+      temperature = features.get("sampling_temp", self.hparams.sampling_temp)
+      samples = multinomial_squeeze(logits, temperature)
 
     return samples, logits, losses
 

From d62e2ee1b069d3d9b327d4d2dd6f9e50b7e62bb3 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sun, 19 Jan 2020 16:08:25 -0800
Subject: [PATCH 2629/2720] Expose sampling temperate as additional serving
 signature.

PiperOrigin-RevId: 290535097
---
 tensor2tensor/data_generators/problem.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 9ad0ae410..bd2bcb7da 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -940,9 +940,12 @@ def serving_input_fn(self, hparams, decode_hparams=None, use_tpu=False):
 
     return tf.estimator.export.ServingInputReceiver(
         features=features,
-        receiver_tensors={
-            "input": serialized_example,
-            "sampling_temp": sampling_temp
+        receiver_tensors=serialized_example,
+        receiver_tensors_alternatives={
+            "sample": {
+                "input": serialized_example,
+                "sampling_temp": sampling_temp
+            }
         })
 
 

From e1f0e3a746bb322f4bf3975fad2c8105b3a43a49 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Thu, 23 Jan 2020 15:00:20 -0800
Subject: [PATCH 2630/2720] Remove unknown flag from t2t_trainer.

PiperOrigin-RevId: 291251966
---
 tensor2tensor/__init__.py                                     | 2 +-
 tensor2tensor/bin/__init__.py                                 | 2 +-
 tensor2tensor/bin/build_vocab.py                              | 2 +-
 tensor2tensor/bin/make_tf_configs.py                          | 2 +-
 tensor2tensor/bin/t2t_attack.py                               | 2 +-
 tensor2tensor/bin/t2t_avg_all.py                              | 2 +-
 tensor2tensor/bin/t2t_bleu.py                                 | 2 +-
 tensor2tensor/bin/t2t_datagen.py                              | 2 +-
 tensor2tensor/bin/t2t_decoder.py                              | 2 +-
 tensor2tensor/bin/t2t_distill.py                              | 2 +-
 tensor2tensor/bin/t2t_eval.py                                 | 2 +-
 tensor2tensor/bin/t2t_prune.py                                | 2 +-
 tensor2tensor/bin/t2t_trainer.py                              | 4 ++--
 tensor2tensor/bin/t2t_trainer_test.py                         | 2 +-
 tensor2tensor/bin/t2t_translate_all.py                        | 2 +-
 tensor2tensor/data_generators/__init__.py                     | 2 +-
 tensor2tensor/data_generators/algorithmic.py                  | 2 +-
 tensor2tensor/data_generators/algorithmic_math.py             | 2 +-
 tensor2tensor/data_generators/algorithmic_math_deepmind.py    | 2 +-
 tensor2tensor/data_generators/algorithmic_math_test.py        | 2 +-
 .../data_generators/algorithmic_math_two_variables.py         | 2 +-
 tensor2tensor/data_generators/algorithmic_test.py             | 2 +-
 tensor2tensor/data_generators/all_problems.py                 | 2 +-
 tensor2tensor/data_generators/allen_brain.py                  | 2 +-
 tensor2tensor/data_generators/allen_brain_test.py             | 2 +-
 tensor2tensor/data_generators/audio.py                        | 2 +-
 tensor2tensor/data_generators/audio_encoder.py                | 2 +-
 tensor2tensor/data_generators/audio_test.py                   | 2 +-
 tensor2tensor/data_generators/babi_qa.py                      | 2 +-
 tensor2tensor/data_generators/bair_robot_pushing.py           | 2 +-
 tensor2tensor/data_generators/celeba.py                       | 2 +-
 tensor2tensor/data_generators/celeba_test.py                  | 2 +-
 tensor2tensor/data_generators/celebahq.py                     | 2 +-
 tensor2tensor/data_generators/cifar.py                        | 2 +-
 tensor2tensor/data_generators/cipher.py                       | 2 +-
 tensor2tensor/data_generators/cleaner_en_xx.py                | 2 +-
 tensor2tensor/data_generators/cnn_dailymail.py                | 2 +-
 tensor2tensor/data_generators/cola.py                         | 2 +-
 tensor2tensor/data_generators/common_voice.py                 | 2 +-
 tensor2tensor/data_generators/common_voice_test.py            | 2 +-
 tensor2tensor/data_generators/conll_ner.py                    | 2 +-
 tensor2tensor/data_generators/desc2code.py                    | 2 +-
 tensor2tensor/data_generators/desc2code_test.py               | 2 +-
 tensor2tensor/data_generators/dialog_abstract.py              | 2 +-
 tensor2tensor/data_generators/dialog_cornell.py               | 2 +-
 tensor2tensor/data_generators/dialog_dailydialog.py           | 2 +-
 tensor2tensor/data_generators/dialog_opensubtitles.py         | 2 +-
 tensor2tensor/data_generators/dialog_personachat.py           | 2 +-
 tensor2tensor/data_generators/dna_encoder.py                  | 2 +-
 tensor2tensor/data_generators/dna_encoder_test.py             | 2 +-
 tensor2tensor/data_generators/enwik8.py                       | 2 +-
 tensor2tensor/data_generators/fsns.py                         | 2 +-
 tensor2tensor/data_generators/function_docstring.py           | 2 +-
 tensor2tensor/data_generators/gene_expression.py              | 2 +-
 tensor2tensor/data_generators/gene_expression_test.py         | 2 +-
 tensor2tensor/data_generators/generator_utils.py              | 2 +-
 tensor2tensor/data_generators/generator_utils_test.py         | 2 +-
 tensor2tensor/data_generators/google_robot_pushing.py         | 2 +-
 tensor2tensor/data_generators/gym_env.py                      | 2 +-
 tensor2tensor/data_generators/gym_env_test.py                 | 2 +-
 tensor2tensor/data_generators/ice_parsing.py                  | 2 +-
 tensor2tensor/data_generators/image_lsun.py                   | 2 +-
 tensor2tensor/data_generators/image_utils.py                  | 2 +-
 tensor2tensor/data_generators/image_utils_test.py             | 2 +-
 tensor2tensor/data_generators/imagenet.py                     | 2 +-
 tensor2tensor/data_generators/imagenet_test.py                | 2 +-
 tensor2tensor/data_generators/imdb.py                         | 2 +-
 tensor2tensor/data_generators/inspect_tfrecord.py             | 2 +-
 tensor2tensor/data_generators/lambada.py                      | 2 +-
 tensor2tensor/data_generators/librispeech.py                  | 2 +-
 tensor2tensor/data_generators/lm1b.py                         | 2 +-
 tensor2tensor/data_generators/lm1b_imdb.py                    | 2 +-
 tensor2tensor/data_generators/lm1b_mnli.py                    | 2 +-
 tensor2tensor/data_generators/mnist.py                        | 2 +-
 tensor2tensor/data_generators/moving_mnist.py                 | 2 +-
 tensor2tensor/data_generators/mrpc.py                         | 2 +-
 tensor2tensor/data_generators/mscoco.py                       | 2 +-
 tensor2tensor/data_generators/mscoco_test.py                  | 2 +-
 tensor2tensor/data_generators/multi_problem.py                | 2 +-
 tensor2tensor/data_generators/multi_problem_v2.py             | 2 +-
 tensor2tensor/data_generators/multi_problem_v2_test.py        | 2 +-
 tensor2tensor/data_generators/multinli.py                     | 2 +-
 tensor2tensor/data_generators/ocr.py                          | 2 +-
 tensor2tensor/data_generators/ops/pack_sequences_ops_test.py  | 2 +-
 .../data_generators/ops/subword_text_encoder_ops_test.py      | 2 +-
 tensor2tensor/data_generators/paraphrase_ms_coco.py           | 2 +-
 tensor2tensor/data_generators/paraphrase_ms_coco_test.py      | 2 +-
 tensor2tensor/data_generators/pointer_generator_word.py       | 2 +-
 tensor2tensor/data_generators/problem.py                      | 2 +-
 tensor2tensor/data_generators/problem_hparams.py              | 2 +-
 tensor2tensor/data_generators/problem_test.py                 | 2 +-
 tensor2tensor/data_generators/program_search.py               | 2 +-
 tensor2tensor/data_generators/program_search_test.py          | 2 +-
 tensor2tensor/data_generators/ptb.py                          | 2 +-
 tensor2tensor/data_generators/qnli.py                         | 2 +-
 tensor2tensor/data_generators/quora_qpairs.py                 | 2 +-
 tensor2tensor/data_generators/rte.py                          | 2 +-
 tensor2tensor/data_generators/scitail.py                      | 2 +-
 tensor2tensor/data_generators/snli.py                         | 2 +-
 tensor2tensor/data_generators/speech_recognition.py           | 2 +-
 tensor2tensor/data_generators/squad.py                        | 2 +-
 tensor2tensor/data_generators/sst_binary.py                   | 2 +-
 tensor2tensor/data_generators/stanford_nli.py                 | 2 +-
 tensor2tensor/data_generators/style_transfer.py               | 2 +-
 tensor2tensor/data_generators/style_transfer_test.py          | 2 +-
 tensor2tensor/data_generators/subject_verb_agreement.py       | 2 +-
 tensor2tensor/data_generators/text_encoder.py                 | 2 +-
 tensor2tensor/data_generators/text_encoder_build_subword.py   | 2 +-
 tensor2tensor/data_generators/text_encoder_test.py            | 2 +-
 tensor2tensor/data_generators/text_problems.py                | 2 +-
 tensor2tensor/data_generators/text_problems_test.py           | 2 +-
 tensor2tensor/data_generators/timeseries.py                   | 2 +-
 tensor2tensor/data_generators/timeseries_data_generator.py    | 2 +-
 .../data_generators/timeseries_data_generator_test.py         | 2 +-
 tensor2tensor/data_generators/timeseries_test.py              | 2 +-
 tensor2tensor/data_generators/tokenizer.py                    | 2 +-
 tensor2tensor/data_generators/tokenizer_test.py               | 2 +-
 tensor2tensor/data_generators/transduction_problems.py        | 2 +-
 tensor2tensor/data_generators/transduction_problems_test.py   | 2 +-
 tensor2tensor/data_generators/translate.py                    | 2 +-
 tensor2tensor/data_generators/translate_encs.py               | 2 +-
 tensor2tensor/data_generators/translate_ende.py               | 2 +-
 tensor2tensor/data_generators/translate_ende_test.py          | 2 +-
 tensor2tensor/data_generators/translate_enes.py               | 2 +-
 tensor2tensor/data_generators/translate_enet.py               | 2 +-
 tensor2tensor/data_generators/translate_enfr.py               | 2 +-
 tensor2tensor/data_generators/translate_enid.py               | 2 +-
 tensor2tensor/data_generators/translate_enmk.py               | 2 +-
 tensor2tensor/data_generators/translate_enro.py               | 2 +-
 tensor2tensor/data_generators/translate_entn.py               | 2 +-
 tensor2tensor/data_generators/translate_envi.py               | 2 +-
 tensor2tensor/data_generators/translate_enzh.py               | 2 +-
 tensor2tensor/data_generators/translate_test.py               | 2 +-
 tensor2tensor/data_generators/video_generated.py              | 2 +-
 tensor2tensor/data_generators/video_utils.py                  | 2 +-
 tensor2tensor/data_generators/video_utils_test.py             | 2 +-
 tensor2tensor/data_generators/vqa.py                          | 2 +-
 tensor2tensor/data_generators/vqa_utils.py                    | 2 +-
 tensor2tensor/data_generators/wiki.py                         | 2 +-
 tensor2tensor/data_generators/wiki_lm.py                      | 2 +-
 tensor2tensor/data_generators/wiki_multi_problems.py          | 2 +-
 tensor2tensor/data_generators/wiki_revision.py                | 2 +-
 tensor2tensor/data_generators/wiki_revision_utils.py          | 2 +-
 tensor2tensor/data_generators/wikisum/__init__.py             | 2 +-
 tensor2tensor/data_generators/wikisum/generate_vocab.py       | 2 +-
 .../data_generators/wikisum/get_references_commoncrawl.py     | 2 +-
 tensor2tensor/data_generators/wikisum/get_references_web.py   | 2 +-
 .../wikisum/get_references_web_single_group.py                | 2 +-
 tensor2tensor/data_generators/wikisum/html.py                 | 2 +-
 tensor2tensor/data_generators/wikisum/parallel_launch.py      | 2 +-
 tensor2tensor/data_generators/wikisum/produce_examples.py     | 2 +-
 tensor2tensor/data_generators/wikisum/utils.py                | 2 +-
 tensor2tensor/data_generators/wikisum/utils_test.py           | 2 +-
 tensor2tensor/data_generators/wikisum/validate_data.py        | 2 +-
 tensor2tensor/data_generators/wikisum/wikisum.py              | 2 +-
 tensor2tensor/data_generators/wikitext103.py                  | 2 +-
 tensor2tensor/data_generators/wnli.py                         | 2 +-
 tensor2tensor/data_generators/wsj_parsing.py                  | 2 +-
 tensor2tensor/data_generators/yelp_full.py                    | 2 +-
 tensor2tensor/data_generators/yelp_polarity.py                | 2 +-
 tensor2tensor/envs/__init__.py                                | 2 +-
 tensor2tensor/envs/env_problem.py                             | 2 +-
 tensor2tensor/envs/env_problem_utils.py                       | 2 +-
 tensor2tensor/envs/env_problem_utils_test.py                  | 2 +-
 tensor2tensor/envs/gym_env_problem.py                         | 2 +-
 tensor2tensor/envs/gym_env_problem_test.py                    | 2 +-
 tensor2tensor/envs/gym_spaces_utils.py                        | 2 +-
 tensor2tensor/envs/gym_spaces_utils_test.py                   | 2 +-
 tensor2tensor/envs/mujoco_problems.py                         | 2 +-
 tensor2tensor/envs/mujoco_problems_test.py                    | 2 +-
 tensor2tensor/envs/rendered_env_problem.py                    | 2 +-
 tensor2tensor/envs/rendered_env_problem_test.py               | 2 +-
 tensor2tensor/envs/tic_tac_toe_env.py                         | 2 +-
 tensor2tensor/envs/tic_tac_toe_env_problem.py                 | 2 +-
 tensor2tensor/envs/tic_tac_toe_env_problem_test.py            | 2 +-
 tensor2tensor/envs/tic_tac_toe_env_test.py                    | 2 +-
 tensor2tensor/envs/time_step.py                               | 2 +-
 tensor2tensor/envs/time_step_test.py                          | 2 +-
 tensor2tensor/envs/trajectory.py                              | 2 +-
 tensor2tensor/envs/trajectory_test.py                         | 2 +-
 tensor2tensor/insights/__init__.py                            | 2 +-
 tensor2tensor/insights/graph.py                               | 2 +-
 tensor2tensor/insights/query_processor.py                     | 2 +-
 tensor2tensor/insights/server.py                              | 2 +-
 tensor2tensor/insights/transformer_model.py                   | 2 +-
 tensor2tensor/layers/__init__.py                              | 2 +-
 tensor2tensor/layers/area_attention.py                        | 2 +-
 tensor2tensor/layers/area_attention_test.py                   | 2 +-
 tensor2tensor/layers/common_attention.py                      | 2 +-
 tensor2tensor/layers/common_attention_test.py                 | 2 +-
 tensor2tensor/layers/common_audio.py                          | 2 +-
 tensor2tensor/layers/common_hparams.py                        | 2 +-
 tensor2tensor/layers/common_image_attention.py                | 2 +-
 tensor2tensor/layers/common_image_attention_test.py           | 2 +-
 tensor2tensor/layers/common_layers.py                         | 2 +-
 tensor2tensor/layers/common_layers_test.py                    | 2 +-
 tensor2tensor/layers/common_video.py                          | 2 +-
 tensor2tensor/layers/common_video_test.py                     | 2 +-
 tensor2tensor/layers/discretization.py                        | 2 +-
 tensor2tensor/layers/discretization_test.py                   | 2 +-
 tensor2tensor/layers/latent_layers.py                         | 2 +-
 tensor2tensor/layers/latent_layers_test.py                    | 2 +-
 tensor2tensor/layers/message_passing_attention.py             | 2 +-
 tensor2tensor/layers/modalities.py                            | 2 +-
 tensor2tensor/layers/modalities_test.py                       | 2 +-
 tensor2tensor/layers/ngram.py                                 | 2 +-
 tensor2tensor/layers/ngram_test.py                            | 2 +-
 tensor2tensor/layers/transformer_layers.py                    | 2 +-
 tensor2tensor/layers/transformer_memory.py                    | 2 +-
 tensor2tensor/layers/transformer_memory_test.py               | 2 +-
 tensor2tensor/layers/vq_discrete.py                           | 2 +-
 tensor2tensor/layers/vqa_layers.py                            | 2 +-
 tensor2tensor/metrics/__init__.py                             | 2 +-
 tensor2tensor/metrics/video_conditional_fvd.py                | 2 +-
 tensor2tensor/metrics/video_conditional_fvd_test.py           | 2 +-
 tensor2tensor/models/__init__.py                              | 2 +-
 tensor2tensor/models/basic.py                                 | 2 +-
 tensor2tensor/models/basic_test.py                            | 2 +-
 tensor2tensor/models/bytenet.py                               | 2 +-
 tensor2tensor/models/bytenet_test.py                          | 2 +-
 tensor2tensor/models/distillation.py                          | 2 +-
 tensor2tensor/models/evolved_transformer.py                   | 2 +-
 tensor2tensor/models/evolved_transformer_test.py              | 2 +-
 tensor2tensor/models/image_transformer.py                     | 2 +-
 tensor2tensor/models/image_transformer_2d.py                  | 2 +-
 tensor2tensor/models/image_transformer_2d_test.py             | 2 +-
 tensor2tensor/models/image_transformer_test.py                | 2 +-
 tensor2tensor/models/lstm.py                                  | 2 +-
 tensor2tensor/models/lstm_test.py                             | 2 +-
 tensor2tensor/models/mtf_image_transformer.py                 | 2 +-
 tensor2tensor/models/mtf_image_transformer_test.py            | 2 +-
 tensor2tensor/models/mtf_resnet.py                            | 2 +-
 tensor2tensor/models/mtf_transformer.py                       | 2 +-
 tensor2tensor/models/mtf_transformer2.py                      | 2 +-
 tensor2tensor/models/mtf_transformer_test.py                  | 2 +-
 tensor2tensor/models/neural_architecture_search/__init__.py   | 2 +-
 tensor2tensor/models/neural_architecture_search/nas_layers.py | 2 +-
 .../models/neural_architecture_search/nas_layers_test.py      | 2 +-
 tensor2tensor/models/neural_architecture_search/nas_model.py  | 2 +-
 .../models/neural_architecture_search/nas_model_test.py       | 2 +-
 tensor2tensor/models/neural_assistant.py                      | 2 +-
 tensor2tensor/models/neural_gpu.py                            | 2 +-
 tensor2tensor/models/neural_gpu_test.py                       | 2 +-
 tensor2tensor/models/research/__init__.py                     | 2 +-
 tensor2tensor/models/research/adafactor_experiments.py        | 2 +-
 tensor2tensor/models/research/aligned.py                      | 2 +-
 tensor2tensor/models/research/attention_lm.py                 | 2 +-
 tensor2tensor/models/research/attention_lm_moe.py             | 2 +-
 tensor2tensor/models/research/autoencoders.py                 | 2 +-
 tensor2tensor/models/research/autoencoders_test.py            | 2 +-
 tensor2tensor/models/research/cycle_gan.py                    | 2 +-
 tensor2tensor/models/research/gene_expression.py              | 2 +-
 tensor2tensor/models/research/gene_expression_test.py         | 2 +-
 tensor2tensor/models/research/glow.py                         | 2 +-
 tensor2tensor/models/research/glow_init_hook.py               | 2 +-
 tensor2tensor/models/research/glow_ops.py                     | 2 +-
 tensor2tensor/models/research/glow_ops_test.py                | 2 +-
 tensor2tensor/models/research/glow_test.py                    | 2 +-
 tensor2tensor/models/research/lm_experiments.py               | 2 +-
 tensor2tensor/models/research/moe.py                          | 2 +-
 tensor2tensor/models/research/moe_experiments.py              | 2 +-
 tensor2tensor/models/research/multiquery_paper.py             | 2 +-
 tensor2tensor/models/research/neural_stack.py                 | 2 +-
 tensor2tensor/models/research/neural_stack_test.py            | 2 +-
 tensor2tensor/models/research/rl.py                           | 2 +-
 tensor2tensor/models/research/shuffle_network.py              | 2 +-
 tensor2tensor/models/research/similarity_transformer.py       | 2 +-
 tensor2tensor/models/research/super_lm.py                     | 2 +-
 tensor2tensor/models/research/transformer_aux.py              | 2 +-
 tensor2tensor/models/research/transformer_aux_test.py         | 2 +-
 tensor2tensor/models/research/transformer_moe.py              | 2 +-
 tensor2tensor/models/research/transformer_nat.py              | 2 +-
 tensor2tensor/models/research/transformer_parallel.py         | 2 +-
 tensor2tensor/models/research/transformer_revnet.py           | 2 +-
 tensor2tensor/models/research/transformer_revnet_test.py      | 2 +-
 tensor2tensor/models/research/transformer_sketch.py           | 2 +-
 tensor2tensor/models/research/transformer_symshard.py         | 2 +-
 tensor2tensor/models/research/transformer_vae.py              | 2 +-
 tensor2tensor/models/research/transformer_vae_test.py         | 2 +-
 tensor2tensor/models/research/universal_transformer.py        | 2 +-
 tensor2tensor/models/research/universal_transformer_test.py   | 2 +-
 tensor2tensor/models/research/universal_transformer_util.py   | 2 +-
 tensor2tensor/models/research/vqa_attention.py                | 2 +-
 tensor2tensor/models/research/vqa_attention_test.py           | 2 +-
 tensor2tensor/models/research/vqa_recurrent_self_attention.py | 2 +-
 tensor2tensor/models/research/vqa_self_attention.py           | 2 +-
 tensor2tensor/models/resnet.py                                | 2 +-
 tensor2tensor/models/resnet_test.py                           | 2 +-
 tensor2tensor/models/revnet.py                                | 2 +-
 tensor2tensor/models/revnet_test.py                           | 2 +-
 tensor2tensor/models/shake_shake.py                           | 2 +-
 tensor2tensor/models/slicenet.py                              | 2 +-
 tensor2tensor/models/slicenet_test.py                         | 2 +-
 tensor2tensor/models/text_cnn.py                              | 2 +-
 tensor2tensor/models/transformer.py                           | 2 +-
 tensor2tensor/models/transformer_test.py                      | 2 +-
 tensor2tensor/models/vanilla_gan.py                           | 2 +-
 tensor2tensor/models/video/__init__.py                        | 2 +-
 tensor2tensor/models/video/base.py                            | 2 +-
 tensor2tensor/models/video/base_vae.py                        | 2 +-
 tensor2tensor/models/video/basic_deterministic.py             | 2 +-
 tensor2tensor/models/video/basic_deterministic_params.py      | 2 +-
 tensor2tensor/models/video/basic_deterministic_test.py        | 2 +-
 tensor2tensor/models/video/basic_recurrent.py                 | 2 +-
 tensor2tensor/models/video/basic_recurrent_test.py            | 2 +-
 tensor2tensor/models/video/basic_stochastic.py                | 2 +-
 tensor2tensor/models/video/basic_stochastic_test.py           | 2 +-
 tensor2tensor/models/video/emily.py                           | 2 +-
 tensor2tensor/models/video/emily_test.py                      | 2 +-
 tensor2tensor/models/video/epva.py                            | 2 +-
 tensor2tensor/models/video/epva_params.py                     | 2 +-
 tensor2tensor/models/video/next_frame_glow.py                 | 2 +-
 tensor2tensor/models/video/nfg_conv3d_test.py                 | 2 +-
 tensor2tensor/models/video/nfg_conv_lstm_test.py              | 2 +-
 tensor2tensor/models/video/nfg_conv_test.py                   | 2 +-
 tensor2tensor/models/video/nfg_interpolate.py                 | 2 +-
 tensor2tensor/models/video/nfg_test_utils.py                  | 2 +-
 tensor2tensor/models/video/nfg_uncond_test.py                 | 2 +-
 tensor2tensor/models/video/savp.py                            | 2 +-
 tensor2tensor/models/video/savp_params.py                     | 2 +-
 tensor2tensor/models/video/savp_test.py                       | 2 +-
 tensor2tensor/models/video/sv2p.py                            | 2 +-
 tensor2tensor/models/video/sv2p_params.py                     | 2 +-
 tensor2tensor/models/video/sv2p_test.py                       | 2 +-
 tensor2tensor/models/video/tests_utils.py                     | 2 +-
 tensor2tensor/models/xception.py                              | 2 +-
 tensor2tensor/models/xception_test.py                         | 2 +-
 tensor2tensor/problems.py                                     | 2 +-
 tensor2tensor/problems_colab.py                               | 2 +-
 tensor2tensor/problems_test.py                                | 2 +-
 tensor2tensor/rl/__init__.py                                  | 2 +-
 tensor2tensor/rl/batch_dqn_agent_test.py                      | 2 +-
 tensor2tensor/rl/batch_runner_test.py                         | 2 +-
 tensor2tensor/rl/datagen_with_agent.py                        | 2 +-
 tensor2tensor/rl/dopamine_connector.py                        | 2 +-
 tensor2tensor/rl/envs/__init__.py                             | 2 +-
 tensor2tensor/rl/envs/in_graph_batch_env.py                   | 2 +-
 tensor2tensor/rl/envs/py_func_batch_env.py                    | 2 +-
 tensor2tensor/rl/envs/simulated_batch_env.py                  | 2 +-
 tensor2tensor/rl/envs/simulated_batch_gym_env.py              | 2 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py                    | 2 +-
 tensor2tensor/rl/evaluator.py                                 | 2 +-
 tensor2tensor/rl/evaluator_test.py                            | 2 +-
 tensor2tensor/rl/gym_utils.py                                 | 2 +-
 tensor2tensor/rl/gym_utils_test.py                            | 2 +-
 tensor2tensor/rl/player.py                                    | 2 +-
 tensor2tensor/rl/player_utils.py                              | 2 +-
 tensor2tensor/rl/policy_learner.py                            | 2 +-
 tensor2tensor/rl/ppo.py                                       | 2 +-
 tensor2tensor/rl/ppo_learner.py                               | 2 +-
 tensor2tensor/rl/restarter.py                                 | 2 +-
 tensor2tensor/rl/restarter_test.py                            | 2 +-
 tensor2tensor/rl/rl_utils.py                                  | 2 +-
 tensor2tensor/rl/trainer_model_based.py                       | 2 +-
 tensor2tensor/rl/trainer_model_based_agent_only.py            | 2 +-
 tensor2tensor/rl/trainer_model_based_params.py                | 2 +-
 tensor2tensor/rl/trainer_model_based_recurrent_test.py        | 2 +-
 tensor2tensor/rl/trainer_model_based_stochastic_test.py       | 2 +-
 tensor2tensor/rl/trainer_model_based_sv2p_test.py             | 2 +-
 tensor2tensor/rl/trainer_model_based_test.py                  | 2 +-
 tensor2tensor/rl/trainer_model_free.py                        | 2 +-
 tensor2tensor/rl/trainer_model_free_test.py                   | 2 +-
 tensor2tensor/rl/trainer_model_free_tictactoe_test.py         | 2 +-
 tensor2tensor/serving/__init__.py                             | 2 +-
 tensor2tensor/serving/export.py                               | 2 +-
 tensor2tensor/serving/query.py                                | 2 +-
 tensor2tensor/serving/serving_utils.py                        | 2 +-
 tensor2tensor/test_data/example_usr_dir/__init__.py           | 2 +-
 tensor2tensor/test_data/example_usr_dir/my_submodule.py       | 2 +-
 tensor2tensor/utils/__init__.py                               | 2 +-
 tensor2tensor/utils/adafactor.py                              | 2 +-
 tensor2tensor/utils/adv_attack_utils.py                       | 2 +-
 tensor2tensor/utils/avg_checkpoints.py                        | 2 +-
 tensor2tensor/utils/beam_search.py                            | 2 +-
 tensor2tensor/utils/beam_search_test.py                       | 2 +-
 tensor2tensor/utils/bleu_hook.py                              | 2 +-
 tensor2tensor/utils/bleu_hook_test.py                         | 2 +-
 tensor2tensor/utils/checkpoint_compatibility_test.py          | 2 +-
 tensor2tensor/utils/cloud_mlengine.py                         | 2 +-
 tensor2tensor/utils/compute_video_metrics.py                  | 2 +-
 tensor2tensor/utils/contrib.py                                | 2 +-
 tensor2tensor/utils/data_reader.py                            | 2 +-
 tensor2tensor/utils/data_reader_test.py                       | 2 +-
 tensor2tensor/utils/decoding.py                               | 2 +-
 tensor2tensor/utils/devices.py                                | 2 +-
 tensor2tensor/utils/diet.py                                   | 2 +-
 tensor2tensor/utils/diet_test.py                              | 2 +-
 tensor2tensor/utils/expert_utils.py                           | 2 +-
 tensor2tensor/utils/expert_utils_test.py                      | 2 +-
 tensor2tensor/utils/flags.py                                  | 2 +-
 tensor2tensor/utils/get_rouge.py                              | 2 +-
 tensor2tensor/utils/hparam.py                                 | 2 +-
 tensor2tensor/utils/hparam_test.py                            | 2 +-
 tensor2tensor/utils/hparams_lib.py                            | 2 +-
 tensor2tensor/utils/hparams_lib_test.py                       | 2 +-
 tensor2tensor/utils/learning_rate.py                          | 2 +-
 tensor2tensor/utils/metrics.py                                | 2 +-
 tensor2tensor/utils/metrics_hook.py                           | 2 +-
 tensor2tensor/utils/metrics_hook_test.py                      | 2 +-
 tensor2tensor/utils/metrics_test.py                           | 2 +-
 tensor2tensor/utils/misc_utils.py                             | 2 +-
 tensor2tensor/utils/misc_utils_test.py                        | 2 +-
 tensor2tensor/utils/mlperf_log.py                             | 2 +-
 tensor2tensor/utils/mlperf_tags.py                            | 2 +-
 tensor2tensor/utils/mtf_model.py                              | 2 +-
 tensor2tensor/utils/multistep_optimizer.py                    | 2 +-
 tensor2tensor/utils/multistep_optimizer_test.py               | 2 +-
 tensor2tensor/utils/optimize.py                               | 2 +-
 tensor2tensor/utils/optimize_test.py                          | 2 +-
 tensor2tensor/utils/partial_checkpoint_load_hook.py           | 2 +-
 tensor2tensor/utils/pruning_utils.py                          | 2 +-
 tensor2tensor/utils/quantization.py                           | 2 +-
 tensor2tensor/utils/registry.py                               | 2 +-
 tensor2tensor/utils/registry_test.py                          | 2 +-
 tensor2tensor/utils/restore_hook.py                           | 2 +-
 tensor2tensor/utils/rouge.py                                  | 2 +-
 tensor2tensor/utils/rouge_test.py                             | 2 +-
 tensor2tensor/utils/sari_hook.py                              | 2 +-
 tensor2tensor/utils/sari_hook_test.py                         | 2 +-
 tensor2tensor/utils/scheduled_sampling.py                     | 2 +-
 tensor2tensor/utils/t2t_model.py                              | 2 +-
 tensor2tensor/utils/t2t_model_test.py                         | 2 +-
 tensor2tensor/utils/test_utils.py                             | 2 +-
 tensor2tensor/utils/test_utils_test.py                        | 2 +-
 tensor2tensor/utils/trainer_lib.py                            | 2 +-
 tensor2tensor/utils/trainer_lib_test.py                       | 2 +-
 tensor2tensor/utils/update_ops_hook.py                        | 2 +-
 tensor2tensor/utils/usr_dir.py                                | 2 +-
 tensor2tensor/utils/video/prediction2gif.py                   | 2 +-
 tensor2tensor/utils/video/reward_confusion.py                 | 2 +-
 tensor2tensor/utils/video2gif.py                              | 2 +-
 tensor2tensor/utils/video_metrics.py                          | 2 +-
 tensor2tensor/utils/video_metrics_test.py                     | 2 +-
 tensor2tensor/utils/yellowfin.py                              | 2 +-
 tensor2tensor/utils/yellowfin_test.py                         | 2 +-
 tensor2tensor/visualization/__init__.py                       | 2 +-
 tensor2tensor/visualization/attention.py                      | 2 +-
 tensor2tensor/visualization/visualization.py                  | 2 +-
 tensor2tensor/visualization/visualization_test.py             | 2 +-
 439 files changed, 440 insertions(+), 440 deletions(-)

diff --git a/tensor2tensor/__init__.py b/tensor2tensor/__init__.py
index 4872e5d5d..798e88593 100644
--- a/tensor2tensor/__init__.py
+++ b/tensor2tensor/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/__init__.py b/tensor2tensor/bin/__init__.py
index 4872e5d5d..798e88593 100644
--- a/tensor2tensor/bin/__init__.py
+++ b/tensor2tensor/bin/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/build_vocab.py b/tensor2tensor/bin/build_vocab.py
index 7b5be823d..36515c2b2 100644
--- a/tensor2tensor/bin/build_vocab.py
+++ b/tensor2tensor/bin/build_vocab.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/make_tf_configs.py b/tensor2tensor/bin/make_tf_configs.py
index a27e3e376..3d978d79d 100644
--- a/tensor2tensor/bin/make_tf_configs.py
+++ b/tensor2tensor/bin/make_tf_configs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 6dbcc627e..a674a8104 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_avg_all.py b/tensor2tensor/bin/t2t_avg_all.py
index 4bb9f9949..608a8705b 100644
--- a/tensor2tensor/bin/t2t_avg_all.py
+++ b/tensor2tensor/bin/t2t_avg_all.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_bleu.py b/tensor2tensor/bin/t2t_bleu.py
index 61fcb2438..dd7ef73f8 100644
--- a/tensor2tensor/bin/t2t_bleu.py
+++ b/tensor2tensor/bin/t2t_bleu.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index 9f5acb2e5..1d5d7c104 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 2013d8218..b311011d9 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_distill.py b/tensor2tensor/bin/t2t_distill.py
index 1815d0b10..348868f18 100644
--- a/tensor2tensor/bin/t2t_distill.py
+++ b/tensor2tensor/bin/t2t_distill.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_eval.py b/tensor2tensor/bin/t2t_eval.py
index 4791972c2..24167c1ac 100644
--- a/tensor2tensor/bin/t2t_eval.py
+++ b/tensor2tensor/bin/t2t_eval.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
index 1d2c668fd..d9f0b4939 100644
--- a/tensor2tensor/bin/t2t_prune.py
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index dc4ad34a6..4ca0c1f62 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -233,7 +233,7 @@ def create_run_config(hp, output_dir=None):
   save_ckpt_secs = FLAGS.save_checkpoints_secs or None
   if save_ckpt_secs:
     save_ckpt_steps = None
-  assert FLAGS.output_dir or FLAGS.checkpoint_path
+  assert FLAGS.output_dir
   tpu_config_extra_kwargs = {}
   if FLAGS.tpu_job_name is not None:
     tpu_config_extra_kwargs["tpu_job_name"] = FLAGS.tpu_job_name
diff --git a/tensor2tensor/bin/t2t_trainer_test.py b/tensor2tensor/bin/t2t_trainer_test.py
index 9ed90dc63..c7bfbdfbf 100644
--- a/tensor2tensor/bin/t2t_trainer_test.py
+++ b/tensor2tensor/bin/t2t_trainer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_translate_all.py b/tensor2tensor/bin/t2t_translate_all.py
index 401004fa0..26dad8e7a 100644
--- a/tensor2tensor/bin/t2t_translate_all.py
+++ b/tensor2tensor/bin/t2t_translate_all.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/__init__.py b/tensor2tensor/data_generators/__init__.py
index 4872e5d5d..798e88593 100644
--- a/tensor2tensor/data_generators/__init__.py
+++ b/tensor2tensor/data_generators/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index fe4f8ee0e..74ae810e9 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math.py b/tensor2tensor/data_generators/algorithmic_math.py
index 42767d421..d3e5ee373 100644
--- a/tensor2tensor/data_generators/algorithmic_math.py
+++ b/tensor2tensor/data_generators/algorithmic_math.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math_deepmind.py b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
index 1d849c075..3f8a9c015 100644
--- a/tensor2tensor/data_generators/algorithmic_math_deepmind.py
+++ b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math_test.py b/tensor2tensor/data_generators/algorithmic_math_test.py
index 006a64ded..ff00e4e9b 100644
--- a/tensor2tensor/data_generators/algorithmic_math_test.py
+++ b/tensor2tensor/data_generators/algorithmic_math_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math_two_variables.py b/tensor2tensor/data_generators/algorithmic_math_two_variables.py
index bad119cbf..d6c992685 100644
--- a/tensor2tensor/data_generators/algorithmic_math_two_variables.py
+++ b/tensor2tensor/data_generators/algorithmic_math_two_variables.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_test.py b/tensor2tensor/data_generators/algorithmic_test.py
index ab6448399..cf6bf1b33 100644
--- a/tensor2tensor/data_generators/algorithmic_test.py
+++ b/tensor2tensor/data_generators/algorithmic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index f39addd71..58bb75e7b 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index 056f6a8f2..19951e435 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/allen_brain_test.py b/tensor2tensor/data_generators/allen_brain_test.py
index 487517ab6..c19874d54 100644
--- a/tensor2tensor/data_generators/allen_brain_test.py
+++ b/tensor2tensor/data_generators/allen_brain_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio.py b/tensor2tensor/data_generators/audio.py
index b37a4d6e9..7bff3e84f 100644
--- a/tensor2tensor/data_generators/audio.py
+++ b/tensor2tensor/data_generators/audio.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio_encoder.py b/tensor2tensor/data_generators/audio_encoder.py
index 9ad4494eb..21d60de5e 100644
--- a/tensor2tensor/data_generators/audio_encoder.py
+++ b/tensor2tensor/data_generators/audio_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio_test.py b/tensor2tensor/data_generators/audio_test.py
index 329c8d77b..ff94e4430 100644
--- a/tensor2tensor/data_generators/audio_test.py
+++ b/tensor2tensor/data_generators/audio_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index 0d49655e3..ba1831dbd 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index ac3ae4152..fae556be2 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index c4d5776c5..b9d0848bf 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celeba_test.py b/tensor2tensor/data_generators/celeba_test.py
index deb3472c1..59f932c09 100644
--- a/tensor2tensor/data_generators/celeba_test.py
+++ b/tensor2tensor/data_generators/celeba_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celebahq.py b/tensor2tensor/data_generators/celebahq.py
index d170883ed..55cbf42b6 100644
--- a/tensor2tensor/data_generators/celebahq.py
+++ b/tensor2tensor/data_generators/celebahq.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index e644ed550..181cdb991 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cipher.py b/tensor2tensor/data_generators/cipher.py
index fc8a59022..93cb6b5f9 100644
--- a/tensor2tensor/data_generators/cipher.py
+++ b/tensor2tensor/data_generators/cipher.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cleaner_en_xx.py b/tensor2tensor/data_generators/cleaner_en_xx.py
index 56737701e..1d690d712 100644
--- a/tensor2tensor/data_generators/cleaner_en_xx.py
+++ b/tensor2tensor/data_generators/cleaner_en_xx.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index c792d5b06..da00e1a3e 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cola.py b/tensor2tensor/data_generators/cola.py
index 7b08178b0..c9ea13094 100644
--- a/tensor2tensor/data_generators/cola.py
+++ b/tensor2tensor/data_generators/cola.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/common_voice.py b/tensor2tensor/data_generators/common_voice.py
index cc4867150..a0806659a 100644
--- a/tensor2tensor/data_generators/common_voice.py
+++ b/tensor2tensor/data_generators/common_voice.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/common_voice_test.py b/tensor2tensor/data_generators/common_voice_test.py
index 005c4034a..302b7f929 100644
--- a/tensor2tensor/data_generators/common_voice_test.py
+++ b/tensor2tensor/data_generators/common_voice_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/conll_ner.py b/tensor2tensor/data_generators/conll_ner.py
index 0e2b536b7..fa4e42034 100644
--- a/tensor2tensor/data_generators/conll_ner.py
+++ b/tensor2tensor/data_generators/conll_ner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/desc2code.py b/tensor2tensor/data_generators/desc2code.py
index 1524765d4..6619130c0 100644
--- a/tensor2tensor/data_generators/desc2code.py
+++ b/tensor2tensor/data_generators/desc2code.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/desc2code_test.py b/tensor2tensor/data_generators/desc2code_test.py
index f33016ab9..882db43db 100644
--- a/tensor2tensor/data_generators/desc2code_test.py
+++ b/tensor2tensor/data_generators/desc2code_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_abstract.py b/tensor2tensor/data_generators/dialog_abstract.py
index 9346910e4..4541b8bd7 100644
--- a/tensor2tensor/data_generators/dialog_abstract.py
+++ b/tensor2tensor/data_generators/dialog_abstract.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_cornell.py b/tensor2tensor/data_generators/dialog_cornell.py
index 12363dd36..408ad1c6a 100644
--- a/tensor2tensor/data_generators/dialog_cornell.py
+++ b/tensor2tensor/data_generators/dialog_cornell.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_dailydialog.py b/tensor2tensor/data_generators/dialog_dailydialog.py
index ea16ff3bb..44634f7be 100644
--- a/tensor2tensor/data_generators/dialog_dailydialog.py
+++ b/tensor2tensor/data_generators/dialog_dailydialog.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_opensubtitles.py b/tensor2tensor/data_generators/dialog_opensubtitles.py
index 172bf0bb3..62cdce509 100644
--- a/tensor2tensor/data_generators/dialog_opensubtitles.py
+++ b/tensor2tensor/data_generators/dialog_opensubtitles.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_personachat.py b/tensor2tensor/data_generators/dialog_personachat.py
index ed085badc..c356a86ec 100644
--- a/tensor2tensor/data_generators/dialog_personachat.py
+++ b/tensor2tensor/data_generators/dialog_personachat.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dna_encoder.py b/tensor2tensor/data_generators/dna_encoder.py
index c155dc9e2..c9b0e87ea 100644
--- a/tensor2tensor/data_generators/dna_encoder.py
+++ b/tensor2tensor/data_generators/dna_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dna_encoder_test.py b/tensor2tensor/data_generators/dna_encoder_test.py
index b7739fae7..e354f2ef4 100644
--- a/tensor2tensor/data_generators/dna_encoder_test.py
+++ b/tensor2tensor/data_generators/dna_encoder_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/enwik8.py b/tensor2tensor/data_generators/enwik8.py
index c752230ce..e4da9a8d8 100644
--- a/tensor2tensor/data_generators/enwik8.py
+++ b/tensor2tensor/data_generators/enwik8.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
index 2ae9a2734..5f3dab07e 100644
--- a/tensor2tensor/data_generators/fsns.py
+++ b/tensor2tensor/data_generators/fsns.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
index 8e17e7d32..0cbdccbe0 100644
--- a/tensor2tensor/data_generators/function_docstring.py
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gene_expression.py b/tensor2tensor/data_generators/gene_expression.py
index ab84ff790..105f5d0c7 100644
--- a/tensor2tensor/data_generators/gene_expression.py
+++ b/tensor2tensor/data_generators/gene_expression.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gene_expression_test.py b/tensor2tensor/data_generators/gene_expression_test.py
index c07221a03..bfe819d0c 100644
--- a/tensor2tensor/data_generators/gene_expression_test.py
+++ b/tensor2tensor/data_generators/gene_expression_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 0b532caf9..36f54d47e 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/generator_utils_test.py b/tensor2tensor/data_generators/generator_utils_test.py
index b729d6076..0cb378a1e 100644
--- a/tensor2tensor/data_generators/generator_utils_test.py
+++ b/tensor2tensor/data_generators/generator_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index d4bbe0793..03f7fdf53 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index fa6ad4d17..78ca8186f 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index dd252f394..81a7c8462 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ice_parsing.py b/tensor2tensor/data_generators/ice_parsing.py
index d5c42903a..8ab9fa8d4 100644
--- a/tensor2tensor/data_generators/ice_parsing.py
+++ b/tensor2tensor/data_generators/ice_parsing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_lsun.py b/tensor2tensor/data_generators/image_lsun.py
index 653fab2f1..2a847f1fb 100644
--- a/tensor2tensor/data_generators/image_lsun.py
+++ b/tensor2tensor/data_generators/image_lsun.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index eed55809c..edd8e61a9 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_utils_test.py b/tensor2tensor/data_generators/image_utils_test.py
index 330d88768..b4434ef61 100644
--- a/tensor2tensor/data_generators/image_utils_test.py
+++ b/tensor2tensor/data_generators/image_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 58a98062b..18550a466 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imagenet_test.py b/tensor2tensor/data_generators/imagenet_test.py
index 6a9299589..89220ec4d 100644
--- a/tensor2tensor/data_generators/imagenet_test.py
+++ b/tensor2tensor/data_generators/imagenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imdb.py b/tensor2tensor/data_generators/imdb.py
index 68f5111a3..f7db47a0c 100644
--- a/tensor2tensor/data_generators/imdb.py
+++ b/tensor2tensor/data_generators/imdb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/inspect_tfrecord.py b/tensor2tensor/data_generators/inspect_tfrecord.py
index ce2025ca2..0cc905c1b 100644
--- a/tensor2tensor/data_generators/inspect_tfrecord.py
+++ b/tensor2tensor/data_generators/inspect_tfrecord.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lambada.py b/tensor2tensor/data_generators/lambada.py
index 605e4a1e3..d26311eb5 100644
--- a/tensor2tensor/data_generators/lambada.py
+++ b/tensor2tensor/data_generators/lambada.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index 97c0c6588..bb3a6552c 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b.py b/tensor2tensor/data_generators/lm1b.py
index a98545438..e608d1799 100644
--- a/tensor2tensor/data_generators/lm1b.py
+++ b/tensor2tensor/data_generators/lm1b.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b_imdb.py b/tensor2tensor/data_generators/lm1b_imdb.py
index 49aeac9b9..2aca37a5c 100644
--- a/tensor2tensor/data_generators/lm1b_imdb.py
+++ b/tensor2tensor/data_generators/lm1b_imdb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b_mnli.py b/tensor2tensor/data_generators/lm1b_mnli.py
index 8ab1af68c..6c4a1b865 100644
--- a/tensor2tensor/data_generators/lm1b_mnli.py
+++ b/tensor2tensor/data_generators/lm1b_mnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mnist.py b/tensor2tensor/data_generators/mnist.py
index 67ca5b9e2..0e5100ba5 100644
--- a/tensor2tensor/data_generators/mnist.py
+++ b/tensor2tensor/data_generators/mnist.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/moving_mnist.py b/tensor2tensor/data_generators/moving_mnist.py
index 1dfc82313..257ba7ad8 100644
--- a/tensor2tensor/data_generators/moving_mnist.py
+++ b/tensor2tensor/data_generators/moving_mnist.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mrpc.py b/tensor2tensor/data_generators/mrpc.py
index 2eb37c4df..2c19f3d41 100644
--- a/tensor2tensor/data_generators/mrpc.py
+++ b/tensor2tensor/data_generators/mrpc.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index 28b77c07a..41ab4bdee 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mscoco_test.py b/tensor2tensor/data_generators/mscoco_test.py
index 8443fa914..5c66f390e 100644
--- a/tensor2tensor/data_generators/mscoco_test.py
+++ b/tensor2tensor/data_generators/mscoco_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 7d51dc002..9ac90b874 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem_v2.py b/tensor2tensor/data_generators/multi_problem_v2.py
index 7d5c8f9ef..82e37f7f3 100644
--- a/tensor2tensor/data_generators/multi_problem_v2.py
+++ b/tensor2tensor/data_generators/multi_problem_v2.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem_v2_test.py b/tensor2tensor/data_generators/multi_problem_v2_test.py
index 56b3d67e5..6cbdba815 100644
--- a/tensor2tensor/data_generators/multi_problem_v2_test.py
+++ b/tensor2tensor/data_generators/multi_problem_v2_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index 8dfaab72c..daeb54f49 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ocr.py b/tensor2tensor/data_generators/ocr.py
index 4d682c4fa..ed1b916fe 100644
--- a/tensor2tensor/data_generators/ocr.py
+++ b/tensor2tensor/data_generators/ocr.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
index c0c033635..443dae16d 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py b/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
index 28bb135ac..c3b9273e3 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/paraphrase_ms_coco.py b/tensor2tensor/data_generators/paraphrase_ms_coco.py
index 36fbde85b..a7a30b710 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/paraphrase_ms_coco_test.py b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
index f9dfdf52c..ac14e8b52 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/pointer_generator_word.py b/tensor2tensor/data_generators/pointer_generator_word.py
index 093ea03cd..ce31b551e 100644
--- a/tensor2tensor/data_generators/pointer_generator_word.py
+++ b/tensor2tensor/data_generators/pointer_generator_word.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index bd2bcb7da..a6c43dbce 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem_hparams.py b/tensor2tensor/data_generators/problem_hparams.py
index 41bbb47e0..7727a0cb6 100644
--- a/tensor2tensor/data_generators/problem_hparams.py
+++ b/tensor2tensor/data_generators/problem_hparams.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index f0f7b99f4..820cbbf27 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/program_search.py b/tensor2tensor/data_generators/program_search.py
index 06f6a8d56..b4c0cd229 100644
--- a/tensor2tensor/data_generators/program_search.py
+++ b/tensor2tensor/data_generators/program_search.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/program_search_test.py b/tensor2tensor/data_generators/program_search_test.py
index 14b1d3291..0bd3687c9 100644
--- a/tensor2tensor/data_generators/program_search_test.py
+++ b/tensor2tensor/data_generators/program_search_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ptb.py b/tensor2tensor/data_generators/ptb.py
index 6e9dfd5b8..e02aed4a2 100644
--- a/tensor2tensor/data_generators/ptb.py
+++ b/tensor2tensor/data_generators/ptb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/qnli.py b/tensor2tensor/data_generators/qnli.py
index 9889ee3f6..a6b7f6079 100644
--- a/tensor2tensor/data_generators/qnli.py
+++ b/tensor2tensor/data_generators/qnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/quora_qpairs.py b/tensor2tensor/data_generators/quora_qpairs.py
index 81d6c1a13..a298af4b5 100644
--- a/tensor2tensor/data_generators/quora_qpairs.py
+++ b/tensor2tensor/data_generators/quora_qpairs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/rte.py b/tensor2tensor/data_generators/rte.py
index e39f5c4ac..2dcae1284 100644
--- a/tensor2tensor/data_generators/rte.py
+++ b/tensor2tensor/data_generators/rte.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/scitail.py b/tensor2tensor/data_generators/scitail.py
index e044a3101..a041fec5f 100644
--- a/tensor2tensor/data_generators/scitail.py
+++ b/tensor2tensor/data_generators/scitail.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/snli.py b/tensor2tensor/data_generators/snli.py
index 70b1a2139..010ace7f0 100644
--- a/tensor2tensor/data_generators/snli.py
+++ b/tensor2tensor/data_generators/snli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index b105bd55c..9cda94072 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/squad.py b/tensor2tensor/data_generators/squad.py
index cc3eb5c81..9b9958adc 100644
--- a/tensor2tensor/data_generators/squad.py
+++ b/tensor2tensor/data_generators/squad.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/sst_binary.py b/tensor2tensor/data_generators/sst_binary.py
index 3b61d5ef0..2e6410a3e 100644
--- a/tensor2tensor/data_generators/sst_binary.py
+++ b/tensor2tensor/data_generators/sst_binary.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/stanford_nli.py b/tensor2tensor/data_generators/stanford_nli.py
index b3f038651..a086ed61b 100644
--- a/tensor2tensor/data_generators/stanford_nli.py
+++ b/tensor2tensor/data_generators/stanford_nli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/style_transfer.py b/tensor2tensor/data_generators/style_transfer.py
index 146793666..33cd352c0 100644
--- a/tensor2tensor/data_generators/style_transfer.py
+++ b/tensor2tensor/data_generators/style_transfer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/style_transfer_test.py b/tensor2tensor/data_generators/style_transfer_test.py
index 4f0947556..b212f6125 100644
--- a/tensor2tensor/data_generators/style_transfer_test.py
+++ b/tensor2tensor/data_generators/style_transfer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/subject_verb_agreement.py b/tensor2tensor/data_generators/subject_verb_agreement.py
index 7f2630b59..55c025b6b 100644
--- a/tensor2tensor/data_generators/subject_verb_agreement.py
+++ b/tensor2tensor/data_generators/subject_verb_agreement.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index a9f51f1be..3c17045ca 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_encoder_build_subword.py b/tensor2tensor/data_generators/text_encoder_build_subword.py
index a7c9c37b3..fa3bddfda 100644
--- a/tensor2tensor/data_generators/text_encoder_build_subword.py
+++ b/tensor2tensor/data_generators/text_encoder_build_subword.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_encoder_test.py b/tensor2tensor/data_generators/text_encoder_test.py
index e37203d53..92cc9cecc 100644
--- a/tensor2tensor/data_generators/text_encoder_test.py
+++ b/tensor2tensor/data_generators/text_encoder_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index d868c58b6..6ce5a5d87 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_problems_test.py b/tensor2tensor/data_generators/text_problems_test.py
index 7cf895e6c..28ea85d20 100644
--- a/tensor2tensor/data_generators/text_problems_test.py
+++ b/tensor2tensor/data_generators/text_problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index e0f92fe8b..e84af305c 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_data_generator.py b/tensor2tensor/data_generators/timeseries_data_generator.py
index 94c2c17c3..8d697184c 100644
--- a/tensor2tensor/data_generators/timeseries_data_generator.py
+++ b/tensor2tensor/data_generators/timeseries_data_generator.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_data_generator_test.py b/tensor2tensor/data_generators/timeseries_data_generator_test.py
index 85356bc97..98a386102 100644
--- a/tensor2tensor/data_generators/timeseries_data_generator_test.py
+++ b/tensor2tensor/data_generators/timeseries_data_generator_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index 41c149837..9c8edd98a 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/tokenizer.py b/tensor2tensor/data_generators/tokenizer.py
index f00b5845b..46631a41b 100644
--- a/tensor2tensor/data_generators/tokenizer.py
+++ b/tensor2tensor/data_generators/tokenizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/tokenizer_test.py b/tensor2tensor/data_generators/tokenizer_test.py
index 60ffd0823..c4d93b723 100644
--- a/tensor2tensor/data_generators/tokenizer_test.py
+++ b/tensor2tensor/data_generators/tokenizer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/transduction_problems.py b/tensor2tensor/data_generators/transduction_problems.py
index a3aa3984b..2471dd343 100644
--- a/tensor2tensor/data_generators/transduction_problems.py
+++ b/tensor2tensor/data_generators/transduction_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/transduction_problems_test.py b/tensor2tensor/data_generators/transduction_problems_test.py
index 033ca3922..c3b8f57ca 100644
--- a/tensor2tensor/data_generators/transduction_problems_test.py
+++ b/tensor2tensor/data_generators/transduction_problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 487c6f143..527b0aafa 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_encs.py b/tensor2tensor/data_generators/translate_encs.py
index 7bd0fffc0..95819486d 100644
--- a/tensor2tensor/data_generators/translate_encs.py
+++ b/tensor2tensor/data_generators/translate_encs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index 1247d0814..140def85d 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_ende_test.py b/tensor2tensor/data_generators/translate_ende_test.py
index bd9467af0..74acc0f4d 100644
--- a/tensor2tensor/data_generators/translate_ende_test.py
+++ b/tensor2tensor/data_generators/translate_ende_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enes.py b/tensor2tensor/data_generators/translate_enes.py
index 1619a5adc..3e3fcf203 100644
--- a/tensor2tensor/data_generators/translate_enes.py
+++ b/tensor2tensor/data_generators/translate_enes.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enet.py b/tensor2tensor/data_generators/translate_enet.py
index 77ad70f07..19335beb1 100644
--- a/tensor2tensor/data_generators/translate_enet.py
+++ b/tensor2tensor/data_generators/translate_enet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index 81fd5e4b1..16067adf8 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enid.py b/tensor2tensor/data_generators/translate_enid.py
index 1aa2d2a2b..4367fd0a0 100644
--- a/tensor2tensor/data_generators/translate_enid.py
+++ b/tensor2tensor/data_generators/translate_enid.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enmk.py b/tensor2tensor/data_generators/translate_enmk.py
index 07947bf04..c897e4e11 100644
--- a/tensor2tensor/data_generators/translate_enmk.py
+++ b/tensor2tensor/data_generators/translate_enmk.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enro.py b/tensor2tensor/data_generators/translate_enro.py
index f03e96729..83e4a1557 100644
--- a/tensor2tensor/data_generators/translate_enro.py
+++ b/tensor2tensor/data_generators/translate_enro.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_entn.py b/tensor2tensor/data_generators/translate_entn.py
index ab78df4d1..d9d2de7e2 100644
--- a/tensor2tensor/data_generators/translate_entn.py
+++ b/tensor2tensor/data_generators/translate_entn.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_envi.py b/tensor2tensor/data_generators/translate_envi.py
index 5e333fd55..6ebef9eab 100644
--- a/tensor2tensor/data_generators/translate_envi.py
+++ b/tensor2tensor/data_generators/translate_envi.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index 02acc5381..8d6719975 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_test.py b/tensor2tensor/data_generators/translate_test.py
index 949c5e161..05e2fc8a8 100644
--- a/tensor2tensor/data_generators/translate_test.py
+++ b/tensor2tensor/data_generators/translate_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index c7a5a6afa..3411475e7 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 4ea7e244c..b8c11e6df 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index 0af54dad6..cc3b2087d 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
index 783200f54..d528f4779 100644
--- a/tensor2tensor/data_generators/vqa.py
+++ b/tensor2tensor/data_generators/vqa.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/vqa_utils.py b/tensor2tensor/data_generators/vqa_utils.py
index 74aa8a5b6..90889f93a 100644
--- a/tensor2tensor/data_generators/vqa_utils.py
+++ b/tensor2tensor/data_generators/vqa_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki.py b/tensor2tensor/data_generators/wiki.py
index 17a7f6330..149b31574 100644
--- a/tensor2tensor/data_generators/wiki.py
+++ b/tensor2tensor/data_generators/wiki.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_lm.py b/tensor2tensor/data_generators/wiki_lm.py
index c64fe9fd9..27b67dc65 100644
--- a/tensor2tensor/data_generators/wiki_lm.py
+++ b/tensor2tensor/data_generators/wiki_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index 4c140f139..a3fcd6a06 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_revision.py b/tensor2tensor/data_generators/wiki_revision.py
index 2a925cc67..b3f283b9f 100644
--- a/tensor2tensor/data_generators/wiki_revision.py
+++ b/tensor2tensor/data_generators/wiki_revision.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index eb027203c..73dc3ea13 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/__init__.py b/tensor2tensor/data_generators/wikisum/__init__.py
index 4872e5d5d..798e88593 100644
--- a/tensor2tensor/data_generators/wikisum/__init__.py
+++ b/tensor2tensor/data_generators/wikisum/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/generate_vocab.py b/tensor2tensor/data_generators/wikisum/generate_vocab.py
index 431b2373f..27d03ccd1 100644
--- a/tensor2tensor/data_generators/wikisum/generate_vocab.py
+++ b/tensor2tensor/data_generators/wikisum/generate_vocab.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py b/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
index d8f346662..3c316bd28 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_web.py b/tensor2tensor/data_generators/wikisum/get_references_web.py
index 9c6524b1a..3f16787f8 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py b/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
index 49d04a70a..1a74e76bc 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/html.py b/tensor2tensor/data_generators/wikisum/html.py
index 5a70dcb2d..c3749385d 100644
--- a/tensor2tensor/data_generators/wikisum/html.py
+++ b/tensor2tensor/data_generators/wikisum/html.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/parallel_launch.py b/tensor2tensor/data_generators/wikisum/parallel_launch.py
index 08b3b172b..bfe0836db 100644
--- a/tensor2tensor/data_generators/wikisum/parallel_launch.py
+++ b/tensor2tensor/data_generators/wikisum/parallel_launch.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/produce_examples.py b/tensor2tensor/data_generators/wikisum/produce_examples.py
index a31afe14e..47fa5b1b7 100644
--- a/tensor2tensor/data_generators/wikisum/produce_examples.py
+++ b/tensor2tensor/data_generators/wikisum/produce_examples.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/utils.py b/tensor2tensor/data_generators/wikisum/utils.py
index 541d2c268..f63115d3f 100644
--- a/tensor2tensor/data_generators/wikisum/utils.py
+++ b/tensor2tensor/data_generators/wikisum/utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/utils_test.py b/tensor2tensor/data_generators/wikisum/utils_test.py
index 3bf6ef71b..9fb87f441 100644
--- a/tensor2tensor/data_generators/wikisum/utils_test.py
+++ b/tensor2tensor/data_generators/wikisum/utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/validate_data.py b/tensor2tensor/data_generators/wikisum/validate_data.py
index 8ebf0c4ae..37295a03b 100644
--- a/tensor2tensor/data_generators/wikisum/validate_data.py
+++ b/tensor2tensor/data_generators/wikisum/validate_data.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/wikisum.py b/tensor2tensor/data_generators/wikisum/wikisum.py
index 8be4a59a6..018a9dc75 100644
--- a/tensor2tensor/data_generators/wikisum/wikisum.py
+++ b/tensor2tensor/data_generators/wikisum/wikisum.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikitext103.py b/tensor2tensor/data_generators/wikitext103.py
index 968fc3fea..338f25912 100644
--- a/tensor2tensor/data_generators/wikitext103.py
+++ b/tensor2tensor/data_generators/wikitext103.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wnli.py b/tensor2tensor/data_generators/wnli.py
index c6267403e..741704732 100644
--- a/tensor2tensor/data_generators/wnli.py
+++ b/tensor2tensor/data_generators/wnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wsj_parsing.py b/tensor2tensor/data_generators/wsj_parsing.py
index 228c80656..02d1f64de 100644
--- a/tensor2tensor/data_generators/wsj_parsing.py
+++ b/tensor2tensor/data_generators/wsj_parsing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/yelp_full.py b/tensor2tensor/data_generators/yelp_full.py
index 88e61c840..11d526ac5 100644
--- a/tensor2tensor/data_generators/yelp_full.py
+++ b/tensor2tensor/data_generators/yelp_full.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/yelp_polarity.py b/tensor2tensor/data_generators/yelp_polarity.py
index bf2faa099..1111e1294 100644
--- a/tensor2tensor/data_generators/yelp_polarity.py
+++ b/tensor2tensor/data_generators/yelp_polarity.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/__init__.py b/tensor2tensor/envs/__init__.py
index 8c036a176..01ea5565b 100644
--- a/tensor2tensor/envs/__init__.py
+++ b/tensor2tensor/envs/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index f85fa1f48..68caada1f 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 2520cc74e..180716c85 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index 07324f30e..c94e1623e 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_env_problem.py b/tensor2tensor/envs/gym_env_problem.py
index eb4f3b619..853f7eef8 100644
--- a/tensor2tensor/envs/gym_env_problem.py
+++ b/tensor2tensor/envs/gym_env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_env_problem_test.py b/tensor2tensor/envs/gym_env_problem_test.py
index ea4b94167..c62705a2e 100644
--- a/tensor2tensor/envs/gym_env_problem_test.py
+++ b/tensor2tensor/envs/gym_env_problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_spaces_utils.py b/tensor2tensor/envs/gym_spaces_utils.py
index 66fb3825a..30dc239e8 100644
--- a/tensor2tensor/envs/gym_spaces_utils.py
+++ b/tensor2tensor/envs/gym_spaces_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_spaces_utils_test.py b/tensor2tensor/envs/gym_spaces_utils_test.py
index 26cd6f574..1b637e185 100644
--- a/tensor2tensor/envs/gym_spaces_utils_test.py
+++ b/tensor2tensor/envs/gym_spaces_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/mujoco_problems.py b/tensor2tensor/envs/mujoco_problems.py
index 928b392e3..2eedeca0b 100644
--- a/tensor2tensor/envs/mujoco_problems.py
+++ b/tensor2tensor/envs/mujoco_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/mujoco_problems_test.py b/tensor2tensor/envs/mujoco_problems_test.py
index 4bad4590f..d2ddc4d1a 100644
--- a/tensor2tensor/envs/mujoco_problems_test.py
+++ b/tensor2tensor/envs/mujoco_problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
index 34d03a357..f0344905c 100644
--- a/tensor2tensor/envs/rendered_env_problem.py
+++ b/tensor2tensor/envs/rendered_env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/rendered_env_problem_test.py b/tensor2tensor/envs/rendered_env_problem_test.py
index d1ec67cc1..bc1c1e745 100644
--- a/tensor2tensor/envs/rendered_env_problem_test.py
+++ b/tensor2tensor/envs/rendered_env_problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env.py b/tensor2tensor/envs/tic_tac_toe_env.py
index 4a6b1a1e7..44fa7f29c 100644
--- a/tensor2tensor/envs/tic_tac_toe_env.py
+++ b/tensor2tensor/envs/tic_tac_toe_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_problem.py b/tensor2tensor/envs/tic_tac_toe_env_problem.py
index c609245f8..a92c2c7a1 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_problem_test.py b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
index 89517e8ba..70cc91a7c 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_test.py b/tensor2tensor/envs/tic_tac_toe_env_test.py
index c00b53683..289916210 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/time_step.py b/tensor2tensor/envs/time_step.py
index 7f6bcc233..b75603ef2 100644
--- a/tensor2tensor/envs/time_step.py
+++ b/tensor2tensor/envs/time_step.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/time_step_test.py b/tensor2tensor/envs/time_step_test.py
index ec75c32a7..80dfcdeec 100644
--- a/tensor2tensor/envs/time_step_test.py
+++ b/tensor2tensor/envs/time_step_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index 5eec464da..05a51b715 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index e98d33a22..148d29de4 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/__init__.py b/tensor2tensor/insights/__init__.py
index 4872e5d5d..798e88593 100644
--- a/tensor2tensor/insights/__init__.py
+++ b/tensor2tensor/insights/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/graph.py b/tensor2tensor/insights/graph.py
index 2ca6bd91c..6e13c1eb6 100644
--- a/tensor2tensor/insights/graph.py
+++ b/tensor2tensor/insights/graph.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/query_processor.py b/tensor2tensor/insights/query_processor.py
index 3456b0725..47aefdace 100644
--- a/tensor2tensor/insights/query_processor.py
+++ b/tensor2tensor/insights/query_processor.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/server.py b/tensor2tensor/insights/server.py
index 122064f97..fea2b2dbb 100644
--- a/tensor2tensor/insights/server.py
+++ b/tensor2tensor/insights/server.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/transformer_model.py b/tensor2tensor/insights/transformer_model.py
index 8cad91251..660a6c8d3 100644
--- a/tensor2tensor/insights/transformer_model.py
+++ b/tensor2tensor/insights/transformer_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/__init__.py b/tensor2tensor/layers/__init__.py
index 4872e5d5d..798e88593 100644
--- a/tensor2tensor/layers/__init__.py
+++ b/tensor2tensor/layers/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/area_attention.py b/tensor2tensor/layers/area_attention.py
index 7e1338fbc..70b92d0fa 100644
--- a/tensor2tensor/layers/area_attention.py
+++ b/tensor2tensor/layers/area_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/area_attention_test.py b/tensor2tensor/layers/area_attention_test.py
index 2a2ebc9b4..e8ffcc9ef 100644
--- a/tensor2tensor/layers/area_attention_test.py
+++ b/tensor2tensor/layers/area_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index bb90f1e9d..f2c8b1cf9 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index aab56b5bc..d18436904 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_audio.py b/tensor2tensor/layers/common_audio.py
index 5cb114495..c4efd96c8 100644
--- a/tensor2tensor/layers/common_audio.py
+++ b/tensor2tensor/layers/common_audio.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 7aec40c15..3049d7fa2 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index a64f71e50..f93cce9ea 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index c9e140e5a..80dda5c8d 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 4a4496d50..992b48f23 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index c7e7c7a2f..19ec82355 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index eb05ebb3c..45a2ccc20 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index 90ce75b7d..e04a7305e 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 6bcdd7db5..73796250e 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index d0454e2e1..aa70fe620 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index f23ba00a7..94c005c3a 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
index 59e6ca506..0f7efa145 100644
--- a/tensor2tensor/layers/latent_layers_test.py
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/message_passing_attention.py b/tensor2tensor/layers/message_passing_attention.py
index 887d88b78..6b444d0e3 100644
--- a/tensor2tensor/layers/message_passing_attention.py
+++ b/tensor2tensor/layers/message_passing_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 36a5afefa..2fdd24eb1 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index adbb86414..de597e1d3 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/ngram.py b/tensor2tensor/layers/ngram.py
index 8372f85c8..ad0bd7ec3 100644
--- a/tensor2tensor/layers/ngram.py
+++ b/tensor2tensor/layers/ngram.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/ngram_test.py b/tensor2tensor/layers/ngram_test.py
index 958788f23..608884556 100644
--- a/tensor2tensor/layers/ngram_test.py
+++ b/tensor2tensor/layers/ngram_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index e0e6a7f54..826220180 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index 63e2b1d66..85ab70dbb 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_memory_test.py b/tensor2tensor/layers/transformer_memory_test.py
index c98c9790c..f990863a0 100644
--- a/tensor2tensor/layers/transformer_memory_test.py
+++ b/tensor2tensor/layers/transformer_memory_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index 7799469dd..be455994e 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/vqa_layers.py b/tensor2tensor/layers/vqa_layers.py
index 341fbf463..9d82822b4 100644
--- a/tensor2tensor/layers/vqa_layers.py
+++ b/tensor2tensor/layers/vqa_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/__init__.py b/tensor2tensor/metrics/__init__.py
index 4872e5d5d..798e88593 100644
--- a/tensor2tensor/metrics/__init__.py
+++ b/tensor2tensor/metrics/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/video_conditional_fvd.py b/tensor2tensor/metrics/video_conditional_fvd.py
index d97b946d2..240d87b2a 100644
--- a/tensor2tensor/metrics/video_conditional_fvd.py
+++ b/tensor2tensor/metrics/video_conditional_fvd.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/video_conditional_fvd_test.py b/tensor2tensor/metrics/video_conditional_fvd_test.py
index 2d9b847f8..67b40df5a 100644
--- a/tensor2tensor/metrics/video_conditional_fvd_test.py
+++ b/tensor2tensor/metrics/video_conditional_fvd_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index bf4a2e3d7..e64b42d2d 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/basic.py b/tensor2tensor/models/basic.py
index 3840e5a59..d79008ee2 100644
--- a/tensor2tensor/models/basic.py
+++ b/tensor2tensor/models/basic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/basic_test.py b/tensor2tensor/models/basic_test.py
index f3b29cad3..a59880d13 100644
--- a/tensor2tensor/models/basic_test.py
+++ b/tensor2tensor/models/basic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/bytenet.py b/tensor2tensor/models/bytenet.py
index 23f9c0962..ae3bef291 100644
--- a/tensor2tensor/models/bytenet.py
+++ b/tensor2tensor/models/bytenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/bytenet_test.py b/tensor2tensor/models/bytenet_test.py
index e65d2022b..c43b5e73d 100644
--- a/tensor2tensor/models/bytenet_test.py
+++ b/tensor2tensor/models/bytenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/distillation.py b/tensor2tensor/models/distillation.py
index 4d90166bc..2e0e3fb56 100644
--- a/tensor2tensor/models/distillation.py
+++ b/tensor2tensor/models/distillation.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 87253ce61..8c993a537 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/evolved_transformer_test.py b/tensor2tensor/models/evolved_transformer_test.py
index 3e6db316d..364280deb 100644
--- a/tensor2tensor/models/evolved_transformer_test.py
+++ b/tensor2tensor/models/evolved_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 9bd066e67..fd1affec5 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 783a44ce0..443ed6d51 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_2d_test.py b/tensor2tensor/models/image_transformer_2d_test.py
index 9499a24b0..9c7b43725 100644
--- a/tensor2tensor/models/image_transformer_2d_test.py
+++ b/tensor2tensor/models/image_transformer_2d_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_test.py b/tensor2tensor/models/image_transformer_test.py
index 87485b9e5..eb88edca9 100644
--- a/tensor2tensor/models/image_transformer_test.py
+++ b/tensor2tensor/models/image_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 246a5be54..60d362d39 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/lstm_test.py b/tensor2tensor/models/lstm_test.py
index a4dc22d0d..b49c7465a 100644
--- a/tensor2tensor/models/lstm_test.py
+++ b/tensor2tensor/models/lstm_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_image_transformer.py b/tensor2tensor/models/mtf_image_transformer.py
index 1a71f673c..ecd8152cf 100644
--- a/tensor2tensor/models/mtf_image_transformer.py
+++ b/tensor2tensor/models/mtf_image_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_image_transformer_test.py b/tensor2tensor/models/mtf_image_transformer_test.py
index 81fca2c9d..58711c2fd 100644
--- a/tensor2tensor/models/mtf_image_transformer_test.py
+++ b/tensor2tensor/models/mtf_image_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_resnet.py b/tensor2tensor/models/mtf_resnet.py
index 2c2baf98f..492cf4b28 100644
--- a/tensor2tensor/models/mtf_resnet.py
+++ b/tensor2tensor/models/mtf_resnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 526565ec1..7dd32d36e 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index a4262a960..ae075601f 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer_test.py b/tensor2tensor/models/mtf_transformer_test.py
index c7153be72..f097353f5 100644
--- a/tensor2tensor/models/mtf_transformer_test.py
+++ b/tensor2tensor/models/mtf_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/__init__.py b/tensor2tensor/models/neural_architecture_search/__init__.py
index b775a72bd..7688494ed 100644
--- a/tensor2tensor/models/neural_architecture_search/__init__.py
+++ b/tensor2tensor/models/neural_architecture_search/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_layers.py b/tensor2tensor/models/neural_architecture_search/nas_layers.py
index f86b8e742..fcdfca5f4 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_layers.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_layers_test.py b/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
index 657b09a31..ae936c021 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_model.py b/tensor2tensor/models/neural_architecture_search/nas_model.py
index dd2073004..4c72c6b84 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_model_test.py b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
index 86d706e9b..da89868ab 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model_test.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_assistant.py b/tensor2tensor/models/neural_assistant.py
index 6cc357739..98f20d88a 100644
--- a/tensor2tensor/models/neural_assistant.py
+++ b/tensor2tensor/models/neural_assistant.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_gpu.py b/tensor2tensor/models/neural_gpu.py
index 1fb255455..48bff0067 100644
--- a/tensor2tensor/models/neural_gpu.py
+++ b/tensor2tensor/models/neural_gpu.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_gpu_test.py b/tensor2tensor/models/neural_gpu_test.py
index e8961f287..c650d1c08 100644
--- a/tensor2tensor/models/neural_gpu_test.py
+++ b/tensor2tensor/models/neural_gpu_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/__init__.py b/tensor2tensor/models/research/__init__.py
index 4872e5d5d..798e88593 100644
--- a/tensor2tensor/models/research/__init__.py
+++ b/tensor2tensor/models/research/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/adafactor_experiments.py b/tensor2tensor/models/research/adafactor_experiments.py
index 9fe1c38e9..d1de970dc 100644
--- a/tensor2tensor/models/research/adafactor_experiments.py
+++ b/tensor2tensor/models/research/adafactor_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/aligned.py b/tensor2tensor/models/research/aligned.py
index 625270ffd..1d7ac1413 100644
--- a/tensor2tensor/models/research/aligned.py
+++ b/tensor2tensor/models/research/aligned.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/attention_lm.py b/tensor2tensor/models/research/attention_lm.py
index b44afdaea..cf876b4ba 100644
--- a/tensor2tensor/models/research/attention_lm.py
+++ b/tensor2tensor/models/research/attention_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/attention_lm_moe.py b/tensor2tensor/models/research/attention_lm_moe.py
index 04a9cfa91..5e2149bae 100644
--- a/tensor2tensor/models/research/attention_lm_moe.py
+++ b/tensor2tensor/models/research/attention_lm_moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 1908548da..1233f10ec 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index 54a13cfe6..07133da8f 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index 78f2260c3..ad0aaac8a 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/gene_expression.py b/tensor2tensor/models/research/gene_expression.py
index dbf843ebf..3173140ca 100644
--- a/tensor2tensor/models/research/gene_expression.py
+++ b/tensor2tensor/models/research/gene_expression.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/gene_expression_test.py b/tensor2tensor/models/research/gene_expression_test.py
index b58883e36..4c36af768 100644
--- a/tensor2tensor/models/research/gene_expression_test.py
+++ b/tensor2tensor/models/research/gene_expression_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index f795e4ef2..f744e4c39 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_init_hook.py b/tensor2tensor/models/research/glow_init_hook.py
index 51a679d9d..dc87ea94c 100644
--- a/tensor2tensor/models/research/glow_init_hook.py
+++ b/tensor2tensor/models/research/glow_init_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index ee21fa5e0..e09159881 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 6fb2d4e9c..f6a39e6e2 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
index 67f6309f9..2063be0cb 100644
--- a/tensor2tensor/models/research/glow_test.py
+++ b/tensor2tensor/models/research/glow_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/lm_experiments.py b/tensor2tensor/models/research/lm_experiments.py
index cf9f2a75f..9a465fae6 100644
--- a/tensor2tensor/models/research/lm_experiments.py
+++ b/tensor2tensor/models/research/lm_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/moe.py b/tensor2tensor/models/research/moe.py
index 370d2e2eb..9b4d448cb 100644
--- a/tensor2tensor/models/research/moe.py
+++ b/tensor2tensor/models/research/moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/moe_experiments.py b/tensor2tensor/models/research/moe_experiments.py
index a817a6314..503c23527 100644
--- a/tensor2tensor/models/research/moe_experiments.py
+++ b/tensor2tensor/models/research/moe_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/multiquery_paper.py b/tensor2tensor/models/research/multiquery_paper.py
index 4fdc07b6f..7c54a108a 100644
--- a/tensor2tensor/models/research/multiquery_paper.py
+++ b/tensor2tensor/models/research/multiquery_paper.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/neural_stack.py b/tensor2tensor/models/research/neural_stack.py
index 9b6942b68..51d27e6c1 100644
--- a/tensor2tensor/models/research/neural_stack.py
+++ b/tensor2tensor/models/research/neural_stack.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/neural_stack_test.py b/tensor2tensor/models/research/neural_stack_test.py
index 132315f5f..e67070bd6 100644
--- a/tensor2tensor/models/research/neural_stack_test.py
+++ b/tensor2tensor/models/research/neural_stack_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index a6f1b5c3a..7c0471ec5 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/shuffle_network.py b/tensor2tensor/models/research/shuffle_network.py
index ca0f0d156..201cb198b 100644
--- a/tensor2tensor/models/research/shuffle_network.py
+++ b/tensor2tensor/models/research/shuffle_network.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/similarity_transformer.py b/tensor2tensor/models/research/similarity_transformer.py
index 9b5d63f18..597a1ebb6 100644
--- a/tensor2tensor/models/research/similarity_transformer.py
+++ b/tensor2tensor/models/research/similarity_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/super_lm.py b/tensor2tensor/models/research/super_lm.py
index 9bf69ebed..b3307306b 100644
--- a/tensor2tensor/models/research/super_lm.py
+++ b/tensor2tensor/models/research/super_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_aux.py b/tensor2tensor/models/research/transformer_aux.py
index 097cf6a1a..1dc1e3137 100644
--- a/tensor2tensor/models/research/transformer_aux.py
+++ b/tensor2tensor/models/research/transformer_aux.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_aux_test.py b/tensor2tensor/models/research/transformer_aux_test.py
index a748f7ffb..a172fba91 100644
--- a/tensor2tensor/models/research/transformer_aux_test.py
+++ b/tensor2tensor/models/research/transformer_aux_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index 418ddddb2..57999eca3 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index fddcaacc9..638335681 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_parallel.py b/tensor2tensor/models/research/transformer_parallel.py
index 6bca08dfe..ebbe72295 100644
--- a/tensor2tensor/models/research/transformer_parallel.py
+++ b/tensor2tensor/models/research/transformer_parallel.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_revnet.py b/tensor2tensor/models/research/transformer_revnet.py
index 7c3ba2c14..1a50ed3aa 100644
--- a/tensor2tensor/models/research/transformer_revnet.py
+++ b/tensor2tensor/models/research/transformer_revnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_revnet_test.py b/tensor2tensor/models/research/transformer_revnet_test.py
index da0ddf258..8466e6d68 100644
--- a/tensor2tensor/models/research/transformer_revnet_test.py
+++ b/tensor2tensor/models/research/transformer_revnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_sketch.py b/tensor2tensor/models/research/transformer_sketch.py
index 0677fe1b5..040683401 100644
--- a/tensor2tensor/models/research/transformer_sketch.py
+++ b/tensor2tensor/models/research/transformer_sketch.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_symshard.py b/tensor2tensor/models/research/transformer_symshard.py
index 1e4e1bf17..e408e1f26 100644
--- a/tensor2tensor/models/research/transformer_symshard.py
+++ b/tensor2tensor/models/research/transformer_symshard.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index a50fe1313..f566b3572 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae_test.py b/tensor2tensor/models/research/transformer_vae_test.py
index 18c639f1a..bd57d4e5e 100644
--- a/tensor2tensor/models/research/transformer_vae_test.py
+++ b/tensor2tensor/models/research/transformer_vae_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 878828921..1352c3b4d 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
index 1b9d9f4cf..1b79a370a 100644
--- a/tensor2tensor/models/research/universal_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 5e1c4cc69..e58bfed5f 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index c09875c7f..a6f428a88 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_attention_test.py b/tensor2tensor/models/research/vqa_attention_test.py
index 86922f7d0..101bf9891 100644
--- a/tensor2tensor/models/research/vqa_attention_test.py
+++ b/tensor2tensor/models/research/vqa_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_recurrent_self_attention.py b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
index ca6976e64..8c2a4b515 100644
--- a/tensor2tensor/models/research/vqa_recurrent_self_attention.py
+++ b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
index efe08ff38..b39b88df5 100644
--- a/tensor2tensor/models/research/vqa_self_attention.py
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index 1560c13fe..da38fb313 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/resnet_test.py b/tensor2tensor/models/resnet_test.py
index cd2481d7c..7c6e4e0d9 100644
--- a/tensor2tensor/models/resnet_test.py
+++ b/tensor2tensor/models/resnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/revnet.py b/tensor2tensor/models/revnet.py
index 193cbb3b2..446a6da08 100644
--- a/tensor2tensor/models/revnet.py
+++ b/tensor2tensor/models/revnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/revnet_test.py b/tensor2tensor/models/revnet_test.py
index 7bad272ff..ace5a5cfc 100644
--- a/tensor2tensor/models/revnet_test.py
+++ b/tensor2tensor/models/revnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index ba87619c2..d35a6af94 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/slicenet.py b/tensor2tensor/models/slicenet.py
index 75e73243d..4cc795eb8 100644
--- a/tensor2tensor/models/slicenet.py
+++ b/tensor2tensor/models/slicenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/slicenet_test.py b/tensor2tensor/models/slicenet_test.py
index 1e20fc00a..4de4be0bf 100644
--- a/tensor2tensor/models/slicenet_test.py
+++ b/tensor2tensor/models/slicenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/text_cnn.py b/tensor2tensor/models/text_cnn.py
index 7b01565d1..bc8b7b504 100644
--- a/tensor2tensor/models/text_cnn.py
+++ b/tensor2tensor/models/text_cnn.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 927a4e9bf..051b74859 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 50e613d4e..1d6a344be 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index 79fc4f37d..59c5f0ba0 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/__init__.py b/tensor2tensor/models/video/__init__.py
index 4872e5d5d..798e88593 100644
--- a/tensor2tensor/models/video/__init__.py
+++ b/tensor2tensor/models/video/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 55ca3dcdd..f6565bd6a 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/base_vae.py b/tensor2tensor/models/video/base_vae.py
index 069397034..345810823 100644
--- a/tensor2tensor/models/video/base_vae.py
+++ b/tensor2tensor/models/video/base_vae.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index cfd592ae9..ef2b8d8ca 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index cebe148c9..de5893da0 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic_test.py b/tensor2tensor/models/video/basic_deterministic_test.py
index 9f9d3a216..90df4cc81 100644
--- a/tensor2tensor/models/video/basic_deterministic_test.py
+++ b/tensor2tensor/models/video/basic_deterministic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_recurrent.py b/tensor2tensor/models/video/basic_recurrent.py
index e2af2747a..8b982811e 100644
--- a/tensor2tensor/models/video/basic_recurrent.py
+++ b/tensor2tensor/models/video/basic_recurrent.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_recurrent_test.py b/tensor2tensor/models/video/basic_recurrent_test.py
index f19f3ac7a..56aeea4de 100644
--- a/tensor2tensor/models/video/basic_recurrent_test.py
+++ b/tensor2tensor/models/video/basic_recurrent_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index e46b142bf..0ebf67ae7 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_stochastic_test.py b/tensor2tensor/models/video/basic_stochastic_test.py
index c0e27b76d..dff04fb1f 100644
--- a/tensor2tensor/models/video/basic_stochastic_test.py
+++ b/tensor2tensor/models/video/basic_stochastic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index 58b990a70..0a0c58790 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/emily_test.py b/tensor2tensor/models/video/emily_test.py
index 6c35d5951..c9660e1b2 100644
--- a/tensor2tensor/models/video/emily_test.py
+++ b/tensor2tensor/models/video/emily_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index 735002fff..c6f94f0a4 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
index e42c94d9d..7dd9e6957 100644
--- a/tensor2tensor/models/video/epva_params.py
+++ b/tensor2tensor/models/video/epva_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index e37044420..73cadd3c3 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv3d_test.py b/tensor2tensor/models/video/nfg_conv3d_test.py
index 9050ea955..0230d454b 100644
--- a/tensor2tensor/models/video/nfg_conv3d_test.py
+++ b/tensor2tensor/models/video/nfg_conv3d_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv_lstm_test.py b/tensor2tensor/models/video/nfg_conv_lstm_test.py
index 607b445e3..5abda7fe8 100644
--- a/tensor2tensor/models/video/nfg_conv_lstm_test.py
+++ b/tensor2tensor/models/video/nfg_conv_lstm_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv_test.py b/tensor2tensor/models/video/nfg_conv_test.py
index b155967db..a6bab8baa 100644
--- a/tensor2tensor/models/video/nfg_conv_test.py
+++ b/tensor2tensor/models/video/nfg_conv_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_interpolate.py b/tensor2tensor/models/video/nfg_interpolate.py
index 535a6c61c..761317f57 100644
--- a/tensor2tensor/models/video/nfg_interpolate.py
+++ b/tensor2tensor/models/video/nfg_interpolate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_test_utils.py b/tensor2tensor/models/video/nfg_test_utils.py
index e135dd6ea..09a2057af 100644
--- a/tensor2tensor/models/video/nfg_test_utils.py
+++ b/tensor2tensor/models/video/nfg_test_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_uncond_test.py b/tensor2tensor/models/video/nfg_uncond_test.py
index a533a995c..b0cb855cc 100644
--- a/tensor2tensor/models/video/nfg_uncond_test.py
+++ b/tensor2tensor/models/video/nfg_uncond_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index 186378bb7..8b06270d4 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index c28c4378a..2c884a6ea 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp_test.py b/tensor2tensor/models/video/savp_test.py
index 49d14d5a8..4294c0f8a 100644
--- a/tensor2tensor/models/video/savp_test.py
+++ b/tensor2tensor/models/video/savp_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 2aacc8da7..10309be08 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index f7352c714..33cf08fbe 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p_test.py b/tensor2tensor/models/video/sv2p_test.py
index 39815b5b1..a43cf3715 100644
--- a/tensor2tensor/models/video/sv2p_test.py
+++ b/tensor2tensor/models/video/sv2p_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/tests_utils.py b/tensor2tensor/models/video/tests_utils.py
index e9d2d843c..ff834522a 100644
--- a/tensor2tensor/models/video/tests_utils.py
+++ b/tensor2tensor/models/video/tests_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/xception.py b/tensor2tensor/models/xception.py
index d89fad897..99828ca7c 100644
--- a/tensor2tensor/models/xception.py
+++ b/tensor2tensor/models/xception.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/xception_test.py b/tensor2tensor/models/xception_test.py
index 0ac84d5d0..4eb0f8ba0 100644
--- a/tensor2tensor/models/xception_test.py
+++ b/tensor2tensor/models/xception_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/problems.py b/tensor2tensor/problems.py
index b6f6f8bdd..e83b2a1a4 100644
--- a/tensor2tensor/problems.py
+++ b/tensor2tensor/problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/problems_colab.py b/tensor2tensor/problems_colab.py
index 12d679659..4ab203c67 100644
--- a/tensor2tensor/problems_colab.py
+++ b/tensor2tensor/problems_colab.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/problems_test.py b/tensor2tensor/problems_test.py
index bab90968a..d9b6143f7 100644
--- a/tensor2tensor/problems_test.py
+++ b/tensor2tensor/problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/__init__.py b/tensor2tensor/rl/__init__.py
index 4872e5d5d..798e88593 100644
--- a/tensor2tensor/rl/__init__.py
+++ b/tensor2tensor/rl/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/batch_dqn_agent_test.py b/tensor2tensor/rl/batch_dqn_agent_test.py
index 7e2821788..9790e31a3 100644
--- a/tensor2tensor/rl/batch_dqn_agent_test.py
+++ b/tensor2tensor/rl/batch_dqn_agent_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/batch_runner_test.py b/tensor2tensor/rl/batch_runner_test.py
index a260d9d97..cfb415ff1 100644
--- a/tensor2tensor/rl/batch_runner_test.py
+++ b/tensor2tensor/rl/batch_runner_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/datagen_with_agent.py b/tensor2tensor/rl/datagen_with_agent.py
index 43c52186f..2b69fcf55 100644
--- a/tensor2tensor/rl/datagen_with_agent.py
+++ b/tensor2tensor/rl/datagen_with_agent.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 136d837c1..3b9ac95b3 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/__init__.py b/tensor2tensor/rl/envs/__init__.py
index 4872e5d5d..798e88593 100644
--- a/tensor2tensor/rl/envs/__init__.py
+++ b/tensor2tensor/rl/envs/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index f8ba6967c..39807fd82 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index ca455b4ac..8533b8545 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index bca298bcd..9688bad94 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index 6a69a045f..3a6ac4d90 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 03567cad4..b79933b3a 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 3db1ec411..a1df89c7c 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/evaluator_test.py b/tensor2tensor/rl/evaluator_test.py
index e8adeeb30..b6ea5a0db 100644
--- a/tensor2tensor/rl/evaluator_test.py
+++ b/tensor2tensor/rl/evaluator_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index f0fc2e219..2737242f5 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index b347ba906..6755ac6dd 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 5a66d3d35..3304f20ab 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index d53066091..630bb83c5 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
index ec59e4c8a..4bf3fb015 100644
--- a/tensor2tensor/rl/policy_learner.py
+++ b/tensor2tensor/rl/policy_learner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index a9f476ca0..e0f0bed50 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 04c053098..6efd381ea 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/restarter.py b/tensor2tensor/rl/restarter.py
index 72e8628e1..9fad2536c 100644
--- a/tensor2tensor/rl/restarter.py
+++ b/tensor2tensor/rl/restarter.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/restarter_test.py b/tensor2tensor/rl/restarter_test.py
index d28da52b7..a447bf6ce 100644
--- a/tensor2tensor/rl/restarter_test.py
+++ b/tensor2tensor/rl/restarter_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index aebabf2a1..b34ba9ba1 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 5a57f875a..24d7fba3e 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_agent_only.py b/tensor2tensor/rl/trainer_model_based_agent_only.py
index e221106a7..da3b17e57 100644
--- a/tensor2tensor/rl/trainer_model_based_agent_only.py
+++ b/tensor2tensor/rl/trainer_model_based_agent_only.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 32553bb10..4b3bd06bc 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_recurrent_test.py b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
index beda282a6..5642bcc56 100644
--- a/tensor2tensor/rl/trainer_model_based_recurrent_test.py
+++ b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_stochastic_test.py b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
index b0598a320..732ef7327 100644
--- a/tensor2tensor/rl/trainer_model_based_stochastic_test.py
+++ b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_sv2p_test.py b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
index 905b2cfc8..0df918c74 100644
--- a/tensor2tensor/rl/trainer_model_based_sv2p_test.py
+++ b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_test.py b/tensor2tensor/rl/trainer_model_based_test.py
index 90e6ad260..39caf4d26 100644
--- a/tensor2tensor/rl/trainer_model_based_test.py
+++ b/tensor2tensor/rl/trainer_model_based_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index a6cbafb70..4f56dc3ba 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index 3e21a1387..ad2a15ab8 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free_tictactoe_test.py b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
index f57e46546..b03195165 100644
--- a/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
+++ b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/__init__.py b/tensor2tensor/serving/__init__.py
index 4872e5d5d..798e88593 100644
--- a/tensor2tensor/serving/__init__.py
+++ b/tensor2tensor/serving/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 2737f6961..cc993b83d 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/query.py b/tensor2tensor/serving/query.py
index 424cd5581..b266b8cf2 100644
--- a/tensor2tensor/serving/query.py
+++ b/tensor2tensor/serving/query.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 9facc7eeb..8d5e327e6 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/test_data/example_usr_dir/__init__.py b/tensor2tensor/test_data/example_usr_dir/__init__.py
index ed4c6aca1..b47898717 100644
--- a/tensor2tensor/test_data/example_usr_dir/__init__.py
+++ b/tensor2tensor/test_data/example_usr_dir/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/test_data/example_usr_dir/my_submodule.py b/tensor2tensor/test_data/example_usr_dir/my_submodule.py
index 1ba2b439a..593c9dc38 100644
--- a/tensor2tensor/test_data/example_usr_dir/my_submodule.py
+++ b/tensor2tensor/test_data/example_usr_dir/my_submodule.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/__init__.py b/tensor2tensor/utils/__init__.py
index 4872e5d5d..798e88593 100644
--- a/tensor2tensor/utils/__init__.py
+++ b/tensor2tensor/utils/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index 87617ecd9..0d6f41176 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/adv_attack_utils.py b/tensor2tensor/utils/adv_attack_utils.py
index ddad56f59..c0aee2b8d 100644
--- a/tensor2tensor/utils/adv_attack_utils.py
+++ b/tensor2tensor/utils/adv_attack_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index 453ee7cc0..146fa21cf 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index 53a931ba6..61a2aa1f7 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/beam_search_test.py b/tensor2tensor/utils/beam_search_test.py
index 88bc9b43e..4aa0f3cb5 100644
--- a/tensor2tensor/utils/beam_search_test.py
+++ b/tensor2tensor/utils/beam_search_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/bleu_hook.py b/tensor2tensor/utils/bleu_hook.py
index 7346b2b74..9f37226a8 100644
--- a/tensor2tensor/utils/bleu_hook.py
+++ b/tensor2tensor/utils/bleu_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/bleu_hook_test.py b/tensor2tensor/utils/bleu_hook_test.py
index e7ee9c264..173f6e0a5 100644
--- a/tensor2tensor/utils/bleu_hook_test.py
+++ b/tensor2tensor/utils/bleu_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/checkpoint_compatibility_test.py b/tensor2tensor/utils/checkpoint_compatibility_test.py
index 91e7ff363..05bb4dd10 100644
--- a/tensor2tensor/utils/checkpoint_compatibility_test.py
+++ b/tensor2tensor/utils/checkpoint_compatibility_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index c8cf08670..110ca093c 100755
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/compute_video_metrics.py b/tensor2tensor/utils/compute_video_metrics.py
index f41888fdd..887fd410a 100644
--- a/tensor2tensor/utils/compute_video_metrics.py
+++ b/tensor2tensor/utils/compute_video_metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/contrib.py b/tensor2tensor/utils/contrib.py
index e3b34e758..671c005a5 100644
--- a/tensor2tensor/utils/contrib.py
+++ b/tensor2tensor/utils/contrib.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index a979d37d7..3f9f0e386 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/data_reader_test.py b/tensor2tensor/utils/data_reader_test.py
index 4c60b556e..4b2a1e797 100644
--- a/tensor2tensor/utils/data_reader_test.py
+++ b/tensor2tensor/utils/data_reader_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 405cca8a6..91b3476d4 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index bed966566..d689bd58d 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/diet.py b/tensor2tensor/utils/diet.py
index 61a20e3aa..ad15583e6 100644
--- a/tensor2tensor/utils/diet.py
+++ b/tensor2tensor/utils/diet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/diet_test.py b/tensor2tensor/utils/diet_test.py
index e523cfde0..c3592d9aa 100644
--- a/tensor2tensor/utils/diet_test.py
+++ b/tensor2tensor/utils/diet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index b7231fe97..c0be890da 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/expert_utils_test.py b/tensor2tensor/utils/expert_utils_test.py
index 48ab95e1a..8ea49450c 100644
--- a/tensor2tensor/utils/expert_utils_test.py
+++ b/tensor2tensor/utils/expert_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index 1359fd140..ab4ab6e05 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/get_rouge.py b/tensor2tensor/utils/get_rouge.py
index ac8ea0761..3d64acd79 100644
--- a/tensor2tensor/utils/get_rouge.py
+++ b/tensor2tensor/utils/get_rouge.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparam.py b/tensor2tensor/utils/hparam.py
index 5bd58035b..1b073348f 100644
--- a/tensor2tensor/utils/hparam.py
+++ b/tensor2tensor/utils/hparam.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparam_test.py b/tensor2tensor/utils/hparam_test.py
index 637299ad5..8d91861fb 100644
--- a/tensor2tensor/utils/hparam_test.py
+++ b/tensor2tensor/utils/hparam_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparams_lib.py b/tensor2tensor/utils/hparams_lib.py
index 891fee943..2994add29 100644
--- a/tensor2tensor/utils/hparams_lib.py
+++ b/tensor2tensor/utils/hparams_lib.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparams_lib_test.py b/tensor2tensor/utils/hparams_lib_test.py
index e98cff0aa..5941dd52c 100644
--- a/tensor2tensor/utils/hparams_lib_test.py
+++ b/tensor2tensor/utils/hparams_lib_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index d8bcd6a2d..b2a87e8ce 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 32c23751d..2a0570dad 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_hook.py b/tensor2tensor/utils/metrics_hook.py
index 640904fc1..72d298fad 100644
--- a/tensor2tensor/utils/metrics_hook.py
+++ b/tensor2tensor/utils/metrics_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_hook_test.py b/tensor2tensor/utils/metrics_hook_test.py
index 74739730f..ae9d41b20 100644
--- a/tensor2tensor/utils/metrics_hook_test.py
+++ b/tensor2tensor/utils/metrics_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 877590d1a..55e8cee46 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/misc_utils.py b/tensor2tensor/utils/misc_utils.py
index 90d640221..bb5d7a2d7 100644
--- a/tensor2tensor/utils/misc_utils.py
+++ b/tensor2tensor/utils/misc_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/misc_utils_test.py b/tensor2tensor/utils/misc_utils_test.py
index 282cef0c3..ab7003f2e 100644
--- a/tensor2tensor/utils/misc_utils_test.py
+++ b/tensor2tensor/utils/misc_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mlperf_log.py b/tensor2tensor/utils/mlperf_log.py
index 304d60b6d..84a67bc7a 100644
--- a/tensor2tensor/utils/mlperf_log.py
+++ b/tensor2tensor/utils/mlperf_log.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mlperf_tags.py b/tensor2tensor/utils/mlperf_tags.py
index 6de1fd71d..ce153228a 100644
--- a/tensor2tensor/utils/mlperf_tags.py
+++ b/tensor2tensor/utils/mlperf_tags.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mtf_model.py b/tensor2tensor/utils/mtf_model.py
index 607ead673..4825c4116 100644
--- a/tensor2tensor/utils/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
index 2367c8437..37ad1137b 100644
--- a/tensor2tensor/utils/multistep_optimizer.py
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_optimizer_test.py b/tensor2tensor/utils/multistep_optimizer_test.py
index 45cc7b15e..6a1531a41 100644
--- a/tensor2tensor/utils/multistep_optimizer_test.py
+++ b/tensor2tensor/utils/multistep_optimizer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index a88f31ad5..689821e40 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/optimize_test.py b/tensor2tensor/utils/optimize_test.py
index e882b4263..326a45f51 100644
--- a/tensor2tensor/utils/optimize_test.py
+++ b/tensor2tensor/utils/optimize_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/partial_checkpoint_load_hook.py b/tensor2tensor/utils/partial_checkpoint_load_hook.py
index 0e8f42e66..2ff9e9da0 100644
--- a/tensor2tensor/utils/partial_checkpoint_load_hook.py
+++ b/tensor2tensor/utils/partial_checkpoint_load_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/pruning_utils.py b/tensor2tensor/utils/pruning_utils.py
index 835c727f8..4e6e5ab3e 100644
--- a/tensor2tensor/utils/pruning_utils.py
+++ b/tensor2tensor/utils/pruning_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/quantization.py b/tensor2tensor/utils/quantization.py
index 95000fadb..a8011b1ed 100644
--- a/tensor2tensor/utils/quantization.py
+++ b/tensor2tensor/utils/quantization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 603b73e0e..911ace7d9 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index f5f0b7d09..632f7c35f 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/restore_hook.py b/tensor2tensor/utils/restore_hook.py
index 5960b26da..28aa238cc 100644
--- a/tensor2tensor/utils/restore_hook.py
+++ b/tensor2tensor/utils/restore_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/rouge.py b/tensor2tensor/utils/rouge.py
index d18504d42..fb5be8ec8 100644
--- a/tensor2tensor/utils/rouge.py
+++ b/tensor2tensor/utils/rouge.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/rouge_test.py b/tensor2tensor/utils/rouge_test.py
index 08005984e..d91dce4c8 100644
--- a/tensor2tensor/utils/rouge_test.py
+++ b/tensor2tensor/utils/rouge_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/sari_hook.py b/tensor2tensor/utils/sari_hook.py
index df611dec7..9cdacb561 100644
--- a/tensor2tensor/utils/sari_hook.py
+++ b/tensor2tensor/utils/sari_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/sari_hook_test.py b/tensor2tensor/utils/sari_hook_test.py
index 3be960534..1a9c88ce3 100644
--- a/tensor2tensor/utils/sari_hook_test.py
+++ b/tensor2tensor/utils/sari_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/scheduled_sampling.py b/tensor2tensor/utils/scheduled_sampling.py
index a26d20d66..339047f88 100644
--- a/tensor2tensor/utils/scheduled_sampling.py
+++ b/tensor2tensor/utils/scheduled_sampling.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index d4bf75c93..69cdbe0e7 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index 7a8423996..ab9760abb 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/test_utils.py b/tensor2tensor/utils/test_utils.py
index 56912b439..5f9cbcd14 100644
--- a/tensor2tensor/utils/test_utils.py
+++ b/tensor2tensor/utils/test_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/test_utils_test.py b/tensor2tensor/utils/test_utils_test.py
index f5f949701..d671b5114 100644
--- a/tensor2tensor/utils/test_utils_test.py
+++ b/tensor2tensor/utils/test_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index b71ddac8d..40c0efacb 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index faf6d8090..6784894c5 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/update_ops_hook.py b/tensor2tensor/utils/update_ops_hook.py
index 46ac18c84..f03ee29c9 100644
--- a/tensor2tensor/utils/update_ops_hook.py
+++ b/tensor2tensor/utils/update_ops_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/usr_dir.py b/tensor2tensor/utils/usr_dir.py
index c349c4f2c..aa7775c99 100644
--- a/tensor2tensor/utils/usr_dir.py
+++ b/tensor2tensor/utils/usr_dir.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video/prediction2gif.py b/tensor2tensor/utils/video/prediction2gif.py
index cd2d15f53..ae59a00d1 100644
--- a/tensor2tensor/utils/video/prediction2gif.py
+++ b/tensor2tensor/utils/video/prediction2gif.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video/reward_confusion.py b/tensor2tensor/utils/video/reward_confusion.py
index 115303934..2234f94f3 100644
--- a/tensor2tensor/utils/video/reward_confusion.py
+++ b/tensor2tensor/utils/video/reward_confusion.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video2gif.py b/tensor2tensor/utils/video2gif.py
index 2e09df0df..f0d52c783 100644
--- a/tensor2tensor/utils/video2gif.py
+++ b/tensor2tensor/utils/video2gif.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index fe4d8638a..1f77c1a46 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video_metrics_test.py b/tensor2tensor/utils/video_metrics_test.py
index 24da20c46..476ebd1f9 100644
--- a/tensor2tensor/utils/video_metrics_test.py
+++ b/tensor2tensor/utils/video_metrics_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/yellowfin.py b/tensor2tensor/utils/yellowfin.py
index d95820eae..aa383d323 100644
--- a/tensor2tensor/utils/yellowfin.py
+++ b/tensor2tensor/utils/yellowfin.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/yellowfin_test.py b/tensor2tensor/utils/yellowfin_test.py
index 693394362..0c77e60b2 100644
--- a/tensor2tensor/utils/yellowfin_test.py
+++ b/tensor2tensor/utils/yellowfin_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/__init__.py b/tensor2tensor/visualization/__init__.py
index b775a72bd..7688494ed 100644
--- a/tensor2tensor/visualization/__init__.py
+++ b/tensor2tensor/visualization/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/attention.py b/tensor2tensor/visualization/attention.py
index 2b3b2e894..2a1ebf071 100644
--- a/tensor2tensor/visualization/attention.py
+++ b/tensor2tensor/visualization/attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/visualization.py b/tensor2tensor/visualization/visualization.py
index a476f1ea3..7ebe48297 100644
--- a/tensor2tensor/visualization/visualization.py
+++ b/tensor2tensor/visualization/visualization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/visualization_test.py b/tensor2tensor/visualization/visualization_test.py
index 06167d34a..e054c8ebd 100644
--- a/tensor2tensor/visualization/visualization_test.py
+++ b/tensor2tensor/visualization/visualization_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

From 90326159f8465ea0414422e2b7ed1505f9e1af5d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 28 Jan 2020 07:27:21 -0800
Subject: [PATCH 2631/2720] Adding discretization wrapper to tensor2tensor.

The wrapper converts continuous actions to multidiscrete actions, treating each action independently.

PiperOrigin-RevId: 291928405
---
 tensor2tensor/envs/env_problem_utils.py    |  13 +-
 tensor2tensor/envs/gym_env_problem.py      |   3 +-
 tensor2tensor/envs/mujoco_problems.py      |   1 +
 tensor2tensor/envs/rendered_env_problem.py |   4 +-
 tensor2tensor/rl/gym_utils.py              | 127 ++++++++++++++++++-
 tensor2tensor/rl/gym_utils_test.py         | 137 ++++++++++++++++++++-
 6 files changed, 278 insertions(+), 7 deletions(-)

diff --git a/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 180716c85..43859f08d 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -226,6 +226,8 @@ def make_env(batch_size=1,
              clip_rewards=True,
              parallelism=1,
              use_tpu=False,
+             num_actions=None,
+             rendered_env=True,
              **env_kwargs):
   """Creates the env."""
 
@@ -234,6 +236,13 @@ def make_env(batch_size=1,
   else:
     env_kwargs.update({"discrete_rewards": False})
 
+  # TODO(henrykm) - below someone linked "resize" with "abnormality"
+  # Probably we need more nuanced concept of "abnormality"
+  # decoupled from "resize". Currently the resize flag implies
+  # that we switch from a generic env to a wrapped env.
+  # Overall this file and gym_utils.py look like good candidates
+  # for a refactor.
+
   # No resizing needed, so let's be on the normal EnvProblem.
   if not resize:  # None or False
     return gym_env_problem.GymEnvProblem(
@@ -251,15 +260,17 @@ def make_env(batch_size=1,
       gym_utils.gym_env_wrapper, **{
           "rl_env_max_episode_steps": max_timestep,
           "maxskip_env": True,
-          "rendered_env": True,
+          "rendered_env": rendered_env,
           "rendered_env_resize_to": resize_dims,
           "sticky_actions": False,
           "output_dtype": np.int32 if use_tpu else None,
+          "num_actions": num_actions,
       })
 
   return rendered_env_problem.RenderedEnvProblem(
       base_env_name=env_problem_name,
       batch_size=batch_size,
       parallelism=parallelism,
+      rendered_env=rendered_env,
       env_wrapper_fn=wrapper_fn,
       **env_kwargs)
diff --git a/tensor2tensor/envs/gym_env_problem.py b/tensor2tensor/envs/gym_env_problem.py
index 853f7eef8..cf7e02dc9 100644
--- a/tensor2tensor/envs/gym_env_problem.py
+++ b/tensor2tensor/envs/gym_env_problem.py
@@ -195,7 +195,8 @@ def union_dicts(dict1, dict2):
       return copy_dict1
 
     self._envs = [
-        gym.make(self.base_env_name, **union_dicts(kwargs, env_kwarg))
+        gym.make(self.base_env_name,
+                 **union_dicts(kwargs, env_kwarg))
         for env_kwarg in per_env_kwargs
     ]
     self._parallelism = parallelism
diff --git a/tensor2tensor/envs/mujoco_problems.py b/tensor2tensor/envs/mujoco_problems.py
index 2eedeca0b..b58174f9c 100644
--- a/tensor2tensor/envs/mujoco_problems.py
+++ b/tensor2tensor/envs/mujoco_problems.py
@@ -41,6 +41,7 @@ def __init__(self):
             "rendered_env_resize_to": None,  # Do not resize frames
             "sticky_actions": False,
             "output_dtype": None,
+            "num_actions": None,
         })
     super(ReacherEnvProblem, self).__init__(
         base_env_name=base_env_name, env_wrapper_fn=wrapper_fn)
diff --git a/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
index f0344905c..69536a095 100644
--- a/tensor2tensor/envs/rendered_env_problem.py
+++ b/tensor2tensor/envs/rendered_env_problem.py
@@ -59,13 +59,15 @@ def __init__(self, *args, **kwargs):
   def initialize_environments(self,
                               batch_size=1,
                               parallelism=1,
+                              rendered_env=True,
                               per_env_kwargs=None,
                               **kwargs):
     gym_env_problem.GymEnvProblem.initialize_environments(
         self, batch_size=batch_size, parallelism=parallelism,
         per_env_kwargs=per_env_kwargs, **kwargs)
     # Assert the underlying gym environment has correct observation space
-    assert len(self.observation_spec.shape) == 3
+    if rendered_env:
+      assert len(self.observation_spec.shape) == 3
 
   def example_reading_spec(self):
     """Return a mix of env and video data fields and decoders."""
diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index 2737242f5..7fdf857f0 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import math
+
 from absl import logging
 import gym
 import gym.wrappers
@@ -81,6 +83,117 @@ def reset(self, **kwargs):
     return self.env.reset(**kwargs)
 
 
+class ActionDiscretizeWrapper(gym.ActionWrapper):
+  """Wraps an environment with continuous actions and discretizes them.
+
+  This is a simplified adaptation of ActionDiscretizeWrapper
+  from tf_agents.
+  """
+
+  def __init__(self, env, num_actions):
+    """Constructs a wrapper for discretizing the action space.
+
+    Args:
+      env: environment to wrap.
+      num_actions: A np.array of the same shape as the environment's
+        action_spec. Elements in the array specify the number of actions to
+        discretize to for each dimension.
+
+    Raises:
+      ValueError: IF the action_spec shape and the limits shape are not equal.
+    """
+
+    if not isinstance(env.action_space, gym.spaces.box.Box):
+      raise ValueError(
+          "The action space is {}, but gym.spaces.box.Box is expected".format(
+              env.action_space))
+
+    gym.Wrapper.__init__(self, env)
+
+    # We convert a scalar num_actions to array [num_actions, num_actions, ...]
+    self._num_actions = np.broadcast_to(num_actions, env.action_space.shape)
+
+    if env.action_space.shape != self._num_actions.shape:
+      raise ValueError("Spec {} and limit shape do not match. Got {}".format(
+          env.action_space.shape, self._num_actions.shape))
+    self.action_space = gym.spaces.MultiDiscrete(nvec=self._num_actions)
+    self._action_map = self._discretize_env(env)
+
+  def _discretize_env(self, env):
+    """Generates a discrete bounded spec and a linspace for the given limits.
+
+    Args:
+      env: An array to discretize.
+
+    Returns:
+      Tuple with the discrete_spec along with a list of lists mapping actions.
+    Raises:
+      ValueError: If not all limits value are >=2 or maximum or minimum of boxes
+      is equal to +- infinity.
+    """
+    if not np.all(self._num_actions >= 2):
+      raise ValueError("num_actions should all be at least size 2.")
+
+    if (math.isinf(np.min(env.action_space.low)) or
+        math.isinf(np.max(env.action_space.high))):
+      raise ValueError(
+          """Minimum of boxes is {} and maximum of boxes is {},
+          but we expect that finite values are provided.""".
+          format(np.min(env.action_space.low),
+                 np.max(env.action_space.high)))
+
+    limits = np.broadcast_to(self._num_actions,
+                             env.action_space.shape)
+    minimum = np.broadcast_to(np.min(env.action_space.low),
+                              env.action_space.shape)
+    maximum = np.broadcast_to(np.max(env.action_space.high),
+                              env.action_space.shape)
+
+    action_map = [
+        np.linspace(env_min, env_max, num=n_actions)
+        for env_min, env_max, n_actions in zip(
+            np.nditer(minimum), np.nditer(maximum), np.nditer(limits))
+    ]
+
+    return action_map
+
+  def _map_actions(self, action):
+    """Maps the given discrete action to the corresponding continuous action.
+
+    Args:
+      action: Discrete action to map.
+
+    Returns:
+      Numpy array with the mapped continuous actions.
+    Raises:
+      ValueError: If the given action's shpe does not match the action_spec
+      shape.
+    """
+    action = np.asarray(action)
+    if action.shape != self.action_space.shape:
+      raise ValueError(
+          "Received action with incorrect shape. Got {}, expected {}".format(
+              action.shape, self.action_space.shape))
+
+    mapped_action = [self._action_map[i][a]
+                     for i, a in enumerate(action.flatten())]
+    return np.reshape(mapped_action, newshape=action.shape)
+
+  def action(self, action):
+    """Steps the environment while remapping the actions.
+
+    Args:
+      action: Action to take.
+
+    Returns:
+      The next time_step from the environment.
+    """
+    return self._map_actions(action)
+
+  def reverse_action(self, action):
+    raise NotImplementedError
+
+
 class RenderedEnv(gym.Wrapper):
   """Simple Env wrapper to override observations with rendered rgb values."""
 
@@ -158,7 +271,8 @@ def remove_time_limit_wrapper(env):
 
 
 def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env,
-                    rendered_env_resize_to, sticky_actions, output_dtype):
+                    rendered_env_resize_to, sticky_actions, output_dtype,
+                    num_actions):
   """Wraps a gym environment. see make_gym_env for details."""
   # rl_env_max_episode_steps is None or int.
   assert ((not rl_env_max_episode_steps) or
@@ -170,6 +284,10 @@ def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env,
   if wrap_with_time_limit:
     env = remove_time_limit_wrapper(env)
 
+  logging.info("Number of actions: %d", num_actions)
+  if num_actions is not None:
+    env = ActionDiscretizeWrapper(env, num_actions=num_actions)
+
   if sticky_actions:
     env = StickyActionEnv(env)
 
@@ -192,7 +310,8 @@ def make_gym_env(name,
                  rendered_env=False,
                  rendered_env_resize_to=None,
                  sticky_actions=False,
-                 output_dtype=None):
+                 output_dtype=None,
+                 num_actions=None):
   """Create a gym env optionally with a time limit and maxskip wrapper.
 
   NOTE: The returned env may already be wrapped with TimeLimit!
@@ -211,6 +330,8 @@ def make_gym_env(name,
     output_dtype: numpy datatype that we want the observation to be in, if None
       this defaults to the env's observation dtype. Useful for TPUs since they
       don't support uint8 which is a default observation type for a lot of envs.
+    num_actions: None if we do not need discretization and the number of
+      discrete actions per continuous action.
 
   Returns:
     An instance of `gym.Env` or `gym.Wrapper`.
@@ -218,7 +339,7 @@ def make_gym_env(name,
   env = gym.make(name)
   return gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env,
                          rendered_env, rendered_env_resize_to, sticky_actions,
-                         output_dtype)
+                         output_dtype, num_actions)
 
 
 def register_gym_env(class_entry_point, version="v0", kwargs=None):
diff --git a/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index 6755ac6dd..6074bcaa8 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import unittest
+
 import gym
 from gym import spaces
 import numpy as np
@@ -48,6 +50,28 @@ def render(self, mode="human"):
     return np.zeros([640, 480, 3], np.uint8)
 
 
+class SimpleContinuousActionsEnv(gym.Env):
+  """A simple environment with a 3x3 observation space, is done on action=1."""
+
+  def __init__(self, dimensions):
+    self.reward_range = (-1.0, 1.0)
+    self.action_space = spaces.Box(low=-1, high=1, shape=(dimensions,))
+    self.observation_space = spaces.Box(low=0, high=255, shape=(3, 3))
+
+  def reset(self):
+    return self.observation_space.low
+
+  def step(self, action):
+    if action == 0:
+      return self.reset(), -1.0, False, {}
+    else:
+      return self.observation_space.high, +1.0, True, {}
+
+  def render(self, mode="human"):
+    del mode  # Unused
+    return np.zeros([640, 480, 3], np.uint8)
+
+
 class EnvWithOptions(SimpleEnv):
   """A simple env that takes arguments on init."""
 
@@ -92,6 +116,101 @@ def test_rendered_env(self):
     obs, _, _, _ = env.step(1)
     self.assertTrue(np.allclose(np.zeros([64, 12, 3], np.float32), obs))
 
+  def test_rendered_env_continuous_1d(self):
+    env = gym_utils.RenderedEnv(
+        SimpleContinuousActionsEnv(dimensions=1),
+        resize_to=(64, 12))
+    obs, _, _, _ = env.step(0.5)
+    self.assertTrue(np.allclose(np.zeros([64, 12, 3], np.uint8), obs))
+
+    env = gym_utils.RenderedEnv(
+        SimpleContinuousActionsEnv(dimensions=1),
+        resize_to=(64, 12),
+        output_dtype=np.float32)
+    obs, _, _, _ = env.step(1)
+    self.assertTrue(np.allclose(np.zeros([64, 12, 3], np.float32), obs))
+
+  def test_rendered_env_continuous_2d(self):
+    env = gym_utils.RenderedEnv(
+        SimpleContinuousActionsEnv(dimensions=2),
+        resize_to=(64, 12))
+    obs, _, _, _ = env.step(0.5)
+    self.assertTrue(np.allclose(np.zeros([64, 12, 3], np.uint8), obs))
+
+    env = gym_utils.RenderedEnv(
+        SimpleContinuousActionsEnv(dimensions=2),
+        resize_to=(64, 12),
+        output_dtype=np.float32)
+    obs, _, _, _ = env.step(1)
+    self.assertTrue(np.allclose(np.zeros([64, 12, 3], np.float32), obs))
+
+  def test_correct_number_of_discrete_actions_1d(self):
+    """The env should become discrete whenever we pass num_action."""
+    env_discrete = gym_utils.ActionDiscretizeWrapper(
+        gym_utils.RenderedEnv(SimpleContinuousActionsEnv(dimensions=1)),
+        num_actions=4)
+
+    expected_action_space = gym.spaces.MultiDiscrete([4,])
+    self.assertEqual(env_discrete.action_space, expected_action_space)
+
+  def test_correct_number_of_discrete_actions_2d(self):
+    env_discrete = gym_utils.ActionDiscretizeWrapper(
+        gym_utils.RenderedEnv(SimpleContinuousActionsEnv(dimensions=2)),
+        num_actions=4)
+
+    expected_action_space = gym.spaces.MultiDiscrete([4, 4])
+    self.assertEqual(env_discrete.action_space, expected_action_space)
+
+  def test_action_mapping_1d(self):
+    """Testing discretization with a mock environment.
+
+    In the mock call we get access to the argument of the
+    SimpleContinuousActionsEnv.step method which we check against
+    precomputed values of continuous actions.
+    """
+    num_actions = 4
+
+    with unittest.mock.patch.object(
+        gym_utils.RenderedEnv, "step", autospec=True) as mock_step_method:
+      env = gym_utils.RenderedEnv(SimpleContinuousActionsEnv(dimensions=1))
+      expected_continuous_actions = np.linspace(
+          np.min(env.action_space.low),
+          np.min(env.action_space.high),
+          num=num_actions).flatten()
+
+      env_discrete = gym_utils.ActionDiscretizeWrapper(env, num_actions)
+      for discrete_action in range(num_actions):
+        env_discrete.step([discrete_action])
+        mock_step_method.assert_called_with(
+            unittest.mock.ANY,
+            expected_continuous_actions[discrete_action])
+
+  def test_action_mapping_2d(self):
+    num_actions = 8
+
+    def expected_continuous_actions(discrete_action):
+      if discrete_action == [0, 0]:
+        return np.array([-1, -1])
+      elif discrete_action == [0, 3]:
+        return np.array([-1, -0.14285714])
+      elif discrete_action == [4, 4]:
+        return np.array([0.14285714, 0.14285714])
+      elif discrete_action == [7, 7]:
+        return np.array([1, 1])
+
+    discrete_actions = [[0, 0], [0, 3], [4, 4], [7, 7]]
+
+    with unittest.mock.patch.object(
+        gym_utils.RenderedEnv, "step", autospec=True) as mock_step_method:
+      env = gym_utils.RenderedEnv(SimpleContinuousActionsEnv(dimensions=2))
+
+      env_discrete = gym_utils.ActionDiscretizeWrapper(env, num_actions)
+      for discrete_action in discrete_actions:
+        env_discrete.step(discrete_action)
+        mock_args, _ = mock_step_method.call_args
+        np.testing.assert_array_almost_equal(
+            mock_args[1], expected_continuous_actions(discrete_action))
+
   def test_gym_registration(self):
     reg_id, env = gym_utils.register_gym_env(
         "tensor2tensor.rl.gym_utils_test:SimpleEnv")
@@ -108,6 +227,23 @@ def test_gym_registration(self):
     _, _, done, _ = env.step(1)
     self.assertTrue(done)
 
+  def test_gym_registration_continuous(self):
+    reg_id, env = gym_utils.register_gym_env(
+        "tensor2tensor.rl.gym_utils_test:SimpleContinuousActionsEnv",
+        kwargs={"dimensions": 2})
+
+    self.assertEqual("T2TEnv-SimpleContinuousActionsEnv-v0", reg_id)
+
+    # Most basic check.
+    self.assertIsInstance(env, gym.Env)
+
+    # Just make sure we got the same environment.
+    self.assertTrue(
+        np.allclose(env.reset(), np.zeros(shape=(3, 3), dtype=np.uint8)))
+
+    _, _, done, _ = env.step(1)
+    self.assertTrue(done)
+
   def test_gym_registration_with_kwargs(self):
     reg_id, env = gym_utils.register_gym_env(
         "tensor2tensor.rl.gym_utils_test:EnvWithOptions",
@@ -149,6 +285,5 @@ def test_gym_registration_with_kwargs(self):
     _, _, done, _ = env.step(1)
     self.assertTrue(done)
 
-
 if __name__ == "__main__":
   tf.test.main()

From c8503f7f6c943d21b80b2231d984793221af3206 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 1 Feb 2020 01:50:42 -0800
Subject: [PATCH 2632/2720] %d fails with the arg is None.

PiperOrigin-RevId: 292686878
---
 tensor2tensor/rl/gym_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index 7fdf857f0..d1dc7840a 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -284,8 +284,8 @@ def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env,
   if wrap_with_time_limit:
     env = remove_time_limit_wrapper(env)
 
-  logging.info("Number of actions: %d", num_actions)
   if num_actions is not None:
+    logging.info("Number of discretized actions: %d", num_actions)
     env = ActionDiscretizeWrapper(env, num_actions=num_actions)
 
   if sticky_actions:

From ab9fb79b834a69433fe4d82f98ecd73d9ed9f853 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Em=C4=ABls=20Ozoli=C5=86=C5=A1?= <ozolinsemils@gmail.com>
Date: Tue, 4 Feb 2020 18:34:54 +0200
Subject: [PATCH 2633/2720] Inference support for Neural Shuffle-Exchange
 network (#1784)

* Inference support for Neural Shuffle-Exchange network.

When target dimensions are not known use hparams.max_length. max_length should be power of 2.

Max length can be enforced using hparams.force_max_length=True.

Auto-Regressive decoding not supported at this moment.

* Remove unused variable from Shuffle Network

* Remove type annotations and debug print calls

* Fix docstring lint errors
---
 .../models/research/shuffle_network.py        | 92 ++++++++++++++++---
 1 file changed, 81 insertions(+), 11 deletions(-)

diff --git a/tensor2tensor/models/research/shuffle_network.py b/tensor2tensor/models/research/shuffle_network.py
index 201cb198b..a79023c8a 100644
--- a/tensor2tensor/models/research/shuffle_network.py
+++ b/tensor2tensor/models/research/shuffle_network.py
@@ -28,6 +28,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import math
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
@@ -348,21 +349,88 @@ def bottom(self, features):
       dictionary: Inputs and targets padded with 0 to the length of power of 2.
       Both are same length.
     """
-    inputs = features["inputs"]
-    targets = features["targets"]
-    inputs_length = tf.shape(inputs)[1]
-    targets_length = tf.shape(targets)[1]
+    pad_len = self.max_pad_length(features)
+    features["inputs"] = self.pad(features["inputs"], pad_len)
+
+    if features.get("targets") is not None:
+      features["targets"] = self.pad(features["targets"], pad_len)
+
+    return super(ShuffleNetwork, self).bottom(features)
+
+  @staticmethod
+  def pad(tensor, pad_len):
+    """Pad tensor on first dimension to pad_len.
+
+    Args:
+      tensor: input tensor of shape length >= 2
+      pad_len: pad length
+
+    Returns:
+      tf.Tensor: Padded input tensor.
+    """
+
+    assert len(tensor.shape) >= 2  # tensor of shape [batch, length, ...]
+    length = tf.shape(tensor)[1]
+
+    padding = [[0, 0], [0, pad_len - length]]
+    padding += [[0, 0]] * (len(tensor.shape) - 2)
+    return tf.pad(tensor, padding)
+
+  def max_pad_length(self, features):
+    """Finds max padding length.
+
+    If target length not specified use fixed padding
+    length from hparams.max_length.
+
+    Args:
+      features: Dictionary with input and target tensors
+
+    Returns:
+      tf.Tensor:  Length of input and output sequence. Length is power of 2.
+    """
+
+    if self.hparams.force_max_length or features.get("targets") is None:
+      assert math.log(self.hparams.max_length, 2).is_integer(), \
+        "hparams.max_length should be power of w"
+
+      return self.hparams.max_length
+
+    length = tf.shape(features["inputs"])[1]
+    targets_length = tf.shape(features["targets"])[1]
+    length = tf.maximum(length, targets_length)
 
-    length = tf.maximum(inputs_length, targets_length)
     p = tf.log(tf.cast(length, tf.float32)) / tf.log(2.0)
     p = tf.cast(tf.ceil(p), tf.int32)
-    pad_len = tf.pow(2, p)
+    return tf.pow(2, p)
 
-    input_padding = [[0, 0], [0, pad_len - inputs_length], [0, 0], [0, 0]]
-    features["inputs"] = tf.pad(inputs, input_padding)
-    target_padding = [[0, 0], [0, pad_len - targets_length], [0, 0], [0, 0]]
-    features["targets"] = tf.pad(targets, target_padding)
-    return super(ShuffleNetwork, self).bottom(features)
+  def infer(self, features=None, **kwargs):
+    """Custom infer method for Shuffle-Exchange network.
+
+    Args:
+      features: Dictionary of inputs and targets
+      **kwargs: SE network currently doesn't support auto-regressive output
+
+    Returns:
+      dict: Dictionary of outputs.
+    """
+
+    del kwargs
+    targets = features.get("targets")
+    infer_targets = features.get("infer_targets")
+
+    if targets is None and infer_targets is not None:
+      features["targets"] = infer_targets
+
+    # Run the model
+    self.hparams.force_full_predict = True
+    with tf.variable_scope(self.name):
+      logits, _ = self.model_fn(features)
+
+    assert len(logits.shape) == 5  # [batch, time, 1, 1, vocab]
+    logits = tf.squeeze(logits, [2, 3])
+    outputs = tf.argmax(logits, axis=2)
+
+    return {"outputs": outputs, "logits": logits, "scores": None}
 
   def loss(self, logits, features):
     """Loss function for Neural Shuffle-Exchange network.
@@ -421,6 +489,8 @@ def shuffle_network_baseline():
   hparams.initializer = "uniform_unit_scaling"
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.999
+  hparams.add_hparam("force_max_length", False)  # use fixed max length
+  hparams.max_length = 256  # use when targets are not known
 
   hparams.dropout = 0.1
   hparams.label_smoothing = 0.

From 6fa217a49a78468c47ab52ea1a92683470c85a7d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 6 Feb 2020 19:15:03 -0800
Subject: [PATCH 2634/2720] Allow setting sampling temperature per example
 instead of having the same temperature for the entire batch.

PiperOrigin-RevId: 293729106
---
 tensor2tensor/data_generators/problem.py   | 20 +++++--------
 tensor2tensor/layers/common_layers.py      | 34 ++++++++++++++++++++++
 tensor2tensor/layers/common_layers_test.py | 14 +++++++++
 tensor2tensor/models/transformer.py        | 24 ++++++++++-----
 tensor2tensor/utils/t2t_model.py           |  4 +--
 5 files changed, 73 insertions(+), 23 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index a6c43dbce..c3ccf7653 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -710,6 +710,11 @@ def decode_example(self, serialized_example):
     # Necessary to rejoin examples in the correct order with the Cloud ML Engine
     # batch prediction API.
     data_fields["batch_prediction_key"] = tf.FixedLenFeature([1], tf.int64, 0)
+
+    if getattr(self._hparams, "sampling_method", "") == "random_per_example":
+      data_fields["sampling_temp"] = tf.FixedLenFeature(
+          [1], tf.float32, getattr(self._hparams, "sampling_temp", 1.0))
+
     if data_items_to_decoders is None:
       data_items_to_decoders = {
           field: contrib.slim().tfexample_decoder.Tensor(field)
@@ -905,13 +910,10 @@ def export_assets(self):
 
   def serving_input_fn(self, hparams, decode_hparams=None, use_tpu=False):
     """Input fn for serving export, starting from serialized example."""
+    self._hparams = hparams
     mode = tf.estimator.ModeKeys.PREDICT
     serialized_example = tf.placeholder(
         dtype=tf.string, shape=[None], name="serialized_example")
-    sampling_temp = tf.placeholder_with_default(
-        tf.constant(getattr(hparams, "sampling_temp", 0.0), dtype=tf.float32),
-        shape=[],
-        name="sampling_temp")
     dataset = tf.data.Dataset.from_tensor_slices(serialized_example)
     dataset = dataset.map(self.decode_example)
     dataset = dataset.map(lambda ex: self.preprocess_example(ex, mode, hparams))
@@ -933,20 +935,12 @@ def serving_input_fn(self, hparams, decode_hparams=None, use_tpu=False):
 
     dataset = dataset.map(data_reader.standardize_shapes)
     features = tf.data.experimental.get_single_element(dataset)
-    features["sampling_temp"] = sampling_temp
 
     if self.has_inputs:
       features.pop("targets", None)
 
     return tf.estimator.export.ServingInputReceiver(
-        features=features,
-        receiver_tensors=serialized_example,
-        receiver_tensors_alternatives={
-            "sample": {
-                "input": serialized_example,
-                "sampling_temp": sampling_temp
-            }
-        })
+        features=features, receiver_tensors=serialized_example)
 
 
 class FeatureInfo(object):
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 992b48f23..aeff78bec 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2974,6 +2974,40 @@ def sample_with_temperature(logits, temperature, sampling_keep_top_k=-1):
     return choices
 
 
+# TODO(bosma): vectorize top k as well
+def sample_temperature_per_example(logits, temperature, sampling_keep_top_k=-1):
+  """Either random sampling with different temperature per example.
+
+  Args:
+    logits: a Tensor.
+    temperature: a float vector of same size as logits.
+    sampling_keep_top_k: If not -1, only sample from the top k logits.
+  Returns:
+    a Tensor with one fewer dimension than logits.
+  """
+  if sampling_keep_top_k != -1:
+    if sampling_keep_top_k <= 0:
+      raise ValueError("sampling_keep_top_k must either be -1 or positive.")
+
+    vocab_size = shape_list(logits)[1]
+
+    k_largest = contrib.nn().nth_element(
+        logits, n=sampling_keep_top_k, reverse=True)
+    k_largest = tf.tile(tf.reshape(k_largest, [-1, 1]), [1, vocab_size])
+
+    # Force every position that is not in the top k to have probability near
+    # 0 by setting the logit to be very negative.
+    logits = tf.where(tf.less_equal(logits, k_largest),
+                      tf.ones_like(logits)*-1e6, logits)
+
+  logits /= tf.reshape(temperature, [-1] + [1] * (len(logits.shape) - 1))
+  reshaped_logits = tf.reshape(logits, [-1, shape_list(logits)[-1]])
+  choices = tf.multinomial(reshaped_logits, 1)
+  choices = tf.reshape(choices,
+                       shape_list(logits)[:logits.get_shape().ndims - 1])
+  return choices
+
+
 def ones_matrix_band_part(rows, cols, num_lower, num_upper, out_shape=None):
   """Matrix band part of ones.
 
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index 19ec82355..00bdaffec 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -677,6 +677,20 @@ def testConvHiddenReluMemoryEfficient(self):
     self.assertAllClose(dnorm_bias, dnorm_bias_f)
     self.assertAllClose(dx, dx_f)
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testSampleTemperaturePerExample(self):
+    batch_size = 3
+    seq_len = 5
+    vocab_size = 7
+
+    logits = np.random.randn(batch_size, seq_len, 1, 1, vocab_size)
+    temperature = np.random.rand(batch_size)
+
+    out = common_layers.sample_temperature_per_example(logits, temperature)
+
+    self.assertAllEqual(
+        self.evaluate(tf.shape(out)), [batch_size, seq_len, 1, 1])
+
   @test_utils.run_in_graph_and_eager_modes()
   def testCycleGANUpsampleNnUpsampleConv(self):
     batch = 8
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 051b74859..129df36d1 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1082,10 +1082,14 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       log_probs = common_layers.log_prob_from_logits(logits)
       temperature = sampling_temperature
       keep_top = getattr(hparams, "sampling_keep_top_k", -1)
-      if hparams.sampling_method == "argmax":
-        temperature = 0.0
-      next_id = common_layers.sample_with_temperature(
-          logits, temperature, keep_top)
+      if hparams.sampling_method == "random_per_example":
+        next_id = common_layers.sample_temperature_per_example(
+            logits, temperature, keep_top)
+      else:
+        if hparams.sampling_method == "argmax":
+          temperature = 0.0
+        next_id = common_layers.sample_with_temperature(
+            logits, temperature, keep_top)
 
       log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],
                                   axis=1)
@@ -1229,10 +1233,14 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       log_probs = common_layers.log_prob_from_logits(logits)
       temperature = sampling_temperature
       keep_top = getattr(hparams, "sampling_keep_top_k", -1)
-      if hparams.sampling_method == "argmax":
-        temperature = 0.0
-      next_id = common_layers.sample_with_temperature(
-          logits, temperature, keep_top)
+      if hparams.sampling_method == "random_per_example":
+        next_id = common_layers.sample_temperature_per_example(
+            logits, temperature, keep_top)
+      else:
+        if hparams.sampling_method == "argmax":
+          temperature = 0.0
+        next_id = common_layers.sample_with_temperature(
+            logits, temperature, keep_top)
 
       log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],
                                   axis=1)
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 69cdbe0e7..78b5cd77a 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1351,8 +1351,8 @@ def sample(self, features):
 
       def multinomial_squeeze(logits, temperature=1.0):
         logits_shape = common_layers.shape_list(logits)
-        reshaped_logits = (
-            tf.reshape(logits, [-1, logits_shape[-1]]) / temperature)
+        logits /= tf.reshape(temperature, [-1] + [1] * (len(logits_shape) - 1))
+        reshaped_logits = tf.reshape(logits, [-1, logits_shape[-1]])
         choices = tf.multinomial(reshaped_logits, 1)
         choices = tf.reshape(choices, logits_shape[:-1])
         return choices

From 27b0b53f44e356954420c3c2184e0cbf3f2bfb29 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sun, 9 Feb 2020 17:18:44 -0800
Subject: [PATCH 2635/2720] Allow changing sampling top k during inference.

PiperOrigin-RevId: 294124311
---
 tensor2tensor/data_generators/problem.py   |  2 +
 tensor2tensor/layers/common_layers.py      | 69 +++++++++++++++++-----
 tensor2tensor/layers/common_layers_test.py | 18 ++++++
 tensor2tensor/models/transformer.py        | 28 +++++----
 4 files changed, 93 insertions(+), 24 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index c3ccf7653..ccbdf1681 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -714,6 +714,8 @@ def decode_example(self, serialized_example):
     if getattr(self._hparams, "sampling_method", "") == "random_per_example":
       data_fields["sampling_temp"] = tf.FixedLenFeature(
           [1], tf.float32, getattr(self._hparams, "sampling_temp", 1.0))
+      data_fields["sampling_keep_top_k"] = tf.FixedLenFeature(
+          [1], tf.int64, getattr(self._hparams, "sampling_keep_top_k", -1))
 
     if data_items_to_decoders is None:
       data_items_to_decoders = {
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index aeff78bec..0d5717bcc 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2974,7 +2974,60 @@ def sample_with_temperature(logits, temperature, sampling_keep_top_k=-1):
     return choices
 
 
-# TODO(bosma): vectorize top k as well
+def _to_nd_indices(indices):
+  """Returns indices used for tf.gather_nd or tf.scatter_nd.
+
+  Args:
+    indices: A `Tensor` of shape [batch_size, size] with integer values. The
+      values are the indices of another `Tensor`. For example, `indices` is the
+      output of tf.argsort or tf.math.top_k.
+
+  Returns:
+    A `Tensor` with shape [batch_size, size, 2] that can be used by tf.gather_nd
+    or tf.scatter_nd.
+
+  """
+  indices.get_shape().assert_has_rank(2)
+  batch_ids = tf.ones_like(indices) * tf.expand_dims(
+      tf.range(tf.shape(input=indices)[0]), 1)
+  return tf.stack([batch_ids, indices], axis=-1)
+
+
+def _select_top_k(logits, top_k):
+  """Replaces logits, expect the top k highest values, with small number (-1e6).
+
+  If k is -1 don't replace anything.
+
+  Args:
+    logits: A `Tensor` of shape [batch_size, ..., vocab_size]
+    top_k: vector of batch size.
+
+  Returns:
+    A `Tensor` with same shape  as logits.
+  """
+  vocab_size = logits.shape[-1]
+  flat_logits = tf.reshape(logits, [-1, vocab_size])
+  top_k = tf.where(
+      tf.not_equal(top_k, -1), top_k,
+      tf.ones_like(top_k) * vocab_size)
+  values, idx = tf.math.top_k(flat_logits, k=vocab_size, sorted=False)
+  nd_idx = _to_nd_indices(idx)
+
+  mask_idx = tf.reshape(
+      tf.range(vocab_size), [1] * (len(logits.shape) - 1) + [-1])
+  for i, size in enumerate(logits.shape[:-1]):
+    mask_idx = tf.repeat(mask_idx, size, axis=i)
+  mask = tf.reshape(
+      mask_idx < tf.reshape(top_k, [-1] + [1] * (len(logits.shape) - 1)), [-1])
+
+  topk_logits = tf.tensor_scatter_nd_update(
+      tf.ones_like(flat_logits) * -1e6,
+      tf.reshape(nd_idx, [-1, 2])[mask],
+      tf.reshape(values, [-1])[mask])
+
+  return tf.reshape(topk_logits, logits.shape)
+
+
 def sample_temperature_per_example(logits, temperature, sampling_keep_top_k=-1):
   """Either random sampling with different temperature per example.
 
@@ -2986,19 +3039,7 @@ def sample_temperature_per_example(logits, temperature, sampling_keep_top_k=-1):
     a Tensor with one fewer dimension than logits.
   """
   if sampling_keep_top_k != -1:
-    if sampling_keep_top_k <= 0:
-      raise ValueError("sampling_keep_top_k must either be -1 or positive.")
-
-    vocab_size = shape_list(logits)[1]
-
-    k_largest = contrib.nn().nth_element(
-        logits, n=sampling_keep_top_k, reverse=True)
-    k_largest = tf.tile(tf.reshape(k_largest, [-1, 1]), [1, vocab_size])
-
-    # Force every position that is not in the top k to have probability near
-    # 0 by setting the logit to be very negative.
-    logits = tf.where(tf.less_equal(logits, k_largest),
-                      tf.ones_like(logits)*-1e6, logits)
+    logits = _select_top_k(logits, sampling_keep_top_k)
 
   logits /= tf.reshape(temperature, [-1] + [1] * (len(logits.shape) - 1))
   reshaped_logits = tf.reshape(logits, [-1, shape_list(logits)[-1]])
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index 00bdaffec..bcd19b202 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -677,6 +677,24 @@ def testConvHiddenReluMemoryEfficient(self):
     self.assertAllClose(dnorm_bias, dnorm_bias_f)
     self.assertAllClose(dx, dx_f)
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testTopk(self):
+    batch_size = 3
+    seq_len = 5
+    vocab_size = 7
+
+    top_k = [3, 2, -1]
+    logits = np.random.rand(batch_size, seq_len, 1, 1, vocab_size) + 0.001
+    topk_logits = common_layers._select_top_k(logits, top_k)
+
+    self.evaluate(tf.global_variables_initializer())
+    topk_logits = self.evaluate(topk_logits)
+
+    for i, k in enumerate(top_k):
+      for j in range(seq_len):
+        self.assertEqual((topk_logits[i, j, 0, 0, :] > -1e6).sum(),
+                         k if k != -1 else vocab_size)
+
   @test_utils.run_in_graph_and_eager_modes()
   def testSampleTemperaturePerExample(self):
     batch_size = 3
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 129df36d1..afbd40a18 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -625,6 +625,8 @@ def forced_logits():
     eos_id = self.get_decode_end_id() or beam_search.EOS_ID
     temperature = features.get("sampling_temp",
                                getattr(hparams, "sampling_temp", 0.0))
+    top_k = features.get("sampling_keep_top_k",
+                         getattr(hparams, "sampling_keep_top_k", -1))
 
     ret = fast_decode_tpu(
         encoder_output=encoder_output,
@@ -640,7 +642,8 @@ def forced_logits():
         batch_size=batch_size,
         force_decode_length=self._decode_hparams.force_decode_length,
         eos_id=eos_id,
-        sampling_temperature=temperature)
+        sampling_temperature=temperature,
+        top_k=top_k)
     if partial_targets is not None:
       if beam_size <= 1 or top_beams <= 1:
         ret["outputs"] = ret["outputs"][:, partial_targets_length:]
@@ -889,6 +892,8 @@ def forced_logits():
     eos_id = self.get_decode_end_id() or beam_search.EOS_ID
     temperature = features.get("sampling_temp",
                                getattr(hparams, "sampling_temp", 0.0))
+    top_k = features.get("sampling_keep_top_k",
+                         getattr(hparams, "sampling_keep_top_k", -1))
 
     ret = fast_decode(
         encoder_output=encoder_output,
@@ -906,6 +911,7 @@ def forced_logits():
         sos_id=sos_id,
         eos_id=eos_id,
         sampling_temperature=temperature,
+        top_k=top_k,
         cache=att_cache)
     if partial_targets is not None:
       if beam_size <= 1 or top_beams <= 1:
@@ -994,7 +1000,8 @@ def fast_decode_tpu(encoder_output,
                     force_decode_length=False,
                     scope_prefix="body/",
                     use_top_k_with_unique=True,
-                    sampling_temperature=0.0):
+                    sampling_temperature=0.0,
+                    top_k=-1):
   """Given encoder output and a symbols to logits function, does fast decoding.
 
   Implements both greedy and beam search decoding for TPU, uses beam search iff
@@ -1023,6 +1030,7 @@ def fast_decode_tpu(encoder_output,
     use_top_k_with_unique: bool, whether to use a fast (but decreased precision)
       top_k during beam search.
     sampling_temperature: scalar, temperature with which to sample.
+    top_k: scalar, sample only top k.
 
   Returns:
     A dict of decoding results {
@@ -1081,15 +1089,14 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       logits, cache = symbols_to_logits_fn(next_id, i, cache)
       log_probs = common_layers.log_prob_from_logits(logits)
       temperature = sampling_temperature
-      keep_top = getattr(hparams, "sampling_keep_top_k", -1)
       if hparams.sampling_method == "random_per_example":
         next_id = common_layers.sample_temperature_per_example(
-            logits, temperature, keep_top)
+            logits, temperature, top_k)
       else:
         if hparams.sampling_method == "argmax":
           temperature = 0.0
-        next_id = common_layers.sample_with_temperature(
-            logits, temperature, keep_top)
+        next_id = common_layers.sample_with_temperature(logits, temperature,
+                                                        top_k)
 
       log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],
                                   axis=1)
@@ -1156,6 +1163,7 @@ def fast_decode(encoder_output,
                 force_decode_length=False,
                 scope_prefix="body/",
                 sampling_temperature=0.0,
+                top_k=-1,
                 cache=None):
   """Given encoder output and a symbols to logits function, does fast decoding.
 
@@ -1183,6 +1191,7 @@ def fast_decode(encoder_output,
       False, stop when all beams hit eos_id.
     scope_prefix: str, prefix for decoder layer variable scopes.
     sampling_temperature: scalar, temperature with which to sample.
+    top_k: scalar, sample only top k.
     cache: cache dictionary for additional predictions.
 
   Returns:
@@ -1232,15 +1241,14 @@ def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
       logits, cache = symbols_to_logits_fn(next_id, i, cache)
       log_probs = common_layers.log_prob_from_logits(logits)
       temperature = sampling_temperature
-      keep_top = getattr(hparams, "sampling_keep_top_k", -1)
       if hparams.sampling_method == "random_per_example":
         next_id = common_layers.sample_temperature_per_example(
-            logits, temperature, keep_top)
+            logits, temperature, top_k)
       else:
         if hparams.sampling_method == "argmax":
           temperature = 0.0
-        next_id = common_layers.sample_with_temperature(
-            logits, temperature, keep_top)
+        next_id = common_layers.sample_with_temperature(logits, temperature,
+                                                        top_k)
 
       log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],
                                   axis=1)

From cb5fe528dffb9f13496e33dd1e8c1686d4430d3b Mon Sep 17 00:00:00 2001
From: Rohan Jain <rohanj@google.com>
Date: Mon, 10 Feb 2020 22:23:40 -0800
Subject: [PATCH 2636/2720] Explicitly replace "import tensorflow" with
 "tensorflow.compat.v1" for TF2.x migration

PiperOrigin-RevId: 294375603
---
 tensor2tensor/rl/datagen_with_agent.py                  | 2 +-
 tensor2tensor/rl/evaluator.py                           | 2 +-
 tensor2tensor/rl/player.py                              | 2 +-
 tensor2tensor/rl/player_utils.py                        | 2 +-
 tensor2tensor/rl/trainer_model_based.py                 | 2 +-
 tensor2tensor/rl/trainer_model_based_agent_only.py      | 2 +-
 tensor2tensor/rl/trainer_model_based_params.py          | 2 +-
 tensor2tensor/rl/trainer_model_based_recurrent_test.py  | 2 +-
 tensor2tensor/rl/trainer_model_based_stochastic_test.py | 2 +-
 tensor2tensor/rl/trainer_model_based_sv2p_test.py       | 2 +-
 tensor2tensor/rl/trainer_model_based_test.py            | 2 +-
 tensor2tensor/rl/trainer_model_free.py                  | 2 +-
 tensor2tensor/rl/trainer_model_free_test.py             | 2 +-
 tensor2tensor/rl/trainer_model_free_tictactoe_test.py   | 2 +-
 14 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/rl/datagen_with_agent.py b/tensor2tensor/rl/datagen_with_agent.py
index 2b69fcf55..8199a1486 100644
--- a/tensor2tensor/rl/datagen_with_agent.py
+++ b/tensor2tensor/rl/datagen_with_agent.py
@@ -25,7 +25,7 @@
 from tensor2tensor.data_generators import gym_env
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index a1df89c7c..1f16c3699 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -41,7 +41,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 flags = tf.flags
diff --git a/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 3304f20ab..8ce9b49f9 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -67,7 +67,7 @@
 import tensor2tensor.rl.trainer_model_based_params  # pylint: disable=unused-import
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 flags = tf.flags
diff --git a/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index 630bb83c5..0be5e19d0 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -35,7 +35,7 @@
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils.misc_utils import camelcase_to_snakecase
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 flags = tf.flags
diff --git a/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index 24d7fba3e..f82defbfc 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -43,7 +43,7 @@
 from tensor2tensor.rl.restarter import Restarter
 from tensor2tensor.utils import trainer_lib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 flags = tf.flags
diff --git a/tensor2tensor/rl/trainer_model_based_agent_only.py b/tensor2tensor/rl/trainer_model_based_agent_only.py
index da3b17e57..55ca3dbc6 100644
--- a/tensor2tensor/rl/trainer_model_based_agent_only.py
+++ b/tensor2tensor/rl/trainer_model_based_agent_only.py
@@ -33,7 +33,7 @@
 from tensor2tensor.rl import trainer_model_based_params
 
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 flags = tf.flags
diff --git a/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 4b3bd06bc..04a1309e3 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -25,7 +25,7 @@
 from tensor2tensor.utils import hparam
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 flags = tf.flags
diff --git a/tensor2tensor/rl/trainer_model_based_recurrent_test.py b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
index 5642bcc56..67844c610 100644
--- a/tensor2tensor/rl/trainer_model_based_recurrent_test.py
+++ b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
@@ -20,7 +20,7 @@
 
 from tensor2tensor.rl import trainer_model_based
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 FLAGS = tf.flags.FLAGS
 
diff --git a/tensor2tensor/rl/trainer_model_based_stochastic_test.py b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
index 732ef7327..2e0159796 100644
--- a/tensor2tensor/rl/trainer_model_based_stochastic_test.py
+++ b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
@@ -20,7 +20,7 @@
 
 from tensor2tensor.rl import trainer_model_based
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 FLAGS = tf.flags.FLAGS
 
diff --git a/tensor2tensor/rl/trainer_model_based_sv2p_test.py b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
index 0df918c74..148c47983 100644
--- a/tensor2tensor/rl/trainer_model_based_sv2p_test.py
+++ b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
@@ -20,7 +20,7 @@
 
 from tensor2tensor.rl import trainer_model_based
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 FLAGS = tf.flags.FLAGS
 
diff --git a/tensor2tensor/rl/trainer_model_based_test.py b/tensor2tensor/rl/trainer_model_based_test.py
index 39caf4d26..68dfd1ac9 100644
--- a/tensor2tensor/rl/trainer_model_based_test.py
+++ b/tensor2tensor/rl/trainer_model_based_test.py
@@ -20,7 +20,7 @@
 
 from tensor2tensor.rl import trainer_model_based
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 FLAGS = tf.flags.FLAGS
 
diff --git a/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 4f56dc3ba..8956666a4 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -47,7 +47,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 flags = tf.flags
diff --git a/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index ad2a15ab8..f654a1f63 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -21,7 +21,7 @@
 from tensor2tensor.rl import trainer_model_free
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 FLAGS = tf.flags.FLAGS
 
diff --git a/tensor2tensor/rl/trainer_model_free_tictactoe_test.py b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
index b03195165..245336495 100644
--- a/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
+++ b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
@@ -22,7 +22,7 @@
 from tensor2tensor.rl import trainer_model_free
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 FLAGS = tf.flags.FLAGS
 

From 99dca5108015926b245193b7e374577960868b7e Mon Sep 17 00:00:00 2001
From: Rohan Jain <rohanj@google.com>
Date: Tue, 11 Feb 2020 20:07:21 -0800
Subject: [PATCH 2637/2720] Explicitly replace "import tensorflow" with
 "tensorflow.compat.v1" for TF2.x migration

PiperOrigin-RevId: 294580218
---
 tensor2tensor/data_generators/allen_brain.py                | 2 +-
 tensor2tensor/data_generators/allen_brain_test.py           | 2 +-
 tensor2tensor/data_generators/bair_robot_pushing.py         | 2 +-
 tensor2tensor/data_generators/fsns.py                       | 2 +-
 tensor2tensor/data_generators/gym_env.py                    | 2 +-
 tensor2tensor/data_generators/image_utils.py                | 2 +-
 tensor2tensor/data_generators/inspect_tfrecord.py           | 2 +-
 tensor2tensor/data_generators/moving_mnist.py               | 2 +-
 tensor2tensor/data_generators/problem.py                    | 2 +-
 tensor2tensor/data_generators/text_encoder_build_subword.py | 2 +-
 tensor2tensor/data_generators/translate.py                  | 2 +-
 tensor2tensor/data_generators/video_generated.py            | 2 +-
 tensor2tensor/data_generators/video_utils.py                | 2 +-
 tensor2tensor/data_generators/vqa.py                        | 2 +-
 14 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index 19951e435..aef529a3e 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -44,7 +44,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 _BASE_EXAMPLE_IMAGE_SIZE = 64
 
diff --git a/tensor2tensor/data_generators/allen_brain_test.py b/tensor2tensor/data_generators/allen_brain_test.py
index c19874d54..2c8a120a6 100644
--- a/tensor2tensor/data_generators/allen_brain_test.py
+++ b/tensor2tensor/data_generators/allen_brain_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.models import image_transformer_2d
 from tensor2tensor.utils import contrib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 tfe = contrib.eager()
diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index fae556be2..f9ccc5a6a 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -36,7 +36,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 DATA_URL = (
     "http://rail.eecs.berkeley.edu/datasets/bair_robot_pushing_dataset_v0.tar")
diff --git a/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
index 5f3dab07e..c1d3b60d5 100644
--- a/tensor2tensor/data_generators/fsns.py
+++ b/tensor2tensor/data_generators/fsns.py
@@ -28,7 +28,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_problem
diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 78ca8186f..e48f59fbb 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -36,7 +36,7 @@
 from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 Frame = collections.namedtuple(
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index edd8e61a9..3f439fa47 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -31,7 +31,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import metrics
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def matplotlib_pyplot():
diff --git a/tensor2tensor/data_generators/inspect_tfrecord.py b/tensor2tensor/data_generators/inspect_tfrecord.py
index 0cc905c1b..6b4239221 100644
--- a/tensor2tensor/data_generators/inspect_tfrecord.py
+++ b/tensor2tensor/data_generators/inspect_tfrecord.py
@@ -29,7 +29,7 @@
 
 from tensor2tensor.data_generators import text_encoder
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 tf.flags.DEFINE_string("subword_text_encoder_filename", "",
diff --git a/tensor2tensor/data_generators/moving_mnist.py b/tensor2tensor/data_generators/moving_mnist.py
index 257ba7ad8..c80e980fe 100644
--- a/tensor2tensor/data_generators/moving_mnist.py
+++ b/tensor2tensor/data_generators/moving_mnist.py
@@ -35,7 +35,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow_datasets as tfds
 from tensorflow_datasets.video import moving_sequence
 
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index ccbdf1681..92783b361 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -32,7 +32,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import mlperf_log
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 # pylint: disable=g-import-not-at-top
 try:
   from tensorflow.contrib.tpu.python.tpu import tpu_config
diff --git a/tensor2tensor/data_generators/text_encoder_build_subword.py b/tensor2tensor/data_generators/text_encoder_build_subword.py
index fa3bddfda..826072ee8 100644
--- a/tensor2tensor/data_generators/text_encoder_build_subword.py
+++ b/tensor2tensor/data_generators/text_encoder_build_subword.py
@@ -34,7 +34,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import tokenizer
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 tf.flags.DEFINE_string('output_filename', '/tmp/my.subword_text_encoder',
                        'where to store the SubwordTextEncoder')
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 527b0aafa..668a3a5c0 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -32,7 +32,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import mlperf_log
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TranslateProblem(text_problems.Text2TextProblem):
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 3411475e7..5bf66fbc2 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -28,7 +28,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 try:
   import matplotlib  # pylint: disable=g-import-not-at-top
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index b8c11e6df..22ef3545c 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -35,7 +35,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import video_metrics
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
index d528f4779..f85f57f81 100644
--- a/tensor2tensor/data_generators/vqa.py
+++ b/tensor2tensor/data_generators/vqa.py
@@ -39,7 +39,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def _get_vqa_v2_annotations(directory,

From 2d7812b8dadb1e9fb1b82af29bd6aef4cfd4c3e8 Mon Sep 17 00:00:00 2001
From: Rohan Jain <rohanj@google.com>
Date: Tue, 11 Feb 2020 20:45:39 -0800
Subject: [PATCH 2638/2720] Explicitly replace "import tensorflow" with
 "tensorflow.compat.v1" for TF2.x migration

PiperOrigin-RevId: 294583762
---
 tensor2tensor/data_generators/wikisum/generate_vocab.py         | 2 +-
 .../data_generators/wikisum/get_references_commoncrawl.py       | 2 +-
 tensor2tensor/data_generators/wikisum/get_references_web.py     | 2 +-
 tensor2tensor/data_generators/wikisum/parallel_launch.py        | 2 +-
 tensor2tensor/data_generators/wikisum/produce_examples.py       | 2 +-
 tensor2tensor/data_generators/wikisum/validate_data.py          | 2 +-
 tensor2tensor/layers/common_attention.py                        | 2 +-
 tensor2tensor/layers/common_attention_test.py                   | 2 +-
 tensor2tensor/layers/common_layers.py                           | 2 +-
 tensor2tensor/layers/common_layers_test.py                      | 2 +-
 tensor2tensor/layers/common_video.py                            | 2 +-
 tensor2tensor/layers/vqa_layers.py                              | 2 +-
 tensor2tensor/models/video/emily.py                             | 2 +-
 tensor2tensor/models/video/epva.py                              | 2 +-
 tensor2tensor/models/video/next_frame_glow.py                   | 2 +-
 tensor2tensor/models/video/nfg_interpolate.py                   | 2 +-
 tensor2tensor/models/video/savp.py                              | 2 +-
 tensor2tensor/models/video/sv2p.py                              | 2 +-
 18 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/tensor2tensor/data_generators/wikisum/generate_vocab.py b/tensor2tensor/data_generators/wikisum/generate_vocab.py
index 27d03ccd1..e2a61927b 100644
--- a/tensor2tensor/data_generators/wikisum/generate_vocab.py
+++ b/tensor2tensor/data_generators/wikisum/generate_vocab.py
@@ -20,7 +20,7 @@
 
 from tensor2tensor.data_generators.wikisum import wikisum
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py b/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
index 3c316bd28..2749de41f 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
@@ -24,7 +24,7 @@
 from tensor2tensor.data_generators.wikisum import utils
 from tensor2tensor.data_generators.wikisum import wikisum
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/data_generators/wikisum/get_references_web.py b/tensor2tensor/data_generators/wikisum/get_references_web.py
index 3f16787f8..ef098b228 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web.py
@@ -44,7 +44,7 @@
 from tensor2tensor.data_generators.wikisum import get_references_web_single_group as fetch
 from tensor2tensor.data_generators.wikisum import utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 flags = tf.flags
diff --git a/tensor2tensor/data_generators/wikisum/parallel_launch.py b/tensor2tensor/data_generators/wikisum/parallel_launch.py
index bfe0836db..dfdd2305e 100644
--- a/tensor2tensor/data_generators/wikisum/parallel_launch.py
+++ b/tensor2tensor/data_generators/wikisum/parallel_launch.py
@@ -53,7 +53,7 @@
 import time
 
 from tensor2tensor.utils import cloud_mlengine as cloud
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/data_generators/wikisum/produce_examples.py b/tensor2tensor/data_generators/wikisum/produce_examples.py
index 47fa5b1b7..1f3c20ef9 100644
--- a/tensor2tensor/data_generators/wikisum/produce_examples.py
+++ b/tensor2tensor/data_generators/wikisum/produce_examples.py
@@ -24,7 +24,7 @@
 from tensor2tensor.data_generators.wikisum import utils
 from tensor2tensor.data_generators.wikisum import wikisum
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/data_generators/wikisum/validate_data.py b/tensor2tensor/data_generators/wikisum/validate_data.py
index 37295a03b..691ca73fb 100644
--- a/tensor2tensor/data_generators/wikisum/validate_data.py
+++ b/tensor2tensor/data_generators/wikisum/validate_data.py
@@ -25,7 +25,7 @@
 
 from tensor2tensor.data_generators.wikisum import wikisum
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index f2c8b1cf9..1713a23cf 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -34,7 +34,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import expert_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow_probability as tfp
 
 # pylint: disable=g-direct-tensorflow-import
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index d18436904..f11b283e5 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -29,7 +29,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import test_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 tfe = contrib.tfe()
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 0d5717bcc..c0193fbb4 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -28,7 +28,7 @@
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensor2tensor.utils import contrib
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow_probability as tfp
 
 from tensorflow.python.framework import function
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index bcd19b202..d67e49650 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -26,7 +26,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import test_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 tf.compat.v1.enable_eager_execution()
 
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 45a2ccc20..db772ee8b 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -23,7 +23,7 @@
 
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import contrib
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.python.ops import summary_op_util  # pylint: disable=g-direct-tensorflow-import
 
diff --git a/tensor2tensor/layers/vqa_layers.py b/tensor2tensor/layers/vqa_layers.py
index 9d82822b4..4e59e7c90 100644
--- a/tensor2tensor/layers/vqa_layers.py
+++ b/tensor2tensor/layers/vqa_layers.py
@@ -23,7 +23,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.utils import contrib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.contrib import slim
 from tensorflow.contrib.slim.python.slim.nets.resnet_v1 import resnet_v1_152
diff --git a/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index 0a0c58790..e74e1c090 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -35,7 +35,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 tfl = tf.layers
 tfcl = contrib.layers()
diff --git a/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index c6f94f0a4..1234ab7db 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -37,7 +37,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.contrib.framework.python.ops import arg_scope
 from tensorflow.contrib.slim.python.slim.nets import vgg
diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index 73cadd3c3..d1f32b479 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -27,7 +27,7 @@
 from tensor2tensor.models.research import glow_ops
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow_probability as tfp
 
 arg_scope = contrib.framework().arg_scope
diff --git a/tensor2tensor/models/video/nfg_interpolate.py b/tensor2tensor/models/video/nfg_interpolate.py
index 761317f57..d076ad579 100644
--- a/tensor2tensor/models/video/nfg_interpolate.py
+++ b/tensor2tensor/models/video/nfg_interpolate.py
@@ -29,7 +29,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import trainer_lib
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # Flags placeholders.
 flags.DEFINE_string("checkpoint_path", None,
diff --git a/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index 8b06270d4..87b71f2aa 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -32,7 +32,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import update_ops_hook
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow_gan as tfgan
 
 gan_losses = tfgan.losses.wargs
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 10309be08..2a3710cd8 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -34,7 +34,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 tfl = tf.layers
 tfcl = contrib.layers()

From 11f9a63b4297b4de70d8bcdee2f7ed8e332de097 Mon Sep 17 00:00:00 2001
From: Rohan Jain <rohanj@google.com>
Date: Tue, 11 Feb 2020 21:51:14 -0800
Subject: [PATCH 2639/2720] Explicitly replace "import tensorflow" with
 "tensorflow.compat.v1" for TF2.x migration

PiperOrigin-RevId: 294590634
---
 tensor2tensor/models/research/attention_lm.py               | 2 +-
 tensor2tensor/models/research/gene_expression.py            | 2 +-
 tensor2tensor/models/research/glow.py                       | 2 +-
 tensor2tensor/models/research/glow_ops.py                   | 2 +-
 tensor2tensor/models/research/glow_ops_test.py              | 2 +-
 tensor2tensor/models/research/neural_stack.py               | 2 +-
 tensor2tensor/models/research/neural_stack_test.py          | 2 +-
 tensor2tensor/models/research/shuffle_network.py            | 2 +-
 tensor2tensor/models/research/transformer_revnet.py         | 2 +-
 tensor2tensor/models/research/transformer_vae.py            | 2 +-
 tensor2tensor/models/research/universal_transformer.py      | 2 +-
 tensor2tensor/models/research/universal_transformer_util.py | 2 +-
 tensor2tensor/models/research/vqa_attention.py              | 2 +-
 13 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/models/research/attention_lm.py b/tensor2tensor/models/research/attention_lm.py
index cf876b4ba..cf234ce43 100644
--- a/tensor2tensor/models/research/attention_lm.py
+++ b/tensor2tensor/models/research/attention_lm.py
@@ -35,7 +35,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 framework = contrib.framework(msg="warn")
 
diff --git a/tensor2tensor/models/research/gene_expression.py b/tensor2tensor/models/research/gene_expression.py
index 3173140ca..bc66d241e 100644
--- a/tensor2tensor/models/research/gene_expression.py
+++ b/tensor2tensor/models/research/gene_expression.py
@@ -25,7 +25,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index f744e4c39..7799624c8 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -27,7 +27,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 arg_scope = contrib.framework().arg_scope
 add_arg_scope = contrib.framework().add_arg_scope
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index e09159881..54578643a 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -25,7 +25,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 from tensor2tensor.utils import contrib
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow_probability as tfp
 
 arg_scope = contrib.framework().arg_scope
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index f6a39e6e2..9ff4239a8 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -30,7 +30,7 @@
 from tensor2tensor.models.research import glow_ops
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import hparam
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 arg_scope = contrib.framework().arg_scope
 add_arg_scope = contrib.framework().add_arg_scope
diff --git a/tensor2tensor/models/research/neural_stack.py b/tensor2tensor/models/research/neural_stack.py
index 51d27e6c1..5f0fe6790 100644
--- a/tensor2tensor/models/research/neural_stack.py
+++ b/tensor2tensor/models/research/neural_stack.py
@@ -35,7 +35,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # This is the interface between the RNN controller and the neural stack.
 NeuralStackControllerInterface = collections.namedtuple(
diff --git a/tensor2tensor/models/research/neural_stack_test.py b/tensor2tensor/models/research/neural_stack_test.py
index e67070bd6..1365282a7 100644
--- a/tensor2tensor/models/research/neural_stack_test.py
+++ b/tensor2tensor/models/research/neural_stack_test.py
@@ -25,7 +25,7 @@
 from tensor2tensor.models.research import neural_stack
 from tensor2tensor.utils import contrib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def build_fake_controller(cell):
diff --git a/tensor2tensor/models/research/shuffle_network.py b/tensor2tensor/models/research/shuffle_network.py
index a79023c8a..7da420bdc 100644
--- a/tensor2tensor/models/research/shuffle_network.py
+++ b/tensor2tensor/models/research/shuffle_network.py
@@ -32,7 +32,7 @@
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def ror(x, n, p=1):
diff --git a/tensor2tensor/models/research/transformer_revnet.py b/tensor2tensor/models/research/transformer_revnet.py
index 1a50ed3aa..ee0531b8b 100644
--- a/tensor2tensor/models/research/transformer_revnet.py
+++ b/tensor2tensor/models/research/transformer_revnet.py
@@ -24,7 +24,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index f566b3572..244dcbfb4 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -37,7 +37,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 _DO_SUMMARIES = True
diff --git a/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 1352c3b4d..ae8ad3f29 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -35,7 +35,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 @registry.register_model
diff --git a/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index e58bfed5f..b8534f02f 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -57,7 +57,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import expert_utils
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def universal_transformer_encoder(encoder_input,
diff --git a/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index a6f428a88..bb961d24f 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -27,7 +27,7 @@
 # from tensor2tensor.utils import restore_hook
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from tensorflow.contrib import rnn as contrib_rnn
 
 # pylint: disable=unused-import

From c1165f67966b86d9fa304ef8d1b745f70a7b9f75 Mon Sep 17 00:00:00 2001
From: Rohan Jain <rohanj@google.com>
Date: Tue, 11 Feb 2020 22:36:38 -0800
Subject: [PATCH 2640/2720] Explicitly replace "import tensorflow" with
 "tensorflow.compat.v1" for TF2.x migration

PiperOrigin-RevId: 294596106
---
 tensor2tensor/bin/build_vocab.py       | 2 +-
 tensor2tensor/bin/make_tf_configs.py   | 2 +-
 tensor2tensor/bin/t2t_attack.py        | 2 +-
 tensor2tensor/bin/t2t_avg_all.py       | 2 +-
 tensor2tensor/bin/t2t_bleu.py          | 2 +-
 tensor2tensor/bin/t2t_datagen.py       | 2 +-
 tensor2tensor/bin/t2t_decoder.py       | 2 +-
 tensor2tensor/bin/t2t_distill.py       | 2 +-
 tensor2tensor/bin/t2t_eval.py          | 2 +-
 tensor2tensor/bin/t2t_prune.py         | 2 +-
 tensor2tensor/bin/t2t_trainer.py       | 2 +-
 tensor2tensor/bin/t2t_trainer_test.py  | 2 +-
 tensor2tensor/bin/t2t_translate_all.py | 2 +-
 13 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/tensor2tensor/bin/build_vocab.py b/tensor2tensor/bin/build_vocab.py
index 36515c2b2..703bde0f5 100644
--- a/tensor2tensor/bin/build_vocab.py
+++ b/tensor2tensor/bin/build_vocab.py
@@ -30,7 +30,7 @@
 from tensor2tensor import problems as problems_lib  # pylint: disable=unused-import
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/bin/make_tf_configs.py b/tensor2tensor/bin/make_tf_configs.py
index 3d978d79d..4156c2108 100644
--- a/tensor2tensor/bin/make_tf_configs.py
+++ b/tensor2tensor/bin/make_tf_configs.py
@@ -30,7 +30,7 @@
 from __future__ import print_function
 
 import json
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index a674a8104..8a7bce8f0 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -48,7 +48,7 @@
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/bin/t2t_avg_all.py b/tensor2tensor/bin/t2t_avg_all.py
index 608a8705b..986d83cfd 100644
--- a/tensor2tensor/bin/t2t_avg_all.py
+++ b/tensor2tensor/bin/t2t_avg_all.py
@@ -25,7 +25,7 @@
 import six
 from six.moves import zip  # pylint: disable=redefined-builtin
 from tensor2tensor.utils import bleu_hook
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/bin/t2t_bleu.py b/tensor2tensor/bin/t2t_bleu.py
index dd7ef73f8..2eccda85a 100644
--- a/tensor2tensor/bin/t2t_bleu.py
+++ b/tensor2tensor/bin/t2t_bleu.py
@@ -59,7 +59,7 @@
 import os
 import time
 from tensor2tensor.utils import bleu_hook
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 flags = tf.flags
diff --git a/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index 1d5d7c104..c974acdfd 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -53,7 +53,7 @@
   pass
 
 # Improrting here to prevent pylint from ungrouped-imports warning.
-import tensorflow as tf  # pylint: disable=g-import-not-at-top
+import tensorflow.compat.v1 as tf  # pylint: disable=g-import-not-at-top
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index b311011d9..9d6b0bfca 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -41,7 +41,7 @@
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/bin/t2t_distill.py b/tensor2tensor/bin/t2t_distill.py
index 348868f18..7e69ff9a4 100644
--- a/tensor2tensor/bin/t2t_distill.py
+++ b/tensor2tensor/bin/t2t_distill.py
@@ -38,7 +38,7 @@
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/bin/t2t_eval.py b/tensor2tensor/bin/t2t_eval.py
index 24167c1ac..7d2f1385b 100644
--- a/tensor2tensor/bin/t2t_eval.py
+++ b/tensor2tensor/bin/t2t_eval.py
@@ -23,7 +23,7 @@
 from tensor2tensor.data_generators import problem  # pylint: disable=unused-import
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
index d9f0b4939..d1e75f8e2 100644
--- a/tensor2tensor/bin/t2t_prune.py
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -39,7 +39,7 @@
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 4ca0c1f62..424e10378 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -34,7 +34,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.contrib.tpu.python.tpu import tpu_config
 
diff --git a/tensor2tensor/bin/t2t_trainer_test.py b/tensor2tensor/bin/t2t_trainer_test.py
index c7bfbdfbf..0a7fee5a6 100644
--- a/tensor2tensor/bin/t2t_trainer_test.py
+++ b/tensor2tensor/bin/t2t_trainer_test.py
@@ -21,7 +21,7 @@
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.utils import trainer_lib_test
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 FLAGS = tf.flags.FLAGS
 
diff --git a/tensor2tensor/bin/t2t_translate_all.py b/tensor2tensor/bin/t2t_translate_all.py
index 26dad8e7a..4fba3de98 100644
--- a/tensor2tensor/bin/t2t_translate_all.py
+++ b/tensor2tensor/bin/t2t_translate_all.py
@@ -32,7 +32,7 @@
 import shutil
 from tensor2tensor.utils import bleu_hook
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS

From cca1766d367a920f2d6c9fcc21b275b63f0f635b Mon Sep 17 00:00:00 2001
From: Rohan Jain <rohanj@google.com>
Date: Tue, 11 Feb 2020 22:58:14 -0800
Subject: [PATCH 2641/2720] Explicitly replace "import tensorflow" with
 "tensorflow.compat.v1" for TF2.x migration

PiperOrigin-RevId: 294598765
---
 tensor2tensor/utils/avg_checkpoints.py       | 2 +-
 tensor2tensor/utils/bleu_hook.py             | 2 +-
 tensor2tensor/utils/cloud_mlengine.py        | 2 +-
 tensor2tensor/utils/compute_video_metrics.py | 2 +-
 tensor2tensor/utils/data_reader.py           | 2 +-
 tensor2tensor/utils/decoding.py              | 2 +-
 tensor2tensor/utils/devices.py               | 2 +-
 tensor2tensor/utils/flags.py                 | 2 +-
 tensor2tensor/utils/get_rouge.py             | 2 +-
 tensor2tensor/utils/metrics.py               | 2 +-
 tensor2tensor/utils/optimize.py              | 2 +-
 tensor2tensor/utils/restore_hook.py          | 2 +-
 tensor2tensor/utils/t2t_model.py             | 2 +-
 tensor2tensor/utils/trainer_lib.py           | 2 +-
 tensor2tensor/utils/video2gif.py             | 2 +-
 15 files changed, 15 insertions(+), 15 deletions(-)
 mode change 100755 => 100644 tensor2tensor/utils/cloud_mlengine.py

diff --git a/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index 146fa21cf..e4bf366be 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -22,7 +22,7 @@
 import numpy as np
 import six
 from six.moves import zip  # pylint: disable=redefined-builtin
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/utils/bleu_hook.py b/tensor2tensor/utils/bleu_hook.py
index 9f37226a8..d840eb232 100644
--- a/tensor2tensor/utils/bleu_hook.py
+++ b/tensor2tensor/utils/bleu_hook.py
@@ -34,7 +34,7 @@
 
 from tensor2tensor.data_generators import text_encoder
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def _get_ngrams(segment, max_order):
diff --git a/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
old mode 100755
new mode 100644
index 110ca093c..678955df5
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -30,7 +30,7 @@
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import usr_dir as usr_dir_lib
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 FLAGS = tf.flags.FLAGS
 
diff --git a/tensor2tensor/utils/compute_video_metrics.py b/tensor2tensor/utils/compute_video_metrics.py
index 887fd410a..012ca5b82 100644
--- a/tensor2tensor/utils/compute_video_metrics.py
+++ b/tensor2tensor/utils/compute_video_metrics.py
@@ -23,7 +23,7 @@
 
 from tensor2tensor.bin import t2t_decoder
 from tensor2tensor.utils import video_metrics
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 FLAGS = tf.flags.FLAGS
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index 3f9f0e386..bcf1a19e1 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -28,7 +28,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import mlperf_log
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def cast_ints_to_int32(features):
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 91b3476d4..2a7ca2fcb 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -37,7 +37,7 @@
 from tensor2tensor.utils import hparam
 from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import registry
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 FLAGS = tf.flags.FLAGS
 
diff --git a/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index d689bd58d..80cacd20e 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -19,7 +19,7 @@
 from __future__ import print_function
 
 from tensor2tensor.utils import expert_utils as eu
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from tensorflow.python.util import tf_inspect as inspect
 
 
diff --git a/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index ab4ab6e05..9e3dc5ab8 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -19,7 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/utils/get_rouge.py b/tensor2tensor/utils/get_rouge.py
index 3d64acd79..c753c0abd 100644
--- a/tensor2tensor/utils/get_rouge.py
+++ b/tensor2tensor/utils/get_rouge.py
@@ -24,7 +24,7 @@
 import shutil
 from tempfile import mkdtemp
 import pyrouge
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 FLAGS = tf.flags.FLAGS
 
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 2a0570dad..be166a3d8 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -28,7 +28,7 @@
 from tensor2tensor.utils import rouge
 from tensor2tensor.utils import sari_hook
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from tensorflow.python.util import tf_inspect as inspect
 
 
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 689821e40..0e1f79fb9 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -28,7 +28,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import yellowfin
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 from tensorflow.python.framework import dtypes  # pylint: disable=g-direct-tensorflow-import
diff --git a/tensor2tensor/utils/restore_hook.py b/tensor2tensor/utils/restore_hook.py
index 28aa238cc..2346da188 100644
--- a/tensor2tensor/utils/restore_hook.py
+++ b/tensor2tensor/utils/restore_hook.py
@@ -22,7 +22,7 @@
 import six
 
 from tensor2tensor.utils import contrib
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class RestoreHook(tf.train.SessionRunHook):
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 78b5cd77a..1fd95d931 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -46,7 +46,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import scheduled_sampling
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.python.layers import base
 from tensorflow.python.ops import inplace_ops
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 40c0efacb..cfd516c4f 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -35,7 +35,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.core.protobuf import rewriter_config_pb2
 from tensorflow.python import debug
diff --git a/tensor2tensor/utils/video2gif.py b/tensor2tensor/utils/video2gif.py
index f0d52c783..819c5aa39 100644
--- a/tensor2tensor/utils/video2gif.py
+++ b/tensor2tensor/utils/video2gif.py
@@ -41,7 +41,7 @@
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import trainer_lib
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS

From 12ac945c1ece684d85d98d8908b5450827aa16e0 Mon Sep 17 00:00:00 2001
From: Rohan Jain <rohanj@google.com>
Date: Tue, 11 Feb 2020 23:12:45 -0800
Subject: [PATCH 2642/2720] Explicitly replace "import tensorflow" with
 "tensorflow.compat.v1" for TF2.x migration

PiperOrigin-RevId: 294600382
---
 tensor2tensor/envs/env_problem.py                            | 2 +-
 tensor2tensor/envs/rendered_env_problem.py                   | 2 +-
 tensor2tensor/insights/server.py                             | 2 +-
 tensor2tensor/insights/transformer_model.py                  | 2 +-
 tensor2tensor/models/lstm.py                                 | 2 +-
 tensor2tensor/models/neural_architecture_search/nas_model.py | 2 +-
 tensor2tensor/models/revnet.py                               | 2 +-
 tensor2tensor/serving/export.py                              | 2 +-
 tensor2tensor/serving/query.py                               | 2 +-
 tensor2tensor/serving/serving_utils.py                       | 2 +-
 tensor2tensor/utils/video/prediction2gif.py                  | 2 +-
 tensor2tensor/utils/video/reward_confusion.py                | 2 +-
 12 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 68caada1f..0984c4314 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -33,7 +33,7 @@
 from tensor2tensor.envs import trajectory
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import contrib
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # Names for data fields in stored tf.Examples.
 TIMESTEP_FIELD = "timestep"
diff --git a/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
index 69536a095..7eb25c7d8 100644
--- a/tensor2tensor/envs/rendered_env_problem.py
+++ b/tensor2tensor/envs/rendered_env_problem.py
@@ -25,7 +25,7 @@
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import gym_env_problem
 from tensor2tensor.utils import contrib
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 _IMAGE_ENCODED_FIELD = "image/encoded"
 _IMAGE_FORMAT_FIELD = "image/format"
diff --git a/tensor2tensor/insights/server.py b/tensor2tensor/insights/server.py
index fea2b2dbb..44aadef58 100644
--- a/tensor2tensor/insights/server.py
+++ b/tensor2tensor/insights/server.py
@@ -25,7 +25,7 @@
 from gunicorn.six import iteritems
 import numpy as np
 from tensor2tensor.insights import transformer_model
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/insights/transformer_model.py b/tensor2tensor/insights/transformer_model.py
index 660a6c8d3..8a26ca1f0 100644
--- a/tensor2tensor/insights/transformer_model.py
+++ b/tensor2tensor/insights/transformer_model.py
@@ -32,7 +32,7 @@
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from tensorflow.python import debug as tfdbg
 
 flags = tf.flags
diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 60d362d39..156c903f4 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -27,7 +27,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def _dropout_lstm_cell(hparams, train):
diff --git a/tensor2tensor/models/neural_architecture_search/nas_model.py b/tensor2tensor/models/neural_architecture_search/nas_model.py
index 4c72c6b84..c8bcb23fe 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model.py
@@ -39,7 +39,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 # Keys for the activation map.
diff --git a/tensor2tensor/models/revnet.py b/tensor2tensor/models/revnet.py
index 446a6da08..d1072ea68 100644
--- a/tensor2tensor/models/revnet.py
+++ b/tensor2tensor/models/revnet.py
@@ -40,7 +40,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def wrapped_partial(fn, *args, **kwargs):
diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index cc993b83d..0cb7da38f 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -25,7 +25,7 @@
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow_hub as hub
 
 FLAGS = tf.flags.FLAGS
diff --git a/tensor2tensor/serving/query.py b/tensor2tensor/serving/query.py
index b266b8cf2..69a7aadff 100644
--- a/tensor2tensor/serving/query.py
+++ b/tensor2tensor/serving/query.py
@@ -28,7 +28,7 @@
 from tensor2tensor.utils import hparam
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import usr_dir
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS
diff --git a/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 8d5e327e6..a1b437282 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -29,7 +29,7 @@
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.utils import cloud_mlengine as cloud
 from tensor2tensor.utils import contrib
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow_serving.apis import predict_pb2
 from tensorflow_serving.apis import prediction_service_pb2_grpc
diff --git a/tensor2tensor/utils/video/prediction2gif.py b/tensor2tensor/utils/video/prediction2gif.py
index ae59a00d1..45483c0a0 100644
--- a/tensor2tensor/utils/video/prediction2gif.py
+++ b/tensor2tensor/utils/video/prediction2gif.py
@@ -42,7 +42,7 @@
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 mpl.use("Agg")
 flags = tf.flags
diff --git a/tensor2tensor/utils/video/reward_confusion.py b/tensor2tensor/utils/video/reward_confusion.py
index 2234f94f3..431717f3e 100644
--- a/tensor2tensor/utils/video/reward_confusion.py
+++ b/tensor2tensor/utils/video/reward_confusion.py
@@ -36,7 +36,7 @@
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 flags = tf.flags
 FLAGS = flags.FLAGS

From d08409f3e860141c0d0b0669d73cd73b952702c2 Mon Sep 17 00:00:00 2001
From: Jason Lee <jasonylee@google.com>
Date: Thu, 13 Feb 2020 10:48:20 -0800
Subject: [PATCH 2643/2720] First changelist of code to replicate the
 experiments from ``On the Discrepancy between Density Estimation and Sequence
 Generation''.

PiperOrigin-RevId: 294950205
---
 .../layers/transformer_glow_layers.py         |  444 +++++++
 .../layers/transformer_glow_layers_ops.py     |  297 +++++
 .../transformer_glow_layers_ops_test.py       |  103 ++
 .../layers/transformer_glow_layers_test.py    |  333 +++++
 .../research/transformer_vae_flow_prior.py    | 1135 +++++++++++++++++
 .../transformer_vae_flow_prior_ops.py         |  364 ++++++
 6 files changed, 2676 insertions(+)
 create mode 100644 tensor2tensor/layers/transformer_glow_layers.py
 create mode 100644 tensor2tensor/layers/transformer_glow_layers_ops.py
 create mode 100644 tensor2tensor/layers/transformer_glow_layers_ops_test.py
 create mode 100644 tensor2tensor/layers/transformer_glow_layers_test.py
 create mode 100644 tensor2tensor/models/research/transformer_vae_flow_prior.py
 create mode 100644 tensor2tensor/models/research/transformer_vae_flow_prior_ops.py

diff --git a/tensor2tensor/layers/transformer_glow_layers.py b/tensor2tensor/layers/transformer_glow_layers.py
new file mode 100644
index 000000000..6cfece84e
--- /dev/null
+++ b/tensor2tensor/layers/transformer_glow_layers.py
@@ -0,0 +1,444 @@
+# coding=utf-8
+# Copyright 2020 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Glow operations for text.
+
+Adapted glow operations from tensor2tensor.models.research.glow_ops to be used
+as a prior in Text VAEs (specifically for MT). Supports:
+1. Log determinant Jacobian computation with variable length data and masking.
+2. Transformer instead of convolution as a basic transformation.
+3. Every transformation (affine, split) conditions on the source
+  sentence.
+4. Three different split functions in affine coupling.
+5. Multi-head 1x1 convolution.
+6. Actnorm with weight normalization.
+
+Implementation based on Ma et al., 2019: https://arxiv.org/pdf/1909.02480.pdf
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+import numpy as np
+import scipy
+from tensor2tensor.layers import common_layers
+import tensor2tensor.layers.transformer_glow_layers_ops as gops
+import tensorflow.compat.v1 as tf
+
+
+def actnorm(name, x, x_mask, inverse, init, logscale_factor=3.0):
+  """Activation normalization, returns logabsdet of shape [B]."""
+  eps = tf.keras.backend.epsilon()
+  n_channels = common_layers.shape_list(x)[2]
+
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    x_mean, x_var = gops.moments_over_bl(x, x_mask)
+    b = gops.get_variable_ddi(
+        "b", (n_channels), -x_mean, init, tf.zeros_initializer)
+    log_w_init = -0.5 * tf.log(x_var + eps) / logscale_factor
+    log_w = gops.get_variable_ddi(
+        "log_w", (n_channels), log_w_init, init,
+        tf.zeros_initializer) * logscale_factor
+
+    if not inverse:
+      x = (x + b) * tf.exp(log_w)
+    else:
+      x = x * tf.exp(-log_w) - b
+
+    x_length = tf.reduce_sum(x_mask, -1)
+    logabsdet = x_length * tf.reduce_sum(log_w)
+    if inverse:
+      logabsdet *= -1
+    return x, logabsdet
+
+
+def multihead_invertible_1x1_conv_np(
+    name, x, x_mask, multihead_split, inverse, dtype):
+  """Multi-head 1X1 convolution on x."""
+  batch_size, length, n_channels_all = common_layers.shape_list(x)
+  assert n_channels_all % 32 == 0
+  n_channels = 32
+  n_1x1_heads = n_channels_all // n_channels
+
+  def get_init_np():
+    """Initializer function for multihead 1x1 parameters using numpy."""
+    results = []
+    for _ in range(n_1x1_heads):
+      random_matrix = np.random.rand(n_channels, n_channels)
+      np_w = scipy.linalg.qr(random_matrix)[0].astype("float32")
+      np_p, np_l, np_u = scipy.linalg.lu(np_w)
+      np_s = np.diag(np_u)
+      np_sign_s = np.sign(np_s)[np.newaxis, :]
+      np_log_s = np.log(np.abs(np_s))[np.newaxis, :]
+      np_u = np.triu(np_u, k=1)
+      results.append(
+          np.concatenate([np_p, np_l, np_u, np_sign_s, np_log_s], axis=0))
+    return tf.convert_to_tensor(np.stack(results, axis=0))
+
+  def get_mask_init():
+    ones = tf.ones([n_1x1_heads, n_channels, n_channels], dtype=dtype)
+    l_mask = tf.matrix_band_part(ones, -1, 0) - tf.matrix_band_part(ones, 0, 0)
+    u_mask = tf.matrix_band_part(ones, 0, -1) - tf.matrix_band_part(ones, 0, 0)
+    return tf.stack([l_mask, u_mask], axis=0)
+
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    params = tf.get_variable("params", initializer=get_init_np, dtype=dtype)
+    mask_params = tf.get_variable(
+        "mask_params", initializer=get_mask_init, dtype=dtype, trainable=False)
+
+    p = tf.stop_gradient(params[:, :n_channels, :])
+    l = params[:, n_channels : 2*n_channels, :]
+    u = params[:, 2*n_channels : 3*n_channels, :]
+    sign_s = tf.stop_gradient(params[:, 3*n_channels, :])
+    log_s = params[:, 3*n_channels+1, :]
+
+    l_mask = mask_params[0]
+    u_mask = mask_params[1]
+
+    l_diag = l * l_mask + (
+        tf.eye(n_channels, n_channels, [n_1x1_heads], dtype=dtype))
+    u_diag = u * u_mask + (
+        tf.matrix_diag(sign_s * tf.exp(log_s)))
+    w = tf.matmul(p, tf.matmul(l_diag, u_diag))
+
+    if multihead_split == "a":
+      x = tf.reshape(x, [batch_size, length, n_channels, n_1x1_heads])
+      x = tf.transpose(x, [3, 0, 1, 2])
+    elif multihead_split == "c":
+      x = tf.reshape(x, [batch_size, length, n_1x1_heads, n_channels])
+      x = tf.transpose(x, [2, 0, 1, 3])
+    else:
+      raise ValueError("Multihead split not supported.")
+    # [n_1x1_heads, batch_size, length, n_channels]
+
+    if not inverse:
+      # [n_1x1_heads, 1, n_channels, n_channels]
+      x = tf.matmul(x, w[:, tf.newaxis, :, :])
+    else:
+      w_inv = tf.matrix_inverse(w)
+      x = tf.matmul(x, w_inv[:, tf.newaxis, :, :])
+
+    if multihead_split == "a":
+      x = tf.transpose(x, [1, 2, 3, 0])
+      x = tf.reshape(x, [batch_size, length, n_channels * n_1x1_heads])
+    elif multihead_split == "c":
+      x = tf.transpose(x, [1, 2, 0, 3])
+      x = tf.reshape(x, [batch_size, length, n_1x1_heads * n_channels])
+    else:
+      raise ValueError("Multihead split not supported.")
+
+    x_length = tf.reduce_sum(x_mask, -1)
+    logabsdet = x_length * tf.reduce_sum(log_s)
+    if inverse:
+      logabsdet *= -1
+  return x, logabsdet
+
+
+def coupling(*args, **kwargs):
+  """Coupling transform layer."""
+  prior_type = kwargs["hparams"].prior_type
+  posterior_type = kwargs["hparams"].posterior_type
+  if prior_type == "affine" or posterior_type == "affine":
+    return affine_coupling(*args, **kwargs)
+  elif prior_type == "additive" or posterior_type == "additive":
+    return additive_coupling(*args, **kwargs)
+
+
+def additive_coupling(
+    name, x, x_mask, inverse, split_dim, identity_first, init,
+    decoder_self_attention_bias=None, **kwargs):
+  """Additive coupling transform layer."""
+  hparams = kwargs["hparams"]
+  batch_size, length, n_channels = common_layers.shape_list(x)
+  assert hparams.scale_width > 0.0 and hparams.scale_width < 1.0
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    x_id, x_tr, _, n_transform, bias, mask = gops.split_coupling(
+        x, x_mask, split_dim, identity_first, decoder_self_attention_bias)
+    z_id = x_id
+
+    loc = gops.transformer_decoder_block(
+        "theta_tr",
+        n_layers=hparams.n_layers_transform_params,
+        x=x_id,
+        x_mask=mask,
+        output_size=n_transform,
+        init=init,
+        decoder_self_attention_bias=bias,
+        **kwargs)
+    if not inverse:
+      z_tr = x_tr + loc
+    else:
+      z_tr = x_tr - loc
+    logabsdet = tf.constant(0.0, dtype=tf.float32)
+
+    tf.summary.histogram("_loc", tf.boolean_mask(loc, mask))
+    result = gops.join_coupling(z_id, z_tr, split_dim, identity_first)
+    result = tf.reshape(result, [batch_size, length, n_channels])
+    return result, logabsdet
+
+
+def affine_coupling(
+    name, x, x_mask, inverse, split_dim, identity_first, init,
+    decoder_self_attention_bias=None, **kwargs):
+  """Affine coupling transform layer.
+
+  Args:
+    name: variable scope.
+    x: 3-D Tensor, shape=[B, L, C].
+    x_mask : 2-D Tensor, shape=[B, L].
+    inverse: Forward or inverse pass.
+    split_dim: which dimension to split
+      (time, channel_continuous, channel_alternate).
+    identity_first: True means the first half remains constant. False for 2nd.
+    init: init.
+    decoder_self_attention_bias: bias.
+    **kwargs: additional arguments. Contains hparams, encoder_output and
+      encoder_decoder_attention_bias.
+
+  Returns:
+    z: data transformed by the affine coupling layer. shape=[B, L, C]
+    logabsdets: Log absolute determinant Jacobian. shape=[B]
+  """
+  hparams = kwargs["hparams"]
+  batch_size, length, n_channels = common_layers.shape_list(x)
+  assert hparams.scale_width > 0.0 and hparams.scale_width < 1.0
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    x_id, x_tr, _, n_transform, bias, mask = gops.split_coupling(
+        x, x_mask, split_dim, identity_first, decoder_self_attention_bias)
+    z_id = x_id
+
+    transform_params = gops.transformer_decoder_block(
+        "theta_tr",
+        n_layers=hparams.n_layers_transform_params,
+        x=x_id,
+        x_mask=mask,
+        output_size=n_transform*2,
+        init=init,
+        decoder_self_attention_bias=bias,
+        **kwargs)
+    loc, unconstrained_scale = tf.split(transform_params, 2, axis=-1)
+    scale = tf.sigmoid(unconstrained_scale + 2.0)
+    if not inverse:
+      z_tr = (x_tr + loc) * scale
+    else:
+      z_tr = x_tr / scale - loc
+
+    logabsdet = gops.reduce_sum_over_lc(tf.log(scale), mask)  # [B]
+    if inverse:
+      logabsdet *= -1
+
+    tf.summary.histogram("_loc", tf.boolean_mask(loc, mask))
+    tf.summary.histogram("_scale", tf.boolean_mask(scale, mask))
+    result = gops.join_coupling(z_id, z_tr, split_dim, identity_first)
+    result = tf.reshape(result, [batch_size, length, n_channels])
+    return result, logabsdet
+
+
+def flow_step_glow(name, x, x_mask, split_dims, inverse, init, dtype, **kwargs):
+  """One step of flow."""
+  conv_fn = multihead_invertible_1x1_conv_np
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    reversible_ops = []
+    for _, split_dim in enumerate(split_dims):
+      identity_first = True
+      reversible_ops += [functools.partial(actnorm, name="actnorm", init=init)]
+      if split_dim in "ca":
+        multihead_split = "a" if split_dim == "c" else "c"
+        reversible_ops += [functools.partial(
+            conv_fn, name="conv_{}".format(multihead_split),
+            multihead_split=multihead_split, dtype=dtype)]
+      reversible_ops += [functools.partial(
+          coupling, name="coupling_{}".format(split_dim),
+          split_dim=split_dim, identity_first=identity_first, init=init,
+          **kwargs)]
+    if inverse:
+      reversible_ops = reversible_ops[::-1]
+
+    logabsdets = tf.constant(0.0, dtype=dtype)
+    for reversible_op in reversible_ops:
+      x, logabsdet = reversible_op(x=x, x_mask=x_mask, inverse=inverse)
+      logabsdets += logabsdet
+    return x, logabsdets
+
+
+def flow_level(
+    name, x, x_mask, depth, split_dims, prior, inverse, init, dtype, **kwargs):
+  """One level of flow."""
+  flow_step_fn = flow_step_glow
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    reversible_ops = []
+    for step in np.arange(depth):
+      reversible_ops += [functools.partial(
+          flow_step_fn, name="{}_step".format(step), split_dims=split_dims,
+          init=init, dtype=dtype, **kwargs)]
+    if prior:
+      reversible_ops += [functools.partial(
+          coupling, name="{}_prior".format(depth), split_dim="c",
+          identity_first=True, init=init, **kwargs)]
+    if inverse:
+      reversible_ops = reversible_ops[::-1]
+
+    logabsdets = tf.constant(0.0, dtype=dtype)
+    for reversible_op in reversible_ops:
+      x, logabsdet = reversible_op(x=x, x_mask=x_mask, inverse=inverse)
+      logabsdets += logabsdet
+    return x, logabsdets
+
+
+def split(name, x, x_mask, inverse, temp=1.0, dtype=tf.float32, z=None):
+  """Splits / concatenates x into x1 and x2 across number of channels.
+
+  x2 is modelled with a standard gaussian distribution.
+  Args:
+    name: variable scope.
+    x: 3-D Tensor, shape=[B, L, C].
+    x_mask: 2-D Tensor, shape=[B, L].
+    inverse: forward or inverse pass.
+    temp: only used for inverse pass. temperature for sampling.
+    dtype: dtype
+    z: used in inverse pass to check invertibility.
+
+  Returns:
+    x: if forward, returns the 1st half of the channel dimensions.
+      if inverse, return concat[input, N(0,1)]
+    z: second half of the channel dimensions. modelled as standard normal.
+    log_p: log p(x2; N(0,1)), shape=[B]
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    if not inverse:
+      x1, x2 = tf.split(x, 2, axis=-1)
+      log_p = gops.standard_normal_density(x2, x_mask)
+      return x1, x2, log_p
+    else:
+      if z is None:
+        x2 = tf.random.normal(
+            common_layers.shape_list(x), stddev=temp, dtype=dtype)
+      else:
+        x2 = z
+      log_p = gops.standard_normal_density(x2, x_mask)
+      return tf.concat([x, x2], 2), None, log_p
+
+
+def squeeze(name, x, factor, inverse):
+  """Temporal squeezing of x to increase the number of channels."""
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    if factor == 1:
+      return x
+    batch_size, length, n_channels = common_layers.shape_list(x)
+    if not inverse:
+      x = tf.reshape(x, [batch_size, length//factor, factor, n_channels])
+      # transposing groups neighbouring elements together.
+      x = tf.transpose(x, [0, 1, 3, 2])
+      x = tf.reshape(x, [batch_size, length//factor, n_channels*factor])
+    else:
+      x = tf.reshape(x, (batch_size, length, n_channels//factor, factor))
+      x = tf.transpose(x, [0, 1, 3, 2])
+      x = tf.reshape(x, (batch_size, length*factor, n_channels//factor))
+    return x
+
+
+def glow(
+    name, x, max_x_mask, max_self_attn_bias, inverse, init, dtype=tf.float32,
+    split_zs=None, temp=1.0, **kwargs):
+  """Multi-scale glow model. Flow + (n_levels-1)*(Split + Squeeze + Flow).
+
+  Note the original glow's ordering is Squeeze + Flow + Split.
+
+  Args:
+    name: variable scope.
+    x: 3-D Tensor, shape=[B, L, C]. The length dimension is padded to the
+      closest multiple of factor**n_levels.
+    max_x_mask : 2-D Tensor, shape=[B, L]. Binary mask indicating padding.
+    max_self_attn_bias : 4-D Tensor, shape=[B, 1, 1, L].
+    inverse: forward or inverse pass.
+    init: init.
+    dtype: dtype.
+    split_zs: intermediate latents modelled as a standard normal.
+    temp: Only used in inverse. Temperature for sampling.
+    **kwargs: additional arguments. Contains hparams, disable_dropout,
+      encoder_output and encoder_decoder_attention_bias.
+
+  Returns:
+    x: if forward, data transformed to the base distribution.
+      if inverse, base transformed to the data (latent) distribution.
+    logabsdets: log absolute determinant Jacobian. [B]
+    log_ps: log probability in the base distribution. [B]
+    split_zs: all intermediate latents (only used to check invertibility.)
+  """
+  assert x.shape.rank == 3
+  hparams = kwargs["hparams"]
+  factor = hparams.factor
+  if hparams.depths:
+    depths = [int(depth_str) for depth_str in hparams.depths.split("/")]
+  else:
+    depths = []
+  split_plans = hparams.split_plans.split("/")
+  n_levels = len(depths)
+  logabsdets = tf.constant(0.0, dtype=dtype)
+  log_ps = tf.constant(0.0, dtype=dtype)
+  with tf.variable_scope(name, use_resource=True, reuse=tf.AUTO_REUSE):
+    if not inverse:  # z -> e (density estimation)
+      x_mask, self_attn_bias = max_x_mask, max_self_attn_bias
+      split_zs = []
+      for level in range(n_levels):
+        if level > 0:
+          x, z, log_p_z = split(
+              "{}_split".format(level), x, x_mask, inverse, dtype)
+          log_ps += log_p_z
+          split_zs.append(z)
+
+          x = squeeze("{}_squeeze".format(level), x, factor, inverse)
+          x_mask = max_x_mask[:, ::factor**level]
+          self_attn_bias = max_self_attn_bias[..., ::factor**level]
+
+        prior = level < n_levels - 1
+        x, logabsdet = flow_level(
+            "{}_level".format(level), x, x_mask, depths[level],
+            split_plans[level], prior, inverse, init, dtype,
+            decoder_self_attention_bias=self_attn_bias, **kwargs)
+        logabsdets += logabsdet  # (B)
+
+      log_p_base = gops.standard_normal_density(x, x_mask)
+      log_ps += log_p_base
+      return x, logabsdets, log_ps, split_zs
+
+    else:  # e -> z (sampling)
+      x_mask = max_x_mask[:, ::factor**(n_levels-1)]
+      log_p_base = gops.standard_normal_density(x, x_mask)
+      log_ps += log_p_base
+      if split_zs is None:
+        split_zs = [None] * (n_levels-1)
+
+      for level in reversed(range(n_levels)):
+        x_mask = max_x_mask[:, ::factor**level]
+        self_attn_bias = max_self_attn_bias[..., ::factor**level]
+        prior = level < n_levels - 1
+        x, logabsdet = flow_level(
+            "{}_level".format(level), x, x_mask, depths[level],
+            split_plans[level], prior, inverse, init, dtype,
+            decoder_self_attention_bias=self_attn_bias, **kwargs)
+        logabsdets += logabsdet
+
+        if level > 0:
+          x = squeeze("{}_squeeze".format(level), x, factor, inverse)
+          x_mask = max_x_mask[:, ::factor**(level-1)]
+          x, _, log_p_z = split(
+              "{}_split".format(level), x, x_mask, inverse, temp=temp,
+              dtype=dtype, z=split_zs[level-1])
+          log_ps += log_p_z
+
+      return x, logabsdets, log_ps, None
diff --git a/tensor2tensor/layers/transformer_glow_layers_ops.py b/tensor2tensor/layers/transformer_glow_layers_ops.py
new file mode 100644
index 000000000..ee1acd850
--- /dev/null
+++ b/tensor2tensor/layers/transformer_glow_layers_ops.py
@@ -0,0 +1,297 @@
+# coding=utf-8
+# Copyright 2020 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Additional operations for transformer_glow_layers.py."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import copy
+import math
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_layers
+from tensor2tensor.models.transformer import transformer_decoder_layer
+import tensorflow.compat.v1 as tf
+import tensorflow_probability as tfp
+
+
+def dense(name, x, n_out, dtype=tf.float32, init_w=0.05):
+  """Dense layer."""
+  n_in = common_layers.shape_list(x)[2]
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    w = tf.get_variable(
+        "w", [n_in, n_out], dtype,
+        initializer=tf.random_normal_initializer(0.0, init_w), trainable=True)
+    b = tf.get_variable(
+        "b", [n_out,], dtype, initializer=tf.zeros_initializer, trainable=True)
+    x = tf.matmul(x, w) + b
+    return x
+
+
+def dense_weightnorm(
+    name, x, n_out, x_mask, init_scale, init, dtype=tf.float32):
+  """Dense layer with weight normalization."""
+  n_in = common_layers.shape_list(x)[2]
+  eps = tf.keras.backend.epsilon()
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    v = tf.get_variable(
+        "v", [n_in, n_out], dtype,
+        initializer=tf.random_normal_initializer(0, 0.05), trainable=True)
+    v = v / tf.norm(v, axis=0, keepdims=True)
+    t = tf.matmul(x, v)  # [B, L, n_out]
+    mean, var = moments_over_bl(t, x_mask)
+    g_init = init_scale / (tf.sqrt(var) + eps)
+    g = get_variable_ddi(
+        "g", [n_out], g_init, init,
+        initializer=tf.zeros_initializer, dtype=dtype, trainable=True)
+    b = get_variable_ddi(
+        "b", [n_out], -mean*g_init, init,
+        initializer=tf.zeros_initializer, dtype=dtype, trainable=True)
+    w = g * v
+    y = tf.matmul(x, w) + b
+    tf.summary.histogram("_g", g)
+    return y
+
+
+def transformer_decoder_block(name,
+                              n_layers,
+                              x,
+                              x_mask,
+                              output_size,
+                              init,
+                              **kwargs):
+  """A transformation block composed of transformer decoder layers.
+
+  Args:
+    name: variable scope.
+    n_layers: number of transformer layers.
+    x: input to transformation.
+    x_mask: mask.
+    output_size: output dimensionality.
+    init: data-dependent init for weightnorm parameters.
+    **kwargs: Constains hparams, encoder_output,
+      encoder_decoder_attention_bias and decoder_self_attention_bias
+
+  Returns:
+    outputs: Tensor of shape [batch_size, length, output_size].
+  """
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    hparams = kwargs.pop("hparams")
+    disable_dropout = kwargs.pop("disable_dropout")
+    if disable_dropout:
+      hparams = copy.deepcopy(hparams)
+      hparams.attention_dropout = 0.0
+      hparams.layer_prepostprocess_dropout = 0.0
+      hparams.relu_dropout = 0.0
+    n_channels = common_layers.shape_list(x)[-1]
+    if n_channels != hparams.hidden_size:
+      hparams = copy.deepcopy(hparams)
+      hparams.hidden_size = n_channels
+
+    outputs = common_attention.add_timing_signal_1d(x)
+    with tf.variable_scope("decoder", reuse=tf.AUTO_REUSE):
+      for layer_idx in range(n_layers):
+        outputs = transformer_decoder_layer(
+            decoder_input=outputs,
+            layer_idx=layer_idx,
+            hparams=hparams,
+            **kwargs)
+    outputs = common_layers.layer_preprocess(outputs, hparams)
+    outputs = dense_weightnorm(
+        "h2o", outputs, output_size, x_mask, init_scale=0.0, init=init)
+    return outputs
+
+
+def reduce_sum_over_lc(x, x_mask):
+  """Returns sum of x (over L and C) given the actual length and pad.
+
+  Args:
+    x: input. (B,L,C)
+    x_mask: binary padding mask. (B,L)
+
+  Returns:
+    sum of x. (B)
+  """
+
+  if x.shape.rank == 3 and x_mask.shape.rank == 2:
+    x_mask = x_mask[..., tf.newaxis]
+  else:
+    tf.logging.info("x: {}, x_mask: {}".format(x.shape.rank, x_mask.shape.rank))
+    raise ValueError("Dimension not supported.")
+
+  mean = x * x_mask
+  return tf.reduce_sum(mean, axis=[1, 2])  # sum over L, C
+
+
+def reduce_sum_over_l(x, x_mask):
+  """Returns sum of x (over L) given the actual length and pad.
+
+  Args:
+    x: input. (B,L,C)
+    x_mask: binary padding mask. (B,L)
+
+  Returns:
+    sum of x. (B,C)
+  """
+
+  if x.shape.rank == 3 and x_mask.shape.rank == 2:
+    x_mask = x_mask[..., tf.newaxis]
+  else:
+    tf.logging.info("x: {}, x_mask: {}".format(x.shape.rank, x_mask.shape.rank))
+    raise ValueError("Dimension not supported.")
+
+  mean = x * x_mask
+  return tf.reduce_sum(mean, axis=-2)  # sum over L
+
+
+def reduce_mean_over_l(x, x_mask):
+  """Returns mean of x (over L) given the actual length and pad."""
+  return reduce_sum_over_l(x, x_mask) / tf.reduce_sum(x_mask, 1, keepdims=True)
+
+
+def reduce_mean_over_bl(x, x_mask):
+  """Returns average of x (over B and L) given the actual length and pad.
+
+  Args:
+    x: input. (B,L,C)
+    x_mask: binary padding mask. (B,L)
+
+  Returns:
+    mean of x. (C)
+  """
+
+  if x.shape.rank == 3 and x_mask.shape.rank == 2:
+    x_mask = x_mask[..., tf.newaxis]
+  else:
+    tf.logging.info("x: {}, x_mask: {}".format(x.shape.rank, x_mask.shape.rank))
+    raise ValueError("Dimension not supported.")
+
+  mean = x * x_mask
+  mean = tf.reduce_sum(mean, axis=[0, 1])  # sum over B, L
+  return mean / tf.reduce_sum(x_mask)
+
+
+def reduce_mean_over_l_sum_over_c(x, x_mask):
+  """Returns mean of x over L and sum over C."""
+  mean = reduce_sum_over_lc(x, x_mask)
+  return mean / tf.reduce_sum(x_mask, 1)
+
+
+def reduce_mean_over_bl_sum_over_c(x, x_mask):
+  """Returns mean of x over B and L and sum over C."""
+  mean = reduce_mean_over_bl(x, x_mask)
+  return tf.reduce_sum(mean)
+
+
+def moments_over_bl(x, x_mask):
+  """Returns mean and var of x over B and L."""
+  mean = reduce_mean_over_bl(x, x_mask)
+  var = reduce_mean_over_bl((x-mean)**2, x_mask)
+  return mean, var
+
+
+def standard_normal_density(x, x_mask, reduce_sum=False):
+  """Return standard normal distribution with same shape as x."""
+  log_probs = -0.5 * (x**2 + math.log(math.pi * 2.0))
+  if reduce_sum:
+    log_probs = reduce_mean_over_bl_sum_over_c(log_probs, x_mask)
+  else:
+    log_probs = reduce_sum_over_lc(log_probs, x_mask)
+  return log_probs
+
+
+def standard_normal(x, name="normal"):
+  """Return standard normal distribution with same shape as x."""
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    dist = tfp.distributions.Normal(
+        loc=tf.zeros_like(x),
+        scale=tf.ones_like(x),
+        allow_nan_stats=False)
+    return dist
+
+
+def diagonal_normal(outputs, name="normal"):
+  """Split outputs into mu and log_sigma and return z."""
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    loc, log_scale = tf.split(outputs, 2, axis=-1)
+    scale = tf.exp(log_scale)
+    dist = tfp.distributions.Normal(
+        loc=loc,
+        scale=scale + tf.keras.backend.epsilon(),
+        allow_nan_stats=False)
+    return dist
+
+
+def split_coupling(
+    x, x_mask, split_dim, identity_first, decoder_self_attention_bias):
+  """Split function used in coupling flows."""
+  n_channels = common_layers.shape_list(x)[-1]
+  if split_dim == "c":
+    n_transform = n_identity = n_channels // 2
+    x_id = x[..., :n_identity] if identity_first else x[..., n_transform:]
+    x_tr = x[..., n_identity:] if identity_first else x[..., :n_transform]
+    bias, mask = decoder_self_attention_bias, x_mask
+
+  elif split_dim == "a":
+    n_transform = n_identity = n_channels // 2
+    x_id = x[..., 0::2] if identity_first else x[..., 1::2]
+    x_tr = x[..., 1::2] if identity_first else x[..., 0::2]
+    bias, mask = decoder_self_attention_bias, x_mask
+
+  elif split_dim == "t":
+    n_transform = n_identity = n_channels
+    x_id = x[:, 0::2, :] if identity_first else x[:, 1::2, :]
+    x_tr = x[:, 1::2, :] if identity_first else x[:, 0::2, :]
+    bias, mask = decoder_self_attention_bias[..., 0::2], x_mask[..., 0::2]
+
+  return x_id, x_tr, n_identity, n_transform, bias, mask
+
+
+def join_coupling(z_id, z_tr, split_dim, identity_first):
+  """Reverse split function used in coupling flows."""
+  assert z_id.shape.rank == 3 and z_tr.shape.rank == 3
+  result = [z_id, z_tr] if identity_first else [z_tr, z_id]
+  if split_dim == "c":
+    result = tf.concat(result, axis=2)  # concat in the channel dimension
+  elif split_dim == "a":
+    result = tf.stack(result, axis=3)  # stack in the channel dimension
+  elif split_dim == "t":
+    result = tf.stack(result, axis=2)  # stack in the time dimension
+  return result
+
+
+def assign(w, initial_value):
+  w = w.assign(initial_value)
+  with tf.control_dependencies([w]):
+    return w
+
+
+def get_variable_ddi(
+    name, shape, value, init, initializer=None, dtype=tf.float32,
+    regularizer=None, trainable=True):
+  """Wrapper for data-dependent initialization."""
+  kwargs = {"trainable": trainable}
+  if initializer:
+    kwargs["initializer"] = initializer
+  if regularizer:
+    kwargs["regularizer"] = regularizer
+  w = tf.get_variable(name, shape, dtype, **kwargs)
+  if isinstance(init, bool):
+    if init:
+      return assign(w, value)
+    return w
+  else:
+    return tf.cond(init, lambda: assign(w, value), lambda: w)
diff --git a/tensor2tensor/layers/transformer_glow_layers_ops_test.py b/tensor2tensor/layers/transformer_glow_layers_ops_test.py
new file mode 100644
index 000000000..fbf862899
--- /dev/null
+++ b/tensor2tensor/layers/transformer_glow_layers_ops_test.py
@@ -0,0 +1,103 @@
+# coding=utf-8
+# Copyright 2020 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.layers.transformer_flow_ops."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensor2tensor.layers import transformer_glow_layers_ops as gops
+from tensor2tensor.models import transformer
+import tensorflow.compat.v1 as tf
+
+BATCH_SIZE = 10
+INPUT_LENGTH = 3
+TARGET_LENGTH = 16
+N_CHANNELS = 24
+HIDDEN_SIZE = 64
+N_1X1_HEADS = 4
+
+
+class TransformerFlowOpsTest(parameterized.TestCase, tf.test.TestCase):
+
+  def get_data(self):
+    x = tf.random_normal((BATCH_SIZE, TARGET_LENGTH, N_CHANNELS),
+                         mean=0.0, stddev=1.0)
+    x_lengths = np.random.randint(low=1, high=TARGET_LENGTH+1, size=BATCH_SIZE)
+    x_mask = tf.sequence_mask(x_lengths, maxlen=TARGET_LENGTH, dtype=tf.float32)
+    return x, x_mask
+
+  def get_hparams(self):
+    hparams = transformer.transformer_small()
+    hparams.add_hparam("prior_type", "affine")
+    hparams.add_hparam("depths", "12")  # infer n_levels from depths
+    hparams.add_hparam("split_plans", "tca")
+    hparams.add_hparam("factor", 2)  # squeezing factor
+    hparams.add_hparam("n_layers_transform_params", 1)
+    hparams.add_hparam("n_layers_multiscale_prior", 3)
+    hparams.add_hparam("flow_num_heads", 4)
+    hparams.add_hparam("flow_num_1x1_heads", N_1X1_HEADS)
+    hparams.add_hparam("flow_hidden_size", 64)
+    hparams.add_hparam("flow_filter_size", 128)
+    hparams.add_hparam("cond_prior_on_src", True)
+    hparams.add_hparam("bottom_prior_std", False)
+    hparams.add_hparam("latent_size", N_CHANNELS)
+    hparams.add_hparam("scale_width", 0.999)
+    hparams.add_hparam("coupling_transform_ratio", 0.5)
+    hparams.add_hparam("actnorm_type", "actnorm")
+    hparams.add_hparam("actnorm_weightnorm", True)
+    hparams.add_hparam("perm_type", "1x1")
+    hparams.add_hparam("init_permutation", True)
+    hparams.causal_decoder_self_attention = False
+    hparams.hidden_size = HIDDEN_SIZE
+    return hparams
+
+  def get_kwargs(self, hparams=None):
+    if hparams is None:
+      hparams = self.get_hparams()
+    encoder_output = tf.random.uniform(
+        (BATCH_SIZE, INPUT_LENGTH, HIDDEN_SIZE))
+    encoder_decoder_attention_bias = tf.random.uniform(
+        (BATCH_SIZE, 1, 1, INPUT_LENGTH))
+    decoder_self_attention_bias = tf.random.uniform(
+        (BATCH_SIZE, 1, 1, TARGET_LENGTH))
+    kwargs = {"hparams": hparams,
+              "encoder_output": encoder_output,
+              "encoder_decoder_attention_bias": encoder_decoder_attention_bias,
+              "decoder_self_attention_bias": decoder_self_attention_bias}
+    return kwargs
+
+  def test_dense_weightnorm(self):
+    x, x_mask = self.get_data()
+    x = tf.random_normal((BATCH_SIZE, TARGET_LENGTH, HIDDEN_SIZE),
+                         mean=0.0, stddev=1.0)
+    y = gops.dense_weightnorm("wn", x, N_CHANNELS, x_mask,
+                              init_scale=1.0, init=True)
+
+    y_nopad = tf.boolean_mask(y, x_mask)
+    mean, var = tf.nn.moments(y_nopad, axes=[0])
+    self.evaluate(tf.global_variables_initializer())
+    x, x_mask, y, y_nopad, mean, var = (
+        self.evaluate([x, x_mask, y, y_nopad, mean, var]))
+    self.assertEqual(y.shape, (BATCH_SIZE, TARGET_LENGTH, N_CHANNELS))
+    self.assertTrue(np.allclose(mean, 0.0, atol=1e-5))
+    self.assertTrue(np.allclose(var, 1.0, atol=1e-5))
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/layers/transformer_glow_layers_test.py b/tensor2tensor/layers/transformer_glow_layers_test.py
new file mode 100644
index 000000000..6b20a21b8
--- /dev/null
+++ b/tensor2tensor/layers/transformer_glow_layers_test.py
@@ -0,0 +1,333 @@
+# coding=utf-8
+# Copyright 2020 The Tensor2Tensor Authors.
+#
+# 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 tensor2tensor.layers.transformer_glow_layers.
+
+1. Actnorm test (zero mean and unit variance).
+2. Invertibility tests for:
+  * actnorm
+  * actnorm with weight normalization
+  * 1x1 invertible convolution
+  * multi-head 1x1 invertible convolution
+  * affine coupling
+  * split
+  * 1 step of flow
+  * k steps of flow
+  * entire pipeline (tested up to 3 levels, 32 steps: tca/tca/ca, 12/12/8)
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import tempfile
+from absl.testing import parameterized
+import numpy as np
+
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import transformer_glow_layers as glow
+from tensor2tensor.layers import transformer_glow_layers_ops as gops
+from tensor2tensor.models import transformer
+import tensorflow.compat.v1 as tf
+
+BATCH_SIZE = 20
+INPUT_LENGTH = 3
+TARGET_LENGTH = 16
+N_CHANNELS = 256
+HIDDEN_SIZE = 64
+N_1X1_HEADS = 4
+DTYPE = tf.float32
+
+
+def float32_bottleneck(x):
+  return tf.cast(tf.cast(x, tf.float32), tf.float64)
+
+
+def get_diff(l1, l2):
+  l2 = l2[::-1]
+  for i1, i2 in zip(l1, l2):
+    print (i1 - i2)
+  for i1, i2 in zip(l1, l2):
+    print (np.max(np.abs(i1 - i2)))
+
+
+class TransformerGlowLayersTest(parameterized.TestCase, tf.test.TestCase):
+
+  def get_hparams(self):
+    hparams = transformer.transformer_small()
+    hparams.add_hparam("prior_type", "affine")
+    hparams.add_hparam("factor", 2)  # squeezing factor
+    hparams.add_hparam("n_layers_transform_params", 1)
+    hparams.add_hparam("n_1x1_heads", N_1X1_HEADS)
+    hparams.add_hparam("flow_num_1x1_heads", 4)
+    hparams.add_hparam("flow_num_heads", 4)
+    hparams.add_hparam("flow_hidden_size", 64)
+    hparams.add_hparam("flow_filter_size", 128)
+    hparams.add_hparam("flow_layer_prepostprocess_dropout", 0.0)
+    hparams.add_hparam("flow_attention_dropout", 0.0)
+    hparams.add_hparam("flow_relu_dropout", 0.0)
+    hparams.add_hparam("latent_size", N_CHANNELS)
+    hparams.add_hparam("use_weightnorm", True)
+    hparams.add_hparam("kl_startup_steps", 2000)
+    hparams.add_hparam("affine_scale", "glow")
+    hparams.add_hparam("scale_width", 0.999)
+    hparams.add_hparam("step_fn", "glow")  # glow / chunting
+    hparams.add_hparam("conv_fn", "np")  # np / tf
+    hparams.add_hparam("posterior_type", "diagonal_normal")
+    hparams.causal_decoder_self_attention = False
+    hparams.hidden_size = HIDDEN_SIZE
+    hparams.weight_dtype = "float32"
+    hparams.add_hparam("pos_attn", False)
+    return hparams
+
+  def get_data(self):
+    x = tf.random_normal(
+        (BATCH_SIZE, TARGET_LENGTH, N_CHANNELS), dtype=DTYPE)
+    x_lengths = np.random.randint(
+        low=1, high=TARGET_LENGTH+1, size=BATCH_SIZE)
+    x_lengths = np.ceil(x_lengths / 4.0) * 4.0
+    x_lengths = x_lengths.astype(int)
+    x_mask = tf.sequence_mask(x_lengths, maxlen=TARGET_LENGTH, dtype=DTYPE)
+    return x, x_mask, x_lengths
+
+  def get_kwargs(self, x_mask, hparams=None):
+    if hparams is None:
+      hparams = self.get_hparams()
+    encoder_output = tf.random.uniform(
+        (BATCH_SIZE, INPUT_LENGTH, HIDDEN_SIZE), dtype=DTYPE)
+    encoder_decoder_attention_bias = tf.zeros(
+        (BATCH_SIZE, 1, 1, INPUT_LENGTH), dtype=DTYPE)
+    decoder_self_attention_bias = 1.0 - x_mask[:, tf.newaxis, tf.newaxis, :]
+    decoder_self_attention_bias *= -1e9
+    kwargs = {"hparams": hparams,
+              "encoder_output": encoder_output,
+              "encoder_decoder_attention_bias": encoder_decoder_attention_bias,
+              "decoder_self_attention_bias": decoder_self_attention_bias}
+    return kwargs
+
+  def test_actnorm(self):
+    _, x_mask, _ = self.get_data()
+    x = tf.random_normal((BATCH_SIZE, TARGET_LENGTH, N_CHANNELS),
+                         mean=50.0, stddev=10.0, dtype=DTYPE)
+    x_act, logabsdet = glow.actnorm(
+        "actnorm", x, x_mask, inverse=False, init=True)
+
+    x_act_nopad = tf.boolean_mask(x_act, x_mask)
+    x_mean, x_var = tf.nn.moments(x_act_nopad, axes=[0])
+    self.evaluate(tf.global_variables_initializer())
+    x, x_act, logabsdet, x_mean, x_var = (
+        self.evaluate([x, x_act, logabsdet, x_mean, x_var]))
+    self.assertEqual(x_act.shape, (BATCH_SIZE, TARGET_LENGTH, N_CHANNELS))
+    self.assertEqual(logabsdet.shape, (BATCH_SIZE,))
+    self.assertTrue(np.allclose(x_mean, 0.0, atol=1e-5))
+    self.assertTrue(np.allclose(x_var, 1.0, atol=1e-5))
+
+  def test_actnorm_invertibility(self):
+    name = "actnorm"
+    x, x_mask, _ = self.get_data()
+
+    x_inv, logabsdet = glow.actnorm(
+        name, x, x_mask, inverse=False, init=False)
+    x_inv_inv, logabsdet_inv = glow.actnorm(
+        name, x_inv, x_mask, inverse=True, init=False)
+    self.evaluate(tf.global_variables_initializer())
+    x, x_inv, x_inv_inv, x_mask, logabsdet, logabsdet_inv = (
+        self.evaluate(
+            [x, x_inv, x_inv_inv, x_mask, logabsdet, logabsdet_inv]))
+    diff = x - x_inv_inv
+    logabsdet_sum = logabsdet + logabsdet_inv
+    self.assertEqual(x.shape, (BATCH_SIZE, TARGET_LENGTH, N_CHANNELS))
+    self.assertEqual(x_inv.shape, (BATCH_SIZE, TARGET_LENGTH, N_CHANNELS))
+    self.assertEqual(x_inv_inv.shape, (BATCH_SIZE, TARGET_LENGTH, N_CHANNELS))
+    self.assertTrue(np.allclose(diff, 0.0, atol=1e-5))
+    self.assertTrue(np.allclose(logabsdet_sum, 0.0, atol=1e-5))
+
+  @parameterized.parameters(
+      (glow.multihead_invertible_1x1_conv_np, "a"),
+      (glow.multihead_invertible_1x1_conv_np, "c"),
+      )
+  def test_multi_1x1_invertibility(
+      self, func, multihead_split):
+    name = "multi_1x1"
+    x, x_mask, _ = self.get_data()
+
+    x_inv, logabsdet = func(
+        name, x, x_mask, multihead_split, inverse=False, dtype=DTYPE)
+    x_inv_inv, logabsdet_inv = func(
+        name, x_inv, x_mask, multihead_split, inverse=True, dtype=DTYPE)
+    self.evaluate(tf.global_variables_initializer())
+    x, x_mask, x_inv, x_inv_inv, logabsdet, logabsdet_inv = (
+        self.evaluate(
+            [x, x_mask, x_inv, x_inv_inv, logabsdet, logabsdet_inv]))
+    diff = x - x_inv_inv
+    logabsdet_sum = logabsdet + logabsdet_inv
+    logabsdet_ = logabsdet / np.sum(x_mask, -1)
+    self.assertTrue(np.allclose(diff, 0.0, atol=1e-5))
+    self.assertTrue(np.allclose(logabsdet_, 0.0, atol=1e-5))
+    self.assertTrue(np.allclose(logabsdet_sum, 0.0, atol=1e-5))
+
+  @parameterized.parameters(
+      (glow.additive_coupling, "c"),
+      (glow.additive_coupling, "t"),
+      (glow.additive_coupling, "a"),
+      (glow.affine_coupling, "c"),
+      (glow.affine_coupling, "t"),
+      (glow.affine_coupling, "a"),
+      )
+  def test_coupling_invertibility(self, func, split_dim):
+    name = "affine"
+    x, x_mask, _ = self.get_data()
+    kwargs = self.get_kwargs(x_mask)
+
+    x_inv, logabsdet = func(
+        name, x, x_mask, split_dim=split_dim,
+        identity_first=True, inverse=False, init=False, disable_dropout=True,
+        **kwargs)
+    x_inv_inv, logabsdet_inv = func(
+        name, x_inv, x_mask, split_dim=split_dim,
+        identity_first=True, inverse=True, init=False, disable_dropout=True,
+        **kwargs)
+    self.evaluate(tf.global_variables_initializer())
+    x, x_mask, x_inv, x_inv_inv, logabsdet, logabsdet_inv = (
+        self.evaluate(
+            [x, x_mask, x_inv, x_inv_inv, logabsdet, logabsdet_inv]))
+    diff = x - x_inv_inv
+    logabsdet_sum = logabsdet + logabsdet_inv
+    self.assertTrue(np.allclose(diff, 0.0, atol=1e-5))
+    self.assertTrue(np.allclose(logabsdet_sum, 0.0, atol=1e-5))
+
+  def test_split(self):
+    x, x_mask, _ = self.get_data()
+
+    x_inv, z, log_p = glow.split(
+        "split", x, x_mask, inverse=False)
+    x_inv_inv, _, log_p_inv = glow.split(
+        "split", x_inv, x_mask, z=z, inverse=True)
+    self.evaluate(tf.global_variables_initializer())
+    x, x_inv, x_inv_inv, z, log_p, log_p_inv = self.evaluate(
+        [x, x_inv, x_inv_inv, z, log_p, log_p_inv])
+    diff = x - x_inv_inv
+    log_p_diff = log_p - log_p_inv
+    self.assertEqual(
+        x_inv.shape, (BATCH_SIZE, TARGET_LENGTH, N_CHANNELS//2))
+    self.assertEqual(
+        z.shape, (BATCH_SIZE, TARGET_LENGTH, N_CHANNELS//2))
+    self.assertTrue(np.allclose(diff, 0.0, atol=1e-5))
+    self.assertTrue(np.allclose(log_p_diff, 0.0, atol=1e-5))
+
+  def test_flow_invertibility(self):
+    name = "flow_step"
+    split_dims = "cat"
+    x, x_mask, _ = self.get_data()
+    kwargs = self.get_kwargs(x_mask)
+    x_inv, logabsdet = glow.flow_step_glow(
+        name, x, x_mask, split_dims, inverse=False, init=False, dtype=DTYPE,
+        disable_dropout=True, **kwargs)
+    x_inv_inv, logabsdet_inv = glow.flow_step_glow(
+        name, x_inv, x_mask, split_dims, inverse=True, init=False,
+        dtype=DTYPE, disable_dropout=True, **kwargs)
+    self.evaluate(tf.global_variables_initializer())
+    x, x_mask, x_inv, x_inv_inv, logabsdet, logabsdet_inv = (
+        self.evaluate(
+            [x, x_mask, x_inv, x_inv_inv, logabsdet, logabsdet_inv]))
+    diff = x - x_inv_inv
+    logabsdet_sum = logabsdet + logabsdet_inv
+    self.assertTrue(np.allclose(diff, 0.0, atol=1e-5))
+    self.assertTrue(np.allclose(logabsdet_sum, 0.0, atol=1e-5))
+
+  @parameterized.parameters(
+      ("1", "cat", "affine"),
+      ("1/1", "cat/cat", "affine"),
+      ("1/1/1", "cat/cat/ca", "affine"),
+      )
+  def test_aaa_glow_training(self, depths, split_plans, prior_type):
+    with tf.Graph().as_default():
+      _, x_mask, _ = self.get_data()
+      x = tf.random_normal((BATCH_SIZE, TARGET_LENGTH, N_CHANNELS),
+                           mean=10.0, stddev=3.0, dtype=DTYPE)
+      bias = common_attention.attention_bias_ignore_padding(1.0 - x_mask)
+      hparams = self.get_hparams()
+      hparams.prior_type = prior_type
+      hparams.depths = depths
+      hparams.split_plans = split_plans
+      n_levels = len(hparams.depths.split("/"))
+      kwargs = self.get_kwargs(x_mask, hparams)
+      _ = kwargs.pop("decoder_self_attention_bias")
+
+      x_inv, _, _, _ = glow.glow(
+          "glow", x, x_mask, bias, inverse=False, init=True,
+          disable_dropout=True, **kwargs)
+      curr_dir = tempfile.mkdtemp()
+      model_path = os.path.join(curr_dir, "model")
+
+      with tf.Session() as session:
+        saver = tf.train.Saver()
+        session.run(tf.global_variables_initializer())
+        session.run(x_inv)
+        saver.save(session, model_path)
+
+    with tf.Graph().as_default():
+      _, x_mask, _ = self.get_data()
+      x = tf.random_normal((BATCH_SIZE, TARGET_LENGTH, N_CHANNELS),
+                           mean=10.0, stddev=3.0, dtype=DTYPE)
+      bias = common_attention.attention_bias_ignore_padding(1.0 - x_mask)
+      hparams = self.get_hparams()
+      hparams.depths = depths
+      hparams.split_plans = split_plans
+      kwargs = self.get_kwargs(x_mask, hparams)
+      _ = kwargs.pop("decoder_self_attention_bias")
+      log_q_z = gops.standard_normal_density(x, x_mask)
+      log_q_z = tf.reduce_sum(log_q_z) / tf.reduce_sum(x_mask)
+
+      x_inv, logabsdets, log_ps, zs = glow.glow(
+          "glow", x, x_mask, bias, inverse=False, init=False,
+          disable_dropout=True, **kwargs)
+      x_inv_inv, logabsdets_inv, log_ps_inv, _ = glow.glow(
+          "glow", x_inv, x_mask, bias, inverse=True, split_zs=zs, init=False,
+          disable_dropout=True, **kwargs)
+      logabsdets = tf.reduce_sum(
+          logabsdets, axis=0) / tf.reduce_sum(x_mask)
+      logabsdets_inv = tf.reduce_sum(
+          logabsdets_inv, axis=0) / tf.reduce_sum(x_mask)
+      log_ps = tf.reduce_sum(log_ps, axis=0) / tf.reduce_sum(x_mask)
+      log_ps_inv = tf.reduce_sum(log_ps_inv, axis=0) / tf.reduce_sum(x_mask)
+
+      with tf.Session() as session:
+        saver = tf.train.Saver()
+        saver.restore(session, model_path)
+        (x, x_inv, x_inv_inv, log_q_z, logabsdets, log_ps,
+         logabsdets_inv, log_ps_inv) = session.run([
+             x, x_inv, x_inv_inv, log_q_z, logabsdets, log_ps,
+             logabsdets_inv, log_ps_inv])
+        diff = x - x_inv_inv
+        log_ps_diff = log_ps - log_ps_inv
+        logabsdets_sum = logabsdets + logabsdets_inv
+        self.assertEqual(
+            x_inv.shape,
+            (BATCH_SIZE, TARGET_LENGTH//(2**(n_levels-1)), N_CHANNELS))
+        print (np.max(np.abs(diff)))
+        print (np.max(np.abs(log_ps_diff)))
+        print (np.max(np.abs(logabsdets_sum)))
+        self.assertTrue(np.allclose(diff, 0.0, atol=1e-4),
+                        msg=np.max(np.abs(diff)))
+        self.assertTrue(np.allclose(log_ps_diff, 0.0, atol=1e-4),
+                        msg=np.max(np.abs(log_ps_diff)))
+        self.assertTrue(np.allclose(logabsdets_sum, 0.0, atol=1e-4),
+                        msg=np.max(np.abs(logabsdets_sum)))
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/tensor2tensor/models/research/transformer_vae_flow_prior.py b/tensor2tensor/models/research/transformer_vae_flow_prior.py
new file mode 100644
index 000000000..e6147ff64
--- /dev/null
+++ b/tensor2tensor/models/research/transformer_vae_flow_prior.py
@@ -0,0 +1,1135 @@
+# coding=utf-8
+# Copyright 2020 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Transformer VAE with Flow Priors for Non-Autoregressive MT."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import copy
+import inspect
+import math
+import six
+
+from tensor2tensor.data_generators import multi_problem
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
+from tensor2tensor.layers import transformer_glow_layers as glow
+from tensor2tensor.layers import transformer_glow_layers_ops as gops
+from tensor2tensor.models import transformer
+from tensor2tensor.research.models import transformer_vae_flow_prior_ops as ops
+from tensor2tensor.utils import contrib
+from tensor2tensor.utils import optimize
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
+import tensorflow.compat.v1 as tf
+
+
+@registry.register_model
+class TransformerVaeFlowPrior(t2t_model.T2TModel):
+  """Transformer VAE using flow priors."""
+
+  def __init__(self, *args, **kwargs):
+    super(TransformerVaeFlowPrior, self).__init__(*args, **kwargs)
+    hparams = self._hparams
+    if hparams.prior_type in ["affine", "additive", "rq"]:
+      self._fparams = contrib.training.HParams(**hparams.values())
+      for key, value in self._fparams.values().items():
+        if key.startswith("flow_"):
+          setattr(self._fparams, key[5:], value)
+
+  @property
+  def is_training(self):
+    return self.hparams.mode == tf.estimator.ModeKeys.TRAIN
+
+  @property
+  def is_evaluating(self):
+    return self._hparams.mode == tf.estimator.ModeKeys.EVAL
+
+  @property
+  def is_predicting(self):
+    return self._hparams.mode == tf.estimator.ModeKeys.PREDICT
+
+  def loss_iw(self, logits, features):
+    if isinstance(logits, dict):
+      losses = {}
+      for k, v in six.iteritems(logits):
+        losses[k] = self._loss_single_iw(
+            v,
+            k,
+            features[k],
+            weights=features.get(k + "_mask"))
+
+        n, d = losses[k]
+        if common_layers.should_generate_summaries():
+          tf.summary.scalar(k + "_loss", n / d)
+          tf.summary.scalar(k + "_loss_num", n)
+          tf.summary.scalar(k + "_loss_den", d)
+          if getattr(self.hparams, "visualize_logits_histogram", False):
+            hist = tf.summary.histogram
+            hist(k + "_predict", tf.argmax(tf.squeeze(v), axis=-1))
+            hist(k + "_targets", features[k])
+
+      return tf.add_n([n / d for n, d in losses.values()])
+    else:
+      return self._loss_single_iw(
+          logits,
+          "targets",
+          features["targets"],
+          weights=features.get("targets_mask"))
+
+  def _loss_single_iw(self, logits, feature_name, feature, weights=None):
+    # The current bfloat16 version still uses float32 for most parts of backward
+    # propagation to keep model quality, so cast back before computing the loss
+    # value.
+    no_problem_err_str = (
+        "The default implementation of %s requires that the "
+        "model be used with a Problem. If using a Problem, augment the "
+        "hparams object with trainer_lib.add_problem_hparams. If not, "
+        "override %s.")
+    no_problem_err = (
+        lambda method_name: no_problem_err_str % (method_name, method_name))
+    if not self._problem_hparams:
+      t2t_model.log_warn(no_problem_err("loss"))
+      return (tf.constant(0., dtype=tf.float32),
+              tf.constant(1., dtype=tf.float32))
+
+    # Calculate loss contribution.
+    modality = self._problem_hparams.modality[feature_name]
+    vocab_size = self._problem_hparams.vocab_size[feature_name]
+    if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
+      vocab_size += (-vocab_size) % self._hparams.vocab_divisor
+    # loss = self._hparams.loss.get(feature_name, modalities.get_loss(modality))
+    loss = ops.generic_loss
+    targets_weights_fn = self._hparams.weights_fn.get(
+        "targets", modalities.get_weights_fn(modality))
+    if weights is None:
+      loss_num, loss_den = loss(logits, feature, self._hparams, vocab_size,
+                                weights_fn=targets_weights_fn)
+    else:
+
+      def weights_fn(labels):
+        """Per-token weights for loss."""
+        # Use target_weights_fn() given by modality as well as explicitly given
+        # weights.
+        modality_weights = targets_weights_fn(labels)
+
+        # Broadcast 'weights' along minor dimensions (TF's default is major).
+        explicit_weights = weights
+        if len(explicit_weights.shape) < len(modality_weights.shape):
+          explicit_weights = common_layers.expand_squeeze_to_nd(
+              weights, modality_weights.shape.ndims)
+
+        return explicit_weights * modality_weights
+
+      # Ensure that target.modality_loss() supports "weights_fn" keyword
+      # argument. If it doesn't and "weights" is specified, raise an exception.
+      argument_names = inspect.getargspec(loss).args
+      if "weights_fn" not in argument_names:
+        raise ValueError(
+            "Explicit 'weights' given but default loss for modality doesn't "
+            "support 'weights_fn' keyword argument: %s.loss(%s)." %
+            (modality, ", ".join(argument_names)))
+
+      loss_num, loss_den = loss(
+          logits, feature, self._hparams, vocab_size, weights_fn=weights_fn)
+
+    loss_num *= self._problem_hparams.loss_multiplier
+
+    if hasattr(self.hparams, "problem") and hasattr(
+        self.hparams.problem, "task_list"):
+      if weights is not None:
+        raise NotImplementedError("weights not yet implemented in "
+                                  "multitask setting.")
+      loss_num, loss_den, summaries = multi_problem.aggregate_task_losses(
+          self.hparams,
+          self._problem_hparams,
+          logits,
+          feature_name,
+          feature
+      )
+
+      for key, val in summaries:
+        tf.summary.scalar(key, val)
+
+    return loss_num, loss_den
+
+  def internal(self, features, real_features):
+    """Main procedure for both training and inference."""
+    inputs = common_layers.flatten4d3d(features["inputs"])
+    targets = common_layers.flatten4d3d(features["targets"])
+    target_space = features["target_space_id"]
+    hparams = self._hparams
+    inputs_mask = ops.embedding_to_non_padding(inputs)
+    inputs_length = tf.reduce_sum(inputs_mask, axis=-1)
+
+    encoder_output, encoder_decoder_attention_bias = (
+        ops.encoder("encoder", hparams, inputs, target_space))
+    kwargs = {"encoder_output": encoder_output,
+              "encoder_decoder_attention_bias": encoder_decoder_attention_bias}
+    losses, monitor = {}, {}
+    log_abs_det = tf.constant(0.0)
+
+    if not self.is_predicting:
+      # Training
+      targets_mask = ops.embedding_to_non_padding(targets)
+      targets_length = tf.reduce_sum(targets_mask, axis=-1)
+      length_diff = targets_length - inputs_length
+      decoder_self_attention_bias = (
+          common_attention.attention_bias_ignore_padding(1.0 - targets_mask))
+      z_q, log_q_z, q_dist = self.sample_q(
+          targets, targets_mask, decoder_self_attention_bias, n_samples=1,
+          temp=1.0, **kwargs)
+
+      body_output = ops.decoder(
+          "decoder", z_q, hparams, decoder_self_attention_bias, **kwargs)
+      logits = self.top(body_output, real_features)
+      numerator, denominator = self.loss(logits, real_features)
+
+      if not (self.is_evaluating and (
+          hparams.compute_kl_refinement or hparams.compute_iw_marginal)):
+        targets_length_pred, lenpred_loss = ops.predict_target_lengths(
+            encoder_output, inputs_mask, hparams, length_diff)
+        log_p_z_base, log_abs_det = self.compute_prior_log_prob(
+            z_q, targets_mask, decoder_self_attention_bias,
+            check_invertibility=False, **kwargs)
+        losses, monitor = ops.save_log_loss(
+            hparams, targets_mask, numerator, denominator, log_q_z, log_abs_det,
+            log_p_z_base, z_q, lenpred_loss, targets_length_pred,
+            targets_length)
+
+      if self.is_evaluating:
+        if hparams.compute_kl_refinement:
+          z_p, _ = self.sample_p(
+              targets_length, temp=self._decode_hparams.temp,
+              check_invertibility=False, targets_mask=targets_mask, **kwargs)
+          z_dq = self.delta_posterior(
+              z_p, targets_mask, decoder_self_attention_bias,
+              self._decode_hparams.n_gibbs_steps, **kwargs)
+          log_q_z_ = q_dist.log_prob(z_dq)
+          log_q_z_ = gops.reduce_mean_over_bl_sum_over_c(log_q_z_, targets_mask)
+          losses = {"training": log_q_z_}
+
+        if hparams.compute_iw_marginal:
+        # if True:
+          log_p_y_x = self.compute_iw_marginal(
+              targets, targets_mask, decoder_self_attention_bias,
+              real_features, self._decode_hparams.n_samples, **kwargs)
+              # real_features, 1, **kwargs)
+          losses = {"training": log_p_y_x}
+
+      return logits, losses, monitor, targets_mask
+
+    else:
+      # Inference
+      targets_length, _ = ops.predict_target_lengths(
+          encoder_output, inputs_mask, hparams)
+      targets_mask = ops.sequence_mask(targets_length, hparams)
+      decoder_self_attention_bias = (
+          common_attention.attention_bias_ignore_padding(1.0 - targets_mask))
+      z_p, _ = self.sample_p(
+          targets_length, temp=self._decode_hparams.temp,
+          check_invertibility=False, **kwargs)
+      z_q = self.delta_posterior(
+          z_p, targets_mask, decoder_self_attention_bias,
+          self._decode_hparams.n_gibbs_steps, **kwargs)
+          # 0, **kwargs)
+
+      body_output = ops.decoder(
+          "decoder", z_q, hparams, decoder_self_attention_bias, **kwargs)
+      return body_output, losses, monitor, targets_mask
+
+  def sample_q(
+      self, targets, targets_mask, decoder_self_attention_bias, n_samples,
+      temp, **kwargs):
+    hparams = self._hparams
+    batch_size, targets_max_length = common_layers.shape_list(targets_mask)[:2]
+    q_params = ops.posterior("posterior", hparams, targets, targets_mask,
+                             decoder_self_attention_bias, **kwargs)
+    q_dist = gops.diagonal_normal(q_params, "posterior")
+    loc, scale = q_dist.loc, q_dist.scale
+    z_shape = [batch_size, targets_max_length, hparams.latent_size]
+    iw_z_shape = [n_samples*batch_size, targets_max_length, hparams.latent_size]
+    if n_samples == 1:
+      noise = tf.random_normal(z_shape, stddev=temp)
+      z_q = loc + scale * noise
+      log_q_z = q_dist.log_prob(z_q)  # [B, L, C]
+    else:
+      noise = tf.random_normal([n_samples] + z_shape, stddev=temp)
+      z_q = loc[tf.newaxis, ...] + scale[tf.newaxis, ...] * noise
+      log_q_z = q_dist.log_prob(z_q)  # [K, B, L, C]
+      z_q = tf.reshape(z_q, iw_z_shape)
+      log_q_z = tf.reshape(log_q_z, iw_z_shape)
+    return z_q, log_q_z, q_dist
+
+  def compute_iw_marginal(
+      self, targets, targets_mask, decoder_self_attention_bias, features,
+      n_samples, reduce_mean=True, **kwargs):
+    hparams = self._hparams
+    z_q, log_q_z, _ = self.sample_q(
+        targets, targets_mask, decoder_self_attention_bias,
+        n_samples=n_samples, temp=1.0, **kwargs)  # [K*B, L, C]
+    iw_kwargs = {key: ops.prepare_for_iw(value, n_samples) for (
+        key, value) in kwargs.items()}
+    iw_targets_mask = ops.prepare_for_iw(targets_mask, n_samples)
+    iw_decoder_self_attention_bias = (
+        common_attention.attention_bias_ignore_padding(1.0 - iw_targets_mask))
+    iw_features = copy.copy(features)
+    iw_features["targets"] = ops.prepare_for_iw(
+        features["targets"], n_samples)
+
+    log_p_z_base, log_abs_det = self.compute_prior_log_prob(
+        z_q, iw_targets_mask, iw_decoder_self_attention_bias,
+        check_invertibility=False, **iw_kwargs)
+    log_p_z = log_p_z_base + log_abs_det
+
+    body_output = ops.decoder(
+        "decoder", z_q, hparams, iw_decoder_self_attention_bias, **iw_kwargs)
+    logits = self.top(body_output, iw_features)
+    numerator, denominator = self.loss_iw(logits, iw_features)
+    numerator = tf.reduce_sum(numerator[..., 0, 0], 1)  # [K*B]
+    denominator = tf.reduce_sum(denominator[..., 0, 0], 1)  # [K*B]
+    log_p_x = -1 * numerator / denominator
+    log_q_z = gops.reduce_mean_over_l_sum_over_c(log_q_z, iw_targets_mask)
+    log_p_z = log_p_z / tf.reduce_sum(iw_targets_mask, 1)
+
+    log_p_x, log_q_z, log_p_z = [ops.unprepare_for_iw(ii, n_samples) for ii in [
+        log_p_x, log_q_z, log_p_z]]
+
+    log_w_n = log_p_z - log_q_z
+    log_w_n = tf.nn.log_softmax(log_w_n, axis=0)  # [K, B]
+
+    iw_marginal = log_p_x + log_w_n
+    iw_marginal = tf.reduce_logsumexp(iw_marginal, 0)  # [B]
+
+    if reduce_mean:
+      iw_marginal = tf.cast(tf.reduce_mean(iw_marginal, 0), tf.float32)  # [1]
+    else:
+      iw_marginal = tf.cast(iw_marginal, tf.float32)  # [1]
+    return iw_marginal
+
+  def argmax_decode(self, z, decoder_self_attention_bias, **kwargs):
+    hparams = self._hparams
+    body_output = ops.decoder(
+        "decoder", z, hparams, decoder_self_attention_bias, **kwargs)
+    logits = self.top(body_output, {"targets": None})
+    targets = tf.argmax(logits, axis=-1)
+    targets_emb = self.bottom({"targets": targets})["targets"][..., 0, :]
+    return targets, targets_emb
+
+  def delta_posterior(
+      self, z, targets_mask, decoder_self_attention_bias, n_gibbs_steps,
+      **kwargs):
+    hparams = self._hparams
+    for _ in range(n_gibbs_steps):
+      _, targets_emb = self.argmax_decode(
+          z, decoder_self_attention_bias, **kwargs)
+      q_params = ops.posterior(
+          "posterior", hparams, targets_emb, targets_mask,
+          decoder_self_attention_bias, **kwargs)
+      q_dist = gops.diagonal_normal(q_params, "posterior")
+      z = q_dist.loc  # [B, L, C]
+    return z
+
+  def compute_prior_log_prob(
+      self, z_q, targets_mask, decoder_self_attention_bias,
+      check_invertibility=False, **kwargs):
+    hparams = self._hparams
+    batch_size, targets_max_length = (
+        common_layers.shape_list(targets_mask)[:2])
+    prior_shape = [batch_size, targets_max_length, hparams.latent_size]
+    log_abs_det = tf.zeros([batch_size])
+
+    if hparams.prior_type == "standard_normal":
+      log_p_z_base = gops.standard_normal_density(z_q, targets_mask)
+    elif hparams.prior_type == "diagonal_normal":
+      diag_prior_params = ops.cond_prior(
+          "diag_prior", hparams, tf.zeros(prior_shape), targets_mask,
+          hparams.latent_size*2, decoder_self_attention_bias, **kwargs)
+      p_dist = gops.diagonal_normal(diag_prior_params, "diag_prior")
+      log_p_z_base = p_dist.log_prob(z_q)  # [B, L, C]
+      log_p_z_base = gops.reduce_sum_over_lc(log_p_z_base, targets_mask)  # [B]
+    elif hparams.prior_type in ["affine", "additive", "rq"]:
+      if self.is_evaluating:
+        disable_dropout = True
+        init = False
+      elif self.is_training:
+        disable_dropout = False
+        init = tf.equal(hparams.kl_startup_steps,
+                        tf.cast(tf.train.get_global_step(), tf.int32))
+      else:
+        raise ValueError("compute_prior shouldn't be used in decoding.")
+
+      z_inv, log_abs_det, log_p_z_base, zs = glow.glow(
+          "glow", z_q, targets_mask, decoder_self_attention_bias,
+          inverse=False, init=init, hparams=self._fparams,
+          disable_dropout=disable_dropout, **kwargs)
+      if self.is_evaluating and check_invertibility:
+        z_inv_inv, _, _, _ = glow.glow(
+            "glow", z_inv, targets_mask, decoder_self_attention_bias,
+            inverse=True, split_zs=zs, init=False, hparams=self._fparams,
+            disable_dropout=True, **kwargs)
+        z_diff = z_q - z_inv_inv
+        tf.summary.scalar("flow_recon_forward", tf.reduce_max(tf.abs(z_diff)))
+    return log_p_z_base, log_abs_det
+
+  def sample_p(
+      self, targets_length, temp, check_invertibility=False, targets_mask=None,
+      **kwargs):
+    hparams = self._hparams
+    if targets_mask is None:
+      targets_mask = ops.sequence_mask(targets_length, hparams)
+    decoder_self_attention_bias = (
+        common_attention.attention_bias_ignore_padding(1.0 - targets_mask))
+    batch_size, targets_max_length = (
+        common_layers.shape_list(targets_mask)[:2])
+    prior_shape = [batch_size, targets_max_length, hparams.latent_size]
+    noise = tf.random.normal(prior_shape, stddev=temp)
+    p_dist = None
+
+    if hparams.prior_type == "standard_normal":
+      z_p = noise
+    elif hparams.prior_type == "diagonal_normal":
+      diag_prior_params = ops.cond_prior(
+          "diag_prior", hparams, tf.zeros(prior_shape), targets_mask,
+          hparams.latent_size*2, decoder_self_attention_bias, **kwargs)
+      p_dist = gops.diagonal_normal(diag_prior_params, "diag_prior")
+      z_p = p_dist.loc + p_dist.scale * noise
+    elif hparams.prior_type in ["affine", "additive", "rq"]:
+      n_levels = len(hparams.depths.split("/"))
+      divi = max(1, hparams.factor**(n_levels-1))
+      flow_prior_shape = [
+          batch_size, targets_max_length//divi, hparams.latent_size]
+      noise = tf.random_normal(flow_prior_shape, stddev=temp)
+      z_p, _, _, _ = glow.glow(
+          "glow", noise, targets_mask, decoder_self_attention_bias,
+          inverse=True, init=False, hparams=self._fparams,
+          disable_dropout=True, temp=temp, **kwargs)
+      if self.is_evaluating and check_invertibility:
+        noise_inv, _, _, _ = glow.glow(
+            "glow", z_p, targets_mask, decoder_self_attention_bias,
+            inverse=False, init=False, hparams=self._fparams,
+            disable_dropout=True, **kwargs)
+        z_diff = noise - noise_inv
+        tf.summary.scalar("flow_recon_inverse", tf.reduce_max(tf.abs(z_diff)))
+    return z_p, p_dist
+
+  def optimize(self, loss, num_async_replicas=1, use_tpu=False, variables=None):
+    """Return a training op minimizing loss."""
+    lr = ops.learning_rate_schedule(self.hparams)
+    if num_async_replicas > 1:
+      t2t_model.log_info("Dividing learning rate by num_async_replicas: %d",
+                         num_async_replicas)
+    lr /= math.sqrt(float(num_async_replicas))
+    train_op = optimize.optimize(
+        loss, lr, self.hparams, use_tpu=use_tpu, variables=variables)
+    return train_op
+
+  def body(self, features, real_features):
+    return self.internal(features, real_features)
+
+  def infer(self,
+            features,
+            *args,
+            **kwargs):
+    """Produce predictions from the model."""
+    del args, kwargs
+    inputs_old = None
+    if "inputs" in features and len(features["inputs"].shape) < 4:
+      inputs_old = features["inputs"]
+      features["inputs"] = tf.expand_dims(features["inputs"], 2)
+    features["targets"] = tf.identity(features["inputs"])
+
+    # logits, _ = self(features)
+    t2t_model.set_custom_getter_compose(self._custom_getter)
+    tf.get_variable_scope().set_initializer(
+        optimize.get_variable_initializer(self.hparams))
+    with self._eager_var_store.as_default():
+      self._fill_problem_hparams_features(features)
+      # intentionally disable sharding during inference (in multi GPU)
+      with tf.variable_scope(self.name):
+        logits, _, _, targets_mask = self.model_fn(features)
+
+    samples = tf.argmax(logits, axis=-1)
+    samples = tf.where(
+        tf.cast(targets_mask[..., tf.newaxis, tf.newaxis], tf.bool),
+        samples, tf.ones_like(samples))
+    if inputs_old is not None:  # Restore to not confuse Estimator.
+      features["inputs"] = inputs_old
+    return samples
+
+  def model_fn(self, features):
+    with tf.variable_scope(
+        tf.get_variable_scope(), use_resource=True, reuse=tf.AUTO_REUSE):
+      transformed_features = self.bottom(features)
+
+      if self.hparams.activation_dtype == "bfloat16":
+        for k, v in sorted(six.iteritems(transformed_features)):
+          if v.dtype == tf.float32:
+            transformed_features[k] = tf.cast(v, tf.bfloat16)
+
+      t2t_model.log_info("Building model body")
+      output, losses, monitor, targets_mask = self.body(
+          transformed_features, features)
+      output, losses = self._normalize_body_output((output, losses))
+
+      if "training" in losses:
+        t2t_model.log_info(
+            "Skipping T2TModel top and loss because training loss "
+            "returned from body")
+        logits = output
+      else:
+        logits = self.top(output, features)
+        losses["training"] = 0.0
+        if (self._hparams.mode != tf.estimator.ModeKeys.PREDICT and
+            self._hparams.mode != "attack"):
+          losses["training"] = self.loss(logits, features)
+
+    return logits, losses, monitor, targets_mask
+
+  def model_fn_sharded(self, sharded_features):
+    """Estimator model_fn sharded along batch dimension.
+
+    Args:
+      sharded_features: {str: [Tensor]}. Features sharded along batch dimension.
+        Each list is the same length (== number of shards).
+
+    Returns:
+      sharded_logits: [Tensor]. Logits for each shard of examples.
+      losses: {str: 0-D Tensor}. Loss averaged across shards.
+    """
+    dp = self._data_parallelism
+
+    # [{str: Tensor}]. Transpose of 'sharded_features'.
+    datashard_to_features = self._to_features_per_datashard(sharded_features)
+    sharded_logits, sharded_losses, sharded_monitors, _ = (
+        dp(self.model_fn, datashard_to_features))
+    sharded_logits, sharded_losses = dp(
+        self.maybe_scheduled_sampling,
+        datashard_to_features, sharded_logits, sharded_losses)
+    if isinstance(sharded_logits[0], dict):
+      temp_dict = {k: [] for k, _ in six.iteritems(sharded_logits[0])}
+      for k, _ in six.iteritems(sharded_logits[0]):
+        for l in sharded_logits:
+          temp_dict[k].append(l[k])
+      sharded_logits = temp_dict
+    losses = t2t_model.average_sharded_losses(sharded_losses)
+    monitor = {}
+    for key in list(sharded_monitors[0].keys()):
+      monitor[key] = (
+          tf.add_n([m[key] for m in sharded_monitors]) / len(sharded_monitors))
+    ops.save_summary(monitor, "monitor")
+
+    return sharded_logits, losses
+
+
+@registry.register_hparams
+def wmt_enro_tpu():
+  """HParams for Transformer model on TPU."""
+  hparams = transformer.transformer_base()
+  hparams = transformer.update_hparams_for_tpu(hparams)
+  hparams.batch_size = 512
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_baseline_gpu():
+  """HParams for Transformer model on TPU."""
+  hparams = transformer.transformer_base()
+  hparams.hidden_size = 256
+  hparams.filter_size = 1024
+  hparams.num_hidden_layers = 5
+  hparams.num_heads = 2
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.attention_dropout = 0.1
+  hparams.relu_dropout = 0.1
+  hparams.dropout = 0.1
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_baseline_single_gpu():
+  """HParams for Transformer model on TPU."""
+  hparams = iwslt_baseline_gpu()
+  hparams.batch_size = 1024
+  hparams.learning_rate_schedule = "constant*linear_warmup*rsqrt_decay"
+  hparams.learning_rate_constant = 0.1
+  hparams.learning_rate_warmup_steps = 16000
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_baseline_tpu():
+  """HParams for Transformer model on TPU."""
+  hparams = transformer.transformer_base()
+  transformer.update_hparams_for_tpu(hparams)
+  hparams.hidden_size = 256
+  hparams.filter_size = 1024
+  hparams.num_hidden_layers = 5
+  hparams.num_heads = 2
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.attention_dropout = 0.1
+  hparams.relu_dropout = 0.1
+  hparams.dropout = 0.1
+  hparams.add_hparam("pos_attn", False)
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_base():
+  """Set of hyperparameters."""
+  # Model architecture flags.
+  hparams = transformer.transformer_base()
+  hparams.num_hidden_layers = 5
+  hparams.hidden_size = 256
+  hparams.filter_size = 1024
+  hparams.num_heads = 4
+  # Other flags.
+  hparams.summarize_grads = False
+  hparams.summarize_vars = False
+  # Optimization-related flags.
+  hparams.clip_grad_norm = 1.0
+  hparams.learning_rate_decay_scheme = "noam"
+  hparams.learning_rate_warmup_steps = 8000
+  hparams.learning_rate = 0.2
+  hparams.learning_rate_schedule = (
+      "constant*linear_warmup*rsqrt_decay*rsqrt_hidden_size")
+  hparams.learning_rate_constant = 2.0
+  hparams.add_hparam("predict_target_length", True)
+  hparams.add_hparam("lendiff_bound", 30)
+  hparams = update_hparams_for_tpu(hparams)
+  hparams.add_hparam("pos_attn", False)
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_diag():
+  """Set of hyperparameters."""
+  hparams = iwslt_base()
+  hparams.batch_size = 4096
+  # Other flags.
+  hparams.force_full_predict = True
+  hparams.causal_decoder_self_attention = False
+  # VAE-related flags.
+  hparams.add_hparam("latent_size", 256)
+  hparams.add_hparam("anneal_min_value", 0.0)
+  hparams.add_hparam("kl_startup_steps", 5000)
+  hparams.add_hparam("kl_anneal_steps", 20000)
+  hparams.add_hparam("n_posterior_layers", 3)
+  hparams.add_hparam("n_decoder_layers", 3)
+  hparams.add_hparam("posterior_2d_dropout", 0.20)
+  # diagonal_normal / affine / additive / rq
+  hparams.add_hparam("posterior_type", "diagonal_normal")
+  # standard_normal / diagonal_normal
+  hparams.add_hparam("prior_type", "diagonal_normal")
+  hparams.add_hparam("decoder_2d_dropout", 0.00)
+  # Optimization-related flags.
+  hparams.learning_rate_warmup_steps = 8000
+  hparams.learning_rate_constant = 2.0
+  hparams.layer_prepostprocess_dropout = 0.2
+  hparams.attention_dropout = 0.2
+  hparams.relu_dropout = 0.2
+  hparams.dropout = 0.2
+  # Optimization-related flags.
+  hparams.add_hparam("kl_reg", 0.0)
+  hparams.add_hparam("n_gibbs_steps", 0)
+  hparams.add_hparam("compute_kl_refinement", False)
+  hparams.add_hparam("compute_iw_marginal", False)
+  hparams.add_hparam("n_samples", 1)
+  return hparams
+
+
+@registry.register_hparams
+def wmt_diag_base():
+  """Set of hyperparameters."""
+  hparams = iwslt_diag()
+  hparams.batch_size = 4096
+  hparams.num_hidden_layers = 6
+  hparams.hidden_size = 512
+  hparams.filter_size = 2048
+  hparams.num_heads = 8
+  # VAE-related flags.
+  hparams.latent_size = 512
+  hparams.n_posterior_layers = 4
+  hparams.n_decoder_layers = 6
+  hparams.dropout = 0.1
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.attention_dropout = 0.1
+  hparams.relu_dropout = 0.1
+  return hparams
+
+
+@registry.register_hparams
+def wmt_diag_small():
+  """Set of hyperparameters."""
+  hparams = wmt_diag_base()
+  hparams.n_posterior_layers = 3
+  hparams.n_decoder_layers = 3
+  hparams.kl_reg = 1e-4
+  return hparams
+
+
+@registry.register_hparams
+def wmt_diag_small_trueadam():
+  """Set of hyperparameters."""
+  hparams = wmt_diag_small()
+  hparams.optimizer = "true_adam"
+  return hparams
+
+
+@registry.register_hparams
+def wmt_diag_small_trueadam_longer():
+  """Set of hyperparameters."""
+  hparams = wmt_diag_small_trueadam()
+  hparams.learning_rate_constant = 4.0
+  hparams.learning_rate_warmup_steps = 20000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_diag_small_trueadam_shorter():
+  """Set of hyperparameters."""
+  hparams = wmt_diag_small_trueadam()
+  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_warmup_steps = 4000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_diag_base_trueadam_1e4():
+  """Set of hyperparameters."""
+  hparams = wmt_diag_base()
+  hparams.kl_reg = 1e-4
+  hparams.optimizer = "true_adam"
+  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_warmup_steps = 8000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_diag_base_trueadam_longer_1e4():
+  """Set of hyperparameters."""
+  hparams = wmt_diag_base_trueadam_1e4()
+  hparams.learning_rate_constant = 4.0
+  hparams.learning_rate_warmup_steps = 20000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_diag_base_trueadam_shorter_1e4():
+  """Set of hyperparameters."""
+  hparams = wmt_diag_base_trueadam_1e4()
+  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_warmup_steps = 4000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_diag_base_1e4_trueadam():
+  """Set of hyperparameters."""
+  hparams = wmt_diag_base()
+  hparams.kl_reg = 1e-4
+  hparams.optimizer = "true_adam"
+  return hparams
+
+
+@registry.register_hparams
+def wmt_diag_base_1e4_trueadam_longer():
+  """Set of hyperparameters."""
+  hparams = wmt_diag_base_1e4_trueadam()
+  hparams.learning_rate_constant = 4.0
+  hparams.learning_rate_warmup_steps = 20000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_diag_base_1e4_trueadam_shorter():
+  """Set of hyperparameters."""
+  hparams = wmt_diag_base_1e4_trueadam()
+  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_warmup_steps = 4000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_diag_base_1e4():
+  """Set of hyperparameters."""
+  hparams = wmt_diag_base()
+  hparams.kl_reg = 1e-4
+  return hparams
+
+
+@registry.register_hparams
+def wmt_diag_base_longer_1e4():
+  """Set of hyperparameters."""
+  hparams = wmt_diag_base_1e4()
+  hparams.learning_rate_constant = 4.0
+  hparams.learning_rate_warmup_steps = 20000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_diag_base_shorter_1e4():
+  """Set of hyperparameters."""
+  hparams = wmt_diag_base_1e4()
+  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_warmup_steps = 4000
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_diag_1e5():
+  """Set of hyperparameters."""
+  hparams = iwslt_diag()
+  hparams.kl_reg = 1e-5
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_diag_1e4():
+  """Set of hyperparameters."""
+  hparams = iwslt_diag()
+  hparams.kl_reg = 1e-4
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_affine():
+  """Set of hyperparameters."""
+  hparams = iwslt_diag()
+  hparams.prior_type = "affine"
+  hparams.batch_size = 2048
+  hparams.latent_size = 256
+  # Glow-related flags.
+  hparams.add_hparam("depths", "4/8/8")  # infer n_levels from depths
+  hparams.add_hparam("step_fn", "glow")  # glow / chunting
+  hparams.add_hparam("affine_scale", "glow")  # glow / jason
+  hparams.add_hparam("conv_fn", "np")  # np / tf
+  hparams.add_hparam("split_plans", "cat/cat/ca")
+  hparams.add_hparam("factor", 2)  # squeezing factor
+  hparams.add_hparam("n_layers_transform_params", 1)
+  hparams.add_hparam("n_1x1_heads", 4)
+  hparams.add_hparam("flow_num_heads", 4)
+  hparams.add_hparam("flow_hidden_size", 256)
+  hparams.add_hparam("flow_filter_size", 512)
+  # Control max scale change.
+  hparams.add_hparam("scale_width", 0.999)
+  # Optimization-related flags.
+  # hparams.learning_rate_warmup_steps = 20000
+  hparams.add_hparam("flow_layer_prepostprocess_dropout", 0.0)
+  hparams.add_hparam("flow_attention_dropout", 0.0)
+  hparams.add_hparam("flow_relu_dropout", 0.0)
+  # hparams.optimizer_adam_beta1 = 0.9
+  # hparams.optimizer_adam_beta2 = 0.999
+  # hparams.optimizer_adam_epsilon = 1e-8
+  # Precision-related flags.
+  hparams.activation_dtype = "float32"
+  hparams.weight_dtype = "float32"
+
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine():
+  """Set of hyperparameters."""
+  hparams = iwslt_affine()
+  hparams.batch_size = 2048  # TODO(jason) : address this later.
+  hparams.num_hidden_layers = 6
+  hparams.hidden_size = 256
+  hparams.filter_size = 1024
+  hparams.num_heads = 8
+  # VAE-related flags.
+  hparams.latent_size = 256
+  hparams.n_posterior_layers = 4
+  hparams.n_decoder_layers = 4
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.attention_dropout = 0.1
+  hparams.relu_dropout = 0.1
+  # Glow-related flags.
+  hparams.flow_num_heads = 8
+  hparams.flow_filter_size = 512
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine_base():
+  """Set of hyperparameters."""
+  hparams = wmt_affine()
+  hparams.batch_size = 2048
+  hparams.hidden_size = 320
+  hparams.latent_size = 320
+  hparams.flow_filter_size = 640
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine_base_small():
+  """Set of hyperparameters."""
+  hparams = wmt_affine_base()
+  hparams.depths = "4/4/4"
+  hparams.kl_reg = 1e-4
+  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_warmup_steps = 8000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine_base_trueadam_small():
+  """Set of hyperparameters."""
+  hparams = wmt_affine_base_small()
+  hparams.optimizer = "true_adam"
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine_base_trueadam_longer_small():
+  """Set of hyperparameters."""
+  hparams = wmt_affine_base_trueadam_small()
+  hparams.learning_rate_constant = 4.0
+  hparams.learning_rate_warmup_steps = 20000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine_base_trueadam_shorter_small():
+  """Set of hyperparameters."""
+  hparams = wmt_affine_base_trueadam_small()
+  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_warmup_steps = 4000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine_base_trueadam():
+  """Set of hyperparameters."""
+  hparams = wmt_affine_base()
+  hparams.optimizer = "true_adam"
+  # hparams.optimizer_adam_beta1 = 0.9
+  # hparams.optimizer_adam_beta2 = 0.999
+  # hparams.optimizer_adam_epsilon = 1e-8
+  hparams.kl_reg = 1e-4
+  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_warmup_steps = 8000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine_base_trueadam_longer():
+  """Set of hyperparameters."""
+  hparams = wmt_affine_base_trueadam()
+  hparams.learning_rate_constant = 4.0
+  hparams.learning_rate_warmup_steps = 20000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine_base_trueadam_shorter():
+  """Set of hyperparameters."""
+  hparams = wmt_affine_base_trueadam()
+  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_warmup_steps = 4000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine_base_1e4():
+  """Set of hyperparameters."""
+  hparams = wmt_affine_base()
+  hparams.kl_reg = 1e-4
+  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_warmup_steps = 8000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine_base_longer_1e4():
+  """Set of hyperparameters."""
+  hparams = wmt_affine_base_1e4()
+  hparams.learning_rate_constant = 4.0
+  hparams.learning_rate_warmup_steps = 20000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine_base_shorter_1e4():
+  """Set of hyperparameters."""
+  hparams = wmt_affine_base_1e4()
+  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_warmup_steps = 4000
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine_1e4():
+  """Set of hyperparameters."""
+  hparams = wmt_affine()
+  hparams.kl_reg = 1e-4
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine_large():
+  """Set of hyperparameters."""
+  hparams = iwslt_affine()
+  hparams.batch_size = 2048
+  hparams.num_hidden_layers = 6
+  hparams.hidden_size = 512
+  hparams.filter_size = 1024
+  hparams.num_heads = 8
+  # VAE-related flags.
+  hparams.latent_size = 512
+  hparams.n_posterior_layers = 4
+  hparams.n_decoder_layers = 4
+  hparams.layer_prepostprocess_dropout = 0.1
+  hparams.attention_dropout = 0.1
+  hparams.relu_dropout = 0.1
+  # Glow-related flags.
+  hparams.flow_num_heads = 8
+  hparams.flow_filter_size = 1024
+  return hparams
+
+
+@registry.register_hparams
+def wmt_affine_large_1e4():
+  """Set of hyperparameters."""
+  hparams = wmt_affine_large()
+  hparams.kl_reg = 1e-4
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_affine_tiny():
+  """Set of hyperparameters."""
+  hparams = iwslt_affine()
+  hparams.depths = "1"
+  hparams.split_plans = "c"
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_affine_small():
+  """Set of hyperparameters."""
+  hparams = iwslt_affine()
+  hparams.depths = "4/4/4"
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_affine_small_1e4_trueadam():
+  """Set of hyperparameters."""
+  hparams = iwslt_affine_small_1e4()
+  hparams.optimizer = "true_adam"
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_affine_small_1e4_trueadam_longer():
+  """Set of hyperparameters."""
+  hparams = iwslt_affine_small_1e4_trueadam()
+  hparams.learning_rate_constant = 4.0
+  hparams.learning_rate_warmup_steps = 20000
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_affine_small_1e4_trueadam_shorter():
+  """Set of hyperparameters."""
+  hparams = iwslt_affine_small_1e4_trueadam()
+  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_warmup_steps = 4000
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_affine_small_1e4():
+  """Set of hyperparameters."""
+  hparams = iwslt_affine_small()
+  hparams.kl_reg = 1e-4
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_affine_tpu_glow_glow_np_1e4_trueadam():
+  """Set of hyperparameters."""
+  hparams = iwslt_affine_tpu_glow_glow_np_1e4()
+  hparams.optimizer = "true_adam"
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_affine_tpu_glow_glow_np_1e4_trueadam_longer():
+  """Set of hyperparameters."""
+  hparams = iwslt_affine_tpu_glow_glow_np_1e4_trueadam()
+  hparams.learning_rate_constant = 4.0
+  hparams.learning_rate_warmup_steps = 20000
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_affine_tpu_glow_glow_np_1e4_trueadam_shorter():
+  """Set of hyperparameters."""
+  hparams = iwslt_affine_tpu_glow_glow_np_1e4_trueadam()
+  hparams.learning_rate_constant = 2.0
+  hparams.learning_rate_warmup_steps = 4000
+  return hparams
+
+
+@registry.register_hparams
+def iwslt_affine_tpu_glow_glow_np_1e4():
+  """Set of hyperparameters."""
+  hparams = iwslt_affine()
+  hparams.conv_fn = "np"
+  hparams.kl_reg = 1e-4
+  return hparams
+
+
+def update_hparams_for_tpu(hparams):
+  """Change hparams to be compatible with TPU training."""
+
+  # Adafactor uses less memory than Adam.
+  # switch to Adafactor with its recommended learning rate scheme.
+  # hparams.optimizer = "Adafactor"
+  # hparams.learning_rate_schedule = "rsqrt_decay"
+  # hparams.learning_rate_warmup_steps = 10000
+
+  # Avoid an expensive concat on TPU.
+  # >1 shards helps with faster parameter distribution on multi-GPU machines
+  hparams.symbol_modality_num_shards = 1
+
+  # Adaptive batch sizes and sequence lengths are not supported on TPU.
+  # Instead, every batch has the same sequence length and the same batch size.
+  # Longer sequences are dropped and shorter ones are padded.
+  #
+  # It is therefore suggested to use a problem where examples have been combined
+  # to a longer length, e.g. the "_packed" problems.
+  #
+  # For problems with variable sequence lengths, this parameter controls the
+  # maximum sequence length.  Shorter sequences are dropped and longer ones
+  # are padded.
+  #
+  # For problems with fixed sequence lengths - e.g. the "_packed" problems,
+  # this hyperparameter is ignored.
+  hparams.max_length = 64
+
+  # TPUs have less memory than GPUs, so decrease the batch size if it's too high
+  if hparams.batch_size > 2048:
+    hparams.batch_size = 2048
+
+  # Using noise broadcast in the dropout layers saves memory during training.
+  hparams.attention_dropout_broadcast_dims = "0,1"  # batch, heads
+  hparams.relu_dropout_broadcast_dims = "1"  # length
+  hparams.layer_prepostprocess_dropout_broadcast_dims = "1"  # length
+  return hparams
diff --git a/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py b/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
new file mode 100644
index 000000000..ddf742690
--- /dev/null
+++ b/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
@@ -0,0 +1,364 @@
+# coding=utf-8
+# Copyright 2020 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Various ops for TransformerVaeFlowPrior."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import transformer_glow_layers_ops as gops
+from tensor2tensor.models.transformer import transformer_decoder_layer
+from tensor2tensor.models.transformer import transformer_encoder
+from tensor2tensor.models.transformer import transformer_prepare_encoder
+from tensor2tensor.utils import learning_rate as lr
+from tensor2tensor.utils import mlperf_log
+import tensorflow.compat.v1 as tf
+
+
+def _mixed_precision_is_enabled(hparams):
+  """Should be the same as in common_attention, avoiding import."""
+  activation_dtype = hparams.activation_dtype
+  weight_dtype = hparams.weight_dtype
+  return activation_dtype == tf.float16 and weight_dtype == tf.float32
+
+
+def encoder(name, hparams, inputs, target_space):
+  """Compute encoder outputs and attention bias."""
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    (encoder_input,
+     encoder_self_attention_bias,
+     encoder_decoder_attention_bias) = (
+         transformer_prepare_encoder(inputs, target_space, hparams))
+    encoder_input = tf.nn.dropout(encoder_input,
+                                  rate=hparams.layer_prepostprocess_dropout)
+    encoder_output = transformer_encoder(encoder_input,
+                                         encoder_self_attention_bias,
+                                         hparams)
+    return encoder_output, encoder_decoder_attention_bias
+
+
+def transformer_decoder_layers(name,
+                               n_layers,
+                               decoder_input,
+                               **kwargs):
+  """A transformation block composed of transformer decoder layers."""
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    hparams = kwargs["hparams"]
+    outputs = decoder_input
+    with tf.variable_scope("decoder", reuse=tf.AUTO_REUSE):
+      for layer_idx in range(n_layers):
+        outputs = transformer_decoder_layer(
+            decoder_input=outputs,
+            layer_idx=layer_idx,
+            **kwargs)
+      outputs = common_layers.layer_preprocess(outputs, hparams)
+    return outputs
+
+
+def posterior(
+    name, hparams, targets, targets_mask, decoder_self_attention_bias,
+    **kwargs):
+  """Compute mu and sigma for diagonal normal posterior q(z|x,y)."""
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    decoder_input = drop_2d(targets, hparams.mode, hparams.posterior_2d_dropout)
+    decoder_input = common_attention.add_timing_signal_1d(decoder_input)
+    decoder_input = tf.nn.dropout(decoder_input,
+                                  rate=hparams.layer_prepostprocess_dropout)
+    decoder_output = transformer_decoder_layers(
+        "block",
+        n_layers=hparams.n_posterior_layers,
+        decoder_input=decoder_input,
+        hparams=hparams,
+        decoder_self_attention_bias=decoder_self_attention_bias,
+        **kwargs)
+    decoder_output = gops.dense_weightnorm(
+        "h2o_out", decoder_output, hparams.latent_size * 2, targets_mask,
+        init_scale=0.0, init=False)
+    return decoder_output
+
+
+def cond_prior(
+    name, hparams, decoder_input, targets_mask, output_size,
+    decoder_self_attention_bias, init_scale=0.0, **kwargs):
+  """Compute hidden states for parameters for conditional prior."""
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    decoder_input = common_attention.add_timing_signal_1d(decoder_input)
+    decoder_input = tf.nn.dropout(decoder_input,
+                                  rate=hparams.layer_prepostprocess_dropout)
+    decoder_output = transformer_decoder_layers(
+        "block",
+        n_layers=hparams.n_posterior_layers,
+        decoder_input=decoder_input,
+        hparams=hparams,
+        decoder_self_attention_bias=decoder_self_attention_bias,
+        **kwargs)
+    decoder_output = gops.dense_weightnorm(
+        "h2o_out", decoder_output, output_size, targets_mask,
+        init_scale=init_scale, init=False)
+    return decoder_output
+
+
+def decoder(name, latents, hparams, decoder_self_attention_bias, **kwargs):
+  """Compute final hidden states for p(y|z,x)."""
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    decoder_input = drop_2d(latents, hparams.mode, hparams.decoder_2d_dropout)
+    if hparams.pos_attn:
+      decoder_input = gops.positional_attention(
+          "pos_attn", decoder_input, decoder_self_attention_bias, hparams)
+    else:
+      decoder_input = common_attention.add_timing_signal_1d(decoder_input)
+    if common_layers.shape_list(latents)[-1] != hparams.hidden_size:
+      decoder_input = gops.dense("lat2hid", latents, hparams.hidden_size)
+    decoder_output = transformer_decoder_layers(
+        "block",
+        n_layers=hparams.n_decoder_layers,
+        decoder_input=decoder_input,
+        hparams=hparams,
+        decoder_self_attention_bias=decoder_self_attention_bias,
+        **kwargs)
+    batch_size, targets_length = common_layers.shape_list(decoder_output)[:2]
+    decoder_output = tf.reshape(
+        decoder_output, [batch_size, targets_length, 1, hparams.hidden_size])
+    # Expand since t2t expects 4d tensors.
+    return decoder_output
+
+
+def drop_2d(targets, mode, dropout_p):
+  """Dropout in 2D."""
+  if dropout_p > 0 and mode == tf.estimator.ModeKeys.TRAIN:
+    batch_size, targets_length, hidden_size = common_layers.shape_list(targets)
+    mask_prob = tf.random_uniform(
+        shape=(batch_size, targets_length), minval=0.0, maxval=1.0)
+    mask_prob = tf.tile(mask_prob[..., tf.newaxis], [1, 1, hidden_size])
+    scale = 1 / (1 - dropout_p)
+    targets_noisy = tf.where(
+        mask_prob > dropout_p, targets * scale, tf.zeros_like(targets))
+    return targets_noisy
+  return targets
+
+
+def sequence_mask(length, hparams):
+  dtype = get_dtype(hparams)
+  return tf.sequence_mask(length, dtype=dtype)
+
+
+def get_padding(mask, hparams):
+  dtype = get_dtype(hparams)
+  return tf.cast(tf.equal(mask, 0.0), dtype=dtype)
+
+
+def get_dtype(hparams):
+  if hparams.activation_dtype == "float32":
+    return tf.float32
+  elif hparams.activation_dtype == "float64":
+    return tf.float64
+  elif hparams.activation_dtype == "bfloat16":
+    return tf.bfloat16
+  else:
+    return None
+
+
+def lenpred_mlp(name, logits, hidden_size, bound):
+  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+    logits = tf.layers.dense(logits, hidden_size)
+    logits = tf.nn.elu(logits)
+    logits = tf.layers.dense(logits, hidden_size)
+    logits = tf.nn.elu(logits)
+    logits = tf.layers.dense(logits, bound * 2 + 1)
+  return logits
+
+
+def predict_target_lengths(
+    encoder_output, inputs_mask, hparams, length_diff=None):
+  """Predict target lengths."""
+  bound = hparams.lendiff_bound
+  inputs_length = tf.cast(tf.reduce_sum(inputs_mask, 1), tf.int32)
+  targets_length = inputs_length
+  loss = None
+  if hparams.predict_target_length:
+    encoder_output = gops.reduce_mean_over_l(encoder_output, inputs_mask)
+    logits = tf.stop_gradient(encoder_output)
+    logits = lenpred_mlp("lenpred", logits, hparams.hidden_size, bound)
+    if length_diff is not None:
+      labels = tf.maximum(tf.minimum(length_diff, bound), -bound)
+      labels = tf.cast(labels + bound, tf.int32)
+      labels = tf.stop_gradient(labels)
+      loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
+          labels=labels, logits=logits)
+      loss = tf.reduce_mean(loss)
+    diff_pred = tf.argmax(logits, 1)
+    diff_pred = tf.cast(diff_pred - bound, tf.int32)
+    targets_length = inputs_length + diff_pred
+    targets_length = tf.maximum(targets_length, 1)
+  divi = 4
+  targets_length = tf.ceil(targets_length / divi) * divi
+  targets_length = tf.cast(targets_length, tf.int32)
+  return targets_length, loss
+
+
+def lenpred_stats(targets_length_pred, targets_length):
+  lenpred_diff = tf.abs(targets_length_pred - tf.cast(targets_length, tf.int32))
+  lenpred_acc = tf.cast(tf.equal(lenpred_diff, 0), tf.float32)
+  lenpred_acc = tf.reduce_mean(lenpred_acc)
+  lenpred_acc5 = tf.cast(tf.less_equal(lenpred_diff, 5), tf.float32)
+  lenpred_acc5 = tf.reduce_mean(lenpred_acc5)
+  return lenpred_acc, lenpred_acc5
+
+
+def save_log_loss(
+    hparams, targets_mask, numerator, denominator, log_q_z, log_abs_det,
+    log_p_z_base, z_q, lenpred_loss, targets_length_pred, targets_length):
+  """Populate loss dictionary and summary."""
+  anneal, kl_mask = get_anneal_mask(hparams)
+  lenpred_acc, lenpred_acc5 = (
+      lenpred_stats(targets_length_pred, targets_length))
+  batch_length = tf.reduce_sum(targets_mask)
+
+  z_q_norm = gops.reduce_mean_over_bl(
+      tf.norm(z_q, axis=2, keepdims=True), targets_mask)[0]
+
+  log_q_z = gops.reduce_mean_over_bl_sum_over_c(log_q_z, targets_mask)
+  log_p_z_base = tf.reduce_sum(log_p_z_base, axis=0) / batch_length
+  log_abs_det = tf.reduce_sum(log_abs_det, axis=0) / batch_length
+  log_p_z_reg = gops.standard_normal_density(z_q, targets_mask, reduce_sum=True)
+
+  log_p_x = -1 * numerator / denominator
+  log_p_z = log_p_z_base + log_abs_det
+  kl = log_q_z - log_p_z
+  kl_reg = log_p_z - log_p_z_reg
+  elbo = log_p_x - kl
+  monitor = {
+      "elbo": elbo,
+      "kl": kl,
+      "kl_reg": kl_reg,
+      "log_p_x": log_p_x,
+      "log_q_z": log_q_z,
+      "log_p_z": log_p_z,
+      "log_p_z_base": log_p_z_base,
+      "log_abs_det": log_abs_det,
+      "anneal": anneal,
+      "z_q_norm": z_q_norm,
+      "lenpred_acc": lenpred_acc,
+      "lenpred_acc5": lenpred_acc5,
+  }
+
+  kl = kl * anneal
+  kl_reg = hparams.kl_reg * kl_reg * anneal
+  loss_dict = {
+      "training": -1 * log_p_x,
+      "kl": kl * kl_mask,
+      "kl_reg": kl_reg * kl_mask,
+  }
+  if lenpred_loss is not None:
+    monitor["lenpred_loss"] = lenpred_loss
+    loss_dict["lenpred_loss"] = lenpred_loss
+  return loss_dict, monitor
+
+
+def get_anneal_mask(hparams):
+  """Get anneal and kl mask."""
+  startup = hparams.kl_startup_steps
+  anneal = hparams.kl_anneal_steps
+  global_step = tf.train.get_global_step()
+  min_value = hparams.anneal_min_value
+  step = tf.maximum(global_step - startup, 0)
+  anneal = common_layers.inverse_lin_decay(
+      anneal, min_value=min_value, step=step)
+  kl_mask = tf.less(startup, tf.to_int32(global_step))
+  kl_mask = tf.cast(kl_mask, tf.float32)
+  return anneal, kl_mask
+
+
+def embedding_to_non_padding(emb, dtype=tf.float32):
+  """Calculates the padding mask based on which embeddings are not zero."""
+  emb_sum = tf.reduce_sum(tf.abs(emb), axis=-1)
+  return tf.cast(tf.not_equal(emb_sum, 0.0), dtype=dtype)
+
+
+def save_summary(monitor, name):
+  with tf.name_scope(name):
+    for key in list(monitor.keys()):
+      tf.summary.scalar(key, monitor[key])
+
+
+def _global_step(hparams):
+  """Adjust global step if a multi-step optimizer is used."""
+  step = tf.cast(tf.train.get_or_create_global_step(), tf.float32)
+  multiplier = hparams.optimizer_multistep_accumulate_steps
+  if not multiplier:
+    return step
+
+  tf.logging.info("Dividing global step by %d for multi-step optimizer."
+                  % multiplier)
+  return step / tf.cast(multiplier, tf.float32)
+
+
+def learning_rate_schedule(hparams):
+  """Learning rate schedule based on hparams."""
+  mlperf_log.transformer_print(key=mlperf_log.OPT_LR, deferred=True)
+  mlperf_log.transformer_print(
+      key=mlperf_log.OPT_LR_WARMUP_STEPS,
+      value=hparams.learning_rate_warmup_steps)
+  step_num = _global_step(hparams)
+  # Simulate pretraining the encoder, decoder and posterior with the same
+  # learning rate schedule, and then restoring the parameters.
+  # using `warm_start_from` is not compatible with actnorm DDI on TPUs.
+  step_num = tf.where(
+      step_num < hparams.kl_startup_steps,
+      step_num,
+      step_num - hparams.kl_startup_steps)
+  schedule_string = hparams.learning_rate_schedule
+  names = schedule_string.split("*")
+  names = [name.strip() for name in names if name.strip()]
+  ret = tf.constant(1.0)
+  for name in names:
+    ret *= lr.learning_rate_factor(name, step_num, hparams)
+  return ret
+
+
+def prepare_for_iw(x, k):
+  """Prepare feature for importance sampling."""
+  batch_size = common_layers.shape_list(x)[0]
+  remaining_shape = common_layers.shape_list(x)[1:]
+
+  multiplier = [1] * x.shape.rank
+  x = tf.tile(x[tf.newaxis, ...], [k] + multiplier)
+  x = tf.reshape(x, [k * batch_size] + remaining_shape)
+  return x
+
+
+def unprepare_for_iw(x, k):
+  """Unprepare feature for importance sampling."""
+  batch_size_times_k = common_layers.shape_list(x)[0]
+  remaining_shape = common_layers.shape_list(x)[1:]
+  x = tf.reshape(x, [k, batch_size_times_k // k] + remaining_shape)
+  return x
+
+
+def generic_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
+  """Compute loss numerator and denominator for one shard of output."""
+  del vocab_size  # unused arg
+  logits = top_out
+  logits = common_attention.maybe_upcast(logits, hparams=model_hparams)
+  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.0)
+  return common_layers.padded_cross_entropy(
+      logits,
+      targets,
+      model_hparams.label_smoothing,
+      cutoff=cutoff,
+      weights_fn=weights_fn,
+      reduce_sum=False)

From f19ca736d06b0b1dd3d83655544afdb14f912bdb Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 14 Feb 2020 10:26:17 -0800
Subject: [PATCH 2644/2720] Internal change

PiperOrigin-RevId: 295173851
---
 tensor2tensor/utils/checkpoint_compatibility_test.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/utils/checkpoint_compatibility_test.py b/tensor2tensor/utils/checkpoint_compatibility_test.py
index 05bb4dd10..cd466a641 100644
--- a/tensor2tensor/utils/checkpoint_compatibility_test.py
+++ b/tensor2tensor/utils/checkpoint_compatibility_test.py
@@ -32,6 +32,7 @@
 import os
 import numpy as np
 
+from six.moves import range
 from tensor2tensor import models  # pylint: disable=unused-import
 from tensor2tensor import problems  # pylint: disable=unused-import
 from tensor2tensor.utils import data_reader

From 78ba8019847426e988294fd58f8953d7990a8db7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 20 Feb 2020 18:13:01 -0800
Subject: [PATCH 2645/2720] Internal change

PiperOrigin-RevId: 296341441
---
 .../data_generators/ops/subword_text_encoder.cc      | 12 ++++++------
 .../data_generators/ops/subword_text_encoder.h       |  6 +++---
 .../data_generators/ops/subword_text_encoder_ops.cc  |  2 +-
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder.cc b/tensor2tensor/data_generators/ops/subword_text_encoder.cc
index bc3d22134..dfa2470cc 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder.cc
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder.cc
@@ -18,9 +18,9 @@ constexpr int64 kEosTokenId = 1;
 
 }  // namespace
 
-SubwordTextEncoder::SubwordTextEncoder(const string& vocab_filename) {
+SubwordTextEncoder::SubwordTextEncoder(const std::string& vocab_filename) {
   // TODO(ormandi): Add a unified vocabulary reader function.
-  string vocab_contents;
+  std::string vocab_contents;
   TF_CHECK_OK(
       ReadFileToString(Env::Default(), vocab_filename, &vocab_contents));
   std::vector<absl::string_view> vocab_list =
@@ -42,7 +42,7 @@ SubwordTextEncoder::SubwordTextEncoder(const string& vocab_filename) {
       CHECK_GE(c, 0);
       alphabet_.insert(c);
     } while (char_index < token.length());
-    vocab_.insert({string(token), i});
+    vocab_.insert({std::string(token), i});
   }
 }
 
@@ -76,7 +76,7 @@ void SubwordTextEncoder::Encode(absl::string_view text, std::vector<int>* ids) {
 
 void SubwordTextEncoder::EncodeSubtokens(
     absl::string_view token, std::vector<int> *ids) {
-  string token_s = EscapeToken(token);
+  std::string token_s = EscapeToken(token);
   token = token_s;
   int subtoken_start = 0;
   // TODO(noam): this algorithm is quadratic in the length of the token.
@@ -102,8 +102,8 @@ void SubwordTextEncoder::EncodeSubtokens(
   }
 }
 
-string SubwordTextEncoder::EscapeToken(absl::string_view token) {
-  string token_s;
+std::string SubwordTextEncoder::EscapeToken(absl::string_view token) {
+  std::string token_s;
   int i = 0;
   do {
     int prev = i;
diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder.h b/tensor2tensor/data_generators/ops/subword_text_encoder.h
index a06cdc2d2..17cb2bf66 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder.h
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder.h
@@ -18,7 +18,7 @@ namespace tensor2tensor {
 
 class SubwordTextEncoder {
  public:
-  explicit SubwordTextEncoder(const string& vocab_filename);
+  explicit SubwordTextEncoder(const std::string& vocab_filename);
   virtual ~SubwordTextEncoder() {}
 
   // Breaks up input text into subtokens.
@@ -30,10 +30,10 @@ class SubwordTextEncoder {
   void EncodeSubtokens(absl::string_view token, std::vector<int>* ids);
 
   // Escapes a token so unencodable characters are replaced by escape sequences.
-  string EscapeToken(absl::string_view token);
+  std::string EscapeToken(absl::string_view token);
 
   // Maps subword tokens to IDs.
-  absl::flat_hash_map<string, int64> vocab_;
+  absl::flat_hash_map<std::string, int64> vocab_;
   // A set containing all valid unicode code points that can be encoded without
   // being escaped.
   absl::flat_hash_set<UChar32> alphabet_;
diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc b/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
index 9e8959f96..afa8813bd 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
@@ -30,7 +30,7 @@ class SubwordTextEncoderEncodeOp : public OpKernel {
  public:
   explicit SubwordTextEncoderEncodeOp(
       OpKernelConstruction* ctx) : OpKernel(ctx) {
-    string vocab_filename;
+    std::string vocab_filename;
     OP_REQUIRES_OK(ctx, ctx->GetAttr("vocab_filename", &vocab_filename));
     encoder_ = absl::make_unique<SubwordTextEncoder>(vocab_filename);
   }

From 281ded6e7313f3ac18d112e5d53dd98b0085881d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 24 Feb 2020 13:43:19 -0800
Subject: [PATCH 2646/2720] Internal change

PiperOrigin-RevId: 296964212
---
 tensor2tensor/models/research/transformer_vae_flow_prior_ops.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py b/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
index ddf742690..d23a50a57 100644
--- a/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
+++ b/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from six.moves import range
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import transformer_glow_layers_ops as gops

From d0f00e1e30dcf3383ca3a7274cbaab5abe2f4d5a Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 26 Feb 2020 14:36:33 -0800
Subject: [PATCH 2647/2720] Internal change

PiperOrigin-RevId: 297453797
---
 tensor2tensor/utils/compute_video_metrics.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/utils/compute_video_metrics.py b/tensor2tensor/utils/compute_video_metrics.py
index 012ca5b82..9a025d5e5 100644
--- a/tensor2tensor/utils/compute_video_metrics.py
+++ b/tensor2tensor/utils/compute_video_metrics.py
@@ -21,6 +21,7 @@
 
 import os
 
+from six.moves import range
 from tensor2tensor.bin import t2t_decoder
 from tensor2tensor.utils import video_metrics
 import tensorflow.compat.v1 as tf

From f2265b2b4b15af3f8f6d8b7f4b19defb8a4ddb3b Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 26 Feb 2020 16:42:22 -0800
Subject: [PATCH 2648/2720] Internal change

PiperOrigin-RevId: 297471304
---
 tensor2tensor/models/video/next_frame_glow.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index d1f32b479..f95872df3 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import numpy as np
+from six.moves import range
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import common_video
 from tensor2tensor.layers import modalities

From 709525aad57fb2c23600f79c2d4289f0cba1a714 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 27 Feb 2020 10:55:11 -0800
Subject: [PATCH 2649/2720] Internal change

PiperOrigin-RevId: 297637517
---
 tensor2tensor/models/video/nfg_interpolate.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensor2tensor/models/video/nfg_interpolate.py b/tensor2tensor/models/video/nfg_interpolate.py
index d076ad579..3c3bb7d3f 100644
--- a/tensor2tensor/models/video/nfg_interpolate.py
+++ b/tensor2tensor/models/video/nfg_interpolate.py
@@ -21,6 +21,7 @@
 import os
 from absl import flags
 import numpy as np
+from six.moves import zip
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
 from tensor2tensor.data_generators import image_utils
 from tensor2tensor.layers import common_layers

From 06cd379c0204084297ab3dacaa37526019624865 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 2 Mar 2020 12:13:57 -0800
Subject: [PATCH 2650/2720] Fix an off-by-one Tensor size error when using
 'pos=timing_from_features' hparam.

PiperOrigin-RevId: 298413010
---
 tensor2tensor/models/transformer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index afbd40a18..89a983556 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -761,7 +761,7 @@ def _fast_decode(self,
           decode_length + 1, hparams.hidden_size)
     elif hparams.pos == "timing_from_features":
       positional_encoding = common_attention.add_timing_signals_from_features(
-          tf.zeros([1, decode_length + 1, hparams.hidden_size]), features,
+          tf.zeros([1, decode_length, hparams.hidden_size]), features,
           hparams.position_features)
     elif hparams.pos == "emb":
       positional_encoding = common_attention.add_positional_embedding(

From c6a09ead7619c7a9b0b6bdc5e6678438b153caec Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 9 Mar 2020 11:15:53 -0700
Subject: [PATCH 2651/2720] Internal.

PiperOrigin-RevId: 299886679
---
 tensor2tensor/models/transformer.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 89a983556..5713bda28 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -675,7 +675,8 @@ def _fast_decode(self,
                    decode_length,
                    beam_size=1,
                    top_beams=1,
-                   alpha=1.0):
+                   alpha=1.0,
+                   preprocess_targets_method=None):
     """Fast decoding.
 
     Implements both greedy and beam search decoding, uses beam search iff
@@ -688,6 +689,8 @@ def _fast_decode(self,
       top_beams: an integer. How many of the beams to return.
       alpha: Float that controls the length penalty. larger the alpha, stronger
         the preference for longer translations.
+      preprocess_targets_method: method used to preprocess targets. If None,
+      uses method "preprocess_targets" defined inside this method.
 
     Returns:
       A dict of decoding results {
@@ -842,12 +845,14 @@ def update_decoder_attention_history(cache):
               [cache["attention_history"][layer_nbr],
                self.attention_weights[k]],
               axis=2)
+    if not preprocess_targets_method:
+      preprocess_targets_method = preprocess_targets
 
     def symbols_to_logits_fn(ids, i, cache):
       """Go from ids to logits for next symbol."""
       ids = ids[:, -1:]
       targets = tf.expand_dims(tf.expand_dims(ids, axis=2), axis=3)
-      targets = preprocess_targets(targets, i)
+      targets = preprocess_targets_method(targets, i)
 
       bias = decoder_self_attention_bias[:, :, i:i + 1, :i + 1]
       with tf.variable_scope("body"):

From f1901df7aed13ef94d3fd52488e61f3ecf3f2095 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Mon, 9 Mar 2020 13:00:57 -0700
Subject: [PATCH 2652/2720] Log the number of discretized actions only once,
 regarless of the train/eval batch size.

PiperOrigin-RevId: 299913269
---
 tensor2tensor/rl/gym_utils.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index d1dc7840a..44abdc1c5 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -285,7 +285,8 @@ def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env,
     env = remove_time_limit_wrapper(env)
 
   if num_actions is not None:
-    logging.info("Number of discretized actions: %d", num_actions)
+    logging.log_every_n(
+        logging.INFO, "Number of discretized actions: %d", 1, num_actions)
     env = ActionDiscretizeWrapper(env, num_actions=num_actions)
 
   if sticky_actions:

From 65516f687d9cca9081c94e023693807cd05ac5dd Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 10 Mar 2020 11:33:25 -0700
Subject: [PATCH 2653/2720] Log the number of discretized actions only once,
 regarless of the train/eval batch size.

PiperOrigin-RevId: 300139035
---
 tensor2tensor/rl/gym_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index 44abdc1c5..1d134f10b 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -285,7 +285,7 @@ def gym_env_wrapper(env, rl_env_max_episode_steps, maxskip_env, rendered_env,
     env = remove_time_limit_wrapper(env)
 
   if num_actions is not None:
-    logging.log_every_n(
+    logging.log_first_n(
         logging.INFO, "Number of discretized actions: %d", 1, num_actions)
     env = ActionDiscretizeWrapper(env, num_actions=num_actions)
 

From 543293cfb490121f67c6d87287b9b0c7c3670286 Mon Sep 17 00:00:00 2001
From: Daniel Furrer <danielfurrer@google.com>
Date: Wed, 18 Mar 2020 08:05:18 -0700
Subject: [PATCH 2654/2720] Make it clear in requirements that tensor2tensor
 doesn't support TF2.

PiperOrigin-RevId: 301590259
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 1482f9f7a..5948d8918 100644
--- a/setup.py
+++ b/setup.py
@@ -69,7 +69,7 @@
         'tqdm',
     ],
     extras_require={
-        'tensorflow': ['tensorflow>=1.15.0'],
+        'tensorflow': ['tensorflow>=1.15.0,<2.0'],
         'tensorflow-hub': ['tensorflow-hub>=0.1.1'],
         'tests': [
             # Needed to fix a Travis pytest error.

From 1c965c47d1bb3bfb54478453005cae77b1637f75 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 23 Mar 2020 15:46:29 -0700
Subject: [PATCH 2655/2720] Fix two tensor2tensor Colab notebooks.

PiperOrigin-RevId: 302535311
---
 tensor2tensor/notebooks/t2t_problem.ipynb | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/notebooks/t2t_problem.ipynb b/tensor2tensor/notebooks/t2t_problem.ipynb
index 6e1eb1311..98e07fcbb 100644
--- a/tensor2tensor/notebooks/t2t_problem.ipynb
+++ b/tensor2tensor/notebooks/t2t_problem.ipynb
@@ -127,7 +127,7 @@
         "import tensorflow as tf\n",
         "\n",
         "# Enable Eager execution - useful for seeing the generated data.\n",
-        "tf.enable_eager_execution()"
+        "tf.compat.v1.enable_eager_execution()"
       ]
     },
     {
@@ -171,9 +171,9 @@
         "TMP_DIR = os.path.expanduser(\"/tmp/t2t/tmp\")\n",
         "\n",
         "# Create them.\n",
-        "tf.gfile.MakeDirs(DATA_DIR)\n",
-        "tf.gfile.MakeDirs(OUTPUT_DIR)\n",
-        "tf.gfile.MakeDirs(TMP_DIR)"
+        "tf.io.gfile.makedirs(DATA_DIR)\n",
+        "tf.io.gfile.makedirs(OUTPUT_DIR)\n",
+        "tf.io.gfile.makedirs(TMP_DIR)"
       ]
     },
     {

From 35789ec177ef5e9de739a4ceb7b5b0e37a73cef3 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 23 Mar 2020 16:15:21 -0700
Subject: [PATCH 2656/2720] Allow for the export of extra outputs from Infer().

PiperOrigin-RevId: 302541178
---
 tensor2tensor/utils/decoding.py  | 4 +++-
 tensor2tensor/utils/t2t_model.py | 4 ++++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 2a7ca2fcb..99ffdb0fb 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -97,7 +97,9 @@ def decode_hparams(overrides=""):
       # Used for MLPerf compliance logging.
       mlperf_decode_step=0.0,
       mlperf_threshold=25.0,
-      mlperf_success=False)
+      mlperf_success=False,
+      # A comma-delimited list of additional infer() outputs to be exported.
+      export_extra_infer_outputs="")
   hp.parse(overrides)
   return hp
 
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 1fd95d931..94ae4bf3b 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1744,6 +1744,10 @@ def estimator_spec_predict(self, features, use_tpu=False):
     if "scores" in predictions:
       export_out["scores"] = predictions["scores"]
 
+    if decode_hparams.get("export_extra_infer_outputs"):
+      for output in decode_hparams.export_extra_infer_outputs.split(","):
+        export_out[output] = infer_out[output]
+
     # Necessary to rejoin examples in the correct order with the Cloud ML Engine
     # batch prediction API.
     if "batch_prediction_key" in predictions:

From 598eace7707d77b7c5117a2235ac3c4924796ca2 Mon Sep 17 00:00:00 2001
From: Mark Daoust <markdaoust@google.com>
Date: Mon, 23 Mar 2020 20:51:54 -0700
Subject: [PATCH 2657/2720] Delete references to
 github.com/tensorflow/models/blob/master/tutorials

The directory is not maintained, and is no longer being used for tensorflow.org.

I am deleting it in https://github.com/tensorflow/models/pull/8276

PiperOrigin-RevId: 302583413
---
 tensor2tensor/data_generators/ptb.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tensor2tensor/data_generators/ptb.py b/tensor2tensor/data_generators/ptb.py
index e02aed4a2..7283ce23f 100644
--- a/tensor2tensor/data_generators/ptb.py
+++ b/tensor2tensor/data_generators/ptb.py
@@ -49,8 +49,6 @@ def _build_vocab(filename, vocab_path, vocab_size):
   """Reads a file to build a vocabulary of `vocab_size` most common words.
 
    The vocabulary is sorted by occurrence count and has one word per line.
-   Originally from:
-   https://github.com/tensorflow/models/blob/master/tutorials/rnn/ptb/reader.py
 
   Args:
     filename: file to read list of words from.

From c6c31096ef931aeb18bd006ac92a87baedbcb65e Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Fri, 17 Apr 2020 07:04:54 -0700
Subject: [PATCH 2658/2720] Set an older gym version to correct problems with
 OSS TF/gym interaction and segfaults.

PiperOrigin-RevId: 307036899
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 5948d8918..5a76eb641 100644
--- a/setup.py
+++ b/setup.py
@@ -50,7 +50,7 @@
         'gin-config',
         'google-api-python-client',
         'gunicorn',
-        'gym',
+        'gym==0.14.0',
         'h5py',
         'kfac',
         'mesh-tensorflow',

From 12d63a3b39e19967ccf7dd8c68a33bf0cec9c0ad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mat=C3=BA=C5=A1=20=C5=BDilinec?= <zilinec.m@gmail.com>
Date: Sat, 18 Apr 2020 03:12:01 +0200
Subject: [PATCH 2659/2720] Add option to select a different TPU zone in
 create_run_config. (#1788)

---
 tensor2tensor/utils/trainer_lib.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index cfd516c4f..8f758fa0c 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -178,7 +178,8 @@ def create_run_config(model_name,
                       log_step_count_steps=100,
                       intra_op_parallelism_threads=0,
                       tpu_config_extra_kwargs=None,
-                      cloud_tpu_name=""):
+                      cloud_tpu_name="",
+                      cloud_tpu_zone=None):
   """Create RunConfig, TPUConfig, and Parallelism object."""
   session_config = create_session_config(
       log_device_placement=log_device_placement,
@@ -229,7 +230,7 @@ def create_run_config(model_name,
       # Update run_config to use cluster instead of master/evaluation_master
       # as we need the cluster spec to use Cloud Pods
       tpu_cluster_resolver = contrib.cluster_resolver().TPUClusterResolver(
-          cloud_tpu_name)
+          tpu=cloud_tpu_name, zone=cloud_tpu_zone)
       run_config_args["cluster"] = tpu_cluster_resolver
       del run_config_args["master"]
       del run_config_args["evaluation_master"]

From f5d73746f7a46dc18fdd541b1f9265c7f3ad2918 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Sat, 18 Apr 2020 09:31:22 -0700
Subject: [PATCH 2660/2720] Bump T2T version to 1.15.5

With this change Travis will turn green.

When we import compat.v1, we shouldn't do tf.compat.v1 anymore.

Also the following to unclog Travis:
six>=1.12.0 is required by tf-hub 0.8.0, one of the core dependencies seems to
be installing tf-hub, so pinning tf-hub to 0.7.0 (which is in the extras
section) doesn't seem to help.

Locally the previous setup.py installs fine, but only errors on Travis for some
reason.

PiperOrigin-RevId: 307203166
---
 setup.py                                      | 4 ++--
 tensor2tensor/layers/common_attention_test.py | 2 +-
 tensor2tensor/layers/common_layers.py         | 8 ++++----
 tensor2tensor/layers/common_layers_test.py    | 2 +-
 tensor2tensor/utils/avg_checkpoints.py        | 2 +-
 5 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/setup.py b/setup.py
index 5a76eb641..36e7e4a86 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.15.4',
+    version='1.15.5',
     description='Tensor2Tensor',
     long_description=(
         'Tensor2Tensor, or T2T for short, is a library of '
@@ -61,7 +61,7 @@
         'pypng',
         'requests',
         'scipy',
-        'six',
+        'six>=1.12.0',
         'sympy',
         'tensorflow-datasets',
         'tensorflow-gan',
diff --git a/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index f11b283e5..79aa383ca 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -34,7 +34,7 @@
 
 tfe = contrib.tfe()
 # from tensorflow.contrib.eager.python import tfe as tfe
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class CommonAttentionTest(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index c0193fbb4..f6d680bb3 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2743,7 +2743,7 @@ def _fn_with_custom_grad(fn, inputs, grad_fn, use_global_vars=False):
   Returns:
     fn(*inputs)
   """
-  vs = tf.compat.v1.get_variable_scope()
+  vs = tf.get_variable_scope()
   get_vars_fn = (
       vs.global_variables if use_global_vars else vs.trainable_variables)
   len_before_vars = len(get_vars_fn())
@@ -3145,7 +3145,7 @@ def grad_fn(inputs, variables, outputs, output_grads):
 
   @fn_with_custom_grad(grad_fn)
   def fn_with_recompute(*args):
-    cached_vs.append(tf.compat.v1.get_variable_scope())
+    cached_vs.append(tf.get_variable_scope())
     cached_arg_scope.append(contrib.framework().current_arg_scope())
     return fn(*args)
 
@@ -3160,7 +3160,7 @@ def dense(x, units, **kwargs):
     # We need to find the layer parameters using scope name for the layer, so
     # check that the layer is named. Otherwise parameters for different layers
     # may get mixed up.
-    layer_name = tf.compat.v1.get_variable_scope().name
+    layer_name = tf.get_variable_scope().name
     if (not layer_name) or ("name" not in kwargs):
       raise ValueError(
           "Variable scope and layer name cannot be empty. Actual: "
@@ -3491,7 +3491,7 @@ def should_generate_summaries():
   if name_scope and "while/" in name_scope:
     # Summaries don't work well within tf.while_loop()
     return False
-  if tf.compat.v1.get_variable_scope().reuse:
+  if tf.get_variable_scope().reuse:
     # Avoid generating separate summaries for different data shards
     return False
   return True
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index d67e49650..1cc3c1b18 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -28,7 +28,7 @@
 
 import tensorflow.compat.v1 as tf
 
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class CommonLayersTest(parameterized.TestCase, tf.test.TestCase):
diff --git a/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index e4bf366be..c2c134670 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -114,4 +114,4 @@ def main(_):
 
 
 if __name__ == "__main__":
-  tf.compat.v1.app.run()
+  tf.app.run()

From b1bebfb92ce347ad28c112189156256681ff9f1e Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 22 Apr 2020 09:57:14 -0700
Subject: [PATCH 2661/2720] Minor nits in README.md

PiperOrigin-RevId: 307838203
---
 README.md           | 2 ++
 docs/walkthrough.md | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/README.md b/README.md
index 4e704b130..dcb4af46b 100644
--- a/README.md
+++ b/README.md
@@ -16,6 +16,8 @@ welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CO
 of deep learning models and datasets designed to make deep learning more
 accessible and [accelerate ML
 research](https://research.googleblog.com/2017/06/accelerating-deep-learning-research.html).
+
+
 T2T was developed by researchers and engineers in the
 [Google Brain team](https://research.google.com/teams/brain/) and a community
 of users. It is now in maintenance mode &mdash; we keep it running and welcome
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index 4e704b130..dcb4af46b 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -16,6 +16,8 @@ welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CO
 of deep learning models and datasets designed to make deep learning more
 accessible and [accelerate ML
 research](https://research.googleblog.com/2017/06/accelerating-deep-learning-research.html).
+
+
 T2T was developed by researchers and engineers in the
 [Google Brain team](https://research.google.com/teams/brain/) and a community
 of users. It is now in maintenance mode &mdash; we keep it running and welcome

From 65c2178f65661a61e697f753008bcd9b03d114b9 Mon Sep 17 00:00:00 2001
From: Gaurav Jain <gjn@google.com>
Date: Wed, 22 Apr 2020 19:50:19 -0700
Subject: [PATCH 2662/2720] Remove functional_ops based sru

PiperOrigin-RevId: 307958075
---
 tensor2tensor/layers/common_layers.py | 100 ++------------------------
 1 file changed, 6 insertions(+), 94 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index f6d680bb3..ac46036cb 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2212,12 +2212,12 @@ def gated_linear_unit_layer(x, name=None):
     return x * tf.nn.sigmoid(gating_x)
 
 
-def sru_with_scan(x,
-                  num_layers=2,
-                  activation=None,
-                  initial_state=None,
-                  name=None,
-                  reuse=None):
+def sru(x,
+        num_layers=2,
+        activation=None,
+        initial_state=None,
+        name=None,
+        reuse=None):
   """SRU cell as in https://arxiv.org/abs/1709.02755.
 
   This implementation uses tf.scan and can incur overhead, see the full SRU
@@ -2275,94 +2275,6 @@ def next_state(cur_state, args_tup):
     return tf.reshape(x, x_shape)
 
 
-class CumsumprodCell(object):
-  """Cumulative sum and product object for use with functional_rnn API."""
-
-  def __init__(self, initializer):
-    self._initializer = initializer
-
-  @property
-  def output_size(self):
-    return int(shape_list(self._initializer)[-1])
-
-  def zero_state(self, batch_size, dtype):
-    dtype = dtype or tf.float32
-    return tf.zeros([batch_size, self.output_size], dtype=dtype)
-
-  def __call__(self, inputs_t, state_t):
-    cur_x_times_one_minus_f, cur_f = tf.split(inputs_t, 2, axis=-1)
-    state_next = cur_f * state_t + cur_x_times_one_minus_f
-    outputs_t = state_next
-    return outputs_t, state_next
-
-
-def sru(x,
-        num_layers=2,
-        activation=None,
-        initial_state=None,
-        name=None,
-        reuse=None):
-  """SRU cell as in https://arxiv.org/abs/1709.02755.
-
-  As defined in the paper:
-  (1) x'_t = W x_t
-  (2) f_t = sigmoid(Wf x_t + bf)
-  (3) r_t = sigmoid(Wr x_t + br)
-  (4) c_t = f_t * c_{t-1} + (1 - f_t) * x'_t
-  (5) h_t = r_t * activation(c_t) + (1 - r_t) * x_t
-
-  This version uses functional ops to be faster on GPUs with TF-1.9+.
-
-  Args:
-    x: A tensor of shape [batch, ..., channels] ; ... is treated as time.
-    num_layers: How many SRU layers; default is 2 as results for 1 disappoint.
-    activation: Optional activation function, try tf.nn.tanh or tf.nn.relu.
-    initial_state: Optional initial c-state, set to zeros if None.
-    name: Optional name, "sru" by default.
-    reuse: Optional reuse.
-
-  Returns:
-    A tensor of the same shape as x.
-
-  Raises:
-    ValueError: if num_layers is not positive.
-  """
-  if num_layers < 1:
-    raise ValueError("Number of layers must be positive: %d" % num_layers)
-  if is_xla_compiled():  # On TPU the XLA does a good job with while.
-    return sru_with_scan(x, num_layers, activation, initial_state, name, reuse)
-  try:
-    from tensorflow.contrib.recurrent.python.ops import functional_rnn  # pylint: disable=g-import-not-at-top
-  except ImportError:
-    tf.logging.info("functional_rnn not found, using sru_with_scan instead")
-    return sru_with_scan(x, num_layers, activation, initial_state, name, reuse)
-
-  with tf.variable_scope(name, default_name="sru", values=[x], reuse=reuse):
-    # We assume x is [batch, ..., channels] and treat all ... as time.
-    x_shape = shape_list(x)
-    x = tf.reshape(x, [x_shape[0], -1, x_shape[-1]])
-    initial_state = initial_state or tf.zeros([x_shape[0], x_shape[-1]])
-    cell = CumsumprodCell(initial_state)
-    # Calculate SRU on each layer.
-    for i in range(num_layers):
-      # The parallel part of the SRU.
-      x_orig = x
-      x, f, r = tf.split(
-          layers().Dense(3 * x_shape[-1], name="kernel_%d" % i)(x), 3, axis=-1)
-      f, r = tf.sigmoid(f), tf.sigmoid(r)
-      x_times_one_minus_f = x * (1.0 - f)  # Compute in parallel for speed.
-      # Calculate states.
-      concat = tf.concat([x_times_one_minus_f, f], axis=-1)
-      c_states, _ = functional_rnn.functional_rnn(
-          cell, concat, time_major=False)
-      # Final output.
-      if activation is not None:
-        c_states = activation(c_states)
-      h = c_states * r + (1.0 - r) * x_orig
-      x = h  # Next layer.
-    return tf.reshape(x, x_shape)
-
-
 def linear_set_layer(layer_size,
                      inputs,
                      context=None,

From 022387c21e9e41fabb54175e40931302205db3a0 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sat, 25 Apr 2020 10:44:33 -0700
Subject: [PATCH 2663/2720] Automated refactoring to make code Python 3
 compatible.

PiperOrigin-RevId: 308425160
---
 tensor2tensor/data_generators/wikisum/validate_data.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/wikisum/validate_data.py b/tensor2tensor/data_generators/wikisum/validate_data.py
index 691ca73fb..6e4e8f0b7 100644
--- a/tensor2tensor/data_generators/wikisum/validate_data.py
+++ b/tensor2tensor/data_generators/wikisum/validate_data.py
@@ -23,6 +23,8 @@
 
 import numpy as np
 
+import six
+from six.moves import zip
 from tensor2tensor.data_generators.wikisum import wikisum
 
 import tensorflow.compat.v1 as tf
@@ -44,7 +46,7 @@ def aggregate_stats(stats_files):
   for fname in stats_files:
     with tf.gfile.Open(fname) as f:
       stats = json.loads(f.read())
-      for k, v in stats.iteritems():
+      for k, v in six.iteritems(stats):
         if k not in all_stats:
           if isinstance(v, list):
             all_stats[k] = []

From 72b8f7b1c47d2054ac9eef8dc83c3dc127c024b3 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 28 Apr 2020 19:21:16 -0700
Subject: [PATCH 2664/2720] Simplify implementation of
 sample_temperature_per_example and make it work with dynamic shapes.

PiperOrigin-RevId: 308944400
---
 tensor2tensor/layers/common_layers.py      | 44 +++------------------
 tensor2tensor/layers/common_layers_test.py | 46 +++++++++++++++++++++-
 2 files changed, 51 insertions(+), 39 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index ac46036cb..eeeefe584 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -2886,25 +2886,6 @@ def sample_with_temperature(logits, temperature, sampling_keep_top_k=-1):
     return choices
 
 
-def _to_nd_indices(indices):
-  """Returns indices used for tf.gather_nd or tf.scatter_nd.
-
-  Args:
-    indices: A `Tensor` of shape [batch_size, size] with integer values. The
-      values are the indices of another `Tensor`. For example, `indices` is the
-      output of tf.argsort or tf.math.top_k.
-
-  Returns:
-    A `Tensor` with shape [batch_size, size, 2] that can be used by tf.gather_nd
-    or tf.scatter_nd.
-
-  """
-  indices.get_shape().assert_has_rank(2)
-  batch_ids = tf.ones_like(indices) * tf.expand_dims(
-      tf.range(tf.shape(input=indices)[0]), 1)
-  return tf.stack([batch_ids, indices], axis=-1)
-
-
 def _select_top_k(logits, top_k):
   """Replaces logits, expect the top k highest values, with small number (-1e6).
 
@@ -2918,26 +2899,15 @@ def _select_top_k(logits, top_k):
     A `Tensor` with same shape  as logits.
   """
   vocab_size = logits.shape[-1]
-  flat_logits = tf.reshape(logits, [-1, vocab_size])
+
   top_k = tf.where(
       tf.not_equal(top_k, -1), top_k,
       tf.ones_like(top_k) * vocab_size)
-  values, idx = tf.math.top_k(flat_logits, k=vocab_size, sorted=False)
-  nd_idx = _to_nd_indices(idx)
 
-  mask_idx = tf.reshape(
-      tf.range(vocab_size), [1] * (len(logits.shape) - 1) + [-1])
-  for i, size in enumerate(logits.shape[:-1]):
-    mask_idx = tf.repeat(mask_idx, size, axis=i)
-  mask = tf.reshape(
-      mask_idx < tf.reshape(top_k, [-1] + [1] * (len(logits.shape) - 1)), [-1])
-
-  topk_logits = tf.tensor_scatter_nd_update(
-      tf.ones_like(flat_logits) * -1e6,
-      tf.reshape(nd_idx, [-1, 2])[mask],
-      tf.reshape(values, [-1])[mask])
-
-  return tf.reshape(topk_logits, logits.shape)
+  return tf.where(
+      tf.argsort(logits) < tf.reshape(top_k, [-1] + [1] *
+                                      (len(logits.shape) - 1)), logits,
+      tf.ones_like(logits) * -1e6)
 
 
 def sample_temperature_per_example(logits, temperature, sampling_keep_top_k=-1):
@@ -2950,9 +2920,7 @@ def sample_temperature_per_example(logits, temperature, sampling_keep_top_k=-1):
   Returns:
     a Tensor with one fewer dimension than logits.
   """
-  if sampling_keep_top_k != -1:
-    logits = _select_top_k(logits, sampling_keep_top_k)
-
+  logits = _select_top_k(logits, sampling_keep_top_k)
   logits /= tf.reshape(temperature, [-1] + [1] * (len(logits.shape) - 1))
   reshaped_logits = tf.reshape(logits, [-1, shape_list(logits)[-1]])
   choices = tf.multinomial(reshaped_logits, 1)
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index 1cc3c1b18..27129dddc 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -704,11 +704,55 @@ def testSampleTemperaturePerExample(self):
     logits = np.random.randn(batch_size, seq_len, 1, 1, vocab_size)
     temperature = np.random.rand(batch_size)
 
-    out = common_layers.sample_temperature_per_example(logits, temperature)
+    out = common_layers.sample_temperature_per_example(logits, temperature, -1)
 
     self.assertAllEqual(
         self.evaluate(tf.shape(out)), [batch_size, seq_len, 1, 1])
 
+  @test_utils.run_in_graph_and_eager_modes()
+  def testSampleTemperaturePerExampleWithTopK(self):
+    batch_size = 3
+    seq_len = 5
+    vocab_size = 7
+
+    logits = np.random.randn(batch_size, seq_len, 1, 1, vocab_size)
+    temperature = np.random.rand(batch_size)
+    top_k = np.array([3, -1, 4], dtype=np.int32)
+
+    out = common_layers.sample_temperature_per_example(logits, temperature,
+                                                       top_k)
+
+    self.assertAllEqual(
+        self.evaluate(tf.shape(out)), [batch_size, seq_len, 1, 1])
+
+  @test_utils.run_in_graph_and_eager_modes()
+  def testSampleTemperaturePerExampleWithTopK2(self):
+    batch_size = 3
+    vocab_size = 7
+
+    logits = np.random.randn(batch_size, vocab_size)
+    temperature = np.random.rand(batch_size)
+    top_k = np.array([3, -1, 4], dtype=np.int32)
+
+    out = common_layers.sample_temperature_per_example(logits, temperature,
+                                                       top_k)
+
+    self.assertAllEqual(self.evaluate(tf.shape(out)), [batch_size])
+
+  @test_utils.run_in_graph_mode_only()
+  def testSampleTemperaturePerExampleDynamicBatchSize(self):
+    batch_size = None
+    vocab_size = 7
+
+    logits = tf.placeholder(tf.float32, shape=(batch_size, vocab_size))
+    temperature = tf.placeholder(tf.float32, shape=(batch_size, 1))
+    sampling_keep_top_k = tf.placeholder(tf.int32, shape=(batch_size, 1))
+
+    out = common_layers.sample_temperature_per_example(logits, temperature,
+                                                       sampling_keep_top_k)
+
+    self.assertAllEqual(out.shape.as_list(), [batch_size])
+
   @test_utils.run_in_graph_and_eager_modes()
   def testCycleGANUpsampleNnUpsampleConv(self):
     batch = 8

From a9da9635917814af890a31a060c5b29d31b2f906 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Tue, 5 May 2020 11:04:12 -0700
Subject: [PATCH 2665/2720] Automated refactoring to make code Python 3
 compatible.

PiperOrigin-RevId: 309980818
---
 tensor2tensor/insights/transformer_model.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensor2tensor/insights/transformer_model.py b/tensor2tensor/insights/transformer_model.py
index 8a26ca1f0..68f86af13 100644
--- a/tensor2tensor/insights/transformer_model.py
+++ b/tensor2tensor/insights/transformer_model.py
@@ -15,6 +15,10 @@
 
 """A QueryProcessor using the Transformer framework."""
 
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
 from collections import deque
 
 import glob
@@ -24,6 +28,7 @@
 
 import numpy as np
 
+from six.moves import range
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.insights import graph

From c7629546054c2015309133602a3fead48a9f63b5 Mon Sep 17 00:00:00 2001
From: Gregory Clark <gregoryclark@google.com>
Date: Thu, 14 May 2020 15:15:31 -0700
Subject: [PATCH 2666/2720] Documentation update.

PiperOrigin-RevId: 311614763
---
 docs/new_model.md | 100 +++++++++++++++++++++++-----------------------
 1 file changed, 51 insertions(+), 49 deletions(-)

diff --git a/docs/new_model.md b/docs/new_model.md
index 861f83bb0..d0df86f74 100644
--- a/docs/new_model.md
+++ b/docs/new_model.md
@@ -15,43 +15,42 @@ Here we show how to create your own model in T2T.
 
   `T2TModel` has three typical usages:
 
-  1. Estimator: The method `make_estimator_model_fn` builds a `model_fn` for
-     the tf.Estimator workflow of training, evaluation, and prediction.
-     It performs the method `call`, which performs the core computation,
-     followed by `estimator_spec_train`, `estimator_spec_eval`, or
-     `estimator_spec_predict` depending on the tf.Estimator mode.
-  2. Layer: The method `call` enables `T2TModel` to be used a callable by
-     itself. It calls the following methods:
-
-     * `bottom`, which transforms features according to `problem_hparams`' input
-       and target `Modality`s;
-     * `body`, which takes features and performs the core model computation to
+1.  Estimator: The method `make_estimator_model_fn` builds a `model_fn` for the
+    tf.Estimator workflow of training, evaluation, and prediction. It performs
+    the method `call`, which performs the core computation, followed by
+    `estimator_spec_train`, `estimator_spec_eval`, or `estimator_spec_predict`
+    depending on the tf.Estimator mode.
+2.  Layer: The method `call` enables `T2TModel` to be used a callable by itself.
+    It calls the following methods:
+
+    *   `bottom`, which transforms features according to `problem_hparams`'
+        input and target `Modality`s;
+    *   `body`, which takes features and performs the core model computation to
         return output and any auxiliary loss terms;
-     * `top`, which takes features and the body output, and transforms them
-       according to `problem_hparams`' input and target `Modality`s to return
-       the final logits;
-     * `loss`, which takes the logits, forms any missing training loss, and sums
-       all loss terms.
-  3. Inference: The method `infer` enables `T2TModel` to make sequence
-     predictions by itself.
+    *   `top`, which takes features and the body output, and transforms them
+        according to `problem_hparams`' input and target `Modality`s to return
+        the final logits;
+    *   `loss`, which takes the logits, forms any missing training loss, and
+        sums all loss terms.
 
+3.  Inference: The method `infer` enables `T2TModel` to make sequence
+    predictions by itself.
 
 ## Creating your own model
 
-1. Create class that extends T2TModel
-    in this example it will be a copy of existing basic fully connected network:
+1.  Create a class that extends `T2TModel`. This example creates a copy of an
+    existing basic fully-connected network:
 
-```python
+    ```python
     from tensor2tensor.utils import t2t_model
 
     class MyFC(t2t_model.T2TModel):
         pass
-```
+    ```
 
+2.  Implement the `body` method:
 
-2. Implement body method:
-
-```python
+    ```python
     class MyFC(t2t_model.T2TModel):
       def body(self, features):
         hparams = self.hparams
@@ -63,43 +62,46 @@ Here we show how to create your own model in T2T.
           x = tf.nn.dropout(x, keep_prob=1.0 - hparams.dropout)
           x = tf.nn.relu(x)
         return tf.expand_dims(tf.expand_dims(x, axis=1), axis=1)  # 4D For T2T.
-```
+    ```
+
+    Method Signature:
 
+    *   Args:
 
-Method Signature:
+        *   features: dict of str to Tensor, where each Tensor has shape
+            [batch_size, ..., hidden_size]. It typically contains keys `inputs`
+            and `targets`.
 
-  * Args:
-      * features: dict of str to Tensor, where each Tensor has shape [batch_size,
-     ..., hidden_size]. It typically contains keys `inputs` and `targets`.
+    *   Returns one of:
 
-  * Returns one of:
-    * output: Tensor of pre-logit activations with shape [batch_size, ...,
-           hidden_size].
-    * losses: Either single loss as a scalar, a list, a Tensor (to be averaged),
-           or a dictionary of losses. If losses is a dictionary with the key
-           "training", losses["training"] is considered the final training
-           loss and output is considered logits; self.top and self.loss will
-           be skipped.
+        *   output: Tensor of pre-logit activations with shape [batch_size, ...,
+            hidden_size].
+        *   losses: Either single loss as a scalar, a list, a Tensor (to be
+            averaged), or a dictionary of losses. If losses is a dictionary with
+            the key "training", losses["training"] is considered the final
+            training loss and output is considered logits; self.top and
+            self.loss will be skipped.
 
-3. Register your model
+3.  Register your model:
 
-```python
+    ```python
     from tensor2tensor.utils import registry
 
     @registry.register_model
     class MyFC(t2t_model.T2TModel):
        # ...
-```
-
+    ```
 
-3. Use it with t2t tools as any other model
+4.  Use it with t2t tools as any other model:
 
-    Have in mind that names are translated from camel case to snake_case `MyFC` -> `my_fc`
-    and that you need to point t2t to directory containing your model with `t2t_usr_dir` switch. 
-    For example if you want to train model on gcloud with 1 GPU worker on IMDB sentiment task you can run your model
-    by executing following command from your model class directory. 
+    Have in mind that names are translated from camel case to snake_case `MyFC`
+    -> `my_fc` and that you need to point t2t to the directory containing your
+    model with the `--t2t_usr_dir` flag. For example if you want to train a
+    model on gcloud with 1 GPU worker on the IMDB sentiment task, you can run
+    your model by executing the following command from your model class
+    directory.
 
-```bash
+    ```bash
     t2t-trainer \
       --model=my_fc \
       --t2t_usr_dir=.
@@ -111,4 +113,4 @@ Method Signature:
       --hparams_set=basic_fc_small \
       --train_steps=10000 \
       --eval_steps=10 \
-```
+    ```

From f65b5e4e0be50b284f9b21d56d3d2a46792cdecf Mon Sep 17 00:00:00 2001
From: Russell Power <power@google.com>
Date: Wed, 20 May 2020 14:36:27 -0700
Subject: [PATCH 2667/2720] Add basic support for TF2 modeling.

This is not complete, but can be extended to add support for TPUs and more
models as required. Tested on CPU/GPU with the following configuration:

PROBLEM=translate_envi_iwslt32k
MODEL=transformer
HPARAMS=transformer_base_single_gpu

DATA_DIR=$HOME/t2t_data
TMP_DIR=/tmp/t2t_datagen
TRAIN_DIR=$HOME/t2t_train/$PROBLEM/$MODEL-$HPARAMS

t2t-trainer   --data_dir=$DATA_DIR   --problem=$PROBLEM   --model=$MODEL   --hparams_set=$HPARAMS   --output_dir=$TRAIN_DIR

Verified the loss decreases as expected and checkpoints etc work.

PiperOrigin-RevId: 312557333
---
 setup.py                           |  4 +-
 tensor2tensor/bin/t2t-datagen      |  2 +-
 tensor2tensor/bin/t2t-trainer      |  4 +-
 tensor2tensor/bin/t2t_trainer.py   |  8 +--
 tensor2tensor/models/__init__.py   | 35 ++++++-----
 tensor2tensor/utils/contrib.py     | 96 +++++++++++++++++++++++++-----
 tensor2tensor/utils/optimize.py    | 35 ++++++++---
 tensor2tensor/utils/trainer_lib.py |  8 ++-
 8 files changed, 144 insertions(+), 48 deletions(-)

diff --git a/setup.py b/setup.py
index 36e7e4a86..6fdf684ba 100644
--- a/setup.py
+++ b/setup.py
@@ -63,13 +63,15 @@
         'scipy',
         'six>=1.12.0',
         'sympy',
+        'tensorflow-addons',
         'tensorflow-datasets',
         'tensorflow-gan',
         'tensorflow-probability==0.7.0',
+        'tf_slim',
         'tqdm',
     ],
     extras_require={
-        'tensorflow': ['tensorflow>=1.15.0,<2.0'],
+        'tensorflow': ['tensorflow>=1.15.0'],
         'tensorflow-hub': ['tensorflow-hub>=0.1.1'],
         'tests': [
             # Needed to fix a Travis pytest error.
diff --git a/tensor2tensor/bin/t2t-datagen b/tensor2tensor/bin/t2t-datagen
index a49ed5615..2150027af 100755
--- a/tensor2tensor/bin/t2t-datagen
+++ b/tensor2tensor/bin/t2t-datagen
@@ -17,7 +17,7 @@ from __future__ import print_function
 
 from tensor2tensor.bin import t2t_datagen
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 def main(argv):
   t2t_datagen.main(argv)
diff --git a/tensor2tensor/bin/t2t-trainer b/tensor2tensor/bin/t2t-trainer
index c2f129409..850990bb5 100755
--- a/tensor2tensor/bin/t2t-trainer
+++ b/tensor2tensor/bin/t2t-trainer
@@ -22,7 +22,7 @@ from __future__ import print_function
 
 from tensor2tensor.bin import t2t_trainer
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 def main(argv):
   t2t_trainer.main(argv)
@@ -30,4 +30,4 @@ def main(argv):
 
 if __name__ == "__main__":
   tf.logging.set_verbosity(tf.logging.INFO)
-  tf.app.run()
+  tf.app.run(main)
diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 424e10378..7ef63c856 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -36,8 +36,6 @@
 from tensor2tensor.utils import usr_dir
 import tensorflow.compat.v1 as tf
 
-from tensorflow.contrib.tpu.python.tpu import tpu_config
-
 
 flags = tf.flags
 FLAGS = flags.FLAGS
@@ -242,8 +240,10 @@ def create_run_config(hp, output_dir=None):
     save_ckpt_steps = None  # Disable the default saver
     save_ckpt_secs = None  # Disable the default saver
     tpu_config_extra_kwargs = {
-        "num_cores_per_replica": 1,
-        "per_host_input_for_training": tpu_config.InputPipelineConfig.BROADCAST,
+        "num_cores_per_replica":
+            1,
+        "per_host_input_for_training":
+            tf.estimator.tpu.InputPipelineConfig.BROADCAST,
     }
 
   # the various custom getters we have written do not play well together yet.
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index e64b42d2d..82bf6fa96 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -30,10 +30,6 @@
 from tensor2tensor.models import image_transformer
 from tensor2tensor.models import image_transformer_2d
 from tensor2tensor.models import lstm
-from tensor2tensor.models import mtf_image_transformer
-from tensor2tensor.models import mtf_resnet
-from tensor2tensor.models import mtf_transformer
-from tensor2tensor.models import mtf_transformer2
 from tensor2tensor.models import neural_assistant
 from tensor2tensor.models import neural_gpu
 from tensor2tensor.models import resnet
@@ -47,15 +43,9 @@
 from tensor2tensor.models.neural_architecture_search import nas_model
 from tensor2tensor.models.research import adafactor_experiments
 from tensor2tensor.models.research import aligned
-from tensor2tensor.models.research import attention_lm
-from tensor2tensor.models.research import attention_lm_moe
 from tensor2tensor.models.research import autoencoders
 from tensor2tensor.models.research import cycle_gan
 from tensor2tensor.models.research import gene_expression
-from tensor2tensor.models.research import glow
-from tensor2tensor.models.research import lm_experiments
-from tensor2tensor.models.research import moe_experiments
-from tensor2tensor.models.research import multiquery_paper
 from tensor2tensor.models.research import neural_stack
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research import shuffle_network
@@ -69,19 +59,34 @@
 from tensor2tensor.models.research import transformer_symshard
 from tensor2tensor.models.research import transformer_vae
 from tensor2tensor.models.research import universal_transformer
-from tensor2tensor.models.research import vqa_attention
-from tensor2tensor.models.research import vqa_recurrent_self_attention
-from tensor2tensor.models.research import vqa_self_attention
 from tensor2tensor.models.video import basic_deterministic
 from tensor2tensor.models.video import basic_recurrent
 from tensor2tensor.models.video import basic_stochastic
 from tensor2tensor.models.video import emily
-from tensor2tensor.models.video import epva
-from tensor2tensor.models.video import next_frame_glow
 from tensor2tensor.models.video import savp
 from tensor2tensor.models.video import sv2p
+from tensor2tensor.utils import contrib
 from tensor2tensor.utils import registry
 
+# The following models can't be imported under TF2
+if not contrib.is_tf2:
+  # pylint: disable=g-import-not-at-top
+  from tensor2tensor.models.research import attention_lm
+  from tensor2tensor.models.research import attention_lm_moe
+  from tensor2tensor.models.research import glow
+  from tensor2tensor.models.research import lm_experiments
+  from tensor2tensor.models.research import moe_experiments
+  from tensor2tensor.models.research import multiquery_paper
+  from tensor2tensor.models import mtf_image_transformer
+  from tensor2tensor.models import mtf_resnet
+  from tensor2tensor.models import mtf_transformer
+  from tensor2tensor.models import mtf_transformer2
+  from tensor2tensor.models.research import vqa_attention
+  from tensor2tensor.models.research import vqa_recurrent_self_attention
+  from tensor2tensor.models.research import vqa_self_attention
+  from tensor2tensor.models.video import epva
+  from tensor2tensor.models.video import next_frame_glow
+  # pylint: enable=g-import-not-at-top
 
 # pylint: disable=unused-import
 
diff --git a/tensor2tensor/utils/contrib.py b/tensor2tensor/utils/contrib.py
index 671c005a5..b99875d7c 100644
--- a/tensor2tensor/utils/contrib.py
+++ b/tensor2tensor/utils/contrib.py
@@ -23,23 +23,40 @@
 from __future__ import print_function  # Not necessary in a Python 3-only module
 
 from absl import logging
-from tensorflow.python import tf2  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
-is_tf2 = tf2.enabled()
+import tensorflow.compat.v1 as tf
+
+# Check if we have contrib available
+try:
+  from tensorflow.contrib import slim as tf_slim  # pylint: disable=g-import-not-at-top
+  is_tf2 = False
+except:  # pylint: disable=bare-except
+  # tf.contrib, including slim and certain optimizers are not available in TF2
+  # Some features are now available in separate packages. We shim support for
+  # these as needed.
+  import tensorflow_addons as tfa  # pylint: disable=g-import-not-at-top
+  import tf_slim  # pylint: disable=g-import-not-at-top
+  is_tf2 = True
 
 
 def err_if_tf2(msg='err'):
   if is_tf2:
-    msg = 'contrib is unavailable in tf2.'
     if msg == 'err':
+      msg = 'contrib is unavailable in tf2.'
       raise ImportError(msg)
     else:
+      msg = 'contrib is unavailable in tf2.'
       logging.info(msg)
 
 
+class DummyModule(object):
+
+  def __init__(self, **kw):
+    for k, v in kw.items():
+      setattr(self, k, v)
+
+
 def slim():
-  err_if_tf2()
-  from tensorflow.contrib import slim as contrib_slim  # pylint: disable=g-import-not-at-top
-  return contrib_slim
+  return tf_slim
 
 
 def util():
@@ -54,8 +71,26 @@ def tfe():
   return contrib_eager
 
 
+def deprecated(reason, date):
+  del reason
+  del date
+  def decorator(fn):
+    return fn
+  return decorator
+
+
 def framework(msg='err'):
-  err_if_tf2(msg=msg)
+  """Return framework module or dummy version."""
+  del msg
+  if is_tf2:
+    return DummyModule(
+        arg_scope=None,
+        get_name_scope=lambda: tf.get_default_graph().get_name_scope(),
+        name_scope=tf.name_scope,
+        deprecated=deprecated,
+        nest=tf.nest,
+        argsort=tf.argsort)
+
   from tensorflow.contrib import framework as contrib_framework  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
   return contrib_framework
 
@@ -67,9 +102,13 @@ def nn():
 
 
 def layers():
-  err_if_tf2(msg='err')
-  from tensorflow.contrib import layers as contrib_layers  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
-  return contrib_layers
+  """Return layers module or dummy version."""
+  try:
+    from tensorflow.contrib import layers as contrib_layers  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+    return contrib_layers
+  except:  # pylint: disable=bare-except
+    return DummyModule(
+        OPTIMIZER_CLS_NAMES={}, optimize_loss=tf_slim.optimize_loss)
 
 
 def rnn():
@@ -109,9 +148,13 @@ def metrics():
 
 
 def opt():
-  err_if_tf2(msg='err')
-  from tensorflow.contrib import opt as contrib_opt  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
-  return contrib_opt
+  if not is_tf2:
+    from tensorflow.contrib import opt as contrib_opt  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+    return contrib_opt
+  return DummyModule(
+      LazyAdam=tfa.optimizers.LazyAdam,
+      LazyAdamOptimizer=tfa.optimizers.LazyAdam,
+  )
 
 
 def mixed_precision():
@@ -132,10 +175,31 @@ def distribute():
   return contrib_distribute
 
 
+def replace_monitors_with_hooks(monitors_or_hooks, estimator):
+  """Stub for missing function."""
+  del estimator
+  monitors_or_hooks = monitors_or_hooks or []
+  hooks = [
+      m for m in monitors_or_hooks if isinstance(m, tf.estimator.SessionRunHook)
+  ]
+  deprecated_monitors = [
+      m for m in monitors_or_hooks
+      if not isinstance(m, tf.estimator.SessionRunHook)
+  ]
+  assert not deprecated_monitors
+  return hooks
+
+
 def learn():
-  err_if_tf2(msg='err')
-  from tensorflow.contrib import learn as contrib_learn  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
-  return contrib_learn
+  """Return tf.contrib.learn module or dummy version."""
+  if not is_tf2:
+    from tensorflow.contrib import learn as contrib_learn  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+    return contrib_learn
+  return DummyModule(
+      RunConfig=tf.estimator.RunConfig,
+      monitors=DummyModule(
+          replace_monitors_with_hooks=replace_monitors_with_hooks),
+  )
 
 
 def tf_prof():
diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 0e1f79fb9..07f3ff5d8 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -112,13 +112,22 @@ def optimize(loss,
 
 @registry.register_optimizer
 def adam(learning_rate, hparams):
+  """Return adam optimizer for the given params."""
   # We change the default epsilon for Adam.
   # Using LazyAdam as it's much faster for large vocabulary embeddings.
-  return contrib.opt().LazyAdamOptimizer(
-      learning_rate,
-      beta1=hparams.optimizer_adam_beta1,
-      beta2=hparams.optimizer_adam_beta2,
-      epsilon=hparams.optimizer_adam_epsilon)
+  if contrib.is_tf2:
+    # in TF2 beta1 -> beta_1 :/
+    return contrib.opt().LazyAdamOptimizer(
+        learning_rate,
+        beta_1=hparams.optimizer_adam_beta1,
+        beta_2=hparams.optimizer_adam_beta2,
+        epsilon=hparams.optimizer_adam_epsilon)
+  else:
+    return contrib.opt().LazyAdamOptimizer(
+        learning_rate,
+        beta1=hparams.optimizer_adam_beta1,
+        beta2=hparams.optimizer_adam_beta2,
+        epsilon=hparams.optimizer_adam_epsilon)
 
 
 @registry.register_optimizer
@@ -229,7 +238,12 @@ def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disab
     self._zero_grads = hparams.optimizer_zero_grads
 
   def compute_gradients(self, loss, var_list=None, **kwargs):  # pylint: disable=arguments-differ
-    gradients = self._opt.compute_gradients(loss, var_list, **kwargs)
+    if contrib.is_tf2:
+      gradients = self._opt.get_gradients(loss, var_list)
+      gradients = zip(gradients, var_list)
+    else:
+      gradients = self._opt.compute_gradients(loss, var_list, **kwargs)
+
     def cast_grad(g, v):
       if v is not None and g is not None:
         g = common_layers.cast_like(g, v)
@@ -240,8 +254,13 @@ def cast_grad(g, v):
     return gradients
 
   def apply_gradients(self, grads_and_vars, global_step=None, name=None):
-    return self._opt.apply_gradients(
-        grads_and_vars, global_step=global_step, name=name)
+    if contrib.is_tf2:
+      with tf.control_dependencies(
+          [tf.assign_add(tf.train.get_or_create_global_step(), 1)]):
+        return self._opt.apply_gradients(grads_and_vars, name=name)
+    else:
+      return self._opt.apply_gradients(
+          grads_and_vars, global_step=global_step, name=name)
 
 
 def weight_decay_and_noise(loss, hparams, learning_rate, var_list=None):
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 8f758fa0c..ae8cae5ef 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -200,7 +200,7 @@ def create_run_config(model_name,
       "keep_checkpoint_max": keep_checkpoint_max,
       "keep_checkpoint_every_n_hours": keep_checkpoint_every_n_hours,
       "tf_random_seed": random_seed,
-      "log_step_count_steps": log_step_count_steps
+      "log_step_count_steps": log_step_count_steps,
   }
   if save_checkpoints_secs:
     del run_config_args["save_checkpoints_steps"]
@@ -239,6 +239,12 @@ def create_run_config(model_name,
     del run_config_args["master"]
     del run_config_args["evaluation_master"]
 
+  # tf.estimator RunConfig construction got totally broken in TF2.
+  # we now have to specify master in a global environment variable
+  if contrib.is_tf2:
+    del run_config_args["evaluation_master"]
+    del run_config_args["master"]
+
   config = run_config_cls(**run_config_args)
 
   # If not using TPU, add device info for data_parallelism

From c1049767fe294db19b6543a043d80cb3a7b88786 Mon Sep 17 00:00:00 2001
From: Daniel De Freitas Adiwardana <adiwardana@google.com>
Date: Thu, 28 May 2020 19:54:46 -0700
Subject: [PATCH 2668/2720] adding hparam to make encoder self-attention
 optional.

PiperOrigin-RevId: 313707269
---
 tensor2tensor/models/evolved_transformer.py | 53 +++++++++++----------
 1 file changed, 27 insertions(+), 26 deletions(-)

diff --git a/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 8c993a537..9174a9bbb 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -223,34 +223,35 @@ def evolved_transformer_encoder(encoder_input,
           hidden_state = common_layers.layer_postprocess(
               residual_state, hidden_state, hparams)
 
-        with tf.variable_scope("self_attention"):
-          residual_state = hidden_state
-          hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
+        if hparams.get("et_encoder_self_attention", True):
+          with tf.variable_scope("self_attention"):
+            residual_state = hidden_state
+            hidden_state = common_layers.layer_preprocess(hidden_state, hparams)
 
-          hidden_state = common_attention.multihead_attention(
-              hidden_state,
-              None,
-              encoder_self_attention_bias,
-              hparams.attention_key_channels or hparams.hidden_size,
-              hparams.attention_value_channels or hparams.hidden_size,
-              hparams.hidden_size,
-              hparams.num_heads,
-              hparams.attention_dropout,
-              attention_type=hparams.self_attention_type,
-              max_relative_position=hparams.max_relative_position,
-              heads_share_relative_embedding=(
-                  hparams.heads_share_relative_embedding),
-              add_relative_to_values=hparams.add_relative_to_values,
-              save_weights_to=save_weights_to,
-              make_image_summary=make_image_summary,
-              dropout_broadcast_dims=attention_dropout_broadcast_dims,
-              max_length=hparams.get("max_length"),
-              vars_3d=hparams.get("attention_variables_3d"),
-              activation_dtype=hparams.get("activation_dtype", "float32"),
-              weight_dtype=hparams.get("weight_dtype", "float32"))
+            hidden_state = common_attention.multihead_attention(
+                hidden_state,
+                None,
+                encoder_self_attention_bias,
+                hparams.attention_key_channels or hparams.hidden_size,
+                hparams.attention_value_channels or hparams.hidden_size,
+                hparams.hidden_size,
+                hparams.num_heads,
+                hparams.attention_dropout,
+                attention_type=hparams.self_attention_type,
+                max_relative_position=hparams.max_relative_position,
+                heads_share_relative_embedding=(
+                    hparams.heads_share_relative_embedding),
+                add_relative_to_values=hparams.add_relative_to_values,
+                save_weights_to=save_weights_to,
+                make_image_summary=make_image_summary,
+                dropout_broadcast_dims=attention_dropout_broadcast_dims,
+                max_length=hparams.get("max_length"),
+                vars_3d=hparams.get("attention_variables_3d"),
+                activation_dtype=hparams.get("activation_dtype", "float32"),
+                weight_dtype=hparams.get("weight_dtype", "float32"))
 
-          hidden_state = common_layers.layer_postprocess(
-              residual_state, hidden_state, hparams)
+            hidden_state = common_layers.layer_postprocess(
+                residual_state, hidden_state, hparams)
 
         with tf.variable_scope("dense_layers"):
           residual_state = hidden_state

From 2c48f8946352bcbb3e944a0dac6cdcf3a60d3368 Mon Sep 17 00:00:00 2001
From: Daniel De Freitas Adiwardana <adiwardana@google.com>
Date: Mon, 1 Jun 2020 14:52:12 -0700
Subject: [PATCH 2669/2720] Fixing feature encoder for tf.string variable
 length features.

PiperOrigin-RevId: 314208560
---
 tensor2tensor/data_generators/problem.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 92783b361..0e4a4f2be 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -718,10 +718,14 @@ def decode_example(self, serialized_example):
           [1], tf.int64, getattr(self._hparams, "sampling_keep_top_k", -1))
 
     if data_items_to_decoders is None:
-      data_items_to_decoders = {
-          field: contrib.slim().tfexample_decoder.Tensor(field)
-          for field in data_fields
-      }
+      data_items_to_decoders = {}
+      for field in data_fields:
+        if data_fields[field].dtype is tf.string:
+          default_value = b""
+        else:
+          default_value = 0
+        data_items_to_decoders[field] = contrib.slim().tfexample_decoder.Tensor(
+            field, default_value=default_value)
 
     decoder = contrib.slim().tfexample_decoder.TFExampleDecoder(
         data_fields, data_items_to_decoders)

From fee90f8eb7b0f604a59487a5fd544a6e53fed9ca Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Tue, 2 Jun 2020 07:48:19 -0700
Subject: [PATCH 2670/2720] [T2T] Bump version to 1.15.6

PiperOrigin-RevId: 314329134
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 6fdf684ba..622ee8d5f 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.15.5',
+    version='1.15.6',
     description='Tensor2Tensor',
     long_description=(
         'Tensor2Tensor, or T2T for short, is a library of '

From b66e054da3dfbb2b6262360c9712ad5c37c5e4c7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 12 Jun 2020 08:14:35 -0700
Subject: [PATCH 2671/2720] Allow action vectors and images as input for
 training Emily. When using images, action visualization is added to tf
 summary.

PiperOrigin-RevId: 316106772
---
 tensor2tensor/models/video/emily.py       | 75 +++++++++++++++++++----
 tensor2tensor/models/video/sv2p.py        | 14 ++++-
 tensor2tensor/models/video/sv2p_params.py |  1 +
 3 files changed, 75 insertions(+), 15 deletions(-)

diff --git a/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index e74e1c090..f157b15d3 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -245,7 +245,7 @@ def construct_model(self, images, actions, rewards):
 
     Args:
       images: tensor of ground truth image sequences
-      actions: NOT used list of action tensors
+      actions: list of action tensors
       rewards: NOT used list of reward tensors
 
     Returns:
@@ -256,7 +256,19 @@ def construct_model(self, images, actions, rewards):
     """
     # model does not support action conditioned and reward prediction
     fake_reward_prediction = rewards
-    del actions, rewards
+    del rewards
+    action_repeat = self.hparams.action_repeat
+    action_type = self.hparams.action_type
+
+    assert action_type in ["", "image", "vector"], "Invalid action type."
+    if not action_type:
+      a_dim = 0
+    elif action_type == "image":
+      a_dim = self.hparams.g_dim
+    else:
+      assert action_repeat > 0, "Action repeat has to be positive integer."
+      actions = tf.tile(actions, (1, 1, action_repeat))
+      a_dim = actions.shape[-1]
 
     z_dim = self.hparams.z_dim
     g_dim = self.hparams.g_dim
@@ -277,13 +289,20 @@ def construct_model(self, images, actions, rewards):
     tf.logging.info(">>>> Encoding")
     # Encoding:
     enc_images, enc_skips = [], []
+    enc_actions = []
     images = tf.unstack(images, axis=0)
+    actions = tf.unstack(actions, axis=0)
     for i, image in enumerate(images):
       with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
         enc, skips = self.encoder(image, g_dim, has_batchnorm=has_batchnorm)
         enc = tfl.flatten(enc)
         enc_images.append(enc)
         enc_skips.append(skips)
+        if action_type == "image":
+          enc_action, _ = self.encoder(
+              actions[i], g_dim, has_batchnorm=has_batchnorm)
+          enc_action = tfl.flatten(enc_action)
+          enc_actions.append(enc_action)
 
     tf.logging.info(">>>> Prediction")
     # Prediction
@@ -304,6 +323,13 @@ def construct_model(self, images, actions, rewards):
         # target encoding
         h_target = enc_images[i]
 
+        if action_type == "image":
+          h_current = tf.concat([h_current, enc_actions[i - 1]], axis=1)
+          h_target = tf.concat([h_target, enc_actions[i]], axis=1)
+        elif action_type == "vector":
+          h_current = tf.concat([h_current, actions[i - 1]], axis=1)
+          h_target = tf.concat([h_target, actions[i]], axis=1)
+
       with tf.variable_scope("prediction", reuse=tf.AUTO_REUSE):
         # Prior parameters
         if self.hparams.learned_prior:
@@ -315,7 +341,8 @@ def construct_model(self, images, actions, rewards):
           logvar_prior = tf.zeros((batch_size, z_dim))
 
         # Only use Posterior if it's training time
-        if self.is_training or len(gen_images) < context_frames:
+        if self.hparams.stochastic_model and \
+            (self.is_training or len(gen_images) < context_frames):
           mu_pos, logvar_pos, posterior_states = self.lstm_gaussian(
               h_target, posterior_states, rnn_size, z_dim, posterior_rnn_layers,
               "posterior")
@@ -338,7 +365,11 @@ def construct_model(self, images, actions, rewards):
 
       with tf.variable_scope("decoding", reuse=tf.AUTO_REUSE):
         skip_index = min(context_frames-1, i-1)
-        h_pred = tf.reshape(h_pred, [batch_size, 1, 1, g_dim])
+        if action_type == "vector":
+          h_pred = tf.concat([h_pred, actions[i - 1]], axis=-1)
+        elif action_type == "image":
+          h_pred = tf.concat([h_pred, enc_actions[i - 1]], axis=-1)
+        h_pred = tf.reshape(h_pred, [batch_size, 1, 1, g_dim + a_dim])
         if self.hparams.has_skips:
           x_pred = self.decoder(
               h_pred, color_channels,
@@ -373,22 +404,37 @@ def body(self, features):
     input_frames = common_video.swap_time_and_batch_axes(features["inputs"])
     target_frames = common_video.swap_time_and_batch_axes(features["targets"])
 
-    # Get actions if exist otherwise use zeros
-    input_actions = self.get_input_if_exists(
-        features, "input_action", batch_size, hparams.video_num_input_frames)
-    target_actions = self.get_input_if_exists(
-        features, "target_action", batch_size, hparams.video_num_target_frames)
-
     # Get rewards if exist otherwise use zeros
     input_rewards = self.get_input_if_exists(
         features, "input_reward", batch_size, hparams.video_num_input_frames)
     target_rewards = self.get_input_if_exists(
         features, "target_reward", batch_size, hparams.video_num_target_frames)
 
-    all_actions = tf.concat([input_actions, target_actions], axis=0)
     all_rewards = tf.concat([input_rewards, target_rewards], axis=0)
     all_frames = tf.concat([input_frames, target_frames], axis=0)
 
+    # Get actions if exist otherwise use zeros
+    visualization_kwargs = {}
+    if hparams.action_type == "image":
+      input_actions = common_video.swap_time_and_batch_axes(
+          features["input_action"])
+      target_actions = common_video.swap_time_and_batch_axes(
+          features["target_action"])
+      all_actions = tf.concat([input_actions, target_actions], axis=0)
+      time, _, h, w, c = all_frames.shape
+      all_actions = tf.reshape(all_actions, (time, -1, h, w, c))
+      if self.hparams.action_normalize:
+        all_actions /= 255.
+      visualization_kwargs["actions"] = all_actions[:-1]
+    else:
+      input_actions = self.get_input_if_exists(features, "input_action",
+                                               batch_size,
+                                               hparams.video_num_input_frames)
+      target_actions = self.get_input_if_exists(features, "target_action",
+                                                batch_size,
+                                                hparams.video_num_target_frames)
+      all_actions = tf.concat([input_actions, target_actions], axis=0)
+
     # Each image is being used twice, in latent tower and main tower.
     # This is to make sure we are using the *same* image for both, ...
     # ... given how TF queues work.
@@ -414,7 +460,8 @@ def body(self, features):
 
     # Visualize predictions in Tensorboard
     if self.is_training:
-      self.visualize_predictions(all_frames[1:], gen_images)
+      self.visualize_predictions(all_frames[1:], gen_images,
+                                 **visualization_kwargs)
 
     # Ignore the predictions from the input frames.
     # This is NOT the same as original paper/implementation.
@@ -473,4 +520,8 @@ def next_frame_emily():
   hparams.add_hparam("predictor_rnn_layers", 2)
   hparams.add_hparam("has_skips", True)
   hparams.add_hparam("has_batchnorm", True)
+  # Repeat actions to signify gradients.
+  # Action type can be '', 'image' or 'vector'.
+  hparams.add_hparam("action_repeat", 40)
+  hparams.add_hparam("action_type", "")
   return hparams
diff --git a/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 2a3710cd8..403d7c936 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -509,14 +509,16 @@ def save_internal_states_ops(self, internal_states):
 class NextFrameSv2pLegacy(NextFrameSv2p):
   """Old SV2P code. Only for legacy reasons."""
 
-  def visualize_predictions(self, real_frames, gen_frames):
+  def visualize_predictions(self, real_frames, gen_frames, actions=None):
+
     def concat_on_y_axis(x):
       x = tf.unstack(x, axis=1)
       x = tf.concat(x, axis=1)
       return x
-
     frames_gd = common_video.swap_time_and_batch_axes(real_frames)
     frames_pd = common_video.swap_time_and_batch_axes(gen_frames)
+    if actions is not None:
+      actions = common_video.swap_time_and_batch_axes(actions)
 
     if self.is_per_pixel_softmax:
       frames_pd_shape = common_layers.shape_list(frames_pd)
@@ -526,7 +528,13 @@ def concat_on_y_axis(x):
 
     frames_gd = concat_on_y_axis(frames_gd)
     frames_pd = concat_on_y_axis(frames_pd)
-    side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
+    if actions is not None:
+      actions = tf.clip_by_value(actions, 0, 1)
+      summary("action_vid", tf.cast(actions * 255, tf.uint8))
+      actions = concat_on_y_axis(actions)
+      side_by_side_video = tf.concat([frames_gd, frames_pd, actions], axis=2)
+    else:
+      side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
     tf.summary.image("full_video", side_by_side_video)
 
   def get_input_if_exists(self, features, key, batch_size, num_frames):
diff --git a/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 33cf08fbe..5c0405992 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -57,6 +57,7 @@ def next_frame_sv2p():
   hparams.add_hparam("upsample_method", "conv2d_transpose")
   hparams.add_hparam("reward_model", "basic")
   hparams.add_hparam("visualize_logits_histogram", True)
+  hparams.add_hparam("action_normalize", False)
   return hparams
 
 

From 4e172eea1ed4bcb044c51643a4099d8ccca8959a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Tue, 16 Jun 2020 09:43:22 -0700
Subject: [PATCH 2672/2720] Tell autograph to not convert data loaders, do not
 force gym versions to improve downstream imports.

PiperOrigin-RevId: 316699386
---
 setup.py                                 | 2 +-
 tensor2tensor/data_generators/problem.py | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 622ee8d5f..37984e8cd 100644
--- a/setup.py
+++ b/setup.py
@@ -50,7 +50,7 @@
         'gin-config',
         'google-api-python-client',
         'gunicorn',
-        'gym==0.14.0',
+        'gym',
         'h5py',
         'kfac',
         'mesh-tensorflow',
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 0e4a4f2be..85539a9fa 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -399,6 +399,7 @@ def set_task_id(self, new_task_id):
   # END SUBCLASS INTERFACE
   # ============================================================================
 
+  @tf.autograph.experimental.do_not_convert()
   def preprocess(self, dataset, mode, hparams, interleave=True):
     """Runtime preprocessing on the whole dataset.
 
@@ -587,6 +588,7 @@ def maybe_reverse_and_copy(self, example):
     self.maybe_copy_features(example)
     return example
 
+  @tf.autograph.experimental.do_not_convert()
   def dataset(self,
               mode,
               data_dir=None,

From ea1c77136dfc2f7f2af0a25e77fffac2ca4682e8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Em=C4=ABls=20Ozoli=C5=86=C5=A1?= <ozolinsemils@gmail.com>
Date: Tue, 16 Jun 2020 20:45:58 +0300
Subject: [PATCH 2673/2720] Residual Shuffle-Exchange network (#1805)

* Fix issue #1802

tensorflow._api.v1.compat.v1.compat' has no attribute 'v1.
tf already imported as tf.compat.v1 and there is
no need to use it explicitly.

* T2T implementation of Residual Shuffle-Exchange networks.

Publication: https://arxiv.org/abs/2004.04662
Original code: https://github.com/LUMII-Syslab/RSE
---
 tensor2tensor/models/__init__.py              |   1 +
 .../research/residual_shuffle_exchange.py     | 279 ++++++++++++++++++
 2 files changed, 280 insertions(+)
 create mode 100644 tensor2tensor/models/research/residual_shuffle_exchange.py

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 82bf6fa96..b6927a2d4 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -49,6 +49,7 @@
 from tensor2tensor.models.research import neural_stack
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research import shuffle_network
+from tensor2tensor.models.research import residual_shuffle_exchange
 from tensor2tensor.models.research import similarity_transformer
 from tensor2tensor.models.research import super_lm
 from tensor2tensor.models.research import transformer_moe
diff --git a/tensor2tensor/models/research/residual_shuffle_exchange.py b/tensor2tensor/models/research/residual_shuffle_exchange.py
new file mode 100644
index 000000000..00cece00e
--- /dev/null
+++ b/tensor2tensor/models/research/residual_shuffle_exchange.py
@@ -0,0 +1,279 @@
+# coding=utf-8
+# Copyright 2020 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Residual Shuffle-Exchange Network.
+
+Implementation of
+"Residual Shuffle-Exchange Networks for Fast Processing of Long Sequences"
+paper by A.Draguns, E.Ozolins, A.Sostaks, M.Apinis, K.Freivalds.
+
+Paper: https://arxiv.org/abs/2004.04662
+Original code: https://github.com/LUMII-Syslab/RSE
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.models.research.shuffle_network import ShuffleNetwork
+from tensor2tensor.models.research.shuffle_network import shuffle_layer
+from tensor2tensor.models.research.shuffle_network import reverse_shuffle_layer
+from tensor2tensor.layers.common_layers import gelu
+from tensor2tensor.utils import registry
+
+import numpy as np
+import tensorflow.compat.v1 as tf
+
+
+class LayerNormalization(tf.keras.layers.Layer):
+  """Layer Normalization (LayerNorm) without output bias and gain."""
+
+  def __init__(self, axis=1, epsilon=1e-10, **kwargs):
+    """Initialize Layer Normalization layer.
+
+    Args:
+      axis: Tuple or number of axis for calculating mean and variance
+      epsilon: Small epsilon to avoid division by zero
+    """
+    self.axis = axis
+    self.epsilon = epsilon
+    self.bias = None
+    super(LayerNormalization, self).__init__(**kwargs)
+
+  def build(self, input_shape):
+    """ Initialize bias weights for layer normalization.
+    Args:
+      input_shape: shape of input tensor
+    """
+    num_units = input_shape.as_list()[-1]
+    self.bias = self.add_weight("bias", [1, 1, num_units],
+                                initializer=tf.zeros_initializer)
+    super(LayerNormalization, self).build(input_shape)
+
+  def call(self, inputs, **kwargs):
+    """ Apply Layer Normalization without output bias and gain.
+
+    Args:
+      inputs: tensor to be normalized. Axis should be smaller than input
+      tensor dimensions.
+      **kwargs: more arguments (unused)
+    """
+    inputs -= tf.reduce_mean(inputs, axis=self.axis, keepdims=True)
+    inputs += self.bias
+    variance = tf.reduce_mean(tf.square(inputs), self.axis, keepdims=True)
+    return inputs * tf.math.rsqrt(variance + self.epsilon)
+
+
+def inv_sigmoid(y):
+  """Inverse sigmoid function.
+
+  Args:
+    y: float in range 0 to 1
+  """
+  return np.log(y / (1 - y))
+
+
+class RSU(tf.keras.layers.Layer):
+  """Residual Switch Unit of Residual Shuffle-Exchange network."""
+
+  def __init__(self, prefix, dropout, mode, **kwargs):
+    """Initialize Switch Layer.
+
+    Args:
+      prefix: Name prefix for switch layer
+      dropout: Dropout rate
+      mode: Training mode
+      **kwargs: more arguments (unused)
+    """
+    super().__init__(**kwargs)
+    self.prefix = prefix
+    self.dropout = dropout
+    self.mode = mode
+    self.first_linear = None
+    self.second_linear = None
+    self.layer_norm = None
+    self.residual_scale = None
+
+    residual_weight = 0.9
+    self.candidate_weight = np.sqrt(1 - residual_weight ** 2) * 0.25
+    self.init_value = inv_sigmoid(residual_weight)
+
+  def build(self, input_shape):
+    """Initialize layer weights and sublayers.
+
+    Args:
+      input_shape: shape of inputs
+    """
+    in_units = input_shape[-1]
+    middle_units = in_units * 4
+    out_units = in_units * 2
+    init = tf.variance_scaling_initializer(scale=1.0, mode="fan_avg",
+                                           distribution="uniform")
+
+    self.first_linear = tf.keras.layers.Dense(middle_units,
+                                              use_bias=False,
+                                              kernel_initializer=init,
+                                              name=self.prefix + "/cand1")
+
+    self.second_linear = tf.keras.layers.Dense(out_units,
+                                               kernel_initializer=init,
+                                               name=self.prefix + "/cand2")
+    self.layer_norm = LayerNormalization()
+
+    init = tf.constant_initializer(self.init_value)
+    self.residual_scale = self.add_weight(self.prefix + "/residual",
+                                          [out_units], initializer=init)
+    super(RSU, self).build(input_shape)
+
+  def call(self, inputs, **kwargs):
+    """Apply Residual Switch Layer to inputs.
+
+    Args:
+      inputs: Input tensor
+
+    Returns:
+      tf.Tensor: New candidate value
+    """
+    input_shape = tf.shape(inputs)
+    batch_size = input_shape[0]
+    length = input_shape[1]
+    num_units = inputs.shape.as_list()[2]
+
+    n_bits = tf.log(tf.cast(length - 1, tf.float32)) / tf.log(2.0)
+    n_bits = tf.floor(n_bits) + 1
+
+    reshape_shape = [batch_size, length // 2, num_units * 2]
+    reshaped_inputs = tf.reshape(inputs, reshape_shape)
+
+    first_linear = self.first_linear(reshaped_inputs)
+    first_linear = self.layer_norm(first_linear)
+    first_linear = gelu(first_linear)
+    candidate = self.second_linear(first_linear)
+
+    residual = tf.sigmoid(self.residual_scale) * reshaped_inputs
+    candidate = residual + candidate * self.candidate_weight
+    candidate = tf.reshape(candidate, input_shape)
+
+    if self.dropout > 0:
+      candidate = tf.nn.dropout(candidate, rate=self.dropout / n_bits)
+    if self.dropout != 0.0 and self.mode == tf.estimator.ModeKeys.TRAIN:
+      noise = tf.random_normal(tf.shape(candidate), mean=1.0, stddev=0.001)
+      candidate = candidate * noise
+
+    return candidate
+
+
+def residual_shuffle_network(inputs, hparams):
+  """Residual Shuffle-Exchange network with weight sharing.
+
+  Args:
+    inputs: inputs to the Shuffle-Exchange network. Should be in length of power
+      of 2.
+    hparams: Model configuration
+
+  Returns:
+    tf.Tensor: Outputs of the Shuffle-Exchange last layer
+  """
+  input_shape = tf.shape(inputs)
+  n_bits = tf.log(tf.cast(input_shape[1] - 1, tf.float32)) / tf.log(2.0)
+  n_bits = tf.cast(n_bits, tf.int32) + 1
+
+  block_out = inputs
+
+  for k in range(hparams.num_hidden_layers):
+    with tf.variable_scope("benes_block_" + str(k), reuse=tf.AUTO_REUSE):
+      forward_output = forward_part(block_out, hparams, n_bits)
+      block_out = reverse_part(forward_output, hparams, n_bits)
+
+  return RSU("last_layer", hparams.dropout, hparams.mode)(block_out)
+
+
+def reverse_part(inputs, hparams, n_bits):
+  """ Reverse part of Beneš block.
+
+  Repeatably applies interleaved Residual Switch layer and Reverse Shuffle
+  Layer. One set of weights used for all Switch layers.
+
+  Args:
+    inputs: inputs for reverse part. Should be outputs from forward part.
+    hparams: params of the network.
+    n_bits: count of repeated layer applications.
+
+  Returns:
+    tf.Tensor: output of reverse part.
+  """
+  reverse_rsu = RSU("reverse_switch", hparams.dropout, hparams.mode)
+
+  def reverse_step(state, _):
+    with tf.variable_scope("reverse"):
+      new_state = reverse_rsu(state)
+      return reverse_shuffle_layer(new_state)
+
+  reverse_outputs = tf.scan(
+    reverse_step,
+    tf.range(n_bits, n_bits * 2),
+    initializer=inputs,
+    parallel_iterations=1,
+    swap_memory=True)
+
+  return reverse_outputs[-1, :, :, :]
+
+
+def forward_part(block_out, hparams, n_bits):
+  """ Forward part of Beneš block.
+
+  Repeatably applies interleaved Residual Switch layer and Shuffle
+  Layer. One set of weights used for all Switch layers.
+
+  Args:
+    inputs: inputs for forward part. Should be inputs from previous layers
+    or Beneš block.
+    hparams: params of the network.
+    n_bits: count of repeated layer applications.
+
+  Returns:
+    tf.Tensor: output of forward part.
+  """
+  forward_rsu = RSU("switch", hparams.dropout, hparams.mode)
+
+  def forward_step(state, _):
+    with tf.variable_scope("forward"):
+      new_state = forward_rsu(state)
+      return shuffle_layer(new_state)
+
+  forward_outputs = tf.scan(
+    forward_step,
+    tf.range(0, n_bits),
+    initializer=block_out,
+    parallel_iterations=1,
+    swap_memory=True)
+
+  return forward_outputs[-1, :, :, :]
+
+
+@registry.register_model
+class ResidualShuffleExchange(ShuffleNetwork):
+  """T2T implementation of Residual Shuffle-Exchange network."""
+
+  def body(self, features):
+    """Body of Residual Shuffle-Exchange network.
+
+    Args:
+      features: dictionary of inputs and targets
+    """
+
+    inputs = tf.squeeze(features["inputs"], axis=2)
+    logits = residual_shuffle_network(inputs, self._hparams)
+    return tf.expand_dims(logits, axis=2)

From cc9d0b6f94558035eeb17d51786843eb59006404 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Em=C4=ABls=20Ozoli=C5=86=C5=A1?= <ozolinsemils@gmail.com>
Date: Tue, 16 Jun 2020 11:15:52 -0700
Subject: [PATCH 2674/2720] Merge of PR #1805

PiperOrigin-RevId: 316719305
---
 tensor2tensor/models/__init__.py              |  2 +-
 .../research/residual_shuffle_exchange.py     | 87 +++++++++++--------
 2 files changed, 50 insertions(+), 39 deletions(-)

diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index b6927a2d4..e106745a7 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -47,9 +47,9 @@
 from tensor2tensor.models.research import cycle_gan
 from tensor2tensor.models.research import gene_expression
 from tensor2tensor.models.research import neural_stack
+from tensor2tensor.models.research import residual_shuffle_exchange
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research import shuffle_network
-from tensor2tensor.models.research import residual_shuffle_exchange
 from tensor2tensor.models.research import similarity_transformer
 from tensor2tensor.models.research import super_lm
 from tensor2tensor.models.research import transformer_moe
diff --git a/tensor2tensor/models/research/residual_shuffle_exchange.py b/tensor2tensor/models/research/residual_shuffle_exchange.py
index 00cece00e..c4107d680 100644
--- a/tensor2tensor/models/research/residual_shuffle_exchange.py
+++ b/tensor2tensor/models/research/residual_shuffle_exchange.py
@@ -27,13 +27,12 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensor2tensor.models.research.shuffle_network import ShuffleNetwork
-from tensor2tensor.models.research.shuffle_network import shuffle_layer
-from tensor2tensor.models.research.shuffle_network import reverse_shuffle_layer
+import numpy as np
 from tensor2tensor.layers.common_layers import gelu
+from tensor2tensor.models.research.shuffle_network import reverse_shuffle_layer
+from tensor2tensor.models.research.shuffle_network import shuffle_layer
+from tensor2tensor.models.research.shuffle_network import ShuffleNetwork
 from tensor2tensor.utils import registry
-
-import numpy as np
 import tensorflow.compat.v1 as tf
 
 
@@ -46,6 +45,7 @@ def __init__(self, axis=1, epsilon=1e-10, **kwargs):
     Args:
       axis: Tuple or number of axis for calculating mean and variance
       epsilon: Small epsilon to avoid division by zero
+      **kwargs: keyword args passed to super.
     """
     self.axis = axis
     self.epsilon = epsilon
@@ -53,22 +53,26 @@ def __init__(self, axis=1, epsilon=1e-10, **kwargs):
     super(LayerNormalization, self).__init__(**kwargs)
 
   def build(self, input_shape):
-    """ Initialize bias weights for layer normalization.
+    """Initialize bias weights for layer normalization.
+
     Args:
       input_shape: shape of input tensor
     """
     num_units = input_shape.as_list()[-1]
-    self.bias = self.add_weight("bias", [1, 1, num_units],
-                                initializer=tf.zeros_initializer)
+    self.bias = self.add_weight(
+        "bias", [1, 1, num_units], initializer=tf.zeros_initializer)
     super(LayerNormalization, self).build(input_shape)
 
   def call(self, inputs, **kwargs):
-    """ Apply Layer Normalization without output bias and gain.
+    """Apply Layer Normalization without output bias and gain.
 
     Args:
-      inputs: tensor to be normalized. Axis should be smaller than input
-      tensor dimensions.
+      inputs: tensor to be normalized. Axis should be smaller than input tensor
+        dimensions.
       **kwargs: more arguments (unused)
+
+    Returns:
+      tensor output.
     """
     inputs -= tf.reduce_mean(inputs, axis=self.axis, keepdims=True)
     inputs += self.bias
@@ -81,6 +85,9 @@ def inv_sigmoid(y):
 
   Args:
     y: float in range 0 to 1
+
+  Returns:
+    the inverse sigmoid.
   """
   return np.log(y / (1 - y))
 
@@ -107,7 +114,7 @@ def __init__(self, prefix, dropout, mode, **kwargs):
     self.residual_scale = None
 
     residual_weight = 0.9
-    self.candidate_weight = np.sqrt(1 - residual_weight ** 2) * 0.25
+    self.candidate_weight = np.sqrt(1 - residual_weight**2) * 0.25
     self.init_value = inv_sigmoid(residual_weight)
 
   def build(self, input_shape):
@@ -119,33 +126,35 @@ def build(self, input_shape):
     in_units = input_shape[-1]
     middle_units = in_units * 4
     out_units = in_units * 2
-    init = tf.variance_scaling_initializer(scale=1.0, mode="fan_avg",
-                                           distribution="uniform")
+    init = tf.variance_scaling_initializer(
+        scale=1.0, mode="fan_avg", distribution="uniform")
 
-    self.first_linear = tf.keras.layers.Dense(middle_units,
-                                              use_bias=False,
-                                              kernel_initializer=init,
-                                              name=self.prefix + "/cand1")
+    self.first_linear = tf.keras.layers.Dense(
+        middle_units,
+        use_bias=False,
+        kernel_initializer=init,
+        name=self.prefix + "/cand1")
 
-    self.second_linear = tf.keras.layers.Dense(out_units,
-                                               kernel_initializer=init,
-                                               name=self.prefix + "/cand2")
+    self.second_linear = tf.keras.layers.Dense(
+        out_units, kernel_initializer=init, name=self.prefix + "/cand2")
     self.layer_norm = LayerNormalization()
 
     init = tf.constant_initializer(self.init_value)
-    self.residual_scale = self.add_weight(self.prefix + "/residual",
-                                          [out_units], initializer=init)
+    self.residual_scale = self.add_weight(
+        self.prefix + "/residual", [out_units], initializer=init)
     super(RSU, self).build(input_shape)
 
   def call(self, inputs, **kwargs):
     """Apply Residual Switch Layer to inputs.
 
     Args:
-      inputs: Input tensor
+      inputs: Input tensor.
+      **kwargs: unused kwargs.
 
     Returns:
       tf.Tensor: New candidate value
     """
+    del kwargs
     input_shape = tf.shape(inputs)
     batch_size = input_shape[0]
     length = input_shape[1]
@@ -201,7 +210,7 @@ def residual_shuffle_network(inputs, hparams):
 
 
 def reverse_part(inputs, hparams, n_bits):
-  """ Reverse part of Beneš block.
+  """Reverse part of Benes block.
 
   Repeatably applies interleaved Residual Switch layer and Reverse Shuffle
   Layer. One set of weights used for all Switch layers.
@@ -222,24 +231,23 @@ def reverse_step(state, _):
       return reverse_shuffle_layer(new_state)
 
   reverse_outputs = tf.scan(
-    reverse_step,
-    tf.range(n_bits, n_bits * 2),
-    initializer=inputs,
-    parallel_iterations=1,
-    swap_memory=True)
+      reverse_step,
+      tf.range(n_bits, n_bits * 2),
+      initializer=inputs,
+      parallel_iterations=1,
+      swap_memory=True)
 
   return reverse_outputs[-1, :, :, :]
 
 
 def forward_part(block_out, hparams, n_bits):
-  """ Forward part of Beneš block.
+  """Forward part of Benes block.
 
   Repeatably applies interleaved Residual Switch layer and Shuffle
   Layer. One set of weights used for all Switch layers.
 
   Args:
-    inputs: inputs for forward part. Should be inputs from previous layers
-    or Beneš block.
+    block_out: TODO(authors) document.
     hparams: params of the network.
     n_bits: count of repeated layer applications.
 
@@ -254,11 +262,11 @@ def forward_step(state, _):
       return shuffle_layer(new_state)
 
   forward_outputs = tf.scan(
-    forward_step,
-    tf.range(0, n_bits),
-    initializer=block_out,
-    parallel_iterations=1,
-    swap_memory=True)
+      forward_step,
+      tf.range(0, n_bits),
+      initializer=block_out,
+      parallel_iterations=1,
+      swap_memory=True)
 
   return forward_outputs[-1, :, :, :]
 
@@ -272,6 +280,9 @@ def body(self, features):
 
     Args:
       features: dictionary of inputs and targets
+
+    Returns:
+      the network output.
     """
 
     inputs = tf.squeeze(features["inputs"], axis=2)

From 94a3c0ee1dc9ff247d2605ef6b195bd716718b67 Mon Sep 17 00:00:00 2001
From: AgoloCuongHoang <38926056+AgoloCuongHoang@users.noreply.github.com>
Date: Wed, 17 Jun 2020 03:26:30 +0700
Subject: [PATCH 2675/2720]  update multistep_optimizer for tensorflow gpu
 (#1773)

* update multistep

* add backend

* add dtypes

* correct indent

* add cond in functions

* fix a type of function name

* rename files

* upload original files
---
 .../utils/multistep_with_adamoptimizer.py     | 229 ++++++++++++++++++
 .../multistep_with_adamoptimizer_test.py      | 108 +++++++++
 2 files changed, 337 insertions(+)
 create mode 100644 tensor2tensor/utils/multistep_with_adamoptimizer.py
 create mode 100644 tensor2tensor/utils/multistep_with_adamoptimizer_test.py

diff --git a/tensor2tensor/utils/multistep_with_adamoptimizer.py b/tensor2tensor/utils/multistep_with_adamoptimizer.py
new file mode 100644
index 000000000..1720b290e
--- /dev/null
+++ b/tensor2tensor/utils/multistep_with_adamoptimizer.py
@@ -0,0 +1,229 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Multi-step optimizers simulating large batches.
+
+Optimizer variants which make it possible to use very large batch sizes with
+limited GPU memory. Optimizers in this module accumulate the gradients for n
+batches, and call the optimizer's update rule every n batches with the
+accumulated gradients.
+
+See [Saunders et al., 2018](https://arxiv.org/abs/1805.00456) for details.
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.training import optimizer
+from tensorflow.python.training import training_ops
+from tensorflow.python.util.tf_export import tf_export
+from tensorflow.keras import backend as K
+
+
+class MultistepAdamOptimizer(optimizer.Optimizer):
+  """Adam with SGD updates every n steps with accumulated gradients."""
+
+  def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8,
+               use_locking=False, name="Adam", n=1):
+    super(MultistepAdamOptimizer, self).__init__(use_locking=use_locking, name=name)
+    self._lr = learning_rate
+    self._beta1 = beta1
+    self._beta2 = beta2
+    self._epsilon = epsilon
+    # Tensor versions of the constructor arguments, created in _prepare().
+    self._lr_t = None
+    self._beta1_t = None
+    self._beta2_t = None
+    self._epsilon_t = None
+    self._n = n  # Call Adam optimizer every n batches with accumulated grads
+    self._n_t = None  # n as tensor
+
+  def _get_beta_accumulators(self):
+    with ops.init_scope():
+      if context.executing_eagerly():
+        graph = None
+      else:
+        graph = ops.get_default_graph()
+      return (self._get_non_slot_variable("beta1_power", graph=graph),
+              self._get_non_slot_variable("beta2_power", graph=graph))
+
+  def _create_slots(self, var_list):
+    """Create slot variables for Adam with accumulated gradients."""
+    first_var = min(var_list, key=lambda x: x.name)
+    self._create_non_slot_variable(initial_value=self._beta1, name="beta1_power", colocate_with=first_var)
+    self._create_non_slot_variable(initial_value=self._beta2, name="beta2_power", colocate_with=first_var)
+    #if iter is initialized as an int32, this optimizer could not run 
+    #with tensorflow_hub with a tensorflow-gpu version
+    self._create_non_slot_variable(initial_value=0.0 if self._n == 1 else 1.0, name="iter", colocate_with=first_var)
+    # Create slots for the first and second moments, as well as grad_acc.
+    for v in var_list:
+      self._zeros_slot(v, "m", self._name)
+      self._zeros_slot(v, "v", self._name)
+      self._zeros_slot(v, "grad_acc", self._name)
+
+
+  def _get_iter_variable(self):
+    graph = (
+        None if tf.executing_eagerly() else tf.get_default_graph())
+    return self._get_non_slot_variable("iter", graph=graph)
+
+  def _prepare(self):
+    lr = self._call_if_callable(self._lr)
+    beta1 = self._call_if_callable(self._beta1)
+    beta2 = self._call_if_callable(self._beta2)
+    epsilon = self._call_if_callable(self._epsilon)
+    self._beta1_t = ops.convert_to_tensor(beta1, name="beta1")
+    self._beta2_t = ops.convert_to_tensor(beta2, name="beta2")
+    self._lr_t = ops.convert_to_tensor(lr, name="learning_rate")
+    self._epsilon_t = ops.convert_to_tensor(epsilon, name="epsilon")
+    self._n_t = tf.convert_to_tensor(self._n, name="n")
+
+  def _apply_cond(self, apply_fn, grad, var, *args, **kwargs):
+    """Apply conditionally if counter is zero."""
+    grad_acc = self.get_slot(var, "grad_acc")
+
+    def apply_adam(grad_acc, apply_fn, grad, var, *args, **kwargs):
+      total_grad = (grad_acc + grad) / tf.cast(self._n_t, grad.dtype)
+      adam_op = apply_fn(total_grad, var, *args, **kwargs)
+      with tf.control_dependencies([adam_op]):
+        grad_acc_to_zero_op = grad_acc.assign(tf.zeros_like(grad_acc),
+                                              use_locking=self._use_locking)
+      return tf.group(adam_op, grad_acc_to_zero_op)
+
+    def accumulate_gradient(grad_acc, grad):
+      assign_op = tf.assign_add(grad_acc, grad, use_locking=self._use_locking)
+      return tf.group(assign_op)  # Strip return value
+
+    return tf.cond(
+        tf.equal(self._get_iter_variable(), 0),
+        lambda: apply_adam(grad_acc, apply_fn, grad, var, *args, **kwargs),
+        lambda: accumulate_gradient(grad_acc, grad))
+
+  def _apply_dense(self, grad, var):
+    return self._apply_cond(self._apply_dense_in_action, grad, var)
+
+  def _apply_dense_in_action(self, grad, var):
+    m = self.get_slot(var, "m")
+    v = self.get_slot(var, "v")
+    beta1_power, beta2_power = self._get_beta_accumulators()
+    return training_ops.apply_adam(var, m, v, 
+        math_ops.cast(beta1_power, var.dtype.base_dtype), 
+        math_ops.cast(beta2_power, var.dtype.base_dtype), 
+        math_ops.cast(self._lr_t, var.dtype.base_dtype), 
+        math_ops.cast(self._beta1_t, var.dtype.base_dtype), 
+        math_ops.cast(self._beta2_t, var.dtype.base_dtype), 
+        math_ops.cast(self._epsilon_t, var.dtype.base_dtype), 
+        grad, 
+        use_locking=self._use_locking).op
+
+  def _resource_apply_dense(self, grad, var):
+    return self._apply_cond(self._resource_apply_dense_in_action, grad, var)
+
+  def _resource_apply_dense_in_action(self, grad, var):
+    m = self.get_slot(var, "m")
+    v = self.get_slot(var, "v")
+    beta1_power, beta2_power = self._get_beta_accumulators()
+    return training_ops.resource_apply_adam(var.handle,
+        m.handle, 
+        v.handle,
+        math_ops.cast(beta1_power, grad.dtype.base_dtype), 
+        math_ops.cast(beta2_power, grad.dtype.base_dtype), 
+        math_ops.cast(self._lr_t, var.dtype.base_dtype), 
+        math_ops.cast(self._beta1_t, grad.dtype.base_dtype), 
+        math_ops.cast(self._beta2_t, grad.dtype.base_dtype), 
+        math_ops.cast(self._epsilon_t, grad.dtype.base_dtype), 
+        grad, use_locking=self._use_locking)
+
+  def _apply_sparse_shared(self, grad, var, indices, scatter_add):
+    beta1_power, beta2_power = self._get_beta_accumulators()
+    beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype)
+    beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype)
+    lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype)
+    beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype)
+    beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype)
+    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_t = beta1 * m + (1 - beta1) * g_t
+    m = self.get_slot(var, "m")
+    m_scaled_g_values = grad * (1 - beta1_t)
+    m_t = state_ops.assign(m, m * beta1_t, use_locking=self._use_locking)
+    with ops.control_dependencies([m_t]):
+      m_t = scatter_add(m, indices, m_scaled_g_values)
+    # v_t = beta2 * v + (1 - beta2) * (g_t * g_t)
+    v = self.get_slot(var, "v")
+    v_scaled_g_values = (grad * grad) * (1 - beta2_t)
+    v_t = state_ops.assign(v, v * beta2_t, use_locking=self._use_locking)
+    with ops.control_dependencies([v_t]):
+      v_t = scatter_add(v, indices, v_scaled_g_values)
+    v_sqrt = math_ops.sqrt(v_t)
+    var_update = state_ops.assign_sub(var, lr * m_t / (v_sqrt + epsilon_t), use_locking=self._use_locking)
+    return control_flow_ops.group(*[var_update, m_t, v_t])
+
+  def _apply_sparse(self, grad, var):
+    # TODO(fstahlberg): Implement a sparse version
+    tf.logging.warning("MultistepAdamOptimizer does not support sparse updates")
+    dense_grad = tf.convert_to_tensor(grad)
+    return self._apply_cond(self._apply_dense_in_action, dense_grad, var)
+
+  def _resource_apply_sparse_duplicate_indices(self, grad, var, indices):
+    tf.logging.warning("MultistepAdamOptimizer does not support sparse updates")
+    # Note that conversion to a dense Tensor handles duplicate `indices`
+    # correctly (summing them). A real sparse implementation will probably want
+    # to override _resource_apply_sparse instead so it gets them de-duplicated
+    # automatically.
+    dense_grad = tf.convert_to_tensor(tf.IndexedSlices(values=grad, 
+        indices=indices, dense_shape=tf.shape(var)))
+    return self._apply_cond(self._resource_apply_dense_in_action, dense_grad, var)
+
+  def _resource_scatter_add(self, x, i, v):
+    with ops.control_dependencies(
+        [resource_variable_ops.resource_scatter_add(x.handle, i, v)]):
+      return x.value()
+
+  def _resource_apply_sparse(self, grad, var, indices):
+    return self._apply_sparse_shared(grad, var, indices, self._resource_scatter_add)
+
+  def _finish(self, update_ops, name_scope):
+    """Updates beta_power variables every n batches and incrs counter."""
+    iter_ = self._get_iter_variable()
+    beta1_power, beta2_power = self._get_beta_accumulators()
+    with tf.control_dependencies(update_ops):
+      with tf.colocate_with(iter_):
+        def update_beta_op():
+          update_beta1 = beta1_power.assign(
+              beta1_power * self._beta1_t,
+              use_locking=self._use_locking)
+          update_beta2 = beta2_power.assign(
+              beta2_power * self._beta2_t,
+              use_locking=self._use_locking)
+          return tf.group(update_beta1, update_beta2)
+        maybe_update_beta = tf.cond(
+            tf.equal(iter_, 0), update_beta_op, tf.no_op)
+        with tf.control_dependencies([maybe_update_beta]):
+          #TODO(Cuong): It is suboptimal here because we have to cast twice (float to int, 
+          #and then int to float)
+          update_iter = iter_.assign(K.cast(tf.mod(K.cast(iter_ + 1.0, dtype=dtypes.int32), self._n_t), dtype=dtypes.float32),
+                                     use_locking=self._use_locking)
+    return tf.group(
+        *update_ops + [update_iter, maybe_update_beta], name=name_scope)
+
diff --git a/tensor2tensor/utils/multistep_with_adamoptimizer_test.py b/tensor2tensor/utils/multistep_with_adamoptimizer_test.py
new file mode 100644
index 000000000..c11fb21b6
--- /dev/null
+++ b/tensor2tensor/utils/multistep_with_adamoptimizer_test.py
@@ -0,0 +1,108 @@
+# coding=utf-8
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Multi-step Optimizer Test Module for TensorFlow."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from tensor2tensor.utils import multistep_optimizer
+import tensorflow as tf
+
+
+class MultistepAdamOptimizerTest(tf.test.TestCase):
+
+  def testMultistep(self):
+    dtype = tf.float32
+    beta1 = 0.2
+    beta2 = 0.99
+    alpha = 10.0
+    grads0_np_lst = [
+        np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype),
+        np.array([0.2, -0.1], dtype=dtype.as_numpy_dtype),
+        np.array([0.3, 0.1], dtype=dtype.as_numpy_dtype),
+        np.array([0.4, -0.1], dtype=dtype.as_numpy_dtype)
+    ]
+    grads1_np_lst = [
+        np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype),
+        np.array([0.02, 0.02], dtype=dtype.as_numpy_dtype),
+        np.array([-0.04, 0.04], dtype=dtype.as_numpy_dtype),
+        np.array([-0.04, 0.06], dtype=dtype.as_numpy_dtype)
+    ]
+    var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+    var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+    # Test accumulating gradients for n=1..4 steps
+    for n in range(1, 5):
+      with tf.Graph().as_default():
+        with tf.Session():
+          singlestep_var0 = tf.Variable(var0_np)
+          singlestep_var1 = tf.Variable(var1_np)
+
+          multistep_var0 = tf.Variable(var0_np)
+          multistep_var1 = tf.Variable(var1_np)
+
+          singlestep_opt = tf.train.AdamOptimizer(
+              beta1=beta1, beta2=beta2, learning_rate=alpha)
+          multistep_opt = multistep_optimizer.MultistepAdamOptimizer(
+              n=n, beta1=beta1, beta2=beta2, learning_rate=alpha)
+
+          singlestep_update = singlestep_opt.apply_gradients([
+              (tf.constant(sum(grads0_np_lst[:n]) / n), singlestep_var0),
+              (tf.constant(sum(grads1_np_lst[:n]) / n), singlestep_var1)])
+          multistep_updates = [
+              multistep_opt.apply_gradients([(tf.constant(g0), multistep_var0),
+                                             (tf.constant(g1), multistep_var1)])
+              for g0, g1 in zip(grads0_np_lst, grads1_np_lst)][:n]
+
+          self.evaluate(tf.global_variables_initializer())
+          (singlestep_beta1_power,
+           singlestep_beta2_power) = singlestep_opt._get_beta_accumulators()
+          (multistep_beta1_power,
+           multistep_beta2_power) = multistep_opt._get_beta_accumulators()
+
+          # Run 3 steps of Adam
+          for _ in range(1, 4):
+            self.evaluate(singlestep_update)
+            for multistep_update in multistep_updates:
+              self.evaluate(multistep_update)
+
+            self.assertAllCloseAccordingToType(
+                self.evaluate(singlestep_beta1_power),
+                self.evaluate(multistep_beta1_power))
+            self.assertAllCloseAccordingToType(
+                self.evaluate(singlestep_beta2_power),
+                self.evaluate(multistep_beta2_power))
+            # Validate updated params
+            self.assertAllCloseAccordingToType(
+                self.evaluate(singlestep_var0),
+                self.evaluate(multistep_var0))
+            self.assertAllCloseAccordingToType(
+                self.evaluate(singlestep_var1),
+                self.evaluate(multistep_var1))
+
+  def testResourceVariables(self):
+    v1 = tf.Variable([1., 2.], use_resource=True)
+    v2 = tf.Variable([3., 4.], use_resource=True)
+    with tf.GradientTape() as tape:
+      tape.watch([v1, v2])
+      loss = tf.reduce_sum(tf.gather(params=v1, indices=[0]) + v2)
+    v1_grad, v2_grad = tape.gradient(loss, [v1, v2])
+    multistep_opt = multistep_optimizer.MultistepAdamOptimizer(0.1)
+    multistep_opt.apply_gradients(((v1_grad, v1), (v2_grad, v2)))
+
+
+if __name__ == '__main__':
+  tf.test.main()

From ceba6653844db38cba6484f2e82096440628aced Mon Sep 17 00:00:00 2001
From: AgoloCuongHoang <38926056+AgoloCuongHoang@users.noreply.github.com>
Date: Tue, 16 Jun 2020 13:26:57 -0700
Subject: [PATCH 2676/2720] Merge of PR #1773

PiperOrigin-RevId: 316746422
---
 .../utils/multistep_with_adamoptimizer.py     | 181 ++++++++++--------
 .../multistep_with_adamoptimizer_test.py      |  22 ++-
 2 files changed, 122 insertions(+), 81 deletions(-)

diff --git a/tensor2tensor/utils/multistep_with_adamoptimizer.py b/tensor2tensor/utils/multistep_with_adamoptimizer.py
index 1720b290e..956a3503d 100644
--- a/tensor2tensor/utils/multistep_with_adamoptimizer.py
+++ b/tensor2tensor/utils/multistep_with_adamoptimizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2019 The Tensor2Tensor Authors.
+# Copyright 2020 The Tensor2Tensor Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,6 +13,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Copyright 2019 The Tensor2Tensor Authors.
+#
+# 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.
 """Multi-step optimizers simulating large batches.
 
 Optimizer variants which make it possible to use very large batch sizes with
@@ -26,26 +39,26 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
-from tensorflow.python.eager import context
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import ops
-from tensorflow.python.ops import control_flow_ops
-from tensorflow.python.ops import math_ops
+import tensorflow.compat.v1 as tf
+# pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.ops import resource_variable_ops
-from tensorflow.python.ops import state_ops
-from tensorflow.python.training import optimizer
 from tensorflow.python.training import training_ops
-from tensorflow.python.util.tf_export import tf_export
-from tensorflow.keras import backend as K
+# pylint: enable=g-direct-tensorflow-import
 
 
-class MultistepAdamOptimizer(optimizer.Optimizer):
+class MultistepAdamOptimizer(tf.train.Optimizer):
   """Adam with SGD updates every n steps with accumulated gradients."""
 
-  def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8,
-               use_locking=False, name="Adam", n=1):
-    super(MultistepAdamOptimizer, self).__init__(use_locking=use_locking, name=name)
+  def __init__(self,
+               learning_rate=0.001,
+               beta1=0.9,
+               beta2=0.999,
+               epsilon=1e-8,
+               use_locking=False,
+               name="Adam",
+               n=1):
+    super(MultistepAdamOptimizer, self).__init__(
+        use_locking=use_locking, name=name)
     self._lr = learning_rate
     self._beta1 = beta1
     self._beta2 = beta2
@@ -59,32 +72,35 @@ def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8,
     self._n_t = None  # n as tensor
 
   def _get_beta_accumulators(self):
-    with ops.init_scope():
-      if context.executing_eagerly():
+    with tf.init_scope():
+      if tf.executing_eagerly():
         graph = None
       else:
-        graph = ops.get_default_graph()
+        graph = tf.get_default_graph()
       return (self._get_non_slot_variable("beta1_power", graph=graph),
               self._get_non_slot_variable("beta2_power", graph=graph))
 
   def _create_slots(self, var_list):
     """Create slot variables for Adam with accumulated gradients."""
     first_var = min(var_list, key=lambda x: x.name)
-    self._create_non_slot_variable(initial_value=self._beta1, name="beta1_power", colocate_with=first_var)
-    self._create_non_slot_variable(initial_value=self._beta2, name="beta2_power", colocate_with=first_var)
-    #if iter is initialized as an int32, this optimizer could not run 
-    #with tensorflow_hub with a tensorflow-gpu version
-    self._create_non_slot_variable(initial_value=0.0 if self._n == 1 else 1.0, name="iter", colocate_with=first_var)
+    self._create_non_slot_variable(
+        initial_value=self._beta1, name="beta1_power", colocate_with=first_var)
+    self._create_non_slot_variable(
+        initial_value=self._beta2, name="beta2_power", colocate_with=first_var)
+    # if iter is initialized as an int32, this optimizer could not run
+    # with tensorflow_hub with a tensorflow-gpu version
+    self._create_non_slot_variable(
+        initial_value=0.0 if self._n == 1 else 1.0,
+        name="iter",
+        colocate_with=first_var)
     # Create slots for the first and second moments, as well as grad_acc.
     for v in var_list:
       self._zeros_slot(v, "m", self._name)
       self._zeros_slot(v, "v", self._name)
       self._zeros_slot(v, "grad_acc", self._name)
 
-
   def _get_iter_variable(self):
-    graph = (
-        None if tf.executing_eagerly() else tf.get_default_graph())
+    graph = (None if tf.executing_eagerly() else tf.get_default_graph())
     return self._get_non_slot_variable("iter", graph=graph)
 
   def _prepare(self):
@@ -92,10 +108,10 @@ def _prepare(self):
     beta1 = self._call_if_callable(self._beta1)
     beta2 = self._call_if_callable(self._beta2)
     epsilon = self._call_if_callable(self._epsilon)
-    self._beta1_t = ops.convert_to_tensor(beta1, name="beta1")
-    self._beta2_t = ops.convert_to_tensor(beta2, name="beta2")
-    self._lr_t = ops.convert_to_tensor(lr, name="learning_rate")
-    self._epsilon_t = ops.convert_to_tensor(epsilon, name="epsilon")
+    self._beta1_t = tf.convert_to_tensor(beta1, name="beta1")
+    self._beta2_t = tf.convert_to_tensor(beta2, name="beta2")
+    self._lr_t = tf.convert_to_tensor(lr, name="learning_rate")
+    self._epsilon_t = tf.convert_to_tensor(epsilon, name="epsilon")
     self._n_t = tf.convert_to_tensor(self._n, name="n")
 
   def _apply_cond(self, apply_fn, grad, var, *args, **kwargs):
@@ -106,8 +122,8 @@ def apply_adam(grad_acc, apply_fn, grad, var, *args, **kwargs):
       total_grad = (grad_acc + grad) / tf.cast(self._n_t, grad.dtype)
       adam_op = apply_fn(total_grad, var, *args, **kwargs)
       with tf.control_dependencies([adam_op]):
-        grad_acc_to_zero_op = grad_acc.assign(tf.zeros_like(grad_acc),
-                                              use_locking=self._use_locking)
+        grad_acc_to_zero_op = grad_acc.assign(
+            tf.zeros_like(grad_acc), use_locking=self._use_locking)
       return tf.group(adam_op, grad_acc_to_zero_op)
 
     def accumulate_gradient(grad_acc, grad):
@@ -126,14 +142,17 @@ def _apply_dense_in_action(self, grad, var):
     m = self.get_slot(var, "m")
     v = self.get_slot(var, "v")
     beta1_power, beta2_power = self._get_beta_accumulators()
-    return training_ops.apply_adam(var, m, v, 
-        math_ops.cast(beta1_power, var.dtype.base_dtype), 
-        math_ops.cast(beta2_power, var.dtype.base_dtype), 
-        math_ops.cast(self._lr_t, var.dtype.base_dtype), 
-        math_ops.cast(self._beta1_t, var.dtype.base_dtype), 
-        math_ops.cast(self._beta2_t, var.dtype.base_dtype), 
-        math_ops.cast(self._epsilon_t, var.dtype.base_dtype), 
-        grad, 
+    return training_ops.apply_adam(
+        var,
+        m,
+        v,
+        tf.cast(beta1_power, var.dtype.base_dtype),
+        tf.cast(beta2_power, var.dtype.base_dtype),
+        tf.cast(self._lr_t, var.dtype.base_dtype),
+        tf.cast(self._beta1_t, var.dtype.base_dtype),
+        tf.cast(self._beta2_t, var.dtype.base_dtype),
+        tf.cast(self._epsilon_t, var.dtype.base_dtype),
+        grad,
         use_locking=self._use_locking).op
 
   def _resource_apply_dense(self, grad, var):
@@ -143,41 +162,44 @@ def _resource_apply_dense_in_action(self, grad, var):
     m = self.get_slot(var, "m")
     v = self.get_slot(var, "v")
     beta1_power, beta2_power = self._get_beta_accumulators()
-    return training_ops.resource_apply_adam(var.handle,
-        m.handle, 
+    return training_ops.resource_apply_adam(
+        var.handle,
+        m.handle,
         v.handle,
-        math_ops.cast(beta1_power, grad.dtype.base_dtype), 
-        math_ops.cast(beta2_power, grad.dtype.base_dtype), 
-        math_ops.cast(self._lr_t, var.dtype.base_dtype), 
-        math_ops.cast(self._beta1_t, grad.dtype.base_dtype), 
-        math_ops.cast(self._beta2_t, grad.dtype.base_dtype), 
-        math_ops.cast(self._epsilon_t, grad.dtype.base_dtype), 
-        grad, use_locking=self._use_locking)
+        tf.cast(beta1_power, grad.dtype.base_dtype),
+        tf.cast(beta2_power, grad.dtype.base_dtype),
+        tf.cast(self._lr_t, var.dtype.base_dtype),
+        tf.cast(self._beta1_t, grad.dtype.base_dtype),
+        tf.cast(self._beta2_t, grad.dtype.base_dtype),
+        tf.cast(self._epsilon_t, grad.dtype.base_dtype),
+        grad,
+        use_locking=self._use_locking)
 
   def _apply_sparse_shared(self, grad, var, indices, scatter_add):
     beta1_power, beta2_power = self._get_beta_accumulators()
-    beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype)
-    beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype)
-    lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype)
-    beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype)
-    beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype)
-    epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype)
-    lr = (lr_t * math_ops.sqrt(1 - beta2_power) / (1 - beta1_power))
+    beta1_power = tf.cast(beta1_power, var.dtype.base_dtype)
+    beta2_power = tf.cast(beta2_power, var.dtype.base_dtype)
+    lr_t = tf.cast(self._lr_t, var.dtype.base_dtype)
+    beta1_t = tf.cast(self._beta1_t, var.dtype.base_dtype)
+    beta2_t = tf.cast(self._beta2_t, var.dtype.base_dtype)
+    epsilon_t = tf.cast(self._epsilon_t, var.dtype.base_dtype)
+    lr = (lr_t * tf.sqrt(1 - beta2_power) / (1 - beta1_power))
     # m_t = beta1 * m + (1 - beta1) * g_t
     m = self.get_slot(var, "m")
     m_scaled_g_values = grad * (1 - beta1_t)
-    m_t = state_ops.assign(m, m * beta1_t, use_locking=self._use_locking)
-    with ops.control_dependencies([m_t]):
+    m_t = tf.assign(m, m * beta1_t, use_locking=self._use_locking)
+    with tf.control_dependencies([m_t]):
       m_t = scatter_add(m, indices, m_scaled_g_values)
     # v_t = beta2 * v + (1 - beta2) * (g_t * g_t)
     v = self.get_slot(var, "v")
     v_scaled_g_values = (grad * grad) * (1 - beta2_t)
-    v_t = state_ops.assign(v, v * beta2_t, use_locking=self._use_locking)
-    with ops.control_dependencies([v_t]):
+    v_t = tf.assign(v, v * beta2_t, use_locking=self._use_locking)
+    with tf.control_dependencies([v_t]):
       v_t = scatter_add(v, indices, v_scaled_g_values)
-    v_sqrt = math_ops.sqrt(v_t)
-    var_update = state_ops.assign_sub(var, lr * m_t / (v_sqrt + epsilon_t), use_locking=self._use_locking)
-    return control_flow_ops.group(*[var_update, m_t, v_t])
+    v_sqrt = tf.sqrt(v_t)
+    var_update = tf.assign_sub(
+        var, lr * m_t / (v_sqrt + epsilon_t), use_locking=self._use_locking)
+    return tf.group(*[var_update, m_t, v_t])
 
   def _apply_sparse(self, grad, var):
     # TODO(fstahlberg): Implement a sparse version
@@ -191,17 +213,20 @@ def _resource_apply_sparse_duplicate_indices(self, grad, var, indices):
     # correctly (summing them). A real sparse implementation will probably want
     # to override _resource_apply_sparse instead so it gets them de-duplicated
     # automatically.
-    dense_grad = tf.convert_to_tensor(tf.IndexedSlices(values=grad, 
-        indices=indices, dense_shape=tf.shape(var)))
-    return self._apply_cond(self._resource_apply_dense_in_action, dense_grad, var)
+    dense_grad = tf.convert_to_tensor(
+        tf.IndexedSlices(
+            values=grad, indices=indices, dense_shape=tf.shape(var)))
+    return self._apply_cond(self._resource_apply_dense_in_action, dense_grad,
+                            var)
 
   def _resource_scatter_add(self, x, i, v):
-    with ops.control_dependencies(
+    with tf.control_dependencies(
         [resource_variable_ops.resource_scatter_add(x.handle, i, v)]):
       return x.value()
 
   def _resource_apply_sparse(self, grad, var, indices):
-    return self._apply_sparse_shared(grad, var, indices, self._resource_scatter_add)
+    return self._apply_sparse_shared(grad, var, indices,
+                                     self._resource_scatter_add)
 
   def _finish(self, update_ops, name_scope):
     """Updates beta_power variables every n batches and incrs counter."""
@@ -209,21 +234,23 @@ def _finish(self, update_ops, name_scope):
     beta1_power, beta2_power = self._get_beta_accumulators()
     with tf.control_dependencies(update_ops):
       with tf.colocate_with(iter_):
+
         def update_beta_op():
           update_beta1 = beta1_power.assign(
-              beta1_power * self._beta1_t,
-              use_locking=self._use_locking)
+              beta1_power * self._beta1_t, use_locking=self._use_locking)
           update_beta2 = beta2_power.assign(
-              beta2_power * self._beta2_t,
-              use_locking=self._use_locking)
+              beta2_power * self._beta2_t, use_locking=self._use_locking)
           return tf.group(update_beta1, update_beta2)
+
         maybe_update_beta = tf.cond(
             tf.equal(iter_, 0), update_beta_op, tf.no_op)
         with tf.control_dependencies([maybe_update_beta]):
-          #TODO(Cuong): It is suboptimal here because we have to cast twice (float to int, 
-          #and then int to float)
-          update_iter = iter_.assign(K.cast(tf.mod(K.cast(iter_ + 1.0, dtype=dtypes.int32), self._n_t), dtype=dtypes.float32),
-                                     use_locking=self._use_locking)
+          # TODO(cuong): It is suboptimal here because we have to cast twice
+          # (float to int, and then int to float)
+          update_iter = iter_.assign(
+              tf.cast(
+                  tf.mod(tf.cast(iter_ + 1.0, tf.int32), self._n_t),
+                  tf.float32),
+              use_locking=self._use_locking)
     return tf.group(
         *update_ops + [update_iter, maybe_update_beta], name=name_scope)
-
diff --git a/tensor2tensor/utils/multistep_with_adamoptimizer_test.py b/tensor2tensor/utils/multistep_with_adamoptimizer_test.py
index c11fb21b6..a7b538cd4 100644
--- a/tensor2tensor/utils/multistep_with_adamoptimizer_test.py
+++ b/tensor2tensor/utils/multistep_with_adamoptimizer_test.py
@@ -1,4 +1,18 @@
 # coding=utf-8
+# Copyright 2020 The Tensor2Tensor Authors.
+#
+# 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.
+
 # Copyright 2019 The Tensor2Tensor Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -19,8 +33,8 @@
 from __future__ import print_function
 
 import numpy as np
-from tensor2tensor.utils import multistep_optimizer
-import tensorflow as tf
+from tensor2tensor.utils import multistep_with_adamoptimizer
+import tensorflow.compat.v1 as tf
 
 
 class MultistepAdamOptimizerTest(tf.test.TestCase):
@@ -56,7 +70,7 @@ def testMultistep(self):
 
           singlestep_opt = tf.train.AdamOptimizer(
               beta1=beta1, beta2=beta2, learning_rate=alpha)
-          multistep_opt = multistep_optimizer.MultistepAdamOptimizer(
+          multistep_opt = multistep_with_adamoptimizer.MultistepAdamOptimizer(
               n=n, beta1=beta1, beta2=beta2, learning_rate=alpha)
 
           singlestep_update = singlestep_opt.apply_gradients([
@@ -100,7 +114,7 @@ def testResourceVariables(self):
       tape.watch([v1, v2])
       loss = tf.reduce_sum(tf.gather(params=v1, indices=[0]) + v2)
     v1_grad, v2_grad = tape.gradient(loss, [v1, v2])
-    multistep_opt = multistep_optimizer.MultistepAdamOptimizer(0.1)
+    multistep_opt = multistep_with_adamoptimizer.MultistepAdamOptimizer(0.1)
     multistep_opt.apply_gradients(((v1_grad, v1), (v2_grad, v2)))
 
 

From 15ea434070b341e36b356ff7224ba7d247cac8a4 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Wed, 17 Jun 2020 01:18:34 -0700
Subject: [PATCH 2677/2720] Bump T2T to 1.15.7

PiperOrigin-RevId: 316843874
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 37984e8cd..236867727 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.15.6',
+    version='1.15.7',
     description='Tensor2Tensor',
     long_description=(
         'Tensor2Tensor, or T2T for short, is a library of '

From 05f02d8942c1f4a48ad5ee54553e446710658ae7 Mon Sep 17 00:00:00 2001
From: Aurko Roy <aurkor@google.com>
Date: Wed, 17 Jun 2020 17:54:30 -0700
Subject: [PATCH 2678/2720] Fix a bug that causes a crash when decoding for
 language models because initially targets is empty and only "infer_targets"
 is populated.

PiperOrigin-RevId: 317003338
---
 tensor2tensor/utils/t2t_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 94ae4bf3b..397ed57da 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1718,7 +1718,7 @@ def estimator_spec_predict(self, features, use_tpu=False):
 
     inputs = features.get("inputs")
     if inputs is None:
-      inputs = features["targets"]
+      inputs = features.get("targets")
 
     predictions = {
         "outputs": outputs,

From 302a04f0e47252458984a7b3fb95057792a06f9a Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser <lukaszkaiser@google.com>
Date: Sat, 27 Jun 2020 02:08:07 -0700
Subject: [PATCH 2679/2720] Use separate markers for cached weights/state to
 improve clarity, correct state corruption bug and move and add comments to a
 few other functions for clarity.

PiperOrigin-RevId: 318608989
---
 README.md           | 2 +-
 docs/walkthrough.md | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index dcb4af46b..7a3e115e2 100644
--- a/README.md
+++ b/README.md
@@ -20,7 +20,7 @@ research](https://research.googleblog.com/2017/06/accelerating-deep-learning-res
 
 T2T was developed by researchers and engineers in the
 [Google Brain team](https://research.google.com/teams/brain/) and a community
-of users. It is now in maintenance mode &mdash; we keep it running and welcome
+of users. It is now deprecated &mdash; we keep it running and welcome
 bug-fixes, but encourage users to use the successor library [Trax](https://github.com/google/trax).
 
 ### Quick Start
diff --git a/docs/walkthrough.md b/docs/walkthrough.md
index dcb4af46b..7a3e115e2 100644
--- a/docs/walkthrough.md
+++ b/docs/walkthrough.md
@@ -20,7 +20,7 @@ research](https://research.googleblog.com/2017/06/accelerating-deep-learning-res
 
 T2T was developed by researchers and engineers in the
 [Google Brain team](https://research.google.com/teams/brain/) and a community
-of users. It is now in maintenance mode &mdash; we keep it running and welcome
+of users. It is now deprecated &mdash; we keep it running and welcome
 bug-fixes, but encourage users to use the successor library [Trax](https://github.com/google/trax).
 
 ### Quick Start

From b63f2d9803385c494176f7eb5acf728c152d00f8 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 3 Jul 2020 13:09:05 -0700
Subject: [PATCH 2680/2720] Fixes typo.

PiperOrigin-RevId: 319541275
---
 tensor2tensor/layers/common_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 1713a23cf..6026fb76c 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -2451,7 +2451,7 @@ def _extract_blocks(x, block_h, block_w):
     block_h: An integer. block height
     block_w: An inteter. block width
 
-  returns:
+  Returns:
     a [batch, num_heads, height/block_h, width/block_w, depth] tensor
   """
   (_, height, width, depth) = common_layers.shape_list(x)

From 2954749ddf5be3c72b2adabb2b10228b6232672b Mon Sep 17 00:00:00 2001
From: Martin Popel <popel@ufal.mff.cuni.cz>
Date: Sat, 18 Jul 2020 05:40:50 +0200
Subject: [PATCH 2681/2720] add a backtranslation-based en-cs NMT problem and
 hparams (#1834)

* backtranslation-based (CUBBITT) en-cs NMT problem

* transformer_cubbitt hparams
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 .../data_generators/translate_encs_cubbitt.py | 111 ++++++++++++++++++
 2 files changed, 112 insertions(+)
 create mode 100644 tensor2tensor/data_generators/translate_encs_cubbitt.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 58bb75e7b..20c4d02dc 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -81,6 +81,7 @@
     "tensor2tensor.data_generators.subject_verb_agreement",
     "tensor2tensor.data_generators.timeseries",
     "tensor2tensor.data_generators.transduction_problems",
+    "tensor2tensor.data_generators.translate_encs_cubbitt",
     "tensor2tensor.data_generators.translate_encs",
     "tensor2tensor.data_generators.translate_ende",
     "tensor2tensor.data_generators.translate_enes",
diff --git a/tensor2tensor/data_generators/translate_encs_cubbitt.py b/tensor2tensor/data_generators/translate_encs_cubbitt.py
new file mode 100644
index 000000000..b59063dd7
--- /dev/null
+++ b/tensor2tensor/data_generators/translate_encs_cubbitt.py
@@ -0,0 +1,111 @@
+# coding=utf-8
+# Copyright 2020 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Data generators for English-Czech backtranslation NMT data-sets.
+
+To use this problem you need to provide backtranslated (synthetic) data to the tmp_dir
+(cs_mono_{en,cs}.txt{0,1,2} - each file of a similar size to the authentic training data).
+You can either translate the monolingual data yourself
+or you can download "csmono" data from CzEng2.0 (http://ufal.mff.cuni.cz/czeng, registration needed),
+which comes with synthetic translations into English using a backtranslation-trained model,
+thus the final model will be using "iterated" backtranslation.
+
+To get the best results out of the Block-Backtranslation
+(where blocks of synthetic and authentic training data are concatenated without shuffling),
+you should use checkpoint averaging (see t2t-avg-all).
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators import translate
+from tensor2tensor.data_generators import translate_encs
+from tensor2tensor.models.transformer import transformer_big_single_gpu
+
+from tensor2tensor.utils import registry
+
+@registry.register_problem
+class TranslateEncsCubbitt(translate_encs.TranslateEncsWmt32k):
+  """Problem spec for English-Czech CUBBITT (CUni Block-Backtranslation-Improved Transformer Translation)."""
+
+  @property
+  def use_vocab_from_other_problem(self):
+    return translate_encs.TranslateEncsWmt32k()
+
+  @property
+  def already_shuffled(self):
+    return True
+
+  @property
+  def skip_random_fraction_when_training(self):
+    return False
+
+  @property
+  def backtranslate_data_filenames(self):
+    """List of pairs of files with matched back-translated data."""
+    # Files must be placed in tmp_dir, each similar size to authentic data.
+    return [("cs_mono_en.txt%d" % i, "cs_mono_cs.txt%d" % i) for i in [0, 1, 2]]
+
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each."""
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 1,  # Use just 1 shard so as to not mix data.
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    datasets = self.source_data_files(dataset_split)
+    tag = "train" if dataset_split == problem.DatasetSplit.TRAIN else "dev"
+    data_path = translate.compile_data(
+        tmp_dir, datasets, "%s-compiled-%s" % (self.name, tag))
+    # For eval, use authentic data.
+    if dataset_split != problem.DatasetSplit.TRAIN:
+      for example in text_problems.text2text_txt_iterator(
+          data_path + ".lang1", data_path + ".lang2"):
+        yield example
+    else:  # For training, mix synthetic and authentic data as follows.
+      for (file1, file2) in self.backtranslate_data_filenames:
+        path1 = os.path.join(tmp_dir, file1)
+        path2 = os.path.join(tmp_dir, file2)
+        # Synthetic data first.
+        for example in text_problems.text2text_txt_iterator(path1, path2):
+          yield example
+        # Now authentic data.
+        for example in text_problems.text2text_txt_iterator(
+            data_path + ".lang1", data_path + ".lang2"):
+          yield example
+
+
+@registry.register_hparams
+def transformer_cubbitt():
+  """Transformer hyperparameters used in CUBBITT experiments."""
+  hparams = transformer_big_single_gpu()
+  hparams.learning_rate_schedule = "rsqrt_decay"
+  hparams.batch_size = 2900
+  hparams.learning_rate_warmup_steps = 8000
+  hparams.max_length = 150
+  hparams.layer_prepostprocess_dropout = 0
+  hparams.optimizer = "Adafactor"
+  return hparams

From 2ea8ec15eee4d04bc16fe6059cc7a38da65e64ed Mon Sep 17 00:00:00 2001
From: Henry Sudhof <hsudhof@google.com>
Date: Fri, 10 Jul 2020 15:24:34 -0700
Subject: [PATCH 2682/2720] Purely Google refactor

PiperOrigin-RevId: 320687148
---
 tensor2tensor/data_generators/all_problems.py |   1 -
 .../data_generators/translate_encs_cubbitt.py | 111 ------------------
 2 files changed, 112 deletions(-)
 delete mode 100644 tensor2tensor/data_generators/translate_encs_cubbitt.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 20c4d02dc..58bb75e7b 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -81,7 +81,6 @@
     "tensor2tensor.data_generators.subject_verb_agreement",
     "tensor2tensor.data_generators.timeseries",
     "tensor2tensor.data_generators.transduction_problems",
-    "tensor2tensor.data_generators.translate_encs_cubbitt",
     "tensor2tensor.data_generators.translate_encs",
     "tensor2tensor.data_generators.translate_ende",
     "tensor2tensor.data_generators.translate_enes",
diff --git a/tensor2tensor/data_generators/translate_encs_cubbitt.py b/tensor2tensor/data_generators/translate_encs_cubbitt.py
deleted file mode 100644
index b59063dd7..000000000
--- a/tensor2tensor/data_generators/translate_encs_cubbitt.py
+++ /dev/null
@@ -1,111 +0,0 @@
-# coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
-#
-# 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.
-
-"""Data generators for English-Czech backtranslation NMT data-sets.
-
-To use this problem you need to provide backtranslated (synthetic) data to the tmp_dir
-(cs_mono_{en,cs}.txt{0,1,2} - each file of a similar size to the authentic training data).
-You can either translate the monolingual data yourself
-or you can download "csmono" data from CzEng2.0 (http://ufal.mff.cuni.cz/czeng, registration needed),
-which comes with synthetic translations into English using a backtranslation-trained model,
-thus the final model will be using "iterated" backtranslation.
-
-To get the best results out of the Block-Backtranslation
-(where blocks of synthetic and authentic training data are concatenated without shuffling),
-you should use checkpoint averaging (see t2t-avg-all).
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-
-from tensor2tensor.data_generators import problem
-from tensor2tensor.data_generators import text_encoder
-from tensor2tensor.data_generators import text_problems
-from tensor2tensor.data_generators import translate
-from tensor2tensor.data_generators import translate_encs
-from tensor2tensor.models.transformer import transformer_big_single_gpu
-
-from tensor2tensor.utils import registry
-
-@registry.register_problem
-class TranslateEncsCubbitt(translate_encs.TranslateEncsWmt32k):
-  """Problem spec for English-Czech CUBBITT (CUni Block-Backtranslation-Improved Transformer Translation)."""
-
-  @property
-  def use_vocab_from_other_problem(self):
-    return translate_encs.TranslateEncsWmt32k()
-
-  @property
-  def already_shuffled(self):
-    return True
-
-  @property
-  def skip_random_fraction_when_training(self):
-    return False
-
-  @property
-  def backtranslate_data_filenames(self):
-    """List of pairs of files with matched back-translated data."""
-    # Files must be placed in tmp_dir, each similar size to authentic data.
-    return [("cs_mono_en.txt%d" % i, "cs_mono_cs.txt%d" % i) for i in [0, 1, 2]]
-
-  @property
-  def dataset_splits(self):
-    """Splits of data to produce and number of output shards for each."""
-    return [{
-        "split": problem.DatasetSplit.TRAIN,
-        "shards": 1,  # Use just 1 shard so as to not mix data.
-    }, {
-        "split": problem.DatasetSplit.EVAL,
-        "shards": 1,
-    }]
-
-  def generate_samples(self, data_dir, tmp_dir, dataset_split):
-    datasets = self.source_data_files(dataset_split)
-    tag = "train" if dataset_split == problem.DatasetSplit.TRAIN else "dev"
-    data_path = translate.compile_data(
-        tmp_dir, datasets, "%s-compiled-%s" % (self.name, tag))
-    # For eval, use authentic data.
-    if dataset_split != problem.DatasetSplit.TRAIN:
-      for example in text_problems.text2text_txt_iterator(
-          data_path + ".lang1", data_path + ".lang2"):
-        yield example
-    else:  # For training, mix synthetic and authentic data as follows.
-      for (file1, file2) in self.backtranslate_data_filenames:
-        path1 = os.path.join(tmp_dir, file1)
-        path2 = os.path.join(tmp_dir, file2)
-        # Synthetic data first.
-        for example in text_problems.text2text_txt_iterator(path1, path2):
-          yield example
-        # Now authentic data.
-        for example in text_problems.text2text_txt_iterator(
-            data_path + ".lang1", data_path + ".lang2"):
-          yield example
-
-
-@registry.register_hparams
-def transformer_cubbitt():
-  """Transformer hyperparameters used in CUBBITT experiments."""
-  hparams = transformer_big_single_gpu()
-  hparams.learning_rate_schedule = "rsqrt_decay"
-  hparams.batch_size = 2900
-  hparams.learning_rate_warmup_steps = 8000
-  hparams.max_length = 150
-  hparams.layer_prepostprocess_dropout = 0
-  hparams.optimizer = "Adafactor"
-  return hparams

From 95d021477272c10af15cd62f25b595ad16ad514e Mon Sep 17 00:00:00 2001
From: Martin Popel <popel@ufal.mff.cuni.cz>
Date: Fri, 17 Jul 2020 20:50:16 -0700
Subject: [PATCH 2683/2720] Merge of PR #1834

PiperOrigin-RevId: 321903600
---
 tensor2tensor/data_generators/all_problems.py |  1 +
 .../data_generators/translate_encs_cubbitt.py | 98 +++++++++++++++++++
 tensor2tensor/models/transformer.py           | 13 +++
 3 files changed, 112 insertions(+)
 create mode 100644 tensor2tensor/data_generators/translate_encs_cubbitt.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 58bb75e7b..20c4d02dc 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -81,6 +81,7 @@
     "tensor2tensor.data_generators.subject_verb_agreement",
     "tensor2tensor.data_generators.timeseries",
     "tensor2tensor.data_generators.transduction_problems",
+    "tensor2tensor.data_generators.translate_encs_cubbitt",
     "tensor2tensor.data_generators.translate_encs",
     "tensor2tensor.data_generators.translate_ende",
     "tensor2tensor.data_generators.translate_enes",
diff --git a/tensor2tensor/data_generators/translate_encs_cubbitt.py b/tensor2tensor/data_generators/translate_encs_cubbitt.py
new file mode 100644
index 000000000..1a24178ff
--- /dev/null
+++ b/tensor2tensor/data_generators/translate_encs_cubbitt.py
@@ -0,0 +1,98 @@
+# coding=utf-8
+# Copyright 2020 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Data generators for English-Czech backtranslation NMT data-sets.
+
+To use this problem you need to provide backtranslated (synthetic) data to
+tmp_dir (cs_mono_{en,cs}.txt{0,1,2} - each file of a similar size to the
+authentic training data).
+You can either translate the monolingual data yourself or you can download
+"csmono" data from CzEng2.0 (http://ufal.mff.cuni.cz/czeng, registration needed)
+which comes with synthetic translations into English using a
+backtranslation-trained model, thus the final model will be using
+"iterated" backtranslation.
+
+To get the best results out of the Block-Backtranslation
+(where blocks of synthetic and authentic training data are concatenated
+without shuffling), you should use checkpoint averaging (see t2t-avg-all).
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.data_generators import translate
+from tensor2tensor.data_generators import translate_encs
+from tensor2tensor.utils import registry
+
+
+@registry.register_problem
+class TranslateEncsCubbitt(translate_encs.TranslateEncsWmt32k):
+  """Problem spec for English-Czech CUBBITT (CUni Block-Backtranslation-Improved Transformer Translation)."""
+
+  @property
+  def use_vocab_from_other_problem(self):
+    return translate_encs.TranslateEncsWmt32k()
+
+  @property
+  def already_shuffled(self):
+    return True
+
+  @property
+  def skip_random_fraction_when_training(self):
+    return False
+
+  @property
+  def backtranslate_data_filenames(self):
+    """List of pairs of files with matched back-translated data."""
+    # Files must be placed in tmp_dir, each similar size to authentic data.
+    return [("cs_mono_en.txt%d" % i, "cs_mono_cs.txt%d" % i) for i in [0, 1, 2]]
+
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each."""
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 1,  # Use just 1 shard so as to not mix data.
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    datasets = self.source_data_files(dataset_split)
+    tag = "train" if dataset_split == problem.DatasetSplit.TRAIN else "dev"
+    data_path = translate.compile_data(
+        tmp_dir, datasets, "%s-compiled-%s" % (self.name, tag))
+    # For eval, use authentic data.
+    if dataset_split != problem.DatasetSplit.TRAIN:
+      for example in text_problems.text2text_txt_iterator(
+          data_path + ".lang1", data_path + ".lang2"):
+        yield example
+    else:  # For training, mix synthetic and authentic data as follows.
+      for (file1, file2) in self.backtranslate_data_filenames:
+        path1 = os.path.join(tmp_dir, file1)
+        path2 = os.path.join(tmp_dir, file2)
+        # Synthetic data first.
+        for example in text_problems.text2text_txt_iterator(path1, path2):
+          yield example
+        # Now authentic data.
+        for example in text_problems.text2text_txt_iterator(
+            data_path + ".lang1", data_path + ".lang2"):
+          yield example
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 5713bda28..7ce5f31e2 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -2186,6 +2186,19 @@ def transformer_base_multistep8():
   return hparams
 
 
+@registry.register_hparams
+def transformer_cubbitt():
+  """Transformer hyperparameters used in CUBBITT experiments."""
+  hparams = transformer_big_single_gpu()
+  hparams.learning_rate_schedule = "rsqrt_decay"
+  hparams.batch_size = 2900
+  hparams.learning_rate_warmup_steps = 8000
+  hparams.max_length = 150
+  hparams.layer_prepostprocess_dropout = 0
+  hparams.optimizer = "Adafactor"
+  return hparams
+
+
 @registry.register_hparams
 def transformer_parsing_base():
   """HParams for parsing on WSJ only."""

From ba8c10d770eda18594520dc91f84e54fe15a3fa6 Mon Sep 17 00:00:00 2001
From: Dmitry Tsarkov <tsar@google.com>
Date: Sat, 25 Jul 2020 14:04:01 -0700
Subject: [PATCH 2684/2720] Internal change.

PiperOrigin-RevId: 323182955
---
 tensor2tensor/data_generators/all_problems.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 20c4d02dc..697afd7bb 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -113,9 +113,11 @@
 def _is_import_err_msg(err_str, module):
   parts = module.split(".")
   suffixes = [".".join(parts[i:]) for i in range(len(parts))]
-  return err_str in (
-      ["No module named %s" % suffix for suffix in suffixes] +
-      ["No module named '%s'" % suffix for suffix in suffixes])
+  prefixes = [".".join(parts[:i]) for i in range(len(parts))]
+  return err_str in (["No module named %s" % suffix for suffix in suffixes] +
+                     ["No module named '%s'" % suffix for suffix in suffixes] +
+                     ["No module named %s" % prefix for prefix in prefixes] +
+                     ["No module named '%s'" % prefix for prefix in prefixes])
 
 
 def _handle_errors(errors):

From 5f9dd2db6d7797162e53adf152310ed13e9fc711 Mon Sep 17 00:00:00 2001
From: Trevor Gale <tgale@google.com>
Date: Tue, 29 Sep 2020 06:26:33 -0700
Subject: [PATCH 2685/2720] Removing layers cache to avoid issues with TF2
 check on initialization.

PiperOrigin-RevId: 334365253
---
 tensor2tensor/layers/common_layers.py | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index eeeefe584..722c29a0a 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -37,15 +37,9 @@
 from tensorflow.python.ops import inplace_ops
 
 
-_cached_layers = None
-
-
 # TODO(lukaszkaiser): remove this function when not needed any more.
 def layers():
   """Get the layers module good for TF 1 and TF 2 work for now."""
-  global _cached_layers
-  if _cached_layers is not None:
-    return _cached_layers
   layers_module = None
   try:
     layers_module = tf.layers
@@ -58,7 +52,6 @@ def layers():
       layers_module = tf.keras.layers
   except ImportError:
     pass
-  _cached_layers = layers_module
   return layers_module
 
 

From 21dba2c1bdcc7ab582a2bfd8c0885c217963bb4f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 16 Nov 2020 06:27:41 -0800
Subject: [PATCH 2686/2720] Add Seq2Edits (go/seq2edits-paper) to T2T.

PiperOrigin-RevId: 342622759
---
 tensor2tensor/data_generators/all_problems.py |   1 +
 tensor2tensor/data_generators/seq2edits.py    | 266 +++++++++
 tensor2tensor/models/__init__.py              |   1 +
 .../models/research/transformer_seq2edits.py  | 543 ++++++++++++++++++
 4 files changed, 811 insertions(+)
 create mode 100644 tensor2tensor/data_generators/seq2edits.py
 create mode 100644 tensor2tensor/models/research/transformer_seq2edits.py

diff --git a/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 697afd7bb..50bf98417 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -73,6 +73,7 @@
     "tensor2tensor.data_generators.quora_qpairs",
     "tensor2tensor.data_generators.rte",
     "tensor2tensor.data_generators.scitail",
+    "tensor2tensor.data_generators.seq2edits",
     "tensor2tensor.data_generators.snli",
     "tensor2tensor.data_generators.stanford_nli",
     "tensor2tensor.data_generators.style_transfer",
diff --git a/tensor2tensor/data_generators/seq2edits.py b/tensor2tensor/data_generators/seq2edits.py
new file mode 100644
index 000000000..35b08ec08
--- /dev/null
+++ b/tensor2tensor/data_generators/seq2edits.py
@@ -0,0 +1,266 @@
+# coding=utf-8
+# Copyright 2020 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""Problems for Seq2Edits (see models/research/transformer_seq2edits.py)."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from tensor2tensor.data_generators import text_encoder
+from tensor2tensor.data_generators import text_problems
+from tensor2tensor.layers import modalities
+from tensor2tensor.utils import registry
+
+import tensorflow.compat.v1 as tf
+
+
+@modalities.is_pointwise
+def pointer_top(body_output, targets, model_hparams, vocab_size):
+  """Like identity_top() with is_pointwise annotation."""
+  del targets, model_hparams, vocab_size  # unused arg
+  return body_output
+
+
+def pointer_bottom(x, model_hparams, vocab_size):
+  """Like identity_bottom() without converting to float."""
+  del model_hparams, vocab_size  # unused arg
+  return x
+
+
+@registry.register_problem
+class Seq2editsGec(text_problems.Text2TextProblem):
+  """Seq2Edits for grammatical error correction."""
+
+  def dataset_filename(self):
+    return "edit_ops_gec"
+
+  @property
+  def vocab_file(self):
+    return "vocab.subwords"
+
+  @property
+  def vocab_filename(self):
+    return "vocab.subwords"
+
+  @property
+  def error_tag_vocab_file(self):
+    return "vocab.error_tags"
+
+  def feature_encoders(self, data_dir):
+    subword_encoder = text_encoder.SubwordTextEncoder(
+        os.path.join(data_dir, self.vocab_file))
+    error_tag_encoder = text_encoder.TokenTextEncoder(
+        os.path.join(data_dir, self.error_tag_vocab_file))
+    return {
+        "inputs": subword_encoder,
+        "targets": subword_encoder,
+        "targets_error_tag": error_tag_encoder
+    }
+
+  def hparams(self, defaults, model_hparams):
+    super(Seq2editsGec, self).hparams(defaults, model_hparams)
+
+    for pointer_feat in ["targets_start_token", "targets_end_token"]:
+      defaults.modality[pointer_feat] = modalities.ModalityType.IDENTITY
+      defaults.vocab_size[pointer_feat] = None
+      model_hparams.bottom[pointer_feat] = pointer_bottom
+      model_hparams.top[pointer_feat] = pointer_top
+    # Whether to use tags.
+    if "use_error_tags" not in model_hparams:
+      model_hparams.add_hparam("use_error_tags", True)
+    # If true, span and tag prediction is in the middle of the decoder layer
+    # stack. Otherwise, they are at the end of the decoder layer stack.
+    if "middle_prediction" not in model_hparams:
+      model_hparams.add_hparam("middle_prediction", True)
+    # If middle_prediction=True, divide num_decoder_layers by this to get the
+    # number of layers before and after the middle prediction.
+    if "middle_prediction_layer_factor" not in model_hparams:
+      model_hparams.add_hparam("middle_prediction_layer_factor", 2)
+    # Whether to predict the targets_start_token feature. If this is false, use
+    # the previous end token as implicit start token.
+    if "use_start_token" not in model_hparams:
+      model_hparams.add_hparam("use_start_token", False)
+    # Whether to feed back targets_end_token to the next time step. If false,
+    # only feed back targets_start_token.
+    if "feedback_end_token" not in model_hparams:
+      model_hparams.add_hparam("feedback_end_token", False)
+    # Number of feedforward layers between prediction layers in the cascade.
+    if "ffn_in_prediction_cascade" not in model_hparams:
+      model_hparams.add_hparam("ffn_in_prediction_cascade", 1)
+    # Embedding size for error tags.
+    if "error_tag_embed_size" not in model_hparams:
+      model_hparams.add_hparam("error_tag_embed_size", 6)
+    if model_hparams.use_error_tags:
+      defaults.modality["targets_error_tag"] = modalities.ModalityType.SYMBOL
+      error_tag_vocab_size = self._encoders["targets_error_tag"].vocab_size
+      defaults.vocab_size["targets_error_tag"] = error_tag_vocab_size
+      model_hparams.top["targets_error_tag"] = pointer_top
+
+  def example_reading_spec(self):
+    data_fields, _ = super(Seq2editsGec, self).example_reading_spec()
+    data_fields["targets_start_token"] = tf.VarLenFeature(tf.int64)
+    data_fields["targets_end_token"] = tf.VarLenFeature(tf.int64)
+    data_fields["targets_error_tag"] = tf.VarLenFeature(tf.int64)
+    return data_fields, None
+
+
+@registry.register_problem
+class Seq2editsGecPacked256(Seq2editsGec):
+  """Packed version for TPU."""
+
+  def dataset_filename(self):
+    return "edit_ops_gec_packed256"
+
+  @property
+  def packed_length(self):
+    return 256
+
+  @property
+  def max_segment_length(self):
+    return 256
+
+
+@registry.register_problem
+class Seq2editsGecNoTags(Seq2editsGec):
+  """Seq2Edits for grammatical error correction without tags."""
+
+  def dataset_filename(self):
+    return "edit_ops_gec"
+
+  def hparams(self, defaults, model_hparams):
+    super(Seq2editsGecNoTags, self).hparams(defaults, model_hparams)
+    model_hparams.use_error_tags = False
+
+
+@registry.register_problem
+class Seq2editsGecNoTagsPacked256(Seq2editsGecPacked256):
+  """Packed version for TPU."""
+
+  def dataset_filename(self):
+    return "edit_ops_gec_packed256"
+
+  def hparams(self, defaults, model_hparams):
+    super(Seq2editsGecNoTagsPacked256, self).hparams(defaults, model_hparams)
+    model_hparams.use_error_tags = False
+
+
+@registry.register_problem
+class Seq2editsGecDeep(Seq2editsGec):
+  """Seq2Edits for grammatical error correction with deeper decoder."""
+
+  def hparams(self, defaults, model_hparams):
+    super(Seq2editsGecDeep, self).hparams(defaults, model_hparams)
+    model_hparams.middle_prediction_layer_factor = 1.5
+
+
+@registry.register_problem
+class Seq2editsGecDeepPacked256(Seq2editsGecPacked256):
+  """Packed version for TPU."""
+
+  def hparams(self, defaults, model_hparams):
+    super(Seq2editsGecDeepPacked256, self).hparams(defaults, model_hparams)
+    model_hparams.middle_prediction_layer_factor = 1.5
+
+
+@registry.register_problem
+class Seq2editsGecDeepNoTags(Seq2editsGec):
+  """Deep Seq2Edits model for grammatical error correction without tags."""
+
+  def hparams(self, defaults, model_hparams):
+    super(Seq2editsGecDeepNoTags, self).hparams(defaults, model_hparams)
+    model_hparams.middle_prediction_layer_factor = 1.5
+    model_hparams.use_error_tags = False
+
+
+@registry.register_problem
+class Seq2editsGecDeepNoTagsPacked256(Seq2editsGecPacked256):
+  """Packed version for TPU."""
+
+  def hparams(self, defaults, model_hparams):
+    super(Seq2editsGecDeepNoTagsPacked256, self).hparams(
+        defaults, model_hparams)
+    model_hparams.middle_prediction_layer_factor = 1.5
+    model_hparams.use_error_tags = False
+
+
+@registry.register_problem
+class Seq2editsTextnorm(Seq2editsGec):
+  """Seq2Edits for text normalization."""
+
+  def dataset_filename(self):
+    return "edit_ops_textnorm"
+
+  @property
+  def source_vocab_file(self):
+    return "vocab.source"
+
+  @property
+  def target_vocab_file(self):
+    return "vocab.target"
+
+  @property
+  def error_tag_vocab_file(self):
+    return "vocab.error_tags"
+
+  def feature_encoders(self, data_dir):
+    source_encoder = text_encoder.TokenTextEncoder(
+        os.path.join(data_dir, self.source_vocab_file))
+    target_encoder = text_encoder.TokenTextEncoder(
+        os.path.join(data_dir, self.target_vocab_file))
+    error_tag_encoder = text_encoder.TokenTextEncoder(
+        os.path.join(data_dir, self.error_tag_vocab_file))
+    return {
+        "inputs": source_encoder,
+        "targets": target_encoder,
+        "targets_error_tag": error_tag_encoder
+    }
+
+
+@registry.register_problem
+class Seq2editsTextnormPacked256(Seq2editsTextnorm):
+  """Packed version for TPU."""
+
+  def dataset_filename(self):
+    return "edit_ops_textnorm_packed256"
+
+  @property
+  def packed_length(self):
+    return 256
+
+  @property
+  def max_segment_length(self):
+    return 256
+
+
+@registry.register_problem
+class Seq2editsTextnormNoTags(Seq2editsTextnorm):
+  """Seq2Edits for text normalization without tags."""
+
+  def hparams(self, defaults, model_hparams):
+    super(Seq2editsTextnormNoTags, self).hparams(defaults, model_hparams)
+    model_hparams.use_error_tags = False
+
+
+@registry.register_problem
+class Seq2editsTextnormNoTagsPacked256(Seq2editsTextnormPacked256):
+  """Packed version for TPU."""
+
+  def hparams(self, defaults, model_hparams):
+    super(Seq2editsTextnormNoTagsPacked256, self).hparams(
+        defaults, model_hparams)
+    model_hparams.use_error_tags = False
diff --git a/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index e106745a7..c6c692bd8 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -56,6 +56,7 @@
 from tensor2tensor.models.research import transformer_nat
 from tensor2tensor.models.research import transformer_parallel
 from tensor2tensor.models.research import transformer_revnet
+from tensor2tensor.models.research import transformer_seq2edits
 from tensor2tensor.models.research import transformer_sketch
 from tensor2tensor.models.research import transformer_symshard
 from tensor2tensor.models.research import transformer_vae
diff --git a/tensor2tensor/models/research/transformer_seq2edits.py b/tensor2tensor/models/research/transformer_seq2edits.py
new file mode 100644
index 000000000..4276b7838
--- /dev/null
+++ b/tensor2tensor/models/research/transformer_seq2edits.py
@@ -0,0 +1,543 @@
+# coding=utf-8
+# Copyright 2020 The Tensor2Tensor Authors.
+#
+# 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.
+
+"""The Seq2Edits model.
+
+Seq2Edits is an adaptation of the Transformer that predicts span level edits
+and pairs them with tags. The Seq2Edits model is described in
+
+  Stahlberg, Felix, and Kumar, Shankar. "Seq2Edits: Sequence Transduction Using
+  Span-level Edit Operations." Proceedings of the 2020 Conference on Empirical
+  Methods in Natural Language Processing (EMNLP). 2020.
+  https://www.aclweb.org/anthology/2020.emnlp-main.418/
+
+T2T problem definitions for Seq2Edits are in data_generators/seq2edits.py.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+
+from tensor2tensor.layers import common_attention
+from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import transformer_layers
+from tensor2tensor.models import transformer
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
+
+import tensorflow.compat.v1 as tf
+
+
+def maybe_flatten4d3d(x):
+  """Flatten if tensor has 4 dimensions.
+
+  Pass through otherwise.
+
+  This is useful since additional dimensions are sometimes removed on the TPU,
+  see e.g.
+    https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/models/transformer.py?l=159&rcl=279807999
+
+  Args:
+    x: a tensor
+
+  Returns:
+    A 3D tensor if x is 4D, unmodified x otherwise.
+  """
+  xshape = common_layers.shape_list(x)
+  return common_layers.flatten4d3d(x) if len(xshape) == 4 else x
+
+
+def maybe_flatten3d2d(x):
+  """Flatten if tensor has 3 dimensions, similar to maybe_flatten4d3d()."""
+  xshape = common_layers.shape_list(x)
+  if len(xshape) != 3:
+    return x
+  return tf.reshape(x, [xshape[0], xshape[1] * xshape[2]])
+
+
+def maybe_flatten4d2d(x):
+  return maybe_flatten3d2d(maybe_flatten4d3d(x))
+
+
+def features_to_nonpadding(features, inputs_or_targets="inputs"):
+  """See transformer.features_to_nonpadding."""
+  key = inputs_or_targets + "_segmentation"
+  if features and key in features:
+    return tf.minimum(tf.to_float(features[key]), 1.0)
+  return None
+
+
+def gather_2d(params, indices):
+  """2D version of tf.gather.
+
+  This is a batched version of tf.gather(), i.e. it applies tf.gather() to
+  each batch separately.
+  Example:
+    params = [[10, 11, 12, 13, 14],
+              [20, 21, 22, 23, 24]]
+    indices = [[0, 0, 1, 1, 1, 2],
+               [1, 3, 0, 0, 2, 2]]
+    result = [[10, 10, 11, 11, 11, 12],
+              [21, 23, 20, 20, 22, 22]]
+  This method is copied from
+    https://github.com/fstahlberg/tensor2tensor-usr/blob/master/usr/utils.py
+  which is published under Apache 2.
+
+  Args:
+    params: A [batch_size, n, ...] tensor with data
+    indices: A [batch_size, num_indices] int32 tensor with indices into params.
+      Entries must be smaller than n
+
+  Returns:
+    The result of tf.gather() on each entry of the batch.
+  """
+  batch_size = tf.shape(params)[0]
+  num_indices = tf.shape(indices)[1]
+  batch_indices = tf.tile(
+      tf.expand_dims(tf.range(batch_size), 1), [1, num_indices])
+  # batch_indices is [[0,0,0,0,...],[1,1,1,1,...],...]
+  gather_nd_indices = tf.stack([batch_indices, indices], axis=2)
+  return tf.gather_nd(params, gather_nd_indices)
+
+
+@registry.register_model
+class TransformerSeq2edits(t2t_model.T2TModel):
+  """The Seq2Edits model. See file docstring."""
+
+  def __init__(self, *args, **kwargs):
+    super(TransformerSeq2edits, self).__init__(*args, **kwargs)
+    self.attention_weights = {}  # For visualizing attention heads.
+    self._encoder_function = transformer_layers.transformer_encoder
+    self._decoder_function = transformer.transformer_decoder
+    self._prepare_encoder_fn = transformer_layers.transformer_prepare_encoder
+    self._prepare_decoder_fn = transformer.transformer_prepare_decoder
+    self.loss_num = {}
+    self.logits = {}
+    self.loss_den = None
+
+  def encode(self, inputs, target_space, hparams, features=None, losses=None):
+    """Encodes transformer inputs, see transformer.transformer_encode()."""
+    return transformer.transformer_encode(
+        self._encoder_function,
+        inputs,
+        target_space,
+        hparams,
+        attention_weights=self.attention_weights,
+        features=features,
+        losses=losses,
+        prepare_encoder_fn=self._prepare_encoder_fn)
+
+  def decode(self,
+             decoder_input,
+             encoder_output,
+             encoder_decoder_attention_bias,
+             decoder_self_attention_bias,
+             hparams,
+             cache=None,
+             decode_loop_step=None,
+             nonpadding=None,
+             losses=None,
+             **kwargs):
+    """Decodes Transformer outputs, see transformer.transformer_decode()."""
+    return transformer.transformer_decode(
+        self._decoder_function,
+        decoder_input,
+        encoder_output,
+        encoder_decoder_attention_bias,
+        decoder_self_attention_bias,
+        hparams,
+        attention_weights=self.attention_weights,
+        cache=cache,
+        decode_loop_step=decode_loop_step,
+        nonpadding=nonpadding,
+        losses=losses,
+        **kwargs)
+
+  def body(self, features):
+    """Seq2Edits main model_fn.
+
+    Args:
+      features: Feature dictionary. Should contain the following fields:
+          "inputs": [batch_size, input_length, 1, hidden_dim] float tensor with
+            input token embeddings.
+          "targets": [batch_size, target_length, 1, hidden_dim] float tensor
+            with target token embeddings.
+          "targets_error_tag": [batch_size, target_length, 1, hidden_dim] float
+            tensor with target error tag embeddings.
+          "targets_start_token": [batch_size, target_length] int tensor with
+            start token positions.
+          "targets_end_token": [batch_size, target_length] int tensor with end
+            token positions.
+          "target_space_id": A scalar int from data_generators.problem.SpaceID.
+
+    Returns:
+      Final decoder representation. Dictionary containing the following fields:
+        "targets": [batch_size, target_length, hidden_dim] float tensor with
+          decoder outputs
+        "targets_error_tag": [batch_size, target_length, hidden_dim] float
+          tensor with decoder outputs
+        "targets_start_token": [batch_size, input_length, target_length] float
+          tensor with start token position logits
+        "targets_end_token": [batch_size, input_length, target_length] float
+          tensor with end token position logits
+    """
+    hparams = self._hparams
+
+    losses = []
+
+    if self.has_input:
+      target_space = features["target_space_id"]
+      encoder_output, encoder_decoder_attention_bias = self.encode(
+          features["inputs"],
+          target_space,
+          hparams,
+          features=features,
+          losses=losses)
+    else:
+      encoder_output, encoder_decoder_attention_bias = (None, None)
+
+    targets = features["targets"]
+    targets_shape = common_layers.shape_list(targets)
+    targets = common_layers.flatten4d3d(targets)
+    decoder_input, decoder_self_attention_bias = self._prepare_decoder_fn(
+        targets, hparams, features=features)
+
+    nonpadding = features_to_nonpadding(features, "targets")
+
+    # Add edit ops layer to condition on start_token, end_token, and error_tag
+    decoder_input = transformer_edit_ops_layer(
+        decoder_input,
+        hparams,
+        encoder_output,
+        features,
+        nonpadding=nonpadding,
+        losses=losses)
+
+    if hparams.middle_prediction:
+      num_decoder_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
+      hparams.num_decoder_layers = int(
+          num_decoder_layers / hparams.middle_prediction_layer_factor)
+
+    decode_kwargs = {}
+    decoder_output = self.decode(
+        decoder_input,
+        encoder_output,
+        encoder_decoder_attention_bias,
+        decoder_self_attention_bias,
+        hparams,
+        nonpadding=nonpadding,
+        losses=losses,
+        **decode_kwargs)
+
+    loss_mask = common_layers.weights_nonzero(
+        maybe_flatten4d2d(features["targets_raw"]))
+    self.loss_den = tf.reduce_sum(loss_mask)
+    decoder_output = self._prediction_cascade(
+        hparams=hparams,
+        features=features,
+        losses=losses,
+        loss_mask=loss_mask,
+        nonpadding=nonpadding,
+        encoder_decoder_attention_bias=encoder_decoder_attention_bias,
+        encoder_output=encoder_output,
+        decoder_output=decoder_output)
+
+    if hparams.middle_prediction:
+      with tf.variable_scope("after_prediction"):
+        decoder_output = self.decode(
+            decoder_input + decoder_output,
+            encoder_output,
+            encoder_decoder_attention_bias,
+            decoder_self_attention_bias,
+            hparams,
+            nonpadding=nonpadding,
+            losses=losses,
+            **decode_kwargs)
+
+    ret = {"targets": tf.reshape(decoder_output, targets_shape)}
+    ret.update(self.logits)
+    if losses:
+      return ret, {"extra_loss": tf.add_n(losses)}
+    else:
+      return ret
+
+  def _prediction_cascade(self, hparams, features, losses, loss_mask,
+                          nonpadding, encoder_decoder_attention_bias,
+                          encoder_output, decoder_output):
+    if hparams.use_error_tags:
+      (decoder_output, error_tag_logits,
+       error_tag_loss) = transformer_error_tag_prediction_layer(
+           decoder_output, hparams, features, loss_mask=loss_mask)
+      self.logits["targets_error_tag"] = error_tag_logits
+      self.loss_num["targets_error_tag"] = error_tag_loss
+      decoder_output = transformer_between_predictions_layer(
+          decoder_output,
+          hparams,
+          name="post_error_tag",
+          nonpadding=nonpadding,
+          losses=losses)
+
+    pos_feat_names = []
+    if hparams.use_start_token:
+      pos_feat_names.append("targets_start_token")
+    pos_feat_names.append("targets_end_token")
+    for pos_feat_name in pos_feat_names:
+      (decoder_output, pos_logits,
+       pos_loss) = transformer_pointer_prediction_layer(
+           pos_feat_name,
+           encoder_output,
+           decoder_output,
+           encoder_decoder_attention_bias,
+           hparams,
+           features,
+           loss_mask=loss_mask)
+      self.logits[pos_feat_name] = pos_logits
+      self.loss_num[pos_feat_name] = pos_loss
+      decoder_output = transformer_between_predictions_layer(
+          decoder_output,
+          hparams,
+          name="post_%s" % pos_feat_name,
+          nonpadding=nonpadding,
+          losses=losses)
+    return decoder_output
+
+  def _loss_single(self, logits, feature_name, feature, weights=None):
+    """Prevents modality loss computation for targets_*."""
+    if feature_name in [
+        "targets_start_token", "targets_end_token", "targets_error_tag"
+    ]:
+      loss_num = self.loss_num[feature_name]
+      loss_num *= self._problem_hparams.loss_multiplier
+      loss_den = self.loss_den
+    else:
+      loss_num, loss_den = super(TransformerSeq2edits,
+                                 self)._loss_single(logits, feature_name,
+                                                    feature, weights)
+    tf.summary.scalar("loss/%s" % feature_name, loss_num / loss_den)
+    return loss_num, loss_den
+
+  def top(self, body_output, features):
+    """Adds additional dimensions and then calls super class implementation."""
+    exp_features = features
+    for feat in body_output.keys():
+      while len(body_output[feat].shape) < 4:
+        logging.warning("Expanding body output %s...", feat)
+        body_output[feat] = tf.expand_dims(body_output[feat], -2)
+      if feat in exp_features:
+        while len(exp_features[feat].shape) < 4:
+          exp_features[feat] = tf.expand_dims(exp_features[feat], -1)
+          logging.warning("Expanding feature %s...", feat)
+    return super(TransformerSeq2edits, self).top(body_output, exp_features)
+
+
+def _pointer_feedback(pointers, encoder_output, shift=True):
+  """Feedback loop for pointer networks.
+
+  Args:
+    pointers: [batch_size, target_length] int tensor with pointers into the
+      source sentence.
+    encoder_output: [batch_size, input_length, hidden_size] tensor with encoder
+      outputs.
+    shift: Whether to shift the pointers to the right.
+
+  Returns:
+    A [batch_size, target_length, hidden_size] tensor with encoder outputs.
+  """
+  if shift:
+    pointers = common_layers.shift_right_2d(pointers)
+  return gather_2d(encoder_output, pointers)
+
+
+def transformer_edit_ops_layer(decoder_input,
+                               hparams,
+                               encoder_output,
+                               features,
+                               cache=None,
+                               decode_loop_step=None,
+                               nonpadding=None,
+                               losses=None,
+                               layer_collection=None):
+  """Layer that conditions on the error tag and start and end token pointers."""
+  if isinstance(encoder_output, list):  # Select forward encoder
+    encoder_output = encoder_output[0]
+  with tf.variable_scope("edit_ops_layer"):
+    with tf.variable_scope("ffn"):
+      x = decoder_input
+      # Shorthand for layer preprocessing
+      # pylint: disable=g-long-lambda
+      preproc = lambda z: common_layers.layer_preprocess(
+          z, hparams, layer_collection=layer_collection)
+      # pylint: enable=g-long-lambda
+
+      feedback_start_token = (hparams.use_start_token or
+                              not hparams.feedback_end_token)
+      if feedback_start_token:
+        start_token = _pointer_feedback(
+            features["targets_start_token"],
+            encoder_output,
+            shift=hparams.feedback_end_token)
+      if hparams.feedback_end_token:
+        end_token = _pointer_feedback(features["targets_end_token"],
+                                      encoder_output)
+      layer_inputs = [preproc(x)]
+      if hparams.use_error_tags:
+        error_tags = common_layers.shift_right_3d(
+            common_layers.flatten4d3d(features["targets_error_tag"]))
+        layer_inputs.append(preproc(error_tags))
+      if feedback_start_token:
+        layer_inputs.append(start_token)
+      if hparams.feedback_end_token:
+        layer_inputs.append(end_token)
+      y = transformer_layers.transformer_ffn_layer(
+          tf.concat(layer_inputs, axis=2),
+          hparams,
+          conv_padding="LEFT",
+          nonpadding_mask=nonpadding,
+          losses=losses,
+          cache=cache,
+          decode_loop_step=decode_loop_step,
+          layer_collection=layer_collection)
+      x = common_layers.layer_postprocess(x, y, hparams)
+      return x
+
+
+def transformer_between_predictions_layer(x,
+                                          hparams,
+                                          name,
+                                          cache=None,
+                                          decode_loop_step=None,
+                                          nonpadding=None,
+                                          losses=None,
+                                          layer_collection=None):
+  """Stack between prediction layers."""
+  with tf.variable_scope(name):
+    for i in range(hparams.ffn_in_prediction_cascade):
+      with tf.variable_scope("layer_%d" % i):
+        y = transformer_layers.transformer_ffn_layer(
+            common_layers.layer_preprocess(
+                x, hparams, layer_collection=layer_collection),
+            hparams,
+            conv_padding="LEFT",
+            nonpadding_mask=nonpadding,
+            losses=losses,
+            cache=cache,
+            decode_loop_step=decode_loop_step,
+            layer_collection=layer_collection)
+        x = common_layers.layer_postprocess(x, y, hparams)
+  return x
+
+
+def get_error_tag_embedding_matrix():
+  candidates = [
+      var for var in tf.global_variables() if "targets_error_tag" in var.op.name
+  ]
+  if len(candidates) != 1:
+    raise ValueError("Could not identify error tag embedding matrix! "
+                     "Matching variable names: %s" % candidates)
+  embed_mat = candidates[0]
+  return embed_mat
+
+
+def transformer_error_tag_prediction_layer(x,
+                                           hparams,
+                                           features,
+                                           loss_mask,
+                                           layer_collection=None):
+  """Layer that predicts the error tag."""
+  with tf.variable_scope("error_tag_prediction"):
+    x = maybe_flatten4d3d(x)
+    vocab_size = hparams.problem.feature_info["targets_error_tag"].vocab_size
+    labels = features["targets_error_tag_raw"]
+    with tf.variable_scope("projection"):
+      bottleneck = common_layers.dense(
+          x,
+          hparams.error_tag_embed_size,
+          layer_collection=layer_collection,
+          name="bottleneck")
+      logits = common_layers.dense(
+          bottleneck,
+          vocab_size,
+          use_bias=False,
+          layer_collection=layer_collection,
+          name="logits")
+      xent = tf.nn.sparse_softmax_cross_entropy_with_logits(
+          logits=logits, labels=labels)
+      loss = tf.reduce_sum(xent * loss_mask)
+    with tf.variable_scope("embedding"):
+      embed_mat = get_error_tag_embedding_matrix()
+      y = common_layers.layer_preprocess(
+          common_layers.embedding(
+              labels, vocab_size, hparams.hidden_size, embedding_var=embed_mat),
+          hparams,
+          layer_collection=layer_collection)
+      x = common_layers.layer_postprocess(x, y, hparams)
+    return x, logits, loss
+
+
+def transformer_pointer_prediction_layer(feature_name,
+                                         encoder_output,
+                                         x,
+                                         encoder_decoder_attention_bias,
+                                         hparams,
+                                         features,
+                                         loss_mask,
+                                         layer_collection=None):
+  """Layer that predicts the start or end token position.
+
+  Args:
+    feature_name: 'targets_start_token' or 'targets_end_token'
+    encoder_output: [batch_size, input_length, hidden_size] tensor with encoder
+      outputs
+    x: [batch_size, target_length, 1, hidden_size] tensor with decoder outputs
+    encoder_decoder_attention_bias: [batch_size, input_length, target_length]
+      attention mask
+    hparams: Hyper parameters
+    features: Feature dictionary
+    loss_mask: [batch_size, target_length] mask for loss computation.
+    layer_collection: Layer collection
+
+  Returns:
+    (x, logits, loss)
+  """
+  if isinstance(encoder_output, list):
+    pointer_encoder_output = encoder_output[1]
+    encoder_output = sum(encoder_output)
+  else:
+    pointer_encoder_output = encoder_output
+  with tf.variable_scope("%s_prediction" % feature_name):
+    x = maybe_flatten4d3d(x)
+    encoder_decoder_attention_bias = common_layers.flatten4d3d(
+        encoder_decoder_attention_bias)
+    q = common_attention.compute_attention_component(x, hparams.hidden_size)
+    k = common_attention.compute_attention_component(encoder_output,
+                                                     hparams.hidden_size)
+    # Scaled dot-product attention
+    scalar = tf.rsqrt(tf.to_float(common_layers.shape_list(q)[2]))
+    logits = tf.matmul(q * scalar, k, transpose_b=True)
+
+    logits += encoder_decoder_attention_bias
+
+    labels = features["%s_raw" % feature_name]
+    xent = tf.nn.sparse_softmax_cross_entropy_with_logits(
+        logits=logits, labels=labels)
+    loss = tf.reduce_sum(xent * loss_mask)
+
+    pointer_out = gather_2d(pointer_encoder_output, labels)
+    y = common_layers.layer_preprocess(
+        pointer_out, hparams, layer_collection=layer_collection)
+    x = common_layers.layer_postprocess(x, y, hparams)
+    return x, logits, loss

From 9902e885fff48d610034b52153f7769081fcf2e0 Mon Sep 17 00:00:00 2001
From: Tiago Quelhas <tjgq@google.com>
Date: Wed, 10 Feb 2021 10:44:04 -0800
Subject: [PATCH 2687/2720] Internal change

PiperOrigin-RevId: 356771804
---
 tensor2tensor/__init__.py                                       | 2 +-
 tensor2tensor/bin/__init__.py                                   | 2 +-
 tensor2tensor/bin/build_vocab.py                                | 2 +-
 tensor2tensor/bin/make_tf_configs.py                            | 2 +-
 tensor2tensor/bin/t2t_attack.py                                 | 2 +-
 tensor2tensor/bin/t2t_avg_all.py                                | 2 +-
 tensor2tensor/bin/t2t_bleu.py                                   | 2 +-
 tensor2tensor/bin/t2t_datagen.py                                | 2 +-
 tensor2tensor/bin/t2t_decoder.py                                | 2 +-
 tensor2tensor/bin/t2t_distill.py                                | 2 +-
 tensor2tensor/bin/t2t_eval.py                                   | 2 +-
 tensor2tensor/bin/t2t_prune.py                                  | 2 +-
 tensor2tensor/bin/t2t_trainer.py                                | 2 +-
 tensor2tensor/bin/t2t_trainer_test.py                           | 2 +-
 tensor2tensor/bin/t2t_translate_all.py                          | 2 +-
 tensor2tensor/data_generators/__init__.py                       | 2 +-
 tensor2tensor/data_generators/algorithmic.py                    | 2 +-
 tensor2tensor/data_generators/algorithmic_math.py               | 2 +-
 tensor2tensor/data_generators/algorithmic_math_deepmind.py      | 2 +-
 tensor2tensor/data_generators/algorithmic_math_test.py          | 2 +-
 tensor2tensor/data_generators/algorithmic_math_two_variables.py | 2 +-
 tensor2tensor/data_generators/algorithmic_test.py               | 2 +-
 tensor2tensor/data_generators/all_problems.py                   | 2 +-
 tensor2tensor/data_generators/allen_brain.py                    | 2 +-
 tensor2tensor/data_generators/allen_brain_test.py               | 2 +-
 tensor2tensor/data_generators/audio.py                          | 2 +-
 tensor2tensor/data_generators/audio_encoder.py                  | 2 +-
 tensor2tensor/data_generators/audio_test.py                     | 2 +-
 tensor2tensor/data_generators/babi_qa.py                        | 2 +-
 tensor2tensor/data_generators/bair_robot_pushing.py             | 2 +-
 tensor2tensor/data_generators/celeba.py                         | 2 +-
 tensor2tensor/data_generators/celeba_test.py                    | 2 +-
 tensor2tensor/data_generators/celebahq.py                       | 2 +-
 tensor2tensor/data_generators/cifar.py                          | 2 +-
 tensor2tensor/data_generators/cipher.py                         | 2 +-
 tensor2tensor/data_generators/cleaner_en_xx.py                  | 2 +-
 tensor2tensor/data_generators/cnn_dailymail.py                  | 2 +-
 tensor2tensor/data_generators/cola.py                           | 2 +-
 tensor2tensor/data_generators/common_voice.py                   | 2 +-
 tensor2tensor/data_generators/common_voice_test.py              | 2 +-
 tensor2tensor/data_generators/conll_ner.py                      | 2 +-
 tensor2tensor/data_generators/desc2code.py                      | 2 +-
 tensor2tensor/data_generators/desc2code_test.py                 | 2 +-
 tensor2tensor/data_generators/dialog_abstract.py                | 2 +-
 tensor2tensor/data_generators/dialog_cornell.py                 | 2 +-
 tensor2tensor/data_generators/dialog_dailydialog.py             | 2 +-
 tensor2tensor/data_generators/dialog_opensubtitles.py           | 2 +-
 tensor2tensor/data_generators/dialog_personachat.py             | 2 +-
 tensor2tensor/data_generators/dna_encoder.py                    | 2 +-
 tensor2tensor/data_generators/dna_encoder_test.py               | 2 +-
 tensor2tensor/data_generators/enwik8.py                         | 2 +-
 tensor2tensor/data_generators/fsns.py                           | 2 +-
 tensor2tensor/data_generators/function_docstring.py             | 2 +-
 tensor2tensor/data_generators/gene_expression.py                | 2 +-
 tensor2tensor/data_generators/gene_expression_test.py           | 2 +-
 tensor2tensor/data_generators/generator_utils.py                | 2 +-
 tensor2tensor/data_generators/generator_utils_test.py           | 2 +-
 tensor2tensor/data_generators/google_robot_pushing.py           | 2 +-
 tensor2tensor/data_generators/gym_env.py                        | 2 +-
 tensor2tensor/data_generators/gym_env_test.py                   | 2 +-
 tensor2tensor/data_generators/ice_parsing.py                    | 2 +-
 tensor2tensor/data_generators/image_lsun.py                     | 2 +-
 tensor2tensor/data_generators/image_utils.py                    | 2 +-
 tensor2tensor/data_generators/image_utils_test.py               | 2 +-
 tensor2tensor/data_generators/imagenet.py                       | 2 +-
 tensor2tensor/data_generators/imagenet_test.py                  | 2 +-
 tensor2tensor/data_generators/imdb.py                           | 2 +-
 tensor2tensor/data_generators/inspect_tfrecord.py               | 2 +-
 tensor2tensor/data_generators/lambada.py                        | 2 +-
 tensor2tensor/data_generators/librispeech.py                    | 2 +-
 tensor2tensor/data_generators/lm1b.py                           | 2 +-
 tensor2tensor/data_generators/lm1b_imdb.py                      | 2 +-
 tensor2tensor/data_generators/lm1b_mnli.py                      | 2 +-
 tensor2tensor/data_generators/mnist.py                          | 2 +-
 tensor2tensor/data_generators/moving_mnist.py                   | 2 +-
 tensor2tensor/data_generators/mrpc.py                           | 2 +-
 tensor2tensor/data_generators/mscoco.py                         | 2 +-
 tensor2tensor/data_generators/mscoco_test.py                    | 2 +-
 tensor2tensor/data_generators/multi_problem.py                  | 2 +-
 tensor2tensor/data_generators/multi_problem_v2.py               | 2 +-
 tensor2tensor/data_generators/multi_problem_v2_test.py          | 2 +-
 tensor2tensor/data_generators/multinli.py                       | 2 +-
 tensor2tensor/data_generators/ocr.py                            | 2 +-
 tensor2tensor/data_generators/ops/pack_sequences_ops_test.py    | 2 +-
 .../data_generators/ops/subword_text_encoder_ops_test.py        | 2 +-
 tensor2tensor/data_generators/paraphrase_ms_coco.py             | 2 +-
 tensor2tensor/data_generators/paraphrase_ms_coco_test.py        | 2 +-
 tensor2tensor/data_generators/pointer_generator_word.py         | 2 +-
 tensor2tensor/data_generators/problem.py                        | 2 +-
 tensor2tensor/data_generators/problem_hparams.py                | 2 +-
 tensor2tensor/data_generators/problem_test.py                   | 2 +-
 tensor2tensor/data_generators/program_search.py                 | 2 +-
 tensor2tensor/data_generators/program_search_test.py            | 2 +-
 tensor2tensor/data_generators/ptb.py                            | 2 +-
 tensor2tensor/data_generators/qnli.py                           | 2 +-
 tensor2tensor/data_generators/quora_qpairs.py                   | 2 +-
 tensor2tensor/data_generators/rte.py                            | 2 +-
 tensor2tensor/data_generators/scitail.py                        | 2 +-
 tensor2tensor/data_generators/seq2edits.py                      | 2 +-
 tensor2tensor/data_generators/snli.py                           | 2 +-
 tensor2tensor/data_generators/speech_recognition.py             | 2 +-
 tensor2tensor/data_generators/squad.py                          | 2 +-
 tensor2tensor/data_generators/sst_binary.py                     | 2 +-
 tensor2tensor/data_generators/stanford_nli.py                   | 2 +-
 tensor2tensor/data_generators/style_transfer.py                 | 2 +-
 tensor2tensor/data_generators/style_transfer_test.py            | 2 +-
 tensor2tensor/data_generators/subject_verb_agreement.py         | 2 +-
 tensor2tensor/data_generators/text_encoder.py                   | 2 +-
 tensor2tensor/data_generators/text_encoder_build_subword.py     | 2 +-
 tensor2tensor/data_generators/text_encoder_test.py              | 2 +-
 tensor2tensor/data_generators/text_problems.py                  | 2 +-
 tensor2tensor/data_generators/text_problems_test.py             | 2 +-
 tensor2tensor/data_generators/timeseries.py                     | 2 +-
 tensor2tensor/data_generators/timeseries_data_generator.py      | 2 +-
 tensor2tensor/data_generators/timeseries_data_generator_test.py | 2 +-
 tensor2tensor/data_generators/timeseries_test.py                | 2 +-
 tensor2tensor/data_generators/tokenizer.py                      | 2 +-
 tensor2tensor/data_generators/tokenizer_test.py                 | 2 +-
 tensor2tensor/data_generators/transduction_problems.py          | 2 +-
 tensor2tensor/data_generators/transduction_problems_test.py     | 2 +-
 tensor2tensor/data_generators/translate.py                      | 2 +-
 tensor2tensor/data_generators/translate_encs.py                 | 2 +-
 tensor2tensor/data_generators/translate_encs_cubbitt.py         | 2 +-
 tensor2tensor/data_generators/translate_ende.py                 | 2 +-
 tensor2tensor/data_generators/translate_ende_test.py            | 2 +-
 tensor2tensor/data_generators/translate_enes.py                 | 2 +-
 tensor2tensor/data_generators/translate_enet.py                 | 2 +-
 tensor2tensor/data_generators/translate_enfr.py                 | 2 +-
 tensor2tensor/data_generators/translate_enid.py                 | 2 +-
 tensor2tensor/data_generators/translate_enmk.py                 | 2 +-
 tensor2tensor/data_generators/translate_enro.py                 | 2 +-
 tensor2tensor/data_generators/translate_entn.py                 | 2 +-
 tensor2tensor/data_generators/translate_envi.py                 | 2 +-
 tensor2tensor/data_generators/translate_enzh.py                 | 2 +-
 tensor2tensor/data_generators/translate_test.py                 | 2 +-
 tensor2tensor/data_generators/video_generated.py                | 2 +-
 tensor2tensor/data_generators/video_utils.py                    | 2 +-
 tensor2tensor/data_generators/video_utils_test.py               | 2 +-
 tensor2tensor/data_generators/vqa.py                            | 2 +-
 tensor2tensor/data_generators/vqa_utils.py                      | 2 +-
 tensor2tensor/data_generators/wiki.py                           | 2 +-
 tensor2tensor/data_generators/wiki_lm.py                        | 2 +-
 tensor2tensor/data_generators/wiki_multi_problems.py            | 2 +-
 tensor2tensor/data_generators/wiki_revision.py                  | 2 +-
 tensor2tensor/data_generators/wiki_revision_utils.py            | 2 +-
 tensor2tensor/data_generators/wikisum/__init__.py               | 2 +-
 tensor2tensor/data_generators/wikisum/generate_vocab.py         | 2 +-
 .../data_generators/wikisum/get_references_commoncrawl.py       | 2 +-
 tensor2tensor/data_generators/wikisum/get_references_web.py     | 2 +-
 .../data_generators/wikisum/get_references_web_single_group.py  | 2 +-
 tensor2tensor/data_generators/wikisum/html.py                   | 2 +-
 tensor2tensor/data_generators/wikisum/parallel_launch.py        | 2 +-
 tensor2tensor/data_generators/wikisum/produce_examples.py       | 2 +-
 tensor2tensor/data_generators/wikisum/utils.py                  | 2 +-
 tensor2tensor/data_generators/wikisum/utils_test.py             | 2 +-
 tensor2tensor/data_generators/wikisum/validate_data.py          | 2 +-
 tensor2tensor/data_generators/wikisum/wikisum.py                | 2 +-
 tensor2tensor/data_generators/wikitext103.py                    | 2 +-
 tensor2tensor/data_generators/wnli.py                           | 2 +-
 tensor2tensor/data_generators/wsj_parsing.py                    | 2 +-
 tensor2tensor/data_generators/yelp_full.py                      | 2 +-
 tensor2tensor/data_generators/yelp_polarity.py                  | 2 +-
 tensor2tensor/envs/__init__.py                                  | 2 +-
 tensor2tensor/envs/env_problem.py                               | 2 +-
 tensor2tensor/envs/env_problem_utils.py                         | 2 +-
 tensor2tensor/envs/env_problem_utils_test.py                    | 2 +-
 tensor2tensor/envs/gym_env_problem.py                           | 2 +-
 tensor2tensor/envs/gym_env_problem_test.py                      | 2 +-
 tensor2tensor/envs/gym_spaces_utils.py                          | 2 +-
 tensor2tensor/envs/gym_spaces_utils_test.py                     | 2 +-
 tensor2tensor/envs/mujoco_problems.py                           | 2 +-
 tensor2tensor/envs/mujoco_problems_test.py                      | 2 +-
 tensor2tensor/envs/rendered_env_problem.py                      | 2 +-
 tensor2tensor/envs/rendered_env_problem_test.py                 | 2 +-
 tensor2tensor/envs/tic_tac_toe_env.py                           | 2 +-
 tensor2tensor/envs/tic_tac_toe_env_problem.py                   | 2 +-
 tensor2tensor/envs/tic_tac_toe_env_problem_test.py              | 2 +-
 tensor2tensor/envs/tic_tac_toe_env_test.py                      | 2 +-
 tensor2tensor/envs/time_step.py                                 | 2 +-
 tensor2tensor/envs/time_step_test.py                            | 2 +-
 tensor2tensor/envs/trajectory.py                                | 2 +-
 tensor2tensor/envs/trajectory_test.py                           | 2 +-
 tensor2tensor/insights/__init__.py                              | 2 +-
 tensor2tensor/insights/graph.py                                 | 2 +-
 tensor2tensor/insights/polymer/common-types.js                  | 1 +
 tensor2tensor/insights/query_processor.py                       | 2 +-
 tensor2tensor/insights/server.py                                | 2 +-
 tensor2tensor/insights/transformer_model.py                     | 2 +-
 tensor2tensor/layers/__init__.py                                | 2 +-
 tensor2tensor/layers/area_attention.py                          | 2 +-
 tensor2tensor/layers/area_attention_test.py                     | 2 +-
 tensor2tensor/layers/common_attention.py                        | 2 +-
 tensor2tensor/layers/common_attention_test.py                   | 2 +-
 tensor2tensor/layers/common_audio.py                            | 2 +-
 tensor2tensor/layers/common_hparams.py                          | 2 +-
 tensor2tensor/layers/common_image_attention.py                  | 2 +-
 tensor2tensor/layers/common_image_attention_test.py             | 2 +-
 tensor2tensor/layers/common_layers.py                           | 2 +-
 tensor2tensor/layers/common_layers_test.py                      | 2 +-
 tensor2tensor/layers/common_video.py                            | 2 +-
 tensor2tensor/layers/common_video_test.py                       | 2 +-
 tensor2tensor/layers/discretization.py                          | 2 +-
 tensor2tensor/layers/discretization_test.py                     | 2 +-
 tensor2tensor/layers/latent_layers.py                           | 2 +-
 tensor2tensor/layers/latent_layers_test.py                      | 2 +-
 tensor2tensor/layers/message_passing_attention.py               | 2 +-
 tensor2tensor/layers/modalities.py                              | 2 +-
 tensor2tensor/layers/modalities_test.py                         | 2 +-
 tensor2tensor/layers/ngram.py                                   | 2 +-
 tensor2tensor/layers/ngram_test.py                              | 2 +-
 tensor2tensor/layers/transformer_glow_layers.py                 | 2 +-
 tensor2tensor/layers/transformer_glow_layers_ops.py             | 2 +-
 tensor2tensor/layers/transformer_glow_layers_ops_test.py        | 2 +-
 tensor2tensor/layers/transformer_glow_layers_test.py            | 2 +-
 tensor2tensor/layers/transformer_layers.py                      | 2 +-
 tensor2tensor/layers/transformer_memory.py                      | 2 +-
 tensor2tensor/layers/transformer_memory_test.py                 | 2 +-
 tensor2tensor/layers/vq_discrete.py                             | 2 +-
 tensor2tensor/layers/vqa_layers.py                              | 2 +-
 tensor2tensor/metrics/__init__.py                               | 2 +-
 tensor2tensor/metrics/video_conditional_fvd.py                  | 2 +-
 tensor2tensor/metrics/video_conditional_fvd_test.py             | 2 +-
 tensor2tensor/models/__init__.py                                | 2 +-
 tensor2tensor/models/basic.py                                   | 2 +-
 tensor2tensor/models/basic_test.py                              | 2 +-
 tensor2tensor/models/bytenet.py                                 | 2 +-
 tensor2tensor/models/bytenet_test.py                            | 2 +-
 tensor2tensor/models/distillation.py                            | 2 +-
 tensor2tensor/models/evolved_transformer.py                     | 2 +-
 tensor2tensor/models/evolved_transformer_test.py                | 2 +-
 tensor2tensor/models/image_transformer.py                       | 2 +-
 tensor2tensor/models/image_transformer_2d.py                    | 2 +-
 tensor2tensor/models/image_transformer_2d_test.py               | 2 +-
 tensor2tensor/models/image_transformer_test.py                  | 2 +-
 tensor2tensor/models/lstm.py                                    | 2 +-
 tensor2tensor/models/lstm_test.py                               | 2 +-
 tensor2tensor/models/mtf_image_transformer.py                   | 2 +-
 tensor2tensor/models/mtf_image_transformer_test.py              | 2 +-
 tensor2tensor/models/mtf_resnet.py                              | 2 +-
 tensor2tensor/models/mtf_transformer.py                         | 2 +-
 tensor2tensor/models/mtf_transformer2.py                        | 2 +-
 tensor2tensor/models/mtf_transformer_test.py                    | 2 +-
 tensor2tensor/models/neural_architecture_search/__init__.py     | 2 +-
 tensor2tensor/models/neural_architecture_search/nas_layers.py   | 2 +-
 .../models/neural_architecture_search/nas_layers_test.py        | 2 +-
 tensor2tensor/models/neural_architecture_search/nas_model.py    | 2 +-
 .../models/neural_architecture_search/nas_model_test.py         | 2 +-
 tensor2tensor/models/neural_assistant.py                        | 2 +-
 tensor2tensor/models/neural_gpu.py                              | 2 +-
 tensor2tensor/models/neural_gpu_test.py                         | 2 +-
 tensor2tensor/models/research/__init__.py                       | 2 +-
 tensor2tensor/models/research/adafactor_experiments.py          | 2 +-
 tensor2tensor/models/research/aligned.py                        | 2 +-
 tensor2tensor/models/research/attention_lm.py                   | 2 +-
 tensor2tensor/models/research/attention_lm_moe.py               | 2 +-
 tensor2tensor/models/research/autoencoders.py                   | 2 +-
 tensor2tensor/models/research/autoencoders_test.py              | 2 +-
 tensor2tensor/models/research/cycle_gan.py                      | 2 +-
 tensor2tensor/models/research/gene_expression.py                | 2 +-
 tensor2tensor/models/research/gene_expression_test.py           | 2 +-
 tensor2tensor/models/research/glow.py                           | 2 +-
 tensor2tensor/models/research/glow_init_hook.py                 | 2 +-
 tensor2tensor/models/research/glow_ops.py                       | 2 +-
 tensor2tensor/models/research/glow_ops_test.py                  | 2 +-
 tensor2tensor/models/research/glow_test.py                      | 2 +-
 tensor2tensor/models/research/lm_experiments.py                 | 2 +-
 tensor2tensor/models/research/moe.py                            | 2 +-
 tensor2tensor/models/research/moe_experiments.py                | 2 +-
 tensor2tensor/models/research/multiquery_paper.py               | 2 +-
 tensor2tensor/models/research/neural_stack.py                   | 2 +-
 tensor2tensor/models/research/neural_stack_test.py              | 2 +-
 tensor2tensor/models/research/residual_shuffle_exchange.py      | 2 +-
 tensor2tensor/models/research/rl.py                             | 2 +-
 tensor2tensor/models/research/shuffle_network.py                | 2 +-
 tensor2tensor/models/research/similarity_transformer.py         | 2 +-
 tensor2tensor/models/research/super_lm.py                       | 2 +-
 tensor2tensor/models/research/transformer_aux.py                | 2 +-
 tensor2tensor/models/research/transformer_aux_test.py           | 2 +-
 tensor2tensor/models/research/transformer_moe.py                | 2 +-
 tensor2tensor/models/research/transformer_nat.py                | 2 +-
 tensor2tensor/models/research/transformer_parallel.py           | 2 +-
 tensor2tensor/models/research/transformer_revnet.py             | 2 +-
 tensor2tensor/models/research/transformer_revnet_test.py        | 2 +-
 tensor2tensor/models/research/transformer_seq2edits.py          | 2 +-
 tensor2tensor/models/research/transformer_sketch.py             | 2 +-
 tensor2tensor/models/research/transformer_symshard.py           | 2 +-
 tensor2tensor/models/research/transformer_vae.py                | 2 +-
 tensor2tensor/models/research/transformer_vae_flow_prior.py     | 2 +-
 tensor2tensor/models/research/transformer_vae_flow_prior_ops.py | 2 +-
 tensor2tensor/models/research/transformer_vae_test.py           | 2 +-
 tensor2tensor/models/research/universal_transformer.py          | 2 +-
 tensor2tensor/models/research/universal_transformer_test.py     | 2 +-
 tensor2tensor/models/research/universal_transformer_util.py     | 2 +-
 tensor2tensor/models/research/vqa_attention.py                  | 2 +-
 tensor2tensor/models/research/vqa_attention_test.py             | 2 +-
 tensor2tensor/models/research/vqa_recurrent_self_attention.py   | 2 +-
 tensor2tensor/models/research/vqa_self_attention.py             | 2 +-
 tensor2tensor/models/resnet.py                                  | 2 +-
 tensor2tensor/models/resnet_test.py                             | 2 +-
 tensor2tensor/models/revnet.py                                  | 2 +-
 tensor2tensor/models/revnet_test.py                             | 2 +-
 tensor2tensor/models/shake_shake.py                             | 2 +-
 tensor2tensor/models/slicenet.py                                | 2 +-
 tensor2tensor/models/slicenet_test.py                           | 2 +-
 tensor2tensor/models/text_cnn.py                                | 2 +-
 tensor2tensor/models/transformer.py                             | 2 +-
 tensor2tensor/models/transformer_test.py                        | 2 +-
 tensor2tensor/models/vanilla_gan.py                             | 2 +-
 tensor2tensor/models/video/__init__.py                          | 2 +-
 tensor2tensor/models/video/base.py                              | 2 +-
 tensor2tensor/models/video/base_vae.py                          | 2 +-
 tensor2tensor/models/video/basic_deterministic.py               | 2 +-
 tensor2tensor/models/video/basic_deterministic_params.py        | 2 +-
 tensor2tensor/models/video/basic_deterministic_test.py          | 2 +-
 tensor2tensor/models/video/basic_recurrent.py                   | 2 +-
 tensor2tensor/models/video/basic_recurrent_test.py              | 2 +-
 tensor2tensor/models/video/basic_stochastic.py                  | 2 +-
 tensor2tensor/models/video/basic_stochastic_test.py             | 2 +-
 tensor2tensor/models/video/emily.py                             | 2 +-
 tensor2tensor/models/video/emily_test.py                        | 2 +-
 tensor2tensor/models/video/epva.py                              | 2 +-
 tensor2tensor/models/video/epva_params.py                       | 2 +-
 tensor2tensor/models/video/next_frame_glow.py                   | 2 +-
 tensor2tensor/models/video/nfg_conv3d_test.py                   | 2 +-
 tensor2tensor/models/video/nfg_conv_lstm_test.py                | 2 +-
 tensor2tensor/models/video/nfg_conv_test.py                     | 2 +-
 tensor2tensor/models/video/nfg_interpolate.py                   | 2 +-
 tensor2tensor/models/video/nfg_test_utils.py                    | 2 +-
 tensor2tensor/models/video/nfg_uncond_test.py                   | 2 +-
 tensor2tensor/models/video/savp.py                              | 2 +-
 tensor2tensor/models/video/savp_params.py                       | 2 +-
 tensor2tensor/models/video/savp_test.py                         | 2 +-
 tensor2tensor/models/video/sv2p.py                              | 2 +-
 tensor2tensor/models/video/sv2p_params.py                       | 2 +-
 tensor2tensor/models/video/sv2p_test.py                         | 2 +-
 tensor2tensor/models/video/tests_utils.py                       | 2 +-
 tensor2tensor/models/xception.py                                | 2 +-
 tensor2tensor/models/xception_test.py                           | 2 +-
 tensor2tensor/problems.py                                       | 2 +-
 tensor2tensor/problems_colab.py                                 | 2 +-
 tensor2tensor/problems_test.py                                  | 2 +-
 tensor2tensor/rl/__init__.py                                    | 2 +-
 tensor2tensor/rl/batch_dqn_agent_test.py                        | 2 +-
 tensor2tensor/rl/batch_runner_test.py                           | 2 +-
 tensor2tensor/rl/datagen_with_agent.py                          | 2 +-
 tensor2tensor/rl/dopamine_connector.py                          | 2 +-
 tensor2tensor/rl/envs/__init__.py                               | 2 +-
 tensor2tensor/rl/envs/in_graph_batch_env.py                     | 2 +-
 tensor2tensor/rl/envs/py_func_batch_env.py                      | 2 +-
 tensor2tensor/rl/envs/simulated_batch_env.py                    | 2 +-
 tensor2tensor/rl/envs/simulated_batch_gym_env.py                | 2 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py                      | 2 +-
 tensor2tensor/rl/evaluator.py                                   | 2 +-
 tensor2tensor/rl/evaluator_test.py                              | 2 +-
 tensor2tensor/rl/gym_utils.py                                   | 2 +-
 tensor2tensor/rl/gym_utils_test.py                              | 2 +-
 tensor2tensor/rl/player.py                                      | 2 +-
 tensor2tensor/rl/player_utils.py                                | 2 +-
 tensor2tensor/rl/policy_learner.py                              | 2 +-
 tensor2tensor/rl/ppo.py                                         | 2 +-
 tensor2tensor/rl/ppo_learner.py                                 | 2 +-
 tensor2tensor/rl/restarter.py                                   | 2 +-
 tensor2tensor/rl/restarter_test.py                              | 2 +-
 tensor2tensor/rl/rl_utils.py                                    | 2 +-
 tensor2tensor/rl/trainer_model_based.py                         | 2 +-
 tensor2tensor/rl/trainer_model_based_agent_only.py              | 2 +-
 tensor2tensor/rl/trainer_model_based_params.py                  | 2 +-
 tensor2tensor/rl/trainer_model_based_recurrent_test.py          | 2 +-
 tensor2tensor/rl/trainer_model_based_stochastic_test.py         | 2 +-
 tensor2tensor/rl/trainer_model_based_sv2p_test.py               | 2 +-
 tensor2tensor/rl/trainer_model_based_test.py                    | 2 +-
 tensor2tensor/rl/trainer_model_free.py                          | 2 +-
 tensor2tensor/rl/trainer_model_free_test.py                     | 2 +-
 tensor2tensor/rl/trainer_model_free_tictactoe_test.py           | 2 +-
 tensor2tensor/serving/__init__.py                               | 2 +-
 tensor2tensor/serving/export.py                                 | 2 +-
 tensor2tensor/serving/query.py                                  | 2 +-
 tensor2tensor/serving/serving_utils.py                          | 2 +-
 tensor2tensor/test_data/example_usr_dir/__init__.py             | 2 +-
 tensor2tensor/test_data/example_usr_dir/my_submodule.py         | 2 +-
 tensor2tensor/utils/__init__.py                                 | 2 +-
 tensor2tensor/utils/adafactor.py                                | 2 +-
 tensor2tensor/utils/adv_attack_utils.py                         | 2 +-
 tensor2tensor/utils/avg_checkpoints.py                          | 2 +-
 tensor2tensor/utils/beam_search.py                              | 2 +-
 tensor2tensor/utils/beam_search_test.py                         | 2 +-
 tensor2tensor/utils/bleu_hook.py                                | 2 +-
 tensor2tensor/utils/bleu_hook_test.py                           | 2 +-
 tensor2tensor/utils/checkpoint_compatibility_test.py            | 2 +-
 tensor2tensor/utils/cloud_mlengine.py                           | 2 +-
 tensor2tensor/utils/compute_video_metrics.py                    | 2 +-
 tensor2tensor/utils/contrib.py                                  | 2 +-
 tensor2tensor/utils/data_reader.py                              | 2 +-
 tensor2tensor/utils/data_reader_test.py                         | 2 +-
 tensor2tensor/utils/decoding.py                                 | 2 +-
 tensor2tensor/utils/devices.py                                  | 2 +-
 tensor2tensor/utils/diet.py                                     | 2 +-
 tensor2tensor/utils/diet_test.py                                | 2 +-
 tensor2tensor/utils/expert_utils.py                             | 2 +-
 tensor2tensor/utils/expert_utils_test.py                        | 2 +-
 tensor2tensor/utils/flags.py                                    | 2 +-
 tensor2tensor/utils/get_rouge.py                                | 2 +-
 tensor2tensor/utils/hparam.py                                   | 2 +-
 tensor2tensor/utils/hparam_test.py                              | 2 +-
 tensor2tensor/utils/hparams_lib.py                              | 2 +-
 tensor2tensor/utils/hparams_lib_test.py                         | 2 +-
 tensor2tensor/utils/learning_rate.py                            | 2 +-
 tensor2tensor/utils/metrics.py                                  | 2 +-
 tensor2tensor/utils/metrics_hook.py                             | 2 +-
 tensor2tensor/utils/metrics_hook_test.py                        | 2 +-
 tensor2tensor/utils/metrics_test.py                             | 2 +-
 tensor2tensor/utils/misc_utils.py                               | 2 +-
 tensor2tensor/utils/misc_utils_test.py                          | 2 +-
 tensor2tensor/utils/mlperf_log.py                               | 2 +-
 tensor2tensor/utils/mlperf_tags.py                              | 2 +-
 tensor2tensor/utils/mtf_model.py                                | 2 +-
 tensor2tensor/utils/multistep_optimizer.py                      | 2 +-
 tensor2tensor/utils/multistep_optimizer_test.py                 | 2 +-
 tensor2tensor/utils/multistep_with_adamoptimizer.py             | 2 +-
 tensor2tensor/utils/multistep_with_adamoptimizer_test.py        | 2 +-
 tensor2tensor/utils/optimize.py                                 | 2 +-
 tensor2tensor/utils/optimize_test.py                            | 2 +-
 tensor2tensor/utils/partial_checkpoint_load_hook.py             | 2 +-
 tensor2tensor/utils/pruning_utils.py                            | 2 +-
 tensor2tensor/utils/quantization.py                             | 2 +-
 tensor2tensor/utils/registry.py                                 | 2 +-
 tensor2tensor/utils/registry_test.py                            | 2 +-
 tensor2tensor/utils/restore_hook.py                             | 2 +-
 tensor2tensor/utils/rouge.py                                    | 2 +-
 tensor2tensor/utils/rouge_test.py                               | 2 +-
 tensor2tensor/utils/sari_hook.py                                | 2 +-
 tensor2tensor/utils/sari_hook_test.py                           | 2 +-
 tensor2tensor/utils/scheduled_sampling.py                       | 2 +-
 tensor2tensor/utils/t2t_model.py                                | 2 +-
 tensor2tensor/utils/t2t_model_test.py                           | 2 +-
 tensor2tensor/utils/test_utils.py                               | 2 +-
 tensor2tensor/utils/test_utils_test.py                          | 2 +-
 tensor2tensor/utils/trainer_lib.py                              | 2 +-
 tensor2tensor/utils/trainer_lib_test.py                         | 2 +-
 tensor2tensor/utils/update_ops_hook.py                          | 2 +-
 tensor2tensor/utils/usr_dir.py                                  | 2 +-
 tensor2tensor/utils/video/prediction2gif.py                     | 2 +-
 tensor2tensor/utils/video/reward_confusion.py                   | 2 +-
 tensor2tensor/utils/video2gif.py                                | 2 +-
 tensor2tensor/utils/video_metrics.py                            | 2 +-
 tensor2tensor/utils/video_metrics_test.py                       | 2 +-
 tensor2tensor/utils/yellowfin.py                                | 2 +-
 tensor2tensor/utils/yellowfin_test.py                           | 2 +-
 tensor2tensor/visualization/__init__.py                         | 2 +-
 tensor2tensor/visualization/attention.py                        | 2 +-
 tensor2tensor/visualization/visualization.py                    | 2 +-
 tensor2tensor/visualization/visualization_test.py               | 2 +-
 452 files changed, 452 insertions(+), 451 deletions(-)

diff --git a/tensor2tensor/__init__.py b/tensor2tensor/__init__.py
index 798e88593..7b0a2368a 100644
--- a/tensor2tensor/__init__.py
+++ b/tensor2tensor/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/__init__.py b/tensor2tensor/bin/__init__.py
index 798e88593..7b0a2368a 100644
--- a/tensor2tensor/bin/__init__.py
+++ b/tensor2tensor/bin/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/build_vocab.py b/tensor2tensor/bin/build_vocab.py
index 703bde0f5..da7d923dd 100644
--- a/tensor2tensor/bin/build_vocab.py
+++ b/tensor2tensor/bin/build_vocab.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/make_tf_configs.py b/tensor2tensor/bin/make_tf_configs.py
index 4156c2108..6d482e330 100644
--- a/tensor2tensor/bin/make_tf_configs.py
+++ b/tensor2tensor/bin/make_tf_configs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 8a7bce8f0..6a4c63b1f 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_avg_all.py b/tensor2tensor/bin/t2t_avg_all.py
index 986d83cfd..543c58ce2 100644
--- a/tensor2tensor/bin/t2t_avg_all.py
+++ b/tensor2tensor/bin/t2t_avg_all.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_bleu.py b/tensor2tensor/bin/t2t_bleu.py
index 2eccda85a..1f07078ee 100644
--- a/tensor2tensor/bin/t2t_bleu.py
+++ b/tensor2tensor/bin/t2t_bleu.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index c974acdfd..9826dc86f 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 9d6b0bfca..6a40a6a81 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_distill.py b/tensor2tensor/bin/t2t_distill.py
index 7e69ff9a4..383d2fc6b 100644
--- a/tensor2tensor/bin/t2t_distill.py
+++ b/tensor2tensor/bin/t2t_distill.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_eval.py b/tensor2tensor/bin/t2t_eval.py
index 7d2f1385b..536069a56 100644
--- a/tensor2tensor/bin/t2t_eval.py
+++ b/tensor2tensor/bin/t2t_eval.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
index d1e75f8e2..385574f45 100644
--- a/tensor2tensor/bin/t2t_prune.py
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 7ef63c856..60c7b4f11 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_trainer_test.py b/tensor2tensor/bin/t2t_trainer_test.py
index 0a7fee5a6..817e35c33 100644
--- a/tensor2tensor/bin/t2t_trainer_test.py
+++ b/tensor2tensor/bin/t2t_trainer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_translate_all.py b/tensor2tensor/bin/t2t_translate_all.py
index 4fba3de98..03eb56c0a 100644
--- a/tensor2tensor/bin/t2t_translate_all.py
+++ b/tensor2tensor/bin/t2t_translate_all.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/__init__.py b/tensor2tensor/data_generators/__init__.py
index 798e88593..7b0a2368a 100644
--- a/tensor2tensor/data_generators/__init__.py
+++ b/tensor2tensor/data_generators/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index 74ae810e9..36b2df49b 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math.py b/tensor2tensor/data_generators/algorithmic_math.py
index d3e5ee373..b71c40ab5 100644
--- a/tensor2tensor/data_generators/algorithmic_math.py
+++ b/tensor2tensor/data_generators/algorithmic_math.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math_deepmind.py b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
index 3f8a9c015..f8a44e4eb 100644
--- a/tensor2tensor/data_generators/algorithmic_math_deepmind.py
+++ b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math_test.py b/tensor2tensor/data_generators/algorithmic_math_test.py
index ff00e4e9b..c582e534f 100644
--- a/tensor2tensor/data_generators/algorithmic_math_test.py
+++ b/tensor2tensor/data_generators/algorithmic_math_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math_two_variables.py b/tensor2tensor/data_generators/algorithmic_math_two_variables.py
index d6c992685..d64175e11 100644
--- a/tensor2tensor/data_generators/algorithmic_math_two_variables.py
+++ b/tensor2tensor/data_generators/algorithmic_math_two_variables.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_test.py b/tensor2tensor/data_generators/algorithmic_test.py
index cf6bf1b33..523afcde6 100644
--- a/tensor2tensor/data_generators/algorithmic_test.py
+++ b/tensor2tensor/data_generators/algorithmic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 50bf98417..920a2a6c8 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index aef529a3e..b961c8f87 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/allen_brain_test.py b/tensor2tensor/data_generators/allen_brain_test.py
index 2c8a120a6..4c118ca9b 100644
--- a/tensor2tensor/data_generators/allen_brain_test.py
+++ b/tensor2tensor/data_generators/allen_brain_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio.py b/tensor2tensor/data_generators/audio.py
index 7bff3e84f..674fde652 100644
--- a/tensor2tensor/data_generators/audio.py
+++ b/tensor2tensor/data_generators/audio.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio_encoder.py b/tensor2tensor/data_generators/audio_encoder.py
index 21d60de5e..d23e5274a 100644
--- a/tensor2tensor/data_generators/audio_encoder.py
+++ b/tensor2tensor/data_generators/audio_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio_test.py b/tensor2tensor/data_generators/audio_test.py
index ff94e4430..ff3e21e09 100644
--- a/tensor2tensor/data_generators/audio_test.py
+++ b/tensor2tensor/data_generators/audio_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index ba1831dbd..20ff4f1eb 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index f9ccc5a6a..eecf55d66 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index b9d0848bf..237bfe92e 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celeba_test.py b/tensor2tensor/data_generators/celeba_test.py
index 59f932c09..08d837728 100644
--- a/tensor2tensor/data_generators/celeba_test.py
+++ b/tensor2tensor/data_generators/celeba_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celebahq.py b/tensor2tensor/data_generators/celebahq.py
index 55cbf42b6..dd1425c16 100644
--- a/tensor2tensor/data_generators/celebahq.py
+++ b/tensor2tensor/data_generators/celebahq.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index 181cdb991..bf00c3b54 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cipher.py b/tensor2tensor/data_generators/cipher.py
index 93cb6b5f9..f7d5e880d 100644
--- a/tensor2tensor/data_generators/cipher.py
+++ b/tensor2tensor/data_generators/cipher.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cleaner_en_xx.py b/tensor2tensor/data_generators/cleaner_en_xx.py
index 1d690d712..141da35b5 100644
--- a/tensor2tensor/data_generators/cleaner_en_xx.py
+++ b/tensor2tensor/data_generators/cleaner_en_xx.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index da00e1a3e..8ba01182a 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cola.py b/tensor2tensor/data_generators/cola.py
index c9ea13094..12b8a2530 100644
--- a/tensor2tensor/data_generators/cola.py
+++ b/tensor2tensor/data_generators/cola.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/common_voice.py b/tensor2tensor/data_generators/common_voice.py
index a0806659a..57e1b8728 100644
--- a/tensor2tensor/data_generators/common_voice.py
+++ b/tensor2tensor/data_generators/common_voice.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/common_voice_test.py b/tensor2tensor/data_generators/common_voice_test.py
index 302b7f929..a07e5c10e 100644
--- a/tensor2tensor/data_generators/common_voice_test.py
+++ b/tensor2tensor/data_generators/common_voice_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/conll_ner.py b/tensor2tensor/data_generators/conll_ner.py
index fa4e42034..3c4095b95 100644
--- a/tensor2tensor/data_generators/conll_ner.py
+++ b/tensor2tensor/data_generators/conll_ner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/desc2code.py b/tensor2tensor/data_generators/desc2code.py
index 6619130c0..e85b253d7 100644
--- a/tensor2tensor/data_generators/desc2code.py
+++ b/tensor2tensor/data_generators/desc2code.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/desc2code_test.py b/tensor2tensor/data_generators/desc2code_test.py
index 882db43db..3bc68db6e 100644
--- a/tensor2tensor/data_generators/desc2code_test.py
+++ b/tensor2tensor/data_generators/desc2code_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_abstract.py b/tensor2tensor/data_generators/dialog_abstract.py
index 4541b8bd7..eb4408742 100644
--- a/tensor2tensor/data_generators/dialog_abstract.py
+++ b/tensor2tensor/data_generators/dialog_abstract.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_cornell.py b/tensor2tensor/data_generators/dialog_cornell.py
index 408ad1c6a..4423cd06c 100644
--- a/tensor2tensor/data_generators/dialog_cornell.py
+++ b/tensor2tensor/data_generators/dialog_cornell.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_dailydialog.py b/tensor2tensor/data_generators/dialog_dailydialog.py
index 44634f7be..80083f532 100644
--- a/tensor2tensor/data_generators/dialog_dailydialog.py
+++ b/tensor2tensor/data_generators/dialog_dailydialog.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_opensubtitles.py b/tensor2tensor/data_generators/dialog_opensubtitles.py
index 62cdce509..9ee843da4 100644
--- a/tensor2tensor/data_generators/dialog_opensubtitles.py
+++ b/tensor2tensor/data_generators/dialog_opensubtitles.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_personachat.py b/tensor2tensor/data_generators/dialog_personachat.py
index c356a86ec..719f71fd6 100644
--- a/tensor2tensor/data_generators/dialog_personachat.py
+++ b/tensor2tensor/data_generators/dialog_personachat.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dna_encoder.py b/tensor2tensor/data_generators/dna_encoder.py
index c9b0e87ea..8db26b305 100644
--- a/tensor2tensor/data_generators/dna_encoder.py
+++ b/tensor2tensor/data_generators/dna_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dna_encoder_test.py b/tensor2tensor/data_generators/dna_encoder_test.py
index e354f2ef4..ada9ff37a 100644
--- a/tensor2tensor/data_generators/dna_encoder_test.py
+++ b/tensor2tensor/data_generators/dna_encoder_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/enwik8.py b/tensor2tensor/data_generators/enwik8.py
index e4da9a8d8..31ae913c0 100644
--- a/tensor2tensor/data_generators/enwik8.py
+++ b/tensor2tensor/data_generators/enwik8.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
index c1d3b60d5..d9649e8b7 100644
--- a/tensor2tensor/data_generators/fsns.py
+++ b/tensor2tensor/data_generators/fsns.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
index 0cbdccbe0..2e58ace24 100644
--- a/tensor2tensor/data_generators/function_docstring.py
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gene_expression.py b/tensor2tensor/data_generators/gene_expression.py
index 105f5d0c7..d2e35d8f9 100644
--- a/tensor2tensor/data_generators/gene_expression.py
+++ b/tensor2tensor/data_generators/gene_expression.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gene_expression_test.py b/tensor2tensor/data_generators/gene_expression_test.py
index bfe819d0c..a4820d8fb 100644
--- a/tensor2tensor/data_generators/gene_expression_test.py
+++ b/tensor2tensor/data_generators/gene_expression_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 36f54d47e..4d4b1a324 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/generator_utils_test.py b/tensor2tensor/data_generators/generator_utils_test.py
index 0cb378a1e..1778616cf 100644
--- a/tensor2tensor/data_generators/generator_utils_test.py
+++ b/tensor2tensor/data_generators/generator_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index 03f7fdf53..de51e5b59 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index e48f59fbb..9597c4abb 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 81a7c8462..f7d897620 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ice_parsing.py b/tensor2tensor/data_generators/ice_parsing.py
index 8ab9fa8d4..3a2c7b49d 100644
--- a/tensor2tensor/data_generators/ice_parsing.py
+++ b/tensor2tensor/data_generators/ice_parsing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_lsun.py b/tensor2tensor/data_generators/image_lsun.py
index 2a847f1fb..63403cb0c 100644
--- a/tensor2tensor/data_generators/image_lsun.py
+++ b/tensor2tensor/data_generators/image_lsun.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 3f439fa47..909464616 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_utils_test.py b/tensor2tensor/data_generators/image_utils_test.py
index b4434ef61..f8578db99 100644
--- a/tensor2tensor/data_generators/image_utils_test.py
+++ b/tensor2tensor/data_generators/image_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 18550a466..bd759d251 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imagenet_test.py b/tensor2tensor/data_generators/imagenet_test.py
index 89220ec4d..4ce37ac12 100644
--- a/tensor2tensor/data_generators/imagenet_test.py
+++ b/tensor2tensor/data_generators/imagenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imdb.py b/tensor2tensor/data_generators/imdb.py
index f7db47a0c..03b158bb3 100644
--- a/tensor2tensor/data_generators/imdb.py
+++ b/tensor2tensor/data_generators/imdb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/inspect_tfrecord.py b/tensor2tensor/data_generators/inspect_tfrecord.py
index 6b4239221..8fa9299c4 100644
--- a/tensor2tensor/data_generators/inspect_tfrecord.py
+++ b/tensor2tensor/data_generators/inspect_tfrecord.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lambada.py b/tensor2tensor/data_generators/lambada.py
index d26311eb5..3d76724b4 100644
--- a/tensor2tensor/data_generators/lambada.py
+++ b/tensor2tensor/data_generators/lambada.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index bb3a6552c..a4faa3d3d 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b.py b/tensor2tensor/data_generators/lm1b.py
index e608d1799..cd1a03a7f 100644
--- a/tensor2tensor/data_generators/lm1b.py
+++ b/tensor2tensor/data_generators/lm1b.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b_imdb.py b/tensor2tensor/data_generators/lm1b_imdb.py
index 2aca37a5c..dc719fc5d 100644
--- a/tensor2tensor/data_generators/lm1b_imdb.py
+++ b/tensor2tensor/data_generators/lm1b_imdb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b_mnli.py b/tensor2tensor/data_generators/lm1b_mnli.py
index 6c4a1b865..6cba7b9ae 100644
--- a/tensor2tensor/data_generators/lm1b_mnli.py
+++ b/tensor2tensor/data_generators/lm1b_mnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mnist.py b/tensor2tensor/data_generators/mnist.py
index 0e5100ba5..3570e3913 100644
--- a/tensor2tensor/data_generators/mnist.py
+++ b/tensor2tensor/data_generators/mnist.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/moving_mnist.py b/tensor2tensor/data_generators/moving_mnist.py
index c80e980fe..c48673350 100644
--- a/tensor2tensor/data_generators/moving_mnist.py
+++ b/tensor2tensor/data_generators/moving_mnist.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mrpc.py b/tensor2tensor/data_generators/mrpc.py
index 2c19f3d41..8031a2a89 100644
--- a/tensor2tensor/data_generators/mrpc.py
+++ b/tensor2tensor/data_generators/mrpc.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index 41ab4bdee..22be4c4e0 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mscoco_test.py b/tensor2tensor/data_generators/mscoco_test.py
index 5c66f390e..6542a51dc 100644
--- a/tensor2tensor/data_generators/mscoco_test.py
+++ b/tensor2tensor/data_generators/mscoco_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 9ac90b874..3f4449778 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem_v2.py b/tensor2tensor/data_generators/multi_problem_v2.py
index 82e37f7f3..d7fc02297 100644
--- a/tensor2tensor/data_generators/multi_problem_v2.py
+++ b/tensor2tensor/data_generators/multi_problem_v2.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem_v2_test.py b/tensor2tensor/data_generators/multi_problem_v2_test.py
index 6cbdba815..e8f1c41dc 100644
--- a/tensor2tensor/data_generators/multi_problem_v2_test.py
+++ b/tensor2tensor/data_generators/multi_problem_v2_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index daeb54f49..5d7a619fc 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ocr.py b/tensor2tensor/data_generators/ocr.py
index ed1b916fe..6cb4b9a21 100644
--- a/tensor2tensor/data_generators/ocr.py
+++ b/tensor2tensor/data_generators/ocr.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
index 443dae16d..4f7c58772 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py b/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
index c3b9273e3..b0a757138 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/paraphrase_ms_coco.py b/tensor2tensor/data_generators/paraphrase_ms_coco.py
index a7a30b710..9ce52fda5 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/paraphrase_ms_coco_test.py b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
index ac14e8b52..f58b33667 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/pointer_generator_word.py b/tensor2tensor/data_generators/pointer_generator_word.py
index ce31b551e..bbd018d69 100644
--- a/tensor2tensor/data_generators/pointer_generator_word.py
+++ b/tensor2tensor/data_generators/pointer_generator_word.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 85539a9fa..af211180d 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem_hparams.py b/tensor2tensor/data_generators/problem_hparams.py
index 7727a0cb6..ffdfabe77 100644
--- a/tensor2tensor/data_generators/problem_hparams.py
+++ b/tensor2tensor/data_generators/problem_hparams.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 820cbbf27..7751704a0 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/program_search.py b/tensor2tensor/data_generators/program_search.py
index b4c0cd229..4b1b46a7d 100644
--- a/tensor2tensor/data_generators/program_search.py
+++ b/tensor2tensor/data_generators/program_search.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/program_search_test.py b/tensor2tensor/data_generators/program_search_test.py
index 0bd3687c9..a4af895af 100644
--- a/tensor2tensor/data_generators/program_search_test.py
+++ b/tensor2tensor/data_generators/program_search_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ptb.py b/tensor2tensor/data_generators/ptb.py
index 7283ce23f..a39d06afb 100644
--- a/tensor2tensor/data_generators/ptb.py
+++ b/tensor2tensor/data_generators/ptb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/qnli.py b/tensor2tensor/data_generators/qnli.py
index a6b7f6079..af1765007 100644
--- a/tensor2tensor/data_generators/qnli.py
+++ b/tensor2tensor/data_generators/qnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/quora_qpairs.py b/tensor2tensor/data_generators/quora_qpairs.py
index a298af4b5..447263db6 100644
--- a/tensor2tensor/data_generators/quora_qpairs.py
+++ b/tensor2tensor/data_generators/quora_qpairs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/rte.py b/tensor2tensor/data_generators/rte.py
index 2dcae1284..e47536540 100644
--- a/tensor2tensor/data_generators/rte.py
+++ b/tensor2tensor/data_generators/rte.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/scitail.py b/tensor2tensor/data_generators/scitail.py
index a041fec5f..26f389bcf 100644
--- a/tensor2tensor/data_generators/scitail.py
+++ b/tensor2tensor/data_generators/scitail.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/seq2edits.py b/tensor2tensor/data_generators/seq2edits.py
index 35b08ec08..3b1afe287 100644
--- a/tensor2tensor/data_generators/seq2edits.py
+++ b/tensor2tensor/data_generators/seq2edits.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/snli.py b/tensor2tensor/data_generators/snli.py
index 010ace7f0..977fbc005 100644
--- a/tensor2tensor/data_generators/snli.py
+++ b/tensor2tensor/data_generators/snli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index 9cda94072..da5911cb8 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/squad.py b/tensor2tensor/data_generators/squad.py
index 9b9958adc..646028a4e 100644
--- a/tensor2tensor/data_generators/squad.py
+++ b/tensor2tensor/data_generators/squad.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/sst_binary.py b/tensor2tensor/data_generators/sst_binary.py
index 2e6410a3e..14909adc4 100644
--- a/tensor2tensor/data_generators/sst_binary.py
+++ b/tensor2tensor/data_generators/sst_binary.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/stanford_nli.py b/tensor2tensor/data_generators/stanford_nli.py
index a086ed61b..dbd7832e5 100644
--- a/tensor2tensor/data_generators/stanford_nli.py
+++ b/tensor2tensor/data_generators/stanford_nli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/style_transfer.py b/tensor2tensor/data_generators/style_transfer.py
index 33cd352c0..291b87ff4 100644
--- a/tensor2tensor/data_generators/style_transfer.py
+++ b/tensor2tensor/data_generators/style_transfer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/style_transfer_test.py b/tensor2tensor/data_generators/style_transfer_test.py
index b212f6125..572a9915b 100644
--- a/tensor2tensor/data_generators/style_transfer_test.py
+++ b/tensor2tensor/data_generators/style_transfer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/subject_verb_agreement.py b/tensor2tensor/data_generators/subject_verb_agreement.py
index 55c025b6b..27ee080c6 100644
--- a/tensor2tensor/data_generators/subject_verb_agreement.py
+++ b/tensor2tensor/data_generators/subject_verb_agreement.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 3c17045ca..0ab6fe0b3 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_encoder_build_subword.py b/tensor2tensor/data_generators/text_encoder_build_subword.py
index 826072ee8..97e52cd9b 100644
--- a/tensor2tensor/data_generators/text_encoder_build_subword.py
+++ b/tensor2tensor/data_generators/text_encoder_build_subword.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_encoder_test.py b/tensor2tensor/data_generators/text_encoder_test.py
index 92cc9cecc..6ca9c7648 100644
--- a/tensor2tensor/data_generators/text_encoder_test.py
+++ b/tensor2tensor/data_generators/text_encoder_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 6ce5a5d87..ab51c5918 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_problems_test.py b/tensor2tensor/data_generators/text_problems_test.py
index 28ea85d20..dd897bf7e 100644
--- a/tensor2tensor/data_generators/text_problems_test.py
+++ b/tensor2tensor/data_generators/text_problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index e84af305c..87e176cf5 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_data_generator.py b/tensor2tensor/data_generators/timeseries_data_generator.py
index 8d697184c..e5476423c 100644
--- a/tensor2tensor/data_generators/timeseries_data_generator.py
+++ b/tensor2tensor/data_generators/timeseries_data_generator.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_data_generator_test.py b/tensor2tensor/data_generators/timeseries_data_generator_test.py
index 98a386102..e32bc5a81 100644
--- a/tensor2tensor/data_generators/timeseries_data_generator_test.py
+++ b/tensor2tensor/data_generators/timeseries_data_generator_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index 9c8edd98a..02157944a 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/tokenizer.py b/tensor2tensor/data_generators/tokenizer.py
index 46631a41b..21b4dddaf 100644
--- a/tensor2tensor/data_generators/tokenizer.py
+++ b/tensor2tensor/data_generators/tokenizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/tokenizer_test.py b/tensor2tensor/data_generators/tokenizer_test.py
index c4d93b723..bd7272d5d 100644
--- a/tensor2tensor/data_generators/tokenizer_test.py
+++ b/tensor2tensor/data_generators/tokenizer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/transduction_problems.py b/tensor2tensor/data_generators/transduction_problems.py
index 2471dd343..e9d4c8626 100644
--- a/tensor2tensor/data_generators/transduction_problems.py
+++ b/tensor2tensor/data_generators/transduction_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/transduction_problems_test.py b/tensor2tensor/data_generators/transduction_problems_test.py
index c3b8f57ca..e35006f25 100644
--- a/tensor2tensor/data_generators/transduction_problems_test.py
+++ b/tensor2tensor/data_generators/transduction_problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 668a3a5c0..2827daf44 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_encs.py b/tensor2tensor/data_generators/translate_encs.py
index 95819486d..657586440 100644
--- a/tensor2tensor/data_generators/translate_encs.py
+++ b/tensor2tensor/data_generators/translate_encs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_encs_cubbitt.py b/tensor2tensor/data_generators/translate_encs_cubbitt.py
index 1a24178ff..7caf37303 100644
--- a/tensor2tensor/data_generators/translate_encs_cubbitt.py
+++ b/tensor2tensor/data_generators/translate_encs_cubbitt.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index 140def85d..da3705f19 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_ende_test.py b/tensor2tensor/data_generators/translate_ende_test.py
index 74acc0f4d..48c6753c6 100644
--- a/tensor2tensor/data_generators/translate_ende_test.py
+++ b/tensor2tensor/data_generators/translate_ende_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enes.py b/tensor2tensor/data_generators/translate_enes.py
index 3e3fcf203..8feb459e0 100644
--- a/tensor2tensor/data_generators/translate_enes.py
+++ b/tensor2tensor/data_generators/translate_enes.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enet.py b/tensor2tensor/data_generators/translate_enet.py
index 19335beb1..8256b0d9f 100644
--- a/tensor2tensor/data_generators/translate_enet.py
+++ b/tensor2tensor/data_generators/translate_enet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index 16067adf8..63a98bdab 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enid.py b/tensor2tensor/data_generators/translate_enid.py
index 4367fd0a0..bfa112dce 100644
--- a/tensor2tensor/data_generators/translate_enid.py
+++ b/tensor2tensor/data_generators/translate_enid.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enmk.py b/tensor2tensor/data_generators/translate_enmk.py
index c897e4e11..2b4e74f9e 100644
--- a/tensor2tensor/data_generators/translate_enmk.py
+++ b/tensor2tensor/data_generators/translate_enmk.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enro.py b/tensor2tensor/data_generators/translate_enro.py
index 83e4a1557..ea2a9e6cf 100644
--- a/tensor2tensor/data_generators/translate_enro.py
+++ b/tensor2tensor/data_generators/translate_enro.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_entn.py b/tensor2tensor/data_generators/translate_entn.py
index d9d2de7e2..d1ea587fc 100644
--- a/tensor2tensor/data_generators/translate_entn.py
+++ b/tensor2tensor/data_generators/translate_entn.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_envi.py b/tensor2tensor/data_generators/translate_envi.py
index 6ebef9eab..fad6f850d 100644
--- a/tensor2tensor/data_generators/translate_envi.py
+++ b/tensor2tensor/data_generators/translate_envi.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index 8d6719975..bbbf661dc 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_test.py b/tensor2tensor/data_generators/translate_test.py
index 05e2fc8a8..3d5928bcf 100644
--- a/tensor2tensor/data_generators/translate_test.py
+++ b/tensor2tensor/data_generators/translate_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 5bf66fbc2..18d621990 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 22ef3545c..6e80102ec 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index cc3b2087d..659a555da 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
index f85f57f81..e027c4f7a 100644
--- a/tensor2tensor/data_generators/vqa.py
+++ b/tensor2tensor/data_generators/vqa.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/vqa_utils.py b/tensor2tensor/data_generators/vqa_utils.py
index 90889f93a..a4e363645 100644
--- a/tensor2tensor/data_generators/vqa_utils.py
+++ b/tensor2tensor/data_generators/vqa_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki.py b/tensor2tensor/data_generators/wiki.py
index 149b31574..affadd6de 100644
--- a/tensor2tensor/data_generators/wiki.py
+++ b/tensor2tensor/data_generators/wiki.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_lm.py b/tensor2tensor/data_generators/wiki_lm.py
index 27b67dc65..7c2fceb3f 100644
--- a/tensor2tensor/data_generators/wiki_lm.py
+++ b/tensor2tensor/data_generators/wiki_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index a3fcd6a06..2f1cac9f1 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_revision.py b/tensor2tensor/data_generators/wiki_revision.py
index b3f283b9f..77942f84f 100644
--- a/tensor2tensor/data_generators/wiki_revision.py
+++ b/tensor2tensor/data_generators/wiki_revision.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index 73dc3ea13..f938abb77 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/__init__.py b/tensor2tensor/data_generators/wikisum/__init__.py
index 798e88593..7b0a2368a 100644
--- a/tensor2tensor/data_generators/wikisum/__init__.py
+++ b/tensor2tensor/data_generators/wikisum/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/generate_vocab.py b/tensor2tensor/data_generators/wikisum/generate_vocab.py
index e2a61927b..6665d1c02 100644
--- a/tensor2tensor/data_generators/wikisum/generate_vocab.py
+++ b/tensor2tensor/data_generators/wikisum/generate_vocab.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py b/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
index 2749de41f..7925715a6 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_web.py b/tensor2tensor/data_generators/wikisum/get_references_web.py
index ef098b228..d3deadd31 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py b/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
index 1a74e76bc..1fcce1010 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/html.py b/tensor2tensor/data_generators/wikisum/html.py
index c3749385d..38ecd194f 100644
--- a/tensor2tensor/data_generators/wikisum/html.py
+++ b/tensor2tensor/data_generators/wikisum/html.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/parallel_launch.py b/tensor2tensor/data_generators/wikisum/parallel_launch.py
index dfdd2305e..c69cb866e 100644
--- a/tensor2tensor/data_generators/wikisum/parallel_launch.py
+++ b/tensor2tensor/data_generators/wikisum/parallel_launch.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/produce_examples.py b/tensor2tensor/data_generators/wikisum/produce_examples.py
index 1f3c20ef9..7a689978b 100644
--- a/tensor2tensor/data_generators/wikisum/produce_examples.py
+++ b/tensor2tensor/data_generators/wikisum/produce_examples.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/utils.py b/tensor2tensor/data_generators/wikisum/utils.py
index f63115d3f..b3c124775 100644
--- a/tensor2tensor/data_generators/wikisum/utils.py
+++ b/tensor2tensor/data_generators/wikisum/utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/utils_test.py b/tensor2tensor/data_generators/wikisum/utils_test.py
index 9fb87f441..59de1a192 100644
--- a/tensor2tensor/data_generators/wikisum/utils_test.py
+++ b/tensor2tensor/data_generators/wikisum/utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/validate_data.py b/tensor2tensor/data_generators/wikisum/validate_data.py
index 6e4e8f0b7..41eb83713 100644
--- a/tensor2tensor/data_generators/wikisum/validate_data.py
+++ b/tensor2tensor/data_generators/wikisum/validate_data.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/wikisum.py b/tensor2tensor/data_generators/wikisum/wikisum.py
index 018a9dc75..293f45444 100644
--- a/tensor2tensor/data_generators/wikisum/wikisum.py
+++ b/tensor2tensor/data_generators/wikisum/wikisum.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikitext103.py b/tensor2tensor/data_generators/wikitext103.py
index 338f25912..977eff511 100644
--- a/tensor2tensor/data_generators/wikitext103.py
+++ b/tensor2tensor/data_generators/wikitext103.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wnli.py b/tensor2tensor/data_generators/wnli.py
index 741704732..42e246633 100644
--- a/tensor2tensor/data_generators/wnli.py
+++ b/tensor2tensor/data_generators/wnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wsj_parsing.py b/tensor2tensor/data_generators/wsj_parsing.py
index 02d1f64de..2a0a86650 100644
--- a/tensor2tensor/data_generators/wsj_parsing.py
+++ b/tensor2tensor/data_generators/wsj_parsing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/yelp_full.py b/tensor2tensor/data_generators/yelp_full.py
index 11d526ac5..7aa130090 100644
--- a/tensor2tensor/data_generators/yelp_full.py
+++ b/tensor2tensor/data_generators/yelp_full.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/yelp_polarity.py b/tensor2tensor/data_generators/yelp_polarity.py
index 1111e1294..5f16a16c0 100644
--- a/tensor2tensor/data_generators/yelp_polarity.py
+++ b/tensor2tensor/data_generators/yelp_polarity.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/__init__.py b/tensor2tensor/envs/__init__.py
index 01ea5565b..858a8238f 100644
--- a/tensor2tensor/envs/__init__.py
+++ b/tensor2tensor/envs/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 0984c4314..049a5c3c7 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 43859f08d..9a67a7128 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index c94e1623e..d9051356e 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_env_problem.py b/tensor2tensor/envs/gym_env_problem.py
index cf7e02dc9..e4714a1ce 100644
--- a/tensor2tensor/envs/gym_env_problem.py
+++ b/tensor2tensor/envs/gym_env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_env_problem_test.py b/tensor2tensor/envs/gym_env_problem_test.py
index c62705a2e..9d5e8f693 100644
--- a/tensor2tensor/envs/gym_env_problem_test.py
+++ b/tensor2tensor/envs/gym_env_problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_spaces_utils.py b/tensor2tensor/envs/gym_spaces_utils.py
index 30dc239e8..dbff5f646 100644
--- a/tensor2tensor/envs/gym_spaces_utils.py
+++ b/tensor2tensor/envs/gym_spaces_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_spaces_utils_test.py b/tensor2tensor/envs/gym_spaces_utils_test.py
index 1b637e185..2998d3e70 100644
--- a/tensor2tensor/envs/gym_spaces_utils_test.py
+++ b/tensor2tensor/envs/gym_spaces_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/mujoco_problems.py b/tensor2tensor/envs/mujoco_problems.py
index b58174f9c..0a8ddef97 100644
--- a/tensor2tensor/envs/mujoco_problems.py
+++ b/tensor2tensor/envs/mujoco_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/mujoco_problems_test.py b/tensor2tensor/envs/mujoco_problems_test.py
index d2ddc4d1a..813810260 100644
--- a/tensor2tensor/envs/mujoco_problems_test.py
+++ b/tensor2tensor/envs/mujoco_problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
index 7eb25c7d8..6ddb3caee 100644
--- a/tensor2tensor/envs/rendered_env_problem.py
+++ b/tensor2tensor/envs/rendered_env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/rendered_env_problem_test.py b/tensor2tensor/envs/rendered_env_problem_test.py
index bc1c1e745..852fc5b6d 100644
--- a/tensor2tensor/envs/rendered_env_problem_test.py
+++ b/tensor2tensor/envs/rendered_env_problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env.py b/tensor2tensor/envs/tic_tac_toe_env.py
index 44fa7f29c..479717283 100644
--- a/tensor2tensor/envs/tic_tac_toe_env.py
+++ b/tensor2tensor/envs/tic_tac_toe_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_problem.py b/tensor2tensor/envs/tic_tac_toe_env_problem.py
index a92c2c7a1..d9aa5236e 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_problem_test.py b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
index 70cc91a7c..2fce45651 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_test.py b/tensor2tensor/envs/tic_tac_toe_env_test.py
index 289916210..854e039d5 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/time_step.py b/tensor2tensor/envs/time_step.py
index b75603ef2..165e05d04 100644
--- a/tensor2tensor/envs/time_step.py
+++ b/tensor2tensor/envs/time_step.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/time_step_test.py b/tensor2tensor/envs/time_step_test.py
index 80dfcdeec..54042d7a2 100644
--- a/tensor2tensor/envs/time_step_test.py
+++ b/tensor2tensor/envs/time_step_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index 05a51b715..d56ab47e6 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index 148d29de4..dd2ee6d21 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/__init__.py b/tensor2tensor/insights/__init__.py
index 798e88593..7b0a2368a 100644
--- a/tensor2tensor/insights/__init__.py
+++ b/tensor2tensor/insights/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/graph.py b/tensor2tensor/insights/graph.py
index 6e13c1eb6..28bb562e0 100644
--- a/tensor2tensor/insights/graph.py
+++ b/tensor2tensor/insights/graph.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/polymer/common-types.js b/tensor2tensor/insights/polymer/common-types.js
index a206c3b9b..9abdfa9af 100644
--- a/tensor2tensor/insights/polymer/common-types.js
+++ b/tensor2tensor/insights/polymer/common-types.js
@@ -17,6 +17,7 @@
 
 /**
  * @fileoverview A set of shared types that will be replaced by js proto types.
+ * @externs
  */
 
 /**
diff --git a/tensor2tensor/insights/query_processor.py b/tensor2tensor/insights/query_processor.py
index 47aefdace..74b716e01 100644
--- a/tensor2tensor/insights/query_processor.py
+++ b/tensor2tensor/insights/query_processor.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/server.py b/tensor2tensor/insights/server.py
index 44aadef58..07f7b1dd8 100644
--- a/tensor2tensor/insights/server.py
+++ b/tensor2tensor/insights/server.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/transformer_model.py b/tensor2tensor/insights/transformer_model.py
index 68f86af13..267a0b299 100644
--- a/tensor2tensor/insights/transformer_model.py
+++ b/tensor2tensor/insights/transformer_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/__init__.py b/tensor2tensor/layers/__init__.py
index 798e88593..7b0a2368a 100644
--- a/tensor2tensor/layers/__init__.py
+++ b/tensor2tensor/layers/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/area_attention.py b/tensor2tensor/layers/area_attention.py
index 70b92d0fa..c4c8fc430 100644
--- a/tensor2tensor/layers/area_attention.py
+++ b/tensor2tensor/layers/area_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/area_attention_test.py b/tensor2tensor/layers/area_attention_test.py
index e8ffcc9ef..78005b8ab 100644
--- a/tensor2tensor/layers/area_attention_test.py
+++ b/tensor2tensor/layers/area_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 6026fb76c..027449854 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index 79aa383ca..ebc286c6d 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_audio.py b/tensor2tensor/layers/common_audio.py
index c4efd96c8..a4a2f493c 100644
--- a/tensor2tensor/layers/common_audio.py
+++ b/tensor2tensor/layers/common_audio.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 3049d7fa2..057016dab 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index f93cce9ea..9aaabe480 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index 80dda5c8d..6b3d89bd9 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 722c29a0a..f6a2105fe 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index 27129dddc..59c1b3a19 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index db772ee8b..b5119700f 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index e04a7305e..7a39b3584 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 73796250e..e8e1be203 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index aa70fe620..c17f23bc4 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 94c005c3a..3c8020306 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
index 0f7efa145..956e080da 100644
--- a/tensor2tensor/layers/latent_layers_test.py
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/message_passing_attention.py b/tensor2tensor/layers/message_passing_attention.py
index 6b444d0e3..9f9b5f14f 100644
--- a/tensor2tensor/layers/message_passing_attention.py
+++ b/tensor2tensor/layers/message_passing_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index 2fdd24eb1..c22fb698c 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index de597e1d3..1d934ca4f 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/ngram.py b/tensor2tensor/layers/ngram.py
index ad0bd7ec3..46086aa13 100644
--- a/tensor2tensor/layers/ngram.py
+++ b/tensor2tensor/layers/ngram.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/ngram_test.py b/tensor2tensor/layers/ngram_test.py
index 608884556..1493b90c1 100644
--- a/tensor2tensor/layers/ngram_test.py
+++ b/tensor2tensor/layers/ngram_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_glow_layers.py b/tensor2tensor/layers/transformer_glow_layers.py
index 6cfece84e..1687fc998 100644
--- a/tensor2tensor/layers/transformer_glow_layers.py
+++ b/tensor2tensor/layers/transformer_glow_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_glow_layers_ops.py b/tensor2tensor/layers/transformer_glow_layers_ops.py
index ee1acd850..64a24eab7 100644
--- a/tensor2tensor/layers/transformer_glow_layers_ops.py
+++ b/tensor2tensor/layers/transformer_glow_layers_ops.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_glow_layers_ops_test.py b/tensor2tensor/layers/transformer_glow_layers_ops_test.py
index fbf862899..9e01c2c67 100644
--- a/tensor2tensor/layers/transformer_glow_layers_ops_test.py
+++ b/tensor2tensor/layers/transformer_glow_layers_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_glow_layers_test.py b/tensor2tensor/layers/transformer_glow_layers_test.py
index 6b20a21b8..6c922e718 100644
--- a/tensor2tensor/layers/transformer_glow_layers_test.py
+++ b/tensor2tensor/layers/transformer_glow_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 826220180..fec79a166 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index 85ab70dbb..aa0e8d908 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_memory_test.py b/tensor2tensor/layers/transformer_memory_test.py
index f990863a0..fa39591b4 100644
--- a/tensor2tensor/layers/transformer_memory_test.py
+++ b/tensor2tensor/layers/transformer_memory_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index be455994e..a37c21801 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/vqa_layers.py b/tensor2tensor/layers/vqa_layers.py
index 4e59e7c90..6e802d662 100644
--- a/tensor2tensor/layers/vqa_layers.py
+++ b/tensor2tensor/layers/vqa_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/__init__.py b/tensor2tensor/metrics/__init__.py
index 798e88593..7b0a2368a 100644
--- a/tensor2tensor/metrics/__init__.py
+++ b/tensor2tensor/metrics/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/video_conditional_fvd.py b/tensor2tensor/metrics/video_conditional_fvd.py
index 240d87b2a..3b6a0f058 100644
--- a/tensor2tensor/metrics/video_conditional_fvd.py
+++ b/tensor2tensor/metrics/video_conditional_fvd.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/video_conditional_fvd_test.py b/tensor2tensor/metrics/video_conditional_fvd_test.py
index 67b40df5a..41bb5ce3c 100644
--- a/tensor2tensor/metrics/video_conditional_fvd_test.py
+++ b/tensor2tensor/metrics/video_conditional_fvd_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index c6c692bd8..579d10c25 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/basic.py b/tensor2tensor/models/basic.py
index d79008ee2..3aaf8ee0c 100644
--- a/tensor2tensor/models/basic.py
+++ b/tensor2tensor/models/basic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/basic_test.py b/tensor2tensor/models/basic_test.py
index a59880d13..556e8acfa 100644
--- a/tensor2tensor/models/basic_test.py
+++ b/tensor2tensor/models/basic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/bytenet.py b/tensor2tensor/models/bytenet.py
index ae3bef291..6c2312b03 100644
--- a/tensor2tensor/models/bytenet.py
+++ b/tensor2tensor/models/bytenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/bytenet_test.py b/tensor2tensor/models/bytenet_test.py
index c43b5e73d..c8f07b7e0 100644
--- a/tensor2tensor/models/bytenet_test.py
+++ b/tensor2tensor/models/bytenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/distillation.py b/tensor2tensor/models/distillation.py
index 2e0e3fb56..62fadd9f6 100644
--- a/tensor2tensor/models/distillation.py
+++ b/tensor2tensor/models/distillation.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 9174a9bbb..dd2053ba0 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/evolved_transformer_test.py b/tensor2tensor/models/evolved_transformer_test.py
index 364280deb..0d71b2f96 100644
--- a/tensor2tensor/models/evolved_transformer_test.py
+++ b/tensor2tensor/models/evolved_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index fd1affec5..67505ca41 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 443ed6d51..c27bd83e9 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_2d_test.py b/tensor2tensor/models/image_transformer_2d_test.py
index 9c7b43725..aa75032f9 100644
--- a/tensor2tensor/models/image_transformer_2d_test.py
+++ b/tensor2tensor/models/image_transformer_2d_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_test.py b/tensor2tensor/models/image_transformer_test.py
index eb88edca9..0c555f1c2 100644
--- a/tensor2tensor/models/image_transformer_test.py
+++ b/tensor2tensor/models/image_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 156c903f4..f6d408ae3 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/lstm_test.py b/tensor2tensor/models/lstm_test.py
index b49c7465a..a54018233 100644
--- a/tensor2tensor/models/lstm_test.py
+++ b/tensor2tensor/models/lstm_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_image_transformer.py b/tensor2tensor/models/mtf_image_transformer.py
index ecd8152cf..aabce6a47 100644
--- a/tensor2tensor/models/mtf_image_transformer.py
+++ b/tensor2tensor/models/mtf_image_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_image_transformer_test.py b/tensor2tensor/models/mtf_image_transformer_test.py
index 58711c2fd..aed3f0954 100644
--- a/tensor2tensor/models/mtf_image_transformer_test.py
+++ b/tensor2tensor/models/mtf_image_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_resnet.py b/tensor2tensor/models/mtf_resnet.py
index 492cf4b28..449629b7c 100644
--- a/tensor2tensor/models/mtf_resnet.py
+++ b/tensor2tensor/models/mtf_resnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 7dd32d36e..044170ef9 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index ae075601f..c94f2bed4 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer_test.py b/tensor2tensor/models/mtf_transformer_test.py
index f097353f5..36c469b9f 100644
--- a/tensor2tensor/models/mtf_transformer_test.py
+++ b/tensor2tensor/models/mtf_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/__init__.py b/tensor2tensor/models/neural_architecture_search/__init__.py
index 7688494ed..4a53b5e30 100644
--- a/tensor2tensor/models/neural_architecture_search/__init__.py
+++ b/tensor2tensor/models/neural_architecture_search/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_layers.py b/tensor2tensor/models/neural_architecture_search/nas_layers.py
index fcdfca5f4..4e0afc3c2 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_layers.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_layers_test.py b/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
index ae936c021..34f109c03 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_model.py b/tensor2tensor/models/neural_architecture_search/nas_model.py
index c8bcb23fe..40c114ef7 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_model_test.py b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
index da89868ab..d1bca5891 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model_test.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_assistant.py b/tensor2tensor/models/neural_assistant.py
index 98f20d88a..a2f02f675 100644
--- a/tensor2tensor/models/neural_assistant.py
+++ b/tensor2tensor/models/neural_assistant.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_gpu.py b/tensor2tensor/models/neural_gpu.py
index 48bff0067..bd5522ad4 100644
--- a/tensor2tensor/models/neural_gpu.py
+++ b/tensor2tensor/models/neural_gpu.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_gpu_test.py b/tensor2tensor/models/neural_gpu_test.py
index c650d1c08..17dc8ebfa 100644
--- a/tensor2tensor/models/neural_gpu_test.py
+++ b/tensor2tensor/models/neural_gpu_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/__init__.py b/tensor2tensor/models/research/__init__.py
index 798e88593..7b0a2368a 100644
--- a/tensor2tensor/models/research/__init__.py
+++ b/tensor2tensor/models/research/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/adafactor_experiments.py b/tensor2tensor/models/research/adafactor_experiments.py
index d1de970dc..6875f45e7 100644
--- a/tensor2tensor/models/research/adafactor_experiments.py
+++ b/tensor2tensor/models/research/adafactor_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/aligned.py b/tensor2tensor/models/research/aligned.py
index 1d7ac1413..470e5b119 100644
--- a/tensor2tensor/models/research/aligned.py
+++ b/tensor2tensor/models/research/aligned.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/attention_lm.py b/tensor2tensor/models/research/attention_lm.py
index cf234ce43..6a442f1f1 100644
--- a/tensor2tensor/models/research/attention_lm.py
+++ b/tensor2tensor/models/research/attention_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/attention_lm_moe.py b/tensor2tensor/models/research/attention_lm_moe.py
index 5e2149bae..581fdb407 100644
--- a/tensor2tensor/models/research/attention_lm_moe.py
+++ b/tensor2tensor/models/research/attention_lm_moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 1233f10ec..52a368f73 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index 07133da8f..0dc6f30c2 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index ad0aaac8a..6107a373c 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/gene_expression.py b/tensor2tensor/models/research/gene_expression.py
index bc66d241e..04d365e10 100644
--- a/tensor2tensor/models/research/gene_expression.py
+++ b/tensor2tensor/models/research/gene_expression.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/gene_expression_test.py b/tensor2tensor/models/research/gene_expression_test.py
index 4c36af768..99f7e35d6 100644
--- a/tensor2tensor/models/research/gene_expression_test.py
+++ b/tensor2tensor/models/research/gene_expression_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 7799624c8..1647e520c 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_init_hook.py b/tensor2tensor/models/research/glow_init_hook.py
index dc87ea94c..824ef826b 100644
--- a/tensor2tensor/models/research/glow_init_hook.py
+++ b/tensor2tensor/models/research/glow_init_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 54578643a..6c4ed96f4 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 9ff4239a8..5020ac430 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
index 2063be0cb..1263075f9 100644
--- a/tensor2tensor/models/research/glow_test.py
+++ b/tensor2tensor/models/research/glow_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/lm_experiments.py b/tensor2tensor/models/research/lm_experiments.py
index 9a465fae6..bd29e8b16 100644
--- a/tensor2tensor/models/research/lm_experiments.py
+++ b/tensor2tensor/models/research/lm_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/moe.py b/tensor2tensor/models/research/moe.py
index 9b4d448cb..3a32429d4 100644
--- a/tensor2tensor/models/research/moe.py
+++ b/tensor2tensor/models/research/moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/moe_experiments.py b/tensor2tensor/models/research/moe_experiments.py
index 503c23527..879e80bd2 100644
--- a/tensor2tensor/models/research/moe_experiments.py
+++ b/tensor2tensor/models/research/moe_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/multiquery_paper.py b/tensor2tensor/models/research/multiquery_paper.py
index 7c54a108a..9cc7eb6e5 100644
--- a/tensor2tensor/models/research/multiquery_paper.py
+++ b/tensor2tensor/models/research/multiquery_paper.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/neural_stack.py b/tensor2tensor/models/research/neural_stack.py
index 5f0fe6790..7857a5c77 100644
--- a/tensor2tensor/models/research/neural_stack.py
+++ b/tensor2tensor/models/research/neural_stack.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/neural_stack_test.py b/tensor2tensor/models/research/neural_stack_test.py
index 1365282a7..a5497d52e 100644
--- a/tensor2tensor/models/research/neural_stack_test.py
+++ b/tensor2tensor/models/research/neural_stack_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/residual_shuffle_exchange.py b/tensor2tensor/models/research/residual_shuffle_exchange.py
index c4107d680..5c70b7b50 100644
--- a/tensor2tensor/models/research/residual_shuffle_exchange.py
+++ b/tensor2tensor/models/research/residual_shuffle_exchange.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 7c0471ec5..8964f1be5 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/shuffle_network.py b/tensor2tensor/models/research/shuffle_network.py
index 7da420bdc..39f817fa1 100644
--- a/tensor2tensor/models/research/shuffle_network.py
+++ b/tensor2tensor/models/research/shuffle_network.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/similarity_transformer.py b/tensor2tensor/models/research/similarity_transformer.py
index 597a1ebb6..5da061d3e 100644
--- a/tensor2tensor/models/research/similarity_transformer.py
+++ b/tensor2tensor/models/research/similarity_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/super_lm.py b/tensor2tensor/models/research/super_lm.py
index b3307306b..05f7bfe21 100644
--- a/tensor2tensor/models/research/super_lm.py
+++ b/tensor2tensor/models/research/super_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_aux.py b/tensor2tensor/models/research/transformer_aux.py
index 1dc1e3137..bd1ca52a6 100644
--- a/tensor2tensor/models/research/transformer_aux.py
+++ b/tensor2tensor/models/research/transformer_aux.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_aux_test.py b/tensor2tensor/models/research/transformer_aux_test.py
index a172fba91..ee94837ad 100644
--- a/tensor2tensor/models/research/transformer_aux_test.py
+++ b/tensor2tensor/models/research/transformer_aux_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index 57999eca3..993aa08d8 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 638335681..c521947ea 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_parallel.py b/tensor2tensor/models/research/transformer_parallel.py
index ebbe72295..5adcac39a 100644
--- a/tensor2tensor/models/research/transformer_parallel.py
+++ b/tensor2tensor/models/research/transformer_parallel.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_revnet.py b/tensor2tensor/models/research/transformer_revnet.py
index ee0531b8b..dd792bd0c 100644
--- a/tensor2tensor/models/research/transformer_revnet.py
+++ b/tensor2tensor/models/research/transformer_revnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_revnet_test.py b/tensor2tensor/models/research/transformer_revnet_test.py
index 8466e6d68..09d9b4c4c 100644
--- a/tensor2tensor/models/research/transformer_revnet_test.py
+++ b/tensor2tensor/models/research/transformer_revnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_seq2edits.py b/tensor2tensor/models/research/transformer_seq2edits.py
index 4276b7838..81e9a4c48 100644
--- a/tensor2tensor/models/research/transformer_seq2edits.py
+++ b/tensor2tensor/models/research/transformer_seq2edits.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_sketch.py b/tensor2tensor/models/research/transformer_sketch.py
index 040683401..fa0e83173 100644
--- a/tensor2tensor/models/research/transformer_sketch.py
+++ b/tensor2tensor/models/research/transformer_sketch.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_symshard.py b/tensor2tensor/models/research/transformer_symshard.py
index e408e1f26..caed36a0b 100644
--- a/tensor2tensor/models/research/transformer_symshard.py
+++ b/tensor2tensor/models/research/transformer_symshard.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 244dcbfb4..11e52d106 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae_flow_prior.py b/tensor2tensor/models/research/transformer_vae_flow_prior.py
index e6147ff64..1de0dd4fc 100644
--- a/tensor2tensor/models/research/transformer_vae_flow_prior.py
+++ b/tensor2tensor/models/research/transformer_vae_flow_prior.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py b/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
index d23a50a57..04deb42ca 100644
--- a/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
+++ b/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae_test.py b/tensor2tensor/models/research/transformer_vae_test.py
index bd57d4e5e..0ab3d4e54 100644
--- a/tensor2tensor/models/research/transformer_vae_test.py
+++ b/tensor2tensor/models/research/transformer_vae_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index ae8ad3f29..bbc1d7277 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
index 1b79a370a..bc1856d11 100644
--- a/tensor2tensor/models/research/universal_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index b8534f02f..c3fc4669b 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index bb961d24f..cacbd1623 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_attention_test.py b/tensor2tensor/models/research/vqa_attention_test.py
index 101bf9891..d70a7d78c 100644
--- a/tensor2tensor/models/research/vqa_attention_test.py
+++ b/tensor2tensor/models/research/vqa_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_recurrent_self_attention.py b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
index 8c2a4b515..7857789b8 100644
--- a/tensor2tensor/models/research/vqa_recurrent_self_attention.py
+++ b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
index b39b88df5..bbb9d7250 100644
--- a/tensor2tensor/models/research/vqa_self_attention.py
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index da38fb313..02174b79a 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/resnet_test.py b/tensor2tensor/models/resnet_test.py
index 7c6e4e0d9..796458bdb 100644
--- a/tensor2tensor/models/resnet_test.py
+++ b/tensor2tensor/models/resnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/revnet.py b/tensor2tensor/models/revnet.py
index d1072ea68..8c923f9fe 100644
--- a/tensor2tensor/models/revnet.py
+++ b/tensor2tensor/models/revnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/revnet_test.py b/tensor2tensor/models/revnet_test.py
index ace5a5cfc..026406f15 100644
--- a/tensor2tensor/models/revnet_test.py
+++ b/tensor2tensor/models/revnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index d35a6af94..423c7d76c 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/slicenet.py b/tensor2tensor/models/slicenet.py
index 4cc795eb8..3d4930d84 100644
--- a/tensor2tensor/models/slicenet.py
+++ b/tensor2tensor/models/slicenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/slicenet_test.py b/tensor2tensor/models/slicenet_test.py
index 4de4be0bf..5a522c88b 100644
--- a/tensor2tensor/models/slicenet_test.py
+++ b/tensor2tensor/models/slicenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/text_cnn.py b/tensor2tensor/models/text_cnn.py
index bc8b7b504..5e299557d 100644
--- a/tensor2tensor/models/text_cnn.py
+++ b/tensor2tensor/models/text_cnn.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 7ce5f31e2..77321d213 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 1d6a344be..720295904 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index 59c5f0ba0..576b707bf 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/__init__.py b/tensor2tensor/models/video/__init__.py
index 798e88593..7b0a2368a 100644
--- a/tensor2tensor/models/video/__init__.py
+++ b/tensor2tensor/models/video/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index f6565bd6a..93d7f3dcf 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/base_vae.py b/tensor2tensor/models/video/base_vae.py
index 345810823..5922ce63b 100644
--- a/tensor2tensor/models/video/base_vae.py
+++ b/tensor2tensor/models/video/base_vae.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index ef2b8d8ca..603ebc254 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index de5893da0..03048882d 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic_test.py b/tensor2tensor/models/video/basic_deterministic_test.py
index 90df4cc81..6b2794cc0 100644
--- a/tensor2tensor/models/video/basic_deterministic_test.py
+++ b/tensor2tensor/models/video/basic_deterministic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_recurrent.py b/tensor2tensor/models/video/basic_recurrent.py
index 8b982811e..b2bcf5d98 100644
--- a/tensor2tensor/models/video/basic_recurrent.py
+++ b/tensor2tensor/models/video/basic_recurrent.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_recurrent_test.py b/tensor2tensor/models/video/basic_recurrent_test.py
index 56aeea4de..4af4fa672 100644
--- a/tensor2tensor/models/video/basic_recurrent_test.py
+++ b/tensor2tensor/models/video/basic_recurrent_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 0ebf67ae7..14510950f 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_stochastic_test.py b/tensor2tensor/models/video/basic_stochastic_test.py
index dff04fb1f..e0dbba8f4 100644
--- a/tensor2tensor/models/video/basic_stochastic_test.py
+++ b/tensor2tensor/models/video/basic_stochastic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index f157b15d3..49d7889cb 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/emily_test.py b/tensor2tensor/models/video/emily_test.py
index c9660e1b2..269dae573 100644
--- a/tensor2tensor/models/video/emily_test.py
+++ b/tensor2tensor/models/video/emily_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index 1234ab7db..72198cbb8 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
index 7dd9e6957..4adc7f93d 100644
--- a/tensor2tensor/models/video/epva_params.py
+++ b/tensor2tensor/models/video/epva_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index f95872df3..8d2b1890f 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv3d_test.py b/tensor2tensor/models/video/nfg_conv3d_test.py
index 0230d454b..8fa3f7f09 100644
--- a/tensor2tensor/models/video/nfg_conv3d_test.py
+++ b/tensor2tensor/models/video/nfg_conv3d_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv_lstm_test.py b/tensor2tensor/models/video/nfg_conv_lstm_test.py
index 5abda7fe8..76fde1aed 100644
--- a/tensor2tensor/models/video/nfg_conv_lstm_test.py
+++ b/tensor2tensor/models/video/nfg_conv_lstm_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv_test.py b/tensor2tensor/models/video/nfg_conv_test.py
index a6bab8baa..af87739b2 100644
--- a/tensor2tensor/models/video/nfg_conv_test.py
+++ b/tensor2tensor/models/video/nfg_conv_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_interpolate.py b/tensor2tensor/models/video/nfg_interpolate.py
index 3c3bb7d3f..c0498ca96 100644
--- a/tensor2tensor/models/video/nfg_interpolate.py
+++ b/tensor2tensor/models/video/nfg_interpolate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_test_utils.py b/tensor2tensor/models/video/nfg_test_utils.py
index 09a2057af..23baa7285 100644
--- a/tensor2tensor/models/video/nfg_test_utils.py
+++ b/tensor2tensor/models/video/nfg_test_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_uncond_test.py b/tensor2tensor/models/video/nfg_uncond_test.py
index b0cb855cc..5472b5454 100644
--- a/tensor2tensor/models/video/nfg_uncond_test.py
+++ b/tensor2tensor/models/video/nfg_uncond_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index 87b71f2aa..6c52944b4 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index 2c884a6ea..bf249b2df 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp_test.py b/tensor2tensor/models/video/savp_test.py
index 4294c0f8a..5baacaea8 100644
--- a/tensor2tensor/models/video/savp_test.py
+++ b/tensor2tensor/models/video/savp_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index 403d7c936..b26c8f403 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 5c0405992..732d25da4 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p_test.py b/tensor2tensor/models/video/sv2p_test.py
index a43cf3715..18b9835b3 100644
--- a/tensor2tensor/models/video/sv2p_test.py
+++ b/tensor2tensor/models/video/sv2p_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/tests_utils.py b/tensor2tensor/models/video/tests_utils.py
index ff834522a..2028f280a 100644
--- a/tensor2tensor/models/video/tests_utils.py
+++ b/tensor2tensor/models/video/tests_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/xception.py b/tensor2tensor/models/xception.py
index 99828ca7c..da9b05ad3 100644
--- a/tensor2tensor/models/xception.py
+++ b/tensor2tensor/models/xception.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/xception_test.py b/tensor2tensor/models/xception_test.py
index 4eb0f8ba0..bc345e41f 100644
--- a/tensor2tensor/models/xception_test.py
+++ b/tensor2tensor/models/xception_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/problems.py b/tensor2tensor/problems.py
index e83b2a1a4..8a5ec1d42 100644
--- a/tensor2tensor/problems.py
+++ b/tensor2tensor/problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/problems_colab.py b/tensor2tensor/problems_colab.py
index 4ab203c67..1b4f3e2f5 100644
--- a/tensor2tensor/problems_colab.py
+++ b/tensor2tensor/problems_colab.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/problems_test.py b/tensor2tensor/problems_test.py
index d9b6143f7..f05538953 100644
--- a/tensor2tensor/problems_test.py
+++ b/tensor2tensor/problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/__init__.py b/tensor2tensor/rl/__init__.py
index 798e88593..7b0a2368a 100644
--- a/tensor2tensor/rl/__init__.py
+++ b/tensor2tensor/rl/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/batch_dqn_agent_test.py b/tensor2tensor/rl/batch_dqn_agent_test.py
index 9790e31a3..4a3fae838 100644
--- a/tensor2tensor/rl/batch_dqn_agent_test.py
+++ b/tensor2tensor/rl/batch_dqn_agent_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/batch_runner_test.py b/tensor2tensor/rl/batch_runner_test.py
index cfb415ff1..efc9a5f4c 100644
--- a/tensor2tensor/rl/batch_runner_test.py
+++ b/tensor2tensor/rl/batch_runner_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/datagen_with_agent.py b/tensor2tensor/rl/datagen_with_agent.py
index 8199a1486..e41889f9f 100644
--- a/tensor2tensor/rl/datagen_with_agent.py
+++ b/tensor2tensor/rl/datagen_with_agent.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 3b9ac95b3..10976d41b 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/__init__.py b/tensor2tensor/rl/envs/__init__.py
index 798e88593..7b0a2368a 100644
--- a/tensor2tensor/rl/envs/__init__.py
+++ b/tensor2tensor/rl/envs/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index 39807fd82..19db0274e 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index 8533b8545..c662d9519 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 9688bad94..c1215ba2d 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index 3a6ac4d90..16de501fe 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index b79933b3a..72f120932 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 1f16c3699..ea38366e5 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/evaluator_test.py b/tensor2tensor/rl/evaluator_test.py
index b6ea5a0db..c49999811 100644
--- a/tensor2tensor/rl/evaluator_test.py
+++ b/tensor2tensor/rl/evaluator_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index 1d134f10b..81b9df7ed 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index 6074bcaa8..a5066f3db 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 8ce9b49f9..ba8c1db2e 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index 0be5e19d0..1648f6fc4 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
index 4bf3fb015..f17f9429c 100644
--- a/tensor2tensor/rl/policy_learner.py
+++ b/tensor2tensor/rl/policy_learner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index e0f0bed50..bec485515 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 6efd381ea..bd2642176 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/restarter.py b/tensor2tensor/rl/restarter.py
index 9fad2536c..223f7fe9b 100644
--- a/tensor2tensor/rl/restarter.py
+++ b/tensor2tensor/rl/restarter.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/restarter_test.py b/tensor2tensor/rl/restarter_test.py
index a447bf6ce..71e1a292c 100644
--- a/tensor2tensor/rl/restarter_test.py
+++ b/tensor2tensor/rl/restarter_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index b34ba9ba1..eaabb7b7f 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index f82defbfc..b07ea6835 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_agent_only.py b/tensor2tensor/rl/trainer_model_based_agent_only.py
index 55ca3dbc6..85c2ccdc7 100644
--- a/tensor2tensor/rl/trainer_model_based_agent_only.py
+++ b/tensor2tensor/rl/trainer_model_based_agent_only.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 04a1309e3..f314db913 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_recurrent_test.py b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
index 67844c610..1ace96b42 100644
--- a/tensor2tensor/rl/trainer_model_based_recurrent_test.py
+++ b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_stochastic_test.py b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
index 2e0159796..355175c74 100644
--- a/tensor2tensor/rl/trainer_model_based_stochastic_test.py
+++ b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_sv2p_test.py b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
index 148c47983..14cb10d1e 100644
--- a/tensor2tensor/rl/trainer_model_based_sv2p_test.py
+++ b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_test.py b/tensor2tensor/rl/trainer_model_based_test.py
index 68dfd1ac9..3eb71091b 100644
--- a/tensor2tensor/rl/trainer_model_based_test.py
+++ b/tensor2tensor/rl/trainer_model_based_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 8956666a4..964cbaa4f 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index f654a1f63..8cb73f173 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free_tictactoe_test.py b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
index 245336495..2a8bae6c1 100644
--- a/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
+++ b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/__init__.py b/tensor2tensor/serving/__init__.py
index 798e88593..7b0a2368a 100644
--- a/tensor2tensor/serving/__init__.py
+++ b/tensor2tensor/serving/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 0cb7da38f..8aa0b3d10 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/query.py b/tensor2tensor/serving/query.py
index 69a7aadff..e5d557eb1 100644
--- a/tensor2tensor/serving/query.py
+++ b/tensor2tensor/serving/query.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index a1b437282..3dc3d2fee 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/test_data/example_usr_dir/__init__.py b/tensor2tensor/test_data/example_usr_dir/__init__.py
index b47898717..aa146a2b5 100644
--- a/tensor2tensor/test_data/example_usr_dir/__init__.py
+++ b/tensor2tensor/test_data/example_usr_dir/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/test_data/example_usr_dir/my_submodule.py b/tensor2tensor/test_data/example_usr_dir/my_submodule.py
index 593c9dc38..88f8730e3 100644
--- a/tensor2tensor/test_data/example_usr_dir/my_submodule.py
+++ b/tensor2tensor/test_data/example_usr_dir/my_submodule.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/__init__.py b/tensor2tensor/utils/__init__.py
index 798e88593..7b0a2368a 100644
--- a/tensor2tensor/utils/__init__.py
+++ b/tensor2tensor/utils/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index 0d6f41176..7fa139c5b 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/adv_attack_utils.py b/tensor2tensor/utils/adv_attack_utils.py
index c0aee2b8d..b179812f1 100644
--- a/tensor2tensor/utils/adv_attack_utils.py
+++ b/tensor2tensor/utils/adv_attack_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index c2c134670..ae65c143a 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index 61a2aa1f7..d4a11594e 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/beam_search_test.py b/tensor2tensor/utils/beam_search_test.py
index 4aa0f3cb5..b1cc0ec2d 100644
--- a/tensor2tensor/utils/beam_search_test.py
+++ b/tensor2tensor/utils/beam_search_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/bleu_hook.py b/tensor2tensor/utils/bleu_hook.py
index d840eb232..a36f5e027 100644
--- a/tensor2tensor/utils/bleu_hook.py
+++ b/tensor2tensor/utils/bleu_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/bleu_hook_test.py b/tensor2tensor/utils/bleu_hook_test.py
index 173f6e0a5..477df16dc 100644
--- a/tensor2tensor/utils/bleu_hook_test.py
+++ b/tensor2tensor/utils/bleu_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/checkpoint_compatibility_test.py b/tensor2tensor/utils/checkpoint_compatibility_test.py
index cd466a641..b7a4a7bfe 100644
--- a/tensor2tensor/utils/checkpoint_compatibility_test.py
+++ b/tensor2tensor/utils/checkpoint_compatibility_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 678955df5..40165b736 100644
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/compute_video_metrics.py b/tensor2tensor/utils/compute_video_metrics.py
index 9a025d5e5..5dfa2ba2a 100644
--- a/tensor2tensor/utils/compute_video_metrics.py
+++ b/tensor2tensor/utils/compute_video_metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/contrib.py b/tensor2tensor/utils/contrib.py
index b99875d7c..f5b7883bb 100644
--- a/tensor2tensor/utils/contrib.py
+++ b/tensor2tensor/utils/contrib.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index bcf1a19e1..47891ec8b 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/data_reader_test.py b/tensor2tensor/utils/data_reader_test.py
index 4b2a1e797..2a4b2c75a 100644
--- a/tensor2tensor/utils/data_reader_test.py
+++ b/tensor2tensor/utils/data_reader_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 99ffdb0fb..a84749413 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index 80cacd20e..f03ef7ada 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/diet.py b/tensor2tensor/utils/diet.py
index ad15583e6..018063f9f 100644
--- a/tensor2tensor/utils/diet.py
+++ b/tensor2tensor/utils/diet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/diet_test.py b/tensor2tensor/utils/diet_test.py
index c3592d9aa..fe808d53b 100644
--- a/tensor2tensor/utils/diet_test.py
+++ b/tensor2tensor/utils/diet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index c0be890da..26282a92c 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/expert_utils_test.py b/tensor2tensor/utils/expert_utils_test.py
index 8ea49450c..deb8581c3 100644
--- a/tensor2tensor/utils/expert_utils_test.py
+++ b/tensor2tensor/utils/expert_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index 9e3dc5ab8..44e43fa62 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/get_rouge.py b/tensor2tensor/utils/get_rouge.py
index c753c0abd..d896e21af 100644
--- a/tensor2tensor/utils/get_rouge.py
+++ b/tensor2tensor/utils/get_rouge.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparam.py b/tensor2tensor/utils/hparam.py
index 1b073348f..6e09ed3e5 100644
--- a/tensor2tensor/utils/hparam.py
+++ b/tensor2tensor/utils/hparam.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparam_test.py b/tensor2tensor/utils/hparam_test.py
index 8d91861fb..3c31efdf6 100644
--- a/tensor2tensor/utils/hparam_test.py
+++ b/tensor2tensor/utils/hparam_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparams_lib.py b/tensor2tensor/utils/hparams_lib.py
index 2994add29..868f6081f 100644
--- a/tensor2tensor/utils/hparams_lib.py
+++ b/tensor2tensor/utils/hparams_lib.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparams_lib_test.py b/tensor2tensor/utils/hparams_lib_test.py
index 5941dd52c..a4d1f394d 100644
--- a/tensor2tensor/utils/hparams_lib_test.py
+++ b/tensor2tensor/utils/hparams_lib_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index b2a87e8ce..362356b69 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index be166a3d8..d47cbf08e 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_hook.py b/tensor2tensor/utils/metrics_hook.py
index 72d298fad..3b738b82d 100644
--- a/tensor2tensor/utils/metrics_hook.py
+++ b/tensor2tensor/utils/metrics_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_hook_test.py b/tensor2tensor/utils/metrics_hook_test.py
index ae9d41b20..97efab7bf 100644
--- a/tensor2tensor/utils/metrics_hook_test.py
+++ b/tensor2tensor/utils/metrics_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 55e8cee46..6ffafd8b1 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/misc_utils.py b/tensor2tensor/utils/misc_utils.py
index bb5d7a2d7..d26a586e8 100644
--- a/tensor2tensor/utils/misc_utils.py
+++ b/tensor2tensor/utils/misc_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/misc_utils_test.py b/tensor2tensor/utils/misc_utils_test.py
index ab7003f2e..b05a822bc 100644
--- a/tensor2tensor/utils/misc_utils_test.py
+++ b/tensor2tensor/utils/misc_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mlperf_log.py b/tensor2tensor/utils/mlperf_log.py
index 84a67bc7a..c8c08f075 100644
--- a/tensor2tensor/utils/mlperf_log.py
+++ b/tensor2tensor/utils/mlperf_log.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mlperf_tags.py b/tensor2tensor/utils/mlperf_tags.py
index ce153228a..bf57c8201 100644
--- a/tensor2tensor/utils/mlperf_tags.py
+++ b/tensor2tensor/utils/mlperf_tags.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mtf_model.py b/tensor2tensor/utils/mtf_model.py
index 4825c4116..7a2f2412d 100644
--- a/tensor2tensor/utils/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
index 37ad1137b..43172d0dc 100644
--- a/tensor2tensor/utils/multistep_optimizer.py
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_optimizer_test.py b/tensor2tensor/utils/multistep_optimizer_test.py
index 6a1531a41..cf5995473 100644
--- a/tensor2tensor/utils/multistep_optimizer_test.py
+++ b/tensor2tensor/utils/multistep_optimizer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_with_adamoptimizer.py b/tensor2tensor/utils/multistep_with_adamoptimizer.py
index 956a3503d..792028e99 100644
--- a/tensor2tensor/utils/multistep_with_adamoptimizer.py
+++ b/tensor2tensor/utils/multistep_with_adamoptimizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_with_adamoptimizer_test.py b/tensor2tensor/utils/multistep_with_adamoptimizer_test.py
index a7b538cd4..6ac29809a 100644
--- a/tensor2tensor/utils/multistep_with_adamoptimizer_test.py
+++ b/tensor2tensor/utils/multistep_with_adamoptimizer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 07f3ff5d8..078de3e03 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/optimize_test.py b/tensor2tensor/utils/optimize_test.py
index 326a45f51..ab4233a5b 100644
--- a/tensor2tensor/utils/optimize_test.py
+++ b/tensor2tensor/utils/optimize_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/partial_checkpoint_load_hook.py b/tensor2tensor/utils/partial_checkpoint_load_hook.py
index 2ff9e9da0..5f436fdf4 100644
--- a/tensor2tensor/utils/partial_checkpoint_load_hook.py
+++ b/tensor2tensor/utils/partial_checkpoint_load_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/pruning_utils.py b/tensor2tensor/utils/pruning_utils.py
index 4e6e5ab3e..476c41b56 100644
--- a/tensor2tensor/utils/pruning_utils.py
+++ b/tensor2tensor/utils/pruning_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/quantization.py b/tensor2tensor/utils/quantization.py
index a8011b1ed..849d8b5a3 100644
--- a/tensor2tensor/utils/quantization.py
+++ b/tensor2tensor/utils/quantization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 911ace7d9..fd6ec5ec5 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index 632f7c35f..5b8acab2f 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/restore_hook.py b/tensor2tensor/utils/restore_hook.py
index 2346da188..f4c5d86f7 100644
--- a/tensor2tensor/utils/restore_hook.py
+++ b/tensor2tensor/utils/restore_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/rouge.py b/tensor2tensor/utils/rouge.py
index fb5be8ec8..706154e2c 100644
--- a/tensor2tensor/utils/rouge.py
+++ b/tensor2tensor/utils/rouge.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/rouge_test.py b/tensor2tensor/utils/rouge_test.py
index d91dce4c8..e13872523 100644
--- a/tensor2tensor/utils/rouge_test.py
+++ b/tensor2tensor/utils/rouge_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/sari_hook.py b/tensor2tensor/utils/sari_hook.py
index 9cdacb561..e10ac7fb9 100644
--- a/tensor2tensor/utils/sari_hook.py
+++ b/tensor2tensor/utils/sari_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/sari_hook_test.py b/tensor2tensor/utils/sari_hook_test.py
index 1a9c88ce3..a063167e2 100644
--- a/tensor2tensor/utils/sari_hook_test.py
+++ b/tensor2tensor/utils/sari_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/scheduled_sampling.py b/tensor2tensor/utils/scheduled_sampling.py
index 339047f88..09a8089a7 100644
--- a/tensor2tensor/utils/scheduled_sampling.py
+++ b/tensor2tensor/utils/scheduled_sampling.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 397ed57da..b1f097f67 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index ab9760abb..a6a9d7008 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/test_utils.py b/tensor2tensor/utils/test_utils.py
index 5f9cbcd14..5f9215609 100644
--- a/tensor2tensor/utils/test_utils.py
+++ b/tensor2tensor/utils/test_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/test_utils_test.py b/tensor2tensor/utils/test_utils_test.py
index d671b5114..2de94b877 100644
--- a/tensor2tensor/utils/test_utils_test.py
+++ b/tensor2tensor/utils/test_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index ae8cae5ef..a3662e35e 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index 6784894c5..8f7e69af2 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/update_ops_hook.py b/tensor2tensor/utils/update_ops_hook.py
index f03ee29c9..5136858ff 100644
--- a/tensor2tensor/utils/update_ops_hook.py
+++ b/tensor2tensor/utils/update_ops_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/usr_dir.py b/tensor2tensor/utils/usr_dir.py
index aa7775c99..bc23ee1d6 100644
--- a/tensor2tensor/utils/usr_dir.py
+++ b/tensor2tensor/utils/usr_dir.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video/prediction2gif.py b/tensor2tensor/utils/video/prediction2gif.py
index 45483c0a0..db3b4e37d 100644
--- a/tensor2tensor/utils/video/prediction2gif.py
+++ b/tensor2tensor/utils/video/prediction2gif.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video/reward_confusion.py b/tensor2tensor/utils/video/reward_confusion.py
index 431717f3e..41ba42d33 100644
--- a/tensor2tensor/utils/video/reward_confusion.py
+++ b/tensor2tensor/utils/video/reward_confusion.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video2gif.py b/tensor2tensor/utils/video2gif.py
index 819c5aa39..13f7fdd43 100644
--- a/tensor2tensor/utils/video2gif.py
+++ b/tensor2tensor/utils/video2gif.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index 1f77c1a46..c1b3c0929 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video_metrics_test.py b/tensor2tensor/utils/video_metrics_test.py
index 476ebd1f9..455135d27 100644
--- a/tensor2tensor/utils/video_metrics_test.py
+++ b/tensor2tensor/utils/video_metrics_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/yellowfin.py b/tensor2tensor/utils/yellowfin.py
index aa383d323..a7ec6bcb7 100644
--- a/tensor2tensor/utils/yellowfin.py
+++ b/tensor2tensor/utils/yellowfin.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/yellowfin_test.py b/tensor2tensor/utils/yellowfin_test.py
index 0c77e60b2..c06f4563b 100644
--- a/tensor2tensor/utils/yellowfin_test.py
+++ b/tensor2tensor/utils/yellowfin_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/__init__.py b/tensor2tensor/visualization/__init__.py
index 7688494ed..4a53b5e30 100644
--- a/tensor2tensor/visualization/__init__.py
+++ b/tensor2tensor/visualization/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/attention.py b/tensor2tensor/visualization/attention.py
index 2a1ebf071..498967326 100644
--- a/tensor2tensor/visualization/attention.py
+++ b/tensor2tensor/visualization/attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/visualization.py b/tensor2tensor/visualization/visualization.py
index 7ebe48297..514e61bd7 100644
--- a/tensor2tensor/visualization/visualization.py
+++ b/tensor2tensor/visualization/visualization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/visualization_test.py b/tensor2tensor/visualization/visualization_test.py
index e054c8ebd..4770c1eb1 100644
--- a/tensor2tensor/visualization/visualization_test.py
+++ b/tensor2tensor/visualization/visualization_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2020 The Tensor2Tensor Authors.
+# Copyright 2021 The Tensor2Tensor Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

From e19130b7dd6a4122ce65b29c64bdac5f0036effb Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 25 Feb 2021 11:35:22 -0800
Subject: [PATCH 2688/2720] Don't overspecify WeightNorm input_spec, match
 input_spec of wrapped

PiperOrigin-RevId: 359573317
---
 tensor2tensor/layers/common_layers.py      |  4 +---
 tensor2tensor/layers/common_layers_test.py | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index f6a2105fe..af1fcb83b 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -4043,9 +4043,6 @@ def _data_dep_init(self, inputs):
 
   def build(self, input_shape=None):
     """Build `Layer`."""
-    input_shape = tf.TensorShape(input_shape).as_list()
-    self.input_spec = layers().InputSpec(shape=input_shape)
-
     if not self.layer.built:
       self.layer.build(input_shape)
       self.layer.built = False
@@ -4072,6 +4069,7 @@ def build(self, input_shape=None):
       self._compute_weights()
 
       self.layer.built = True
+    self.input_spec = self.layer.input_spec
 
     super(WeightNorm, self).build()
     self.built = True
diff --git a/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index 59c1b3a19..f649b24a5 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -965,5 +965,20 @@ def fn_recompute(x):
         self.assertAllClose(g1, g2)
 
 
+class WeightNormTest(tf.test.TestCase):
+
+  def testInputSpec(self):
+    """Test that WeighNorm does not overspecify the input_spec."""
+    conv = common_layers.WeightNorm(
+        tf.keras.layers.Conv1D(filters=8, kernel_size=3))
+    # Call with one batch size:
+    conv(tf.zeros([1, 16, 2]))
+    # Should allow call with another batch size.
+    conv(tf.zeros([2, 16, 2]))
+    # Input spec does detect incorrect input feature dim.
+    with self.assertRaises(ValueError):
+      conv(tf.zeros([2, 16, 3]))
+
+
 if __name__ == "__main__":
   tf.test.main()

From 5623deb79cfcd28f8f8c5463b58b5bd76a81fd0d Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 5 Mar 2021 00:14:45 -0800
Subject: [PATCH 2689/2720] [Mesh-TF] Add is_training as an arg to mtf.dropout

PiperOrigin-RevId: 361088273
---
 tensor2tensor/models/mtf_image_transformer.py | 14 +++++++++++++-
 tensor2tensor/models/mtf_transformer.py       | 18 +++++++++++++-----
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/tensor2tensor/models/mtf_image_transformer.py b/tensor2tensor/models/mtf_image_transformer.py
index aabce6a47..fdcee7a23 100644
--- a/tensor2tensor/models/mtf_image_transformer.py
+++ b/tensor2tensor/models/mtf_image_transformer.py
@@ -243,8 +243,10 @@ def import_to_batch_by_length(x, name):
 def layer_prepostprocess_dropout(x, hparams):
   batch_dim = x.shape.dims[0]
   model_dim = x.shape.dims[-1]
+  mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN)
+  is_training = mode == tf.estimator.ModeKeys.TRAIN
   return mtf.dropout(
-      x,
+      x, is_training,
       keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
       noise_shape=mtf.Shape([batch_dim, model_dim]))
 
@@ -259,6 +261,8 @@ def local_attention1d_spatial_decoder(x, kv_dim, heads_dim,
   x = mtf.reshape(
       x, mtf.Shape([batch_dim, num_w_blocks_dim, blocks_w_dim, model_dim]))
   # [ self attention - ffn - residual + dropout] x n
+  mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN)
+  is_training = mode == tf.estimator.ModeKeys.TRAIN
   for layer in range(hparams.num_decoder_layers):
     layer_name = "decoder_layer_%d" % layer
     with tf.variable_scope(layer_name):
@@ -268,6 +272,7 @@ def local_attention1d_spatial_decoder(x, kv_dim, heads_dim,
               mtf.layers.layer_norm(x, model_dim, name="layer_norm_att"),
               kv_dim,
               heads_dim,
+              is_training,
               memory_w_dim=blocks_w_dim,
               mask_right=True,
               name="self_att"), hparams)
@@ -276,6 +281,7 @@ def local_attention1d_spatial_decoder(x, kv_dim, heads_dim,
           mtf.layers.dense_relu_dense(
               mtf.layers.layer_norm(x, model_dim, name="layer_norm_ffn"),
               feedforward_dim,
+              is_training,
               hparams.dropout,
               dropout_broadcast_dims=[length_dim]), hparams)
 
@@ -305,6 +311,8 @@ def local_attention2d_spatial_decoder(x, kv_dim, heads_dim,
           batch_dim, num_h_blocks_dim, num_w_blocks_dim,
           blocks_h_dim, blocks_w_dim, model_dim
       ]))
+  mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN)
+  is_training = mode == tf.estimator.ModeKeys.TRAIN
   # Image Transformer Decoder
   # [ self attention - ffn - residual + dropout] x n
   for layer in range(hparams.num_decoder_layers):
@@ -316,6 +324,7 @@ def local_attention2d_spatial_decoder(x, kv_dim, heads_dim,
               mtf.layers.layer_norm(x, model_dim, name="layer_norm_att"),
               kv_dim,
               heads_dim,
+              is_training,
               memory_h_dim=num_h_blocks_dim,
               memory_w_dim=num_w_blocks_dim,
               name="self_att"), hparams)
@@ -336,6 +345,8 @@ def local_attention1d_masked_decoder(x, kv_dim, heads_dim,
   """Image Transformer decoder with local1D masked layers."""
   print(x)
   _, length_dim, model_dim = x.shape.dims
+  mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN)
+  is_training = mode == tf.estimator.ModeKeys.TRAIN
   for layer in range(hparams.num_decoder_layers):
     layer_name = "decoder_layer_%d" % layer
     with tf.variable_scope(layer_name):
@@ -347,6 +358,7 @@ def local_attention1d_masked_decoder(x, kv_dim, heads_dim,
               mtf.layers.layer_norm(x, model_dim, name="layer_norm_att"),
               kv_dim,
               heads_dim,
+              is_training,
               window_size=hparams.block_length,
               length_per_split=length_per_split,
               name="self_att"), hparams)
diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 044170ef9..5ac5e091a 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -242,6 +242,8 @@ def _mtf_model_fn(self, features, mesh):
     hparams = self._hparams
     extra_losses = []
     targets = tf.to_int32(features["targets"])
+    mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN)
+    is_training = mode == tf.estimator.ModeKeys.TRAIN
     if len(targets.get_shape()) > 2:
       tf.logging.info("targets = %s" % targets)
       targets = tf.squeeze(targets, [2, 3])
@@ -289,7 +291,7 @@ def pad_to_max_length(x):
 
     def layer_prepostprocess_dropout(x):
       return mtf.dropout(
-          x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
+          x, is_training, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
           noise_shape=mtf.Shape(self.batch_dims + [self.model_dim]))
 
     (inputs_embedding_var,
@@ -426,10 +428,11 @@ def _feedforward_layer(self, x, layer_type, losses=None):
       ValueError: if hparams make no sense
     """
     hparams = self._hparams
-
+    mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN)
+    is_training = mode == tf.estimator.ModeKeys.TRAIN
     if layer_type == "drd":
       return mtf.layers.dense_relu_dense(
-          x, self.feedforward_dim, dropout=hparams.relu_dropout,
+          x, self.feedforward_dim, is_training, dropout=hparams.relu_dropout,
           dropout_broadcast_dims=[self.length_dim],
           master_dtype=self.master_dtype,
           slice_dtype=self.slice_dtype)
@@ -493,11 +496,13 @@ def _layer_stack(self,
     """
     hparams = self._hparams
     is_incremental = (step_num is not None)
+    mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN)
+    is_training = mode == tf.estimator.ModeKeys.TRAIN
     def layer_prepostprocess_dropout(x):
       if is_incremental:
         return x
       return mtf.dropout(
-          x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
+          x, is_training, keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
           noise_shape=mtf.Shape(self.batch_dims + [self.model_dim]))
     num_layers = len(layers)
     num_layer_norms = num_layers + 1
@@ -540,6 +545,7 @@ def normalize(x):
                 mtf.layers.multihead_attention(
                     normalize(x), None,
                     self_attention_mask, self.kv_dim, self.heads_dim,
+                    is_training,
                     dropout=hparams.attention_dropout,
                     dropout_broadcast_dims=[self.length_dim],
                     master_dtype=self.master_dtype,
@@ -560,6 +566,7 @@ def normalize(x):
                 mtf.layers.multihead_attention(
                     normalize(x), encoder_output,
                     encdec_attention_mask, self.kv_dim, self.heads_dim,
+                    is_training,
                     dropout=hparams.attention_dropout,
                     dropout_broadcast_dims=[self.length_dim],
                     master_dtype=self.master_dtype,
@@ -582,7 +589,7 @@ def normalize(x):
             x += layer_prepostprocess_dropout(
                 mtf.layers.masked_local_attention_1d(
                     normalize(x),
-                    self.kv_dim, self.heads_dim,
+                    self.kv_dim, self.heads_dim, is_training,
                     window_size=hparams.local_attention_window_size,
                     master_dtype=self.master_dtype,
                     slice_dtype=self.slice_dtype,
@@ -601,6 +608,7 @@ def normalize(x):
                     compression_factor=hparams.compression_factor,
                     kv_channels=self.kv_dim,
                     heads=self.heads_dim,
+                    is_training=is_training,
                     dropout=hparams.attention_dropout,
                     dropout_broadcast_dims=[self.length_dim],
                     master_dtype=self.master_dtype,

From 51256a221aa87fcebf882938397a8fcc814b7134 Mon Sep 17 00:00:00 2001
From: Sebastian Goodman <seabass@google.com>
Date: Mon, 3 May 2021 14:49:50 -0700
Subject: [PATCH 2690/2720] A more generic packing op

PiperOrigin-RevId: 371783025
---
 .../data_generators/ops/pack_sequences_ops.cc | 402 +++++++++++++++++-
 .../ops/pack_sequences_ops_test.py            | 177 +++++++-
 2 files changed, 576 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
index 76ed97b4f..34f634c6a 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
@@ -1,22 +1,31 @@
+#include "base/integral_types.h"
 #include "third_party/tensorflow/core/framework/op_kernel.h"
 #include "third_party/tensorflow/core/framework/shape_inference.h"
 #include "third_party/tensorflow/core/framework/tensor.h"
 #include "third_party/tensorflow/core/framework/types.h"
+#include "third_party/tensorflow/core/framework/types.proto.h"
+#include "third_party/tensorflow/core/platform/errors.h"
 
 namespace tensor2tensor {
 namespace {
 
+using ::tensorflow::bfloat16;
+using ::tensorflow::DataTypeVector;
 using ::tensorflow::DEVICE_CPU;
+using ::tensorflow::OpInputList;
 using ::tensorflow::OpKernel;
 using ::tensorflow::OpKernelConstruction;
 using ::tensorflow::OpKernelContext;
+using ::tensorflow::OpOutputList;
 using ::tensorflow::Status;
 using ::tensorflow::Tensor;
 using ::tensorflow::TensorShape;
+using ::tensorflow::TTypes;
+using ::tensorflow::errors::InvalidArgument;
+using ::tensorflow::shape_inference::DimensionHandle;
 using ::tensorflow::shape_inference::InferenceContext;
+using ::tensorflow::shape_inference::ShapeHandle;
 
-// TODO(noam): this op packs a dataset of pairs of sequences (inputs, targets)
-// Generalize later to an arbitrary number of sequences.
 REGISTER_OP("PackSequences2")
     .Input("inputs: int64")
     .Input("targets: int64")
@@ -161,8 +170,397 @@ class PackSequences2Op : public OpKernel {
   }
 };
 
+REGISTER_OP("PackSequencesK")
+    .Input("inputs: Tinput_types")
+    .Input("max_lengths: Tinput_count * int32")
+    .Attr("Tinput_types: list(type)")
+    .Attr("Tinput_count: int")
+    .Output("outputs_packed: Tinput_types")
+    .Output("outputs_segmentation: Tinput_count * int32")
+    .Output("outputs_position: Tinput_count * int32")
+    .SetShapeFn([](InferenceContext* ctx) {
+      DataTypeVector input_types;
+      int input_count;
+      TF_RETURN_IF_ERROR(ctx->GetAttr("Tinput_types", &input_types));
+      TF_RETURN_IF_ERROR(ctx->GetAttr("Tinput_count", &input_count));
+      if (input_types.size() != input_count) {
+        return InvalidArgument(
+            "`inputs` and `max_lengths` had different numbers of elements");
+      }
+      auto inputs = ctx->input_handle_shapes_and_types(0);
+      std::vector<ShapeHandle> output_dims(inputs->size());
+      std::vector<ShapeHandle> segmentation_dims(inputs->size());
+      std::vector<ShapeHandle> position_dims(inputs->size());
+      for (int i = 0; i < inputs->size(); i++) {
+        auto input = inputs->at(i);
+        int rank = ctx->Rank(input.shape);
+        std::vector<DimensionHandle> dims(rank);
+        for (int r = 0; r < rank; r++) {
+          dims.push_back(ctx->UnknownDim());
+        }
+        output_dims.push_back(ctx->MakeShape(dims));
+        segmentation_dims.push_back(
+            ctx->Matrix(ctx->UnknownDim(), ctx->UnknownDim()));
+        position_dims.push_back(
+            ctx->Matrix(ctx->UnknownDim(), ctx->UnknownDim()));
+      }
+      TF_RETURN_IF_ERROR(ctx->set_output("outputs_packed", output_dims));
+      TF_RETURN_IF_ERROR(
+          ctx->set_output("outputs_segmentation", segmentation_dims));
+      TF_RETURN_IF_ERROR(ctx->set_output("outputs_position", position_dims));
+      return Status::OK();
+    });
+
+typedef int InputIndex;
+typedef int BatchIndex;
+typedef int SeqIndex;
+
+struct PackingSpec {
+  SeqIndex seq_id;
+  BatchIndex batch_pos;
+  int seq_length;
+  int offset;
+  int segment_id;
+};
+
+class PackSequencesKOp : public OpKernel {
+ public:
+  explicit PackSequencesKOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("Tinput_types", &input_types_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("Tinput_count", &input_count_));
+    OP_REQUIRES(
+        ctx, input_types_.size() == input_count_,
+        InvalidArgument(
+            "`inputs` and `max_lengths` had different numbers of elements"));
+  }
+
+  void Compute(OpKernelContext* ctx) override {
+    OpInputList inputs;
+    OpInputList max_lengths_list;
+
+    OP_REQUIRES_OK(ctx, ctx->input_list("inputs", &inputs));
+    OP_REQUIRES_OK(ctx, ctx->input_list("max_lengths", &max_lengths_list));
+    OP_REQUIRES(
+        ctx, inputs.size() == max_lengths_list.size(),
+        InvalidArgument(
+            "`inputs` and `max_lengths` had different numbers of elements"));
+
+    std::map<InputIndex, int> max_lengths;
+    for (InputIndex i = 0; i < max_lengths_list.size(); i++) {
+      max_lengths[i] = max_lengths_list[i].scalar<int32>()();
+    }
+
+    int n = inputs.begin()->dim_size(0);
+    for (const auto& input : inputs) {
+      OP_REQUIRES(ctx, input.dim_size(0) == n,
+                  InvalidArgument("`inputs` had different batch sizes"));
+    }
+
+    std::map<InputIndex, int> padded_inputs_lengths;
+    for (InputIndex i = 0; i < inputs.size(); i++) {
+      padded_inputs_lengths[i] =
+          std::min(static_cast<int>(inputs[i].dim_size(1)), max_lengths[i]);
+    }
+
+    std::map<InputIndex, std::vector<int>> inputs_lengths;
+    for (InputIndex i = 0; i < inputs.size(); i++) {
+      inputs_lengths[i] =
+          GetInputLengths(ctx, inputs[i], padded_inputs_lengths[i]);
+    }
+
+    int num_combined = 0;
+    std::map<InputIndex, std::map<BatchIndex, int>> combined_inputs_lengths;
+    std::map<InputIndex, std::map<SeqIndex, PackingSpec>> packing_specs;
+    std::map<BatchIndex, int> segment_counter;
+
+    for (SeqIndex seq_id = 0; seq_id < n; seq_id++) {
+      for (BatchIndex b = std::max(0, num_combined - 1000); b < n; b++) {
+        bool enough_room = true;
+        for (InputIndex i = 0; i < inputs.size(); i++) {
+          int cur_seq_len = combined_inputs_lengths[i][b];
+          if (cur_seq_len + inputs_lengths[i][seq_id] > max_lengths[i]) {
+            enough_room = false;
+            break;
+          }
+        }
+        if (enough_room) {
+          num_combined = std::max(num_combined, b + 1);
+          for (InputIndex i = 0; i < inputs.size(); i++) {
+            packing_specs[i][seq_id] = {
+              .seq_id = seq_id,
+              .batch_pos = b,
+              .seq_length = inputs_lengths[i][seq_id],
+              .offset = combined_inputs_lengths[i][b],
+              .segment_id = (segment_counter[b] + 1)  // Add 1 because zero=pad
+            };
+            combined_inputs_lengths[i][b] += inputs_lengths[i][seq_id];
+          }
+          segment_counter[b]++;
+          break;
+        }
+      }
+      for (InputIndex i = 0; i < inputs.size(); i++) {
+        if (packing_specs[i].find(seq_id) == packing_specs[i].end()) {
+          ctx->CtxFailure(InvalidArgument(tensorflow::strings::StrCat(
+              "failed to pack example=", seq_id, " into input=", i)));
+        }
+      }
+    }
+
+    OpOutputList outputs_packed;
+    OpOutputList outputs_segmentation;
+    OpOutputList outputs_position;
+
+    OP_REQUIRES_OK(
+        ctx, ctx->output_list("outputs_packed", &outputs_packed));
+    OP_REQUIRES_OK(
+        ctx, ctx->output_list("outputs_segmentation", &outputs_segmentation));
+    OP_REQUIRES_OK(
+        ctx, ctx->output_list("outputs_position", &outputs_position));
+
+    for (InputIndex i = 0; i < inputs.size(); i++) {
+      TensorShape output_shape_2d = {
+        static_cast<int64>(num_combined),
+        static_cast<int64>(max_lengths[i])};
+
+      TensorShape output_shape = output_shape_2d;
+      if (inputs[i].dims() == 3) {
+        output_shape.AddDim(inputs[i].dim_size(2));
+      } else if (inputs[i].dims() != 2) {
+        ctx->CtxFailure(InvalidArgument("invalid rank"));
+      }
+
+      Tensor* packed;
+      Tensor* segmentation;
+      Tensor* position;
+
+      OP_REQUIRES_OK(ctx, outputs_packed.allocate(i, output_shape, &packed));
+      OP_REQUIRES_OK(ctx, outputs_segmentation.allocate(i, output_shape_2d,
+                                                        &segmentation));
+      OP_REQUIRES_OK(ctx,
+                     outputs_position.allocate(i, output_shape_2d, &position));
+
+      auto segmentation_eigen = segmentation->matrix<int32>();
+      auto position_eigen = position->matrix<int32>();
+
+      SetZero(ctx, packed);
+      segmentation_eigen.setZero();
+      position_eigen.setZero();
+
+      for (const auto& pair : packing_specs.at(i)) {
+        PackSequence(ctx, inputs[i], packed, segmentation_eigen,
+                     position_eigen, pair.second);
+      }
+    }
+  }
+
+ private:
+  std::vector<int> GetInputLengths(
+      OpKernelContext* ctx,
+      const Tensor& input,
+      const int padded_input_length) {
+    switch (input.dtype()) {
+      case tensorflow::DT_BFLOAT16:
+        return GetInputLengths<bfloat16>(ctx, input, padded_input_length);
+      case tensorflow::DT_FLOAT:
+        return GetInputLengths<float>(ctx, input, padded_input_length);
+      case tensorflow::DT_INT32:
+        return GetInputLengths<int32>(ctx, input, padded_input_length);
+      case tensorflow::DT_INT64:
+        return GetInputLengths<int64>(ctx, input, padded_input_length);
+      default:
+        ctx->CtxFailure(
+            tensorflow::errors::InvalidArgument("unsupported input dtype"));
+        return {};
+    }
+  }
+
+  template <typename T>
+  std::vector<int> GetInputLengths(
+      OpKernelContext* ctx,
+      const Tensor& input,
+      const int padded_input_length) {
+    if (input.dims() == 2) {
+      return GetInputLengths<const T>(
+          input.tensor<T, 2>(), padded_input_length);
+    } else if (input.dims() == 3) {
+      return GetInputLengths<const T>(
+          input.tensor<T, 3>(), padded_input_length);
+    } else {
+      ctx->CtxFailure(
+          tensorflow::errors::InvalidArgument("unsupported input rank"));
+      return {};
+    }
+  }
+
+  template <typename T>
+  std::vector<int> GetInputLengths(
+      const typename TTypes<T, 2>::Tensor& input,
+      const int padded_input_length) {
+    std::vector<int> input_lengths;
+    for (int i = 0; i < input.dimension(0); i++) {
+      int input_length = 0;
+      for (int j = 0; j < padded_input_length; j++) {
+        if (input(i, j) != 0) {
+          input_length++;
+        }
+      }
+      input_lengths.push_back(input_length);
+    }
+    return input_lengths;
+  }
+
+  template <typename T>
+  std::vector<int> GetInputLengths(
+      const typename TTypes<T, 3>::Tensor& input,
+      const int padded_input_length) {
+    std::vector<int> input_lengths;
+    for (int i = 0; i < input.dimension(0); i++) {
+      int input_length = 0;
+      for (int j = 0; j < padded_input_length; j++) {
+        for (int k = 0; k < input.dimension(2); k++) {
+          if (input(i, j, k) != 0) {
+            input_length++;
+            break;
+          }
+        }
+      }
+      input_lengths.push_back(input_length);
+    }
+    return input_lengths;
+  }
+
+  void SetZero(OpKernelContext* ctx, Tensor* inputs) {
+    switch (inputs->dtype()) {
+      case tensorflow::DT_BFLOAT16:
+        SetZero<bfloat16>(ctx, inputs);
+        break;
+      case tensorflow::DT_FLOAT:
+        SetZero<float>(ctx, inputs);
+        break;
+      case tensorflow::DT_INT32:
+        SetZero<int32>(ctx, inputs);
+        break;
+      case tensorflow::DT_INT64:
+        SetZero<int64>(ctx, inputs);
+        break;
+      default:
+        ctx->CtxFailure(
+            tensorflow::errors::InvalidArgument("unsupported input dtype"));
+    }
+  }
+
+  template <typename T>
+  void SetZero(OpKernelContext* ctx, Tensor* inputs) {
+    switch (inputs->dims()) {
+      case 2:
+        inputs->tensor<T, 2>().setZero();
+        break;
+      case 3:
+        inputs->tensor<T, 3>().setZero();
+        break;
+      default:
+        ctx->CtxFailure(
+            tensorflow::errors::InvalidArgument("unsupported input rank"));
+    }
+  }
+
+  void PackSequence(OpKernelContext* ctx, const Tensor& inputs, Tensor* packed,
+                    TTypes<int32, 2>::Tensor segmentation,
+                    TTypes<int32, 2>::Tensor position,
+                    const PackingSpec& spec) {
+    switch (inputs.dtype()) {
+      case tensorflow::DT_FLOAT:
+        PackSequence<float>(
+            ctx, inputs, packed, segmentation, position, spec);
+        break;
+      case tensorflow::DT_BFLOAT16:
+        PackSequence<bfloat16>(
+            ctx, inputs, packed, segmentation, position, spec);
+        break;
+      case tensorflow::DT_INT32:
+        PackSequence<int32>(
+            ctx, inputs, packed, segmentation, position, spec);
+        break;
+      case tensorflow::DT_INT64:
+        PackSequence<int64>(
+            ctx, inputs, packed, segmentation, position, spec);
+        break;
+      default:
+        ctx->CtxFailure(
+            tensorflow::errors::InvalidArgument("unsupported input dtype"));
+    }
+  }
+
+  template <typename T>
+  void PackSequence(OpKernelContext* ctx, const Tensor& inputs, Tensor* packed,
+                    TTypes<int32, 2>::Tensor segmentation,
+                    TTypes<int32, 2>::Tensor position,
+                    const PackingSpec& spec) {
+    switch (inputs.dims()) {
+      case 2:
+        PackSequence<T>(
+            ctx,
+            inputs.tensor<T, 2>(),
+            packed->tensor<T, 2>(),  // TensorMap is pass-by-ref.
+            segmentation,
+            position,
+            spec);
+        break;
+      case 3:
+        PackSequence<T>(
+            ctx,
+            inputs.tensor<T, 3>(),
+            packed->tensor<T, 3>(),  // TensorMap is pass-by-ref.
+            segmentation,
+            position,
+            spec);
+        break;
+      default:
+        ctx->CtxFailure(
+            tensorflow::errors::InvalidArgument("unsupported input rank"));
+    }
+  }
+
+  template <typename T>
+  void PackSequence(OpKernelContext* ctx,
+                    const typename TTypes<const T, 2>::Tensor& inputs,
+                    typename TTypes<T, 2>::Tensor packed,
+                    TTypes<int32, 2>::Tensor segmentation,
+                    TTypes<int32, 2>::Tensor position,
+                    const PackingSpec& spec) {
+    for (int i = 0; i < spec.seq_length; i++) {
+      packed(spec.batch_pos, spec.offset + i) = inputs(spec.seq_id, i);
+      segmentation(spec.batch_pos, spec.offset + i) = spec.segment_id;
+      position(spec.batch_pos, spec.offset + i) = i;
+    }
+  }
+
+  template <typename T>
+  void PackSequence(OpKernelContext* ctx,
+                    const typename TTypes<const T, 3>::Tensor& inputs,
+                    typename TTypes<T, 3>::Tensor packed,
+                    TTypes<int32, 2>::Tensor segmentation,
+                    TTypes<int32, 2>::Tensor position,
+                    const PackingSpec& spec) {
+    for (int i = 0; i < spec.seq_length; i++) {
+      for (int k = 0; k < inputs.dimension(2); k++) {
+        packed(spec.batch_pos, spec.offset + i, k) = inputs(spec.seq_id, i, k);
+      }
+      segmentation(spec.batch_pos, spec.offset + i) = spec.segment_id;
+      position(spec.batch_pos, spec.offset + i) = i;
+    }
+  }
+
+  DataTypeVector input_types_;
+  int input_count_;
+};
+
 REGISTER_KERNEL_BUILDER(Name("PackSequences2").Device(DEVICE_CPU),
                         PackSequences2Op);
 
+REGISTER_KERNEL_BUILDER(Name("PackSequencesK").Device(DEVICE_CPU),
+                        PackSequencesKOp);
+
 }  // namespace
 }  // namespace tensor2tensor
diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
index 4f7c58772..3592683cf 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
@@ -25,7 +25,7 @@
 
 class PackSequencesOpsTest(tf.test.TestCase):
 
-  def test_pack_sequences(self):
+  def test_pack_sequences2(self):
     inputs = [
         [1, 2, 3],
         [4, 5, 0],
@@ -72,6 +72,181 @@ def test_pack_sequences(self):
             [0, 1, 0, 0, 0],
         ])
 
+  def test_pack_sequences_k(self):
+    inputs = tf.convert_to_tensor([
+        [1, 2, 3],
+        [4, 5, 0],
+        [6, 0, 0],
+    ], dtype=tf.int32)
+    targets = tf.convert_to_tensor([
+        [10, 0, 0],
+        [20, 30, 40],
+        [50, 60, 0],
+    ], dtype=tf.int32)
+    max_length = tf.convert_to_tensor(5, dtype=tf.int32)
+    (packed, segmentation, position) = pack_sequences_ops.pack_sequences_k(
+        [inputs, targets], [max_length, max_length])
+    (inputs_packed, targets_packed) = packed
+    (inputs_segmentation, targets_segmentation) = segmentation
+    (inputs_position, targets_position) = position
+    self.assertAllEqual(
+        inputs_packed, [
+            [1, 2, 3, 4, 5],
+            [6, 0, 0, 0, 0],
+        ])
+    self.assertAllEqual(
+        inputs_segmentation, [
+            [1, 1, 1, 2, 2],
+            [1, 0, 0, 0, 0],
+        ])
+    self.assertAllEqual(
+        inputs_position, [
+            [0, 1, 2, 0, 1],
+            [0, 0, 0, 0, 0],
+        ])
+    self.assertAllEqual(
+        targets_packed, [
+            [10, 20, 30, 40, 0],
+            [50, 60, 0, 0, 0],
+        ])
+    self.assertAllEqual(
+        targets_segmentation, [
+            [1, 2, 2, 2, 0],
+            [1, 1, 0, 0, 0],
+        ])
+    self.assertAllEqual(
+        targets_position, [
+            [0, 0, 1, 2, 0],
+            [0, 1, 0, 0, 0],
+        ])
+
+  def test_pack_sequences_k_multi_input(self):
+    input_tokens = tf.convert_to_tensor([
+        [1, 2, 3],
+        [4, 5, 0],
+        [6, 0, 0],
+    ], dtype=tf.int32)
+    input_vectors = tf.convert_to_tensor([
+        [[0, 1, 2], [1, 2, 3], [2, 3, 4]],
+        [[3, 4, 5], [4, 5, 6], [0, 0, 0]],
+        [[5, 6, 7], [0, 0, 0], [0, 0, 0]],
+    ], dtype=tf.float32)
+    targets = tf.convert_to_tensor([
+        [10, 0, 0],
+        [20, 30, 40],
+        [50, 60, 0],
+    ], dtype=tf.int32)
+    (packed, segmentation, position) = pack_sequences_ops.pack_sequences_k(
+        [input_tokens, input_vectors, targets],
+        [5, 3, 5])
+    (input_tokens_packed, input_vectors_packed, targets_packed) = packed
+    (input_tokens_segmentation, input_vectors_segmentation,
+     targets_segmentation) = segmentation
+    (input_tokens_position, input_vectors_position, targets_position) = position
+    self.assertAllEqual(
+        input_tokens_packed, [
+            [1, 2, 3, 0, 0],
+            [4, 5, 6, 0, 0],
+        ])
+    self.assertAllEqual(
+        input_vectors_packed, [
+            [[0, 1, 2], [1, 2, 3], [2, 3, 4]],
+            [[3, 4, 5], [4, 5, 6], [5, 6, 7]],
+        ])
+    self.assertAllEqual(
+        input_tokens_segmentation, [
+            [1, 1, 1, 0, 0],
+            [1, 1, 2, 0, 0],
+        ])
+    self.assertAllEqual(
+        input_vectors_segmentation, [
+            [1, 1, 1],
+            [1, 1, 2],
+        ])
+    self.assertAllEqual(
+        input_tokens_position, [
+            [0, 1, 2, 0, 0],
+            [0, 1, 0, 0, 0],
+        ])
+    self.assertAllEqual(
+        input_vectors_position, [
+            [0, 1, 2],
+            [0, 1, 0],
+        ])
+    self.assertAllEqual(
+        targets_packed, [
+            [10, 0, 0, 0, 0],
+            [20, 30, 40, 50, 60],
+        ])
+    self.assertAllEqual(
+        targets_segmentation, [
+            [1, 0, 0, 0, 0],
+            [1, 1, 1, 2, 2],
+        ])
+    self.assertAllEqual(
+        targets_position, [
+            [0, 0, 0, 0, 0],
+            [0, 1, 2, 0, 1],
+        ])
+
+  def test_pack_sequences_k_int64(self):
+    inputs = tf.convert_to_tensor([
+        [1, 2, 3],
+        [4, 5, 0],
+        [6, 0, 0],
+    ], dtype=tf.int64)
+    max_length = tf.convert_to_tensor(5, dtype=tf.int32)
+    (packed, segmentation, position) = pack_sequences_ops.pack_sequences_k(
+        [inputs], [max_length])
+    (inputs_packed,) = packed
+    (inputs_segmentation,) = segmentation
+    (inputs_position,) = position
+    self.assertAllEqual(
+        inputs_packed, [
+            [1, 2, 3, 4, 5],
+            [6, 0, 0, 0, 0],
+        ])
+    self.assertEqual(inputs_packed.dtype, tf.int64)
+    self.assertAllEqual(
+        inputs_segmentation, [
+            [1, 1, 1, 2, 2],
+            [1, 0, 0, 0, 0],
+        ])
+    self.assertAllEqual(
+        inputs_position, [
+            [0, 1, 2, 0, 1],
+            [0, 0, 0, 0, 0],
+        ])
+
+  def test_pack_sequences_k_bfloat16(self):
+    inputs = tf.convert_to_tensor([
+        [1, 2, 3],
+        [4, 5, 0],
+        [6, 0, 0],
+    ], dtype=tf.bfloat16)
+    max_length = tf.convert_to_tensor(5, dtype=tf.int32)
+    (packed, segmentation, position) = pack_sequences_ops.pack_sequences_k(
+        [inputs], [max_length])
+    (inputs_packed,) = packed
+    (inputs_segmentation,) = segmentation
+    (inputs_position,) = position
+    self.assertAllEqual(
+        inputs_packed, [
+            [1, 2, 3, 4, 5],
+            [6, 0, 0, 0, 0],
+        ])
+    self.assertEqual(inputs_packed.dtype, tf.bfloat16)
+    self.assertAllEqual(
+        inputs_segmentation, [
+            [1, 1, 1, 2, 2],
+            [1, 0, 0, 0, 0],
+        ])
+    self.assertAllEqual(
+        inputs_position, [
+            [0, 1, 2, 0, 1],
+            [0, 0, 0, 0, 0],
+        ])
+
 
 if __name__ == "__main__":
   tf.enable_eager_execution()

From 3f12173b19c1bad2a7c37eb390f3ad46baee0c19 Mon Sep 17 00:00:00 2001
From: Sebastian Goodman <seabass@google.com>
Date: Tue, 11 May 2021 09:43:18 -0700
Subject: [PATCH 2691/2720] A more generic packing op for seqio

PiperOrigin-RevId: 373172087
---
 .../data_generators/ops/pack_sequences_ops.cc | 41 +++++++++++--------
 1 file changed, 24 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
index 34f634c6a..02f49a562 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
@@ -187,27 +187,34 @@ REGISTER_OP("PackSequencesK")
         return InvalidArgument(
             "`inputs` and `max_lengths` had different numbers of elements");
       }
-      auto inputs = ctx->input_handle_shapes_and_types(0);
-      std::vector<ShapeHandle> output_dims(inputs->size());
-      std::vector<ShapeHandle> segmentation_dims(inputs->size());
-      std::vector<ShapeHandle> position_dims(inputs->size());
-      for (int i = 0; i < inputs->size(); i++) {
-        auto input = inputs->at(i);
-        int rank = ctx->Rank(input.shape);
-        std::vector<DimensionHandle> dims(rank);
-        for (int r = 0; r < rank; r++) {
-          dims.push_back(ctx->UnknownDim());
-        }
-        output_dims.push_back(ctx->MakeShape(dims));
-        segmentation_dims.push_back(
+      std::vector<ShapeHandle> input_shapes;
+      TF_RETURN_IF_ERROR(ctx->input("inputs", &input_shapes));
+      std::vector<ShapeHandle> output_shapes;
+      std::vector<ShapeHandle> segmentation_shapes;
+      std::vector<ShapeHandle> position_shapes;
+      for (int i = 0; i < input_shapes.size(); i++) {
+        const auto& input_shape = input_shapes.at(i);
+        int rank = ctx->Rank(input_shape);
+        segmentation_shapes.push_back(
             ctx->Matrix(ctx->UnknownDim(), ctx->UnknownDim()));
-        position_dims.push_back(
+        position_shapes.push_back(
             ctx->Matrix(ctx->UnknownDim(), ctx->UnknownDim()));
+        if (rank == 2) {
+          output_shapes.push_back(
+              ctx->MakeShape({ctx->UnknownDim(), ctx->UnknownDim()}));
+        } else if (rank == 3) {
+          output_shapes.push_back(
+              ctx->MakeShape({ctx->UnknownDim(), ctx->UnknownDim(),
+                              ctx->Value(ctx->Dim(input_shape, 2))}));
+        } else {
+          return InvalidArgument(
+              "Only rank 2 and rank 3 inputs are supported");
+        }
       }
-      TF_RETURN_IF_ERROR(ctx->set_output("outputs_packed", output_dims));
+      TF_RETURN_IF_ERROR(ctx->set_output("outputs_packed", output_shapes));
       TF_RETURN_IF_ERROR(
-          ctx->set_output("outputs_segmentation", segmentation_dims));
-      TF_RETURN_IF_ERROR(ctx->set_output("outputs_position", position_dims));
+          ctx->set_output("outputs_segmentation", segmentation_shapes));
+      TF_RETURN_IF_ERROR(ctx->set_output("outputs_position", position_shapes));
       return Status::OK();
     });
 

From ae042f66e013494eb2c4c2b50963da5a3d3fc828 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Mon, 21 Jun 2021 19:46:22 -0700
Subject: [PATCH 2692/2720] Internal

PiperOrigin-RevId: 380716688
---
 .../data_generators/text_problems.py          | 27 ++++++++++++++
 tensor2tensor/data_generators/translate.py    | 35 +++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index ab51c5918..ef4af2ce8 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -799,6 +799,32 @@ def text2real_txt_iterator(source_txt_path, target_txt_path):
     yield {"inputs": inputs, "targets": targets}
 
 
+def txt_line_sharded_iterator(txt_pattern):
+  """Iterate through lines of sharded file."""
+  all_files = tf.gfile.Glob(txt_pattern)
+  for txt_path in all_files:
+    with tf.gfile.Open(txt_path) as f:
+      for line in f:
+        yield line.strip()
+
+
+def text2text_txt_sharded_iterator(source_txt_pattern, target_txt_pattern):
+  """Yield dicts for Text2TextProblem.generate_samples from lines of files.
+
+  Args:
+    source_txt_pattern: path to the sharded source file
+    target_txt_pattern: path to the sharded target file
+
+  Yields:
+    {"inputs": inputs, "targets": targets}
+
+  """
+  for inputs, targets in zip(
+      txt_line_sharded_iterator(source_txt_pattern),
+      txt_line_sharded_iterator(target_txt_pattern)):
+    yield {"inputs": inputs, "targets": targets}
+
+
 def text2text_txt_tab_iterator(txt_path):
   """Yield dicts for Text2TextProblem.generate_samples from lines of txt_path.
 
@@ -848,6 +874,7 @@ class Text2textTmpdir(Text2TextProblem):
   TRAIN_FILES = ("inputs.train.txt", "targets.train.txt")
   EVAL_FILES = ("inputs.eval.txt", "targets.eval.txt")
 
+  @property
   def is_generate_per_split(self):
     return True
 
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 2827daf44..3c32512a3 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -266,6 +266,7 @@ def compile_data(tmp_dir, datasets, filename, datatypes_to_clean=None):
 class TranslateDistillProblem(TranslateProblem):
   """Base class for translation problems."""
 
+  @property
   def is_generate_per_split(self):
     return True
 
@@ -311,3 +312,37 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
     return text_problems.text2text_distill_iterator(data_path + "inputs",
                                                     data_path + "gold",
                                                     data_path + "prediction")
+
+
+class TranslateWmt20Problem(TranslateProblem):
+  """Base class for WMT20 Datasets."""
+
+  @property
+  def is_generate_per_split(self):
+    return True
+
+  def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
+    generator = self.generate_samples(data_dir, tmp_dir, dataset_split)
+    vocab = self.get_or_create_vocab(data_dir, tmp_dir)
+    # For each example, encode the text and append EOS ID.
+    for sample in generator:
+      if self.has_inputs:
+        sample["inputs"] = vocab.encode(sample["inputs"])
+        sample["inputs"].append(text_encoder.EOS_ID)
+        sample["targets"] = vocab.encode(sample["targets"])
+        sample["targets"].append(text_encoder.EOS_ID)
+        yield sample
+
+  def generate_text_for_vocab(self, data_dir, tmp_dir):
+    for i, sample in enumerate(
+        self.generate_samples(data_dir, tmp_dir, problem.DatasetSplit.TRAIN)):
+      if self.has_inputs:
+        yield sample["inputs"]
+      yield sample["targets"]
+      if self.max_samples_for_vocab and (i + 1) >= self.max_samples_for_vocab:
+        break
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    data_path = self.source_data_files(dataset_split)[0]
+    assert tf.gfile.Exists(data_path)
+    return text_problems.text2text_txt_tab_iterator(data_path)

From 874389bf5ce8b1c6e33063f887883603ad5c8fd2 Mon Sep 17 00:00:00 2001
From: Szymon Tworkowski <44476995+syzymon@users.noreply.github.com>
Date: Fri, 20 Aug 2021 21:27:31 +0200
Subject: [PATCH 2693/2720] Add enwik8 with different lengths + binary read
 (#1895)

---
 tensor2tensor/data_generators/enwik8.py | 75 ++++++++++++++++++++++++-
 1 file changed, 74 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/enwik8.py b/tensor2tensor/data_generators/enwik8.py
index 31ae913c0..837e981ad 100644
--- a/tensor2tensor/data_generators/enwik8.py
+++ b/tensor2tensor/data_generators/enwik8.py
@@ -55,6 +55,7 @@ def _maybe_download_corpus(tmp_dir):
 class Enwik8L65k(text_problems.Text2SelfProblem):
   """Enwiki8, with examples up to 65,536 characters long."""
 
+  READ_MODE = "r"
   DUPE_FACTOR = 4
 
   @property
@@ -92,7 +93,7 @@ def sequence_length(self):
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
     filepath = _maybe_download_corpus(tmp_dir)
-    with tf.io.gfile.GFile(filepath) as f:
+    with tf.io.gfile.GFile(filepath, mode=self.READ_MODE) as f:
       data = f.read()
 
     tf.logging.info("Length of enwik8 = %d", len(data))
@@ -126,3 +127,75 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
     for sample in generator:
       sample["targets"] = vocab.encode(sample["targets"])
       yield sample
+
+
+@registry.register_problem
+class Enwik8L2k(Enwik8L65k):
+  """Enwiki8, with examples up to 2048 characters long. Reads the input
+  byte-wise and chunks it into fragments of maximum length of 2048. Does not
+  shift byte indices (we do not assume cls or pad are used),
+  unlike the base class!"""
+
+  READ_MODE = "rb"
+
+  @property
+  def sequence_length(self):
+    """Length of each example (number of characters)."""
+    return 2048
+
+  def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
+    return self.generate_samples(data_dir, tmp_dir, dataset_split)
+
+
+@registry.register_problem
+class Enwik8L32k(Enwik8L2k):
+
+  @property
+  def sequence_length(self):
+    """Length of each example (in tokens)."""
+    return 32768
+
+
+@registry.register_problem
+class Enwik8L16k(Enwik8L2k):
+
+  @property
+  def sequence_length(self):
+    """Length of each example (in tokens)."""
+    return 16384
+
+
+@registry.register_problem
+class Enwik8L8k(Enwik8L2k):
+
+  @property
+  def sequence_length(self):
+    """Length of each example (in tokens)."""
+    return 8192
+
+
+@registry.register_problem
+class Enwik8L4k(Enwik8L2k):
+
+  @property
+  def sequence_length(self):
+    """Length of each example (in tokens)."""
+    return 4096
+
+
+@registry.register_problem
+class Enwik8L1k(Enwik8L2k):
+
+  @property
+  def sequence_length(self):
+    """Length of each example (in tokens)."""
+    return 1024
+
+
+@registry.register_problem
+class Enwik8L512(Enwik8L2k):
+
+  @property
+  def sequence_length(self):
+    """Length of each example (in tokens)."""
+    return 512

From 7ae6d2896b48969f006b3b51aa4aff8655d855de Mon Sep 17 00:00:00 2001
From: Szymon Tworkowski <44476995+syzymon@users.noreply.github.com>
Date: Fri, 20 Aug 2021 14:34:01 -0700
Subject: [PATCH 2694/2720] Merge of PR #1895

PiperOrigin-RevId: 392071163
---
 tensor2tensor/data_generators/enwik8.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/enwik8.py b/tensor2tensor/data_generators/enwik8.py
index 837e981ad..6f26b86e0 100644
--- a/tensor2tensor/data_generators/enwik8.py
+++ b/tensor2tensor/data_generators/enwik8.py
@@ -131,10 +131,12 @@ def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
 
 @registry.register_problem
 class Enwik8L2k(Enwik8L65k):
-  """Enwiki8, with examples up to 2048 characters long. Reads the input
-  byte-wise and chunks it into fragments of maximum length of 2048. Does not
-  shift byte indices (we do not assume cls or pad are used),
-  unlike the base class!"""
+  """Enwiki8, with examples up to 2048 characters long.
+
+  Reads the input byte-wise and chunks it into fragments of maximum
+  length of 2048. Does not shift byte indices (we do not assume cls or
+  pad are used), unlike the base class!
+  """
 
   READ_MODE = "rb"
 

From 5ca66de8e708abcf5681ef3a2e7c53f825e00e4c Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 20 Aug 2021 20:59:40 -0700
Subject: [PATCH 2695/2720] Internal

PiperOrigin-RevId: 392122923
---
 tensor2tensor/data_generators/text_problems.py |  6 +++++-
 tensor2tensor/data_generators/translate.py     | 10 +++++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index ef4af2ce8..dd0233ab9 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -835,7 +835,11 @@ def text2text_txt_tab_iterator(txt_path):
   Yields:
     {"inputs": inputs, "targets": targets}
   """
-  for line in txt_line_iterator(txt_path):
+  if txt_path.endswith(".tsv*"):
+    data_iterator = txt_line_sharded_iterator(txt_path)
+  else:
+    data_iterator = txt_line_iterator(txt_path)
+  for line in data_iterator:
     if line and "\t" in line:
       parts = line.split("\t", 1)
       inputs, targets = parts[:2]
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 3c32512a3..aae476ec0 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -344,5 +344,13 @@ def generate_text_for_vocab(self, data_dir, tmp_dir):
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
     data_path = self.source_data_files(dataset_split)[0]
-    assert tf.gfile.Exists(data_path)
     return text_problems.text2text_txt_tab_iterator(data_path)
+
+
+class TranslateSamanantarProblem(TranslateWmt20Problem):
+  """Base class for Samanantar Datasets."""
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    src_data_path = self.source_data_files(dataset_split)[0]
+    tgt_data_path = self.source_data_files(dataset_split)[1]
+    return text_problems.text2text_txt_iterator(src_data_path, tgt_data_path)

From 0dd16b9488aa0ae0bf873bd591d8f1b4b399f30f Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Mon, 30 Aug 2021 13:26:45 -0700
Subject: [PATCH 2696/2720] Delete .travis.yml (#1897)

---
 .travis.yml | 33 ---------------------------------
 1 file changed, 33 deletions(-)
 delete mode 100644 .travis.yml

diff --git a/.travis.yml b/.travis.yml
deleted file mode 100644
index 0bc70f988..000000000
--- a/.travis.yml
+++ /dev/null
@@ -1,33 +0,0 @@
-sudo: required
-language: python
-cache: pip
-git:
-  depth: 3
-  quiet: true
-services:
-  - docker
-python:
-  - "3.6"
-env:
-  global:
-    - T2T_PROBLEM=algorithmic_reverse_binary40_test
-    - T2T_DATA_DIR=/tmp/t2t-data
-    - T2T_TRAIN_DIR=/tmp/t2t-train
-    - TF_LATEST="1.15.*"
-    # This is necessary to have gsutil work with Python 2.7
-    - BOTO_CONFIG=/dev/null
-  matrix:
-    - TF_VERSION="1.15.*"
-install:
-  - ./oss_scripts/oss_pip_install.sh
-script:
-  - ./oss_scripts/oss_tests.sh
-  - ./oss_scripts/oss_integration_test.sh
-
-  # Conditional commands should each be in a separate block to get proper
-  # errors on Travis.
-  #
-  # TODO(afrozm): Re-enable if this becomes an issue.
-  # - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]]; then
-  #       pylint -j 2 tensor2tensor;
-  #   fi

From c22a226704e5887862bf9edd9f269892c9016ad4 Mon Sep 17 00:00:00 2001
From: Afroz Mohiuddin <afrozm@google.com>
Date: Fri, 17 Sep 2021 10:07:29 -0700
Subject: [PATCH 2697/2720] Fix glow invertibility test with latest lapack
 version.

PiperOrigin-RevId: 397344660
---
 .travis.yml                                   | 33 +++++++++++++++++++
 .../layers/transformer_glow_layers_test.py    |  4 +--
 2 files changed, 35 insertions(+), 2 deletions(-)
 create mode 100644 .travis.yml

diff --git a/.travis.yml b/.travis.yml
new file mode 100644
index 000000000..0bc70f988
--- /dev/null
+++ b/.travis.yml
@@ -0,0 +1,33 @@
+sudo: required
+language: python
+cache: pip
+git:
+  depth: 3
+  quiet: true
+services:
+  - docker
+python:
+  - "3.6"
+env:
+  global:
+    - T2T_PROBLEM=algorithmic_reverse_binary40_test
+    - T2T_DATA_DIR=/tmp/t2t-data
+    - T2T_TRAIN_DIR=/tmp/t2t-train
+    - TF_LATEST="1.15.*"
+    # This is necessary to have gsutil work with Python 2.7
+    - BOTO_CONFIG=/dev/null
+  matrix:
+    - TF_VERSION="1.15.*"
+install:
+  - ./oss_scripts/oss_pip_install.sh
+script:
+  - ./oss_scripts/oss_tests.sh
+  - ./oss_scripts/oss_integration_test.sh
+
+  # Conditional commands should each be in a separate block to get proper
+  # errors on Travis.
+  #
+  # TODO(afrozm): Re-enable if this becomes an issue.
+  # - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]]; then
+  #       pylint -j 2 tensor2tensor;
+  #   fi
diff --git a/tensor2tensor/layers/transformer_glow_layers_test.py b/tensor2tensor/layers/transformer_glow_layers_test.py
index 6c922e718..a65b77d89 100644
--- a/tensor2tensor/layers/transformer_glow_layers_test.py
+++ b/tensor2tensor/layers/transformer_glow_layers_test.py
@@ -245,8 +245,8 @@ def test_flow_invertibility(self):
             [x, x_mask, x_inv, x_inv_inv, logabsdet, logabsdet_inv]))
     diff = x - x_inv_inv
     logabsdet_sum = logabsdet + logabsdet_inv
-    self.assertTrue(np.allclose(diff, 0.0, atol=1e-5))
-    self.assertTrue(np.allclose(logabsdet_sum, 0.0, atol=1e-5))
+    self.assertTrue(np.allclose(diff, 0.0, atol=2e-5))
+    self.assertTrue(np.allclose(logabsdet_sum, 0.0, atol=2e-5))
 
   @parameterized.parameters(
       ("1", "cat", "affine"),

From 79b920fc83c996a0376125683424fad5a5d50232 Mon Sep 17 00:00:00 2001
From: Sebastian Goodman <seabass@google.com>
Date: Thu, 21 Oct 2021 07:27:52 -0700
Subject: [PATCH 2698/2720] Parity tests for `pack_sequences_2` and
 `pack_sequences_k`.

PiperOrigin-RevId: 404790223
---
 .../ops/pack_sequences_ops_test.py            | 319 +++++++++++++++---
 1 file changed, 273 insertions(+), 46 deletions(-)

diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
index 3592683cf..ef10510e0 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
@@ -19,13 +19,117 @@
 from __future__ import division
 from __future__ import print_function
 
+import numpy as np
 from tensor2tensor.data_generators.ops import pack_sequences_ops
 import tensorflow.compat.v1 as tf
 
 
+def _pack_sequences_k(inputs, targets, input_max_length, target_max_length):
+  """Wrapper for pack_sequences_k with same interface as pack_sequences_2."""
+  inputs = tf.convert_to_tensor(inputs, tf.int32)
+  targets = tf.convert_to_tensor(targets, tf.int32)
+  input_max_length = tf.convert_to_tensor(input_max_length, dtype=tf.int32)
+  target_max_length = tf.convert_to_tensor(target_max_length, dtype=tf.int32)
+  (packed, segmentation, position) = pack_sequences_ops.pack_sequences_k(
+      [inputs, targets], [input_max_length, target_max_length])
+  (inputs_packed, targets_packed) = packed
+  (inputs_segmentation, targets_segmentation) = segmentation
+  (inputs_position, targets_position) = position
+  return (inputs_packed, inputs_segmentation, inputs_position, targets_packed,
+          targets_segmentation, targets_position)
+
+
 class PackSequencesOpsTest(tf.test.TestCase):
 
-  def test_pack_sequences2(self):
+  def do_test_pack_sequences_length3(self, pack_fn):
+    inputs = [
+        [1, 2, 3],
+        [4, 5, 0],
+        [6, 0, 0],
+    ]
+    targets = [
+        [10, 0, 0],
+        [20, 30, 40],
+        [50, 60, 0],
+    ]
+    inputs_max_length = 3
+    targets_max_length = 3
+    (inputs_packed, inputs_segmentation, inputs_position, targets_packed,
+     targets_segmentation, targets_position) = (
+         pack_fn(inputs, targets, inputs_max_length, targets_max_length))
+    self.assertAllEqual(inputs_packed, [
+        [1, 2, 3],
+        [4, 5, 0],
+        [6, 0, 0],
+    ])
+    self.assertAllEqual(inputs_segmentation, [
+        [1, 1, 1],
+        [1, 1, 0],
+        [1, 0, 0],
+    ])
+    self.assertAllEqual(inputs_position, [
+        [0, 1, 2],
+        [0, 1, 0],
+        [0, 0, 0],
+    ])
+    self.assertAllEqual(targets_packed, [
+        [10, 0, 0],
+        [20, 30, 40],
+        [50, 60, 0],
+    ])
+    self.assertAllEqual(targets_segmentation, [
+        [1, 0, 0],
+        [1, 1, 1],
+        [1, 1, 0],
+    ])
+    self.assertAllEqual(targets_position, [
+        [0, 0, 0],
+        [0, 1, 2],
+        [0, 1, 0],
+    ])
+
+  def do_test_pack_sequences_length4(self, pack_fn):
+    inputs = [
+        [1, 2, 3],
+        [4, 5, 0],
+        [6, 0, 0],
+    ]
+    targets = [
+        [10, 0, 0],
+        [20, 30, 40],
+        [50, 60, 0],
+    ]
+    inputs_max_length = 4
+    targets_max_length = 4
+    (inputs_packed, inputs_segmentation, inputs_position, targets_packed,
+     targets_segmentation, targets_position) = (
+         pack_fn(inputs, targets, inputs_max_length, targets_max_length))
+    self.assertAllEqual(inputs_packed, [
+        [1, 2, 3, 6],
+        [4, 5, 0, 0],
+    ])
+    self.assertAllEqual(inputs_segmentation, [
+        [1, 1, 1, 2],
+        [1, 1, 0, 0],
+    ])
+    self.assertAllEqual(inputs_position, [
+        [0, 1, 2, 0],
+        [0, 1, 0, 0],
+    ])
+    self.assertAllEqual(targets_packed, [
+        [10, 50, 60, 0],
+        [20, 30, 40, 0],
+    ])
+    self.assertAllEqual(targets_segmentation, [
+        [1, 2, 2, 0],
+        [1, 1, 1, 0],
+    ])
+    self.assertAllEqual(targets_position, [
+        [0, 0, 1, 0],
+        [0, 1, 2, 0],
+    ])
+
+  def do_test_pack_sequences_length5(self, pack_fn):
     inputs = [
         [1, 2, 3],
         [4, 5, 0],
@@ -37,10 +141,9 @@ def test_pack_sequences2(self):
         [50, 60, 0],
     ]
     max_length = 5
-    (inputs_packed, inputs_segmentation, inputs_position,
-     targets_packed, targets_segmentation, targets_position) = (
-         pack_sequences_ops.pack_sequences2(
-             inputs, targets, max_length, max_length))
+    (inputs_packed, inputs_segmentation, inputs_position, targets_packed,
+     targets_segmentation, targets_position) = (
+         pack_fn(inputs, targets, max_length, max_length))
     self.assertAllEqual(
         inputs_packed, [
             [1, 2, 3, 4, 5],
@@ -72,53 +175,177 @@ def test_pack_sequences2(self):
             [0, 1, 0, 0, 0],
         ])
 
-  def test_pack_sequences_k(self):
-    inputs = tf.convert_to_tensor([
+  def do_test_pack_sequences_length6(self, pack_fn):
+    inputs = [
         [1, 2, 3],
         [4, 5, 0],
         [6, 0, 0],
-    ], dtype=tf.int32)
-    targets = tf.convert_to_tensor([
+    ]
+    targets = [
         [10, 0, 0],
         [20, 30, 40],
         [50, 60, 0],
-    ], dtype=tf.int32)
-    max_length = tf.convert_to_tensor(5, dtype=tf.int32)
-    (packed, segmentation, position) = pack_sequences_ops.pack_sequences_k(
-        [inputs, targets], [max_length, max_length])
-    (inputs_packed, targets_packed) = packed
-    (inputs_segmentation, targets_segmentation) = segmentation
-    (inputs_position, targets_position) = position
-    self.assertAllEqual(
-        inputs_packed, [
-            [1, 2, 3, 4, 5],
-            [6, 0, 0, 0, 0],
-        ])
-    self.assertAllEqual(
-        inputs_segmentation, [
-            [1, 1, 1, 2, 2],
-            [1, 0, 0, 0, 0],
-        ])
-    self.assertAllEqual(
-        inputs_position, [
-            [0, 1, 2, 0, 1],
-            [0, 0, 0, 0, 0],
-        ])
-    self.assertAllEqual(
-        targets_packed, [
-            [10, 20, 30, 40, 0],
-            [50, 60, 0, 0, 0],
-        ])
-    self.assertAllEqual(
-        targets_segmentation, [
-            [1, 2, 2, 2, 0],
-            [1, 1, 0, 0, 0],
-        ])
-    self.assertAllEqual(
-        targets_position, [
-            [0, 0, 1, 2, 0],
-            [0, 1, 0, 0, 0],
-        ])
+    ]
+    max_length = 6
+    (inputs_packed, inputs_segmentation, inputs_position, targets_packed,
+     targets_segmentation, targets_position) = (
+         pack_fn(inputs, targets, max_length, max_length))
+    self.assertAllEqual(inputs_packed, [
+        [1, 2, 3, 4, 5, 6],
+    ])
+    self.assertAllEqual(inputs_segmentation, [
+        [1, 1, 1, 2, 2, 3],
+    ])
+    self.assertAllEqual(inputs_position, [
+        [0, 1, 2, 0, 1, 0],
+    ])
+    self.assertAllEqual(targets_packed, [
+        [10, 20, 30, 40, 50, 60],
+    ])
+    self.assertAllEqual(targets_segmentation, [
+        [1, 2, 2, 2, 3, 3],
+    ])
+    self.assertAllEqual(targets_position, [
+        [0, 0, 1, 2, 0, 1],
+    ])
+
+  def do_test_pack_sequences_length7(self, pack_fn):
+    inputs = [
+        [1, 2, 3],
+        [4, 5, 0],
+        [6, 0, 0],
+    ]
+    targets = [
+        [10, 0, 0],
+        [20, 30, 40],
+        [50, 60, 0],
+    ]
+    max_length = 7
+    (inputs_packed, inputs_segmentation, inputs_position, targets_packed,
+     targets_segmentation, targets_position) = (
+         pack_fn(inputs, targets, max_length, max_length))
+    self.assertAllEqual(inputs_packed, [
+        [1, 2, 3, 4, 5, 6, 0],
+    ])
+    self.assertAllEqual(inputs_segmentation, [
+        [1, 1, 1, 2, 2, 3, 0],
+    ])
+    self.assertAllEqual(inputs_position, [
+        [0, 1, 2, 0, 1, 0, 0],
+    ])
+    self.assertAllEqual(targets_packed, [
+        [10, 20, 30, 40, 50, 60, 0],
+    ])
+    self.assertAllEqual(targets_segmentation, [
+        [1, 2, 2, 2, 3, 3, 0],
+    ])
+    self.assertAllEqual(targets_position, [
+        [0, 0, 1, 2, 0, 1, 0],
+    ])
+
+  def do_test_pack_sequences_length_different_lengths(self, pack_fn):
+    inputs = [
+        [1, 2, 3],
+        [4, 5, 0],
+        [6, 0, 0],
+    ]
+    targets = [
+        [10, 0, 0],
+        [20, 30, 40],
+        [50, 60, 0],
+    ]
+    input_max_length = 3
+    target_max_length = 4
+    (inputs_packed, inputs_segmentation, inputs_position, targets_packed,
+     targets_segmentation, targets_position) = (
+         pack_fn(inputs, targets, input_max_length, target_max_length))
+    self.assertAllEqual(inputs_packed, [
+        [1, 2, 3],
+        [4, 5, 0],
+        [6, 0, 0],
+    ])
+    self.assertAllEqual(inputs_segmentation, [
+        [1, 1, 1],
+        [1, 1, 0],
+        [1, 0, 0],
+    ])
+    self.assertAllEqual(inputs_position, [
+        [0, 1, 2],
+        [0, 1, 0],
+        [0, 0, 0],
+    ])
+    self.assertAllEqual(targets_packed, [
+        [10, 0, 0, 0],
+        [20, 30, 40, 0],
+        [50, 60, 0, 0],
+    ])
+    self.assertAllEqual(targets_segmentation, [
+        [1, 0, 0, 0],
+        [1, 1, 1, 0],
+        [1, 1, 0, 0],
+    ])
+    self.assertAllEqual(targets_position, [
+        [0, 0, 0, 0],
+        [0, 1, 2, 0],
+        [0, 1, 0, 0],
+    ])
+
+  def test_pack_sequences2(self):
+    self.do_test_pack_sequences_length3(pack_sequences_ops.pack_sequences2)
+    self.do_test_pack_sequences_length4(pack_sequences_ops.pack_sequences2)
+    self.do_test_pack_sequences_length5(pack_sequences_ops.pack_sequences2)
+    self.do_test_pack_sequences_length6(pack_sequences_ops.pack_sequences2)
+    self.do_test_pack_sequences_length7(pack_sequences_ops.pack_sequences2)
+    self.do_test_pack_sequences_length_different_lengths(
+        pack_sequences_ops.pack_sequences2)
+
+  def test_pack_sequences_k(self):
+    self.do_test_pack_sequences_length3(_pack_sequences_k)
+    self.do_test_pack_sequences_length4(_pack_sequences_k)
+    self.do_test_pack_sequences_length5(_pack_sequences_k)
+    self.do_test_pack_sequences_length6(_pack_sequences_k)
+    self.do_test_pack_sequences_length7(_pack_sequences_k)
+    self.do_test_pack_sequences_length_different_lengths(_pack_sequences_k)
+
+  def test_random_inputs(self):
+    for _ in range(10):
+      batch_size = np.random.randint(900, 1100, size=[])
+      input_seqlen = np.random.randint(1, 10, size=[])
+      target_seqlen = np.random.randint(1, 10, size=[])
+      inputs_list = []
+      targets_list = []
+      for _ in range(batch_size):
+        input_num_pads = np.random.randint(0, input_seqlen, size=[])
+        input_pads = np.full([input_num_pads], 0, dtype=np.int32)
+        inputs = np.random.randint(1, 10, size=[input_seqlen - input_num_pads])
+        inputs = np.concatenate([inputs, input_pads], axis=0)
+
+        target_num_pads = np.random.randint(0, target_seqlen, size=[])
+        target_pads = np.full([target_num_pads], 0, dtype=np.int32)
+        targets = np.random.randint(
+            1, 10, size=[target_seqlen - target_num_pads])
+        targets = np.concatenate([targets, target_pads], axis=0)
+
+        inputs_list.append(inputs)
+        targets_list.append(targets)
+      input_maxlen = np.random.randint(input_seqlen, input_seqlen + 10, size=[])
+      target_maxlen = np.random.randint(
+          target_seqlen, target_seqlen + 10, size=[])
+      (inputs_packed2, inputs_segmentation2, inputs_positions2, targets_packed2,
+       targets_segmentation2, targets_positions2) = (
+           pack_sequences_ops.pack_sequences2(inputs_list, targets_list,
+                                              input_maxlen, target_maxlen))
+      (inputs_packed_k, inputs_segmentation_k, inputs_positions_k,
+       targets_packed_k, targets_segmentation_k, targets_positions_k) = (
+           _pack_sequences_k(inputs_list, targets_list, input_maxlen,
+                             target_maxlen))
+
+      self.assertAllEqual(inputs_packed2, inputs_packed_k)
+      self.assertAllEqual(inputs_segmentation2, inputs_segmentation_k)
+      self.assertAllEqual(inputs_positions2, inputs_positions_k)
+      self.assertAllEqual(targets_packed2, targets_packed_k)
+      self.assertAllEqual(targets_segmentation2, targets_segmentation_k)
+      self.assertAllEqual(targets_positions2, targets_positions_k)
 
   def test_pack_sequences_k_multi_input(self):
     input_tokens = tf.convert_to_tensor([

From 0b3ee79c460ac5b0d5bb32b9e48001a933cdc5db Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 21 Oct 2021 13:40:36 -0700
Subject: [PATCH 2699/2720] Fixes python3 related unicode errors in
 wiki_revision_utils.

PiperOrigin-RevId: 404870305
---
 tensor2tensor/data_generators/wiki_revision_utils.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index f938abb77..aae6010ab 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -33,6 +33,10 @@
 import tensorflow.compat.v1 as tf
 
 
+def to_unicode(s):
+  return s.decode("utf-8")
+
+
 def include_revision(revision_num, skip_factor=1.1):
   """Decide whether to include a revision.
 
@@ -76,6 +80,10 @@ def file_page_generator(my_file, max_page_size=2**28):
   leftovers = ""
   while True:
     chunk = my_file.read(chunk_size)
+    try:
+      chunk = to_unicode(chunk)
+    except UnicodeDecodeError:
+      chunk = ""
     if not chunk:
       break
     chunk = leftovers + chunk
@@ -112,7 +120,7 @@ def get_title(page):
   assert start_pos != -1
   assert end_pos != -1
   start_pos += len("<title>")
-  return text_encoder.to_unicode_utf8(page[start_pos:end_pos])
+  return page[start_pos:end_pos]
 
 
 def get_id(page):
@@ -251,7 +259,6 @@ def get_text(revision, strip=True):
     ret = revision[end_tag_pos:end_pos]
   if strip:
     ret = strip_text(ret)
-  ret = text_encoder.to_unicode_utf8(ret)
   return ret
 
 

From 0dcb66520ca2f46630f0b384683e15ca82ae73e7 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 3 Nov 2021 14:34:02 -0700
Subject: [PATCH 2700/2720] Update Eigen to
 https://gitlab.com/libeigen/eigen/-/commit/9cf34ee0aed25a7464e6ec14f977cfa940f48f1b

PiperOrigin-RevId: 407427880
---
 tensor2tensor/layers/transformer_glow_layers_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/transformer_glow_layers_test.py b/tensor2tensor/layers/transformer_glow_layers_test.py
index a65b77d89..92bdf28ae 100644
--- a/tensor2tensor/layers/transformer_glow_layers_test.py
+++ b/tensor2tensor/layers/transformer_glow_layers_test.py
@@ -246,7 +246,7 @@ def test_flow_invertibility(self):
     diff = x - x_inv_inv
     logabsdet_sum = logabsdet + logabsdet_inv
     self.assertTrue(np.allclose(diff, 0.0, atol=2e-5))
-    self.assertTrue(np.allclose(logabsdet_sum, 0.0, atol=2e-5))
+    self.assertTrue(np.allclose(logabsdet_sum, 0.0, atol=7e-5))
 
   @parameterized.parameters(
       ("1", "cat", "affine"),

From 86caf01fd174c2ea2d3c789abb27a43f8948df1f Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Fri, 5 Nov 2021 13:06:56 -0700
Subject: [PATCH 2701/2720] Update gym to 0.19.0

PiperOrigin-RevId: 407891929
---
 tensor2tensor/envs/gym_env_problem_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/envs/gym_env_problem_test.py b/tensor2tensor/envs/gym_env_problem_test.py
index 9d5e8f693..93d8093e3 100644
--- a/tensor2tensor/envs/gym_env_problem_test.py
+++ b/tensor2tensor/envs/gym_env_problem_test.py
@@ -68,7 +68,7 @@ def test_reward_range(self):
     # Passing reward_range=None means take the reward range of the underlying
     # environment as the reward range.
     ep = gym_env_problem.GymEnvProblem(
-        base_env_name="FrozenLake-v0", batch_size=5, reward_range=None)
+        base_env_name="FrozenLake-v1", batch_size=5, reward_range=None)
     ep.assert_common_preconditions()
 
     # Assert reward range is finite here.

From 81c2b2eddbf09e0c963d32e6eb7a3027d736d808 Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 2 Dec 2021 10:49:52 -0800
Subject: [PATCH 2702/2720] Update adafactor so it can accept a callable
 learning rate.

PiperOrigin-RevId: 413718409
---
 tensor2tensor/utils/adafactor.py      |  7 +++--
 tensor2tensor/utils/adafactor_test.py | 44 +++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100644 tensor2tensor/utils/adafactor_test.py

diff --git a/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index 7fa139c5b..12ed63143 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -122,7 +122,7 @@ def __init__(self,
 
     Args:
       multiply_by_parameter_scale: a boolean
-      learning_rate: an optional Scalar.
+      learning_rate: an optional Scalar or callable.
       decay_rate: an optional Scalar.
       beta1: a float value between 0 and 1
       clipping_threshold: an optional float >= 1
@@ -218,7 +218,9 @@ def _resource_apply_dense(self, grad, handle):
     grad_squared = tf.square(grad) + self._epsilon1
     grad_squared_mean = tf.reduce_mean(grad_squared)
     decay_rate = self._decay_rate
-    update_scale = self._learning_rate
+    update_scale = self._call_if_callable(self._learning_rate)
+    update_scale = tf.convert_to_tensor(update_scale, name="update_scale")
+    update_scale = tf.cast(update_scale, grad_squared_mean.dtype.base_dtype)
     old_val = var
     if var.dtype.base_dtype == tf.bfloat16:
       old_val = tf.to_float(self._parameter_encoding.decode(old_val))
@@ -272,6 +274,7 @@ def _resource_apply_dense(self, grad, handle):
       new_val = quantization.simulated_quantize(
           var - subtrahend, self._simulated_quantize_bits,
           self._quantization_noise)
+    new_val = tf.cast(new_val, var.dtype)
     var_update = tf.assign(var, new_val, use_locking=self._use_locking)
     updates = [var_update] + updates
     return tf.group(*updates)
diff --git a/tensor2tensor/utils/adafactor_test.py b/tensor2tensor/utils/adafactor_test.py
new file mode 100644
index 000000000..56befa985
--- /dev/null
+++ b/tensor2tensor/utils/adafactor_test.py
@@ -0,0 +1,44 @@
+# coding=utf-8
+# Copyright 2021 The Tensor2Tensor Authors.
+#
+# 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 adafactor."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensor2tensor.utils import adafactor
+
+import tensorflow as tf
+
+
+class AdafactorTest(tf.test.TestCase):
+
+  def testCallableLearningRate(self):
+    def lr():
+      return 0.01
+
+    opt = adafactor.AdafactorOptimizer(learning_rate=lr)
+    v1 = tf.Variable([1., 2.])
+    v2 = tf.Variable([3., 4.])
+    with tf.GradientTape() as tape:
+      tape.watch([v1, v2])
+      loss = v1 * v2
+    v1_grad, v2_grad = tape.gradient(loss, [v1, v2])
+    opt.apply_gradients(((v1_grad, v1), (v2_grad, v2)))
+
+
+if __name__ == '__main__':
+  tf.test.main()

From 2a33b152d7835af66a6d20afe7961751047e28dd Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Wed, 12 Jan 2022 10:23:57 -0800
Subject: [PATCH 2703/2720] Support callable decay rates in Adafactor

PiperOrigin-RevId: 421333767
---
 tensor2tensor/utils/adafactor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index 12ed63143..1383fc271 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -217,7 +217,7 @@ def _resource_apply_dense(self, grad, handle):
     grad = tf.to_float(grad)
     grad_squared = tf.square(grad) + self._epsilon1
     grad_squared_mean = tf.reduce_mean(grad_squared)
-    decay_rate = self._decay_rate
+    decay_rate = self._call_if_callable(self._decay_rate)
     update_scale = self._call_if_callable(self._learning_rate)
     update_scale = tf.convert_to_tensor(update_scale, name="update_scale")
     update_scale = tf.cast(update_scale, grad_squared_mean.dtype.base_dtype)

From ef1fccebe8d2c0cf482f41f9d940e2938c816c78 Mon Sep 17 00:00:00 2001
From: Dan Zheng <danielzheng@google.com>
Date: Thu, 24 Feb 2022 19:30:23 -0800
Subject: [PATCH 2704/2720] Update out-of-vocabulary error message to include
 the OOV token.

This facilitates debugging.

PiperOrigin-RevId: 430848881
---
 tensor2tensor/__init__.py                                   | 2 +-
 tensor2tensor/bin/__init__.py                               | 2 +-
 tensor2tensor/bin/build_vocab.py                            | 2 +-
 tensor2tensor/bin/make_tf_configs.py                        | 2 +-
 tensor2tensor/bin/t2t_attack.py                             | 2 +-
 tensor2tensor/bin/t2t_avg_all.py                            | 2 +-
 tensor2tensor/bin/t2t_bleu.py                               | 2 +-
 tensor2tensor/bin/t2t_datagen.py                            | 2 +-
 tensor2tensor/bin/t2t_decoder.py                            | 2 +-
 tensor2tensor/bin/t2t_distill.py                            | 2 +-
 tensor2tensor/bin/t2t_eval.py                               | 2 +-
 tensor2tensor/bin/t2t_prune.py                              | 2 +-
 tensor2tensor/bin/t2t_trainer.py                            | 2 +-
 tensor2tensor/bin/t2t_trainer_test.py                       | 2 +-
 tensor2tensor/bin/t2t_translate_all.py                      | 2 +-
 tensor2tensor/data_generators/__init__.py                   | 2 +-
 tensor2tensor/data_generators/algorithmic.py                | 2 +-
 tensor2tensor/data_generators/algorithmic_math.py           | 2 +-
 tensor2tensor/data_generators/algorithmic_math_deepmind.py  | 2 +-
 tensor2tensor/data_generators/algorithmic_math_test.py      | 2 +-
 .../data_generators/algorithmic_math_two_variables.py       | 2 +-
 tensor2tensor/data_generators/algorithmic_test.py           | 2 +-
 tensor2tensor/data_generators/all_problems.py               | 2 +-
 tensor2tensor/data_generators/allen_brain.py                | 2 +-
 tensor2tensor/data_generators/allen_brain_test.py           | 2 +-
 tensor2tensor/data_generators/audio.py                      | 2 +-
 tensor2tensor/data_generators/audio_encoder.py              | 2 +-
 tensor2tensor/data_generators/audio_test.py                 | 2 +-
 tensor2tensor/data_generators/babi_qa.py                    | 2 +-
 tensor2tensor/data_generators/bair_robot_pushing.py         | 2 +-
 tensor2tensor/data_generators/celeba.py                     | 2 +-
 tensor2tensor/data_generators/celeba_test.py                | 2 +-
 tensor2tensor/data_generators/celebahq.py                   | 2 +-
 tensor2tensor/data_generators/cifar.py                      | 2 +-
 tensor2tensor/data_generators/cipher.py                     | 2 +-
 tensor2tensor/data_generators/cleaner_en_xx.py              | 2 +-
 tensor2tensor/data_generators/cnn_dailymail.py              | 2 +-
 tensor2tensor/data_generators/cola.py                       | 2 +-
 tensor2tensor/data_generators/common_voice.py               | 2 +-
 tensor2tensor/data_generators/common_voice_test.py          | 2 +-
 tensor2tensor/data_generators/conll_ner.py                  | 2 +-
 tensor2tensor/data_generators/desc2code.py                  | 2 +-
 tensor2tensor/data_generators/desc2code_test.py             | 2 +-
 tensor2tensor/data_generators/dialog_abstract.py            | 2 +-
 tensor2tensor/data_generators/dialog_cornell.py             | 2 +-
 tensor2tensor/data_generators/dialog_dailydialog.py         | 2 +-
 tensor2tensor/data_generators/dialog_opensubtitles.py       | 2 +-
 tensor2tensor/data_generators/dialog_personachat.py         | 2 +-
 tensor2tensor/data_generators/dna_encoder.py                | 2 +-
 tensor2tensor/data_generators/dna_encoder_test.py           | 2 +-
 tensor2tensor/data_generators/enwik8.py                     | 2 +-
 tensor2tensor/data_generators/fsns.py                       | 2 +-
 tensor2tensor/data_generators/function_docstring.py         | 2 +-
 tensor2tensor/data_generators/gene_expression.py            | 2 +-
 tensor2tensor/data_generators/gene_expression_test.py       | 2 +-
 tensor2tensor/data_generators/generator_utils.py            | 2 +-
 tensor2tensor/data_generators/generator_utils_test.py       | 2 +-
 tensor2tensor/data_generators/google_robot_pushing.py       | 2 +-
 tensor2tensor/data_generators/gym_env.py                    | 2 +-
 tensor2tensor/data_generators/gym_env_test.py               | 2 +-
 tensor2tensor/data_generators/ice_parsing.py                | 2 +-
 tensor2tensor/data_generators/image_lsun.py                 | 2 +-
 tensor2tensor/data_generators/image_utils.py                | 2 +-
 tensor2tensor/data_generators/image_utils_test.py           | 2 +-
 tensor2tensor/data_generators/imagenet.py                   | 2 +-
 tensor2tensor/data_generators/imagenet_test.py              | 2 +-
 tensor2tensor/data_generators/imdb.py                       | 2 +-
 tensor2tensor/data_generators/inspect_tfrecord.py           | 2 +-
 tensor2tensor/data_generators/lambada.py                    | 2 +-
 tensor2tensor/data_generators/librispeech.py                | 2 +-
 tensor2tensor/data_generators/lm1b.py                       | 2 +-
 tensor2tensor/data_generators/lm1b_imdb.py                  | 2 +-
 tensor2tensor/data_generators/lm1b_mnli.py                  | 2 +-
 tensor2tensor/data_generators/mnist.py                      | 2 +-
 tensor2tensor/data_generators/moving_mnist.py               | 2 +-
 tensor2tensor/data_generators/mrpc.py                       | 2 +-
 tensor2tensor/data_generators/mscoco.py                     | 2 +-
 tensor2tensor/data_generators/mscoco_test.py                | 2 +-
 tensor2tensor/data_generators/multi_problem.py              | 2 +-
 tensor2tensor/data_generators/multi_problem_v2.py           | 2 +-
 tensor2tensor/data_generators/multi_problem_v2_test.py      | 2 +-
 tensor2tensor/data_generators/multinli.py                   | 2 +-
 tensor2tensor/data_generators/ocr.py                        | 2 +-
 .../data_generators/ops/pack_sequences_ops_test.py          | 2 +-
 .../data_generators/ops/subword_text_encoder_ops_test.py    | 2 +-
 tensor2tensor/data_generators/paraphrase_ms_coco.py         | 2 +-
 tensor2tensor/data_generators/paraphrase_ms_coco_test.py    | 2 +-
 tensor2tensor/data_generators/pointer_generator_word.py     | 2 +-
 tensor2tensor/data_generators/problem.py                    | 2 +-
 tensor2tensor/data_generators/problem_hparams.py            | 2 +-
 tensor2tensor/data_generators/problem_test.py               | 2 +-
 tensor2tensor/data_generators/program_search.py             | 2 +-
 tensor2tensor/data_generators/program_search_test.py        | 2 +-
 tensor2tensor/data_generators/ptb.py                        | 2 +-
 tensor2tensor/data_generators/qnli.py                       | 2 +-
 tensor2tensor/data_generators/quora_qpairs.py               | 2 +-
 tensor2tensor/data_generators/rte.py                        | 2 +-
 tensor2tensor/data_generators/scitail.py                    | 2 +-
 tensor2tensor/data_generators/seq2edits.py                  | 2 +-
 tensor2tensor/data_generators/snli.py                       | 2 +-
 tensor2tensor/data_generators/speech_recognition.py         | 2 +-
 tensor2tensor/data_generators/squad.py                      | 2 +-
 tensor2tensor/data_generators/sst_binary.py                 | 2 +-
 tensor2tensor/data_generators/stanford_nli.py               | 2 +-
 tensor2tensor/data_generators/style_transfer.py             | 2 +-
 tensor2tensor/data_generators/style_transfer_test.py        | 2 +-
 tensor2tensor/data_generators/subject_verb_agreement.py     | 2 +-
 tensor2tensor/data_generators/text_encoder.py               | 6 ++++--
 tensor2tensor/data_generators/text_encoder_build_subword.py | 2 +-
 tensor2tensor/data_generators/text_encoder_test.py          | 4 ++--
 tensor2tensor/data_generators/text_problems.py              | 2 +-
 tensor2tensor/data_generators/text_problems_test.py         | 2 +-
 tensor2tensor/data_generators/timeseries.py                 | 2 +-
 tensor2tensor/data_generators/timeseries_data_generator.py  | 2 +-
 .../data_generators/timeseries_data_generator_test.py       | 2 +-
 tensor2tensor/data_generators/timeseries_test.py            | 2 +-
 tensor2tensor/data_generators/tokenizer.py                  | 2 +-
 tensor2tensor/data_generators/tokenizer_test.py             | 2 +-
 tensor2tensor/data_generators/transduction_problems.py      | 2 +-
 tensor2tensor/data_generators/transduction_problems_test.py | 2 +-
 tensor2tensor/data_generators/translate.py                  | 2 +-
 tensor2tensor/data_generators/translate_encs.py             | 2 +-
 tensor2tensor/data_generators/translate_encs_cubbitt.py     | 2 +-
 tensor2tensor/data_generators/translate_ende.py             | 2 +-
 tensor2tensor/data_generators/translate_ende_test.py        | 2 +-
 tensor2tensor/data_generators/translate_enes.py             | 2 +-
 tensor2tensor/data_generators/translate_enet.py             | 2 +-
 tensor2tensor/data_generators/translate_enfr.py             | 2 +-
 tensor2tensor/data_generators/translate_enid.py             | 2 +-
 tensor2tensor/data_generators/translate_enmk.py             | 2 +-
 tensor2tensor/data_generators/translate_enro.py             | 2 +-
 tensor2tensor/data_generators/translate_entn.py             | 2 +-
 tensor2tensor/data_generators/translate_envi.py             | 2 +-
 tensor2tensor/data_generators/translate_enzh.py             | 2 +-
 tensor2tensor/data_generators/translate_test.py             | 2 +-
 tensor2tensor/data_generators/video_generated.py            | 2 +-
 tensor2tensor/data_generators/video_utils.py                | 2 +-
 tensor2tensor/data_generators/video_utils_test.py           | 2 +-
 tensor2tensor/data_generators/vqa.py                        | 2 +-
 tensor2tensor/data_generators/vqa_utils.py                  | 2 +-
 tensor2tensor/data_generators/wiki.py                       | 2 +-
 tensor2tensor/data_generators/wiki_lm.py                    | 2 +-
 tensor2tensor/data_generators/wiki_multi_problems.py        | 2 +-
 tensor2tensor/data_generators/wiki_revision.py              | 2 +-
 tensor2tensor/data_generators/wiki_revision_utils.py        | 2 +-
 tensor2tensor/data_generators/wikisum/__init__.py           | 2 +-
 tensor2tensor/data_generators/wikisum/generate_vocab.py     | 2 +-
 .../data_generators/wikisum/get_references_commoncrawl.py   | 2 +-
 tensor2tensor/data_generators/wikisum/get_references_web.py | 2 +-
 .../wikisum/get_references_web_single_group.py              | 2 +-
 tensor2tensor/data_generators/wikisum/html.py               | 2 +-
 tensor2tensor/data_generators/wikisum/parallel_launch.py    | 2 +-
 tensor2tensor/data_generators/wikisum/produce_examples.py   | 2 +-
 tensor2tensor/data_generators/wikisum/utils.py              | 2 +-
 tensor2tensor/data_generators/wikisum/utils_test.py         | 2 +-
 tensor2tensor/data_generators/wikisum/validate_data.py      | 2 +-
 tensor2tensor/data_generators/wikisum/wikisum.py            | 2 +-
 tensor2tensor/data_generators/wikitext103.py                | 2 +-
 tensor2tensor/data_generators/wnli.py                       | 2 +-
 tensor2tensor/data_generators/wsj_parsing.py                | 2 +-
 tensor2tensor/data_generators/yelp_full.py                  | 2 +-
 tensor2tensor/data_generators/yelp_polarity.py              | 2 +-
 tensor2tensor/envs/__init__.py                              | 2 +-
 tensor2tensor/envs/env_problem.py                           | 2 +-
 tensor2tensor/envs/env_problem_utils.py                     | 2 +-
 tensor2tensor/envs/env_problem_utils_test.py                | 2 +-
 tensor2tensor/envs/gym_env_problem.py                       | 2 +-
 tensor2tensor/envs/gym_env_problem_test.py                  | 2 +-
 tensor2tensor/envs/gym_spaces_utils.py                      | 2 +-
 tensor2tensor/envs/gym_spaces_utils_test.py                 | 2 +-
 tensor2tensor/envs/mujoco_problems.py                       | 2 +-
 tensor2tensor/envs/mujoco_problems_test.py                  | 2 +-
 tensor2tensor/envs/rendered_env_problem.py                  | 2 +-
 tensor2tensor/envs/rendered_env_problem_test.py             | 2 +-
 tensor2tensor/envs/tic_tac_toe_env.py                       | 2 +-
 tensor2tensor/envs/tic_tac_toe_env_problem.py               | 2 +-
 tensor2tensor/envs/tic_tac_toe_env_problem_test.py          | 2 +-
 tensor2tensor/envs/tic_tac_toe_env_test.py                  | 2 +-
 tensor2tensor/envs/time_step.py                             | 2 +-
 tensor2tensor/envs/time_step_test.py                        | 2 +-
 tensor2tensor/envs/trajectory.py                            | 2 +-
 tensor2tensor/envs/trajectory_test.py                       | 2 +-
 tensor2tensor/insights/__init__.py                          | 2 +-
 tensor2tensor/insights/graph.py                             | 2 +-
 tensor2tensor/insights/query_processor.py                   | 2 +-
 tensor2tensor/insights/server.py                            | 2 +-
 tensor2tensor/insights/transformer_model.py                 | 2 +-
 tensor2tensor/layers/__init__.py                            | 2 +-
 tensor2tensor/layers/area_attention.py                      | 2 +-
 tensor2tensor/layers/area_attention_test.py                 | 2 +-
 tensor2tensor/layers/common_attention.py                    | 2 +-
 tensor2tensor/layers/common_attention_test.py               | 2 +-
 tensor2tensor/layers/common_audio.py                        | 2 +-
 tensor2tensor/layers/common_hparams.py                      | 2 +-
 tensor2tensor/layers/common_image_attention.py              | 2 +-
 tensor2tensor/layers/common_image_attention_test.py         | 2 +-
 tensor2tensor/layers/common_layers.py                       | 2 +-
 tensor2tensor/layers/common_layers_test.py                  | 2 +-
 tensor2tensor/layers/common_video.py                        | 2 +-
 tensor2tensor/layers/common_video_test.py                   | 2 +-
 tensor2tensor/layers/discretization.py                      | 2 +-
 tensor2tensor/layers/discretization_test.py                 | 2 +-
 tensor2tensor/layers/latent_layers.py                       | 2 +-
 tensor2tensor/layers/latent_layers_test.py                  | 2 +-
 tensor2tensor/layers/message_passing_attention.py           | 2 +-
 tensor2tensor/layers/modalities.py                          | 2 +-
 tensor2tensor/layers/modalities_test.py                     | 2 +-
 tensor2tensor/layers/ngram.py                               | 2 +-
 tensor2tensor/layers/ngram_test.py                          | 2 +-
 tensor2tensor/layers/transformer_glow_layers.py             | 2 +-
 tensor2tensor/layers/transformer_glow_layers_ops.py         | 2 +-
 tensor2tensor/layers/transformer_glow_layers_ops_test.py    | 2 +-
 tensor2tensor/layers/transformer_glow_layers_test.py        | 2 +-
 tensor2tensor/layers/transformer_layers.py                  | 2 +-
 tensor2tensor/layers/transformer_memory.py                  | 2 +-
 tensor2tensor/layers/transformer_memory_test.py             | 2 +-
 tensor2tensor/layers/vq_discrete.py                         | 2 +-
 tensor2tensor/layers/vqa_layers.py                          | 2 +-
 tensor2tensor/metrics/__init__.py                           | 2 +-
 tensor2tensor/metrics/video_conditional_fvd.py              | 2 +-
 tensor2tensor/metrics/video_conditional_fvd_test.py         | 2 +-
 tensor2tensor/models/__init__.py                            | 2 +-
 tensor2tensor/models/basic.py                               | 2 +-
 tensor2tensor/models/basic_test.py                          | 2 +-
 tensor2tensor/models/bytenet.py                             | 2 +-
 tensor2tensor/models/bytenet_test.py                        | 2 +-
 tensor2tensor/models/distillation.py                        | 2 +-
 tensor2tensor/models/evolved_transformer.py                 | 2 +-
 tensor2tensor/models/evolved_transformer_test.py            | 2 +-
 tensor2tensor/models/image_transformer.py                   | 2 +-
 tensor2tensor/models/image_transformer_2d.py                | 2 +-
 tensor2tensor/models/image_transformer_2d_test.py           | 2 +-
 tensor2tensor/models/image_transformer_test.py              | 2 +-
 tensor2tensor/models/lstm.py                                | 2 +-
 tensor2tensor/models/lstm_test.py                           | 2 +-
 tensor2tensor/models/mtf_image_transformer.py               | 2 +-
 tensor2tensor/models/mtf_image_transformer_test.py          | 2 +-
 tensor2tensor/models/mtf_resnet.py                          | 2 +-
 tensor2tensor/models/mtf_transformer.py                     | 2 +-
 tensor2tensor/models/mtf_transformer2.py                    | 2 +-
 tensor2tensor/models/mtf_transformer_test.py                | 2 +-
 tensor2tensor/models/neural_architecture_search/__init__.py | 2 +-
 .../models/neural_architecture_search/nas_layers.py         | 2 +-
 .../models/neural_architecture_search/nas_layers_test.py    | 2 +-
 .../models/neural_architecture_search/nas_model.py          | 2 +-
 .../models/neural_architecture_search/nas_model_test.py     | 2 +-
 tensor2tensor/models/neural_assistant.py                    | 2 +-
 tensor2tensor/models/neural_gpu.py                          | 2 +-
 tensor2tensor/models/neural_gpu_test.py                     | 2 +-
 tensor2tensor/models/research/__init__.py                   | 2 +-
 tensor2tensor/models/research/adafactor_experiments.py      | 2 +-
 tensor2tensor/models/research/aligned.py                    | 2 +-
 tensor2tensor/models/research/attention_lm.py               | 2 +-
 tensor2tensor/models/research/attention_lm_moe.py           | 2 +-
 tensor2tensor/models/research/autoencoders.py               | 2 +-
 tensor2tensor/models/research/autoencoders_test.py          | 2 +-
 tensor2tensor/models/research/cycle_gan.py                  | 2 +-
 tensor2tensor/models/research/gene_expression.py            | 2 +-
 tensor2tensor/models/research/gene_expression_test.py       | 2 +-
 tensor2tensor/models/research/glow.py                       | 2 +-
 tensor2tensor/models/research/glow_init_hook.py             | 2 +-
 tensor2tensor/models/research/glow_ops.py                   | 2 +-
 tensor2tensor/models/research/glow_ops_test.py              | 2 +-
 tensor2tensor/models/research/glow_test.py                  | 2 +-
 tensor2tensor/models/research/lm_experiments.py             | 2 +-
 tensor2tensor/models/research/moe.py                        | 2 +-
 tensor2tensor/models/research/moe_experiments.py            | 2 +-
 tensor2tensor/models/research/multiquery_paper.py           | 2 +-
 tensor2tensor/models/research/neural_stack.py               | 2 +-
 tensor2tensor/models/research/neural_stack_test.py          | 2 +-
 tensor2tensor/models/research/residual_shuffle_exchange.py  | 2 +-
 tensor2tensor/models/research/rl.py                         | 2 +-
 tensor2tensor/models/research/shuffle_network.py            | 2 +-
 tensor2tensor/models/research/similarity_transformer.py     | 2 +-
 tensor2tensor/models/research/super_lm.py                   | 2 +-
 tensor2tensor/models/research/transformer_aux.py            | 2 +-
 tensor2tensor/models/research/transformer_aux_test.py       | 2 +-
 tensor2tensor/models/research/transformer_moe.py            | 2 +-
 tensor2tensor/models/research/transformer_nat.py            | 2 +-
 tensor2tensor/models/research/transformer_parallel.py       | 2 +-
 tensor2tensor/models/research/transformer_revnet.py         | 2 +-
 tensor2tensor/models/research/transformer_revnet_test.py    | 2 +-
 tensor2tensor/models/research/transformer_seq2edits.py      | 2 +-
 tensor2tensor/models/research/transformer_sketch.py         | 2 +-
 tensor2tensor/models/research/transformer_symshard.py       | 2 +-
 tensor2tensor/models/research/transformer_vae.py            | 2 +-
 tensor2tensor/models/research/transformer_vae_flow_prior.py | 2 +-
 .../models/research/transformer_vae_flow_prior_ops.py       | 2 +-
 tensor2tensor/models/research/transformer_vae_test.py       | 2 +-
 tensor2tensor/models/research/universal_transformer.py      | 2 +-
 tensor2tensor/models/research/universal_transformer_test.py | 2 +-
 tensor2tensor/models/research/universal_transformer_util.py | 2 +-
 tensor2tensor/models/research/vqa_attention.py              | 2 +-
 tensor2tensor/models/research/vqa_attention_test.py         | 2 +-
 .../models/research/vqa_recurrent_self_attention.py         | 2 +-
 tensor2tensor/models/research/vqa_self_attention.py         | 2 +-
 tensor2tensor/models/resnet.py                              | 2 +-
 tensor2tensor/models/resnet_test.py                         | 2 +-
 tensor2tensor/models/revnet.py                              | 2 +-
 tensor2tensor/models/revnet_test.py                         | 2 +-
 tensor2tensor/models/shake_shake.py                         | 2 +-
 tensor2tensor/models/slicenet.py                            | 2 +-
 tensor2tensor/models/slicenet_test.py                       | 2 +-
 tensor2tensor/models/text_cnn.py                            | 2 +-
 tensor2tensor/models/transformer.py                         | 2 +-
 tensor2tensor/models/transformer_test.py                    | 2 +-
 tensor2tensor/models/vanilla_gan.py                         | 2 +-
 tensor2tensor/models/video/__init__.py                      | 2 +-
 tensor2tensor/models/video/base.py                          | 2 +-
 tensor2tensor/models/video/base_vae.py                      | 2 +-
 tensor2tensor/models/video/basic_deterministic.py           | 2 +-
 tensor2tensor/models/video/basic_deterministic_params.py    | 2 +-
 tensor2tensor/models/video/basic_deterministic_test.py      | 2 +-
 tensor2tensor/models/video/basic_recurrent.py               | 2 +-
 tensor2tensor/models/video/basic_recurrent_test.py          | 2 +-
 tensor2tensor/models/video/basic_stochastic.py              | 2 +-
 tensor2tensor/models/video/basic_stochastic_test.py         | 2 +-
 tensor2tensor/models/video/emily.py                         | 2 +-
 tensor2tensor/models/video/emily_test.py                    | 2 +-
 tensor2tensor/models/video/epva.py                          | 2 +-
 tensor2tensor/models/video/epva_params.py                   | 2 +-
 tensor2tensor/models/video/next_frame_glow.py               | 2 +-
 tensor2tensor/models/video/nfg_conv3d_test.py               | 2 +-
 tensor2tensor/models/video/nfg_conv_lstm_test.py            | 2 +-
 tensor2tensor/models/video/nfg_conv_test.py                 | 2 +-
 tensor2tensor/models/video/nfg_interpolate.py               | 2 +-
 tensor2tensor/models/video/nfg_test_utils.py                | 2 +-
 tensor2tensor/models/video/nfg_uncond_test.py               | 2 +-
 tensor2tensor/models/video/savp.py                          | 2 +-
 tensor2tensor/models/video/savp_params.py                   | 2 +-
 tensor2tensor/models/video/savp_test.py                     | 2 +-
 tensor2tensor/models/video/sv2p.py                          | 2 +-
 tensor2tensor/models/video/sv2p_params.py                   | 2 +-
 tensor2tensor/models/video/sv2p_test.py                     | 2 +-
 tensor2tensor/models/video/tests_utils.py                   | 2 +-
 tensor2tensor/models/xception.py                            | 2 +-
 tensor2tensor/models/xception_test.py                       | 2 +-
 tensor2tensor/problems.py                                   | 2 +-
 tensor2tensor/problems_colab.py                             | 2 +-
 tensor2tensor/problems_test.py                              | 2 +-
 tensor2tensor/rl/__init__.py                                | 2 +-
 tensor2tensor/rl/batch_dqn_agent_test.py                    | 2 +-
 tensor2tensor/rl/batch_runner_test.py                       | 2 +-
 tensor2tensor/rl/datagen_with_agent.py                      | 2 +-
 tensor2tensor/rl/dopamine_connector.py                      | 2 +-
 tensor2tensor/rl/envs/__init__.py                           | 2 +-
 tensor2tensor/rl/envs/in_graph_batch_env.py                 | 2 +-
 tensor2tensor/rl/envs/py_func_batch_env.py                  | 2 +-
 tensor2tensor/rl/envs/simulated_batch_env.py                | 2 +-
 tensor2tensor/rl/envs/simulated_batch_gym_env.py            | 2 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py                  | 2 +-
 tensor2tensor/rl/evaluator.py                               | 2 +-
 tensor2tensor/rl/evaluator_test.py                          | 2 +-
 tensor2tensor/rl/gym_utils.py                               | 2 +-
 tensor2tensor/rl/gym_utils_test.py                          | 2 +-
 tensor2tensor/rl/player.py                                  | 2 +-
 tensor2tensor/rl/player_utils.py                            | 2 +-
 tensor2tensor/rl/policy_learner.py                          | 2 +-
 tensor2tensor/rl/ppo.py                                     | 2 +-
 tensor2tensor/rl/ppo_learner.py                             | 2 +-
 tensor2tensor/rl/restarter.py                               | 2 +-
 tensor2tensor/rl/restarter_test.py                          | 2 +-
 tensor2tensor/rl/rl_utils.py                                | 2 +-
 tensor2tensor/rl/trainer_model_based.py                     | 2 +-
 tensor2tensor/rl/trainer_model_based_agent_only.py          | 2 +-
 tensor2tensor/rl/trainer_model_based_params.py              | 2 +-
 tensor2tensor/rl/trainer_model_based_recurrent_test.py      | 2 +-
 tensor2tensor/rl/trainer_model_based_stochastic_test.py     | 2 +-
 tensor2tensor/rl/trainer_model_based_sv2p_test.py           | 2 +-
 tensor2tensor/rl/trainer_model_based_test.py                | 2 +-
 tensor2tensor/rl/trainer_model_free.py                      | 2 +-
 tensor2tensor/rl/trainer_model_free_test.py                 | 2 +-
 tensor2tensor/rl/trainer_model_free_tictactoe_test.py       | 2 +-
 tensor2tensor/serving/__init__.py                           | 2 +-
 tensor2tensor/serving/export.py                             | 2 +-
 tensor2tensor/serving/query.py                              | 2 +-
 tensor2tensor/serving/serving_utils.py                      | 2 +-
 tensor2tensor/test_data/example_usr_dir/__init__.py         | 2 +-
 tensor2tensor/test_data/example_usr_dir/my_submodule.py     | 2 +-
 tensor2tensor/utils/__init__.py                             | 2 +-
 tensor2tensor/utils/adafactor.py                            | 2 +-
 tensor2tensor/utils/adafactor_test.py                       | 2 +-
 tensor2tensor/utils/adv_attack_utils.py                     | 2 +-
 tensor2tensor/utils/avg_checkpoints.py                      | 2 +-
 tensor2tensor/utils/beam_search.py                          | 2 +-
 tensor2tensor/utils/beam_search_test.py                     | 2 +-
 tensor2tensor/utils/bleu_hook.py                            | 2 +-
 tensor2tensor/utils/bleu_hook_test.py                       | 2 +-
 tensor2tensor/utils/checkpoint_compatibility_test.py        | 2 +-
 tensor2tensor/utils/cloud_mlengine.py                       | 2 +-
 tensor2tensor/utils/compute_video_metrics.py                | 2 +-
 tensor2tensor/utils/contrib.py                              | 2 +-
 tensor2tensor/utils/data_reader.py                          | 2 +-
 tensor2tensor/utils/data_reader_test.py                     | 2 +-
 tensor2tensor/utils/decoding.py                             | 2 +-
 tensor2tensor/utils/devices.py                              | 2 +-
 tensor2tensor/utils/diet.py                                 | 2 +-
 tensor2tensor/utils/diet_test.py                            | 2 +-
 tensor2tensor/utils/expert_utils.py                         | 2 +-
 tensor2tensor/utils/expert_utils_test.py                    | 2 +-
 tensor2tensor/utils/flags.py                                | 2 +-
 tensor2tensor/utils/get_rouge.py                            | 2 +-
 tensor2tensor/utils/hparam.py                               | 2 +-
 tensor2tensor/utils/hparam_test.py                          | 2 +-
 tensor2tensor/utils/hparams_lib.py                          | 2 +-
 tensor2tensor/utils/hparams_lib_test.py                     | 2 +-
 tensor2tensor/utils/learning_rate.py                        | 2 +-
 tensor2tensor/utils/metrics.py                              | 2 +-
 tensor2tensor/utils/metrics_hook.py                         | 2 +-
 tensor2tensor/utils/metrics_hook_test.py                    | 2 +-
 tensor2tensor/utils/metrics_test.py                         | 2 +-
 tensor2tensor/utils/misc_utils.py                           | 2 +-
 tensor2tensor/utils/misc_utils_test.py                      | 2 +-
 tensor2tensor/utils/mlperf_log.py                           | 2 +-
 tensor2tensor/utils/mlperf_tags.py                          | 2 +-
 tensor2tensor/utils/mtf_model.py                            | 2 +-
 tensor2tensor/utils/multistep_optimizer.py                  | 2 +-
 tensor2tensor/utils/multistep_optimizer_test.py             | 2 +-
 tensor2tensor/utils/multistep_with_adamoptimizer.py         | 2 +-
 tensor2tensor/utils/multistep_with_adamoptimizer_test.py    | 2 +-
 tensor2tensor/utils/optimize.py                             | 2 +-
 tensor2tensor/utils/optimize_test.py                        | 2 +-
 tensor2tensor/utils/partial_checkpoint_load_hook.py         | 2 +-
 tensor2tensor/utils/pruning_utils.py                        | 2 +-
 tensor2tensor/utils/quantization.py                         | 2 +-
 tensor2tensor/utils/registry.py                             | 2 +-
 tensor2tensor/utils/registry_test.py                        | 2 +-
 tensor2tensor/utils/restore_hook.py                         | 2 +-
 tensor2tensor/utils/rouge.py                                | 2 +-
 tensor2tensor/utils/rouge_test.py                           | 2 +-
 tensor2tensor/utils/sari_hook.py                            | 2 +-
 tensor2tensor/utils/sari_hook_test.py                       | 2 +-
 tensor2tensor/utils/scheduled_sampling.py                   | 2 +-
 tensor2tensor/utils/t2t_model.py                            | 2 +-
 tensor2tensor/utils/t2t_model_test.py                       | 2 +-
 tensor2tensor/utils/test_utils.py                           | 2 +-
 tensor2tensor/utils/test_utils_test.py                      | 2 +-
 tensor2tensor/utils/trainer_lib.py                          | 2 +-
 tensor2tensor/utils/trainer_lib_test.py                     | 2 +-
 tensor2tensor/utils/update_ops_hook.py                      | 2 +-
 tensor2tensor/utils/usr_dir.py                              | 2 +-
 tensor2tensor/utils/video/prediction2gif.py                 | 2 +-
 tensor2tensor/utils/video/reward_confusion.py               | 2 +-
 tensor2tensor/utils/video2gif.py                            | 2 +-
 tensor2tensor/utils/video_metrics.py                        | 2 +-
 tensor2tensor/utils/video_metrics_test.py                   | 2 +-
 tensor2tensor/utils/yellowfin.py                            | 2 +-
 tensor2tensor/utils/yellowfin_test.py                       | 2 +-
 tensor2tensor/visualization/__init__.py                     | 2 +-
 tensor2tensor/visualization/attention.py                    | 2 +-
 tensor2tensor/visualization/visualization.py                | 2 +-
 tensor2tensor/visualization/visualization_test.py           | 2 +-
 452 files changed, 456 insertions(+), 454 deletions(-)

diff --git a/tensor2tensor/__init__.py b/tensor2tensor/__init__.py
index 7b0a2368a..ffe5892e3 100644
--- a/tensor2tensor/__init__.py
+++ b/tensor2tensor/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/__init__.py b/tensor2tensor/bin/__init__.py
index 7b0a2368a..ffe5892e3 100644
--- a/tensor2tensor/bin/__init__.py
+++ b/tensor2tensor/bin/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/build_vocab.py b/tensor2tensor/bin/build_vocab.py
index da7d923dd..1b3847f06 100644
--- a/tensor2tensor/bin/build_vocab.py
+++ b/tensor2tensor/bin/build_vocab.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/make_tf_configs.py b/tensor2tensor/bin/make_tf_configs.py
index 6d482e330..e8ec19c08 100644
--- a/tensor2tensor/bin/make_tf_configs.py
+++ b/tensor2tensor/bin/make_tf_configs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 6a4c63b1f..125bed0eb 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_avg_all.py b/tensor2tensor/bin/t2t_avg_all.py
index 543c58ce2..4cd702105 100644
--- a/tensor2tensor/bin/t2t_avg_all.py
+++ b/tensor2tensor/bin/t2t_avg_all.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_bleu.py b/tensor2tensor/bin/t2t_bleu.py
index 1f07078ee..2015c79a4 100644
--- a/tensor2tensor/bin/t2t_bleu.py
+++ b/tensor2tensor/bin/t2t_bleu.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index 9826dc86f..d299bf27c 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 6a40a6a81..47e6344f3 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_distill.py b/tensor2tensor/bin/t2t_distill.py
index 383d2fc6b..ee1526995 100644
--- a/tensor2tensor/bin/t2t_distill.py
+++ b/tensor2tensor/bin/t2t_distill.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_eval.py b/tensor2tensor/bin/t2t_eval.py
index 536069a56..81ce7d443 100644
--- a/tensor2tensor/bin/t2t_eval.py
+++ b/tensor2tensor/bin/t2t_eval.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
index 385574f45..c4c6f411e 100644
--- a/tensor2tensor/bin/t2t_prune.py
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 60c7b4f11..4955cbceb 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_trainer_test.py b/tensor2tensor/bin/t2t_trainer_test.py
index 817e35c33..64cc7bc2c 100644
--- a/tensor2tensor/bin/t2t_trainer_test.py
+++ b/tensor2tensor/bin/t2t_trainer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_translate_all.py b/tensor2tensor/bin/t2t_translate_all.py
index 03eb56c0a..7c78045b8 100644
--- a/tensor2tensor/bin/t2t_translate_all.py
+++ b/tensor2tensor/bin/t2t_translate_all.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/__init__.py b/tensor2tensor/data_generators/__init__.py
index 7b0a2368a..ffe5892e3 100644
--- a/tensor2tensor/data_generators/__init__.py
+++ b/tensor2tensor/data_generators/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index 36b2df49b..91ffcfe1a 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math.py b/tensor2tensor/data_generators/algorithmic_math.py
index b71c40ab5..6b0329dca 100644
--- a/tensor2tensor/data_generators/algorithmic_math.py
+++ b/tensor2tensor/data_generators/algorithmic_math.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math_deepmind.py b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
index f8a44e4eb..86a6c0603 100644
--- a/tensor2tensor/data_generators/algorithmic_math_deepmind.py
+++ b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math_test.py b/tensor2tensor/data_generators/algorithmic_math_test.py
index c582e534f..b9d89a575 100644
--- a/tensor2tensor/data_generators/algorithmic_math_test.py
+++ b/tensor2tensor/data_generators/algorithmic_math_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math_two_variables.py b/tensor2tensor/data_generators/algorithmic_math_two_variables.py
index d64175e11..18d7704ca 100644
--- a/tensor2tensor/data_generators/algorithmic_math_two_variables.py
+++ b/tensor2tensor/data_generators/algorithmic_math_two_variables.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_test.py b/tensor2tensor/data_generators/algorithmic_test.py
index 523afcde6..99b1cb4d5 100644
--- a/tensor2tensor/data_generators/algorithmic_test.py
+++ b/tensor2tensor/data_generators/algorithmic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 920a2a6c8..4f6a1c494 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index b961c8f87..9a6341921 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/allen_brain_test.py b/tensor2tensor/data_generators/allen_brain_test.py
index 4c118ca9b..036e2a744 100644
--- a/tensor2tensor/data_generators/allen_brain_test.py
+++ b/tensor2tensor/data_generators/allen_brain_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio.py b/tensor2tensor/data_generators/audio.py
index 674fde652..10c6bd9f2 100644
--- a/tensor2tensor/data_generators/audio.py
+++ b/tensor2tensor/data_generators/audio.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio_encoder.py b/tensor2tensor/data_generators/audio_encoder.py
index d23e5274a..e8b76dde4 100644
--- a/tensor2tensor/data_generators/audio_encoder.py
+++ b/tensor2tensor/data_generators/audio_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio_test.py b/tensor2tensor/data_generators/audio_test.py
index ff3e21e09..52e9ca32e 100644
--- a/tensor2tensor/data_generators/audio_test.py
+++ b/tensor2tensor/data_generators/audio_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index 20ff4f1eb..1f37a8021 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index eecf55d66..9b5ca4885 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index 237bfe92e..9ca8a5877 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celeba_test.py b/tensor2tensor/data_generators/celeba_test.py
index 08d837728..11119714a 100644
--- a/tensor2tensor/data_generators/celeba_test.py
+++ b/tensor2tensor/data_generators/celeba_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celebahq.py b/tensor2tensor/data_generators/celebahq.py
index dd1425c16..cabd2e819 100644
--- a/tensor2tensor/data_generators/celebahq.py
+++ b/tensor2tensor/data_generators/celebahq.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index bf00c3b54..762c72f35 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cipher.py b/tensor2tensor/data_generators/cipher.py
index f7d5e880d..c1b6d2139 100644
--- a/tensor2tensor/data_generators/cipher.py
+++ b/tensor2tensor/data_generators/cipher.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cleaner_en_xx.py b/tensor2tensor/data_generators/cleaner_en_xx.py
index 141da35b5..4f2905a9d 100644
--- a/tensor2tensor/data_generators/cleaner_en_xx.py
+++ b/tensor2tensor/data_generators/cleaner_en_xx.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index 8ba01182a..08dbf5123 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cola.py b/tensor2tensor/data_generators/cola.py
index 12b8a2530..558f9b05d 100644
--- a/tensor2tensor/data_generators/cola.py
+++ b/tensor2tensor/data_generators/cola.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/common_voice.py b/tensor2tensor/data_generators/common_voice.py
index 57e1b8728..e35a8d50f 100644
--- a/tensor2tensor/data_generators/common_voice.py
+++ b/tensor2tensor/data_generators/common_voice.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/common_voice_test.py b/tensor2tensor/data_generators/common_voice_test.py
index a07e5c10e..675c71b5d 100644
--- a/tensor2tensor/data_generators/common_voice_test.py
+++ b/tensor2tensor/data_generators/common_voice_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/conll_ner.py b/tensor2tensor/data_generators/conll_ner.py
index 3c4095b95..0d80a7f91 100644
--- a/tensor2tensor/data_generators/conll_ner.py
+++ b/tensor2tensor/data_generators/conll_ner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/desc2code.py b/tensor2tensor/data_generators/desc2code.py
index e85b253d7..efd527ff0 100644
--- a/tensor2tensor/data_generators/desc2code.py
+++ b/tensor2tensor/data_generators/desc2code.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/desc2code_test.py b/tensor2tensor/data_generators/desc2code_test.py
index 3bc68db6e..dc8ea0cfd 100644
--- a/tensor2tensor/data_generators/desc2code_test.py
+++ b/tensor2tensor/data_generators/desc2code_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_abstract.py b/tensor2tensor/data_generators/dialog_abstract.py
index eb4408742..16ae87159 100644
--- a/tensor2tensor/data_generators/dialog_abstract.py
+++ b/tensor2tensor/data_generators/dialog_abstract.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_cornell.py b/tensor2tensor/data_generators/dialog_cornell.py
index 4423cd06c..a917e4e0b 100644
--- a/tensor2tensor/data_generators/dialog_cornell.py
+++ b/tensor2tensor/data_generators/dialog_cornell.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_dailydialog.py b/tensor2tensor/data_generators/dialog_dailydialog.py
index 80083f532..9bd62bce8 100644
--- a/tensor2tensor/data_generators/dialog_dailydialog.py
+++ b/tensor2tensor/data_generators/dialog_dailydialog.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_opensubtitles.py b/tensor2tensor/data_generators/dialog_opensubtitles.py
index 9ee843da4..2a8adc90b 100644
--- a/tensor2tensor/data_generators/dialog_opensubtitles.py
+++ b/tensor2tensor/data_generators/dialog_opensubtitles.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_personachat.py b/tensor2tensor/data_generators/dialog_personachat.py
index 719f71fd6..d6561c96c 100644
--- a/tensor2tensor/data_generators/dialog_personachat.py
+++ b/tensor2tensor/data_generators/dialog_personachat.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dna_encoder.py b/tensor2tensor/data_generators/dna_encoder.py
index 8db26b305..38dbfbb5e 100644
--- a/tensor2tensor/data_generators/dna_encoder.py
+++ b/tensor2tensor/data_generators/dna_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dna_encoder_test.py b/tensor2tensor/data_generators/dna_encoder_test.py
index ada9ff37a..65202ebd8 100644
--- a/tensor2tensor/data_generators/dna_encoder_test.py
+++ b/tensor2tensor/data_generators/dna_encoder_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/enwik8.py b/tensor2tensor/data_generators/enwik8.py
index 6f26b86e0..772380f34 100644
--- a/tensor2tensor/data_generators/enwik8.py
+++ b/tensor2tensor/data_generators/enwik8.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
index d9649e8b7..a9c0c2210 100644
--- a/tensor2tensor/data_generators/fsns.py
+++ b/tensor2tensor/data_generators/fsns.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
index 2e58ace24..e3a0e8a02 100644
--- a/tensor2tensor/data_generators/function_docstring.py
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gene_expression.py b/tensor2tensor/data_generators/gene_expression.py
index d2e35d8f9..cd0593142 100644
--- a/tensor2tensor/data_generators/gene_expression.py
+++ b/tensor2tensor/data_generators/gene_expression.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gene_expression_test.py b/tensor2tensor/data_generators/gene_expression_test.py
index a4820d8fb..74f7b5d66 100644
--- a/tensor2tensor/data_generators/gene_expression_test.py
+++ b/tensor2tensor/data_generators/gene_expression_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 4d4b1a324..7081e1f23 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/generator_utils_test.py b/tensor2tensor/data_generators/generator_utils_test.py
index 1778616cf..55ee6e7ad 100644
--- a/tensor2tensor/data_generators/generator_utils_test.py
+++ b/tensor2tensor/data_generators/generator_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index de51e5b59..9c4b9f49d 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index 9597c4abb..b8db3a19e 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index f7d897620..197134eb3 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ice_parsing.py b/tensor2tensor/data_generators/ice_parsing.py
index 3a2c7b49d..f97c7fde2 100644
--- a/tensor2tensor/data_generators/ice_parsing.py
+++ b/tensor2tensor/data_generators/ice_parsing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_lsun.py b/tensor2tensor/data_generators/image_lsun.py
index 63403cb0c..0755f2b78 100644
--- a/tensor2tensor/data_generators/image_lsun.py
+++ b/tensor2tensor/data_generators/image_lsun.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 909464616..04955f671 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_utils_test.py b/tensor2tensor/data_generators/image_utils_test.py
index f8578db99..43dfee321 100644
--- a/tensor2tensor/data_generators/image_utils_test.py
+++ b/tensor2tensor/data_generators/image_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index bd759d251..7c6c4d52f 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imagenet_test.py b/tensor2tensor/data_generators/imagenet_test.py
index 4ce37ac12..130283a8f 100644
--- a/tensor2tensor/data_generators/imagenet_test.py
+++ b/tensor2tensor/data_generators/imagenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imdb.py b/tensor2tensor/data_generators/imdb.py
index 03b158bb3..93b20bfd3 100644
--- a/tensor2tensor/data_generators/imdb.py
+++ b/tensor2tensor/data_generators/imdb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/inspect_tfrecord.py b/tensor2tensor/data_generators/inspect_tfrecord.py
index 8fa9299c4..0322a4bde 100644
--- a/tensor2tensor/data_generators/inspect_tfrecord.py
+++ b/tensor2tensor/data_generators/inspect_tfrecord.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lambada.py b/tensor2tensor/data_generators/lambada.py
index 3d76724b4..425ac2268 100644
--- a/tensor2tensor/data_generators/lambada.py
+++ b/tensor2tensor/data_generators/lambada.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index a4faa3d3d..cbf342e2f 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b.py b/tensor2tensor/data_generators/lm1b.py
index cd1a03a7f..75ab01632 100644
--- a/tensor2tensor/data_generators/lm1b.py
+++ b/tensor2tensor/data_generators/lm1b.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b_imdb.py b/tensor2tensor/data_generators/lm1b_imdb.py
index dc719fc5d..6d50ec7ed 100644
--- a/tensor2tensor/data_generators/lm1b_imdb.py
+++ b/tensor2tensor/data_generators/lm1b_imdb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b_mnli.py b/tensor2tensor/data_generators/lm1b_mnli.py
index 6cba7b9ae..c28c51ac3 100644
--- a/tensor2tensor/data_generators/lm1b_mnli.py
+++ b/tensor2tensor/data_generators/lm1b_mnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mnist.py b/tensor2tensor/data_generators/mnist.py
index 3570e3913..f49e53664 100644
--- a/tensor2tensor/data_generators/mnist.py
+++ b/tensor2tensor/data_generators/mnist.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/moving_mnist.py b/tensor2tensor/data_generators/moving_mnist.py
index c48673350..1f4600804 100644
--- a/tensor2tensor/data_generators/moving_mnist.py
+++ b/tensor2tensor/data_generators/moving_mnist.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mrpc.py b/tensor2tensor/data_generators/mrpc.py
index 8031a2a89..428632596 100644
--- a/tensor2tensor/data_generators/mrpc.py
+++ b/tensor2tensor/data_generators/mrpc.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index 22be4c4e0..c45c477e6 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mscoco_test.py b/tensor2tensor/data_generators/mscoco_test.py
index 6542a51dc..40f447711 100644
--- a/tensor2tensor/data_generators/mscoco_test.py
+++ b/tensor2tensor/data_generators/mscoco_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 3f4449778..e6ca24a31 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem_v2.py b/tensor2tensor/data_generators/multi_problem_v2.py
index d7fc02297..ee79b4e8f 100644
--- a/tensor2tensor/data_generators/multi_problem_v2.py
+++ b/tensor2tensor/data_generators/multi_problem_v2.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem_v2_test.py b/tensor2tensor/data_generators/multi_problem_v2_test.py
index e8f1c41dc..d28f298e4 100644
--- a/tensor2tensor/data_generators/multi_problem_v2_test.py
+++ b/tensor2tensor/data_generators/multi_problem_v2_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index 5d7a619fc..b09dcbe18 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ocr.py b/tensor2tensor/data_generators/ocr.py
index 6cb4b9a21..d01646318 100644
--- a/tensor2tensor/data_generators/ocr.py
+++ b/tensor2tensor/data_generators/ocr.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
index ef10510e0..f46eee526 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py b/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
index b0a757138..b654f5867 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/paraphrase_ms_coco.py b/tensor2tensor/data_generators/paraphrase_ms_coco.py
index 9ce52fda5..9713bfab1 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/paraphrase_ms_coco_test.py b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
index f58b33667..fb40709a2 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/pointer_generator_word.py b/tensor2tensor/data_generators/pointer_generator_word.py
index bbd018d69..c3f2549df 100644
--- a/tensor2tensor/data_generators/pointer_generator_word.py
+++ b/tensor2tensor/data_generators/pointer_generator_word.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index af211180d..68b86d653 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem_hparams.py b/tensor2tensor/data_generators/problem_hparams.py
index ffdfabe77..09a625be7 100644
--- a/tensor2tensor/data_generators/problem_hparams.py
+++ b/tensor2tensor/data_generators/problem_hparams.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 7751704a0..882f8c42a 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/program_search.py b/tensor2tensor/data_generators/program_search.py
index 4b1b46a7d..5fc1e96e2 100644
--- a/tensor2tensor/data_generators/program_search.py
+++ b/tensor2tensor/data_generators/program_search.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/program_search_test.py b/tensor2tensor/data_generators/program_search_test.py
index a4af895af..340b39391 100644
--- a/tensor2tensor/data_generators/program_search_test.py
+++ b/tensor2tensor/data_generators/program_search_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ptb.py b/tensor2tensor/data_generators/ptb.py
index a39d06afb..26186a7b6 100644
--- a/tensor2tensor/data_generators/ptb.py
+++ b/tensor2tensor/data_generators/ptb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/qnli.py b/tensor2tensor/data_generators/qnli.py
index af1765007..1d788eb91 100644
--- a/tensor2tensor/data_generators/qnli.py
+++ b/tensor2tensor/data_generators/qnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/quora_qpairs.py b/tensor2tensor/data_generators/quora_qpairs.py
index 447263db6..e7a3ff4be 100644
--- a/tensor2tensor/data_generators/quora_qpairs.py
+++ b/tensor2tensor/data_generators/quora_qpairs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/rte.py b/tensor2tensor/data_generators/rte.py
index e47536540..03236f54c 100644
--- a/tensor2tensor/data_generators/rte.py
+++ b/tensor2tensor/data_generators/rte.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/scitail.py b/tensor2tensor/data_generators/scitail.py
index 26f389bcf..0bd1865f0 100644
--- a/tensor2tensor/data_generators/scitail.py
+++ b/tensor2tensor/data_generators/scitail.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/seq2edits.py b/tensor2tensor/data_generators/seq2edits.py
index 3b1afe287..aecfa1897 100644
--- a/tensor2tensor/data_generators/seq2edits.py
+++ b/tensor2tensor/data_generators/seq2edits.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/snli.py b/tensor2tensor/data_generators/snli.py
index 977fbc005..b14215792 100644
--- a/tensor2tensor/data_generators/snli.py
+++ b/tensor2tensor/data_generators/snli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index da5911cb8..ae99187b7 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/squad.py b/tensor2tensor/data_generators/squad.py
index 646028a4e..2ade5f91a 100644
--- a/tensor2tensor/data_generators/squad.py
+++ b/tensor2tensor/data_generators/squad.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/sst_binary.py b/tensor2tensor/data_generators/sst_binary.py
index 14909adc4..9624a4153 100644
--- a/tensor2tensor/data_generators/sst_binary.py
+++ b/tensor2tensor/data_generators/sst_binary.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/stanford_nli.py b/tensor2tensor/data_generators/stanford_nli.py
index dbd7832e5..0820ed5ac 100644
--- a/tensor2tensor/data_generators/stanford_nli.py
+++ b/tensor2tensor/data_generators/stanford_nli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/style_transfer.py b/tensor2tensor/data_generators/style_transfer.py
index 291b87ff4..7f5c7e61e 100644
--- a/tensor2tensor/data_generators/style_transfer.py
+++ b/tensor2tensor/data_generators/style_transfer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/style_transfer_test.py b/tensor2tensor/data_generators/style_transfer_test.py
index 572a9915b..97e5267ae 100644
--- a/tensor2tensor/data_generators/style_transfer_test.py
+++ b/tensor2tensor/data_generators/style_transfer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/subject_verb_agreement.py b/tensor2tensor/data_generators/subject_verb_agreement.py
index 27ee080c6..8e5beb138 100644
--- a/tensor2tensor/data_generators/subject_verb_agreement.py
+++ b/tensor2tensor/data_generators/subject_verb_agreement.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 0ab6fe0b3..769dd9b77 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -624,7 +624,9 @@ def _escaped_token_to_subtoken_strings(self, escaped_token):
         # If there is no possible encoding of the escaped token then one of the
         # characters in the token is not in the alphabet. This should be
         # impossible and would be indicative of a bug.
-        assert False, "Token substring not found in subtoken vocabulary."
+        raise ValueError(
+            "Token substring '%s' not found in subtoken vocabulary." %
+            escaped_token)
 
     return ret
 
diff --git a/tensor2tensor/data_generators/text_encoder_build_subword.py b/tensor2tensor/data_generators/text_encoder_build_subword.py
index 97e52cd9b..03d437c36 100644
--- a/tensor2tensor/data_generators/text_encoder_build_subword.py
+++ b/tensor2tensor/data_generators/text_encoder_build_subword.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_encoder_test.py b/tensor2tensor/data_generators/text_encoder_test.py
index 6ca9c7648..1d0b489be 100644
--- a/tensor2tensor/data_generators/text_encoder_test.py
+++ b/tensor2tensor/data_generators/text_encoder_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -259,7 +259,7 @@ def test_raises_exception_when_not_encodable(self):
     original = "This has UPPER CASE letters that are out of alphabet"
 
     # Previously there was a bug which produced an infinite loop in this case.
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(ValueError):
       encoder.encode(original)
 
   def test_load_from_file(self):
diff --git a/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index dd0233ab9..361c69028 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_problems_test.py b/tensor2tensor/data_generators/text_problems_test.py
index dd897bf7e..4dd17bdf3 100644
--- a/tensor2tensor/data_generators/text_problems_test.py
+++ b/tensor2tensor/data_generators/text_problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index 87e176cf5..99a8e3b22 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_data_generator.py b/tensor2tensor/data_generators/timeseries_data_generator.py
index e5476423c..126250504 100644
--- a/tensor2tensor/data_generators/timeseries_data_generator.py
+++ b/tensor2tensor/data_generators/timeseries_data_generator.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_data_generator_test.py b/tensor2tensor/data_generators/timeseries_data_generator_test.py
index e32bc5a81..c9bbf95e0 100644
--- a/tensor2tensor/data_generators/timeseries_data_generator_test.py
+++ b/tensor2tensor/data_generators/timeseries_data_generator_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index 02157944a..aa6d20cbe 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/tokenizer.py b/tensor2tensor/data_generators/tokenizer.py
index 21b4dddaf..9a8c4c6df 100644
--- a/tensor2tensor/data_generators/tokenizer.py
+++ b/tensor2tensor/data_generators/tokenizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/tokenizer_test.py b/tensor2tensor/data_generators/tokenizer_test.py
index bd7272d5d..20da16115 100644
--- a/tensor2tensor/data_generators/tokenizer_test.py
+++ b/tensor2tensor/data_generators/tokenizer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/transduction_problems.py b/tensor2tensor/data_generators/transduction_problems.py
index e9d4c8626..02755088a 100644
--- a/tensor2tensor/data_generators/transduction_problems.py
+++ b/tensor2tensor/data_generators/transduction_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/transduction_problems_test.py b/tensor2tensor/data_generators/transduction_problems_test.py
index e35006f25..391835554 100644
--- a/tensor2tensor/data_generators/transduction_problems_test.py
+++ b/tensor2tensor/data_generators/transduction_problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index aae476ec0..1ae57641b 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_encs.py b/tensor2tensor/data_generators/translate_encs.py
index 657586440..f709ad811 100644
--- a/tensor2tensor/data_generators/translate_encs.py
+++ b/tensor2tensor/data_generators/translate_encs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_encs_cubbitt.py b/tensor2tensor/data_generators/translate_encs_cubbitt.py
index 7caf37303..acc1bf610 100644
--- a/tensor2tensor/data_generators/translate_encs_cubbitt.py
+++ b/tensor2tensor/data_generators/translate_encs_cubbitt.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index da3705f19..be5282cc1 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_ende_test.py b/tensor2tensor/data_generators/translate_ende_test.py
index 48c6753c6..23644c322 100644
--- a/tensor2tensor/data_generators/translate_ende_test.py
+++ b/tensor2tensor/data_generators/translate_ende_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enes.py b/tensor2tensor/data_generators/translate_enes.py
index 8feb459e0..f4a7f2199 100644
--- a/tensor2tensor/data_generators/translate_enes.py
+++ b/tensor2tensor/data_generators/translate_enes.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enet.py b/tensor2tensor/data_generators/translate_enet.py
index 8256b0d9f..f10677000 100644
--- a/tensor2tensor/data_generators/translate_enet.py
+++ b/tensor2tensor/data_generators/translate_enet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index 63a98bdab..c3dfc39a5 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enid.py b/tensor2tensor/data_generators/translate_enid.py
index bfa112dce..b2b800c93 100644
--- a/tensor2tensor/data_generators/translate_enid.py
+++ b/tensor2tensor/data_generators/translate_enid.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enmk.py b/tensor2tensor/data_generators/translate_enmk.py
index 2b4e74f9e..f41cc4e99 100644
--- a/tensor2tensor/data_generators/translate_enmk.py
+++ b/tensor2tensor/data_generators/translate_enmk.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enro.py b/tensor2tensor/data_generators/translate_enro.py
index ea2a9e6cf..a81894848 100644
--- a/tensor2tensor/data_generators/translate_enro.py
+++ b/tensor2tensor/data_generators/translate_enro.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_entn.py b/tensor2tensor/data_generators/translate_entn.py
index d1ea587fc..b622fda2b 100644
--- a/tensor2tensor/data_generators/translate_entn.py
+++ b/tensor2tensor/data_generators/translate_entn.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_envi.py b/tensor2tensor/data_generators/translate_envi.py
index fad6f850d..c3c68c58a 100644
--- a/tensor2tensor/data_generators/translate_envi.py
+++ b/tensor2tensor/data_generators/translate_envi.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index bbbf661dc..cd8d952e1 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_test.py b/tensor2tensor/data_generators/translate_test.py
index 3d5928bcf..1b6925cc2 100644
--- a/tensor2tensor/data_generators/translate_test.py
+++ b/tensor2tensor/data_generators/translate_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 18d621990..9a097e368 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 6e80102ec..9ad8c8dfc 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index 659a555da..9cb9b91db 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
index e027c4f7a..cd932845a 100644
--- a/tensor2tensor/data_generators/vqa.py
+++ b/tensor2tensor/data_generators/vqa.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/vqa_utils.py b/tensor2tensor/data_generators/vqa_utils.py
index a4e363645..e4557e00b 100644
--- a/tensor2tensor/data_generators/vqa_utils.py
+++ b/tensor2tensor/data_generators/vqa_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki.py b/tensor2tensor/data_generators/wiki.py
index affadd6de..e189a6955 100644
--- a/tensor2tensor/data_generators/wiki.py
+++ b/tensor2tensor/data_generators/wiki.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_lm.py b/tensor2tensor/data_generators/wiki_lm.py
index 7c2fceb3f..d29e27026 100644
--- a/tensor2tensor/data_generators/wiki_lm.py
+++ b/tensor2tensor/data_generators/wiki_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index 2f1cac9f1..501b8b92a 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_revision.py b/tensor2tensor/data_generators/wiki_revision.py
index 77942f84f..d65a21a7f 100644
--- a/tensor2tensor/data_generators/wiki_revision.py
+++ b/tensor2tensor/data_generators/wiki_revision.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index aae6010ab..3641e887a 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/__init__.py b/tensor2tensor/data_generators/wikisum/__init__.py
index 7b0a2368a..ffe5892e3 100644
--- a/tensor2tensor/data_generators/wikisum/__init__.py
+++ b/tensor2tensor/data_generators/wikisum/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/generate_vocab.py b/tensor2tensor/data_generators/wikisum/generate_vocab.py
index 6665d1c02..517de7fb1 100644
--- a/tensor2tensor/data_generators/wikisum/generate_vocab.py
+++ b/tensor2tensor/data_generators/wikisum/generate_vocab.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py b/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
index 7925715a6..be7aa0e80 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_web.py b/tensor2tensor/data_generators/wikisum/get_references_web.py
index d3deadd31..1a2afdda0 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py b/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
index 1fcce1010..86399371a 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/html.py b/tensor2tensor/data_generators/wikisum/html.py
index 38ecd194f..62c752df9 100644
--- a/tensor2tensor/data_generators/wikisum/html.py
+++ b/tensor2tensor/data_generators/wikisum/html.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/parallel_launch.py b/tensor2tensor/data_generators/wikisum/parallel_launch.py
index c69cb866e..13a50e355 100644
--- a/tensor2tensor/data_generators/wikisum/parallel_launch.py
+++ b/tensor2tensor/data_generators/wikisum/parallel_launch.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/produce_examples.py b/tensor2tensor/data_generators/wikisum/produce_examples.py
index 7a689978b..435151fb2 100644
--- a/tensor2tensor/data_generators/wikisum/produce_examples.py
+++ b/tensor2tensor/data_generators/wikisum/produce_examples.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/utils.py b/tensor2tensor/data_generators/wikisum/utils.py
index b3c124775..e60b579d9 100644
--- a/tensor2tensor/data_generators/wikisum/utils.py
+++ b/tensor2tensor/data_generators/wikisum/utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/utils_test.py b/tensor2tensor/data_generators/wikisum/utils_test.py
index 59de1a192..eb2c989c7 100644
--- a/tensor2tensor/data_generators/wikisum/utils_test.py
+++ b/tensor2tensor/data_generators/wikisum/utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/validate_data.py b/tensor2tensor/data_generators/wikisum/validate_data.py
index 41eb83713..4a8494c22 100644
--- a/tensor2tensor/data_generators/wikisum/validate_data.py
+++ b/tensor2tensor/data_generators/wikisum/validate_data.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/wikisum.py b/tensor2tensor/data_generators/wikisum/wikisum.py
index 293f45444..9003d5d8c 100644
--- a/tensor2tensor/data_generators/wikisum/wikisum.py
+++ b/tensor2tensor/data_generators/wikisum/wikisum.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikitext103.py b/tensor2tensor/data_generators/wikitext103.py
index 977eff511..3d9ea2a61 100644
--- a/tensor2tensor/data_generators/wikitext103.py
+++ b/tensor2tensor/data_generators/wikitext103.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wnli.py b/tensor2tensor/data_generators/wnli.py
index 42e246633..7d2af8b7e 100644
--- a/tensor2tensor/data_generators/wnli.py
+++ b/tensor2tensor/data_generators/wnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wsj_parsing.py b/tensor2tensor/data_generators/wsj_parsing.py
index 2a0a86650..f4f1ee892 100644
--- a/tensor2tensor/data_generators/wsj_parsing.py
+++ b/tensor2tensor/data_generators/wsj_parsing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/yelp_full.py b/tensor2tensor/data_generators/yelp_full.py
index 7aa130090..02d1b826d 100644
--- a/tensor2tensor/data_generators/yelp_full.py
+++ b/tensor2tensor/data_generators/yelp_full.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/yelp_polarity.py b/tensor2tensor/data_generators/yelp_polarity.py
index 5f16a16c0..60d0d7dcd 100644
--- a/tensor2tensor/data_generators/yelp_polarity.py
+++ b/tensor2tensor/data_generators/yelp_polarity.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/__init__.py b/tensor2tensor/envs/__init__.py
index 858a8238f..869af5e33 100644
--- a/tensor2tensor/envs/__init__.py
+++ b/tensor2tensor/envs/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 049a5c3c7..4efebddb6 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 9a67a7128..467422dc1 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index d9051356e..c592e6bc8 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_env_problem.py b/tensor2tensor/envs/gym_env_problem.py
index e4714a1ce..492df2d9b 100644
--- a/tensor2tensor/envs/gym_env_problem.py
+++ b/tensor2tensor/envs/gym_env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_env_problem_test.py b/tensor2tensor/envs/gym_env_problem_test.py
index 93d8093e3..0b4904a92 100644
--- a/tensor2tensor/envs/gym_env_problem_test.py
+++ b/tensor2tensor/envs/gym_env_problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_spaces_utils.py b/tensor2tensor/envs/gym_spaces_utils.py
index dbff5f646..c4256196f 100644
--- a/tensor2tensor/envs/gym_spaces_utils.py
+++ b/tensor2tensor/envs/gym_spaces_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_spaces_utils_test.py b/tensor2tensor/envs/gym_spaces_utils_test.py
index 2998d3e70..6234939b5 100644
--- a/tensor2tensor/envs/gym_spaces_utils_test.py
+++ b/tensor2tensor/envs/gym_spaces_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/mujoco_problems.py b/tensor2tensor/envs/mujoco_problems.py
index 0a8ddef97..51cb951a7 100644
--- a/tensor2tensor/envs/mujoco_problems.py
+++ b/tensor2tensor/envs/mujoco_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/mujoco_problems_test.py b/tensor2tensor/envs/mujoco_problems_test.py
index 813810260..bd4edce91 100644
--- a/tensor2tensor/envs/mujoco_problems_test.py
+++ b/tensor2tensor/envs/mujoco_problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
index 6ddb3caee..e08bd99f2 100644
--- a/tensor2tensor/envs/rendered_env_problem.py
+++ b/tensor2tensor/envs/rendered_env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/rendered_env_problem_test.py b/tensor2tensor/envs/rendered_env_problem_test.py
index 852fc5b6d..4a6273e0c 100644
--- a/tensor2tensor/envs/rendered_env_problem_test.py
+++ b/tensor2tensor/envs/rendered_env_problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env.py b/tensor2tensor/envs/tic_tac_toe_env.py
index 479717283..65f4d4d4e 100644
--- a/tensor2tensor/envs/tic_tac_toe_env.py
+++ b/tensor2tensor/envs/tic_tac_toe_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_problem.py b/tensor2tensor/envs/tic_tac_toe_env_problem.py
index d9aa5236e..c85301ca2 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_problem_test.py b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
index 2fce45651..78acab86e 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_test.py b/tensor2tensor/envs/tic_tac_toe_env_test.py
index 854e039d5..9a7f77d5b 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/time_step.py b/tensor2tensor/envs/time_step.py
index 165e05d04..675d5e4d6 100644
--- a/tensor2tensor/envs/time_step.py
+++ b/tensor2tensor/envs/time_step.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/time_step_test.py b/tensor2tensor/envs/time_step_test.py
index 54042d7a2..98476bb21 100644
--- a/tensor2tensor/envs/time_step_test.py
+++ b/tensor2tensor/envs/time_step_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index d56ab47e6..dcb9fcc92 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index dd2ee6d21..23d94dea8 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/__init__.py b/tensor2tensor/insights/__init__.py
index 7b0a2368a..ffe5892e3 100644
--- a/tensor2tensor/insights/__init__.py
+++ b/tensor2tensor/insights/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/graph.py b/tensor2tensor/insights/graph.py
index 28bb562e0..afa4b8452 100644
--- a/tensor2tensor/insights/graph.py
+++ b/tensor2tensor/insights/graph.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/query_processor.py b/tensor2tensor/insights/query_processor.py
index 74b716e01..a19213592 100644
--- a/tensor2tensor/insights/query_processor.py
+++ b/tensor2tensor/insights/query_processor.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/server.py b/tensor2tensor/insights/server.py
index 07f7b1dd8..23c34c485 100644
--- a/tensor2tensor/insights/server.py
+++ b/tensor2tensor/insights/server.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/transformer_model.py b/tensor2tensor/insights/transformer_model.py
index 267a0b299..d44703962 100644
--- a/tensor2tensor/insights/transformer_model.py
+++ b/tensor2tensor/insights/transformer_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/__init__.py b/tensor2tensor/layers/__init__.py
index 7b0a2368a..ffe5892e3 100644
--- a/tensor2tensor/layers/__init__.py
+++ b/tensor2tensor/layers/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/area_attention.py b/tensor2tensor/layers/area_attention.py
index c4c8fc430..46e6a695b 100644
--- a/tensor2tensor/layers/area_attention.py
+++ b/tensor2tensor/layers/area_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/area_attention_test.py b/tensor2tensor/layers/area_attention_test.py
index 78005b8ab..79191b07b 100644
--- a/tensor2tensor/layers/area_attention_test.py
+++ b/tensor2tensor/layers/area_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 027449854..2c6929402 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index ebc286c6d..dbee72c1b 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_audio.py b/tensor2tensor/layers/common_audio.py
index a4a2f493c..f090fc5d4 100644
--- a/tensor2tensor/layers/common_audio.py
+++ b/tensor2tensor/layers/common_audio.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index 057016dab..c3fd8cfbf 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 9aaabe480..754749620 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index 6b3d89bd9..739113c78 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index af1fcb83b..7bedb1af2 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index f649b24a5..530d9104d 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index b5119700f..0ef067f83 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index 7a39b3584..ed241b3ec 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index e8e1be203..d11d11994 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index c17f23bc4..283081388 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 3c8020306..fc0eaffe7 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
index 956e080da..178d880d2 100644
--- a/tensor2tensor/layers/latent_layers_test.py
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/message_passing_attention.py b/tensor2tensor/layers/message_passing_attention.py
index 9f9b5f14f..65c3912d5 100644
--- a/tensor2tensor/layers/message_passing_attention.py
+++ b/tensor2tensor/layers/message_passing_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index c22fb698c..c36a5db19 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 1d934ca4f..bbb17dfc2 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/ngram.py b/tensor2tensor/layers/ngram.py
index 46086aa13..26c08e241 100644
--- a/tensor2tensor/layers/ngram.py
+++ b/tensor2tensor/layers/ngram.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/ngram_test.py b/tensor2tensor/layers/ngram_test.py
index 1493b90c1..a7b8e8787 100644
--- a/tensor2tensor/layers/ngram_test.py
+++ b/tensor2tensor/layers/ngram_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_glow_layers.py b/tensor2tensor/layers/transformer_glow_layers.py
index 1687fc998..eb46c6f6a 100644
--- a/tensor2tensor/layers/transformer_glow_layers.py
+++ b/tensor2tensor/layers/transformer_glow_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_glow_layers_ops.py b/tensor2tensor/layers/transformer_glow_layers_ops.py
index 64a24eab7..3b7c0ee15 100644
--- a/tensor2tensor/layers/transformer_glow_layers_ops.py
+++ b/tensor2tensor/layers/transformer_glow_layers_ops.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_glow_layers_ops_test.py b/tensor2tensor/layers/transformer_glow_layers_ops_test.py
index 9e01c2c67..7c5404c7e 100644
--- a/tensor2tensor/layers/transformer_glow_layers_ops_test.py
+++ b/tensor2tensor/layers/transformer_glow_layers_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_glow_layers_test.py b/tensor2tensor/layers/transformer_glow_layers_test.py
index 92bdf28ae..7658aaaf4 100644
--- a/tensor2tensor/layers/transformer_glow_layers_test.py
+++ b/tensor2tensor/layers/transformer_glow_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index fec79a166..97739299e 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index aa0e8d908..cadee628b 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_memory_test.py b/tensor2tensor/layers/transformer_memory_test.py
index fa39591b4..a03769a12 100644
--- a/tensor2tensor/layers/transformer_memory_test.py
+++ b/tensor2tensor/layers/transformer_memory_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index a37c21801..182127b03 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/vqa_layers.py b/tensor2tensor/layers/vqa_layers.py
index 6e802d662..eed7b24aa 100644
--- a/tensor2tensor/layers/vqa_layers.py
+++ b/tensor2tensor/layers/vqa_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/__init__.py b/tensor2tensor/metrics/__init__.py
index 7b0a2368a..ffe5892e3 100644
--- a/tensor2tensor/metrics/__init__.py
+++ b/tensor2tensor/metrics/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/video_conditional_fvd.py b/tensor2tensor/metrics/video_conditional_fvd.py
index 3b6a0f058..6c29fd4a8 100644
--- a/tensor2tensor/metrics/video_conditional_fvd.py
+++ b/tensor2tensor/metrics/video_conditional_fvd.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/video_conditional_fvd_test.py b/tensor2tensor/metrics/video_conditional_fvd_test.py
index 41bb5ce3c..bf1743c45 100644
--- a/tensor2tensor/metrics/video_conditional_fvd_test.py
+++ b/tensor2tensor/metrics/video_conditional_fvd_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index 579d10c25..b06ecb04b 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/basic.py b/tensor2tensor/models/basic.py
index 3aaf8ee0c..424b3abe9 100644
--- a/tensor2tensor/models/basic.py
+++ b/tensor2tensor/models/basic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/basic_test.py b/tensor2tensor/models/basic_test.py
index 556e8acfa..c7098e5f0 100644
--- a/tensor2tensor/models/basic_test.py
+++ b/tensor2tensor/models/basic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/bytenet.py b/tensor2tensor/models/bytenet.py
index 6c2312b03..aac98cdd0 100644
--- a/tensor2tensor/models/bytenet.py
+++ b/tensor2tensor/models/bytenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/bytenet_test.py b/tensor2tensor/models/bytenet_test.py
index c8f07b7e0..1ba972e9f 100644
--- a/tensor2tensor/models/bytenet_test.py
+++ b/tensor2tensor/models/bytenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/distillation.py b/tensor2tensor/models/distillation.py
index 62fadd9f6..206fccd30 100644
--- a/tensor2tensor/models/distillation.py
+++ b/tensor2tensor/models/distillation.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index dd2053ba0..1240bd2de 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/evolved_transformer_test.py b/tensor2tensor/models/evolved_transformer_test.py
index 0d71b2f96..bf8c21f8a 100644
--- a/tensor2tensor/models/evolved_transformer_test.py
+++ b/tensor2tensor/models/evolved_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 67505ca41..885ec7db1 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index c27bd83e9..291a44b99 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_2d_test.py b/tensor2tensor/models/image_transformer_2d_test.py
index aa75032f9..56a00d4ef 100644
--- a/tensor2tensor/models/image_transformer_2d_test.py
+++ b/tensor2tensor/models/image_transformer_2d_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_test.py b/tensor2tensor/models/image_transformer_test.py
index 0c555f1c2..1f5c04395 100644
--- a/tensor2tensor/models/image_transformer_test.py
+++ b/tensor2tensor/models/image_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index f6d408ae3..90d4b67a7 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/lstm_test.py b/tensor2tensor/models/lstm_test.py
index a54018233..92615e167 100644
--- a/tensor2tensor/models/lstm_test.py
+++ b/tensor2tensor/models/lstm_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_image_transformer.py b/tensor2tensor/models/mtf_image_transformer.py
index fdcee7a23..c12b1b23f 100644
--- a/tensor2tensor/models/mtf_image_transformer.py
+++ b/tensor2tensor/models/mtf_image_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_image_transformer_test.py b/tensor2tensor/models/mtf_image_transformer_test.py
index aed3f0954..77a9f902e 100644
--- a/tensor2tensor/models/mtf_image_transformer_test.py
+++ b/tensor2tensor/models/mtf_image_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_resnet.py b/tensor2tensor/models/mtf_resnet.py
index 449629b7c..3c42f6eba 100644
--- a/tensor2tensor/models/mtf_resnet.py
+++ b/tensor2tensor/models/mtf_resnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 5ac5e091a..8e1ba2a7b 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index c94f2bed4..baf6d2c1a 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer_test.py b/tensor2tensor/models/mtf_transformer_test.py
index 36c469b9f..345bc8e30 100644
--- a/tensor2tensor/models/mtf_transformer_test.py
+++ b/tensor2tensor/models/mtf_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/__init__.py b/tensor2tensor/models/neural_architecture_search/__init__.py
index 4a53b5e30..aa007c728 100644
--- a/tensor2tensor/models/neural_architecture_search/__init__.py
+++ b/tensor2tensor/models/neural_architecture_search/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_layers.py b/tensor2tensor/models/neural_architecture_search/nas_layers.py
index 4e0afc3c2..d198a3d62 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_layers.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_layers_test.py b/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
index 34f109c03..a9db2679a 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_model.py b/tensor2tensor/models/neural_architecture_search/nas_model.py
index 40c114ef7..d3405250a 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_model_test.py b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
index d1bca5891..39a4d6a05 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model_test.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_assistant.py b/tensor2tensor/models/neural_assistant.py
index a2f02f675..c96c1e243 100644
--- a/tensor2tensor/models/neural_assistant.py
+++ b/tensor2tensor/models/neural_assistant.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_gpu.py b/tensor2tensor/models/neural_gpu.py
index bd5522ad4..964b00bf6 100644
--- a/tensor2tensor/models/neural_gpu.py
+++ b/tensor2tensor/models/neural_gpu.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_gpu_test.py b/tensor2tensor/models/neural_gpu_test.py
index 17dc8ebfa..6ae190e65 100644
--- a/tensor2tensor/models/neural_gpu_test.py
+++ b/tensor2tensor/models/neural_gpu_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/__init__.py b/tensor2tensor/models/research/__init__.py
index 7b0a2368a..ffe5892e3 100644
--- a/tensor2tensor/models/research/__init__.py
+++ b/tensor2tensor/models/research/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/adafactor_experiments.py b/tensor2tensor/models/research/adafactor_experiments.py
index 6875f45e7..69e63824c 100644
--- a/tensor2tensor/models/research/adafactor_experiments.py
+++ b/tensor2tensor/models/research/adafactor_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/aligned.py b/tensor2tensor/models/research/aligned.py
index 470e5b119..88510cff1 100644
--- a/tensor2tensor/models/research/aligned.py
+++ b/tensor2tensor/models/research/aligned.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/attention_lm.py b/tensor2tensor/models/research/attention_lm.py
index 6a442f1f1..1d71b2d97 100644
--- a/tensor2tensor/models/research/attention_lm.py
+++ b/tensor2tensor/models/research/attention_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/attention_lm_moe.py b/tensor2tensor/models/research/attention_lm_moe.py
index 581fdb407..fa7b28055 100644
--- a/tensor2tensor/models/research/attention_lm_moe.py
+++ b/tensor2tensor/models/research/attention_lm_moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 52a368f73..617523343 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index 0dc6f30c2..54a2c48f9 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index 6107a373c..3174e84c2 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/gene_expression.py b/tensor2tensor/models/research/gene_expression.py
index 04d365e10..e67eac9b3 100644
--- a/tensor2tensor/models/research/gene_expression.py
+++ b/tensor2tensor/models/research/gene_expression.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/gene_expression_test.py b/tensor2tensor/models/research/gene_expression_test.py
index 99f7e35d6..9a93f118a 100644
--- a/tensor2tensor/models/research/gene_expression_test.py
+++ b/tensor2tensor/models/research/gene_expression_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index 1647e520c..aa12689a4 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_init_hook.py b/tensor2tensor/models/research/glow_init_hook.py
index 824ef826b..5b1f18789 100644
--- a/tensor2tensor/models/research/glow_init_hook.py
+++ b/tensor2tensor/models/research/glow_init_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 6c4ed96f4..5b955cd35 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 5020ac430..c00d13827 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
index 1263075f9..90967272f 100644
--- a/tensor2tensor/models/research/glow_test.py
+++ b/tensor2tensor/models/research/glow_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/lm_experiments.py b/tensor2tensor/models/research/lm_experiments.py
index bd29e8b16..9e50090f9 100644
--- a/tensor2tensor/models/research/lm_experiments.py
+++ b/tensor2tensor/models/research/lm_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/moe.py b/tensor2tensor/models/research/moe.py
index 3a32429d4..e6195b22f 100644
--- a/tensor2tensor/models/research/moe.py
+++ b/tensor2tensor/models/research/moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/moe_experiments.py b/tensor2tensor/models/research/moe_experiments.py
index 879e80bd2..3bef4be56 100644
--- a/tensor2tensor/models/research/moe_experiments.py
+++ b/tensor2tensor/models/research/moe_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/multiquery_paper.py b/tensor2tensor/models/research/multiquery_paper.py
index 9cc7eb6e5..af79c2454 100644
--- a/tensor2tensor/models/research/multiquery_paper.py
+++ b/tensor2tensor/models/research/multiquery_paper.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/neural_stack.py b/tensor2tensor/models/research/neural_stack.py
index 7857a5c77..8ed4b14b0 100644
--- a/tensor2tensor/models/research/neural_stack.py
+++ b/tensor2tensor/models/research/neural_stack.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/neural_stack_test.py b/tensor2tensor/models/research/neural_stack_test.py
index a5497d52e..4571af15d 100644
--- a/tensor2tensor/models/research/neural_stack_test.py
+++ b/tensor2tensor/models/research/neural_stack_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/residual_shuffle_exchange.py b/tensor2tensor/models/research/residual_shuffle_exchange.py
index 5c70b7b50..217defb76 100644
--- a/tensor2tensor/models/research/residual_shuffle_exchange.py
+++ b/tensor2tensor/models/research/residual_shuffle_exchange.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index 8964f1be5..a9e848f0e 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/shuffle_network.py b/tensor2tensor/models/research/shuffle_network.py
index 39f817fa1..4e09a2fde 100644
--- a/tensor2tensor/models/research/shuffle_network.py
+++ b/tensor2tensor/models/research/shuffle_network.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/similarity_transformer.py b/tensor2tensor/models/research/similarity_transformer.py
index 5da061d3e..973e711f6 100644
--- a/tensor2tensor/models/research/similarity_transformer.py
+++ b/tensor2tensor/models/research/similarity_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/super_lm.py b/tensor2tensor/models/research/super_lm.py
index 05f7bfe21..90eafc309 100644
--- a/tensor2tensor/models/research/super_lm.py
+++ b/tensor2tensor/models/research/super_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_aux.py b/tensor2tensor/models/research/transformer_aux.py
index bd1ca52a6..d2e60d3a7 100644
--- a/tensor2tensor/models/research/transformer_aux.py
+++ b/tensor2tensor/models/research/transformer_aux.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_aux_test.py b/tensor2tensor/models/research/transformer_aux_test.py
index ee94837ad..2a9b3eefc 100644
--- a/tensor2tensor/models/research/transformer_aux_test.py
+++ b/tensor2tensor/models/research/transformer_aux_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index 993aa08d8..b5ba233e3 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index c521947ea..0438ebb17 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_parallel.py b/tensor2tensor/models/research/transformer_parallel.py
index 5adcac39a..aff322027 100644
--- a/tensor2tensor/models/research/transformer_parallel.py
+++ b/tensor2tensor/models/research/transformer_parallel.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_revnet.py b/tensor2tensor/models/research/transformer_revnet.py
index dd792bd0c..03c6d0a89 100644
--- a/tensor2tensor/models/research/transformer_revnet.py
+++ b/tensor2tensor/models/research/transformer_revnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_revnet_test.py b/tensor2tensor/models/research/transformer_revnet_test.py
index 09d9b4c4c..f2dd3c599 100644
--- a/tensor2tensor/models/research/transformer_revnet_test.py
+++ b/tensor2tensor/models/research/transformer_revnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_seq2edits.py b/tensor2tensor/models/research/transformer_seq2edits.py
index 81e9a4c48..bb18ba845 100644
--- a/tensor2tensor/models/research/transformer_seq2edits.py
+++ b/tensor2tensor/models/research/transformer_seq2edits.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_sketch.py b/tensor2tensor/models/research/transformer_sketch.py
index fa0e83173..4f27d4fc1 100644
--- a/tensor2tensor/models/research/transformer_sketch.py
+++ b/tensor2tensor/models/research/transformer_sketch.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_symshard.py b/tensor2tensor/models/research/transformer_symshard.py
index caed36a0b..030bcfb02 100644
--- a/tensor2tensor/models/research/transformer_symshard.py
+++ b/tensor2tensor/models/research/transformer_symshard.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 11e52d106..8df933e97 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae_flow_prior.py b/tensor2tensor/models/research/transformer_vae_flow_prior.py
index 1de0dd4fc..9a254567b 100644
--- a/tensor2tensor/models/research/transformer_vae_flow_prior.py
+++ b/tensor2tensor/models/research/transformer_vae_flow_prior.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py b/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
index 04deb42ca..824ad6e99 100644
--- a/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
+++ b/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae_test.py b/tensor2tensor/models/research/transformer_vae_test.py
index 0ab3d4e54..27a1bef72 100644
--- a/tensor2tensor/models/research/transformer_vae_test.py
+++ b/tensor2tensor/models/research/transformer_vae_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index bbc1d7277..2fd339659 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
index bc1856d11..350c8481a 100644
--- a/tensor2tensor/models/research/universal_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index c3fc4669b..353acf811 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index cacbd1623..8d745c7f9 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_attention_test.py b/tensor2tensor/models/research/vqa_attention_test.py
index d70a7d78c..34116bf71 100644
--- a/tensor2tensor/models/research/vqa_attention_test.py
+++ b/tensor2tensor/models/research/vqa_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_recurrent_self_attention.py b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
index 7857789b8..0ae482051 100644
--- a/tensor2tensor/models/research/vqa_recurrent_self_attention.py
+++ b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
index bbb9d7250..d10842a01 100644
--- a/tensor2tensor/models/research/vqa_self_attention.py
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index 02174b79a..288241bca 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/resnet_test.py b/tensor2tensor/models/resnet_test.py
index 796458bdb..14d881fe5 100644
--- a/tensor2tensor/models/resnet_test.py
+++ b/tensor2tensor/models/resnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/revnet.py b/tensor2tensor/models/revnet.py
index 8c923f9fe..d20c6b378 100644
--- a/tensor2tensor/models/revnet.py
+++ b/tensor2tensor/models/revnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/revnet_test.py b/tensor2tensor/models/revnet_test.py
index 026406f15..af3ddce7c 100644
--- a/tensor2tensor/models/revnet_test.py
+++ b/tensor2tensor/models/revnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index 423c7d76c..47224bcd1 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/slicenet.py b/tensor2tensor/models/slicenet.py
index 3d4930d84..c3e0af60a 100644
--- a/tensor2tensor/models/slicenet.py
+++ b/tensor2tensor/models/slicenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/slicenet_test.py b/tensor2tensor/models/slicenet_test.py
index 5a522c88b..cc9101125 100644
--- a/tensor2tensor/models/slicenet_test.py
+++ b/tensor2tensor/models/slicenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/text_cnn.py b/tensor2tensor/models/text_cnn.py
index 5e299557d..bfc791e30 100644
--- a/tensor2tensor/models/text_cnn.py
+++ b/tensor2tensor/models/text_cnn.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 77321d213..0cc8b5f2a 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 720295904..fe6fc0ceb 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index 576b707bf..03c1a9354 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/__init__.py b/tensor2tensor/models/video/__init__.py
index 7b0a2368a..ffe5892e3 100644
--- a/tensor2tensor/models/video/__init__.py
+++ b/tensor2tensor/models/video/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index 93d7f3dcf..b860518e3 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/base_vae.py b/tensor2tensor/models/video/base_vae.py
index 5922ce63b..56d7b1325 100644
--- a/tensor2tensor/models/video/base_vae.py
+++ b/tensor2tensor/models/video/base_vae.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 603ebc254..3832c2cbd 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 03048882d..10431a978 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic_test.py b/tensor2tensor/models/video/basic_deterministic_test.py
index 6b2794cc0..c66aebb4c 100644
--- a/tensor2tensor/models/video/basic_deterministic_test.py
+++ b/tensor2tensor/models/video/basic_deterministic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_recurrent.py b/tensor2tensor/models/video/basic_recurrent.py
index b2bcf5d98..bd6436e30 100644
--- a/tensor2tensor/models/video/basic_recurrent.py
+++ b/tensor2tensor/models/video/basic_recurrent.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_recurrent_test.py b/tensor2tensor/models/video/basic_recurrent_test.py
index 4af4fa672..7618d19f4 100644
--- a/tensor2tensor/models/video/basic_recurrent_test.py
+++ b/tensor2tensor/models/video/basic_recurrent_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 14510950f..4c4426b5e 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_stochastic_test.py b/tensor2tensor/models/video/basic_stochastic_test.py
index e0dbba8f4..03c68358e 100644
--- a/tensor2tensor/models/video/basic_stochastic_test.py
+++ b/tensor2tensor/models/video/basic_stochastic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index 49d7889cb..30636709b 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/emily_test.py b/tensor2tensor/models/video/emily_test.py
index 269dae573..9b72810a2 100644
--- a/tensor2tensor/models/video/emily_test.py
+++ b/tensor2tensor/models/video/emily_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index 72198cbb8..79c640e27 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
index 4adc7f93d..7d8bec8e7 100644
--- a/tensor2tensor/models/video/epva_params.py
+++ b/tensor2tensor/models/video/epva_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index 8d2b1890f..d3719798a 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv3d_test.py b/tensor2tensor/models/video/nfg_conv3d_test.py
index 8fa3f7f09..3b7d68b64 100644
--- a/tensor2tensor/models/video/nfg_conv3d_test.py
+++ b/tensor2tensor/models/video/nfg_conv3d_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv_lstm_test.py b/tensor2tensor/models/video/nfg_conv_lstm_test.py
index 76fde1aed..ae300b0f6 100644
--- a/tensor2tensor/models/video/nfg_conv_lstm_test.py
+++ b/tensor2tensor/models/video/nfg_conv_lstm_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv_test.py b/tensor2tensor/models/video/nfg_conv_test.py
index af87739b2..1c1b1ca5b 100644
--- a/tensor2tensor/models/video/nfg_conv_test.py
+++ b/tensor2tensor/models/video/nfg_conv_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_interpolate.py b/tensor2tensor/models/video/nfg_interpolate.py
index c0498ca96..8cbcfe54c 100644
--- a/tensor2tensor/models/video/nfg_interpolate.py
+++ b/tensor2tensor/models/video/nfg_interpolate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_test_utils.py b/tensor2tensor/models/video/nfg_test_utils.py
index 23baa7285..9ca0d4a4c 100644
--- a/tensor2tensor/models/video/nfg_test_utils.py
+++ b/tensor2tensor/models/video/nfg_test_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_uncond_test.py b/tensor2tensor/models/video/nfg_uncond_test.py
index 5472b5454..44763d7ff 100644
--- a/tensor2tensor/models/video/nfg_uncond_test.py
+++ b/tensor2tensor/models/video/nfg_uncond_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index 6c52944b4..48accd97f 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index bf249b2df..c1a34a5fa 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp_test.py b/tensor2tensor/models/video/savp_test.py
index 5baacaea8..b353f1f70 100644
--- a/tensor2tensor/models/video/savp_test.py
+++ b/tensor2tensor/models/video/savp_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index b26c8f403..debcee4ac 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 732d25da4..0f21c9248 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p_test.py b/tensor2tensor/models/video/sv2p_test.py
index 18b9835b3..dacbcb87b 100644
--- a/tensor2tensor/models/video/sv2p_test.py
+++ b/tensor2tensor/models/video/sv2p_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/tests_utils.py b/tensor2tensor/models/video/tests_utils.py
index 2028f280a..e2487bd80 100644
--- a/tensor2tensor/models/video/tests_utils.py
+++ b/tensor2tensor/models/video/tests_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/xception.py b/tensor2tensor/models/xception.py
index da9b05ad3..a5fb9c1cf 100644
--- a/tensor2tensor/models/xception.py
+++ b/tensor2tensor/models/xception.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/xception_test.py b/tensor2tensor/models/xception_test.py
index bc345e41f..90d49b8f2 100644
--- a/tensor2tensor/models/xception_test.py
+++ b/tensor2tensor/models/xception_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/problems.py b/tensor2tensor/problems.py
index 8a5ec1d42..f8432957d 100644
--- a/tensor2tensor/problems.py
+++ b/tensor2tensor/problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/problems_colab.py b/tensor2tensor/problems_colab.py
index 1b4f3e2f5..67563b387 100644
--- a/tensor2tensor/problems_colab.py
+++ b/tensor2tensor/problems_colab.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/problems_test.py b/tensor2tensor/problems_test.py
index f05538953..b4cc7ccb4 100644
--- a/tensor2tensor/problems_test.py
+++ b/tensor2tensor/problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/__init__.py b/tensor2tensor/rl/__init__.py
index 7b0a2368a..ffe5892e3 100644
--- a/tensor2tensor/rl/__init__.py
+++ b/tensor2tensor/rl/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/batch_dqn_agent_test.py b/tensor2tensor/rl/batch_dqn_agent_test.py
index 4a3fae838..c751e02a9 100644
--- a/tensor2tensor/rl/batch_dqn_agent_test.py
+++ b/tensor2tensor/rl/batch_dqn_agent_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/batch_runner_test.py b/tensor2tensor/rl/batch_runner_test.py
index efc9a5f4c..61c04bac5 100644
--- a/tensor2tensor/rl/batch_runner_test.py
+++ b/tensor2tensor/rl/batch_runner_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/datagen_with_agent.py b/tensor2tensor/rl/datagen_with_agent.py
index e41889f9f..a001ffc9f 100644
--- a/tensor2tensor/rl/datagen_with_agent.py
+++ b/tensor2tensor/rl/datagen_with_agent.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index 10976d41b..d8a4c2728 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/__init__.py b/tensor2tensor/rl/envs/__init__.py
index 7b0a2368a..ffe5892e3 100644
--- a/tensor2tensor/rl/envs/__init__.py
+++ b/tensor2tensor/rl/envs/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index 19db0274e..c16e209e6 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index c662d9519..984edf08e 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index c1215ba2d..a679f99ce 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index 16de501fe..c045ee7f6 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index 72f120932..a3c0c1918 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index ea38366e5..87335f83b 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/evaluator_test.py b/tensor2tensor/rl/evaluator_test.py
index c49999811..6f90ad36e 100644
--- a/tensor2tensor/rl/evaluator_test.py
+++ b/tensor2tensor/rl/evaluator_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index 81b9df7ed..3fb2a708e 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index a5066f3db..fece87b74 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index ba8c1db2e..22af30be1 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index 1648f6fc4..0fbc3270e 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
index f17f9429c..fb628a493 100644
--- a/tensor2tensor/rl/policy_learner.py
+++ b/tensor2tensor/rl/policy_learner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index bec485515..c458957cc 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index bd2642176..3abf4255d 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/restarter.py b/tensor2tensor/rl/restarter.py
index 223f7fe9b..e0dcbbe0d 100644
--- a/tensor2tensor/rl/restarter.py
+++ b/tensor2tensor/rl/restarter.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/restarter_test.py b/tensor2tensor/rl/restarter_test.py
index 71e1a292c..38fbe9eab 100644
--- a/tensor2tensor/rl/restarter_test.py
+++ b/tensor2tensor/rl/restarter_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index eaabb7b7f..b3eee7232 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index b07ea6835..c1ddb4d41 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_agent_only.py b/tensor2tensor/rl/trainer_model_based_agent_only.py
index 85c2ccdc7..804f0264b 100644
--- a/tensor2tensor/rl/trainer_model_based_agent_only.py
+++ b/tensor2tensor/rl/trainer_model_based_agent_only.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index f314db913..91c3d08a2 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_recurrent_test.py b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
index 1ace96b42..7e7e8b268 100644
--- a/tensor2tensor/rl/trainer_model_based_recurrent_test.py
+++ b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_stochastic_test.py b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
index 355175c74..dd978ac98 100644
--- a/tensor2tensor/rl/trainer_model_based_stochastic_test.py
+++ b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_sv2p_test.py b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
index 14cb10d1e..24610310f 100644
--- a/tensor2tensor/rl/trainer_model_based_sv2p_test.py
+++ b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_test.py b/tensor2tensor/rl/trainer_model_based_test.py
index 3eb71091b..2600e3a25 100644
--- a/tensor2tensor/rl/trainer_model_based_test.py
+++ b/tensor2tensor/rl/trainer_model_based_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 964cbaa4f..3bd83b027 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index 8cb73f173..b98f4c658 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free_tictactoe_test.py b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
index 2a8bae6c1..66d545257 100644
--- a/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
+++ b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/__init__.py b/tensor2tensor/serving/__init__.py
index 7b0a2368a..ffe5892e3 100644
--- a/tensor2tensor/serving/__init__.py
+++ b/tensor2tensor/serving/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 8aa0b3d10..754df03ed 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/query.py b/tensor2tensor/serving/query.py
index e5d557eb1..7910d30f7 100644
--- a/tensor2tensor/serving/query.py
+++ b/tensor2tensor/serving/query.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 3dc3d2fee..821f7c086 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/test_data/example_usr_dir/__init__.py b/tensor2tensor/test_data/example_usr_dir/__init__.py
index aa146a2b5..3f8cf243b 100644
--- a/tensor2tensor/test_data/example_usr_dir/__init__.py
+++ b/tensor2tensor/test_data/example_usr_dir/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/test_data/example_usr_dir/my_submodule.py b/tensor2tensor/test_data/example_usr_dir/my_submodule.py
index 88f8730e3..36609fa4a 100644
--- a/tensor2tensor/test_data/example_usr_dir/my_submodule.py
+++ b/tensor2tensor/test_data/example_usr_dir/my_submodule.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/__init__.py b/tensor2tensor/utils/__init__.py
index 7b0a2368a..ffe5892e3 100644
--- a/tensor2tensor/utils/__init__.py
+++ b/tensor2tensor/utils/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index 1383fc271..c4b5039cd 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/adafactor_test.py b/tensor2tensor/utils/adafactor_test.py
index 56befa985..d6832fe7d 100644
--- a/tensor2tensor/utils/adafactor_test.py
+++ b/tensor2tensor/utils/adafactor_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/adv_attack_utils.py b/tensor2tensor/utils/adv_attack_utils.py
index b179812f1..c31d4bfc7 100644
--- a/tensor2tensor/utils/adv_attack_utils.py
+++ b/tensor2tensor/utils/adv_attack_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index ae65c143a..f91b96ccd 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index d4a11594e..2a2e6178f 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/beam_search_test.py b/tensor2tensor/utils/beam_search_test.py
index b1cc0ec2d..60fcc09ef 100644
--- a/tensor2tensor/utils/beam_search_test.py
+++ b/tensor2tensor/utils/beam_search_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/bleu_hook.py b/tensor2tensor/utils/bleu_hook.py
index a36f5e027..0e92b8e41 100644
--- a/tensor2tensor/utils/bleu_hook.py
+++ b/tensor2tensor/utils/bleu_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/bleu_hook_test.py b/tensor2tensor/utils/bleu_hook_test.py
index 477df16dc..4f59ad1dd 100644
--- a/tensor2tensor/utils/bleu_hook_test.py
+++ b/tensor2tensor/utils/bleu_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/checkpoint_compatibility_test.py b/tensor2tensor/utils/checkpoint_compatibility_test.py
index b7a4a7bfe..0676ddc10 100644
--- a/tensor2tensor/utils/checkpoint_compatibility_test.py
+++ b/tensor2tensor/utils/checkpoint_compatibility_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 40165b736..26301b7e3 100644
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/compute_video_metrics.py b/tensor2tensor/utils/compute_video_metrics.py
index 5dfa2ba2a..dbb793dbc 100644
--- a/tensor2tensor/utils/compute_video_metrics.py
+++ b/tensor2tensor/utils/compute_video_metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/contrib.py b/tensor2tensor/utils/contrib.py
index f5b7883bb..efa22c9f6 100644
--- a/tensor2tensor/utils/contrib.py
+++ b/tensor2tensor/utils/contrib.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index 47891ec8b..aea085f28 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/data_reader_test.py b/tensor2tensor/utils/data_reader_test.py
index 2a4b2c75a..3d98d314d 100644
--- a/tensor2tensor/utils/data_reader_test.py
+++ b/tensor2tensor/utils/data_reader_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index a84749413..7a7d2b130 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index f03ef7ada..b595d7a3c 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/diet.py b/tensor2tensor/utils/diet.py
index 018063f9f..9039ced6a 100644
--- a/tensor2tensor/utils/diet.py
+++ b/tensor2tensor/utils/diet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/diet_test.py b/tensor2tensor/utils/diet_test.py
index fe808d53b..ef7c5cf68 100644
--- a/tensor2tensor/utils/diet_test.py
+++ b/tensor2tensor/utils/diet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index 26282a92c..ddb4139d8 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/expert_utils_test.py b/tensor2tensor/utils/expert_utils_test.py
index deb8581c3..79607b158 100644
--- a/tensor2tensor/utils/expert_utils_test.py
+++ b/tensor2tensor/utils/expert_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index 44e43fa62..c0f3a7079 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/get_rouge.py b/tensor2tensor/utils/get_rouge.py
index d896e21af..614e5e37b 100644
--- a/tensor2tensor/utils/get_rouge.py
+++ b/tensor2tensor/utils/get_rouge.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparam.py b/tensor2tensor/utils/hparam.py
index 6e09ed3e5..f198a5f66 100644
--- a/tensor2tensor/utils/hparam.py
+++ b/tensor2tensor/utils/hparam.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparam_test.py b/tensor2tensor/utils/hparam_test.py
index 3c31efdf6..01778191c 100644
--- a/tensor2tensor/utils/hparam_test.py
+++ b/tensor2tensor/utils/hparam_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparams_lib.py b/tensor2tensor/utils/hparams_lib.py
index 868f6081f..4882cd03f 100644
--- a/tensor2tensor/utils/hparams_lib.py
+++ b/tensor2tensor/utils/hparams_lib.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparams_lib_test.py b/tensor2tensor/utils/hparams_lib_test.py
index a4d1f394d..a068f910d 100644
--- a/tensor2tensor/utils/hparams_lib_test.py
+++ b/tensor2tensor/utils/hparams_lib_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index 362356b69..6151ad014 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index d47cbf08e..0da9889c0 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_hook.py b/tensor2tensor/utils/metrics_hook.py
index 3b738b82d..93fa924f4 100644
--- a/tensor2tensor/utils/metrics_hook.py
+++ b/tensor2tensor/utils/metrics_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_hook_test.py b/tensor2tensor/utils/metrics_hook_test.py
index 97efab7bf..a2c9d290c 100644
--- a/tensor2tensor/utils/metrics_hook_test.py
+++ b/tensor2tensor/utils/metrics_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index 6ffafd8b1..b9081edbd 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/misc_utils.py b/tensor2tensor/utils/misc_utils.py
index d26a586e8..b5da5d26a 100644
--- a/tensor2tensor/utils/misc_utils.py
+++ b/tensor2tensor/utils/misc_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/misc_utils_test.py b/tensor2tensor/utils/misc_utils_test.py
index b05a822bc..ccb453bb1 100644
--- a/tensor2tensor/utils/misc_utils_test.py
+++ b/tensor2tensor/utils/misc_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mlperf_log.py b/tensor2tensor/utils/mlperf_log.py
index c8c08f075..8fa42d802 100644
--- a/tensor2tensor/utils/mlperf_log.py
+++ b/tensor2tensor/utils/mlperf_log.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mlperf_tags.py b/tensor2tensor/utils/mlperf_tags.py
index bf57c8201..76f852185 100644
--- a/tensor2tensor/utils/mlperf_tags.py
+++ b/tensor2tensor/utils/mlperf_tags.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mtf_model.py b/tensor2tensor/utils/mtf_model.py
index 7a2f2412d..9f3e707b8 100644
--- a/tensor2tensor/utils/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
index 43172d0dc..438d4cc64 100644
--- a/tensor2tensor/utils/multistep_optimizer.py
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_optimizer_test.py b/tensor2tensor/utils/multistep_optimizer_test.py
index cf5995473..9d3f6990e 100644
--- a/tensor2tensor/utils/multistep_optimizer_test.py
+++ b/tensor2tensor/utils/multistep_optimizer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_with_adamoptimizer.py b/tensor2tensor/utils/multistep_with_adamoptimizer.py
index 792028e99..8f00cfbdf 100644
--- a/tensor2tensor/utils/multistep_with_adamoptimizer.py
+++ b/tensor2tensor/utils/multistep_with_adamoptimizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_with_adamoptimizer_test.py b/tensor2tensor/utils/multistep_with_adamoptimizer_test.py
index 6ac29809a..a76bdaca1 100644
--- a/tensor2tensor/utils/multistep_with_adamoptimizer_test.py
+++ b/tensor2tensor/utils/multistep_with_adamoptimizer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 078de3e03..6a11c943a 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/optimize_test.py b/tensor2tensor/utils/optimize_test.py
index ab4233a5b..938cd5c3d 100644
--- a/tensor2tensor/utils/optimize_test.py
+++ b/tensor2tensor/utils/optimize_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/partial_checkpoint_load_hook.py b/tensor2tensor/utils/partial_checkpoint_load_hook.py
index 5f436fdf4..a1242ab76 100644
--- a/tensor2tensor/utils/partial_checkpoint_load_hook.py
+++ b/tensor2tensor/utils/partial_checkpoint_load_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/pruning_utils.py b/tensor2tensor/utils/pruning_utils.py
index 476c41b56..d22b1537c 100644
--- a/tensor2tensor/utils/pruning_utils.py
+++ b/tensor2tensor/utils/pruning_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/quantization.py b/tensor2tensor/utils/quantization.py
index 849d8b5a3..d5c920b79 100644
--- a/tensor2tensor/utils/quantization.py
+++ b/tensor2tensor/utils/quantization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index fd6ec5ec5..31b0e70d8 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index 5b8acab2f..bb81fcf75 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/restore_hook.py b/tensor2tensor/utils/restore_hook.py
index f4c5d86f7..28d959f0f 100644
--- a/tensor2tensor/utils/restore_hook.py
+++ b/tensor2tensor/utils/restore_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/rouge.py b/tensor2tensor/utils/rouge.py
index 706154e2c..246bf80f7 100644
--- a/tensor2tensor/utils/rouge.py
+++ b/tensor2tensor/utils/rouge.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/rouge_test.py b/tensor2tensor/utils/rouge_test.py
index e13872523..3b27391b4 100644
--- a/tensor2tensor/utils/rouge_test.py
+++ b/tensor2tensor/utils/rouge_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/sari_hook.py b/tensor2tensor/utils/sari_hook.py
index e10ac7fb9..1579ebe99 100644
--- a/tensor2tensor/utils/sari_hook.py
+++ b/tensor2tensor/utils/sari_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/sari_hook_test.py b/tensor2tensor/utils/sari_hook_test.py
index a063167e2..250d0d6a2 100644
--- a/tensor2tensor/utils/sari_hook_test.py
+++ b/tensor2tensor/utils/sari_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/scheduled_sampling.py b/tensor2tensor/utils/scheduled_sampling.py
index 09a8089a7..4be7c9e3c 100644
--- a/tensor2tensor/utils/scheduled_sampling.py
+++ b/tensor2tensor/utils/scheduled_sampling.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index b1f097f67..1f3726a9a 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index a6a9d7008..578d22588 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/test_utils.py b/tensor2tensor/utils/test_utils.py
index 5f9215609..2adf5cdc3 100644
--- a/tensor2tensor/utils/test_utils.py
+++ b/tensor2tensor/utils/test_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/test_utils_test.py b/tensor2tensor/utils/test_utils_test.py
index 2de94b877..ba8a3cdc0 100644
--- a/tensor2tensor/utils/test_utils_test.py
+++ b/tensor2tensor/utils/test_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index a3662e35e..3ebe09d1c 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index 8f7e69af2..2c319fed7 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/update_ops_hook.py b/tensor2tensor/utils/update_ops_hook.py
index 5136858ff..a76671dfe 100644
--- a/tensor2tensor/utils/update_ops_hook.py
+++ b/tensor2tensor/utils/update_ops_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/usr_dir.py b/tensor2tensor/utils/usr_dir.py
index bc23ee1d6..5ca60a883 100644
--- a/tensor2tensor/utils/usr_dir.py
+++ b/tensor2tensor/utils/usr_dir.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video/prediction2gif.py b/tensor2tensor/utils/video/prediction2gif.py
index db3b4e37d..26e404362 100644
--- a/tensor2tensor/utils/video/prediction2gif.py
+++ b/tensor2tensor/utils/video/prediction2gif.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video/reward_confusion.py b/tensor2tensor/utils/video/reward_confusion.py
index 41ba42d33..854c1f899 100644
--- a/tensor2tensor/utils/video/reward_confusion.py
+++ b/tensor2tensor/utils/video/reward_confusion.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video2gif.py b/tensor2tensor/utils/video2gif.py
index 13f7fdd43..7a87d7a42 100644
--- a/tensor2tensor/utils/video2gif.py
+++ b/tensor2tensor/utils/video2gif.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index c1b3c0929..aad285336 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video_metrics_test.py b/tensor2tensor/utils/video_metrics_test.py
index 455135d27..7cd22012f 100644
--- a/tensor2tensor/utils/video_metrics_test.py
+++ b/tensor2tensor/utils/video_metrics_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/yellowfin.py b/tensor2tensor/utils/yellowfin.py
index a7ec6bcb7..333d93c0a 100644
--- a/tensor2tensor/utils/yellowfin.py
+++ b/tensor2tensor/utils/yellowfin.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/yellowfin_test.py b/tensor2tensor/utils/yellowfin_test.py
index c06f4563b..09ff51123 100644
--- a/tensor2tensor/utils/yellowfin_test.py
+++ b/tensor2tensor/utils/yellowfin_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/__init__.py b/tensor2tensor/visualization/__init__.py
index 4a53b5e30..aa007c728 100644
--- a/tensor2tensor/visualization/__init__.py
+++ b/tensor2tensor/visualization/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/attention.py b/tensor2tensor/visualization/attention.py
index 498967326..aa90649b9 100644
--- a/tensor2tensor/visualization/attention.py
+++ b/tensor2tensor/visualization/attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/visualization.py b/tensor2tensor/visualization/visualization.py
index 514e61bd7..f3ff216d1 100644
--- a/tensor2tensor/visualization/visualization.py
+++ b/tensor2tensor/visualization/visualization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/visualization_test.py b/tensor2tensor/visualization/visualization_test.py
index 4770c1eb1..e7ad86171 100644
--- a/tensor2tensor/visualization/visualization_test.py
+++ b/tensor2tensor/visualization/visualization_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Tensor2Tensor Authors.
+# Copyright 2022 The Tensor2Tensor Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

From cf72b2827460df92faf8e1c337e04a13e21901e2 Mon Sep 17 00:00:00 2001
From: Sergei Lebedev <slebedev@google.com>
Date: Mon, 7 Mar 2022 16:47:00 -0800
Subject: [PATCH 2705/2720] Removed references to g-no-augmented-assignment in
 pylint directives

PiperOrigin-RevId: 433073495
---
 tensor2tensor/layers/discretization.py           | 2 --
 tensor2tensor/models/research/transformer_nat.py | 2 --
 2 files changed, 4 deletions(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index d11d11994..93ae86acd 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -689,10 +689,8 @@ def discrete_bottleneck(inputs,
           n = tf.reduce_sum(updated_ema_count_res, axis=-1, keep_dims=True)
           updated_ema_count_res = (
               (updated_ema_count_res + epsilon) / (n + 2**z_size * epsilon) * n)
-          # pylint: disable=g-no-augmented-assignment
           updated_ema_means_res = updated_ema_means_res / tf.expand_dims(
               updated_ema_count_res, axis=-1)
-          # pylint: enable=g-no-augmented-assignment
 
           with tf.control_dependencies([e_loss_res]):
             update_means_res = tf.assign(means[i],
diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 0438ebb17..1eec30cde 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -94,10 +94,8 @@ def vq_discrete_bottleneck(x, hparams):
   updated_ema_count = (
       (updated_ema_count + hparams.epsilon) /
       (n + bottleneck_size * hparams.epsilon) * n)
-  # pylint: disable=g-no-augmented-assignment
   updated_ema_means = updated_ema_means / tf.expand_dims(
       updated_ema_count, axis=-1)
-  # pylint: enable=g-no-augmented-assignment
   with tf.control_dependencies([e_loss]):
     update_means = tf.assign(means, updated_ema_means)
     with tf.control_dependencies([update_means]):

From 61415a53a47b1d88685b49e135aa3be5d637cadb Mon Sep 17 00:00:00 2001
From: Fabien Hertschuh <fhertschuh@google.com>
Date: Tue, 22 Mar 2022 13:38:01 -0700
Subject: [PATCH 2706/2720] Explicitly import estimator from tensorflow as a
 separate import instead of accessing it via tf.estimator and depend on the
 tensorflow estimator target.

PiperOrigin-RevId: 436558187
---
 tensor2tensor/envs/gym_env_problem_test.py            |  3 ++-
 .../models/neural_architecture_search/nas_model.py    | 11 ++++++-----
 .../neural_architecture_search/nas_model_test.py      |  3 ++-
 tensor2tensor/rl/envs/simulated_batch_env.py          |  3 ++-
 tensor2tensor/rl/rl_utils.py                          |  5 +++--
 tensor2tensor/serving/export.py                       |  5 +++--
 tensor2tensor/utils/video/prediction2gif.py           |  5 +++--
 tensor2tensor/utils/video/reward_confusion.py         |  5 +++--
 tensor2tensor/visualization/visualization.py          |  3 ++-
 9 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/tensor2tensor/envs/gym_env_problem_test.py b/tensor2tensor/envs/gym_env_problem_test.py
index 0b4904a92..5e9c48ea3 100644
--- a/tensor2tensor/envs/gym_env_problem_test.py
+++ b/tensor2tensor/envs/gym_env_problem_test.py
@@ -31,6 +31,7 @@
 from tensor2tensor.envs import gym_env_problem
 from tensor2tensor.layers import modalities
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class GymEnvProblemTest(tf.test.TestCase):
@@ -324,7 +325,7 @@ def reward_modality(self):
         dev_filenames, ep)
 
     # Count them using a tf.data.Dataset.
-    dev_dataset = ep.dataset(tf.estimator.ModeKeys.EVAL, data_dir=self.tmp_dir)
+    dev_dataset = ep.dataset(tf_estimator.ModeKeys.EVAL, data_dir=self.tmp_dir)
 
     last_timestep = -1
     dev_timesteps_ds = 0
diff --git a/tensor2tensor/models/neural_architecture_search/nas_model.py b/tensor2tensor/models/neural_architecture_search/nas_model.py
index d3405250a..741504df1 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model.py
@@ -40,6 +40,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 # Keys for the activation map.
@@ -358,7 +359,7 @@ def decode(self,
         save_weights_to=self.attention_weights)
 
     if (common_layers.is_xla_compiled() and
-        hparams.mode == tf.estimator.ModeKeys.TRAIN):
+        hparams.mode == tf_estimator.ModeKeys.TRAIN):
       # TPU does not react kindly to extra dimensions.
       return decoder_output
 
@@ -421,8 +422,8 @@ def _gpu_estimator_spec_eval(self, features, logits, labels, loss,
         eval_metrics[metric_name] = metric_fn(logits, features,
                                               features["targets"])
 
-    return tf.estimator.EstimatorSpec(
-        tf.estimator.ModeKeys.EVAL,
+    return tf_estimator.EstimatorSpec(
+        tf_estimator.ModeKeys.EVAL,
         predictions={"predictions": logits},
         eval_metric_ops=eval_metrics,
         loss=loss)
@@ -446,12 +447,12 @@ def _tpu_estimator_spec_eval(self, features, logits, labels, loss,
       # eval_metrics_fn. Here we add the labels to those arguments.
       logits.update({"labels": labels})
       return contrib.tpu().TPUEstimatorSpec(
-          tf.estimator.ModeKeys.EVAL,
+          tf_estimator.ModeKeys.EVAL,
           eval_metrics=(eval_metrics_fn, logits),
           loss=loss)
     else:
       return contrib.tpu().TPUEstimatorSpec(
-          tf.estimator.ModeKeys.EVAL,
+          tf_estimator.ModeKeys.EVAL,
           eval_metrics=(eval_metrics_fn, [logits, labels]),
           loss=loss)
 
diff --git a/tensor2tensor/models/neural_architecture_search/nas_model_test.py b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
index 39a4d6a05..f3b05f6ce 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model_test.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.models.neural_architecture_search import nas_layers as layers
 from tensor2tensor.models.neural_architecture_search import nas_model as translation_nas_net
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 _BATCH_SIZE = 5
 _INPUT_LENGTH = 5
@@ -185,7 +186,7 @@ def _test_model(self, model_cls, hparams):
         "target_space_id": tf.constant(1, dtype=tf.int32)
     }
 
-    model = model_cls(hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
+    model = model_cls(hparams, tf_estimator.ModeKeys.TRAIN, p_hparams)
     logits, _ = model(features)
     with self.test_session() as session:
       session.run(tf.global_variables_initializer())
diff --git a/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index a679f99ce..881432581 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -35,6 +35,7 @@
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 # Lazy load PIL.Image
@@ -140,7 +141,7 @@ def __init__(
     trainer_lib.add_problem_hparams(model_hparams, problem)
     model_hparams.force_full_predict = True
     self._model = registry.model(model_name)(
-        model_hparams, tf.estimator.ModeKeys.PREDICT
+        model_hparams, tf_estimator.ModeKeys.PREDICT
     )
 
     self.history_buffer = HistoryBuffer(
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index b3eee7232..d3ed8680c 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -40,6 +40,7 @@
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 def compute_mean_reward(rollouts, clipped):
@@ -119,7 +120,7 @@ def evaluate_all_configs(
 
 def evaluate_world_model(
     real_env, hparams, world_model_dir, debug_video_path,
-    split=tf.estimator.ModeKeys.EVAL,
+    split=tf_estimator.ModeKeys.EVAL,
 ):
   """Evaluate the world model (reward accuracy)."""
   frame_stack_size = hparams.frame_stack_size
@@ -339,7 +340,7 @@ def choose_subsequence():
 def make_initial_frame_chooser(
     real_env, frame_stack_size, simulation_random_starts,
     simulation_flip_first_random_for_beginning,
-    split=tf.estimator.ModeKeys.TRAIN,
+    split=tf_estimator.ModeKeys.TRAIN,
 ):
   """Make frame chooser.
 
diff --git a/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 754df03ed..53d656f62 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -26,6 +26,7 @@
 from tensor2tensor.utils import usr_dir
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 import tensorflow_hub as hub
 
 FLAGS = tf.flags.FLAGS
@@ -155,7 +156,7 @@ def hub_module_fn():
     # we must do a copy of the features, as the model_fn can add additional
     # entries there (like hyperparameter settings etc).
     original_features = features.copy()
-    spec = model_fn(features, labels=None, mode=tf.estimator.ModeKeys.PREDICT)
+    spec = model_fn(features, labels=None, mode=tf_estimator.ModeKeys.PREDICT)
 
     hub.add_signature(
         inputs=original_features,
@@ -204,7 +205,7 @@ def main(_):
 
   estimator = create_estimator(run_config, hparams)
 
-  exporter = tf.estimator.FinalExporter(
+  exporter = tf_estimator.FinalExporter(
       "exporter",
       lambda: problem.serving_input_fn(hparams, decode_hparams, FLAGS.use_tpu),
       as_text=FLAGS.as_text)
diff --git a/tensor2tensor/utils/video/prediction2gif.py b/tensor2tensor/utils/video/prediction2gif.py
index 26e404362..50cf20662 100644
--- a/tensor2tensor/utils/video/prediction2gif.py
+++ b/tensor2tensor/utils/video/prediction2gif.py
@@ -43,6 +43,7 @@
 from tensor2tensor.utils import usr_dir
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 mpl.use("Agg")
 flags = tf.flags
@@ -81,7 +82,7 @@ def main(_):
     frame_shape += [hparams.problem.num_channels]
 
   dataset = registry.problem(FLAGS.problem).dataset(
-      tf.estimator.ModeKeys.TRAIN,
+      tf_estimator.ModeKeys.TRAIN,
       shuffle_files=True,
       data_dir=os.path.expanduser(FLAGS.data_dir),
       hparams=hparams)
@@ -103,7 +104,7 @@ def main(_):
     }
   # Create model.
   model_cls = registry.model(FLAGS.model)
-  model = model_cls(hparams, tf.estimator.ModeKeys.PREDICT)
+  model = model_cls(hparams, tf_estimator.ModeKeys.PREDICT)
   prediction_ops = model.infer(placeholders)
 
   states_q = Queue(maxsize=hparams.video_num_input_frames)
diff --git a/tensor2tensor/utils/video/reward_confusion.py b/tensor2tensor/utils/video/reward_confusion.py
index 854c1f899..11f69e6de 100644
--- a/tensor2tensor/utils/video/reward_confusion.py
+++ b/tensor2tensor/utils/video/reward_confusion.py
@@ -37,6 +37,7 @@
 from tensor2tensor.utils import usr_dir
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 flags = tf.flags
 FLAGS = flags.FLAGS
@@ -64,7 +65,7 @@ def main(_):
   # Iterating over dev/test partition of the data.
   # Change the data partition if necessary.
   dataset = registry.problem(FLAGS.problem).dataset(
-      tf.estimator.ModeKeys.PREDICT,
+      tf_estimator.ModeKeys.PREDICT,
       shuffle_files=False,
       hparams=hparams)
 
@@ -74,7 +75,7 @@ def main(_):
 
   # Creat model
   model_cls = registry.model(FLAGS.model)
-  model = model_cls(hparams, tf.estimator.ModeKeys.PREDICT)
+  model = model_cls(hparams, tf_estimator.ModeKeys.PREDICT)
   prediction_ops = model.infer(input_data)
 
   # Confusion Matrix
diff --git a/tensor2tensor/visualization/visualization.py b/tensor2tensor/visualization/visualization.py
index f3ff216d1..8c1914a37 100644
--- a/tensor2tensor/visualization/visualization.py
+++ b/tensor2tensor/visualization/visualization.py
@@ -27,6 +27,7 @@
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 EOS_ID = 1
 
@@ -138,7 +139,7 @@ def build_model(hparams_set, model_name, data_dir, problem_name, beam_size=1):
   hparams = trainer_lib.create_hparams(
       hparams_set, data_dir=data_dir, problem_name=problem_name)
   translate_model = registry.model(model_name)(
-      hparams, tf.estimator.ModeKeys.EVAL)
+      hparams, tf_estimator.ModeKeys.EVAL)
 
   inputs = tf.placeholder(tf.int32, shape=(1, None, 1, 1), name="inputs")
   targets = tf.placeholder(tf.int32, shape=(1, None, 1, 1), name="targets")

From c81d7700d7519fcf5bb87081ae117e351687ab37 Mon Sep 17 00:00:00 2001
From: Fabien Hertschuh <fhertschuh@google.com>
Date: Tue, 22 Mar 2022 13:38:29 -0700
Subject: [PATCH 2707/2720] Explicitly import estimator from tensorflow as a
 separate import instead of accessing it via tf.estimator and depend on the
 tensorflow estimator target.

PiperOrigin-RevId: 436558317
---
 tensor2tensor/models/research/aligned.py      |  3 +-
 .../models/research/attention_lm_moe.py       |  3 +-
 tensor2tensor/models/research/autoencoders.py | 49 ++++++++++---------
 .../models/research/autoencoders_test.py      |  3 +-
 .../models/research/gene_expression_test.py   |  3 +-
 tensor2tensor/models/research/glow.py         |  5 +-
 tensor2tensor/models/research/glow_ops.py     |  3 +-
 .../models/research/glow_ops_test.py          |  3 +-
 tensor2tensor/models/research/glow_test.py    |  9 ++--
 .../research/residual_shuffle_exchange.py     |  3 +-
 tensor2tensor/models/research/rl.py           |  3 +-
 .../models/research/shuffle_network.py        |  3 +-
 .../models/research/similarity_transformer.py |  3 +-
 tensor2tensor/models/research/super_lm.py     |  5 +-
 .../models/research/transformer_aux_test.py   |  3 +-
 .../models/research/transformer_moe.py        |  3 +-
 .../models/research/transformer_nat.py        |  5 +-
 .../models/research/transformer_parallel.py   | 13 ++---
 .../models/research/transformer_revnet.py     |  5 +-
 .../research/transformer_revnet_test.py       |  3 +-
 .../models/research/transformer_vae.py        | 11 +++--
 .../research/transformer_vae_flow_prior.py    |  9 ++--
 .../transformer_vae_flow_prior_ops.py         |  3 +-
 .../models/research/transformer_vae_test.py   |  3 +-
 .../research/universal_transformer_test.py    |  3 +-
 .../models/research/vqa_attention.py          |  5 +-
 .../models/research/vqa_attention_test.py     |  3 +-
 .../research/vqa_recurrent_self_attention.py  |  3 +-
 .../models/research/vqa_self_attention.py     |  7 +--
 29 files changed, 103 insertions(+), 74 deletions(-)

diff --git a/tensor2tensor/models/research/aligned.py b/tensor2tensor/models/research/aligned.py
index 88510cff1..d5071f64b 100644
--- a/tensor2tensor/models/research/aligned.py
+++ b/tensor2tensor/models/research/aligned.py
@@ -34,8 +34,9 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
-ModeKeys = tf.estimator.ModeKeys  # pylint: disable=invalid-name
+ModeKeys = tf_estimator.ModeKeys  # pylint: disable=invalid-name
 
 
 def _should_preprocess(layer_type):
diff --git a/tensor2tensor/models/research/attention_lm_moe.py b/tensor2tensor/models/research/attention_lm_moe.py
index fa7b28055..2af44b896 100644
--- a/tensor2tensor/models/research/attention_lm_moe.py
+++ b/tensor2tensor/models/research/attention_lm_moe.py
@@ -36,9 +36,10 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
-ModeKeys = tf.estimator.ModeKeys  # pylint: disable=invalid-name
+ModeKeys = tf_estimator.ModeKeys  # pylint: disable=invalid-name
 
 
 class AttentionType(object):
diff --git a/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 617523343..3298bc123 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -29,6 +29,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 def reverse_gradient(x, lr=1.0):
@@ -97,7 +98,7 @@ def bottleneck(self, x):
     with tf.variable_scope("bottleneck"):
       hparams = self.hparams
       x = tf.layers.dense(x, hparams.bottleneck_bits, name="bottleneck")
-      if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+      if hparams.mode == tf_estimator.ModeKeys.TRAIN:
         noise = 2.0 * tf.random_uniform(common_layers.shape_list(x)) - 1.0
         return tf.tanh(x) + noise * hparams.bottleneck_noise, 0.0
       return tf.tanh(x), 0.0
@@ -158,7 +159,7 @@ def decoder(self, x, encoder_layers):
 
   def gumbel_sample(self, reconstr_gan):
     hparams = self.hparams
-    is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+    is_training = hparams.mode == tf_estimator.ModeKeys.TRAIN
     vocab_size = self._problem_hparams.vocab_size["targets"]
     if hasattr(self._hparams, "vocab_divisor"):
       vocab_size += (-vocab_size) % self._hparams.vocab_divisor
@@ -181,13 +182,13 @@ def gumbel_sample(self, reconstr_gan):
 
   def body(self, features):
     hparams = self.hparams
-    is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+    is_training = hparams.mode == tf_estimator.ModeKeys.TRAIN
     vocab_size = self._problem_hparams.vocab_size["targets"]
     if hasattr(self._hparams, "vocab_divisor"):
       vocab_size += (-vocab_size) % self._hparams.vocab_divisor
     encoder_layers = None
     self.is1d = hparams.sample_width == 1
-    if (hparams.mode != tf.estimator.ModeKeys.PREDICT
+    if (hparams.mode != tf_estimator.ModeKeys.PREDICT
         or self._encode_on_predict):
       labels = features["targets_raw"]
       labels_shape = common_layers.shape_list(labels)
@@ -252,7 +253,7 @@ def body(self, features):
 
     # Cut to the right size and mix before returning.
     res = x
-    if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+    if hparams.mode != tf_estimator.ModeKeys.PREDICT:
       res = x[:, :shape[1], :shape[2], :]
 
     # Final dense layer.
@@ -264,7 +265,7 @@ def body(self, features):
     ]
     res = tf.reshape(res, output_shape)
 
-    if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+    if hparams.mode == tf_estimator.ModeKeys.PREDICT:
       if hparams.use_vq_loss:
         (reconstr, _, _, _, _) = discretization.vq_loss(res, labels, vocab_size)
       else:
@@ -284,7 +285,7 @@ def body(self, features):
       vq_temperature = hparams.vq_temperature / common_layers.inverse_exp_decay(
           hparams.gan_codes_warmup_steps * 1.2,
           min_value=hparams.vq_temperature * 2)
-      if hparams.mode != tf.estimator.ModeKeys.TRAIN:
+      if hparams.mode != tf_estimator.ModeKeys.TRAIN:
         vq_temperature = None
       with tf.variable_scope("vq_loss"):
         (reconstr, _, target_codes, code_loss,
@@ -471,7 +472,7 @@ def body(self, features):
     # Prepare inputs for autoregressive modes.
     if common_layers.shape_list(features["targets"])[1] == 1:
       # This happens on the first step of predicitions.
-      assert hparams.mode == tf.estimator.ModeKeys.PREDICT
+      assert hparams.mode == tf_estimator.ModeKeys.PREDICT
       targets = tf.zeros_like(basic_result)
     targets = self.embed(targets)
     if hparams.autoregressive_gumbel_sample:
@@ -483,14 +484,14 @@ def body(self, features):
     basic1d = tf.reshape(basic_result, [shape[0], -1, shape[-1]])
     targets = tf.reshape(targets, common_layers.shape_list(basic_result))
     # During autoregressive inference, don't resample.
-    if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+    if hparams.mode == tf_estimator.ModeKeys.PREDICT:
       if hasattr(hparams, "sampled_basic1d_tensor"):
         basic1d = hparams.sampled_basic1d_tensor
       else:
         hparams.sampled_basic1d_tensor = basic1d
     # Sometimes it's useful to look at non-autoregressive evals.
     targets_dropout = targets
-    if (hparams.mode == tf.estimator.ModeKeys.EVAL and
+    if (hparams.mode == tf_estimator.ModeKeys.EVAL and
         hparams.autoregressive_eval_pure_autoencoder):
       targets_dropout = tf.zeros_like(basic_result)
     # Now combine the basic reconstruction with shifted targets.
@@ -590,7 +591,7 @@ class AutoencoderResidual(AutoencoderAutoregressive):
   """Residual autoencoder."""
 
   def dropout(self, x):
-    is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
+    is_training = self.hparams.mode == tf_estimator.ModeKeys.TRAIN
     hparams = self.hparams
     if hparams.dropout <= 0.0 or not is_training:
       return x
@@ -649,7 +650,7 @@ def encoder(self, x):
   def decoder(self, x, encoder_layers=None):
     with tf.variable_scope("decoder"):
       hparams = self.hparams
-      is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
+      is_training = self.hparams.mode == tf_estimator.ModeKeys.TRAIN
       kernel, strides = self._get_kernel_and_strides()
       residual_kernel = (hparams.residual_kernel_height,
                          hparams.residual_kernel_width)
@@ -719,7 +720,7 @@ def bottleneck(self, x):
     x_shape = common_layers.shape_list(x)
     with tf.variable_scope("vae"):
       mu = tf.layers.dense(x, z_size, name="mu")
-      if hparams.mode != tf.estimator.ModeKeys.TRAIN:
+      if hparams.mode != tf_estimator.ModeKeys.TRAIN:
         return mu, 0.0  # No sampling or kl loss on eval.
       log_sigma = tf.layers.dense(x, z_size, name="log_sigma")
       epsilon = tf.random_normal(x_shape[:-1] + [z_size])
@@ -751,12 +752,12 @@ def bottleneck(self, x):
     hparams = self.hparams
     x = tf.tanh(tf.layers.dense(x, hparams.bottleneck_bits, name="bottleneck"))
     d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
-    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    if hparams.mode == tf_estimator.ModeKeys.TRAIN:
       noise = tf.random_uniform(common_layers.shape_list(x))
       noise = 2.0 * tf.to_float(tf.less(hparams.bottleneck_noise, noise)) - 1.0
       d *= noise
     x = common_layers.mix(d, x, hparams.discretize_warmup_steps,
-                          hparams.mode == tf.estimator.ModeKeys.TRAIN)
+                          hparams.mode == tf_estimator.ModeKeys.TRAIN)
     return x, 0.0
 
   def sample(self, features=None, shape=None):
@@ -828,7 +829,7 @@ def bottleneck(self, x):  # pylint: disable=arguments-differ
     hparams.bottleneck_noise = 0.0  # We'll add noise below.
     x, loss = discretization.parametrized_bottleneck(x, hparams)
     hparams.bottleneck_noise = noise
-    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    if hparams.mode == tf_estimator.ModeKeys.TRAIN:
       # We want a number p such that p^bottleneck_bits = 1 - noise.
       # So log(p) * bottleneck_bits = log(noise)
       log_p = tf.log1p(-float(noise) / 2) / float(hparams.bottleneck_bits)
@@ -848,14 +849,14 @@ class AutoencoderDualDiscrete(AutoencoderResidualDiscrete):
   """Dual discrete autoencoder."""
 
   def body(self, features):
-    if self.hparams.mode != tf.estimator.ModeKeys.EVAL:
+    if self.hparams.mode != tf_estimator.ModeKeys.EVAL:
       t, i = features["targets_raw"], features["inputs_raw"]
       t, i = common_layers.pad_to_same_length(t, i)
       features["targets_raw"] = tf.concat([t, i], axis=0)
     return super(AutoencoderDualDiscrete, self).body(features)
 
   def embed(self, x, name="embedding"):
-    if self.hparams.mode == tf.estimator.ModeKeys.EVAL:
+    if self.hparams.mode == tf_estimator.ModeKeys.EVAL:
       return super(AutoencoderDualDiscrete, self).embed(x, name=name + "_t")
     xt, xi = tf.split(x, 2, axis=0)
     xte = super(AutoencoderDualDiscrete, self).embed(xt, name=name + "_t")
@@ -865,10 +866,10 @@ def embed(self, x, name="embedding"):
   def bottleneck(self, x):
     hparams = self.hparams
     b, _ = super(AutoencoderDualDiscrete, self).bottleneck(x)
-    if hparams.mode == tf.estimator.ModeKeys.EVAL:
+    if hparams.mode == tf_estimator.ModeKeys.EVAL:
       return b, 0.0
     bt, bi = tf.split(b, 2, axis=0)
-    if self.hparams.mode != tf.estimator.ModeKeys.TRAIN:
+    if self.hparams.mode != tf_estimator.ModeKeys.TRAIN:
       return tf.concat([bi, bi], axis=0), 0.0
     # Share the first hparams.bottleneck_shared_bits.
     shared = (bt + bi) / 2  # -1 if both -1, 1 if both were 1, 0 if disagree.
@@ -894,7 +895,7 @@ def bottleneck(self, x):
   def unbottleneck(self, b, res_size, reuse=None):
     x = super(AutoencoderDualDiscrete, self).unbottleneck(
         b, res_size, reuse=reuse)
-    if self.hparams.mode == tf.estimator.ModeKeys.EVAL:
+    if self.hparams.mode == tf_estimator.ModeKeys.EVAL:
       return tf.layers.dense(x, res_size, name="dual_unbottleneck_t")
     xt, xi = tf.split(x, 2, axis=0)
     xt = tf.layers.dense(xt, res_size, name="dual_unbottleneck_t")
@@ -984,8 +985,8 @@ def body(self, features):
     hparams = self.hparams
     num_stacks = hparams.num_hidden_layers
     hparams.num_hidden_layers = 1
-    is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
-    if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+    is_training = hparams.mode == tf_estimator.ModeKeys.TRAIN
+    if hparams.mode != tf_estimator.ModeKeys.PREDICT:
       x = features["targets"]
       shape = common_layers.shape_list(x)
       is1d = shape[2] == 1
@@ -1013,7 +1014,7 @@ def body(self, features):
       x = self.unbottleneck(b, res_size)
     # Run decoder.
     x = self.decoder(x)
-    if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+    if hparams.mode == tf_estimator.ModeKeys.PREDICT:
       return x
     # Cut to the right size and mix before returning.
     res = x[:, :shape[1], :shape[2], :]
diff --git a/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index 54a2c48f9..08f25a118 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -26,12 +26,13 @@
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class AutoencoderTest(tf.test.TestCase):
 
   def get_mnist_random_output(self, model_name, hparams_set=None,
-                              mode=tf.estimator.ModeKeys.TRAIN):
+                              mode=tf_estimator.ModeKeys.TRAIN):
     hparams_set = hparams_set or model_name
     x = np.random.randint(256, size=(1, 28, 28, 1))
     y = np.random.randint(10, size=(1, 1))
diff --git a/tensor2tensor/models/research/gene_expression_test.py b/tensor2tensor/models/research/gene_expression_test.py
index 9a93f118a..dfd0469fb 100644
--- a/tensor2tensor/models/research/gene_expression_test.py
+++ b/tensor2tensor/models/research/gene_expression_test.py
@@ -24,6 +24,7 @@
 from tensor2tensor.models.research import gene_expression
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 def gene_expression_conv_test():
@@ -53,7 +54,7 @@ def _test_model(self, hparams, model_cls):
     }
     p_hparams = hparams.problem_hparams
     logits, _ = model_cls(
-        hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)(features)
+        hparams, tf_estimator.ModeKeys.TRAIN, p_hparams)(features)
 
     with self.test_session() as sess:
       sess.run(tf.global_variables_initializer())
diff --git a/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index aa12689a4..b76aa973a 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -28,6 +28,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 arg_scope = contrib.framework().arg_scope
 add_arg_scope = contrib.framework().add_arg_scope
@@ -101,7 +102,7 @@ def temperature(self):
 
   @property
   def is_training(self):
-    return self.hparams.mode == tf.estimator.ModeKeys.TRAIN
+    return self.hparams.mode == tf_estimator.ModeKeys.TRAIN
 
   def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
     del args, kwargs
@@ -129,7 +130,7 @@ def create_init_batch(self, features):
       init_features: initialization features.
     """
     train_dataset = self.hparams.problem.dataset(
-        tf.estimator.ModeKeys.TRAIN, hparams=self.hparams)
+        tf_estimator.ModeKeys.TRAIN, hparams=self.hparams)
     train_dataset = train_dataset.batch(self.hparams.init_batch_size)
     train_dataset = self.init_preprocess(train_dataset)
     return train_dataset.make_one_shot_iterator().get_next()
diff --git a/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 5b955cd35..4c904c3cd 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -26,6 +26,7 @@
 from tensor2tensor.layers import common_video
 from tensor2tensor.utils import contrib
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 import tensorflow_probability as tfp
 
 arg_scope = contrib.framework().arg_scope
@@ -931,7 +932,7 @@ def noise_op(latents, hparams):
   Returns:
     latents: latents with isotropic gaussian noise appended.
   """
-  if hparams.latent_noise == 0 or hparams.mode != tf.estimator.ModeKeys.TRAIN:
+  if hparams.latent_noise == 0 or hparams.mode != tf_estimator.ModeKeys.TRAIN:
     return latents
   latent_shape = common_layers.shape_list(latents)
   return latents + tf.random_normal(latent_shape, stddev=hparams.latent_noise)
diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index c00d13827..10db10080 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -31,6 +31,7 @@
 from tensor2tensor.utils import contrib
 from tensor2tensor.utils import hparam
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 arg_scope = contrib.framework().arg_scope
 add_arg_scope = contrib.framework().add_arg_scope
@@ -40,7 +41,7 @@ class GlowOpsTest(parameterized.TestCase, tf.test.TestCase):
 
   def get_glow_hparams(self):
     hparams = glow.glow_hparams()
-    hparams.add_hparam("mode", tf.estimator.ModeKeys.TRAIN)
+    hparams.add_hparam("mode", tf_estimator.ModeKeys.TRAIN)
     hparams.add_hparam("num_cond_latents", 1)
     hparams.add_hparam("latent_architecture", "glow_resnet")
     # Use latent skip connections
diff --git a/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
index 90967272f..fa21e53e2 100644
--- a/tensor2tensor/models/research/glow_test.py
+++ b/tensor2tensor/models/research/glow_test.py
@@ -29,8 +29,9 @@
 from tensor2tensor.models.research import glow
 from tensor2tensor.utils import registry  # pylint: disable=unused-import
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
-MODES = tf.estimator.ModeKeys
+MODES = tf_estimator.ModeKeys
 
 
 class GlowModelTest(tf.test.TestCase):
@@ -53,7 +54,7 @@ def test_glow(self):
       hparams.data_dir = ''
       cifar_problem = problems.problem('image_cifar10_plain_random_shift')
       hparams.problem = cifar_problem
-      model = glow.Glow(hparams, tf.estimator.ModeKeys.TRAIN)
+      model = glow.Glow(hparams, tf_estimator.ModeKeys.TRAIN)
       train_dataset = cifar_problem.dataset(MODES.TRAIN)
       one_shot = train_dataset.make_one_shot_iterator()
       x_batch, y_batch = self.batch(one_shot)
@@ -86,7 +87,7 @@ def test_glow_inference(self):
     with tf.Graph().as_default():
       cifar_problem = problems.problem('image_cifar10_plain_random_shift')
       hparams.problem = cifar_problem
-      model = glow.Glow(hparams, tf.estimator.ModeKeys.TRAIN)
+      model = glow.Glow(hparams, tf_estimator.ModeKeys.TRAIN)
       train_dataset = cifar_problem.dataset(MODES.TRAIN)
       one_shot = train_dataset.make_one_shot_iterator()
       x_batch, y_batch = self.batch(one_shot)
@@ -110,7 +111,7 @@ def test_glow_inference(self):
     with tf.Graph().as_default():
       cifar_problem = problems.problem('image_cifar10_plain_random_shift')
       hparams.problem = cifar_problem
-      model = glow.Glow(hparams, tf.estimator.ModeKeys.PREDICT)
+      model = glow.Glow(hparams, tf_estimator.ModeKeys.PREDICT)
       test_dataset = cifar_problem.dataset(MODES.EVAL)
       one_shot = test_dataset.make_one_shot_iterator()
       x_batch, y_batch = self.batch(one_shot)
diff --git a/tensor2tensor/models/research/residual_shuffle_exchange.py b/tensor2tensor/models/research/residual_shuffle_exchange.py
index 217defb76..745537bd5 100644
--- a/tensor2tensor/models/research/residual_shuffle_exchange.py
+++ b/tensor2tensor/models/research/residual_shuffle_exchange.py
@@ -34,6 +34,7 @@
 from tensor2tensor.models.research.shuffle_network import ShuffleNetwork
 from tensor2tensor.utils import registry
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class LayerNormalization(tf.keras.layers.Layer):
@@ -177,7 +178,7 @@ def call(self, inputs, **kwargs):
 
     if self.dropout > 0:
       candidate = tf.nn.dropout(candidate, rate=self.dropout / n_bits)
-    if self.dropout != 0.0 and self.mode == tf.estimator.ModeKeys.TRAIN:
+    if self.dropout != 0.0 and self.mode == tf_estimator.ModeKeys.TRAIN:
       noise = tf.random_normal(tf.shape(candidate), mean=1.0, stddev=0.001)
       candidate = candidate * noise
 
diff --git a/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index a9e848f0e..ab1a5f2f0 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -40,6 +40,7 @@
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 import tensorflow_probability as tfp
 
 
@@ -319,7 +320,7 @@ def get_policy(observations, hparams, action_space,
   trainer_lib.add_problem_hparams(hparams, policy_problem)
   hparams.force_full_predict = True
   model = registry.model(hparams.policy_network)(
-      hparams, tf.estimator.ModeKeys.TRAIN
+      hparams, tf_estimator.ModeKeys.TRAIN
   )
   try:
     num_target_frames = hparams.video_num_target_frames
diff --git a/tensor2tensor/models/research/shuffle_network.py b/tensor2tensor/models/research/shuffle_network.py
index 4e09a2fde..dc284f87f 100644
--- a/tensor2tensor/models/research/shuffle_network.py
+++ b/tensor2tensor/models/research/shuffle_network.py
@@ -33,6 +33,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 def ror(x, n, p=1):
@@ -249,7 +250,7 @@ def __call__(self, inputs, residual_inputs):
 
     if self.dropout > 0:
       candidate = tf.nn.dropout(candidate, rate=self.dropout / self.n_bits)
-    if self.dropout != 0.0 and self.mode == tf.estimator.ModeKeys.TRAIN:
+    if self.dropout != 0.0 and self.mode == tf_estimator.ModeKeys.TRAIN:
       noise = tf.random_normal(tf.shape(candidate), mean=1.0, stddev=0.001)
       candidate = candidate * noise
 
diff --git a/tensor2tensor/models/research/similarity_transformer.py b/tensor2tensor/models/research/similarity_transformer.py
index 973e711f6..0e2216b93 100644
--- a/tensor2tensor/models/research/similarity_transformer.py
+++ b/tensor2tensor/models/research/similarity_transformer.py
@@ -20,6 +20,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 @registry.register_model
@@ -37,7 +38,7 @@ def top(self, body_output, _):
     return body_output
 
   def body(self, features):
-    if self.hparams.mode != tf.estimator.ModeKeys.PREDICT:
+    if self.hparams.mode != tf_estimator.ModeKeys.PREDICT:
       # In training mode we need to embed both the queries and the code
       # using the inputs and targets respectively.
       with tf.variable_scope('string_embedding'):
diff --git a/tensor2tensor/models/research/super_lm.py b/tensor2tensor/models/research/super_lm.py
index 90eafc309..aa6c39375 100644
--- a/tensor2tensor/models/research/super_lm.py
+++ b/tensor2tensor/models/research/super_lm.py
@@ -39,8 +39,9 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
-ModeKeys = tf.estimator.ModeKeys  # pylint: disable=invalid-name
+ModeKeys = tf_estimator.ModeKeys  # pylint: disable=invalid-name
 
 
 @registry.register_model
@@ -222,7 +223,7 @@ def _split(t):
         x, loss = mp(
             expert_utils.local_moe,
             x,
-            train=hparams.mode == tf.estimator.ModeKeys.TRAIN,
+            train=hparams.mode == tf_estimator.ModeKeys.TRAIN,
             expert_fn=expert_fn,
             num_experts=hparams.moe_num_experts,
             k=hparams.moe_k,
diff --git a/tensor2tensor/models/research/transformer_aux_test.py b/tensor2tensor/models/research/transformer_aux_test.py
index 2a9b3eefc..0e7bd6568 100644
--- a/tensor2tensor/models/research/transformer_aux_test.py
+++ b/tensor2tensor/models/research/transformer_aux_test.py
@@ -25,6 +25,7 @@
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.models.research import transformer_aux
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class TransformerAuxTest(parameterized.TestCase, tf.test.TestCase):
@@ -100,7 +101,7 @@ def test_transformer_aux_body(self):
         "target_space_id": tf.constant(1, dtype=tf.int32),
     }
     tf.train.create_global_step()
-    model = transformer_aux.TransformerAux(hparams, tf.estimator.ModeKeys.TRAIN,
+    model = transformer_aux.TransformerAux(hparams, tf_estimator.ModeKeys.TRAIN,
                                            p_hparams)
     logits, losses = model(features)
 
diff --git a/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index b5ba233e3..8a2e1b327 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -29,6 +29,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 # The transformer architecture can be defined using the layer_types hparams.
@@ -108,7 +109,7 @@ def decorated(x, *args, **kwargs):
         dp=dp,
     )
 
-    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    if hparams.mode == tf_estimator.ModeKeys.TRAIN:
 
       # Display the encoder-decoder architecture
       def print_layer(name, layers):
diff --git a/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 1eec30cde..29b56c8f6 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -25,6 +25,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 from tensorflow.python.training import moving_averages  # pylint: disable=g-direct-tensorflow-import
 
 
@@ -255,7 +256,7 @@ def ae_transformer_internal(inputs, targets, target_space, hparams, cache=None):
       max_targets_len_from_inputs,
       final_length_divisible_by=2**hparams.num_compress_steps)
   targets_c = compress(targets, hparams, "compress")
-  if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+  if hparams.mode != tf_estimator.ModeKeys.PREDICT:
     # Compress and bottleneck.
     latents_discrete_hot, extra_loss = vq_discrete_bottleneck(
         x=targets_c, hparams=hparams)
@@ -298,7 +299,7 @@ def ae_transformer_internal(inputs, targets, target_space, hparams, cache=None):
   masking *= common_layers.inverse_exp_decay(
       hparams.mask_startup_steps // 4)  # Not much at start.
   masking = tf.minimum(tf.maximum(masking, 0.0), 1.0)
-  if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+  if hparams.mode == tf_estimator.ModeKeys.PREDICT:
     masking = 1.0
   mask = tf.less(masking,
                  tf.random_uniform(common_layers.shape_list(targets)[:-1]))
diff --git a/tensor2tensor/models/research/transformer_parallel.py b/tensor2tensor/models/research/transformer_parallel.py
index aff322027..5909eee4c 100644
--- a/tensor2tensor/models/research/transformer_parallel.py
+++ b/tensor2tensor/models/research/transformer_parallel.py
@@ -25,6 +25,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 @registry.register_model
@@ -68,9 +69,9 @@ def body(self, features):
   def top(self, body_output, features):
     assert self._hparams.block_size > 0
 
-    if (self._hparams.mode == tf.estimator.ModeKeys.TRAIN or
-        self._hparams.mode == tf.estimator.ModeKeys.EVAL):
-      if self._hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    if (self._hparams.mode == tf_estimator.ModeKeys.TRAIN or
+        self._hparams.mode == tf_estimator.ModeKeys.EVAL):
+      if self._hparams.mode == tf_estimator.ModeKeys.TRAIN:
         features["block_index"] = tf.random_uniform(
             shape=[], minval=0, maxval=self._hparams.block_size, dtype=tf.int64)
       else:
@@ -94,8 +95,8 @@ def shift_left_4d(x, k):
         for i in range(self._hparams.block_size)
     ], axis=2)
 
-    if (self._hparams.mode == tf.estimator.ModeKeys.TRAIN or
-        self._hparams.mode == tf.estimator.ModeKeys.EVAL):
+    if (self._hparams.mode == tf_estimator.ModeKeys.TRAIN or
+        self._hparams.mode == tf_estimator.ModeKeys.EVAL):
       assert "block_index" in features
       k = features["block_index"]
       targets = targets[:, :, k:k + 1, :]
@@ -104,7 +105,7 @@ def shift_left_4d(x, k):
 
     loss = super(TransformerBlockParallel, self).loss(logits, features)
 
-    if self._hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    if self._hparams.mode == tf_estimator.ModeKeys.TRAIN:
       loss_num, loss_den = loss
       loss_val = loss_num / loss_den
       for i in range(self._hparams.block_size):
diff --git a/tensor2tensor/models/research/transformer_revnet.py b/tensor2tensor/models/research/transformer_revnet.py
index 03c6d0a89..21af70e10 100644
--- a/tensor2tensor/models/research/transformer_revnet.py
+++ b/tensor2tensor/models/research/transformer_revnet.py
@@ -25,6 +25,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 @registry.register_model
@@ -128,7 +129,7 @@ def g(x):
         g,
         num_layers=hparams.num_hidden_layers,
         f_side_input=[encoder_self_attention_bias],
-        is_training=hparams.mode == tf.estimator.ModeKeys.TRAIN)
+        is_training=hparams.mode == tf_estimator.ModeKeys.TRAIN)
     y = tf.concat([y1, y2], axis=-1)
 
   return common_layers.layer_preprocess(y, hparams)
@@ -209,7 +210,7 @@ def g(x):
             decoder_self_attention_bias, encoder_decoder_attention_bias,
             encoder_output
         ],
-        is_training=hparams.mode == tf.estimator.ModeKeys.TRAIN)
+        is_training=hparams.mode == tf_estimator.ModeKeys.TRAIN)
     y = tf.concat([y1, y2], axis=-1)
     return common_layers.layer_preprocess(y, hparams)
 
diff --git a/tensor2tensor/models/research/transformer_revnet_test.py b/tensor2tensor/models/research/transformer_revnet_test.py
index f2dd3c599..3c66ab616 100644
--- a/tensor2tensor/models/research/transformer_revnet_test.py
+++ b/tensor2tensor/models/research/transformer_revnet_test.py
@@ -24,6 +24,7 @@
 from tensor2tensor.models.research import transformer_revnet
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 def transformer_revnet_test():
@@ -57,7 +58,7 @@ def testTransformer(self):
         "target_space_id": tf.constant(1, dtype=tf.int32),
     }
     model = transformer_revnet.TransformerRevnet(
-        hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
+        hparams, tf_estimator.ModeKeys.TRAIN, p_hparams)
     logits, _ = model(features)
     grads = tf.gradients(
         tf.reduce_mean(logits), [features["inputs"]] + tf.global_variables())
diff --git a/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 8df933e97..021b361c3 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -38,6 +38,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 _DO_SUMMARIES = True
@@ -104,7 +105,7 @@ def top_k_softmax(x, k):
 def top_k_experts(x, k, hparams):
   x_shape = common_layers.shape_list(x)
   x_flat = tf.reshape(x, [-1, common_layers.shape_list(x)[-1]])
-  is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+  is_training = hparams.mode == tf_estimator.ModeKeys.TRAIN
   gates, load = expert_utils.noisy_top_k_gating(
       x_flat, 2 ** hparams.z_size, is_training, k)
   gates_shape = [x_shape[0], x_shape[1], x_shape[2], 2 ** hparams.z_size]
@@ -400,7 +401,7 @@ def ae_transformer_internal(inputs,
       targets_noisy = targets
 
     targets_c = compress(targets_noisy, inputs, False, hparams, "compress")
-    if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+    if hparams.mode != tf_estimator.ModeKeys.PREDICT:
       # Compress and bottleneck.
       latents_dense, latents_discrete, extra_loss, embed, neg_q_entropy = (
           hparams.bottleneck(inputs=targets_c,
@@ -410,7 +411,7 @@ def ae_transformer_internal(inputs,
       if _DO_SUMMARIES:
         tf.summary.histogram("b0", tf.reshape(latents_discrete[:, 0, :], [-1]))
       pc = common_layers.inverse_exp_decay(hparams.startup_steps)
-      pc = pc if hparams.mode == tf.estimator.ModeKeys.TRAIN else 1.0
+      pc = pc if hparams.mode == tf_estimator.ModeKeys.TRAIN else 1.0
       cond = tf.less(tf.random_uniform([batch_size]), pc)
       latents_dense = tf.where(cond, latents_dense, targets_c)
       # TODO(lukaszkaiser): return extra losses batchwise, multiply before mean.
@@ -449,7 +450,7 @@ def bn_inputs():
           return bn
         inputs_c = bn_inputs()
         ptc = 1.0 - common_layers.inverse_lin_decay(200000) * 0.5
-        ptc = ptc if hparams.mode == tf.estimator.ModeKeys.TRAIN else 1.0
+        ptc = ptc if hparams.mode == tf_estimator.ModeKeys.TRAIN else 1.0
         latents_dense = tf.where(tf.less(tf.random_uniform([batch_size]), ptc),
                                  latents_dense, inputs_c)
     else:
@@ -497,7 +498,7 @@ def bn_inputs():
       masking = tf.minimum(tf.maximum(masking, 0.0), 1.0)
       if hparams.use_predict_mask:
         masking = predict_mask
-      if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+      if hparams.mode == tf_estimator.ModeKeys.PREDICT:
         masking = predict_mask
       mask = tf.less(masking, tf.random_uniform(
           common_layers.shape_list(targets)[:-1]))
diff --git a/tensor2tensor/models/research/transformer_vae_flow_prior.py b/tensor2tensor/models/research/transformer_vae_flow_prior.py
index 9a254567b..64c1d88e3 100644
--- a/tensor2tensor/models/research/transformer_vae_flow_prior.py
+++ b/tensor2tensor/models/research/transformer_vae_flow_prior.py
@@ -36,6 +36,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 @registry.register_model
@@ -53,15 +54,15 @@ def __init__(self, *args, **kwargs):
 
   @property
   def is_training(self):
-    return self.hparams.mode == tf.estimator.ModeKeys.TRAIN
+    return self.hparams.mode == tf_estimator.ModeKeys.TRAIN
 
   @property
   def is_evaluating(self):
-    return self._hparams.mode == tf.estimator.ModeKeys.EVAL
+    return self._hparams.mode == tf_estimator.ModeKeys.EVAL
 
   @property
   def is_predicting(self):
-    return self._hparams.mode == tf.estimator.ModeKeys.PREDICT
+    return self._hparams.mode == tf_estimator.ModeKeys.PREDICT
 
   def loss_iw(self, logits, features):
     if isinstance(logits, dict):
@@ -494,7 +495,7 @@ def model_fn(self, features):
       else:
         logits = self.top(output, features)
         losses["training"] = 0.0
-        if (self._hparams.mode != tf.estimator.ModeKeys.PREDICT and
+        if (self._hparams.mode != tf_estimator.ModeKeys.PREDICT and
             self._hparams.mode != "attack"):
           losses["training"] = self.loss(logits, features)
 
diff --git a/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py b/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
index 824ad6e99..735c8bfc3 100644
--- a/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
+++ b/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
@@ -28,6 +28,7 @@
 from tensor2tensor.utils import learning_rate as lr
 from tensor2tensor.utils import mlperf_log
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 def _mixed_precision_is_enabled(hparams):
@@ -140,7 +141,7 @@ def decoder(name, latents, hparams, decoder_self_attention_bias, **kwargs):
 
 def drop_2d(targets, mode, dropout_p):
   """Dropout in 2D."""
-  if dropout_p > 0 and mode == tf.estimator.ModeKeys.TRAIN:
+  if dropout_p > 0 and mode == tf_estimator.ModeKeys.TRAIN:
     batch_size, targets_length, hidden_size = common_layers.shape_list(targets)
     mask_prob = tf.random_uniform(
         shape=(batch_size, targets_length), minval=0.0, maxval=1.0)
diff --git a/tensor2tensor/models/research/transformer_vae_test.py b/tensor2tensor/models/research/transformer_vae_test.py
index 27a1bef72..a9f38ee43 100644
--- a/tensor2tensor/models/research/transformer_vae_test.py
+++ b/tensor2tensor/models/research/transformer_vae_test.py
@@ -21,6 +21,7 @@
 from tensor2tensor.data_generators import problem_hparams
 from tensor2tensor.models.research import transformer_vae
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class TransformerVaeTest(tf.test.TestCase):
@@ -47,7 +48,7 @@ def testTransformerAEOnDVQ(self):
         "target_space_id": tf.constant(1, dtype=tf.int32),
     }
     tf.train.create_global_step()
-    model = transformer_vae.TransformerAE(hparams, tf.estimator.ModeKeys.TRAIN,
+    model = transformer_vae.TransformerAE(hparams, tf_estimator.ModeKeys.TRAIN,
                                           p_hparams)
     logits, _ = model(features)
     with self.test_session() as session:
diff --git a/tensor2tensor/models/research/universal_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
index 350c8481a..28f348c06 100644
--- a/tensor2tensor/models/research/universal_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.models.research import universal_transformer
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 BATCH_SIZE = 3
 INPUT_LENGTH = 5
@@ -36,7 +37,7 @@
 class UniversalTransformerTest(tf.test.TestCase):
 
   def get_model(self,
-                hparams, mode=tf.estimator.ModeKeys.TRAIN, has_input=True):
+                hparams, mode=tf_estimator.ModeKeys.TRAIN, has_input=True):
     hparams.hidden_size = 8
     hparams.filter_size = 32
     hparams.num_heads = 1
diff --git a/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index 8d745c7f9..5c07af376 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -28,6 +28,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 from tensorflow.contrib import rnn as contrib_rnn
 
 # pylint: disable=unused-import
@@ -60,7 +61,7 @@ def body(self, features):
           features["inputs"],
           model_fn=model_fn,
           trainable=hp.train_resnet,
-          is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
+          is_training=hp.mode == tf_estimator.ModeKeys.TRAIN)
     else:
       image_feat = features["inputs"]
 
@@ -133,7 +134,7 @@ def body(self, features):
           features["inputs"],
           model_fn=eval(hp.image_model_fn),
           trainable=hp.train_resnet,
-          is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
+          is_training=hp.mode == tf_estimator.ModeKeys.TRAIN)
     else:
       image_feat = features["inputs"]
 
diff --git a/tensor2tensor/models/research/vqa_attention_test.py b/tensor2tensor/models/research/vqa_attention_test.py
index 34116bf71..0512a766c 100644
--- a/tensor2tensor/models/research/vqa_attention_test.py
+++ b/tensor2tensor/models/research/vqa_attention_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.models.research import vqa_attention
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class VqaAttentionBaselineTest(tf.test.TestCase):
@@ -58,7 +59,7 @@ def testVqaAttentionBaseline(self):
           "targets": tf.constant(a, dtype=tf.int32),
       }
       model = vqa_attention.VqaAttentionBaseline(
-          hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
+          hparams, tf_estimator.ModeKeys.TRAIN, p_hparams)
       logits, losses = model(features)
       session.run(tf.global_variables_initializer())
       logits_, losses_ = session.run([logits, losses])
diff --git a/tensor2tensor/models/research/vqa_recurrent_self_attention.py b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
index 0ae482051..c85495c28 100644
--- a/tensor2tensor/models/research/vqa_recurrent_self_attention.py
+++ b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
@@ -31,6 +31,7 @@
 # from tensor2tensor.utils import restore_hook
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 from tensorflow.contrib.layers.python.layers import utils
 
@@ -57,7 +58,7 @@ def body(self, features):
           features["inputs"],
           model_fn=eval(hp.image_model_fn),
           trainable=hp.train_resnet,
-          is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
+          is_training=hp.mode == tf_estimator.ModeKeys.TRAIN)
     else:
       image_feat = features["inputs"]
 
diff --git a/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
index d10842a01..eb56afe1d 100644
--- a/tensor2tensor/models/research/vqa_self_attention.py
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -30,6 +30,7 @@
 # from tensor2tensor.utils import restore_hook
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 from tensorflow.contrib.layers.python.layers import utils
 
@@ -56,7 +57,7 @@ def body(self, features):
           features["inputs"],
           model_fn=eval(hp.image_model_fn),
           trainable=hp.train_resnet,
-          is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
+          is_training=hp.mode == tf_estimator.ModeKeys.TRAIN)
     else:
       image_feat = features["inputs"]
 
@@ -150,7 +151,7 @@ def body(self, features):
           features["inputs"],
           model_fn=eval(hp.image_model_fn),
           trainable=hp.train_resnet,
-          is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
+          is_training=hp.mode == tf_estimator.ModeKeys.TRAIN)
     else:
       image_feat = features["inputs"]
 
@@ -215,7 +216,7 @@ def body(self, features):
           features["inputs"],
           model_fn=eval(hp.image_model_fn),
           trainable=hp.train_resnet,
-          is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
+          is_training=hp.mode == tf_estimator.ModeKeys.TRAIN)
     else:
       image_feat = features["inputs"]
 

From a8e50c0364071ca596bc2e4a617e1f4174b2941b Mon Sep 17 00:00:00 2001
From: Fabien Hertschuh <fhertschuh@google.com>
Date: Tue, 22 Mar 2022 13:38:29 -0700
Subject: [PATCH 2708/2720] Explicitly import estimator from tensorflow as a
 separate import instead of accessing it via tf.estimator and depend on the
 tensorflow estimator target.

PiperOrigin-RevId: 436558319
---
 .../data_generators/allen_brain_test.py       |  5 +++--
 tensor2tensor/data_generators/celeba_test.py  |  3 ++-
 tensor2tensor/data_generators/celebahq.py     |  3 ++-
 tensor2tensor/data_generators/cifar.py        |  9 +++++----
 tensor2tensor/data_generators/common_voice.py |  4 ++--
 .../data_generators/function_docstring.py     |  3 ++-
 .../data_generators/generator_utils.py        |  3 ++-
 tensor2tensor/data_generators/imagenet.py     |  3 ++-
 .../data_generators/imagenet_test.py          |  3 ++-
 tensor2tensor/data_generators/librispeech.py  |  6 +++---
 tensor2tensor/data_generators/mscoco_test.py  |  3 ++-
 .../data_generators/multi_problem.py          |  5 +++--
 tensor2tensor/data_generators/problem.py      | 19 ++++++++++---------
 tensor2tensor/data_generators/problem_test.py |  7 ++++---
 .../data_generators/text_problems_test.py     |  3 ++-
 .../data_generators/timeseries_test.py        |  5 +++--
 tensor2tensor/data_generators/video_utils.py  | 11 ++++++-----
 tensor2tensor/data_generators/vqa_utils.py    |  5 +++--
 18 files changed, 58 insertions(+), 42 deletions(-)

diff --git a/tensor2tensor/data_generators/allen_brain_test.py b/tensor2tensor/data_generators/allen_brain_test.py
index 036e2a744..8ad44025e 100644
--- a/tensor2tensor/data_generators/allen_brain_test.py
+++ b/tensor2tensor/data_generators/allen_brain_test.py
@@ -26,11 +26,12 @@
 from tensor2tensor.utils import contrib
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 tfe = contrib.eager()
 tfe.enable_eager_execution()
-Modes = tf.estimator.ModeKeys  # pylint: disable=invalid-name
+Modes = tf_estimator.ModeKeys  # pylint: disable=invalid-name
 
 
 def mock_raw_image(x_dim=1024, y_dim=1024, num_channels=3,
@@ -195,7 +196,7 @@ def test_transformer2d_single_step_e2e(self):
         p_hparams = problem_object.get_hparams(hparams)
 
         model = image_transformer_2d.Img2imgTransformer(
-            hparams, tf.estimator.ModeKeys.TRAIN, p_hparams
+            hparams, tf_estimator.ModeKeys.TRAIN, p_hparams
         )
 
         @tfe.implicit_value_and_gradients
diff --git a/tensor2tensor/data_generators/celeba_test.py b/tensor2tensor/data_generators/celeba_test.py
index 11119714a..8206e38a7 100644
--- a/tensor2tensor/data_generators/celeba_test.py
+++ b/tensor2tensor/data_generators/celeba_test.py
@@ -24,6 +24,7 @@
 from tensor2tensor.utils import hparam
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class CelebaTest(parameterized.TestCase, tf.test.TestCase):
@@ -34,7 +35,7 @@ class CelebaTest(parameterized.TestCase, tf.test.TestCase):
       ("Dilated", "DILATED"))
   def testCelebaMultiResolutionPreprocessExample(self, resize_method):
     example = {"inputs": tf.random_uniform([218, 178, 3], minval=-1.)}
-    mode = tf.estimator.ModeKeys.TRAIN
+    mode = tf_estimator.ModeKeys.TRAIN
     hparams = hparam.HParams(resolutions=[8, 16, 32])
     if resize_method is not None:
       hparams.resize_method = resize_method
diff --git a/tensor2tensor/data_generators/celebahq.py b/tensor2tensor/data_generators/celebahq.py
index cabd2e819..5a590efc7 100644
--- a/tensor2tensor/data_generators/celebahq.py
+++ b/tensor2tensor/data_generators/celebahq.py
@@ -26,6 +26,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 @registry.register_problem
@@ -63,7 +64,7 @@ def filepattern(self, data_dir, mode, shard=None):
       shard_str = "000[0-8]"
     else:
       assert mode in [problem.DatasetSplit.EVAL,
-                      tf.estimator.ModeKeys.PREDICT,
+                      tf_estimator.ModeKeys.PREDICT,
                       problem.DatasetSplit.TEST]
       # Use the last 10 shards.
       shard_str = "0009"
diff --git a/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index 762c72f35..dd26c0ac5 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -35,6 +35,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 # URLs and filenames for CIFAR data.
 _CIFAR10_URL = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
@@ -131,7 +132,7 @@ def class_labels(self):
   def preprocess_example(self, example, mode, unused_hparams):
     image = example["inputs"]
     image.set_shape([_CIFAR10_IMAGE_SIZE, _CIFAR10_IMAGE_SIZE, 3])
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
       image = image_utils.cifar_image_augmentation(image)
     if not self._was_reversed:
       image = tf.image.per_image_standardization(image)
@@ -211,7 +212,7 @@ def dataset_filename(self):
   def preprocess_example(self, example, mode, unused_hparams):
     example["inputs"].set_shape([_CIFAR10_IMAGE_SIZE, _CIFAR10_IMAGE_SIZE, 3])
     example["inputs"] = tf.to_int64(example["inputs"])
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
       example["inputs"] = image_utils.random_shift(
           example["inputs"], wsr=0.1, hsr=0.1)
     return example
@@ -391,7 +392,7 @@ def class_labels(self):
   def preprocess_example(self, example, mode, unused_hparams):
     image = example["inputs"]
     image.set_shape([_CIFAR100_IMAGE_SIZE, _CIFAR100_IMAGE_SIZE, 3])
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
       image = image_utils.cifar_image_augmentation(image)
     if not self._was_reversed:
       image = tf.image.per_image_standardization(image)
@@ -522,7 +523,7 @@ def class_labels(self):
   def preprocess_example(self, example, mode, unused_hparams):
     image = example["inputs"]
     image.set_shape([_CIFAR100_IMAGE_SIZE, _CIFAR100_IMAGE_SIZE, 3])
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
       image = image_utils.cifar_image_augmentation(image)
     if not self._was_reversed:
       image = tf.image.per_image_standardization(image)
diff --git a/tensor2tensor/data_generators/common_voice.py b/tensor2tensor/data_generators/common_voice.py
index e35a8d50f..26df843d1 100644
--- a/tensor2tensor/data_generators/common_voice.py
+++ b/tensor2tensor/data_generators/common_voice.py
@@ -24,13 +24,13 @@
 import csv
 import os
 import tarfile
+from tensorflow.compat.v1 import estimator as tf_estimator
 import tqdm  # pylint: disable=g-bad-import-order
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import speech_recognition
 from tensor2tensor.utils import registry
 
-import tensorflow.compat.v1 as tf
 
 _COMMONVOICE_URL = "https://common-voice-data-download.s3.amazonaws.com/cv_corpus_v1.tar.gz"  # pylint: disable=line-too-long
 
@@ -213,7 +213,7 @@ def filepattern(self, data_dir, mode, shard=None):
     if mode == problem.DatasetSplit.TRAIN:
       path = os.path.join(data_dir, "common_voice")
       suffix = "train"
-    elif mode in [problem.DatasetSplit.EVAL, tf.estimator.ModeKeys.PREDICT]:
+    elif mode in [problem.DatasetSplit.EVAL, tf_estimator.ModeKeys.PREDICT]:
       path = os.path.join(data_dir, "common_voice_clean")
       suffix = "dev"
     else:
diff --git a/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
index e3a0e8a02..e604dfe57 100644
--- a/tensor2tensor/data_generators/function_docstring.py
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -21,6 +21,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 @registry.register_problem
@@ -99,7 +100,7 @@ def generate_samples(self, data_dir, tmp_dir, dataset_split):
             }
 
   def preprocess_example(self, example, mode, unused_hparams):
-    if mode != tf.estimator.ModeKeys.TRAIN:
+    if mode != tf_estimator.ModeKeys.TRAIN:
       example["embed_code"] = [0]
     return example
 
diff --git a/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 7081e1f23..89f182a65 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -39,6 +39,7 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 UNSHUFFLED_SUFFIX = "-unshuffled"
 
@@ -1168,7 +1169,7 @@ def make_tmp_dir(suffix="", prefix="tmp", dir=None):  # pylint: disable=redefine
 
 
 def tfrecord_iterator_for_problem(problem, data_dir,
-                                  dataset_split=tf.estimator.ModeKeys.TRAIN):
+                                  dataset_split=tf_estimator.ModeKeys.TRAIN):
   """Iterate over the records on disk for the Problem."""
   filenames = tf.gfile.Glob(problem.filepattern(data_dir, mode=dataset_split))
   example_spec = problem.example_reading_spec()[0]
diff --git a/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index 7c6c4d52f..af0b130b1 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -28,6 +28,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 # URLs and filenames for IMAGENET 32x32 data from
 # https://arxiv.org/abs/1601.06759.
@@ -105,7 +106,7 @@ def imagenet_preprocess_example(example, mode, resize_size=None,
   assert resize_size[0] == resize_size[1]
 
   image = example["inputs"]
-  if mode == tf.estimator.ModeKeys.TRAIN:
+  if mode == tf_estimator.ModeKeys.TRAIN:
     image = preprocess_for_train(image, image_size=resize_size[0],
                                  normalize=normalize)
   else:
diff --git a/tensor2tensor/data_generators/imagenet_test.py b/tensor2tensor/data_generators/imagenet_test.py
index 130283a8f..034ffad70 100644
--- a/tensor2tensor/data_generators/imagenet_test.py
+++ b/tensor2tensor/data_generators/imagenet_test.py
@@ -24,6 +24,7 @@
 from tensor2tensor.utils import hparam
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class ImagenetTest(parameterized.TestCase, tf.test.TestCase):
@@ -34,7 +35,7 @@ class ImagenetTest(parameterized.TestCase, tf.test.TestCase):
       ("Dilated", "DILATED"))
   def testImagenetMultiResolutionPreprocessExample(self, resize_method):
     example = {"inputs": tf.random_uniform([64, 64, 3], minval=-1.)}
-    mode = tf.estimator.ModeKeys.TRAIN
+    mode = tf_estimator.ModeKeys.TRAIN
     hparams = hparam.HParams(resolutions=[8, 16, 32])
     if resize_method is not None:
       hparams.resize_method = resize_method
diff --git a/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index cbf342e2f..8e5062cef 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -22,7 +22,7 @@
 from tensor2tensor.data_generators import speech_recognition
 from tensor2tensor.utils import registry
 
-import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 _LIBRISPEECH_TRAIN_DATASETS = [
     [
@@ -216,7 +216,7 @@ def filepattern(self, data_dir, mode, shard=None):
     if mode == problem.DatasetSplit.TRAIN:
       path = os.path.join(data_dir, "librispeech")
       suffix = "train"
-    elif mode in [problem.DatasetSplit.EVAL, tf.estimator.ModeKeys.PREDICT]:
+    elif mode in [problem.DatasetSplit.EVAL, tf_estimator.ModeKeys.PREDICT]:
       path = os.path.join(data_dir, "librispeech_clean")
       suffix = "dev"
     else:
@@ -264,7 +264,7 @@ def filepattern(self, data_dir, mode, shard=None):
     if mode == problem.DatasetSplit.TRAIN:
       path = os.path.join(data_dir, "librispeech")
       suffix = "train"
-    elif mode in [problem.DatasetSplit.EVAL, tf.estimator.ModeKeys.PREDICT]:
+    elif mode in [problem.DatasetSplit.EVAL, tf_estimator.ModeKeys.PREDICT]:
       path = os.path.join(data_dir, "librispeech_noisy")
       suffix = "dev"
     else:
diff --git a/tensor2tensor/data_generators/mscoco_test.py b/tensor2tensor/data_generators/mscoco_test.py
index 40f447711..3e5ab687a 100644
--- a/tensor2tensor/data_generators/mscoco_test.py
+++ b/tensor2tensor/data_generators/mscoco_test.py
@@ -24,6 +24,7 @@
 from tensor2tensor.utils import hparam
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class MscocoTest(parameterized.TestCase, tf.test.TestCase):
@@ -34,7 +35,7 @@ class MscocoTest(parameterized.TestCase, tf.test.TestCase):
       ("Dilated", "DILATED"))
   def testMsCocoMultiResolutionPreprocessExample(self, resize_method):
     example = {"inputs": tf.random_uniform([400, 400, 3], minval=-1.)}
-    mode = tf.estimator.ModeKeys.TRAIN
+    mode = tf_estimator.ModeKeys.TRAIN
     hparams = hparam.HParams(resolutions=[8, 16, 32])
     if resize_method is not None:
       hparams.resize_method = resize_method
diff --git a/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index e6ca24a31..4e2ea12d7 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -26,6 +26,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import metrics
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class MixingSchedule(object):
@@ -192,8 +193,8 @@ def dataset(self,
     # A list of datasets corresponding to the tasks in the task_list object
     # that need to be mixed.
     datasets = []
-    is_training = mode == tf.estimator.ModeKeys.TRAIN
-    is_infer = mode == tf.estimator.ModeKeys.PREDICT
+    is_training = mode == tf_estimator.ModeKeys.TRAIN
+    is_infer = mode == tf_estimator.ModeKeys.PREDICT
     enc = self.task_list[0].feature_encoders(data_dir=data_dir)["targets"]
     self.update_task_ids(enc.vocab_size)
 
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 68b86d653..30d63f3bb 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -33,6 +33,7 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 # pylint: disable=g-import-not-at-top
 try:
   from tensorflow.contrib.tpu.python.tpu import tpu_config
@@ -44,8 +45,8 @@
 
 
 class DatasetSplit(object):
-  TRAIN = tf.estimator.ModeKeys.TRAIN
-  EVAL = tf.estimator.ModeKeys.EVAL
+  TRAIN = tf_estimator.ModeKeys.TRAIN
+  EVAL = tf_estimator.ModeKeys.EVAL
   TEST = "test"
 
 
@@ -151,7 +152,7 @@ def preprocess_example_common(example, mode, hparams):
   if "inputs" in example and hparams.max_input_seq_length > 0:
     example["inputs"] = example["inputs"][:hparams.max_input_seq_length]
   if hparams.prepend_mode != "none":
-    if mode == tf.estimator.ModeKeys.PREDICT:
+    if mode == tf_estimator.ModeKeys.PREDICT:
       example["partial_targets"] = tf.concat([example["inputs"], [0]], 0)
     else:
       example["targets"] = tf.concat(
@@ -484,7 +485,7 @@ def filepattern(self, data_dir, mode, shard=None):
     shard_str = "-%05d" % shard if shard is not None else ""
     if mode == DatasetSplit.TRAIN:
       suffix = "train"
-    elif mode in [DatasetSplit.EVAL, tf.estimator.ModeKeys.PREDICT]:
+    elif mode in [DatasetSplit.EVAL, tf_estimator.ModeKeys.PREDICT]:
       suffix = "dev"
     else:
       assert mode == DatasetSplit.TEST
@@ -633,7 +634,7 @@ def dataset(self,
     Raises:
       ValueError: if num_partitions is greater than the number of data files.
     """
-    is_training = mode == tf.estimator.ModeKeys.TRAIN
+    is_training = mode == tf_estimator.ModeKeys.TRAIN
     shuffle_files = shuffle_files or shuffle_files is None and is_training
 
     dataset_split = dataset_split or mode
@@ -825,7 +826,7 @@ def _dataset_partition(self, mode, config, params):
       partition_id: an integer
       num_partitions: an integer
     """
-    if mode != tf.estimator.ModeKeys.TRAIN or not hasattr(config, "tpu_config"):
+    if mode != tf_estimator.ModeKeys.TRAIN or not hasattr(config, "tpu_config"):
       # Reset in the case when using TPU but alternating TRAIN and EVAL.
       self._next_partition_id = 0
       return 0, 1
@@ -875,7 +876,7 @@ def input_fn(self,
       (features_dict<str name, Tensor feature>, Tensor targets)
     """
     partition_id, num_partitions = self._dataset_partition(mode, config, params)
-    is_training = mode == tf.estimator.ModeKeys.TRAIN
+    is_training = mode == tf_estimator.ModeKeys.TRAIN
     if config and config.use_tpu:
       num_threads = 64
     else:
@@ -919,7 +920,7 @@ def export_assets(self):
   def serving_input_fn(self, hparams, decode_hparams=None, use_tpu=False):
     """Input fn for serving export, starting from serialized example."""
     self._hparams = hparams
-    mode = tf.estimator.ModeKeys.PREDICT
+    mode = tf_estimator.ModeKeys.PREDICT
     serialized_example = tf.placeholder(
         dtype=tf.string, shape=[None], name="serialized_example")
     dataset = tf.data.Dataset.from_tensor_slices(serialized_example)
@@ -947,7 +948,7 @@ def serving_input_fn(self, hparams, decode_hparams=None, use_tpu=False):
     if self.has_inputs:
       features.pop("targets", None)
 
-    return tf.estimator.export.ServingInputReceiver(
+    return tf_estimator.export.ServingInputReceiver(
         features=features, receiver_tensors=serialized_example)
 
 
diff --git a/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 882f8c42a..4a907c639 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -30,6 +30,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 tf.enable_eager_execution()
 
 
@@ -58,7 +59,7 @@ def setUpClass(cls):
   @test_utils.run_in_graph_mode_only()
   def testNoShuffleDeterministic(self):
     problem = algorithmic.TinyAlgo()
-    dataset = problem.dataset(mode=tf.estimator.ModeKeys.TRAIN,
+    dataset = problem.dataset(mode=tf_estimator.ModeKeys.TRAIN,
                               data_dir=algorithmic.TinyAlgo.data_dir,
                               shuffle_files=False)
 
@@ -72,10 +73,10 @@ def testNoShuffleDeterministic(self):
   def testNoShufflePreprocess(self):
 
     problem = algorithmic.TinyAlgo()
-    dataset1 = problem.dataset(mode=tf.estimator.ModeKeys.TRAIN,
+    dataset1 = problem.dataset(mode=tf_estimator.ModeKeys.TRAIN,
                                data_dir=algorithmic.TinyAlgo.data_dir,
                                shuffle_files=False, preprocess=False)
-    dataset2 = problem.dataset(mode=tf.estimator.ModeKeys.TRAIN,
+    dataset2 = problem.dataset(mode=tf_estimator.ModeKeys.TRAIN,
                                data_dir=algorithmic.TinyAlgo.data_dir,
                                shuffle_files=False, preprocess=True)
 
diff --git a/tensor2tensor/data_generators/text_problems_test.py b/tensor2tensor/data_generators/text_problems_test.py
index 4dd17bdf3..c7986dce0 100644
--- a/tensor2tensor/data_generators/text_problems_test.py
+++ b/tensor2tensor/data_generators/text_problems_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.data_generators import text_problems
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class Test1(text_problems.Text2textTmpdir):
@@ -171,7 +172,7 @@ def testText2TextTmpDir(self):
     self.assertTrue(tf.gfile.Exists(train_file))
     self.assertTrue(tf.gfile.Exists(eval_file))
 
-    dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, self.tmp_dir)
+    dataset = problem.dataset(tf_estimator.ModeKeys.TRAIN, self.tmp_dir)
     features = dataset.make_one_shot_iterator().get_next()
 
     examples = []
diff --git a/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index aa6d20cbe..0f9e6678d 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -25,6 +25,7 @@
 from tensor2tensor.data_generators import timeseries
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class TimeseriesTest(tf.test.TestCase):
@@ -39,7 +40,7 @@ def testTimeseriesToyProblem(self):
     problem = timeseries.TimeseriesToyProblem()
     problem.generate_data(self.tmp_dir, self.tmp_dir)
 
-    dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, self.tmp_dir)
+    dataset = problem.dataset(tf_estimator.ModeKeys.TRAIN, self.tmp_dir)
     features = dataset.make_one_shot_iterator().get_next()
 
     examples = []
@@ -65,7 +66,7 @@ def testTimeseriesToyProblemNoInputs(self):
     problem = timeseries.TimeseriesToyProblemNoInputs()
     problem.generate_data(self.tmp_dir, self.tmp_dir)
 
-    dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, self.tmp_dir)
+    dataset = problem.dataset(tf_estimator.ModeKeys.TRAIN, self.tmp_dir)
     features = dataset.make_one_shot_iterator().get_next()
 
     examples = []
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 9ad8c8dfc..4d0026eeb 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -36,6 +36,7 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import video_metrics
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 FLAGS = flags.FLAGS
@@ -405,7 +406,7 @@ def serving_input_fn(self, hparams):
         ])
 
     # TODO(michalski): add support for passing input_action and input_reward.
-    return tf.estimator.export.ServingInputReceiver(
+    return tf_estimator.export.ServingInputReceiver(
         features={"inputs": video_input_frames},
         receiver_tensors=video_input_frames)
 
@@ -514,12 +515,12 @@ def check_integrity_and_batch(*datasets):
 
     num_frames = (
         hparams.video_num_input_frames + hparams.video_num_target_frames)
-    if mode == tf.estimator.ModeKeys.PREDICT:
+    if mode == tf_estimator.ModeKeys.PREDICT:
       num_frames = min(self.max_frames_per_video(hparams), num_frames)
 
     # We jump by a random position at the beginning to add variety.
     if (self.random_skip and self.settable_random_skip and interleave and
-        mode == tf.estimator.ModeKeys.TRAIN):
+        mode == tf_estimator.ModeKeys.TRAIN):
       random_skip = tf.random_uniform([], maxval=num_frames, dtype=tf.int64)
       preprocessed_dataset = preprocessed_dataset.skip(random_skip)
     if (self.use_not_breaking_batching and
@@ -529,7 +530,7 @@ def check_integrity_and_batch(*datasets):
       batch_dataset = preprocessed_dataset.batch(num_frames,
                                                  drop_remainder=True)
     dataset = batch_dataset.map(features_from_batch)
-    if self.shuffle and interleave and mode == tf.estimator.ModeKeys.TRAIN:
+    if self.shuffle and interleave and mode == tf_estimator.ModeKeys.TRAIN:
       dataset = dataset.shuffle(hparams.get("shuffle_buffer_size", 128))
     return dataset
 
@@ -719,7 +720,7 @@ def preprocess(self, dataset, mode, hparams, interleave=True):
     video_augment_func = functools.partial(
         video_augmentation, hue=self.hue, contrast=self.contrast,
         saturate=self.saturate)
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
       dataset = dataset.map(video_augment_func)
     return dataset
 
diff --git a/tensor2tensor/data_generators/vqa_utils.py b/tensor2tensor/data_generators/vqa_utils.py
index e4557e00b..d2811e889 100644
--- a/tensor2tensor/data_generators/vqa_utils.py
+++ b/tensor2tensor/data_generators/vqa_utils.py
@@ -21,6 +21,7 @@
 
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 from tensorflow.python.ops import control_flow_ops
 
@@ -208,7 +209,7 @@ def vqa_v2_preprocess_image(
   assert resize_side > 0
   if resize_side:
     image = _aspect_preserving_resize(image, resize_side)
-  if mode == tf.estimator.ModeKeys.TRAIN:
+  if mode == tf_estimator.ModeKeys.TRAIN:
     image = tf.random_crop(image, [height, width, 3])
   else:
     # Central crop, assuming resize_height > height, resize_width > width.
@@ -216,7 +217,7 @@ def vqa_v2_preprocess_image(
 
   image = tf.clip_by_value(image, 0.0, 1.0)
 
-  if mode == tf.estimator.ModeKeys.TRAIN and distort:
+  if mode == tf_estimator.ModeKeys.TRAIN and distort:
     image = _flip(image)
     num_distort_cases = 4
     # pylint: disable=unnecessary-lambda

From 316c9ce2f2b2373f44f5be0da712dda3e5861a75 Mon Sep 17 00:00:00 2001
From: Fabien Hertschuh <fhertschuh@google.com>
Date: Wed, 23 Mar 2022 12:53:11 -0700
Subject: [PATCH 2709/2720] Explicitly import estimator from tensorflow as a
 separate import instead of accessing it via tf.estimator and depend on the
 tensorflow estimator target.

PiperOrigin-RevId: 436808246
---
 tensor2tensor/bin/t2t_attack.py               |  3 ++-
 tensor2tensor/bin/t2t_decoder.py              |  3 ++-
 tensor2tensor/bin/t2t_eval.py                 |  3 ++-
 tensor2tensor/bin/t2t_prune.py                |  5 +++--
 tensor2tensor/bin/t2t_trainer.py              |  3 ++-
 .../layers/common_image_attention.py          |  9 +++++----
 .../layers/common_image_attention_test.py     |  7 ++++---
 tensor2tensor/layers/discretization.py        | 19 ++++++++++---------
 tensor2tensor/layers/latent_layers.py         |  9 +++++----
 tensor2tensor/layers/latent_layers_test.py    |  3 ++-
 tensor2tensor/layers/modalities.py            |  7 ++++---
 tensor2tensor/layers/modalities_test.py       |  7 ++++---
 tensor2tensor/layers/transformer_layers.py    |  9 +++++----
 tensor2tensor/models/video/nfg_interpolate.py |  3 ++-
 tensor2tensor/models/video/nfg_test_utils.py  |  5 +++--
 tensor2tensor/models/video/savp.py            |  3 ++-
 tensor2tensor/models/video/tests_utils.py     |  5 +++--
 17 files changed, 60 insertions(+), 43 deletions(-)

diff --git a/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 125bed0eb..55b10e18a 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -49,6 +49,7 @@
 from tensor2tensor.utils import usr_dir
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 flags = tf.flags
 FLAGS = flags.FLAGS
@@ -134,7 +135,7 @@ def create_surrogate_run_config(hp):
 def prepare_data(problem, hparams, params, config):
   """Construct input pipeline."""
   input_fn = problem.make_estimator_input_fn(
-      tf.estimator.ModeKeys.EVAL, hparams, force_repeat=True)
+      tf_estimator.ModeKeys.EVAL, hparams, force_repeat=True)
   dataset = input_fn(params, config)
   features, _ = dataset.make_one_shot_iterator().get_next()
   inputs, labels = features["targets"], features["inputs"]
diff --git a/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index 47e6344f3..ee4036f3d 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -42,6 +42,7 @@
 from tensor2tensor.utils import usr_dir
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 flags = tf.flags
 FLAGS = flags.FLAGS
@@ -129,7 +130,7 @@ def score_file(filename):
     features = {"targets": batch_targets}
 
   # Prepare the model and the graph when model runs on features.
-  model = registry.model(FLAGS.model)(hparams, tf.estimator.ModeKeys.EVAL)
+  model = registry.model(FLAGS.model)(hparams, tf_estimator.ModeKeys.EVAL)
   _, losses = model(features)
   saver = tf.train.Saver()
 
diff --git a/tensor2tensor/bin/t2t_eval.py b/tensor2tensor/bin/t2t_eval.py
index 81ce7d443..2ea5608b0 100644
--- a/tensor2tensor/bin/t2t_eval.py
+++ b/tensor2tensor/bin/t2t_eval.py
@@ -24,6 +24,7 @@
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 flags = tf.flags
 FLAGS = flags.FLAGS
@@ -42,7 +43,7 @@ def main(_):
   dataset_split = "test" if FLAGS.eval_use_test_set else None
   dataset_kwargs = {"dataset_split": dataset_split}
   eval_input_fn = hparams.problem.make_estimator_input_fn(
-      tf.estimator.ModeKeys.EVAL, hparams, dataset_kwargs=dataset_kwargs)
+      tf_estimator.ModeKeys.EVAL, hparams, dataset_kwargs=dataset_kwargs)
   config = t2t_trainer.create_run_config(hparams)
 
   # summary-hook in tf.estimator.EstimatorSpec requires
diff --git a/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
index c4c6f411e..d66aea769 100644
--- a/tensor2tensor/bin/t2t_prune.py
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -40,6 +40,7 @@
 from tensor2tensor.utils import usr_dir
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 flags = tf.flags
 FLAGS = flags.FLAGS
@@ -79,7 +80,7 @@ def main(argv):
 
   # add "_rev" as a hack to avoid image standardization
   problem = registry.problem(FLAGS.problem)
-  input_fn = problem.make_estimator_input_fn(tf.estimator.ModeKeys.EVAL,
+  input_fn = problem.make_estimator_input_fn(tf_estimator.ModeKeys.EVAL,
                                              hparams)
   dataset = input_fn(params, config).repeat()
   features, labels = dataset.make_one_shot_iterator().get_next()
@@ -91,7 +92,7 @@ def main(argv):
   spec = model_fn(
       features,
       labels,
-      tf.estimator.ModeKeys.EVAL,
+      tf_estimator.ModeKeys.EVAL,
       params=hparams,
       config=config)
 
diff --git a/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 4955cbceb..4ff922ece 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -35,6 +35,7 @@
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 flags = tf.flags
@@ -243,7 +244,7 @@ def create_run_config(hp, output_dir=None):
         "num_cores_per_replica":
             1,
         "per_host_input_for_training":
-            tf.estimator.tpu.InputPipelineConfig.BROADCAST,
+            tf_estimator.tpu.InputPipelineConfig.BROADCAST,
     }
 
   # the various custom getters we have written do not play well together yet.
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 754749620..5bb532f64 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -27,6 +27,7 @@
 from tensor2tensor.utils import expert_utils
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class AttentionType(object):
@@ -460,7 +461,7 @@ def ffn_layer(x, hparams, losses=None):
       y = tf.reshape(y, x_shape)
     elif hparams.ffn_layer == "local_moe_tpu":
       overhead = (hparams.moe_overhead_train
-                  if hparams.mode == tf.estimator.ModeKeys.TRAIN
+                  if hparams.mode == tf_estimator.ModeKeys.TRAIN
                   else hparams.moe_overhead_eval)
       x, x_shape, is_4d = maybe_reshape_4d_to_3d(x)
       y, loss = expert_utils.local_moe_tpu(
@@ -531,7 +532,7 @@ def postprocess_image(x, rows, cols, hparams):
                               use_bias=True,
                               activation=None,
                               name="output_conv")
-  if (hparams.mode == tf.estimator.ModeKeys.PREDICT and
+  if (hparams.mode == tf_estimator.ModeKeys.PREDICT and
       hparams.block_raster_scan):
     y = targets
     yshape = common_layers.shape_list(y)
@@ -577,7 +578,7 @@ def prepare_decoder(targets, hparams):
 
   # during training, images are [batch, IMG_LEN, IMG_LEN, 3].
   # At inference, they are [batch, curr_infer_length, 1, 1]
-  if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+  if hparams.mode == tf_estimator.ModeKeys.PREDICT:
     curr_infer_length = targets_shape[1]
     if hparams.block_raster_scan:
       assert hparams.img_len*channels % hparams.query_shape[1] == 0
@@ -659,7 +660,7 @@ def create_output(decoder_output, rows, cols, targets, hparams):
   batch = common_layers.shape_list(decoded_image)[0]
   depth = common_layers.shape_list(decoded_image)[-1]
   likelihood = getattr(hparams, "likelihood", DistributionType.CAT)
-  if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+  if hparams.mode == tf_estimator.ModeKeys.PREDICT:
     y = tf.reshape(decoded_image, [batch, -1, 1, 1, depth])
     output = y[:, :rows, :, :, :]
   elif likelihood == DistributionType.CAT:
diff --git a/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index 739113c78..16b51c5da 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -25,6 +25,7 @@
 from tensor2tensor.utils import hparam
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class CommonImageAttentionTest(parameterized.TestCase, tf.test.TestCase):
@@ -40,7 +41,7 @@ def testPostProcessImageTrainMode(self, likelihood, num_mixtures, depth):
     hparams = hparam.HParams(
         hidden_size=2,
         likelihood=likelihood,
-        mode=tf.estimator.ModeKeys.TRAIN,
+        mode=tf_estimator.ModeKeys.TRAIN,
         num_mixtures=num_mixtures,
     )
     inputs = tf.random_uniform([batch, rows, cols, hparams.hidden_size],
@@ -63,7 +64,7 @@ def testPostProcessImageInferMode(self, likelihood, num_mixtures, depth):
         block_raster_scan=True,
         hidden_size=2,
         likelihood=likelihood,
-        mode=tf.estimator.ModeKeys.PREDICT,
+        mode=tf_estimator.ModeKeys.PREDICT,
         num_mixtures=num_mixtures,
         query_shape=[block_length, block_width],
     )
@@ -95,7 +96,7 @@ def testCreateOutputTrainMode(self, likelihood, num_mixtures, depth):
         hidden_size=2,
         likelihood=likelihood,
         num_channels=channels,
-        mode=tf.estimator.ModeKeys.TRAIN,
+        mode=tf_estimator.ModeKeys.TRAIN,
         num_mixtures=num_mixtures,
     )
     decoder_output = tf.random_normal([batch, rows, cols, hparams.hidden_size])
diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 93ae86acd..b0ec46772 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -25,6 +25,7 @@
 from tensor2tensor.layers import common_layers
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 import tensorflow_probability as tfp
 
 from tensorflow.python.training import moving_averages  # pylint: disable=g-direct-tensorflow-import
@@ -472,7 +473,7 @@ def gumbel_softmax(x,
     d_dev = -tf.reduce_mean(d_variance)
     ret = s
 
-    if mode != tf.estimator.ModeKeys.TRAIN:
+    if mode != tf_estimator.ModeKeys.TRAIN:
       ret = tf.reshape(maxvhot, common_layers.shape_list(s))  # Just hot @eval.
     return m, ret, d_dev * 5.0 + tf.reduce_mean(kl) * 0.002
 
@@ -754,7 +755,7 @@ def discrete_bottleneck(inputs,
       y_clean = common_layers.saturating_sigmoid(outputs_discrete)
       if summary:
         tf.summary.histogram("y_clean", tf.reshape(y_clean, [-1]))
-      if noise_dev > 0 and mode == tf.estimator.ModeKeys.TRAIN:
+      if noise_dev > 0 and mode == tf_estimator.ModeKeys.TRAIN:
         noise = tf.truncated_normal(
             common_layers.shape_list(outputs_discrete),
             mean=0.0,
@@ -766,7 +767,7 @@ def discrete_bottleneck(inputs,
       y_discrete = tf.stop_gradient(d) + y - tf.stop_gradient(y)
       pd = common_layers.inverse_exp_decay(startup_steps * 2)
       pd *= discrete_mix
-      pd = pd if mode == tf.estimator.ModeKeys.TRAIN else 1.0
+      pd = pd if mode == tf_estimator.ModeKeys.TRAIN else 1.0
       c = tf.where(
           tf.less(tf.random_uniform([common_layers.shape_list(y)[0]]), pd),
           y_discrete, y)
@@ -1379,17 +1380,17 @@ def tanh_discrete_bottleneck(x, bottleneck_bits, bottleneck_noise,
   """Simple discretization through tanh, flip bottleneck_noise many bits."""
   x = tf.layers.dense(x, bottleneck_bits, name="tanh_discrete_bottleneck")
   d0 = tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x))) - 1.0
-  if mode == tf.estimator.ModeKeys.TRAIN:
+  if mode == tf_estimator.ModeKeys.TRAIN:
     x += tf.truncated_normal(
         common_layers.shape_list(x), mean=0.0, stddev=0.2)
   x = tf.tanh(x)
   d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
-  if mode == tf.estimator.ModeKeys.TRAIN:
+  if mode == tf_estimator.ModeKeys.TRAIN:
     noise = tf.random_uniform(common_layers.shape_list(x))
     noise = 2.0 * tf.to_float(tf.less(bottleneck_noise, noise)) - 1.0
     d *= noise
   d = common_layers.mix(d, x, discretize_warmup_steps,
-                        mode == tf.estimator.ModeKeys.TRAIN)
+                        mode == tf_estimator.ModeKeys.TRAIN)
   return d, d0
 
 
@@ -1410,13 +1411,13 @@ def isemhash_bottleneck(x,
   with tf.variable_scope("isemhash_bottleneck"):
     x = tf.layers.dense(x, bottleneck_bits, name="dense")
     y = common_layers.saturating_sigmoid(x)
-    if isemhash_noise_dev > 0 and mode == tf.estimator.ModeKeys.TRAIN:
+    if isemhash_noise_dev > 0 and mode == tf_estimator.ModeKeys.TRAIN:
       noise = tf.truncated_normal(
           common_layers.shape_list(x), mean=0.0, stddev=isemhash_noise_dev)
       y = common_layers.saturating_sigmoid(x + noise)
     d = tf.to_float(tf.less(0.5, y)) + y - tf.stop_gradient(y)
     d = 2.0 * d - 1.0  # Move from [0, 1] to [-1, 1].
-    if mode == tf.estimator.ModeKeys.TRAIN:  # Flip some bits.
+    if mode == tf_estimator.ModeKeys.TRAIN:  # Flip some bits.
       noise = tf.random_uniform(common_layers.shape_list(x))
       noise = 2.0 * tf.to_float(tf.less(bottleneck_noise, noise)) - 1.0
       d *= noise
@@ -1424,7 +1425,7 @@ def isemhash_bottleneck(x,
           d,
           2.0 * y - 1.0,
           discretize_warmup_steps,
-          mode == tf.estimator.ModeKeys.TRAIN,
+          mode == tf_estimator.ModeKeys.TRAIN,
           max_prob=isemhash_mix_prob)
     return d, 0.0
 
diff --git a/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index fc0eaffe7..025a8217d 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -27,6 +27,7 @@
 from tensor2tensor.utils import beam_search
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 import tensorflow_probability as tfp
 
 DO_SUMMARIES = True
@@ -556,7 +557,7 @@ def latent_prediction_model(inputs,
     latents_pred_loss: Tensor of shape [batch, length_q].
   """
   with tf.variable_scope(name, default_name="latent_prediction"):
-    if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+    if hparams.mode != tf_estimator.ModeKeys.PREDICT:
       latents_pred = transformer_latent_decoder(tf.stop_gradient(latents_dense),
                                                 inputs,
                                                 ed_attention_bias,
@@ -617,10 +618,10 @@ def transformer_autoencoder(inputs,
   losses = {"extra": 0.,
             "extra_loss": 0.,
             "latent_pred": 0.}
-  if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+  if hparams.mode != tf_estimator.ModeKeys.PREDICT:
     targets_compressed = compress_fn(targets, hparams, name="compress")
 
-    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    if hparams.mode == tf_estimator.ModeKeys.TRAIN:
       scale = common_layers.inverse_exp_decay(hparams.startup_steps)
     else:
       scale = 1.0
@@ -681,7 +682,7 @@ def transformer_autoencoder(inputs,
       [-1, hparams.img_len, hparams.img_len, hparams.hidden_size])
 
   if hparams.use_gold_targets:
-    if hparams.mode == tf.estimator.ModeKeys.PREDICT:
+    if hparams.mode == tf_estimator.ModeKeys.PREDICT:
       masking = predict_mask
     else:
       masking = common_layers.inverse_exp_decay(hparams.mask_startup_steps)
diff --git a/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
index 178d880d2..5abbcfd05 100644
--- a/tensor2tensor/layers/latent_layers_test.py
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -28,6 +28,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 tf.enable_eager_execution()
 
 
@@ -110,7 +111,7 @@ def testComputeBitsAndNats(self):
   @test_utils.run_in_graph_and_eager_modes()
   def testTransformerAutoencoder(self):
     hparams = imagetransformer_latent_tiny()
-    hparams.mode = tf.estimator.ModeKeys.TRAIN
+    hparams.mode = tf_estimator.ModeKeys.TRAIN
     block_dim = int(hparams.hidden_size // hparams.num_blocks)
     block_v_size = 2**(hparams.bottleneck_bits /
                        (hparams.num_residuals * hparams.num_blocks))
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index c36a5db19..ab842c55b 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -33,6 +33,7 @@
 from tensor2tensor.layers import discretization
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 import tensorflow_probability as tfp
 
 
@@ -309,7 +310,7 @@ def _image_channel_compress_bottom(inputs, model_hparams, name="bottom"):
   with tf.variable_scope(name):
     inputs = tf.to_float(inputs)
     hp = model_hparams
-    if hp.mode != tf.estimator.ModeKeys.PREDICT:
+    if hp.mode != tf_estimator.ModeKeys.PREDICT:
       tf.summary.image(
           "inputs",
           common_layers.tpu_safe_image_summary(inputs),
@@ -600,7 +601,7 @@ def video_pixel_noise_bottom(x, model_hparams, vocab_size):
   """Bottom transformation for video."""
   input_noise = getattr(model_hparams, "video_modality_input_noise", 0.25)
   inputs = x
-  if model_hparams.mode == tf.estimator.ModeKeys.TRAIN:
+  if model_hparams.mode == tf_estimator.ModeKeys.TRAIN:
     background = tfp.stats.percentile(inputs, 50., axis=[0, 1, 2, 3])
     input_shape = common_layers.shape_list(inputs)
     input_size = tf.reduce_prod(input_shape[:-1])
@@ -1126,7 +1127,7 @@ def symbol_top(body_output, targets, model_hparams, vocab_size):
     body_output_shape = common_layers.shape_list(body_output)
     var = get_weights(model_hparams, vocab_size, body_output_shape[-1])
     if (model_hparams.factored_logits and
-        model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
+        model_hparams.mode == tf_estimator.ModeKeys.TRAIN):
       # insert channels dimension
       body_output = tf.expand_dims(body_output, 3)
       return common_layers.FactoredTensor(body_output, var)
diff --git a/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index bbb17dfc2..0a5420890 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 tf.enable_eager_execution()
 
 
@@ -60,7 +61,7 @@ def testSymbolModalityInputs(self):
     hidden_size = 9
     model_hparams = common_hparams.basic_params1()
     model_hparams.hidden_size = hidden_size
-    model_hparams.mode = tf.estimator.ModeKeys.TRAIN
+    model_hparams.mode = tf_estimator.ModeKeys.TRAIN
     x = np.random.randint(
         vocab_size, size=(batch_size, length, 1, 1))
     data_parallelism = expert_utils.Parallelism(
@@ -86,7 +87,7 @@ def testSymbolModalityTargets(self):
     vocab_size = 11
     model_hparams = common_hparams.basic_params1()
     model_hparams.hidden_size = hidden_size
-    model_hparams.mode = tf.estimator.ModeKeys.TRAIN
+    model_hparams.mode = tf_estimator.ModeKeys.TRAIN
     body_output = np.random.randint(
         100, size=(batch_size, length, height, hidden_size))
     targets = np.random.randint(
@@ -127,7 +128,7 @@ def testSymbolModalityTargetsFactored(self):
     model_hparams = common_hparams.basic_params1()
     model_hparams.factored_logits = True
     model_hparams.hidden_size = hidden_size
-    model_hparams.mode = tf.estimator.ModeKeys.TRAIN
+    model_hparams.mode = tf_estimator.ModeKeys.TRAIN
     body_output = np.random.randint(
         100, size=(batch_size, length, height, hidden_size))
     targets = np.random.randint(
diff --git a/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 97739299e..662ba70a7 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -25,6 +25,7 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 # TODO(lukaszkaiser): remove this function when not needed any more.
@@ -237,8 +238,8 @@ def transformer_encoder(encoder_input,
               memory_height=memory_height,
               area_key_mode=hparams.get("area_key_mode", "none"),
               area_value_mode=hparams.get("area_value_mode", "none"),
-              training=(hparams.get("mode", tf.estimator.ModeKeys.TRAIN)
-                        == tf.estimator.ModeKeys.TRAIN))
+              training=(hparams.get("mode", tf_estimator.ModeKeys.TRAIN)
+                        == tf_estimator.ModeKeys.TRAIN))
           x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(
@@ -371,7 +372,7 @@ def transformer_ffn_layer(x,
     return common_layers.sru(x)
   elif ffn_layer == "local_moe_tpu":
     overhead = hparams.moe_overhead_eval
-    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    if hparams.mode == tf_estimator.ModeKeys.TRAIN:
       overhead = hparams.moe_overhead_train
     ret, loss = expert_utils.local_moe_tpu(
         x,
@@ -382,7 +383,7 @@ def transformer_ffn_layer(x,
         loss_coef=hparams.moe_loss_coef)
   elif ffn_layer == "local_moe":
     overhead = hparams.moe_overhead_eval
-    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    if hparams.mode == tf_estimator.ModeKeys.TRAIN:
       overhead = hparams.moe_overhead_train
     ret, loss = expert_utils.local_moe(
         x,
diff --git a/tensor2tensor/models/video/nfg_interpolate.py b/tensor2tensor/models/video/nfg_interpolate.py
index 8cbcfe54c..c9424d0b9 100644
--- a/tensor2tensor/models/video/nfg_interpolate.py
+++ b/tensor2tensor/models/video/nfg_interpolate.py
@@ -31,6 +31,7 @@
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import trainer_lib
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 # Flags placeholders.
 flags.DEFINE_string("checkpoint_path", None,
@@ -222,7 +223,7 @@ def main(_):
   # prepare dataset using Predict mode.
   dataset_split = "test" if FLAGS.eval_use_test_set else None
   dataset = hparams.problem.dataset(
-      tf.estimator.ModeKeys.PREDICT, shuffle_files=False, hparams=hparams,
+      tf_estimator.ModeKeys.PREDICT, shuffle_files=False, hparams=hparams,
       data_dir=FLAGS.data_dir, dataset_split=dataset_split)
   dataset = dataset.batch(hparams.batch_size)
   dataset = dataset.make_one_shot_iterator().get_next()
diff --git a/tensor2tensor/models/video/nfg_test_utils.py b/tensor2tensor/models/video/nfg_test_utils.py
index 9ca0d4a4c..fd679dd5d 100644
--- a/tensor2tensor/models/video/nfg_test_utils.py
+++ b/tensor2tensor/models/video/nfg_test_utils.py
@@ -26,7 +26,8 @@
 from tensor2tensor.models.video import next_frame_glow
 from tensor2tensor.utils import registry
 import tensorflow.compat.v1 as tf
-MODES = tf.estimator.ModeKeys
+from tensorflow.compat.v1 import estimator as tf_estimator
+MODES = tf_estimator.ModeKeys
 
 
 # TODO(mechcoder): Refactor or merge tests with the other next_frame_tests when
@@ -165,7 +166,7 @@ def GlowTrainAndDecode(self, in_frames=1, out_frames=1,
                              apply_dilations, activation)
       features = create_basic_features(hparams)
       model = next_frame_glow.NextFrameGlow(
-          hparams, tf.estimator.ModeKeys.PREDICT)
+          hparams, tf_estimator.ModeKeys.PREDICT)
       predictions = model.infer(features)
       outputs = predictions["outputs"]
       model_path = os.path.join(curr_dir, "model")
diff --git a/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index 48accd97f..2036af698 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -33,6 +33,7 @@
 from tensor2tensor.utils import update_ops_hook
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 import tensorflow_gan as tfgan
 
 gan_losses = tfgan.losses.wargs
@@ -398,7 +399,7 @@ def construct_model(self, images, actions, rewards):
       [], [], [], [], []
     pred_image = tf.zeros_like(images[0])
     prior_latent_state, cond_latent_state = None, None
-    train_mode = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
+    train_mode = self.hparams.mode == tf_estimator.ModeKeys.TRAIN
 
     # Create scheduled sampling function
     ss_func = self.get_scheduled_sample_func(batch_size)
diff --git a/tensor2tensor/models/video/tests_utils.py b/tensor2tensor/models/video/tests_utils.py
index e2487bd80..170ca484b 100644
--- a/tensor2tensor/models/video/tests_utils.py
+++ b/tensor2tensor/models/video/tests_utils.py
@@ -26,6 +26,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 def fill_hparams(hparams, in_frames, out_frames):
@@ -116,7 +117,7 @@ class BaseNextFrameTest(tf.test.TestCase):
 
   def RunModel(self, model, hparams, features):
     with tf.Session() as session:
-      model = model(hparams, tf.estimator.ModeKeys.TRAIN)
+      model = model(hparams, tf_estimator.ModeKeys.TRAIN)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
@@ -124,7 +125,7 @@ def RunModel(self, model, hparams, features):
 
   def InferModel(self, model, hparams, features):
     with tf.Session() as session:
-      model = model(hparams, tf.estimator.ModeKeys.PREDICT)
+      model = model(hparams, tf_estimator.ModeKeys.PREDICT)
       output = model.infer(features)
       session.run(tf.global_variables_initializer())
       res = session.run(output)

From 3817e96deda6f3fdada4fedcd5efe33ed0438485 Mon Sep 17 00:00:00 2001
From: Fabien Hertschuh <fhertschuh@google.com>
Date: Thu, 24 Mar 2022 12:33:41 -0700
Subject: [PATCH 2710/2720] Explicitly import estimator from tensorflow as a
 separate import instead of accessing it via tf.estimator and depend on the
 tensorflow estimator target.

PiperOrigin-RevId: 437060146
---
 tensor2tensor/models/basic_test.py            |  3 +-
 tensor2tensor/models/bytenet_test.py          |  3 +-
 tensor2tensor/models/distillation.py          |  3 +-
 .../models/evolved_transformer_test.py        | 11 +++----
 tensor2tensor/models/image_transformer.py     |  3 +-
 tensor2tensor/models/image_transformer_2d.py  | 21 +++++++-------
 .../models/image_transformer_2d_test.py       |  5 ++--
 .../models/image_transformer_test.py          |  3 +-
 tensor2tensor/models/lstm.py                  | 13 +++++----
 tensor2tensor/models/lstm_test.py             |  9 +++---
 tensor2tensor/models/mtf_image_transformer.py | 17 ++++++-----
 .../models/mtf_image_transformer_test.py      |  3 +-
 tensor2tensor/models/mtf_resnet.py            |  3 +-
 tensor2tensor/models/mtf_transformer.py       | 29 ++++++++++---------
 tensor2tensor/models/mtf_transformer_test.py  |  3 +-
 tensor2tensor/models/neural_assistant.py      |  5 ++--
 tensor2tensor/models/neural_gpu_test.py       |  3 +-
 tensor2tensor/models/resnet.py                |  3 +-
 tensor2tensor/models/resnet_test.py           |  3 +-
 tensor2tensor/models/revnet.py                |  3 +-
 tensor2tensor/models/revnet_test.py           |  5 ++--
 tensor2tensor/models/shake_shake.py           |  7 +++--
 tensor2tensor/models/slicenet_test.py         |  5 ++--
 tensor2tensor/models/transformer.py           |  7 +++--
 tensor2tensor/models/transformer_test.py      | 21 +++++++-------
 tensor2tensor/models/vanilla_gan.py           |  5 ++--
 tensor2tensor/models/xception_test.py         |  3 +-
 27 files changed, 113 insertions(+), 86 deletions(-)

diff --git a/tensor2tensor/models/basic_test.py b/tensor2tensor/models/basic_test.py
index c7098e5f0..df77a707c 100644
--- a/tensor2tensor/models/basic_test.py
+++ b/tensor2tensor/models/basic_test.py
@@ -25,6 +25,7 @@
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class BasicTest(tf.test.TestCase):
@@ -39,7 +40,7 @@ def testBasicFcRelu(self):
           "inputs": tf.constant(x, dtype=tf.int32),
           "targets": tf.constant(y, dtype=tf.int32),
       }
-      model = basic.BasicFcRelu(hparams, tf.estimator.ModeKeys.TRAIN)
+      model = basic.BasicFcRelu(hparams, tf_estimator.ModeKeys.TRAIN)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
diff --git a/tensor2tensor/models/bytenet_test.py b/tensor2tensor/models/bytenet_test.py
index 1ba972e9f..3dd087289 100644
--- a/tensor2tensor/models/bytenet_test.py
+++ b/tensor2tensor/models/bytenet_test.py
@@ -24,6 +24,7 @@
 from tensor2tensor.models import bytenet
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class ByteNetTest(tf.test.TestCase):
@@ -42,7 +43,7 @@ def testByteNet(self):
           "targets": tf.constant(y, dtype=tf.int32),
       }
       model = bytenet.ByteNet(
-          hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
+          hparams, tf_estimator.ModeKeys.TRAIN, p_hparams)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
diff --git a/tensor2tensor/models/distillation.py b/tensor2tensor/models/distillation.py
index 206fccd30..40c4adcf8 100644
--- a/tensor2tensor/models/distillation.py
+++ b/tensor2tensor/models/distillation.py
@@ -23,6 +23,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 @registry.register_model
@@ -44,7 +45,7 @@ class Distillation(t2t_model.T2TModel):
 
   def __init__(self,
                hparams,
-               mode=tf.estimator.ModeKeys.TRAIN,
+               mode=tf_estimator.ModeKeys.TRAIN,
                problem_hparams=None,
                data_parallelism=None,
                decode_hparams=None,
diff --git a/tensor2tensor/models/evolved_transformer_test.py b/tensor2tensor/models/evolved_transformer_test.py
index bf8c21f8a..6888e5c24 100644
--- a/tensor2tensor/models/evolved_transformer_test.py
+++ b/tensor2tensor/models/evolved_transformer_test.py
@@ -24,6 +24,7 @@
 from tensor2tensor.models import transformer
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 BATCH_SIZE = 3
 INPUT_LENGTH = 5
@@ -88,7 +89,7 @@ def get_model(hparams, has_input=True, num_decoder_layers=1):
     features["inputs"] = tf.constant(inputs, dtype=tf.int32, name="inputs")
 
   return (evolved_transformer.EvolvedTransformer(hparams,
-                                                 tf.estimator.ModeKeys.TRAIN,
+                                                 tf_estimator.ModeKeys.TRAIN,
                                                  p_hparams), features)
 
 
@@ -121,7 +122,7 @@ def testSlowVsFast(self):
       for _ in range(10):
         apply_grad.run()
 
-    model.set_mode(tf.estimator.ModeKeys.PREDICT)
+    model.set_mode(tf_estimator.ModeKeys.PREDICT)
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
       greedy_result = model._slow_greedy_infer(features,
@@ -155,7 +156,7 @@ def testSlowVsFastNoInput(self):
       for _ in range(10):
         apply_grad.run()
 
-    model.set_mode(tf.estimator.ModeKeys.PREDICT)
+    model.set_mode(tf_estimator.ModeKeys.PREDICT)
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
       slow_result = model._slow_greedy_infer(features, decode_length)["outputs"]
@@ -188,7 +189,7 @@ def testBeamVsFast(self):
       for _ in range(10):
         apply_grad.run()
 
-    model.set_mode(tf.estimator.ModeKeys.PREDICT)
+    model.set_mode(tf_estimator.ModeKeys.PREDICT)
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
       beam_result = model._beam_decode_slow(
@@ -227,7 +228,7 @@ def _create_greedy_infer_model(self):
       for _ in range(10):
         apply_grad.run()
 
-    model.set_mode(tf.estimator.ModeKeys.PREDICT)
+    model.set_mode(tf_estimator.ModeKeys.PREDICT)
 
     return model, features
 
diff --git a/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index 885ec7db1..bb8b5b87e 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -33,6 +33,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 @registry.register_model
@@ -52,7 +53,7 @@ def body(self, features):
       raise ValueError("When using DMOL for the likelihood, bottom function "
                        " must be identity and num_channels must be 1.")
     if (not tf.get_variable_scope().reuse and
-        hparams.mode != tf.estimator.ModeKeys.PREDICT):
+        hparams.mode != tf_estimator.ModeKeys.PREDICT):
       tf.summary.image("targets", tf.to_float(targets), max_outputs=1)
 
     # Extra losses list if we want to use moe.
diff --git a/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index 291a44b99..d0b1c1981 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -34,6 +34,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 @registry.register_model
@@ -46,7 +47,7 @@ def body(self, features):
     targets = features["targets"]
     targets_shape = common_layers.shape_list(targets)
     if not (tf.get_variable_scope().reuse or
-            hparams.mode == tf.estimator.ModeKeys.PREDICT):
+            hparams.mode == tf_estimator.ModeKeys.PREDICT):
       tf.summary.image("targets", targets, max_outputs=1)
 
     decoder_input, rows, cols = cia.prepare_decoder(
@@ -76,7 +77,7 @@ def body(self, features):
     targets = features["targets"]
     inputs = features["inputs"]
     if not (tf.get_variable_scope().reuse or
-            hparams.mode == tf.estimator.ModeKeys.PREDICT):
+            hparams.mode == tf_estimator.ModeKeys.PREDICT):
       tf.summary.image("inputs", inputs, max_outputs=1)
       tf.summary.image("targets", targets, max_outputs=1)
 
@@ -112,7 +113,7 @@ def body(self, features):
     targets = features["targets"]
     inputs = features["inputs"]
     if not (tf.get_variable_scope().reuse or
-            hparams.mode == tf.estimator.ModeKeys.PREDICT):
+            hparams.mode == tf_estimator.ModeKeys.PREDICT):
       tf.summary.image("inputs", inputs, max_outputs=1)
       tf.summary.image("targets", targets, max_outputs=1)
 
@@ -174,11 +175,11 @@ def top(self, body_output, features):
     assert self._hparams.block_size > 0
 
     train_or_eval = (
-        self._hparams.mode == tf.estimator.ModeKeys.TRAIN or
-        self._hparams.mode == tf.estimator.ModeKeys.EVAL)
+        self._hparams.mode == tf_estimator.ModeKeys.TRAIN or
+        self._hparams.mode == tf_estimator.ModeKeys.EVAL)
 
     if train_or_eval:
-      if self._hparams.mode == tf.estimator.ModeKeys.TRAIN:
+      if self._hparams.mode == tf_estimator.ModeKeys.TRAIN:
         features["block_index"] = tf.random_uniform(
             shape=[], minval=0, maxval=self._hparams.block_size, dtype=tf.int64)
       else:
@@ -203,7 +204,7 @@ def top(self, body_output, features):
   def loss(self, logits, features):
     assert self._hparams.block_size > 0
 
-    if self._hparams.mode == tf.estimator.ModeKeys.PREDICT:
+    if self._hparams.mode == tf_estimator.ModeKeys.PREDICT:
       return 0.0
 
     def shift_left_2d(x, k):
@@ -222,8 +223,8 @@ def shift_left_4d_raster_scan(x, k):
         for i in range(self._hparams.block_size)
     ], axis=4)
 
-    if (self._hparams.mode == tf.estimator.ModeKeys.TRAIN or
-        self._hparams.mode == tf.estimator.ModeKeys.EVAL):
+    if (self._hparams.mode == tf_estimator.ModeKeys.TRAIN or
+        self._hparams.mode == tf_estimator.ModeKeys.EVAL):
       assert "block_index" in features
       targets = targets[:, :, :, :, features["block_index"]]
 
@@ -231,7 +232,7 @@ def shift_left_4d_raster_scan(x, k):
 
     loss = super(Img2imgTransformerBlockParallel, self).loss(logits, features)
 
-    if self._hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    if self._hparams.mode == tf_estimator.ModeKeys.TRAIN:
       k = features["block_index"]
       loss_num, loss_den = loss
       loss_val = loss_num / loss_den
diff --git a/tensor2tensor/models/image_transformer_2d_test.py b/tensor2tensor/models/image_transformer_2d_test.py
index 56a00d4ef..13f6e799a 100644
--- a/tensor2tensor/models/image_transformer_2d_test.py
+++ b/tensor2tensor/models/image_transformer_2d_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class Img2imgTransformerTest(tf.test.TestCase):
@@ -43,7 +44,7 @@ def _test_img2img_transformer(self, net):
           "targets": tf.constant(targets, dtype=tf.int32),
           "target_space_id": tf.constant(1, dtype=tf.int32),
       }
-      model = net(hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
+      model = net(hparams, tf_estimator.ModeKeys.TRAIN, p_hparams)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
@@ -73,7 +74,7 @@ def _test_imagetransformer_2d(self, net):
           "targets": tf.constant(targets, dtype=tf.int32),
           "target_space_id": tf.constant(1, dtype=tf.int32),
       }
-      model = net(hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
+      model = net(hparams, tf_estimator.ModeKeys.TRAIN, p_hparams)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
diff --git a/tensor2tensor/models/image_transformer_test.py b/tensor2tensor/models/image_transformer_test.py
index 1f5c04395..142db24b3 100644
--- a/tensor2tensor/models/image_transformer_test.py
+++ b/tensor2tensor/models/image_transformer_test.py
@@ -27,6 +27,7 @@
 from tensor2tensor.models import image_transformer
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class ImagetransformerTest(parameterized.TestCase, tf.test.TestCase):
@@ -56,7 +57,7 @@ def testImagetransformer(self, net, hparams):
           "targets": tf.constant(targets, dtype=tf.int32),
           "target_space_id": tf.constant(1, dtype=tf.int32),
       }
-      model = net(hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
+      model = net(hparams, tf_estimator.ModeKeys.TRAIN, p_hparams)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
diff --git a/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 90d4b67a7..4f7447357 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -28,6 +28,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 def _dropout_lstm_cell(hparams, train):
@@ -114,7 +115,7 @@ def _area_key_value_fn(keys, values):
       keys = area_attention.compute_area_key(
           keys, max_area_width=hparams.get("max_area_width", 1),
           mode=hparams.get("area_key_mode", "none"), name="decoder_encoder",
-          training=(hparams.mode == tf.estimator.ModeKeys.TRAIN))
+          training=(hparams.mode == tf_estimator.ModeKeys.TRAIN))
       if hparams.get("area_value_mode", "none") == "sum":
         _, _, values, _, _ = area_attention.compute_area_features(
             values, max_area_width=hparams.get("max_area_width", 1))
@@ -333,7 +334,7 @@ class LSTMEncoder(t2t_model.T2TModel):
   def body(self, features):
     if self._hparams.initializer == "orthogonal":
       raise ValueError("LSTM models fail with orthogonal initializer.")
-    train = self._hparams.mode == tf.estimator.ModeKeys.TRAIN
+    train = self._hparams.mode == tf_estimator.ModeKeys.TRAIN
     inputs = features.get("inputs")
     inputs_length = common_layers.length_from_embedding(inputs)
     # Flatten inputs.
@@ -352,7 +353,7 @@ def body(self, features):
     # TODO(lukaszkaiser): investigate this issue and repair.
     if self._hparams.initializer == "orthogonal":
       raise ValueError("LSTM models fail with orthogonal initializer.")
-    train = self._hparams.mode == tf.estimator.ModeKeys.TRAIN
+    train = self._hparams.mode == tf_estimator.ModeKeys.TRAIN
     return lstm_seq2seq_internal(features.get("inputs"), features["targets"],
                                  self._hparams, train)
 
@@ -365,7 +366,7 @@ def body(self, features):
     # TODO(lukaszkaiser): investigate this issue and repair.
     if self._hparams.initializer == "orthogonal":
       raise ValueError("LSTM models fail with orthogonal initializer.")
-    train = self._hparams.mode == tf.estimator.ModeKeys.TRAIN
+    train = self._hparams.mode == tf_estimator.ModeKeys.TRAIN
     # This is a temporary fix for varying-length sequences within in a batch.
     # A more complete fix should pass a length tensor from outside so that
     # all the lstm variants can use it.
@@ -390,7 +391,7 @@ def body(self, features):
     # TODO(lukaszkaiser): investigate this issue and repair.
     if self._hparams.initializer == "orthogonal":
       raise ValueError("LSTM models fail with orthogonal initializer.")
-    train = self._hparams.mode == tf.estimator.ModeKeys.TRAIN
+    train = self._hparams.mode == tf_estimator.ModeKeys.TRAIN
     return lstm_seq2seq_internal_bid_encoder(
         features.get("inputs"), features["targets"], self._hparams, train)
 
@@ -402,7 +403,7 @@ def body(self, features):
     # TODO(lukaszkaiser): investigate this issue and repair.
     if self._hparams.initializer == "orthogonal":
       raise ValueError("LSTM models fail with orthogonal initializer.")
-    train = self._hparams.mode == tf.estimator.ModeKeys.TRAIN
+    train = self._hparams.mode == tf_estimator.ModeKeys.TRAIN
     return lstm_seq2seq_internal_attention_bid_encoder(
         features.get("inputs"), features["targets"], self._hparams, train)
 
diff --git a/tensor2tensor/models/lstm_test.py b/tensor2tensor/models/lstm_test.py
index 92615e167..4af344ee5 100644
--- a/tensor2tensor/models/lstm_test.py
+++ b/tensor2tensor/models/lstm_test.py
@@ -24,6 +24,7 @@
 from tensor2tensor.models import lstm
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class LSTMTest(tf.test.TestCase):
@@ -41,7 +42,7 @@ def testLSTMSeq2Seq(self):
           "inputs": tf.constant(x, dtype=tf.int32),
           "targets": tf.constant(y, dtype=tf.int32),
       }
-      model = lstm.LSTMSeq2seq(hparams, tf.estimator.ModeKeys.TRAIN,
+      model = lstm.LSTMSeq2seq(hparams, tf_estimator.ModeKeys.TRAIN,
                                p_hparams)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
@@ -66,7 +67,7 @@ def testLSTMSeq2SeqAttention(self):
           "targets": tf.constant(y, dtype=tf.int32),
       }
       model = lstm.LSTMSeq2seqAttention(
-          hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
+          hparams, tf_estimator.ModeKeys.TRAIN, p_hparams)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
@@ -86,7 +87,7 @@ def testLSTMSeq2seqBidirectionalEncoder(self):
           "targets": tf.constant(y, dtype=tf.int32),
       }
       model = lstm.LSTMSeq2seqBidirectionalEncoder(
-          hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
+          hparams, tf_estimator.ModeKeys.TRAIN, p_hparams)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
@@ -108,7 +109,7 @@ def testLSTMSeq2seqAttentionBidirectionalEncoder(self):
           "targets": tf.constant(y, dtype=tf.int32),
       }
       model = lstm.LSTMSeq2seqAttentionBidirectionalEncoder(
-          hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
+          hparams, tf_estimator.ModeKeys.TRAIN, p_hparams)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
diff --git a/tensor2tensor/models/mtf_image_transformer.py b/tensor2tensor/models/mtf_image_transformer.py
index c12b1b23f..6db8e0d38 100644
--- a/tensor2tensor/models/mtf_image_transformer.py
+++ b/tensor2tensor/models/mtf_image_transformer.py
@@ -33,6 +33,7 @@
 from tensor2tensor.utils import mtf_model
 from tensor2tensor.utils import registry
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 @registry.register_model
@@ -243,8 +244,8 @@ def import_to_batch_by_length(x, name):
 def layer_prepostprocess_dropout(x, hparams):
   batch_dim = x.shape.dims[0]
   model_dim = x.shape.dims[-1]
-  mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN)
-  is_training = mode == tf.estimator.ModeKeys.TRAIN
+  mode = getattr(hparams, "mode", tf_estimator.ModeKeys.TRAIN)
+  is_training = mode == tf_estimator.ModeKeys.TRAIN
   return mtf.dropout(
       x, is_training,
       keep_prob=1.0 - hparams.layer_prepostprocess_dropout,
@@ -261,8 +262,8 @@ def local_attention1d_spatial_decoder(x, kv_dim, heads_dim,
   x = mtf.reshape(
       x, mtf.Shape([batch_dim, num_w_blocks_dim, blocks_w_dim, model_dim]))
   # [ self attention - ffn - residual + dropout] x n
-  mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN)
-  is_training = mode == tf.estimator.ModeKeys.TRAIN
+  mode = getattr(hparams, "mode", tf_estimator.ModeKeys.TRAIN)
+  is_training = mode == tf_estimator.ModeKeys.TRAIN
   for layer in range(hparams.num_decoder_layers):
     layer_name = "decoder_layer_%d" % layer
     with tf.variable_scope(layer_name):
@@ -311,8 +312,8 @@ def local_attention2d_spatial_decoder(x, kv_dim, heads_dim,
           batch_dim, num_h_blocks_dim, num_w_blocks_dim,
           blocks_h_dim, blocks_w_dim, model_dim
       ]))
-  mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN)
-  is_training = mode == tf.estimator.ModeKeys.TRAIN
+  mode = getattr(hparams, "mode", tf_estimator.ModeKeys.TRAIN)
+  is_training = mode == tf_estimator.ModeKeys.TRAIN
   # Image Transformer Decoder
   # [ self attention - ffn - residual + dropout] x n
   for layer in range(hparams.num_decoder_layers):
@@ -345,8 +346,8 @@ def local_attention1d_masked_decoder(x, kv_dim, heads_dim,
   """Image Transformer decoder with local1D masked layers."""
   print(x)
   _, length_dim, model_dim = x.shape.dims
-  mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN)
-  is_training = mode == tf.estimator.ModeKeys.TRAIN
+  mode = getattr(hparams, "mode", tf_estimator.ModeKeys.TRAIN)
+  is_training = mode == tf_estimator.ModeKeys.TRAIN
   for layer in range(hparams.num_decoder_layers):
     layer_name = "decoder_layer_%d" % layer
     with tf.variable_scope(layer_name):
diff --git a/tensor2tensor/models/mtf_image_transformer_test.py b/tensor2tensor/models/mtf_image_transformer_test.py
index 77a9f902e..fd028f331 100644
--- a/tensor2tensor/models/mtf_image_transformer_test.py
+++ b/tensor2tensor/models/mtf_image_transformer_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.models import mtf_image_transformer
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 # Constants shared between all functions.
 BATCH_SIZE = 8
@@ -34,7 +35,7 @@
 
 
 def get_model(hparams=None,
-              mode=tf.estimator.ModeKeys.TRAIN,
+              mode=tf_estimator.ModeKeys.TRAIN,
               model_cls=mtf_image_transformer.MtfImageTransformer):
   if hparams is None:
     hparams = mtf_image_transformer.mtf_image_transformer_single()
diff --git a/tensor2tensor/models/mtf_resnet.py b/tensor2tensor/models/mtf_resnet.py
index 3c42f6eba..84a95254a 100644
--- a/tensor2tensor/models/mtf_resnet.py
+++ b/tensor2tensor/models/mtf_resnet.py
@@ -29,6 +29,7 @@
 from tensor2tensor.utils import mtf_model
 from tensor2tensor.utils import registry
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 BATCH_NORM_DECAY = 0.9
@@ -214,7 +215,7 @@ def mtf_model_fn(self, features, mesh):
     tf.logging.info("features = %s" % features)
     hparams = self._hparams
     activation_dtype = self.set_activation_type()
-    is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+    is_training = hparams.mode == tf_estimator.ModeKeys.TRAIN
 
     # Declare all the dimensions
     batch_dim = mtf.Dimension("batch", hparams.batch_size)
diff --git a/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 8e1ba2a7b..18015ea00 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -29,6 +29,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 @registry.register_model
@@ -37,7 +38,7 @@ class MtfTransformer(mtf_model.MtfModel):
 
   def __init__(self,
                hparams,
-               mode=tf.estimator.ModeKeys.TRAIN,
+               mode=tf_estimator.ModeKeys.TRAIN,
                problem_hparams=None,
                data_parallelism=None,
                decode_hparams=None,
@@ -194,7 +195,7 @@ def _noisy_targets_from_spec(self, targets, noising_spec, losses=None):
       # Train a small transformer to fill in masked out values, then
       # sample from it.
       hparams = self._hparams
-      if hparams.mode != tf.estimator.ModeKeys.TRAIN:
+      if hparams.mode != tf_estimator.ModeKeys.TRAIN:
         raise NotImplementedError("Not implemented")
       noiser_hparams = copy.copy(self._hparams)
       noiser_hparams.del_hparam("mode")
@@ -222,7 +223,7 @@ def _noisy_targets(self, targets, losses=None):
       a Tensor the same dtype and shape as Targets
     """
     hparams = self._hparams
-    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    if hparams.mode == tf_estimator.ModeKeys.TRAIN:
       nt_train = self._noisy_targets_from_spec(
           targets, hparams.noising_spec_train, losses=losses)
       if hparams.noising_use_eval_during_train > 0:
@@ -242,8 +243,8 @@ def _mtf_model_fn(self, features, mesh):
     hparams = self._hparams
     extra_losses = []
     targets = tf.to_int32(features["targets"])
-    mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN)
-    is_training = mode == tf.estimator.ModeKeys.TRAIN
+    mode = getattr(hparams, "mode", tf_estimator.ModeKeys.TRAIN)
+    is_training = mode == tf_estimator.ModeKeys.TRAIN
     if len(targets.get_shape()) > 2:
       tf.logging.info("targets = %s" % targets)
       targets = tf.squeeze(targets, [2, 3])
@@ -359,7 +360,7 @@ def layer_prepostprocess_dropout(x):
             encdec_attention_mask=encoder_decoder_attention_mask,
             losses=extra_losses)
     if (hparams.reshape_logits_hack and
-        hparams.mode == tf.estimator.ModeKeys.TRAIN):
+        hparams.mode == tf_estimator.ModeKeys.TRAIN):
       # For some reason, the logits computation is extremely slow on TPU
       # in some cases where the batch size per core is 1.  Reshape the logits
       # and the targets to double the batch size and halve the length.
@@ -373,7 +374,7 @@ def layer_prepostprocess_dropout(x):
       targets = mtf.reshape(targets, new_dims)
 
     logits = mtf.matmul(x, softmax_var)
-    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
+    if hparams.mode == tf_estimator.ModeKeys.TRAIN:
       logits = mtf.layers.multiplicative_jitter(logits, epsilon=1e-2)
     off_value = hparams.label_smoothing / self._targets_vocab_size
     on_value = 1.0 - hparams.label_smoothing + off_value
@@ -387,7 +388,7 @@ def layer_prepostprocess_dropout(x):
     for l in extra_losses:
       loss += l
     if (hparams.reshape_logits_hack and
-        hparams.mode == tf.estimator.ModeKeys.TRAIN):
+        hparams.mode == tf_estimator.ModeKeys.TRAIN):
       logits = mtf.reshape(logits, old_dims + [self.targets_vocab_dim])
     logits = mtf.to_float(logits)
     return logits, loss
@@ -428,8 +429,8 @@ def _feedforward_layer(self, x, layer_type, losses=None):
       ValueError: if hparams make no sense
     """
     hparams = self._hparams
-    mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN)
-    is_training = mode == tf.estimator.ModeKeys.TRAIN
+    mode = getattr(hparams, "mode", tf_estimator.ModeKeys.TRAIN)
+    is_training = mode == tf_estimator.ModeKeys.TRAIN
     if layer_type == "drd":
       return mtf.layers.dense_relu_dense(
           x, self.feedforward_dim, is_training, dropout=hparams.relu_dropout,
@@ -443,7 +444,7 @@ def _feedforward_layer(self, x, layer_type, losses=None):
           x,
           self.model_dim,
           hparams,
-          hparams.mode == tf.estimator.ModeKeys.TRAIN,
+          hparams.mode == tf_estimator.ModeKeys.TRAIN,
           master_dtype=self.master_dtype,
           slice_dtype=self.slice_dtype)
       if losses is not None:
@@ -454,7 +455,7 @@ def _feedforward_layer(self, x, layer_type, losses=None):
           x,
           self.model_dim,
           hparams,
-          hparams.mode == tf.estimator.ModeKeys.TRAIN,
+          hparams.mode == tf_estimator.ModeKeys.TRAIN,
           master_dtype=self.master_dtype,
           slice_dtype=self.slice_dtype)
       if losses is not None:
@@ -496,8 +497,8 @@ def _layer_stack(self,
     """
     hparams = self._hparams
     is_incremental = (step_num is not None)
-    mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN)
-    is_training = mode == tf.estimator.ModeKeys.TRAIN
+    mode = getattr(hparams, "mode", tf_estimator.ModeKeys.TRAIN)
+    is_training = mode == tf_estimator.ModeKeys.TRAIN
     def layer_prepostprocess_dropout(x):
       if is_incremental:
         return x
diff --git a/tensor2tensor/models/mtf_transformer_test.py b/tensor2tensor/models/mtf_transformer_test.py
index 345bc8e30..cb32f805c 100644
--- a/tensor2tensor/models/mtf_transformer_test.py
+++ b/tensor2tensor/models/mtf_transformer_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.models import mtf_transformer
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 # Constants shared between all functions.
 BATCH_SIZE = 2
@@ -34,7 +35,7 @@
 VOCAB_SIZE = 128
 
 
-def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
+def get_model(hparams=None, mode=tf_estimator.ModeKeys.TRAIN,
               has_input=True, model_cls=mtf_transformer.MtfTransformer):
   if hparams is None:
     hparams = mtf_transformer.mtf_transformer_single()
diff --git a/tensor2tensor/models/neural_assistant.py b/tensor2tensor/models/neural_assistant.py
index c96c1e243..17f6ae954 100644
--- a/tensor2tensor/models/neural_assistant.py
+++ b/tensor2tensor/models/neural_assistant.py
@@ -24,6 +24,7 @@
 from tensor2tensor.models import transformer
 from tensor2tensor.utils import registry
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 @registry.register_model
@@ -68,7 +69,7 @@ def model_fn(self, features):
         kb_train_weight = self._hparams.kb_train_weight
         cur_lm_loss_weight = 1.0 - cur_kb_loss_weight
         # Finalize loss
-        if (self._hparams.mode != tf.estimator.ModeKeys.PREDICT and
+        if (self._hparams.mode != tf_estimator.ModeKeys.PREDICT and
             self._hparams.mode != "attack"):
           lm_loss_num, lm_loss_denom = self.loss(logits, features)
           total_loss = (kb_train_weight) * cur_knowledge_training_loss + (
@@ -254,7 +255,7 @@ def compute_knowledge_selection_and_loss(self, features, encoder_output,
         margin + positive_loss - negative_loss,
         clip_value_min=0,
         clip_value_max=100)
-    if hparams.mode != tf.estimator.ModeKeys.PREDICT:
+    if hparams.mode != tf_estimator.ModeKeys.PREDICT:
       triple_losses = tf.nn.weighted_cross_entropy_with_logits(
           labels=triple_labels,
           logits=triple_logits,
diff --git a/tensor2tensor/models/neural_gpu_test.py b/tensor2tensor/models/neural_gpu_test.py
index 6ae190e65..07a7c5b83 100644
--- a/tensor2tensor/models/neural_gpu_test.py
+++ b/tensor2tensor/models/neural_gpu_test.py
@@ -25,6 +25,7 @@
 from tensor2tensor.models import neural_gpu
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class NeuralGPUTest(tf.test.TestCase):
@@ -48,7 +49,7 @@ def testNeuralGPU(self):
           "inputs": tf.constant(inputs, dtype=tf.int32),
           "targets": tf.constant(targets, dtype=tf.int32)
       }
-      model = neural_gpu.NeuralGPU(hparams, tf.estimator.ModeKeys.TRAIN,
+      model = neural_gpu.NeuralGPU(hparams, tf_estimator.ModeKeys.TRAIN,
                                    p_hparams)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
diff --git a/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index 288241bca..e36e79b15 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -27,6 +27,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 BATCH_NORM_DECAY = 0.9
@@ -542,7 +543,7 @@ def body(self, features):
         "bottleneck": bottleneck_block,
     }
     assert hp.block_fn in block_fns
-    is_training = hp.mode == tf.estimator.ModeKeys.TRAIN
+    is_training = hp.mode == tf_estimator.ModeKeys.TRAIN
     if is_training:
       targets = features["targets_raw"]
 
diff --git a/tensor2tensor/models/resnet_test.py b/tensor2tensor/models/resnet_test.py
index 14d881fe5..d982153ad 100644
--- a/tensor2tensor/models/resnet_test.py
+++ b/tensor2tensor/models/resnet_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.models import resnet
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 def resnet_tiny_cpu():
@@ -55,7 +56,7 @@ def _test_resnet(self, img_size, output_size):
           "inputs": tf.constant(x, dtype=tf.int32),
           "targets": tf.constant(y, dtype=tf.int32),
       }
-      model = resnet.Resnet(hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
+      model = resnet.Resnet(hparams, tf_estimator.ModeKeys.TRAIN, p_hparams)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
diff --git a/tensor2tensor/models/revnet.py b/tensor2tensor/models/revnet.py
index d20c6b378..05b21b0f7 100644
--- a/tensor2tensor/models/revnet.py
+++ b/tensor2tensor/models/revnet.py
@@ -41,6 +41,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 def wrapped_partial(fn, *args, **kwargs):
@@ -312,7 +313,7 @@ def revnet(inputs, hparams, reuse=None):
   Returns:
     [batch_size, hidden_dim] pre-logits tensor from the bottleneck RevNet.
   """
-  training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+  training = hparams.mode == tf_estimator.ModeKeys.TRAIN
   with tf.variable_scope('RevNet', reuse=reuse):
     x1, x2 = init(inputs,
                   num_channels=hparams.num_channels_init_block,
diff --git a/tensor2tensor/models/revnet_test.py b/tensor2tensor/models/revnet_test.py
index af3ddce7c..1aef79886 100644
--- a/tensor2tensor/models/revnet_test.py
+++ b/tensor2tensor/models/revnet_test.py
@@ -17,6 +17,7 @@
 
 from tensor2tensor.models import revnet
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class RevnetTest(tf.test.TestCase):
@@ -99,7 +100,7 @@ def testFinalBlock3D(self):
   def testEndToEnd(self):
     images = tf.random_uniform([1, 299, 299, 3])
     hparams = revnet.revnet_base()
-    hparams.mode = tf.estimator.ModeKeys.TRAIN
+    hparams.mode = tf_estimator.ModeKeys.TRAIN
     logits = revnet.revnet(images, hparams)
     self.assertEqual(logits.shape, [1, 1, 1, 3328])
 
@@ -107,7 +108,7 @@ def testEndToEnd3D(self):
     images = tf.random_uniform([1, 299, 299, 299, 3])
     hparams = revnet.revnet_base()
     hparams.dim = '3d'
-    hparams.mode = tf.estimator.ModeKeys.TRAIN
+    hparams.mode = tf_estimator.ModeKeys.TRAIN
     logits = revnet.revnet(images, hparams)
     self.assertEqual(logits.shape, [1, 1, 1, 1, 3328])
 
diff --git a/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index 47224bcd1..f63bf3329 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -25,6 +25,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 def shake_shake_skip_connection(x, output_filters, stride, is_training):
@@ -55,7 +56,7 @@ def shake_shake_skip_connection(x, output_filters, stride, is_training):
 def shake_shake_branch(x, output_filters, stride, rand_forward, rand_backward,
                        hparams):
   """Building a 2 branching convnet."""
-  is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+  is_training = hparams.mode == tf_estimator.ModeKeys.TRAIN
   x = tf.nn.relu(x)
   x = tf.layers.conv2d(
       x,
@@ -77,7 +78,7 @@ def shake_shake_branch(x, output_filters, stride, rand_forward, rand_backward,
 
 def shake_shake_block(x, output_filters, stride, hparams):
   """Builds a full shake-shake sub layer."""
-  is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+  is_training = hparams.mode == tf_estimator.ModeKeys.TRAIN
   batch_size = common_layers.shape_list(x)[0]
 
   # Generate random numbers for scaling the branches.
@@ -139,7 +140,7 @@ class ShakeShake(t2t_model.T2TModel):
 
   def body(self, features):
     hparams = self._hparams
-    is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
+    is_training = hparams.mode == tf_estimator.ModeKeys.TRAIN
     inputs = features["inputs"]
     assert (hparams.num_hidden_layers - 2) % 6 == 0
     assert hparams.hidden_size % 16 == 0
diff --git a/tensor2tensor/models/slicenet_test.py b/tensor2tensor/models/slicenet_test.py
index cc9101125..d4e43f7a4 100644
--- a/tensor2tensor/models/slicenet_test.py
+++ b/tensor2tensor/models/slicenet_test.py
@@ -27,6 +27,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class SliceNetTest(tf.test.TestCase):
@@ -45,7 +46,7 @@ def testSliceNet(self):
           "targets": tf.constant(y, dtype=tf.int32),
           "target_space_id": tf.constant(1, dtype=tf.int32),
       }
-      model = slicenet.SliceNet(hparams, tf.estimator.ModeKeys.TRAIN,
+      model = slicenet.SliceNet(hparams, tf_estimator.ModeKeys.TRAIN,
                                 p_hparams)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
@@ -66,7 +67,7 @@ def testSliceNetImageToText(self):
           "targets": tf.constant(y, dtype=tf.int32),
           "target_space_id": tf.constant(1, dtype=tf.int32),
       }
-      model = slicenet.SliceNet(hparams, tf.estimator.ModeKeys.TRAIN,
+      model = slicenet.SliceNet(hparams, tf_estimator.ModeKeys.TRAIN,
                                 p_hparams)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
diff --git a/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index 0cc8b5f2a..e5294e89f 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -42,6 +42,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 # pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.ops import inplace_ops
@@ -172,7 +173,7 @@ def transformer_decode(decoder_function,
       **kwargs)
 
   if (common_layers.is_xla_compiled() and
-      hparams.mode == tf.estimator.ModeKeys.TRAIN):
+      hparams.mode == tf_estimator.ModeKeys.TRAIN):
     # TPU does not react kindly to extra dimensions.
     # TODO(noam): remove this once TPU is more forgiving of extra dims.
     return decoder_output
@@ -1530,7 +1531,7 @@ def transformer_self_attention_layer(decoder_input,
           area_value_mode=hparams.get("area_value_mode", "none"),
           training=(hparams.get(
               "mode",
-              tf.estimator.ModeKeys.TRAIN) == tf.estimator.ModeKeys.TRAIN))
+              tf_estimator.ModeKeys.TRAIN) == tf_estimator.ModeKeys.TRAIN))
       x = common_layers.layer_postprocess(x, y, hparams)
     if encoder_output is not None:
       if not isinstance(encoder_output, (list,)):
@@ -1569,7 +1570,7 @@ def transformer_self_attention_layer(decoder_input,
               area_value_mode=hparams.get("area_value_mode", "none"),
               training=(hparams.get(
                   "mode",
-                  tf.estimator.ModeKeys.TRAIN) == tf.estimator.ModeKeys.TRAIN))
+                  tf_estimator.ModeKeys.TRAIN) == tf_estimator.ModeKeys.TRAIN))
           x = common_layers.layer_postprocess(x, y, hparams)
     return x, layer_cache
 
diff --git a/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index fe6fc0ceb..397f75b56 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -25,6 +25,7 @@
 from tensor2tensor.models import transformer
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 BATCH_SIZE = 3
@@ -33,7 +34,7 @@
 VOCAB_SIZE = 10
 
 
-def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
+def get_model(hparams=None, mode=tf_estimator.ModeKeys.TRAIN,
               has_input=True, model_cls=transformer.Transformer):
   if hparams is None:
     hparams = transformer.transformer_tiny()
@@ -150,7 +151,7 @@ def testSlowVsFast(self, get_model_fn=None, p=None):
       for _ in range(100):
         apply_grad.run()
 
-    model.set_mode(tf.estimator.ModeKeys.PREDICT)
+    model.set_mode(tf_estimator.ModeKeys.PREDICT)
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
       greedy_result = model._slow_greedy_infer(
@@ -185,7 +186,7 @@ def testSlowVsFastNoInput(self):
       for _ in range(100):
         apply_grad.run()
 
-    model.set_mode(tf.estimator.ModeKeys.PREDICT)
+    model.set_mode(tf_estimator.ModeKeys.PREDICT)
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
       slow_result = model._slow_greedy_infer(
@@ -204,7 +205,7 @@ def testSlowVsFastNoInput(self):
   def testBeamDecodeWithRelativeAttention(self):
     decode_length = 2
     model, features = get_model(transformer.transformer_relative_tiny())
-    model.set_mode(tf.estimator.ModeKeys.PREDICT)
+    model.set_mode(tf_estimator.ModeKeys.PREDICT)
 
     beam_result = model._beam_decode(
         features, decode_length, beam_size=4, top_beams=1,
@@ -237,7 +238,7 @@ def testBeamVsFast(self):
       for _ in range(100):
         apply_grad.run()
 
-    model.set_mode(tf.estimator.ModeKeys.PREDICT)
+    model.set_mode(tf_estimator.ModeKeys.PREDICT)
 
     with tf.variable_scope(tf.get_variable_scope(), reuse=True):
       beam_result = model._beam_decode_slow(
@@ -315,7 +316,7 @@ def _create_greedy_infer_model(self):
       for _ in range(100):
         apply_grad.run()
 
-    model.set_mode(tf.estimator.ModeKeys.PREDICT)
+    model.set_mode(tf_estimator.ModeKeys.PREDICT)
 
     return model, features
 
@@ -387,7 +388,7 @@ class TransformerScorerTest(tf.test.TestCase):
 
   def testReturnsScores(self):
     model, features = get_model(
-        mode=tf.estimator.ModeKeys.PREDICT,
+        mode=tf_estimator.ModeKeys.PREDICT,
         model_cls=transformer.TransformerScorer)
     infer_out = model.infer(features)
     self.assertTrue("outputs" in infer_out)
@@ -402,21 +403,21 @@ def testReturnsScores(self):
   def testVarNames(self):
     with tf.Graph().as_default():
       model, features = get_model(
-          mode=tf.estimator.ModeKeys.PREDICT,
+          mode=tf_estimator.ModeKeys.PREDICT,
           model_cls=transformer.TransformerScorer)
       _ = model.infer(features)
       scorer_vars = [v.name for v in tf.global_variables()]
 
     with tf.Graph().as_default():
       model, features = get_model(
-          mode=tf.estimator.ModeKeys.EVAL,
+          mode=tf_estimator.ModeKeys.EVAL,
           model_cls=transformer.TransformerScorer)
       _ = model(features)
       scorer_eval_vars = [v.name for v in tf.global_variables()]
 
     with tf.Graph().as_default():
       model, features = get_model(
-          mode=tf.estimator.ModeKeys.EVAL,
+          mode=tf_estimator.ModeKeys.EVAL,
           model_cls=transformer.Transformer)
       _ = model(features)
       transformer_vars = [v.name for v in tf.global_variables()]
diff --git a/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index 03c1a9354..371a5031e 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -28,6 +28,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 def lrelu(input_, leak=0.2, name="lrelu"):
@@ -135,7 +136,7 @@ def body(self, features):
       and losses is a dictionary of losses (that get added for the final loss).
     """
     features["targets"] = features["inputs"]
-    is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
+    is_training = self.hparams.mode == tf_estimator.ModeKeys.TRAIN
 
     # Input images.
     inputs = tf.to_float(features["targets_raw"])
@@ -170,7 +171,7 @@ class SlicedGan(AbstractGAN):
 
   def losses(self, inputs, generated):
     """Losses in the sliced case."""
-    is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN
+    is_training = self.hparams.mode == tf_estimator.ModeKeys.TRAIN
     def discriminate(x):
       return self.discriminator(x, is_training=is_training, reuse=False)
     generator_loss = common_layers.sliced_gan_loss(
diff --git a/tensor2tensor/models/xception_test.py b/tensor2tensor/models/xception_test.py
index 90d49b8f2..3f7df88b5 100644
--- a/tensor2tensor/models/xception_test.py
+++ b/tensor2tensor/models/xception_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.models import xception
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class XceptionTest(tf.test.TestCase):
@@ -48,7 +49,7 @@ def _test_xception(self, img_size):
           "inputs": tf.constant(x, dtype=tf.int32),
           "targets": tf.constant(y, dtype=tf.int32),
       }
-      model = xception.Xception(hparams, tf.estimator.ModeKeys.TRAIN, p_hparams)
+      model = xception.Xception(hparams, tf_estimator.ModeKeys.TRAIN, p_hparams)
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)

From 23c36f6404775e19c086d7dbe7c6e90e47cc8e14 Mon Sep 17 00:00:00 2001
From: Fabien Hertschuh <fhertschuh@google.com>
Date: Fri, 25 Mar 2022 15:13:09 -0700
Subject: [PATCH 2711/2720] Explicitly import estimator from tensorflow as a
 separate import instead of accessing it via tf.estimator and depend on the
 tensorflow estimator target.

PiperOrigin-RevId: 437344902
---
 tensor2tensor/utils/contrib.py            |  7 ++--
 tensor2tensor/utils/data_reader.py        |  7 ++--
 tensor2tensor/utils/data_reader_test.py   |  9 ++--
 tensor2tensor/utils/decoding.py           |  3 +-
 tensor2tensor/utils/mtf_model.py          | 31 +++++++-------
 tensor2tensor/utils/scheduled_sampling.py |  3 +-
 tensor2tensor/utils/t2t_model.py          | 51 ++++++++++++-----------
 tensor2tensor/utils/trainer_lib.py        | 27 ++++++------
 tensor2tensor/utils/trainer_lib_test.py   |  9 ++--
 tensor2tensor/utils/video2gif.py          |  3 +-
 10 files changed, 80 insertions(+), 70 deletions(-)

diff --git a/tensor2tensor/utils/contrib.py b/tensor2tensor/utils/contrib.py
index efa22c9f6..78a3acd02 100644
--- a/tensor2tensor/utils/contrib.py
+++ b/tensor2tensor/utils/contrib.py
@@ -24,6 +24,7 @@
 
 from absl import logging
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 # Check if we have contrib available
 try:
@@ -180,11 +181,11 @@ def replace_monitors_with_hooks(monitors_or_hooks, estimator):
   del estimator
   monitors_or_hooks = monitors_or_hooks or []
   hooks = [
-      m for m in monitors_or_hooks if isinstance(m, tf.estimator.SessionRunHook)
+      m for m in monitors_or_hooks if isinstance(m, tf_estimator.SessionRunHook)
   ]
   deprecated_monitors = [
       m for m in monitors_or_hooks
-      if not isinstance(m, tf.estimator.SessionRunHook)
+      if not isinstance(m, tf_estimator.SessionRunHook)
   ]
   assert not deprecated_monitors
   return hooks
@@ -196,7 +197,7 @@ def learn():
     from tensorflow.contrib import learn as contrib_learn  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
     return contrib_learn
   return DummyModule(
-      RunConfig=tf.estimator.RunConfig,
+      RunConfig=tf_estimator.RunConfig,
       monitors=DummyModule(
           replace_monitors_with_hooks=replace_monitors_with_hooks),
   )
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index aea085f28..25cc58868 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -29,6 +29,7 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 def cast_ints_to_int32(features):
@@ -345,7 +346,7 @@ def input_fn(dataset,
   Returns:
     (features_dict<str name, Tensor feature>, Tensor targets)
   """
-  is_training = mode == tf.estimator.ModeKeys.TRAIN
+  is_training = mode == tf_estimator.ModeKeys.TRAIN
   if config and config.use_tpu:
     num_threads = 64
   else:
@@ -555,7 +556,7 @@ def collapse_nested_datasets(example):
   def prepare_for_output(example):
     if not config or not config.use_tpu:
       _summarize_features(example, num_shards)
-    if mode == tf.estimator.ModeKeys.PREDICT:
+    if mode == tf_estimator.ModeKeys.PREDICT:
       example["infer_targets"] = example.pop("targets")
       return example
     else:
@@ -565,7 +566,7 @@ def prepare_for_output(example):
   dataset = dataset.map(prepare_for_output, num_parallel_calls=num_threads)
   dataset = dataset.prefetch(2)
 
-  if mode == tf.estimator.ModeKeys.PREDICT:
+  if mode == tf_estimator.ModeKeys.PREDICT:
     # This is because of a bug in the Estimator that short-circuits prediction
     # if it doesn't see a QueueRunner. DummyQueueRunner implements the
     # minimal expected interface but does nothing.
diff --git a/tensor2tensor/utils/data_reader_test.py b/tensor2tensor/utils/data_reader_test.py
index 3d98d314d..451ba824c 100644
--- a/tensor2tensor/utils/data_reader_test.py
+++ b/tensor2tensor/utils/data_reader_test.py
@@ -31,6 +31,7 @@
 from tensor2tensor.utils import registry
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 @registry.register_problem
@@ -72,7 +73,7 @@ def preprocess_example(self, example, unused_mode, unused_hparams):
 
 def generate_test_data(problem, tmp_dir):
   problem.generate_data(tmp_dir, tmp_dir)
-  return [problem.filepattern(tmp_dir, tf.estimator.ModeKeys.TRAIN)]
+  return [problem.filepattern(tmp_dir, tf_estimator.ModeKeys.TRAIN)]
 
 
 class DataReaderTest(tf.test.TestCase):
@@ -94,7 +95,7 @@ def tearDownClass(cls):
 
   def testBasicExampleReading(self):
     dataset = self.problem.dataset(
-        tf.estimator.ModeKeys.TRAIN,
+        tf_estimator.ModeKeys.TRAIN,
         data_dir=self.data_dir,
         shuffle_files=False)
     examples = dataset.make_one_shot_iterator().get_next()
@@ -113,7 +114,7 @@ def testBasicExampleReading(self):
 
   def testPreprocess(self):
     dataset = self.problem.dataset(
-        tf.estimator.ModeKeys.TRAIN,
+        tf_estimator.ModeKeys.TRAIN,
         data_dir=self.data_dir,
         shuffle_files=False)
     examples = dataset.make_one_shot_iterator().get_next()
@@ -125,7 +126,7 @@ def testPreprocess(self):
   def testLengthFilter(self):
     max_len = 15
     dataset = self.problem.dataset(
-        tf.estimator.ModeKeys.TRAIN,
+        tf_estimator.ModeKeys.TRAIN,
         data_dir=self.data_dir,
         shuffle_files=False)
     dataset = dataset.filter(
diff --git a/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 7a7d2b130..87631b1e3 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -38,6 +38,7 @@
 from tensor2tensor.utils import mlperf_log
 from tensor2tensor.utils import registry
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 FLAGS = tf.flags.FLAGS
 
@@ -204,7 +205,7 @@ def decode_from_dataset(estimator,
   # Build the inference input function
   problem = hparams.problem
   infer_input_fn = problem.make_estimator_input_fn(
-      tf.estimator.ModeKeys.PREDICT, hparams, dataset_kwargs=dataset_kwargs)
+      tf_estimator.ModeKeys.PREDICT, hparams, dataset_kwargs=dataset_kwargs)
 
   predictions, output_dirs = [], []
   for decode_id in range(decode_hp.num_decodes):
diff --git a/tensor2tensor/utils/mtf_model.py b/tensor2tensor/utils/mtf_model.py
index 9f3e707b8..b96d2b80c 100644
--- a/tensor2tensor/utils/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -29,6 +29,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 from tensorflow.contrib.tpu.python.tpu import tpu_estimator
 
@@ -49,7 +50,7 @@ def estimator_model_fn(cls,
     hparams = hparams_lib.copy_hparams(hparams)
     hparams.use_tpu = use_tpu
     # merge decode_hparams into hparams if present
-    if mode == tf.estimator.ModeKeys.PREDICT and decode_hparams is not None:
+    if mode == tf_estimator.ModeKeys.PREDICT and decode_hparams is not None:
       for k, v in six.iteritems(decode_hparams.values()):
         if hasattr(hparams, k) and getattr(hparams, k) != v:
           tf.logging.warning("Overriding hparams.%s with %s from decode_hparams"
@@ -98,7 +99,7 @@ def estimator_model_fn(cls,
     graph = mtf.Graph()
     mesh = mtf.Mesh(graph, "my_mesh", var_placer)
     # PREDICT mode
-    if mode == tf.estimator.ModeKeys.PREDICT:
+    if mode == tf_estimator.ModeKeys.PREDICT:
       return model.estimator_spec_predict(features, mesh, mesh_impl, use_tpu)
 
     logits, loss = model.mtf_model_fn(features, mesh)
@@ -106,7 +107,7 @@ def estimator_model_fn(cls,
       logits = mtf.anonymize(logits)
 
     # TRAIN mode
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
       var_grads = mtf.gradients(
           [loss], [v.outputs[0] for v in graph.trainable_variables])
       lr = learning_rate.learning_rate_schedule(hparams)
@@ -120,10 +121,10 @@ def estimator_model_fn(cls,
 
     tf_loss = lowering.export_to_tf_tensor(loss)
     tf_loss = tf.to_float(tf_loss)
-    if logits and mode != tf.estimator.ModeKeys.TRAIN:
+    if logits and mode != tf_estimator.ModeKeys.TRAIN:
       tf_logits = lowering.export_to_tf_tensor(logits)
 
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
       tf_update_ops = [lowering.lowered_operation(op) for op in update_ops]
       tf_update_ops.append(tf.assign_add(global_step, 1))
       # tf.logging.info("tf_update_ops: {}".format(tf_update_ops))
@@ -148,7 +149,7 @@ def estimator_model_fn(cls,
           listeners=[saver_listener])
 
     # EVAL mode
-    if mode == tf.estimator.ModeKeys.EVAL:
+    if mode == tf_estimator.ModeKeys.EVAL:
       tf_logits = lowering.export_to_tf_tensor(logits)
       return model.estimator_spec_eval(features, tf_logits, labels, tf_loss,
                                        restore_hook, use_tpu)
@@ -171,7 +172,7 @@ def scaffold_fn():
 
       t2t_model.remove_summaries()
       return tpu_estimator.TPUEstimatorSpec(
-          mode=tf.estimator.ModeKeys.TRAIN,
+          mode=tf_estimator.ModeKeys.TRAIN,
           loss=tf_loss,
           train_op=train_op,
           host_call=host_call,
@@ -181,8 +182,8 @@ def scaffold_fn():
       if hparams.warm_start_from:
         t2t_model.initialize_from_ckpt(
             ckpt_dir=hparams.warm_start_from, hparams=hparams)
-      return tf.estimator.EstimatorSpec(
-          tf.estimator.ModeKeys.TRAIN, loss=tf_loss, train_op=train_op,
+      return tf_estimator.EstimatorSpec(
+          tf_estimator.ModeKeys.TRAIN, loss=tf_loss, train_op=train_op,
           training_chief_hooks=[restore_hook, saver_hook])
 
   def estimator_spec_eval(
@@ -210,7 +211,7 @@ def metric_fn(tf_logits, labels):
                   tf_logits, None, tf.identity(labels))
           return eval_metrics
       return tpu_estimator.TPUEstimatorSpec(
-          tf.estimator.ModeKeys.EVAL,
+          tf_estimator.ModeKeys.EVAL,
           evaluation_hooks=[restore_hook],
           loss=loss,
           eval_metrics=(metric_fn, [logits, labels]))
@@ -221,8 +222,8 @@ def metric_fn(tf_logits, labels):
         eval_metrics[metric_name] = metric_fn(logits, features,
                                               features["targets"])
 
-      return tf.estimator.EstimatorSpec(
-          tf.estimator.ModeKeys.EVAL,
+      return tf_estimator.EstimatorSpec(
+          tf_estimator.ModeKeys.EVAL,
           predictions=predictions,
           eval_metric_ops=eval_metrics,
           evaluation_hooks=[restore_hook],
@@ -249,12 +250,12 @@ def estimator_spec_predict(self, features, mesh, mesh_impl, use_tpu):
     if use_tpu:
       t2t_model.remove_summaries()
       return tpu_estimator.TPUEstimatorSpec(
-          mode=tf.estimator.ModeKeys.PREDICT,
+          mode=tf_estimator.ModeKeys.PREDICT,
           predictions=predictions,
           prediction_hooks=[mtf.MtfRestoreHook(lowering)])
     else:
-      return tf.estimator.EstimatorSpec(
-          tf.estimator.ModeKeys.PREDICT,
+      return tf_estimator.EstimatorSpec(
+          tf_estimator.ModeKeys.PREDICT,
           predictions=predictions,
           prediction_hooks=[mtf.MtfRestoreHook(lowering)])
 
diff --git a/tensor2tensor/utils/scheduled_sampling.py b/tensor2tensor/utils/scheduled_sampling.py
index 4be7c9e3c..5c38f527f 100644
--- a/tensor2tensor/utils/scheduled_sampling.py
+++ b/tensor2tensor/utils/scheduled_sampling.py
@@ -37,6 +37,7 @@
 
 from tensor2tensor.layers import common_layers
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 from tensorflow.python.ops import inplace_ops  # pylint: disable=g-direct-tensorflow-import
 
@@ -210,7 +211,7 @@ def __init__(self, t2tmodel, features):
     self._features = features
 
     hparams = self._t2tmodel.hparams
-    assert hparams.mode == tf.estimator.ModeKeys.TRAIN, hparams.mode
+    assert hparams.mode == tf_estimator.ModeKeys.TRAIN, hparams.mode
 
   def infer_fn(self, partial_targets):
     """Computes logits for all timesteps.
diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 1f3726a9a..542f83d81 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -47,6 +47,7 @@
 from tensor2tensor.utils import scheduled_sampling
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 from tensorflow.python.layers import base
 from tensorflow.python.ops import inplace_ops
@@ -150,7 +151,7 @@ class T2TModel(base.Layer):
 
   def __init__(self,
                hparams,
-               mode=tf.estimator.ModeKeys.TRAIN,
+               mode=tf_estimator.ModeKeys.TRAIN,
                problem_hparams=None,
                data_parallelism=None,
                decode_hparams=None,
@@ -174,7 +175,7 @@ def __init__(self,
     default_name = registry.default_name(type(self))
     name = self.REGISTERED_NAME or default_name
     super(T2TModel, self).__init__(
-        trainable=mode == tf.estimator.ModeKeys.TRAIN, name=name, **kwargs)
+        trainable=mode == tf_estimator.ModeKeys.TRAIN, name=name, **kwargs)
 
     if not problem_hparams and hasattr(hparams, "problem_hparams"):
       problem_hparams = hparams.problem_hparams
@@ -217,7 +218,7 @@ def __init__(self,
                          modalities.ModalityType.SYMBOL_ONE_HOT)):
           if (hparams.prepend_mode == "prepend_inputs_full_attention" or
               (hparams.prepend_mode == "prepend_inputs_masked_attention" and
-               mode != tf.estimator.ModeKeys.TRAIN)):
+               mode != tf_estimator.ModeKeys.TRAIN)):
             weights_fn = common_layers.weights_prepend_inputs_to_targets
             hparams.weights_fn[feature_name] = weights_fn
 
@@ -269,11 +270,11 @@ def problem_hparams(self):
 
   @property
   def is_training(self):
-    return self._hparams.mode == tf.estimator.ModeKeys.TRAIN
+    return self._hparams.mode == tf_estimator.ModeKeys.TRAIN
 
   @property
   def is_predicting(self):
-    return self._hparams.mode == tf.estimator.ModeKeys.PREDICT
+    return self._hparams.mode == tf_estimator.ModeKeys.PREDICT
 
   @property
   def has_input(self):
@@ -436,7 +437,7 @@ def model_fn(self, features):
       else:
         logits = self.top(output, features)
         losses["training"] = 0.0
-        if (self._hparams.mode != tf.estimator.ModeKeys.PREDICT and
+        if (self._hparams.mode != tf_estimator.ModeKeys.PREDICT and
             self._hparams.mode != "attack"):
           losses["training"] = self.loss(logits, features)
 
@@ -555,7 +556,7 @@ def _top_single(self, body_output, feature_name, features):
       top = self._hparams.top.get(feature_name, modalities.get_top(modality))
       top_is_pointwise = getattr(top, "pointwise", False)
       last_only = (top_is_pointwise and
-                   self.hparams.mode == tf.estimator.ModeKeys.PREDICT and
+                   self.hparams.mode == tf_estimator.ModeKeys.PREDICT and
                    not self.hparams.force_full_predict)
       if not last_only:
         logits = top(body_output, features.get("targets"),
@@ -726,7 +727,7 @@ def set_mode(self, mode):
     hparams = hparams_lib.copy_hparams(self._original_hparams)
     hparams.add_hparam("mode", mode)
     # When not in training mode, set all forms of dropout to zero.
-    if mode != tf.estimator.ModeKeys.TRAIN:
+    if mode != tf_estimator.ModeKeys.TRAIN:
       for key in hparams.values():
         if key.endswith("dropout") or key == "label_smoothing":
           log_info("Setting hparams.%s to 0.0", key)
@@ -1447,7 +1448,7 @@ def estimator_model_fn(cls,
     Returns:
       TPUEstimatorSpec if use tpu else EstimatorSpec
     """
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
       create_dummy_vars()
     hparams = hparams_lib.copy_hparams(hparams)
 
@@ -1464,7 +1465,7 @@ def estimator_model_fn(cls,
         _reuse=reuse)
 
     # PREDICT mode
-    if mode == tf.estimator.ModeKeys.PREDICT:
+    if mode == tf_estimator.ModeKeys.PREDICT:
       if use_tpu:
         inputs = features.get("inputs")
         if inputs is None:
@@ -1480,7 +1481,7 @@ def estimator_model_fn(cls,
       return model.estimator_spec_predict(features, use_tpu=use_tpu)
 
     # TRAIN and EVAL modes
-    if hparams.eval_run_autoregressive and mode == tf.estimator.ModeKeys.EVAL:
+    if hparams.eval_run_autoregressive and mode == tf_estimator.ModeKeys.EVAL:
       logits, losses_dict = model.eval_autoregressive(features)
     else:
       logits, losses_dict = model(features)  # pylint: disable=not-callable
@@ -1535,12 +1536,12 @@ def estimator_model_fn(cls,
     loss = sum(losses_dict[key] for key in sorted(losses_dict.keys()))
 
     # EVAL mode
-    if mode == tf.estimator.ModeKeys.EVAL:
+    if mode == tf_estimator.ModeKeys.EVAL:
       return model.estimator_spec_eval(features, logits, labels, loss,
                                        losses_dict)
 
     # TRAIN mode
-    assert mode == tf.estimator.ModeKeys.TRAIN
+    assert mode == tf_estimator.ModeKeys.TRAIN
     num_async_replicas = 1
     if config and not use_tpu:
       num_async_replicas = config.t2t_device_info["num_async_replicas"]
@@ -1581,7 +1582,7 @@ def scaffold_fn():
       remove_summaries()
 
       return contrib.tpu().TPUEstimatorSpec(
-          tf.estimator.ModeKeys.TRAIN,
+          tf_estimator.ModeKeys.TRAIN,
           loss=loss,
           train_op=train_op,
           host_call=host_call,
@@ -1595,8 +1596,8 @@ def scaffold_fn():
       if self._hparams.warm_start_from_second:
         self.initialize_from_ckpt(self._hparams.warm_start_from_second)
 
-      return tf.estimator.EstimatorSpec(
-          tf.estimator.ModeKeys.TRAIN, loss=loss, train_op=train_op)
+      return tf_estimator.EstimatorSpec(
+          tf_estimator.ModeKeys.TRAIN, loss=loss, train_op=train_op)
 
   def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
     """Constructs `tf.estimator.EstimatorSpec` for EVAL (evaluation) mode."""
@@ -1638,7 +1639,7 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
 
       eval_metrics_fn_flat_args = _flatten_dict(eval_metrics_fn_args)
       return contrib.tpu().TPUEstimatorSpec(
-          tf.estimator.ModeKeys.EVAL,
+          tf_estimator.ModeKeys.EVAL,
           eval_metrics=(eval_metrics_fn, eval_metrics_fn_flat_args),
           host_call=host_call,
           loss=loss)
@@ -1684,8 +1685,8 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
 
       evaluation_hooks += problem.eval_hooks(features, logits, hparams)
 
-      return tf.estimator.EstimatorSpec(
-          tf.estimator.ModeKeys.EVAL,
+      return tf_estimator.EstimatorSpec(
+          tf_estimator.ModeKeys.EVAL,
           predictions=predictions,
           eval_metric_ops=eval_metrics,
           evaluation_hooks=evaluation_hooks,
@@ -1755,7 +1756,7 @@ def estimator_spec_predict(self, features, use_tpu=False):
 
     export_outputs = {
         tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
-            tf.estimator.export.PredictOutput(export_out)
+            tf_estimator.export.PredictOutput(export_out)
     }
     if use_tpu:
       # Note: important to call this before remove_summaries()
@@ -1767,13 +1768,13 @@ def estimator_spec_predict(self, features, use_tpu=False):
       remove_summaries()
 
       return contrib.tpu().TPUEstimatorSpec(
-          tf.estimator.ModeKeys.PREDICT,
+          tf_estimator.ModeKeys.PREDICT,
           predictions=predictions,
           host_call=host_call,
           export_outputs=export_outputs)
     else:
-      return tf.estimator.EstimatorSpec(
-          tf.estimator.ModeKeys.PREDICT,
+      return tf_estimator.EstimatorSpec(
+          tf_estimator.ModeKeys.PREDICT,
           predictions=predictions,
           export_outputs=export_outputs)
 
@@ -1846,7 +1847,7 @@ def maybe_scheduled_sampling(self, features, logits, losses):
       return (logits, losses)
 
     # Only do scheduled sampling when training.
-    is_training = (hparams.mode == tf.estimator.ModeKeys.TRAIN)
+    is_training = (hparams.mode == tf_estimator.ModeKeys.TRAIN)
     if not is_training:
       tf.logging.info("Running in %s mode. Not using scheduled sampling.",
                       hparams.mode)
@@ -1929,7 +1930,7 @@ def parallel_scheduled_sampling_pass(
         new_logits = self.top(new_body_outputs, new_features)
 
         # Compute loss. Use original features (== labels).
-        if (hparams.mode != tf.estimator.ModeKeys.PREDICT and
+        if (hparams.mode != tf_estimator.ModeKeys.PREDICT and
             hparams.mode != "attack"):
           new_losses["training"] = self.loss(new_logits, features)
         else:
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 3ebe09d1c..99bdeedcd 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -36,6 +36,7 @@
 from tensor2tensor.utils import t2t_model
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 from tensorflow.core.protobuf import rewriter_config_pb2
 from tensorflow.python import debug
@@ -235,7 +236,7 @@ def create_run_config(model_name,
       del run_config_args["master"]
       del run_config_args["evaluation_master"]
   elif is_cloud_async_distributed():
-    run_config_cls = tf.estimator.RunConfig
+    run_config_cls = tf_estimator.RunConfig
     del run_config_args["master"]
     del run_config_args["evaluation_master"]
 
@@ -355,7 +356,7 @@ def guarantee_const_scope():
 
       def tpu_model_fn(features, labels, mode, params):
         """Wrapper model_fn with tpu.rewrite / TPUPartitionedCall."""
-        if mode == tf.estimator.ModeKeys.PREDICT and params["use_tpu"]:
+        if mode == tf_estimator.ModeKeys.PREDICT and params["use_tpu"]:
           batch_config = tpu_estimator.BatchConfig(
               num_batch_threads=2,
               max_batch_size=predict_batch_size,
@@ -384,7 +385,7 @@ def tpu_model_fn(features, labels, mode, params):
         predict_batch_size=predict_batch_size,
         export_saved_model_api_version=api_version_enum_name)
   else:
-    estimator = tf.estimator.Estimator(
+    estimator = tf_estimator.Estimator(
         model_fn=model_fn,
         model_dir=run_config.model_dir,
         config=run_config,
@@ -464,7 +465,7 @@ def eval_steps(self):
 
   def continuous_train_and_eval(self, continuous_eval_predicate_fn=None):
     del continuous_eval_predicate_fn
-    tf.estimator.train_and_evaluate(self._estimator, self._train_spec,
+    tf_estimator.train_and_evaluate(self._estimator, self._train_spec,
                                     self._eval_spec)
     return self.evaluate()
 
@@ -515,7 +516,7 @@ def train_eval_and_decode(self):
       mlperf_log.transformer_print(key=mlperf_log.EVAL_START)
       if self._hparams.mlperf_mode:
         self._decode_hparams.mlperf_decode_step = i + eval_steps
-      self.decode(dataset_split=tf.estimator.ModeKeys.EVAL)
+      self.decode(dataset_split=tf_estimator.ModeKeys.EVAL)
       d_hparams = self._decode_hparams
       if self._hparams.mlperf_mode and d_hparams.mlperf_success:
         mlperf_log.transformer_print(
@@ -594,7 +595,7 @@ def run_std_server(self):
       ValueError: if not enough information is available in the estimator's
         config to create a server.
     """
-    config = tf.estimator.RunConfig()
+    config = tf_estimator.RunConfig()
     server = tf.train.Server(
         config.cluster_spec,
         job_name=config.task_type,
@@ -632,7 +633,7 @@ def continuous_decode_on_train_data(self):
     """Decode from dataset on new checkpoint."""
     for _ in next_checkpoint(self._hparams.model_dir,
                              self._decode_hparams.decode_timeout_mins):
-      self.decode(dataset_split=tf.estimator.ModeKeys.TRAIN)
+      self.decode(dataset_split=tf_estimator.ModeKeys.TRAIN)
 
   def continuous_decode_on_eval_data(self):
     """Decode from dataset on new checkpoint."""
@@ -657,7 +658,7 @@ def continuous_decode_on_eval_data(self):
 
       mlperf_log.transformer_print(key=mlperf_log.EVAL_START)
       self.decode(
-          dataset_split=tf.estimator.ModeKeys.EVAL,
+          dataset_split=tf_estimator.ModeKeys.EVAL,
           checkpoint_path=checkpoint_path)
       d_hparams = self._decode_hparams
       if self._hparams.mlperf_mode and d_hparams.mlperf_success:
@@ -744,12 +745,12 @@ def create_experiment(
 
   # Input fns from Problem
   problem = hparams.problem
-  train_input_fn = problem.make_estimator_input_fn(tf.estimator.ModeKeys.TRAIN,
+  train_input_fn = problem.make_estimator_input_fn(tf_estimator.ModeKeys.TRAIN,
                                                    hparams)
 
   dataset_split = "test" if eval_use_test_set else None
   dataset_kwargs = {"dataset_split": dataset_split}
-  eval_input_fn = problem.make_estimator_input_fn(tf.estimator.ModeKeys.EVAL,
+  eval_input_fn = problem.make_estimator_input_fn(tf_estimator.ModeKeys.EVAL,
                                                   hparams,
                                                   dataset_kwargs=dataset_kwargs)
 
@@ -763,7 +764,7 @@ def compare_fn(best_eval_result, current_eval_result):
     def serving_input_receiver_fn(hparams, decode_hparams, use_tpu):
       return problem.serving_input_fn(hparams, decode_hparams, use_tpu)
 
-    exporter = tf.estimator.BestExporter(
+    exporter = tf_estimator.BestExporter(
         name="best",
         serving_input_receiver_fn=serving_input_receiver_fn,
         compare_fn=compare_fn,
@@ -824,9 +825,9 @@ def serving_input_receiver_fn(hparams, decode_hparams, use_tpu):
   eval_hooks = contrib.learn().monitors.replace_monitors_with_hooks(
       eval_hooks, estimator)
 
-  train_spec = tf.estimator.TrainSpec(
+  train_spec = tf_estimator.TrainSpec(
       train_input_fn, max_steps=train_steps, hooks=train_hooks)
-  eval_spec = tf.estimator.EvalSpec(
+  eval_spec = tf_estimator.EvalSpec(
       eval_input_fn,
       steps=eval_steps,
       hooks=eval_hooks,
diff --git a/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index 2c319fed7..7d1458325 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -26,6 +26,7 @@
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import trainer_lib
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
 class TrainerLibTest(tf.test.TestCase):
@@ -78,14 +79,14 @@ def testModel(self):
 
     # Dataset
     problem = hparams.problem
-    dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN,
+    dataset = problem.dataset(tf_estimator.ModeKeys.TRAIN,
                               algorithmic.TinyAlgo.data_dir)
     dataset = dataset.repeat(None).padded_batch(10, dataset.output_shapes)
     features = dataset.make_one_shot_iterator().get_next()
     features = data_reader.standardize_shapes(features)
 
     # Model
-    model = registry.model("transformer")(hparams, tf.estimator.ModeKeys.TRAIN)
+    model = registry.model("transformer")(hparams, tf_estimator.ModeKeys.TRAIN)
     logits, losses = model(features)
 
     self.assertTrue("training" in losses)
@@ -120,7 +121,7 @@ def testMultipleTargetModalities(self):
 
     # Dataset
     problem = hparams.problem
-    dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN,
+    dataset = problem.dataset(tf_estimator.ModeKeys.TRAIN,
                               algorithmic.TinyAlgo.data_dir)
     dataset = dataset.repeat(None).padded_batch(10, dataset.output_shapes)
     features = dataset.make_one_shot_iterator().get_next()
@@ -128,7 +129,7 @@ def testMultipleTargetModalities(self):
     features["targets_A"] = features["targets_B"] = features["targets"]
 
     # Model
-    model = registry.model("transformer")(hparams, tf.estimator.ModeKeys.TRAIN)
+    model = registry.model("transformer")(hparams, tf_estimator.ModeKeys.TRAIN)
 
     def body(args, mb=model.body):
       out = mb(args)
diff --git a/tensor2tensor/utils/video2gif.py b/tensor2tensor/utils/video2gif.py
index 7a87d7a42..d9ae0fd1f 100644
--- a/tensor2tensor/utils/video2gif.py
+++ b/tensor2tensor/utils/video2gif.py
@@ -42,6 +42,7 @@
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 flags = tf.flags
 FLAGS = flags.FLAGS
@@ -60,7 +61,7 @@ def main(_):
     print("This tool only works for video problems.")
     return
 
-  mode = tf.estimator.ModeKeys.TRAIN
+  mode = tf_estimator.ModeKeys.TRAIN
   hparams = trainer_lib.create_hparams(
       FLAGS.hparams_set,
       FLAGS.hparams,

From c8fe559e0b357389d8754474e1306b6ca9afc4f3 Mon Sep 17 00:00:00 2001
From: Yilei Yang <yileiyang@google.com>
Date: Fri, 15 Apr 2022 11:01:53 -0700
Subject: [PATCH 2712/2720] Remove unused comments related to Python 2
 compatibility.

PiperOrigin-RevId: 442053350
---
 tensor2tensor/models/research/glow_ops_test.py | 1 -
 tensor2tensor/models/research/glow_test.py     | 1 -
 2 files changed, 2 deletions(-)

diff --git a/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 10db10080..5caf324eb 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Lint as: python2, python3
 """Tests for tensor2tensor.models.research.glow_ops."""
 
 from __future__ import absolute_import
diff --git a/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
index fa21e53e2..ffdca2a05 100644
--- a/tensor2tensor/models/research/glow_test.py
+++ b/tensor2tensor/models/research/glow_test.py
@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Lint as: python2, python3
 """Tests for tensor2tensor.models.research.glow_model."""
 
 from __future__ import absolute_import

From 53a1be68727b5d5c3a0d0bf18721013843a49041 Mon Sep 17 00:00:00 2001
From: Santiago Ontanon <santiontanon@google.com>
Date: Tue, 9 Aug 2022 12:07:05 -0700
Subject: [PATCH 2713/2720] Add some documentation to the example packing ops.

PiperOrigin-RevId: 466429011
---
 .../data_generators/ops/pack_sequences_ops.cc | 26 +++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
index 02f49a562..cea7d3674 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
@@ -45,6 +45,9 @@ REGISTER_OP("PackSequences2")
                   return Status::OK();
                 });
 
+// Given a collection of examples, each of which consists of two sequences
+// ('inputs' and 'targets') this op packs them into as few packed/combined
+// examples as possible, to try to minimize padding.
 class PackSequences2Op : public OpKernel {
  public:
   explicit PackSequences2Op(
@@ -56,9 +59,11 @@ class PackSequences2Op : public OpKernel {
     auto targets = ctx->input(1).matrix<int64>();
     int inputs_max_length = ctx->input(2).scalar<int32>()();
     int targets_max_length = ctx->input(3).scalar<int32>()();
-    int n = inputs.dimension(0);
+    int n = inputs.dimension(0);  // Number of examples in the input.
     std::vector<int> inputs_lengths(n);
     std::vector<int> targets_lengths(n);
+    // Calculate, in 'inputs_lengths', the actual length of each input sequence
+    // in "inputs", ignoring padding:
     int padded_inputs_length =
         std::min(static_cast<int>(inputs.dimension(1)), inputs_max_length);
     for (int i = 0; i < n; i++) {
@@ -67,6 +72,8 @@ class PackSequences2Op : public OpKernel {
             inputs_lengths[i]++;
       }
     }
+    // Calculate, in 'targets_lengths', the actual length of each target
+    // sequence in "targets", ignoring padding:
     int padded_targets_length =
         std::min(static_cast<int>(targets.dimension(1)), targets_max_length);
     for (int i = 0; i < n; i++) {
@@ -75,16 +82,24 @@ class PackSequences2Op : public OpKernel {
             targets_lengths[i]++;
       }
     }
-    int num_combined = 0;
+    int num_combined = 0;  // Number of combined examples currently generated.
     std::vector<int> combined_inputs_length;
     std::vector<int> combined_targets_length;
     std::vector<std::vector<int> > combined_sequence_ids;
     for (int seq_id = 0; seq_id < n; seq_id++) {
       int inputs_length = inputs_lengths[seq_id];
       int targets_length = targets_lengths[seq_id];
+      // Try to fit the current example, 'seq_id', into one of the existing
+      // packed examples. The code checks to see if the current example fits in
+      // any of the last 1000 packed examples already generated. If it fits in
+      // any, then the example if packed there. Otherwise, a new packed example
+      // is generated with the new example, and 'num_combined' is increased to
+      // reflect this:
       for (int combined_id = std::max(0, num_combined - 1000); true;
            combined_id++) {
         if (combined_id == num_combined) {
+          // The current example, 'seq_id', did not fit in any of the current
+          // packed examples, so, we generate a new packed example:
           combined_inputs_length.push_back(inputs_length);
           combined_targets_length.push_back(targets_length);
           combined_sequence_ids.push_back(std::vector<int>(1, seq_id));
@@ -95,6 +110,8 @@ class PackSequences2Op : public OpKernel {
              <= inputs_max_length) &&
             (combined_targets_length[combined_id] + targets_length
              <= targets_max_length)) {
+          // The current example, 'seq_id', fits in one of the current packed
+          // examples, 'combined_id', so, we just add it there,
           combined_inputs_length[combined_id] += inputs_length;
           combined_targets_length[combined_id] += targets_length;
           combined_sequence_ids[combined_id].push_back(seq_id);
@@ -148,6 +165,8 @@ class PackSequences2Op : public OpKernel {
     auto targets_position_m = targets_position->matrix<int32>();
     targets_position_m.setZero();
 
+    // Copy the actual sequences from 'inputs' and 'targets' into the
+    // packed/combined examples:
     for (int combined_id = 0; combined_id < num_combined; combined_id++) {
       int inputs_pos = 0;
       int targets_pos = 0;
@@ -230,6 +249,9 @@ struct PackingSpec {
   int segment_id;
 };
 
+// This op generalizes PackSequences2Op to examples that contain an arbitrary
+// number of sequences (rather than assuming there are just inputs and targets).
+// The packing logic is the same.
 class PackSequencesKOp : public OpKernel {
  public:
   explicit PackSequencesKOp(OpKernelConstruction* ctx) : OpKernel(ctx) {

From e18775d084e65eb34e21e237fe2d188589a013c7 Mon Sep 17 00:00:00 2001
From: Jean-Baptiste Lespiau <jblespiau@google.com>
Date: Tue, 13 Sep 2022 08:23:49 -0700
Subject: [PATCH 2714/2720] Replace `tensorflow::Status::OK()` with`
 tensorflow::OkStatus()`.

PiperOrigin-RevId: 474024418
---
 tensor2tensor/data_generators/ops/pack_sequences_ops.cc       | 4 ++--
 tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
index cea7d3674..14357643b 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
@@ -42,7 +42,7 @@ REGISTER_OP("PackSequences2")
                     ctx->set_output(i, ctx->Matrix(ctx->UnknownDim(),
                                                    ctx->UnknownDim()));
                   }
-                  return Status::OK();
+                  return tensorflow::Status();
                 });
 
 // Given a collection of examples, each of which consists of two sequences
@@ -234,7 +234,7 @@ REGISTER_OP("PackSequencesK")
       TF_RETURN_IF_ERROR(
           ctx->set_output("outputs_segmentation", segmentation_shapes));
       TF_RETURN_IF_ERROR(ctx->set_output("outputs_position", position_shapes));
-      return Status::OK();
+      return tensorflow::Status();
     });
 
 typedef int InputIndex;
diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc b/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
index afa8813bd..f5e3c8108 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
@@ -23,7 +23,7 @@ REGISTER_OP("SubwordTextEncoderEncode")
     .Attr("vocab_filename: string")
     .SetShapeFn([](InferenceContext* ctx) {
       ctx->set_output(0, ctx->Vector(ctx->UnknownDim()));
-      return Status::OK();
+      return tensorflow::Status();
     });
 
 class SubwordTextEncoderEncodeOp : public OpKernel {

From 384d527cbef3e3917f52a22f6892c5ea4cb125fd Mon Sep 17 00:00:00 2001
From: Derek Mauro <dmauro@google.com>
Date: Mon, 24 Oct 2022 08:05:06 -0700
Subject: [PATCH 2715/2720] Internal change

PiperOrigin-RevId: 483383652
---
 tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc b/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
index f5e3c8108..ca49ca14f 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
@@ -1,3 +1,5 @@
+#include <memory>
+
 #include "third_party/py/tensor2tensor/data_generators/ops/subword_text_encoder.h"
 #include "third_party/tensorflow/core/framework/op_kernel.h"
 #include "third_party/tensorflow/core/framework/shape_inference.h"
@@ -32,7 +34,7 @@ class SubwordTextEncoderEncodeOp : public OpKernel {
       OpKernelConstruction* ctx) : OpKernel(ctx) {
     std::string vocab_filename;
     OP_REQUIRES_OK(ctx, ctx->GetAttr("vocab_filename", &vocab_filename));
-    encoder_ = absl::make_unique<SubwordTextEncoder>(vocab_filename);
+    encoder_ = std::make_unique<SubwordTextEncoder>(vocab_filename);
   }
 
   void Compute(OpKernelContext* ctx) override {

From 6d25001cd49a7f2408de185fc0e2699f1b8ea70f Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Wed, 21 Dec 2022 11:03:47 -0800
Subject: [PATCH 2716/2720] [NumPy] Remove references to deprecated NumPy type
 aliases.

This change replaces references to a number of deprecated NumPy type aliases (np.bool, np.int, np.float, np.complex, np.object, np.str) with their recommended replacement (bool, int, float, complex, object, str).

NumPy 1.24 drops the deprecated aliases, so we must remove uses before updating NumPy.

PiperOrigin-RevId: 496961661
---
 tensor2tensor/data_generators/allen_brain.py   | 2 +-
 tensor2tensor/data_generators/video_utils.py   | 4 ++--
 tensor2tensor/layers/common_image_attention.py | 2 +-
 tensor2tensor/rl/rl_utils.py                   | 2 +-
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index 9a6341921..9df0e3d2f 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -178,7 +178,7 @@ def random_square_mask(shape, fraction):
   mask = np.ones(shape)
 
   patch_area = shape[0]*shape[1]*fraction
-  patch_dim = np.int(math.floor(math.sqrt(patch_area)))
+  patch_dim = int(math.floor(math.sqrt(patch_area)))
   if patch_area == 0 or patch_dim == 0:
     return mask
 
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 4d0026eeb..d20b074d4 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -96,8 +96,8 @@ def create_border(video, color="blue", border_percent=2):
   color_to_axis = {"blue": 2, "red": 0, "green": 1}
   axis = color_to_axis[color]
   _, _, height, width, _ = video.shape
-  border_height = np.ceil(border_percent * height / 100.0).astype(np.int)
-  border_width = np.ceil(border_percent * width / 100.0).astype(np.int)
+  border_height = np.ceil(border_percent * height / 100.0).astype(int)
+  border_width = np.ceil(border_percent * width / 100.0).astype(int)
   video[:, :, :border_height, :, axis] = 255
   video[:, :, -border_height:, :, axis] = 255
   video[:, :, :, :border_width, axis] = 255
diff --git a/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index 5bb532f64..c7d28162d 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -167,7 +167,7 @@ def get_dilated_1d_attention_mask(
     num_blocks, memory_size, gap_size,
     name="dilated_mask"):
   """Dilated attention with a masking strategy."""
-  mask = np.ones((num_heads, block_size, 2*block_size), np.bool)
+  mask = np.ones((num_heads, block_size, 2*block_size), bool)
 
   # now going over every row to do the right assignment of
   # memory blocks
diff --git a/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index d3ed8680c..5277c36c6 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -395,7 +395,7 @@ def absolute_hinge_difference(arr1, arr2, min_diff=10, dtype=np.uint8):
   Returns:
     array
   """
-  diff = np.abs(arr1.astype(np.int) - arr2, dtype=np.int)
+  diff = np.abs(arr1.astype(int) - arr2, dtype=int)
   return np.maximum(diff - min_diff, 0).astype(dtype)
 
 

From fb75dcd7b7d8379e63c8ab9923c630e12d154228 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Wed, 18 Jan 2023 20:17:29 -0800
Subject: [PATCH 2717/2720] Fix test failures under NumPy 1.24.

NumPy 1.24 release notes: https://numpy.org/devdocs/release/1.24.0-notes.html

The fixes vary, but there are three particularly common changes:
* NumPy 1.24 removes a number of deprecated NumPy type aliases references (np.bool, np.int, np.float, np.complex, np.object, np.str, np.unicode, np.long). This change replaces them with their recommended replacements (bool, int, float, complex, object, str, str, int).
* Under NumPy 1.24 no longer automatically infers dtype=object when ragged sequences are passed to np.array(). See https://numpy.org/neps/nep-0034-infer-dtype-is-object.html . In most cases the fix is to pass dtype=object explicitly, but in some cases where the raggedness seems accidental other fixes were used.
* NumPy 1.24 is pickier about the dtype= option passed to comparison ufuncs.

PiperOrigin-RevId: 503050203
---
 tensor2tensor/envs/gym_env_problem.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/tensor2tensor/envs/gym_env_problem.py b/tensor2tensor/envs/gym_env_problem.py
index 492df2d9b..f85e61781 100644
--- a/tensor2tensor/envs/gym_env_problem.py
+++ b/tensor2tensor/envs/gym_env_problem.py
@@ -33,6 +33,15 @@
 from tensor2tensor.envs import trajectory
 
 
+# This is a compatibility shim introduced to support NumPy 1.24. See:
+# https://numpy.org/neps/nep-0034-infer-dtype-is-object.html
+def _stack(xs):
+  try:
+    return np.stack(xs)
+  except ValueError:
+    return np.stack(np.asarray(xs, dtype=object))
+
+
 class GymEnvProblem(env_problem.EnvProblem):
   """An EnvProblem implemented as a batch of gym envs.
 
@@ -289,7 +298,7 @@ def reset_at(idx):
       for i in range(num_envs_to_reset):
         reset_at(i)
 
-    return np.stack(observations)
+    return _stack(observations)
 
   def _step(self, actions):
     """Takes a step in all environments, shouldn't pre-process or record.
@@ -322,4 +331,4 @@ def apply_step(i):
 
     # Convert each list (observations, rewards, ...) into np.array and return a
     # tuple.
-    return tuple(map(np.stack, [observations, rewards, dones, infos]))
+    return tuple(map(_stack, [observations, rewards, dones, infos]))

From d083bdd7af1d0f54224806e177c11bb69870af2d Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Thu, 16 Feb 2023 18:16:03 -0800
Subject: [PATCH 2718/2720] internal

PiperOrigin-RevId: 510296114
---
 .../data_generators/ops/subword_text_encoder.cc   | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder.cc b/tensor2tensor/data_generators/ops/subword_text_encoder.cc
index dfa2470cc..03150925b 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder.cc
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder.cc
@@ -47,22 +47,29 @@ SubwordTextEncoder::SubwordTextEncoder(const std::string& vocab_filename) {
 }
 
 void SubwordTextEncoder::Encode(absl::string_view text, std::vector<int>* ids) {
+  // Subsequent code can read characters beyond the bound of the string_view
+  // in "text".  For example, U8_NEXT requires that the offset should be
+  // strictly smaller than the length, but this is possible with the code
+  // below.  Ideally, this should not happen, but work around this issue by
+  // using the pointer to circumvent bounds checking until the code or tests
+  // are fixed.
+  const char* ptr = text.data();
+
   ids->clear();
   int token_start = 0;
   int token_end = 0;
   UChar32 c;
   UChar32 next_c;
-  U8_NEXT(text, token_end, text.length(), c);
+  U8_NEXT(ptr, token_end, text.length(), c);
   CHECK_GE(c, 0);
   while (token_end <= text.length()) {
     int next_end = token_end;
-    U8_NEXT(text, next_end, text.length(), next_c);
+    U8_NEXT(ptr, next_end, text.length(), next_c);
     CHECK_GE(next_c, 0);
     // Subtoken break when switching from non-alphanum to alphanum, or when
     // reaching the end of the original token.
     if (u_isalnum(next_c) != u_isalnum(c) || token_end >= text.length()) {
-      absl::string_view next_token =
-          text.substr(token_start, token_end - token_start);
+      absl::string_view next_token(ptr + token_start, token_end - token_start);
       if (next_token != " ") {
         EncodeSubtokens(next_token, ids);
       }

From ba283d94091fc8c5a6dd83da4201e48527f0e5ba Mon Sep 17 00:00:00 2001
From: Fiona Lang <flang@google.com>
Date: Thu, 23 Mar 2023 18:28:18 -0700
Subject: [PATCH 2719/2720] Update some isinstance checks of
 `tf.compat.v1.Variable` to `tf.compat.v2.Variable`.

Also fix calls to the deprecated Variable.initialized_value implementation by directly copying it to the use cases.

These changes are in preparation for changing `tensorflow/python/ops/resource_variable_ops.BaseResourceVariable` to inherit from `tensorflow/python/ops/variables.Variable` instead of `tensorflow/python/ops/variables.VariableV1`.

Since `tensorflow/python/ops/variables.VariableV1` inherits from `tensorflow/python/ops/variables.Variable`, these changes are backwards compatible.

PiperOrigin-RevId: 519016593
---
 tensor2tensor/layers/discretization.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index b0ec46772..99364b4ae 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -903,7 +903,10 @@ def get_vq_codebook(codebook_size, hidden_size):
     with tf.colocate_with(means):
       ema_means = tf.get_variable(
           name="ema_means",
-          initializer=means.initialized_value(),
+          initializer=tf.cond(
+              tf.is_variable_initialized(means),
+              means.read_value,
+              lambda: means.initial_value),
           trainable=False)
 
   return means, ema_means, ema_count

From bafdc1b67730430d38d6ab802cbd51f9d053ba2e Mon Sep 17 00:00:00 2001
From: T2T Team <no-reply@google.com>
Date: Sat, 1 Apr 2023 03:18:41 -0700
Subject: [PATCH 2720/2720] Internal change

PiperOrigin-RevId: 521121284
---
 tensor2tensor/__init__.py                     |  2 +-
 tensor2tensor/bin/__init__.py                 |  2 +-
 tensor2tensor/bin/build_vocab.py              |  2 +-
 tensor2tensor/bin/make_tf_configs.py          |  2 +-
 tensor2tensor/bin/t2t_attack.py               |  2 +-
 tensor2tensor/bin/t2t_avg_all.py              |  2 +-
 tensor2tensor/bin/t2t_bleu.py                 |  2 +-
 tensor2tensor/bin/t2t_datagen.py              |  2 +-
 tensor2tensor/bin/t2t_decoder.py              |  2 +-
 tensor2tensor/bin/t2t_distill.py              |  2 +-
 tensor2tensor/bin/t2t_eval.py                 |  2 +-
 tensor2tensor/bin/t2t_prune.py                |  2 +-
 tensor2tensor/bin/t2t_trainer.py              |  2 +-
 tensor2tensor/bin/t2t_trainer_test.py         |  2 +-
 tensor2tensor/bin/t2t_translate_all.py        |  2 +-
 tensor2tensor/data_generators/__init__.py     |  2 +-
 tensor2tensor/data_generators/algorithmic.py  |  2 +-
 .../data_generators/algorithmic_math.py       |  2 +-
 .../algorithmic_math_deepmind.py              |  2 +-
 .../data_generators/algorithmic_math_test.py  |  2 +-
 .../algorithmic_math_two_variables.py         |  2 +-
 .../data_generators/algorithmic_test.py       |  2 +-
 tensor2tensor/data_generators/all_problems.py |  2 +-
 tensor2tensor/data_generators/allen_brain.py  |  2 +-
 .../data_generators/allen_brain_test.py       |  2 +-
 tensor2tensor/data_generators/audio.py        |  2 +-
 .../data_generators/audio_encoder.py          |  2 +-
 tensor2tensor/data_generators/audio_test.py   |  2 +-
 tensor2tensor/data_generators/babi_qa.py      |  2 +-
 .../data_generators/bair_robot_pushing.py     |  2 +-
 tensor2tensor/data_generators/celeba.py       |  2 +-
 tensor2tensor/data_generators/celeba_test.py  |  2 +-
 tensor2tensor/data_generators/celebahq.py     |  2 +-
 tensor2tensor/data_generators/cifar.py        |  2 +-
 tensor2tensor/data_generators/cipher.py       |  2 +-
 .../data_generators/cleaner_en_xx.py          |  2 +-
 .../data_generators/cnn_dailymail.py          |  2 +-
 tensor2tensor/data_generators/cola.py         |  2 +-
 tensor2tensor/data_generators/common_voice.py |  2 +-
 .../data_generators/common_voice_test.py      |  2 +-
 tensor2tensor/data_generators/conll_ner.py    |  2 +-
 tensor2tensor/data_generators/desc2code.py    |  2 +-
 .../data_generators/desc2code_test.py         |  2 +-
 .../data_generators/dialog_abstract.py        |  2 +-
 .../data_generators/dialog_cornell.py         |  2 +-
 .../data_generators/dialog_dailydialog.py     |  2 +-
 .../data_generators/dialog_opensubtitles.py   |  2 +-
 .../data_generators/dialog_personachat.py     |  2 +-
 tensor2tensor/data_generators/dna_encoder.py  |  2 +-
 .../data_generators/dna_encoder_test.py       |  2 +-
 tensor2tensor/data_generators/enwik8.py       |  2 +-
 tensor2tensor/data_generators/fsns.py         |  2 +-
 .../data_generators/function_docstring.py     |  2 +-
 .../data_generators/gene_expression.py        |  2 +-
 .../data_generators/gene_expression_test.py   |  2 +-
 .../data_generators/generator_utils.py        |  2 +-
 .../data_generators/generator_utils_test.py   |  2 +-
 .../data_generators/google_robot_pushing.py   |  2 +-
 tensor2tensor/data_generators/gym_env.py      |  2 +-
 tensor2tensor/data_generators/gym_env_test.py |  2 +-
 tensor2tensor/data_generators/ice_parsing.py  |  2 +-
 tensor2tensor/data_generators/image_lsun.py   |  2 +-
 tensor2tensor/data_generators/image_utils.py  |  2 +-
 .../data_generators/image_utils_test.py       |  2 +-
 tensor2tensor/data_generators/imagenet.py     |  2 +-
 .../data_generators/imagenet_test.py          |  2 +-
 tensor2tensor/data_generators/imdb.py         |  2 +-
 .../data_generators/inspect_tfrecord.py       |  2 +-
 tensor2tensor/data_generators/lambada.py      |  2 +-
 tensor2tensor/data_generators/librispeech.py  |  2 +-
 tensor2tensor/data_generators/lm1b.py         |  2 +-
 tensor2tensor/data_generators/lm1b_imdb.py    |  2 +-
 tensor2tensor/data_generators/lm1b_mnli.py    |  2 +-
 tensor2tensor/data_generators/mnist.py        |  2 +-
 tensor2tensor/data_generators/moving_mnist.py |  2 +-
 tensor2tensor/data_generators/mrpc.py         |  2 +-
 tensor2tensor/data_generators/mscoco.py       |  2 +-
 tensor2tensor/data_generators/mscoco_test.py  |  2 +-
 .../data_generators/multi_problem.py          |  2 +-
 .../data_generators/multi_problem_v2.py       |  2 +-
 .../data_generators/multi_problem_v2_test.py  |  2 +-
 tensor2tensor/data_generators/multinli.py     |  2 +-
 tensor2tensor/data_generators/ocr.py          |  2 +-
 .../data_generators/ops/pack_sequences_ops.cc | 75 +++++++++----------
 .../ops/pack_sequences_ops_test.py            |  2 +-
 .../ops/subword_text_encoder.cc               |  2 +-
 .../ops/subword_text_encoder_ops.cc           |  8 +-
 .../ops/subword_text_encoder_ops_test.py      |  2 +-
 .../data_generators/paraphrase_ms_coco.py     |  2 +-
 .../paraphrase_ms_coco_test.py                |  2 +-
 .../data_generators/pointer_generator_word.py |  2 +-
 tensor2tensor/data_generators/problem.py      |  2 +-
 .../data_generators/problem_hparams.py        |  2 +-
 tensor2tensor/data_generators/problem_test.py |  2 +-
 .../data_generators/program_search.py         |  2 +-
 .../data_generators/program_search_test.py    |  2 +-
 tensor2tensor/data_generators/ptb.py          |  2 +-
 tensor2tensor/data_generators/qnli.py         |  2 +-
 tensor2tensor/data_generators/quora_qpairs.py |  2 +-
 tensor2tensor/data_generators/rte.py          |  2 +-
 tensor2tensor/data_generators/scitail.py      |  2 +-
 tensor2tensor/data_generators/seq2edits.py    |  2 +-
 tensor2tensor/data_generators/snli.py         |  2 +-
 .../data_generators/speech_recognition.py     |  2 +-
 tensor2tensor/data_generators/squad.py        |  2 +-
 tensor2tensor/data_generators/sst_binary.py   |  2 +-
 tensor2tensor/data_generators/stanford_nli.py |  2 +-
 .../data_generators/style_transfer.py         |  2 +-
 .../data_generators/style_transfer_test.py    |  2 +-
 .../data_generators/subject_verb_agreement.py |  2 +-
 tensor2tensor/data_generators/text_encoder.py |  2 +-
 .../text_encoder_build_subword.py             |  2 +-
 .../data_generators/text_encoder_test.py      |  2 +-
 .../data_generators/text_problems.py          |  2 +-
 .../data_generators/text_problems_test.py     |  2 +-
 tensor2tensor/data_generators/timeseries.py   |  2 +-
 .../timeseries_data_generator.py              |  2 +-
 .../timeseries_data_generator_test.py         |  2 +-
 .../data_generators/timeseries_test.py        |  2 +-
 tensor2tensor/data_generators/tokenizer.py    |  2 +-
 .../data_generators/tokenizer_test.py         |  2 +-
 .../data_generators/transduction_problems.py  |  2 +-
 .../transduction_problems_test.py             |  2 +-
 tensor2tensor/data_generators/translate.py    |  2 +-
 .../data_generators/translate_encs.py         |  2 +-
 .../data_generators/translate_encs_cubbitt.py |  2 +-
 .../data_generators/translate_ende.py         |  2 +-
 .../data_generators/translate_ende_test.py    |  2 +-
 .../data_generators/translate_enes.py         |  2 +-
 .../data_generators/translate_enet.py         |  2 +-
 .../data_generators/translate_enfr.py         |  2 +-
 .../data_generators/translate_enid.py         |  2 +-
 .../data_generators/translate_enmk.py         |  2 +-
 .../data_generators/translate_enro.py         |  2 +-
 .../data_generators/translate_entn.py         |  2 +-
 .../data_generators/translate_envi.py         |  2 +-
 .../data_generators/translate_enzh.py         |  2 +-
 .../data_generators/translate_test.py         |  2 +-
 .../data_generators/video_generated.py        |  2 +-
 tensor2tensor/data_generators/video_utils.py  |  2 +-
 .../data_generators/video_utils_test.py       |  2 +-
 tensor2tensor/data_generators/vqa.py          |  2 +-
 tensor2tensor/data_generators/vqa_utils.py    |  2 +-
 tensor2tensor/data_generators/wiki.py         |  2 +-
 tensor2tensor/data_generators/wiki_lm.py      |  2 +-
 .../data_generators/wiki_multi_problems.py    |  2 +-
 .../data_generators/wiki_revision.py          |  2 +-
 .../data_generators/wiki_revision_utils.py    |  2 +-
 .../data_generators/wikisum/__init__.py       |  2 +-
 .../data_generators/wikisum/generate_vocab.py |  2 +-
 .../wikisum/get_references_commoncrawl.py     |  2 +-
 .../wikisum/get_references_web.py             |  2 +-
 .../get_references_web_single_group.py        |  2 +-
 tensor2tensor/data_generators/wikisum/html.py |  2 +-
 .../wikisum/parallel_launch.py                |  2 +-
 .../wikisum/produce_examples.py               |  2 +-
 .../data_generators/wikisum/utils.py          |  2 +-
 .../data_generators/wikisum/utils_test.py     |  2 +-
 .../data_generators/wikisum/validate_data.py  |  2 +-
 .../data_generators/wikisum/wikisum.py        |  2 +-
 tensor2tensor/data_generators/wikitext103.py  |  2 +-
 tensor2tensor/data_generators/wnli.py         |  2 +-
 tensor2tensor/data_generators/wsj_parsing.py  |  2 +-
 tensor2tensor/data_generators/yelp_full.py    |  2 +-
 .../data_generators/yelp_polarity.py          |  2 +-
 tensor2tensor/envs/__init__.py                |  2 +-
 tensor2tensor/envs/env_problem.py             |  2 +-
 tensor2tensor/envs/env_problem_utils.py       |  2 +-
 tensor2tensor/envs/env_problem_utils_test.py  |  2 +-
 tensor2tensor/envs/gym_env_problem.py         |  2 +-
 tensor2tensor/envs/gym_env_problem_test.py    |  2 +-
 tensor2tensor/envs/gym_spaces_utils.py        |  2 +-
 tensor2tensor/envs/gym_spaces_utils_test.py   |  2 +-
 tensor2tensor/envs/mujoco_problems.py         |  2 +-
 tensor2tensor/envs/mujoco_problems_test.py    |  2 +-
 tensor2tensor/envs/rendered_env_problem.py    |  2 +-
 .../envs/rendered_env_problem_test.py         |  2 +-
 tensor2tensor/envs/tic_tac_toe_env.py         |  2 +-
 tensor2tensor/envs/tic_tac_toe_env_problem.py |  2 +-
 .../envs/tic_tac_toe_env_problem_test.py      |  2 +-
 tensor2tensor/envs/tic_tac_toe_env_test.py    |  2 +-
 tensor2tensor/envs/time_step.py               |  2 +-
 tensor2tensor/envs/time_step_test.py          |  2 +-
 tensor2tensor/envs/trajectory.py              |  2 +-
 tensor2tensor/envs/trajectory_test.py         |  2 +-
 tensor2tensor/insights/__init__.py            |  2 +-
 tensor2tensor/insights/graph.py               |  2 +-
 tensor2tensor/insights/query_processor.py     |  2 +-
 tensor2tensor/insights/server.py              |  2 +-
 tensor2tensor/insights/transformer_model.py   |  2 +-
 tensor2tensor/layers/__init__.py              |  2 +-
 tensor2tensor/layers/area_attention.py        |  2 +-
 tensor2tensor/layers/area_attention_test.py   |  2 +-
 tensor2tensor/layers/common_attention.py      |  2 +-
 tensor2tensor/layers/common_attention_test.py |  2 +-
 tensor2tensor/layers/common_audio.py          |  2 +-
 tensor2tensor/layers/common_hparams.py        |  2 +-
 .../layers/common_image_attention.py          |  2 +-
 .../layers/common_image_attention_test.py     |  2 +-
 tensor2tensor/layers/common_layers.py         |  2 +-
 tensor2tensor/layers/common_layers_test.py    |  2 +-
 tensor2tensor/layers/common_video.py          |  2 +-
 tensor2tensor/layers/common_video_test.py     |  2 +-
 tensor2tensor/layers/discretization.py        |  2 +-
 tensor2tensor/layers/discretization_test.py   |  2 +-
 tensor2tensor/layers/latent_layers.py         |  2 +-
 tensor2tensor/layers/latent_layers_test.py    |  2 +-
 .../layers/message_passing_attention.py       |  2 +-
 tensor2tensor/layers/modalities.py            |  2 +-
 tensor2tensor/layers/modalities_test.py       |  2 +-
 tensor2tensor/layers/ngram.py                 |  2 +-
 tensor2tensor/layers/ngram_test.py            |  2 +-
 .../layers/transformer_glow_layers.py         |  2 +-
 .../layers/transformer_glow_layers_ops.py     |  2 +-
 .../transformer_glow_layers_ops_test.py       |  2 +-
 .../layers/transformer_glow_layers_test.py    |  2 +-
 tensor2tensor/layers/transformer_layers.py    |  2 +-
 tensor2tensor/layers/transformer_memory.py    |  2 +-
 .../layers/transformer_memory_test.py         |  2 +-
 tensor2tensor/layers/vq_discrete.py           |  2 +-
 tensor2tensor/layers/vqa_layers.py            |  2 +-
 tensor2tensor/metrics/__init__.py             |  2 +-
 .../metrics/video_conditional_fvd.py          |  2 +-
 .../metrics/video_conditional_fvd_test.py     |  2 +-
 tensor2tensor/models/__init__.py              |  2 +-
 tensor2tensor/models/basic.py                 |  2 +-
 tensor2tensor/models/basic_test.py            |  2 +-
 tensor2tensor/models/bytenet.py               |  2 +-
 tensor2tensor/models/bytenet_test.py          |  2 +-
 tensor2tensor/models/distillation.py          |  2 +-
 tensor2tensor/models/evolved_transformer.py   |  2 +-
 .../models/evolved_transformer_test.py        |  2 +-
 tensor2tensor/models/image_transformer.py     |  2 +-
 tensor2tensor/models/image_transformer_2d.py  |  2 +-
 .../models/image_transformer_2d_test.py       |  2 +-
 .../models/image_transformer_test.py          |  2 +-
 tensor2tensor/models/lstm.py                  |  2 +-
 tensor2tensor/models/lstm_test.py             |  2 +-
 tensor2tensor/models/mtf_image_transformer.py |  2 +-
 .../models/mtf_image_transformer_test.py      |  2 +-
 tensor2tensor/models/mtf_resnet.py            |  2 +-
 tensor2tensor/models/mtf_transformer.py       |  2 +-
 tensor2tensor/models/mtf_transformer2.py      |  2 +-
 tensor2tensor/models/mtf_transformer_test.py  |  2 +-
 .../neural_architecture_search/__init__.py    |  2 +-
 .../neural_architecture_search/nas_layers.py  |  2 +-
 .../nas_layers_test.py                        |  2 +-
 .../neural_architecture_search/nas_model.py   |  2 +-
 .../nas_model_test.py                         |  2 +-
 tensor2tensor/models/neural_assistant.py      |  2 +-
 tensor2tensor/models/neural_gpu.py            |  2 +-
 tensor2tensor/models/neural_gpu_test.py       |  2 +-
 tensor2tensor/models/research/__init__.py     |  2 +-
 .../models/research/adafactor_experiments.py  |  2 +-
 tensor2tensor/models/research/aligned.py      |  2 +-
 tensor2tensor/models/research/attention_lm.py |  2 +-
 .../models/research/attention_lm_moe.py       |  2 +-
 tensor2tensor/models/research/autoencoders.py |  2 +-
 .../models/research/autoencoders_test.py      |  2 +-
 tensor2tensor/models/research/cycle_gan.py    |  2 +-
 .../models/research/gene_expression.py        |  2 +-
 .../models/research/gene_expression_test.py   |  2 +-
 tensor2tensor/models/research/glow.py         |  2 +-
 .../models/research/glow_init_hook.py         |  2 +-
 tensor2tensor/models/research/glow_ops.py     |  2 +-
 .../models/research/glow_ops_test.py          |  2 +-
 tensor2tensor/models/research/glow_test.py    |  2 +-
 .../models/research/lm_experiments.py         |  2 +-
 tensor2tensor/models/research/moe.py          |  2 +-
 .../models/research/moe_experiments.py        |  2 +-
 .../models/research/multiquery_paper.py       |  2 +-
 tensor2tensor/models/research/neural_stack.py |  2 +-
 .../models/research/neural_stack_test.py      |  2 +-
 .../research/residual_shuffle_exchange.py     |  2 +-
 tensor2tensor/models/research/rl.py           |  2 +-
 .../models/research/shuffle_network.py        |  2 +-
 .../models/research/similarity_transformer.py |  2 +-
 tensor2tensor/models/research/super_lm.py     |  2 +-
 .../models/research/transformer_aux.py        |  2 +-
 .../models/research/transformer_aux_test.py   |  2 +-
 .../models/research/transformer_moe.py        |  2 +-
 .../models/research/transformer_nat.py        |  2 +-
 .../models/research/transformer_parallel.py   |  2 +-
 .../models/research/transformer_revnet.py     |  2 +-
 .../research/transformer_revnet_test.py       |  2 +-
 .../models/research/transformer_seq2edits.py  |  2 +-
 .../models/research/transformer_sketch.py     |  2 +-
 .../models/research/transformer_symshard.py   |  2 +-
 .../models/research/transformer_vae.py        |  2 +-
 .../research/transformer_vae_flow_prior.py    |  2 +-
 .../transformer_vae_flow_prior_ops.py         |  2 +-
 .../models/research/transformer_vae_test.py   |  2 +-
 .../models/research/universal_transformer.py  |  2 +-
 .../research/universal_transformer_test.py    |  2 +-
 .../research/universal_transformer_util.py    |  2 +-
 .../models/research/vqa_attention.py          |  2 +-
 .../models/research/vqa_attention_test.py     |  2 +-
 .../research/vqa_recurrent_self_attention.py  |  2 +-
 .../models/research/vqa_self_attention.py     |  2 +-
 tensor2tensor/models/resnet.py                |  2 +-
 tensor2tensor/models/resnet_test.py           |  2 +-
 tensor2tensor/models/revnet.py                |  2 +-
 tensor2tensor/models/revnet_test.py           |  2 +-
 tensor2tensor/models/shake_shake.py           |  2 +-
 tensor2tensor/models/slicenet.py              |  2 +-
 tensor2tensor/models/slicenet_test.py         |  2 +-
 tensor2tensor/models/text_cnn.py              |  2 +-
 tensor2tensor/models/transformer.py           |  2 +-
 tensor2tensor/models/transformer_test.py      |  2 +-
 tensor2tensor/models/vanilla_gan.py           |  2 +-
 tensor2tensor/models/video/__init__.py        |  2 +-
 tensor2tensor/models/video/base.py            |  2 +-
 tensor2tensor/models/video/base_vae.py        |  2 +-
 .../models/video/basic_deterministic.py       |  2 +-
 .../video/basic_deterministic_params.py       |  2 +-
 .../models/video/basic_deterministic_test.py  |  2 +-
 tensor2tensor/models/video/basic_recurrent.py |  2 +-
 .../models/video/basic_recurrent_test.py      |  2 +-
 .../models/video/basic_stochastic.py          |  2 +-
 .../models/video/basic_stochastic_test.py     |  2 +-
 tensor2tensor/models/video/emily.py           |  2 +-
 tensor2tensor/models/video/emily_test.py      |  2 +-
 tensor2tensor/models/video/epva.py            |  2 +-
 tensor2tensor/models/video/epva_params.py     |  2 +-
 tensor2tensor/models/video/next_frame_glow.py |  2 +-
 tensor2tensor/models/video/nfg_conv3d_test.py |  2 +-
 .../models/video/nfg_conv_lstm_test.py        |  2 +-
 tensor2tensor/models/video/nfg_conv_test.py   |  2 +-
 tensor2tensor/models/video/nfg_interpolate.py |  2 +-
 tensor2tensor/models/video/nfg_test_utils.py  |  2 +-
 tensor2tensor/models/video/nfg_uncond_test.py |  2 +-
 tensor2tensor/models/video/savp.py            |  2 +-
 tensor2tensor/models/video/savp_params.py     |  2 +-
 tensor2tensor/models/video/savp_test.py       |  2 +-
 tensor2tensor/models/video/sv2p.py            |  2 +-
 tensor2tensor/models/video/sv2p_params.py     |  2 +-
 tensor2tensor/models/video/sv2p_test.py       |  2 +-
 tensor2tensor/models/video/tests_utils.py     |  2 +-
 tensor2tensor/models/xception.py              |  2 +-
 tensor2tensor/models/xception_test.py         |  2 +-
 tensor2tensor/problems.py                     |  2 +-
 tensor2tensor/problems_colab.py               |  2 +-
 tensor2tensor/problems_test.py                |  2 +-
 tensor2tensor/rl/__init__.py                  |  2 +-
 tensor2tensor/rl/batch_dqn_agent_test.py      |  2 +-
 tensor2tensor/rl/batch_runner_test.py         |  2 +-
 tensor2tensor/rl/datagen_with_agent.py        |  2 +-
 tensor2tensor/rl/dopamine_connector.py        |  2 +-
 tensor2tensor/rl/envs/__init__.py             |  2 +-
 tensor2tensor/rl/envs/in_graph_batch_env.py   |  2 +-
 tensor2tensor/rl/envs/py_func_batch_env.py    |  2 +-
 tensor2tensor/rl/envs/simulated_batch_env.py  |  2 +-
 .../rl/envs/simulated_batch_gym_env.py        |  2 +-
 tensor2tensor/rl/envs/tf_atari_wrappers.py    |  2 +-
 tensor2tensor/rl/evaluator.py                 |  2 +-
 tensor2tensor/rl/evaluator_test.py            |  2 +-
 tensor2tensor/rl/gym_utils.py                 |  2 +-
 tensor2tensor/rl/gym_utils_test.py            |  2 +-
 tensor2tensor/rl/player.py                    |  2 +-
 tensor2tensor/rl/player_utils.py              |  2 +-
 tensor2tensor/rl/policy_learner.py            |  2 +-
 tensor2tensor/rl/ppo.py                       |  2 +-
 tensor2tensor/rl/ppo_learner.py               |  2 +-
 tensor2tensor/rl/restarter.py                 |  2 +-
 tensor2tensor/rl/restarter_test.py            |  2 +-
 tensor2tensor/rl/rl_utils.py                  |  2 +-
 tensor2tensor/rl/trainer_model_based.py       |  2 +-
 .../rl/trainer_model_based_agent_only.py      |  2 +-
 .../rl/trainer_model_based_params.py          |  2 +-
 .../rl/trainer_model_based_recurrent_test.py  |  2 +-
 .../rl/trainer_model_based_stochastic_test.py |  2 +-
 .../rl/trainer_model_based_sv2p_test.py       |  2 +-
 tensor2tensor/rl/trainer_model_based_test.py  |  2 +-
 tensor2tensor/rl/trainer_model_free.py        |  2 +-
 tensor2tensor/rl/trainer_model_free_test.py   |  2 +-
 .../rl/trainer_model_free_tictactoe_test.py   |  2 +-
 tensor2tensor/serving/__init__.py             |  2 +-
 tensor2tensor/serving/export.py               |  2 +-
 tensor2tensor/serving/query.py                |  2 +-
 tensor2tensor/serving/serving_utils.py        |  2 +-
 .../test_data/example_usr_dir/__init__.py     |  2 +-
 .../test_data/example_usr_dir/my_submodule.py |  2 +-
 tensor2tensor/utils/__init__.py               |  2 +-
 tensor2tensor/utils/adafactor.py              |  2 +-
 tensor2tensor/utils/adafactor_test.py         |  2 +-
 tensor2tensor/utils/adv_attack_utils.py       |  2 +-
 tensor2tensor/utils/avg_checkpoints.py        |  2 +-
 tensor2tensor/utils/beam_search.py            |  2 +-
 tensor2tensor/utils/beam_search_test.py       |  2 +-
 tensor2tensor/utils/bleu_hook.py              |  2 +-
 tensor2tensor/utils/bleu_hook_test.py         |  2 +-
 .../utils/checkpoint_compatibility_test.py    |  2 +-
 tensor2tensor/utils/cloud_mlengine.py         |  2 +-
 tensor2tensor/utils/compute_video_metrics.py  |  2 +-
 tensor2tensor/utils/contrib.py                |  2 +-
 tensor2tensor/utils/data_reader.py            |  2 +-
 tensor2tensor/utils/data_reader_test.py       |  2 +-
 tensor2tensor/utils/decoding.py               |  2 +-
 tensor2tensor/utils/devices.py                |  2 +-
 tensor2tensor/utils/diet.py                   |  2 +-
 tensor2tensor/utils/diet_test.py              |  2 +-
 tensor2tensor/utils/expert_utils.py           |  2 +-
 tensor2tensor/utils/expert_utils_test.py      |  2 +-
 tensor2tensor/utils/flags.py                  |  2 +-
 tensor2tensor/utils/get_rouge.py              |  2 +-
 tensor2tensor/utils/hparam.py                 |  2 +-
 tensor2tensor/utils/hparam_test.py            |  2 +-
 tensor2tensor/utils/hparams_lib.py            |  2 +-
 tensor2tensor/utils/hparams_lib_test.py       |  2 +-
 tensor2tensor/utils/learning_rate.py          |  2 +-
 tensor2tensor/utils/metrics.py                |  2 +-
 tensor2tensor/utils/metrics_hook.py           |  2 +-
 tensor2tensor/utils/metrics_hook_test.py      |  2 +-
 tensor2tensor/utils/metrics_test.py           |  2 +-
 tensor2tensor/utils/misc_utils.py             |  2 +-
 tensor2tensor/utils/misc_utils_test.py        |  2 +-
 tensor2tensor/utils/mlperf_log.py             |  2 +-
 tensor2tensor/utils/mlperf_tags.py            |  2 +-
 tensor2tensor/utils/mtf_model.py              |  2 +-
 tensor2tensor/utils/multistep_optimizer.py    |  2 +-
 .../utils/multistep_optimizer_test.py         |  2 +-
 .../utils/multistep_with_adamoptimizer.py     |  2 +-
 .../multistep_with_adamoptimizer_test.py      |  2 +-
 tensor2tensor/utils/optimize.py               |  2 +-
 tensor2tensor/utils/optimize_test.py          |  2 +-
 .../utils/partial_checkpoint_load_hook.py     |  2 +-
 tensor2tensor/utils/pruning_utils.py          |  2 +-
 tensor2tensor/utils/quantization.py           |  2 +-
 tensor2tensor/utils/registry.py               |  2 +-
 tensor2tensor/utils/registry_test.py          |  2 +-
 tensor2tensor/utils/restore_hook.py           |  2 +-
 tensor2tensor/utils/rouge.py                  |  2 +-
 tensor2tensor/utils/rouge_test.py             |  2 +-
 tensor2tensor/utils/sari_hook.py              |  2 +-
 tensor2tensor/utils/sari_hook_test.py         |  2 +-
 tensor2tensor/utils/scheduled_sampling.py     |  2 +-
 tensor2tensor/utils/t2t_model.py              |  2 +-
 tensor2tensor/utils/t2t_model_test.py         |  2 +-
 tensor2tensor/utils/test_utils.py             |  2 +-
 tensor2tensor/utils/test_utils_test.py        |  2 +-
 tensor2tensor/utils/trainer_lib.py            |  2 +-
 tensor2tensor/utils/trainer_lib_test.py       |  2 +-
 tensor2tensor/utils/update_ops_hook.py        |  2 +-
 tensor2tensor/utils/usr_dir.py                |  2 +-
 tensor2tensor/utils/video/prediction2gif.py   |  2 +-
 tensor2tensor/utils/video/reward_confusion.py |  2 +-
 tensor2tensor/utils/video2gif.py              |  2 +-
 tensor2tensor/utils/video_metrics.py          |  2 +-
 tensor2tensor/utils/video_metrics_test.py     |  2 +-
 tensor2tensor/utils/yellowfin.py              |  2 +-
 tensor2tensor/utils/yellowfin_test.py         |  2 +-
 tensor2tensor/visualization/__init__.py       |  2 +-
 tensor2tensor/visualization/attention.py      |  2 +-
 tensor2tensor/visualization/visualization.py  |  2 +-
 .../visualization/visualization_test.py       |  2 +-
 455 files changed, 494 insertions(+), 495 deletions(-)

diff --git a/tensor2tensor/__init__.py b/tensor2tensor/__init__.py
index ffe5892e3..ff174dd63 100644
--- a/tensor2tensor/__init__.py
+++ b/tensor2tensor/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/__init__.py b/tensor2tensor/bin/__init__.py
index ffe5892e3..ff174dd63 100644
--- a/tensor2tensor/bin/__init__.py
+++ b/tensor2tensor/bin/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/build_vocab.py b/tensor2tensor/bin/build_vocab.py
index 1b3847f06..be3b59fe7 100644
--- a/tensor2tensor/bin/build_vocab.py
+++ b/tensor2tensor/bin/build_vocab.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/make_tf_configs.py b/tensor2tensor/bin/make_tf_configs.py
index e8ec19c08..77b80f450 100644
--- a/tensor2tensor/bin/make_tf_configs.py
+++ b/tensor2tensor/bin/make_tf_configs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_attack.py b/tensor2tensor/bin/t2t_attack.py
index 55b10e18a..336419cfd 100644
--- a/tensor2tensor/bin/t2t_attack.py
+++ b/tensor2tensor/bin/t2t_attack.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_avg_all.py b/tensor2tensor/bin/t2t_avg_all.py
index 4cd702105..6589e84c3 100644
--- a/tensor2tensor/bin/t2t_avg_all.py
+++ b/tensor2tensor/bin/t2t_avg_all.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_bleu.py b/tensor2tensor/bin/t2t_bleu.py
index 2015c79a4..40d691ad6 100644
--- a/tensor2tensor/bin/t2t_bleu.py
+++ b/tensor2tensor/bin/t2t_bleu.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_datagen.py b/tensor2tensor/bin/t2t_datagen.py
index d299bf27c..91c11cd8e 100644
--- a/tensor2tensor/bin/t2t_datagen.py
+++ b/tensor2tensor/bin/t2t_datagen.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_decoder.py b/tensor2tensor/bin/t2t_decoder.py
index ee4036f3d..3e9e41389 100644
--- a/tensor2tensor/bin/t2t_decoder.py
+++ b/tensor2tensor/bin/t2t_decoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_distill.py b/tensor2tensor/bin/t2t_distill.py
index ee1526995..3c86a80e5 100644
--- a/tensor2tensor/bin/t2t_distill.py
+++ b/tensor2tensor/bin/t2t_distill.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_eval.py b/tensor2tensor/bin/t2t_eval.py
index 2ea5608b0..77ca8d7a0 100644
--- a/tensor2tensor/bin/t2t_eval.py
+++ b/tensor2tensor/bin/t2t_eval.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_prune.py b/tensor2tensor/bin/t2t_prune.py
index d66aea769..e43872d32 100644
--- a/tensor2tensor/bin/t2t_prune.py
+++ b/tensor2tensor/bin/t2t_prune.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_trainer.py b/tensor2tensor/bin/t2t_trainer.py
index 4ff922ece..290b45a83 100644
--- a/tensor2tensor/bin/t2t_trainer.py
+++ b/tensor2tensor/bin/t2t_trainer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_trainer_test.py b/tensor2tensor/bin/t2t_trainer_test.py
index 64cc7bc2c..48826748d 100644
--- a/tensor2tensor/bin/t2t_trainer_test.py
+++ b/tensor2tensor/bin/t2t_trainer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/bin/t2t_translate_all.py b/tensor2tensor/bin/t2t_translate_all.py
index 7c78045b8..c938e6aa5 100644
--- a/tensor2tensor/bin/t2t_translate_all.py
+++ b/tensor2tensor/bin/t2t_translate_all.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/__init__.py b/tensor2tensor/data_generators/__init__.py
index ffe5892e3..ff174dd63 100644
--- a/tensor2tensor/data_generators/__init__.py
+++ b/tensor2tensor/data_generators/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic.py b/tensor2tensor/data_generators/algorithmic.py
index 91ffcfe1a..42699a42a 100644
--- a/tensor2tensor/data_generators/algorithmic.py
+++ b/tensor2tensor/data_generators/algorithmic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math.py b/tensor2tensor/data_generators/algorithmic_math.py
index 6b0329dca..85b77e31c 100644
--- a/tensor2tensor/data_generators/algorithmic_math.py
+++ b/tensor2tensor/data_generators/algorithmic_math.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math_deepmind.py b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
index 86a6c0603..9069e6e00 100644
--- a/tensor2tensor/data_generators/algorithmic_math_deepmind.py
+++ b/tensor2tensor/data_generators/algorithmic_math_deepmind.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math_test.py b/tensor2tensor/data_generators/algorithmic_math_test.py
index b9d89a575..e1583f3f4 100644
--- a/tensor2tensor/data_generators/algorithmic_math_test.py
+++ b/tensor2tensor/data_generators/algorithmic_math_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_math_two_variables.py b/tensor2tensor/data_generators/algorithmic_math_two_variables.py
index 18d7704ca..594ef06dc 100644
--- a/tensor2tensor/data_generators/algorithmic_math_two_variables.py
+++ b/tensor2tensor/data_generators/algorithmic_math_two_variables.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/algorithmic_test.py b/tensor2tensor/data_generators/algorithmic_test.py
index 99b1cb4d5..28bc21923 100644
--- a/tensor2tensor/data_generators/algorithmic_test.py
+++ b/tensor2tensor/data_generators/algorithmic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/all_problems.py b/tensor2tensor/data_generators/all_problems.py
index 4f6a1c494..253a8e331 100644
--- a/tensor2tensor/data_generators/all_problems.py
+++ b/tensor2tensor/data_generators/all_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index 9df0e3d2f..cc05b4599 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/allen_brain_test.py b/tensor2tensor/data_generators/allen_brain_test.py
index 8ad44025e..90f87741b 100644
--- a/tensor2tensor/data_generators/allen_brain_test.py
+++ b/tensor2tensor/data_generators/allen_brain_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio.py b/tensor2tensor/data_generators/audio.py
index 10c6bd9f2..0543a902c 100644
--- a/tensor2tensor/data_generators/audio.py
+++ b/tensor2tensor/data_generators/audio.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio_encoder.py b/tensor2tensor/data_generators/audio_encoder.py
index e8b76dde4..1d8de1a05 100644
--- a/tensor2tensor/data_generators/audio_encoder.py
+++ b/tensor2tensor/data_generators/audio_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/audio_test.py b/tensor2tensor/data_generators/audio_test.py
index 52e9ca32e..adf86fbfb 100644
--- a/tensor2tensor/data_generators/audio_test.py
+++ b/tensor2tensor/data_generators/audio_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/babi_qa.py b/tensor2tensor/data_generators/babi_qa.py
index 1f37a8021..db5c8ef1f 100644
--- a/tensor2tensor/data_generators/babi_qa.py
+++ b/tensor2tensor/data_generators/babi_qa.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index 9b5ca4885..9ceb834cf 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celeba.py b/tensor2tensor/data_generators/celeba.py
index 9ca8a5877..8fe0547a6 100644
--- a/tensor2tensor/data_generators/celeba.py
+++ b/tensor2tensor/data_generators/celeba.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celeba_test.py b/tensor2tensor/data_generators/celeba_test.py
index 8206e38a7..f9c5a62b7 100644
--- a/tensor2tensor/data_generators/celeba_test.py
+++ b/tensor2tensor/data_generators/celeba_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/celebahq.py b/tensor2tensor/data_generators/celebahq.py
index 5a590efc7..9383f0da0 100644
--- a/tensor2tensor/data_generators/celebahq.py
+++ b/tensor2tensor/data_generators/celebahq.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cifar.py b/tensor2tensor/data_generators/cifar.py
index dd26c0ac5..a5dfc01af 100644
--- a/tensor2tensor/data_generators/cifar.py
+++ b/tensor2tensor/data_generators/cifar.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cipher.py b/tensor2tensor/data_generators/cipher.py
index c1b6d2139..a6b55a2c9 100644
--- a/tensor2tensor/data_generators/cipher.py
+++ b/tensor2tensor/data_generators/cipher.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cleaner_en_xx.py b/tensor2tensor/data_generators/cleaner_en_xx.py
index 4f2905a9d..2d95b6045 100644
--- a/tensor2tensor/data_generators/cleaner_en_xx.py
+++ b/tensor2tensor/data_generators/cleaner_en_xx.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cnn_dailymail.py b/tensor2tensor/data_generators/cnn_dailymail.py
index 08dbf5123..e87d17d14 100644
--- a/tensor2tensor/data_generators/cnn_dailymail.py
+++ b/tensor2tensor/data_generators/cnn_dailymail.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/cola.py b/tensor2tensor/data_generators/cola.py
index 558f9b05d..366e7db8d 100644
--- a/tensor2tensor/data_generators/cola.py
+++ b/tensor2tensor/data_generators/cola.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/common_voice.py b/tensor2tensor/data_generators/common_voice.py
index 26df843d1..ef3808f65 100644
--- a/tensor2tensor/data_generators/common_voice.py
+++ b/tensor2tensor/data_generators/common_voice.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/common_voice_test.py b/tensor2tensor/data_generators/common_voice_test.py
index 675c71b5d..3798b4240 100644
--- a/tensor2tensor/data_generators/common_voice_test.py
+++ b/tensor2tensor/data_generators/common_voice_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/conll_ner.py b/tensor2tensor/data_generators/conll_ner.py
index 0d80a7f91..52a054223 100644
--- a/tensor2tensor/data_generators/conll_ner.py
+++ b/tensor2tensor/data_generators/conll_ner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/desc2code.py b/tensor2tensor/data_generators/desc2code.py
index efd527ff0..3e0de47b9 100644
--- a/tensor2tensor/data_generators/desc2code.py
+++ b/tensor2tensor/data_generators/desc2code.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/desc2code_test.py b/tensor2tensor/data_generators/desc2code_test.py
index dc8ea0cfd..5c2b0635c 100644
--- a/tensor2tensor/data_generators/desc2code_test.py
+++ b/tensor2tensor/data_generators/desc2code_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_abstract.py b/tensor2tensor/data_generators/dialog_abstract.py
index 16ae87159..266a900ed 100644
--- a/tensor2tensor/data_generators/dialog_abstract.py
+++ b/tensor2tensor/data_generators/dialog_abstract.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_cornell.py b/tensor2tensor/data_generators/dialog_cornell.py
index a917e4e0b..5b60afd1f 100644
--- a/tensor2tensor/data_generators/dialog_cornell.py
+++ b/tensor2tensor/data_generators/dialog_cornell.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_dailydialog.py b/tensor2tensor/data_generators/dialog_dailydialog.py
index 9bd62bce8..6f885ea94 100644
--- a/tensor2tensor/data_generators/dialog_dailydialog.py
+++ b/tensor2tensor/data_generators/dialog_dailydialog.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_opensubtitles.py b/tensor2tensor/data_generators/dialog_opensubtitles.py
index 2a8adc90b..1584fd984 100644
--- a/tensor2tensor/data_generators/dialog_opensubtitles.py
+++ b/tensor2tensor/data_generators/dialog_opensubtitles.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dialog_personachat.py b/tensor2tensor/data_generators/dialog_personachat.py
index d6561c96c..921e8cc57 100644
--- a/tensor2tensor/data_generators/dialog_personachat.py
+++ b/tensor2tensor/data_generators/dialog_personachat.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dna_encoder.py b/tensor2tensor/data_generators/dna_encoder.py
index 38dbfbb5e..e79006920 100644
--- a/tensor2tensor/data_generators/dna_encoder.py
+++ b/tensor2tensor/data_generators/dna_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/dna_encoder_test.py b/tensor2tensor/data_generators/dna_encoder_test.py
index 65202ebd8..46f9b54aa 100644
--- a/tensor2tensor/data_generators/dna_encoder_test.py
+++ b/tensor2tensor/data_generators/dna_encoder_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/enwik8.py b/tensor2tensor/data_generators/enwik8.py
index 772380f34..58e262b07 100644
--- a/tensor2tensor/data_generators/enwik8.py
+++ b/tensor2tensor/data_generators/enwik8.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
index a9c0c2210..ac012e8ed 100644
--- a/tensor2tensor/data_generators/fsns.py
+++ b/tensor2tensor/data_generators/fsns.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/function_docstring.py b/tensor2tensor/data_generators/function_docstring.py
index e604dfe57..93a5e830d 100644
--- a/tensor2tensor/data_generators/function_docstring.py
+++ b/tensor2tensor/data_generators/function_docstring.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gene_expression.py b/tensor2tensor/data_generators/gene_expression.py
index cd0593142..0244edda8 100644
--- a/tensor2tensor/data_generators/gene_expression.py
+++ b/tensor2tensor/data_generators/gene_expression.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gene_expression_test.py b/tensor2tensor/data_generators/gene_expression_test.py
index 74f7b5d66..c23c869c0 100644
--- a/tensor2tensor/data_generators/gene_expression_test.py
+++ b/tensor2tensor/data_generators/gene_expression_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/generator_utils.py b/tensor2tensor/data_generators/generator_utils.py
index 89f182a65..0a56c7ae0 100644
--- a/tensor2tensor/data_generators/generator_utils.py
+++ b/tensor2tensor/data_generators/generator_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/generator_utils_test.py b/tensor2tensor/data_generators/generator_utils_test.py
index 55ee6e7ad..f8e1926e1 100644
--- a/tensor2tensor/data_generators/generator_utils_test.py
+++ b/tensor2tensor/data_generators/generator_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/google_robot_pushing.py b/tensor2tensor/data_generators/google_robot_pushing.py
index 9c4b9f49d..a35670ef8 100644
--- a/tensor2tensor/data_generators/google_robot_pushing.py
+++ b/tensor2tensor/data_generators/google_robot_pushing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index b8db3a19e..b176ab857 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/gym_env_test.py b/tensor2tensor/data_generators/gym_env_test.py
index 197134eb3..29d6afc68 100644
--- a/tensor2tensor/data_generators/gym_env_test.py
+++ b/tensor2tensor/data_generators/gym_env_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ice_parsing.py b/tensor2tensor/data_generators/ice_parsing.py
index f97c7fde2..1f7380515 100644
--- a/tensor2tensor/data_generators/ice_parsing.py
+++ b/tensor2tensor/data_generators/ice_parsing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_lsun.py b/tensor2tensor/data_generators/image_lsun.py
index 0755f2b78..7ad16d897 100644
--- a/tensor2tensor/data_generators/image_lsun.py
+++ b/tensor2tensor/data_generators/image_lsun.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 04955f671..be402162d 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/image_utils_test.py b/tensor2tensor/data_generators/image_utils_test.py
index 43dfee321..d275f2844 100644
--- a/tensor2tensor/data_generators/image_utils_test.py
+++ b/tensor2tensor/data_generators/image_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imagenet.py b/tensor2tensor/data_generators/imagenet.py
index af0b130b1..175016b86 100644
--- a/tensor2tensor/data_generators/imagenet.py
+++ b/tensor2tensor/data_generators/imagenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imagenet_test.py b/tensor2tensor/data_generators/imagenet_test.py
index 034ffad70..688153590 100644
--- a/tensor2tensor/data_generators/imagenet_test.py
+++ b/tensor2tensor/data_generators/imagenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/imdb.py b/tensor2tensor/data_generators/imdb.py
index 93b20bfd3..b7c62b4e3 100644
--- a/tensor2tensor/data_generators/imdb.py
+++ b/tensor2tensor/data_generators/imdb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/inspect_tfrecord.py b/tensor2tensor/data_generators/inspect_tfrecord.py
index 0322a4bde..592215bc4 100644
--- a/tensor2tensor/data_generators/inspect_tfrecord.py
+++ b/tensor2tensor/data_generators/inspect_tfrecord.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lambada.py b/tensor2tensor/data_generators/lambada.py
index 425ac2268..8b00f6a4e 100644
--- a/tensor2tensor/data_generators/lambada.py
+++ b/tensor2tensor/data_generators/lambada.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/librispeech.py b/tensor2tensor/data_generators/librispeech.py
index 8e5062cef..90d073222 100644
--- a/tensor2tensor/data_generators/librispeech.py
+++ b/tensor2tensor/data_generators/librispeech.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b.py b/tensor2tensor/data_generators/lm1b.py
index 75ab01632..8f209ed47 100644
--- a/tensor2tensor/data_generators/lm1b.py
+++ b/tensor2tensor/data_generators/lm1b.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b_imdb.py b/tensor2tensor/data_generators/lm1b_imdb.py
index 6d50ec7ed..e4789e822 100644
--- a/tensor2tensor/data_generators/lm1b_imdb.py
+++ b/tensor2tensor/data_generators/lm1b_imdb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/lm1b_mnli.py b/tensor2tensor/data_generators/lm1b_mnli.py
index c28c51ac3..6e7b385f7 100644
--- a/tensor2tensor/data_generators/lm1b_mnli.py
+++ b/tensor2tensor/data_generators/lm1b_mnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mnist.py b/tensor2tensor/data_generators/mnist.py
index f49e53664..ba5f85b4b 100644
--- a/tensor2tensor/data_generators/mnist.py
+++ b/tensor2tensor/data_generators/mnist.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/moving_mnist.py b/tensor2tensor/data_generators/moving_mnist.py
index 1f4600804..445b45b8e 100644
--- a/tensor2tensor/data_generators/moving_mnist.py
+++ b/tensor2tensor/data_generators/moving_mnist.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mrpc.py b/tensor2tensor/data_generators/mrpc.py
index 428632596..bdf4a3b5c 100644
--- a/tensor2tensor/data_generators/mrpc.py
+++ b/tensor2tensor/data_generators/mrpc.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mscoco.py b/tensor2tensor/data_generators/mscoco.py
index c45c477e6..8e8ed8686 100644
--- a/tensor2tensor/data_generators/mscoco.py
+++ b/tensor2tensor/data_generators/mscoco.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/mscoco_test.py b/tensor2tensor/data_generators/mscoco_test.py
index 3e5ab687a..dc2e74d40 100644
--- a/tensor2tensor/data_generators/mscoco_test.py
+++ b/tensor2tensor/data_generators/mscoco_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem.py b/tensor2tensor/data_generators/multi_problem.py
index 4e2ea12d7..b5cb4e638 100644
--- a/tensor2tensor/data_generators/multi_problem.py
+++ b/tensor2tensor/data_generators/multi_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem_v2.py b/tensor2tensor/data_generators/multi_problem_v2.py
index ee79b4e8f..dfdf66ece 100644
--- a/tensor2tensor/data_generators/multi_problem_v2.py
+++ b/tensor2tensor/data_generators/multi_problem_v2.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multi_problem_v2_test.py b/tensor2tensor/data_generators/multi_problem_v2_test.py
index d28f298e4..d21b98ecb 100644
--- a/tensor2tensor/data_generators/multi_problem_v2_test.py
+++ b/tensor2tensor/data_generators/multi_problem_v2_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/multinli.py b/tensor2tensor/data_generators/multinli.py
index b09dcbe18..d784157ae 100644
--- a/tensor2tensor/data_generators/multinli.py
+++ b/tensor2tensor/data_generators/multinli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ocr.py b/tensor2tensor/data_generators/ocr.py
index d01646318..395bbf7e6 100644
--- a/tensor2tensor/data_generators/ocr.py
+++ b/tensor2tensor/data_generators/ocr.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ops/pack_sequences_ops.cc b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
index 14357643b..211f18d2a 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops.cc
@@ -55,10 +55,10 @@ class PackSequences2Op : public OpKernel {
   }
 
   void Compute(OpKernelContext* ctx) override {
-    auto inputs = ctx->input(0).matrix<int64>();
-    auto targets = ctx->input(1).matrix<int64>();
-    int inputs_max_length = ctx->input(2).scalar<int32>()();
-    int targets_max_length = ctx->input(3).scalar<int32>()();
+    auto inputs = ctx->input(0).matrix<int64_t>();
+    auto targets = ctx->input(1).matrix<int64_t>();
+    int inputs_max_length = ctx->input(2).scalar<int32_t>()();
+    int targets_max_length = ctx->input(3).scalar<int32_t>()();
     int n = inputs.dimension(0);  // Number of examples in the input.
     std::vector<int> inputs_lengths(n);
     std::vector<int> targets_lengths(n);
@@ -120,49 +120,49 @@ class PackSequences2Op : public OpKernel {
       }
     }
 
-    auto output_shape_inputs = TensorShape(
-        {static_cast<int64>(num_combined),
-         static_cast<int64>(inputs_max_length)});
-    auto output_shape_targets = TensorShape(
-        {static_cast<int64>(num_combined),
-         static_cast<int64>(targets_max_length)});
+    auto output_shape_inputs =
+        TensorShape({static_cast<int64_t>(num_combined),
+                     static_cast<int64_t>(inputs_max_length)});
+    auto output_shape_targets =
+        TensorShape({static_cast<int64_t>(num_combined),
+                     static_cast<int64_t>(targets_max_length)});
 
     Tensor* inputs_packed;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(
         0, output_shape_inputs, &inputs_packed));
-    auto inputs_packed_m = inputs_packed->matrix<int64>();
+    auto inputs_packed_m = inputs_packed->matrix<int64_t>();
     inputs_packed_m.setZero();
 
     Tensor* inputs_segmentation;
     OP_REQUIRES_OK(
         ctx, ctx->allocate_output(
             1, output_shape_inputs, &inputs_segmentation));
-    auto inputs_segmentation_m = inputs_segmentation->matrix<int32>();
+    auto inputs_segmentation_m = inputs_segmentation->matrix<int32_t>();
     inputs_segmentation_m.setZero();
 
     Tensor* inputs_position;
     OP_REQUIRES_OK(
         ctx, ctx->allocate_output(2, output_shape_inputs, &inputs_position));
-    auto inputs_position_m = inputs_position->matrix<int32>();
+    auto inputs_position_m = inputs_position->matrix<int32_t>();
     inputs_position_m.setZero();
 
     Tensor* targets_packed;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(
         3, output_shape_targets, &targets_packed));
-    auto targets_packed_m = targets_packed->matrix<int64>();
+    auto targets_packed_m = targets_packed->matrix<int64_t>();
     targets_packed_m.setZero();
 
     Tensor* targets_segmentation;
     OP_REQUIRES_OK(
         ctx, ctx->allocate_output(
             4, output_shape_targets, &targets_segmentation));
-    auto targets_segmentation_m = targets_segmentation->matrix<int32>();
+    auto targets_segmentation_m = targets_segmentation->matrix<int32_t>();
     targets_segmentation_m.setZero();
 
     Tensor* targets_position;
     OP_REQUIRES_OK(
         ctx, ctx->allocate_output(5, output_shape_targets, &targets_position));
-    auto targets_position_m = targets_position->matrix<int32>();
+    auto targets_position_m = targets_position->matrix<int32_t>();
     targets_position_m.setZero();
 
     // Copy the actual sequences from 'inputs' and 'targets' into the
@@ -276,7 +276,7 @@ class PackSequencesKOp : public OpKernel {
 
     std::map<InputIndex, int> max_lengths;
     for (InputIndex i = 0; i < max_lengths_list.size(); i++) {
-      max_lengths[i] = max_lengths_list[i].scalar<int32>()();
+      max_lengths[i] = max_lengths_list[i].scalar<int32_t>()();
     }
 
     int n = inputs.begin()->dim_size(0);
@@ -348,9 +348,8 @@ class PackSequencesKOp : public OpKernel {
         ctx, ctx->output_list("outputs_position", &outputs_position));
 
     for (InputIndex i = 0; i < inputs.size(); i++) {
-      TensorShape output_shape_2d = {
-        static_cast<int64>(num_combined),
-        static_cast<int64>(max_lengths[i])};
+      TensorShape output_shape_2d = {static_cast<int64_t>(num_combined),
+                                     static_cast<int64_t>(max_lengths[i])};
 
       TensorShape output_shape = output_shape_2d;
       if (inputs[i].dims() == 3) {
@@ -369,8 +368,8 @@ class PackSequencesKOp : public OpKernel {
       OP_REQUIRES_OK(ctx,
                      outputs_position.allocate(i, output_shape_2d, &position));
 
-      auto segmentation_eigen = segmentation->matrix<int32>();
-      auto position_eigen = position->matrix<int32>();
+      auto segmentation_eigen = segmentation->matrix<int32_t>();
+      auto position_eigen = position->matrix<int32_t>();
 
       SetZero(ctx, packed);
       segmentation_eigen.setZero();
@@ -394,9 +393,9 @@ class PackSequencesKOp : public OpKernel {
       case tensorflow::DT_FLOAT:
         return GetInputLengths<float>(ctx, input, padded_input_length);
       case tensorflow::DT_INT32:
-        return GetInputLengths<int32>(ctx, input, padded_input_length);
+        return GetInputLengths<int32_t>(ctx, input, padded_input_length);
       case tensorflow::DT_INT64:
-        return GetInputLengths<int64>(ctx, input, padded_input_length);
+        return GetInputLengths<int64_t>(ctx, input, padded_input_length);
       default:
         ctx->CtxFailure(
             tensorflow::errors::InvalidArgument("unsupported input dtype"));
@@ -468,10 +467,10 @@ class PackSequencesKOp : public OpKernel {
         SetZero<float>(ctx, inputs);
         break;
       case tensorflow::DT_INT32:
-        SetZero<int32>(ctx, inputs);
+        SetZero<int32_t>(ctx, inputs);
         break;
       case tensorflow::DT_INT64:
-        SetZero<int64>(ctx, inputs);
+        SetZero<int64_t>(ctx, inputs);
         break;
       default:
         ctx->CtxFailure(
@@ -495,8 +494,8 @@ class PackSequencesKOp : public OpKernel {
   }
 
   void PackSequence(OpKernelContext* ctx, const Tensor& inputs, Tensor* packed,
-                    TTypes<int32, 2>::Tensor segmentation,
-                    TTypes<int32, 2>::Tensor position,
+                    TTypes<int32_t, 2>::Tensor segmentation,
+                    TTypes<int32_t, 2>::Tensor position,
                     const PackingSpec& spec) {
     switch (inputs.dtype()) {
       case tensorflow::DT_FLOAT:
@@ -508,12 +507,12 @@ class PackSequencesKOp : public OpKernel {
             ctx, inputs, packed, segmentation, position, spec);
         break;
       case tensorflow::DT_INT32:
-        PackSequence<int32>(
-            ctx, inputs, packed, segmentation, position, spec);
+        PackSequence<int32_t>(ctx, inputs, packed, segmentation, position,
+                              spec);
         break;
       case tensorflow::DT_INT64:
-        PackSequence<int64>(
-            ctx, inputs, packed, segmentation, position, spec);
+        PackSequence<int64_t>(ctx, inputs, packed, segmentation, position,
+                              spec);
         break;
       default:
         ctx->CtxFailure(
@@ -523,8 +522,8 @@ class PackSequencesKOp : public OpKernel {
 
   template <typename T>
   void PackSequence(OpKernelContext* ctx, const Tensor& inputs, Tensor* packed,
-                    TTypes<int32, 2>::Tensor segmentation,
-                    TTypes<int32, 2>::Tensor position,
+                    TTypes<int32_t, 2>::Tensor segmentation,
+                    TTypes<int32_t, 2>::Tensor position,
                     const PackingSpec& spec) {
     switch (inputs.dims()) {
       case 2:
@@ -555,8 +554,8 @@ class PackSequencesKOp : public OpKernel {
   void PackSequence(OpKernelContext* ctx,
                     const typename TTypes<const T, 2>::Tensor& inputs,
                     typename TTypes<T, 2>::Tensor packed,
-                    TTypes<int32, 2>::Tensor segmentation,
-                    TTypes<int32, 2>::Tensor position,
+                    TTypes<int32_t, 2>::Tensor segmentation,
+                    TTypes<int32_t, 2>::Tensor position,
                     const PackingSpec& spec) {
     for (int i = 0; i < spec.seq_length; i++) {
       packed(spec.batch_pos, spec.offset + i) = inputs(spec.seq_id, i);
@@ -569,8 +568,8 @@ class PackSequencesKOp : public OpKernel {
   void PackSequence(OpKernelContext* ctx,
                     const typename TTypes<const T, 3>::Tensor& inputs,
                     typename TTypes<T, 3>::Tensor packed,
-                    TTypes<int32, 2>::Tensor segmentation,
-                    TTypes<int32, 2>::Tensor position,
+                    TTypes<int32_t, 2>::Tensor segmentation,
+                    TTypes<int32_t, 2>::Tensor position,
                     const PackingSpec& spec) {
     for (int i = 0; i < spec.seq_length; i++) {
       for (int k = 0; k < inputs.dimension(2); k++) {
diff --git a/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
index f46eee526..48b676f4e 100644
--- a/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
+++ b/tensor2tensor/data_generators/ops/pack_sequences_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ops/subword_text_encoder.cc b/tensor2tensor/data_generators/ops/subword_text_encoder.cc
index 03150925b..fba7ba2a4 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder.cc
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder.cc
@@ -14,7 +14,7 @@ namespace {
 using ::tensorflow::Env;
 
 // End of Sequence token ID to insert at end of encoded text.
-constexpr int64 kEosTokenId = 1;
+constexpr int64_t kEosTokenId = 1;
 
 }  // namespace
 
diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc b/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
index ca49ca14f..7d89b7d65 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder_ops.cc
@@ -46,10 +46,10 @@ class SubwordTextEncoderEncodeOp : public OpKernel {
     encoder_->Encode(s, &encoded_ids);
     Tensor* encoded;
     OP_REQUIRES_OK(
-        ctx,
-        ctx->allocate_output(0, TensorShape(
-            {static_cast<int64>(encoded_ids.size())}), &encoded));
-    auto encoded_vec = encoded->vec<int64>();
+        ctx, ctx->allocate_output(
+                 0, TensorShape({static_cast<int64_t>(encoded_ids.size())}),
+                 &encoded));
+    auto encoded_vec = encoded->vec<int64_t>();
     // TODO(noam): find someone who remembers c++ eigen and ask the proper way
     // to copy a std::Vector to an Eigen whatever-this-is
     for (int i = 0; i < encoded_ids.size(); i++) {
diff --git a/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py b/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
index b654f5867..952838f3a 100644
--- a/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
+++ b/tensor2tensor/data_generators/ops/subword_text_encoder_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/paraphrase_ms_coco.py b/tensor2tensor/data_generators/paraphrase_ms_coco.py
index 9713bfab1..84de12859 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/paraphrase_ms_coco_test.py b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
index fb40709a2..95649d80c 100644
--- a/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
+++ b/tensor2tensor/data_generators/paraphrase_ms_coco_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/pointer_generator_word.py b/tensor2tensor/data_generators/pointer_generator_word.py
index c3f2549df..257b3d7bb 100644
--- a/tensor2tensor/data_generators/pointer_generator_word.py
+++ b/tensor2tensor/data_generators/pointer_generator_word.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index 30d63f3bb..20c421ef9 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem_hparams.py b/tensor2tensor/data_generators/problem_hparams.py
index 09a625be7..6ad656cbd 100644
--- a/tensor2tensor/data_generators/problem_hparams.py
+++ b/tensor2tensor/data_generators/problem_hparams.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/problem_test.py b/tensor2tensor/data_generators/problem_test.py
index 4a907c639..de1c26bea 100644
--- a/tensor2tensor/data_generators/problem_test.py
+++ b/tensor2tensor/data_generators/problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/program_search.py b/tensor2tensor/data_generators/program_search.py
index 5fc1e96e2..694434f96 100644
--- a/tensor2tensor/data_generators/program_search.py
+++ b/tensor2tensor/data_generators/program_search.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/program_search_test.py b/tensor2tensor/data_generators/program_search_test.py
index 340b39391..85f87ac98 100644
--- a/tensor2tensor/data_generators/program_search_test.py
+++ b/tensor2tensor/data_generators/program_search_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/ptb.py b/tensor2tensor/data_generators/ptb.py
index 26186a7b6..2f436c577 100644
--- a/tensor2tensor/data_generators/ptb.py
+++ b/tensor2tensor/data_generators/ptb.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/qnli.py b/tensor2tensor/data_generators/qnli.py
index 1d788eb91..38be61ee5 100644
--- a/tensor2tensor/data_generators/qnli.py
+++ b/tensor2tensor/data_generators/qnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/quora_qpairs.py b/tensor2tensor/data_generators/quora_qpairs.py
index e7a3ff4be..47d8541e8 100644
--- a/tensor2tensor/data_generators/quora_qpairs.py
+++ b/tensor2tensor/data_generators/quora_qpairs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/rte.py b/tensor2tensor/data_generators/rte.py
index 03236f54c..0fc5e56af 100644
--- a/tensor2tensor/data_generators/rte.py
+++ b/tensor2tensor/data_generators/rte.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/scitail.py b/tensor2tensor/data_generators/scitail.py
index 0bd1865f0..f56a2f84c 100644
--- a/tensor2tensor/data_generators/scitail.py
+++ b/tensor2tensor/data_generators/scitail.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/seq2edits.py b/tensor2tensor/data_generators/seq2edits.py
index aecfa1897..40495454a 100644
--- a/tensor2tensor/data_generators/seq2edits.py
+++ b/tensor2tensor/data_generators/seq2edits.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/snli.py b/tensor2tensor/data_generators/snli.py
index b14215792..cdfceb8ad 100644
--- a/tensor2tensor/data_generators/snli.py
+++ b/tensor2tensor/data_generators/snli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/speech_recognition.py b/tensor2tensor/data_generators/speech_recognition.py
index ae99187b7..6b253f0d0 100644
--- a/tensor2tensor/data_generators/speech_recognition.py
+++ b/tensor2tensor/data_generators/speech_recognition.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/squad.py b/tensor2tensor/data_generators/squad.py
index 2ade5f91a..2e1141c3a 100644
--- a/tensor2tensor/data_generators/squad.py
+++ b/tensor2tensor/data_generators/squad.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/sst_binary.py b/tensor2tensor/data_generators/sst_binary.py
index 9624a4153..c45e2d046 100644
--- a/tensor2tensor/data_generators/sst_binary.py
+++ b/tensor2tensor/data_generators/sst_binary.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/stanford_nli.py b/tensor2tensor/data_generators/stanford_nli.py
index 0820ed5ac..f6be6b6e6 100644
--- a/tensor2tensor/data_generators/stanford_nli.py
+++ b/tensor2tensor/data_generators/stanford_nli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/style_transfer.py b/tensor2tensor/data_generators/style_transfer.py
index 7f5c7e61e..5067d8dfa 100644
--- a/tensor2tensor/data_generators/style_transfer.py
+++ b/tensor2tensor/data_generators/style_transfer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/style_transfer_test.py b/tensor2tensor/data_generators/style_transfer_test.py
index 97e5267ae..0397fd942 100644
--- a/tensor2tensor/data_generators/style_transfer_test.py
+++ b/tensor2tensor/data_generators/style_transfer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/subject_verb_agreement.py b/tensor2tensor/data_generators/subject_verb_agreement.py
index 8e5beb138..aa9bf012d 100644
--- a/tensor2tensor/data_generators/subject_verb_agreement.py
+++ b/tensor2tensor/data_generators/subject_verb_agreement.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_encoder.py b/tensor2tensor/data_generators/text_encoder.py
index 769dd9b77..636f5f1f2 100644
--- a/tensor2tensor/data_generators/text_encoder.py
+++ b/tensor2tensor/data_generators/text_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_encoder_build_subword.py b/tensor2tensor/data_generators/text_encoder_build_subword.py
index 03d437c36..2f5bca643 100644
--- a/tensor2tensor/data_generators/text_encoder_build_subword.py
+++ b/tensor2tensor/data_generators/text_encoder_build_subword.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_encoder_test.py b/tensor2tensor/data_generators/text_encoder_test.py
index 1d0b489be..e2ed4e985 100644
--- a/tensor2tensor/data_generators/text_encoder_test.py
+++ b/tensor2tensor/data_generators/text_encoder_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_problems.py b/tensor2tensor/data_generators/text_problems.py
index 361c69028..8e4693f22 100644
--- a/tensor2tensor/data_generators/text_problems.py
+++ b/tensor2tensor/data_generators/text_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/text_problems_test.py b/tensor2tensor/data_generators/text_problems_test.py
index c7986dce0..331d65e6f 100644
--- a/tensor2tensor/data_generators/text_problems_test.py
+++ b/tensor2tensor/data_generators/text_problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries.py b/tensor2tensor/data_generators/timeseries.py
index 99a8e3b22..78048bc8a 100644
--- a/tensor2tensor/data_generators/timeseries.py
+++ b/tensor2tensor/data_generators/timeseries.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_data_generator.py b/tensor2tensor/data_generators/timeseries_data_generator.py
index 126250504..d0bb165f9 100644
--- a/tensor2tensor/data_generators/timeseries_data_generator.py
+++ b/tensor2tensor/data_generators/timeseries_data_generator.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_data_generator_test.py b/tensor2tensor/data_generators/timeseries_data_generator_test.py
index c9bbf95e0..256038b10 100644
--- a/tensor2tensor/data_generators/timeseries_data_generator_test.py
+++ b/tensor2tensor/data_generators/timeseries_data_generator_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/timeseries_test.py b/tensor2tensor/data_generators/timeseries_test.py
index 0f9e6678d..441e3dac7 100644
--- a/tensor2tensor/data_generators/timeseries_test.py
+++ b/tensor2tensor/data_generators/timeseries_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/tokenizer.py b/tensor2tensor/data_generators/tokenizer.py
index 9a8c4c6df..9e5eb5108 100644
--- a/tensor2tensor/data_generators/tokenizer.py
+++ b/tensor2tensor/data_generators/tokenizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/tokenizer_test.py b/tensor2tensor/data_generators/tokenizer_test.py
index 20da16115..1f4f955e2 100644
--- a/tensor2tensor/data_generators/tokenizer_test.py
+++ b/tensor2tensor/data_generators/tokenizer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/transduction_problems.py b/tensor2tensor/data_generators/transduction_problems.py
index 02755088a..e0b729548 100644
--- a/tensor2tensor/data_generators/transduction_problems.py
+++ b/tensor2tensor/data_generators/transduction_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/transduction_problems_test.py b/tensor2tensor/data_generators/transduction_problems_test.py
index 391835554..20fa474f1 100644
--- a/tensor2tensor/data_generators/transduction_problems_test.py
+++ b/tensor2tensor/data_generators/transduction_problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 1ae57641b..81baa4315 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_encs.py b/tensor2tensor/data_generators/translate_encs.py
index f709ad811..bce8a243e 100644
--- a/tensor2tensor/data_generators/translate_encs.py
+++ b/tensor2tensor/data_generators/translate_encs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_encs_cubbitt.py b/tensor2tensor/data_generators/translate_encs_cubbitt.py
index acc1bf610..f2a813b5b 100644
--- a/tensor2tensor/data_generators/translate_encs_cubbitt.py
+++ b/tensor2tensor/data_generators/translate_encs_cubbitt.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_ende.py b/tensor2tensor/data_generators/translate_ende.py
index be5282cc1..98182aec8 100644
--- a/tensor2tensor/data_generators/translate_ende.py
+++ b/tensor2tensor/data_generators/translate_ende.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_ende_test.py b/tensor2tensor/data_generators/translate_ende_test.py
index 23644c322..37443b620 100644
--- a/tensor2tensor/data_generators/translate_ende_test.py
+++ b/tensor2tensor/data_generators/translate_ende_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enes.py b/tensor2tensor/data_generators/translate_enes.py
index f4a7f2199..4a971e43f 100644
--- a/tensor2tensor/data_generators/translate_enes.py
+++ b/tensor2tensor/data_generators/translate_enes.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enet.py b/tensor2tensor/data_generators/translate_enet.py
index f10677000..be91cb3c3 100644
--- a/tensor2tensor/data_generators/translate_enet.py
+++ b/tensor2tensor/data_generators/translate_enet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enfr.py b/tensor2tensor/data_generators/translate_enfr.py
index c3dfc39a5..ed05e3e1d 100644
--- a/tensor2tensor/data_generators/translate_enfr.py
+++ b/tensor2tensor/data_generators/translate_enfr.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enid.py b/tensor2tensor/data_generators/translate_enid.py
index b2b800c93..936827f5c 100644
--- a/tensor2tensor/data_generators/translate_enid.py
+++ b/tensor2tensor/data_generators/translate_enid.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enmk.py b/tensor2tensor/data_generators/translate_enmk.py
index f41cc4e99..60b23b179 100644
--- a/tensor2tensor/data_generators/translate_enmk.py
+++ b/tensor2tensor/data_generators/translate_enmk.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enro.py b/tensor2tensor/data_generators/translate_enro.py
index a81894848..1c17b77c7 100644
--- a/tensor2tensor/data_generators/translate_enro.py
+++ b/tensor2tensor/data_generators/translate_enro.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_entn.py b/tensor2tensor/data_generators/translate_entn.py
index b622fda2b..fa11081e3 100644
--- a/tensor2tensor/data_generators/translate_entn.py
+++ b/tensor2tensor/data_generators/translate_entn.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_envi.py b/tensor2tensor/data_generators/translate_envi.py
index c3c68c58a..479883e36 100644
--- a/tensor2tensor/data_generators/translate_envi.py
+++ b/tensor2tensor/data_generators/translate_envi.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_enzh.py b/tensor2tensor/data_generators/translate_enzh.py
index cd8d952e1..14e0d8a63 100644
--- a/tensor2tensor/data_generators/translate_enzh.py
+++ b/tensor2tensor/data_generators/translate_enzh.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/translate_test.py b/tensor2tensor/data_generators/translate_test.py
index 1b6925cc2..2148ab3ed 100644
--- a/tensor2tensor/data_generators/translate_test.py
+++ b/tensor2tensor/data_generators/translate_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index 9a097e368..e57eafd1f 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index d20b074d4..8bdfbdf0c 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/video_utils_test.py b/tensor2tensor/data_generators/video_utils_test.py
index 9cb9b91db..7091a65ae 100644
--- a/tensor2tensor/data_generators/video_utils_test.py
+++ b/tensor2tensor/data_generators/video_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
index cd932845a..8d1eb40dd 100644
--- a/tensor2tensor/data_generators/vqa.py
+++ b/tensor2tensor/data_generators/vqa.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/vqa_utils.py b/tensor2tensor/data_generators/vqa_utils.py
index d2811e889..38042b139 100644
--- a/tensor2tensor/data_generators/vqa_utils.py
+++ b/tensor2tensor/data_generators/vqa_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki.py b/tensor2tensor/data_generators/wiki.py
index e189a6955..892f189b5 100644
--- a/tensor2tensor/data_generators/wiki.py
+++ b/tensor2tensor/data_generators/wiki.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_lm.py b/tensor2tensor/data_generators/wiki_lm.py
index d29e27026..61713bab5 100644
--- a/tensor2tensor/data_generators/wiki_lm.py
+++ b/tensor2tensor/data_generators/wiki_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_multi_problems.py b/tensor2tensor/data_generators/wiki_multi_problems.py
index 501b8b92a..135648d34 100644
--- a/tensor2tensor/data_generators/wiki_multi_problems.py
+++ b/tensor2tensor/data_generators/wiki_multi_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_revision.py b/tensor2tensor/data_generators/wiki_revision.py
index d65a21a7f..ffc737a6a 100644
--- a/tensor2tensor/data_generators/wiki_revision.py
+++ b/tensor2tensor/data_generators/wiki_revision.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
index 3641e887a..4df263d6b 100644
--- a/tensor2tensor/data_generators/wiki_revision_utils.py
+++ b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/__init__.py b/tensor2tensor/data_generators/wikisum/__init__.py
index ffe5892e3..ff174dd63 100644
--- a/tensor2tensor/data_generators/wikisum/__init__.py
+++ b/tensor2tensor/data_generators/wikisum/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/generate_vocab.py b/tensor2tensor/data_generators/wikisum/generate_vocab.py
index 517de7fb1..5769a650d 100644
--- a/tensor2tensor/data_generators/wikisum/generate_vocab.py
+++ b/tensor2tensor/data_generators/wikisum/generate_vocab.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py b/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
index be7aa0e80..2f19ca09d 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_web.py b/tensor2tensor/data_generators/wikisum/get_references_web.py
index 1a2afdda0..3bd8a69a4 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py b/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
index 86399371a..18208d742 100644
--- a/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
+++ b/tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/html.py b/tensor2tensor/data_generators/wikisum/html.py
index 62c752df9..2cf300b1d 100644
--- a/tensor2tensor/data_generators/wikisum/html.py
+++ b/tensor2tensor/data_generators/wikisum/html.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/parallel_launch.py b/tensor2tensor/data_generators/wikisum/parallel_launch.py
index 13a50e355..0cdc8403a 100644
--- a/tensor2tensor/data_generators/wikisum/parallel_launch.py
+++ b/tensor2tensor/data_generators/wikisum/parallel_launch.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/produce_examples.py b/tensor2tensor/data_generators/wikisum/produce_examples.py
index 435151fb2..25cad27a5 100644
--- a/tensor2tensor/data_generators/wikisum/produce_examples.py
+++ b/tensor2tensor/data_generators/wikisum/produce_examples.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/utils.py b/tensor2tensor/data_generators/wikisum/utils.py
index e60b579d9..f45566e22 100644
--- a/tensor2tensor/data_generators/wikisum/utils.py
+++ b/tensor2tensor/data_generators/wikisum/utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/utils_test.py b/tensor2tensor/data_generators/wikisum/utils_test.py
index eb2c989c7..559889d39 100644
--- a/tensor2tensor/data_generators/wikisum/utils_test.py
+++ b/tensor2tensor/data_generators/wikisum/utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/validate_data.py b/tensor2tensor/data_generators/wikisum/validate_data.py
index 4a8494c22..b03bef990 100644
--- a/tensor2tensor/data_generators/wikisum/validate_data.py
+++ b/tensor2tensor/data_generators/wikisum/validate_data.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikisum/wikisum.py b/tensor2tensor/data_generators/wikisum/wikisum.py
index 9003d5d8c..f2e2c8c13 100644
--- a/tensor2tensor/data_generators/wikisum/wikisum.py
+++ b/tensor2tensor/data_generators/wikisum/wikisum.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wikitext103.py b/tensor2tensor/data_generators/wikitext103.py
index 3d9ea2a61..4d8ec3957 100644
--- a/tensor2tensor/data_generators/wikitext103.py
+++ b/tensor2tensor/data_generators/wikitext103.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wnli.py b/tensor2tensor/data_generators/wnli.py
index 7d2af8b7e..b56746ba0 100644
--- a/tensor2tensor/data_generators/wnli.py
+++ b/tensor2tensor/data_generators/wnli.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/wsj_parsing.py b/tensor2tensor/data_generators/wsj_parsing.py
index f4f1ee892..86107e0a8 100644
--- a/tensor2tensor/data_generators/wsj_parsing.py
+++ b/tensor2tensor/data_generators/wsj_parsing.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/yelp_full.py b/tensor2tensor/data_generators/yelp_full.py
index 02d1b826d..5606fd745 100644
--- a/tensor2tensor/data_generators/yelp_full.py
+++ b/tensor2tensor/data_generators/yelp_full.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/data_generators/yelp_polarity.py b/tensor2tensor/data_generators/yelp_polarity.py
index 60d0d7dcd..e108d6ce8 100644
--- a/tensor2tensor/data_generators/yelp_polarity.py
+++ b/tensor2tensor/data_generators/yelp_polarity.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/__init__.py b/tensor2tensor/envs/__init__.py
index 869af5e33..9833a8ba3 100644
--- a/tensor2tensor/envs/__init__.py
+++ b/tensor2tensor/envs/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 4efebddb6..4d25e9459 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/env_problem_utils.py b/tensor2tensor/envs/env_problem_utils.py
index 467422dc1..32d5499ca 100644
--- a/tensor2tensor/envs/env_problem_utils.py
+++ b/tensor2tensor/envs/env_problem_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/env_problem_utils_test.py b/tensor2tensor/envs/env_problem_utils_test.py
index c592e6bc8..333863d08 100644
--- a/tensor2tensor/envs/env_problem_utils_test.py
+++ b/tensor2tensor/envs/env_problem_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_env_problem.py b/tensor2tensor/envs/gym_env_problem.py
index f85e61781..f5856a485 100644
--- a/tensor2tensor/envs/gym_env_problem.py
+++ b/tensor2tensor/envs/gym_env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_env_problem_test.py b/tensor2tensor/envs/gym_env_problem_test.py
index 5e9c48ea3..5ac72b6e3 100644
--- a/tensor2tensor/envs/gym_env_problem_test.py
+++ b/tensor2tensor/envs/gym_env_problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_spaces_utils.py b/tensor2tensor/envs/gym_spaces_utils.py
index c4256196f..ceca77dab 100644
--- a/tensor2tensor/envs/gym_spaces_utils.py
+++ b/tensor2tensor/envs/gym_spaces_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/gym_spaces_utils_test.py b/tensor2tensor/envs/gym_spaces_utils_test.py
index 6234939b5..b68f62723 100644
--- a/tensor2tensor/envs/gym_spaces_utils_test.py
+++ b/tensor2tensor/envs/gym_spaces_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/mujoco_problems.py b/tensor2tensor/envs/mujoco_problems.py
index 51cb951a7..53f82887a 100644
--- a/tensor2tensor/envs/mujoco_problems.py
+++ b/tensor2tensor/envs/mujoco_problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/mujoco_problems_test.py b/tensor2tensor/envs/mujoco_problems_test.py
index bd4edce91..a9c53c608 100644
--- a/tensor2tensor/envs/mujoco_problems_test.py
+++ b/tensor2tensor/envs/mujoco_problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
index e08bd99f2..b32c48226 100644
--- a/tensor2tensor/envs/rendered_env_problem.py
+++ b/tensor2tensor/envs/rendered_env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/rendered_env_problem_test.py b/tensor2tensor/envs/rendered_env_problem_test.py
index 4a6273e0c..c7e816434 100644
--- a/tensor2tensor/envs/rendered_env_problem_test.py
+++ b/tensor2tensor/envs/rendered_env_problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env.py b/tensor2tensor/envs/tic_tac_toe_env.py
index 65f4d4d4e..795bcd1b2 100644
--- a/tensor2tensor/envs/tic_tac_toe_env.py
+++ b/tensor2tensor/envs/tic_tac_toe_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_problem.py b/tensor2tensor/envs/tic_tac_toe_env_problem.py
index c85301ca2..4f99516ff 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_problem_test.py b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
index 78acab86e..6aac41db2 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_problem_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/tic_tac_toe_env_test.py b/tensor2tensor/envs/tic_tac_toe_env_test.py
index 9a7f77d5b..f277d94a2 100644
--- a/tensor2tensor/envs/tic_tac_toe_env_test.py
+++ b/tensor2tensor/envs/tic_tac_toe_env_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/time_step.py b/tensor2tensor/envs/time_step.py
index 675d5e4d6..58649652a 100644
--- a/tensor2tensor/envs/time_step.py
+++ b/tensor2tensor/envs/time_step.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/time_step_test.py b/tensor2tensor/envs/time_step_test.py
index 98476bb21..2b67639ef 100644
--- a/tensor2tensor/envs/time_step_test.py
+++ b/tensor2tensor/envs/time_step_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index dcb9fcc92..91b580893 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/envs/trajectory_test.py b/tensor2tensor/envs/trajectory_test.py
index 23d94dea8..af0652d57 100644
--- a/tensor2tensor/envs/trajectory_test.py
+++ b/tensor2tensor/envs/trajectory_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/__init__.py b/tensor2tensor/insights/__init__.py
index ffe5892e3..ff174dd63 100644
--- a/tensor2tensor/insights/__init__.py
+++ b/tensor2tensor/insights/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/graph.py b/tensor2tensor/insights/graph.py
index afa4b8452..fc2eb577d 100644
--- a/tensor2tensor/insights/graph.py
+++ b/tensor2tensor/insights/graph.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/query_processor.py b/tensor2tensor/insights/query_processor.py
index a19213592..7500b8467 100644
--- a/tensor2tensor/insights/query_processor.py
+++ b/tensor2tensor/insights/query_processor.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/server.py b/tensor2tensor/insights/server.py
index 23c34c485..e666580c4 100644
--- a/tensor2tensor/insights/server.py
+++ b/tensor2tensor/insights/server.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/insights/transformer_model.py b/tensor2tensor/insights/transformer_model.py
index d44703962..f0b4ac097 100644
--- a/tensor2tensor/insights/transformer_model.py
+++ b/tensor2tensor/insights/transformer_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/__init__.py b/tensor2tensor/layers/__init__.py
index ffe5892e3..ff174dd63 100644
--- a/tensor2tensor/layers/__init__.py
+++ b/tensor2tensor/layers/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/area_attention.py b/tensor2tensor/layers/area_attention.py
index 46e6a695b..88ced00f2 100644
--- a/tensor2tensor/layers/area_attention.py
+++ b/tensor2tensor/layers/area_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/area_attention_test.py b/tensor2tensor/layers/area_attention_test.py
index 79191b07b..dfc13eb73 100644
--- a/tensor2tensor/layers/area_attention_test.py
+++ b/tensor2tensor/layers/area_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index 2c6929402..11ce57fcf 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_attention_test.py b/tensor2tensor/layers/common_attention_test.py
index dbee72c1b..108cd754c 100644
--- a/tensor2tensor/layers/common_attention_test.py
+++ b/tensor2tensor/layers/common_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_audio.py b/tensor2tensor/layers/common_audio.py
index f090fc5d4..27614d7da 100644
--- a/tensor2tensor/layers/common_audio.py
+++ b/tensor2tensor/layers/common_audio.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_hparams.py b/tensor2tensor/layers/common_hparams.py
index c3fd8cfbf..65e349b24 100644
--- a/tensor2tensor/layers/common_hparams.py
+++ b/tensor2tensor/layers/common_hparams.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_image_attention.py b/tensor2tensor/layers/common_image_attention.py
index c7d28162d..f19be8c71 100644
--- a/tensor2tensor/layers/common_image_attention.py
+++ b/tensor2tensor/layers/common_image_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_image_attention_test.py b/tensor2tensor/layers/common_image_attention_test.py
index 16b51c5da..4ece5779d 100644
--- a/tensor2tensor/layers/common_image_attention_test.py
+++ b/tensor2tensor/layers/common_image_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 7bedb1af2..79b393c86 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_layers_test.py b/tensor2tensor/layers/common_layers_test.py
index 530d9104d..897aae2db 100644
--- a/tensor2tensor/layers/common_layers_test.py
+++ b/tensor2tensor/layers/common_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 0ef067f83..103fb08b3 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/common_video_test.py b/tensor2tensor/layers/common_video_test.py
index ed241b3ec..d17b67ca1 100644
--- a/tensor2tensor/layers/common_video_test.py
+++ b/tensor2tensor/layers/common_video_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/discretization.py b/tensor2tensor/layers/discretization.py
index 99364b4ae..feb338dd4 100644
--- a/tensor2tensor/layers/discretization.py
+++ b/tensor2tensor/layers/discretization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/discretization_test.py b/tensor2tensor/layers/discretization_test.py
index 283081388..d0957a03f 100644
--- a/tensor2tensor/layers/discretization_test.py
+++ b/tensor2tensor/layers/discretization_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/latent_layers.py b/tensor2tensor/layers/latent_layers.py
index 025a8217d..bf1629d2f 100644
--- a/tensor2tensor/layers/latent_layers.py
+++ b/tensor2tensor/layers/latent_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/latent_layers_test.py b/tensor2tensor/layers/latent_layers_test.py
index 5abbcfd05..7a2c59275 100644
--- a/tensor2tensor/layers/latent_layers_test.py
+++ b/tensor2tensor/layers/latent_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/message_passing_attention.py b/tensor2tensor/layers/message_passing_attention.py
index 65c3912d5..e2db7246a 100644
--- a/tensor2tensor/layers/message_passing_attention.py
+++ b/tensor2tensor/layers/message_passing_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
index ab842c55b..0d9894997 100644
--- a/tensor2tensor/layers/modalities.py
+++ b/tensor2tensor/layers/modalities.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/modalities_test.py b/tensor2tensor/layers/modalities_test.py
index 0a5420890..dbaf3c68e 100644
--- a/tensor2tensor/layers/modalities_test.py
+++ b/tensor2tensor/layers/modalities_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/ngram.py b/tensor2tensor/layers/ngram.py
index 26c08e241..05dcd54d8 100644
--- a/tensor2tensor/layers/ngram.py
+++ b/tensor2tensor/layers/ngram.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/ngram_test.py b/tensor2tensor/layers/ngram_test.py
index a7b8e8787..3dce37268 100644
--- a/tensor2tensor/layers/ngram_test.py
+++ b/tensor2tensor/layers/ngram_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_glow_layers.py b/tensor2tensor/layers/transformer_glow_layers.py
index eb46c6f6a..12c963fcb 100644
--- a/tensor2tensor/layers/transformer_glow_layers.py
+++ b/tensor2tensor/layers/transformer_glow_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_glow_layers_ops.py b/tensor2tensor/layers/transformer_glow_layers_ops.py
index 3b7c0ee15..58bae3831 100644
--- a/tensor2tensor/layers/transformer_glow_layers_ops.py
+++ b/tensor2tensor/layers/transformer_glow_layers_ops.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_glow_layers_ops_test.py b/tensor2tensor/layers/transformer_glow_layers_ops_test.py
index 7c5404c7e..9bed12aca 100644
--- a/tensor2tensor/layers/transformer_glow_layers_ops_test.py
+++ b/tensor2tensor/layers/transformer_glow_layers_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_glow_layers_test.py b/tensor2tensor/layers/transformer_glow_layers_test.py
index 7658aaaf4..c4f97050a 100644
--- a/tensor2tensor/layers/transformer_glow_layers_test.py
+++ b/tensor2tensor/layers/transformer_glow_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_layers.py b/tensor2tensor/layers/transformer_layers.py
index 662ba70a7..327f8e591 100644
--- a/tensor2tensor/layers/transformer_layers.py
+++ b/tensor2tensor/layers/transformer_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_memory.py b/tensor2tensor/layers/transformer_memory.py
index cadee628b..b35707128 100644
--- a/tensor2tensor/layers/transformer_memory.py
+++ b/tensor2tensor/layers/transformer_memory.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/transformer_memory_test.py b/tensor2tensor/layers/transformer_memory_test.py
index a03769a12..10155eb9c 100644
--- a/tensor2tensor/layers/transformer_memory_test.py
+++ b/tensor2tensor/layers/transformer_memory_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/vq_discrete.py b/tensor2tensor/layers/vq_discrete.py
index 182127b03..31c9b1215 100644
--- a/tensor2tensor/layers/vq_discrete.py
+++ b/tensor2tensor/layers/vq_discrete.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/layers/vqa_layers.py b/tensor2tensor/layers/vqa_layers.py
index eed7b24aa..7949eb662 100644
--- a/tensor2tensor/layers/vqa_layers.py
+++ b/tensor2tensor/layers/vqa_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/__init__.py b/tensor2tensor/metrics/__init__.py
index ffe5892e3..ff174dd63 100644
--- a/tensor2tensor/metrics/__init__.py
+++ b/tensor2tensor/metrics/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/video_conditional_fvd.py b/tensor2tensor/metrics/video_conditional_fvd.py
index 6c29fd4a8..f6a089651 100644
--- a/tensor2tensor/metrics/video_conditional_fvd.py
+++ b/tensor2tensor/metrics/video_conditional_fvd.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/metrics/video_conditional_fvd_test.py b/tensor2tensor/metrics/video_conditional_fvd_test.py
index bf1743c45..3f2723ec1 100644
--- a/tensor2tensor/metrics/video_conditional_fvd_test.py
+++ b/tensor2tensor/metrics/video_conditional_fvd_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/__init__.py b/tensor2tensor/models/__init__.py
index b06ecb04b..62d059134 100644
--- a/tensor2tensor/models/__init__.py
+++ b/tensor2tensor/models/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/basic.py b/tensor2tensor/models/basic.py
index 424b3abe9..4a3209022 100644
--- a/tensor2tensor/models/basic.py
+++ b/tensor2tensor/models/basic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/basic_test.py b/tensor2tensor/models/basic_test.py
index df77a707c..3f6b4affd 100644
--- a/tensor2tensor/models/basic_test.py
+++ b/tensor2tensor/models/basic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/bytenet.py b/tensor2tensor/models/bytenet.py
index aac98cdd0..84594f36a 100644
--- a/tensor2tensor/models/bytenet.py
+++ b/tensor2tensor/models/bytenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/bytenet_test.py b/tensor2tensor/models/bytenet_test.py
index 3dd087289..204d54bc1 100644
--- a/tensor2tensor/models/bytenet_test.py
+++ b/tensor2tensor/models/bytenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/distillation.py b/tensor2tensor/models/distillation.py
index 40c4adcf8..9d8ccb849 100644
--- a/tensor2tensor/models/distillation.py
+++ b/tensor2tensor/models/distillation.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/evolved_transformer.py b/tensor2tensor/models/evolved_transformer.py
index 1240bd2de..bac01a3cf 100644
--- a/tensor2tensor/models/evolved_transformer.py
+++ b/tensor2tensor/models/evolved_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/evolved_transformer_test.py b/tensor2tensor/models/evolved_transformer_test.py
index 6888e5c24..388769918 100644
--- a/tensor2tensor/models/evolved_transformer_test.py
+++ b/tensor2tensor/models/evolved_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer.py b/tensor2tensor/models/image_transformer.py
index bb8b5b87e..dd7c2d882 100644
--- a/tensor2tensor/models/image_transformer.py
+++ b/tensor2tensor/models/image_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_2d.py b/tensor2tensor/models/image_transformer_2d.py
index d0b1c1981..32c4aa59a 100644
--- a/tensor2tensor/models/image_transformer_2d.py
+++ b/tensor2tensor/models/image_transformer_2d.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_2d_test.py b/tensor2tensor/models/image_transformer_2d_test.py
index 13f6e799a..de3e73837 100644
--- a/tensor2tensor/models/image_transformer_2d_test.py
+++ b/tensor2tensor/models/image_transformer_2d_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/image_transformer_test.py b/tensor2tensor/models/image_transformer_test.py
index 142db24b3..6dde81d5e 100644
--- a/tensor2tensor/models/image_transformer_test.py
+++ b/tensor2tensor/models/image_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/lstm.py b/tensor2tensor/models/lstm.py
index 4f7447357..f59dabb19 100644
--- a/tensor2tensor/models/lstm.py
+++ b/tensor2tensor/models/lstm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/lstm_test.py b/tensor2tensor/models/lstm_test.py
index 4af344ee5..4723998db 100644
--- a/tensor2tensor/models/lstm_test.py
+++ b/tensor2tensor/models/lstm_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_image_transformer.py b/tensor2tensor/models/mtf_image_transformer.py
index 6db8e0d38..dffe8c66b 100644
--- a/tensor2tensor/models/mtf_image_transformer.py
+++ b/tensor2tensor/models/mtf_image_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_image_transformer_test.py b/tensor2tensor/models/mtf_image_transformer_test.py
index fd028f331..4737d16ea 100644
--- a/tensor2tensor/models/mtf_image_transformer_test.py
+++ b/tensor2tensor/models/mtf_image_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_resnet.py b/tensor2tensor/models/mtf_resnet.py
index 84a95254a..4ad14ee63 100644
--- a/tensor2tensor/models/mtf_resnet.py
+++ b/tensor2tensor/models/mtf_resnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer.py b/tensor2tensor/models/mtf_transformer.py
index 18015ea00..42bb88705 100644
--- a/tensor2tensor/models/mtf_transformer.py
+++ b/tensor2tensor/models/mtf_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer2.py b/tensor2tensor/models/mtf_transformer2.py
index baf6d2c1a..ed3ffa88d 100644
--- a/tensor2tensor/models/mtf_transformer2.py
+++ b/tensor2tensor/models/mtf_transformer2.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/mtf_transformer_test.py b/tensor2tensor/models/mtf_transformer_test.py
index cb32f805c..f411e078b 100644
--- a/tensor2tensor/models/mtf_transformer_test.py
+++ b/tensor2tensor/models/mtf_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/__init__.py b/tensor2tensor/models/neural_architecture_search/__init__.py
index aa007c728..06080ebe9 100644
--- a/tensor2tensor/models/neural_architecture_search/__init__.py
+++ b/tensor2tensor/models/neural_architecture_search/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_layers.py b/tensor2tensor/models/neural_architecture_search/nas_layers.py
index d198a3d62..c89fc1f78 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_layers.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_layers_test.py b/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
index a9db2679a..11b13324a 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_layers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_model.py b/tensor2tensor/models/neural_architecture_search/nas_model.py
index 741504df1..7e405c870 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_architecture_search/nas_model_test.py b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
index f3b05f6ce..8d2a9b446 100644
--- a/tensor2tensor/models/neural_architecture_search/nas_model_test.py
+++ b/tensor2tensor/models/neural_architecture_search/nas_model_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_assistant.py b/tensor2tensor/models/neural_assistant.py
index 17f6ae954..53f87eb1d 100644
--- a/tensor2tensor/models/neural_assistant.py
+++ b/tensor2tensor/models/neural_assistant.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_gpu.py b/tensor2tensor/models/neural_gpu.py
index 964b00bf6..953855172 100644
--- a/tensor2tensor/models/neural_gpu.py
+++ b/tensor2tensor/models/neural_gpu.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/neural_gpu_test.py b/tensor2tensor/models/neural_gpu_test.py
index 07a7c5b83..57a4a1f36 100644
--- a/tensor2tensor/models/neural_gpu_test.py
+++ b/tensor2tensor/models/neural_gpu_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/__init__.py b/tensor2tensor/models/research/__init__.py
index ffe5892e3..ff174dd63 100644
--- a/tensor2tensor/models/research/__init__.py
+++ b/tensor2tensor/models/research/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/adafactor_experiments.py b/tensor2tensor/models/research/adafactor_experiments.py
index 69e63824c..60daee121 100644
--- a/tensor2tensor/models/research/adafactor_experiments.py
+++ b/tensor2tensor/models/research/adafactor_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/aligned.py b/tensor2tensor/models/research/aligned.py
index d5071f64b..41dda38ac 100644
--- a/tensor2tensor/models/research/aligned.py
+++ b/tensor2tensor/models/research/aligned.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/attention_lm.py b/tensor2tensor/models/research/attention_lm.py
index 1d71b2d97..e6a456ef6 100644
--- a/tensor2tensor/models/research/attention_lm.py
+++ b/tensor2tensor/models/research/attention_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/attention_lm_moe.py b/tensor2tensor/models/research/attention_lm_moe.py
index 2af44b896..c385dacf4 100644
--- a/tensor2tensor/models/research/attention_lm_moe.py
+++ b/tensor2tensor/models/research/attention_lm_moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/autoencoders.py b/tensor2tensor/models/research/autoencoders.py
index 3298bc123..8a3da53dd 100644
--- a/tensor2tensor/models/research/autoencoders.py
+++ b/tensor2tensor/models/research/autoencoders.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/autoencoders_test.py b/tensor2tensor/models/research/autoencoders_test.py
index 08f25a118..f2c1afbdf 100644
--- a/tensor2tensor/models/research/autoencoders_test.py
+++ b/tensor2tensor/models/research/autoencoders_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/cycle_gan.py b/tensor2tensor/models/research/cycle_gan.py
index 3174e84c2..7146a4ee8 100644
--- a/tensor2tensor/models/research/cycle_gan.py
+++ b/tensor2tensor/models/research/cycle_gan.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/gene_expression.py b/tensor2tensor/models/research/gene_expression.py
index e67eac9b3..999f26edc 100644
--- a/tensor2tensor/models/research/gene_expression.py
+++ b/tensor2tensor/models/research/gene_expression.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/gene_expression_test.py b/tensor2tensor/models/research/gene_expression_test.py
index dfd0469fb..a20b82e26 100644
--- a/tensor2tensor/models/research/gene_expression_test.py
+++ b/tensor2tensor/models/research/gene_expression_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow.py b/tensor2tensor/models/research/glow.py
index b76aa973a..8ebb189d8 100644
--- a/tensor2tensor/models/research/glow.py
+++ b/tensor2tensor/models/research/glow.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_init_hook.py b/tensor2tensor/models/research/glow_init_hook.py
index 5b1f18789..34dd4fe5b 100644
--- a/tensor2tensor/models/research/glow_init_hook.py
+++ b/tensor2tensor/models/research/glow_init_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_ops.py b/tensor2tensor/models/research/glow_ops.py
index 4c904c3cd..c5ee371cb 100644
--- a/tensor2tensor/models/research/glow_ops.py
+++ b/tensor2tensor/models/research/glow_ops.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_ops_test.py b/tensor2tensor/models/research/glow_ops_test.py
index 5caf324eb..baab0fd1a 100644
--- a/tensor2tensor/models/research/glow_ops_test.py
+++ b/tensor2tensor/models/research/glow_ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/glow_test.py b/tensor2tensor/models/research/glow_test.py
index ffdca2a05..ef14224b5 100644
--- a/tensor2tensor/models/research/glow_test.py
+++ b/tensor2tensor/models/research/glow_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/lm_experiments.py b/tensor2tensor/models/research/lm_experiments.py
index 9e50090f9..91c074ac4 100644
--- a/tensor2tensor/models/research/lm_experiments.py
+++ b/tensor2tensor/models/research/lm_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/moe.py b/tensor2tensor/models/research/moe.py
index e6195b22f..89e26a174 100644
--- a/tensor2tensor/models/research/moe.py
+++ b/tensor2tensor/models/research/moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/moe_experiments.py b/tensor2tensor/models/research/moe_experiments.py
index 3bef4be56..73f3f4f50 100644
--- a/tensor2tensor/models/research/moe_experiments.py
+++ b/tensor2tensor/models/research/moe_experiments.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/multiquery_paper.py b/tensor2tensor/models/research/multiquery_paper.py
index af79c2454..9157a5177 100644
--- a/tensor2tensor/models/research/multiquery_paper.py
+++ b/tensor2tensor/models/research/multiquery_paper.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/neural_stack.py b/tensor2tensor/models/research/neural_stack.py
index 8ed4b14b0..0c5464bd9 100644
--- a/tensor2tensor/models/research/neural_stack.py
+++ b/tensor2tensor/models/research/neural_stack.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/neural_stack_test.py b/tensor2tensor/models/research/neural_stack_test.py
index 4571af15d..83f45ab68 100644
--- a/tensor2tensor/models/research/neural_stack_test.py
+++ b/tensor2tensor/models/research/neural_stack_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/residual_shuffle_exchange.py b/tensor2tensor/models/research/residual_shuffle_exchange.py
index 745537bd5..7c22c2880 100644
--- a/tensor2tensor/models/research/residual_shuffle_exchange.py
+++ b/tensor2tensor/models/research/residual_shuffle_exchange.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/rl.py b/tensor2tensor/models/research/rl.py
index ab1a5f2f0..16ee162c6 100644
--- a/tensor2tensor/models/research/rl.py
+++ b/tensor2tensor/models/research/rl.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/shuffle_network.py b/tensor2tensor/models/research/shuffle_network.py
index dc284f87f..3b0117010 100644
--- a/tensor2tensor/models/research/shuffle_network.py
+++ b/tensor2tensor/models/research/shuffle_network.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/similarity_transformer.py b/tensor2tensor/models/research/similarity_transformer.py
index 0e2216b93..8d596a808 100644
--- a/tensor2tensor/models/research/similarity_transformer.py
+++ b/tensor2tensor/models/research/similarity_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/super_lm.py b/tensor2tensor/models/research/super_lm.py
index aa6c39375..ea45c67b7 100644
--- a/tensor2tensor/models/research/super_lm.py
+++ b/tensor2tensor/models/research/super_lm.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_aux.py b/tensor2tensor/models/research/transformer_aux.py
index d2e60d3a7..6bb0f48e7 100644
--- a/tensor2tensor/models/research/transformer_aux.py
+++ b/tensor2tensor/models/research/transformer_aux.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_aux_test.py b/tensor2tensor/models/research/transformer_aux_test.py
index 0e7bd6568..0d4ae332a 100644
--- a/tensor2tensor/models/research/transformer_aux_test.py
+++ b/tensor2tensor/models/research/transformer_aux_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_moe.py b/tensor2tensor/models/research/transformer_moe.py
index 8a2e1b327..9e3eadcea 100644
--- a/tensor2tensor/models/research/transformer_moe.py
+++ b/tensor2tensor/models/research/transformer_moe.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_nat.py b/tensor2tensor/models/research/transformer_nat.py
index 29b56c8f6..7ad67c8b1 100644
--- a/tensor2tensor/models/research/transformer_nat.py
+++ b/tensor2tensor/models/research/transformer_nat.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_parallel.py b/tensor2tensor/models/research/transformer_parallel.py
index 5909eee4c..8e73e673c 100644
--- a/tensor2tensor/models/research/transformer_parallel.py
+++ b/tensor2tensor/models/research/transformer_parallel.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_revnet.py b/tensor2tensor/models/research/transformer_revnet.py
index 21af70e10..7b9ee347d 100644
--- a/tensor2tensor/models/research/transformer_revnet.py
+++ b/tensor2tensor/models/research/transformer_revnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_revnet_test.py b/tensor2tensor/models/research/transformer_revnet_test.py
index 3c66ab616..fca42b54e 100644
--- a/tensor2tensor/models/research/transformer_revnet_test.py
+++ b/tensor2tensor/models/research/transformer_revnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_seq2edits.py b/tensor2tensor/models/research/transformer_seq2edits.py
index bb18ba845..75ce9b02a 100644
--- a/tensor2tensor/models/research/transformer_seq2edits.py
+++ b/tensor2tensor/models/research/transformer_seq2edits.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_sketch.py b/tensor2tensor/models/research/transformer_sketch.py
index 4f27d4fc1..79aa21aec 100644
--- a/tensor2tensor/models/research/transformer_sketch.py
+++ b/tensor2tensor/models/research/transformer_sketch.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_symshard.py b/tensor2tensor/models/research/transformer_symshard.py
index 030bcfb02..aed4b5e96 100644
--- a/tensor2tensor/models/research/transformer_symshard.py
+++ b/tensor2tensor/models/research/transformer_symshard.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae.py b/tensor2tensor/models/research/transformer_vae.py
index 021b361c3..84644446d 100644
--- a/tensor2tensor/models/research/transformer_vae.py
+++ b/tensor2tensor/models/research/transformer_vae.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae_flow_prior.py b/tensor2tensor/models/research/transformer_vae_flow_prior.py
index 64c1d88e3..343a457bc 100644
--- a/tensor2tensor/models/research/transformer_vae_flow_prior.py
+++ b/tensor2tensor/models/research/transformer_vae_flow_prior.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py b/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
index 735c8bfc3..ad176f174 100644
--- a/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
+++ b/tensor2tensor/models/research/transformer_vae_flow_prior_ops.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/transformer_vae_test.py b/tensor2tensor/models/research/transformer_vae_test.py
index a9f38ee43..8ac40854d 100644
--- a/tensor2tensor/models/research/transformer_vae_test.py
+++ b/tensor2tensor/models/research/transformer_vae_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer.py b/tensor2tensor/models/research/universal_transformer.py
index 2fd339659..f480bf2da 100644
--- a/tensor2tensor/models/research/universal_transformer.py
+++ b/tensor2tensor/models/research/universal_transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer_test.py b/tensor2tensor/models/research/universal_transformer_test.py
index 28f348c06..143d756fe 100644
--- a/tensor2tensor/models/research/universal_transformer_test.py
+++ b/tensor2tensor/models/research/universal_transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/universal_transformer_util.py b/tensor2tensor/models/research/universal_transformer_util.py
index 353acf811..daa5ab578 100644
--- a/tensor2tensor/models/research/universal_transformer_util.py
+++ b/tensor2tensor/models/research/universal_transformer_util.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_attention.py b/tensor2tensor/models/research/vqa_attention.py
index 5c07af376..7a006e1e1 100644
--- a/tensor2tensor/models/research/vqa_attention.py
+++ b/tensor2tensor/models/research/vqa_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_attention_test.py b/tensor2tensor/models/research/vqa_attention_test.py
index 0512a766c..87646fdd3 100644
--- a/tensor2tensor/models/research/vqa_attention_test.py
+++ b/tensor2tensor/models/research/vqa_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_recurrent_self_attention.py b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
index c85495c28..175e76d15 100644
--- a/tensor2tensor/models/research/vqa_recurrent_self_attention.py
+++ b/tensor2tensor/models/research/vqa_recurrent_self_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/research/vqa_self_attention.py b/tensor2tensor/models/research/vqa_self_attention.py
index eb56afe1d..eb1b86948 100644
--- a/tensor2tensor/models/research/vqa_self_attention.py
+++ b/tensor2tensor/models/research/vqa_self_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/resnet.py b/tensor2tensor/models/resnet.py
index e36e79b15..5eeb4792f 100644
--- a/tensor2tensor/models/resnet.py
+++ b/tensor2tensor/models/resnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/resnet_test.py b/tensor2tensor/models/resnet_test.py
index d982153ad..3b629fa48 100644
--- a/tensor2tensor/models/resnet_test.py
+++ b/tensor2tensor/models/resnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/revnet.py b/tensor2tensor/models/revnet.py
index 05b21b0f7..e841652af 100644
--- a/tensor2tensor/models/revnet.py
+++ b/tensor2tensor/models/revnet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/revnet_test.py b/tensor2tensor/models/revnet_test.py
index 1aef79886..234752514 100644
--- a/tensor2tensor/models/revnet_test.py
+++ b/tensor2tensor/models/revnet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/shake_shake.py b/tensor2tensor/models/shake_shake.py
index f63bf3329..378f86c97 100644
--- a/tensor2tensor/models/shake_shake.py
+++ b/tensor2tensor/models/shake_shake.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/slicenet.py b/tensor2tensor/models/slicenet.py
index c3e0af60a..e20786f31 100644
--- a/tensor2tensor/models/slicenet.py
+++ b/tensor2tensor/models/slicenet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/slicenet_test.py b/tensor2tensor/models/slicenet_test.py
index d4e43f7a4..944a78234 100644
--- a/tensor2tensor/models/slicenet_test.py
+++ b/tensor2tensor/models/slicenet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/text_cnn.py b/tensor2tensor/models/text_cnn.py
index bfc791e30..ee6434d3e 100644
--- a/tensor2tensor/models/text_cnn.py
+++ b/tensor2tensor/models/text_cnn.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/transformer.py b/tensor2tensor/models/transformer.py
index e5294e89f..2bc8f33d1 100644
--- a/tensor2tensor/models/transformer.py
+++ b/tensor2tensor/models/transformer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/transformer_test.py b/tensor2tensor/models/transformer_test.py
index 397f75b56..96cdae359 100644
--- a/tensor2tensor/models/transformer_test.py
+++ b/tensor2tensor/models/transformer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/vanilla_gan.py b/tensor2tensor/models/vanilla_gan.py
index 371a5031e..a79a7575f 100644
--- a/tensor2tensor/models/vanilla_gan.py
+++ b/tensor2tensor/models/vanilla_gan.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/__init__.py b/tensor2tensor/models/video/__init__.py
index ffe5892e3..ff174dd63 100644
--- a/tensor2tensor/models/video/__init__.py
+++ b/tensor2tensor/models/video/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/base.py b/tensor2tensor/models/video/base.py
index b860518e3..9fffc50dc 100644
--- a/tensor2tensor/models/video/base.py
+++ b/tensor2tensor/models/video/base.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/base_vae.py b/tensor2tensor/models/video/base_vae.py
index 56d7b1325..9f6e11fa2 100644
--- a/tensor2tensor/models/video/base_vae.py
+++ b/tensor2tensor/models/video/base_vae.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic.py b/tensor2tensor/models/video/basic_deterministic.py
index 3832c2cbd..624acac29 100644
--- a/tensor2tensor/models/video/basic_deterministic.py
+++ b/tensor2tensor/models/video/basic_deterministic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic_params.py b/tensor2tensor/models/video/basic_deterministic_params.py
index 10431a978..bb86c866b 100644
--- a/tensor2tensor/models/video/basic_deterministic_params.py
+++ b/tensor2tensor/models/video/basic_deterministic_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_deterministic_test.py b/tensor2tensor/models/video/basic_deterministic_test.py
index c66aebb4c..b89c54a49 100644
--- a/tensor2tensor/models/video/basic_deterministic_test.py
+++ b/tensor2tensor/models/video/basic_deterministic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_recurrent.py b/tensor2tensor/models/video/basic_recurrent.py
index bd6436e30..e82cc6f15 100644
--- a/tensor2tensor/models/video/basic_recurrent.py
+++ b/tensor2tensor/models/video/basic_recurrent.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_recurrent_test.py b/tensor2tensor/models/video/basic_recurrent_test.py
index 7618d19f4..0ec66a753 100644
--- a/tensor2tensor/models/video/basic_recurrent_test.py
+++ b/tensor2tensor/models/video/basic_recurrent_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_stochastic.py b/tensor2tensor/models/video/basic_stochastic.py
index 4c4426b5e..68bc73589 100644
--- a/tensor2tensor/models/video/basic_stochastic.py
+++ b/tensor2tensor/models/video/basic_stochastic.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/basic_stochastic_test.py b/tensor2tensor/models/video/basic_stochastic_test.py
index 03c68358e..c9c7b865e 100644
--- a/tensor2tensor/models/video/basic_stochastic_test.py
+++ b/tensor2tensor/models/video/basic_stochastic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/emily.py b/tensor2tensor/models/video/emily.py
index 30636709b..1a8b1e4ef 100644
--- a/tensor2tensor/models/video/emily.py
+++ b/tensor2tensor/models/video/emily.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/emily_test.py b/tensor2tensor/models/video/emily_test.py
index 9b72810a2..4d95500bf 100644
--- a/tensor2tensor/models/video/emily_test.py
+++ b/tensor2tensor/models/video/emily_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/epva.py b/tensor2tensor/models/video/epva.py
index 79c640e27..2a3b895bc 100644
--- a/tensor2tensor/models/video/epva.py
+++ b/tensor2tensor/models/video/epva.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/epva_params.py b/tensor2tensor/models/video/epva_params.py
index 7d8bec8e7..23c70fdd9 100644
--- a/tensor2tensor/models/video/epva_params.py
+++ b/tensor2tensor/models/video/epva_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/next_frame_glow.py b/tensor2tensor/models/video/next_frame_glow.py
index d3719798a..7186d180c 100644
--- a/tensor2tensor/models/video/next_frame_glow.py
+++ b/tensor2tensor/models/video/next_frame_glow.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv3d_test.py b/tensor2tensor/models/video/nfg_conv3d_test.py
index 3b7d68b64..e3434cd04 100644
--- a/tensor2tensor/models/video/nfg_conv3d_test.py
+++ b/tensor2tensor/models/video/nfg_conv3d_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv_lstm_test.py b/tensor2tensor/models/video/nfg_conv_lstm_test.py
index ae300b0f6..982245464 100644
--- a/tensor2tensor/models/video/nfg_conv_lstm_test.py
+++ b/tensor2tensor/models/video/nfg_conv_lstm_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_conv_test.py b/tensor2tensor/models/video/nfg_conv_test.py
index 1c1b1ca5b..5a07b812c 100644
--- a/tensor2tensor/models/video/nfg_conv_test.py
+++ b/tensor2tensor/models/video/nfg_conv_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_interpolate.py b/tensor2tensor/models/video/nfg_interpolate.py
index c9424d0b9..9294785b5 100644
--- a/tensor2tensor/models/video/nfg_interpolate.py
+++ b/tensor2tensor/models/video/nfg_interpolate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_test_utils.py b/tensor2tensor/models/video/nfg_test_utils.py
index fd679dd5d..e81269e03 100644
--- a/tensor2tensor/models/video/nfg_test_utils.py
+++ b/tensor2tensor/models/video/nfg_test_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/nfg_uncond_test.py b/tensor2tensor/models/video/nfg_uncond_test.py
index 44763d7ff..041062ab7 100644
--- a/tensor2tensor/models/video/nfg_uncond_test.py
+++ b/tensor2tensor/models/video/nfg_uncond_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp.py b/tensor2tensor/models/video/savp.py
index 2036af698..bf70623df 100644
--- a/tensor2tensor/models/video/savp.py
+++ b/tensor2tensor/models/video/savp.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp_params.py b/tensor2tensor/models/video/savp_params.py
index c1a34a5fa..b5705f43e 100644
--- a/tensor2tensor/models/video/savp_params.py
+++ b/tensor2tensor/models/video/savp_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/savp_test.py b/tensor2tensor/models/video/savp_test.py
index b353f1f70..94bf0c056 100644
--- a/tensor2tensor/models/video/savp_test.py
+++ b/tensor2tensor/models/video/savp_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p.py b/tensor2tensor/models/video/sv2p.py
index debcee4ac..09cc3f586 100644
--- a/tensor2tensor/models/video/sv2p.py
+++ b/tensor2tensor/models/video/sv2p.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p_params.py b/tensor2tensor/models/video/sv2p_params.py
index 0f21c9248..6a151dcab 100644
--- a/tensor2tensor/models/video/sv2p_params.py
+++ b/tensor2tensor/models/video/sv2p_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/sv2p_test.py b/tensor2tensor/models/video/sv2p_test.py
index dacbcb87b..c0e40e96c 100644
--- a/tensor2tensor/models/video/sv2p_test.py
+++ b/tensor2tensor/models/video/sv2p_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/video/tests_utils.py b/tensor2tensor/models/video/tests_utils.py
index 170ca484b..2a38b1cf3 100644
--- a/tensor2tensor/models/video/tests_utils.py
+++ b/tensor2tensor/models/video/tests_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/xception.py b/tensor2tensor/models/xception.py
index a5fb9c1cf..83b6697c6 100644
--- a/tensor2tensor/models/xception.py
+++ b/tensor2tensor/models/xception.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/models/xception_test.py b/tensor2tensor/models/xception_test.py
index 3f7df88b5..36ca2d1be 100644
--- a/tensor2tensor/models/xception_test.py
+++ b/tensor2tensor/models/xception_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/problems.py b/tensor2tensor/problems.py
index f8432957d..a16a1a717 100644
--- a/tensor2tensor/problems.py
+++ b/tensor2tensor/problems.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/problems_colab.py b/tensor2tensor/problems_colab.py
index 67563b387..20a7fa8b4 100644
--- a/tensor2tensor/problems_colab.py
+++ b/tensor2tensor/problems_colab.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/problems_test.py b/tensor2tensor/problems_test.py
index b4cc7ccb4..5753cb8fe 100644
--- a/tensor2tensor/problems_test.py
+++ b/tensor2tensor/problems_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/__init__.py b/tensor2tensor/rl/__init__.py
index ffe5892e3..ff174dd63 100644
--- a/tensor2tensor/rl/__init__.py
+++ b/tensor2tensor/rl/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/batch_dqn_agent_test.py b/tensor2tensor/rl/batch_dqn_agent_test.py
index c751e02a9..fc826f0d0 100644
--- a/tensor2tensor/rl/batch_dqn_agent_test.py
+++ b/tensor2tensor/rl/batch_dqn_agent_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/batch_runner_test.py b/tensor2tensor/rl/batch_runner_test.py
index 61c04bac5..648e7ff6e 100644
--- a/tensor2tensor/rl/batch_runner_test.py
+++ b/tensor2tensor/rl/batch_runner_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/datagen_with_agent.py b/tensor2tensor/rl/datagen_with_agent.py
index a001ffc9f..66a68e780 100644
--- a/tensor2tensor/rl/datagen_with_agent.py
+++ b/tensor2tensor/rl/datagen_with_agent.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/dopamine_connector.py b/tensor2tensor/rl/dopamine_connector.py
index d8a4c2728..e51a0c6c9 100644
--- a/tensor2tensor/rl/dopamine_connector.py
+++ b/tensor2tensor/rl/dopamine_connector.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/__init__.py b/tensor2tensor/rl/envs/__init__.py
index ffe5892e3..ff174dd63 100644
--- a/tensor2tensor/rl/envs/__init__.py
+++ b/tensor2tensor/rl/envs/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/in_graph_batch_env.py b/tensor2tensor/rl/envs/in_graph_batch_env.py
index c16e209e6..91f44afa4 100644
--- a/tensor2tensor/rl/envs/in_graph_batch_env.py
+++ b/tensor2tensor/rl/envs/in_graph_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/py_func_batch_env.py b/tensor2tensor/rl/envs/py_func_batch_env.py
index 984edf08e..e009eeb3a 100644
--- a/tensor2tensor/rl/envs/py_func_batch_env.py
+++ b/tensor2tensor/rl/envs/py_func_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/simulated_batch_env.py b/tensor2tensor/rl/envs/simulated_batch_env.py
index 881432581..b0048f198 100644
--- a/tensor2tensor/rl/envs/simulated_batch_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/simulated_batch_gym_env.py b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
index c045ee7f6..74b568339 100644
--- a/tensor2tensor/rl/envs/simulated_batch_gym_env.py
+++ b/tensor2tensor/rl/envs/simulated_batch_gym_env.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/envs/tf_atari_wrappers.py b/tensor2tensor/rl/envs/tf_atari_wrappers.py
index a3c0c1918..36838a62a 100644
--- a/tensor2tensor/rl/envs/tf_atari_wrappers.py
+++ b/tensor2tensor/rl/envs/tf_atari_wrappers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/evaluator.py b/tensor2tensor/rl/evaluator.py
index 87335f83b..f2f3b980a 100644
--- a/tensor2tensor/rl/evaluator.py
+++ b/tensor2tensor/rl/evaluator.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/evaluator_test.py b/tensor2tensor/rl/evaluator_test.py
index 6f90ad36e..76ebfb912 100644
--- a/tensor2tensor/rl/evaluator_test.py
+++ b/tensor2tensor/rl/evaluator_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/gym_utils.py b/tensor2tensor/rl/gym_utils.py
index 3fb2a708e..ef1f511ed 100644
--- a/tensor2tensor/rl/gym_utils.py
+++ b/tensor2tensor/rl/gym_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/gym_utils_test.py b/tensor2tensor/rl/gym_utils_test.py
index fece87b74..dbc85a024 100644
--- a/tensor2tensor/rl/gym_utils_test.py
+++ b/tensor2tensor/rl/gym_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/player.py b/tensor2tensor/rl/player.py
index 22af30be1..8d81ec474 100644
--- a/tensor2tensor/rl/player.py
+++ b/tensor2tensor/rl/player.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/player_utils.py b/tensor2tensor/rl/player_utils.py
index 0fbc3270e..6265533e3 100644
--- a/tensor2tensor/rl/player_utils.py
+++ b/tensor2tensor/rl/player_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/policy_learner.py b/tensor2tensor/rl/policy_learner.py
index fb628a493..67d037cf9 100644
--- a/tensor2tensor/rl/policy_learner.py
+++ b/tensor2tensor/rl/policy_learner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/ppo.py b/tensor2tensor/rl/ppo.py
index c458957cc..cf21a5b7d 100644
--- a/tensor2tensor/rl/ppo.py
+++ b/tensor2tensor/rl/ppo.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/ppo_learner.py b/tensor2tensor/rl/ppo_learner.py
index 3abf4255d..fb963232c 100644
--- a/tensor2tensor/rl/ppo_learner.py
+++ b/tensor2tensor/rl/ppo_learner.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/restarter.py b/tensor2tensor/rl/restarter.py
index e0dcbbe0d..88b2e6312 100644
--- a/tensor2tensor/rl/restarter.py
+++ b/tensor2tensor/rl/restarter.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/restarter_test.py b/tensor2tensor/rl/restarter_test.py
index 38fbe9eab..5e0826c23 100644
--- a/tensor2tensor/rl/restarter_test.py
+++ b/tensor2tensor/rl/restarter_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/rl_utils.py b/tensor2tensor/rl/rl_utils.py
index 5277c36c6..082b806bd 100644
--- a/tensor2tensor/rl/rl_utils.py
+++ b/tensor2tensor/rl/rl_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based.py b/tensor2tensor/rl/trainer_model_based.py
index c1ddb4d41..2d5e48dd7 100644
--- a/tensor2tensor/rl/trainer_model_based.py
+++ b/tensor2tensor/rl/trainer_model_based.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_agent_only.py b/tensor2tensor/rl/trainer_model_based_agent_only.py
index 804f0264b..9475f3227 100644
--- a/tensor2tensor/rl/trainer_model_based_agent_only.py
+++ b/tensor2tensor/rl/trainer_model_based_agent_only.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_params.py b/tensor2tensor/rl/trainer_model_based_params.py
index 91c3d08a2..0c5ee7120 100644
--- a/tensor2tensor/rl/trainer_model_based_params.py
+++ b/tensor2tensor/rl/trainer_model_based_params.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_recurrent_test.py b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
index 7e7e8b268..74fc24220 100644
--- a/tensor2tensor/rl/trainer_model_based_recurrent_test.py
+++ b/tensor2tensor/rl/trainer_model_based_recurrent_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_stochastic_test.py b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
index dd978ac98..0e2c1fe31 100644
--- a/tensor2tensor/rl/trainer_model_based_stochastic_test.py
+++ b/tensor2tensor/rl/trainer_model_based_stochastic_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_sv2p_test.py b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
index 24610310f..420331a4f 100644
--- a/tensor2tensor/rl/trainer_model_based_sv2p_test.py
+++ b/tensor2tensor/rl/trainer_model_based_sv2p_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_based_test.py b/tensor2tensor/rl/trainer_model_based_test.py
index 2600e3a25..20fecc190 100644
--- a/tensor2tensor/rl/trainer_model_based_test.py
+++ b/tensor2tensor/rl/trainer_model_based_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free.py b/tensor2tensor/rl/trainer_model_free.py
index 3bd83b027..543ac0655 100644
--- a/tensor2tensor/rl/trainer_model_free.py
+++ b/tensor2tensor/rl/trainer_model_free.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free_test.py b/tensor2tensor/rl/trainer_model_free_test.py
index b98f4c658..372bd6497 100644
--- a/tensor2tensor/rl/trainer_model_free_test.py
+++ b/tensor2tensor/rl/trainer_model_free_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/rl/trainer_model_free_tictactoe_test.py b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
index 66d545257..c7f429d46 100644
--- a/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
+++ b/tensor2tensor/rl/trainer_model_free_tictactoe_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/__init__.py b/tensor2tensor/serving/__init__.py
index ffe5892e3..ff174dd63 100644
--- a/tensor2tensor/serving/__init__.py
+++ b/tensor2tensor/serving/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/export.py b/tensor2tensor/serving/export.py
index 53d656f62..49e8dfb5b 100644
--- a/tensor2tensor/serving/export.py
+++ b/tensor2tensor/serving/export.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/query.py b/tensor2tensor/serving/query.py
index 7910d30f7..f4d05ffd1 100644
--- a/tensor2tensor/serving/query.py
+++ b/tensor2tensor/serving/query.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/serving/serving_utils.py b/tensor2tensor/serving/serving_utils.py
index 821f7c086..264c51d68 100644
--- a/tensor2tensor/serving/serving_utils.py
+++ b/tensor2tensor/serving/serving_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/test_data/example_usr_dir/__init__.py b/tensor2tensor/test_data/example_usr_dir/__init__.py
index 3f8cf243b..334f2b12b 100644
--- a/tensor2tensor/test_data/example_usr_dir/__init__.py
+++ b/tensor2tensor/test_data/example_usr_dir/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/test_data/example_usr_dir/my_submodule.py b/tensor2tensor/test_data/example_usr_dir/my_submodule.py
index 36609fa4a..c1133c895 100644
--- a/tensor2tensor/test_data/example_usr_dir/my_submodule.py
+++ b/tensor2tensor/test_data/example_usr_dir/my_submodule.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/__init__.py b/tensor2tensor/utils/__init__.py
index ffe5892e3..ff174dd63 100644
--- a/tensor2tensor/utils/__init__.py
+++ b/tensor2tensor/utils/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/adafactor.py b/tensor2tensor/utils/adafactor.py
index c4b5039cd..d54d9adf9 100644
--- a/tensor2tensor/utils/adafactor.py
+++ b/tensor2tensor/utils/adafactor.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/adafactor_test.py b/tensor2tensor/utils/adafactor_test.py
index d6832fe7d..924296866 100644
--- a/tensor2tensor/utils/adafactor_test.py
+++ b/tensor2tensor/utils/adafactor_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/adv_attack_utils.py b/tensor2tensor/utils/adv_attack_utils.py
index c31d4bfc7..472bf8ecd 100644
--- a/tensor2tensor/utils/adv_attack_utils.py
+++ b/tensor2tensor/utils/adv_attack_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/avg_checkpoints.py b/tensor2tensor/utils/avg_checkpoints.py
index f91b96ccd..e0be08ed4 100644
--- a/tensor2tensor/utils/avg_checkpoints.py
+++ b/tensor2tensor/utils/avg_checkpoints.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/beam_search.py b/tensor2tensor/utils/beam_search.py
index 2a2e6178f..43df9db90 100644
--- a/tensor2tensor/utils/beam_search.py
+++ b/tensor2tensor/utils/beam_search.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/beam_search_test.py b/tensor2tensor/utils/beam_search_test.py
index 60fcc09ef..d83f7c94d 100644
--- a/tensor2tensor/utils/beam_search_test.py
+++ b/tensor2tensor/utils/beam_search_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/bleu_hook.py b/tensor2tensor/utils/bleu_hook.py
index 0e92b8e41..9baadc4b9 100644
--- a/tensor2tensor/utils/bleu_hook.py
+++ b/tensor2tensor/utils/bleu_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/bleu_hook_test.py b/tensor2tensor/utils/bleu_hook_test.py
index 4f59ad1dd..47b607a3c 100644
--- a/tensor2tensor/utils/bleu_hook_test.py
+++ b/tensor2tensor/utils/bleu_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/checkpoint_compatibility_test.py b/tensor2tensor/utils/checkpoint_compatibility_test.py
index 0676ddc10..9a3e1d3eb 100644
--- a/tensor2tensor/utils/checkpoint_compatibility_test.py
+++ b/tensor2tensor/utils/checkpoint_compatibility_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/cloud_mlengine.py b/tensor2tensor/utils/cloud_mlengine.py
index 26301b7e3..337ffa249 100644
--- a/tensor2tensor/utils/cloud_mlengine.py
+++ b/tensor2tensor/utils/cloud_mlengine.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/compute_video_metrics.py b/tensor2tensor/utils/compute_video_metrics.py
index dbb793dbc..4e1765d6d 100644
--- a/tensor2tensor/utils/compute_video_metrics.py
+++ b/tensor2tensor/utils/compute_video_metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/contrib.py b/tensor2tensor/utils/contrib.py
index 78a3acd02..23441aeab 100644
--- a/tensor2tensor/utils/contrib.py
+++ b/tensor2tensor/utils/contrib.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index 25cc58868..c23e71083 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/data_reader_test.py b/tensor2tensor/utils/data_reader_test.py
index 451ba824c..ec93f6ac1 100644
--- a/tensor2tensor/utils/data_reader_test.py
+++ b/tensor2tensor/utils/data_reader_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/decoding.py b/tensor2tensor/utils/decoding.py
index 87631b1e3..383451012 100644
--- a/tensor2tensor/utils/decoding.py
+++ b/tensor2tensor/utils/decoding.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/devices.py b/tensor2tensor/utils/devices.py
index b595d7a3c..6c869984d 100644
--- a/tensor2tensor/utils/devices.py
+++ b/tensor2tensor/utils/devices.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/diet.py b/tensor2tensor/utils/diet.py
index 9039ced6a..67bd94afc 100644
--- a/tensor2tensor/utils/diet.py
+++ b/tensor2tensor/utils/diet.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/diet_test.py b/tensor2tensor/utils/diet_test.py
index ef7c5cf68..98df97fd9 100644
--- a/tensor2tensor/utils/diet_test.py
+++ b/tensor2tensor/utils/diet_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/expert_utils.py b/tensor2tensor/utils/expert_utils.py
index ddb4139d8..469ba5362 100644
--- a/tensor2tensor/utils/expert_utils.py
+++ b/tensor2tensor/utils/expert_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/expert_utils_test.py b/tensor2tensor/utils/expert_utils_test.py
index 79607b158..cc09ea404 100644
--- a/tensor2tensor/utils/expert_utils_test.py
+++ b/tensor2tensor/utils/expert_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/flags.py b/tensor2tensor/utils/flags.py
index c0f3a7079..e7e4fd8ba 100644
--- a/tensor2tensor/utils/flags.py
+++ b/tensor2tensor/utils/flags.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/get_rouge.py b/tensor2tensor/utils/get_rouge.py
index 614e5e37b..a042f6aed 100644
--- a/tensor2tensor/utils/get_rouge.py
+++ b/tensor2tensor/utils/get_rouge.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparam.py b/tensor2tensor/utils/hparam.py
index f198a5f66..b45634399 100644
--- a/tensor2tensor/utils/hparam.py
+++ b/tensor2tensor/utils/hparam.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparam_test.py b/tensor2tensor/utils/hparam_test.py
index 01778191c..51c07dba4 100644
--- a/tensor2tensor/utils/hparam_test.py
+++ b/tensor2tensor/utils/hparam_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparams_lib.py b/tensor2tensor/utils/hparams_lib.py
index 4882cd03f..14f60739b 100644
--- a/tensor2tensor/utils/hparams_lib.py
+++ b/tensor2tensor/utils/hparams_lib.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/hparams_lib_test.py b/tensor2tensor/utils/hparams_lib_test.py
index a068f910d..0ac0864b0 100644
--- a/tensor2tensor/utils/hparams_lib_test.py
+++ b/tensor2tensor/utils/hparams_lib_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/learning_rate.py b/tensor2tensor/utils/learning_rate.py
index 6151ad014..a3c700468 100644
--- a/tensor2tensor/utils/learning_rate.py
+++ b/tensor2tensor/utils/learning_rate.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index 0da9889c0..bfacc7184 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_hook.py b/tensor2tensor/utils/metrics_hook.py
index 93fa924f4..f5f006935 100644
--- a/tensor2tensor/utils/metrics_hook.py
+++ b/tensor2tensor/utils/metrics_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_hook_test.py b/tensor2tensor/utils/metrics_hook_test.py
index a2c9d290c..4744ec118 100644
--- a/tensor2tensor/utils/metrics_hook_test.py
+++ b/tensor2tensor/utils/metrics_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/metrics_test.py b/tensor2tensor/utils/metrics_test.py
index b9081edbd..1057b141f 100644
--- a/tensor2tensor/utils/metrics_test.py
+++ b/tensor2tensor/utils/metrics_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/misc_utils.py b/tensor2tensor/utils/misc_utils.py
index b5da5d26a..1e452b3fb 100644
--- a/tensor2tensor/utils/misc_utils.py
+++ b/tensor2tensor/utils/misc_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/misc_utils_test.py b/tensor2tensor/utils/misc_utils_test.py
index ccb453bb1..c988384c3 100644
--- a/tensor2tensor/utils/misc_utils_test.py
+++ b/tensor2tensor/utils/misc_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mlperf_log.py b/tensor2tensor/utils/mlperf_log.py
index 8fa42d802..0e89aabf9 100644
--- a/tensor2tensor/utils/mlperf_log.py
+++ b/tensor2tensor/utils/mlperf_log.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mlperf_tags.py b/tensor2tensor/utils/mlperf_tags.py
index 76f852185..1e882ec47 100644
--- a/tensor2tensor/utils/mlperf_tags.py
+++ b/tensor2tensor/utils/mlperf_tags.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/mtf_model.py b/tensor2tensor/utils/mtf_model.py
index b96d2b80c..08dfbf979 100644
--- a/tensor2tensor/utils/mtf_model.py
+++ b/tensor2tensor/utils/mtf_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_optimizer.py b/tensor2tensor/utils/multistep_optimizer.py
index 438d4cc64..916455b8d 100644
--- a/tensor2tensor/utils/multistep_optimizer.py
+++ b/tensor2tensor/utils/multistep_optimizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_optimizer_test.py b/tensor2tensor/utils/multistep_optimizer_test.py
index 9d3f6990e..5655acfcf 100644
--- a/tensor2tensor/utils/multistep_optimizer_test.py
+++ b/tensor2tensor/utils/multistep_optimizer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_with_adamoptimizer.py b/tensor2tensor/utils/multistep_with_adamoptimizer.py
index 8f00cfbdf..667f068f9 100644
--- a/tensor2tensor/utils/multistep_with_adamoptimizer.py
+++ b/tensor2tensor/utils/multistep_with_adamoptimizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/multistep_with_adamoptimizer_test.py b/tensor2tensor/utils/multistep_with_adamoptimizer_test.py
index a76bdaca1..5411b2d6c 100644
--- a/tensor2tensor/utils/multistep_with_adamoptimizer_test.py
+++ b/tensor2tensor/utils/multistep_with_adamoptimizer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py
index 6a11c943a..0725c52f0 100644
--- a/tensor2tensor/utils/optimize.py
+++ b/tensor2tensor/utils/optimize.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/optimize_test.py b/tensor2tensor/utils/optimize_test.py
index 938cd5c3d..3b191d6de 100644
--- a/tensor2tensor/utils/optimize_test.py
+++ b/tensor2tensor/utils/optimize_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/partial_checkpoint_load_hook.py b/tensor2tensor/utils/partial_checkpoint_load_hook.py
index a1242ab76..c2795e6cf 100644
--- a/tensor2tensor/utils/partial_checkpoint_load_hook.py
+++ b/tensor2tensor/utils/partial_checkpoint_load_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/pruning_utils.py b/tensor2tensor/utils/pruning_utils.py
index d22b1537c..b3624743a 100644
--- a/tensor2tensor/utils/pruning_utils.py
+++ b/tensor2tensor/utils/pruning_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/quantization.py b/tensor2tensor/utils/quantization.py
index d5c920b79..9c2eb748b 100644
--- a/tensor2tensor/utils/quantization.py
+++ b/tensor2tensor/utils/quantization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/registry.py b/tensor2tensor/utils/registry.py
index 31b0e70d8..cfe256366 100644
--- a/tensor2tensor/utils/registry.py
+++ b/tensor2tensor/utils/registry.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/registry_test.py b/tensor2tensor/utils/registry_test.py
index bb81fcf75..cee46aba2 100644
--- a/tensor2tensor/utils/registry_test.py
+++ b/tensor2tensor/utils/registry_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/restore_hook.py b/tensor2tensor/utils/restore_hook.py
index 28d959f0f..f8ad7e0fc 100644
--- a/tensor2tensor/utils/restore_hook.py
+++ b/tensor2tensor/utils/restore_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/rouge.py b/tensor2tensor/utils/rouge.py
index 246bf80f7..766bac018 100644
--- a/tensor2tensor/utils/rouge.py
+++ b/tensor2tensor/utils/rouge.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/rouge_test.py b/tensor2tensor/utils/rouge_test.py
index 3b27391b4..c2e3ca37b 100644
--- a/tensor2tensor/utils/rouge_test.py
+++ b/tensor2tensor/utils/rouge_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/sari_hook.py b/tensor2tensor/utils/sari_hook.py
index 1579ebe99..b992a1ae4 100644
--- a/tensor2tensor/utils/sari_hook.py
+++ b/tensor2tensor/utils/sari_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/sari_hook_test.py b/tensor2tensor/utils/sari_hook_test.py
index 250d0d6a2..caff0a07e 100644
--- a/tensor2tensor/utils/sari_hook_test.py
+++ b/tensor2tensor/utils/sari_hook_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/scheduled_sampling.py b/tensor2tensor/utils/scheduled_sampling.py
index 5c38f527f..8f556bd86 100644
--- a/tensor2tensor/utils/scheduled_sampling.py
+++ b/tensor2tensor/utils/scheduled_sampling.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index 542f83d81..995fb5982 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/t2t_model_test.py b/tensor2tensor/utils/t2t_model_test.py
index 578d22588..4ce17dd67 100644
--- a/tensor2tensor/utils/t2t_model_test.py
+++ b/tensor2tensor/utils/t2t_model_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/test_utils.py b/tensor2tensor/utils/test_utils.py
index 2adf5cdc3..4a823dea3 100644
--- a/tensor2tensor/utils/test_utils.py
+++ b/tensor2tensor/utils/test_utils.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/test_utils_test.py b/tensor2tensor/utils/test_utils_test.py
index ba8a3cdc0..4c79e7ceb 100644
--- a/tensor2tensor/utils/test_utils_test.py
+++ b/tensor2tensor/utils/test_utils_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 99bdeedcd..0dcbe7a81 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/trainer_lib_test.py b/tensor2tensor/utils/trainer_lib_test.py
index 7d1458325..a5bbecb9e 100644
--- a/tensor2tensor/utils/trainer_lib_test.py
+++ b/tensor2tensor/utils/trainer_lib_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/update_ops_hook.py b/tensor2tensor/utils/update_ops_hook.py
index a76671dfe..849216fe0 100644
--- a/tensor2tensor/utils/update_ops_hook.py
+++ b/tensor2tensor/utils/update_ops_hook.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/usr_dir.py b/tensor2tensor/utils/usr_dir.py
index 5ca60a883..b7a54ebcd 100644
--- a/tensor2tensor/utils/usr_dir.py
+++ b/tensor2tensor/utils/usr_dir.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video/prediction2gif.py b/tensor2tensor/utils/video/prediction2gif.py
index 50cf20662..f28674b4c 100644
--- a/tensor2tensor/utils/video/prediction2gif.py
+++ b/tensor2tensor/utils/video/prediction2gif.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video/reward_confusion.py b/tensor2tensor/utils/video/reward_confusion.py
index 11f69e6de..2b01f08c8 100644
--- a/tensor2tensor/utils/video/reward_confusion.py
+++ b/tensor2tensor/utils/video/reward_confusion.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video2gif.py b/tensor2tensor/utils/video2gif.py
index d9ae0fd1f..afbf694d2 100644
--- a/tensor2tensor/utils/video2gif.py
+++ b/tensor2tensor/utils/video2gif.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video_metrics.py b/tensor2tensor/utils/video_metrics.py
index aad285336..b22474c97 100644
--- a/tensor2tensor/utils/video_metrics.py
+++ b/tensor2tensor/utils/video_metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/video_metrics_test.py b/tensor2tensor/utils/video_metrics_test.py
index 7cd22012f..a7619d486 100644
--- a/tensor2tensor/utils/video_metrics_test.py
+++ b/tensor2tensor/utils/video_metrics_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/yellowfin.py b/tensor2tensor/utils/yellowfin.py
index 333d93c0a..d89641ec1 100644
--- a/tensor2tensor/utils/yellowfin.py
+++ b/tensor2tensor/utils/yellowfin.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/utils/yellowfin_test.py b/tensor2tensor/utils/yellowfin_test.py
index 09ff51123..fb0fc57f9 100644
--- a/tensor2tensor/utils/yellowfin_test.py
+++ b/tensor2tensor/utils/yellowfin_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/__init__.py b/tensor2tensor/visualization/__init__.py
index aa007c728..06080ebe9 100644
--- a/tensor2tensor/visualization/__init__.py
+++ b/tensor2tensor/visualization/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/attention.py b/tensor2tensor/visualization/attention.py
index aa90649b9..7ab271653 100644
--- a/tensor2tensor/visualization/attention.py
+++ b/tensor2tensor/visualization/attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/visualization.py b/tensor2tensor/visualization/visualization.py
index 8c1914a37..b700093fc 100644
--- a/tensor2tensor/visualization/visualization.py
+++ b/tensor2tensor/visualization/visualization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # 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/tensor2tensor/visualization/visualization_test.py b/tensor2tensor/visualization/visualization_test.py
index e7ad86171..ba78a8f9c 100644
--- a/tensor2tensor/visualization/visualization_test.py
+++ b/tensor2tensor/visualization/visualization_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Tensor2Tensor Authors.
+# Copyright 2023 The Tensor2Tensor Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.